mirror of
https://github.com/Significant-Gravitas/AutoGPT.git
synced 2026-11-10 12:51:43 -10:00
Merge remote-tracking branch 'origin/dev' into abhi/ds-theming-a-colours
Brings in #16513 (Text, Icon and Button variants). Each conflicting call site keeps #16513's markup and props, with its colours renamed to roles by scripts/ds/alias-to-role.ts and palette-to-role.ts. The Text tones moved to src/components/styles/tones.ts, so the success-light exemption follows them there. Co-authored-by: Claude Opus 5.5 (Claude Code) <[email protected]>
This commit is contained in:
commit
1ae437b0e5
653 files changed
+10788
-12351
No files matched your search
@@ -14,6 +14,7 @@ SEED_VERIFY_CACHED_STEP = (
|
||||
"Set up tests - Verify cached seed credentials with fresh encryption key"
|
||||
)
|
||||
FRONTEND_WAIT_STEP = "Wait for frontend image build"
|
||||
SERVICES_START_STEP = "Set up Platform - Start (all other services)"
|
||||
SEED_SAVE_STEP = "Push E2E test data to GHCR"
|
||||
PLAYWRIGHT_STEP = "Run Playwright E2E suite"
|
||||
PLAYWRIGHT_REPORT_STEP = "Validate Playwright JSON report"
|
||||
@@ -133,9 +134,11 @@ class E2ESeedCacheContractTests(unittest.TestCase):
|
||||
for name in (SEED_VERIFY_CACHED_STEP, SEED_VERIFY_STEP):
|
||||
self.assertLess(steps.index(name), steps.index(PLAYWRIGHT_STEP))
|
||||
|
||||
def test_seed_waits_for_the_frontend_build_and_is_time_boxed(self):
|
||||
# Seeding beside the Next.js build froze the runner for up to 48 minutes
|
||||
# (OPEN-3734); a restored dump's check still overlaps the build.
|
||||
def test_services_and_seed_wait_for_the_frontend_build_and_seed_is_time_boxed(
|
||||
self,
|
||||
):
|
||||
# Beside the Next.js build, seeding froze the runner for up to 48 minutes
|
||||
# (OPEN-3734) and the started backend services lost it (SECRT-3254).
|
||||
workflow = (
|
||||
Path(__file__).resolve().parents[1] / "workflows/platform-fullstack-ci.yml"
|
||||
).read_text(encoding="utf-8")
|
||||
@@ -147,10 +150,9 @@ class E2ESeedCacheContractTests(unittest.TestCase):
|
||||
self.assertIn(
|
||||
"wait: [platform-build-frontend]", _step(workflow, FRONTEND_WAIT_STEP)
|
||||
)
|
||||
self.assertLess(
|
||||
steps.index(SEED_VERIFY_CACHED_STEP), steps.index(FRONTEND_WAIT_STEP)
|
||||
)
|
||||
self.assertLess(steps.index(FRONTEND_WAIT_STEP), steps.index(SEED_CREATE_STEP))
|
||||
for name in (SERVICES_START_STEP, SEED_VERIFY_CACHED_STEP, SEED_CREATE_STEP):
|
||||
with self.subTest(step=name):
|
||||
self.assertLess(steps.index(FRONTEND_WAIT_STEP), steps.index(name))
|
||||
|
||||
timeout = re.search(
|
||||
r"^\s+timeout-minutes: (\d+)$",
|
||||
|
||||
@@ -413,8 +413,30 @@ jobs:
|
||||
cp -a /seed/. /app/.next/cache/
|
||||
EOF
|
||||
|
||||
# The runner has 15.6 GiB and a cold `next build` peaks near 8 GiB; a runner
|
||||
# that runs out leaves no other memory readings (SECRT-3254).
|
||||
- name: Set up Platform - Log runner memory
|
||||
id: memory-log
|
||||
background: true
|
||||
run: |
|
||||
lowest=
|
||||
# e2e_test sets no job-level timeout-minutes, so GitHub's 360-minute
|
||||
# default is the real ceiling; sample until then or the stop file,
|
||||
# whichever comes first, so a long build+seed can't outlast the log.
|
||||
for i in $(seq 2160); do
|
||||
[ -f /tmp/stop-memory-log ] && break
|
||||
read -r used available < <(free -m | awk '/^Mem:/ {print $3, $7}')
|
||||
swap=$(free -m | awk '/^Swap:/ {print $3}')
|
||||
echo "$(date -u +%T) used=${used}M available=${available}M swap=${swap}M"
|
||||
if [ -z "$lowest" ] || [ "$available" -lt "$lowest" ]; then lowest=$available; fi
|
||||
# Once a minute, which containers hold it
|
||||
if (( i % 6 == 0 )); then docker stats --no-stream --format ' {{.Name}} {{.MemUsage}}' || true; fi
|
||||
sleep 10
|
||||
done
|
||||
echo "Lowest available memory: ${lowest}M"
|
||||
|
||||
# The backend images finish minutes before the Next.js build, so they are
|
||||
# baked on their own and the backend stack starts while the frontend builds.
|
||||
# baked on their own and the database is migrated while the frontend builds.
|
||||
- name: Set up Platform - Build backend Docker images
|
||||
id: platform-build-backend
|
||||
working-directory: autogpt_platform
|
||||
@@ -449,7 +471,11 @@ jobs:
|
||||
for image in $images; do
|
||||
bash "$GITHUB_WORKSPACE/.github/scripts/docker-pull-with-retry.sh" "$image"
|
||||
done
|
||||
docker compose -f ../docker-compose.resolved.yml up -d db --no-build
|
||||
# --no-deps: db depends_on rabbitmq (startup-ordering only, see
|
||||
# docker-compose.yml), so without it this would start rabbitmq here
|
||||
# too, overlapping the frontend build. "Start (all other services)"
|
||||
# brings rabbitmq up with everything else once the build is done.
|
||||
docker compose -f ../docker-compose.resolved.yml up -d db --no-build --no-deps
|
||||
echo "Waiting for database to be ready..."
|
||||
timeout 60 sh -c 'until docker compose -f ../docker-compose.resolved.yml exec -T db pg_isready -U postgres 2>/dev/null; do sleep 2; done'
|
||||
|
||||
@@ -497,7 +523,10 @@ jobs:
|
||||
- name: Set up Platform - Run migrations
|
||||
run: |
|
||||
echo "Running migrations..."
|
||||
docker compose -f ../docker-compose.resolved.yml run --rm migrate
|
||||
# --no-deps: `run` resolves migrate's full dependency closure, so
|
||||
# without it this starts rabbitmq too (db's own dependency, even
|
||||
# though db is already healthy from the step above) during the build.
|
||||
docker compose -f ../docker-compose.resolved.yml run --rm --no-deps migrate
|
||||
echo "✅ Migrations completed"
|
||||
env:
|
||||
NEXT_PUBLIC_PW_TEST: true
|
||||
@@ -524,9 +553,14 @@ jobs:
|
||||
echo "restored=true" >> "$GITHUB_OUTPUT"
|
||||
echo "✅ E2E test data restored from cache"
|
||||
|
||||
# Beside the Next.js build, the backend services and the seed run the
|
||||
# runner out of memory and freeze or lose it (OPEN-3734, SECRT-3254).
|
||||
- name: Wait for frontend image build
|
||||
wait: [platform-build-frontend]
|
||||
|
||||
- name: Set up Platform - Start (all other services)
|
||||
run: |
|
||||
# Everything but the frontend, whose image may still be building
|
||||
# Everything but the frontend, which starts once the seed is in
|
||||
docker compose -f ../docker-compose.resolved.yml up -d --no-build \
|
||||
$(docker compose -f ../docker-compose.resolved.yml config --services | grep -vx frontend)
|
||||
echo "Waiting for rest_server to be ready..."
|
||||
@@ -546,7 +580,6 @@ jobs:
|
||||
env:
|
||||
NEXT_PUBLIC_PW_TEST: true
|
||||
|
||||
# A restored dump's check runs before the wait below, so it still overlaps the frontend build.
|
||||
- name: Set up tests - Verify cached seed credentials with fresh encryption key
|
||||
id: e2e-seed-verify-cached
|
||||
if: steps.e2e-seed-restore.outputs.restored == 'true'
|
||||
@@ -555,10 +588,6 @@ jobs:
|
||||
docker compose -f ../docker-compose.resolved.yml exec -T rest_server \
|
||||
sh -c "cd /app/autogpt_platform && python /tmp/e2e_test_data.py --refresh-credentials-only"
|
||||
|
||||
# Seeding while the Next.js build runs can freeze the whole runner (OPEN-3734).
|
||||
- name: Wait for frontend image build
|
||||
wait: [platform-build-frontend]
|
||||
|
||||
- name: Set up tests - Create E2E test data
|
||||
id: e2e-seed-create
|
||||
if: matrix.seed && steps.e2e-seed-restore.outputs.restored != 'true'
|
||||
@@ -615,6 +644,13 @@ jobs:
|
||||
exit 1
|
||||
}
|
||||
|
||||
- name: Stop runner memory log
|
||||
if: always()
|
||||
run: touch /tmp/stop-memory-log
|
||||
|
||||
- name: Wait for runner memory log
|
||||
wait: [memory-log]
|
||||
|
||||
- name: Copy source maps from Docker for E2E coverage
|
||||
run: |
|
||||
FRONTEND_CONTAINER=$(docker compose -f ../docker-compose.resolved.yml ps -q frontend)
|
||||
|
||||
@@ -25,7 +25,7 @@ jobs:
|
||||
pull-requests: write
|
||||
steps:
|
||||
- name: Update PRs with conflict labels
|
||||
uses: eps1lon/actions-label-merge-conflict@fd1f295ee7443d13745804bc49fe158e240f6c6e # v2.1.0
|
||||
uses: eps1lon/actions-label-merge-conflict@0273be72a0bbd58fcd71d0d6c02c209b50d1e5e1 # v3.1.0
|
||||
with:
|
||||
dirtyLabel: "conflicts"
|
||||
#removeOnDirtyLabel: "PR: ready to ship"
|
||||
|
||||
@@ -227,10 +227,10 @@ cd /src/autogpt_platform/frontend
|
||||
then sign in again (`POST http://localhost:3000/api/auth/sign-in/email` with the email and password) and fetch a
|
||||
token from the new session; the old session keeps the old role. An admin can give a user credits with
|
||||
`POST /api/credits/admin/add_credits`.
|
||||
- **Feature flags** come from PostHog (`FEATURE_FLAG_BACKEND=posthog`), which isn't configured here, so every
|
||||
flag reads as its default in the code. `FORCE_FLAG_<KEY>=true` in the backend's environment turns one on (the
|
||||
flag key upper-cased, with `-` as `_`), and `FORCE_ALL_FLAGS=true` turns them all on. Code behind a flag is in
|
||||
scope: the hosted service turns flags on for some or all users.
|
||||
- **Feature flags** come from PostHog, which isn't configured here, so every flag reads as its default in the
|
||||
code. `FORCE_FLAG_<KEY>=true` in the backend's environment turns one on (the flag key upper-cased, with `-` as
|
||||
`_`), and `FORCE_ALL_FLAGS=true` turns them all on. Code behind a flag is in scope: the hosted service turns flags
|
||||
on for some or all users.
|
||||
- `next start` warns that it doesn't work with standalone output; it serves everything anyway.
|
||||
- Anything that needs an API key logs errors and fails offline (LLM calls, embeddings, third-party integrations).
|
||||
|
||||
|
||||
@@ -7,17 +7,14 @@
|
||||
-- Tracks real API spend at the call level: provider, model,
|
||||
-- token counts (including Anthropic cache tokens), cost in
|
||||
-- microdollars, and the block/execution that incurred the cost.
|
||||
-- Joins the User table to provide email for per-user breakdowns.
|
||||
--
|
||||
-- SOURCE TABLES
|
||||
-- platform.PlatformCostLog — Per-call cost records
|
||||
-- platform.User — User email
|
||||
--
|
||||
-- OUTPUT COLUMNS
|
||||
-- id TEXT Log entry UUID
|
||||
-- createdAt TIMESTAMPTZ When the cost was recorded
|
||||
-- userId TEXT User who incurred the cost (nullable)
|
||||
-- email TEXT User email (nullable)
|
||||
-- graphExecId TEXT Graph execution UUID (nullable)
|
||||
-- nodeExecId TEXT Node execution UUID (nullable)
|
||||
-- blockName TEXT Block that made the API call (nullable)
|
||||
@@ -51,10 +48,10 @@
|
||||
-- GROUP BY 1, 2 ORDER BY total_usd DESC;
|
||||
--
|
||||
-- -- Top 20 users by spend
|
||||
-- SELECT "userId", email, SUM("costUsd") AS total_usd, COUNT(*) AS calls
|
||||
-- SELECT "userId", SUM("costUsd") AS total_usd, COUNT(*) AS calls
|
||||
-- FROM analytics.platform_cost_log
|
||||
-- WHERE "userId" IS NOT NULL
|
||||
-- GROUP BY 1, 2 ORDER BY total_usd DESC LIMIT 20;
|
||||
-- GROUP BY 1 ORDER BY total_usd DESC LIMIT 20;
|
||||
--
|
||||
-- -- Daily spend trend
|
||||
-- SELECT DATE_TRUNC('day', "createdAt") AS day,
|
||||
@@ -76,7 +73,6 @@ SELECT
|
||||
p."id" AS id,
|
||||
p."createdAt" AS createdAt,
|
||||
p."userId" AS userId,
|
||||
u."email" AS email,
|
||||
p."graphExecId" AS graphExecId,
|
||||
p."nodeExecId" AS nodeExecId,
|
||||
p."blockName" AS blockName,
|
||||
@@ -96,5 +92,4 @@ SELECT
|
||||
END AS totalTokens,
|
||||
p."duration" AS duration
|
||||
FROM platform."PlatformCostLog" p
|
||||
LEFT JOIN platform."User" u ON u."id" = p."userId"
|
||||
WHERE p."createdAt" > CURRENT_DATE - INTERVAL '90 days'
|
||||
@@ -27,7 +27,7 @@
|
||||
-- platform.PlatformCostLog, platform.CreditTransaction, platform.User
|
||||
--
|
||||
-- OUTPUT COLUMNS
|
||||
-- user_id, email, subscription_tier, month (DATE, first of month),
|
||||
-- user_id, subscription_tier, month (DATE, first of month),
|
||||
-- tasks_human, tasks_automated, agent_runs, autopilot_turns, expert_turns,
|
||||
-- active_days, platform_cost_usd, agent_cost_usd, copilot_cost_usd,
|
||||
-- background_cost_usd, credits_spent_usd, credits_purchased_usd,
|
||||
@@ -164,7 +164,6 @@ keys AS (
|
||||
assembled AS (
|
||||
SELECT
|
||||
k.user_id,
|
||||
u."email" AS email,
|
||||
u."subscriptionTier"::text AS subscription_tier,
|
||||
k.month,
|
||||
COALESCE(r.agent_runs_human, 0) - COALESCE(r.agent_runs_copilot, 0)
|
||||
|
||||
@@ -6,9 +6,11 @@
|
||||
-- One row per user: where they came from, captured once around
|
||||
-- signup by the browser that created the account. Carries the ids
|
||||
-- the analytics tools knew the user by before they had a user id,
|
||||
-- so channel (DataFast), product analytics (PostHog) and flag
|
||||
-- bucketing (LaunchDarkly, via the shared anonymous id) can all be
|
||||
-- joined to activation, retention and revenue on user_id.
|
||||
-- so channel (DataFast) and product analytics (PostHog, via the
|
||||
-- shared anonymous id) can be joined to activation, retention and
|
||||
-- revenue on user_id. No feature-flag bucket is recorded here:
|
||||
-- anonymous flag reads send no exposure event, and experiment arms
|
||||
-- are stored per user by the experiments API.
|
||||
--
|
||||
-- Join to analytics.user_lifecycle for "which channel brings users
|
||||
-- who activate". Rows exist only for users who signed up (or first
|
||||
@@ -20,7 +22,7 @@
|
||||
-- OUTPUT COLUMNS
|
||||
-- user_id TEXT User UUID
|
||||
-- captured_at TIMESTAMPTZ When the row was first written
|
||||
-- anonymous_id TEXT First-party anonymous id shared with PostHog + LaunchDarkly
|
||||
-- anonymous_id TEXT First-party anonymous id shared with PostHog
|
||||
-- posthog_distinct_id TEXT PostHog device id at signup, when different
|
||||
-- datafast_visitor_id TEXT DataFast visitor (join key into DataFast exports)
|
||||
-- datafast_session_id TEXT DataFast session that contained the signup
|
||||
|
||||
@@ -34,7 +34,7 @@
|
||||
-- (org-billed runs, attributed to the user who initiated them)
|
||||
--
|
||||
-- OUTPUT COLUMNS
|
||||
-- Identity: user_id, email, signup_at, subscription_tier, timezone,
|
||||
-- Identity: user_id, signup_at, subscription_tier, timezone,
|
||||
-- usage_reason, onboarding_completed, onboarding_integrations_selected,
|
||||
-- integrations_connected_total, integration_providers_connected,
|
||||
-- first_integration_connected_at
|
||||
@@ -85,7 +85,6 @@
|
||||
WITH users AS (
|
||||
SELECT
|
||||
u."id" AS user_id,
|
||||
u."email" AS email,
|
||||
u."createdAt" AS signup_at,
|
||||
u."subscriptionTier"::text AS subscription_tier,
|
||||
u."timezone" AS timezone
|
||||
@@ -268,7 +267,7 @@ integrations AS (
|
||||
),
|
||||
assembled AS (
|
||||
SELECT
|
||||
u.user_id, u.email, u.signup_at, u.subscription_tier, u.timezone,
|
||||
u.user_id, u.signup_at, u.subscription_tier, u.timezone,
|
||||
o.usage_reason,
|
||||
COALESCE(o.onboarding_completed, FALSE) AS onboarding_completed,
|
||||
COALESCE(o.onboarding_integrations_selected, 0) AS onboarding_integrations_selected,
|
||||
|
||||
@@ -139,7 +139,7 @@ NVIDIA_API_KEY=
|
||||
# Note: on the managed cloud platform (BEHAVE_AS=cloud) CHAT_*_MODEL is
|
||||
# the BOTTOM layer of AutoPilot model resolution — the LLM catalog's
|
||||
# copilot routing cells (backend/data/llm_registry/catalog.py) and the
|
||||
# per-user LaunchDarkly "copilot-model-routing" flag take precedence
|
||||
# per-user PostHog "copilot-model-routing" flag take precedence
|
||||
# when set. On self-hosted installs the catalog cells are skipped and
|
||||
# these env vars stay authoritative. See
|
||||
# docs/platform/contributing/managing-llm-models.md.
|
||||
@@ -176,7 +176,9 @@ GRAPHITI_EMBEDDER_MODEL=text-embedding-3-small
|
||||
GRAPHITI_SEMAPHORE_LIMIT=5
|
||||
|
||||
# Langfuse Prompt Management
|
||||
# Used for managing the CoPilot system prompt externally
|
||||
# Used only to fetch the CoPilot system prompt and the onboarding greeting
|
||||
# prompt. The backend sends no traces to Langfuse; AutoPilot traces go to
|
||||
# PostHog (if OpenRouter Broadcast still has a Langfuse destination, turn it off).
|
||||
# Get credentials from https://cloud.langfuse.com or your self-hosted instance
|
||||
LANGFUSE_PUBLIC_KEY=
|
||||
LANGFUSE_SECRET_KEY=
|
||||
@@ -318,17 +320,15 @@ SENTRY_DSN=
|
||||
# SENTRY_ENVIRONMENT=
|
||||
|
||||
# Feature Flags
|
||||
LAUNCH_DARKLY_SDK_KEY=
|
||||
# Which vendor answers a flag read: posthog (default), launchdarkly, or dual
|
||||
# (evaluate both, always serve LaunchDarkly here, log every disagreement).
|
||||
# Without POSTHOG_API_KEY every flag reads as its code default.
|
||||
FEATURE_FLAG_BACKEND=posthog
|
||||
# Flags are read from PostHog (POSTHOG_API_KEY below; POSTHOG_SECRET_KEY lets
|
||||
# them evaluate in-process). Without POSTHOG_API_KEY every flag reads as its
|
||||
# code default.
|
||||
# Local-dev master switch: turn EVERY boolean feature flag on at once, without
|
||||
# listing them. Off by default; local dev only, where no flag vendor is
|
||||
# listing them. Off by default; local dev only, where PostHog is usually not
|
||||
# configured and every flag is otherwise its default. A per-flag FORCE_FLAG_<NAME>
|
||||
# below still wins, so one flag can be excluded with =false. Non-boolean flags
|
||||
# (stripe-product-id-topup, copilot-model-routing, the copilot tier/cost JSON
|
||||
# flags) are left to the flag vendor. Ignored, with an error log, unless
|
||||
# flags) are left to PostHog. Ignored, with an error log, unless
|
||||
# APP_ENV=local -- `dev` and `prod` are real deployments, not a developer's
|
||||
# machine. That includes the single-container image, whose entrypoint exports
|
||||
# APP_ENV=dev; use the per-flag FORCE_FLAG_<NAME> below there instead. Set the
|
||||
|
||||
@@ -0,0 +1,198 @@
|
||||
"""Per-principal rate limiting for the AutoGPT external API.
|
||||
|
||||
The external API authenticates with API keys / OAuth tokens, so one valid
|
||||
credential can hammer store endpoints or spam graph executions. This module
|
||||
provides an atomic fixed-window counter in Redis (Lua ``INCR`` + first-hit
|
||||
``EXPIRE``, mirroring ``backend/api/features/billing/credits_rate_limit.py``) keyed
|
||||
per (user, scope), with per-scope tiers, returning HTTP 429 with
|
||||
``Retry-After`` / ``X-RateLimit-*`` headers once the window counter is
|
||||
exhausted.
|
||||
|
||||
The bucket is deliberately keyed on the *user*, not on the credential: an
|
||||
abuse control must not be resettable by minting another API key or by
|
||||
refreshing an OAuth access token, both of which produce a new credential id
|
||||
for the same principal.
|
||||
|
||||
Availability: the check is strictly best-effort and **fails open** on any
|
||||
Redis trouble — being unable to prove a principal is under its cap must never
|
||||
take the whole external API down. The whole Redis interaction is bounded by a
|
||||
short deadline so an outage falls through fast instead of parking ASGI worker
|
||||
slots (same reasoning as ``credits_rate_limit``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import UTC, datetime
|
||||
from typing import Annotated, Any, cast
|
||||
|
||||
from fastapi import Depends, HTTPException, Response, Security, status
|
||||
from redis.exceptions import RedisClusterException, RedisError
|
||||
|
||||
from backend.api.external.middleware import require_auth
|
||||
from backend.data.auth.base import APIAuthorizationInfo
|
||||
from backend.data.redis_client import get_redis_async
|
||||
from backend.monitoring.instrumentation import record_rate_limit_hit
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Per-scope tiers: (max requests, window seconds).
|
||||
# Cheap reads (identity, block listing) get a generous cap; anything that
|
||||
# executes work (block execution, graph runs) gets a tight one.
|
||||
READ_MAX_REQUESTS = 300
|
||||
READ_WINDOW_SECONDS = 60
|
||||
EXECUTION_MAX_REQUESTS = 30
|
||||
EXECUTION_WINDOW_SECONDS = 60
|
||||
|
||||
# Cancellation deadline on the Redis interaction. Without it "fail open" is not
|
||||
# "fail fast": a cold client runs redis-py's own connect retry ladder, so
|
||||
# during an outage a request would park an ASGI worker slot for minutes
|
||||
# instead of falling through.
|
||||
RATE_LIMIT_REDIS_TIMEOUT_SECONDS = 0.25
|
||||
|
||||
# Atomic fixed-window counter: INCR the key, and set the TTL only on the INCR
|
||||
# that opened the window (count == 1). Doing both in one server-side script
|
||||
# means a freshly-created key always gets its expiry — the key can never
|
||||
# linger without a TTL, and the window stays fixed (the TTL is not refreshed
|
||||
# on later hits).
|
||||
_INCR_OPEN_WINDOW = """
|
||||
local count = redis.call('INCR', KEYS[1])
|
||||
if count == 1 then
|
||||
redis.call('EXPIRE', KEYS[1], ARGV[1])
|
||||
end
|
||||
return count
|
||||
"""
|
||||
|
||||
|
||||
def _rate_limit_key(auth: APIAuthorizationInfo, scope: str) -> str:
|
||||
"""One bucket per (user, scope).
|
||||
|
||||
The user is the unit that cannot be multiplied, so the budget must follow
|
||||
the user rather than the credential presented with the request. Minting a
|
||||
second API key, or letting an OAuth access token refresh, both yield a new
|
||||
credential ``id`` for the same principal; keying on that id would hand out
|
||||
a fresh budget each time and make the limit trivially bypassable.
|
||||
"""
|
||||
return f"external_api:rl:{auth.user_id}:{scope}"
|
||||
|
||||
|
||||
def _window_start(now: datetime, window_seconds: int) -> int:
|
||||
"""Start of the wall-clock window ``now`` falls in.
|
||||
|
||||
The window is aligned to absolute time so the key, ``X-RateLimit-Reset``
|
||||
and ``Retry-After`` all describe the same instant. A client that waits for
|
||||
``Retry-After`` therefore lands in the next window and is admitted, instead
|
||||
of being refused again because the TTL runs from its first hit.
|
||||
"""
|
||||
return int(now.timestamp()) // window_seconds * window_seconds
|
||||
|
||||
|
||||
async def _incr_window(key: str, window_seconds: int) -> int:
|
||||
"""Run the atomic counter, returning the request's position in the window.
|
||||
|
||||
Lua ARGV values are strings; ``EXPIRE`` coerces ``"60"`` back to an int.
|
||||
The cast mirrors the other ``eval()`` call sites (e.g.
|
||||
``credits_rate_limit``): the cluster client types ``eval()``'s return as
|
||||
``str``.
|
||||
"""
|
||||
redis = await get_redis_async()
|
||||
return await cast(
|
||||
Any,
|
||||
redis.eval(_INCR_OPEN_WINDOW, 1, key, str(window_seconds)),
|
||||
)
|
||||
|
||||
|
||||
def require_rate_limit(
|
||||
*,
|
||||
max_requests: int,
|
||||
window_seconds: int,
|
||||
scope: str,
|
||||
):
|
||||
"""FastAPI dependency enforcing a per-principal rate limit on a route.
|
||||
|
||||
The dependency injects the ``Response`` so it can set ``X-RateLimit-*``
|
||||
headers on the way out, and raises HTTP 429 (with ``Retry-After``) once
|
||||
the fixed-window counter is exhausted. On any Redis trouble it fails
|
||||
open: logs and lets the request through.
|
||||
"""
|
||||
|
||||
async def dependency(
|
||||
response: Response,
|
||||
auth: Annotated[APIAuthorizationInfo, Security(require_auth)],
|
||||
) -> APIAuthorizationInfo:
|
||||
now = datetime.now(UTC)
|
||||
window_start = _window_start(now, window_seconds)
|
||||
# A new key at the reported reset admits the next window immediately;
|
||||
# the previous key's TTL only controls its eventual cleanup.
|
||||
key = f"{_rate_limit_key(auth, scope)}:{window_start}"
|
||||
try:
|
||||
count = await asyncio.wait_for(
|
||||
_incr_window(key, window_seconds),
|
||||
timeout=RATE_LIMIT_REDIS_TIMEOUT_SECONDS,
|
||||
)
|
||||
except (
|
||||
RedisError,
|
||||
RedisClusterException,
|
||||
ConnectionError,
|
||||
OSError,
|
||||
TimeoutError,
|
||||
ValueError,
|
||||
) as e:
|
||||
logger.warning(
|
||||
"External API rate-limit check failed open for user %s scope %s: %s",
|
||||
auth.user_id,
|
||||
scope,
|
||||
e,
|
||||
)
|
||||
return auth
|
||||
|
||||
remaining = max(0, max_requests - int(count))
|
||||
reset_at = window_start + window_seconds
|
||||
retry_after = max(0, reset_at - int(now.timestamp()))
|
||||
|
||||
response.headers["X-RateLimit-Limit"] = str(max_requests)
|
||||
response.headers["X-RateLimit-Remaining"] = str(remaining)
|
||||
response.headers["X-RateLimit-Reset"] = str(reset_at)
|
||||
|
||||
if int(count) > max_requests:
|
||||
logger.info(
|
||||
"External API rate limit hit for user %s scope %s (count=%s)",
|
||||
auth.user_id,
|
||||
scope,
|
||||
count,
|
||||
)
|
||||
record_rate_limit_hit(f"external_api:{scope}", auth.user_id)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
||||
detail=f"Rate limit exceeded for scope '{scope}'. "
|
||||
f"Retry after {retry_after}s.",
|
||||
headers={
|
||||
"Retry-After": str(retry_after),
|
||||
"X-RateLimit-Limit": str(max_requests),
|
||||
"X-RateLimit-Remaining": "0",
|
||||
"X-RateLimit-Reset": str(reset_at),
|
||||
},
|
||||
)
|
||||
return auth
|
||||
|
||||
return dependency
|
||||
|
||||
|
||||
# Ready-made dependencies for the two tiers. Every ``/v1`` route carries one
|
||||
# of these, so they are defined once here rather than repeating the eight-line
|
||||
# ``Depends(require_rate_limit(...))`` block at each call site.
|
||||
READ_LIMIT = Depends(
|
||||
require_rate_limit(
|
||||
max_requests=READ_MAX_REQUESTS,
|
||||
window_seconds=READ_WINDOW_SECONDS,
|
||||
scope="read",
|
||||
)
|
||||
)
|
||||
EXECUTION_LIMIT = Depends(
|
||||
require_rate_limit(
|
||||
max_requests=EXECUTION_MAX_REQUESTS,
|
||||
window_seconds=EXECUTION_WINDOW_SECONDS,
|
||||
scope="execution",
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,280 @@
|
||||
"""Unit + integration tests for the external API rate limiter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import fastapi
|
||||
import pytest
|
||||
from fastapi import Response
|
||||
from fastapi.testclient import TestClient
|
||||
from prisma.enums import APIKeyPermission, APIKeyStatus
|
||||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||||
from redis.exceptions import RedisClusterException
|
||||
|
||||
from backend.api.external import rate_limit
|
||||
from backend.api.external.middleware import require_auth
|
||||
from backend.api.external.v1.routes import v1_router
|
||||
from backend.data import redis_client
|
||||
from backend.data.auth.api_key import APIKeyInfo
|
||||
from backend.data.auth.base import APIAuthorizationInfo
|
||||
from backend.data.auth.oauth import OAuthAccessTokenInfo
|
||||
from backend.data.redis_client import get_redis_async
|
||||
from backend.util.testing import is_tcp_port_reachable
|
||||
|
||||
|
||||
class FakeRedis:
|
||||
"""In-memory counter fake; the live Redis test below validates the Lua."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.counters: dict[str, int] = {}
|
||||
self.expires_at: dict[str, float] = {}
|
||||
self._now = 0.0
|
||||
|
||||
async def eval(self, script: str, numkeys: int, key: str, ttl: str) -> int:
|
||||
assert script == rate_limit._INCR_OPEN_WINDOW
|
||||
assert numkeys == 1
|
||||
count = self.counters.get(key, 0) + 1
|
||||
self.counters[key] = count
|
||||
if count == 1:
|
||||
self.expires_at[key] = self._now + int(ttl)
|
||||
return count
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def fixed_window(mocker):
|
||||
clock = mocker.patch.object(rate_limit, "datetime")
|
||||
clock.now.return_value = datetime(2026, 10, 30, 12, 0, 30, tzinfo=UTC)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_redis(mocker):
|
||||
redis = FakeRedis()
|
||||
mocker.patch(
|
||||
"backend.api.external.rate_limit.get_redis_async",
|
||||
new=AsyncMock(return_value=redis),
|
||||
)
|
||||
return redis
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_redis(mocker):
|
||||
"""MagicMock whose ``eval`` return value each test sets directly."""
|
||||
redis = MagicMock()
|
||||
redis.eval = AsyncMock(return_value=1)
|
||||
mocker.patch(
|
||||
"backend.api.external.rate_limit.get_redis_async",
|
||||
new=AsyncMock(return_value=redis),
|
||||
)
|
||||
return redis
|
||||
|
||||
|
||||
def _auth(user_id: str = "u1", key_id: str | None = "k1") -> APIAuthorizationInfo:
|
||||
if key_id is None:
|
||||
return OAuthAccessTokenInfo(
|
||||
id=uuid.uuid4().hex,
|
||||
application_id="test-app",
|
||||
user_id=user_id,
|
||||
scopes=list(APIKeyPermission),
|
||||
created_at=datetime.now(UTC),
|
||||
expires_at=datetime.now(UTC) + timedelta(hours=1),
|
||||
)
|
||||
return APIKeyInfo(
|
||||
id=key_id,
|
||||
name="test-key",
|
||||
head="abcd",
|
||||
tail="wxyz",
|
||||
status=APIKeyStatus.ACTIVE,
|
||||
user_id=user_id,
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
|
||||
|
||||
async def _call_dependency(**kwargs):
|
||||
"""Call ``require_rate_limit``'s dependency directly with a Response."""
|
||||
dep = rate_limit.require_rate_limit(
|
||||
max_requests=rate_limit.READ_MAX_REQUESTS,
|
||||
window_seconds=rate_limit.READ_WINDOW_SECONDS,
|
||||
scope="read",
|
||||
)
|
||||
response = Response()
|
||||
auth = kwargs.get("auth", _auth())
|
||||
ret = await dep(response=response, auth=auth)
|
||||
return ret, response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_first_hit_sets_ttl_and_headers(fake_redis):
|
||||
"""The window-opening request must set the TTL and expose X-RateLimit-*
|
||||
headers on the response."""
|
||||
_, response = await _call_dependency()
|
||||
|
||||
assert response.headers["X-RateLimit-Limit"] == str(rate_limit.READ_MAX_REQUESTS)
|
||||
assert response.headers["X-RateLimit-Remaining"] == str(
|
||||
rate_limit.READ_MAX_REQUESTS - 1
|
||||
)
|
||||
assert int(response.headers["X-RateLimit-Reset"]) > 0
|
||||
(key,) = fake_redis.expires_at
|
||||
assert key.startswith("external_api:rl:u1:read:")
|
||||
# The key ends with the wall-clock window start, so the counter, the TTL
|
||||
# and the reported reset all describe the same window.
|
||||
window_start = int(key.rsplit(":", 1)[1])
|
||||
assert window_start % rate_limit.READ_WINDOW_SECONDS == 0
|
||||
assert int(response.headers["X-RateLimit-Reset"]) == (
|
||||
window_start + rate_limit.READ_WINDOW_SECONDS
|
||||
)
|
||||
assert (
|
||||
fake_redis.expires_at[key] - fake_redis._now == rate_limit.READ_WINDOW_SECONDS
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_keys_of_one_user_share_a_bucket(fake_redis):
|
||||
"""A second API key must not buy a fresh budget.
|
||||
|
||||
The user is the unit that cannot be multiplied: if the bucket were keyed on
|
||||
the credential id, minting another key would reset the limit.
|
||||
"""
|
||||
dep = rate_limit.require_rate_limit(
|
||||
max_requests=rate_limit.READ_MAX_REQUESTS,
|
||||
window_seconds=rate_limit.READ_WINDOW_SECONDS,
|
||||
scope="read",
|
||||
)
|
||||
response = Response()
|
||||
await dep(response=response, auth=_auth(user_id="u1", key_id="k1"))
|
||||
await dep(response=response, auth=_auth(user_id="u1", key_id="k2"))
|
||||
|
||||
assert len(fake_redis.counters) == 1
|
||||
(key,) = fake_redis.counters
|
||||
assert key.startswith("external_api:rl:u1:read:")
|
||||
assert fake_redis.counters[key] == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oauth_refresh_shares_the_bucket(fake_redis):
|
||||
"""Refreshing an OAuth token must not reset the budget either.
|
||||
|
||||
OAuth principals carry an access-token id, which changes on every refresh;
|
||||
the bucket has to follow the user instead.
|
||||
"""
|
||||
dep = rate_limit.require_rate_limit(
|
||||
max_requests=rate_limit.READ_MAX_REQUESTS,
|
||||
window_seconds=rate_limit.READ_WINDOW_SECONDS,
|
||||
scope="read",
|
||||
)
|
||||
response = Response()
|
||||
await dep(response=response, auth=_auth(user_id="u9", key_id=None))
|
||||
await dep(response=response, auth=_auth(user_id="u9", key_id=None))
|
||||
|
||||
assert len(fake_redis.counters) == 1
|
||||
(key,) = fake_redis.counters
|
||||
assert key.startswith("external_api:rl:u9:read:")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_over_limit_returns_429(mock_redis):
|
||||
"""Once the counter exceeds the cap the dependency raises HTTP 429 with
|
||||
Retry-After and the rate-limit headers."""
|
||||
mock_redis.eval = AsyncMock(return_value=rate_limit.READ_MAX_REQUESTS + 1)
|
||||
dep = rate_limit.require_rate_limit(
|
||||
max_requests=rate_limit.READ_MAX_REQUESTS,
|
||||
window_seconds=rate_limit.READ_WINDOW_SECONDS,
|
||||
scope="read",
|
||||
)
|
||||
response = Response()
|
||||
with pytest.raises(fastapi.HTTPException) as exc_info:
|
||||
await dep(response=response, auth=_auth())
|
||||
|
||||
exc = exc_info.value
|
||||
assert exc.status_code == 429
|
||||
assert exc.headers["Retry-After"] is not None
|
||||
assert exc.headers["X-RateLimit-Remaining"] == "0"
|
||||
assert exc.headers["X-RateLimit-Limit"] == str(rate_limit.READ_MAX_REQUESTS)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_failure_fails_open(mock_redis):
|
||||
"""A Redis blip must not take the request down: log and let it through."""
|
||||
mock_redis.eval = AsyncMock(side_effect=RedisConnectionError("connection refused"))
|
||||
ret, response = await _call_dependency()
|
||||
|
||||
assert ret is not None
|
||||
assert "X-RateLimit-Limit" not in response.headers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cluster_exception_fails_open(mock_redis):
|
||||
"""RedisClusterException does not inherit from RedisError and would
|
||||
otherwise surface as 500 — it must fail open too."""
|
||||
mock_redis.eval = AsyncMock(side_effect=RedisClusterException("slot missing"))
|
||||
ret, _ = await _call_dependency()
|
||||
|
||||
assert ret is not None
|
||||
|
||||
|
||||
def test_integration_429_via_route(mocker, mock_redis):
|
||||
"""A real request through the router gets a 429 once the limit is hit."""
|
||||
app = fastapi.FastAPI()
|
||||
app.include_router(v1_router, prefix="/v1")
|
||||
|
||||
async def fake_require_auth() -> APIAuthorizationInfo:
|
||||
return _auth()
|
||||
|
||||
app.dependency_overrides[require_auth] = fake_require_auth
|
||||
mock_redis.eval = AsyncMock(return_value=rate_limit.READ_MAX_REQUESTS + 1)
|
||||
|
||||
with TestClient(app) as client:
|
||||
res = client.get("/v1/me")
|
||||
assert res.status_code == 429
|
||||
assert res.headers["Retry-After"] is not None
|
||||
assert res.headers["X-RateLimit-Limit"] == str(rate_limit.READ_MAX_REQUESTS)
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skipif(
|
||||
not is_tcp_port_reachable(redis_client.HOST, redis_client.PORT),
|
||||
reason="needs Redis",
|
||||
)
|
||||
async def test_the_script_opens_a_window_once_on_real_redis():
|
||||
"""FakeRedis re-implements the Lua, so only real Redis can catch a
|
||||
script that never expires the key or re-opens the window on every hit."""
|
||||
key = f"external_api:rl:test:{uuid.uuid4().hex}"
|
||||
redis = await get_redis_async()
|
||||
try:
|
||||
assert await rate_limit._incr_window(key, 60) == 1
|
||||
assert 0 < await redis.ttl(key) <= 60
|
||||
await redis.expire(key, 30)
|
||||
assert await rate_limit._incr_window(key, 60) == 2
|
||||
assert 0 < await redis.ttl(key) <= 30
|
||||
finally:
|
||||
await redis.delete(key)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_hung_redis_fails_open_within_the_deadline(mocker, mock_redis):
|
||||
"""A Redis that never answers must not park the request: the 0.25s
|
||||
deadline fires and the request is let through."""
|
||||
|
||||
async def hang(*_):
|
||||
await asyncio.Event().wait()
|
||||
|
||||
mock_redis.eval = hang
|
||||
wait_for = mocker.spy(asyncio, "wait_for")
|
||||
ret, response = await asyncio.wait_for(_call_dependency(), timeout=2)
|
||||
|
||||
assert ret is not None
|
||||
assert "X-RateLimit-Limit" not in response.headers
|
||||
assert any(call.kwargs.get("timeout") == 0.25 for call in wait_for.call_args_list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_cancellation_propagates(mock_redis):
|
||||
mock_redis.eval.side_effect = asyncio.CancelledError
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await _call_dependency()
|
||||
@@ -0,0 +1,127 @@
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from hashlib import sha256
|
||||
from unittest.mock import AsyncMock
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
from prisma.enums import APIKeyPermission, APIKeyStatus
|
||||
from prisma.models import OAuthAccessToken, OAuthApplication
|
||||
|
||||
from backend.api.external import rate_limit
|
||||
from backend.api.external.fastapi_app import external_api
|
||||
from backend.api.external.v1.routes import UserInfoResponse
|
||||
from backend.data.auth import api_key
|
||||
from backend.data.auth.api_key import APIKeyInfoWithHash
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def v1_generated_keys():
|
||||
return [api_key.keysmith.generate_key() for _ in range(3)]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def v1_key_rows(v1_generated_keys):
|
||||
owner = uuid4().hex
|
||||
return [
|
||||
APIKeyInfoWithHash(
|
||||
id=f"key-{i}",
|
||||
name="rate-limit-test",
|
||||
head=key.head,
|
||||
tail=key.tail,
|
||||
hash=key.hash,
|
||||
salt=key.salt,
|
||||
user_id=owner if i < 2 else uuid4().hex,
|
||||
scopes=list(APIKeyPermission),
|
||||
status=APIKeyStatus.ACTIVE,
|
||||
created_at=datetime.now(UTC),
|
||||
last_used_at=datetime.now(UTC),
|
||||
)
|
||||
for i, key in enumerate(v1_generated_keys)
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def v1_tokens(v1_key_rows):
|
||||
now = datetime.now(UTC)
|
||||
application = OAuthApplication(
|
||||
id="app-1",
|
||||
name="rate-limit-test",
|
||||
clientId="client-1",
|
||||
clientSecret="unused",
|
||||
clientSecretSalt="unused",
|
||||
redirectUris=[],
|
||||
grantTypes=["authorization_code", "refresh_token"],
|
||||
scopes=list(APIKeyPermission),
|
||||
ownerId=v1_key_rows[0].user_id,
|
||||
isActive=True,
|
||||
createdAt=now,
|
||||
updatedAt=now,
|
||||
)
|
||||
return {
|
||||
f"agpt_xt_test-{i}": OAuthAccessToken(
|
||||
id=f"token-{i}",
|
||||
token=sha256(f"agpt_xt_test-{i}".encode()).hexdigest(),
|
||||
applicationId=application.id,
|
||||
Application=application,
|
||||
userId=row.user_id,
|
||||
scopes=list(APIKeyPermission),
|
||||
createdAt=now,
|
||||
expiresAt=now + timedelta(hours=1),
|
||||
)
|
||||
for i, row in enumerate(v1_key_rows)
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def v1_auth_storage(mocker, v1_key_rows, v1_tokens):
|
||||
async def keys_by_head(head):
|
||||
return [row for row in v1_key_rows if row.head == head]
|
||||
|
||||
async def token_by_hash(*, where, include):
|
||||
return next(
|
||||
(row for row in v1_tokens.values() if row.token == where["token"]), None
|
||||
)
|
||||
|
||||
key_lookup = mocker.patch.object(
|
||||
api_key, "get_active_api_keys_by_head", side_effect=keys_by_head
|
||||
)
|
||||
token_client = AsyncMock()
|
||||
token_client.find_unique.side_effect = token_by_hash
|
||||
mocker.patch.object(OAuthAccessToken, "prisma", return_value=token_client)
|
||||
return key_lookup, token_client.find_unique
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def v1_clock(mocker):
|
||||
clock = mocker.patch.object(rate_limit, "datetime")
|
||||
clock.now.return_value = datetime(2026, 10, 30, 12, 0, 57, tzinfo=UTC)
|
||||
return clock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def v1_client(mocker, v1_auth_storage, v1_clock):
|
||||
async def user_by_id(user_id):
|
||||
return UserInfoResponse(
|
||||
id=user_id, name="Test user", email="[email protected]", timezone="UTC"
|
||||
)
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.external.v1.routes.user_db.get_user_by_id",
|
||||
side_effect=user_by_id,
|
||||
)
|
||||
async with AsyncClient(
|
||||
transport=ASGITransport(app=external_api), base_url="http://test"
|
||||
) as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def v1_counter(mocker):
|
||||
counts: dict[str, int] = {}
|
||||
|
||||
async def increment(key, window_seconds):
|
||||
counts[key] = counts.get(key, 0) + 1
|
||||
return counts[key]
|
||||
|
||||
return mocker.patch.object(rate_limit, "_incr_window", side_effect=increment)
|
||||
@@ -17,6 +17,7 @@ from prisma.enums import APIKeyPermission
|
||||
from pydantic import BaseModel, Field, SecretStr
|
||||
|
||||
from backend.api.external.middleware import require_permission
|
||||
from backend.api.external.rate_limit import EXECUTION_LIMIT, READ_LIMIT
|
||||
from backend.api.features.integrations.models import get_all_provider_names
|
||||
from backend.api.features.integrations.router import (
|
||||
CredentialsMetaResponse,
|
||||
@@ -266,7 +267,11 @@ def _get_oauth_handler_for_external(
|
||||
# ==================== Endpoints ==================== #
|
||||
|
||||
|
||||
@integrations_router.get("/providers", response_model=list[ProviderInfo])
|
||||
@integrations_router.get(
|
||||
"/providers",
|
||||
response_model=list[ProviderInfo],
|
||||
dependencies=[READ_LIMIT],
|
||||
)
|
||||
async def list_providers(
|
||||
auth: APIAuthorizationInfo = Security(
|
||||
require_permission(APIKeyPermission.READ_INTEGRATIONS)
|
||||
@@ -331,6 +336,7 @@ async def list_providers(
|
||||
"/{provider}/oauth/initiate",
|
||||
response_model=OAuthInitiateResponse,
|
||||
summary="Initiate OAuth flow",
|
||||
dependencies=[EXECUTION_LIMIT],
|
||||
)
|
||||
async def initiate_oauth(
|
||||
provider: Annotated[str, Path(title="The OAuth provider")],
|
||||
@@ -410,6 +416,7 @@ async def initiate_oauth(
|
||||
"/{provider}/oauth/complete",
|
||||
response_model=OAuthCompleteResponse,
|
||||
summary="Complete OAuth flow",
|
||||
dependencies=[EXECUTION_LIMIT],
|
||||
)
|
||||
async def complete_oauth(
|
||||
provider: Annotated[str, Path(title="The OAuth provider")],
|
||||
@@ -489,7 +496,11 @@ async def complete_oauth(
|
||||
)
|
||||
|
||||
|
||||
@integrations_router.get("/credentials", response_model=list[CredentialsMetaResponse])
|
||||
@integrations_router.get(
|
||||
"/credentials",
|
||||
response_model=list[CredentialsMetaResponse],
|
||||
dependencies=[READ_LIMIT],
|
||||
)
|
||||
async def list_credentials(
|
||||
auth: APIAuthorizationInfo = Security(
|
||||
require_permission(APIKeyPermission.READ_INTEGRATIONS)
|
||||
@@ -507,7 +518,9 @@ async def list_credentials(
|
||||
|
||||
|
||||
@integrations_router.get(
|
||||
"/{provider}/credentials", response_model=list[CredentialsMetaResponse]
|
||||
"/{provider}/credentials",
|
||||
response_model=list[CredentialsMetaResponse],
|
||||
dependencies=[READ_LIMIT],
|
||||
)
|
||||
async def list_credentials_by_provider(
|
||||
provider: Annotated[str, Path(title="The provider to list credentials for")],
|
||||
@@ -531,6 +544,7 @@ async def list_credentials_by_provider(
|
||||
response_model=CreateCredentialResponse,
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
summary="Create credentials",
|
||||
dependencies=[EXECUTION_LIMIT],
|
||||
)
|
||||
async def create_credential(
|
||||
provider: Annotated[str, Path(title="The provider to create credentials for")],
|
||||
@@ -630,6 +644,7 @@ class DeleteCredentialResponse(BaseModel):
|
||||
@integrations_router.delete(
|
||||
"/{provider}/credentials/{cred_id}",
|
||||
response_model=DeleteCredentialResponse,
|
||||
dependencies=[EXECUTION_LIMIT],
|
||||
)
|
||||
async def delete_credential(
|
||||
provider: Annotated[str, Path(title="The provider")],
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from threading import get_ident
|
||||
|
||||
import pytest
|
||||
from prisma.enums import APIKeyPermission
|
||||
|
||||
from backend.api.external import middleware
|
||||
from backend.data.auth import api_key
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bearer", [False, True], ids=["api-key", "bearer-key"])
|
||||
async def test_real_key_verification_runs_off_the_route_event_loop(
|
||||
mocker, v1_client, v1_counter, v1_generated_keys, v1_key_rows, bearer
|
||||
):
|
||||
verification_threads: list[int] = []
|
||||
verify = api_key.keysmith.verify_key
|
||||
|
||||
def traced_verify(*args, **kwargs):
|
||||
verification_threads.append(get_ident())
|
||||
return verify(*args, **kwargs)
|
||||
|
||||
mocker.patch.object(api_key.keysmith, "verify_key", side_effect=traced_verify)
|
||||
key = v1_generated_keys[0].key
|
||||
headers = {"Authorization": f"Bearer {key}"} if bearer else {"X-API-Key": key}
|
||||
response = await v1_client.get("/v1/me", headers=headers)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["id"] == v1_key_rows[0].user_id
|
||||
assert verification_threads
|
||||
assert get_ident() not in verification_threads
|
||||
v1_counter.assert_awaited_once()
|
||||
assert "no-store" in response.headers["cache-control"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"headers",
|
||||
[{}, {"X-API-Key": "invalid"}, {"Authorization": "Bearer agpt_xt_invalid"}],
|
||||
ids=["missing", "invalid-key", "invalid-token"],
|
||||
)
|
||||
async def test_authentication_precedes_the_counter(v1_client, v1_counter, headers):
|
||||
response = await v1_client.get("/v1/me", headers=headers)
|
||||
|
||||
assert response.status_code == 401
|
||||
v1_counter.assert_not_awaited()
|
||||
assert "x-ratelimit-limit" not in response.headers
|
||||
|
||||
|
||||
async def test_wrong_key_with_a_valid_head_is_rejected(
|
||||
v1_client, v1_counter, v1_generated_keys
|
||||
):
|
||||
key = v1_generated_keys[0].key
|
||||
wrong = key[:-1] + ("a" if key[-1] != "a" else "b")
|
||||
|
||||
response = await v1_client.get("/v1/me", headers={"X-API-Key": wrong})
|
||||
|
||||
assert response.status_code == 401
|
||||
v1_counter.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_invalid_key_does_not_fall_back_to_a_valid_bearer(
|
||||
v1_client, v1_counter, v1_auth_storage
|
||||
):
|
||||
response = await v1_client.get(
|
||||
"/v1/me",
|
||||
headers={"X-API-Key": "invalid", "Authorization": "Bearer agpt_xt_test-0"},
|
||||
)
|
||||
|
||||
assert response.status_code == 401
|
||||
v1_counter.assert_not_awaited()
|
||||
v1_auth_storage[1].assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("credential", ["key", "token"])
|
||||
async def test_banned_owner_is_rejected_before_the_counter(
|
||||
mocker, v1_client, v1_counter, v1_generated_keys, credential
|
||||
):
|
||||
mocker.patch.object(middleware, "is_user_banned", return_value=True)
|
||||
headers = (
|
||||
{"X-API-Key": v1_generated_keys[0].key}
|
||||
if credential == "key"
|
||||
else {"Authorization": "Bearer agpt_xt_test-0"}
|
||||
)
|
||||
|
||||
response = await v1_client.get("/v1/me", headers=headers)
|
||||
|
||||
assert response.status_code == 403
|
||||
assert response.json()["detail"]["code"] == "account_banned"
|
||||
v1_counter.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("state", ["expired", "revoked", "disabled-client"])
|
||||
async def test_invalid_oauth_state_precedes_the_counter(
|
||||
v1_client, v1_counter, v1_tokens, state
|
||||
):
|
||||
token = v1_tokens["agpt_xt_test-0"]
|
||||
if state == "expired":
|
||||
token.expiresAt = datetime.now(UTC) - timedelta(hours=1)
|
||||
elif state == "revoked":
|
||||
token.revokedAt = datetime.now(UTC)
|
||||
else:
|
||||
token.Application.isActive = False
|
||||
|
||||
response = await v1_client.get(
|
||||
"/v1/me", headers={"Authorization": "Bearer agpt_xt_test-0"}
|
||||
)
|
||||
|
||||
assert response.status_code == 401
|
||||
v1_counter.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", ["/me", "/blocks", "/integrations/credentials"])
|
||||
async def test_missing_permissions_never_reach_the_handler(
|
||||
mocker, v1_client, v1_counter, v1_tokens, path
|
||||
):
|
||||
v1_tokens["agpt_xt_test-0"].scopes = []
|
||||
user = mocker.patch("backend.api.external.v1.routes.user_db.get_user_by_id")
|
||||
blocks = mocker.patch("backend.blocks.get_blocks")
|
||||
credentials = mocker.patch(
|
||||
"backend.api.external.v1.integrations.creds_manager.store.get_all_creds"
|
||||
)
|
||||
|
||||
response = await v1_client.get(
|
||||
f"/v1{path}", headers={"Authorization": "Bearer agpt_xt_test-0"}
|
||||
)
|
||||
|
||||
assert response.status_code == 403
|
||||
assert "Missing required permission" in response.json()["detail"]
|
||||
user.assert_not_called()
|
||||
blocks.assert_not_called()
|
||||
credentials.assert_not_called()
|
||||
|
||||
|
||||
async def test_verification_and_permissions_are_not_cached_across_requests(
|
||||
v1_client, v1_counter, v1_tokens
|
||||
):
|
||||
headers = {"Authorization": "Bearer agpt_xt_test-0"}
|
||||
assert (await v1_client.get("/v1/me", headers=headers)).status_code == 200
|
||||
v1_tokens["agpt_xt_test-0"].scopes.remove(APIKeyPermission.IDENTITY)
|
||||
assert (await v1_client.get("/v1/me", headers=headers)).status_code == 403
|
||||
v1_tokens["agpt_xt_test-0"].revokedAt = datetime.now(UTC)
|
||||
assert (await v1_client.get("/v1/me", headers=headers)).status_code == 401
|
||||
@@ -0,0 +1,44 @@
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from backend.api.external import rate_limit
|
||||
from backend.data import redis_client
|
||||
from backend.util.testing import is_tcp_port_reachable
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not is_tcp_port_reachable(redis_client.HOST, redis_client.PORT),
|
||||
reason="needs Redis",
|
||||
)
|
||||
|
||||
|
||||
async def test_concurrent_routes_share_an_atomic_budget(
|
||||
v1_client, v1_key_rows, v1_clock
|
||||
):
|
||||
redis = await redis_client.get_redis_async()
|
||||
window = rate_limit._window_start(v1_clock.now.return_value, 60)
|
||||
key = f"external_api:rl:{v1_key_rows[0].user_id}:read:{window}"
|
||||
gate = asyncio.Barrier(20)
|
||||
|
||||
async def request(index):
|
||||
await gate.wait()
|
||||
return await v1_client.get(
|
||||
"/v1/me", headers={"Authorization": f"Bearer agpt_xt_test-{index % 2}"}
|
||||
)
|
||||
|
||||
try:
|
||||
await redis.set(key, 290, ex=60)
|
||||
responses = await asyncio.wait_for(
|
||||
asyncio.gather(*(request(i) for i in range(20))), timeout=10
|
||||
)
|
||||
assert sum(r.status_code == 200 for r in responses) == 10
|
||||
assert sum(r.status_code == 429 for r in responses) == 10
|
||||
assert sorted(
|
||||
int(r.headers["x-ratelimit-remaining"])
|
||||
for r in responses
|
||||
if r.status_code == 200
|
||||
) == list(range(10))
|
||||
assert int(await redis.get(key)) == 310
|
||||
assert 0 < await redis.ttl(key) <= 60
|
||||
finally:
|
||||
await redis.delete(key)
|
||||
@@ -0,0 +1,160 @@
|
||||
from datetime import timedelta
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi.routing import APIRoute
|
||||
from redis.exceptions import RedisError
|
||||
|
||||
from backend.api.external import rate_limit
|
||||
from backend.api.external.v1 import routes
|
||||
from backend.api.external.v1.app import v1_app
|
||||
from backend.blocks._base import BlockType
|
||||
|
||||
ROUTES = [
|
||||
("GET", "/me", "read"),
|
||||
("GET", "/blocks", "read"),
|
||||
("POST", "/blocks/{block_id}/execute", "execution"),
|
||||
("POST", "/graphs", "execution"),
|
||||
("POST", "/graphs/{graph_id}/execute/{graph_version}", "execution"),
|
||||
("GET", "/graphs/{graph_id}/executions/{graph_exec_id}/results", "read"),
|
||||
("GET", "/store/agents", "read"),
|
||||
("GET", "/store/agents/{username}/{agent_name}", "read"),
|
||||
("GET", "/store/creators", "read"),
|
||||
("GET", "/store/creators/{username}", "read"),
|
||||
("GET", "/integrations/providers", "read"),
|
||||
("GET", "/integrations/credentials", "read"),
|
||||
("GET", "/integrations/{provider}/credentials", "read"),
|
||||
("POST", "/integrations/{provider}/oauth/initiate", "execution"),
|
||||
("POST", "/integrations/{provider}/oauth/complete", "execution"),
|
||||
("POST", "/integrations/{provider}/credentials", "execution"),
|
||||
("DELETE", "/integrations/{provider}/credentials/{cred_id}", "execution"),
|
||||
("POST", "/tools/find-agent", "execution"),
|
||||
("POST", "/tools/run-agent", "execution"),
|
||||
]
|
||||
|
||||
|
||||
async def test_read_cap_is_shared_across_credentials_and_resets_at_reported_time(
|
||||
v1_client, v1_counter, v1_clock, v1_generated_keys, v1_key_rows
|
||||
):
|
||||
for index in range(rate_limit.READ_MAX_REQUESTS):
|
||||
headers = (
|
||||
{"X-API-Key": v1_generated_keys[index].key}
|
||||
if index < 2
|
||||
else {"Authorization": "Bearer agpt_xt_test-0"}
|
||||
)
|
||||
response = await v1_client.get("/v1/me", headers=headers)
|
||||
assert response.status_code == 200
|
||||
assert int(response.headers["x-ratelimit-remaining"]) == 299 - index
|
||||
|
||||
headers = {"Authorization": "Bearer agpt_xt_test-1"}
|
||||
rejected = await v1_client.get("/v1/me", headers=headers)
|
||||
assert rejected.status_code == 429
|
||||
assert rejected.headers["x-ratelimit-limit"] == "300"
|
||||
assert rejected.headers["x-ratelimit-remaining"] == "0"
|
||||
assert rejected.headers["retry-after"] == "3"
|
||||
assert "no-store" in rejected.headers["cache-control"]
|
||||
assert (
|
||||
int(rejected.headers["x-ratelimit-reset"])
|
||||
== int(v1_clock.now.return_value.timestamp()) + 3
|
||||
)
|
||||
assert len({call.args[0] for call in v1_counter.await_args_list}) == 1
|
||||
|
||||
other = await v1_client.get(
|
||||
"/v1/me", headers={"Authorization": "Bearer agpt_xt_test-2"}
|
||||
)
|
||||
assert other.status_code == 200
|
||||
assert other.json()["id"] == v1_key_rows[2].user_id
|
||||
assert other.headers["x-ratelimit-remaining"] == "299"
|
||||
|
||||
v1_clock.now.return_value += timedelta(seconds=3)
|
||||
retried = await v1_client.get("/v1/me", headers=headers)
|
||||
assert retried.status_code == 200
|
||||
assert retried.headers["x-ratelimit-remaining"] == "299"
|
||||
|
||||
|
||||
async def test_execution_cap_blocks_work_without_consuming_the_read_budget(
|
||||
mocker, v1_client, v1_counter
|
||||
):
|
||||
async def execute(*args, **kwargs):
|
||||
yield "result", "ok"
|
||||
|
||||
block = MagicMock(
|
||||
id="test-block", block_type=BlockType.STANDARD, disabled=False, execute=execute
|
||||
)
|
||||
get_block = mocker.patch("backend.blocks.get_block", return_value=block)
|
||||
charge = mocker.patch.object(routes, "charge_for_direct_block_execution")
|
||||
headers = {"Authorization": "Bearer agpt_xt_test-0"}
|
||||
for index in range(rate_limit.EXECUTION_MAX_REQUESTS):
|
||||
response = await v1_client.post(
|
||||
"/v1/blocks/test-block/execute", headers=headers, json={}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"result": ["ok"]}
|
||||
assert int(response.headers["x-ratelimit-remaining"]) == 29 - index
|
||||
|
||||
rejected = await v1_client.post(
|
||||
"/v1/blocks/test-block/execute", headers=headers, json={}
|
||||
)
|
||||
assert rejected.status_code == 429
|
||||
assert rejected.headers["x-ratelimit-limit"] == "30"
|
||||
assert get_block.call_count == charge.await_count == 30
|
||||
read = await v1_client.get("/v1/me", headers=headers)
|
||||
assert read.status_code == 200
|
||||
assert read.headers["x-ratelimit-remaining"] == "299"
|
||||
|
||||
|
||||
def test_route_matrix_covers_every_v1_operation():
|
||||
assert {(method, path) for method, path, _ in ROUTES} == {
|
||||
(method, route.path)
|
||||
for route in v1_app.routes
|
||||
if isinstance(route, APIRoute)
|
||||
for method in route.methods
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method,path,tier", ROUTES)
|
||||
async def test_every_route_enforces_exactly_its_declared_tier(
|
||||
mocker, v1_client, v1_counter, v1_key_rows, method, path, tier
|
||||
):
|
||||
limit = 300 if tier == "read" else 30
|
||||
v1_counter.side_effect = None
|
||||
v1_counter.return_value = limit + 1
|
||||
path = path.format(
|
||||
block_id="block",
|
||||
graph_id="graph",
|
||||
graph_version=1,
|
||||
graph_exec_id="run",
|
||||
username="user",
|
||||
agent_name="agent",
|
||||
provider="github",
|
||||
cred_id="cred",
|
||||
)
|
||||
metric = mocker.patch.object(rate_limit, "record_rate_limit_hit")
|
||||
response = await v1_client.request(
|
||||
method,
|
||||
f"/v1{path}",
|
||||
json={},
|
||||
headers={"Authorization": "Bearer agpt_xt_test-0"},
|
||||
)
|
||||
|
||||
assert response.status_code == 429
|
||||
assert response.headers["x-ratelimit-limit"] == str(limit)
|
||||
assert response.headers["x-ratelimit-remaining"] == "0"
|
||||
v1_counter.assert_awaited_once()
|
||||
assert f":{v1_key_rows[0].user_id}:{tier}:" in v1_counter.await_args.args[0]
|
||||
metric.assert_called_once_with(f"external_api:{tier}", v1_key_rows[0].user_id)
|
||||
|
||||
|
||||
async def test_redis_outage_preserves_authentication_and_permissions(
|
||||
v1_client, v1_counter, v1_tokens
|
||||
):
|
||||
v1_counter.side_effect = RedisError("unavailable")
|
||||
headers = {"Authorization": "Bearer agpt_xt_test-0"}
|
||||
allowed = await v1_client.get("/v1/me", headers=headers)
|
||||
assert allowed.status_code == 200
|
||||
assert "x-ratelimit-limit" not in allowed.headers
|
||||
v1_tokens["agpt_xt_test-0"].scopes = []
|
||||
assert (await v1_client.get("/v1/me", headers=headers)).status_code == 403
|
||||
v1_counter.reset_mock()
|
||||
assert (await v1_client.get("/v1/me")).status_code == 401
|
||||
v1_counter.assert_not_awaited()
|
||||
+41
-7
@@ -12,6 +12,7 @@ import backend.api.features.store.db as store_db
|
||||
import backend.api.features.store.model as store_model
|
||||
import backend.blocks
|
||||
from backend.api.external.middleware import require_auth, require_permission
|
||||
from backend.api.external.rate_limit import EXECUTION_LIMIT, READ_LIMIT
|
||||
from backend.blocks._base import AnyBlockSchema
|
||||
from backend.copilot.capabilities.block_meta import is_graph_only_block
|
||||
from backend.copilot.rate_limit import UserPaywalledError, enforce_payment_paywall
|
||||
@@ -47,6 +48,11 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
v1_router = APIRouter()
|
||||
|
||||
# Each integrations route declares its own tier in ``integrations.py``: the
|
||||
# reads take READ_LIMIT, while credential writes and the two OAuth hops (which
|
||||
# call out to a provider) take EXECUTION_LIMIT. A router-level dependency
|
||||
# cannot express that split — FastAPI would AND it with every route's own
|
||||
# dependency, so a write would be counted against both budgets.
|
||||
v1_router.include_router(integrations_router)
|
||||
v1_router.include_router(tools_router)
|
||||
|
||||
@@ -64,6 +70,9 @@ class UserInfoResponse(BaseModel):
|
||||
@v1_router.get(
|
||||
path="/me",
|
||||
tags=["user", "meta"],
|
||||
dependencies=[
|
||||
READ_LIMIT,
|
||||
],
|
||||
)
|
||||
async def get_user_info(
|
||||
auth: APIAuthorizationInfo = Security(
|
||||
@@ -83,7 +92,10 @@ async def get_user_info(
|
||||
@v1_router.get(
|
||||
path="/blocks",
|
||||
tags=["blocks"],
|
||||
dependencies=[Security(require_permission(APIKeyPermission.READ_BLOCK))],
|
||||
dependencies=[
|
||||
Security(require_permission(APIKeyPermission.READ_BLOCK)),
|
||||
READ_LIMIT,
|
||||
],
|
||||
)
|
||||
async def get_graph_blocks() -> Sequence[dict[Any, Any]]:
|
||||
blocks = [block() for block in backend.blocks.get_blocks().values()]
|
||||
@@ -93,7 +105,10 @@ async def get_graph_blocks() -> Sequence[dict[Any, Any]]:
|
||||
@v1_router.post(
|
||||
path="/blocks/{block_id}/execute",
|
||||
tags=["blocks"],
|
||||
dependencies=[Security(require_permission(APIKeyPermission.EXECUTE_BLOCK))],
|
||||
dependencies=[
|
||||
Security(require_permission(APIKeyPermission.EXECUTE_BLOCK)),
|
||||
EXECUTION_LIMIT,
|
||||
],
|
||||
)
|
||||
async def execute_graph_block(
|
||||
block_id: str,
|
||||
@@ -186,7 +201,8 @@ async def _run_direct_block(
|
||||
require_permission(
|
||||
APIKeyPermission.WRITE_GRAPH, APIKeyPermission.WRITE_LIBRARY
|
||||
)
|
||||
)
|
||||
),
|
||||
EXECUTION_LIMIT,
|
||||
],
|
||||
)
|
||||
async def create_graph(
|
||||
@@ -233,6 +249,9 @@ async def create_graph(
|
||||
@v1_router.post(
|
||||
path="/graphs/{graph_id}/execute/{graph_version}",
|
||||
tags=["graphs"],
|
||||
dependencies=[
|
||||
EXECUTION_LIMIT,
|
||||
],
|
||||
)
|
||||
async def execute_graph(
|
||||
graph_id: str,
|
||||
@@ -309,6 +328,9 @@ class GraphExecutionResult(TypedDict):
|
||||
@v1_router.get(
|
||||
path="/graphs/{graph_id}/executions/{graph_exec_id}/results",
|
||||
tags=["graphs"],
|
||||
dependencies=[
|
||||
READ_LIMIT,
|
||||
],
|
||||
)
|
||||
async def get_graph_execution_results(
|
||||
graph_id: str,
|
||||
@@ -365,7 +387,10 @@ async def get_graph_execution_results(
|
||||
@v1_router.get(
|
||||
path="/store/agents",
|
||||
tags=["store"],
|
||||
dependencies=[Security(require_auth)], # data is public; auth required as anti-DDoS
|
||||
dependencies=[
|
||||
Security(require_auth), # data is public; auth required as anti-DDoS
|
||||
READ_LIMIT,
|
||||
],
|
||||
response_model=store_model.StoreAgentsResponse,
|
||||
)
|
||||
async def get_store_agents(
|
||||
@@ -413,7 +438,10 @@ async def get_store_agents(
|
||||
@v1_router.get(
|
||||
path="/store/agents/{username}/{agent_name}",
|
||||
tags=["store"],
|
||||
dependencies=[Security(require_auth)], # data is public; auth required as anti-DDoS
|
||||
dependencies=[
|
||||
Security(require_auth), # data is public; auth required as anti-DDoS
|
||||
READ_LIMIT,
|
||||
],
|
||||
response_model=store_model.StoreAgentDetails,
|
||||
)
|
||||
async def get_store_agent(
|
||||
@@ -441,7 +469,10 @@ async def get_store_agent(
|
||||
@v1_router.get(
|
||||
path="/store/creators",
|
||||
tags=["store"],
|
||||
dependencies=[Security(require_auth)], # data is public; auth required as anti-DDoS
|
||||
dependencies=[
|
||||
Security(require_auth), # data is public; auth required as anti-DDoS
|
||||
READ_LIMIT,
|
||||
],
|
||||
response_model=store_model.CreatorsResponse,
|
||||
)
|
||||
async def get_store_creators(
|
||||
@@ -483,7 +514,10 @@ async def get_store_creators(
|
||||
@v1_router.get(
|
||||
path="/store/creators/{username}",
|
||||
tags=["store"],
|
||||
dependencies=[Security(require_auth)], # data is public; auth required as anti-DDoS
|
||||
dependencies=[
|
||||
Security(require_auth), # data is public; auth required as anti-DDoS
|
||||
READ_LIMIT,
|
||||
],
|
||||
response_model=store_model.CreatorDetails,
|
||||
)
|
||||
async def get_store_creator(
|
||||
|
||||
@@ -24,8 +24,8 @@ client = fastapi.testclient.TestClient(app)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup_auth(test_user_id):
|
||||
"""Override require_auth to return a synthetic API-key principal with all scopes."""
|
||||
def setup_auth(test_user_id, v1_counter):
|
||||
"""Give each case an authenticated principal and a fresh quota counter."""
|
||||
|
||||
async def fake_require_auth() -> APIAuthorizationInfo:
|
||||
return APIAuthorizationInfo(
|
||||
@@ -272,8 +272,11 @@ def test_disabled_block_still_403(monkeypatch: pytest.MonkeyPatch):
|
||||
|
||||
def test_registry_blocks_needing_a_workflow_are_refused(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
v1_counter: AsyncMock,
|
||||
):
|
||||
"""The registry's own blocks are refused before they are charged or run."""
|
||||
"""Check graph-only blocks before charging or running, independent of quota."""
|
||||
v1_counter.side_effect = None
|
||||
v1_counter.return_value = 1
|
||||
ran: list[str] = []
|
||||
|
||||
async def _record_run(self, _data, **_kwargs):
|
||||
|
||||
@@ -3,8 +3,8 @@
|
||||
Note: These endpoints use ephemeral sessions that are not persisted to Redis.
|
||||
As a result, session-based rate limiting (max_agent_runs, max_agent_schedules)
|
||||
is not enforced for external API calls. Each request creates a fresh session
|
||||
with zeroed counters. Rate limiting for external API consumers should be
|
||||
handled separately (e.g., via API key quotas).
|
||||
with zeroed counters. Both routes use rate_limit.EXECUTION_LIMIT to share the
|
||||
per-user execution budget with the other v1 write operations.
|
||||
"""
|
||||
|
||||
import logging
|
||||
@@ -15,6 +15,7 @@ from prisma.enums import APIKeyPermission
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from backend.api.external.middleware import require_permission
|
||||
from backend.api.external.rate_limit import EXECUTION_LIMIT
|
||||
from backend.copilot.model import ChatSession
|
||||
from backend.copilot.tools import find_agent_tool, run_agent_tool
|
||||
from backend.copilot.tools.models import ToolResponseBase
|
||||
@@ -84,6 +85,7 @@ def _create_ephemeral_session(
|
||||
|
||||
@tools_router.post(
|
||||
path="/find-agent",
|
||||
dependencies=[EXECUTION_LIMIT],
|
||||
)
|
||||
async def find_agent(
|
||||
request: FindAgentRequest,
|
||||
@@ -111,6 +113,7 @@ async def find_agent(
|
||||
|
||||
@tools_router.post(
|
||||
path="/run-agent",
|
||||
dependencies=[EXECUTION_LIMIT],
|
||||
)
|
||||
async def run_agent(
|
||||
request: RunAgentRequest,
|
||||
|
||||
@@ -22,6 +22,7 @@ open, and never crashes the service.
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Optional
|
||||
from urllib.parse import quote
|
||||
|
||||
from autogpt_libs.auth import User, requires_admin_user
|
||||
from fastapi import APIRouter, Body, HTTPException, Security
|
||||
@@ -110,14 +111,17 @@ async def notify_impersonation_start(
|
||||
async def _send_impersonation_alert(admin: User, target_user_id: str) -> None:
|
||||
"""Send the Discord alert. Raises if it was not actually delivered."""
|
||||
# Prefer the email carried in the JWT; fall back to a DB lookup if absent.
|
||||
# The target is a user, so the alert names them by ID only, never by email.
|
||||
admin_email = admin.email or await _resolve_email(admin.user_id)
|
||||
target_email = await _resolve_email(target_user_id)
|
||||
base_url = settings.config.frontend_base_url or settings.config.platform_base_url
|
||||
admin_link = f"{base_url}/admin/spending?search={quote(target_user_id, safe='')}"
|
||||
|
||||
content = (
|
||||
"🕵️ **Admin impersonation started**\n"
|
||||
f"Admin: {admin_email or 'unknown'} (`{admin.user_id}`)\n"
|
||||
f"Now viewing as: {target_email or 'unknown'} (`{target_user_id}`)\n"
|
||||
"Source: admin dashboard"
|
||||
f"Now viewing as: `{target_user_id}`\n"
|
||||
"Source: admin dashboard\n"
|
||||
f"[View User Details]({admin_link})"
|
||||
)
|
||||
|
||||
status = await asyncio.wait_for(
|
||||
|
||||
+35
-8
@@ -22,6 +22,7 @@ _MOCK_MODULE = "backend.api.features.admin.impersonation_admin_routes"
|
||||
_TARGET_USER_ID = "target-user-id"
|
||||
_ADMIN_EMAIL = "[email protected]"
|
||||
_TARGET_EMAIL = "[email protected]"
|
||||
_BASE_URL = "https://platform.test"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
@@ -36,7 +37,12 @@ def _set_token(mocker: pytest_mock.MockerFixture, token: str) -> None:
|
||||
"""Swap the module's settings for a tiny stub exposing only what's read."""
|
||||
mocker.patch(
|
||||
f"{_MOCK_MODULE}.settings",
|
||||
SimpleNamespace(secrets=SimpleNamespace(discord_bot_token=token)),
|
||||
SimpleNamespace(
|
||||
secrets=SimpleNamespace(discord_bot_token=token),
|
||||
config=SimpleNamespace(
|
||||
frontend_base_url=_BASE_URL, platform_base_url="unused"
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@@ -176,10 +182,33 @@ def test_notify_missing_target_user_id_returns_422() -> None:
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
def test_notify_email_lookup_failure_is_non_fatal(
|
||||
def test_notify_names_target_by_id_not_email(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
) -> None:
|
||||
"""A failed email lookup degrades to 'unknown' but still delivers (200)."""
|
||||
"""The target user is shown by ID with an admin link; their email never is."""
|
||||
_set_token(mocker, "fake-token")
|
||||
_patch_emails(mocker)
|
||||
mock_alert = _patch_alert(mocker, return_value="Message sent")
|
||||
|
||||
response = client.post(
|
||||
"/admin/impersonation/notify", json={"target_user_id": _TARGET_USER_ID}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
content = mock_alert.await_args.args[0]
|
||||
assert _TARGET_EMAIL not in content
|
||||
assert f"`{_TARGET_USER_ID}`" in content
|
||||
assert f"{_BASE_URL}/admin/spending?search={_TARGET_USER_ID}" in content
|
||||
|
||||
|
||||
def test_notify_admin_lookup_failure_is_non_fatal(
|
||||
mocker: pytest_mock.MockerFixture, admin_user_id: str
|
||||
) -> None:
|
||||
"""A JWT without an email and a failed lookup degrade to 'unknown' (200)."""
|
||||
app.dependency_overrides[get_jwt_payload] = lambda: {
|
||||
"sub": admin_user_id,
|
||||
"role": "admin",
|
||||
}
|
||||
_set_token(mocker, "fake-token")
|
||||
mocker.patch(
|
||||
f"{_MOCK_MODULE}.get_user_email_by_id",
|
||||
@@ -195,14 +224,13 @@ def test_notify_email_lookup_failure_is_non_fatal(
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"alerted": True}
|
||||
content = mock_alert.await_args.args[0]
|
||||
assert "[email protected]" in content # admin email from the JWT
|
||||
assert "unknown" in content # target lookup failed, degraded gracefully
|
||||
assert f"Admin: unknown (`{admin_user_id}`)" in content
|
||||
|
||||
|
||||
def test_notify_admin_email_comes_from_jwt_without_db_lookup(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
) -> None:
|
||||
"""The admin email is read from the JWT; only the target is looked up in DB."""
|
||||
"""The admin email is read from the JWT, so no email is looked up in the DB."""
|
||||
_set_token(mocker, "fake-token")
|
||||
mock_email = mocker.patch(
|
||||
f"{_MOCK_MODULE}.get_user_email_by_id",
|
||||
@@ -218,5 +246,4 @@ def test_notify_admin_email_comes_from_jwt_without_db_lookup(
|
||||
assert response.status_code == 200
|
||||
content = mock_alert.await_args.args[0]
|
||||
assert "[email protected]" in content
|
||||
# Admin email came from the JWT, so the DB is only hit for the target.
|
||||
mock_email.assert_awaited_once_with(_TARGET_USER_ID)
|
||||
mock_email.assert_not_awaited()
|
||||
@@ -27,6 +27,7 @@ from backend.api.features.billing.credits_rate_limit import (
|
||||
from backend.copilot.rate_limit import get_tier_multipliers
|
||||
from backend.data import checkout_audience
|
||||
from backend.data.credit import (
|
||||
BillingCycle,
|
||||
PendingChangeUnknown,
|
||||
UserCredit,
|
||||
cancel_stripe_subscription,
|
||||
@@ -119,7 +120,7 @@ class SubscriptionStatusResponse(BaseModel):
|
||||
default_factory=dict,
|
||||
description=(
|
||||
"Tier → yearly amount in cents. Populated only for tiers with a"
|
||||
" yearly Stripe price configured in LaunchDarkly. Empty for"
|
||||
" yearly Stripe price configured in PostHog. Empty for"
|
||||
" monthly-only configurations."
|
||||
),
|
||||
)
|
||||
@@ -281,21 +282,22 @@ async def get_subscription_status(
|
||||
SubscriptionTier.MAX,
|
||||
SubscriptionTier.BUSINESS,
|
||||
]
|
||||
monthly_price_ids, yearly_price_ids = await asyncio.gather(
|
||||
asyncio.gather(
|
||||
*[get_subscription_price_id(t, "monthly") for t in priceable_tiers]
|
||||
),
|
||||
asyncio.gather(
|
||||
*[get_subscription_price_id(t, "yearly") for t in priceable_tiers]
|
||||
),
|
||||
)
|
||||
|
||||
async def _cost(pid: str | None) -> int:
|
||||
return (await _get_stripe_price_amount(pid) or 0) if pid else 0
|
||||
|
||||
monthly_costs, yearly_costs = await asyncio.gather(
|
||||
asyncio.gather(*[_cost(pid) for pid in monthly_price_ids]),
|
||||
asyncio.gather(*[_cost(pid) for pid in yearly_price_ids]),
|
||||
# None of these reads needs another's answer, so they run side by side;
|
||||
# only the proration below needs the current plan's cost.
|
||||
(
|
||||
(monthly_price_ids, yearly_price_ids, monthly_costs, yearly_costs),
|
||||
multipliers,
|
||||
billing_cycle,
|
||||
current_period_end,
|
||||
paused_workflow_schedules,
|
||||
pending,
|
||||
) = await asyncio.gather(
|
||||
_tier_prices(priceable_tiers),
|
||||
get_tier_multipliers(),
|
||||
get_user_billing_cycle(user_id),
|
||||
get_active_subscription_period_end(user_id),
|
||||
_count_paused_workflow_schedules(user_id, tier),
|
||||
_pending_subscription_change(user_id),
|
||||
)
|
||||
|
||||
# Row visibility: include a tier if EITHER cycle is configured. Monthly
|
||||
@@ -320,7 +322,6 @@ async def get_subscription_status(
|
||||
# tier without hard-coding backend defaults. Only emit entries for tiers
|
||||
# that land in ``tier_costs`` — rows hidden at the price layer must stay
|
||||
# hidden in the multiplier layer too.
|
||||
multipliers = await get_tier_multipliers()
|
||||
# get_tier_multipliers() keys by tier string value (see its docstring),
|
||||
# so the lookup must use t.value — passing the enum t silently misses
|
||||
# every tier and falls back to 1.0, ignoring LD-configured multipliers.
|
||||
@@ -330,34 +331,12 @@ async def get_subscription_status(
|
||||
if t.value in tier_costs
|
||||
}
|
||||
|
||||
user_cycle = await get_user_billing_cycle(user_id) or "monthly"
|
||||
user_cycle = billing_cycle or "monthly"
|
||||
if user_cycle == "yearly":
|
||||
current_monthly_cost = tier_costs_yearly.get(tier.value, 0)
|
||||
else:
|
||||
current_monthly_cost = tier_costs.get(tier.value, 0)
|
||||
(
|
||||
proration_credit,
|
||||
current_period_end,
|
||||
paused_workflow_schedules,
|
||||
) = await asyncio.gather(
|
||||
get_proration_credit_cents(user_id, current_monthly_cost),
|
||||
get_active_subscription_period_end(user_id),
|
||||
_count_paused_workflow_schedules(user_id, tier),
|
||||
)
|
||||
|
||||
try:
|
||||
pending = await get_pending_subscription_change(user_id)
|
||||
except (stripe.StripeError, PendingChangeUnknown):
|
||||
# Swallow Stripe-side failures (rate limits, transient network) AND
|
||||
# PendingChangeUnknown (LaunchDarkly price-id lookup failed). Both
|
||||
# propagate past the cache so the next request retries fresh instead
|
||||
# of serving a stale None for the TTL window. Let real bugs (KeyError,
|
||||
# AttributeError, etc.) propagate so they surface in Sentry.
|
||||
logger.exception(
|
||||
"get_subscription_status: failed to resolve pending change for user %s",
|
||||
user_id,
|
||||
)
|
||||
pending = None
|
||||
proration_credit = await get_proration_credit_cents(user_id, current_monthly_cost)
|
||||
|
||||
response = SubscriptionStatusResponse(
|
||||
tier=tier.value,
|
||||
@@ -402,6 +381,48 @@ async def _count_paused_workflow_schedules(user_id: str, tier: SubscriptionTier)
|
||||
return sum(schedule.kind == "graph" for schedule in schedules)
|
||||
|
||||
|
||||
async def _tier_prices(
|
||||
tiers: list[SubscriptionTier],
|
||||
) -> tuple[list[str | None], list[str | None], list[int], list[int]]:
|
||||
"""Each tier's monthly and yearly price ids, then what each one costs."""
|
||||
monthly_price_ids, yearly_price_ids = await asyncio.gather(
|
||||
asyncio.gather(*[get_subscription_price_id(t, "monthly") for t in tiers]),
|
||||
asyncio.gather(*[get_subscription_price_id(t, "yearly") for t in tiers]),
|
||||
)
|
||||
monthly_costs, yearly_costs = await asyncio.gather(
|
||||
asyncio.gather(*[_price_cost(pid) for pid in monthly_price_ids]),
|
||||
asyncio.gather(*[_price_cost(pid) for pid in yearly_price_ids]),
|
||||
)
|
||||
return (
|
||||
list(monthly_price_ids),
|
||||
list(yearly_price_ids),
|
||||
list(monthly_costs),
|
||||
list(yearly_costs),
|
||||
)
|
||||
|
||||
|
||||
async def _price_cost(price_id: str | None) -> int:
|
||||
return (await _get_stripe_price_amount(price_id) or 0) if price_id else 0
|
||||
|
||||
|
||||
async def _pending_subscription_change(
|
||||
user_id: str,
|
||||
) -> tuple[SubscriptionTier, datetime, BillingCycle | None] | None:
|
||||
try:
|
||||
return await get_pending_subscription_change(user_id)
|
||||
except (stripe.StripeError, PendingChangeUnknown):
|
||||
# Swallow Stripe-side failures (rate limits, transient network) AND
|
||||
# PendingChangeUnknown (PostHog price-id lookup failed). Both
|
||||
# propagate past the cache so the next request retries fresh instead
|
||||
# of serving a stale None for the TTL window. Let real bugs (KeyError,
|
||||
# AttributeError, etc.) propagate so they surface in Sentry.
|
||||
logger.exception(
|
||||
"get_subscription_status: failed to resolve pending change for user %s",
|
||||
user_id,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
@router.post(
|
||||
path="/credits/subscription/preview",
|
||||
summary="Review the prorated charge for a paid Pro-to-Max upgrade",
|
||||
|
||||
@@ -77,11 +77,8 @@ async def submit_message_feedback(
|
||||
"""Rate a reply in one of the caller's chats.
|
||||
|
||||
Saves the rating, one per reply and score name so rating again
|
||||
overwrites it, and scores it in Langfuse on the trace of the turn that
|
||||
wrote the reply, or on the chat's Langfuse session when that trace is
|
||||
not known. It also goes to PostHog LLM analytics as an ``$ai_metric``
|
||||
event. Langfuse or PostHog being down does not fail the request; the
|
||||
response says where the Langfuse score went.
|
||||
overwrites it, and sends it to PostHog LLM analytics as an
|
||||
``$ai_metric`` event. PostHog being down does not fail the request.
|
||||
"""
|
||||
feedback = await record_message_feedback(
|
||||
user_id=user_id,
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
"""Route tests for POST /sessions/{session_id}/feedback (thumbs up/down, copy).
|
||||
|
||||
The database layer and the Langfuse client are mocked at the boundary;
|
||||
The database layer and PostHog are mocked at the boundary;
|
||||
``backend/copilot/feedback_db_test.py`` proves the owner scoping against a
|
||||
real Postgres.
|
||||
"""
|
||||
@@ -13,7 +13,6 @@ import fastapi.testclient
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from autogpt_libs.auth.jwt_utils import get_jwt_payload
|
||||
from langfuse.api import ScoreDataType
|
||||
from prisma.models import ChatMessage as PrismaChatMessage
|
||||
|
||||
from backend.api.features.chat import feedback as feedback_routes
|
||||
@@ -26,7 +25,6 @@ client = fastapi.testclient.TestClient(app)
|
||||
|
||||
SESSION_ID = "8d2f5c1e-2b7a-4c55-9d0e-3f6a1b2c4d5e"
|
||||
MESSAGE_ROW_ID = "0f9e8d7c-6b5a-4938-8271-605f4e3d2c1b"
|
||||
TRACE_ID = "1edf31f11b1693cc6103f358c1481694"
|
||||
FEEDBACK_ID = "c3a1b2d4-e5f6-4789-8abc-def012345678"
|
||||
|
||||
|
||||
@@ -37,14 +35,6 @@ def setup_app_auth(mock_jwt_user):
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def langfuse_settings(mocker: pytest_mock.MockerFixture):
|
||||
secrets = feedback_service.settings.secrets
|
||||
mocker.patch.object(secrets, "langfuse_public_key", "pk-lf-test")
|
||||
mocker.patch.object(secrets, "langfuse_secret_key", "sk-lf-test")
|
||||
mocker.patch.object(secrets, "langfuse_tracing_environment", "dev")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def posthog_capture(mocker: pytest_mock.MockerFixture) -> MagicMock:
|
||||
"""PostHog is covered in ``backend/copilot/feedback_test.py``; here it
|
||||
@@ -52,15 +42,6 @@ def posthog_capture(mocker: pytest_mock.MockerFixture) -> MagicMock:
|
||||
return mocker.patch.object(feedback_service, "capture")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def score_create(mocker: pytest_mock.MockerFixture) -> AsyncMock:
|
||||
"""The Langfuse score API call, with the client it hangs off mocked."""
|
||||
langfuse = MagicMock()
|
||||
langfuse.async_api.score.create = AsyncMock()
|
||||
mocker.patch.object(feedback_service, "get_client", return_value=langfuse)
|
||||
return langfuse.async_api.score.create
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def upsert(mocker: pytest_mock.MockerFixture) -> AsyncMock:
|
||||
return mocker.patch.object(
|
||||
@@ -70,7 +51,7 @@ def upsert(mocker: pytest_mock.MockerFixture) -> AsyncMock:
|
||||
)
|
||||
|
||||
|
||||
def _reply(*, trace_id: str | None, sequence: int = 7) -> PrismaChatMessage:
|
||||
def _reply(*, sequence: int = 7) -> PrismaChatMessage:
|
||||
return PrismaChatMessage(
|
||||
id=MESSAGE_ROW_ID,
|
||||
createdAt=datetime.now(UTC),
|
||||
@@ -78,7 +59,6 @@ def _reply(*, trace_id: str | None, sequence: int = 7) -> PrismaChatMessage:
|
||||
role="assistant",
|
||||
content="Here is the report.",
|
||||
sequence=sequence,
|
||||
langfuseTraceId=trace_id,
|
||||
)
|
||||
|
||||
|
||||
@@ -122,14 +102,13 @@ def test_feedback_on_another_users_session_is_not_found(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
test_user_id: str,
|
||||
upsert: AsyncMock,
|
||||
score_create: AsyncMock,
|
||||
) -> None:
|
||||
"""The lookup is scoped to the caller, so a session owned by someone else
|
||||
yields no message and the same 404 an unknown session would."""
|
||||
owner_id = "someone-else"
|
||||
|
||||
async def owner_only(user_id: str, session_id: str, **_):
|
||||
return _reply(trace_id=TRACE_ID) if user_id == owner_id else None
|
||||
return _reply() if user_id == owner_id else None
|
||||
|
||||
lookup = mocker.patch.object(
|
||||
feedback_db, "get_rateable_message", new=AsyncMock(side_effect=owner_only)
|
||||
@@ -141,13 +120,11 @@ def test_feedback_on_another_users_session_is_not_found(
|
||||
assert lookup.await_args is not None
|
||||
assert lookup.await_args.args == (test_user_id, SESSION_ID)
|
||||
upsert.assert_not_awaited()
|
||||
score_create.assert_not_awaited()
|
||||
|
||||
|
||||
def test_feedback_on_unknown_message_is_not_found(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
upsert: AsyncMock,
|
||||
score_create: AsyncMock,
|
||||
) -> None:
|
||||
_mock_message(mocker, None)
|
||||
|
||||
@@ -156,24 +133,22 @@ def test_feedback_on_unknown_message_is_not_found(
|
||||
assert response.status_code == 404
|
||||
assert "not found" in response.json()["detail"]
|
||||
upsert.assert_not_awaited()
|
||||
score_create.assert_not_awaited()
|
||||
|
||||
|
||||
# ─── Langfuse score ────────────────────────────────────────────────────
|
||||
# ─── Rating ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_thumbs_down_scores_the_turn_trace(
|
||||
def test_thumbs_down_is_saved(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
test_user_id: str,
|
||||
upsert: AsyncMock,
|
||||
score_create: AsyncMock,
|
||||
) -> None:
|
||||
lookup = _mock_message(mocker, _reply(trace_id=TRACE_ID))
|
||||
lookup = _mock_message(mocker, _reply())
|
||||
|
||||
response = _post(_thumbs_down())
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"id": FEEDBACK_ID, "langfuse_target": "trace"}
|
||||
assert response.json() == {"id": FEEDBACK_ID}
|
||||
lookup.assert_awaited_once_with(test_user_id, SESSION_ID, sequence=7)
|
||||
upsert.assert_awaited_once_with(
|
||||
user_id=test_user_id,
|
||||
@@ -182,61 +157,14 @@ def test_thumbs_down_scores_the_turn_trace(
|
||||
score_name="user-feedback",
|
||||
score_value=0,
|
||||
comment="It edited the wrong file",
|
||||
langfuse_trace_id=TRACE_ID,
|
||||
)
|
||||
score_create.assert_awaited_once()
|
||||
request = score_create.await_args.kwargs["request"]
|
||||
assert request.id == FEEDBACK_ID
|
||||
assert request.trace_id == TRACE_ID
|
||||
assert request.session_id is None
|
||||
assert request.name == "user-feedback"
|
||||
assert request.value == 0
|
||||
assert request.data_type == ScoreDataType.NUMERIC
|
||||
assert request.comment == "It edited the wrong file"
|
||||
assert request.metadata == {
|
||||
"message_id": MESSAGE_ROW_ID,
|
||||
"message_sequence": 7,
|
||||
"session_id": SESSION_ID,
|
||||
}
|
||||
assert request.environment == "dev"
|
||||
|
||||
|
||||
def test_reply_without_a_trace_falls_back_to_a_session_score(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
upsert: AsyncMock,
|
||||
score_create: AsyncMock,
|
||||
) -> None:
|
||||
"""Baseline turns and replies from before the trace stamp have no trace:
|
||||
the score goes on the chat's Langfuse session, naming the message."""
|
||||
_mock_message(mocker, _reply(trace_id=None))
|
||||
|
||||
response = _post(
|
||||
{
|
||||
"message_id": f"{SESSION_ID}-seq-7",
|
||||
"score_name": "user-feedback",
|
||||
"score_value": 1,
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"id": FEEDBACK_ID, "langfuse_target": "session"}
|
||||
assert upsert.await_args is not None
|
||||
assert upsert.await_args.kwargs["langfuse_trace_id"] is None
|
||||
request = score_create.await_args.kwargs["request"]
|
||||
assert request.trace_id is None
|
||||
assert request.session_id == SESSION_ID
|
||||
assert request.value == 1
|
||||
assert request.comment is None
|
||||
assert request.metadata["message_id"] == MESSAGE_ROW_ID
|
||||
assert request.metadata["session_id"] == SESSION_ID
|
||||
|
||||
|
||||
def test_copy_is_scored_under_its_own_name(
|
||||
def test_copy_is_saved_under_its_own_name(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
upsert: AsyncMock,
|
||||
score_create: AsyncMock,
|
||||
) -> None:
|
||||
_mock_message(mocker, _reply(trace_id=TRACE_ID))
|
||||
_mock_message(mocker, _reply())
|
||||
|
||||
response = _post(
|
||||
{
|
||||
@@ -247,19 +175,17 @@ def test_copy_is_scored_under_its_own_name(
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
request = score_create.await_args.kwargs["request"]
|
||||
assert request.name == "copy"
|
||||
assert request.value == 1
|
||||
assert request.trace_id == TRACE_ID
|
||||
assert upsert.await_args is not None
|
||||
assert upsert.await_args.kwargs["score_name"] == "copy"
|
||||
assert upsert.await_args.kwargs["score_value"] == 1
|
||||
|
||||
|
||||
def test_raw_row_id_is_looked_up_by_id(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
test_user_id: str,
|
||||
upsert: AsyncMock,
|
||||
score_create: AsyncMock,
|
||||
) -> None:
|
||||
lookup = _mock_message(mocker, _reply(trace_id=TRACE_ID))
|
||||
lookup = _mock_message(mocker, _reply())
|
||||
|
||||
response = _post({**_thumbs_down(), "message_id": MESSAGE_ROW_ID})
|
||||
|
||||
@@ -267,47 +193,12 @@ def test_raw_row_id_is_looked_up_by_id(
|
||||
lookup.assert_awaited_once_with(test_user_id, SESSION_ID, message_id=MESSAGE_ROW_ID)
|
||||
|
||||
|
||||
def test_langfuse_failure_keeps_the_rating(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
upsert: AsyncMock,
|
||||
score_create: AsyncMock,
|
||||
) -> None:
|
||||
"""Postgres already holds the rating, so a Langfuse outage is logged and
|
||||
reported but does not fail the click."""
|
||||
_mock_message(mocker, _reply(trace_id=TRACE_ID))
|
||||
score_create.side_effect = RuntimeError("langfuse is down")
|
||||
|
||||
response = _post(_thumbs_down())
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"id": FEEDBACK_ID, "langfuse_target": None}
|
||||
upsert.assert_awaited_once()
|
||||
|
||||
|
||||
def test_unconfigured_langfuse_saves_the_rating_only(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
upsert: AsyncMock,
|
||||
) -> None:
|
||||
"""Self-hosted installs without Langfuse still record ratings."""
|
||||
_mock_message(mocker, _reply(trace_id=TRACE_ID))
|
||||
mocker.patch.object(feedback_service.settings.secrets, "langfuse_public_key", "")
|
||||
get_client = mocker.patch.object(feedback_service, "get_client")
|
||||
|
||||
response = _post(_thumbs_down())
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"id": FEEDBACK_ID, "langfuse_target": None}
|
||||
upsert.assert_awaited_once()
|
||||
get_client.assert_not_called()
|
||||
|
||||
|
||||
def test_database_failure_is_an_error_the_ui_can_show(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
score_create: AsyncMock,
|
||||
) -> None:
|
||||
"""With the rating unsaved the request fails rather than reporting a
|
||||
success the UI would thank the user for."""
|
||||
_mock_message(mocker, _reply(trace_id=TRACE_ID))
|
||||
_mock_message(mocker, _reply())
|
||||
mocker.patch.object(
|
||||
feedback_db,
|
||||
"upsert_message_feedback",
|
||||
@@ -320,7 +211,6 @@ def test_database_failure_is_an_error_the_ui_can_show(
|
||||
)
|
||||
|
||||
assert response.status_code == 500
|
||||
score_create.assert_not_awaited()
|
||||
|
||||
|
||||
# ─── Validation ────────────────────────────────────────────────────────
|
||||
@@ -352,7 +242,7 @@ def test_invalid_ratings_are_rejected(
|
||||
overrides: dict,
|
||||
upsert: AsyncMock,
|
||||
) -> None:
|
||||
lookup = _mock_message(mocker, _reply(trace_id=TRACE_ID))
|
||||
lookup = _mock_message(mocker, _reply())
|
||||
|
||||
response = _post({**_thumbs_down(), **overrides})
|
||||
|
||||
@@ -365,10 +255,9 @@ def test_rating_is_sent_to_posthog_as_the_user(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
test_user_id: str,
|
||||
upsert: AsyncMock,
|
||||
score_create: AsyncMock,
|
||||
posthog_capture: MagicMock,
|
||||
) -> None:
|
||||
_mock_message(mocker, _reply(trace_id=TRACE_ID))
|
||||
_mock_message(mocker, _reply())
|
||||
|
||||
response = _post(_thumbs_down())
|
||||
|
||||
@@ -377,7 +266,6 @@ def test_rating_is_sent_to_posthog_as_the_user(
|
||||
assert distinct_id == test_user_id
|
||||
assert event == "$ai_metric"
|
||||
assert properties["impersonated"] is False
|
||||
score_create.assert_awaited_once()
|
||||
|
||||
|
||||
def test_rating_by_an_impersonating_admin_is_marked(
|
||||
@@ -385,11 +273,10 @@ def test_rating_by_an_impersonating_admin_is_marked(
|
||||
mock_jwt_admin,
|
||||
target_user_id: str,
|
||||
upsert: AsyncMock,
|
||||
score_create: AsyncMock,
|
||||
posthog_capture: MagicMock,
|
||||
) -> None:
|
||||
app.dependency_overrides[get_jwt_payload] = mock_jwt_admin["get_jwt_payload"]
|
||||
_mock_message(mocker, _reply(trace_id=TRACE_ID))
|
||||
_mock_message(mocker, _reply())
|
||||
|
||||
response = client.post(
|
||||
f"/sessions/{SESSION_ID}/feedback",
|
||||
@@ -406,9 +293,8 @@ def test_rating_by_an_impersonating_admin_is_marked(
|
||||
def test_blank_comment_is_saved_as_none(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
upsert: AsyncMock,
|
||||
score_create: AsyncMock,
|
||||
) -> None:
|
||||
_mock_message(mocker, _reply(trace_id=TRACE_ID))
|
||||
_mock_message(mocker, _reply())
|
||||
|
||||
response = _post(_thumbs_down(comment=" "))
|
||||
|
||||
|
||||
@@ -1467,7 +1467,7 @@ async def get_copilot_usage(
|
||||
|
||||
Returns the percentage of the daily/weekly allowance used — not the
|
||||
raw spend or cap — so clients cannot derive per-turn cost or platform
|
||||
margins. Global defaults sourced from LaunchDarkly (falling back to
|
||||
margins. Global defaults sourced from PostHog (falling back to
|
||||
config). Includes the user's rate-limit tier.
|
||||
"""
|
||||
daily_limit, weekly_limit, tier = await get_global_rate_limits(
|
||||
|
||||
@@ -1545,7 +1545,7 @@ def test_stream_chat_returns_503_with_retry_after_when_rate_limit_unavailable(
|
||||
mocker.patch.object(chat_routes.config, "daily_cost_limit_microdollars", 10000)
|
||||
mocker.patch.object(chat_routes.config, "weekly_cost_limit_microdollars", 50000)
|
||||
# Patch the limit-resolution helper so the test does not exercise the
|
||||
# real LaunchDarkly / tier-lookup path — keeps the assertion focused on
|
||||
# real PostHog / tier-lookup path — keeps the assertion focused on
|
||||
# the RateLimitUnavailable → 503 mapping.
|
||||
mocker.patch(
|
||||
"backend.api.features.chat.routes.get_global_rate_limits",
|
||||
@@ -1611,7 +1611,7 @@ def _mock_usage(
|
||||
|
||||
Mocks both ``get_global_rate_limits`` (returns the given limits + tier) and
|
||||
``get_usage_status`` so that tests exercise the endpoint without hitting
|
||||
LaunchDarkly or Prisma.
|
||||
PostHog or Prisma.
|
||||
"""
|
||||
from backend.copilot.rate_limit import CoPilotUsageStatus, UsageWindow
|
||||
|
||||
|
||||
@@ -133,11 +133,11 @@ class TestEnableChatSharing:
|
||||
# backend can auto-link existing + future runs.
|
||||
assert captured["auto_share_executions"] is True
|
||||
|
||||
def test_enables_share_without_launchdarkly(self, client):
|
||||
def test_enables_share_without_posthog(self, client):
|
||||
with (
|
||||
patch(
|
||||
"backend.util.feature_flag.get_client",
|
||||
side_effect=RuntimeError("LaunchDarkly unavailable"),
|
||||
"backend.util.feature_flag.posthog.get_flag_client",
|
||||
side_effect=RuntimeError("PostHog unavailable"),
|
||||
),
|
||||
patch(
|
||||
"backend.api.features.chat.share.share_db.enable_chat_session_share",
|
||||
|
||||
@@ -1,23 +1,19 @@
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
from typing import Any, ClassVar
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import fastapi
|
||||
import httpx
|
||||
import pytest
|
||||
from autogpt_libs.auth import get_optional_user_id
|
||||
from ldclient import Context, LDClient
|
||||
from ldclient.config import Config as LDConfig
|
||||
from ldclient.integrations.test_data import TestData
|
||||
from ldclient.interfaces import DataSourceState, EventProcessor, UpdateProcessor
|
||||
from posthog import Posthog
|
||||
from posthog.types import FeatureFlag
|
||||
|
||||
import backend.util.feature_flag as ff
|
||||
from backend.util.feature_flag import (
|
||||
Flag,
|
||||
_UserContext,
|
||||
frontend,
|
||||
get_feature_flag_value,
|
||||
is_feature_enabled,
|
||||
@@ -28,7 +24,6 @@ from backend.util.feature_flag.frontend import (
|
||||
_anonymous_flags,
|
||||
_defined_flags,
|
||||
)
|
||||
from backend.util.settings import FeatureFlagBackend
|
||||
|
||||
from .routes import router
|
||||
|
||||
@@ -46,64 +41,26 @@ app.include_router(router, prefix="/api")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def launchdarkly(mocker, monkeypatch: pytest.MonkeyPatch):
|
||||
"""A real LaunchDarkly client serving fixed targeting, and no env overrides.
|
||||
|
||||
Yields every analytics event the client would have sent LaunchDarkly.
|
||||
"""
|
||||
def posthog_flags(mocker, monkeypatch: pytest.MonkeyPatch):
|
||||
"""A real PostHog client serving fixed targeting from local definitions,
|
||||
with no network and no env overrides."""
|
||||
for name in list(os.environ):
|
||||
if name.startswith(
|
||||
("FORCE_FLAG_", "NEXT_PUBLIC_FORCE_FLAG_", "FORCE_ALL_FLAGS")
|
||||
):
|
||||
monkeypatch.delenv(name)
|
||||
td = TestData.data_source()
|
||||
td.update(
|
||||
td.flag(FrontendFlag.COPILOT_VOICE_MODE.value)
|
||||
.boolean_flag()
|
||||
.variation_for_user("u-staff", True)
|
||||
.fallthrough_variation(False)
|
||||
)
|
||||
td.update(
|
||||
td.flag(FrontendFlag.GENERIC_TRIGGER_AGENTS.value)
|
||||
.boolean_flag()
|
||||
.variation_for_all(True)
|
||||
)
|
||||
td.update(
|
||||
td.flag(FrontendFlag.MARKETPLACE_SEARCH_TERMS.value)
|
||||
.variations(SEARCH_TERMS, VISITOR_TERMS)
|
||||
.variation_for_user(f"visitor:{VISITOR}", 1)
|
||||
.fallthrough_variation(0)
|
||||
)
|
||||
td.update(
|
||||
td.flag(PINNING)
|
||||
.boolean_flag()
|
||||
.variation_for_user(f"visitor:{VISITOR}", True)
|
||||
.variation_for_user(f"visitor:{'v' * 65}", True)
|
||||
.variation_for_user("visitor:visitor:bad$id", True)
|
||||
.variation_for_user(USER_ID, True)
|
||||
.if_match("email_domain", "agpt.co")
|
||||
.then_return(True)
|
||||
.fallthrough_variation(False)
|
||||
)
|
||||
_RecordedEvents.sent.clear()
|
||||
client = LDClient(
|
||||
LDConfig(
|
||||
"sdk-test", update_processor_class=td, event_processor_class=_RecordedEvents
|
||||
)
|
||||
)
|
||||
mocker.patch.object(
|
||||
ff.settings.config, "feature_flag_backend", FeatureFlagBackend.LAUNCHDARKLY
|
||||
)
|
||||
mocker.patch.object(ff.settings.secrets, "launch_darkly_sdk_key", "sdk-test")
|
||||
monkeypatch.setattr(ff, "_init_attempted", True)
|
||||
mocker.patch("backend.util.feature_flag.ldclient.get", return_value=client)
|
||||
mocker.patch.object(ph.settings.secrets, "posthog_api_key", "phc-test")
|
||||
client = _posthog_with_fixed_targeting()
|
||||
mocker.patch.object(client, "_get_flags_decision", side_effect=ConnectionError)
|
||||
mocker.patch.object(client, "capture")
|
||||
mocker.patch.object(ph, "get_flag_client", return_value=client)
|
||||
_anonymous_flags.cache_clear()
|
||||
_defined_flags.cache_clear()
|
||||
yield _RecordedEvents.sent
|
||||
yield client
|
||||
_anonymous_flags.cache_clear()
|
||||
_defined_flags.cache_clear()
|
||||
app.dependency_overrides.clear()
|
||||
client.close()
|
||||
client.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("caller", ["u-staff", "u-other", None])
|
||||
@@ -128,7 +85,12 @@ async def test_retired_rollout_flags_are_not_served_or_restored_by_overrides(
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.util.feature_flag._fetch_user_context_status",
|
||||
AsyncMock(return_value=(Context.builder(caller or "anonymous").build(), True)),
|
||||
AsyncMock(
|
||||
return_value=(
|
||||
_UserContext(key=caller or "anonymous", anonymous=caller is None),
|
||||
True,
|
||||
)
|
||||
),
|
||||
)
|
||||
retired = {
|
||||
"artifacts-page",
|
||||
@@ -165,9 +127,7 @@ async def test_a_defined_flag_the_vendor_could_not_answer_is_left_out(mocker):
|
||||
# A failed user lookup makes every answer a guess about someone else.
|
||||
mocker.patch(
|
||||
"backend.util.feature_flag._fetch_user_context_status",
|
||||
AsyncMock(
|
||||
return_value=(Context.builder(USER_ID).anonymous(True).build(), False)
|
||||
),
|
||||
AsyncMock(return_value=(_UserContext(key=USER_ID), False)),
|
||||
)
|
||||
|
||||
flags = (await _get_flags(USER_ID))["flags"]
|
||||
@@ -181,24 +141,24 @@ async def test_a_defined_flag_the_vendor_could_not_answer_is_left_out(mocker):
|
||||
|
||||
@pytest.mark.parametrize("ready_first", [False, True])
|
||||
async def test_a_signed_out_answer_is_kept_long_only_when_complete(
|
||||
mocker, monkeypatch, ready_first
|
||||
mocker, monkeypatch, posthog_flags, ready_first
|
||||
):
|
||||
"""A vendor that is down or starting is asked again within seconds, while a
|
||||
complete answer keeps serving every signed-out visitor from one read."""
|
||||
monkeypatch.setattr(frontend, "PARTIAL_TTL_SECONDS", 0.01)
|
||||
ready = ff.ldclient.get()
|
||||
not_ready = LDClient(LDConfig("sdk-test", offline=True))
|
||||
get = mocker.patch(
|
||||
"backend.util.feature_flag.ldclient.get",
|
||||
return_value=ready if ready_first else not_ready,
|
||||
not_ready = _posthog_still_loading(mocker)
|
||||
get = mocker.patch.object(
|
||||
ph,
|
||||
"get_flag_client",
|
||||
return_value=posthog_flags if ready_first else not_ready,
|
||||
)
|
||||
try:
|
||||
first = (await _get_flags(None))["flags"]
|
||||
get.return_value = not_ready if ready_first else ready
|
||||
get.return_value = not_ready if ready_first else posthog_flags
|
||||
await asyncio.sleep(0.05)
|
||||
second = (await _get_flags(None))["flags"]
|
||||
finally:
|
||||
not_ready.close()
|
||||
not_ready.shutdown()
|
||||
|
||||
# Either way the second read sees the first, complete answer.
|
||||
assert (GENERIC_TRIGGER_AGENTS in first, second.get(GENERIC_TRIGGER_AGENTS)) == (
|
||||
@@ -207,51 +167,23 @@ async def test_a_signed_out_answer_is_kept_long_only_when_complete(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data_source,body",
|
||||
[
|
||||
# A rejected key ends LaunchDarkly's attempts; the route gates then serve
|
||||
# defaults, and the page must too rather than poll forever.
|
||||
("_RejectedKey", {"flags": {}, "configured": False}),
|
||||
# Still connecting past the start-up wait: the page keeps asking.
|
||||
("_StillConnecting", {"flags": {}, "configured": True}),
|
||||
],
|
||||
)
|
||||
async def test_a_vendor_that_cannot_answer_says_whether_it_ever_will(
|
||||
mocker, data_source, body
|
||||
async def test_a_vendor_still_loading_leaves_every_flag_out_and_says_it_will_answer(
|
||||
mocker,
|
||||
):
|
||||
client = LDClient(
|
||||
LDConfig(
|
||||
"sdk-test", update_processor_class=globals()[data_source], send_events=False
|
||||
),
|
||||
start_wait=0,
|
||||
)
|
||||
mocker.patch("backend.util.feature_flag.ldclient.get", return_value=client)
|
||||
# Definitions not loaded and /flags unreachable, as while PostHog starts.
|
||||
client = _posthog_still_loading(mocker)
|
||||
mocker.patch.object(ph, "get_flag_client", return_value=client)
|
||||
try:
|
||||
answer = await _get_flags("u-staff")
|
||||
finally:
|
||||
client.close()
|
||||
client.shutdown()
|
||||
|
||||
assert answer == body
|
||||
|
||||
|
||||
async def test_a_launchdarkly_client_not_ready_yet_leaves_every_flag_out(mocker):
|
||||
# Offline, the client answers every read as CLIENT_NOT_READY, as while it starts.
|
||||
client = LDClient(LDConfig("sdk-test", offline=True))
|
||||
mocker.patch("backend.util.feature_flag.ldclient.get", return_value=client)
|
||||
try:
|
||||
flags = (await _get_flags("u-staff"))["flags"]
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
assert flags == {}
|
||||
# PostHog keeps retrying, so the page keeps asking rather than settling on defaults.
|
||||
assert answer == {"flags": {}, "configured": True}
|
||||
|
||||
|
||||
async def test_a_signed_in_read_asks_posthog_nothing_about_flags_it_lacks(mocker):
|
||||
"""Each would be a remote /flags call on the executor every flag read shares."""
|
||||
mocker.patch.object(
|
||||
ff.settings.config, "feature_flag_backend", FeatureFlagBackend.POSTHOG
|
||||
)
|
||||
client = _posthog_splitting_pinning_and_gating_hire_on_email()
|
||||
remote = mocker.patch.object(
|
||||
client, "_get_flags_decision", side_effect=ConnectionError
|
||||
@@ -271,9 +203,6 @@ async def test_a_signed_in_read_asks_posthog_nothing_about_flags_it_lacks(mocker
|
||||
async def test_on_posthog_only_loaded_definitions_can_rule_a_flag_undefined(
|
||||
mocker, loaded
|
||||
):
|
||||
mocker.patch.object(
|
||||
ff.settings.config, "feature_flag_backend", FeatureFlagBackend.POSTHOG
|
||||
)
|
||||
client = _posthog_splitting_pinning_and_gating_hire_on_email()
|
||||
if not loaded:
|
||||
client.feature_flags = None
|
||||
@@ -293,23 +222,10 @@ async def test_on_posthog_only_loaded_definitions_can_rule_a_flag_undefined(
|
||||
) == expected
|
||||
|
||||
|
||||
async def test_with_no_vendor_configured_every_flag_is_the_frontends_default(mocker):
|
||||
mocker.patch.object(ff.settings.secrets, "launch_darkly_sdk_key", "")
|
||||
|
||||
body = await _get_flags("u-staff")
|
||||
|
||||
assert body["configured"] is False
|
||||
|
||||
|
||||
async def test_posthog_with_no_key_leaves_every_flag_to_the_frontend(
|
||||
mocker, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
mocker.patch.object(
|
||||
ff.settings.config, "feature_flag_backend", FeatureFlagBackend.POSTHOG
|
||||
)
|
||||
async def test_posthog_with_no_key_leaves_every_flag_to_the_frontend(mocker):
|
||||
mocker.patch.object(ph.settings.secrets, "posthog_api_key", "")
|
||||
monkeypatch.setattr(ph, "_client", None)
|
||||
monkeypatch.setattr(ph, "_init_attempted", False)
|
||||
# What the real client getter answers without a key.
|
||||
mocker.patch.object(ph, "get_flag_client", return_value=None)
|
||||
|
||||
assert await _get_flags(USER_ID) == {"flags": {}, "configured": False}
|
||||
|
||||
@@ -357,11 +273,11 @@ async def test_a_malformed_visitor_id_counts_as_none(malformed):
|
||||
async def test_a_users_id_sent_as_a_visitor_id_is_never_looked_up_as_that_user(
|
||||
mocker,
|
||||
):
|
||||
# The lookup would hand back an agpt.co user, whom the flag's rule turns on.
|
||||
employee = Context.builder("u-staff").set("email_domain", "agpt.co").build()
|
||||
# The lookup would hand back the user, whom the flag targets.
|
||||
user = _UserContext(key=USER_ID, anonymous=False)
|
||||
lookup = mocker.patch(
|
||||
"backend.util.feature_flag._fetch_user_context_status",
|
||||
AsyncMock(return_value=(employee, True)),
|
||||
AsyncMock(return_value=(user, True)),
|
||||
)
|
||||
|
||||
flags = (await _get_flags(None, USER_ID))["flags"]
|
||||
@@ -375,7 +291,7 @@ async def test_a_users_id_sent_as_a_visitor_id_never_matches_that_users_targets(
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.util.feature_flag._fetch_user_context_status",
|
||||
AsyncMock(return_value=(Context.builder(USER_ID).build(), True)),
|
||||
AsyncMock(return_value=(_UserContext(key=USER_ID, anonymous=False), True)),
|
||||
)
|
||||
|
||||
as_user = (await _get_flags(USER_ID))["flags"]
|
||||
@@ -384,22 +300,9 @@ async def test_a_users_id_sent_as_a_visitor_id_never_matches_that_users_targets(
|
||||
assert (as_user[PINNING], as_visitor[PINNING]) == (True, False)
|
||||
|
||||
|
||||
async def test_fresh_visitor_ids_send_launchdarkly_nothing(launchdarkly):
|
||||
await _get_flags(None)
|
||||
launchdarkly.clear()
|
||||
|
||||
for _ in range(20):
|
||||
await _get_flags(None, str(uuid.uuid4()))
|
||||
|
||||
assert launchdarkly == []
|
||||
|
||||
|
||||
async def test_fresh_visitor_ids_cost_posthog_nothing(mocker):
|
||||
"""No billed event, no remote /flags call and no turn on the executor every
|
||||
flag read shares: the id is the caller's to pick, so a new one must be free."""
|
||||
mocker.patch.object(
|
||||
ff.settings.config, "feature_flag_backend", FeatureFlagBackend.POSTHOG
|
||||
)
|
||||
client = _posthog_splitting_pinning_and_gating_hire_on_email()
|
||||
mocker.patch.object(ph, "get_flag_client", return_value=client)
|
||||
remote = mocker.patch.object(
|
||||
@@ -434,48 +337,89 @@ async def test_fresh_visitor_ids_cost_posthog_nothing(mocker):
|
||||
assert all(v.get(GENERIC_TRIGGER_AGENTS) is True for v in visitors)
|
||||
|
||||
|
||||
class _StillConnecting(UpdateProcessor):
|
||||
"""A data source that has not received flags yet."""
|
||||
|
||||
def __init__(self, config: LDConfig, store: Any, ready: Any):
|
||||
self.config = config
|
||||
|
||||
def start(self) -> None:
|
||||
pass
|
||||
|
||||
def stop(self) -> None:
|
||||
pass
|
||||
|
||||
def is_alive(self) -> bool:
|
||||
return True
|
||||
|
||||
def initialized(self) -> bool:
|
||||
return False
|
||||
def _posthog_with_fixed_targeting() -> Posthog:
|
||||
"""Local definitions every read decides in-process, without /flags."""
|
||||
client = Posthog("phc-test", secret_key="phx-test", enable_local_evaluation=False)
|
||||
client.feature_flags = [
|
||||
_targeted(1, FrontendFlag.COPILOT_VOICE_MODE.value, ["u-staff"]),
|
||||
{
|
||||
"id": 2,
|
||||
"key": GENERIC_TRIGGER_AGENTS,
|
||||
"active": True,
|
||||
"filters": {"groups": [{"properties": [], "rollout_percentage": 100}]},
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"key": FrontendFlag.MARKETPLACE_SEARCH_TERMS.value,
|
||||
"active": True,
|
||||
"filters": {
|
||||
"groups": [
|
||||
{
|
||||
"properties": [_distinct_id_in([f"visitor:{VISITOR}"])],
|
||||
"rollout_percentage": 100,
|
||||
"variant": "visitor",
|
||||
},
|
||||
{"properties": [], "rollout_percentage": 100, "variant": "all"},
|
||||
],
|
||||
"multivariate": {
|
||||
"variants": [
|
||||
{"key": "all", "rollout_percentage": 100},
|
||||
{"key": "visitor", "rollout_percentage": 0},
|
||||
]
|
||||
},
|
||||
"payloads": {
|
||||
"all": json.dumps(SEARCH_TERMS),
|
||||
"visitor": json.dumps(VISITOR_TERMS),
|
||||
},
|
||||
},
|
||||
},
|
||||
_targeted(
|
||||
4,
|
||||
PINNING,
|
||||
[
|
||||
f"visitor:{VISITOR}",
|
||||
f"visitor:{'v' * 65}",
|
||||
"visitor:visitor:bad$id",
|
||||
USER_ID,
|
||||
],
|
||||
),
|
||||
]
|
||||
return client
|
||||
|
||||
|
||||
class _RejectedKey(_StillConnecting):
|
||||
"""Gives up for good, as LaunchDarkly's own data sources do on a 401."""
|
||||
|
||||
def start(self) -> None:
|
||||
sink = self.config.data_source_update_sink
|
||||
assert sink is not None
|
||||
sink.update_status(DataSourceState.OFF, None)
|
||||
def _targeted(flag_id: int, key: str, distinct_ids: list[str]) -> dict:
|
||||
"""A boolean flag on for exactly these distinct ids, off for everyone else."""
|
||||
return {
|
||||
"id": flag_id,
|
||||
"key": key,
|
||||
"active": True,
|
||||
"filters": {
|
||||
"groups": [
|
||||
{
|
||||
"properties": [_distinct_id_in(distinct_ids)],
|
||||
"rollout_percentage": 100,
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class _RecordedEvents(EventProcessor):
|
||||
sent: ClassVar[list[Any]] = []
|
||||
def _distinct_id_in(distinct_ids: list[str]) -> dict:
|
||||
return {
|
||||
"key": "distinct_id",
|
||||
"type": "person",
|
||||
"operator": "exact",
|
||||
"value": distinct_ids,
|
||||
}
|
||||
|
||||
def __init__(self, config: LDConfig):
|
||||
pass
|
||||
|
||||
def send_event(self, event: Any) -> None:
|
||||
self.sent.append(event)
|
||||
|
||||
def flush(self) -> None:
|
||||
pass
|
||||
|
||||
def stop(self) -> None:
|
||||
pass
|
||||
def _posthog_still_loading(mocker) -> Posthog:
|
||||
"""No definitions yet and /flags unreachable: every read goes unanswered."""
|
||||
client = Posthog("phc-test", secret_key="phx-test", enable_local_evaluation=False)
|
||||
mocker.patch.object(client, "load_feature_flags")
|
||||
mocker.patch.object(client, "_get_flags_decision", side_effect=ConnectionError)
|
||||
mocker.patch.object(client, "capture")
|
||||
return client
|
||||
|
||||
|
||||
def _posthog_splitting_pinning_and_gating_hire_on_email() -> Posthog:
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
import json
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import fastapi
|
||||
@@ -1179,3 +1180,39 @@ def test_execute_graph_records_a_failure_and_re_raises(
|
||||
client.post("/graphs/graph-1/execute/1", json={})
|
||||
|
||||
assert record.call_args.kwargs["status"] == "error"
|
||||
|
||||
|
||||
def test_execute_graph_refuses_an_input_nested_too_deep(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""#15648: an input nested ~1000 levels deep answered 422 but still saved
|
||||
an INCOMPLETE run whose ``inputs`` column held a string, which broke every
|
||||
later listing of the user's runs. The real ``add_graph_execution`` must
|
||||
refuse it as a validation error before anything is written."""
|
||||
credit_model = Mock()
|
||||
credit_model.get_credits = AsyncMock(return_value=100)
|
||||
mocker.patch(
|
||||
"backend.api.features.graphs.routes.get_credit_model",
|
||||
AsyncMock(return_value=credit_model),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.executor.utils.is_user_paywalled", AsyncMock(return_value=False)
|
||||
)
|
||||
mocker.patch("backend.executor.utils.is_user_banned", AsyncMock(return_value=False))
|
||||
created = mocker.patch(
|
||||
"backend.executor.utils.execution_db.create_graph_execution", AsyncMock()
|
||||
)
|
||||
mocker.patch("backend.api.features.graphs.routes.record_graph_operation", Mock())
|
||||
value: Any = "abyss"
|
||||
for _ in range(300):
|
||||
value = [value]
|
||||
|
||||
response = client.post(
|
||||
"/graphs/graph-1/execute/1", json={"inputs": {"shift_tags": value}}
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
detail = response.json()["detail"]
|
||||
assert detail["type"] == "validation_error"
|
||||
assert "shift_tags" in detail["message"]
|
||||
created.assert_not_awaited()
|
||||
@@ -13,8 +13,6 @@ import asyncio
|
||||
import logging
|
||||
import os
|
||||
|
||||
from langfuse import get_client
|
||||
|
||||
from backend.api.features.onboarding_dump.models import SuggestedPrompt
|
||||
from backend.api.features.onboarding_dump.parsing import parse_response_json
|
||||
from backend.api.features.onboarding_dump.providers import (
|
||||
@@ -22,12 +20,13 @@ from backend.api.features.onboarding_dump.providers import (
|
||||
provider_lines,
|
||||
)
|
||||
from backend.util.clients import get_openai_client
|
||||
from backend.util.settings import Settings
|
||||
from backend.util.langfuse_prompts import (
|
||||
get_langfuse_prompt_client,
|
||||
langfuse_prompts_configured,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
settings = Settings()
|
||||
|
||||
_MODEL = os.environ.get("BRAIN_DUMP_GREETING_MODEL", "anthropic/claude-sonnet-5-5")
|
||||
_TIMEOUT_SECONDS = 45
|
||||
# The model thinks adaptively and the thinking counts against this: at 3,000
|
||||
@@ -250,13 +249,11 @@ async def _fetch_langfuse_prompt() -> str | None:
|
||||
unconfigured or the fetch fails — the caller falls back to
|
||||
``_LOCAL_PROMPT`` so a Langfuse outage never blocks onboarding.
|
||||
"""
|
||||
if not (
|
||||
settings.secrets.langfuse_public_key and settings.secrets.langfuse_secret_key
|
||||
):
|
||||
if not langfuse_prompts_configured():
|
||||
return None
|
||||
try:
|
||||
prompt = await asyncio.to_thread(
|
||||
get_client().get_prompt,
|
||||
get_langfuse_prompt_client().get_prompt,
|
||||
LANGFUSE_PROMPT_NAME,
|
||||
cache_ttl_seconds=LANGFUSE_PROMPT_CACHE_TTL_SECONDS,
|
||||
)
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
A skill is a directory, so upload and download carry it as an archive. The caps
|
||||
and the path rules live in :mod:`backend.copilot.tools.skills`; what is here is
|
||||
only what an archive can carry that a path check cannot see — a symlink, an
|
||||
encrypted member, a declared size that would unpack to gigabytes.
|
||||
encrypted member, a member past the size caps.
|
||||
|
||||
Both entry points are CPU-bound and synchronous; callers run them off the event
|
||||
loop.
|
||||
@@ -21,6 +21,7 @@ from backend.copilot.tools.skills import (
|
||||
SkillPackageError,
|
||||
validate_package,
|
||||
)
|
||||
from backend.util.zip_member import ZipMemberError, ZipMemberTooLarge, read_member
|
||||
|
||||
ROOT_SKILL_MD = "SKILL.md"
|
||||
# A compressed archive already past the uncompressed cap cannot hold a package
|
||||
@@ -32,8 +33,8 @@ def package_from_zip(data: bytes) -> SkillPackage:
|
||||
"""Unpack an uploaded archive into a validated :class:`SkillPackage`.
|
||||
|
||||
Every refusal is a :class:`SkillPackageError`, ``over_limit`` set where the
|
||||
REST edge owes a 413. Sizes are checked from the central directory, so a
|
||||
zip bomb is refused without a member being decompressed.
|
||||
REST edge owes a 413. Sizes are checked from the central directory before
|
||||
a member is read, and each member is read under ``MAX_PACKAGE_FILE_BYTES``.
|
||||
"""
|
||||
try:
|
||||
archive = zipfile.ZipFile(io.BytesIO(data))
|
||||
@@ -44,12 +45,27 @@ def package_from_zip(data: bytes) -> SkillPackage:
|
||||
root = members.pop(ROOT_SKILL_MD)
|
||||
# A member's CRC is verified as it decompresses, so a corrupt archive
|
||||
# opens cleanly and fails here.
|
||||
# The declared sizes are the archive's own claim, so the package total
|
||||
# is held against the bytes each read actually produced.
|
||||
unpacked = 0
|
||||
|
||||
def read_counted(info: zipfile.ZipInfo) -> bytes:
|
||||
nonlocal unpacked
|
||||
content = _read(archive, info)
|
||||
unpacked += len(content)
|
||||
if unpacked > MAX_PACKAGE_BYTES:
|
||||
raise SkillPackageError(
|
||||
f"archive unpacks to more than {MAX_PACKAGE_BYTES} bytes",
|
||||
over_limit=True,
|
||||
)
|
||||
return content
|
||||
|
||||
try:
|
||||
skill_md = _decoded(archive.read(root))
|
||||
skill_md = _decoded(read_counted(root))
|
||||
files = [
|
||||
SkillFile(
|
||||
relative_path=path,
|
||||
content=archive.read(info),
|
||||
content=read_counted(info),
|
||||
is_executable=bool(_mode(info) & 0o111),
|
||||
)
|
||||
for path, info in members.items()
|
||||
@@ -73,6 +89,24 @@ def zip_from_package(package: SkillPackage) -> bytes:
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def _read(archive: zipfile.ZipFile, info: zipfile.ZipInfo) -> bytes:
|
||||
"""One member's bytes, read under ``MAX_PACKAGE_FILE_BYTES``."""
|
||||
named = f"member '{info.filename[:120]}'"
|
||||
try:
|
||||
return read_member(archive, info, max_bytes=MAX_PACKAGE_FILE_BYTES)
|
||||
except ZipMemberTooLarge:
|
||||
raise SkillPackageError(
|
||||
f"{named} unpacks to more than the limit of "
|
||||
f"{MAX_PACKAGE_FILE_BYTES} bytes",
|
||||
over_limit=True,
|
||||
)
|
||||
except ZipMemberError:
|
||||
raise SkillPackageError(
|
||||
f"{named} uses an unsupported compression method; "
|
||||
"zip it again with deflate"
|
||||
)
|
||||
|
||||
|
||||
_ENCRYPTED_FLAG = 0x1
|
||||
_FILE_TYPE_MASK = 0o170000
|
||||
_SYMLINK_TYPE = 0o120000
|
||||
|
||||
@@ -162,8 +162,7 @@ def _with_encryption_flag(data: bytes) -> bytes:
|
||||
|
||||
|
||||
def test_a_bomb_is_refused_without_a_member_being_read(monkeypatch):
|
||||
"""The declared sizes are the whole check: 22 MiB of zeros compresses to a
|
||||
few KB, so reading first is how a bomb lands in memory."""
|
||||
"""22 MiB of zeros compresses to a few KB, so it is refused unread."""
|
||||
member = b"\0" * MAX_PACKAGE_FILE_BYTES
|
||||
data = _zip({"SKILL.md": SKILL_MD, **{f"f{i}.bin": member for i in range(11)}})
|
||||
assert len(data) < MAX_PACKAGE_BYTES
|
||||
@@ -177,6 +176,28 @@ def test_a_bomb_is_refused_without_a_member_being_read(monkeypatch):
|
||||
opened.assert_not_called()
|
||||
|
||||
|
||||
def test_reading_stops_once_the_real_bytes_pass_the_total(monkeypatch):
|
||||
"""The declared sizes are small, so the check before reading passes; the
|
||||
bytes the reader returns are what the package total is held against, and
|
||||
the rest of the members are never decompressed."""
|
||||
count = 30
|
||||
data = _zip({"SKILL.md": SKILL_MD, **{f"f{i}.bin": b"x" for i in range(count)}})
|
||||
produced = b"\0" * MAX_PACKAGE_FILE_BYTES
|
||||
reads: list[str] = []
|
||||
|
||||
def read_member(archive, info, *, max_bytes):
|
||||
reads.append(info.filename)
|
||||
return produced
|
||||
|
||||
monkeypatch.setattr("backend.api.features.skill_zip.read_member", read_member)
|
||||
|
||||
with pytest.raises(SkillPackageError) as exc:
|
||||
package_from_zip(data)
|
||||
assert exc.value.over_limit
|
||||
assert str(MAX_PACKAGE_BYTES) in str(exc.value)
|
||||
assert len(reads) == MAX_PACKAGE_BYTES // MAX_PACKAGE_FILE_BYTES + 1
|
||||
|
||||
|
||||
def test_a_file_over_the_per_file_cap_is_refused_unread(monkeypatch):
|
||||
"""The package validator would also catch the size — but only after the
|
||||
member has been decompressed, which is the cost this check avoids."""
|
||||
@@ -294,6 +315,39 @@ def _with_a_corrupt_member(data: bytes) -> bytes:
|
||||
return bytes(out)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", [zipfile.ZIP_BZIP2, zipfile.ZIP_LZMA])
|
||||
def test_a_bzip2_or_lzma_member_is_refused(method: int, monkeypatch):
|
||||
buffer = io.BytesIO()
|
||||
with zipfile.ZipFile(buffer, "w") as archive:
|
||||
archive.writestr("SKILL.md", SKILL_MD, zipfile.ZIP_DEFLATED)
|
||||
archive.writestr("a.bin", b"x", method)
|
||||
|
||||
decompressor = zipfile._get_decompressor
|
||||
|
||||
def deflate_only(compress_type: int):
|
||||
assert compress_type == zipfile.ZIP_DEFLATED, "a.bin was decompressed"
|
||||
return decompressor(compress_type)
|
||||
|
||||
monkeypatch.setattr(zipfile, "_get_decompressor", deflate_only)
|
||||
|
||||
with pytest.raises(SkillPackageError, match="unsupported compression method"):
|
||||
package_from_zip(buffer.getvalue())
|
||||
|
||||
|
||||
def test_a_stored_package_is_read():
|
||||
buffer = io.BytesIO()
|
||||
with zipfile.ZipFile(buffer, "w", zipfile.ZIP_STORED) as archive:
|
||||
archive.writestr("SKILL.md", SKILL_MD)
|
||||
archive.writestr("references/API.md", "# API\n")
|
||||
|
||||
package = package_from_zip(buffer.getvalue())
|
||||
|
||||
assert package.skill_md == SKILL_MD
|
||||
assert [(f.relative_path, f.content) for f in package.files] == [
|
||||
("references/API.md", b"# API\n")
|
||||
]
|
||||
|
||||
|
||||
def test_something_that_is_not_a_zip_is_refused():
|
||||
with pytest.raises(SkillPackageError, match="not a readable zip"):
|
||||
package_from_zip(b"---\nname: demo\n---\n")
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Tests for subscription tier API endpoints."""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import fastapi
|
||||
@@ -10,8 +11,10 @@ import stripe
|
||||
from autogpt_libs.auth.jwt_utils import get_jwt_payload
|
||||
from prisma.enums import SubscriptionTier
|
||||
|
||||
from backend.data.credit import PendingChangeUnknown
|
||||
from backend.data.pro_activation import UsageActivationState
|
||||
from backend.executor.scheduler import CopilotTurnJobInfo, GraphExecutionJobInfo
|
||||
from backend.util.testing import Rendezvous
|
||||
|
||||
from .billing.credits_rate_limit import enforce_subscription_status_rate_limit
|
||||
from .billing.subscriptions import routes as subscription_routes
|
||||
@@ -187,7 +190,7 @@ def _stub_subscription_status_lookups(mocker: pytest_mock.MockFixture) -> None:
|
||||
return_value=0,
|
||||
)
|
||||
# Default tier-multiplier resolver to the backend defaults so the endpoint
|
||||
# never reaches LaunchDarkly during tests. Individual tests override for
|
||||
# never reaches PostHog during tests. Individual tests override for
|
||||
# LD-override scenarios. get_tier_multipliers returns a string-keyed dict
|
||||
# (per its docstring) so we mirror that shape here — passing the
|
||||
# enum-keyed _DEFAULT_TIER_MULTIPLIERS directly would silently mismatch
|
||||
@@ -391,11 +394,73 @@ def test_get_subscription_status_pro(
|
||||
assert data["tier_multipliers"]["BUSINESS"] == 60.0
|
||||
|
||||
|
||||
def test_subscription_status_reads_its_independent_parts_side_by_side(
|
||||
client: fastapi.testclient.TestClient,
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""Multipliers, billing cycle, period end and the pending change each wait
|
||||
for none of the others; only the proration needs the cycle's cost."""
|
||||
meeting = Rendezvous(4)
|
||||
mocker.patch.object(
|
||||
subscription_routes,
|
||||
"get_user_by_id",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(subscription_tier=SubscriptionTier.PRO),
|
||||
)
|
||||
amounts = {"price_pro": 1999, "price_pro_yearly": 19990, "price_max": 4999}
|
||||
|
||||
async def price_amount(price_id: str) -> int:
|
||||
return amounts.get(price_id, 0)
|
||||
|
||||
mocker.patch.object(
|
||||
subscription_routes, "_get_stripe_price_amount", side_effect=price_amount
|
||||
)
|
||||
mocker.patch.object(
|
||||
subscription_routes,
|
||||
"get_tier_multipliers",
|
||||
side_effect=meeting.returning({"PRO": 2.0, "MAX": 5.0}),
|
||||
)
|
||||
mocker.patch.object(
|
||||
subscription_routes,
|
||||
"get_user_billing_cycle",
|
||||
side_effect=meeting.returning("yearly"),
|
||||
)
|
||||
mocker.patch.object(
|
||||
subscription_routes,
|
||||
"get_active_subscription_period_end",
|
||||
side_effect=meeting.returning(1_900_000_000),
|
||||
)
|
||||
pending_at = datetime(2030, 1, 1, tzinfo=UTC)
|
||||
mocker.patch.object(
|
||||
subscription_routes,
|
||||
"get_pending_subscription_change",
|
||||
side_effect=meeting.returning((SubscriptionTier.MAX, pending_at, "monthly")),
|
||||
)
|
||||
proration = mocker.patch.object(
|
||||
subscription_routes,
|
||||
"get_proration_credit_cents",
|
||||
new_callable=AsyncMock,
|
||||
return_value=321,
|
||||
)
|
||||
|
||||
response = client.get("/credits/subscription")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["billing_cycle"] == "yearly"
|
||||
assert data["monthly_cost"] == 19990
|
||||
assert data["tier_multipliers"]["PRO"] == 2.0
|
||||
assert data["current_period_end"] == 1_900_000_000
|
||||
assert data["pending_tier"] == "MAX"
|
||||
assert data["proration_credit_cents"] == 321
|
||||
proration.assert_awaited_once_with(TEST_USER_ID, 19990)
|
||||
|
||||
|
||||
def test_get_subscription_status_tier_multipliers_ld_override(
|
||||
client: fastapi.testclient.TestClient,
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""A LaunchDarkly-overridden tier multiplier flows through the response."""
|
||||
"""A PostHog-overridden tier multiplier flows through the response."""
|
||||
mock_user = Mock()
|
||||
mock_user.subscription_tier = SubscriptionTier.BASIC
|
||||
|
||||
@@ -2174,6 +2239,47 @@ def test_get_subscription_status_no_pending_tier(
|
||||
assert data["pending_tier_effective_at"] is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"error",
|
||||
[stripe.StripeError("rate limited"), PendingChangeUnknown("no price ids")],
|
||||
ids=["stripe", "price-lookup"],
|
||||
)
|
||||
def test_get_subscription_status_reads_a_failed_pending_change_as_none(
|
||||
client: fastapi.testclient.TestClient,
|
||||
mocker: pytest_mock.MockFixture,
|
||||
error: Exception,
|
||||
) -> None:
|
||||
"""A pending change that can't be resolved leaves the status without one
|
||||
instead of failing it, and the failure is logged."""
|
||||
mocker.patch(
|
||||
"backend.api.features.billing.subscriptions.routes.get_user_by_id",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(subscription_tier=SubscriptionTier.PRO),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.billing.subscriptions.routes.get_subscription_price_id",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.billing.subscriptions.routes.get_proration_credit_cents",
|
||||
new_callable=AsyncMock,
|
||||
return_value=0,
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.billing.subscriptions.routes.get_pending_subscription_change",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=error,
|
||||
)
|
||||
logged = mocker.spy(subscription_routes.logger, "exception")
|
||||
|
||||
response = client.get("/credits/subscription")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["pending_tier"] is None
|
||||
logged.assert_called_once()
|
||||
|
||||
|
||||
def test_update_subscription_tier_downgrade_paid_to_paid_schedules(
|
||||
client: fastapi.testclient.TestClient,
|
||||
mocker: pytest_mock.MockFixture,
|
||||
|
||||
+1
-1
@@ -282,7 +282,7 @@ def test_update_subscription_tier_cancel_pending_trial_converts_without_a_curren
|
||||
trial_conversion: MagicMock,
|
||||
) -> None:
|
||||
"""The conversion bills the price the trial accepted, so it needs no
|
||||
current LaunchDarkly price for the plan."""
|
||||
current PostHog price for the plan."""
|
||||
mocker.patch(
|
||||
"backend.api.features.billing.subscriptions.routes.get_subscription_price_id",
|
||||
new_callable=AsyncMock,
|
||||
|
||||
@@ -29,6 +29,8 @@ from typing import Optional, cast
|
||||
import pypdfium2
|
||||
from PIL import Image
|
||||
|
||||
from backend.util.zip_member import ZipMemberTooLarge, read_member
|
||||
|
||||
if sys.platform != "win32":
|
||||
import resource
|
||||
|
||||
@@ -39,6 +41,8 @@ RENDER_MAX_MEMORY_BYTES = 768 * 1024 * 1024
|
||||
# A thumbnail is at most 1024 px wide; anything bigger did not come from a
|
||||
# renderer here.
|
||||
MAX_THUMBNAIL_BYTES = 16 * 1024 * 1024
|
||||
# The embedded thumbnail of an Office file: the same cap as an image preview.
|
||||
MAX_EMBEDDED_THUMBNAIL_BYTES = 10_000_000
|
||||
WEBP_QUALITY = 80
|
||||
# A preview is a small thumbnail, so rasterising a source larger than this spends
|
||||
# memory and CPU for no visible gain. Every renderer stays within this budget:
|
||||
@@ -344,7 +348,13 @@ def _extract_zip_thumbnail(content: bytes) -> Optional[bytes]:
|
||||
names = {name.lower(): name for name in archive.namelist()}
|
||||
for candidate in _EMBEDDED_THUMBNAILS:
|
||||
if candidate in names:
|
||||
return archive.read(names[candidate])
|
||||
info = archive.getinfo(names[candidate])
|
||||
try:
|
||||
return read_member(
|
||||
archive, info, max_bytes=MAX_EMBEDDED_THUMBNAIL_BYTES
|
||||
)
|
||||
except ZipMemberTooLarge as e:
|
||||
raise PreviewTooLarge() from e
|
||||
return None
|
||||
|
||||
|
||||
|
||||
@@ -443,6 +443,53 @@ def test_office_preview_declines_an_oversized_thumbnail(mocker):
|
||||
assert response.status_code == 413
|
||||
|
||||
|
||||
# The embedded thumbnail cap: 10 MB, the same as an image preview.
|
||||
_THUMBNAIL_CAP = 10_000_000
|
||||
|
||||
|
||||
def _office_with_thumbnail(thumbnail_bytes: int, method: int) -> bytes:
|
||||
buf = io.BytesIO()
|
||||
with zipfile.ZipFile(buf, "w", method) as archive:
|
||||
with archive.open("docProps/thumbnail.jpeg", "w") as member:
|
||||
# Written in pieces so the test never holds the whole member.
|
||||
for start in range(0, thumbnail_bytes, 1_000_000):
|
||||
member.write(b"\0" * min(1_000_000, thumbnail_bytes - start))
|
||||
archive.writestr("word/document.xml", "<w/>")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def test_office_thumbnail_unpacking_past_the_cap_is_refused():
|
||||
content = _office_with_thumbnail(_THUMBNAIL_CAP + 1, zipfile.ZIP_DEFLATED)
|
||||
assert len(content) < 100_000
|
||||
|
||||
with pytest.raises(preview_render.PreviewTooLarge):
|
||||
preview_render._extract_zip_thumbnail(content)
|
||||
|
||||
|
||||
def test_office_thumbnail_at_the_cap_is_read():
|
||||
content = _office_with_thumbnail(_THUMBNAIL_CAP, zipfile.ZIP_DEFLATED)
|
||||
|
||||
thumbnail = preview_render._extract_zip_thumbnail(content)
|
||||
assert thumbnail is not None and len(thumbnail) == _THUMBNAIL_CAP
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"method, status", [(zipfile.ZIP_DEFLATED, 413), (zipfile.ZIP_BZIP2, 415)]
|
||||
)
|
||||
def test_office_preview_declines_a_thumbnail_that_unpacks_past_the_cap(
|
||||
mocker, method, status
|
||||
):
|
||||
# Rendered in the child: an unbounded read would hand 10 MB of zeros to the
|
||||
# image decoder and fail as a 415, so a 413 means the cap stopped the read.
|
||||
_mock_lookups(mocker, _make_file(name="report.docx", path="/report.docx"))
|
||||
_mock_storage(mocker, _office_with_thumbnail(_THUMBNAIL_CAP + 1, method))
|
||||
_mock_redis(mocker)
|
||||
|
||||
response = client.get("/files/file-001/preview")
|
||||
|
||||
assert response.status_code == status
|
||||
|
||||
|
||||
def _mock_render_io(mocker, content: bytes) -> None:
|
||||
storage = mocker.MagicMock()
|
||||
storage.retrieve = AsyncMock(return_value=content)
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
import json
|
||||
from decimal import Decimal, InvalidOperation
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
@@ -97,7 +99,7 @@ class ConditionBlock(Block):
|
||||
value2 = value2.strip()
|
||||
|
||||
yes_value = input_data.yes_value if input_data.yes_value is not None else value1
|
||||
no_value = input_data.no_value if input_data.no_value is not None else value2
|
||||
no_value = input_data.no_value if input_data.no_value is not None else value1
|
||||
|
||||
comparison_funcs = {
|
||||
ComparisonOperator.EQUAL: lambda a, b: a == b,
|
||||
@@ -121,6 +123,74 @@ class ConditionBlock(Block):
|
||||
yield "no_output", no_value
|
||||
|
||||
|
||||
def _parse_exact_float(text: str) -> float | Decimal:
|
||||
"""``parse_float`` hook for ``json.loads`` that keeps decimal values exact.
|
||||
|
||||
A float is returned only when it represents *text* exactly; otherwise the
|
||||
exact ``Decimal`` is kept, which never equals a float. This applies the
|
||||
scalar rule to numbers nested in lists and dicts, so "[0.10000000000000001]"
|
||||
no longer matches [0.1].
|
||||
"""
|
||||
number = float(text)
|
||||
if Decimal(text) == Decimal(repr(number)):
|
||||
return number
|
||||
return Decimal(text)
|
||||
|
||||
|
||||
def _convert_for_match(value: Any, target: Any) -> Any:
|
||||
"""Convert *value* to the type of *target* only when the conversion is lossless.
|
||||
|
||||
``convert`` is deliberately lenient (non-numeric text becomes 0, a list
|
||||
becomes its length, any text other than "true"/"1" becomes False), which
|
||||
is fine for filling block inputs but makes unrelated values compare equal.
|
||||
Values that cannot be converted faithfully are returned unchanged.
|
||||
"""
|
||||
input_type = type(target)
|
||||
if input_type in (Any, object):
|
||||
return value
|
||||
if issubclass(input_type, bool):
|
||||
# Preserve the boolean spellings accepted by convert().
|
||||
if isinstance(value, str):
|
||||
text = value.strip().lower()
|
||||
if text in ("true", "1"):
|
||||
return True
|
||||
if text in ("false", "0"):
|
||||
return False
|
||||
return value
|
||||
if issubclass(input_type, (int, float)):
|
||||
if isinstance(value, str):
|
||||
# Compare exact decimal values: float() would round, so that
|
||||
# "9007199254740993" matched 9007199254740992.
|
||||
text = value.strip()
|
||||
try:
|
||||
number = Decimal(text)
|
||||
except InvalidOperation:
|
||||
return value
|
||||
if number.is_finite() and number == Decimal(str(target)):
|
||||
return target
|
||||
return value
|
||||
if issubclass(input_type, (list, tuple, set, dict)):
|
||||
# convert() wraps a scalar into a container (["x"], {"value": "x"}),
|
||||
# which made "x" match ["x"]. Accept only text that parses to the
|
||||
# same shape, or a list/tuple holding the same items.
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
parsed = json.loads(value, parse_float=_parse_exact_float)
|
||||
except ValueError:
|
||||
return value
|
||||
if issubclass(input_type, dict):
|
||||
return parsed if isinstance(parsed, dict) else value
|
||||
if not isinstance(parsed, list):
|
||||
return value
|
||||
converted = input_type(parsed)
|
||||
# A set drops duplicates, so "[1, 1]" is not the set {1}.
|
||||
return converted if len(converted) == len(parsed) else value
|
||||
if isinstance(value, (list, tuple)) and issubclass(input_type, (list, tuple)):
|
||||
return input_type(value)
|
||||
return value
|
||||
return convert(value, input_type)
|
||||
|
||||
|
||||
class IfInputMatchesBlock(Block):
|
||||
class Input(BlockSchemaInput):
|
||||
input: Any = SchemaField(
|
||||
@@ -203,10 +273,9 @@ class IfInputMatchesBlock(Block):
|
||||
try:
|
||||
# Only attempt conversion if input is not None and value is not None
|
||||
if input_data.input is not None and input_data.value is not None:
|
||||
input_type = type(input_data.input)
|
||||
# Avoid converting if input_type is Any or object
|
||||
if input_type not in (Any, object):
|
||||
input_data.value = convert(input_data.value, input_type)
|
||||
input_data.value = _convert_for_match(
|
||||
input_data.value, input_data.input
|
||||
)
|
||||
except Exception:
|
||||
pass # If conversion fails, just leave value as is
|
||||
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
"""ConditionBlock passes value1 through when a branch value is omitted."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from backend.blocks import get_block
|
||||
from backend.blocks.branching import ComparisonOperator, ConditionBlock
|
||||
from backend.data.execution import ExecutionContext
|
||||
from backend.data.graph import Link, Node
|
||||
from backend.executor.utils import validate_exec
|
||||
from backend.util.exceptions import BlockExecutionError
|
||||
|
||||
|
||||
async def _run(**fields: object) -> dict[str, object]:
|
||||
node = Node(block_id=ConditionBlock().id, input_default=fields)
|
||||
return await _execute_node(node, {})
|
||||
|
||||
|
||||
async def _execute_node(node: Node, data: dict[str, object]) -> dict[str, object]:
|
||||
block = get_block(node.block_id)
|
||||
assert type(block) is ConditionBlock
|
||||
validated, error = validate_exec(node, data)
|
||||
assert validated is not None, error
|
||||
return {
|
||||
name: value
|
||||
async for name, value in block.execute(
|
||||
validated, execution_context=ExecutionContext()
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
async def test_omitted_false_value_passes_through_value1():
|
||||
outputs = await _run(
|
||||
value1=3,
|
||||
operator=ComparisonOperator.GREATER_THAN.value,
|
||||
value2=5,
|
||||
)
|
||||
|
||||
assert outputs["result"] is False
|
||||
assert outputs["no_output"] == 3
|
||||
assert "yes_output" not in outputs
|
||||
|
||||
|
||||
async def test_omitted_true_value_passes_through_value1():
|
||||
outputs = await _run(
|
||||
value1=10,
|
||||
operator=ComparisonOperator.GREATER_THAN.value,
|
||||
value2=5,
|
||||
)
|
||||
|
||||
assert outputs["result"] is True
|
||||
assert outputs["yes_output"] == 10
|
||||
|
||||
|
||||
async def test_explicit_false_value_is_kept():
|
||||
outputs = await _run(
|
||||
value1=3,
|
||||
operator=ComparisonOperator.GREATER_THAN.value,
|
||||
value2=5,
|
||||
no_value="Not greater",
|
||||
)
|
||||
|
||||
assert outputs["no_output"] == "Not greater"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("overrides", [{}, {"yes_value": None, "no_value": None}])
|
||||
@pytest.mark.parametrize("operator,result", [("==", False), ("!=", True)])
|
||||
@pytest.mark.parametrize(
|
||||
"value1,value2,expected",
|
||||
[
|
||||
(" 3 ", " 5 ", 3.0),
|
||||
(" left ", "right", "left"),
|
||||
(False, True, False),
|
||||
(0, 1, 0),
|
||||
("", "other", ""),
|
||||
([], [1], []),
|
||||
({}, {"key": "value"}, {}),
|
||||
],
|
||||
)
|
||||
async def test_missing_or_null_override_uses_normalized_value1(
|
||||
overrides: dict[str, object],
|
||||
operator: str,
|
||||
result: bool,
|
||||
value1: object,
|
||||
value2: object,
|
||||
expected: object,
|
||||
):
|
||||
outputs = await _run(value1=value1, operator=operator, value2=value2, **overrides)
|
||||
|
||||
output_name = "yes_output" if result else "no_output"
|
||||
assert outputs == {"result": result, output_name: expected}
|
||||
assert outputs["result"] is result
|
||||
assert type(outputs[output_name]) is type(expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("override", [False, 0, "", [], {}, " 3 ", "Not greater"])
|
||||
@pytest.mark.parametrize("operator,result", [("==", False), ("!=", True)])
|
||||
async def test_explicit_overrides_are_preserved(
|
||||
override: object, operator: str, result: bool
|
||||
):
|
||||
outputs = await _run(
|
||||
value1=3,
|
||||
operator=operator,
|
||||
value2=5,
|
||||
yes_value=override,
|
||||
no_value=override,
|
||||
)
|
||||
|
||||
output_name = "yes_output" if result else "no_output"
|
||||
assert outputs == {"result": result, output_name: override}
|
||||
assert outputs["result"] is result
|
||||
assert type(outputs[output_name]) is type(override)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("override", [None, False, 0, "", [], {}, " 3 "])
|
||||
async def test_saved_workflow_link_overrides_default(override: object):
|
||||
node = Node(
|
||||
id="condition-node",
|
||||
block_id="715696a0-e1da-45c8-b209-c2fa9c3b0be6",
|
||||
input_default={"value1": 3, "operator": ">", "value2": 5, "no_value": "saved"},
|
||||
input_links=[
|
||||
Link(
|
||||
source_id="upstream-node",
|
||||
sink_id="condition-node",
|
||||
source_name="output",
|
||||
sink_name="no_value",
|
||||
)
|
||||
],
|
||||
)
|
||||
restored = Node.model_validate_json(node.model_dump_json())
|
||||
missing, error = validate_exec(restored, {})
|
||||
assert missing is None
|
||||
assert "unpopulated links" in error
|
||||
|
||||
outputs = await _execute_node(restored, {"no_value": override})
|
||||
|
||||
expected = 3 if override is None else override
|
||||
assert outputs == {"result": False, "no_output": expected}
|
||||
assert type(outputs["no_output"]) is type(expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"fields",
|
||||
[
|
||||
{"value1": 3, "operator": "invalid", "value2": 5},
|
||||
{"operator": ">", "value2": 5},
|
||||
{"value1": 3, "operator": ">"},
|
||||
],
|
||||
)
|
||||
def test_executor_rejects_invalid_inputs(fields: dict[str, object]):
|
||||
data, error = validate_exec(
|
||||
Node(block_id=ConditionBlock().id, input_default=fields), {}
|
||||
)
|
||||
assert data is None
|
||||
assert "Input data missing or mismatch" in error
|
||||
|
||||
|
||||
async def test_incomparable_values_raise_block_execution_error():
|
||||
with pytest.raises(BlockExecutionError, match="Comparison failed"):
|
||||
await _run(value1=[1], operator=">", value2=3)
|
||||
@@ -0,0 +1,223 @@
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from backend.blocks.branching import ConditionBlock, IfInputMatchesBlock
|
||||
from backend.data.execution import ExecutionContext
|
||||
from backend.data.graph import Link, Node
|
||||
from backend.executor.utils import validate_exec
|
||||
from backend.util.exceptions import BlockInputError
|
||||
from backend.util.type import coerce_inputs_to_schema
|
||||
|
||||
|
||||
@pytest.fixture(scope="session", autouse=True)
|
||||
def graph_cleanup():
|
||||
"""These pure block tests do not create persisted graphs or need services."""
|
||||
yield
|
||||
|
||||
|
||||
async def _matches(input_value: Any, value: Any) -> bool:
|
||||
block = IfInputMatchesBlock()
|
||||
outputs = [
|
||||
output
|
||||
async for output in block.run(
|
||||
IfInputMatchesBlock.Input(
|
||||
input=input_value, value=value, yes_value="yes", no_value="no"
|
||||
)
|
||||
)
|
||||
]
|
||||
return dict(outputs)["result"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"input_value, value",
|
||||
[
|
||||
(0, "hello"), # non-numeric text used to coerce to 0
|
||||
(0, "None"),
|
||||
(0.0, "abc"),
|
||||
(10, "10.9"), # used to truncate to 10
|
||||
(3, [1, 2, 3]), # used to compare against len(value)
|
||||
(2, {"a": 1, "b": 2}),
|
||||
(False, "no"), # any string other than "true"/"1" used to become False
|
||||
(False, ""),
|
||||
(9007199254740992, "9007199254740993"), # float() rounded these equal
|
||||
(0.1, "0.10000000000000001"),
|
||||
([0.1], "[0.10000000000000001]"), # nested numbers are exact too
|
||||
({"a": 0.1}, '{"a": 0.10000000000000001}'),
|
||||
(1.0, "nan"),
|
||||
(["hello"], "hello"), # convert() wrapped scalars into containers
|
||||
([3], 3),
|
||||
({"value": "x"}, "x"),
|
||||
({"value": 5}, "5"),
|
||||
(("x",), "x"),
|
||||
({"x"}, "x"),
|
||||
({1}, "[1, 1]"), # set() would drop the duplicate
|
||||
([1, 2], '{"a": 1}'),
|
||||
(True, "yes"), # only the spellings convert() treats as booleans
|
||||
(False, "off"),
|
||||
(True, 2), # used to compare against bool(value)
|
||||
([], {}),
|
||||
],
|
||||
)
|
||||
async def test_lossy_conversion_does_not_report_a_match(input_value, value):
|
||||
assert await _matches(input_value, value) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"input_value, value",
|
||||
[
|
||||
(10, 10),
|
||||
(10, "10"),
|
||||
(10, " 10 "),
|
||||
(10, "10.0"),
|
||||
(2.5, "2.5"),
|
||||
(0.1, "0.1"),
|
||||
(1000, "1e3"),
|
||||
(1.5, "1.50"),
|
||||
(0, "0"),
|
||||
(True, "true"),
|
||||
(True, "True"),
|
||||
(False, "false"),
|
||||
("hello", "hello"),
|
||||
("5", 5),
|
||||
([1, 2], "[1, 2]"),
|
||||
([1, 2], (1, 2)),
|
||||
((1, 2), [1, 2]),
|
||||
({"a": 1}, '{"a": 1}'),
|
||||
(True, "1"), # saved graphs compare toggles against "1"/"0"
|
||||
(True, " 1 "),
|
||||
(False, "0"),
|
||||
(10, "1_0"),
|
||||
(1000, "1_000"),
|
||||
(9007199254740993, "9007199254740993"), # float() used to miss this
|
||||
(True, 1), # plain equality between bool and number is unchanged
|
||||
([True], "[1]"),
|
||||
([0.1], "[0.10]"),
|
||||
({"a": [1.5]}, '{"a": [1.50]}'),
|
||||
([100.0], "[1e2]"),
|
||||
({1, 2}, "[2, 1]"),
|
||||
([1, 1], "[1, 1]"),
|
||||
],
|
||||
)
|
||||
async def test_equivalent_values_still_match(input_value, value):
|
||||
assert await _matches(input_value, value) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"input_value, value",
|
||||
[
|
||||
(10, 20),
|
||||
(10, "None"),
|
||||
(True, "false"),
|
||||
("hello", "world"),
|
||||
],
|
||||
)
|
||||
async def test_different_values_do_not_match(input_value, value):
|
||||
assert await _matches(input_value, value) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"input_value, value, expected",
|
||||
[
|
||||
(True, " 1 ", True),
|
||||
(False, " 0 ", True),
|
||||
(True, " TRUE ", True),
|
||||
(False, " FALSE ", True),
|
||||
(True, 1, True),
|
||||
(False, 0, True),
|
||||
(10, "1_0", True),
|
||||
(1000.0, "1_000", True),
|
||||
(9007199254740993, "9007199254740993", True),
|
||||
(12345678901234567890, "12345678901234567890", True),
|
||||
([True], "[1]", True),
|
||||
({"a": 1}, '{"a": 1}', True),
|
||||
("example.com", "example.com", True),
|
||||
("example.com", "other.example", False),
|
||||
(0, "hello", False),
|
||||
(False, "off", False),
|
||||
(True, 2, False),
|
||||
(10, "10.9", False),
|
||||
(10, 10.9, False),
|
||||
(9007199254740992, "9007199254740993", False),
|
||||
(0.1, "0.10000000000000001", False),
|
||||
(0, "1e999999999999999999999", False),
|
||||
(0, "nan", False),
|
||||
(0, "Infinity", False),
|
||||
(3, [1, 2, 3], False),
|
||||
(2, {"a": 1, "b": 2}, False),
|
||||
(["x"], "x", False),
|
||||
({"value": "x"}, "x", False),
|
||||
([], {}, False),
|
||||
],
|
||||
)
|
||||
async def test_saved_node_inputs_keep_their_types_and_route_correctly(
|
||||
input_value, value, expected
|
||||
):
|
||||
block = IfInputMatchesBlock()
|
||||
fields = json.loads(json.dumps({"input": input_value, "value": value}))
|
||||
coerce_inputs_to_schema(fields, block.input_schema)
|
||||
assert fields["input"] == input_value
|
||||
assert type(fields["input"]) is type(input_value)
|
||||
assert fields["value"] == value
|
||||
assert type(fields["value"]) is type(value)
|
||||
|
||||
node = Node(
|
||||
id="match-node",
|
||||
block_id=block.id,
|
||||
input_default={"value": fields["value"], "yes_value": "yes", "no_value": "no"},
|
||||
input_links=[
|
||||
Link(
|
||||
id="input-link",
|
||||
source_id="upstream-node",
|
||||
sink_id="match-node",
|
||||
source_name="output",
|
||||
sink_name="input",
|
||||
)
|
||||
],
|
||||
)
|
||||
validated, message = validate_exec(node, {"input": fields["input"]})
|
||||
assert validated is not None, message
|
||||
assert validated["input"] == input_value
|
||||
assert type(validated["input"]) is type(input_value)
|
||||
assert validated["value"] == value
|
||||
assert type(validated["value"]) is type(value)
|
||||
outputs = [
|
||||
output
|
||||
async for output in block.execute(
|
||||
validated, execution_context=ExecutionContext()
|
||||
)
|
||||
]
|
||||
branch = ("yes_output", "yes") if expected else ("no_output", "no")
|
||||
assert outputs == [("result", expected), branch]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("missing", ["input", "value"])
|
||||
async def test_null_required_inputs_remain_invalid_at_execution(missing):
|
||||
block = IfInputMatchesBlock()
|
||||
fields = {"input": 0, "value": 0, missing: None}
|
||||
coerce_inputs_to_schema(fields, block.input_schema)
|
||||
assert fields[missing] is None
|
||||
with pytest.raises(BlockInputError):
|
||||
async for _ in block.execute(fields, execution_context=ExecutionContext()):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.parametrize("block_class", [ConditionBlock, IfInputMatchesBlock])
|
||||
async def test_builtin_branching_vectors(block_class):
|
||||
block = block_class()
|
||||
inputs = (
|
||||
block.test_input if isinstance(block.test_input, list) else [block.test_input]
|
||||
)
|
||||
outputs = []
|
||||
for fields in inputs:
|
||||
assert fields is not None
|
||||
outputs.extend(
|
||||
[
|
||||
output
|
||||
async for output in block.execute(
|
||||
fields, execution_context=ExecutionContext()
|
||||
)
|
||||
]
|
||||
)
|
||||
assert outputs == block.test_output
|
||||
@@ -1854,11 +1854,8 @@ class OrchestratorBlock(Block):
|
||||
# We must NOT cancel __anext__() mid-flight — doing so corrupts
|
||||
# the SDK's internal anyio memory stream (same pattern as
|
||||
# copilot/sdk/service.py:_iter_sdk_messages). Every fetch task
|
||||
# runs in one shared context: once configure_claude_agent_sdk()
|
||||
# has wrapped the client, langsmith's tracing wrapper keeps the
|
||||
# conversation run in a ContextVar set during the first fetch,
|
||||
# and a fresh context per task would drop every later reply's
|
||||
# span from the trace.
|
||||
# runs in one shared context, so a ContextVar set while fetching
|
||||
# one message is still set for the next.
|
||||
|
||||
_HEARTBEAT_INTERVAL = 10.0 # seconds
|
||||
options.cli_path = await asyncio.to_thread(
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
import asyncio
|
||||
import io
|
||||
from pathlib import Path
|
||||
|
||||
from backend.blocks._base import (
|
||||
@@ -11,6 +13,7 @@ from backend.blocks._base import (
|
||||
from backend.data.execution import ExecutionContext
|
||||
from backend.data.model import ContributorDetails, SchemaField
|
||||
from backend.util.file import get_exec_file_path, store_media_file
|
||||
from backend.util.file_content_parser import check_xlsx_unpacked_size
|
||||
from backend.util.type import MediaFileType
|
||||
|
||||
|
||||
@@ -134,16 +137,23 @@ class ReadSpreadsheetBlock(Block):
|
||||
try:
|
||||
import pandas as pd
|
||||
|
||||
# dtype=object and keep_default_na=False keep the sheet's own
|
||||
# values: "N/A"/"NULL"/"NaN" stay text instead of being read as
|
||||
# missing, and a gap in a numeric column no longer upcasts the
|
||||
# whole column to float (turning 12 into "12.0").
|
||||
df = pd.read_excel(
|
||||
file_path,
|
||||
sheet_name=input_data.sheet_name or 0,
|
||||
dtype=object,
|
||||
keep_default_na=False,
|
||||
)
|
||||
def _read_xlsx() -> "pd.DataFrame":
|
||||
workbook = Path(file_path).read_bytes()
|
||||
check_xlsx_unpacked_size(
|
||||
io.BytesIO(workbook), source="spreadsheet_block"
|
||||
)
|
||||
# dtype=object and keep_default_na=False keep the sheet's own
|
||||
# values: "N/A"/"NULL"/"NaN" stay text instead of being read as
|
||||
# missing, and a gap in a numeric column no longer upcasts the
|
||||
# whole column to float (turning 12 into "12.0").
|
||||
return pd.read_excel(
|
||||
io.BytesIO(workbook),
|
||||
sheet_name=input_data.sheet_name or 0,
|
||||
dtype=object,
|
||||
keep_default_na=False,
|
||||
)
|
||||
|
||||
df = await asyncio.to_thread(_read_xlsx)
|
||||
|
||||
# Write the CSV with the same dialect it is parsed with below,
|
||||
# otherwise a non-default delimiter yields a single column.
|
||||
|
||||
@@ -14,6 +14,7 @@ backend test suite has no binary fixtures, and building it makes the exact
|
||||
difference between a literal "N/A" cell and a genuinely empty one explicit.
|
||||
"""
|
||||
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
@@ -144,6 +145,107 @@ async def test_excel_cell_containing_the_delimiter_survives(tmp_path):
|
||||
assert rows == [{"name": "widget", "note": "a;b"}]
|
||||
|
||||
|
||||
# The .xlsx caps: 50 MiB per member, 200 MiB per workbook.
|
||||
XLSX_MEMBER_CAP = 50 * 1024 * 1024
|
||||
|
||||
|
||||
def add_member(workbook: Path, name: str, size: int) -> None:
|
||||
"""Append a member of ``size`` zero bytes, written in pieces."""
|
||||
with (
|
||||
zipfile.ZipFile(workbook, "a", zipfile.ZIP_DEFLATED) as archive,
|
||||
archive.open(name, "w") as member,
|
||||
):
|
||||
for start in range(0, size, 4 * 1024 * 1024):
|
||||
member.write(b"\0" * min(4 * 1024 * 1024, size - start))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_excel_member_unpacking_past_the_cap_is_refused_unread(workbook_path):
|
||||
add_member(workbook_path, "xl/media/image1.png", XLSX_MEMBER_CAP + 1)
|
||||
assert workbook_path.stat().st_size < 1_000_000
|
||||
|
||||
with (
|
||||
patch(
|
||||
"pandas.read_excel",
|
||||
side_effect=AssertionError("the workbook was read"),
|
||||
) as read_excel,
|
||||
pytest.raises(ValueError, match="unpacks to more than"),
|
||||
):
|
||||
await read_rows(file_input=MediaFileType("book.xlsx"))
|
||||
read_excel.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_excel_is_read_from_the_bytes_that_were_checked(workbook_path):
|
||||
import io
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from backend.blocks import spreadsheet
|
||||
from backend.util.file_content_parser import check_xlsx_unpacked_size
|
||||
|
||||
checked: list[bytes] = []
|
||||
read: list[bytes] = []
|
||||
real_read_excel = pd.read_excel
|
||||
original = workbook_path.read_bytes()
|
||||
|
||||
def check_then_change_file(archive, **kwargs):
|
||||
checked.append(archive.getvalue())
|
||||
check_xlsx_unpacked_size(archive, **kwargs)
|
||||
workbook_path.write_bytes(b"not the workbook that was checked")
|
||||
|
||||
def record_read(source, *args, **kwargs):
|
||||
assert isinstance(source, io.BytesIO)
|
||||
read.append(source.getvalue())
|
||||
return real_read_excel(source, *args, **kwargs)
|
||||
|
||||
with (
|
||||
patch.object(spreadsheet, "check_xlsx_unpacked_size", check_then_change_file),
|
||||
patch("pandas.read_excel", record_read),
|
||||
):
|
||||
rows = await read_rows(file_input=MediaFileType("book.xlsx"))
|
||||
|
||||
assert checked == [original]
|
||||
assert read == checked
|
||||
assert rows[0]["status"] == "N/A"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_excel_member_at_the_cap_is_read(workbook_path):
|
||||
add_member(workbook_path, "xl/media/image1.png", XLSX_MEMBER_CAP)
|
||||
|
||||
rows = await read_rows(file_input=MediaFileType("book.xlsx"))
|
||||
|
||||
assert rows[0]["status"] == "N/A"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_excel_member_past_the_cap_is_reported_once(workbook_path):
|
||||
add_member(workbook_path, "xl/media/image1.png", XLSX_MEMBER_CAP + 1)
|
||||
|
||||
with (
|
||||
patch("backend.util.file_content_parser.sentry_sdk.capture_message") as report,
|
||||
pytest.raises(ValueError, match="unpacks to more than"),
|
||||
):
|
||||
await read_rows(file_input=MediaFileType("book.xlsx"))
|
||||
|
||||
report.assert_called_once()
|
||||
assert report.call_args.args == ("xlsx workbook over unpack limit",)
|
||||
assert report.call_args.kwargs["level"] == "warning"
|
||||
assert report.call_args.kwargs["tags"] == {
|
||||
"xlsx_limit": "MAX_XLSX_MEMBER_BYTES",
|
||||
"xlsx_source": "spreadsheet_block",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_excel_within_the_caps_is_not_reported(workbook_path):
|
||||
with patch("backend.util.file_content_parser.sentry_sdk.capture_message") as report:
|
||||
await read_rows(file_input=MediaFileType("book.xlsx"))
|
||||
|
||||
report.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skip_columns_skips_by_positional_index():
|
||||
rows = await read_rows(contents="a,b,c\n1,2,3\n4,5,6", skip_columns=["0"])
|
||||
|
||||
@@ -1,140 +0,0 @@
|
||||
"""langsmith's Claude Agent SDK tracing must survive the orchestrator's SDK mode.
|
||||
|
||||
The tracing wrapper around ``receive_response()`` keeps the
|
||||
``claude.conversation`` run in a ContextVar that it sets while the first
|
||||
message is fetched, and parents every ``claude.assistant.turn`` run on it.
|
||||
``_execute_tools_sdk_mode`` fetches each message in its own task, so unless
|
||||
every fetch runs in the same context, replies after the first message find no
|
||||
parent and are silently dropped from the trace.
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import claude_agent_sdk
|
||||
import pytest
|
||||
from claude_agent_sdk import (
|
||||
AssistantMessage,
|
||||
ClaudeAgentOptions,
|
||||
ResultMessage,
|
||||
SystemMessage,
|
||||
TextBlock,
|
||||
ToolResultBlock,
|
||||
ToolUseBlock,
|
||||
UserMessage,
|
||||
)
|
||||
from langsmith import tracing_context
|
||||
from langsmith.integrations.claude_agent_sdk import configure_claude_agent_sdk
|
||||
from langsmith.run_trees import RunTree
|
||||
from pydantic import SecretStr
|
||||
|
||||
from backend.blocks.orchestrator import OrchestratorBlock
|
||||
from backend.data.model import APIKeyCredentials
|
||||
|
||||
|
||||
class _ScriptedClient:
|
||||
"""Stands in for ``ClaudeSDKClient`` and replays one agent turn."""
|
||||
|
||||
def __init__(self, options: ClaudeAgentOptions) -> None:
|
||||
self.options = options
|
||||
|
||||
async def __aenter__(self) -> "_ScriptedClient":
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_: object) -> None:
|
||||
return None
|
||||
|
||||
async def query(self, prompt: str, session_id: str = "default") -> None:
|
||||
return None
|
||||
|
||||
async def receive_response(self) -> AsyncGenerator[Any, None]:
|
||||
yield SystemMessage(subtype="init", data={})
|
||||
yield AssistantMessage(
|
||||
content=[ToolUseBlock(id="toolu_1", name="lookup", input={})],
|
||||
model="claude-test",
|
||||
message_id="msg_1",
|
||||
)
|
||||
yield UserMessage(content=[ToolResultBlock(tool_use_id="toolu_1")])
|
||||
yield AssistantMessage(
|
||||
content=[TextBlock(text="done")], model="claude-test", message_id="msg_2"
|
||||
)
|
||||
yield ResultMessage(
|
||||
subtype="success",
|
||||
duration_ms=1,
|
||||
duration_api_ms=1,
|
||||
is_error=False,
|
||||
num_turns=2,
|
||||
session_id="sess-1",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def traced_client() -> type[_ScriptedClient]:
|
||||
"""Wrap the scripted client exactly as production wraps ``ClaudeSDKClient``.
|
||||
|
||||
``configure_claude_agent_sdk`` patches whichever class
|
||||
``claude_agent_sdk.ClaudeSDKClient`` names, so pointing that name at the
|
||||
scripted client for the call wraps it through the public API. The real
|
||||
tool class and the global trace config are left untouched.
|
||||
"""
|
||||
with (
|
||||
patch.object(claude_agent_sdk, "ClaudeSDKClient", _ScriptedClient),
|
||||
patch.object(claude_agent_sdk, "SdkMcpTool", None),
|
||||
patch("langsmith.integrations.claude_agent_sdk.set_tracing_config"),
|
||||
):
|
||||
assert configure_claude_agent_sdk()
|
||||
return _ScriptedClient
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def posted_runs(monkeypatch: pytest.MonkeyPatch) -> list[RunTree]:
|
||||
"""Record every run the tracing wrapper posts instead of exporting it."""
|
||||
posted: list[RunTree] = []
|
||||
monkeypatch.setattr(RunTree, "post", lambda self, *_, **__: posted.append(self))
|
||||
monkeypatch.setattr(RunTree, "patch", lambda self, *_, **__: None)
|
||||
return posted
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sdk_mode_replies_get_assistant_turn_runs(
|
||||
traced_client: type[_ScriptedClient],
|
||||
posted_runs: list[RunTree],
|
||||
) -> None:
|
||||
input_data = MagicMock(sys_prompt="", prompt="list the files")
|
||||
input_data.model.value = "claude-test"
|
||||
credentials = APIKeyCredentials(
|
||||
id="test-anthropic-id",
|
||||
provider="anthropic",
|
||||
api_key=SecretStr("mock-anthropic-key"),
|
||||
title="Mock Anthropic key",
|
||||
expires_at=None,
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(claude_agent_sdk, "ClaudeSDKClient", traced_client),
|
||||
patch(
|
||||
"backend.blocks.orchestrator.require_verbatim_prompts_cli",
|
||||
return_value="claude",
|
||||
),
|
||||
tracing_context(enabled=True),
|
||||
):
|
||||
outputs = [
|
||||
output
|
||||
async for output in OrchestratorBlock()._execute_tools_sdk_mode(
|
||||
input_data=input_data,
|
||||
credentials=credentials,
|
||||
tool_functions=[],
|
||||
prompt=[{"role": "user", "content": "list the files"}],
|
||||
execution_params=MagicMock(graph_exec_id="graph-exec-1"),
|
||||
execution_processor=MagicMock(),
|
||||
)
|
||||
]
|
||||
|
||||
assert ("finished", "done") in outputs
|
||||
conversation = next(run for run in posted_runs if run.name == "claude.conversation")
|
||||
turns = [run for run in posted_runs if run.name == "claude.assistant.turn"]
|
||||
assert [turn.parent_run_id for turn in turns] == [
|
||||
conversation.id,
|
||||
conversation.id,
|
||||
], "each assistant reply should get a turn run under the conversation"
|
||||
@@ -21,12 +21,10 @@ from functools import partial
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
import orjson
|
||||
from langfuse import propagate_attributes
|
||||
from openai import APIConnectionError
|
||||
from openai import omit as openai_omit
|
||||
from openai.types import CompletionUsage
|
||||
from openai.types.chat import ChatCompletionMessageParam, ChatCompletionToolParam
|
||||
from opentelemetry import trace as otel_trace
|
||||
|
||||
from backend.blocks.desktop._common import workspace_volume_mounts
|
||||
from backend.copilot import engine_switch
|
||||
@@ -1016,8 +1014,8 @@ async def _baseline_llm_caller(
|
||||
# Route through the shared providers helper so future provider
|
||||
# work (streaming flex tier, new SDK upgrades) propagates here
|
||||
# without a parallel migration. The pre-built ``client`` is
|
||||
# passed through so the module-level Langfuse-wrapped client
|
||||
# keeps its HTTP connection pool warm across turns.
|
||||
# passed through so the module-level client keeps its HTTP
|
||||
# connection pool warm across turns.
|
||||
response = await call_provider_stream(
|
||||
client=client,
|
||||
model=state.model,
|
||||
@@ -2582,20 +2580,6 @@ async def stream_chat_completion_baseline(
|
||||
format_pending_as_user_message(pm) for pm in held_results
|
||||
)
|
||||
|
||||
# Propagate user/session context to Langfuse so all LLM calls within
|
||||
# this request are grouped under a single trace with proper attribution.
|
||||
_trace_ctx: Any = None
|
||||
try:
|
||||
_trace_ctx = propagate_attributes(
|
||||
user_id=user_id,
|
||||
session_id=session_id,
|
||||
trace_name="copilot-baseline",
|
||||
tags=["baseline", "tool_surface:registry"],
|
||||
)
|
||||
_trace_ctx.__enter__()
|
||||
except Exception:
|
||||
logger.warning("[Baseline] Langfuse trace context setup failed")
|
||||
|
||||
_stream_error = False # Track whether an error occurred during streaming
|
||||
state = _BaselineStreamState(
|
||||
model=active_model,
|
||||
@@ -3013,27 +2997,6 @@ async def stream_chat_completion_baseline(
|
||||
# Any message pushed after the final drain window stays in the
|
||||
# buffer and gets picked up at the start of the next turn.
|
||||
|
||||
# Set cost attributes on OTEL span before closing
|
||||
if _trace_ctx is not None:
|
||||
try:
|
||||
span = otel_trace.get_current_span()
|
||||
if span and span.is_recording():
|
||||
span.set_attribute(
|
||||
"gen_ai.usage.prompt_tokens", state.turn_prompt_tokens
|
||||
)
|
||||
span.set_attribute(
|
||||
"gen_ai.usage.completion_tokens",
|
||||
state.turn_completion_tokens,
|
||||
)
|
||||
if state.cost_usd is not None:
|
||||
span.set_attribute("gen_ai.usage.cost_usd", state.cost_usd)
|
||||
except Exception:
|
||||
logger.debug("[Baseline] Failed to set OTEL cost attributes")
|
||||
try:
|
||||
_trace_ctx.__exit__(None, None, None)
|
||||
except Exception:
|
||||
logger.warning("[Baseline] Langfuse trace context teardown failed")
|
||||
|
||||
# Fallback: estimate tokens via tiktoken when the provider does
|
||||
# not honour stream_options={"include_usage": True}.
|
||||
# ``openai_messages`` is mutated in-place by tool_call_loop so it
|
||||
|
||||
@@ -64,7 +64,6 @@ def _paying_user(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"get_redis_async",
|
||||
AsyncMock(return_value=MagicMock(get=AsyncMock(return_value=None))),
|
||||
)
|
||||
monkeypatch.setattr(sdk_service, "get_user_tier", AsyncMock(return_value=None))
|
||||
monkeypatch.setattr(sdk_service, "_graphiti_enabled", lambda _user_id: False)
|
||||
monkeypatch.setattr(
|
||||
sdk_service,
|
||||
|
||||
@@ -134,7 +134,6 @@ class SdkTurnsViaExecutor(SdkTurns):
|
||||
|
||||
def _stubs(self, session: ChatSession) -> dict[str, Any]:
|
||||
stubs = super()._stubs(session)
|
||||
del stubs["get_user_tier"]
|
||||
del stubs["resolve_turn_budget"]
|
||||
if self._tools:
|
||||
del stubs["create_copilot_mcp_server"]
|
||||
|
||||
@@ -136,7 +136,6 @@ class SdkTurns:
|
||||
"delete_transcript": AsyncMock(side_effect=self._delete),
|
||||
"compact_transcript": AsyncMock(return_value=None),
|
||||
"get_redis_async": AsyncMock(),
|
||||
"propagate_attributes": MagicMock(),
|
||||
"_build_system_prompt": AsyncMock(
|
||||
return_value=("You are a test assistant. Reply briefly.", None)
|
||||
),
|
||||
@@ -148,7 +147,6 @@ class SdkTurns:
|
||||
"create_security_hooks": MagicMock(),
|
||||
"get_copilot_tool_names": MagicMock(return_value=[]),
|
||||
"get_sdk_disallowed_tools": MagicMock(return_value=[]),
|
||||
"get_user_tier": AsyncMock(return_value=None),
|
||||
# Well above any turn here, so the governor never lands one.
|
||||
"resolve_turn_budget": AsyncMock(
|
||||
return_value=TurnBudget(usd=100.0, limiting="task")
|
||||
|
||||
@@ -43,7 +43,6 @@ from backend.copilot.inference.context import (
|
||||
)
|
||||
from backend.copilot.inference.record import record
|
||||
from backend.copilot.inference.routing import resolve_route
|
||||
from backend.copilot.inference.trace import trace
|
||||
from backend.copilot.trial_cost_context import attributed_usage
|
||||
|
||||
from .models import BriefingContent, BriefingRunItem
|
||||
@@ -104,7 +103,7 @@ async def compose_narrative(user_id: str, content: BriefingContent) -> str | Non
|
||||
{"role": "system", "content": _system_prompt()},
|
||||
{"role": "user", "content": _facts_block(content)},
|
||||
]
|
||||
# Both attempts at one briefing share a correlation id (their trace session).
|
||||
# Both attempts at one briefing share a correlation id.
|
||||
correlation_id = str(uuid.uuid4())
|
||||
loop = asyncio.get_running_loop()
|
||||
deadline = loop.time() + _TOTAL_BUDGET_SECONDS
|
||||
@@ -147,26 +146,24 @@ async def _attempt(
|
||||
async def _complete(
|
||||
user_id: str, job: InferenceJob, messages: list[dict[str, str]]
|
||||
) -> StructuredCompletion[NarrativeResponse]:
|
||||
"""The call, traced, on Otto's own scope: the briefing is never an
|
||||
expert's. Its spend is recorded whether or not the answer parsed."""
|
||||
"""The call, on Otto's own scope: the briefing is never an expert's. Its
|
||||
spend is recorded whether or not the answer parsed."""
|
||||
scope = InferenceScope(user_id=user_id)
|
||||
ctx = InferenceContext(
|
||||
scope=scope, job=job, route=resolve_route(scope, job, config=config)
|
||||
)
|
||||
async with trace(ctx) as call:
|
||||
try:
|
||||
completion = await structured_complete(
|
||||
call.ctx,
|
||||
messages,
|
||||
NarrativeResponse,
|
||||
max_output_tokens=_MAX_OUTPUT_TOKENS,
|
||||
)
|
||||
except InferenceError as e:
|
||||
# A response that arrived but didn't parse was still billed; the
|
||||
# trace reports it as priced, like a call that parsed.
|
||||
call.usage = await _record_spend(call.ctx, e.usage)
|
||||
raise
|
||||
call.usage = await _record_spend(call.ctx, completion.usage)
|
||||
try:
|
||||
completion = await structured_complete(
|
||||
ctx,
|
||||
messages,
|
||||
NarrativeResponse,
|
||||
max_output_tokens=_MAX_OUTPUT_TOKENS,
|
||||
)
|
||||
except InferenceError as e:
|
||||
# A response that arrived but didn't parse was still billed.
|
||||
await _record_spend(ctx, e.usage)
|
||||
raise
|
||||
await _record_spend(ctx, completion.usage)
|
||||
return completion
|
||||
|
||||
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
from contextlib import ExitStack
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
@@ -359,32 +358,6 @@ async def test_failed_attempt_that_the_provider_billed_is_still_recorded(cost_lo
|
||||
assert cost_log.await_args.kwargs["prompt_tokens"] == 500
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_billed_failure_reaches_the_trace_priced(cost_log):
|
||||
"""The span of an attempt whose answer did not parse carries the cost the
|
||||
ledger charged, not the provider's unpriced usage."""
|
||||
usage = InferenceUsage(
|
||||
model="claude-haiku-4-5-20251001",
|
||||
input_tokens=1_000_000,
|
||||
payer="platform_allowance",
|
||||
)
|
||||
written: list[InferenceUsage | None] = []
|
||||
with patch_llm(side_effect=InferenceError("bad json", usage)), patch(
|
||||
"backend.copilot.inference.trace._open_spans", return_value=ExitStack()
|
||||
), patch(
|
||||
"backend.copilot.inference.trace._current_trace_id", return_value=None
|
||||
), patch(
|
||||
"backend.copilot.inference.trace._write_usage", side_effect=written.append
|
||||
):
|
||||
assert await compose_narrative(USER, make_content()) is None
|
||||
|
||||
assert cost_log.await_count == 2
|
||||
assert [u.cost_usd if u else None for u in written] == [
|
||||
pytest.approx(1.0),
|
||||
pytest.approx(1.0),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_transport_failure_is_not_billed(cost_log):
|
||||
with patch_llm(side_effect=InferenceError("no api key")):
|
||||
|
||||
@@ -173,7 +173,7 @@ _TRANSPORT_PROFILES: dict[TransportName, TransportProfile] = {
|
||||
# - 'extended_thinking': route to the Claude Agent SDK path with the default
|
||||
# (opus) model.
|
||||
# ``None`` means "no override"; the server falls back to the Claude Code
|
||||
# subscription flag → LaunchDarkly COPILOT_SDK → config.use_claude_agent_sdk.
|
||||
# subscription flag → the COPILOT_SDK flag → config.use_claude_agent_sdk.
|
||||
CopilotMode = Literal["fast", "extended_thinking"]
|
||||
|
||||
# Per-request model tier set by the frontend model toggle.
|
||||
@@ -426,15 +426,15 @@ class ChatConfig(BaseSettings):
|
||||
# User.subscriptionTier DB column and resolved inside
|
||||
# get_global_rate_limits().
|
||||
#
|
||||
# These defaults act as the ceiling when LaunchDarkly is unreachable;
|
||||
# These defaults act as the ceiling when PostHog is unreachable;
|
||||
# the live per-tier values come from the COPILOT_*_COST_LIMIT flags.
|
||||
#
|
||||
# A negative value disables that window's cap. Self-hosted distributions
|
||||
# (the single-container image and the unraid template) export -1 for both
|
||||
# because the operator pays the model provider directly and, without
|
||||
# LaunchDarkly, every account resolves to NO_TIER → BASIC multiplier and
|
||||
# PostHog, every account resolves to NO_TIER → BASIC multiplier and
|
||||
# would otherwise inherit these cloud ceilings. The defaults here stay
|
||||
# positive so a LaunchDarkly outage on cloud never removes the cap.
|
||||
# positive so a PostHog outage on cloud never removes the cap.
|
||||
daily_cost_limit_microdollars: int = Field(
|
||||
default=2_500_000,
|
||||
description="Max cost per day in microdollars, resets at midnight UTC. "
|
||||
|
||||
@@ -719,8 +719,6 @@ async def add_chat_messages_batch(
|
||||
data["llmAuthProvider"] = msg["llm_auth_provider"]
|
||||
if msg.get("llm_credential_id") is not None:
|
||||
data["llmCredentialId"] = msg["llm_credential_id"]
|
||||
if msg.get("langfuse_trace_id") is not None:
|
||||
data["langfuseTraceId"] = msg["langfuse_trace_id"]
|
||||
if msg.get("trace_id") is not None:
|
||||
data["traceId"] = msg["trace_id"]
|
||||
|
||||
@@ -1402,7 +1400,6 @@ async def update_chat_message_stamps(
|
||||
routing_source: str | None,
|
||||
llm_auth_provider: str | None = None,
|
||||
llm_credential_id: str | None = None,
|
||||
langfuse_trace_id: str | None = None,
|
||||
trace_id: str | None = None,
|
||||
) -> bool:
|
||||
"""Back-fill the execution stamps on an already-persisted message row.
|
||||
@@ -1425,8 +1422,6 @@ async def update_chat_message_stamps(
|
||||
data["llmAuthProvider"] = llm_auth_provider
|
||||
if llm_credential_id is not None:
|
||||
data["llmCredentialId"] = llm_credential_id
|
||||
if langfuse_trace_id is not None:
|
||||
data["langfuseTraceId"] = langfuse_trace_id
|
||||
if trace_id is not None:
|
||||
data["traceId"] = trace_id
|
||||
result = await PrismaChatMessage.prisma().update(
|
||||
|
||||
@@ -1105,7 +1105,6 @@ async def test_batch_persist_maps_stamp_columns_to_prisma_names():
|
||||
"content": "hi",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"routing_source": "catalog",
|
||||
"langfuse_trace_id": "1edf31f11b1693cc6103f358c1481694",
|
||||
"trace_id": "turn-1",
|
||||
}
|
||||
],
|
||||
@@ -1115,8 +1114,6 @@ async def test_batch_persist_maps_stamp_columns_to_prisma_names():
|
||||
assert row["model"] == "claude-sonnet-4-6"
|
||||
assert row["routingSource"] == "catalog"
|
||||
assert "routing_source" not in row
|
||||
assert row["langfuseTraceId"] == "1edf31f11b1693cc6103f358c1481694"
|
||||
assert "langfuse_trace_id" not in row
|
||||
assert row["traceId"] == "turn-1"
|
||||
assert "trace_id" not in row
|
||||
|
||||
@@ -1133,17 +1130,14 @@ def test_from_db_restores_stamp_columns():
|
||||
content="hi",
|
||||
model="claude-sonnet-4-6",
|
||||
routingSource="catalog",
|
||||
langfuseTraceId="1edf31f11b1693cc6103f358c1481694",
|
||||
traceId="turn-1",
|
||||
createdAt=datetime.now(UTC),
|
||||
)
|
||||
restored = ChatMessage.from_db(prisma_msg)
|
||||
assert restored.model == "claude-sonnet-4-6"
|
||||
assert restored.routing_source == "catalog"
|
||||
assert restored.langfuse_trace_id == "1edf31f11b1693cc6103f358c1481694"
|
||||
# Internal: the trace id never reaches the session payload clients get.
|
||||
assert "langfuse_trace_id" not in restored.model_dump()
|
||||
assert restored.trace_id == "turn-1"
|
||||
# Internal: the trace id never reaches the session payload clients get.
|
||||
assert "trace_id" not in restored.model_dump()
|
||||
|
||||
|
||||
@@ -1206,7 +1200,7 @@ async def test_update_chat_message_stamps_writes_the_trace_only_when_known():
|
||||
sequence=7,
|
||||
model="claude-sonnet-4-6",
|
||||
routing_source="env",
|
||||
langfuse_trace_id="1edf31f11b1693cc6103f358c1481694",
|
||||
trace_id="turn-1",
|
||||
)
|
||||
with_trace = update.call_args.kwargs["data"]
|
||||
await update_chat_message_stamps(
|
||||
@@ -1217,8 +1211,8 @@ async def test_update_chat_message_stamps_writes_the_trace_only_when_known():
|
||||
)
|
||||
without_trace = update.call_args.kwargs["data"]
|
||||
|
||||
assert with_trace["langfuseTraceId"] == "1edf31f11b1693cc6103f358c1481694"
|
||||
assert "langfuseTraceId" not in without_trace
|
||||
assert with_trace["traceId"] == "turn-1"
|
||||
assert "traceId" not in without_trace
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -5,8 +5,8 @@ Walks a user's recent memory window through the consolidate → recombine
|
||||
(clamped by ``clamp.py``) to Graphiti + Postgres.
|
||||
|
||||
Each phase here is one ``inference.complete.structured_complete`` call
|
||||
on the chat transport's provider, traced and recorded through the
|
||||
inference package. When ``routing.resolve_dream_execution_path``
|
||||
on the chat transport's provider, recorded through the inference
|
||||
package. When ``routing.resolve_dream_execution_path``
|
||||
picks ``anthropic_batch``, ``batch_handoff.submit_dream_pass_batch``
|
||||
submits the first phase to Anthropic's Message Batches API instead and
|
||||
``batch_callbacks`` runs the later phases and the apply step as the
|
||||
@@ -45,7 +45,6 @@ from backend.copilot.inference.context import (
|
||||
InferenceScope,
|
||||
)
|
||||
from backend.copilot.inference.routing import resolve_route
|
||||
from backend.copilot.inference.trace import trace
|
||||
from backend.copilot.trial_cost_context import attributed_usage
|
||||
from backend.util.feature_flag import Flag, is_feature_enabled
|
||||
|
||||
@@ -284,28 +283,26 @@ async def _run_phase(
|
||||
max_output_tokens: int,
|
||||
timeout_seconds: float,
|
||||
) -> tuple[_Output, PhaseUsage]:
|
||||
"""One phase's call on the pass's route, traced, with its cost recorded.
|
||||
"""One phase's call on the pass's route, with its cost recorded.
|
||||
|
||||
The usage comes back priced: the provider's cost when it reported one
|
||||
(OpenRouter's ``usage.cost``, what we were billed), else the model's
|
||||
catalog list rate; unknown when the response reported no usage at all.
|
||||
|
||||
Raises ``InferenceError`` when the phase got no usable answer. Failed
|
||||
phases are absorbed; their usage remains on the error for tracing.
|
||||
phases are absorbed; their usage remains on the error.
|
||||
"""
|
||||
job = phase_job(phase, run.pass_id, timeout_seconds=timeout_seconds)
|
||||
route = resolve_route(run.scope, job, config=run.config)
|
||||
ctx = InferenceContext(scope=run.scope, job=job, route=route)
|
||||
async with trace(ctx) as call:
|
||||
completion = await structured_complete(
|
||||
call.ctx,
|
||||
messages,
|
||||
response_model,
|
||||
temperature=temperature,
|
||||
max_output_tokens=max_output_tokens,
|
||||
)
|
||||
usage = await record_phase_cost(call.ctx, completion.usage)
|
||||
call.usage = usage
|
||||
completion = await structured_complete(
|
||||
ctx,
|
||||
messages,
|
||||
response_model,
|
||||
temperature=temperature,
|
||||
max_output_tokens=max_output_tokens,
|
||||
)
|
||||
usage = await record_phase_cost(ctx, completion.usage)
|
||||
return completion.value, phase_usage(phase, usage)
|
||||
|
||||
|
||||
|
||||
@@ -33,7 +33,6 @@ from backend.copilot.inference.context import (
|
||||
InferenceUsage,
|
||||
RouteDecision,
|
||||
)
|
||||
from backend.copilot.inference.trace import TracedCall
|
||||
from backend.copilot.trial_cost_context import get_trial_cost_context
|
||||
from backend.data.dream_pass_models import DreamPassDraft
|
||||
from backend.data.pro_activation import UsageActivationState
|
||||
@@ -211,7 +210,7 @@ def _charge_spy() -> AsyncMock:
|
||||
@pytest.fixture(autouse=True)
|
||||
def _stub_batch_flag(mocker):
|
||||
"""Default the batch-path LD flag off so tests stay on sync + hermetic
|
||||
(no LaunchDarkly calls). The wiring test re-patches to assert flag flow."""
|
||||
(no PostHog calls). The wiring test re-patches to assert flag flow."""
|
||||
mocker.patch.object(
|
||||
orchestrator_mod, "is_feature_enabled", AsyncMock(return_value=False)
|
||||
)
|
||||
@@ -1163,24 +1162,8 @@ async def test_partial_failure_still_charges_completed_phases(
|
||||
apply_mock.assert_not_awaited()
|
||||
|
||||
|
||||
def _tracing_failures(seen: list[InferenceUsage | None]):
|
||||
"""A stand-in for ``trace`` that keeps the usage the real one would write
|
||||
on a span closing on an error (``call.usage or exc.usage``)."""
|
||||
|
||||
@asynccontextmanager
|
||||
async def fake_trace(ctx):
|
||||
call = TracedCall(ctx=ctx)
|
||||
try:
|
||||
yield call
|
||||
except InferenceError as exc:
|
||||
seen.append(call.usage or exc.usage)
|
||||
raise
|
||||
|
||||
return fake_trace
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_phase_preserves_trace_usage_without_charging(mocker):
|
||||
async def test_a_failed_phase_is_not_charged(mocker):
|
||||
mocker.patch.object(
|
||||
orchestrator_mod,
|
||||
"gather_dream_input",
|
||||
@@ -1207,8 +1190,6 @@ async def test_a_failed_phase_preserves_trace_usage_without_charging(mocker):
|
||||
persist = mocker.patch(
|
||||
"backend.copilot.inference.record.persist_and_record_usage", AsyncMock()
|
||||
)
|
||||
traced: list[InferenceUsage | None] = []
|
||||
mocker.patch.object(orchestrator_mod, "trace", _tracing_failures(traced))
|
||||
apply_mock = mocker.patch.object(orchestrator_mod, "apply_operations", AsyncMock())
|
||||
|
||||
result = await orchestrator_mod.execute_dream_pass("u", expert_id="expert-1")
|
||||
@@ -1216,8 +1197,6 @@ async def test_a_failed_phase_preserves_trace_usage_without_charging(mocker):
|
||||
assert result.error == "consolidate: schema_mismatch", "a code, never the answer"
|
||||
apply_mock.assert_not_awaited()
|
||||
persist.assert_not_awaited()
|
||||
(seen,) = traced
|
||||
assert seen == billed
|
||||
assert result.usage is not None
|
||||
assert result.usage.phases == []
|
||||
assert result.usage.total_cost_usd == 0
|
||||
|
||||
@@ -23,7 +23,7 @@ async def resolve_use_sdk(
|
||||
"""Pick the SDK vs baseline engine.
|
||||
|
||||
Entirely the server's call: the Claude Code subscription override, then
|
||||
the ``COPILOT_SDK`` LaunchDarkly flag, then the config default. Callers
|
||||
the ``COPILOT_SDK`` feature flag, then the config default. Callers
|
||||
used to be able to name an engine per request; nothing can now, so this
|
||||
is one decision rather than a request overriding it — which is also what
|
||||
makes it answerable ahead of a turn.
|
||||
|
||||
@@ -163,8 +163,8 @@ def chat_system_prompt(expert: Expert | None) -> str:
|
||||
|
||||
async def resolve_chat_model(config: ChatConfig) -> RoutedModel:
|
||||
"""The model an expert turn runs on: the engine decision, then the
|
||||
router's ``(mode, standard)`` cell. No user id, so LaunchDarkly is
|
||||
skipped and this is the catalog/env layer; a proposed LD value is
|
||||
router's ``(mode, standard)`` cell. No user id, so the per-user flag is
|
||||
skipped and this is the catalog/env layer; a proposed flag value is
|
||||
checked by hand with ``--model``."""
|
||||
use_sdk = await resolve_use_sdk(
|
||||
None,
|
||||
|
||||
@@ -174,7 +174,7 @@ def test_the_baseline_records_which_delegation_state_it_measured():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_model_comes_from_the_router_without_launchdarkly():
|
||||
async def test_chat_model_comes_from_the_router_without_posthog():
|
||||
config = ChatConfig(
|
||||
use_claude_agent_sdk=True,
|
||||
thinking_standard_model="anthropic/claude-style-test",
|
||||
|
||||
@@ -11,11 +11,10 @@ from backend.copilot.inference.context import (
|
||||
InferenceScope,
|
||||
)
|
||||
from backend.copilot.inference.routing import resolve_route
|
||||
from backend.copilot.inference.trace import trace
|
||||
|
||||
from .models import Judgement, Rubric, Usage
|
||||
|
||||
# The eval runs outside any account; its judge traces go under this user.
|
||||
# The eval runs outside any account; its judge calls go under this user.
|
||||
EVAL_USER_ID = "expert-style-eval"
|
||||
JUDGE_TIMEOUT_SECONDS = 60.0
|
||||
# Haiku ran out of room mid-JSON writing long evidence quotes on 5 of run A's
|
||||
@@ -35,15 +34,13 @@ async def judge_response(
|
||||
run_id: str,
|
||||
) -> tuple[Judgement, Usage]:
|
||||
ctx = judge_context(expert, model=model, run_id=run_id)
|
||||
async with trace(ctx) as call:
|
||||
completion = await structured_complete(
|
||||
call.ctx,
|
||||
judge_messages(expert, rubric, kind=kind, prompt=prompt, response=response),
|
||||
Judgement,
|
||||
temperature=0.0,
|
||||
max_output_tokens=JUDGE_MAX_OUTPUT_TOKENS,
|
||||
)
|
||||
call.usage = completion.usage
|
||||
completion = await structured_complete(
|
||||
ctx,
|
||||
judge_messages(expert, rubric, kind=kind, prompt=prompt, response=response),
|
||||
Judgement,
|
||||
temperature=0.0,
|
||||
max_output_tokens=JUDGE_MAX_OUTPUT_TOKENS,
|
||||
)
|
||||
judgement = completion.value
|
||||
missing = {d.key for d in rubric.dimensions} - judgement.scores.keys()
|
||||
if missing:
|
||||
|
||||
@@ -797,7 +797,7 @@ class CoPilotProcessor:
|
||||
else:
|
||||
if entry.llm_credential_id is not None:
|
||||
raise RuntimeError("codex_session_route_mismatch")
|
||||
# Choose service based on LaunchDarkly flag.
|
||||
# Choose service based on PostHog flag.
|
||||
# Claude Code subscription forces SDK mode (CLI subprocess auth).
|
||||
config = ChatConfig()
|
||||
|
||||
|
||||
@@ -1,18 +1,8 @@
|
||||
"""Thumbs up/down and copy ratings on AutoPilot replies.
|
||||
|
||||
Postgres (``ChatMessageFeedback``) keeps every rating. Each one is also sent
|
||||
to Langfuse as a score named after the rating ("user-feedback": 1 = up,
|
||||
0 = down; "copy": 1), in the environment the copilot traces use:
|
||||
|
||||
- on the trace of the SDK turn that wrote the reply, which the turn stamps
|
||||
on its assistant rows (``ChatMessage.langfuseTraceId``);
|
||||
- otherwise on the chat's Langfuse session, with the message in the score's
|
||||
metadata: baseline turns have no turn-level trace, and replies written
|
||||
before the stamp existed have none recorded.
|
||||
|
||||
The score reuses the rating's id, so rating a reply again updates its score
|
||||
instead of adding one. A Langfuse failure is logged and shows in the
|
||||
response; the rating itself is already saved.
|
||||
Postgres (``ChatMessageFeedback``) keeps every rating, named after what it
|
||||
rates ("user-feedback": 1 = up, 0 = down; "copy": 1). Rating a reply again
|
||||
overwrites its rating and keeps its id.
|
||||
|
||||
Each rating also goes to PostHog LLM analytics as an ``$ai_metric`` event on
|
||||
the user (distinct id = user id), for every user regardless of cookie
|
||||
@@ -27,33 +17,20 @@ the column existed has none and attaches to the chat only. Rating again sends
|
||||
another event; the latest per ``feedback_id`` is the rating that stands.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Literal
|
||||
|
||||
from langfuse import get_client
|
||||
from langfuse.api import CreateScoreRequest, ScoreDataType
|
||||
from prisma.models import ChatMessage as PrismaChatMessage
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from backend.copilot import feedback_db
|
||||
from backend.util.posthog_client import capture
|
||||
from backend.util.posthog_events import PostHogEvent
|
||||
from backend.util.settings import Settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
settings = Settings()
|
||||
|
||||
FeedbackScoreName = Literal["user-feedback", "copy"]
|
||||
FeedbackScoreValue = Literal[0, 1]
|
||||
LangfuseScoreTarget = Literal["trace", "session"]
|
||||
|
||||
MAX_FEEDBACK_COMMENT_CHARS = 2000
|
||||
|
||||
# The rating waits on Langfuse, so an unresponsive Langfuse must not hold the
|
||||
# click for long; the rating is saved regardless.
|
||||
_LANGFUSE_TIMEOUT_SECONDS = 5
|
||||
|
||||
# Sequences are Postgres ``int4``: more digits than this could overflow the
|
||||
# lookup, and no chat is that long.
|
||||
_MAX_SEQUENCE_DIGITS = 9
|
||||
@@ -63,16 +40,8 @@ class MessageFeedbackResponse(BaseModel):
|
||||
"""A saved rating of an AutoPilot reply."""
|
||||
|
||||
id: str = Field(
|
||||
description="The rating's id, which its Langfuse score shares.",
|
||||
)
|
||||
langfuse_target: LangfuseScoreTarget | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Where the Langfuse score went: the trace of the turn that wrote "
|
||||
"the reply, or the chat's Langfuse session when that trace is not "
|
||||
"known. Null when Langfuse is not configured or did not take the "
|
||||
"score; the rating is saved either way."
|
||||
),
|
||||
description="The rating's id, which its PostHog event carries as "
|
||||
"``feedback_id``.",
|
||||
)
|
||||
|
||||
|
||||
@@ -86,8 +55,7 @@ async def record_message_feedback(
|
||||
comment: str | None,
|
||||
impersonated: bool = False,
|
||||
) -> MessageFeedbackResponse | None:
|
||||
"""Save a rating of one reply in *user_id*'s chat and score it in Langfuse
|
||||
and PostHog.
|
||||
"""Save a rating of one reply in *user_id*'s chat and send it to PostHog.
|
||||
|
||||
*message_id* is the id the chat UI knows the reply by: ``<session>-seq-<N>``
|
||||
for a reply loaded from the chat history, or the message row's own id.
|
||||
@@ -106,7 +74,6 @@ async def record_message_feedback(
|
||||
score_name=score_name,
|
||||
score_value=score_value,
|
||||
comment=comment,
|
||||
langfuse_trace_id=message.langfuseTraceId,
|
||||
)
|
||||
_send_to_posthog(
|
||||
user_id=user_id,
|
||||
@@ -118,15 +85,7 @@ async def record_message_feedback(
|
||||
comment=comment,
|
||||
impersonated=impersonated,
|
||||
)
|
||||
target = await _score_in_langfuse(
|
||||
feedback_id=feedback_id,
|
||||
session_id=session_id,
|
||||
message=message,
|
||||
score_name=score_name,
|
||||
score_value=score_value,
|
||||
comment=comment,
|
||||
)
|
||||
return MessageFeedbackResponse(id=feedback_id, langfuse_target=target)
|
||||
return MessageFeedbackResponse(id=feedback_id)
|
||||
|
||||
|
||||
def sequence_from_ui_message_id(session_id: str, message_id: str) -> int | None:
|
||||
@@ -157,52 +116,6 @@ async def _find_rated_message(
|
||||
)
|
||||
|
||||
|
||||
async def _score_in_langfuse(
|
||||
*,
|
||||
feedback_id: str,
|
||||
session_id: str,
|
||||
message: PrismaChatMessage,
|
||||
score_name: FeedbackScoreName,
|
||||
score_value: FeedbackScoreValue,
|
||||
comment: str | None,
|
||||
) -> LangfuseScoreTarget | None:
|
||||
"""Send the rating to Langfuse and say where it landed, or None."""
|
||||
if not (
|
||||
settings.secrets.langfuse_public_key and settings.secrets.langfuse_secret_key
|
||||
):
|
||||
return None
|
||||
trace_id = message.langfuseTraceId
|
||||
request = CreateScoreRequest(
|
||||
id=feedback_id,
|
||||
# A score belongs to a trace or to a session, not both.
|
||||
traceId=trace_id,
|
||||
sessionId=None if trace_id else session_id,
|
||||
name=score_name,
|
||||
value=score_value,
|
||||
dataType=ScoreDataType.NUMERIC,
|
||||
comment=comment,
|
||||
metadata={
|
||||
"message_id": message.id,
|
||||
"message_sequence": message.sequence,
|
||||
"session_id": session_id,
|
||||
},
|
||||
environment=settings.secrets.langfuse_tracing_environment,
|
||||
)
|
||||
try:
|
||||
await get_client().async_api.score.create(
|
||||
request=request,
|
||||
request_options={"timeout_in_seconds": _LANGFUSE_TIMEOUT_SECONDS},
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
f"Langfuse did not take {score_name} score {feedback_id} for "
|
||||
f"session {session_id}; the rating is saved in Postgres only",
|
||||
exc_info=True,
|
||||
)
|
||||
return None
|
||||
return "trace" if trace_id else "session"
|
||||
|
||||
|
||||
def _send_to_posthog(
|
||||
*,
|
||||
user_id: str,
|
||||
|
||||
@@ -59,12 +59,11 @@ async def upsert_message_feedback(
|
||||
score_name: str,
|
||||
score_value: int,
|
||||
comment: str | None,
|
||||
langfuse_trace_id: str | None,
|
||||
) -> str:
|
||||
"""Save *user_id*'s rating of a reply and return the rating's id.
|
||||
|
||||
One row per user, message and score name: rating again overwrites the
|
||||
value and comment and keeps the id, which is also the Langfuse score id.
|
||||
value and comment and keeps the id.
|
||||
The caller must already have checked that the message is in a session
|
||||
*user_id* owns (see ``get_rateable_message``).
|
||||
"""
|
||||
@@ -83,12 +82,10 @@ async def upsert_message_feedback(
|
||||
"scoreName": score_name,
|
||||
"scoreValue": score_value,
|
||||
"comment": comment,
|
||||
"langfuseTraceId": langfuse_trace_id,
|
||||
}
|
||||
update: ChatMessageFeedbackUpdateInput = {
|
||||
"scoreValue": score_value,
|
||||
"comment": comment,
|
||||
"langfuseTraceId": langfuse_trace_id,
|
||||
}
|
||||
data: ChatMessageFeedbackUpsertInput = {"create": create, "update": update}
|
||||
try:
|
||||
|
||||
@@ -71,7 +71,6 @@ async def test_rating_again_updates_the_same_row(setup_test_user, test_user_id):
|
||||
score_name="user-feedback",
|
||||
score_value=0,
|
||||
comment="Wrong numbers",
|
||||
langfuse_trace_id=None,
|
||||
)
|
||||
second = await upsert_message_feedback(
|
||||
user_id=test_user_id,
|
||||
@@ -80,7 +79,6 @@ async def test_rating_again_updates_the_same_row(setup_test_user, test_user_id):
|
||||
score_name="user-feedback",
|
||||
score_value=1,
|
||||
comment=None,
|
||||
langfuse_trace_id="1edf31f11b1693cc6103f358c1481694",
|
||||
)
|
||||
|
||||
assert second == first
|
||||
@@ -88,7 +86,6 @@ async def test_rating_again_updates_the_same_row(setup_test_user, test_user_id):
|
||||
assert len(rows) == 1
|
||||
assert rows[0].scoreValue == 1
|
||||
assert rows[0].comment is None
|
||||
assert rows[0].langfuseTraceId == "1edf31f11b1693cc6103f358c1481694"
|
||||
assert rows[0].sessionId == session_id
|
||||
assert rows[0].userId == test_user_id
|
||||
|
||||
@@ -104,7 +101,6 @@ async def test_each_score_name_keeps_its_own_rating(setup_test_user, test_user_i
|
||||
score_name="user-feedback",
|
||||
score_value=1,
|
||||
comment=None,
|
||||
langfuse_trace_id=None,
|
||||
)
|
||||
copy = await upsert_message_feedback(
|
||||
user_id=test_user_id,
|
||||
@@ -113,7 +109,6 @@ async def test_each_score_name_keeps_its_own_rating(setup_test_user, test_user_i
|
||||
score_name="copy",
|
||||
score_value=1,
|
||||
comment=None,
|
||||
langfuse_trace_id=None,
|
||||
)
|
||||
|
||||
assert vote != copy
|
||||
@@ -132,7 +127,6 @@ async def test_comment_is_stripped_of_postgres_control_characters(
|
||||
score_name="user-feedback",
|
||||
score_value=0,
|
||||
comment="bad\x00 answer",
|
||||
langfuse_trace_id=None,
|
||||
)
|
||||
|
||||
row = await ChatMessageFeedback.prisma().find_unique(where={"id": feedback_id})
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
sending a rating to PostHog LLM analytics.
|
||||
|
||||
The route tests (``backend/api/features/chat/feedback_test.py``) cover the
|
||||
Postgres and Langfuse halves.
|
||||
Postgres half.
|
||||
"""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
@@ -35,7 +35,6 @@ def reply(mocker: pytest_mock.MockerFixture) -> PrismaChatMessage:
|
||||
role="assistant",
|
||||
content="Here is the report.",
|
||||
sequence=7,
|
||||
langfuseTraceId="1edf31f11b1693cc6103f358c1481694",
|
||||
)
|
||||
mocker.patch.object(
|
||||
feedback_db, "get_rateable_message", new=AsyncMock(return_value=message)
|
||||
@@ -48,17 +47,6 @@ def reply(mocker: pytest_mock.MockerFixture) -> PrismaChatMessage:
|
||||
return message
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def langfuse_score(mocker: pytest_mock.MockerFixture) -> AsyncMock:
|
||||
secrets = feedback.settings.secrets
|
||||
mocker.patch.object(secrets, "langfuse_public_key", "pk-lf-test")
|
||||
mocker.patch.object(secrets, "langfuse_secret_key", "sk-lf-test")
|
||||
langfuse = MagicMock()
|
||||
langfuse.async_api.score.create = AsyncMock()
|
||||
mocker.patch.object(feedback, "get_client", return_value=langfuse)
|
||||
return langfuse.async_api.score.create
|
||||
|
||||
|
||||
async def _rate(**overrides):
|
||||
kwargs = {
|
||||
"user_id": USER_ID,
|
||||
@@ -75,7 +63,7 @@ async def _rate(**overrides):
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("reply")
|
||||
async def test_thumbs_down_is_an_ai_metric_on_the_user_and_chat(
|
||||
mocker: pytest_mock.MockerFixture, langfuse_score: AsyncMock
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
) -> None:
|
||||
capture = mocker.patch.object(feedback, "capture")
|
||||
|
||||
@@ -97,12 +85,11 @@ async def test_thumbs_down_is_an_ai_metric_on_the_user_and_chat(
|
||||
"$ai_feedback_text": "It edited the wrong file",
|
||||
},
|
||||
)
|
||||
langfuse_score.assert_awaited_once()
|
||||
assert response.langfuse_target == "trace"
|
||||
assert response.id == FEEDBACK_ID
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("reply", "langfuse_score")
|
||||
@pytest.mark.usefixtures("reply")
|
||||
async def test_copy_carries_no_feedback_text(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
) -> None:
|
||||
@@ -117,7 +104,7 @@ async def test_copy_carries_no_feedback_text(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("reply", "langfuse_score")
|
||||
@pytest.mark.usefixtures("reply")
|
||||
async def test_impersonated_rating_is_marked(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
) -> None:
|
||||
@@ -129,7 +116,6 @@ async def test_impersonated_rating_is_marked(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("langfuse_score")
|
||||
async def test_a_rating_names_the_trace_of_the_turn_that_wrote_the_reply(
|
||||
mocker: pytest_mock.MockerFixture, reply: PrismaChatMessage
|
||||
) -> None:
|
||||
@@ -142,7 +128,7 @@ async def test_a_rating_names_the_trace_of_the_turn_that_wrote_the_reply(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("reply", "langfuse_score")
|
||||
@pytest.mark.usefixtures("reply")
|
||||
async def test_a_reply_with_no_recorded_trace_sends_no_trace_id(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
) -> None:
|
||||
@@ -156,8 +142,8 @@ async def test_a_reply_with_no_recorded_trace_sends_no_trace_id(
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("reply")
|
||||
async def test_posthog_failure_keeps_the_rating_and_the_langfuse_score(
|
||||
mocker: pytest_mock.MockerFixture, langfuse_score: AsyncMock
|
||||
async def test_posthog_failure_keeps_the_rating(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
) -> None:
|
||||
"""The real ``capture`` runs against a client that raises."""
|
||||
client = MagicMock()
|
||||
@@ -170,7 +156,6 @@ async def test_posthog_failure_keeps_the_rating_and_the_langfuse_score(
|
||||
assert client.capture.call_args.kwargs["distinct_id"] == USER_ID
|
||||
assert response is not None
|
||||
assert response.id == FEEDBACK_ID
|
||||
langfuse_score.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -43,8 +43,8 @@ Every step is idempotent, so a re-run only picks up what is left.
|
||||
What it does not reach:
|
||||
|
||||
- Text derived from the block before it was removed: compaction summaries,
|
||||
the assistant's own replies and tool results in the session, and Langfuse
|
||||
traces.
|
||||
the assistant's own replies and tool results in the session, and legacy
|
||||
Langfuse traces (no longer produced; SECRT-2711).
|
||||
- A request that loaded an idle session before its row was written and saves
|
||||
it after the run's final eviction: it caches the old first message again,
|
||||
until the entry expires or is evicted. Model input still reads it without
|
||||
|
||||
@@ -1364,11 +1364,21 @@ async def test_a_sandbox_files_lines_are_judged_and_an_empty_grep_is_not(
|
||||
|
||||
|
||||
async def _agent_without_credentials() -> ToolResponseBase | None:
|
||||
from backend.copilot.tools.expert_scope import CredentialScopeSnapshot
|
||||
from backend.copilot.tools.run_agent import RunAgentInput, RunAgentTool
|
||||
|
||||
graph = MagicMock(id="graph-1", version=1, input_schema={})
|
||||
graph = MagicMock(
|
||||
id="graph-1",
|
||||
version=1,
|
||||
input_schema={},
|
||||
regular_credentials_inputs={"credentials": MagicMock()},
|
||||
)
|
||||
graph.name = "Daily Scraper"
|
||||
with (
|
||||
patch(
|
||||
"backend.copilot.tools.run_agent.resolve_credential_scope",
|
||||
AsyncMock(return_value=CredentialScopeSnapshot([], None)),
|
||||
),
|
||||
patch(
|
||||
"backend.copilot.tools.run_agent.match_user_credentials_to_graph",
|
||||
AsyncMock(return_value=({}, ["credentials"])),
|
||||
|
||||
@@ -1,13 +1,11 @@
|
||||
"""Keep Google user data out of the LLM traces sent to Langfuse.
|
||||
"""Tell Google user data apart, and redact it from anything traced.
|
||||
|
||||
Google's Limited Use policy bars staff from reading what the Gmail, Drive,
|
||||
Docs, Sheets and Calendar blocks return, so neither their output nor the
|
||||
arguments AutoPilot passed them may reach a trace. Every traced payload goes
|
||||
through :func:`redact_google_data`: the Langfuse SDK applies it as its
|
||||
``mask`` (baseline turns and every call through the Langfuse-wrapped OpenAI
|
||||
client), and LangSmith's Claude Agent SDK tracing, which exports SDK turns to
|
||||
Langfuse over OTLP, applies it as the client's ``hide_inputs`` and
|
||||
``hide_outputs``.
|
||||
arguments AutoPilot passed them may reach a trace. :func:`redact_google_data`
|
||||
is the redaction for a traced payload. Nothing in the backend applies it: the
|
||||
backend no longer traces to Langfuse (SECRT-2711), and the PostHog spans it
|
||||
sends (``copilot/ai_spans.py``) carry no tool input or output.
|
||||
|
||||
Data counts as Google's by the block that produced it, never by what it says:
|
||||
|
||||
@@ -44,15 +42,12 @@ from collections.abc import Iterator
|
||||
from contextvars import ContextVar
|
||||
from typing import Any
|
||||
|
||||
import langsmith
|
||||
from langfuse import Langfuse
|
||||
from pydantic import BaseModel
|
||||
|
||||
from backend.blocks import get_blocks
|
||||
from backend.copilot.capabilities.block_meta import get_block_provider
|
||||
from backend.copilot.constants import COPILOT_NODE_PREFIX, parse_node_id_from_exec_id
|
||||
from backend.integrations.providers import ProviderName
|
||||
from backend.util.settings import Settings
|
||||
|
||||
REDACTED = "[redacted: Google data]"
|
||||
_IMAGE_REDACTED = "[redacted: image data]"
|
||||
@@ -102,50 +97,6 @@ def tracing_google_data(active: bool = True) -> Iterator[None]:
|
||||
_redacting_every_text.reset(token)
|
||||
|
||||
|
||||
def mask_langfuse_data(*, data: Any, **kwargs: Any) -> Any:
|
||||
"""The Langfuse SDK's ``mask`` hook."""
|
||||
return redact_google_data(data)
|
||||
|
||||
|
||||
def install_trace_redaction() -> None:
|
||||
"""Apply the redaction to both tracers.
|
||||
|
||||
Langfuse first: its client keeps the settings of the first one built in
|
||||
the process, and building it sets the OpenTelemetry tracer provider the
|
||||
LangSmith client then exports through, as it would have lazily.
|
||||
"""
|
||||
install_langfuse_mask()
|
||||
langsmith.configure(client=langsmith.Client(hide_inputs=_hide, hide_outputs=_hide))
|
||||
|
||||
|
||||
def install_langfuse_mask() -> None:
|
||||
"""Build the process's Langfuse client with the mask.
|
||||
|
||||
Has to run before anything calls ``langfuse.get_client()``, which hands
|
||||
out copies of the first client built: the SDK keeps that one's settings
|
||||
and drops a mask passed later. Raises when the client it gets back does
|
||||
not mask, so the caller can leave tracing off rather than trace unmasked.
|
||||
"""
|
||||
# The configured credentials, which ``.env`` puts in ``Settings`` but not
|
||||
# in the environment the SDK would otherwise read them from. ``host``
|
||||
# still yields to ``LANGFUSE_BASE_URL``, as it would by default.
|
||||
secrets = Settings().secrets
|
||||
client = Langfuse(
|
||||
public_key=secrets.langfuse_public_key or None,
|
||||
secret_key=secrets.langfuse_secret_key or None,
|
||||
host=secrets.langfuse_host or None,
|
||||
mask=mask_langfuse_data,
|
||||
)
|
||||
# ``_resources`` is None when the client is disabled (no keys): it sends
|
||||
# nothing, so there is nothing to mask.
|
||||
resources = client._resources
|
||||
if resources is not None and resources.mask is not mask_langfuse_data:
|
||||
raise RuntimeError(
|
||||
"A Langfuse client was built before the Google data mask; "
|
||||
"its traces would not be redacted"
|
||||
)
|
||||
|
||||
|
||||
def remember_google_output(output: str) -> None:
|
||||
"""Record a tool output built from a Google block's result, so it is
|
||||
redacted wherever it is traced even when it no longer parses."""
|
||||
@@ -210,10 +161,6 @@ def _reviewed_block_id(review_id: object) -> str | None:
|
||||
return parse_node_id_from_exec_id(review_id).removeprefix(COPILOT_NODE_PREFIX)
|
||||
|
||||
|
||||
def _hide(data: dict) -> dict:
|
||||
return redact_google_data(data)
|
||||
|
||||
|
||||
def _google_call_ids(data: Any) -> set[str]:
|
||||
"""Ids of the tool calls in *data* that ran a Google block: by the block
|
||||
the call names, or by the Google data its result carries."""
|
||||
@@ -321,7 +268,7 @@ def _call_arguments(node: dict[str, Any]) -> dict[str, Any] | None:
|
||||
|
||||
def _answered_call_id(node: dict[str, Any]) -> str | None:
|
||||
"""The call a tool result answers: Anthropic ``tool_use_id`` or OpenAI
|
||||
(and LangSmith history) ``tool_call_id``."""
|
||||
``tool_call_id``."""
|
||||
answered = node.get("tool_use_id") or node.get("tool_call_id")
|
||||
return answered if isinstance(answered, str) else None
|
||||
|
||||
|
||||
@@ -1,17 +1,10 @@
|
||||
"""What reaches Langfuse must not carry what a Google block returned."""
|
||||
"""A traced payload must not carry what a Google block returned."""
|
||||
|
||||
import json
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import patch
|
||||
|
||||
import langsmith
|
||||
import pytest
|
||||
from langfuse import Langfuse, get_client
|
||||
from langfuse._client.resource_manager import LangfuseResourceManager
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
|
||||
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
||||
|
||||
from backend.blocks.google.gmail import GmailReadBlock
|
||||
from backend.blocks.text import MatchTextPatternBlock
|
||||
@@ -21,10 +14,7 @@ from backend.copilot.tools.models import BlockOutputResponse
|
||||
from . import google_data_redaction
|
||||
from .google_data_redaction import (
|
||||
REDACTED,
|
||||
install_langfuse_mask,
|
||||
install_trace_redaction,
|
||||
is_google_output,
|
||||
mask_langfuse_data,
|
||||
redact_google_data,
|
||||
remember_google_output,
|
||||
tracing_google_data,
|
||||
@@ -50,7 +40,7 @@ TEXT_OUTPUT = _output(TEXT_BLOCK_ID, None, "a harmless match")
|
||||
|
||||
|
||||
def _sdk_history(call: dict[str, Any], result: str) -> dict[str, Any]:
|
||||
"""An SDK turn's model input, as LangSmith records it."""
|
||||
"""An SDK turn's model input, in the Anthropic shape."""
|
||||
return {
|
||||
"messages": [
|
||||
{"role": "user", "content": "what did finance send me?"},
|
||||
@@ -71,7 +61,7 @@ def _sdk_history(call: dict[str, Any], result: str) -> dict[str, Any]:
|
||||
|
||||
|
||||
def _baseline_history(call: dict[str, Any], result: str) -> list[dict[str, Any]]:
|
||||
"""A baseline turn's model input, as the Langfuse OpenAI wrapper records it."""
|
||||
"""A baseline turn's model input, in the OpenAI shape."""
|
||||
return [
|
||||
{"role": "user", "content": "what did finance send me?"},
|
||||
{
|
||||
@@ -158,7 +148,7 @@ def test_google_call_without_its_result_redacts_the_result():
|
||||
|
||||
|
||||
def test_tool_run_output_redacted_by_what_the_tool_produced():
|
||||
"""LangSmith records a tool run's output without its call."""
|
||||
"""A tool run's output traced without its call."""
|
||||
preview = f"<tool-output-truncated total_chars=90000>\n{EMAIL[:20]}"
|
||||
remember_google_output(preview)
|
||||
|
||||
@@ -193,8 +183,8 @@ def test_tool_run_input_redacted():
|
||||
],
|
||||
)
|
||||
def test_resumed_google_tool_run_input_redacted(call: dict[str, Any]):
|
||||
"""A LangSmith tool run records the call without its result; the review
|
||||
id names the block it resumes."""
|
||||
"""A tool run's call traced without its result; the review id names the
|
||||
block it resumes."""
|
||||
redacted = redact_google_data({"input": call})
|
||||
|
||||
assert "the new email" not in json.dumps(redacted)
|
||||
@@ -220,142 +210,17 @@ def test_google_block_output_response_is_remembered():
|
||||
assert response.model_copy(update={"provider": "google"}).holds_google_data
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def langfuse_spans() -> Iterator[InMemorySpanExporter]:
|
||||
"""A Langfuse client built the way production builds it, exporting to
|
||||
memory instead of to Langfuse."""
|
||||
exporter = InMemorySpanExporter()
|
||||
provider = TracerProvider()
|
||||
provider.add_span_processor(SimpleSpanProcessor(exporter))
|
||||
LangfuseResourceManager.reset()
|
||||
Langfuse(
|
||||
public_key="pk-lf-test",
|
||||
secret_key="sk-lf-test",
|
||||
base_url="http://127.0.0.1:9",
|
||||
tracer_provider=provider,
|
||||
mask=mask_langfuse_data,
|
||||
)
|
||||
try:
|
||||
yield exporter
|
||||
finally:
|
||||
LangfuseResourceManager.reset()
|
||||
|
||||
|
||||
def test_langfuse_generation_input_is_masked(langfuse_spans: InMemorySpanExporter):
|
||||
call = {"block_id": GMAIL_BLOCK_ID, "input_data": {"query": "from:finance"}}
|
||||
history = _baseline_history(call, GMAIL_OUTPUT)
|
||||
client = get_client()
|
||||
|
||||
generation = client.start_observation(
|
||||
name="OpenAI-generation", as_type="generation", input=history
|
||||
)
|
||||
generation.update(output={"role": "assistant", "content": "Done."})
|
||||
generation.end()
|
||||
|
||||
sent = json.dumps(
|
||||
[dict(span.attributes or {}) for span in langfuse_spans.get_finished_spans()]
|
||||
)
|
||||
assert EMAIL not in sent
|
||||
assert "from:finance" not in sent
|
||||
assert REDACTED in sent
|
||||
assert "run_block" in sent
|
||||
|
||||
|
||||
def _build_langfuse(monkeypatch: pytest.MonkeyPatch, **kwargs: Any) -> None:
|
||||
monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-lf-test")
|
||||
monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-lf-test")
|
||||
monkeypatch.setenv("LANGFUSE_BASE_URL", "http://127.0.0.1:9")
|
||||
Langfuse(tracer_provider=TracerProvider(), **kwargs)
|
||||
|
||||
|
||||
def test_install_langfuse_mask_refuses_a_client_built_without_it(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""The SDK hands out the first client built; a mask passed later is
|
||||
dropped, so installing it must fail rather than trace unmasked."""
|
||||
LangfuseResourceManager.reset()
|
||||
try:
|
||||
_build_langfuse(monkeypatch)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
install_langfuse_mask()
|
||||
finally:
|
||||
LangfuseResourceManager.reset()
|
||||
|
||||
|
||||
def test_install_langfuse_mask_accepts_a_client_built_with_it(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
LangfuseResourceManager.reset()
|
||||
try:
|
||||
_build_langfuse(monkeypatch, mask=mask_langfuse_data)
|
||||
|
||||
install_langfuse_mask()
|
||||
finally:
|
||||
LangfuseResourceManager.reset()
|
||||
|
||||
|
||||
def test_install_langfuse_mask_uses_the_configured_credentials():
|
||||
"""Credentials read from ``.env`` reach ``Settings`` but not the process
|
||||
environment the SDK falls back to; without them the first client, which
|
||||
every later ``get_client()`` reuses, would be disabled."""
|
||||
secrets = MagicMock(
|
||||
langfuse_public_key="pk-from-settings",
|
||||
langfuse_secret_key="sk-from-settings",
|
||||
langfuse_host="https://langfuse.example.com",
|
||||
)
|
||||
with (
|
||||
patch.object(google_data_redaction, "Settings") as settings,
|
||||
patch.object(google_data_redaction, "Langfuse") as langfuse,
|
||||
):
|
||||
settings.return_value.secrets = secrets
|
||||
langfuse.return_value._resources.mask = mask_langfuse_data
|
||||
install_langfuse_mask()
|
||||
|
||||
langfuse.assert_called_once_with(
|
||||
public_key="pk-from-settings",
|
||||
secret_key="sk-from-settings",
|
||||
host="https://langfuse.example.com",
|
||||
mask=mask_langfuse_data,
|
||||
)
|
||||
|
||||
|
||||
def test_install_trace_redaction_masks_both_tracers():
|
||||
with (
|
||||
patch.object(google_data_redaction, "Langfuse") as langfuse,
|
||||
patch.object(google_data_redaction.langsmith, "configure") as configure,
|
||||
):
|
||||
langfuse.return_value._resources.mask = mask_langfuse_data
|
||||
install_trace_redaction()
|
||||
|
||||
langfuse.assert_called_once()
|
||||
assert langfuse.call_args.kwargs["mask"] is mask_langfuse_data
|
||||
client = configure.call_args.kwargs["client"]
|
||||
assert isinstance(client, langsmith.Client)
|
||||
inputs = {"input": {"block_id": GMAIL_BLOCK_ID, "input_data": {"q": "x"}}}
|
||||
assert client._hide_run_inputs(inputs)["input"]["input_data"] == REDACTED
|
||||
assert client._hide_run_outputs({"output": GMAIL_OUTPUT}) == {"output": REDACTED}
|
||||
|
||||
|
||||
def test_tracing_google_data_redacts_every_text(langfuse_spans: InMemorySpanExporter):
|
||||
def test_tracing_google_data_redacts_every_text():
|
||||
"""A gate prompt quoting a Google result has no structure to find it by."""
|
||||
prompt = [{"role": "user", "content": f"Judge this read:\n{EMAIL}"}]
|
||||
client = get_client()
|
||||
|
||||
# An inner scope that is not a Google call keeps the outer one's redaction.
|
||||
with tracing_google_data(), tracing_google_data(False):
|
||||
generation = client.start_observation(
|
||||
name="OpenAI-generation", as_type="generation", input=prompt
|
||||
)
|
||||
generation.update(output={"role": "assistant", "content": EMAIL})
|
||||
generation.end()
|
||||
traced = json.dumps(redact_google_data(prompt))
|
||||
assert redact_google_data(prompt) == prompt
|
||||
|
||||
sent = json.dumps(
|
||||
[dict(span.attributes or {}) for span in langfuse_spans.get_finished_spans()]
|
||||
)
|
||||
assert "quarterly numbers" not in sent
|
||||
assert REDACTED in sent
|
||||
assert "quarterly numbers" not in traced
|
||||
assert REDACTED in traced
|
||||
|
||||
|
||||
def test_a_non_string_id_is_left_alone():
|
||||
|
||||
@@ -25,7 +25,7 @@ Runtime prompt instructions live in `prompting.py:get_graphiti_supplement()`.
|
||||
- Recall goes through `recall.py` (what may be read back) and `recall_render.py` (how it is written out). Read facts with `search_facts` and episodes with `recent_episodes`; what the assistant is shown is read again by uuid as the last graph read before rendering (`search_facts` ends with `live_now`, and warm context and `memory_search`, which hold both lists while the slower read finishes, call `recall_recheck.recheck`). Ingestion hands graphiti `recall_ingest.previous_episode_uuids`. Code with its own Cypher uses the predicates there: `live_fact_predicate` for facts (the settings and admin fact lists, the dream's demotion writers), `forgotten_facts_clause` with `recallable_episode_predicate` for episodes (the dream gather and its session filter, the forget's redaction).
|
||||
- A fact is live while `expired_at` and `forgotten_at` are unset and `status` is `active`, `tentative` or unset. It is forgotten once a forget has stamped `forgotten_at`. The marker is written only by a forget (`recall_forget.py` and the cascade it runs on what the dream derived from the fact, `recall_cascade.py`; the legacy-forget backfill and the derivation backfill's cascade for older forgets), and no ingestion, dream demotion or other status change can make a forgotten fact look remembered, nor a remembered one look forgotten: a `status` of `retracted` or an `expiration_reason` of `user_signal` is no forget without the marker, since graphiti's attribute extraction fills both from an episode's text and a dream demotion picks its reason. The extraction schema (`types.MemoryFact`) leaves both out, and an extracted `retracted` is put back to `active` as the episode is written (`ingest._undo_extracted_retractions`). Nor can extraction write the marker: graphiti-core 0.30.2 saves every key its attribute model answers, so `recall_ingest.ForgetAwareLLMClient` keeps only what the edge or entity type declares (`types.extracted_attributes`), and no extraction model declares the marker or the audit copies (`types.POLICY_EDGE_FIELDS`). Forgets made before the marker are recognised by the pre-policy shape (`expired_at` set with no `invalid_at` and no reason). `migrations/backfill_legacy_forgets.py` restamps that shape, after which its clause can go. It restamps each graph holding the graph's write lock (below), so an ingestion running meanwhile cannot save an older copy of an edge over the restamp; a graph whose lock another writer keeps past 60 s, or that cannot be locked because Redis is unreachable, is skipped with nothing written and counted busy, a graph whose backfill raises is counted failed, and either makes the script exit 1 so it is run again. Drop the clause only after an `--apply` run that exits 0 everywhere (a dry run exits 0 too, having written nothing). An episode is recallable while it has no `redacted_at`, none of its `entity_edges` is forgotten and it is not `write_pending` (a dream write's episode graphiti has not saved yet, below). `recall.forgotten_fact_predicate` is never null (it only asks whether properties are set), so `NOT` it is the not-forgotten test on any edge, a live one included.
|
||||
- What a forget guarantees: the fact and every episode it came from are hidden from warm context (a session's first turn and every follow-up turn's refresh, `context.py` and `context_refresh.py`; no turn stores its block, `context_marker.py`), `memory_search`, normal `memory_forget_search` calls (without `include_forgotten`), the settings fact list and counts, the dream's input (its episodes and facts, and every chat session a hidden episode came from, left out whole) and ingestion's extraction context (the earlier episodes graphiti shows its prompts). So is every fact the dream derived from it, and the dream text that derived it, level by level (the cascade, below). A read already under way when the forget answers does not show it either, once its last check has run: right before rendering, warm context and `memory_search` check every fact and episode they are about to show again, by uuid, in one statement (`recall_recheck.recheck`), and `memory_forget_search` its facts (`search_facts` ends with `live_now`). Every item shown passed that check, its own last graph read. A forget writes each fact's marker before anything else, so a check that began after the marker committed drops the fact and every episode citing it; a response whose check began earlier can still carry them while it is delivered (rendering, scheduling, network), and no later read will. The sentence also leaves what graphiti keeps for its own prompts (`recall_hide.py`). The edge's `fact` and `name` read `[forgotten]` (`recall.FORGOTTEN_FACT`), the originals moved to `fact_redacted` / `name_redacted`; no edge type has that name, so graphiti never runs its attribute prompt, which lists an edge's stored properties, on a forgotten edge. Every entity the fact joins, and every entity an episode citing it mentions, loses its `summary` and every property but its identity (uuid, name, group_id, labels, created_at, name_embedding): typed attributes such as `Person.role` and any `attributes` map go. The `summary` of each community one of them belongs to is blanked too. They are blanked, not rewritten: graphiti extracts them again from the next episode that mentions the entity, and the weekly rebuild restores the communities.
|
||||
- What it does not: admin audit surfaces (status `any`, the graph view) keep showing retracted facts, with the text from `fact_redacted` (blank for a derived fact a hard forget erased), and redacted episodes, on purpose, and show a tombstone's title and `hard_deleted_at` (never its content). `memory_forget_search` with `include_forgotten` returns only redacted placeholders and metadata for forgotten facts, so a soft forget can be made permanent with a hard one, but never what they said: each is `[forgotten]` with its uuid and when it was stored and forgotten (each read as a timestamp and written out anew, else `unknown`), without its text, audit copy, scope or valid dates. `recall.search_forgotten_facts` finds them by the embedding a forget keeps, ranking the forgotten facts alone, so no live fact takes their slots. Not scrubbed: community names, the fact's embedding vector (a hard forget removes a derived fact's), and a redacted episode's name and source description (a hard forget trims a dream episode's to its kind). A hard forget keeps, redacted, the text of an episode another edge still cites, and a tombstone keeps its name (for a stored memory, the model's short title). Two windows have no bound. When Redis cannot be reached, forgets and ingestion go ahead without the write lock (below); and a holder that loses its lease (its renewals fail, or its event loop stalls past the five-minute expiry) keeps writing while another writer takes the lock, since the renewal only warns and nothing fences the graph write. Either way an ingestion that read a fact before a forget of it can save its older copy over the forget. On a soft-forgotten fact the driver saves nothing but its episode list (`fact_attributes.py`): the fact stays forgotten, scrubbed and out of recall, and a later statement of it is a new fact. But graphiti's `add_episode` saves the fact's entities whole, so they get back the summaries and attributes the forget blanked, and a hard forget deletes the edge, so the save writes the fact again, live, with its sentence. Either stays until the fact is forgotten again: Redis coming back repairs nothing. A marker older builds left on unscrubbed text shows up in `migrations/backfill_legacy_forgets.py --scan-markers`, and forgetting that fact again scrubs it. And order is by when an episode is written, not when it was said, by design: an episode still queued when a forget lands is written after it, and if it states the fact again it teaches it again, as a message sent after the forget would. The weekly community rebuild does not take the lock either, so one running while a forget lands can write a community summary from entity summaries the forget then blanks. Nor does the settings page's erase-all, which deletes the whole graph: an ingestion running meanwhile can save into the emptied graph what it read before the erase. And a forget does not reach what a chat session already holds besides warm context, a boundary proposed as a known limit that still needs the owner's acceptance: the assistant's own earlier replies and tool results (a `memory_search` result, say) in a session that saw the fact keep it, and so do a compaction summary written from a turn's input, its block included (the CLI's own, or the history compression's), Langfuse traces of the turns that sent it, and, within one turn, the block that turn already sent its provider, which the model reads until the turn ends whatever a `memory_forget` in it does, and a building-mode restart after such a forget: it resumes the native CLI session, which still holds the turn's query and its block until the upload scrub runs at the end of the turn.
|
||||
- What it does not: admin audit surfaces (status `any`, the graph view) keep showing retracted facts, with the text from `fact_redacted` (blank for a derived fact a hard forget erased), and redacted episodes, on purpose, and show a tombstone's title and `hard_deleted_at` (never its content). `memory_forget_search` with `include_forgotten` returns only redacted placeholders and metadata for forgotten facts, so a soft forget can be made permanent with a hard one, but never what they said: each is `[forgotten]` with its uuid and when it was stored and forgotten (each read as a timestamp and written out anew, else `unknown`), without its text, audit copy, scope or valid dates. `recall.search_forgotten_facts` finds them by the embedding a forget keeps, ranking the forgotten facts alone, so no live fact takes their slots. Not scrubbed: community names, the fact's embedding vector (a hard forget removes a derived fact's), and a redacted episode's name and source description (a hard forget trims a dream episode's to its kind). A hard forget keeps, redacted, the text of an episode another edge still cites, and a tombstone keeps its name (for a stored memory, the model's short title). Two windows have no bound. When Redis cannot be reached, forgets and ingestion go ahead without the write lock (below); and a holder that loses its lease (its renewals fail, or its event loop stalls past the five-minute expiry) keeps writing while another writer takes the lock, since the renewal only warns and nothing fences the graph write. Either way an ingestion that read a fact before a forget of it can save its older copy over the forget. On a soft-forgotten fact the driver saves nothing but its episode list (`fact_attributes.py`): the fact stays forgotten, scrubbed and out of recall, and a later statement of it is a new fact. But graphiti's `add_episode` saves the fact's entities whole, so they get back the summaries and attributes the forget blanked, and a hard forget deletes the edge, so the save writes the fact again, live, with its sentence. Either stays until the fact is forgotten again: Redis coming back repairs nothing. A marker older builds left on unscrubbed text shows up in `migrations/backfill_legacy_forgets.py --scan-markers`, and forgetting that fact again scrubs it. And order is by when an episode is written, not when it was said, by design: an episode still queued when a forget lands is written after it, and if it states the fact again it teaches it again, as a message sent after the forget would. The weekly community rebuild does not take the lock either, so one running while a forget lands can write a community summary from entity summaries the forget then blanks. Nor does the settings page's erase-all, which deletes the whole graph: an ingestion running meanwhile can save into the emptied graph what it read before the erase. And a forget does not reach what a chat session already holds besides warm context, a boundary proposed as a known limit that still needs the owner's acceptance: the assistant's own earlier replies and tool results (a `memory_search` result, say) in a session that saw the fact keep it, and so do a compaction summary written from a turn's input, its block included (the CLI's own, or the history compression's), Langfuse traces of the turns that sent it (recorded until Langfuse tracing was removed), and, within one turn, the block that turn already sent its provider, which the model reads until the turn ends whatever a `memory_forget` in it does, and a building-mode restart after such a forget: it resumes the native CLI session, which still holds the turn's query and its block until the upload scrub runs at the end of the turn.
|
||||
- Every forget goes through `recall_forget.retract`, which holds the graph's write lock (`scope_lock.py`) for the whole forget: if an ingestion holds it for all of `FORGET_LOCK_WAIT_SECONDS` (20 s), every uuid fails as `busy` and nothing is written; the chat tool waits five seconds and tries once more, and the settings page answers 409. It first completes the dream records still pending in the graph (`recall_reconcile.py`, below). Soft stamps `forgotten_at`, sets `status='retracted'` and `expiration_reason`, keeps an earlier `expired_at`, leaves `invalid_at` alone, then scrubs the sentence, entities and communities and stamps `redacted_at` on every episode citing the fact (`recall_hide.py`); then, still under the lock, it retracts what the dream derived from the forgotten facts (`recall_cascade.py`, below). An audit copy, once written, is never replaced by the placeholder. Hard does all of that first, stamping `hard_forgotten_at` on each fact as it retracts it (so a dream write settling before the purge erases, below), the cascade too (softly, erasing the derived text it reaches, below), then (`recall_orphans.py`) turns each episode no other edge of any status cites into a tombstone (`content` emptied, `hard_deleted_at` and `redacted_at` set, its name, source description and envelope `provenance` kept, so the dream still leaves its chat session out), and last, in one query per edge, deletes the edge, drops the tombstones' mentions and deletes each entity left with no fact and no mention. Every step before the delete is idempotent and the delete finds the edge again, so repeating a forget whose clean-up failed finishes it, a purged fact's cascade included (below). A step that fails is reported per edge as `cleanup_error`; recall hides the text regardless. A tombstone is never recalled, counted on the settings page or used as extraction context.
|
||||
- While Redis answers and the holder's lease holds, a forget and an ingestion never overlap. graphiti's `add_episode` saves the edges and entities it read when it began, so a forget landing mid-episode would be written over. The ingestion worker therefore holds the graph's write lock around `add_episode` and the dream metadata stamp after it (`ingest._write_locked`), waiting up to `INGEST_LOCK_WAIT_SECONDS` (60 s) for the writer holding it (a forget, or an ingestion in another process) to finish; if it cannot, the episode goes to the back of the queue after a pause (one second, doubling with each wait up to 30 s) and waits again, as often as it takes: contention never drops an episode while the process runs, and every third wait a warning names the episode still unwritten. The queue lives in the process: what it holds when the process stops is lost. Waiters take turns: a writer that finds the lock held claims the next turn (the lock's key plus `:next`, renewed at every poll, lapsing two seconds after its waiter stops), and while it holds that claim nobody else takes the lock, the writer that just let it go included, so a worker writing a backlog hands the lock over between two episodes. The lock (`scope_lock.py`) is one Redis key per graph (`MemoryScope.redis_key("write_lock")`), set NX with a token and a five-minute expiry that the holder renews, and released by the dream lock's compare-and-delete script, so a holder whose key expired never frees a newer holder's. When Redis cannot be reached both writers go ahead without it and log a warning (the unbounded windows above). graphiti also resolves each new statement against every edge between the same entities, forgotten ones included: if its model named a forgotten edge a duplicate, graphiti would append the episode to it, hiding the statement under the forget, and a contradiction would stamp its `invalid_at`. So every graphiti client is built with `recall_ingest.ForgetAwareLLMClient` (`client._build_graphiti`), which keeps a forgotten edge (its text is `[forgotten]`) out of graphiti's edge-dedup answers: a statement said again after a forget is saved by graphiti as a new live edge in the same `add_episode`, one merged into a live fact stays merged, and graphiti's model-decided dedup never rewrites a forgotten edge. There is no second extraction and nothing to repair afterwards. One path runs before the model and so past the guard: graphiti's exact-text match reuses an edge whose text equals the new statement's (lower-cased, whitespace collapsed), so a statement that reads `[forgotten]` appends its episode's uuid to the forgotten edge's `episodes`. The edge keeps its marker and audit copies and stays out of recall, and the episode, citing it, is hidden. The guard reads graphiti's prompt, so it fails closed: a candidate printed with the placeholder is forgotten wherever it appears, in the two candidate lists or not, and when the prompt shows the placeholder, anything short of both lists, once each and readable, makes the answer name no edge at all, so the statement is saved as new (at worst a duplicate of a live fact). A prompt whose lists it cannot read is logged at error once per process for each shape, so a graphiti upgrade that changes the prompt shows up before a forgotten edge meets it.
|
||||
- The dream keeps its own Cypher. Its gather and session filter (`dream/fetch.py`, `dream/hidden_sessions.py`) build on those predicates, and so does its fact gather (`live_fact_predicate`). Its writers only write over live facts: its demotions (`supersede_unless_recalled`) and single-hop neighbour invalidation (`invalidate_entity_direct_neighbors`), both in `guarded_writes.py`, match `live_fact_predicate`, and ratification (`dream/ratification.py`) supersedes with `supersede_unless_recalled` and `expected_status='tentative'` and promotes only a still-tentative, unexpired, unforgotten edge, so none of them can overwrite a forget, and none writes a new fact. Its consolidated facts and proposals are new statements, written minutes or hours after its read, so each goes to the ingestion worker with what it cites (`recall_citations.Citations`). Every one must cite the facts and episodes it was drawn from: apply keeps a citation only when it names something the pass read, filed under the kind the pass read it as, and drops before it is queued a write left citing nothing or citing a fact of another scope than its own (a fact's `scope`, else `real:global`: the prompts' "group by scope" rule, enforced), counting it in `uncited_writes_dropped` and each such citation in `cross_scope_citations_dropped` (`dream/citations.py`; both counts through the apply stats, `DreamPassResult`, the DreamPass record and the job status). The whole write goes, not only its citations of another scope: trimmed, it could restate the other fact and escape that fact's forget. Episode citations are not scoped: a chat turn can hold facts of any scope, and the cascade reaches a fact consolidated from one through `derived_from_episodes` when a forget hides that turn. Under the write lock, right before `add_episode`, a write resting on a forget made since the pass read the graph is dropped unwritten: one citing a fact forgotten or deleted, an episode no longer recallable, or a derived fact no longer live (superseded, contradicted) that rests on one of those. The last is found walking up the records (`derived_from_facts`, `derived_from_episodes`) from each cited derived fact no longer live, as far as derived facts no longer live go, the cascade's walk-through rule upward (`recall_sources.ancestry`): the cascade walks down through such a fact to what rests on it, and the write would rest on it the same way, carrying that content. A live fact, a user's fact (a source without a record states it) or a recallable episode ends a branch clean. The walk is bounded like the cascade, at most `CASCADE_MAX_ROUNDS` (10) levels up and `CASCADE_MAX_ITEMS` (500) facts read, and a walk stopped by either with facts left to read drops the write too (it fails closed). A dropped write is logged at info, and counted in the pass's `dropped_forgotten` (through the same places as `uncited_writes_dropped`; a record written before it reads 0; it counts the drops the worker made before apply returned, which is all of them on a drained pass; the batch path does not wait, so its count covers only the drops made before it returned). The pass counts a write consolidated, or a proposal proposed, only once graphiti has saved it: those still queued or being written when apply returns are reported as `pending_writes` (every one on the batch path), and those not made, their citation marker unwritable or their write raising, as `failed_writes`. The check trusts the model's citations: a write that cites one source but restates another, forgotten, in other words is not recognised, and no check can tell whether a source in the write's own scope truly supports it. (The check still compares the statement of a write citing nothing with the sentences forgotten facts keep, as defence in depth: apply no longer sends one.)
|
||||
|
||||
@@ -7,8 +7,9 @@ community. The audit's default recommendation was to leave communities
|
||||
off in P0 and revisit only if retrieval relevance plateaued.
|
||||
|
||||
The user direction for P-1 overrides that default: enable communities
|
||||
now, behind a LaunchDarkly flag, run rebuilds during off-peak windows
|
||||
to bound cost. The scheduler hook in ``backend/executor/scheduler.py``
|
||||
now, and run rebuilds as a weekly memory-maintenance job at an off-peak
|
||||
hour (owner-local 04:00, see ``dream/scheduling.py``) to bound cost. No
|
||||
feature flag gates them. The scheduler hook in ``backend/executor/scheduler.py``
|
||||
calls into this module.
|
||||
|
||||
Operational notes:
|
||||
|
||||
+1
-1
@@ -210,7 +210,7 @@ def stub_boundaries(
|
||||
* ``get_graphiti_client`` — patched on ``ingest`` (where it is used) to
|
||||
return a Graphiti wired to the test's real driver with the LLM stubbed.
|
||||
* ``ensure_scope_scheduled`` — patched on ``ingest`` (where the first
|
||||
write per group calls it). Left live it would fire LaunchDarkly,
|
||||
write per group calls it). Left live it would fire PostHog,
|
||||
registry and scheduler calls from a background task.
|
||||
* ``get_redis_async`` — likewise imported lazily by the hit tracker.
|
||||
"""
|
||||
|
||||
@@ -2,20 +2,14 @@ import io
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from langfuse import get_client
|
||||
from PIL import Image
|
||||
|
||||
from backend.copilot.google_data_redaction import redact_google_data
|
||||
from backend.copilot.google_data_redaction_test import langfuse_spans
|
||||
from backend.copilot.tool_image import prepare_image
|
||||
|
||||
__all__ = ["langfuse_spans"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("format", ["anthropic", "mcp", "openai", "responses"])
|
||||
def test_image_pixels_stay_out_of_traces_without_changing_model_input(
|
||||
format, langfuse_spans
|
||||
):
|
||||
def test_image_pixels_stay_out_of_traces_without_changing_model_input(format):
|
||||
output = io.BytesIO()
|
||||
Image.new("RGB", (20, 20), "blue").save(output, format="PNG")
|
||||
image = prepare_image(output.getvalue())
|
||||
@@ -36,15 +30,8 @@ def test_image_pixels_stay_out_of_traces_without_changing_model_input(
|
||||
],
|
||||
}
|
||||
before = json.dumps(payload)
|
||||
generation = get_client().start_observation(
|
||||
name="image-inspection", as_type="generation", input=payload
|
||||
)
|
||||
generation.end()
|
||||
traced = json.dumps(
|
||||
[dict(span.attributes or {}) for span in langfuse_spans.get_finished_spans()]
|
||||
)
|
||||
traced = json.dumps(redact_google_data(payload))
|
||||
assert image.data_base64 not in traced
|
||||
assert image.data_base64 not in json.dumps(redact_google_data(json.loads(before)))
|
||||
assert "Inspect this page" in traced
|
||||
assert json.dumps(payload) == before
|
||||
assert image.data_base64 in before
|
||||
|
||||
@@ -116,14 +116,13 @@ class InferenceUsage(BaseModel):
|
||||
|
||||
|
||||
class InferenceContext(BaseModel):
|
||||
"""One call's scope, job and route; ``trace_id`` once it is traced."""
|
||||
"""One call's scope, job and route."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
scope: InferenceScope
|
||||
job: InferenceJob
|
||||
route: RouteDecision
|
||||
trace_id: str | None = None
|
||||
|
||||
|
||||
# How a background call failed, naming it without quoting the model.
|
||||
|
||||
@@ -60,7 +60,7 @@ def test_models_are_frozen():
|
||||
scope=InferenceScope(user_id="u1"), job=_job(), route=_route()
|
||||
)
|
||||
with pytest.raises(ValidationError):
|
||||
ctx.trace_id = "t"
|
||||
ctx.scope = InferenceScope(user_id="u2")
|
||||
|
||||
|
||||
def test_route_keeps_the_platform_credential_implicit():
|
||||
|
||||
@@ -187,8 +187,6 @@ def _row_metadata(ctx: InferenceContext, accounting: _KindAccounting) -> dict[st
|
||||
"discount_applied": batch_discount(ctx.route.execution_path),
|
||||
"expert_id": ctx.scope.expert_id,
|
||||
}
|
||||
if ctx.trace_id:
|
||||
metadata["langfuse_trace_id"] = ctx.trace_id
|
||||
return metadata
|
||||
|
||||
|
||||
|
||||
@@ -41,7 +41,6 @@ def _ctx(
|
||||
*,
|
||||
route: RouteDecision | None = None,
|
||||
expert_id: str | None = "e1",
|
||||
trace_id: str | None = None,
|
||||
) -> InferenceContext:
|
||||
return InferenceContext(
|
||||
scope=InferenceScope(user_id="u1", expert_id=expert_id),
|
||||
@@ -53,7 +52,6 @@ def _ctx(
|
||||
tier="standard",
|
||||
),
|
||||
route=route or _sync_route(),
|
||||
trace_id=trace_id,
|
||||
)
|
||||
|
||||
|
||||
@@ -115,12 +113,6 @@ async def test_row_metadata_names_the_job_the_path_and_who_pays(persist):
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_row_links_the_langfuse_trace_when_there_is_one(persist):
|
||||
await record(_ctx(trace_id="t-123"), _usage(0.01), block_name="b")
|
||||
assert persist.await_args.kwargs["extra_metadata"]["langfuse_trace_id"] == "t-123"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_account_scope_leaves_the_expert_empty(persist):
|
||||
await record(_ctx(expert_id=None), _usage(0.01), block_name="b")
|
||||
|
||||
@@ -1,164 +0,0 @@
|
||||
"""A Langfuse trace around one background LLM call, the way chat traces a turn.
|
||||
|
||||
Chat opens a parent span per turn, propagates the user, session, trace name,
|
||||
tags and metadata onto it (``sdk/service.py``, ``baseline/service.py``) and
|
||||
writes the turn's ``gen_ai.usage.*`` onto the span as it closes. ``trace(ctx)``
|
||||
does the same for a background call: a span named ``kind:phase``, the scope's
|
||||
user, the job's correlation id as the session, tags for the kind, the
|
||||
execution path and the expert, and the usage the caller hands back on
|
||||
``TracedCall.usage`` (or the billed usage an ``InferenceError`` carries).
|
||||
|
||||
A call made inside a span someone else opened (a consult inside the chat turn
|
||||
that asked for it) becomes a child span and leaves that trace's own
|
||||
attributes alone. Without Langfuse keys ``trace`` does nothing, and a tracing
|
||||
failure never fails the call.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import ExitStack, asynccontextmanager
|
||||
from types import TracebackType
|
||||
|
||||
from langfuse import get_client, propagate_attributes
|
||||
from opentelemetry import trace as otel_trace
|
||||
from pydantic import BaseModel
|
||||
|
||||
from backend.util.settings import Settings
|
||||
|
||||
from .context import InferenceContext, InferenceError, InferenceUsage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
settings = Settings()
|
||||
|
||||
_ExcInfo = tuple[type[BaseException] | None, BaseException | None, TracebackType | None]
|
||||
_NO_EXCEPTION: _ExcInfo = (None, None, None)
|
||||
|
||||
|
||||
class TracedCall(BaseModel):
|
||||
"""One traced call: its context (with the trace id once one is open) and
|
||||
the usage to write on the span when it closes."""
|
||||
|
||||
ctx: InferenceContext
|
||||
usage: InferenceUsage | None = None
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def trace(ctx: InferenceContext) -> AsyncIterator[TracedCall]:
|
||||
"""Trace one call; set ``usage`` on what this yields once the call returns."""
|
||||
spans = _open_spans(ctx)
|
||||
if spans is None:
|
||||
yield TracedCall(ctx=ctx)
|
||||
return
|
||||
call = TracedCall(ctx=ctx.model_copy(update={"trace_id": _current_trace_id()}))
|
||||
exc_info = _NO_EXCEPTION
|
||||
try:
|
||||
yield call
|
||||
except InferenceError as exc:
|
||||
call.usage = call.usage or exc.usage
|
||||
exc_info = sys.exc_info()
|
||||
raise
|
||||
except BaseException:
|
||||
exc_info = sys.exc_info()
|
||||
raise
|
||||
finally:
|
||||
_write_usage(call.usage)
|
||||
_close_spans(spans, exc_info)
|
||||
|
||||
|
||||
def _open_spans(ctx: InferenceContext) -> ExitStack | None:
|
||||
"""The call's span, plus the trace's attributes when the call starts the
|
||||
trace; ``None`` when Langfuse is off or the span would not open."""
|
||||
if not _langfuse_configured():
|
||||
return None
|
||||
spans = ExitStack()
|
||||
try:
|
||||
starts_trace = not otel_trace.get_current_span().get_span_context().is_valid
|
||||
spans.enter_context(
|
||||
get_client().start_as_current_span(
|
||||
name=ctx.job.label, metadata=_metadata(ctx)
|
||||
)
|
||||
)
|
||||
if starts_trace:
|
||||
spans.enter_context(
|
||||
propagate_attributes(
|
||||
user_id=ctx.scope.user_id,
|
||||
session_id=ctx.job.correlation_id,
|
||||
trace_name=ctx.job.label,
|
||||
tags=_tags(ctx),
|
||||
metadata=_metadata(ctx),
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("Langfuse span for %s did not open", ctx.job.label, exc_info=True)
|
||||
_close_spans(spans, _NO_EXCEPTION)
|
||||
return None
|
||||
return spans
|
||||
|
||||
|
||||
def _close_spans(spans: ExitStack, exc_info: _ExcInfo) -> None:
|
||||
try:
|
||||
spans.__exit__(*exc_info)
|
||||
except Exception:
|
||||
logger.warning("Langfuse span teardown failed", exc_info=True)
|
||||
|
||||
|
||||
def _write_usage(usage: InferenceUsage | None) -> None:
|
||||
"""The ``gen_ai.usage.*`` attributes chat writes on its turn span."""
|
||||
if usage is None:
|
||||
return
|
||||
try:
|
||||
span = otel_trace.get_current_span()
|
||||
if not span.is_recording():
|
||||
return
|
||||
span.set_attribute("gen_ai.usage.prompt_tokens", usage.input_tokens)
|
||||
span.set_attribute("gen_ai.usage.completion_tokens", usage.output_tokens)
|
||||
span.set_attribute("gen_ai.usage.cache_read_tokens", usage.cache_read_tokens)
|
||||
span.set_attribute(
|
||||
"gen_ai.usage.cache_creation_tokens", usage.cache_creation_tokens
|
||||
)
|
||||
if usage.cost_usd is not None:
|
||||
span.set_attribute("gen_ai.usage.cost_usd", usage.cost_usd)
|
||||
except Exception:
|
||||
logger.debug("Failed to set OTEL usage attributes", exc_info=True)
|
||||
|
||||
|
||||
def _current_trace_id() -> str | None:
|
||||
try:
|
||||
return get_client().get_current_trace_id()
|
||||
except Exception:
|
||||
logger.debug("Failed to read the Langfuse trace id", exc_info=True)
|
||||
return None
|
||||
|
||||
|
||||
def _tags(ctx: InferenceContext) -> list[str]:
|
||||
tags = [ctx.job.kind, f"execution_path:{ctx.route.execution_path}"]
|
||||
if ctx.scope.expert_id:
|
||||
tags.append(f"expert:{ctx.scope.expert_id}")
|
||||
return tags
|
||||
|
||||
|
||||
def _metadata(ctx: InferenceContext) -> dict[str, str]:
|
||||
"""String values only: Langfuse drops any other kind from propagated
|
||||
metadata."""
|
||||
metadata = {
|
||||
"job_kind": ctx.job.kind,
|
||||
"phase": ctx.job.phase,
|
||||
"correlation_id": ctx.job.correlation_id,
|
||||
"latency_class": ctx.job.latency_class,
|
||||
"tier": ctx.job.tier,
|
||||
"model": ctx.route.model,
|
||||
"provider": ctx.route.provider,
|
||||
"payer": ctx.route.payer,
|
||||
"execution_path": ctx.route.execution_path,
|
||||
"route_reason": ctx.route.reason,
|
||||
"expert_id": ctx.scope.expert_id,
|
||||
}
|
||||
return {key: value for key, value in metadata.items() if value is not None}
|
||||
|
||||
|
||||
def _langfuse_configured() -> bool:
|
||||
return bool(
|
||||
settings.secrets.langfuse_public_key and settings.secrets.langfuse_secret_key
|
||||
)
|
||||
@@ -1,198 +0,0 @@
|
||||
"""``trace`` gives a background call the Langfuse trace chat gives a turn."""
|
||||
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from . import trace as trace_mod
|
||||
from .context import (
|
||||
InferenceContext,
|
||||
InferenceError,
|
||||
InferenceJob,
|
||||
InferenceScope,
|
||||
InferenceUsage,
|
||||
RouteDecision,
|
||||
)
|
||||
from .trace import trace
|
||||
|
||||
_CTX = InferenceContext(
|
||||
scope=InferenceScope(user_id="u1", expert_id="e1"),
|
||||
job=InferenceJob(
|
||||
kind="dream",
|
||||
phase="recombine",
|
||||
correlation_id="pass-1",
|
||||
latency_class="deferred",
|
||||
tier="advanced",
|
||||
),
|
||||
route=RouteDecision(
|
||||
engine="provider_sync",
|
||||
auth_provider="platform",
|
||||
provider="open_router",
|
||||
model="anthropic/claude-opus-5-5",
|
||||
payer="platform_allowance",
|
||||
execution_path="sync_baseline",
|
||||
cost_log_provider="open_router",
|
||||
reason="open_router chat transport, advanced tier (fast_advanced_model)",
|
||||
),
|
||||
)
|
||||
|
||||
_USAGE = InferenceUsage(
|
||||
model="anthropic/claude-opus-5-5",
|
||||
input_tokens=120,
|
||||
output_tokens=30,
|
||||
cache_read_tokens=7,
|
||||
cache_creation_tokens=3,
|
||||
cost_usd=0.0123,
|
||||
cost_source="provider",
|
||||
payer="platform_allowance",
|
||||
)
|
||||
|
||||
|
||||
class _FakeSpan:
|
||||
"""The current OTEL span: inside someone else's trace or not."""
|
||||
|
||||
def __init__(self, *, inside_a_trace: bool) -> None:
|
||||
self.attributes: dict[str, object] = {}
|
||||
self._context = SimpleNamespace(is_valid=inside_a_trace)
|
||||
|
||||
def get_span_context(self):
|
||||
return self._context
|
||||
|
||||
def is_recording(self) -> bool:
|
||||
return True
|
||||
|
||||
def set_attribute(self, key: str, value: object) -> None:
|
||||
self.attributes[key] = value
|
||||
|
||||
|
||||
class _FakeLangfuse:
|
||||
def __init__(self) -> None:
|
||||
self.spans: list[dict] = []
|
||||
self.exits: list[type[BaseException] | None] = []
|
||||
|
||||
@contextmanager
|
||||
def start_as_current_span(self, *, name, metadata):
|
||||
self.spans.append({"name": name, "metadata": metadata})
|
||||
try:
|
||||
yield SimpleNamespace()
|
||||
except BaseException as exc:
|
||||
self.exits.append(type(exc))
|
||||
raise
|
||||
else:
|
||||
self.exits.append(None)
|
||||
|
||||
def get_current_trace_id(self) -> str:
|
||||
return "trace-1"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def langfuse(monkeypatch):
|
||||
"""Langfuse configured, with a fake client, tracer and propagation."""
|
||||
client = _FakeLangfuse()
|
||||
propagated: list[dict] = []
|
||||
|
||||
@contextmanager
|
||||
def fake_propagate(**kwargs):
|
||||
propagated.append(kwargs)
|
||||
yield
|
||||
|
||||
span = _FakeSpan(inside_a_trace=False)
|
||||
monkeypatch.setattr(trace_mod, "_langfuse_configured", lambda: True)
|
||||
monkeypatch.setattr(trace_mod, "get_client", lambda: client)
|
||||
monkeypatch.setattr(trace_mod, "propagate_attributes", fake_propagate)
|
||||
monkeypatch.setattr(
|
||||
trace_mod, "otel_trace", SimpleNamespace(get_current_span=lambda: span)
|
||||
)
|
||||
return SimpleNamespace(client=client, propagated=propagated, span=span)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_without_langfuse_it_does_nothing(monkeypatch):
|
||||
monkeypatch.setattr(trace_mod, "_langfuse_configured", lambda: False)
|
||||
get_client = MagicMock()
|
||||
with patch.object(trace_mod, "get_client", get_client):
|
||||
async with trace(_CTX) as call:
|
||||
call.usage = _USAGE
|
||||
get_client.assert_not_called()
|
||||
assert call.ctx is _CTX
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_without_langfuse_errors_still_propagate(monkeypatch):
|
||||
monkeypatch.setattr(trace_mod, "_langfuse_configured", lambda: False)
|
||||
with pytest.raises(InferenceError):
|
||||
async with trace(_CTX):
|
||||
raise InferenceError("bad json")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_background_call_starts_a_trace_like_a_chat_turn(langfuse):
|
||||
async with trace(_CTX) as call:
|
||||
call.usage = _USAGE
|
||||
|
||||
(span,) = langfuse.client.spans
|
||||
assert span["name"] == "dream:recombine"
|
||||
(attributes,) = langfuse.propagated
|
||||
assert attributes["user_id"] == "u1"
|
||||
assert attributes["session_id"] == "pass-1"
|
||||
assert attributes["trace_name"] == "dream:recombine"
|
||||
assert attributes["tags"] == ["dream", "execution_path:sync_baseline", "expert:e1"]
|
||||
assert attributes["metadata"]["model"] == "anthropic/claude-opus-5-5"
|
||||
assert attributes["metadata"]["payer"] == "platform_allowance"
|
||||
# Langfuse drops non-string propagated metadata, so none is sent.
|
||||
assert all(isinstance(v, str) for v in attributes["metadata"].values())
|
||||
assert call.ctx.trace_id == "trace-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_lands_on_the_span_as_chat_writes_it(langfuse):
|
||||
async with trace(_CTX) as call:
|
||||
call.usage = _USAGE
|
||||
|
||||
assert langfuse.span.attributes == {
|
||||
"gen_ai.usage.prompt_tokens": 120,
|
||||
"gen_ai.usage.completion_tokens": 30,
|
||||
"gen_ai.usage.cache_read_tokens": 7,
|
||||
"gen_ai.usage.cache_creation_tokens": 3,
|
||||
"gen_ai.usage.cost_usd": 0.0123,
|
||||
}
|
||||
assert langfuse.client.exits == [None]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_billed_failure_still_reports_its_usage(langfuse):
|
||||
with pytest.raises(InferenceError):
|
||||
async with trace(_CTX):
|
||||
raise InferenceError("did not parse", _USAGE)
|
||||
|
||||
assert langfuse.span.attributes["gen_ai.usage.prompt_tokens"] == 120
|
||||
# The span closes with the error, so the trace shows the call failed.
|
||||
assert langfuse.client.exits == [InferenceError]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inside_a_chat_turn_it_is_a_child_span_only(langfuse):
|
||||
"""A consult runs inside the asking turn's trace; renaming that trace or
|
||||
re-tagging it would clobber the chat's own attributes."""
|
||||
langfuse.span._context = SimpleNamespace(is_valid=True)
|
||||
|
||||
async with trace(_CTX):
|
||||
pass
|
||||
|
||||
assert [s["name"] for s in langfuse.client.spans] == ["dream:recombine"]
|
||||
assert langfuse.propagated == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_span_that_will_not_open_never_fails_the_call(langfuse, monkeypatch):
|
||||
def broken_client():
|
||||
raise RuntimeError("langfuse down")
|
||||
|
||||
monkeypatch.setattr(trace_mod, "get_client", broken_client)
|
||||
|
||||
async with trace(_CTX) as call:
|
||||
call.usage = _USAGE
|
||||
|
||||
assert call.ctx.trace_id is None
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
Goes through the inference layer every background call takes
|
||||
(``copilot/inference``): the chat transport's provider and the platform's
|
||||
key, on the reviewer's model, the call traced, and its spend recorded in the
|
||||
key, on the reviewer's model, with its spend recorded in the
|
||||
usage ledger so learning appears next to chat spend under one provider label.
|
||||
"""
|
||||
|
||||
@@ -21,7 +21,6 @@ from backend.copilot.inference.context import (
|
||||
)
|
||||
from backend.copilot.inference.record import price, record
|
||||
from backend.copilot.inference.routing import resolve_route
|
||||
from backend.copilot.inference.trace import trace
|
||||
from backend.copilot.tools.skills import ParsedSkill
|
||||
|
||||
from .contract import EvidenceBundle
|
||||
@@ -57,16 +56,13 @@ async def review_evidence(
|
||||
``DreamLLMError`` on provider/parse failure; its spend is the caller's to
|
||||
record (``record_review_cost``)."""
|
||||
ctx = _review_context(user_id, run_id, reviewer_model(config), config)
|
||||
async with trace(ctx) as call:
|
||||
completion = await structured_complete(
|
||||
call.ctx,
|
||||
build_review_messages(bundle, existing_skills, packages),
|
||||
LearningProposal,
|
||||
temperature=REVIEW_TEMPERATURE,
|
||||
max_output_tokens=REVIEW_MAX_OUTPUT_TOKENS,
|
||||
)
|
||||
call.usage = completion.usage
|
||||
return completion
|
||||
return await structured_complete(
|
||||
ctx,
|
||||
build_review_messages(bundle, existing_skills, packages),
|
||||
LearningProposal,
|
||||
temperature=REVIEW_TEMPERATURE,
|
||||
max_output_tokens=REVIEW_MAX_OUTPUT_TOKENS,
|
||||
)
|
||||
|
||||
|
||||
async def record_review_cost(
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"""Who an AutoPilot LLM call is for, in the fields OpenRouter Broadcast reads.
|
||||
|
||||
OpenRouter Broadcast sends every generation to PostHog LLM analytics as an
|
||||
``$ai_generation`` event (and to Langfuse). It maps these request body fields
|
||||
``$ai_generation`` event. It maps these request body fields
|
||||
(https://openrouter.ai/docs/guides/features/broadcast/posthog):
|
||||
|
||||
- ``user`` becomes the event's ``distinct_id``. It is the platform user id,
|
||||
|
||||
@@ -240,7 +240,7 @@ class ChatMessage(BaseModel):
|
||||
# model. None on user/tool rows.
|
||||
#
|
||||
# routing_source is excluded from serialized payloads: "ld" vs "env"
|
||||
# would let a client infer LaunchDarkly cohort membership. It still
|
||||
# would let a client infer flag cohort membership. It still
|
||||
# persists — the DB save path maps fields explicitly (not via
|
||||
# model_dump) and the back-fill reads attributes.
|
||||
model: str | None = None
|
||||
@@ -254,14 +254,10 @@ class ChatMessage(BaseModel):
|
||||
llm_auth_provider: CopilotLlmAuthProvider | None = None
|
||||
llm_credential_id: str | None = None
|
||||
|
||||
# Langfuse trace of the SDK turn that wrote this assistant or reasoning
|
||||
# row; a rating of the reply is scored against it. Internal, so excluded
|
||||
# from payloads like routing_source, and persisted the same way.
|
||||
langfuse_trace_id: str | None = Field(default=None, exclude=True)
|
||||
|
||||
# PostHog trace of the turn that wrote this assistant or reasoning row
|
||||
# (``copilot.ai_spans``): the ``$ai_trace_id`` of its generations, which a
|
||||
# rating of the reply carries. Internal and persisted like the above.
|
||||
# rating of the reply carries. Internal, so excluded from payloads like
|
||||
# routing_source, and persisted the same way.
|
||||
trace_id: str | None = Field(default=None, exclude=True)
|
||||
|
||||
stamps_pending_save: bool = Field(default=False, exclude=True)
|
||||
@@ -331,7 +327,6 @@ class ChatMessage(BaseModel):
|
||||
"CopilotLlmAuthProvider | None", prisma_message.llmAuthProvider
|
||||
),
|
||||
llm_credential_id=prisma_message.llmCredentialId,
|
||||
langfuse_trace_id=prisma_message.langfuseTraceId,
|
||||
trace_id=prisma_message.traceId,
|
||||
)
|
||||
|
||||
@@ -1223,7 +1218,6 @@ async def _save_session_to_db(
|
||||
"routing_source": msg.routing_source,
|
||||
"llm_auth_provider": msg.llm_auth_provider,
|
||||
"llm_credential_id": msg.llm_credential_id,
|
||||
"langfuse_trace_id": msg.langfuse_trace_id,
|
||||
"trace_id": msg.trace_id,
|
||||
"metadata": msg.metadata,
|
||||
}
|
||||
@@ -1278,7 +1272,6 @@ async def _save_session_to_db(
|
||||
routing_source=msg.routing_source,
|
||||
llm_auth_provider=msg.llm_auth_provider,
|
||||
llm_credential_id=msg.llm_credential_id,
|
||||
langfuse_trace_id=msg.langfuse_trace_id,
|
||||
trace_id=msg.trace_id,
|
||||
)
|
||||
except Exception as e:
|
||||
@@ -1799,17 +1792,19 @@ async def get_user_sessions(
|
||||
number of sessions for the user (not just the current page).
|
||||
"""
|
||||
db = chat_db()
|
||||
sessions = await db.get_user_chat_sessions(
|
||||
user_id,
|
||||
limit,
|
||||
offset,
|
||||
organization_id=organization_id,
|
||||
title_contains=title_contains,
|
||||
expert_id=expert_id,
|
||||
pinned_first=pinned_first,
|
||||
)
|
||||
total_count = await db.get_user_session_count(
|
||||
user_id, organization_id=organization_id, expert_id=expert_id
|
||||
sessions, total_count = await asyncio.gather(
|
||||
db.get_user_chat_sessions(
|
||||
user_id,
|
||||
limit,
|
||||
offset,
|
||||
organization_id=organization_id,
|
||||
title_contains=title_contains,
|
||||
expert_id=expert_id,
|
||||
pinned_first=pinned_first,
|
||||
),
|
||||
db.get_user_session_count(
|
||||
user_id, organization_id=organization_id, expert_id=expert_id
|
||||
),
|
||||
)
|
||||
|
||||
return sessions, total_count
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
"""Model selection for the copilot: LaunchDarkly → registry cell → env.
|
||||
"""Model selection for the copilot: feature flag → registry cell → env.
|
||||
|
||||
Each cell of the ``(mode, tier)`` matrix resolves through three layers:
|
||||
|
||||
1. The JSON-valued LaunchDarkly flag ``copilot-model-routing`` (per-user —
|
||||
1. The JSON-valued PostHog flag ``copilot-model-routing`` (per-user —
|
||||
cohort experiments and rollouts live here, and the flag returns model
|
||||
slugs directly).
|
||||
2. The catalog's routing cell (``catalog.py`` ``routing`` section, surface
|
||||
@@ -110,7 +110,7 @@ def catalog_lookup(slug: str) -> llm_registry.RegistryModel | None:
|
||||
"""Look up *slug* in the catalog, tolerating transport spellings.
|
||||
|
||||
The catalog registers Claude models under bare canonical enum slugs
|
||||
(``claude-opus-4-7``) while LaunchDarkly payloads and env defaults use
|
||||
(``claude-opus-4-7``) while flag payloads and env defaults use
|
||||
OpenRouter forms (``anthropic/claude-opus-4.6``). Gate on the model,
|
||||
not the spelling: try the exact slug, then the ``anthropic/``-stripped
|
||||
tail, then its dots→dashes form.
|
||||
|
||||
@@ -2226,7 +2226,6 @@ async def test_save_session_to_db_backfills_stamps_on_flushed_rows(
|
||||
routing_source="env",
|
||||
llm_auth_provider=None,
|
||||
llm_credential_id=None,
|
||||
langfuse_trace_id=None,
|
||||
trace_id=None,
|
||||
)
|
||||
assert flushed.stamps_pending_save is False
|
||||
@@ -2236,17 +2235,15 @@ async def test_save_session_to_db_backfills_stamps_on_flushed_rows(
|
||||
async def test_save_session_to_db_persists_the_turn_trace(
|
||||
mocker: MockerFixture,
|
||||
) -> None:
|
||||
"""The Langfuse and PostHog traces stamped at end of turn reach the DB
|
||||
"""The PostHog trace stamped at end of turn reaches the DB
|
||||
both on rows flushed mid-turn (back-fill) and on rows inserted by this
|
||||
save."""
|
||||
trace_id = "1edf31f11b1693cc6103f358c1481694"
|
||||
flushed = ChatMessage(
|
||||
role="assistant",
|
||||
content="working on it",
|
||||
sequence=7,
|
||||
model="claude-sonnet-4-6",
|
||||
routing_source="env",
|
||||
langfuse_trace_id=trace_id,
|
||||
trace_id="turn-1",
|
||||
stamps_pending_save=True,
|
||||
)
|
||||
@@ -2255,7 +2252,6 @@ async def test_save_session_to_db_persists_the_turn_trace(
|
||||
content="done",
|
||||
model="claude-sonnet-4-6",
|
||||
routing_source="env",
|
||||
langfuse_trace_id=trace_id,
|
||||
trace_id="turn-1",
|
||||
)
|
||||
session = _make_session_with_messages(flushed, unsaved)
|
||||
@@ -2272,10 +2268,8 @@ async def test_save_session_to_db_persists_the_turn_trace(
|
||||
|
||||
backfill = mock_db.update_chat_message_stamps.await_args
|
||||
assert backfill.kwargs["sequence"] == 7
|
||||
assert backfill.kwargs["langfuse_trace_id"] == trace_id
|
||||
assert backfill.kwargs["trace_id"] == "turn-1"
|
||||
(inserted,) = mock_db.add_chat_messages_batch.await_args.kwargs["messages"]
|
||||
assert inserted["langfuse_trace_id"] == trace_id
|
||||
assert inserted["trace_id"] == "turn-1"
|
||||
|
||||
|
||||
|
||||
@@ -39,7 +39,7 @@ class ConnectionTier(BaseModel):
|
||||
"""One quality level on a connection.
|
||||
|
||||
``display_model`` is the model this tier resolves to, run through the
|
||||
same router the turn will use — LaunchDarkly cell, then registry, then
|
||||
same router the turn will use — feature-flag cell, then registry, then
|
||||
config — so it cannot drift from what actually answers.
|
||||
|
||||
On a ChatGPT connection it is the model the catalog pins for that cell,
|
||||
|
||||
@@ -157,7 +157,7 @@ DEFAULT_TIER = SubscriptionTier.NO_TIER
|
||||
|
||||
# Per-tier workspace storage caps in MB. NO_TIER keeps the same baseline as
|
||||
# BASIC so users who cancel retain a small quota and see a real overage cap,
|
||||
# while LaunchDarkly can still tune tiers without a deploy.
|
||||
# while PostHog can still tune tiers without a deploy.
|
||||
_DEFAULT_TIER_WORKSPACE_STORAGE_MB: dict[SubscriptionTier, int] = {
|
||||
SubscriptionTier.NO_TIER: 250, # 250 MB
|
||||
SubscriptionTier.TRIAL: 250,
|
||||
@@ -1464,9 +1464,9 @@ async def get_global_rate_limits(
|
||||
config_daily: int,
|
||||
config_weekly: int,
|
||||
) -> tuple[int, int, SubscriptionTier]:
|
||||
"""Resolve global rate limits from LaunchDarkly, falling back to config.
|
||||
"""Resolve global rate limits from PostHog, falling back to config.
|
||||
|
||||
Values are microdollars. The base limits (from LD or config) are
|
||||
Values are microdollars. The base limits (from the flag or config) are
|
||||
multiplied by the user's tier multiplier so that higher tiers receive
|
||||
proportionally larger allowances.
|
||||
|
||||
|
||||
@@ -1715,8 +1715,8 @@ class TestSetUserTier:
|
||||
assert tier_after == SubscriptionTier.ENTERPRISE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drift_check_swallows_launchdarkly_failure(self):
|
||||
"""LaunchDarkly price-id lookup failures inside the drift check must
|
||||
async def test_drift_check_swallows_posthog_failure(self):
|
||||
"""PostHog price-id lookup failures inside the drift check must
|
||||
never bubble up and 500 the admin tier write — the DB update is
|
||||
already committed by the time we check drift."""
|
||||
mock_prisma = AsyncMock()
|
||||
|
||||
@@ -166,6 +166,12 @@ class _DuplicateSubmission(_GatewayState):
|
||||
replay_key: str
|
||||
|
||||
|
||||
class _Relaunch(_GatewayState):
|
||||
"""A request from the attempt that replaced the one holding these records."""
|
||||
|
||||
fingerprint: str
|
||||
|
||||
|
||||
class _ToolCallRecord(_GatewayState):
|
||||
gateway_call_id: str
|
||||
raw_call_id: str
|
||||
@@ -186,6 +192,12 @@ class _DuplicateToolResultError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
# A refusal of ours is a defect on our side, not an answer, so the turn fails
|
||||
# with the bot's own wording rather than reading as a successful reply.
|
||||
GATEWAY_REFUSED_CODE = "gateway_refused"
|
||||
GATEWAY_REFUSED_MESSAGE = "The model connection failed on our side — please try again"
|
||||
|
||||
|
||||
class CodexAnthropicGateway:
|
||||
def __init__(
|
||||
self,
|
||||
@@ -210,6 +222,11 @@ class CodexAnthropicGateway:
|
||||
# is an HTTP status and a sentence; the typed exception only exists
|
||||
# in here.
|
||||
self.last_failure: ProviderFailure | None = None
|
||||
# The last request this gateway refused itself, as opposed to a
|
||||
# provider failure it relays. The service layer only ever sees the
|
||||
# CLI's rendering of it ("API Error: 409 <message>"), which it would
|
||||
# otherwise forward to the user as the assistant's answer.
|
||||
self.last_refusal: str | None = None
|
||||
self.model = model
|
||||
self.effort: CodexReasoningEffort | None = effort
|
||||
self._transport = cast(
|
||||
@@ -226,6 +243,10 @@ class CodexAnthropicGateway:
|
||||
self._auth_token = secrets.token_urlsafe(32)
|
||||
self._conversations: dict[str, _Conversation] = {}
|
||||
self._tool_calls: dict[str, _ToolCallRecord] = {}
|
||||
# The conversation each relaunch opened, by its fingerprint. Nothing
|
||||
# claimable is left behind a relaunch, so this is the only thing that
|
||||
# can answer its resend.
|
||||
self._relaunches: dict[str, _Conversation] = {}
|
||||
self._results: list[CodexInvocationResult] = []
|
||||
self._closed = False
|
||||
self._close_lock = asyncio.Lock()
|
||||
@@ -359,6 +380,11 @@ class CodexAnthropicGateway:
|
||||
if cleanup_error is not None:
|
||||
raise cleanup_error
|
||||
|
||||
def _refuse(self, status: int, error_type: str, message: str) -> web.Response:
|
||||
"""Answer a request this gateway rejects itself, and remember it."""
|
||||
self.last_refusal = message
|
||||
return _anthropic_error(status, error_type, message)
|
||||
|
||||
async def _handle_health(self, request: web.Request) -> web.Response:
|
||||
if not self._is_authorized(request):
|
||||
return _anthropic_error(401, "authentication_error", "Unauthorized")
|
||||
@@ -383,21 +409,19 @@ class CodexAnthropicGateway:
|
||||
|
||||
async def _handle_messages(self, request: web.Request) -> web.StreamResponse:
|
||||
if not self._is_authorized(request):
|
||||
return _anthropic_error(401, "authentication_error", "Unauthorized")
|
||||
return self._refuse(401, "authentication_error", "Unauthorized")
|
||||
try:
|
||||
payload = await request.json()
|
||||
except web.HTTPRequestEntityTooLarge:
|
||||
return _anthropic_error(
|
||||
413, "request_too_large", _REQUEST_TOO_LARGE_MESSAGE
|
||||
)
|
||||
return self._refuse(413, "request_too_large", _REQUEST_TOO_LARGE_MESSAGE)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return _anthropic_error(
|
||||
return self._refuse(
|
||||
400,
|
||||
"invalid_request_error",
|
||||
"Request body must be JSON",
|
||||
)
|
||||
if not isinstance(payload, dict):
|
||||
return _anthropic_error(
|
||||
return self._refuse(
|
||||
400,
|
||||
"invalid_request_error",
|
||||
"Request body must be an object",
|
||||
@@ -406,43 +430,60 @@ class CodexAnthropicGateway:
|
||||
try:
|
||||
continuation = self._continue_conversation(payload)
|
||||
except _DuplicateToolResultError as exc:
|
||||
return _anthropic_error(
|
||||
return self._refuse(
|
||||
409,
|
||||
"invalid_request_error",
|
||||
str(exc),
|
||||
)
|
||||
except (TypeError, ValueError) as exc:
|
||||
return _anthropic_error(
|
||||
return self._refuse(
|
||||
400,
|
||||
"invalid_request_error",
|
||||
str(exc),
|
||||
)
|
||||
|
||||
streamed = payload.get("stream") is True
|
||||
relaunch_key: str | None = None
|
||||
if isinstance(continuation, _DuplicateSubmission):
|
||||
return await self._duplicate_response(request, continuation, streamed)
|
||||
if continuation is None:
|
||||
try:
|
||||
conversation, replay_key = self._start_conversation(payload), None
|
||||
except (TypeError, ValueError) as exc:
|
||||
return _anthropic_error(400, "invalid_request_error", str(exc))
|
||||
else:
|
||||
if isinstance(continuation, _Continuation):
|
||||
conversation = continuation.conversation
|
||||
replay_key = continuation.replay_key
|
||||
else:
|
||||
try:
|
||||
conversation = self._start_conversation(payload)
|
||||
except (TypeError, ValueError) as exc:
|
||||
return self._refuse(400, "invalid_request_error", str(exc))
|
||||
replay_key = None
|
||||
if isinstance(continuation, _Relaunch):
|
||||
# Registered before the answer is written, so a resend that
|
||||
# arrives meanwhile waits on this conversation's lock and
|
||||
# replays that answer instead of invoking Codex again.
|
||||
relaunch_key = replay_key = continuation.fingerprint
|
||||
self._relaunches[relaunch_key] = conversation
|
||||
|
||||
input_tokens = _estimate_input_tokens(payload)
|
||||
if streamed:
|
||||
return await self._streaming_response(
|
||||
response = await self._streaming_response(
|
||||
request,
|
||||
conversation,
|
||||
input_tokens,
|
||||
replay_key,
|
||||
)
|
||||
return await self._nonstreaming_response(
|
||||
conversation,
|
||||
input_tokens,
|
||||
replay_key,
|
||||
)
|
||||
else:
|
||||
response = await self._nonstreaming_response(
|
||||
conversation,
|
||||
input_tokens,
|
||||
replay_key,
|
||||
)
|
||||
if relaunch_key is not None and relaunch_key not in conversation.replays:
|
||||
# The attempt answered with nothing replayable -- it failed before
|
||||
# its first event, or its answer was too large to keep -- so the
|
||||
# registration would turn the CLI's next retry into the very 409
|
||||
# this branch exists to avoid. Let that retry open an attempt of
|
||||
# its own, as it did before any of this was registered.
|
||||
self._relaunches.pop(relaunch_key, None)
|
||||
return response
|
||||
|
||||
async def _duplicate_response(
|
||||
self,
|
||||
@@ -459,7 +500,7 @@ class CodexAnthropicGateway:
|
||||
streamed=streamed,
|
||||
)
|
||||
if replay is None:
|
||||
return _anthropic_error(
|
||||
return self._refuse(
|
||||
409,
|
||||
"invalid_request_error",
|
||||
"This tool-result request was already accepted",
|
||||
@@ -479,7 +520,7 @@ class CodexAnthropicGateway:
|
||||
def _continue_conversation(
|
||||
self,
|
||||
payload: dict[str, object],
|
||||
) -> _Continuation | _DuplicateSubmission | None:
|
||||
) -> _Continuation | _DuplicateSubmission | _Relaunch | None:
|
||||
tool_results = _extract_tool_results(payload.get("messages"))
|
||||
known = [
|
||||
(self._tool_calls[call_id], result)
|
||||
@@ -512,13 +553,21 @@ class CodexAnthropicGateway:
|
||||
conversation=claimed_by[0].conversation,
|
||||
replay_key=fingerprint,
|
||||
)
|
||||
if any(record.result is None for record, _ in known):
|
||||
raise _DuplicateToolResultError(
|
||||
"This tool-result request refers to a closed model call"
|
||||
relaunched = self._relaunches.get(fingerprint)
|
||||
if relaunched is not None:
|
||||
# The CLI resends a relaunch verbatim too, and each resend
|
||||
# would otherwise cost another subscription invocation.
|
||||
return _DuplicateSubmission(
|
||||
conversation=relaunched,
|
||||
replay_key=fingerprint,
|
||||
)
|
||||
raise _DuplicateToolResultError(
|
||||
"This tool-result request refers to a completed model call"
|
||||
)
|
||||
# Nothing is waiting for these results, so there is no
|
||||
# conversation left to continue. The attempt that made the calls
|
||||
# was abandoned, and the one replacing it resumes the CLI session,
|
||||
# so it replays them with a message appended and can never match a
|
||||
# claim. /v1/messages is stateless: answer it as a new
|
||||
# conversation, because refusing it killed the whole turn.
|
||||
return _Relaunch(fingerprint=fingerprint)
|
||||
|
||||
conversation = claimable[0][0].conversation
|
||||
if any(record.conversation is not conversation for record, _ in claimable):
|
||||
|
||||
@@ -82,6 +82,37 @@ class _FailingAgentSession(_FakeAgentSession):
|
||||
raise self.error
|
||||
|
||||
|
||||
class _FailingOnTheRelaunchSession(_FakeAgentSession):
|
||||
"""A tool call, then an attempt that dies, then one that works."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(use_tool=True)
|
||||
|
||||
async def invoke(
|
||||
self,
|
||||
request,
|
||||
dynamic_tools,
|
||||
tool_handler,
|
||||
event_handler=None,
|
||||
) -> CodexInvocationResult:
|
||||
self.requests.append(request)
|
||||
self.tools.append(dynamic_tools)
|
||||
if len(self.requests) == 1:
|
||||
self.tool_result = await tool_handler(
|
||||
CodexDynamicToolCall(
|
||||
thread_id="thread-1",
|
||||
turn_id="turn-1",
|
||||
call_id="call-1",
|
||||
tool=dynamic_tools[0].name,
|
||||
arguments={"query": "status"},
|
||||
)
|
||||
)
|
||||
return _result("done")
|
||||
if len(self.requests) == 2:
|
||||
raise RuntimeError("codex transport blew up")
|
||||
return _result("recovered")
|
||||
|
||||
|
||||
class _FinalWithoutDeltaSession:
|
||||
async def invoke(
|
||||
self,
|
||||
@@ -862,7 +893,15 @@ async def test_duplicate_tool_result_request_is_claimed_once_without_new_turn()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_tool_result_cannot_start_a_new_conversation() -> None:
|
||||
async def test_a_relaunched_attempt_starts_a_new_conversation() -> None:
|
||||
"""What a building-mode relaunch sends after the model used some tools.
|
||||
|
||||
The service interrupts the CLI and resumes the same session with the
|
||||
guide added, so the replayed request carries results the abandoned
|
||||
conversation already consumed, plus the continuation message -- which no
|
||||
claim fingerprint can match. Refusing it killed the whole turn and sent
|
||||
the refusal to the user as the reply (SECRT-3333).
|
||||
"""
|
||||
agent_session = _FakeAgentSession(use_tool=True)
|
||||
transport = _FakeTransport(agent_session)
|
||||
async with CodexAnthropicGateway(
|
||||
@@ -876,34 +915,213 @@ async def test_completed_tool_result_cannot_start_a_new_conversation() -> None:
|
||||
headers=_headers(gateway),
|
||||
json=_tool_request("run"),
|
||||
)
|
||||
gateway_call_id = _tool_use_id(await first.json())
|
||||
continuation = _tool_result_request(gateway_call_id, "one result")
|
||||
completed = await client.post(
|
||||
abandoned_call_id = _tool_use_id(await first.json())
|
||||
answered = await client.post(
|
||||
f"{gateway.base_url}/v1/messages",
|
||||
headers=_headers(gateway),
|
||||
json=_tool_result_request(abandoned_call_id, "one result"),
|
||||
)
|
||||
assert answered.status == 200
|
||||
await answered.read()
|
||||
|
||||
relaunch = _tool_result_request(abandoned_call_id, "one result")
|
||||
relaunch_messages = relaunch["messages"]
|
||||
assert isinstance(relaunch_messages, list)
|
||||
replayed_result = relaunch_messages[0]
|
||||
relaunch_messages.append(
|
||||
{"role": "user", "content": "Building mode is now active"}
|
||||
)
|
||||
# Every request carries the tool list, this one included: it is
|
||||
# the one that has to open a conversation of its own.
|
||||
relaunch["tools"] = _tool_request("run")["tools"]
|
||||
resumed = await client.post(
|
||||
f"{gateway.base_url}/v1/messages",
|
||||
headers=_headers(gateway),
|
||||
json=relaunch,
|
||||
)
|
||||
assert resumed.status == 200
|
||||
resumed_call_id = _tool_use_id(await resumed.json())
|
||||
|
||||
# The resumed session replays the abandoned result alongside the
|
||||
# new one on every later round, so they must not refuse it either.
|
||||
continuation = _tool_result_request(resumed_call_id, "second result")
|
||||
continuation_messages = continuation["messages"]
|
||||
assert isinstance(continuation_messages, list)
|
||||
continuation_messages.insert(0, replayed_result)
|
||||
continued = await client.post(
|
||||
f"{gateway.base_url}/v1/messages",
|
||||
headers=_headers(gateway),
|
||||
json=continuation,
|
||||
)
|
||||
assert completed.status == 200
|
||||
await completed.read()
|
||||
assert continued.status == 200
|
||||
await continued.read()
|
||||
|
||||
replay = _tool_result_request(gateway_call_id, "one result")
|
||||
replay_messages = replay["messages"]
|
||||
assert isinstance(replay_messages, list)
|
||||
replay_messages.insert(
|
||||
0,
|
||||
{"role": "user", "content": "different request fingerprint"},
|
||||
)
|
||||
duplicate = await client.post(
|
||||
assert resumed_call_id != abandoned_call_id
|
||||
assert len(agent_session.requests) == 2
|
||||
assert "Building mode is now active" in agent_session.requests[1].prompt
|
||||
assert agent_session.tool_result == CodexDynamicToolResult(
|
||||
content="second result",
|
||||
success=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_resent_relaunch_is_replayed_without_a_second_invocation() -> None:
|
||||
"""A relaunch is resent verbatim like any other request, and costs quota.
|
||||
|
||||
Nothing claimable is left behind a relaunch, so its fingerprint is the
|
||||
only thing that can answer its own resend: unregistered, every CLI retry
|
||||
opens another Codex invocation and leaves it waiting on tools.
|
||||
"""
|
||||
agent_session = _FakeAgentSession(use_tool=True)
|
||||
transport = _FakeTransport(agent_session)
|
||||
async with CodexAnthropicGateway(
|
||||
credential_lease=_lease(),
|
||||
model="gpt-5.6-terra",
|
||||
transport=transport,
|
||||
) as gateway:
|
||||
async with ClientSession() as client:
|
||||
first = await client.post(
|
||||
f"{gateway.base_url}/v1/messages",
|
||||
headers=_headers(gateway),
|
||||
json=replay,
|
||||
json=_tool_request("run"),
|
||||
)
|
||||
duplicate_payload = await duplicate.json()
|
||||
abandoned_call_id = _tool_use_id(await first.json())
|
||||
answered = await client.post(
|
||||
f"{gateway.base_url}/v1/messages",
|
||||
headers=_headers(gateway),
|
||||
json=_tool_result_request(abandoned_call_id, "one result"),
|
||||
)
|
||||
assert answered.status == 200
|
||||
await answered.read()
|
||||
|
||||
assert duplicate.status == 409
|
||||
assert duplicate_payload["error"]["message"] == (
|
||||
"This tool-result request refers to a completed model call"
|
||||
)
|
||||
relaunch = _tool_result_request(abandoned_call_id, "one result")
|
||||
relaunch_messages = relaunch["messages"]
|
||||
assert isinstance(relaunch_messages, list)
|
||||
relaunch_messages.append(
|
||||
{"role": "user", "content": "Building mode is now active"}
|
||||
)
|
||||
relaunch["tools"] = _tool_request("run")["tools"]
|
||||
resumed = await client.post(
|
||||
f"{gateway.base_url}/v1/messages",
|
||||
headers=_headers(gateway),
|
||||
json=relaunch,
|
||||
)
|
||||
assert resumed.status == 200
|
||||
resumed_body = await resumed.json()
|
||||
|
||||
# What the CLI sends when it times out or disconnects waiting on
|
||||
# the relaunch: the same bytes again.
|
||||
resent = await client.post(
|
||||
f"{gateway.base_url}/v1/messages",
|
||||
headers=_headers(gateway),
|
||||
json=relaunch,
|
||||
)
|
||||
assert resent.status == 200
|
||||
resent_body = await resent.json()
|
||||
|
||||
assert resent_body == resumed_body
|
||||
assert gateway.last_refusal is None
|
||||
assert len(agent_session.requests) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_retry_after_a_failed_relaunch_opens_a_new_attempt() -> None:
|
||||
"""A relaunch that dies leaves nothing to replay, so its retry must run.
|
||||
|
||||
Holding the fingerprint past a failure would answer the CLI's retry with
|
||||
409 "already accepted" -- the refusal that killed the turn in SECRT-3333,
|
||||
reintroduced one layer down.
|
||||
"""
|
||||
agent_session = _FailingOnTheRelaunchSession()
|
||||
transport = _FakeTransport(agent_session)
|
||||
async with CodexAnthropicGateway(
|
||||
credential_lease=_lease(),
|
||||
model="gpt-5.6-terra",
|
||||
transport=transport,
|
||||
) as gateway:
|
||||
async with ClientSession() as client:
|
||||
first = await client.post(
|
||||
f"{gateway.base_url}/v1/messages",
|
||||
headers=_headers(gateway),
|
||||
json=_tool_request("run"),
|
||||
)
|
||||
abandoned_call_id = _tool_use_id(await first.json())
|
||||
answered = await client.post(
|
||||
f"{gateway.base_url}/v1/messages",
|
||||
headers=_headers(gateway),
|
||||
json=_tool_result_request(abandoned_call_id, "one result"),
|
||||
)
|
||||
assert answered.status == 200
|
||||
await answered.read()
|
||||
|
||||
relaunch = _tool_result_request(abandoned_call_id, "one result")
|
||||
relaunch_messages = relaunch["messages"]
|
||||
assert isinstance(relaunch_messages, list)
|
||||
relaunch_messages.append(
|
||||
{"role": "user", "content": "Building mode is now active"}
|
||||
)
|
||||
relaunch["tools"] = _tool_request("run")["tools"]
|
||||
failed = await client.post(
|
||||
f"{gateway.base_url}/v1/messages",
|
||||
headers=_headers(gateway),
|
||||
json=relaunch,
|
||||
)
|
||||
assert failed.status >= 500
|
||||
await failed.read()
|
||||
|
||||
retried = await client.post(
|
||||
f"{gateway.base_url}/v1/messages",
|
||||
headers=_headers(gateway),
|
||||
json=relaunch,
|
||||
)
|
||||
assert retried.status == 200
|
||||
recovered = await retried.json()
|
||||
|
||||
assert recovered["content"][0]["text"] == "recovered"
|
||||
assert len(agent_session.requests) == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_two_results_for_one_call_are_refused_and_remembered() -> None:
|
||||
"""The guard a relaunch must not cost: one call, one answer.
|
||||
|
||||
``last_refusal`` is how the service layer tells our own 4xx, which the
|
||||
CLI reports as an assistant-message error, from the model's answer.
|
||||
"""
|
||||
agent_session = _FakeAgentSession(use_tool=True)
|
||||
transport = _FakeTransport(agent_session)
|
||||
async with CodexAnthropicGateway(
|
||||
credential_lease=_lease(),
|
||||
model="gpt-5.6-terra",
|
||||
transport=transport,
|
||||
) as gateway:
|
||||
assert gateway.last_refusal is None
|
||||
async with ClientSession() as client:
|
||||
first = await client.post(
|
||||
f"{gateway.base_url}/v1/messages",
|
||||
headers=_headers(gateway),
|
||||
json=_tool_request("run"),
|
||||
)
|
||||
gateway_call_id = _tool_use_id(await first.json())
|
||||
answered = await client.post(
|
||||
f"{gateway.base_url}/v1/messages",
|
||||
headers=_headers(gateway),
|
||||
json=_tool_result_request(gateway_call_id, "one result"),
|
||||
)
|
||||
assert answered.status == 200
|
||||
await answered.read()
|
||||
|
||||
conflicting = await client.post(
|
||||
f"{gateway.base_url}/v1/messages",
|
||||
headers=_headers(gateway),
|
||||
json=_tool_result_request(gateway_call_id, "another result"),
|
||||
)
|
||||
refused = await conflicting.json()
|
||||
|
||||
assert conflicting.status == 409
|
||||
assert refused["error"]["message"].startswith("Conflicting result for tool_use_id")
|
||||
assert gateway.last_refusal == refused["error"]["message"]
|
||||
assert len(agent_session.requests) == 1
|
||||
|
||||
|
||||
|
||||
@@ -121,7 +121,7 @@ def build_sdk_env(
|
||||
route through the proxy, with the headers and body fields OpenRouter
|
||||
Broadcast attributes generations by (user, chat session,
|
||||
environment). Only this mode reaches Broadcast, so only its
|
||||
generations show in PostHog LLM analytics and Langfuse broadcast.
|
||||
generations show in PostHog LLM analytics.
|
||||
|
||||
All modes receive workspace isolation (``CLAUDE_CODE_TMPDIR``) and
|
||||
security hardening env vars to prevent .claude.md loading, prompt
|
||||
|
||||
@@ -66,10 +66,6 @@ def _make_patches():
|
||||
(f"{_SVC}.AsyncClusterLock", dict(side_effect=_make_lock_mock())),
|
||||
(f"{_SVC}._make_sdk_cwd", dict(return_value="/tmp/test-sdk-cwd")),
|
||||
("os.makedirs", {}),
|
||||
(
|
||||
f"{_SVC}.propagate_attributes",
|
||||
dict(return_value=MagicMock(__enter__=MagicMock(), __exit__=MagicMock())),
|
||||
),
|
||||
(
|
||||
f"{_SVC}._build_system_prompt",
|
||||
dict(new_callable=AsyncMock, return_value=("system prompt", None)),
|
||||
@@ -124,7 +120,6 @@ def _make_patches():
|
||||
thinking_standard_model="anthropic/claude-sonnet-4-6",
|
||||
),
|
||||
),
|
||||
(f"{_SVC}.get_user_tier", dict(new_callable=AsyncMock, return_value=None)),
|
||||
(
|
||||
f"{_SVC}.resolve_turn_budget",
|
||||
dict(
|
||||
|
||||
@@ -31,6 +31,7 @@ Examples
|
||||
@@agptfile:/home/user/script.sh
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import itertools
|
||||
import logging
|
||||
import re
|
||||
@@ -602,7 +603,9 @@ async def _expand_bare_ref(
|
||||
# file) instead of silently returning garbled bytes.
|
||||
strict = fmt in BINARY_FORMATS
|
||||
try:
|
||||
parsed = parse_file_content(content, fmt, strict=strict)
|
||||
parsed = await asyncio.to_thread(
|
||||
parse_file_content, content, fmt, strict=strict
|
||||
)
|
||||
except PARSE_EXCEPTIONS as exc:
|
||||
raise FileRefExpansionError(f"Failed to parse {fmt} file: {exc}") from exc
|
||||
# Normalize bytes fallback to str so tools never
|
||||
|
||||
@@ -1,119 +0,0 @@
|
||||
"""langsmith's Claude Agent SDK tracing must survive ``_iter_sdk_messages``.
|
||||
|
||||
The tracing wrapper around ``receive_response()`` keeps the
|
||||
``claude.conversation`` run in a ContextVar that it sets while the first
|
||||
message is fetched, and parents every ``claude.assistant.turn`` run on it.
|
||||
``_iter_sdk_messages`` fetches each message in its own task, so unless every
|
||||
fetch runs in the same context, main-agent replies after the first message
|
||||
find no parent and are silently dropped from the trace.
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Any, cast
|
||||
from unittest.mock import patch
|
||||
|
||||
import claude_agent_sdk
|
||||
import pytest
|
||||
from claude_agent_sdk import (
|
||||
AssistantMessage,
|
||||
ClaudeAgentOptions,
|
||||
ClaudeSDKClient,
|
||||
ResultMessage,
|
||||
SystemMessage,
|
||||
TextBlock,
|
||||
ToolResultBlock,
|
||||
ToolUseBlock,
|
||||
UserMessage,
|
||||
)
|
||||
from langsmith import tracing_context
|
||||
from langsmith.integrations.claude_agent_sdk import configure_claude_agent_sdk
|
||||
from langsmith.run_trees import RunTree
|
||||
|
||||
from .service import _iter_sdk_messages
|
||||
|
||||
|
||||
class _ScriptedClient:
|
||||
"""Stands in for ``ClaudeSDKClient`` and replays one main-agent turn."""
|
||||
|
||||
def __init__(self, options: ClaudeAgentOptions) -> None:
|
||||
self.options = options
|
||||
|
||||
async def query(self, prompt: str) -> None:
|
||||
return None
|
||||
|
||||
async def receive_response(self) -> AsyncGenerator[Any, None]:
|
||||
yield SystemMessage(subtype="init", data={})
|
||||
yield AssistantMessage(
|
||||
content=[ToolUseBlock(id="toolu_1", name="bash_exec", input={})],
|
||||
model="claude-test",
|
||||
message_id="msg_1",
|
||||
)
|
||||
yield UserMessage(content=[ToolResultBlock(tool_use_id="toolu_1")])
|
||||
yield AssistantMessage(
|
||||
content=[TextBlock(text="done")], model="claude-test", message_id="msg_2"
|
||||
)
|
||||
yield ResultMessage(
|
||||
subtype="success",
|
||||
duration_ms=1,
|
||||
duration_api_ms=1,
|
||||
is_error=False,
|
||||
num_turns=2,
|
||||
session_id="sess-1",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def traced_client() -> type[_ScriptedClient]:
|
||||
"""Wrap the scripted client exactly as production wraps ``ClaudeSDKClient``.
|
||||
|
||||
``configure_claude_agent_sdk`` patches whichever class
|
||||
``claude_agent_sdk.ClaudeSDKClient`` names, so pointing that name at the
|
||||
scripted client for the call wraps it through the public API. The real
|
||||
tool class and the global trace config are left untouched.
|
||||
"""
|
||||
with (
|
||||
patch.object(claude_agent_sdk, "ClaudeSDKClient", _ScriptedClient),
|
||||
patch.object(claude_agent_sdk, "SdkMcpTool", None),
|
||||
patch("langsmith.integrations.claude_agent_sdk.set_tracing_config"),
|
||||
):
|
||||
assert configure_claude_agent_sdk()
|
||||
return _ScriptedClient
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def posted_runs(monkeypatch: pytest.MonkeyPatch) -> list[RunTree]:
|
||||
"""Record every run the tracing wrapper posts instead of exporting it."""
|
||||
posted: list[RunTree] = []
|
||||
monkeypatch.setattr(RunTree, "post", lambda self, *_, **__: posted.append(self))
|
||||
monkeypatch.setattr(RunTree, "patch", lambda self, *_, **__: None)
|
||||
return posted
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_main_agent_replies_get_assistant_turn_runs(
|
||||
traced_client: type[_ScriptedClient],
|
||||
posted_runs: list[RunTree],
|
||||
) -> None:
|
||||
client = traced_client(options=ClaudeAgentOptions())
|
||||
await client.query("list the files")
|
||||
|
||||
with tracing_context(enabled=True):
|
||||
received = [
|
||||
msg
|
||||
async for msg in _iter_sdk_messages(cast(ClaudeSDKClient, client))
|
||||
if msg is not None
|
||||
]
|
||||
|
||||
assert [type(msg).__name__ for msg in received] == [
|
||||
"SystemMessage",
|
||||
"AssistantMessage",
|
||||
"UserMessage",
|
||||
"AssistantMessage",
|
||||
"ResultMessage",
|
||||
]
|
||||
conversation = next(run for run in posted_runs if run.name == "claude.conversation")
|
||||
turns = [run for run in posted_runs if run.name == "claude.assistant.turn"]
|
||||
assert [turn.parent_run_id for turn in turns] == [
|
||||
conversation.id,
|
||||
conversation.id,
|
||||
], "each main-agent reply should get a turn run under the conversation"
|
||||
@@ -37,7 +37,6 @@ from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import httpx
|
||||
from langfuse import get_client
|
||||
|
||||
from backend.copilot.token_tracking import persist_and_record_usage
|
||||
from backend.util import json
|
||||
@@ -336,7 +335,6 @@ async def record_turn_cost_from_openrouter(
|
||||
fallback_cost_usd: float | None,
|
||||
api_key: str | None,
|
||||
log_prefix: str,
|
||||
langfuse_trace_id: str | None = None,
|
||||
absorbed_generation_ids: Collection[str] = (),
|
||||
) -> None:
|
||||
"""Persist turn cost from OpenRouter's authoritative ``/generation``.
|
||||
@@ -430,10 +428,8 @@ async def record_turn_cost_from_openrouter(
|
||||
results = []
|
||||
|
||||
fetched = [r for r in results if isinstance(r, (int, float))]
|
||||
cost_source = "fallback"
|
||||
if fetched and len(fetched) == len(generation_ids):
|
||||
real_cost: float | None = sum(fetched)
|
||||
cost_source = "openrouter"
|
||||
# Log real (OpenRouter billed) vs CLI rate-card estimate so an
|
||||
# operator can spot divergence without querying OpenRouter by
|
||||
# hand. Under-count typically means a gen-ID source we don't
|
||||
@@ -493,36 +489,3 @@ async def record_turn_cost_from_openrouter(
|
||||
log_prefix,
|
||||
exc,
|
||||
)
|
||||
|
||||
# Backfill the Langfuse trace with reconciled cost + token usage. The
|
||||
# OTel span for the turn closes before this background task runs, so the
|
||||
# Langfuse trace UI otherwise shows the SDK-CLI rate-card estimate (which
|
||||
# for non-Anthropic OpenRouter routes is wildly wrong — Sonnet pricing on
|
||||
# Kimi tokens, ~5x too high). Emitting a child event with the real
|
||||
# numbers gives operators a single Langfuse view per turn instead of
|
||||
# cross-referencing pod logs.
|
||||
if langfuse_trace_id and real_cost is not None:
|
||||
try:
|
||||
get_client().create_event(
|
||||
trace_context={"trace_id": langfuse_trace_id},
|
||||
name="openrouter-cost-reconcile",
|
||||
metadata={
|
||||
"cost_usd": real_cost,
|
||||
"cost_source": cost_source,
|
||||
"fallback_cost_usd": fallback_cost_usd,
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"cache_read_tokens": cache_read_tokens,
|
||||
"cache_creation_tokens": cache_creation_tokens,
|
||||
"resolved_generation_id_count": len(fetched),
|
||||
"generation_id_count": len(generation_ids),
|
||||
"model": model,
|
||||
"provider": "open_router",
|
||||
},
|
||||
)
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"%s[cost-record] Langfuse event emit failed",
|
||||
log_prefix,
|
||||
exc_info=True,
|
||||
)
|
||||
@@ -3,7 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
@@ -518,198 +518,3 @@ class TestRecordTurnCostFromOpenRouter:
|
||||
assert mock_persist.call_args.kwargs["cost_usd"] == pytest.approx(
|
||||
sum(costs_by_id.values()), rel=1e-9
|
||||
)
|
||||
|
||||
|
||||
class TestLangfuseTraceBackfill:
|
||||
"""Reconciled cost is mirrored back onto the Langfuse trace as a
|
||||
child event so operators see the real number alongside the SDK-CLI
|
||||
rate-card estimate the OTel bridge already wrote."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_event_emitted_with_real_cost_when_trace_id_supplied(self):
|
||||
real_cost = 0.025
|
||||
mock_lf = MagicMock()
|
||||
|
||||
async def _get(_self, _url, **_kwargs):
|
||||
return httpx.Response(200, json=_mock_generation_response(real_cost))
|
||||
|
||||
with (
|
||||
patch(
|
||||
"backend.copilot.sdk.openrouter_cost.persist_and_record_usage",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch("httpx.AsyncClient.get", new=_get),
|
||||
patch(
|
||||
"backend.copilot.sdk.openrouter_cost.get_client",
|
||||
return_value=mock_lf,
|
||||
),
|
||||
):
|
||||
await record_turn_cost_from_openrouter(
|
||||
session=_session(),
|
||||
user_id="u1",
|
||||
model="moonshotai/kimi-k2.6",
|
||||
prompt_tokens=29669,
|
||||
completion_tokens=280,
|
||||
cache_read_tokens=0,
|
||||
cache_creation_tokens=0,
|
||||
generation_ids=["gen-1"],
|
||||
cli_project_dir=None,
|
||||
cli_session_id=None,
|
||||
turn_start_ts=None,
|
||||
fallback_cost_usd=0.018,
|
||||
api_key="sk-or-test",
|
||||
log_prefix="[test]",
|
||||
langfuse_trace_id="trace-abc",
|
||||
)
|
||||
|
||||
mock_lf.create_event.assert_called_once()
|
||||
kwargs = mock_lf.create_event.call_args.kwargs
|
||||
assert kwargs["trace_context"] == {"trace_id": "trace-abc"}
|
||||
assert kwargs["name"] == "openrouter-cost-reconcile"
|
||||
meta = kwargs["metadata"]
|
||||
assert meta["cost_usd"] == pytest.approx(real_cost, rel=1e-9)
|
||||
assert meta["cost_source"] == "openrouter"
|
||||
assert meta["fallback_cost_usd"] == pytest.approx(0.018, rel=1e-9)
|
||||
assert meta["resolved_generation_id_count"] == 1
|
||||
assert meta["generation_id_count"] == 1
|
||||
assert meta["prompt_tokens"] == 29669
|
||||
assert meta["completion_tokens"] == 280
|
||||
assert meta["model"] == "moonshotai/kimi-k2.6"
|
||||
assert meta["provider"] == "open_router"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_event_marks_fallback_when_lookup_partial(self):
|
||||
"""When some gen-ID lookups fail, real_cost falls back to the
|
||||
rate-card estimate. The Langfuse event must mark cost_source as
|
||||
``"fallback"`` so operators don't mistake it for an authoritative
|
||||
OpenRouter reconciliation.
|
||||
"""
|
||||
mock_lf = MagicMock()
|
||||
call_count = {"n": 0}
|
||||
|
||||
async def _get(_self, _url, **_kwargs):
|
||||
call_count["n"] += 1
|
||||
if call_count["n"] == 1:
|
||||
return httpx.Response(200, json=_mock_generation_response(0.012))
|
||||
return httpx.Response(404, json={"error": "not found"})
|
||||
|
||||
with (
|
||||
patch(
|
||||
"backend.copilot.sdk.openrouter_cost.persist_and_record_usage",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch("httpx.AsyncClient.get", new=_get),
|
||||
patch(
|
||||
"backend.copilot.sdk.openrouter_cost.asyncio.sleep",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"backend.copilot.sdk.openrouter_cost.get_client",
|
||||
return_value=mock_lf,
|
||||
),
|
||||
):
|
||||
await record_turn_cost_from_openrouter(
|
||||
session=_session(),
|
||||
user_id="u1",
|
||||
model="moonshotai/kimi-k2.6",
|
||||
prompt_tokens=100,
|
||||
completion_tokens=10,
|
||||
cache_read_tokens=0,
|
||||
cache_creation_tokens=0,
|
||||
generation_ids=["gen-1", "gen-2"],
|
||||
cli_project_dir=None,
|
||||
cli_session_id=None,
|
||||
turn_start_ts=None,
|
||||
fallback_cost_usd=0.018,
|
||||
api_key="sk-or-test",
|
||||
log_prefix="[test]",
|
||||
langfuse_trace_id="trace-partial",
|
||||
)
|
||||
|
||||
mock_lf.create_event.assert_called_once()
|
||||
meta = mock_lf.create_event.call_args.kwargs["metadata"]
|
||||
assert meta["cost_source"] == "fallback"
|
||||
assert meta["cost_usd"] == pytest.approx(0.018, rel=1e-9)
|
||||
assert meta["resolved_generation_id_count"] == 1
|
||||
assert meta["generation_id_count"] == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_event_when_trace_id_missing(self):
|
||||
"""Anthropic-direct turns and any path without an active Langfuse
|
||||
OTel context don't have a trace_id — backfill is a no-op."""
|
||||
mock_lf = MagicMock()
|
||||
|
||||
async def _get(_self, _url, **_kwargs):
|
||||
return httpx.Response(200, json=_mock_generation_response(0.01))
|
||||
|
||||
with (
|
||||
patch(
|
||||
"backend.copilot.sdk.openrouter_cost.persist_and_record_usage",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch("httpx.AsyncClient.get", new=_get),
|
||||
patch(
|
||||
"backend.copilot.sdk.openrouter_cost.get_client",
|
||||
return_value=mock_lf,
|
||||
),
|
||||
):
|
||||
await record_turn_cost_from_openrouter(
|
||||
session=_session(),
|
||||
user_id="u1",
|
||||
model="moonshotai/kimi-k2.6",
|
||||
prompt_tokens=100,
|
||||
completion_tokens=10,
|
||||
cache_read_tokens=0,
|
||||
cache_creation_tokens=0,
|
||||
generation_ids=["gen-1"],
|
||||
cli_project_dir=None,
|
||||
cli_session_id=None,
|
||||
turn_start_ts=None,
|
||||
fallback_cost_usd=0.005,
|
||||
api_key="sk-or-test",
|
||||
log_prefix="[test]",
|
||||
langfuse_trace_id=None,
|
||||
)
|
||||
|
||||
mock_lf.create_event.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_backfill_failure_swallowed_does_not_break_persist(self):
|
||||
"""If Langfuse is unavailable, the reconcile still persists the
|
||||
cost row — backfill is best-effort and must not raise."""
|
||||
mock_lf = MagicMock()
|
||||
mock_lf.create_event = MagicMock(side_effect=RuntimeError("network down"))
|
||||
|
||||
async def _get(_self, _url, **_kwargs):
|
||||
return httpx.Response(200, json=_mock_generation_response(0.02))
|
||||
|
||||
with (
|
||||
patch(
|
||||
"backend.copilot.sdk.openrouter_cost.persist_and_record_usage",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_persist,
|
||||
patch("httpx.AsyncClient.get", new=_get),
|
||||
patch(
|
||||
"backend.copilot.sdk.openrouter_cost.get_client",
|
||||
return_value=mock_lf,
|
||||
),
|
||||
):
|
||||
await record_turn_cost_from_openrouter(
|
||||
session=_session(),
|
||||
user_id="u1",
|
||||
model="moonshotai/kimi-k2.6",
|
||||
prompt_tokens=100,
|
||||
completion_tokens=10,
|
||||
cache_read_tokens=0,
|
||||
cache_creation_tokens=0,
|
||||
generation_ids=["gen-1"],
|
||||
cli_project_dir=None,
|
||||
cli_session_id=None,
|
||||
turn_start_ts=None,
|
||||
fallback_cost_usd=0.018,
|
||||
api_key="sk-or-test",
|
||||
log_prefix="[test]",
|
||||
langfuse_trace_id="trace-xyz",
|
||||
)
|
||||
|
||||
mock_persist.assert_called_once()
|
||||
@@ -1,274 +0,0 @@
|
||||
"""Tests for OTEL tracing setup in the SDK copilot path."""
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
class TestSetupLangfuseOtel:
|
||||
"""Tests for _setup_langfuse_otel()."""
|
||||
|
||||
def test_noop_when_langfuse_not_configured(self):
|
||||
"""No env vars should be set when Langfuse credentials are missing."""
|
||||
with patch(
|
||||
"backend.copilot.sdk.service._is_langfuse_configured", return_value=False
|
||||
):
|
||||
from backend.copilot.sdk.service import _setup_langfuse_otel
|
||||
|
||||
# Clear any previously set env vars
|
||||
env_keys = [
|
||||
"LANGSMITH_OTEL_ENABLED",
|
||||
"LANGSMITH_OTEL_ONLY",
|
||||
"LANGSMITH_TRACING",
|
||||
"OTEL_EXPORTER_OTLP_ENDPOINT",
|
||||
"OTEL_EXPORTER_OTLP_HEADERS",
|
||||
"OTEL_PYTHON_EXPORTER_OTLP_HTTP_TRACES_CREDENTIAL_PROVIDER",
|
||||
]
|
||||
saved = {k: os.environ.pop(k, None) for k in env_keys}
|
||||
try:
|
||||
_setup_langfuse_otel()
|
||||
for key in env_keys:
|
||||
assert key not in os.environ, f"{key} should not be set"
|
||||
finally:
|
||||
for k, v in saved.items():
|
||||
if v is not None:
|
||||
os.environ[k] = v
|
||||
|
||||
def test_sets_env_vars_when_langfuse_configured(self):
|
||||
"""OTEL env vars should be set when Langfuse credentials exist."""
|
||||
mock_settings = MagicMock()
|
||||
mock_settings.secrets.langfuse_public_key = "pk-test-123"
|
||||
mock_settings.secrets.langfuse_secret_key = "sk-test-456" # pragma: allowlist secret # noqa: E501; fmt: skip
|
||||
mock_settings.secrets.langfuse_host = "https://langfuse.example.com"
|
||||
mock_settings.secrets.langfuse_tracing_environment = "test"
|
||||
|
||||
with (
|
||||
patch(
|
||||
"backend.copilot.sdk.service._is_langfuse_configured",
|
||||
return_value=True,
|
||||
),
|
||||
patch("backend.copilot.sdk.service.Settings", return_value=mock_settings),
|
||||
patch(
|
||||
"backend.copilot.sdk.service.install_trace_redaction"
|
||||
) as mock_redaction,
|
||||
patch(
|
||||
"backend.copilot.sdk.service.configure_claude_agent_sdk",
|
||||
return_value=True,
|
||||
) as mock_configure,
|
||||
):
|
||||
from backend.copilot.sdk.service import _setup_langfuse_otel
|
||||
|
||||
# Clear env vars so setdefault works
|
||||
env_keys = [
|
||||
"LANGSMITH_OTEL_ENABLED",
|
||||
"LANGSMITH_OTEL_ONLY",
|
||||
"LANGSMITH_TRACING",
|
||||
"OTEL_EXPORTER_OTLP_ENDPOINT",
|
||||
"OTEL_EXPORTER_OTLP_HEADERS",
|
||||
"OTEL_PYTHON_EXPORTER_OTLP_HTTP_TRACES_CREDENTIAL_PROVIDER",
|
||||
"OTEL_RESOURCE_ATTRIBUTES",
|
||||
]
|
||||
saved = {k: os.environ.pop(k, None) for k in env_keys}
|
||||
try:
|
||||
_setup_langfuse_otel()
|
||||
|
||||
assert os.environ["LANGSMITH_OTEL_ENABLED"] == "true"
|
||||
assert os.environ["LANGSMITH_OTEL_ONLY"] == "true"
|
||||
assert os.environ["LANGSMITH_TRACING"] == "true"
|
||||
assert (
|
||||
os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"]
|
||||
== "https://langfuse.example.com/api/public/otel"
|
||||
)
|
||||
assert "Authorization=Basic" in os.environ["OTEL_EXPORTER_OTLP_HEADERS"]
|
||||
assert (
|
||||
os.environ[
|
||||
"OTEL_PYTHON_EXPORTER_OTLP_HTTP_TRACES_CREDENTIAL_PROVIDER"
|
||||
]
|
||||
== "autogpt-copilot-traces"
|
||||
)
|
||||
assert (
|
||||
os.environ["OTEL_RESOURCE_ATTRIBUTES"]
|
||||
== "langfuse.environment=test"
|
||||
)
|
||||
|
||||
mock_configure.assert_called_once_with(tags=["sdk"])
|
||||
mock_redaction.assert_called_once_with()
|
||||
finally:
|
||||
for k, v in saved.items():
|
||||
if v is not None:
|
||||
os.environ[k] = v
|
||||
elif k in os.environ:
|
||||
del os.environ[k]
|
||||
|
||||
def test_existing_env_vars_not_overwritten(self):
|
||||
"""Explicit env-var overrides should not be clobbered."""
|
||||
mock_settings = MagicMock()
|
||||
mock_settings.secrets.langfuse_public_key = "pk-test"
|
||||
mock_settings.secrets.langfuse_secret_key = "sk-test" # pragma: allowlist secret # noqa: E501; fmt: skip
|
||||
mock_settings.secrets.langfuse_host = "https://langfuse.example.com"
|
||||
|
||||
with (
|
||||
patch(
|
||||
"backend.copilot.sdk.service._is_langfuse_configured",
|
||||
return_value=True,
|
||||
),
|
||||
patch("backend.copilot.sdk.service.Settings", return_value=mock_settings),
|
||||
patch("backend.copilot.sdk.service.install_trace_redaction"),
|
||||
patch(
|
||||
"backend.copilot.sdk.service.configure_claude_agent_sdk",
|
||||
return_value=True,
|
||||
),
|
||||
):
|
||||
from backend.copilot.sdk.service import _setup_langfuse_otel
|
||||
|
||||
env_keys = [
|
||||
"OTEL_EXPORTER_OTLP_ENDPOINT",
|
||||
"OTEL_PYTHON_EXPORTER_OTLP_HTTP_TRACES_CREDENTIAL_PROVIDER",
|
||||
]
|
||||
saved = {key: os.environ.get(key) for key in env_keys}
|
||||
try:
|
||||
os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = "https://custom.endpoint/v1"
|
||||
os.environ[
|
||||
"OTEL_PYTHON_EXPORTER_OTLP_HTTP_TRACES_CREDENTIAL_PROVIDER"
|
||||
] = "custom-provider"
|
||||
_setup_langfuse_otel()
|
||||
assert (
|
||||
os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"]
|
||||
== "https://custom.endpoint/v1"
|
||||
)
|
||||
assert (
|
||||
os.environ[
|
||||
"OTEL_PYTHON_EXPORTER_OTLP_HTTP_TRACES_CREDENTIAL_PROVIDER"
|
||||
]
|
||||
== "custom-provider"
|
||||
)
|
||||
finally:
|
||||
for key, value in saved.items():
|
||||
if value is not None:
|
||||
os.environ[key] = value
|
||||
else:
|
||||
os.environ.pop(key, None)
|
||||
|
||||
def test_sdk_left_untraced_when_redaction_cannot_install(self):
|
||||
"""Tracing without the Google data redaction would export it."""
|
||||
mock_settings = MagicMock()
|
||||
mock_settings.secrets.langfuse_public_key = "pk-test"
|
||||
mock_settings.secrets.langfuse_secret_key = "sk-test"
|
||||
mock_settings.secrets.langfuse_host = "https://langfuse.example.com"
|
||||
mock_settings.secrets.langfuse_tracing_environment = "test"
|
||||
|
||||
with (
|
||||
patch(
|
||||
"backend.copilot.sdk.service._is_langfuse_configured",
|
||||
return_value=True,
|
||||
),
|
||||
patch("backend.copilot.sdk.service.Settings", return_value=mock_settings),
|
||||
patch(
|
||||
"backend.copilot.sdk.service.install_trace_redaction",
|
||||
side_effect=RuntimeError("no client"),
|
||||
),
|
||||
patch(
|
||||
"backend.copilot.sdk.service.configure_claude_agent_sdk"
|
||||
) as mock_configure,
|
||||
patch.dict(os.environ, {}, clear=False),
|
||||
):
|
||||
from backend.copilot.sdk.service import _setup_langfuse_otel
|
||||
|
||||
_setup_langfuse_otel()
|
||||
|
||||
mock_configure.assert_not_called()
|
||||
|
||||
def test_tracing_left_off_when_redaction_cannot_install(self):
|
||||
"""LangSmith reads these to decide whether to trace at all, so they
|
||||
must not outlive a redaction that failed to install."""
|
||||
mock_settings = MagicMock()
|
||||
mock_settings.secrets.langfuse_public_key = "pk-test"
|
||||
mock_settings.secrets.langfuse_secret_key = "sk-test"
|
||||
mock_settings.secrets.langfuse_host = "https://langfuse.example.com"
|
||||
mock_settings.secrets.langfuse_tracing_environment = "test"
|
||||
keys = [
|
||||
"LANGSMITH_OTEL_ENABLED",
|
||||
"LANGSMITH_OTEL_ONLY",
|
||||
"LANGSMITH_TRACING",
|
||||
"OTEL_EXPORTER_OTLP_HEADERS",
|
||||
]
|
||||
seen_at_install: dict[str, str | None] = {}
|
||||
|
||||
def failing_install():
|
||||
seen_at_install.update({key: os.environ.get(key) for key in keys})
|
||||
raise RuntimeError("no client")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"backend.copilot.sdk.service._is_langfuse_configured",
|
||||
return_value=True,
|
||||
),
|
||||
patch("backend.copilot.sdk.service.Settings", return_value=mock_settings),
|
||||
patch(
|
||||
"backend.copilot.sdk.service.install_trace_redaction",
|
||||
side_effect=failing_install,
|
||||
),
|
||||
patch("backend.copilot.sdk.service.configure_claude_agent_sdk"),
|
||||
patch.dict(os.environ, {"OTEL_EXPORTER_OTLP_ENDPOINT": "https://mine"}),
|
||||
):
|
||||
for key in keys:
|
||||
os.environ.pop(key, None)
|
||||
from backend.copilot.sdk.service import _setup_langfuse_otel
|
||||
|
||||
_setup_langfuse_otel()
|
||||
|
||||
# The LangSmith client picks its OTEL mode up as it is built.
|
||||
assert seen_at_install["LANGSMITH_OTEL_ONLY"] == "true"
|
||||
assert seen_at_install["LANGSMITH_TRACING"] is None
|
||||
for key in keys:
|
||||
assert key not in os.environ, f"{key} outlived the failed install"
|
||||
assert os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] == "https://mine"
|
||||
|
||||
def test_graceful_failure_on_exception(self):
|
||||
"""Setup should not raise even if internal code fails."""
|
||||
with (
|
||||
patch(
|
||||
"backend.copilot.sdk.service._is_langfuse_configured",
|
||||
return_value=True,
|
||||
),
|
||||
patch(
|
||||
"backend.copilot.sdk.service.Settings",
|
||||
side_effect=RuntimeError("settings unavailable"),
|
||||
),
|
||||
):
|
||||
from backend.copilot.sdk.service import _setup_langfuse_otel
|
||||
|
||||
# Should not raise — just logs and returns
|
||||
_setup_langfuse_otel()
|
||||
|
||||
|
||||
class TestPropagateAttributesImport:
|
||||
"""Verify langfuse.propagate_attributes is available."""
|
||||
|
||||
def test_propagate_attributes_is_importable(self):
|
||||
from langfuse import propagate_attributes
|
||||
|
||||
assert callable(propagate_attributes)
|
||||
|
||||
def test_propagate_attributes_returns_context_manager(self):
|
||||
from langfuse import propagate_attributes
|
||||
|
||||
ctx = propagate_attributes(user_id="u1", session_id="s1", tags=["test"])
|
||||
assert hasattr(ctx, "__enter__")
|
||||
assert hasattr(ctx, "__exit__")
|
||||
|
||||
|
||||
class TestReceiveResponseCompat:
|
||||
"""Verify ClaudeSDKClient.receive_response() exists (langsmith patches it)."""
|
||||
|
||||
def test_receive_response_exists(self):
|
||||
from claude_agent_sdk import ClaudeSDKClient
|
||||
|
||||
assert hasattr(ClaudeSDKClient, "receive_response")
|
||||
|
||||
def test_receive_response_is_async_generator(self):
|
||||
import inspect
|
||||
|
||||
from claude_agent_sdk import ClaudeSDKClient
|
||||
|
||||
method = getattr(ClaudeSDKClient, "receive_response")
|
||||
assert inspect.isfunction(method) or inspect.ismethod(method)
|
||||
@@ -1061,10 +1061,6 @@ def _make_sdk_patches(
|
||||
),
|
||||
(f"{_SVC}._make_sdk_cwd", dict(return_value="/tmp/test-sdk-cwd")),
|
||||
("os.makedirs", {}),
|
||||
(
|
||||
f"{_SVC}.propagate_attributes",
|
||||
dict(return_value=MagicMock(__enter__=MagicMock(), __exit__=MagicMock())),
|
||||
),
|
||||
(
|
||||
f"{_SVC}._build_system_prompt",
|
||||
dict(new_callable=AsyncMock, return_value=("system prompt", None)),
|
||||
@@ -1130,7 +1126,6 @@ def _make_sdk_patches(
|
||||
thinking_standard_model="anthropic/claude-sonnet-4-6",
|
||||
),
|
||||
),
|
||||
(f"{_SVC}.get_user_tier", dict(new_callable=AsyncMock, return_value=None)),
|
||||
(
|
||||
"backend.copilot.gate.held.answered",
|
||||
dict(new_callable=AsyncMock, return_value=[]),
|
||||
|
||||
@@ -41,9 +41,6 @@ from claude_agent_sdk import (
|
||||
ToolUseBlock,
|
||||
)
|
||||
from claude_agent_sdk.types import SystemPromptPreset
|
||||
from langfuse import get_client, propagate_attributes
|
||||
from langsmith.integrations.claude_agent_sdk import configure_claude_agent_sdk
|
||||
from opentelemetry import trace as otel_trace
|
||||
from pydantic import BaseModel
|
||||
|
||||
from backend.blocks.desktop._common import workspace_volume_mounts
|
||||
@@ -59,7 +56,6 @@ from backend.copilot.budget_gate import gate_auto_continue
|
||||
from backend.copilot.budget_signal import build_turn_budget_block
|
||||
from backend.copilot.sandbox_status import build_sandbox_replaced_notice
|
||||
from backend.copilot.continuation import as_model_prompt
|
||||
from backend.copilot.feedback_db import RATEABLE_ROLES
|
||||
from backend.copilot.cli_session_entry import rewrite_user_entry
|
||||
from backend.copilot.legacy_first_turn_memory import without_stored_first_turn_memory
|
||||
from backend.copilot.graphiti.context import fetch_warm_context
|
||||
@@ -68,7 +64,6 @@ from backend.copilot.graphiti.context_marker import (
|
||||
append_injected_memory_block,
|
||||
strip_injected_memory_text,
|
||||
)
|
||||
from backend.copilot.google_data_redaction import install_trace_redaction
|
||||
from backend.copilot.host_fs import cli_tmpdir, save_new_file
|
||||
from backend.copilot.markers import append_error_marker
|
||||
from backend.copilot.provider_failure import ProviderFailure
|
||||
@@ -78,6 +73,8 @@ from backend.copilot.graphiti.ingest import enqueue_conversation_turn
|
||||
from backend.copilot.learning.capture import capture_chat_turn, find_initiating_user_row
|
||||
from backend.copilot.sdk.cli_prelaunch import CLIPrelaunch
|
||||
from backend.copilot.sdk.codex_compat_gateway import (
|
||||
GATEWAY_REFUSED_CODE,
|
||||
GATEWAY_REFUSED_MESSAGE,
|
||||
CodexAnthropicGateway,
|
||||
gateway_for_lease,
|
||||
)
|
||||
@@ -109,7 +106,6 @@ from backend.util.prompt import (
|
||||
estimate_token_count,
|
||||
get_compression_target,
|
||||
)
|
||||
from backend.util.settings import Settings
|
||||
|
||||
from ..config import ChatConfig, CopilotLLMModel
|
||||
from ..constants import (
|
||||
@@ -170,7 +166,6 @@ from ..prompting import (
|
||||
get_graphiti_supplement,
|
||||
get_sdk_supplement,
|
||||
)
|
||||
from ..rate_limit import get_user_tier
|
||||
from ..response_model import (
|
||||
StreamBaseResponse,
|
||||
StreamCheckpoint,
|
||||
@@ -204,7 +199,6 @@ from ..expert_context import build_expert_identity_suffix
|
||||
from ..expert_kickoff import is_expert_kickoff_turn
|
||||
from ..service import (
|
||||
_build_system_prompt,
|
||||
_is_langfuse_configured,
|
||||
_update_title_async,
|
||||
first_message_expert_context,
|
||||
inject_user_context,
|
||||
@@ -619,6 +613,32 @@ async def _consume_sdk_until_done(
|
||||
loop_state.ended_with_stream_error = True
|
||||
break
|
||||
|
||||
# Our own gateway refusing a request is a defect on our side, so
|
||||
# its sentence is not an answer: forwarded, the turn reads as a
|
||||
# successful reply of "API Error: 409 …" and nothing retries it.
|
||||
# Judged after the transient check, which owns a provider failure
|
||||
# the gateway relayed.
|
||||
refusal = (
|
||||
ctx.codex_gateway.last_refusal
|
||||
if ctx.codex_gateway is not None
|
||||
else None
|
||||
)
|
||||
# The preview is a truncated ``repr`` of the block list, so a
|
||||
# refusal behind a long block, or carrying an escaped quote, is
|
||||
# not in it: read the blocks' own text.
|
||||
refusal_text = "".join(
|
||||
block.text for block in sdk_msg.content if isinstance(block, TextBlock)
|
||||
)
|
||||
if refusal is not None and refusal in refusal_text:
|
||||
logger.error(
|
||||
f"{ctx.log_prefix} Codex gateway refused the request — "
|
||||
f"failing the turn: {refusal}"
|
||||
)
|
||||
loop_state.stream_error_msg = GATEWAY_REFUSED_MESSAGE
|
||||
loop_state.stream_error_code = GATEWAY_REFUSED_CODE
|
||||
loop_state.ended_with_stream_error = True
|
||||
break
|
||||
|
||||
# Determine if the message is a tool-only batch (all content
|
||||
# items are ToolUseBlocks) — such messages have no text output yet,
|
||||
# so we skip the wait_for_stash flush below.
|
||||
@@ -1272,7 +1292,7 @@ _EXPECTED_STOP_CODES: frozenset[str] = frozenset(
|
||||
# Handled-error codes whose StreamError the outer retry loop yields itself,
|
||||
# after it has persisted the history marker, rather than the attempt.
|
||||
_OUTER_LOOP_YIELDS_ERROR_CODES: frozenset[str] = frozenset(
|
||||
{"transient_api_error", PROVIDER_UNAVAILABLE_CODE}
|
||||
{"transient_api_error", PROVIDER_UNAVAILABLE_CODE, GATEWAY_REFUSED_CODE}
|
||||
)
|
||||
|
||||
|
||||
@@ -2135,75 +2155,6 @@ def _classify_final_failure(
|
||||
return None
|
||||
|
||||
|
||||
def _setup_langfuse_otel() -> None:
|
||||
"""Configure OTEL tracing for the Claude Agent SDK → Langfuse.
|
||||
|
||||
This uses LangSmith's built-in Claude Agent SDK integration to monkey-patch
|
||||
`ClaudeSDKClient`, capturing every tool call and model turn as OTEL spans.
|
||||
Spans are exported via OTLP to Langfuse (or any OTEL-compatible backend).
|
||||
|
||||
To route traces elsewhere, override `OTEL_EXPORTER_OTLP_ENDPOINT` and
|
||||
`OTEL_EXPORTER_OTLP_HEADERS` environment variables — no code changes needed.
|
||||
|
||||
Every run is passed through the Google data redaction before it is
|
||||
exported; if that cannot be installed, the SDK is left untraced.
|
||||
"""
|
||||
if not _is_langfuse_configured():
|
||||
return
|
||||
|
||||
# Ours to take back if the setup fails, so LangSmith does not trace
|
||||
# without the redaction.
|
||||
set_here: list[str] = []
|
||||
try:
|
||||
settings = Settings()
|
||||
pk = settings.secrets.langfuse_public_key
|
||||
sk = settings.secrets.langfuse_secret_key
|
||||
host = settings.secrets.langfuse_host
|
||||
tracing_env = settings.secrets.langfuse_tracing_environment
|
||||
|
||||
# OTEL exporter config — these are only set if not already present,
|
||||
# so explicit env-var overrides always win. The LangSmith client
|
||||
# reads its OTEL mode as the redaction builds it, so these go first.
|
||||
creds = base64.b64encode(f"{pk}:{sk}".encode()).decode()
|
||||
set_here = _set_env_defaults(
|
||||
{
|
||||
"LANGSMITH_OTEL_ENABLED": "true",
|
||||
"LANGSMITH_OTEL_ONLY": "true",
|
||||
"OTEL_EXPORTER_OTLP_ENDPOINT": f"{host}/api/public/otel",
|
||||
"OTEL_EXPORTER_OTLP_HEADERS": f"Authorization=Basic {creds}",
|
||||
"OTEL_PYTHON_EXPORTER_OTLP_HTTP_TRACES_CREDENTIAL_PROVIDER": (
|
||||
"autogpt-copilot-traces"
|
||||
),
|
||||
# The Langfuse server maps this to its first-class
|
||||
# environment field.
|
||||
"OTEL_RESOURCE_ATTRIBUTES": f"langfuse.environment={tracing_env}",
|
||||
}
|
||||
)
|
||||
|
||||
install_trace_redaction()
|
||||
# Switched on only once every run it traces is redacted.
|
||||
set_here += _set_env_defaults({"LANGSMITH_TRACING": "true"})
|
||||
configure_claude_agent_sdk(tags=["sdk"])
|
||||
logger.info(
|
||||
"OTEL tracing configured for Claude Agent SDK → %s [%s]", host, tracing_env
|
||||
)
|
||||
except Exception:
|
||||
for key in set_here:
|
||||
os.environ.pop(key, None)
|
||||
logger.warning("OTEL setup skipped — failed to configure", exc_info=True)
|
||||
|
||||
|
||||
def _set_env_defaults(defaults: dict[str, str]) -> list[str]:
|
||||
"""Set each variable not already set; return the ones this set."""
|
||||
added = [key for key in defaults if key not in os.environ]
|
||||
for key in added:
|
||||
os.environ[key] = defaults[key]
|
||||
return added
|
||||
|
||||
|
||||
_setup_langfuse_otel()
|
||||
|
||||
|
||||
# Set to hold background tasks to prevent garbage collection
|
||||
_background_tasks: set[asyncio.Task[Any]] = set()
|
||||
|
||||
@@ -2297,13 +2248,10 @@ async def _iter_sdk_messages(
|
||||
timeout. On timeout we yield a heartbeat sentinel but keep the Task
|
||||
alive so it can deliver the next message.
|
||||
|
||||
Every fetch Task runs in one context, copied once up front. The
|
||||
langsmith tracing wrapper around ``receive_response()`` stores the
|
||||
``claude.conversation`` run in a ContextVar during the first fetch and
|
||||
parents each reply's ``claude.assistant.turn`` span on it; a Task given
|
||||
a fresh copy of the caller's context per message never sees that run,
|
||||
so every reply after the first message would drop out of the trace.
|
||||
Only one fetch is in flight at a time, so sharing the context is safe.
|
||||
Every fetch Task runs in one context, copied once up front, so a
|
||||
ContextVar the SDK sets while fetching one message is still set for
|
||||
the next. Only one fetch is in flight at a time, so sharing the
|
||||
context is safe.
|
||||
|
||||
Yields `None` on heartbeat timeout (caller should refresh locks and
|
||||
emit heartbeat events). Yields the raw SDK message otherwise.
|
||||
@@ -2383,7 +2331,7 @@ def _resolve_sdk_model() -> str | None:
|
||||
user's subscription plan. Otherwise derives from
|
||||
``config.thinking_standard_model``.
|
||||
|
||||
For per-user routing (LaunchDarkly overrides), see
|
||||
For per-user routing (PostHog overrides), see
|
||||
:func:`_resolve_sdk_model_for_request`.
|
||||
"""
|
||||
if config.claude_agent_model:
|
||||
@@ -2438,7 +2386,7 @@ async def _resolve_sdk_model_for_request(
|
||||
|
||||
Priority (highest first):
|
||||
1. ``config.claude_agent_model`` — unconditional override, bypasses LD.
|
||||
2. LaunchDarkly ``copilot-model-routing[thinking][{tier}]`` if it
|
||||
2. PostHog ``copilot-model-routing[thinking][{tier}]`` if it
|
||||
serves a value different from the config default for *user_id*.
|
||||
An LD-served override wins over subscription mode so admins can
|
||||
route specific users to a specific model without flipping
|
||||
@@ -5049,7 +4997,7 @@ async def stream_chat_completion_sdk( # pyright: ignore[reportGeneralTypeIssues
|
||||
The SDK path does not currently branch on this value.
|
||||
model: Per-request model preference from the frontend toggle.
|
||||
'advanced' → Claude Opus; 'standard' → global config default.
|
||||
Takes priority over per-user LaunchDarkly targeting.
|
||||
Takes priority over per-user PostHog targeting.
|
||||
"""
|
||||
if session is None:
|
||||
session = await get_chat_session(session_id, user_id)
|
||||
@@ -5240,16 +5188,8 @@ async def stream_chat_completion_sdk( # pyright: ignore[reportGeneralTypeIssues
|
||||
)
|
||||
return
|
||||
|
||||
# OTEL context manager — initialized inside the try and cleaned up in finally.
|
||||
_otel_ctx: Any = None
|
||||
# The CLI spawned ahead of the query — stopped in finally if no launch used it.
|
||||
_cli_prelaunch: CLIPrelaunch | None = None
|
||||
# Parent Langfuse span for the turn — created so that the
|
||||
# ``openrouter-cost-reconcile`` backfill event has a stable trace_id to
|
||||
# attach to even though it fires after the SDK-emitted spans end.
|
||||
# ``propagate_attributes`` alone doesn't create a span, so without this
|
||||
# wrapper ``get_current_trace_id()`` returns None at the finally block.
|
||||
_lf_span: Any = None
|
||||
skip_transcript_upload = False
|
||||
has_history = len(session.messages) > 1
|
||||
transcript_content: str = ""
|
||||
@@ -5311,7 +5251,8 @@ async def stream_chat_completion_sdk( # pyright: ignore[reportGeneralTypeIssues
|
||||
|
||||
# Make sure there is no more code between the lock acquisition and try-block.
|
||||
try:
|
||||
# Build system prompt (reuses non-SDK path with Langfuse support).
|
||||
# Build system prompt (reuses non-SDK path with Langfuse prompt
|
||||
# management).
|
||||
# Pre-compute the cwd here so the exact working directory path can be
|
||||
# injected into the supplement instead of the generic placeholder.
|
||||
# Catch ValueError early so the failure yields a clean StreamError rather
|
||||
@@ -5338,7 +5279,7 @@ async def stream_chat_completion_sdk( # pyright: ignore[reportGeneralTypeIssues
|
||||
)
|
||||
|
||||
# --- Run independent async I/O operations in parallel ---
|
||||
# E2B sandbox setup, system prompt build (Langfuse + DB), Graphiti
|
||||
# E2B sandbox setup, system prompt build (Langfuse prompt + DB), Graphiti
|
||||
# warm-context, and CLI session restore are all independent network
|
||||
# calls. Running them concurrently saves ~500-1000ms vs sequential.
|
||||
|
||||
@@ -5794,41 +5735,9 @@ async def stream_chat_completion_sdk( # pyright: ignore[reportGeneralTypeIssues
|
||||
)
|
||||
adapter.governor = governor
|
||||
|
||||
# Propagate user_id/session_id as OTEL context attributes so the
|
||||
# langsmith tracing integration attaches them to every span. This
|
||||
# is what Langfuse (or any OTEL backend) maps to its native
|
||||
# user/session fields.
|
||||
_user_tier = await get_user_tier(user_id) if user_id else None
|
||||
_otel_metadata: dict[str, str] = {
|
||||
"resume": str(use_resume),
|
||||
"conversation_turn": str(turn),
|
||||
"tool_surface": "registry",
|
||||
}
|
||||
if _user_tier:
|
||||
_otel_metadata["subscription_tier"] = _user_tier.value
|
||||
|
||||
# Open a Langfuse parent span so the trace_id is observable from
|
||||
# the finally block — ``propagate_attributes`` only annotates an
|
||||
# existing span, it does not create one.
|
||||
try:
|
||||
_lf_span = get_client().start_as_current_span(name="copilot-sdk-turn")
|
||||
_lf_span.__enter__()
|
||||
except Exception:
|
||||
logger.debug("Failed to open Langfuse parent span", exc_info=True)
|
||||
_lf_span = None
|
||||
|
||||
_otel_ctx = propagate_attributes(
|
||||
user_id=user_id,
|
||||
session_id=session_id,
|
||||
trace_name="copilot-sdk",
|
||||
tags=["sdk"],
|
||||
metadata=_otel_metadata,
|
||||
)
|
||||
_otel_ctx.__enter__()
|
||||
|
||||
# The CLI's options are final and the query isn't built yet: spawn it
|
||||
# now, under the turn's trace context, so it boots while the query is
|
||||
# assembled instead of after (``CLIPrelaunch``).
|
||||
# now, so it boots while the query is assembled instead of after
|
||||
# (``CLIPrelaunch``).
|
||||
_cli_prelaunch = await CLIPrelaunch.start(sdk_options, log_prefix)
|
||||
|
||||
current_message = message or ""
|
||||
@@ -6806,48 +6715,6 @@ async def stream_chat_completion_sdk( # pyright: ignore[reportGeneralTypeIssues
|
||||
if _cli_prelaunch is not None:
|
||||
await _cli_prelaunch.close()
|
||||
|
||||
# --- Close OTEL context (with cost attributes) ---
|
||||
# Captured before __exit__ so the reconcile task (launched below,
|
||||
# after the span closes) can attach a backfill event to this turn's
|
||||
# Langfuse trace. Without it, Langfuse shows the rate-card estimate
|
||||
# only — for non-Anthropic OpenRouter routes that's wildly wrong.
|
||||
langfuse_trace_id: str | None = None
|
||||
if _otel_ctx is not None:
|
||||
try:
|
||||
span = otel_trace.get_current_span()
|
||||
if span and span.is_recording():
|
||||
span.set_attribute("gen_ai.usage.prompt_tokens", turn_prompt_tokens)
|
||||
span.set_attribute(
|
||||
"gen_ai.usage.completion_tokens", turn_completion_tokens
|
||||
)
|
||||
span.set_attribute(
|
||||
"gen_ai.usage.cache_read_tokens", turn_cache_read_tokens
|
||||
)
|
||||
span.set_attribute(
|
||||
"gen_ai.usage.cache_creation_tokens",
|
||||
turn_cache_creation_tokens,
|
||||
)
|
||||
if turn_cost_usd is not None:
|
||||
span.set_attribute("gen_ai.usage.cost_usd", turn_cost_usd)
|
||||
except Exception:
|
||||
logger.debug("Failed to set OTEL cost attributes", exc_info=True)
|
||||
try:
|
||||
_otel_ctx.__exit__(*sys.exc_info())
|
||||
except Exception:
|
||||
logger.warning("OTEL context teardown failed", exc_info=True)
|
||||
if _lf_span is not None:
|
||||
# Capture from our Langfuse parent span before tearing it down;
|
||||
# tracks the lifetime of ``_lf_span`` so the trace id is still
|
||||
# available if ``_otel_ctx`` was never entered.
|
||||
try:
|
||||
langfuse_trace_id = get_client().get_current_trace_id()
|
||||
except Exception:
|
||||
logger.debug("Failed to capture Langfuse trace_id", exc_info=True)
|
||||
try:
|
||||
_lf_span.__exit__(*sys.exc_info())
|
||||
except Exception:
|
||||
logger.warning("Langfuse parent span teardown failed", exc_info=True)
|
||||
|
||||
# --- Persist token usage to session + rate-limit counters ---
|
||||
# Both must live in finally so they stay consistent even when an
|
||||
# exception interrupts the try block after StreamUsage was yielded.
|
||||
@@ -6925,7 +6792,6 @@ async def stream_chat_completion_sdk( # pyright: ignore[reportGeneralTypeIssues
|
||||
fallback_cost_usd=turn_cost_usd,
|
||||
api_key=config.api_key,
|
||||
log_prefix=log_prefix,
|
||||
langfuse_trace_id=langfuse_trace_id,
|
||||
absorbed_generation_ids=absorbed_gen_ids,
|
||||
)
|
||||
)
|
||||
@@ -6975,11 +6841,6 @@ async def stream_chat_completion_sdk( # pyright: ignore[reportGeneralTypeIssues
|
||||
actual_model=state.observed_model if state is not None else None,
|
||||
routing_source=routing_source,
|
||||
)
|
||||
_stamp_turn_trace_id(
|
||||
session.messages,
|
||||
start_index=pre_turn_message_count,
|
||||
trace_id=langfuse_trace_id,
|
||||
)
|
||||
stamp_turn_trace_id(session.messages, pre_turn_message_count)
|
||||
# What this turn ran on, recorded on the turn rather than read
|
||||
# back off the session later, so a route change cannot rewrite it.
|
||||
@@ -7460,30 +7321,6 @@ def _stamp_turn_messages(
|
||||
msg.stamps_pending_save = True
|
||||
|
||||
|
||||
def _stamp_turn_trace_id(
|
||||
messages: list[ChatMessage],
|
||||
*,
|
||||
start_index: int,
|
||||
trace_id: str | None,
|
||||
) -> None:
|
||||
"""Record the turn's Langfuse trace on the reply rows it wrote.
|
||||
|
||||
A thumbs up/down on the reply is scored against this trace (see
|
||||
``backend.copilot.feedback``). Every rateable role is stamped: the UI
|
||||
names a reply bubble after its last assistant *or* reasoning row, and a
|
||||
rating of either must find the trace. Bounded to the turn and never
|
||||
overwriting, exactly like ``_stamp_turn_messages``; rows flushed mid-turn
|
||||
ride the same stamps back-fill.
|
||||
"""
|
||||
if not trace_id:
|
||||
return
|
||||
for msg in messages[start_index:]:
|
||||
if msg.role in RATEABLE_ROLES and msg.langfuse_trace_id is None:
|
||||
msg.langfuse_trace_id = trace_id
|
||||
if msg.sequence is not None:
|
||||
msg.stamps_pending_save = True
|
||||
|
||||
|
||||
def _make_cli_tmpdir(sdk_cwd: str) -> str:
|
||||
"""Create the CLI's temp directory with mode 0700, or reuse one an earlier
|
||||
turn of this session left, which must be a 0700 directory of this user."""
|
||||
|
||||
@@ -1690,84 +1690,6 @@ class TestStampTurnMessages:
|
||||
), "pre_turn_message_count must be captured AFTER error-marker cleanup"
|
||||
|
||||
|
||||
class TestStampTurnTraceId:
|
||||
"""A thumbs up/down on a reply is scored against the trace recorded on
|
||||
it, so the stamp must name this turn's trace on this turn's replies."""
|
||||
|
||||
TRACE_ID = "1edf31f11b1693cc6103f358c1481694"
|
||||
|
||||
def _messages(self):
|
||||
from backend.copilot.model import ChatMessage
|
||||
|
||||
return [
|
||||
ChatMessage(role="assistant", content="old", sequence=0),
|
||||
ChatMessage(role="user", content="q", sequence=1),
|
||||
ChatMessage(role="assistant", content="working", sequence=2),
|
||||
ChatMessage(role="tool", content="{}", tool_call_id="t1"),
|
||||
ChatMessage(role="assistant", content="done"),
|
||||
# The UI can name the reply bubble after a trailing reasoning row.
|
||||
ChatMessage(role="reasoning", content="checking the totals"),
|
||||
]
|
||||
|
||||
def test_stamps_this_turns_replies_only(self):
|
||||
from backend.copilot.sdk.service import _stamp_turn_trace_id
|
||||
|
||||
msgs = self._messages()
|
||||
_stamp_turn_trace_id(msgs, start_index=1, trace_id=self.TRACE_ID)
|
||||
|
||||
assert msgs[0].langfuse_trace_id is None # an earlier turn's reply
|
||||
assert msgs[1].langfuse_trace_id is None # user row
|
||||
assert msgs[2].langfuse_trace_id == self.TRACE_ID
|
||||
assert msgs[3].langfuse_trace_id is None # tool row
|
||||
assert msgs[4].langfuse_trace_id == self.TRACE_ID
|
||||
assert msgs[5].langfuse_trace_id == self.TRACE_ID # reasoning row
|
||||
|
||||
def test_flushed_rows_flagged_for_backfill(self):
|
||||
from backend.copilot.sdk.service import _stamp_turn_trace_id
|
||||
|
||||
msgs = self._messages()
|
||||
_stamp_turn_trace_id(msgs, start_index=1, trace_id=self.TRACE_ID)
|
||||
|
||||
assert msgs[2].stamps_pending_save is True # flushed mid-turn
|
||||
assert msgs[4].stamps_pending_save is False # inserted by this save
|
||||
|
||||
def test_never_overwrites_an_existing_trace(self):
|
||||
from backend.copilot.sdk.service import _stamp_turn_trace_id
|
||||
|
||||
msgs = self._messages()
|
||||
msgs[2].langfuse_trace_id = "a" * 32
|
||||
_stamp_turn_trace_id(msgs, start_index=0, trace_id=self.TRACE_ID)
|
||||
|
||||
assert msgs[2].langfuse_trace_id == "a" * 32
|
||||
assert msgs[2].stamps_pending_save is False
|
||||
|
||||
def test_no_trace_stamps_nothing(self):
|
||||
"""Langfuse off, or its span failed to open: nothing to record."""
|
||||
from backend.copilot.sdk.service import _stamp_turn_trace_id
|
||||
|
||||
msgs = self._messages()
|
||||
_stamp_turn_trace_id(msgs, start_index=0, trace_id=None)
|
||||
|
||||
assert all(m.langfuse_trace_id is None for m in msgs)
|
||||
assert not any(m.stamps_pending_save for m in msgs)
|
||||
|
||||
def test_stamped_in_the_finally_after_the_trace_is_captured(self):
|
||||
"""The turn handler is too large to drive in a unit test, so pin at
|
||||
the source that the stamp runs after the trace id is captured and
|
||||
before the end-of-turn save that persists it."""
|
||||
import inspect
|
||||
|
||||
from backend.copilot.sdk.service import stream_chat_completion_sdk
|
||||
|
||||
src = inspect.getsource(stream_chat_completion_sdk)
|
||||
captured_at = src.index(
|
||||
"langfuse_trace_id = get_client().get_current_trace_id()"
|
||||
)
|
||||
stamped_at = src.index("_stamp_turn_trace_id(")
|
||||
saved_at = src.index("await asyncio.shield(upsert_chat_session(session))")
|
||||
assert captured_at < stamped_at < saved_at
|
||||
|
||||
|
||||
class TestSdkServingSegment:
|
||||
def test_platform_turn_is_not_read_from_mutable_session_metadata(self):
|
||||
from backend.copilot.sdk.service import _sdk_serving_segment
|
||||
|
||||
@@ -1610,7 +1610,7 @@ class TestConsumeSdkUntilDone:
|
||||
rig — just verify that the moved-from-while-True body still dispatches
|
||||
the SDK message stream correctly across the common branches."""
|
||||
|
||||
def _ctx(self, session_id="s1"):
|
||||
def _ctx(self, session_id="s1", codex_gateway=None):
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
@@ -1642,6 +1642,7 @@ class TestConsumeSdkUntilDone:
|
||||
attachments=attachments,
|
||||
compaction=CompactionTracker(),
|
||||
lock=lock,
|
||||
codex_gateway=codex_gateway,
|
||||
)
|
||||
|
||||
def _state(self, session_id="s1"):
|
||||
@@ -1728,6 +1729,96 @@ class TestConsumeSdkUntilDone:
|
||||
assert acc.stream_completed is True
|
||||
assert loop_state.ended_with_stream_error is False
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"preceding_text",
|
||||
["", "a line the model wrote before the error happened. " * 20],
|
||||
ids=["error-alone", "error-behind-a-long-block"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_refusal_from_our_own_gateway_fails_the_turn(self, preceding_text):
|
||||
"""The CLI reports a 4xx of ours as an assistant message with an error.
|
||||
|
||||
Forwarded, it becomes the turn's reply: a user read "API Error: 409
|
||||
This tool-result request refers to a completed model call" off a turn
|
||||
that ended as a success with no output (SECRT-3333). It must fail the
|
||||
turn instead, so the chat bot can say why.
|
||||
|
||||
The second case puts the error behind a block long enough to push it
|
||||
out of the 500-character ``repr`` preview the surrounding checks use.
|
||||
"""
|
||||
from aiohttp import ClientSession
|
||||
from claude_agent_sdk import AssistantMessage, ResultMessage, TextBlock
|
||||
|
||||
from backend.copilot.response_model import StreamTextDelta
|
||||
from backend.copilot.sdk.codex_compat_gateway import (
|
||||
GATEWAY_REFUSED_CODE,
|
||||
GATEWAY_REFUSED_MESSAGE,
|
||||
CodexAnthropicGateway,
|
||||
)
|
||||
from backend.copilot.sdk.service import _consume_sdk_until_done
|
||||
|
||||
class _NeverInvoked:
|
||||
async def invoke(self, *_args, **_kwargs):
|
||||
raise AssertionError("a refused request must not reach the model")
|
||||
|
||||
async with CodexAnthropicGateway(
|
||||
agent_session=_NeverInvoked(),
|
||||
model="gpt-5.6-terra",
|
||||
) as gateway:
|
||||
async with ClientSession() as client:
|
||||
refused = await client.post(
|
||||
f"{gateway.base_url}/v1/messages",
|
||||
headers={"Authorization": f"Bearer {gateway.auth_token}"},
|
||||
json=["not an object"],
|
||||
)
|
||||
body = await refused.json()
|
||||
# The gateway's own sentence, in the wording the CLI hands us.
|
||||
reported = f"API Error: {refused.status} {body['error']['message']}"
|
||||
|
||||
ctx = self._ctx(codex_gateway=gateway)
|
||||
state = self._state()
|
||||
acc = self._acc()
|
||||
loop_state = self._loop_state()
|
||||
|
||||
async def fake_iter(client, wake=None, tool_display_wake=None):
|
||||
blocks = [TextBlock(text=reported)]
|
||||
if preceding_text:
|
||||
blocks.insert(0, TextBlock(text=preceding_text))
|
||||
yield AssistantMessage(
|
||||
content=blocks,
|
||||
model="test",
|
||||
error="unknown",
|
||||
)
|
||||
yield ResultMessage(
|
||||
subtype="success",
|
||||
duration_ms=1,
|
||||
duration_api_ms=1,
|
||||
is_error=False,
|
||||
num_turns=1,
|
||||
session_id="s1",
|
||||
result=reported,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"backend.copilot.sdk.service._iter_sdk_messages",
|
||||
new=fake_iter,
|
||||
):
|
||||
events = []
|
||||
async for ev in _consume_sdk_until_done(
|
||||
MagicMock(), ctx, state, acc, loop_state
|
||||
):
|
||||
events.append(ev)
|
||||
|
||||
assert gateway.last_refusal == body["error"]["message"]
|
||||
assert [e for e in events if isinstance(e, StreamTextDelta)] == []
|
||||
assert all(reported not in str(event) for event in events)
|
||||
assert all(
|
||||
reported not in (message.content or "") for message in ctx.session.messages
|
||||
)
|
||||
assert loop_state.ended_with_stream_error is True
|
||||
assert loop_state.stream_error_code == GATEWAY_REFUSED_CODE
|
||||
assert loop_state.stream_error_msg == GATEWAY_REFUSED_MESSAGE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_heartbeat_sentinel_yields_heartbeat_event(self):
|
||||
"""``None`` from ``_iter_sdk_messages`` is the heartbeat sentinel —
|
||||
|
||||
@@ -1,47 +0,0 @@
|
||||
from random import uniform
|
||||
from time import monotonic, sleep
|
||||
from typing import Any
|
||||
|
||||
from requests import Response, Session
|
||||
from requests.exceptions import ConnectionError, RequestException, SSLError
|
||||
from requests.exceptions import Timeout as RequestTimeout
|
||||
from urllib3.util import Timeout
|
||||
|
||||
_MAX_ATTEMPTS = 3
|
||||
|
||||
|
||||
def create_trace_export_session() -> Session:
|
||||
"""Retry OTLP requests within a window of at most three request timeouts."""
|
||||
return _TraceExportSession()
|
||||
|
||||
|
||||
class _TraceExportSession(Session):
|
||||
def post(
|
||||
self,
|
||||
url: str | bytes,
|
||||
data: Any = None,
|
||||
json: Any = None,
|
||||
**kwargs: Any,
|
||||
) -> Response:
|
||||
request_timeout = float(kwargs["timeout"])
|
||||
deadline = monotonic() + request_timeout * _MAX_ATTEMPTS
|
||||
last_error: RequestException | None = None
|
||||
for attempt in range(_MAX_ATTEMPTS):
|
||||
remaining = deadline - monotonic()
|
||||
if remaining <= 0:
|
||||
break
|
||||
kwargs["timeout"] = Timeout(total=min(request_timeout, remaining))
|
||||
try:
|
||||
return super().post(url, data=data, json=json, **kwargs)
|
||||
except SSLError as error:
|
||||
raise RequestException("OTLP trace export TLS failure") from error
|
||||
except (ConnectionError, RequestTimeout) as error:
|
||||
last_error = error
|
||||
remaining = deadline - monotonic()
|
||||
if attempt + 1 < _MAX_ATTEMPTS and remaining > 0:
|
||||
sleep(uniform(0, min(0.1 * 2**attempt, remaining / 2)))
|
||||
|
||||
# ConnectionError would make the OTLP exporter retry the exhausted batch again.
|
||||
raise RequestException(
|
||||
"OTLP trace export transport retry budget exhausted"
|
||||
) from last_error
|
||||
@@ -1,243 +0,0 @@
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from typing import cast
|
||||
from unittest.mock import MagicMock, patch
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from langfuse import Langfuse
|
||||
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
|
||||
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
|
||||
ExportTraceServiceRequest,
|
||||
)
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
from opentelemetry.sdk.trace.export import SpanExportResult
|
||||
from requests.exceptions import ConnectionError, ReadTimeout, RequestException, SSLError
|
||||
|
||||
from backend.copilot.sdk.service import _setup_langfuse_otel
|
||||
from backend.copilot.sdk.trace_export import create_trace_export_session
|
||||
|
||||
|
||||
class _TraceServer(ThreadingHTTPServer):
|
||||
daemon_threads = True
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(("127.0.0.1", 0), _TraceHandler)
|
||||
self.bodies: list[bytes] = []
|
||||
self.accepted: list[bytes] = []
|
||||
self.fail_requests = 1
|
||||
self.fail_body: bytes | None = None
|
||||
self.response_delay = 1.2
|
||||
self.success_delay = 0.0
|
||||
self.response_status = 200
|
||||
|
||||
|
||||
class _TraceHandler(BaseHTTPRequestHandler):
|
||||
def do_POST(self) -> None:
|
||||
receiver = cast(_TraceServer, self.server)
|
||||
body = self.rfile.read(int(self.headers["Content-Length"]))
|
||||
receiver.bodies.append(body)
|
||||
status = receiver.response_status
|
||||
if len(receiver.bodies) <= receiver.fail_requests or body == receiver.fail_body:
|
||||
time.sleep(receiver.response_delay)
|
||||
status = 504
|
||||
elif status == 200:
|
||||
time.sleep(receiver.success_delay)
|
||||
receiver.accepted.append(body)
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Length", "0")
|
||||
self.send_header("Retry-After", "1")
|
||||
self.end_headers()
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def trace_receiver() -> Iterator[_TraceServer]:
|
||||
receiver = _TraceServer()
|
||||
worker = threading.Thread(target=receiver.serve_forever, daemon=True)
|
||||
worker.start()
|
||||
try:
|
||||
yield receiver
|
||||
finally:
|
||||
receiver.shutdown()
|
||||
receiver.server_close()
|
||||
worker.join()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def trace_export_environment() -> Iterator[None]:
|
||||
with patch.dict(os.environ, {"LANGFUSE_TRACING_ENABLED": "true"}):
|
||||
for key in (
|
||||
"OTEL_PYTHON_EXPORTER_OTLP_HTTP_CREDENTIAL_PROVIDER",
|
||||
"OTEL_PYTHON_EXPORTER_OTLP_HTTP_TRACES_CREDENTIAL_PROVIDER",
|
||||
):
|
||||
os.environ.pop(key, None)
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def langfuse_client(
|
||||
trace_receiver: _TraceServer, trace_export_environment: None
|
||||
) -> Iterator[Langfuse]:
|
||||
settings = MagicMock()
|
||||
settings.secrets.langfuse_public_key = f"pk-test-{uuid4()}"
|
||||
settings.secrets.langfuse_secret_key = "sk-test"
|
||||
settings.secrets.langfuse_host = f"http://127.0.0.1:{trace_receiver.server_port}"
|
||||
settings.secrets.langfuse_tracing_environment = "test"
|
||||
with (
|
||||
patch("backend.copilot.sdk.service._is_langfuse_configured", return_value=True),
|
||||
patch("backend.copilot.sdk.service.Settings", return_value=settings),
|
||||
patch("backend.copilot.sdk.service.configure_claude_agent_sdk"),
|
||||
):
|
||||
_setup_langfuse_otel()
|
||||
|
||||
provider = TracerProvider()
|
||||
client = Langfuse(
|
||||
public_key=settings.secrets.langfuse_public_key,
|
||||
secret_key=settings.secrets.langfuse_secret_key,
|
||||
base_url=settings.secrets.langfuse_host,
|
||||
timeout=1,
|
||||
tracer_provider=provider,
|
||||
tracing_enabled=True,
|
||||
)
|
||||
try:
|
||||
yield client
|
||||
finally:
|
||||
client.shutdown()
|
||||
provider.shutdown()
|
||||
|
||||
|
||||
def test_langfuse_retries_timed_out_span_batch(
|
||||
trace_receiver: _TraceServer, langfuse_client: Langfuse
|
||||
) -> None:
|
||||
with langfuse_client.start_as_current_span(name="copilot-sdk-turn") as span:
|
||||
span.update(input="test", output="done")
|
||||
langfuse_client.flush()
|
||||
|
||||
assert len(trace_receiver.bodies) == 2
|
||||
assert trace_receiver.bodies[0] == trace_receiver.bodies[1]
|
||||
assert len(trace_receiver.accepted) == 1
|
||||
request = ExportTraceServiceRequest.FromString(trace_receiver.accepted[0])
|
||||
exported = request.resource_spans[0].scope_spans[0].spans
|
||||
assert [item.name for item in exported] == ["copilot-sdk-turn"]
|
||||
assert exported[0].span_id.hex() == span.id
|
||||
assert exported[0].trace_id.hex() == span.trace_id
|
||||
|
||||
|
||||
def test_langfuse_allows_slow_success_within_request_timeout(
|
||||
trace_receiver: _TraceServer, langfuse_client: Langfuse
|
||||
) -> None:
|
||||
trace_receiver.fail_requests = 0
|
||||
trace_receiver.success_delay = 0.6
|
||||
with langfuse_client.start_as_current_span(name="copilot-sdk-turn"):
|
||||
pass
|
||||
langfuse_client.flush()
|
||||
|
||||
assert len(trace_receiver.bodies) == 1
|
||||
assert len(trace_receiver.accepted) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("error_type", [ConnectionError, ReadTimeout])
|
||||
def test_trace_export_session_bounds_transport_retries(
|
||||
error_type: type[RequestException],
|
||||
) -> None:
|
||||
with (
|
||||
create_trace_export_session() as session,
|
||||
patch(
|
||||
"backend.copilot.sdk.trace_export.Session.post",
|
||||
side_effect=error_type("connection failed"),
|
||||
) as post,
|
||||
patch("backend.copilot.sdk.trace_export.sleep"),
|
||||
):
|
||||
exporter = OTLPSpanExporter(session=session, timeout=1)
|
||||
with pytest.raises(RequestException, match="retry budget exhausted") as error:
|
||||
exporter.export([])
|
||||
exporter.shutdown()
|
||||
|
||||
assert post.call_count == 3
|
||||
assert error.type is RequestException
|
||||
|
||||
|
||||
def test_exporter_bounds_retry_window_and_recovers_next_batch(
|
||||
trace_receiver: _TraceServer,
|
||||
) -> None:
|
||||
# Fail the exhausted batch by its body: an attempt the client gave up on can
|
||||
# still reach the server after the next batch has started.
|
||||
trace_receiver.fail_requests = 0
|
||||
trace_receiver.fail_body = b""
|
||||
trace_receiver.response_delay = 0.4
|
||||
next_batch = TracerProvider().get_tracer(__name__).start_span("next-batch")
|
||||
next_batch.end()
|
||||
with create_trace_export_session() as session:
|
||||
exporter = OTLPSpanExporter(
|
||||
endpoint=f"http://127.0.0.1:{trace_receiver.server_port}/v1/traces",
|
||||
timeout=0.1,
|
||||
session=session,
|
||||
)
|
||||
started = time.monotonic()
|
||||
with pytest.raises(RequestException):
|
||||
exporter.export([])
|
||||
elapsed = time.monotonic() - started
|
||||
assert 1 <= len(trace_receiver.bodies) <= 3
|
||||
assert not trace_receiver.accepted
|
||||
|
||||
assert exporter.export([next_batch]) == SpanExportResult.SUCCESS
|
||||
exporter.shutdown()
|
||||
|
||||
assert elapsed < 0.7
|
||||
assert len(trace_receiver.accepted) == 1
|
||||
|
||||
|
||||
def test_exporter_does_not_retry_after_timeout_budget_is_spent() -> None:
|
||||
with (
|
||||
create_trace_export_session() as session,
|
||||
patch(
|
||||
"backend.copilot.sdk.trace_export.Session.post",
|
||||
side_effect=ReadTimeout("request timed out"),
|
||||
) as post,
|
||||
patch("backend.copilot.sdk.trace_export.monotonic", side_effect=[0, 0, 4, 4]),
|
||||
patch("backend.copilot.sdk.trace_export.sleep") as sleep,
|
||||
):
|
||||
with pytest.raises(RequestException, match="retry budget exhausted"):
|
||||
session.post("https://example.com/v1/traces", data=b"batch", timeout=1)
|
||||
|
||||
post.assert_called_once()
|
||||
sleep.assert_not_called()
|
||||
|
||||
|
||||
def test_exporter_does_not_retry_tls_errors() -> None:
|
||||
with (
|
||||
create_trace_export_session() as session,
|
||||
patch(
|
||||
"backend.copilot.sdk.trace_export.Session.post",
|
||||
side_effect=SSLError("certificate rejected"),
|
||||
) as post,
|
||||
):
|
||||
exporter = OTLPSpanExporter(session=session, timeout=1)
|
||||
with pytest.raises(RequestException, match="TLS failure"):
|
||||
exporter.export([])
|
||||
exporter.shutdown()
|
||||
|
||||
post.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", [401, 429, 503])
|
||||
def test_trace_export_session_leaves_status_handling_to_exporter(
|
||||
trace_receiver: _TraceServer, status: int
|
||||
) -> None:
|
||||
trace_receiver.fail_requests = 0
|
||||
trace_receiver.response_status = status
|
||||
with create_trace_export_session() as session:
|
||||
response = session.post(
|
||||
f"http://127.0.0.1:{trace_receiver.server_port}/v1/traces",
|
||||
data=b"span batch",
|
||||
timeout=1,
|
||||
)
|
||||
|
||||
assert response.status_code == status
|
||||
assert trace_receiver.bodies == [b"span batch"]
|
||||
Loaded 100 of 653 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user