From 1784fa742f8fb0e3b7259c9f298a591d6bd2b326 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Wed, 16 Sep 2026 22:02:05 +0200 Subject: [PATCH 1/2] feat(cloud-agent-next): emit point-in-time Cloud Agent open stock metric The outcome aggregate counts only turns whose terminal_at falls in a rolling window, so a tick with no settlement reports settled: 0. An operator cannot tell an idle system from one with work in flight. Add the missing read of the non-terminal population. A single grouped statement reads cloud_agent_session_runs filtered to status in (queued, accepted) and terminal_at is null, joined to retained cloud_agent_sessions, and emits one cloud_agent_open_stock record per tick with per-generation queued/accepted turns and sessions, the oldest age in row state and counts of rows missing the age timestamp. A database failure emits collectionStatus failed, never zeros. It reuses the outcome collector's generation expression, coverage label, retention cutoff and envelope pattern, and dispatches under the existing 3-minute cron in finally after the outcome collection so a query outage in one metric does not blank the other. Admission, retry, cancellation, cleanup and terminalization are untouched. Verified against a real local Worker scheduled tick over Postgres, plus Postgres tests for null timestamps, terminal exclusion, generation split, retention and the failure path. EXPLAIN (ANALYZE, BUFFERS) on a 125k-row representative dataset uses IDX_cloud_agent_session_runs_status_terminal and examines 12 rows. --- services/cloud-agent-next/src/server.test.ts | 29 +- services/cloud-agent-next/src/server.ts | 7 +- .../src/telemetry/open-stock.test.ts | 291 +++++++++++++++ .../src/telemetry/open-stock.ts | 165 +++++++++ .../src/telemetry/outcome-aggregate.ts | 2 +- .../test/postgres/open-stock.postgres.test.ts | 344 ++++++++++++++++++ 6 files changed, 835 insertions(+), 3 deletions(-) create mode 100644 services/cloud-agent-next/src/telemetry/open-stock.test.ts create mode 100644 services/cloud-agent-next/src/telemetry/open-stock.ts create mode 100644 services/cloud-agent-next/test/postgres/open-stock.postgres.test.ts diff --git a/services/cloud-agent-next/src/server.test.ts b/services/cloud-agent-next/src/server.test.ts index 5a34dd18fbe..539e8af1ea6 100644 --- a/services/cloud-agent-next/src/server.test.ts +++ b/services/cloud-agent-next/src/server.test.ts @@ -22,6 +22,7 @@ const { consumeCloudAgentReportBatchMock, removeExpiredCloudAgentReportDataMock, runCloudAgentOutcomeCollectionMock, + runCloudAgentOpenStockCollectionMock, requireCurrentSessionAccessMock, getPgDbMock, loggerWarnMock, @@ -30,6 +31,7 @@ const { consumeCloudAgentReportBatchMock: vi.fn().mockResolvedValue(undefined), removeExpiredCloudAgentReportDataMock: vi.fn().mockResolvedValue(undefined), runCloudAgentOutcomeCollectionMock: vi.fn().mockResolvedValue(undefined), + runCloudAgentOpenStockCollectionMock: vi.fn().mockResolvedValue(undefined), requireCurrentSessionAccessMock: vi.fn(), getPgDbMock: vi.fn(), loggerWarnMock: vi.fn(), @@ -92,6 +94,10 @@ vi.mock('./telemetry/outcome-aggregate.js', () => ({ runCloudAgentOutcomeCollection: runCloudAgentOutcomeCollectionMock, })); +vi.mock('./telemetry/open-stock.js', () => ({ + runCloudAgentOpenStockCollection: runCloudAgentOpenStockCollectionMock, +})); + vi.mock('./middleware/auth.js', () => ({ authMiddleware: vi.fn(), })); @@ -282,6 +288,7 @@ beforeEach(() => { consumeCloudAgentReportBatchMock.mockClear(); removeExpiredCloudAgentReportDataMock.mockClear(); runCloudAgentOutcomeCollectionMock.mockClear(); + runCloudAgentOpenStockCollectionMock.mockClear(); loggerWarnMock.mockClear(); getPgDbMock.mockReset(); requireCurrentSessionAccessMock.mockReset().mockResolvedValue({ @@ -394,9 +401,10 @@ describe('server background reporting', () => { expect(removeExpiredCloudAgentReportDataMock).toHaveBeenCalledTimes(1); expect(removeExpiredCloudAgentReportDataMock).toHaveBeenCalledWith(env); expect(runCloudAgentOutcomeCollectionMock).not.toHaveBeenCalled(); + expect(runCloudAgentOpenStockCollectionMock).not.toHaveBeenCalled(); }); - it('runs only the outcome collection on the 3-minute cron', async () => { + it('runs the outcome and open-stock collections on the 3-minute cron', async () => { const env = createEnv(); await worker.scheduled( @@ -406,9 +414,27 @@ describe('server background reporting', () => { expect(runCloudAgentOutcomeCollectionMock).toHaveBeenCalledTimes(1); expect(runCloudAgentOutcomeCollectionMock).toHaveBeenCalledWith(env); + expect(runCloudAgentOpenStockCollectionMock).toHaveBeenCalledTimes(1); + expect(runCloudAgentOpenStockCollectionMock).toHaveBeenCalledWith(env); expect(removeExpiredCloudAgentReportDataMock).not.toHaveBeenCalled(); }); + it('still runs the open-stock collection and preserves the outcome error when outcome collection rejects', async () => { + const env = createEnv(); + const outcomeError = new Error('outcome collection failed'); + runCloudAgentOutcomeCollectionMock.mockRejectedValueOnce(outcomeError); + + await expect( + worker.scheduled( + { cron: OUTCOME_AGGREGATE_CRON } as ScheduledController, + env as unknown as Env + ) + ).rejects.toBe(outcomeError); + + expect(runCloudAgentOpenStockCollectionMock).toHaveBeenCalledTimes(1); + expect(runCloudAgentOpenStockCollectionMock).toHaveBeenCalledWith(env); + }); + it('runs neither branch for an unrecognized cron and warns', async () => { const env = createEnv(); @@ -416,6 +442,7 @@ describe('server background reporting', () => { expect(removeExpiredCloudAgentReportDataMock).not.toHaveBeenCalled(); expect(runCloudAgentOutcomeCollectionMock).not.toHaveBeenCalled(); + expect(runCloudAgentOpenStockCollectionMock).not.toHaveBeenCalled(); expect(loggerWarnMock).toHaveBeenCalledWith( 'Cloud Agent scheduled handler received an unrecognized cron', { cron: '0 0 * * 0' } diff --git a/services/cloud-agent-next/src/server.ts b/services/cloud-agent-next/src/server.ts index 7a782a1b66c..ec425459d39 100644 --- a/services/cloud-agent-next/src/server.ts +++ b/services/cloud-agent-next/src/server.ts @@ -24,6 +24,7 @@ import { removeExpiredCloudAgentReportData, } from './telemetry/report-consumer.js'; import { runCloudAgentOutcomeCollection } from './telemetry/outcome-aggregate.js'; +import { runCloudAgentOpenStockCollection } from './telemetry/open-stock.js'; import { authMiddleware } from './middleware/auth.js'; import { balanceMiddleware } from './middleware/balance.js'; import { resolveTerminalWrapperClient } from './terminal/access.js'; @@ -1122,7 +1123,11 @@ export default { return; } if (controller.cron === OUTCOME_AGGREGATE_CRON) { - await runCloudAgentOutcomeCollection(env); + try { + await runCloudAgentOutcomeCollection(env); + } finally { + await runCloudAgentOpenStockCollection(env); + } return; } logger.warn('Cloud Agent scheduled handler received an unrecognized cron', { diff --git a/services/cloud-agent-next/src/telemetry/open-stock.test.ts b/services/cloud-agent-next/src/telemetry/open-stock.test.ts new file mode 100644 index 00000000000..7e19b134c2f --- /dev/null +++ b/services/cloud-agent-next/src/telemetry/open-stock.test.ts @@ -0,0 +1,291 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { eq, sql, type SQL } from 'drizzle-orm'; +import { getWorkerDb } from '@kilocode/db/client'; +import { cloud_agent_session_runs, cloud_agent_sessions } from '@kilocode/db/schema'; + +const { getPgDbMock, loggerMock } = vi.hoisted(() => ({ + getPgDbMock: vi.fn(), + loggerMock: { + withFields: vi.fn(), + info: vi.fn(), + error: vi.fn(), + warn: vi.fn(), + }, +})); + +vi.mock('../db/pg.js', () => ({ getPgDb: getPgDbMock })); +vi.mock('../logger.js', () => ({ logger: loggerMock })); + +import { COVERAGE_LABEL } from './outcome-aggregate.js'; +import { + OPEN_STOCK_CONTRACT_VERSION, + OPEN_STOCK_LIMITATIONS, + OPEN_STOCK_METRIC, + assembleOpenStock, + readOpenStock, + runCloudAgentOpenStockCollection, + type OpenStockQueryRow, +} from './open-stock.js'; + +const retentionCutoffIso = '2025-11-03T00:10:00.000Z'; +const now = '2026-02-01T00:10:00.000Z'; + +type QueryResult = OpenStockQueryRow[] | Error; + +function queryRow( + input: Partial & { + generation: OpenStockQueryRow['generation']; + status: OpenStockQueryRow['status']; + } +): OpenStockQueryRow { + return { + turns: 0, + sessions: 0, + oldestQueuedEpochMs: null, + oldestAcceptedEpochMs: null, + queuedMissingAgeTurns: 0, + acceptedMissingAgeTurns: 0, + ...input, + }; +} + +function makeDb(result: QueryResult) { + const state: { + fields?: Record; + where?: unknown; + groupBy?: unknown[]; + } = {}; + + const chain = { + from: () => chain, + innerJoin: () => chain, + where: (arg: unknown) => { + state.where = arg; + return chain; + }, + groupBy: (...args: unknown[]) => { + state.groupBy = args; + return chain; + }, + then(resolve: (value: unknown) => unknown, reject: (reason: unknown) => unknown): unknown { + if (result instanceof Error) return reject(result); + return resolve(result); + }, + }; + + const db = { + select: (fields: Record) => { + state.fields = fields; + return chain; + }, + }; + + return { db, state }; +} + +const usageDb = () => getWorkerDb('postgres://unused:unused@localhost:0/unused'); + +function renderSelect(fields: Record) { + return usageDb() + .select(fields as never) + .from(cloud_agent_session_runs) + .innerJoin( + cloud_agent_sessions, + eq( + cloud_agent_sessions.cloud_agent_session_id, + cloud_agent_session_runs.cloud_agent_session_id + ) + ) + .toSQL(); +} + +function renderWhere(where: unknown, groupBy: unknown[]) { + return usageDb() + .select({ marker: sql`1` }) + .from(cloud_agent_session_runs) + .innerJoin( + cloud_agent_sessions, + eq( + cloud_agent_sessions.cloud_agent_session_id, + cloud_agent_session_runs.cloud_agent_session_id + ) + ) + .where(where as SQL) + .groupBy(...(groupBy as SQL[])) + .toSQL(); +} + +beforeEach(() => { + vi.clearAllMocks(); + loggerMock.withFields.mockReturnValue(loggerMock); +}); + +describe('cloud agent open stock assembly', () => { + it('reports both generations as empty stock with null ages', () => { + const generations = assembleOpenStock([], now); + + expect(generations.map(entry => entry.generation)).toEqual(['legacy', 'control']); + for (const entry of generations) { + expect(entry).toEqual({ + generation: entry.generation, + queuedTurns: 0, + queuedSessions: 0, + acceptedTurns: 0, + acceptedSessions: 0, + oldestQueuedAgeMs: null, + oldestAcceptedAgeMs: null, + queuedMissingAgeTurns: 0, + acceptedMissingAgeTurns: 0, + }); + } + }); + + it('preserves queued counts and the missing-age count without inventing an age', () => { + const generations = assembleOpenStock( + [ + queryRow({ + generation: 'legacy', + status: 'queued', + turns: 3, + sessions: 2, + oldestQueuedEpochMs: null, + queuedMissingAgeTurns: 3, + }), + ], + now + ); + + const legacy = generations[0]; + expect(legacy.queuedTurns).toBe(3); + expect(legacy.queuedSessions).toBe(2); + expect(legacy.oldestQueuedAgeMs).toBeNull(); + expect(legacy.queuedMissingAgeTurns).toBe(3); + expect(legacy.acceptedTurns).toBe(0); + expect(legacy.oldestAcceptedAgeMs).toBeNull(); + }); + + it('derives queued and accepted ages from their own timestamps', () => { + const queuedEpochMs = Date.parse('2026-02-01T00:00:00.000Z'); + const acceptedEpochMs = Date.parse('2026-01-31T23:55:00.000Z'); + const generations = assembleOpenStock( + [ + queryRow({ + generation: 'legacy', + status: 'queued', + turns: 1, + sessions: 1, + oldestQueuedEpochMs: queuedEpochMs, + }), + queryRow({ + generation: 'legacy', + status: 'accepted', + turns: 2, + sessions: 2, + oldestAcceptedEpochMs: acceptedEpochMs, + }), + ], + now + ); + + const legacy = generations[0]; + expect(legacy.oldestQueuedAgeMs).toBe(Date.parse(now) - queuedEpochMs); + expect(legacy.oldestAcceptedAgeMs).toBe(Date.parse(now) - acceptedEpochMs); + expect(legacy.oldestAcceptedAgeMs).not.toBe(legacy.oldestQueuedAgeMs); + }); +}); + +describe('cloud agent open stock query shape', () => { + it('narrows to non-terminal queued/accepted runs of retained sessions', async () => { + const fake = makeDb([]); + await readOpenStock(fake.db as never, { retentionCutoff: retentionCutoffIso }); + + const fields = fake.state.fields as Record; + const selectSql = renderSelect(fields); + const whereSql = renderWhere(fake.state.where, fake.state.groupBy ?? []); + + expect(selectSql.params).toEqual([10, 'workspace_']); + expect(selectSql.sql).toMatch( + /left\("cloud_agent_session_runs"\."cloud_agent_session_id", \$1\) = \$2/ + ); + expect(selectSql.sql).toContain( + 'count(distinct "cloud_agent_session_runs"."cloud_agent_session_id"))::int' + ); + expect(selectSql.sql).toContain('extract(epoch from'); + expect(selectSql.sql).toContain('::double precision'); + expect(selectSql.sql).toContain( + `count(*) filter (where "cloud_agent_session_runs"."status" = 'queued' and "cloud_agent_session_runs"."queued_at" is null)` + ); + expect(selectSql.sql).toContain( + `count(*) filter (where "cloud_agent_session_runs"."status" = 'accepted' and "cloud_agent_session_runs"."dispatch_accepted_at" is null)` + ); + + expect(whereSql.sql).toContain('"status" in ($1, $2)'); + expect(whereSql.sql).toContain('"terminal_at" is null'); + expect(whereSql.sql).toContain('created_at" > $3'); + expect(whereSql.sql).toContain('group by 1, 2'); + expect(whereSql.params).toEqual(['queued', 'accepted', retentionCutoffIso]); + }); +}); + +describe('runCloudAgentOpenStockCollection emission', () => { + it('emits one complete record with the stock envelope and two generations', async () => { + getPgDbMock.mockReturnValue(makeDb([]).db); + + await runCloudAgentOpenStockCollection({} as never, new Date(now)); + + expect(loggerMock.info).toHaveBeenCalledTimes(1); + expect(loggerMock.error).not.toHaveBeenCalled(); + const record = loggerMock.withFields.mock.calls[0]?.[0] as Record; + expect(record.metric).toBe(OPEN_STOCK_METRIC); + expect(record.logTag).toBe(OPEN_STOCK_METRIC); + expect(record.contractVersion).toBe(OPEN_STOCK_CONTRACT_VERSION); + expect(record.service).toBe('cloud-agent-next'); + expect(record.environment).toBeNull(); + expect(record.evaluationId).toBe(`${OPEN_STOCK_METRIC}:${now}`); + expect(record.collectionTimestamp).toBe(now); + expect(typeof record.queryElapsedMs).toBe('number'); + expect(record.coverage).toBe(COVERAGE_LABEL); + expect(record.retentionCutoff).toBe(retentionCutoffIso); + expect(record.expectedGenerations).toEqual(['legacy', 'control']); + expect(record.limitations).toEqual([...OPEN_STOCK_LIMITATIONS]); + expect(record.collectionStatus).toBe('complete'); + expect(record.generations).toHaveLength(2); + expect(record.windowMinutes).toBeUndefined(); + expect(record.windowStart).toBeUndefined(); + expect(record.windowEnd).toBeUndefined(); + expect(record.reportingDelayMs).toBeUndefined(); + }); + + it('emits exactly one failed record with no generations when the query rejects', async () => { + getPgDbMock.mockReturnValue(makeDb(new Error('stock query failed')).db); + + await expect( + runCloudAgentOpenStockCollection({} as never, new Date(now)) + ).resolves.toBeUndefined(); + + expect(loggerMock.error).toHaveBeenCalledTimes(1); + expect(loggerMock.info).not.toHaveBeenCalled(); + const record = loggerMock.withFields.mock.calls.at(-1)?.[0] as Record; + expect(record.collectionStatus).toBe('failed'); + expect(record.failureKind).toBe('db_query_failed'); + expect(record.generations).toBeUndefined(); + expect(record.limitations).toEqual([...OPEN_STOCK_LIMITATIONS]); + }); + + it('emits a failed record when the database binding throws', async () => { + getPgDbMock.mockImplementation(() => { + throw new Error('HYPERDRIVE binding missing'); + }); + + await expect( + runCloudAgentOpenStockCollection({} as never, new Date(now)) + ).resolves.toBeUndefined(); + + expect(loggerMock.error).toHaveBeenCalledTimes(1); + expect(loggerMock.info).not.toHaveBeenCalled(); + const record = loggerMock.withFields.mock.calls.at(-1)?.[0] as Record; + expect(record.collectionStatus).toBe('failed'); + expect(record.failureKind).toBe('db_query_failed'); + expect(record.generations).toBeUndefined(); + }); +}); diff --git a/services/cloud-agent-next/src/telemetry/open-stock.ts b/services/cloud-agent-next/src/telemetry/open-stock.ts new file mode 100644 index 00000000000..2a3ed584708 --- /dev/null +++ b/services/cloud-agent-next/src/telemetry/open-stock.ts @@ -0,0 +1,165 @@ +import { and, eq, gt, inArray, isNull, sql } from 'drizzle-orm'; +import type { WorkerDb } from '@kilocode/db/client'; +import { + cloud_agent_session_runs, + cloud_agent_sessions, + type CloudAgentSessionRunStatus, +} from '@kilocode/db/schema'; +import { getPgDb } from '../db/pg.js'; +import type { Env } from '../types.js'; +import { logger } from '../logger.js'; +import { retentionCutoff } from './report-store.js'; +import { + COVERAGE_LABEL, + EXPECTED_GENERATIONS, + generationExpression, + type OutcomeGeneration, +} from './outcome-aggregate.js'; + +export const OPEN_STOCK_METRIC = 'cloud_agent_open_stock'; +export const OPEN_STOCK_CONTRACT_VERSION = 1; +export const OPEN_STOCK_LIMITATIONS = [ + 'row_state_age_not_stall', + 'accepted_not_sandbox_alive', + 'lost_reports_undercount', + 'age_alerting_deferred', +] as const; + +const OPEN_STOCK_SERVICE = 'cloud-agent-next'; + +export type OpenStockQueryRow = { + generation: OutcomeGeneration; + status: CloudAgentSessionRunStatus; + turns: number; + sessions: number; + oldestQueuedEpochMs: number | null; + oldestAcceptedEpochMs: number | null; + queuedMissingAgeTurns: number; + acceptedMissingAgeTurns: number; +}; + +export type GenerationOpenStock = { + generation: OutcomeGeneration; + queuedTurns: number; + queuedSessions: number; + acceptedTurns: number; + acceptedSessions: number; + oldestQueuedAgeMs: number | null; + oldestAcceptedAgeMs: number | null; + queuedMissingAgeTurns: number; + acceptedMissingAgeTurns: number; +}; + +export async function readOpenStock( + db: WorkerDb, + input: { retentionCutoff: string } +): Promise { + const generation = generationExpression(cloud_agent_session_runs.cloud_agent_session_id); + return db + .select({ + generation, + status: cloud_agent_session_runs.status, + turns: sql`count(*)::int`, + sessions: sql`(count(distinct ${cloud_agent_session_runs.cloud_agent_session_id}))::int`, + oldestQueuedEpochMs: sql< + number | null + >`(extract(epoch from (min(${cloud_agent_session_runs.queued_at}) filter (where ${cloud_agent_session_runs.status} = 'queued'))) * 1000)::double precision`, + oldestAcceptedEpochMs: sql< + number | null + >`(extract(epoch from (min(${cloud_agent_session_runs.dispatch_accepted_at}) filter (where ${cloud_agent_session_runs.status} = 'accepted'))) * 1000)::double precision`, + queuedMissingAgeTurns: sql`(count(*) filter (where ${cloud_agent_session_runs.status} = 'queued' and ${cloud_agent_session_runs.queued_at} is null))::int`, + acceptedMissingAgeTurns: sql`(count(*) filter (where ${cloud_agent_session_runs.status} = 'accepted' and ${cloud_agent_session_runs.dispatch_accepted_at} is null))::int`, + }) + .from(cloud_agent_session_runs) + .innerJoin( + cloud_agent_sessions, + eq( + cloud_agent_sessions.cloud_agent_session_id, + cloud_agent_session_runs.cloud_agent_session_id + ) + ) + .where( + and( + inArray(cloud_agent_session_runs.status, ['queued', 'accepted']), + isNull(cloud_agent_session_runs.terminal_at), + gt(cloud_agent_sessions.created_at, input.retentionCutoff) + ) + ) + .groupBy(sql`1`, sql`2`); +} + +export function assembleOpenStock(rows: OpenStockQueryRow[], now: string): GenerationOpenStock[] { + const nowMs = Date.parse(now); + return EXPECTED_GENERATIONS.map(generation => { + const queued = rows.find(row => row.generation === generation && row.status === 'queued'); + const accepted = rows.find(row => row.generation === generation && row.status === 'accepted'); + const oldestQueuedEpochMs = queued?.oldestQueuedEpochMs ?? null; + const oldestAcceptedEpochMs = accepted?.oldestAcceptedEpochMs ?? null; + return { + generation, + queuedTurns: queued?.turns ?? 0, + queuedSessions: queued?.sessions ?? 0, + acceptedTurns: accepted?.turns ?? 0, + acceptedSessions: accepted?.sessions ?? 0, + oldestQueuedAgeMs: oldestQueuedEpochMs === null ? null : nowMs - oldestQueuedEpochMs, + oldestAcceptedAgeMs: oldestAcceptedEpochMs === null ? null : nowMs - oldestAcceptedEpochMs, + queuedMissingAgeTurns: queued?.queuedMissingAgeTurns ?? 0, + acceptedMissingAgeTurns: accepted?.acceptedMissingAgeTurns ?? 0, + }; + }); +} + +function openStockEnvelope(input: { + collectionTimestamp: string; + queryElapsedMs: number; + retentionCutoff: string; +}): Record { + return { + metric: OPEN_STOCK_METRIC, + logTag: OPEN_STOCK_METRIC, + contractVersion: OPEN_STOCK_CONTRACT_VERSION, + service: OPEN_STOCK_SERVICE, + environment: null, + evaluationId: `${OPEN_STOCK_METRIC}:${input.collectionTimestamp}`, + collectionTimestamp: input.collectionTimestamp, + queryElapsedMs: input.queryElapsedMs, + coverage: COVERAGE_LABEL, + retentionCutoff: input.retentionCutoff, + expectedGenerations: [...EXPECTED_GENERATIONS], + limitations: [...OPEN_STOCK_LIMITATIONS], + }; +} + +export async function runCloudAgentOpenStockCollection(env: Env, now = new Date()): Promise { + const collectionTimestamp = now.toISOString(); + const retentionCutoffIso = retentionCutoff(collectionTimestamp); + const evaluatedAt = Date.now(); + + try { + const db = getPgDb(env); + const rows = await readOpenStock(db, { retentionCutoff: retentionCutoffIso }); + logger + .withFields({ + ...openStockEnvelope({ + collectionTimestamp, + queryElapsedMs: Date.now() - evaluatedAt, + retentionCutoff: retentionCutoffIso, + }), + collectionStatus: 'complete', + generations: assembleOpenStock(rows, collectionTimestamp), + }) + .info('Cloud Agent open stock'); + } catch { + logger + .withFields({ + ...openStockEnvelope({ + collectionTimestamp, + queryElapsedMs: Date.now() - evaluatedAt, + retentionCutoff: retentionCutoffIso, + }), + collectionStatus: 'failed', + failureKind: 'db_query_failed', + }) + .error('Cloud Agent open stock failed'); + } +} diff --git a/services/cloud-agent-next/src/telemetry/outcome-aggregate.ts b/services/cloud-agent-next/src/telemetry/outcome-aggregate.ts index baaf04692b8..0e8e7366e50 100644 --- a/services/cloud-agent-next/src/telemetry/outcome-aggregate.ts +++ b/services/cloud-agent-next/src/telemetry/outcome-aggregate.ts @@ -90,7 +90,7 @@ type DatabaseTransaction = Parameters[0]>[0] const unknownResponsibilityCondition: SQL = sql`(${cloud_agent_session_runs.failure_responsibility} is null or ${cloud_agent_session_runs.failure_responsibility} not in ('platform', 'user'))`; -function generationExpression(sessionId: AnyColumn): SQL { +export function generationExpression(sessionId: AnyColumn): SQL { return sql`case when left(${sessionId}, ${CONTROL_PLANE_SESSION_PREFIX.length}) = ${CONTROL_PLANE_SESSION_PREFIX} then 'control' else 'legacy' end`; } diff --git a/services/cloud-agent-next/test/postgres/open-stock.postgres.test.ts b/services/cloud-agent-next/test/postgres/open-stock.postgres.test.ts new file mode 100644 index 00000000000..a2463c5e8df --- /dev/null +++ b/services/cloud-agent-next/test/postgres/open-stock.postgres.test.ts @@ -0,0 +1,344 @@ +import { randomUUID } from 'node:crypto'; +import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from 'vitest'; +import { createDrizzleClient, getWorkerDb, type WorkerDb } from '@kilocode/db/client'; +import { eq } from 'drizzle-orm'; +import { cloud_agent_session_runs, cloud_agent_sessions } from '@kilocode/db/schema'; +import { logger } from '../../src/logger.js'; +import { + OPEN_STOCK_METRIC, + assembleOpenStock, + readOpenStock, + runCloudAgentOpenStockCollection, + type OpenStockQueryRow, +} from '../../src/telemetry/open-stock.js'; + +const { getPgDbMock } = vi.hoisted(() => ({ getPgDbMock: vi.fn() })); + +vi.mock('../../src/db/pg.js', () => ({ getPgDb: getPgDbMock })); + +const STOCK_RETENTION_CUTOFF = '2020-01-01T00:00:00.000Z'; +const SESSION_CREATED_AT = '2026-02-01T00:00:00.000Z'; + +let connectionString: string; +let reader: WorkerDb; +let writer: ReturnType; +const trackedSessionIds: string[] = []; + +function requirePostgresUrl(): string { + const url = process.env.POSTGRES_URL; + if (!url) throw new Error('Set POSTGRES_URL to a migrated test database'); + return url; +} + +const uniqueSessionId = (plane: 'agent' | 'workspace') => `${plane}_${randomUUID()}`; +const uniqueMessageId = () => `msg_${randomUUID().replace(/-/g, '')}`; + +async function insertSession(sessionId: string, createdAt: string): Promise { + trackedSessionIds.push(sessionId); + const initialMessageId = uniqueMessageId(); + await writer.db.insert(cloud_agent_sessions).values({ + cloud_agent_session_id: sessionId, + kilo_session_id: `ses_${randomUUID().replace(/-/g, '').slice(0, 26)}`, + initial_message_id: initialMessageId, + created_at: createdAt, + }); + return initialMessageId; +} + +async function insertRun(values: { + cloudAgentSessionId: string; + messageId?: string; + status: 'queued' | 'accepted' | 'completed' | 'failed'; + queuedAt?: string | null; + dispatchAcceptedAt?: string | null; + terminalAt?: string | null; +}): Promise { + await writer.db.insert(cloud_agent_session_runs).values({ + cloud_agent_session_id: values.cloudAgentSessionId, + message_id: values.messageId ?? uniqueMessageId(), + status: values.status, + queued_at: values.queuedAt ?? null, + dispatch_accepted_at: values.dispatchAcceptedAt ?? null, + terminal_at: values.terminalAt ?? null, + }); +} + +async function readStock(cutoff = STOCK_RETENTION_CUTOFF): Promise { + return readOpenStock(reader, { retentionCutoff: cutoff }); +} + +function stockRow( + rows: OpenStockQueryRow[], + generation: OpenStockQueryRow['generation'], + status: OpenStockQueryRow['status'] +): OpenStockQueryRow | undefined { + return rows.find(row => row.generation === generation && row.status === status); +} + +function failAtSelect( + db: WorkerDb, + error: Error +): { db: WorkerDb; state: { selectCalls: number } } { + const state = { selectCalls: 0 }; + const instrumented = new Proxy(db, { + get(target, property) { + if (property !== 'select') return Reflect.get(target, property, target); + return () => { + state.selectCalls += 1; + throw error; + }; + }, + }) as WorkerDb; + return { db: instrumented, state }; +} + +beforeAll(() => { + connectionString = requirePostgresUrl(); + reader = getWorkerDb(connectionString); + writer = createDrizzleClient({ connectionString, ssl: false }); +}); + +afterEach(async () => { + const ids = trackedSessionIds.splice(0); + if (ids.length === 0) return; + for (const id of ids) { + await writer.db + .delete(cloud_agent_sessions) + .where(eq(cloud_agent_sessions.cloud_agent_session_id, id)); + } +}); + +afterAll(async () => { + await writer.pool.end(); + const readerPool = (reader as unknown as { $client?: { end: () => Promise } }).$client; + await readerPool?.end(); +}); + +describe('cloud agent open stock against PostgreSQL', () => { + it('counts a null-age queued row beside a timestamped one in the same session', async () => { + const sessionId = uniqueSessionId('agent'); + await insertSession(sessionId, SESSION_CREATED_AT); + const queuedAt = '2026-02-01T00:05:00.000Z'; + await insertRun({ cloudAgentSessionId: sessionId, status: 'queued', queuedAt }); + await insertRun({ cloudAgentSessionId: sessionId, status: 'queued', queuedAt: null }); + + const rows = await readStock(); + const queued = stockRow(rows, 'legacy', 'queued'); + expect(queued?.turns).toBe(2); + expect(queued?.sessions).toBe(1); + expect(queued?.oldestQueuedEpochMs).toBe(Date.parse(queuedAt)); + expect(queued?.queuedMissingAgeTurns).toBe(1); + }); + + it('keeps an all-null queued group and an all-null accepted group with null ages', async () => { + const queuedSession = uniqueSessionId('agent'); + await insertSession(queuedSession, SESSION_CREATED_AT); + await insertRun({ cloudAgentSessionId: queuedSession, status: 'queued', queuedAt: null }); + + const acceptedSession = uniqueSessionId('agent'); + await insertSession(acceptedSession, SESSION_CREATED_AT); + await insertRun({ + cloudAgentSessionId: acceptedSession, + status: 'accepted', + dispatchAcceptedAt: null, + }); + + const rows = await readStock(); + const queued = stockRow(rows, 'legacy', 'queued'); + expect(queued?.turns).toBe(1); + expect(queued?.sessions).toBe(1); + expect(queued?.oldestQueuedEpochMs).toBeNull(); + expect(queued?.queuedMissingAgeTurns).toBe(1); + + const accepted = stockRow(rows, 'legacy', 'accepted'); + expect(accepted?.turns).toBe(1); + expect(accepted?.sessions).toBe(1); + expect(accepted?.oldestAcceptedEpochMs).toBeNull(); + expect(accepted?.acceptedMissingAgeTurns).toBe(1); + + const assembled = assembleOpenStock(rows, '2026-02-01T00:10:00.000Z'); + const legacy = assembled.find(entry => entry.generation === 'legacy'); + expect(legacy?.oldestQueuedAgeMs).toBeNull(); + expect(legacy?.oldestAcceptedAgeMs).toBeNull(); + }); + + it('excludes terminal statuses and non-terminal rows that carry a terminal_at', async () => { + const sessionId = uniqueSessionId('agent'); + await insertSession(sessionId, SESSION_CREATED_AT); + const terminalAt = '2026-02-01T00:05:00.000Z'; + await insertRun({ cloudAgentSessionId: sessionId, status: 'completed', terminalAt }); + await insertRun({ cloudAgentSessionId: sessionId, status: 'failed', terminalAt }); + await insertRun({ + cloudAgentSessionId: sessionId, + status: 'queued', + queuedAt: terminalAt, + terminalAt, + }); + await insertRun({ + cloudAgentSessionId: sessionId, + status: 'accepted', + dispatchAcceptedAt: terminalAt, + terminalAt, + }); + + expect(await readStock()).toEqual([]); + }); + + it('splits generations by the workspace_ prefix', async () => { + const sessionIds = [ + uniqueSessionId('workspace'), + uniqueSessionId('agent'), + 'workspace', + `Workspace_${randomUUID()}`, + ]; + for (const sessionId of sessionIds) { + await insertSession(sessionId, SESSION_CREATED_AT); + await insertRun({ + cloudAgentSessionId: sessionId, + status: 'queued', + queuedAt: SESSION_CREATED_AT, + }); + } + + const rows = await readStock(); + expect(stockRow(rows, 'legacy', 'queued')?.turns).toBe(3); + expect(stockRow(rows, 'legacy', 'queued')?.sessions).toBe(3); + expect(stockRow(rows, 'control', 'queued')?.turns).toBe(1); + expect(stockRow(rows, 'control', 'queued')?.sessions).toBe(1); + }); + + it('keeps the queued and accepted ages separate with discriminating timestamps', async () => { + const sessionId = uniqueSessionId('agent'); + await insertSession(sessionId, SESSION_CREATED_AT); + const queuedAt = '2026-02-01T00:00:00.000Z'; + const acceptedQueuedAt = '2026-01-01T00:00:00.000Z'; + const dispatchAcceptedAt = '2026-02-01T00:05:00.000Z'; + await insertRun({ cloudAgentSessionId: sessionId, status: 'queued', queuedAt }); + await insertRun({ + cloudAgentSessionId: sessionId, + status: 'accepted', + queuedAt: acceptedQueuedAt, + dispatchAcceptedAt, + }); + + const rows = await readStock(); + expect(stockRow(rows, 'legacy', 'queued')?.turns).toBe(1); + expect(stockRow(rows, 'legacy', 'accepted')?.turns).toBe(1); + expect(stockRow(rows, 'legacy', 'queued')?.oldestQueuedEpochMs).toBe(Date.parse(queuedAt)); + expect(stockRow(rows, 'legacy', 'accepted')?.oldestAcceptedEpochMs).toBe( + Date.parse(dispatchAcceptedAt) + ); + }); + + it('counts two open turns in one session as two turns and one session', async () => { + const sessionId = uniqueSessionId('agent'); + await insertSession(sessionId, SESSION_CREATED_AT); + await insertRun({ + cloudAgentSessionId: sessionId, + status: 'queued', + queuedAt: '2026-02-01T00:00:00.000Z', + }); + await insertRun({ + cloudAgentSessionId: sessionId, + status: 'queued', + queuedAt: '2026-02-01T00:01:00.000Z', + }); + + const rows = await readStock(); + expect(stockRow(rows, 'legacy', 'queued')?.turns).toBe(2); + expect(stockRow(rows, 'legacy', 'queued')?.sessions).toBe(1); + expect(stockRow(rows, 'legacy', 'queued')?.oldestQueuedEpochMs).toBe( + Date.parse('2026-02-01T00:00:00.000Z') + ); + }); + + it('retains only sessions created strictly after the cutoff', async () => { + const cutoff = '2026-01-15T00:00:00.000Z'; + const atCutoff = uniqueSessionId('agent'); + const older = uniqueSessionId('agent'); + const newer = uniqueSessionId('agent'); + await insertSession(atCutoff, cutoff); + await insertSession(older, new Date(Date.parse(cutoff) - 1000).toISOString()); + await insertSession(newer, new Date(Date.parse(cutoff) + 1000).toISOString()); + for (const sessionId of [atCutoff, older, newer]) { + await insertRun({ + cloudAgentSessionId: sessionId, + status: 'queued', + queuedAt: SESSION_CREATED_AT, + }); + } + + const rows = await readStock(cutoff); + expect(stockRow(rows, 'legacy', 'queued')?.turns).toBe(1); + expect(stockRow(rows, 'legacy', 'queued')?.sessions).toBe(1); + }); + + it('returns exact finite epoch milliseconds and numeric counts from the raw reader', async () => { + const sessionId = uniqueSessionId('agent'); + await insertSession(sessionId, SESSION_CREATED_AT); + const queuedAt = '2026-02-01T00:02:00.000Z'; + const dispatchAcceptedAt = '2026-02-01T00:07:00.000Z'; + await insertRun({ cloudAgentSessionId: sessionId, status: 'queued', queuedAt }); + await insertRun({ + cloudAgentSessionId: sessionId, + status: 'accepted', + dispatchAcceptedAt, + }); + + const controlSessionId = uniqueSessionId('workspace'); + await insertSession(controlSessionId, SESSION_CREATED_AT); + await insertRun({ + cloudAgentSessionId: controlSessionId, + status: 'accepted', + dispatchAcceptedAt: null, + }); + + const rows = await readStock(); + const queued = stockRow(rows, 'legacy', 'queued'); + expect(typeof queued?.turns).toBe('number'); + expect(typeof queued?.sessions).toBe('number'); + expect(typeof queued?.oldestQueuedEpochMs).toBe('number'); + expect(Number.isFinite(queued?.oldestQueuedEpochMs)).toBe(true); + expect(queued?.oldestQueuedEpochMs).toBe(Date.parse(queuedAt)); + + const accepted = stockRow(rows, 'legacy', 'accepted'); + expect(typeof accepted?.oldestAcceptedEpochMs).toBe('number'); + expect(Number.isFinite(accepted?.oldestAcceptedEpochMs)).toBe(true); + expect(accepted?.oldestAcceptedEpochMs).toBe(Date.parse(dispatchAcceptedAt)); + + const controlAccepted = stockRow(rows, 'control', 'accepted'); + expect(controlAccepted?.turns).toBe(1); + expect(controlAccepted?.oldestAcceptedEpochMs).toBeNull(); + expect(controlAccepted?.acceptedMissingAgeTurns).toBe(1); + }); + + it('emits a failed record instead of zeros when the query is injected to fail', async () => { + const injected = new Error('injected open-stock select failure'); + const instrumented = failAtSelect(reader, injected); + + const errorMock = vi.fn(); + const infoMock = vi.fn(); + const withFieldsSpy = vi + .spyOn(logger, 'withFields') + .mockReturnValue({ info: infoMock, error: errorMock } as never); + getPgDbMock.mockReturnValue(instrumented.db); + + let records: Record[] = []; + try { + await runCloudAgentOpenStockCollection({} as never, new Date('2026-02-01T00:10:00.000Z')); + records = withFieldsSpy.mock.calls.map(([fields]) => fields as Record); + } finally { + getPgDbMock.mockReset(); + withFieldsSpy.mockRestore(); + } + + expect(instrumented.state.selectCalls).toBe(1); + expect(records).toHaveLength(1); + expect(records[0].metric).toBe(OPEN_STOCK_METRIC); + expect(records[0].collectionStatus).toBe('failed'); + expect(records[0].failureKind).toBe('db_query_failed'); + expect(records[0].generations).toBeUndefined(); + expect(errorMock).toHaveBeenCalledTimes(1); + expect(infoMock).not.toHaveBeenCalled(); + }); +}); From 8748e10c1a4e26d825d95d5b2d8b43154984d4a2 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Wed, 16 Sep 2026 22:17:31 +0200 Subject: [PATCH 2/2] test(cloud-agent-next): isolate open stock postgres cases from pre-existing runs The suite aggregated the whole retained population, so its exact counts and its empty-result assertion only held against a database with no open runs. Eight of nine cases failed once the shared local database held seven open runs. Scope every case through the reader's own retention predicate: fixture sessions are created above the real retained population and reads use a cutoff beneath them, so pre-existing rows are excluded by the same predicate the metric uses. The retention case derives its own cutoffs around the fixture base. No production code or reader signature changes. --- .../test/postgres/open-stock.postgres.test.ts | 38 +++++++++++-------- 1 file changed, 23 insertions(+), 15 deletions(-) diff --git a/services/cloud-agent-next/test/postgres/open-stock.postgres.test.ts b/services/cloud-agent-next/test/postgres/open-stock.postgres.test.ts index a2463c5e8df..ce24d3b9fca 100644 --- a/services/cloud-agent-next/test/postgres/open-stock.postgres.test.ts +++ b/services/cloud-agent-next/test/postgres/open-stock.postgres.test.ts @@ -16,8 +16,16 @@ const { getPgDbMock } = vi.hoisted(() => ({ getPgDbMock: vi.fn() })); vi.mock('../../src/db/pg.js', () => ({ getPgDb: getPgDbMock })); -const STOCK_RETENTION_CUTOFF = '2020-01-01T00:00:00.000Z'; -const SESSION_CREATED_AT = '2026-02-01T00:00:00.000Z'; +// Fixture timestamps sit above the real retained reporting population, so +// pre-existing rows are excluded by the reader's own retention predicate rather +// than by an empty table. Exact counts assume nothing else creates sessions in +// this reserved post-2095 range and that two runs of this suite do not share a +// database concurrently. These fixtures are cleaned up per test, but the +// terminal-exclusion fixture still carries a real 2026 terminal_at, so normal +// serialized execution with per-test cleanup is what keeps the sibling outcome +// suite unaffected. +const FIXTURE_RETENTION_CUTOFF = '2095-01-01T00:00:00.000Z'; +const FIXTURE_CREATED_AT = '2096-02-01T00:00:00.000Z'; let connectionString: string; let reader: WorkerDb; @@ -63,7 +71,7 @@ async function insertRun(values: { }); } -async function readStock(cutoff = STOCK_RETENTION_CUTOFF): Promise { +async function readStock(cutoff = FIXTURE_RETENTION_CUTOFF): Promise { return readOpenStock(reader, { retentionCutoff: cutoff }); } @@ -117,7 +125,7 @@ afterAll(async () => { describe('cloud agent open stock against PostgreSQL', () => { it('counts a null-age queued row beside a timestamped one in the same session', async () => { const sessionId = uniqueSessionId('agent'); - await insertSession(sessionId, SESSION_CREATED_AT); + await insertSession(sessionId, FIXTURE_CREATED_AT); const queuedAt = '2026-02-01T00:05:00.000Z'; await insertRun({ cloudAgentSessionId: sessionId, status: 'queued', queuedAt }); await insertRun({ cloudAgentSessionId: sessionId, status: 'queued', queuedAt: null }); @@ -132,11 +140,11 @@ describe('cloud agent open stock against PostgreSQL', () => { it('keeps an all-null queued group and an all-null accepted group with null ages', async () => { const queuedSession = uniqueSessionId('agent'); - await insertSession(queuedSession, SESSION_CREATED_AT); + await insertSession(queuedSession, FIXTURE_CREATED_AT); await insertRun({ cloudAgentSessionId: queuedSession, status: 'queued', queuedAt: null }); const acceptedSession = uniqueSessionId('agent'); - await insertSession(acceptedSession, SESSION_CREATED_AT); + await insertSession(acceptedSession, FIXTURE_CREATED_AT); await insertRun({ cloudAgentSessionId: acceptedSession, status: 'accepted', @@ -164,7 +172,7 @@ describe('cloud agent open stock against PostgreSQL', () => { it('excludes terminal statuses and non-terminal rows that carry a terminal_at', async () => { const sessionId = uniqueSessionId('agent'); - await insertSession(sessionId, SESSION_CREATED_AT); + await insertSession(sessionId, FIXTURE_CREATED_AT); const terminalAt = '2026-02-01T00:05:00.000Z'; await insertRun({ cloudAgentSessionId: sessionId, status: 'completed', terminalAt }); await insertRun({ cloudAgentSessionId: sessionId, status: 'failed', terminalAt }); @@ -192,11 +200,11 @@ describe('cloud agent open stock against PostgreSQL', () => { `Workspace_${randomUUID()}`, ]; for (const sessionId of sessionIds) { - await insertSession(sessionId, SESSION_CREATED_AT); + await insertSession(sessionId, FIXTURE_CREATED_AT); await insertRun({ cloudAgentSessionId: sessionId, status: 'queued', - queuedAt: SESSION_CREATED_AT, + queuedAt: FIXTURE_CREATED_AT, }); } @@ -209,7 +217,7 @@ describe('cloud agent open stock against PostgreSQL', () => { it('keeps the queued and accepted ages separate with discriminating timestamps', async () => { const sessionId = uniqueSessionId('agent'); - await insertSession(sessionId, SESSION_CREATED_AT); + await insertSession(sessionId, FIXTURE_CREATED_AT); const queuedAt = '2026-02-01T00:00:00.000Z'; const acceptedQueuedAt = '2026-01-01T00:00:00.000Z'; const dispatchAcceptedAt = '2026-02-01T00:05:00.000Z'; @@ -232,7 +240,7 @@ describe('cloud agent open stock against PostgreSQL', () => { it('counts two open turns in one session as two turns and one session', async () => { const sessionId = uniqueSessionId('agent'); - await insertSession(sessionId, SESSION_CREATED_AT); + await insertSession(sessionId, FIXTURE_CREATED_AT); await insertRun({ cloudAgentSessionId: sessionId, status: 'queued', @@ -253,7 +261,7 @@ describe('cloud agent open stock against PostgreSQL', () => { }); it('retains only sessions created strictly after the cutoff', async () => { - const cutoff = '2026-01-15T00:00:00.000Z'; + const cutoff = '2096-01-15T00:00:00.000Z'; const atCutoff = uniqueSessionId('agent'); const older = uniqueSessionId('agent'); const newer = uniqueSessionId('agent'); @@ -264,7 +272,7 @@ describe('cloud agent open stock against PostgreSQL', () => { await insertRun({ cloudAgentSessionId: sessionId, status: 'queued', - queuedAt: SESSION_CREATED_AT, + queuedAt: FIXTURE_CREATED_AT, }); } @@ -275,7 +283,7 @@ describe('cloud agent open stock against PostgreSQL', () => { it('returns exact finite epoch milliseconds and numeric counts from the raw reader', async () => { const sessionId = uniqueSessionId('agent'); - await insertSession(sessionId, SESSION_CREATED_AT); + await insertSession(sessionId, FIXTURE_CREATED_AT); const queuedAt = '2026-02-01T00:02:00.000Z'; const dispatchAcceptedAt = '2026-02-01T00:07:00.000Z'; await insertRun({ cloudAgentSessionId: sessionId, status: 'queued', queuedAt }); @@ -286,7 +294,7 @@ describe('cloud agent open stock against PostgreSQL', () => { }); const controlSessionId = uniqueSessionId('workspace'); - await insertSession(controlSessionId, SESSION_CREATED_AT); + await insertSession(controlSessionId, FIXTURE_CREATED_AT); await insertRun({ cloudAgentSessionId: controlSessionId, status: 'accepted',