gitmost/apps/server/src/core/ai-chat/ai-chat-run.service.spec.ts

import { Logger } from '@nestjs/common';
import {
  AiChatRunService,
  RunAlreadyActiveError,
  ONE_ACTIVE_RUN_PER_CHAT_INDEX,
  mapTurnStatusToRun,
} from './ai-chat-run.service';

/** Shape a Postgres unique-violation the way the postgres.js driver surfaces it:
 *  SQLSTATE 23505 + the offending index in `constraint_name`. */
function uniqueViolation(constraintName: string): Error & {
  code: string;
  constraint_name: string;
} {
  return Object.assign(
    new Error('duplicate key value violates unique constraint'),
    {
      code: '23505',
      constraint_name: constraintName,
    },
  );
}

/**
 * Unit coverage for the #184 phase-1 run lifecycle (AiChatRunService) with a
 * hand-rolled mock repo — no Nest graph, no DB. The invariant under test is the
 * one that makes a run "autonomous": a run keeps going when its SUBSCRIBER (the
 * browser) detaches, and ONLY an explicit stop aborts it. We assert that at the
 * abort-signal level (the signal the agent loop actually consumes).
 */

function makeRepo(overrides: Record<string, jest.Mock> = {}) {
  return {
    insert: jest.fn(async (v: any) => ({
      id: 'run-1',
      status: v.status ?? 'running',
      chatId: v.chatId,
      workspaceId: v.workspaceId,
    })),
    update: jest.fn(async () => ({ id: 'run-1' })),
    markStopRequested: jest.fn(async () => ({ id: 'run-1' })),
    findActiveByChat: jest.fn(async () => undefined),
    findLatestByChat: jest.fn(async () => undefined),
    findById: jest.fn(async () => undefined),
    sweepRunning: jest.fn(async () => 0),
    ...overrides,
  };
}

describe('mapTurnStatusToRun', () => {
  it('maps the turn terminal status to the run terminal status', () => {
    expect(mapTurnStatusToRun('completed')).toBe('succeeded');
    expect(mapTurnStatusToRun('error')).toBe('failed');
    expect(mapTurnStatusToRun('aborted')).toBe('aborted');
  });
});

describe('AiChatRunService.onModuleInit (startup sweep)', () => {
  afterEach(() => jest.restoreAllMocks());

  it('calls sweepRunning and resolves; logs when > 0', async () => {
    const repo = makeRepo({ sweepRunning: jest.fn(async () => 2) });
    const logSpy = jest
      .spyOn(Logger.prototype, 'log')
      .mockImplementation(() => undefined);
    const svc = new AiChatRunService(repo as never);
    await expect(svc.onModuleInit()).resolves.toBeUndefined();
    expect(repo.sweepRunning).toHaveBeenCalledTimes(1);
    expect(logSpy).toHaveBeenCalledTimes(1);
    expect(String(logSpy.mock.calls[0][0])).toContain('2');
  });

  it('a sweep failure is swallowed (never blocks startup)', async () => {
    const repo = makeRepo({
      sweepRunning: jest.fn(async () => {
        throw new Error('db down');
      }),
    });
    const warnSpy = jest
      .spyOn(Logger.prototype, 'warn')
      .mockImplementation(() => undefined);
    const svc = new AiChatRunService(repo as never);
    await expect(svc.onModuleInit()).resolves.toBeUndefined();
    expect(warnSpy).toHaveBeenCalledTimes(1);
    expect(String(warnSpy.mock.calls[0][0])).toContain('db down');
  });
});

