From f75ff582de6a4235534648816da79731d9f5b958 Mon Sep 17 00:00:00 2001
From: JSONbored <49853598+JSONbored@users.noreply.github.com>
Date: Tue, 7 Jul 2026 00:21:29 -0700
Subject: [PATCH] feat(selfhost): add retention for agent_context_snapshots

Every table originally flagged as unbounded in #3896 except one turned
out already handled: audit_events/ai_usage_events/product_usage_events
already have cron-scheduled retention via RETENTION_POLICY (the
original audit's grep for literal .delete() calls missed the dynamic
DELETE FROM ${table} SQL this module uses), and webhook_events is
deliberately excluded per a prior fix (#666) since GitHub can redeliver
a webhook long after the original event and this table is the
idempotency record that detects that -- adding it back would
reintroduce that fixed bug.

agent_context_snapshots (one payloadJson blob per agent run, capped
only by a per-runId .limit(50) READ, never deleted) was genuinely the
only unaddressed table. Add it to RETENTION_POLICY with a 30-day
window, reusing the existing generic prune mechanism.

Closes #3896
---
 src/db/retention.ts         |  3 +++
 test/unit/retention.test.ts | 20 +++++++++++++++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/src/db/retention.ts b/src/db/retention.ts
index 459f8dc7b3..0dcf6293f0 100644
--- a/src/db/retention.ts
+++ b/src/db/retention.ts
@@ -21,6 +21,9 @@ export const RETENTION_POLICY: readonly RetentionRule[] = [
   { table: "signal_snapshots", column: "generated_at", days: 90 },
   { table: "score_previews", column: "generated_at", days: 90 },
   { table: "repo_snapshots", column: "fetched_at", days: 90 },
+  // One payloadJson blob per agent run (#3896); a per-run diagnostic snapshot with no cross-run rollup
+  // depending on it, so a shorter window than the audit/usage-log tables above is appropriate.
+  { table: "agent_context_snapshots", column: "created_at", days: 30 },
 ];
 
 export type PruneResult = { table: string; column: string; cutoff: string; deleted: number };
diff --git a/test/unit/retention.test.ts b/test/unit/retention.test.ts
index 187930c9ae..4d69a338a8 100644
--- a/test/unit/retention.test.ts
+++ b/test/unit/retention.test.ts
@@ -2,7 +2,7 @@ import { describe, expect, it } from "vitest";
 import { createApp } from "../../src/api/routes";
 import { getDb } from "../../src/db/client";
 import { dedupeSignalSnapshots, pruneExpiredRecords, RETENTION_POLICY } from "../../src/db/retention";
-import { aiUsageEvents, webhookEvents } from "../../src/db/schema";
+import { agentContextSnapshots, aiUsageEvents, webhookEvents } from "../../src/db/schema";
 import { processJob, runRetentionPrune } from "../../src/queue/processors";
 import { createTestEnv } from "../helpers/d1";
 
@@ -102,6 +102,24 @@ describe("pruneExpiredRecords", () => {
     await expect(pruneExpiredRecords(env, { policy: [{ table: "webhook_events; DROP TABLE x", column: "received_at", days: 1 }] })).rejects.toThrow("Unsafe retention identifier");
   });
 
+  it("prunes agent_context_snapshots older than its window and keeps recent runs (#3896)", async () => {
+    const env = createTestEnv();
+    const db = getDb(env.DB);
+    await db.insert(agentContextSnapshots).values([
+      { id: "ctx-old", runId: "run-old", createdAt: daysAgo(40) },
+      { id: "ctx-recent", runId: "run-recent", createdAt: daysAgo(2) },
+    ]);
+
+    const results = await pruneExpiredRecords(env, {
+      nowMs: NOW,
+      policy: [{ table: "agent_context_snapshots", column: "created_at", days: 30 }],
+    });
+
+    expect(results[0]?.deleted).toBe(1);
+    const rows = await env.DB.prepare("SELECT id FROM agent_context_snapshots").all<{ id: string }>();
+    expect(rows.results.map((row) => row.id)).toEqual(["ctx-recent"]);
+  });
+
   it("the policy only targets append-only/log/snapshot tables (no current-state tables)", () => {
     const tables = RETENTION_POLICY.map((r) => r.table);
     for (const protectedTable of ["webhook_events", "repositories", "repository_settings", "pull_requests", "issues", "repository_ai_keys", "contributors"]) {