diff --git a/AGENTS.md b/AGENTS.md index 0019264f5..991d40289 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -285,7 +285,7 @@ handlers must pass `auth.org_id`; never accept an organization selector from the client. A user cost drilldown returns 404 when the requested user belongs to another org. Deployment-wide diagnostics or mutations (global queue status/health and slot topology, model concurrency, shared-channel Slack alert -settings, and the global cost-excluded LLM-key list) additionally require the active org to match +settings, and the global cost-excluded LLM-key and experiment lists) additionally require the active org to match `ODDISH_OPERATOR_ORG_ID`, which fails closed when unset; the frontend discovers that capability through `GET /admin/operator-access` and hides those controls for other orgs. diff --git a/backend/api/app.py b/backend/api/app.py index 9c2ac30a4..7ec1bac71 100644 --- a/backend/api/app.py +++ b/backend/api/app.py @@ -270,6 +270,7 @@ async def add_server_timing_header(request: Request, call_next): api_keys, byok, clerk_webhooks, + cost_excluded_experiments, cost_excluded_keys, dashboard, documents, @@ -310,6 +311,7 @@ async def add_server_timing_header(request: Request, call_next): api.include_router(slack.router) api.include_router(admin.router) api.include_router(cost_excluded_keys.router) + api.include_router(cost_excluded_experiments.router) api.include_router(tags.router) api.include_router(reports.router) diff --git a/backend/api/routers/cost_excluded_experiments.py b/backend/api/routers/cost_excluded_experiments.py new file mode 100644 index 000000000..66cef9f2e --- /dev/null +++ b/backend/api/routers/cost_excluded_experiments.py @@ -0,0 +1,154 @@ +"""Admin API for excluding experiments from cost accounting.""" + +from __future__ import annotations + +from typing import Annotated + +from fastapi import APIRouter, Depends, HTTPException +from pydantic import BaseModel +from sqlalchemy import select +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession + +from auth import AuthContext, can_manage_api_keys, require_admin +from auth.permissions import require_operator_org +from oddish.db import ( + CostExcludedExperimentModel, + ExperimentModel, + get_session, + utcnow, +) + +router = APIRouter(prefix="/admin/cost-excluded-experiments", tags=["Admin"]) + + +class CostExcludedExperimentResponse(BaseModel): + id: str + experiment_id: str + experiment_name: str + label: str + created_by: str | None + created_at: str + + +class CreateCostExcludedExperimentRequest(BaseModel): + experiment: str + label: str = "" + + +def _response(row: CostExcludedExperimentModel) -> CostExcludedExperimentResponse: + return CostExcludedExperimentResponse( + id=row.id, + experiment_id=row.experiment_id, + experiment_name=row.experiment_name, + label=row.label, + created_by=row.created_by_user_id, + created_at=row.created_at.isoformat(), + ) + + +def _require_manage(auth: AuthContext) -> None: + require_operator_org(auth) + if not can_manage_api_keys(auth): + raise HTTPException( + status_code=403, + detail="Only organization admins may edit the cost-exclusion list", + ) + + +async def _resolve_experiment(session: AsyncSession, ref: str) -> ExperimentModel: + # An exact id wins (include_deleted: spend from a soft-deleted experiment + # still shows on cost surfaces, so it must stay excludable). Names resolve + # among live experiments only and must be unambiguous -- experiment names + # are not unique. + by_id = await session.execute( + select(ExperimentModel) + .where(ExperimentModel.id == ref) + .execution_options(include_deleted=True) + ) + experiment = by_id.scalar_one_or_none() + if experiment is not None: + return experiment + by_name = await session.execute( + select(ExperimentModel).where(ExperimentModel.name == ref).limit(2) + ) + matches = by_name.scalars().all() + if len(matches) > 1: + raise HTTPException( + status_code=409, + detail="experiment name is ambiguous; use the experiment id", + ) + if not matches: + raise HTTPException(status_code=404, detail="experiment not found") + return matches[0] + + +@router.get("", response_model=list[CostExcludedExperimentResponse]) +async def list_cost_excluded_experiments( + auth: Annotated[AuthContext, Depends(require_admin)], +) -> list[CostExcludedExperimentResponse]: + require_operator_org(auth) + async with get_session() as session: + result = await session.execute( + select(CostExcludedExperimentModel).order_by( + CostExcludedExperimentModel.created_at.desc() + ) + ) + return [_response(row) for row in result.scalars().all()] + + +@router.post("", response_model=CostExcludedExperimentResponse) +async def add_cost_excluded_experiment( + request: CreateCostExcludedExperimentRequest, + auth: Annotated[AuthContext, Depends(require_admin)], +) -> CostExcludedExperimentResponse: + _require_manage(auth) + + ref = request.experiment.strip() + if not ref: + raise HTTPException(status_code=400, detail="experiment must not be empty") + + async with get_session() as session: + experiment = await _resolve_experiment(session, ref) + existing = await session.execute( + select(CostExcludedExperimentModel).where( + CostExcludedExperimentModel.experiment_id == experiment.id + ) + ) + if existing.scalar_one_or_none() is not None: + raise HTTPException(status_code=409, detail="experiment is already excluded") + + row = CostExcludedExperimentModel( + experiment_id=experiment.id, + experiment_name=experiment.name, + label=request.label.strip(), + created_by_user_id=auth.user_id, + ) + session.add(row) + try: + await session.commit() + except IntegrityError: + raise HTTPException(status_code=409, detail="experiment is already excluded") + return _response(row) + + +@router.delete("/{row_id}") +async def remove_cost_excluded_experiment( + row_id: str, + auth: Annotated[AuthContext, Depends(require_admin)], +) -> dict: + _require_manage(auth) + async with get_session() as session: + result = await session.execute( + select(CostExcludedExperimentModel).where( + CostExcludedExperimentModel.id == row_id + ) + ) + row = result.scalar_one_or_none() + if row is None: + raise HTTPException( + status_code=404, detail="cost-excluded experiment not found" + ) + row.deleted_at = utcnow() + await session.commit() + return {"deleted": row_id} diff --git a/backend/tests/test_cost_excluded_experiments_router.py b/backend/tests/test_cost_excluded_experiments_router.py new file mode 100644 index 000000000..2ccdac6a2 --- /dev/null +++ b/backend/tests/test_cost_excluded_experiments_router.py @@ -0,0 +1,307 @@ +from __future__ import annotations + +import pytest +import pytest_asyncio +from httpx import ASGITransport, AsyncClient + +from api.app import create_app +from api.routers import cost_excluded_experiments as router_mod +from auth import AuthContext, AuthMethod, require_auth +from models import APIKeyScope, UserRole +from oddish.db import ExperimentModel + +pytestmark = pytest.mark.asyncio + + +@pytest.fixture(autouse=True) +def operator_org(monkeypatch): + monkeypatch.setenv("ODDISH_OPERATOR_ORG_ID", "org_1") + + +class _FakeResult: + def __init__(self, rows): + self._rows = list(rows) + + def scalar_one_or_none(self): + return self._rows[0] if self._rows else None + + def scalars(self): + return self + + def all(self): + return list(self._rows) + + +class FakeSession: + """Returns one queued result list per execute() call, in order.""" + + def __init__(self, results=()): + self.results = [list(rows) for rows in results] + self.added: list[object] = [] + self.committed = False + + async def execute(self, _stmt): + return _FakeResult(self.results.pop(0) if self.results else []) + + def add(self, obj): + self.added.append(obj) + + async def commit(self): + # Simulate the Python-side column defaults a real flush would apply. + from oddish.db import generate_id, utcnow + + for obj in self.added: + if getattr(obj, "id", None) is None: + obj.id = generate_id() + if getattr(obj, "created_at", None) is None: + obj.created_at = utcnow() + self.committed = True + + +class _FakeSessionCtx: + def __init__(self, session): + self._session = session + + async def __aenter__(self): + return self._session + + async def __aexit__(self, *exc_info): + return False + + +def _install_fake_get_session(monkeypatch, session): + monkeypatch.setattr(router_mod, "get_session", lambda: _FakeSessionCtx(session)) + + +def _experiment(exp_id="exp_1", name="glm sweep") -> ExperimentModel: + return ExperimentModel(id=exp_id, name=name) + + +def _admin_jwt() -> AuthContext: + return AuthContext( + method=AuthMethod.CLERK_JWT, + org_id="org_1", + user_id="admin_1", + user_role=UserRole.ADMIN, + ) + + +def _member_jwt() -> AuthContext: + return AuthContext( + method=AuthMethod.CLERK_JWT, + org_id="org_1", + user_id="member_1", + user_role=UserRole.MEMBER, + ) + + +def _other_admin_jwt() -> AuthContext: + return AuthContext( + method=AuthMethod.CLERK_JWT, + org_id="org_2", + user_id="admin_2", + user_role=UserRole.ADMIN, + ) + + +def _full_api_key() -> AuthContext: + return AuthContext( + method=AuthMethod.API_KEY, + org_id="org_1", + user_id="key_1", + user_role=UserRole.ADMIN, + scope=APIKeyScope.FULL, + ) + + +@pytest.fixture +def app(): + return create_app() + + +def _client(app, auth): + app.dependency_overrides[require_auth] = lambda: auth + transport = ASGITransport(app=app) + return AsyncClient(transport=transport, base_url="http://test") + + +@pytest_asyncio.fixture +async def admin_client(app): + app.dependency_overrides[require_auth] = lambda: _admin_jwt() + try: + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + yield client + finally: + app.dependency_overrides.pop(require_auth, None) + + +async def test_add_by_id_snapshots_name(admin_client, monkeypatch): + # Queries: resolve by id (hit), duplicate check (miss). + session = FakeSession(results=[[_experiment()], []]) + _install_fake_get_session(monkeypatch, session) + + resp = await admin_client.post( + "/admin/cost-excluded-experiments", + json={"experiment": "exp_1", "label": "sponsored"}, + ) + assert resp.status_code == 200, resp.text + body = resp.json() + assert body["experiment_id"] == "exp_1" + assert body["experiment_name"] == "glm sweep" + assert body["label"] == "sponsored" + + assert len(session.added) == 1 + row = session.added[0] + assert row.experiment_id == "exp_1" + assert row.experiment_name == "glm sweep" + assert session.committed + + +async def test_add_by_name_resolves(admin_client, monkeypatch): + # Queries: resolve by id (miss), by name (hit), duplicate check (miss). + session = FakeSession(results=[[], [_experiment()], []]) + _install_fake_get_session(monkeypatch, session) + + resp = await admin_client.post( + "/admin/cost-excluded-experiments", json={"experiment": "glm sweep"} + ) + assert resp.status_code == 200, resp.text + assert resp.json()["experiment_id"] == "exp_1" + + +async def test_add_ambiguous_name_is_409(admin_client, monkeypatch): + session = FakeSession( + results=[[], [_experiment("exp_1"), _experiment("exp_2")]] + ) + _install_fake_get_session(monkeypatch, session) + + resp = await admin_client.post( + "/admin/cost-excluded-experiments", json={"experiment": "glm sweep"} + ) + assert resp.status_code == 409 + assert "ambiguous" in resp.json()["detail"] + + +async def test_add_unknown_experiment_is_404(admin_client, monkeypatch): + _install_fake_get_session(monkeypatch, FakeSession(results=[[], []])) + resp = await admin_client.post( + "/admin/cost-excluded-experiments", json={"experiment": "missing"} + ) + assert resp.status_code == 404 + + +async def test_add_duplicate_is_409(admin_client, monkeypatch): + session = FakeSession(results=[[_experiment()], [object()]]) + _install_fake_get_session(monkeypatch, session) + resp = await admin_client.post( + "/admin/cost-excluded-experiments", json={"experiment": "exp_1"} + ) + assert resp.status_code == 409 + + +async def test_add_race_integrity_error_is_409(admin_client, monkeypatch): + from sqlalchemy.exc import IntegrityError + + class RacingSession(FakeSession): + async def commit(self): + raise IntegrityError("INSERT", {}, Exception("duplicate key")) + + _install_fake_get_session( + monkeypatch, RacingSession(results=[[_experiment()], []]) + ) + resp = await admin_client.post( + "/admin/cost-excluded-experiments", json={"experiment": "exp_1"} + ) + assert resp.status_code == 409 + + +async def test_add_empty_experiment_is_400(admin_client, monkeypatch): + _install_fake_get_session(monkeypatch, FakeSession()) + resp = await admin_client.post( + "/admin/cost-excluded-experiments", json={"experiment": " "} + ) + assert resp.status_code == 400 + + +async def test_list_returns_rows(admin_client, monkeypatch): + from oddish.db import CostExcludedExperimentModel, utcnow + + row = CostExcludedExperimentModel( + id="x1", + experiment_id="exp_1", + experiment_name="glm sweep", + label="sponsored", + created_at=utcnow(), + ) + _install_fake_get_session(monkeypatch, FakeSession(results=[[row]])) + resp = await admin_client.get("/admin/cost-excluded-experiments") + assert resp.status_code == 200 + body = resp.json() + assert body[0]["experiment_id"] == "exp_1" + assert body[0]["experiment_name"] == "glm sweep" + + +async def test_delete_soft_deletes(admin_client, monkeypatch): + from oddish.db import CostExcludedExperimentModel + + row = CostExcludedExperimentModel( + id="x1", experiment_id="exp_1", experiment_name="glm sweep", label="" + ) + session = FakeSession(results=[[row]]) + _install_fake_get_session(monkeypatch, session) + resp = await admin_client.delete("/admin/cost-excluded-experiments/x1") + assert resp.status_code == 200 + assert row.deleted_at is not None + assert session.committed + + +async def test_delete_not_found_is_404(admin_client, monkeypatch): + _install_fake_get_session(monkeypatch, FakeSession(results=[[]])) + resp = await admin_client.delete("/admin/cost-excluded-experiments/missing") + assert resp.status_code == 404 + + +async def test_member_jwt_cannot_add(app, monkeypatch): + _install_fake_get_session(monkeypatch, FakeSession()) + client = _client(app, _member_jwt()) + try: + resp = await client.post( + "/admin/cost-excluded-experiments", json={"experiment": "exp_1"} + ) + assert resp.status_code == 403 + finally: + await client.aclose() + app.dependency_overrides.pop(require_auth, None) + + +async def test_full_api_key_cannot_add(app, monkeypatch): + _install_fake_get_session(monkeypatch, FakeSession()) + client = _client(app, _full_api_key()) + try: + resp = await client.post( + "/admin/cost-excluded-experiments", json={"experiment": "exp_1"} + ) + assert resp.status_code == 403 + finally: + await client.aclose() + app.dependency_overrides.pop(require_auth, None) + + +@pytest.mark.parametrize("method", ["get", "post", "delete"]) +async def test_non_operator_admin_cannot_access_list(app, monkeypatch, method): + _install_fake_get_session(monkeypatch, FakeSession()) + client = _client(app, _other_admin_jwt()) + try: + if method == "get": + resp = await client.get("/admin/cost-excluded-experiments") + elif method == "post": + resp = await client.post( + "/admin/cost-excluded-experiments", json={"experiment": "exp_1"} + ) + else: + resp = await client.delete("/admin/cost-excluded-experiments/x1") + assert resp.status_code == 403 + finally: + await client.aclose() + app.dependency_overrides.pop(require_auth, None) diff --git a/frontend/src/app/(app)/admin/page.tsx b/frontend/src/app/(app)/admin/page.tsx index 6bb440a5c..b619c6984 100644 --- a/frontend/src/app/(app)/admin/page.tsx +++ b/frontend/src/app/(app)/admin/page.tsx @@ -38,6 +38,7 @@ import { UsagePanel } from "@/components/usage-panel"; import { QueueHealthOverviewCard } from "@/components/queue-health-overview-card"; import { CostBreakdownCard } from "@/components/cost-breakdown-card"; import { CostExcludedKeysCard } from "@/components/cost-excluded-keys-card"; +import { CostExcludedExperimentsCard } from "@/components/cost-excluded-experiments-card"; import { SlackAlertSettingsForm } from "@/components/slack-alert-settings-form"; import { RefreshCw, Server, Clock, AlertCircle } from "lucide-react"; @@ -764,6 +765,7 @@ function AdminPageContent() { + )} diff --git a/frontend/src/app/api/admin/cost-excluded-experiments/[id]/route.ts b/frontend/src/app/api/admin/cost-excluded-experiments/[id]/route.ts new file mode 100644 index 000000000..f8ada624a --- /dev/null +++ b/frontend/src/app/api/admin/cost-excluded-experiments/[id]/route.ts @@ -0,0 +1,50 @@ +import { NextRequest, NextResponse } from "next/server"; +import { auth } from "@clerk/nextjs/server"; +import { + getAuthHeaders, + getBackendUrl, + getClerkToken, +} from "@/lib/backend-config"; + +export async function DELETE( + request: NextRequest, + { params }: { params: Promise<{ id: string }> } +) { + try { + const { getToken } = await auth(); + const token = await getClerkToken(getToken); + + if (!token) { + return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); + } + + const { id } = await params; + const url = getBackendUrl("admin/cost-excluded-experiments", `/${id}`); + + const res = await fetch(url, { + method: "DELETE", + headers: getAuthHeaders(token), + }); + + if (!res.ok) { + const errorText = await res.text(); + console.error( + `[cost-excluded-experiments] Backend error: ${res.status} - ${errorText}` + ); + return NextResponse.json( + { + error: "Failed to remove cost-excluded experiment", + details: errorText, + }, + { status: res.status } + ); + } + + return NextResponse.json(await res.json()); + } catch (error) { + return NextResponse.json( + { error: error instanceof Error ? error.message : "Unknown error" }, + { status: 503 } + ); + } +} diff --git a/frontend/src/app/api/admin/cost-excluded-experiments/route.ts b/frontend/src/app/api/admin/cost-excluded-experiments/route.ts new file mode 100644 index 000000000..74ccab42f --- /dev/null +++ b/frontend/src/app/api/admin/cost-excluded-experiments/route.ts @@ -0,0 +1,89 @@ +import { NextRequest, NextResponse } from "next/server"; +import { auth } from "@clerk/nextjs/server"; +import { + getAuthHeaders, + getBackendUrl, + getClerkToken, +} from "@/lib/backend-config"; + +export async function GET() { + try { + const { getToken } = await auth(); + const token = await getClerkToken(getToken); + + if (!token) { + return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); + } + + const url = getBackendUrl("admin/cost-excluded-experiments"); + const res = await fetch(url, { + cache: "no-store", + headers: getAuthHeaders(token), + }); + + if (!res.ok) { + const errorText = await res.text(); + console.error( + `[cost-excluded-experiments] Backend error: ${res.status} - ${errorText}` + ); + return NextResponse.json( + { + error: "Failed to fetch cost-excluded experiments", + details: errorText, + }, + { status: res.status } + ); + } + + return NextResponse.json(await res.json()); + } catch (error) { + return NextResponse.json( + { error: error instanceof Error ? error.message : "Unknown error" }, + { status: 503 } + ); + } +} + +export async function POST(request: NextRequest) { + try { + const { getToken } = await auth(); + const token = await getClerkToken(getToken); + + if (!token) { + return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); + } + + const body = await request.json(); + const url = getBackendUrl("admin/cost-excluded-experiments"); + + const res = await fetch(url, { + method: "POST", + headers: { + "Content-Type": "application/json", + ...getAuthHeaders(token), + }, + body: JSON.stringify(body), + }); + + if (!res.ok) { + const errorText = await res.text(); + console.error( + `[cost-excluded-experiments] Backend error: ${res.status} - ${errorText}` + ); + return NextResponse.json( + { + error: "Failed to add cost-excluded experiment", + details: errorText, + }, + { status: res.status } + ); + } + + return NextResponse.json(await res.json()); + } catch (error) { + return NextResponse.json( + { error: error instanceof Error ? error.message : "Unknown error" }, + { status: 503 } + ); + } +} diff --git a/frontend/src/components/cost-excluded-experiments-card.tsx b/frontend/src/components/cost-excluded-experiments-card.tsx new file mode 100644 index 000000000..d7e66063b --- /dev/null +++ b/frontend/src/components/cost-excluded-experiments-card.tsx @@ -0,0 +1,162 @@ +"use client"; + +import { useState } from "react"; +import useSWR from "swr"; +import { Card, CardHeader, CardTitle, CardContent } from "@/components/ui/card"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { fetcher } from "@/lib/api"; + +type ExcludedExperiment = { + id: string; + experiment_id: string; + experiment_name: string; + label: string; + created_by: string | null; + created_at: string; +}; + +// The proxy wraps backend failures as { error, details } where details is the +// raw FastAPI response text; unwrap its "detail" message for display. +function backendDetail(body: unknown): string | null { + const details = (body as { details?: string } | null)?.details; + if (!details) return null; + try { + const parsed = JSON.parse(details); + if (typeof parsed?.detail === "string") return parsed.detail; + } catch { + // not JSON — fall through to the raw text + } + return details; +} + +export function CostExcludedExperimentsCard() { + const { data, error, isLoading, mutate } = useSWR( + "/api/admin/cost-excluded-experiments", + fetcher + ); + + const [experiment, setExperiment] = useState(""); + const [label, setLabel] = useState(""); + const [busy, setBusy] = useState(false); + const [formError, setFormError] = useState(null); + + async function add() { + if (!experiment.trim()) return; + setBusy(true); + setFormError(null); + try { + const res = await fetch("/api/admin/cost-excluded-experiments", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + experiment: experiment.trim(), + label: label.trim(), + }), + }); + if (!res.ok) { + const body = await res.json().catch(() => ({})); + setFormError( + backendDetail(body) || body?.error || "Failed to add experiment." + ); + return; + } + setExperiment(""); + setLabel(""); + await mutate(); + } finally { + setBusy(false); + } + } + + async function remove(id: string) { + setFormError(null); + const res = await fetch(`/api/admin/cost-excluded-experiments/${id}`, { + method: "DELETE", + }); + if (!res.ok) { + const body = await res.json().catch(() => ({})); + setFormError( + backendDetail(body) || body?.error || "Failed to remove experiment." + ); + } + await mutate(); + } + + return ( + + + Cost-excluded experiments +

