Skip to content

src/storage.ts

236 lines · 6.4 KB · TypeScript
Raw
import { kody, packageStorage } from 'kody:runtime'
import { readFetchQuota, type FetchQuota } from './quota.ts'
import type { Disposition } from './status.ts'

/** Decisions and ship events expire after at most 7 days. */
export const DEFAULT_TTL_DAYS = 7
export const MAX_TTL_DAYS = 7

export const decisionKinds = ['thread', 'check', 'step'] as const
export type DecisionKind = (typeof decisionKinds)[number]

export type Decision = {
	prKey: string
	kind: DecisionKind
	target: string
	disposition: Disposition
	reason: string
	headSha: string | null
	agentId: string | null
	createdAt: string
	expiresAt: string
}

export type ShipEventKind = 'bugbot_trigger' | 'invalid_reply' | 'merge' | 'summary_sent'

export type ShipEvent = {
	prKey: string
	kind: ShipEventKind
	ref: string
	data: Record<string, unknown>
	createdAt: string
	expiresAt: string
}

type Row = Record<string, unknown>

function rowsOf(result: unknown): Array<Row> {
	if (!result) return []
	if (Array.isArray(result)) return result as Array<Row>
	const rows = (result as { rows?: unknown }).rows
	return Array.isArray(rows) ? (rows as Array<Row>) : []
}

async function sql(query: string, bindings: Array<unknown> = []) {
	return await packageStorage().sql(query, bindings)
}

let schemaReady = false

export async function initSchema() {
	if (schemaReady) return
	await sql(`
		CREATE TABLE IF NOT EXISTS decisions (
			pr_key TEXT NOT NULL,
			kind TEXT NOT NULL,
			target TEXT NOT NULL,
			disposition TEXT NOT NULL,
			reason TEXT NOT NULL,
			head_sha TEXT,
			agent_id TEXT,
			created_at TEXT NOT NULL,
			expires_at TEXT NOT NULL,
			PRIMARY KEY (pr_key, kind, target)
		)
	`)
	await sql(`
		CREATE TABLE IF NOT EXISTS events (
			pr_key TEXT NOT NULL,
			kind TEXT NOT NULL,
			ref TEXT NOT NULL,
			data_json TEXT NOT NULL,
			created_at TEXT NOT NULL,
			expires_at TEXT NOT NULL,
			PRIMARY KEY (pr_key, kind, ref)
		)
	`)
	await sql(`
		CREATE TABLE IF NOT EXISTS usage_cache (
			key TEXT PRIMARY KEY,
			json TEXT NOT NULL,
			fetched_at INTEGER NOT NULL
		)
	`)
	schemaReady = true
}

async function purgeExpired(now: string) {
	await sql('DELETE FROM decisions WHERE expires_at <= ?', [now])
	await sql('DELETE FROM events WHERE expires_at <= ?', [now])
}

export function expiryFrom(nowMs: number, ttlDays: number) {
	return new Date(nowMs + ttlDays * 24 * 60 * 60 * 1000).toISOString()
}

function toDecision(row: Row): Decision {
	return {
		prKey: String(row.pr_key),
		kind: String(row.kind) as DecisionKind,
		target: String(row.target),
		disposition: String(row.disposition) as Disposition,
		reason: String(row.reason),
		headSha: row.head_sha == null ? null : String(row.head_sha),
		agentId: row.agent_id == null ? null : String(row.agent_id),
		createdAt: String(row.created_at),
		expiresAt: String(row.expires_at),
	}
}

export async function listDecisions(prKey: string): Promise<Array<Decision>> {
	await initSchema()
	const now = new Date().toISOString()
	await purgeExpired(now)
	return rowsOf(
		await sql(
			'SELECT * FROM decisions WHERE pr_key = ? AND expires_at > ? ORDER BY created_at ASC',
			[prKey, now],
		),
	).map(toDecision)
}

