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:
abhi1992002andClaude Opus 5.5 committed 2026-11-03 05:27:40 +00:00
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+)$",
+45 -9
View File
@@ -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)
+1 -1
View File
@@ -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"
+4 -4
View File
@@ -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,
+9 -9
View File
@@ -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
View File
@@ -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):
+5 -2
View File
@@ -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(
@@ -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,
@@ -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:
@@ -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