Improve job status notifications

2026-03-09 18:53:41 -04:00
parent 1a18ac825d
commit 12a9741eca
22 changed files with 2243 additions and 302 deletions
--- a/lib/server/repos/tasks.test.ts
+++ b/lib/server/repos/tasks.test.ts
@@ -0,0 +1,222 @@
+import {
+  afterAll,
+  beforeAll,
+  beforeEach,
+  describe,
+  expect,
+  it
+} from 'bun:test';
+import { mkdtempSync, readFileSync, rmSync } from 'node:fs';
+import { tmpdir } from 'node:os';
+import { join } from 'node:path';
+import { Database } from 'bun:sqlite';
+
+const TEST_USER_ID = 'task-test-user';
+
+let tempDir: string | null = null;
+let sqliteClient: Database | null = null;
+let tasksRepo: typeof import('./tasks') | null = null;
+
+function resetDbSingletons() {
+  const globalState = globalThis as typeof globalThis & {
+    __fiscalSqliteClient?: { close?: () => void };
+    __fiscalDrizzleDb?: unknown;
+  };
+
+  globalState.__fiscalSqliteClient?.close?.();
+  globalState.__fiscalSqliteClient = undefined;
+  globalState.__fiscalDrizzleDb = undefined;
+}
+
+function applyMigration(client: Database, fileName: string) {
+  const sql = readFileSync(join(process.cwd(), 'drizzle', fileName), 'utf8');
+  client.exec(sql);
+}
+
+function ensureUser(client: Database) {
+  const now = Date.now();
+  client.exec(`
+    INSERT OR REPLACE INTO user (id, name, email, emailVerified, image, createdAt, updatedAt, role, banned, banReason, banExpires)
+    VALUES ('${TEST_USER_ID}', 'Task Test User', 'tasks@example.com', 1, NULL, ${now}, ${now}, NULL, 0, NULL, NULL);
+  `);
+}
+
+function clearTasks(client: Database) {
+  client.exec('DELETE FROM task_stage_event;');
+  client.exec('DELETE FROM task_run;');
+}
+
+describe('task repos', () => {
+  beforeAll(async () => {
+    tempDir = mkdtempSync(join(tmpdir(), 'fiscal-task-repo-'));
+    const env = process.env as Record<string, string | undefined>;
+    env.DATABASE_URL = `file:${join(tempDir, 'repo.sqlite')}`;
+    env.NODE_ENV = 'test';
+
+    resetDbSingletons();
+
+    sqliteClient = new Database(join(tempDir, 'repo.sqlite'), { create: true });
+    sqliteClient.exec('PRAGMA foreign_keys = ON;');
+    for (const file of [
+      '0000_cold_silver_centurion.sql',
+      '0001_glossy_statement_snapshots.sql',
+      '0002_workflow_task_projection_metadata.sql',
+      '0003_task_stage_event_timeline.sql',
+      '0004_watchlist_company_taxonomy.sql',
+      '0005_financial_taxonomy_v3.sql',
+      '0006_coverage_journal_tracking.sql',
+      '0007_company_financial_bundles.sql',
+      '0008_research_workspace.sql',
+      '0009_task_notification_context.sql'
+    ]) {
+      applyMigration(sqliteClient, file);
+    }
+    ensureUser(sqliteClient);
+
+    tasksRepo = await import('./tasks');
+  });
+
+  afterAll(() => {
+    sqliteClient?.close();
+    resetDbSingletons();
+    if (tempDir) {
+      rmSync(tempDir, { recursive: true, force: true });
+    }
+  });
+
+  beforeEach(() => {
+    if (!sqliteClient) {
+      throw new Error('sqlite client not initialized');
+    }
+
+    clearTasks(sqliteClient);
+  });
+
+  it('updates same-stage progress without duplicating stage events', async () => {
+    if (!tasksRepo) {
+      throw new Error('tasks repo not initialized');
+    }
+
+    const task = await tasksRepo.createTaskRunRecord({
+      id: 'task-progress',
+      user_id: TEST_USER_ID,
+      task_type: 'index_search',
+      payload: { ticker: 'AAPL' },
+      priority: 50,
+      max_attempts: 3,
+      resource_key: 'index_search:ticker:AAPL'
+    });
+
+    await tasksRepo.markTaskRunning(task.id);
+    await tasksRepo.updateTaskStage(task.id, 'search.embed', 'Embedding 1 of 3 sources', {
+      progress: { current: 1, total: 3, unit: 'sources' },
+      counters: { chunksEmbedded: 12 },
+      subject: { ticker: 'AAPL', label: 'doc-1' }
+    });
+    await tasksRepo.updateTaskStage(task.id, 'search.embed', 'Embedding 2 of 3 sources', {
+      progress: { current: 2, total: 3, unit: 'sources' },
+      counters: { chunksEmbedded: 24 },
+      subject: { ticker: 'AAPL', label: 'doc-2' }
+    });
+
+    const current = await tasksRepo.getTaskByIdForUser(task.id, TEST_USER_ID);
+    const events = await tasksRepo.listTaskStageEventsForTask(task.id, TEST_USER_ID);
+
+    expect(current?.stage_detail).toBe('Embedding 2 of 3 sources');
+    expect(current?.stage_context?.progress?.current).toBe(2);
+    expect(events.filter((event) => event.stage === 'search.embed')).toHaveLength(1);
+  });
+
+  it('lists recent tasks by updated_at descending', async () => {
+    if (!tasksRepo) {
+      throw new Error('tasks repo not initialized');
+    }
+
+    const first = await tasksRepo.createTaskRunRecord({
+      id: 'task-first',
+      user_id: TEST_USER_ID,
+      task_type: 'refresh_prices',
+      payload: {},
+      priority: 50,
+      max_attempts: 3,
+      resource_key: 'refresh_prices:portfolio'
+    });
+    await Bun.sleep(5);
+    const second = await tasksRepo.createTaskRunRecord({
+      id: 'task-second',
+      user_id: TEST_USER_ID,
+      task_type: 'portfolio_insights',
+      payload: {},
+      priority: 50,
+      max_attempts: 3,
+      resource_key: 'portfolio_insights:portfolio'
+    });
+    await Bun.sleep(5);
+    await tasksRepo.updateTaskStage(first.id, 'refresh.fetch_quotes', 'Fetching quotes', {
+      progress: { current: 1, total: 3, unit: 'tickers' }
+    });
+
+    const tasks = await tasksRepo.listRecentTasksForUser(TEST_USER_ID, 10);
+
+    expect(tasks[0]?.id).toBe(first.id);
+    expect(tasks[1]?.id).toBe(second.id);
+  });
+
+  it('preserves completion and failure detail/context on terminal tasks', async () => {
+    if (!tasksRepo) {
+      throw new Error('tasks repo not initialized');
+    }
+
+    const completedTask = await tasksRepo.createTaskRunRecord({
+      id: 'task-completed',
+      user_id: TEST_USER_ID,
+      task_type: 'analyze_filing',
+      payload: { accessionNumber: '0000320193-26-000001' },
+      priority: 50,
+      max_attempts: 3,
+      resource_key: 'analyze_filing:0000320193-26-000001'
+    });
+
+    await tasksRepo.completeTask(completedTask.id, {
+      ticker: 'AAPL',
+      accessionNumber: '0000320193-26-000001',
+      filingType: '10-Q'
+    }, {
+      detail: 'Analysis report generated for AAPL 10-Q 0000320193-26-000001.',
+      context: {
+        subject: {
+          ticker: 'AAPL',
+          accessionNumber: '0000320193-26-000001',
+          label: '10-Q'
+        }
+      }
+    });
+
+    const failedTask = await tasksRepo.createTaskRunRecord({
+      id: 'task-failed',
+      user_id: TEST_USER_ID,
+      task_type: 'index_search',
+      payload: { ticker: 'AAPL' },
+      priority: 50,
+      max_attempts: 3,
+      resource_key: 'index_search:ticker:AAPL'
+    });
+
+    await tasksRepo.markTaskFailure(failedTask.id, 'Embedding request failed', 'failed', {
+      detail: 'Embedding request failed',
+      context: {
+        progress: { current: 2, total: 5, unit: 'sources' },
+        counters: { chunksEmbedded: 20 },
+        subject: { ticker: 'AAPL', label: 'doc-2' }
+      }
+    });
+
+    const completed = await tasksRepo.getTaskByIdForUser(completedTask.id, TEST_USER_ID);
+    const failed = await tasksRepo.getTaskByIdForUser(failedTask.id, TEST_USER_ID);
+
+    expect(completed?.stage_detail).toContain('Analysis report generated');
+    expect(completed?.stage_context?.subject?.ticker).toBe('AAPL');
+    expect(failed?.stage_detail).toBe('Embedding request failed');
+    expect(failed?.stage_context?.progress?.current).toBe(2);
+  });
+});
--- a/lib/server/repos/tasks.ts
+++ b/lib/server/repos/tasks.ts
@@ -1,7 +1,8 @@
 import { and, asc, desc, eq, inArray, sql } from 'drizzle-orm';