export async function upsertDecision(decision: Decision) {
	await initSchema()
	await sql(
		`INSERT INTO decisions (pr_key, kind, target, disposition, reason, head_sha, agent_id, created_at, expires_at)
		 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
		 ON CONFLICT(pr_key, kind, target) DO UPDATE SET
			disposition = excluded.disposition,
			reason = excluded.reason,
			head_sha = excluded.head_sha,
			agent_id = excluded.agent_id,
			created_at = excluded.created_at,
			expires_at = excluded.expires_at`,
		[
			decision.prKey,
			decision.kind,
			decision.target,
			decision.disposition,
			decision.reason,
			decision.headSha,
			decision.agentId,
			decision.createdAt,
			decision.expiresAt,
		],
	)
}

export async function removeDecision(prKey: string, kind: DecisionKind, target: string) {
	await initSchema()
	await sql('DELETE FROM decisions WHERE pr_key = ? AND kind = ? AND target = ?', [
		prKey,
		kind,
		target,
	])
}

export async function listEvents(prKey: string): Promise<Array<ShipEvent>> {
	await initSchema()
	const now = new Date().toISOString()
	return rowsOf(
		await sql(
			'SELECT * FROM events WHERE pr_key = ? AND expires_at > ? ORDER BY created_at ASC',
			[prKey, now],
		),
	).map((row) => {
		let data: Record<string, unknown> = {}
		try {
			data = JSON.parse(String(row.data_json || '{}'))
		} catch {
			data = {}
		}
		return {
			prKey: String(row.pr_key),
			kind: String(row.kind) as ShipEventKind,
			ref: String(row.ref),
			data,
			createdAt: String(row.created_at),
			expiresAt: String(row.expires_at),
		}
	})
}

export async function recordEvent(
	prKey: string,
	kind: ShipEventKind,
	ref: string,
	data: Record<string, unknown> = {},
) {
	await initSchema()
	const nowMs = Date.now()
	await sql(
		`INSERT INTO events (pr_key, kind, ref, data_json, created_at, expires_at)
		 VALUES (?, ?, ?, ?, ?, ?)
		 ON CONFLICT(pr_key, kind, ref) DO UPDATE SET
			data_json = excluded.data_json,
			created_at = excluded.created_at,
			expires_at = excluded.expires_at`,
		[
			prKey,
			kind,
			ref,
			JSON.stringify(data),
			new Date(nowMs).toISOString(),
			expiryFrom(nowMs, DEFAULT_TTL_DAYS),
		],
	)
}

/** How long one `usageGet` read is reused across ticks (all PRs share it). */
export const QUOTA_CACHE_MS = 5 * 60 * 1000

/**
 * Today's `outbound_fetches_per_day` usage, cached in package storage for 5
 * minutes so a fleet of ticks costs one usage read, not one per tick. `usageGet`
 * is a Kody platform call, not a sandbox outbound fetch. Null when unavailable.
 */
export async function getFetchQuota(nowMs = Date.now()): Promise<FetchQuota | null> {
	try {
		await initSchema()
		const day = new Date(nowMs).toISOString().slice(0, 10)
		const cached = rowsOf(await sql('SELECT json, fetched_at FROM usage_cache WHERE key = ?', ['fetch_quota']))[0]
		if (cached && nowMs - Number(cached.fetched_at) < QUOTA_CACHE_MS) {
			const quota = JSON.parse(String(cached.json)) as FetchQuota | null
			if (!quota || quota.day == null || quota.day === day) return quota
		}
		const quota = readFetchQuota(await kody.usageGet({}))
		await sql(
			`INSERT INTO usage_cache (key, json, fetched_at) VALUES (?, ?, ?)
			 ON CONFLICT(key) DO UPDATE SET json = excluded.json, fetched_at = excluded.fetched_at`,
			['fetch_quota', JSON.stringify(quota), nowMs],
		)
		return quota
	} catch {
		return null
	}
}