diff --git a/.env.template b/.env.template index d9488574..92bdf93d 100644 --- a/.env.template +++ b/.env.template @@ -8,6 +8,7 @@ # Application Settings # ============================================================================= LOG_LEVEL=INFO +PERFORMANCE_LOG_FORMAT=compact # compact|rich # SESSION_OBSERVERS_LIMIT=10 # GET_CONTEXT_MAX_TOKENS=100000 # MAX_FILE_SIZE=5242880 # Bytes @@ -25,6 +26,7 @@ LOG_LEVEL=INFO # LANGFUSE_HOST= # LANGFUSE_PUBLIC_KEY= +# LANGFUSE_SECRET_KEY= # COLLECT_METRICS_LOCAL=false # LOCAL_METRICS_FILE=metrics.jsonl @@ -44,12 +46,15 @@ DB_CONNECTION_URI=postgresql+psycopg://postgres:postgres@localhost:5432/postgres # DB_POOL_CLASS=default # DB_POOL_SIZE=10 # DB_MAX_OVERFLOW=20 -# DB_POOL_TIMEOUT=30 +# DB_POOL_TIMEOUT=5 # seconds a pooled checkout waits for a free connection (QueuePool only) # DB_POOL_RECYCLE=300 # DB_POOL_PRE_PING=true # DB_POOL_USE_LIFO=true # DB_SQL_DEBUG=false # DB_TRACING=false +# Per-connection establish timeout (seconds) so a single connection attempt +# fails fast instead of hanging when the server/pooler is unreachable. +# DB_CONNECT_TIMEOUT_SECONDS=2 # ============================================================================= # Authentication Settings @@ -104,18 +109,31 @@ LLM_OPENAI_API_KEY=your-api-key-here # DERIVER_MODEL_CONFIG__OVERRIDES__BASE_URL=https://openrouter.ai/api/v1 # DERIVER_WORKERS=1 # DERIVER_POLLING_SLEEP_INTERVAL_SECONDS=1.0 +# Adaptive polling: grows the idle/error sleep from the base toward the max by +# the multiplier each cycle, snapping back to base when work is found. +# DERIVER_POLLING_BACKOFF_ENABLED=true +# DERIVER_POLLING_SLEEP_MAX_INTERVAL_SECONDS=30.0 +# DERIVER_POLLING_BACKOFF_MULTIPLIER=2.0 +# Jitter so instances that start together don't poll in lockstep. Startup: sleep +# a random delay in [0, value] before the first poll (0.0 disables). Per-cycle: +# multiply every poll sleep by a random factor in [1-ratio, 1+ratio] (0.0 disables). +# DERIVER_POLLING_STARTUP_JITTER_SECONDS=30.0 +# DERIVER_POLLING_JITTER_RATIO=0.5 # DERIVER_STALE_SESSION_TIMEOUT_MINUTES=5 # DERIVER_QUEUE_ERROR_RETENTION_SECONDS=2592000 # 30 days # DERIVER_MODEL_CONFIG__TEMPERATURE= # DERIVER_MODEL_CONFIG__THINKING_EFFORT=minimal # DERIVER_MODEL_CONFIG__THINKING_BUDGET_TOKENS=1024 # Gemini/Anthropic only +# DERIVER_MODEL_CONFIG__STRUCTURED_OUTPUT_MODE=json_object # for providers without json_schema support # DERIVER_DEDUPLICATE=true # DERIVER_MODEL_CONFIG__MAX_OUTPUT_TOKENS=4096 # DERIVER_LOG_OBSERVATIONS=false # DERIVER_MAX_INPUT_TOKENS=25000 # DERIVER_MAX_CUSTOM_INSTRUCTIONS_TOKENS=2000 # DERIVER_WORKING_REPRESENTATION_MAX_OBSERVATIONS=100 -# DERIVER_REPRESENTATION_BATCH_MAX_TOKENS=1024 +# DERIVER_REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS=512 # Min tokens a work unit accumulates before the deriver claims it; 0 disables the gate +# DERIVER_REPRESENTATION_BATCH_TARGET_INPUT_TOKENS=1024 # Max context-window tokens per deriver LLM call +# DERIVER_REPRESENTATION_BATCH_MAX_AGE_SECONDS=1800 # DERIVER_FLUSH_ENABLED=false # Bypass batch token threshold, process work immediately # DERIVER_MODEL_CONFIG__FALLBACK__MODEL= # DERIVER_MODEL_CONFIG__FALLBACK__TRANSPORT= @@ -157,12 +175,17 @@ LLM_OPENAI_API_KEY=your-api-key-here # DIALECTIC_LEVELS__max__MODEL_CONFIG__TRANSPORT=openai # DIALECTIC_LEVELS__max__MODEL_CONFIG__MODEL=gpt-5.4-mini # DIALECTIC_LEVELS__max__MAX_TOOL_ITERATIONS=10 -# Optional overrides: +# Optional overrides (model and OpenAI-compatible base URL are per-level): # DIALECTIC_LEVELS__minimal__MODEL_CONFIG__MODEL=your-model-here +# DIALECTIC_LEVELS__minimal__MODEL_CONFIG__OVERRIDES__BASE_URL=https://openrouter.ai/api/v1 # DIALECTIC_LEVELS__low__MODEL_CONFIG__MODEL=your-model-here +# DIALECTIC_LEVELS__low__MODEL_CONFIG__OVERRIDES__BASE_URL=https://openrouter.ai/api/v1 # DIALECTIC_LEVELS__medium__MODEL_CONFIG__MODEL=your-model-here +# DIALECTIC_LEVELS__medium__MODEL_CONFIG__OVERRIDES__BASE_URL=https://openrouter.ai/api/v1 # DIALECTIC_LEVELS__high__MODEL_CONFIG__MODEL=your-model-here +# DIALECTIC_LEVELS__high__MODEL_CONFIG__OVERRIDES__BASE_URL=https://openrouter.ai/api/v1 # DIALECTIC_LEVELS__max__MODEL_CONFIG__MODEL=your-model-here +# DIALECTIC_LEVELS__max__MODEL_CONFIG__OVERRIDES__BASE_URL=https://openrouter.ai/api/v1 # DIALECTIC_LEVELS__max__MODEL_CONFIG__THINKING_EFFORT=medium # DIALECTIC_LEVELS__max__MODEL_CONFIG__THINKING_BUDGET_TOKENS=1024 # Optional backup per level (must set both or neither): @@ -254,14 +277,28 @@ LLM_OPENAI_API_KEY=your-api-key-here # TELEMETRY_MAX_BUFFER_SIZE=10000 # TELEMETRY_NAMESPACE=honcho # Inherits from NAMESPACE if not set +# Full-fidelity payload tracing (llm.call.traced / trace.content). Default-off +# TELEMETRY_TRACE_PAYLOADS_ENABLED=false # Trace events ship to TELEMETRY_ENDPOINT +# TELEMETRY_TRACE_MAX_BYTES=262144 # Per-message cap; oversized content is clipped +# TELEMETRY_TRACE_PURPOSES=[] # JSON list of CallPurpose values to capture; empty = all + # ============================================================================= # Cache # ============================================================================= # CACHE_ENABLED=false # CACHE_URL="redis://localhost:6379/0?suppress=true" +# CACHE_CLUSTER=false # true when CACHE_URL is a Redis Cluster (e.g. Memorystore for Redis Cluster) # CACHE_NAMESPACE="honcho" # Inherits from NAMESPACE if not set # CACHE_DEFAULT_TTL_SECONDS=300 # CACHE_DEFAULT_LOCK_TTL_SECONDS=5 +# CACHE_LOCK_WAIT_CHECK_INTERVAL_SECONDS=0.1 + +# ============================================================================= +# CORS Settings +# ============================================================================= +# JSON array of origins allowed by the FastAPI CORSMiddleware. Defaults match +# the previously hardcoded list: localhost, 127.0.0.1:8000 and api.honcho.dev. +# CORS_ORIGINS=["http://localhost","http://127.0.0.1:8000","https://api.honcho.dev"] # ============================================================================= # Vector Store Settings diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 00000000..039d4e6e --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,19 @@ +# Code owners for Honcho. +# +# Beyond review routing, this file is the allowlist for manually triggering +# the live-llm-tests and unified-tests workflows (via their PR labels and +# workflow_dispatch). The gate jobs grep every @username in this file — +# regardless of which path pattern it sits on — and read it from `main`, +# never from the PR branch, so additions only take effect once merged. +# +# The workflow gates only understand individual @usernames (no @org/team +# entries). + +# Reviewers auto-requested on changes under .github/ (workflows, this file, +# templates). +/.github/ @akattelu @eisene @Rajat-Ahuja1997 @VVoruganti + +# CI-trigger allowlist only: this path matches no real file, so these people +# are never auto-requested for review, but the workflow gates still pick +# them up. +/ci-trigger-allowlist @3un01a @adavyas @ajspig @courtlandleer @erosika @lowyelling @matthewlanders @vintrocode diff --git a/.github/actions/load-staging-secrets/action.yml b/.github/actions/load-staging-secrets/action.yml new file mode 100644 index 00000000..64665ec5 --- /dev/null +++ b/.github/actions/load-staging-secrets/action.yml @@ -0,0 +1,84 @@ +name: Load staging secrets +description: >- + Resolve staging secret ids from the two newest v git tags, + then load the newest fetchable secret's keys into the job environment from + AWS Secrets Manager. Falls back to the second-latest tag when the latest + tag's secret isn't published yet, and fails the job loudly when neither can + be fetched. Requires the repository to be checked out and AWS credentials to + be configured beforehand. + +inputs: + secret-prefix: + description: >- + Secret-name prefix combined with a resolved tag version to form the full + secret id. Masked so it stays out of public CI logs. + required: true + +runs: + using: composite + steps: + - name: Resolve secret ids from latest git tags + id: resolve-secret + shell: bash + env: + SECRET_PREFIX: ${{ inputs.secret-prefix }} + run: | + set -euo pipefail + : "${SECRET_PREFIX:?secret-prefix input is empty — is the STAGING_SECRET_PREFIX secret set for this environment?}" + # Keep the secret-name prefix out of public CI logs. + echo "::add-mask::${SECRET_PREFIX}" + + # Two newest v tags, highest first (tags are public). + versions="$(git ls-remote --tags origin 'v*' \ + | sed -n 's#.*refs/tags/v\([0-9][0-9]*\.[0-9][0-9]*\.[0-9][0-9]*\)$#\1#p' \ + | sort -t. -k1,1nr -k2,2nr -k3,3nr -u)" + latest="$(printf '%s\n' "$versions" | sed -n '1p')" + second="$(printf '%s\n' "$versions" | sed -n '2p')" + if [ -z "${latest:-}" ]; then + echo "::error::No v git tags found to resolve a secret version" + exit 1 + fi + + latest_id="${SECRET_PREFIX}${latest}" + echo "::add-mask::${latest_id}" + echo "latest-id=${latest_id}" >> "$GITHUB_OUTPUT" + echo "Latest version: ${latest}" + if [ -n "${second:-}" ]; then + second_id="${SECRET_PREFIX}${second}" + echo "::add-mask::${second_id}" + echo "second-id=${second_id}" >> "$GITHUB_OUTPUT" + echo "Fallback version: ${second}" + fi + + # Fetch the latest tag's secret. continue-on-error so a not-yet-published + # latest falls through to the second-latest instead of failing the job. + - name: Fetch staging secret (latest) + id: fetch-latest + continue-on-error: true + uses: aws-actions/aws-secretsmanager-get-secrets@v2 + with: + secret-ids: | + ,${{ steps.resolve-secret.outputs.latest-id }} + parse-json-secrets: true + + # Runs only if the latest fetch failed; this one is NOT continue-on-error, + # so if the fallback also fails the job fails loudly. + - name: Fetch staging secret (fallback to second-latest) + id: fetch-fallback + if: steps.fetch-latest.outcome == 'failure' && steps.resolve-secret.outputs.second-id != '' + uses: aws-actions/aws-secretsmanager-get-secrets@v2 + with: + secret-ids: | + ,${{ steps.resolve-secret.outputs.second-id }} + parse-json-secrets: true + + # If the latest fetch failed and the fallback was skipped (no second tag), + # no secret keys were loaded — fail here instead of letting the job run + # without staging config (e.g. every live LLM test would silently skip via + # require_provider_key and the run would go green). + - name: Verify staging secrets were loaded + if: steps.fetch-latest.outcome != 'success' && steps.fetch-fallback.outcome != 'success' + shell: bash + run: | + echo "::error::No staging secret could be fetched (latest failed; fallback skipped or failed)" + exit 1 diff --git a/.github/workflows/live-llm-tests.yml b/.github/workflows/live-llm-tests.yml new file mode 100644 index 00000000..ee1b5b2b --- /dev/null +++ b/.github/workflows/live-llm-tests.yml @@ -0,0 +1,123 @@ +name: Live LLM Tests + +on: + # Runs on main pushes that can affect the LLM transport (narrower than + # unified-tests' src/** — live provider calls aren't worth burning on + # changes that can't reach the backends). + push: + branches: [main] + paths: + - 'src/llm/**' + - 'src/config.py' + - 'tests/live_llm/**' + - 'pyproject.toml' + - 'uv.lock' + - '.github/workflows/live-llm-tests.yml' + # Manual trigger for PRs: add the `run-live-llm` label to run the suite + # against the PR's merge commit. The label is purged as soon as the run + # starts so it can be re-added to trigger another run. + pull_request: + types: [labeled] + workflow_dispatch: + +# Cap spend: at most one active run per PR (per ref for push/dispatch). +# Re-triggering a PR run cancels the in-flight one instead of stacking live +# provider calls; pushes to main queue instead of cancelling so main CI +# results aren't lost. +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: ${{ github.event_name != 'push' }} + +permissions: + contents: read + +jobs: + # Only code owners (.github/CODEOWNERS) may trigger the suite manually via + # the label or workflow_dispatch; the gate also purges the trigger label so + # it can be re-added to trigger another run. + gate: + name: Gate manual trigger + permissions: + contents: read + pull-requests: write + uses: ./.github/workflows/manual-trigger-gate.yml + with: + label: run-live-llm + allow-workflow-dispatch: true + + live-llm-tests: + name: Run Live LLM Tests + needs: gate + # always() lets this run on push events, where the gate's jobs are skipped. + # Manual triggers (label / workflow_dispatch) additionally require the + # gate's CODEOWNERS check to have passed. + if: >- + always() && + (github.event_name == 'push' || + ((github.event_name == 'workflow_dispatch' || + github.event.label.name == 'run-live-llm') && + needs.gate.outputs.authorized == 'true')) + runs-on: ubuntu-latest + timeout-minutes: 20 + environment: unified-tests + permissions: + id-token: write # Required for OIDC authentication with AWS + contents: read + env: + PYTHONUNBUFFERED: "1" + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + ref: ${{ github.sha }} + + - name: Configure AWS credentials + uses: aws-actions/configure-aws-credentials@v4 + with: + role-to-assume: ${{ vars.AWS_OIDC_ROLE_ARN }} + aws-region: us-east-1 + role-duration-seconds: 3600 + + # Resolves secret ids from the newest release tags, fetches the newest + # available staging secret into the job env, and fails if none loaded. + - name: Load staging secrets + uses: ./.github/actions/load-staging-secrets + with: + secret-prefix: ${{ secrets.STAGING_SECRET_PREFIX }} + + # Configure the test environment. Sentry/CloudEvents endpoints aren't + # reachable from CI. LIVE_LLM_ANTHROPIC_45_PLUS_MODELS must be set for + # the Anthropic tests to materialize — the claude_4_5_plus family has no + # default models, so with only the API key they'd silently collect as + # empty parameter sets. + - name: Configure test environment + run: | + { + # The staging dotenv carries AUTH_USE_AUTH=true without a usable + # JWT secret; src/config.py validates the pair at import time, so + # disable auth (this suite never runs the API server anyway). + echo "AUTH_USE_AUTH=false" + echo "SENTRY_ENABLED=false" + echo "TELEMETRY_ENABLED=false" + echo "LIVE_LLM_ANTHROPIC_45_PLUS_MODELS=claude-sonnet-4-5" + } >> "$GITHUB_ENV" + + - name: Install uv + uses: astral-sh/setup-uv@v2 + with: + enable-cache: true + cache-dependency-glob: "uv.lock" + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version-file: "pyproject.toml" + + - name: Install the project + run: uv sync --all-extras + + # -n 0 overrides the `-n auto` xdist default from pyproject: ~15 short + # tests gain nothing from parallelism, serial execution avoids bursting + # every provider at once, and flake diagnosis gets ordered output. + - name: Run live LLM tests + run: uv run --frozen pytest tests/live_llm/ --live-llm -n 0 -v diff --git a/.github/workflows/manual-trigger-gate.yml b/.github/workflows/manual-trigger-gate.yml new file mode 100644 index 00000000..6cd8b209 --- /dev/null +++ b/.github/workflows/manual-trigger-gate.yml @@ -0,0 +1,81 @@ +name: Manual Trigger Gate + +# Shared gate for workflows that can be triggered manually on PRs by adding a +# label (and optionally via workflow_dispatch): verifies the actor is a code +# owner and purges the trigger label so it can be re-added for another run. +# +# Callers must grant `pull-requests: write` on the calling job so the +# remove-label job can delete the label, and should gate downstream jobs on +# the `authorized` output rather than this workflow's conclusion. + +on: + workflow_call: + inputs: + label: + description: PR label that triggers the calling workflow + required: true + type: string + allow-workflow-dispatch: + description: Whether workflow_dispatch events may pass the gate + required: false + default: false + type: boolean + outputs: + authorized: + description: >- + 'true' when the manual trigger's actor passed the CODEOWNERS check. + Empty on events where the check did not run (e.g. push). + value: ${{ jobs.check-actor.outputs.authorized }} + +jobs: + # Only code owners (.github/CODEOWNERS) may trigger the calling workflow + # manually. + check-actor: + name: Verify actor is a code owner + if: >- + (inputs.allow-workflow-dispatch && github.event_name == 'workflow_dispatch') || + (github.event_name == 'pull_request' && github.event.label.name == inputs.label) + runs-on: ubuntu-latest + permissions: + contents: read + outputs: + authorized: ${{ steps.codeowners.outputs.authorized }} + steps: + - name: Check actor against CODEOWNERS on main + id: codeowners + env: + GH_TOKEN: ${{ github.token }} + ACTOR: ${{ github.actor }} + run: | + set -euo pipefail + # Usernames are case-insensitive on GitHub; compare lowercased. + owners="$(gh api -H "Accept: application/vnd.github.raw" \ + "repos/${{ github.repository }}/contents/.github/CODEOWNERS?ref=main" \ + | sed 's/#.*//' | grep -oE '@[A-Za-z0-9-]+' | tr -d '@' \ + | tr '[:upper:]' '[:lower:]' | sort -u)" + actor_lc="$(printf '%s' "$ACTOR" | tr '[:upper:]' '[:lower:]')" + if printf '%s\n' "$owners" | grep -qxF "$actor_lc"; then + echo "@${ACTOR} is a code owner; proceeding" + echo "authorized=true" >> "$GITHUB_OUTPUT" + else + echo "::error::@${ACTOR} is not listed in .github/CODEOWNERS on main — only code owners may trigger this workflow manually" + exit 1 + fi + + # Purge the trigger label first thing. Best-effort: failing to remove the + # label (e.g. read-only token on a fork PR) doesn't block the tests. + remove-label: + name: Remove trigger label + if: github.event_name == 'pull_request' && github.event.label.name == inputs.label + runs-on: ubuntu-latest + permissions: + pull-requests: write + steps: + - name: Remove trigger label + env: + GH_TOKEN: ${{ github.token }} + run: | + if ! gh api --method DELETE \ + "repos/${{ github.repository }}/issues/${{ github.event.pull_request.number }}/labels/${{ inputs.label }}"; then + echo "::warning::Could not remove the ${{ inputs.label }} label (it may have been removed already)" + fi diff --git a/.github/workflows/push-gcp-registry-prod.yml b/.github/workflows/push-gcp-registry-prod.yml new file mode 100644 index 00000000..d5d9fc42 --- /dev/null +++ b/.github/workflows/push-gcp-registry-prod.yml @@ -0,0 +1,56 @@ +name: Build and Push to GCP Artifact Registry (production) + +permissions: + contents: read + +on: + push: + tags: + - v* + workflow_dispatch: + inputs: + version: + description: "Version to deploy (without v prefix)" + required: true + type: string + default: "manual" + +env: + GCP_PROJECT_ID: ${{ secrets.PROD_GCP_PROJECT_ID }} + GCP_AR_LOCATION: ${{ secrets.PROD_GCP_AR_LOCATION }} + GCP_AR_REPO: ${{ secrets.PROD_GCP_AR_REPO }} + IMAGE_NAME: ${{ secrets.PROD_IMAGE_NAME }} + GCP_SA_KEY: ${{ secrets.PROD_GCP_SA_KEY }} + +jobs: + build-and-push: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Authenticate to GCP + uses: google-github-actions/auth@v2 + with: + credentials_json: ${{ env.GCP_SA_KEY }} + + - name: Set up Cloud SDK + uses: google-github-actions/setup-gcloud@v2 + + - name: Configure Docker for Artifact Registry + run: gcloud auth configure-docker ${{ env.GCP_AR_LOCATION }}-docker.pkg.dev --quiet + + - name: Build and push image + env: + VERSION: ${{ github.event.inputs.version }} + run: | + # Determine the image label based on trigger type + if [[ "$GITHUB_EVENT_NAME" == "workflow_dispatch" ]]; then + IMAGE_LABEL="deployment-${VERSION}" + else + IMAGE_LABEL="deployment-${GITHUB_REF_NAME}" + fi + BASE="${{ env.GCP_AR_LOCATION }}-docker.pkg.dev/${{ env.GCP_PROJECT_ID }}/${{ env.GCP_AR_REPO }}/${{ env.IMAGE_NAME }}" + TAG="$BASE:$IMAGE_LABEL" + docker build -t "$TAG" . + docker push "$TAG" diff --git a/.github/workflows/push-gcp-registry-staging.yml b/.github/workflows/push-gcp-registry-staging.yml new file mode 100644 index 00000000..f41a38de --- /dev/null +++ b/.github/workflows/push-gcp-registry-staging.yml @@ -0,0 +1,56 @@ +name: Build and Push to GCP Artifact Registry (Staging) + +permissions: + contents: read + +on: + push: + tags: + - v* + workflow_dispatch: + inputs: + version: + description: "Version to deploy (without v prefix)" + required: true + type: string + default: "manual" + +env: + GCP_PROJECT_ID: ${{ secrets.STAGING_GCP_PROJECT_ID }} + GCP_AR_LOCATION: ${{ secrets.STAGING_GCP_AR_LOCATION }} + GCP_AR_REPO: ${{ secrets.STAGING_GCP_AR_REPO }} + IMAGE_NAME: ${{ secrets.STAGING_IMAGE_NAME }} + GCP_SA_KEY: ${{ secrets.STAGING_GCP_SA_KEY }} + +jobs: + build-and-push: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Authenticate to GCP + uses: google-github-actions/auth@v2 + with: + credentials_json: ${{ env.GCP_SA_KEY }} + + - name: Set up Cloud SDK + uses: google-github-actions/setup-gcloud@v2 + + - name: Configure Docker for Artifact Registry + run: gcloud auth configure-docker ${{ env.GCP_AR_LOCATION }}-docker.pkg.dev --quiet + + - name: Build and push image + env: + VERSION: ${{ github.event.inputs.version }} + run: | + # Determine the image label based on trigger type + if [[ "$GITHUB_EVENT_NAME" == "workflow_dispatch" ]]; then + IMAGE_LABEL="deployment-${VERSION}" + else + IMAGE_LABEL="deployment-${GITHUB_REF_NAME}" + fi + BASE="${{ env.GCP_AR_LOCATION }}-docker.pkg.dev/${{ env.GCP_PROJECT_ID }}/${{ env.GCP_AR_REPO }}/${{ env.IMAGE_NAME }}" + TAG="$BASE:$IMAGE_LABEL" + docker build -t "$TAG" . + docker push "$TAG" diff --git a/.github/workflows/unified-tests.yml b/.github/workflows/unified-tests.yml index bfc27018..4813a126 100644 --- a/.github/workflows/unified-tests.yml +++ b/.github/workflows/unified-tests.yml @@ -6,14 +6,50 @@ on: paths: - 'src/**' - 'tests/**' + # Manual trigger for PRs: add the `run-unified-tests` label to run the suite + # against the PR's merge commit. The label is purged as soon as the run + # starts so it can be re-added to trigger another run. + pull_request: + types: [labeled] + +# Cap spend: at most one active run per PR (per ref for push). Re-triggering +# a PR run cancels the in-flight one instead of stacking Fly machines; pushes +# to main queue instead of cancelling so main CI results aren't lost. The +# cleanup-machine job runs `if: always()`, which still executes on cancelled +# runs, so a cancelled run's Fly machine and runner are still torn down. +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: ${{ github.event_name != 'push' }} permissions: contents: read actions: read jobs: + # Only code owners (.github/CODEOWNERS) may trigger the suite manually via + # the label; the gate also purges the trigger label so it can be re-added + # to trigger another run. + gate: + name: Gate manual trigger + permissions: + contents: read + pull-requests: write + uses: ./.github/workflows/manual-trigger-gate.yml + with: + label: run-unified-tests + start-runner: name: Start Fly Runner + needs: gate + # always() lets this run on push events, where the gate's jobs are skipped. + # Label adds other than run-unified-tests trigger the workflow but skip + # every job here; the run-unified-tests label additionally requires the + # gate's CODEOWNERS check to have passed. + if: >- + always() && + (github.event_name == 'push' || + (github.event.label.name == 'run-unified-tests' && + needs.gate.outputs.authorized == 'true')) uses: ./.github/workflows/start-fly-runner.yml secrets: inherit @@ -40,17 +76,72 @@ jobs: - name: Configure AWS credentials uses: aws-actions/configure-aws-credentials@v4 with: - role-to-assume: arn:aws:iam::444554165670:role/GitHubActionsS3Role + role-to-assume: ${{ vars.AWS_OIDC_ROLE_ARN }} aws-region: us-east-1 role-duration-seconds: 43200 # 12 hours - - name: Fetch secrets from AWS Secrets Manager + # Resolves secret ids from the newest release tags, fetches the newest + # available staging secret into the job env, and fails if none loaded. + - name: Load staging secrets + uses: ./.github/actions/load-staging-secrets + with: + secret-prefix: ${{ secrets.STAGING_SECRET_PREFIX }} + + # Layer test-specific overrides on top of the staging secret. The staging + # dotenv tracks the deployed release and can drift from what main's config + # expects; the TESTING_SECRET_ID secret holds only the keys (flat JSON, + # exact env var names) the unified tests need to pin. The get-secrets + # action refuses to inject an env var that already exists, so the + # overrides are fetched under a prefix alias here and promoted over the + # staging values in the next step. + - name: Fetch testing secret overrides uses: aws-actions/aws-secretsmanager-get-secrets@v2 with: secret-ids: | - ,testing/unified/tests + HONCHO_TEST_OVERRIDE,${{ secrets.TESTING_SECRET_ID }} parse-json-secrets: true + # Re-export each HONCHO_TEST_OVERRIDE_* var under its real name; the + # later $GITHUB_ENV write wins over the value loaded from the staging + # secret. Values are already masked by the fetch step above. + - name: Apply testing secret overrides + run: | + set -euo pipefail + applied=0 + while IFS= read -r -d '' entry; do + name="${entry%%=*}" + value="${entry#*=}" + case "$name" in + HONCHO_TEST_OVERRIDE_*) + target="${name#HONCHO_TEST_OVERRIDE_}" + { + echo "${target}<<__HONCHO_OVERRIDE_EOF__" + printf '%s\n' "$value" + echo "__HONCHO_OVERRIDE_EOF__" + } >> "$GITHUB_ENV" + echo "Overriding ${target}" + applied=$((applied + 1)) + ;; + esac + done < <(env -0) + echo "Applied ${applied} override(s)" + + # Configure the test environment. Disables auth/Sentry/CloudEvents telemetry + # (their endpoints aren't reachable from CI), and points REASONING_TRACES_FILE + # at a shared path so the API + deriver record full LLM I/O for auditing — the + # runner uploads it to S3. Written after the fetch steps so these win over the + # values loaded from Secrets Manager (last $GITHUB_ENV write wins). Stale + # config keys loaded from the staging secret (e.g. settings that have since + # been renamed or removed on main) must always be ignored by the app config. + - name: Configure test environment + run: | + { + echo "AUTH_USE_AUTH=false" + echo "SENTRY_ENABLED=false" + echo "TELEMETRY_ENABLED=false" + echo "REASONING_TRACES_FILE=unified-reasoning-traces.jsonl" + } >> "$GITHUB_ENV" + - name: Verify Docker is available run: docker info @@ -110,7 +201,7 @@ jobs: exit 0 fi - RUNNER_ID=""  + RUNNER_ID="" if [ -n "$RUNNER_NAME" ]; then RUNNER_ID=$(echo "$RUNNERS_RESPONSE" | jq -r --arg name "$RUNNER_NAME" '.runners[]? | select(.name == $name) | .id') fi diff --git a/.gitignore b/.gitignore index e6cbf1d7..9fa90b3e 100644 --- a/.gitignore +++ b/.gitignore @@ -193,3 +193,6 @@ metrics.jsonl AGENTS.md lancedb_data/ grafana-data/ + +# Claude Code addon stuff +.omc diff --git a/CHANGELOG.md b/CHANGELOG.md index 6fb883e2..7a12356e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,94 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](http://keepachangelog.com/) and this project adheres to [Semantic Versioning](http://semver.org/). +## [3.0.11] - 2026-06-24 + +### Added + +- `api_request_duration_seconds` Prometheus histogram tracking per-route request latency, labeled by method and endpoint (#837) +- LLM `provider_params` passthroughs (`extra_body` / `extra_headers` / `extra_query`) are now forwarded to the underlying provider transport across all backends, with shape validation that rejects non-mapping values (#821) +- `structured_output_mode` model-config option to use `json_object` mode for OpenAI-compatible providers that lack native Structured Outputs support (used by the deriver) (#820) +- OpenRouter app-attribution headers (`HTTP-Referer` / `X-Openrouter-Title`) are now sent on OpenAI-compatible clients when the configured base URL is OpenRouter, so requests are attributed to "Honcho" in OpenRouter's dashboard (#805) +- Langfuse traces are now tagged with user and session IDs for easier trace filtering (#814) +- `DERIVER_REPRESENTATION_BATCH_MAX_AGE_SECONDS` (default 1800s) lets sub-threshold representation work units flush once their oldest unprocessed queue item ages out. Set it to `0` to keep the legacy behavior where sub-threshold tails wait indefinitely unless `DERIVER_FLUSH_ENABLED=true` (#826) +- Conclusion responses now include a `level` field (`explicit`, `deductive`, `inductive`, `contradiction`); list/query endpoints support filtering by `level` via `filters`, with reserved filter keys protected from being overridden by user-supplied filters (#851) + +### Changed + +- Peer-scoped JWTs now get read-only access to the sessions their peer is an active member of (session context, summaries, peers, their own per-session config, search, and message reads). Session-scoped JWTs remain confined to their session and cannot reach peer routes (#679) +- Compacted Honcho's log output, with guarded ms/s metric formatting that falls back to a plain string for non-numeric values (#836) +- Sentry now drops noisy infra/scrape transactions: the reconciler opens a transaction only once a batch has rows (idle cycles emit none), and a `traces_sampler` returns `0.0` for `/metrics`, `/health`, `/openapi.json`, `/docs`, `/redoc`, and the deriver metrics server. `SENTRY.TRACES_SAMPLE_RATE` still governs real traffic (#834) + +### Fixed + +- Peer- and session-scoped JWTs were effectively workspace-scoped: authorization walked the route's declared scope and fell through to a workspace match, so a `{w, p: alice}` token could act on any peer in the workspace. JWTs are now authorized by their narrowest claim and never widen to workspace access (#679) +- The keys API now rejects creating a peer- or session-scoped key without a workspace. Such keys were minted successfully but failed verification on every request (#679) +- Agent-supplied observation IDs carrying the display-format `id:` prefix are now normalized (prefix and trailing whitespace stripped) before `source_ids` are stored and on `get_reasoning_chain` lookups, fixing corrupted provenance links and broken reasoning-chain traversal (#795) +- Fixed a `create_tree` keyword-argument mismatch in the Dreamer's surprisal tree construction (#749) +- Providers that omit output-token counts (observed with Gemini on tool-loop completions) returned `output_tokens=None`, which raised a Pydantic validation error that aborted the call and crashed the Dreamer's induction phase before inductive conclusions were persisted. `None` is now coerced to `0` so token accounting degrades gracefully (#809) +- Document creation now performs exact (case-insensitive, whitespace-trimmed) content deduplication before the existing semantic dedup step: exact duplicates within a batch collapse to a single insert, and an exact match against a live document reinforces it (atomic `times_derived` increment) instead of creating a new row (#861) + +## [3.0.10] - 2026-06-15 + +### Added + +- Messages are now embedded via a background task rather than blocking API request +- Read-only DB session mode (`get_read_db` / `tracked_db(..., read_only=True)`) so reads don't hold a transaction open across the work +- `CORS_ORIGINS` env var to configure CORS allowed origins without editing source; defaults match the prior hardcoded list, so self-hosted deployments behind custom domains can whitelist their frontend (#697) +- `scripts/generate_jwt.py` — utility for minting scoped or admin Honcho JWTs (`--admin`, `--workspace`/`--peer`/`--session`, `--expires` with human-friendly durations, `--print-only`) without calling the keys API (#757) +- `STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS` (default 60s) — minimum jittered spacing between deriver stale-work-unit cleanup runs, so cleanup no longer runs on every seconds-scale poll (`0.0` keeps the legacy every-poll behavior) (#773) + +### Changed + +- Optimized the deriver and dreamer prompt cache prefixes to improve prompt-cache hit rates (#806) + +### Fixed + +- `times_derived` is now properly reinforced when a duplicate conclusion is detected. It had been pinned at 1 for nearly every conclusion (the reject-new branch dropped the increment and the new-wins branch reset the count to 1), so `ORDER BY times_derived DESC` fell back to arbitrary heap order and froze stale conclusions to the front of injected context. Reinforcement is now an atomic increment and both most-derived queries gained a `created_at DESC` recency tiebreaker (#768) +- Webhook creation now correctly rejects private/internal IP addresses (#793) + +## [3.0.9] - 2026-06-02 + +### Changed + +- Connection acquisition is now a single attempt with no server-side retry, on a vanilla `AsyncSession`. A new `DB_CONNECT_TIMEOUT_SECONDS` (default 2s) bounds the attempt so a saturated or unreachable pooler fails fast instead of holding a client connection open to re-knock. A saturated DB now surfaces to the caller — the API returns an error and the deriver backs off and retries on a later poll — which lets the pooler drain rather than amplifying saturation. + +### Added + +- Deriver poll jitter so instances that start together don't poll in lockstep: `DERIVER_POLLING_STARTUP_JITTER_SECONDS` (random delay before the first poll, default 30s) and `DERIVER_POLLING_JITTER_RATIO` (±fraction applied to every poll sleep, default 0.5). Both disable at `0.0`; the underlying backoff schedule is unchanged. + +### Removed + +- Reverted the connection-checkout retry and `HonchoAsyncSession` custom session introduced in 3.0.8. Removed the `DB_CONNECTION_RETRY_ENABLED` / `DB_CONNECTION_RETRY_MAX_DELAY_SECONDS` / `DB_CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS` / `DB_CONNECTION_RETRY_BACKOFF_MAX_SECONDS` settings, the `db_connection_acquisitions{outcome=...}` Prometheus counter, and the `db.pool.acquire` Sentry span. Alerting built on `db_connection_acquisitions` should migrate to `db_pool_connections` / `db_queries_in_flight`. + +## [3.0.8] - 2026-06-01 + +### Added + +- Connection-checkout retry with bounded exponential backoff (tenacity) on `get_db`/`tracked_db`: transient transaction-pooler (Supavisor) rejections — SQLAlchemy `TimeoutError` and `OperationalError` — now retry with backoff instead of surfacing as 500s under client-connection saturation. Gated by + `DB_CONNECTION_RETRY_ENABLED` with configurable delay/backoff knobs; ~10s default budget (#758) +- `HonchoAsyncSession` — a lazy `AsyncSession` that checks out its pooled connection (with retry) on the first DB-touching call rather than at construction. Request handlers doing non-DB work (embedding, file, LLM) before their first query no longer pin a pooler connection across it. Only the checkout is retried; + the statement still runs exactly once, so writes are never duplicated (#758) +- Adaptive deriver queue polling: the poll interval backs off when the queue is idle or erroring (base → max, doubling each cycle) and snaps back to base the moment work is claimed, cutting steady-state query load against the DB. Gated by `DERIVER_POLLING_BACKOFF_ENABLED` with configurable max/multiplier (#758) +- New Prometheus `db_pool_connections` gauge (checked_out / checked_in / size / overflow), labeled `api`|`deriver`, registered in both the API lifespan and the deriver metrics server (#758) +- New Prometheus `db_connection_acquisitions{outcome=ok|retried|exhausted}` counter — the alertable early-warning signal that connection checkouts are retrying through pooler rejection, before requests start failing (#758) +- New Prometheus `db_queries_in_flight` gauge — statements actually executing on the wire (via SQLAlchemy cursor-execute events). Paired with `checked_out`, the gap reveals connections held but parked (the "idle in transaction during an external call" antipattern). Gated on `METRICS.ENABLED` for zero overhead when + off (#758) +- Explicit `SqlalchemyIntegration` in both the API and deriver Sentry inits; connection acquisition wrapped in a `db.pool.acquire` span with live pool stats captured on retry exhaustion (#758) + +### Changed + +- Default `POOL_TIMEOUT` lowered to 5s, with validation that it stays under the connection-retry budget when a pooled (non-null) `POOL_CLASS` is configured; `config.toml.example` and the v2/v3 configuration docs updated to match (#758) +- `HonchoAsyncSession` wraps every DB-touching session method (execute / scalar / scalars / flush / merge / refresh / commit / get / get_one / stream / stream_scalars / delete) so the lazy-checkout-with-retry guarantee has no holes; the acquired flag resets on `close()`/`reset()` so a reused session re-acquires on + next use (#758) + +### Fixed + +- Roll the session back on a retryable checkout failure before retrying — a failed autobegin could otherwise leave it pending-rollback, making the next connection attempt raise instead of cleanly re-checking-out (#758) +- Guard `DBPoolCollector.collect()` so a pool-read/import hiccup can't raise and abort the entire `/metrics` scrape (Prometheus drops all metrics if any collector raises) (#758) +- Clamp the pool overflow gauge to ≥ 0 (it could report negative before the pool fills) (#758) +- Removed a double-sleep in the deriver idle poll so the backoff cap is a true cap rather than 2× (#758) + ## [3.0.7] - 2026-05-21 ### Added diff --git a/CLAUDE.md b/CLAUDE.md index fbd05871..c9862b59 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -116,6 +116,12 @@ cd sdks/typescript && bun run tsc --noEmit - Explicit error handling with appropriate exception types - Docstrings: Use Google style docstrings - **Never hold a DB session during external calls** (LLM, embedding, HTTP). If a function needs both a DB session and an external call result, compute the external result first and pass it as a parameter. This avoids tying up DB connections during slow network I/O. Use `tracked_db` for short-lived, DB-only operations; pass a shared session when multiple DB-only calls can reuse one connection. +- **Never write through a read-only session** (`tracked_db(..., read_only=True)`, `get_read_db`, `ReadSessionLocal`). These run in AUTOCOMMIT mode with no transaction: writes are NOT blocked by the database — they silently commit immediately, and `begin_nested()` savepoints break. There is no runtime guard; this is enforced by convention only. Use `read_only=True` strictly for SELECT-only windows; anything that mutates (including get-or-create paths) must use a regular write session. + +#### Auth scoping + +- **`allow_member_read=True` (in `require_auth(...)`) is read-only — NEVER set it on a route that mutates state.** It lets a peer-scoped key reach a session route when its peer is an active member of the session, so on a mutating route it would hand any session member write access (message injection, config mutation, deletion). HTTP method is not a reliable read/write signal here (some read routes use POST for a richer body), so this is enforced by an explicit allowlist in `tests/routes/test_auth_route_policy.py` — adding the flag to a new route fails that test until you consciously add the route to `EXPECTED_MEMBER_READ_ROUTES`, and you must never add a mutating method there. +- **When a member-read route is keyed by another sub-resource** (e.g. `peers/{peer_id}/config`), the handler must additionally confirm a peer-scoped caller only reads its OWN resource (`jwt_params.p == peer_id`, else raise `AuthenticationException`). Membership grants session access, not access to a co-member's data. See `get_peer_config` in `src/routers/sessions.py`. ### Runtime Architecture diff --git a/README.md b/README.md index 4ad0518a..3ecd6f7f 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ --- -![Static Badge](https://img.shields.io/badge/Server-3.0.7-blue) +![Static Badge](https://img.shields.io/badge/Server-3.0.9-blue) [![PyPI version](https://img.shields.io/pypi/v/honcho-ai.svg)](https://pypi.org/project/honcho-ai/) [![NPM version](https://img.shields.io/npm/v/@honcho-ai/sdk.svg)](https://npmjs.org/package/@honcho-ai/sdk) [![Discord](https://img.shields.io/discord/1016845111637839922?style=flat&logo=discord&logoColor=23ffffff&label=Plastic%20Labs&labelColor=235865F2)](https://discord.gg/honcho) @@ -394,6 +394,26 @@ the `AUTH_JWT_SECRET` environment variable. This is required for `AUTH_USE_AUTH` AUTH_JWT_SECRET= ``` +Once auth is enabled, use `scripts/generate_jwt.py` to mint tokens for local +development and scripting: + +```bash +# Admin token (full access, no expiry) +uv run python scripts/generate_jwt.py --admin + +# Admin token expiring in 24 hours +uv run python scripts/generate_jwt.py --admin --expires 24h + +# Workspace-scoped token +uv run python scripts/generate_jwt.py --workspace my-workspace --expires 30d + +# Capture a token for use in curl/scripts +TOKEN=$(uv run python scripts/generate_jwt.py --admin --print-only) +curl -H "Authorization: Bearer $TOKEN" http://localhost:8000/v3/workspaces +``` + +Duration units: `s` (seconds), `m` (minutes), `h` (hours), `d` (days), `w` (weeks), `y` (years). + 5. **Run database migrations** With the database set up and environment variables configured, run the migrations diff --git a/config.toml.example b/config.toml.example index 5f96ab32..aa4ddb41 100644 --- a/config.toml.example +++ b/config.toml.example @@ -6,6 +6,7 @@ # Application-level settings [app] LOG_LEVEL = "INFO" +PERFORMANCE_LOG_FORMAT = "compact" # "compact" for single-line logs, "rich" for local panels SESSION_OBSERVERS_LIMIT = 10 GET_CONTEXT_MAX_TOKENS = 100000 MAX_FILE_SIZE = 5242880 # 5MB @@ -26,11 +27,14 @@ POOL_CLASS = "default" POOL_PRE_PING = true POOL_SIZE = 10 MAX_OVERFLOW = 20 -POOL_TIMEOUT = 30 # seconds +POOL_TIMEOUT = 5 # seconds a pooled checkout waits for a free connection (QueuePool only) POOL_RECYCLE = 300 # seconds POOL_USE_LIFO = true SQL_DEBUG = false TRACING = false +# Per-connection establish timeout (seconds) so a single connection attempt +# fails fast instead of hanging when the server/pooler is unreachable. +CONNECT_TIMEOUT_SECONDS = 2 # Authentication settings [auth] @@ -81,14 +85,33 @@ model = "text-embedding-3-small" ENABLED = true WORKERS = 1 POLLING_SLEEP_INTERVAL_SECONDS = 1.0 +# Adaptive polling: when idle/erroring, the sleep interval grows from +# POLLING_SLEEP_INTERVAL_SECONDS toward POLLING_SLEEP_MAX_INTERVAL_SECONDS by +# POLLING_BACKOFF_MULTIPLIER each cycle, then snaps back to base when work is +# found. Cuts steady-state query load against the shared DB/pooler. +POLLING_BACKOFF_ENABLED = true +POLLING_SLEEP_MAX_INTERVAL_SECONDS = 30.0 +POLLING_BACKOFF_MULTIPLIER = 2.0 +# Jitter so instances that start together don't poll in lockstep. Startup: +# sleep a random delay in [0, POLLING_STARTUP_JITTER_SECONDS] before the first +# poll (0.0 disables). Per-cycle: multiply every poll sleep by a random factor +# in [1 - ratio, 1 + ratio] (0.5 -> [0.5x, 1.5x]; 0.0 disables). +POLLING_STARTUP_JITTER_SECONDS = 30.0 +POLLING_JITTER_RATIO = 0.5 STALE_SESSION_TIMEOUT_MINUTES = 5 +# Minimum (jittered) spacing between stale-work-unit cleanup runs per instance. +# Staleness is a minutes-timescale condition, so cleanup doesn't need to run on +# every seconds-scale poll (0.0 = run every poll, legacy behavior). +STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS = 60.0 # QUEUE_ERROR_RETENTION_SECONDS = 2592000 # 30 days DEDUPLICATE = true LOG_OBSERVATIONS = false MAX_INPUT_TOKENS = 25000 MAX_CUSTOM_INSTRUCTIONS_TOKENS = 2000 WORKING_REPRESENTATION_MAX_OBSERVATIONS = 100 -REPRESENTATION_BATCH_MAX_TOKENS = 1024 +REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS = 512 # Min tokens a work unit accumulates before the deriver claims it; 0 disables the gate +REPRESENTATION_BATCH_TARGET_INPUT_TOKENS = 1024 # Max context-window tokens per deriver LLM call +REPRESENTATION_BATCH_MAX_AGE_SECONDS = 1800 FLUSH_ENABLED = false # Bypass batch token threshold, process work immediately [deriver.model_config] diff --git a/docs/changelog/compatibility-guide.mdx b/docs/changelog/compatibility-guide.mdx index f12d3cee..42b2f373 100644 --- a/docs/changelog/compatibility-guide.mdx +++ b/docs/changelog/compatibility-guide.mdx @@ -30,7 +30,11 @@ This guide helps you match the right SDK version to your Honcho API version. New | Honcho API Version | TypeScript SDK | Python SDK | |-------------------|---------------|------------| -| v3.0.7 (Current) | v2.1.2 | v2.1.2 | +| v3.0.11 (Current) | v2.1.2 | v2.1.2 | +| v3.0.10 | v2.1.2 | v2.1.2 | +| v3.0.9 | v2.1.2 | v2.1.2 | +| v3.0.8 | v2.1.2 | v2.1.2 | +| v3.0.7 | v2.1.2 | v2.1.2 | | v3.0.6 | v2.1.1 | v2.1.1 | | v3.0.5 | v2.1.0 | v2.1.0 | | v3.0.4 | v2.1.0 | v2.1.0 | diff --git a/docs/changelog/introduction.mdx b/docs/changelog/introduction.mdx index 7a0dfbaa..6e741563 100644 --- a/docs/changelog/introduction.mdx +++ b/docs/changelog/introduction.mdx @@ -27,7 +27,95 @@ Welcome to the Honcho changelog! This section documents all notable changes to t ### Honcho API and SDK Changelogs - + + ### Added + + - `api_request_duration_seconds` Prometheus histogram tracking per-route request latency, labeled by method and endpoint (#837) + - LLM `provider_params` passthroughs (`extra_body` / `extra_headers` / `extra_query`) are now forwarded to the underlying provider transport across all backends, with shape validation that rejects non-mapping values (#821) + - `structured_output_mode` model-config option to use `json_object` mode for OpenAI-compatible providers that lack native Structured Outputs support (used by the deriver) (#820) + - OpenRouter app-attribution headers (`HTTP-Referer` / `X-Openrouter-Title`) are now sent on OpenAI-compatible clients when the configured base URL is OpenRouter, so requests are attributed to "Honcho" in OpenRouter's dashboard (#805) + - Langfuse traces are now tagged with user and session IDs for easier trace filtering (#814) + - `DERIVER_REPRESENTATION_BATCH_MAX_AGE_SECONDS` (default 1800s) lets sub-threshold representation work units flush once their oldest unprocessed queue item ages out. Set it to `0` to keep the legacy behavior where sub-threshold tails wait indefinitely unless `DERIVER_FLUSH_ENABLED=true` (#826) + - Conclusion responses now include a `level` field (`explicit`, `deductive`, `inductive`, `contradiction`); list/query endpoints support filtering by `level` via `filters`, with reserved filter keys protected from being overridden by user-supplied filters (#851) + + ### Changed + + - Peer-scoped JWTs now get read-only access to the sessions their peer is an active member of (session context, summaries, peers, their own per-session config, search, and message reads). Session-scoped JWTs remain confined to their session and cannot reach peer routes (#679) + - Compacted Honcho's log output, with guarded ms/s metric formatting that falls back to a plain string for non-numeric values (#836) + - Sentry now drops noisy infra/scrape transactions: the reconciler opens a transaction only once a batch has rows (idle cycles emit none), and a `traces_sampler` returns `0.0` for `/metrics`, `/health`, `/openapi.json`, `/docs`, `/redoc`, and the deriver metrics server. `SENTRY.TRACES_SAMPLE_RATE` still governs real traffic (#834) + + ### Fixed + + - Peer- and session-scoped JWTs were effectively workspace-scoped: authorization walked the route's declared scope and fell through to a workspace match, so a `{w, p: alice}` token could act on any peer in the workspace. JWTs are now authorized by their narrowest claim and never widen to workspace access (#679) + - The keys API now rejects creating a peer- or session-scoped key without a workspace. Such keys were minted successfully but failed verification on every request (#679) + - Agent-supplied observation IDs carrying the display-format `id:` prefix are now normalized (prefix and trailing whitespace stripped) before `source_ids` are stored and on `get_reasoning_chain` lookups, fixing corrupted provenance links and broken reasoning-chain traversal (#795) + - Fixed a `create_tree` keyword-argument mismatch in the Dreamer's surprisal tree construction (#749) + - Providers that omit output-token counts (observed with Gemini on tool-loop completions) returned `output_tokens=None`, which raised a Pydantic validation error that aborted the call and crashed the Dreamer's induction phase before inductive conclusions were persisted. `None` is now coerced to `0` so token accounting degrades gracefully (#809) + - Document creation now performs exact (case-insensitive, whitespace-trimmed) content deduplication before the existing semantic dedup step: exact duplicates within a batch collapse to a single insert, and an exact match against a live document reinforces it (atomic `times_derived` increment) instead of creating a new row (#861) + + + + ### Added + + - Messages are now embedded via a background task rather than blocking API request + - Read-only DB session mode (`get_read_db` / `tracked_db(..., read_only=True)`) so reads don't hold a transaction open across the work + - `CORS_ORIGINS` env var to configure CORS allowed origins without editing source; defaults match the prior hardcoded list, so self-hosted deployments behind custom domains can whitelist their frontend (#697) + - `scripts/generate_jwt.py` — utility for minting scoped or admin Honcho JWTs (`--admin`, `--workspace`/`--peer`/`--session`, `--expires` with human-friendly durations, `--print-only`) without calling the keys API (#757) + - `STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS` (default 60s) — minimum jittered spacing between deriver stale-work-unit cleanup runs, so cleanup no longer runs on every seconds-scale poll (`0.0` keeps the legacy every-poll behavior) (#773) + + ### Changed + + - Optimized the deriver and dreamer prompt cache prefixes to improve prompt-cache hit rates (#806) + + ### Fixed + + - `times_derived` is now properly reinforced when a duplicate conclusion is detected. It had been pinned at 1 for nearly every conclusion (the reject-new branch dropped the increment and the new-wins branch reset the count to 1), so `ORDER BY times_derived DESC` fell back to arbitrary heap order and froze stale conclusions to the front of injected context. Reinforcement is now an atomic increment and both most-derived queries gained a `created_at DESC` recency tiebreaker (#768) + - Webhook creation now correctly rejects private/internal IP addresses (#793) + + + + ### Changed + + - Connection acquisition is now a single attempt with no server-side retry, on a vanilla `AsyncSession`. A new `DB_CONNECT_TIMEOUT_SECONDS` (default 2s) bounds the attempt so a saturated or unreachable pooler fails fast instead of holding a client connection open to re-knock. A saturated DB now surfaces to the caller — the API returns an error and the deriver backs off and retries on a later poll — which lets the pooler drain rather than amplifying saturation. + + ### Added + + - Deriver poll jitter so instances that start together don't poll in lockstep: `DERIVER_POLLING_STARTUP_JITTER_SECONDS` (random delay before the first poll, default 30s) and `DERIVER_POLLING_JITTER_RATIO` (±fraction applied to every poll sleep, default 0.5). Both disable at `0.0`; the underlying backoff schedule is unchanged. + + ### Removed + + - Reverted the connection-checkout retry and `HonchoAsyncSession` custom session introduced in 3.0.8. Removed the `DB_CONNECTION_RETRY_ENABLED` / `DB_CONNECTION_RETRY_MAX_DELAY_SECONDS` / `DB_CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS` / `DB_CONNECTION_RETRY_BACKOFF_MAX_SECONDS` settings, the `db_connection_acquisitions{outcome=...}` Prometheus counter, and the `db.pool.acquire` Sentry span. Alerting built on `db_connection_acquisitions` should migrate to `db_pool_connections` / `db_queries_in_flight`. + + + + ### Added + + - Connection-checkout retry with bounded exponential backoff (tenacity) on `get_db`/`tracked_db`: transient transaction-pooler (Supavisor) rejections — SQLAlchemy `TimeoutError` and `OperationalError` — now retry with backoff instead of surfacing as 500s under client-connection saturation. Gated by + `DB_CONNECTION_RETRY_ENABLED` with configurable delay/backoff knobs; ~10s default budget (#758) + - `HonchoAsyncSession` — a lazy `AsyncSession` that checks out its pooled connection (with retry) on the first DB-touching call rather than at construction. Request handlers doing non-DB work (embedding, file, LLM) before their first query no longer pin a pooler connection across it. Only the checkout is retried; + the statement still runs exactly once, so writes are never duplicated (#758) + - Adaptive deriver queue polling: the poll interval backs off when the queue is idle or erroring (base → max, doubling each cycle) and snaps back to base the moment work is claimed, cutting steady-state query load against the DB. Gated by `DERIVER_POLLING_BACKOFF_ENABLED` with configurable max/multiplier (#758) + - New Prometheus `db_pool_connections` gauge (checked_out / checked_in / size / overflow), labeled `api`|`deriver`, registered in both the API lifespan and the deriver metrics server (#758) + - New Prometheus `db_connection_acquisitions{outcome=ok|retried|exhausted}` counter — the alertable early-warning signal that connection checkouts are retrying through pooler rejection, before requests start failing (#758) + - New Prometheus `db_queries_in_flight` gauge — statements actually executing on the wire (via SQLAlchemy cursor-execute events). Paired with `checked_out`, the gap reveals connections held but parked (the "idle in transaction during an external call" antipattern). Gated on `METRICS.ENABLED` for zero overhead when + off (#758) + - Explicit `SqlalchemyIntegration` in both the API and deriver Sentry inits; connection acquisition wrapped in a `db.pool.acquire` span with live pool stats captured on retry exhaustion (#758) + + ### Changed + + - Default `POOL_TIMEOUT` lowered to 5s, with validation that it stays under the connection-retry budget when a pooled (non-null) `POOL_CLASS` is configured; `config.toml.example` and the v2/v3 configuration docs updated to match (#758) + - `HonchoAsyncSession` wraps every DB-touching session method (execute / scalar / scalars / flush / merge / refresh / commit / get / get_one / stream / stream_scalars / delete) so the lazy-checkout-with-retry guarantee has no holes; the acquired flag resets on `close()`/`reset()` so a reused session re-acquires on + next use (#758) + + ### Fixed + + - Roll the session back on a retryable checkout failure before retrying — a failed autobegin could otherwise leave it pending-rollback, making the next connection attempt raise instead of cleanly re-checking-out (#758) + - Guard `DBPoolCollector.collect()` so a pool-read/import hiccup can't raise and abort the entire `/metrics` scrape (Prometheus drops all metrics if any collector raises) (#758) + - Clamp the pool overflow gauge to ≥ 0 (it could report negative before the pool fills) (#758) + - Removed a double-sleep in the deriver idle poll so the backoff cap is a true cap rather than 2× (#758) + + + ### Added - New `src/llm/` package as the single owner of provider runtime: clients, backends, history adapters, tool loop, request builder, credentials, and caching policy (#459) @@ -610,7 +698,17 @@ Welcome to the Honcho changelog! This section documents all notable changes to t [Python SDK](https://pypi.org/project/honcho-ai/) - + + ### Added + + - `ConclusionLevel` type (`explicit`, `deductive`, `inductive`, `contradiction`) and a `level` field on `Conclusion`, exposing the reasoning level the server already tracked but previously stripped from responses. + - `filters` parameter on `ConclusionScope.list()` and `ConclusionScope.query()` (sync and async), passed through to the same dynamic server-side filter logic as `peers()`/`sessions()`/`messages()`. Filter explicit-only conclusions with `filters={"level": "explicit"}`, or by any other supported field/operator. Requires a Honcho server with the matching API support (Honcho v3.0.11+). + + ### Fixed + + - Scope-managed filter keys (`observer`, `observed`, `session`) are now rejected with a clear `ValueError` if passed in `filters`, instead of silently overriding the scope and returning conclusions from a different peer pair. Use `peer.conclusions` / `conclusions_of(target)` and the `session=` parameter instead. `session_id` remains a valid filter on `query()`. + + ### Added - `page`, `size`, and `reverse` pagination parameters on `Honcho.workspaces()` and `HonchoAio.workspaces()`, closing the gap from 2.1.0 which added these to other list methods but not to `workspaces()`. Honoring `reverse` on the workspace/peer/session list routes also requires a Honcho server with the matching API fix; older servers silently ignore the parameter. @@ -762,7 +860,17 @@ Welcome to the Honcho changelog! This section documents all notable changes to t [TypeScript SDK](https://www.npmjs.com/package/@honcho-ai/sdk) - + + ### Added + + - `ConclusionLevel` type (`explicit`, `deductive`, `inductive`, `contradiction`) and a `level` field on `Conclusion`, exposing the reasoning level the server already tracked but previously stripped from responses. + - `filters` option on `conclusions.list()` and `conclusions.query()`, passed through to the same dynamic server-side filter logic as the other list endpoints. Filter explicit-only conclusions with `{ filters: { level: 'explicit' } }`, or by any other supported field/operator. Requires a Honcho server with the matching API support (Honcho v3.0.11+). + + ### Fixed + + - Scope-managed filter keys (`observer`, `observed`, `session`) are now rejected with a clear error if passed in `filters`, instead of silently overriding the scope and returning conclusions from a different peer pair. Use `peer.conclusions` / `peer.conclusionsOf(target)` and the dedicated `session` option instead. `session_id` remains a valid filter on `query()`. + + ### Added - `peers` option on `Honcho.session()` — attach peers to a session at creation time instead of needing a follow-up `session.addPeers()` call. Accepts the same `PeerAddition` shape as `session.addPeers()` (peer ID strings, `Peer` objects, arrays of either, or a record with per-peer `observe_me`/`observe_others` config). diff --git a/docs/docs.json b/docs/docs.json index 64b56d86..e5bb31e8 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -24,7 +24,7 @@ "navigation": { "versions": [ { - "version": "v3.0.7", + "version": "v3.0.11", "api": { "openapi": ["v3/openapi.json"] }, @@ -67,6 +67,7 @@ "v3/documentation/features/advanced/queue-status", "v3/documentation/features/advanced/search", "v3/documentation/features/advanced/using-filters", + "v3/documentation/features/advanced/structured-outputs", "v3/documentation/features/advanced/streaming-response", "v3/documentation/features/advanced/file-uploads" ] @@ -95,6 +96,7 @@ "pages": [ "v3/guides/integrations/claude-code", "v3/guides/integrations/opencode", + "v3/guides/integrations/codex", "v3/guides/integrations/vercel-ai-sdk", "v3/guides/integrations/crewai", "v3/guides/integrations/langgraph", diff --git a/docs/images/honcho-system-diagram.png b/docs/images/honcho-system-diagram.png new file mode 100644 index 00000000..1606c803 Binary files /dev/null and b/docs/images/honcho-system-diagram.png differ diff --git a/docs/v2/contributing/configuration.mdx b/docs/v2/contributing/configuration.mdx index c172369c..f2548f94 100644 --- a/docs/v2/contributing/configuration.mdx +++ b/docs/v2/contributing/configuration.mdx @@ -161,7 +161,7 @@ DB_CONNECTION_URI=postgresql+psycopg://honcho_user:secure_password@db.example.co DB_SCHEMA=public DB_POOL_SIZE=10 DB_MAX_OVERFLOW=20 -DB_POOL_TIMEOUT=30 +DB_POOL_TIMEOUT=5 DB_POOL_RECYCLE=300 DB_POOL_PRE_PING=true DB_SQL_DEBUG=false diff --git a/docs/v3/api-reference/endpoint/health-check.mdx b/docs/v3/api-reference/endpoint/health-check.mdx new file mode 100644 index 00000000..35ceba0d --- /dev/null +++ b/docs/v3/api-reference/endpoint/health-check.mdx @@ -0,0 +1,3 @@ +--- +openapi: get /health +--- diff --git a/docs/v3/contributing/configuration.mdx b/docs/v3/contributing/configuration.mdx index 707ff282..8f4c3eb6 100644 --- a/docs/v3/contributing/configuration.mdx +++ b/docs/v3/contributing/configuration.mdx @@ -60,6 +60,12 @@ You can mix providers freely — for example, use Gemini for the deriver and Cla For OpenAI-compatible proxies (OpenRouter, vLLM, Ollama, etc.), use `transport = "openai"` and set `MODEL_CONFIG__OVERRIDES__BASE_URL` on each feature to point at your endpoint. + +Some OpenAI-compatible providers don't support OpenAI Structured Outputs (`json_schema`). Set `DERIVER_MODEL_CONFIG__STRUCTURED_OUTPUT_MODE=json_object` to request loose JSON mode and inject the schema into the prompt instead. + +This setting only applies to the **deriver** on the **`openai`** transport — it is the only feature that uses structured output. The dialectic, summarizer, and dreamer don't request structured output, so the setting has no effect there, and the anthropic/gemini transports reject it. + + ### Tiered Model Setup Once you're past initial setup, you can assign different models per feature for better cost/quality tradeoffs. This example uses OpenRouter with light/medium/heavy tiers: @@ -185,6 +191,33 @@ Each model config supports an `overrides.provider_params` dict for passing arbit verbosity = "low" ``` +#### Transport passthrough keys + +Three keys inside `provider_params` are recognized as request-level escape hatches and forwarded to the underlying transport. Where a transport actually validates and merges one of these keys, its value must be a mapping — a non-mapping value raises a configuration error (see the per-transport behavior below; a key a transport ignores is not validated): + +- **`extra_body`** — merged into the request body +- **`extra_headers`** — extra HTTP headers +- **`extra_query`** — extra URL query parameters + +How each transport forwards them differs: + +- **OpenAI and Anthropic** forward all three as identically-named SDK kwargs (`extra_body`, `extra_headers`, `extra_query`). +- **Gemini** has no SDK kwargs for these. It merges `extra_body` into the `GenerateContentConfig` dict and folds `extra_headers` into `http_options.headers`; `extra_query` is **unsupported and silently ignored**. + +The merge is shallow and **operator-wins**: if Honcho and your config both set the same top-level key inside `extra_body`, your value replaces Honcho's. You are responsible for choosing a coherent combination — e.g. unset `thinking_budget_tokens` when supplying an `extra_body.thinking` for Anthropic-via-proxy, since Honcho will not translate between the two shapes. + +Because Gemini merges `extra_body` directly into `GenerateContentConfig` (rather than a nested request body), an `extra_body` written for OpenAI/Anthropic generally will not transfer to Gemini unchanged — and a key collision there can overwrite a field Honcho manages (`thinking_config`, `response_schema`, `tools`, …). + +```toml +# Example: route an OpenAI-compatible proxy and tag requests for tracing +[deriver.model_config.overrides.provider_params.extra_headers] +X-Proxy-Route = "vertex" + +[deriver.model_config.overrides.provider_params.extra_body] +# Provider-native body fields the standard config doesn't expose +anthropic_beta = ["context-1m-2025-01-15"] +``` + ### Changing Transport When changing a feature's `transport`, always specify `model` explicitly. Partial overrides that change transport without model will keep the previous model name, which may not be valid for the new provider. @@ -346,6 +379,7 @@ DERIVER_MAX_CUSTOM_INSTRUCTIONS_TOKENS=2000 # DERIVER_MODEL_CONFIG__THINKING_EFFORT=minimal # DERIVER_MODEL_CONFIG__THINKING_BUDGET_TOKENS=1024 # DERIVER_MODEL_CONFIG__TEMPERATURE=0.7 # Optional temperature override +# DERIVER_MODEL_CONFIG__STRUCTURED_OUTPUT_MODE=json_object # for providers without json_schema support # Backup model (optional) # DERIVER_MODEL_CONFIG__FALLBACK__MODEL=claude-haiku-4-5 @@ -354,6 +388,17 @@ DERIVER_MAX_CUSTOM_INSTRUCTIONS_TOKENS=2000 # Worker settings DERIVER_WORKERS=1 # Increase for higher throughput DERIVER_POLLING_SLEEP_INTERVAL_SECONDS=1.0 +# Adaptive polling: when idle/erroring, the sleep interval grows from the base +# toward DERIVER_POLLING_SLEEP_MAX_INTERVAL_SECONDS by the multiplier each cycle, +# then snaps back to base when work is found. Cuts steady-state query load. +DERIVER_POLLING_BACKOFF_ENABLED=true +DERIVER_POLLING_SLEEP_MAX_INTERVAL_SECONDS=30.0 +DERIVER_POLLING_BACKOFF_MULTIPLIER=2.0 +# Jitter so instances that start together don't poll in lockstep. Startup: sleep +# a random delay in [0, value] before the first poll (0.0 disables). Per-cycle: +# multiply every poll sleep by a random factor in [1-ratio, 1+ratio] (0.0 disables). +DERIVER_POLLING_STARTUP_JITTER_SECONDS=30.0 +DERIVER_POLLING_JITTER_RATIO=0.5 DERIVER_STALE_SESSION_TIMEOUT_MINUTES=5 # Queue management @@ -363,7 +408,9 @@ DERIVER_QUEUE_ERROR_RETENTION_SECONDS=2592000 # 30 days DERIVER_DEDUPLICATE=true DERIVER_LOG_OBSERVATIONS=false DERIVER_WORKING_REPRESENTATION_MAX_OBSERVATIONS=100 -DERIVER_REPRESENTATION_BATCH_MAX_TOKENS=1024 +DERIVER_REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS=512 +DERIVER_REPRESENTATION_BATCH_TARGET_INPUT_TOKENS=1024 +DERIVER_REPRESENTATION_BATCH_MAX_AGE_SECONDS=1800 ``` **Peer Card:** @@ -459,10 +506,13 @@ DB_SCHEMA=public DB_POOL_PRE_PING=true DB_POOL_SIZE=10 DB_MAX_OVERFLOW=20 -DB_POOL_TIMEOUT=30 +DB_POOL_TIMEOUT=5 DB_POOL_RECYCLE=300 DB_POOL_USE_LIFO=true DB_SQL_DEBUG=false +# Per-connection establish timeout (seconds) so a single connection attempt +# fails fast instead of hanging when the server/pooler is unreachable. +DB_CONNECT_TIMEOUT_SECONDS=2 ``` ### Authentication diff --git a/docs/v3/contributing/troubleshooting.mdx b/docs/v3/contributing/troubleshooting.mdx index bb71a475..d9b745a1 100644 --- a/docs/v3/contributing/troubleshooting.mdx +++ b/docs/v3/contributing/troubleshooting.mdx @@ -109,7 +109,7 @@ Messages are stored but no observations, summaries, or representations are being ```bash DERIVER_WORKERS=4 ``` -5. **Representation Batch Max** — By default the deriver is set to buffer its operations until there are enough tokens for a given representation in a session. This is set via the `REPRESENTATION_BATCH_MAX_TOKENS` environment variable. If you aren't seeing tasks continue it may be that the batch size is set too high or enough data hasn't flowed into to the session yet. See [token batching](/v3/documentation/core-concepts/reasoning#token-batching) for more details +5. **Representation Batching** — By default the deriver buffers representation work until a work unit has accumulated enough tokens, set via `DERIVER_REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS` (`0` disables the accumulation gate). A separate setting, `DERIVER_REPRESENTATION_BATCH_TARGET_INPUT_TOKENS`, caps the conversation window fed to each deriver LLM call when draining a claimed work unit. Sub-threshold tails become eligible after `DERIVER_REPRESENTATION_BATCH_MAX_AGE_SECONDS` (default 1800 seconds), so quiet sessions eventually flush without disabling batching globally. Set the age to `0` for legacy behavior where sub-threshold tails wait indefinitely. See [token batching](/v3/documentation/core-concepts/reasoning#token-batching) for more details ## Alternative Provider Issues @@ -142,7 +142,19 @@ If calls to an OpenAI-compatible proxy fail: DERIVER_MODEL_CONFIG__OVERRIDES__BASE_URL=http://host.docker.internal:8000/v1 ``` -3. **Structured output failures** — vLLM's structured output support is limited to certain response formats. If you see JSON parsing errors, check the deriver/dream logs for the raw response. +3. **Structured output failures** — vLLM's structured output support is limited to certain response formats. If you see JSON parsing errors, check the deriver/dream logs for the raw response. See [Deriver produces no observations](#deriver-produces-no-observations) below. + +### Deriver produces no observations + +If messages are processed (the queue drains, no errors in logs) but peers never accumulate observations — and you're using an OpenAI-compatible provider — the likely cause is that the provider doesn't support OpenAI Structured Outputs (`json_schema`). The OpenAI backend requests `json_schema` by default; providers like **Z.AI GLM** and some **Ollama/vLLM** deployments either reject it or silently ignore it and return prose, which the deriver can't parse into observations. + +**Fix:** set `STRUCTURED_OUTPUT_MODE=json_object` on the deriver's model config to request loose JSON mode, which injects the schema into the prompt instead: + +```bash +DERIVER_MODEL_CONFIG__STRUCTURED_OUTPUT_MODE=json_object +``` + +This is a per-model-config setting on the OpenAI transport; set it on whichever features use the affected provider (e.g. `DREAM_DEDUCTION_MODEL_CONFIG__STRUCTURED_OUTPUT_MODE`). ### Thinking budget errors with non-Anthropic providers diff --git a/docs/v3/documentation/core-concepts/architecture.mdx b/docs/v3/documentation/core-concepts/architecture.mdx index 429fcb0a..b3f0f7a1 100644 --- a/docs/v3/documentation/core-concepts/architecture.mdx +++ b/docs/v3/documentation/core-concepts/architecture.mdx @@ -76,6 +76,20 @@ When you need context from Honcho, you query through the "Chat" endpoint or "Get The diagram above shows how agents write messages to Honcho, which triggers reasoning that updates peer representations. Agents can then query representations to get additional context for their next response. Black arrows represent read/write of regular data (messages, storage), while red arrows represent read/write of reasoned-over data (logic, peer representations). +### Under the Hood: Write, Reasoning, and Query Paths + +Honcho runs as two cooperating processes: an **API server** that handles requests and enqueues background work, and a **worker** that consumes that work off the queue. This split is what keeps the write path fast--your request never waits on an LLM call. + +![Honcho Detailed Internals](/images/honcho-system-diagram.png) + +**Write path (synchronous).** A message is stored and a reasoning task is enqueued in the same request; the API returns immediately. Nothing about the reasoning that follows blocks the caller. + +**Deriver + Summarizer (async, per-message).** The worker picks up queued tasks in small batches. The Deriver reads new messages and extracts conclusions about the peer--explicit statements and direct deductions. In parallel, the Summarizer periodically rolls up recent messages into short- and long-form session summaries. Both run per-message (well, per-batch) rather than on a schedule. + +**Dreamer (periodic).** On a schedule (or triggered on demand), the Dreamer revisits existing conclusions to consolidate and deepen them: removing redundant or stale ones, drawing inductive conclusions across patterns that span multiple messages, and updating peer cards--compact biographical summaries of a peer. This is where memory gets richer over time, not just larger. + +**Query path (Dialectic).** A `chat()` call spawns a Dialectic agent that answers your question by exploring memory--searching conclusions semantically, pulling supporting messages, and tracing a conclusion back to the premises it was drawn from--before synthesizing a grounded answer. This all happens inline during the request, since answering well is worth the latency that write-path reasoning is designed to avoid. + ## Configuration & Extensibility Honcho is designed to be flexible. Settings cascade hierarchically from workspace to peer to session, so you can set defaults at the workspace level and override them for specific peers or sessions. Feature flags let you enable or disable reasoning modes, perspective tracking, and other capabilities. You can bring your own LLM provider--OpenAI, Anthropic, or custom endpoints--and metadata fields let you extend any primitive with custom JSON data. Batch operations let you create up to 100 messages in a single API call for efficient bulk ingestion. diff --git a/docs/v3/documentation/core-concepts/design-patterns.mdx b/docs/v3/documentation/core-concepts/design-patterns.mdx index 06582d80..9893d3f4 100644 --- a/docs/v3/documentation/core-concepts/design-patterns.mdx +++ b/docs/v3/documentation/core-concepts/design-patterns.mdx @@ -63,7 +63,7 @@ Sessions define the temporal boundaries of an interaction. How you scope them af Create a **new** session when context resets (new conversation, new day, new topic); **reuse** one when context should keep accumulating (ongoing channel, persistent thread). -**Don't scope sessions too thin.** Honcho only reasons over a peer once it accumulates ~1,000 tokens *within a single session* ([token batching](/v3/documentation/core-concepts/reasoning#token-batching)). Many tiny sessions each stall below that threshold, so low-volume or trickle inputs should append to one ongoing session rather than fragment across many (nothing is lost — it just waits). +**Don't scope sessions too thin.** Honcho batches reasoning until a peer accumulates ~1,000 tokens *within a single session*, with a default age-based flush for quiet tails ([token batching](/v3/documentation/core-concepts/reasoning#token-batching)). Low-volume or trickle inputs should still append to one ongoing session rather than fragment across many, so reasoning runs with useful context instead of many small delayed batches. **How cross-session reasoning works** diff --git a/docs/v3/documentation/core-concepts/reasoning.mdx b/docs/v3/documentation/core-concepts/reasoning.mdx index 6d198c5a..aa4de900 100644 --- a/docs/v3/documentation/core-concepts/reasoning.mdx +++ b/docs/v3/documentation/core-concepts/reasoning.mdx @@ -54,6 +54,8 @@ The explicit reasoning model ([Neuromancer XR](https://blog.plasticlabs.ai/resea The reasoning that Honcho does is something we're constantly iterating and improving on. Our goal is simple--provide the richest, most relevant context in the fastest, cheapest way possible in order to simulate statefulness in whatever setting you need. +Two components produce this logic: the **Deriver** extracts explicit and deductive conclusions from incoming messages as they arrive, and the **Dreamer** periodically revisits stored conclusions to consolidate them and draw inductive conclusions across patterns spanning multiple messages. See [Architecture](/v3/documentation/core-concepts/architecture) for how these fit into the request/background split. + ## How It Works When you write messages to Honcho, they're stored immediately and enqueued for background processing. Reasoning asynchronously ensures fast writes while still providing rich reasoning capabilities. Messages are stored immediately without blocking, and session-based queues maintain chronological consistency so reasoning tasks affecting the same peer representation are always processed in order. diff --git a/docs/v3/documentation/features/advanced/search.mdx b/docs/v3/documentation/features/advanced/search.mdx index 90f155a2..a95fff15 100644 --- a/docs/v3/documentation/features/advanced/search.mdx +++ b/docs/v3/documentation/features/advanced/search.mdx @@ -6,6 +6,10 @@ icon: 'magnifying-glass' Honcho's search functionality allows you to find relevant messages and conversations across different scopes - from entire workspaces down to specific peers or sessions. + +Search is hybrid: it combines full-text (keyword) matching with semantic (vector) similarity. Keyword matches are available the instant a message is created. Semantic matches depend on the message's embedding, which is generated in the background, so a freshly created message may take a few seconds to surface in semantic results. If you need to assert on semantic results immediately after writing (for example in tests), wait briefly or poll. + + ## Search Scopes ### Workspace Search diff --git a/docs/v3/documentation/features/advanced/structured-outputs.mdx b/docs/v3/documentation/features/advanced/structured-outputs.mdx new file mode 100644 index 00000000..8e8f7e9a --- /dev/null +++ b/docs/v3/documentation/features/advanced/structured-outputs.mdx @@ -0,0 +1,239 @@ +--- +title: "Structured Outputs" +description: "Get chat endpoint answers as typed, machine-readable JSON" +icon: "brackets-curly" +--- + +By default, the [chat endpoint](/v3/documentation/features/chat) returns a free-form natural language answer. When your application needs machine-readable output, parsing that string yourself can be fragile and model-dependent. In this case you can use structured Dialectic outputs: pass a schema with your query, and the answer is guaranteed to conform to it. The agent still runs its full reasoning loop and only the final synthesized answer is formatted to your schema. + +## Basic Usage + +Pass a Pydantic model (Python) or Zod schema (TypeScript) as `response_format`, and the SDK returns a parsed, typed instance: + + +```python Python +from typing import Literal +from pydantic import BaseModel, Field +from honcho import Honcho + +class FoodPreference(BaseModel): + food: str + sentiment: Literal["loves", "likes", "neutral", "dislikes", "hates"] + confidence: float = Field(description="0-1, how certain the evidence is") + +class FoodPreferences(BaseModel): + preferences: list[FoodPreference] + summary: str + +honcho = Honcho() +peer = honcho.peer("user-123") + +result = peer.chat( + "What are this user's top 3 food preferences?", + response_format=FoodPreferences, +) + +# result is a FoodPreferences instance (or None if no relevant information) +if result: + for pref in result.preferences: + print(f"{pref.food}: {pref.sentiment} ({pref.confidence})") +``` + +```typescript TypeScript +import { z } from 'zod'; +import { Honcho } from '@honcho-ai/sdk'; + +const FoodPreferences = z.object({ + preferences: z.array(z.object({ + food: z.string(), + sentiment: z.enum(["loves", "likes", "neutral", "dislikes", "hates"]), + confidence: z.number(), + })), + summary: z.string(), +}); + +const honcho = new Honcho({}); +const peer = await honcho.peer("user-123"); + +const result = await peer.chat( + "What are this user's top 3 food preferences?", + { responseFormat: FoodPreferences }, +); + +// result is typed as z.infer (or null) +if (result) { + console.log(result.summary); +} +``` + + +## Using a Raw JSON Schema + +You can also pass a plain JSON Schema object instead of a Pydantic/Zod schema. In that case the SDK returns the answer as a JSON **string** and leaves parsing to you. This is also the shape the REST API accepts directly: + + +```python Python +result = peer.chat( + "What are this user's food preferences?", + response_format={ + "type": "object", + "properties": { + "foods": {"type": "array", "items": {"type": "string"}}, + }, + "required": ["foods"], + }, +) +# result is a JSON string, e.g. '{"foods": ["dark roast coffee", "sushi"]}' +``` + +```bash cURL +curl -X POST "$HONCHO_URL/v3/workspaces/my-app/peers/user-123/chat" \ + -H "Authorization: Bearer $HONCHO_API_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "query": "What are this user'\''s food preferences?", + "response_format": { + "type": "object", + "properties": { + "foods": { "type": "array", "items": { "type": "string" } } + }, + "required": ["foods"] + } + }' +``` + + +At the API level, `content` in the response is always a string. When `response_format` is set, it is a JSON-encoded object conforming to your schema. + +## Streaming + +`response_format` works with streaming. The stream emits the JSON answer incrementally as raw text chunks; the accumulated text is a valid JSON string once the stream completes. To enable streaming the SDKs cannot parse streamed responses for you, you parse the final string yourself: + + +```python Python +response_stream = peer.chat( + "What are this user's food preferences?", + stream=True, + response_format=FoodPreferences, +) + +chunks = [] +for chunk in response_stream.iter_text(): + chunks.append(chunk) + +result = FoodPreferences.model_validate_json("".join(chunks)) +``` + +```typescript TypeScript +const responseStream = await peer.chat( + "What are this user's food preferences?", + { stream: true, responseFormat: FoodPreferences }, +); + +let text = ""; +for await (const chunk of responseStream.iter_text()) { + text += chunk; +} + +const result = FoodPreferences.parse(JSON.parse(text)); +``` + + +## Supported Schema Subset + +Honcho supports a conservative subset of JSON Schema that enables the kind of Pydantic models used for structured LLM outputs. Schemas outside this subset are rejected with a `422` validation error before any reasoning runs. + +The root of the schema must be `"type": "object"`. + +| Construct | Support | +|-----------|---------| +| `string`, `number`, `integer`, `boolean`, `null` | Supported | +| `object` with `properties` (nested recursively) | Supported | +| `array` with `items` (missing `items` yields an untyped list) | Supported | +| `enum` of strings, integers, booleans, or null | Supported | +| `anyOf` / `oneOf` unions (a `null` member makes the field optional) | Supported | +| `type` given as a list (e.g. `["string", "null"]`) | Supported | +| `required`, `default`, `description` | Supported | +| Boolean `additionalProperties` | Accepted and ignored | +| `$ref` into root-level `$defs` / `definitions` | Supported — resolved by inlining (this is what Pydantic and Zod emit) | +| Recursive `$ref` (a definition that references itself, directly or indirectly) | Rejected (422) — the error will identify the cycle | +| Other `$ref` forms (external URLs, arbitrary JSON pointers) | Rejected (422) | +| `allOf`, `not`, `if` / `then` / `else` | Rejected (422) | +| `patternProperties`, schema-valued `additionalProperties` | Rejected (422) | + +Schemas may nest at most 20 levels deep and contain at most 500 total nodes. + + +Constraint keywords like `minItems`, `maxLength`, `minimum`, `pattern`, and `format` are passed through to the model as hints but are **not enforced server-side**. If you need hard guarantees on these, validate the returned object in your application. + + + +**Recursive schemas are not supported.** A self-referential Pydantic model (`Node.children: list[Node]`) or a recursive Zod schema (`z.lazy(...)`) produces a recursive `$ref`, which is rejected with a 422 naming the cycle. Restructure recursive shapes as explicit nesting with a fixed depth. + + +## Optional Fields and Unions + +Two distinct mechanisms control "optionality" in a raw JSON Schema: + +- **Omission** is controlled by `required`. A property not listed in `required` may be left out by the model entirely; the parsed answer will contain it as `null`. +- **Nullability** is controlled by the field's type. An `anyOf`/`oneOf` with a `{"type": "null"}` member (or the shorthand `"type": ["string", "null"]`) means the field's *value* may be `null` even when the field itself is required. + +`anyOf` and `oneOf` are treated identically: a plain union of the member schemas. Unions of non-null types (e.g. a string-or-integer field) are also supported. + +```json +{ + "type": "object", + "properties": { + "favorite_food": { "type": "string" }, + "dietary_restriction": { + "anyOf": [{ "type": "string" }, { "type": "null" }], + "description": "The user's dietary restriction, or null if none is known" + }, + "years_vegetarian": { "type": ["integer", "null"] }, + "confidence": { "type": "number" } + }, + "required": ["favorite_food", "dietary_restriction", "years_vegetarian"] +} +``` + +In this schema: + +- `favorite_food` is required and must be a string. +- `dietary_restriction` is required but **nullable**: the key is always present in the answer, and the model can answer `null` when it has no evidence. This is the recommended way to give the model an escape hatch (see [Best Practices](#model-uncertainty-explicitly)). +- `years_vegetarian` is the same thing written with the `type`-list shorthand (`["integer", "null"]` is equivalent to an `anyOf` of the two). +- `confidence` is not in `required`, so the model may omit it; if it does, the field comes back as `null`. + +If a property declares a `default`, that default is used whenever the model omits the field _even if_ the property is listed in `required`. + +Pydantic and Zod produce these shapes for you: `str | None` in Pydantic emits the `anyOf` form above, and `z.string().nullable()` does the same in Zod (`z.string().optional()` controls presence in `required`). + +## Error Handling + +| Condition | Result | +|-----------|--------| +| `response_format` is not a valid JSON Schema object | `422` validation error | +| Root type is not `"object"` | `422` validation error | +| Schema uses an unsupported construct | `422` identifying the construct and its path | +| Schema contains a recursive `$ref` | `422` identifying the cycle (e.g. `cycle: Node -> Node`) | +| Model fails to produce valid structured output after retries | `500`, same as any LLM failure | + +## How It Works + +Structured output constrains the final synthesis step. The reasoning itself works the same in both settings. + +1. The dialectic agent runs its normal tool loop in free-form text. It will search conclusions, grep messages, and traverse reasoning chains +2. Once the agent has gathered enough context, the final answer generation is constrained to your schema using the provider's native structured output support. +3. The conforming JSON is returned as the response `content` and parsed into a typed object by the SDK when you passed a Pydantic model or Zod schema. + +This means answer *quality* is unaffected by the schema: the agent reasons exactly as it would for a free-form answer, and reasoning levels (`minimal` through `max`) work the same way alongside `response_format`. + +## Best Practices + +### Add descriptions to your fields +Field `description`s are visible to the model when it formats the answer. `confidence: float` with a provided description of "score how certain the evidence is from 0-5" gets meaningfully better output than a bare field. + +### Model uncertainty explicitly +The chat endpoint returns `None`/`null` when it has no relevant information. With a schema, you can force an answer even when evidence is thin. To avoid hallucinations, consider including an escape hatch as an optional field, a `"confidence"` score, or an enum member like `"unknown"` so the model isn't forced to fabricate. + +### Keep schemas focused +A schema with three well-described fields outperforms one with twenty. If you need many distinct insights, consider making separate chat calls. diff --git a/docs/v3/documentation/features/advanced/using-filters.mdx b/docs/v3/documentation/features/advanced/using-filters.mdx index 306312ce..2a015c03 100644 --- a/docs/v3/documentation/features/advanced/using-filters.mdx +++ b/docs/v3/documentation/features/advanced/using-filters.mdx @@ -614,6 +614,65 @@ messages = session.messages(filters={ ``` +### Filtering Conclusions + +Conclusions are scoped to an observer/observed peer pair (accessed via +`peer.conclusions` for self-conclusions or `peer.conclusions_of(target)` for +conclusions about another peer). The observer and observed are filled in +automatically by the scope, so the `filters` you pass add to them. + +The most useful conclusion-specific field is `level`, the reasoning level: + +- `explicit` — extracted directly from messages +- `deductive` / `inductive` / `contradiction` — derived later during dreaming + +A common request is to surface only the directly-stated facts and exclude +anything inferred during dreaming — filter `level` to `explicit`: + + +```python Python +# Only conclusions extracted directly from messages (exclude dream-derived) +explicit = peer.conclusions.list(filters={"level": "explicit"}) + +# Only dream-derived conclusions +derived = peer.conclusions.list(filters={"level": {"in": ["deductive", "inductive"]}}) + +# Same filtering on semantic search +results = peer.conclusions.query( + "food preferences", + filters={"level": "deductive"}, +) + +# Conclusions about another peer, explicit only +bob_explicit = peer.conclusions_of("bob").list(filters={"level": "explicit"}) +``` + +```typescript TypeScript +(async () => { + // Only conclusions extracted directly from messages (exclude dream-derived) + const explicit = await peer.conclusions.list({ filters: { level: "explicit" } }); + + // Only dream-derived conclusions + const derived = await peer.conclusions.list({ + filters: { level: { in: ["deductive", "inductive"] } } + }); + + // Same filtering on semantic search (query, topK, distance, filters) + const results = await peer.conclusions.query( + "food preferences", + 10, + undefined, + { level: "deductive" } + ); + + // Conclusions about another peer, explicit only + const bobExplicit = await peer.conclusionsOf("bob").list({ + filters: { level: "explicit" } + }); +})(); +``` + + ## Error Handling Handle filter errors gracefully: diff --git a/docs/v3/documentation/features/chat.mdx b/docs/v3/documentation/features/chat.mdx index 7493e791..6aab6996 100644 --- a/docs/v3/documentation/features/chat.mdx +++ b/docs/v3/documentation/features/chat.mdx @@ -94,6 +94,43 @@ for await (const chunk of responseStream.iter_text()) { Streaming is useful for displaying real-time responses in chat interfaces or when asking complex questions that require longer answers. +## Structured Outputs + +When your application needs a machine-readable answer instead of prose, pass a schema as `response_format` and the answer is guaranteed to conform to it: + + +```python Python +from pydantic import BaseModel + +class OnboardingStatus(BaseModel): + completed: bool + remaining_steps: list[str] + +status = peer.chat( + "Has the user completed the onboarding flow?", + response_format=OnboardingStatus, +) +# status is a parsed OnboardingStatus instance +``` + +```typescript TypeScript +import { z } from 'zod'; + +const OnboardingStatus = z.object({ + completed: z.boolean(), + remainingSteps: z.array(z.string()), +}); + +const status = await peer.chat( + "Has the user completed the onboarding flow?", + { responseFormat: OnboardingStatus }, +); +// status is typed as z.infer +``` + + +The agent runs its full reasoning loop either way — only the final answer is formatted to your schema. See [Structured Outputs](/v3/documentation/features/advanced/structured-outputs) for the supported schema subset, streaming behavior, and best practices. + ## Integration Patterns ### Dynamic Prompt Enhancement diff --git a/docs/v3/documentation/reference/platform.mdx b/docs/v3/documentation/reference/platform.mdx index 75fc9bd5..abf3c646 100644 --- a/docs/v3/documentation/reference/platform.mdx +++ b/docs/v3/documentation/reference/platform.mdx @@ -60,7 +60,13 @@ The **Performance** page provides comprehensive monitoring with usage metrics, h ## 3. Manage API Keys -The [API Keys](https://app.honcho.dev/api-keys) page allows you to create and manage authentication tokens for different environments. You can create admin-level keys with full instance access or scope keys to specific `Workspaces`, `Peers`, or `Sessions`. +The [API Keys](https://app.honcho.dev/api-keys) page allows you to create and manage authentication tokens for different environments. You can create admin-level keys with full instance access or scope keys to a specific `Workspace`, `Peer`, or `Session`. + +Scoped keys are authorized by their narrowest claim and never widen to the whole workspace: + +- A **peer-scoped** key acts on its own peer, plus **read-only** access to the sessions its peer is an active member of (context, summaries, peers, its own per-session config, search, and message reads). It cannot write to those sessions or act on other peers. +- A **session-scoped** key is confined to its own session and cannot reach peer routes. +- Peer- and session-scoped keys **must carry their parent workspace** — creating one without a workspace is rejected. API Key Management Dashboard diff --git a/docs/v3/guides/integrations/codex.mdx b/docs/v3/guides/integrations/codex.mdx new file mode 100644 index 00000000..8187ca6e --- /dev/null +++ b/docs/v3/guides/integrations/codex.mdx @@ -0,0 +1,203 @@ +--- +title: "Codex" +icon: 'square-terminal' +description: "Add AI-native memory to OpenAI Codex" +sidebarTitle: 'Codex' +--- + +Give Codex long-term memory that survives context resets, session restarts, and fresh conversations. Codex remembers what you're working on, your preferences, and the decisions you've made — across every project. Lifecycle hooks capture each session to Honcho and inject the relevant context back at session start, so you never have to repeat yourself. + +## Prerequisites + +- **[Codex](https://developers.openai.com/codex) ≥ 0.136.0** +- **[Node](https://nodejs.org)** on your `PATH` (runs the installer and the hooks) + +## Quick Start + +### Step 1: Get Your Honcho API Key + +1. Go to **[app.honcho.dev](https://app.honcho.dev)** +2. Sign up or log in +3. Copy your API key (starts with `hch-`) + +### Step 2: Save Your API Key + +Your key lives in **`~/.honcho/config.json`** — the single config file every Honcho integration reads. codex-honcho takes the key straight from there. + + +**Already have your key in `~/.honcho/config.json`?** If another Honcho integration already wrote it there, there's nothing to do — skip to Step 3 and `install` picks it up automatically. + + +**First time?** Create it with the Honcho CLI: + +```bash +honcho init # prompts for your key, writes ~/.honcho/config.json + # no CLI yet? uv tool install honcho-cli && honcho init +``` + +If you'd rather write the file yourself: + +```jsonc +// ~/.honcho/config.json +{ "apiKey": "hch-your-api-key-here" } +``` + +### Step 3: Install the Plugin + +```bash +npm install -g @honcho-ai/codex-honcho +codex-honcho install # registers hooks + MCP + skill in ~/.codex +``` + +`install` copies your resolved key into `~/.codex/config.toml` so the Honcho MCP server authenticates with no environment variable to set. If you ever rotate your key, re-run `codex-honcho install` to refresh it. + +### Step 4: Restart Codex + +Restart Codex (or start a new session) to load the hooks and the `[features].hooks` flag. On your next session start you'll see Honcho memory load into context. + +### Step 5: (Optional) Tell Codex to use its memory + +The bundled `honcho-memory` skill already nudges Codex to recall and save actively. To reinforce it, add a short directive to your global Codex instructions (`~/.codex/AGENTS.md`): + +```markdown +# Honcho Memory + +You have persistent memory via Honcho. Context about me is loaded at the start +of every session — trust it and act on it; don't ask me what you already know. +Use the Honcho MCP tools (`search`, `chat`) to recall more mid-task, and +`create_conclusions` to save new preferences, decisions, and patterns as you learn them. +``` + +## What You Get + +- **Persistent Memory** — Codex remembers your preferences, projects, and context across sessions +- **Survives Context Resets** — Memory persists through `/clear`, compaction, and restarts +- **Active Recall** — Codex can search your history and query what Honcho knows about you mid-task, not just at startup +- **Git Awareness** — Optionally scope memory per branch, so feature work keeps its own context +- **Flexible Sessions** — Map memory per directory, per git branch, or per chat instance +- **Local-First Capture** — Conversations are queued to disk instantly and uploaded in the background — capture never blocks your turn or hits the network mid-conversation +- **Cross-Tool Context** — Shares `~/.honcho/config.json` with other Honcho integrations (Claude Code, Cursor, …), so context can follow you between tools + +## Configuration + +All settings live in `~/.honcho/config.json` (shared with other Honcho integrations). Codex-specific settings go under `hosts.codex`, falling back to the root fields. The hooks only ever read this file; `install` is the only writer. + +```jsonc +{ + "apiKey": "hch-…", + "peerName": "alice", // your identity (default: $USER) + "hosts": { + "codex": { + "workspace": "codex", // Honcho workspace for Codex memory + "sessionStrategy": "per-directory", + "injectPerPrompt": false, // re-inject context every turn (off by default) + "saveMessages": true // false = read memory but never write + } + } +} +``` + +### Session Strategies + +Controls how Codex conversations map to Honcho sessions: + +| Strategy | Session name | Best for | +| --- | --- | --- | +| `per-directory` (default) | `my-app` | Most users — each project accumulates its own memory | +| `git-branch` | `my-app-main` | Feature-branch workflows where context per branch matters | +| `chat-instance` | `my-app-019ea7df` | Ephemeral usage — a clean slate per conversation | + +An explicit `sessions[cwd]` mapping overrides all strategies. Environment overrides: `HONCHO_API_KEY`, `HONCHO_PEER_NAME`, `HONCHO_CONFIG_DIR`. + +## Building with Teammates + +Because `~/.honcho/config.json` is shared across Honcho hosts, teammates can collaborate by pointing at the same workspace while keeping their own identities. Each person uses their own `peerName`, so their contributions are attributed to distinct peers even when they work in the same repo. + +**Alice** (`~/.honcho/config.json`): +```json +{ + "apiKey": "hch-team-key...", + "peerName": "alice", + "hosts": { + "codex": { "workspace": "team-acme" } + } +} +``` + +**Bob** (`~/.honcho/config.json`): +```json +{ + "apiKey": "hch-team-key...", + "peerName": "bob", + "hosts": { + "codex": { "workspace": "team-acme" } + } +} +``` + +Both Alice and Bob write to the `team-acme` workspace. Working in the same repo, they share a session (named by directory, e.g. `my-app`) but appear in it as separate peers — so Honcho's dialectic reasoning can draw on context from both. + +## MCP Tools + +Once installed, Codex can call these Honcho tools directly: + +| Tool | Description | +| --- | --- | +| `search` | Semantic search across your session messages | +| `chat` | Ask Honcho a natural-language question about you | +| `get_peer_context` | Fetch the current model of you (representation + peer card) | +| `get_representation` | Lightweight representation string | +| `create_conclusions` | Save durable insights to memory | +| `list_conclusions` | List saved conclusions | +| `query_conclusions` | Semantic search across derived conclusions | +| `delete_conclusion` | Remove a conclusion by ID | + +## Commands + +| Command | Effect | +| --- | --- | +| `codex-honcho install` | Install hooks + MCP + skill | +| `codex-honcho status` | Installed components, pending queue depth, GUI link | +| `codex-honcho remove` | Strip only what this installs | + +## What Install Writes + +| Path | Change | +| --- | --- | +| `~/.codex/honcho/` | staged copy of the bundle the hooks run (kept stable across npm/npx cache eviction) | +| `~/.codex/hooks.json` | adds the four hook entries (merged; your own hooks untouched) | +| `~/.codex/config.toml` | sets `[features].hooks = true`; registers `[mcp_servers.honcho]` → `mcp.honcho.dev` (native HTTP) | +| `~/.codex/skills/honcho-memory/` | the active-recall skill | +| `~/.honcho/config.json` | persists the resolved `apiKey` + `peerName` (other fields and `hosts.*` blocks preserved) | + +`codex-honcho remove` reverses exactly these. + +## Troubleshooting + +**No memory loading / MCP not registered.** Confirm your key is in `~/.honcho/config.json` (`codex-honcho status` shows `honcho config: found`). If it's missing, run `honcho init` (or add `{ "apiKey": "hch-…" }` to the file yourself), then re-run `codex-honcho install` — without a key, install registers the hooks and skill but skips the MCP server. + +**Hooks aren't firing.** Restart Codex after installing so it loads `hooks.json` and the `[features].hooks` flag. Check `codex-honcho status` for installed components and pending queue depth. + +**Memory not persisting.** Make sure `saveMessages` isn't set to `false` under `hosts.codex`. + +## Install from a GitHub Clone (no npm) + +```bash +git clone https://github.com/plastic-labs/codex-honcho +cd codex-honcho +./install.sh # bun install + bun run bin/codex-honcho.ts install +``` + +The clone path runs the TypeScript source directly and so requires **[bun](https://bun.sh)**; it wires the hooks to `bun run /bin/codex-honcho.ts`, so keep the clone in place. The npm install instead stages the bundled `dist/codex-honcho.mjs` to `~/.codex/honcho/` and wires hooks to `node` — node-only, and stable across `npm update`, npx cache eviction, or removing the package. + +## Next Steps + + + + Source code, issues, and README. + + + + Learn about peers, sessions, and dialectic reasoning. + + diff --git a/docs/v3/guides/integrations/mcp.mdx b/docs/v3/guides/integrations/mcp.mdx index 37d2f977..10bfdfd0 100644 --- a/docs/v3/guides/integrations/mcp.mdx +++ b/docs/v3/guides/integrations/mcp.mdx @@ -224,6 +224,32 @@ Add to `~/.config/zed/settings.json`: Zed uses `context_servers` instead of `mcpServers`. Native HTTP support requires Zed v0.214.5 or later. +### Goose + +[Goose](https://goose-docs.ai/) supports remote MCP servers natively over Streamable HTTP. + +The easiest way is to run `goose configure`, choose **Add Extension → Remote Extension (Streamable HTTP)**, and enter the name `honcho`, the URI `https://mcp.honcho.dev`, and the headers `Authorization: Bearer hch-your-key-here` and `X-Honcho-User-Name: YourName`. + +Or edit your `config.yaml` directly (on Linux, `~/.config/goose/config.yaml`): + +```yaml +extensions: + honcho: + enabled: true + type: streamable_http + name: honcho + description: Honcho persistent memory & personalization + uri: https://mcp.honcho.dev + headers: + Authorization: "Bearer hch-your-key-here" + X-Honcho-User-Name: "YourName" + timeout: 60 +``` + + +To teach Goose the recommended memory flow, save the [instructions](https://raw.githubusercontent.com/plastic-labs/honcho/refs/heads/main/mcp/instructions.md) into a `.goosehints` file in your Goose config directory (or a project root). This is Goose's equivalent of Claude Desktop's "Project Instructions". Not sure of your config path? Run `goose info`. + + --- ## Optional Configuration diff --git a/docs/v3/guides/recipes/unified-memory-setup.mdx b/docs/v3/guides/recipes/unified-memory-setup.mdx index 1d547db2..ea921cc3 100644 --- a/docs/v3/guides/recipes/unified-memory-setup.mdx +++ b/docs/v3/guides/recipes/unified-memory-setup.mdx @@ -131,16 +131,17 @@ for i in range(0, len(messages), 100): session.add_messages(messages[i:i + 100]) ``` -Honcho only reasons over a peer once it accumulates ~1,000 tokens *within a single session* -([token batching](/v3/documentation/core-concepts/reasoning#token-batching)). Scope -the session to the volume you ingest: +Honcho batches reasoning until a peer accumulates ~1,000 tokens *within a single session*, +with a default age-based flush for quiet tails +([token batching](/v3/documentation/core-concepts/reasoning#token-batching)). Scope the +session to the volume you ingest: - **High-volume runs** (a day of emails, a CRM export) clear the threshold easily — a per-run session like `email-import-{date}` is fine. - **Low-volume or trickle imports** (a few short records at a time) should append to one **ongoing per-source session** (e.g. `email-import-gmail`), so content - accumulates across runs instead of fragmenting into thin sessions that each stall - below the threshold (nothing is lost — it just waits). + accumulates across runs instead of fragmenting into thin sessions that each flush + later with little context. The [Gmail](/v3/guides/gmail) and [Granola](/v3/guides/granola) guides are related import examples. diff --git a/docs/v3/openapi.json b/docs/v3/openapi.json index adcd9d29..85e4891a 100644 --- a/docs/v3/openapi.json +++ b/docs/v3/openapi.json @@ -9,7 +9,7 @@ "url": "https://honcho.dev/", "email": "hello@plasticlabs.ai" }, - "version": "3.0.7" + "version": "3.0.11" }, "servers": [ { @@ -1574,7 +1574,7 @@ "get": { "tags": ["sessions"], "summary": "Get Peer Config", - "description": "Get the configuration for a Peer in a Session.", + "description": "Get the configuration for a Peer in a Session.\n\nMember-read lets a peer-scoped key reach this route, but a peer may only\nread its own per-session config — not a co-member's. Workspace/admin and\nsession-scoped tokens (which already span the whole session) are unaffected.", "operationId": "get_peer_config_v3_workspaces__workspace_id__sessions__session_id__peers__peer_id__config_get", "security": [{ "HTTPBearer": [] }], "parameters": [ @@ -2783,6 +2783,13 @@ "anyOf": [{ "type": "string" }, { "type": "null" }], "title": "Session Id" }, + "level": { + "type": "string", + "enum": ["explicit", "deductive", "inductive", "contradiction"], + "title": "Level", + "description": "Reasoning level of the conclusion: 'explicit' (directly extracted from messages) or 'deductive'/'inductive'/'contradiction' (derived during dreaming).", + "default": "explicit" + }, "created_at": { "type": "string", "format": "date-time", diff --git a/honcho-cli/CHANGELOG.md b/honcho-cli/CHANGELOG.md new file mode 100644 index 00000000..3c2c383a --- /dev/null +++ b/honcho-cli/CHANGELOG.md @@ -0,0 +1,29 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](http://keepachangelog.com/) +and this project adheres to [Semantic Versioning](http://semver.org/). + +## [0.1.2] - 2026-07-20 + +### Added + +- Device-code OAuth login for managed Honcho servers. `honcho init` now offers browser-based login (RFC 8628 device authorization grant) when the host advertises the device grant in its OAuth authorization-server metadata; tokens are persisted to `~/.honcho/config.json` and auto-refreshed (#891) + +## [0.1.1] - 2026-06-15 + +### Fixed + +- Declare `click` as an explicit dependency. The CLI imported `click` directly but relied on it being pulled in transitively, so installs without it on the path could fail at runtime (#787) + +## [0.1.0] - 2026-04-20 + +### Added + +- Initial release of `honcho-cli` — a terminal for inspecting and managing a Honcho deployment (#424) +- `workspace`, `peer`, `session`, `message`, `conclusion`, and `config` command groups for managing resources against any Honcho server +- `init` onboarding flow that prompts for and persists connection settings, with flag/env-var pre-seeding for non-interactive use +- Per-command flags, environment variables, and a config file for pointing the CLI at different servers (local, self-hosted, or hosted) +- Rich terminal output and an agent-usage mode for scripting against the CLI +- Documentation and an agent skill for the CLI (#589) diff --git a/honcho-cli/pyproject.toml b/honcho-cli/pyproject.toml index c06f22e0..5eb859fd 100644 --- a/honcho-cli/pyproject.toml +++ b/honcho-cli/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "honcho-cli" -version = "0.1.0" +version = "0.1.2" description = "A terminal for Honcho — memory that reasons." readme = "README.md" requires-python = ">=3.11" @@ -17,6 +17,7 @@ classifiers = [ "Topic :: Software Development :: Libraries", ] dependencies = [ + "click>=8.0.0", "typer>=0.15.0", "honcho-ai>=2.0.0", "rich>=13.0.0", diff --git a/honcho-cli/src/honcho_cli/__init__.py b/honcho-cli/src/honcho_cli/__init__.py index f38546a9..81efe6f5 100644 --- a/honcho-cli/src/honcho_cli/__init__.py +++ b/honcho-cli/src/honcho_cli/__init__.py @@ -1,3 +1,3 @@ """Honcho CLI — a terminal for Honcho.""" -__version__ = "0.1.0" +__version__ = "0.1.2" diff --git a/honcho-cli/src/honcho_cli/commands/setup.py b/honcho-cli/src/honcho_cli/commands/setup.py index 6073b3e4..85dc4511 100644 --- a/honcho-cli/src/honcho_cli/commands/setup.py +++ b/honcho-cli/src/honcho_cli/commands/setup.py @@ -7,6 +7,8 @@ from __future__ import annotations import json +import time +import webbrowser import typer from honcho import ( @@ -19,13 +21,14 @@ from honcho import ( from rich.console import Console from rich.panel import Panel -from honcho_cli import __version__ +from honcho_cli import __version__, oauth from honcho_cli.branding import BANNER, BRAND, ICON_FAIL, ICON_OK, ICON_RUN -from honcho_cli.common import get_resolved_config +from honcho_cli.common import get_resolved_config, maybe_refresh_token from honcho_cli.config import ( CONFIG_FILE, DEFAULT_BASE_URL, CLIConfig, + OAuthTokens, ) from honcho_cli.output import print_error, print_result, set_json_mode, use_json @@ -117,21 +120,142 @@ def init( _console.print() _console.print() + # Non-interactive (JSON/piped) or an explicit --api-key: manual-key path. + # Device login needs a human at a browser, so it's TTY-only. + if use_json() or api_key: + _init_manual_key(key_val, url_val, file_key, file_url) + else: + _init_interactive(key_val, url_val, file_url) + + +def _init_manual_key(key_val: str, url_val: str, file_key: str, file_url: str) -> None: + """Non-interactive path: confirm/save apiKey + URL, no device login.""" final_key = _prompt_api_key(key_val) final_url = _prompt_url(url_val) - - # Persist if anything changed or if the value came from env/flag. if final_key != file_key or final_url != file_url: CLIConfig(base_url=final_url, api_key=final_key).save() if not use_json(): _console.print(f" {ICON_OK} [dim]Saved to {CONFIG_FILE}[/dim]") - _check_connection(final_url, final_key) - if use_json(): print_result({"apiKey": _redact(final_key), "baseUrl": final_url}) +def _init_interactive(key_val: str, url_val: str, file_url: str) -> None: + """Interactive path: URL first (device flow needs the host), then auth method.""" + final_url = _prompt_url(url_val) + existing = CLIConfig.load() + has_creds = bool(key_val) or bool(existing.oauth and existing.oauth.access_token) + # only offer browser login if the host advertises the device grant (managed) + device_available = oauth.supports_device_login(final_url) + method = _prompt_auth_method(has_creds, device_available) + + if method == "keep": + if final_url != file_url: + existing.base_url = final_url + existing.save() + _console.print(f" {ICON_OK} [dim]Saved to {CONFIG_FILE}[/dim]") + # refresh an expired token so "keep" behaves like every live command; + # a failed refresh surfaces as the connectivity check below, not an abort + try: + maybe_refresh_token(existing) + except typer.Exit: + pass + _check_connection(final_url, existing.resolved_api_key()) + return + + if method == "device": + tokens = _device_login(final_url) + CLIConfig(base_url=final_url, oauth=tokens).save() + _console.print(f" {ICON_OK} [dim]Saved to {CONFIG_FILE}[/dim]") + _check_connection(final_url, tokens.access_token) + return + + # paste a key + final_key = _prompt_api_key("") + CLIConfig(base_url=final_url, api_key=final_key).save() + _console.print(f" {ICON_OK} [dim]Saved to {CONFIG_FILE}[/dim]") + _check_connection(final_url, final_key) + + +def _prompt_auth_method(has_creds: bool, device_available: bool) -> str: + """Ask how to authenticate. Returns ``device`` / ``key`` / ``keep``. + + ``device`` is only offered when the host advertises the device grant; when + it doesn't, pasting a key is the only login path. + """ + _console.print(" [dim]How do you want to authenticate?[/dim]") + options: list[str] = [] + if device_available: + options.append("device") + _console.print(f" [dim]({len(options)})[/dim] Log in with your browser (device code)") + options.append("key") + _console.print(f" [dim]({len(options)})[/dim] Paste an API key") + if has_creds: + options.append("keep") + _console.print(f" [dim]({len(options)})[/dim] Keep current credentials") + # default to keeping existing creds so a returning user pressing Enter doesn't + # get dropped into an unwanted browser login that overwrites them + default = str(options.index("keep") + 1) if "keep" in options else "1" + choice = typer.prompt(" Choice", default=default, show_default=True, prompt_suffix=": ").strip() + try: + idx = int(choice) + except ValueError: + return options[0] + # explicit 1..len bounds — bare `options[idx - 1]` would let "0"/negatives + # wrap to the tail of the list via Python's negative indexing + if 1 <= idx <= len(options): + return options[idx - 1] + return options[0] + + +def _device_login(base_url: str) -> OAuthTokens: + """Run the device-authorization flow and return the minted tokens. + + Prints the user code + verification URL, opens the browser best-effort, and + blocks on the poll loop until the user approves. Exits non-zero on denial, + expiry, or interrupt. + """ + endpoints = oauth.resolve_endpoints(base_url) + try: + device = oauth.request_device_code(endpoints) + except oauth.OAuthFlowError as e: + _console.print(f" {ICON_FAIL} [red]Could not start device login[/red]: {e}") + raise typer.Exit(1) + + _console.print() + _console.print(f" Enter this code to authorize: [bold {BRAND}]{device.user_code}[/bold {BRAND}]") + _console.print(f" [dim]at[/dim] {device.verification_uri}") + _console.print() + try: + webbrowser.open(device.verification_uri_complete) + except Exception: + pass # headless is expected — the URL is printed above + + try: + with _console.status("Waiting for approval…", spinner="dots"): + tokens = oauth.poll_for_token(endpoints, device) + except oauth.AccessDenied: + _console.print(f" {ICON_FAIL} [red]Authorization denied[/red]") + raise typer.Exit(1) + except (oauth.DeviceCodeExpired, oauth.AuthorizationTimeout): + _console.print(f" {ICON_FAIL} [red]Code expired[/red] — run `honcho init` to try again") + raise typer.Exit(1) + except oauth.OAuthFlowError as e: + _console.print(f" {ICON_FAIL} [red]Login failed[/red]: {e}") + raise typer.Exit(1) + except KeyboardInterrupt: + _console.print(f" {ICON_FAIL} [red]Cancelled[/red]") + raise typer.Exit(1) + + return OAuthTokens.from_response( + tokens, + client_id=endpoints.client_id, + scope_fallback=endpoints.scope, + host=base_url, + ) + + def _prompt_api_key(value: str) -> str: """Prompt for API key. @@ -206,6 +330,21 @@ def _check_connection(base_url: str, api_key: str) -> None: # --------------------------------------------------------------------------- # # honcho doctor +def _auth_mode_detail(config: CLIConfig) -> str: + """Human summary of which credential the CLI will use.""" + tokens = config.usable_oauth() + if tokens is not None: + if tokens.access_valid(): + secs = max(int(tokens.access_expires_at - time.time()), 0) + return f"OAuth device token (expires in {secs // 60}m)" + if config.api_key: + return "API key (OAuth token expired)" + return "OAuth device token (expired — will refresh)" + if config.api_key: + return "API key" + return "missing — run `honcho init`" + + def doctor( json_output: bool = typer.Option(False, "--json", help="Force JSON output"), ) -> None: @@ -230,23 +369,30 @@ def doctor( _console.print(f"\n[bold {BRAND}]Honcho Doctor[/bold {BRAND}]\n") config = get_resolved_config() + # Refresh an expired OAuth token if we can; a failure surfaces as a failed + # connectivity check below rather than aborting the diagnostic. + try: + maybe_refresh_token(config) + except typer.Exit: + pass + key = config.resolved_api_key() + _add("Config file", CONFIG_FILE.exists(), str(CONFIG_FILE) if CONFIG_FILE.exists() else f"{CONFIG_FILE} not found") - _add("API key configured", bool(config.api_key), - "set" if config.api_key else "missing — run `honcho init`") + _add("Credentials configured", bool(key), _auth_mode_detail(config)) - if config.base_url and config.api_key: - _add("API connectivity", *_test_connection(config.base_url, config.api_key)) + if config.base_url and key: + _add("API connectivity", *_test_connection(config.base_url, key)) else: - _add("API connectivity", False, "skipped — no base_url or api_key") + _add("API connectivity", False, "skipped — no base_url or credentials") # Workspace / peer / queue run only when scoped via -w / -p. ws_ok, client = False, None - if config.workspace_id and config.api_key: + if config.workspace_id and key: try: - client = Honcho(base_url=config.base_url, api_key=config.api_key, workspace_id=config.workspace_id) + client = Honcho(base_url=config.base_url, api_key=key, workspace_id=config.workspace_id) client.get_configuration() ws_ok = True _add("Workspace reachable", True, config.workspace_id) @@ -280,7 +426,7 @@ def doctor( _console.print(f"\n [{color}]{passed}/{total}[/{color}] checks passed{hint}\n") # Config file + API connectivity are hard requirements. - critical = {"Config file", "API key configured", "API connectivity"} + critical = {"Config file", "Credentials configured", "API connectivity"} if config.workspace_id: critical.add("Workspace reachable") if any(not c["ok"] for c in checks if c["check"] in critical): diff --git a/honcho-cli/src/honcho_cli/common.py b/honcho-cli/src/honcho_cli/common.py index 2680a869..d87a4be7 100644 --- a/honcho-cli/src/honcho_cli/common.py +++ b/honcho-cli/src/honcho_cli/common.py @@ -12,13 +12,15 @@ no-op if the same flag was already set at an outer level. from __future__ import annotations +from dataclasses import replace from typing import Optional import typer from honcho import Honcho -from honcho_cli.config import CLIConfig, get_client_kwargs +from honcho_cli import oauth +from honcho_cli.config import CLIConfig, OAuthTokens, get_client_kwargs from honcho_cli.output import print_error, set_json_mode from honcho_cli.validation import validate_resource_id @@ -50,6 +52,50 @@ def get_resolved_config(): return config +def maybe_refresh_token(config: CLIConfig) -> None: + """Refresh an expired OAuth access token in place and persist it. + + No-op when there is no grant for the current host or the token is still + valid. A dead grant degrades to the saved apiKey with a warning; exits + only when nothing is left to authenticate with. + """ + tokens = config.usable_oauth() + if tokens is None or tokens.access_valid(): + return + + if tokens.refresh_token: + endpoints = oauth.resolve_endpoints(config.base_url) + if tokens.client_id: + endpoints = replace(endpoints, client_id=tokens.client_id) + try: + refreshed = oauth.refresh_access_token(endpoints, tokens.refresh_token) + except oauth.OAuthFlowError: + refreshed = None + if refreshed is not None: + # rotation-safe: persist the (possibly new) refresh token before + # it's reused; keep the old one if the server didn't rotate + # (refresh_token is optional) + config.oauth = OAuthTokens.from_response( + refreshed, + client_id=tokens.client_id, + scope_fallback=tokens.scope, + refresh_fallback=tokens.refresh_token, + host=tokens.host, + ) + config.save() + return + + if config.api_key: + typer.echo( + "OAuth session expired; using the saved API key. " + "Run `honcho init` to log in again.", + err=True, + ) + return + print_error("SESSION_EXPIRED", "OAuth session expired. Run `honcho init` to log in again.") + raise typer.Exit(1) + + def get_client(*, require_workspace: bool = True): """Create a Honcho client from resolved config. @@ -65,6 +111,7 @@ def get_client(*, require_workspace: bool = True): "No workspace scoped. Pass --workspace/-w or set HONCHO_WORKSPACE_ID.", ) raise typer.Exit(1) + maybe_refresh_token(config) return Honcho(**get_client_kwargs(config)), config diff --git a/honcho-cli/src/honcho_cli/config.py b/honcho-cli/src/honcho_cli/config.py index a0c64864..7cd103bd 100644 --- a/honcho-cli/src/honcho_cli/config.py +++ b/honcho-cli/src/honcho_cli/config.py @@ -1,11 +1,18 @@ """Configuration management for Honcho CLI. -Config stored at ``~/.honcho/config.json`` with env var overrides. +Config stored at ``~/.honcho/config.json`` with env var overrides. The config +directory defaults to ``~/.honcho`` and can be relocated with `HONCHO_CONFIG_DIR` -The CLI owns exactly two top-level keys in that file: +The CLI owns these top-level keys in that file: - apiKey -- Honcho admin JWT environmentUrl -- Honcho API URL (full URL, e.g. https://api.honcho.dev) + oauth -- OAuth device-grant tokens (accessToken, refreshToken, + accessExpiresAt, clientId, scope, host), written by + device login + +``apiKey`` (manual admin JWT) is shared with sibling tools: the CLI writes it +on paste-key login and reads it as a fallback, but never deletes it. A live +OAuth token takes precedence over ``apiKey`` for the CLI's own calls. All other top-level keys (``hosts``, ``sessions``, ``saveMessages``, ``sessionStrategy``, …) are written by sibling Honcho tools and are @@ -20,14 +27,44 @@ from __future__ import annotations import json import os +import time from dataclasses import dataclass, fields from pathlib import Path +from typing import TYPE_CHECKING -CONFIG_DIR = Path.home() / ".honcho" +if TYPE_CHECKING: + from honcho_cli.oauth import TokenResponse + +def _config_dir() -> Path: + """Config directory: ``$HONCHO_CONFIG_DIR`` if set, else ``~/.honcho``.""" + override = os.environ.get("HONCHO_CONFIG_DIR") + return Path(override).expanduser() if override else Path.home() / ".honcho" + + +CONFIG_DIR = _config_dir() CONFIG_FILE = CONFIG_DIR / "config.json" DEFAULT_BASE_URL = "https://api.honcho.dev" + +def _redact_token(token: str) -> str: + """Show ``***`` — enough to compare tokens without leaking the body.""" + if not token: + return "" + return "***" + token[-4:] if len(token) > 4 else "***" + + +def _coerce_epoch(value: object) -> float: + """Parse a persisted epoch-seconds value, treating garbage as expired (0).""" + if isinstance(value, (int, float)): + return float(value) + if isinstance(value, str): + try: + return float(value) + except ValueError: + return 0.0 + return 0.0 + # Env var mapping for runtime overrides. # # Resolution order: flag > env var > config file > default. @@ -40,6 +77,59 @@ ENV_MAP: dict[str, str] = { } +@dataclass +class OAuthTokens: + """Device-grant tokens persisted under the config ``oauth`` key.""" + + access_token: str = "" + refresh_token: str = "" + access_expires_at: float = 0.0 # epoch seconds + client_id: str = "" + scope: str = "" + host: str = "" # base_url the grant was minted against + + def matches_host(self, base_url: str) -> bool: + """True when the grant belongs to ``base_url``. + + Tokens are host-scoped — a staging grant must not be sent to prod. + Legacy blocks with no recorded host are trusted. + """ + return not self.host or self.host.rstrip("/") == base_url.rstrip("/") + + def access_valid(self, skew: int = 60) -> bool: + """True while the access token is present and not within ``skew`` of expiry. + + Checks the expiry timestamp recorded at mint time, not the token + itself — the server is the real authority, so a wrong answer here + costs at most an extra refresh or a 401. + """ + return bool(self.access_token) and time.time() < self.access_expires_at - skew + + @classmethod + def from_response( + cls, + resp: TokenResponse, + *, + client_id: str, + scope_fallback: str = "", + refresh_fallback: str = "", + host: str = "", + ) -> OAuthTokens: + """Build persisted tokens from a token response. + + ``refresh_fallback`` keeps the prior refresh token when the server + doesn't rotate one (optional on the refresh grant, RFC 6749 §5.1). + """ + return cls( + access_token=resp.access_token, + refresh_token=resp.refresh_token or refresh_fallback, + access_expires_at=time.time() + resp.expires_in, + client_id=client_id, + scope=resp.scope or scope_fallback, + host=host, + ) + + @dataclass class CLIConfig: """CLI configuration with layered resolution: flag > env > file > default. @@ -54,6 +144,33 @@ class CLIConfig: workspace_id: str = "" peer_id: str = "" session_id: str = "" + oauth: OAuthTokens | None = None + + def usable_oauth(self) -> OAuthTokens | None: + """The OAuth grant, if present and bound to the current host.""" + if ( + self.oauth + and self.oauth.access_token + and self.oauth.matches_host(self.base_url) + ): + return self.oauth + return None + + def resolved_api_key(self) -> str: + """The key handed to the SDK: a live OAuth token wins, else apiKey. + + An expired grant loses to a saved apiKey (a dead grant degrades to the + shared key) but still wins over nothing, since the server is the final + judge. + """ + tokens = self.usable_oauth() + if tokens and tokens.access_valid(): + return tokens.access_token + if self.api_key: + return self.api_key + if tokens: + return tokens.access_token + return "" @classmethod def load(cls) -> CLIConfig: @@ -74,6 +191,16 @@ class CLIConfig: key = data.get("apiKey") if isinstance(key, str): config.api_key = key + oauth = data.get("oauth") + if isinstance(oauth, dict) and oauth.get("accessToken"): + config.oauth = OAuthTokens( + access_token=str(oauth.get("accessToken", "")), + refresh_token=str(oauth.get("refreshToken", "")), + access_expires_at=_coerce_epoch(oauth.get("accessExpiresAt")), + client_id=str(oauth.get("clientId", "")), + scope=str(oauth.get("scope", "")), + host=str(oauth.get("host", "")), + ) for fld_name, env_var in ENV_MAP.items(): val = os.environ.get(env_var) @@ -88,10 +215,12 @@ class CLIConfig: return config def save(self) -> None: - """Write ``apiKey`` + ``environmentUrl`` to config.json. + """Write ``environmentUrl`` + credentials to config.json. Preserves unrelated top-level keys (``hosts``, ``sessions``, ``saveMessages``, ``sessionStrategy``, …) that other tools write. + ``apiKey`` is written when set but never removed — sibling tools read + it. The ``oauth`` block is CLI-owned and dropped when empty. """ CONFIG_DIR.mkdir(parents=True, exist_ok=True) @@ -108,8 +237,18 @@ class CLIConfig: data["environmentUrl"] = self.base_url if self.api_key: data["apiKey"] = self.api_key + + if self.oauth and self.oauth.access_token: + data["oauth"] = { + "accessToken": self.oauth.access_token, + "refreshToken": self.oauth.refresh_token, + "accessExpiresAt": self.oauth.access_expires_at, + "clientId": self.oauth.client_id, + "scope": self.oauth.scope, + "host": self.oauth.host, + } else: - data.pop("apiKey", None) + data.pop("oauth", None) CONFIG_FILE.write_text(json.dumps(data, indent=2) + "\n") # API key in plaintext — restrict to the owner on multi-user hosts. @@ -124,7 +263,7 @@ class CLIConfig: Only includes fields that have a value set — per-command fields (workspace_id, peer_id, session_id) are omitted when empty. """ - d: dict[str, str] = {} + result: dict[str, str] = {} for fld in fields(self): val = getattr(self, fld.name) if not val: @@ -132,10 +271,12 @@ class CLIConfig: if fld.name == "api_key": # Show ``***`` only — enough to compare keys without # leaking the header or body of the JWT. - d[fld.name] = "***" + val[-4:] if len(val) > 4 else "***" + result[fld.name] = _redact_token(val) + elif fld.name == "oauth": + result[fld.name] = _redact_token(val.access_token) else: - d[fld.name] = val - return d + result[fld.name] = val + return result def get_client_kwargs(config: CLIConfig) -> dict: @@ -143,8 +284,9 @@ def get_client_kwargs(config: CLIConfig) -> dict: kwargs: dict = {} if config.base_url: kwargs["base_url"] = config.base_url - if config.api_key: - kwargs["api_key"] = config.api_key + api_key = config.resolved_api_key() + if api_key: + kwargs["api_key"] = api_key if config.workspace_id: kwargs["workspace_id"] = config.workspace_id return kwargs diff --git a/honcho-cli/src/honcho_cli/oauth.py b/honcho-cli/src/honcho_cli/oauth.py new file mode 100644 index 00000000..37e50d36 --- /dev/null +++ b/honcho-cli/src/honcho_cli/oauth.py @@ -0,0 +1,260 @@ +"""OAuth 2.0 Device Authorization Grant (RFC 8628) client for the CLI. + +Transport-only: HTTP calls plus the poll loop, no Typer or config writes, so it +can be unit-tested by mocking httpx. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +import httpx + +DEVICE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:device_code" + +DEFAULT_CLIENT_ID = "honcho-cli" +DEFAULT_SCOPE = "write" + +# self-declared requesting surface; tells the consent screen not to offer config +# delivery (a CLI has nowhere to write it) +DEVICE_SOURCE = "honcho-cli" + +# extra seconds added to the poll interval on a slow_down response (RFC 8628 §3.5) +SLOW_DOWN_STEP = 5 + + +class OAuthFlowError(Exception): + """A device-flow request failed. ``error`` is the RFC error code when known.""" + + def __init__(self, error: str, description: str | None = None): + self.error: str = error + self.description: str | None = description + super().__init__(description or error) + + +class AccessDenied(OAuthFlowError): + """The user denied the authorization request.""" + + +class DeviceCodeExpired(OAuthFlowError): + """The device code expired before the user approved it.""" + + +class AuthorizationTimeout(OAuthFlowError): + """Polling ran past the device code's lifetime with no decision.""" + + +@dataclass(frozen=True) +class Endpoints: + """Resolved authorization-server URLs and client identity.""" + + device_auth_url: str + token_url: str + client_id: str + scope: str + + +@dataclass(frozen=True) +class DeviceCode: + """RFC 8628 §3.2 device authorization response.""" + + device_code: str + user_code: str + verification_uri: str + verification_uri_complete: str + expires_in: int + interval: int + + +@dataclass(frozen=True) +class TokenResponse: + """An access/refresh token pair minted for a grant.""" + + access_token: str + refresh_token: str + expires_in: int + scope: str + config: dict[str, Any] = field(default_factory=dict) + + +def resolve_endpoints(base_url: str) -> Endpoints: + """Derive OAuth endpoints and client identity from the API ``base_url``.""" + host = base_url.rstrip("/") + return Endpoints( + device_auth_url=f"{host}/oauth/device_authorization", + token_url=f"{host}/oauth/token", + client_id=DEFAULT_CLIENT_ID, + scope=DEFAULT_SCOPE, + ) + + +# RFC 8414 authorization-server metadata; presence of the device grant tells us +# whether this host can do browser login at all (managed only, not core) +AUTH_SERVER_METADATA_PATH = "/.well-known/oauth-authorization-server" + + +def supports_device_login(base_url: str, *, timeout: float = 5.0) -> bool: + """Whether the host advertises the device grant in its RFC 8414 metadata. + + Fails closed: any connection error, non-200, unparseable body, or missing + capability returns False, so self-hosted / non-managed instances simply + don't offer device login. + """ + host = base_url.rstrip("/") + try: + resp = httpx.get(f"{host}{AUTH_SERVER_METADATA_PATH}", timeout=timeout) + except httpx.HTTPError: + return False + if resp.status_code != 200: + return False + try: + body = resp.json() + except ValueError: + return False + grants = body.get("grant_types_supported") if isinstance(body, dict) else None + return isinstance(grants, list) and DEVICE_GRANT_TYPE in grants + + +def _post(url: str, data: dict[str, str]) -> httpx.Response: + """POST form data, surfacing transport failures as ``OAuthFlowError``. + + Connection refusals, DNS failures, and timeouts would otherwise escape as + raw ``httpx.HTTPError`` past callers that only catch ``OAuthFlowError``. + """ + try: + return httpx.post(url, data=data) + except httpx.HTTPError as e: + raise OAuthFlowError("connection_error", f"could not reach {url}: {e}") from e + + +def _error_from_response(resp: httpx.Response) -> tuple[str, str | None]: + """Pull ``(error, error_description)`` out of an OAuth error body.""" + try: + body = resp.json() + except ValueError: + return "invalid_response", resp.text[:200] or None + if isinstance(body, dict) and body.get("error"): + return str(body["error"]), body.get("error_description") + return "invalid_response", None + + +def request_device_code(endpoints: Endpoints) -> DeviceCode: + """Request a device + user code pair (RFC 8628 §3.1).""" + resp = _post( + endpoints.device_auth_url, + { + "client_id": endpoints.client_id, + "scope": endpoints.scope, + "source": DEVICE_SOURCE, + }, + ) + if resp.status_code != 200: + error, desc = _error_from_response(resp) + raise OAuthFlowError(error, desc) + try: + body = resp.json() + return DeviceCode( + device_code=body["device_code"], + user_code=body["user_code"], + verification_uri=body["verification_uri"], + verification_uri_complete=body.get( + "verification_uri_complete", body["verification_uri"] + ), + expires_in=int(body["expires_in"]), + interval=int(body["interval"]), + ) + except (KeyError, TypeError, ValueError) as e: + raise OAuthFlowError( + "invalid_response", f"malformed device authorization response: {e}" + ) from e + + +def _token_from_body(body: dict[str, Any]) -> TokenResponse: + # refresh_token is optional on the refresh grant (RFC 6749 §5.1); a + # malformed/missing field is a server fault, surfaced as OAuthFlowError so + # callers' existing handling catches it instead of a raw KeyError/ValueError + try: + return TokenResponse( + access_token=body["access_token"], + refresh_token=body.get("refresh_token", ""), + expires_in=int(body["expires_in"]), + scope=body.get("scope", ""), + config=body.get("config") or {}, + ) + except (KeyError, TypeError, ValueError) as e: + raise OAuthFlowError("invalid_response", f"malformed token response: {e}") from e + + +def poll_for_token( + endpoints: Endpoints, + device: DeviceCode, + *, + sleep: Callable[[float], None] = time.sleep, + monotonic: Callable[[], float] = time.monotonic, +) -> TokenResponse: + """Poll the token endpoint until the grant is approved (RFC 8628 §3.4/§3.5). + + Sleeps ``interval`` between polls, bumping it on ``slow_down``. Raises + ``AccessDenied`` / ``DeviceCodeExpired`` / ``AuthorizationTimeout`` on the + terminal outcomes. ``sleep`` / ``monotonic`` are injectable for tests. + """ + interval = device.interval + deadline = monotonic() + device.expires_in + while True: + if monotonic() >= deadline: + raise AuthorizationTimeout("expired_token", "Timed out waiting for approval") + sleep(interval) + resp = _post( + endpoints.token_url, + { + "grant_type": DEVICE_GRANT_TYPE, + "device_code": device.device_code, + "client_id": endpoints.client_id, + }, + ) + if resp.status_code == 200: + try: + body = resp.json() + except ValueError as e: + raise OAuthFlowError("invalid_response", "non-JSON token response") from e + return _token_from_body(body) + + error, desc = _error_from_response(resp) + if error == "authorization_pending": + continue + if error == "slow_down": + interval += SLOW_DOWN_STEP + continue + if error == "access_denied": + raise AccessDenied(error, desc) + if error == "expired_token": + raise DeviceCodeExpired(error, desc) + raise OAuthFlowError(error, desc) + + +def refresh_access_token(endpoints: Endpoints, refresh_token: str) -> TokenResponse: + """Exchange a refresh token for a fresh access/refresh pair. + + The response may rotate the refresh token; the caller must persist the + returned ``refresh_token`` before reusing it — replaying a superseded one + revokes the grant. + """ + resp = _post( + endpoints.token_url, + { + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": endpoints.client_id, + }, + ) + if resp.status_code != 200: + error, desc = _error_from_response(resp) + raise OAuthFlowError(error, desc) + try: + body = resp.json() + except ValueError as e: + raise OAuthFlowError("invalid_response", "non-JSON token response") from e + return _token_from_body(body) diff --git a/honcho-cli/tests/conftest.py b/honcho-cli/tests/conftest.py new file mode 100644 index 00000000..5f121748 --- /dev/null +++ b/honcho-cli/tests/conftest.py @@ -0,0 +1,21 @@ +"""Shared test fixtures.""" + +from __future__ import annotations + +import pytest + +from honcho_cli import common +from honcho_cli.output import set_json_mode + + +@pytest.fixture(autouse=True) +def _reset_cli_globals(): + """Reset process-global CLI state between tests. + + ``_global_overrides`` (set by ``-w``/``-p``/``-s`` flags) and the JSON-mode + flag are module globals that leak across tests otherwise — a workspace set + by one test would silently satisfy the next test's workspace check. + """ + yield + common._global_overrides.update(workspace=None, peer=None, session=None) + set_json_mode(False) diff --git a/honcho-cli/tests/test_common.py b/honcho-cli/tests/test_common.py new file mode 100644 index 00000000..f5b2ec27 --- /dev/null +++ b/honcho-cli/tests/test_common.py @@ -0,0 +1,126 @@ +"""Tests for the client factory's transparent OAuth refresh.""" + +from __future__ import annotations + +import json +import os +import time +from unittest.mock import patch + +import pytest +import typer +from honcho_cli import common +from honcho_cli.config import CLIConfig, OAuthTokens +from honcho_cli.oauth import OAuthFlowError, TokenResponse + + +@pytest.fixture +def cfg_path(tmp_path, monkeypatch): + f = tmp_path / "config.json" + monkeypatch.setattr("honcho_cli.config.CONFIG_FILE", f) + monkeypatch.setattr("honcho_cli.config.CONFIG_DIR", tmp_path) + for k in [k for k in os.environ if k.startswith("HONCHO_")]: + monkeypatch.delenv(k) + return f + + +def _cfg(expires_at: float) -> CLIConfig: + return CLIConfig( + base_url="http://localhost:8000", + oauth=OAuthTokens( + access_token="old-at", + refresh_token="old-rt", + access_expires_at=expires_at, + client_id="honcho-cli", + scope="write", + ), + ) + + +def test_valid_token_is_not_refreshed(cfg_path): + config = _cfg(time.time() + 3600) + with patch("honcho_cli.oauth.refresh_access_token") as refresh: + common.maybe_refresh_token(config) + refresh.assert_not_called() + + +def test_expired_token_refreshes_despite_manual_key(cfg_path): + """OAuth wins over apiKey now, so the grant is kept alive even with a key set.""" + config = _cfg(time.time() - 100) + config.api_key = "manual" + rotated = TokenResponse( + access_token="new-at", refresh_token="new-rt", expires_in=3600, scope="write" + ) + with patch("honcho_cli.oauth.refresh_access_token", return_value=rotated) as refresh: + common.maybe_refresh_token(config) + refresh.assert_called_once() + assert config.resolved_api_key() == "new-at" + + +def test_host_mismatch_skips_refresh(cfg_path): + """A grant minted for another host is ignored — no refresh, apiKey covers this one.""" + config = _cfg(time.time() - 100) + config.oauth.host = "https://staging.example.com" + config.api_key = "manual" + with patch("honcho_cli.oauth.refresh_access_token") as refresh: + common.maybe_refresh_token(config) + refresh.assert_not_called() + assert config.resolved_api_key() == "manual" + + +def test_expired_token_refreshes_and_persists(cfg_path): + config = _cfg(time.time() - 100) + rotated = TokenResponse( + access_token="new-at", + refresh_token="new-rt", + expires_in=3600, + scope="write", + ) + with patch("honcho_cli.oauth.refresh_access_token", return_value=rotated) as refresh: + common.maybe_refresh_token(config) + + # used the stored refresh token + client_id + _endpoints, sent_rt = refresh.call_args.args + assert sent_rt == "old-rt" + assert _endpoints.client_id == "honcho-cli" + + # in-memory config updated with the rotated pair + assert config.oauth.access_token == "new-at" + assert config.oauth.refresh_token == "new-rt" + assert config.oauth.access_valid() + + # rotation persisted to disk before reuse + on_disk = json.loads(cfg_path.read_text())["oauth"] + assert on_disk["accessToken"] == "new-at" + assert on_disk["refreshToken"] == "new-rt" + + +def test_refresh_failure_exits(cfg_path): + config = _cfg(time.time() - 100) + with patch("honcho_cli.oauth.refresh_access_token", side_effect=OAuthFlowError("invalid_grant")): + with pytest.raises(typer.Exit): + common.maybe_refresh_token(config) + + +def test_refresh_failure_falls_back_to_api_key(cfg_path): + """A dead grant degrades to the saved apiKey instead of aborting.""" + config = _cfg(time.time() - 100) + config.api_key = "manual" + with patch("honcho_cli.oauth.refresh_access_token", side_effect=OAuthFlowError("invalid_grant")): + common.maybe_refresh_token(config) # must not raise + assert config.resolved_api_key() == "manual" + + +def test_missing_refresh_token_exits(cfg_path): + config = _cfg(time.time() - 100) + config.oauth.refresh_token = "" + with pytest.raises(typer.Exit): + common.maybe_refresh_token(config) + + +def test_missing_refresh_token_falls_back_to_api_key(cfg_path): + config = _cfg(time.time() - 100) + config.oauth.refresh_token = "" + config.api_key = "manual" + common.maybe_refresh_token(config) # must not raise + assert config.resolved_api_key() == "manual" diff --git a/honcho-cli/tests/test_config.py b/honcho-cli/tests/test_config.py index bdac2bc4..5b90087d 100644 --- a/honcho-cli/tests/test_config.py +++ b/honcho-cli/tests/test_config.py @@ -2,9 +2,12 @@ import json import os +import time +from pathlib import Path import pytest -from honcho_cli.config import CLIConfig +from honcho_cli.config import CLIConfig, OAuthTokens, _config_dir +from honcho_cli.oauth import TokenResponse @pytest.fixture @@ -18,6 +21,20 @@ def cfg_path(tmp_path, monkeypatch): return f +class TestConfigDir: + def test_defaults_to_dot_honcho(self, monkeypatch): + monkeypatch.delenv("HONCHO_CONFIG_DIR", raising=False) + assert _config_dir() == Path.home() / ".honcho" + + def test_honcho_config_dir_override(self, monkeypatch, tmp_path): + monkeypatch.setenv("HONCHO_CONFIG_DIR", str(tmp_path / "profile")) + assert _config_dir() == tmp_path / "profile" + + def test_expands_user_in_override(self, monkeypatch): + monkeypatch.setenv("HONCHO_CONFIG_DIR", "~/.honcho-test") + assert _config_dir() == Path.home() / ".honcho-test" + + class TestLoad: def test_defaults_when_no_file(self, cfg_path): loaded = CLIConfig.load() @@ -44,6 +61,31 @@ class TestLoad: assert loaded.api_key == "env-key" assert loaded.base_url == "http://localhost:8000" + def test_empty_env_var_popped_from_environ(self, cfg_path, monkeypatch): + """Empty HONCHO_* vars are removed so the SDK doesn't crash on them.""" + cfg_path.write_text(json.dumps({"apiKey": "file-key"})) + monkeypatch.setenv("HONCHO_API_KEY", "") + loaded = CLIConfig.load() + assert "HONCHO_API_KEY" not in os.environ + assert loaded.api_key == "file-key" + + def test_garbage_access_expires_at_treated_as_expired(self, cfg_path): + """Hand-edited/corrupt expiry degrades to the refresh path, not a crash.""" + cfg_path.write_text(json.dumps( + {"oauth": {"accessToken": "x", "accessExpiresAt": "not-a-number"}} + )) + loaded = CLIConfig.load() + assert loaded.oauth is not None + assert loaded.oauth.access_valid() is False + + def test_numeric_string_access_expires_at_parses(self, cfg_path): + cfg_path.write_text(json.dumps( + {"oauth": {"accessToken": "x", "accessExpiresAt": "12345"}} + )) + loaded = CLIConfig.load() + assert loaded.oauth is not None + assert loaded.oauth.access_expires_at == 12345.0 + class TestSave: def test_writes_only_cli_owned_keys(self, cfg_path): @@ -104,6 +146,117 @@ def test_api_key_redaction_empty_omitted(): assert "api_key" not in CLIConfig(api_key="").redacted() +class TestOAuth: + def _tokens(self, expires_at: float) -> OAuthTokens: + return OAuthTokens( + access_token="hch-at-x", + refresh_token="hch-rt-x", + access_expires_at=expires_at, + client_id="honcho-cli", + scope="write", + ) + + def test_round_trips_oauth_block(self, cfg_path): + CLIConfig(base_url="http://localhost:8000", oauth=self._tokens(9999999999)).save() + loaded = CLIConfig.load() + assert loaded.oauth is not None + assert loaded.oauth.access_token == "hch-at-x" + assert loaded.oauth.refresh_token == "hch-rt-x" + assert loaded.oauth.client_id == "honcho-cli" + + def test_oauth_persists_camelcase_keys(self, cfg_path): + CLIConfig(base_url="http://localhost:8000", oauth=self._tokens(1234)).save() + on_disk = json.loads(cfg_path.read_text())["oauth"] + assert set(on_disk) == {"accessToken", "refreshToken", "accessExpiresAt", "clientId", "scope", "host"} + + def test_save_preserves_foreign_keys_with_oauth(self, cfg_path): + cfg_path.write_text(json.dumps({"hosts": {"claude_code": {"peerName": "u"}}})) + CLIConfig(base_url="http://localhost:8000", oauth=self._tokens(1234)).save() + on_disk = json.loads(cfg_path.read_text()) + assert on_disk["hosts"] == {"claude_code": {"peerName": "u"}} + assert "oauth" in on_disk + + def test_empty_oauth_is_dropped(self, cfg_path): + cfg_path.write_text(json.dumps({"oauth": {"accessToken": "old"}})) + CLIConfig(base_url="http://localhost:8000").save() + assert "oauth" not in json.loads(cfg_path.read_text()) + + def test_api_key_preserved_on_device_login(self, cfg_path): + """apiKey is shared with sibling tools — device login must not delete it.""" + cfg_path.write_text(json.dumps({"apiKey": "shared-key"})) + CLIConfig(base_url="http://localhost:8000", oauth=self._tokens(9999999999)).save() + on_disk = json.loads(cfg_path.read_text()) + assert on_disk["apiKey"] == "shared-key" + assert on_disk["oauth"]["accessToken"] == "hch-at-x" + + def test_resolved_api_key_prefers_live_oauth(self, cfg_path): + cfg = CLIConfig(api_key="manual", oauth=self._tokens(9999999999)) + assert cfg.resolved_api_key() == "hch-at-x" + + def test_resolved_api_key_expired_oauth_falls_back_to_api_key(self, cfg_path): + cfg = CLIConfig(api_key="manual", oauth=self._tokens(time.time() - 100)) + assert cfg.resolved_api_key() == "manual" + + def test_resolved_api_key_host_mismatch_falls_back_to_api_key(self, cfg_path): + tokens = self._tokens(9999999999) + tokens.host = "https://staging.example.com" + cfg = CLIConfig( + base_url="https://api.honcho.dev", api_key="manual", oauth=tokens + ) + assert cfg.resolved_api_key() == "manual" + + def test_resolved_api_key_expired_oauth_wins_over_nothing(self, cfg_path): + cfg = CLIConfig(oauth=self._tokens(time.time() - 100)) + assert cfg.resolved_api_key() == "hch-at-x" + + def test_resolved_api_key_falls_back_to_oauth(self, cfg_path): + cfg = CLIConfig(oauth=self._tokens(9999999999)) + assert cfg.resolved_api_key() == "hch-at-x" + + def test_access_valid_expiry_and_skew(self): + assert self._tokens(time.time() + 3600).access_valid() + assert not self._tokens(time.time() - 10).access_valid() + # inside the default 60s skew window → treated as invalid + assert not self._tokens(time.time() + 30).access_valid() + + def test_access_valid_false_without_token(self): + """A missing token is invalid even with a far-future expiry.""" + tokens = OAuthTokens(access_token="", access_expires_at=time.time() + 3600) + assert tokens.access_valid() is False + + def test_from_response_keeps_prior_refresh_token_when_not_rotated(self): + """Refresh-token rotation is optional (RFC 6749 §5.1) — keep the old one.""" + resp = TokenResponse( + access_token="new-at", refresh_token="", expires_in=3600, scope="" + ) + tokens = OAuthTokens.from_response( + resp, + client_id="honcho-cli", + scope_fallback="write", + refresh_fallback="prior-rt", + host="https://staging.example.com", + ) + assert tokens.refresh_token == "prior-rt" + assert tokens.scope == "write" + assert tokens.host == "https://staging.example.com" + + def test_host_round_trips_and_legacy_matches_all(self, cfg_path): + tokens = self._tokens(9999999999) + tokens.host = "https://staging.example.com" + CLIConfig(base_url="https://staging.example.com", oauth=tokens).save() + loaded = CLIConfig.load() + assert loaded.oauth is not None + assert loaded.oauth.host == "https://staging.example.com" + # trailing-slash normalization + legacy blocks (no host) trust any host + assert loaded.oauth.matches_host("https://staging.example.com/") + assert not loaded.oauth.matches_host("https://api.honcho.dev") + assert OAuthTokens(access_token="x").matches_host("https://anything.dev") + + def test_redacted_masks_oauth_token(self): + red = CLIConfig(oauth=self._tokens(1234)).redacted() + assert red["oauth"] == "***at-x" + + def test_save_sets_600_permissions(cfg_path): """Config with plaintext API key must be owner-readable only on POSIX.""" import stat diff --git a/honcho-cli/tests/test_oauth.py b/honcho-cli/tests/test_oauth.py new file mode 100644 index 00000000..45121130 --- /dev/null +++ b/honcho-cli/tests/test_oauth.py @@ -0,0 +1,232 @@ +"""Tests for the device-authorization OAuth engine (transport-only).""" + +from __future__ import annotations + +from unittest.mock import patch + +import httpx +import pytest +from honcho_cli import oauth +from honcho_cli.oauth import ( + AccessDenied, + AuthorizationTimeout, + DeviceCode, + DeviceCodeExpired, + Endpoints, + OAuthFlowError, +) + + +class FakeResponse: + def __init__(self, status_code: int, body): + self.status_code = status_code + self._body = body + self.text = str(body) + + def json(self): + if isinstance(self._body, Exception): + raise self._body + return self._body + + +def _endpoints() -> Endpoints: + return Endpoints( + device_auth_url="https://api.honcho.dev/oauth/device_authorization", + token_url="https://api.honcho.dev/oauth/token", + client_id="honcho-cli", + scope="write", + ) + + +DEVICE = DeviceCode( + device_code="dev-abc", + user_code="WXYZ-1234", + verification_uri="https://app.honcho.dev/device", + verification_uri_complete="https://app.honcho.dev/device?user_code=WXYZ-1234", + expires_in=600, + interval=5, +) + + +# --------------------------------------------------------------------------- # +# resolve_endpoints + +class TestResolveEndpoints: + def test_derives_urls_from_base_url(self): + ep = oauth.resolve_endpoints("https://api.honcho.dev") + assert ep.device_auth_url == "https://api.honcho.dev/oauth/device_authorization" + assert ep.token_url == "https://api.honcho.dev/oauth/token" + assert ep.client_id == "honcho-cli" + assert ep.scope == "write" + + def test_strips_trailing_slash(self): + ep = oauth.resolve_endpoints("http://localhost:8000/") + assert ep.token_url == "http://localhost:8000/oauth/token" + + +# --------------------------------------------------------------------------- # +# supports_device_login + +class TestSupportsDeviceLogin: + def test_true_when_device_grant_advertised(self): + body = {"grant_types_supported": ["authorization_code", oauth.DEVICE_GRANT_TYPE]} + with patch("honcho_cli.oauth.httpx.get", return_value=FakeResponse(200, body)): + assert oauth.supports_device_login("https://api.honcho.dev") is True + + def test_false_when_device_grant_absent(self): + body = {"grant_types_supported": ["authorization_code", "refresh_token"]} + with patch("honcho_cli.oauth.httpx.get", return_value=FakeResponse(200, body)): + assert oauth.supports_device_login("https://api.honcho.dev") is False + + @pytest.mark.parametrize("status", [404, 500]) + def test_false_on_non_200(self, status): + with patch("honcho_cli.oauth.httpx.get", return_value=FakeResponse(status, "")): + assert oauth.supports_device_login("http://localhost:8000") is False + + def test_false_on_connection_error(self): + with patch("honcho_cli.oauth.httpx.get", side_effect=httpx.ConnectError("no route")): + assert oauth.supports_device_login("http://localhost:8000") is False + + def test_false_on_unparseable_body(self): + with patch("honcho_cli.oauth.httpx.get", return_value=FakeResponse(200, ValueError())): + assert oauth.supports_device_login("https://api.honcho.dev") is False + + +# --------------------------------------------------------------------------- # +# request_device_code + +class TestRequestDeviceCode: + def test_success(self): + body = { + "device_code": "dev-abc", + "user_code": "WXYZ-1234", + "verification_uri": "https://app.honcho.dev/device", + "verification_uri_complete": "https://app.honcho.dev/device?user_code=WXYZ-1234", + "expires_in": 600, + "interval": 5, + } + with patch("honcho_cli.oauth.httpx.post", return_value=FakeResponse(200, body)) as post: + dc = oauth.request_device_code(_endpoints()) + assert dc.device_code == "dev-abc" + assert dc.user_code == "WXYZ-1234" + assert dc.interval == 5 + assert post.call_args.kwargs["data"]["source"] == "honcho-cli" + + def test_error_raises(self): + body = {"error": "invalid_client", "error_description": "unknown client"} + with patch("honcho_cli.oauth.httpx.post", return_value=FakeResponse(401, body)): + with pytest.raises(OAuthFlowError) as exc: + oauth.request_device_code(_endpoints()) + assert exc.value.error == "invalid_client" + + def test_transport_failure_wrapped(self): + with patch("honcho_cli.oauth.httpx.post", side_effect=httpx.ConnectError("no route")): + with pytest.raises(OAuthFlowError) as exc: + oauth.request_device_code(_endpoints()) + assert exc.value.error == "connection_error" + + +# --------------------------------------------------------------------------- # +# poll_for_token + +class TestPollForToken: + def _run(self, responses, monotonic_vals=None): + """Poll with a scripted response sequence, capturing sleep durations.""" + sleeps: list[float] = [] + clock = iter(monotonic_vals or [0.0] * (len(responses) + 2)) + with patch("honcho_cli.oauth.httpx.post", side_effect=responses): + token = oauth.poll_for_token( + _endpoints(), + DEVICE, + sleep=sleeps.append, + monotonic=lambda: next(clock), + ) + return token, sleeps + + def test_pending_then_slowdown_then_success(self): + success = { + "access_token": "hch-at-1", + "refresh_token": "hch-rt-1", + "expires_in": 3600, + "scope": "write", + "config": {"k": "v"}, + } + responses = [ + FakeResponse(400, {"error": "authorization_pending"}), + FakeResponse(400, {"error": "slow_down"}), + FakeResponse(200, success), + ] + token, sleeps = self._run(responses) + assert token.access_token == "hch-at-1" + assert token.refresh_token == "hch-rt-1" + assert token.config == {"k": "v"} + # interval starts at 5, bumps by 5 after slow_down → third sleep is 10 + assert sleeps == [5, 5, 10] + + def test_access_denied(self): + responses = [FakeResponse(400, {"error": "access_denied"})] + with pytest.raises(AccessDenied): + self._run(responses) + + def test_expired_token(self): + responses = [FakeResponse(400, {"error": "expired_token"})] + with pytest.raises(DeviceCodeExpired): + self._run(responses) + + def test_unexpected_error_raises_generic(self): + responses = [FakeResponse(400, {"error": "invalid_grant"})] + with pytest.raises(OAuthFlowError) as exc: + self._run(responses) + assert exc.value.error == "invalid_grant" + + def test_transport_failure_wrapped(self): + # an exception in the side_effect list is raised on that poll + responses = [httpx.ReadTimeout("timed out")] + with pytest.raises(OAuthFlowError) as exc: + self._run(responses) + assert exc.value.error == "connection_error" + + def test_times_out_past_deadline(self): + # monotonic jumps past deadline (0 + expires_in) on the first check + with patch("honcho_cli.oauth.httpx.post") as post: + with pytest.raises(AuthorizationTimeout): + oauth.poll_for_token( + _endpoints(), + DEVICE, + sleep=lambda _s: None, + monotonic=iter([0.0, 9999.0]).__next__, + ) + post.assert_not_called() + + +# --------------------------------------------------------------------------- # +# refresh_access_token + +class TestRefresh: + def test_success_returns_rotated_pair(self): + body = { + "access_token": "hch-at-2", + "refresh_token": "hch-rt-2", + "expires_in": 3600, + "scope": "write", + } + with patch("honcho_cli.oauth.httpx.post", return_value=FakeResponse(200, body)) as post: + token = oauth.refresh_access_token(_endpoints(), "hch-rt-1") + assert token.access_token == "hch-at-2" + assert token.refresh_token == "hch-rt-2" + sent = post.call_args.kwargs["data"] + assert sent["grant_type"] == "refresh_token" + assert sent["refresh_token"] == "hch-rt-1" + + def test_error_raises(self): + body = {"error": "invalid_grant", "error_description": "revoked"} + with patch("honcho_cli.oauth.httpx.post", return_value=FakeResponse(400, body)): + with pytest.raises(OAuthFlowError) as exc: + oauth.refresh_access_token(_endpoints(), "stale") + assert exc.value.error == "invalid_grant" + + def test_transport_failure_wrapped(self): + with patch("honcho_cli.oauth.httpx.post", side_effect=httpx.ConnectError("no route")): + with pytest.raises(OAuthFlowError) as exc: + oauth.refresh_access_token(_endpoints(), "hch-rt-1") + assert exc.value.error == "connection_error" diff --git a/honcho-cli/uv.lock b/honcho-cli/uv.lock index dcadb6f3..fce3f13a 100644 --- a/honcho-cli/uv.lock +++ b/honcho-cli/uv.lock @@ -88,9 +88,10 @@ wheels = [ [[package]] name = "honcho-cli" -version = "0.1.0" +version = "0.1.1" source = { editable = "." } dependencies = [ + { name = "click" }, { name = "honcho-ai" }, { name = "httpx" }, { name = "rich" }, @@ -105,6 +106,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "click", specifier = ">=8.0.0" }, { name = "honcho-ai", specifier = ">=0.1.0" }, { name = "httpx", specifier = ">=0.27.0" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0.0" }, diff --git a/mcp/README.md b/mcp/README.md index ba90710b..7ec5963d 100644 --- a/mcp/README.md +++ b/mcp/README.md @@ -16,13 +16,10 @@ A Cloudflare Worker that implements the [Model Context Protocol (MCP)](https://m "mcp-remote", "https://mcp.honcho.dev", "--header", - "Authorization:${AUTH_HEADER}", - "--header", - "X-Honcho-User-Name:${USER_NAME}" + "Authorization:${AUTH_HEADER}" ], "env": { - "AUTH_HEADER": "Bearer ", - "USER_NAME": "" + "AUTH_HEADER": "Bearer " } } } @@ -115,8 +112,7 @@ bun run tsc --noEmit ```bash bunx mcp-remote http://localhost:8787 \ - --header "Authorization:Bearer " \ - --header "X-Honcho-User-Name:test" + --header "Authorization:Bearer " ``` ### Deploy diff --git a/mcp/src/config.ts b/mcp/src/config.ts index 58e4307e..98cf8760 100644 --- a/mcp/src/config.ts +++ b/mcp/src/config.ts @@ -2,8 +2,6 @@ import { Honcho } from "@honcho-ai/sdk"; export interface HonchoConfig { apiKey: string; - userName: string; - assistantName: string; baseUrl: string; workspaceId: string; } @@ -14,7 +12,7 @@ export interface Env { /** * Parse configuration from request headers and Worker env bindings. - * Throws on missing required fields so callers get clear errors. + * Throws only when the Authorization bearer token is missing/empty. * * The Honcho API URL is read from the `HONCHO_API_URL` env var when set, * allowing operators to run this Worker alongside a self-hosted Honcho @@ -24,29 +22,19 @@ export interface Env { */ export function parseConfig(request: Request, env: Env = {}): HonchoConfig { const authHeader = request.headers.get("Authorization"); - const trimmedAuthHeader = authHeader?.trim(); - if (!trimmedAuthHeader?.startsWith("Bearer ")) { + const bearerMatch = authHeader?.trim().match(/^Bearer\s+(.*)$/i); + if (!bearerMatch) { throw new Error( "Missing Authorization header. Provide 'Authorization: Bearer '.", ); } - const apiKey = trimmedAuthHeader.substring(7).trim(); + const apiKey = bearerMatch[1].trim(); if (!apiKey) { throw new Error("Authorization header is empty after 'Bearer '."); } - const rawUserName = request.headers.get("X-Honcho-User-Name"); - const userName = rawUserName?.trim(); - if (!userName) { - throw new Error( - "Missing X-Honcho-User-Name header. Provide 'X-Honcho-User-Name: '.", - ); - } - return { apiKey, - userName, - assistantName: request.headers.get("X-Honcho-Assistant-Name")?.trim() || "Assistant", baseUrl: env.HONCHO_API_URL?.trim() || "https://api.honcho.dev", workspaceId: request.headers.get("X-Honcho-Workspace-ID")?.trim() || "default", }; diff --git a/mcp/src/index.ts b/mcp/src/index.ts index 4e895198..a7e0f107 100644 --- a/mcp/src/index.ts +++ b/mcp/src/index.ts @@ -5,14 +5,25 @@ import { createServer } from "./server.js"; const CORS_ORIGIN = "*"; const CORS_METHODS = "GET, POST, DELETE, OPTIONS"; const CORS_ALLOWED_HEADERS = - "Content-Type, Authorization, X-Honcho-User-Name, X-Honcho-Workspace-ID, X-Honcho-Assistant-Name"; + "Content-Type, Authorization, X-Honcho-Workspace-ID"; const CORS_HEADERS = { "Access-Control-Allow-Origin": CORS_ORIGIN, "Access-Control-Allow-Methods": CORS_METHODS, "Access-Control-Allow-Headers": CORS_ALLOWED_HEADERS, + "Access-Control-Expose-Headers": "WWW-Authenticate", }; +const PROTECTED_RESOURCE_PATH = "/.well-known/oauth-protected-resource"; + +function resourceUrl(request: Request): string { + return new URL(request.url).origin; +} + +function authorizationServer(env: Env): string { + return env.HONCHO_API_URL?.trim() || "https://api.honcho.dev"; +} + export default { async fetch( request: Request, @@ -23,15 +34,34 @@ export default { return new Response(null, { status: 204, headers: CORS_HEADERS }); } + // Protected Resource Metadata (RFC 9728) — served without auth so clients + // can discover the authorization server. + if (new URL(request.url).pathname === PROTECTED_RESOURCE_PATH) { + return Response.json( + { + resource: resourceUrl(request), + authorization_servers: [authorizationServer(env)], + bearer_methods_supported: ["header"], + }, + { headers: CORS_HEADERS }, + ); + } + let config; try { config = parseConfig(request, env); } catch (e) { const message = e instanceof Error ? e.message : "Invalid request"; + // WWW-Authenticate points clients at the metadata so they start the OAuth flow. + const resourceMetadata = `${resourceUrl(request)}${PROTECTED_RESOURCE_PATH}`; return new Response(JSON.stringify({ error: message }), { status: 401, - headers: { "Content-Type": "application/json", ...CORS_HEADERS }, + headers: { + "Content-Type": "application/json", + "WWW-Authenticate": `Bearer resource_metadata="${resourceMetadata}"`, + ...CORS_HEADERS, + }, }); } diff --git a/pyproject.toml b/pyproject.toml index 793d363b..9427cd88 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "honcho" -version = "3.0.7" +version = "3.0.11" description = "Honcho Server" authors = [ {name = "Plastic Labs", email = "hello@plasticlabs.ai"}, @@ -32,11 +32,11 @@ dependencies = [ "typing-extensions>=4.11.0", "json-repair>=0.49.0", "turbopuffer>=1.8.1", - "lancedb>=0.25.3", "qdrant-client>=1.18.0", + "lancedb>=0.25.3; sys_platform != \"darwin\" or platform_machine != \"x86_64\"", "pyarrow>=19.0.0", "redis>=7.0.0,<8.0.0", - "cashews[redis]==7.4.4", + "cashews[redis]==7.5.0", "scikit-learn>=1.6.0", "prometheus_client>=0.21.0", "cloudevents>=1.12.0,<2.0", @@ -60,6 +60,9 @@ dev = [ "pytest-xdist>=3.8.0", ] +[tool.uv] +exclude-newer = "5 days" + [tool.uv.workspace] members = [ "sdks/python", diff --git a/scripts/dialectic_cost_calculator.py b/scripts/dialectic_cost_calculator.py index 0c014628..af0e3275 100644 --- a/scripts/dialectic_cost_calculator.py +++ b/scripts/dialectic_cost_calculator.py @@ -170,10 +170,10 @@ def calculate_level_cost( realistic_final_answer=realistic_final, ) - model = level_config.MODEL + model = level_config.MODEL_CONFIG.model max_iterations = level_config.MAX_TOOL_ITERATIONS - thinking_budget = level_config.THINKING_BUDGET_TOKENS - provider = level_config.PROVIDER + thinking_budget = level_config.MODEL_CONFIG.thinking_budget_tokens or 0 + provider = level_config.MODEL_CONFIG.transport # Get pricing for this model pricing = MODEL_PRICING.get(model, {"input": 0, "output": 0, "cached": 0}) diff --git a/scripts/generate_jwt.py b/scripts/generate_jwt.py new file mode 100644 index 00000000..8d96cd67 --- /dev/null +++ b/scripts/generate_jwt.py @@ -0,0 +1,148 @@ +#!/usr/bin/env uv run python +""" +Utility script to generate scoped JWTs for Honcho. + +Examples: + # Admin JWT (no expiry) + uv run python scripts/generate_jwt.py --admin + + # Admin JWT expiring in 24 hours + uv run python scripts/generate_jwt.py --admin --expires 24h + + # Workspace-scoped JWT expiring in 30 days + uv run python scripts/generate_jwt.py --workspace my-workspace --expires 30d + + # Peer-scoped JWT expiring in 1 year + uv run python scripts/generate_jwt.py --workspace my-workspace --peer my-peer --expires 1y + + # Session-scoped JWT + uv run python scripts/generate_jwt.py --workspace my-workspace --session my-session --expires 8h +""" + +import argparse +import datetime +import os +import re +import sys + +# Allow running from repo root without installing +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from src.security import JWTParams, create_jwt +from src.utils.formatting import format_datetime_utc + +DURATION_UNITS = { + "s": datetime.timedelta(seconds=1), + "m": datetime.timedelta(minutes=1), + "h": datetime.timedelta(hours=1), + "d": datetime.timedelta(days=1), + "w": datetime.timedelta(weeks=1), + "y": datetime.timedelta(days=365), +} + + +def parse_duration(value: str) -> datetime.timedelta: + """Parse a duration string like '5h', '1d', '2w', '1y' into a timedelta.""" + match = re.fullmatch(r"(\d+)([smhdwy])", value.strip().lower()) + if not match: + raise argparse.ArgumentTypeError( + f"Invalid duration '{value}'. Use format like: 30s, 5m, 2h, 7d, 2w, 1y" + ) + amount, unit = int(match.group(1)), match.group(2) + return DURATION_UNITS[unit] * amount + + +def main(): + parser = argparse.ArgumentParser( + description="Generate a scoped JWT for Honcho authentication.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + parser.add_argument( + "--admin", + action="store_true", + help="Generate an admin JWT (full access)", + ) + parser.add_argument( + "--workspace", + "-w", + metavar="NAME", + help="Scope the JWT to a workspace", + ) + parser.add_argument( + "--peer", + "-p", + metavar="NAME", + help="Scope the JWT to a peer (requires --workspace)", + ) + parser.add_argument( + "--session", + "-s", + metavar="NAME", + help="Scope the JWT to a session (requires --workspace)", + ) + parser.add_argument( + "--expires", + "-e", + metavar="DURATION", + type=parse_duration, + help="Token expiry duration. Units: s=seconds, m=minutes, h=hours, d=days, w=weeks, y=years. E.g. 5h, 30d, 1y", + ) + parser.add_argument( + "--print-only", + action="store_true", + help="Only print the token, no labels", + ) + args = parser.parse_args() + + if not args.admin and not any([args.workspace, args.peer, args.session]): + parser.error( + "Specify --admin or at least one of --workspace, --peer, --session" + ) + + if args.admin and any([args.workspace, args.peer, args.session]): + parser.error( + "--admin cannot be combined with --workspace, --peer, or --session" + ) + + if (args.peer or args.session) and not args.workspace: + parser.error("--peer and --session require --workspace") + + exp_str: str | None = None + if args.expires: + expiry = datetime.datetime.now(datetime.timezone.utc) + args.expires + exp_str = format_datetime_utc(expiry) + + params = JWTParams( + ad=True if args.admin else None, + w=args.workspace, + p=args.peer, + s=args.session, + exp=exp_str, + ) + + token = create_jwt(params) + + if args.print_only: + print(token) + else: + scope_parts: list[str] = [] + if args.admin: + scope_parts.append("admin") + if args.workspace: + scope_parts.append(f"workspace={args.workspace}") + if args.peer: + scope_parts.append(f"peer={args.peer}") + if args.session: + scope_parts.append(f"session={args.session}") + + print(f"Scope: {', '.join(scope_parts)}") + if exp_str: + print(f"Expires: {exp_str}") + else: + print("Expires: never") + print(f"Token: {token}") + + +if __name__ == "__main__": + main() diff --git a/scripts/test_reasoning_levels.py b/scripts/test_reasoning_levels.py index 699b5f18..3fd38752 100755 --- a/scripts/test_reasoning_levels.py +++ b/scripts/test_reasoning_levels.py @@ -6,6 +6,7 @@ import json import os import time from datetime import datetime, timedelta, timezone +from typing import Any import httpx from dotenv import load_dotenv @@ -106,7 +107,7 @@ def load_locomo( print(f" Created session: {session_id}") # Build message batch - msg_batch = [] + msg_batch: list[dict[str, Any]] = [] for i, msg in enumerate(messages): msg_time = base_time + timedelta(seconds=i * 2) msg_batch.append( @@ -134,7 +135,7 @@ def load_locomo( def chat( client: httpx.Client, workspace_id: str, peer_id: str, query: str, level: str -) -> dict: +) -> dict[str, Any]: """Call the chat endpoint with a specific reasoning level.""" resp = client.post( f"{BASE_URL}/workspaces/{workspace_id}/peers/{peer_id}/chat", diff --git a/scripts/update_version.py b/scripts/update_version.py index 4d99020b..40416567 100755 --- a/scripts/update_version.py +++ b/scripts/update_version.py @@ -6,6 +6,7 @@ This script helps update version numbers across the Honcho repository. It handles the main API, Python SDK, and TypeScript SDK in a single operation. """ +import argparse import json import os import re @@ -17,11 +18,11 @@ from datetime import datetime class VersionUpdater: def __init__(self, base_path: str): - self.base_path = base_path + self.base_path: str = base_path def get_current_versions(self) -> dict[str, str]: """Get current version numbers from the repository.""" - versions = {} + versions: dict[str, str] = {} # Main API version with open(os.path.join(self.base_path, "pyproject.toml")) as f: @@ -89,7 +90,7 @@ TYPESCRIPT_VERSION= os.unlink(temp_file) # Extract all versions and changelogs - updates = {} + updates: dict[str, dict[str, str]] = {} # Parse API version api_match = re.search(r"^API_VERSION=(.*)$", content, re.MULTILINE) @@ -131,7 +132,7 @@ TYPESCRIPT_VERSION= ) -> str: """Extract changelog content between markers.""" lines = content.split("\n") - changelog_lines = [] + changelog_lines: list[str] = [] in_section = False for line in lines: @@ -158,8 +159,8 @@ TYPESCRIPT_VERSION= """Remove empty changelog sections.""" sections = ["Added", "Changed", "Fixed", "Deprecated", "Removed", "Security"] lines = changelog.split("\n") - cleaned_lines = [] - current_section = None + cleaned_lines: list[str] = [] + current_section: str | None = None section_has_content = False section_start_idx = -1 for i, line in enumerate(lines): @@ -346,28 +347,30 @@ TYPESCRIPT_VERSION= self._update_compatibility_guide("typescript", new_version) def _update_docs_json(self, new_version: str): - """Update docs.json - only update versions with same major version.""" + """Update docs.json version label(s) sharing the new version's major. + + Uses a targeted regex replacement rather than a JSON round-trip so the + file's existing formatting (compact inline arrays) is preserved instead + of being reflowed. + """ file_path = os.path.join(self.base_path, "docs/docs.json") with open(file_path) as f: - data = json.load(f) + content = f.read() # Get major version of new version new_major = new_version.split(".")[0] - # Update only matching major versions - if "navigation" in data and "versions" in data["navigation"]: - for version_entry in data["navigation"]["versions"]: - if "version" in version_entry: - current_version = version_entry["version"].lstrip("v") - current_major = current_version.split(".")[0] + def _replace(match: re.Match[str]) -> str: + # Only update labels whose major version matches the new version's. + if match.group(1) == new_major: + return f'"version": "v{new_version}"' + return match.group(0) - if current_major == new_major: - version_entry["version"] = f"v{new_version}" + content = re.sub(r'"version": "v(\d+)\.\d+\.\d+"', _replace, content) with open(file_path, "w") as f: - json.dump(data, f, indent=2) - f.write("\n") + f.write(content) def _update_sdk_changelog(self, version: str, changelog: str, relative_path: str): """Update an SDK's CHANGELOG.md file.""" @@ -405,12 +408,39 @@ TYPESCRIPT_VERSION= f.write(new_content) def _update_changelog_md(self, version: str, changelog: str): - """Update the main CHANGELOG.md file.""" + """Update the main CHANGELOG.md file. + + If an ``## [Unreleased]`` section is present, it is promoted to the new + version (its contents replaced by ``changelog``, which the caller is + expected to have already merged). Otherwise a new version entry is + prepended above the most recent release, preserving the legacy behavior. + """ file_path = os.path.join(self.base_path, "CHANGELOG.md") with open(file_path) as f: content = f.read() + date = datetime.now().strftime("%Y-%m-%d") + + # Ensure changelog content is properly formatted + if changelog.strip(): + formatted_changelog = changelog.strip() + else: + formatted_changelog = "### Changed\n\n- Updated version" + + # Promote an existing [Unreleased] section if one exists. Match from the + # "## [Unreleased]" header up to (but not including) the next release + # heading, and replace the whole block with the new version section. + unreleased_re = re.compile( + r"\n## \[Unreleased\][\s\S]*?(?=\n## \[)", re.IGNORECASE + ) + if unreleased_re.search(content): + replacement = f"\n## [{version}] - {date}\n\n{formatted_changelog}\n" + new_content = unreleased_re.sub(replacement, content, count=1) + with open(file_path, "w") as f: + f.write(new_content) + return + # Find the position after the header header_end = content.find("\n## [") if header_end == -1: @@ -420,15 +450,6 @@ TYPESCRIPT_VERSION= # No existing entries, add after title header_end = content.find("\n", content.find("# Changelog")) - # Create new entry with proper formatting - date = datetime.now().strftime("%Y-%m-%d") - - # Ensure changelog content is properly formatted - if changelog.strip(): - formatted_changelog = changelog.strip() - else: - formatted_changelog = "### Changed\n\n- Updated version" - new_entry = f"\n\n## [{version}] - {date}\n\n{formatted_changelog}\n" # Insert the new entry @@ -624,7 +645,68 @@ TYPESCRIPT_VERSION= f.write(content) +def _resolve_changelog(value: str | None) -> str: + """Resolve a changelog argument that is either inline text or a file path.""" + if not value: + return "" + if os.path.isfile(value): + with open(value) as f: + return f.read().strip() + return value.strip() + + +def _updates_from_args(args: argparse.Namespace) -> dict[str, dict[str, str]]: + """Build the updates dict from CLI flags (headless mode).""" + updates: dict[str, dict[str, str]] = {} + if args.api_version: + updates["api"] = { + "version": args.api_version, + "changelog": _resolve_changelog(args.api_changelog), + } + if args.python_version: + updates["python_sdk"] = { + "version": args.python_version, + "changelog": _resolve_changelog(args.python_changelog), + } + if args.typescript_version: + updates["typescript_sdk"] = { + "version": args.typescript_version, + "changelog": _resolve_changelog(args.typescript_changelog), + } + return updates + + def main(): + parser = argparse.ArgumentParser( + description=( + "Update Honcho version numbers and changelogs. With no version " + "flags, opens an interactive editor; pass one or more --*-version " + "flags to run headless (agent-friendly)." + ) + ) + parser.add_argument("--api-version", help="New Main API version.") + parser.add_argument("--python-version", help="New Python SDK version.") + parser.add_argument("--typescript-version", help="New TypeScript SDK version.") + parser.add_argument( + "--api-changelog", + help="API changelog markdown, or a path to a file containing it.", + ) + parser.add_argument( + "--python-changelog", + help="Python SDK changelog markdown, or a path to a file containing it.", + ) + parser.add_argument( + "--typescript-changelog", + help="TypeScript SDK changelog markdown, or a path to a file containing it.", + ) + parser.add_argument( + "-y", + "--yes", + action="store_true", + help="Skip the confirmation prompt (implied in headless mode).", + ) + args = parser.parse_args() + # Get the parent directory of the scripts folder (the project root) script_dir = os.path.dirname(os.path.abspath(__file__)) base_path = os.path.dirname(script_dir) @@ -633,6 +715,8 @@ def main(): # Get current versions current_versions = updater.get_current_versions() + headless = any([args.api_version, args.python_version, args.typescript_version]) + print("Honcho Version Updater") print("=" * 50) print("\nCurrent versions:") @@ -640,12 +724,15 @@ def main(): print(f" Python SDK: {current_versions['python_sdk']}") print(f" TypeScript SDK: {current_versions['typescript_sdk']}") print() - print("Opening editor for version updates...") - print("Leave version fields blank to skip updating that component.") - print() - # Get all updates at once - updates = updater.get_all_versions_from_editor(current_versions) + if headless: + updates = _updates_from_args(args) + else: + print("Opening editor for version updates...") + print("Leave version fields blank to skip updating that component.") + print() + # Get all updates at once + updates = updater.get_all_versions_from_editor(current_versions) if not updates: print("No versions specified. Exiting...") @@ -661,11 +748,12 @@ def main(): }[component] print(f" {component_name}: {current_versions[component]} → {info['version']}") - # Confirm - response = input("\nProceed with updates? (y/n): ").strip().lower() - if response != "y": - print("Cancelled.") - sys.exit(0) + # Confirm (skipped in headless mode or with --yes) + if not headless and not args.yes: + response = input("\nProceed with updates? (y/n): ").strip().lower() + if response != "y": + print("Cancelled.") + sys.exit(0) # Apply all updates updater.update_all(updates, current_versions) @@ -673,6 +761,7 @@ def main(): print("\nVersion updates complete!") print("\nDon't forget to:") print(" - Review the changes with `git diff`") + print(" - Run `uv lock` to refresh the lockfile") print(" - Commit the changes") print(" - Create git tags for the new versions") print(" - Push the changes and tags") diff --git a/sdks/python/CHANGELOG.md b/sdks/python/CHANGELOG.md index 2c8d3f91..84843c55 100644 --- a/sdks/python/CHANGELOG.md +++ b/sdks/python/CHANGELOG.md @@ -5,6 +5,17 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](http://keepachangelog.com/) and this project adheres to [Semantic Versioning](http://semver.org/). +## [2.2.0] - 2026-07-02 + +### Added + +- `ConclusionLevel` type (`explicit`, `deductive`, `inductive`, `contradiction`) and a `level` field on `Conclusion`, exposing the reasoning level the server already tracked but previously stripped from responses. +- `filters` parameter on `ConclusionScope.list()` and `ConclusionScope.query()` (sync and async), passed through to the same dynamic server-side filter logic as `peers()`/`sessions()`/`messages()`. Filter explicit-only conclusions with `filters={"level": "explicit"}`, or by any other supported field/operator. Requires a Honcho server with the matching API support (Honcho v3.0.11+). + +### Fixed + +- Scope-managed filter keys (`observer`, `observed`, `session`) are now rejected with a clear `ValueError` if passed in `filters`, instead of silently overriding the scope and returning conclusions from a different peer pair. Use `peer.conclusions` / `conclusions_of(target)` and the `session=` parameter instead. `session_id` remains a valid filter on `query()`. + ## [2.1.2] - 2026-05-21 ### Added diff --git a/sdks/python/pyproject.toml b/sdks/python/pyproject.toml index 0e6ad9ba..6fe801ff 100644 --- a/sdks/python/pyproject.toml +++ b/sdks/python/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "honcho-ai" -version = "2.1.2" +version = "2.2.0" description = "Official DX Optimized Python SDK for Honcho" dynamic = ["readme"] license = "Apache-2.0" diff --git a/sdks/python/src/honcho/aio.py b/sdks/python/src/honcho/aio.py index 9ff27a12..51797c72 100644 --- a/sdks/python/src/honcho/aio.py +++ b/sdks/python/src/honcho/aio.py @@ -26,9 +26,9 @@ import logging import warnings from collections.abc import AsyncGenerator from datetime import datetime -from typing import TYPE_CHECKING, Any, ClassVar, Literal +from typing import TYPE_CHECKING, Any, ClassVar, Literal, overload -from pydantic import ConfigDict, Field, validate_call +from pydantic import BaseModel, ConfigDict, Field, validate_call from .api_types import ( ConclusionResponse, @@ -47,7 +47,11 @@ from .api_types import ( WorkspaceResponse, ) from .base import PeerBase, SessionBase -from .conclusions import Conclusion +from .conclusions import ( + _SCOPE_RESERVED, + Conclusion, + _reject_reserved_filter_keys, +) from .http import routes from .message import Message from .mixins import AsyncMetadataConfigMixin @@ -67,7 +71,7 @@ if TYPE_CHECKING: from .conclusions import ConclusionScope from .conclusions import ConclusionCreateParams -from .peer import Peer +from .peer import Peer, TResponseFormat, serialize_response_format from .session import Session logger = logging.getLogger(__name__) @@ -573,6 +577,30 @@ class PeerAio(AsyncMetadataConfigMixin): ) self._peer._configuration = configuration + @overload + async def chat( + self, + query: str, + *, + target: str | PeerBase | None = None, + session: str | SessionBase | None = None, + reasoning_level: Literal["minimal", "low", "medium", "high", "max"] + | None = None, + response_format: type[TResponseFormat], + ) -> TResponseFormat | None: ... + + @overload + async def chat( + self, + query: str, + *, + target: str | PeerBase | None = None, + session: str | SessionBase | None = None, + reasoning_level: Literal["minimal", "low", "medium", "high", "max"] + | None = None, + response_format: dict[str, Any] | None = None, + ) -> str | None: ... + @validate_call(config=ConfigDict(arbitrary_types_allowed=True)) async def chat( self, @@ -582,8 +610,14 @@ class PeerAio(AsyncMetadataConfigMixin): session: str | SessionBase | None = None, reasoning_level: Literal["minimal", "low", "medium", "high", "max"] | None = None, - ) -> str | None: - """Query the peer's representation asynchronously.""" + response_format: type[BaseModel] | dict[str, Any] | None = None, + ) -> BaseModel | str | None: + """Query the peer's representation asynchronously. + + See Peer.chat for parameter details. When response_format is a Pydantic + model class, the answer is parsed into an instance of it; when it is a + JSON Schema dict, the answer is a JSON string. + """ await self._peer._honcho._ensure_workspace_async() target_id = resolve_id(target) resolved_session_id = resolve_id(session) @@ -595,6 +629,9 @@ class PeerAio(AsyncMetadataConfigMixin): body["session_id"] = resolved_session_id if reasoning_level: body["reasoning_level"] = reasoning_level + response_format_schema = serialize_response_format(response_format) + if response_format_schema is not None: + body["response_format"] = response_format_schema data = await self._peer._honcho._async_http_client.post( routes.peer_chat(self._peer.workspace_id, self._peer.id), @@ -603,6 +640,8 @@ class PeerAio(AsyncMetadataConfigMixin): content = data.get("content") if not content: return None + if isinstance(response_format, type): + return response_format.model_validate_json(content) return content @validate_call(config=ConfigDict(arbitrary_types_allowed=True)) @@ -614,8 +653,14 @@ class PeerAio(AsyncMetadataConfigMixin): session: str | SessionBase | None = None, reasoning_level: Literal["minimal", "low", "medium", "high", "max"] | None = None, + response_format: type[BaseModel] | dict[str, Any] | None = None, ) -> AsyncDialecticStreamResponse: - """Query the peer's representation with streaming asynchronously.""" + """Query the peer's representation with streaming asynchronously. + + See Peer.chat_stream for parameter details. With response_format set, + chunks stay raw text that accumulates to a JSON string; parse it after + the stream completes. + """ await self._peer._honcho._ensure_workspace_async() target_id = resolve_id(target) resolved_session_id = resolve_id(session) @@ -627,6 +672,9 @@ class PeerAio(AsyncMetadataConfigMixin): body["session_id"] = resolved_session_id if reasoning_level: body["reasoning_level"] = reasoning_level + response_format_schema = serialize_response_format(response_format) + if response_format_schema is not None: + body["response_format"] = response_format_schema async def stream_response() -> AsyncGenerator[str, None]: async for content in parse_sse_astream( @@ -1460,17 +1508,28 @@ class ConclusionScopeAio: size: int = 50, session: str | SessionBase | None = None, *, + filters: dict[str, Any] | None = None, reverse: bool = False, ) -> AsyncPage[ConclusionResponse, Conclusion]: - """List conclusions in this scope asynchronously.""" + """List conclusions in this scope asynchronously. + + Pass ``filters`` to add criteria merged with this scope's + observer/observed (and session, if given) — e.g. + ``{"level": "explicit"}`` to get only conclusions extracted directly + from messages (i.e. not derived during dreaming). See + https://honcho.dev/docs/v3/documentation/features/advanced/using-filters + """ + _reject_reserved_filter_keys( + filters, _SCOPE_RESERVED + ("session", "session_id") + ) await self._scope._honcho._ensure_workspace_async() resolved_session_id = resolve_id(session) - filters: dict[str, Any] = { + filters = { "observer_id": self._scope.observer, "observed_id": self._scope.observed, + **({"session_id": resolved_session_id} if resolved_session_id else {}), + **(filters or {}), } - if resolved_session_id: - filters["session_id"] = resolved_session_id query: dict[str, Any] = {"page": page, "size": size} if reverse: @@ -1504,12 +1563,24 @@ class ConclusionScopeAio: query: str, top_k: int = 10, distance: float | None = None, + *, + filters: dict[str, Any] | None = None, ) -> list[Conclusion]: - """Semantic search for conclusions asynchronously.""" + """Semantic search for conclusions asynchronously. + + Args: + query: The search query string + top_k: Maximum number of results to return + distance: Maximum cosine distance threshold (0.0-1.0) + filters: Optional dictionary of additional filter criteria, merged + with this scope's observer/observed (e.g. ``{"level": "deductive"}``). + """ + _reject_reserved_filter_keys(filters, _SCOPE_RESERVED) await self._scope._honcho._ensure_workspace_async() - filters: dict[str, Any] = { + filters = { "observer_id": self._scope.observer, "observed_id": self._scope.observed, + **(filters or {}), } body: dict[str, Any] = { diff --git a/sdks/python/src/honcho/api_types.py b/sdks/python/src/honcho/api_types.py index 897a86fe..f626c5a6 100644 --- a/sdks/python/src/honcho/api_types.py +++ b/sdks/python/src/honcho/api_types.py @@ -10,6 +10,10 @@ from typing import Any, Literal from pydantic import BaseModel, ConfigDict, Field +# Reasoning level of a conclusion. "explicit" conclusions are extracted directly +# from messages; the others are derived during dreaming. +ConclusionLevel = Literal["explicit", "deductive", "inductive", "contradiction"] + # ============================================================================== # Configuration Types # ============================================================================== @@ -414,6 +418,7 @@ class ConclusionResponse(BaseModel): observer_id: str observed_id: str session_id: str | None = None + level: ConclusionLevel = "explicit" created_at: datetime.datetime @@ -498,6 +503,7 @@ class DialecticParams(BaseModel): query: str = Field(min_length=1, max_length=10000) stream: bool = False reasoning_level: ReasoningLevel = "low" + response_format: dict[str, Any] | None = None class DialecticResponse(BaseModel): diff --git a/sdks/python/src/honcho/conclusions.py b/sdks/python/src/honcho/conclusions.py index e76b1900..708cc3ed 100644 --- a/sdks/python/src/honcho/conclusions.py +++ b/sdks/python/src/honcho/conclusions.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Any from pydantic import BaseModel -from .api_types import ConclusionResponse, RepresentationResponse +from .api_types import ConclusionLevel, ConclusionResponse, RepresentationResponse from .base import SessionBase from .http import routes from .pagination import SyncPage @@ -24,6 +24,34 @@ __all__ = [ "ConclusionCreateParams", ] +# Filter keys that define a conclusion scope (the observer/observed peer pair). +# They are set from the scope itself, so a caller must not pass them in `filters`. +_SCOPE_RESERVED = ("observer", "observed", "observer_id", "observed_id") + + +def _reject_reserved_filter_keys( + filters: dict[str, Any] | None, reserved: tuple[str, ...] +) -> None: + """Raise if ``filters`` contains keys managed by the conclusion scope. + + The observer/observed peer pair (and, on ``list``, the session) is fixed by + the scope, so letting a user filter override it would silently return data + from a different scope than requested. Fail loud instead. + """ + if not filters: + return + clash = sorted(k for k in reserved if k in filters) + if clash: + guidance = ( + "Choose the peer pair via peer.conclusions / peer.conclusions_of(target)" + ) + if "session" in reserved or "session_id" in reserved: + guidance += "; use the session= parameter to filter by session" + raise ValueError( + f"Filter key(s) {clash} are managed by this conclusion scope and " + + f"cannot be passed in filters. {guidance}." + ) + class ConclusionCreateParams(BaseModel): content: str @@ -43,6 +71,9 @@ class Conclusion: observer_id: The peer ID who made this conclusion observed_id: The peer ID this conclusion is about session_id: The session this conclusion relates to + level: Reasoning level ("explicit", "deductive", "inductive", + "contradiction"). "explicit" conclusions are extracted directly + from messages; the others are derived during dreaming. created_at: Timestamp for when the conclusion was created """ @@ -51,6 +82,7 @@ class Conclusion: observer_id: str observed_id: str session_id: str | None = None + level: ConclusionLevel = "explicit" created_at: datetime.datetime def __init__( @@ -61,12 +93,14 @@ class Conclusion: observed_id: str, session_id: str | None, created_at: datetime.datetime, + level: ConclusionLevel = "explicit", ) -> None: self.id = id self.content = content self.observer_id = observer_id self.observed_id = observed_id self.session_id = session_id + self.level = level self.created_at = created_at @classmethod @@ -78,6 +112,7 @@ class Conclusion: observer_id=data.observer_id, observed_id=data.observed_id, session_id=data.session_id, + level=data.level, created_at=data.created_at, ) @@ -169,6 +204,7 @@ class ConclusionScope: size: int = 50, session: str | SessionBase | None = None, *, + filters: dict[str, Any] | None = None, reverse: bool = False, ) -> SyncPage[ConclusionResponse, Conclusion]: """ @@ -178,19 +214,28 @@ class ConclusionScope: page: Page number (1-indexed) size: Number of results per page session: Optional session (ID string or Session object) to filter by + filters: Optional dictionary of additional filter criteria, merged + with this scope's observer/observed (and session, if given). + Supports the same operators as other list endpoints — e.g. + ``{"level": "explicit"}`` to get only conclusions extracted + directly from messages (i.e. not derived during dreaming). See + https://honcho.dev/docs/v3/documentation/features/advanced/using-filters reverse: If True, reverses the default ordering. Default: False. Returns: Paginated response containing Conclusion objects """ + _reject_reserved_filter_keys( + filters, _SCOPE_RESERVED + ("session", "session_id") + ) self._honcho._ensure_workspace() resolved_session_id = resolve_id(session) - filters: dict[str, Any] = { + filters = { "observer_id": self.observer, "observed_id": self.observed, + **({"session_id": resolved_session_id} if resolved_session_id else {}), + **(filters or {}), } - if resolved_session_id: - filters["session_id"] = resolved_session_id query: dict[str, Any] = {"page": page, "size": size} if reverse: @@ -224,6 +269,8 @@ class ConclusionScope: query: str, top_k: int = 10, distance: float | None = None, + *, + filters: dict[str, Any] | None = None, ) -> list[Conclusion]: """ Semantic search for conclusions in this scope. @@ -232,14 +279,21 @@ class ConclusionScope: query: The search query string top_k: Maximum number of results to return distance: Maximum cosine distance threshold (0.0-1.0) + filters: Optional dictionary of additional filter criteria, merged + with this scope's observer/observed. Supports the same operators + as the list endpoint — e.g. ``{"level": "deductive"}`` to search + only conclusions derived during dreaming. See + https://honcho.dev/docs/v3/documentation/features/advanced/using-filters Returns: List of matching Conclusion objects """ + _reject_reserved_filter_keys(filters, _SCOPE_RESERVED) self._honcho._ensure_workspace() - filters: dict[str, Any] = { + filters = { "observer_id": self.observer, "observed_id": self.observed, + **(filters or {}), } body: dict[str, Any] = { diff --git a/sdks/python/src/honcho/peer.py b/sdks/python/src/honcho/peer.py index f24357bb..e004db6a 100644 --- a/sdks/python/src/honcho/peer.py +++ b/sdks/python/src/honcho/peer.py @@ -7,9 +7,9 @@ import datetime import logging import warnings from collections.abc import Generator -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload -from pydantic import ConfigDict, Field, PrivateAttr, validate_call +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, validate_call from .api_types import ( MessageCreateParams, @@ -38,6 +38,19 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +TResponseFormat = TypeVar("TResponseFormat", bound=BaseModel) + + +def serialize_response_format( + response_format: type[BaseModel] | dict[str, Any] | None, +) -> dict[str, Any] | None: + """Convert a chat response_format argument to a JSON Schema dict.""" + if response_format is None: + return None + if isinstance(response_format, type): + return response_format.model_json_schema() + return response_format + class Peer(PeerBase, MetadataConfigMixin): """ @@ -221,6 +234,30 @@ class Peer(PeerBase, MetadataConfigMixin): self._configuration = configuration # pyright: ignore[reportIncompatibleVariableOverride] self._created_at = created_at + @overload + def chat( + self, + query: str, + *, + target: str | PeerBase | None = None, + session: str | SessionBase | None = None, + reasoning_level: Literal["minimal", "low", "medium", "high", "max"] + | None = None, + response_format: type[TResponseFormat], + ) -> TResponseFormat | None: ... + + @overload + def chat( + self, + query: str, + *, + target: str | PeerBase | None = None, + session: str | SessionBase | None = None, + reasoning_level: Literal["minimal", "low", "medium", "high", "max"] + | None = None, + response_format: dict[str, Any] | None = None, + ) -> str | None: ... + @validate_call(config=ConfigDict(arbitrary_types_allowed=True)) def chat( self, @@ -230,7 +267,8 @@ class Peer(PeerBase, MetadataConfigMixin): session: str | SessionBase | None = None, reasoning_level: Literal["minimal", "low", "medium", "high", "max"] | None = None, - ) -> str | None: + response_format: type[BaseModel] | dict[str, Any] | None = None, + ) -> BaseModel | str | None: """ Query the peer's representation with a natural language question. @@ -249,9 +287,15 @@ class Peer(PeerBase, MetadataConfigMixin): ID string or a Session object. reasoning_level: Optional reasoning level for the query: "minimal", "low", "medium", "high", or "max". Defaults to "low" if not provided. + response_format: Optional structure for the answer. Pass a Pydantic + model class to get a parsed instance back, or a raw + JSON Schema dict (root type "object") to get the + answer as a JSON string. Returns: - Response string containing the answer, or None if no relevant information + Response string containing the answer (a JSON string when a schema + dict was given), a parsed model instance when a Pydantic model class + was given, or None if no relevant information. """ self._honcho._ensure_workspace() target_id = resolve_id(target) @@ -264,6 +308,9 @@ class Peer(PeerBase, MetadataConfigMixin): body["session_id"] = resolved_session_id if reasoning_level: body["reasoning_level"] = reasoning_level + response_format_schema = serialize_response_format(response_format) + if response_format_schema is not None: + body["response_format"] = response_format_schema data = self._honcho._http.post( routes.peer_chat(self.workspace_id, self.id), @@ -272,6 +319,8 @@ class Peer(PeerBase, MetadataConfigMixin): content = data.get("content") if not content: return None + if isinstance(response_format, type): + return response_format.model_validate_json(content) return content @validate_call(config=ConfigDict(arbitrary_types_allowed=True)) @@ -283,6 +332,7 @@ class Peer(PeerBase, MetadataConfigMixin): session: str | SessionBase | None = None, reasoning_level: Literal["minimal", "low", "medium", "high", "max"] | None = None, + response_format: type[BaseModel] | dict[str, Any] | None = None, ) -> DialecticStreamResponse: """ Query the peer's representation with a natural language question, streaming the response. @@ -302,6 +352,11 @@ class Peer(PeerBase, MetadataConfigMixin): ID string or a Session object. reasoning_level: Optional reasoning level for the query: "minimal", "low", "medium", "high", or "max". Defaults to "low" if not provided. + response_format: Optional structure for the answer: a Pydantic model + class or a JSON Schema dict (root type "object"). + Streamed chunks stay raw text that accumulates to a + JSON string; parse it yourself (e.g. with + Model.model_validate_json) once the stream completes. Returns: DialecticStreamResponse object that can be iterated over and provides final response @@ -317,6 +372,9 @@ class Peer(PeerBase, MetadataConfigMixin): body["session_id"] = resolved_session_id if reasoning_level: body["reasoning_level"] = reasoning_level + response_format_schema = serialize_response_format(response_format) + if response_format_schema is not None: + body["response_format"] = response_format_schema def stream_response() -> Generator[str, None, None]: yield from parse_sse_stream( diff --git a/sdks/typescript/CHANGELOG.md b/sdks/typescript/CHANGELOG.md index bd6dd844..179cca72 100644 --- a/sdks/typescript/CHANGELOG.md +++ b/sdks/typescript/CHANGELOG.md @@ -5,6 +5,17 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](http://keepachangelog.com/) and this project adheres to [Semantic Versioning](http://semver.org/). +## [2.2.0] - 2026-07-02 + +### Added + +- `ConclusionLevel` type (`explicit`, `deductive`, `inductive`, `contradiction`) and a `level` field on `Conclusion`, exposing the reasoning level the server already tracked but previously stripped from responses. +- `filters` option on `conclusions.list()` and `conclusions.query()`, passed through to the same dynamic server-side filter logic as the other list endpoints. Filter explicit-only conclusions with `{ filters: { level: 'explicit' } }`, or by any other supported field/operator. Requires a Honcho server with the matching API support (Honcho v3.0.11+). + +### Fixed + +- Scope-managed filter keys (`observer`, `observed`, `session`) are now rejected with a clear error if passed in `filters`, instead of silently overriding the scope and returning conclusions from a different peer pair. Use `peer.conclusions` / `peer.conclusionsOf(target)` and the dedicated `session` option instead. `session_id` remains a valid filter on `query()`. + ## [2.1.2] - 2026-05-21 ### Added diff --git a/sdks/typescript/__tests__/conclusions.test.ts b/sdks/typescript/__tests__/conclusions.test.ts index 8e7a6a09..7716d3ac 100644 --- a/sdks/typescript/__tests__/conclusions.test.ts +++ b/sdks/typescript/__tests__/conclusions.test.ts @@ -278,6 +278,61 @@ describe('Conclusions', () => { }) }) + // =========================================================================== + // Scope-reserved filter guard + // =========================================================================== + + describe('reserved filter keys', () => { + test('list rejects observer/observed scope keys in filters', async () => { + const peer = await client.peer('reserved-list-peer', { metadata: {} }) + + for (const key of ['observer', 'observed', 'observer_id', 'observed_id']) { + await expect( + peer.conclusions.list({ filters: { [key]: 'someone-else' } }) + ).rejects.toThrow(/managed by this conclusion scope/) + } + }) + + test('list rejects session keys in filters (use the session option)', async () => { + const peer = await client.peer('reserved-list-session-peer', { metadata: {} }) + + await expect( + peer.conclusions.list({ filters: { session_id: 'sess' } }) + ).rejects.toThrow(/managed by this conclusion scope/) + await expect( + peer.conclusions.list({ filters: { session: 'sess' } }) + ).rejects.toThrow(/managed by this conclusion scope/) + }) + + test('query rejects observer/observed scope keys in filters', async () => { + const peer = await client.peer('reserved-query-peer', { metadata: {} }) + + for (const key of ['observer', 'observed', 'observer_id', 'observed_id']) { + await expect( + peer.conclusions.query('q', 10, undefined, { [key]: 'someone-else' }) + ).rejects.toThrow(/managed by this conclusion scope/) + } + }) + + test('query allows session_id in filters (no dedicated session param)', async () => { + const peer = await client.peer('reserved-query-session-peer', { metadata: {} }) + + // Should not throw the reserved-key guard; session_id is a normal filter + // for query. The call may return no matches, which is fine. + await expect( + peer.conclusions.query('q', 10, undefined, { session_id: 'sess' }) + ).resolves.toBeDefined() + }) + + test('non-reserved filters (level) still work on list', async () => { + const peer = await client.peer('reserved-allowed-peer', { metadata: {} }) + + await expect( + peer.conclusions.list({ filters: { level: 'explicit' } }) + ).resolves.toBeDefined() + }) + }) + // =========================================================================== // Conclusion Deletion (DELETE /conclusions/:id) // =========================================================================== diff --git a/sdks/typescript/__tests__/peer.test.ts b/sdks/typescript/__tests__/peer.test.ts index d8a3750c..fca85910 100644 --- a/sdks/typescript/__tests__/peer.test.ts +++ b/sdks/typescript/__tests__/peer.test.ts @@ -16,6 +16,7 @@ */ import { describe, test, expect, beforeAll, afterAll } from 'bun:test' +import { z } from 'zod' import { Honcho, Peer } from '../src' import { createTestClient, generateId, requireServer } from './setup' import { @@ -624,6 +625,44 @@ describe('Peer', () => { expect(response === null || typeof response === 'string').toBe(true) }) + test('chat with responseFormat as JSON schema object returns JSON string', async () => { + const peer = await client.peer('chat-rf-peer') + + const response = await peer.chat('What do you know?', { + responseFormat: { + type: 'object', + properties: { summary: { type: 'string' } }, + }, + }) + + expect(response === null || typeof response === 'string').toBe(true) + if (response !== null) { + expect(() => JSON.parse(response)).not.toThrow() + } + }) + + test('chat with responseFormat as Zod schema returns parsed object', async () => { + const peer = await client.peer('chat-rf-zod-peer') + const ResultSchema = z.object({ summary: z.string().optional() }) + + const response = await peer.chat('What do you know?', { + responseFormat: ResultSchema, + }) + + expect(response === null || typeof response === 'object').toBe(true) + }) + + test('chat with unsupported responseFormat is rejected by the server', async () => { + const peer = await client.peer('chat-rf-invalid-peer') + + // Non-object root is rejected with 422 + await expect( + peer.chat('What do you know?', { + responseFormat: { type: 'string' }, + }) + ).rejects.toThrow() + }) + // Streaming tests are in streaming.test.ts }) diff --git a/sdks/typescript/__tests__/validation.test.ts b/sdks/typescript/__tests__/validation.test.ts index 42d03cf7..97d71f71 100644 --- a/sdks/typescript/__tests__/validation.test.ts +++ b/sdks/typescript/__tests__/validation.test.ts @@ -5,7 +5,7 @@ */ import { describe, test, expect } from 'bun:test' -import { ZodError } from 'zod' +import { z, ZodError } from 'zod' import { ChatQuerySchema, ContextParamsSchema, @@ -67,6 +67,24 @@ describe('ChatQuerySchema', () => { } ) + test('responseFormat as plain JSON schema object is valid', () => { + const schema = { type: 'object', properties: { a: { type: 'string' } } } + const result = ChatQuerySchema.parse({ query: 'hello', responseFormat: schema }) + expect(result.responseFormat).toEqual(schema) + }) + + test('responseFormat as Zod schema instance is valid and passed through', () => { + const schema = z.object({ a: z.string() }) + const result = ChatQuerySchema.parse({ query: 'hello', responseFormat: schema }) + expect(result.responseFormat).toBe(schema) + }) + + test('responseFormat as a non-object throws', () => { + expect(() => + ChatQuerySchema.parse({ query: 'hello', responseFormat: 'not-a-schema' }) + ).toThrow(ZodError) + }) + // --- Missing required fields --- test('missing query throws', () => { diff --git a/sdks/typescript/package.json b/sdks/typescript/package.json index 4ac6cea4..9819ddf9 100644 --- a/sdks/typescript/package.json +++ b/sdks/typescript/package.json @@ -1,6 +1,6 @@ { "name": "@honcho-ai/sdk", - "version": "2.1.2", + "version": "2.2.0", "description": "Official DX Optimized TypeScript SDK for Honcho", "author": "Plastic Labs ", "license": "Apache-2.0", diff --git a/sdks/typescript/src/conclusions.ts b/sdks/typescript/src/conclusions.ts index 7c854b57..c9add4e2 100644 --- a/sdks/typescript/src/conclusions.ts +++ b/sdks/typescript/src/conclusions.ts @@ -3,6 +3,7 @@ import type { HonchoHTTPClient } from './http/client' import { Page } from './pagination' import type { Session } from './session' import type { + ConclusionLevel, ConclusionResponse, PageResponse, RepresentationOptions, @@ -10,6 +11,43 @@ import type { } from './types/api' import { normalizeSearchQuery, RepresentationOptionsSchema } from './validation' +/** + * Filter keys that define a conclusion scope (the observer/observed peer pair). + * They are set from the scope itself, so a caller must not pass them in `filters`. + */ +const SCOPE_RESERVED_KEYS = [ + 'observer', + 'observed', + 'observer_id', + 'observed_id', +] + +/** + * Throw if `filters` contains keys managed by the conclusion scope. + * + * The observer/observed peer pair (and, on `list`, the session) is fixed by the + * scope, so letting a user filter override it would silently return data from a + * different scope than requested. Fail loud instead. + */ +function rejectReservedFilterKeys( + filters: Record | undefined, + reserved: string[] +): void { + if (!filters) return + const clash = reserved.filter((k) => k in filters).sort() + if (clash.length > 0) { + let guidance = + 'Choose the peer pair via peer.conclusions / peer.conclusionsOf(target)' + if (reserved.includes('session') || reserved.includes('session_id')) { + guidance += '; use the session option to filter by session' + } + throw new Error( + `Filter key(s) ${clash.join(', ')} are managed by this conclusion scope ` + + `and cannot be passed in filters. ${guidance}.` + ) + } +} + /** * Parameters for creating a conclusion. */ @@ -32,6 +70,12 @@ export class Conclusion { readonly observerId: string readonly observedId: string readonly sessionId: string | null + /** + * Reasoning level: 'explicit' conclusions are extracted directly from + * messages; 'deductive'/'inductive'/'contradiction' are derived during + * dreaming. + */ + readonly level: ConclusionLevel readonly createdAt: string constructor( @@ -40,13 +84,15 @@ export class Conclusion { observerId: string, observedId: string, sessionId: string | null, - createdAt: string + createdAt: string, + level: ConclusionLevel = 'explicit' ) { this.id = id this.content = content this.observerId = observerId this.observedId = observedId this.sessionId = sessionId + this.level = level this.createdAt = createdAt } @@ -57,7 +103,8 @@ export class Conclusion { data.observer_id, data.observed_id, data.session_id, - data.created_at + data.created_at, + data.level ) } @@ -182,14 +229,26 @@ export class ConclusionScope { * @param options.page - Page number (1-indexed, default: 1) * @param options.size - Number of items per page (default: 50) * @param options.session - Optional session (ID string or Session object) to filter by + * @param options.filters - Optional additional filter criteria, merged with + * this scope's observer/observed (and session, if given). Supports the same + * operators as other list endpoints — e.g. `{ level: 'explicit' }` to get + * only conclusions extracted directly from messages (i.e. not derived during + * dreaming). See + * https://honcho.dev/docs/v3/documentation/features/advanced/using-filters * @returns Promise resolving to a Page of Conclusion objects */ async list(options?: { page?: number size?: number session?: string | Session + filters?: Record reverse?: boolean }): Promise> { + rejectReservedFilterKeys(options?.filters, [ + ...SCOPE_RESERVED_KEYS, + 'session', + 'session_id', + ]) const resolvedSessionId = options?.session ? typeof options.session === 'string' ? options.session @@ -198,9 +257,8 @@ export class ConclusionScope { const filters: Record = { observer_id: this.observer, observed_id: this.observed, - } - if (resolvedSessionId) { - filters.session_id = resolvedSessionId + ...(resolvedSessionId ? { session_id: resolvedSessionId } : {}), + ...options?.filters, } const reverse = options?.reverse @@ -227,22 +285,32 @@ export class ConclusionScope { /** * Semantic search for conclusions in this scope. + * + * @param query - The search query string + * @param topK - Maximum number of results to return (default: 10) + * @param distance - Maximum cosine distance threshold (0.0-1.0) + * @param filters - Optional additional filter criteria, merged with this + * scope's observer/observed. Supports the same operators as the list + * endpoint — e.g. `{ level: 'deductive' }` to search only conclusions + * derived during dreaming. See + * https://honcho.dev/docs/v3/documentation/features/advanced/using-filters */ async query( query: string, topK: number = 10, - distance?: number + distance?: number, + filters?: Record ): Promise { - const filters: Record = { - observer_id: this.observer, - observed_id: this.observed, - } - + rejectReservedFilterKeys(filters, SCOPE_RESERVED_KEYS) const response = await this._query({ query, top_k: topK, distance, - filters, + filters: { + observer_id: this.observer, + observed_id: this.observed, + ...filters, + }, }) return (response ?? []).map((item) => Conclusion.fromApiResponse(item)) diff --git a/sdks/typescript/src/index.ts b/sdks/typescript/src/index.ts index f6b11d26..90bff9f2 100644 --- a/sdks/typescript/src/index.ts +++ b/sdks/typescript/src/index.ts @@ -40,6 +40,7 @@ export { // API types (snake_case, for advanced usage) export type { + ConclusionLevel, ConclusionQueryParams, ConclusionResponse, MessageResponse, diff --git a/sdks/typescript/src/peer.ts b/sdks/typescript/src/peer.ts index b4c1d756..adb67cb2 100644 --- a/sdks/typescript/src/peer.ts +++ b/sdks/typescript/src/peer.ts @@ -1,3 +1,4 @@ +import { ZodType, z } from 'zod' import { API_VERSION } from './api-version' import { ConclusionScope } from './conclusions' import type { HonchoHTTPClient } from './http/client' @@ -229,12 +230,29 @@ export class Peer { ) } + /** + * Convert a responseFormat option (Zod schema or raw JSON Schema object) + * to the JSON Schema dict the API expects. + */ + private static toResponseFormatSchema( + responseFormat: ZodType | Record | undefined + ): Record | undefined { + if (!responseFormat) { + return undefined + } + if (responseFormat instanceof ZodType) { + return z.toJSONSchema(responseFormat) as Record + } + return responseFormat + } + private async _chat(params: { query: string stream?: boolean target?: string session_id?: string reasoning_level?: string + response_format?: Record }): Promise { await this._ensureWorkspace() return this._http.post( @@ -248,6 +266,7 @@ export class Peer { target?: string session_id?: string reasoning_level?: string + response_format?: Record }): Promise { await this._ensureWorkspace() return this._http.stream( @@ -362,14 +381,33 @@ export class Peer { * }) * ``` */ + async chat( + query: string, + options: { + target?: string | Peer + session?: string | Session + reasoningLevel?: string + responseFormat: ZodType + } + ): Promise async chat( query: string, options?: { target?: string | Peer session?: string | Session reasoningLevel?: string + responseFormat?: Record } - ): Promise { + ): Promise + async chat( + query: string, + options?: { + target?: string | Peer + session?: string | Session + reasoningLevel?: string + responseFormat?: ZodType | Record + } + ): Promise { const targetId = options?.target ? typeof options.target === 'string' ? options.target @@ -386,18 +424,28 @@ export class Peer { target: targetId, session: resolvedSessionId, reasoningLevel: options?.reasoningLevel, + responseFormat: options?.responseFormat, }) + const zodSchema = + options?.responseFormat instanceof ZodType + ? options.responseFormat + : undefined + const response = await this._chat({ query: chatParams.query, stream: false, target: chatParams.target, session_id: chatParams.session, reasoning_level: chatParams.reasoningLevel, + response_format: Peer.toResponseFormatSchema(options?.responseFormat), }) if (!response.content) { return null } + if (zodSchema) { + return zodSchema.parse(JSON.parse(response.content)) + } return response.content } @@ -442,6 +490,7 @@ export class Peer { target?: string | Peer session?: string | Session reasoningLevel?: string + responseFormat?: ZodType | Record } ): Promise { const targetId = options?.target @@ -460,6 +509,7 @@ export class Peer { target: targetId, session: resolvedSessionId, reasoningLevel: options?.reasoningLevel, + responseFormat: options?.responseFormat, }) const response = await this._chatStream({ @@ -467,6 +517,7 @@ export class Peer { target: chatParams.target, session_id: chatParams.session, reasoning_level: chatParams.reasoningLevel, + response_format: Peer.toResponseFormatSchema(options?.responseFormat), }) return createDialecticStream(response) diff --git a/sdks/typescript/src/types/api.ts b/sdks/typescript/src/types/api.ts index 65c6d88c..dbe4378d 100644 --- a/sdks/typescript/src/types/api.ts +++ b/sdks/typescript/src/types/api.ts @@ -74,6 +74,7 @@ export interface PeerChatParams { session_id?: string target?: string reasoning_level?: 'minimal' | 'low' | 'medium' | 'high' | 'max' + response_format?: Record } export interface PeerChatResponse { @@ -242,12 +243,23 @@ export interface MessageSearchParams { // Conclusion Types // ============================================================================= +/** + * Reasoning level of a conclusion. "explicit" conclusions are extracted + * directly from messages; the others are derived during dreaming. + */ +export type ConclusionLevel = + | 'explicit' + | 'deductive' + | 'inductive' + | 'contradiction' + export interface ConclusionResponse { id: string content: string observer_id: string observed_id: string session_id: string | null + level: ConclusionLevel created_at: string } diff --git a/sdks/typescript/src/validation.ts b/sdks/typescript/src/validation.ts index fed1b10b..48255833 100644 --- a/sdks/typescript/src/validation.ts +++ b/sdks/typescript/src/validation.ts @@ -312,6 +312,11 @@ export const ChatQuerySchema = z reasoningLevel: z .enum(['minimal', 'low', 'medium', 'high', 'max']) .optional(), + // A Zod schema (checked first — it is itself an object) or a raw JSON + // Schema object describing the desired response structure. + responseFormat: z + .union([z.instanceof(z.ZodType), z.record(z.string(), z.unknown())]) + .optional(), }) .strict() diff --git a/src/cache/client.py b/src/cache/client.py index 319fab6e..1ad8e3f5 100644 --- a/src/cache/client.py +++ b/src/cache/client.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio import logging from typing import Any, cast +from urllib.parse import urlparse, urlunparse import sentry_sdk from cashews import cache @@ -20,10 +21,99 @@ from src.config import settings logger = logging.getLogger(__name__) - _cache_lock = asyncio.Lock() +# Query parameters that carry secrets when configured via URL: +# redis-py accepts ``?password=`` (all querystring options become client +# kwargs) and cashews accepts ``?secret=`` (HMAC key for value signing). +_SENSITIVE_QUERY_PARAMS = frozenset({"password", "secret"}) + + +def _mask_sensitive_query(query: str) -> str: + """Mask values of secret-bearing query parameters. + + Operates on the raw query string (no decode/re-encode round trip) + so non-secret parameters are preserved byte-for-byte. + + Args: + query: The raw query string from a parsed URL. + + Returns: + The query string with sensitive values replaced by ``***``, or + the original string if no sensitive parameter is present. + """ + if not query: + return query + parts: list[str] = [] + changed = False + for part in query.split("&"): + name, sep, _value = part.partition("=") + if sep and name.lower() in _SENSITIVE_QUERY_PARAMS: + parts.append(f"{name}=***") + changed = True + else: + parts.append(part) + return "&".join(parts) if changed else query + + +def _redact_cache_url(url: str) -> str: + """Mask credentials in a Redis connection URL before logging. + + Given ``redis://:password@host:port/db`` returns + ``redis://:***@host:port/db``; secret-bearing query parameters + (``?password=``, ``?secret=``) are masked as well. A URL carrying + no credentials is returned unchanged. This function never raises + and never returns a credential: an invalid port is omitted from + the output, and a URL that cannot be parsed at all is replaced by + a generic placeholder rather than echoed back, so that logging + inside ``except`` blocks can neither crash startup nor leak the + secrets this helper exists to hide. + + Args: + url: The Redis connection URL to redact. + + Returns: + The URL with its credentials masked, the original URL if it + carries none, or ``""`` if parsing + fails entirely. + """ + try: + parsed = urlparse(url) + query = _mask_sensitive_query(parsed.query) + # .password only splits netloc and never raises, unlike .port + if parsed.password is None and query == parsed.query: + # A string with an "@" but no parsed authority (e.g. a URL + # missing its scheme, ":pass@host:6379/0") may still carry + # userinfo that urlparse could not see — never echo it. + if "@" in url and not parsed.netloc: + return "" + return url + netloc = parsed.netloc + if parsed.password is not None: + userinfo = parsed.username or "" + hostname = parsed.hostname or "" + # Preserve IPv6 brackets (urlparse strips them from .hostname) + if hostname and ":" in hostname and not hostname.startswith("["): + hostname = f"[{hostname}]" + netloc = f"{userinfo}:***@{hostname}" + try: + port = parsed.port + except ValueError: + # Invalid or out-of-range port: omit it rather than let + # the outer fallback echo the raw URL (and its password) + # back. + port = None + if port is not None: + netloc += f":{port}" + parsed = parsed._replace(netloc=netloc, query=query) + return urlunparse(parsed) + except (ValueError, TypeError): + # Unparseable URL: never return the raw input — it may contain + # the very password this helper exists to hide. + return "" + + def is_cache_enabled() -> bool: return settings.CACHE.ENABLED @@ -45,17 +135,21 @@ async def init_cache() -> None: cache.setup("mem://", pickle_type=PicklerType.SQLALCHEMY) return - # Setup cache with Redis backend + # Setup cache with Redis backend. CACHE_CLUSTER selects the + # cluster-aware client, which follows the MOVED redirects a Redis + # Cluster returns for keys hashed to another shard; the standalone + # client treats those as command errors. try: cache.setup( settings.CACHE.URL, pickle_type=PicklerType.SQLALCHEMY, + cluster=settings.CACHE.CLUSTER, ) except Exception as setup_err: logger.warning( "Cache setup failed for %s: %s. Falling back to in-memory cache", - settings.CACHE.URL, + _redact_cache_url(settings.CACHE.URL), setup_err, ) if settings.SENTRY.ENABLED: @@ -83,7 +177,10 @@ async def init_cache() -> None: with attempt: async with asyncio.timeout(2): await cache.ping() - logger.info("Connected to cache at %s", settings.CACHE.URL) + logger.info( + "Connected to cache at %s", + _redact_cache_url(settings.CACHE.URL), + ) except ( redis_exc.TimeoutError, redis_exc.ConnectionError, @@ -92,7 +189,7 @@ async def init_cache() -> None: ) as e: logger.warning( "Failed to connect to cache at %s: %s. Falling back to in-memory cache", - settings.CACHE.URL, + _redact_cache_url(settings.CACHE.URL), e, ) if settings.SENTRY.ENABLED: @@ -103,7 +200,7 @@ async def init_cache() -> None: except Exception as e: logger.warning( "Unexpected cache error at %s: %s. Falling back to in-memory cache", - settings.CACHE.URL, + _redact_cache_url(settings.CACHE.URL), e, ) if settings.SENTRY.ENABLED: diff --git a/src/config.py b/src/config.py index 29210478..91e5c84d 100644 --- a/src/config.py +++ b/src/config.py @@ -61,6 +61,10 @@ ThinkingEffortLevel = Literal[ "none", "minimal", "low", "medium", "high", "xhigh", "max" ] +# "json_object" injects the schema into the prompt for OpenAI-compatible +# providers that don't support json_schema (Structured Outputs). +StructuredOutputMode = Literal["json_schema", "json_object"] + class ModelOverrideSettings(BaseModel): """Advanced module-level transport overrides.""" @@ -69,7 +73,23 @@ class ModelOverrideSettings(BaseModel): api_key_env: str | None = None base_url: str | None = None - provider_params: dict[str, Any] = Field(default_factory=dict) + provider_params: dict[str, Any] = Field( + default_factory=dict, + description=( + "Operator escape hatch for provider-specific request fields. " + "Three recognized keys: `extra_body` (merged into the request body), " + "`extra_headers` (HTTP headers), `extra_query` (URL query params). " + "OpenAI and Anthropic transports forward these as identically-named " + "SDK kwargs. The Gemini transport merges `extra_body` into the " + "GenerateContentConfig dict and folds `extra_headers` into " + "`http_options.headers`; `extra_query` is unsupported. Shallow merge " + "with operator-wins — if Honcho and the operator both set the same " + "key inside `extra_body`, the operator's value replaces Honcho's. " + "Operators are responsible for picking a coherent combination of " + "this and other config (e.g. unset `thinking_budget_tokens` when " + "supplying an `extra_body.thinking` for Anthropic-via-proxy)." + ), + ) class PromptCachePolicy(BaseModel): @@ -117,6 +137,23 @@ def _validate_thinking_constraints( raise ValueError("thinking_budget_tokens must be >= 1024 for Anthropic models") +def _validate_structured_output_mode( + transport: ModelTransport, structured_output_mode: StructuredOutputMode | None +) -> None: + """Reject ``structured_output_mode`` on transports that ignore it. + + Only the OpenAI backend honors this setting (it controls the json_schema vs + json_object structured-output path). On the anthropic/gemini transports it is + a silent no-op, so a value set there is a misconfiguration — fail fast at + startup rather than letting the operator wonder why it has no effect. + """ + if structured_output_mode is not None and transport != "openai": + raise ValueError( + "structured_output_mode is only supported on the 'openai' transport; " + + f"remove it from the '{transport}' model config" + ) + + class FallbackModelSettings(BaseModel): """Independent fallback model configuration. No inheritance from primary.""" @@ -136,6 +173,8 @@ class FallbackModelSettings(BaseModel): ) thinking_budget_tokens: int | None = None + structured_output_mode: StructuredOutputMode | None = None + max_output_tokens: int | None = None stop_sequences: list[str] | None = None @@ -155,6 +194,7 @@ class FallbackModelSettings(BaseModel): @model_validator(mode="after") def _validate_runtime_shape(self) -> "FallbackModelSettings": _validate_thinking_constraints(self.transport, self.thinking_budget_tokens) + _validate_structured_output_mode(self.transport, self.structured_output_mode) return self @@ -179,6 +219,8 @@ class ConfiguredModelSettings(BaseModel): ) thinking_budget_tokens: int | None = None + structured_output_mode: StructuredOutputMode | None = None + max_output_tokens: int | None = None stop_sequences: list[str] | None = None @@ -199,6 +241,7 @@ class ConfiguredModelSettings(BaseModel): @model_validator(mode="after") def _validate_runtime_shape(self) -> "ConfiguredModelSettings": _validate_thinking_constraints(self.transport, self.thinking_budget_tokens) + _validate_structured_output_mode(self.transport, self.structured_output_mode) return self @@ -223,6 +266,7 @@ class ResolvedFallbackConfig(BaseModel): validation_alias=AliasChoices("thinking_effort", "reasoning_effort"), ) thinking_budget_tokens: int | None = None + structured_output_mode: StructuredOutputMode | None = None provider_params: dict[str, Any] = Field(default_factory=dict) max_output_tokens: int | None = None @@ -258,6 +302,7 @@ class ModelConfig(BaseModel): validation_alias=AliasChoices("thinking_effort", "reasoning_effort"), ) thinking_budget_tokens: int | None = None + structured_output_mode: StructuredOutputMode | None = None provider_params: dict[str, Any] = Field(default_factory=dict) max_output_tokens: int | None = None @@ -394,6 +439,7 @@ def _resolve_fallback_config( seed=fallback.seed, thinking_effort=fallback.thinking_effort, thinking_budget_tokens=fallback.thinking_budget_tokens, + structured_output_mode=fallback.structured_output_mode, provider_params=fallback.overrides.provider_params, max_output_tokens=fallback.max_output_tokens, stop_sequences=fallback.stop_sequences, @@ -427,6 +473,7 @@ def resolve_model_config(configured: ConfiguredModelSettings) -> ModelConfig: seed=configured.seed, thinking_effort=configured.thinking_effort, thinking_budget_tokens=configured.thinking_budget_tokens, + structured_output_mode=configured.structured_output_mode, provider_params=configured.overrides.provider_params, max_output_tokens=configured.max_output_tokens, stop_sequences=configured.stop_sequences, @@ -612,8 +659,9 @@ class DBSettings(HonchoSettings): POOL_PRE_PING: bool = True POOL_SIZE: Annotated[int, Field(default=10, gt=0, le=1000)] = 10 MAX_OVERFLOW: Annotated[int, Field(default=20, ge=0, le=1000)] = 20 - POOL_TIMEOUT: Annotated[int, Field(default=30, gt=0, le=300)] = ( - 30 # seconds (max 5 minutes) + POOL_TIMEOUT: Annotated[int, Field(default=5, gt=0, le=300)] = ( + 5 # seconds a pooled checkout may wait for a free connection (QueuePool + # only; NullPool has no local queue wait) ) POOL_RECYCLE: Annotated[int, Field(default=300, gt=0, le=7200)] = ( 300 # seconds (max 2 hours) @@ -622,6 +670,13 @@ class DBSettings(HonchoSettings): SQL_DEBUG: bool = False TRACING: bool = False + # Per-connection establish timeout (seconds) passed to the driver, so a + # single connection attempt fails fast instead of hanging when the server or + # pooler is unreachable or stalled. Connection acquisition is a single + # attempt with no retry; callers handle failure (the API surfaces it, the + # deriver backs off and retries on a later poll). + CONNECT_TIMEOUT_SECONDS: Annotated[int, Field(default=2, gt=0, le=60)] = 2 + class AuthSettings(HonchoSettings): model_config = SettingsConfigDict(env_prefix="AUTH_", extra="ignore") # pyright: ignore @@ -696,6 +751,13 @@ class EmbeddingSettings(HonchoSettings): VECTOR_DIMENSIONS: Annotated[int, Field(default=1536, gt=0)] = 1536 MAX_INPUT_TOKENS: Annotated[int, Field(default=8192, gt=0)] = 8192 MAX_TOKENS_PER_REQUEST: Annotated[int, Field(default=300_000, gt=0)] = 300_000 + # Caps concurrent message-embedding fan-out on the API request path (the + # immediate-embed background task). The reconciler is unaffected. + MAX_CONCURRENT_EMBEDDINGS: Annotated[int, Field(default=10, gt=0, le=100)] = 10 + # Caps in-flight immediate-embed background tasks per API process. When + # saturated, message creation skips the fast path entirely and the + # reconciler embeds on its next cycle. 0 disables the fast path. + MAX_PENDING_EMBED_TASKS: Annotated[int, Field(default=50, ge=0)] = 50 @model_validator(mode="before") @classmethod @@ -737,7 +799,34 @@ class DeriverSettings(HonchoSettings): POLLING_SLEEP_INTERVAL_SECONDS: Annotated[ float, Field(default=1.0, gt=0.0, le=60.0) ] = 1.0 + # Adaptive polling: when the queue is idle (or the loop is erroring) the + # sleep interval grows from POLLING_SLEEP_INTERVAL_SECONDS toward + # POLLING_SLEEP_MAX_INTERVAL_SECONDS by POLLING_BACKOFF_MULTIPLIER each + # cycle, then snaps back to the base interval as soon as work is found. + # Reduces steady-state query load against the (shared) DB/pooler. + POLLING_BACKOFF_ENABLED: bool = True + POLLING_SLEEP_MAX_INTERVAL_SECONDS: Annotated[ + float, Field(default=30.0, gt=0.0, le=300.0) + ] = 30.0 + POLLING_BACKOFF_MULTIPLIER: Annotated[ + float, Field(default=2.0, ge=1.0, le=10.0) + ] = 2.0 + # Sleep a uniform-random delay in [0, POLLING_STARTUP_JITTER_SECONDS] before + # the first poll so instances that start together don't poll in lockstep. + # Set to 0.0 to disable. + POLLING_STARTUP_JITTER_SECONDS: Annotated[ + float, Field(default=30.0, ge=0.0, le=300.0) + ] = 30.0 + # Multiply every poll sleep by a random factor in [1 - ratio, 1 + ratio] + # (0.5 -> [0.5x, 1.5x]) so poll loops don't re-converge over time. The + # backoff schedule is unchanged; only the returned sleep is scattered. Set + # to 0.0 to disable. + POLLING_JITTER_RATIO: Annotated[float, Field(default=0.5, ge=0.0, le=1.0)] = 0.5 STALE_SESSION_TIMEOUT_MINUTES: Annotated[int, Field(default=5, gt=0, le=1440)] = 5 + # Minimum (jittered) spacing between stale-work-unit cleanup runs + STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS: Annotated[ + float, Field(default=60.0, ge=0.0, le=3600.0) + ] = 60.0 # Retention window (seconds) for keeping errored items in the queue QUEUE_ERROR_RETENTION_SECONDS: Annotated[ @@ -772,10 +861,28 @@ class DeriverSettings(HonchoSettings): int, Field(default=100, gt=0, le=1000) ] = 100 - REPRESENTATION_BATCH_MAX_TOKENS: Annotated[ + # Minimum tokens a representation work unit must accumulate (summed over + # its own unprocessed messages) before it becomes claimable. Bypassed by + # FLUSH_ENABLED and by REPRESENTATION_BATCH_MAX_AGE_SECONDS age-flushing. + # 0 disables the accumulation gate entirely (equivalent to FLUSH_ENABLED + # for claiming): work units are claimable as soon as anything is pending. + REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS: Annotated[ + int, + Field(default=512, ge=0, le=16_384), + ] = 512 + # Cumulative-token cap on the conversation window (queued messages plus + # interleaved context) fed to a single deriver LLM call when draining a + # claimed work unit. The first unprocessed message is always included, + # even if it alone exceeds the cap. + REPRESENTATION_BATCH_TARGET_INPUT_TOKENS: Annotated[ int, Field(default=1024, ge=128, le=16_384), ] = 1024 + # Sub-threshold work units become eligible once their oldest unprocessed + # item exceeds this age. 0 disables age-based flushing. + REPRESENTATION_BATCH_MAX_AGE_SECONDS: Annotated[int, Field(default=1800, ge=0)] = ( + 1800 + ) # When enabled, bypasses the batch token threshold and processes work immediately FLUSH_ENABLED: bool = False @@ -793,9 +900,9 @@ class DeriverSettings(HonchoSettings): @model_validator(mode="after") def validate_batch_tokens_vs_context_limit(self): - if self.REPRESENTATION_BATCH_MAX_TOKENS > self.MAX_INPUT_TOKENS: + if self.REPRESENTATION_BATCH_TARGET_INPUT_TOKENS > self.MAX_INPUT_TOKENS: raise ValueError( - f"REPRESENTATION_BATCH_MAX_TOKENS ({self.REPRESENTATION_BATCH_MAX_TOKENS}) cannot exceed max deriver input tokens ({self.MAX_INPUT_TOKENS})" + f"REPRESENTATION_BATCH_TARGET_INPUT_TOKENS ({self.REPRESENTATION_BATCH_TARGET_INPUT_TOKENS}) cannot exceed max deriver input tokens ({self.MAX_INPUT_TOKENS})" ) return self @@ -1086,12 +1193,29 @@ class TelemetrySettings(HonchoSettings): # that join high-volume events to aggregate envelopes first. HIGH_VOLUME_SAMPLE_RATE: Annotated[float, Field(default=1.0, ge=0.0, le=1.0)] = 1.0 + # --- Full-fidelity payload tracing (llm.call.traced / trace.content) --- + # Master toggle for replay-grade content capture. Default-off. + TRACE_PAYLOADS_ENABLED: bool = False + + # Per-message cap (bytes) for captured content; oversized string content is + # clipped (with a marker) and the call is flagged was_truncated. + TRACE_MAX_BYTES: Annotated[int, Field(default=262144, gt=0)] = 262144 + + # Allowlist of CallPurpose values to capture; empty = all. Typed as str to + # keep the enum out of config (validated against CallPurpose at the producer, + # same pattern as LLMTelemetryContext.call_purpose). + TRACE_PURPOSES: list[str] = Field(default_factory=list) + class CacheSettings(HonchoSettings): model_config = SettingsConfigDict(env_prefix="CACHE_", extra="ignore") # pyright: ignore ENABLED: bool = False URL: str = "redis://localhost:6379/0?suppress=true" + # URL points at a Redis Cluster (OSS cluster protocol, e.g. GCP Memorystore + # for Redis Cluster). A standalone client cannot follow the MOVED redirects + # such deployments return for keys hashed to another shard. + CLUSTER: bool = False NAMESPACE: str | None = None DEFAULT_TTL_SECONDS: Annotated[int, Field(default=300, ge=1, le=86_400)] = ( 300 # how long to keep items in cache @@ -1101,6 +1225,12 @@ class CacheSettings(HonchoSettings): 5 # how long to hold a lock on a resource when fetching DB after cache miss ) + # Polling interval while waiting for another worker's fetch lock. cashews + # defaults to 0, which busy-spins the event loop for the whole wait. + LOCK_WAIT_CHECK_INTERVAL_SECONDS: Annotated[ + float, Field(default=0.1, gt=0, le=5) + ] = 0.1 + class SurprisalSettings(BaseModel): """Settings for tree-based surprisal sampling during dreams.""" @@ -1264,6 +1394,17 @@ class VectorStoreSettings(HonchoSettings): return self +class TraceViewerSettings(HonchoSettings): + model_config = SettingsConfigDict(env_prefix="TRACE_VIEWER_", extra="ignore") # pyright: ignore + + ENABLED: bool = False + HOST: str = "127.0.0.1" + PORT: int = 8002 + STORAGE_DIR: str = "./traces" + MAX_REQUEST_BYTES: int = 10 * 1024 * 1024 # 10 MB + VENDOR_CDN_BASE: str = "https://cdn.jsdelivr.net/npm" + + class AppSettings(HonchoSettings): # No env_prefix for app-level settings model_config = SettingsConfigDict( # pyright: ignore @@ -1272,6 +1413,7 @@ class AppSettings(HonchoSettings): # Application-wide settings LOG_LEVEL: str = "INFO" + PERFORMANCE_LOG_FORMAT: str = "compact" SESSION_OBSERVERS_LIMIT: Annotated[int, Field(default=10, gt=0)] = 10 MAX_FILE_SIZE: Annotated[int, Field(default=5_242_880, gt=0)] = 5_242_880 # 5MB GET_CONTEXT_MAX_TOKENS: Annotated[int, Field(default=100_000, gt=0, le=250_000)] = ( @@ -1282,6 +1424,36 @@ class AppSettings(HonchoSettings): EMBED_MESSAGES: bool = True LANGFUSE_HOST: str | None = None LANGFUSE_PUBLIC_KEY: str | None = None + # How Langfuse traces are produced: + # "exporter" (default) — Langfuse is a projection over the captured + # CapturedLLMCall stream (LangfuseExporter), the same source of truth as + # the CloudEvents trace stream. + # "inline" — legacy live instrumentation (@observe + propagate_attributes + # spans during execution). Kept one release for side-by-side validation. + LANGFUSE_EXPORTER_MODE: Literal["inline", "exporter"] = "exporter" + + @property + def langfuse_inline_enabled(self) -> bool: + """True when the legacy inline Langfuse instrumentation is active + (keys configured + ``LANGFUSE_EXPORTER_MODE == "inline"``).""" + return ( + bool(self.LANGFUSE_PUBLIC_KEY) and self.LANGFUSE_EXPORTER_MODE == "inline" + ) + + @property + def langfuse_exporter_enabled(self) -> bool: + """True when the Langfuse exporter (a projection over the captured call + stream) is active (keys configured + ``LANGFUSE_EXPORTER_MODE == "exporter"``).""" + return ( + bool(self.LANGFUSE_PUBLIC_KEY) and self.LANGFUSE_EXPORTER_MODE == "exporter" + ) + + # Origins allowed by the FastAPI CORSMiddleware + CORS_ORIGINS: list[str] = [ + "http://localhost", + "http://127.0.0.1:8000", + "https://api.honcho.dev", + ] COLLECT_METRICS_LOCAL: bool = False LOCAL_METRICS_FILE: str = "metrics.jsonl" @@ -1305,6 +1477,7 @@ class AppSettings(HonchoSettings): CACHE: CacheSettings = Field(default_factory=CacheSettings) DREAM: DreamSettings = Field(default_factory=DreamSettings) VECTOR_STORE: VectorStoreSettings = Field(default_factory=VectorStoreSettings) + TRACE_VIEWER: TraceViewerSettings = Field(default_factory=TraceViewerSettings) @field_validator("LOG_LEVEL") def validate_log_level(cls, v: str) -> str: @@ -1313,6 +1486,13 @@ class AppSettings(HonchoSettings): raise ValueError(f"Invalid log level: {v}") return log_level + @field_validator("PERFORMANCE_LOG_FORMAT") + def validate_performance_log_format(cls, v: str) -> str: + log_format = v.lower() + if log_format not in ["compact", "rich"]: + raise ValueError(f"Invalid performance log format: {v}") + return log_format + @model_validator(mode="after") def propagate_namespace(self) -> "AppSettings": """Propagate top-level NAMESPACE to nested settings if not explicitly set.""" diff --git a/src/crud/__init__.py b/src/crud/__init__.py index 58bcd9f0..b34e4d17 100644 --- a/src/crud/__init__.py +++ b/src/crud/__init__.py @@ -5,6 +5,7 @@ from .collection import ( ) from .deriver import get_deriver_status, get_queue_status from .document import ( + CreateDocumentsResult, create_documents, create_observations, delete_document, @@ -83,6 +84,7 @@ __all__ = [ "get_deriver_status", "get_queue_status", # Document + "CreateDocumentsResult", "create_documents", "create_observations", "fetch_documents_by_ids", diff --git a/src/crud/collection.py b/src/crud/collection.py index 2790cec1..63a775e3 100644 --- a/src/crud/collection.py +++ b/src/crud/collection.py @@ -48,6 +48,7 @@ def collection_cache_key(workspace_name: str, observer: str, observed: str) -> s key=COLLECTION_CACHE_KEY_TEMPLATE, ttl=f"{settings.CACHE.DEFAULT_LOCK_TTL_SECONDS}s", prefix=COLLECTION_LOCK_PREFIX, + check_interval=settings.CACHE.LOCK_WAIT_CHECK_INTERVAL_SECONDS, ) async def _fetch_collection( db: AsyncSession, diff --git a/src/crud/document.py b/src/crud/document.py index fd068d6a..6b7810ca 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -1,5 +1,7 @@ import datetime from collections.abc import Sequence +from dataclasses import dataclass, field +from enum import Enum from logging import getLogger from typing import Any, cast @@ -176,7 +178,8 @@ async def query_documents_most_derived( limit: Maximum number of documents to return Returns: - Sequence of documents ordered by times_derived descending + Sequence of documents ordered by times_derived descending, + ties broken by created_at descending (most recent first) """ stmt = ( select(models.Document) @@ -186,7 +189,13 @@ async def query_documents_most_derived( models.Document.observed == observed, models.Document.deleted_at.is_(None), ) - .order_by(models.Document.times_derived.desc()) + .order_by( + models.Document.times_derived.desc(), + models.Document.created_at.desc(), + # created_at is the transaction timestamp, so documents created in + # the same batch share it -- id keeps the order deterministic. + models.Document.id, + ) .limit(limit) ) @@ -369,7 +378,7 @@ async def query_documents( max_distance, top_k, ) - async with tracked_db("query_documents.pgvector") as managed_db: + async with tracked_db("query_documents.pgvector", read_only=True) as managed_db: docs = await _query_documents_pgvector( managed_db, workspace_name, @@ -407,7 +416,7 @@ async def query_documents( document_ids=document_ids, filters=filters, ) - async with tracked_db("query_documents.fetch") as managed_db: + async with tracked_db("query_documents.fetch", read_only=True) as managed_db: docs = await fetch_documents_by_ids( db=managed_db, workspace_name=workspace_name, @@ -421,6 +430,52 @@ async def query_documents( return docs +def _normalize_content(content: str) -> str: + """Normalize document content for exact-match deduplication. + + Content is compared after trimming surrounding whitespace and lowercasing + + The SQL filter in ``create_documents`` must stay in sync with this: + ``lower(regexp_replace(content, '^\\s+|\\s+$', '', 'g'))``. Postgres' + ``trim()`` only strips spaces, so a regex is used to match Python's + ``str.strip()`` across all whitespace. + """ + return content.strip().lower() + + +def _dedup_key( + content: str, level: str, session_name: str | None +) -> tuple[str, str, str | None]: + """Build the exact-match dedup key for a document. + + Dedup never crosses levels: a same-content document at a different level is + a different kind of record (an explicit fact is not interchangeable with a + deductive conclusion that happens to share its text). + + For **explicit** documents dedup additionally never crosses sessions. + Explicit documents are session-pure records of what was derived from that + session's messages — the Scopes copy-by-session model depends on this — so + a repeat of the same fact in a different session must produce a new + document in that session rather than reinforce another session's row. + Derived levels (deductive/inductive/contradiction) are consolidations and + may still dedup across sessions. + """ + return ( + _normalize_content(content), + level, + session_name if level == "explicit" else None, + ) + + +@dataclass +class CreateDocumentsResult: + created_documents: list[schemas.DocumentCreate] = field(default_factory=list) + exact_dup_in_batch_count: int = 0 + exact_dup_existing_count: int = 0 + semantic_dup_rejected_count: int = 0 + semantic_dup_replaced_count: int = 0 + + async def create_documents( db: AsyncSession, documents: list[schemas.DocumentCreate], @@ -429,16 +484,21 @@ async def create_documents( observer: str, observed: str, deduplicate: bool = False, -) -> list[schemas.DocumentCreate]: +) -> CreateDocumentsResult: """ Create multiple documents with optional duplicate detection. + The ``deduplicate`` flag additionally enables semantic (cosine-similarity) + dedup via ``is_rejected_duplicate`` for documents that survive the exact + deduplication check. + Args: db: Database session documents: List of document creation schemas workspace_name: Name of the workspace observer: Name of the observing peer observed: Name of the observed peer + deduplicate: Enable semantic duplicate detection Returns: List of DocumentCreate schemas that were actually inserted (excludes @@ -449,16 +509,116 @@ async def create_documents( # Store (document_model, embedding) pairs - IDs aren't available until after commit docs_with_embeddings: list[tuple[models.Document, list[float]]] = [] + # exact-content dedup (independent of `deduplicate`): pre-fetch + # existing live documents whose normalized content matches anything in this + # batch, scoped to (workspace, observer, observed). The SQL normalization must + # mirror _normalize_content. Matching is further scoped per-document by + # level (always) and session (for explicit documents) via _dedup_key. + batch_normalized: set[str] = {_normalize_content(d.content) for d in documents} + existing_by_key: dict[tuple[str, str, str | None], models.Document] = {} + if batch_normalized: + # The `normalized_content_sql.in_(...)` filter below narrows to the + # (workspace, observer, observed) partition via the single-column indexes, + # then evaluates lower(regexp_replace(...)) per row. + # TODO: add a partial expression index matching + # this filter exactly + # CREATE INDEX ix_documents_normalized_content + # ON documents ( + # workspace_name, + # observer, + # observed, + # (lower(regexp_replace(content, '^\s+|\s+$', '', 'g'))) + # ) + # WHERE deleted_at IS NULL; + normalized_content_sql = func.lower( + func.regexp_replace(models.Document.content, r"^\s+|\s+$", "", "g") + ) + existing_result = await db.execute( + select(models.Document).where( + models.Document.workspace_name == workspace_name, + models.Document.observer == observer, + models.Document.observed == observed, + models.Document.deleted_at.is_(None), + normalized_content_sql.in_(batch_normalized), + ) + ) + for existing_doc in existing_result.scalars(): + # If multiple historical rows share a dedup key, reinforcing + # one is sufficient; keep the first. + existing_by_key.setdefault( + _dedup_key( + existing_doc.content, + existing_doc.level, + existing_doc.session_name, + ), + existing_doc, + ) + + # Tracks dedup keys already accepted from this batch so exact + # duplicates within a single inference call collapse to one document. + seen_in_batch: set[tuple[str, str, str | None]] = set() + + exact_dup_existing_count = 0 + exact_dup_in_batch_count = 0 + semantic_dup_rejected_count = 0 + semantic_dup_replaced_count = 0 for doc in documents: try: + # Session-purity invariant: an explicit document must always carry + # the session it was derived from. Refuse to write session-less + # explicit documents rather than silently minting global explicit + # memory (the Scopes copy-by-session model depends on explicit + # documents staying session-pure). + if doc.level == "explicit" and doc.session_name is None: + logger.error( + "Refusing to create explicit document without session_name in %s/%s/%s (session-purity invariant): %r", + workspace_name, + observer, + observed, + doc.content[:80], + ) + continue + + dedup_key = _dedup_key(doc.content, doc.level, doc.session_name) + + # Exact-match dedup, always on: + # 1) collapse exact duplicates within this batch (drop silently). + if dedup_key in seen_in_batch: + exact_dup_in_batch_count += 1 + continue + seen_in_batch.add(dedup_key) + + # 2) drop exact duplicates of an existing live document, recording + # the re-derivation as reinforcement on the existing row. + existing_match = existing_by_key.get(dedup_key) + if existing_match is not None: + # Reinforce the existing row. greatest(...) keeps the bump atomic + # server-side (concurrent workers can't lose an increment) while + # still honoring an incoming doc that already carries accumulated + # reinforcement (times_derived > 1, e.g. a future re-ingestion or + # collection-merge path). Mirrors the superior-replacement branch + # in is_rejected_duplicate. + existing_match.times_derived = func.greatest( + models.Document.times_derived + 1, + doc.times_derived, + ) + await db.flush() + exact_dup_existing_count += 1 + continue + # for each document, if deduplicate is True, perform a process # that checks against existing documents and either rejects this document # as a duplicate OR deletes an existing document that is a duplicate. if deduplicate: - is_duplicate = await is_rejected_duplicate( + duplicate_result = await is_rejected_duplicate( db, doc, workspace_name, observer=observer, observed=observed ) - if is_duplicate: + if duplicate_result is SemanticRejectionResult.REPLACED_EXISTING: + # Existing doc was soft-deleted in favor of this one; the + # new doc still gets inserted below. + semantic_dup_replaced_count += 1 + elif duplicate_result is SemanticRejectionResult.REJECTED: + semantic_dup_rejected_count += 1 continue metadata_dict = doc.metadata.model_dump(exclude_none=True) @@ -610,7 +770,13 @@ async def create_documents( "Failed to create documents due to integrity constraint violation" ) from e - return accepted_documents + return CreateDocumentsResult( + created_documents=accepted_documents, + exact_dup_existing_count=exact_dup_existing_count, + exact_dup_in_batch_count=exact_dup_in_batch_count, + semantic_dup_rejected_count=semantic_dup_rejected_count, + semantic_dup_replaced_count=semantic_dup_replaced_count, + ) async def delete_document( @@ -960,6 +1126,12 @@ async def create_observations( return honcho_documents +class SemanticRejectionResult(Enum): + NOT_DUPLICATE = 0 + REPLACED_EXISTING = 1 + REJECTED = 2 + + async def is_rejected_duplicate( db: AsyncSession, doc: schemas.DocumentCreate, @@ -967,7 +1139,7 @@ async def is_rejected_duplicate( *, observer: str, observed: str, -) -> bool: +) -> SemanticRejectionResult: """ Check if a document is a duplicate of an existing document. @@ -980,8 +1152,27 @@ async def is_rejected_duplicate( If the document is not a duplicate, returns False. If the document is a duplicate AND the new document is superior, - deletes the existing document and returns False. + deletes the existing document and returns False. In this case + ``doc.times_derived`` is updated in place to carry the replaced + document's reinforcement count forward. + + If the document is a duplicate AND the existing document is superior, + increments the existing document's ``times_derived`` to record the + reinforcement, then returns True. + + Merges are scoped so they never cross document levels, and never cross + sessions for explicit-level documents (session-purity invariant: an + explicit document records what was derived from exactly one session, so + a near-duplicate from another session must not reinforce or replace it). """ + filters: dict[str, Any] = {"level": doc.level} + if doc.level == "explicit": + if doc.session_name is None: + # create_documents refuses session-less explicit documents; if one + # reaches here anyway it has no valid merge partner. + return SemanticRejectionResult.NOT_DUPLICATE + filters["session_name"] = doc.session_name + # Step 1: Find potential duplicates using cosine similarity similar_docs = await query_documents( db=db, @@ -989,13 +1180,14 @@ async def is_rejected_duplicate( query=doc.content, observer=observer, observed=observed, + filters=filters, max_distance=0.05, top_k=1, embedding=doc.embedding, ) if not similar_docs: - return False + return SemanticRejectionResult.NOT_DUPLICATE existing_doc = similar_docs[0] @@ -1011,19 +1203,37 @@ async def is_rejected_duplicate( # If new document has more or equal information, keep it and delete existing if score_new >= score_existing: - logger.warning( - f"[DUPLICATE DETECTION] Deleting existing in favor of new. new='{doc.content}', existing='{existing_doc.content}'." + logger.debug( + "[DUPLICATE DETECTION] Deleting existing in favor of new. new=%r, existing=%r.", + doc.content, + existing_doc.content, ) + # Carry the reinforcement count forward so replacing a duplicate counts as + # another derivation rather than resetting times_derived to 1. + doc.times_derived = max(doc.times_derived, existing_doc.times_derived + 1) # Soft-delete the existing document - reconciliation will clean up vectors and hard-delete existing_doc.deleted_at = datetime.datetime.now(datetime.timezone.utc) await db.flush() - return False # Don't reject the new document + return ( + SemanticRejectionResult.REPLACED_EXISTING + ) # Don't reject the new document - # Existing document has more information, reject the new one - logger.warning( - f"[DUPLICATE DETECTION] Rejecting new in favor of existing. new='{doc.content}', existing='{existing_doc.content}'." + # Existing document has more information, reject the new one but record the + # reinforcement: a semantic duplicate was derived again. greatest(...) keeps + # the increment atomic server-side -- concurrent workers reinforcing the same + # document must not lose updates -- while still honoring an incoming doc that + # already carries accumulated reinforcement (times_derived > 1). + existing_doc.times_derived = func.greatest( + models.Document.times_derived + 1, + doc.times_derived, ) - return True + await db.flush() + logger.debug( + "[DUPLICATE DETECTION] Rejecting new in favor of existing. new=%r, existing=%r.", + doc.content, + existing_doc.content, + ) + return SemanticRejectionResult.REJECTED async def cleanup_soft_deleted_documents( diff --git a/src/crud/message.py b/src/crud/message.py index ec673bb0..684ddb4d 100644 --- a/src/crud/message.py +++ b/src/crud/message.py @@ -4,19 +4,18 @@ from logging import getLogger from typing import Any from nanoid import generate as generate_nanoid -from sqlalchemy import ColumnElement, Select, and_, func, or_, select, text, update +from sqlalchemy import ColumnElement, Select, and_, func, or_, select, text from sqlalchemy.ext.asyncio import AsyncSession from src import models, schemas from src.config import settings from src.dependencies import tracked_db from src.embedding_client import embedding_client -from src.exceptions import VectorStoreError from src.telemetry.events import EmbeddingCallPurpose from src.utils.filter import apply_filter from src.utils.formatting import ILIKE_ESCAPE_CHAR, escape_ilike_pattern from src.utils.types import embedding_call_purpose -from src.vector_store import VectorRecord, get_external_vector_store +from src.vector_store import get_external_vector_store from .session import get_or_create_session @@ -56,11 +55,28 @@ async def get_peer_session_names( db: AsyncSession, workspace_name: str, peer_name: str, + *, + active_only: bool = False, ) -> list[str]: - """Get all session names where a peer has any membership record. + """Get all session names where a peer has a membership record. - Any membership record (regardless of joined_at/left_at) grants visibility - to all messages in that session. + By default any membership record (regardless of joined_at/left_at) grants + visibility to all messages in that session — this is the loose definition + recall scoping uses. + + Pass ``active_only=True`` for the strict definition (``left_at IS NULL``), + matching :func:`src.crud.session.is_peer_in_session`. The auth layer must + use the strict one so that a single peer-scoped key gets the same answer + whether it names a session directly or via a filter allowlist. + + Args: + db: Database session + workspace_name: Name of the workspace + peer_name: Name of the peer + active_only: Restrict to sessions the peer has not left + + Returns: + Distinct session names the peer has a matching membership record in. """ stmt = ( select(models.session_peers_table.c.session_name) @@ -68,10 +84,80 @@ async def get_peer_session_names( .where(models.session_peers_table.c.peer_name == peer_name) .distinct() ) + if active_only: + stmt = stmt.where(models.session_peers_table.c.left_at.is_(None)) result = await db.execute(stmt) return [row[0] for row in result.all()] +async def resolve_session_scope( + db: AsyncSession | None, + workspace_name: str, + session_name: str | None, + session_allowlist: list[str] | None, + observer: str | None, + *, + operation_name: str = "resolve_session_scope", +) -> tuple[list[str] | None, bool]: + """Resolve the effective session scope for a message query. + + Returns ``(allowed_session_names, deny)``: + + - ``allowed_session_names is None`` — apply no allowlist filter. Either the + query is unrestricted, or ``session_name`` already pins it to one session. + - a populated list — restrict the query to exactly these sessions. + - ``deny=True`` — the caller must return an empty result *without* querying. + + The distinction between ``None`` and an empty list is load-bearing: the + external vector stores drop an empty ``IN`` clause rather than matching + nothing, so collapsing the two would fail open. This function therefore + never returns an empty list — it returns ``deny=True`` instead. + + Touches the database only when an observer lookup is actually required, so + callers on the external-vector-store path don't check out a connection + before their network call. + + Args: + db: Database session to reuse. Pass None to let this function open its + own short-lived read-only session if (and only if) it needs one. + workspace_name: Name of the workspace + session_name: A single pinned session, if the caller named one + session_allowlist: Optional session allowlist. ``None`` is unrestricted; + an empty list fails closed. + observer: When set, scope is limited to this peer's sessions and then + intersected with ``session_allowlist`` + operation_name: Label for the self-managed DB session, when one is opened + + Returns: + Tuple of (allowlist to filter on or None, whether to deny outright). + """ + if session_name: + # A specific session was requested. Fail closed when the allowlist + # forbids it — routes guard this too, but other CRUD callers (the + # dialectic tools) don't, so enforce it at the boundary. + if session_allowlist is not None and session_name not in session_allowlist: + return None, True + return None, False + + if observer is None: + if session_allowlist is None: + return None, False + allowed = list(session_allowlist) + return (allowed, False) if allowed else (None, True) + + if db is not None: + allowed = await get_peer_session_names(db, workspace_name, observer) + else: + async with tracked_db(f"{operation_name}.peer_scope", read_only=True) as own_db: + allowed = await get_peer_session_names(own_db, workspace_name, observer) + + if session_allowlist is not None: + scope = set(session_allowlist) + allowed = [s for s in allowed if s in scope] + + return (allowed, False) if allowed else (None, True) + + def _apply_token_limit( base_conditions: list[ColumnElement[Any]], token_limit: int ) -> Select[tuple[models.Message]]: @@ -276,158 +362,37 @@ async def create_messages( db.add_all(message_objects) - # Commit here to release the advisory lock before generating embeddings - await db.commit() - try: - if settings.EMBED_MESSAGES: - id_resource_dict = { - message.public_id: message.content - for message in message_objects - if message.content and message.content.strip() - } - if id_resource_dict: - with embedding_call_purpose( - EmbeddingCallPurpose.MESSAGE_CREATE.value, - workspace_name=workspace_name, - parent_category="api", - ): - embedding_dict = await embedding_client.batch_embed( - id_resource_dict - ) - else: - embedding_dict = {} - - external_vector_store = get_external_vector_store() - - # Determine if we need to persist embeddings to postgres - # True when: TYPE=pgvector OR still migrating (dual-write to both stores) - store_embeddings_in_postgres = ( - settings.VECTOR_STORE.TYPE == "pgvector" - or not settings.VECTOR_STORE.MIGRATED - ) - - # Create MessageEmbedding entries - embedding_objects: list[models.MessageEmbedding] = [] - # Maps emb index -> (chunk_position, embedding vector) - pending_embedding_data: dict[int, tuple[int, list[float]]] = {} + # If embedding is enabled, locally chunk the content and insert + # one pending MessageEmbedding row per chunk in chunk order. The actual + # embedding work is deferred to the reconciler + if settings.EMBED_MESSAGES: + id_resource_dict = { + message_obj.public_id: message_obj.content + for message_obj in message_objects + if message_obj.content and message_obj.content.strip() + } + if id_resource_dict: + chunks_by_id = embedding_client.prepare_chunks(id_resource_dict) + peer_by_id = {m.public_id: m.peer_name for m in message_objects} + pending_rows: list[models.MessageEmbedding] = [] for message_obj in message_objects: - embeddings = embedding_dict.get(message_obj.public_id, []) - for chunk_position, embedding in enumerate(embeddings): - embedding_obj = models.MessageEmbedding( - content=message_obj.content, - message_id=message_obj.public_id, - workspace_name=workspace_name, - session_name=session_name, - peer_name=message_obj.peer_name, - sync_state="pending", - embedding=embedding if store_embeddings_in_postgres else None, - ) - emb_idx = len(embedding_objects) - pending_embedding_data[emb_idx] = (chunk_position, embedding) - embedding_objects.append(embedding_obj) - - # Always create MessageEmbedding rows so reconciliation can track sync state - # even when embeddings aren't stored in postgres - embedding_ids: list[int] = [] - if embedding_objects: - db.add_all(embedding_objects) - await db.flush() - embedding_ids = [emb.id for emb in embedding_objects] - - await db.commit() - - # If no external vector store (pgvector-only mode), mark as synced immediately - if external_vector_store is None: - if embedding_ids: - await db.execute( - update(models.MessageEmbedding) - .where(models.MessageEmbedding.id.in_(embedding_ids)) - .values( - sync_state="synced", - last_sync_at=func.now(), - sync_attempts=0, + chunks = chunks_by_id.get(message_obj.public_id, []) + for chunk_text in chunks: + pending_rows.append( + models.MessageEmbedding( + content=chunk_text, + message_id=message_obj.public_id, + workspace_name=workspace_name, + session_name=session_name, + peer_name=peer_by_id[message_obj.public_id], + sync_state="pending", + embedding=None, ) ) - await db.commit() - else: - # External vector store - build and upsert vector records - namespace = external_vector_store.get_vector_namespace( - "message", workspace_name - ) + if pending_rows: + db.add_all(pending_rows) - # Build vector records with {message_id}_{chunk_position} as vector ID - vector_records: list[VectorRecord] = [] - for emb_idx, emb in enumerate(embedding_objects): - chunk_position, embedding = pending_embedding_data[emb_idx] - vector_id = f"{emb.message_id}_{chunk_position}" - vector_records.append( - VectorRecord( - id=vector_id, - embedding=list(embedding), - metadata={ - "message_id": emb.message_id, - "session_name": emb.session_name, - "peer_name": emb.peer_name, - }, - ) - ) - - # Upsert to external vector store and update sync state - if vector_records: - try: - await external_vector_store.upsert_many( - namespace, vector_records - ) - # Success: mark as synced if we have DB rows - if embedding_ids: - await db.execute( - update(models.MessageEmbedding) - .where(models.MessageEmbedding.id.in_(embedding_ids)) - .values( - sync_state="synced", - last_sync_at=func.now(), - sync_attempts=0, - ) - ) - await db.commit() - - except VectorStoreError: - logger.warning( - "Vector store unavailable; leaving message vectors unsynced" - ) - if embedding_ids: - await db.execute( - update(models.MessageEmbedding) - .where(models.MessageEmbedding.id.in_(embedding_ids)) - .values( - sync_attempts=models.MessageEmbedding.sync_attempts - + 1, - last_sync_at=func.now(), - ) - ) - await db.commit() - - except Exception: - logger.exception("Unexpected error upserting message vectors") - if embedding_ids: - await db.execute( - update(models.MessageEmbedding) - .where(models.MessageEmbedding.id.in_(embedding_ids)) - .values( - sync_attempts=models.MessageEmbedding.sync_attempts - + 1, - last_sync_at=func.now(), - ) - ) - await db.commit() - - except Exception: - logger.exception( - "Failed to generate message embeddings for %s messages in workspace %s and session %s.", - len(message_objects), - workspace_name, - session_name, - ) + await db.commit() return message_objects @@ -770,6 +735,9 @@ async def _search_messages_pgvector( models.MessageEmbedding, models.Message.public_id == models.MessageEmbedding.message_id, ) + # Exclude pending rows that haven't been embedded yet: their NULL + # distance sorts last and would pad the window with unranked messages. + .where(models.MessageEmbedding.embedding.isnot(None)) .where(models.MessageEmbedding.workspace_name == workspace_name) .order_by(models.MessageEmbedding.embedding.cosine_distance(query_embedding)) .limit(limit * 2) @@ -808,21 +776,29 @@ async def _semantic_search_messages( after_date: datetime | None = None, before_date: datetime | None = None, observer: str | None = None, + session_allowlist: list[str] | None = None, ) -> list[tuple[list[models.Message], list[models.Message]]]: """Run semantic message search with optional temporal filters. When observer is provided and session_name is None, results are - scoped to sessions the observer has any membership record in. + scoped to sessions the observer has any membership record in. When + session_allowlist is provided, that membership scope is further + intersected with the allowlist (fail-closed: empty result on empty + intersection). """ - # Pre-fetch peer session scope if needed (short-lived DB session) - allowed_session_names: list[str] | None = None - if observer and not session_name: - async with tracked_db(f"{operation_name}.peer_scope") as db: - allowed_session_names = await get_peer_session_names( - db, workspace_name, observer - ) - if not allowed_session_names: - return [] + # db=None: the helper opens its own short-lived session only if it needs + # an observer lookup, so the external-store path below stays the first + # thing that happens when no observer scoping applies. + allowed_session_names, deny = await resolve_session_scope( + None, + workspace_name, + session_name, + session_allowlist, + observer, + operation_name=operation_name, + ) + if deny: + return [] if settings.VECTOR_STORE.TYPE != "pgvector" and settings.VECTOR_STORE.MIGRATED: message_ids = await _search_messages_external( @@ -837,7 +813,7 @@ async def _semantic_search_messages( if not message_ids: return [] - async with tracked_db(operation_name) as db: + async with tracked_db(operation_name, read_only=True) as db: matched_messages = ( await _fetch_messages_by_ids( db, @@ -853,7 +829,7 @@ async def _semantic_search_messages( _expunge_snippets(db, snippets) return snippets - async with tracked_db(operation_name) as db: + async with tracked_db(operation_name, read_only=True) as db: snippets = await _search_messages_pgvector( db, workspace_name, @@ -877,6 +853,7 @@ async def search_messages( context_window: int = 2, embedding: list[float] | None = None, observer: str | None = None, + session_allowlist: list[str] | None = None, ) -> list[tuple[list[models.Message], list[models.Message]]]: """ Search for messages using semantic similarity and return conversation snippets. @@ -887,12 +864,19 @@ async def search_messages( Args: workspace_name: Name of the workspace session_name: Name of the session (optional) + Deprecated for *scoping*: prefer session_allowlist, which + intersects with observer membership. This parameter also pins + the query to one session and bypasses observer scoping, so it + is not a drop-in equivalent and is not removed. query: Search query text limit: Maximum number of matching messages to return context_window: Number of messages before/after each match to include embedding: Optional pre-computed embedding observer: When provided and session_name is None, scope results to sessions this peer belongs to + session_allowlist: Optional session allowlist. None is unrestricted; an + empty list fails closed (empty result); a populated list is + intersected with the observer's session scope when observer is set Returns: List of tuples: (matched_messages, context_messages) @@ -918,6 +902,7 @@ async def search_messages( context_window=context_window, operation_name="message.search_messages", observer=observer, + session_allowlist=session_allowlist, ) @@ -965,6 +950,7 @@ async def grep_messages( limit: int = 10, context_window: int = 2, observer: str | None = None, + session_allowlist: list[str] | None = None, ) -> list[tuple[list[models.Message], list[models.Message]]]: """ Search for messages containing specific text (case-insensitive substring match). @@ -975,25 +961,29 @@ async def grep_messages( Args: workspace_name: Name of the workspace session_name: Name of the session (optional - searches all sessions if None) + Deprecated for *scoping*: prefer session_allowlist, which + intersects with observer membership. This parameter also pins + the query to one session and bypasses observer scoping, so it + is not a drop-in equivalent and is not removed. text: Text to search for (case-insensitive) limit: Maximum number of matching messages to return context_window: Number of messages before/after each match to include observer: When provided and session_name is None, scope results to sessions this peer belongs to + session_allowlist: Optional session allowlist. None is unrestricted; an + empty list fails closed (empty result); a populated list is + intersected with the observer's session scope when observer is set Returns: List of tuples: (matched_messages, context_messages) Each snippet may contain multiple matches if they were close together. """ - async with tracked_db("message.grep_messages") as db: - # Pre-fetch peer session scope if needed - allowed_session_names = None - if observer and not session_name: - allowed_session_names = await get_peer_session_names( - db, workspace_name, observer - ) - if not allowed_session_names: - return [] + async with tracked_db("message.grep_messages", read_only=True) as db: + allowed_session_names, deny = await resolve_session_scope( + db, workspace_name, session_name, session_allowlist, observer + ) + if deny: + return [] snippets = await _grep_messages_internal( db, @@ -1017,6 +1007,7 @@ async def get_messages_by_date_range( limit: int = 20, order: str = "desc", observer: str | None = None, + session_allowlist: list[str] | None = None, ) -> list[models.Message]: """ Get messages within a date range. @@ -1025,24 +1016,28 @@ async def get_messages_by_date_range( db: Database session workspace_name: Name of the workspace session_name: Name of the session (optional - searches all sessions if None) + Deprecated for *scoping*: prefer session_allowlist, which + intersects with observer membership. This parameter also pins + the query to one session and bypasses observer scoping, so it + is not a drop-in equivalent and is not removed. after_date: Return messages after this datetime before_date: Return messages before this datetime limit: Maximum messages to return order: Sort order - 'asc' for oldest first, 'desc' for newest first observer: When provided and session_name is None, scope results to sessions this peer belongs to + session_allowlist: Optional session allowlist. None is unrestricted; an + empty list fails closed (empty result); a populated list is + intersected with the observer's session scope when observer is set Returns: List of messages within the date range """ - # Pre-fetch peer session scope if needed - allowed_session_names = None - if observer and not session_name: - allowed_session_names = await get_peer_session_names( - db, workspace_name, observer - ) - if not allowed_session_names: - return [] + allowed_session_names, deny = await resolve_session_scope( + db, workspace_name, session_name, session_allowlist, observer + ) + if deny: + return [] stmt = select(models.Message).where(models.Message.workspace_name == workspace_name) @@ -1076,6 +1071,7 @@ async def search_messages_temporal( context_window: int = 2, embedding: list[float] | None = None, observer: str | None = None, + session_allowlist: list[str] | None = None, ) -> list[tuple[list[models.Message], list[models.Message]]]: """ Search for messages using semantic similarity with optional date filtering. @@ -1086,6 +1082,10 @@ async def search_messages_temporal( Args: workspace_name: Name of the workspace session_name: Name of the session (optional) + Deprecated for *scoping*: prefer session_allowlist, which + intersects with observer membership. This parameter also pins + the query to one session and bypasses observer scoping, so it + is not a drop-in equivalent and is not removed. query: Search query text after_date: Only return messages after this datetime before_date: Only return messages before this datetime @@ -1094,6 +1094,9 @@ async def search_messages_temporal( embedding: Optional pre-computed embedding for the query observer: When provided and session_name is None, scope results to sessions this peer belongs to + session_allowlist: Optional session allowlist. None is unrestricted; an + empty list fails closed (empty result); a populated list is + intersected with the observer's session scope when observer is set Returns: List of tuples: (matched_messages, context_messages) @@ -1120,4 +1123,5 @@ async def search_messages_temporal( context_window=context_window, operation_name="message.search_messages_temporal", observer=observer, + session_allowlist=session_allowlist, ) diff --git a/src/crud/peer.py b/src/crud/peer.py index 21792b0f..f7936761 100644 --- a/src/crud/peer.py +++ b/src/crud/peer.py @@ -152,6 +152,7 @@ async def get_or_create_peers( key=PEER_CACHE_KEY_TEMPLATE, ttl=f"{settings.CACHE.DEFAULT_LOCK_TTL_SECONDS}s", prefix=PEER_LOCK_PREFIX, + check_interval=settings.CACHE.LOCK_WAIT_CHECK_INTERVAL_SECONDS, ) async def _fetch_peer( db: AsyncSession, diff --git a/src/crud/representation.py b/src/crud/representation.py index bb969120..d4b86ffb 100644 --- a/src/crud/representation.py +++ b/src/crud/representation.py @@ -19,9 +19,11 @@ from src.telemetry.events import EmbeddingCallPurpose from src.telemetry.logging import accumulate_metric from src.utils.formatting import format_datetime_utc from src.utils.representation import ( + ALLOWLIST_SAFE_LEVELS, DeductiveObservation, ExplicitObservation, Representation, + allowlist_safe_levels, ) from src.utils.types import embedding_call_purpose @@ -64,7 +66,7 @@ class RepresentationManager: session_name: str, message_created_at: datetime.datetime, message_level_configuration: ResolvedConfiguration, - ) -> int: + ) -> crud.CreateDocumentsResult: """ Save Representation objects to the collection as a set of documents. @@ -75,14 +77,15 @@ class RepresentationManager: message_created_at: Timestamp when the message was created Returns: - The number of *new documents saved* + The result of document creation, including saved documents and + deduplication counts. """ - new_documents = 0 + empty_result = crud.CreateDocumentsResult() if not representation.deductive and not representation.explicit: logger.debug("No observations to save") - return new_documents + return empty_result all_observations = [ _normalized_observation(obs) @@ -91,7 +94,7 @@ class RepresentationManager: ] if not all_observations: logger.debug("No non-empty observations to save") - return new_documents + return empty_result # Batch embed all observations batch_embed_start = time.perf_counter() @@ -123,7 +126,7 @@ class RepresentationManager: # Batch create document objects create_document_start = time.perf_counter() async with tracked_db("representation_manager.save_representation") as db: - new_documents = await self._save_representation_internal( + new_documents_result = await self._save_representation_internal( db, all_observations, embeddings, @@ -141,7 +144,7 @@ class RepresentationManager: "ms", ) - return new_documents + return new_documents_result async def _save_representation_internal( self, @@ -152,7 +155,7 @@ class RepresentationManager: session_name: str, message_created_at: datetime.datetime, message_level_configuration: ResolvedConfiguration, - ) -> int: + ) -> crud.CreateDocumentsResult: # get_or_create_collection already handles IntegrityError with rollback and a retry collection = await crud.get_or_create_collection( db, @@ -191,7 +194,7 @@ class RepresentationManager: ) # Use bulk creation with optional duplicate detection - accepted_documents = await crud.create_documents( + accepted_documents_result = await crud.create_documents( db, documents_to_create, self.workspace_name, @@ -206,13 +209,13 @@ class RepresentationManager: except Exception as e: logger.warning(f"Failed to check dream scheduling: {e}") - return len(accepted_documents) + return accepted_documents_result async def get_working_representation( self, *, db: AsyncSession | None = None, - session_name: str | None = None, + session_allowlist: list[str] | None = None, include_semantic_query: str | None = None, embedding: list[float] | None = None, semantic_search_top_k: int | None = None, @@ -228,7 +231,10 @@ class RepresentationManager: Args: db: Optional database session. If provided, uses it directly; otherwise creates a new session via tracked_db. - session_name: Optional session to filter by + session_allowlist: Optional session allowlist to filter by. Applied + uniformly to every query path (semantic, most-derived, and + recent). None means no session restriction; an empty list + fail-closes to an empty representation. include_semantic_query: Query for semantic search embedding: Pre-computed embedding for the semantic query. semantic_search_top_k: Number of semantic results @@ -266,7 +272,7 @@ class RepresentationManager: if db is not None: return await self._get_working_representation_internal( db, - session_name=session_name, + session_allowlist=session_allowlist, include_semantic_query=include_semantic_query, embedding=embedding, semantic_search_top_k=semantic_search_top_k, @@ -276,11 +282,11 @@ class RepresentationManager: ) async with tracked_db( - "representation_manager.get_working_representation" + "representation_manager.get_working_representation", read_only=True ) as new_db: return await self._get_working_representation_internal( new_db, - session_name=session_name, + session_allowlist=session_allowlist, include_semantic_query=include_semantic_query, embedding=embedding, semantic_search_top_k=semantic_search_top_k, @@ -295,7 +301,7 @@ class RepresentationManager: self, db: AsyncSession, *, - session_name: str | None = None, + session_allowlist: list[str] | None = None, include_semantic_query: str | None = None, embedding: list[float] | None = None, semantic_search_top_k: int | None = None, @@ -304,6 +310,12 @@ class RepresentationManager: max_observations: int = settings.DERIVER.WORKING_REPRESENTATION_MAX_OBSERVATIONS, ) -> Representation: """Internal implementation of get_working_representation.""" + # Fail closed on an empty allowlist. This must short-circuit before + # any query: downstream stores drop an `IN ()` clause with an empty + # list (lancedb), which would silently widen the scope instead. + if session_allowlist is not None and not session_allowlist: + return Representation() + total = max_observations # Calculate how many observations to get from each source @@ -344,6 +356,7 @@ class RepresentationManager: top_k=semantic_observations, max_distance=semantic_search_max_distance, embedding=embedding, + session_allowlist=session_allowlist, ) representation.merge_representation( Representation.from_documents(semantic_docs) @@ -352,7 +365,7 @@ class RepresentationManager: # Get most derived observations if requested if include_most_derived: derived_docs = await self._query_documents_most_derived( - db, top_k=top_observations + db, top_k=top_observations, session_allowlist=session_allowlist ) representation.merge_representation( Representation.from_documents(derived_docs) @@ -360,7 +373,7 @@ class RepresentationManager: # Get recent observations recent_docs = await self._query_documents_recent( - db, top_k=recent_observations, session_name=session_name + db, top_k=recent_observations, session_allowlist=session_allowlist ) representation.merge_representation(Representation.from_documents(recent_docs)) @@ -375,6 +388,7 @@ class RepresentationManager: max_distance: float | None = None, level: str | None = None, embedding: list[float] | None = None, + session_allowlist: list[str] | None = None, ) -> list[models.Document]: """Query documents by semantic similarity.""" try: @@ -386,6 +400,7 @@ class RepresentationManager: top_k, max_distance, embedding=embedding, + session_allowlist=session_allowlist, ) else: documents = await crud.query_documents( @@ -397,6 +412,10 @@ class RepresentationManager: max_distance=max_distance, top_k=top_k, embedding=embedding, + filters=self._build_filter_conditions( + session_allowlist=session_allowlist + ) + or None, ) db.expunge_all() return list(documents) @@ -406,7 +425,7 @@ class RepresentationManager: return [] async def _query_documents_recent( - self, db: AsyncSession, top_k: int, session_name: str | None = None + self, db: AsyncSession, top_k: int, session_allowlist: list[str] | None = None ) -> list[models.Document]: """Query most recent documents.""" stmt = ( @@ -418,8 +437,13 @@ class RepresentationManager: models.Document.observed == self.observed, models.Document.deleted_at.is_(None), *( - [models.Document.session_name == session_name] - if session_name is not None + [ + models.Document.session_name.in_(session_allowlist), + # Only levels with a trustworthy session stamp are + # scopeable — see ALLOWLIST_SAFE_LEVELS. + models.Document.level.in_(ALLOWLIST_SAFE_LEVELS), + ] + if session_allowlist is not None else [] ), ) @@ -432,7 +456,7 @@ class RepresentationManager: return list(documents) async def _query_documents_most_derived( - self, db: AsyncSession, top_k: int + self, db: AsyncSession, top_k: int, session_allowlist: list[str] | None = None ) -> list[models.Document]: """Query most derived documents.""" stmt = ( @@ -443,8 +467,24 @@ class RepresentationManager: models.Document.observer == self.observer, models.Document.observed == self.observed, models.Document.deleted_at.is_(None), + *( + [ + models.Document.session_name.in_(session_allowlist), + # Only levels with a trustworthy session stamp are + # scopeable — see ALLOWLIST_SAFE_LEVELS. + models.Document.level.in_(ALLOWLIST_SAFE_LEVELS), + ] + if session_allowlist is not None + else [] + ), + ) + .order_by( + models.Document.times_derived.desc(), + models.Document.created_at.desc(), + # created_at is the transaction timestamp, so documents created + # in the same batch share it -- id keeps the order deterministic. + models.Document.id, ) - .order_by(models.Document.times_derived.desc()) ) result = await db.execute(stmt) @@ -473,6 +513,7 @@ class RepresentationManager: count: int, max_distance: float | None = None, embedding: list[float] | None = None, + session_allowlist: list[str] | None = None, ) -> list[models.Document]: """Query documents for a specific level.""" documents = await crud.query_documents( @@ -483,7 +524,9 @@ class RepresentationManager: query=query, max_distance=max_distance, top_k=count, - filters=self._build_filter_conditions(level), + filters=self._build_filter_conditions( + level, session_allowlist=session_allowlist + ), embedding=embedding, ) @@ -496,17 +539,32 @@ class RepresentationManager: def _build_filter_conditions( self, level: str | None = None, + session_allowlist: list[str] | None = None, ) -> dict[str, Any]: """ Build filter conditions for document queries. Returns a flat dict of key-value pairs for vector store filtering. + Callers must not pass an empty session_allowlist list — empty allowlists + fail closed before any query is issued (see + _get_working_representation_internal). """ filters: dict[str, Any] = {} if level: filters["level"] = level + # `is not None` (not truthiness): an explicit empty allowlist must emit + # an empty `in` so downstream stores fail closed, matching + # _query_documents_recent / _query_documents_most_derived. Truthiness + # here would silently drop the filter and widen scope. + if session_allowlist is not None: + filters["session_name"] = {"in": session_allowlist} + # Only levels with a trustworthy session stamp are scopeable. This + # overrides any narrower `level` above; an empty intersection emits + # `{"in": []}`, which matches nothing rather than everything. + filters["level"] = {"in": allowlist_safe_levels([level] if level else None)} + return filters @@ -519,7 +577,7 @@ async def get_working_representation( db: AsyncSession | None = None, observer: str, observed: str, - session_name: str | None = None, + session_allowlist: list[str] | None = None, include_semantic_query: str | None = None, embedding: list[float] | None = None, semantic_search_top_k: int | None = None, @@ -552,7 +610,7 @@ async def get_working_representation( ) return await manager.get_working_representation( db=db, - session_name=session_name, + session_allowlist=session_allowlist, include_semantic_query=include_semantic_query, embedding=embedding, semantic_search_top_k=semantic_search_top_k, diff --git a/src/crud/session.py b/src/crud/session.py index 712188bc..4310cb3a 100644 --- a/src/crud/session.py +++ b/src/crud/session.py @@ -72,6 +72,7 @@ def session_cache_key(workspace_name: str, session_name: str) -> str: key=SESSION_CACHE_KEY_TEMPLATE, ttl=f"{settings.CACHE.DEFAULT_LOCK_TTL_SECONDS}s", prefix=SESSION_LOCK_PREFIX, + check_interval=settings.CACHE.LOCK_WAIT_CHECK_INTERVAL_SECONDS, ) async def _fetch_session( db: AsyncSession, @@ -834,6 +835,38 @@ async def get_peers_from_session( ) +async def is_peer_in_session( + db: AsyncSession, + workspace_name: str, + session_name: str, + peer_name: str, +) -> bool: + """Return whether a peer is an active member of a session. + + Active membership means a `SessionPeer` row exists with `left_at IS NULL`. + Used by the auth layer to grant a peer-scoped key read access to the + sessions that peer belongs to. + + Args: + db: Database session + workspace_name: Name of the workspace + session_name: Name of the session + peer_name: Name of the peer + + Returns: + True if the peer is currently a member of the session. + """ + result = await db.scalar( + select(models.SessionPeer.peer_name) + .where(models.SessionPeer.workspace_name == workspace_name) + .where(models.SessionPeer.session_name == session_name) + .where(models.SessionPeer.peer_name == peer_name) + .where(models.SessionPeer.left_at.is_(None)) + .limit(1) + ) + return result is not None + + async def get_session_peer_configuration( workspace_name: str, session_name: str, diff --git a/src/crud/workspace.py b/src/crud/workspace.py index 52662c79..c9040f55 100644 --- a/src/crud/workspace.py +++ b/src/crud/workspace.py @@ -60,6 +60,7 @@ def workspace_cache_key(workspace_name: str) -> str: key=WORKSPACE_CACHE_KEY_TEMPLATE, ttl=f"{settings.CACHE.DEFAULT_LOCK_TTL_SECONDS}s", prefix=WORKSPACE_LOCK_PREFIX, + check_interval=settings.CACHE.LOCK_WAIT_CHECK_INTERVAL_SECONDS, ) async def _fetch_workspace( db: AsyncSession, workspace_name: str diff --git a/src/db.py b/src/db.py index 4b656175..005d52eb 100644 --- a/src/db.py +++ b/src/db.py @@ -1,13 +1,27 @@ import contextvars +import logging +from typing import Any -from sqlalchemy import MetaData, text -from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine +from sqlalchemy import MetaData, event, text +from sqlalchemy.ext.asyncio import ( + AsyncSession, + async_sessionmaker, + create_async_engine, +) from sqlalchemy.orm import declarative_base -from sqlalchemy.pool import NullPool +from sqlalchemy.pool import NullPool, QueuePool from src.config import settings +from src.telemetry.prometheus.metrics import db_queries_in_flight_gauge -connect_args = {"prepare_threshold": None} +logger = logging.getLogger(__name__) + +connect_args = { + "prepare_threshold": None, + # Bound a single connection attempt so it fails fast instead of hanging when + # the server/pooler is unreachable or stalled (psycopg, seconds). + "connect_timeout": settings.DB.CONNECT_TIMEOUT_SECONDS, +} # Context variable to store request context request_context: contextvars.ContextVar[str | None] = contextvars.ContextVar( @@ -38,13 +52,174 @@ engine = create_async_engine( **engine_kwargs, ) +# A vanilla AsyncSession is lazy: it checks out a pooled connection on the first +# DB-touching call (not at construction) and couples the checkout to the +# statement, so a handler doing non-DB work (embedding/file/LLM) before its +# first query does not pin a connection across it. Connection acquisition is a +# single attempt with no retry — callers handle a saturated/unreachable DB (the +# API surfaces the error; the deriver backs off and retries on a later poll). SessionLocal = async_sessionmaker( autocommit=False, autoflush=False, expire_on_commit=False, bind=engine, + class_=AsyncSession, ) +# Read-only engine: shares `engine`'s pool, but checks connections out in DBAPI +# AUTOCOMMIT mode, so psycopg emits NO BEGIN — a SELECT never autobegins a +# transaction. The backend therefore returns to state 'idle' (not 'idle in +# transaction') the moment a statement completes. +read_engine = engine.execution_options(isolation_level="AUTOCOMMIT") + +# Sessions for SELECT-only work (same lazy-checkout semantics as SessionLocal). +# MUST NOT be used for writes: with no enclosing transaction, begin_nested() +# savepoints (see the crud get-or-create paths) break, and every flush would +# commit immediately. Use SessionLocal for anything that mutates. +ReadSessionLocal = async_sessionmaker( + autocommit=False, + autoflush=False, + expire_on_commit=False, + bind=read_engine, + class_=AsyncSession, +) + + +def _set_application_name_on_checkout( + dbapi_connection: Any, _connection_record: Any, _connection_proxy: Any +) -> None: + """Tag each checked-out connection with the current request context. + + Registered only when ``DB.TRACING`` is on. Fires on every pool checkout (so a + reused pooled connection is re-tagged for the new caller), reading the + per-task ``request_context`` the request/task scope has already set. + Best-effort: a failure here must never break the checkout. + + Runs in autocommit so it never leaves the connection 'idle in transaction' + at checkout: this hook fires BEFORE the dialect applies execution-option + isolation levels, and psycopg refuses to switch a connection into AUTOCOMMIT + (which the read engine does) while a transaction opened by this statement is + still in progress. set_config(..., is_local=false) is session-scoped, so it + persists past the autocommit boundary. + """ + context = request_context.get() or "unknown" + try: + previous_autocommit = dbapi_connection.autocommit + if not previous_autocommit: + dbapi_connection.autocommit = True + try: + cursor = dbapi_connection.cursor() + try: + cursor.execute( + "SELECT set_config('application_name', %s, false)", (context,) + ) + finally: + cursor.close() + finally: + if not previous_autocommit: + dbapi_connection.autocommit = False + except Exception: + logger.debug("setting application_name on checkout failed", exc_info=True) + + +if settings.DB.TRACING: + event.listen(engine.sync_engine, "checkout", _set_application_name_on_checkout) + + +def get_pool_stats() -> dict[str, int]: + """Return live connection-pool stats for this process. + + ``engine.pool`` is the AsyncEngine's pool (the same object as + ``engine.sync_engine.pool``); its stat methods are synchronous counter + reads with no I/O, so they are safe to call without ``await``. Returns + zeros for pools that do not track connections (e.g. ``NullPool``). + """ + zeros = {"checked_out": 0, "checked_in": 0, "size": 0, "overflow": 0} + pool = engine.pool + # Only QueuePool (and its AsyncAdaptedQueuePool subclass) tracks connection + # counts; NullPool and others have no meaningful stats. + if not isinstance(pool, QueuePool): + return zeros + try: + # overflow() is negative until the base pool fills (it starts at + # -pool_size); clamp to the count of overflow connections actually open. + return { + "checked_out": pool.checkedout(), + "checked_in": pool.checkedin(), + "size": pool.size(), + "overflow": max(0, pool.overflow()), + } + except Exception: + return zeros + + +class DBQueryInflightTracker: + """Tracks statements executing on the wire via SQLAlchemy cursor events. + + Drift-proof: marks ``Connection.info`` when a statement starts and clears it + on completion OR error, so the gauge can't leak upward (an errored statement + skips ``after_cursor_execute``) or go negative (a connect-time error has no + matching start). Bound to a pre-resolved labeled gauge child so the + per-statement hot path does no label resolution. + """ + + # Marker on Connection.info recording that we incremented for the current + # statement, so we decrement exactly once on completion or error. + INFLIGHT_KEY: str = "_honcho_inflight" + + def __init__(self, gauge_child: Any) -> None: + self._child: Any = gauge_child + + def on_before(self, conn: Any, *_: Any) -> None: + try: + conn.info[self.INFLIGHT_KEY] = True + self._child.inc() + except Exception: + logger.debug("in-flight gauge inc failed", exc_info=True) + + def on_after(self, conn: Any, *_: Any) -> None: + try: + if conn.info.pop(self.INFLIGHT_KEY, False): + self._child.dec() + except Exception: + logger.debug("in-flight gauge dec failed", exc_info=True) + + def on_error(self, exception_context: Any) -> None: + try: + conn = exception_context.connection + if conn is not None and conn.info.pop(self.INFLIGHT_KEY, False): + self._child.dec() + except Exception: + logger.debug("in-flight gauge error-path dec failed", exc_info=True) + + +# Process-wide tracker, created at registration (None until then / if metrics off). +_inflight_tracker: DBQueryInflightTracker | None = None + + +_db_query_instrumentation_registered = False + + +def register_db_query_instrumentation(instance_type: str) -> None: + """Attach per-statement in-flight tracking to the engine (no-op if off). + + Gated on METRICS.ENABLED so there is zero overhead — not even attached event + listeners — when metrics are disabled. Idempotent: repeated calls (e.g. a + re-run lifespan or test startup) won't attach duplicate listeners, which + would double-count in-flight statements. + """ + global _inflight_tracker, _db_query_instrumentation_registered + if not settings.METRICS.ENABLED or _db_query_instrumentation_registered: + return + child = db_queries_in_flight_gauge.labels(instance_type=instance_type) + _inflight_tracker = DBQueryInflightTracker(child) + sync_engine = engine.sync_engine + event.listen(sync_engine, "before_cursor_execute", _inflight_tracker.on_before) + event.listen(sync_engine, "after_cursor_execute", _inflight_tracker.on_after) + event.listen(sync_engine, "handle_error", _inflight_tracker.on_error) + _db_query_instrumentation_registered = True + + # Define your naming convention convention = { "ix": "ix_%(table_name)s_%(column_0_N_name)s", # Index - supports multi-column diff --git a/src/dependencies.py b/src/dependencies.py index 66b1931e..060186b4 100644 --- a/src/dependencies.py +++ b/src/dependencies.py @@ -2,25 +2,21 @@ import uuid from contextlib import asynccontextmanager from fastapi import Depends -from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession -from src.config import settings -from src.db import SessionLocal, request_context +from src.db import ReadSessionLocal, SessionLocal, request_context async def get_db(): - """FastAPI Dependency Generator for Database""" - - context = request_context.get() or "unknown" + """FastAPI Dependency Generator for Database. + The session is lazy: it does NOT check out a pooled connection here. The + AsyncSession checks one out on the first DB-touching call, so a handler doing + non-DB work (embedding/file/LLM) before its first query does not pin a + connection across it. + """ db: AsyncSession = SessionLocal() try: - if settings.DB.TRACING: - await db.execute( - text("SELECT set_config('application_name', :name, false)"), - {"name": context}, - ) yield db except Exception: await db.rollback() @@ -30,14 +26,46 @@ async def get_db(): # is closed before the TCP connection drops. Supavisor v2 does NOT # clean up orphaned transactions on client disconnect in transaction- # pooling mode, so relying on `in_transaction()` (Python-side state) - # can leave the backend pinned with an open BEGIN. + # can leave the backend pinned with an open BEGIN. (Cheap no-op if the + # lazy session never checked out a connection.) + await db.rollback() + await db.close() + + +async def get_read_db(): + """FastAPI Dependency Generator for SELECT-only handlers. + + Same lazy-checkout semantics as get_db, but the session is bound to the + AUTOCOMMIT read engine: no BEGIN is ever emitted, so the connection can not + sit 'idle in transaction' between the query and this teardown — a delayed + finally here is harmless (the backend is plain 'idle'). close() is still + required to release the connection itself back to the pool. + + MUST only be used by handlers that never mutate; see ReadSessionLocal. + """ + db: AsyncSession = ReadSessionLocal() + try: + yield db + finally: + # rollback is a wire-level no-op under AUTOCOMMIT; kept to reset any + # Python-side session state before close, mirroring get_db. await db.rollback() await db.close() @asynccontextmanager -async def tracked_db(operation_name: str | None = None): - """Context manager for tracked database sessions""" +async def tracked_db(operation_name: str | None = None, *, read_only: bool = False): + """Context manager for tracked database sessions. + + Sets a task-scoped request_context so the lazy session picks it up for + tracing/attribution, then yields a lazy session (see get_db). + + Pass read_only=True for SELECT-only windows: the session is then bound to + the AUTOCOMMIT read engine, so the work inside the block never holds an + open transaction (no idle-in-transaction parking; the pooler can reclaim + the backend between statements). Never use read_only=True on a path that + mutates — see ReadSessionLocal. + """ # Get request ID if available, or create operation-specific one context = request_context.get() token = None @@ -46,22 +74,15 @@ async def tracked_db(operation_name: str | None = None): context = f"task:{operation_name}:{str(uuid.uuid4())[:8]}" token = request_context.set(context) - # Create session with tracking info - db = SessionLocal() - + db = (ReadSessionLocal if read_only else SessionLocal)() try: - if settings.DB.TRACING: - await db.execute( - text("SELECT set_config('application_name', :name, false)"), - {"name": context or f"task:{operation_name}"}, - ) - yield db except Exception: await db.rollback() raise finally: - # Always send ROLLBACK unconditionally — see get_db() comment. + # Always send ROLLBACK unconditionally — see get_db() comment. (Under + # read_only/AUTOCOMMIT it is a wire-level no-op.) await db.rollback() await db.close() if token: # Only reset if we set it @@ -69,3 +90,4 @@ async def tracked_db(operation_name: str | None = None): db: AsyncSession = Depends(get_db) +read_db: AsyncSession = Depends(get_read_db) diff --git a/src/deriver/__main__.py b/src/deriver/__main__.py index 20ffbb66..c56ed6a0 100644 --- a/src/deriver/__main__.py +++ b/src/deriver/__main__.py @@ -6,9 +6,13 @@ import uvloop from prometheus_client import start_http_server from src.config import settings -from src.db import engine +from src.db import engine, register_db_query_instrumentation from src.startup import validate_embedding_schema -from src.telemetry import initialize_telemetry_async, shutdown_telemetry +from src.telemetry import ( + initialize_telemetry_async, + register_db_pool_collector, + shutdown_telemetry, +) from .queue_manager import main @@ -18,6 +22,9 @@ logger = logging.getLogger(__name__) def start_metrics_server() -> None: """Start the Prometheus metrics HTTP server on port 9090.""" start_http_server(9090) + # Expose DB connection-pool stats for this deriver instance. + register_db_pool_collector("deriver") + register_db_query_instrumentation("deriver") logger.info("Prometheus metrics server started on port 9090") diff --git a/src/deriver/consumer.py b/src/deriver/consumer.py index 8bbd9359..708d05f6 100644 --- a/src/deriver/consumer.py +++ b/src/deriver/consumer.py @@ -43,17 +43,20 @@ async def process_item(queue_item: models.QueueItem) -> None: # Handle reconciler first - it's the only task type that doesn't require workspace_name if task_type == "reconciler": - with sentry_sdk.start_transaction(name="process_reconciler_task", op="deriver"): - try: - validated = ReconcilerPayload(**queue_payload) - except ValidationError as e: - logger.error( - "Invalid reconciler payload received: %s. Payload: %s", - str(e), - queue_payload, - ) - raise ValueError(f"Invalid payload structure: {str(e)}") from e - await process_reconciler(validated) + # No top-level transaction here: reconciler tasks poll on a fixed + # interval and usually find no work. Tracing is started per-batch + # inside the reconciler only when actual work is found, so idle + # cycles don't consume Sentry tracing/profiling quota. + try: + validated = ReconcilerPayload(**queue_payload) + except ValidationError as e: + logger.error( + "Invalid reconciler payload received: %s. Payload: %s", + str(e), + queue_payload, + ) + raise ValueError(f"Invalid payload structure: {str(e)}") from e + await process_reconciler(validated) return # All other task types require a workspace_name @@ -173,7 +176,7 @@ async def process_representation_batch( queue_item_message_ids: Message IDs from queue items hit_batch_token_cap: whether the queue batcher clamped this batch to fit was_flush_enabled: snapshot of DERIVER.FLUSH_ENABLED at fetch time - batch_max_tokens: DERIVER.REPRESENTATION_BATCH_MAX_TOKENS snapshot + batch_max_tokens: DERIVER.REPRESENTATION_BATCH_TARGET_INPUT_TOKENS snapshot """ if not messages or not messages[0]: logger.debug("process_representation_batch received no messages") diff --git a/src/deriver/deriver.py b/src/deriver/deriver.py index 303a6ec2..2ad7d14f 100644 --- a/src/deriver/deriver.py +++ b/src/deriver/deriver.py @@ -1,6 +1,8 @@ import logging import time +from nanoid import generate as generate_nanoid + from src import crud from src.config import ConfiguredModelSettings, settings from src.crud.representation import RepresentationManager @@ -56,7 +58,7 @@ async def process_representation_tasks_batch( queue_item_message_ids: Message IDs from queue items being processed hit_batch_token_cap: queue batcher clamped this batch to fit was_flush_enabled: DERIVER.FLUSH_ENABLED snapshot at batch time - batch_max_tokens: DERIVER.REPRESENTATION_BATCH_MAX_TOKENS snapshot + batch_max_tokens: DERIVER.REPRESENTATION_BATCH_TARGET_INPUT_TOKENS snapshot """ if not messages: return @@ -142,12 +144,12 @@ async def process_representation_tasks_batch( model_config = base_model_config # Single LLM call + trace_id = generate_nanoid() llm_start = time.perf_counter() response = await honcho_llm_call( model_config=model_config, prompt=prompt, max_tokens=max_tokens, - track_name="Minimal Deriver", response_model=PromptRepresentation, json_mode=True, max_input_tokens=settings.DERIVER.MAX_INPUT_TOKENS, @@ -159,6 +161,9 @@ async def process_representation_tasks_batch( call_purpose=CallPurpose.DERIVER_REPRESENTATION.value, parent_category="representation", observed=observed, + track_name="Minimal Deriver", + trace_id=trace_id, + span_id=trace_id, ), ) llm_duration = (time.perf_counter() - llm_start) * 1000 @@ -189,6 +194,7 @@ async def process_representation_tasks_batch( latest_message.created_at, ) + agg_representation_result = crud.CreateDocumentsResult() successful_observer_count = 0 if observations.is_empty() or not message_ids: logger.warning( @@ -208,12 +214,26 @@ async def process_representation_tasks_batch( ) try: - await representation_manager.save_representation( - observations, - message_ids, - latest_message.session_name, - latest_message.created_at, - message_level_configuration, + representation_result = ( + await representation_manager.save_representation( + observations, + message_ids, + latest_message.session_name, + latest_message.created_at, + message_level_configuration, + ) + ) + agg_representation_result.exact_dup_existing_count += ( + representation_result.exact_dup_existing_count + ) + agg_representation_result.exact_dup_in_batch_count += ( + representation_result.exact_dup_in_batch_count + ) + agg_representation_result.semantic_dup_rejected_count += ( + representation_result.semantic_dup_rejected_count + ) + agg_representation_result.semantic_dup_replaced_count += ( + representation_result.semantic_dup_replaced_count ) successful_observer_count += 1 except Exception as e: @@ -313,5 +333,9 @@ async def process_representation_tasks_batch( hit_batch_token_cap=hit_batch_token_cap, hit_input_token_cap=response.hit_input_token_cap, observer_count=successful_observer_count, + exact_dup_existing_count=agg_representation_result.exact_dup_existing_count, + exact_dup_in_batch_count=agg_representation_result.exact_dup_in_batch_count, + semantic_dup_rejected_count=agg_representation_result.semantic_dup_rejected_count, + semantic_dup_replaced_count=agg_representation_result.semantic_dup_replaced_count, ) ) diff --git a/src/deriver/enqueue.py b/src/deriver/enqueue.py index aaeda415..7cd42aa3 100644 --- a/src/deriver/enqueue.py +++ b/src/deriver/enqueue.py @@ -403,6 +403,7 @@ def create_dream_record( delay_reason: str | None = None, documents_since_last_dream_at_schedule: int | None = None, document_threshold: int | None = None, + rebuild: bool = False, ) -> dict[str, Any]: """ Create a queue record for a dream task. @@ -417,6 +418,7 @@ def create_dream_record( delay_reason: what governed when it fires documents_since_last_dream_at_schedule: count snapshot at schedule time document_threshold: DOCUMENT_THRESHOLD snapshot at schedule time + rebuild: card_refresh only — rebuild the card without the prior card Returns: Queue record dictionary with workspace_name and other fields @@ -430,6 +432,7 @@ def create_dream_record( delay_reason=delay_reason, documents_since_last_dream_at_schedule=documents_since_last_dream_at_schedule, document_threshold=document_threshold, + rebuild=rebuild, ) return { @@ -452,6 +455,7 @@ async def enqueue_dream( delay_reason: str | None = None, documents_since_last_dream_at_schedule: int | None = None, document_threshold: int | None = None, + rebuild: bool = False, ) -> None: """ Enqueue a dream task for immediate processing by the deriver. @@ -461,6 +465,8 @@ async def enqueue_dream( Deduplication: If a dream with the same work_unit_key is already in-progress (has an ActiveQueueSession) or pending in the queue, the enqueue is skipped. + The work unit key includes the dream type, so e.g. a card_refresh dream + never collides with a pending omni dream for the same collection. Args: workspace_name: Name of the workspace @@ -468,6 +474,7 @@ async def enqueue_dream( observed: Name of the observed peer dream_type: Type of dream to execute session_name: Name of the session to scope the dream to if specified + rebuild: card_refresh only — rebuild the card without the prior card """ async with tracked_db("dream_enqueue") as db_session: try: @@ -481,6 +488,7 @@ async def enqueue_dream( delay_reason=delay_reason, documents_since_last_dream_at_schedule=documents_since_last_dream_at_schedule, document_threshold=document_threshold, + rebuild=rebuild, ) work_unit_key = dream_record["work_unit_key"] @@ -495,7 +503,7 @@ async def enqueue_dream( is_in_progress = await db_session.scalar(in_progress_check) if is_in_progress: - logger.info( + logger.debug( "Skipping dream enqueue - already in progress: %s/%s/%s (type: %s)", workspace_name, observer, @@ -515,7 +523,7 @@ async def enqueue_dream( is_pending = await db_session.scalar(pending_check) if is_pending: - logger.info( + logger.debug( "Dream already pending in queue: %s/%s/%s (type: %s)", workspace_name, observer, diff --git a/src/deriver/prompts.py b/src/deriver/prompts.py index 683834bf..402bcae7 100644 --- a/src/deriver/prompts.py +++ b/src/deriver/prompts.py @@ -31,6 +31,7 @@ def _custom_instructions_section(custom_instructions: str | None) -> str: return c( f""" CUSTOM INSTRUCTIONS: + These instructions apply to the target peer identified below. {normalized_custom_instructions} """ ) @@ -54,26 +55,32 @@ def minimal_deriver_prompt( custom_instructions_section = _custom_instructions_section(custom_instructions) return c( f""" -Analyze messages from {peer_id} to extract **explicit atomic facts** about them. +Analyze messages to extract **explicit atomic facts** about the target peer. -[EXPLICIT] DEFINITION: Facts about {peer_id} that can be derived directly from their messages. +[EXPLICIT] DEFINITION: Facts about the target peer that can be derived directly from their messages. - Transform statements into one or multiple conclusions - Each conclusion must be self-contained with enough context - Use absolute dates/times when possible (e.g. "June 26, 2025" not "yesterday") RULES: -- Properly attribute observations to the correct subject: if it is about {peer_id}, say so. If {peer_id} is referencing someone or something else, make that clear. -- Observations should make sense on their own. Each observation will be used in the future to better understand {peer_id}. -- Extract ALL observations from {peer_id} messages, using others as context. +- The target peer is the peer identified below under `Target peer:`. +- A peer can be a human user, AI agent, bot, service, or other actor. +- Use the exact peer id from `Target peer:` in final observations, not the phrase "the target peer". +- Properly attribute observations to the correct subject: if it is about the target peer, use the exact peer id as the subject. If the target peer is referencing someone or something else, make that clear. +- Observations should make sense on their own. Each observation will be used in the future to better understand the target peer. +- Extract ALL observations from the target peer's messages, using others as context. - Contextualize each observation sufficiently (e.g. "Ann is nervous about the job interview at the pharmacy" not just "Ann is nervous") -EXAMPLES: -- EXPLICIT: "I just had my 25th birthday last Saturday" → "{peer_id} is 25 years old", "{peer_id}'s birthday is June 21st" -- EXPLICIT: "I took my dog for a walk in NYC" → "{peer_id} has a dog", "{peer_id} lives in NYC" -- EXPLICIT: "{peer_id} attended college" + general knowledge → "{peer_id} completed high school or equivalent" +EXAMPLES (using `alice` as the target peer id): +- EXPLICIT: "I just had my 25th birthday last Saturday" → "alice is 25 years old", "alice's birthday is June 21st" +- EXPLICIT: "I took my dog for a walk in NYC" → "alice has a dog", "alice lives in NYC" +- EXPLICIT: "alice attended college" + general knowledge → "alice completed high school or equivalent" {custom_instructions_section} +Target peer: +{peer_id} + Messages to analyze: {messages} diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index 8c2b5850..498372ed 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -1,5 +1,8 @@ import asyncio +import contextlib +import random import signal +import time from asyncio import Task from collections.abc import Sequence from dataclasses import dataclass, field @@ -11,6 +14,7 @@ import sentry_sdk from dotenv import load_dotenv from nanoid import generate as generate_nanoid from sentry_sdk.integrations.asyncio import AsyncioIntegration +from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration from sqlalchemy import and_, delete, or_, select, update from sqlalchemy.dialects.postgresql import insert from sqlalchemy.engine import CursorResult @@ -125,6 +129,22 @@ class QueueManager: self.worker_ownership: dict[str, WorkerOwnership] = {} self.queue_empty_flag: asyncio.Event = asyncio.Event() + # Current adaptive polling interval; grows while idle/erroring and + # resets to the base interval as soon as work is claimed. + self._current_poll_interval: float = ( + settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS + ) + + # Monotonic timestamp of the last stale-work-unit cleanup ATTEMPT. + # None -> the first poll always runs cleanup (recovers rows left stale + # by a crashed predecessor immediately). + self._last_stale_cleanup_attempt: float | None = None + # Jittered gate width (seconds) sampled ONCE per attempt, so the deadline + # for the next run is fixed when the timestamp is set rather than + # re-rolled on every poll (which would make the effective spacing a + # random walk and untestable at non-zero jitter ratios). + self._stale_cleanup_gate_seconds: float = 0.0 + # Initialize from settings self.workers: int = settings.DERIVER.WORKERS self.semaphore: asyncio.Semaphore = asyncio.Semaphore(self.workers) @@ -147,7 +167,9 @@ class QueueManager: # Initialize Sentry if enabled, using settings if settings.SENTRY.ENABLED: - initialize_sentry(integrations=[AsyncioIntegration()]) + initialize_sentry( + integrations=[AsyncioIntegration(), SqlalchemyIntegration()] + ) def add_task(self, task: asyncio.Task[None]) -> None: """Track a new task""" @@ -196,6 +218,7 @@ class QueueManager: # Run the polling loop directly in this task logger.debug("Starting polling loop directly") try: + await self._sleep_startup_jitter() await self.polling_loop() finally: await self.cleanup() @@ -246,6 +269,35 @@ class QueueManager: # Polling and Scheduling # ########################## + async def _maybe_cleanup_stale_work_units(self) -> None: + """Run stale-work-unit cleanup at most once per (jittered) interval. + + Staleness is a minutes-timescale condition (STALE_SESSION_TIMEOUT_MINUTES), + but the polling loop fires on a seconds timescale on every deriver + instance — running cleanup unconditionally per poll multiplies into + unnecessary write transactions. Gate it locally: + concurrent cleaners on other instances remain safe via FOR UPDATE SKIP + LOCKED, so no cross-instance coordination is required, and the jittered + gate (sampled once per attempt) keeps instances from re-synchronizing + their cleanup runs. The gate tracks the last ATTEMPT (set before + running), so a failing cleanup waits a full interval instead of retrying + every poll against a DB that is already struggling. An interval of 0 + preserves run-every-poll behavior. + """ + interval = settings.DERIVER.STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS + if ( + interval > 0.0 + and self._last_stale_cleanup_attempt is not None + and time.monotonic() - self._last_stale_cleanup_attempt + < self._stale_cleanup_gate_seconds + ): + return + # Record the attempt and fix the next deadline before running, so the + # gate width is stable for this cycle and a failing cleanup still waits. + self._last_stale_cleanup_attempt = time.monotonic() + self._stale_cleanup_gate_seconds = self._jitter(interval) + await self.cleanup_stale_work_units() + async def cleanup_stale_work_units(self) -> None: """Clean up stale work units""" async with tracked_db("cleanup_stale_work_units") as db: @@ -278,15 +330,19 @@ class QueueManager: async def get_and_claim_work_units(self) -> dict[str, str]: """ Get available work units that aren't being processed. - For representation tasks, only returns work units with accumulated tokens - >= REPRESENTATION_BATCH_MAX_TOKENS (forced batching), unless FLUSH_ENABLED is True. + For representation tasks, only returns work units whose accumulated + tokens reach REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS or whose + oldest pending item exceeds REPRESENTATION_BATCH_MAX_AGE_SECONDS, + unless FLUSH_ENABLED is True. Returns a dict mapping work_unit_key to aqs_id. """ limit: int = max(0, self.workers - self.get_total_owned_work_units()) if limit == 0: return {} - batch_max_tokens = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + work_unit_target_tokens = ( + settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS + ) async with tracked_db("get_available_work_units") as db: representation_prefix = "representation:" @@ -294,6 +350,7 @@ class QueueManager: select( models.QueueItem.work_unit_key, func.sum(models.Message.token_count).label("total_tokens"), + func.min(models.QueueItem.created_at).label("oldest_created_at"), ) .join( models.Message, @@ -306,15 +363,21 @@ class QueueManager: ) work_units_subq = ( - select(models.QueueItem.work_unit_key) + select( + models.QueueItem.work_unit_key, + func.min(models.QueueItem.created_at).label("oldest_created_at"), + ) .where(~models.QueueItem.processed) .group_by(models.QueueItem.work_unit_key) .subquery() ) query = ( - select(work_units_subq.c.work_unit_key) - .limit(limit) + select( + work_units_subq.c.work_unit_key, + token_stats_subq.c.total_tokens, + token_stats_subq.c.oldest_created_at, + ) .outerjoin( token_stats_subq, work_units_subq.c.work_unit_key == token_stats_subq.c.work_unit_key, @@ -327,22 +390,53 @@ class QueueManager: ) .exists() ) + .order_by( + work_units_subq.c.oldest_created_at.asc(), + work_units_subq.c.work_unit_key.asc(), + ) + .limit(limit) ) # Apply batch threshold filter (skip if FLUSH_ENABLED is True) - if not settings.DERIVER.FLUSH_ENABLED and batch_max_tokens > 0: + if not settings.DERIVER.FLUSH_ENABLED and work_unit_target_tokens > 0: + max_age_seconds = settings.DERIVER.REPRESENTATION_BATCH_MAX_AGE_SECONDS + threshold_clause = ( + func.coalesce(token_stats_subq.c.total_tokens, 0) + >= work_unit_target_tokens + ) + if max_age_seconds > 0: + threshold_clause = or_( + threshold_clause, + token_stats_subq.c.oldest_created_at + <= func.now() - timedelta(seconds=max_age_seconds), + ) query = query.where( or_( ~work_units_subq.c.work_unit_key.startswith( representation_prefix ), - func.coalesce(token_stats_subq.c.total_tokens, 0) - >= batch_max_tokens, + threshold_clause, ) ) result = await db.execute(query) - available_units = result.scalars().all() + available_rows = result.all() + available_units: list[str] = [] + for work_unit_key, total_tokens, oldest_created_at in available_rows: + available_units.append(work_unit_key) + if ( + not settings.DERIVER.FLUSH_ENABLED + and settings.DERIVER.REPRESENTATION_BATCH_MAX_AGE_SECONDS > 0 + and work_unit_key.startswith(representation_prefix) + and int(total_tokens or 0) < work_unit_target_tokens + ): + logger.info( + "age-flushing work unit %s (tokens=%s < %s, oldest=%s)", + work_unit_key, + total_tokens or 0, + work_unit_target_tokens, + oldest_created_at, + ) if not available_units: await db.commit() return {} @@ -378,27 +472,77 @@ class QueueManager: ) return claimed_mapping + def _reset_poll_interval(self) -> None: + """Snap the polling interval back to the base after finding work.""" + self._current_poll_interval = settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS + + def _jitter(self, seconds: float) -> float: + """Scatter a sleep by +/- POLLING_JITTER_RATIO to avoid lockstep polling. + + Returns a uniform-random value in [(1-ratio)*seconds, (1+ratio)*seconds]. + Only the returned sleep is scattered; the underlying backoff schedule is + left unchanged. A ratio of 0.0 returns ``seconds`` unchanged. + """ + ratio = settings.DERIVER.POLLING_JITTER_RATIO + if ratio <= 0.0: + return seconds + # Scheduling jitter, not security/crypto — stdlib random is appropriate. + return seconds * random.uniform(1.0 - ratio, 1.0 + ratio) # nosec B311 + + async def _sleep_startup_jitter(self) -> None: + """Sleep a random delay before the first poll so instances that start + together don't poll in lockstep. Interruptible by shutdown so a signal + during the delay exits promptly. No-op when the window is 0.0. + """ + window = settings.DERIVER.POLLING_STARTUP_JITTER_SECONDS + if window <= 0.0: + return + # Scheduling jitter, not security/crypto — stdlib random is appropriate. + delay = random.uniform(0.0, window) # nosec B311 + logger.debug(f"Startup poll jitter: sleeping {delay:.1f}s before first poll") + # Timeout (slept the full delay without a shutdown) is the normal path; + # an early return means shutdown fired and polling_loop will exit at once. + with contextlib.suppress(asyncio.TimeoutError): + await asyncio.wait_for(self.shutdown_event.wait(), timeout=delay) + + def _advance_poll_interval(self) -> float: + """Return the current idle/backoff sleep, then grow it toward the cap.""" + interval = self._current_poll_interval + if settings.DERIVER.POLLING_BACKOFF_ENABLED: + self._current_poll_interval = min( + self._current_poll_interval + * settings.DERIVER.POLLING_BACKOFF_MULTIPLIER, + settings.DERIVER.POLLING_SLEEP_MAX_INTERVAL_SECONDS, + ) + return self._jitter(interval) + async def polling_loop(self) -> None: """Main polling loop to find and process new work units""" logger.debug("Starting polling loop") try: while not self.shutdown_event.is_set(): if self.queue_empty_flag.is_set(): - # logger.debug("Queue empty flag set, waiting") - await asyncio.sleep(settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS) + # The empty-poll branch below already slept this cycle's + # interval; just clear the flag and re-query (no second + # sleep — that would double the effective idle interval). self.queue_empty_flag.clear() continue - # Check if we have capacity before querying + # Check if we have capacity before querying. There is work to do + # (workers are busy), so keep the base interval for fast pickup + # when capacity frees rather than backing off. if self.semaphore.locked(): # logger.debug("All workers busy, waiting") - await asyncio.sleep(settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS) + await asyncio.sleep( + self._jitter(settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS) + ) continue try: - await self.cleanup_stale_work_units() + await self._maybe_cleanup_stale_work_units() claimed_work_units = await self.get_and_claim_work_units() if claimed_work_units: + self._reset_poll_interval() for work_unit_key, aqs_id in claimed_work_units.items(): # Create a new task for processing this work unit if not self.shutdown_event.is_set(): @@ -414,15 +558,14 @@ class QueueManager: self.add_task(task) else: self.queue_empty_flag.set() - await asyncio.sleep( - settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS - ) + await asyncio.sleep(self._advance_poll_interval()) except Exception as e: logger.exception("Error in polling loop") if settings.SENTRY.ENABLED: sentry_sdk.capture_exception(e) - # Note: rollback is handled by tracked_db dependency - await asyncio.sleep(settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS) + # Note: rollback is handled by tracked_db dependency. + # Back off so a down/saturated DB isn't hammered every cycle. + await asyncio.sleep(self._advance_poll_interval()) finally: logger.info("Polling loop stopped") @@ -675,7 +818,7 @@ class QueueManager: f"{task_type} tasks are not supported for get_queue_item_batch" ) - batch_max_tokens = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + batch_max_tokens = settings.DERIVER.REPRESENTATION_BATCH_TARGET_INPUT_TOKENS was_flush_enabled = settings.DERIVER.FLUSH_ENABLED parsed_key = parse_work_unit_key(work_unit_key) messages_context: list[models.Message] = [] diff --git a/src/dialectic/chat.py b/src/dialectic/chat.py index e9659a21..d2f7e741 100644 --- a/src/dialectic/chat.py +++ b/src/dialectic/chat.py @@ -8,6 +8,8 @@ using the DialecticAgent. import logging from collections.abc import AsyncIterator +from pydantic import BaseModel + from src import crud, schemas from src.config import ReasoningLevel from src.dependencies import tracked_db @@ -24,6 +26,8 @@ async def agentic_chat( observer: str, observed: str, reasoning_level: ReasoningLevel = "low", + session_allowlist: list[str] | None = None, + response_model: type[BaseModel] | None = None, ) -> str: """ Answer a query about a peer using the agentic dialectic. @@ -35,12 +39,15 @@ async def agentic_chat( observer: The peer making the query observed: The peer being queried about reasoning_level: Level of reasoning to apply + session_allowlist: Optional session allowlist restricting all recall + response_model: Optional Pydantic model the answer must conform to. + When set, the returned string is JSON matching the model's schema. Returns: The synthesized answer string """ # Short-lived DB session for validation + config - async with tracked_db("dialectic.preflight") as db: + async with tracked_db("dialectic.preflight", read_only=True) as db: await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observer)) if observer != observed: await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observed)) @@ -50,6 +57,9 @@ async def agentic_chat( session = await crud.get_session( db, workspace_name=workspace_name, session_name=session_name ) + # Read the opaque Session.id while the instance is still bound; the ORM + # object detaches once this read-only session closes below. + session_id = session.id if session else None workspace = await crud.get_workspace(db, workspace_name=workspace_name) configuration = get_configuration(None, session, workspace) @@ -68,14 +78,16 @@ async def agentic_chat( agent = DialecticAgent( workspace_name=workspace_name, session_name=session_name, + session_id=session_id, observer=observer, observed=observed, observer_peer_card=observer_peer_card, observed_peer_card=observed_peer_card, reasoning_level=reasoning_level, + session_allowlist=session_allowlist, ) - return await agent.answer(query) + return await agent.answer(query, response_model=response_model) async def agentic_chat_stream( @@ -85,6 +97,8 @@ async def agentic_chat_stream( observer: str, observed: str, reasoning_level: ReasoningLevel = "low", + session_allowlist: list[str] | None = None, + response_model: type[BaseModel] | None = None, ) -> AsyncIterator[str]: """ Stream an answer to a query about a peer using the agentic dialectic. @@ -96,12 +110,16 @@ async def agentic_chat_stream( observer: The peer making the query observed: The peer being queried about reasoning_level: Level of reasoning to apply + session_allowlist: Optional session allowlist restricting all recall + response_model: Optional Pydantic model the answer must conform to. + When set, the streamed text accumulates to JSON matching the + model's schema. Yields: Chunks of the response text as they are generated """ # Short-lived DB session for validation + config - async with tracked_db("dialectic.preflight") as db: + async with tracked_db("dialectic.preflight", read_only=True) as db: await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observer)) if observer != observed: await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observed)) @@ -111,6 +129,9 @@ async def agentic_chat_stream( session = await crud.get_session( db, workspace_name=workspace_name, session_name=session_name ) + # Read the opaque Session.id while the instance is still bound; the ORM + # object detaches once this read-only session closes below. + session_id = session.id if session else None workspace = await crud.get_workspace(db, workspace_name=workspace_name) configuration = get_configuration(None, session, workspace) @@ -129,12 +150,14 @@ async def agentic_chat_stream( agent = DialecticAgent( workspace_name=workspace_name, session_name=session_name, + session_id=session_id, observer=observer, observed=observed, observer_peer_card=observer_peer_card, observed_peer_card=observed_peer_card, reasoning_level=reasoning_level, + session_allowlist=session_allowlist, ) - async for chunk in agent.answer_stream(query): + async for chunk in agent.answer_stream(query, response_model=response_model): yield chunk diff --git a/src/dialectic/core.py b/src/dialectic/core.py index 5a5f690b..9580b94c 100644 --- a/src/dialectic/core.py +++ b/src/dialectic/core.py @@ -11,6 +11,7 @@ from collections.abc import AsyncIterator, Callable from typing import Any, cast from nanoid import generate as generate_nanoid +from pydantic import BaseModel from src import crud from src.config import ConfiguredModelSettings, ReasoningLevel, settings @@ -68,6 +69,8 @@ class DialecticAgent: observed_peer_card: list[str] | None = None, metric_key: str | None = None, reasoning_level: ReasoningLevel = "low", + session_id: str | None = None, + session_allowlist: list[str] | None = None, ): """ Initialize the dialectic agent. @@ -81,9 +84,15 @@ class DialecticAgent: observed_peer_card: Biographical information about the observed peer metric_key: Optional key for logging metrics (if provided, agent won't log separately) reasoning_level: Level of reasoning to apply + session_id: ID used for grouping traces (not session_name) + session_allowlist: Optional session allowlist restricting all recall + (conclusions and messages) to these sessions; empty list + fails closed """ self.workspace_name: str = workspace_name self.session_name: str | None = session_name + self.session_allowlist: list[str] | None = session_allowlist + self.session_id: str | None = session_id self.observer: str = observer self.observed: str = observed self.observer_peer_card: list[str] | None = observer_peer_card @@ -104,6 +113,24 @@ class DialecticAgent: self._prefetched_conclusion_count: int = 0 self._run_id: str = generate_nanoid() # Always generate for event correlation + def _select_tools(self) -> list[dict[str, Any]]: + """Pick the toolset for this query. + + Minimal reasoning uses a reduced set to reduce cost. Under a session + allowlist `get_reasoning_chain` is dropped entirely rather than left in + to fail at call time: chains traverse provenance across sessions, so it + can't be scoped, and offering it costs both the schema in context and a + wasted turn when the model tries it. + """ + tools = ( + DIALECTIC_TOOLS_MINIMAL + if self.reasoning_level == "minimal" + else DIALECTIC_TOOLS + ) + if self.session_allowlist is not None: + tools = [t for t in tools if t.get("name") != "get_reasoning_chain"] + return tools + async def _initialize_session_history(self) -> None: """Fetch and inject session history into the system prompt if configured.""" if self._session_history_initialized: @@ -121,7 +148,7 @@ class DialecticAgent: token_limit=max_tokens, reverse=False, # chronological order ) - async with tracked_db("dialectic.session_history") as db: + async with tracked_db("dialectic.session_history", read_only=True) as db: result = await db.execute(stmt) messages = result.scalars().all() @@ -179,6 +206,7 @@ class DialecticAgent: workspace_name=self.workspace_name, run_id=self._run_id, parent_category="dialectic", + session_id=self.session_id, ): query_embedding = await embedding_client.embed(query) @@ -192,6 +220,7 @@ class DialecticAgent: limit=prefetch_limit, levels=["explicit"], embedding=query_embedding, + session_allowlist=self.session_allowlist, ) derived_repr = await search_memory( @@ -202,6 +231,7 @@ class DialecticAgent: limit=prefetch_limit, levels=["deductive", "inductive", "contradiction"], embedding=query_embedding, + session_allowlist=self.session_allowlist, ) if explicit_repr.is_empty() and derived_repr.is_empty(): @@ -291,6 +321,7 @@ class DialecticAgent: ] = await create_tool_executor( workspace_name=self.workspace_name, session_name=self.session_name, + session_allowlist=self.session_allowlist, observer=self.observer, observed=self.observed, history_token_limit=settings.DIALECTIC.HISTORY_TOKEN_LIMIT, @@ -301,13 +332,14 @@ class DialecticAgent: return tool_executor, task_name, run_id, start_time - def _telemetry_context(self) -> LLMTelemetryContext: + def _telemetry_context(self, track_name: str | None = None) -> LLMTelemetryContext: """Build the LLMTelemetryContext shared by answer() and answer_stream(). Carries the instance's `_run_id` (always set in __init__) + workspace + peer identifiers so LLMCallCompletedEvent and 's AgentIterationEvent can attribute every per-iteration LLM call back to - this dialectic invocation. + this dialectic invocation. `track_name` names the Langfuse trace/step + (e.g. "Dialectic Agent" vs "Dialectic Agent Stream"). """ return LLMTelemetryContext( workspace_name=self.workspace_name, @@ -315,7 +347,11 @@ class DialecticAgent: parent_category="dialectic", agent_type="dialectic", run_id=self._run_id, + trace_id=self._run_id, + span_id=self._run_id, + session_id=self.session_id, peer_name=self.observed, + track_name=track_name, ) def _log_response_metrics( @@ -404,7 +440,9 @@ class DialecticAgent: ) ) - async def answer(self, query: str) -> str: + async def answer( + self, query: str, response_model: type[BaseModel] | None = None + ) -> str: """ Answer a query about the peer using agentic tool calling. @@ -415,6 +453,8 @@ class DialecticAgent: Args: query: The question to answer about the peer + response_model: Optional Pydantic model the final synthesis must + conform to. When set, the returned string is JSON. Returns: The synthesized answer string @@ -424,12 +464,7 @@ class DialecticAgent: # Get level-specific settings level_settings = settings.DIALECTIC.LEVELS[self.reasoning_level] - # Use minimal tools for minimal reasoning to reduce cost - tools = ( - DIALECTIC_TOOLS_MINIMAL - if self.reasoning_level == "minimal" - else DIALECTIC_TOOLS - ) + tools = self._select_tools() # Use level-specific max_output_tokens if set, otherwise global default max_tokens = ( level_settings.MAX_OUTPUT_TOKENS @@ -437,26 +472,38 @@ class DialecticAgent: else settings.DIALECTIC.MAX_OUTPUT_TOKENS ) - response: HonchoLLMCallResponse[str] = await honcho_llm_call( - model_config=_get_dialectic_level_model_config(self.reasoning_level), - prompt="", # Ignored since we pass messages - max_tokens=max_tokens, - tools=tools, - tool_choice=level_settings.TOOL_CHOICE, - tool_executor=tool_executor, - max_tool_iterations=level_settings.MAX_TOOL_ITERATIONS, - messages=self.messages, - track_name="Dialectic Agent", - max_input_tokens=settings.DIALECTIC.MAX_INPUT_TOKENS, - trace_name="dialectic_chat", - telemetry=self._telemetry_context(), + # cast: `type[BaseModel] | None` matches neither the parsed nor the + # plain-text overload statically, so pyright resolves the stream + # overload — but without stream=True the call is non-streaming. + response = cast( # pyright: ignore[reportInvalidCast] + HonchoLLMCallResponse[Any], + await honcho_llm_call( + model_config=_get_dialectic_level_model_config(self.reasoning_level), + prompt="", # Ignored since we pass messages + max_tokens=max_tokens, + tools=tools, + tool_choice=level_settings.TOOL_CHOICE, + tool_executor=tool_executor, + max_tool_iterations=level_settings.MAX_TOOL_ITERATIONS, + messages=self.messages, + max_input_tokens=settings.DIALECTIC.MAX_INPUT_TOKENS, + trace_name="dialectic_chat", + telemetry=self._telemetry_context(track_name="Dialectic Agent"), + response_model=response_model, + ), ) + # With response_model, the backend parses content into a model + # instance; the API contract is a JSON string. + content = response.content + if isinstance(content, BaseModel): + content = content.model_dump_json(by_alias=True) + self._log_response_metrics( task_name=task_name, run_id=run_id, start_time=start_time, - response_content=response.content, + response_content=content, input_tokens=response.input_tokens, output_tokens=response.output_tokens, cache_read_input_tokens=response.cache_read_input_tokens, @@ -467,9 +514,11 @@ class DialecticAgent: hit_input_token_cap=response.hit_input_token_cap, ) - return response.content + return content - async def answer_stream(self, query: str) -> AsyncIterator[str]: + async def answer_stream( + self, query: str, response_model: type[BaseModel] | None = None + ) -> AsyncIterator[str]: """ Answer a query about the peer using agentic tool calling, streaming the response. @@ -480,6 +529,9 @@ class DialecticAgent: Args: query: The question to answer about the peer + response_model: Optional Pydantic model the final synthesis must + conform to. When set, the streamed text accumulates to JSON + (chunks are raw text; no parsing happens on the stream path). Yields: Chunks of the response text as they are generated @@ -489,12 +541,7 @@ class DialecticAgent: # Get level-specific settings level_settings = settings.DIALECTIC.LEVELS[self.reasoning_level] - # Use minimal tools for minimal reasoning to reduce cost - tools = ( - DIALECTIC_TOOLS_MINIMAL - if self.reasoning_level == "minimal" - else DIALECTIC_TOOLS - ) + tools = self._select_tools() # Use level-specific max_output_tokens if set, otherwise global default max_tokens = ( level_settings.MAX_OUTPUT_TOKENS @@ -515,10 +562,10 @@ class DialecticAgent: tool_executor=tool_executor, max_tool_iterations=level_settings.MAX_TOOL_ITERATIONS, messages=self.messages, - track_name="Dialectic Agent Stream", max_input_tokens=settings.DIALECTIC.MAX_INPUT_TOKENS, trace_name="dialectic_chat", - telemetry=self._telemetry_context(), + telemetry=self._telemetry_context(track_name="Dialectic Agent Stream"), + response_model=response_model, ), ) diff --git a/src/dreamer/dream_scheduler.py b/src/dreamer/dream_scheduler.py index bad850f4..43bba428 100644 --- a/src/dreamer/dream_scheduler.py +++ b/src/dreamer/dream_scheduler.py @@ -218,7 +218,7 @@ class DreamScheduler: configuration = get_configuration(None, session, workspace) if not configuration.dream.enabled: - logger.info( + logger.debug( f"Dreams disabled for {workspace_name}/{session_name}, skipping dream" ) return @@ -319,7 +319,7 @@ async def check_and_schedule_dream( ).total_seconds() / 3600 if hours_since_last_dream < settings.DREAM.MIN_HOURS_BETWEEN_DREAMS: - logger.info( + logger.debug( f"Skipping dream for {collection.observer}/{collection.observed}: only {hours_since_last_dream:.1f} hours " + f"since last dream (minimum: {settings.DREAM.MIN_HOURS_BETWEEN_DREAMS})" ) @@ -359,7 +359,7 @@ async def check_and_schedule_dream( ) ) if pending_exists: - logger.info( + logger.debug( "Skipping dream schedule for %s/%s: pending dream already in queue", collection.observer, collection.observed, diff --git a/src/dreamer/orchestrator.py b/src/dreamer/orchestrator.py index 0b529cc6..00f001f9 100644 --- a/src/dreamer/orchestrator.py +++ b/src/dreamer/orchestrator.py @@ -26,7 +26,11 @@ from sqlalchemy import func, select from src import crud, models from src.config import settings from src.dependencies import tracked_db -from src.dreamer.specialists import SPECIALISTS, SpecialistResult +from src.dreamer.specialists import ( + SPECIALISTS, + CardRefreshSpecialist, + SpecialistResult, +) from src.dreamer.surprisal import SurprisalScore # type: ignore from src.exceptions import SurprisalError from src.schemas import DreamType @@ -307,6 +311,151 @@ async def run_dream( ) +async def run_card_refresh_dream( + workspace_name: str, + observer: str, + observed: str, + session_name: str | None = None, + *, + rebuild: bool = False, + dream_type: str | None = None, + trigger_reason: str | None = None, + delay_reason: str | None = None, +) -> DreamResult | None: + """ + Run a lightweight card-only refresh dream. + + Runs a single CardRefreshSpecialist restricted to peer-card tools + (get_recent_observations, search_memory, update_peer_card) with a low + tool-iteration cap. It never creates or deletes observations. + + Args: + workspace_name: Workspace identifier + observer: Observer peer name + observed: Observed peer name + session_name: Session identifier if specified + rebuild: When True the existing card is NOT injected into the prompt + and the specialist rebuilds it solely from observations present in + the collection (used after removals). + """ + if not settings.DREAM.ENABLED: + return None + + run_id = generate_nanoid() + task_name = f"dream_orchestrator_{run_id}" + start_time = time.perf_counter() + + logger.info( + f"[{run_id}] Starting card-refresh dream for {workspace_name}/{observer}/{observed} (rebuild={rebuild})" + ) + + # Short-lived DB session for config resolution + async with tracked_db("dream.config") as db: + if session_name is not None: + session = await crud.get_session( + db, workspace_name=workspace_name, session_name=session_name + ) + else: + session = None + + workspace = await crud.get_workspace(db, workspace_name=workspace_name) + configuration = get_configuration(None, session, workspace) + if not configuration.dream.enabled: + logger.info( + f"[{run_id}] Dreams disabled for {workspace_name}/{session_name}, skipping card refresh" + ) + return None + if not configuration.peer_card.create: + logger.info( + f"[{run_id}] Peer card creation disabled for {workspace_name}, skipping card refresh" + ) + return None + + specialist_success = False + specialist_result: SpecialistResult | None = None + duration_ms = 0.0 + try: + specialist = CardRefreshSpecialist(rebuild=rebuild) + try: + specialist_result = await specialist.run( + workspace_name=workspace_name, + observer=observer, + observed=observed, + session_name=session_name, + configuration=configuration, + parent_run_id=run_id, + ) + logger.info( + f"[{run_id}] Card refresh completed: {specialist_result.content[:200]}..." + ) + accumulate_metric( + task_name, "card_refresh_result", specialist_result.content, "blob" + ) + specialist_success = specialist_result.success + except Exception as e: + # Exception (not BaseException) — CancelledError must propagate so + # the worker can shut down; the finally still emits the run event. + logger.error( + f"[{run_id}] Card refresh specialist failed: {e}", exc_info=True + ) + accumulate_metric(task_name, "card_refresh_error", str(e), "blob") + + duration_ms = (time.perf_counter() - start_time) * 1000 + accumulate_metric(task_name, "total_duration", duration_ms, "ms") + logger.info(f"[{run_id}] Card-refresh dream completed in {duration_ms:.0f}ms") + log_performance_metrics("dream_orchestrator", run_id) + finally: + # Emit DreamRunEvent unconditionally so analytics see a parent for the + # specialist event, mirroring run_dream. Card refresh is a + # deduction-family run, so its outcome rides on deduction_success. + if duration_ms == 0.0: + duration_ms = (time.perf_counter() - start_time) * 1000 + try: + emit( + DreamRunEvent( + run_id=run_id, + workspace_name=workspace_name, + session_name=session_name, + observer=observer, + observed=observed, + specialists_run=["card_refresh"], + deduction_success=specialist_success, + induction_success=False, + surprisal_enabled=False, + surprisal_conclusion_count=0, + total_iterations=( + specialist_result.iterations if specialist_result else 0 + ), + total_input_tokens=( + specialist_result.input_tokens if specialist_result else 0 + ), + total_output_tokens=( + specialist_result.output_tokens if specialist_result else 0 + ), + total_duration_ms=duration_ms, + dream_type=dream_type, + enabled_types_count=len(settings.DREAM.ENABLED_TYPES), + trigger_reason=trigger_reason, + delay_reason=delay_reason, + ) + ) + except Exception: # pragma: no cover - telemetry must not raise + logger.debug("Failed to emit DreamRunEvent", exc_info=True) + + return DreamResult( + run_id=run_id, + specialists_run=["card_refresh"], + deduction_success=specialist_success, + induction_success=False, + surprisal_enabled=False, + surprisal_conclusion_count=0, + total_iterations=specialist_result.iterations if specialist_result else 0, + total_duration_ms=duration_ms, + input_tokens=specialist_result.input_tokens if specialist_result else 0, + output_tokens=specialist_result.output_tokens if specialist_result else 0, + ) + + def _create_queries_from_surprisal( high_surprisal_obs: list[SurprisalScore], ) -> list[str]: @@ -401,6 +550,28 @@ DREAM: {payload.dream_type} documents for {workspace_name}/{payload.observer}/{p update_data={"dream": dream_meta}, ) + case DreamType.CARD_REFRESH: + # Card-only refresh: never touches observations and never + # advances the omni dream guard pair (last_dream_at / + # last_dream_document_count) — a card refresh must not delay + # or satisfy consolidation scheduling. + result = await run_card_refresh_dream( + workspace_name=workspace_name, + observer=payload.observer, + observed=payload.observed, + session_name=payload.session_name, + rebuild=payload.rebuild, + dream_type=payload.dream_type.value, + trigger_reason=payload.trigger_reason, + delay_reason=payload.delay_reason, + ) + if result is not None: + logger.info( + f"Card-refresh dream completed: run_id={result.run_id}, " + + f"iterations={result.total_iterations}, " + + f"duration={result.total_duration_ms:.0f}ms" + ) + except Exception as e: logger.error( f"Error processing dream task {payload.dream_type} for {payload.observer}/{payload.observed}: {str(e)}", diff --git a/src/dreamer/specialists.py b/src/dreamer/specialists.py index af67b989..61e8a9d3 100644 --- a/src/dreamer/specialists.py +++ b/src/dreamer/specialists.py @@ -32,6 +32,7 @@ from src.telemetry.events import DreamSpecialistEvent, emit from src.telemetry.logging import accumulate_metric, log_performance_metrics from src.telemetry.prometheus.metrics import TokenTypes from src.utils.agent_tools import ( + CARD_REFRESH_SPECIALIST_TOOLS, DEDUCTION_SPECIALIST_TOOLS, INDUCTION_SPECIALIST_TOOLS, create_tool_executor, @@ -70,6 +71,64 @@ class SpecialistResult: # Tool names to exclude when peer card creation is disabled PEER_CARD_TOOL_NAMES = {"update_peer_card"} +# Shared PEER CARD system-prompt section (identity-store taxonomy + rules). +# Used verbatim by DeductionSpecialist and CardRefreshSpecialist. +PEER_CARD_SYSTEM_SECTION = """ + +## PEER CARD (REQUIRED) + +The peer card is the target observee's identity store: stable identity markers that distinguish this entity from others and persist across interactions. Behavior, tendencies, transient state, and episodic facts belong in observations, not on the peer card. + +A peer can be anything with identity that changes over time — a human, an agent, a codebase, a team, an organization. Do not assume the target observee is human. Do not require any field; empty is the correct output when evidence is absent. + +### Allowed entry kinds + +Each entry must start with one of these four prefixes (exact case, followed by a space): + +- `IDENTITY: ...` — canonical name, kind, aliases, IDs + - `IDENTITY: Name: Alice` + - `IDENTITY: Kind: Python monorepo` + - `IDENTITY: Version: 4.2` + - `IDENTITY: Aliases: alice@example.com` +- `ATTRIBUTE: ...` — stable durable property of the entity (including explicitly stated standing preferences) + - `ATTRIBUTE: Location: NYC` + - `ATTRIBUTE: Language: Python` + - `ATTRIBUTE: Prefers tea` + - `ATTRIBUTE: Charter: ship Honcho infrastructure` +- `RELATIONSHIP: ...` — durable link to another entity + - `RELATIONSHIP: Spouse: Bob` + - `RELATIONSHIP: Maintainer: vineeth` + - `RELATIONSHIP: Members: vineeth, rajat` +- `INSTRUCTION: ...` — standing rule of engagement that the target observee has explicitly stated (do/don't for the observer). Only when explicit; never inferred from behavior. + - `INSTRUCTION: Call me Vee` + - `INSTRUCTION: Never push to main without review` + +### Rules + +1. **Stable.** If the value plausibly changes within six months absent a deliberate announcement, it does not belong on the card. Prefer leaving the card empty over filling it with volatile content. +2. **Subject is the target observee.** Every entry must be a fact about the target observee, not about another participant in the session. Never write facts about co-occurring peers into the card, no matter how frequently they appear in the messages. +3. **Evidence-grounded.** Only write what the target observee has explicitly stated, or what another participant has explicitly stated about the target observee with the target observee's assent. No "general knowledge" inferences (`"co-founder"` does not imply an age; mentioning a colleague does not imply a family relationship). +4. **Type-agnostic.** The target observee may not be human. Do not require name/age/location/family/occupation fields. +5. **No behavioral content.** TRAITs, behavioral tendencies, patterns, and inferred preferences belong in observations, not on the peer card. Do not write `TRAIT:` entries or behavioral `PREFERENCE:` entries — they will be rejected. +6. **No evidence bundles.** Each entry is one concise fact. No `e.g.` clauses, no parenthetical example lists, no semicolon-separated value dumps. + +### Migrating an existing peer card + +The CURRENT PEER CARD shown in the user message may contain entries from an older format that do not start with an allowed prefix (e.g. `Name: Alice`, `Lives in NYC`, `TRAIT: Analytical`, `PREFERENCE: Detailed explanations`). When you call `update_peer_card`, you are responsible for re-emitting the entries you want to keep — entries you omit are dropped, and entries without an allowed prefix are silently rejected. + +For each legacy entry: + +- If it is still a valid identity marker, re-emit it under the correct prefix and keep the original content where reasonable. Examples: + - `Name: Alice` → `IDENTITY: Name: Alice` + - `Lives in NYC` → `ATTRIBUTE: Location: NYC` + - `Works at Google` → `ATTRIBUTE: Employer: Google` + - `INSTRUCTION: Call me Vee` → keep as is (already correctly prefixed) +- Drop entries that violate the rules above: behavioral `TRAIT:` lines, inferred behavioral `PREFERENCE:` lines, one-off events, transient state, evidence bundles. Do not re-prefix them — they are not identity markers. + +When in doubt about a specific legacy entry, prefer migrating it (so valid info isn't lost) over dropping it. Splitting one dense legacy entry into multiple correctly-prefixed entries is fine and encouraged (e.g. a semicolon-separated `Tech Stack:` dump can become several `ATTRIBUTE:` lines, one per durable tool/platform). + +Call `update_peer_card` with the complete deduplicated list when there is a durable identity update to record, or when the existing card needs migration. Entries that do not start with one of the four allowed prefixes will be rejected. Keep concise (max 40 entries).""" + class BaseSpecialist(ABC): """Base class for agentic specialists.""" @@ -78,6 +137,10 @@ class BaseSpecialist(ABC): # Whether this specialist is allowed to write to the peer card. Defaults to True; # specialists that should never touch the card (e.g., induction) override to False. can_update_peer_card: bool = True + # Whether the current peer card is fetched and injected into the user prompt. + # Card-refresh runs in rebuild mode set this to False so the card is + # reconstructed solely from observations present in the collection. + inject_peer_card: bool = True # Subclasses can override to customize the peer card update instruction peer_card_update_instruction: str = ( "Only update this with durable identity markers via `update_peer_card`." @@ -111,12 +174,21 @@ class BaseSpecialist(ABC): @abstractmethod def build_user_prompt( self, + observed: str, hints: list[str] | None, peer_card: list[str] | None = None, ) -> str: """Build the user prompt with optional exploration hints and current peer card.""" ... + def _build_target_observee_context(self, observed: str) -> str: + return f"""Target observee: +{observed} + +The target observee is the peer identified above. When created observations need to name this subject, use the exact observee id above, not the phrase "the target observee". + +""" + def _build_peer_card_context(self, peer_card: list[str] | None) -> str: """Build the peer card context section for user prompts.""" if not peer_card: @@ -160,6 +232,10 @@ If you update it, send the full deduplicated list and remove stale entries. SpecialistResult with metrics and content """ run_id = parent_run_id or generate_nanoid() + # Specialists sharing the orchestrator's run_id (one dream trace) each get a + # distinct span_id so their CloudEvents trace resource ids don't collide; + # trace_id stays run_id so Langfuse still groups them (keyed by agent_type). + span_id = generate_nanoid() if parent_run_id is not None else run_id task_name = f"dreamer_{self.name}_{run_id}" start_time = time.perf_counter() @@ -205,9 +281,11 @@ If you update it, send the full deduplicated list and remove stale entries. configuration is None or configuration.peer_card.create ) - # Fetch current peer card to inject into prompt (saves a tool call) + # Fetch current peer card to inject into prompt (saves a tool call). + # Skipped when inject_peer_card is False (card-refresh rebuild + # mode): the card must be reconstructed from observations only. current_peer_card: list[str] | None = None - if peer_card_enabled: + if peer_card_enabled and self.inject_peer_card: current_peer_card = await crud.get_peer_card( db, workspace_name=workspace_name, @@ -226,7 +304,11 @@ If you update it, send the full deduplicated list and remove stale entries. }, { "role": "user", - "content": self.build_user_prompt(hints, current_peer_card), + "content": self.build_user_prompt( + observed=observed, + hints=hints, + peer_card=current_peer_card, + ), }, ] @@ -273,15 +355,20 @@ If you update it, send the full deduplicated list and remove stale entries. tool_executor=tool_executor, max_tool_iterations=self.get_max_iterations(), messages=messages, - track_name=f"Dreamer/{self.name}", telemetry=LLMTelemetryContext( workspace_name=workspace_name, call_purpose=call_purpose_slug, parent_category="dream", agent_type=self.name, run_id=run_id, + # Root span per specialist run (distinct span_id, see above). + # parent_span_id stays None for now; wiring specialists as + # children of a dream-level trace is forking (out of scope). + trace_id=run_id, + span_id=span_id, observer=observer, observed=observed, + track_name=f"Dreamer/{self.name}", ), ) @@ -465,65 +552,12 @@ class DeductionSpecialist(BaseSpecialist): def build_system_prompt( self, observed: str, *, peer_card_enabled: bool = True ) -> str: + _ = observed peer_card_section = "" if peer_card_enabled: - peer_card_section = f""" + peer_card_section = PEER_CARD_SYSTEM_SECTION -## PEER CARD (REQUIRED) - -The peer card is {observed}'s identity store: stable identity markers that distinguish this entity from others and persist across interactions. Behavior, tendencies, transient state, and episodic facts belong in observations, not on the peer card. - -A peer can be anything with identity that changes over time — a human, an agent, a codebase, a team, an organization. Do not assume {observed} is human. Do not require any field; empty is the correct output when evidence is absent. - -### Allowed entry kinds - -Each entry must start with one of these four prefixes (exact case, followed by a space): - -- `IDENTITY: ...` — canonical name, kind, aliases, IDs - - `IDENTITY: Name: Alice` - - `IDENTITY: Kind: Python monorepo` - - `IDENTITY: Version: 4.2` - - `IDENTITY: Aliases: alice@example.com` -- `ATTRIBUTE: ...` — stable durable property of the entity (including explicitly stated standing preferences) - - `ATTRIBUTE: Location: NYC` - - `ATTRIBUTE: Language: Python` - - `ATTRIBUTE: Prefers tea` - - `ATTRIBUTE: Charter: ship Honcho infrastructure` -- `RELATIONSHIP: ...` — durable link to another entity - - `RELATIONSHIP: Spouse: Bob` - - `RELATIONSHIP: Maintainer: vineeth` - - `RELATIONSHIP: Members: vineeth, rajat` -- `INSTRUCTION: ...` — standing rule of engagement that {observed} has explicitly stated (do/don't for the observer). Only when explicit; never inferred from behavior. - - `INSTRUCTION: Call me Vee` - - `INSTRUCTION: Never push to main without review` - -### Rules - -1. **Stable.** If the value plausibly changes within six months absent a deliberate announcement, it does not belong on the card. Prefer leaving the card empty over filling it with volatile content. -2. **Subject is {observed}.** Every entry must be a fact about {observed}, not about another participant in the session. Never write facts about co-occurring peers into the card, no matter how frequently they appear in the messages. -3. **Evidence-grounded.** Only write what {observed} has explicitly stated, or what another participant has explicitly stated about {observed} with {observed}'s assent. No "general knowledge" inferences (`"co-founder"` does not imply an age; mentioning a colleague does not imply a family relationship). -4. **Type-agnostic.** {observed} may not be human. Do not require name/age/location/family/occupation fields. -5. **No behavioral content.** TRAITs, behavioral tendencies, patterns, and inferred preferences belong in observations, not on the peer card. Do not write `TRAIT:` entries or behavioral `PREFERENCE:` entries — they will be rejected. -6. **No evidence bundles.** Each entry is one concise fact. No `e.g.` clauses, no parenthetical example lists, no semicolon-separated value dumps. - -### Migrating an existing peer card - -The CURRENT PEER CARD shown in the user message may contain entries from an older format that do not start with an allowed prefix (e.g. `Name: Alice`, `Lives in NYC`, `TRAIT: Analytical`, `PREFERENCE: Detailed explanations`). When you call `update_peer_card`, you are responsible for re-emitting the entries you want to keep — entries you omit are dropped, and entries without an allowed prefix are silently rejected. - -For each legacy entry: - -- If it is still a valid identity marker, re-emit it under the correct prefix and keep the original content where reasonable. Examples: - - `Name: Alice` → `IDENTITY: Name: Alice` - - `Lives in NYC` → `ATTRIBUTE: Location: NYC` - - `Works at Google` → `ATTRIBUTE: Employer: Google` - - `INSTRUCTION: Call me Vee` → keep as is (already correctly prefixed) -- Drop entries that violate the rules above: behavioral `TRAIT:` lines, inferred behavioral `PREFERENCE:` lines, one-off events, transient state, evidence bundles. Do not re-prefix them — they are not identity markers. - -When in doubt about a specific legacy entry, prefer migrating it (so valid info isn't lost) over dropping it. Splitting one dense legacy entry into multiple correctly-prefixed entries is fine and encouraged (e.g. a semicolon-separated `Tech Stack:` dump can become several `ATTRIBUTE:` lines, one per durable tool/platform). - -Call `update_peer_card` with the complete deduplicated list when there is a durable identity update to record, or when the existing card needs migration. Entries that do not start with one of the four allowed prefixes will be rejected. Keep concise (max 40 entries).""" - - return f"""You are a deductive reasoning agent analyzing observations about {observed}. + return f"""You are a deductive reasoning agent analyzing observations about the target observee. ## YOUR JOB @@ -579,18 +613,21 @@ Use `create_observations_deductive`. 3. Always include source_ids linking to the observations you're synthesizing 4. Empty or missing source_ids will be rejected 5. Delete outdated observations - don't leave duplicates -6. Quality over quantity - fewer good deductions beat many weak ones""" +6. Quality over quantity - fewer good deductions beat many weak ones +7. When you are finished, do not output a summary of what you did - output only the token DONE""" def build_user_prompt( self, + observed: str, hints: list[str] | None, peer_card: list[str] | None = None, ) -> str: + target_observee_context = self._build_target_observee_context(observed) peer_card_context = self._build_peer_card_context(peer_card) if hints: hints_str = "\n".join(f"- {q}" for q in hints[:5]) - return f"""{peer_card_context}Start by exploring recent observations and messages. These topics may be worth investigating: + return f"""{target_observee_context}{peer_card_context}Start by exploring recent observations and messages. These topics may be worth investigating: {hints_str} @@ -598,7 +635,7 @@ But follow the evidence - if you find something more interesting, pursue that in Begin with `get_recent_observations` to see what's there.""" - return f"""{peer_card_context}Explore the observation space and create deductive observations. + return f"""{target_observee_context}{peer_card_context}Explore the observation space and create deductive observations. Start with `get_recent_observations` to see what's been learned recently, then investigate whatever seems most promising. @@ -647,8 +684,9 @@ class InductionSpecialist(BaseSpecialist): def build_system_prompt( self, observed: str, *, peer_card_enabled: bool = True ) -> str: + _ = observed _ = peer_card_enabled - return f"""You are an inductive reasoning agent identifying patterns about {observed}. + return """You are an inductive reasoning agent identifying patterns about the target observee. ## YOUR JOB @@ -707,20 +745,23 @@ Use `create_observations_inductive`. 3. Confidence based on evidence count: 2=low, 3-4=medium, 5+=high 4. Look for HOW things change over time, not just static facts 5. Include source_ids - always link back to evidence -6. Empty or missing source_ids will be rejected""" +6. Empty or missing source_ids will be rejected +7. When you are finished, do not output a summary of what you did - output only the token DONE""" def build_user_prompt( self, + observed: str, hints: list[str] | None, peer_card: list[str] | None = None, ) -> str: + target_observee_context = self._build_target_observee_context(observed) # Induction does not consume peer card context — it produces inductive # observations, not identity-marker updates. _ = peer_card if hints: hints_str = "\n".join(f"- {q}" for q in hints[:5]) - return f"""Explore and find patterns. These areas may be worth investigating: + return f"""{target_observee_context}Explore and find patterns. These areas may be worth investigating: {hints_str} @@ -728,13 +769,120 @@ But follow the evidence - if you find patterns elsewhere, pursue those. Start with `get_recent_observations`.""" - return """Explore the observation space and identify patterns. + return f"""{target_observee_context}Explore the observation space and identify patterns. Remember: patterns need 2+ sources. Look for tendencies, preferences, and behavioral regularities. Go.""" +class CardRefreshSpecialist(BaseSpecialist): + """ + Card-only maintenance specialist for the ``card_refresh`` dream type. + + Restricted to peer-card work: it may discover observations + (get_recent_observations, search_memory) and rewrite the peer card + (update_peer_card). It has NO observation-mutating tools — a card refresh + must never create or delete observations. + + Two modes: + - refresh (default): the current card is injected into the prompt and the + specialist folds in new identity markers. + - rebuild: the current card is NOT injected; the specialist reconstructs + the card solely from observations present in the collection. Used after + removals, where the old card may contain facts whose support was deleted. + + Not a singleton — instantiated per run because ``rebuild`` is per-dream + state. + """ + + name: str = "card_refresh" + peer_card_update_instruction: str = "Update this with `update_peer_card`. See the PEER CARD section in the system prompt for the allowed entry kinds and rules." + + # Low iteration ceiling for this lightweight, single-purpose run. + MAX_ITERATIONS_CEILING: int = 6 + + def __init__(self, *, rebuild: bool = False) -> None: + self.rebuild: bool = rebuild + # In rebuild mode the existing card is withheld from the prompt. + self.inject_peer_card: bool = not rebuild + + def get_tools(self, *, peer_card_enabled: bool = True) -> list[dict[str, Any]]: + if peer_card_enabled: + return CARD_REFRESH_SPECIALIST_TOOLS + # Defensive: a card refresh without card write access is a no-op, and + # the orchestrator skips the run entirely when peer cards are disabled. + return [ + t + for t in CARD_REFRESH_SPECIALIST_TOOLS + if t["name"] not in PEER_CARD_TOOL_NAMES + ] + + def get_model_config(self) -> ConfiguredModelSettings: + # Card refresh is a deduction-family task; reuse its model config. + return _require_specialist_model_config( + settings.DREAM.DEDUCTION_MODEL_CONFIG, + specialist_name="DREAM CARD_REFRESH", + ) + + def get_max_tokens(self) -> int: + return 8192 + + def get_max_iterations(self) -> int: + return min(self.MAX_ITERATIONS_CEILING, settings.DREAM.MAX_TOOL_ITERATIONS) + + def build_system_prompt( + self, observed: str, *, peer_card_enabled: bool = True + ) -> str: + _ = observed + _ = peer_card_enabled + rebuild_section = "" + if self.rebuild: + rebuild_section = """ + +## REBUILD MODE + +The existing peer card is deliberately NOT shown to you: it may contain entries whose supporting observations have since been removed. Build the card solely from the observations you find in the collection right now. Do not carry over or guess at prior card content — if an identity marker is not supported by a current observation, it does not go on the card.""" + + return f"""You are a peer-card maintenance agent for the target observee. + +## YOUR JOB + +Refresh the peer card and nothing else. You cannot create or delete observations — you have no tools for that. Your only write operation is `update_peer_card`. + +## PROCESS + +1. Survey the observation space: start with `get_recent_observations`, then use `search_memory` for targeted follow-ups (names, roles, locations, standing instructions). +2. Extract stable identity markers supported by the observations you found. +3. Call `update_peer_card` once with the complete deduplicated list. + +Keep it short — a handful of tool calls at most.{rebuild_section} +{PEER_CARD_SYSTEM_SECTION}""" + + def build_user_prompt( + self, + observed: str, + hints: list[str] | None, + peer_card: list[str] | None = None, + ) -> str: + _ = hints + target_observee_context = self._build_target_observee_context(observed) + peer_card_context = self._build_peer_card_context(peer_card) + + if self.rebuild: + return f"""{target_observee_context}Rebuild the peer card from scratch. + +The previous card is not shown and must not be assumed: reconstruct the card solely from observations currently in the collection. Start with `get_recent_observations`, verify with `search_memory` where needed, then call `update_peer_card` with the complete list. + +Go.""" + + return f"""{target_observee_context}{peer_card_context}Refresh the peer card. + +Review recent observations with `get_recent_observations` (and `search_memory` for targeted checks), then call `update_peer_card` with the complete deduplicated list if there is anything to add, correct, or migrate. If the card is already accurate and complete, finish without updating it. + +Go.""" + + # Singleton instances SPECIALISTS: dict[str, BaseSpecialist] = { "deduction": DeductionSpecialist(), diff --git a/src/dreamer/trees/__init__.py b/src/dreamer/trees/__init__.py index e1ac6acf..eba2ff3d 100644 --- a/src/dreamer/trees/__init__.py +++ b/src/dreamer/trees/__init__.py @@ -29,6 +29,15 @@ def create_tree(tree_type: str, **kwargs: Any) -> SurprisalTree: Raises: ValueError: If tree_type is not recognized """ + # `surprisal.py` calls this factory with a uniform `k=settings.DREAM.SURPRISAL.TREE_K` kwarg for every tree type, + # but `k` is only meaningful for the kNN-based trees (kdtree, balltree, graph). + # The other 4 use different tunables and raise TypeError if `k` is passed. + # Drop it here so the factory accepts a uniform kwargs dict. + + trees_without_k = {"rptree", "covertree", "lsh", "prototype"} + if tree_type in trees_without_k: + kwargs.pop("k", None) + if tree_type == "rptree": return RPTree(**kwargs) elif tree_type == "kdtree": diff --git a/src/embedding_client.py b/src/embedding_client.py index 60516bc5..07197f43 100644 --- a/src/embedding_client.py +++ b/src/embedding_client.py @@ -9,6 +9,7 @@ from typing import Any, Literal, NamedTuple, TypeVar import tiktoken from google import genai from google.genai import types as genai_types +from nanoid import generate as generate_nanoid from openai import AsyncOpenAI from .config import EmbeddingModelConfig, resolve_embedding_model_config, settings @@ -88,6 +89,7 @@ def _publish_embedding_event( get_embedding_call_purpose, get_embedding_parent_category, get_embedding_run_id, + get_embedding_session_id, get_embedding_workspace_name, ) @@ -121,6 +123,30 @@ def _publish_embedding_event( run_id=get_embedding_run_id(), ) ) + + # Trace stream (ground-truth) — gated on payload tracing. Each embedding + # gets its own span nested under the driving agent run (parent_span_id = + # run_id), so multiple embeddings in one run don't share a span id. + if settings.TELEMETRY.TRACE_PAYLOADS_ENABLED: + from src.telemetry.events import EmbeddingCallTracedEvent, emit_trace + + run_id = get_embedding_run_id() + span_id = generate_nanoid() + emit_trace( + EmbeddingCallTracedEvent( + trace_id=run_id or span_id, + span_id=span_id, + parent_span_id=run_id, + session_id=get_embedding_session_id(), + call_purpose=purpose_slug, + parent_category=get_embedding_parent_category(), + provider=provider, + model=model, + provider_input_tokens=input_tokens_estimate, + provider_output_tokens=0, + input_count=input_count, + ) + ) except Exception: # pragma: no cover - telemetry must not raise logger.debug("Failed to emit EmbeddingCallCompletedEvent", exc_info=True) @@ -250,76 +276,61 @@ class _EmbeddingClient: async def simple_batch_embed(self, texts: list[str]) -> list[list[float]]: """ - Simple batch embedding for a list of text strings. + Batch-embed a list of text strings. Each input must already fit within + `max_embedding_tokens`; this method does not sub-chunk oversized inputs. + + Internally goes through the same token-aware batching pipeline as + `batch_embed()` so the per-request token cap is respected. Args: texts: List of text strings to embed Returns: - List of embedding vectors corresponding to input texts + List of embedding vectors, one per input text (in order) Raises: ValueError: If any text exceeds token limits """ - embeddings: list[list[float]] = [] + if not texts: + return [] - for i in range(0, len(texts), self.max_batch_size): - batch = texts[i : i + self.max_batch_size] - - async def _embed_batch(batch: list[str] = batch) -> list[list[float]]: - """One provider call for one batch. Lifted into a closure so - _emit_embedding_call can time + emit + propagate errors.""" - batch_embeddings: list[list[float]] = [] - if isinstance(self.client, genai.Client): - # Type cast needed due to genai type signature complexity - response = await self.client.aio.models.embed_content( - model=self.model, - contents=batch, # pyright: ignore[reportArgumentType] - config={"output_dimensionality": self.vector_dimensions}, - ) - if response.embeddings: - for emb in response.embeddings: - if emb.values: - batch_embeddings.append( - self._validate_embedding_dimensions(emb.values) - ) - else: # openai - openai_kwargs: dict[str, Any] = { - "input": batch, - "model": self.model, - } - if self.send_dimensions: - openai_kwargs["dimensions"] = self.vector_dimensions - response = await self.client.embeddings.create(**openai_kwargs) - batch_embeddings.extend( - [ - self._validate_embedding_dimensions(data.embedding) - for data in response.data - ] - ) - return batch_embeddings - - try: - # Pre-compute the tiktoken estimate ONCE for telemetry; the - # batch contents don't change between attempts. - tokens_estimate = sum(len(self.encoding.encode(t)) for t in batch) - batch_embeddings = await _emit_embedding_call( - provider=self.transport, - model=self.model, - texts=batch, - input_tokens_estimate=tokens_estimate, - fn=_embed_batch, + # Validate per-input token limit and collect token counts for batching + token_counts: list[int] = [] + for idx, text in enumerate(texts): + tokens = len(self.encoding.encode(text)) + if tokens > self.max_embedding_tokens: + raise ValueError( + f"Text at index {idx} exceeds maximum token limit of {self.max_embedding_tokens} tokens (got {tokens} tokens)" ) - embeddings.extend(batch_embeddings) - except Exception as e: - # Check if it's a token limit error and re-raise as ValueError for consistency - if "token" in str(e).lower(): - raise ValueError( - f"Text content exceeds maximum token limit of {self.max_embedding_tokens}." - ) from e - raise + token_counts.append(tokens) - return embeddings + # Use positional indices as text_ids so we can reassemble in input order. + text_chunks: dict[str, list[tuple[str, int]]] = { + str(i): [(text, token_counts[i])] for i, text in enumerate(texts) + } + + batches = self._create_batches(text_chunks) + batch_results = await asyncio.gather( + *[self._process_batch(batch) for batch in batches], + ) + + combined: dict[str, list[list[float]]] = self._accumulate_embeddings( + batch_results + ) + return [combined[str(i)][0] for i in range(len(texts))] + + def prepare_chunks(self, id_resource_dict: dict[str, str]) -> dict[str, list[str]]: + """ + Public helper: tokenize and chunk texts using the same rules as + `batch_embed()`. Returns ordered chunk texts per input id. + + Intended for callers that want to persist embeddable chunks + before later embedding them off the request path. + """ + return { + text_id: [chunk_text for chunk_text, _ in chunks] + for text_id, chunks in self._prepare_chunks(id_resource_dict).items() + } async def batch_embed( self, id_resource_dict: dict[str, str] @@ -623,9 +634,13 @@ class EmbeddingClient: return await self._get_client().embed(query) async def simple_batch_embed(self, texts: list[str]) -> list[list[float]]: - """Simple batch embedding for a list of text strings.""" + """Batch embed a list of text strings (each must fit token limit).""" return await self._get_client().simple_batch_embed(texts) + def prepare_chunks(self, id_resource_dict: dict[str, str]) -> dict[str, list[str]]: + """Chunk texts using the same rules as `batch_embed` (no network).""" + return self._get_client().prepare_chunks(id_resource_dict) + async def batch_embed( self, id_resource_dict: dict[str, str] ) -> dict[str, list[list[float]]]: diff --git a/src/llm/api.py b/src/llm/api.py index 9c9a628d..13c3a5e5 100644 --- a/src/llm/api.py +++ b/src/llm/api.py @@ -21,7 +21,6 @@ from tenacity import retry, stop_after_attempt, wait_exponential from src.config import ConfiguredModelSettings, ModelConfig from src.exceptions import ValidationException -from src.telemetry.logging import conditional_observe from src.telemetry.reasoning_traces import log_reasoning_trace from .executor import honcho_llm_call_inner @@ -31,7 +30,7 @@ from .runtime import ( effective_temperature, plan_attempt, resolve_runtime_model_config, - update_current_langfuse_observation, + start_langfuse_agent_run, ) from .tool_loop import execute_tool_loop from .types import ( @@ -54,7 +53,6 @@ async def honcho_llm_call( model_config: ModelConfig | ConfiguredModelSettings, prompt: str, max_tokens: int, - track_name: str | None = None, response_model: type[M], json_mode: bool = False, temperature: float | None = None, @@ -84,7 +82,6 @@ async def honcho_llm_call( model_config: ModelConfig | ConfiguredModelSettings, prompt: str, max_tokens: int, - track_name: str | None = None, response_model: None = None, json_mode: bool = False, temperature: float | None = None, @@ -114,7 +111,6 @@ async def honcho_llm_call( model_config: ModelConfig | ConfiguredModelSettings, prompt: str, max_tokens: int, - track_name: str | None = None, response_model: type[BaseModel] | None = None, json_mode: bool = False, temperature: float | None = None, @@ -138,13 +134,11 @@ async def honcho_llm_call( ) -> AsyncIterator[HonchoLLMCallStreamChunk] | StreamingResponseWithMetadata: ... -@conditional_observe(name="LLM Call") async def honcho_llm_call( *, model_config: ModelConfig | ConfiguredModelSettings, prompt: str, max_tokens: int, - track_name: str | None = None, response_model: type[BaseModel] | None = None, json_mode: bool = False, temperature: float | None = None, @@ -206,11 +200,6 @@ async def honcho_llm_call( call_thinking_budget_tokens=thinking_budget_tokens, call_reasoning_effort=reasoning_effort, ) - update_current_langfuse_observation( - plan.provider, - plan.model, - name=track_name, - ) return plan async def _call_with_provider_selection() -> ( @@ -267,8 +256,9 @@ async def honcho_llm_call( decorated = _call_with_provider_selection - if track_name: - decorated = ai_track(track_name)(decorated) + sentry_track_name = telemetry.track_name if telemetry is not None else None + if sentry_track_name: + decorated = ai_track(sentry_track_name)(decorated) def before_retry_callback(retry_state: Any) -> None: """Update attempt counter before each retry + log transient failures. @@ -397,8 +387,8 @@ async def honcho_llm_call( ) wrapped = _toolless_call - if track_name: - wrapped = ai_track(track_name)(wrapped) + if sentry_track_name: + wrapped = ai_track(sentry_track_name)(wrapped) if enable_retry: wrapped = retry( stop=stop_after_attempt(retry_attempts), @@ -406,7 +396,9 @@ async def honcho_llm_call( before_sleep=before_retry_callback, )(wrapped) result: ( - HonchoLLMCallResponse[Any] | AsyncIterator[HonchoLLMCallStreamChunk] + HonchoLLMCallResponse[Any] + | AsyncIterator[HonchoLLMCallStreamChunk] + | StreamingResponseWithMetadata ) = await wrapped() else: result = await decorated() @@ -429,30 +421,59 @@ async def honcho_llm_call( ) return result - # execute_tool_loop raises ValidationException on out-of-range - # max_tool_iterations; fail-fast is cheaper than silent clamping here. - result = await execute_tool_loop( - prompt=prompt, - max_tokens=max_tokens, - messages=messages, - tools=tools, - tool_choice=tool_choice, - tool_executor=tool_executor, - max_tool_iterations=max_tool_iterations, - response_model=response_model, - json_mode=json_mode, - temperature=temperature, - stop_seqs=stop_seqs, - verbosity=verbosity, - enable_retry=enable_retry, - retry_attempts=retry_attempts, - max_input_tokens=max_input_tokens, - get_attempt_plan=_get_attempt_plan, - before_retry_callback=before_retry_callback, - stream_final=stream_final_only, - iteration_callback=iteration_callback, - telemetry=telemetry, - ) + # One run-level Langfuse trace wraps the whole run; step/LLM/tool spans + # nest under it (the run handle keeps `start_as_current_observation` open + # via ExitStack, so the run span stays current OTel-wise even though we + # never use a `with` block here). The handle is passed into + # `execute_tool_loop` so streaming results own it from construction and + # close the span after drain — that's how the streamed text shows up as + # the trace's output instead of blank. Non-streaming results: we end in + # the `finally`. + run_label = (telemetry.track_name if telemetry else None) or "Agent" + run_handle = start_langfuse_agent_run(run_label, telemetry) + if run_handle is not None: + # Mirror execute_tool_loop's prompt-only handling: when messages is + # omitted it seeds the conversation with a single user message built + # from prompt. Record that same effective input so the run span isn't + # blank for prompt-only calls. + run_handle.update( + input=messages if messages else [{"role": "user", "content": prompt}] + ) + try: + # execute_tool_loop raises ValidationException on out-of-range + # max_tool_iterations; fail-fast is cheaper than silent clamping here. + result = await execute_tool_loop( + prompt=prompt, + max_tokens=max_tokens, + messages=messages, + tools=tools, + tool_choice=tool_choice, + tool_executor=tool_executor, + max_tool_iterations=max_tool_iterations, + response_model=response_model, + json_mode=json_mode, + temperature=temperature, + stop_seqs=stop_seqs, + verbosity=verbosity, + enable_retry=enable_retry, + retry_attempts=retry_attempts, + max_input_tokens=max_input_tokens, + get_attempt_plan=_get_attempt_plan, + before_retry_callback=before_retry_callback, + stream_final=stream_final_only, + iteration_callback=iteration_callback, + telemetry=telemetry, + langfuse_run_handle=run_handle, + ) + except BaseException: + if run_handle is not None: + run_handle.end() + raise + # Streaming wrapper owns the handle and closes it after drain; + # non-streaming paths (always a HonchoLLMCallResponse here) close it now + # with the final content as output. + if run_handle is not None and isinstance(result, HonchoLLMCallResponse): + run_handle.end(output=result.content) if trace_name and isinstance(result, HonchoLLMCallResponse): log_reasoning_trace( task_type=trace_name, diff --git a/src/llm/backend.py b/src/llm/backend.py index 5645998c..380911d2 100644 --- a/src/llm/backend.py +++ b/src/llm/backend.py @@ -14,7 +14,8 @@ class ToolCallResult: id: str name: str input: dict[str, Any] - thought_signature: str | None = None + # Gemini returns this as raw bytes; other providers omit it. + thought_signature: str | bytes | None = None @dataclass(slots=True) diff --git a/src/llm/backends/anthropic.py b/src/llm/backends/anthropic.py index 17138583..614a2e32 100644 --- a/src/llm/backends/anthropic.py +++ b/src/llm/backends/anthropic.py @@ -9,7 +9,8 @@ from anthropic.types import TextBlock, ThinkingBlock, ToolUseBlock from pydantic import BaseModel, ValidationError from src.llm.backend import CompletionResult, StreamChunk, ToolCallResult -from src.llm.structured_output import repair_response_model_json +from src.llm.request_builder import apply_sdk_passthroughs +from src.llm.structured_output import repair_response_model_json, schema_instruction class AnthropicBackend: @@ -69,27 +70,32 @@ class AnthropicBackend: for key in ("top_p", "top_k"): if key in extra_params: params[key] = extra_params[key] + # Operator escape hatch: forward Anthropic SDK passthrough kwargs + # from ModelConfig.provider_params. Shallow merge with operator-wins. + apply_sdk_passthroughs(params, extra_params) + # The '{' prefill forces a JSON-first response, which suppresses + # tool_use blocks — skip it when tools are available and rely on the + # conditional instruction + repair fallback instead. use_json_prefill = ( bool(response_format or self._json_mode(extra_params)) and not thinking_budget_tokens + and not tools and self._supports_assistant_prefill(model) ) if use_json_prefill and params["messages"]: if response_format and isinstance(response_format, type): - schema_json = json.dumps(response_format.model_json_schema(), indent=2) self._append_text_to_last_message( params["messages"], - f"\n\nRespond with valid JSON matching this schema:\n{schema_json}", + schema_instruction(response_format, tools_present=False), ) params["messages"].append({"role": "assistant", "content": "{"}) elif ( response_format and isinstance(response_format, type) and params["messages"] ): - schema_json = json.dumps(response_format.model_json_schema(), indent=2) self._append_text_to_last_message( params["messages"], - f"\n\nRespond with valid JSON matching this schema:\n{schema_json}", + schema_instruction(response_format, tools_present=bool(tools)), ) response = await self._client.messages.create(**params) @@ -148,26 +154,30 @@ class AnthropicBackend: for key in ("top_p", "top_k"): if key in extra_params: params[key] = extra_params[key] + # Operator escape hatch: forward Anthropic SDK passthrough kwargs + # from ModelConfig.provider_params. Shallow merge with operator-wins. + apply_sdk_passthroughs(params, extra_params) + # See complete(): no '{' prefill when tools are available, so + # tool_use blocks stay reachable on the streamed path too. use_json_prefill = ( bool(response_format or is_json_mode) and not thinking_budget_tokens + and not tools and self._supports_assistant_prefill(model) ) if use_json_prefill and params["messages"]: if response_format and isinstance(response_format, type): - schema_json = json.dumps(response_format.model_json_schema(), indent=2) self._append_text_to_last_message( params["messages"], - f"\n\nRespond with valid JSON matching this schema:\n{schema_json}", + schema_instruction(response_format, tools_present=False), ) params["messages"].append({"role": "assistant", "content": "{"}) elif ( response_format and isinstance(response_format, type) and params["messages"] ): - schema_json = json.dumps(response_format.model_json_schema(), indent=2) self._append_text_to_last_message( params["messages"], - f"\n\nRespond with valid JSON matching this schema:\n{schema_json}", + schema_instruction(response_format, tools_present=bool(tools)), ) if thinking_budget_tokens: params["thinking"] = { @@ -244,7 +254,8 @@ class AnthropicBackend: ) content: Any = text_content - if response_format is not None: + # Tool-call turns carry no consumable content + if response_format is not None and not tool_calls: raw_content = f"{{{text_content}" if prefilled_json else text_content try: if prefilled_json: diff --git a/src/llm/backends/gemini.py b/src/llm/backends/gemini.py index b14cefe4..c114196e 100644 --- a/src/llm/backends/gemini.py +++ b/src/llm/backends/gemini.py @@ -14,7 +14,8 @@ from src.llm.caching import ( build_cache_key, gemini_cache_store, ) -from src.llm.structured_output import repair_response_model_json +from src.llm.request_builder import coerce_passthrough_mapping +from src.llm.structured_output import repair_response_model_json, schema_instruction GEMINI_BLOCKED_FINISH_REASONS = { "SAFETY", @@ -30,6 +31,29 @@ class GeminiBackend: def __init__(self, client: Any) -> None: self._client: Any = client + @staticmethod + def _append_schema_instruction( + contents: list[dict[str, Any]] | str, + response_format: type[BaseModel], + ) -> list[dict[str, Any]] | str: + """Append the schema instruction to the final turn. + + Used when tools accompany a response_format: native response_schema + + function calling is a Gemini 3 preview feature and earlier models + reject the pairing, so instruct the model and rely on parse + repair. + Returns a new list — _convert_messages shallow-copies parts-style + messages, so in-place appends would leak into the caller's history + and accumulate across tool-loop iterations. + """ + instruction = schema_instruction(response_format, tools_present=True) + if isinstance(contents, str): + return contents + instruction + if not contents: + return contents + last = contents[-1] + parts: list[Any] = [*(last.get("parts") or []), {"text": instruction}] + return [*contents[:-1], {**last, "parts": parts}] + async def complete( self, *, @@ -60,6 +84,10 @@ class GeminiBackend: ) if system_instruction: config["system_instruction"] = system_instruction + if tools and isinstance(response_format, type): + # The final turn is never part of the cached prefix, so this is + # safe to do before cache attachment. + contents = self._append_schema_instruction(contents, response_format) cache_policy = ( extra_params.get("cache_policy") @@ -129,6 +157,10 @@ class GeminiBackend: ) if system_instruction: config["system_instruction"] = system_instruction + if tools and isinstance(response_format, type): + # The final turn is never part of the cached prefix, so this is + # safe to do before cache attachment. + contents = self._append_schema_instruction(contents, response_format) cache_policy = ( extra_params.get("cache_policy") @@ -227,7 +259,11 @@ class GeminiBackend: config["tools"] = self._convert_tools(tools) if tool_choice: config["tool_config"] = self._convert_tool_choice(tool_choice) - if response_format is not None: + # Native structured output combined with function calling is a + # Gemini 3 preview feature; earlier models reject the pairing. With + # tools present, callers inject a schema instruction instead (see + # _append_schema_instruction) and rely on parse + repair downstream. + if response_format is not None and not tools: config["response_mime_type"] = "application/json" config["response_schema"] = response_format elif extra_params and extra_params.get("json_mode") and not tools: @@ -246,6 +282,26 @@ class GeminiBackend: for key in ("top_p", "top_k", "frequency_penalty", "presence_penalty", "seed"): if extra_params and key in extra_params: config[key] = extra_params[key] + # Operator escape hatch: forward provider_params into the google-genai + # config dict. The Gemini SDK doesn't expose extra_body/extra_headers + # as kwargs (unlike OpenAI/Anthropic) — body-shaped fields live on + # GenerateContentConfig and headers live under config.http_options. + # extra_query has no SDK-level equivalent and is ignored. Shallow + # merge with operator-wins. Operators are responsible for not setting + # unknown fields that google-genai's validation will reject. + if extra_params: + operator_extra_body = extra_params.get("extra_body") + if operator_extra_body: + config.update( + coerce_passthrough_mapping("extra_body", operator_extra_body) + ) + operator_extra_headers = extra_params.get("extra_headers") + if operator_extra_headers: + http_options = config.setdefault("http_options", {}) + existing_headers = http_options.setdefault("headers", {}) + existing_headers.update( + coerce_passthrough_mapping("extra_headers", operator_extra_headers) + ) return config def _normalize_response( @@ -314,7 +370,10 @@ class GeminiBackend: ) content: Any = "\n".join(text_parts) if text_parts else "" - if response_format is not None: + # Tool-call turns carry no consumable content — the tool loop ignores + # it — and parsing their (empty) text would raise through the repair + # fallback, failing the iteration. + if response_format is not None and not tool_calls: parsed_response = getattr(response, "parsed", None) if isinstance(parsed_response, response_format): content = parsed_response diff --git a/src/llm/backends/openai.py b/src/llm/backends/openai.py index b2d82d91..d5d0ed73 100644 --- a/src/llm/backends/openai.py +++ b/src/llm/backends/openai.py @@ -2,6 +2,7 @@ from __future__ import annotations import json import logging +import weakref from collections.abc import AsyncIterator from typing import Any, cast @@ -10,7 +11,10 @@ from pydantic import BaseModel, ValidationError from src.exceptions import ValidationException from src.llm.backend import CompletionResult, StreamChunk, ToolCallResult +from src.llm.request_builder import apply_sdk_passthroughs from src.llm.structured_output import ( + StructuredOutputError, + empty_structured_output, repair_response_model_json, validate_structured_output, ) @@ -18,6 +22,37 @@ from src.llm.structured_output import ( logger = logging.getLogger(__name__) +# The point of this being a WeakKeyDictionary, as opposed to a regular dict, is that it +# does not hold a reference to the keyed BaseModel so that when a dynamically created +# type is no longer referenced it becomes eligible for garbage collection. This avoids a +# memory leak. +_json_object_instruction_cache: weakref.WeakKeyDictionary[type[BaseModel], str] = ( + weakref.WeakKeyDictionary() +) + + +def _json_object_instruction(response_format: type[BaseModel]) -> str: + """Schema-injection instruction for json_object mode. + + The JSON schema is static per response_format class, so cache the serialized + instruction — the deriver would otherwise re-walk the schema + re-serialize + it every call. + """ + cached = _json_object_instruction_cache.get(response_format) + if cached is not None: + return cached + # Some OpenAI-compatible providers enforce this JSON-object precondition with + # a case-sensitive substring check, so include lowercase "json" explicitly. + instruction = ( + "You must respond with a single JSON object (json) that conforms " + "exactly to the following JSON schema. Do not include any text, " + "markdown, or code fences outside the JSON object.\n\nJSON schema:\n" + f"{json.dumps(response_format.model_json_schema())}" + ) + _json_object_instruction_cache[response_format] = instruction + return instruction + + def _uses_max_completion_tokens(model: str) -> bool: """OpenAI reasoning models (gpt-5 family + o-series) require ``max_completion_tokens`` instead of the classic ``max_tokens`` parameter. @@ -142,10 +177,40 @@ class OpenAIBackend: ) if isinstance(response_format, type): + if self._structured_output_mode(extra_params) == "json_object": + self._apply_json_object_mode(params, response_format) + response = await self._client.chat.completions.create(**params) + # A loose provider that returns nothing shouldn't crash the call. + content = self._parse_or_repair_structured_content( + response, response_format, model, empty_on_missing=True + ) + return self._normalize_response(response, content_override=content) + if tools: + # parse() refuses non-strict function tools, and our agent tool + # schemas are deliberately non-strict (see _convert_tools), so + # tool-loop iterations use create() with an explicit json_schema + # response_format — same server-side schema enforcement, no + # strict-tools requirement — mirroring the streaming path. + params["response_format"] = self._json_schema_response_format( + response_format + ) + response = await self._client.chat.completions.create(**params) + # Tool-call turns carry no consumable content — the tool loop + # ignores it — and parsing their empty text would raise. + if getattr(response.choices[0].message, "tool_calls", None): + return self._normalize_response(response) + content = self._parse_or_repair_structured_content( + response, response_format, model, empty_on_missing=False + ) + return self._normalize_response(response, content_override=content) params["response_format"] = response_format try: response = await self._client.chat.completions.parse(**params) except LengthFinishReasonError as exc: + # Truncated output: repair the partial content directly. repair + # handles empty/unrepairable JSON with its own model-aware fallback + # (PromptRepresentation -> empty, others -> raise), which differs + # from the parse-fallback terminal below, so it stays a direct call. truncated = exc.completion raw_content = truncated.choices[0].message.content or "" content = repair_response_model_json( @@ -157,41 +222,42 @@ class OpenAIBackend: truncated, content_override=content, ) - except (BadRequestError, json.JSONDecodeError, ValidationError): - fallback_response = await self._create_structured_response( - params=params, - response_format=response_format, - ) - content = self._parse_or_repair_structured_content( - fallback_response, - response_format, + except BadRequestError: + # A 400 means the provider rejected the request shape — most + # often it doesn't support OpenAI Structured Outputs (json_schema). + # Retrying or re-requesting won't help (it rejects the same shape + # again, the latency trap of #797), so return empty rather than + # erroring existing flows. The warning is the signal to set + # structured_output_mode=json_object. There is no response body to + # account for, so token usage is legitimately zero here. + logger.warning( + "Structured output via json_schema rejected by model %s; " + + "set structured_output_mode=json_object if the provider does " + + "not support OpenAI Structured Outputs.", model, ) - return self._normalize_response( - fallback_response, - content_override=content, - ) + # empty_structured_output() validates {} against the model, which + # itself raises if the model has required fields. Fall back to + # empty string content rather than letting that escape the handler. + try: + fallback_content: Any = empty_structured_output(response_format) + except ValidationError: + fallback_content = "" + return CompletionResult(content=fallback_content) parsed = response.choices[0].message.parsed - raw_content = response.choices[0].message.content or "" - if parsed is None and raw_content: - content = repair_response_model_json( - raw_content, - response_format, - model, + if parsed is not None: + return self._normalize_response( + response, + content_override=validate_structured_output( + parsed, response_format + ), ) - return self._normalize_response(response, content_override=content) - if parsed is None: - refusal = getattr(response.choices[0].message, "refusal", None) - if refusal: - return self._normalize_response( - response, - content_override=refusal, - ) - raise ValidationException("No parsed content in structured response") - return self._normalize_response( - response, - content_override=validate_structured_output(parsed, response_format), + # parse() returned no model: repair raw content, surface a refusal, + # or raise so the retry/fallback chain engages on a junk response. + content = self._parse_or_repair_structured_content( + response, response_format, model, empty_on_missing=False ) + return self._normalize_response(response, content_override=content) if response_format is not None: params["response_format"] = response_format @@ -232,15 +298,16 @@ class OpenAIBackend: params["stream"] = True params["stream_options"] = {"include_usage": True} if isinstance(response_format, type): - # parse() supports BaseModel types but streaming create() does not — - # convert to a json_schema dict so the streaming path works. - params["response_format"] = { - "type": "json_schema", - "json_schema": { - "name": response_format.__name__, - "schema": response_format.model_json_schema(), - }, - } + if self._structured_output_mode(extra_params) == "json_object": + # Inject the schema into the prompt for providers without + # json_schema support; repair happens downstream. + self._apply_json_object_mode(params, response_format) + else: + # Streaming create() can't take a BaseModel like parse() does; + # convert to a json_schema dict. + params["response_format"] = self._json_schema_response_format( + response_format + ) elif response_format is not None: params["response_format"] = response_format elif extra_params and extra_params.get("json_mode"): @@ -300,8 +367,10 @@ class OpenAIBackend: # Token-budget style thinking is not part of the native OpenAI API, but # OpenAI-compatible proxies (OpenRouter, etc.) accept a `reasoning` object # on the request body. Pass through via extra_body so it reaches those - # backends; operators on providers that need a different shape (vLLM, - # Fireworks, ...) can override via ModelConfig.provider_params. + # backends. Operators on providers that need a different shape (e.g. + # Anthropic-via-Vertex behind litellm wants `thinking`, not `reasoning`) + # supply that shape via ModelConfig.provider_params.extra_body and unset + # thinking_budget_tokens themselves — Honcho does not try to translate. if thinking_budget_tokens is not None and thinking_budget_tokens > 0: params.setdefault("extra_body", {}).setdefault("reasoning", {})[ "max_tokens" @@ -311,8 +380,9 @@ class OpenAIBackend: params["stop"] = stop if tools: params["tools"] = self._convert_tools(tools) - if tool_choice is not None: - params["tool_choice"] = tool_choice + converted_tool_choice = self._convert_tool_choice(tool_choice) + if converted_tool_choice is not None: + params["tool_choice"] = converted_tool_choice if extra_params: for key in ( "top_p", @@ -322,6 +392,11 @@ class OpenAIBackend: ): if key in extra_params: params[key] = extra_params[key] + # Operator escape hatch: forward OpenAI SDK passthrough kwargs from + # ModelConfig.provider_params. Shallow merge with operator-wins — + # if the operator supplies `extra_body.reasoning`, it replaces any + # value Honcho auto-injected above. + apply_sdk_passthroughs(params, extra_params) return params def _normalize_response( @@ -374,37 +449,130 @@ class OpenAIBackend: raw_response=response, ) - async def _create_structured_response( - self, - *, - params: dict[str, Any], + @staticmethod + def _json_schema_response_format( response_format: type[BaseModel], - ) -> Any: - structured_params = dict(params) - structured_params["response_format"] = { + ) -> dict[str, Any]: + """Build the response_format param for create() calls that can't use + parse(): streaming, and requests carrying non-strict function tools.""" + return { "type": "json_schema", "json_schema": { "name": response_format.__name__, "schema": response_format.model_json_schema(), }, } - return await self._client.chat.completions.create(**structured_params) + + @staticmethod + def _structured_output_mode(extra_params: dict[str, Any] | None) -> str | None: + # Threaded in via extra_params (see build_config_extra_params). + if not extra_params: + return None + return extra_params.get("structured_output_mode") + + def _apply_json_object_mode( + self, + params: dict[str, Any], + response_format: type[BaseModel], + ) -> None: + """Configure params for json_object mode in place (shared by complete/stream). + + Injects the schema into the prompt and requests loose JSON, so providers + without OpenAI Structured Outputs (json_schema) support still return JSON. + """ + params["messages"] = self._with_json_schema_instructions( + params["messages"], response_format + ) + params["response_format"] = {"type": "json_object"} + + @staticmethod + def _with_json_schema_instructions( + messages: list[dict[str, Any]], + response_format: type[BaseModel], + ) -> list[dict[str, Any]]: + """Add JSON-schema instructions to a copy of messages for json_object mode. + + The Anthropic backend has its own schema-into-prompt injection + (``_append_text_to_last_message``); the two are intentionally kept + separate since the providers want different placement and wording. + """ + instruction = _json_object_instruction(response_format) + new_messages = [dict(message) for message in messages] + first = new_messages[0] if new_messages else None + # Only merge into a leading system message when its content is a plain + # string; non-string content (e.g. a list of content parts) would be + # corrupted by f-string coercion, so prepend a fresh system message. + if ( + first + and first.get("role") == "system" + and isinstance(first.get("content"), str) + ): + first["content"] = f"{first['content']}\n\n{instruction}".strip() + else: + new_messages.insert(0, {"role": "system", "content": instruction}) + return new_messages @staticmethod def _parse_or_repair_structured_content( response: Any, response_format: type[BaseModel], model: str, + *, + empty_on_missing: bool, ) -> BaseModel | str: - raw_content = response.choices[0].message.content or "" + """Validate (or repair) the raw structured content of a response. + + Shared by the json_object path and the json_schema parse() fallbacks + (truncation, parsed=None). On a contentless response with no refusal, + ``empty_on_missing`` selects the terminal behavior: json_object returns a + graceful empty so a loose provider can't crash the call, while json_schema + raises so the retry/fallback chain engages on a junk response. + """ + message = response.choices[0].message + raw_content = message.content or "" if raw_content: - return repair_response_model_json(raw_content, response_format, model) - refusal = getattr(response.choices[0].message, "refusal", None) + # Fast path: clean JSON validates directly. Only fall back to the + # repair pipeline when validation fails — repair is comparatively + # expensive and silently degrades malformed input to an empty model. + try: + return validate_structured_output(raw_content, response_format) + except (StructuredOutputError, ValidationError): + return repair_response_model_json(raw_content, response_format, model) + refusal = getattr(message, "refusal", None) if refusal: return refusal - raise ValidationException( - "No raw content available for structured output repair" - ) + if not empty_on_missing: + raise ValidationException("No parsed content in structured response") + # empty_structured_output() validates {} against the model, which itself + # raises if the model has required fields. Fall back to empty string + # content rather than letting that escape the handler. + try: + return empty_structured_output(response_format) + except ValidationError: + return "" + + @staticmethod + def _convert_tool_choice( + tool_choice: str | dict[str, Any] | None, + ) -> str | dict[str, Any] | None: + # Translate Honcho's canonical tool_choice vocabulary to OpenAI's. This + # mirrors the Anthropic/Gemini backends so a single TOOL_CHOICE value + # works regardless of which provider a fallback chain lands on. Notably + # OpenAI has no "any" — it spells the same intent "required". + if tool_choice is None: + return None + if isinstance(tool_choice, dict): + if "name" in tool_choice: + return { + "type": "function", + "function": {"name": tool_choice["name"]}, + } + return tool_choice + if tool_choice in {"any", "required"}: + return "required" + if tool_choice in {"auto", "none"}: + return tool_choice + return {"type": "function", "function": {"name": tool_choice}} @staticmethod def _convert_tools(tools: list[dict[str, Any]]) -> list[dict[str, Any]]: diff --git a/src/llm/capture.py b/src/llm/capture.py new file mode 100644 index 00000000..ff0d1b0e --- /dev/null +++ b/src/llm/capture.py @@ -0,0 +1,403 @@ +"""Structures for data captured from LLM calls via telemetry. + +All capture is best-effort: `dispatch_captured_call` swallows exporter exceptions +so telemetry can never break the LLM call path. +""" + +from __future__ import annotations + +import base64 +import contextlib +import hashlib +import json +import logging +from dataclasses import dataclass, field +from typing import Any, Protocol, cast, runtime_checkable + +from src.config import settings + +from .backend import CompletionResult as BackendCompletionResult +from .backend import ToolCallResult +from .types import LLMTelemetryContext + +logger = logging.getLogger(__name__) + +# Sentinel roles for non-message content stored in the shared content store so +# the same hash+dedup machinery covers them. They never collide with real +# conversation roles ("user"/"assistant"/"system"/"tool"). +ROLE_OUTPUT = "assistant" +ROLE_TOOL_SCHEMA = "__tool_schema__" +ROLE_THINKING = "__thinking__" + + +def canonical_json(obj: Any) -> str: + """Deterministic JSON encoding used for every content hash.""" + return json.dumps( + obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str + ) + + +def compute_content_hash( + role: str, + content: Any, + tool_call_id: str | None, + tool_calls: list[dict[str, Any]] | None = None, +) -> str: + """Content hash covering the FULL message identity, not just the text. + + Includes `tool_calls` so two assistant turns with identical (often empty) + content but different tool calls don't collide in the dedup store. + """ + digest = hashlib.sha256( + canonical_json( + { + "role": role, + "content": content, + "tool_call_id": tool_call_id, + "tool_calls": tool_calls or [], + } + ).encode("utf-8") + ).hexdigest() + return f"sha256:{digest}" + + +def clip_for_trace(content: Any) -> tuple[Any, bool]: + """Clip a content value to `TELEMETRY.TRACE_MAX_BYTES`, returning (content, truncated). + + Only oversized string content is clipped (with a marker); non-string + structured content is left intact. Returns the input unchanged when it + fits or when the cap is non-positive. + """ + max_bytes = settings.TELEMETRY.TRACE_MAX_BYTES + if max_bytes <= 0 or not isinstance(content, str): + return content, False + encoded = content.encode("utf-8") + if len(encoded) <= max_bytes: + return content, False + marker = "…[truncated]" + keep = max(0, max_bytes - len(marker.encode("utf-8"))) + clipped = encoded[:keep].decode("utf-8", errors="ignore") + marker + return clipped, True + + +@dataclass(slots=True) +class CapturedMessage: + """One input message, normalized to a provider-agnostic shape. + + `content` is the message text; `tool_calls` holds any tool calls in a + unified `{id, name, input}` shape regardless of provider. `content_hash` + covers all identity fields so the ref and the shipped `trace.content` agree. + """ + + role: str + content: Any + tool_call_id: str | None + content_hash: str + truncated: bool = False + tool_calls: list[dict[str, Any]] = field(default_factory=list) + + +@dataclass(slots=True) +class CapturedLLMCall: + """Everything one LLM call needs to be reconstructed, captured once.""" + + # Correlation (span tree) + trace_id: str | None + span_id: str | None + parent_span_id: str | None + iteration: int | None + step_seq: int + attempt: int + was_fallback: bool + run_id: str | None + # Path identity + workspace_name: str | None + call_purpose: str | None + parent_category: str | None + agent_type: str | None + # unique session ID for grouping traces + session_id: str | None + observer: str | None + observed: str | None + peer_name: str | None + track_name: str | None + transport: str + provider_label: str | None + model: str + # Context window + input_messages: list[CapturedMessage] + tool_schemas: list[dict[str, Any]] + tool_choice: Any + # Output (replay-grade) + output_content: Any + output_tool_calls: list[dict[str, Any]] + thinking_content: str | None + thinking_blocks: list[dict[str, Any]] + reasoning_details: list[dict[str, Any]] + finish_reason: str | None + # Accounting copy (so the trace stream stands alone) + input_tokens: int + output_tokens: int + cache_read_tokens: int + cache_creation_tokens: int + was_stream: bool + # True when any input message was clipped to TRACE_MAX_BYTES. + input_truncated: bool = False + + +def _normalize_message( + message: dict[str, Any], transport: str | None +) -> tuple[Any, str | None, list[dict[str, Any]]]: + """Normalize a provider-native message to (content, tool_call_id, tool_calls). + + Providers stash tool calls and results outside `content` (openai's + `tool_calls`, gemini's `parts`), so a naive `content` read loses them. This + lifts them into a unified shape: `content` becomes text, `tool_calls` is a + list of `{id, name, input}`, and tool results surface as `content` keyed by + `tool_call_id`. + """ + content: Any = message.get("content") + tool_call_id: str | None = message.get("tool_call_id") + tool_calls: list[dict[str, Any]] = [] + + if transport == "openai": + for tc in cast("list[dict[str, Any]]", message.get("tool_calls") or []): + fn = cast("dict[str, Any]", tc.get("function") or {}) + args = fn.get("arguments") + if isinstance(args, str): + with contextlib.suppress(json.JSONDecodeError): + args = json.loads(args) + tool_calls.append( + {"id": tc.get("id"), "name": fn.get("name"), "input": args} + ) + + elif transport == "gemini": + parts = message.get("parts") + if isinstance(parts, list): + texts: list[str] = [] + results: list[Any] = [] + for raw_part in cast("list[Any]", parts): + if not isinstance(raw_part, dict): + continue + part = cast("dict[str, Any]", raw_part) + text = part.get("text") + if isinstance(text, str): + texts.append(text) + elif "function_call" in part: + fc = cast("dict[str, Any]", part["function_call"] or {}) + tool_calls.append( + {"id": None, "name": fc.get("name"), "input": fc.get("args")} + ) + elif "function_response" in part: + fr = cast("dict[str, Any]", part["function_response"] or {}) + resp = fr.get("response") + if isinstance(resp, dict): + results.append(cast("dict[str, Any]", resp).get("result")) + else: + results.append(resp) + if tool_call_id is None: + tool_call_id = fr.get("name") + content = "\n".join(texts) if texts else (results[0] if results else None) + + elif transport == "anthropic" and isinstance(content, list): + texts = [] + for raw_block in cast("list[Any]", content): + if not isinstance(raw_block, dict): + continue + block = cast("dict[str, Any]", raw_block) + btype = block.get("type") + text = block.get("text") + if btype == "text" and isinstance(text, str): + texts.append(text) + elif btype == "tool_use": + tool_calls.append( + { + "id": block.get("id"), + "name": block.get("name"), + "input": block.get("input"), + } + ) + elif btype == "tool_result": + if tool_call_id is None: + tool_call_id = block.get("tool_use_id") + inner = block.get("content") + texts.append(inner if isinstance(inner, str) else canonical_json(inner)) + content = "\n".join(texts) if texts else None + + return content, tool_call_id, tool_calls + + +def build_captured_messages( + messages: list[dict[str, Any]], + memo: dict[int, CapturedMessage] | None, + transport: str | None = None, +) -> tuple[list[CapturedMessage], bool]: + """Create a list of CapturedMessage from LLM response messages. + + Conversation is append-only. Uses hashed message content to deduplicate + across turns. Messages are normalized per provider, then content is + truncated and hashed. + """ + captured: list[CapturedMessage] = [] + any_truncated = False + for message in messages: + key = id(message) + cached = memo.get(key) if memo is not None else None + if cached is not None: + captured.append(cached) + any_truncated = any_truncated or cached.truncated + continue + role = str(message.get("role", "")) + raw_content, tool_call_id, tool_calls = _normalize_message(message, transport) + content, truncated = clip_for_trace(raw_content) + any_truncated = any_truncated or truncated + captured_message = CapturedMessage( + role=role, + content=content, + tool_call_id=tool_call_id, + content_hash=compute_content_hash(role, content, tool_call_id, tool_calls), + truncated=truncated, + tool_calls=tool_calls, + ) + if memo is not None: + memo[key] = captured_message + captured.append(captured_message) + return captured, any_truncated + + +def build_captured_call( + *, + telemetry: LLMTelemetryContext | None, + transport: str, + provider_label: str | None, + model: str, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None, + tool_choice: Any, + result: BackendCompletionResult | None, + attempt: int, + was_fallback: bool, + was_stream: bool, + finish_reason: str | None, +) -> CapturedLLMCall: + """Assemble a `CapturedLLMCall` from telemetry + the provider result.""" + memo = telemetry.hash_memo if telemetry is not None else None + captured_messages, input_truncated = build_captured_messages( + messages, memo, transport + ) + + output_tool_calls = [ + _tool_call_to_dict(tc) for tc in (result.tool_calls if result else []) + ] + + return CapturedLLMCall( + trace_id=telemetry.trace_id if telemetry else None, + span_id=telemetry.span_id if telemetry else None, + parent_span_id=telemetry.exported_parent_span_id() if telemetry else None, + iteration=telemetry.iteration if telemetry else None, + step_seq=telemetry.step_seq if telemetry else 0, + attempt=attempt, + was_fallback=was_fallback, + run_id=telemetry.run_id if telemetry else None, + workspace_name=telemetry.workspace_name if telemetry else None, + call_purpose=telemetry.call_purpose if telemetry else None, + parent_category=telemetry.parent_category if telemetry else None, + agent_type=telemetry.agent_type if telemetry else None, + session_id=telemetry.session_id if telemetry else None, + observer=telemetry.observer if telemetry else None, + observed=telemetry.observed if telemetry else None, + peer_name=telemetry.peer_name if telemetry else None, + track_name=telemetry.track_name if telemetry else None, + transport=transport, + provider_label=provider_label, + model=model, + input_messages=captured_messages, + tool_schemas=list(tools) if tools else [], + tool_choice=tool_choice, + output_content=result.content if result else None, + output_tool_calls=output_tool_calls, + thinking_content=result.thinking_content if result else None, + thinking_blocks=result.thinking_blocks if result else [], + reasoning_details=result.reasoning_details if result else [], + finish_reason=finish_reason, + input_tokens=result.input_tokens if result else 0, + output_tokens=result.output_tokens if result else 0, + cache_read_tokens=result.cache_read_input_tokens if result else 0, + cache_creation_tokens=result.cache_creation_input_tokens if result else 0, + was_stream=was_stream, + input_truncated=input_truncated, + ) + + +def _tool_call_to_dict(tool_call: ToolCallResult) -> dict[str, Any]: + """Normalize a ToolCallResult to a JSON-safe dict for the trace stream. + + `thought_signature` arrives as raw bytes from Gemini; base64-encode it so + CloudEvents JSON serialization can't choke on non-UTF8 bytes (which would + silently drop the whole event via the best-effort emit path). + """ + out: dict[str, Any] = { + "id": tool_call.id, + "name": tool_call.name, + "input": tool_call.input, + } + sig = tool_call.thought_signature + if sig is not None: + out["thought_signature"] = ( + base64.b64encode(sig).decode("ascii") if isinstance(sig, bytes) else sig + ) + return out + + +@runtime_checkable +class LLMCallExporter(Protocol): + """A sink that consumes a `CapturedLLMCall`""" + + def export(self, call: CapturedLLMCall) -> None: ... + + +_EXPORTERS: list[LLMCallExporter] = [] + + +def register_exporter(exporter: LLMCallExporter) -> None: + """Register an exporter (idempotent on identity). Called at startup.""" + if exporter not in _EXPORTERS: + _EXPORTERS.append(exporter) + + +def clear_exporters() -> None: + """Drop all exporters — used on shutdown and in tests.""" + _EXPORTERS.clear() + + +def has_exporters() -> bool: + """True when at least one exporter is registered.""" + return bool(_EXPORTERS) + + +def dispatch_captured_call(call: CapturedLLMCall) -> None: + """Fan a captured call out to every exporter.""" + for exporter in _EXPORTERS: + try: + exporter.export(call) + except Exception: # pragma: no cover - best-effort telemetry + logger.debug("LLM call exporter failed", exc_info=True) + + +__all__ = [ + "ROLE_OUTPUT", + "ROLE_THINKING", + "ROLE_TOOL_SCHEMA", + "CapturedLLMCall", + "CapturedMessage", + "LLMCallExporter", + "build_captured_call", + "build_captured_messages", + "canonical_json", + "clear_exporters", + "clip_for_trace", + "compute_content_hash", + "dispatch_captured_call", + "has_exporters", + "register_exporter", +] diff --git a/src/llm/executor.py b/src/llm/executor.py index 87db8dd0..956bc673 100644 --- a/src/llm/executor.py +++ b/src/llm/executor.py @@ -19,14 +19,21 @@ from typing import Any, Literal, TypeVar, overload from pydantic import BaseModel -from src.config import ModelConfig, ModelTransport +from src.config import ModelConfig, ModelTransport, settings +from src.telemetry.logging import conditional_observe from .backend import CompletionResult as BackendCompletionResult from .backend import StreamChunk as BackendStreamChunk from .backend import ToolCallResult +from .capture import build_captured_call, dispatch_captured_call, has_exporters from .registry import CLIENTS, backend_for_provider from .request_builder import execute_completion, execute_stream -from .runtime import AttemptPlan, effective_config_for_call +from .runtime import ( + AttemptPlan, + annotate_current_generation_io, + annotate_current_langfuse_trace, + effective_config_for_call, +) from .types import ( HonchoLLMCallResponse, HonchoLLMCallStreamChunk, @@ -40,6 +47,82 @@ logger = logging.getLogger(__name__) M = TypeVar("M", bound=BaseModel) +# ModelConfig fields that must NEVER reach a trace: secrets and nested holders +# of secrets. Everything else on the config is a safe tuning knob and is dumped +# automatically — so new knobs get traced without touching this code. Keep this +# a deny-list (small, stable) rather than an allow-list (drifts with the model). +_UNSAFE_CONFIG_FIELDS = frozenset( + { + "api_key", # provider secret + "base_url", # may embed credentials / private host + "fallback", # ResolvedFallbackConfig carries its own api_key/base_url + "provider_params", # opaque dict; can carry auth headers/keys + } +) + + +def _langfuse_model_parameters( + *, + max_tokens: int, + config: ModelConfig, + json_mode: bool, + verbosity: str | None, + stream: bool, + tools: list[dict[str, Any]] | None, + tool_choice: str | dict[str, Any] | None, + response_model: type[BaseModel] | None, +) -> dict[str, Any]: + """Serializable tuning knobs for the Langfuse generation. + + Surfaces everything @observe auto-capture used to show (temperature, tools, + ...) MINUS the live client and the secret-bearing config fields. We dump the + resolved `effective_config` and deny-list only `_UNSAFE_CONFIG_FIELDS`, so a + new ModelConfig knob is traced automatically — no allow-list to keep in sync. + `mode="json"` coerces enums/sub-models to JSON-safe values. See HONCHO-4HA. + """ + params: dict[str, Any] = config.model_dump( + exclude=set(_UNSAFE_CONFIG_FIELDS), exclude_none=True, mode="json" + ) + # Per-call extras that live outside ModelConfig. + params["max_tokens"] = max_tokens + params["stream"] = stream + params["json_mode"] = json_mode + if verbosity is not None: + params["verbosity"] = verbosity + if response_model is not None: + params["response_format"] = response_model.__name__ + if tools: + params["tools"] = [ + t.get("name") or t.get("function", {}).get("name") or "unknown" + for t in tools + ] + if tool_choice is not None: + params["tool_choice"] = ( + tool_choice if isinstance(tool_choice, str) else str(tool_choice) + ) + return params + + +def _langfuse_usage_details(response: HonchoLLMCallResponse[Any]) -> dict[str, int]: + """Token usage duplicated onto the Langfuse generation. + + These counts are also emitted via CloudEvents (LLMCallCompletedEvent), but + we mirror them here so Langfuse renders per-call tokens + cost natively, + including Anthropic-style prompt-cache reads/writes. Zero-valued cache keys + are dropped so non-cached calls stay tidy. Stream calls don't surface token + totals at this layer, so usage is set only on the non-stream path. + """ + usage: dict[str, int] = { + "input": response.input_tokens, + "output": response.output_tokens, + } + if response.cache_read_input_tokens: + usage["cache_read_input_tokens"] = response.cache_read_input_tokens + if response.cache_creation_input_tokens: + usage["cache_creation_input_tokens"] = response.cache_creation_input_tokens + return usage + + def _outcome_from_error( err: BaseException | None, ) -> Literal["success", "error", "cancelled"]: @@ -56,7 +139,7 @@ def _outcome_from_error( def _tool_call_result_to_dict(tool_call: ToolCallResult) -> dict[str, Any]: - result = { + result: dict[str, Any] = { "id": tool_call.id, "name": tool_call.name, "input": tool_call.input, @@ -107,7 +190,7 @@ def _emit_llm_call_completed( call_purpose=call_purpose, parent_category=(telemetry.parent_category if telemetry else None), transport=provider, - provider_label=_infer_provider_label(provider, model, plan), + provider_label=infer_provider_label(provider, model, plan), model=model, effective_max_output_tokens=max_tokens, provider_input_tokens=(result.input_tokens if result else 0), @@ -135,7 +218,7 @@ def _emit_llm_call_completed( logger.debug("Failed to emit LLMCallCompletedEvent", exc_info=True) -def _infer_provider_label( +def infer_provider_label( _transport: ModelTransport, model: str, plan: AttemptPlan | None ) -> str | None: """Best-effort vendor inference for relay setups. @@ -161,13 +244,56 @@ def _infer_provider_label( return None +def _maybe_dispatch_capture( + *, + plan: AttemptPlan | None, + telemetry: LLMTelemetryContext | None, + provider: ModelTransport, + model: str, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None, + tool_choice: Any, + result: BackendCompletionResult | None, + error: BaseException | None, +) -> None: + """Build a CapturedLLMCall and fan it out to registered exporters. + + No-op when payload capture is off + `has_exporters()` is checked BEFORE building. + Best-effort: never raises into the call path. + """ + if not has_exporters(): + return + try: + outcome = _outcome_from_error(error) + finish_reason = result.finish_reason if result is not None else outcome + dispatch_captured_call( + build_captured_call( + telemetry=telemetry, + transport=str(provider), + provider_label=infer_provider_label(provider, model, plan), + model=model, + messages=messages, + tools=tools, + tool_choice=tool_choice, + result=result, + attempt=plan.attempt if plan is not None else 1, + was_fallback=plan.is_fallback if plan is not None else False, + was_stream=False, + finish_reason=finish_reason, + ) + ) + except Exception: # pragma: no cover - best-effort telemetry + logger.debug("Failed to dispatch CapturedLLMCall", exc_info=True) + + def completion_result_to_response( result: BackendCompletionResult, ) -> HonchoLLMCallResponse[Any]: return HonchoLLMCallResponse( content=result.content, input_tokens=result.input_tokens, - output_tokens=result.output_tokens, + output_tokens=result.output_tokens or 0, cache_creation_input_tokens=result.cache_creation_input_tokens, cache_read_input_tokens=result.cache_read_input_tokens, finish_reasons=[result.finish_reason] if result.finish_reason else [], @@ -261,6 +387,18 @@ async def honcho_llm_call_inner( ) -> AsyncIterator[HonchoLLMCallStreamChunk]: ... +@conditional_observe( + name="LLM Call", + as_type="generation", + # Disable @observe auto-capture: it would serialize `client_override` (a + # live AsyncOpenAI/genai client) and `selected_config` (carries api_key) + # into the span. Auto-capture deep-copies the client into a half-built + # object whose teardown raises `_state`/`_http_options` AttributeErrors + # (HONCHO-4HA) and leaks the key. We set curated input/output explicitly + # below via `annotate_current_generation_io`, preserving full fidelity. + capture_input=False, + capture_output=False, +) async def honcho_llm_call_inner( provider: ModelTransport, model: str, @@ -284,6 +422,11 @@ async def honcho_llm_call_inner( ) -> HonchoLLMCallResponse[Any] | AsyncIterator[HonchoLLMCallStreamChunk]: """One backend call. No retry, no fallback, no tool loop. + This is the Langfuse trace boundary (``@conditional_observe``): every + provider call is its own trace. Multi-turn agents thread a shared + ``run_id`` through ``telemetry`` so their per-iteration traces roll up into + one Langfuse session (see ``annotate_current_langfuse_trace``). + The outer src/llm/api.py `honcho_llm_call` handles retry + fallback + tool orchestration on top of this. @@ -300,6 +443,12 @@ async def honcho_llm_call_inner( if client is None: raise ValueError(f"Missing client for {provider}") + # Stamp this trace (user_id/session_id/metadata) now that the @observe + # span is open and the resolved provider/model are known. Set early so the + # annotation lands even on the stream path, where the span closes once the + # generator is returned (before chunks drain). + annotate_current_langfuse_trace(provider, model, telemetry=telemetry) + if messages is None: messages = [{"role": "user", "content": prompt}] @@ -314,6 +463,27 @@ async def honcho_llm_call_inner( thinking_budget_tokens=thinking_budget_tokens, reasoning_effort=reasoning_effort, ) + + # Explicit generation input + tuning knobs (replaces @observe auto-capture, + # which would serialize the live client / api key). Set before the stream + # branch so it lands on the generation span for both paths. Guard on inline + # mode (matching annotate_current_generation_io's own gate) so we don't + # build the (model_dump-backed) payload when the helper would no-op — in + # exporter mode there's no active generation span to stamp. + if settings.langfuse_inline_enabled: + annotate_current_generation_io( + input=messages, + model_parameters=_langfuse_model_parameters( + max_tokens=max_tokens, + config=effective_config, + json_mode=json_mode, + verbosity=verbosity, + stream=stream, + tools=tools, + tool_choice=tool_choice, + response_model=response_model, + ), + ) # json_mode + verbosity are per-call transport toggles, not ModelConfig # knobs — they pass through extra_params. execute_completion merges # build_config_extra_params(effective_config) on top for top_p/seed/etc. @@ -399,7 +569,16 @@ async def honcho_llm_call_inner( cache_policy=effective_config.cache_policy, extra_params=call_extras, ) - return completion_result_to_response(backend_result) + response = completion_result_to_response(backend_result) + # Explicit generation output + token usage (replaces @observe + # auto-capture). The stream path closes this span before drain, so its + # output is stamped on the run-level span instead + if settings.langfuse_inline_enabled: + annotate_current_generation_io( + output=response, + usage_details=_langfuse_usage_details(response), + ) + return response except BaseException as exc: error = exc raise @@ -417,6 +596,17 @@ async def honcho_llm_call_inner( result=backend_result, error=error, ) + _maybe_dispatch_capture( + plan=plan, + telemetry=telemetry, + provider=provider, + model=model, + messages=messages, + tools=tools, + tool_choice=tool_choice, + result=backend_result, + error=error, + ) __all__ = [ diff --git a/src/llm/registry.py b/src/llm/registry.py index d89f088d..6dd5eecc 100644 --- a/src/llm/registry.py +++ b/src/llm/registry.py @@ -32,6 +32,27 @@ from .history_adapters import ( ) from .types import ProviderClient +# Client-level ``default_headers`` applied to OpenAI-compatible clients, keyed by +# base-URL prefix. Currently only OpenRouter, which uses them for app attribution +# (https://openrouter.ai/docs/app-attribution); add a prefix here to tag another +# provider. Other OpenAI-compatible backends ignore unrecognized headers. +_DEFAULT_HEADERS_BY_BASE_URL: dict[str, dict[str, str]] = { + "https://openrouter.ai": { + "HTTP-Referer": "https://honcho.dev", + "X-Openrouter-Title": "Honcho", + }, +} + + +def _default_headers_for(base_url: str | None) -> dict[str, str]: + """Default headers for ``base_url`` (prefix match); these merge under any + per-request ``extra_headers`` passthrough, which wins on key collision.""" + if base_url: + for prefix, headers in _DEFAULT_HEADERS_BY_BASE_URL.items(): + if base_url.startswith(prefix): + return headers + return {} + @lru_cache(maxsize=1) def get_anthropic_client() -> AsyncAnthropic: @@ -49,6 +70,7 @@ def get_openai_client() -> AsyncOpenAI: return AsyncOpenAI( api_key=settings.LLM.OPENAI_API_KEY, base_url=settings.LLM.OPENAI_BASE_URL, + default_headers=_default_headers_for(settings.LLM.OPENAI_BASE_URL), ) @@ -70,7 +92,11 @@ def get_openai_override_client( base_url: str | None, api_key: str | None ) -> AsyncOpenAI: """OpenAI client for a specific (base_url, api_key) pair. Cached by key.""" - return AsyncOpenAI(api_key=api_key, base_url=base_url) + return AsyncOpenAI( + api_key=api_key, + base_url=base_url, + default_headers=_default_headers_for(base_url), + ) @lru_cache(maxsize=128) @@ -106,6 +132,7 @@ if settings.LLM.OPENAI_API_KEY: CLIENTS["openai"] = AsyncOpenAI( api_key=settings.LLM.OPENAI_API_KEY, base_url=settings.LLM.OPENAI_BASE_URL, + default_headers=_default_headers_for(settings.LLM.OPENAI_BASE_URL), ) if settings.LLM.GEMINI_API_KEY: diff --git a/src/llm/request_builder.py b/src/llm/request_builder.py index d6be5a22..3da437b5 100644 --- a/src/llm/request_builder.py +++ b/src/llm/request_builder.py @@ -7,14 +7,70 @@ src/llm/api.py, src/llm/tool_loop.py, src/llm/runtime.py. from __future__ import annotations from collections.abc import AsyncIterator -from typing import Any +from typing import Any, cast from pydantic import BaseModel from src.config import ModelConfig, PromptCachePolicy +from src.exceptions import ValidationException from .backend import CompletionResult, ProviderBackend, StreamChunk +# Operator escape-hatch keys recognized inside ModelConfig.provider_params. +PASSTHROUGH_KEYS = ("extra_body", "extra_headers", "extra_query") + + +def coerce_passthrough_mapping(key: str, value: Any) -> dict[str, Any]: + """Validate an operator-supplied provider_params passthrough is a mapping. + + ``provider_params`` is typed ``dict[str, Any]`` with no nested schema, so an + operator can supply a non-mapping (e.g. a list or string) for one of the + passthrough keys. Catch that here with a clear error instead of letting a + later ``dict.update()`` raise an opaque ``TypeError`` deep in the transport. + + Args: + key: The passthrough key name, used only for the error message. + value: The operator-supplied value to validate. + + Returns: + The value, narrowed to ``dict[str, Any]``. + + Raises: + ValidationException: If ``value`` is not a mapping. + """ + if not isinstance(value, dict): + raise ValidationException( + f"provider_params.{key} must be a mapping, got {type(value).__name__}" + ) + return cast(dict[str, Any], value) + + +def apply_sdk_passthroughs( + params: dict[str, Any], extra_params: dict[str, Any] +) -> None: + """Forward operator provider_params passthroughs onto an SDK call dict. + + OpenAI and Anthropic both accept ``extra_body`` / ``extra_headers`` / + ``extra_query`` as identically-named SDK kwargs, so they share this merge. + Operator values shallow-merge onto ``params`` in place, winning over any + value Honcho already set under the same top-level key (e.g. an auto-injected + ``extra_body.reasoning``). Gemini handles passthroughs separately because the + google-genai SDK does not expose these as kwargs. + + Args: + params: The SDK call kwargs being assembled; mutated in place. + extra_params: Flattened per-call params (see build_config_extra_params). + + Raises: + ValidationException: If a passthrough value is not a mapping. + """ + for passthrough_key in PASSTHROUGH_KEYS: + operator_value = extra_params.get(passthrough_key) + if not operator_value: + continue + existing = params.setdefault(passthrough_key, {}) + existing.update(coerce_passthrough_mapping(passthrough_key, operator_value)) + def build_config_extra_params(config: ModelConfig) -> dict[str, Any]: """Flatten ModelConfig's optional knobs and provider_params into extra_params. @@ -34,6 +90,8 @@ def build_config_extra_params(config: ModelConfig) -> dict[str, Any]: extra_params["presence_penalty"] = config.presence_penalty if config.seed is not None: extra_params["seed"] = config.seed + if config.structured_output_mode is not None: + extra_params["structured_output_mode"] = config.structured_output_mode if config.provider_params: extra_params.update(config.provider_params) diff --git a/src/llm/runtime.py b/src/llm/runtime.py index ae551378..93e961c9 100644 --- a/src/llm/runtime.py +++ b/src/llm/runtime.py @@ -12,8 +12,9 @@ Owns: from __future__ import annotations import logging +from contextlib import ExitStack from contextvars import ContextVar -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any from src.config import ( @@ -25,7 +26,7 @@ from src.config import ( ) from .registry import backend_for_provider, client_for_model_config -from .types import ProviderClient, ReasoningEffortType +from .types import LLMTelemetryContext, ProviderClient, ReasoningEffortType logger = logging.getLogger(__name__) @@ -33,31 +34,316 @@ logger = logging.getLogger(__name__) current_attempt: ContextVar[int] = ContextVar("current_attempt", default=0) -def update_current_langfuse_observation( +def annotate_current_langfuse_trace( provider: ModelTransport, model: str, *, - name: str | None = None, + telemetry: LLMTelemetryContext | None = None, ) -> None: - """Best-effort annotation of the current Langfuse span with LLM routing.""" - if not settings.LANGFUSE_PUBLIC_KEY: + """Stamp provider/model + step metadata on the current Langfuse generation. + + Inside an active agent run, `propagate_attributes` already stamped + user_id/session_id/trace_name on the run span; this call only needs to + decorate the per-iteration generation. Outside a run (single-shot + callers — deriver, summarizer), this generation IS the trace root, so we + also stamp the trace attrs. + + `model`/`metadata` are set on every call regardless of `inside_run`, so + every multi-turn iteration carries provider/model attribution. + """ + if not settings.langfuse_inline_enabled: return + try: + from langfuse import get_client, propagate_attributes + + inside_run = telemetry is not None and telemetry.parent_span_id is not None + gen_metadata = _step_metadata(telemetry) if telemetry is not None else {} + gen_metadata["provider"] = str(provider) + gen_metadata["model"] = str(model) + gen_name = ( + f"{telemetry.track_name} LLM call" + if telemetry is not None and telemetry.track_name + else None + ) + + if not inside_run: + session_id = telemetry.span_identity() if telemetry is not None else None + trace_name = telemetry.track_name if telemetry is not None else None + trace_metadata: dict[str, str] = dict(gen_metadata) + if telemetry is None: + trace_metadata.setdefault("namespace", str(settings.NAMESPACE)) + # Empty body is intentional: propagate_attributes stamps the active + # @observe generation (this trace root, for single-shot callers) at + # __enter__; there are no child spans to scope here. Don't delete as + # dead code — the enter-time side effect is the point. + with propagate_attributes( + user_id=str(settings.NAMESPACE), + session_id=session_id, + trace_name=trace_name, + metadata=trace_metadata, + ): + pass + + get_client().update_current_generation( + name=gen_name, + model=str(model), + metadata=gen_metadata, + ) + except Exception as exc: # pragma: no cover - best-effort telemetry + logger.debug("Failed to update Langfuse trace metadata: %s", exc) + + +def annotate_current_generation_io( + *, + input: Any = None, # noqa: A002 - mirrors langfuse's `input` kwarg name + output: Any = None, + model_parameters: dict[str, Any] | None = None, + usage_details: dict[str, Any] | None = None, +) -> None: + """Set explicit input/output/model_parameters/usage on the current generation. + + Used in place of ``@observe``'s auto-capture (disabled on + ``honcho_llm_call_inner``) so the provider client and api-key-bearing + ``ModelConfig`` arguments are never serialized into traces. Auto-capture + deep-copies those args, producing half-constructed clients whose teardown + raised ``AsyncHttpxClientWrapper ... no attribute '_state'`` / + ``BaseApiClient ... no attribute '_http_options'`` (HONCHO-4HA) and leaked + ``ModelConfig.api_key``. We instead hand Langfuse curated, serializable + values: ``messages`` in, response out, and the call's tuning knobs as + ``model_parameters`` — preserving (and tidying) full trace fidelity. + + Best-effort: telemetry must never fail the LLM call. + """ + # Gated on inline mode (NOT just key presence): this writes to the *active* + # @observe generation span, which only exists in inline mode. In exporter + # mode `conditional_observe` applies no decorator. + if not settings.langfuse_inline_enabled: + return + payload: dict[str, Any] = {} + if input is not None: + payload["input"] = input + if output is not None: + payload["output"] = output + if model_parameters: + payload["model_parameters"] = model_parameters + if usage_details: + payload["usage_details"] = usage_details + if not payload: + return try: from langfuse import get_client - update_kwargs: dict[str, Any] = { - "metadata": { - "namespace": settings.NAMESPACE, - "provider": provider, - "model": model, - } - } - if name is not None: - update_kwargs["name"] = name - get_client().update_current_span(**update_kwargs) + get_client().update_current_generation(**payload) except Exception as exc: # pragma: no cover - best-effort telemetry - logger.debug("Failed to update Langfuse span metadata: %s", exc) + logger.debug("Failed to set Langfuse generation IO: %s", exc) + + +def _base_metadata(telemetry: LLMTelemetryContext) -> dict[str, str]: + """Static routing/attribution metadata (everything except ``iteration``). + + Rebuilt per run (cheap); callers that need ``iteration`` copy and add it. + """ + metadata: dict[str, str] = {"namespace": str(settings.NAMESPACE)} + for key, value in ( + ("workspace_name", telemetry.workspace_name), + ("call_purpose", telemetry.call_purpose), + ("agent_type", telemetry.agent_type), + ("observer", telemetry.observer), + ("observed", telemetry.observed), + ("peer_name", telemetry.peer_name), + ("trace_id", telemetry.trace_id), + ("span_id", telemetry.span_id), + ("parent_span_id", telemetry.exported_parent_span_id()), + ): + if value is not None: + metadata[key] = str(value) + return metadata + + +def _step_metadata( + telemetry: LLMTelemetryContext, + base: dict[str, str] | None = None, +) -> dict[str, str]: + """Per-step metadata: ``base`` (or freshly computed) plus the per-step + ``iteration`` / ``step_seq`` / ``attempt`` counters.""" + metadata = dict(base) if base is not None else _base_metadata(telemetry) + if telemetry.iteration is not None: + metadata["iteration"] = str(telemetry.iteration) + metadata["step_seq"] = str(telemetry.step_seq) + metadata["attempt"] = str(telemetry.attempt) + return metadata + + +@dataclass +class LangfuseAgentRun: + """Imperative handle for the run-level Langfuse span. + + Owns an ``ExitStack`` that keeps ``start_as_current_observation`` and + ``propagate_attributes`` open until ``.end()``. This lets the run span + outlive the function that created it — streaming flows transfer the + handle to the response wrapper, which calls ``.end(output=...)`` after + the stream drains. While the handle is alive, the run span is the + current OTel observation, so step spans and auto-instrumented LLM + generations nest under it without any ContextVar choreography. + + Use ``start_langfuse_agent_run`` to construct; never instantiate directly. + """ + + span: Any # LangfuseSpan; opaque to keep src/llm/ free of langfuse imports. + _stack: ExitStack + _ended: bool = field(default=False) + + def update(self, **kwargs: Any) -> None: + """Set input/output/metadata on the run span (best-effort, no-op if ended).""" + if self._ended or self.span is None: + return + try: + self.span.update(**kwargs) + except Exception as exc: # pragma: no cover - best-effort telemetry + logger.debug("Failed to update Langfuse run span: %s", exc) + + def end(self, *, output: Any = None) -> None: + """Stamp final output (optional) and close the run span. Idempotent.""" + if self._ended: + return + self._ended = True + try: + if self.span is not None and output is not None: + self.span.update(output=output) + except Exception as exc: # pragma: no cover - best-effort telemetry + logger.debug("Failed to set Langfuse run output: %s", exc) + try: + self._stack.close() + except Exception as exc: # pragma: no cover - best-effort telemetry + logger.debug("Failed to close Langfuse run span: %s", exc) + + +def start_langfuse_agent_run( + name: str, telemetry: LLMTelemetryContext | None +) -> LangfuseAgentRun | None: + """Open the one run-level Langfuse trace per agentic run, imperatively. + + Returns ``None`` when Langfuse is disabled or there's no span identity + (single-shot callers without a ``span_id``/``run_id`` — those self-stamp + via ``annotate_current_langfuse_trace``). When non-None, the caller MUST + eventually call ``.end()`` — typically in a ``finally`` block, or by + transferring ownership to the streaming wrapper. + """ + if not settings.langfuse_inline_enabled or telemetry is None: + return None + session_id = telemetry.span_identity() + if not session_id: + return None + stack = ExitStack() + try: + from langfuse import get_client, propagate_attributes + + span = stack.enter_context( + get_client().start_as_current_observation(as_type="span", name=name) + ) + stack.enter_context( + propagate_attributes( + user_id=str(settings.NAMESPACE), + session_id=session_id, + trace_name=name, + metadata=_base_metadata(telemetry), + ) + ) + except Exception as exc: # pragma: no cover - best-effort telemetry + logger.debug("Failed to open Langfuse agent run: %s", exc) + stack.close() + return None + + return LangfuseAgentRun(span=span, _stack=stack) + + +@dataclass +class LangfuseAgentStep: + """Imperative handle for a per-iteration step span under the run root. + + Owns an ``ExitStack`` holding ``start_as_current_observation`` open until + ``.end()``. While alive the step span is the current OTel observation, + so the LLM generation (auto-instrumented or otherwise) nests under it. + No trace attrs (the run root carries them); just the per-step + ``iteration`` metadata. + """ + + span: Any + _stack: ExitStack + _ended: bool = field(default=False) + + def update(self, **kwargs: Any) -> None: + """Set input/output/metadata on the step span (best-effort, no-op if ended).""" + if self._ended or self.span is None: + return + try: + self.span.update(**kwargs) + except Exception as exc: # pragma: no cover - best-effort telemetry + logger.debug("Failed to update Langfuse step span: %s", exc) + + def annotate_io( + self, + messages: list[dict[str, Any]], + content: Any, + tool_calls: list[dict[str, Any]], + ) -> None: + """Stamp this turn's messages-in / content-or-tool-summary-out. + + On a tool-calling turn the model returns no text yet, so we summarize + the tool calls for the step output preview; otherwise the assistant + text is used. + """ + if self._ended or self.span is None: + return + if isinstance(content, str) and content.strip(): + output: Any = content + elif tool_calls: + output = {"tool_calls": [tc.get("name") for tc in tool_calls]} + else: + output = content + self.update(input=messages, output=output) + + def end(self, *, output: Any = None) -> None: + """Stamp final output (optional) and close the step span. Idempotent.""" + if self._ended: + return + self._ended = True + try: + if self.span is not None and output is not None: + self.span.update(output=output) + except Exception as exc: # pragma: no cover - best-effort telemetry + logger.debug("Failed to set Langfuse step output: %s", exc) + try: + self._stack.close() + except Exception as exc: # pragma: no cover - best-effort telemetry + logger.debug("Failed to close Langfuse step span: %s", exc) + + +def start_langfuse_agent_step( + name: str, telemetry: LLMTelemetryContext | None +) -> LangfuseAgentStep | None: + """Open a per-iteration step span, imperatively. Returns ``None`` when + Langfuse is disabled or there's no span identity (no agent run to nest under). + """ + if not settings.langfuse_inline_enabled or telemetry is None: + return None + if not telemetry.span_identity(): + return None + stack = ExitStack() + try: + from langfuse import get_client + + span = stack.enter_context( + get_client().start_as_current_observation( + as_type="span", name=name, metadata=_step_metadata(telemetry) + ) + ) + except Exception as exc: # pragma: no cover - best-effort telemetry + logger.debug("Failed to open Langfuse agent step: %s", exc) + stack.close() + return None + return LangfuseAgentStep(span=span, _stack=stack) @dataclass(frozen=True) @@ -118,6 +404,7 @@ def select_model_config_for_attempt( seed=fb.seed, thinking_effort=fb.thinking_effort, thinking_budget_tokens=fb.thinking_budget_tokens, + structured_output_mode=fb.structured_output_mode, provider_params=fb.provider_params, max_output_tokens=fb.max_output_tokens, stop_sequences=fb.stop_sequences, @@ -231,6 +518,10 @@ def resolve_backend_for_plan(plan: AttemptPlan) -> Any: __all__ = [ "AttemptPlan", + "LangfuseAgentRun", + "LangfuseAgentStep", + "annotate_current_generation_io", + "annotate_current_langfuse_trace", "current_attempt", "effective_config_for_call", "effective_temperature", @@ -238,5 +529,6 @@ __all__ = [ "resolve_backend_for_plan", "resolve_runtime_model_config", "select_model_config_for_attempt", - "update_current_langfuse_observation", + "start_langfuse_agent_run", + "start_langfuse_agent_step", ] diff --git a/src/llm/structured_output.py b/src/llm/structured_output.py index 76c0690a..37d6b982 100644 --- a/src/llm/structured_output.py +++ b/src/llm/structured_output.py @@ -1,27 +1,33 @@ from __future__ import annotations import json -from collections.abc import Awaitable, Callable -from typing import Literal from pydantic import BaseModel, ValidationError from src.utils.json_parser import validate_and_repair_json from src.utils.representation import PromptRepresentation -from .backend import CompletionResult - -StructuredOutputFailurePolicy = Literal[ - "raise", - "repair_then_raise", - "repair_then_empty", -] - class StructuredOutputError(ValueError): """Raised when structured output cannot be validated or repaired.""" +def schema_instruction(response_format: type[BaseModel], *, tools_present: bool) -> str: + """Structured-output instruction appended to the conversation for + providers without native (or tools-compatible) schema enforcement. + + When tools are in play the wording is conditional so the model remains + free to emit tool calls; validation then relies on parse + repair. + """ + schema_json = json.dumps(response_format.model_json_schema(), indent=2) + if tools_present: + return ( + "\n\nIf not responding with a tool call, respond with valid JSON " + f"matching this schema:\n{schema_json}" + ) + return f"\n\nRespond with valid JSON matching this schema:\n{schema_json}" + + def repair_response_model_json( raw_content: str, response_model: type[BaseModel], @@ -79,54 +85,7 @@ def validate_structured_output( ) -def attempt_structured_output_repair( - content: object, - response_model: type[BaseModel], - model: str, -) -> BaseModel | None: - if not isinstance(content, str): - return None - try: - return repair_response_model_json(content, response_model, model) - except (StructuredOutputError, ValidationError): - return None - - def empty_structured_output(response_model: type[BaseModel]) -> BaseModel: if response_model is PromptRepresentation: return PromptRepresentation(explicit=[]) return response_model.model_validate({}) - - -async def execute_structured_output_call( - executor: Callable[[], Awaitable[CompletionResult]], - *, - response_model: type[BaseModel], - model_name: str, - failure_policy: StructuredOutputFailurePolicy = "repair_then_raise", -) -> CompletionResult: - result = await executor() - - try: - result.content = validate_structured_output(result.content, response_model) - return result - except (StructuredOutputError, ValidationError): - if failure_policy == "raise": - raise - - repaired = attempt_structured_output_repair( - result.content, - response_model, - model_name, - ) - if repaired is not None: - result.content = repaired - return result - - if failure_policy == "repair_then_empty": - result.content = empty_structured_output(response_model) - return result - - raise StructuredOutputError( - f"Failed to produce valid structured output for {model_name}" - ) diff --git a/src/llm/tool_loop.py b/src/llm/tool_loop.py index 734af8b3..0feec1d8 100644 --- a/src/llm/tool_loop.py +++ b/src/llm/tool_loop.py @@ -30,12 +30,18 @@ from src.utils.types import ( set_last_tool_metadata, ) -from .executor import honcho_llm_call_inner +from .capture import ( + build_captured_call, + dispatch_captured_call, + has_exporters, +) +from .executor import honcho_llm_call_inner, infer_provider_label from .registry import history_adapter_for_provider from .runtime import ( AttemptPlan, current_attempt, effective_temperature, + start_langfuse_agent_step, ) from .types import ( HonchoLLMCallResponse, @@ -68,30 +74,81 @@ def _with_iteration_scope( return wrapper +def _step_label(base: LLMTelemetryContext | None) -> str: + """Stable per-agent step-span name, e.g. "Dialectic Agent step". + + No step number — Langfuse aggregates by name; the index rides on the + ``iteration`` metadata. The " step" suffix distinguishes it from the bare + agent name, which names the enclosing run trace (see + `start_langfuse_agent_run`). + """ + return f"{(base.track_name if base else None) or 'Agent'} step" + + def _telemetry_for_iteration( - base: LLMTelemetryContext | None, iteration: int + base: LLMTelemetryContext | None, + iteration: int, + *, + step_seq: int, ) -> LLMTelemetryContext | None: - """Return a copy of `base` with `iteration` set, or None if no base. + """Return a copy of `base` with per-step correlation set, or None if no base. We always copy rather than mutate the caller-supplied context so callers that pass the same context into multiple `honcho_llm_call` invocations - don't see drift across concurrent runs. + don't see drift across concurrent runs. `parent_span_id` is set to the + span being looped over (`base.span_id`) so nested generations are correctly + treated as children of the run span. """ if base is None: return None - return LLMTelemetryContext( - workspace_name=base.workspace_name, - call_purpose=base.call_purpose, - parent_category=base.parent_category, - run_id=base.run_id, + return dataclasses.replace( + base, iteration=iteration, - observer=base.observer, - observed=base.observed, - peer_name=base.peer_name, - agent_type=base.agent_type, + step_seq=step_seq, + parent_span_id=base.span_identity(), ) +def _make_stream_capture_finalizer( + telemetry: LLMTelemetryContext | None, + plan: AttemptPlan, + messages: list[dict[str, Any]], +) -> Callable[[str, str], None] | None: + """Build the streamed-call capture finalizer, or None when capture is off. + + Snapshots the input messages now and returns a closure the streaming wrapper + calls on drain with `(streamed_text, finish_reason)`. Tool calls already ran + in the loop, so the final streamed turn is text-only. Returns None when no + exporter is registered. + """ + if not has_exporters(): + return None + captured_messages = list(messages) + + def _finalize(text: str, finish_reason: str) -> None: + from .backend import CompletionResult as BackendCompletionResult + + result = BackendCompletionResult(content=text, finish_reason=finish_reason) + dispatch_captured_call( + build_captured_call( + telemetry=telemetry, + transport=str(plan.provider), + provider_label=infer_provider_label(plan.provider, plan.model, plan), + model=plan.model, + messages=captured_messages, + tools=None, + tool_choice=None, + result=result, + attempt=plan.attempt, + was_fallback=plan.is_fallback, + was_stream=True, + finish_reason=finish_reason, + ) + ) + + return _finalize + + def _emit_agent_iteration( telemetry: LLMTelemetryContext | None, iteration: int, @@ -221,6 +278,13 @@ async def stream_final_response( # value — telemetry can't tell the retry sequence apart. stream_attempt = 0 + # No ContextVar gymnastics around `_in_agent_run` here: the run handle + # is alive for the lifetime of the stream (owned by + # `StreamingResponseWithMetadata` and closed on drain), so this streamed + # generation correctly nests under the run span as the current OTel + # observation. The previous code had to flip `_in_agent_run` to escape + # the run; with imperative handles the run isn't going anywhere. + async def _setup_stream() -> AsyncIterator[HonchoLLMCallStreamChunk]: nonlocal stream_attempt stream_attempt += 1 @@ -292,6 +356,7 @@ async def execute_tool_loop( stream_final: bool = False, iteration_callback: IterationCallback | None = None, telemetry: LLMTelemetryContext | None = None, + langfuse_run_handle: Any | None = None, ) -> HonchoLLMCallResponse[Any] | StreamingResponseWithMetadata: """Run the iterative tool calling loop for agentic LLM interactions. @@ -318,6 +383,12 @@ async def execute_tool_loop( messages.copy() if messages else [{"role": "user", "content": prompt}] ) + # Seed one hash memo for the whole span. dataclasses.replace copies the dict reference into + # every per-iteration telemetry copy, so each appended message is content-hashed exactly once + # across the span. + if telemetry is not None and telemetry.hash_memo is None: + telemetry = dataclasses.replace(telemetry, hash_memo={}) + iteration = 0 all_tool_calls: list[dict[str, Any]] = [] total_input_tokens = 0 @@ -337,195 +408,232 @@ async def execute_tool_loop( effective_tool_choice = tool_choice while iteration < max_tool_iterations: - # Reset attempt counter so each iteration starts with the primary provider. - current_attempt.set(1) - logger.debug(f"Tool execution iteration {iteration + 1}/{max_tool_iterations}") - - if max_input_tokens is not None: - if count_message_tokens(conversation_messages) > max_input_tokens: - hit_input_token_cap = True - conversation_messages = truncate_messages_to_fit( - conversation_messages, max_input_tokens - ) - - async def _call_with_messages( - effective_tool_choice: str | dict[str, Any] | None = effective_tool_choice, - conversation_messages: list[dict[str, Any]] = conversation_messages, - iteration_for_call: int = iteration + 1, - ) -> HonchoLLMCallResponse[Any]: - plan = get_attempt_plan() - return await honcho_llm_call_inner( - plan.provider, - plan.model, - prompt, # ignored when messages is passed - max_tokens, - response_model, - json_mode, - effective_temperature(temperature), - stop_seqs, - plan.reasoning_effort, - verbosity, - plan.thinking_budget_tokens, - stream=False, - client_override=plan.client, - tools=tools, - tool_choice=effective_tool_choice, - messages=conversation_messages, - selected_config=plan.selected_config, - plan=plan, - telemetry=_telemetry_for_iteration(telemetry, iteration_for_call), - ) - - if enable_retry: - call_func = retry( - stop=stop_after_attempt(retry_attempts), - wait=wait_exponential(multiplier=1, min=4, max=10), - before_sleep=before_retry_callback, - )(_call_with_messages) - else: - call_func = _call_with_messages - - response = await call_func() - - total_input_tokens += response.input_tokens - total_output_tokens += response.output_tokens - total_cache_creation_tokens += response.cache_creation_input_tokens - total_cache_read_tokens += response.cache_read_input_tokens - - # emit one AgentIterationEvent per LLM response BEFORE the - # no-tool early return. The terminating iteration counts too — it has - # an empty tool_calls list and is essential for cost calibration. - _emit_agent_iteration(telemetry, iteration + 1, response) - - if not response.tool_calls_made: - logger.debug("No tool calls in response, finishing") - - if ( - isinstance(response.content, str) - and not response.content.strip() - and empty_response_retries < 1 - and iteration < max_tool_iterations - 1 - ): - empty_response_retries += 1 - conversation_messages.append( - { - "role": "user", - "content": ( - "Your last response was empty. Provide a concise answer " - "to the original query using the available context." - ), - } - ) - iteration += 1 - continue - - if stream_final: - # Snapshot the plan that just succeeded — streaming retries - # pin to this exact client/model so we don't bounce back to - # primary after the tool loop settled on fallback. - winning_plan = get_attempt_plan() - stream = stream_final_response( - winning_plan=winning_plan, - prompt=prompt, - max_tokens=max_tokens, - conversation_messages=conversation_messages, - response_model=response_model, - json_mode=json_mode, - temperature=temperature, - stop_seqs=stop_seqs, - verbosity=verbosity, - enable_retry=enable_retry, - retry_attempts=retry_attempts, - before_retry_callback=before_retry_callback, - telemetry=_telemetry_for_iteration(telemetry, iteration + 1), - ) - return StreamingResponseWithMetadata( - stream=stream, - tool_calls_made=all_tool_calls, - input_tokens=total_input_tokens, - output_tokens=total_output_tokens, - cache_creation_input_tokens=total_cache_creation_tokens, - cache_read_input_tokens=total_cache_read_tokens, - thinking_content=response.thinking_content, - iterations=iteration + 1, - hit_input_token_cap=hit_input_token_cap, - ) - - response.tool_calls_made = all_tool_calls - response.input_tokens = total_input_tokens - response.output_tokens = total_output_tokens - response.cache_creation_input_tokens = total_cache_creation_tokens - response.cache_read_input_tokens = total_cache_read_tokens - response.iterations = iteration + 1 - response.hit_input_token_cap = ( - response.hit_input_token_cap or hit_input_token_cap - ) - return response - - current_provider = get_attempt_plan().provider - - assistant_message = format_assistant_tool_message( - current_provider, - response.content, - response.tool_calls_made, - response.thinking_blocks, - response.reasoning_details, + step = start_langfuse_agent_step( + _step_label(telemetry), + _telemetry_for_iteration(telemetry, iteration + 1, step_seq=iteration + 1), ) - conversation_messages.append(assistant_message) + try: + # Reset attempt counter so each iteration starts with the primary provider. + current_attempt.set(1) + logger.debug( + f"Tool execution iteration {iteration + 1}/{max_tool_iterations}" + ) - # Telemetry context — 1-indexed iteration. - set_current_iteration(iteration + 1) - - tool_results: list[dict[str, Any]] = [] - for seq, tool_call in enumerate(response.tool_calls_made): - tool_name = tool_call["name"] - tool_input = tool_call["input"] - tool_id = tool_call.get("id", "") - - logger.debug(f"Executing tool: {tool_name}") - - # the executor closure reads these from - # ContextVars to populate AgentToolCallCompletedEvent. Set BEFORE - # the executor call so two calls to the same tool in one iteration - # get distinct seq values. Reset last-tool metadata so we never - # observe stale state from a prior call. - set_current_tool_call_seq(seq, tool_id or None) - set_last_tool_metadata({}) - - try: - tool_result = await tool_executor(tool_name, tool_input) - # Stash ToolResult.metadata on all_tool_calls so - # specialist rollups can read created/deleted observation - # counts without round-tripping through the event store. - tool_result_metadata = get_last_tool_metadata() - tool_results.append( - { - "tool_id": tool_id, - "tool_name": tool_name, - "result": tool_result, - } - ) - all_tool_calls.append( - { - "tool_name": tool_name, - "tool_input": tool_input, - "tool_result": tool_result, - "tool_result_metadata": tool_result_metadata, - } - ) - except Exception as e: - logger.error(f"Tool execution failed for {tool_name}: {e}") - tool_results.append( - { - "tool_id": tool_id, - "tool_name": tool_name, - "result": f"Error: {str(e)}", - "is_error": True, - } + if max_input_tokens is not None: + if count_message_tokens(conversation_messages) > max_input_tokens: + hit_input_token_cap = True + conversation_messages = truncate_messages_to_fit( + conversation_messages, max_input_tokens ) - append_tool_results(current_provider, tool_results, conversation_messages) + async def _call_with_messages( + tool_choice_for_call: str + | dict[str, Any] + | None = effective_tool_choice, + captured_messages: list[dict[str, Any]] = conversation_messages, + iteration_for_call: int = iteration + 1, + ) -> HonchoLLMCallResponse[Any]: + plan = get_attempt_plan() + return await honcho_llm_call_inner( + plan.provider, + plan.model, + prompt, # ignored when messages is passed + max_tokens, + response_model, + json_mode, + effective_temperature(temperature), + stop_seqs, + plan.reasoning_effort, + verbosity, + plan.thinking_budget_tokens, + stream=False, + client_override=plan.client, + tools=tools, + tool_choice=tool_choice_for_call, + messages=captured_messages, + selected_config=plan.selected_config, + plan=plan, + telemetry=_telemetry_for_iteration( + telemetry, iteration_for_call, step_seq=iteration_for_call + ), + ) + call_func: Callable[[], Awaitable[HonchoLLMCallResponse[Any]]] + if enable_retry: + call_func = retry( + stop=stop_after_attempt(retry_attempts), + wait=wait_exponential(multiplier=1, min=4, max=10), + before_sleep=before_retry_callback, + )(_call_with_messages) + else: + call_func = _call_with_messages # pyright: ignore[reportGeneralTypeIssues] + + response = await call_func() + + total_input_tokens += response.input_tokens + total_output_tokens += response.output_tokens + total_cache_creation_tokens += response.cache_creation_input_tokens + total_cache_read_tokens += response.cache_read_input_tokens + + # emit one AgentIterationEvent per LLM response BEFORE the + # no-tool early return. The terminating iteration counts too — it has + # an empty tool_calls list and is essential for cost calibration. + _emit_agent_iteration(telemetry, iteration + 1, response) + + # Step span is current again (the generation closed); stamp this + # turn's I/O so it isn't blank. + if step is not None: + step.annotate_io( + conversation_messages, + response.content, + response.tool_calls_made, + ) + + if not response.tool_calls_made: + logger.debug("No tool calls in response, finishing") + + if ( + isinstance(response.content, str) + and not response.content.strip() + and empty_response_retries < 1 + and iteration < max_tool_iterations - 1 + ): + empty_response_retries += 1 + conversation_messages.append( + { + "role": "user", + "content": ( + "Your last response was empty. Provide a concise answer " + "to the original query using the available context." + ), + } + ) + iteration += 1 + continue + + if stream_final: + # Snapshot the plan that just succeeded — streaming retries + # pin to this exact client/model so we don't bounce back to + # primary after the tool loop settled on fallback. + winning_plan = get_attempt_plan() + # +2 (not +1): the in-loop call we just made used iteration+1, + # so the streamed tail needs the next ordinal — otherwise its + # trace resource id collides with that call's. Mirrors the + # synthesis path's distinct-next-value behavior. + stream_telemetry = _telemetry_for_iteration( + telemetry, iteration + 2, step_seq=iteration + 2 + ) + stream = stream_final_response( + winning_plan=winning_plan, + prompt=prompt, + max_tokens=max_tokens, + conversation_messages=conversation_messages, + response_model=response_model, + json_mode=json_mode, + temperature=temperature, + stop_seqs=stop_seqs, + verbosity=verbosity, + enable_retry=enable_retry, + retry_attempts=retry_attempts, + before_retry_callback=before_retry_callback, + telemetry=stream_telemetry, + ) + return StreamingResponseWithMetadata( + stream=stream, + tool_calls_made=all_tool_calls, + input_tokens=total_input_tokens, + output_tokens=total_output_tokens, + cache_creation_input_tokens=total_cache_creation_tokens, + cache_read_input_tokens=total_cache_read_tokens, + thinking_content=response.thinking_content, + iterations=iteration + 1, + hit_input_token_cap=hit_input_token_cap, + langfuse_run_handle=langfuse_run_handle, + capture_finalizer=_make_stream_capture_finalizer( + stream_telemetry, winning_plan, conversation_messages + ), + ) + + response.tool_calls_made = all_tool_calls + response.input_tokens = total_input_tokens + response.output_tokens = total_output_tokens + response.cache_creation_input_tokens = total_cache_creation_tokens + response.cache_read_input_tokens = total_cache_read_tokens + response.iterations = iteration + 1 + response.hit_input_token_cap = ( + response.hit_input_token_cap or hit_input_token_cap + ) + return response + + current_provider = get_attempt_plan().provider + + assistant_message = format_assistant_tool_message( + current_provider, + response.content, + response.tool_calls_made, + response.thinking_blocks, + response.reasoning_details, + ) + conversation_messages.append(assistant_message) + + # Telemetry context — 1-indexed iteration. + set_current_iteration(iteration + 1) + + tool_results: list[dict[str, Any]] = [] + for seq, tool_call in enumerate(response.tool_calls_made): + tool_name = tool_call["name"] + tool_input = tool_call["input"] + tool_id = tool_call.get("id", "") + + logger.debug(f"Executing tool: {tool_name}") + + # the executor closure reads these from + # ContextVars to populate AgentToolCallCompletedEvent. Set BEFORE + # the executor call so two calls to the same tool in one iteration + # get distinct seq values. Reset last-tool metadata so we never + # observe stale state from a prior call. + set_current_tool_call_seq(seq, tool_id or None) + set_last_tool_metadata({}) + + try: + tool_result = await tool_executor(tool_name, tool_input) + # Stash ToolResult.metadata on all_tool_calls so + # specialist rollups can read created/deleted observation + # counts without round-tripping through the event store. + tool_result_metadata = get_last_tool_metadata() + tool_results.append( + { + "tool_id": tool_id, + "tool_name": tool_name, + "result": tool_result, + } + ) + all_tool_calls.append( + { + "tool_name": tool_name, + "tool_input": tool_input, + "tool_result": tool_result, + "tool_result_metadata": tool_result_metadata, + } + ) + except Exception as e: + logger.error(f"Tool execution failed for {tool_name}: {e}") + tool_results.append( + { + "tool_id": tool_id, + "tool_name": tool_name, + "result": f"Error: {str(e)}", + "is_error": True, + } + ) + + append_tool_results(current_provider, tool_results, conversation_messages) + finally: + if step is not None: + step.end() + + # Between-turn bookkeeping lives outside the step span — the span + # scopes the LLM call + its tools, not the iteration accounting. if iteration_callback is not None: try: iteration_data = IterationData( @@ -577,6 +685,9 @@ async def execute_tool_loop( # Snapshot the plan the loop settled on — streaming retries pin to # this exact client/model rather than re-running provider selection. winning_plan = get_attempt_plan() + stream_telemetry = _telemetry_for_iteration( + telemetry, synthesis_iteration, step_seq=synthesis_iteration + ) stream = stream_final_response( winning_plan=winning_plan, prompt=prompt, @@ -590,7 +701,7 @@ async def execute_tool_loop( enable_retry=enable_retry, retry_attempts=retry_attempts, before_retry_callback=before_retry_callback, - telemetry=_telemetry_for_iteration(telemetry, synthesis_iteration), + telemetry=stream_telemetry, ) return StreamingResponseWithMetadata( stream=stream, @@ -602,6 +713,10 @@ async def execute_tool_loop( thinking_content=None, iterations=iteration + 1, hit_input_token_cap=hit_input_token_cap, + langfuse_run_handle=langfuse_run_handle, + capture_finalizer=_make_stream_capture_finalizer( + stream_telemetry, winning_plan, conversation_messages + ), ) current_attempt.set(1) @@ -627,7 +742,9 @@ async def execute_tool_loop( messages=conversation_messages, selected_config=plan.selected_config, plan=plan, - telemetry=_telemetry_for_iteration(telemetry, synthesis_iteration), + telemetry=_telemetry_for_iteration( + telemetry, synthesis_iteration, step_seq=synthesis_iteration + ), ) if enable_retry: @@ -639,13 +756,28 @@ async def execute_tool_loop( else: final_call_func = _final_call - final_response = await final_call_func() + # Step span around the synthesis call — same shape as in-loop iterations + # so the generation nests under the run root instead of dangling at the + # trace. Imperative pair with a try/finally for the .end(). + synthesis_step = start_langfuse_agent_step( + _step_label(telemetry), + _telemetry_for_iteration( + telemetry, synthesis_iteration, step_seq=synthesis_iteration + ), + ) + try: + final_response = await final_call_func() + finally: + if synthesis_step is not None: + synthesis_step.end() # emit the synthesis-call iteration event BEFORE merging cumulative # totals onto final_response below — otherwise the event's per-iteration # token counts would double-count the running totals. _emit_agent_iteration( - _telemetry_for_iteration(telemetry, synthesis_iteration), + _telemetry_for_iteration( + telemetry, synthesis_iteration, step_seq=synthesis_iteration + ), synthesis_iteration, final_response, ) diff --git a/src/llm/types.py b/src/llm/types.py index a81e6e43..33058e09 100644 --- a/src/llm/types.py +++ b/src/llm/types.py @@ -6,15 +6,22 @@ of the migration toward src/llm/ owning all non-embedding LLM orchestration. from __future__ import annotations +import asyncio +import logging from collections.abc import AsyncIterator, Callable -from dataclasses import dataclass -from typing import Any, Generic, Literal, TypeVar +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar from anthropic import AsyncAnthropic from google import genai from openai import AsyncOpenAI from pydantic import BaseModel, Field +if TYPE_CHECKING: + from src.llm.capture import CapturedMessage + +logger = logging.getLogger(__name__) + T = TypeVar("T") # OpenAI GPT-5 specific reasoning levels. @@ -65,16 +72,46 @@ class LLMTelemetryContext: parent_category: str | None = None run_id: str | None = None iteration: int | None = None + # OpenTelemetry-style span-tree correlation. + trace_id: str | None = None + span_id: str | None = None + parent_span_id: str | None = None + # Monotonic executor-call ordinal WITHIN a span (total ordering of its + # steps). + step_seq: int = 0 + # Retry/fallback attempt within an iteration. + attempt: int = 1 # Optional peer context (dream agents pass observer/observed; dialectic # passes peer_name). Kept here so AgentIterationEvent can populate # them without a separate threading path. observer: str | None = None observed: str | None = None peer_name: str | None = None + # Used to group traces (should not use session_name because it is not unique) + session_id: str | None = None # Tool-related context: agent_type is the human-readable identifier of the # agent — dialectic/deduction/induction. Used by agent iteration # event and tool call event. agent_type: str | None = None + # Human-readable name for the Langfuse trace + per-call generation + # (e.g. "Dialectic Agent", "Minimal Deriver"). Sole home for this name — + # callers set it here; `honcho_llm_call` no longer takes a separate kwarg. + # Also used to label the sentry `ai_track` decorator and as the source for + # the run-level `langfuse_agent_run` label. + track_name: str | None = None + # Per-span memo for O(N) message capture in CapturedLLMCall + hash_memo: dict[int, CapturedMessage] | None = field( + default=None, compare=False, repr=False + ) + + def span_identity(self) -> str | None: + """Effective span id: the new `span_id`, falling back to legacy `run_id`.""" + return self.span_id or self.run_id + + def exported_parent_span_id(self) -> str | None: + """`parent_span_id` for EXPORT, collapsing the self-parent sentinel to None.""" + pid = self.parent_span_id + return None if pid is not None and pid == self.span_id else pid IterationCallback = Callable[[IterationData], None] @@ -134,6 +171,20 @@ class StreamingResponseWithMetadata: reflects tool-loop output + final-stream output. Callers that read `output_tokens` AFTER fully iterating the stream get the true total; callers that read it before drain see only the tool-loop portion. + + `langfuse_run_handle` (optional) is the run-level Langfuse span handle + transferred from `honcho_llm_call` when streaming. The wrapper owns it + after construction: on drain, the accumulated streamed text is stamped + as the run span's output and the span is closed. Without this transfer, + streaming traces would show blank output because the synchronous return + happens before any chunks arrive. + + `capture_finalizer` (optional) closes the replay-grade content capture for + a streamed call. The synchronous return happens before any chunks arrive, + so the streamed text only exists once the stream drains — the wrapper calls + the finalizer with `(accumulated_text, finish_reason)` in its `finally`. + A partial/aborted stream still finalizes, with `finish_reason` = + "cancelled"/"error". """ _stream: AsyncIterator[HonchoLLMCallStreamChunk] @@ -145,6 +196,8 @@ class StreamingResponseWithMetadata: thinking_content: str | None iterations: int hit_input_token_cap: bool + _langfuse_run_handle: Any | None + _capture_finalizer: Callable[[str, str], None] | None def __init__( self, @@ -157,6 +210,8 @@ class StreamingResponseWithMetadata: thinking_content: str | None = None, iterations: int = 0, hit_input_token_cap: bool = False, + langfuse_run_handle: Any | None = None, + capture_finalizer: Callable[[str, str], None] | None = None, ): self._stream = stream self.tool_calls_made = tool_calls_made @@ -167,6 +222,8 @@ class StreamingResponseWithMetadata: self.thinking_content = thinking_content self.iterations = iterations self.hit_input_token_cap = hit_input_token_cap + self._langfuse_run_handle = langfuse_run_handle + self._capture_finalizer = capture_finalizer def __aiter__(self) -> AsyncIterator[HonchoLLMCallStreamChunk]: # Wrap the underlying iterator to capture final-stream output_tokens @@ -180,20 +237,61 @@ class StreamingResponseWithMetadata: self, ) -> AsyncIterator[HonchoLLMCallStreamChunk]: final_stream_output_tokens = 0 - async for chunk in self._stream: - if chunk.output_tokens is not None: - # Take the LATEST value, not the sum — providers report - # the cumulative usage in the final chunk, not deltas. - final_stream_output_tokens = chunk.output_tokens - yield chunk - # Stream drained — fold the final-stream output tokens into the - # tool-loop totals so DialecticCompletedEvent / downstream readers - # see the true cost. - if final_stream_output_tokens > 0: - self.output_tokens += final_stream_output_tokens - - async def __anext__(self) -> HonchoLLMCallStreamChunk: - return await self._stream.__anext__() + # Accumulate the streamed text when either consumer needs it: the + # Langfuse run span (stamped as output on drain) or the content-capture + # finalizer. + accumulate = ( + self._langfuse_run_handle is not None or self._capture_finalizer is not None + ) + accumulated_text: list[str] = [] + last_finish_reason: str | None = None + stream_error: BaseException | None = None + try: + async for chunk in self._stream: + if chunk.output_tokens is not None: + # Take the LATEST value, not the sum — providers report + # the cumulative usage in the final chunk, not deltas. + final_stream_output_tokens = chunk.output_tokens + if chunk.finish_reasons: + last_finish_reason = chunk.finish_reasons[-1] + if accumulate and chunk.content: + accumulated_text.append(chunk.content) + yield chunk + # Stream drained — fold the final-stream output tokens into the + # tool-loop totals so DialecticCompletedEvent / downstream readers + # see the true cost. + if final_stream_output_tokens > 0: + self.output_tokens += final_stream_output_tokens + except BaseException as exc: + stream_error = exc + raise + finally: + text = "".join(accumulated_text) + # Close the run span once, stamping the streamed text as its + # output. In `finally` so an early-exit caller still closes + # the span rather than leaking it. + handle = self._langfuse_run_handle + if handle is not None: + self._langfuse_run_handle = None + handle.end(output=text or None) + # Finalize the content capture with the full streamed text. Even a + # partial/aborted stream captures, tagged with the right outcome. + finalizer = self._capture_finalizer + if finalizer is not None: + self._capture_finalizer = None + finish_reason = ( + (last_finish_reason or "stop") + if stream_error is None + else ( + "cancelled" + if isinstance(stream_error, asyncio.CancelledError) + else "error" + ) + ) + try: + finalizer(text, finish_reason) + except Exception: # pragma: no cover - best-effort telemetry + logger.debug("Stream capture finalizer failed", exc_info=True) __all__ = [ diff --git a/src/main.py b/src/main.py index d08d1164..3c8bf3e8 100644 --- a/src/main.py +++ b/src/main.py @@ -1,24 +1,23 @@ import logging import re +import time import uuid from collections.abc import Awaitable, Callable from contextlib import asynccontextmanager -from typing import TYPE_CHECKING import sentry_sdk from fastapi import FastAPI, Request, Response -from fastapi.exceptions import RequestValidationError from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse from fastapi_pagination import add_pagination -from pydantic import ValidationError from sentry_sdk.integrations.fastapi import FastApiIntegration +from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration from sentry_sdk.integrations.starlette import StarletteIntegration from src._version import HONCHO_VERSION from src.cache.client import close_cache, init_cache from src.config import settings -from src.db import engine, request_context +from src.db import engine, register_db_query_instrumentation, request_context from src.exceptions import HonchoException from src.routers import ( conclusions, @@ -34,14 +33,12 @@ from src.telemetry import ( initialize_telemetry_async, metrics_endpoint, prometheus_metrics, + register_db_pool_collector, shutdown_telemetry, ) from src.telemetry.logging import get_route_template from src.telemetry.sentry import initialize_sentry -if TYPE_CHECKING: - from sentry_sdk._types import Event, Hint - def get_log_level() -> int: """ @@ -85,30 +82,10 @@ class MetricsAccessFilter(logging.Filter): logging.getLogger("uvicorn.access").addFilter(MetricsAccessFilter()) -def before_send(event: "Event", hint: "Hint | None") -> "Event | None": - """Filter out events raised from known non-actionable exceptions before Sentry sees them.""" - if not hint: - return event - - exc_info = hint.get("exc_info") - if not exc_info: - return event - - _, exc_value, _ = exc_info - if isinstance(exc_value, HonchoException): - return None - - # Filters out ValidationErrors and RequestValidationErrors (typically coming from Pydantic) - if isinstance(exc_value, ValidationError | RequestValidationError): - logger.info(f"Filtering out validation error from Sentry: {exc_value}") - return None - - return event - - # Sentry Setup SENTRY_ENABLED = settings.SENTRY.ENABLED if SENTRY_ENABLED: + # before_send defaults to sentry.default_before_send (shared with the deriver). initialize_sentry( integrations=[ StarletteIntegration( @@ -117,8 +94,9 @@ if SENTRY_ENABLED: FastApiIntegration( transaction_style="endpoint", ), + # Explicit so DB-query spans are not reliant on auto-enabling. + SqlalchemyIntegration(), ], - before_send=before_send, ) @@ -127,6 +105,10 @@ async def lifespan(_: FastAPI): # Initialize CloudEvents telemetry await initialize_telemetry_async() + # Expose DB connection-pool stats for this API instance (no-op if metrics off) + register_db_pool_collector("api") + register_db_query_instrumentation("api") + # Validate embedding schema before serving any traffic. Fails closed: if # the configured EMBEDDING_VECTOR_DIMENSIONS does not match the physical # pgvector columns, the process refuses to start rather than silently @@ -175,15 +157,9 @@ app = FastAPI( }, ) -origins = [ - "http://localhost", - "http://127.0.0.1:8000", - "https://api.honcho.dev", -] - app.add_middleware( CORSMiddleware, - allow_origins=origins, + allow_origins=settings.CORS_ORIGINS, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], @@ -248,6 +224,7 @@ async def track_request( token = request_context.set(f"api:{request_id}") try: + start_time = time.perf_counter() response = await call_next(request) # Track metrics if enabled @@ -257,6 +234,7 @@ async def track_request( method=request.method, endpoint=template, status_code=str(response.status_code), + duration_seconds=time.perf_counter() - start_time, ) return response diff --git a/src/reconciler/embed_now.py b/src/reconciler/embed_now.py new file mode 100644 index 00000000..f76fa760 --- /dev/null +++ b/src/reconciler/embed_now.py @@ -0,0 +1,389 @@ +""" +Immediate message-embedding fast path. + +``create_messages`` writes ``MessageEmbedding`` rows as ``sync_state='pending'`` +with no vector and defers embedding to the reconciler, which runs on a fixed +interval. To keep freshly created messages searchable within seconds (not +minutes), the message routers schedule ``embed_messages_now`` as a FastAPI +background task right after the response is sent. The reconciler remains the +fallback for anything this path leaves pending (failures, process restarts, or +rows it could not claim). + +The fast path never holds a DB session across a network call (embedding or +external vector store): it claims and leases rows in one short transaction, +embeds with no session open, then persists in short transactions with any +external-store upserts running between them, not inside them. Running concurrently with the +reconciler is safe because the claim uses ``FOR UPDATE SKIP LOCKED`` and leases +rows by stamping ``last_sync_at``, which the reconciler's backoff filter then +skips. +""" + +import asyncio +import logging +from dataclasses import dataclass +from typing import Any + +from fastapi import BackgroundTasks +from sqlalchemy import and_, func, select, update +from sqlalchemy.ext.asyncio import AsyncSession + +from src import models +from src.config import settings +from src.dependencies import tracked_db +from src.embedding_client import embedding_client +from src.exceptions import VectorStoreError +from src.reconciler.sync_vectors import ( + _backoff_eligible, # pyright: ignore[reportPrivateUsage] + build_message_vector_record, + compute_chunk_positions, +) +from src.telemetry import prometheus_metrics +from src.telemetry.events import EmbeddingCallPurpose +from src.utils.types import embedding_call_purpose +from src.vector_store import VectorRecord, VectorStore, get_external_vector_store + +logger = logging.getLogger(__name__) + +_embed_semaphore: asyncio.Semaphore | None = None + + +def _get_embed_semaphore() -> asyncio.Semaphore: + """Lazily create the embed-concurrency semaphore. + + Built on first use (not at import time) so it binds to the running event + loop rather than whatever loop happened to exist at import. + """ + global _embed_semaphore + if _embed_semaphore is None: + _embed_semaphore = asyncio.Semaphore( + settings.EMBEDDING.MAX_CONCURRENT_EMBEDDINGS + ) + return _embed_semaphore + + +def reset_embed_semaphore() -> None: + """Test hook: drop the cached semaphore so the next call rebuilds it on the + current event loop and current config.""" + global _embed_semaphore + _embed_semaphore = None + + +class EmbedTaskGate: + """Non-blocking admission gate for immediate-embed background tasks. + + Bounds the number of in-flight tasks per API process at + ``EMBEDDING.MAX_PENDING_EMBED_TASKS``. When saturated, nothing is scheduled: + the rows are already ``sync_state='pending'``, so the reconciler embeds them + on its next cycle. The count is taken at schedule time (not task start) + because background tasks only run after the response is sent — a request + burst would otherwise stack up unbounded scheduled-but-not-started tasks. + + ``in_flight`` is mutated only from the event loop (request handlers and the + tracked task), so a plain int is race-free. + """ + + def __init__(self) -> None: + self.in_flight: int = 0 + + def try_schedule( + self, background_tasks: BackgroundTasks, message_ids: list[str] + ) -> bool: + """Schedule ``embed_messages_now`` if the cap allows it; return whether + the task was scheduled.""" + if self.in_flight >= settings.EMBEDDING.MAX_PENDING_EMBED_TASKS: + if settings.METRICS.ENABLED: + prometheus_metrics.record_embed_now_task_shed() + return False + self.in_flight += 1 + if settings.METRICS.ENABLED: + prometheus_metrics.set_embed_now_tasks_in_flight(self.in_flight) + background_tasks.add_task(self._run, message_ids) + return True + + async def _run(self, message_ids: list[str]) -> None: + try: + await embed_messages_now(message_ids) + finally: + self.in_flight -= 1 + if settings.METRICS.ENABLED: + prometheus_metrics.set_embed_now_tasks_in_flight(self.in_flight) + + +embed_task_gate = EmbedTaskGate() + + +@dataclass(frozen=True) +class _ClaimedChunk: + """Plain snapshot of a claimed ``MessageEmbedding`` row. + + Captured before the claim transaction commits — after commit the ORM object + is detached and attribute access would lazy-load against a closed session. + """ + + id: int + message_id: str + content: str + workspace_name: str + session_name: str | None + peer_name: str | None + + +async def embed_messages_now(message_ids: list[str]) -> None: + """Embed freshly created messages immediately, leaving the reconciler as the + fallback for anything left pending. + + Args: + message_ids: ``Message.public_id`` values (what + ``MessageEmbedding.message_id`` references). Messages without + embeddable content simply have no pending rows to claim. + """ + if not message_ids: + return + + # Runs as a fire-and-forget background task, so guard the whole flow: an + # unhandled error here would escape into the server's task runner and be + # lost. Any failure just leaves rows pending (claimed rows stay leased), + # and the reconciler heals them on its next cycle. + try: + claimed = await _claim_and_lease(message_ids) + if not claimed: + return + + vectors = await _embed_chunks(claimed) + if vectors is None: + # Embedding failed; rows stay pending + leased, reconciler will retry. + return + + await _persist(message_ids, claimed, vectors) + except Exception: + logger.exception( + "Immediate embed failed for %s message(s); reconciler will retry", + len(message_ids), + ) + + +async def _claim_and_lease(message_ids: list[str]) -> list[_ClaimedChunk]: + """Phase 1 (short txn): claim eligible pending rows with FOR UPDATE SKIP + LOCKED, lease them by stamping ``last_sync_at``, and snapshot their data. + + ``sync_attempts`` is intentionally left untouched: the reconciler owns retry + accounting and the eventual ``sync_state='failed'`` backstop, so a transient + embedding failure on this best-effort path never burns that budget. + """ + async with tracked_db("embed_now_claim") as db: + rows_stmt = ( + select(models.MessageEmbedding) + .where( + and_( + models.MessageEmbedding.message_id.in_(message_ids), + models.MessageEmbedding.sync_state == "pending", + _backoff_eligible(models.MessageEmbedding.last_sync_at), + ) + ) + .order_by(models.MessageEmbedding.message_id, models.MessageEmbedding.id) + .with_for_update(skip_locked=True) + ) + rows = list((await db.execute(rows_stmt)).scalars().all()) + if not rows: + return [] + + claimed = [ + _ClaimedChunk( + id=row.id, + message_id=row.message_id, + content=row.content, + workspace_name=row.workspace_name, + session_name=row.session_name, + peer_name=row.peer_name, + ) + for row in rows + ] + + await db.execute( + update(models.MessageEmbedding) + .where(models.MessageEmbedding.id.in_([c.id for c in claimed])) + .values(last_sync_at=func.now()) + ) + await db.commit() + return claimed + + +async def _embed_chunks(claimed: list[_ClaimedChunk]) -> list[list[float]] | None: + """Phase 2 (no DB session): embed the claimed chunk contents under the + concurrency semaphore. Returns vectors in input order, or None on failure.""" + workspaces = {c.workspace_name for c in claimed} + try: + async with _get_embed_semaphore(): + with embedding_call_purpose( + EmbeddingCallPurpose.MESSAGE_CREATE.value, + workspace_name=workspaces.pop() if len(workspaces) == 1 else None, + parent_category="api", + ): + return await embedding_client.simple_batch_embed( + [c.content for c in claimed] + ) + except Exception: + logger.exception( + "Immediate embedding failed for %s chunk(s); reconciler will retry", + len(claimed), + ) + return None + + +async def _persist( + message_ids: list[str], + claimed: list[_ClaimedChunk], + vectors: list[list[float]], +) -> None: + """Phase 3: persist vectors and mark rows synced. On failure, rows stay + pending (already leased) and the reconciler heals them. + + pgvector mode is one short transaction. External-store mode never holds a + DB session across the vector-store network call: positions are read in one + short transaction, the upserts run with no session open, and the surviving + rows are marked synced in a second short transaction.""" + if len(vectors) != len(claimed): + logger.warning( + "Embedding count %s != claimed chunk count %s; skipping immediate persist, reconciler will heal", + len(vectors), + len(claimed), + ) + return + + vector_by_id = {c.id: vec for c, vec in zip(claimed, vectors, strict=True)} + # True for pgvector OR during migration (dual-write to both stores). + store_in_postgres = ( + settings.VECTOR_STORE.TYPE == "pgvector" or not settings.VECTOR_STORE.MIGRATED + ) + external = get_external_vector_store() + + if external is None: + async with tracked_db("embed_now_persist") as db: + await _persist_pgvector(db, claimed, vector_by_id) + await db.commit() + return + + synced = await _upsert_external(message_ids, claimed, vector_by_id, external) + if not synced: + return + + async with tracked_db("embed_now_persist") as db: + await _mark_synced(db, synced, vector_by_id, store_in_postgres) + await db.commit() + + +async def _persist_pgvector( + db: AsyncSession, + claimed: list[_ClaimedChunk], + vector_by_id: dict[int, list[float]], +) -> None: + """pgvector-only mode: write the vector and mark synced per row. The + ``sync_state='pending'`` guard keeps us idempotent if the reconciler synced + a row in the gap.""" + for c in claimed: + await db.execute( + update(models.MessageEmbedding) + .where( + and_( + models.MessageEmbedding.id == c.id, + models.MessageEmbedding.sync_state == "pending", + ) + ) + .values( + sync_state="synced", + last_sync_at=func.now(), + sync_attempts=0, + embedding=vector_by_id[c.id], + ) + ) + + +async def _upsert_external( + message_ids: list[str], + claimed: list[_ClaimedChunk], + vector_by_id: dict[int, list[float]], + external: VectorStore, +) -> list[_ClaimedChunk]: + """External-store mode: upsert vectors per namespace with no DB session + open, returning the chunks whose namespaces upserted successfully. + + Chunk positions come from the shared helper (full sibling ordering) so vector + ids match whatever the reconciler writes for any chunk we skipped; reading + them is the only DB work here, done in its own short transaction before any + network call.""" + async with tracked_db("embed_now_positions") as db: + chunk_position = await compute_chunk_positions(db, message_ids) + + by_namespace: dict[str, list[_ClaimedChunk]] = {} + for c in claimed: + ns = external.get_vector_namespace("message", c.workspace_name) + by_namespace.setdefault(ns, []).append(c) + + synced: list[_ClaimedChunk] = [] + for namespace, chunks in by_namespace.items(): + records: list[VectorRecord] = [] + synced_chunks: list[_ClaimedChunk] = [] + for c in chunks: + pos = chunk_position.get(c.id) + if pos is None: + continue + records.append( + build_message_vector_record( + message_id=c.message_id, + chunk_position=pos, + session_name=c.session_name, + peer_name=c.peer_name, + embedding=vector_by_id[c.id], + ) + ) + synced_chunks.append(c) + + if not records: + continue + + try: + await external.upsert_many(namespace, records) + except VectorStoreError: + logger.warning( + "Vector store unavailable during immediate embed of namespace %s; reconciler will retry", + namespace, + ) + continue + except Exception: + logger.exception( + "Unexpected error during immediate embed of namespace %s; reconciler will retry", + namespace, + ) + continue + + synced.extend(synced_chunks) + + return synced + + +async def _mark_synced( + db: AsyncSession, + chunks: list[_ClaimedChunk], + vector_by_id: dict[int, list[float]], + store_in_postgres: bool, +) -> None: + """Mark upserted chunks synced (DB-only). The ``sync_state='pending'`` + guard keeps us idempotent if the reconciler synced a row in the gap.""" + for c in chunks: + values: dict[str, Any] = { + "sync_state": "synced", + "last_sync_at": func.now(), + "sync_attempts": 0, + } + if store_in_postgres: + values["embedding"] = vector_by_id[c.id] + await db.execute( + update(models.MessageEmbedding) + .where( + and_( + models.MessageEmbedding.id == c.id, + models.MessageEmbedding.sync_state == "pending", + ) + ) + .values(**values) + ) diff --git a/src/reconciler/scheduler.py b/src/reconciler/scheduler.py index d4f314a3..e171c4fa 100644 --- a/src/reconciler/scheduler.py +++ b/src/reconciler/scheduler.py @@ -264,5 +264,5 @@ class ReconcilerScheduler: ) return False - logger.info("Enqueued reconciler task: %s", task.name) + logger.debug("Enqueued reconciler task: %s", task.name) return True diff --git a/src/reconciler/sync_vectors.py b/src/reconciler/sync_vectors.py index 0dada25d..68e9ac59 100644 --- a/src/reconciler/sync_vectors.py +++ b/src/reconciler/sync_vectors.py @@ -9,8 +9,9 @@ import datetime import logging import time from dataclasses import dataclass -from typing import cast +from typing import Any, cast +import sentry_sdk from sqlalchemy import and_, delete, or_, select, update from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm.attributes import InstrumentedAttribute @@ -109,28 +110,55 @@ async def _get_message_embeddings_needing_sync( """ Get pending message embeddings that need to be synced to the vector store. - Returns only pending embeddings (with full data including embedding vectors). - The batch_size limits the number of embeddings returned. + Claims up to `batch_size` distinct message_ids that have at least one + eligible pending row, then loads ALL pending rows for those message_ids. + This guarantees a single message's chunks are always processed together in + one batch, which keeps vector-ID assignment (`{message_id}_{chunk_index}`, + derived from row-id ordering) stable across reconciler cycles. - Uses FOR UPDATE SKIP LOCKED to prevent concurrent processing and - orders by last_sync_at (nulls first) to prioritize never-synced records. + Uses FOR UPDATE SKIP LOCKED on the per-row claim so concurrent reconcilers + don't double-process the same chunks. Note: "synced" = done forever, "failed" = permanent failure (manual intervention) """ - stmt = ( - select(models.MessageEmbedding) + # Step 1: pick distinct message_ids with at least one eligible pending row, + # prioritizing those with the oldest last_sync_at. + msg_id_stmt = ( + select( + models.MessageEmbedding.message_id, + func.min(models.MessageEmbedding.last_sync_at).label("oldest_attempt"), + ) .where( and_( models.MessageEmbedding.sync_state == "pending", _backoff_eligible(models.MessageEmbedding.last_sync_at), ) ) - .order_by(models.MessageEmbedding.last_sync_at.asc().nullsfirst()) + .group_by(models.MessageEmbedding.message_id) + .order_by(func.min(models.MessageEmbedding.last_sync_at).asc().nullsfirst()) .limit(batch_size) + ) + msg_id_rows = (await db.execute(msg_id_stmt)).all() + message_ids = [row[0] for row in msg_id_rows] + if not message_ids: + return [] + + # Step 2: claim all pending rows for those messages. Skip rows another + # reconciler holds; if we can't claim every chunk of a message right now, + # the message will be retried next cycle. + rows_stmt = ( + select(models.MessageEmbedding) + .where( + and_( + models.MessageEmbedding.message_id.in_(message_ids), + models.MessageEmbedding.sync_state == "pending", + _backoff_eligible(models.MessageEmbedding.last_sync_at), + ) + ) + .order_by(models.MessageEmbedding.message_id, models.MessageEmbedding.id) .with_for_update(skip_locked=True) ) - - result = await db.execute(stmt) + result = await db.execute(rows_stmt) return list(result.scalars().all()) @@ -177,6 +205,63 @@ async def _bump_message_embedding_sync_attempts( ) +async def compute_chunk_positions( + db: AsyncSession, message_ids: list[str] +) -> dict[int, int]: + """Map each MessageEmbedding row id to its 0-indexed chunk position within + its message. + + Positions are derived from the full set of sibling rows for each message, + ordered by ``(message_id, id)`` — never from a partial subset — so the + ``{message_id}_{chunk_position}`` vector id stays stable no matter which + rows a given caller claimed. Shared by the reconciler and the immediate + embed path so the two writers always agree on vector ids. + """ + if not message_ids: + return {} + + sibling_stmt = ( + select(models.MessageEmbedding.id, models.MessageEmbedding.message_id) + .where(models.MessageEmbedding.message_id.in_(message_ids)) + .order_by(models.MessageEmbedding.message_id, models.MessageEmbedding.id) + ) + sibling_rows = (await db.execute(sibling_stmt)).all() + + embs_by_message: dict[str, list[int]] = {} + for emb_id, msg_id in sibling_rows: + embs_by_message.setdefault(msg_id, []).append(emb_id) + + chunk_position: dict[int, int] = {} + for emb_ids in embs_by_message.values(): + for pos, emb_id in enumerate(emb_ids): + chunk_position[emb_id] = pos + return chunk_position + + +def build_message_vector_record( + *, + message_id: str, + chunk_position: int, + session_name: str | None, + peer_name: str | None, + embedding: list[float], +) -> VectorRecord: + """Build the external-store record for one message-embedding chunk. + + Single source of the ``{message_id}_{chunk_position}`` vector id and the + metadata shape, shared by the reconciler and the immediate embed path. + """ + return VectorRecord( + id=f"{message_id}_{chunk_position}", + embedding=[float(x) for x in embedding], + metadata={ + "message_id": message_id, + "session_name": session_name, + "peer_name": peer_name, + }, + ) + + async def _sync_documents( db: AsyncSession, documents: list[models.Document], @@ -306,16 +391,20 @@ async def _sync_documents( async def _sync_message_embeddings( db: AsyncSession, embeddings: list[models.MessageEmbedding], - external_vector_store: VectorStore, + external_vector_store: VectorStore | None, ) -> tuple[int, int]: """ - Sync a batch of pending message embeddings to the external vector store. + Sync a batch of pending message embeddings. - Handles three cases for each embedding: + When `external_vector_store` is provided, handles three cases per embedding: 1. Embedding exists in postgres → use it for external upsert 2. Embedding missing + need postgres storage → re-embed, write to both stores 3. Embedding missing + external-only mode → re-embed, write to external only + When `external_vector_store` is None (pgvector-only mode), re-embeds any + pending row missing a vector, writes the vector to postgres, and marks + sync_state='synced'. No external upsert is performed. + Returns (synced_count, failed_count). """ if not embeddings: @@ -338,8 +427,12 @@ async def _sync_message_embeddings( if embs_needing_embed: try: contents = [emb.content for emb in embs_needing_embed] + # MESSAGE_CREATE (not VECTOR_SYNC): these rows come from create_messages + # as pending chunks; document re-embeds stay on VECTOR_SYNC below. + workspaces = {emb.workspace_name for emb in embs_needing_embed} with embedding_call_purpose( - EmbeddingCallPurpose.VECTOR_SYNC.value, + EmbeddingCallPurpose.MESSAGE_CREATE.value, + workspace_name=workspaces.pop() if len(workspaces) == 1 else None, parent_category="reconciliation", ): new_embeddings = await embedding_client.simple_batch_embed(contents) @@ -368,6 +461,32 @@ async def _sync_message_embeddings( await _bump_message_embedding_sync_attempts(db, failed_to_embed) failed_count += len(failed_to_embed) + # pgvector-only mode: no external store to upsert to. Any row that now + # has an embedding (either pre-existing or freshly embedded) is fully + # synced. Write embeddings via per-row UPDATE so the vector is persisted + # alongside sync_state in a single statement (session has autoflush=False, + # so the ORM mutation above isn't enough on its own). + if external_vector_store is None: + embs_done: list[models.MessageEmbedding] = [] + for emb in embeddings: + new_emb = freshly_embedded.get(emb.id) + existing = emb.embedding + if new_emb is None and existing is None: + continue + await db.execute( + update(models.MessageEmbedding) + .where(models.MessageEmbedding.id == emb.id) + .values( + sync_state="synced", + last_sync_at=func.now(), + sync_attempts=0, + **({"embedding": new_emb} if new_emb is not None else {}), + ) + ) + embs_done.append(emb) + synced_count += len(embs_done) + return synced_count, failed_count + # Step 2: Compute chunk positions for vector IDs # Messages can be split into multiple chunks; we need {message_id}_{chunk_position} # @@ -380,21 +499,7 @@ async def _sync_message_embeddings( # 2. Removing MessageEmbedding table entirely if it becomes unnecessary # See: https://github.com/plastic-labs/honcho/issues/XXX message_ids = list({emb.message_id for emb in embeddings}) - sibling_stmt = ( - select(models.MessageEmbedding.id, models.MessageEmbedding.message_id) - .where(models.MessageEmbedding.message_id.in_(message_ids)) - .order_by(models.MessageEmbedding.message_id, models.MessageEmbedding.id) - ) - sibling_rows = (await db.execute(sibling_stmt)).all() - - embs_by_message: dict[str, list[int]] = {} - for emb_id, msg_id in sibling_rows: - embs_by_message.setdefault(msg_id, []).append(emb_id) - - chunk_position: dict[int, int] = {} - for emb_ids in embs_by_message.values(): - for pos, emb_id in enumerate(emb_ids): - chunk_position[emb_id] = pos + chunk_position = await compute_chunk_positions(db, message_ids) # Step 3: Build vector records and upsert to external store (all cases) by_namespace: dict[str, list[models.MessageEmbedding]] = {} @@ -416,14 +521,12 @@ async def _sync_message_embeddings( continue vector_records.append( - VectorRecord( - id=f"{emb.message_id}_{chunk_position[emb.id]}", - embedding=[float(x) for x in embedding], - metadata={ - "message_id": emb.message_id, - "session_name": emb.session_name, - "peer_name": emb.peer_name, - }, + build_message_vector_record( + message_id=emb.message_id, + chunk_position=chunk_position[emb.id], + session_name=emb.session_name, + peer_name=emb.peer_name, + embedding=embedding, ) ) embs_to_sync.append(emb) @@ -433,11 +536,23 @@ async def _sync_message_embeddings( try: await external_vector_store.upsert_many(namespace, vector_records) - await db.execute( - update(models.MessageEmbedding) - .where(models.MessageEmbedding.id.in_([e.id for e in embs_to_sync])) - .values(sync_state="synced", last_sync_at=func.now(), sync_attempts=0) - ) + # Per-row UPDATEs so freshly-embedded rows persist the vector + # alongside sync_state. Session has autoflush=False so the ORM + # mutation above isn't sufficient on its own. + for emb in embs_to_sync: + new_emb = freshly_embedded.get(emb.id) + values: dict[str, Any] = { + "sync_state": "synced", + "last_sync_at": func.now(), + "sync_attempts": 0, + } + if new_emb is not None and store_in_postgres: + values["embedding"] = new_emb + await db.execute( + update(models.MessageEmbedding) + .where(models.MessageEmbedding.id == emb.id) + .values(**values) + ) synced_count += len(embs_to_sync) except VectorStoreError: logger.warning( @@ -504,15 +619,18 @@ async def _reconcile_documents_batch( if not docs: return False - synced, failed = await _sync_documents(db, docs, external_vector_store) - metrics.documents_synced += synced - metrics.documents_failed += failed - await db.commit() + with sentry_sdk.start_transaction( + name="reconcile_documents_batch", op="reconciler" + ): + synced, failed = await _sync_documents(db, docs, external_vector_store) + metrics.documents_synced += synced + metrics.documents_failed += failed + await db.commit() return True async def _reconcile_message_embeddings_batch( - external_vector_store: VectorStore, + external_vector_store: VectorStore | None, metrics: ReconciliationMetrics, ) -> bool: """ @@ -525,10 +643,15 @@ async def _reconcile_message_embeddings_batch( if not embs: return False - synced, failed = await _sync_message_embeddings(db, embs, external_vector_store) - metrics.message_embeddings_synced += synced - metrics.message_embeddings_failed += failed - await db.commit() + with sentry_sdk.start_transaction( + name="reconcile_message_embeddings_batch", op="reconciler" + ): + synced, failed = await _sync_message_embeddings( + db, embs, external_vector_store + ) + metrics.message_embeddings_synced += synced + metrics.message_embeddings_failed += failed + await db.commit() return True @@ -592,13 +715,20 @@ async def run_vector_reconciliation_cycle() -> ReconciliationMetrics: external_vector_store = get_external_vector_store() deadline = time.monotonic() + RECONCILIATION_TIME_BUDGET_SECONDS - # If no external vector store (pgvector mode), only clean up soft-deleted documents + # pgvector-only mode: still need to embed pending MessageEmbedding rows + # (create_messages defers embedding to the reconciler), then clean up. if external_vector_store is None: while time.monotonic() < deadline: - did_work = await _cleanup_pgvector_batch(metrics) - if not did_work: + embs_work = await _reconcile_message_embeddings_batch(None, metrics) + + if time.monotonic() >= deadline: break - logger.info("Vector reconciliation cycle completed (pgvector mode)") + + cleanup_work = await _cleanup_pgvector_batch(metrics) + + if not (embs_work or cleanup_work): + break + logger.debug("Vector reconciliation cycle completed (pgvector mode)") return metrics # External vector store mode - reconcile documents, embeddings, and cleanup @@ -625,5 +755,5 @@ async def run_vector_reconciliation_cycle() -> ReconciliationMetrics: logger.debug("No work done, breaking reconciliation loop") break - logger.info("Vector reconciliation cycle completed") + logger.debug("Vector reconciliation cycle completed") return metrics diff --git a/src/routers/conclusions.py b/src/routers/conclusions.py index 20aadb75..3a25a5d7 100644 --- a/src/routers/conclusions.py +++ b/src/routers/conclusions.py @@ -6,7 +6,7 @@ from fastapi_pagination.ext.sqlalchemy import apaginate from sqlalchemy.ext.asyncio import AsyncSession from src import crud, schemas -from src.dependencies import db +from src.dependencies import db, read_db from src.exceptions import ResourceNotFoundException, ValidationException from src.security import require_auth from src.telemetry.events import EmbeddingCallPurpose @@ -67,7 +67,7 @@ async def list_conclusions( False, description="Whether to reverse the order of results", ), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """ List Conclusions using optional filters, ordered by recency unless `reverse` is true. Results are paginated. @@ -97,7 +97,7 @@ async def query_conclusions( ..., description="Semantic search parameters for Conclusions", ), - db: AsyncSession = db, + db: AsyncSession = read_db, ) -> list[schemas.Conclusion]: """ Query Conclusions using semantic search. Use `top_k` to control the number of results returned. diff --git a/src/routers/keys.py b/src/routers/keys.py index f4a50375..0051db90 100644 --- a/src/routers/keys.py +++ b/src/routers/keys.py @@ -9,6 +9,7 @@ from src.security import ( JWTParams, create_jwt, require_auth, + scope_requires_workspace, ) from src.utils.formatting import format_datetime_utc @@ -42,6 +43,17 @@ async def create_key( "At least one of workspace_id, peer_id, or session_id must be provided" ) + # A peer- or session-scoped key must carry its parent workspace, otherwise + # verify_jwt rejects it on every request (the workspace is required to rule + # out cross-workspace use). Shares the predicate with verify_jwt so the + # creation-time guard and the verification-time invariant cannot drift. + if scope_requires_workspace( + peer=peer_id, session=session_id, workspace=workspace_id + ): + raise ValidationException( + "workspace_id is required when scoping a key to a peer or session" + ) + key_str = create_jwt( JWTParams( exp=format_datetime_utc(expires_at) if expires_at else None, diff --git a/src/routers/messages.py b/src/routers/messages.py index 917ca713..b0dd8752 100644 --- a/src/routers/messages.py +++ b/src/routers/messages.py @@ -18,9 +18,10 @@ from sqlalchemy.orm.attributes import flag_modified from src import crud, schemas from src.config import settings -from src.dependencies import db +from src.dependencies import db, read_db from src.deriver import enqueue from src.exceptions import FileTooLargeError, ResourceNotFoundException +from src.reconciler.embed_now import embed_task_gate from src.security import require_auth from src.telemetry import prometheus_metrics from src.telemetry.events import FileUploadedEvent, MessageCreatedEvent, emit @@ -31,9 +32,19 @@ logger = logging.getLogger(__name__) router = APIRouter( prefix="/workspaces/{workspace_id}/sessions/{session_id}/messages", tags=["messages"], - dependencies=[ - Depends(require_auth(workspace_name="workspace_id", session_name="session_id")) - ], +) + +# Read routes additionally allow a peer-scoped key whose peer is a member of the +# session; write routes stay session-scoped only. Applied per-route rather than +# on the router so the two policies can differ. +require_session_read = require_auth( + workspace_name="workspace_id", + session_name="session_id", + allow_member_read=True, +) +require_session_write = require_auth( + workspace_name="workspace_id", + session_name="session_id", ) @@ -81,9 +92,18 @@ async def parse_upload_form( ) -@router.post("", response_model=list[schemas.Message], status_code=201) @router.post( - "/", response_model=list[schemas.Message], status_code=201, include_in_schema=False + "", + response_model=list[schemas.Message], + status_code=201, + dependencies=[Depends(require_session_write)], +) +@router.post( + "/", + response_model=list[schemas.Message], + status_code=201, + include_in_schema=False, + dependencies=[Depends(require_session_write)], ) # backwards compatibility with pre-2.6.0 faulty route endpoint async def create_messages_for_session( background_tasks: BackgroundTasks, @@ -140,13 +160,31 @@ async def create_messages_for_session( # Enqueue all messages in one call background_tasks.add_task(enqueue, payloads) + # Embed immediately so messages are searchable within seconds; the + # reconciler is the fallback for anything left pending. Scheduling is + # capped per process — when saturated, the reconciler picks them up. + if settings.EMBED_MESSAGES and created_messages: + scheduled = embed_task_gate.try_schedule( + background_tasks, [m.public_id for m in created_messages] + ) + if not scheduled: + logger.debug( + "Immediate-embed tasks saturated; deferring %s message(s) to reconciler", + len(created_messages), + ) + return created_messages except ValueError as e: logger.warning(f"Failed to create messages for session {session_id}: {str(e)}") raise -@router.post("/upload", response_model=list[schemas.Message], status_code=201) +@router.post( + "/upload", + response_model=list[schemas.Message], + status_code=201, + dependencies=[Depends(require_session_write)], +) async def create_messages_with_file( background_tasks: BackgroundTasks, workspace_id: str = Path(...), @@ -206,6 +244,20 @@ async def create_messages_with_file( ] background_tasks.add_task(enqueue, payloads) + + # Embed immediately so messages are searchable within seconds; the + # reconciler is the fallback for anything left pending. Scheduling is + # capped per process — when saturated, the reconciler picks them up. + if settings.EMBED_MESSAGES and created_messages: + scheduled = embed_task_gate.try_schedule( + background_tasks, [m.public_id for m in created_messages] + ) + if not scheduled: + logger.debug( + "Immediate-embed tasks saturated; deferring %s message(s) to reconciler", + len(created_messages), + ) + logger.debug( "Batch of %s messages created from file uploads and queued for processing", len(created_messages), @@ -250,7 +302,11 @@ async def create_messages_with_file( return created_messages -@router.post("/list", response_model=Page[schemas.Message]) +@router.post( + "/list", + response_model=Page[schemas.Message], + dependencies=[Depends(require_session_read)], +) async def get_messages( workspace_id: str = Path(...), session_id: str = Path(...), @@ -260,7 +316,7 @@ async def get_messages( reverse: bool | None = Query( False, description="Whether to reverse the order of results" ), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get all messages for a Session with optional filters. Results are paginated.""" try: @@ -283,12 +339,16 @@ async def get_messages( raise ResourceNotFoundException("Session not found") from e -@router.get("/{message_id}", response_model=schemas.Message) +@router.get( + "/{message_id}", + response_model=schemas.Message, + dependencies=[Depends(require_session_read)], +) async def get_message( workspace_id: str = Path(...), session_id: str = Path(...), message_id: str = Path(...), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get a single message by ID from a Session.""" honcho_message = await crud.get_message( @@ -300,7 +360,11 @@ async def get_message( return honcho_message -@router.put("/{message_id}", response_model=schemas.Message) +@router.put( + "/{message_id}", + response_model=schemas.Message, + dependencies=[Depends(require_session_write)], +) async def update_message( workspace_id: str = Path(...), session_id: str = Path(...), diff --git a/src/routers/peers.py b/src/routers/peers.py index efa19a78..bccb2dd0 100644 --- a/src/routers/peers.py +++ b/src/routers/peers.py @@ -10,17 +10,26 @@ from fastapi import APIRouter, Body, Depends, Path, Query, Response from fastapi.responses import StreamingResponse from fastapi_pagination import Page from fastapi_pagination.ext.sqlalchemy import apaginate +from pydantic import BaseModel from sqlalchemy.ext.asyncio import AsyncSession from src import crud, schemas from src.config import settings -from src.dependencies import db, tracked_db +from src.crud.message import get_peer_session_names +from src.crud.session import is_peer_in_session +from src.dependencies import db, read_db, tracked_db from src.dialectic.chat import agentic_chat, agentic_chat_stream from src.embedding_client import embedding_client -from src.exceptions import AuthenticationException, ResourceNotFoundException +from src.exceptions import ( + AuthenticationException, + ResourceNotFoundException, + ValidationException, +) from src.security import JWTParams, require_auth from src.telemetry import prometheus_metrics from src.telemetry.events import EmbeddingCallPurpose, GetContextEvent, emit +from src.utils.filter import extract_session_allowlist +from src.utils.schema_conversion import json_response_schema_to_pydantic from src.utils.search import search from src.utils.types import embedding_call_purpose @@ -43,7 +52,7 @@ async def get_peers( None, description="Filtering options for the peers list" ), reverse: bool = Query(False, description="Whether to reverse the order of results"), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get all Peers for a Workspace, paginated with optional filters.""" filter_param = None @@ -134,7 +143,7 @@ async def get_sessions_for_peer( None, description="Filtering options for the sessions list" ), reverse: bool = Query(False, description="Whether to reverse the order of results"), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get all Sessions for a Peer, paginated with optional filters.""" filter_param = None @@ -167,19 +176,58 @@ async def get_sessions_for_peer( }, }, }, - dependencies=[ - Depends(require_auth(workspace_name="workspace_id", peer_name="peer_id")) - ], ) async def chat( workspace_id: str = Path(...), peer_id: str = Path(...), options: schemas.DialecticOptions = Body(...), + jwt_params: JWTParams = Depends( + require_auth(workspace_name="workspace_id", peer_name="peer_id") + ), ): """ Query a Peer's representation using natural language. Performs agentic search and reasoning to comprehensively answer the query based on all latent knowledge gathered about the peer from their messages and conclusions. """ + # The session id arrives in the body, so require_auth can't gate on it. A + # peer-scoped key may only scope a chat to a session its peer belongs to; + # without this check it could read any session's messages (the dialectic + # injects session history) by naming it here. Workspace/admin tokens + # (jwt_params.p is None) are unaffected. + if jwt_params.p is not None and options.session_id: + async with tracked_db("peers.chat.is_peer_in_session", read_only=True) as s_db: + if not await is_peer_in_session( + s_db, workspace_id, options.session_id, jwt_params.p + ): + raise AuthenticationException("JWT not permissioned for this resource") + + # Parse the session allowlist from filters (422 on unsupported keys/shapes, + # and on a session_id the allowlist doesn't cover). + session_allowlist = extract_session_allowlist( + options.filters, must_include=options.session_id + ) + # A peer-scoped key may only name sessions its peer belongs to — the + # allowlist reaches message recall the same way session_id does above. + # `active_only` matches the is_peer_in_session check above, so both gates + # answer the same question for a peer that has left a session. + if jwt_params.p is not None and session_allowlist is not None: + async with tracked_db("peers.chat.session_scope_auth", read_only=True) as s_db: + member_sessions = set( + await get_peer_session_names( + s_db, workspace_id, jwt_params.p, active_only=True + ) + ) + if not set(session_allowlist) <= member_sessions: + raise AuthenticationException("JWT not permissioned for this resource") + + # Convert the caller's JSON Schema so malformed schemas fail immediately with 422 + response_model: type[BaseModel] | None = None + if options.response_format is not None: + try: + response_model = json_response_schema_to_pydantic(options.response_format) + except ValueError as e: + raise ValidationException(f"Invalid response_format: {e}") from None + # Get or create the peer to ensure it exists async with tracked_db("peers.chat.get_or_create_peer") as peer_db: peers_result = await crud.get_or_create_peers( @@ -217,6 +265,8 @@ async def chat( observer=peer_id, observed=options.target if options.target is not None else peer_id, reasoning_level=options.reasoning_level, + session_allowlist=session_allowlist, + response_model=response_model, ) ), media_type="text/event-stream", @@ -231,6 +281,8 @@ async def chat( # and it's answered from the omniscient Honcho perspective observed=options.target if options.target is not None else peer_id, reasoning_level=options.reasoning_level, + session_allowlist=session_allowlist, + response_model=response_model, ) # Prometheus metrics @@ -265,6 +317,12 @@ async def get_representation( If a target is provided, we get the Representation of the target from the perspective of the Peer. If no target is provided, we get the omniscient Honcho Representation of the Peer. """ + # Parse the session allowlist from filters (422 on unsupported keys/shapes, + # and on a session_id the allowlist doesn't cover). + session_allowlist = extract_session_allowlist( + options.filters, must_include=options.session_id + ) + try: embedding: list[float] | None = None if options.search_query: @@ -283,7 +341,9 @@ async def get_representation( workspace_id, observer=peer_id, observed=options.target if options.target is not None else peer_id, - session_name=options.session_id, + session_allowlist=[options.session_id] + if options.session_id is not None + else session_allowlist, include_semantic_query=options.search_query, embedding=embedding, semantic_search_top_k=options.search_top_k, @@ -318,7 +378,7 @@ async def get_peer_card( None, description="Optional target peer to retrieve a card for, from the observer's perspective. If not provided, returns the observer's own card", ), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get a peer card for a specific peer relationship. @@ -412,7 +472,6 @@ async def get_peer_context( le=100, description="Maximum number of conclusions to include in the representation", ), - db: AsyncSession = db, ): """ Get context for a peer, including their representation and peer card. @@ -447,7 +506,7 @@ async def get_peer_context( workspace_id, observer=peer_id, observed=observed, - session_name=None, # Peer context is global, not session-scoped + session_allowlist=None, # Peer context is global, not session-scoped include_semantic_query=search_query, embedding=embedding, semantic_search_top_k=search_top_k, @@ -459,10 +518,12 @@ async def get_peer_context( parent_category="api", ) - # Get the peer card - peer_card = await crud.get_peer_card( - db, workspace_id, observer=peer_id, observed=observed - ) + async with tracked_db( + "peers.get_peer_context.peer_card", read_only=True + ) as card_db: + peer_card = await crud.get_peer_card( + card_db, workspace_id, observer=peer_id, observed=observed + ) response = schemas.PeerContext( peer_id=peer_id, diff --git a/src/routers/sessions.py b/src/routers/sessions.py index 98a8714a..191748d4 100644 --- a/src/routers/sessions.py +++ b/src/routers/sessions.py @@ -12,7 +12,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import config, crud, schemas from src.cache.client import safe_cache_delete from src.crud.session import session_cache_key -from src.dependencies import db +from src.dependencies import db, read_db from src.deriver.enqueue import enqueue_deletion from src.embedding_client import embedding_client from src.exceptions import ( @@ -43,7 +43,7 @@ async def _get_working_representation_task( *, observer: str, observed: str, - session_name: str | None, + session_allowlist: list[str] | None, search_top_k: int | None, search_max_distance: float | None, include_most_derived: bool, @@ -59,7 +59,7 @@ async def _get_working_representation_task( last_message: Optional last message for semantic query observer: Name of the observer peer observed: Name of the observed peer - session_name: Optional session to filter by + session_allowlist: Optional session allowlist to filter by search_top_k: Number of semantic-search-retrieved observations to include in the representation search_max_distance: Maximum distance to search for semantically relevant observations include_most_derived: Whether to include the most derived observations in the representation @@ -74,7 +74,7 @@ async def _get_working_representation_task( db=db, observer=observer, observed=observed, - session_name=session_name, + session_allowlist=session_allowlist, include_semantic_query=last_message, semantic_search_top_k=search_top_k, semantic_search_max_distance=search_max_distance, @@ -251,7 +251,7 @@ async def get_sessions( None, description="Filtering and pagination options for the sessions list" ), reverse: bool = Query(False, description="Whether to reverse the order of results"), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get all Sessions for a Workspace, paginated with optional filters.""" filter_param = None @@ -536,17 +536,28 @@ async def remove_peers_from_session( @router.get( "/{session_id}/peers/{peer_id}/config", response_model=schemas.SessionPeerConfig, - dependencies=[ - Depends(require_auth(workspace_name="workspace_id", session_name="session_id")) - ], ) async def get_peer_config( workspace_id: str = Path(...), session_id: str = Path(...), peer_id: str = Path(...), - db: AsyncSession = db, + jwt_params: JWTParams = Depends( + require_auth( + workspace_name="workspace_id", + session_name="session_id", + allow_member_read=True, + ) + ), + db: AsyncSession = read_db, ): - """Get the configuration for a Peer in a Session.""" + """Get the configuration for a Peer in a Session. + + Member-read lets a peer-scoped key reach this route, but a peer may only + read its own per-session config — not a co-member's. Workspace/admin and + session-scoped tokens (which already span the whole session) are unaffected. + """ + if jwt_params.p is not None and jwt_params.p != peer_id: + raise AuthenticationException("JWT not permissioned for this resource") return await crud.get_peer_config( db, workspace_name=workspace_id, @@ -593,13 +604,19 @@ async def set_peer_config( "/{session_id}/peers", response_model=Page[schemas.Peer], dependencies=[ - Depends(require_auth(workspace_name="workspace_id", session_name="session_id")) + Depends( + require_auth( + workspace_name="workspace_id", + session_name="session_id", + allow_member_read=True, + ) + ) ], ) async def get_session_peers( workspace_id: str = Path(...), session_id: str = Path(...), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get all Peers in a Session. Results are paginated.""" try: @@ -616,13 +633,19 @@ async def get_session_peers( "/{session_id}/context", response_model=schemas.SessionContext, dependencies=[ - Depends(require_auth(workspace_name="workspace_id", session_name="session_id")) + Depends( + require_auth( + workspace_name="workspace_id", + session_name="session_id", + allow_member_read=True, + ) + ) ], ) async def get_session_context( workspace_id: str = Path(...), session_id: str = Path(...), - db: AsyncSession = db, + db: AsyncSession = read_db, tokens: int | None = Query( None, le=config.settings.GET_CONTEXT_MAX_TOKENS, @@ -742,7 +765,7 @@ async def get_session_context( search_query, observer=observer, observed=observed, - session_name=session_id if limit_to_session else None, + session_allowlist=[session_id] if limit_to_session else None, search_top_k=search_top_k, search_max_distance=search_max_distance, include_most_derived=include_most_frequent, @@ -808,13 +831,19 @@ async def get_session_context( "/{session_id}/summaries", response_model=schemas.SessionSummaries, dependencies=[ - Depends(require_auth(workspace_name="workspace_id", session_name="session_id")) + Depends( + require_auth( + workspace_name="workspace_id", + session_name="session_id", + allow_member_read=True, + ) + ) ], ) async def get_session_summaries( workspace_id: str = Path(...), session_id: str = Path(...), - db: AsyncSession = db, + db: AsyncSession = read_db, ) -> schemas.SessionSummaries: """ Get available summaries for a Session. @@ -849,7 +878,13 @@ async def get_session_summaries( "/{session_id}/search", response_model=list[schemas.Message], dependencies=[ - Depends(require_auth(workspace_name="workspace_id", session_name="session_id")) + Depends( + require_auth( + workspace_name="workspace_id", + session_name="session_id", + allow_member_read=True, + ) + ) ], ) async def search_session( diff --git a/src/routers/webhooks.py b/src/routers/webhooks.py index f66164e9..91fc63f9 100644 --- a/src/routers/webhooks.py +++ b/src/routers/webhooks.py @@ -31,7 +31,7 @@ async def get_or_create_webhook_endpoint( webhook: schemas.WebhookEndpointCreate = Body( ..., description="Webhook endpoint parameters" ), - jwt_params: JWTParams = Depends(require_auth()), + jwt_params: JWTParams = Depends(require_auth(workspace_name="workspace_id")), db: AsyncSession = db, ) -> schemas.WebhookEndpoint: """ @@ -55,7 +55,7 @@ async def get_or_create_webhook_endpoint( @router.get("", response_model=Page[schemas.WebhookEndpoint]) async def list_webhook_endpoints( workspace_id: str = Path(..., description="Workspace ID"), - jwt_params: JWTParams = Depends(require_auth()), + jwt_params: JWTParams = Depends(require_auth(workspace_name="workspace_id")), db: AsyncSession = db, ) -> Page[schemas.WebhookEndpoint]: """ @@ -72,7 +72,7 @@ async def list_webhook_endpoints( async def delete_webhook_endpoint( workspace_id: str = Path(..., description="Workspace ID"), endpoint_id: str = Path(..., description="Webhook endpoint ID"), - jwt_params: JWTParams = Depends(require_auth()), + jwt_params: JWTParams = Depends(require_auth(workspace_name="workspace_id")), db: AsyncSession = db, ) -> None: """ @@ -88,7 +88,7 @@ async def delete_webhook_endpoint( @router.get("/test") async def test_emit( workspace_id: str = Path(..., description="Workspace ID"), - jwt_params: JWTParams = Depends(require_auth()), + jwt_params: JWTParams = Depends(require_auth(workspace_name="workspace_id")), ) -> None: """ Test publishing a webhook event. diff --git a/src/routers/workspaces.py b/src/routers/workspaces.py index a2298480..b7e15a71 100644 --- a/src/routers/workspaces.py +++ b/src/routers/workspaces.py @@ -9,7 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import crud, schemas from src.config import settings -from src.dependencies import db +from src.dependencies import db, read_db from src.deriver.enqueue import enqueue_deletion, enqueue_dream from src.exceptions import AuthenticationException from src.security import JWTParams, require_auth @@ -67,7 +67,7 @@ async def get_all_workspaces( None, description="Filtering and pagination options for the workspaces list" ), reverse: bool = Query(False, description="Whether to reverse the order of results"), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """Get all Workspaces, paginated with optional filters.""" filter_param = None @@ -169,7 +169,7 @@ async def get_queue_status( session_id: str | None = Query( None, description="Optional session ID to filter by" ), - db: AsyncSession = db, + db: AsyncSession = read_db, ): """ Get the processing queue status for a Workspace, optionally scoped to an observer, sender, @@ -231,6 +231,7 @@ async def schedule_dream( observed=observed, dream_type=dream_type, session_name=request.session_id, + rebuild=request.rebuild, # Manual route — explicit sentinels for the DreamRunEvent # scheduling-context fields. Auto-schedule threads concrete # threshold/delay reasons (see src/dreamer/dream_scheduler.py); diff --git a/src/schemas/api.py b/src/schemas/api.py index 237c4e14..78e5a125 100644 --- a/src/schemas/api.py +++ b/src/schemas/api.py @@ -29,6 +29,7 @@ from src.schemas.configuration import ( SessionPeerConfig, WorkspaceConfiguration, ) +from src.utils.types import DocumentLevel # --------------------------------------------------------------------------- # Metadata validation helpers @@ -176,6 +177,15 @@ class PeerRepresentationGet(BaseModel): session_id: str | None = Field( None, description="Optional session ID within which to scope the representation" ) + filters: dict[str, Any] | None = Field( + None, + description=( + "Optional filters to scope the representation. This endpoint " + "supports only the 'session_id' key: a session id, a list of " + 'session ids, or {"in": [...]}. When session_id is also set, it ' + "must be included in the allowlist." + ), + ) target: str | None = Field( None, description="Optional peer ID to get the representation for, from the perspective of this peer", @@ -446,6 +456,14 @@ class Conclusion(BaseModel): serialization_alias="observed_id", ) session_name: str | None = Field(default=None, serialization_alias="session_id") + level: DocumentLevel = Field( + default="explicit", + description=( + "Reasoning level of the conclusion: 'explicit' (directly extracted " + "from messages) or 'deductive'/'inductive'/'contradiction' (derived " + "during dreaming)." + ), + ) created_at: datetime.datetime model_config = ConfigDict( # pyright: ignore @@ -553,6 +571,16 @@ class DialecticOptions(BaseModel): session_id: str | None = Field( None, description="ID of the session to scope the representation to" ) + filters: dict[str, Any] | None = Field( + None, + description=( + "Optional filters to scope recall. This endpoint supports only the " + "'session_id' key: a session id, a list of session ids, or " + '{"in": [...]}. Recall (conclusions and messages) is restricted to ' + "the allowlist; unsupported keys are rejected. When session_id is " + "also set, it must be included in the allowlist." + ), + ) target: str | None = Field( None, description="Optional peer to get the representation for, from the perspective of this peer", @@ -565,6 +593,16 @@ class DialecticOptions(BaseModel): default="low", description="Level of reasoning to apply: minimal, low, medium, high, or max", ) + response_format: dict[str, Any] | None = Field( + None, + description=( + "Optional JSON Schema (root type 'object') the response must conform" + " to. When provided, `content` is a JSON string matching this schema." + " Only a conservative subset of JSON Schema is supported; unsupported" + " schemas are rejected with 422. Constraint keywords (minItems, " + " maxLength, ...) are hints to the model, not enforced server-side." + ), + ) @field_validator("query", mode="after") @classmethod @@ -651,6 +689,14 @@ class ScheduleDreamRequest(BaseModel): session_id: str | None = Field( None, description="Session ID to scope the dream to if specified" ) + rebuild: bool = Field( + False, + description=( + "card_refresh dreams only: rebuild the peer card solely from " + "observations currently in the collection, without injecting the " + "existing card (use after removals)" + ), + ) # --------------------------------------------------------------------------- @@ -681,10 +727,17 @@ class WebhookEndpointCreate(WebhookEndpointBase): if parsed.hostname: try: ip_address = ipaddress.ip_address(parsed.hostname) - if ip_address.is_private: - raise ValueError("Private IP addresses are not allowed") - except ValueError: # Not an IP address, might be a hostname - pass + except ValueError: # Not an IP literal — a hostname, leave it alone + ip_address = None + if ip_address is not None and ( + ip_address.is_private + or ip_address.is_loopback + or ip_address.is_link_local + or ip_address.is_reserved + or ip_address.is_multicast + or ip_address.is_unspecified + ): + raise ValueError("Private/internal IP addresses are not allowed") return v diff --git a/src/schemas/configuration.py b/src/schemas/configuration.py index b8291cab..53a8c11a 100644 --- a/src/schemas/configuration.py +++ b/src/schemas/configuration.py @@ -17,6 +17,10 @@ class DreamType(str, Enum): """Types of dreams that can be triggered.""" OMNI = "omni" + # Lightweight card-only refresh: runs a single specialist restricted to + # peer-card tools. Used for event-driven refreshes (scope membership + # changes, cold starts) — never creates or deletes observations. + CARD_REFRESH = "card_refresh" class ReasoningConfiguration(BaseModel): diff --git a/src/security.py b/src/security.py index 6988919e..3a60d8c1 100644 --- a/src/security.py +++ b/src/security.py @@ -80,6 +80,28 @@ def create_jwt(params: JWTParams) -> str: ) +def scope_requires_workspace( + *, peer: str | None, session: str | None, workspace: str | None +) -> bool: + """Return whether a peer- or session-scoped claim lacks its parent workspace. + + A peer or session scope is meaningless without a workspace: the route-level + check cannot rule out cross-workspace use (a ``{p: "alice"}`` token would + match ``alice`` in any workspace). Truthiness-based so empty-string claims + count as absent. Shared by `verify_jwt` (the token-shape invariant) and the + keys API (the creation-time guard) so the two rules cannot drift apart. + + Args: + peer: The peer claim, if any. + session: The session claim, if any. + workspace: The workspace claim, if any. + + Returns: + True when a peer/session scope is present but the workspace is not. + """ + return bool(peer or session) and not workspace + + def verify_jwt(token: str) -> JWTParams: """Verify a JWT and return the decoded parameters.""" @@ -101,12 +123,23 @@ def verify_jwt(token: str) -> JWTParams: raise AuthenticationException("JWT expired") if "ad" in decoded: params.ad = decoded["ad"] + # Normalize empty-string scope claims to None so a blank `w`/`p`/`s` + # cannot masquerade as a present claim in the checks below. if "w" in decoded: - params.w = decoded["w"] + params.w = decoded["w"] or None if "p" in decoded: - params.p = decoded["p"] + params.p = decoded["p"] or None if "s" in decoded: - params.s = decoded["s"] + params.s = decoded["s"] or None + # Token-shape invariant: a peer- or session-scoped token MUST also + # carry its parent workspace, otherwise the route-level check cannot + # rule out cross-workspace use. + if scope_requires_workspace( + peer=params.p, session=params.s, workspace=params.w + ): + raise AuthenticationException( + "Invalid JWT scope: peer/session token missing workspace" + ) return params except jwt.PyJWTError: raise AuthenticationException("Invalid JWT") from None @@ -117,9 +150,14 @@ def require_auth( workspace_name: str | None = None, peer_name: str | None = None, session_name: str | None = None, + allow_member_read: bool = False, ): """ Generate a dependency that requires authentication for the given parameters. + + Set `allow_member_read=True` on read-only session routes to additionally + grant access to peer-scoped keys whose peer is an active member of the + session. Never set it on routes that mutate state. """ async def auth_dependency( @@ -150,8 +188,14 @@ def require_auth( workspace_name=workspace_name_param, peer_name=peer_name_param, session_name=session_name_param, + allow_member_read=allow_member_read, ) + # Tag the closure so route-policy tests can introspect which routes opt into + # member read without re-deriving it from HTTP method (an unreliable + # read/write signal here — some read routes use POST for a richer body). + auth_dependency.honcho_allow_member_read = allow_member_read # pyright: ignore[reportFunctionMemberAccess] + return auth_dependency @@ -161,6 +205,7 @@ async def auth( workspace_name: str | None = None, peer_name: str | None = None, session_name: str | None = None, + allow_member_read: bool = False, ) -> JWTParams: """Authenticate the given JWT and return the decoded parameters.""" if not settings.AUTH.USE_AUTH: @@ -171,30 +216,66 @@ async def auth( jwt_params = verify_jwt(credentials.credentials) - # based on api operation, verify api key based on that key's permissions + # Authorize by the token's narrowest scope, not by the route's. A + # narrower-than-workspace token must NOT fall back to workspace access: + # `{w: ws, p: alice}` may only act on `alice`, never on a sibling peer. if jwt_params.ad: return jwt_params if admin: raise AuthenticationException("Resource requires admin privileges") - # For session level access - if session_name and jwt_params.s == session_name: - if workspace_name and jwt_params.w != workspace_name: - raise AuthenticationException("JWT not permissioned for this resource") + if not any([session_name, peer_name, workspace_name]): + # Self-authorizing routes decode the token here and compare the claims + # against body/path data inside the handler. This is needed for routes + # whose resource identifier is not available to require_auth(). return jwt_params - # For peer level access - if peer_name and jwt_params.p == peer_name: - if workspace_name and jwt_params.w != workspace_name: - raise AuthenticationException("JWT not permissioned for this resource") - return jwt_params - - # For workspace level access - can access all peers/sessions under this workspace - if workspace_name and jwt_params.w == workspace_name: - return jwt_params - - if any([session_name, peer_name, workspace_name]): + # Every scoped, non-admin path requires the token's workspace to match the + # route's. Check it once here so no individual branch below can forget it + # and silently re-open cross-workspace access (the bug this module fixes). + if workspace_name and jwt_params.w != workspace_name: raise AuthenticationException("JWT not permissioned for this resource") - # Route did not specify any parameters, so it should parse parameters itself - return jwt_params + if jwt_params.s is not None: + # Session-scoped token: confined to its own session. It gets no + # cross-scope access to peer routes. + if not session_name or jwt_params.s != session_name: + raise AuthenticationException("JWT not permissioned for this resource") + return jwt_params + + if jwt_params.p is not None: + # Peer-scoped token: its own peer routes... + if peer_name and jwt_params.p == peer_name: + return jwt_params + # ...plus read-only access to the sessions the peer is a member of. + # Gated on `allow_member_read` so only read routes opt in; writes stay + # denied. Requires the route's workspace so the membership lookup is + # scoped (every session route declares workspace_name); the workspace + # match itself was already verified above. + if allow_member_read and session_name and workspace_name: + # Lazy imports avoid an import cycle with the crud/db layers and + # keep this DB round-trip off the common (same-scope) auth path. + from src.crud.session import is_peer_in_session + from src.dependencies import tracked_db + + # Membership is read on a separate committed-only (read_only) + # connection, so a peer added to the session in a not-yet-committed + # transaction reads as a non-member: writes must commit before a + # member-scoped read. Fails closed. + async with tracked_db( + "auth.is_peer_in_session", read_only=True + ) as member_db: + is_member = await is_peer_in_session( + member_db, workspace_name, session_name, jwt_params.p + ) + if is_member: + return jwt_params + raise AuthenticationException("JWT not permissioned for this resource") + + if jwt_params.w is not None: + # Workspace tokens reach any route inside their workspace (the workspace + # match was verified above). Routes without a declared workspace (e.g. + # POST /v3/workspaces) self-authorize by reading jwt_params.w themselves. + return jwt_params + + raise AuthenticationException("JWT not permissioned for this resource") diff --git a/src/telemetry/__init__.py b/src/telemetry/__init__.py index 401be523..03cc9e0c 100644 --- a/src/telemetry/__init__.py +++ b/src/telemetry/__init__.py @@ -12,13 +12,18 @@ This module consolidates all telemetry, metrics, and observability functionality """ from src.telemetry.events import emit -from src.telemetry.prometheus import metrics_endpoint, prometheus_metrics +from src.telemetry.prometheus import ( + metrics_endpoint, + prometheus_metrics, + register_db_pool_collector, +) __all__ = [ "emit", "initialize_telemetry_async", "metrics_endpoint", "prometheus_metrics", + "register_db_pool_collector", "shutdown_telemetry", ] @@ -37,6 +42,8 @@ async def initialize_telemetry_async() -> None: from src.config import settings from src.telemetry.events import initialize_telemetry_events + # Master switch for every trace sink, Langfuse included: telemetry off + # initializes nothing. if settings.TELEMETRY.ENABLED: await initialize_telemetry_events() @@ -47,8 +54,17 @@ async def shutdown_telemetry() -> None: This should be called during application shutdown to ensure: - CloudEvents buffer is flushed + - Langfuse's buffered observations are flushed (its background exporter and + atexit hook don't fire reliably on SIGTERM) """ from src.telemetry.events import shutdown_telemetry_events + from src.telemetry.logging import flush_langfuse - # Shutdown CloudEvents emitter (flushes buffer) - await shutdown_telemetry_events() + # Flush Langfuse even if the CloudEvents shutdown raises, so the final + # batch of spans isn't dropped on a noisy shutdown. + try: + # Shutdown CloudEvents emitter (flushes buffer) + await shutdown_telemetry_events() + finally: + # Flush any buffered Langfuse spans before the process exits. + flush_langfuse() diff --git a/src/telemetry/emitter.py b/src/telemetry/emitter.py index 2413c7c2..201d16aa 100644 --- a/src/telemetry/emitter.py +++ b/src/telemetry/emitter.py @@ -100,6 +100,7 @@ class TelemetryEmitter: max_retries: int max_buffer_size: int enabled: bool + drop_reason_prefix: str _buffer: deque[CloudEvent] _flush_task: asyncio.Task[None] | None _client: httpx.AsyncClient | None @@ -118,6 +119,7 @@ class TelemetryEmitter: max_retries: int = 3, max_buffer_size: int = 10000, enabled: bool = True, + drop_reason_prefix: str = "", ): """Initialize the telemetry emitter. @@ -130,6 +132,9 @@ class TelemetryEmitter: max_retries: Maximum retry attempts on failure max_buffer_size: Maximum events to buffer (oldest dropped if exceeded) enabled: Whether emission is enabled + drop_reason_prefix: Prefix for the dropped-event metric reason label + (e.g. "trace_") so a second emitter's drops are distinguishable + from the primary metrics emitter's in Prometheus. """ self.endpoint = endpoint self.headers = headers or {} @@ -139,6 +144,7 @@ class TelemetryEmitter: self.max_retries = max_retries self.max_buffer_size = max_buffer_size self.enabled = enabled and endpoint is not None + self.drop_reason_prefix = drop_reason_prefix self._buffer = deque(maxlen=max_buffer_size) self._flush_task = None @@ -292,7 +298,9 @@ class TelemetryEmitter: cloud_event = CloudEvent(attributes, body) if will_drop_oldest: - prometheus_metrics.record_telemetry_event_dropped(reason="buffer_full") + prometheus_metrics.record_telemetry_event_dropped( + reason=f"{self.drop_reason_prefix}buffer_full" + ) self._buffer.append(cloud_event) buffer_size = len(self._buffer) @@ -375,7 +383,7 @@ class TelemetryEmitter: for event in reversed(batch): if len(self._buffer) >= self.max_buffer_size: prometheus_metrics.record_telemetry_event_dropped( - reason="send_failed" + reason=f"{self.drop_reason_prefix}send_failed" ) self._buffer.appendleft(event) logger.warning( @@ -527,3 +535,54 @@ async def shutdown_emitter() -> None: if _emitter is not None: await _emitter.shutdown() _emitter = None + + +# Separate emitter for the full-fidelity trace stream (llm.call.traced / +# trace.content). Kept distinct from the metrics `_emitter` so a trace burst +# can never evict billing events from the metrics buffer. +_trace_emitter: TelemetryEmitter | None = None + + +def get_trace_emitter() -> TelemetryEmitter | None: + """Get the global trace-stream emitter instance (None when payload tracing off).""" + return _trace_emitter + + +async def initialize_trace_emitter( + endpoint: str | None = None, + headers: dict[str, str] | None = None, + batch_size: int = 100, + flush_interval_seconds: float = 1.0, + flush_threshold: int = 50, + max_retries: int = 3, + max_buffer_size: int = 10000, + enabled: bool = True, +) -> TelemetryEmitter: + """Initialize and start the global trace-stream emitter. + + Drops are recorded under the ``trace_`` reason prefix so they're + distinguishable from the metrics emitter's drops in Prometheus. + """ + global _trace_emitter + + _trace_emitter = TelemetryEmitter( + endpoint=endpoint, + headers=headers, + batch_size=batch_size, + flush_interval_seconds=flush_interval_seconds, + flush_threshold=flush_threshold, + max_retries=max_retries, + max_buffer_size=max_buffer_size, + enabled=enabled, + drop_reason_prefix="trace_", + ) + await _trace_emitter.start() + return _trace_emitter + + +async def shutdown_trace_emitter() -> None: + """Shutdown the global trace-stream emitter.""" + global _trace_emitter + if _trace_emitter is not None: + await _trace_emitter.shutdown() + _trace_emitter = None diff --git a/src/telemetry/events/__init__.py b/src/telemetry/events/__init__.py index 3da745dd..000e70b4 100644 --- a/src/telemetry/events/__init__.py +++ b/src/telemetry/events/__init__.py @@ -89,6 +89,11 @@ from src.telemetry.events.reconciliation import ( SyncVectorsCompletedEvent, ) from src.telemetry.events.representation import RepresentationCompletedEvent +from src.telemetry.events.trace import ( + EmbeddingCallTracedEvent, + LLMCallTracedEvent, + TraceContentEvent, +) logger = logging.getLogger(__name__) @@ -120,6 +125,11 @@ __all__ = [ "CallPurpose", "EmbeddingCallCompletedEvent", "EmbeddingCallPurpose", + # Trace (full-fidelity payload) events + "emit_trace", + "EmbeddingCallTracedEvent", + "LLMCallTracedEvent", + "TraceContentEvent", # Reconciliation events "SyncVectorsCompletedEvent", "CleanupStaleItemsCompletedEvent", @@ -172,6 +182,27 @@ def emit(event: BaseEvent) -> None: ) +def emit_trace(event: BaseEvent) -> None: + """Queue a payload-trace event on the SEPARATE trace emitter. + + Distinct from `emit()` so a trace burst can never evict billing events from + the metrics buffer. No-op when payload tracing is off (trace emitter None). + Best-effort — swallows failures so telemetry never breaks the LLM path. + """ + try: + from src.telemetry.emitter import get_trace_emitter + + emitter = get_trace_emitter() + if emitter is None: + logger.debug("Trace emitter not initialized, dropping trace event") + return + emitter.emit(event) + except Exception: # pragma: no cover - best-effort telemetry + logger.debug( + "Failed to emit trace event %s", type(event).__name__, exc_info=True + ) + + async def initialize_telemetry_events() -> None: """Initialize the telemetry events system based on configuration. @@ -188,6 +219,17 @@ async def initialize_telemetry_events() -> None: logger.info("CloudEvents telemetry disabled") return + # Langfuse as a projection over the captured LLM stream. Gated behind the + # telemetry master switch above, so disabling telemetry disables Langfuse + # too. Registering the exporter makes has_exporters() true, which is what + # turns on CapturedLLMCall building. + if settings.langfuse_exporter_enabled: + from src.llm.capture import register_exporter + from src.telemetry.langfuse_exporter import LangfuseExporter + + register_exporter(LangfuseExporter()) + logger.info("Langfuse exporter registered (LANGFUSE_EXPORTER_MODE=exporter)") + await initialize_emitter( endpoint=settings.TELEMETRY.ENDPOINT, headers=settings.TELEMETRY.HEADERS, @@ -203,6 +245,27 @@ async def initialize_telemetry_events() -> None: "CloudEvents telemetry initialized, endpoint: %s", settings.TELEMETRY.ENDPOINT ) + # Full-fidelity payload tracing — opt-in, separate emitter + content exporter. + if settings.TELEMETRY.TRACE_PAYLOADS_ENABLED: + from src.llm.capture import register_exporter + from src.telemetry.emitter import initialize_trace_emitter + from src.telemetry.trace_exporter import TraceExporter + + await initialize_trace_emitter( + endpoint=settings.TELEMETRY.ENDPOINT, + headers=settings.TELEMETRY.HEADERS, + batch_size=settings.TELEMETRY.BATCH_SIZE, + flush_interval_seconds=settings.TELEMETRY.FLUSH_INTERVAL_SECONDS, + flush_threshold=settings.TELEMETRY.FLUSH_THRESHOLD, + max_retries=settings.TELEMETRY.MAX_RETRIES, + max_buffer_size=settings.TELEMETRY.MAX_BUFFER_SIZE, + enabled=True, + ) + register_exporter(TraceExporter()) + logger.info( + "Payload tracing initialized, endpoint: %s", settings.TELEMETRY.ENDPOINT + ) + async def shutdown_telemetry_events() -> None: """Shutdown the telemetry events system. @@ -210,7 +273,16 @@ async def shutdown_telemetry_events() -> None: This should be called during application shutdown to ensure all buffered events are flushed before exit. """ - from src.telemetry.emitter import shutdown_emitter + # Tear down the trace path first (flush its buffer, drop exporters + dedup + # state) before the primary emitter, so a late capture can't re-register work. + from src.llm.capture import clear_exporters + from src.telemetry import langfuse_session, trace_session + from src.telemetry.emitter import shutdown_emitter, shutdown_trace_emitter + + clear_exporters() + await shutdown_trace_emitter() + trace_session.reset() + langfuse_session.reset() await shutdown_emitter() logger.info("CloudEvents telemetry shutdown complete") diff --git a/src/telemetry/events/agent.py b/src/telemetry/events/agent.py index df771709..95145608 100644 --- a/src/telemetry/events/agent.py +++ b/src/telemetry/events/agent.py @@ -191,12 +191,18 @@ class AgentToolSummaryCreatedEvent(BaseEvent): """ _event_type: ClassVar[str] = "agent.tool.summary.created" - _schema_version: ClassVar[int] = 2 + _schema_version: ClassVar[int] = 3 _category: ClassVar[str] = "agent" - # Run identification (may be placeholder if not from an agentic loop) - run_id: str = Field(..., description="Nanoid for run correlation") - iteration: int = Field(..., description="Iteration number when this occurred") + # Run identification. + run_id: str | None = Field( + default=None, + description="Run id for agentic correlation; None when not in a run", + ) + iteration: int | None = Field( + default=None, + description="Iteration within an agentic loop; None when not in one", + ) # Context parent_category: str = Field(..., description="Parent category") @@ -258,8 +264,9 @@ class AgentToolSummaryCreatedEvent(BaseEvent): ) def get_resource_id(self) -> str: - """Resource ID includes run_id and iteration for uniqueness.""" - return f"{self.run_id}:{self.iteration}:summary_created" + """Idempotency key. A summary is unique per (message it covers up to, + tier)""" + return f"{self.message_id}:{self.summary_type}:summary_created" class AgentToolCallCompletedEvent(BaseEvent): diff --git a/src/telemetry/events/llm.py b/src/telemetry/events/llm.py index 9cbe7901..1761e333 100644 --- a/src/telemetry/events/llm.py +++ b/src/telemetry/events/llm.py @@ -34,6 +34,7 @@ class CallPurpose(str, Enum): DIALECTIC_ANSWER = "dialectic.answer" DREAM_DEDUCTION = "dream.deduction" DREAM_INDUCTION = "dream.induction" + DREAM_CARD_REFRESH = "dream.card_refresh" SUMMARY_SHORT = "summary.short" SUMMARY_LONG = "summary.long" @@ -164,6 +165,9 @@ class EmbeddingCallPurpose(str, Enum): CREATE_OBSERVATIONS = "create_observations" VECTOR_SYNC = "vector_sync" SUMMARY = "summary" + # Pending MessageEmbedding rows from create_messages; embedding runs in the + # reconciler (not inline on the API path). Distinct from VECTOR_SYNC, which + # covers document re-embeds and other vector-store healing work. MESSAGE_CREATE = "message_create" # Added so previously-unattributed call sites land on a distinct slug # instead of None. Closed taxonomy — coordinate with analytics before diff --git a/src/telemetry/events/representation.py b/src/telemetry/events/representation.py index 77d6b292..ad6dfca4 100644 --- a/src/telemetry/events/representation.py +++ b/src/telemetry/events/representation.py @@ -100,11 +100,36 @@ class RepresentationCompletedEvent(BaseEvent): default=0, description="Estimated tokens for the system/scaffold portion of the prompt", ) + exact_dup_in_batch_count: int = Field( + default=0, + description="Number of documents produced in this representation that had the same normalized content", + ) + exact_dup_existing_count: int = Field( + default=0, + description=( + "Number of documents previously written that had a representation that had the same normalized " + "content as a document in this representation" + ), + ) + semantic_dup_rejected_count: int = Field( + default=0, + description=( + "Number of documents in this representation rejected because their cosine-similarity was high " + "for an existing document but were worse than the corresponding existing document" + ), + ) + semantic_dup_replaced_count: int = Field( + default=0, + description=( + "Number of documents in this representation that replaced existing documents because their " + "cosine-similarity was high and they were better than the corresponding existing document" + ), + ) # Cap configuration + hit flags () batch_max_tokens: int = Field( default=0, - description="settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS at fetch time", + description="settings.DERIVER.REPRESENTATION_BATCH_TARGET_INPUT_TOKENS at fetch time", ) max_input_tokens: int = Field( default=0, description="settings.DERIVER.MAX_INPUT_TOKENS at call time" diff --git a/src/telemetry/events/trace.py b/src/telemetry/events/trace.py new file mode 100644 index 00000000..37383bc0 --- /dev/null +++ b/src/telemetry/events/trace.py @@ -0,0 +1,161 @@ +"""Replay-grade payload events for full-fidelity LLM tracing. + +Two ground-truth (never-sampled) events that make the CloudEvents stream carry +the exact context a model saw, content-addressed to keep payload O(N): + +- ``LLMCallTracedEvent`` (``llm.call.traced``) — one per LLM call. Carries + span-tree correlation, path identity, content *references* (hashes, not bytes) + for the context window and the replay-grade output, plus a self-contained + accounting copy. It deliberately does NOT claim a join to + ``llm.call.completed`` — cost is computable from the trace alone, so the + billing and audit streams stay decoupled. +- ``TraceContentEvent`` (``trace.content``) — one per unique message, emitted + once per run and referenced by hash. ``content_hash`` covers the full message + identity ({role, content, tool_call_id}) so identical text under different + roles can't collide. ``generate_id()`` is overridden to derive the CloudEvent + id from the hash with NO timestamp, so accidental re-sends of identical + content dedupe at the transport layer too. +""" + +from __future__ import annotations + +from typing import Any, ClassVar + +from pydantic import Field + +from src.config import ModelTransport +from src.telemetry.events.base import BaseEvent + +__all__ = ["EmbeddingCallTracedEvent", "LLMCallTracedEvent", "TraceContentEvent"] + + +class LLMCallTracedEvent(BaseEvent): + """One replay-grade record per LLM call. Ground-truth, never sampled. + + Content fields are *references* (content hashes) into the ``trace.content`` + store, never inline bytes + """ + + _event_type: ClassVar[str] = "llm.call.traced" + _schema_version: ClassVar[int] = 1 + _category: ClassVar[str] = "trace" + _volume_class: ClassVar[str] = "ground_truth" + + # --- Correlation (span tree) --- + trace_id: str | None = None + span_id: str | None = None + parent_span_id: str | None = None + iteration: int | None = None + step_seq: int = 0 + attempt: int = 1 + was_fallback: bool = False + parent_event_id: str | None = None + + # --- Path identity --- + call_purpose: str | None = None + parent_category: str | None = None + # Used for grouping traces + session_id: str | None = None + transport: ModelTransport + provider_label: str | None = None + model: str + + # --- Context window (content-addressed) --- + input_message_refs: list[str] = Field(default_factory=list) + system_prompt_ref: str | None = None + tool_schema_refs: list[str] = Field(default_factory=list) + tool_choice: Any = None + + # --- Output (replay-grade) --- + output_content_ref: str | None = None + output_tool_calls: list[dict[str, Any]] = Field(default_factory=list) + output_thinking_ref: str | None = None + output_signatures: list[str] = Field(default_factory=list) + # Reserved: Honcho captures the normalized request/response, not wire bytes. + raw_response_ref: str | None = None + finish_reason: str | None = None + + # --- Accounting copy (stream stands alone; NOT joined to llm.call.completed) --- + provider_input_tokens: int = 0 + provider_output_tokens: int = 0 + cache_read_tokens: int = 0 + cache_creation_tokens: int = 0 + was_truncated: bool = False + + def get_resource_id(self) -> str: + """Idempotency key. ``tool_call_seq`` is deliberately absent — it indexed + tool *executions*, not LLM calls, and was never a valid join field.""" + return f"{self.span_id}:{self.iteration}:{self.attempt}:{self.step_seq}" + + +class EmbeddingCallTracedEvent(BaseEvent): + """One trace-stream record per embedding-provider call.""" + + _event_type: ClassVar[str] = "embedding.call.traced" + _schema_version: ClassVar[int] = 1 + _category: ClassVar[str] = "trace" + _volume_class: ClassVar[str] = "ground_truth" + + # --- Correlation (span tree) --- + trace_id: str | None = None + span_id: str | None = None + parent_span_id: str | None = None + iteration: int | None = None + step_seq: int = 0 + attempt: int = 1 + session_id: str | None = None + + # --- Path identity --- + call_purpose: str | None = None + parent_category: str | None = None + provider: str + model: str + + # --- Accounting copy --- + # v1: input tokens are a tiktoken ESTIMATE (no authoritative provider count is + # plumbed yet); output tokens are always 0 (embeddings produce none). + provider_input_tokens: int = 0 + provider_output_tokens: int = 0 + input_count: int = 0 + was_truncated: bool = False + + def get_resource_id(self) -> str: + return f"{self.span_id}:embedding:{self.call_purpose}:{self.input_count}" + + +class TraceContentEvent(BaseEvent): + """One unique message in the content store. Ground-truth, never sampled. + + The hash covers the full message identity, and the event id derives from the + hash with no timestamp. + """ + + _event_type: ClassVar[str] = "trace.content" + _schema_version: ClassVar[int] = 1 + _category: ClassVar[str] = "trace" + _volume_class: ClassVar[str] = "ground_truth" + + content_hash: str + role: str + # Message text, normalized across providers. + content: Any = None + tool_call_id: str | None = None + # Tool calls in a unified {id, name, input} shape (provider-agnostic). + tool_calls: list[dict[str, Any]] = Field(default_factory=list) + # Tags Honcho-authored content (system prompts, scaffold) so tenant-facing + # views can withhold globally-shared content (the §6.3 access invariant — + # dedup is global, the content store has no tenant column). + honcho_authored: bool = False + + def get_resource_id(self) -> str: + return self.content_hash + + def generate_id(self) -> str: + """Content-addressed id with NO timestamp/version. + + Overrides the base (which folds timestamp + honcho_version) so any + cross-process or cross-retry re-send of the same content collides on + the same id and dedupes at the transport layer. + """ + digest = self.content_hash.split(":", 1)[-1] + return f"content_{digest[:22]}" diff --git a/src/telemetry/langfuse_exporter.py b/src/telemetry/langfuse_exporter.py new file mode 100644 index 00000000..e66eb8ac --- /dev/null +++ b/src/telemetry/langfuse_exporter.py @@ -0,0 +1,398 @@ +"""Langfuse projection over the captured LLM trace stream. + +`LangfuseExporter` is an `LLMCallExporter`, active when +`LANGFUSE_EXPORTER_MODE == "exporter"`. It receives one `CapturedLLMCall` at a +time and rebuilds the Langfuse trace tree from the ids on each call, since there +is no live span nesting to inherit: + + Trace (id = create_trace_id(seed=honcho trace_id)) + └─ [dream root] (multi-specialist agents only — one "Dream" span per trace) + └─ run span (one per (run_id, agent_type); name = track_name) + └─ step span (one per (agent_type, iteration); name = " step") + ├─ generation (one per CapturedLLMCall; name = " generation") + └─ tool span (one per requested tool call; sibling of generation) + +Run and step spans are created once per trace and reused as the `parent_span_id` +of later calls (tracked in `langfuse_session`). Single-shot callers +(deriver/summarizer, `run_id is None`) skip the run/step wrappers and put the +generation at the trace root. + +The Dreamer runs two specialists (deduction + induction) under one run_id, so its +branches hang off a single synthetic "Dream" root to keep the trace +single-rooted; single-specialist agents (dialectic) let their run span be the +root. Keeping exactly one root is also why child spans are demoted from the SDK's +auto-root flag (see `_demote_from_root`). + +Best-effort throughout: every export is wrapped so telemetry can never break the +LLM call path. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from src.config import settings +from src.llm.capture import CapturedLLMCall +from src.telemetry import langfuse_session + +logger = logging.getLogger(__name__) + +# finish_reason values that mark the generation as failed. +_ERROR_FINISHES = frozenset({"error", "cancelled"}) + + +class LangfuseExporter: + """`LLMCallExporter` that projects captured calls onto Langfuse traces.""" + + def export(self, call: CapturedLLMCall) -> None: + if not settings.langfuse_exporter_enabled: + return + try: + self._export(call) + except Exception: # pragma: no cover - best-effort telemetry + logger.debug("Langfuse exporter failed", exc_info=True) + + def _export(self, call: CapturedLLMCall) -> None: + from langfuse import get_client + + client = get_client() + seed = call.trace_id or call.run_id or call.span_id + if not seed: + return + lf_trace_id = client.create_trace_id(seed=seed) + + # Agentic runs (run_id set: dialectic / dreamer) get a run span and + # per-iteration step spans; single-shot calls put the generation at root. + # Branch = agent_type so co-trace specialists (dreamer) don't collide. + parent_span_id: str | None = None + if call.run_id is not None: + branch = call.agent_type or "_" + # Multi-specialist agents (the Dreamer runs deduction + induction in + # ONE trace) hang every branch off a single synthetic trace root, so + # the trace has one root instead of one per specialist. That root + # also stamps the trace attrs. Single-specialist agents (dialectic) + # get None here and let their run span be the root. + root_span_id = self._ensure_trace_root(client, lf_trace_id, call) + run_span_id = langfuse_session.ensure_run_span( + lf_trace_id, + branch, + lambda should_stamp: self._create_span( + client, + lf_trace_id, + parent_span_id=root_span_id, + name=call.track_name or "LLM run", + metadata=self._metadata(call), + # The synthetic root stamps the trace attrs when present; + # otherwise the first branch's run span does. + stamp_trace=should_stamp and root_span_id is None, + call=call, + ), + ) + parent_span_id = run_span_id + if call.iteration is not None and run_span_id is not None: + parent_span_id = langfuse_session.ensure_step_span( + lf_trace_id, + branch, + call.iteration, + lambda: self._create_span( + client, + lf_trace_id, + parent_span_id=run_span_id, + name=self._step_name(call), + metadata=self._step_metadata(call), + stamp_trace=False, + call=call, + ), + ) + + self._create_generation( + client, + lf_trace_id, + parent_span_id=parent_span_id, + # Single-shot: the generation is the trace root, so it stamps the + # trace attrs. Agentic: the first branch's run span already did. + stamp_trace=call.run_id is None, + call=call, + ) + + # Tool calls the model requested this iteration: siblings of the + # generation under the step span. Skipped at the trace root (single-shot + # callers don't use tools) since there's no step to anchor them. + if parent_span_id is not None and call.output_tool_calls: + for seq, tool_call in enumerate(call.output_tool_calls): + self._create_tool_span( + client, lf_trace_id, parent_span_id, seq, tool_call, call + ) + + # -- observation builders ------------------------------------------------ + + def _ensure_trace_root( + self, client: Any, lf_trace_id: str, call: CapturedLLMCall + ) -> str | None: + """Single branch-agnostic trace root for multi-specialist agents. + + The Dreamer's deduction + induction specialists share one trace (same + run_id) but each builds its own run span with no parent — so Langfuse + sees two roots, races the trace name between them, and renders the + specialists as separate sub-traces. One synthetic "Dream" root (which + also stamps the trace attrs) gives the trace a single root with both + specialists nested beneath. Single-specialist agents (dialectic) return + None and let their run span be the root. + """ + if call.parent_category != "dream": + return None + return langfuse_session.ensure_trace_root( + lf_trace_id, + lambda: self._create_span( + client, + lf_trace_id, + parent_span_id=None, + name=self._trace_name(call) or "Dream", + metadata=self._root_metadata(call), + stamp_trace=True, + call=call, + ), + ) + + def _create_span( + self, + client: Any, + lf_trace_id: str, + *, + parent_span_id: str | None, + name: str, + metadata: dict[str, str], + stamp_trace: bool, + call: CapturedLLMCall, + ) -> str | None: + """Create a (run or step) span, returning its OTEL span id. + + Created-and-ended immediately: nesting is by id, so children link fine to + an already-ended parent. Span durations are therefore approximate — an + accepted v1 trade for not having a 'run finished' signal in the stream. + """ + obs = client.start_observation( + trace_context=self._trace_context(lf_trace_id, parent_span_id), + name=name, + as_type="span", + metadata=metadata, + ) + if stamp_trace: + self._stamp_trace_attrs(obs, call) + if parent_span_id is not None: + self._demote_from_root(obs) + obs.end() + return getattr(obs, "id", None) + + def _create_generation( + self, + client: Any, + lf_trace_id: str, + *, + parent_span_id: str | None, + stamp_trace: bool, + call: CapturedLLMCall, + ) -> None: + level = "ERROR" if (call.finish_reason in _ERROR_FINISHES) else None + obs = client.start_observation( + trace_context=self._trace_context(lf_trace_id, parent_span_id), + name=self._gen_name(call), + as_type="generation", + model=call.model, + input=self._input(call), + output=self._output(call), + metadata=self._step_metadata(call), + usage_details=self._usage(call), + level=level, + ) + if stamp_trace: + self._stamp_trace_attrs(obs, call) + if parent_span_id is not None: + self._demote_from_root(obs) + obs.end() + + def _create_tool_span( + self, + client: Any, + lf_trace_id: str, + parent_span_id: str, + seq: int, + tool_call: dict[str, Any], + call: CapturedLLMCall, + ) -> None: + """Create a tool span for one requested tool call, under the step span. + + Built from the model's request (`output_tool_calls`): tool name + input + args. Result/duration/error aren't on the captured call (they live on + AgentToolCallCompletedEvent) — a later enrichment, not v1. + """ + obs = client.start_observation( + trace_context=self._trace_context(lf_trace_id, parent_span_id), + name=str(tool_call.get("name") or "tool"), + as_type="tool", + input=tool_call.get("input"), + metadata=self._tool_metadata(call, seq), + ) + self._demote_from_root(obs) # always a child of the step span + obs.end() + + @staticmethod + def _demote_from_root(obs: Any) -> None: + """Clear the AS_ROOT flag the SDK auto-stamps on a child observation. + + `start_observation(trace_context={"trace_id": ...})` marks EVERY span it + mints with `AS_ROOT=True` (langfuse `_client/client.py`) — including the + step/generation/tool spans we link under a run span by id. With several + root-flagged spans in one trace, Langfuse resolves the trace's root (and + therefore its name) from whichever it ingests first: a race that names a + dialectic trace after a child ("... step"/"... generation") and renders + children as if each were its own trace. Demoting every span that has a + real parent leaves exactly one root, making name + nesting deterministic. + Verified empirically against Langfuse cloud (the dangling remote-parent + id on the surviving root is benign and unavoidable — it's present even + with native context nesting). + """ + span = getattr(obs, "_otel_span", None) + if span is None: + return + from langfuse import LangfuseOtelSpanAttributes as Attr + + span.set_attribute(Attr.AS_ROOT, False) + + @staticmethod + def _trace_context(lf_trace_id: str, parent_span_id: str | None) -> dict[str, str]: + ctx: dict[str, str] = {"trace_id": lf_trace_id} + if parent_span_id is not None: + ctx["parent_span_id"] = parent_span_id + return ctx + + def _stamp_trace_attrs(self, obs: Any, call: CapturedLLMCall) -> None: + """Stamp user/name on the trace via the root observation's span. + + Called once per trace, on the first branch's run span (decided by + `langfuse_session.ensure_run_span`) or, for single-shot calls, on the + generation (its own trace). + + Deliberately does NOT set a Langfuse session: no Honcho construct is a + conversation thread. A dialectic chat is a one-shot query scoped to a + session, not a turn in a multi-turn dialectic exchange (no such primitive + exists), so grouping independent queries under one Langfuse session would + invent a conversation that isn't there. The Honcho session rides in + metadata (`honcho_session`) instead — a correlation key, not a group.""" + span = getattr(obs, "_otel_span", None) + if span is None: + return + from langfuse import LangfuseOtelSpanAttributes as Attr + + span.set_attribute(Attr.TRACE_USER_ID, str(settings.NAMESPACE)) + trace_name = self._trace_name(call) + if trace_name: + span.set_attribute(Attr.TRACE_NAME, trace_name) + + # -- field mappers (port of runtime._base_metadata/_step_metadata) ------- + + @staticmethod + def _metadata(call: CapturedLLMCall) -> dict[str, str]: + # `trace_id` is the run grouping key (also handy for cross-referencing the + # CloudEvents stream). `span_id`/`parent_span_id` are intentionally omitted + # until the source mints distinct per-call span ids: today every call in a + # run shares span_id == trace_id == run_id, so surfacing them here only + # duplicates trace_id and misleads. Re-add once the source differentiates. + md: dict[str, str] = {"namespace": str(settings.NAMESPACE)} + for key, value in ( + ("workspace_name", call.workspace_name), + ("call_purpose", call.call_purpose), + ("agent_type", call.agent_type), + ("observer", call.observer), + ("observed", call.observed), + ("peer_name", call.peer_name), + ("trace_id", call.trace_id), + # Honcho session as a correlation key, NOT a Langfuse session — see + # `_stamp_trace_attrs`. Lets you filter "queries scoped to session X" + # without falsely grouping one-shot dialectic queries as a thread. + ("honcho_session", call.session_id), + ): + if value is not None: + md[key] = str(value) + return md + + @staticmethod + def _root_metadata(call: CapturedLLMCall) -> dict[str, str]: + # Branch-agnostic: the synthetic dream root spans both specialists, so it + # carries only trace-level fields — not a single specialist's agent_type/ + # observer/observed/call_purpose. + md: dict[str, str] = {"namespace": str(settings.NAMESPACE)} + for key, value in ( + ("workspace_name", call.workspace_name), + ("trace_id", call.trace_id), + ): + if value is not None: + md[key] = str(value) + return md + + def _step_metadata(self, call: CapturedLLMCall) -> dict[str, str]: + md = self._metadata(call) + if call.iteration is not None: + md["iteration"] = str(call.iteration) + md["step_seq"] = str(call.step_seq) + md["attempt"] = str(call.attempt) + md["provider"] = str(call.transport) + md["model"] = str(call.model) + return md + + def _tool_metadata(self, call: CapturedLLMCall, seq: int) -> dict[str, str]: + md = self._step_metadata(call) + md["tool_call_seq"] = str(seq) + return md + + @staticmethod + def _trace_name(call: CapturedLLMCall) -> str | None: + # Branch-agnostic trace label: the Dreamer's two specialists share one + # trace, so the trace name must not be pinned to whichever specialist's + # run span stamped it first. Per-branch identity stays on the run spans. + if call.parent_category == "dream": + return "Dream" + return call.track_name + + @staticmethod + def _step_name(call: CapturedLLMCall) -> str: + # Canonical, index-free name: Langfuse aggregates step spans by name and + # the iteration/step_seq/attempt ride on metadata (see _step_metadata). + return f"{call.track_name} step" if call.track_name else "Agent step" + + @staticmethod + def _gen_name(call: CapturedLLMCall) -> str: + return f"{call.track_name} generation" if call.track_name else "generation" + + @staticmethod + def _input(call: CapturedLLMCall) -> list[dict[str, Any]]: + out: list[dict[str, Any]] = [] + for message in call.input_messages: + entry: dict[str, Any] = {"role": message.role, "content": message.content} + if message.tool_call_id is not None: + entry["tool_call_id"] = message.tool_call_id + if message.tool_calls: + entry["tool_calls"] = message.tool_calls + out.append(entry) + return out + + @staticmethod + def _output(call: CapturedLLMCall) -> Any: + if isinstance(call.output_content, str) and call.output_content.strip(): + return call.output_content + if call.output_tool_calls: + return {"tool_calls": [tc.get("name") for tc in call.output_tool_calls]} + return call.output_content + + @staticmethod + def _usage(call: CapturedLLMCall) -> dict[str, int]: + return { + "input": call.input_tokens, + "output": call.output_tokens, + "cache_read_input_tokens": call.cache_read_tokens, + "cache_creation_input_tokens": call.cache_creation_tokens, + } + + +__all__ = ["LangfuseExporter"] diff --git a/src/telemetry/langfuse_session.py b/src/telemetry/langfuse_session.py new file mode 100644 index 00000000..3dc2c13e --- /dev/null +++ b/src/telemetry/langfuse_session.py @@ -0,0 +1,122 @@ +"""Per-trace span registry backing the `LangfuseExporter`. + +The exporter sees one `CapturedLLMCall` at a time, but a single agentic run fans +out into many calls that must nest under one run span with per-iteration step +spans. Langfuse links observations by OTEL span id, and each id is minted fresh +and unpredictable — so this module remembers the run/step span ids created for a +trace and hands them back as the `parent_span_id` of later calls. + +Spans are keyed per branch (the `agent_type`) within a trace. The Dreamer's +deduction and induction specialists share one trace but are separate sub-trees; +without the branch key their iterations and generations would collide. + +Per trace, it holds each branch's run span id, the per-(branch, iteration) step +span ids, and whether trace-level attrs have been stamped — so each is created +once. The stamp decision is made inside `ensure_run_span` under the lock so it +can't double-fire across branches. + +Bounded by an LRU over traces (`_MAX_TRACES`), lock-guarded, best-effort. +""" + +from __future__ import annotations + +import logging +import threading +from collections import OrderedDict +from collections.abc import Callable +from dataclasses import dataclass, field + +logger = logging.getLogger(__name__) + +# LRU window / runaway backstop — far more than the traces ever live at once; the +# least-recently-used trace is evicted past this. Not a tuning knob. +_MAX_TRACES = 4096 + + +@dataclass +class _TraceState: + root_span_id: str | None = None # synthetic trace root (multi-specialist agents) + run_span_ids: dict[str, str] = field(default_factory=dict) # branch -> span id + step_span_ids: dict[tuple[str, int], str] = field( + default_factory=dict + ) # (branch, iteration) -> span id + attrs_stamped: bool = False + + +_traces: OrderedDict[str, _TraceState] = OrderedDict() +_lock = threading.Lock() + + +def _get_or_create_state(trace_key: str) -> _TraceState: + """Return the `_TraceState` for `trace_key`, creating it if new and marking it + most-recently-used. Caller MUST hold `_lock`. Bounded by an LRU: a new trace + past `_MAX_TRACES` evicts the least-recently-used (almost always finished) one. + """ + state = _traces.get(trace_key) + if state is None: + if len(_traces) >= _MAX_TRACES: + _traces.popitem(last=False) # evict the least-recently-used trace + state = _traces[trace_key] = _TraceState() + else: + _traces.move_to_end(trace_key) # mark most-recently-used + return state + + +def ensure_trace_root(trace_key: str, create: Callable[[], str | None]) -> str | None: + """Return the single trace-root span id for `trace_key`, creating it once. + + Used by multi-specialist agents (the Dreamer) whose branches share one trace + but must all hang off ONE root span. Single-specialist agents don't call + this. Mirrors `ensure_run_span`'s retry-on-None: a failed create just yields + None (the caller then roots the branch directly) and is retried next call. + """ + with _lock: + state = _get_or_create_state(trace_key) + if state.root_span_id is None: + state.root_span_id = create() + return state.root_span_id + + +def ensure_run_span( + trace_key: str, branch: str, create: Callable[[bool], str | None] +) -> str | None: + """Return the run span id for `(trace_key, branch)`, creating it once. + + `create` receives `should_stamp` — True exactly once per trace, on the first + branch's run span — and builds the Langfuse run span (stamping trace-level + attrs iff asked), returning its span id (or None on failure). The stamp + decision is computed here, under the lock, so it can't double-fire across the + Dreamer's two specialist branches; `create` must not re-enter this module. + """ + with _lock: + state = _get_or_create_state(trace_key) + existing = state.run_span_ids.get(branch) + if existing is None: + should_stamp = not state.attrs_stamped + existing = create(should_stamp) + if existing is not None: + state.run_span_ids[branch] = existing + if should_stamp: + state.attrs_stamped = True + return existing + + +def ensure_step_span( + trace_key: str, branch: str, iteration: int, create: Callable[[], str | None] +) -> str | None: + """Return the step span id for `(trace_key, branch, iteration)`, creating once.""" + with _lock: + state = _get_or_create_state(trace_key) + key = (branch, iteration) + existing = state.step_span_ids.get(key) + if existing is None: + existing = create() + if existing is not None: + state.step_span_ids[key] = existing + return existing + + +def reset() -> None: + """Drop all tracked traces — used on shutdown and in tests.""" + with _lock: + _traces.clear() diff --git a/src/telemetry/logging.py b/src/telemetry/logging.py index 065e1ccf..9c5f6648 100644 --- a/src/telemetry/logging.py +++ b/src/telemetry/logging.py @@ -5,9 +5,10 @@ and a conditional observe decorator that only applies when Langfuse is configure """ import datetime +import logging from collections import OrderedDict from collections.abc import Callable -from typing import ParamSpec, TypeVar, overload +from typing import Literal, ParamSpec, TypeVar, overload from fastapi import Request from langfuse import observe @@ -16,7 +17,6 @@ from rich.console import Console, Group, RenderableType from rich.panel import Panel from rich.table import Table from rich.text import Text -from rich.tree import Tree from src.config import settings from src.telemetry.metrics_collector import append_metrics_to_file @@ -24,6 +24,8 @@ from src.utils.representation import ( Representation, ) +logger = logging.getLogger(__name__) + # Global console instance for consistent formatting console = Console(markup=True) @@ -32,6 +34,21 @@ COLLECT_METRICS_LOCAL = settings.COLLECT_METRICS_LOCAL P = ParamSpec("P") R = TypeVar("R") +# Langfuse observation types accepted by `@observe(as_type=...)`. Mirrors the +# literal union the SDK exposes; kept local so callers don't import langfuse +# internals just to name an observation type. +ObserveAsType = Literal[ + "generation", + "embedding", + "span", + "agent", + "tool", + "chain", + "retriever", + "evaluator", + "guardrail", +] + @overload def conditional_observe( @@ -42,7 +59,10 @@ def conditional_observe( @overload def conditional_observe( *, - name: str, + name: str | None = None, + as_type: ObserveAsType | None = None, + capture_input: bool | None = None, + capture_output: bool | None = None, ) -> Callable[[Callable[P, R]], Callable[P, R]]: ... @@ -50,28 +70,58 @@ def conditional_observe( func: Callable[P, R] | None = None, *, name: str | None = None, + as_type: ObserveAsType | None = None, + capture_input: bool | None = None, + capture_output: bool | None = None, ) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]: """ - Conditionally apply the @observe decorator only when LANGFUSE_PUBLIC_KEY is present. + Conditionally apply the @observe decorator only in legacy inline mode + (``langfuse_inline_enabled`` — i.e. a key is set AND + ``LANGFUSE_EXPORTER_MODE == "inline"``). In exporter mode the LangfuseExporter + rebuilds every observation from the captured trace stream, so a live + @observe span here would double-emit. Can be used in two ways: 1. As a decorator: @conditional_observe - 2. As a decorator factory: @conditional_observe(name="...") + 2. As a decorator factory: @conditional_observe(name="...", as_type="generation") Args: func: The function to potentially decorate (when used as @conditional_observe) name: Optional name for the observation (when used as @conditional_observe(name="...")) + as_type: Optional Langfuse observation type (e.g. "generation", "tool"). When + omitted, Langfuse infers a default span. + capture_input: When ``False``, Langfuse does NOT auto-serialize the + function's arguments into the span input. Set this on functions that + receive live SDK clients or secret-bearing config as parameters + (e.g. the LLM executor): auto-capture would deep-copy those clients + into throwaway, half-constructed objects whose GC raises + ``AsyncHttpxClientWrapper ... no attribute '_state'`` / + ``BaseApiClient ... no attribute '_http_options'`` (see HONCHO-4HA), + and would also leak ``ModelConfig.api_key`` into traces. Pair with an + explicit ``update_current_generation(input=...)`` call to keep + full-fidelity input. ``None`` leaves the SDK default (capture on). + capture_output: When ``False``, Langfuse does NOT auto-serialize the + return value. Pair with an explicit + ``update_current_generation(output=...)``. ``None`` = SDK default. Returns: The decorated function if Langfuse is configured, otherwise the original function """ def decorator(f: Callable[P, R]) -> Callable[P, R]: - if settings.LANGFUSE_PUBLIC_KEY: - observe_name = name if name is not None else f.__name__ - return observe(name=observe_name)(f) - else: + # Only auto-instrument with @observe in legacy inline mode. In exporter + # mode the LangfuseExporter produces every observation from the captured + # trace stream, so a live @observe span here would double-emit. + if not settings.langfuse_inline_enabled: return f + # `observe` treats None as "use SDK default", so passing the optionals + # straight through is equivalent to omitting them. + return observe( + name=name if name is not None else f.__name__, + as_type=as_type, + capture_input=capture_input, + capture_output=capture_output, + )(f) if func is not None: # Used as @conditional_observe (without parentheses) @@ -81,6 +131,22 @@ def conditional_observe( return decorator +def flush_langfuse() -> None: + """Flush buffered Langfuse spans on shutdown. + + The SDK's background timer/atexit hook don't fire reliably on SIGTERM, so + the final batch is dropped without this. No-op when Langfuse is unconfigured. + """ + if not settings.LANGFUSE_PUBLIC_KEY: + return + try: + from langfuse import get_client + + get_client().flush() + except Exception: + logger.debug("Failed to flush Langfuse on shutdown", exc_info=True) + + # Bounded OrderedDict for accumulated metrics to prevent memory leaks. # If an exception occurs between accumulate_metric() and log_performance_metrics(), # the metrics would stay in memory forever. Using OrderedDict allows us to evict @@ -136,28 +202,6 @@ def format_reasoning_inputs_as_markdown( return "\n".join(parts) -def log_representation( - representation: Representation, -) -> None: - """ - Log representation in a tree structure. - Args: - representation: Representation to log - """ - tree = Tree("📊 REPRESENTATION") - - type_branch = tree.add(f"[bold cyan]EXPLICIT[/] ({len(representation.explicit)})") - for i, obs in enumerate(representation.explicit, 1): - type_branch.add(f"[dim]{i}.[/] {obs}") - - type_branch = tree.add(f"[bold cyan]DEDUCTIVE[/] ({len(representation.deductive)})") - for i, obs in enumerate(representation.deductive, 1): - type_branch.add(f"[dim]{i}.[/] {obs}") - - console.print(tree) - console.print() - - def accumulate_metric( task_name: str, label: str, @@ -236,16 +280,20 @@ def log_performance_metrics( task_slug: str, task_name: str, metrics: list[tuple[str, str | int | float, str]] | None = None, - title: str = "⚡ PERFORMANCE", + title: str = "PERFORMANCE", ) -> None: """ - Log performance metrics in a clean table and optionally send to global collector. + Log performance metrics and optionally send them to the global collector. + + PERFORMANCE_LOG_FORMAT=compact emits numeric metrics on one INFO line and + keeps large "blob" metrics at DEBUG. PERFORMANCE_LOG_FORMAT=rich prints the + local Rich panel, including blob metrics, for interactive readability. Args: task_slug: Slug of the task that generated these metrics task_name: Name of the task that generated these metrics - metrics: Dictionary of metric names and (value, unit) tuples - title: Table title + metrics: List of (metric_name, value, unit) tuples + title: Prefix for the log line """ task_name = f"{task_slug}_{task_name}" # No-op if metrics were evicted (due to MAX_ACCUMULATED_TASKS limit) and no @@ -260,7 +308,41 @@ def log_performance_metrics( if COLLECT_METRICS_LOCAL: append_metrics_to_file(task_slug, task_name, metrics) - # Remove metrics with "blob" unit type. They get printed separately below the table. + if settings.PERFORMANCE_LOG_FORMAT == "rich": + _log_performance_metrics_rich(task_name, metrics, title) + return + + # Keep large text payloads out of the compact INFO summary. + blob_metrics: list[tuple[str, str | int | float, str]] = [] + summary_parts: list[str] = [] + for metric, value, unit in metrics: + if unit == "blob": + blob_metrics.append((metric, value, unit)) + continue + if unit == "ms" and isinstance(value, int | float): + formatted_value = f"{value:.0f}ms" + elif unit == "s" and isinstance(value, int | float): + formatted_value = f"{value:.3f}s" + elif unit in ("", "tokens", "count", "id"): + formatted_value = str(value) + else: + formatted_value = f"{value}{unit}" + summary_parts.append(f"{metric}={formatted_value}") + + if summary_parts: + logger.info("%s %s | %s", title, task_name, " | ".join(summary_parts)) + else: + logger.info("%s %s", title, task_name) + + for metric, value, _unit in blob_metrics: + logger.debug("%s %s :: %s\n%s", title, task_name, metric, value) + + +def _log_performance_metrics_rich( + task_name: str, + metrics: list[tuple[str, str | int | float, str]], + title: str, +) -> None: blob_metrics: list[tuple[str, str | int | float, str]] = [] non_blob_metrics: list[tuple[str, str | int | float, str]] = [] for metric in metrics: @@ -277,24 +359,20 @@ def log_performance_metrics( table.add_column("Unit", style="dim", width=8) for metric, value, unit in non_blob_metrics: - if unit == "ms": + if unit == "ms" and isinstance(value, int | float): formatted_value = f"{value:.0f}" - elif unit == "s": + elif unit == "s" and isinstance(value, int | float): formatted_value = f"{value:.3f}" else: formatted_value = str(value) table.add_row(metric.replace("_", " ").title(), formatted_value, unit) - # Build content for the panel content_items: list[RenderableType] = [table] - if blob_metrics: - for metric, value, _unit in blob_metrics: - content_items.append( - Text.assemble(" ", (f"\n{metric}:", "bold"), " ") - ) - content_items.append(Text(str(value))) + for metric, value, _unit in blob_metrics: + content_items.append(Text.assemble(" ", (f"\n{metric}:", "bold"), " ")) + content_items.append(Text(str(value))) panel = Panel( Group(*content_items), diff --git a/src/telemetry/prometheus/__init__.py b/src/telemetry/prometheus/__init__.py index ab0a2616..876ea110 100644 --- a/src/telemetry/prometheus/__init__.py +++ b/src/telemetry/prometheus/__init__.py @@ -14,6 +14,7 @@ from src.telemetry.prometheus.metrics import ( TokenTypes, metrics_endpoint, prometheus_metrics, + register_db_pool_collector, ) __all__ = [ @@ -23,4 +24,5 @@ __all__ = [ "TokenTypes", "metrics_endpoint", "prometheus_metrics", + "register_db_pool_collector", ] diff --git a/src/telemetry/prometheus/metrics.py b/src/telemetry/prometheus/metrics.py index 90082b85..01d7be4f 100644 --- a/src/telemetry/prometheus/metrics.py +++ b/src/telemetry/prometheus/metrics.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +from collections.abc import Iterator from enum import Enum from typing import cast, final @@ -11,9 +12,11 @@ from prometheus_client import ( REGISTRY, Counter, Gauge, + Histogram, disable_created_metrics, generate_latest, ) +from prometheus_client.core import GaugeMetricFamily from starlette.requests import Request from starlette.responses import Response @@ -36,6 +39,12 @@ class NamespacedGauge(Gauge): return super().labels(**kwargs) # type: ignore[return-value] +class NamespacedHistogram(Histogram): + def labels(self, **kwargs: str) -> NamespacedHistogram: + kwargs["namespace"] = cast(str, settings.METRICS.NAMESPACE) + return super().labels(**kwargs) # type: ignore[return-value] + + class TokenTypes(Enum): INPUT = "input" OUTPUT = "output" @@ -63,12 +72,33 @@ api_requests_counter = NamespacedCounter( ["namespace", "method", "endpoint", "status_code"], ) +# Per-route latency. Buckets are a geometric ladder spanning +# the full range of API classes +api_request_duration_seconds = NamespacedHistogram( + "api_request_duration_seconds", + "API request latency in seconds", + ["namespace", "method", "endpoint"], + buckets=(0.05, 0.1, 0.25, 0.5, 0.75, 1, 2, 5, 10, 20, 30, 60, 120), +) + messages_created_counter = NamespacedCounter( "messages_created", "Total messages created", ["namespace", "workspace_name"], ) +embed_now_tasks_shed_counter = NamespacedCounter( + "embed_now_tasks_shed", + "Immediate-embed background tasks skipped because MAX_PENDING_EMBED_TASKS was reached", + ["namespace"], +) + +embed_now_tasks_in_flight_gauge = NamespacedGauge( + "embed_now_tasks_in_flight", + "Immediate-embed background tasks currently in flight for this process", + ["namespace"], +) + dialectic_calls_counter = NamespacedCounter( "dialectic_calls", "Total dialectic calls", @@ -125,6 +155,15 @@ telemetry_buffer_size_gauge = NamespacedGauge( ["namespace"], ) +# DB connection-pool health. The in-flight gauge counts statements actually +# executing on the wire, so checked_out minus in_flight reveals connections held +# but parked (the "idle in transaction during an external call" antipattern). +db_queries_in_flight_gauge = NamespacedGauge( + "db_queries_in_flight", + "DB statements currently executing on a connection for this instance", + ["namespace", "instance_type"], +) + @final class PrometheusMetrics: @@ -149,6 +188,7 @@ class PrometheusMetrics: method: str, endpoint: str, status_code: str, + duration_seconds: float, ) -> None: try: api_requests_counter.labels( @@ -156,6 +196,10 @@ class PrometheusMetrics: endpoint=endpoint, status_code=status_code, ).inc() + api_request_duration_seconds.labels( + method=method, + endpoint=endpoint, + ).observe(duration_seconds) except Exception as e: self._handle_metric_error("record_api_request", e) @@ -172,6 +216,18 @@ class PrometheusMetrics: except Exception as e: self._handle_metric_error("record_messages_created", e) + def record_embed_now_task_shed(self) -> None: + try: + embed_now_tasks_shed_counter.labels().inc() + except Exception as e: + self._handle_metric_error("record_embed_now_task_shed", e) + + def set_embed_now_tasks_in_flight(self, count: int) -> None: + try: + embed_now_tasks_in_flight_gauge.labels().set(count) + except Exception as e: + self._handle_metric_error("set_embed_now_tasks_in_flight", e) + def record_dialectic_call( self, *, @@ -279,6 +335,54 @@ class PrometheusMetrics: prometheus_metrics = PrometheusMetrics() +class DBPoolCollector: + """Scrape-time collector for SQLAlchemy connection-pool stats. + + Computed live on each /metrics scrape from the async engine's pool, so it + is always current with no background task or sampling lag. One instance is + registered per process (the API server or a deriver worker). + """ + + def __init__(self, instance_type: str) -> None: + # instance_type: "api" | "deriver" + self.instance_type: str = instance_type + + def collect(self) -> Iterator[GaugeMetricFamily]: + namespace = settings.METRICS.NAMESPACE or "" + gauge = GaugeMetricFamily( + "db_pool_connections", + "DB connections held by this instance, by pool state", + labels=["namespace", "instance_type", "state"], + ) + # Fail soft: Prometheus aborts the entire scrape (dropping ALL metrics) + # if any collector raises, so never let a pool/import hiccup here sink + # the whole /metrics response. + try: + # Lazy import to avoid an import cycle at module load (db imports + # config, telemetry is imported widely). Reads engine.pool directly. + from src.db import get_pool_stats + + stats = get_pool_stats() + except Exception: + logger.warning("Failed to collect DB pool stats", exc_info=True) + stats = {} + for state, value in stats.items(): + gauge.add_metric([namespace, self.instance_type, state], value) + yield gauge + + +_db_pool_collector_registered = False + + +def register_db_pool_collector(instance_type: str) -> None: + """Register the DB pool collector once per process (no-op if metrics off).""" + global _db_pool_collector_registered + if _db_pool_collector_registered or not settings.METRICS.ENABLED: + return + REGISTRY.register(DBPoolCollector(instance_type)) + _db_pool_collector_registered = True + + async def metrics_endpoint(_request: Request) -> Response: if not settings.METRICS.ENABLED: return Response("Metrics are disabled", status_code=404) diff --git a/src/telemetry/sentry.py b/src/telemetry/sentry.py index 1d16b1e7..7b17e156 100644 --- a/src/telemetry/sentry.py +++ b/src/telemetry/sentry.py @@ -6,22 +6,111 @@ import inspect import logging from collections.abc import Callable, Sequence from functools import wraps -from typing import TYPE_CHECKING, ParamSpec, TypeVar, cast +from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, cast import sentry_sdk +from fastapi.exceptions import RequestValidationError +from pydantic import ValidationError +from sqlalchemy.exc import OperationalError from src.config import settings +from src.exceptions import HonchoException P = ParamSpec("P") T = TypeVar("T") if TYPE_CHECKING: - from sentry_sdk._types import EventProcessor + from sentry_sdk._types import Event, EventProcessor, Hint from sentry_sdk.integrations import Integration logger = logging.getLogger(__name__) +def default_before_send(event: Event, hint: Hint | None) -> Event | None: + """Filter/regroup known non-actionable events before Sentry ingests them. + + Shared by every entrypoint (API + deriver) so filtering is process-agnostic. + """ + if not hint: + return event + + exc_info = hint.get("exc_info") + if not exc_info: + return event + + _, exc_value, _ = exc_info + if isinstance(exc_value, HonchoException): + return None + + # Filters out ValidationErrors and RequestValidationErrors (typically from Pydantic) + if isinstance(exc_value, ValidationError | RequestValidationError): + logger.info(f"Filtering out validation error from Sentry: {exc_value}") + return None + + # DB connection-pool checkout timeouts are a fleet-wide saturation symptom, not a + # per-transaction bug. Collapse every occurrence into one issue (Sentry would otherwise + # split by transaction/endpoint) and drop to warning so it stops tripping error alerts. + # Watch it via a rate/spike metric alert instead. Root cause tracked in DEV-1852. + if isinstance(exc_value, OperationalError) and "connection timeout expired" in str( + exc_value + ): + event["fingerprint"] = ["honcho-db-connection-timeout"] + event["level"] = "warning" + return event + + return event + + +# Paths whose transactions carry no debugging value but are hit constantly +# (health checks, Prometheus scrapes, OpenAPI schema, docs). Tracing them at the +# same rate as real traffic drowns the signal and burns tracing/profiling quota. +# Note: /docs and /redoc are disabled in production but are listed for safety. +_UNSAMPLED_PATHS = frozenset( + {"/metrics", "/health", "/openapi.json", "/docs", "/redoc"} +) + + +def _is_unsampled_transaction_name(name: str | None) -> bool: + """Match infra/scrape transactions by name. + + Fallback for transactions that don't expose an ASGI scope path (e.g. the + deriver's metrics server) or whose endpoint-style name encodes the route. + """ + if not name: + return False + return ( + name.endswith("openapi") + or name.endswith("metrics_endpoint") + or "prometheus.metrics" in name + ) + + +def traces_sampler(sampling_context: dict[str, Any]) -> float: + """Drop infra/scrape transactions; sample everything else at the default rate. + + Using a sampler (rather than ``before_send_transaction``) means dropped + transactions are never recorded or profiled, and the decision propagates to + child spans. ``SENTRY.TRACES_SAMPLE_RATE`` remains the rate for real traffic. + """ + asgi_scope = cast("dict[str, Any] | None", sampling_context.get("asgi_scope")) + if asgi_scope is not None and asgi_scope.get("path") in _UNSAMPLED_PATHS: + return 0.0 + + transaction_context = cast( + "dict[str, Any] | None", sampling_context.get("transaction_context") + ) + name = transaction_context.get("name") if transaction_context else None + if _is_unsampled_transaction_name(name if isinstance(name, str) else None): + return 0.0 + + # Respect an upstream sampling decision when continuing a distributed trace. + parent_sampled = sampling_context.get("parent_sampled") + if parent_sampled is not None: + return float(parent_sampled) + + return settings.SENTRY.TRACES_SAMPLE_RATE + + # Sentry SDK's default behavior: # - Captures INFO+ level logs as breadcrumbs # - Captures ERROR+ level logs as Sentry events @@ -31,25 +120,33 @@ logger = logging.getLogger(__name__) def initialize_sentry( *, integrations: Sequence[Integration], - before_send: EventProcessor | None = None, + before_send: EventProcessor | None = default_before_send, ) -> None: """Initialize Sentry SDK with project settings. Args: integrations: Sentry SDK integrations to enable (e.g., Starlette, FastAPI). - before_send: Optional event filter callback to suppress specific exceptions. + before_send: Event filter override. Defaults to ``default_before_send`` so + every entrypoint gets the shared filters; pass ``None`` to opt out. """ sentry_sdk.init( dsn=settings.SENTRY.DSN, enable_tracing=True, release=settings.SENTRY.RELEASE, environment=settings.SENTRY.ENVIRONMENT, - traces_sample_rate=settings.SENTRY.TRACES_SAMPLE_RATE, + # traces_sampler supersedes traces_sample_rate; it returns the configured + # rate for real traffic and 0.0 for infra/scrape endpoints (see above). + traces_sampler=traces_sampler, profiles_sample_rate=settings.SENTRY.PROFILES_SAMPLE_RATE, before_send=before_send, integrations=integrations, ) + # Tag every event with the configured namespace so errors can be filtered by + # instance. Set on the global scope so it applies regardless of the current + # isolation/task scope. + sentry_sdk.get_global_scope().set_tag("namespace", settings.NAMESPACE) + def with_sentry_transaction( name: str, op: str diff --git a/src/telemetry/trace_exporter.py b/src/telemetry/trace_exporter.py new file mode 100644 index 00000000..0023f97a --- /dev/null +++ b/src/telemetry/trace_exporter.py @@ -0,0 +1,169 @@ +"""CloudEvents exporter: turns a CapturedLLMCall into trace events. + +Registered into the `src/llm/capture.py` exporter registry at startup when +`TELEMETRY.TRACE_PAYLOADS_ENABLED` is on. For each captured call it emits: +- one `trace.content` per unique message/output/thinking/tool-schema (deduped + per run so each ships once), and +- one `llm.call.traced` carrying the span-tree correlation + content refs + + a self-contained accounting copy. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from src.config import settings +from src.llm.capture import ( + ROLE_OUTPUT, + ROLE_THINKING, + ROLE_TOOL_SCHEMA, + CapturedLLMCall, + clip_for_trace, + compute_content_hash, +) +from src.telemetry import trace_session +from src.telemetry.events import emit_trace +from src.telemetry.events.trace import LLMCallTracedEvent, TraceContentEvent + +logger = logging.getLogger(__name__) + + +class TraceExporter: + """`LLMCallExporter` that ships replay-grade content to the trace stream.""" + + def export(self, call: CapturedLLMCall) -> None: + # Double-gate (the exporter is only registered when on, but a config + # flip or a stray registration shouldn't leak payloads). + if not settings.TELEMETRY.TRACE_PAYLOADS_ENABLED: + return + purposes = settings.TELEMETRY.TRACE_PURPOSES + if purposes and call.call_purpose not in purposes: + return + + run_key = call.trace_id or call.span_id or call.run_id or "" + was_truncated = call.input_truncated + + # --- Context window: reuse precomputed input-message hashes --- + input_message_refs: list[str] = [] + for message in call.input_messages: + input_message_refs.append(message.content_hash) + self._emit_content( + run_key, + content_hash=message.content_hash, + role=message.role, + content=message.content, + tool_call_id=message.tool_call_id, + honcho_authored=message.role == "system", + tool_calls=message.tool_calls, + ) + + # --- Tool schemas (Honcho-authored, content-addressed) --- + tool_schema_refs: list[str] = [] + for schema in call.tool_schemas: + ref, truncated = self._emit_hashed_content( + run_key, ROLE_TOOL_SCHEMA, schema, honcho_authored=True + ) + was_truncated = was_truncated or truncated + tool_schema_refs.append(ref) + + # --- Output content / thinking --- + output_content_ref: str | None = None + if call.output_content not in (None, ""): + output_content_ref, truncated = self._emit_hashed_content( + run_key, ROLE_OUTPUT, call.output_content + ) + was_truncated = was_truncated or truncated + + output_thinking_ref: str | None = None + if call.thinking_content: + output_thinking_ref, truncated = self._emit_hashed_content( + run_key, ROLE_THINKING, call.thinking_content + ) + was_truncated = was_truncated or truncated + + signatures = [ + block["signature"] + for block in call.thinking_blocks + if block.get("signature") + ] + + emit_trace( + LLMCallTracedEvent( + trace_id=call.trace_id, + span_id=call.span_id, + parent_span_id=call.parent_span_id, + iteration=call.iteration, + step_seq=call.step_seq, + attempt=call.attempt, + was_fallback=call.was_fallback, + call_purpose=call.call_purpose, + parent_category=call.parent_category, + session_id=call.session_id, + transport=call.transport, # pyright: ignore[reportArgumentType] + provider_label=call.provider_label, + model=call.model, + input_message_refs=input_message_refs, + tool_schema_refs=tool_schema_refs, + tool_choice=call.tool_choice, + output_content_ref=output_content_ref, + output_tool_calls=call.output_tool_calls, + output_thinking_ref=output_thinking_ref, + output_signatures=signatures, + finish_reason=call.finish_reason, + provider_input_tokens=call.input_tokens, + provider_output_tokens=call.output_tokens, + cache_read_tokens=call.cache_read_tokens, + cache_creation_tokens=call.cache_creation_tokens, + was_truncated=was_truncated, + ) + ) + + def _emit_hashed_content( + self, + run_key: str, + role: str, + raw_content: Any, + *, + honcho_authored: bool = False, + ) -> tuple[str, bool]: + """Clip + hash a non-message content value and emit it. Returns (hash, truncated).""" + content, truncated = clip_for_trace(raw_content) + content_hash = compute_content_hash(role, content, None) + self._emit_content( + run_key, + content_hash=content_hash, + role=role, + content=content, + tool_call_id=None, + honcho_authored=honcho_authored, + ) + return content_hash, truncated + + def _emit_content( + self, + run_key: str, + *, + content_hash: str, + role: str, + content: Any, + tool_call_id: str | None, + honcho_authored: bool, + tool_calls: list[dict[str, Any]] | None = None, + ) -> None: + """Emit one trace.content, deduped per run (skip if already shipped).""" + if not trace_session.mark_emitted(run_key, content_hash): + return + emit_trace( + TraceContentEvent( + content_hash=content_hash, + role=role, + content=content, + tool_call_id=tool_call_id, + honcho_authored=honcho_authored, + tool_calls=tool_calls or [], + ) + ) + + +__all__ = ["TraceExporter"] diff --git a/src/telemetry/trace_session.py b/src/telemetry/trace_session.py new file mode 100644 index 00000000..a08ca9d7 --- /dev/null +++ b/src/telemetry/trace_session.py @@ -0,0 +1,73 @@ +"""Per-run content dedup for the trace stream — makes bandwidth O(N). + +Tracks the set of content hashes a run has already shipped, so each unique +message ships its `trace.content` exactly once per run. + +The set is bounded by an LRU over runs: once more than `_MAX_RUNS` runs are +tracked, the least-recently-touched one is evicted (almost always a run that has +already finished), so dedup keeps working for active runs no matter how many the +process has handled. A single run exceeding `_MAX_HASHES_PER_RUN` unique messages +stops deduping and emits-anyway, bumping a metric — that loss is measured. +""" + +from __future__ import annotations + +import logging +import threading +from collections import OrderedDict + +logger = logging.getLogger(__name__) + +# LRU window over runs + per-run hash cap. Generous — a run rarely has more than +# a few hundred unique messages, and far fewer than _MAX_RUNS are ever live at +# once; both are runaway backstops, not tuning knobs. +_MAX_RUNS = 4096 +_MAX_HASHES_PER_RUN = 8192 + +# trace_id (fallback span_id) → set of content hashes already shipped this run. +# OrderedDict so we can evict the least-recently-used run when over _MAX_RUNS. +_runs: OrderedDict[str, set[str]] = OrderedDict() +_lock = threading.Lock() + + +def mark_emitted(run_key: str, content_hash: str) -> bool: + """Return True if this hash should be shipped for ``run_key`` (first time), + False if already shipped this run (skip the ``trace.content``). + + A run exceeding `_MAX_HASHES_PER_RUN` returns True (emit-anyway) and records a + drop of the dedup *guarantee* — the event still ships, we just stopped + tracking. Tracking a new run past `_MAX_RUNS` evicts the LRU run instead (a + routine, lossless bound — the evicted run is almost always already finished). + """ + with _lock: + seen = _runs.get(run_key) + if seen is None: + if len(_runs) >= _MAX_RUNS: + _runs.popitem(last=False) # evict the least-recently-used run + seen = _runs[run_key] = set() + else: + _runs.move_to_end(run_key) # mark most-recently-used + if content_hash in seen: + return False + if len(seen) >= _MAX_HASHES_PER_RUN: + _record_overflow("max_hashes") + return True + seen.add(content_hash) + return True + + +def reset() -> None: + """Drop all tracked runs — used on shutdown and in tests.""" + with _lock: + _runs.clear() + + +def _record_overflow(reason: str) -> None: + try: + from src.telemetry import prometheus_metrics + + prometheus_metrics.record_telemetry_event_dropped( + reason=f"trace_dedup_{reason}" + ) + except Exception: # pragma: no cover - best-effort telemetry + logger.debug("trace dedup overflow (%s)", reason) diff --git a/src/utils/agent_tools.py b/src/utils/agent_tools.py index de5b2b09..9f214009 100644 --- a/src/utils/agent_tools.py +++ b/src/utils/agent_tools.py @@ -1,7 +1,7 @@ import asyncio import logging import weakref -from collections.abc import Callable +from collections.abc import Callable, Sequence from dataclasses import dataclass from datetime import datetime from typing import Any, cast @@ -24,8 +24,17 @@ from src.telemetry.events import ( emit, ) from src.utils import summarizer -from src.utils.formatting import format_new_turn_with_timestamp, utc_now_iso -from src.utils.representation import Representation +from src.utils.formatting import ( + format_datetime_utc, + format_new_turn_with_timestamp, + parse_datetime_iso, + utc_now_iso, +) +from src.utils.representation import ( + ALLOWLIST_SAFE_LEVELS, + Representation, + allowlist_safe_levels, +) from src.utils.types import ToolResult, embedding_call_purpose, get_current_iteration logger = logging.getLogger(__name__) @@ -845,6 +854,18 @@ INDUCTION_SPECIALIST_TOOLS: list[dict[str, Any]] = [ TOOLS["create_observations_inductive"], ] +# Tools for the card-refresh specialist (card_refresh dream type). +# Card-only maintenance: discovery plus update_peer_card. Deliberately +# excludes every observation-mutating tool (create_observations*, +# delete_observations) — a card refresh must never touch observations. +CARD_REFRESH_SPECIALIST_TOOLS: list[dict[str, Any]] = [ + # Discovery tools + TOOLS["get_recent_observations"], + TOOLS["search_memory"], + # Action tool + TOOLS["update_peer_card"], +] + async def create_observations( observations: list[schemas.ObservationInput], @@ -982,14 +1003,16 @@ async def create_observations( accepted: list[schemas.DocumentCreate] = [] if documents: async with tracked_db("create_observations.save") as db: - accepted = await crud.create_documents( - db, - documents=documents, - workspace_name=workspace_name, - observer=observer, - observed=observed, - deduplicate=True, - ) + accepted = ( + await crud.create_documents( + db, + documents=documents, + workspace_name=workspace_name, + observer=observer, + observed=observed, + deduplicate=True, + ) + ).created_documents logger.info( "Created %d observations in %s/%s/%s", len(accepted), @@ -1011,6 +1034,7 @@ async def get_recent_history( session_name: str | None, observed: str | None = None, token_limit: int = 8192, + session_allowlist: list[str] | None = None, ) -> list[models.Message]: """ Retrieve recent conversation history. @@ -1023,6 +1047,10 @@ async def get_recent_history( db: Database session workspace_name: Workspace identifier session_name: Session identifier (optional) + Deprecated for *scoping*: prefer session_allowlist, which + intersects with observer membership. This parameter also pins + the query to one session and bypasses observer scoping, so it + is not a drop-in equivalent and is not removed. observed: Peer name to filter by when no session specified (optional) token_limit: Maximum tokens to retrieve (default: 8192) @@ -1030,6 +1058,9 @@ async def get_recent_history( List of messages in chronological order """ if session_name: + # Fail closed: a specific session outside the allowlist is not readable. + if session_allowlist is not None and session_name not in session_allowlist: + return [] # Get messages from a specific session messages_stmt = await crud.get_messages( workspace_name=workspace_name, @@ -1042,7 +1073,11 @@ async def get_recent_history( # Return in chronological order return list(reversed(messages)) elif observed: + # Fail closed on an empty allowlist + if session_allowlist is not None and not session_allowlist: + return [] # Get recent messages from the observed peer across all sessions + # (restricted to the session allowlist when one is provided) stmt = ( select(models.Message) .where(models.Message.workspace_name == workspace_name) @@ -1050,6 +1085,8 @@ async def get_recent_history( .order_by(models.Message.created_at.desc()) .limit(50) # Limit to recent messages ) + if session_allowlist is not None: + stmt = stmt.where(models.Message.session_name.in_(session_allowlist)) result = await db.execute(stmt) messages = list(result.scalars().all()) # Return in chronological order @@ -1067,6 +1104,7 @@ async def search_memory( limit: int, levels: list[str] | None = None, embedding: list[float] | None = None, + session_allowlist: list[str] | None = None, ) -> Representation: """ Search for observations in memory using semantic similarity. @@ -1087,10 +1125,22 @@ async def search_memory( Returns: Representation object containing relevant observations """ - # Build filter for levels if specified - filters: dict[str, Any] | None = None + # Fail closed on an empty allowlist — downstream stores drop empty IN + # clauses, which would silently widen scope. + if session_allowlist is not None and not session_allowlist: + return Representation() + + if session_allowlist is not None: + levels = allowlist_safe_levels(levels) + if not levels: + return Representation() + + # Build filters for levels / session allowlist if specified + filters: dict[str, Any] = {} if levels: - filters = {"level": {"in": levels}} + filters["level"] = {"in": levels} + if session_allowlist is not None: + filters["session_name"] = {"in": session_allowlist} documents = await crud.query_documents( db=None, @@ -1099,7 +1149,7 @@ async def search_memory( observed=observed, query=query, top_k=limit, - filters=filters, + filters=filters or None, embedding=embedding, ) @@ -1112,6 +1162,7 @@ async def get_observation_context( session_name: str | None, message_ids: list[str], observer: str | None = None, + session_allowlist: list[str] | None = None, ) -> list[models.Message]: """ Retrieve messages for given message IDs along with surrounding context. @@ -1124,9 +1175,16 @@ async def get_observation_context( db: Database session workspace_name: Workspace identifier session_name: Session identifier (optional) + Deprecated for *scoping*: prefer session_allowlist, which + intersects with observer membership. This parameter also pins + the query to one session and bypasses observer scoping, so it + is not a drop-in equivalent and is not removed. message_ids: List of message IDs to retrieve observer: When provided and session_name is None, scope results to sessions this peer belongs to + session_allowlist: Optional session allowlist. None is unrestricted; an + empty list fails closed (empty result); a populated list is + intersected with the observer's session scope when observer is set Returns: List of messages in chronological order, including the requested messages and surrounding context @@ -1134,16 +1192,13 @@ async def get_observation_context( if not message_ids: return [] - # Pre-fetch peer session scope if needed - allowed_session_names: list[str] | None = None - if observer and not session_name: - from src.crud.message import get_peer_session_names + from src.crud.message import resolve_session_scope - allowed_session_names = await get_peer_session_names( - db, workspace_name, observer - ) - if not allowed_session_names: - return [] + allowed_session_names, deny = await resolve_session_scope( + db, workspace_name, session_name, session_allowlist, observer + ) + if deny: + return [] # Use a CTE to get seq_in_session values for target messages stmt = ( @@ -1202,9 +1257,16 @@ async def extract_preferences( Args: workspace_name: Workspace identifier session_name: Session identifier (optional) + Deprecated for *scoping*: prefer session_allowlist, which + intersects with observer membership. This parameter also pins + the query to one session and bypasses observer scoping, so it + is not a drop-in equivalent and is not removed. observed: The peer whose preferences to extract observer: When provided and session_name is None, scope results to sessions this peer belongs to + session_allowlist: Optional session allowlist. None is unrestricted; an + empty list fails closed (empty result); a populated list is + intersected with the observer's session scope when observer is set Returns: Dict with 'messages' list containing potentially relevant messages @@ -1284,12 +1346,78 @@ class ToolContext: db_lock: asyncio.Lock # Optional resolved configuration for checking feature flags configuration: ResolvedConfiguration | None = None + # Optional session allowlist (dialectic filters). When set, message and + # conclusion recall is restricted to these sessions (intersected with + # observer membership); empty list fails closed. + session_allowlist: list[str] | None = None # Telemetry context fields run_id: str | None = None agent_type: str | None = None # "dialectic", "deriver", "dreamer" parent_category: str | None = None # Parent category for CloudEvents +def _normalize_observation_id(obs_id: str) -> str: + """Strip the display-format ``id:`` prefix from a model-supplied observation ID. + + Observations are presented to agents as ``[id:xxx]`` (see + ``Representation.str_with_ids``), and despite tool-schema instructions to + pass the bare ID, models sometimes copy the prefix verbatim. Since document + IDs are nanoids whose alphabet includes ``-`` and ``_``, only the ``id:`` + prefix and surrounding whitespace are stripped — anything more aggressive + could mangle legitimate IDs. + """ + obs_id = obs_id.strip() + if obs_id.lower().startswith("id:"): + obs_id = obs_id[3:] + return obs_id.strip() + + +async def _latest_source_timestamp( + ctx: ToolContext, + observations: list[schemas.ObservationInput], +) -> str | None: + """Latest ``message_created_at`` across all source observations in the batch. + + Dreamer conclusions (deductive/inductive) are derived from existing + observations referenced by ``source_ids`` rather than from live messages. + Their logical timestamp is the point when the conclusion became possible + from its evidence, not when the dreamer happened to run, so we date + ``internal_metadata["message_created_at"]`` to the most recent source + observation. The physical ``Document.created_at`` column remains the insert + time. Returns None if no source_ids resolve to a usable timestamp (caller + falls back to now). + """ + source_ids: list[str] = [] + for obs in observations: + if obs.source_ids: + source_ids.extend(obs.source_ids) + if not source_ids: + return None + + latest: datetime | None = None + async with tracked_db("create_observations.source_ts", read_only=True) as db: + docs = await crud.fetch_documents_by_ids( + db, + workspace_name=ctx.workspace_name, + observer=ctx.observer, + observed=ctx.observed, + document_ids=list(set(source_ids)), + ) + for doc in docs: + raw = doc.internal_metadata.get("message_created_at") + if not isinstance(raw, str): + continue + try: + # always tz-aware, so the comparison below can't crash on mixed formats + parsed = parse_datetime_iso(raw) + except ValueError: + continue + if latest is None or parsed > latest: + latest = parsed + + return format_datetime_utc(latest) if latest is not None else None + + async def _handle_create_observations_impl( ctx: ToolContext, tool_input: dict[str, Any], @@ -1309,7 +1437,15 @@ async def _handle_create_observations_impl( obs["level"] = forced_level else: obs.setdefault("level", default_level) - + # Models sometimes copy the display-format "id:" prefix into source_ids; + # normalize so provenance links reference real document IDs. + source_ids = obs.get("source_ids") + if isinstance(source_ids, list): + normalized_source_ids: list[str] = [] + for source_id in cast(list[Any], source_ids): + if isinstance(source_id, str): + normalized_source_ids.append(_normalize_observation_id(source_id)) + obs["source_ids"] = normalized_source_ids # Validate observations individually so valid ones are still processed observations: list[schemas.ObservationInput] = [] validation_failures: list[ObservationFailure] = [] @@ -1333,6 +1469,21 @@ async def _handle_create_observations_impl( ) ) continue + # Session-purity invariant: explicit observations record what was + # directly derived from a session's messages. Agents that are not + # processing messages (dreamer specialists, dialectic) must not mint + # them — consolidation output belongs at a derived level. + if not ctx.current_messages and validated.level == "explicit": + validation_failures.append( + ObservationFailure( + content_preview=validated.content[:50], + error=( + "Only message ingestion can create 'explicit' observations; " + "use a derived level (deductive/inductive/contradiction)" + ), + ) + ) + continue observations.append(validated) if not observations: @@ -1344,10 +1495,16 @@ async def _handle_create_observations_impl( # Determine message context if ctx.current_messages: message_ids = [msg.id for msg in ctx.current_messages] - message_created_at = str(ctx.current_messages[-1].created_at) + # same ISO-8601 Z format as the dreamer path below + message_created_at = format_datetime_utc(ctx.current_messages[-1].created_at) else: + # Dreamer path: no current messages. Backdate the conclusion to the + # latest source observation, which is when the inference became possible. + message_ids = [] - message_created_at = utc_now_iso() + message_created_at = ( + await _latest_source_timestamp(ctx, observations) + ) or utc_now_iso() # Use lock to serialize database writes (prevents concurrent commit issues) async with ctx.db_lock: @@ -1607,13 +1764,14 @@ async def _handle_get_recent_history( ) -> "str | ToolResult": """Handle get_recent_history tool.""" _ = tool_input - async with tracked_db("tool.get_recent_history") as db: + async with tracked_db("tool.get_recent_history", read_only=True) as db: history: list[models.Message] = await get_recent_history( db, workspace_name=ctx.workspace_name, session_name=ctx.session_name, observed=ctx.observed, token_limit=ctx.history_token_limit, + session_allowlist=ctx.session_allowlist, ) if not history: return "No conversation history available" @@ -1659,15 +1817,28 @@ async def _handle_search_memory( "query_tokens": _estimate_tokens_safe(query), } - documents = await crud.query_documents( - db=None, - workspace_name=ctx.workspace_name, - observer=ctx.observer, - observed=ctx.observed, - query=query, - top_k=top_k, - embedding=query_embedding, - ) + # Restrict conclusion recall to the session allowlist when one is set. + # Empty allowlist fails closed (downstream stores drop empty IN clauses), + # and only levels with a trustworthy session stamp are served. + documents: Sequence[models.Document] + if ctx.session_allowlist is not None and not ctx.session_allowlist: + documents = [] + else: + documents = await crud.query_documents( + db=None, + workspace_name=ctx.workspace_name, + observer=ctx.observer, + observed=ctx.observed, + query=query, + top_k=top_k, + embedding=query_embedding, + filters={ + "session_name": {"in": ctx.session_allowlist}, + "level": {"in": list(ALLOWLIST_SAFE_LEVELS)}, + } + if ctx.session_allowlist is not None + else None, + ) mem = Representation.from_documents(documents) total_count = mem.len() if total_count == 0: @@ -1689,6 +1860,7 @@ async def _handle_search_memory( context_window=0, embedding=query_embedding, observer=ctx.observer, + session_allowlist=ctx.session_allowlist, ) if snippets: message_output = _format_message_snippets( @@ -1723,13 +1895,14 @@ async def _handle_get_observation_context( ctx: ToolContext, tool_input: dict[str, Any] ) -> "str | ToolResult": """Handle get_observation_context tool.""" - async with tracked_db("tool.get_observation_context") as db: + async with tracked_db("tool.get_observation_context", read_only=True) as db: messages = await get_observation_context( db, workspace_name=ctx.workspace_name, session_name=ctx.session_name, message_ids=tool_input["message_ids"], observer=ctx.observer, + session_allowlist=ctx.session_allowlist, ) if not messages: return f"No messages found for IDs {tool_input['message_ids']}" @@ -1772,6 +1945,7 @@ async def _handle_search_messages( context_window=2, embedding=query_embedding, observer=ctx.observer, + session_allowlist=ctx.session_allowlist, ) search_meta: dict[str, Any] = { "top_k": limit, @@ -1808,6 +1982,7 @@ async def _handle_grep_messages( limit=limit, context_window=context_window, observer=ctx.observer, + session_allowlist=ctx.session_allowlist, ) if not snippets: return f"No messages found containing '{text}'" @@ -1862,7 +2037,7 @@ async def _handle_get_messages_by_date_range( if isinstance(before_date, str): return before_date # Error message - async with tracked_db("tool.get_messages_by_date_range") as db: + async with tracked_db("tool.get_messages_by_date_range", read_only=True) as db: messages = await crud.get_messages_by_date_range( db, workspace_name=ctx.workspace_name, @@ -1872,6 +2047,7 @@ async def _handle_get_messages_by_date_range( limit=limit, order=order, observer=ctx.observer, + session_allowlist=ctx.session_allowlist, ) msg_count = len(messages) messages_text = ( @@ -1944,6 +2120,7 @@ async def _handle_search_messages_temporal( before_date=before_date, limit=limit, context_window=context_window, + session_allowlist=ctx.session_allowlist, embedding=query_embedding, observer=ctx.observer, ) @@ -1980,7 +2157,7 @@ async def _handle_get_recent_observations( ) -> str: """Handle get_recent_observations tool.""" session_only = tool_input.get("session_only", False) - async with tracked_db("tool.get_recent_observations") as db: + async with tracked_db("tool.get_recent_observations", read_only=True) as db: documents = await crud.query_documents_recent( db=db, workspace_name=ctx.workspace_name, @@ -2006,7 +2183,7 @@ async def _handle_get_most_derived_observations( ctx: ToolContext, tool_input: dict[str, Any] ) -> str: """Handle get_most_derived_observations tool.""" - async with tracked_db("tool.get_most_derived_observations") as db: + async with tracked_db("tool.get_most_derived_observations", read_only=True) as db: documents = await crud.query_documents_most_derived( db=db, workspace_name=ctx.workspace_name, @@ -2038,7 +2215,7 @@ async def _handle_get_session_summary( if summary_type == "long" else summarizer.SummaryType.SHORT ) - async with tracked_db("tool.get_session_summary") as db: + async with tracked_db("tool.get_session_summary", read_only=True) as db: summary = await summarizer.get_summary( db, ctx.workspace_name, ctx.session_name, st ) @@ -2050,7 +2227,7 @@ async def _handle_get_session_summary( async def _handle_get_peer_card(ctx: ToolContext, tool_input: dict[str, Any]) -> str: """Handle get_peer_card tool.""" _ = tool_input - async with tracked_db("tool.get_peer_card") as db: + async with tracked_db("tool.get_peer_card", read_only=True) as db: peer_card = await crud.get_peer_card( db, workspace_name=ctx.workspace_name, @@ -2200,16 +2377,25 @@ async def _handle_get_reasoning_chain( ctx: ToolContext, tool_input: dict[str, Any] ) -> str: """Handle get_reasoning_chain tool.""" + # Reasoning chains traverse provenance across sessions by design, so a + # session allowlist cannot be enforced on the traversal without exposing + # out-of-scope premises/conclusions. Fail closed rather than leak. + if ctx.session_allowlist is not None: + return ( + "Reasoning-chain traversal is unavailable for session-scoped " + "queries. Use search_memory and message tools instead." + ) observation_id = tool_input.get("observation_id") if not observation_id: return "ERROR: 'observation_id' is required" + observation_id = _normalize_observation_id(observation_id) direction = tool_input.get("direction", "both") if direction not in ("premises", "conclusions", "both"): return f"ERROR: Invalid direction '{direction}'. Must be 'premises', 'conclusions', or 'both'" # Get the observation itself - async with tracked_db("tool.get_reasoning_chain") as db: + async with tracked_db("tool.get_reasoning_chain", read_only=True) as db: docs = await crud.get_documents_by_ids(db, ctx.workspace_name, [observation_id]) if not docs or not docs[0]: return f"ERROR: Observation '{observation_id}' not found" @@ -2324,6 +2510,7 @@ async def create_tool_executor( run_id: str | None = None, agent_type: str | None = None, parent_category: str | None = None, + session_allowlist: list[str] | None = None, ) -> Callable[[str, dict[str, Any]], Any]: """ Create a unified tool executor function for all agent operations. @@ -2364,6 +2551,7 @@ async def create_tool_executor( history_token_limit=history_token_limit, db_lock=shared_lock, configuration=configuration, + session_allowlist=session_allowlist, run_id=run_id, agent_type=agent_type, parent_category=parent_category, @@ -2403,6 +2591,10 @@ async def create_tool_executor( metadata: dict[str, Any] = {} is_error: bool = False + # Langfuse tool observation; auto-parents under the active step span. + # Closed in the finally below with output + level. + tool_obs = _begin_tool_observation(tool_name, tool_input) + try: handler = _TOOL_HANDLERS.get(tool_name) if handler: @@ -2481,11 +2673,49 @@ async def create_tool_executor( provider_tool_call_id=get_current_provider_tool_call_id(), ) + _finish_tool_observation(tool_obs, result_str, is_error) + return result_str return execute_tool +def _begin_tool_observation(tool_name: str, tool_input: dict[str, Any]) -> Any: + """Open a non-current Langfuse "tool" observation for one tool execution. + + Auto-parents under the active step span (else standalone). Returns a handle + (closed by `_finish_tool_observation`) or None when disabled/setup fails. + All tools are ``as_type="tool"`` — they share one generic dispatcher. + + Only fires in legacy *inline* mode. In exporter mode there's no live span + context to parent under, so this would emit a rootless tool trace per call; + the LangfuseExporter already projects tool spans (from ``output_tool_calls``) + nested under the step span, so a live observation here just double-emits. + """ + if not settings.langfuse_inline_enabled: + return None + try: + from langfuse import get_client + + return get_client().start_observation( + as_type="tool", name=tool_name, input=tool_input + ) + except Exception: # pragma: no cover - best-effort telemetry + logger.debug("Failed to open Langfuse tool observation", exc_info=True) + return None + + +def _finish_tool_observation(tool_obs: Any, result_str: str, is_error: bool) -> None: + """Close a Langfuse tool observation opened by `_begin_tool_observation`.""" + if tool_obs is None: + return + try: + tool_obs.update(output=result_str, level="ERROR" if is_error else None) + tool_obs.end() + except Exception: # pragma: no cover - best-effort telemetry + logger.debug("Failed to close Langfuse tool observation", exc_info=True) + + def _emit_agent_tool_call_completed( *, ctx: "ToolContext", diff --git a/src/utils/filter.py b/src/utils/filter.py index 1394a8c6..8df3590c 100644 --- a/src/utils/filter.py +++ b/src/utils/filter.py @@ -1,7 +1,8 @@ import datetime -from collections.abc import Callable +from collections.abc import Callable, Sequence from logging import getLogger from typing import Any, TypeVar +from typing import cast as typing_cast from sqlalchemy import ColumnElement, Select, and_, case, cast, literal, not_, or_ from sqlalchemy.types import Numeric @@ -28,6 +29,10 @@ COMPARISON_OPERATORS = { NUMERIC_OPERATORS = {"gte", "lte", "gt", "lt", "ne"} +# JSONB columns keep containment semantics: bare lists are not membership +# sugar, and dict values map to nested-metadata conditions rather than IN/Eq. +JSONB_COLUMNS = ("h_metadata", "configuration", "internal_metadata") + ALLOWED_EXTERNAL_TO_INTERNAL_COLUMN_MAPPING = { "id": "name", "created_at": "created_at", @@ -56,6 +61,83 @@ ALLOWED_EXTERNAL_TO_INTERNAL_COLUMN_MAPPING_DOCUMENTS = { } +MAX_SESSION_ALLOWLIST_ENTRIES = 1000 + + +def extract_session_allowlist( + filters: dict[str, Any] | None, + must_include: str | None = None, +) -> list[str] | None: + """Parse a recall-path ``filters`` body into a session allowlist. + + The dialectic and representation endpoints accept a constrained subset of + the filter DSL: only the ``session_id`` key, valued as a single id, a bare + list of ids, or ``{"in": [...]}``. Unsupported keys or shapes raise + FilterError (422) rather than being silently ignored — a dropped filter + on these endpoints would widen recall scope. + + Args: + filters: The raw ``filters`` body, or None. + must_include: A session id that must appear in the parsed allowlist — + used by routes that also accept a top-level ``session_id``, so the + two can't contradict each other. Ignored when filters is None. + + Returns: + None when filters is None. An explicit empty list is preserved so + downstream consumers fail closed. + + Raises: + FilterError: On an unsupported key or shape, an over-cap list, or a + ``must_include`` session missing from the allowlist. + """ + if filters is None: + return None + + unsupported = set(filters) - {"session_id"} + if unsupported: + raise FilterError( + f"Unsupported filter key(s) for this endpoint: {sorted(unsupported)}. Only 'session_id' is supported." + ) + if "session_id" not in filters: + raise FilterError("filters must contain 'session_id'") + + value = filters["session_id"] + entries: list[Any] + if isinstance(value, str): + entries = [value] + elif isinstance(value, list): + entries = list(typing_cast(Sequence[Any], value)) + elif ( + isinstance(value, dict) + and set(typing_cast("dict[str, Any]", value)) == {"in"} + and isinstance(value["in"], list) + ): + entries = list(typing_cast(Sequence[Any], value["in"])) + else: + raise FilterError( + 'filters.session_id must be a session id, a list of session ids, or {"in": [...]}' + ) + + if len(entries) > MAX_SESSION_ALLOWLIST_ENTRIES: + raise FilterError( + f"filters.session_id supports at most {MAX_SESSION_ALLOWLIST_ENTRIES} sessions per request" + ) + + allowlist: list[str] = [] + seen: set[str] = set() + for entry in entries: + if not isinstance(entry, str) or not entry: + raise FilterError("filters.session_id entries must be non-empty strings") + if entry not in seen: + seen.add(entry) + allowlist.append(entry) + + if must_include is not None and must_include not in seen: + raise FilterError("session_id must be included in filters.session_id") + + return allowlist + + def apply_filter( stmt: Select[tuple[T]], model_class: type[T], filters: dict[str, Any] | None = None ) -> Select[tuple[T]]: @@ -216,6 +298,13 @@ def _build_field_condition( if model_class.__name__ == "Message": column_name = ALLOWED_EXTERNAL_TO_INTERNAL_COLUMN_MAPPING_MESSAGES.get(key) elif model_class.__name__ == "Document": + # NOTE: unlike Message/Workspace, Document falls back to the raw key so + # internal callers can filter on internal column names. The session + # allowlist depends on this: recall passes {"session_name": {"in": ...}} + # (see search_memory in utils/agent_tools.py and RepresentationManager), + # and "session_name" is deliberately absent from the mapping below. + # Tightening this to a strict allowlist would break session scoping — + # fail-closed, since an unmapped key raises, but silently. column_name = ALLOWED_EXTERNAL_TO_INTERNAL_COLUMN_MAPPING_DOCUMENTS.get( key, key, # fallback to the key itself if not found in the mapping for internal use here @@ -238,6 +327,12 @@ def _build_field_condition( if value == "*": return None + # Bare-list sugar on regular columns: {"session_id": ["a", "b"]} is + # shorthand for {"session_id": {"in": ["a", "b"]}}. JSONB columns are + # excluded — a bare list there keeps JSONB containment semantics. + if isinstance(value, list | tuple | set) and column_name not in JSONB_COLUMNS: + value = {"in": list(typing_cast(Sequence[Any], value))} + # Handle comparison operators vs regular values if isinstance(value, dict): # Check if this is a comparison operators dict by looking for known operators @@ -248,12 +343,12 @@ def _build_field_condition( else: # This is a regular value that happens to be a dict # For JSONB fields (metadata, configuration), check if it contains nested comparison operators - if column_name in ("h_metadata", "configuration", "internal_metadata"): + if column_name in JSONB_COLUMNS: return _build_nested_metadata_conditions(column, value) # pyright: ignore else: return column == value else: - if column_name in ("h_metadata", "configuration", "internal_metadata"): + if column_name in JSONB_COLUMNS: return column.contains(value) else: return column == value diff --git a/src/utils/queue_payload.py b/src/utils/queue_payload.py index 605cf615..5d616b1e 100644 --- a/src/utils/queue_payload.py +++ b/src/utils/queue_payload.py @@ -66,6 +66,11 @@ class DreamPayload(BasePayload): delay_reason: str | None = None documents_since_last_dream_at_schedule: int | None = None document_threshold: int | None = None + # card_refresh only: when True the existing peer card is NOT injected into + # the specialist prompt and the card is rebuilt solely from observations + # currently in the collection (used after removals, where the old card may + # contain facts whose support was deleted). + rebuild: bool = False class DeletionPayload(BasePayload): @@ -103,6 +108,7 @@ def create_dream_payload( delay_reason: str | None = None, documents_since_last_dream_at_schedule: int | None = None, document_threshold: int | None = None, + rebuild: bool = False, ) -> dict[str, Any]: """Create a dream payload.""" return DreamPayload( @@ -114,6 +120,7 @@ def create_dream_payload( delay_reason=delay_reason, documents_since_last_dream_at_schedule=documents_since_last_dream_at_schedule, document_threshold=document_threshold, + rebuild=rebuild, ).model_dump(mode="json", exclude_none=True) diff --git a/src/utils/representation.py b/src/utils/representation.py index 674b25bb..01e4b70a 100644 --- a/src/utils/representation.py +++ b/src/utils/representation.py @@ -7,6 +7,35 @@ from pydantic import BaseModel, Field, field_validator from src import models from src.utils.formatting import parse_datetime_iso +# Conclusion levels whose `session_name` stamp is trustworthy enough to scope on. +# +# Explicit conclusions come from the deriver over a single session's message +# batch, so their stamp is authoritative. Deductive/inductive conclusions are +# produced by the dreamer, which reads across *all* sessions (its discovery +# tools default to session_only=False) but stamps its output with one session — +# whichever holds the most recent explicit conclusion, see +# dreamer/dream_scheduler.py. Serving those under a session allowlist would leak +# conclusions synthesized from sessions outside it. +# +# ponytail: whole-level exclusion rather than per-conclusion provenance. The +# reasoning trees already link each conclusion to its premises, so the real fix +# is an authoritative source-session set per conclusion; until that exists this +# fails closed. Tracked in DEV-2201. +ALLOWLIST_SAFE_LEVELS = ("explicit",) + + +def allowlist_safe_levels(levels: list[str] | None) -> list[str]: + """Narrow a level filter to those safe to serve under a session allowlist. + + Returns the intersection with :data:`ALLOWLIST_SAFE_LEVELS`; ``None`` means + "no level filter requested" and yields the full safe set. An empty result + means the caller asked only for levels we can't scope, and should receive + nothing rather than unscoped conclusions. + """ + if levels is None: + return list(ALLOWLIST_SAFE_LEVELS) + return [level for level in levels if level in ALLOWLIST_SAFE_LEVELS] + def _strip_microseconds_and_timezone(timestamp: datetime) -> datetime: """ diff --git a/src/utils/schema_conversion.py b/src/utils/schema_conversion.py new file mode 100644 index 00000000..3e715f4b --- /dev/null +++ b/src/utils/schema_conversion.py @@ -0,0 +1,505 @@ +"""Convert caller-supplied JSON Schema objects into dynamic Pydantic models. + +Used by the dialectic chat endpoint's ``response_format`` option: the caller +sends a JSON Schema dict, and the resulting model is passed as +``response_model`` to ``honcho_llm_call()`` so providers return conforming +JSON. + +Only a conservative subset of JSON Schema is supported (see +``json_response_schema_to_pydantic``). Conversion doubles as validation: any +unsupported construct raises ``ValueError`` with the offending path, which the +router surfaces as a 422. +""" + +import re +from dataclasses import dataclass, field +from typing import Any, Literal, NoReturn, cast + +from pydantic import BaseModel, ConfigDict, Field, create_model + +# "$defs"/"definitions" are extracted at the root before the walk and are +# rejected anywhere else. "$ref" nodes are resolved by _resolve_ref before +# node validation, so they never reach this check. +_UNSUPPORTED_KEYS = ( + "$defs", + "definitions", + "allOf", + "not", + "if", + "then", + "else", + "patternProperties", +) + +_REF_PREFIXES = ("#/$defs/", "#/definitions/") + +_PRIMITIVE_TYPES: dict[str, Any] = { + "string": str, + "number": float, + "integer": int, + "boolean": bool, + "null": type(None), +} + +# Constraint keywords are forwarded to the model's json_schema_extra so the +# LLM sees them, but Pydantic does not enforce them (they are hints only). +_HINT_KEYS = ( + "minItems", + "maxItems", + "minLength", + "maxLength", + "minimum", + "maximum", + "exclusiveMinimum", + "exclusiveMaximum", + "multipleOf", + "pattern", + "format", + "minProperties", + "maxProperties", + "uniqueItems", +) + +_IDENTIFIER_RE = re.compile(r"^[A-Za-z][A-Za-z0-9_]*$") + + +@dataclass +class _Ctx: + """Mutable state shared across one conversion walk.""" + + max_depth: int + max_nodes: int + defs: dict[str, Any] = field(default_factory=dict) + used_names: set[str] = field(default_factory=set) + ref_stack: list[str] = field(default_factory=list) + node_count: int = 0 + + +def json_response_schema_to_pydantic( + schema: dict[str, Any], + *, + model_name: str = "ResponseFormat", + max_depth: int = 20, + max_nodes: int = 500, +) -> type[BaseModel]: + """Convert a JSON Schema dict (root type ``object``) into a Pydantic model. + + Supported constructs: primitive types (``string``/``number``/``integer``/ + ``boolean``/``null``), nested ``object`` with ``properties``, ``array`` + with ``items`` (missing ``items`` yields ``list[Any]``), ``enum`` of + strings/integers/booleans/null, ``anyOf``/``oneOf`` unions (a + ``{"type": "null"}`` member yields an optional), ``type`` given as a list, + ``required``, ``default``, and ``description``. A root-level ``$schema`` + key is ignored. Boolean ``additionalProperties`` is accepted and ignored; + extra keys in LLM output are silently dropped (``extra="ignore"``). + + ``$ref`` is supported for references of the form ``#/$defs/`` or + ``#/definitions/`` into root-level ``$defs``/``definitions`` (this + is what Pydantic's ``model_json_schema()`` and Zod's ``toJSONSchema`` + emit). References are resolved by inlining; sibling keys next to ``$ref`` + overlay the referenced definition (siblings win). Recursive references + are rejected — the error names the cycle. Unreferenced definitions are + ignored without validation. + + Constraint keywords (``minItems``, ``maxLength``, ``minimum``, + ``pattern``, ...) are passed through to the generated schema as hints but + are not enforced by Pydantic. + + Args: + schema: The JSON Schema object. Root must resolve to type ``object``. + model_name: Name for the generated root model class. + max_depth: Maximum nesting depth. Guard against excessive or + malicious schemas (e.g. pathologically deep nesting); exceeding + it raises ``ValueError``. + max_nodes: Maximum total nodes visited across the whole schema. + Guard against excessive or malicious schemas (e.g. enormous + property fan-out); exceeding it raises ``ValueError``. + + Returns: + A dynamically created Pydantic model class. + + Raises: + ValueError: If the schema is malformed or uses an unsupported + construct (``allOf``, ``not``, ``if``/``then``/``else``, + ``patternProperties``, schema-valued ``additionalProperties``, + boolean schemas, unknown types, a non-object root, a ``$ref`` + that is recursive, malformed, or targets an unknown definition, + or ``$defs``/``definitions`` anywhere but the root). The message + names the construct and its path. + """ + schema_obj: Any = schema + if not isinstance(schema_obj, dict): + raise ValueError("response_format must be a JSON Schema object") + + root = {k: v for k, v in schema.items() if k != "$schema"} + defs = _extract_defs(root) + root_type = root.get("type") + # A "$ref" root is allowed through here; the post-conversion check below + # still enforces that it resolves to an object. + is_object_root = root_type == "object" or ( + root_type is None and ("properties" in root or "$ref" in root) + ) + if not is_object_root: + raise ValueError("root schema must have type 'object'") + + ctx = _Ctx(max_depth=max_depth, max_nodes=max_nodes, defs=defs) + annotation = _convert_schema(root, "", model_name, ctx, depth=0) + # An object root always converts to a model class; this is a safety net. + if not (isinstance(annotation, type) and issubclass(annotation, BaseModel)): + raise ValueError("root schema must have type 'object'") + return annotation + + +def _fail(msg: str, path: str) -> NoReturn: + raise ValueError(f"{msg} at {path or 'root'}") + + +def _union(members: tuple[Any, ...]) -> Any: + """Build ``A | B | ...`` from a dynamic tuple of annotations.""" + result: Any = members[0] + for member in members[1:]: + result = result | member + return result + + +def _convert_schema( + raw_node: Any, path: str, name_hint: str, ctx: _Ctx, depth: int +) -> Any: + """Convert one schema node into a type annotation. + + Dispatch order matters: $ref resolution comes first (a ref node is + replaced by its target before anything else looks at it), then enum (a + value constraint) wins over unions, which win over "type"-based + conversion. + """ + if isinstance(raw_node, dict) and "$ref" in raw_node: + return _resolve_ref(cast(dict[str, Any], raw_node), path, ctx, depth) + + node = _validate_node(raw_node, path, ctx, depth) + + if "enum" in node: + return _convert_enum(node["enum"], path) + + if "anyOf" in node or "oneOf" in node: + return _convert_union(node, path, name_hint, ctx, depth) + + node_type = node.get("type") + if isinstance(node_type, list): + return _convert_type_list( + node, cast(list[Any], node_type), path, name_hint, ctx, depth + ) + + # Tolerate an omitted "type" when "properties" makes the intent clear. + if node_type is None and "properties" in node: + node_type = "object" + + if node_type == "object": + return _build_object_model(node, path, name_hint, ctx, depth) + + if node_type == "array": + return _convert_array(node, path, name_hint, ctx, depth) + + if node_type in _PRIMITIVE_TYPES: + return _PRIMITIVE_TYPES[node_type] + + if node_type is None: + _fail("schema has no recognizable type", path) + _fail(f"unsupported type '{node_type}'", path) + + +def _extract_defs(root: dict[str, Any]) -> dict[str, Any]: + """Pop root-level ``$defs``/``definitions`` and merge them into one + registry. Entries are validated lazily, when (and only when) referenced.""" + defs: dict[str, Any] = {} + for key in ("$defs", "definitions"): + raw = root.pop(key, None) + if raw is None: + continue + if not isinstance(raw, dict): + _fail(f"'{key}' must be an object", "") + for name, definition in cast(dict[Any, Any], raw).items(): + if not isinstance(name, str) or not name: + _fail(f"'{key}' definition names must be non-empty strings", "") + if name in defs: + _fail( + f"definition '{name}' appears in both '$defs' and 'definitions'", + "", + ) + defs[name] = definition + return defs + + +def _resolve_ref(node: dict[str, Any], path: str, ctx: _Ctx, depth: int) -> Any: + """Inline a ``$ref`` node: resolve the target definition, overlay any + sibling keys (siblings win), and convert the result in place. + + Only root-relative refs into ``$defs``/``definitions`` are supported. + Cycles are rejected — recursion cannot be inlined. The resolved node is + converted at the same depth (replacement semantics); the node budget in + ``_validate_node`` still counts every expansion, so a definition that is + referenced many times cannot blow up the walk. + """ + ref: Any = node["$ref"] + if not isinstance(ref, str): + _fail("'$ref' must be a string", path) + name: str | None = None + for prefix in _REF_PREFIXES: + if ref.startswith(prefix): + name = ref[len(prefix) :] + break + # "/", "~", and "%" would make the remainder a deeper or escaped JSON + # pointer (e.g. "#/$defs/a/b", "~1" escapes, %-encoding) rather than a + # plain definition name, so their presence means an unsupported form. + if not name or "/" in name or "~" in name or "%" in name: + _fail( + f"unsupported $ref '{ref}': only '#/$defs/' or " + + "'#/definitions/' references are supported", + path, + ) + if name not in ctx.defs: + _fail(f"$ref '{ref}' points to an unknown definition", path) + # ref_stack holds the definitions currently being expanded on this branch + # of the walk, so membership means the definition (transitively) + # references itself. Slicing from the first occurrence yields the cycle + # for the error message, e.g. stack [A, B] + name A -> "A -> B -> A". + if name in ctx.ref_stack: + cycle = " -> ".join([*ctx.ref_stack[ctx.ref_stack.index(name) :], name]) + _fail( + f"recursive $ref is not supported (cycle: {cycle}); " + + "restructure the schema so definitions do not reference themselves", + path, + ) + target: Any = ctx.defs[name] + siblings = {k: v for k, v in node.items() if k != "$ref"} + resolved: Any = target + if siblings and isinstance(target, dict): + resolved = {**cast(dict[str, Any], target), **siblings} + # Pop after converting (not just on success) so the stack tracks only the + # current branch: a diamond — the same definition referenced from two + # sibling nodes — is legitimate reuse, not a cycle. + ctx.ref_stack.append(name) + try: + return _convert_schema(resolved, path, name, ctx, depth) + finally: + ctx.ref_stack.pop() + + +def _validate_node(raw_node: Any, path: str, ctx: _Ctx, depth: int) -> dict[str, Any]: + """Enforce size budgets and node shape; reject unsupported constructs.""" + # node_count is cumulative across the whole walk; depth tracks only the + # current branch. + ctx.node_count += 1 + if ctx.node_count > ctx.max_nodes: + raise ValueError(f"schema exceeds the maximum of {ctx.max_nodes} nodes") + if depth > ctx.max_depth: + raise ValueError(f"schema nesting exceeds the maximum depth of {ctx.max_depth}") + if isinstance(raw_node, bool): + # A special case of the object requirement, with its own message: + # boolean schemas are legal JSON Schema, just deliberately unsupported. + _fail("boolean schemas are not supported", path) + if not isinstance(raw_node, dict): + _fail("schema must be an object", path) + node = cast(dict[str, Any], raw_node) + for key in _UNSUPPORTED_KEYS: + if key in node: + _fail(f"unsupported construct '{key}'", path) + if isinstance(node.get("additionalProperties"), dict): + _fail("additionalProperties with a schema is not supported", path) + return node + + +def _convert_union( + node: dict[str, Any], path: str, name_hint: str, ctx: _Ctx, depth: int +) -> Any: + """Convert anyOf/oneOf, treated identically: a plain union of the member + schemas (a {"type": "null"} member makes the union optional).""" + union_key = "anyOf" if "anyOf" in node else "oneOf" + members = node[union_key] + if not isinstance(members, list) or not members: + _fail(f"'{union_key}' must be a non-empty array", path) + converted = tuple( + _convert_schema( + member, + _child_path(path, f"{union_key}[{i}]"), + f"{name_hint}Option{i}", + ctx, + depth + 1, + ) + for i, member in enumerate(cast(list[Any], members)) + ) + return _union(converted) + + +def _convert_type_list( + node: dict[str, Any], + types: list[Any], + path: str, + name_hint: str, + ctx: _Ctx, + depth: int, +) -> Any: + """Convert the type: ["string", "null"] sugar — re-convert the node once + per entry (keeping sibling keys, at the same depth since it's the same + source node) and union the results.""" + if not types: + _fail("'type' array must not be empty", path) + variants = tuple( + _convert_schema( + {**{k: v for k, v in node.items() if k != "type"}, "type": t}, + path, + name_hint, + ctx, + depth, + ) + for t in types + ) + return _union(variants) + + +def _convert_array( + node: dict[str, Any], path: str, name_hint: str, ctx: _Ctx, depth: int +) -> Any: + """Convert an array schema into a list annotation.""" + items = node.get("items") + # A missing "items" constraint means any element type is allowed. + if items is None: + return list[Any] + item_annotation = _convert_schema( + items, _child_path(path, "items"), f"{name_hint}Item", ctx, depth + 1 + ) + return list[item_annotation] + + +def _convert_enum(raw_values: Any, path: str) -> Any: + if not isinstance(raw_values, list) or not raw_values: + _fail("'enum' must be a non-empty array", path) + literal_values: list[Any] = [] + # None is not Literal-legal, so collect it separately and union NoneType + # back in at the end (an all-null enum degenerates to NoneType). + has_null = False + for value in cast(list[Any], raw_values): + if value is None: + has_null = True + elif isinstance(value, str | int | bool): + literal_values.append(value) + else: + _fail("enum values must be strings, integers, booleans, or null", path) + if not literal_values: + return type(None) + annotation: Any = Literal[tuple(literal_values)] + return _union((annotation, type(None))) if has_null else annotation + + +def _build_object_model( + node: dict[str, Any], path: str, name_hint: str, ctx: _Ctx, depth: int +) -> type[BaseModel]: + raw_properties: Any = node.get("properties", {}) + if not isinstance(raw_properties, dict): + _fail("'properties' must be an object", path) + properties = cast(dict[Any, Any], raw_properties) + raw_required: Any = node.get("required", []) + if not isinstance(raw_required, list) or not all( + isinstance(entry, str) for entry in cast(list[Any], raw_required) + ): + _fail("'required' must be an array of strings", path) + # Entries naming properties that don't exist are tolerated; they just + # have no effect. + required = set(cast(list[str], raw_required)) + + fields: dict[str, tuple[Any, Any]] = {} + for prop_key, prop_schema in properties.items(): + if not isinstance(prop_key, str) or not prop_key: + _fail("property names must be non-empty strings", path) + annotation = _convert_schema( + prop_schema, + _child_path(path, f"properties.{prop_key}"), + f"{name_hint} {prop_key}", + ctx, + depth + 1, + ) + # Non-identifier keys ("my-key") get a sanitized field name, with the + # original key preserved as the alias for validation/serialization. + field_name = _field_name(prop_key, fields) + alias = prop_key if field_name != prop_key else None + fields[field_name] = _make_field( + cast(dict[str, Any], prop_schema) if isinstance(prop_schema, dict) else {}, + annotation, + is_required=prop_key in required, + alias=alias, + ) + + # create_model's overloads can't type dynamic **fields; the values are + # (annotation, FieldInfo) tuples, which is the documented calling form. + model: Any = create_model( # pyright: ignore[reportCallIssue, reportUnknownVariableType] + _unique_model_name(name_hint, ctx), + __config__=ConfigDict(extra="ignore", populate_by_name=True), + **fields, # pyright: ignore[reportArgumentType] + ) + return cast(type[BaseModel], model) + + +def _make_field( + prop_schema: dict[str, Any], + annotation: Any, + *, + is_required: bool, + alias: str | None, +) -> tuple[Any, Any]: + kwargs: dict[str, Any] = {} + if alias is not None: + kwargs["alias"] = alias + description = prop_schema.get("description") + if isinstance(description, str): + kwargs["description"] = description + hints = {key: prop_schema[key] for key in _HINT_KEYS if key in prop_schema} + if hints: + kwargs["json_schema_extra"] = hints + + # Precedence: an explicit default wins (even for required fields), then + # required, then optional — which widens to `T | None` defaulting to None. + if "default" in prop_schema: + return annotation, Field(default=prop_schema["default"], **kwargs) + if is_required: + return annotation, Field(**kwargs) + return annotation | None, Field(default=None, **kwargs) + + +def _child_path(path: str, segment: str) -> str: + return f"{path}.{segment}" if path else segment + + +def _field_name(prop_key: str, existing: dict[str, Any]) -> str: + """Return a valid, unique Python field name for a JSON property key. + + Keys that aren't valid identifiers (or would be Pydantic-private via a + leading underscore) are sanitized; the original key is preserved as the + field alias by the caller. + """ + if _IDENTIFIER_RE.match(prop_key) and prop_key not in existing: + return prop_key + sanitized = re.sub(r"[^A-Za-z0-9_]", "_", prop_key).lstrip("_") + if not sanitized or sanitized[0].isdigit(): + sanitized = f"field_{sanitized}" + candidate = sanitized + suffix = 2 + while candidate in existing: + candidate = f"{sanitized}_{suffix}" + suffix += 1 + return candidate + + +def _unique_model_name(name_hint: str, ctx: _Ctx) -> str: + # PascalCase the hint (e.g. "ResponseFormat address geo" -> "ResponseFormatAddressGeo"). + parts = re.split(r"[^A-Za-z0-9]+", name_hint) + name = "".join(part[:1].upper() + part[1:] for part in parts if part) + # Class names can't be empty or start with a digit. + if not name or name[0].isdigit(): + name = f"Model{name}" + # Distinct hints can sanitize to the same name; suffix _2, _3, ... to disambiguate. + candidate = name + suffix = 2 + while candidate in ctx.used_names: + candidate = f"{name}_{suffix}" + suffix += 1 + ctx.used_names.add(candidate) + return candidate diff --git a/src/utils/search.py b/src/utils/search.py index 15721933..761b63e0 100644 --- a/src/utils/search.py +++ b/src/utils/search.py @@ -448,7 +448,7 @@ async def search( return search_results[0][:limit] return [] - async with tracked_db("search.messages") as managed_db: + async with tracked_db("search.messages", read_only=True) as managed_db: combined_results = await _run_search(managed_db) for message in combined_results: managed_db.expunge(message) diff --git a/src/utils/summarizer.py b/src/utils/summarizer.py index 42d18bcc..2abdb0f9 100644 --- a/src/utils/summarizer.py +++ b/src/utils/summarizer.py @@ -6,6 +6,7 @@ from functools import cache from inspect import cleandoc as c from typing import TypedDict +from nanoid import generate as generate_nanoid from sqlalchemy import update from sqlalchemy.ext.asyncio import AsyncSession @@ -219,6 +220,9 @@ async def create_short_summary( formatted_messages, output_words, previous_summary_text ) + # Mint a root span id. + # No session_id or run_id for tracing + trace_id = generate_nanoid() return await honcho_llm_call( model_config=_get_summary_model_config(), prompt=prompt, @@ -227,6 +231,9 @@ async def create_short_summary( workspace_name=workspace_name, call_purpose=CallPurpose.SUMMARY_SHORT.value, parent_category="summary", + trace_id=trace_id, + span_id=trace_id, + track_name="Short Summary", ), ) @@ -251,6 +258,9 @@ async def create_long_summary( formatted_messages, output_words, previous_summary_text ) + # Mint a root span id. + # No session_id or run_id for tracing + trace_id = generate_nanoid() return await honcho_llm_call( model_config=_get_summary_model_config(), prompt=prompt, @@ -259,6 +269,9 @@ async def create_long_summary( workspace_name=workspace_name, call_purpose=CallPurpose.SUMMARY_LONG.value, parent_category="summary", + trace_id=trace_id, + span_id=trace_id, + track_name="Long Summary", ), ) @@ -514,17 +527,13 @@ async def _create_and_save_summary( "ms", ) - # Emit telemetry event (only for non-fallback summaries) - # Note: Using AgentToolSummaryCreatedEvent with dummy run_id/iteration since - # this is called from the deriver, not from an agentic loop + # Emit telemetry event (only for non-fallback summaries). if not is_fallback: # `prompt_tokens` is set in the `if not is_fallback` block above for # both SHORT and LONG summary types — we're inside the same branch, so # it's guaranteed bound here. emit( AgentToolSummaryCreatedEvent( - run_id="deriver", # Placeholder - not from an agentic run - iteration=0, # Placeholder - not from an agentic loop parent_category="deriver", agent_type="summarizer", workspace_name=workspace_name, diff --git a/src/utils/types.py b/src/utils/types.py index 2a470d20..33a2993d 100644 --- a/src/utils/types.py +++ b/src/utils/types.py @@ -114,6 +114,12 @@ _embedding_run_id: ContextVar[str | None] = ContextVar("embedding_run_id", defau _embedding_parent_category: ContextVar[str | None] = ContextVar( "embedding_parent_category", default=None ) +# Honcho Session.id for the embedding's trace grouping (e.g. a dialectic +# prefetch embedding shares the dialectic invocation's session). None when the +# embedding isn't scoped to a session. +_embedding_session_id: ContextVar[str | None] = ContextVar( + "embedding_session_id", default=None +) def get_embedding_call_purpose() -> str | None: @@ -136,6 +142,11 @@ def get_embedding_parent_category() -> str | None: return _embedding_parent_category.get() +def get_embedding_session_id() -> str | None: + """Read the Honcho Session.id attached to the current embedding call scope.""" + return _embedding_session_id.get() + + @contextmanager def embedding_call_purpose( purpose: str, @@ -143,6 +154,7 @@ def embedding_call_purpose( workspace_name: str | None = None, run_id: str | None = None, parent_category: str | None = None, + session_id: str | None = None, ) -> Generator[None]: """Tag any embedding calls made inside this `with` block. @@ -172,6 +184,9 @@ def embedding_call_purpose( if parent_category is not None else None ) + session_id_token = ( + _embedding_session_id.set(session_id) if session_id is not None else None + ) try: yield finally: @@ -182,6 +197,8 @@ def embedding_call_purpose( _embedding_run_id.reset(run_id_token) if parent_category_token is not None: _embedding_parent_category.reset(parent_category_token) + if session_id_token is not None: + _embedding_session_id.reset(session_id_token) @dataclass diff --git a/src/vector_store/__init__.py b/src/vector_store/__init__.py index 7ea6455f..96ae172d 100644 --- a/src/vector_store/__init__.py +++ b/src/vector_store/__init__.py @@ -202,7 +202,16 @@ def _create_store_by_type(store_type: str) -> VectorStore: return TurbopufferVectorStore() elif store_type == "lancedb": - from src.vector_store.lancedb import LanceDBVectorStore + try: + from src.vector_store.lancedb import LanceDBVectorStore + except ImportError as exc: + raise RuntimeError( + "VECTOR_STORE.TYPE is set to 'lancedb', but the 'lancedb' package " + + "is not installed (for example on macOS Intel, where it is omitted " + + "from dependencies because PyPI has no wheel). " + + "Use TYPE 'pgvector' or 'turbopuffer', or install lancedb manually. " + + f"Original import error: {exc}" + ) from exc return LanceDBVectorStore() elif store_type == "qdrant": diff --git a/src/vector_store/lancedb.py b/src/vector_store/lancedb.py index 0b621971..1b4c4880 100644 --- a/src/vector_store/lancedb.py +++ b/src/vector_store/lancedb.py @@ -298,10 +298,15 @@ class LanceDBVectorStore(VectorStore): if not _VALID_IDENTIFIER_PATTERN.match(key): raise ValueError(f"Invalid filter key: {key!r}") - # Check if value is a dict with "in" operator - if isinstance(value, dict) and "in" in value: - # IN clause for list membership - in_values = cast(Sequence[Any], value["in"]) + # Membership: dict form {"in": [...]} or bare-list sugar + if (isinstance(value, dict) and "in" in value) or isinstance( + value, list | tuple | set + ): + in_values = ( + cast(Sequence[Any], value["in"]) + if isinstance(value, dict) + else list(cast(Sequence[Any], value)) + ) if in_values: escaped_values = [ f"'{str(v).replace(chr(39), chr(39) + chr(39))}'" @@ -310,6 +315,11 @@ class LanceDBVectorStore(VectorStore): for v in in_values ] conditions.append(f"{key} IN ({', '.join(escaped_values)})") + else: + # An empty membership list matches nothing. Emitting no + # condition would silently widen the result set + # (fail-open); force an always-false condition instead. + conditions.append("1 = 0") # Handle string values with proper quoting elif isinstance(value, str): # Escape single quotes in the value diff --git a/src/vector_store/turbopuffer.py b/src/vector_store/turbopuffer.py index ded0c691..e7825310 100644 --- a/src/vector_store/turbopuffer.py +++ b/src/vector_store/turbopuffer.py @@ -20,9 +20,7 @@ from . import VectorQueryResult, VectorRecord, VectorStore logger = logging.getLogger(__name__) -# Type aliases for Turbopuffer's filter formats -EqFilter = tuple[str, Literal["Eq"], Any] -InFilter = tuple[str, Literal["In"], Sequence[Any]] +# Type alias for Turbopuffer's AND filter format AndFilter = tuple[Literal["And"], Sequence[Filter]] DISTANCE_METRIC = "cosine_distance" @@ -245,13 +243,17 @@ class TurbopufferVectorStore(VectorStore): if not filters: return None - filter_list: list[EqFilter | InFilter] = [] + filter_list: list[Filter] = [] for key, value in filters.items(): # Check if value is a dict with "in" operator if isinstance(value, dict) and "in" in value: # Membership filter using "In" operator - in_values = cast(Sequence[Any], value["in"]) - filter_list.append((key, "In", in_values)) + in_values = list(cast(Sequence[Any], value["in"])) + filter_list.append(self._membership_filter(key, in_values)) + elif isinstance(value, list | tuple | set): + # Bare-list sugar: same membership semantics as {"in": [...]} + in_values = list(cast(Sequence[Any], value)) + filter_list.append(self._membership_filter(key, in_values)) else: # Simple equality filter using "Eq" operator filter_list.append((key, "Eq", cast(Any, value))) @@ -266,6 +268,20 @@ class TurbopufferVectorStore(VectorStore): and_filter: AndFilter = ("And", filter_list) return and_filter + @staticmethod + def _membership_filter(key: str, values: list[Any]) -> Filter: + """Build an "In" membership filter, failing closed on an empty list. + + Turbopuffer's empty-"In" semantics are undocumented, so an empty + allowlist emits an explicit contradiction (`Eq(x) AND NotEq(x)` is + false for every document) rather than risk a fail-open widening. + Mirrors lancedb's `1 = 0` guard. + """ + if not values: + never: AndFilter = ("And", [(key, "Eq", ""), (key, "NotEq", "")]) + return never + return (key, "In", values) + async def delete_many(self, namespace: str, ids: list[str]) -> None: """ Delete multiple vectors from Turbopuffer. diff --git a/tests/bench/calculate_expected_events.py b/tests/bench/calculate_expected_events.py index 348cd95c..b9696262 100644 --- a/tests/bench/calculate_expected_events.py +++ b/tests/bench/calculate_expected_events.py @@ -246,7 +246,7 @@ def calculate_question_events( # Calculate representation events # Each unique (session, observed) pair generates one representation event - # (assuming messages fit within REPRESENTATION_BATCH_MAX_TOKENS) + # (assuming messages fit within REPRESENTATION_BATCH_TARGET_INPUT_TOKENS) # When merge_sessions=True, all messages go into one session if merge_sessions: # One merged session = one representation event diff --git a/tests/conftest.py b/tests/conftest.py index 06a31ac8..1ec64055 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,8 @@ import logging +import os +import re +import time +import uuid from collections.abc import AsyncGenerator, Callable from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -13,7 +17,7 @@ from fastapi import Request from fastapi.responses import JSONResponse from fastapi.testclient import TestClient from nanoid import generate as generate_nanoid -from sqlalchemy import text +from sqlalchemy import create_engine, text from sqlalchemy.engine.url import URL, make_url from sqlalchemy.exc import OperationalError, ProgrammingError from sqlalchemy.ext.asyncio import ( @@ -25,19 +29,25 @@ from sqlalchemy.ext.asyncio import ( from sqlalchemy_utils import ( create_database, # pyright: ignore[reportUnknownVariableType] database_exists, # pyright: ignore[reportUnknownVariableType] - drop_database, # pyright: ignore[reportUnknownVariableType] ) from src import models from src.cache.client import cache from src.config import settings from src.db import Base -from src.dependencies import get_db +from src.dependencies import get_db, get_read_db from src.exceptions import HonchoException -from src.main import app from src.models import Peer, Workspace from src.security import JWTParams, create_admin_jwt, create_jwt +# Disable Langfuse for the whole suite before importing src.main: @conditional_observe +# binds to settings.LANGFUSE_PUBLIC_KEY at import time, so blanking it here keeps mocked +# test calls from emitting traces to a configured Langfuse backend. Tests that exercise +# Langfuse patch settings.LANGFUSE_PUBLIC_KEY themselves. +settings.LANGFUSE_PUBLIC_KEY = None + +from src.main import app # noqa: E402 + # Create a custom handler that doesn't get closed prematurely class TestHandler(logging.Handler): @@ -78,6 +88,9 @@ _RUNTIME_MOCK_TEST_BLOCKLIST_PREFIXES = ( # LLM transport tests mock providers directly and don't need database/runtime setup. "tests/utils/test_length_finish_reason.py", "tests/utils/test_clients.py", + # Pure JWT scope tests — operate on src.security directly, no DB needed. + "tests/test_security.py", + "tests/test_generate_jwt_script.py", ) _LIVE_LLM_MARKER = "live_llm" @@ -118,11 +131,122 @@ def pytest_collection_modifyitems( item.add_marker(skip_live) +_RUN_ID_ENV_VAR = "HONCHO_TEST_RUN_ID" +_RUN_ID_TIME_FORMAT = "%Y%m%d%H%M%S" + +# Only a database whose name carries a run-id timestamp this old is swept. Long +# enough that no live suite is ever this stale, short enough that a leak from the +# morning is gone by the afternoon. +_STALE_DB_AGE_SECONDS = 2 * 60 * 60 + +# test_db_<14-digit timestamp>_<4 hex>[_gwN] -- only names this function minted. +# A pinned HONCHO_TEST_RUN_ID deliberately won't match, so it's never swept. +_SWEEPABLE_DB_NAME = re.compile(r"^test_db_(\d{14})_[0-9a-f]{4}(?:_gw\d+)?$") + + +def pytest_configure(config: pytest.Config) -> None: # pyright: ignore[reportUnusedParameter] + """Stamp this pytest run with an id so its databases can't collide with another run's. + + The xdist controller runs this first and its environment is inherited by the + workers it spawns, so `setdefault` gives every worker in a run the same id + while separate runs (concurrent worktrees, two agents, a local run alongside + CI) each get their own. Set the env var yourself to pin a stable name. + + The id leads with a sortable local-time timestamp so leaked databases can be + aged out (see `_sweep_stale_test_databases`); the random tail keeps two runs + starting in the same second apart. + """ + + os.environ.setdefault( + _RUN_ID_ENV_VAR, + f"{time.strftime(_RUN_ID_TIME_FORMAT)}_{uuid.uuid4().hex[:4]}", + ) + + # Workers inherit the controller's env and would each redo this. + if os.environ.get("PYTEST_XDIST_WORKER") is None: + _sweep_stale_test_databases() + + def _get_test_db_url(worker_id: str) -> URL: """Get a worker-specific test database URL for pytest-xdist parallelism.""" - db_name = "test_db" if worker_id == "master" else f"test_db_{worker_id}" - return CONNECTION_URI.set(database=db_name) + run_id = os.environ.get(_RUN_ID_ENV_VAR, "local") + suffix = "" if worker_id == "master" else f"_{worker_id}" + return CONNECTION_URI.set(database=f"test_db_{run_id}{suffix}") + + +def _drop_database(db_url: URL) -> None: + """Drop a test database, evicting any connections still holding it open. + + WITH (FORCE) (pg13+) is what makes this reliable: a pooled connection that + outlives engine disposal, or an xdist worker killed mid-query, otherwise + leaves the drop failing with "database is being accessed by other users". + """ + + name = db_url.database + if not name: + return + + # Maintenance connection: you cannot drop the database you're connected to. + engine = create_engine( + db_url.set(database="postgres"), isolation_level="AUTOCOMMIT" + ) + try: + with engine.connect() as conn: + conn.exec_driver_sql(f'DROP DATABASE IF EXISTS "{name}" WITH (FORCE)') + finally: + engine.dispose() + + +def _sweep_stale_test_databases() -> None: + """Reclaim test databases left behind by runs that died before teardown. + + A run killed by SIGKILL, an IDE stop button, an OOM'd worker or `-x` on a hang + never reaches the `db_engine` teardown, and since every run mints its own + database name nothing later reuses (and thus cleans) it. + + Two guards keep this from touching a suite that is currently running, which is + the whole point of per-run names: + + - the run-id timestamp in the name must be older than `_STALE_DB_AGE_SECONDS` + - the database must have no backends connected to it right now + + Each covers the other's blind spot: the age check is immune to the race where + a database has been created but its first worker hasn't connected yet, and the + connection check catches a genuinely long-running suite. Failure to sweep is + logged and ignored -- it must never fail a test session. + """ + + cutoff = time.strftime( + _RUN_ID_TIME_FORMAT, time.localtime(time.time() - _STALE_DB_AGE_SECONDS) + ) + + try: + engine = create_engine( + CONNECTION_URI.set(database="postgres"), isolation_level="AUTOCOMMIT" + ) + try: + with engine.connect() as conn: + names = [ + row[0] + for row in conn.exec_driver_sql( + "SELECT datname FROM pg_database d " + + "WHERE NOT EXISTS (" + + " SELECT 1 FROM pg_stat_activity WHERE datname = d.datname" + + ")" + ) + ] + finally: + engine.dispose() + + for name in names: + match = _SWEEPABLE_DB_NAME.match(name) + if match is None or match.group(1) >= cutoff: + continue + logger.info(f"Dropping stale test database: {name}") + _drop_database(CONNECTION_URI.set(database=name)) + except Exception as e: + logger.warning(f"Could not sweep stale test databases: {e}") # Test API authorization - no longer needed as module-level constants @@ -183,11 +307,21 @@ async def setup_test_database(db_url: URL): return engine -async def _truncate_all_tables(engine: AsyncEngine) -> None: - """Remove all data from every mapped table while resetting identities.""" +async def _clear_all_tables(engine: AsyncEngine) -> None: + """Remove all data from every mapped table between tests. + + Uses DELETE rather than TRUNCATE: TRUNCATE rewrites the relfilenode of every + table and index it touches, so it costs a flat ~33ms for this schema's 11 + tables / 41 indexes no matter how few rows a test actually wrote. DELETE of + the same (near-empty) tables, batched into one round trip, is ~3ms. Tables go + in reverse dependency order so foreign keys are satisfied without CASCADE. + + This does not reset identity sequences, so tests must not assert on absolute + generated id values -- compare against the ids the test itself created. + """ table_names: list[str] = [] - for table in Base.metadata.sorted_tables: + for table in reversed(Base.metadata.sorted_tables): if table.schema: table_names.append(f'"{table.schema}"."{table.name}"') else: @@ -196,9 +330,9 @@ async def _truncate_all_tables(engine: AsyncEngine) -> None: if not table_names: return - joined_names = ", ".join(table_names) + statement = "; ".join(f"DELETE FROM {name}" for name in table_names) async with engine.begin() as conn: - await conn.execute(text(f"TRUNCATE {joined_names} RESTART IDENTITY CASCADE")) + await conn.exec_driver_sql(statement) @pytest_asyncio.fixture(scope="session") @@ -232,7 +366,7 @@ async def db_engine(worker_id: str): for table in Base.metadata.tables.values(): table.schema = original_schema - drop_database(test_db_url) + _drop_database(test_db_url) @pytest_asyncio.fixture(scope="function") @@ -246,7 +380,7 @@ async def db_session(db_engine: AsyncEngine): finally: await session.rollback() finally: - await _truncate_all_tables(db_engine) + await _clear_all_tables(db_engine) @pytest_asyncio.fixture(scope="session") @@ -339,6 +473,10 @@ async def client( yield db_session app.dependency_overrides[get_db] = override_get_db + # Read-only routes use get_read_db (AUTOCOMMIT engine) in production; in + # tests they must see the same per-test database/session as writes, both + # for isolation and so data written by a test is visible to its reads. + app.dependency_overrides[get_read_db] = override_get_db # No-op the startup embedding-schema validator inside the lifespan. The # global `engine` it would inspect points to a DB that isn't migrated in @@ -416,7 +554,7 @@ async def sample_data( db_session.add(test_peer) # Commit so data is visible to independent tracked_db sessions. - # _truncate_all_tables handles cleanup between tests. + # _clear_all_tables handles cleanup between tests. await db_session.commit() yield test_workspace, test_peer @@ -479,6 +617,9 @@ def mock_openai_embeddings(request: pytest.FixtureRequest): patch( "src.embedding_client.embedding_client.simple_batch_embed" ) as mock_simple_batch_embed, + patch( + "src.embedding_client.embedding_client.prepare_chunks" + ) as mock_prepare_chunks, patch("src.embedding_client.embedding_client.batch_embed") as mock_batch_embed, ): # Mock the embed method to return content-dependent embedding @@ -492,6 +633,14 @@ def mock_openai_embeddings(request: pytest.FixtureRequest): mock_simple_batch_embed.side_effect = mock_simple_batch_embed_func + def mock_prepare_chunks_func( + id_resource_dict: dict[str, str], + ) -> dict[str, list[str]]: + # No real tokenizer in mocks: treat each input as a single chunk. + return {text_id: [text] for text_id, text in id_resource_dict.items()} + + mock_prepare_chunks.side_effect = mock_prepare_chunks_func + # Mock the batch_embed method to return content-dependent embeddings async def mock_batch_embed_func( id_resource_dict: dict[str, str], @@ -506,6 +655,7 @@ def mock_openai_embeddings(request: pytest.FixtureRequest): yield { "embed": mock_embed, "simple_batch_embed": mock_simple_batch_embed, + "prepare_chunks": mock_prepare_chunks, "batch_embed": mock_batch_embed, } @@ -646,8 +796,15 @@ def mock_llm_call_functions(request: pytest.FixtureRequest): mock_short_summary.return_value = "Test short summary content" mock_long_summary.return_value = "Test long summary content" - # Mock agentic_chat to return a string (matching actual return type) - mock_agentic_chat.return_value = "Test dialectic response" + # Mock agentic_chat to return a string (matching actual return type). + # With a response_model (structured output) the real function returns + # a JSON string, so mirror that for SDK clients that parse content. + async def _agentic_chat_response(*_args: object, **kwargs: object) -> str: + if kwargs.get("response_model") is not None: + return "{}" + return "Test dialectic response" + + mock_agentic_chat.side_effect = _agentic_chat_response yield { "short_summary": mock_short_summary, @@ -785,38 +942,48 @@ def mock_tracked_db(request: pytest.FixtureRequest): yield return - from contextlib import asynccontextmanager + from contextlib import ExitStack, asynccontextmanager db_engine = request.getfixturevalue("db_engine") session_factory = async_sessionmaker(bind=db_engine, expire_on_commit=False) @asynccontextmanager - async def mock_tracked_db_context(_: str | None = None): + async def mock_tracked_db_context(_: str | None = None, *, read_only: bool = False): + # read_only is accepted (and ignored): in tests both engines resolve to + # the same per-test database session. + del read_only async with session_factory() as session: yield session - with ( - patch("src.dependencies.tracked_db", mock_tracked_db_context), - patch("src.deriver.queue_manager.tracked_db", mock_tracked_db_context), - patch("src.deriver.consumer.tracked_db", mock_tracked_db_context), - patch("src.deriver.enqueue.tracked_db", mock_tracked_db_context), - patch("src.routers.peers.tracked_db", mock_tracked_db_context), - patch("src.crud.representation.tracked_db", mock_tracked_db_context), - patch("src.dreamer.orchestrator.tracked_db", mock_tracked_db_context), - patch("src.dreamer.dream_scheduler.tracked_db", mock_tracked_db_context), - patch("src.dialectic.chat.tracked_db", mock_tracked_db_context), - patch("src.utils.summarizer.tracked_db", mock_tracked_db_context), - patch("src.webhooks.events.tracked_db", mock_tracked_db_context), - patch("src.webhooks.webhook_delivery.tracked_db", mock_tracked_db_context), - patch("src.utils.agent_tools.tracked_db", mock_tracked_db_context), - patch("src.utils.search.tracked_db", mock_tracked_db_context), - patch("src.crud.document.tracked_db", mock_tracked_db_context), - patch("src.crud.message.tracked_db", mock_tracked_db_context), - patch("src.reconciler.sync_vectors.tracked_db", mock_tracked_db_context), - patch("src.dialectic.core.tracked_db", mock_tracked_db_context), - patch("src.dreamer.specialists.tracked_db", mock_tracked_db_context), - patch("src.dreamer.surprisal.tracked_db", mock_tracked_db_context), - ): + # Each module imports tracked_db by name, so patch every import site. + # Use ExitStack (not a parenthesized `with`) to stay under CPython's + # 20-statically-nested-block limit as this list grows. + tracked_db_targets = [ + "src.dependencies.tracked_db", + "src.deriver.queue_manager.tracked_db", + "src.deriver.consumer.tracked_db", + "src.deriver.enqueue.tracked_db", + "src.routers.peers.tracked_db", + "src.crud.representation.tracked_db", + "src.dreamer.orchestrator.tracked_db", + "src.dreamer.dream_scheduler.tracked_db", + "src.dialectic.chat.tracked_db", + "src.utils.summarizer.tracked_db", + "src.webhooks.events.tracked_db", + "src.webhooks.webhook_delivery.tracked_db", + "src.utils.agent_tools.tracked_db", + "src.utils.search.tracked_db", + "src.crud.document.tracked_db", + "src.crud.message.tracked_db", + "src.reconciler.sync_vectors.tracked_db", + "src.reconciler.embed_now.tracked_db", + "src.dialectic.core.tracked_db", + "src.dreamer.specialists.tracked_db", + "src.dreamer.surprisal.tracked_db", + ] + with ExitStack() as stack: + for target in tracked_db_targets: + stack.enter_context(patch(target, mock_tracked_db_context)) yield diff --git a/tests/crud/test_document.py b/tests/crud/test_document.py index f25b75d2..5f03c80b 100644 --- a/tests/crud/test_document.py +++ b/tests/crud/test_document.py @@ -1,4 +1,5 @@ import datetime +from unittest.mock import AsyncMock, patch import pytest from nanoid import generate as generate_nanoid @@ -6,6 +7,7 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from src import crud, models, schemas +from src.crud.document import SemanticRejectionResult, is_rejected_duplicate from src.exceptions import ResourceNotFoundException @@ -274,6 +276,606 @@ class TestDocumentCRUD: assert len(results) == 1 assert results[0].id == times_derived_map[2] + @pytest.mark.asyncio + async def test_most_derived_orders_by_recency_when_reinforcement_ties( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """Regression: when times_derived ties, most-derived must fall back to + recency, not insertion order. Otherwise stale conclusions stick to the + front of the injected representation (the mid-Jan stickiness bug).""" + test_workspace, test_peer = sample_data + test_peer2, test_session, _ = await self._setup_test_data( + db_session, test_workspace, test_peer + ) + + base = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) + # Three conclusions, all reinforced once -- the real-world steady state + # before the fix -- inserted oldest-first. + for i in range(3): + db_session.add( + models.Document( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + content=f"tie {i}", + session_name=test_session.name, + times_derived=1, + created_at=base + datetime.timedelta(days=i), + ) + ) + # A genuinely reinforced conclusion that is also the oldest of all. + db_session.add( + models.Document( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + content="hot", + session_name=test_session.name, + times_derived=5, + created_at=base - datetime.timedelta(days=10), + ) + ) + await db_session.flush() + + docs = await crud.query_documents_most_derived( + db_session, + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + limit=10, + ) + contents = [d.content for d in docs] + # Primary sort still wins: the actually-reinforced conclusion leads. + assert contents[0] == "hot" + # Ties break toward most-recent, not oldest-inserted. + assert contents[1:] == ["tie 2", "tie 1", "tie 0"] + + @pytest.mark.asyncio + async def test_duplicate_rejection_reinforces_existing( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """Rejecting a new duplicate must bump the surviving doc's times_derived.""" + test_workspace, test_peer = sample_data + test_peer2, test_session, _ = await self._setup_test_data( + db_session, test_workspace, test_peer + ) + + await crud.create_documents( + db_session, + [ + schemas.DocumentCreate( + content="eri loves cats and dogs and birds and snakes", + embedding=[0.5] * 1536, + session_name=test_session.name, + times_derived=1, + metadata=schemas.DocumentMetadata( + message_ids=[1], + message_created_at="2026-01-01T00:00:00Z", + ), + ) + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + + # Fewer unique tokens -> existing wins -> new doc is rejected. + new_doc = schemas.DocumentCreate( + content="eri loves cats", + embedding=[0.5] * 1536, + session_name=test_session.name, + times_derived=1, + metadata=schemas.DocumentMetadata( + message_ids=[2], + message_created_at="2026-01-02T00:00:00Z", + ), + ) + rejected = await is_rejected_duplicate( + db_session, + new_doc, + test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + + assert rejected is SemanticRejectionResult.REJECTED + surviving = ( + await db_session.execute( + select(models.Document).where( + models.Document.workspace_name == test_workspace.name, + models.Document.observer == test_peer.name, + models.Document.observed == test_peer2.name, + models.Document.deleted_at.is_(None), + ) + ) + ).scalar_one() + assert surviving.times_derived == 2 + + @pytest.mark.asyncio + async def test_duplicate_replacement_carries_count_forward( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """When a new duplicate wins, it must inherit the replaced doc's count + 1 + rather than resetting reinforcement to 1.""" + test_workspace, test_peer = sample_data + test_peer2, test_session, _ = await self._setup_test_data( + db_session, test_workspace, test_peer + ) + + await crud.create_documents( + db_session, + [ + schemas.DocumentCreate( + content="eri loves cats", + embedding=[0.5] * 1536, + session_name=test_session.name, + times_derived=3, + metadata=schemas.DocumentMetadata( + message_ids=[1], + message_created_at="2026-01-01T00:00:00Z", + ), + ) + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + + # More information -> new wins -> existing is soft-deleted. + new_doc = schemas.DocumentCreate( + content="eri loves cats and dogs", + embedding=[0.5] * 1536, + session_name=test_session.name, + times_derived=1, + metadata=schemas.DocumentMetadata( + message_ids=[2], + message_created_at="2026-01-02T00:00:00Z", + ), + ) + rejected = await is_rejected_duplicate( + db_session, + new_doc, + test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + + assert rejected is SemanticRejectionResult.REPLACED_EXISTING + # Count carried forward onto the replacement (3 -> 4), not reset to 1. + assert new_doc.times_derived == 4 + live = ( + ( + await db_session.execute( + select(models.Document).where( + models.Document.workspace_name == test_workspace.name, + models.Document.observer == test_peer.name, + models.Document.observed == test_peer2.name, + models.Document.deleted_at.is_(None), + ) + ) + ) + .scalars() + .all() + ) + # Original is soft-deleted; replacement isn't inserted until create_documents runs. + assert len(live) == 0 + + @pytest.mark.asyncio + async def test_exact_dedup_within_batch_drops_repeat( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """Exact (case/whitespace-insensitive) duplicates within a single batch + collapse to one document, even with semantic dedup disabled.""" + test_workspace, test_peer = sample_data + test_peer2, test_session, _ = await self._setup_test_data( + db_session, test_workspace, test_peer + ) + + # Three "exact" matches that differ only by case/surrounding whitespace. + doc_schemas = [ + schemas.DocumentCreate( + content="User likes coffee", + embedding=[0.1] * 1536, + session_name=test_session.name, + metadata=schemas.DocumentMetadata( + message_ids=[1], + message_created_at="2026-01-01T00:00:00Z", + ), + ), + schemas.DocumentCreate( + content="user likes coffee", + embedding=[0.2] * 1536, + session_name=test_session.name, + metadata=schemas.DocumentMetadata( + message_ids=[2], + message_created_at="2026-01-01T00:01:00Z", + ), + ), + schemas.DocumentCreate( + content=" User likes coffee\n", + embedding=[0.3] * 1536, + session_name=test_session.name, + metadata=schemas.DocumentMetadata( + message_ids=[3], + message_created_at="2026-01-01T00:02:00Z", + ), + ), + ] + + result = await crud.create_documents( + db_session, + documents=doc_schemas, + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + deduplicate=False, + ) + accepted = result.created_documents + + assert len(accepted) == 1 + assert result.exact_dup_in_batch_count == 2 + assert result.exact_dup_existing_count == 0 + assert result.semantic_dup_rejected_count == 0 + assert result.semantic_dup_replaced_count == 0 + live = ( + ( + await db_session.execute( + select(models.Document).where( + models.Document.workspace_name == test_workspace.name, + models.Document.observer == test_peer.name, + models.Document.observed == test_peer2.name, + models.Document.deleted_at.is_(None), + ) + ) + ) + .scalars() + .all() + ) + assert len(live) == 1 + # Within-batch repeats are dropped silently, no reinforcement. + assert live[0].times_derived == 1 + + @pytest.mark.asyncio + async def test_exact_dedup_against_existing_reinforces( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """An exact match of an existing live document is rejected and reinforces + the existing row, even with semantic dedup disabled.""" + test_workspace, test_peer = sample_data + test_peer2, test_session, _ = await self._setup_test_data( + db_session, test_workspace, test_peer + ) + + await crud.create_documents( + db_session, + [ + schemas.DocumentCreate( + content="User likes coffee", + embedding=[0.1] * 1536, + session_name=test_session.name, + times_derived=1, + metadata=schemas.DocumentMetadata( + message_ids=[1], + message_created_at="2026-01-01T00:00:00Z", + ), + ) + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + deduplicate=False, + ) + + # Case/whitespace variant of the existing content -> exact match. + result = await crud.create_documents( + db_session, + [ + schemas.DocumentCreate( + content="user likes coffee ", + embedding=[0.9] * 1536, + session_name=test_session.name, + times_derived=1, + metadata=schemas.DocumentMetadata( + message_ids=[2], + message_created_at="2026-01-02T00:00:00Z", + ), + ) + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + deduplicate=False, + ) + accepted = result.created_documents + + assert len(accepted) == 0 + assert result.exact_dup_existing_count == 1 + assert result.exact_dup_in_batch_count == 0 + assert result.semantic_dup_rejected_count == 0 + assert result.semantic_dup_replaced_count == 0 + surviving = ( + ( + await db_session.execute( + select(models.Document).where( + models.Document.workspace_name == test_workspace.name, + models.Document.observer == test_peer.name, + models.Document.observed == test_peer2.name, + models.Document.deleted_at.is_(None), + ) + ) + ) + .scalars() + .all() + ) + assert len(surviving) == 1 + assert surviving[0].content == "User likes coffee" + assert surviving[0].times_derived == 2 + + @pytest.mark.asyncio + async def test_exact_dedup_honors_incoming_times_derived( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """Reinforcement folds in an incoming doc that already carries + accumulated reinforcement: the existing row becomes + ``greatest(existing + 1, incoming)``.""" + test_workspace, test_peer = sample_data + test_peer2, test_session, _ = await self._setup_test_data( + db_session, test_workspace, test_peer + ) + + async def _live() -> list[models.Document]: + return list( + ( + await db_session.execute( + select(models.Document).where( + models.Document.workspace_name == test_workspace.name, + models.Document.observer == test_peer.name, + models.Document.observed == test_peer2.name, + models.Document.deleted_at.is_(None), + ) + ) + ) + .scalars() + .all() + ) + + # Existing row already reinforced twice. + await crud.create_documents( + db_session, + [ + schemas.DocumentCreate( + content="User likes coffee", + embedding=[0.1] * 1536, + session_name=test_session.name, + times_derived=2, + metadata=schemas.DocumentMetadata( + message_ids=[1], + message_created_at="2026-01-01T00:00:00Z", + ), + ) + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + deduplicate=False, + ) + + # Incoming exact match claims more accumulated reinforcement (5) than + # existing + 1 (3) -> incoming wins. + accepted = ( + await crud.create_documents( + db_session, + [ + schemas.DocumentCreate( + content="user likes coffee ", + embedding=[0.9] * 1536, + session_name=test_session.name, + times_derived=5, + metadata=schemas.DocumentMetadata( + message_ids=[2], + message_created_at="2026-01-02T00:00:00Z", + ), + ) + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + deduplicate=False, + ) + ).created_documents + assert len(accepted) == 0 + live = await _live() + assert len(live) == 1 + assert live[0].times_derived == 5 + + # A normal re-derivation (times_derived defaults to 1) now bumps by one: + # greatest(existing + 1, 1) -> existing + 1. + accepted = ( + await crud.create_documents( + db_session, + [ + schemas.DocumentCreate( + content="USER LIKES COFFEE", + embedding=[0.4] * 1536, + session_name=test_session.name, + metadata=schemas.DocumentMetadata( + message_ids=[3], + message_created_at="2026-01-03T00:00:00Z", + ), + ) + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + deduplicate=False, + ) + ).created_documents + assert len(accepted) == 0 + live = await _live() + assert len(live) == 1 + assert live[0].times_derived == 6 + + @pytest.mark.asyncio + async def test_exact_dedup_flushes_before_semantic_replacement( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """An exact-match reinforcement in a batch must be visible to a later + semantic replacement of the same existing row when autoflush is off.""" + test_workspace, test_peer = sample_data + test_peer2, test_session, _ = await self._setup_test_data( + db_session, test_workspace, test_peer + ) + + await crud.create_documents( + db_session, + [ + schemas.DocumentCreate( + content="User likes coffee", + embedding=[0.5] * 1536, + session_name=test_session.name, + times_derived=1, + metadata=schemas.DocumentMetadata( + message_ids=[1], + message_created_at="2026-01-01T00:00:00Z", + ), + ) + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + deduplicate=False, + ) + + db_session.autoflush = False + result = await crud.create_documents( + db_session, + [ + schemas.DocumentCreate( + content=" user likes coffee ", + embedding=[0.5] * 1536, + session_name=test_session.name, + times_derived=1, + metadata=schemas.DocumentMetadata( + message_ids=[2], + message_created_at="2026-01-02T00:00:00Z", + ), + ), + schemas.DocumentCreate( + content="User likes coffee and tea", + embedding=[0.5] * 1536, + session_name=test_session.name, + times_derived=1, + metadata=schemas.DocumentMetadata( + message_ids=[3], + message_created_at="2026-01-03T00:00:00Z", + ), + ), + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + deduplicate=True, + ) + accepted = result.created_documents + + assert len(accepted) == 1 + assert accepted[0].content == "User likes coffee and tea" + assert result.exact_dup_existing_count == 1 + assert result.semantic_dup_replaced_count == 1 + assert result.exact_dup_in_batch_count == 0 + assert result.semantic_dup_rejected_count == 0 + + surviving = ( + ( + await db_session.execute( + select(models.Document).where( + models.Document.workspace_name == test_workspace.name, + models.Document.observer == test_peer.name, + models.Document.observed == test_peer2.name, + models.Document.deleted_at.is_(None), + ) + ) + ) + .scalars() + .all() + ) + assert len(surviving) == 1 + assert surviving[0].content == "User likes coffee and tea" + assert surviving[0].times_derived == 3 + + @pytest.mark.asyncio + async def test_semantic_dedup_rejected_counts( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """A semantically-similar doc with less information than the existing one + is rejected, and the rejection is counted on the result.""" + test_workspace, test_peer = sample_data + test_peer2, test_session, _ = await self._setup_test_data( + db_session, test_workspace, test_peer + ) + + await crud.create_documents( + db_session, + [ + schemas.DocumentCreate( + content="eri loves cats and dogs and birds and snakes", + embedding=[0.5] * 1536, + session_name=test_session.name, + times_derived=1, + metadata=schemas.DocumentMetadata( + message_ids=[1], + message_created_at="2026-01-01T00:00:00Z", + ), + ) + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + + # Fewer unique tokens -> existing wins -> new doc is rejected. + result = await crud.create_documents( + db_session, + [ + schemas.DocumentCreate( + content="eri loves cats", + embedding=[0.5] * 1536, + session_name=test_session.name, + times_derived=1, + metadata=schemas.DocumentMetadata( + message_ids=[2], + message_created_at="2026-01-02T00:00:00Z", + ), + ) + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + deduplicate=True, + ) + + assert len(result.created_documents) == 0 + assert result.semantic_dup_rejected_count == 1 + assert result.exact_dup_in_batch_count == 0 + assert result.exact_dup_existing_count == 0 + assert result.semantic_dup_replaced_count == 0 + @pytest.mark.asyncio async def test_delete_document_success( self, @@ -379,15 +981,17 @@ class TestDocumentCRUD: ] # Create documents - count = await crud.create_documents( - db_session, - documents=doc_schemas, - workspace_name=test_workspace.name, - observer=test_peer.name, - observed=test_peer2.name, - ) + created_documents = ( + await crud.create_documents( + db_session, + documents=doc_schemas, + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + ).created_documents - assert len(count) == 2 + assert len(created_documents) == 2 # Verify documents were created stmt = select(models.Document).where( @@ -401,3 +1005,332 @@ class TestDocumentCRUD: assert len(documents) == 2 assert documents[0].content in ["Observation 1", "Observation 2"] assert documents[1].content in ["Observation 1", "Observation 2"] + + +class TestSessionPurityInvariant: + """Regression tests for the explicit-document session-purity invariant. + + Explicit documents are session-pure records of what was derived from one + session's messages (the Scopes copy-by-session model depends on this): + + - an explicit document must always carry a non-null session_name + - dedup/merge (exact and semantic) must never cross document levels + - dedup/merge must never cross sessions for explicit documents + """ + + async def _setup( + self, + db_session: AsyncSession, + test_workspace: models.Workspace, + test_peer: models.Peer, + ) -> tuple[models.Peer, models.Session, models.Session]: + """Create an observed peer, two sessions, and the collection.""" + test_peer2 = models.Peer( + name=str(generate_nanoid()), workspace_name=test_workspace.name + ) + db_session.add(test_peer2) + session_a = models.Session( + name=str(generate_nanoid()), workspace_name=test_workspace.name + ) + session_b = models.Session( + name=str(generate_nanoid()), workspace_name=test_workspace.name + ) + db_session.add_all([session_a, session_b]) + await db_session.flush() + + collection = models.Collection( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + db_session.add(collection) + await db_session.flush() + return test_peer2, session_a, session_b + + def _doc( + self, + content: str, + *, + session_name: str | None, + level: str = "explicit", + message_id: int = 1, + ) -> schemas.DocumentCreate: + return schemas.DocumentCreate( + content=content, + embedding=[0.1] * 1536, + session_name=session_name, + level=level, # pyright: ignore[reportArgumentType] + metadata=schemas.DocumentMetadata( + message_ids=[message_id], + message_created_at="2026-01-01T00:00:00Z", + ), + ) + + async def _live_docs( + self, + db_session: AsyncSession, + workspace_name: str, + observer: str, + observed: str, + ) -> list[models.Document]: + return list( + ( + await db_session.execute( + select(models.Document).where( + models.Document.workspace_name == workspace_name, + models.Document.observer == observer, + models.Document.observed == observed, + models.Document.deleted_at.is_(None), + ) + ) + ) + .scalars() + .all() + ) + + @pytest.mark.asyncio + async def test_explicit_without_session_is_refused( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """An explicit document with session_name=None must not be written; + derived levels remain allowed without a session (dream output).""" + test_workspace, test_peer = sample_data + test_peer2, _, _ = await self._setup(db_session, test_workspace, test_peer) + + accepted = ( + await crud.create_documents( + db_session, + [ + self._doc("Global explicit fact", session_name=None), + self._doc( + "Dream-derived conclusion", + session_name=None, + level="deductive", + message_id=2, + ), + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + ).created_documents + + assert [d.content for d in accepted] == ["Dream-derived conclusion"] + live = await self._live_docs( + db_session, test_workspace.name, test_peer.name, test_peer2.name + ) + assert len(live) == 1 + assert live[0].level == "deductive" + + @pytest.mark.asyncio + async def test_exact_dedup_never_merges_explicit_across_sessions( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """The same explicit fact stated in two sessions produces two + session-pure documents; the other session's row is not reinforced.""" + test_workspace, test_peer = sample_data + test_peer2, session_a, session_b = await self._setup( + db_session, test_workspace, test_peer + ) + + await crud.create_documents( + db_session, + [self._doc("User likes coffee", session_name=session_a.name)], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + accepted = ( + await crud.create_documents( + db_session, + [ + self._doc( + "user likes coffee ", session_name=session_b.name, message_id=2 + ) + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + ).created_documents + + assert len(accepted) == 1 + live = await self._live_docs( + db_session, test_workspace.name, test_peer.name, test_peer2.name + ) + assert len(live) == 2 + assert {doc.session_name for doc in live} == {session_a.name, session_b.name} + assert all(doc.times_derived == 1 for doc in live) + + @pytest.mark.asyncio + async def test_exact_dedup_never_merges_across_levels( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """An explicit fact must not be dropped/reinforced against a derived + document that happens to share its content.""" + test_workspace, test_peer = sample_data + test_peer2, session_a, _ = await self._setup( + db_session, test_workspace, test_peer + ) + + await crud.create_documents( + db_session, + [ + self._doc( + "User likes coffee", session_name=session_a.name, level="deductive" + ) + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + accepted = ( + await crud.create_documents( + db_session, + [ + self._doc( + "User likes coffee", session_name=session_a.name, message_id=2 + ) + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + ).created_documents + + assert len(accepted) == 1 + live = await self._live_docs( + db_session, test_workspace.name, test_peer.name, test_peer2.name + ) + assert len(live) == 2 + assert {doc.level for doc in live} == {"explicit", "deductive"} + assert all(doc.times_derived == 1 for doc in live) + + @pytest.mark.asyncio + async def test_exact_dedup_still_merges_derived_levels_across_sessions( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """Derived levels are consolidations, not session-pure records: + cross-session exact dedup still reinforces the existing row.""" + test_workspace, test_peer = sample_data + test_peer2, session_a, session_b = await self._setup( + db_session, test_workspace, test_peer + ) + + await crud.create_documents( + db_session, + [ + self._doc( + "Probably a morning person", + session_name=session_a.name, + level="deductive", + ) + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + accepted = ( + await crud.create_documents( + db_session, + [ + self._doc( + "probably a morning person", + session_name=session_b.name, + level="deductive", + message_id=2, + ) + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + ).created_documents + + assert len(accepted) == 0 + live = await self._live_docs( + db_session, test_workspace.name, test_peer.name, test_peer2.name + ) + assert len(live) == 1 + assert live[0].times_derived == 2 + + @pytest.mark.asyncio + async def test_semantic_dedup_scoped_to_level_and_session( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """is_rejected_duplicate must constrain candidate search to the same + level, and to the same session for explicit documents.""" + test_workspace, test_peer = sample_data + test_peer2, session_a, _ = await self._setup( + db_session, test_workspace, test_peer + ) + + explicit_doc = self._doc("User likes coffee", session_name=session_a.name) + with patch( + "src.crud.document.query_documents", new=AsyncMock(return_value=[]) + ) as mock_query: + rejected = await is_rejected_duplicate( + db_session, + explicit_doc, + test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + assert rejected is SemanticRejectionResult.NOT_DUPLICATE + assert mock_query.await_args is not None + assert mock_query.await_args.kwargs["filters"] == { + "level": "explicit", + "session_name": session_a.name, + } + + deductive_doc = self._doc( + "User likes coffee", session_name=None, level="deductive" + ) + with patch( + "src.crud.document.query_documents", new=AsyncMock(return_value=[]) + ) as mock_query: + rejected = await is_rejected_duplicate( + db_session, + deductive_doc, + test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + assert rejected is SemanticRejectionResult.NOT_DUPLICATE + assert mock_query.await_args is not None + assert mock_query.await_args.kwargs["filters"] == {"level": "deductive"} + + @pytest.mark.asyncio + async def test_semantic_dedup_refuses_sessionless_explicit( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """A session-less explicit document has no valid merge partner: it is + never treated as a duplicate and no candidate search runs.""" + test_workspace, test_peer = sample_data + test_peer2, _, _ = await self._setup(db_session, test_workspace, test_peer) + + doc = self._doc("User likes coffee", session_name=None) + with patch( + "src.crud.document.query_documents", new=AsyncMock(return_value=[]) + ) as mock_query: + rejected = await is_rejected_duplicate( + db_session, + doc, + test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + assert rejected is SemanticRejectionResult.NOT_DUPLICATE + mock_query.assert_not_awaited() diff --git a/tests/crud/test_representation_manager.py b/tests/crud/test_representation_manager.py index 141b61c7..f551de76 100644 --- a/tests/crud/test_representation_manager.py +++ b/tests/crud/test_representation_manager.py @@ -1,6 +1,6 @@ from contextlib import asynccontextmanager -from datetime import datetime, timezone -from unittest.mock import AsyncMock, patch +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock, patch import pytest from nanoid import generate as generate_nanoid @@ -8,6 +8,7 @@ from sqlalchemy import func, update from sqlalchemy.ext.asyncio import AsyncSession from src import models +from src.crud.document import CreateDocumentsResult from src.crud.representation import RepresentationManager from src.schemas.configuration import ( ResolvedConfiguration, @@ -180,6 +181,294 @@ class TestRepresentationManagerSoftDelete: assert doc_live.id in result_ids assert doc_deleted.id not in result_ids + @pytest.mark.asyncio + async def test_query_documents_most_derived_ties_break_by_recency( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """Regression: when times_derived ties, the manager's most-derived query + must fall back to recency, not insertion order. Mirrors the equivalent + test on crud.query_documents_most_derived -- the query is duplicated in + both modules and must not drift.""" + test_workspace, test_peer = sample_data + test_peer2, test_session, _, manager = await self._setup( + db_session, test_workspace, test_peer + ) + + base = datetime(2026, 1, 1, tzinfo=timezone.utc) + # Three conclusions, all reinforced once, inserted oldest-first. + for i in range(3): + db_session.add( + models.Document( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + content=f"tie {i}", + session_name=test_session.name, + times_derived=1, + created_at=base + timedelta(days=i), + ) + ) + # A genuinely reinforced conclusion that is also the oldest of all. + db_session.add( + models.Document( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + content="hot", + session_name=test_session.name, + times_derived=5, + created_at=base - timedelta(days=10), + ) + ) + await db_session.flush() + + results = await manager._query_documents_most_derived(db_session, top_k=10) # pyright: ignore[reportPrivateUsage] + + contents = [doc.content for doc in results] + # Primary sort still wins: the actually-reinforced conclusion leads. + assert contents[0] == "hot" + # Ties break toward most-recent, not oldest-inserted. + assert contents[1:] == ["tie 2", "tie 1", "tie 0"] + + +class TestRepresentationManagerSessionScoping: + """Tests that the session allowlist is applied uniformly to every query path. + + Regression for DEV-1994: session_name used to be applied only to the + recent-documents query; the semantic and most-derived paths ignored it, + so limit_to_session leaked cross-session conclusions. + """ + + async def _setup( + self, + db_session: AsyncSession, + test_workspace: models.Workspace, + test_peer: models.Peer, + ) -> tuple[models.Session, models.Session, RepresentationManager]: + """Create two sessions and documents in each, plus a session-less doc.""" + test_peer2 = models.Peer( + name=str(generate_nanoid()), workspace_name=test_workspace.name + ) + db_session.add(test_peer2) + await db_session.flush() + + session_a = models.Session( + name=str(generate_nanoid()), workspace_name=test_workspace.name + ) + session_b = models.Session( + name=str(generate_nanoid()), workspace_name=test_workspace.name + ) + db_session.add_all([session_a, session_b]) + await db_session.flush() + + collection = models.Collection( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + db_session.add(collection) + await db_session.flush() + + db_session.add_all( + [ + models.Document( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + content="in-scope observation", + session_name=session_a.name, + times_derived=1, + ), + models.Document( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + content="out-of-scope observation", + session_name=session_b.name, + times_derived=100, + ), + # Dream-produced documents have no session_name; a session + # allowlist must exclude them (fail-closed). + models.Document( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + content="sessionless dream observation", + session_name=None, + times_derived=50, + ), + ] + ) + await db_session.flush() + + manager = RepresentationManager( + test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + return session_a, session_b, manager + + @pytest.mark.asyncio + async def test_recent_respects_session_allowlist( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + test_workspace, test_peer = sample_data + session_a, _, manager = await self._setup(db_session, test_workspace, test_peer) + + results = await manager._query_documents_recent( # pyright: ignore[reportPrivateUsage] + db_session, top_k=10, session_allowlist=[session_a.name] + ) + + contents = [doc.content for doc in results] + assert contents == ["in-scope observation"] + + @pytest.mark.asyncio + async def test_most_derived_respects_session_allowlist( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """The out-of-scope doc has far higher times_derived; it must still be excluded.""" + test_workspace, test_peer = sample_data + session_a, _, manager = await self._setup(db_session, test_workspace, test_peer) + + results = await manager._query_documents_most_derived( # pyright: ignore[reportPrivateUsage] + db_session, top_k=10, session_allowlist=[session_a.name] + ) + + contents = [doc.content for doc in results] + assert contents == ["in-scope observation"] + + @pytest.mark.asyncio + async def test_semantic_passes_session_allowlist_as_filters( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """The semantic path must push the allowlist down to query_documents.""" + test_workspace, test_peer = sample_data + session_a, _, manager = await self._setup(db_session, test_workspace, test_peer) + + with patch( + "src.crud.query_documents", new=AsyncMock(return_value=[]) + ) as mock_query: + await manager._query_documents_semantic( # pyright: ignore[reportPrivateUsage] + db_session, + query="anything", + top_k=5, + embedding=[0.1], + session_allowlist=[session_a.name], + ) + + assert mock_query.await_args is not None + assert mock_query.await_args.kwargs["filters"] == { + "session_name": {"in": [session_a.name]}, + # Scoped recall serves only levels with a trustworthy session + # stamp (ALLOWLIST_SAFE_LEVELS / DEV-2201). + "level": {"in": ["explicit"]}, + } + + @pytest.mark.asyncio + async def test_semantic_passes_no_filters_when_unscoped( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + test_workspace, test_peer = sample_data + _, _, manager = await self._setup(db_session, test_workspace, test_peer) + + with patch( + "src.crud.query_documents", new=AsyncMock(return_value=[]) + ) as mock_query: + await manager._query_documents_semantic( # pyright: ignore[reportPrivateUsage] + db_session, + query="anything", + top_k=5, + embedding=[0.1], + ) + + assert mock_query.await_args is not None + assert mock_query.await_args.kwargs["filters"] is None + + @pytest.mark.asyncio + async def test_working_representation_scoped_end_to_end( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """All blended paths active: only in-scope content may appear.""" + test_workspace, test_peer = sample_data + session_a, _, manager = await self._setup(db_session, test_workspace, test_peer) + + representation = await manager.get_working_representation( + db=db_session, + session_allowlist=[session_a.name], + include_most_derived=True, + ) + + contents = [obs.content for obs in representation.explicit] + assert "in-scope observation" in contents + assert "out-of-scope observation" not in contents + assert "sessionless dream observation" not in contents + + @pytest.mark.asyncio + async def test_empty_allowlist_fails_closed( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """An empty allowlist must return an empty representation, not fall + back to unscoped behavior (downstream stores drop empty IN clauses).""" + test_workspace, test_peer = sample_data + _, _, manager = await self._setup(db_session, test_workspace, test_peer) + + representation = await manager.get_working_representation( + db=db_session, + session_allowlist=[], + include_most_derived=True, + ) + + assert representation.explicit == [] + assert representation.deductive == [] + + def test_build_filter_conditions_empty_allowlist_fails_closed(self): + """The filter-builder layer itself must fail closed, independent of the + early-return guard in _get_working_representation_internal. An empty + allowlist emits an empty `in` (renders as always-false downstream), not + an omitted filter.""" + manager = RepresentationManager( + "workspace", observer="observer", observed="observed" + ) + + # Scoping also narrows to levels whose session stamp is trustworthy + # (see ALLOWLIST_SAFE_LEVELS / DEV-2201). + assert manager._build_filter_conditions(session_allowlist=[]) == { # pyright: ignore[reportPrivateUsage] + "session_name": {"in": []}, + "level": {"in": ["explicit"]}, + } + # None means unscoped — no session filter and no level narrowing. + assert manager._build_filter_conditions(session_allowlist=None) == {} # pyright: ignore[reportPrivateUsage] + assert manager._build_filter_conditions(session_allowlist=["s1"]) == { # pyright: ignore[reportPrivateUsage] + "session_name": {"in": ["s1"]}, + "level": {"in": ["explicit"]}, + } + # A requested level outside the safe set yields an empty `in`, which + # matches nothing rather than falling back to unscoped recall. + assert manager._build_filter_conditions( # pyright: ignore[reportPrivateUsage] + level="inductive", session_allowlist=["s1"] + ) == { + "session_name": {"in": ["s1"]}, + "level": {"in": []}, + } + # ...while an unscoped level filter is left exactly as asked. + assert manager._build_filter_conditions(level="inductive") == { # pyright: ignore[reportPrivateUsage] + "level": "inductive" + } + class TestRepresentationManagerSave: @pytest.mark.asyncio @@ -217,7 +506,9 @@ class TestRepresentationManagerSave: patch.object( manager, "_save_representation_internal", - new=AsyncMock(return_value=1), + new=AsyncMock( + return_value=CreateDocumentsResult(created_documents=[MagicMock()]) + ), ) as mock_save, ): saved = await manager.save_representation( @@ -228,7 +519,7 @@ class TestRepresentationManagerSave: message_level_configuration=_resolved_config(), ) - assert saved == 1 + assert len(saved.created_documents) == 1 mock_embed.assert_awaited_once_with(["useful observation"]) saved_observations = _saved_observations(mock_save) assert len(saved_observations) == 1 @@ -271,7 +562,9 @@ class TestRepresentationManagerSave: patch.object( manager, "_save_representation_internal", - new=AsyncMock(return_value=1), + new=AsyncMock( + return_value=CreateDocumentsResult(created_documents=[MagicMock()]) + ), ) as mock_save, ): saved = await manager.save_representation( @@ -282,7 +575,7 @@ class TestRepresentationManagerSave: message_level_configuration=_resolved_config(), ) - assert saved == 1 + assert len(saved.created_documents) == 1 mock_embed.assert_awaited_once_with(["inferred conclusion"]) saved_observations = _saved_observations(mock_save) assert len(saved_observations) == 1 @@ -333,6 +626,6 @@ class TestRepresentationManagerSave: message_level_configuration=_resolved_config(), ) - assert saved == 0 + assert len(saved.created_documents) == 0 mock_embed.assert_not_awaited() mock_save.assert_not_awaited() diff --git a/tests/deriver/test_deriver_processing.py b/tests/deriver/test_deriver_processing.py index a2058bf9..6785d159 100644 --- a/tests/deriver/test_deriver_processing.py +++ b/tests/deriver/test_deriver_processing.py @@ -5,11 +5,15 @@ from unittest.mock import AsyncMock, Mock, patch import pytest -from src import models +from src import crud, models from src.config import settings from src.deriver.deriver import process_representation_tasks_batch from src.llm import HonchoLLMCallResponse -from src.utils.representation import PromptRepresentation, Representation +from src.utils.representation import ( + ExplicitObservationBase, + PromptRepresentation, + Representation, +) from src.utils.work_unit import construct_work_unit_key, parse_work_unit_key @@ -316,6 +320,78 @@ class TestDeriverProcessing: for record in caplog.records ) + async def test_emits_dedup_counts_summed_across_observers(self) -> None: + """RepresentationCompletedEvent dedup counts must be the sum across all + observer collections, not the last observer's result.""" + message = Mock( + id=1, + public_id="msg_dedup", + session_name="session-1", + workspace_name="workspace-1", + peer_name="alice", + content="hello", + token_count=5, + created_at=datetime.now(timezone.utc), + ) + configuration = Mock() + configuration.reasoning.enabled = True + + mock_response = HonchoLLMCallResponse( + content=PromptRepresentation( + explicit=[ExplicitObservationBase(content="alice says hello")] + ), + input_tokens=10, + output_tokens=5, + finish_reasons=["STOP"], + ) + + manager = Mock() + manager.save_representation = AsyncMock( + side_effect=[ + crud.CreateDocumentsResult( + exact_dup_in_batch_count=1, + exact_dup_existing_count=2, + semantic_dup_rejected_count=3, + semantic_dup_replaced_count=4, + ), + crud.CreateDocumentsResult( + exact_dup_in_batch_count=10, + exact_dup_existing_count=20, + semantic_dup_rejected_count=30, + semantic_dup_replaced_count=40, + ), + ] + ) + emitted: list[Any] = [] + + with ( + patch( + "src.deriver.deriver.honcho_llm_call", + new_callable=AsyncMock, + return_value=mock_response, + ), + patch( + "src.deriver.deriver.RepresentationManager", + return_value=manager, + ), + patch("src.deriver.deriver.emit", side_effect=emitted.append), + ): + await process_representation_tasks_batch( + messages=[message], + message_level_configuration=configuration, + observers=["bob", "carol"], + observed="alice", + queue_item_message_ids=[1], + ) + + assert len(emitted) == 1 + event = emitted[0] + assert event.observer_count == 2 + assert event.exact_dup_in_batch_count == 11 + assert event.exact_dup_existing_count == 22 + assert event.semantic_dup_rejected_count == 33 + assert event.semantic_dup_replaced_count == 44 + class TestBackwardsCompatibility: """Test backwards compatibility for queue items created before the deduplication change.""" diff --git a/tests/deriver/test_embed_now.py b/tests/deriver/test_embed_now.py new file mode 100644 index 00000000..382571b8 --- /dev/null +++ b/tests/deriver/test_embed_now.py @@ -0,0 +1,352 @@ +""" +Tests for the immediate message-embedding fast path (src/reconciler/embed_now.py). + +These exercise embed_messages_now end-to-end against the test database: it opens +its own tracked_db sessions (patched to the test engine in conftest), so each test +creates committed fixture rows and asserts on the result via the provided session. +""" + +from unittest.mock import AsyncMock, patch + +import pytest +from fastapi import BackgroundTasks +from nanoid import generate as generate_nanoid +from prometheus_client import REGISTRY +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker + +from src import models +from src.config import settings +from src.reconciler.embed_now import ( + embed_messages_now, + embed_task_gate, + reset_embed_semaphore, +) +from src.vector_store import VectorStore + + +async def _create_message_with_pending_chunks( + db_session: AsyncSession, + workspace: models.Workspace, + peer: models.Peer, + chunk_contents: list[str], +) -> tuple[str, list[int]]: + """Create a message plus one pending MessageEmbedding row per chunk. + + Returns (message public_id, ordered embedding row ids). + """ + session = models.Session(name=str(generate_nanoid()), workspace_name=workspace.name) + db_session.add(session) + await db_session.commit() + + message_id = str(generate_nanoid()) + message = models.Message( + public_id=message_id, + session_name=session.name, + workspace_name=workspace.name, + peer_name=peer.name, + content=" ".join(chunk_contents), + seq_in_session=1, + ) + db_session.add(message) + await db_session.commit() + + rows = [ + models.MessageEmbedding( + content=chunk, + message_id=message_id, + workspace_name=workspace.name, + session_name=session.name, + peer_name=peer.name, + sync_state="pending", + embedding=None, + ) + for chunk in chunk_contents + ] + db_session.add_all(rows) + await db_session.commit() + for row in rows: + await db_session.refresh(row) + return message_id, [row.id for row in rows] + + +@pytest.fixture(autouse=True) +def reset_semaphore_fixture(): + """Rebuild the module semaphore per test so it binds to the active loop, + and clear the admission gate's in-flight count.""" + reset_embed_semaphore() + embed_task_gate.in_flight = 0 + yield + reset_embed_semaphore() + embed_task_gate.in_flight = 0 + + +@pytest.mark.asyncio +class TestEmbedTaskGate: + """Admission gate for immediate-embed background tasks + (EMBEDDING.MAX_PENDING_EMBED_TASKS).""" + + async def test_admits_under_cap_and_releases_slot(self) -> None: + """Under the cap, the task is scheduled; running it embeds the given ids + and releases the slot.""" + tasks = BackgroundTasks() + with ( + patch.object(settings.EMBEDDING, "MAX_PENDING_EMBED_TASKS", 2), + patch( + "src.reconciler.embed_now.embed_messages_now", new=AsyncMock() + ) as mock_embed, + ): + assert embed_task_gate.try_schedule(tasks, ["msg_1"]) is True + assert embed_task_gate.in_flight == 1 + await tasks() + mock_embed.assert_awaited_once_with(["msg_1"]) + assert embed_task_gate.in_flight == 0 + + async def test_rejects_at_cap_without_scheduling(self) -> None: + """At the cap, nothing is scheduled and False is returned.""" + tasks = BackgroundTasks() + with patch.object(settings.EMBEDDING, "MAX_PENDING_EMBED_TASKS", 1): + assert embed_task_gate.try_schedule(tasks, ["msg_1"]) is True + assert embed_task_gate.try_schedule(tasks, ["msg_2"]) is False + assert len(tasks.tasks) == 1 + assert embed_task_gate.in_flight == 1 + + async def test_slot_released_when_task_raises(self) -> None: + """A failing task still releases its slot (finally path).""" + tasks = BackgroundTasks() + with ( + patch.object(settings.EMBEDDING, "MAX_PENDING_EMBED_TASKS", 1), + patch( + "src.reconciler.embed_now.embed_messages_now", + new=AsyncMock(side_effect=RuntimeError("boom")), + ), + ): + assert embed_task_gate.try_schedule(tasks, ["msg_1"]) is True + with pytest.raises(RuntimeError): + await tasks() + assert embed_task_gate.in_flight == 0 + + async def test_gauge_mirrors_in_flight_count(self) -> None: + """With metrics enabled, the Prometheus gauge tracks the gate's + in-flight count through schedule and release.""" + + def gauge_value() -> float | None: + return REGISTRY.get_sample_value( + "embed_now_tasks_in_flight", {"namespace": "test"} + ) + + tasks = BackgroundTasks() + with ( + patch.object(settings.EMBEDDING, "MAX_PENDING_EMBED_TASKS", 2), + patch.object(settings.METRICS, "ENABLED", True), + patch.object(settings.METRICS, "NAMESPACE", "test"), + patch("src.reconciler.embed_now.embed_messages_now", new=AsyncMock()), + ): + assert embed_task_gate.try_schedule(tasks, ["msg_1"]) is True + assert gauge_value() == 1 + await tasks() + assert gauge_value() == 0 + + async def test_zero_cap_disables_fast_path(self) -> None: + """MAX_PENDING_EMBED_TASKS=0 rejects every schedule attempt.""" + tasks = BackgroundTasks() + with patch.object(settings.EMBEDDING, "MAX_PENDING_EMBED_TASKS", 0): + assert embed_task_gate.try_schedule(tasks, ["msg_1"]) is False + assert len(tasks.tasks) == 0 + + +@pytest.mark.asyncio +class TestEmbedMessagesNow: + async def test_pgvector_happy_path_marks_synced( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ) -> None: + """pgvector-only mode: rows get a vector and flip to synced immediately.""" + workspace, peer = sample_data + message_id, emb_ids = await _create_message_with_pending_chunks( + db_session, workspace, peer, ["hello world"] + ) + + await embed_messages_now([message_id]) + + for emb_id in emb_ids: + row = await db_session.get(models.MessageEmbedding, emb_id) + assert row is not None + await db_session.refresh(row) + assert row.sync_state == "synced" + assert row.embedding is not None + assert row.sync_attempts == 0 + + async def test_no_message_ids_is_noop(self) -> None: + """Empty input returns without touching the DB or embedding.""" + with patch( + "src.embedding_client.embedding_client.simple_batch_embed" + ) as mock_embed: + await embed_messages_now([]) + mock_embed.assert_not_called() + + async def test_already_synced_rows_not_reclaimed( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ) -> None: + """A second run finds no pending rows and does not re-embed.""" + workspace, peer = sample_data + message_id, _ = await _create_message_with_pending_chunks( + db_session, workspace, peer, ["first content"] + ) + await embed_messages_now([message_id]) + + with patch( + "src.embedding_client.embedding_client.simple_batch_embed" + ) as mock_embed: + await embed_messages_now([message_id]) + mock_embed.assert_not_called() + + async def test_embed_failure_leaves_rows_pending_and_leased( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ) -> None: + """Embedding failure must leave rows pending + leased, attempts untouched, + so the reconciler owns retry accounting.""" + workspace, peer = sample_data + message_id, emb_ids = await _create_message_with_pending_chunks( + db_session, workspace, peer, ["will fail"] + ) + + with patch( + "src.embedding_client.embedding_client.simple_batch_embed", + new=AsyncMock(side_effect=RuntimeError("provider down")), + ): + await embed_messages_now([message_id]) + + for emb_id in emb_ids: + row = await db_session.get(models.MessageEmbedding, emb_id) + assert row is not None + await db_session.refresh(row) + assert row.sync_state == "pending" + assert row.embedding is None + assert row.sync_attempts == 0 # lease only, no attempt bump + assert row.last_sync_at is not None # leased + + async def test_external_store_upserts_with_chunk_positioned_ids( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + mock_vector_store: VectorStore, + ) -> None: + """External-store mode: upsert each chunk with id {message_id}_{position} + and mark rows synced.""" + workspace, peer = sample_data + message_id, emb_ids = await _create_message_with_pending_chunks( + db_session, workspace, peer, ["chunk a", "chunk b", "chunk c"] + ) + + with patch( + "src.reconciler.embed_now.get_external_vector_store", + return_value=mock_vector_store, + ): + await embed_messages_now([message_id]) + + upsert_mock: AsyncMock = mock_vector_store.upsert_many # pyright: ignore[reportAssignmentType] + upsert_mock.assert_awaited() + upserted_ids = { + record.id for call in upsert_mock.await_args_list for record in call.args[1] + } + assert upserted_ids == { + f"{message_id}_0", + f"{message_id}_1", + f"{message_id}_2", + } + + for emb_id in emb_ids: + row = await db_session.get(models.MessageEmbedding, emb_id) + assert row is not None + await db_session.refresh(row) + assert row.sync_state == "synced" + + async def test_locked_chunk_skipped_keeps_positions_stable( + self, + db_session: AsyncSession, + db_engine: AsyncEngine, + sample_data: tuple[models.Workspace, models.Peer], + mock_vector_store: VectorStore, + ) -> None: + """If a sibling chunk is locked by another txn, SKIP LOCKED skips it but + chunk positions still come from the full sibling ordering — so the claimed + chunks keep their {message_id}_0 / _2 ids (not _0 / _1).""" + workspace, peer = sample_data + message_id, emb_ids = await _create_message_with_pending_chunks( + db_session, workspace, peer, ["chunk a", "chunk b", "chunk c"] + ) + locked_id = emb_ids[1] # middle chunk -> position 1 + + # Hold a row lock on the middle chunk from an independent transaction. + lock_factory = async_sessionmaker(bind=db_engine, expire_on_commit=False) + lock_session = lock_factory() + await lock_session.execute( + select(models.MessageEmbedding) + .where(models.MessageEmbedding.id == locked_id) + .with_for_update() + ) + try: + with patch( + "src.reconciler.embed_now.get_external_vector_store", + return_value=mock_vector_store, + ): + await embed_messages_now([message_id]) + finally: + await lock_session.rollback() + await lock_session.close() + + upsert_mock: AsyncMock = mock_vector_store.upsert_many # pyright: ignore[reportAssignmentType] + upserted_ids = { + record.id for call in upsert_mock.await_args_list for record in call.args[1] + } + assert upserted_ids == {f"{message_id}_0", f"{message_id}_2"} + + # The locked chunk stays pending; the other two are synced. + locked_row = await db_session.get(models.MessageEmbedding, locked_id) + assert locked_row is not None + await db_session.refresh(locked_row) + assert locked_row.sync_state == "pending" + for emb_id in (emb_ids[0], emb_ids[2]): + row = await db_session.get(models.MessageEmbedding, emb_id) + assert row is not None + await db_session.refresh(row) + assert row.sync_state == "synced" + + async def test_external_store_unavailable_leaves_rows_pending( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + mock_vector_store: VectorStore, + ) -> None: + """External-store mode: if upsert_many raises VectorStoreError, rows must + stay pending with no vector and untouched attempts, so the reconciler + heals them. embed_now never bumps sync_attempts.""" + from src.exceptions import VectorStoreError + + workspace, peer = sample_data + message_id, emb_ids = await _create_message_with_pending_chunks( + db_session, workspace, peer, ["chunk a", "chunk b"] + ) + + upsert_mock: AsyncMock = mock_vector_store.upsert_many # pyright: ignore[reportAssignmentType] + upsert_mock.side_effect = VectorStoreError("vector store down") + + with patch( + "src.reconciler.embed_now.get_external_vector_store", + return_value=mock_vector_store, + ): + await embed_messages_now([message_id]) + + for emb_id in emb_ids: + row = await db_session.get(models.MessageEmbedding, emb_id) + assert row is not None + await db_session.refresh(row) + assert row.sync_state == "pending" + assert row.embedding is None + assert row.sync_attempts == 0 # embed_now never bumps attempts diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py index c293f52c..81dd0632 100644 --- a/tests/deriver/test_queue_processing.py +++ b/tests/deriver/test_queue_processing.py @@ -1,4 +1,6 @@ +import asyncio from collections.abc import Callable +from datetime import datetime, timedelta, timezone from typing import Any from unittest.mock import patch @@ -17,6 +19,71 @@ from src.utils.work_unit import construct_work_unit_key class TestQueueProcessing: """Test suite for queue processing functionality""" + async def _add_representation_work_unit( + self, + *, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + token_counts: list[int], + created_ats: list[datetime] | None = None, + ) -> tuple[str, list[models.QueueItem]]: + session, peers = sample_session_with_peers + peer = peers[0] + + messages: list[models.Message] = [] + for index, token_count in enumerate(token_counts): + message = models.Message( + session_name=session.name, + workspace_name=session.workspace_name, + peer_name=peer.name, + content=f"Message {index}", + token_count=token_count, + seq_in_session=index + 1, + ) + db_session.add(message) + messages.append(message) + + await db_session.commit() + for message in messages: + await db_session.refresh(message) + + work_unit_key = "" + queue_items: list[models.QueueItem] = [] + for index, message in enumerate(messages): + payload = create_queue_payload( + message=message, + task_type="representation", + observed=peer.name, + observer=peer.name, + ) + work_unit_key = work_unit_key or construct_work_unit_key( + session.workspace_name, payload + ) + + queue_item_kwargs: dict[str, Any] = {} + if created_ats: + queue_item_kwargs["created_at"] = created_ats[index] + + queue_item = models.QueueItem( + session_id=session.id, + task_type="representation", + work_unit_key=work_unit_key, + payload=payload, + processed=False, + workspace_name=session.workspace_name, + message_id=message.id, + **queue_item_kwargs, + ) + db_session.add(queue_item) + queue_items.append(queue_item) + + await db_session.commit() + for queue_item in queue_items: + await db_session.refresh(queue_item) + + return work_unit_key, queue_items + async def test_get_and_claim_work_units( self, db_session: AsyncSession, @@ -292,7 +359,7 @@ class TestQueueProcessing: peer = peers[0] # Create messages with token counts that exceed batch limit - limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + limit = settings.DERIVER.REPRESENTATION_BATCH_TARGET_INPUT_TOKENS token_counts = [limit // 2, limit // 2, limit // 2] # Create and save messages to the database first @@ -412,7 +479,7 @@ class TestQueueProcessing: session, peers = sample_session_with_peers peer = peers[0] - cap = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + cap = settings.DERIVER.REPRESENTATION_BATCH_TARGET_INPUT_TOKENS # M1 + M2 sum to exactly the cap; M3 pushes over it. After SQL, # messages_context = [M1, M2]; the cap is genuinely binding *on the @@ -514,7 +581,7 @@ class TestQueueProcessing: session, peers = sample_session_with_peers peer_a = peers[0] peer_b = peers[1] if len(peers) > 1 else peers[0] - cap = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + cap = settings.DERIVER.REPRESENTATION_BATCH_TARGET_INPUT_TOKENS # Layout: 4 messages, ordered. # M1 (peer_a, queue, 200) @@ -677,7 +744,9 @@ class TestQueueProcessing: qm = QueueManager() # Mock the token limit to 2000 for this test - with patch.object(settings.DERIVER, "REPRESENTATION_BATCH_MAX_TOKENS", 2000): + with patch.object( + settings.DERIVER, "REPRESENTATION_BATCH_TARGET_INPUT_TOKENS", 2000 + ): # Test alice's work unit alice_work_unit_key = alice_queue_items[0].work_unit_key alice_aqs = models.ActiveQueueSession(work_unit_key=alice_work_unit_key) @@ -852,7 +921,9 @@ class TestQueueProcessing: qm = QueueManager() # Mock the token limit to 1500 for this test - with patch.object(settings.DERIVER, "REPRESENTATION_BATCH_MAX_TOKENS", 1500): + with patch.object( + settings.DERIVER, "REPRESENTATION_BATCH_TARGET_INPUT_TOKENS", 1500 + ): # Test alice's work unit # With per-work-unit anchoring + preceding context: # Alice starts at message 3, includes preceding message 2 (steve) for context @@ -1066,7 +1137,7 @@ class TestQueueProcessing: peer = peers[0] # Create messages where first message exceeds the batch limit - limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + limit = settings.DERIVER.REPRESENTATION_BATCH_TARGET_INPUT_TOKENS token_counts = [limit + 1000, 100, 200] # First message way over limit # Create and save messages to the database first @@ -1183,7 +1254,7 @@ class TestQueueProcessing: peer = peers[0] # Create messages that test the exact boundary - limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + limit = settings.DERIVER.REPRESENTATION_BATCH_TARGET_INPUT_TOKENS token_counts = [ limit // 2, limit // 2, @@ -1314,7 +1385,7 @@ class TestQueueProcessing: peer = peers[0] # Create messages with tokens BELOW the threshold - limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + limit = settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS token_counts = [100, 100, 100] # Total 300, way below 4096 messages: list[models.Message] = [] @@ -1412,6 +1483,167 @@ class TestQueueProcessing: claimed2 = await qm.get_and_claim_work_units() assert rep_work_unit_key in claimed2 + @pytest.mark.asyncio + async def test_age_flush_waits_for_fresh_sub_threshold_items( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(settings.DERIVER, "FLUSH_ENABLED", False) + monkeypatch.setattr( + settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800 + ) + + work_unit_key, _queue_items = await self._add_representation_work_unit( + db_session=db_session, + sample_session_with_peers=sample_session_with_peers, + create_queue_payload=create_queue_payload, + token_counts=[100, 100, 100], + ) + + claimed = await QueueManager().get_and_claim_work_units() + + assert work_unit_key not in claimed + + @pytest.mark.asyncio + async def test_age_flush_claims_old_sub_threshold_items_and_fetches_tail( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(settings.DERIVER, "FLUSH_ENABLED", False) + monkeypatch.setattr( + settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800 + ) + old_timestamp = datetime.now(timezone.utc) - timedelta(hours=2) + + work_unit_key, queue_items = await self._add_representation_work_unit( + db_session=db_session, + sample_session_with_peers=sample_session_with_peers, + create_queue_payload=create_queue_payload, + token_counts=[100, 100, 100], + created_ats=[old_timestamp, old_timestamp, old_timestamp], + ) + + qm = QueueManager() + claimed = await qm.get_and_claim_work_units() + + assert work_unit_key in claimed + batch = await qm.get_queue_item_batch( + task_type="representation", + work_unit_key=work_unit_key, + aqs_id=claimed[work_unit_key], + ) + assert [item.id for item in batch.items_to_process] == [ + item.id for item in queue_items + ] + + @pytest.mark.asyncio + async def test_age_flush_zero_preserves_legacy_wait_for_old_items( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(settings.DERIVER, "FLUSH_ENABLED", False) + monkeypatch.setattr(settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 0) + old_timestamp = datetime.now(timezone.utc) - timedelta(hours=2) + + work_unit_key, _queue_items = await self._add_representation_work_unit( + db_session=db_session, + sample_session_with_peers=sample_session_with_peers, + create_queue_payload=create_queue_payload, + token_counts=[100, 100], + created_ats=[old_timestamp, old_timestamp], + ) + + claimed = await QueueManager().get_and_claim_work_units() + + assert work_unit_key not in claimed + + @pytest.mark.asyncio + async def test_flush_enabled_bypasses_age_and_token_thresholds( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(settings.DERIVER, "FLUSH_ENABLED", True) + monkeypatch.setattr( + settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800 + ) + + work_unit_key, _queue_items = await self._add_representation_work_unit( + db_session=db_session, + sample_session_with_peers=sample_session_with_peers, + create_queue_payload=create_queue_payload, + token_counts=[100, 100], + ) + + claimed = await QueueManager().get_and_claim_work_units() + + assert work_unit_key in claimed + + @pytest.mark.asyncio + async def test_age_flush_uses_oldest_unprocessed_item( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(settings.DERIVER, "FLUSH_ENABLED", False) + monkeypatch.setattr( + settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800 + ) + now = datetime.now(timezone.utc) + + work_unit_key, _queue_items = await self._add_representation_work_unit( + db_session=db_session, + sample_session_with_peers=sample_session_with_peers, + create_queue_payload=create_queue_payload, + token_counts=[100, 100], + created_ats=[now - timedelta(hours=2), now], + ) + + claimed = await QueueManager().get_and_claim_work_units() + + assert work_unit_key in claimed + + @pytest.mark.asyncio + async def test_age_flush_ignores_old_processed_items( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(settings.DERIVER, "FLUSH_ENABLED", False) + monkeypatch.setattr( + settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800 + ) + now = datetime.now(timezone.utc) + + work_unit_key, queue_items = await self._add_representation_work_unit( + db_session=db_session, + sample_session_with_peers=sample_session_with_peers, + create_queue_payload=create_queue_payload, + token_counts=[100, 100], + created_ats=[now - timedelta(hours=2), now], + ) + queue_items[0].processed = True + await db_session.commit() + + claimed = await QueueManager().get_and_claim_work_units() + + assert work_unit_key not in claimed + @pytest.mark.asyncio async def test_forced_batching_single_large_message( self, @@ -1424,7 +1656,7 @@ class TestQueueProcessing: session, peers = sample_session_with_peers peer = peers[0] - limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + limit = settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS # Create a single message that exceeds the threshold message = models.Message( @@ -1530,7 +1762,7 @@ class TestQueueProcessing: session, peers = sample_session_with_peers peer = peers[0] - limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + limit = settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS # Create messages that sum to exactly the threshold token_counts = [limit // 2, limit // 2] @@ -1583,3 +1815,62 @@ class TestQueueProcessing: claimed = await qm.get_and_claim_work_units() assert work_unit_key is not None assert work_unit_key in claimed + + +class TestPollingJitter: + """Polling jitter: desynchronize poll loops without changing the schedule.""" + + def test_jitter_stays_within_ratio_bounds( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.5) + qm = QueueManager() + samples = [qm._jitter(10.0) for _ in range(1000)] # pyright: ignore[reportPrivateUsage] + assert all(5.0 <= s <= 15.0 for s in samples) + # A 0.5 ratio over 1000 samples should produce real spread, not a constant. + assert max(samples) - min(samples) > 1.0 + + def test_jitter_ratio_zero_is_identity( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) + qm = QueueManager() + assert all(qm._jitter(7.5) == 7.5 for _ in range(50)) # pyright: ignore[reportPrivateUsage] + + def test_advance_jitters_return_but_keeps_deterministic_schedule( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.5) + monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_ENABLED", True) + monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_MULTIPLIER", 2.0) + monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0) + monkeypatch.setattr( + settings.DERIVER, "POLLING_SLEEP_MAX_INTERVAL_SECONDS", 30.0 + ) + qm = QueueManager() + + # The underlying schedule advances 1 -> 2 -> 4 -> ... -> 30 deterministically; + # each returned sleep is jittered within [0.5x, 1.5x] of the pre-advance step. + expected_schedule = [1.0, 2.0, 4.0, 8.0, 16.0, 30.0, 30.0] + for step in expected_schedule: + returned = qm._advance_poll_interval() # pyright: ignore[reportPrivateUsage] + assert 0.5 * step <= returned <= 1.5 * step + + @pytest.mark.asyncio + async def test_startup_jitter_disabled_returns_immediately( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(settings.DERIVER, "POLLING_STARTUP_JITTER_SECONDS", 0.0) + qm = QueueManager() + # Window 0.0 must not sleep at all. + await qm._sleep_startup_jitter() # pyright: ignore[reportPrivateUsage] + + @pytest.mark.asyncio + async def test_startup_jitter_interrupted_by_shutdown( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(settings.DERIVER, "POLLING_STARTUP_JITTER_SECONDS", 300.0) + qm = QueueManager() + qm.shutdown_event.set() + # A shutdown already signalled must short-circuit the (long) jitter sleep. + await asyncio.wait_for(qm._sleep_startup_jitter(), timeout=1.0) # pyright: ignore[reportPrivateUsage] diff --git a/tests/deriver/test_vector_reconciliation.py b/tests/deriver/test_vector_reconciliation.py index 748b84df..5e2e4712 100644 --- a/tests/deriver/test_vector_reconciliation.py +++ b/tests/deriver/test_vector_reconciliation.py @@ -6,6 +6,7 @@ message embeddings to the vector store, handling failures and retries. """ import datetime +from contextlib import asynccontextmanager from typing import cast from unittest.mock import AsyncMock, MagicMock, patch @@ -20,9 +21,12 @@ from src.reconciler.sync_vectors import ( ReconciliationMetrics, _get_documents_needing_sync, # pyright: ignore[reportPrivateUsage] _get_message_embeddings_needing_sync, # pyright: ignore[reportPrivateUsage] + _reconcile_documents_batch, # pyright: ignore[reportPrivateUsage] _reconcile_message_embeddings_batch, # pyright: ignore[reportPrivateUsage] _sync_documents, # pyright: ignore[reportPrivateUsage] _sync_message_embeddings, # pyright: ignore[reportPrivateUsage] + build_message_vector_record, + compute_chunk_positions, run_vector_reconciliation_cycle, ) from src.vector_store import ( @@ -873,6 +877,82 @@ class TestMessageEmbeddings: assert pending_emb.sync_attempts == 0 assert pending_emb.last_sync_at is None + async def test_pgvector_only_mode_embeds_and_marks_synced( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ) -> None: + """In pgvector-only mode, the reconciler must still embed pending rows.""" + workspace, peer = sample_data + pending_emb = await self._create_pending_message_embedding( + db_session, workspace, peer + ) + + # external_vector_store=None == pgvector-only mode. The reconciler should + # re-embed the pending row, write the vector to postgres, and mark synced. + synced, failed = await _sync_message_embeddings(db_session, [pending_emb], None) + + await db_session.commit() + await db_session.refresh(pending_emb) + + assert synced == 1 + assert failed == 0 + assert pending_emb.sync_state == "synced" + assert pending_emb.sync_attempts == 0 + assert pending_emb.embedding is not None + + async def test_all_chunks_of_a_message_claimed_together( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ) -> None: + """A single message's chunks must always be claimed in one batch. + + Selecting by message_id (not row) keeps `{message_id}_{chunk_index}` + vector IDs stable across reconciler cycles. + """ + workspace, peer = sample_data + + # Create one message with 5 chunks. + session = models.Session( + name=str(generate_nanoid()), workspace_name=workspace.name + ) + db_session.add(session) + await db_session.commit() + + message_id = str(generate_nanoid()) + message = models.Message( + public_id=message_id, + session_name=session.name, + workspace_name=workspace.name, + peer_name=peer.name, + content="full message content", + seq_in_session=1, + ) + db_session.add(message) + await db_session.commit() + + chunk_count = 5 + for i in range(chunk_count): + db_session.add( + models.MessageEmbedding( + content=f"chunk-{i}", + message_id=message_id, + workspace_name=workspace.name, + session_name=session.name, + peer_name=peer.name, + sync_state="pending", + embedding=None, + ) + ) + await db_session.commit() + + # Even with batch_size=1, all 5 chunks for the message should be claimed + # together because the query selects by distinct message_id first. + claimed = await _get_message_embeddings_needing_sync(db_session, batch_size=1) + assert len(claimed) == chunk_count + assert all(emb.message_id == message_id for emb in claimed) + @pytest.mark.asyncio class TestEndToEndReconciliation: @@ -913,3 +993,176 @@ class TestEndToEndReconciliation: mock_reconcile_docs.assert_awaited_once() mock_reconcile_embs.assert_awaited_once() mock_cleanup_docs.assert_awaited_once() + + +@pytest.mark.asyncio +class TestReconcilerTracing: + """A Sentry transaction is started only when a sync batch finds real work. + + Reconciler tasks poll on a fixed interval and usually find nothing; an idle + cycle must create zero transactions so it doesn't drain Sentry quota. + """ + + @staticmethod + def _fake_tracked_db(db: AsyncMock): + @asynccontextmanager + async def _cm(*_args: object, **_kwargs: object): + yield db + + return _cm + + async def test_no_transaction_when_no_embeddings_to_sync(self) -> None: + """The no-work path returns before starting a transaction.""" + metrics = ReconciliationMetrics() + with ( + patch( + "src.reconciler.sync_vectors.tracked_db", + self._fake_tracked_db(AsyncMock()), + ), + patch( + "src.reconciler.sync_vectors._get_message_embeddings_needing_sync", + new_callable=AsyncMock, + return_value=[], + ), + patch("src.reconciler.sync_vectors.sentry_sdk.start_transaction") as txn, + ): + worked = await _reconcile_message_embeddings_batch(None, metrics) + + assert worked is False + txn.assert_not_called() + + async def test_transaction_started_when_embeddings_present(self) -> None: + """A batch with real work starts its own named transaction.""" + metrics = ReconciliationMetrics() + with ( + patch( + "src.reconciler.sync_vectors.tracked_db", + self._fake_tracked_db(AsyncMock()), + ), + patch( + "src.reconciler.sync_vectors._get_message_embeddings_needing_sync", + new_callable=AsyncMock, + return_value=[MagicMock()], + ), + patch( + "src.reconciler.sync_vectors._sync_message_embeddings", + new_callable=AsyncMock, + return_value=(1, 0), + ), + patch("src.reconciler.sync_vectors.sentry_sdk.start_transaction") as txn, + ): + worked = await _reconcile_message_embeddings_batch(None, metrics) + + assert worked is True + assert metrics.message_embeddings_synced == 1 + txn.assert_called_once() + assert txn.call_args.kwargs.get("name") == "reconcile_message_embeddings_batch" + + async def test_no_transaction_when_no_documents_to_sync(self) -> None: + """The document batch also skips tracing when there is nothing to sync.""" + metrics = ReconciliationMetrics() + with ( + patch( + "src.reconciler.sync_vectors.tracked_db", + self._fake_tracked_db(AsyncMock()), + ), + patch( + "src.reconciler.sync_vectors._get_documents_needing_sync", + new_callable=AsyncMock, + return_value=[], + ), + patch("src.reconciler.sync_vectors.sentry_sdk.start_transaction") as txn, + ): + worked = await _reconcile_documents_batch( + MagicMock(spec=VectorStore), metrics + ) + + assert worked is False + txn.assert_not_called() + + +def test_build_message_vector_record() -> None: + """The shared vector-id/metadata builder: id is {message_id}_{position}, + embeddings are coerced to float, metadata shape is fixed.""" + record = build_message_vector_record( + message_id="msg_abc", + chunk_position=2, + session_name="sess", + peer_name="peer", + embedding=[1, 2, 3], # ints, must be coerced + ) + assert record.id == "msg_abc_2" + assert record.embedding == [1.0, 2.0, 3.0] + assert all(isinstance(x, float) for x in record.embedding) + assert record.metadata == { + "message_id": "msg_abc", + "session_name": "sess", + "peer_name": "peer", + } + + +@pytest.mark.asyncio +class TestComputeChunkPositions: + """Direct coverage for compute_chunk_positions, the source of truth for + {message_id}_{position} vector ids shared by the reconciler and embed_now.""" + + async def test_empty_input_returns_empty(self, db_session: AsyncSession) -> None: + assert await compute_chunk_positions(db_session, []) == {} + + async def test_positions_are_per_message_zero_indexed( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ) -> None: + """Each message's rows are numbered from 0 in (message_id, id) order, + independent of how rows from other messages interleave.""" + workspace, peer = sample_data + session = models.Session( + name=str(generate_nanoid()), workspace_name=workspace.name + ) + db_session.add(session) + await db_session.commit() + + # msg_a has 2 chunks, msg_b has 1 chunk. + msg_a = str(generate_nanoid()) + msg_b = str(generate_nanoid()) + for seq, mid in enumerate((msg_a, msg_b), start=1): + db_session.add( + models.Message( + public_id=mid, + session_name=session.name, + workspace_name=workspace.name, + peer_name=peer.name, + content="content", + seq_in_session=seq, + ) + ) + await db_session.commit() + + rows = [ + models.MessageEmbedding( + content=content, + message_id=mid, + workspace_name=workspace.name, + session_name=session.name, + peer_name=peer.name, + sync_state="pending", + embedding=None, + ) + for mid, content in ( + (msg_a, "a0"), + (msg_a, "a1"), + (msg_b, "b0"), + ) + ] + db_session.add_all(rows) + await db_session.commit() + for row in rows: + await db_session.refresh(row) + a0, a1, b0 = (row.id for row in rows) + + positions = await compute_chunk_positions(db_session, [msg_a, msg_b]) + + assert positions[a0] == 0 + assert positions[a1] == 1 + assert positions[b0] == 0 diff --git a/tests/dialectic/test_structured_output.py b/tests/dialectic/test_structured_output.py new file mode 100644 index 00000000..10b32da8 --- /dev/null +++ b/tests/dialectic/test_structured_output.py @@ -0,0 +1,126 @@ +"""Tests for response_model threading through the DialecticAgent.""" + +import time +from unittest.mock import AsyncMock, patch + +import pytest +from pydantic import BaseModel + +from src.dialectic.core import DialecticAgent +from src.llm import ( + HonchoLLMCallResponse, + HonchoLLMCallStreamChunk, + StreamingResponseWithMetadata, +) + + +class FoodPreferences(BaseModel): + favorite: str + confidence: float + + +def _make_agent() -> DialecticAgent: + return DialecticAgent( + workspace_name="workspace", + session_name="session", + observer="observer", + observed="observed", + reasoning_level="low", + ) + + +def _patches(mock_llm_call: AsyncMock): + return ( + patch.object( + DialecticAgent, + "_prepare_query", + new=AsyncMock( + return_value=(AsyncMock(), "task", "run", time.perf_counter()) + ), + ), + patch.object(DialecticAgent, "_log_response_metrics"), + patch("src.dialectic.core.honcho_llm_call", new=mock_llm_call), + ) + + +@pytest.mark.asyncio +async def test_answer_passes_response_model_and_serializes() -> None: + """answer() threads response_model to the LLM call and serializes the + parsed model instance back to a JSON string.""" + agent = _make_agent() + parsed = FoodPreferences(favorite="sushi", confidence=0.9) + mock_llm_call = AsyncMock( + return_value=HonchoLLMCallResponse( + content=parsed, + input_tokens=10, + output_tokens=5, + finish_reasons=["stop"], + ) + ) + + p1, p2, p3 = _patches(mock_llm_call) + with p1, p2, p3: + result = await agent.answer("query", response_model=FoodPreferences) + + kwargs = mock_llm_call.await_args.kwargs # pyright: ignore + assert kwargs["response_model"] is FoodPreferences + assert isinstance(result, str) + assert FoodPreferences.model_validate_json(result) == parsed + + +@pytest.mark.asyncio +async def test_answer_without_response_model_returns_plain_text() -> None: + agent = _make_agent() + mock_llm_call = AsyncMock( + return_value=HonchoLLMCallResponse( + content="plain answer", + input_tokens=10, + output_tokens=5, + finish_reasons=["stop"], + ) + ) + + p1, p2, p3 = _patches(mock_llm_call) + with p1, p2, p3: + result = await agent.answer("query") + + assert result == "plain answer" + assert mock_llm_call.await_args.kwargs["response_model"] is None # pyright: ignore + + +@pytest.mark.asyncio +async def test_answer_stream_passes_response_model() -> None: + """answer_stream() threads response_model; chunks stay raw text.""" + agent = _make_agent() + + async def _stream(): + yield HonchoLLMCallStreamChunk(content='{"favorite":"sushi",') + yield HonchoLLMCallStreamChunk(content='"confidence":0.9}') + yield HonchoLLMCallStreamChunk(content="", is_done=True) + + mock_llm_call = AsyncMock( + return_value=StreamingResponseWithMetadata( + _stream(), + tool_calls_made=[], + input_tokens=10, + output_tokens=5, + cache_creation_input_tokens=0, + cache_read_input_tokens=0, + iterations=1, + ) + ) + + p1, p2, p3 = _patches(mock_llm_call) + with p1, p2, p3: + chunks = [ + chunk + async for chunk in agent.answer_stream( + "query", response_model=FoodPreferences + ) + ] + + kwargs = mock_llm_call.await_args.kwargs # pyright: ignore + assert kwargs["response_model"] is FoodPreferences + assert kwargs["stream_final_only"] is True + accumulated = "".join(chunks) + assert FoodPreferences.model_validate_json(accumulated).favorite == "sushi" diff --git a/tests/dreamer/test_card_refresh.py b/tests/dreamer/test_card_refresh.py new file mode 100644 index 00000000..43d1edea --- /dev/null +++ b/tests/dreamer/test_card_refresh.py @@ -0,0 +1,335 @@ +"""Tests for the card_refresh dream type (DEV-2000, Scopes RFC prerequisite). + +Covers: +- queue plumbing: payload roundtrip, work-unit key isolation from omni, + enqueue alongside a pending omni dream +- process_dream dispatch of DreamType.CARD_REFRESH (and that it does NOT + advance the omni dream guard pair) +- specialist tool restriction (no observation-mutating tools) +- the low tool-iteration cap +- rebuild mode omitting the prior peer card from the prompt +""" + +from typing import Any +from unittest.mock import AsyncMock, patch + +import pytest +import pytest_asyncio +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from src import models +from src.config import settings +from src.deriver.enqueue import enqueue_dream +from src.dreamer.orchestrator import DreamResult, process_dream +from src.dreamer.specialists import CardRefreshSpecialist +from src.llm import HonchoLLMCallResponse +from src.schemas import DreamType +from src.utils.queue_payload import DreamPayload, create_dream_payload +from src.utils.work_unit import construct_work_unit_key, parse_work_unit_key + +OBSERVATION_MUTATION_TOOLS = { + "create_observations", + "create_observations_deductive", + "create_observations_inductive", + "delete_observations", +} + + +@pytest_asyncio.fixture +async def seeded_collection( + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], +) -> models.Collection: + """Create a Collection with an empty dream metadata dict.""" + workspace, peer = sample_data + collection = models.Collection( + observer=peer.name, + observed=peer.name, + workspace_name=workspace.name, + internal_metadata={}, + ) + db_session.add(collection) + await db_session.commit() + await db_session.refresh(collection) + return collection + + +def _make_card_refresh_result() -> DreamResult: + return DreamResult( + run_id="test_run_card", + specialists_run=["card_refresh"], + deduction_success=True, + induction_success=False, + surprisal_enabled=False, + surprisal_conclusion_count=0, + total_iterations=2, + total_duration_ms=42.0, + input_tokens=10, + output_tokens=5, + ) + + +class TestQueuePlumbing: + def test_payload_roundtrip_carries_rebuild(self): + payload_dict = create_dream_payload( + DreamType.CARD_REFRESH, + observer="alice", + observed="bob", + rebuild=True, + ) + validated = DreamPayload(**payload_dict) + assert validated.dream_type == DreamType.CARD_REFRESH + assert validated.rebuild is True + + # Default is False, including for older payloads missing the field. + assert ( + DreamPayload(dream_type=DreamType.OMNI, observer="a", observed="b").rebuild + is False + ) + + def test_work_unit_key_does_not_collide_with_omni(self): + base = {"task_type": "dream", "observer": "alice", "observed": "bob"} + omni_key = construct_work_unit_key("ws", {**base, "dream_type": "omni"}) + card_key = construct_work_unit_key("ws", {**base, "dream_type": "card_refresh"}) + + assert omni_key != card_key + parsed = parse_work_unit_key(card_key) + assert parsed.task_type == "dream" + assert parsed.dream_type == "card_refresh" + assert parsed.observer == "alice" + assert parsed.observed == "bob" + + @pytest.mark.asyncio + async def test_enqueue_alongside_pending_omni( + self, + db_session: AsyncSession, + seeded_collection: models.Collection, + ): + """A pending omni dream must not dedupe away a card_refresh enqueue — + the work-unit keys differ by dream type.""" + await enqueue_dream( + seeded_collection.workspace_name, + observer=seeded_collection.observer, + observed=seeded_collection.observed, + dream_type=DreamType.OMNI, + ) + await enqueue_dream( + seeded_collection.workspace_name, + observer=seeded_collection.observer, + observed=seeded_collection.observed, + dream_type=DreamType.CARD_REFRESH, + rebuild=True, + ) + + items = ( + ( + await db_session.execute( + select(models.QueueItem).where( + models.QueueItem.workspace_name + == seeded_collection.workspace_name, + models.QueueItem.task_type == "dream", + models.QueueItem.processed == False, # noqa: E712 + ) + ) + ) + .scalars() + .all() + ) + assert len(items) == 2 + dream_types = {item.payload["dream_type"] for item in items} + assert dream_types == {"omni", "card_refresh"} + card_item = next( + item for item in items if item.payload["dream_type"] == "card_refresh" + ) + assert card_item.payload["rebuild"] is True + + +class TestProcessDreamDispatch: + @pytest.mark.asyncio + async def test_dispatches_card_refresh( + self, + seeded_collection: models.Collection, + ): + payload = DreamPayload( + dream_type=DreamType.CARD_REFRESH, + observer=seeded_collection.observer, + observed=seeded_collection.observed, + rebuild=True, + trigger_reason="manual", + ) + + with patch( + "src.dreamer.orchestrator.run_card_refresh_dream", + new=AsyncMock(return_value=_make_card_refresh_result()), + ) as mock_run: + await process_dream(payload, seeded_collection.workspace_name) + + assert mock_run.await_args is not None + kwargs = mock_run.await_args.kwargs + assert kwargs["workspace_name"] == seeded_collection.workspace_name + assert kwargs["observer"] == seeded_collection.observer + assert kwargs["observed"] == seeded_collection.observed + assert kwargs["rebuild"] is True + assert kwargs["dream_type"] == "card_refresh" + assert kwargs["trigger_reason"] == "manual" + + @pytest.mark.asyncio + async def test_card_refresh_does_not_advance_dream_guard( + self, + db_session: AsyncSession, + seeded_collection: models.Collection, + ): + """The omni guard pair (last_dream_at / last_dream_document_count) + must not move on a card refresh — it would delay real consolidation.""" + payload = DreamPayload( + dream_type=DreamType.CARD_REFRESH, + observer=seeded_collection.observer, + observed=seeded_collection.observed, + ) + + with patch( + "src.dreamer.orchestrator.run_card_refresh_dream", + new=AsyncMock(return_value=_make_card_refresh_result()), + ): + await process_dream(payload, seeded_collection.workspace_name) + + await db_session.refresh(seeded_collection) + dream_meta: dict[str, Any] = seeded_collection.internal_metadata.get( + "dream", {} + ) + assert "last_dream_at" not in dream_meta + assert "last_dream_document_count" not in dream_meta + + +class TestCardRefreshSpecialist: + def test_tools_exclude_observation_mutation(self): + for rebuild in (False, True): + specialist = CardRefreshSpecialist(rebuild=rebuild) + tool_names = {t["name"] for t in specialist.get_tools()} + assert tool_names == { + "get_recent_observations", + "search_memory", + "update_peer_card", + } + assert not tool_names & OBSERVATION_MUTATION_TOOLS + + def test_tools_without_peer_card_strip_update(self): + specialist = CardRefreshSpecialist() + tool_names = {t["name"] for t in specialist.get_tools(peer_card_enabled=False)} + assert "update_peer_card" not in tool_names + assert not tool_names & OBSERVATION_MUTATION_TOOLS + + def test_low_iteration_cap(self, monkeypatch: pytest.MonkeyPatch): + specialist = CardRefreshSpecialist() + assert specialist.get_max_iterations() == min( + 6, settings.DREAM.MAX_TOOL_ITERATIONS + ) + + monkeypatch.setattr(settings.DREAM, "MAX_TOOL_ITERATIONS", 4) + assert specialist.get_max_iterations() == 4 + + monkeypatch.setattr(settings.DREAM, "MAX_TOOL_ITERATIONS", 30) + assert specialist.get_max_iterations() == 6 + + def test_rebuild_flag_controls_card_injection(self): + assert CardRefreshSpecialist(rebuild=False).inject_peer_card is True + assert CardRefreshSpecialist(rebuild=True).inject_peer_card is False + + def test_rebuild_prompts_instruct_observation_only_build(self): + specialist = CardRefreshSpecialist(rebuild=True) + system_prompt = specialist.build_system_prompt("alice") + assert "REBUILD MODE" in system_prompt + assert "solely from the observations" in system_prompt + + user_prompt = specialist.build_user_prompt("alice", hints=None, peer_card=None) + assert "Rebuild the peer card" in user_prompt + assert "CURRENT PEER CARD" not in user_prompt + + async def _run_specialist( + self, specialist: CardRefreshSpecialist, stored_card: list[str] + ) -> tuple[AsyncMock, AsyncMock]: + """Run the specialist with a fully mocked LLM layer; returns the + (get_peer_card, honcho_llm_call) mocks for inspection.""" + mock_response = HonchoLLMCallResponse( + content="done", + input_tokens=10, + output_tokens=5, + finish_reasons=["stop"], + ) + mock_get_peer_card = AsyncMock(return_value=stored_card) + mock_llm_call = AsyncMock(return_value=mock_response) + + with ( + patch("src.dreamer.specialists.crud.get_peer", new=AsyncMock()), + patch( + "src.dreamer.specialists.crud.get_peer_card", + new=mock_get_peer_card, + ), + patch( + "src.dreamer.specialists.create_tool_executor", + new=AsyncMock(return_value=AsyncMock()), + ), + patch( + "src.dreamer.specialists.honcho_llm_call", + new=mock_llm_call, + ), + ): + result = await specialist.run( + workspace_name="workspace", + observer="alice", + observed="alice", + session_name=None, + ) + assert result.success is True + return mock_get_peer_card, mock_llm_call + + # Sentinel card entry that cannot collide with the prompt's own examples + # (the shared PEER CARD section contains e.g. "IDENTITY: Name: Alice"). + STORED_CARD: list[str] = [ + "IDENTITY: Name: Zorblax-Prime", + "ATTRIBUTE: Location: Ganymede", + ] + + @pytest.mark.asyncio + async def test_refresh_mode_injects_existing_card( + self, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setattr(settings.METRICS, "ENABLED", False) + + mock_get_peer_card, mock_llm_call = await self._run_specialist( + CardRefreshSpecialist(rebuild=False), self.STORED_CARD + ) + + mock_get_peer_card.assert_awaited_once() + assert mock_llm_call.await_args is not None + kwargs = mock_llm_call.await_args.kwargs + user_message = kwargs["messages"][1]["content"] + assert "IDENTITY: Name: Zorblax-Prime" in user_message + assert "CURRENT PEER CARD" in user_message + # Restricted tool offering and low iteration cap reach the LLM call. + tool_names = {t["name"] for t in kwargs["tools"]} + assert not tool_names & OBSERVATION_MUTATION_TOOLS + assert kwargs["max_tool_iterations"] == min( + 6, settings.DREAM.MAX_TOOL_ITERATIONS + ) + + @pytest.mark.asyncio + async def test_rebuild_mode_omits_existing_card( + self, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setattr(settings.METRICS, "ENABLED", False) + + mock_get_peer_card, mock_llm_call = await self._run_specialist( + CardRefreshSpecialist(rebuild=True), self.STORED_CARD + ) + + # The stored card is never even fetched, let alone injected. + mock_get_peer_card.assert_not_awaited() + assert mock_llm_call.await_args is not None + kwargs = mock_llm_call.await_args.kwargs + for message in kwargs["messages"]: + assert "IDENTITY: Name: Zorblax-Prime" not in message["content"] + # No CURRENT PEER CARD block in the user prompt (the system prompt's + # shared taxonomy section legitimately mentions the phrase). + assert "CURRENT PEER CARD" not in kwargs["messages"][1]["content"] diff --git a/tests/dreamer/test_model_config_usage.py b/tests/dreamer/test_model_config_usage.py index 90892632..18614043 100644 --- a/tests/dreamer/test_model_config_usage.py +++ b/tests/dreamer/test_model_config_usage.py @@ -32,6 +32,26 @@ def test_deduction_prompt_omits_peer_card_when_disabled() -> None: assert "IDENTITY:" not in prompt +def test_dreamer_system_prompts_delay_observed_observee_for_cache_prefix() -> None: + for specialist in (DeductionSpecialist(), InductionSpecialist()): + prompt = specialist.build_system_prompt("alice", peer_card_enabled=True) + other_prompt = specialist.build_system_prompt("bob", peer_card_enabled=True) + + assert prompt == other_prompt + assert "the target observee" in prompt + + +def test_dreamer_user_prompts_include_target_observee() -> None: + for specialist in (DeductionSpecialist(), InductionSpecialist()): + prompt = specialist.build_user_prompt( + observed="alice", + hints=None, + peer_card=None, + ) + + assert "Target observee:\nalice" in prompt + + def test_induction_prompt_has_no_peer_card_section() -> None: """Induction no longer writes to the peer card; its prompt must not reference it.""" prompt = InductionSpecialist().build_system_prompt("alice", peer_card_enabled=True) diff --git a/tests/dreamer/test_trees.py b/tests/dreamer/test_trees.py new file mode 100644 index 00000000..5c09a241 --- /dev/null +++ b/tests/dreamer/test_trees.py @@ -0,0 +1,51 @@ +import pytest + +from src.dreamer.trees import ( + CoverTree, + GraphSurprisal, + LSHSurprisal, + PrototypeSurprisal, + RPTree, + SklearnTreeWrapper, + SurprisalTree, + create_tree, +) + +ALL_TREE_TYPES = [ + "kdtree", + "balltree", + "rptree", + "covertree", + "lsh", + "graph", + "prototype", +] + +EXPECTED_CLASS = { + "kdtree": SklearnTreeWrapper, + "balltree": SklearnTreeWrapper, + "rptree": RPTree, + "covertree": CoverTree, + "lsh": LSHSurprisal, + "graph": GraphSurprisal, + "prototype": PrototypeSurprisal, +} + + +@pytest.mark.parametrize("tree_type", ALL_TREE_TYPES) +def test_create_tree_accepts_uniform_k_kwarg(tree_type: str): + tree = create_tree(tree_type=tree_type, k=5) + assert isinstance(tree, SurprisalTree) + assert isinstance(tree, EXPECTED_CLASS[tree_type]) + + +@pytest.mark.parametrize("tree_type", ALL_TREE_TYPES) +def test_create_tree_without_k(tree_type: str): + """The factory should also work when no ``k`` is supplied.""" + tree = create_tree(tree_type=tree_type) + assert isinstance(tree, EXPECTED_CLASS[tree_type]) + + +def test_create_tree_unknown_type_raises(): + with pytest.raises(ValueError, match="Unknown tree type"): + create_tree(tree_type="not_a_tree") diff --git a/tests/integration/test_message_embeddings.py b/tests/integration/test_message_embeddings.py index 198f2d36..b2f1b032 100644 --- a/tests/integration/test_message_embeddings.py +++ b/tests/integration/test_message_embeddings.py @@ -18,6 +18,7 @@ from src.config import settings from src.crud import create_messages from src.crud import message as message_crud from src.models import Message, Peer, Workspace +from src.reconciler.sync_vectors import run_vector_reconciliation_cycle from src.schemas import MessageCreate from src.utils.search import search @@ -161,9 +162,12 @@ async def test_blank_messages_are_not_sent_for_embedding( nonblank_content, ] - mock_openai_embeddings["batch_embed"].assert_awaited_once() - batch_arg = mock_openai_embeddings["batch_embed"].await_args.args[0] - assert batch_arg == {created_messages[1].public_id: nonblank_content} + # Inline embedding is gone: create_messages should chunk via prepare_chunks + # (no network) and never call batch_embed. + mock_openai_embeddings["batch_embed"].assert_not_awaited() + mock_openai_embeddings["prepare_chunks"].assert_called_once() + prepare_arg = mock_openai_embeddings["prepare_chunks"].call_args.args[0] + assert prepare_arg == {created_messages[1].public_id: nonblank_content} stmt = select(models.MessageEmbedding).where( models.MessageEmbedding.message_id.in_( @@ -176,6 +180,8 @@ async def test_blank_messages_are_not_sent_for_embedding( assert len(embedding_records) == 1 assert embedding_records[0].message_id == created_messages[1].public_id assert embedding_records[0].content == nonblank_content + assert embedding_records[0].sync_state == "pending" + assert embedding_records[0].embedding is None @pytest.mark.asyncio @@ -327,13 +333,23 @@ async def test_semantic_search_when_embeddings_enabled( assert len(created_messages) == 1 created_message = created_messages[0] - # Verify the embedding was created - stmt = select(models.MessageEmbedding).where( - models.MessageEmbedding.message_id == created_message.public_id + # The pending row exists, but the embedding is generated by the reconciler. + # Drive a reconciliation cycle so the row gets an embedding before search. + await db_session.commit() + await run_vector_reconciliation_cycle() + + # Verify the row was created and reconciled. expire_on_commit=False keeps + # stale cached ORM rows, so use populate_existing() to force a reload from + # the DB (the reconciler wrote in a different session). + stmt = ( + select(models.MessageEmbedding) + .where(models.MessageEmbedding.message_id == created_message.public_id) + .execution_options(populate_existing=True) ) result = await db_session.execute(stmt) embedding_record = result.scalar_one_or_none() assert embedding_record is not None + assert embedding_record.sync_state == "synced" # Now test semantic search without explicitly setting semantic=True # This should use semantic search because EMBED_MESSAGES is True @@ -361,6 +377,77 @@ async def test_semantic_search_when_embeddings_enabled( assert created_message.public_id in found_message_ids +@pytest.mark.asyncio +async def test_pgvector_search_excludes_pending_unembedded_rows( + db_session: AsyncSession, + sample_data: tuple[Workspace, Peer], +): + """Pending MessageEmbedding rows (embedding=None, awaiting the immediate + path or reconciler) must not appear in pgvector semantic search results: + their NULL distance sorts last and would pad the window with unranked + messages.""" + test_workspace, test_peer = sample_data + test_session = models.Session( + workspace_name=test_workspace.name, name=str(generate_nanoid()) + ) + db_session.add(test_session) + await db_session.commit() + + embedded_id = str(generate_nanoid()) + pending_id = str(generate_nanoid()) + for seq, (mid, content) in enumerate( + ((embedded_id, "embedded message"), (pending_id, "pending message")), start=1 + ): + db_session.add( + models.Message( + public_id=mid, + session_name=test_session.name, + workspace_name=test_workspace.name, + peer_name=test_peer.name, + content=content, + seq_in_session=seq, + ) + ) + await db_session.commit() + + dims = settings.EMBEDDING.VECTOR_DIMENSIONS + db_session.add_all( + [ + models.MessageEmbedding( + content="embedded message", + message_id=embedded_id, + workspace_name=test_workspace.name, + session_name=test_session.name, + peer_name=test_peer.name, + sync_state="synced", + embedding=[0.1] * dims, + ), + models.MessageEmbedding( + content="pending message", + message_id=pending_id, + workspace_name=test_workspace.name, + session_name=test_session.name, + peer_name=test_peer.name, + sync_state="pending", + embedding=None, + ), + ] + ) + await db_session.commit() + + snippets = await message_crud._search_messages_pgvector( # pyright: ignore[reportPrivateUsage] + db_session, + test_workspace.name, + test_session.name, + query_embedding=[0.1] * dims, + limit=10, + ) + + matched_ids = {msg.public_id for matched, _context in snippets for msg in matched} + assert embedded_id in matched_ids + assert pending_id not in matched_ids + + @pytest.mark.asyncio async def test_build_merged_snippets_batches_context_query_across_sessions(): """Context expansion should not issue one DB query per matched session.""" @@ -474,7 +561,10 @@ async def test_search_messages_external_lookup_happens_before_tracked_db( return [([message], [message])] @asynccontextmanager - async def fake_tracked_db(_operation_name: str | None = None): + async def fake_tracked_db( + _operation_name: str | None = None, *, read_only: bool = False + ): + del read_only call_order.append("enter") yield fake_db call_order.append("exit") @@ -578,7 +668,10 @@ async def test_search_messages_temporal_external_lookup_happens_before_tracked_d return [([message], [message])] @asynccontextmanager - async def fake_tracked_db(_operation_name: str | None = None): + async def fake_tracked_db( + _operation_name: str | None = None, *, read_only: bool = False + ): + del read_only call_order.append("enter") yield fake_db call_order.append("exit") @@ -632,15 +725,14 @@ async def test_message_chunking_creates_multiple_embeddings( test_message_content = "This is a very long message that should be chunked into multiple pieces because it exceeds the token limit that we set for testing purposes. This message contains many words and should definitely be split into multiple chunks." - def mock_batch_embed_chunked( - id_resource_dict: dict[str, str], - ) -> dict[str, list[list[float]]]: - return { - text_id: [[0.1] * 1536, [0.2] * 1536, [0.3] * 1536] # 3 chunks per message - for text_id in id_resource_dict - } + chunk_texts = ["chunk-a", "chunk-b", "chunk-c"] - mock_openai_embeddings["batch_embed"].side_effect = mock_batch_embed_chunked + def mock_prepare_chunks_chunked( + id_resource_dict: dict[str, str], + ) -> dict[str, list[str]]: + return {text_id: list(chunk_texts) for text_id in id_resource_dict} + + mock_openai_embeddings["prepare_chunks"].side_effect = mock_prepare_chunks_chunked messages = [ MessageCreate( @@ -660,22 +752,26 @@ async def test_message_chunking_creates_multiple_embeddings( assert len(created_messages) == 1 created_message = created_messages[0] - # Query the MessageEmbedding table to verify multiple embeddings were created - stmt = select(models.MessageEmbedding).where( - models.MessageEmbedding.message_id == created_message.public_id + # batch_embed is no longer called inline; embedding is deferred to reconciler. + mock_openai_embeddings["batch_embed"].assert_not_awaited() + + # Query the MessageEmbedding table to verify multiple pending rows were created, + # one per chunk, in chunk order (id ascending). + stmt = ( + select(models.MessageEmbedding) + .where(models.MessageEmbedding.message_id == created_message.public_id) + .order_by(models.MessageEmbedding.id) ) result = await db_session.execute(stmt) embedding_records = list(result.scalars().all()) - # Verify multiple embedding records were created (one per chunk) - # Embedding vectors are now stored externally in the vector store - assert len(embedding_records) == 3 # Should have 3 embeddings for 3 chunks + assert len(embedding_records) == 3 + assert [r.content for r in embedding_records] == chunk_texts - for _, embedding_record in enumerate(embedding_records): + for embedding_record in embedding_records: assert embedding_record.message_id == created_message.public_id - assert ( - embedding_record.content == test_message_content - ) # Full content stored in each assert embedding_record.workspace_name == test_workspace.name assert embedding_record.session_name == test_session.name assert embedding_record.peer_name == test_peer.name + assert embedding_record.sync_state == "pending" + assert embedding_record.embedding is None diff --git a/tests/live_llm/model_matrix.py b/tests/live_llm/model_matrix.py index a9abf0f9..2478c6c3 100644 --- a/tests/live_llm/model_matrix.py +++ b/tests/live_llm/model_matrix.py @@ -84,6 +84,18 @@ MODEL_FAMILIES: tuple[LiveModelFamily, ...] = ( supports_caching=False, docs_url="https://openrouter.ai/models", ), + # OpenAI-compatible providers that don't support OpenAI Structured Outputs + # (json_schema) and need structured_output_mode="json_object" on the + # ModelConfig. Point LLM_OPENAI_BASE_URL/API_KEY at the target (Z.AI GLM is + # the canonical #797 repro; vLLM/Ollama are self-hostable equivalents) and + # set the model via this env var. Empty default_models → skipped unless set. + LiveModelFamily( + provider="openai", + family="openai_json_object", + env_var="LIVE_LLM_OPENAI_JSON_OBJECT_MODELS", + supports_structured_output=True, + docs_url="https://docs.z.ai/guides/llm/glm-4.6", + ), LiveModelFamily( provider="gemini", family="gemini_2_5_class", diff --git a/tests/live_llm/test_live_openai.py b/tests/live_llm/test_live_openai.py index 60d89161..8b0544ec 100644 --- a/tests/live_llm/test_live_openai.py +++ b/tests/live_llm/test_live_openai.py @@ -25,6 +25,11 @@ _GPT5_SPECS = tuple( for spec in get_live_model_specs(provider="openai") if spec.family == "gpt_5_class" ) +_JSON_OBJECT_SPECS = tuple( + spec + for spec in get_live_model_specs(provider="openai") + if spec.family == "openai_json_object" +) @pytest.mark.asyncio @@ -88,7 +93,11 @@ async def test_live_openai_gpt5_reasoning_structured_output_and_prefix_caching( monkeypatch: pytest.MonkeyPatch, ) -> None: require_provider_key(model_spec) - backend, config = make_backend(model_spec, reasoning_effort="minimal") + # Only the original gpt-5 generation accepts 'minimal'; gpt-5.1+ replaced + # it with 'none'. 'low' is valid everywhere else, including future models. + is_base_gpt5 = model_spec.model == "gpt-5" or model_spec.model.startswith("gpt-5-") + reasoning_effort = "minimal" if is_base_gpt5 else "low" + backend, config = make_backend(model_spec, reasoning_effort=reasoning_effort) parse_calls = wrap_async_method( monkeypatch, backend._client.chat.completions, @@ -131,6 +140,57 @@ async def test_live_openai_gpt5_reasoning_structured_output_and_prefix_caching( assert second.cache_read_input_tokens > 0 assert parse_calls[0]["kwargs"]["response_format"] is StructuredLiveResponse - assert parse_calls[0]["kwargs"]["reasoning_effort"] == "minimal" + assert parse_calls[0]["kwargs"]["reasoning_effort"] == reasoning_effort assert "max_completion_tokens" in parse_calls[0]["kwargs"] assert "max_tokens" not in parse_calls[0]["kwargs"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_spec", _JSON_OBJECT_SPECS, ids=lambda spec: spec.id) +async def test_live_openai_json_object_structured_output( + model_spec: LiveModelSpec, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """For OpenAI-compatible providers without json_schema support, json_object + mode must skip parse(), request {"type": "json_object"}, and still produce a + valid structured object (the #797 fix, proven against a real provider). + + Configure: LLM_OPENAI_BASE_URL + LLM_OPENAI_API_KEY pointed at the target + provider, and LIVE_LLM_OPENAI_JSON_OBJECT_MODELS=. + """ + require_provider_key(model_spec) + backend, config = make_backend(model_spec, structured_output_mode="json_object") + parse_calls = wrap_async_method( + monkeypatch, backend._client.chat.completions, "parse" + ) + create_calls = wrap_async_method( + monkeypatch, backend._client.chat.completions, "create" + ) + + messages = [ + { + "role": "system", + "content": "You answer questions about a test run.", + }, + { + "role": "user", + "content": ( + "Return provider='openai', " + f"family='{model_spec.family}', and answer='json-object-ok'." + ), + }, + ] + + result = await execute_completion( + backend, + config, + messages=messages, + max_tokens=512, + response_format=StructuredLiveResponse, + ) + + assert isinstance(result.content, StructuredLiveResponse) + assert result.content.provider == "openai" + assert parse_calls == [] + assert create_calls, "expected a chat.completions.create call" + assert create_calls[0]["kwargs"]["response_format"] == {"type": "json_object"} diff --git a/tests/live_llm/test_live_structured_output_unions.py b/tests/live_llm/test_live_structured_output_unions.py new file mode 100644 index 00000000..bd69f8f1 --- /dev/null +++ b/tests/live_llm/test_live_structured_output_unions.py @@ -0,0 +1,162 @@ +"""Live coverage for union-bearing structured output without tools. + +The dialectic's final synthesis call carries response_format but tools=None +(see src/llm/tool_loop.py), and the model class is created dynamically from a +caller-supplied JSON Schema (src/utils/schema_conversion.py). That call shape +differs from test_live_tools_structured_output.py in one important way: with +no tools attached, the Gemini backend uses its NATIVE response_schema config +instead of injecting a schema instruction — and Gemini's response_schema +historically rejected anyOf. These tests drive a schema that exercises every +union-ish construct the converter supports (anyOf with null, a type list, an +enum, a $defs reference) through that exact call shape on all three +providers. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from pydantic import BaseModel, ValidationError + +from src.exceptions import LLMError +from src.llm.backend import CompletionResult +from src.llm.request_builder import execute_completion +from src.llm.structured_output import StructuredOutputError +from src.utils.schema_conversion import json_response_schema_to_pydantic + +from .conftest import make_backend, require_provider_key, wrap_async_method +from .model_matrix import LiveModelSpec, get_live_model_specs + +pytestmark = [pytest.mark.live_llm] + +_USER_FACTS_SCHEMA: dict[str, Any] = { + "type": "object", + "properties": { + "favorite_food": {"$ref": "#/$defs/Food"}, + "sentiment": {"enum": ["loves", "likes", "neutral", "dislikes", "hates"]}, + "years_vegetarian": {"type": ["integer", "null"]}, + "salient_fact": { + "anyOf": [{"type": "string"}, {"type": "null"}], + "description": "One short salient fact about the user, or null", + }, + }, + "required": ["favorite_food", "sentiment", "years_vegetarian"], + "$defs": { + "Food": { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + } + }, +} + +_PROMPT = ( + "The user said: 'I love sushi. I've been vegetarian for 3 years.' " + "Report the user's favorite food, their sentiment toward it, how many " + "years they have been vegetarian, and optionally one salient fact." +) + + +async def run_union_structured_flow(backend: Any, config: Any) -> CompletionResult: + """One no-tools turn that must return a schema-conforming answer. + + A fresh model class is created per call, matching how the dialectic + converts the caller's schema on every request. Retries mirror + test_live_tools_structured_output.py: empty/unparseable candidates are + absorbed by the executor's retry layer in production, but these tests + call the backend directly. + """ + response_model = json_response_schema_to_pydantic( + _USER_FACTS_SCHEMA, model_name="UserFactsReport" + ) + + result: CompletionResult | None = None + last_error: Exception | None = None + for _ in range(3): + try: + result = await execute_completion( + backend, + config, + messages=[{"role": "user", "content": _PROMPT}], + max_tokens=4096, + response_format=response_model, + ) + except (ValidationError, LLMError, StructuredOutputError) as exc: + last_error = exc + continue + break + if result is None: + raise AssertionError( + "union structured turn failed on all attempts" + ) from last_error + + content = result.content + assert isinstance(content, BaseModel), f"expected parsed model, got {content!r}" + # The model class is dynamic, so field access is untyped by construction. + report: Any = content + assert "sushi" in report.favorite_food.name.lower() + assert report.sentiment == "loves" + assert report.years_vegetarian == 3 + assert report.salient_fact is None or isinstance(report.salient_fact, str) + return result + + +@pytest.mark.asyncio +@pytest.mark.requires_anthropic +@pytest.mark.parametrize( + "model_spec", + get_live_model_specs(provider="anthropic", feature="structured_output"), + ids=lambda spec: spec.id, +) +async def test_live_anthropic_union_structured_output( + model_spec: LiveModelSpec, +) -> None: + require_provider_key(model_spec) + backend, config = make_backend(model_spec) + await run_union_structured_flow(backend, config) + + +@pytest.mark.asyncio +@pytest.mark.requires_openai +@pytest.mark.parametrize( + "model_spec", + get_live_model_specs(provider="openai", feature="structured_output"), + ids=lambda spec: spec.id, +) +async def test_live_openai_union_structured_output( + model_spec: LiveModelSpec, +) -> None: + require_provider_key(model_spec) + backend, config = make_backend(model_spec) + await run_union_structured_flow(backend, config) + + +@pytest.mark.asyncio +@pytest.mark.requires_gemini +@pytest.mark.parametrize( + "model_spec", + get_live_model_specs(provider="gemini", feature="structured_output"), + ids=lambda spec: spec.id, +) +async def test_live_gemini_union_structured_output( + model_spec: LiveModelSpec, + monkeypatch: pytest.MonkeyPatch, +) -> None: + require_provider_key(model_spec) + backend, config = make_backend(model_spec) + generate_calls = wrap_async_method( + monkeypatch, + backend._client.aio.models, + "generate_content", + ) + + await run_union_structured_flow(backend, config) + + # The point of this test: with no tools, Gemini must take the NATIVE + # response_schema path (the one that historically rejected anyOf), not + # the prompt-injection workaround used when tools are attached. + assert generate_calls + gen_config = generate_calls[-1]["kwargs"]["config"] + assert "response_schema" in gen_config + assert gen_config["response_mime_type"] == "application/json" diff --git a/tests/live_llm/test_live_tools_structured_output.py b/tests/live_llm/test_live_tools_structured_output.py new file mode 100644 index 00000000..33aac21b --- /dev/null +++ b/tests/live_llm/test_live_tools_structured_output.py @@ -0,0 +1,246 @@ +"""Live coverage for combining tool calling with structured output. + +Each provider needs a different workaround when a request carries both +function tools and a response_format (see src/llm/backends/): + +- OpenAI: parse() rejects non-strict function tools with a 500, so + tool-carrying structured requests must go through create() with an + explicit json_schema response_format. +- Anthropic: the '{' assistant prefill suppresses tool_use blocks, so it + must be skipped when tools are present (conditional instruction + + parse/repair instead). +- Gemini: native response_schema + function calling is rejected before + Gemini 3, so a schema instruction is injected into the final turn. + +The flow below drives both halves of the combination against real APIs: +a first turn that must produce a tool call (structured parsing skipped), +and a replay turn that must produce a schema-conforming final answer +while tools are still attached. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from pydantic import BaseModel, ValidationError + +from src.exceptions import LLMError +from src.llm.backend import CompletionResult +from src.llm.history_adapters import ( + AnthropicHistoryAdapter, + GeminiHistoryAdapter, + HistoryAdapter, + OpenAIHistoryAdapter, +) +from src.llm.request_builder import execute_completion +from src.llm.structured_output import StructuredOutputError + +from .conftest import ( + execute_local_tool, + favorite_prime_tools, + make_backend, + require_provider_key, + wrap_async_method, +) +from .model_matrix import LiveModelSpec, get_live_model_specs + +pytestmark = [pytest.mark.live_llm] + +_OPENAI_TOOL_SPECS = tuple( + spec + for spec in get_live_model_specs(provider="openai", feature="structured_output") + if spec.family in {"gpt_4_class", "gpt_5_class"} +) + + +class FavoritePrimeReport(BaseModel): + number: int + is_prime: bool + summary: str + + +_INITIAL_PROMPT = ( + "Before answering, call the get_favorite_prime tool exactly once. " + "Do not answer with plain text or JSON on this turn. " + "After you receive the tool result, answer with JSON where number is " + "the tool's number, is_prime says whether that number is prime, and " + "summary is one short sentence." +) + + +async def run_tool_then_structured_flow( + backend: Any, + config: Any, + adapter: HistoryAdapter, +) -> tuple[CompletionResult, CompletionResult]: + """First turn must tool-call (parsing skipped); replay turn must return + a schema-conforming answer with tools still attached.""" + initial_messages = [{"role": "user", "content": _INITIAL_PROMPT}] + tools = favorite_prime_tools() + + first = await execute_completion( + backend, + config, + messages=initial_messages, + max_tokens=4096, + tools=tools, + tool_choice="required", + response_format=FavoritePrimeReport, + ) + + assert first.tool_calls, "first turn should issue a tool call" + assert not isinstance( + first.content, FavoritePrimeReport + ), "tool-call turns carry no consumable content and must not be parsed" + + tool_call = first.tool_calls[0] + tool_result = execute_local_tool(tool_call.name, tool_call.input) + replay_messages = initial_messages + [ + adapter.format_assistant_tool_message(first), + *adapter.format_tool_results( + [ + { + "tool_id": tool_call.id, + "tool_name": tool_call.name, + "result": tool_result, + } + ] + ), + ] + + # tool_choice stays "auto" on the replay turn — the production tool loop + # (dialectic) never forces "none", and gemini-2.5-flash is prone to + # returning empty candidates under NONE mode. Retry the turn on empty / + # unparseable candidates (in production the executor's retry layer + # absorbs those; this calls the backend directly) and on the rare run + # where the model chooses to tool-call again instead of answering. + second: CompletionResult | None = None + last_error: Exception | None = None + for _ in range(3): + try: + candidate = await execute_completion( + backend, + config, + messages=replay_messages, + max_tokens=4096, + tools=tools, + tool_choice="auto", + response_format=FavoritePrimeReport, + ) + except (ValidationError, LLMError, StructuredOutputError) as exc: + last_error = exc + continue + if candidate.tool_calls: + last_error = AssertionError( + "model issued another tool call instead of answering" + ) + continue + second = candidate + break + if second is None: + raise AssertionError( + "structured replay turn failed on all attempts" + ) from last_error + + assert isinstance(second.content, FavoritePrimeReport) + assert second.content.number == 13 + assert second.content.is_prime is True + return first, second + + +@pytest.mark.asyncio +@pytest.mark.requires_anthropic +@pytest.mark.parametrize( + "model_spec", + get_live_model_specs(provider="anthropic", feature="structured_output"), + ids=lambda spec: spec.id, +) +async def test_live_anthropic_tools_with_structured_output( + model_spec: LiveModelSpec, + monkeypatch: pytest.MonkeyPatch, +) -> None: + require_provider_key(model_spec) + backend, config = make_backend(model_spec) + create_calls = wrap_async_method(monkeypatch, backend._client.messages, "create") + + await run_tool_then_structured_flow(backend, config, AnthropicHistoryAdapter()) + + assert len(create_calls) >= 2 + for call in create_calls: + messages = call["kwargs"]["messages"] + assert messages[-1] != { + "role": "assistant", + "content": "{", + }, "the '{' prefill would suppress tool_use blocks" + assert "If not responding with a tool call" in str( + messages[-1] + ), "schema instruction should use the conditional wording with tools" + + +@pytest.mark.asyncio +@pytest.mark.requires_openai +@pytest.mark.parametrize("model_spec", _OPENAI_TOOL_SPECS, ids=lambda spec: spec.id) +async def test_live_openai_tools_with_structured_output( + model_spec: LiveModelSpec, + monkeypatch: pytest.MonkeyPatch, +) -> None: + require_provider_key(model_spec) + # gpt-5.4 rejects function tools combined with any explicit + # reasoning_effort other than 'none' on /v1/chat/completions, so leave + # the parameter unset and let the server default apply. + backend, config = make_backend(model_spec) + parse_calls = wrap_async_method( + monkeypatch, backend._client.chat.completions, "parse" + ) + create_calls = wrap_async_method( + monkeypatch, backend._client.chat.completions, "create" + ) + + await run_tool_then_structured_flow(backend, config, OpenAIHistoryAdapter()) + + # favorite_prime_tools() is deliberately non-strict, the exact shape + # parse() refuses with a 500. + assert not parse_calls, "tool-carrying structured requests must avoid parse()" + assert len(create_calls) >= 2 + for call in create_calls: + response_format = call["kwargs"]["response_format"] + assert response_format["type"] == "json_schema" + assert response_format["json_schema"]["name"] == "FavoritePrimeReport" + + +@pytest.mark.asyncio +@pytest.mark.requires_gemini +@pytest.mark.parametrize( + "model_spec", + get_live_model_specs(provider="gemini", feature="structured_output"), + ids=lambda spec: spec.id, +) +async def test_live_gemini_tools_with_structured_output( + model_spec: LiveModelSpec, + monkeypatch: pytest.MonkeyPatch, +) -> None: + require_provider_key(model_spec) + # No temperature pin: gemini-2.5-flash occasionally returns an empty + # candidate on the replay turn, and at temperature=0 the retry re-sends + # a deterministic request — default sampling gives retries a real chance. + backend, config = make_backend(model_spec) + generate_calls = wrap_async_method( + monkeypatch, + backend._client.aio.models, + "generate_content", + ) + + await run_tool_then_structured_flow(backend, config, GeminiHistoryAdapter()) + + assert len(generate_calls) >= 2 + for call in generate_calls: + gen_config = call["kwargs"]["config"] + assert ( + "response_schema" not in gen_config + ), "native response_schema + function calling is rejected pre-Gemini 3" + assert "response_mime_type" not in gen_config + contents = call["kwargs"]["contents"] + assert "matching this schema" in str( + contents[-1] + ), "schema instruction should be injected into the final turn" diff --git a/tests/llm/test_backends/test_anthropic.py b/tests/llm/test_backends/test_anthropic.py index 52de0fa2..13255ec7 100644 --- a/tests/llm/test_backends/test_anthropic.py +++ b/tests/llm/test_backends/test_anthropic.py @@ -1,4 +1,5 @@ from types import SimpleNamespace +from typing import Any from unittest.mock import AsyncMock, Mock import pytest @@ -122,6 +123,99 @@ async def test_anthropic_backend_skips_assistant_prefill_for_claude_4_models() - assert call["messages"][0]["content"].startswith("Hello\n\nRespond with valid JSON") +@pytest.mark.asyncio +async def test_anthropic_backend_forwards_provider_params_passthroughs() -> None: + """provider_params.extra_body/extra_headers/extra_query reach the Anthropic + SDK call as kwargs of the same name (the SDK's documented passthrough). + """ + client = Mock() + client.messages.create = AsyncMock( + return_value=SimpleNamespace( + content=[TextBlock(type="text", text="ok")], + usage=SimpleNamespace( + input_tokens=10, + output_tokens=5, + cache_creation_input_tokens=0, + cache_read_input_tokens=0, + ), + stop_reason="end_turn", + ) + ) + + backend = AnthropicBackend(client) + await backend.complete( + model="claude-haiku-4-5", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + extra_params={ + "extra_body": {"anthropic_beta": ["context-1m-2025-01-15"]}, + "extra_headers": {"X-Proxy-Route": "vertex"}, + "extra_query": {"trace_id": "abc123"}, + }, + ) + + await_args = client.messages.create.await_args + if await_args is None: + raise AssertionError("Expected Anthropic client call") + call = await_args.kwargs + assert call["extra_body"] == {"anthropic_beta": ["context-1m-2025-01-15"]} + assert call["extra_headers"] == {"X-Proxy-Route": "vertex"} + assert call["extra_query"] == {"trace_id": "abc123"} + + +@pytest.mark.asyncio +async def test_anthropic_backend_stream_forwards_provider_params_passthroughs() -> None: + """The stream() path forwards provider_params passthroughs to the SDK the + same way complete() does — it has its own merge block, so cover it too. + """ + + class _FakeStream: + async def __aenter__(self) -> "_FakeStream": + return self + + async def __aexit__(self, *_: object) -> bool: + return False + + def __aiter__(self) -> "_FakeStream": + return self + + async def __anext__(self) -> object: + raise StopAsyncIteration + + async def get_final_message(self) -> SimpleNamespace: + return SimpleNamespace( + usage=SimpleNamespace(output_tokens=5), + stop_reason="end_turn", + ) + + client = Mock() + client.messages.stream = Mock(return_value=_FakeStream()) + + backend = AnthropicBackend(client) + chunks = [ + chunk + async for chunk in backend.stream( + model="claude-haiku-4-5", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + extra_params={ + "extra_body": {"anthropic_beta": ["context-1m-2025-01-15"]}, + "extra_headers": {"X-Proxy-Route": "vertex"}, + "extra_query": {"trace_id": "abc123"}, + }, + ) + ] + + assert chunks # the terminal is_done chunk is always emitted + call = client.messages.stream.call_args + if call is None: + raise AssertionError("Expected Anthropic stream call") + kwargs = call.kwargs + assert kwargs["extra_body"] == {"anthropic_beta": ["context-1m-2025-01-15"]} + assert kwargs["extra_headers"] == {"X-Proxy-Route": "vertex"} + assert kwargs["extra_query"] == {"trace_id": "abc123"} + + @pytest.mark.asyncio async def test_anthropic_backend_ignores_thinking_effort() -> None: client = Mock() @@ -152,3 +246,206 @@ async def test_anthropic_backend_ignores_thinking_effort() -> None: call = await_args.kwargs assert "thinking" not in call assert "reasoning_effort" not in call + + +def _make_client(content_blocks: list[Any]) -> Mock: + client = Mock() + client.messages.create = AsyncMock( + return_value=SimpleNamespace( + content=content_blocks, + usage=SimpleNamespace( + input_tokens=10, + output_tokens=5, + cache_creation_input_tokens=0, + cache_read_input_tokens=0, + ), + stop_reason="end_turn", + ) + ) + return client + + +SEARCH_TOOL = { + "name": "search", + "description": "Search for information", + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, +} + + +@pytest.mark.asyncio +async def test_anthropic_backend_no_prefill_when_tools_present() -> None: + """With tools + response_format, the '{' prefill must be skipped and the + schema instruction must be conditional, so tool_use blocks stay reachable.""" + client = _make_client([TextBlock(type="text", text='{"answer":"ok"}')]) + + backend = AnthropicBackend(client) + result = await backend.complete( + # claude-3-5 supports prefill, so only the tools guard prevents it here + model="claude-3-5-sonnet-latest", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + tools=[SEARCH_TOOL], + response_format=StructuredResponse, + ) + + assert isinstance(result.content, StructuredResponse) + call = client.messages.create.await_args.kwargs + assert call["messages"][-1]["role"] == "user" # no assistant '{' prefill + assert ( + "If not responding with a tool call, respond with valid JSON" + in call["messages"][0]["content"] + ) + + +@pytest.mark.asyncio +async def test_anthropic_backend_prefill_unchanged_without_tools() -> None: + """Tool-less structured calls keep the prefill + unconditional wording.""" + client = _make_client([TextBlock(type="text", text='"answer":"ok"}')]) + + backend = AnthropicBackend(client) + result = await backend.complete( + model="claude-3-5-sonnet-latest", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=StructuredResponse, + ) + + assert isinstance(result.content, StructuredResponse) + call = client.messages.create.await_args.kwargs + assert call["messages"][-1] == {"role": "assistant", "content": "{"} + instruction = call["messages"][0]["content"] + assert "\n\nRespond with valid JSON matching this schema:" in instruction + assert "If not responding with a tool call" not in instruction + + +@pytest.mark.asyncio +async def test_anthropic_backend_repairs_malformed_structured_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Structured output that fails JSON parsing falls back to + repair_response_model_json, whose result becomes the response content.""" + # After the '{' prefill is prepended this is still invalid JSON. + client = _make_client([TextBlock(type="text", text='"answer": not-json')]) + + repaired = StructuredResponse(answer="fixed") + repair_calls: list[tuple[str, type[BaseModel], str]] = [] + + def _fake_repair( + raw: str, response_format: type[BaseModel], model_name: str + ) -> StructuredResponse: + repair_calls.append((raw, response_format, model_name)) + return repaired + + monkeypatch.setattr( + "src.llm.backends.anthropic.repair_response_model_json", _fake_repair + ) + + backend = AnthropicBackend(client) + result = await backend.complete( + model="claude-3-5-sonnet-latest", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=StructuredResponse, + ) + + assert result.content is repaired + assert repair_calls == [ + ('{"answer": not-json', StructuredResponse, "claude-3-5-sonnet-latest") + ] + + +@pytest.mark.asyncio +async def test_anthropic_backend_skips_parsing_on_tool_call_turns( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A tool-call response with response_format set must not attempt JSON + parsing — the repair fallback raises on the empty text of tool-call + turns, which would fail every intermediate tool iteration.""" + client = _make_client( + [ + ToolUseBlock( + type="tool_use", + id="tool_1", + name="search", + input={"query": "honcho"}, + ) + ] + ) + + def _fail_repair(*_args: object, **_kwargs: object) -> None: + raise AssertionError("repair must not run for tool-call turns") + + monkeypatch.setattr( + "src.llm.backends.anthropic.repair_response_model_json", _fail_repair + ) + + backend = AnthropicBackend(client) + result = await backend.complete( + model="claude-3-5-sonnet-latest", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + tools=[SEARCH_TOOL], + response_format=StructuredResponse, + ) + + assert result.content == "" # raw (empty) text, not a parsed model + assert result.tool_calls[0].name == "search" + + +@pytest.mark.asyncio +async def test_anthropic_backend_stream_no_prefill_when_tools_present() -> None: + """The streaming path applies the same tools guard.""" + + class _FakeStream: + def __init__(self) -> None: + self._chunks: list[SimpleNamespace] = [ + SimpleNamespace( + type="content_block_delta", + delta=SimpleNamespace(text='{"answer":"ok"}'), + ) + ] + + async def __aenter__(self): + return self + + async def __aexit__(self, *_args: object) -> bool: + return False + + def __aiter__(self): + return self + + async def __anext__(self): + if self._chunks: + return self._chunks.pop(0) + raise StopAsyncIteration + + async def get_final_message(self): + return SimpleNamespace( + stop_reason="end_turn", usage=SimpleNamespace(output_tokens=5) + ) + + client = Mock() + client.messages.stream = Mock(return_value=_FakeStream()) + + backend = AnthropicBackend(client) + chunks = [ + chunk + async for chunk in backend.stream( + model="claude-3-5-sonnet-latest", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + tools=[SEARCH_TOOL], + response_format=StructuredResponse, + ) + ] + + assert chunks[0].content == '{"answer":"ok"}' + call = client.messages.stream.call_args.kwargs + assert call["messages"][-1]["role"] == "user" # no assistant '{' prefill + assert ( + "If not responding with a tool call, respond with valid JSON" + in call["messages"][0]["content"] + ) diff --git a/tests/llm/test_backends/test_gemini.py b/tests/llm/test_backends/test_gemini.py index b327c8e4..1e3933b4 100644 --- a/tests/llm/test_backends/test_gemini.py +++ b/tests/llm/test_backends/test_gemini.py @@ -315,6 +315,119 @@ async def test_gemini_backend_strips_system_and_tools_when_using_cached_content( assert "tool_config" not in call["config"] +@pytest.mark.asyncio +async def test_gemini_backend_forwards_provider_params_extra_body() -> None: + """provider_params.extra_body merges into the GenerateContentConfig dict + (Gemini's body-shaped fields live there, not as an SDK kwarg). + """ + client = Mock() + client.aio.models.generate_content = AsyncMock( + return_value=SimpleNamespace( + candidates=[ + SimpleNamespace( + finish_reason=SimpleNamespace(name="STOP"), + content=SimpleNamespace(parts=[SimpleNamespace(text="ok")]), + ) + ], + usage_metadata=SimpleNamespace( + prompt_token_count=12, + candidates_token_count=6, + ), + parsed=None, + ) + ) + + backend = GeminiBackend(client) + await backend.complete( + model="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + extra_params={"extra_body": {"candidate_count": 3, "seed": 42}}, + ) + + await_args = client.aio.models.generate_content.await_args + if await_args is None: + raise AssertionError("Expected Gemini generate_content call") + call = await_args.kwargs + assert call["config"]["candidate_count"] == 3 + assert call["config"]["seed"] == 42 + + +@pytest.mark.asyncio +async def test_gemini_backend_forwards_provider_params_extra_headers() -> None: + """provider_params.extra_headers folds into config.http_options.headers.""" + client = Mock() + client.aio.models.generate_content = AsyncMock( + return_value=SimpleNamespace( + candidates=[ + SimpleNamespace( + finish_reason=SimpleNamespace(name="STOP"), + content=SimpleNamespace(parts=[SimpleNamespace(text="ok")]), + ) + ], + usage_metadata=SimpleNamespace( + prompt_token_count=12, + candidates_token_count=6, + ), + parsed=None, + ) + ) + + backend = GeminiBackend(client) + await backend.complete( + model="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + extra_params={"extra_headers": {"X-Trace-Id": "abc123"}}, + ) + + await_args = client.aio.models.generate_content.await_args + if await_args is None: + raise AssertionError("Expected Gemini generate_content call") + call = await_args.kwargs + assert call["config"]["http_options"]["headers"] == {"X-Trace-Id": "abc123"} + + +@pytest.mark.asyncio +async def test_gemini_backend_silently_ignores_extra_query() -> None: + """extra_query has no google-genai SDK equivalent. The backend drops it + rather than crashing or surfacing it somewhere unexpected. + """ + client = Mock() + client.aio.models.generate_content = AsyncMock( + return_value=SimpleNamespace( + candidates=[ + SimpleNamespace( + finish_reason=SimpleNamespace(name="STOP"), + content=SimpleNamespace(parts=[SimpleNamespace(text="ok")]), + ) + ], + usage_metadata=SimpleNamespace( + prompt_token_count=12, + candidates_token_count=6, + ), + parsed=None, + ) + ) + + backend = GeminiBackend(client) + await backend.complete( + model="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + extra_params={"extra_query": {"trace_id": "abc123"}}, + ) + + await_args = client.aio.models.generate_content.await_args + if await_args is None: + raise AssertionError("Expected Gemini generate_content call") + call = await_args.kwargs + # extra_query is not surfaced anywhere in the request — neither at the top + # level of the SDK call nor inside config. + assert "extra_query" not in call + assert "extra_query" not in call["config"] + + def test_gemini_sanitize_schema_strips_unsupported_keywords() -> None: """Gemini's function-declarations validator rejects JSON-Schema keywords outside its narrow allowlist (additionalProperties, allOf, if/then, $ref, @@ -389,3 +502,126 @@ def test_gemini_convert_tools_sanitizes_parameters_schema() -> None: params = converted[0]["function_declarations"][0]["parameters"] assert "additionalProperties" not in params assert "additionalProperties" not in params["properties"]["observations"]["items"] + + +class _GeminiStructured(BaseModel): + answer: str + + +GEMINI_AGENT_TOOL = { + "name": "search", + "description": "Search for information", + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, +} + + +def _gemini_response( + parts: list[SimpleNamespace], parsed: object = None +) -> SimpleNamespace: + return SimpleNamespace( + candidates=[ + SimpleNamespace( + finish_reason=SimpleNamespace(name="STOP"), + content=SimpleNamespace(parts=parts), + ) + ], + usage_metadata=SimpleNamespace( + prompt_token_count=12, + candidates_token_count=6, + ), + parsed=parsed, + ) + + +@pytest.mark.asyncio +async def test_gemini_backend_structured_with_tools_skips_native_schema() -> None: + """Native response_schema + function calling is Gemini-3-preview-only, so + with tools present the schema must be delivered as an instruction on the + final turn instead, and the answer parsed from raw text.""" + client = Mock() + client.aio.models.generate_content = AsyncMock( + return_value=_gemini_response([SimpleNamespace(text='{"answer":"ok"}')]) + ) + + backend = GeminiBackend(client) + result = await backend.complete( + model="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + tools=[GEMINI_AGENT_TOOL], + response_format=_GeminiStructured, + ) + + assert isinstance(result.content, _GeminiStructured) + assert result.content.answer == "ok" + await_args = client.aio.models.generate_content.await_args + if await_args is None: + raise AssertionError("Expected Gemini generate_content call") + call = await_args.kwargs + assert "response_schema" not in call["config"] + assert "response_mime_type" not in call["config"] + assert "tools" in call["config"] + last_part_text = call["contents"][-1]["parts"][-1]["text"] + assert "If not responding with a tool call" in last_part_text + + +@pytest.mark.asyncio +async def test_gemini_backend_structured_tool_call_turn_not_parsed() -> None: + """A tool-call turn under tools + response_format must not attempt JSON + parsing (its text is empty and the repair fallback raises on that).""" + client = Mock() + client.aio.models.generate_content = AsyncMock( + return_value=_gemini_response( + [ + SimpleNamespace( + function_call=SimpleNamespace( + name="search", args={"query": "honcho"} + ), + thought_signature=None, + ) + ] + ) + ) + + backend = GeminiBackend(client) + result = await backend.complete( + model="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + tools=[GEMINI_AGENT_TOOL], + response_format=_GeminiStructured, + ) + + assert result.content == "" # raw (empty) text, not a parsed model + assert result.tool_calls[0].name == "search" + assert result.tool_calls[0].input == {"query": "honcho"} + + +@pytest.mark.asyncio +async def test_gemini_backend_structured_without_tools_uses_native_schema() -> None: + """Tool-less structured calls keep native response_schema enforcement.""" + client = Mock() + client.aio.models.generate_content = AsyncMock( + return_value=_gemini_response( + [SimpleNamespace(text='{"answer":"ok"}')], + parsed={"answer": "ok"}, + ) + ) + + backend = GeminiBackend(client) + result = await backend.complete( + model="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=_GeminiStructured, + ) + + assert isinstance(result.content, _GeminiStructured) + call = client.aio.models.generate_content.await_args.kwargs # pyright: ignore + assert call["config"]["response_schema"] is _GeminiStructured + assert call["config"]["response_mime_type"] == "application/json" + # No instruction injected on the tool-less path. + assert call["contents"][-1]["parts"][-1]["text"] == "Hello" diff --git a/tests/llm/test_backends/test_openai.py b/tests/llm/test_backends/test_openai.py index 695b12cd..8567e034 100644 --- a/tests/llm/test_backends/test_openai.py +++ b/tests/llm/test_backends/test_openai.py @@ -1,9 +1,71 @@ +import gc +import json +import weakref +from collections.abc import AsyncIterator from types import SimpleNamespace +from typing import Any from unittest.mock import AsyncMock, Mock +import httpx import pytest +from openai import BadRequestError +from pydantic import BaseModel -from src.llm.backends.openai import OpenAIBackend +from src.exceptions import ValidationException +from src.llm.backends.openai import ( + OpenAIBackend, + _json_object_instruction, # pyright: ignore[reportPrivateUsage] +) +from src.utils.representation import PromptRepresentation +from src.utils.schema_conversion import json_response_schema_to_pydantic + + +def _await_kwargs(mock_method: Any) -> dict[str, Any]: + await_args = mock_method.await_args + if await_args is None: + raise AssertionError("Expected the mocked method to have been awaited") + return await_args.kwargs + + +def _bad_request_error() -> BadRequestError: + """A 400 like a provider that doesn't support json_schema would return.""" + request = httpx.Request("POST", "https://example.test/v1/chat/completions") + response = httpx.Response(400, request=request) + return BadRequestError( + "response_format json_schema is not supported", response=response, body=None + ) + + +async def _empty_stream() -> AsyncIterator[Any]: + chunks: list[Any] = [] # async generator that yields nothing + for chunk in chunks: + yield chunk + + +class _StructuredResponse(BaseModel): + answer: str + + +def _structured_create_return(content: str, parsed: Any = None) -> SimpleNamespace: + return SimpleNamespace( + choices=[ + SimpleNamespace( + finish_reason="stop", + message=SimpleNamespace( + content=content, + parsed=parsed, + tool_calls=[], + reasoning_details=[], + refusal=None, + ), + ) + ], + usage=SimpleNamespace( + prompt_tokens=10, + completion_tokens=5, + prompt_tokens_details=None, + ), + ) @pytest.mark.asyncio @@ -227,6 +289,164 @@ async def test_openai_backend_skips_extra_body_when_thinking_budget_zero() -> No assert "extra_body" not in call +@pytest.mark.asyncio +async def test_openai_backend_forwards_provider_params_extra_body() -> None: + """Operator-supplied extra_body in provider_params reaches the OpenAI SDK call. + + This is the escape hatch for OpenAI-compatible proxies that translate to + other providers (litellm → Vertex AI Anthropic) and need provider-native + body fields (e.g. Anthropic's `thinking`). + """ + client = Mock() + client.chat.completions.create = AsyncMock( + return_value=SimpleNamespace( + choices=[ + SimpleNamespace( + finish_reason="stop", + message=SimpleNamespace( + content="ok", + tool_calls=[], + reasoning_details=[], + ), + ) + ], + usage=SimpleNamespace( + prompt_tokens=10, + completion_tokens=5, + prompt_tokens_details=None, + ), + ) + ) + + backend = OpenAIBackend(client) + await backend.complete( + model="claude-haiku-4-5", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + extra_params={ + "extra_body": {"thinking": {"type": "enabled", "budget_tokens": 4096}} + }, + ) + + await_args = client.chat.completions.create.await_args + if await_args is None: + raise AssertionError("Expected OpenAI create call") + call = await_args.kwargs + assert call["extra_body"] == { + "thinking": {"type": "enabled", "budget_tokens": 4096} + } + + +@pytest.mark.asyncio +async def test_openai_backend_forwards_provider_params_extra_headers_and_query() -> ( + None +): + client = Mock() + client.chat.completions.create = AsyncMock( + return_value=SimpleNamespace( + choices=[ + SimpleNamespace( + finish_reason="stop", + message=SimpleNamespace( + content="ok", + tool_calls=[], + reasoning_details=[], + ), + ) + ], + usage=SimpleNamespace( + prompt_tokens=10, + completion_tokens=5, + prompt_tokens_details=None, + ), + ) + ) + + backend = OpenAIBackend(client) + await backend.complete( + model="gpt-4.1", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + extra_params={ + "extra_headers": {"X-Proxy-Route": "vertex"}, + "extra_query": {"trace_id": "abc123"}, + }, + ) + + await_args = client.chat.completions.create.await_args + if await_args is None: + raise AssertionError("Expected OpenAI create call") + call = await_args.kwargs + assert call["extra_headers"] == {"X-Proxy-Route": "vertex"} + assert call["extra_query"] == {"trace_id": "abc123"} + + +@pytest.mark.asyncio +async def test_openai_backend_operator_extra_body_wins_over_auto_injection() -> None: + """When the operator supplies extra_body.reasoning, it must replace the + value that thinking_budget_tokens would otherwise auto-inject (operator-wins + shallow merge). + """ + client = Mock() + client.chat.completions.create = AsyncMock( + return_value=SimpleNamespace( + choices=[ + SimpleNamespace( + finish_reason="stop", + message=SimpleNamespace( + content="ok", + tool_calls=[], + reasoning_details=[], + ), + ) + ], + usage=SimpleNamespace( + prompt_tokens=10, + completion_tokens=5, + prompt_tokens_details=None, + ), + ) + ) + + backend = OpenAIBackend(client) + await backend.complete( + model="x-ai/grok-4.1-fast", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + thinking_budget_tokens=256, + extra_params={ + "extra_body": {"reasoning": {"effort": "high", "max_tokens": 9999}} + }, + ) + + await_args = client.chat.completions.create.await_args + if await_args is None: + raise AssertionError("Expected OpenAI create call") + call = await_args.kwargs + # Operator's whole `reasoning` dict replaces Honcho's auto-injected one. + assert call["extra_body"] == {"reasoning": {"effort": "high", "max_tokens": 9999}} + + +@pytest.mark.asyncio +async def test_openai_backend_rejects_non_mapping_passthrough() -> None: + """A non-mapping passthrough (operator misconfiguration) raises a clear + ValidationException instead of an opaque TypeError deep in the transport. + """ + client = Mock() + client.chat.completions.create = AsyncMock() + + backend = OpenAIBackend(client) + with pytest.raises(ValidationException, match=r"provider_params\.extra_headers"): + await backend.complete( + model="gpt-4.1", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + extra_params={"extra_headers": [["X-Foo", "bar"]]}, + ) + + client.chat.completions.create.assert_not_awaited() + + @pytest.mark.asyncio async def test_openai_backend_converts_anthropic_style_tools() -> None: client = Mock() @@ -290,6 +510,67 @@ async def test_openai_backend_converts_anthropic_style_tools() -> None: assert call["tool_choice"] == "required" +async def test_openai_backend_translates_canonical_any_tool_choice_to_required() -> ( + None +): + """Regression: a Gemini→OpenAI fallback passes canonical "any", which OpenAI + rejects as an invalid param. The backend must translate it to "required".""" + client = Mock() + client.chat.completions.create = AsyncMock( + return_value=SimpleNamespace( + choices=[ + SimpleNamespace( + finish_reason="stop", + message=SimpleNamespace( + content="ok", + tool_calls=[], + reasoning_details=[], + ), + ) + ], + usage=SimpleNamespace( + prompt_tokens=10, + completion_tokens=5, + prompt_tokens_details=None, + ), + ) + ) + + backend = OpenAIBackend(client) + await backend.complete( + model="gpt-4.1", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + tools=[ + { + "name": "get_weather", + "description": "Lookup weather", + "input_schema": {"type": "object", "properties": {}}, + } + ], + tool_choice="any", + ) + + assert _await_kwargs(client.chat.completions.create)["tool_choice"] == "required" + + +@pytest.mark.parametrize( + ("canonical", "expected"), + [ + ("any", "required"), + ("required", "required"), + ("auto", "auto"), + ("none", "none"), + (None, None), + ("search", {"type": "function", "function": {"name": "search"}}), + ({"name": "search"}, {"type": "function", "function": {"name": "search"}}), + ({"type": "function"}, {"type": "function"}), + ], +) +def test_openai_convert_tool_choice(canonical: Any, expected: Any) -> None: + assert OpenAIBackend._convert_tool_choice(canonical) == expected # pyright: ignore[reportPrivateUsage] + + @pytest.mark.parametrize( "model", [ @@ -334,3 +615,525 @@ def test_openai_classic_models_use_max_tokens(model: str) -> None: ) assert _uses_max_completion_tokens(model) is False + + +@pytest.mark.asyncio +async def test_structured_output_parsed_none_with_raw_content_repairs() -> None: + """parse() returning parsed=None but with raw content repairs that content.""" + client = Mock() + client.chat.completions.parse = AsyncMock( + return_value=_structured_create_return('{"answer": "ok"}', parsed=None) + ) + + backend = OpenAIBackend(client) + result = await backend.complete( + model="gpt-4.1", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=_StructuredResponse, + ) + + assert isinstance(result.content, _StructuredResponse) + assert result.content.answer == "ok" + + +@pytest.mark.asyncio +async def test_structured_output_parsed_none_returns_refusal() -> None: + """parse() returning parsed=None with no content surfaces the refusal.""" + client = Mock() + response = _structured_create_return("", parsed=None) + response.choices[0].message.refusal = "I can't help with that" + client.chat.completions.parse = AsyncMock(return_value=response) + + backend = OpenAIBackend(client) + result = await backend.complete( + model="gpt-4.1", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=_StructuredResponse, + ) + + assert result.content == "I can't help with that" + + +@pytest.mark.asyncio +async def test_structured_output_parsed_none_no_content_raises() -> None: + """json_schema with no parsed model, content, or refusal raises so the + retry/fallback chain engages — it must NOT silently empty like json_object.""" + from src.exceptions import ValidationException + + client = Mock() + client.chat.completions.parse = AsyncMock( + return_value=_structured_create_return("", parsed=None) + ) + + backend = OpenAIBackend(client) + with pytest.raises(ValidationException): + await backend.complete( + model="gpt-4.1", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=_StructuredResponse, + ) + + +@pytest.mark.asyncio +async def test_structured_output_default_mode_uses_parse() -> None: + """Default json_schema mode uses parse(), not the json_object path.""" + client = Mock() + client.chat.completions.parse = AsyncMock( + return_value=_structured_create_return( + '{"answer": "ok"}', parsed=_StructuredResponse(answer="ok") + ) + ) + client.chat.completions.create = AsyncMock() + + backend = OpenAIBackend(client) + result = await backend.complete( + model="gpt-4.1", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=_StructuredResponse, + ) + + assert client.chat.completions.parse.await_count == 1 + assert client.chat.completions.create.await_count == 0 + parse_call = _await_kwargs(client.chat.completions.parse) + assert parse_call["response_format"] is _StructuredResponse + assert isinstance(result.content, _StructuredResponse) + + +@pytest.mark.asyncio +async def test_structured_output_json_schema_rejected_returns_empty_without_second_request() -> ( + None +): + """A provider that rejects json_schema (400) returns empty, no second request. + + Retrying or re-requesting the same shape is pointless (#797), so a + BadRequestError is swallowed to an empty representation rather than erroring. + """ + client = Mock() + client.chat.completions.parse = AsyncMock(side_effect=_bad_request_error()) + client.chat.completions.create = AsyncMock() + + backend = OpenAIBackend(client) + result = await backend.complete( + model="glm-4.6", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=PromptRepresentation, + ) + + assert client.chat.completions.parse.await_count == 1 + assert client.chat.completions.create.await_count == 0 # no second request + assert isinstance(result.content, PromptRepresentation) + assert result.content.explicit == [] + + +@pytest.mark.asyncio +async def test_structured_output_json_schema_rejected_with_required_fields_does_not_raise() -> ( + None +): + """A json_schema rejection must not raise even when the response model has + required fields (empty_structured_output() can't build an empty instance) — + it falls back to empty content instead of escaping the handler.""" + client = Mock() + client.chat.completions.parse = AsyncMock(side_effect=_bad_request_error()) + client.chat.completions.create = AsyncMock() + + backend = OpenAIBackend(client) + result = await backend.complete( + model="glm-4.6", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=_StructuredResponse, # has a required field + ) + + assert client.chat.completions.parse.await_count == 1 + assert client.chat.completions.create.await_count == 0 # no second request + assert result.content == "" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "exc", + [ + json.JSONDecodeError("Expecting value", "not json", 0), + ValueError("transient parse glitch"), + ], +) +async def test_structured_output_transient_parse_error_propagates_for_retry( + exc: Exception, +) -> None: + """Non-400 parse failures propagate so the retry/fallback chain can engage. + + Only a 400 (json_schema rejection) is treated as terminal-empty; a transient + decode/validation glitch must re-raise so tenacity retries and the fallback + model gets a chance — not be silently swallowed to empty on the first try. + """ + client = Mock() + client.chat.completions.parse = AsyncMock(side_effect=exc) + client.chat.completions.create = AsyncMock() + + backend = OpenAIBackend(client) + with pytest.raises(type(exc)): + await backend.complete( + model="glm-4.6", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=PromptRepresentation, + ) + + assert client.chat.completions.create.await_count == 0 # no second request + + +@pytest.mark.asyncio +async def test_structured_output_json_object_mode_request_shape() -> None: + """json_object mode skips parse(), requests json_object, injects the schema.""" + client = Mock() + client.chat.completions.parse = AsyncMock() + client.chat.completions.create = AsyncMock( + return_value=_structured_create_return('{"answer": "ok"}') + ) + + backend = OpenAIBackend(client) + result = await backend.complete( + model="glm-4.6", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=_StructuredResponse, + extra_params={"structured_output_mode": "json_object"}, + ) + + assert client.chat.completions.parse.await_count == 0 + assert client.chat.completions.create.await_count == 1 + call = _await_kwargs(client.chat.completions.create) + assert call["response_format"] == {"type": "json_object"} + + system_messages = [m for m in call["messages"] if m["role"] == "system"] + assert system_messages, "expected a system message carrying the schema" + system_content = system_messages[0]["content"] + assert "JSON" in system_content + assert "json" in system_content + assert "answer" in system_content # schema property serialized in + assert isinstance(result.content, _StructuredResponse) + assert result.content.answer == "ok" + + +@pytest.mark.asyncio +async def test_structured_output_json_object_mode_repairs_markdown() -> None: + """A provider that ignores json_object and returns prose must not crash — + PromptRepresentation repairs to an empty representation, not an exception.""" + client = Mock() + client.chat.completions.parse = AsyncMock() + client.chat.completions.create = AsyncMock( + return_value=_structured_create_return( + "Sure! Here are the facts:\n- the user likes coffee" + ) + ) + + backend = OpenAIBackend(client) + result = await backend.complete( + model="glm-4.6", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=PromptRepresentation, + extra_params={"structured_output_mode": "json_object"}, + ) + + assert isinstance(result.content, PromptRepresentation) + + +@pytest.mark.asyncio +async def test_structured_output_json_object_mode_with_dynamic_model() -> None: + """json_object mode composes with a caller-supplied schema converted at + request time (the dialectic's response_format path): the generated class's + schema — unions included — is injected into the prompt and the JSON body + parses back through the class.""" + response_model = json_response_schema_to_pydantic( + { + "type": "object", + "properties": { + "answer": {"type": "string"}, + "years": {"type": ["integer", "null"]}, + }, + "required": ["answer", "years"], + } + ) + client = Mock() + client.chat.completions.parse = AsyncMock() + client.chat.completions.create = AsyncMock( + return_value=_structured_create_return('{"answer": "ok", "years": 3}') + ) + + backend = OpenAIBackend(client) + result = await backend.complete( + model="glm-4.6", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=response_model, + extra_params={"structured_output_mode": "json_object"}, + ) + + assert client.chat.completions.parse.await_count == 0 + call = _await_kwargs(client.chat.completions.create) + assert call["response_format"] == {"type": "json_object"} + system_messages = [m for m in call["messages"] if m["role"] == "system"] + assert system_messages, "expected a system message carrying the schema" + # The dynamic model's schema (union field included) made it into the prompt. + assert "years" in system_messages[0]["content"] + assert "anyOf" in system_messages[0]["content"] + + assert isinstance(result.content, response_model) + # The model class is dynamic, so field access is untyped by construction. + content: Any = result.content + assert content.answer == "ok" + assert content.years == 3 + + +def test_json_object_instruction_does_not_pin_dynamic_models() -> None: + """The instruction cache must hold its model-class keys weakly: the + dialectic creates a fresh response_format class per request (see + src/utils/schema_conversion.py), and a strong-keyed cache would grow by + one pinned class per structured chat call.""" + model = json_response_schema_to_pydantic( + {"type": "object", "properties": {"answer": {"type": "string"}}} + ) + first = _json_object_instruction(model) + assert _json_object_instruction(model) is first # cached while alive + + ref = weakref.ref(model) + del model + gc.collect() + assert ref() is None, "instruction cache must not keep dynamic classes alive" + + +@pytest.mark.asyncio +async def test_structured_output_json_object_empty_content_returns_empty() -> None: + """An empty body with no refusal must produce a graceful empty result, not + raise — matching the json_schema path's behavior on a contentless response. + """ + client = Mock() + client.chat.completions.create = AsyncMock( + return_value=_structured_create_return("") + ) + + backend = OpenAIBackend(client) + result = await backend.complete( + model="glm-4.6", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=PromptRepresentation, + extra_params={"structured_output_mode": "json_object"}, + ) + + assert isinstance(result.content, PromptRepresentation) + assert result.content.explicit == [] + # Usage from the (empty) response is preserved, not zeroed. + assert result.input_tokens == 10 + + +@pytest.mark.asyncio +async def test_structured_output_json_object_empty_content_required_fields() -> None: + """Empty content for a required-field model falls back to empty string content + instead of raising (empty_structured_output can't build the instance).""" + client = Mock() + client.chat.completions.create = AsyncMock( + return_value=_structured_create_return("") + ) + + backend = OpenAIBackend(client) + result = await backend.complete( + model="glm-4.6", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=_StructuredResponse, # has a required field + extra_params={"structured_output_mode": "json_object"}, + ) + + assert result.content == "" + + +@pytest.mark.asyncio +async def test_structured_output_json_object_mode_does_not_mutate_messages() -> None: + """The schema-injection helper must copy, never mutate the caller's list.""" + client = Mock() + client.chat.completions.create = AsyncMock( + return_value=_structured_create_return('{"answer": "ok"}') + ) + + backend = OpenAIBackend(client) + original_messages = [{"role": "user", "content": "Hello"}] + await backend.complete( + model="glm-4.6", + messages=original_messages, + max_tokens=100, + response_format=_StructuredResponse, + extra_params={"structured_output_mode": "json_object"}, + ) + + assert original_messages == [{"role": "user", "content": "Hello"}] + + +@pytest.mark.asyncio +async def test_stream_structured_output_default_mode_uses_json_schema() -> None: + """Streaming in default mode converts the model to a json_schema dict.""" + client = Mock() + client.chat.completions.create = AsyncMock(return_value=_empty_stream()) + + backend = OpenAIBackend(client) + async for _ in backend.stream( + model="gpt-4.1", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=_StructuredResponse, + ): + pass + + call = _await_kwargs(client.chat.completions.create) + assert call["response_format"]["type"] == "json_schema" + assert call["response_format"]["json_schema"]["name"] == "_StructuredResponse" + + +@pytest.mark.asyncio +async def test_stream_structured_output_json_object_mode() -> None: + """Streaming in json_object mode requests json_object + injects the schema.""" + client = Mock() + client.chat.completions.create = AsyncMock(return_value=_empty_stream()) + + backend = OpenAIBackend(client) + async for _ in backend.stream( + model="glm-4.6", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=_StructuredResponse, + extra_params={"structured_output_mode": "json_object"}, + ): + pass + + call = _await_kwargs(client.chat.completions.create) + assert call["response_format"] == {"type": "json_object"} + system_messages = [m for m in call["messages"] if m["role"] == "system"] + assert system_messages + assert "JSON" in system_messages[0]["content"] + assert "json" in system_messages[0]["content"] + + +AGENT_TOOL = { + "name": "search", + "description": "Search for information", + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, +} + + +@pytest.mark.asyncio +async def test_openai_backend_structured_with_tools_uses_create_not_parse() -> None: + """With tools + response_format, complete() must route through create() + with an explicit json_schema response_format: parse() raises client-side + on non-strict function tools, and our agent tools are deliberately + non-strict (see _convert_tools).""" + client = Mock() + client.chat.completions.create = AsyncMock( + return_value=_structured_create_return('{"answer":"ok"}') + ) + client.chat.completions.parse = AsyncMock( + side_effect=AssertionError("parse() must not be called with tools") + ) + + backend = OpenAIBackend(client) + result = await backend.complete( + model="gpt-5.4-mini", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + tools=[AGENT_TOOL], + response_format=_StructuredResponse, + ) + + assert isinstance(result.content, _StructuredResponse) + assert result.content.answer == "ok" + client.chat.completions.parse.assert_not_awaited() + call = _await_kwargs(client.chat.completions.create) + assert call["response_format"]["type"] == "json_schema" + assert ( + call["response_format"]["json_schema"]["schema"] + == _StructuredResponse.model_json_schema() + ) + # Tools stay non-strict; strictness was the whole reason to avoid parse(). + assert "strict" not in call["tools"][0]["function"] + + +@pytest.mark.asyncio +async def test_openai_backend_structured_with_tools_skips_parsing_tool_call_turn() -> ( + None +): + """A tool-call turn under tools + response_format must not attempt JSON + parsing (its content is empty and _parse_or_repair raises on that).""" + client = Mock() + client.chat.completions.create = AsyncMock( + return_value=SimpleNamespace( + choices=[ + SimpleNamespace( + finish_reason="tool_calls", + message=SimpleNamespace( + content=None, + tool_calls=[ + SimpleNamespace( + id="tool_1", + function=SimpleNamespace( + name="search", + arguments='{"query": "honcho"}', + ), + ) + ], + reasoning_details=[], + refusal=None, + ), + ) + ], + usage=SimpleNamespace( + prompt_tokens=10, + completion_tokens=5, + prompt_tokens_details=None, + ), + ) + ) + + backend = OpenAIBackend(client) + result = await backend.complete( + model="gpt-5.4-mini", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + tools=[AGENT_TOOL], + response_format=_StructuredResponse, + ) + + assert result.content == "" # raw empty text, not a parsed model + assert result.tool_calls[0].name == "search" + assert result.tool_calls[0].input == {"query": "honcho"} + + +@pytest.mark.asyncio +async def test_openai_backend_structured_without_tools_still_uses_parse() -> None: + """Tool-less structured calls (deriver, final synthesis) keep parse().""" + parsed = _StructuredResponse(answer="ok") + client = Mock() + client.chat.completions.parse = AsyncMock( + return_value=_structured_create_return('{"answer":"ok"}', parsed=parsed) + ) + client.chat.completions.create = AsyncMock( + side_effect=AssertionError("create() must not be called without tools") + ) + + backend = OpenAIBackend(client) + result = await backend.complete( + model="gpt-5.4-mini", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=100, + response_format=_StructuredResponse, + ) + + assert result.content is parsed + client.chat.completions.create.assert_not_awaited() diff --git a/tests/llm/test_capture.py b/tests/llm/test_capture.py new file mode 100644 index 00000000..7dc22e95 --- /dev/null +++ b/tests/llm/test_capture.py @@ -0,0 +1,488 @@ +"""Tests for the single-capture content layer (src/llm/capture.py).""" + +from __future__ import annotations + +from collections.abc import AsyncIterator + +import pytest + +from src.llm import capture +from src.llm.backend import CompletionResult, ToolCallResult +from src.llm.capture import ( + CapturedLLMCall, + build_captured_call, + build_captured_messages, + canonical_json, + clip_for_trace, + compute_content_hash, +) +from src.llm.types import ( + HonchoLLMCallStreamChunk, + LLMTelemetryContext, + StreamingResponseWithMetadata, +) + + +async def _chunks( + texts: list[str], *, raise_after: BaseException | None = None +) -> AsyncIterator[HonchoLLMCallStreamChunk]: + for text in texts: + yield HonchoLLMCallStreamChunk(content=text) + if raise_after is not None: + raise raise_after + + +def _wrapper( + stream: AsyncIterator[HonchoLLMCallStreamChunk], + recorder: list[tuple[str, str]], +): + return StreamingResponseWithMetadata( + stream=stream, + tool_calls_made=[], + input_tokens=0, + output_tokens=0, + cache_creation_input_tokens=0, + cache_read_input_tokens=0, + capture_finalizer=lambda text, reason: recorder.append((text, reason)), + ) + + +class TestStreamingCaptureFinalizer: + async def test_clean_drain_captures_stop(self): + recorded: list[tuple[str, str]] = [] + wrapper = _wrapper(_chunks(["hel", "lo"]), recorded) + async for _ in wrapper: + pass + assert recorded == [("hello", "stop")] + + async def test_error_drain_captures_error_and_partial_text(self): + recorded: list[tuple[str, str]] = [] + wrapper = _wrapper(_chunks(["par"], raise_after=RuntimeError("boom")), recorded) + with pytest.raises(RuntimeError): + async for _ in wrapper: + pass + # Partial text still captured, tagged error. + assert recorded == [("par", "error")] + + async def test_cancelled_drain_captures_cancelled(self): + import asyncio + + recorded: list[tuple[str, str]] = [] + wrapper = _wrapper( + _chunks(["x"], raise_after=asyncio.CancelledError()), recorded + ) + with pytest.raises(asyncio.CancelledError): + async for _ in wrapper: + pass + assert recorded == [("x", "cancelled")] + + +class TestContentHash: + def test_is_deterministic_and_prefixed(self): + h1 = compute_content_hash("user", "hello", None) + h2 = compute_content_hash("user", "hello", None) + assert h1 == h2 + assert h1.startswith("sha256:") + + def test_role_is_inside_the_hash(self): + # Identical text under different roles must never collide — role lives + # inside the hash, closing the role-in-hash collision bug. + assert compute_content_hash("user", "hi", None) != compute_content_hash( + "assistant", "hi", None + ) + + def test_tool_call_id_is_inside_the_hash(self): + assert compute_content_hash("tool", "ok", "call_1") != compute_content_hash( + "tool", "ok", "call_2" + ) + + def test_canonical_json_is_order_independent(self): + assert canonical_json({"a": 1, "b": 2}) == canonical_json({"b": 2, "a": 1}) + + +class TestClipForTrace: + def test_leaves_small_content_untouched(self): + content, truncated = clip_for_trace("short") + assert content == "short" + assert truncated is False + + def test_clips_oversized_string(self, monkeypatch: pytest.MonkeyPatch): + from src.config import settings + + monkeypatch.setattr(settings.TELEMETRY, "TRACE_MAX_BYTES", 32) + content, truncated = clip_for_trace("x" * 1000) + assert truncated is True + assert content.endswith("…[truncated]") + assert len(content.encode("utf-8")) <= settings.TELEMETRY.TRACE_MAX_BYTES + + def test_leaves_structured_content_intact(self, monkeypatch: pytest.MonkeyPatch): + from src.config import settings + + monkeypatch.setattr(settings.TELEMETRY, "TRACE_MAX_BYTES", 4) + blocks = [{"type": "text", "text": "a long block of structured content"}] + content, truncated = clip_for_trace(blocks) + assert content == blocks + assert truncated is False + + +class TestBuildCapturedMessages: + def test_hashes_each_message(self): + messages = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + ] + captured, truncated = build_captured_messages(messages, memo=None) + assert [m.role for m in captured] == ["user", "assistant"] + assert all(m.content_hash.startswith("sha256:") for m in captured) + assert truncated is False + + def test_memo_makes_hashing_on(self, monkeypatch: pytest.MonkeyPatch): + # The conversation is append-only and message dicts are reused, so with + # a shared memo each message is hashed exactly once across iterations. + calls = {"n": 0} + real = compute_content_hash + + def counting( + role: str, + content: object, + tool_call_id: str | None, + tool_calls: list[dict[str, object]] | None = None, + ) -> str: + calls["n"] += 1 + return real(role, content, tool_call_id, tool_calls) + + monkeypatch.setattr(capture, "compute_content_hash", counting) + + m1 = {"role": "user", "content": "q1"} + m2 = {"role": "assistant", "content": "a1"} + m3 = {"role": "user", "content": "q2"} + memo: dict[int, capture.CapturedMessage] = {} + + build_captured_messages([m1, m2], memo) + assert calls["n"] == 2 # both hashed + build_captured_messages([m1, m2, m3], memo) + assert calls["n"] == 3 # only the newly-appended m3 hashed (not re-hashed) + + +class TestNormalizeToolCalls: + """Tool calls live outside `content` for openai/gemini — capture must lift + them into the unified `tool_calls` shape (the PR concern).""" + + def test_openai_assistant_tool_calls_captured(self): + msg = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "search_memory", + "arguments": '{"query": "coffee"}', + }, + } + ], + } + captured, _ = build_captured_messages([msg], memo=None, transport="openai") + assert captured[0].tool_calls == [ + {"id": "call_1", "name": "search_memory", "input": {"query": "coffee"}} + ] + + def test_gemini_model_parts_captured(self): + msg = { + "role": "model", + "parts": [ + {"text": "let me look"}, + {"function_call": {"name": "grep_messages", "args": {"text": "x"}}}, + ], + } + captured, _ = build_captured_messages([msg], memo=None, transport="gemini") + assert captured[0].content == "let me look" + assert captured[0].tool_calls == [ + {"id": None, "name": "grep_messages", "input": {"text": "x"}} + ] + + def test_gemini_tool_result_recovered(self): + # Gemini tool results live in `parts` (no `content` key) and were dropped. + msg = { + "role": "user", + "parts": [ + { + "function_response": { + "name": "grep_messages", + "response": {"result": "3 hits"}, + } + } + ], + } + captured, _ = build_captured_messages([msg], memo=None, transport="gemini") + assert captured[0].content == "3 hits" + assert captured[0].tool_call_id == "grep_messages" + + def test_anthropic_tool_use_blocks_normalized(self): + msg = { + "role": "assistant", + "content": [ + {"type": "text", "text": "searching"}, + { + "type": "tool_use", + "id": "tu_1", + "name": "search_memory", + "input": {"q": "x"}, + }, + ], + } + captured, _ = build_captured_messages([msg], memo=None, transport="anthropic") + assert captured[0].content == "searching" + assert captured[0].tool_calls == [ + {"id": "tu_1", "name": "search_memory", "input": {"q": "x"}} + ] + + def test_hash_distinguishes_tool_calls(self): + # Two empty-content assistant turns with different tool calls must not + # collide in the dedup store (they did before tool_calls entered the hash). + base = {"role": "assistant", "content": None} + a = { + **base, + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "search_memory", "arguments": "{}"}, + } + ], + } + b = { + **base, + "tool_calls": [ + { + "id": "c2", + "type": "function", + "function": {"name": "search_messages", "arguments": "{}"}, + } + ], + } + (ca,), _ = build_captured_messages([a], memo=None, transport="openai") + (cb,), _ = build_captured_messages([b], memo=None, transport="openai") + assert ca.content_hash != cb.content_hash + + +class TestBuildCapturedCall: + def test_maps_telemetry_and_result(self): + telemetry = LLMTelemetryContext( + workspace_name="ws", + call_purpose="dialectic.answer", + parent_category="dialectic", + run_id="r1", + trace_id="r1", + span_id="r1", + session_id="sess_abc", + iteration=2, + step_seq=2, + ) + result = CompletionResult( + content="answer", + input_tokens=10, + output_tokens=5, + finish_reason="stop", + tool_calls=[ToolCallResult(id="t1", name="search", input={"q": "x"})], + ) + call = build_captured_call( + telemetry=telemetry, + transport="anthropic", + provider_label=None, + model="claude-x", + messages=[{"role": "user", "content": "q"}], + tools=None, + tool_choice=None, + result=result, + attempt=1, + was_fallback=False, + was_stream=False, + finish_reason="stop", + ) + assert isinstance(call, CapturedLLMCall) + assert call.trace_id == "r1" and call.span_id == "r1" + assert call.iteration == 2 and call.step_seq == 2 + assert call.output_content == "answer" + assert call.output_tool_calls == [ + {"id": "t1", "name": "search", "input": {"q": "x"}} + ] + assert call.input_tokens == 10 and call.output_tokens == 5 + assert call.session_id == "sess_abc" + assert len(call.input_messages) == 1 + assert call.input_messages[0].content_hash.startswith("sha256:") + + def test_session_id_defaults_none_without_telemetry(self): + # Sessionless calls (and the no-telemetry path) carry session_id=None so + # the Langfuse projection emits no session grouping for them. + telemetry = LLMTelemetryContext(run_id="r1", trace_id="r1", span_id="r1") + call = build_captured_call( + telemetry=telemetry, + transport="anthropic", + provider_label=None, + model="claude-x", + messages=[{"role": "user", "content": "q"}], + tools=None, + tool_choice=None, + result=CompletionResult(content="a", finish_reason="stop"), + attempt=1, + was_fallback=False, + was_stream=False, + finish_reason="stop", + ) + assert call.session_id is None + + def test_self_parent_is_normalized_to_none(self): + # The tool loop sets parent_span_id == span_id on the run span (it + # doubles as the Langfuse "inside a run" signal). A span that is its own + # parent is a root, so the EXPORTED parent_span_id must be None — else + # span-tree consumers file the root as a child of itself. + telemetry = LLMTelemetryContext( + run_id="r1", trace_id="r1", span_id="r1", parent_span_id="r1" + ) + call = build_captured_call( + telemetry=telemetry, + transport="anthropic", + provider_label=None, + model="claude-x", + messages=[{"role": "user", "content": "q"}], + tools=None, + tool_choice=None, + result=CompletionResult(content="a", finish_reason="stop"), + attempt=1, + was_fallback=False, + was_stream=False, + finish_reason="stop", + ) + assert call.span_id == "r1" + assert call.parent_span_id is None + # A genuine distinct parent is preserved. + telemetry.parent_span_id = "parent-span" + assert telemetry.exported_parent_span_id() == "parent-span" + + def test_error_path_collapses_output(self): + call = build_captured_call( + telemetry=None, + transport="anthropic", + provider_label=None, + model="claude-x", + messages=[{"role": "user", "content": "q"}], + tools=None, + tool_choice=None, + result=None, + attempt=2, + was_fallback=True, + was_stream=False, + finish_reason="error", + ) + assert call.output_content is None + assert call.output_tool_calls == [] + assert call.finish_reason == "error" + assert call.attempt == 2 and call.was_fallback is True + + +class TestExporterRegistry: + def test_register_dispatch_and_clear(self): + capture.clear_exporters() + assert capture.has_exporters() is False + seen: list[CapturedLLMCall] = [] + + class _Spy: + def export(self, call: CapturedLLMCall) -> None: + seen.append(call) + + capture.register_exporter(_Spy()) + assert capture.has_exporters() is True + + call = build_captured_call( + telemetry=None, + transport="anthropic", + provider_label=None, + model="m", + messages=[], + tools=None, + tool_choice=None, + result=None, + attempt=1, + was_fallback=False, + was_stream=False, + finish_reason="stop", + ) + capture.dispatch_captured_call(call) + assert seen == [call] + capture.clear_exporters() + assert capture.has_exporters() is False + + def test_dispatch_swallows_exporter_errors(self): + capture.clear_exporters() + + class _Boom: + def export(self, call: CapturedLLMCall) -> None: + raise RuntimeError(f"nope: {call.model}") + + capture.register_exporter(_Boom()) + call = build_captured_call( + telemetry=None, + transport="anthropic", + provider_label=None, + model="m", + messages=[], + tools=None, + tool_choice=None, + result=None, + attempt=1, + was_fallback=False, + was_stream=False, + finish_reason="stop", + ) + # Must not raise — telemetry never breaks the LLM path. + capture.dispatch_captured_call(call) + capture.clear_exporters() + + +class TestThoughtSignatureSerialization: + """Gemini `thought_signature` is bytes; it must not break trace serialization.""" + + def test_bytes_signature_base64_encoded_and_serializes(self): + import json + + from src.telemetry.events.trace import LLMCallTracedEvent + + result = CompletionResult( + content=None, + finish_reason="STOP", + tool_calls=[ + ToolCallResult( + id="call_1", + name="grep_messages", + input={"text": "coffee"}, + thought_signature=b"\x0a\x1f\x88\xff\x00sig", + ) + ], + ) + call = build_captured_call( + telemetry=LLMTelemetryContext(trace_id="t1", span_id="s1"), + transport="gemini", + provider_label=None, + model="gemini-2.5-flash", + messages=[{"role": "user", "content": "q"}], + tools=None, + tool_choice=None, + result=result, + attempt=1, + was_fallback=False, + was_stream=False, + finish_reason="STOP", + ) + sig = call.output_tool_calls[0]["thought_signature"] + assert isinstance(sig, str) # base64, not raw bytes + + # The traced event must serialize to JSON without raising (the emit path + # calls model_dump(mode="json"), which threw UnicodeDecodeError on bytes). + event = LLMCallTracedEvent( + model="gemini-2.5-flash", + transport="gemini", + output_tool_calls=call.output_tool_calls, + ) + json.dumps(event.model_dump(mode="json")) diff --git a/tests/llm/test_embedding_client.py b/tests/llm/test_embedding_client.py index fcee66d9..a4642159 100644 --- a/tests/llm/test_embedding_client.py +++ b/tests/llm/test_embedding_client.py @@ -339,3 +339,108 @@ def test_resolve_send_dimensions_never_returns_false_regardless( monkeypatch, ) assert s.resolve_send_dimensions() is False + + +@pytest.mark.asyncio +async def test_simple_batch_embed_respects_token_budget_per_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """simple_batch_embed must split inputs across requests so per-request token cap holds.""" + fake_embeddings = FakeOpenAIEmbeddingsAPI([0.5] * 4) + + class FakeOpenAIClient: + def __init__(self, *, api_key: str | None, base_url: str | None) -> None: + self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings + + monkeypatch.setattr("src.embedding_client.AsyncOpenAI", FakeOpenAIClient) + + # max_input_tokens=100 per single input; max_tokens_per_request=120 total, + # so two ~80-token inputs must end up in *separate* requests. + client = _EmbeddingClient( + EmbeddingModelConfig( + transport="openai", + model="text-embedding-3-small", + api_key="test-key", + base_url=None, + ), + vector_dimensions=4, + max_input_tokens=100, + max_tokens_per_request=120, + send_dimensions=False, + ) + + # "word " * 80 produces ~80 tokens with cl100k_base/the model encoding. + long_a = ("alpha " * 80).strip() + long_b = ("beta " * 80).strip() + + out = await client.simple_batch_embed([long_a, long_b]) + assert len(out) == 2 + # Per-request token cap forces two separate requests. + assert len(fake_embeddings.calls) == 2 + + +@pytest.mark.asyncio +async def test_simple_batch_embed_rejects_oversized_input( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Inputs that exceed max_embedding_tokens must raise ValueError immediately.""" + fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 4) + + class FakeOpenAIClient: + def __init__(self, *, api_key: str | None, base_url: str | None) -> None: + self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings + + monkeypatch.setattr("src.embedding_client.AsyncOpenAI", FakeOpenAIClient) + + client = _EmbeddingClient( + EmbeddingModelConfig( + transport="openai", + model="text-embedding-3-small", + api_key="test-key", + base_url=None, + ), + vector_dimensions=4, + max_input_tokens=10, + max_tokens_per_request=1000, + send_dimensions=False, + ) + + too_long = ("word " * 50).strip() + with pytest.raises(ValueError, match="maximum token limit"): + await client.simple_batch_embed([too_long]) + + +def test_prepare_chunks_returns_ordered_chunks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """prepare_chunks must split oversized inputs using the same rules as batch_embed.""" + fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 4) + + class FakeOpenAIClient: + def __init__(self, *, api_key: str | None, base_url: str | None) -> None: + self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings + + monkeypatch.setattr("src.embedding_client.AsyncOpenAI", FakeOpenAIClient) + + client = _EmbeddingClient( + EmbeddingModelConfig( + transport="openai", + model="text-embedding-3-small", + api_key="test-key", + base_url=None, + ), + vector_dimensions=4, + max_input_tokens=10, + max_tokens_per_request=1000, + send_dimensions=False, + ) + + short_text = "hello" + long_text = ("word " * 50).strip() + + out = client.prepare_chunks({"short": short_text, "long": long_text}) + + assert out["short"] == [short_text] + assert len(out["long"]) > 1 + # Order preserved + assert isinstance(out["long"][0], str) diff --git a/tests/llm/test_langfuse_trace_annotation.py b/tests/llm/test_langfuse_trace_annotation.py new file mode 100644 index 00000000..2f1ec50a --- /dev/null +++ b/tests/llm/test_langfuse_trace_annotation.py @@ -0,0 +1,562 @@ +# pyright: reportPrivateUsage=false, reportUnusedParameter=false +"""Tests for the Langfuse session/trace wiring in `src/llm/runtime.py`. + +One agentic run = one trace; `session_id = run_id` (globally unique, so it's +conflict-free across tenants — unlike the Honcho session name). The run handle +(`start_langfuse_agent_run`) opens an `as_type="span"` root and keeps it +current via an ``ExitStack`` until `.end()`. Step spans + nested generations +nest under the run while it's open. A single-call agent (deriver, summarizer) +gets no run handle and self-stamps its lone generation as the trace root. +Disabled (no LANGFUSE_PUBLIC_KEY) → no Langfuse calls at all. +""" + +from __future__ import annotations + +import contextlib +from typing import Any + +import pytest + +from src.config import settings +from src.llm import runtime +from src.llm.types import LLMTelemetryContext + + +@pytest.fixture +def capture_propagate(monkeypatch: pytest.MonkeyPatch): + """Stub `langfuse.propagate_attributes`, capturing the kwargs it's called with. + + Returns a dict that's empty until propagate_attributes is invoked. + """ + captured: dict[str, Any] = {} + + @contextlib.contextmanager + def fake_propagate(**kwargs: Any): + captured.clear() + captured.update(kwargs) + yield + + import langfuse + + monkeypatch.setattr(langfuse, "propagate_attributes", fake_propagate) + return captured + + +@pytest.fixture +def langfuse_client(monkeypatch: pytest.MonkeyPatch): + """Stub `langfuse.get_client()`, capturing observation/generation/span calls. + + Returns ``{"observation", "generation", "span", "run_span"}`` — each sub-dict + stays empty until the corresponding call is made: + - ``observation``: the span opened via `start_as_current_observation` + (run root or step span). + - ``generation``: the rename via `update_current_generation`. + - ``span``: the step I/O via `update_current_span`. + - ``run_span``: I/O written onto the span object handed back by + `start_as_current_observation` (merged, mirroring Langfuse's update + semantics). + """ + captured: dict[str, dict[str, Any]] = { + "observation": {}, + "generation": {}, + "span": {}, + "run_span": {}, + } + + class FakeSpan: + def update(self, **kwargs: Any) -> None: + # Merge (don't clear): Langfuse's span.update accumulates, so input + # set at run start and output set at run end coexist. + captured["run_span"].update(kwargs) + + @contextlib.contextmanager + def fake_observation(**kwargs: Any): + captured["observation"].clear() + captured["observation"].update(kwargs) + yield FakeSpan() + + class FakeClient: + def start_as_current_observation(self, **kwargs: Any): + return fake_observation(**kwargs) + + def update_current_generation(self, **kwargs: Any) -> None: + captured["generation"].clear() + captured["generation"].update(kwargs) + + def update_current_span(self, **kwargs: Any) -> None: + captured["span"].clear() + captured["span"].update(kwargs) + + import langfuse + + monkeypatch.setattr(langfuse, "get_client", lambda: FakeClient()) + return captured + + +@pytest.fixture +def langfuse_enabled(monkeypatch: pytest.MonkeyPatch): + """Turn the integration on with a known NAMESPACE (the tenant / user_id). + + Pins LANGFUSE_EXPORTER_MODE='inline' — this module tests the legacy inline + span machinery, which is gated to inline mode (the default is now 'exporter', + where these functions no-op in favor of the LangfuseExporter).""" + monkeypatch.setattr(settings, "LANGFUSE_PUBLIC_KEY", "pk-test") + monkeypatch.setattr(settings, "LANGFUSE_EXPORTER_MODE", "inline") + monkeypatch.setattr(settings, "NAMESPACE", "acme-tenant") + + +class TestAnnotateDisabled: + def test_noop_when_key_unset( + self, monkeypatch: pytest.MonkeyPatch, capture_propagate: dict[str, Any] + ): + monkeypatch.setattr(settings, "LANGFUSE_PUBLIC_KEY", None) + + runtime.annotate_current_langfuse_trace( + "anthropic", + "claude-x", + telemetry=LLMTelemetryContext(run_id="run-abc"), + ) + + assert capture_propagate == {} + + +class TestAnnotateInsideRun: + """A generation nested under a run (it carries a `parent_span_id`): the run + owns the trace attrs, so this call must NOT propagate. It still stamps model + + per-step metadata + name on the generation (the multi-turn regression fix — + provider/model used to be dropped on every iteration after the first). + Nesting is derived from the explicit `parent_span_id`, not a contextvar.""" + + def test_nested_generation_does_not_propagate( + self, + langfuse_enabled: None, + capture_propagate: dict[str, Any], + langfuse_client: dict[str, dict[str, Any]], + ): + telemetry = LLMTelemetryContext( + workspace_name="ws1", + call_purpose="dialectic.answer", + agent_type="dialectic", + run_id="run-abc", + span_id="run-abc", + # A non-null parent_span_id is what marks this generation as nested + # under the run span (replaces the old `_in_agent_run` contextvar). + parent_span_id="run-abc", + iteration=2, + peer_name="alice", + track_name="Dialectic Agent", + ) + + runtime.annotate_current_langfuse_trace( + "anthropic", "claude-x", telemetry=telemetry + ) + + # Run handle owns user_id/session_id/trace_name — re-propagating here + # would clobber the run's session, so we don't propagate at all. + assert capture_propagate == {} + # Per-call generation: name + model + step metadata stamped every + # iteration (formerly only name was stamped, dropping provider/model). + gen = langfuse_client["generation"] + assert gen["name"] == "Dialectic Agent LLM call" + assert gen["model"] == "claude-x" + assert gen["metadata"]["provider"] == "anthropic" + assert gen["metadata"]["model"] == "claude-x" + assert gen["metadata"]["iteration"] == "2" + assert gen["metadata"]["agent_type"] == "dialectic" + + +class TestAnnotateOwnTraceRoot: + """A generation that IS its own trace root stamps the trace attributes: + single calls (no run_id → no session) and the no-telemetry case.""" + + def test_single_call_has_no_session_and_names_trace( + self, + langfuse_enabled: None, + capture_propagate: dict[str, Any], + langfuse_client: dict[str, dict[str, Any]], + ): + telemetry = LLMTelemetryContext( + workspace_name="ws1", + call_purpose="deriver.representation", + observed="bob", + track_name="Minimal Deriver", + ) + + runtime.annotate_current_langfuse_trace( + "gemini", "gemini-x", telemetry=telemetry + ) + + assert capture_propagate["session_id"] is None + assert capture_propagate["user_id"] == "acme-tenant" + # Single-call: this generation IS the trace root, so it names the trace. + assert capture_propagate["trace_name"] == "Minimal Deriver" + assert capture_propagate["metadata"]["observed"] == "bob" + # The generation observation is still named per agent+action. + assert langfuse_client["generation"]["name"] == "Minimal Deriver LLM call" + + def test_no_telemetry_still_stamps_user_id( + self, + langfuse_enabled: None, + capture_propagate: dict[str, Any], + langfuse_client: dict[str, dict[str, Any]], + ): + runtime.annotate_current_langfuse_trace("openai", "gpt-x", telemetry=None) + + assert capture_propagate["user_id"] == "acme-tenant" + assert capture_propagate["session_id"] is None + assert capture_propagate["trace_name"] is None + assert capture_propagate["metadata"]["provider"] == "openai" + # No telemetry → no per-agent generation name, but provider/model still set. + gen = langfuse_client["generation"] + assert gen["name"] is None + assert gen["model"] == "gpt-x" + + +class TestAgentRun: + """`start_langfuse_agent_run` returns an imperative handle: opens an + ``as_type="span"`` root, stamps ``session_id = run_id`` via + ``propagate_attributes``, and keeps the span open until ``.end()``. + Only fires for multi-turn runs (run_id present).""" + + def test_noop_when_disabled( + self, + monkeypatch: pytest.MonkeyPatch, + langfuse_client: dict[str, dict[str, Any]], + capture_propagate: dict[str, Any], + ): + monkeypatch.setattr(settings, "LANGFUSE_PUBLIC_KEY", None) + + handle = runtime.start_langfuse_agent_run( + "Dialectic Agent", LLMTelemetryContext(run_id="r1") + ) + + assert handle is None + assert langfuse_client["observation"] == {} + assert capture_propagate == {} + + def test_noop_without_run_id( + self, + langfuse_enabled: None, + langfuse_client: dict[str, dict[str, Any]], + capture_propagate: dict[str, Any], + ): + # Single-call agents (deriver/summarizer) have no run_id → no run root, + # so their LLM calls stay standalone, sessionless traces. + handle = runtime.start_langfuse_agent_run( + "Minimal Deriver", LLMTelemetryContext(workspace_name="ws1") + ) + + assert handle is None + assert langfuse_client["observation"] == {} + assert capture_propagate == {} + + def test_noop_without_telemetry( + self, + langfuse_enabled: None, + langfuse_client: dict[str, dict[str, Any]], + ): + handle = runtime.start_langfuse_agent_run("anything", None) + assert handle is None + assert langfuse_client["observation"] == {} + + def test_opens_run_root_and_owns_trace_attrs( + self, + langfuse_enabled: None, + langfuse_client: dict[str, dict[str, Any]], + capture_propagate: dict[str, Any], + ): + tele = LLMTelemetryContext( + workspace_name="ws1", + call_purpose="dialectic.answer", + agent_type="dialectic", + run_id="run-abc", + observed="bob", + track_name="Dialectic Agent", + ) + + handle = runtime.start_langfuse_agent_run("Dialectic Agent", tele) + assert handle is not None + try: + # The run IS the trace root: an as_type="span" observation whose + # name is STABLE (no step number) so Langfuse aggregates by name. + observation = langfuse_client["observation"] + assert observation["as_type"] == "span" + assert observation["name"] == "Dialectic Agent" + # Trace grouping: one Langfuse session per run, drillable per tenant. + assert capture_propagate["session_id"] == "run-abc" + assert capture_propagate["user_id"] == "acme-tenant" + assert capture_propagate["trace_name"] == "Dialectic Agent" + md = capture_propagate["metadata"] + assert md["workspace_name"] == "ws1" + assert md["agent_type"] == "dialectic" + assert md["observed"] == "bob" + # Honcho's Session is deliberately NOT the grouping key. + assert "honcho_session_id" not in md + finally: + handle.end() + + def test_run_keyed_on_span_id_and_nesting_via_parent_span_id( + self, + langfuse_enabled: None, + langfuse_client: dict[str, dict[str, Any]], + capture_propagate: dict[str, Any], + ): + # The `_in_agent_run` contextvar is retired — nesting is now derived + # from the explicit `parent_span_id` field on the telemetry context. + assert not hasattr(runtime, "_in_agent_run") + + # The run handle opens keyed on span_id (falling back to run_id). + handle = runtime.start_langfuse_agent_run( + "Dialectic Agent", + LLMTelemetryContext( + run_id="r1", span_id="r1", track_name="Dialectic Agent" + ), + ) + assert handle is not None + handle.end() + + # A root call (no parent_span_id) propagates trace attrs; a nested call + # (parent_span_id set) stays silent. + capture_propagate.clear() + runtime.annotate_current_langfuse_trace( + "anthropic", + "claude-x", + telemetry=LLMTelemetryContext(run_id="r2", span_id="r2"), + ) + assert capture_propagate.get("session_id") == "r2" + + capture_propagate.clear() + runtime.annotate_current_langfuse_trace( + "anthropic", + "claude-x", + telemetry=LLMTelemetryContext( + run_id="r2", span_id="r2", parent_span_id="r2" + ), + ) + assert capture_propagate == {} + + def test_end_is_idempotent( + self, + langfuse_enabled: None, + langfuse_client: dict[str, dict[str, Any]], + capture_propagate: dict[str, Any], + ): + # The streaming wrapper may call .end() after the api.py finally already + # called it (or vice-versa); the handle has to tolerate that. + handle = runtime.start_langfuse_agent_run( + "Dialectic Agent", LLMTelemetryContext(run_id="r1") + ) + assert handle is not None + handle.end() + handle.end() # must not raise + + +class TestAgentRunIO: + """The run handle exposes `.update(input=..., output=...)` for stamping + the run-root span — the trace's input/output preview in the Langfuse UI. + A second call merges into the first (Langfuse's update semantics).""" + + def test_sets_input_then_output_on_handle( + self, + langfuse_enabled: None, + langfuse_client: dict[str, dict[str, Any]], + capture_propagate: dict[str, Any], + ): + messages = [{"role": "user", "content": "How many coffees?"}] + handle = runtime.start_langfuse_agent_run( + "Dialectic Agent", LLMTelemetryContext(run_id="run-abc") + ) + assert handle is not None + try: + handle.update(input=messages) + finally: + handle.end(output="You bought 4 coffees.") + + assert langfuse_client["run_span"]["input"] == messages + assert langfuse_client["run_span"]["output"] == "You bought 4 coffees." + + def test_end_without_output_leaves_output_unset( + self, + langfuse_enabled: None, + langfuse_client: dict[str, dict[str, Any]], + capture_propagate: dict[str, Any], + ): + handle = runtime.start_langfuse_agent_run( + "Dialectic Agent", LLMTelemetryContext(run_id="run-abc") + ) + assert handle is not None + handle.update(input=[{"role": "user"}]) + handle.end() + + assert "input" in langfuse_client["run_span"] + # Only input was passed → output is not written (so it isn't blanked). + assert "output" not in langfuse_client["run_span"] + + +class TestAgentStep: + """`start_langfuse_agent_step` opens a per-iteration child span under the + run root (one reasoning turn). Unlike the run handle, it does NOT touch + trace attributes.""" + + def test_noop_when_disabled( + self, + monkeypatch: pytest.MonkeyPatch, + langfuse_client: dict[str, dict[str, Any]], + capture_propagate: dict[str, Any], + ): + monkeypatch.setattr(settings, "LANGFUSE_PUBLIC_KEY", None) + + step = runtime.start_langfuse_agent_step( + "Dialectic Agent step", LLMTelemetryContext(run_id="r1") + ) + + assert step is None + assert langfuse_client["observation"] == {} + assert capture_propagate == {} + + def test_noop_without_run_id( + self, + langfuse_enabled: None, + langfuse_client: dict[str, dict[str, Any]], + capture_propagate: dict[str, Any], + ): + step = runtime.start_langfuse_agent_step( + "Minimal Deriver step", LLMTelemetryContext(workspace_name="ws1") + ) + + assert step is None + assert langfuse_client["observation"] == {} + assert capture_propagate == {} + + def test_noop_without_telemetry( + self, + langfuse_enabled: None, + langfuse_client: dict[str, dict[str, Any]], + ): + step = runtime.start_langfuse_agent_step("anything", None) + assert step is None + assert langfuse_client["observation"] == {} + + def test_opens_child_span_without_propagating( + self, + langfuse_enabled: None, + langfuse_client: dict[str, dict[str, Any]], + capture_propagate: dict[str, Any], + ): + tele = LLMTelemetryContext( + workspace_name="ws1", + agent_type="dialectic", + run_id="run-abc", + iteration=2, + observed="bob", + track_name="Dialectic Agent", + ) + + step = runtime.start_langfuse_agent_step("Dialectic Agent step", tele) + assert step is not None + try: + observation = langfuse_client["observation"] + assert observation["as_type"] == "span" + assert observation["name"] == "Dialectic Agent step" + # The per-step index rides on the span's metadata (str-coerced). + assert observation["metadata"]["iteration"] == "2" + assert observation["metadata"]["observed"] == "bob" + # The run root owns the trace attrs — the step must NOT propagate. + assert capture_propagate == {} + finally: + step.end() + + +class TestStepIO: + """`step.annotate_io` stamps this turn's I/O on the step span — without it, + only the nested generation would carry I/O and the step would show blank.""" + + def test_text_answer_sets_input_and_output( + self, + langfuse_enabled: None, + langfuse_client: dict[str, dict[str, Any]], + ): + messages = [{"role": "user", "content": "What is the user's name?"}] + step = runtime.start_langfuse_agent_step( + "Dialectic Agent step", LLMTelemetryContext(run_id="run-abc") + ) + assert step is not None + try: + step.annotate_io(messages, "The user's name is Jordan.", []) + finally: + step.end() + + # The step span's input/output is stamped via the handle's underlying + # span.update() — captured on the run_span fixture key. + assert langfuse_client["run_span"]["input"] == messages + assert langfuse_client["run_span"]["output"] == "The user's name is Jordan." + + def test_tool_calling_turn_summarizes_tools_as_output( + self, + langfuse_enabled: None, + langfuse_client: dict[str, dict[str, Any]], + ): + # A tool-calling turn has no text yet — the step's "output" is the set + # of tools it chose, by name (per-tool I/O lives on the tool children). + step = runtime.start_langfuse_agent_step( + "Dialectic Agent step", LLMTelemetryContext(run_id="run-abc") + ) + assert step is not None + try: + step.annotate_io( + [{"role": "user", "content": "how many coffees?"}], + "", + [{"name": "grep_messages"}, {"name": "search_memory"}], + ) + finally: + step.end() + + assert langfuse_client["run_span"]["output"] == { + "tool_calls": ["grep_messages", "search_memory"] + } + + +class TestAnnotateGenerationIOGating: + """`annotate_current_generation_io` writes to the ACTIVE @observe generation + span (the `conditional_observe` wrapper), which only exists in inline mode. + In exporter mode — the default — there is no active span, so calling + `update_current_generation()` would make the Langfuse SDK log "No active span + in current context" on every LLM call. The helper must therefore no-op in + exporter mode (the LangfuseExporter projects I/O from the captured stream). + Regression guard for the gate that was on LANGFUSE_PUBLIC_KEY instead of + langfuse_inline_enabled.""" + + def test_noops_in_exporter_mode_even_with_key( + self, + monkeypatch: pytest.MonkeyPatch, + langfuse_client: dict[str, dict[str, Any]], + ): + # Key present but exporter mode (the production default). + monkeypatch.setattr(settings, "LANGFUSE_PUBLIC_KEY", "pk-test") + monkeypatch.setattr(settings, "LANGFUSE_EXPORTER_MODE", "exporter") + + runtime.annotate_current_generation_io( + input=[{"role": "user", "content": "hi"}], + output="hello", + usage_details={"input": 1, "output": 1}, + ) + + # No active generation span in exporter mode → must not touch it. + assert langfuse_client["generation"] == {} + + def test_writes_in_inline_mode( + self, + langfuse_enabled: None, # pins inline mode + a key + langfuse_client: dict[str, dict[str, Any]], + ): + messages = [{"role": "user", "content": "hi"}] + runtime.annotate_current_generation_io( + input=messages, + output="hello", + usage_details={"input": 1, "output": 1}, + ) + + gen = langfuse_client["generation"] + assert gen["input"] == messages + assert gen["output"] == "hello" + assert gen["usage_details"] == {"input": 1, "output": 1} diff --git a/tests/llm/test_model_config.py b/tests/llm/test_model_config.py index bfc1f392..2c31448c 100644 --- a/tests/llm/test_model_config.py +++ b/tests/llm/test_model_config.py @@ -41,6 +41,62 @@ def test_fallback_config_is_independent() -> None: assert config.fallback.base_url == "https://example.com/v1" +def test_select_model_config_for_attempt_preserves_structured_output_mode() -> None: + """The per-attempt fallback config must carry structured_output_mode. + + The operator sets it on the (independent) fallback; dropping it on the final + attempt would silently send json_schema to a provider that can't parse it. + """ + from src.config import ResolvedFallbackConfig + from src.llm.runtime import select_model_config_for_attempt + + config = ModelConfig( + model="gpt-5.4-mini", + transport="openai", + fallback=ResolvedFallbackConfig( + model="glm-4.6", + transport="openai", + structured_output_mode="json_object", + ), + ) + + # Final attempt swaps to the fallback. + selected = select_model_config_for_attempt(config, attempt=3, retry_attempts=3) + + assert selected.model == "glm-4.6" + assert selected.structured_output_mode == "json_object" + + +def test_structured_output_mode_rejected_on_non_openai_transport() -> None: + """structured_output_mode is a no-op off the openai transport — reject it.""" + with pytest.raises(ValueError, match="structured_output_mode is only supported"): + ConfiguredModelSettings( + model="claude-haiku-4-5", + transport="anthropic", + structured_output_mode="json_object", + ) + + +def test_structured_output_mode_rejected_on_non_openai_fallback() -> None: + from src.config import FallbackModelSettings + + with pytest.raises(ValueError, match="structured_output_mode is only supported"): + FallbackModelSettings( + model="gemini-2.5-pro", + transport="gemini", + structured_output_mode="json_object", + ) + + +def test_structured_output_mode_allowed_on_openai_transport() -> None: + config = ConfiguredModelSettings( + model="glm-4.6", + transport="openai", + structured_output_mode="json_object", + ) + assert config.structured_output_mode == "json_object" + + def test_base_url_is_allowed_for_any_transport() -> None: config = ModelConfig( model="claude-haiku-4-5", diff --git a/tests/llm/test_registry.py b/tests/llm/test_registry.py new file mode 100644 index 00000000..71f3302e --- /dev/null +++ b/tests/llm/test_registry.py @@ -0,0 +1,22 @@ +"""Tests for src.llm.registry helpers.""" + +from __future__ import annotations + +from src.llm.registry import _default_headers_for # pyright: ignore[reportPrivateUsage] + + +def test_default_headers_for_openrouter_base_url() -> None: + """OpenRouter base URLs get the app-attribution headers.""" + headers = _default_headers_for("https://openrouter.ai/api/v1") + assert headers["HTTP-Referer"] == "https://honcho.dev" + assert headers["X-Openrouter-Title"] == "Honcho" + + +def test_default_headers_for_non_openrouter_base_url() -> None: + """Other OpenAI-compatible providers get no extra headers.""" + assert _default_headers_for("https://api.openai.com/v1") == {} + + +def test_default_headers_for_none_base_url() -> None: + """A missing base URL (default OpenAI) gets no extra headers.""" + assert _default_headers_for(None) == {} diff --git a/tests/llm/test_telemetry_agent_iteration.py b/tests/llm/test_telemetry_agent_iteration.py index 382cc2e3..71f5c4ba 100644 --- a/tests/llm/test_telemetry_agent_iteration.py +++ b/tests/llm/test_telemetry_agent_iteration.py @@ -45,7 +45,7 @@ def _response( class TestTelemetryForIteration: def test_returns_none_when_base_is_none(self): - assert _telemetry_for_iteration(None, 1) is None + assert _telemetry_for_iteration(None, 1, step_seq=1) is None def test_returns_fresh_copy_with_iteration_set(self): base = LLMTelemetryContext( @@ -54,12 +54,13 @@ class TestTelemetryForIteration: parent_category="dialectic", agent_type="dialectic", run_id="run-xyz", + span_id="run-xyz", iteration=None, peer_name="user_peer", ) - copy_a = _telemetry_for_iteration(base, 3) - copy_b = _telemetry_for_iteration(base, 4) + copy_a = _telemetry_for_iteration(base, 3, step_seq=3) + copy_b = _telemetry_for_iteration(base, 4, step_seq=4) assert copy_a is not None and copy_b is not None assert copy_a is not base and copy_b is not base @@ -67,6 +68,9 @@ class TestTelemetryForIteration: assert base.iteration is None assert copy_a.iteration == 3 assert copy_b.iteration == 4 + # Per-step correlation is set; parent_span_id is derived from the span. + assert copy_a.step_seq == 3 + assert copy_a.parent_span_id == "run-xyz" # All other fields round-trip. assert copy_a.run_id == "run-xyz" assert copy_a.peer_name == "user_peer" diff --git a/tests/llm/test_telemetry_llm_call.py b/tests/llm/test_telemetry_llm_call.py index 817aa41e..ef3cc21c 100644 --- a/tests/llm/test_telemetry_llm_call.py +++ b/tests/llm/test_telemetry_llm_call.py @@ -13,7 +13,7 @@ Targets: from __future__ import annotations -from typing import Any +from typing import Any, cast from unittest.mock import AsyncMock, patch import pytest @@ -684,3 +684,84 @@ class TestStreamingResponseTokenWriteBack: # And we yielded every chunk to the caller — the wrapper is a # passthrough, not a sink. assert len(chunks) == 3 + + +class TestStreamingResponseRunHandleClose: + """When a `langfuse_run_handle` is transferred to the streaming wrapper, + the wrapper owns it: the accumulated streamed text is stamped as the run + span's output and the span is closed exactly once when the stream drains. + The close lives in a `finally`, so an early-exit caller still closes the + span rather than leaking it. + """ + + class _FakeRunHandle: + def __init__(self) -> None: + self.end_calls: list[Any] = [] + + def end(self, *, output: Any = None) -> None: + self.end_calls.append(output) + + @staticmethod + async def _fake_stream() -> Any: + from src.llm.types import HonchoLLMCallStreamChunk + + yield HonchoLLMCallStreamChunk(content="hel", output_tokens=None) + yield HonchoLLMCallStreamChunk(content="lo", output_tokens=None) + yield HonchoLLMCallStreamChunk(content="", is_done=True, output_tokens=7) + + @pytest.mark.asyncio + async def test_full_drain_stamps_output_and_closes_once(self): + from src.llm.types import StreamingResponseWithMetadata + + handle = self._FakeRunHandle() + wrapper = StreamingResponseWithMetadata( + stream=self._fake_stream(), + tool_calls_made=[], + input_tokens=0, + output_tokens=0, + cache_creation_input_tokens=0, + cache_read_input_tokens=0, + langfuse_run_handle=handle, + ) + + async for _ in wrapper: + pass + + # Closed exactly once, with the concatenated streamed text as output. + assert handle.end_calls == ["hello"] + # Ownership released so a second drain can't double-close. + assert wrapper._langfuse_run_handle is None + + @pytest.mark.asyncio + async def test_abandoned_stream_still_closes_via_finally(self): + from src.llm.types import StreamingResponseWithMetadata + + handle = self._FakeRunHandle() + wrapper = StreamingResponseWithMetadata( + stream=self._fake_stream(), + tool_calls_made=[], + input_tokens=0, + output_tokens=0, + cache_creation_input_tokens=0, + cache_read_input_tokens=0, + langfuse_run_handle=handle, + ) + + # Consume one chunk, then abandon the stream. `aclose()` is what the + # runtime/GC drives when a caller stops iterating early; it throws + # GeneratorExit at the suspended `yield`, firing the `finally`. + from collections.abc import AsyncGenerator + + from src.llm.types import HonchoLLMCallStreamChunk + + agen = cast( + "AsyncGenerator[HonchoLLMCallStreamChunk, None]", wrapper.__aiter__() + ) + first = await agen.__anext__() + assert first.content == "hel" + await agen.aclose() + + # Span closed once with only the text accumulated before abandonment — + # the span is closed, not leaked. + assert handle.end_calls == ["hel"] + assert wrapper._langfuse_run_handle is None diff --git a/tests/routes/test_auth_route_policy.py b/tests/routes/test_auth_route_policy.py new file mode 100644 index 00000000..b31f2cd3 --- /dev/null +++ b/tests/routes/test_auth_route_policy.py @@ -0,0 +1,108 @@ +"""Route-policy regression tests for auth scoping. + +Two invariants this guards: + +1. `allow_member_read=True` grants peer-scoped keys read access to sessions + their peer belongs to. It must appear ONLY on intended read routes — never on + a mutating route, where it would hand session members write access. HTTP + method is not a reliable read/write signal in this codebase (some read + endpoints use POST for a richer request body), so we assert against an + explicit allowlist instead of deriving from the method. + +2. The messages router dropped its router-level auth dependency in favor of + per-route dependencies. Every route on it must still carry auth, or a future + route added without an explicit dependency would serve unauthenticated. +""" + +from fastapi.routing import APIRoute + +from src.main import app + +# (method, path) pairs intentionally granting member peers read access. Adding a +# route here is a deliberate security decision: it must be read-only. Never add +# a mutating route. See CLAUDE.md "Auth scoping" for the rule. +EXPECTED_MEMBER_READ_ROUTES = { + ("POST", "/v3/workspaces/{workspace_id}/sessions/{session_id}/messages/list"), + ( + "GET", + "/v3/workspaces/{workspace_id}/sessions/{session_id}/messages/{message_id}", + ), + ("GET", "/v3/workspaces/{workspace_id}/sessions/{session_id}/context"), + ("GET", "/v3/workspaces/{workspace_id}/sessions/{session_id}/summaries"), + ("GET", "/v3/workspaces/{workspace_id}/sessions/{session_id}/peers"), + ( + "GET", + "/v3/workspaces/{workspace_id}/sessions/{session_id}/peers/{peer_id}/config", + ), + ("POST", "/v3/workspaces/{workspace_id}/sessions/{session_id}/search"), +} + +# Unambiguously mutating methods. POST is intentionally excluded: this codebase +# uses POST for some read endpoints (`/messages/list`, `/search`) to take a +# richer request body, so POST is not a write signal. The allowlist test above +# is the real guard against a write route opting into member read; this test +# additionally catches the clear-cut PUT/PATCH/DELETE mistakes. +MUTATING_METHODS = {"PUT", "PATCH", "DELETE"} + + +def _auth_dependency_calls(route: APIRoute): + """Yield the callables of every honcho auth dependency attached to a route. + + `require_auth(...)` closures are tagged with `honcho_allow_member_read`, so a + dependency is a honcho auth dependency iff its callable has that attribute. + Walks the dependant tree to cover both `dependencies=[Depends(...)]` and + parameter-level `Depends(...)`. + """ + stack = list(route.dependant.dependencies) + while stack: + dep = stack.pop() + if hasattr(dep.call, "honcho_allow_member_read"): + yield dep.call + stack.extend(dep.dependencies) + + +def _method_path_pairs(route: APIRoute): + for method in route.methods or set(): + if method in ("HEAD", "OPTIONS"): + continue + yield (method, route.path) + + +def test_member_read_allowlist_matches_routes(): + """Exactly the allowlisted routes opt into member read — no more, no less.""" + actual: set[tuple[str, str]] = set() + for route in app.routes: + if not isinstance(route, APIRoute): + continue + if any( + getattr(call, "honcho_allow_member_read", False) + for call in _auth_dependency_calls(route) + ): + actual.update(_method_path_pairs(route)) + + assert actual == EXPECTED_MEMBER_READ_ROUTES + + +def test_member_read_never_on_mutating_route(): + """A member-read route must never use a mutating HTTP method.""" + for method, path in EXPECTED_MEMBER_READ_ROUTES: + assert method not in MUTATING_METHODS, ( + f"{method} {path} grants member-read on a mutating method — " + "member peers would gain write access" + ) + + +def test_every_message_route_requires_auth(): + """The messages router has no router-level auth dependency; assert each route + carries its own so a newly added route cannot be silently unauthenticated.""" + prefix = "/v3/workspaces/{workspace_id}/sessions/{session_id}/messages" + message_routes = [ + route + for route in app.routes + if isinstance(route, APIRoute) and route.path.startswith(prefix) + ] + assert message_routes, "expected to find message routes mounted under the prefix" + for route in message_routes: + assert any( + _auth_dependency_calls(route) + ), f"{route.methods} {route.path} has no auth dependency" diff --git a/tests/routes/test_conclusions.py b/tests/routes/test_conclusions.py index a2ef56a0..77cb23fb 100644 --- a/tests/routes/test_conclusions.py +++ b/tests/routes/test_conclusions.py @@ -737,6 +737,7 @@ class TestConclusionRoutes: assert conclusion["observer_id"] == doc.observer assert conclusion["observed_id"] == doc.observed assert conclusion["session_id"] == doc.session_name + assert conclusion["level"] == "explicit" assert "created_at" in conclusion # Verify internal fields are NOT exposed @@ -744,6 +745,80 @@ class TestConclusionRoutes: assert "internal_metadata" not in conclusion assert "collection" not in conclusion + @pytest.mark.asyncio + async def test_list_conclusions_filter_by_level( + self, + client: TestClient, + db_session: AsyncSession, + sample_data: tuple[Workspace, Peer], + ): + """Filtering by `level` returns only conclusions at that reasoning level. + + `level="explicit"` is the "not dreamed on" view — it excludes the + deductive/inductive conclusions produced during dreaming. + """ + test_workspace, test_peer = sample_data + + test_peer2 = models.Peer( + name=str(generate_nanoid()), workspace_name=test_workspace.name + ) + db_session.add(test_peer2) + await db_session.flush() + + test_session = models.Session( + name=str(generate_nanoid()), workspace_name=test_workspace.name + ) + db_session.add(test_session) + await db_session.commit() + + await self._create_collection( + db_session, test_workspace.name, test_peer.name, test_peer2.name + ) + + # Two explicit, one deductive, one inductive + levels = ["explicit", "explicit", "deductive", "inductive"] + for i, level in enumerate(levels): + db_session.add( + models.Document( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + content=f"{level} conclusion {i}", + embedding=[0.1] * 1536, + session_name=test_session.name, + level=level, + ) + ) + await db_session.commit() + + # No level filter -> all four + all_resp = client.post( + f"/v3/workspaces/{test_workspace.name}/conclusions/list", + json={"filters": {"session_id": test_session.name}}, + ) + assert all_resp.status_code == 200 + assert all_resp.json()["total"] == 4 + + # level="explicit" -> only the two non-dreamed conclusions + explicit_resp = client.post( + f"/v3/workspaces/{test_workspace.name}/conclusions/list", + json={"filters": {"session_id": test_session.name, "level": "explicit"}}, + ) + assert explicit_resp.status_code == 200 + explicit_data = explicit_resp.json() + assert explicit_data["total"] == 2 + assert all(item["level"] == "explicit" for item in explicit_data["items"]) + + # level="deductive" -> only the one deductive conclusion + deductive_resp = client.post( + f"/v3/workspaces/{test_workspace.name}/conclusions/list", + json={"filters": {"session_id": test_session.name, "level": "deductive"}}, + ) + assert deductive_resp.status_code == 200 + deductive_data = deductive_resp.json() + assert deductive_data["total"] == 1 + assert deductive_data["items"][0]["level"] == "deductive" + @pytest.mark.asyncio async def test_create_conclusion_success( self, diff --git a/tests/routes/test_messages.py b/tests/routes/test_messages.py index a121cfe0..a0186adc 100644 --- a/tests/routes/test_messages.py +++ b/tests/routes/test_messages.py @@ -1,5 +1,5 @@ import datetime -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest from fastapi.testclient import TestClient @@ -7,7 +7,9 @@ from nanoid import generate as generate_nanoid from sqlalchemy.ext.asyncio import AsyncSession from src import models +from src.config import settings from src.models import Peer, Workspace +from src.security import JWTParams, create_jwt @pytest.mark.asyncio @@ -46,6 +48,122 @@ async def test_create_message( assert "id" in message +@pytest.mark.asyncio +async def test_create_message_schedules_immediate_embed( + client: TestClient, db_session: AsyncSession, sample_data: tuple[Workspace, Peer] +): + """Creating messages should schedule the immediate-embed background task with + the created messages' public ids.""" + test_workspace, test_peer = sample_data + test_session = models.Session( + workspace_name=test_workspace.name, name=str(generate_nanoid()) + ) + db_session.add(test_session) + await db_session.commit() + + with ( + patch("src.config.settings.EMBED_MESSAGES", True), + patch.object(settings.EMBEDDING, "MAX_PENDING_EMBED_TASKS", 50), + patch( + "src.reconciler.embed_now.embed_messages_now", new=AsyncMock() + ) as mock_embed_now, + ): + response = client.post( + f"/v3/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages", + json={"messages": [{"content": "hello", "peer_id": test_peer.name}]}, + ) + assert response.status_code == 201 + public_id = response.json()[0]["id"] + mock_embed_now.assert_awaited_once_with([public_id]) + + +@pytest.mark.asyncio +async def test_create_message_skips_embed_when_disabled( + client: TestClient, db_session: AsyncSession, sample_data: tuple[Workspace, Peer] +): + """When EMBED_MESSAGES is disabled, the immediate-embed task is not scheduled.""" + test_workspace, test_peer = sample_data + test_session = models.Session( + workspace_name=test_workspace.name, name=str(generate_nanoid()) + ) + db_session.add(test_session) + await db_session.commit() + + with ( + patch("src.config.settings.EMBED_MESSAGES", False), + patch( + "src.reconciler.embed_now.embed_messages_now", new=AsyncMock() + ) as mock_embed_now, + ): + response = client.post( + f"/v3/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages", + json={"messages": [{"content": "hello", "peer_id": test_peer.name}]}, + ) + assert response.status_code == 201 + mock_embed_now.assert_not_called() + + +@pytest.mark.asyncio +async def test_create_message_defers_embed_when_saturated( + client: TestClient, db_session: AsyncSession, sample_data: tuple[Workspace, Peer] +): + """When the immediate-embed task cap is saturated, message creation still + succeeds and no embed task runs — rows stay pending for the reconciler.""" + test_workspace, test_peer = sample_data + test_session = models.Session( + workspace_name=test_workspace.name, name=str(generate_nanoid()) + ) + db_session.add(test_session) + await db_session.commit() + + with ( + patch("src.config.settings.EMBED_MESSAGES", True), + patch.object(settings.EMBEDDING, "MAX_PENDING_EMBED_TASKS", 0), + patch( + "src.reconciler.embed_now.embed_messages_now", new=AsyncMock() + ) as mock_embed_now, + ): + response = client.post( + f"/v3/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages", + json={"messages": [{"content": "hello", "peer_id": test_peer.name}]}, + ) + assert response.status_code == 201 + mock_embed_now.assert_not_called() + + +@pytest.mark.asyncio +async def test_file_upload_schedules_immediate_embed( + client: TestClient, db_session: AsyncSession, sample_data: tuple[Workspace, Peer] +): + """The file-upload path schedules the immediate-embed task with the created + messages' public ids, mirroring the session-message path.""" + import io + + test_workspace, test_peer = sample_data + test_session = models.Session( + workspace_name=test_workspace.name, name=str(generate_nanoid()) + ) + db_session.add(test_session) + await db_session.commit() + + with ( + patch("src.config.settings.EMBED_MESSAGES", True), + patch.object(settings.EMBEDDING, "MAX_PENDING_EMBED_TASKS", 50), + patch( + "src.reconciler.embed_now.embed_messages_now", new=AsyncMock() + ) as mock_embed_now, + ): + files = {"file": ("note.txt", io.BytesIO(b"hello world"), "text/plain")} + response = client.post( + f"/v3/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/upload", + files=files, + data={"peer_id": test_peer.name}, + ) + assert response.status_code == 201 + expected_ids = [m["id"] for m in response.json()] + mock_embed_now.assert_awaited_once_with(expected_ids) + + @pytest.mark.asyncio async def test_create_batch_messages_with_metadata( client: TestClient, db_session: AsyncSession, sample_data: tuple[Workspace, Peer] @@ -192,6 +310,96 @@ async def test_get_messages( assert data["items"][0]["metadata"] == {} +@pytest.mark.asyncio +async def test_member_peer_key_reads_session_but_cannot_write( + client: TestClient, + db_session: AsyncSession, + sample_data: tuple[Workspace, Peer], + monkeypatch: pytest.MonkeyPatch, +): + """A peer-scoped key may read sessions its peer belongs to (membership-based + cross-scope read), but not write to them. Non-member peer keys and session + keys on peer routes are denied. Exercises the real session_peers lookup.""" + test_workspace, alice = sample_data + session_name = str(generate_nanoid()) + base = f"/v3/workspaces/{test_workspace.name}/sessions/{session_name}" + + # Setup with auth disabled: create the session with alice as an active + # member, then commit so the independent tracked_db session in auth() (which + # only sees committed rows) can resolve membership. + client.post(f"{base}/peers", json={alice.name: {}}) + await db_session.commit() + + # Enforce auth for the assertions below. + monkeypatch.setattr(settings.AUTH, "USE_AUTH", True) + monkeypatch.setattr(settings.AUTH, "JWT_SECRET", "test-secret") + + # Member peer key: reads allowed. + client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(w=test_workspace.name, p=alice.name))}" + ) + assert client.post(f"{base}/messages/list", json={}).status_code == 200 + assert client.get(f"{base}/context").status_code == 200 + + # Member peer key: writes denied (write routes don't opt into member read). + assert ( + client.post( + f"{base}/messages", + json={"messages": [{"content": "nope", "peer_id": alice.name}]}, + ).status_code + == 401 + ) + + # Non-member peer key: even reads denied. + client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(w=test_workspace.name, p='not-a-member'))}" + ) + assert client.post(f"{base}/messages/list", json={}).status_code == 401 + + # Session key: no cross-scope access to peer routes. + client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(w=test_workspace.name, s=session_name))}" + ) + assert ( + client.get( + f"/v3/workspaces/{test_workspace.name}/peers/{alice.name}/card" + ).status_code + == 401 + ) + + +@pytest.mark.asyncio +async def test_member_peer_key_reads_only_own_session_peer_config( + client: TestClient, + db_session: AsyncSession, + sample_data: tuple[Workspace, Peer], + monkeypatch: pytest.MonkeyPatch, +): + """A member peer key may read its OWN per-session config but not a + co-member's. The route opts into member read, so without the in-handler + self-check alice could read bob's config.""" + test_workspace, alice = sample_data + bob_name = str(generate_nanoid()) + session_name = str(generate_nanoid()) + base = f"/v3/workspaces/{test_workspace.name}/sessions/{session_name}" + + # Create the session with alice and bob as active members; commit so the + # independent read-only tracked_db in auth() can resolve membership. + client.post(f"{base}/peers", json={alice.name: {}, bob_name: {}}) + await db_session.commit() + + monkeypatch.setattr(settings.AUTH, "USE_AUTH", True) + monkeypatch.setattr(settings.AUTH, "JWT_SECRET", "test-secret") + + client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(w=test_workspace.name, p=alice.name))}" + ) + # Own config: allowed. + assert client.get(f"{base}/peers/{alice.name}/config").status_code == 200 + # Co-member's config: denied even though alice is a session member. + assert client.get(f"{base}/peers/{bob_name}/config").status_code == 401 + + @pytest.mark.asyncio async def test_get_messages_with_reverse( client: TestClient, db_session: AsyncSession, sample_data: tuple[Workspace, Peer] diff --git a/tests/routes/test_peers.py b/tests/routes/test_peers.py index dd00eff9..f4b54aa3 100644 --- a/tests/routes/test_peers.py +++ b/tests/routes/test_peers.py @@ -4,10 +4,13 @@ from typing import Any import pytest from fastapi.testclient import TestClient from nanoid import generate as generate_nanoid +from pydantic import BaseModel from sqlalchemy.ext.asyncio import AsyncSession from src import crud, models +from src.config import settings from src.models import Peer, Workspace +from src.security import JWTParams, create_jwt def test_get_or_create_peer(client: TestClient, sample_data: tuple[Workspace, Peer]): @@ -625,6 +628,39 @@ def test_chat( assert "content" in data +@pytest.mark.asyncio +async def test_chat_peer_key_denied_for_non_member_session( + client: TestClient, + db_session: AsyncSession, + sample_data: tuple[Workspace, Peer], + monkeypatch: pytest.MonkeyPatch, +): + """A peer-scoped key cannot chat scoped to a session its peer is not a member + of — the session id is in the body, so the handler checks membership. The + guard fires before the dialectic runs, so no LLM call is made.""" + test_workspace, alice = sample_data + session_id = str(generate_nanoid()) + + # Session exists but alice is NOT a member of it. + client.post( + f"/v3/workspaces/{test_workspace.name}/sessions", + json={"id": session_id}, + ) + await db_session.commit() + + monkeypatch.setattr(settings.AUTH, "USE_AUTH", True) + monkeypatch.setattr(settings.AUTH, "JWT_SECRET", "test-secret") + client.headers["Authorization"] = ( + f"Bearer {create_jwt(JWTParams(w=test_workspace.name, p=alice.name))}" + ) + + response = client.post( + f"/v3/workspaces/{test_workspace.name}/peers/{alice.name}/chat", + json={"query": "what do you know?", "stream": False, "session_id": session_id}, + ) + assert response.status_code == 401 + + def test_chat_with_optional_params( client: TestClient, sample_data: tuple[Workspace, Peer], @@ -1235,3 +1271,114 @@ def test_set_peer_card(client: TestClient, sample_data: tuple[Workspace, Peer]): ) assert response.status_code == 200 assert response.json()["peer_card"] == target_card + + +FOOD_PREFS_SCHEMA = { + "type": "object", + "properties": { + "preferences": { + "type": "array", + "items": { + "type": "object", + "properties": { + "food": {"type": "string"}, + "sentiment": {"enum": ["loves", "likes", "dislikes"]}, + }, + "required": ["food", "sentiment"], + }, + }, + "summary": {"type": "string"}, + }, + "required": ["preferences", "summary"], +} + + +def test_chat_with_response_format( + client: TestClient, + sample_data: tuple[Workspace, Peer], + mock_llm_call_functions: dict[str, Any], +): + """A valid response_format converts to a Pydantic model and is passed to + the dialectic as response_model.""" + test_workspace, test_peer = sample_data + + response = client.post( + f"/v3/workspaces/{test_workspace.name}/peers/{test_peer.name}/chat", + json={ + "query": "What are this user's food preferences?", + "stream": False, + "response_format": FOOD_PREFS_SCHEMA, + }, + ) + assert response.status_code == 200 + assert "content" in response.json() + + kwargs = mock_llm_call_functions["agentic_chat"].await_args.kwargs + response_model = kwargs["response_model"] + assert isinstance(response_model, type) + assert issubclass(response_model, BaseModel) + # The converted model enforces the caller's schema. + instance = response_model.model_validate( + {"preferences": [{"food": "sushi", "sentiment": "loves"}], "summary": "s"} + ) + assert instance.summary == "s" # pyright: ignore + + +def test_chat_with_response_format_streaming( + client: TestClient, + sample_data: tuple[Workspace, Peer], + mock_llm_call_functions: dict[str, Any], +): + test_workspace, test_peer = sample_data + + response = client.post( + f"/v3/workspaces/{test_workspace.name}/peers/{test_peer.name}/chat", + json={ + "query": "What are this user's food preferences?", + "stream": True, + "response_format": FOOD_PREFS_SCHEMA, + }, + ) + assert response.status_code == 200 + assert "data:" in response.text + + kwargs = mock_llm_call_functions["agentic_chat_stream"].call_args.kwargs + response_model = kwargs["response_model"] + assert isinstance(response_model, type) + assert issubclass(response_model, BaseModel) + + +@pytest.mark.parametrize( + "bad_schema", + [ + {"type": "string"}, # non-object root + {"type": "object", "properties": {"a": {"$ref": "#/x"}}}, + {"type": "object", "properties": {"a": {"allOf": [{"type": "string"}]}}}, + { + "type": "object", + "properties": { + "m": {"type": "object", "additionalProperties": {"type": "string"}} + }, + }, + ], +) +def test_chat_with_invalid_response_format( + client: TestClient, + sample_data: tuple[Workspace, Peer], + mock_llm_call_functions: dict[str, Any], + bad_schema: dict[str, Any], +): + """Unsupported schemas are rejected with 422 before the dialectic runs.""" + test_workspace, test_peer = sample_data + + response = client.post( + f"/v3/workspaces/{test_workspace.name}/peers/{test_peer.name}/chat", + json={ + "query": "Hello?", + "stream": False, + "response_format": bad_schema, + }, + ) + assert response.status_code == 422 + assert "Invalid response_format" in response.json()["detail"] + mock_llm_call_functions["agentic_chat"].assert_not_awaited() diff --git a/tests/routes/test_scoped_api.py b/tests/routes/test_scoped_api.py index a4d4677a..7b504a7a 100644 --- a/tests/routes/test_scoped_api.py +++ b/tests/routes/test_scoped_api.py @@ -141,7 +141,7 @@ def test_get_peer_by_name_with_auth( # Test with peer-scoped JWT if auth_client.auth_type == "empty": auth_client.headers["Authorization"] = ( - f"Bearer {create_jwt(JWTParams(p=test_peer.name))}" + f"Bearer {create_jwt(JWTParams(w=test_workspace.name, p=test_peer.name))}" ) # Get specific peer using get_or_create endpoint @@ -218,7 +218,7 @@ def test_create_session_with_auth( # Test with peer-scoped JWT if auth_client.auth_type == "empty": auth_client.headers["Authorization"] = ( - f"Bearer {create_jwt(JWTParams(p=test_peer.name))}" + f"Bearer {create_jwt(JWTParams(w=test_workspace.name, p=test_peer.name))}" ) session_name2 = str(generate_nanoid()) @@ -262,7 +262,7 @@ def test_get_session_by_name_with_auth( if auth_client.auth_type == "empty": # Test with session-scoped JWT auth_client.headers["Authorization"] = ( - f"Bearer {create_jwt(JWTParams(s=session_name))}" + f"Bearer {create_jwt(JWTParams(w=test_workspace.name, s=session_name))}" ) response = auth_client.post( @@ -282,7 +282,7 @@ def test_get_session_by_name_with_auth( # Test with peer-scoped JWT auth_client.headers["Authorization"] = ( - f"Bearer {create_jwt(JWTParams(p=test_peer.name))}" + f"Bearer {create_jwt(JWTParams(w=test_workspace.name, p=test_peer.name))}" ) assert auth_client.post( diff --git a/tests/routes/test_workspaces.py b/tests/routes/test_workspaces.py index 0350b729..5934ccfd 100644 --- a/tests/routes/test_workspaces.py +++ b/tests/routes/test_workspaces.py @@ -9,6 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import models from src.models import Peer, Workspace +from src.schemas import DreamType def test_get_or_create_workspace(client: TestClient): @@ -758,3 +759,52 @@ async def test_schedule_dream_invokes_enqueue_dream( "Loop 4: enqueue_dream no longer accepts document_count; the baseline " "is written atomically with last_dream_at in process_dream." ) + + +@pytest.mark.asyncio +async def test_schedule_dream_card_refresh_forwards_rebuild( + client: TestClient, + db_session: AsyncSession, + sample_data: tuple[Workspace, Peer], +): + """POST /schedule_dream accepts dream_type=card_refresh and forwards the + rebuild flag to enqueue_dream (manual/event-driven card refreshes bypass + the volume gates by design).""" + workspace, peer = sample_data + + collection = models.Collection( + observer=peer.name, + observed=peer.name, + workspace_name=workspace.name, + internal_metadata={}, + ) + db_session.add(collection) + await db_session.commit() + + captured: dict[str, Any] = {} + + async def fake_enqueue_dream(*args: Any, **kwargs: Any) -> None: + captured["args"] = args + captured["kwargs"] = kwargs + + with ( + patch("src.routers.workspaces.settings.DREAM.ENABLED", True), + patch( + "src.routers.workspaces.enqueue_dream", + new=AsyncMock(side_effect=fake_enqueue_dream), + ), + ): + response = client.post( + f"/v3/workspaces/{workspace.name}/schedule_dream", + json={ + "observer": peer.name, + "observed": peer.name, + "dream_type": "card_refresh", + "rebuild": True, + }, + ) + + assert response.status_code == 204, response.text + assert "kwargs" in captured, "enqueue_dream was not called" + assert captured["kwargs"]["dream_type"] == DreamType.CARD_REFRESH + assert captured["kwargs"]["rebuild"] is True diff --git a/tests/sdk/test_conclusions.py b/tests/sdk/test_conclusions.py index 33331de3..9c2cdbc0 100644 --- a/tests/sdk/test_conclusions.py +++ b/tests/sdk/test_conclusions.py @@ -769,3 +769,73 @@ async def test_observation_create_mixed_session_and_sessionless( assert session_obs.session_id == session.id assert global_obs.session_id is None + + +@pytest.mark.asyncio +async def test_list_rejects_reserved_scope_filter_keys( + client_fixture: tuple[Honcho, str], +): + """`list` rejects observer/observed/session filter keys managed by the scope. + + These keys are fixed by the scope (observer/observed) or by the dedicated + ``session=`` parameter, so passing them in ``filters`` would silently return + data from a different scope. The guard raises before any HTTP call. + """ + honcho_client, client_type = client_fixture + reserved = [ + "observer", + "observed", + "observer_id", + "observed_id", + "session_id", + "session", + ] + + if client_type == "async": + observer = await honcho_client.aio.peer(id="test-obs-reserved-list-observer") + target = await honcho_client.aio.peer(id="test-obs-reserved-list-target") + obs_scope = observer.conclusions_of(target) + for key in reserved: + with pytest.raises(ValueError, match="managed by this conclusion scope"): + await obs_scope.aio.list(filters={key: "someone-else"}) + # A non-reserved filter (level) is allowed through. + await obs_scope.aio.list(filters={"level": "explicit"}) + else: + observer = honcho_client.peer(id="test-obs-reserved-list-observer") + target = honcho_client.peer(id="test-obs-reserved-list-target") + obs_scope = observer.conclusions_of(target) + for key in reserved: + with pytest.raises(ValueError, match="managed by this conclusion scope"): + obs_scope.list(filters={key: "someone-else"}) + obs_scope.list(filters={"level": "explicit"}) + + +@pytest.mark.asyncio +async def test_query_rejects_reserved_scope_filter_keys( + client_fixture: tuple[Honcho, str], +): + """`query` rejects observer/observed filter keys but allows session_id. + + Unlike ``list``, ``query`` has no dedicated session parameter, so + ``session_id`` remains a normal filter and must NOT be rejected. + """ + honcho_client, client_type = client_fixture + reserved = ["observer", "observed", "observer_id", "observed_id"] + + if client_type == "async": + observer = await honcho_client.aio.peer(id="test-obs-reserved-query-observer") + target = await honcho_client.aio.peer(id="test-obs-reserved-query-target") + obs_scope = observer.conclusions_of(target) + for key in reserved: + with pytest.raises(ValueError, match="managed by this conclusion scope"): + await obs_scope.aio.query("q", filters={key: "someone-else"}) + # session_id is a normal filter for query (no dedicated param) — allowed. + await obs_scope.aio.query("q", filters={"session_id": "some-session"}) + else: + observer = honcho_client.peer(id="test-obs-reserved-query-observer") + target = honcho_client.peer(id="test-obs-reserved-query-target") + obs_scope = observer.conclusions_of(target) + for key in reserved: + with pytest.raises(ValueError, match="managed by this conclusion scope"): + obs_scope.query("q", filters={key: "someone-else"}) + obs_scope.query("q", filters={"session_id": "some-session"}) diff --git a/tests/sdk/test_peer.py b/tests/sdk/test_peer.py index 258fb989..0c019d9b 100644 --- a/tests/sdk/test_peer.py +++ b/tests/sdk/test_peer.py @@ -2,6 +2,7 @@ from collections.abc import AsyncIterator, Iterator from unittest.mock import patch import pytest +from pydantic import BaseModel from sdks.python.src.honcho.client import Honcho from sdks.python.src.honcho.peer import Peer @@ -626,3 +627,166 @@ async def test_peer_representation_with_all_params( max_conclusions=5, ) assert isinstance(result, str) + + +class ChatFoodPreferences(BaseModel): + favorite: str + confidence: float + + +CHAT_SCHEMA_DICT = { + "type": "object", + "properties": {"items": {"type": "array", "items": {"type": "string"}}}, + "required": ["items"], +} + + +@pytest.mark.asyncio +async def test_peer_chat_response_format_pydantic(client_fixture: tuple[Honcho, str]): + """A Pydantic model class is sent as JSON Schema and the response content + is parsed back into a model instance.""" + honcho_client, client_type = client_fixture + content = '{"favorite": "sushi", "confidence": 0.9}' + + if client_type == "async": + peer = await honcho_client.aio.peer(id="test-rf-async-peer") + + async def mock_post(*_args: object, **_kwargs: object) -> dict[str, object]: + return {"content": content} + + with patch.object( + peer._honcho._async_http_client, # pyright: ignore[reportPrivateUsage] + "post", + side_effect=mock_post, + ) as mock: + result = await peer.aio.chat( + "What do I like?", response_format=ChatFoodPreferences + ) + else: + peer = honcho_client.peer(id="test-rf-peer") + with patch.object( + peer._honcho._http, # pyright: ignore[reportPrivateUsage] + "post", + return_value={"content": content}, + ) as mock: + result = peer.chat("What do I like?", response_format=ChatFoodPreferences) + + body = mock.call_args.kwargs["body"] + assert body["response_format"] == ChatFoodPreferences.model_json_schema() + assert isinstance(result, ChatFoodPreferences) + assert result.favorite == "sushi" + + +@pytest.mark.asyncio +async def test_peer_chat_response_format_dict(client_fixture: tuple[Honcho, str]): + """A raw JSON Schema dict is sent as-is and the response stays a string.""" + honcho_client, client_type = client_fixture + content = '{"items": ["sushi"]}' + + if client_type == "async": + peer = await honcho_client.aio.peer(id="test-rf-dict-async-peer") + + async def mock_post(*_args: object, **_kwargs: object) -> dict[str, object]: + return {"content": content} + + with patch.object( + peer._honcho._async_http_client, # pyright: ignore[reportPrivateUsage] + "post", + side_effect=mock_post, + ) as mock: + result = await peer.aio.chat( + "What do I like?", response_format=CHAT_SCHEMA_DICT + ) + else: + peer = honcho_client.peer(id="test-rf-dict-peer") + with patch.object( + peer._honcho._http, # pyright: ignore[reportPrivateUsage] + "post", + return_value={"content": content}, + ) as mock: + result = peer.chat("What do I like?", response_format=CHAT_SCHEMA_DICT) + + body = mock.call_args.kwargs["body"] + assert body["response_format"] == CHAT_SCHEMA_DICT + assert result == content + + +@pytest.mark.asyncio +async def test_peer_chat_response_format_empty_content( + client_fixture: tuple[Honcho, str], +): + """Empty/None content returns None even when a Pydantic class was given.""" + honcho_client, client_type = client_fixture + + if client_type == "async": + peer = await honcho_client.aio.peer(id="test-rf-empty-async-peer") + + async def mock_post(*_args: object, **_kwargs: object) -> dict[str, object]: + return {"content": None} + + with patch.object( + peer._honcho._async_http_client, # pyright: ignore[reportPrivateUsage] + "post", + side_effect=mock_post, + ): + result = await peer.aio.chat( + "What do I like?", response_format=ChatFoodPreferences + ) + else: + peer = honcho_client.peer(id="test-rf-empty-peer") + with patch.object( + peer._honcho._http, # pyright: ignore[reportPrivateUsage] + "post", + return_value={"content": None}, + ): + result = peer.chat("What do I like?", response_format=ChatFoodPreferences) + + assert result is None + + +@pytest.mark.asyncio +async def test_peer_chat_stream_response_format(client_fixture: tuple[Honcho, str]): + """chat_stream sends the schema in the body; chunks stay raw text.""" + honcho_client, client_type = client_fixture + + if client_type == "async": + peer = await honcho_client.aio.peer(id="test-rf-stream-async-peer") + + async def mock_astream( + *_args: object, **_kwargs: object + ) -> AsyncIterator[bytes]: + yield b'data: {"delta": {"content": "{\\"favorite\\":"}}\n' + yield b'data: {"delta": {"content": "\\"sushi\\",\\"confidence\\":0.9}"}}\n' + yield b'data: {"done": true}\n' + + with patch.object( + peer._honcho._async_http_client, # pyright: ignore[reportPrivateUsage] + "stream", + side_effect=mock_astream, + ) as mock: + result = await peer.aio.chat_stream( + "What do I like?", response_format=ChatFoodPreferences + ) + chunks = [chunk async for chunk in result] + else: + peer = honcho_client.peer(id="test-rf-stream-peer") + + def mock_stream(*_args: object, **_kwargs: object) -> Iterator[bytes]: + yield b'data: {"delta": {"content": "{\\"favorite\\":"}}\n' + yield b'data: {"delta": {"content": "\\"sushi\\",\\"confidence\\":0.9}"}}\n' + yield b'data: {"done": true}\n' + + with patch.object( + peer._honcho._http, # pyright: ignore[reportPrivateUsage] + "stream", + side_effect=mock_stream, + ) as mock: + result = peer.chat_stream( + "What do I like?", response_format=ChatFoodPreferences + ) + chunks = list(result) + + body = mock.call_args.kwargs["body"] + assert body["response_format"] == ChatFoodPreferences.model_json_schema() + accumulated = "".join(chunks) + assert ChatFoodPreferences.model_validate_json(accumulated).favorite == "sushi" diff --git a/tests/sdk_typescript/conftest.py b/tests/sdk_typescript/conftest.py index 74fe148b..15505a0e 100644 --- a/tests/sdk_typescript/conftest.py +++ b/tests/sdk_typescript/conftest.py @@ -20,7 +20,7 @@ from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker from uvicorn.config import Config from uvicorn.server import Server -from src.dependencies import get_db +from src.dependencies import get_db, get_read_db from src.main import app @@ -95,6 +95,9 @@ def ts_test_server( yield session app.dependency_overrides[get_db] = override_get_db + # Read-only routes use get_read_db (AUTOCOMMIT engine) in production; in + # tests they must resolve to the same per-test database. + app.dependency_overrides[get_read_db] = override_get_db # No-op the lifespan's startup embedding-schema validator — same # reasoning as the `client` fixture in tests/conftest.py: the module- @@ -133,7 +136,9 @@ def mock_tracked_db(ts_db_session: async_sessionmaker[AsyncSession]): # Create a tracked_db that uses fresh sessions (not shared) @asynccontextmanager - async def ts_tracked_db(_: str | None = None): + async def ts_tracked_db(_: str | None = None, *, read_only: bool = False): + # read_only accepted (and ignored): tests use one per-test database. + del read_only async with ts_db_session() as session: yield session diff --git a/tests/startup/test_embedding_validator.py b/tests/startup/test_embedding_validator.py index 962e07f0..1d257ba5 100644 --- a/tests/startup/test_embedding_validator.py +++ b/tests/startup/test_embedding_validator.py @@ -14,6 +14,7 @@ from sqlalchemy import text from sqlalchemy.exc import OperationalError from sqlalchemy.ext.asyncio import AsyncEngine +from src.config import settings from src.startup.embedding_validator import ( StartupValidationError, _assert_pgvector_dims_match, # pyright: ignore[reportPrivateUsage] @@ -120,18 +121,24 @@ async def test_validator_fails_closed_when_introspection_keeps_failing( @pytest.mark.asyncio async def test_validator_passes_against_test_database( db_engine: AsyncEngine, + monkeypatch: pytest.MonkeyPatch, ) -> None: """The test DB is provisioned at the default dim (1536); the validator should accept it without raising.""" + # conftest provisions the test tables in `public`; pin the validator to it + # so a developer's local .env DB_SCHEMA can't point it at another schema. + monkeypatch.setattr(settings.DB, "SCHEMA", "public") await validate_embedding_schema(db_engine) @pytest.mark.asyncio async def test_validator_raises_when_schema_dim_diverges_from_settings( db_engine: AsyncEngine, + monkeypatch: pytest.MonkeyPatch, ) -> None: """ALTER one of the embedding columns to a non-1536 dim and confirm the validator raises with an actionable message.""" + monkeypatch.setattr(settings.DB, "SCHEMA", "public") # see test above async with db_engine.begin() as conn: await conn.execute( text( @@ -181,8 +188,13 @@ def test_non_1536_pgvector_without_migrated_no_longer_raises_at_config_time() -> """The dim-vs-MIGRATED guard has been removed. Constructing AppSettings with non-1536 + default pgvector + MIGRATED=false should now succeed (the runtime schema validator at startup is the safety net).""" + # Minimal env, NOT a copy of os.environ: load_dotenv() in the app mutates + # the parent pytest process's environ, so inheriting it would leak a + # developer's local .env (DB_SCHEMA, VECTOR_STORE_MIGRATED, ...) into the + # child despite PYTHON_DOTENV_DISABLED. The child must see pure defaults + # plus exactly the overrides below. env = { - **os.environ, + "PATH": os.environ.get("PATH", ""), "PYTHON_DOTENV_DISABLED": "1", "EMBEDDING_VECTOR_DIMENSIONS": "768", } diff --git a/tests/telemetry/conftest.py b/tests/telemetry/conftest.py index f871e22d..442f6264 100644 --- a/tests/telemetry/conftest.py +++ b/tests/telemetry/conftest.py @@ -35,6 +35,42 @@ from src.telemetry.events.reconciliation import ( ) from src.telemetry.events.representation import RepresentationCompletedEvent +# ============================================================================= +# Global trace-state isolation +# ============================================================================= + + +@pytest.fixture(autouse=True) +def _isolate_trace_globals(): # pyright: ignore[reportUnusedFunction] + """Snapshot/restore process-global trace state around every telemetry test. + + `initialize_telemetry_events()` (exercised in test_emit_function) registers a + real ``TraceExporter`` into ``capture._EXPORTERS`` and starts a real trace + emitter (``emitter._trace_emitter``). Without cleanup that state bleeds into + the trace-exporter tests, which then fail — and because xdist schedules tests + across workers nondeterministically, the failure looks flaky (a different + trace test fails each run depending on who shared its worker). + + Snapshotting these globals (and resetting the per-run dedup) before/after + each test makes the trace tests hermetic regardless of neighbor ordering. + """ + from src.llm import capture + from src.telemetry import emitter as emitter_mod + from src.telemetry import langfuse_session, trace_session + + saved_exporters = list(capture._EXPORTERS) # pyright: ignore[reportPrivateUsage] + saved_trace_emitter = emitter_mod._trace_emitter # pyright: ignore[reportPrivateUsage] + trace_session.reset() + langfuse_session.reset() + try: + yield + finally: + capture._EXPORTERS[:] = saved_exporters # pyright: ignore[reportPrivateUsage] + emitter_mod._trace_emitter = saved_trace_emitter # pyright: ignore[reportPrivateUsage] + trace_session.reset() + langfuse_session.reset() + + # ============================================================================= # Fixed timestamp for deterministic tests # ============================================================================= diff --git a/tests/telemetry/test_cross_agent_trace.py b/tests/telemetry/test_cross_agent_trace.py new file mode 100644 index 00000000..b6364ed4 --- /dev/null +++ b/tests/telemetry/test_cross_agent_trace.py @@ -0,0 +1,215 @@ +# pyright: reportPrivateUsage=false, reportUnannotatedClassAttribute=false, reportUnusedFunction=false, reportUnknownLambdaType=false, reportUnknownArgumentType=false, reportArgumentType=false +"""Cross-agent trace-metadata contract. + +Verifies that each agent's telemetry produces a well-formed `CapturedLLMCall` +with the right correlation/session/identity fields, and that the SAME captured +call fans out to BOTH exporters (CloudEvents + Langfuse) — the "one data model, +two projections" invariant. + +This is the metadata-correctness bar across agents. It drives the real +`DialecticAgent` telemetry and the real `dispatch_captured_call`; the other +agents are represented by the telemetry contexts they construct (cited inline). +Full end-to-end capture through a live stack (real subprocesses feeding a trace +sink) is exercised separately, outside this unit suite. +""" + +from __future__ import annotations + +import pytest + +from src.config import settings +from src.dialectic.core import DialecticAgent +from src.llm import capture as capture_mod +from src.llm.backend import CompletionResult +from src.llm.capture import ( + CapturedLLMCall, + build_captured_call, + dispatch_captured_call, + register_exporter, +) +from src.llm.types import LLMTelemetryContext +from src.telemetry.langfuse_exporter import LangfuseExporter + + +class SpyExporter: + def __init__(self) -> None: + self.calls: list[CapturedLLMCall] = [] + + def export(self, call: CapturedLLMCall) -> None: + self.calls.append(call) + + +@pytest.fixture +def spy() -> SpyExporter: + """Clean exporter registry with a single spy (restored by the telemetry + conftest's _isolate_trace_globals).""" + capture_mod._EXPORTERS.clear() + exporter = SpyExporter() + register_exporter(exporter) + return exporter + + +def _dispatch(telemetry: LLMTelemetryContext, *, content: str = "answer") -> None: + call = build_captured_call( + telemetry=telemetry, + transport="anthropic", + provider_label=None, + model="claude-x", + messages=[{"role": "user", "content": "q"}], + tools=None, + tool_choice=None, + result=CompletionResult(content=content, finish_reason="stop"), + attempt=1, + was_fallback=False, + was_stream=False, + finish_reason="stop", + ) + dispatch_captured_call(call) + + +# --- Dialectic: real agent telemetry -------------------------------------- + + +def test_dialectic_with_session_sets_session_id(spy: SpyExporter): + agent = DialecticAgent( + workspace_name="ws", + session_name="my-session", + session_id="sess-nanoid", + observer="alice", + observed="bob", + ) + _dispatch(agent._telemetry_context("Dialectic Agent")) + + call = spy.calls[-1] + assert call.session_id == "sess-nanoid" + assert call.agent_type == "dialectic" + # Root of the invocation: trace_id == span_id == run_id, parent normalized off. + assert call.trace_id == call.span_id == agent._run_id + assert call.parent_span_id is None + assert call.track_name == "Dialectic Agent" + + +def test_dialectic_global_has_no_session(spy: SpyExporter): + agent = DialecticAgent( + workspace_name="ws", + session_name=None, + session_id=None, + observer="alice", + observed="alice", + ) + _dispatch(agent._telemetry_context("Dialectic Agent")) + assert spy.calls[-1].session_id is None + + +# --- Background agents: sessionless single-shot / shared-tree contracts ----- + + +@pytest.mark.parametrize( + ("call_purpose", "parent_category", "track_name"), + [ + ("deriver.representation", "representation", "Minimal Deriver"), + ("summary.short", "summary", None), + ], +) +def test_background_agents_are_sessionless_single_shot( + spy: SpyExporter, + call_purpose: str, + parent_category: str, + track_name: str | None, +): + # Deriver + summarizer mirror their src/ contexts: trace_id == span_id, no + # run_id/session_id, self-rooted. + tid = f"{parent_category}-trace" + _dispatch( + LLMTelemetryContext( + workspace_name="ws", + call_purpose=call_purpose, + parent_category=parent_category, + track_name=track_name, + trace_id=tid, + span_id=tid, + ) + ) + call = spy.calls[-1] + assert call.session_id is None + assert call.run_id is None + assert call.trace_id == call.span_id == tid + assert call.parent_span_id is None + + +def test_dreamer_specialists_share_one_tree(spy: SpyExporter): + # Single-dream-tree (this PR): both specialists reuse the orchestrator run_id + # as trace_id (src/dreamer/specialists.py), session_id None. + run_id = "dream-run" + for agent_type in ("deduction", "induction"): + _dispatch( + LLMTelemetryContext( + workspace_name="ws", + call_purpose=f"dream.{agent_type}", + parent_category="dream", + agent_type=agent_type, + run_id=run_id, + trace_id=run_id, + span_id=run_id, + observer="assistant", + observed="bob", + iteration=1, + ) + ) + ded, ind = spy.calls[-2], spy.calls[-1] + assert ded.trace_id == ind.trace_id == run_id # one shared tree + assert ded.session_id is None and ind.session_id is None + assert ded.agent_type == "deduction" and ind.agent_type == "induction" + + +# --- One data model, two projections --------------------------------------- + + +def test_same_call_reaches_both_exporters( + spy: SpyExporter, monkeypatch: pytest.MonkeyPatch +): + """A dispatched call fans out to the CloudEvents spy AND the LangfuseExporter.""" + monkeypatch.setattr(settings, "LANGFUSE_PUBLIC_KEY", "pk-test") + monkeypatch.setattr(settings, "LANGFUSE_EXPORTER_MODE", "exporter") + monkeypatch.setattr(settings, "NAMESPACE", "tenant1") + + created: list[dict[str, object]] = [] + + class FakeOtel: + def set_attribute(self, *_a: object) -> None: ... + + class FakeObs: + def __init__(self, **kwargs: object) -> None: + self.id = "obs" + self.kwargs = kwargs + self._otel_span = FakeOtel() + + def end(self) -> None: ... + + class FakeClient: + def create_trace_id(self, *, seed: str | None = None) -> str: + return f"lf-{seed}" + + def start_observation(self, **kwargs: object) -> FakeObs: + created.append(kwargs) + return FakeObs(**kwargs) + + import langfuse + + monkeypatch.setattr(langfuse, "get_client", lambda: FakeClient()) + register_exporter(LangfuseExporter()) + + agent = DialecticAgent( + workspace_name="ws", + session_name="s", + session_id="sess-1", + observer="alice", + observed="bob", + ) + _dispatch(agent._telemetry_context("Dialectic Agent")) + + # CloudEvents projection saw the raw captured call... + assert spy.calls and spy.calls[-1].session_id == "sess-1" + # ...and the Langfuse projection built observations from the SAME call. + assert created, "LangfuseExporter produced no observations" + assert any(o.get("as_type") == "generation" for o in created) diff --git a/tests/telemetry/test_embedding_trace.py b/tests/telemetry/test_embedding_trace.py new file mode 100644 index 00000000..a878e4f3 --- /dev/null +++ b/tests/telemetry/test_embedding_trace.py @@ -0,0 +1,111 @@ +# pyright: reportPrivateUsage=false, reportUnannotatedClassAttribute=false, reportUnusedFunction=false, reportUnknownLambdaType=false, reportUnknownArgumentType=false, reportArgumentType=false +"""Tests for embedding calls joining the trace stream. + +`_publish_embedding_event` emits an `EmbeddingCallTracedEvent` (gated on +TRACE_PAYLOADS_ENABLED) in addition to the metrics-grade completed event, carrying the +span-tree correlation from the embedding ContextVars so an embedding made inside +an agent run nests under that run's trace. +""" + +from __future__ import annotations + +import pytest + +import src.telemetry.events as events_mod +from src.config import settings +from src.embedding_client import _publish_embedding_event +from src.telemetry.events.trace import EmbeddingCallTracedEvent +from src.utils.types import embedding_call_purpose + + +@pytest.fixture +def capture_emits(monkeypatch: pytest.MonkeyPatch): + """Capture emit()/emit_trace() without a live emitter.""" + traced: list[object] = [] + monkeypatch.setattr(events_mod, "emit", lambda _e: None) + monkeypatch.setattr(events_mod, "emit_trace", lambda e: traced.append(e)) + return traced + + +def test_embedding_traced_event_carries_correlation( + capture_emits: list[object], monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(settings.TELEMETRY, "TRACE_PAYLOADS_ENABLED", True) + with embedding_call_purpose( + "dialectic.prefetch", + workspace_name="ws", + run_id="run-1", + parent_category="dialectic", + session_id="sess-1", + ): + _publish_embedding_event( + provider="openai", + model="text-embedding-3", + input_count=1, + input_tokens_estimate=7, + duration_ms=1.0, + outcome="success", + error=None, + is_final_attempt=True, + ) + + assert len(capture_emits) == 1 + ev = capture_emits[0] + assert isinstance(ev, EmbeddingCallTracedEvent) + # The embedding gets its own span under the run: trace_id/parent are the + # run_id, span_id is a fresh id so sibling embeddings don't collide. + assert ev.trace_id == "run-1" + assert ev.parent_span_id == "run-1" + assert ev.span_id and ev.span_id != "run-1" + assert ev.session_id == "sess-1" + assert ev.call_purpose == "dialectic.prefetch" + assert ev.parent_category == "dialectic" + assert ev.provider == "openai" and ev.model == "text-embedding-3" + assert ev.provider_input_tokens == 7 + assert ev.provider_output_tokens == 0 + assert ev.input_count == 1 + + +def test_no_trace_event_when_payloads_off( + capture_emits: list[object], monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(settings.TELEMETRY, "TRACE_PAYLOADS_ENABLED", False) + with embedding_call_purpose("dialectic.prefetch", run_id="run-1"): + _publish_embedding_event( + provider="openai", + model="m", + input_count=2, + input_tokens_estimate=3, + duration_ms=1.0, + outcome="success", + error=None, + is_final_attempt=True, + ) + assert capture_emits == [] + + +def test_sessionless_embedding_has_no_session( + capture_emits: list[object], monkeypatch: pytest.MonkeyPatch +): + # A deriver/reconciler embedding (no session scope) traces with session None. + monkeypatch.setattr(settings.TELEMETRY, "TRACE_PAYLOADS_ENABLED", True) + with embedding_call_purpose( + "deriver", workspace_name="ws", parent_category="deriver" + ): + _publish_embedding_event( + provider="gemini", + model="emb", + input_count=5, + input_tokens_estimate=20, + duration_ms=2.0, + outcome="success", + error=None, + is_final_attempt=True, + ) + assert len(capture_emits) == 1 + ev = capture_emits[0] + assert isinstance(ev, EmbeddingCallTracedEvent) + assert ev.session_id is None + # No run_id → the span self-roots (trace_id == span_id) with no parent. + assert ev.parent_span_id is None + assert ev.span_id and ev.trace_id == ev.span_id diff --git a/tests/telemetry/test_emit_function.py b/tests/telemetry/test_emit_function.py index c4e2c8ad..1e8f78f9 100644 --- a/tests/telemetry/test_emit_function.py +++ b/tests/telemetry/test_emit_function.py @@ -314,6 +314,11 @@ class TestInitializeTelemetryEvents: mock_settings.TELEMETRY.FLUSH_THRESHOLD = 50 mock_settings.TELEMETRY.MAX_RETRIES = 3 mock_settings.TELEMETRY.MAX_BUFFER_SIZE = 10000 + # This test only covers the primary emitter. Pin trace payloads off + # so we don't fall into the trace branch and start a *real* trace + # emitter + register a real TraceExporter (a MagicMock here is + # truthy) — that global state would leak into other tests. + mock_settings.TELEMETRY.TRACE_PAYLOADS_ENABLED = False mock_init.return_value = AsyncMock() await initialize_telemetry_events() @@ -378,8 +383,10 @@ class TestInitializeTelemetryAsync: mock_ce_init.assert_called_once() @pytest.mark.asyncio - async def test_skip_cloudevents_when_disabled(self): - """initialize_telemetry_async() skips CloudEvents when disabled.""" + async def test_skip_when_telemetry_disabled(self): + """TELEMETRY.ENABLED is the master switch — no init when it's off, even + with the Langfuse exporter configured (no traces for open-source users + who leave telemetry off).""" from src.telemetry import initialize_telemetry_async with ( @@ -390,6 +397,7 @@ class TestInitializeTelemetryAsync: ) as mock_ce_init, ): mock_settings.TELEMETRY.ENABLED = False + mock_settings.langfuse_exporter_enabled = True await initialize_telemetry_async() diff --git a/tests/telemetry/test_events.py b/tests/telemetry/test_events.py index ef6f4ffb..a5d959e5 100644 --- a/tests/telemetry/test_events.py +++ b/tests/telemetry/test_events.py @@ -239,6 +239,7 @@ class TestLLMCallCompletedEvent: assert CallPurpose.DIALECTIC_ANSWER.value == "dialectic.answer" assert CallPurpose.DREAM_DEDUCTION.value == "dream.deduction" assert CallPurpose.DREAM_INDUCTION.value == "dream.induction" + assert CallPurpose.DREAM_CARD_REFRESH.value == "dream.card_refresh" assert CallPurpose.SUMMARY_SHORT.value == "summary.short" assert CallPurpose.SUMMARY_LONG.value == "summary.long" @@ -759,10 +760,11 @@ class TestAgentToolSummaryCreatedEvent: def test_get_resource_id( self, sample_summary_created_event: AgentToolSummaryCreatedEvent ): - """get_resource_id() returns run_id:iteration:summary_created format.""" + """get_resource_id() keys on message_id:summary_type (run_id/iteration are + None for the non-agentic summarizer and can't identify the summary).""" assert ( sample_summary_created_event.get_resource_id() - == "ghi11111:1:summary_created" + == "msg_020:short:summary_created" ) def test_summary_type_values(self, fixed_timestamp: datetime): diff --git a/tests/telemetry/test_langfuse_exporter.py b/tests/telemetry/test_langfuse_exporter.py new file mode 100644 index 00000000..36f3c9c2 --- /dev/null +++ b/tests/telemetry/test_langfuse_exporter.py @@ -0,0 +1,438 @@ +# pyright: reportPrivateUsage=false, reportUnannotatedClassAttribute=false, reportUnusedFunction=false, reportUnknownLambdaType=false, reportUnknownArgumentType=false, reportArgumentType=false, reportIndexIssue=false +"""Tests for the Langfuse projection over the captured LLM stream. + +Exercises `LangfuseExporter` with a fake Langfuse client so we can assert the +reconstructed trace tree (trace ids, parent linkage, names, usage, trace-level +user attributes, session-as-metadata) without a real Langfuse backend. +""" + +from __future__ import annotations + +import pytest + +from src.config import settings +from src.llm.backend import CompletionResult, ToolCallResult +from src.llm.capture import build_captured_call +from src.llm.types import LLMTelemetryContext +from src.telemetry import langfuse_session +from src.telemetry.langfuse_exporter import LangfuseExporter + + +class FakeOtelSpan: + def __init__(self) -> None: + self.attributes: dict[str, object] = {} + + def set_attribute(self, key: str, value: object) -> None: + self.attributes[key] = value + + +class FakeObs: + _counter = 0 + + def __init__(self, **kwargs: object) -> None: + FakeObs._counter += 1 + self.id = f"obs-{FakeObs._counter}" + self.kwargs = kwargs + self._otel_span = FakeOtelSpan() + self.ended = False + + def end(self) -> None: + self.ended = True + + +class FakeClient: + def __init__(self) -> None: + self.observations: list[FakeObs] = [] + + def create_trace_id(self, *, seed: str | None = None) -> str: + return f"lf-{seed}" + + def start_observation(self, **kwargs: object) -> FakeObs: + obs = FakeObs(**kwargs) + self.observations.append(obs) + return obs + + +@pytest.fixture(autouse=True) +def _exporter_env(monkeypatch: pytest.MonkeyPatch): + """Enable the exporter and install a fake langfuse client + clean registry.""" + monkeypatch.setattr(settings, "LANGFUSE_PUBLIC_KEY", "pk-test") + monkeypatch.setattr(settings, "LANGFUSE_EXPORTER_MODE", "exporter") + monkeypatch.setattr(settings, "NAMESPACE", "tenant1") + client = FakeClient() + import langfuse + + monkeypatch.setattr(langfuse, "get_client", lambda: client) + langfuse_session.reset() + FakeObs._counter = 0 + yield client + langfuse_session.reset() + + +def _call( + *, + run_id: str | None, + trace_id: str, + iteration: int | None = None, + step_seq: int = 0, + attempt: int = 1, + session_id: str | None = None, + track_name: str | None = None, + agent_type: str = "dialectic", + parent_category: str = "dialectic", + tool_names: list[str] | None = None, + finish_reason: str = "stop", + content: str = "answer", +): + telemetry = LLMTelemetryContext( + workspace_name="ws", + call_purpose="dialectic.answer", + parent_category=parent_category, + agent_type=agent_type, + run_id=run_id, + trace_id=trace_id, + span_id=trace_id, + session_id=session_id, + track_name=track_name, + iteration=iteration, + step_seq=step_seq, + ) + result = CompletionResult( + content=content, + input_tokens=10, + output_tokens=5, + cache_read_input_tokens=2, + finish_reason=finish_reason, + tool_calls=[ + ToolCallResult(id=f"tc-{i}", name=name, input={"q": name}) + for i, name in enumerate(tool_names or []) + ], + ) + return build_captured_call( + telemetry=telemetry, + transport="anthropic", + provider_label=None, + model="claude-x", + messages=[{"role": "user", "content": "q"}], + tools=None, + tool_choice=None, + result=result, + attempt=attempt, + was_fallback=False, + was_stream=False, + finish_reason=finish_reason, + ) + + +def test_single_shot_generation_is_trace_root(_exporter_env: FakeClient): + # Deriver/summarizer style: run_id None → no run/step span, generation is root. + client = _exporter_env + LangfuseExporter().export( + _call(run_id=None, trace_id="t1", track_name="Minimal Deriver") + ) + + assert len(client.observations) == 1 + gen = client.observations[0] + assert gen.kwargs["as_type"] == "generation" + assert gen.kwargs["trace_context"] == {"trace_id": "lf-t1"} + assert gen.kwargs["model"] == "claude-x" + assert gen.kwargs["usage_details"] == { + "input": 10, + "output": 5, + "cache_read_input_tokens": 2, + "cache_creation_input_tokens": 0, + } + # Trace attrs stamped on the root generation; no session (session_id None). + assert gen._otel_span.attributes.get("user.id") == "tenant1" + assert "session.id" not in gen._otel_span.attributes + + +def test_agentic_run_builds_run_step_generation(_exporter_env: FakeClient): + client = _exporter_env + LangfuseExporter().export( + _call( + run_id="r1", + trace_id="r1", + iteration=1, + session_id="sess_abc", + track_name="Dialectic Agent", + ) + ) + + by_type: dict[str, list[FakeObs]] = {} + for obs in client.observations: + by_type.setdefault(str(obs.kwargs["as_type"]), []).append(obs) + assert len(by_type["span"]) == 2 # run span + step span + assert len(by_type["generation"]) == 1 + + run_span, step_span = by_type["span"] + gen = by_type["generation"][0] + assert run_span.kwargs["trace_context"] == {"trace_id": "lf-r1"} + assert step_span.kwargs["trace_context"] == { + "trace_id": "lf-r1", + "parent_span_id": run_span.id, + } + assert gen.kwargs["trace_context"] == { + "trace_id": "lf-r1", + "parent_span_id": step_span.id, + } + # Trace attrs stamped once, on the run span (the root). The Honcho session is + # NOT a Langfuse session (one-shot queries aren't a conversation thread) — it + # rides in metadata as a correlation key instead. + assert "session.id" not in run_span._otel_span.attributes + assert run_span._otel_span.attributes["user.id"] == "tenant1" + assert run_span._otel_span.attributes["langfuse.trace.name"] == "Dialectic Agent" + assert run_span.kwargs["metadata"]["honcho_session"] == "sess_abc" + + +def test_run_span_created_once_across_iterations(_exporter_env: FakeClient): + client = _exporter_env + exporter = LangfuseExporter() + exporter.export(_call(run_id="r1", trace_id="r1", iteration=1, session_id="s")) + exporter.export(_call(run_id="r1", trace_id="r1", iteration=2, session_id="s")) + + spans = [o for o in client.observations if o.kwargs["as_type"] == "span"] + gens = [o for o in client.observations if o.kwargs["as_type"] == "generation"] + # One run span shared, one step span per iteration, one generation per call. + assert len(gens) == 2 + assert len(spans) == 3 # 1 run + 2 step + # Trace attrs (user/name) stamped exactly once across the whole run. + stamped = [o for o in client.observations if "user.id" in o._otel_span.attributes] + assert len(stamped) == 1 + + +def test_langfuse_session_lru_evicts_least_recently_used( + monkeypatch: pytest.MonkeyPatch, +): + """Past _MAX_TRACES the least-recently-touched trace is evicted (not refused), + so an active trace keeps its remembered span ids no matter the run volume.""" + langfuse_session.reset() + monkeypatch.setattr(langfuse_session, "_MAX_TRACES", 2) + + langfuse_session.ensure_run_span("t1", "b", lambda _s: "t1-span") + langfuse_session.ensure_run_span("t2", "b", lambda _s: "t2-span") + # Touch t1 so t2 becomes the least-recently-used trace. + assert ( + langfuse_session.ensure_run_span("t1", "b", lambda _s: "ignored") == "t1-span" + ) + # A third trace evicts the LRU trace (t2), keeping t1. + langfuse_session.ensure_run_span("t3", "b", lambda _s: "t3-span") + + created: list[str] = [] + # t1 still tracked → remembered span returned, create NOT re-invoked. + assert ( + langfuse_session.ensure_run_span( + "t1", "b", lambda _s: created.append("t1") or "new" + ) + == "t1-span" + ) + assert created == [] + # t2 was evicted → fresh state, create IS re-invoked. + assert ( + langfuse_session.ensure_run_span( + "t2", "b", lambda _s: created.append("t2") or "t2-span2" + ) + == "t2-span2" + ) + assert created == ["t2"] + + +def test_error_finish_marks_generation_level(_exporter_env: FakeClient): + client = _exporter_env + LangfuseExporter().export( + _call(run_id=None, trace_id="t1", finish_reason="error", content="") + ) + gen = client.observations[0] + assert gen.kwargs["level"] == "ERROR" + + +@pytest.mark.parametrize( + ("attr", "value"), + [ + ("LANGFUSE_EXPORTER_MODE", "inline"), # exporter off in inline mode + ("LANGFUSE_PUBLIC_KEY", None), # exporter off without a public key + ], +) +def test_exporter_disabled_emits_nothing( + _exporter_env: FakeClient, + monkeypatch: pytest.MonkeyPatch, + attr: str, + value: object, +): + client = _exporter_env + monkeypatch.setattr(settings, attr, value) + LangfuseExporter().export(_call(run_id="r1", trace_id="r1", iteration=1)) + assert client.observations == [] + + +def test_generation_name_uses_generation_suffix(_exporter_env: FakeClient): + client = _exporter_env + LangfuseExporter().export( + _call(run_id="r1", trace_id="r1", iteration=1, track_name="Dialectic Agent") + ) + gen = [o for o in client.observations if o.kwargs["as_type"] == "generation"][0] + assert gen.kwargs["name"] == "Dialectic Agent generation" + step = [o for o in client.observations if o.kwargs["as_type"] == "span"][1] + assert step.kwargs["name"] == "Dialectic Agent step" + + +def test_tool_calls_become_spans_under_the_step(_exporter_env: FakeClient): + client = _exporter_env + LangfuseExporter().export( + _call( + run_id="r1", + trace_id="r1", + iteration=1, + track_name="Dialectic Agent", + tool_names=["search_memory", "search_messages"], + ) + ) + + spans = [o for o in client.observations if o.kwargs["as_type"] == "span"] + gen = [o for o in client.observations if o.kwargs["as_type"] == "generation"][0] + tools = [o for o in client.observations if o.kwargs["as_type"] == "tool"] + step_span = spans[1] # run span, then step span + + assert [t.kwargs["name"] for t in tools] == ["search_memory", "search_messages"] + # Tool spans are siblings of the generation: same parent (the step span). + for t in tools: + assert t.kwargs["trace_context"]["parent_span_id"] == step_span.id + assert gen.kwargs["trace_context"]["parent_span_id"] == step_span.id + # The model's requested input args ride on the tool span. + assert tools[0].kwargs["input"] == {"q": "search_memory"} + + +def test_only_the_root_span_keeps_as_root(_exporter_env: FakeClient): + # The SDK stamps AS_ROOT on every trace_context span; the exporter must + # demote children so exactly one root survives — otherwise Langfuse races to + # pick the trace name/root and names the trace after a child span. + from langfuse import LangfuseOtelSpanAttributes as Attr + + client = _exporter_env + LangfuseExporter().export( + _call( + run_id="r1", + trace_id="r1", + iteration=1, + track_name="Dialectic Agent", + tool_names=["search_memory"], + ) + ) + + def is_demoted(obs: FakeObs) -> bool: + return obs._otel_span.attributes.get(Attr.AS_ROOT) is False + + spans = [o for o in client.observations if o.kwargs["as_type"] == "span"] + run_span, step_span = spans[0], spans[1] + gen = [o for o in client.observations if o.kwargs["as_type"] == "generation"][0] + tools = [o for o in client.observations if o.kwargs["as_type"] == "tool"] + + # Exactly one root: the run span is never demoted; everything with a real + # parent is. + assert not is_demoted(run_span) + assert is_demoted(step_span) + assert is_demoted(gen) + assert all(is_demoted(t) for t in tools) + demoted = [o for o in client.observations if is_demoted(o)] + assert len(demoted) == len(client.observations) - 1 + + +def test_single_shot_generation_keeps_as_root(_exporter_env: FakeClient): + # No parent → the generation is the trace root and must not be demoted. + from langfuse import LangfuseOtelSpanAttributes as Attr + + client = _exporter_env + LangfuseExporter().export( + _call(run_id=None, trace_id="t1", track_name="Minimal Deriver") + ) + gen = client.observations[0] + assert gen._otel_span.attributes.get(Attr.AS_ROOT) is not False + + +def test_single_shot_tool_calls_are_skipped(_exporter_env: FakeClient): + # No step span to anchor to (deriver-style); tools don't orphan to the root. + client = _exporter_env + LangfuseExporter().export( + _call(run_id=None, trace_id="t1", tool_names=["search_memory"]) + ) + assert [o.kwargs["as_type"] for o in client.observations] == ["generation"] + + +def test_dreamer_specialists_nest_under_one_dream_root(_exporter_env: FakeClient): + # Both specialists share ONE dream trace (run_id) and both start at + # iteration 1. They must nest under a single synthetic "Dream" root (so the + # trace has one root, not one per specialist) while staying distinct + # sub-trees (no step-span collision). + from langfuse import LangfuseOtelSpanAttributes as Attr + + client = _exporter_env + exporter = LangfuseExporter() + for agent_type in ("deduction", "induction"): + exporter.export( + _call( + run_id="dream1", + trace_id="dream1", + iteration=1, + agent_type=agent_type, + parent_category="dream", + track_name=f"Dreamer/{agent_type}", + ) + ) + + by_name: dict[str, list[FakeObs]] = {} + for o in client.observations: + by_name.setdefault(str(o.kwargs["name"]), []).append(o) + gens = [o for o in client.observations if o.kwargs["as_type"] == "generation"] + + def is_demoted(o: FakeObs) -> bool: + return o._otel_span.attributes.get(Attr.AS_ROOT) is False + + # Exactly one trace root: the synthetic "Dream" span — no parent, not demoted. + roots = [ + o + for o in client.observations + if o.kwargs["trace_context"] == {"trace_id": "lf-dream1"} + ] + assert len(roots) == 1 + dream_root = roots[0] + assert dream_root.kwargs["name"] == "Dream" + assert dream_root.kwargs["as_type"] == "span" + assert not is_demoted(dream_root) + + # Both specialist run spans hang off the Dream root and are demoted. + run_dd = by_name["Dreamer/deduction"][0] + run_in = by_name["Dreamer/induction"][0] + assert len(by_name["Dreamer/deduction"]) == 1 + assert len(by_name["Dreamer/induction"]) == 1 + for rs in (run_dd, run_in): + assert rs.kwargs["trace_context"] == { + "trace_id": "lf-dream1", + "parent_span_id": dream_root.id, + } + assert is_demoted(rs) + + # One step span per specialist, parented to its own run span; no collapsing. + assert len(by_name["Dreamer/deduction step"]) == 1 + assert len(by_name["Dreamer/induction step"]) == 1 + assert ( + by_name["Dreamer/deduction step"][0].kwargs["trace_context"]["parent_span_id"] + == run_dd.id + ) + assert ( + by_name["Dreamer/induction step"][0].kwargs["trace_context"]["parent_span_id"] + == run_in.id + ) + + # Each generation nests under its OWN specialist's step. + assert len({g.kwargs["trace_context"]["parent_span_id"] for g in gens}) == 2 + + # Trace name is the branch-agnostic "Dream", stamped exactly once — on the + # Dream root, not on a specialist's run span. + named = [ + o + for o in client.observations + if o._otel_span.attributes.get("langfuse.trace.name") + ] + assert len(named) == 1 + assert named[0] is dream_root + assert named[0]._otel_span.attributes["langfuse.trace.name"] == "Dream" diff --git a/tests/telemetry/test_representation_v2_fields.py b/tests/telemetry/test_representation_v2_fields.py index 669da96b..a41bed41 100644 --- a/tests/telemetry/test_representation_v2_fields.py +++ b/tests/telemetry/test_representation_v2_fields.py @@ -2,8 +2,6 @@ """tests for RepresentationCompletedEvent additive fields + truncation. Targets: -- Schema stays at v2 (additive, no bump). Existing `input_tokens` semantics - unchanged. - fields are defaultable (no breakage for callers that ignore them) and round-trip through Pydantic serialization. - `HonchoLLMCallResponse.hit_input_token_cap` defaults to False but can be @@ -17,10 +15,6 @@ from src.telemetry.events.representation import RepresentationCompletedEvent class TestRepresentationV2AdditiveFields: - def test_schema_stays_at_v2(self): - """is additive — schema_version must NOT bump to 3.""" - assert RepresentationCompletedEvent.schema_version() == 2 - def test_new_fields_are_optional(self): """Existing callers must keep working without supplying any new fields. All new fields default.""" @@ -53,6 +47,10 @@ class TestRepresentationV2AdditiveFields: assert event.hit_batch_token_cap is False assert event.hit_input_token_cap is False assert event.observer_count == 0 + assert event.exact_dup_in_batch_count == 0 + assert event.exact_dup_existing_count == 0 + assert event.semantic_dup_rejected_count == 0 + assert event.semantic_dup_replaced_count == 0 def test_input_tokens_semantics_preserved(self): """The downstream metering key must remain 'queued-message tokens'. @@ -159,10 +157,41 @@ class TestRepresentationV2AdditiveFields: "hit_batch_token_cap", "hit_input_token_cap", "observer_count", + "exact_dup_in_batch_count", + "exact_dup_existing_count", + "semantic_dup_rejected_count", + "semantic_dup_replaced_count", ): assert field in data, f"missing field: {field}" assert data["hit_batch_token_cap"] is True + def test_dedup_count_fields_round_trip(self): + event = RepresentationCompletedEvent( + workspace_name="ws", + session_name="s", + observed="user", + queue_items_processed=1, + earliest_message_id="m1", + latest_message_id="m1", + message_count=1, + explicit_conclusion_count=0, + context_preparation_ms=10.0, + llm_call_ms=100.0, + total_duration_ms=110.0, + input_tokens=100, + total_input_tokens=200, + output_tokens=50, + exact_dup_in_batch_count=2, + exact_dup_existing_count=3, + semantic_dup_rejected_count=4, + semantic_dup_replaced_count=5, + ) + data = event.model_dump(mode="json") + assert data["exact_dup_in_batch_count"] == 2 + assert data["exact_dup_existing_count"] == 3 + assert data["semantic_dup_rejected_count"] == 4 + assert data["semantic_dup_replaced_count"] == 5 + class TestHitInputTokenCapFlag: """`HonchoLLMCallResponse.hit_input_token_cap` is the bridge between the diff --git a/tests/telemetry/test_sentry_before_send.py b/tests/telemetry/test_sentry_before_send.py new file mode 100644 index 00000000..7e63b555 --- /dev/null +++ b/tests/telemetry/test_sentry_before_send.py @@ -0,0 +1,82 @@ +"""Tests for the shared Sentry before_send filter. + +default_before_send runs in every entrypoint (API + deriver). It drops known +non-actionable exceptions and collapses DB connection-pool checkout timeouts +into a single warning-level issue so they stop spawning a fresh error issue per +transaction (fleet-wide saturation symptom, tracked in DEV-1852). +""" + +from typing import TYPE_CHECKING, cast + +import pytest +import sentry_sdk +from fastapi.exceptions import RequestValidationError +from pydantic import ValidationError +from sqlalchemy.exc import OperationalError + +from src.exceptions import ResourceNotFoundException +from src.telemetry.sentry import default_before_send, initialize_sentry + +if TYPE_CHECKING: + from sentry_sdk._types import Event, Hint + + +def _hint(exc: BaseException) -> "Hint": + return cast("Hint", {"exc_info": (type(exc), exc, None)}) + + +def _event(**kwargs: object) -> "Event": + return cast("Event", cast(object, dict(kwargs))) + + +def test_connection_timeout_is_consolidated_and_downgraded() -> None: + exc = OperationalError("SELECT 1", {}, Exception("connection timeout expired")) + out = default_before_send({}, _hint(exc)) + assert out == { + "fingerprint": ["honcho-db-connection-timeout"], + "level": "warning", + } + + +def test_unrelated_operational_error_passes_through() -> None: + exc = OperationalError("SELECT 1", {}, Exception("some other db failure")) + event = _event(level="error") + assert default_before_send(event, _hint(exc)) == {"level": "error"} + + +def test_honcho_and_validation_errors_are_dropped() -> None: + assert default_before_send({}, _hint(ResourceNotFoundException("nope"))) is None + assert ( + default_before_send({}, _hint(ValidationError.from_exception_data("x", []))) + is None + ) + assert default_before_send({}, _hint(RequestValidationError([]))) is None + + +def test_events_without_exc_info_pass_through() -> None: + event = _event(release="1.0") + assert default_before_send(event, None) == {"release": "1.0"} + assert default_before_send(event, cast("Hint", {})) == {"release": "1.0"} + + +def _captured_before_send(monkeypatch: pytest.MonkeyPatch, **kwargs: object) -> object: + captured: dict[str, object] = {} + + def fake_init(**init_kwargs: object) -> None: + captured.update(init_kwargs) + + monkeypatch.setattr(sentry_sdk, "init", fake_init) + initialize_sentry(integrations=[], **kwargs) # pyright: ignore[reportArgumentType] + return captured["before_send"] + + +def test_initialize_sentry_defaults_to_shared_filter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + assert _captured_before_send(monkeypatch) is default_before_send + + +def test_initialize_sentry_explicit_none_bypasses_shared_filter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + assert _captured_before_send(monkeypatch, before_send=None) is None diff --git a/tests/telemetry/test_sentry_sampler.py b/tests/telemetry/test_sentry_sampler.py new file mode 100644 index 00000000..c0d08a08 --- /dev/null +++ b/tests/telemetry/test_sentry_sampler.py @@ -0,0 +1,60 @@ +"""Tests for the Sentry traces sampler. + +The sampler must drop high-volume infra/scrape transactions (health checks, +Prometheus scrapes, OpenAPI schema, docs) while sampling real traffic at the +configured rate. These endpoints otherwise dominate transaction + profiling +volume and drown out useful traces. +""" + +import pytest + +from src.config import settings +from src.telemetry.sentry import traces_sampler + + +@pytest.mark.parametrize( + "path", + ["/metrics", "/health", "/openapi.json", "/docs", "/redoc"], +) +def test_infra_paths_are_dropped(path: str) -> None: + """ASGI requests to infra/scrape paths get a 0.0 sample rate.""" + assert traces_sampler({"asgi_scope": {"path": path}}) == 0.0 + + +@pytest.mark.parametrize( + "name", + [ + "src.telemetry.prometheus.metrics.metrics_endpoint", + "src.prometheus.metrics", + "fastapi.applications.FastAPI.setup..openapi", + ], +) +def test_infra_transaction_names_are_dropped(name: str) -> None: + """Transactions without an ASGI path still drop by their endpoint name.""" + assert traces_sampler({"transaction_context": {"name": name}}) == 0.0 + + +def test_real_route_uses_default_rate() -> None: + """A normal API route is sampled at the configured default rate.""" + ctx = { + "asgi_scope": {"path": "/v3/peers/alice/chat"}, + "transaction_context": {"name": "src.routers.peers.chat"}, + } + assert traces_sampler(ctx) == settings.SENTRY.TRACES_SAMPLE_RATE + + +def test_parent_sampling_decision_is_respected() -> None: + """When continuing a distributed trace, inherit the upstream decision.""" + assert traces_sampler({"parent_sampled": True}) == 1.0 + assert traces_sampler({"parent_sampled": False}) == 0.0 + + +def test_infra_path_overrides_parent_decision() -> None: + """Infra paths are dropped even if an upstream trace was sampled in.""" + ctx = {"asgi_scope": {"path": "/metrics"}, "parent_sampled": True} + assert traces_sampler(ctx) == 0.0 + + +def test_empty_context_falls_back_to_default_rate() -> None: + """A context with no scope, name, or parent uses the default rate.""" + assert traces_sampler({}) == settings.SENTRY.TRACES_SAMPLE_RATE diff --git a/tests/telemetry/test_trace_events.py b/tests/telemetry/test_trace_events.py new file mode 100644 index 00000000..35da6492 --- /dev/null +++ b/tests/telemetry/test_trace_events.py @@ -0,0 +1,218 @@ +"""Tests for full-fidelity trace events, dedup, and the CloudEvents exporter.""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import Any + +import pytest + +from src.llm.backend import CompletionResult +from src.llm.capture import CapturedLLMCall, build_captured_call +from src.telemetry import trace_session +from src.telemetry.events.trace import LLMCallTracedEvent, TraceContentEvent + + +class TestLLMCallTracedEvent: + def test_metadata(self): + assert LLMCallTracedEvent.event_type() == "llm.call.traced" + assert LLMCallTracedEvent.category() == "trace" + # Ground-truth — never sampled (the system of record). + assert LLMCallTracedEvent.volume_class() == "ground_truth" + + def test_resource_id_format(self): + event = LLMCallTracedEvent( + span_id="s1", + iteration=2, + attempt=1, + step_seq=3, + transport="anthropic", + model="m", + ) + # {span_id}:{iteration}:{attempt}:{step_seq} — tool_call_seq dropped. + assert event.get_resource_id() == "s1:2:1:3" + + def test_keeps_default_evt_id(self): + event = LLMCallTracedEvent( + span_id="s1", + iteration=1, + attempt=1, + step_seq=1, + transport="anthropic", + model="m", + ) + event_id = event.generate_id() + assert event_id.startswith("evt_") + assert len(event_id) == 26 + + +class TestTraceContentEvent: + def test_metadata(self): + assert TraceContentEvent.event_type() == "trace.content" + assert TraceContentEvent.category() == "trace" + assert TraceContentEvent.volume_class() == "ground_truth" + + def test_resource_id_is_content_hash(self): + event = TraceContentEvent(content_hash="sha256:abc", role="user", content="hi") + assert event.get_resource_id() == "sha256:abc" + + def test_generate_id_is_content_addressed_and_timestamp_free(self): + # Two instances with the same hash but DIFFERENT timestamps must collide + # on id, so cross-process/retry re-sends dedupe at the transport layer. + import datetime + + a = TraceContentEvent( + content_hash="sha256:deadbeefdeadbeefdeadbeef", + role="user", + content="hi", + timestamp=datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC), + ) + b = TraceContentEvent( + content_hash="sha256:deadbeefdeadbeefdeadbeef", + role="user", + content="hi", + timestamp=datetime.datetime(2026, 6, 22, tzinfo=datetime.UTC), + ) + assert a.generate_id() == b.generate_id() + assert a.generate_id().startswith("content_") + # Different content → different id. + c = TraceContentEvent(content_hash="sha256:other", role="user", content="hi") + assert c.generate_id() != a.generate_id() + + +class TestTraceSessionDedup: + def setup_method(self): + trace_session.reset() + + def teardown_method(self): + trace_session.reset() + + def test_first_emit_true_repeat_false(self): + assert trace_session.mark_emitted("run-1", "h1") is True + assert trace_session.mark_emitted("run-1", "h1") is False # already shipped + assert trace_session.mark_emitted("run-1", "h2") is True # new hash + + def test_runs_are_independent(self): + assert trace_session.mark_emitted("run-1", "h1") is True + assert trace_session.mark_emitted("run-2", "h1") is True # different run + + def test_lru_evicts_least_recently_used_run(self, monkeypatch: pytest.MonkeyPatch): + # Shrink the window so eviction is testable without _MAX_RUNS runs. + monkeypatch.setattr(trace_session, "_MAX_RUNS", 2) + trace_session.mark_emitted("run-1", "h1") + trace_session.mark_emitted("run-2", "h1") + # Touch run-1 so run-2 becomes the least-recently-used run. + trace_session.mark_emitted("run-1", "h2") + # A third run evicts the LRU run (run-2), keeping run-1. + trace_session.mark_emitted("run-3", "h1") + # run-1 is still tracked → its already-shipped hash stays deduped. + assert trace_session.mark_emitted("run-1", "h1") is False + # run-2 was evicted → its hash ships again as if a fresh run. + assert trace_session.mark_emitted("run-2", "h1") is True + + +class _FakeTraceEmitter: + """Stand-in for the trace emitter that records emitted events.""" + + def __init__(self) -> None: + self.events: list[object] = [] + + def emit(self, event: object) -> None: + self.events.append(event) + + +@pytest.fixture +def trace_on(monkeypatch: pytest.MonkeyPatch) -> Iterator[_FakeTraceEmitter]: + """Enable payload tracing and route emit_trace at a fake emitter.""" + from src.config import settings + from src.telemetry import emitter as emitter_mod + + monkeypatch.setattr(settings.TELEMETRY, "TRACE_PAYLOADS_ENABLED", True) + fake = _FakeTraceEmitter() + monkeypatch.setattr(emitter_mod, "_trace_emitter", fake) + trace_session.reset() + yield fake + trace_session.reset() + + +def _captured( + messages: list[dict[str, Any]], *, content: str = "answer", run: str = "r1" +) -> CapturedLLMCall: + from src.llm.types import LLMTelemetryContext + + return build_captured_call( + telemetry=LLMTelemetryContext( + workspace_name="ws", + call_purpose="dialectic.answer", + parent_category="dialectic", + run_id=run, + trace_id=run, + span_id=run, + iteration=1, + step_seq=1, + ), + transport="anthropic", + provider_label=None, + model="claude-x", + messages=messages, + tools=None, + tool_choice=None, + result=CompletionResult(content=content, finish_reason="stop"), + attempt=1, + was_fallback=False, + was_stream=False, + finish_reason="stop", + ) + + +class TestTraceExporter: + def test_refs_match_emitted_content(self, trace_on: _FakeTraceEmitter): + from src.telemetry.trace_exporter import TraceExporter + + call = _captured([{"role": "user", "content": "q"}]) + TraceExporter().export(call) + + traced = [e for e in trace_on.events if isinstance(e, LLMCallTracedEvent)] + contents = [e for e in trace_on.events if isinstance(e, TraceContentEvent)] + assert len(traced) == 1 + # input message + output content → two content events. + emitted_hashes = {c.content_hash for c in contents} + # Every input ref points at an emitted trace.content. + for ref in traced[0].input_message_refs: + assert ref in emitted_hashes + assert traced[0].output_content_ref in emitted_hashes + + def test_dedup_across_iterations(self, trace_on: _FakeTraceEmitter): + from src.telemetry.trace_exporter import TraceExporter + + exporter = TraceExporter() + shared = {"role": "user", "content": "system context"} + # Iteration 1: messages [shared]; iteration 2: [shared, follow-up]. + exporter.export(_captured([shared])) + before = sum(isinstance(e, TraceContentEvent) for e in trace_on.events) + exporter.export(_captured([shared, {"role": "user", "content": "more"}])) + after = sum(isinstance(e, TraceContentEvent) for e in trace_on.events) + # `shared` already shipped this run → only the new message (+ output if + # not already seen) emit again; `shared` is NOT re-emitted. + shared_hash = None + for e in trace_on.events: + if isinstance(e, TraceContentEvent) and e.content == "system context": + shared_hash = e.content_hash + emitted_shared = [ + e + for e in trace_on.events + if isinstance(e, TraceContentEvent) and e.content_hash == shared_hash + ] + assert len(emitted_shared) == 1 # shipped once across both iterations + assert after > before # the new message did ship + + def test_purpose_allowlist_filters( + self, trace_on: _FakeTraceEmitter, monkeypatch: pytest.MonkeyPatch + ): + from src.config import settings + from src.telemetry.trace_exporter import TraceExporter + + monkeypatch.setattr(settings.TELEMETRY, "TRACE_PURPOSES", ["summary.short"]) + TraceExporter().export(_captured([{"role": "user", "content": "q"}])) + # call_purpose is dialectic.answer, not in the allowlist → nothing emits. + assert trace_on.events == [] diff --git a/tests/test_advanced_filters.py b/tests/test_advanced_filters.py index 768f29bc..3b187470 100644 --- a/tests/test_advanced_filters.py +++ b/tests/test_advanced_filters.py @@ -235,6 +235,86 @@ async def test_comparison_operators_filters( ), f"Unexpected message '{message_config['content']}' found in results for {description}" +@pytest.mark.asyncio +async def test_bare_list_membership_sugar( + client: TestClient, + sample_data: tuple[Workspace, Peer], +): + """A bare list on a regular column is shorthand for {"in": [...]}. + + JSONB metadata columns are excluded from the sugar: a bare list there + keeps JSONB containment semantics. + """ + test_workspace, test_peer = sample_data + + # Second peer so peer_id membership has something to exclude + peer2_name = str(generate_nanoid()) + client.post( + f"/v3/workspaces/{test_workspace.name}/peers", + json={"id": peer2_name}, + ) + + session_id = str(generate_nanoid()) + session_response = client.post( + f"/v3/workspaces/{test_workspace.name}/sessions", + json={ + "id": session_id, + "peer_names": {test_peer.name: {}, peer2_name: {}}, + }, + ) + assert session_response.status_code == 201 + + message_configs = [ + { + "content": "From peer one", + "peer_id": test_peer.name, + "metadata": {"tags": ["important", "urgent"]}, + }, + { + "content": "From peer two", + "peer_id": peer2_name, + "metadata": {"tags": ["normal"]}, + }, + ] + messages_response = client.post( + f"/v3/workspaces/{test_workspace.name}/sessions/{session_id}/messages", + json={"messages": message_configs}, + ) + assert messages_response.status_code == 201 + + def list_contents(filter_config: dict[str, Any]) -> list[str]: + response = client.post( + f"/v3/workspaces/{test_workspace.name}/sessions/{session_id}/messages/list", + json={"filters": filter_config}, + ) + assert response.status_code == 200 + return [item["content"] for item in response.json()["items"]] + + # Bare list == membership on a regular column + assert list_contents({"peer_id": [test_peer.name]}) == ["From peer one"] + + # Multiple values + assert sorted(list_contents({"peer_id": [test_peer.name, peer2_name]})) == [ + "From peer one", + "From peer two", + ] + + # Equivalent to the explicit {"in": [...]} form + assert list_contents({"peer_id": [test_peer.name]}) == list_contents( + {"peer_id": {"in": [test_peer.name]}} + ) + + # Empty list matches nothing (fail-closed), never everything + assert list_contents({"peer_id": []}) == [] + + # JSONB metadata keeps containment semantics for bare lists: + # matches arrays containing ALL listed elements, not membership. + assert list_contents({"metadata": {"tags": ["important", "urgent"]}}) == [ + "From peer one" + ] + assert list_contents({"metadata": {"tags": ["important", "missing"]}}) == [] + + @pytest.mark.asyncio async def test_wildcard_filters( client: TestClient, sample_data: tuple[Workspace, Peer] diff --git a/tests/test_cache_redaction.py b/tests/test_cache_redaction.py new file mode 100644 index 00000000..3a8e6891 --- /dev/null +++ b/tests/test_cache_redaction.py @@ -0,0 +1,126 @@ +"""Unit tests for the cache client's _redact_cache_url helper.""" + +import pytest + +from src.cache.client import _redact_cache_url # pyright: ignore[reportPrivateUsage] + + +class TestRedactCacheUrl: + """Tests for _redact_cache_url — a security-relevant logging helper + that must never raise and must never leak a password.""" + + # --- Password masking --- + + def test_password_only_userinfo(self): + assert ( + _redact_cache_url("redis://:secret@localhost:6379/0") + == "redis://:***@localhost:6379/0" + ) + + def test_user_and_password(self): + result = _redact_cache_url("redis://user:s3cret@10.0.0.1:6380/2") + assert "***" in result + assert "s3cret" not in result + assert "user" in result + + def test_rediss_protocol(self): + result = _redact_cache_url("rediss://:secret@redis.example.com:6380") + assert result.startswith("rediss://") + assert "***" in result + assert "secret" not in result + + def test_complex_password(self): + result = _redact_cache_url("redis://:p%40ssw0rd!%24@host:6379/0") + assert "***" in result + assert "p%40ssw0rd" not in result + + def test_password_never_leaked(self): + """The original password must never appear in the redacted output.""" + for url in [ + "redis://:hunter2@localhost:6379/0", + "redis://admin:hunter2@localhost:6379/0", + "rediss://:hunter2@[::1]:6380/1", + ]: + assert "hunter2" not in _redact_cache_url(url) + + # --- Secrets in query parameters --- + # redis-py accepts ?password= (querystring options become client + # kwargs) and cashews accepts ?secret= (HMAC signing key), so both + # are real configuration paths that must not reach the logs. + + @pytest.mark.parametrize("param", ["password", "secret", "PASSWORD"]) + def test_query_param_secret_masked(self, param: str): + result = _redact_cache_url(f"redis://host:6379/0?{param}=s3cret") + assert "s3cret" not in result + assert f"{param}=***" in result + + def test_query_param_masking_preserves_other_params(self): + result = _redact_cache_url("redis://host:6379/0?db=1&password=s3cret&ssl=true") + assert "s3cret" not in result + assert "db=1" in result + assert "ssl=true" in result + + def test_userinfo_and_query_secret_both_masked(self): + result = _redact_cache_url("redis://:hunter2@host:6379/0?secret=s3cret") + assert "hunter2" not in result + assert "s3cret" not in result + + def test_non_secret_query_params_unchanged(self): + url = "redis://localhost:6379/0?suppress=true" + assert _redact_cache_url(url) == url + + # --- No-password URLs (returned unchanged) --- + + def test_user_without_password_unchanged(self): + assert ( + _redact_cache_url("redis://user@localhost:6379/0") + == "redis://user@localhost:6379/0" + ) + + def test_in_memory_url_unchanged(self): + assert _redact_cache_url("mem://") == "mem://" + + # --- IPv6 --- + + def test_ipv6_brackets_preserved(self): + result = _redact_cache_url("rediss://:secret@[::1]:6380/1") + assert "[::1]" in result + assert "***" in result + assert "secret" not in result + + # --- Malformed URLs (must NOT raise) --- + + def test_invalid_port_redacts_password(self): + """Regression test for two review findings: accessing + ``parsed.port`` on a URL with a non-numeric port raises + ``ValueError`` (must not crash startup inside an except block), + and the fallback must never echo the raw URL back — the + password has to be masked even when the port is unparseable. + """ + result = _redact_cache_url("redis://:pass@host:notaport/0") + assert "pass" not in result + assert "***" in result + + def test_out_of_range_port_redacts_password(self): + result = _redact_cache_url("redis://:supersecret@host:99999/0") + assert "supersecret" not in result + assert "***" in result + + def test_unparseable_url_never_echoed(self): + # Unbalanced IPv6 bracket makes urlparse itself raise; the + # fallback must return a placeholder, not the raw input. + result = _redact_cache_url("redis://:secret@[::1:6379/0") + assert "secret" not in result + + def test_missing_scheme_never_echoed(self): + # Without "redis://" urlparse sees no netloc, so the userinfo + # (and its password) is invisible to .password — the string + # must not be echoed back. + result = _redact_cache_url(":hunter2@host:6379/0") + assert "hunter2" not in result + + def test_garbage_input_does_not_raise(self): + assert isinstance(_redact_cache_url("not a url at all"), str) + + def test_empty_string_does_not_raise(self): + assert isinstance(_redact_cache_url(""), str) diff --git a/tests/test_config.py b/tests/test_config.py index 86ea0994..7730c751 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -7,7 +7,9 @@ def _make_deriver_settings( *, MAX_INPUT_TOKENS: int = 25000, MAX_CUSTOM_INSTRUCTIONS_TOKENS: int = 2000, - REPRESENTATION_BATCH_MAX_TOKENS: int = 1024, + REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS: int = 512, + REPRESENTATION_BATCH_TARGET_INPUT_TOKENS: int = 1024, + REPRESENTATION_BATCH_MAX_AGE_SECONDS: int = 1800, ) -> DeriverSettings: return DeriverSettings( MODEL_CONFIG=ConfiguredModelSettings( @@ -16,7 +18,9 @@ def _make_deriver_settings( ), MAX_INPUT_TOKENS=MAX_INPUT_TOKENS, MAX_CUSTOM_INSTRUCTIONS_TOKENS=MAX_CUSTOM_INSTRUCTIONS_TOKENS, - REPRESENTATION_BATCH_MAX_TOKENS=REPRESENTATION_BATCH_MAX_TOKENS, + REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS=REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS, + REPRESENTATION_BATCH_TARGET_INPUT_TOKENS=REPRESENTATION_BATCH_TARGET_INPUT_TOKENS, + REPRESENTATION_BATCH_MAX_AGE_SECONDS=REPRESENTATION_BATCH_MAX_AGE_SECONDS, ) @@ -25,6 +29,7 @@ def test_deriver_defaults_enable_custom_instructions_at_supported_cap() -> None: assert settings.MAX_INPUT_TOKENS == 25000 assert settings.MAX_CUSTOM_INSTRUCTIONS_TOKENS == 2000 + assert settings.REPRESENTATION_BATCH_MAX_AGE_SECONDS == 1800 def test_custom_instructions_tokens_can_be_disabled_with_zero() -> None: @@ -36,3 +41,43 @@ def test_custom_instructions_tokens_can_be_disabled_with_zero() -> None: def test_custom_instructions_tokens_cannot_exceed_supported_cap() -> None: with pytest.raises(ValueError, match="less than or equal to 2000"): _make_deriver_settings(MAX_CUSTOM_INSTRUCTIONS_TOKENS=2001) + + +def test_representation_batch_age_can_be_disabled_with_zero() -> None: + settings = _make_deriver_settings(REPRESENTATION_BATCH_MAX_AGE_SECONDS=0) + + assert settings.REPRESENTATION_BATCH_MAX_AGE_SECONDS == 0 + + +def test_representation_batch_age_rejects_negative_values() -> None: + with pytest.raises(ValueError, match="greater than or equal to 0"): + _make_deriver_settings(REPRESENTATION_BATCH_MAX_AGE_SECONDS=-1) + + +def test_representation_batch_work_unit_target_can_be_disabled_with_zero() -> None: + settings = _make_deriver_settings(REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS=0) + + assert settings.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS == 0 + + +def test_representation_batch_work_unit_target_rejects_negative_values() -> None: + with pytest.raises(ValueError, match="greater than or equal to 0"): + _make_deriver_settings(REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS=-1) + + +def test_representation_batch_tokens_can_diverge() -> None: + settings = _make_deriver_settings( + REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS=4096, + REPRESENTATION_BATCH_TARGET_INPUT_TOKENS=1024, + ) + + assert settings.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS == 4096 + assert settings.REPRESENTATION_BATCH_TARGET_INPUT_TOKENS == 1024 + + +def test_representation_batch_target_input_cannot_exceed_max_input_tokens() -> None: + with pytest.raises(ValueError, match="cannot exceed max deriver input tokens"): + _make_deriver_settings( + MAX_INPUT_TOKENS=1000, + REPRESENTATION_BATCH_TARGET_INPUT_TOKENS=2048, + ) diff --git a/tests/test_db_resilience.py b/tests/test_db_resilience.py new file mode 100644 index 00000000..faffa89a --- /dev/null +++ b/tests/test_db_resilience.py @@ -0,0 +1,310 @@ +"""Unit tests for DB connection resilience + observability. + +These are DB-free: they exercise the application_name checkout hook against a +fake DBAPI connection, the deriver polling backoff math, and the in-flight gauge +listeners directly. +""" + +from types import SimpleNamespace +from typing import Any + +import pytest +from sqlalchemy.ext.asyncio import AsyncSession + +import src.db as db_module +from src.config import settings +from src.db import DBQueryInflightTracker +from src.telemetry.prometheus.metrics import db_queries_in_flight_gauge + + +def test_session_local_uses_vanilla_async_session() -> None: + """Regression guard: no custom session subclass / acquisition logic. + + Connection acquisition is a single lazy checkout owned by AsyncSession; there + must be no re-introduced eager-checkout or retry hooks on the session. + """ + session = db_module.SessionLocal() + assert type(session) is AsyncSession + assert not hasattr(session, "_ensure_acquired") + assert not hasattr(session, "_honcho_acquired") + + +# --- application_name checkout hook ------------------------------------------ + + +class _FakeCursor: + def __init__(self, recorder: list[Any], raise_exc: Exception | None) -> None: + self.recorder: list[Any] = recorder + self.raise_exc: Exception | None = raise_exc + self.closed: bool = False + + def execute(self, sql: str, params: Any = None) -> None: + if self.raise_exc is not None: + raise self.raise_exc + self.recorder.append((sql, params)) + + def close(self) -> None: + self.closed = True + + +class _FakeDBAPIConn: + def __init__(self, recorder: list[Any], raise_exc: Exception | None = None) -> None: + self._cursor: _FakeCursor = _FakeCursor(recorder, raise_exc) + # Real pooled connections are checked out in non-autocommit mode; the + # hook flips this to True for its statement then restores it so it never + # leaves an open transaction that would block the read engine's + # AUTOCOMMIT switch. + self.autocommit: bool = False + + def cursor(self) -> _FakeCursor: + return self._cursor + + +def test_checkout_hook_sets_application_name_from_request_context() -> None: + recorder: list[Any] = [] + conn = _FakeDBAPIConn(recorder) + token = db_module.request_context.set("request:trace-ctx") + try: + db_module._set_application_name_on_checkout(conn, None, None) # pyright: ignore[reportPrivateUsage] + finally: + db_module.request_context.reset(token) + + assert len(recorder) == 1 + sql, params = recorder[0] + assert "set_config" in sql and "application_name" in sql + assert params == ("request:trace-ctx",) + assert conn._cursor.closed is True # pyright: ignore[reportPrivateUsage] + # The hook restored the original (non-autocommit) mode after its statement. + assert conn.autocommit is False + + +def test_checkout_hook_defaults_to_unknown_without_context() -> None: + recorder: list[Any] = [] + conn = _FakeDBAPIConn(recorder) + token = db_module.request_context.set(None) + try: + db_module._set_application_name_on_checkout(conn, None, None) # pyright: ignore[reportPrivateUsage] + finally: + db_module.request_context.reset(token) + + assert recorder[0][1] == ("unknown",) + + +def test_checkout_hook_swallows_errors() -> None: + """A failure tagging the connection must never break the checkout.""" + conn = _FakeDBAPIConn([], raise_exc=RuntimeError("boom")) + # Must not raise. + db_module._set_application_name_on_checkout(conn, None, None) # pyright: ignore[reportPrivateUsage] + + +# --- deriver polling backoff math -------------------------------------------- + + +def test_polling_backoff_sequence_and_reset(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_ENABLED", True) + monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0) + monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_MULTIPLIER", 2.0) + monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_MAX_INTERVAL_SECONDS", 30.0) + # Disable jitter so the schedule is asserted exactly (jitter is tested + # separately in tests/deriver/test_queue_processing.py::TestPollingJitter). + monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) + + from src.deriver.queue_manager import QueueManager + + qm = QueueManager() + seq = [qm._advance_poll_interval() for _ in range(8)] # pyright: ignore[reportPrivateUsage] + assert seq == [1.0, 2.0, 4.0, 8.0, 16.0, 30.0, 30.0, 30.0] # caps at max + + qm._reset_poll_interval() # pyright: ignore[reportPrivateUsage] + assert qm._advance_poll_interval() == 1.0 # pyright: ignore[reportPrivateUsage] + + +def test_polling_backoff_disabled_stays_constant( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_ENABLED", False) + monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0) + monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) + + from src.deriver.queue_manager import QueueManager + + qm = QueueManager() + assert [qm._advance_poll_interval() for _ in range(3)] == [1.0, 1.0, 1.0] # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_polling_loop_idle_sleeps_once_per_cycle( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Drive the real loop on an empty queue: exactly one (growing, capped) + sleep per empty poll — no double-sleep from the queue_empty_flag branch.""" + monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_ENABLED", True) + monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0) + monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_MULTIPLIER", 2.0) + monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_MAX_INTERVAL_SECONDS", 8.0) + monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) + + import asyncio + + from src.deriver import queue_manager as qm_mod + + qm = qm_mod.QueueManager() + sleeps: list[float] = [] + polls = {"n": 0} + + async def fake_cleanup() -> None: + return None + + async def fake_claim() -> dict[str, str]: + polls["n"] += 1 + if polls["n"] >= 5: + qm.shutdown_event.set() # stop after 5 empty polls + return {} + + async def fake_sleep(seconds: float) -> None: + sleeps.append(seconds) + + monkeypatch.setattr(qm, "cleanup_stale_work_units", fake_cleanup) + monkeypatch.setattr(qm, "get_and_claim_work_units", fake_claim) + # queue_manager calls asyncio.sleep on the stdlib module; patch it there. + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + await qm.polling_loop() + + # One sleep per empty poll, growing 1->2->4->8 then capped at 8 (not doubled). + assert sleeps == [1.0, 2.0, 4.0, 8.0, 8.0] + + +def test_inflight_gauge_no_drift(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(settings.METRICS, "NAMESPACE", "test") + child: Any = db_queries_in_flight_gauge.labels(instance_type="api") + tracker = DBQueryInflightTracker(child) + key = DBQueryInflightTracker.INFLIGHT_KEY + + def value() -> float: + return float(child._value.get()) + + start = value() + conn = SimpleNamespace(info={}) + + # Normal execute: before -> after returns to baseline. + tracker.on_before(conn) + assert value() == start + 1 + assert conn.info[key] is True + tracker.on_after(conn) + assert value() == start + assert key not in conn.info + + # Errored execute: before -> on_error decrements (after never fires). + tracker.on_before(conn) + assert value() == start + 1 + tracker.on_error(SimpleNamespace(connection=conn)) + assert value() == start + + # on_error without a matching before (e.g. connect error) must not push the + # gauge negative. + tracker.on_error(SimpleNamespace(connection=SimpleNamespace(info={}))) + assert value() == start + + +@pytest.mark.asyncio +async def test_stale_cleanup_time_gate(monkeypatch: pytest.MonkeyPatch) -> None: + """cleanup_stale_work_units runs at most once per gate interval per + instance (staleness is a minutes-timescale condition; per-poll cleanup + multiplies into needless fleet-wide write transactions). First poll always + runs it so a crashed predecessor's stale rows are recovered immediately.""" + monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) + monkeypatch.setattr( + settings.DERIVER, "STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS", 60.0 + ) + + from src.deriver import queue_manager as qm_mod + + qm = qm_mod.QueueManager() + runs = {"n": 0} + + async def fake_cleanup() -> None: + runs["n"] += 1 + + monkeypatch.setattr(qm, "cleanup_stale_work_units", fake_cleanup) + + clock = {"now": 1_000.0} + monkeypatch.setattr( + "src.deriver.queue_manager.time.monotonic", lambda: clock["now"] + ) + + # First call runs (no prior attempt recorded). + await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] + assert runs["n"] == 1 + + # Inside the gate window: skipped. + clock["now"] += 10.0 + await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] + assert runs["n"] == 1 + + # Past the gate window: runs again. + clock["now"] += 60.0 + await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] + assert runs["n"] == 2 + + +@pytest.mark.asyncio +async def test_stale_cleanup_gate_failed_attempt_waits_full_interval( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The gate records the ATTEMPT before running, so a failing cleanup is not + retried on every poll against a DB that is already struggling.""" + monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) + monkeypatch.setattr( + settings.DERIVER, "STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS", 60.0 + ) + + from src.deriver import queue_manager as qm_mod + + qm = qm_mod.QueueManager() + attempts = {"n": 0} + + async def failing_cleanup() -> None: + attempts["n"] += 1 + raise RuntimeError("db unavailable") + + monkeypatch.setattr(qm, "cleanup_stale_work_units", failing_cleanup) + + clock = {"now": 1_000.0} + monkeypatch.setattr( + "src.deriver.queue_manager.time.monotonic", lambda: clock["now"] + ) + + with pytest.raises(RuntimeError): + await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] + assert attempts["n"] == 1 + + # Immediately after the failure: still gated, no hammering. + clock["now"] += 1.0 + await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] + assert attempts["n"] == 1 + + +@pytest.mark.asyncio +async def test_stale_cleanup_gate_zero_interval_runs_every_poll( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Interval 0.0 preserves legacy run-on-every-poll behavior.""" + monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) + monkeypatch.setattr( + settings.DERIVER, "STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS", 0.0 + ) + + from src.deriver import queue_manager as qm_mod + + qm = qm_mod.QueueManager() + runs = {"n": 0} + + async def fake_cleanup() -> None: + runs["n"] += 1 + + monkeypatch.setattr(qm, "cleanup_stale_work_units", fake_cleanup) + + await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] + await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] + assert runs["n"] == 2 diff --git a/tests/test_dependencies.py b/tests/test_dependencies.py index c90643ec..1b5c219b 100644 --- a/tests/test_dependencies.py +++ b/tests/test_dependencies.py @@ -16,6 +16,12 @@ class FakeSession: self.execute_calls: list[tuple[Any, ...]] = [] self.rollback_calls: int = 0 self.close_calls: int = 0 + self.connection_calls: int = 0 + + async def connection(self) -> None: + # Tracks checkout attempts so tests can assert get_db/tracked_db stay + # lazy (they should never force a checkout themselves). + self.connection_calls += 1 async def execute(self, statement: Any, params: Any = None) -> None: self.execute_calls.append((statement, params)) @@ -31,26 +37,24 @@ class FakeSession: @pytest.mark.asyncio -async def test_get_db_sets_application_name_when_tracing_enabled( +async def test_get_db_yields_lazily_without_checkout_or_tracing( monkeypatch: pytest.MonkeyPatch, ) -> None: + # get_db must NOT touch the connection or run set_config itself — checkout + # happens lazily inside the AsyncSession on first DB use, so a handler doing + # non-DB work before its first query never pins a connection. fake_db = FakeSession() monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db) - monkeypatch.setattr(settings.DB, "TRACING", True) + monkeypatch.setattr(settings.DB, "TRACING", True) # still no set_config here - context_token = request_context.set("request:test-ctx") dep_gen = real_get_db() - try: db = await anext(dep_gen) assert db is fake_db - assert len(fake_db.execute_calls) == 1 - stmt, params = fake_db.execute_calls[0] - assert "set_config" in str(stmt) - assert params == {"name": "request:test-ctx"} + assert fake_db.connection_calls == 0 # no eager checkout + assert fake_db.execute_calls == [] # no set_config in get_db finally: await dep_gen.aclose() - request_context.reset(context_token) assert fake_db.rollback_calls == 1 # unconditional rollback in finally assert fake_db.close_calls == 1 @@ -80,7 +84,6 @@ async def test_tracked_db_creates_and_resets_task_context( ) -> None: fake_db = FakeSession() monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db) - monkeypatch.setattr(settings.DB, "TRACING", True) monkeypatch.setattr( uuid, "uuid4", @@ -90,15 +93,12 @@ async def test_tracked_db_creates_and_resets_task_context( clear_token = request_context.set(None) try: async with real_tracked_db("cleanup_job"): + # tracked_db sets the task context so the lazy session can read it. assert request_context.get() == "task:cleanup_job:12345678" finally: request_context.reset(clear_token) assert request_context.get() is None - assert len(fake_db.execute_calls) == 1 - stmt, params = fake_db.execute_calls[0] - assert "set_config" in str(stmt) - assert params == {"name": "task:cleanup_job:12345678"} assert fake_db.rollback_calls == 1 # unconditional rollback in finally assert fake_db.close_calls == 1 @@ -109,7 +109,6 @@ async def test_tracked_db_preserves_existing_request_context( ) -> None: fake_db = FakeSession() monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db) - monkeypatch.setattr(settings.DB, "TRACING", True) context_token = request_context.set("request:existing") try: @@ -118,10 +117,6 @@ async def test_tracked_db_preserves_existing_request_context( finally: request_context.reset(context_token) - assert len(fake_db.execute_calls) == 1 - stmt, params = fake_db.execute_calls[0] - assert "set_config" in str(stmt) - assert params == {"name": "request:existing"} assert fake_db.rollback_calls == 1 # unconditional rollback in finally assert fake_db.close_calls == 1 @@ -155,3 +150,161 @@ async def test_tracked_db_rolls_back_open_transaction_on_exit( assert fake_db.rollback_calls == 1 assert fake_db.close_calls == 1 + + +@pytest.mark.asyncio +async def test_tracked_db_read_only_uses_read_sessionmaker( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # read_only=True must construct the session from ReadSessionLocal (the + # AUTOCOMMIT engine) — and never touch SessionLocal — while keeping the + # same rollback/close teardown. + read_fake = FakeSession() + monkeypatch.setattr(dependencies_module, "ReadSessionLocal", lambda: read_fake) + monkeypatch.setattr( + dependencies_module, + "SessionLocal", + lambda: pytest.fail("read_only window constructed a write session"), + ) + + async with real_tracked_db("read_op", read_only=True) as db: + assert db is read_fake + + assert read_fake.rollback_calls == 1 + assert read_fake.close_calls == 1 + + +@pytest.mark.asyncio +async def test_get_read_db_rolls_back_and_closes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + read_fake = FakeSession() + monkeypatch.setattr(dependencies_module, "ReadSessionLocal", lambda: read_fake) + + dep_gen = dependencies_module.get_read_db() + try: + db = await anext(dep_gen) + assert db is read_fake + assert read_fake.connection_calls == 0 # still lazy, no eager checkout + finally: + await dep_gen.aclose() + + assert read_fake.rollback_calls == 1 + assert read_fake.close_calls == 1 + + +def test_read_engine_is_autocommit_and_shares_pool() -> None: + # The read engine must differ from the write engine ONLY by isolation + # level: AUTOCOMMIT (so reads never autobegin a transaction) on the same + # underlying pool (no second connection budget). + from src.db import engine, read_engine + + assert ( + read_engine.sync_engine._execution_options.get( # pyright: ignore[reportPrivateUsage] + "isolation_level" + ) + == "AUTOCOMMIT" + ) + assert read_engine.sync_engine.pool is engine.sync_engine.pool + + +@pytest.mark.asyncio +async def test_read_only_session_runs_in_autocommit_on_the_wire() -> None: + # Wire-level guarantee behind the whole read-path fix: a read_only session's + # connection has the DBAPI autocommit flag set, so psycopg emits no BEGIN + # and the backend sits in state 'idle' (not 'idle in transaction') after a + # statement returns. NOTE: get_isolation_level() can NOT verify this — it + # reports the server's transaction_isolation GUC (READ COMMITTED), because + # autocommit is a driver behavior, not a server isolation level. + from sqlalchemy import text + + from src.db import read_engine + + async with real_tracked_db("read_op", read_only=True) as db: + pid = (await db.execute(text("SELECT pg_backend_pid()"))).scalar() + conn = await db.connection() + raw = (await conn.get_raw_connection()).driver_connection + assert raw is not None + assert raw.autocommit is True + + # Definitive check, from a second connection: after the SELECT above, + # the session's backend must be plain 'idle' — an open transaction + # would report 'idle in transaction' and be reapable in production. + async with read_engine.connect() as observer: + state = ( + await observer.execute( + text("SELECT state FROM pg_stat_activity WHERE pid = :p"), + {"p": pid}, + ) + ).scalar() + assert state == "idle" + + +@pytest.mark.asyncio +async def test_write_session_holds_idle_in_transaction_after_select() -> None: + # Contrast guard documenting WHY the read engine exists: the default + # (transactional) session autobegins on the first statement and leaves the + # backend 'idle in transaction' until rollback/close — the state that + # Postgres's idle_in_transaction_session_timeout reaps and that pins a + # transaction-mode pooler backend. + from sqlalchemy import text + + from src.db import read_engine + + async with real_tracked_db("write_op") as db: + pid = (await db.execute(text("SELECT pg_backend_pid()"))).scalar() + async with read_engine.connect() as observer: + state = ( + await observer.execute( + text("SELECT state FROM pg_stat_activity WHERE pid = :p"), + {"p": pid}, + ) + ).scalar() + assert state == "idle in transaction" + + +@pytest.mark.asyncio +async def test_read_only_session_works_with_tracing_checkout_hook() -> None: + # Regression: the DB.TRACING checkout hook runs set_config() at pool + # checkout, BEFORE the dialect applies the read engine's AUTOCOMMIT + # isolation level. If that statement is allowed to autobegin a transaction, + # psycopg then refuses to switch the connection into AUTOCOMMIT + # ("can't change 'autocommit' now: connection in transaction") and every + # read_only session 500s under TRACING. The hook must run in autocommit so + # it leaves the connection idle. This combination is otherwise untested + # because DB.TRACING defaults to false. + from sqlalchemy import event, text + + from src.db import ( + _set_application_name_on_checkout, # pyright: ignore[reportPrivateUsage] + engine, + read_engine, + ) + + context_token = request_context.set("tracing-regression") + event.listen(engine.sync_engine, "checkout", _set_application_name_on_checkout) + try: + async with real_tracked_db("read_op", read_only=True) as db: + pid = (await db.execute(text("SELECT pg_backend_pid()"))).scalar() + app_name = (await db.execute(text("SHOW application_name"))).scalar() + conn = await db.connection() + raw = (await conn.get_raw_connection()).driver_connection + assert raw is not None + # AUTOCOMMIT was applied despite the checkout hook running first. + assert raw.autocommit is True + # The hook still tagged the connection (set_config is session-scoped, + # so it survives the autocommit boundary). + assert app_name == "tracing-regression" + # Backend is idle, not idle-in-transaction: the no-BEGIN guarantee + # holds even with the hook firing. + async with read_engine.connect() as observer: + state = ( + await observer.execute( + text("SELECT state FROM pg_stat_activity WHERE pid = :p"), + {"p": pid}, + ) + ).scalar() + assert state == "idle" + finally: + event.remove(engine.sync_engine, "checkout", _set_application_name_on_checkout) + request_context.reset(context_token) diff --git a/tests/test_generate_jwt_script.py b/tests/test_generate_jwt_script.py new file mode 100644 index 00000000..04011974 --- /dev/null +++ b/tests/test_generate_jwt_script.py @@ -0,0 +1,23 @@ +import sys + +import pytest + +from scripts import generate_jwt + + +def test_admin_cannot_be_combined_with_scoped_flags(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + sys, + "argv", + [ + "generate_jwt.py", + "--admin", + "--workspace", + "my-workspace", + ], + ) + + with pytest.raises(SystemExit) as exc_info: + generate_jwt.main() + + assert exc_info.value.code == 2 diff --git a/tests/test_security.py b/tests/test_security.py new file mode 100644 index 00000000..1785be8b --- /dev/null +++ b/tests/test_security.py @@ -0,0 +1,287 @@ +"""Auth scope tests — DEV-1736 regression coverage. + +Prior to this fix `auth()` walked the route's declared scope first and fell +through to a workspace check, so a `{w, p}` token authorized any peer in `w`. +The contract now is: authorize by the token's narrowest claim, never widen. +""" + +from contextlib import asynccontextmanager + +import jwt as pyjwt +import pytest +from fastapi.security import HTTPAuthorizationCredentials + +from src.config import settings +from src.exceptions import AuthenticationException, ValidationException +from src.security import JWTParams, auth, create_jwt, verify_jwt + + +@pytest.fixture(autouse=True) +def _enable_auth(monkeypatch: pytest.MonkeyPatch): # pyright: ignore[reportUnusedFunction] + monkeypatch.setattr(settings.AUTH, "USE_AUTH", True) + monkeypatch.setattr(settings.AUTH, "JWT_SECRET", "test-secret") + + +def _bearer(token: str) -> HTTPAuthorizationCredentials: + return HTTPAuthorizationCredentials(scheme="Bearer", credentials=token) + + +class TestVerifyJWTShape: + def test_peer_token_without_workspace_rejected(self): + token = pyjwt.encode({"p": "alice"}, b"test-secret", algorithm="HS256") + with pytest.raises(AuthenticationException): + verify_jwt(token) + + def test_session_token_without_workspace_rejected(self): + token = pyjwt.encode({"s": "sess-1"}, b"test-secret", algorithm="HS256") + with pytest.raises(AuthenticationException): + verify_jwt(token) + + def test_workspace_only_token_ok(self): + token = create_jwt(JWTParams(w="ws-a")) + params = verify_jwt(token) + assert params.w == "ws-a" + + def test_workspace_peer_token_ok(self): + token = create_jwt(JWTParams(w="ws-a", p="alice")) + params = verify_jwt(token) + assert params.w == "ws-a" + assert params.p == "alice" + + +class TestAuthPeerScope: + """`{w: ws-a, p: alice}` may only act on alice in ws-a.""" + + @pytest.mark.asyncio + async def test_matches_own_peer(self): + creds = _bearer(create_jwt(JWTParams(w="ws-a", p="alice"))) + params = await auth(credentials=creds, workspace_name="ws-a", peer_name="alice") + assert params.p == "alice" + + @pytest.mark.asyncio + async def test_denies_sibling_peer_same_workspace(self): + """The original bug: peer-scoped token fell through to workspace auth.""" + creds = _bearer(create_jwt(JWTParams(w="ws-a", p="alice"))) + with pytest.raises(AuthenticationException): + await auth(credentials=creds, workspace_name="ws-a", peer_name="bob") + + @pytest.mark.asyncio + async def test_denies_workspace_route_with_no_peer(self): + """Peer-scoped token cannot use workspace-listing routes.""" + creds = _bearer(create_jwt(JWTParams(w="ws-a", p="alice"))) + with pytest.raises(AuthenticationException): + await auth(credentials=creds, workspace_name="ws-a") + + @pytest.mark.asyncio + async def test_self_authorizing_route_receives_claims(self): + """Body-scoped routes use require_auth() and compare claims in-handler.""" + creds = _bearer(create_jwt(JWTParams(w="ws-a", p="alice"))) + params = await auth(credentials=creds) + assert params.w == "ws-a" + assert params.p == "alice" + + @pytest.mark.asyncio + async def test_denies_cross_workspace(self): + creds = _bearer(create_jwt(JWTParams(w="ws-a", p="alice"))) + with pytest.raises(AuthenticationException): + await auth(credentials=creds, workspace_name="ws-b", peer_name="alice") + + +class TestAuthSessionScope: + @pytest.mark.asyncio + async def test_matches_own_session(self): + creds = _bearer(create_jwt(JWTParams(w="ws-a", s="sess-1"))) + params = await auth( + credentials=creds, workspace_name="ws-a", session_name="sess-1" + ) + assert params.s == "sess-1" + + @pytest.mark.asyncio + async def test_denies_sibling_session_same_workspace(self): + creds = _bearer(create_jwt(JWTParams(w="ws-a", s="sess-1"))) + with pytest.raises(AuthenticationException): + await auth(credentials=creds, workspace_name="ws-a", session_name="sess-2") + + @pytest.mark.asyncio + async def test_denies_workspace_route_with_no_session(self): + creds = _bearer(create_jwt(JWTParams(w="ws-a", s="sess-1"))) + with pytest.raises(AuthenticationException): + await auth(credentials=creds, workspace_name="ws-a") + + @pytest.mark.asyncio + async def test_self_authorizing_route_receives_claims(self): + creds = _bearer(create_jwt(JWTParams(w="ws-a", s="sess-1"))) + params = await auth(credentials=creds) + assert params.w == "ws-a" + assert params.s == "sess-1" + + +class TestAuthWorkspaceScope: + @pytest.mark.asyncio + async def test_matches_workspace(self): + creds = _bearer(create_jwt(JWTParams(w="ws-a"))) + params = await auth(credentials=creds, workspace_name="ws-a") + assert params.w == "ws-a" + + @pytest.mark.asyncio + async def test_workspace_token_reaches_peer_route(self): + """Workspace tokens still authorize narrower routes inside the workspace.""" + creds = _bearer(create_jwt(JWTParams(w="ws-a"))) + params = await auth(credentials=creds, workspace_name="ws-a", peer_name="alice") + assert params.w == "ws-a" + + @pytest.mark.asyncio + async def test_denies_cross_workspace(self): + creds = _bearer(create_jwt(JWTParams(w="ws-a"))) + with pytest.raises(AuthenticationException): + await auth(credentials=creds, workspace_name="ws-b") + + @pytest.mark.asyncio + async def test_passes_self_authorizing_route(self): + """Routes with no declared scope (e.g. POST /v3/workspaces) self-authorize + on the token's `w`. The auth dependency must let workspace tokens through.""" + creds = _bearer(create_jwt(JWTParams(w="ws-a"))) + params = await auth(credentials=creds) + assert params.w == "ws-a" + + +@asynccontextmanager +async def _fake_tracked_db(*_args: object, **_kwargs: object): + """Stand-in for tracked_db; the membership query itself is monkeypatched.""" + yield None + + +def _patch_membership(monkeypatch: pytest.MonkeyPatch, *, is_member: bool): + async def _is_peer_in_session(*_args: object, **_kwargs: object) -> bool: + return is_member + + # Names are resolved via lazy imports inside auth(), so patch the source + # modules rather than the security namespace. + monkeypatch.setattr("src.dependencies.tracked_db", _fake_tracked_db) + monkeypatch.setattr("src.crud.session.is_peer_in_session", _is_peer_in_session) + + +class TestAuthMemberRead: + """Peer-scoped key gets read-only access to sessions it is a member of.""" + + @pytest.mark.asyncio + async def test_member_peer_allowed_on_read_route( + self, monkeypatch: pytest.MonkeyPatch + ): + _patch_membership(monkeypatch, is_member=True) + creds = _bearer(create_jwt(JWTParams(w="ws-a", p="alice"))) + params = await auth( + credentials=creds, + workspace_name="ws-a", + session_name="sess-1", + allow_member_read=True, + ) + assert params.p == "alice" + + @pytest.mark.asyncio + async def test_non_member_peer_denied_on_read_route( + self, monkeypatch: pytest.MonkeyPatch + ): + _patch_membership(monkeypatch, is_member=False) + creds = _bearer(create_jwt(JWTParams(w="ws-a", p="alice"))) + with pytest.raises(AuthenticationException): + await auth( + credentials=creds, + workspace_name="ws-a", + session_name="sess-1", + allow_member_read=True, + ) + + @pytest.mark.asyncio + async def test_member_peer_denied_on_write_route( + self, monkeypatch: pytest.MonkeyPatch + ): + """Write routes never set allow_member_read, so membership is irrelevant.""" + _patch_membership(monkeypatch, is_member=True) + creds = _bearer(create_jwt(JWTParams(w="ws-a", p="alice"))) + with pytest.raises(AuthenticationException): + await auth( + credentials=creds, + workspace_name="ws-a", + session_name="sess-1", + allow_member_read=False, + ) + + @pytest.mark.asyncio + async def test_member_peer_denied_cross_workspace( + self, monkeypatch: pytest.MonkeyPatch + ): + _patch_membership(monkeypatch, is_member=True) + creds = _bearer(create_jwt(JWTParams(w="ws-a", p="alice"))) + with pytest.raises(AuthenticationException): + await auth( + credentials=creds, + workspace_name="ws-b", + session_name="sess-1", + allow_member_read=True, + ) + + @pytest.mark.asyncio + async def test_session_token_has_no_cross_scope_to_peer_routes(self): + """A session key never reaches peer routes, even with allow_member_read.""" + creds = _bearer(create_jwt(JWTParams(w="ws-a", s="sess-1"))) + with pytest.raises(AuthenticationException): + await auth( + credentials=creds, + workspace_name="ws-a", + peer_name="alice", + allow_member_read=True, + ) + + +class TestCreateKeyValidation: + @pytest.mark.asyncio + async def test_peer_key_without_workspace_rejected(self): + from src.routers.keys import create_key + + with pytest.raises(ValidationException): + await create_key(workspace_id=None, peer_id="alice", session_id=None) + + @pytest.mark.asyncio + async def test_session_key_without_workspace_rejected(self): + from src.routers.keys import create_key + + with pytest.raises(ValidationException): + await create_key(workspace_id=None, peer_id=None, session_id="sess-1") + + @pytest.mark.asyncio + async def test_peer_key_with_workspace_ok(self): + from src.routers.keys import create_key + + result = await create_key(workspace_id="ws-a", peer_id="alice", session_id=None) + assert "key" in result + + +class TestAuthAdminAndUnscoped: + @pytest.mark.asyncio + async def test_admin_passes_any_route(self): + creds = _bearer(create_jwt(JWTParams(ad=True))) + params = await auth(credentials=creds, workspace_name="ws-a", peer_name="alice") + assert params.ad is True + + @pytest.mark.asyncio + async def test_non_admin_token_denied_on_admin_route(self): + creds = _bearer(create_jwt(JWTParams(w="ws-a"))) + with pytest.raises(AuthenticationException): + await auth(credentials=creds, admin=True) + + @pytest.mark.asyncio + async def test_unscoped_token_on_self_authorizing_route(self): + """A token with no scope claims and a route with no declared scope is the + escape hatch for routes that introspect jwt_params themselves.""" + creds = _bearer(create_jwt(JWTParams())) + params = await auth(credentials=creds) + assert params.w is None + assert params.p is None + assert params.s is None + + @pytest.mark.asyncio + async def test_unscoped_token_denied_on_scoped_route(self): + creds = _bearer(create_jwt(JWTParams())) + with pytest.raises(AuthenticationException): + await auth(credentials=creds, workspace_name="ws-a") diff --git a/tests/test_session_allowlist.py b/tests/test_session_allowlist.py new file mode 100644 index 00000000..aa6077cb --- /dev/null +++ b/tests/test_session_allowlist.py @@ -0,0 +1,679 @@ +""" +Tests for the session allowlist (DEV-1995). + +Covers the constrained `filters` surface on dialectic/representation +(extract_session_allowlist), fail-closed conclusion recall (search_memory), +and the strict allowlist ∩ membership intersection in message cruds. +""" + +from typing import Any +from unittest.mock import AsyncMock, patch + +import pytest +from fastapi.testclient import TestClient +from nanoid import generate as generate_nanoid +from sqlalchemy.ext.asyncio import AsyncSession + +from src import crud, models +from src.config import settings +from src.crud.message import resolve_session_scope +from src.exceptions import FilterError +from src.models import Peer, Workspace +from src.security import JWTParams, create_jwt +from src.utils.agent_tools import search_memory +from src.utils.filter import ( + MAX_SESSION_ALLOWLIST_ENTRIES, + extract_session_allowlist, +) + + +class TestExtractSessionAllowlist: + def test_none_passthrough(self): + assert extract_session_allowlist(None) is None + + def test_single_id(self): + assert extract_session_allowlist({"session_id": "s1"}) == ["s1"] + + def test_bare_list(self): + assert extract_session_allowlist({"session_id": ["s1", "s2"]}) == ["s1", "s2"] + + def test_in_operator(self): + assert extract_session_allowlist({"session_id": {"in": ["s1"]}}) == ["s1"] + + def test_dedupes_preserving_order(self): + assert extract_session_allowlist({"session_id": ["s2", "s1", "s2"]}) == [ + "s2", + "s1", + ] + + def test_empty_list_preserved_for_fail_closed(self): + assert extract_session_allowlist({"session_id": []}) == [] + + def test_unsupported_key_rejected(self): + with pytest.raises(FilterError, match="Unsupported filter key"): + extract_session_allowlist({"peer_id": ["a"], "session_id": ["s1"]}) + + def test_missing_session_id_rejected(self): + with pytest.raises(FilterError, match="must contain"): + extract_session_allowlist({}) + + def test_bad_shapes_rejected(self): + for bad in [123, {"gte": "x"}, {"in": "s1"}, [1, 2], [""], None]: + with pytest.raises(FilterError): + extract_session_allowlist({"session_id": bad}) + + def test_cap_enforced(self): + too_many = [f"s{i}" for i in range(MAX_SESSION_ALLOWLIST_ENTRIES + 1)] + with pytest.raises(FilterError, match="at most"): + extract_session_allowlist({"session_id": too_many}) + + def test_must_include_satisfied(self): + assert extract_session_allowlist( + {"session_id": ["s1", "s2"]}, must_include="s2" + ) == ["s1", "s2"] + + def test_must_include_missing_rejected(self): + with pytest.raises(FilterError, match="must be included"): + extract_session_allowlist({"session_id": ["s1"]}, must_include="s2") + + def test_must_include_ignored_without_filters(self): + assert extract_session_allowlist(None, must_include="s1") is None + + def test_must_include_none_is_no_constraint(self): + assert extract_session_allowlist({"session_id": ["s1"]}, must_include=None) == [ + "s1" + ] + + +class TestSearchMemoryAllowlist: + @pytest.mark.asyncio + async def test_allowlist_pushed_down_as_filters(self): + with patch( + "src.crud.query_documents", new=AsyncMock(return_value=[]) + ) as mock_query: + await search_memory( + workspace_name="w", + observer="o", + observed="o", + query="q", + limit=5, + levels=["explicit"], + embedding=[0.1], + session_allowlist=["s1", "s2"], + ) + assert mock_query.await_args is not None + assert mock_query.await_args.kwargs["filters"] == { + "level": {"in": ["explicit"]}, + "session_name": {"in": ["s1", "s2"]}, + } + + @pytest.mark.asyncio + async def test_allowlist_narrows_levels_to_allowlist_safe(self): + """Only levels with a trustworthy session stamp survive scoping. + + Dream-derived levels are stamped with one session but synthesized + across many (DEV-2201), so they can't be served under an allowlist. + """ + with patch( + "src.crud.query_documents", new=AsyncMock(return_value=[]) + ) as mock_query: + await search_memory( + workspace_name="w", + observer="o", + observed="o", + query="q", + limit=5, + levels=["explicit", "inductive"], + embedding=[0.1], + session_allowlist=["s1"], + ) + assert mock_query.await_args is not None + assert mock_query.await_args.kwargs["filters"]["level"] == {"in": ["explicit"]} + + @pytest.mark.asyncio + async def test_allowlist_defaults_to_explicit_when_no_levels_requested(self): + with patch( + "src.crud.query_documents", new=AsyncMock(return_value=[]) + ) as mock_query: + await search_memory( + workspace_name="w", + observer="o", + observed="o", + query="q", + limit=5, + embedding=[0.1], + session_allowlist=["s1"], + ) + assert mock_query.await_args is not None + assert mock_query.await_args.kwargs["filters"]["level"] == {"in": ["explicit"]} + + @pytest.mark.asyncio + async def test_derived_only_request_under_allowlist_returns_empty(self): + """The dialectic's derived prefetch short-circuits instead of querying.""" + with patch( + "src.crud.query_documents", new=AsyncMock(return_value=[]) + ) as mock_query: + result = await search_memory( + workspace_name="w", + observer="o", + observed="o", + query="q", + limit=5, + levels=["deductive", "inductive", "contradiction"], + embedding=[0.1], + session_allowlist=["s1"], + ) + mock_query.assert_not_awaited() + assert result.is_empty() + + @pytest.mark.asyncio + async def test_levels_untouched_without_allowlist(self): + """No allowlist means no level narrowing — unscoped recall is unchanged.""" + with patch( + "src.crud.query_documents", new=AsyncMock(return_value=[]) + ) as mock_query: + await search_memory( + workspace_name="w", + observer="o", + observed="o", + query="q", + limit=5, + levels=["deductive", "inductive"], + embedding=[0.1], + ) + assert mock_query.await_args is not None + assert mock_query.await_args.kwargs["filters"] == { + "level": {"in": ["deductive", "inductive"]} + } + + @pytest.mark.asyncio + async def test_empty_allowlist_fails_closed_without_querying(self): + with patch( + "src.crud.query_documents", new=AsyncMock(return_value=[]) + ) as mock_query: + result = await search_memory( + workspace_name="w", + observer="o", + observed="o", + query="q", + limit=5, + embedding=[0.1], + session_allowlist=[], + ) + mock_query.assert_not_awaited() + assert result.is_empty() + + +class TestMessageCrudAllowlistIntersection: + """allowlist ∩ observer-membership, fail-closed on empty intersection.""" + + async def _setup_two_sessions( + self, + client: TestClient, + workspace: Workspace, + peer: Peer, + ) -> tuple[str, str]: + ids: list[str] = [] + for marker in ("alpha", "beta"): + session_id = str(generate_nanoid()) + resp = client.post( + f"/v3/workspaces/{workspace.name}/sessions", + json={"id": session_id, "peer_names": {peer.name: {}}}, + ) + assert resp.status_code == 201 + resp = client.post( + f"/v3/workspaces/{workspace.name}/sessions/{session_id}/messages", + json={ + "messages": [ + { + "content": f"needle in {marker}", + "peer_id": peer.name, + } + ] + }, + ) + assert resp.status_code == 201 + ids.append(session_id) + return ids[0], ids[1] + + @pytest.mark.asyncio + async def test_grep_messages_intersects_allowlist( + self, + client: TestClient, + sample_data: tuple[Workspace, Peer], + ): + workspace, peer = sample_data + session_a, session_b = await self._setup_two_sessions(client, workspace, peer) + + snippets = await crud.grep_messages( + workspace_name=workspace.name, + session_name=None, + text="needle", + observer=peer.name, + session_allowlist=[session_a], + ) + contents = [m.content for matches, _ in snippets for m in matches] + assert contents == ["needle in alpha"] + + # A session the observer is NOT a member of contributes nothing, + # even when allowlisted (strict intersection). + foreign = str(generate_nanoid()) + snippets = await crud.grep_messages( + workspace_name=workspace.name, + session_name=None, + text="needle", + observer=peer.name, + session_allowlist=[foreign], + ) + assert snippets == [] + + # Both sessions allowlisted -> both found + snippets = await crud.grep_messages( + workspace_name=workspace.name, + session_name=None, + text="needle", + observer=peer.name, + session_allowlist=[session_a, session_b], + ) + assert len(snippets) == 2 + + @pytest.mark.asyncio + async def test_get_messages_by_date_range_intersects_allowlist( + self, + client: TestClient, + db_session: AsyncSession, + sample_data: tuple[Workspace, Peer], + ): + workspace, peer = sample_data + session_a, _session_b = await self._setup_two_sessions(client, workspace, peer) + + messages = await crud.get_messages_by_date_range( + db_session, + workspace_name=workspace.name, + session_name=None, + observer=peer.name, + session_allowlist=[session_a], + ) + assert [m.content for m in messages] == ["needle in alpha"] + + # Empty allowlist fails closed + messages = await crud.get_messages_by_date_range( + db_session, + workspace_name=workspace.name, + session_name=None, + observer=peer.name, + session_allowlist=[], + ) + assert messages == [] + + +class TestPeerScopedJWTAllowlistGate: + """A peer-scoped key may only allowlist sessions its peer actively belongs to. + + The gate uses `active_only=True` so it agrees with the `is_peer_in_session` + check on `options.session_id` — a peer that has left a session is denied by + both, not just one. + """ + + @pytest.fixture(autouse=True) + def _enable_auth(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(settings.AUTH, "USE_AUTH", True) + monkeypatch.setattr(settings.AUTH, "JWT_SECRET", "test-secret") + + def _chat_as( + self, + client: TestClient, + workspace: Workspace, + peer: Peer, + token: str, + body: dict[str, Any], + ): + return client.post( + f"/v3/workspaces/{workspace.name}/peers/{peer.name}/chat", + json={"query": "what do you know?", **body}, + headers={"Authorization": f"Bearer {token}"}, + ) + + async def _session_with( + self, client: TestClient, workspace: Workspace, peer: Peer + ) -> str: + session_id = str(generate_nanoid()) + resp = client.post( + f"/v3/workspaces/{workspace.name}/sessions", + json={"id": session_id, "peer_names": {peer.name: {}}}, + ) + assert resp.status_code == 201 + return session_id + + @pytest.mark.asyncio + async def test_member_sessions_allowed( + self, client: TestClient, sample_data: tuple[Workspace, Peer] + ): + workspace, peer = sample_data + session_id = await self._session_with(client, workspace, peer) + token = create_jwt(JWTParams(w=workspace.name, p=peer.name)) + + with patch( + "src.routers.peers.agentic_chat", new=AsyncMock(return_value="ok") + ) as mock_chat: + resp = self._chat_as( + client, + workspace, + peer, + token, + {"filters": {"session_id": [session_id]}}, + ) + assert resp.status_code == 200 + # The allowlist reaches the agent rather than being dropped at the gate. + assert mock_chat.await_args is not None + assert mock_chat.await_args.kwargs["session_allowlist"] == [session_id] + + @pytest.mark.asyncio + async def test_non_member_session_denied( + self, client: TestClient, sample_data: tuple[Workspace, Peer] + ): + workspace, peer = sample_data + session_id = await self._session_with(client, workspace, peer) + token = create_jwt(JWTParams(w=workspace.name, p=peer.name)) + + # One allowlisted session the peer belongs to, one it doesn't: + # membership must hold for *every* entry. + resp = self._chat_as( + client, + workspace, + peer, + token, + {"filters": {"session_id": [session_id, str(generate_nanoid())]}}, + ) + assert resp.status_code == 401 + + @pytest.mark.asyncio + async def test_left_session_denied( + self, + client: TestClient, + db_session: AsyncSession, + sample_data: tuple[Workspace, Peer], + ): + """The regression this gate's `active_only` flag exists to prevent. + + With the loose membership definition a peer that left a session still + passed here, while the adjacent `session_id` check rejected it. + """ + workspace, peer = sample_data + session_id = await self._session_with(client, workspace, peer) + token = create_jwt(JWTParams(w=workspace.name, p=peer.name)) + + await crud.remove_peers_from_session( + db_session, + workspace_name=workspace.name, + session_name=session_id, + peer_names={peer.name}, + ) + await db_session.commit() + + resp = self._chat_as( + client, workspace, peer, token, {"filters": {"session_id": [session_id]}} + ) + assert resp.status_code == 401 + + # ...and the single-session gate agrees, which is the whole point. + resp = self._chat_as(client, workspace, peer, token, {"session_id": session_id}) + assert resp.status_code == 401 + + @pytest.mark.asyncio + async def test_workspace_scoped_key_bypasses_gate( + self, client: TestClient, sample_data: tuple[Workspace, Peer] + ): + """Workspace keys are trusted callers — the allowlist passes as given.""" + workspace, peer = sample_data + token = create_jwt(JWTParams(w=workspace.name)) + foreign = str(generate_nanoid()) + + with patch("src.routers.peers.agentic_chat", new=AsyncMock(return_value="ok")): + resp = self._chat_as( + client, workspace, peer, token, {"filters": {"session_id": [foreign]}} + ) + assert resp.status_code == 200 + + @pytest.mark.asyncio + async def test_empty_allowlist_still_gated( + self, client: TestClient, sample_data: tuple[Workspace, Peer] + ): + """`filters={"session_id": []}` is a real allowlist, not an absent one. + + It must reach the gate (and pass trivially, since the empty set is a + subset of anything) rather than being skipped by a truthiness check. + """ + workspace, peer = sample_data + token = create_jwt(JWTParams(w=workspace.name, p=peer.name)) + + with patch( + "src.routers.peers.agentic_chat", new=AsyncMock(return_value="ok") + ) as mock_chat: + resp = self._chat_as( + client, workspace, peer, token, {"filters": {"session_id": []}} + ) + assert resp.status_code == 200 + assert mock_chat.await_args is not None + assert mock_chat.await_args.kwargs["session_allowlist"] == [] + + +class TestResolveSessionScope: + """The tri-state contract the four message-crud call sites depend on.""" + + @pytest.mark.asyncio + async def test_unrestricted_when_no_observer_and_no_allowlist( + self, db_session: AsyncSession, sample_data: tuple[Workspace, Peer] + ): + workspace, _ = sample_data + assert await resolve_session_scope( + db_session, workspace.name, None, None, None + ) == (None, False) + + @pytest.mark.asyncio + async def test_pinned_session_inside_allowlist_passes_through( + self, db_session: AsyncSession, sample_data: tuple[Workspace, Peer] + ): + workspace, _ = sample_data + # None (not [s1]) — the query filters on session_name directly. + assert await resolve_session_scope( + db_session, workspace.name, "s1", ["s1", "s2"], None + ) == (None, False) + + @pytest.mark.asyncio + async def test_pinned_session_outside_allowlist_denies( + self, db_session: AsyncSession, sample_data: tuple[Workspace, Peer] + ): + workspace, _ = sample_data + assert await resolve_session_scope( + db_session, workspace.name, "s3", ["s1", "s2"], None + ) == (None, True) + + @pytest.mark.asyncio + async def test_empty_allowlist_denies_rather_than_returning_empty_list( + self, db_session: AsyncSession, sample_data: tuple[Workspace, Peer] + ): + """Never returns [] — downstream stores drop an empty IN clause.""" + workspace, _ = sample_data + allowed, deny = await resolve_session_scope( + db_session, workspace.name, None, [], None + ) + assert (allowed, deny) == (None, True) + + @pytest.mark.asyncio + async def test_no_db_touched_when_no_observer_lookup_needed(self): + """Callers pass db=None on the external-vector-store path. + + The helper must not open a session of its own unless it actually needs + an observer lookup, or the external semantic lookup stops being the + first thing that happens (see + tests/integration/test_message_embeddings.py). + """ + with patch("src.crud.message.tracked_db") as mock_tracked_db: + # No observer: pinned session, unrestricted, and plain allowlist. + assert await resolve_session_scope(None, "w", "s1", None, None) == ( + None, + False, + ) + assert await resolve_session_scope(None, "w", None, None, None) == ( + None, + False, + ) + assert await resolve_session_scope(None, "w", None, ["s1"], None) == ( + ["s1"], + False, + ) + mock_tracked_db.assert_not_called() + + @pytest.mark.asyncio + async def test_observer_scope_intersected_with_allowlist( + self, + client: TestClient, + db_session: AsyncSession, + sample_data: tuple[Workspace, Peer], + ): + workspace, peer = sample_data + session_id = str(generate_nanoid()) + resp = client.post( + f"/v3/workspaces/{workspace.name}/sessions", + json={"id": session_id, "peer_names": {peer.name: {}}}, + ) + assert resp.status_code == 201 + + allowed, deny = await resolve_session_scope( + db_session, workspace.name, None, [session_id], peer.name + ) + assert (allowed, deny) == ([session_id], False) + + # Allowlisting only a session the observer isn't in denies outright. + allowed, deny = await resolve_session_scope( + db_session, workspace.name, None, [str(generate_nanoid())], peer.name + ) + assert (allowed, deny) == (None, True) + + +class TestChatRouteFilterValidation: + """Filter validation happens before any LLM work — safe to exercise.""" + + def _chat( + self, + client: TestClient, + workspace: Workspace, + peer: Peer, + body: dict[str, Any], + ): + return client.post( + f"/v3/workspaces/{workspace.name}/peers/{peer.name}/chat", + json={"query": "what do you know?", **body}, + ) + + def test_unsupported_filter_key_422( + self, client: TestClient, sample_data: tuple[Workspace, Peer] + ): + workspace, peer = sample_data + resp = self._chat(client, workspace, peer, {"filters": {"peer_id": ["x"]}}) + assert resp.status_code == 422 + + def test_bad_filter_shape_422( + self, client: TestClient, sample_data: tuple[Workspace, Peer] + ): + workspace, peer = sample_data + resp = self._chat(client, workspace, peer, {"filters": {"session_id": 42}}) + assert resp.status_code == 422 + + def test_session_id_not_in_allowlist_422( + self, client: TestClient, sample_data: tuple[Workspace, Peer] + ): + workspace, peer = sample_data + resp = self._chat( + client, + workspace, + peer, + {"session_id": "s-outside", "filters": {"session_id": ["s1", "s2"]}}, + ) + assert resp.status_code == 422 + + def test_allowlist_cap_422( + self, client: TestClient, sample_data: tuple[Workspace, Peer] + ): + workspace, peer = sample_data + too_many = [f"s{i}" for i in range(MAX_SESSION_ALLOWLIST_ENTRIES + 1)] + resp = self._chat( + client, workspace, peer, {"filters": {"session_id": too_many}} + ) + assert resp.status_code == 422 + + +class TestRepresentationRouteFilters: + @pytest.mark.asyncio + async def test_representation_scoped_by_filters( + self, + client: TestClient, + db_session: AsyncSession, + sample_data: tuple[Workspace, Peer], + ): + workspace, peer = sample_data + + session_a = models.Session( + name=str(generate_nanoid()), workspace_name=workspace.name + ) + session_b = models.Session( + name=str(generate_nanoid()), workspace_name=workspace.name + ) + db_session.add_all([session_a, session_b]) + await db_session.flush() + + collection = models.Collection( + workspace_name=workspace.name, + observer=peer.name, + observed=peer.name, + ) + db_session.add(collection) + await db_session.flush() + + db_session.add_all( + [ + models.Document( + workspace_name=workspace.name, + observer=peer.name, + observed=peer.name, + content="fact from session a", + session_name=session_a.name, + ), + models.Document( + workspace_name=workspace.name, + observer=peer.name, + observed=peer.name, + content="fact from session b", + session_name=session_b.name, + ), + models.Document( + workspace_name=workspace.name, + observer=peer.name, + observed=peer.name, + content="sessionless dream fact", + session_name=None, + ), + ] + ) + await db_session.commit() + + resp = client.post( + f"/v3/workspaces/{workspace.name}/peers/{peer.name}/representation", + json={"filters": {"session_id": [session_a.name]}}, + ) + assert resp.status_code == 200 + representation = resp.json()["representation"] + assert "fact from session a" in representation + assert "fact from session b" not in representation + assert "sessionless dream fact" not in representation + + def test_session_id_not_in_allowlist_422( + self, client: TestClient, sample_data: tuple[Workspace, Peer] + ): + workspace, peer = sample_data + resp = client.post( + f"/v3/workspaces/{workspace.name}/peers/{peer.name}/representation", + json={"session_id": "s-out", "filters": {"session_id": ["s-in"]}}, + ) + assert resp.status_code == 422 diff --git a/tests/unified/runner.py b/tests/unified/runner.py index 297b7100..b99cd37a 100644 --- a/tests/unified/runner.py +++ b/tests/unified/runner.py @@ -119,6 +119,7 @@ async def save_results_to_s3( # Create comprehensive results object timestamp = datetime.now(timezone.utc).isoformat() github_run_id = os.getenv("GITHUB_RUN_ID", "local") + github_run_attempt = os.getenv("GITHUB_RUN_ATTEMPT", "1") github_sha = os.getenv("GITHUB_SHA", "unknown") github_ref = os.getenv("GITHUB_REF_NAME", "unknown") @@ -132,6 +133,7 @@ async def save_results_to_s3( }, "metadata": { "github_run_id": github_run_id, + "github_run_attempt": github_run_attempt, "github_sha": github_sha, "github_ref": github_ref, }, @@ -145,31 +147,59 @@ async def save_results_to_s3( ], } + # One "folder" per run: /// holding results.json plus + # the reasoning-trace file(s), so a run's summary and full LLM I/O live together. date_str = datetime.now(timezone.utc).strftime("%Y-%m-%d") sha_short = github_sha[:7] if github_sha != "unknown" else "unknown" - ref_name = github_ref if github_ref != "unknown" else "unknown" - key = f"{s3_prefix}/{date_str}-{ref_name}-{sha_short}.json" + ref_slug = github_ref.replace("/", "-") # branch names may contain "/" + run_slug = f"{ref_slug}-{sha_short}-{github_run_id}-{github_run_attempt}" + run_prefix = f"{s3_prefix}/{date_str}/{run_slug}" + results_key = f"{run_prefix}/results.json" s3_client = boto3.client("s3", region_name=aws_region) # pyright: ignore s3_client.put_object( # pyright: ignore Bucket=s3_bucket, - Key=key, + Key=results_key, Body=json.dumps(comprehensive_results, indent=2).encode("utf-8"), ContentType="application/json", ) + logger.info(f"Saved test results to S3 key {results_key}") + + # Upload the reasoning traces (full LLM/deriver I/O) captured this run. The + # API and deriver both append to REASONING_TRACES_FILE (file-locked). Use + # upload_file so large trace files stream via multipart instead of buffering. + traces_path_str = os.getenv("REASONING_TRACES_FILE") + if traces_path_str: + traces_path = Path(traces_path_str) + if traces_path.is_file() and traces_path.stat().st_size > 0: + traces_key = f"{run_prefix}/{traces_path.name}" + try: + s3_client.upload_file( # pyright: ignore + str(traces_path), + s3_bucket, + traces_key, + ExtraArgs={"ContentType": "application/x-ndjson"}, + ) + logger.info(f"Saved reasoning traces to S3 key {traces_key}") + except Exception as e: + logger.error( + f"Failed to upload reasoning traces: {e}", exc_info=True + ) + else: + logger.warning( + f"REASONING_TRACES_FILE={traces_path} is missing or empty; no traces uploaded" + ) try: url: str = s3_client.generate_presigned_url( # pyright: ignore "get_object", - Params={"Bucket": s3_bucket, "Key": key}, + Params={"Bucket": s3_bucket, "Key": results_key}, ExpiresIn=259200, # 3 days ) - logger.info(f"Saved test results to s3://{s3_bucket}/{key}") - return url, key # pyright: ignore + return url, results_key # pyright: ignore except Exception as e: logger.warning(f"Could not generate S3 presigned URL: {e}") - logger.info(f"Saved test results to s3://{s3_bucket}/{key}") - return None, key + return None, results_key except Exception as e: logger.error(f"Failed to save results to S3: {e}", exc_info=True) @@ -327,6 +357,7 @@ class UnifiedTestExecutor: session=step.session_id, target=step.observed_peer_id, reasoning_level=step.reasoning_level, + response_format=step.response_format, ) return response diff --git a/tests/unified/schema.py b/tests/unified/schema.py index aa4c78ad..161d4f08 100644 --- a/tests/unified/schema.py +++ b/tests/unified/schema.py @@ -149,6 +149,9 @@ class QueryAction(TestStep): # for chat - reasoning level reasoning_level: ReasoningLevel | None = None + # for chat - optional JSON Schema the response must conform to + response_format: dict[str, Any] | None = None + assertions: list[ LLMJudgeAssertion | ContainsAssertion diff --git a/tests/unified/test_cases/dialectic_structured_output.json b/tests/unified/test_cases/dialectic_structured_output.json new file mode 100644 index 00000000..59cc1365 --- /dev/null +++ b/tests/unified/test_cases/dialectic_structured_output.json @@ -0,0 +1,135 @@ +{ + "description": "Dialectic chat with a response_format JSON Schema while the agent must use tools (reasoning off + enumeration question forces grep/search calls). Exercises the transport-layer combination of tool calling and structured output on every provider: OpenAI must avoid parse() for non-strict tools, Anthropic must skip the '{' prefill, Gemini must fall back to a schema instruction. The final answer must be a JSON string conforming to the schema.", + "workspace_config": {}, + "steps": [ + { + "step_type": "create_session", + "session_id": "structured_output_test", + "config": { + "reasoning": { + "enabled": false + } + }, + "peer_configs": { + "user": { + "observe_me": true, + "observe_others": false + }, + "assistant": { + "observe_me": false, + "observe_others": true + } + } + }, + { + "step_type": "add_messages", + "session_id": "structured_output_test", + "messages": [ + { + "peer_id": "user", + "content": "Monday I grabbed a $5 latte at Starbucks before my standup.", + "created_at": "2024-03-04T08:30:00" + }, + { + "peer_id": "assistant", + "content": "Nice, a classic way to start the week.", + "created_at": "2024-03-04T08:31:00" + }, + { + "peer_id": "user", + "content": "Tuesday I tried a $4 cold brew from Blue Bottle, really smooth.", + "created_at": "2024-03-05T09:15:00" + }, + { + "peer_id": "assistant", + "content": "Blue Bottle makes a solid cold brew.", + "created_at": "2024-03-05T09:16:00" + }, + { + "peer_id": "user", + "content": "Wednesday was a $6 oat-milk mocha at a little place downtown.", + "created_at": "2024-03-06T08:45:00" + }, + { + "peer_id": "assistant", + "content": "Oat milk mochas are underrated.", + "created_at": "2024-03-06T08:46:00" + }, + { + "peer_id": "user", + "content": "Thursday I skipped coffee and just had tea at home.", + "created_at": "2024-03-07T08:20:00" + }, + { + "peer_id": "assistant", + "content": "A calm morning, sounds good.", + "created_at": "2024-03-07T08:21:00" + }, + { + "peer_id": "user", + "content": "Friday I splurged on a $7 pour-over at the roastery near the office.", + "created_at": "2024-03-08T08:50:00" + }, + { + "peer_id": "assistant", + "content": "Ending the week strong!", + "created_at": "2024-03-08T08:51:00" + } + ] + }, + { + "step_type": "wait", + "target": "queue_empty", + "timeout": 180, + "flush": true + }, + { + "step_type": "query", + "description": "Global query + empty prefetch forces tool calls; response_format forces structured output on the same LLM calls", + "target": "chat", + "observer_peer_id": "assistant", + "observed_peer_id": "user", + "reasoning_level": "max", + "input": "How many separate coffees did I buy this week, and exactly how much did I spend in total across all of them?", + "response_format": { + "type": "object", + "properties": { + "coffee_count": { + "type": "integer", + "description": "How many separate coffees the user bought during the week" + }, + "total_spent_usd": { + "type": "number", + "description": "Total amount in US dollars the user spent on coffee" + }, + "purchases": { + "type": "array", + "items": { + "type": "string" + }, + "description": "One short entry per coffee purchase, including its price" + }, + "summary": { + "type": "string", + "description": "One-sentence answer to the question" + } + }, + "required": ["coffee_count", "total_spent_usd", "purchases", "summary"] + }, + "assertions": [ + { + "assertion_type": "json_match", + "key_value_pairs": { + "coffee_count": 4, + "total_spent_usd": 22 + } + }, + { + "assertion_type": "llm_judge", + "prompt": "The result must be a JSON object whose 'purchases' array enumerates the $5 latte, $4 cold brew, $6 mocha, and $7 pour-over (wording may vary), and whose 'summary' answers that the user bought 4 coffees for $22 total.", + "pass_if": true + } + ] + } + ] +} diff --git a/tests/unified/test_cases/dialectic_tool_calls.json b/tests/unified/test_cases/dialectic_tool_calls.json new file mode 100644 index 00000000..0ae3b923 --- /dev/null +++ b/tests/unified/test_cases/dialectic_tool_calls.json @@ -0,0 +1,103 @@ +{ + "description": "Force the dialectic agent to actually invoke tools (grep/search) by asking an enumeration+aggregation question that the prefetched first turn cannot answer in one shot. Used to verify Langfuse tool-call observations nest under the step span.", + "workspace_config": {}, + "steps": [ + { + "step_type": "create_session", + "session_id": "tool_calls_test", + "config": { + "reasoning": { + "enabled": false + } + }, + "peer_configs": { + "user": { + "observe_me": true, + "observe_others": false + }, + "assistant": { + "observe_me": false, + "observe_others": true + } + } + }, + { + "step_type": "add_messages", + "session_id": "tool_calls_test", + "messages": [ + { + "peer_id": "user", + "content": "Monday I grabbed a $5 latte at Starbucks before my standup.", + "created_at": "2024-03-04T08:30:00" + }, + { + "peer_id": "assistant", + "content": "Nice, a classic way to start the week.", + "created_at": "2024-03-04T08:31:00" + }, + { + "peer_id": "user", + "content": "Tuesday I tried a $4 cold brew from Blue Bottle, really smooth.", + "created_at": "2024-03-05T09:15:00" + }, + { + "peer_id": "assistant", + "content": "Blue Bottle makes a solid cold brew.", + "created_at": "2024-03-05T09:16:00" + }, + { + "peer_id": "user", + "content": "Wednesday was a $6 oat-milk mocha at a little place downtown.", + "created_at": "2024-03-06T08:45:00" + }, + { + "peer_id": "assistant", + "content": "Oat milk mochas are underrated.", + "created_at": "2024-03-06T08:46:00" + }, + { + "peer_id": "user", + "content": "Thursday I skipped coffee and just had tea at home.", + "created_at": "2024-03-07T08:20:00" + }, + { + "peer_id": "assistant", + "content": "A calm morning, sounds good.", + "created_at": "2024-03-07T08:21:00" + }, + { + "peer_id": "user", + "content": "Friday I splurged on a $7 pour-over at the roastery near the office.", + "created_at": "2024-03-08T08:50:00" + }, + { + "peer_id": "assistant", + "content": "Ending the week strong!", + "created_at": "2024-03-08T08:51:00" + } + ] + }, + { + "step_type": "wait", + "target": "queue_empty", + "timeout": 180, + "flush": true + }, + { + "step_type": "query", + "description": "Global query (no session history) + empty prefetch (reasoning off) forces grep/search tool calls", + "target": "chat", + "observer_peer_id": "assistant", + "observed_peer_id": "user", + "reasoning_level": "max", + "input": "How many separate coffees did I buy this week, and exactly how much did I spend in total across all of them?", + "assertions": [ + { + "assertion_type": "llm_judge", + "prompt": "Does the response state that the user bought 4 coffees totaling $22 (or correctly enumerate the $5, $4, $6, and $7 purchases)?", + "pass_if": true + } + ] + } + ] +} diff --git a/tests/utils/test_agent_tools.py b/tests/utils/test_agent_tools.py index 8ddab1bc..c13fb1e7 100644 --- a/tests/utils/test_agent_tools.py +++ b/tests/utils/test_agent_tools.py @@ -34,10 +34,13 @@ from src.utils.agent_tools import ( _handle_search_messages, # pyright: ignore[reportPrivateUsage] _handle_search_messages_temporal, # pyright: ignore[reportPrivateUsage] _handle_update_peer_card, # pyright: ignore[reportPrivateUsage] + _normalize_observation_id, # pyright: ignore[reportPrivateUsage] _validate_peer_card_entry, # pyright: ignore[reportPrivateUsage] create_observations, create_tool_executor, extract_preferences, + get_observation_context, + get_recent_history, ) # ============================================================================= @@ -126,7 +129,7 @@ async def tool_test_data( # Commit so data is visible to independent tracked_db sessions. # Tool handlers no longer share the test's db_session — they open # their own short-lived sessions via tracked_db. - # _truncate_all_tables handles cleanup between tests. + # _clear_all_tables handles cleanup between tests. await db_session.commit() yield workspace, peer1, peer2, session, messages, documents @@ -247,6 +250,70 @@ class TestCreateObservations: assert doc.level == "deductive" assert doc.source_ids == ["premise1", "premise2"] + async def test_non_deriver_context_rejects_explicit( + self, + db_session: AsyncSession, + make_tool_context: Callable[..., ToolContext], + ): + """Session-purity invariant: agents without current_messages (dreamer + specialists, dialectic) must not create explicit-level observations, + even when they pass level='explicit' to the generic tool.""" + ctx = make_tool_context(current_messages=None) + + result = await _handle_create_observations( + ctx, + { + "observations": [ + {"content": "Claims to be a doctor", "level": "explicit"}, + ] + }, + ) + + assert isinstance(result, str) + assert "ERROR" in result + assert "explicit" in result + + # Verify nothing landed in the DB + stmt = select(models.Document).where( + models.Document.content == "Claims to be a doctor" + ) + doc = (await db_session.execute(stmt)).scalar_one_or_none() + assert doc is None + + async def test_source_ids_display_prefix_is_stripped( + self, + db_session: AsyncSession, + make_tool_context: Callable[..., ToolContext], + ): + """Models sometimes copy the '[id:xxx]' display format into source_ids; + the prefix must be stripped so provenance links reference real IDs.""" + ctx = make_tool_context(current_messages=None) + + result = await _handle_create_observations( + ctx, + { + "observations": [ + { + "content": "Inferred preference for early mornings", + "source_ids": ["id:premise1", "ID:premise2"], + "premises": [ + "User schedules meetings before 9am", + "User mentions waking at 5:30", + ], + }, + ] + }, + ) + + assert "Created 1 observations" in result + + stmt = select(models.Document).where( + models.Document.content == "Inferred preference for early mornings" + ) + doc = (await db_session.execute(stmt)).scalar_one_or_none() + assert doc is not None + assert doc.source_ids == ["premise1", "premise2"] + async def test_empty_observations_list_returns_error( self, make_tool_context: Callable[..., ToolContext] ): @@ -283,10 +350,10 @@ class TestCreateObservations: observer: str, observed: str, deduplicate: bool = False, - ) -> list[Any]: + ) -> crud.CreateDocumentsResult: _ = (workspace_name, observer, observed, deduplicate) created_documents.extend(documents) - return documents + return crud.CreateDocumentsResult(created_documents=documents) monkeypatch.setattr( "src.utils.agent_tools.embedding_client.simple_batch_embed", @@ -344,10 +411,10 @@ class TestCreateObservations: observer: str, observed: str, deduplicate: bool = False, - ) -> list[Any]: + ) -> crud.CreateDocumentsResult: _ = (workspace_name, observer, observed, deduplicate) created_documents.extend(documents) - return documents + return crud.CreateDocumentsResult(created_documents=documents) monkeypatch.setattr( "src.utils.agent_tools.embedding_client.simple_batch_embed", @@ -403,10 +470,10 @@ class TestCreateObservations: observer: str, observed: str, deduplicate: bool = False, - ) -> list[Any]: + ) -> crud.CreateDocumentsResult: _ = (workspace_name, observer, observed, deduplicate) created_documents.extend(documents) - return documents + return crud.CreateDocumentsResult(created_documents=documents) monkeypatch.setattr( "src.utils.agent_tools.embedding_client.simple_batch_embed", @@ -475,6 +542,27 @@ class TestCreateObservations: create_documents.assert_not_awaited() +class TestNormalizeObservationId: + """Unit tests for _normalize_observation_id.""" + + @pytest.mark.parametrize( + "raw,expected", + [ + ("doc_abc123", "doc_abc123"), + ("id:doc_abc123", "doc_abc123"), + ("ID:doc_abc123", "doc_abc123"), + (" id:doc_abc123 ", "doc_abc123"), + ("id: doc_abc123", "doc_abc123"), + # nanoid alphabet includes '-' and '_'; these must survive untouched + ("3-bwp1hxCRkRbUh_nrqn0", "3-bwp1hxCRkRbUh_nrqn0"), + ("id:3-bwp1hxCRkRbUh_nrqn0", "3-bwp1hxCRkRbUh_nrqn0"), + ("_leading_underscore", "_leading_underscore"), + ], + ) + def test_normalization(self, raw: str, expected: str): + assert _normalize_observation_id(raw) == expected + + @pytest.mark.asyncio class TestDeleteObservations: """Tests for _handle_delete_observations.""" @@ -715,6 +803,7 @@ class TestSearchMemory: context_window: int = 2, embedding: list[float] | None = None, observer: str | None = None, + **_kwargs: Any, ) -> list[tuple[list[models.Message], list[models.Message]]]: _ = (workspace_name, session_name, query, limit, context_window, observer) fallback_embeddings.append(embedding) @@ -823,6 +912,7 @@ class TestSearchMessagesTemporal: context_window: int = 2, embedding: list[float] | None = None, observer: str | None = None, + **_kwargs: Any, ) -> list[tuple[list[models.Message], list[models.Message]]]: _ = ( workspace_name, @@ -1416,6 +1506,7 @@ class TestExtractPreferences: context_window: int, embedding: list[float] | None, observer: str | None = None, + **_kwargs: Any, ) -> list[tuple[list[models.Message], list[models.Message]]]: _ = (limit, context_window, observer) embedding_args.append(embedding) @@ -1762,3 +1853,63 @@ class TestObserverPeerNameWiring: await _handle_get_messages_by_date_range(ctx, {"after_date": "2024-01-01"}) assert captured_kwargs["observer"] == ctx.observer + + +@pytest.mark.asyncio +class TestSessionAllowlistFailClosed: + """A specific session_name outside the session_allowlist allowlist must fail closed. + + Routes guard this too, but these CRUD/tool functions are reachable directly + from the dialectic loop, so the allowlist is enforced at the boundary. + """ + + async def test_get_recent_history_respects_allowlist( + self, db_session: AsyncSession, tool_test_data: Any + ): + workspace, _peer1, peer2, session, _messages, _ = tool_test_data + + # session IS in the allowlist -> history returned + allowed = await get_recent_history( + db_session, + workspace_name=workspace.name, + session_name=session.name, + observed=peer2.name, + session_allowlist=[session.name], + ) + assert allowed # non-empty + + # session is NOT in the allowlist -> fail closed + blocked = await get_recent_history( + db_session, + workspace_name=workspace.name, + session_name=session.name, + observed=peer2.name, + session_allowlist=["some-other-session"], + ) + assert blocked == [] + + async def test_get_observation_context_fails_closed( + self, db_session: AsyncSession, tool_test_data: Any + ): + workspace, peer1, _peer2, session, messages, _ = tool_test_data + blocked = await get_observation_context( + db_session, + workspace_name=workspace.name, + session_name=session.name, + message_ids=[messages[0].id], + observer=peer1.name, + session_allowlist=["some-other-session"], + ) + assert blocked == [] + + async def test_get_messages_by_date_range_fails_closed( + self, db_session: AsyncSession, tool_test_data: Any + ): + workspace, _peer1, _peer2, session, _messages, _ = tool_test_data + blocked = await crud.get_messages_by_date_range( + db_session, + workspace_name=workspace.name, + session_name=session.name, + session_allowlist=["some-other-session"], + ) + assert blocked == [] diff --git a/tests/utils/test_clients.py b/tests/utils/test_clients.py index f5506790..49b2d75d 100644 --- a/tests/utils/test_clients.py +++ b/tests/utils/test_clients.py @@ -9,6 +9,7 @@ Tests cover: - Provider-specific features """ +import contextlib from typing import Any from unittest.mock import AsyncMock, Mock, patch @@ -37,6 +38,7 @@ from src.llm import ( honcho_llm_call, honcho_llm_call_inner, ) +from src.llm.types import LLMTelemetryContext class SampleTestModel(BaseModel): @@ -910,8 +912,10 @@ class TestMainLLMCallFunction: assert response.content == "No retry response" - async def test_track_name_updates_langfuse_span_name(self): - """track_name should rename the top-level Langfuse span.""" + async def test_track_name_on_telemetry_names_langfuse_trace_and_generation(self): + """track_name on telemetry should name the Langfuse trace + generation + per agent and stamp provider/model metadata (via propagate_attributes + + update_current_generation; see annotate_current_langfuse_trace).""" mock_llm_client = AsyncMock(spec=AsyncAnthropic) mock_response = Mock() @@ -921,11 +925,19 @@ class TestMainLLMCallFunction: mock_llm_client.messages.create = AsyncMock(return_value=mock_response) mock_langfuse_client = Mock() + captured: dict[str, Any] = {} + + @contextlib.contextmanager + def fake_propagate(**kwargs: Any): + captured.update(kwargs) + yield with ( patch.dict(CLIENTS, {"anthropic": mock_llm_client}), patch.object(settings, "LANGFUSE_PUBLIC_KEY", "test-public-key"), + patch.object(settings, "LANGFUSE_EXPORTER_MODE", "inline"), patch("langfuse.get_client", return_value=mock_langfuse_client), + patch("langfuse.propagate_attributes", fake_propagate), ): response = await honcho_llm_call( model_config=ConfiguredModelSettings( @@ -935,19 +947,152 @@ class TestMainLLMCallFunction: prompt="Hello", max_tokens=100, enable_retry=False, - track_name="Dialectic Agent", + telemetry=LLMTelemetryContext( + workspace_name="ws1", track_name="Dialectic Agent" + ), ) assert response.content == "Named response" - mock_langfuse_client.update_current_span.assert_called_once_with( - name="Dialectic Agent", - metadata={ - "namespace": settings.NAMESPACE, - "provider": "anthropic", - "model": "claude-4-sonnet", - }, + # No run_id → this single call IS the trace root: it names the trace + # and stamps metadata via propagate_attributes... + assert captured["trace_name"] == "Dialectic Agent" + assert captured["session_id"] is None + assert captured["metadata"]["namespace"] == settings.NAMESPACE + assert captured["metadata"]["provider"] == "anthropic" + assert captured["metadata"]["model"] == "claude-4-sonnet" + # ...and the generation is named + carries per-call model/metadata. + gen_calls = mock_langfuse_client.update_current_generation.call_args_list + meta_kwargs = next(c.kwargs for c in gen_calls if "model" in c.kwargs) + assert meta_kwargs["name"] == "Dialectic Agent LLM call" + assert meta_kwargs["model"] == "claude-4-sonnet" + assert meta_kwargs["metadata"]["provider"] == "anthropic" + # Input/output are stamped explicitly: @observe auto-capture is + # disabled so the live client / api-key-bearing config never reach + # the trace (HONCHO-4HA), with no loss of trace fidelity. + input_kwargs = next(c.kwargs for c in gen_calls if "input" in c.kwargs) + assert input_kwargs["input"] == [{"role": "user", "content": "Hello"}] + output_kwargs = next(c.kwargs for c in gen_calls if "output" in c.kwargs) + assert output_kwargs["output"].content == "Named response" + # Token usage is duplicated onto the generation (also in CloudEvents) + # so Langfuse renders native per-call tokens + cost. + assert output_kwargs["usage_details"]["input"] == 5 + assert output_kwargs["usage_details"]["output"] == 5 + # Tuning knobs are tracked as model_parameters (not the live client + # or api-key-bearing config). No serialized client/secret anywhere. + params = next(c.kwargs for c in gen_calls if "model_parameters" in c.kwargs) + assert params["model_parameters"]["max_tokens"] == 100 + assert params["model_parameters"]["stream"] is False + assert "client_override" not in params["model_parameters"] + assert "api_key" not in params["model_parameters"] + + async def test_no_telemetry_still_stamps_trace_without_name(self): + """Without telemetry, propagate_attributes still fires with namespace + metadata, but the trace stays unnamed — track_name lives exclusively + on telemetry now. The per-call generation still gets model/metadata + stamped (the multi-turn-regression fix means we always stamp these).""" + + mock_llm_client = AsyncMock(spec=AsyncAnthropic) + mock_response = Mock() + mock_response.content = [TextBlock(text="Unnamed response", type="text")] + mock_response.usage = Usage(input_tokens=5, output_tokens=5) + mock_response.stop_reason = "stop" + mock_llm_client.messages.create = AsyncMock(return_value=mock_response) + + mock_langfuse_client = Mock() + captured: dict[str, Any] = {} + + @contextlib.contextmanager + def fake_propagate(**kwargs: Any): + captured.update(kwargs) + yield + + with ( + patch.dict(CLIENTS, {"anthropic": mock_llm_client}), + patch.object(settings, "LANGFUSE_PUBLIC_KEY", "test-public-key"), + patch.object(settings, "LANGFUSE_EXPORTER_MODE", "inline"), + patch("langfuse.get_client", return_value=mock_langfuse_client), + patch("langfuse.propagate_attributes", fake_propagate), + ): + response = await honcho_llm_call( + model_config=ConfiguredModelSettings( + model="claude-4-sonnet", + transport="anthropic", + ), + prompt="Hello", + max_tokens=100, + enable_retry=False, ) + assert response.content == "Unnamed response" + assert captured["user_id"] == str(settings.NAMESPACE) + assert captured["trace_name"] is None + assert captured["metadata"]["namespace"] == settings.NAMESPACE + assert captured["metadata"]["provider"] == "anthropic" + # Generation gets model + metadata even without a track_name — only + # the name kwarg stays None. + gen_calls = mock_langfuse_client.update_current_generation.call_args_list + meta_kwargs = next(c.kwargs for c in gen_calls if "model" in c.kwargs) + assert meta_kwargs["name"] is None + assert meta_kwargs["model"] == "claude-4-sonnet" + # Input/output stamped explicitly (auto-capture disabled; HONCHO-4HA). + input_kwargs = next(c.kwargs for c in gen_calls if "input" in c.kwargs) + assert input_kwargs["input"] == [{"role": "user", "content": "Hello"}] + output_kwargs = next(c.kwargs for c in gen_calls if "output" in c.kwargs) + assert output_kwargs["output"].content == "Unnamed response" + + +class TestLangfuseModelParameters: + """`_langfuse_model_parameters` is the deny-list seam that keeps secrets and + live clients out of Langfuse traces while still surfacing every tuning knob + (HONCHO-4HA). It dumps the config and excludes only secret-bearing fields, so + new knobs are traced automatically without an allow-list to maintain.""" + + def test_secret_fields_never_leak_but_knobs_do(self): + from src.config import ModelConfig + from src.llm.executor import ( + _langfuse_model_parameters, # pyright: ignore[reportPrivateUsage] + ) + + # A config shaped like the production override path: real api_key / + # base_url / nested fallback / opaque provider_params. + config = ModelConfig( + model="gpt-4o", + transport="openai", + api_key="sk-super-secret", + base_url="https://user:pw@private.host/v1", + temperature=0.7, + provider_params={"x-internal-auth": "leak-me"}, + ) + + params = _langfuse_model_parameters( + max_tokens=256, + config=config, + json_mode=True, + verbosity=None, + stream=False, + tools=[{"name": "search_memory"}], + tool_choice="auto", + response_model=None, + ) + + # Secrets and their nested holders are excluded entirely... + assert "api_key" not in params + assert "base_url" not in params + assert "fallback" not in params + assert "provider_params" not in params + # ...and no value anywhere echoes a secret. + flat = str(params) + assert "sk-super-secret" not in flat + assert "leak-me" not in flat + assert "private.host" not in flat + # Tuning knobs (config-derived + per-call) are still tracked. + assert params["model"] == "gpt-4o" + assert params["temperature"] == 0.7 + assert params["max_tokens"] == 256 + assert params["json_mode"] is True + assert params["tools"] == ["search_memory"] + assert params["tool_choice"] == "auto" + class TestEdgeCases: """Tests for edge cases and boundary conditions""" diff --git a/tests/utils/test_schema_conversion.py b/tests/utils/test_schema_conversion.py new file mode 100644 index 00000000..e6c607f4 --- /dev/null +++ b/tests/utils/test_schema_conversion.py @@ -0,0 +1,897 @@ +"""Unit tests for src/utils/schema_conversion.py.""" + +import json +import re +from typing import Any + +import pytest +from pydantic import BaseModel, ValidationError + +from src.utils.schema_conversion import json_response_schema_to_pydantic + + +def _object(properties: dict[str, Any], **extra: Any) -> dict[str, Any]: + return {"type": "object", "properties": properties, **extra} + + +class TestPrimitives: + def test_flat_object_with_primitives(self): + model = json_response_schema_to_pydantic( + _object( + { + "name": {"type": "string"}, + "age": {"type": "integer"}, + "score": {"type": "number"}, + "active": {"type": "boolean"}, + }, + required=["name", "age"], + ) + ) + instance = model.model_validate( + {"name": "ada", "age": 36, "score": 9.5, "active": True} + ) + assert instance.name == "ada" # pyright: ignore + assert instance.age == 36 # pyright: ignore + + def test_required_field_missing_fails(self): + model = json_response_schema_to_pydantic( + _object({"name": {"type": "string"}}, required=["name"]) + ) + with pytest.raises(ValidationError): + model.model_validate({}) + + def test_optional_field_defaults_to_none(self): + model = json_response_schema_to_pydantic( + _object({"nickname": {"type": "string"}}) + ) + instance = model.model_validate({}) + assert instance.nickname is None # pyright: ignore + + def test_default_value(self): + model = json_response_schema_to_pydantic( + _object({"count": {"type": "integer", "default": 3}}) + ) + assert model.model_validate({}).count == 3 # pyright: ignore + + def test_null_type(self): + model = json_response_schema_to_pydantic( + _object({"nothing": {"type": "null"}}, required=["nothing"]) + ) + assert model.model_validate({"nothing": None}).nothing is None # pyright: ignore + + def test_default_wins_over_required(self): + model = json_response_schema_to_pydantic( + _object({"count": {"type": "integer", "default": 3}}, required=["count"]) + ) + assert model.model_validate({}).count == 3 # pyright: ignore + + +class TestNesting: + def test_nested_object(self): + model = json_response_schema_to_pydantic( + _object( + { + "address": _object( + { + "city": {"type": "string"}, + "geo": _object( + {"lat": {"type": "number"}}, required=["lat"] + ), + }, + required=["city", "geo"], + ) + }, + required=["address"], + ) + ) + instance = model.model_validate( + {"address": {"city": "oakland", "geo": {"lat": 37.8}}} + ) + assert instance.address.geo.lat == 37.8 # pyright: ignore + + def test_array_of_objects(self): + model = json_response_schema_to_pydantic( + _object( + { + "items": { + "type": "array", + "items": _object( + {"food": {"type": "string"}}, required=["food"] + ), + } + }, + required=["items"], + ) + ) + instance = model.model_validate({"items": [{"food": "sushi"}]}) + assert instance.items[0].food == "sushi" # pyright: ignore + + def test_array_without_items_accepts_anything(self): + model = json_response_schema_to_pydantic( + _object({"stuff": {"type": "array"}}, required=["stuff"]) + ) + instance = model.model_validate({"stuff": [1, "two", {"three": 3}]}) + assert len(instance.stuff) == 3 # pyright: ignore + + def test_nested_model_name_collision(self): + # Two sibling objects whose name hints collide must not clash. + model = json_response_schema_to_pydantic( + _object( + { + "a": _object({"x b": _object({"v": {"type": "string"}})}), + "a_x": _object({"b": _object({"v": {"type": "integer"}})}), + } + ) + ) + instance = model.model_validate( + {"a": {"x b": {"v": "s"}}, "a_x": {"b": {"v": 1}}} + ) + assert instance.a_x.b.v == 1 # pyright: ignore + + +class TestEnumsAndUnions: + def test_string_enum(self): + model = json_response_schema_to_pydantic( + _object( + {"sentiment": {"enum": ["loves", "hates"]}}, + required=["sentiment"], + ) + ) + assert model.model_validate({"sentiment": "loves"}).sentiment == "loves" # pyright: ignore + with pytest.raises(ValidationError): + model.model_validate({"sentiment": "meh"}) + + def test_int_enum_and_null_member(self): + model = json_response_schema_to_pydantic( + _object({"level": {"enum": [1, 2, None]}}, required=["level"]) + ) + assert model.model_validate({"level": None}).level is None # pyright: ignore + assert model.model_validate({"level": 2}).level == 2 # pyright: ignore + + def test_invalid_enum_value_type(self): + with pytest.raises(ValueError, match="enum values"): + json_response_schema_to_pydantic(_object({"bad": {"enum": [[1]]}})) + + def test_anyof_with_null_is_optional(self): + model = json_response_schema_to_pydantic( + _object( + {"maybe": {"anyOf": [{"type": "string"}, {"type": "null"}]}}, + required=["maybe"], + ) + ) + assert model.model_validate({"maybe": None}).maybe is None # pyright: ignore + assert model.model_validate({"maybe": "x"}).maybe == "x" # pyright: ignore + + def test_oneof_union(self): + model = json_response_schema_to_pydantic( + _object( + {"value": {"oneOf": [{"type": "integer"}, {"type": "string"}]}}, + required=["value"], + ) + ) + assert model.model_validate({"value": 5}).value == 5 # pyright: ignore + + def test_type_list_form(self): + model = json_response_schema_to_pydantic( + _object({"name": {"type": ["string", "null"]}}, required=["name"]) + ) + assert model.model_validate({"name": None}).name is None # pyright: ignore + + def test_all_null_enum_degenerates_to_none(self): + model = json_response_schema_to_pydantic( + _object({"nothing": {"enum": [None]}}, required=["nothing"]) + ) + assert model.model_validate({"nothing": None}).nothing is None # pyright: ignore + with pytest.raises(ValidationError): + model.model_validate({"nothing": "x"}) + + def test_union_of_objects(self): + model = json_response_schema_to_pydantic( + _object( + { + "pet": { + "anyOf": [ + _object({"meows": {"type": "boolean"}}, required=["meows"]), + _object({"barks": {"type": "boolean"}}, required=["barks"]), + ] + } + }, + required=["pet"], + ) + ) + instance = model.model_validate({"pet": {"barks": True}}) + assert instance.pet.barks is True # pyright: ignore + + +class TestRejections: + @pytest.mark.parametrize( + "construct,schema", + [ + ("$defs", _object({"a": {"type": "object", "$defs": {}}})), + ("definitions", _object({"a": {"type": "object", "definitions": {}}})), + ("allOf", _object({"a": {"allOf": [{"type": "string"}]}})), + ("not", _object({"a": {"not": {"type": "string"}}})), + ("if", _object({"a": {"if": {"type": "string"}}})), + ( + "patternProperties", + _object({"a": {"type": "object", "patternProperties": {}}}), + ), + ], + ) + def test_unsupported_constructs(self, construct: str, schema: dict[str, Any]): + with pytest.raises(ValueError, match=re.escape(construct)): + json_response_schema_to_pydantic(schema) + + def test_error_message_includes_path(self): + with pytest.raises( + ValueError, match=r"unsupported \$ref '#/x'.*at properties\.address" + ): + json_response_schema_to_pydantic(_object({"address": {"$ref": "#/x"}})) + + def test_schema_valued_additional_properties(self): + with pytest.raises(ValueError, match="additionalProperties"): + json_response_schema_to_pydantic( + _object( + { + "map": { + "type": "object", + "additionalProperties": {"type": "string"}, + } + } + ) + ) + + def test_boolean_schema(self): + with pytest.raises(ValueError, match="boolean schemas"): + json_response_schema_to_pydantic(_object({"anything": True})) + + def test_unknown_type(self): + with pytest.raises(ValueError, match="unsupported type 'date'"): + json_response_schema_to_pydantic(_object({"when": {"type": "date"}})) + + +class TestRefs: + def test_ref_into_defs(self): + model = json_response_schema_to_pydantic( + _object( + {"address": {"$ref": "#/$defs/Address"}}, + required=["address"], + **{ + "$defs": { + "Address": _object( + {"city": {"type": "string"}}, required=["city"] + ) + } + }, + ) + ) + instance = model.model_validate({"address": {"city": "Berlin"}}) + assert instance.address.city == "Berlin" # pyright: ignore + + def test_ref_into_definitions_alias(self): + model = json_response_schema_to_pydantic( + _object( + {"item": {"$ref": "#/definitions/Item"}}, + required=["item"], + definitions={"Item": {"type": "string"}}, + ) + ) + assert model.model_validate({"item": "x"}).item == "x" # pyright: ignore + + def test_pydantic_nested_model_schema(self): + """The real-world motivation: model_json_schema() of a nested Pydantic + model emits $defs/$ref and must convert cleanly.""" + + class Preference(BaseModel): + food: str + confidence: float + + class Preferences(BaseModel): + preferences: list[Preference] + summary: str + + model = json_response_schema_to_pydantic(Preferences.model_json_schema()) + instance = model.model_validate( + { + "preferences": [{"food": "sushi", "confidence": 0.9}], + "summary": "likes sushi", + } + ) + assert instance.preferences[0].food == "sushi" # pyright: ignore + + def test_root_ref(self): + model = json_response_schema_to_pydantic( + { + "$ref": "#/$defs/Root", + "$defs": { + "Root": _object({"ok": {"type": "boolean"}}, required=["ok"]) + }, + } + ) + assert model.model_validate({"ok": True}).ok is True # pyright: ignore + + def test_ref_sibling_keys_overlay_target(self): + model = json_response_schema_to_pydantic( + _object( + {"count": {"$ref": "#/$defs/Count", "default": 3}}, + **{"$defs": {"Count": {"type": "integer"}}}, + ) + ) + assert model.model_validate({}).count == 3 # pyright: ignore + + def test_same_def_referenced_twice(self): + model = json_response_schema_to_pydantic( + _object( + { + "home": {"$ref": "#/$defs/Address"}, + "work": {"$ref": "#/$defs/Address"}, + }, + required=["home", "work"], + **{"$defs": {"Address": _object({"city": {"type": "string"}})}}, + ) + ) + instance = model.model_validate( + {"home": {"city": "Berlin"}, "work": {"city": "Kyiv"}} + ) + assert instance.work.city == "Kyiv" # pyright: ignore + + def test_chained_refs(self): + model = json_response_schema_to_pydantic( + _object( + {"a": {"$ref": "#/$defs/A"}}, + required=["a"], + **{"$defs": {"A": {"$ref": "#/$defs/B"}, "B": {"type": "string"}}}, + ) + ) + assert model.model_validate({"a": "x"}).a == "x" # pyright: ignore + + def test_unreferenced_invalid_def_is_ignored(self): + model = json_response_schema_to_pydantic( + _object( + {"name": {"type": "string"}}, + **{"$defs": {"Broken": {"allOf": [{"type": "string"}]}}}, + ) + ) + assert model.model_validate({"name": "x"}).name == "x" # pyright: ignore + + @pytest.mark.parametrize( + "ref", + ["#", "#/x", "#/$defs/a/b", "#/properties/a", "https://x.dev/s.json#/$defs/X"], + ) + def test_unsupported_ref_forms(self, ref: str): + with pytest.raises(ValueError, match=r"unsupported \$ref"): + json_response_schema_to_pydantic( + _object( + {"a": {"$ref": ref}}, + **{"$defs": {"a": {"type": "string"}}}, + ) + ) + + def test_unknown_definition(self): + with pytest.raises(ValueError, match="unknown definition"): + json_response_schema_to_pydantic( + _object({"a": {"$ref": "#/$defs/Missing"}}, **{"$defs": {}}) + ) + + def test_direct_recursion_rejected(self): + with pytest.raises(ValueError, match=r"recursive \$ref.*cycle: Node -> Node"): + json_response_schema_to_pydantic( + _object( + {"tree": {"$ref": "#/$defs/Node"}}, + **{ + "$defs": { + "Node": _object( + { + "children": { + "type": "array", + "items": {"$ref": "#/$defs/Node"}, + } + } + ) + } + }, + ) + ) + + def test_mutual_recursion_rejected(self): + with pytest.raises(ValueError, match=r"cycle: A -> B -> A"): + json_response_schema_to_pydantic( + _object( + {"a": {"$ref": "#/$defs/A"}}, + **{ + "$defs": { + "A": _object({"b": {"$ref": "#/$defs/B"}}), + "B": _object({"a": {"$ref": "#/$defs/A"}}), + } + }, + ) + ) + + def test_recursive_pydantic_model_rejected(self): + class Node(BaseModel): + value: str + children: list["Node"] = [] + + with pytest.raises(ValueError, match=r"recursive \$ref"): + json_response_schema_to_pydantic(Node.model_json_schema()) + + def test_ref_expansion_counts_against_node_budget(self): + """A doubling ref chain (billion laughs) is stopped by max_nodes.""" + defs = { + f"L{i}": _object( + { + "a": {"$ref": f"#/$defs/L{i + 1}"}, + "b": {"$ref": f"#/$defs/L{i + 1}"}, + } + ) + for i in range(10) + } + defs["L10"] = {"type": "string"} + with pytest.raises(ValueError, match="maximum of .* nodes"): + json_response_schema_to_pydantic( + _object({"root": {"$ref": "#/$defs/L0"}}, **{"$defs": defs}) + ) + + def test_duplicate_name_across_defs_and_definitions(self): + with pytest.raises(ValueError, match="appears in both"): + json_response_schema_to_pydantic( + _object( + {"a": {"$ref": "#/$defs/X"}}, + **{ + "$defs": {"X": {"type": "string"}}, + "definitions": {"X": {"type": "integer"}}, + }, + ) + ) + + def test_root_must_be_object(self): + with pytest.raises(ValueError, match="root schema"): + json_response_schema_to_pydantic({"type": "string"}) + + def test_root_must_be_dict(self): + with pytest.raises(ValueError, match="JSON Schema object"): + json_response_schema_to_pydantic(["not", "a", "schema"]) # pyright: ignore + + def test_no_recognizable_type(self): + with pytest.raises(ValueError, match="no recognizable type"): + json_response_schema_to_pydantic(_object({"mystery": {}})) + + def test_depth_limit(self): + schema: dict[str, Any] = {"type": "string"} + for _ in range(25): + schema = _object({"inner": schema}) + with pytest.raises(ValueError, match="maximum depth"): + json_response_schema_to_pydantic(schema) + + def test_node_limit(self): + schema = _object({f"field_{i}": {"type": "string"} for i in range(600)}) + with pytest.raises(ValueError, match="maximum of 500 nodes"): + json_response_schema_to_pydantic(schema) + + def test_property_schema_not_an_object(self): + with pytest.raises(ValueError, match="schema must be an object"): + json_response_schema_to_pydantic(_object({"a": "string"})) + + @pytest.mark.parametrize("members", [[], "not-a-list"]) + def test_malformed_anyof(self, members: Any): + with pytest.raises(ValueError, match="'anyOf' must be a non-empty array"): + json_response_schema_to_pydantic(_object({"a": {"anyOf": members}})) + + def test_empty_type_list(self): + with pytest.raises(ValueError, match="'type' array must not be empty"): + json_response_schema_to_pydantic(_object({"a": {"type": []}})) + + @pytest.mark.parametrize("values", [[], "loves"]) + def test_malformed_enum(self, values: Any): + with pytest.raises(ValueError, match="'enum' must be a non-empty array"): + json_response_schema_to_pydantic(_object({"a": {"enum": values}})) + + def test_properties_not_an_object(self): + with pytest.raises(ValueError, match="'properties' must be an object"): + json_response_schema_to_pydantic({"type": "object", "properties": []}) + + @pytest.mark.parametrize("required", ["a", [1]]) + def test_malformed_required(self, required: Any): + with pytest.raises(ValueError, match="'required' must be an array of strings"): + json_response_schema_to_pydantic( + _object({"a": {"type": "string"}}, required=required) + ) + + def test_empty_property_name(self): + with pytest.raises(ValueError, match="property names"): + json_response_schema_to_pydantic(_object({"": {"type": "string"}})) + + +class TestLenientAcceptance: + def test_additional_properties_false_ignored(self): + model = json_response_schema_to_pydantic( + _object( + {"known": {"type": "string"}}, + required=["known"], + additionalProperties=False, + ) + ) + instance = model.model_validate({"known": "x", "extra": "dropped"}) + assert instance.model_dump() == {"known": "x"} + + def test_root_dollar_schema_ignored(self): + model = json_response_schema_to_pydantic( + { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": {"a": {"type": "string"}}, + } + ) + assert issubclass(model, BaseModel) + + def test_empty_properties(self): + model = json_response_schema_to_pydantic({"type": "object", "properties": {}}) + assert model.model_validate({}).model_dump() == {} + + def test_missing_properties_with_object_type(self): + model = json_response_schema_to_pydantic({"type": "object"}) + assert model.model_validate({"anything": 1}).model_dump() == {} + + def test_missing_type_with_properties_treated_as_object(self): + model = json_response_schema_to_pydantic( + {"properties": {"a": {"type": "string"}}, "required": ["a"]} + ) + assert model.model_validate({"a": "x"}).a == "x" # pyright: ignore + + def test_required_naming_unknown_property_ignored(self): + model = json_response_schema_to_pydantic( + _object({"a": {"type": "string"}}, required=["a", "ghost"]) + ) + assert model.model_validate({"a": "x"}).a == "x" # pyright: ignore + + +class TestFieldMetadata: + def test_description_propagates(self): + model = json_response_schema_to_pydantic( + _object({"food": {"type": "string", "description": "A food item"}}) + ) + generated = model.model_json_schema() + assert generated["properties"]["food"]["description"] == "A food item" + + def test_constraint_hints_pass_through_unenforced(self): + model = json_response_schema_to_pydantic( + _object( + { + "tags": { + "type": "array", + "items": {"type": "string"}, + "maxItems": 3, + } + }, + required=["tags"], + ) + ) + generated = model.model_json_schema() + assert generated["properties"]["tags"]["maxItems"] == 3 + # Not enforced: more than maxItems still validates. + instance = model.model_validate({"tags": ["a", "b", "c", "d"]}) + assert len(instance.tags) == 4 # pyright: ignore + + def test_non_identifier_key_alias_round_trip(self): + model = json_response_schema_to_pydantic( + _object( + {"my-key": {"type": "string"}, "_private": {"type": "integer"}}, + required=["my-key"], + ) + ) + instance = model.model_validate({"my-key": "v", "_private": 7}) + dumped = instance.model_dump_json(by_alias=True) + assert '"my-key":"v"' in dumped + assert '"_private":7' in dumped + + def test_digit_leading_key_gets_field_prefix(self): + model = json_response_schema_to_pydantic( + _object({"123": {"type": "integer"}}, required=["123"]) + ) + instance = model.model_validate({"123": 7}) + assert instance.model_dump(by_alias=True) == {"123": 7} + + def test_sanitized_key_collision_round_trip(self): + # "my-key" sanitizes to "my_key", which then collides with the real + # "my_key" property; both must survive with their original JSON keys. + model = json_response_schema_to_pydantic( + _object( + {"my-key": {"type": "string"}, "my_key": {"type": "integer"}}, + required=["my-key", "my_key"], + ) + ) + instance = model.model_validate({"my-key": "v", "my_key": 7}) + assert instance.model_dump(by_alias=True) == {"my-key": "v", "my_key": 7} + + def test_digit_leading_model_name(self): + model = json_response_schema_to_pydantic( + _object({"a": {"type": "string"}}), model_name="123" + ) + assert model.__name__ == "Model123" + + +class TestZodCompatibility: + def test_zod4_tojsonschema_output_converts(self): + # Captured shape of zod 4's z.toJSONSchema() for + # z.object({ preferences: z.array(z.object({ food: z.string(), + # sentiment: z.enum(["loves","hates"]) })), summary: z.string(), + # note: z.string().optional() }) + schema = { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": { + "preferences": { + "type": "array", + "items": { + "type": "object", + "properties": { + "food": {"type": "string"}, + "sentiment": { + "type": "string", + "enum": ["loves", "hates"], + }, + }, + "required": ["food", "sentiment"], + "additionalProperties": False, + }, + }, + "summary": {"type": "string"}, + "note": {"type": "string"}, + }, + "required": ["preferences", "summary"], + "additionalProperties": False, + } + model = json_response_schema_to_pydantic(schema) + instance = model.model_validate( + { + "preferences": [{"food": "sushi", "sentiment": "loves"}], + "summary": "likes sushi", + } + ) + assert instance.preferences[0].sentiment == "loves" # pyright: ignore + assert instance.note is None # pyright: ignore + + +class TestCustomGuardLimits: + def test_custom_max_depth(self): + schema: dict[str, Any] = {"type": "string"} + for _ in range(5): + schema = _object({"inner": schema}) + with pytest.raises(ValueError, match="maximum depth of 3"): + json_response_schema_to_pydantic(schema, max_depth=3) + + def test_custom_max_nodes(self): + schema = _object({f"f{i}": {"type": "string"} for i in range(20)}) + with pytest.raises(ValueError, match="maximum of 10 nodes"): + json_response_schema_to_pydantic(schema, max_nodes=10) + + def test_depth_exactly_at_limit_allowed(self): + # Leaf sits at depth == max_depth; only depth > max_depth must fail. + schema: dict[str, Any] = {"type": "string"} + for _ in range(3): + schema = _object({"inner": schema}) + model = json_response_schema_to_pydantic(schema, max_depth=3) + assert issubclass(model, BaseModel) + + +# The wiki spec's own request example (dialectic-enhancements §3.A.1). +SPEC_EXAMPLE_SCHEMA = _object( + { + "preferences": { + "type": "array", + "items": _object( + { + "food": {"type": "string"}, + "sentiment": { + "type": "string", + "enum": ["loves", "likes", "neutral", "dislikes", "hates"], + }, + "confidence": {"type": "number", "minimum": 0, "maximum": 1}, + }, + required=["food", "sentiment"], + ), + "maxItems": 3, + }, + "summary": {"type": "string"}, + }, + required=["preferences", "summary"], +) + + +class TestEndToEnd: + """Table tests running the full pipeline the server runs: convert the + caller's schema, validate a payload against the generated model, and + serialize it back with model_dump_json(by_alias=True).""" + + @pytest.mark.parametrize( + "schema,payload,expected", + [ + pytest.param( + SPEC_EXAMPLE_SCHEMA, + { + "preferences": [ + { + "food": "dark roast coffee", + "sentiment": "loves", + "confidence": 0.95, + }, + {"food": "sushi", "sentiment": "likes"}, + ], + "summary": "Coffee enthusiast.", + }, + { + "preferences": [ + { + "food": "dark roast coffee", + "sentiment": "loves", + "confidence": 0.95, + }, + {"food": "sushi", "sentiment": "likes", "confidence": None}, + ], + "summary": "Coffee enthusiast.", + }, + id="spec-example", + ), + pytest.param( + _object( + { + "user": _object( + { + "name": {"type": "string"}, + "location": _object( + { + "lat": {"type": "number"}, + "lon": {"type": "number"}, + }, + required=["lat", "lon"], + ), + }, + required=["name", "location"], + ) + }, + required=["user"], + ), + {"user": {"name": "ada", "location": {"lat": 37.8, "lon": -122.3}}}, + {"user": {"name": "ada", "location": {"lat": 37.8, "lon": -122.3}}}, + id="nested-three-levels", + ), + pytest.param( + _object( + {"my-key": {"type": "string"}, "first name": {"type": "string"}}, + required=["my-key"], + ), + {"my-key": "v", "first name": "Ada"}, + {"my-key": "v", "first name": "Ada"}, + id="alias-keys-round-trip", + ), + pytest.param( + _object( + { + "count": {"type": "integer", "default": 3}, + "tag": {"type": "string", "default": "none"}, + } + ), + {}, + {"count": 3, "tag": "none"}, + id="defaults-fill-omitted-fields", + ), + pytest.param( + _object( + { + "a": {"anyOf": [{"type": "string"}, {"type": "null"}]}, + "b": {"type": ["integer", "null"]}, + }, + required=["a", "b"], + ), + {"a": None, "b": 2}, + {"a": None, "b": 2}, + id="nullable-via-anyof-and-type-list", + ), + pytest.param( + _object( + {"value": {"oneOf": [{"type": "integer"}, {"type": "string"}]}}, + required=["value"], + ), + {"value": "five"}, + {"value": "five"}, + id="oneof-union-string-member", + ), + pytest.param( + _object({"level": {"enum": [1, 2, None]}}, required=["level"]), + {"level": None}, + {"level": None}, + id="enum-with-null-member", + ), + pytest.param( + _object({"stuff": {"type": "array"}}, required=["stuff"]), + {"stuff": [1, "two", {"three": 3}, None]}, + {"stuff": [1, "two", {"three": 3}, None]}, + id="array-without-items-accepts-anything", + ), + pytest.param( + _object( + { + "tags": { + "type": "array", + "items": {"type": "string"}, + "maxItems": 2, + } + }, + required=["tags"], + ), + {"tags": ["a", "b", "c", "d"]}, + {"tags": ["a", "b", "c", "d"]}, + id="constraint-hints-not-enforced", + ), + pytest.param( + _object({"known": {"type": "string"}}, required=["known"]), + {"known": "x", "hallucinated": "dropped"}, + {"known": "x"}, + id="extra-keys-dropped", + ), + pytest.param( + {"type": "object", "properties": {}}, + {}, + {}, + id="empty-object", + ), + ], + ) + def test_construct_validate_serialize( + self, + schema: dict[str, Any], + payload: dict[str, Any], + expected: dict[str, Any], + ): + model = json_response_schema_to_pydantic(schema) + instance = model.model_validate(payload) + # by_alias=True mirrors DialecticAgent.answer's serialization. + assert json.loads(instance.model_dump_json(by_alias=True)) == expected + + @pytest.mark.parametrize( + "schema,payload", + [ + pytest.param( + SPEC_EXAMPLE_SCHEMA, + {"preferences": [{"food": "sushi"}], "summary": "s"}, + id="missing-required-in-array-item", + ), + pytest.param( + SPEC_EXAMPLE_SCHEMA, + { + "preferences": [{"food": "sushi", "sentiment": "adores"}], + "summary": "s", + }, + id="invalid-enum-value", + ), + pytest.param( + SPEC_EXAMPLE_SCHEMA, + {"preferences": [{"food": "sushi", "sentiment": "likes"}]}, + id="missing-required-top-level", + ), + pytest.param( + _object( + {"user": _object({"name": {"type": "string"}}, required=["name"])}, + required=["user"], + ), + {"user": {}}, + id="missing-required-nested", + ), + pytest.param( + _object({"a": {"type": "string"}}, required=["a"]), + {"a": None}, + id="null-for-non-nullable", + ), + pytest.param( + _object({"n": {"type": "integer"}}, required=["n"]), + {"n": {"nested": "dict"}}, + id="wrong-type-for-integer", + ), + ], + ) + def test_rejects_nonconforming_payloads( + self, schema: dict[str, Any], payload: dict[str, Any] + ): + model = json_response_schema_to_pydantic(schema) + with pytest.raises(ValidationError): + model.model_validate(payload) diff --git a/tests/vector_store/test_lancedb.py b/tests/vector_store/test_lancedb.py index 0c502e4c..9b52708a 100644 --- a/tests/vector_store/test_lancedb.py +++ b/tests/vector_store/test_lancedb.py @@ -37,6 +37,27 @@ def store() -> LanceDBVectorStore: return LanceDBVectorStore() +def test_build_where_clause_membership(store: LanceDBVectorStore) -> None: + """Both the dict `in` form and the bare-list sugar produce an IN clause.""" + assert ( + store._build_where_clause({"session_name": {"in": ["s1", "s2"]}}) # pyright: ignore[reportPrivateUsage] + == "session_name IN ('s1', 's2')" + ) + assert ( + store._build_where_clause({"session_name": ["s1", "s2"]}) # pyright: ignore[reportPrivateUsage] + == "session_name IN ('s1', 's2')" + ) + + +def test_build_where_clause_empty_membership_fails_closed( + store: LanceDBVectorStore, +) -> None: + """An empty membership list must emit an always-false predicate, never an + omitted condition that would widen scope (fail-open).""" + assert store._build_where_clause({"session_name": {"in": []}}) == "1 = 0" # pyright: ignore[reportPrivateUsage] + assert store._build_where_clause({"session_name": []}) == "1 = 0" # pyright: ignore[reportPrivateUsage] + + @pytest.mark.asyncio async def test_query_returns_empty_when_table_missing( store: LanceDBVectorStore, diff --git a/tests/vector_store/test_turbopuffer.py b/tests/vector_store/test_turbopuffer.py index ca8c4ff3..cdae44d5 100644 --- a/tests/vector_store/test_turbopuffer.py +++ b/tests/vector_store/test_turbopuffer.py @@ -53,6 +53,30 @@ async def test_upsert_many_raises_vector_store_error_on_5xx( namespace_mock.write.assert_awaited_once() +def test_build_filters_membership(store: TurbopufferVectorStore) -> None: + """Both the dict `in` form and the bare-list sugar produce an In filter.""" + assert store._build_filters({"session_name": {"in": ["s1", "s2"]}}) == ( # pyright: ignore[reportPrivateUsage] + "session_name", + "In", + ["s1", "s2"], + ) + assert store._build_filters({"session_name": ["s1", "s2"]}) == ( # pyright: ignore[reportPrivateUsage] + "session_name", + "In", + ["s1", "s2"], + ) + + +def test_build_filters_empty_membership_fails_closed( + store: TurbopufferVectorStore, +) -> None: + """An empty membership list must produce an always-false filter, never an + omitted/empty In that could widen scope (fail-open).""" + never = ("And", [("session_name", "Eq", ""), ("session_name", "NotEq", "")]) + assert store._build_filters({"session_name": {"in": []}}) == never # pyright: ignore[reportPrivateUsage] + assert store._build_filters({"session_name": []}) == never # pyright: ignore[reportPrivateUsage] + + @pytest.mark.asyncio async def test_upsert_many_short_circuits_on_empty( store: TurbopufferVectorStore, diff --git a/uv.lock b/uv.lock index ce198bbe..0c2a9889 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-05-16T17:58:57.678125Z" +exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P5D" [manifest] @@ -276,11 +276,11 @@ wheels = [ [[package]] name = "cashews" -version = "7.4.4" +version = "7.5.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/8c/5d/26eb556824a7ac9e24f751645961d2078b7b15be105f7fc39eda5308896f/cashews-7.4.4.tar.gz", hash = "sha256:dca761c60192bfe354abd6e9eb98d6f62c817e675df3fbe7d1bdfaa4303d1320", size = 92948, upload-time = "2025-12-06T22:31:56.187Z" } +sdist = { url = "https://files.pythonhosted.org/packages/44/73/31598b352165cd0f0b777df1eb67e33f334e29fcd7eb4f4bb48a41b9affe/cashews-7.5.0.tar.gz", hash = "sha256:3f88b8c5ced0ea4826915a1ff67055b647252dd65ef25f4813316a6341f00b37", size = 97699, upload-time = "2026-03-02T22:28:52.462Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/bf/65/29d94c27dfa3cdb213ae62a328c6efe6cd37b888d334e90ecaa22eadafe9/cashews-7.4.4-py3-none-any.whl", hash = "sha256:d5b8fc3cb590ed388823388b972947fd5659e2a94109af107cb508a3240f5ef0", size = 79893, upload-time = "2025-12-06T22:31:53.918Z" }, + { url = "https://files.pythonhosted.org/packages/0d/14/06cca741567a2ec458fb1db9d053e72477d9da2be387c1d705cb1060b2c6/cashews-7.5.0-py3-none-any.whl", hash = "sha256:e79cb4e5cc164d8f2d2856b166d45dcc2dd8d53b95874d2c6d07dfdb1c9ac3c4", size = 82413, upload-time = "2026-03-02T22:28:50.98Z" }, ] [package.optional-dependencies] @@ -1224,7 +1224,7 @@ wheels = [ [[package]] name = "honcho" -version = "3.0.7" +version = "3.0.11" source = { virtual = "." } dependencies = [ { name = "alembic" }, @@ -1236,7 +1236,7 @@ dependencies = [ { name = "greenlet" }, { name = "httpx" }, { name = "json-repair" }, - { name = "lancedb" }, + { name = "lancedb", marker = "platform_machine != 'x86_64' or sys_platform != 'darwin'" }, { name = "langfuse" }, { name = "nanoid" }, { name = "openai" }, @@ -1283,7 +1283,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "alembic", specifier = ">=1.14.0" }, - { name = "cashews", extras = ["redis"], specifier = "==7.4.4" }, + { name = "cashews", extras = ["redis"], specifier = "==7.5.0" }, { name = "cloudevents", specifier = ">=1.12.0,<2.0" }, { name = "fastapi", extras = ["standard"], specifier = ">=0.131.0" }, { name = "fastapi-pagination", specifier = ">=0.14.2" }, @@ -1291,7 +1291,7 @@ requires-dist = [ { name = "greenlet", specifier = ">=3.0.3" }, { name = "httpx", specifier = ">=0.27.0" }, { name = "json-repair", specifier = ">=0.49.0" }, - { name = "lancedb", specifier = ">=0.25.3" }, + { name = "lancedb", marker = "platform_machine != 'x86_64' or sys_platform != 'darwin'", specifier = ">=0.25.3" }, { name = "langfuse", specifier = ">=3.3.2" }, { name = "nanoid", specifier = ">=2.0.0" }, { name = "openai", specifier = ">=1.99.7" }, @@ -1337,7 +1337,7 @@ dev = [ [[package]] name = "honcho-ai" -version = "2.1.2" +version = "2.2.0" source = { editable = "sdks/python" } dependencies = [ { name = "httpx" }, @@ -1369,9 +1369,10 @@ dev = [{ name = "ruff", specifier = ">=0.11.13" }] [[package]] name = "honcho-cli" -version = "0.1.0" +version = "0.1.2" source = { editable = "honcho-cli" } dependencies = [ + { name = "click" }, { name = "honcho-ai" }, { name = "httpx" }, { name = "rich" }, @@ -1386,6 +1387,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "click", specifier = ">=8.0.0" }, { name = "honcho-ai", editable = "sdks/python" }, { name = "httpx", specifier = ">=0.27.0" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0.0" },