-import type { Task, TaskStage, TaskStageEvent, TaskStatus, TaskType } from '@/lib/types';
+import type { Task, TaskStage, TaskStageContext, TaskStageEvent, TaskStatus, TaskType } from '@/lib/types';
 import { db } from '@/lib/server/db';
 import { taskRun, taskStageEvent } from '@/lib/server/db/schema';
+import { buildTaskNotification } from '@/lib/server/task-notifications';

 type TaskRow = typeof taskRun.$inferSelect;
 type TaskStageEventRow = typeof taskStageEvent.$inferSelect;
@@ -26,20 +27,27 @@ type EventInsertInput = {
  user_id: string;
  stage: TaskStage;
  stage_detail: string | null;
+  stage_context: TaskStageContext | null;
  status: TaskStatus;
  created_at: string;
 };

+type TaskCompletionState = {
+  detail?: string | null;
+  context?: TaskStageContext | null;
+};
+
 type InsertExecutor = Pick<typeof db, 'insert'>;

 function toTask(row: TaskRow): Task {
-  return {
+  const task = {
    id: row.id,
    user_id: row.user_id,
    task_type: row.task_type,
    status: row.status,
    stage: row.stage as TaskStage,
    stage_detail: row.stage_detail,
+    stage_context: row.stage_context ?? null,
    resource_key: row.resource_key,
    notification_read_at: row.notification_read_at,
    notification_silenced_at: row.notification_silenced_at,
@@ -53,6 +61,11 @@ function toTask(row: TaskRow): Task {
    created_at: row.created_at,
    updated_at: row.updated_at,
    finished_at: row.finished_at
+  } satisfies Omit<Task, 'notification'>;
+
+  return {
+    ...task,
+    notification: buildTaskNotification(task)
  };
 }

@@ -63,6 +76,7 @@ function toTaskStageEvent(row: TaskStageEventRow): TaskStageEvent {
    user_id: row.user_id,
    stage: row.stage as TaskStage,
    stage_detail: row.stage_detail,
+    stage_context: row.stage_context ?? null,
    status: row.status as TaskStatus,
    created_at: row.created_at
  };
@@ -89,6 +103,7 @@ async function insertTaskStageEvent(executor: InsertExecutor, input: EventInsert
    user_id: input.user_id,
    stage: input.stage,
    stage_detail: input.stage_detail,
+    stage_context: input.stage_context,
    status: input.status,
    created_at: input.created_at
  });
@@ -107,6 +122,7 @@ export async function createTaskRunRecord(input: CreateTaskInput) {
        status: 'queued',
        stage: 'queued',
        stage_detail: null,
+        stage_context: null,
        resource_key: input.resource_key ?? null,
        notification_read_at: null,
        notification_silenced_at: null,
@@ -128,6 +144,7 @@ export async function createTaskRunRecord(input: CreateTaskInput) {
      user_id: row.user_id,
      stage: row.stage as TaskStage,
      stage_detail: row.stage_detail,
+      stage_context: row.stage_context ?? null,
      status: row.status,
      created_at: now
    });
@@ -168,13 +185,13 @@ export async function listRecentTasksForUser(
      .select()
      .from(taskRun)
      .where(and(eq(taskRun.user_id, userId), inArray(taskRun.status, statuses)))
-      .orderBy(desc(taskRun.created_at))
+      .orderBy(desc(taskRun.updated_at), desc(taskRun.created_at))
      .limit(safeLimit)
    : await db
      .select()
      .from(taskRun)
      .where(eq(taskRun.user_id, userId))
-      .orderBy(desc(taskRun.created_at))
+      .orderBy(desc(taskRun.updated_at), desc(taskRun.created_at))
      .limit(safeLimit);

  return rows.map(toTask);
@@ -212,7 +229,7 @@ export async function findInFlightTaskByResourceKey(
      eq(taskRun.resource_key, resourceKey),
      inArray(taskRun.status, ['queued', 'running'])
    ))
-    .orderBy(desc(taskRun.created_at))
+    .orderBy(desc(taskRun.updated_at), desc(taskRun.created_at))
    .limit(1);

  return row ? toTask(row) : null;
@@ -228,6 +245,7 @@ export async function markTaskRunning(taskId: string) {
        status: 'running',
        stage: 'running',
        stage_detail: 'Workflow task is now running',
+        stage_context: null,
        attempts: sql`${taskRun.attempts} + 1`,
        updated_at: now,
        finished_at: null
@@ -244,6 +262,7 @@ export async function markTaskRunning(taskId: string) {
      user_id: row.user_id,
      stage: row.stage as TaskStage,
      stage_detail: row.stage_detail,
+      stage_context: row.stage_context ?? null,
      status: row.status,
      created_at: now
    });
@@ -252,15 +271,31 @@ export async function markTaskRunning(taskId: string) {
  });
 }

