import type { ServerWebSocket } from 'bun'; import { randomUUID } from 'crypto'; import { EventEmitter } from 'events'; import { createPipelineJob, getPipelineJob, updatePipelineJob, getPipelineJobsForUser, getOldestPendingJob, countPendingJobs, deletePipelineJob, deleteTerminalJobsForUser, markInterruptedJobs, getUserById, } from 'officerdb'; import { unlink } from 'node:fs/promises'; import { toShellUsername } from '../../data-path'; import { executePipeline } from './pipeline-executor'; import type { OutMessage } from './pipeline-executor'; import { executeScript, jobLogPath, type ScriptEvent } from './execute-script'; // Everything a job can stream — pipeline structural events plus script stdout/stderr/exit. type JobEvent = OutMessage | ScriptEvent; type JobMode = 'pipeline' | 'script' | 'agentic'; type WSData = { userId: number; email: string; username: string; }; type LiveJob = { jobId: string; userId: number; abortSignal: { aborted: boolean }; emitter: EventEmitter; eventBuffer: JobEvent[]; viewers: Set>; progressDirty: boolean; lastProgress: unknown; lastCost: unknown; }; const EVENT_BUFFER_SIZE = 500; const PROGRESS_FLUSH_MS = 3000; const liveJobs = new Map(); function sendToViewer(ws: ServerWebSocket, jobId: string, event: JobEvent) { if (ws.readyState === 1) { ws.send(JSON.stringify({ jobId, ...event })); } } function broadcast(job: LiveJob, event: JobEvent) { // Buffer for replay — skip high-frequency streams (deltas / script stdout+stderr). Those are // live-only; a late viewer backfills script output from the persisted log file instead. if (event.type !== 'assistant:delta' && event.type !== 'stdout' && event.type !== 'stderr') { job.eventBuffer.push(event); if (job.eventBuffer.length > EVENT_BUFFER_SIZE) { job.eventBuffer.shift(); } } // Forward to all viewers for (const ws of job.viewers) { sendToViewer(ws, job.jobId, event); } } type StartJobParams = { userId: number; email: string; username: string; mode?: JobMode; // defaults to 'pipeline' for back-compat with the existing pipeline caller taskDirName: string; taskName: string; inputs: Record; cwd?: string; config: unknown; model?: string; startAt?: number; }; export function runningCount(): number { return liveJobs.size; } // Create a job. action 'start' runs it now; 'queue' runs it only if nothing is running, else it stays // 'pending' and gets promoted when the running job finishes. (Single user → one global queue.) export async function enqueueJob( params: StartJobParams, action: 'start' | 'queue', ): Promise<{ jobId: string; status: 'running' | 'pending' }> { const jobId = randomUUID(); const mode: JobMode = params.mode ?? 'pipeline'; const run = action === 'start' || runningCount() === 0; await createPipelineJob({ id: jobId, userId: params.userId, taskDirName: params.taskDirName, taskName: params.taskName, mode, status: run ? 'running' : 'pending', inputs: params.inputs, cwd: params.cwd ?? null, config: params.config ?? {}, // column is NOT NULL; many script tasks have no config block startedAt: run ? new Date() : null, }); if (run) launch(jobId, mode, params); return { jobId, status: run ? 'running' : 'pending' }; } // Back-compat: the pipeline WS caller expects startJob(params) -> jobId, running immediately. export async function startJob(params: StartJobParams): Promise { return (await enqueueJob(params, 'start')).jobId; } // Build live state + run the executor for a job whose DB row is already status 'running'. function launch(jobId: string, mode: JobMode, params: StartJobParams) { // Create live state const job: LiveJob = { jobId, userId: params.userId, abortSignal: { aborted: false }, emitter: new EventEmitter(), eventBuffer: [], viewers: new Set(), progressDirty: false, lastProgress: null, lastCost: null, }; liveJobs.set(jobId, job); // Emit callback — broadcasts to viewers and tracks progress const emit = (event: JobEvent) => { broadcast(job, event); // Track progress for DB persistence (skip streaming messages — only structural events) if (event.type === 'pipeline:init' || event.type.startsWith('step:') || event.type.startsWith('iteration:')) { job.progressDirty = true; job.lastProgress = extractProgress(event, job.lastProgress); } if (event.type === 'step:complete' || event.type === 'iteration:complete') { const cost = 'cost' in event ? event.cost : undefined; if (cost) { const prev = (job.lastCost as { inputTokens: number; outputTokens: number; totalUSD: number } | null) ?? { inputTokens: 0, outputTokens: 0, totalUSD: 0, }; job.lastCost = { inputTokens: prev.inputTokens + cost.inputTokens, outputTokens: prev.outputTokens + cost.outputTokens, totalUSD: prev.totalUSD + cost.totalUSD, }; job.progressDirty = true; } } }; // Periodic progress flush to DB const flushInterval = setInterval(async () => { if (job.progressDirty) { job.progressDirty = false; try { await updatePipelineJob(jobId, { progress: job.lastProgress as Record, totalCost: job.lastCost as Record, }); } catch (err) { console.error(`[pipeline-jobs] failed to flush progress for ${jobId}:`, err); } } }, PROGRESS_FLUSH_MS); // Run the job in the background — dispatch by mode. const runner: Promise = mode === 'script' ? executeScript({ jobId, email: params.email, taskDirName: params.taskDirName, inputs: params.inputs, cwd: params.cwd, abortSignal: job.abortSignal, emit, }) : executePipeline({ userId: params.userId, email: params.email, username: params.username, taskDirName: params.taskDirName, inputs: params.inputs, cwd: params.cwd, model: params.model, startAt: params.startAt, abortSignal: job.abortSignal, emit, }); runner .then(async (result) => { clearInterval(flushInterval); // Script jobs resolve with an exit code — a non-zero exit is a failure. Pipelines resolve void. const exitCode = result && typeof result === 'object' && 'exitCode' in result ? result.exitCode : null; const failed = exitCode !== null && exitCode !== 0; await updatePipelineJob(jobId, { status: failed ? 'failed' : 'completed', exitCode, progress: job.lastProgress as Record, totalCost: job.lastCost as Record, error: failed ? `Script exited with code ${exitCode}` : undefined, completedAt: new Date(), }).catch((err) => console.error(`[pipeline-jobs] failed to finalize ${jobId}:`, err)); liveJobs.delete(jobId); void promoteNext(); }) .catch(async (err) => { clearInterval(flushInterval); const message = err instanceof Error ? err.message : String(err); const isStopped = job.abortSignal.aborted; broadcast(job, isStopped ? { type: 'stopped' } : { type: 'error', message }); await updatePipelineJob(jobId, { status: isStopped ? 'stopped' : 'failed', progress: job.lastProgress as Record, totalCost: job.lastCost as Record, error: isStopped ? undefined : message, completedAt: new Date(), }).catch((e) => console.error(`[pipeline-jobs] failed to finalize ${jobId}:`, e)); liveJobs.delete(jobId); void promoteNext(); }); } // When a job finishes (and nothing else is running), promote the oldest queued job. Also called on // startup to resume a backlog. Rebuilds the user context from userId since there's no request here. async function promoteNext(): Promise { if (runningCount() > 0) return; const next = await getOldestPendingJob(); if (!next) return; const user = await getUserById(next.userId); if (!user) { await updatePipelineJob(next.id, { status: 'failed', error: 'user not found', completedAt: new Date() }).catch( () => {}, ); return promoteNext(); } await updatePipelineJob(next.id, { status: 'running', startedAt: new Date() }); const nextMode = (next.mode as JobMode) ?? 'pipeline'; launch(next.id, nextMode, { userId: next.userId, email: user.email, username: toShellUsername(user.username ?? '', user.email), mode: nextMode, taskDirName: next.taskDirName, taskName: next.taskName, inputs: next.inputs as Record, cwd: next.cwd ?? undefined, config: next.config, }); } export function attachViewer(jobId: string, ws: ServerWebSocket) { const job = liveJobs.get(jobId); if (!job) return false; job.viewers.add(ws); // Replay buffered events so the viewer catches up for (const event of job.eventBuffer) { sendToViewer(ws, jobId, event); } return true; } export function detachViewer(jobId: string, ws: ServerWebSocket) { const job = liveJobs.get(jobId); if (job) job.viewers.delete(ws); } export function detachAllForViewer(ws: ServerWebSocket) { for (const job of liveJobs.values()) { job.viewers.delete(ws); } } export function stopJob(jobId: string): boolean { const job = liveJobs.get(jobId); if (!job) return false; job.abortSignal.aborted = true; return true; } // Stop a running job (cooperative abort) OR cancel a queued one (mark stopped so it won't promote). export async function requestStop(jobId: string): Promise<'stopped' | 'cancelled' | 'noop'> { if (stopJob(jobId)) return 'stopped'; const job = await getPipelineJob(jobId); if (job && job.status === 'pending') { await updatePipelineJob(jobId, { status: 'stopped', completedAt: new Date() }); return 'cancelled'; } return 'noop'; } export function isJobLive(jobId: string): boolean { return liveJobs.has(jobId); } export function getLiveJobForUser(userId: number): string | null { for (const [jobId, job] of liveJobs) { if (job.userId === userId) return jobId; } return null; } export async function getJobsForUser(userId: number) { const jobs = await getPipelineJobsForUser(userId); return jobs.map((j) => ({ ...j, isLive: liveJobs.has(j.id), })); } // Delete a job (row + log). Refuses a running job — it must be stopped first. export async function deleteJob(jobId: string): Promise<'deleted' | 'running' | 'noop'> { if (liveJobs.has(jobId)) return 'running'; const job = await getPipelineJob(jobId); if (!job) return 'noop'; await deletePipelineJob(jobId); await unlink(jobLogPath(jobId)).catch(() => {}); return 'deleted'; } // Clear all finished jobs (rows + logs). Running/queued jobs are untouched. export async function clearHistory(userId: number): Promise { const ids = await deleteTerminalJobsForUser(userId); await Promise.all(ids.map((id) => unlink(jobLogPath(id)).catch(() => {}))); return ids.length; } // Lightweight header-badge summary: how many of the user's jobs are running / queued, and which one // is running (for the "running" badge's link). export async function getCounts( userId: number, ): Promise<{ running: number; runningJobId: string | null; queued: number }> { let running = 0; let runningJobId: string | null = null; for (const [id, job] of liveJobs) { if (job.userId === userId) { running++; if (!runningJobId) runningJobId = id; } } const queued = await countPendingJobs(userId); return { running, runningJobId, queued }; } export async function getJob(jobId: string) { const job = await getPipelineJob(jobId); if (!job) return null; return { ...job, isLive: liveJobs.has(job.id) }; } export async function cleanupOnStartup() { const count = await markInterruptedJobs(); if (count > 0) { console.log(`[pipeline-jobs] marked ${count} interrupted jobs on startup`); } // Resume any queued backlog that survived the restart. await promoteNext(); } /** Extract minimal progress info from events for DB storage. Script events fall through to `prev`. */ function extractProgress(event: JobEvent, prev: unknown): unknown { const p = (prev as Record) ?? {}; switch (event.type) { case 'pipeline:init': return { ...p, steps: event.steps, currentStepIndex: -1 }; case 'step:start': return { ...p, currentStepIndex: event.stepIndex, currentIteration: event.iteration ?? null }; case 'step:complete': return { ...p, currentStepIndex: event.stepIndex }; case 'step:parallel': return { ...p, currentStepIndex: event.stepIndex, parallel: { taskName: event.taskName, concurrency: event.concurrency, iterations: event.iterations.map((l) => ({ label: l, status: 'pending' })), }, }; case 'iteration:start': if (p.parallel && Array.isArray((p.parallel as Record).iterations)) { const parallel = p.parallel as { iterations: Array<{ label: string; status: string }> }; return { ...p, parallel: { ...parallel, iterations: parallel.iterations.map((it) => (it.label === event.label ? { ...it, status: 'running' } : it)), }, }; } return p; case 'iteration:complete': case 'iteration:error': if (p.parallel && Array.isArray((p.parallel as Record).iterations)) { const parallel = p.parallel as { iterations: Array<{ label: string; status: string }> }; const status = event.type === 'iteration:complete' ? 'complete' : 'error'; return { ...p, parallel: { ...parallel, iterations: parallel.iterations.map((it) => (it.label === event.label ? { ...it, status } : it)), }, }; } return p; default: return p; } }