Neon-Desk/lib/server/repos/tasks.test.ts

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',
      '0012_company_overview_cache.sql',
      '0013_task_active_resource_unique.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, 'Search indexing could not generate embeddings for AAPL · doc-2. The AI provider returned an empty response.', 'search.embed', {
      detail: 'Embedding generation 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).toBe('search.embed');
    expect(failed?.stage_detail).toBe('Embedding generation failed.');
    expect(failed?.error).toContain('Search indexing could not generate embeddings');
    expect(failed?.stage_context?.progress?.current).toBe(2);
  });

  it('atomically deduplicates concurrent task creation for same resource', async () => {
    if (!tasksRepo) {
      throw new Error('tasks repo not initialized');
    }

    const resourceKey = 'sync_filings:AAPL';
    const concurrentCount = 10;

    const results = await Promise.all(
      Array.from({ length: concurrentCount }, (_, i) =>
        tasksRepo!.createTaskRunRecordAtomic({
          id: `task-race-${i}`,
          user_id: TEST_USER_ID,
          task_type: 'sync_filings',
          payload: { ticker: 'AAPL' },
          priority: 50,
          max_attempts: 3,
          resource_key: resourceKey
        })
      )
    );

    const createdResults = results.filter((r) => r.created);
    const conflictResults = results.filter((r) => !r.created);

    expect(createdResults.length).toBe(1);
    expect(conflictResults.length).toBe(concurrentCount - 1);
    expect(createdResults[0]?.task.resource_key).toBe(resourceKey);
    expect(createdResults[0]?.task.status).toBe('queued');
  });

  it('allows creating new task for same resource after previous task completes', async () => {
    if (!tasksRepo) {
      throw new Error('tasks repo not initialized');
    }

    const resourceKey = 'sync_filings:MSFT';

    const first = await tasksRepo.createTaskRunRecordAtomic({
      id: 'task-first-msft',
      user_id: TEST_USER_ID,
      task_type: 'sync_filings',
      payload: { ticker: 'MSFT' },
      priority: 50,
      max_attempts: 3,
      resource_key: resourceKey
    });

    expect(first.created).toBe(true);
    if (!first.created) throw new Error('Expected task to be created');

    await tasksRepo.completeTask(first.task.id, { ticker: 'MSFT' });

    const second = await tasksRepo.createTaskRunRecordAtomic({
      id: 'task-second-msft',
      user_id: TEST_USER_ID,
      task_type: 'sync_filings',
      payload: { ticker: 'MSFT' },
      priority: 50,
      max_attempts: 3,
      resource_key: resourceKey
    });

    expect(second.created).toBe(true);
    if (!second.created) throw new Error('Expected second task to be created');
    expect(second.task.id).not.toBe(first.task.id);
  });

  it('allows creating tasks without resource key without deduplication', async () => {
    if (!tasksRepo) {
      throw new Error('tasks repo not initialized');
    }

    const results = await Promise.all([
      tasksRepo.createTaskRunRecordAtomic({
        id: 'task-nokey-1',
        user_id: TEST_USER_ID,
        task_type: 'sync_filings',
        payload: { ticker: 'GOOGL' },
        priority: 50,
        max_attempts: 3,
        resource_key: null
      }),
      tasksRepo.createTaskRunRecordAtomic({
        id: 'task-nokey-2',
        user_id: TEST_USER_ID,
        task_type: 'sync_filings',
        payload: { ticker: 'GOOGL' },
        priority: 50,
        max_attempts: 3,
        resource_key: null
      })
    ]);

    expect(results[0]?.created).toBe(true);
    expect(results[1]?.created).toBe(true);
    if (!results[0]?.created || !results[1]?.created) {
      throw new Error('Expected both tasks to be created');
    }
    expect(results[0].task.id).not.toBe(results[1].task.id);
  });
});