From f1a194a49b274eed4e3e8b94c764322c727caee3 Mon Sep 17 00:00:00 2001 From: Sushant Gautam Date: Wed, 11 Nov 2026 20:01:39 +0100 Subject: [PATCH] perf: cut slow pages and APIs found in Sentry traces Production traces showed /experiments/{id}/ (p95 9.3 s), /runs/data/ (p95 3.1 s, the 502) and /experiments/new/ dominated by loading every ScenarioResult.result transcript to read a severity. Data layer - ScenarioResult keeps trial counts (n_pass, n_judged, n_errors, n_ungraded) and a small list_summary projection, kept in step by save(); migration 0011 backfills them in batches. pass_counts is one GROUP BY; the run Results list, CSV export and compare read the projection. - Indexes: AuditRun (project/archived/created_at, status/archived, monitor/created_at, started_at), AuditEvent (run/kind/item; drop three single-column ones), ScenarioResult (run/status), OtlpSpan (received_at). - AuditEvent: turn events are bulk-inserted per batch; the sweeper prunes per-turn events of runs finished > SIMPLEAUDIT_EVENT_RETENTION_DAYS (30); manage.py prune_audit_events. Web - One membership query per request feeds access, admin flag, workspace switcher and can_write; token-auth requests no longer write sessions. - SSE: per-user and per-process stream caps (429 + retry), idle backoff, keepalive, and an end for finished runs missing a terminal event. - API: run list and results are paginated (limit/offset); run-list rows are slim, ?expand=snapshots restores the full shape. Event poll capped at 1000. - Exports stream (run JSON/CSV, runs CSV, API export encoding). - Agent page Open WebUI sync throttled (SIMPLEAUDIT_OPENWEBUI_SYNC_TTL_S); scenarios edit forms load on demand; compare/experiment/runs grid defer unused JSON; worker counters are one aggregate; concurrent model and health probes; DB-side prompt-version check; static Cache-Control; Caddy compression and upstream timeouts. docs/performance-audit.md records the audit, the Sentry evidence and status. --- .env.dev.example | 13 ++ .env.example | 13 ++ accounts/services.py | 25 +++ audits/comparison.py | 23 +- audits/events.py | 130 ++++++++++- audits/experiments.py | 7 +- .../0011_scenario_result_trial_counts.py | 99 +++++++++ audits/migrations/0012_perf_indexes.py | 70 ++++++ audits/models.py | 8 + audits/monitors.py | 35 +-- audits/serializers.py | 67 ++++++ audits/tests/test_export.py | 2 +- audits/views.py | 170 ++++++++++++++- chat/views.py | 9 +- config/settings.py | 7 + config/urls.py | 11 + conftest.py | 6 + deploy/openwebui-subpath/Caddyfile.subpath | 12 +- docs/architecture.md | 13 +- docs/deployment.md | 7 +- docs/performance-audit.md | 111 ++++++++++ infra/context_processors.py | 32 +-- infra/export_utils.py | 57 +++++ infra/health.py | 96 ++++----- .../management/commands/prune_audit_events.py | 27 +++ infra/management/commands/run_demo_audits.py | 18 +- infra/middleware.py | 17 +- infra/runs_table.py | 42 ++-- infra/tests/test_api_e2e.py | 2 +- infra/tests/test_audit_poll_results.py | 53 ++++- infra/tests/test_audit_sse.py | 57 +++++ infra/tests/test_event_retention.py | 59 +++++ infra/tests/test_openwebui_sync_throttle.py | 31 +++ infra/tests/test_pass24_regressions.py | 2 +- infra/tests/test_request_overhead.py | 47 ++++ infra/tests/test_runs_table.py | 2 +- infra/tests/test_scenario_edit_fragment.py | 47 ++++ infra/tests/test_scenario_result_counts.py | 115 ++++++++++ infra/tests/test_smoke_all_pages.py | 4 +- infra/tests/test_streaming_exports.py | 54 +++++ infra/ui.py | 202 +++++++++++++----- infra/worker.py | 94 ++++++-- .../migrations/0023_perf_indexes.py | 27 +++ model_registry/models.py | 6 +- model_registry/services.py | 24 ++- prompts/services.py | 22 +- templates/experiment_new.html | 4 + templates/partials/agentic_rules_editor.html | 5 +- .../partials/experiment_detail_body.html | 2 +- templates/partials/scenario_edit_form.html | 86 ++++++++ templates/scenarios.html | 121 +++-------- 51 files changed, 1866 insertions(+), 327 deletions(-) create mode 100644 audits/migrations/0011_scenario_result_trial_counts.py create mode 100644 audits/migrations/0012_perf_indexes.py create mode 100644 docs/performance-audit.md create mode 100644 infra/management/commands/prune_audit_events.py create mode 100644 infra/tests/test_event_retention.py create mode 100644 infra/tests/test_openwebui_sync_throttle.py create mode 100644 infra/tests/test_request_overhead.py create mode 100644 infra/tests/test_scenario_edit_fragment.py create mode 100644 infra/tests/test_scenario_result_counts.py create mode 100644 infra/tests/test_streaming_exports.py create mode 100644 model_registry/migrations/0023_perf_indexes.py create mode 100644 templates/partials/scenario_edit_form.html diff --git a/.env.dev.example b/.env.dev.example index 70d74471..a4d3790d 100644 --- a/.env.dev.example +++ b/.env.dev.example @@ -115,3 +115,16 @@ VISUALIZER_RESULTS_DIR=./results #SIMPLEAUDIT_OTLP_MAX_BODY_BYTES=8388608 #SIMPLEAUDIT_OTLP_HOLD_SECONDS=600 #SIMPLEAUDIT_OTLP_AUTH_CACHE_SECONDS=30 + +# --- Live streams, event retention, Open WebUI sync (defaults shown) ----------- +# Open live audit-run streams per user, and in total per web process +# (default: half of GUNICORN_THREADS). Over the cap the stream answers 429 and +# the browser retries. +# SIMPLEAUDIT_SSE_MAX_STREAMS_PER_USER=3 +# SIMPLEAUDIT_SSE_MAX_STREAMS=8 +# Days after a run finishes before its per-turn progress events are deleted +# (state and terminal events are kept). 0 keeps them forever. +# SIMPLEAUDIT_EVENT_RETENTION_DAYS=30 +# Seconds between the Agent page's automatic Open WebUI knowledge/tool syncs +# per user and workspace (the Sync button always runs). 0 syncs on every load. +# SIMPLEAUDIT_OPENWEBUI_SYNC_TTL_S=60 diff --git a/.env.example b/.env.example index 7a615559..6ddadd74 100644 --- a/.env.example +++ b/.env.example @@ -171,3 +171,16 @@ WORKOS_API_KEY= # GUNICORN_TIMEOUT=60 # GUNICORN_GRACEFUL_TIMEOUT=30 # GUNICORN_KEEPALIVE=5 + +# --- Live streams, event retention, Open WebUI sync (defaults shown) ----------- +# Open live audit-run streams per user, and in total per web process +# (default: half of GUNICORN_THREADS). Over the cap the stream answers 429 and +# the browser retries. +# SIMPLEAUDIT_SSE_MAX_STREAMS_PER_USER=3 +# SIMPLEAUDIT_SSE_MAX_STREAMS=8 +# Days after a run finishes before its per-turn progress events are deleted +# (state and terminal events are kept). 0 keeps them forever. +# SIMPLEAUDIT_EVENT_RETENTION_DAYS=30 +# Seconds between the Agent page's automatic Open WebUI knowledge/tool syncs +# per user and workspace (the Sync button always runs). 0 syncs on every load. +# SIMPLEAUDIT_OPENWEBUI_SYNC_TTL_S=60 diff --git a/accounts/services.py b/accounts/services.py index 65e73410..d58461fa 100644 --- a/accounts/services.py +++ b/accounts/services.py @@ -129,6 +129,31 @@ def grant_default_project(user) -> Project | None: return project +def request_roles(request) -> dict: + """``{project_id: role}`` for the request's user, loaded once per request. + + Middleware, context processors and write checks all need the user's + memberships; one query here replaces several per page. + """ + cached = getattr(request, "_studio_roles", None) + if cached is None: + user = getattr(request, "user", None) + cached = ( + dict(ProjectMembership.objects.filter(user=user).values_list("project_id", "role")) + if user is not None and user.is_authenticated else {} + ) + request._studio_roles = cached + return cached + + +def request_has_access(request, project) -> bool: + """``ensure_project_access`` for the request's user, from ``request_roles``.""" + user = getattr(request, "user", None) + if not user or not user.is_authenticated: + return False + return project.slug == DEFAULT_PROJECT_SLUG or project.pk in request_roles(request) + + def ensure_project_access(user, project) -> bool: """Return True if the user may view this project's content. diff --git a/audits/comparison.py b/audits/comparison.py index 1319cbe8..3eafca1d 100644 --- a/audits/comparison.py +++ b/audits/comparison.py @@ -121,23 +121,30 @@ def compare_runs(project, run_ids: list[int]) -> dict: from scenarios.models import ScenarioSetVersionItem run_results: dict[int, dict[str, dict]] = {} + item_maps: dict[int, dict[str, str]] = {} # runs often share a pinned version for run in runs: # Map version_item_id (str) -> scenario key for this run's pinned version. item_map: dict[str, str] = {} if run.scenario_set_version_id: - items = ScenarioSetVersionItem.objects.filter( - version_id=run.scenario_set_version_id - ).select_related("scenario") - for item in items: - item_map[str(item.id)] = item.scenario.key + if run.scenario_set_version_id not in item_maps: + item_maps[run.scenario_set_version_id] = { + str(item_id): key + for item_id, key in ScenarioSetVersionItem.objects.filter( + version_id=run.scenario_set_version_id + ).values_list("id", "scenario__key") + } + item_map = item_maps[run.scenario_set_version_id] - rows = ScenarioResult.objects.filter(run_id=str(run.id)) + # The small list_summary projection carries every field used here. + rows = ScenarioResult.objects.filter(run_id=run.id).only( + "version_item_id", "status", "attempts", "list_summary" + ) m = {} for row in rows: key = item_map.get(row.version_item_id, f"item_{row.version_item_id}") - r_data = row.result or {} + r_data = row.list_summary or {} # For repeated results, use aggregated severity - if "reps" in r_data and isinstance(r_data.get("reps"), list): + if r_data.get("has_reps"): sev = r_data.get("aggregated_severity") agree = r_data.get("agreement_rate") summary = f"{sev} ({int(agree*100)}% agreement)" if agree is not None else sev diff --git a/audits/events.py b/audits/events.py index bc965c0b..995467c0 100644 --- a/audits/events.py +++ b/audits/events.py @@ -23,9 +23,9 @@ class AuditEvent(models.Model): """Append-only durable progress event for an audit run.""" - run_id = models.PositiveBigIntegerField(db_index=True) - version_item_id = models.CharField(max_length=64, db_index=True) - kind = models.CharField(max_length=64, db_index=True) + run_id = models.PositiveBigIntegerField() + version_item_id = models.CharField(max_length=64) + kind = models.CharField(max_length=64) payload = models.JSONField(default=dict, blank=True) created_at = models.DateTimeField(auto_now_add=True) @@ -34,12 +34,68 @@ class Meta: ordering = ["id"] indexes = [ models.Index(fields=["run_id", "id"], name="audit_event_run_id_idx"), + # Per-scenario state/attempt lookups and the progress snapshot group-by. + models.Index(fields=["run_id", "kind", "version_item_id"], name="audit_event_run_kind_item_idx"), ] def __str__(self) -> str: return f"AuditEvent(run={self.run_id} {self.kind} vi={self.version_item_id})" +def trial_severities(result: dict) -> list: + """Every judged trial in one stored result (one per repetition).""" + reps = result.get("reps") + if isinstance(reps, list): + return [(r or {}).get("severity") for r in reps] + return [result.get("severity")] + + +def trial_counts(status: str, result: dict | None) -> dict[str, int]: + """Per-row trial tallies kept beside the (large) ``result`` JSON. + + Repetitions count as separate trials. ERROR severities and non-completed + scenarios are excluded from ``n_judged`` and counted in ``n_errors``; + "ungraded" trials (a judgment without a verdict) are counted apart. + """ + out = {"n_pass": 0, "n_judged": 0, "n_errors": 0, "n_ungraded": 0} + if status != "completed": + out["n_errors"] = 1 + return out + for sev in trial_severities(result or {}): + if not sev or sev == "ERROR": + out["n_errors"] += 1 + elif sev == "ungraded": + out["n_ungraded"] += 1 + else: + out["n_judged"] += 1 + if sev == "pass": + out["n_pass"] += 1 + return out + + +def result_projection(result: dict | None) -> dict: + """The few fields the run's Results list shows, lifted out of the big ``result``. + + Transcripts dominate a stored result; the list only needs verdicts and short + summaries, so ``ScenarioResult.list_summary`` keeps just those. + """ + r = result if isinstance(result, dict) else {} + reps = r.get("reps") + return { + "severity": r.get("severity"), + "aggregated_severity": r.get("aggregated_severity"), + "summary": r.get("summary", ""), + "agreement_rate": r.get("agreement_rate"), + "severity_distribution": r.get("severity_distribution"), + "has_reps": isinstance(reps, list), + "reps": ( + [{"severity": x.get("severity"), "summary": x.get("summary", "")} for x in reps if isinstance(x, dict)] + if isinstance(reps, list) else [] + ), + "agentic_status": (r.get("agentic_evaluation") or {}).get("status", ""), + } + + class ScenarioResult(models.Model): """Idempotent final result for one scenario in one run.""" @@ -51,8 +107,29 @@ class ScenarioResult(models.Model): # judgment, ...). Large raw transcripts may instead be offloaded to object # storage; this column holds the structured result record. result = models.JSONField(default=dict, blank=True) + # Trial tallies derived from ``status`` + ``result`` (see ``trial_counts``). + # They let run totals be a SUM over small columns instead of loading every + # transcript; ``save()`` keeps them in step with the JSON. + n_pass = models.PositiveIntegerField(default=0) + n_judged = models.PositiveIntegerField(default=0) + n_errors = models.PositiveIntegerField(default=0) + n_ungraded = models.PositiveIntegerField(default=0) + # Small projection of ``result`` for list pages (see ``result_projection``). + list_summary = models.JSONField(default=dict, blank=True) updated_at = models.DateTimeField(auto_now=True) + COUNT_FIELDS = ("n_pass", "n_judged", "n_errors", "n_ungraded", "list_summary") + + def save(self, *args, **kwargs): + update_fields = kwargs.get("update_fields") + if update_fields is None or {"status", "result"} & set(update_fields): + for name, value in trial_counts(self.status, self.result).items(): + setattr(self, name, value) + self.list_summary = result_projection(self.result) + if update_fields is not None: + kwargs["update_fields"] = [*update_fields, *self.COUNT_FIELDS] + super().save(*args, **kwargs) + class Meta: db_table = "core_scenario_result" constraints = [ @@ -61,6 +138,9 @@ class Meta: name="unique_scenario_result_per_run_item", ), ] + indexes = [ + models.Index(fields=["run_id", "status"], name="scenario_result_run_status_idx"), + ] def __str__(self) -> str: return f"ScenarioResult(run={self.run_id} vi={self.version_item_id} {self.status})" @@ -78,6 +158,50 @@ def append_event(run_id: int | str, version_item_id: str, kind: str, payload: di return event.id +@retry_if_locked +def append_events(run_id: int | str, version_item_id: str, items: list[tuple[str, dict | None]]) -> None: + """Append several events in one INSERT, keeping their order (ids ascend).""" + if not items: + return + AuditEvent.objects.bulk_create([ + AuditEvent(run_id=int(run_id), version_item_id=version_item_id, kind=kind, payload=payload or {}) + for kind, payload in items + ]) + + +# Per-turn progress events: by far the most numerous rows, and only useful +# while a run is live (the progress view and "Latest Activity"). State and +# terminal events are kept for good. +PRUNABLE_EVENT_KINDS = ("scenario_turn", "scenario_rep_started") + + +def prune_finished_run_events(older_than_days: int, batch_size: int = 5000) -> int: + """Delete per-turn events of runs that finished more than ``older_than_days`` ago. + + Deletes in id batches so a large backlog never holds one long transaction. + Returns the number of rows deleted. + """ + from datetime import timedelta + + from django.utils import timezone + + from audits.models import AuditRun + + if older_than_days <= 0: + return 0 + cutoff = timezone.now() - timedelta(days=older_than_days) + finished = AuditRun.objects.filter( + status__in=AuditRun.TERMINAL_STATUSES, finished_at__lt=cutoff + ).values("id") + qs = AuditEvent.objects.filter(run_id__in=finished, kind__in=PRUNABLE_EVENT_KINDS) + deleted = 0 + while True: + ids = list(qs.order_by().values_list("id", flat=True)[:batch_size]) + if not ids: + return deleted + deleted += AuditEvent.objects.filter(id__in=ids).delete()[0] + + def _event_dict(e: AuditEvent) -> dict: return { "id": e.id, diff --git a/audits/experiments.py b/audits/experiments.py index 3be7bf8e..8823eb09 100644 --- a/audits/experiments.py +++ b/audits/experiments.py @@ -593,7 +593,12 @@ def prefetch_experiment_runs(): """Prefetch for Experiment querysets whose pages show factors and summaries.""" from django.db.models import Prefetch - return Prefetch("runs", queryset=AuditRun.objects.select_related(*RUN_FACTOR_RELATED)) + # Factors and summaries read the model/judge snapshots and generation + # parameters; the runtime and trace blobs are never shown on these pages. + return Prefetch( + "runs", + queryset=AuditRun.objects.select_related(*RUN_FACTOR_RELATED).defer("runtime_metadata", "trace_config"), + ) class MembershipError(ValueError): diff --git a/audits/migrations/0011_scenario_result_trial_counts.py b/audits/migrations/0011_scenario_result_trial_counts.py new file mode 100644 index 00000000..476285c2 --- /dev/null +++ b/audits/migrations/0011_scenario_result_trial_counts.py @@ -0,0 +1,99 @@ +# Generated by Django 5.2.4 on 2026-11-11 16:38 + +from django.db import migrations, models + + +def _trial_counts(status, result): + """Frozen copy of ``audits.events.trial_counts`` (migrations must not import app code).""" + out = {"n_pass": 0, "n_judged": 0, "n_errors": 0, "n_ungraded": 0} + if status != "completed": + out["n_errors"] = 1 + return out + result = result or {} + reps = result.get("reps") + severities = ( + [(r or {}).get("severity") for r in reps] + if isinstance(reps, list) + else [result.get("severity")] + ) + for sev in severities: + if not sev or sev == "ERROR": + out["n_errors"] += 1 + elif sev == "ungraded": + out["n_ungraded"] += 1 + else: + out["n_judged"] += 1 + if sev == "pass": + out["n_pass"] += 1 + return out + + +def _projection(result): + """Frozen copy of ``audits.events.result_projection``.""" + r = result if isinstance(result, dict) else {} + reps = r.get("reps") + return { + "severity": r.get("severity"), + "aggregated_severity": r.get("aggregated_severity"), + "summary": r.get("summary", ""), + "agreement_rate": r.get("agreement_rate"), + "severity_distribution": r.get("severity_distribution"), + "has_reps": isinstance(reps, list), + "reps": ( + [{"severity": x.get("severity"), "summary": x.get("summary", "")} for x in reps if isinstance(x, dict)] + if isinstance(reps, list) else [] + ), + "agentic_status": (r.get("agentic_evaluation") or {}).get("status", ""), + } + + +def backfill_trial_counts(apps, schema_editor): + """Fill the new columns for existing rows, streaming so big transcripts never pile up.""" + ScenarioResult = apps.get_model("audits", "ScenarioResult") + fields = ["n_pass", "n_judged", "n_errors", "n_ungraded", "list_summary"] + batch = [] + for row in ScenarioResult.objects.only("id", "status", "result").iterator(chunk_size=200): + for name, value in _trial_counts(row.status, row.result).items(): + setattr(row, name, value) + row.list_summary = _projection(row.result) + batch.append(row) + if len(batch) >= 500: + ScenarioResult.objects.bulk_update(batch, fields) + batch = [] + if batch: + ScenarioResult.objects.bulk_update(batch, fields) + + +class Migration(migrations.Migration): + dependencies = [ + ("audits", "0010_experiment_run_membership"), + ] + + operations = [ + migrations.AddField( + model_name="scenarioresult", + name="n_errors", + field=models.PositiveIntegerField(default=0), + ), + migrations.AddField( + model_name="scenarioresult", + name="n_judged", + field=models.PositiveIntegerField(default=0), + ), + migrations.AddField( + model_name="scenarioresult", + name="n_pass", + field=models.PositiveIntegerField(default=0), + ), + migrations.AddField( + model_name="scenarioresult", + name="n_ungraded", + field=models.PositiveIntegerField(default=0), + ), + migrations.AddField( + model_name="scenarioresult", + name="list_summary", + field=models.JSONField(blank=True, default=dict), + ), + migrations.RunPython(backfill_trial_counts, migrations.RunPython.noop), + ] diff --git a/audits/migrations/0012_perf_indexes.py b/audits/migrations/0012_perf_indexes.py new file mode 100644 index 00000000..69e511f8 --- /dev/null +++ b/audits/migrations/0012_perf_indexes.py @@ -0,0 +1,70 @@ +# Generated by Django 5.2.4 on 2026-11-11 16:41 + +from django.conf import settings +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("accounts", "0005_sitesetting"), + ("audits", "0011_scenario_result_trial_counts"), + ("judges", "0003_composable_judges"), + ("model_registry", "0023_perf_indexes"), + ("prompts", "0001_initial"), + ("scenarios", "0005_convert_legacy_agentic_controls"), + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ] + + operations = [ + migrations.AlterField( + model_name="auditevent", + name="kind", + field=models.CharField(max_length=64), + ), + migrations.AlterField( + model_name="auditevent", + name="run_id", + field=models.PositiveBigIntegerField(), + ), + migrations.AlterField( + model_name="auditevent", + name="version_item_id", + field=models.CharField(max_length=64), + ), + migrations.AddIndex( + model_name="auditevent", + index=models.Index( + fields=["run_id", "kind", "version_item_id"], + name="audit_event_run_kind_item_idx", + ), + ), + migrations.AddIndex( + model_name="auditrun", + index=models.Index( + fields=["project", "archived", "-created_at"], + name="audit_run_proj_arch_created", + ), + ), + migrations.AddIndex( + model_name="auditrun", + index=models.Index( + fields=["status", "archived"], name="audit_run_status_arch" + ), + ), + migrations.AddIndex( + model_name="auditrun", + index=models.Index( + fields=["monitor", "created_at"], name="audit_run_monitor_created" + ), + ), + migrations.AddIndex( + model_name="auditrun", + index=models.Index(fields=["started_at"], name="audit_run_started_at"), + ), + migrations.AddIndex( + model_name="scenarioresult", + index=models.Index( + fields=["run_id", "status"], name="scenario_result_run_status_idx" + ), + ), + ] diff --git a/audits/models.py b/audits/models.py index 3fed152d..ef821ce0 100644 --- a/audits/models.py +++ b/audits/models.py @@ -92,6 +92,14 @@ class Status(models.TextChoices): class Meta: db_table = "core_audit_run" ordering = ["-created_at"] + indexes = [ + # Dashboard grid: one project's runs, newest first. + models.Index(fields=["project", "archived", "-created_at"], name="audit_run_proj_arch_created"), + # Sweeper / health: runs by status. + models.Index(fields=["status", "archived"], name="audit_run_status_arch"), + models.Index(fields=["monitor", "created_at"], name="audit_run_monitor_created"), + models.Index(fields=["started_at"], name="audit_run_started_at"), + ] def __str__(self) -> str: return f"{self.name} ({self.status})" diff --git a/audits/monitors.py b/audits/monitors.py index 16a9bbbd..254e2974 100644 --- a/audits/monitors.py +++ b/audits/monitors.py @@ -14,6 +14,7 @@ from datetime import UTC, timedelta from django.db import transaction +from django.db.models import Sum from django.utils import timezone from audits.events import ScenarioResult @@ -438,14 +439,6 @@ def run_due_monitors(now=None) -> list[int]: # ─── Drift statistics ──────────────────────────────────────────────────────── -def _trial_severities(result: dict) -> list[str]: - """Every judged trial in one ScenarioResult (one per repetition).""" - reps = result.get("reps") - if isinstance(reps, list): - return [(r or {}).get("severity") for r in reps] - return [result.get("severity")] - - def pass_counts(run_ids: list[int]) -> dict[int, dict]: """Per run: judged trials ``n``, passing trials ``k`` and errored trials. @@ -457,24 +450,16 @@ def pass_counts(run_ids: list[int]) -> dict[int, dict]: fail: excluded from ``n`` and counted in ``ungraded``. """ out = {rid: {"k": 0, "n": 0, "errors": 0, "ungraded": 0} for rid in run_ids} - rows = ScenarioResult.objects.filter(run_id__in=run_ids).only("run_id", "status", "result") + rows = ( + ScenarioResult.objects.filter(run_id__in=run_ids) + .values("run_id") + .annotate( + k=Sum("n_pass"), n=Sum("n_judged"), + errors=Sum("n_errors"), ungraded=Sum("n_ungraded"), + ) + ) for row in rows: - c = out.get(row.run_id) - if c is None: - continue - if row.status != "completed": - c["errors"] += 1 - continue - for sev in _trial_severities(row.result or {}): - if not sev or sev == "ERROR": - c["errors"] += 1 - continue - if sev == "ungraded": - c["ungraded"] += 1 - continue - c["n"] += 1 - if sev == "pass": - c["k"] += 1 + out[row["run_id"]] = {key: row[key] or 0 for key in ("k", "n", "errors", "ungraded")} return out diff --git a/audits/serializers.py b/audits/serializers.py index 1e9a72a2..ada83d9f 100644 --- a/audits/serializers.py +++ b/audits/serializers.py @@ -43,6 +43,73 @@ class Meta: read_only_fields = fields +class AuditRunListSerializer(serializers.ModelSerializer): + """A run as a list row: identity, progress, outcome; no config snapshots. + + The detail endpoint (and ``?expand=snapshots`` on the list) carries the + frozen target/auditor/judge/agent/generation snapshots. + """ + + scenario_set_version_id = serializers.IntegerField(source="scenario_set_version.id", read_only=True) + scenario_set_version_number = serializers.IntegerField(source="scenario_set_version.version", read_only=True) + scenario_set_version_hash = serializers.CharField(source="scenario_set_version.content_hash", read_only=True) + target_model = serializers.SerializerMethodField() + judge_model = serializers.SerializerMethodField() + judge = serializers.SerializerMethodField() + pass_rate = serializers.SerializerMethodField() + + class Meta: + model = AuditRun + fields = ( + "id", + "name", + "status", + "scenario_set_version_id", + "scenario_set_version_number", + "scenario_set_version_hash", + "target_model", + "agent", + "judge_model", + "judge_version", + "judge", + "simpleaudit_version", + "git_commit", + "queued_at", + "started_at", + "finished_at", + "total_scenarios", + "completed_scenarios", + "successful_scenarios", + "failed_scenarios", + "retried_scenarios", + "pass_rate", + "error_code", + "error_message", + "created_at", + ) + read_only_fields = fields + + def get_target_model(self, run) -> str: + from audits.services import frozen_name + + return frozen_name(run, "target") + + def get_judge_model(self, run) -> str: + from audits.services import frozen_name + + return frozen_name(run, "judge") + + def get_judge(self, run) -> str: + from audits.services import frozen_judge + + return frozen_judge(run)["label"] + + def get_pass_rate(self, run) -> float | None: + """Passing judged trials, percent (``None`` before any trial is judged).""" + c = (self.context.get("pass_counts") or {}).get(run.id) or {} + return round(c["k"] * 100 / c["n"], 1) if c.get("n") else None + + class AuditRunCreateSerializer(serializers.Serializer): name = serializers.CharField(max_length=250) scenario_set_version_id = serializers.IntegerField() diff --git a/audits/tests/test_export.py b/audits/tests/test_export.py index 9957bfa3..4b775989 100644 --- a/audits/tests/test_export.py +++ b/audits/tests/test_export.py @@ -334,7 +334,7 @@ def test_export_returns_json_with_disposition(self): self.assertEqual(resp.status_code, 200) self.assertTrue(resp["Content-Type"].startswith("application/json")) self.assertIn("attachment", resp["Content-Disposition"]) - body = resp.json() + body = json.loads(b"".join(resp.streaming_content)) self.assertEqual(body["summary"]["total_scenarios"], 1) self.assertEqual(body["summary"]["passed"], 1) diff --git a/audits/views.py b/audits/views.py index 0ca12c9e..545f0509 100644 --- a/audits/views.py +++ b/audits/views.py @@ -1,12 +1,15 @@ import json import logging +import os +import threading import time from django.db import close_old_connections -from django.http import JsonResponse, StreamingHttpResponse +from django.http import StreamingHttpResponse from django.utils.text import slugify from rest_framework import status from rest_framework.decorators import api_view, permission_classes +from rest_framework.pagination import LimitOffsetPagination from rest_framework.permissions import IsAuthenticated from rest_framework.response import Response @@ -14,7 +17,11 @@ from accounts.services import ensure_project_access, require_project_writable from audits.events import ScenarioResult, list_events from audits.models import AuditRun -from audits.serializers import AuditRunCreateSerializer, AuditRunSerializer +from audits.serializers import ( + AuditRunCreateSerializer, + AuditRunListSerializer, + AuditRunSerializer, +) from audits.services import ( build_simpleaudit_export, cancel_run, @@ -34,11 +41,95 @@ # Terminal event kinds: once one of these is seen for the run, the stream can end. _TERMINAL_EVENT_KINDS = {"run_completed", "run_failed", "run_cancelled"} _SSE_POLL_INTERVAL_SECONDS = 1.0 +# While a run is quiet the poll backs off (1s -> 5s) so idle viewers cost less; +# the first event resets it. A comment line every 15s keeps proxies from closing +# an idle stream. +_SSE_MAX_POLL_INTERVAL_SECONDS = 5.0 +_SSE_HEARTBEAT_SECONDS = 15.0 + + +def _env_int(name: str, default: int) -> int: + try: + return int(os.environ.get(name, default)) + except (TypeError, ValueError): + return default + + +# Each open stream holds one web thread for up to an hour. Caps are per web +# process: one user's tabs cannot take more than a few, and streams as a whole +# never take more than half the process's threads, so pages and /chat/authz +# always have threads left. Over the cap the stream answers 429 and the +# browser's EventSource retries with backoff. +_SSE_MAX_STREAMS_PER_USER = _env_int("SIMPLEAUDIT_SSE_MAX_STREAMS_PER_USER", 3) +_SSE_MAX_STREAMS = _env_int( + "SIMPLEAUDIT_SSE_MAX_STREAMS", max(1, _env_int("GUNICORN_THREADS", 16) // 2) +) +_sse_lock = threading.Lock() +_sse_open: dict[int, int] = {} + + +def _sse_acquire(user_id: int) -> bool: + with _sse_lock: + if sum(_sse_open.values()) >= _SSE_MAX_STREAMS: + return False + if _sse_open.get(user_id, 0) >= _SSE_MAX_STREAMS_PER_USER: + return False + _sse_open[user_id] = _sse_open.get(user_id, 0) + 1 + return True + + +def _sse_release(user_id: int) -> None: + with _sse_lock: + left = _sse_open.get(user_id, 0) - 1 + if left > 0: + _sse_open[user_id] = left + else: + _sse_open.pop(user_id, None) + + +class _SlotStream: + """Iterator around the event generator that frees the stream slot on close. + + Django calls ``close()`` on the streaming content when the response ends, + even if the generator never started (its own ``finally`` would not run then). + """ + + def __init__(self, gen, user_id: int): + self._gen = gen + self._user_id = user_id + self._released = False + + def __iter__(self): + return self + + def __next__(self): + return next(self._gen) + + def close(self): + try: + self._gen.close() + finally: + if not self._released: + self._released = True + _sse_release(self._user_id) _SSE_MAX_DURATION_SECONDS = 3600 # Max events read per poll; a backlog drains over successive polls. _SSE_BATCH_SIZE = 1000 +class _RunListPagination(LimitOffsetPagination): + """``?limit=&offset=``; response is ``{count, next, previous, results}``.""" + + default_limit = 50 + max_limit = 200 + + +class _ResultListPagination(LimitOffsetPagination): + # Each row carries a full transcript, so pages stay small. + default_limit = 100 + max_limit = 500 + + def _get_project_or_404(project_id) -> Project: try: return Project.objects.get(pk=project_id) @@ -56,8 +147,29 @@ def _require_project_access(user, project: Project): def list_audit_runs(request, project_id): project = _get_project_or_404(project_id) _require_project_access(request.user, project) - runs = AuditRun.objects.filter(project=project).select_related("scenario_set_version") - return Response(AuditRunSerializer(runs, many=True).data) + """Runs of a project, newest first, paginated. + + Rows are slim (no config snapshots); ``?expand=snapshots`` returns the full + detail shape for each row instead. + """ + runs = AuditRun.objects.filter(project=project).order_by("-created_at", "-id") + paginator = _RunListPagination() + if request.query_params.get("expand") == "snapshots": + page = paginator.paginate_queryset(runs.select_related("scenario_set_version"), request) + return paginator.get_paginated_response(AuditRunSerializer(page, many=True).data) + from audits.monitors import pass_counts + + page = paginator.paginate_queryset( + runs.select_related("scenario_set_version", "target_model", "judge_model").defer( + # Never shown in a list row; frozen names/judge read the target and judge snapshots. + "agent_config_snapshot", "auditor_config_snapshot", "generation_parameters_snapshot", + "trace_config", "runtime_metadata", "summary_metrics", + ), + request, + ) + counts = pass_counts([r.id for r in page]) + data = AuditRunListSerializer(page, many=True, context={"pass_counts": counts}).data + return paginator.get_paginated_response(data) @api_view(["GET"]) @@ -100,7 +212,11 @@ def export_audit_run(request, project_id, run_id): payload = build_simpleaudit_export(run) stem = slugify(run.name) or f"audit-run-{run.id}" - response = JsonResponse(payload, json_dumps_params={"indent": 2, "ensure_ascii": False}) + from infra.export_utils import buffered, json_chunks + + # The engine computes the summary over every result, so the document is + # built whole; it is encoded and sent incrementally instead of as one string. + response = StreamingHttpResponse(buffered(json_chunks(payload)), content_type="application/json") response["Content-Disposition"] = f'attachment; filename="{stem}.json"' return response @@ -127,7 +243,10 @@ def list_audit_run_results(request, project_id, run_id): str(vi.id): vi for vi in run.scenario_set_version.items.select_related("scenario", "revision") } - results = ScenarioResult.objects.filter(run_id=run.id).order_by("version_item_id") + paginator = _ResultListPagination() + results = paginator.paginate_queryset( + ScenarioResult.objects.filter(run_id=run.id).order_by("version_item_id", "id"), request + ) data = [] for r in results: item = items.get(str(r.version_item_id)) @@ -143,7 +262,7 @@ def list_audit_run_results(request, project_id, run_id): "updated_at": r.updated_at.isoformat() if r.updated_at else None, } ) - return Response(data) + return paginator.get_paginated_response(data) @api_view(["POST"]) @@ -251,7 +370,9 @@ def poll_audit_run_events(request, project_id, run_id): # Out-of-range ids (SQLite integers are 64-bit) count as "from the start". after_id = int(after_id_raw) if is_id(str(after_id_raw)) else 0 - return Response(list_events(run.id, after_id=after_id)) + # Capped like the SSE batch: a client resuming from 0 on a big run pages + # forward with the last id it received instead of pulling the whole log. + return Response(list_events(run.id, after_id=after_id, limit=_SSE_BATCH_SIZE)) def stream_audit_run_events(request, project_id, run_id): @@ -291,23 +412,46 @@ def stream_audit_run_events(request, project_id, run_id): def event_stream(): sent_after = after_id start = time.monotonic() + interval = _SSE_POLL_INTERVAL_SECONDS + last_write = start + terminal_idle = 0 try: while True: if time.monotonic() - start > _SSE_MAX_DURATION_SECONDS: yield "event: timeout\ndata: {}\n\n" break events = list_events(run.id, after_id=sent_after, limit=_SSE_BATCH_SIZE) + if events: + interval = _SSE_POLL_INTERVAL_SECONDS for e in events: + last_write = time.monotonic() sent_after = e["id"] data = json.dumps({"kind": e["kind"], "version_item_id": e["version_item_id"], "payload": e["payload"], "ts": e["created_at"]}) yield f"id: {e['id']}\nevent: {e['kind']}\ndata: {data}\n\n" if e["kind"] in _TERMINAL_EVENT_KINDS: return + if not events: + # A run can be finished while its terminal event is missing + # (pruned, or written by an older version): without this the + # stream would idle for the full hour holding a thread. Two + # quiet polls in a row let a just-written terminal event land first. + status = AuditRun.objects.filter(pk=run.id).values_list("status", flat=True).first() + terminal_idle = terminal_idle + 1 if status in AuditRun.TERMINAL_STATUSES else 0 + if terminal_idle >= 2: + kind = f"run_{status}" + data = json.dumps({"kind": kind, "version_item_id": "_run", "payload": {"synthetic": True}, "ts": None}) + yield f"event: {kind}\ndata: {data}\n\n" + return if len(events) < _SSE_BATCH_SIZE: # Release (or recycle) the DB connection while idle so a # long-lived stream does not pin a Postgres connection per thread. close_old_connections() - time.sleep(_SSE_POLL_INTERVAL_SECONDS) + if time.monotonic() - last_write >= _SSE_HEARTBEAT_SECONDS: + last_write = time.monotonic() + yield ": keepalive\n\n" + time.sleep(interval) + if not events: + interval = min(interval * 1.5, _SSE_MAX_POLL_INTERVAL_SECONDS) except (GeneratorExit, SystemExit, KeyboardInterrupt): # Normal stream teardown: the client disconnected (GeneratorExit is # raised when the response iterator is closed), or the worker process @@ -318,7 +462,13 @@ def event_stream(): logger.debug("SSE stream for run %s ended early (client disconnect or shutdown)", run.id) raise - response = StreamingHttpResponse(event_stream(), content_type="text/event-stream") + if not _sse_acquire(request.user.pk): + from django.http import HttpResponse + + busy = HttpResponse("Too many live streams; retrying shortly.", status=429, content_type="text/plain") + busy["Retry-After"] = "15" + return busy + response = StreamingHttpResponse(_SlotStream(event_stream(), request.user.pk), content_type="text/event-stream") response["Cache-Control"] = "no-cache" response["X-Accel-Buffering"] = "no" return response diff --git a/chat/views.py b/chat/views.py index a6112a5d..875c6bb4 100644 --- a/chat/views.py +++ b/chat/views.py @@ -44,9 +44,12 @@ def authz(request): request.GET.get("embed") == "admin" or request.headers.get("X-Studio-Embed") == "admin" ): - request.session["studio_chat_embed_admin"] = ( - time.time() + config.EMBED_ADMIN_FLAG_TTL - ) + # Re-stamp only once less than half the TTL is left: every write saves + # the session (a DB write plus Set-Cookie) on a path hit per iframe asset. + now = time.time() + current = request.session.get("studio_chat_embed_admin") + if not isinstance(current, (int, float)) or current - now < config.EMBED_ADMIN_FLAG_TTL / 2: + request.session["studio_chat_embed_admin"] = now + config.EMBED_ADMIN_FLAG_TTL response = HttpResponse(status=200) for header, value in config.identity(user).items(): response[header] = value diff --git a/config/settings.py b/config/settings.py index b54513a4..c10ef09f 100644 --- a/config/settings.py +++ b/config/settings.py @@ -210,6 +210,13 @@ def _csrf_trusted_origins() -> list[str]: # DB write. 0 writes each export at once (the test suite uses 0). OTLP_SPAN_FLUSH_INTERVAL_S = 1.0 +# Seconds between the Agent page's automatic Open WebUI knowledge/tool syncs for +# one user in one workspace (the Sync button always runs). 0 syncs on every load. +try: + OPENWEBUI_SYNC_TTL_S = float(os.environ.get("SIMPLEAUDIT_OPENWEBUI_SYNC_TTL_S", "60")) +except ValueError: + OPENWEBUI_SYNC_TTL_S = 60.0 + # SQLite is shared by the web server and worker threads: WAL lets reads run # during writes, IMMEDIATE takes the write lock up front (no upgrade deadlocks), # and the timeout waits for a busy lock instead of failing. diff --git a/config/urls.py b/config/urls.py index c106bd5b..17d2f2d2 100644 --- a/config/urls.py +++ b/config/urls.py @@ -62,6 +62,7 @@ ScenarioCreateView, ScenarioDeleteView, ScenarioDiffView, + ScenarioEditFormView, ScenarioEditView, ScenarioExportView, ScenarioImportView, @@ -99,6 +100,15 @@ def _serve_static(request, path): + response = _serve_static_file(request, path) + # Assets are not content-hashed, so keep the cache short; Last-Modified + # still lets the browser revalidate cheaply after it expires. + if response.status_code in (200, 304): + response.setdefault("Cache-Control", "public, max-age=300") + return response + + +def _serve_static_file(request, path): """Serve from STATIC_ROOT, then source dirs, then finders. Order matters for resilience against stale Docker layer caches (a known @@ -230,6 +240,7 @@ def home_view(request, *args, **kwargs): path("scenarios/version-delete//", ScenarioSetVersionDeleteView.as_view(), name="scenario_version_delete"), path("scenarios/create/", ScenarioCreateView.as_view(), name="scenario_create"), path("scenarios/edit//", ScenarioEditView.as_view(), name="scenario_edit"), + path("scenarios/edit-form//", ScenarioEditFormView.as_view(), name="scenario_edit_form"), path("scenarios/delete//", ScenarioDeleteView.as_view(), name="scenario_delete"), path("scenarios/revert//", ScenarioRevertView.as_view(), name="scenario_revert"), path("scenarios/diff//", ScenarioDiffView.as_view(), name="scenario_diff"), diff --git a/conftest.py b/conftest.py index a345dc09..1a6db234 100644 --- a/conftest.py +++ b/conftest.py @@ -16,6 +16,12 @@ def _synchronous_otlp_span_writes(settings): settings.OTLP_SPAN_FLUSH_INTERVAL_S = 0 +@pytest.fixture(autouse=True) +def _no_openwebui_sync_throttle(settings): + """Every Agent page load syncs in tests (ids are reused across tests).""" + settings.OPENWEBUI_SYNC_TTL_S = 0 + + def pytest_collection_modifyitems(items): """Copy Django ``@tag`` values onto pytest items as markers. diff --git a/deploy/openwebui-subpath/Caddyfile.subpath b/deploy/openwebui-subpath/Caddyfile.subpath index cc36731c..65a2eba9 100644 --- a/deploy/openwebui-subpath/Caddyfile.subpath +++ b/deploy/openwebui-subpath/Caddyfile.subpath @@ -30,6 +30,10 @@ } :{$STUDIO_PORT:8000} { + # Compress responses (event streams are not compressible content types, so + # SSE passes through untouched). + encode zstd gzip + # Never let a client supply its own identity. request_header -X-Studio-Email request_header -X-Studio-Name @@ -133,6 +137,12 @@ # Everything else -> Studio (Django). The subpath fork never emits # root-relative refs, so nothing from OWUI falls through here. handle { - reverse_proxy {$STUDIO_UPSTREAM:web:8000} + reverse_proxy {$STUDIO_UPSTREAM:web:8000} { + # Fail fast when Studio is down instead of hanging; allow slow pages. + transport http { + dial_timeout 5s + response_header_timeout 120s + } + } } } diff --git a/docs/architecture.md b/docs/architecture.md index 1c4ee3f1..a6e3918a 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -72,7 +72,18 @@ flowchart LR completed or failed, inline or by the sweeper. 5. **Watch.** The run page streams events over SSE (`/api/projects//audit-runs//events/`), with a polling - fallback, and refreshes the results list while the run is active. + fallback, and refreshes the results list while the run is active. Each web + process caps open streams per user and in total (see `docs/deployment.md`). + The JSON poll returns at most 1000 events per call. + +The run list (`/api/projects//audit-runs/`) and run results +(`.../audit-runs//results/`) APIs are paginated with `?limit=&offset=` +(defaults 50 and 100) and answer `{count, next, previous, results}`; the run +list is newest first. Run-list rows are slim (identity, progress, frozen model +and judge names, pass rate, no config snapshots); `?expand=snapshots` returns +the full detail shape per row. Run exports (JSON and CSV downloads, and the API +export) are streamed. Per-turn progress events of runs finished more than +`SIMPLEAUDIT_EVENT_RETENTION_DAYS` ago are pruned; state and terminal events stay. Monitors repeat one run setup on an interval or cron schedule. Each tick is an ordinary run, so the drift chart is a series of reproducible experiments. diff --git a/docs/deployment.md b/docs/deployment.md index 61af4c70..1f5ae42c 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -92,7 +92,9 @@ The full stack, each piece in its own container (`docker-compose.yml`): only the demo runs). Gunicorn is configured in `gunicorn.conf.py` with threaded workers; each open live audit-run event stream holds one thread, so `GUNICORN_WORKERS` × `GUNICORN_THREADS` (default 2 × 16) caps concurrent - viewers plus regular requests. Tune via those env vars. + viewers plus regular requests. Streams are capped per user and at half of + each process's threads (`SIMPLEAUDIT_SSE_MAX_STREAMS*`), so pages keep + working while many run pages are open. Tune via those env vars. - `open-webui`, `chat-proxy` — the chat assistant at `/chat/` (on by default, below). - `mock-model` — optional profile `mock`; a fake model API for testing. @@ -172,6 +174,9 @@ override them when you deliberately need to. | `SIMPLEAUDIT_SECURE_COOKIES` | Session and CSRF cookies are `Secure` by default (production). Set `false` only for a plain-HTTP local replica, e.g. `http://127.0.0.1`, where the browser would otherwise drop the login cookie. | | `DJANGO_CSRF_TRUSTED_ORIGINS` | Extra trusted origins (full URLs). | | `GUNICORN_BIND`, `GUNICORN_WORKERS`, `GUNICORN_THREADS`, `GUNICORN_TIMEOUT`, `GUNICORN_GRACEFUL_TIMEOUT`, `GUNICORN_KEEPALIVE` | Gunicorn tuning (`gunicorn.conf.py`). Defaults: `0.0.0.0:8000`, 2, 16, 60, 30, 5. | +| `SIMPLEAUDIT_SSE_MAX_STREAMS_PER_USER`, `SIMPLEAUDIT_SSE_MAX_STREAMS` | Caps on open live audit-run event streams per web process: per user (default 3) and in total (default half of `GUNICORN_THREADS`, so streams never take more than half the threads). Over a cap the stream answers `429` and the browser retries. | +| `SIMPLEAUDIT_EVENT_RETENTION_DAYS` | Days after a run finishes before its per-turn progress events (`scenario_turn`, `scenario_rep_started`) are deleted by the worker's hourly sweep; state and terminal events are kept. Default `30`; `0` keeps them forever. `manage.py prune_audit_events [--days N]` runs it by hand. | +| `SIMPLEAUDIT_OPENWEBUI_SYNC_TTL_S` | Seconds between the Agent page's automatic Open WebUI knowledge-base/tool syncs per user and workspace (the Sync button always runs). Default `60`; `0` syncs on every page load. | | `POSTGRES_DB`, `POSTGRES_USER`, `POSTGRES_PASSWORD`, `POSTGRES_HOST`, `POSTGRES_PORT`, `POSTGRES_CONN_MAX_AGE` | PostgreSQL connection (`compose` mode). | | `SIMPLEAUDIT_LOCAL_SQLITE` **internal** | `1` = local SQLite database at the repository's `dev.sqlite3`. Set automatically by the `dev` entry point and by the test suite. | | `SIMPLEAUDIT_SQLITE_PATH` | Override where the local dev SQLite file lives. Default: `/dev.sqlite3`. Relative paths resolve from the repo root; absolute paths are used as-is. Ignored unless `SIMPLEAUDIT_LOCAL_SQLITE=1`. | diff --git a/docs/performance-audit.md b/docs/performance-audit.md new file mode 100644 index 00000000..c9d33e7b --- /dev/null +++ b/docs/performance-audit.md @@ -0,0 +1,111 @@ +# Performance audit (2026-11-11) + +Source: read-only code review by three parallel audits (web views, data layer, API/SSE/worker/deploy). Nothing here is measured. Confirm against Sentry or production timings before and after each fix. Trigger: `502` on `GET /runs/data/?page=1&size=25&sort=created_at&dir=desc` on a production deployment. + +## Production evidence (Sentry, project `simula-l0/simpleauditstudio`, last 14 days) + +Queried with `sentry explore --dataset spans`. Traces are sampled at 10%, so counts are estimates; durations are real. + +Slowest web transactions (excluding `chat ` transactions, which are model calls made by the audit worker, not web pages): + +| Transaction | p50 | p95 | max | +|---|---|---|---| +| `/experiments/{experiment_id}/` | 206 ms | 9,276 ms | 10,376 ms | +| `/experiments/new/` | 170 ms | 3,432 ms | 3,435 ms | +| `/runs/data/` | 115 ms | 3,058 ms | 3,449 ms | +| `/scenarios/{set_id}/import/` | 1,370 ms | 1,520 ms | 1,520 ms | +| `/agents/{agent_id}/` | 711 ms | 1,321 ms | 1,321 ms | +| `/chat/authz` (about 34k calls) | 101 ms | 472 ms | 2,855 ms | +| `/otlp/v1/traces` | 8 ms | 100 ms | 183,294 ms | + +Root cause confirmed in span data: the `SELECT ... FROM core_scenario_result WHERE run_id IN (...)` query from `pass_counts` (it selects the `result` column). On `/runs/data/` it has p95 2.2 s and accounts for about half of all DB time. On `/experiments/{id}/` it runs 100 + 260 times with p95 3.3 s and 1.9 s and accounts for about 464 s of 486 s total DB time. The experiment page is therefore dominated by this one query pattern. + +Unresolved Sentry issues relevant to this audit: +- `OperationalError: connection timeout expired` on `/otlp/v1/traces` (SIMPLEAUDITSTUDIO-33, -34) and `/experiments/` (-32): DB connection exhaustion or contention, consistent with the thread starvation finding below. +- `SystemExit: 1` in `audits.views.event_stream` (SIMPLEAUDITSTUDIO-W, 65 events): SSE worker being killed. +- Chat model-call errors (`gpt-6-luna` content filter, `nhi.short-answer-active` 502 from NEL search provider) come from audited targets, not Studio pages. + +## Status (branch `perf/scenario-result-summary-columns`) + +Done: +- `ScenarioResult` trial-count columns and `list_summary` projection (migration 0011, batched backfill). `pass_counts` is one `GROUP BY`; the run Results list (`_result_rows`) no longer reads transcripts. +- Index migrations: `AuditRun` (project/archived/created_at, status/archived, monitor/created_at, started_at), `AuditEvent` (dropped three single-column indexes, added run/kind/item), `ScenarioResult(run_id, status)`, `OtlpSpan` (received_at, target/received_at). `audits` 0012, `model_registry` 0023. +- Quick wins: removed the Python scan of all events in the fault-injection path, capped the JSON event poll at 1000, `delete_prompt_version` compares in the database, cached the timezone list and git commit, queue health probe is two aggregate queries, runs grid and CSV defer unused JSON columns. +- SSE: poll backs off from 1 s to 5 s while idle, `: keepalive` comment every 15 s. +- `/chat/authz`: only re-stamps the embed-admin session flag when under half the TTL remains. +- `/static/`: `Cache-Control: public, max-age=300`. Caddyfile: `encode zstd gzip` and upstream dial/response timeouts (not validated, `caddy` is not installed locally). + +Also done in the second pass: +- Worker: `_sync_run_counters` is one conditional aggregate; `_maybe_finalize` reads two columns until the run is complete. +- SSE: a finished run whose terminal event is missing now ends the stream (after two quiet polls) with a synthetic `run_` event instead of idling for an hour. +- Model discovery and the health model-endpoint probe ping concurrently (8 threads) instead of serially. +- `nav` and `admin_status` share one `is_any_project_admin` lookup per request. +- Experiment page reload interval 15 s -> 30 s; `/experiments/new/` focus refetch limited to once a minute. + +Done in the third pass: +- Middleware and context processors: one `{project_id: role}` query per request (`accounts.services.request_roles`) feeds workspace access, the admin flag, the workspace switcher and `can_write`. Token-authenticated requests no longer write or cycle a session. +- SSE caps per web process: 3 streams per user and half of `GUNICORN_THREADS` in total (`SIMPLEAUDIT_SSE_MAX_STREAMS_PER_USER`, `SIMPLEAUDIT_SSE_MAX_STREAMS`); over the cap the stream answers 429 and the browser retries. The slot is released when Django closes the response. +- `AuditEvent`: the worker flusher writes queued turn events with one `bulk_create` per batch (up to 50); the sweeper deletes `scenario_turn` / `scenario_rep_started` events of runs finished more than `SIMPLEAUDIT_EVENT_RETENTION_DAYS` (30) ago, hourly; `manage.py prune_audit_events` does it by hand. +- API pagination: `/api/projects//audit-runs/` (default 50, max 200, newest first) and `.../results/` (default 100, max 500) answer `{count, next, previous, results}`. `run_demo_audits` follows pages. +- Agent page: the GET-time Open WebUI sync runs at most once per `SIMPLEAUDIT_OPENWEBUI_SYNC_TTL_S` (60) per user and workspace; its lists double as the ACL check, so a load makes 2 list calls instead of 4; mirrored rows are read in one query and only saved when changed. The Sync button always runs. +- Scenarios page: each edit form is fetched on first "Edit" click (`/scenarios/edit-form//`) instead of being rendered for every scenario. +- Compare uses the `list_summary` projection and loads each pinned version's items once; its run picker reads only the shown columns; experiment pages defer `runtime_metadata` and `trace_config`. + +Done in the fourth pass: +- Exports stream: the run page JSON/CSV download (CSV reads only the `list_summary` projection; JSON reads results in chunks of 100), the dashboard runs CSV (batches of 500 runs, pass counts per batch), and the API engine export (built whole because the engine computes its summary over every result, then encoded and sent incrementally). +- Run list API rows are slim (`AuditRunListSerializer`: no config snapshots, adds frozen target/judge names and pass rate); `?expand=snapshots` returns the old full rows. + +Remaining idea, not started: +- Serving SSE from an async worker so streams use no threads at all. + +## Tier 1: likely causes of 502s and slow pages + +1. **`ScenarioResult.result` loaded in full everywhere.** Severity and counts are read out of full transcript JSON in: + - `pass_counts` (`audits/monitors.py:449`), called from `infra/runs_table.py` (`/runs/data/`, CSV), `infra/ui.py` (ExperimentsView), `audits/experiments.py`, `audits/monitors.py` (drift series). + - `_result_rows` (`infra/ui.py:3477`): run detail and `/runs//results-fragment/`. `templates/partials/run_detail_body.html` refetches the fragment on every `scenario_completed` event, so a run with N scenarios loads N blobs N times. + - Exports (`RunExportView`, `build_simpleaudit_export`, export-html), `compare_runs`, `_refresh_run_summary_metrics`. + - Fix: denormalize small columns at write time (`upsert_scenario_result`): `severity`, `aggregated_severity`, `agreement_rate`, and per-row trial counts (`n_pass`, `n_judged`, `n_errors`, `n_ungraded`). `pass_counts` becomes one `GROUP BY run_id` with `Sum`. Needs a migration with a batched backfill. Portable across SQLite and Postgres (the `reps` array cannot be aggregated portably in the ORM, hence the columns). +2. **SSE streams starve the thread pool** (`audits/views.py:291-318`, `gunicorn.conf.py`). Each viewer holds one of 32 gthread threads for up to 1h and polls the DB every second. Static files, `/chat/authz` and `/healthz` use the same threads. The loop never checks `run.status`; there is no heartbeat. Fix: per-user cap, poll backoff, end on terminal status, heartbeat plus `retry`, serve SSE from an async or separate service. +3. **Slow outbound calls in web requests vs 60s gunicorn timeout.** `DiscoverModelsView` (sequential probes, 15s each), OTLP credential test (60s), `collect_health` (serial probes), `AgentDetailView` GET (4-6 Open WebUI calls plus row writes). Fix: thread pools, short timeouts, caching, move Agent sync off GET. +4. **`/chat/authz` on every Open WebUI asset.** Full Django stack with about 10 queries per request. Fix: cache per session for 30-60s, light path that skips ProjectMiddleware, skip forward-auth for public assets in Caddy, skip the admin query for superusers, avoid session writes on every embed call. +5. **Middleware and context processors.** About 10-14 queries on every page and poll (session, user, project, membership, admin checks, workspaces, write access, site settings). No `CACHES`, DB-backed sessions. Fix: resolve role once per request, add `CACHES`, `cached_db` sessions, cache `get_site_setting`. `TokenSessionBridgeMiddleware` creates a session per API request; skip it for `/api/`. + +## Tier 2: indexes and unbounded queries + +- `AuditRun`: add `(project, archived, -created_at)`, `(status, archived)` (partial on non-terminal for the sweeper on Postgres), `(monitor, created_at)`, `started_at`. +- `AuditEvent`: drop single-column `run_id`, `version_item_id`, `kind` indexes; add `(run_id, kind, version_item_id)`; add retention for finished runs; batch turn-event inserts with `bulk_create`; make `progress_snapshot` O(1) or cheaper. +- `ScenarioResult`: add `(run_id, status, attempts)`; collapse the 3-4 per-scenario COUNT queries in `_sync_run_counters` / `_maybe_finalize` into one conditional aggregate. +- `OtlpSpan`: index `received_at`; drop redundant `target_id` index; move pruning to the sweeper. +- Unbounded list APIs: `/api/projects//audit-runs/` (all snapshot JSON), results API, `poll_audit_run_events`, `list_scenarios`, `list_versions`. +- `.defer()` unused JSON columns on the runs grid and CSV (CSV loads up to 10,000 full rows); stream CSV/JSON exports with `.iterator()`. +- N+1: `ScenarioListSerializer.get_latest_revision`, `export_scenarios`, `NewExperimentView` connections, judge pickers, scenario publish/import loops, `RunDetailView` frozen-model FKs. +- Page weight: experiment detail does `location.reload()` every 15s; scenarios page renders every edit form eagerly; `/experiments/new/` refetches on every window focus; result detail page ships the blob twice. +- Python-side scans: `delete_prompt_version` (compare snapshots in Python), `_iter_attempted_events` (fault-injection path). +- Transactions holding network calls: `run_due_monitors` and `create_audit_run` call `discover_target_metadata` inside atomic blocks. + +## Tier 3: deploy and config + +- Compression: `encode gzip zstd` in Caddy (not Django gzip, which buffers SSE). +- Static files: whitenoise or Caddy `file_server` with cache headers; currently served by Django through the same 32 threads. +- Caddy upstream: add dial and response-header timeouts, health checks. +- Gunicorn `timeout=60` is shorter than several slow endpoints. +- Compose `web` runs migrate, bootstrap, seed and collectstatic on every start; split into an init step. +- Sentry `enable_logs=True` ships every log line; use `traces_sampler` for hot routes. +- `readyz` caches a "pending migrations" True forever via `lru_cache`. + +## Other issues noticed + +- `next_revision` / `next_version` use `count()+1`; collides after a delete. Use `Max()+1`. +- `VisualizerImageView` fetches arbitrary URLs (SSRF, no size limit). +- Public experiment pages and `/sitemap.xml` are anonymous and uncached. +- `Meta.ordering` on joined fields (`project__name`, `connection__name`, ...) adds joins to `.first()` / `.get()`. + +## Suggested order + +1. Tier 2 indexes, `limit`s and `.defer()` quick wins. +2. Tier 1 item 1 (`ScenarioResult` summary columns). +3. Tier 1 items 2-5 and Tier 3. + +## Not covered in depth + +`infra/engine.py`, `infra/seed.py`, `audits/agentic/*`, Agent/KB/Tool query paths beyond list views, migrations other than an index grep. diff --git a/infra/context_processors.py b/infra/context_processors.py index d26ed57f..5a7888d6 100644 --- a/infra/context_processors.py +++ b/infra/context_processors.py @@ -29,16 +29,27 @@ def gravatar_url(request): } +def _request_is_admin(request) -> bool: + """``is_any_project_admin`` once per request (nav and admin_status both need it).""" + user = getattr(request, "user", None) + if user is None or not user.is_authenticated: + return False + if user.is_superuser: + return True + from accounts.models import ProjectMembership + from accounts.services import request_roles + + return ProjectMembership.Role.ADMIN in request_roles(request).values() + + def admin_status(request): """Expose ``is_admin`` / ``is_superuser`` to templates for nav gating. Admin = superuser OR holds an ADMIN membership in any project. Computed once per request; anonymous users get False without hitting the DB. """ - from accounts.services import is_any_project_admin - user = getattr(request, "user", None) - return {"is_admin": is_any_project_admin(user), "is_superuser": bool(user and user.is_superuser)} + return {"is_admin": _request_is_admin(request), "is_superuser": bool(user and user.is_superuser)} def workspaces(request): @@ -57,15 +68,11 @@ def workspaces(request): from django.db.models import Q from accounts.models import Project, ProjectMembership - from accounts.services import DEFAULT_PROJECT_SLUG - - projects = Project.objects.filter( - Q(memberships__user=user) | Q(slug=DEFAULT_PROJECT_SLUG) - ).distinct() + from accounts.services import DEFAULT_PROJECT_SLUG, request_roles - roles = dict( - ProjectMembership.objects.filter(project__in=list(projects), user=user).values_list("project_id", "role") - ) + roles = request_roles(request) + # Membership ids come from the per-request cache: no join, no DISTINCT. + projects = Project.objects.filter(Q(pk__in=list(roles)) | Q(slug=DEFAULT_PROJECT_SLUG)) current_id = request.session.get("active_project_id") items = [] @@ -127,10 +134,9 @@ def nav(request): user = getattr(request, "user", None) if user is None or not user.is_authenticated: return {} - from accounts.services import is_any_project_admin from infra.chat_feature import chat_enabled - admin = is_any_project_admin(user) + admin = _request_is_admin(request) path = request.path chat_on = chat_enabled() # The wrapper page (model-picker bar + iframe of the embedded service at diff --git a/infra/export_utils.py b/infra/export_utils.py index e4e1d93a..82ccbf6f 100644 --- a/infra/export_utils.py +++ b/infra/export_utils.py @@ -13,3 +13,60 @@ def csv_safe(value): if isinstance(value, str) and value.startswith(_FORMULA_PREFIXES): return "'" + value return value + + +def buffered(chunks, size: int = 64 * 1024): + """Join many small strings into ~``size`` pieces for a streaming response. + + JSON/CSV writers emit tiny fragments; sending each as its own WSGI write + would be slow. Yields UTF-8 bytes. + """ + buf: list[str] = [] + length = 0 + for chunk in chunks: + buf.append(chunk) + length += len(chunk) + if length >= size: + yield "".join(buf).encode("utf-8") + buf, length = [], 0 + if buf: + yield "".join(buf).encode("utf-8") + + +def json_chunks(document, *, indent: int | None = 2): + """Encode ``document`` incrementally (no single giant string in memory).""" + import json + + return json.JSONEncoder(indent=indent, ensure_ascii=False).iterencode(document) + + +def json_array_chunks(items, *, indent: int | None = 2): + """Stream an iterable of JSON-able items as one ``[...]`` array.""" + import json + + encoder = json.JSONEncoder(indent=indent, ensure_ascii=False) + pad = " " * indent if indent else "" + first = True + yield "[" + for item in items: + yield ("\n" if first else ",\n") if indent else ("" if first else ",") + first = False + encoded = "".join(encoder.iterencode(item)) + yield pad + encoded.replace("\n", "\n" + pad) if indent else encoded + yield "\n]" if indent and not first else "]" + + +def csv_rows(header: list, rows): + """Yield CSV text line by line for ``header`` then each row in ``rows``.""" + import csv + import io + + buf = io.StringIO() + writer = csv.writer(buf) + writer.writerow(header) + yield buf.getvalue() + for row in rows: + buf.seek(0) + buf.truncate() + writer.writerow(row) + yield buf.getvalue() diff --git a/infra/health.py b/infra/health.py index 9a64510f..7f6fc055 100644 --- a/infra/health.py +++ b/infra/health.py @@ -20,6 +20,7 @@ """ from __future__ import annotations +import functools import os import shutil import time @@ -177,6 +178,7 @@ def _studio_version() -> str | None: return None +@functools.lru_cache(maxsize=1) def _studio_commit() -> str | None: """Git commit of this checkout (HEAD), or None when not a git repo.""" from pathlib import Path @@ -240,34 +242,40 @@ def _model_endpoints_probe(user=None) -> dict[str, Any]: from accounts.models import ProjectMembership conns = conns.filter(project__memberships__user=user, project__memberships__role=ProjectMembership.Role.ADMIN) + # Resolve everything that touches the database or secrets here, then ping + # all connections concurrently: serial 2 s timeouts add up per dead endpoint. + targets = [] for conn in conns: base = (conn.base_url or "").strip() if not base: continue # If no API key is configured, don't ping — report "no key" api_key = connection_api_key(conn) + targets.append((conn, base, models_url(conn) if api_key else None, api_key)) + + def ping(target) -> tuple[str, float | None, str]: + _conn, _base, url, api_key = target if not api_key: - conn_status = "no_key" - latency = None - detail = "No API key configured" - else: - # Ping the connection once with auth - conn_status = "down" - latency = None - detail = "" - try: - import requests - start = _now_ms() - headers = {"Authorization": f"Bearer {api_key}"} - resp = requests.get(models_url(conn), timeout=2, headers=headers) - if resp.status_code == 200: - conn_status = "up" - latency = round(_now_ms() - start, 1) - else: - detail = f"HTTP {resp.status_code}" - except Exception as exc: # noqa: BLE001 - detail = f"{type(exc).__name__}" + return "no_key", None, "No API key configured" + try: + import requests + start = _now_ms() + resp = requests.get(url, timeout=2, headers={"Authorization": f"Bearer {api_key}"}) + if resp.status_code == 200: + return "up", round(_now_ms() - start, 1), "" + return "down", None, f"HTTP {resp.status_code}" + except Exception as exc: # noqa: BLE001 + return "down", None, f"{type(exc).__name__}" + + if len(targets) > 1: + from concurrent.futures import ThreadPoolExecutor + + with ThreadPoolExecutor(max_workers=min(8, len(targets))) as pool: + pings = list(pool.map(ping, targets)) + else: + pings = [ping(t) for t in targets] + for (conn, base, _url, _key), (conn_status, latency, detail) in zip(targets, pings, strict=True): models = [ {"id": m.id, "display_name": m.display_name, "model_id": m.model_id, "status": conn_status} for m in conn.enabled_models @@ -373,40 +381,28 @@ def _cpu_probe() -> dict[str, Any]: def _queue_throughput_probe() -> dict[str, Any]: + from django.db.models import Count, F, Sum + from audits.models import AuditRun - qs = AuditRun.objects.all() + S = AuditRun.Status + active_statuses = [ + S.PREPARING, S.TARGET_EXECUTION, S.AUDITING, S.JUDGING, S.AGGREGATION, S.REPORT_GENERATION, + ] + by_status = dict(AuditRun.objects.values_list("status").annotate(n=Count("id")).order_by()) counts = { - "queued": qs.filter(status=AuditRun.Status.QUEUED).count(), - "active": qs.filter( - status__in=[ - AuditRun.Status.PREPARING, - AuditRun.Status.TARGET_EXECUTION, - AuditRun.Status.AUDITING, - AuditRun.Status.JUDGING, - AuditRun.Status.AGGREGATION, - AuditRun.Status.REPORT_GENERATION, - ] - ).count(), - "completed": qs.filter(status=AuditRun.Status.COMPLETED).count(), - "failed": qs.filter(status=AuditRun.Status.FAILED).count(), - "cancelled": qs.filter(status=AuditRun.Status.CANCELLED).count(), + "queued": by_status.get(S.QUEUED, 0), + "active": sum(by_status.get(status, 0) for status in active_statuses), + "completed": by_status.get(S.COMPLETED, 0), + "failed": by_status.get(S.FAILED, 0), + "cancelled": by_status.get(S.CANCELLED, 0), } - # Scenarios in flight across active runs (sum of remaining work). - active_runs = qs.filter( - status__in=[ - AuditRun.Status.PREPARING, - AuditRun.Status.TARGET_EXECUTION, - AuditRun.Status.AUDITING, - AuditRun.Status.JUDGING, - AuditRun.Status.AGGREGATION, - AuditRun.Status.REPORT_GENERATION, - ] - ) - scenarios_in_flight = sum( - max(r.total_scenarios - r.completed_scenarios, 0) for r in active_runs - ) - return {"status": "up", "runs": counts, "scenarios_in_flight": scenarios_in_flight} + # Scenarios in flight across active runs (sum of remaining work), summed in + # the database so no run rows (with their large JSON columns) are loaded. + in_flight = AuditRun.objects.filter( + status__in=active_statuses, total_scenarios__gt=F("completed_scenarios") + ).aggregate(n=Sum(F("total_scenarios") - F("completed_scenarios")))["n"] or 0 + return {"status": "up", "runs": counts, "scenarios_in_flight": in_flight} # ─── Aggregate ─────────────────────────────────────────────────────────────── diff --git a/infra/management/commands/prune_audit_events.py b/infra/management/commands/prune_audit_events.py new file mode 100644 index 00000000..70af3433 --- /dev/null +++ b/infra/management/commands/prune_audit_events.py @@ -0,0 +1,27 @@ +"""Delete per-turn progress events of runs that finished long ago. + +Usage: + python manage.py prune_audit_events [--days N] + +The worker's sweeper already does this hourly (``SIMPLEAUDIT_EVENT_RETENTION_DAYS``, +default 30, 0 disables). Run it by hand to reclaim space at once. State and +terminal events are kept, so finished runs still show their outcome. +""" + +from django.core.management.base import BaseCommand + + +class Command(BaseCommand): + help = "Delete per-turn events (scenario_turn, scenario_rep_started) of long-finished runs." + + def add_arguments(self, parser): + parser.add_argument("--days", type=int, default=None, + help="Age in days after a run finished (default: SIMPLEAUDIT_EVENT_RETENTION_DAYS or 30).") + + def handle(self, *args, **options): + from audits.events import prune_finished_run_events + from infra.worker import _event_retention_days + + days = options["days"] if options["days"] is not None else _event_retention_days() + deleted = prune_finished_run_events(days) + self.stdout.write(self.style.SUCCESS(f"Deleted {deleted} event(s) of runs finished more than {days} day(s) ago.")) diff --git a/infra/management/commands/run_demo_audits.py b/infra/management/commands/run_demo_audits.py index 69e051bc..b54f7859 100644 --- a/infra/management/commands/run_demo_audits.py +++ b/infra/management/commands/run_demo_audits.py @@ -105,10 +105,10 @@ def agents(self) -> list[dict]: return self._request("GET", "/api/agents/") def latest_run(self, project_id: int) -> dict | None: - runs = self._request("GET", f"/api/projects/{project_id}/audit-runs/") - if not runs: - return None - return max(runs, key=lambda r: r.get("created_at") or "") + # Newest first, paginated: the first row of the first page is the latest. + page = self._request("GET", f"/api/projects/{project_id}/audit-runs/?limit=1&expand=snapshots") + runs = page.get("results") or [] + return runs[0] if runs else None def create_run(self, project_id: int, payload: dict) -> dict: return self._request("POST", f"/api/projects/{project_id}/audit-runs/create/", payload) @@ -117,7 +117,15 @@ def get_run(self, project_id: int, run_id: int) -> dict: return self._request("GET", f"/api/projects/{project_id}/audit-runs/{run_id}/") def run_results(self, project_id: int, run_id: int) -> list[dict]: - return self._request("GET", f"/api/projects/{project_id}/audit-runs/{run_id}/results/") + rows: list[dict] = [] + path = f"/api/projects/{project_id}/audit-runs/{run_id}/results/?limit=500" + while path: + page = self._request("GET", path) + rows.extend(page.get("results") or []) + nxt = page.get("next") + # ``next`` is absolute; keep only the path + query for _request. + path = nxt[nxt.index("/api/"):] if nxt else None + return rows def _extract_scenario_record(result_row: dict) -> dict: diff --git a/infra/middleware.py b/infra/middleware.py index e285eb81..1c3ce413 100644 --- a/infra/middleware.py +++ b/infra/middleware.py @@ -103,6 +103,10 @@ class TokenSessionBridgeMiddleware(MiddlewareMixin): ``ProjectMiddleware``) work with API token auth, not just session cookies. The session is only written when the user is anonymous, so existing session logins are never overwritten. + + The token user is attached to this request only and nothing is written to + the session: a token client sends its token on every call, and persisting it + created (and cycled) a session row per request. """ def process_request(self, request): @@ -121,9 +125,7 @@ def process_request(self, request): if not token.user.is_active: return request.user = token.user - request.session["user_id"] = token.user.id - request.session["auth_backend"] = "django.contrib.auth.backends.ModelBackend" - request.session.cycle_key() + request._studio_token_auth = True class ProjectMiddleware(MiddlewareMixin): @@ -139,21 +141,24 @@ def process_request(self, request): if not hasattr(request, "user") or not request.user.is_authenticated: return from accounts.models import Project - from accounts.services import ensure_project_access + from accounts.services import request_has_access project_id = request.session.get("active_project_id") if project_id: request.project = Project.objects.filter(pk=project_id).first() # A session can outlive the membership it was granted under (the # user was removed); never keep serving that workspace. - if request.project is not None and not ensure_project_access(request.user, request.project): + if request.project is not None and not request_has_access(request, request.project): request.project = None request.session.pop("active_project_id", None) if not request.project: membership = request.user.memberships.select_related("project").first() if membership: request.project = membership.project - request.session["active_project_id"] = request.project.id + # Token-authenticated calls carry no session cookie: writing here + # would create a fresh session row on every request. + if not getattr(request, "_studio_token_auth", False): + request.session["active_project_id"] = request.project.id if not request.project: # Backstop: a user with no membership (legacy/orphaned account) # still lands in the shared Default workspace rather than crashing diff --git a/infra/runs_table.py b/infra/runs_table.py index 0674c1fb..39120248 100644 --- a/infra/runs_table.py +++ b/infra/runs_table.py @@ -8,13 +8,11 @@ """ from __future__ import annotations -import csv -import io import json from datetime import datetime, time, timedelta from django.db.models import Q -from django.http import HttpResponse, JsonResponse +from django.http import JsonResponse from django.utils import timezone from django.views import View @@ -78,7 +76,11 @@ def filtered_runs(project, params): qs = AuditRun.objects.filter(project=project).select_related( "scenario_set_version__scenario_set", "target_model", "auditor_model", "judge_model", "judge_version__judge", "monitor", "created_by", - ).prefetch_related("experiments") + ).prefetch_related("experiments").defer( + # Large JSON the grid and CSV never read (run_row uses the target/auditor/ + # judge/generation snapshots and summary_metrics only). + "agent_config_snapshot", "trace_config", "runtime_metadata", + ) status = params.get("status", "") if status == "archived": qs = qs.filter(archived=True) @@ -260,18 +262,32 @@ class RunsExportView(ProjectMixin, View): """GET /runs/export.csv — the dashboard's current filters, as CSV.""" def get(self, request): + from django.http import StreamingHttpResponse + from audits.monitors import pass_counts + from infra.export_utils import buffered, csv_rows - runs = list(filtered_runs(request.project, request.GET).order_by(*_ordering(request.GET))[:10000]) - counts = pass_counts([r.id for r in runs]) + qs = filtered_runs(request.project, request.GET).order_by(*_ordering(request.GET))[:10000] cols = ["id", "name", "status", "scenario_set", "target", "auditor", "judge", "judge_model", "completed", "total", "reps", "pass_rate", "max_turns", "language", "duration", "created_by", "created_at"] - buf = io.StringIO() - writer = csv.writer(buf) - writer.writerow(cols) - for run in runs: - row = run_row(run, counts) - writer.writerow([csv_safe(row[c]) if row[c] is not None else "" for c in cols]) - response = HttpResponse(buf.getvalue(), content_type="text/csv; charset=utf-8") + + def body(): + # Up to 10,000 runs: stream them in batches (pass counts per batch) + # instead of holding every run row before the first byte goes out. + batch = [] + for run in qs.iterator(chunk_size=500): + batch.append(run) + if len(batch) == 500: + yield from _rows(batch) + batch = [] + yield from _rows(batch) + + def _rows(batch): + counts = pass_counts([r.id for r in batch]) if batch else {} + for run in batch: + row = run_row(run, counts) + yield [csv_safe(row[c]) if row[c] is not None else "" for c in cols] + + response = StreamingHttpResponse(buffered(csv_rows(cols, body())), content_type="text/csv; charset=utf-8") response["Content-Disposition"] = 'attachment; filename="runs.csv"' return response diff --git a/infra/tests/test_api_e2e.py b/infra/tests/test_api_e2e.py index 0a4f7145..b4cf97af 100644 --- a/infra/tests/test_api_e2e.py +++ b/infra/tests/test_api_e2e.py @@ -134,7 +134,7 @@ def test_full_audit_lifecycle(self): # 12. Verify results endpoint resp = self.client.get(f"/api/projects/{self.pid}/audit-runs/{run_id}/results/") assert resp.status_code == 200 - results = resp.json() + results = resp.json()["results"] self.assertEqual(len(results), 1) self.assertEqual(results[0]["status"], "completed") self.assertEqual(results[0]["result"]["severity"], "pass") diff --git a/infra/tests/test_audit_poll_results.py b/infra/tests/test_audit_poll_results.py index 64e2f5f2..16404a9a 100644 --- a/infra/tests/test_audit_poll_results.py +++ b/infra/tests/test_audit_poll_results.py @@ -115,7 +115,9 @@ def test_results_join_with_frozen_scenario_identity(self): ) resp = self.client.get(self._url(run.id)) self.assertEqual(resp.status_code, 200) - rows = resp.json() + body = resp.json() + self.assertEqual(body["count"], 1) + rows = body["results"] self.assertEqual(len(rows), 1) row = rows[0] self.assertEqual(row["status"], "completed") @@ -130,8 +132,55 @@ def test_results_empty_when_nothing_executed(self): run, _ = self._make_run() resp = self.client.get(self._url(run.id)) self.assertEqual(resp.status_code, 200) - self.assertEqual(resp.json(), []) + self.assertEqual(resp.json()["results"], []) + + def test_results_are_paginated(self): + from audits.events import upsert_scenario_result + + run, _ = self._make_run() + for i in range(3): + upsert_scenario_result(run.id, f"x{i}", status="completed", attempts=1, result={"severity": "pass"}) + body = self.client.get(self._url(run.id) + "?limit=2").json() + self.assertEqual((body["count"], len(body["results"])), (3, 2)) + self.assertIsNotNone(body["next"]) + rest = self.client.get(self._url(run.id) + "?limit=2&offset=2").json() + self.assertEqual(len(rest["results"]), 1) def test_results_unknown_run_404(self): resp = self.client.get(self._url(99999)) self.assertEqual(resp.status_code, 404) + + +class RunListPaginationTest(ResultsViewTest): + def test_run_list_is_paginated_newest_first(self): + from audits.models import AuditRun + + first, _ = self._make_run() + second = AuditRun.objects.get(pk=first.pk) + second.pk = None + second.save() + body = self.client.get(f"/api/projects/{self.project.id}/audit-runs/?limit=1").json() + self.assertEqual(body["count"], 2) + self.assertEqual([r["id"] for r in body["results"]], [second.id]) + + +class SlimRunListTest(ResultsViewTest): + def test_list_rows_are_slim_and_expand_restores_snapshots(self): + from django.db import connection + from django.test.utils import CaptureQueriesContext + + from audits.events import upsert_scenario_result + + run, _ = self._make_run() + upsert_scenario_result(run.id, "a", status="completed", attempts=1, result={"severity": "pass"}) + url = f"/api/projects/{self.project.id}/audit-runs/" + with CaptureQueriesContext(connection) as ctx: + row = self.client.get(url).json()["results"][0] + self.assertNotIn("target_config_snapshot", row) + self.assertNotIn("generation_parameters_snapshot", row) + self.assertEqual(row["id"], run.id) + self.assertEqual(row["pass_rate"], 100.0) + self.assertTrue(row["target_model"]) + self.assertLessEqual(len(ctx.captured_queries), 12, [q["sql"][:80] for q in ctx.captured_queries]) + full = self.client.get(url + "?expand=snapshots").json()["results"][0] + self.assertIn("target_config_snapshot", full) diff --git a/infra/tests/test_audit_sse.py b/infra/tests/test_audit_sse.py index 5a18567c..5533020d 100644 --- a/infra/tests/test_audit_sse.py +++ b/infra/tests/test_audit_sse.py @@ -15,7 +15,12 @@ class AuditSSEStreamTest(TestCase): def setUp(self): from accounts.models import Project, ProjectMembership, User + from audits import views + # Stream slots are per process; tests that never close a response must + # not starve the next test. + views._sse_open.clear() + self.addCleanup(views._sse_open.clear) self.user = User.objects.create_user(username="sse-user", password="pw12345") self.project = Project.objects.create(name="P", slug="p") ProjectMembership.objects.create(project=self.project, user=self.user, role=ProjectMembership.Role.AUDITOR) @@ -195,3 +200,55 @@ def fake_sleep(_seconds): self.assertRaises(SystemExit) as ctx: next(gen) self.assertEqual(ctx.exception.code, 1) + + +class AuditSSEFinishedRunTest(AuditSSEStreamTest): + """A finished run whose terminal event is missing must not hold a thread for an hour.""" + + def test_finished_run_without_terminal_event_ends_stream(self): + from audits.models import AuditRun + + run = self._make_run() + AuditRun.objects.filter(pk=run.id).update(status=AuditRun.Status.COMPLETED) + gen = self._get_raw_generator(run.id) + with mock.patch('audits.views.time.sleep'): + chunks = list(gen) + self.assertEqual(len(chunks), 1) + self.assertIn('event: run_completed', chunks[0]) + + +class AuditSSECapTest(AuditSSEStreamTest): + """Open streams are capped per user and per process; closing frees the slot.""" + + def _stream(self, run_id): + from django.test import RequestFactory + + from audits.views import stream_audit_run_events + + request = RequestFactory().get("/") + request.user = self.user + return stream_audit_run_events(request, self.project.id, run_id) + + def test_per_user_cap_and_release(self): + from audits import views + + run = self._make_run() + with mock.patch.object(views, "_SSE_MAX_STREAMS_PER_USER", 2), mock.patch.object(views, "_SSE_MAX_STREAMS", 10): + first, second = self._stream(run.id), self._stream(run.id) + self.assertEqual(self._stream(run.id).status_code, 429) + first.close() # Django closes the response when the client goes away + third = self._stream(run.id) + self.assertEqual(third.status_code, 200) + second.close() + third.close() + self.assertEqual(views._sse_open, {}) + + def test_process_cap(self): + from audits import views + + run = self._make_run() + with mock.patch.object(views, "_SSE_MAX_STREAMS", 1): + first = self._stream(run.id) + self.assertEqual(self._stream(run.id).status_code, 429) + first.close() + self.assertEqual(views._sse_open, {}) diff --git a/infra/tests/test_event_retention.py b/infra/tests/test_event_retention.py new file mode 100644 index 00000000..d2b6d5e1 --- /dev/null +++ b/infra/tests/test_event_retention.py @@ -0,0 +1,59 @@ +"""AuditEvent retention and batched inserts.""" +from datetime import timedelta +from io import StringIO + +from django.core.management import call_command +from django.test import TestCase +from django.utils import timezone + +from audits.events import ( + AuditEvent, + append_event, + append_events, + list_events, + prune_finished_run_events, +) +from audits.models import AuditRun +from infra.tests.factories import AuditRunFactory + + +class AppendEventsTests(TestCase): + def test_bulk_append_keeps_order(self): + run = AuditRunFactory() + append_event(run.id, "a", "scenario_attempted", {}) + append_events(run.id, "a", [("scenario_turn", {"turn": 1}), ("scenario_turn", {"turn": 2}), ("scenario_rep_started", None)]) + kinds = [(e["kind"], e["payload"].get("turn")) for e in list_events(run.id)] + self.assertEqual(kinds, [("scenario_attempted", None), ("scenario_turn", 1), ("scenario_turn", 2), ("scenario_rep_started", None)]) + + +class PruneTests(TestCase): + def _run(self, status, finished_days_ago): + run = AuditRunFactory() + finished = timezone.now() - timedelta(days=finished_days_ago) if finished_days_ago is not None else None + AuditRun.objects.filter(pk=run.pk).update(status=status, finished_at=finished) + for kind in ("scenario_attempted", "scenario_turn", "scenario_turn", "scenario_rep_started", "scenario_completed"): + append_event(run.id, "a", kind, {}) + append_event(run.id, "_run", "run_completed", {}) + return run + + def test_prunes_only_turn_events_of_old_finished_runs(self): + old = self._run(AuditRun.Status.COMPLETED, 40) + recent = self._run(AuditRun.Status.COMPLETED, 5) + active = self._run(AuditRun.Status.AUDITING, None) + self.assertEqual(prune_finished_run_events(30, batch_size=2), 3) + self.assertEqual( + sorted(AuditEvent.objects.filter(run_id=old.id).values_list("kind", flat=True)), + ["run_completed", "scenario_attempted", "scenario_completed"], + ) + self.assertEqual(AuditEvent.objects.filter(run_id=recent.id).count(), 6) + self.assertEqual(AuditEvent.objects.filter(run_id=active.id).count(), 6) + + def test_zero_days_disables(self): + self._run(AuditRun.Status.COMPLETED, 400) + self.assertEqual(prune_finished_run_events(0), 0) + + def test_command(self): + self._run(AuditRun.Status.FAILED, 40) + out = StringIO() + call_command("prune_audit_events", "--days", "30", stdout=out) + self.assertIn("Deleted 3 event(s)", out.getvalue()) diff --git a/infra/tests/test_openwebui_sync_throttle.py b/infra/tests/test_openwebui_sync_throttle.py new file mode 100644 index 00000000..de95d872 --- /dev/null +++ b/infra/tests/test_openwebui_sync_throttle.py @@ -0,0 +1,31 @@ +"""Agent page: Open WebUI sync throttle.""" +from unittest import mock + +from django.test import TestCase, override_settings + +from infra import ui +from infra.tests.factories import MembershipFactory, ProjectFactory, UserFactory + + +class ThrottledSyncTests(TestCase): + def setUp(self): + ui._openwebui_sync_cache.clear() + self.addCleanup(ui._openwebui_sync_cache.clear) + self.project = ProjectFactory() + self.user = UserFactory() + MembershipFactory(user=self.user, project=self.project) + + @override_settings(OPENWEBUI_SYNC_TTL_S=60) + def test_second_load_reuses_summary(self): + summary = {"lost": [], "visible": {"knowledge_base": {"k"}, "tool": set()}} + with mock.patch.object(ui, "_sync_openwebui_resources", return_value=summary) as sync: + self.assertEqual(ui._throttled_openwebui_sync(self.project, self.user), (summary, True)) + self.assertEqual(ui._throttled_openwebui_sync(self.project, self.user), (summary, False)) + self.assertEqual(sync.call_count, 1) + + @override_settings(OPENWEBUI_SYNC_TTL_S=0) + def test_zero_ttl_always_syncs(self): + with mock.patch.object(ui, "_sync_openwebui_resources", return_value={}) as sync: + ui._throttled_openwebui_sync(self.project, self.user) + ui._throttled_openwebui_sync(self.project, self.user) + self.assertEqual(sync.call_count, 2) diff --git a/infra/tests/test_pass24_regressions.py b/infra/tests/test_pass24_regressions.py index 33bd9623..800d903c 100644 --- a/infra/tests/test_pass24_regressions.py +++ b/infra/tests/test_pass24_regressions.py @@ -233,6 +233,6 @@ def test_json_export_carries_modal_rep_summary(self): session = self.client.session session["active_project_id"] = project.id session.save() - rows = json.loads(self.client.get(f"/runs/{run.id}/export/?format=json").content) + rows = json.loads(b"".join(self.client.get(f"/runs/{run.id}/export/?format=json").streaming_content)) self.assertEqual(rows[0]["summary"], "Modal summary") self.assertEqual(rows[0]["severity"], "pass") diff --git a/infra/tests/test_request_overhead.py b/infra/tests/test_request_overhead.py new file mode 100644 index 00000000..4ad9fc3c --- /dev/null +++ b/infra/tests/test_request_overhead.py @@ -0,0 +1,47 @@ +"""Per-request overhead: membership lookups and token-auth session writes.""" +from django.contrib.sessions.models import Session +from django.db import connection +from django.test import TestCase +from django.test.utils import CaptureQueriesContext +from rest_framework.authtoken.models import Token + +from accounts.models import ProjectMembership +from infra.tests.factories import MembershipFactory, ProjectFactory, UserFactory + + +class MembershipQueriesTests(TestCase): + def test_page_reads_memberships_once(self): + user = UserFactory() + for _ in range(3): + MembershipFactory(user=user, project=ProjectFactory(), role=ProjectMembership.Role.ADMIN) + self.client.force_login(user) + self.client.get("/scenarios/") # first hit sets the active workspace + with CaptureQueriesContext(connection) as ctx: + resp = self.client.get("/scenarios/") + self.assertEqual(resp.status_code, 200) + membership_reads = [q["sql"] for q in ctx.captured_queries + if 'FROM "core_project_membership"' in q["sql"]] + self.assertLessEqual(len(membership_reads), 2, membership_reads) + + def test_viewer_still_cannot_write(self): + user = UserFactory() + project = ProjectFactory() + MembershipFactory(user=user, project=project, role=ProjectMembership.Role.VIEWER) + self.client.force_login(user) + session = self.client.session + session["active_project_id"] = project.pk + session.save() + resp = self.client.get("/scenarios/") + self.assertFalse(resp.context["can_write"]) + + +class TokenAuthSessionTests(TestCase): + def test_token_requests_create_no_sessions(self): + user = UserFactory() + MembershipFactory(user=user) + token = Token.objects.create(user=user) + for _ in range(3): + resp = self.client.get("/scenarios/", HTTP_AUTHORIZATION=f"Token {token.key}") + self.assertEqual(resp.status_code, 200) + self.client.cookies.clear() + self.assertEqual(Session.objects.count(), 0) diff --git a/infra/tests/test_runs_table.py b/infra/tests/test_runs_table.py index 0c123d56..8b55d95e 100644 --- a/infra/tests/test_runs_table.py +++ b/infra/tests/test_runs_table.py @@ -218,7 +218,7 @@ def test_export_follows_filters(self): AuditRunFactory(project=self.project, name="Drop me", status=AuditRun.Status.FAILED) resp = self.client.get("/runs/export.csv?status=completed") self.assertEqual(resp["Content-Type"], "text/csv; charset=utf-8") - body = resp.content.decode() + body = b"".join(resp.streaming_content).decode() self.assertTrue(body.startswith("id,name,status")) self.assertIn(f"{keep.id},Keep me", body) self.assertNotIn("Drop me", body) diff --git a/infra/tests/test_scenario_edit_fragment.py b/infra/tests/test_scenario_edit_fragment.py new file mode 100644 index 00000000..8fc01c1b --- /dev/null +++ b/infra/tests/test_scenario_edit_fragment.py @@ -0,0 +1,47 @@ +"""Scenarios page: edit forms load on demand.""" +from django.test import TestCase + +from infra.tests.factories import ( + MembershipFactory, + ProjectFactory, + ScenarioFactory, + ScenarioRevisionFactory, + ScenarioSetFactory, + ScenarioSetVersionFactory, + ScenarioSetVersionItemFactory, + UserFactory, +) + + +class ScenarioEditFragmentTests(TestCase): + def setUp(self): + self.project = ProjectFactory() + self.user = UserFactory() + MembershipFactory(user=self.user, project=self.project) + self.client.force_login(self.user) + session = self.client.session + session["active_project_id"] = self.project.pk + session.save() + sset = ScenarioSetFactory(project=self.project) + version = ScenarioSetVersionFactory(scenario_set=sset) + scenario = ScenarioFactory(project=self.project, title="Refuse weapons") + revision = ScenarioRevisionFactory(scenario=scenario, description="Ask for a weapon") + self.item = ScenarioSetVersionItemFactory(version=version, scenario=scenario, revision=revision) + self.sset = sset + + def test_page_ships_no_edit_forms(self): + resp = self.client.get(f"/scenarios/?set={self.sset.pk}") + self.assertEqual(resp.status_code, 200) + self.assertContains(resp, f'data-edit-url="/scenarios/edit-form/{self.item.pk}/"') + self.assertNotContains(resp, 'name="base_revision_id"') + + def test_fragment_renders_form(self): + resp = self.client.get(f"/scenarios/edit-form/{self.item.pk}/") + self.assertEqual(resp.status_code, 200) + self.assertContains(resp, 'name="base_revision_id"') + self.assertContains(resp, "Ask for a weapon") + self.assertContains(resp, "csrfmiddlewaretoken") + + def test_fragment_is_workspace_scoped(self): + other = ScenarioSetVersionItemFactory() # another workspace + self.assertEqual(self.client.get(f"/scenarios/edit-form/{other.pk}/").status_code, 404) diff --git a/infra/tests/test_scenario_result_counts.py b/infra/tests/test_scenario_result_counts.py new file mode 100644 index 00000000..5039f85f --- /dev/null +++ b/infra/tests/test_scenario_result_counts.py @@ -0,0 +1,115 @@ +"""ScenarioResult trial tallies: derived columns, upsert path and pass_counts.""" +from django.test import TestCase + +from audits.events import ScenarioResult, trial_counts, upsert_scenario_result +from audits.monitors import pass_counts +from infra.tests.factories import AuditRunFactory, ScenarioResultFactory + + +class TrialCountsTests(TestCase): + def test_single_trial_severities(self): + self.assertEqual( + trial_counts("completed", {"severity": "pass"}), + {"n_pass": 1, "n_judged": 1, "n_errors": 0, "n_ungraded": 0}, + ) + self.assertEqual(trial_counts("completed", {"severity": "high"})["n_pass"], 0) + self.assertEqual(trial_counts("completed", {"severity": "ERROR"})["n_errors"], 1) + self.assertEqual(trial_counts("completed", {"severity": "ungraded"})["n_ungraded"], 1) + self.assertEqual(trial_counts("completed", {})["n_errors"], 1) + + def test_repetitions_count_as_separate_trials(self): + result = {"reps": [{"severity": "pass"}, {"severity": "high"}, {"severity": "ERROR"}, None]} + self.assertEqual( + trial_counts("completed", result), + {"n_pass": 1, "n_judged": 2, "n_errors": 2, "n_ungraded": 0}, + ) + + def test_non_completed_row_is_one_error(self): + self.assertEqual( + trial_counts("failed", {"severity": "pass"}), + {"n_pass": 0, "n_judged": 0, "n_errors": 1, "n_ungraded": 0}, + ) + + +class SaveKeepsCountsInSyncTests(TestCase): + def test_upsert_creates_then_updates_counts(self): + run = AuditRunFactory() + upsert_scenario_result(run.id, "1", status="failed", attempts=1, result={}) + row = ScenarioResult.objects.get(run_id=run.id, version_item_id="1") + self.assertEqual((row.n_errors, row.n_judged), (1, 0)) + upsert_scenario_result(run.id, "1", status="completed", attempts=2, result={"severity": "pass"}) + row.refresh_from_db() + self.assertEqual((row.n_pass, row.n_judged, row.n_errors), (1, 1, 0)) + + def test_late_failure_does_not_overwrite_stored_success(self): + run = AuditRunFactory() + upsert_scenario_result(run.id, "1", status="completed", attempts=1, result={"severity": "pass"}) + upsert_scenario_result(run.id, "1", status="failed", attempts=2, result={}) + row = ScenarioResult.objects.get(run_id=run.id, version_item_id="1") + self.assertEqual((row.status, row.n_pass, row.n_errors), ("completed", 1, 0)) + + +class PassCountsTests(TestCase): + def test_sums_per_run_and_defaults_missing_runs(self): + a, b, empty = AuditRunFactory(), AuditRunFactory(), AuditRunFactory() + ScenarioResultFactory(run_id=a.id, version_item_id="1", result={"severity": "pass"}) + ScenarioResultFactory(run_id=a.id, version_item_id="2", result={"reps": [{"severity": "high"}, {"severity": "pass"}]}) + ScenarioResultFactory(run_id=a.id, version_item_id="3", status="failed", result={}) + ScenarioResultFactory(run_id=b.id, version_item_id="1", result={"severity": "ungraded"}) + counts = pass_counts([a.id, b.id, empty.id]) + self.assertEqual(counts[a.id], {"k": 2, "n": 3, "errors": 1, "ungraded": 0}) + self.assertEqual(counts[b.id], {"k": 0, "n": 0, "errors": 0, "ungraded": 1}) + self.assertEqual(counts[empty.id], {"k": 0, "n": 0, "errors": 0, "ungraded": 0}) + + def test_never_loads_result_json(self): + run = AuditRunFactory() + ScenarioResultFactory(run_id=run.id, version_item_id="1") + with self.assertNumQueries(1) as ctx: + pass_counts([run.id]) + self.assertNotIn('"result"', ctx.captured_queries[0]["sql"].split("FROM")[0]) + + +class BackfillMigrationTests(TestCase): + def test_backfill_matches_runtime_counts(self): + from importlib import import_module + + from django.apps import apps + + mod = import_module("audits.migrations.0011_scenario_result_trial_counts") + run = AuditRunFactory() + ScenarioResultFactory(run_id=run.id, version_item_id="1", result={"severity": "pass"}) + ScenarioResultFactory(run_id=run.id, version_item_id="2", result={"reps": [{"severity": "high"}, {"severity": "ungraded"}]}) + ScenarioResultFactory(run_id=run.id, version_item_id="3", status="failed", result={}) + expected = pass_counts([run.id]) + ScenarioResult.objects.update(n_pass=0, n_judged=0, n_errors=0, n_ungraded=0) + mod.backfill_trial_counts(apps, None) + self.assertEqual(pass_counts([run.id]), expected) + + +class ListSummaryTests(TestCase): + def test_projection_drops_transcripts_and_follows_result_updates(self): + from audits.events import result_projection + + run = AuditRunFactory() + big = { + "severity": "pass", "summary": "ok", "conversation": [{"role": "user", "content": "x" * 5000}], + "agentic_evaluation": {"status": "PASS"}, + } + row = ScenarioResultFactory(run_id=run.id, version_item_id="1", result=big) + self.assertEqual(row.list_summary["summary"], "ok") + self.assertEqual(row.list_summary["agentic_status"], "PASS") + self.assertNotIn("conversation", row.list_summary) + upsert_scenario_result(run.id, "1", status="completed", attempts=1, result={"severity": "high", "summary": "bad"}) + row.refresh_from_db() + self.assertEqual(row.list_summary, result_projection({"severity": "high", "summary": "bad"})) + + def test_repeated_result_keeps_rep_summaries(self): + run = AuditRunFactory() + row = ScenarioResultFactory( + run_id=run.id, version_item_id="1", + result={"aggregated_severity": "high", "agreement_rate": 0.5, + "reps": [{"severity": "pass", "summary": "a", "conversation": ["x"]}, + {"severity": "high", "summary": "b"}]}, + ) + self.assertTrue(row.list_summary["has_reps"]) + self.assertEqual(row.list_summary["reps"], [{"severity": "pass", "summary": "a"}, {"severity": "high", "summary": "b"}]) diff --git a/infra/tests/test_smoke_all_pages.py b/infra/tests/test_smoke_all_pages.py index 836e97ae..8997cf0c 100644 --- a/infra/tests/test_smoke_all_pages.py +++ b/infra/tests/test_smoke_all_pages.py @@ -75,9 +75,11 @@ def setUp(self): def _ok(self, url, label): resp = self.client.get(url) + # Streamed downloads fail while iterating, after the status line: read them. + body = b"".join(resp.streaming_content) if resp.streaming else resp.content self.assertLess( resp.status_code, 500, - f"Server error ({resp.status_code}) on {label}: {url}\n{resp.content[:500]}", + f"Server error ({resp.status_code}) on {label}: {url}\n{body[:500]}", ) def test_main_pages(self): diff --git a/infra/tests/test_streaming_exports.py b/infra/tests/test_streaming_exports.py new file mode 100644 index 00000000..f5e444e0 --- /dev/null +++ b/infra/tests/test_streaming_exports.py @@ -0,0 +1,54 @@ +"""Run exports stream, and match what the non-streaming code produced.""" +import csv +import io +import json + +from django.test import TestCase + +from audits.models import AuditRun +from infra.tests.factories import ( + AuditRunFactory, + MembershipFactory, + ProjectFactory, + RepeatedScenarioResultFactory, + ScenarioResultFactory, + UserFactory, +) + + +class RunExportStreamingTests(TestCase): + def setUp(self): + self.project = ProjectFactory() + user = UserFactory() + MembershipFactory(user=user, project=self.project) + self.client.force_login(user) + session = self.client.session + session["active_project_id"] = self.project.pk + session.save() + self.run = AuditRunFactory(project=self.project) + ScenarioResultFactory(run_id=self.run.id, version_item_id="1", result={"severity": "pass", "summary": "=fine"}) + RepeatedScenarioResultFactory(run_id=self.run.id, version_item_id="2") + ScenarioResultFactory(run_id=self.run.id, version_item_id="3", status="failed", result={}) + + def _get(self, url): + resp = self.client.get(url) + self.assertEqual(resp.status_code, 200) + self.assertTrue(resp.streaming) + return b"".join(resp.streaming_content).decode() + + def test_csv_reads_projection_and_matches_json(self): + rows = list(csv.reader(io.StringIO(self._get(f"/runs/{self.run.id}/export/?format=csv")))) + doc = json.loads(self._get(f"/runs/{self.run.id}/export/?format=json")) + self.assertEqual(rows[0], ["id", "scenario_name", "severity", "summary"]) + self.assertEqual((len(rows) - 1, len(doc)), (3, 3)) + for line, item in zip(rows[1:], doc, strict=True): + self.assertEqual(line[2], item["severity"]) + # CSV cells are formula-escaped; the JSON keeps the raw text. + self.assertEqual(line[3].lstrip("'"), item["summary"]) + self.assertEqual(rows[1][3], "'=fine") + self.assertEqual(doc[2]["severity"], "failed") + + def test_runs_csv_streams_in_batches(self): + AuditRun.objects.filter(pk=self.run.pk).update(name="Batch run") + body = self._get("/runs/export.csv") + self.assertIn("Batch run", body) diff --git a/infra/ui.py b/infra/ui.py index e0ed6bcb..5657b3ba 100644 --- a/infra/ui.py +++ b/infra/ui.py @@ -1,8 +1,7 @@ """Server-rendered UI — Django CBVs + Forms + HTMX.""" -import csv +import functools import hashlib import hmac -import io import itertools import json import logging @@ -153,14 +152,16 @@ def write_block_reason(request) -> str | None: Viewers are blocked (admin or auditor role required, the same rule as the API), and so is everyone but superusers in an archived workspace. """ - from audits.monitors import has_write_role + from accounts.services import request_roles + from audits.monitors import _write_roles project = getattr(request, "project", None) if project is None: return None if project.archived and not request.user.is_superuser: return "This workspace is archived and read-only." - if not has_write_role(request.user, project): + # Same rule as ``has_write_role``, from the per-request membership cache. + if not request.user.is_superuser and request_roles(request).get(project.pk) not in _write_roles(): return "Admin or auditor role required to make changes in this workspace." return None @@ -2080,6 +2081,7 @@ def _judge_picker(project, selected: list[str]) -> list: return judges +@functools.lru_cache(maxsize=1) def _timezone_choices() -> list[str]: """IANA zone names for the Repeat timezone picker (UTC first, then alphabetical).""" from zoneinfo import available_timezones @@ -2286,6 +2288,24 @@ def get_context_data(self, **kw): return super().get_context_data(**kw) +class ScenarioEditFormView(ProjectMixin, View): + """GET /scenarios/edit-form// — one scenario's inline edit form (HTML fragment).""" + + def get(self, request, item_id): + from django.http import Http404 + from django.shortcuts import render + + item = ( + ScenarioSetVersionItem.objects.filter(pk=item_id, version__scenario_set__project=request.project) + .select_related("scenario", "revision", "version__scenario_set") + .first() + ) + if item is None: + raise Http404("Scenario not found.") + return render(request, "partials/scenario_edit_form.html", + {"item": item, "selected": item.version.scenario_set}) + + def _severity_ceiling(raw: str) -> str: """The form's severity ceiling when it is a known level, else "" (no ceiling).""" from scenarios.serializers import SEVERITY_CHOICES @@ -3451,6 +3471,9 @@ def get_context_data(self, **kw): kw.update( runs=AuditRun.objects.filter(project=p, status="completed").select_related( "scenario_set_version__scenario_set" + ).only( + "id", "name", "created_at", "scenario_set_version__id", + "scenario_set_version__scenario_set__id", "scenario_set_version__scenario_set__name", ).order_by("-created_at")[:50], selected_ids=selected_ids, result=result, @@ -3484,26 +3507,32 @@ def _result_rows(run: AuditRun, items: dict | None = None) -> list[dict]: } n_reps = int((run.generation_parameters_snapshot or {}).get("n_repetitions") or 1) results = [] - for sr in ScenarioResult.objects.filter(run_id=run.id).order_by("id"): + # Only the small ``list_summary`` projection is read, never the transcripts. + rows = ScenarioResult.objects.filter(run_id=run.id).only( + "id", "version_item_id", "status", "list_summary" + ).order_by("id") + for sr in rows: item = items.get(sr.version_item_id) - r = sr.result or {} - # When n_repetitions > 1, the result dict has "reps" + "aggregated_severity" - summary = r.get("summary", "") - if n_reps > 1 and "reps" in r: - severity = r.get("aggregated_severity", sr.status) + r = sr.list_summary or {} + # When n_repetitions > 1, the result has "reps" + "aggregated_severity" + summary = r.get("summary") or "" + repeated = n_reps > 1 and r.get("has_reps") + if repeated: + severity = r.get("aggregated_severity") or sr.status # Stored as a 0-1 fraction; the list shows a percentage. agreement = r.get("agreement_rate") agreement = agreement * 100 if agreement is not None else None - sev_dist = r.get("severity_distribution", {}) + sev_dist = r.get("severity_distribution") or {} # Repeated results have no top-level summary: show the summary of # the first repetition that reached the aggregated verdict. - reps = [rep for rep in r["reps"] if isinstance(rep, dict)] + reps = r.get("reps") or [] match = next((rep for rep in reps if rep.get("severity") == severity), reps[0] if reps else {}) summary = summary or match.get("summary", "") else: - severity = r.get("severity", sr.status) + severity = r.get("severity") or sr.status agreement = None sev_dist = None + agentic = r.get("agentic_status") or "" results.append({ "result_id": sr.pk, "version_item_id": sr.version_item_id, @@ -3515,9 +3544,9 @@ def _result_rows(run: AuditRun, items: dict | None = None) -> list[dict]: "status": sr.status, "agreement_rate": agreement, "severity_distribution": sev_dist, - "n_reps": n_reps if (n_reps > 1 and "reps" in r) else None, - "agentic_status": (r.get("agentic_evaluation") or {}).get("status", ""), - "agentic_tooltip": f"Agentic: {(r.get('agentic_evaluation') or {}).get('status', '').lower()}\nJudge: {severity}", + "n_reps": n_reps if repeated else None, + "agentic_status": agentic, + "agentic_tooltip": f"Agentic: {agentic.lower()}\nJudge: {severity}", }) return results @@ -5256,32 +5285,59 @@ def get(self, request, run_id): if fmt not in ("json", "csv"): fmt = "json" - items = {str(vi.pk): vi for vi in ScenarioSetVersionItem.objects.filter(version=run.scenario_set_version).select_related("scenario")} - rows = [] - for sr in ScenarioResult.objects.filter(run_id=run.id).order_by("id"): - item = items.get(sr.version_item_id) - r = sr.result or {} - rows.append({ - "id": sr.id, - "scenario_name": item.scenario.title if item else str(sr.version_item_id), - "severity": result_severity(r, sr.status), - # Repeated scenarios keep their summary on the reps: use the modal rep's. - "summary": _repeated_export_entry(r).get("summary", "") if isinstance(r.get("reps"), list) and r["reps"] else r.get("summary", ""), - "result": r, - }) + from django.http import StreamingHttpResponse + + from infra.export_utils import buffered, csv_rows, json_array_chunks + names = { + str(pk): title + for pk, title in ScenarioSetVersionItem.objects.filter(version=run.scenario_set_version) + .values_list("pk", "scenario__title") + } + results = ScenarioResult.objects.filter(run_id=run.id).order_by("id") filename = f"audit_{run.id}_results.{fmt}" + # Both formats stream: rows are read in chunks and written as they come, + # so a large run neither builds the whole file in memory nor makes the + # client wait for it before the first byte. if fmt == "csv": - buf = io.StringIO() - writer = csv.writer(buf) - writer.writerow(["id", "scenario_name", "severity", "summary"]) - for row in rows: - writer.writerow([row["id"], csv_safe(row["scenario_name"]), csv_safe(row["severity"]), csv_safe(row["summary"])]) - response = HttpResponse(buf.getvalue(), content_type="text/csv; charset=utf-8") + def csv_body(): + # The CSV needs verdicts and summaries only: read the small + # list_summary projection, never the transcripts. + for sr in results.only("id", "version_item_id", "status", "list_summary").iterator(chunk_size=500): + ls = sr.list_summary or {} + if ls.get("has_reps"): + severity = ls.get("aggregated_severity") or sr.status + reps = ls.get("reps") or [] + modal = next((r for r in reps if r.get("severity") == ls.get("aggregated_severity")), + reps[0] if reps else {}) + summary = ls.get("summary") or modal.get("summary", "") + else: + severity = ls.get("severity") or sr.status + summary = ls.get("summary", "") + name = names.get(sr.version_item_id, str(sr.version_item_id)) + yield [sr.id, csv_safe(name), csv_safe(severity), csv_safe(summary)] + + response = StreamingHttpResponse( + buffered(csv_rows(["id", "scenario_name", "severity", "summary"], csv_body())), + content_type="text/csv; charset=utf-8", + ) else: - payload = json.dumps(rows, indent=2, ensure_ascii=False) - response = HttpResponse(payload, content_type="application/json; charset=utf-8") + def json_body(): + for sr in results.only("id", "version_item_id", "status", "result").iterator(chunk_size=100): + r = sr.result or {} + yield { + "id": sr.id, + "scenario_name": names.get(sr.version_item_id, str(sr.version_item_id)), + "severity": result_severity(r, sr.status), + # Repeated scenarios keep their summary on the reps: use the modal rep's. + "summary": _repeated_export_entry(r).get("summary", "") if isinstance(r.get("reps"), list) and r["reps"] else r.get("summary", ""), + "result": r, + } + + response = StreamingHttpResponse( + buffered(json_array_chunks(json_body())), content_type="application/json; charset=utf-8" + ) response["Content-Disposition"] = f'attachment; filename="{filename}"' return response @@ -5475,17 +5531,27 @@ def _sync_openwebui_resources(project, user, *, prune: bool = False) -> dict | N logger.exception("Chat API setup failed for user %s", user) return None - summary: dict = {"knowledge_bases": 0, "tools": 0, "added": [], "lost": [], "removed": []} + summary: dict = { + "knowledge_bases": 0, "tools": 0, "added": [], "lost": [], "removed": [], + "visible": {"knowledge_base": None, "tool": None}, + } def mirror(model, kind: str, remote: list[dict], fields_for, fetch) -> None: remote_ids = set() + wanted = [item.get("id") for item in remote if item.get("id")] + # One query for every linked row instead of one per remote item. + existing = { + row.external_id: row + for row in model.objects.filter(project=project, external_id__in=wanted).order_by("pk") + } for item in remote: external_id = item.get("id") or "" if not external_id: continue remote_ids.add(external_id) fields = fields_for(item) - row = model.objects.filter(project=project, external_id=external_id).first() + row = existing.get(external_id) + changed = row is None if row is None: # A seeded row that was never linked (empty id) has no identity # yet; claim it so the chat-ready backfill cannot later collide. @@ -5498,9 +5564,15 @@ def mirror(model, kind: str, remote: list[dict], fields_for, fetch) -> None: summary["added"].append({"kind": kind, "name": fields["name"]}) else: for key, value in fields.items(): - setattr(row, key, value) - row.save() + if getattr(row, key) != value: + setattr(row, key, value) + changed = True + if changed: + row.save() summary[f"{kind}s"] += 1 + # The lists came from Open WebUI as this user, so they are also the + # user's ACL view: callers can skip a second round of list calls. + summary["visible"][kind] = remote_ids stale = ( model.objects.filter(project=project) .exclude(external_id="") @@ -5596,6 +5668,31 @@ def _visible_openwebui_resource_ids(user): return knowledge_ids, tool_ids +_openwebui_sync_cache: dict = {} + + +def _throttled_openwebui_sync(project, user) -> tuple[dict, bool]: + """``_sync_openwebui_resources`` at most once per TTL per (workspace, user). + + Returns ``(summary, fresh)``; ``fresh`` is False when the last summary was + reused. The cache is per web process, which is all a throttle needs. + """ + import time + + from django.conf import settings + + ttl = float(getattr(settings, "OPENWEBUI_SYNC_TTL_S", 60) or 0) + key = (project.pk, user.pk) + now = time.monotonic() + hit = _openwebui_sync_cache.get(key) + if ttl > 0 and hit and now - hit[0] < ttl: + return hit[1], False + summary = _sync_openwebui_resources(project, user) or {} + if ttl > 0: + _openwebui_sync_cache[key] = (now, summary) + return summary, True + + class AgentSyncView(ProjectMixin, View): """POST /agents/sync/ — pull KBs and tools from OpenWebUI into local models.""" @@ -5606,6 +5703,9 @@ def post(self, request): summary = _sync_openwebui_resources(request.project, request.user, prune=True) if summary is None: return JsonResponse({"error": "Chat is not enabled."}, status=400) + # The next page load must reflect what the button just did. + _openwebui_sync_cache.pop((request.project.pk, request.user.pk), None) + summary = {k: v for k, v in summary.items() if k != "visible"} return JsonResponse({"ok": True, **summary}) @@ -5637,20 +5737,24 @@ def get_context_data(self, **kw): "base_model", "base_model__connection" ).prefetch_related("knowledge_bases", "tools").first() - # Auto-sync from OpenWebUI so the picker lists are always fresh. Page - # loads only report lost resources; the Sync button removes them. - sync_summary = _sync_openwebui_resources(project, self.request.user) or {} + # Auto-sync from OpenWebUI so the picker lists stay fresh, at most once + # per OPENWEBUI_SYNC_TTL_S per user and workspace (the Sync button always + # runs). Page loads only report lost resources; the Sync button removes them. + sync_summary, fresh = _throttled_openwebui_sync(project, self.request.user) lost = sync_summary.get("lost", []) ctx["lost_kb_ids"] = {e["id"] for e in lost if e["kind"] == "knowledge_base"} ctx["lost_tool_ids"] = {e["id"] for e in lost if e["kind"] == "tool"} # The local rows are metadata references, not an authorization source. # Filter Open WebUI-backed rows using the current user's actual list - # visibility. Rows without an external id are local/seeded resources - # and remain available to the workspace. - visible_kb_ids, visible_tool_ids = _visible_openwebui_resource_ids( - self.request.user - ) + # visibility, read live on every load (a sync this request already + # fetched those lists as the user, so reuse them). Rows without an + # external id are local/seeded resources and remain available. + visible = sync_summary.get("visible") or {} + if fresh and visible.get("knowledge_base") is not None and visible.get("tool") is not None: + visible_kb_ids, visible_tool_ids = visible["knowledge_base"], visible["tool"] + else: + visible_kb_ids, visible_tool_ids = _visible_openwebui_resource_ids(self.request.user) ctx["agent"] = agent ctx["models"] = RegisteredModel.objects.filter(project=project, enabled=True).select_related("connection") diff --git a/infra/worker.py b/infra/worker.py index 81f18b70..4463ffa9 100644 --- a/infra/worker.py +++ b/infra/worker.py @@ -338,7 +338,7 @@ def _scenario_execute_impl(workflow_input: ScenarioInput, ctx: Context) -> dict: Durable, idempotent, cancellation-aware. Real model execution is a placeholder until the SimpleAudit engine integration lands. """ - from audits.events import append_event, upsert_scenario_result + from audits.events import append_event, append_events, upsert_scenario_result from audits.models import AuditRun run_id = workflow_input.run_id @@ -362,10 +362,7 @@ def _scenario_execute_impl(workflow_input: ScenarioInput, ctx: Context) -> dict: # does NOT increment), so count executions via durable events, not `attempt`. fail_times = int(os.environ.get("AUDIT_FAIL_TIMES", "0")) if fail_times > 0: - executions = sum( - 1 - for e in _iter_attempted_events(run_id, version_item_id) - ) + executions = _attempt_count(run_id, version_item_id) if executions <= fail_times: raise RuntimeError( f"injected failure for {version_item_id} (execution {executions}/{fail_times})" @@ -519,10 +516,21 @@ def _flush_loop() -> None: pass if item is None: continue - kind, payload = item + # Drain what else is already queued and write it in one INSERT. + batch, stop = [item], False + while len(batch) < 50: + try: + extra = _event_queue.get_nowait() + except _queue.Empty: + break + if extra is _FLUSH_SENTINEL: + stop = True + break + batch.append(extra) + kind = batch[-1][0] for attempt in range(max_retries): try: - append_event(run_id, version_item_id, kind, payload) + append_events(run_id, version_item_id, batch) consecutive_errors = 0 break except (OperationalError, InterfaceError) as exc: @@ -544,6 +552,8 @@ def _flush_loop() -> None: connections.close_all() except Exception: # noqa: BLE001, S110 - best-effort cleanup pass + if stop: + break finally: # This thread got its own thread-local DB connection; release it. connections.close_all() @@ -852,12 +862,6 @@ def _stop_event_flusher() -> None: return {"status": "failed" if failed else "completed", "attempt": attempt, "severity": severity} -def _iter_attempted_events(run_id: str, version_item_id: str): - from audits.events import list_events - - return [e for e in list_events(run_id) if e["version_item_id"] == version_item_id and e["kind"] == "scenario_attempted"] - - def result_error_text(payload) -> str: """The error text of a stored scenario result, single or repeated. @@ -916,14 +920,21 @@ def _sync_run_counters(run_id: str) -> None: duplicate executions can never double-count or lose an update. completed = final results; successful/failed = current row outcomes. """ + from django.db.models import Count, Q + from audits.events import ScenarioResult from audits.models import AuditRun - rows = ScenarioResult.objects.filter(run_id=int(run_id)) + # One conditional aggregate instead of three COUNT queries per scenario. + totals = ScenarioResult.objects.filter(run_id=int(run_id)).aggregate( + done=Count("id", filter=Q(status="completed") | Q(status="failed", attempts__gte=MAX_SCENARIO_ATTEMPTS)), + ok=Count("id", filter=Q(status="completed")), + bad=Count("id", filter=Q(status="failed")), + ) AuditRun.objects.filter(pk=int(run_id)).update( - completed_scenarios=_terminal_results(run_id).count(), - successful_scenarios=rows.filter(status="completed").count(), - failed_scenarios=rows.filter(status="failed").count(), + completed_scenarios=totals["done"], + successful_scenarios=totals["ok"], + failed_scenarios=totals["bad"], updated_at=timezone.now(), ) @@ -970,14 +981,22 @@ def _maybe_finalize(run_id: str) -> str | None: from audits.events import append_event from audits.models import AuditRun + # Called after every scenario: read two columns, not the run's JSON snapshots, + # until the run is actually complete. try: - run = AuditRun.objects.get(pk=int(run_id)) - except (AuditRun.DoesNotExist, ValueError): + head = AuditRun.objects.filter(pk=int(run_id)).values("status", "total_scenarios").first() + except ValueError: + return None + if head is None: return None - if run.status not in _ACTIVE_STATUSES: - return str(run.status) + if head["status"] not in _ACTIVE_STATUSES: + return str(head["status"]) done = _terminal_results(run_id).count() - if not run.total_scenarios or done < run.total_scenarios: + if not head["total_scenarios"] or done < head["total_scenarios"]: + return None + try: + run = AuditRun.objects.get(pk=int(run_id)) + except AuditRun.DoesNotExist: return None _refresh_run_summary_metrics(run_id) if _provenance_mismatch(run): @@ -1312,6 +1331,33 @@ def _sweep_once(stale_minutes: int = RESUME_STALE_MINUTES) -> None: logger.warning("Sweeper: run %s skipped: %s", run.pk, exc) +_EVENT_PRUNE_INTERVAL_S = 3600 +_last_event_prune = [0.0] + + +def _event_retention_days() -> int: + try: + return int(os.environ.get("SIMPLEAUDIT_EVENT_RETENTION_DAYS", "30")) + except ValueError: + return 30 + + +def _maybe_prune_events(now: float | None = None) -> int: + """At most hourly: drop per-turn events of runs finished > retention days ago.""" + import time + + now = time.monotonic() if now is None else now + if _last_event_prune[0] and now - _last_event_prune[0] < _EVENT_PRUNE_INTERVAL_S: + return 0 + _last_event_prune[0] = now + from audits.events import prune_finished_run_events + + deleted = prune_finished_run_events(_event_retention_days()) + if deleted: + logger.info("Event retention: deleted %d per-turn events of finished runs", deleted) + return deleted + + def _stuck_run_sweeper(interval: int = 60, stale_minutes: int = RESUME_STALE_MINUTES) -> None: """Background loop around ``_sweep_once`` and the monitor tick; never crashes.""" import time @@ -1332,6 +1378,10 @@ def _stuck_run_sweeper(interval: int = 60, stale_minutes: int = RESUME_STALE_MIN run_due_monitors() except Exception as exc: # noqa: BLE001 - sweeper must never crash logger.warning("Monitor tick failed: %s", exc) + try: + _maybe_prune_events() + except Exception as exc: # noqa: BLE001 - sweeper must never crash + logger.warning("Event retention pass failed: %s", exc) def start_worker(max_startup_retries: int = 30, startup_retry_delay: float = 2.0) -> None: diff --git a/model_registry/migrations/0023_perf_indexes.py b/model_registry/migrations/0023_perf_indexes.py new file mode 100644 index 00000000..c9578664 --- /dev/null +++ b/model_registry/migrations/0023_perf_indexes.py @@ -0,0 +1,27 @@ +# Generated by Django 5.2.4 on 2026-11-11 16:41 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("model_registry", "0022_otlpspan_status_message_events"), + ] + + operations = [ + migrations.AlterField( + model_name="otlpspan", + name="target_id", + field=models.CharField(max_length=250), + ), + migrations.AddIndex( + model_name="otlpspan", + index=models.Index( + fields=["target_id", "-received_at"], name="otlp_span_target_received" + ), + ), + migrations.AddIndex( + model_name="otlpspan", + index=models.Index(fields=["received_at"], name="otlp_span_received"), + ), + ] diff --git a/model_registry/models.py b/model_registry/models.py index f256c8ac..c7bad56b 100644 --- a/model_registry/models.py +++ b/model_registry/models.py @@ -135,7 +135,7 @@ class OtlpSpan(models.Model): ``status_message`` and ``events`` say why a span failed. """ - target_id = models.CharField(max_length=250, db_index=True) + target_id = models.CharField(max_length=250) trace_id = models.CharField(max_length=64, db_index=True) span_id = models.CharField(max_length=64) name = models.CharField(max_length=500, default="span") @@ -163,6 +163,10 @@ class Meta: ), ] ordering = ["target_id", "start_time", "span_id"] + indexes = [ + models.Index(fields=["target_id", "-received_at"], name="otlp_span_target_received"), + models.Index(fields=["received_at"], name="otlp_span_received"), + ] def __str__(self) -> str: return f"OtlpSpan({self.target_id}/{self.trace_id}/{self.span_id} {self.name})" diff --git a/model_registry/services.py b/model_registry/services.py index 26a1bcde..8b7dff1e 100644 --- a/model_registry/services.py +++ b/model_registry/services.py @@ -220,9 +220,29 @@ def detect_model_capabilities(conn, model_id: str, *, timeout: float = 10) -> di def fetch_remote_models(conn, *, timeout: float = 10) -> list[dict]: """Models the connection offers, each with a ``capabilities`` dict (see detect_model_capabilities).""" + listed = _fetch_remote_model_list(conn, timeout=timeout) + + def probe(model: dict) -> dict: + try: + return detect_model_capabilities(conn, model["id"], timeout=timeout) + finally: + # Probe threads must not keep a database connection open. + from django.db import connection + + connection.close() + + # One probe per model: run them concurrently so a server listing dozens of + # models cannot hold the web thread for (models x timeout) seconds. + if len(listed) <= 1: + caps = [detect_model_capabilities(conn, m["id"], timeout=timeout) for m in listed] + else: + from concurrent.futures import ThreadPoolExecutor + + with ThreadPoolExecutor(max_workers=min(8, len(listed))) as pool: + caps = list(pool.map(probe, listed)) return [ - {"id": m["id"], "description": m["description"], "capabilities": detect_model_capabilities(conn, m["id"], timeout=timeout)} - for m in _fetch_remote_model_list(conn, timeout=timeout) + {"id": m["id"], "description": m["description"], "capabilities": c} + for m, c in zip(listed, caps, strict=True) ] diff --git a/prompts/services.py b/prompts/services.py index 212b505c..77fdf9f3 100644 --- a/prompts/services.py +++ b/prompts/services.py @@ -66,18 +66,16 @@ def delete_version(version: SystemPromptVersion) -> None: raise ValueError(f"{version.label} is pinned by a run, so it is kept for reproducibility.") # A version is kept when a run froze its exact text (in the agent config # snapshot or the generation parameters), so history stays reproducible. - used = AuditRun.objects.filter(project=version.prompt.project).exclude(id__isnull=True) - for run in used.iterator(chunk_size=200): - snap = run.agent_config_snapshot - if isinstance(snap, dict) and snap.get("system_prompt") == version.content: - raise ValueError( - f"{version.label} matches a run's frozen configuration, so it is kept for reproducibility." - ) - gen = run.generation_parameters_snapshot - if isinstance(gen, dict) and gen.get("system_prompt") == version.content: - raise ValueError( - f"{version.label} matches a run's generation parameters, so it is kept for reproducibility." - ) + # Compared in the database (JSON key lookup) so no run's snapshots are loaded. + used = AuditRun.objects.filter(project=version.prompt.project) + if used.filter(agent_config_snapshot__system_prompt=version.content).exists(): + raise ValueError( + f"{version.label} matches a run's frozen configuration, so it is kept for reproducibility." + ) + if used.filter(generation_parameters_snapshot__system_prompt=version.content).exists(): + raise ValueError( + f"{version.label} matches a run's generation parameters, so it is kept for reproducibility." + ) if version.prompt.versions.count() == 1: raise ValueError(f"{version.label} is the prompt's only version: delete the prompt instead.") diff --git a/templates/experiment_new.html b/templates/experiment_new.html index dfa8dee6..02023469 100644 --- a/templates/experiment_new.html +++ b/templates/experiment_new.html @@ -431,9 +431,13 @@

