@axiom-lattice/pg-stores 3.3.12 → 3.4.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/.turbo/turbo-build.log +10 -10
- package/CHANGELOG.md +187 -0
- package/dist/index.d.mts +353 -3
- package/dist/index.d.ts +353 -3
- package/dist/index.js +2323 -93
- package/dist/index.js.map +1 -1
- package/dist/index.mjs +2329 -92
- package/dist/index.mjs.map +1 -1
- package/package.json +3 -3
- package/src/PGVectorStoreProvider.ts +19 -13
- package/src/__tests__/AgentRunRetention.migrations.test.ts +41 -0
- package/src/__tests__/PostgreSQLAgentRunLedger.test.ts +305 -0
- package/src/__tests__/PostgreSQLAgentRunRetention.test.ts +126 -0
- package/src/__tests__/PostgreSQLAgentRunTaskQueue.test.ts +1375 -0
- package/src/__tests__/PostgreSQLChannelReplyJobStore.test.ts +333 -0
- package/src/__tests__/PostgreSQLCheckpointStorage.deleteAfter.test.ts +473 -0
- package/src/__tests__/PostgreSQLCheckpointStorage.test.ts +538 -0
- package/src/__tests__/PostgreSQLStore.test.ts +173 -0
- package/src/__tests__/PostgreSQLThreadStore.test.ts +80 -0
- package/src/__tests__/ThreadMessageQueueStore.integration.test.ts +21 -0
- package/src/__tests__/ThreadMessageQueueStore.test.ts +8 -3
- package/src/__tests__/createPgStoreConfig.test.ts +38 -0
- package/src/createPgStoreConfig.ts +36 -1
- package/src/index.ts +7 -0
- package/src/migrations/agent_run_base_checkpoint_captured_migration.ts +34 -0
- package/src/migrations/agent_run_capability_revision_migration.ts +24 -0
- package/src/migrations/agent_run_ledger_migration.ts +51 -0
- package/src/migrations/agent_run_project_task_migration.ts +34 -0
- package/src/migrations/agent_run_project_task_running_index_migration.ts +33 -0
- package/src/migrations/agent_run_retention_migration.ts +25 -0
- package/src/migrations/agent_run_rollback_failure_migration.ts +28 -0
- package/src/migrations/agent_run_running_index_source_migration.ts +31 -0
- package/src/migrations/agent_run_task_queue_migration.ts +65 -0
- package/src/migrations/channel_reply_job_migration.ts +55 -0
- package/src/migrations/model_provider_migrations.ts +27 -0
- package/src/migrations/store_migrations.ts +50 -0
- package/src/migrations/thread_migrations.ts +25 -0
- package/src/stores/PostgreSQLAgentRunLedger.ts +1210 -0
- package/src/stores/PostgreSQLChannelReplyJobStore.ts +373 -0
- package/src/stores/PostgreSQLCheckpointStorage.ts +548 -0
- package/src/stores/PostgreSQLModelProviderStore.ts +25 -7
- package/src/stores/PostgreSQLStore.ts +167 -0
- package/src/stores/PostgreSQLThreadStore.ts +49 -10
- package/src/stores/ThreadMessageQueueStore.ts +10 -5
|
@@ -0,0 +1,1210 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* PostgreSQL Agent Run Ledger
|
|
3
|
+
*
|
|
4
|
+
* Durable run lifecycle and task queue. `source='run_api'` rows are claimed by
|
|
5
|
+
* the run-task worker; `source='legacy'` rows preserve the existing
|
|
6
|
+
* observational ledger/message-queue behavior.
|
|
7
|
+
*/
|
|
8
|
+
|
|
9
|
+
import { Pool, type PoolConfig } from "pg";
|
|
10
|
+
import {
|
|
11
|
+
AgentRunAlreadyActiveError,
|
|
12
|
+
AgentRunDuplicateMessageIdError,
|
|
13
|
+
AgentRunInputCommandConflictError,
|
|
14
|
+
AgentRunInvalidTaskIdError,
|
|
15
|
+
AgentRunTenantQuotaExceededError,
|
|
16
|
+
assertQueueSources,
|
|
17
|
+
extractAgentRunInputMessageIds,
|
|
18
|
+
normalizeAgentRunMessageIds,
|
|
19
|
+
type AgentRunCounts,
|
|
20
|
+
type AgentRunExpiredOutcome,
|
|
21
|
+
type AgentRunJsonObject,
|
|
22
|
+
type AgentRunRecord,
|
|
23
|
+
type AgentRunRollbackResult,
|
|
24
|
+
type AgentRunSource,
|
|
25
|
+
type AgentRunStatus,
|
|
26
|
+
type BeginAgentRunInput,
|
|
27
|
+
type ClaimAgentRunInput,
|
|
28
|
+
type CreatePendingAgentRunInput,
|
|
29
|
+
type DeleteAgentRunExpected,
|
|
30
|
+
type IAgentRunLedger,
|
|
31
|
+
type ListPendingRunsOptions,
|
|
32
|
+
type ReapedAgentRun,
|
|
33
|
+
type RollbackFailureResult,
|
|
34
|
+
type SettleExpiredRunOptions,
|
|
35
|
+
} from "@axiom-lattice/core";
|
|
36
|
+
import { MigrationManager } from "../migrations/migration";
|
|
37
|
+
import { createAgentRunLedgerTable } from "../migrations/agent_run_ledger_migration";
|
|
38
|
+
import { createAgentRunTaskQueueMigration } from "../migrations/agent_run_task_queue_migration";
|
|
39
|
+
import { createAgentRunCapabilityRevisionMigration } from "../migrations/agent_run_capability_revision_migration";
|
|
40
|
+
import { createAgentRunBaseCheckpointCapturedMigration } from "../migrations/agent_run_base_checkpoint_captured_migration";
|
|
41
|
+
import { createAgentRunRollbackFailureMigration } from "../migrations/agent_run_rollback_failure_migration";
|
|
42
|
+
import { createAgentRunRetentionMigration } from "../migrations/agent_run_retention_migration";
|
|
43
|
+
import { agentRunRunningIndexSourceMigration } from "../migrations/agent_run_running_index_source_migration";
|
|
44
|
+
import { createAgentRunProjectTaskMigration } from "../migrations/agent_run_project_task_migration";
|
|
45
|
+
import { agentRunProjectTaskRunningIndexMigration } from "../migrations/agent_run_project_task_running_index_migration";
|
|
46
|
+
|
|
47
|
+
/** Options accepted by {@link PostgreSQLAgentRunLedger}. */
|
|
48
|
+
export interface PostgreSQLAgentRunLedgerOptions {
|
|
49
|
+
/** Externally managed pool. When supplied, migrations are not run here. */
|
|
50
|
+
pool?: Pool;
|
|
51
|
+
/** Connection string or pool config used to create a pool. */
|
|
52
|
+
poolConfig?: string | PoolConfig;
|
|
53
|
+
/** Set to `false` to skip running this store's migration automatically. */
|
|
54
|
+
autoMigrate?: boolean;
|
|
55
|
+
}
|
|
56
|
+
|
|
57
|
+
/** Raw `lattice_agent_runs` row shape. */
|
|
58
|
+
interface AgentRunRow {
|
|
59
|
+
run_id: string;
|
|
60
|
+
tenant_id: string;
|
|
61
|
+
workspace_id: string | null;
|
|
62
|
+
project_id: string | null;
|
|
63
|
+
task_id: string | null;
|
|
64
|
+
assistant_id: string;
|
|
65
|
+
thread_id: string;
|
|
66
|
+
status: string;
|
|
67
|
+
source: string | null;
|
|
68
|
+
owner: string | null;
|
|
69
|
+
lease_expires_at: Date | null;
|
|
70
|
+
input_message_ids: string[];
|
|
71
|
+
attempts: number;
|
|
72
|
+
error: string | null;
|
|
73
|
+
cancel_requested: boolean;
|
|
74
|
+
rollback_requested: boolean;
|
|
75
|
+
rollback_attempts: number;
|
|
76
|
+
rollback_error: string | null;
|
|
77
|
+
input: unknown;
|
|
78
|
+
command: unknown;
|
|
79
|
+
config: unknown;
|
|
80
|
+
metadata: unknown;
|
|
81
|
+
base_checkpoint_id: string | null;
|
|
82
|
+
base_checkpoint_captured: boolean;
|
|
83
|
+
capability_runtime_revision: string | null;
|
|
84
|
+
available_at: Date;
|
|
85
|
+
created_at: Date;
|
|
86
|
+
updated_at: Date;
|
|
87
|
+
started_at: Date | null;
|
|
88
|
+
finished_at: Date | null;
|
|
89
|
+
}
|
|
90
|
+
|
|
91
|
+
const DEFAULT_LIST_LIMIT = 50;
|
|
92
|
+
const MAX_LIST_LIMIT = 200;
|
|
93
|
+
|
|
94
|
+
const CLAIM_RUN_SQL = `UPDATE lattice_agent_runs AS r
|
|
95
|
+
SET status = 'running',
|
|
96
|
+
owner = $2,
|
|
97
|
+
lease_expires_at = now() + ($3::bigint || ' milliseconds')::interval,
|
|
98
|
+
started_at = now(),
|
|
99
|
+
attempts = r.attempts + 1,
|
|
100
|
+
input_message_ids = COALESCE($4::text[], r.input_message_ids),
|
|
101
|
+
updated_at = now()
|
|
102
|
+
WHERE r.run_id = $1
|
|
103
|
+
AND r.source = ANY($5::text[])
|
|
104
|
+
AND r.status = 'pending'
|
|
105
|
+
AND r.available_at <= now()
|
|
106
|
+
AND NOT EXISTS (
|
|
107
|
+
SELECT 1 FROM lattice_agent_runs AS o
|
|
108
|
+
WHERE o.tenant_id = r.tenant_id
|
|
109
|
+
AND o.thread_id = r.thread_id
|
|
110
|
+
AND o.source IN ('run_api', 'project_task')
|
|
111
|
+
AND o.status = 'running'
|
|
112
|
+
)
|
|
113
|
+
AND NOT EXISTS (
|
|
114
|
+
SELECT 1 FROM lattice_agent_runs AS b
|
|
115
|
+
WHERE b.tenant_id = r.tenant_id
|
|
116
|
+
AND b.thread_id = r.thread_id
|
|
117
|
+
AND b.source = 'run_api'
|
|
118
|
+
AND b.rollback_requested = TRUE
|
|
119
|
+
)
|
|
120
|
+
RETURNING r.*`;
|
|
121
|
+
|
|
122
|
+
/**
|
|
123
|
+
* PostgreSQL implementation of {@link IAgentRunLedger}.
|
|
124
|
+
*
|
|
125
|
+
* @remarks
|
|
126
|
+
* - Run rows are the durable queue for `run_api` and `project_task` work.
|
|
127
|
+
* Claim is an atomic CAS guarded by the same-thread running rule.
|
|
128
|
+
* - Legacy `beginRun` keeps the historical reconcile-then-insert behavior and
|
|
129
|
+
* writes `source='legacy'`.
|
|
130
|
+
* - `reapExpired` reports expired run-queue rows without mutating them;
|
|
131
|
+
* `reapLegacyExpired` preserves the supervisor's legacy settle behavior.
|
|
132
|
+
*/
|
|
133
|
+
export class PostgreSQLAgentRunLedger implements IAgentRunLedger {
|
|
134
|
+
private pool: Pool;
|
|
135
|
+
private ownsPool = true;
|
|
136
|
+
private initialized = false;
|
|
137
|
+
private initPromise: Promise<void> | null = null;
|
|
138
|
+
private migrationManager!: MigrationManager;
|
|
139
|
+
|
|
140
|
+
/**
|
|
141
|
+
* @param options An externally managed pool, or a pool config/connection URL.
|
|
142
|
+
*/
|
|
143
|
+
constructor(options: PostgreSQLAgentRunLedgerOptions) {
|
|
144
|
+
if (options.pool) {
|
|
145
|
+
this.pool = options.pool;
|
|
146
|
+
this.ownsPool = false;
|
|
147
|
+
this.initialized = true;
|
|
148
|
+
return;
|
|
149
|
+
}
|
|
150
|
+
|
|
151
|
+
if (typeof options.poolConfig === "string") {
|
|
152
|
+
this.pool = new Pool({ connectionString: options.poolConfig });
|
|
153
|
+
} else if (options.poolConfig) {
|
|
154
|
+
this.pool = new Pool(options.poolConfig as PoolConfig);
|
|
155
|
+
} else {
|
|
156
|
+
throw new Error("Either pool or poolConfig must be provided");
|
|
157
|
+
}
|
|
158
|
+
|
|
159
|
+
this.migrationManager = new MigrationManager(this.pool);
|
|
160
|
+
this.migrationManager.register(createAgentRunLedgerTable);
|
|
161
|
+
this.migrationManager.register(createAgentRunTaskQueueMigration);
|
|
162
|
+
this.migrationManager.register(createAgentRunCapabilityRevisionMigration);
|
|
163
|
+
this.migrationManager.register(createAgentRunBaseCheckpointCapturedMigration);
|
|
164
|
+
this.migrationManager.register(createAgentRunRollbackFailureMigration);
|
|
165
|
+
this.migrationManager.register(createAgentRunRetentionMigration);
|
|
166
|
+
this.migrationManager.register(agentRunRunningIndexSourceMigration);
|
|
167
|
+
this.migrationManager.register(createAgentRunProjectTaskMigration);
|
|
168
|
+
this.migrationManager.register(agentRunProjectTaskRunningIndexMigration);
|
|
169
|
+
|
|
170
|
+
if (options.autoMigrate !== false) {
|
|
171
|
+
this.initialize().catch((error) => {
|
|
172
|
+
console.error("Failed to initialize PostgreSQLAgentRunLedger:", error);
|
|
173
|
+
throw error;
|
|
174
|
+
});
|
|
175
|
+
}
|
|
176
|
+
}
|
|
177
|
+
|
|
178
|
+
/** Runs pending migrations once. */
|
|
179
|
+
async initialize(): Promise<void> {
|
|
180
|
+
if (this.initialized) return;
|
|
181
|
+
if (this.initPromise) return this.initPromise;
|
|
182
|
+
|
|
183
|
+
this.initPromise = (async () => {
|
|
184
|
+
try {
|
|
185
|
+
await this.migrationManager.migrate();
|
|
186
|
+
this.initialized = true;
|
|
187
|
+
} finally {
|
|
188
|
+
this.initPromise = null;
|
|
189
|
+
}
|
|
190
|
+
})();
|
|
191
|
+
return this.initPromise;
|
|
192
|
+
}
|
|
193
|
+
|
|
194
|
+
/** Ends the pool when this store owns it. */
|
|
195
|
+
async dispose(): Promise<void> {
|
|
196
|
+
if (this.ownsPool && this.pool) {
|
|
197
|
+
await this.pool.end();
|
|
198
|
+
}
|
|
199
|
+
}
|
|
200
|
+
|
|
201
|
+
/**
|
|
202
|
+
* Creates a running legacy run, reconciling any existing running row first.
|
|
203
|
+
*
|
|
204
|
+
* @remarks The caller holds the `thread-run` lease, so a surviving `running`
|
|
205
|
+
* row is stale (its owner can no longer hold the lease) even when its ledger
|
|
206
|
+
* lease has not expired yet.
|
|
207
|
+
*/
|
|
208
|
+
async beginRun(input: BeginAgentRunInput): Promise<AgentRunRecord> {
|
|
209
|
+
const client = await this.pool.connect();
|
|
210
|
+
try {
|
|
211
|
+
await client.query("BEGIN");
|
|
212
|
+
await client.query(
|
|
213
|
+
`UPDATE lattice_agent_runs
|
|
214
|
+
SET status = 'interrupted',
|
|
215
|
+
error = COALESCE(error, 'superseded by a new run on the same thread'),
|
|
216
|
+
finished_at = now(),
|
|
217
|
+
updated_at = now()
|
|
218
|
+
WHERE tenant_id = $1 AND thread_id = $2
|
|
219
|
+
AND status = 'running'`,
|
|
220
|
+
[input.tenantId, input.threadId],
|
|
221
|
+
);
|
|
222
|
+
const result = await client.query<AgentRunRow>(
|
|
223
|
+
`INSERT INTO lattice_agent_runs
|
|
224
|
+
(run_id, tenant_id, workspace_id, project_id, assistant_id, thread_id,
|
|
225
|
+
status, source, owner, lease_expires_at, input_message_ids, attempts,
|
|
226
|
+
started_at)
|
|
227
|
+
VALUES ($1, $2, $3, $4, $5, $6, 'running', 'legacy', $7,
|
|
228
|
+
now() + ($8::bigint || ' milliseconds')::interval, $9, 1, now())
|
|
229
|
+
RETURNING *`,
|
|
230
|
+
[
|
|
231
|
+
input.runId,
|
|
232
|
+
input.tenantId,
|
|
233
|
+
input.workspaceId ?? null,
|
|
234
|
+
input.projectId ?? null,
|
|
235
|
+
input.assistantId,
|
|
236
|
+
input.threadId,
|
|
237
|
+
input.owner,
|
|
238
|
+
input.leaseMs,
|
|
239
|
+
input.inputMessageIds,
|
|
240
|
+
],
|
|
241
|
+
);
|
|
242
|
+
await client.query("COMMIT");
|
|
243
|
+
return rowToRecord(result.rows[0]);
|
|
244
|
+
} catch (error) {
|
|
245
|
+
await client.query("ROLLBACK");
|
|
246
|
+
if (isActiveRunViolation(error)) {
|
|
247
|
+
throw new AgentRunAlreadyActiveError();
|
|
248
|
+
}
|
|
249
|
+
throw error;
|
|
250
|
+
} finally {
|
|
251
|
+
client.release();
|
|
252
|
+
}
|
|
253
|
+
}
|
|
254
|
+
|
|
255
|
+
/** Creates an unowned pending run-queue row. */
|
|
256
|
+
async createPendingRun(
|
|
257
|
+
input: CreatePendingAgentRunInput,
|
|
258
|
+
): Promise<AgentRunRecord> {
|
|
259
|
+
const source = input.source ?? "run_api";
|
|
260
|
+
if (source === "project_task") {
|
|
261
|
+
if (input.taskId === undefined || input.taskId.trim().length === 0) {
|
|
262
|
+
throw new AgentRunInvalidTaskIdError(
|
|
263
|
+
"A project_task run requires a non-empty taskId",
|
|
264
|
+
);
|
|
265
|
+
}
|
|
266
|
+
} else if (input.taskId !== undefined) {
|
|
267
|
+
throw new AgentRunInvalidTaskIdError(
|
|
268
|
+
`A ${source} run cannot carry a taskId`,
|
|
269
|
+
);
|
|
270
|
+
}
|
|
271
|
+
if (input.input != null && input.command != null) {
|
|
272
|
+
throw new AgentRunInputCommandConflictError();
|
|
273
|
+
}
|
|
274
|
+
const inputMessageIds = input.inputMessageIds
|
|
275
|
+
? normalizeAgentRunMessageIds(input.inputMessageIds)
|
|
276
|
+
: extractAgentRunInputMessageIds(input.input);
|
|
277
|
+
const availableAt = input.availableAt ?? Date.now();
|
|
278
|
+
const limit = input.maxActiveRunsForTenant;
|
|
279
|
+
const client = await this.pool.connect();
|
|
280
|
+
try {
|
|
281
|
+
await client.query("BEGIN");
|
|
282
|
+
// Quota admission must serialize every create for one tenant, not just
|
|
283
|
+
// one thread, or concurrent creates could each observe spare capacity.
|
|
284
|
+
if (limit !== undefined && limit > 0) {
|
|
285
|
+
await client.query("SELECT pg_advisory_xact_lock(hashtext($1))", [
|
|
286
|
+
input.tenantId,
|
|
287
|
+
]);
|
|
288
|
+
}
|
|
289
|
+
// Serialize duplicate-id checks and inserts per thread. The lock is
|
|
290
|
+
// transaction-scoped and does not block unrelated threads.
|
|
291
|
+
await client.query("SELECT pg_advisory_xact_lock(hashtext($1))", [
|
|
292
|
+
`agent-run:${input.tenantId}:${input.threadId}`,
|
|
293
|
+
]);
|
|
294
|
+
|
|
295
|
+
if (inputMessageIds.length > 0) {
|
|
296
|
+
const duplicate = await client.query<{ input_message_ids: string[] }>(
|
|
297
|
+
`SELECT input_message_ids
|
|
298
|
+
FROM lattice_agent_runs
|
|
299
|
+
WHERE tenant_id = $1 AND thread_id = $2
|
|
300
|
+
AND status IN ('pending', 'running')
|
|
301
|
+
AND input_message_ids && $3::text[]
|
|
302
|
+
LIMIT 1`,
|
|
303
|
+
[input.tenantId, input.threadId, inputMessageIds],
|
|
304
|
+
);
|
|
305
|
+
if ((duplicate.rowCount ?? 0) > 0) {
|
|
306
|
+
const existing = duplicate.rows[0].input_message_ids;
|
|
307
|
+
const messageId =
|
|
308
|
+
existing.find((id) => inputMessageIds.includes(id)) ??
|
|
309
|
+
inputMessageIds[0];
|
|
310
|
+
throw new AgentRunDuplicateMessageIdError(messageId);
|
|
311
|
+
}
|
|
312
|
+
}
|
|
313
|
+
|
|
314
|
+
if (limit !== undefined && limit > 0) {
|
|
315
|
+
const active = await client.query<{ count: string }>(
|
|
316
|
+
`SELECT count(*)::text AS count
|
|
317
|
+
FROM lattice_agent_runs
|
|
318
|
+
WHERE tenant_id = $1
|
|
319
|
+
AND source = 'run_api'
|
|
320
|
+
AND status IN ('pending', 'running')`,
|
|
321
|
+
[input.tenantId],
|
|
322
|
+
);
|
|
323
|
+
if (Number(active.rows[0].count) >= limit) {
|
|
324
|
+
throw new AgentRunTenantQuotaExceededError(input.tenantId, limit);
|
|
325
|
+
}
|
|
326
|
+
}
|
|
327
|
+
|
|
328
|
+
const result = await client.query<AgentRunRow>(
|
|
329
|
+
`INSERT INTO lattice_agent_runs
|
|
330
|
+
(run_id, tenant_id, workspace_id, project_id, assistant_id, thread_id,
|
|
331
|
+
status, source, owner, lease_expires_at, input_message_ids, attempts,
|
|
332
|
+
input, command, config, metadata, capability_runtime_revision,
|
|
333
|
+
available_at, created_at, updated_at, task_id)
|
|
334
|
+
VALUES ($1, $2, $3, $4, $5, $6, 'pending', $7, NULL, NULL, $8, 0,
|
|
335
|
+
$9::jsonb, $10::jsonb, $11::jsonb, $12::jsonb, $13,
|
|
336
|
+
to_timestamp($14::double precision / 1000.0),
|
|
337
|
+
date_trunc('milliseconds', now()),
|
|
338
|
+
date_trunc('milliseconds', now()), $15)
|
|
339
|
+
RETURNING *`,
|
|
340
|
+
[
|
|
341
|
+
input.runId,
|
|
342
|
+
input.tenantId,
|
|
343
|
+
input.workspaceId ?? null,
|
|
344
|
+
input.projectId ?? null,
|
|
345
|
+
input.assistantId,
|
|
346
|
+
input.threadId,
|
|
347
|
+
source,
|
|
348
|
+
inputMessageIds,
|
|
349
|
+
jsonParam(input.input),
|
|
350
|
+
jsonParam(input.command),
|
|
351
|
+
jsonParam(input.config),
|
|
352
|
+
jsonParam(input.metadata),
|
|
353
|
+
input.capabilityRuntimeRevision ?? null,
|
|
354
|
+
availableAt,
|
|
355
|
+
input.taskId ?? null,
|
|
356
|
+
],
|
|
357
|
+
);
|
|
358
|
+
await client.query("COMMIT");
|
|
359
|
+
return rowToRecord(result.rows[0]);
|
|
360
|
+
} catch (error) {
|
|
361
|
+
await client.query("ROLLBACK");
|
|
362
|
+
throw error;
|
|
363
|
+
} finally {
|
|
364
|
+
client.release();
|
|
365
|
+
}
|
|
366
|
+
}
|
|
367
|
+
|
|
368
|
+
/** Lists claimable run-queue rows, oldest available first. */
|
|
369
|
+
async listPendingRuns(
|
|
370
|
+
options?: ListPendingRunsOptions,
|
|
371
|
+
): Promise<AgentRunRecord[]> {
|
|
372
|
+
const limit = clampLimit(options?.limit);
|
|
373
|
+
const after = options?.after;
|
|
374
|
+
const fairTenants = options?.fairTenants === true;
|
|
375
|
+
if (fairTenants && after !== undefined) {
|
|
376
|
+
throw new RangeError(
|
|
377
|
+
"listPendingRuns fairTenants does not support the after cursor",
|
|
378
|
+
);
|
|
379
|
+
}
|
|
380
|
+
const maxRunning = positiveLimit(options?.maxRunningPerTenant);
|
|
381
|
+
const sources: AgentRunSource[] =
|
|
382
|
+
options?.sources === undefined ? ["run_api"] : [...options.sources];
|
|
383
|
+
assertQueueSources(sources);
|
|
384
|
+
const params: unknown[] = [limit];
|
|
385
|
+
let cursorClause = "";
|
|
386
|
+
if (after) {
|
|
387
|
+
// createPendingRun stores millisecond-precision timestamps, so cursor
|
|
388
|
+
// values round-trip exactly through AgentRunRecord's numeric fields.
|
|
389
|
+
params.push(
|
|
390
|
+
new Date(after.availableAt),
|
|
391
|
+
new Date(after.createdAt),
|
|
392
|
+
after.runId,
|
|
393
|
+
);
|
|
394
|
+
cursorClause = `
|
|
395
|
+
AND (r.available_at, r.created_at, r.run_id)
|
|
396
|
+
> ($2::timestamptz, $3::timestamptz, $4::uuid)`;
|
|
397
|
+
}
|
|
398
|
+
params.push(maxRunning ?? null);
|
|
399
|
+
const capIndex = params.length;
|
|
400
|
+
params.push(sources);
|
|
401
|
+
const sourcesIndex = params.length;
|
|
402
|
+
// Exclude a tenant already at its execution cap so the scan can reach
|
|
403
|
+
// other tenants. This is best-effort fairness, not round-robin scheduling.
|
|
404
|
+
const capClause = `
|
|
405
|
+
AND ($${capIndex}::bigint IS NULL OR (
|
|
406
|
+
SELECT count(*) FROM lattice_agent_runs AS c
|
|
407
|
+
WHERE c.tenant_id = r.tenant_id
|
|
408
|
+
AND c.source = ANY($${sourcesIndex}::text[])
|
|
409
|
+
AND c.status = 'running'
|
|
410
|
+
) < $${capIndex}::bigint)`;
|
|
411
|
+
const filterClause = `WHERE r.source = ANY($${sourcesIndex}::text[]) AND r.status = 'pending'
|
|
412
|
+
AND r.available_at <= now()
|
|
413
|
+
AND NOT EXISTS (
|
|
414
|
+
SELECT 1 FROM lattice_agent_runs AS o
|
|
415
|
+
WHERE o.tenant_id = r.tenant_id
|
|
416
|
+
AND o.thread_id = r.thread_id
|
|
417
|
+
AND o.source IN ('run_api', 'project_task')
|
|
418
|
+
AND o.status = 'running'
|
|
419
|
+
)
|
|
420
|
+
AND NOT EXISTS (
|
|
421
|
+
SELECT 1 FROM lattice_agent_runs AS b
|
|
422
|
+
WHERE b.tenant_id = r.tenant_id
|
|
423
|
+
AND b.thread_id = r.thread_id
|
|
424
|
+
AND b.source = 'run_api'
|
|
425
|
+
AND b.rollback_requested = TRUE
|
|
426
|
+
)${capClause}${cursorClause}`;
|
|
427
|
+
if (fairTenants) {
|
|
428
|
+
// Rank each tenant's pending rows, then interleave by rank. This is
|
|
429
|
+
// best-effort fair ordering at scan time: the first row of every tenant
|
|
430
|
+
// precedes the second row of any tenant, but a later scan can still see
|
|
431
|
+
// a different mix as work is claimed.
|
|
432
|
+
const ranked = await this.pool.query<AgentRunRow>(
|
|
433
|
+
`WITH ranked AS (
|
|
434
|
+
SELECT r.*,
|
|
435
|
+
row_number() OVER (
|
|
436
|
+
PARTITION BY r.tenant_id
|
|
437
|
+
ORDER BY r.available_at ASC, r.created_at ASC, r.run_id ASC
|
|
438
|
+
) AS tenant_rank
|
|
439
|
+
FROM lattice_agent_runs AS r
|
|
440
|
+
${filterClause}
|
|
441
|
+
)
|
|
442
|
+
SELECT * FROM ranked
|
|
443
|
+
ORDER BY tenant_rank ASC, available_at ASC, created_at ASC,
|
|
444
|
+
run_id ASC, tenant_id ASC
|
|
445
|
+
LIMIT $1`,
|
|
446
|
+
params,
|
|
447
|
+
);
|
|
448
|
+
return ranked.rows.map(rowToRecord);
|
|
449
|
+
}
|
|
450
|
+
const result = await this.pool.query<AgentRunRow>(
|
|
451
|
+
`SELECT r.* FROM lattice_agent_runs AS r
|
|
452
|
+
${filterClause}
|
|
453
|
+
ORDER BY r.available_at ASC, r.created_at ASC, r.run_id ASC
|
|
454
|
+
LIMIT $1`,
|
|
455
|
+
params,
|
|
456
|
+
);
|
|
457
|
+
return result.rows.map(rowToRecord);
|
|
458
|
+
}
|
|
459
|
+
|
|
460
|
+
/** Atomically claims a pending run-queue row. */
|
|
461
|
+
async claimRun(
|
|
462
|
+
input: ClaimAgentRunInput,
|
|
463
|
+
): Promise<AgentRunRecord | undefined> {
|
|
464
|
+
const maxRunning = positiveLimit(input.maxRunningPerTenant);
|
|
465
|
+
const sources: AgentRunSource[] =
|
|
466
|
+
input.sources === undefined ? ["run_api"] : [...input.sources];
|
|
467
|
+
assertQueueSources(sources);
|
|
468
|
+
if (maxRunning !== undefined) {
|
|
469
|
+
return this.claimRunWithTenantCap(input, maxRunning, sources);
|
|
470
|
+
}
|
|
471
|
+
try {
|
|
472
|
+
const result = await this.pool.query<AgentRunRow>(
|
|
473
|
+
CLAIM_RUN_SQL,
|
|
474
|
+
[
|
|
475
|
+
input.runId,
|
|
476
|
+
input.owner,
|
|
477
|
+
input.leaseMs,
|
|
478
|
+
input.inputMessageIds ?? null,
|
|
479
|
+
sources,
|
|
480
|
+
],
|
|
481
|
+
);
|
|
482
|
+
return result.rowCount === 0 ? undefined : rowToRecord(result.rows[0]);
|
|
483
|
+
} catch (error) {
|
|
484
|
+
// The partial unique index is the safety net for a concurrent claim.
|
|
485
|
+
if (isActiveRunViolation(error)) return undefined;
|
|
486
|
+
throw error;
|
|
487
|
+
}
|
|
488
|
+
}
|
|
489
|
+
|
|
490
|
+
/**
|
|
491
|
+
* Claims under a per-tenant running cap.
|
|
492
|
+
*
|
|
493
|
+
* @remarks Serializes claims for one tenant with a transaction-scoped
|
|
494
|
+
* advisory lock, counts running rows, then runs the same claim statement.
|
|
495
|
+
* The lock is shared with admission creates so cap checks cannot race.
|
|
496
|
+
*/
|
|
497
|
+
private async claimRunWithTenantCap(
|
|
498
|
+
input: ClaimAgentRunInput,
|
|
499
|
+
maxRunning: number,
|
|
500
|
+
sources: AgentRunSource[],
|
|
501
|
+
): Promise<AgentRunRecord | undefined> {
|
|
502
|
+
const client = await this.pool.connect();
|
|
503
|
+
try {
|
|
504
|
+
await client.query("BEGIN");
|
|
505
|
+
const candidate = await client.query<{ tenant_id: string }>(
|
|
506
|
+
`SELECT tenant_id FROM lattice_agent_runs
|
|
507
|
+
WHERE run_id = $1 AND source = ANY($2::text[]) AND status = 'pending'`,
|
|
508
|
+
[input.runId, sources],
|
|
509
|
+
);
|
|
510
|
+
if (candidate.rowCount === 0) {
|
|
511
|
+
await client.query("COMMIT");
|
|
512
|
+
return undefined;
|
|
513
|
+
}
|
|
514
|
+
const tenantId = candidate.rows[0].tenant_id;
|
|
515
|
+
await client.query("SELECT pg_advisory_xact_lock(hashtext($1))", [
|
|
516
|
+
tenantId,
|
|
517
|
+
]);
|
|
518
|
+
const running = await client.query<{ count: string }>(
|
|
519
|
+
`SELECT count(*)::text AS count
|
|
520
|
+
FROM lattice_agent_runs
|
|
521
|
+
WHERE tenant_id = $1
|
|
522
|
+
AND source = ANY($2::text[])
|
|
523
|
+
AND status = 'running'`,
|
|
524
|
+
[tenantId, sources],
|
|
525
|
+
);
|
|
526
|
+
if (Number(running.rows[0].count) >= maxRunning) {
|
|
527
|
+
await client.query("COMMIT");
|
|
528
|
+
return undefined;
|
|
529
|
+
}
|
|
530
|
+
const result = await client.query<AgentRunRow>(
|
|
531
|
+
CLAIM_RUN_SQL,
|
|
532
|
+
[
|
|
533
|
+
input.runId,
|
|
534
|
+
input.owner,
|
|
535
|
+
input.leaseMs,
|
|
536
|
+
input.inputMessageIds ?? null,
|
|
537
|
+
sources,
|
|
538
|
+
],
|
|
539
|
+
);
|
|
540
|
+
await client.query("COMMIT");
|
|
541
|
+
return result.rowCount === 0 ? undefined : rowToRecord(result.rows[0]);
|
|
542
|
+
} catch (error) {
|
|
543
|
+
await client.query("ROLLBACK");
|
|
544
|
+
if (isActiveRunViolation(error)) return undefined;
|
|
545
|
+
throw error;
|
|
546
|
+
} finally {
|
|
547
|
+
client.release();
|
|
548
|
+
}
|
|
549
|
+
}
|
|
550
|
+
|
|
551
|
+
/** Captures the checkpoint a running run-queue row started from. */
|
|
552
|
+
async setBaseCheckpoint(
|
|
553
|
+
tenantId: string,
|
|
554
|
+
runId: string,
|
|
555
|
+
owner: string,
|
|
556
|
+
baseCheckpointId: string | null,
|
|
557
|
+
): Promise<boolean> {
|
|
558
|
+
const result = await this.pool.query(
|
|
559
|
+
`UPDATE lattice_agent_runs
|
|
560
|
+
SET base_checkpoint_id = $4,
|
|
561
|
+
base_checkpoint_captured = TRUE,
|
|
562
|
+
updated_at = now()
|
|
563
|
+
WHERE tenant_id = $1 AND run_id = $2 AND owner = $3
|
|
564
|
+
AND status = 'running'
|
|
565
|
+
AND source IN ('run_api', 'project_task')
|
|
566
|
+
RETURNING run_id`,
|
|
567
|
+
[tenantId, runId, owner, baseCheckpointId],
|
|
568
|
+
);
|
|
569
|
+
return (result.rowCount ?? 0) > 0;
|
|
570
|
+
}
|
|
571
|
+
|
|
572
|
+
/** Records the resolved capability runtime revision for a running run. */
|
|
573
|
+
async setCapabilityRevision(
|
|
574
|
+
tenantId: string,
|
|
575
|
+
runId: string,
|
|
576
|
+
owner: string,
|
|
577
|
+
revision: string | null,
|
|
578
|
+
): Promise<boolean> {
|
|
579
|
+
const result = await this.pool.query(
|
|
580
|
+
`UPDATE lattice_agent_runs
|
|
581
|
+
SET capability_runtime_revision = $4, updated_at = now()
|
|
582
|
+
WHERE tenant_id = $1 AND run_id = $2 AND owner = $3
|
|
583
|
+
AND status = 'running'
|
|
584
|
+
AND source IN ('run_api', 'project_task')
|
|
585
|
+
RETURNING run_id`,
|
|
586
|
+
[tenantId, runId, owner, revision],
|
|
587
|
+
);
|
|
588
|
+
return (result.rowCount ?? 0) > 0;
|
|
589
|
+
}
|
|
590
|
+
|
|
591
|
+
/** Deletes a run row immediately, optionally guarded by status/owner. */
|
|
592
|
+
async deleteRun(
|
|
593
|
+
tenantId: string,
|
|
594
|
+
runId: string,
|
|
595
|
+
expected?: DeleteAgentRunExpected,
|
|
596
|
+
): Promise<boolean> {
|
|
597
|
+
const params: unknown[] = [tenantId, runId];
|
|
598
|
+
let predicate = "tenant_id = $1 AND run_id = $2";
|
|
599
|
+
if (expected?.status !== undefined) {
|
|
600
|
+
params.push(expected.status);
|
|
601
|
+
predicate += ` AND status = $${params.length}`;
|
|
602
|
+
}
|
|
603
|
+
if (expected?.owner !== undefined) {
|
|
604
|
+
params.push(expected.owner);
|
|
605
|
+
predicate += ` AND owner = $${params.length}`;
|
|
606
|
+
}
|
|
607
|
+
const result = await this.pool.query(
|
|
608
|
+
`DELETE FROM lattice_agent_runs WHERE ${predicate}`,
|
|
609
|
+
params,
|
|
610
|
+
);
|
|
611
|
+
return (result.rowCount ?? 0) > 0;
|
|
612
|
+
}
|
|
613
|
+
|
|
614
|
+
/** Deletes terminal rows finished before `cutoffMs`, oldest first. */
|
|
615
|
+
async deleteTerminalBefore(
|
|
616
|
+
cutoffMs: number,
|
|
617
|
+
limit: number,
|
|
618
|
+
): Promise<number> {
|
|
619
|
+
if (!Number.isFinite(cutoffMs)) return 0;
|
|
620
|
+
const max = Number.isFinite(limit) ? Math.floor(limit) : 0;
|
|
621
|
+
if (max <= 0) return 0;
|
|
622
|
+
|
|
623
|
+
const result = await this.pool.query(
|
|
624
|
+
`DELETE FROM lattice_agent_runs
|
|
625
|
+
WHERE run_id IN (
|
|
626
|
+
SELECT run_id FROM lattice_agent_runs
|
|
627
|
+
WHERE status IN ('success', 'error', 'interrupted', 'cancelled', 'timeout')
|
|
628
|
+
AND finished_at < to_timestamp($1::double precision / 1000.0)
|
|
629
|
+
ORDER BY finished_at ASC
|
|
630
|
+
LIMIT $2
|
|
631
|
+
)`,
|
|
632
|
+
[cutoffMs, max],
|
|
633
|
+
);
|
|
634
|
+
return result.rowCount ?? 0;
|
|
635
|
+
}
|
|
636
|
+
|
|
637
|
+
/** Extends a running run's lease when `owner` still holds it. */
|
|
638
|
+
async heartbeatRun(
|
|
639
|
+
runId: string,
|
|
640
|
+
owner: string,
|
|
641
|
+
leaseMs: number,
|
|
642
|
+
): Promise<boolean> {
|
|
643
|
+
const result = await this.pool.query(
|
|
644
|
+
`UPDATE lattice_agent_runs
|
|
645
|
+
SET lease_expires_at = now() + ($3::bigint || ' milliseconds')::interval,
|
|
646
|
+
updated_at = now()
|
|
647
|
+
WHERE run_id = $1 AND owner = $2 AND status = 'running'
|
|
648
|
+
RETURNING run_id`,
|
|
649
|
+
[runId, owner, leaseMs],
|
|
650
|
+
);
|
|
651
|
+
return (result.rowCount ?? 0) > 0;
|
|
652
|
+
}
|
|
653
|
+
|
|
654
|
+
/** Settles a running run when `owner` still holds it. */
|
|
655
|
+
async settleRun(
|
|
656
|
+
runId: string,
|
|
657
|
+
owner: string,
|
|
658
|
+
status: Exclude<AgentRunStatus, "pending" | "running">,
|
|
659
|
+
error?: string,
|
|
660
|
+
): Promise<boolean> {
|
|
661
|
+
const result = await this.pool.query(
|
|
662
|
+
`UPDATE lattice_agent_runs
|
|
663
|
+
SET status = $3, error = $4, finished_at = now(), updated_at = now()
|
|
664
|
+
WHERE run_id = $1 AND owner = $2 AND status = 'running'
|
|
665
|
+
RETURNING run_id`,
|
|
666
|
+
[runId, owner, status, error ?? null],
|
|
667
|
+
);
|
|
668
|
+
return (result.rowCount ?? 0) > 0;
|
|
669
|
+
}
|
|
670
|
+
|
|
671
|
+
/** Settles a running run only when no rollback is pending. */
|
|
672
|
+
async settleRunUnlessRollbackRequested(
|
|
673
|
+
runId: string,
|
|
674
|
+
owner: string,
|
|
675
|
+
status: Exclude<AgentRunStatus, "pending" | "running">,
|
|
676
|
+
error?: string,
|
|
677
|
+
): Promise<boolean> {
|
|
678
|
+
const result = await this.pool.query(
|
|
679
|
+
`UPDATE lattice_agent_runs
|
|
680
|
+
SET status = $3, error = $4, finished_at = now(), updated_at = now()
|
|
681
|
+
WHERE run_id = $1 AND owner = $2 AND status = 'running'
|
|
682
|
+
AND rollback_requested = FALSE
|
|
683
|
+
RETURNING run_id`,
|
|
684
|
+
[runId, owner, status, error ?? null],
|
|
685
|
+
);
|
|
686
|
+
return (result.rowCount ?? 0) > 0;
|
|
687
|
+
}
|
|
688
|
+
|
|
689
|
+
/** Terminates an uncaptured rollback run with a repair signal. */
|
|
690
|
+
async settleRollbackFailure(
|
|
691
|
+
tenantId: string,
|
|
692
|
+
runId: string,
|
|
693
|
+
error: string,
|
|
694
|
+
): Promise<boolean> {
|
|
695
|
+
const result = await this.pool.query(
|
|
696
|
+
`UPDATE lattice_agent_runs
|
|
697
|
+
SET status = 'error',
|
|
698
|
+
error = $3,
|
|
699
|
+
rollback_requested = FALSE,
|
|
700
|
+
rollback_error = $3,
|
|
701
|
+
finished_at = now(),
|
|
702
|
+
updated_at = now()
|
|
703
|
+
WHERE tenant_id = $1 AND run_id = $2
|
|
704
|
+
AND source = 'run_api'
|
|
705
|
+
AND rollback_requested = TRUE
|
|
706
|
+
AND base_checkpoint_captured = FALSE
|
|
707
|
+
AND status IN ('running', 'interrupted')
|
|
708
|
+
RETURNING run_id`,
|
|
709
|
+
[tenantId, runId, error],
|
|
710
|
+
);
|
|
711
|
+
return (result.rowCount ?? 0) > 0;
|
|
712
|
+
}
|
|
713
|
+
|
|
714
|
+
/** Records one bounded rollback failure and optionally converges. */
|
|
715
|
+
async recordRollbackFailure(
|
|
716
|
+
tenantId: string,
|
|
717
|
+
runId: string,
|
|
718
|
+
error: string,
|
|
719
|
+
maxAttempts: number,
|
|
720
|
+
): Promise<RollbackFailureResult | undefined> {
|
|
721
|
+
const result = await this.pool.query<{
|
|
722
|
+
rollback_attempts: number;
|
|
723
|
+
status: string;
|
|
724
|
+
}>(
|
|
725
|
+
`UPDATE lattice_agent_runs
|
|
726
|
+
SET rollback_attempts = rollback_attempts + 1,
|
|
727
|
+
rollback_error = $3,
|
|
728
|
+
status = CASE
|
|
729
|
+
WHEN rollback_attempts + 1 >= $4 THEN 'error'
|
|
730
|
+
ELSE status
|
|
731
|
+
END,
|
|
732
|
+
error = CASE
|
|
733
|
+
WHEN rollback_attempts + 1 >= $4 THEN $3
|
|
734
|
+
ELSE error
|
|
735
|
+
END,
|
|
736
|
+
rollback_requested = CASE
|
|
737
|
+
WHEN rollback_attempts + 1 >= $4 THEN FALSE
|
|
738
|
+
ELSE rollback_requested
|
|
739
|
+
END,
|
|
740
|
+
finished_at = CASE
|
|
741
|
+
WHEN rollback_attempts + 1 >= $4 THEN now()
|
|
742
|
+
ELSE finished_at
|
|
743
|
+
END,
|
|
744
|
+
updated_at = now()
|
|
745
|
+
WHERE tenant_id = $1
|
|
746
|
+
AND run_id = $2
|
|
747
|
+
AND source = 'run_api'
|
|
748
|
+
AND rollback_requested = TRUE
|
|
749
|
+
AND status IN ('running', 'interrupted')
|
|
750
|
+
RETURNING rollback_attempts, status`,
|
|
751
|
+
[tenantId, runId, error, maxAttempts],
|
|
752
|
+
);
|
|
753
|
+
if (result.rowCount === 0) return undefined;
|
|
754
|
+
const row = result.rows[0];
|
|
755
|
+
return {
|
|
756
|
+
attempts: row.rollback_attempts,
|
|
757
|
+
terminal: row.status === "error",
|
|
758
|
+
};
|
|
759
|
+
}
|
|
760
|
+
|
|
761
|
+
/** Reads a run within a tenant. */
|
|
762
|
+
async getRun(
|
|
763
|
+
tenantId: string,
|
|
764
|
+
runId: string,
|
|
765
|
+
): Promise<AgentRunRecord | undefined> {
|
|
766
|
+
const result = await this.pool.query<AgentRunRow>(
|
|
767
|
+
`SELECT * FROM lattice_agent_runs WHERE tenant_id = $1 AND run_id = $2`,
|
|
768
|
+
[tenantId, runId],
|
|
769
|
+
);
|
|
770
|
+
return result.rowCount === 0 ? undefined : rowToRecord(result.rows[0]);
|
|
771
|
+
}
|
|
772
|
+
|
|
773
|
+
/** Finds the active (pending/running) run for a ProjectTask. */
|
|
774
|
+
async getActiveRunForTask(
|
|
775
|
+
tenantId: string,
|
|
776
|
+
taskId: string,
|
|
777
|
+
): Promise<AgentRunRecord | undefined> {
|
|
778
|
+
const result = await this.pool.query<AgentRunRow>(
|
|
779
|
+
`SELECT * FROM lattice_agent_runs
|
|
780
|
+
WHERE tenant_id = $1 AND task_id = $2
|
|
781
|
+
AND status IN ('pending', 'running')
|
|
782
|
+
ORDER BY created_at DESC, run_id DESC
|
|
783
|
+
LIMIT 1`,
|
|
784
|
+
[tenantId, taskId],
|
|
785
|
+
);
|
|
786
|
+
return result.rowCount === 0 ? undefined : rowToRecord(result.rows[0]);
|
|
787
|
+
}
|
|
788
|
+
|
|
789
|
+
/** Lists a thread's runs, newest first (including pending rows). */
|
|
790
|
+
async listRuns(
|
|
791
|
+
tenantId: string,
|
|
792
|
+
threadId: string,
|
|
793
|
+
options?: { limit?: number; offset?: number },
|
|
794
|
+
): Promise<AgentRunRecord[]> {
|
|
795
|
+
const limit = clampLimit(options?.limit);
|
|
796
|
+
const offset = Math.max(0, Math.floor(options?.offset ?? 0));
|
|
797
|
+
const result = await this.pool.query<AgentRunRow>(
|
|
798
|
+
`SELECT * FROM lattice_agent_runs
|
|
799
|
+
WHERE tenant_id = $1 AND thread_id = $2
|
|
800
|
+
ORDER BY created_at DESC
|
|
801
|
+
LIMIT $3 OFFSET $4`,
|
|
802
|
+
[tenantId, threadId, limit, offset],
|
|
803
|
+
);
|
|
804
|
+
return result.rows.map(rowToRecord);
|
|
805
|
+
}
|
|
806
|
+
|
|
807
|
+
/** Finds the thread's running run, if any. */
|
|
808
|
+
async findRunning(
|
|
809
|
+
tenantId: string,
|
|
810
|
+
threadId: string,
|
|
811
|
+
): Promise<AgentRunRecord | undefined> {
|
|
812
|
+
const result = await this.pool.query<AgentRunRow>(
|
|
813
|
+
`SELECT * FROM lattice_agent_runs
|
|
814
|
+
WHERE tenant_id = $1 AND thread_id = $2 AND status = 'running'
|
|
815
|
+
LIMIT 1`,
|
|
816
|
+
[tenantId, threadId],
|
|
817
|
+
);
|
|
818
|
+
return result.rowCount === 0 ? undefined : rowToRecord(result.rows[0]);
|
|
819
|
+
}
|
|
820
|
+
|
|
821
|
+
/** Requests cancellation, settling pending/expired rows interrupted. */
|
|
822
|
+
async requestCancel(
|
|
823
|
+
tenantId: string,
|
|
824
|
+
runId: string,
|
|
825
|
+
): Promise<"cancelled" | "requested" | "not_found" | "terminal"> {
|
|
826
|
+
const client = await this.pool.connect();
|
|
827
|
+
try {
|
|
828
|
+
await client.query("BEGIN");
|
|
829
|
+
const current = await client.query<{
|
|
830
|
+
status: string;
|
|
831
|
+
expired: boolean;
|
|
832
|
+
}>(
|
|
833
|
+
`SELECT status,
|
|
834
|
+
COALESCE(lease_expires_at < now(), TRUE) AS expired
|
|
835
|
+
FROM lattice_agent_runs
|
|
836
|
+
WHERE tenant_id = $1 AND run_id = $2
|
|
837
|
+
FOR UPDATE`,
|
|
838
|
+
[tenantId, runId],
|
|
839
|
+
);
|
|
840
|
+
|
|
841
|
+
if (current.rowCount === 0) {
|
|
842
|
+
await client.query("COMMIT");
|
|
843
|
+
return "not_found";
|
|
844
|
+
}
|
|
845
|
+
|
|
846
|
+
const row = current.rows[0];
|
|
847
|
+
if (row.status === "pending") {
|
|
848
|
+
await client.query(
|
|
849
|
+
`UPDATE lattice_agent_runs
|
|
850
|
+
SET status = 'interrupted', cancel_requested = TRUE,
|
|
851
|
+
finished_at = now(), updated_at = now()
|
|
852
|
+
WHERE tenant_id = $1 AND run_id = $2`,
|
|
853
|
+
[tenantId, runId],
|
|
854
|
+
);
|
|
855
|
+
await client.query("COMMIT");
|
|
856
|
+
return "cancelled";
|
|
857
|
+
}
|
|
858
|
+
if (row.status !== "running") {
|
|
859
|
+
await client.query("COMMIT");
|
|
860
|
+
return "terminal";
|
|
861
|
+
}
|
|
862
|
+
|
|
863
|
+
if (row.expired) {
|
|
864
|
+
await client.query(
|
|
865
|
+
`UPDATE lattice_agent_runs
|
|
866
|
+
SET status = 'interrupted', cancel_requested = TRUE,
|
|
867
|
+
finished_at = now(), updated_at = now()
|
|
868
|
+
WHERE tenant_id = $1 AND run_id = $2`,
|
|
869
|
+
[tenantId, runId],
|
|
870
|
+
);
|
|
871
|
+
await client.query("COMMIT");
|
|
872
|
+
return "cancelled";
|
|
873
|
+
}
|
|
874
|
+
|
|
875
|
+
await client.query(
|
|
876
|
+
`UPDATE lattice_agent_runs
|
|
877
|
+
SET cancel_requested = TRUE, updated_at = now()
|
|
878
|
+
WHERE tenant_id = $1 AND run_id = $2`,
|
|
879
|
+
[tenantId, runId],
|
|
880
|
+
);
|
|
881
|
+
await client.query("COMMIT");
|
|
882
|
+
return "requested";
|
|
883
|
+
} catch (error) {
|
|
884
|
+
await client.query("ROLLBACK");
|
|
885
|
+
throw error;
|
|
886
|
+
} finally {
|
|
887
|
+
client.release();
|
|
888
|
+
}
|
|
889
|
+
}
|
|
890
|
+
|
|
891
|
+
/** Deletes a pending run or flags a running run for rollback. */
|
|
892
|
+
async requestRollback(
|
|
893
|
+
tenantId: string,
|
|
894
|
+
runId: string,
|
|
895
|
+
): Promise<AgentRunRollbackResult> {
|
|
896
|
+
const client = await this.pool.connect();
|
|
897
|
+
try {
|
|
898
|
+
await client.query("BEGIN");
|
|
899
|
+
const current = await client.query<{ status: string; source: string }>(
|
|
900
|
+
`SELECT status, source FROM lattice_agent_runs
|
|
901
|
+
WHERE tenant_id = $1 AND run_id = $2
|
|
902
|
+
FOR UPDATE`,
|
|
903
|
+
[tenantId, runId],
|
|
904
|
+
);
|
|
905
|
+
if (current.rowCount === 0) {
|
|
906
|
+
await client.query("COMMIT");
|
|
907
|
+
return "not_found";
|
|
908
|
+
}
|
|
909
|
+
|
|
910
|
+
const { status, source } = current.rows[0];
|
|
911
|
+
if (source !== "run_api") {
|
|
912
|
+
await client.query("COMMIT");
|
|
913
|
+
return "terminal";
|
|
914
|
+
}
|
|
915
|
+
if (status === "pending") {
|
|
916
|
+
await client.query(
|
|
917
|
+
`DELETE FROM lattice_agent_runs WHERE tenant_id = $1 AND run_id = $2`,
|
|
918
|
+
[tenantId, runId],
|
|
919
|
+
);
|
|
920
|
+
await client.query("COMMIT");
|
|
921
|
+
return "deleted";
|
|
922
|
+
}
|
|
923
|
+
if (status !== "running") {
|
|
924
|
+
await client.query("COMMIT");
|
|
925
|
+
return "terminal";
|
|
926
|
+
}
|
|
927
|
+
|
|
928
|
+
await client.query(
|
|
929
|
+
`UPDATE lattice_agent_runs
|
|
930
|
+
SET rollback_requested = TRUE, cancel_requested = TRUE,
|
|
931
|
+
updated_at = now()
|
|
932
|
+
WHERE tenant_id = $1 AND run_id = $2`,
|
|
933
|
+
[tenantId, runId],
|
|
934
|
+
);
|
|
935
|
+
await client.query("COMMIT");
|
|
936
|
+
return "requested";
|
|
937
|
+
} catch (error) {
|
|
938
|
+
await client.query("ROLLBACK");
|
|
939
|
+
throw error;
|
|
940
|
+
} finally {
|
|
941
|
+
client.release();
|
|
942
|
+
}
|
|
943
|
+
}
|
|
944
|
+
|
|
945
|
+
/** Atomically reclaims an expired running row without requeueing it. */
|
|
946
|
+
async recoverExpiredRun(
|
|
947
|
+
runId: string,
|
|
948
|
+
owner: string,
|
|
949
|
+
leaseMs: number,
|
|
950
|
+
): Promise<AgentRunRecord | undefined> {
|
|
951
|
+
const result = await this.pool.query<AgentRunRow>(
|
|
952
|
+
`UPDATE lattice_agent_runs
|
|
953
|
+
SET owner = $2,
|
|
954
|
+
lease_expires_at = now() + ($3::bigint || ' milliseconds')::interval,
|
|
955
|
+
attempts = attempts + 1,
|
|
956
|
+
updated_at = now()
|
|
957
|
+
WHERE run_id = $1
|
|
958
|
+
AND source IN ('run_api', 'project_task')
|
|
959
|
+
AND status = 'running'
|
|
960
|
+
AND lease_expires_at < now()
|
|
961
|
+
RETURNING *`,
|
|
962
|
+
[runId, owner, leaseMs],
|
|
963
|
+
);
|
|
964
|
+
return result.rowCount === 0 ? undefined : rowToRecord(result.rows[0]);
|
|
965
|
+
}
|
|
966
|
+
|
|
967
|
+
/** Atomically requeues or interrupts an expired run-queue row. */
|
|
968
|
+
async settleExpiredRun(
|
|
969
|
+
runId: string,
|
|
970
|
+
outcome: AgentRunExpiredOutcome,
|
|
971
|
+
options?: SettleExpiredRunOptions,
|
|
972
|
+
): Promise<boolean> {
|
|
973
|
+
if (outcome === "pending") {
|
|
974
|
+
const availableAt = options?.availableAt ?? Date.now();
|
|
975
|
+
const result = await this.pool.query(
|
|
976
|
+
`UPDATE lattice_agent_runs
|
|
977
|
+
SET status = 'pending',
|
|
978
|
+
owner = NULL,
|
|
979
|
+
lease_expires_at = NULL,
|
|
980
|
+
started_at = NULL,
|
|
981
|
+
available_at = to_timestamp($2::double precision / 1000.0),
|
|
982
|
+
finished_at = NULL,
|
|
983
|
+
updated_at = now()
|
|
984
|
+
WHERE run_id = $1
|
|
985
|
+
AND source IN ('run_api', 'project_task')
|
|
986
|
+
AND status = 'running'
|
|
987
|
+
AND lease_expires_at < now()
|
|
988
|
+
RETURNING run_id`,
|
|
989
|
+
[runId, availableAt],
|
|
990
|
+
);
|
|
991
|
+
return (result.rowCount ?? 0) > 0;
|
|
992
|
+
}
|
|
993
|
+
|
|
994
|
+
const result = await this.pool.query(
|
|
995
|
+
`UPDATE lattice_agent_runs
|
|
996
|
+
SET status = $3,
|
|
997
|
+
error = COALESCE($2, error),
|
|
998
|
+
finished_at = now(),
|
|
999
|
+
updated_at = now()
|
|
1000
|
+
WHERE run_id = $1
|
|
1001
|
+
AND source IN ('run_api', 'project_task')
|
|
1002
|
+
AND status = 'running'
|
|
1003
|
+
AND lease_expires_at < now()
|
|
1004
|
+
RETURNING run_id`,
|
|
1005
|
+
[
|
|
1006
|
+
runId,
|
|
1007
|
+
options?.error ?? null,
|
|
1008
|
+
outcome === "timeout" ? "timeout" : "interrupted",
|
|
1009
|
+
],
|
|
1010
|
+
);
|
|
1011
|
+
return (result.rowCount ?? 0) > 0;
|
|
1012
|
+
}
|
|
1013
|
+
|
|
1014
|
+
/** Counts runs grouped by lifecycle status. */
|
|
1015
|
+
async counts(): Promise<AgentRunCounts> {
|
|
1016
|
+
const counts: AgentRunCounts = {
|
|
1017
|
+
pending: 0,
|
|
1018
|
+
running: 0,
|
|
1019
|
+
success: 0,
|
|
1020
|
+
error: 0,
|
|
1021
|
+
interrupted: 0,
|
|
1022
|
+
cancelled: 0,
|
|
1023
|
+
timeout: 0,
|
|
1024
|
+
};
|
|
1025
|
+
const result = await this.pool.query<{ status: string; count: number }>(
|
|
1026
|
+
`SELECT status, count(*)::int AS count
|
|
1027
|
+
FROM lattice_agent_runs
|
|
1028
|
+
GROUP BY status`,
|
|
1029
|
+
);
|
|
1030
|
+
for (const row of result.rows) {
|
|
1031
|
+
const status = parseStatus(row.status);
|
|
1032
|
+
if (status) counts[status] = row.count;
|
|
1033
|
+
}
|
|
1034
|
+
return counts;
|
|
1035
|
+
}
|
|
1036
|
+
|
|
1037
|
+
/** Reports expired run-queue rows for the requested sources. */
|
|
1038
|
+
async reapExpired(
|
|
1039
|
+
sources: readonly AgentRunSource[] = ["run_api"],
|
|
1040
|
+
): Promise<ReapedAgentRun[]> {
|
|
1041
|
+
assertQueueSources(sources);
|
|
1042
|
+
const result = await this.pool.query<AgentRunRow>(
|
|
1043
|
+
`SELECT * FROM lattice_agent_runs
|
|
1044
|
+
WHERE source = ANY($1::text[]) AND status = 'running'
|
|
1045
|
+
AND lease_expires_at < now()
|
|
1046
|
+
ORDER BY lease_expires_at ASC`,
|
|
1047
|
+
[[...sources]],
|
|
1048
|
+
);
|
|
1049
|
+
return result.rows.map(reapedFromRow);
|
|
1050
|
+
}
|
|
1051
|
+
|
|
1052
|
+
/** Lists rollback work that still needs a checkpoint rewind. */
|
|
1053
|
+
async listRollbackRuns(): Promise<ReapedAgentRun[]> {
|
|
1054
|
+
const result = await this.pool.query<AgentRunRow>(
|
|
1055
|
+
`SELECT * FROM lattice_agent_runs
|
|
1056
|
+
WHERE source = 'run_api'
|
|
1057
|
+
AND rollback_requested = TRUE
|
|
1058
|
+
AND (
|
|
1059
|
+
status = 'interrupted'
|
|
1060
|
+
OR (status = 'running' AND lease_expires_at < now())
|
|
1061
|
+
)
|
|
1062
|
+
ORDER BY updated_at ASC`,
|
|
1063
|
+
);
|
|
1064
|
+
return result.rows.map(reapedFromRow);
|
|
1065
|
+
}
|
|
1066
|
+
|
|
1067
|
+
/** Settles expired legacy rows, preserving the historical supervisor behavior. */
|
|
1068
|
+
async reapLegacyExpired(): Promise<ReapedAgentRun[]> {
|
|
1069
|
+
const result = await this.pool.query<{
|
|
1070
|
+
run_id: string;
|
|
1071
|
+
tenant_id: string;
|
|
1072
|
+
assistant_id: string;
|
|
1073
|
+
thread_id: string;
|
|
1074
|
+
workspace_id: string | null;
|
|
1075
|
+
project_id: string | null;
|
|
1076
|
+
}>(
|
|
1077
|
+
`UPDATE lattice_agent_runs
|
|
1078
|
+
SET status = 'interrupted',
|
|
1079
|
+
error = COALESCE(error, 'run owner lost'),
|
|
1080
|
+
finished_at = now(),
|
|
1081
|
+
updated_at = now()
|
|
1082
|
+
WHERE source = 'legacy' AND status = 'running'
|
|
1083
|
+
AND lease_expires_at < now()
|
|
1084
|
+
RETURNING run_id, tenant_id, assistant_id, thread_id, workspace_id, project_id`,
|
|
1085
|
+
);
|
|
1086
|
+
return result.rows.map((row) => ({
|
|
1087
|
+
runId: row.run_id,
|
|
1088
|
+
tenantId: row.tenant_id,
|
|
1089
|
+
assistantId: row.assistant_id,
|
|
1090
|
+
threadId: row.thread_id,
|
|
1091
|
+
workspaceId: row.workspace_id ?? undefined,
|
|
1092
|
+
projectId: row.project_id ?? undefined,
|
|
1093
|
+
}));
|
|
1094
|
+
}
|
|
1095
|
+
}
|
|
1096
|
+
|
|
1097
|
+
function positiveLimit(limit: number | undefined): number | undefined {
|
|
1098
|
+
return limit !== undefined && Number.isSafeInteger(limit) && limit > 0
|
|
1099
|
+
? limit
|
|
1100
|
+
: undefined;
|
|
1101
|
+
}
|
|
1102
|
+
|
|
1103
|
+
function clampLimit(limit: number | undefined): number {
|
|
1104
|
+
if (limit === undefined) return DEFAULT_LIST_LIMIT;
|
|
1105
|
+
const floored = Math.floor(limit);
|
|
1106
|
+
if (!Number.isFinite(floored)) return DEFAULT_LIST_LIMIT;
|
|
1107
|
+
return Math.min(MAX_LIST_LIMIT, Math.max(1, floored));
|
|
1108
|
+
}
|
|
1109
|
+
|
|
1110
|
+
function jsonParam(value: AgentRunJsonObject | null | undefined): string | null {
|
|
1111
|
+
return value === undefined || value === null ? null : JSON.stringify(value);
|
|
1112
|
+
}
|
|
1113
|
+
|
|
1114
|
+
function parseStatus(value: string): AgentRunStatus | undefined {
|
|
1115
|
+
switch (value) {
|
|
1116
|
+
case "pending":
|
|
1117
|
+
case "running":
|
|
1118
|
+
case "success":
|
|
1119
|
+
case "error":
|
|
1120
|
+
case "interrupted":
|
|
1121
|
+
case "cancelled":
|
|
1122
|
+
case "timeout":
|
|
1123
|
+
return value;
|
|
1124
|
+
default:
|
|
1125
|
+
return undefined;
|
|
1126
|
+
}
|
|
1127
|
+
}
|
|
1128
|
+
|
|
1129
|
+
function parseSource(value: string | null): AgentRunSource {
|
|
1130
|
+
switch (value) {
|
|
1131
|
+
case "run_api":
|
|
1132
|
+
return "run_api";
|
|
1133
|
+
case "project_task":
|
|
1134
|
+
return "project_task";
|
|
1135
|
+
default:
|
|
1136
|
+
return "legacy";
|
|
1137
|
+
}
|
|
1138
|
+
}
|
|
1139
|
+
|
|
1140
|
+
function toJsonObject(value: unknown): AgentRunJsonObject | null {
|
|
1141
|
+
if (value === null || value === undefined) return null;
|
|
1142
|
+
if (typeof value !== "object" || Array.isArray(value)) return null;
|
|
1143
|
+
return value as AgentRunJsonObject;
|
|
1144
|
+
}
|
|
1145
|
+
|
|
1146
|
+
function rowToRecord(row: AgentRunRow): AgentRunRecord {
|
|
1147
|
+
return {
|
|
1148
|
+
runId: row.run_id,
|
|
1149
|
+
tenantId: row.tenant_id,
|
|
1150
|
+
workspaceId: row.workspace_id ?? undefined,
|
|
1151
|
+
projectId: row.project_id ?? undefined,
|
|
1152
|
+
assistantId: row.assistant_id,
|
|
1153
|
+
threadId: row.thread_id,
|
|
1154
|
+
status: parseStatus(row.status) ?? "error",
|
|
1155
|
+
source: parseSource(row.source),
|
|
1156
|
+
taskId: row.task_id ?? undefined,
|
|
1157
|
+
owner: row.owner,
|
|
1158
|
+
leaseExpiresAt:
|
|
1159
|
+
row.lease_expires_at === null ? null : row.lease_expires_at.getTime(),
|
|
1160
|
+
inputMessageIds: row.input_message_ids,
|
|
1161
|
+
attempts: row.attempts,
|
|
1162
|
+
error: row.error,
|
|
1163
|
+
cancelRequested: row.cancel_requested,
|
|
1164
|
+
rollbackRequested: row.rollback_requested,
|
|
1165
|
+
rollbackAttempts: row.rollback_attempts,
|
|
1166
|
+
rollbackError: row.rollback_error,
|
|
1167
|
+
input: toJsonObject(row.input),
|
|
1168
|
+
command: toJsonObject(row.command),
|
|
1169
|
+
config: toJsonObject(row.config),
|
|
1170
|
+
metadata: toJsonObject(row.metadata),
|
|
1171
|
+
baseCheckpointId: row.base_checkpoint_id,
|
|
1172
|
+
baseCheckpointCaptured: row.base_checkpoint_captured,
|
|
1173
|
+
capabilityRuntimeRevision: row.capability_runtime_revision,
|
|
1174
|
+
availableAt: row.available_at.getTime(),
|
|
1175
|
+
createdAt: row.created_at.getTime(),
|
|
1176
|
+
updatedAt: row.updated_at.getTime(),
|
|
1177
|
+
startedAt: row.started_at === null ? null : row.started_at.getTime(),
|
|
1178
|
+
finishedAt: row.finished_at === null ? null : row.finished_at.getTime(),
|
|
1179
|
+
};
|
|
1180
|
+
}
|
|
1181
|
+
|
|
1182
|
+
function reapedFromRow(row: AgentRunRow): ReapedAgentRun {
|
|
1183
|
+
return {
|
|
1184
|
+
runId: row.run_id,
|
|
1185
|
+
tenantId: row.tenant_id,
|
|
1186
|
+
assistantId: row.assistant_id,
|
|
1187
|
+
threadId: row.thread_id,
|
|
1188
|
+
workspaceId: row.workspace_id ?? undefined,
|
|
1189
|
+
projectId: row.project_id ?? undefined,
|
|
1190
|
+
source: parseSource(row.source),
|
|
1191
|
+
taskId: row.task_id ?? undefined,
|
|
1192
|
+
attempts: row.attempts,
|
|
1193
|
+
cancelRequested: row.cancel_requested,
|
|
1194
|
+
rollbackRequested: row.rollback_requested,
|
|
1195
|
+
baseCheckpointId: row.base_checkpoint_id,
|
|
1196
|
+
createdAt: row.created_at.getTime(),
|
|
1197
|
+
startedAt: row.started_at === null ? null : row.started_at.getTime(),
|
|
1198
|
+
leaseExpiresAt:
|
|
1199
|
+
row.lease_expires_at === null ? null : row.lease_expires_at.getTime(),
|
|
1200
|
+
};
|
|
1201
|
+
}
|
|
1202
|
+
|
|
1203
|
+
function isActiveRunViolation(error: unknown): boolean {
|
|
1204
|
+
if (typeof error !== "object" || error === null) return false;
|
|
1205
|
+
const { code, constraint } = error as {
|
|
1206
|
+
code?: unknown;
|
|
1207
|
+
constraint?: unknown;
|
|
1208
|
+
};
|
|
1209
|
+
return code === "23505" && constraint === "uq_lattice_agent_runs_one_running";
|
|
1210
|
+
}
|