Neon-Desk/lib/server/repos/research-copilot.test.ts

import {
  afterAll,
  beforeAll,
  beforeEach,
  describe,
  expect,
  it
} from 'bun:test';
import { mock } 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 = 'copilot-user';

let tempDir: string | null = null;
let sqliteClient: Database | null = null;
let copilotRepo: typeof import('./research-copilot') | null = null;

async function loadRepoModule() {
  const moduleUrl = new URL(`./research-copilot.ts?test=${Date.now()}`, import.meta.url).href;
  return await import(moduleUrl) as typeof import('./research-copilot');
}

function resetDbSingletons() {
  const globalState = globalThis as typeof globalThis & {
    __fiscalSqliteClient?: Database;
    __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}', 'Copilot User', 'copilot@example.com', 1, NULL, ${now}, ${now}, NULL, 0, NULL, NULL);
  `);
}

describe('research copilot repo', () => {
  beforeAll(async () => {
    mock.restore();
    tempDir = mkdtempSync(join(tmpdir(), 'fiscal-copilot-repo-'));
    process.env.DATABASE_URL = `file:${join(tempDir, 'repo.sqlite')}`;
    (process.env as Record<string, string | undefined>).NODE_ENV = 'test';

    resetDbSingletons();
    sqliteClient = new Database(join(tempDir, 'repo.sqlite'), { create: true });
    sqliteClient.exec('PRAGMA foreign_keys = ON;');
    applyMigration(sqliteClient, '0000_cold_silver_centurion.sql');
    applyMigration(sqliteClient, '0008_research_workspace.sql');
    applyMigration(sqliteClient, '0013_research_copilot.sql');
    ensureUser(sqliteClient);

    const globalState = globalThis as typeof globalThis & {
      __fiscalSqliteClient?: Database;
      __fiscalDrizzleDb?: unknown;
    };
    globalState.__fiscalSqliteClient = sqliteClient;
    globalState.__fiscalDrizzleDb = undefined;

    copilotRepo = await loadRepoModule();
  });

  afterAll(() => {
    mock.restore();
    sqliteClient?.close();
    resetDbSingletons();
    if (tempDir) {
      rmSync(tempDir, { recursive: true, force: true });
    }
  });

  beforeEach(() => {
    sqliteClient?.exec('DELETE FROM research_copilot_message;');
    sqliteClient?.exec('DELETE FROM research_copilot_session;');
  });

  it('creates and reloads ticker-scoped sessions', async () => {
    if (!copilotRepo) {
      throw new Error('repo not initialized');
    }

    const session = await copilotRepo.getOrCreateResearchCopilotSession({
      userId: TEST_USER_ID,
      ticker: 'msft',
      selectedSources: ['documents', 'research'],
      pinnedArtifactIds: [2, 2, 5]
    });

    const loaded = await copilotRepo.getResearchCopilotSessionByTicker(TEST_USER_ID, 'MSFT');

    expect(session.ticker).toBe('MSFT');
    expect(session.selected_sources).toEqual(['documents', 'research']);
    expect(session.pinned_artifact_ids).toEqual([2, 5]);
    expect(loaded?.id).toBe(session.id);
  });

  it('appends messages and updates session state', async () => {
    if (!copilotRepo) {
      throw new Error('repo not initialized');
    }

    const session = await copilotRepo.getOrCreateResearchCopilotSession({
      userId: TEST_USER_ID,
      ticker: 'NVDA'
    });

    await copilotRepo.appendResearchCopilotMessage({
      userId: TEST_USER_ID,
      sessionId: session.id,
      role: 'user',
      contentMarkdown: 'What changed in the latest filing?',
      selectedSources: ['filings'],
      pinnedArtifactIds: [7],
      memoSection: 'thesis'
    });

    await copilotRepo.appendResearchCopilotMessage({
      userId: TEST_USER_ID,
      sessionId: session.id,
      role: 'assistant',
      contentMarkdown: 'Demand remained strong [1]',
      citations: [{
        index: 1,
        label: 'NVDA 10-K [1]',
        chunkId: 1,
        href: '/filings?ticker=NVDA',
        source: 'filings',
        sourceKind: 'filing_brief',
        sourceRef: '0001',
        title: '10-K brief',
        ticker: 'NVDA',
        accessionNumber: '0001',
        filingDate: '2026-01-01',
        excerpt: 'Demand remained strong.',
        artifactId: 3
      }]
    });

    const updated = await copilotRepo.upsertResearchCopilotSessionState({
      userId: TEST_USER_ID,
      ticker: 'NVDA',
      title: 'NVDA demand update',
      selectedSources: ['filings'],
      pinnedArtifactIds: [7]
    });

    expect(updated.title).toBe('NVDA demand update');
    expect(updated.messages).toHaveLength(2);
    expect(updated.messages[0]?.selected_sources).toEqual(['filings']);
    expect(updated.messages[0]?.memo_section).toBe('thesis');
    expect(updated.messages[1]?.citations[0]?.artifactId).toBe(3);
  });
});