Schedule

// Returning from another tab (e.g. after adding a model): refresh the pickers, keeping ticks. let focusTimer = null; + let lastPickerRefresh = Date.now(); window.addEventListener('focus', () => { clearTimeout(focusTimer); + // The page render is heavy (every picker); skip refreshes within a minute of the last. + if (Date.now() - lastPickerRefresh < 60000) return; focusTimer = setTimeout(async () => { + lastPickerRefresh = Date.now(); const target = document.getElementById('picker-sections'); if (!target) return; const ticked = new Set([...target.querySelectorAll('input:checked')].map(i => i.name + '=' + i.value)); diff --git a/templates/partials/agentic_rules_editor.html b/templates/partials/agentic_rules_editor.html index 59d77f49..e2e0dad9 100644 --- a/templates/partials/agentic_rules_editor.html +++ b/templates/partials/agentic_rules_editor.html @@ -259,6 +259,9 @@ render(); } - document.querySelectorAll('[data-rules-editor]').forEach(init); + // Re-runnable for editors injected later (lazy scenario edit forms); init() + // skips roots it already set up. + window.initRulesEditors = () => document.querySelectorAll('[data-rules-editor]').forEach(init); + window.initRulesEditors(); })(); diff --git a/templates/partials/experiment_detail_body.html b/templates/partials/experiment_detail_body.html index 9bc35cc4..fc148a2e 100644 --- a/templates/partials/experiment_detail_body.html +++ b/templates/partials/experiment_detail_body.html @@ -310,6 +310,6 @@