describe('AiChatRunService run lifecycle', () => {
  it('beginRun inserts a running row and registers a live abort controller', async () => {
    const repo = makeRepo();
    const svc = new AiChatRunService(repo as never);
    const handle = await svc.beginRun({
      chatId: 'chat-1',
      workspaceId: 'ws-1',
      userId: 'user-1',
    });
    expect(repo.insert).toHaveBeenCalledWith(
      expect.objectContaining({
        chatId: 'chat-1',
        workspaceId: 'ws-1',
        createdBy: 'user-1',
        status: 'running',
        trigger: 'user',
      }),
    );
    expect(handle.runId).toBe('run-1');
    expect(handle.signal.aborted).toBe(false);
    expect(svc.isLocallyActive('run-1')).toBe(true);
  });

  it('beginRun REJECTS the racer: a 23505 on the one-active-per-chat index throws RunAlreadyActiveError (not swallowed) and registers no controller', async () => {
    // The race: the controller's cheap pre-check passed for BOTH concurrent
    // turns, so the loser's INSERT hits the partial unique index. That rejection
    // is the authoritative gate — it must surface, not be swallowed into an
    // untracked turn.
    const repo = makeRepo({
      insert: jest.fn(async () => {
        throw uniqueViolation(ONE_ACTIVE_RUN_PER_CHAT_INDEX);
      }),
    });
    const svc = new AiChatRunService(repo as never);
    await expect(
      svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'user-1' }),
    ).rejects.toBeInstanceOf(RunAlreadyActiveError);
    // No controller leaked for a rejected start.
    expect(svc.isLocallyActive('run-1')).toBe(false);
  });

  it('beginRun does NOT mask an unrelated unique violation as already-active', async () => {
    // A 23505 on some OTHER constraint is a real bug, not the race — it must
    // propagate unchanged so it is never silently treated as "already active".
    const other = uniqueViolation('ai_chat_runs_pkey');
    const repo = makeRepo({
      insert: jest.fn(async () => {
        throw other;
      }),
    });
    const svc = new AiChatRunService(repo as never);
    await expect(
      svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'user-1' }),
    ).rejects.toBe(other);
  });

  it('beginRun propagates a non-unique insert failure unchanged', async () => {
    const boom = new Error('connection reset');
    const repo = makeRepo({
      insert: jest.fn(async () => {
        throw boom;
      }),
    });
    const svc = new AiChatRunService(repo as never);
    await expect(
      svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'user-1' }),
    ).rejects.toBe(boom);
  });

  it('two concurrent begins on one chat: exactly one wins, the other is rejected as already-active', async () => {
    // Integration-style: model the DB partial unique index with a one-shot slot.
    // The first insert claims it; the second hits a 23505 on the active index.
    let slotTaken = false;
    const repo = makeRepo({
      insert: jest.fn(async (v: any) => {
        if (slotTaken) throw uniqueViolation(ONE_ACTIVE_RUN_PER_CHAT_INDEX);
        slotTaken = true;
        return { id: 'run-win', status: v.status, chatId: v.chatId };
      }),
    });
    const svc = new AiChatRunService(repo as never);
    const results = await Promise.allSettled([
      svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'user-1' }),
      svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'user-1' }),
    ]);
    const fulfilled = results.filter((r) => r.status === 'fulfilled');
    const rejected = results.filter((r) => r.status === 'rejected');
    expect(fulfilled).toHaveLength(1);
    expect(rejected).toHaveLength(1);
    expect((rejected[0] as PromiseRejectedResult).reason).toBeInstanceOf(
      RunAlreadyActiveError,
    );
    // Exactly the winner is locally active.
    expect(svc.isLocallyActive('run-win')).toBe(true);
  });

  it('a SUBSCRIBER detaching does NOT abort the run (only an explicit stop does)', async () => {
    const repo = makeRepo();
    const svc = new AiChatRunService(repo as never);
    const handle = await svc.beginRun({
      chatId: 'chat-1',
      workspaceId: 'ws-1',
      userId: 'user-1',
    });
    // Model a browser disconnect: nothing in the run service is told to stop.
    // The signal the agent loop consumes must stay un-aborted and the run stays
    // locally active — i.e. it keeps running server-side.
    expect(handle.signal.aborted).toBe(false);
    expect(svc.isLocallyActive('run-1')).toBe(true);
    // markStopRequested was never called by a mere detach.
    expect(repo.markStopRequested).not.toHaveBeenCalled();
  });

  it('requestStop aborts the live controller, marks the row, and reports true', async () => {
    const repo = makeRepo();
    const svc = new AiChatRunService(repo as never);
    const handle = await svc.beginRun({
      chatId: 'chat-1',
      workspaceId: 'ws-1',
      userId: 'user-1',
    });
    const aborted = jest.fn();
    handle.signal.addEventListener('abort', aborted);

    const result = await svc.requestStop('run-1', 'ws-1');

    expect(result).toBe(true);
    expect(handle.signal.aborted).toBe(true);
    expect(aborted).toHaveBeenCalledTimes(1);
    expect(repo.markStopRequested).toHaveBeenCalledWith('run-1', 'ws-1');
  });

  it('requestStop on a run this replica does NOT hold still marks the row (true)', async () => {
    // e.g. after a restart, or a sibling replica owns the controller. The row is
    // marked so the owning replica/sweep settles it; we report a stop took effect.
    const repo = makeRepo({
      markStopRequested: jest.fn(async () => ({ id: 'run-9' })),
    });
    const svc = new AiChatRunService(repo as never);
    const result = await svc.requestStop('run-9', 'ws-1');
    expect(result).toBe(true);
    expect(svc.isLocallyActive('run-9')).toBe(false);
  });

  it('requestStop on an already-settled run (nothing active) reports false', async () => {
    const repo = makeRepo({
      markStopRequested: jest.fn(async () => undefined),
    });
    const svc = new AiChatRunService(repo as never);
    const result = await svc.requestStop('run-done', 'ws-1');
    expect(result).toBe(false);
  });

  it('finalizeRun settles the row to the mapped status with finishedAt and drops the in-memory entry', async () => {
    const repo = makeRepo();
    const svc = new AiChatRunService(repo as never);
    await svc.beginRun({
      chatId: 'chat-1',
      workspaceId: 'ws-1',
      userId: 'user-1',
    });
    expect(svc.isLocallyActive('run-1')).toBe(true);

    await svc.finalizeRun('run-1', 'ws-1', 'error', 'provider blew up');

    expect(svc.isLocallyActive('run-1')).toBe(false);
    expect(repo.update).toHaveBeenCalledWith(
      'run-1',
      'ws-1',
      expect.objectContaining({
        status: 'failed',
        error: 'provider blew up',
        finishedAt: expect.any(Date),
      }),
    );
  });

  it('finalizeRun is IDEMPOTENT: a second settle no-ops (single terminal write)', async () => {
    // The #184 review fix: AiChatService.stream wraps the turn in a safety-net
    // catch that settles a failed turn AND streamText's terminal callback may
    // also settle — both routes call finalizeRun. Only the FIRST may write the
    // terminal row; the second must no-op so a late settle can never clobber the
    // real terminal status or double-write the row.
    const repo = makeRepo();
    const svc = new AiChatRunService(repo as never);
    await svc.beginRun({
      chatId: 'chat-1',
      workspaceId: 'ws-1',
      userId: 'user-1',
    });

    await svc.finalizeRun('run-1', 'ws-1', 'error', 'first');
    expect(svc.isLocallyActive('run-1')).toBe(false);
    // A second settle (e.g. a streamText callback firing after the catch) no-ops.
    await svc.finalizeRun('run-1', 'ws-1', 'completed', undefined);

    expect(repo.update).toHaveBeenCalledTimes(1);
    expect(repo.update).toHaveBeenCalledWith(
      'run-1',
      'ws-1',
      expect.objectContaining({ status: 'failed', error: 'first' }),
    );
  });

  it('recordStep / linkAssistantMessage are best-effort: a repo failure is swallowed', async () => {
    const repo = makeRepo({
      update: jest.fn(async () => {
        throw new Error('transient');
      }),
    });
    jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined);
    const svc = new AiChatRunService(repo as never);
    await expect(svc.recordStep('run-1', 'ws-1', 3)).resolves.toBeUndefined();
    await expect(
      svc.linkAssistantMessage('run-1', 'ws-1', 'msg-1'),
    ).resolves.toBeUndefined();
  });
});