-export async function updateTaskStage(taskId: string, stage: TaskStage, detail: string | null = null) {
+export async function updateTaskStage(
+  taskId: string,
+  stage: TaskStage,
+  detail: string | null = null,
+  context: TaskStageContext | null = null
+) {
  const now = new Date().toISOString();

  return await db.transaction(async (tx) => {
+    const [current] = await tx
+      .select()
+      .from(taskRun)
+      .where(eq(taskRun.id, taskId))
+      .limit(1);
+
+    if (!current) {
+      return null;
+    }
+
    const [row] = await tx
      .update(taskRun)
      .set({
        stage,
        stage_detail: detail,
+        stage_context: context,
        updated_at: now
      })
      .where(eq(taskRun.id, taskId))
@@ -270,20 +305,27 @@ export async function updateTaskStage(taskId: string, stage: TaskStage, detail:
      return null;
    }

-    await insertTaskStageEvent(tx, {
-      task_id: row.id,
-      user_id: row.user_id,
-      stage: row.stage as TaskStage,
-      stage_detail: row.stage_detail,
-      status: row.status,
-      created_at: now
-    });
+    if (current.stage !== stage) {
+      await insertTaskStageEvent(tx, {
+        task_id: row.id,
+        user_id: row.user_id,
+        stage: row.stage as TaskStage,
+        stage_detail: row.stage_detail,
+        stage_context: row.stage_context ?? null,
+        status: row.status,
+        created_at: now
+      });
+    }

    return toTask(row);
  });
 }

-export async function completeTask(taskId: string, result: Record<string, unknown>) {
+export async function completeTask(
+  taskId: string,
+  result: Record<string, unknown>,
+  completion: TaskCompletionState = {}
+) {
  const now = new Date().toISOString();

  return await db.transaction(async (tx) => {
@@ -292,7 +334,8 @@ export async function completeTask(taskId: string, result: Record<string, unknow
      .set({
        status: 'completed',
        stage: 'completed',
-        stage_detail: null,
+        stage_detail: completion.detail ?? 'Task finished successfully.',
+        stage_context: completion.context ?? null,
        result,
        error: null,
        updated_at: now,
@@ -310,6 +353,7 @@ export async function completeTask(taskId: string, result: Record<string, unknow
      user_id: row.user_id,
      stage: row.stage as TaskStage,
      stage_detail: row.stage_detail,
+      stage_context: row.stage_context ?? null,
      status: row.status,
      created_at: now
    });
@@ -318,7 +362,12 @@ export async function completeTask(taskId: string, result: Record<string, unknow
  });
 }

-export async function markTaskFailure(taskId: string, reason: string, stage: TaskStage = 'failed') {
+export async function markTaskFailure(
+  taskId: string,
+  reason: string,
+  stage: TaskStage = 'failed',
+  failure: TaskCompletionState = {}
+) {
  const now = new Date().toISOString();

  return await db.transaction(async (tx) => {
@@ -327,7 +376,8 @@ export async function markTaskFailure(taskId: string, reason: string, stage: Tas
      .set({
        status: 'failed',
        stage,
-        stage_detail: null,
+        stage_detail: failure.detail ?? reason,
+        stage_context: failure.context ?? null,
        error: reason,
        updated_at: now,
        finished_at: now
@@ -344,6 +394,7 @@ export async function markTaskFailure(taskId: string, reason: string, stage: Tas
      user_id: row.user_id,
      stage: row.stage as TaskStage,
      stage_detail: row.stage_detail,
+      stage_context: row.stage_context ?? null,
      status: row.status,
      created_at: now
    });
@@ -375,7 +426,8 @@ export async function setTaskStatusFromWorkflow(
    const hasNoStateChange = current.status === status
      && current.stage === nextStage
      && (current.error ?? null) === nextError
-      && current.stage_detail === null
+      && (current.stage_detail ?? null) === (nextStatusDetail(status, nextError) ?? null)
+      && (current.stage_context ?? null) === null
      && (isTerminal ? current.finished_at !== null : current.finished_at === null);

    if (hasNoStateChange) {
@@ -388,7 +440,8 @@ export async function setTaskStatusFromWorkflow(
      .set({
        status,
        stage: nextStage,
-        stage_detail: null,
+        stage_detail: nextStatusDetail(status, nextError),
+        stage_context: null,
        error: nextError,
        updated_at: now,
        finished_at: isTerminal ? now : null
@@ -405,6 +458,7 @@ export async function setTaskStatusFromWorkflow(
      user_id: row.user_id,
      stage: row.stage as TaskStage,
      stage_detail: row.stage_detail,
+      stage_context: row.stage_context ?? null,
      status: row.status,
      created_at: now
    });
@@ -452,6 +506,22 @@ export async function updateTaskNotificationState(
  return row ? toTask(row) : null;
 }

+function nextStatusDetail(status: TaskStatus, error?: string | null) {
+  if (status === 'failed') {
+    return error ?? 'Workflow run failed';
+  }
+
+  if (status === 'completed') {
+    return 'Workflow run completed.';
+  }
+
+  if (status === 'running') {
+    return 'Workflow task is now running';
+  }
+
+  return null;
+}
+
 export async function listTaskStageEventsForTask(taskId: string, userId: string) {
  const rows = await db
    .select()