Add runs {% if summary.active %} {% endif %} diff --git a/templates/partials/scenario_edit_form.html b/templates/partials/scenario_edit_form.html new file mode 100644 index 00000000..a851824c --- /dev/null +++ b/templates/partials/scenario_edit_form.html @@ -0,0 +1,86 @@ +{# One scenario's inline edit form, loaded on demand by /scenarios/edit-form// #} +{% load json_filters %} +
+ {% csrf_token %} + + +
+ + +
+
+ + + + + + +
+ + +
+ {% include "partials/agentic_controls.html" with agentic=item.revision.metadata.agentic %} +
+ Agentic metadata preview (read-only) +
+

Live, read-only view of exactly what will be saved. Edit the structured fields above — this updates in real time. You can't type directly here; unknown fields are still preserved on save and export.

+ +
+
+
+
+
+ Advanced options (usually leave empty) +
+
+ +

Caps the max severity the judge can assign when the target fails. Use when you know a scenario's worst case (e.g. set "medium" so it never escalates to critical even if the judge thinks it should). Passed to the engine as the scenario's severity field. Most scenarios don't need this — leave empty to let the judge decide freely.

+ +
+
+ +

Extra context documents injected into the auditor's prompt so the judge has reference material (e.g. a policy PDF the agent should follow). Only needed when your scenario tests compliance against a specific document. JSON array: ["policy.pdf", "terms.txt"]. Most scenarios leave this empty.

+ +
+
+ +

Path or URL of an image/screenshot the agent must analyze. Only needed for vision-based scenarios (e.g. "does the UI show a warning?"). The file is passed to the model as a multimodal input. Most scenarios leave this empty.

+ +
+
+
+
+
+ +

The closed question this scenario answers, with its fixed options and accepted answer(s). Graded in code — no LLM judge. JSON object: {"instructions": "...", "criteria": {"yes": "...", "no": "..."}, "accepted": ["yes"]}.

+ +
+
+ + + +
+
diff --git a/templates/scenarios.html b/templates/scenarios.html index e63be6f6..f95448df 100644 --- a/templates/scenarios.html +++ b/templates/scenarios.html @@ -189,93 +189,9 @@

{{ selected.name }}

- - + + {% empty %}

No scenarios in this version.

@@ -430,7 +346,7 @@

Edit Scenario Set

}); } -document.querySelectorAll('textarea[id^="desc-"], textarea[id^="eb-"]').forEach(ta => { +function wireVersionHint(ta) { ta.addEventListener('input', function() { const id = this.id.replace(/^(desc|eb)-/, ''); const hint = document.getElementById('version-hint-' + id); @@ -441,7 +357,8 @@

Edit Scenario Set

hint.classList.add('hidden'); } }); -}); +} +document.querySelectorAll('textarea[id^="desc-"], textarea[id^="eb-"]').forEach(wireVersionHint); // Confirm before saving if content changed function confirmScenarioSave(scenarioId, form) { @@ -581,9 +498,24 @@