+ Trials homed in these experiments are ignored by cost accounting: + quota enforcement and the admin cost dashboards. Experiment pages + still show the trials' raw compute cost. Add an experiment by + name or id. +

+
+ +
+ setExperiment(e.target.value)} + className="w-72" + /> + setLabel(e.target.value)} + className="w-48" + /> + +
+ {formError &&

{formError}

} + + {isLoading &&

Loading…

} + {error && ( +

+ Failed to load experiments. +

+ )} + {data && data.length === 0 && ( +

+ No experiments are excluded from cost. +

+ )} + {data && data.length > 0 && ( +
+ {data.map((row) => ( +
+
+ + {row.experiment_name || row.experiment_id} + + + {row.experiment_id} + + {row.label && ( + + {row.label} + + )} +
+ +
+ ))} +
+ )} +
+
+ ); +} diff --git a/oddish/alembic/versions/costexclexp01_add_cost_excluded_experiments.py b/oddish/alembic/versions/costexclexp01_add_cost_excluded_experiments.py new file mode 100644 index 000000000..22882c285 --- /dev/null +++ b/oddish/alembic/versions/costexclexp01_add_cost_excluded_experiments.py @@ -0,0 +1,50 @@ +"""add cost_excluded_experiments + +Revision ID: costexclexp01 +Revises: drop_prompt_registry_001 +Create Date: 2026-08-08 00:00:00.000000 + +Admin-managed list of experiments whose trials' spend is excluded from cost +accounting -- the experiment-level sibling of ``cost_excluded_llm_keys``. +Exclusion correlates on ``trials.experiment_id`` (already indexed on trials); +``experiment_name`` is a display snapshot from registration time. ``deleted_at`` +is the soft-delete tombstone; the partial UNIQUE keeps one live row per +experiment so a removed experiment can be re-added. + +""" + +from typing import Sequence, Union + +from alembic import op + +revision: str = "costexclexp01" +down_revision: Union[str, Sequence[str], None] = "drop_prompt_registry_001" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute( + """ + CREATE TABLE IF NOT EXISTS cost_excluded_experiments ( + id VARCHAR(64) PRIMARY KEY, + experiment_id VARCHAR(64) NOT NULL, + experiment_name VARCHAR(255) NOT NULL DEFAULT '', + label VARCHAR(255) NOT NULL DEFAULT '', + created_by_user_id VARCHAR(64), + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + deleted_at TIMESTAMPTZ + ) + """ + ) + op.execute( + """ + CREATE UNIQUE INDEX IF NOT EXISTS idx_cost_excluded_experiments_exp_live + ON cost_excluded_experiments (experiment_id) WHERE deleted_at IS NULL + """ + ) + + +def downgrade() -> None: + op.execute("DROP TABLE IF EXISTS cost_excluded_experiments") diff --git a/oddish/src/oddish/core/cost_basis.py b/oddish/src/oddish/core/cost_basis.py index 5b57e4d4f..b20608ca8 100644 --- a/oddish/src/oddish/core/cost_basis.py +++ b/oddish/src/oddish/core/cost_basis.py @@ -27,7 +27,12 @@ from sqlalchemy import and_, case, func, or_, select from oddish.config import settings -from oddish.db import CostExcludedLlmKeyModel, TrialModel, TrialOrigin +from oddish.db import ( + CostExcludedExperimentModel, + CostExcludedLlmKeyModel, + TrialModel, + TrialOrigin, +) from oddish.model_pricing import estimate_cost_usd # ``harbor_stage='cancelled'`` marks an abandoned trial. Three paths stamp it: @@ -100,20 +105,52 @@ def not_excluded_llm_key_filter(): return ~_cost_excluded_key_spend() +def _cost_excluded_experiment_spend(): + """Trials homed in an experiment that admins flagged cost-excluded. + + A correlated EXISTS on the live ``cost_excluded_experiments`` rows, + matching ``trials.experiment_id`` (the home experiment), so collection + experiments that merely gather trials are unaffected. The ``deleted_at`` + filter is explicit for the same reason as :func:`_cost_excluded_key_spend`: + cost surfaces run with ``include_deleted=True``, and removing an experiment + from the list must re-include its spend. + """ + return ( + select(CostExcludedExperimentModel.id) + .where( + CostExcludedExperimentModel.experiment_id == TrialModel.experiment_id, + CostExcludedExperimentModel.deleted_at.is_(None), + ) + .correlate(TrialModel) + .exists() + ) + + +def not_excluded_experiment_filter(): + """Drop trials homed in an experiment on the admin cost-exclusion list. + + Shared by :func:`first_party_spend_filter` (settled spend) and the quota + inflight reservation, mirroring :func:`not_excluded_llm_key_filter`. + """ + return ~_cost_excluded_experiment_spend() + + def first_party_spend_filter(): """Select actual Oddish executions, excluding non-spend materializations. Imported trials were paid for outside Oddish. Experiment-combine rows copy an existing trial's result and cost, so counting them would charge the same execution twice. Spend stamped with an LLM key on the admin cost-exclusion - list (sponsored/free keys) is deliberately not counted. Keep this - eligibility rule shared by quota accounting and cost reporting so both - surfaces count the same execution population. + list (sponsored/free keys), or homed in an admin cost-excluded experiment, + is deliberately not counted. Keep this eligibility rule shared by quota + accounting and cost reporting so both surfaces count the same execution + population. """ return and_( TrialModel.origin == TrialOrigin.ODDISH, not_combine_copy_filter(), not_excluded_llm_key_filter(), + not_excluded_experiment_filter(), ) diff --git a/oddish/src/oddish/core/quotas.py b/oddish/src/oddish/core/quotas.py index 3e117e032..dc2d6f352 100644 --- a/oddish/src/oddish/core/quotas.py +++ b/oddish/src/oddish/core/quotas.py @@ -11,6 +11,7 @@ from oddish.config import settings from oddish.core.cost_basis import ( first_party_spend_filter, + not_excluded_experiment_filter, not_excluded_llm_key_filter, settled_cost_columns, settled_cost_from_row, @@ -437,9 +438,9 @@ async def _bump_aware_limits_by_org_user_all_orgs( def _inflight_predicates(org_id: str | None, billed_user_id: str) -> list: - # ``not_excluded_llm_key_filter``: a RETRYING attempt already carries its - # settlement stamp while finished_at is still NULL; spend the settled sums - # will drop must not keep reserving against the cap either. + # Exclusion filters: a RETRYING attempt already carries its settlement + # stamp while finished_at is still NULL; spend the settled sums will drop + # must not keep reserving against the cap either. return [ TrialModel.org_id == org_id, TrialModel.billed_user_id == billed_user_id, @@ -448,6 +449,7 @@ def _inflight_predicates(org_id: str | None, billed_user_id: str) -> list: TrialModel.superseded_by_trial_id.is_(None), TrialModel.status.in_(_INFLIGHT_TRIAL_STATUSES), not_excluded_llm_key_filter(), + not_excluded_experiment_filter(), ] @@ -459,6 +461,7 @@ def _org_inflight_predicates(org_id: str | None) -> list: TrialModel.superseded_by_trial_id.is_(None), TrialModel.status.in_(_INFLIGHT_TRIAL_STATUSES), not_excluded_llm_key_filter(), + not_excluded_experiment_filter(), ] @@ -497,6 +500,7 @@ async def inflight_trial_count_by_org_user_all_orgs( TrialModel.superseded_by_trial_id.is_(None), TrialModel.status.in_(_INFLIGHT_TRIAL_STATUSES), not_excluded_llm_key_filter(), + not_excluded_experiment_filter(), ) .group_by(TrialModel.org_id, TrialModel.billed_user_id) ) diff --git a/oddish/src/oddish/db/__init__.py b/oddish/src/oddish/db/__init__.py index ede4300b0..453e5e269 100644 --- a/oddish/src/oddish/db/__init__.py +++ b/oddish/src/oddish/db/__init__.py @@ -27,6 +27,7 @@ AnalysisCostModel, AnalyzerModel, AnalyzerBlockModel, + CostExcludedExperimentModel, CostExcludedLlmKeyModel, DocumentModel, ExperimentModel, @@ -112,6 +113,7 @@ "AnalysisCostModel", "AnalyzerModel", "AnalyzerBlockModel", + "CostExcludedExperimentModel", "CostExcludedLlmKeyModel", "ExperimentModel", "DocumentModel", diff --git a/oddish/src/oddish/db/models.py b/oddish/src/oddish/db/models.py index f2631e040..5cc4d4de1 100644 --- a/oddish/src/oddish/db/models.py +++ b/oddish/src/oddish/db/models.py @@ -2396,6 +2396,37 @@ class CostExcludedLlmKeyModel(TimestampedMixin, Base): created_by_user_id: Mapped[str | None] = mapped_column(String(64), nullable=True) +class CostExcludedExperimentModel(TimestampedMixin, Base): + """An experiment whose trials' spend is excluded from cost accounting. + + The experiment-level sibling of ``CostExcludedLlmKeyModel``. Exclusion is + equality matching against ``trials.experiment_id`` (the home experiment), + so collection experiments that merely gather trials are unaffected. + ``experiment_name`` is a display snapshot from registration time (no FK, so + the row outlives the experiment like excluded keys outlive their trials). + ``deleted_at`` (soft delete) is the live/removed state, and the partial + UNIQUE keeps one live row per experiment so a removed one can be re-added. + """ + + __tablename__ = "cost_excluded_experiments" + __table_args__ = ( + Index( + "idx_cost_excluded_experiments_exp_live", + "experiment_id", + unique=True, + postgresql_where=text("deleted_at IS NULL"), + ), + ) + + id: Mapped[str] = mapped_column(String(64), primary_key=True, default=generate_id) + experiment_id: Mapped[str] = mapped_column(String(64), nullable=False) + experiment_name: Mapped[str] = mapped_column( + String(255), nullable=False, server_default="" + ) + label: Mapped[str] = mapped_column(String(255), nullable=False, server_default="") + created_by_user_id: Mapped[str | None] = mapped_column(String(64), nullable=True) + + from oddish.db.soft_delete import register_soft_delete_models register_soft_delete_models( @@ -2412,4 +2443,5 @@ class CostExcludedLlmKeyModel(TimestampedMixin, Base): SkillModel, DocumentModel, CostExcludedLlmKeyModel, + CostExcludedExperimentModel, ) diff --git a/oddish/tests/test_cost_excluded_experiments.py b/oddish/tests/test_cost_excluded_experiments.py new file mode 100644 index 000000000..52a5db45f --- /dev/null +++ b/oddish/tests/test_cost_excluded_experiments.py @@ -0,0 +1,215 @@ +"""DB-backed tests for excluding flagged experiments' spend from cost accounting. + +Spend from trials homed in an experiment on the ``cost_excluded_experiments`` +list must vanish from every surface that shares ``first_party_spend_filter``: +the admin cost breakdown and the quota sums. Removing an experiment from the +list re-includes its spend. +""" + +from __future__ import annotations + +import sys +import uuid +from datetime import timedelta +from pathlib import Path + +import pytest +import pytest_asyncio +from sqlalchemy import text + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) + +from oddish.core.admin import get_cost_breakdown_core # noqa: E402 +from oddish.core.quotas import sum_org_cost_usd # noqa: E402 +from oddish.db import ( # noqa: E402 + CostExcludedExperimentModel, + ExperimentModel, + TaskModel, + TrialModel, + get_session, + utcnow, +) + +_RUN = uuid.uuid4().hex[:8] + +ORG = f"expcost-org-{_RUN}" +USER = f"expcost-user-{_RUN}" +EXCLUDED_EXP = f"expcost-excluded-{_RUN}" +INCLUDED_EXP = f"expcost-included-{_RUN}" +EXCLUDED_TASK = f"{EXCLUDED_EXP}-task" +INCLUDED_TASK = f"{INCLUDED_EXP}-task" + +EXCLUDED_COST = 7.0 +INCLUDED_COST = 2.0 + + +def _trial(task_id, experiment_id, index, cost_usd, created_at) -> TrialModel: + return TrialModel( + id=f"{task_id}-{index}", + name=f"{task_id}-{index}", + task_id=task_id, + experiment_id=experiment_id, + org_id=ORG, + agent="claude-code", + provider="xai", + queue_key="xai/grok-4", + model="xai/grok-4", + billed_user_id=USER, + cost_usd=cost_usd, + created_at=created_at, + finished_at=created_at, + ) + + +@pytest_asyncio.fixture +async def seeded_data(): + recent = utcnow() - timedelta(hours=1) + excluded = CostExcludedExperimentModel( + experiment_id=EXCLUDED_EXP, experiment_name="expcost-excluded", label="free" + ) + + async with get_session() as session: + await session.execute( + text( + "INSERT INTO organizations " + "(id, name, slug, plan, settings, is_active, created_at, updated_at) " + "VALUES (:id, :id, :id, 'free', '{}'::jsonb, true, NOW(), NOW()) " + "ON CONFLICT (id) DO NOTHING" + ), + {"id": ORG}, + ) + session.add(excluded) + for exp_id, name in ( + (EXCLUDED_EXP, "expcost-excluded"), + (INCLUDED_EXP, "expcost-included"), + ): + session.add( + ExperimentModel( + id=exp_id, + name=name, + org_id=ORG, + owner_user_id=USER, + created_at=recent, + last_activity_at=recent, + ) + ) + for task_id in (EXCLUDED_TASK, INCLUDED_TASK): + session.add( + TaskModel( + id=task_id, + name=task_id, + user="test", + org_id=ORG, + task_path="some/path", + ) + ) + session.add_all( + [ + _trial(EXCLUDED_TASK, EXCLUDED_EXP, 0, EXCLUDED_COST, recent), + _trial(INCLUDED_TASK, INCLUDED_EXP, 0, INCLUDED_COST, recent), + ] + ) + await session.flush() + + yield excluded.id + + async with get_session() as session: + await session.execute( + TrialModel.__table__.delete().where( + TrialModel.experiment_id.in_([EXCLUDED_EXP, INCLUDED_EXP]) + ) + ) + await session.execute( + TaskModel.__table__.delete().where( + TaskModel.id.in_([EXCLUDED_TASK, INCLUDED_TASK]) + ) + ) + await session.execute( + ExperimentModel.__table__.delete().where( + ExperimentModel.id.in_([EXCLUDED_EXP, INCLUDED_EXP]) + ) + ) + await session.execute( + CostExcludedExperimentModel.__table__.delete().where( + CostExcludedExperimentModel.experiment_id == EXCLUDED_EXP + ) + ) + await session.execute( + text("DELETE FROM organizations WHERE id = :o"), {"o": ORG} + ) + + +@pytest.mark.asyncio +async def test_excluded_experiment_spend_dropped_then_reincluded_on_removal( + seeded_data, +): + excluded_id = seeded_data + period_start = utcnow() - timedelta(days=1) + + async with get_session() as session: + result = await get_cost_breakdown_core( + session, window_days=7, experiment_limit=500, user_limit=500 + ) + by_id = {e.experiment_id: e for e in result.experiments} + assert EXCLUDED_EXP not in by_id + included = by_id[INCLUDED_EXP] + assert abs(included.cost_usd - INCLUDED_COST) <= 1e-6, included.cost_usd + + org_total = await sum_org_cost_usd(session, ORG, period_start) + assert abs(float(org_total) - INCLUDED_COST) <= 1e-6, org_total + + # Removing the experiment from the list re-includes its spend: the + # exclusion probe only matches live rows. + await session.execute( + CostExcludedExperimentModel.__table__.update() + .where(CostExcludedExperimentModel.id == excluded_id) + .values(deleted_at=utcnow()) + ) + await session.flush() + + org_total = await sum_org_cost_usd(session, ORG, period_start) + expected = INCLUDED_COST + EXCLUDED_COST + assert abs(float(org_total) - expected) <= 1e-6, org_total + + +@pytest.mark.asyncio +async def test_inflight_reservation_skips_excluded_experiment_spend( + seeded_data, monkeypatch +): + # An in-flight trial homed in an excluded experiment must not reserve + # quota the settled sums will never charge -- and it reserves again once + # the experiment leaves the list. + from oddish.config import settings + from oddish.core.quotas import ( + inflight_reserved_usd, + inflight_trial_count_by_org_user_all_orgs, + ) + from oddish.db import TrialStatus + + monkeypatch.setattr(settings, "pending_trial_reservation_usd", 2.5) + + excluded_id = seeded_data + inflight = _trial(EXCLUDED_TASK, EXCLUDED_EXP, 9, 5.0, utcnow()) + inflight.finished_at = None + inflight.status = TrialStatus.RUNNING + + async with get_session() as session: + session.add(inflight) + await session.flush() + + reserved = await inflight_reserved_usd(session, ORG, USER) + assert float(reserved) == 0.0, reserved + inflight_counts = await inflight_trial_count_by_org_user_all_orgs(session) + assert inflight_counts.get((ORG, USER), 0) == 0 + + await session.execute( + CostExcludedExperimentModel.__table__.update() + .where(CostExcludedExperimentModel.id == excluded_id) + .values(deleted_at=utcnow()) + ) + await session.flush() + + reserved = await inflight_reserved_usd(session, ORG, USER) + assert abs(float(reserved) - 5.0) <= 1e-6, reserved + inflight_counts = await inflight_trial_count_by_org_user_all_orgs(session) + assert inflight_counts[(ORG, USER)] == 1