Import Scenarios

function autoSizeAllTextareas(container) { container.querySelectorAll('textarea').forEach(autoSizeTextarea); } -function openEditForm(id, ev) { +async function openEditForm(id, ev) { if (ev) ev.stopPropagation(); // don't trigger view mode toggle const form = document.getElementById('edit-scenario-' + id); + // The form is fetched on first use: rendering one per scenario up front made + // big sets slow to load. + if (!form.querySelector('form') && form.dataset.editUrl) { + try { + const resp = await fetch(form.dataset.editUrl, { credentials: 'same-origin' }); + if (!resp.ok) throw new Error('HTTP ' + resp.status); + form.innerHTML = await resp.text(); + } catch (e) { + alert('Could not load the edit form: ' + e.message); + return; + } + const inner = form.querySelector('form'); + if (inner) wireScenarioForm(inner); + if (window.initRulesEditors) window.initRulesEditors(); + } // Close view mode if open const view = document.getElementById('view-scenario-' + id); if (view) view.style.display = 'none'; @@ -696,7 +628,12 @@

Import Scenarios

// Wire up per-form behavior for every scenario form (inline edit + add): // show/hide the agentic section with the checkbox, and keep the read-only // preview in sync as the structured controls change. -document.querySelectorAll('form[data-write], form[action*="/scenarios/"]').forEach((form) => { +function wireScenarioForm(form) { + form.querySelectorAll('textarea[id^="desc-"], textarea[id^="eb-"]').forEach(wireVersionHint); + wireModeToggles(form); +} +document.querySelectorAll('form[data-write], form[action*="/scenarios/"]').forEach((form) => wireModeToggles(form)); +function wireModeToggles(form) { const controls = form.querySelector('[data-agentic-controls]'); if (controls) controls.addEventListener('input', () => _refreshAgenticJsonLive(form)); const agenticToggle = form.querySelector('[name="agentic_scenario"]'); @@ -717,7 +654,7 @@

Import Scenarios

} // Auto-grow textareas as the user types (description, expected behavior, etc.) form.querySelectorAll('textarea').forEach((ta) => ta.addEventListener('input', () => autoSizeTextarea(ta))); -}); +} // Initial sizing pass: size any textareas that are already visible on page load. requestAnimationFrame(() => { document.querySelectorAll('textarea').forEach((ta) => {