feat(ai-chat): durable detached agent runs — phase 1 (#184/#234)
Squashed for a clean rebase onto develop (was 19 commits; the reviewer approved
the net diff at fb246080). Detaches an agent run from the HTTP request/browser
window: a run is a first-class lifecycle object (ai_chat_runs), a browser
disconnect no longer kills it, a concurrent-run insert-gate prevents double runs,
and a reopened chat live-follows a still-running run via a polled observer merge.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
agent_coder
@@ -0,0 +1,527 @@
|
||||
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).
|
||||
*/
|
||||
|
||||
/** Minimal EnvironmentService stub. Single-instance (CLOUD unset) by default. */
|
||||
function makeEnv(isCloud = false) {
|
||||
return { isCloud: () => isCloud };
|
||||
}
|
||||
|
||||
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, makeEnv() 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, makeEnv() as never);
|
||||
await expect(svc.onModuleInit()).resolves.toBeUndefined();
|
||||
// The first warn is the sweep failure (the multi-instance warn never fires
|
||||
// single-instance), so the message is the db error.
|
||||
expect(String(warnSpy.mock.calls[0][0])).toContain('db down');
|
||||
});
|
||||
|
||||
it('F1 (DECISION C): the boot sweep is UNCONDITIONAL — sweepRunning is called with NO staleness window, so a fresh running run (updatedAt = now) is settled, not skipped', async () => {
|
||||
// The bug: a fast restart (deploy/OOM within minutes of the last step) left a
|
||||
// run stuck 'running' under the old 10-min window, 409ing every later turn in
|
||||
// the chat. The fix settles ALL pending|running on boot. We assert the service
|
||||
// invokes sweepRunning with no `staleMs` (the unconditional path); the repo's
|
||||
// own spec proves no-window => no updatedAt filter.
|
||||
const repo = makeRepo({ sweepRunning: jest.fn(async () => 1) });
|
||||
jest.spyOn(Logger.prototype, 'log').mockImplementation(() => undefined);
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await svc.onModuleInit();
|
||||
expect(repo.sweepRunning).toHaveBeenCalledTimes(1);
|
||||
const callArgs = repo.sweepRunning.mock.calls[0] as unknown[];
|
||||
const firstArg = callArgs[0] as { staleMs?: number } | undefined;
|
||||
// Either no opts at all, or opts without a staleMs window => unconditional.
|
||||
expect(firstArg?.staleMs).toBeUndefined();
|
||||
});
|
||||
|
||||
it('F2 (DECISION A): warns at startup that autonomousRuns is single-instance-only when a horizontally-scaled deployment (CLOUD) is detected', async () => {
|
||||
const repo = makeRepo();
|
||||
const warnSpy = jest
|
||||
.spyOn(Logger.prototype, 'warn')
|
||||
.mockImplementation(() => undefined);
|
||||
const svc = new AiChatRunService(repo as never, makeEnv(true) as never);
|
||||
await svc.onModuleInit();
|
||||
const warned = warnSpy.mock.calls.some((c) =>
|
||||
/single-instance-only/i.test(String(c[0])),
|
||||
);
|
||||
expect(warned).toBe(true);
|
||||
});
|
||||
|
||||
it('F2: does NOT warn about multi-instance on a single-instance (CLOUD unset) deployment', async () => {
|
||||
const repo = makeRepo();
|
||||
const warnSpy = jest
|
||||
.spyOn(Logger.prototype, 'warn')
|
||||
.mockImplementation(() => undefined);
|
||||
const svc = new AiChatRunService(repo as never, makeEnv(false) as never);
|
||||
await svc.onModuleInit();
|
||||
const warned = warnSpy.mock.calls.some((c) =>
|
||||
/single-instance-only/i.test(String(c[0])),
|
||||
);
|
||||
expect(warned).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
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, makeEnv() 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, makeEnv() 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, makeEnv() 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, makeEnv() 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, makeEnv() 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, makeEnv() 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, makeEnv() 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, makeEnv() as never);
|
||||
const result = await svc.requestStop('run-9', 'ws-1');
|
||||
expect(result).toBe(true);
|
||||
expect(svc.isLocallyActive('run-9')).toBe(false);
|
||||
});
|
||||
|
||||
it('requestStop still aborts the live controller when markStopRequested rejects (transient DB error)', async () => {
|
||||
// F15: the in-memory abort is the ONLY thing that stops a run and must not be
|
||||
// hostage to the audit write of stop_requested_at. A transient failure on
|
||||
// markStopRequested must NOT prevent abort() nor make requestStop throw.
|
||||
const warnSpy = jest
|
||||
.spyOn(Logger.prototype, 'warn')
|
||||
.mockImplementation(() => undefined);
|
||||
const repo = makeRepo({
|
||||
markStopRequested: jest.fn(async () => {
|
||||
throw new Error('pool exhausted');
|
||||
}),
|
||||
});
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() 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);
|
||||
|
||||
// Does NOT throw despite the DB write rejecting.
|
||||
const result = await svc.requestStop('run-1', 'ws-1');
|
||||
|
||||
// The live turn was aborted even though the audit write failed...
|
||||
expect(handle.signal.aborted).toBe(true);
|
||||
expect(aborted).toHaveBeenCalledTimes(1);
|
||||
expect(repo.markStopRequested).toHaveBeenCalledWith('run-1', 'ws-1');
|
||||
// ...the catch branch logged the swallowed failure...
|
||||
expect(warnSpy).toHaveBeenCalledTimes(1);
|
||||
// ...and a stop is reported as having taken effect (the entry existed).
|
||||
expect(result).toBe(true);
|
||||
warnSpy.mockRestore();
|
||||
});
|
||||
|
||||
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, makeEnv() 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, makeEnv() 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, makeEnv() 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('CONCURRENCY: two simultaneous finalizeRun on the same run write the terminal row EXACTLY ONCE (the 2nd caller exits synchronously at the atomic claim)', async () => {
|
||||
// The CRITICAL race: AiChatService.stream's safety-net catch settles the turn
|
||||
// to 'error' while a streamText terminal callback also settles it — both call
|
||||
// finalizeRun for the SAME runId. The once-gate must close ATOMICALLY: a
|
||||
// `settled.has` check alone is read BEFORE the awaited UPDATE, so both callers
|
||||
// would pass it and BOTH write the row (last-write-wins clobber + double
|
||||
// write). The fix claims the run with a SYNCHRONOUS `active.delete` before any
|
||||
// await, so the second caller returns in the same tick, before the UPDATE.
|
||||
//
|
||||
// We force the two calls to overlap by making `update` return a promise we
|
||||
// resolve only AFTER both finalizeRun calls have run their synchronous bodies.
|
||||
let resolveUpdate!: (v: unknown) => void;
|
||||
const updateGate = new Promise((res) => {
|
||||
resolveUpdate = res;
|
||||
});
|
||||
const update = jest.fn(() => updateGate);
|
||||
const repo = makeRepo({ update });
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await svc.beginRun({
|
||||
chatId: 'chat-1',
|
||||
workspaceId: 'ws-1',
|
||||
userId: 'user-1',
|
||||
});
|
||||
|
||||
// Fire both before the (pending) update resolves. The first synchronously
|
||||
// claims the entry (active.delete) and awaits update; the second, started in
|
||||
// the same macrotask, finds the entry already gone and returns at the claim
|
||||
// WITHOUT ever calling update.
|
||||
const p1 = svc.finalizeRun('run-1', 'ws-1', 'completed');
|
||||
const p2 = svc.finalizeRun('run-1', 'ws-1', 'error', 'safety-net');
|
||||
|
||||
// The decisive assertion: exactly one caller reached the terminal UPDATE.
|
||||
expect(update).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Let the single in-flight update land; both calls resolve cleanly.
|
||||
resolveUpdate({ id: 'run-1' });
|
||||
await Promise.all([p1, p2]);
|
||||
|
||||
expect(update).toHaveBeenCalledTimes(1);
|
||||
// The winner is the FIRST caller ('completed' -> 'succeeded'); the late
|
||||
// 'error' settle never wrote, so it could not clobber the real status.
|
||||
expect(update).toHaveBeenCalledWith(
|
||||
'run-1',
|
||||
'ws-1',
|
||||
expect.objectContaining({ status: 'succeeded' }),
|
||||
);
|
||||
expect(svc.isLocallyActive('run-1')).toBe(false);
|
||||
});
|
||||
|
||||
it('F6: a TRANSIENT terminal-write failure is ridden out by the bounded retry — the run is settled, not stranded', async () => {
|
||||
// The bug: finalizeRun used to DROP the in-memory entry BEFORE the terminal
|
||||
// UPDATE, then only warn-log a failure. A single transient blip (pool
|
||||
// exhaustion / deadlock / connection hiccup) on that PK UPDATE left the row
|
||||
// 'running' with nothing left to recover it -> every later turn in that chat
|
||||
// 409s until a restart. The fix updates FIRST and retries.
|
||||
let calls = 0;
|
||||
const repo = makeRepo({
|
||||
update: jest.fn(async () => {
|
||||
calls += 1;
|
||||
if (calls === 1) throw new Error('deadlock detected');
|
||||
return { id: 'run-1' };
|
||||
}),
|
||||
});
|
||||
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined);
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await svc.beginRun({
|
||||
chatId: 'chat-1',
|
||||
workspaceId: 'ws-1',
|
||||
userId: 'user-1',
|
||||
});
|
||||
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'completed');
|
||||
|
||||
// The retry landed the terminal write: the entry is dropped (slot freed) and
|
||||
// the row carries the real terminal status — NOT stranded at 'running'.
|
||||
expect(svc.isLocallyActive('run-1')).toBe(false);
|
||||
expect(repo.update).toHaveBeenCalledTimes(2);
|
||||
expect(repo.update).toHaveBeenLastCalledWith(
|
||||
'run-1',
|
||||
'ws-1',
|
||||
expect.objectContaining({ status: 'succeeded' }),
|
||||
);
|
||||
});
|
||||
|
||||
it('F6: if the terminal write keeps failing, the entry is RETAINED and a LATER settle completes it (chat not permanently 409d)', async () => {
|
||||
// Worst case: the DB is down for the whole first finalize (all attempts fail).
|
||||
// The run must NOT be silently lost — the entry stays so a subsequent settle
|
||||
// (a streamText callback, requestStop -> onAbort, or a future sweep) can retry.
|
||||
let healthy = false;
|
||||
const repo = makeRepo({
|
||||
update: jest.fn(async () => {
|
||||
if (!healthy) throw new Error('pool exhausted');
|
||||
return { id: 'run-1' };
|
||||
}),
|
||||
});
|
||||
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined);
|
||||
const errorSpy = jest
|
||||
.spyOn(Logger.prototype, 'error')
|
||||
.mockImplementation(() => undefined);
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await svc.beginRun({
|
||||
chatId: 'chat-1',
|
||||
workspaceId: 'ws-1',
|
||||
userId: 'user-1',
|
||||
});
|
||||
|
||||
// First settle: every bounded attempt fails -> entry retained, NOT settled.
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'completed');
|
||||
expect(svc.isLocallyActive('run-1')).toBe(true);
|
||||
// F12: the give-up emits ONE explicit, greppable ERROR (run + chat context)
|
||||
// so an operator can tell "gave up, run held in memory" from a per-attempt
|
||||
// blip — distinct from the per-attempt warns.
|
||||
const gaveUp = errorSpy.mock.calls.some(
|
||||
(c) =>
|
||||
/NON-TERMINAL/.test(String(c[0])) &&
|
||||
/run-1/.test(String(c[0])) &&
|
||||
/chat-1/.test(String(c[0])),
|
||||
);
|
||||
expect(gaveUp).toBe(true);
|
||||
|
||||
// The DB recovers; a later settle now succeeds and frees the slot.
|
||||
healthy = true;
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'completed');
|
||||
expect(svc.isLocallyActive('run-1')).toBe(false);
|
||||
expect(repo.update).toHaveBeenLastCalledWith(
|
||||
'run-1',
|
||||
'ws-1',
|
||||
expect.objectContaining({ status: 'succeeded' }),
|
||||
);
|
||||
|
||||
// And it is now idempotent: a further settle no-ops (terminal row already
|
||||
// written), so a double-settle can never clobber the real status.
|
||||
const callsBefore = repo.update.mock.calls.length;
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'error', 'late');
|
||||
expect(repo.update).toHaveBeenCalledTimes(callsBefore);
|
||||
});
|
||||
|
||||
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, makeEnv() 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();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,452 @@
|
||||
import { Injectable, Logger, OnModuleInit } from '@nestjs/common';
|
||||
import { AiChatRunRepo } from '@docmost/db/repos/ai-chat/ai-chat-run.repo';
|
||||
import { AiChatRun } from '@docmost/db/types/entity.types';
|
||||
import { isUniqueViolation, violatedConstraint } from '@docmost/db/utils';
|
||||
import { EnvironmentService } from '../../integrations/environment/environment.service';
|
||||
|
||||
/** Name of the partial unique index enforcing "one active run per chat" (see the
|
||||
* ai_chat_runs migration). A 23505 on THIS constraint is the race-safe signal
|
||||
* that a concurrent turn already owns the chat — distinct from any other unique
|
||||
* collision, which must NOT be silently treated as "already active". */
|
||||
export const ONE_ACTIVE_RUN_PER_CHAT_INDEX = 'ai_chat_runs_one_active_per_chat';
|
||||
|
||||
/**
|
||||
* Thrown by {@link AiChatRunService.beginRun} when the run-row INSERT loses the
|
||||
* race for a chat's single active slot (the partial unique index rejects it with
|
||||
* a 23505). This is the AUTHORITATIVE concurrency gate: the controller's cheap
|
||||
* pre-check is only a fast-path, and a request that slips past it must NOT run
|
||||
* untracked. The caller (AiChatService.stream) translates this into a 409 and
|
||||
* aborts the turn BEFORE any AI/provider call.
|
||||
*/
|
||||
export class RunAlreadyActiveError extends Error {
|
||||
constructor(public readonly chatId: string) {
|
||||
super(`An agent run is already in progress for chat ${chatId}`);
|
||||
this.name = 'RunAlreadyActiveError';
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* The terminal status of a TURN (the #183 assistant-row lifecycle) maps onto the
|
||||
* terminal status of a RUN (#184). A turn that completed -> the run succeeded; a
|
||||
* turn that errored -> the run failed; a turn aborted (explicit user stop) -> the
|
||||
* run aborted. Pure + unit-testable.
|
||||
*/
|
||||
export type TurnTerminalStatus = 'completed' | 'error' | 'aborted';
|
||||
export type RunTerminalStatus = 'succeeded' | 'failed' | 'aborted';
|
||||
|
||||
export function mapTurnStatusToRun(
|
||||
status: TurnTerminalStatus,
|
||||
): RunTerminalStatus {
|
||||
switch (status) {
|
||||
case 'completed':
|
||||
return 'succeeded';
|
||||
case 'error':
|
||||
return 'failed';
|
||||
case 'aborted':
|
||||
return 'aborted';
|
||||
}
|
||||
}
|
||||
|
||||
/** An in-flight run held in process memory: its AbortController is the ONLY thing
|
||||
* that can stop the turn (an explicit user stop), independent of the browser
|
||||
* socket. A mere disconnect never touches it, so the run keeps going. */
|
||||
interface ActiveRun {
|
||||
controller: AbortController;
|
||||
chatId: string;
|
||||
workspaceId: string;
|
||||
}
|
||||
|
||||
/** The live handle the streaming path drives a run through (returned by
|
||||
* {@link AiChatRunService.beginRun}). The `signal` governs the agent loop's
|
||||
* abort — wired to the run, NOT to the HTTP socket. */
|
||||
export interface RunHandle {
|
||||
runId: string;
|
||||
signal: AbortSignal;
|
||||
}
|
||||
|
||||
/**
|
||||
* AiChatRunService (#184 phase 1) — owns the agent RUN as a first-class,
|
||||
* server-side lifecycle object detached from the HTTP request / browser window.
|
||||
*
|
||||
* Responsibilities:
|
||||
* - create a run row when a turn starts (inserted directly as 'running'; the
|
||||
* 'pending' status is only the column default + a reserved value, never
|
||||
* written by code in phase 1) and register an in-memory AbortController for it
|
||||
* (the explicit-stop lever);
|
||||
* - finalize the run row (succeeded / failed / aborted) and unregister it;
|
||||
* - service an EXPLICIT user stop (`requestStop`) — the ONLY thing that aborts a
|
||||
* run; a browser disconnect deliberately does NOT;
|
||||
* - crash-recovery sweep of dangling runs on startup.
|
||||
*
|
||||
* The agent loop itself still runs in AiChatService.stream (reusing #183's
|
||||
* step-granular durable write path, `consumeStream` already drains it independent
|
||||
* of the socket); this service only wraps it in a durable lifecycle and an
|
||||
* abort handle that outlives the subscriber.
|
||||
*/
|
||||
@Injectable()
|
||||
export class AiChatRunService implements OnModuleInit {
|
||||
private readonly logger = new Logger(AiChatRunService.name);
|
||||
|
||||
// runId -> ActiveRun. Process-local on purpose (phase 1 is single-process /
|
||||
// in-memory transport; a cross-process BullMQ runner + Redis stop-signal is
|
||||
// deferred to phase 2). A stop for a runId not in this map (e.g. after a
|
||||
// restart) still records `stop_requested_at` on the row.
|
||||
private readonly active = new Map<string, ActiveRun>();
|
||||
|
||||
// runIds whose TERMINAL row write has SUCCEEDED — the idempotency once-gate
|
||||
// (F6). A finalize must short-circuit only AFTER the terminal write has landed,
|
||||
// NOT merely after the in-memory entry was dropped: a transient UPDATE failure
|
||||
// has to stay retryable, so "already settled" means "row already terminal", not
|
||||
// "entry already gone". Grows by one short UUID per finished run over process
|
||||
// uptime — negligible in phase 1's single process.
|
||||
private readonly settled = new Set<string>();
|
||||
|
||||
// Bounded retry for the terminal write (F6): a single PK UPDATE can fail
|
||||
// transiently under many fire-and-forget writes (pool exhaustion, deadlock, a
|
||||
// brief connection blip). Riding out that blip in-place matters because the
|
||||
// dominant success path (streamText onFinish) settles exactly ONCE — if that
|
||||
// write is dropped and never retried, the row is stranded 'running' and the
|
||||
// one-active-run gate 409s every future turn in the chat until a restart (no
|
||||
// periodic sweep in phase 1).
|
||||
private static readonly FINALIZE_MAX_ATTEMPTS = 3;
|
||||
private static readonly FINALIZE_RETRY_BASE_MS = 50;
|
||||
|
||||
constructor(
|
||||
private readonly runRepo: AiChatRunRepo,
|
||||
private readonly environment: EnvironmentService,
|
||||
) {}
|
||||
|
||||
/**
|
||||
* Crash-recovery sweep on server start: settle EVERY run still left
|
||||
* pending/running to 'aborted' (F1 / DECISION C). The boot sweep is
|
||||
* UNCONDITIONAL — no staleness window — because phase 1 is single-process: on a
|
||||
* fresh boot any pending|running run is definitionally hung (no live runner owns
|
||||
* it), so even a fast restart (deploy/OOM within minutes of the last step) can
|
||||
* no longer leave a run stuck 'running' forever (which would make the
|
||||
* one-active-run gate 409 every future turn in that chat). The staleness window
|
||||
* is reintroduced only for the phase-2 multi-instance timer sweep, where a
|
||||
* booting replica must not abort a run another replica is actively executing.
|
||||
* Best-effort — a sweep failure is logged but MUST NOT block startup (mirrors
|
||||
* AiChatService.onModuleInit for #183).
|
||||
*/
|
||||
async onModuleInit(): Promise<void> {
|
||||
this.warnIfMultiInstance();
|
||||
try {
|
||||
// No `staleMs`: unconditional boot sweep (F1). See AiChatRunRepo.sweepRunning.
|
||||
const swept = await this.runRepo.sweepRunning();
|
||||
if (swept > 0) {
|
||||
this.logger.log(
|
||||
`Startup sweep: marked ${swept} dangling agent run(s) as 'aborted'.`,
|
||||
);
|
||||
}
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`Startup sweep of dangling runs failed: ${
|
||||
err instanceof Error ? err.message : 'unknown error'
|
||||
}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* F2 (DECISION A): autonomous runs are SINGLE-INSTANCE-ONLY in phase 1. An
|
||||
* explicit Stop, and the in-memory AbortController that backs it, are
|
||||
* process-local: a Stop only aborts the live turn if it lands on the SAME
|
||||
* replica that owns the run (it still stamps `stop_requested_at` cross-instance,
|
||||
* but nothing reads that flag during an active run yet). Cross-instance pub/sub
|
||||
* stop is phase 2. So if the deployment is horizontally scaled, warn loudly at
|
||||
* startup that a Stop may not reach a run executing on another replica.
|
||||
*
|
||||
* DETECTION: this codebase always wires the socket.io Redis adapter (REDIS_URL
|
||||
* is mandatory), so the adapter alone is NOT a horizontal-scaling signal. The
|
||||
* authoritative signal the codebase has is `CLOUD=true` (EnvironmentService
|
||||
* .isCloud()), the Docmost-cloud multi-replica deployment. We warn whenever that
|
||||
* is set, because any workspace could enable settings.ai.autonomousRuns. A
|
||||
* self-hosted operator running multiple replicas behind a load balancer is also
|
||||
* multi-instance; the deploy docs (.env.example / AGENTS.md) spell out the
|
||||
* single-instance constraint for that case.
|
||||
*/
|
||||
private warnIfMultiInstance(): void {
|
||||
if (this.environment.isCloud()) {
|
||||
this.logger.warn(
|
||||
'Autonomous agent runs (settings.ai.autonomousRuns) are SINGLE-INSTANCE-ONLY ' +
|
||||
'in phase 1: a horizontally-scaled deployment was detected (CLOUD=true). ' +
|
||||
'An explicit Stop only aborts a run executing on the same replica that owns ' +
|
||||
'it (cross-instance Stop is not yet reliable — phase 2). Run a single ' +
|
||||
'instance if you enable autonomousRuns, or keep the flag off.',
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Start a run for a turn: insert the run row (status 'running', startedAt now),
|
||||
* register a fresh AbortController for it, and return a {@link RunHandle} whose
|
||||
* `signal` the agent loop uses. The DB partial unique index guarantees at most
|
||||
* one active run per chat — a second concurrent start on the same chat REJECTS
|
||||
* at the insert (a 23505 on {@link ONE_ACTIVE_RUN_PER_CHAT_INDEX}). That
|
||||
* rejection is the AUTHORITATIVE race gate: it is surfaced as a distinct
|
||||
* {@link RunAlreadyActiveError} (NOT swallowed), so the caller turns it into a
|
||||
* 409 and never streams an untracked turn. The controller is registered AFTER a
|
||||
* successful insert so a rejected start leaks nothing.
|
||||
*/
|
||||
async beginRun(args: {
|
||||
chatId: string;
|
||||
workspaceId: string;
|
||||
userId: string;
|
||||
trigger?: string;
|
||||
}): Promise<RunHandle> {
|
||||
let run: AiChatRun;
|
||||
try {
|
||||
run = await this.runRepo.insert({
|
||||
chatId: args.chatId,
|
||||
workspaceId: args.workspaceId,
|
||||
createdBy: args.userId,
|
||||
trigger: args.trigger ?? 'user',
|
||||
status: 'running',
|
||||
startedAt: new Date(),
|
||||
});
|
||||
} catch (err) {
|
||||
// The race backstop: a concurrent turn already holds this chat's single
|
||||
// active slot, so the partial unique index rejected our insert. Surface a
|
||||
// distinct signal — the caller MUST reject this turn (409), not run it
|
||||
// untracked. Any OTHER error propagates unchanged.
|
||||
if (
|
||||
isUniqueViolation(err) &&
|
||||
violatedConstraint(err) === ONE_ACTIVE_RUN_PER_CHAT_INDEX
|
||||
) {
|
||||
throw new RunAlreadyActiveError(args.chatId);
|
||||
}
|
||||
throw err;
|
||||
}
|
||||
const controller = new AbortController();
|
||||
this.active.set(run.id, {
|
||||
controller,
|
||||
chatId: args.chatId,
|
||||
workspaceId: args.workspaceId,
|
||||
});
|
||||
return { runId: run.id, signal: controller.signal };
|
||||
}
|
||||
|
||||
/** Link the assistant message (the #183 projection) to its run. Best-effort. */
|
||||
async linkAssistantMessage(
|
||||
runId: string,
|
||||
workspaceId: string,
|
||||
assistantMessageId: string,
|
||||
): Promise<void> {
|
||||
try {
|
||||
await this.runRepo.update(runId, workspaceId, { assistantMessageId });
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`Failed to link assistant message to run ${runId}: ${
|
||||
err instanceof Error ? err.message : 'unknown error'
|
||||
}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/** Persist progress: bump the run's finished-step count. Best-effort (never
|
||||
* blocks or breaks the stream). */
|
||||
async recordStep(
|
||||
runId: string,
|
||||
workspaceId: string,
|
||||
stepCount: number,
|
||||
): Promise<void> {
|
||||
try {
|
||||
await this.runRepo.update(runId, workspaceId, { stepCount });
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`Failed to record step for run ${runId}: ${
|
||||
err instanceof Error ? err.message : 'unknown error'
|
||||
}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Finalize a run to its terminal status (succeeded / failed / aborted),
|
||||
* stamping finishedAt + any error. Best-effort, but ROBUST against a transient
|
||||
* terminal-write failure (F6) AND atomically safe against a concurrent settle.
|
||||
*
|
||||
* ATOMIC ONCE-CLAIM (the gate must close in ONE synchronous tick): two
|
||||
* finalizeRun calls for the SAME run can race — the documented real path is
|
||||
* AiChatService.stream's safety-net catch settling the turn to 'error' while a
|
||||
* streamText terminal callback (onFinish/onAbort/onError) ALSO settles it. The
|
||||
* `settled.has` check alone is NOT a gate: it is read BEFORE the awaited UPDATE,
|
||||
* so two callers can both see `false` and both write the row (last-write-wins
|
||||
* clobbers the real terminal status, and the bounded retry only widens that
|
||||
* window). The claim therefore happens via `active.delete`, a SYNCHRONOUS
|
||||
* check-and-clear with NO await between the gate and the entry removal: the
|
||||
* second concurrent caller finds the entry already gone and returns in the same
|
||||
* tick, before any UPDATE. The transition "nobody is finalizing" -> "I am
|
||||
* finalizing" is thus a single atomic step.
|
||||
*
|
||||
* ORDER MATTERS (F6): once we own the claim, the terminal UPDATE happens FIRST;
|
||||
* only once it SUCCEEDS do we record the run as settled. If the UPDATE fails on
|
||||
* every bounded attempt we RESTORE the in-memory entry, leave the run UNsettled,
|
||||
* and emit an ERROR signal that the row is left non-terminal 'running' (which
|
||||
* would 409 every future turn in the chat until recovery). An in-process retry
|
||||
* by a LATER settle is only POSSIBLE, never guaranteed: it needs (a) the entry
|
||||
* to have been restored at the give-up path AND (b) a fresh settler to arrive
|
||||
* AFTER that restore. A concurrent settler that arrives DURING the retry window
|
||||
* — while the entry is deleted for backoff and not yet restored — is consumed at
|
||||
* the synchronous `active.delete` claim (it finds nothing to delete and returns
|
||||
* a no-op), so it does NOT become an in-process retrier. The NO-streamText path
|
||||
* (the turn threw before streamText was wired, so ONLY the safety-net ever
|
||||
* settles) likewise has no second in-process settler at all. The UNCONDITIONAL
|
||||
* backstop in every case is the boot sweep on the next restart (phase 1 has no
|
||||
* periodic in-process sweep); the retained entry is bounded (cleared on restart)
|
||||
* and harmless meanwhile.
|
||||
*
|
||||
* IDEMPOTENT on SUCCESS (#184 review): the terminal write happens AT MOST ONCE
|
||||
* per run. After a successful write the once-gate keys off {@link settled} (the
|
||||
* terminal row already written) so a settle arriving AFTER the entry was already
|
||||
* dropped-and-settled returns early; a settle racing the in-flight write is
|
||||
* stopped earlier still, by the `active.delete` claim. Either way a genuine
|
||||
* double-settle collapses to a single write and a late settle can never clobber
|
||||
* the real terminal status or double-write the row.
|
||||
*/
|
||||
async finalizeRun(
|
||||
runId: string,
|
||||
workspaceId: string,
|
||||
turnStatus: TurnTerminalStatus,
|
||||
error?: string,
|
||||
): Promise<void> {
|
||||
// ---- Atomic once-claim (synchronous; NO await before the gate closes) ----
|
||||
// Already terminally written -> idempotent no-op.
|
||||
if (this.settled.has(runId)) return;
|
||||
// Capture the entry BEFORE the delete so a total-failure path can restore it.
|
||||
const entry = this.active.get(runId);
|
||||
// SYNCHRONOUS check-and-clear: the FIRST caller deletes (claims) the entry;
|
||||
// any concurrent SECOND caller finds nothing to delete and returns HERE, in
|
||||
// the same tick, before any await — so it can never reach the UPDATE.
|
||||
if (!this.active.delete(runId)) return;
|
||||
|
||||
let lastError: unknown;
|
||||
for (
|
||||
let attempt = 1;
|
||||
attempt <= AiChatRunService.FINALIZE_MAX_ATTEMPTS;
|
||||
attempt++
|
||||
) {
|
||||
try {
|
||||
await this.runRepo.update(runId, workspaceId, {
|
||||
status: mapTurnStatusToRun(turnStatus),
|
||||
finishedAt: new Date(),
|
||||
error: error ?? null,
|
||||
});
|
||||
// Terminal write landed: arm the once-gate. The entry is already gone
|
||||
// (claimed above); we do NOT restore it. The slot is now free.
|
||||
this.settled.add(runId);
|
||||
return;
|
||||
} catch (err) {
|
||||
lastError = err;
|
||||
this.logger.warn(
|
||||
`Failed to finalize run ${runId} (attempt ${attempt}/${
|
||||
AiChatRunService.FINALIZE_MAX_ATTEMPTS
|
||||
}): ${err instanceof Error ? err.message : 'unknown error'}`,
|
||||
);
|
||||
if (attempt < AiChatRunService.FINALIZE_MAX_ATTEMPTS) {
|
||||
await this.delay(AiChatRunService.FINALIZE_RETRY_BASE_MS * attempt);
|
||||
}
|
||||
}
|
||||
}
|
||||
// Every attempt failed: this is a give-up, materially worse than a per-attempt
|
||||
// blip — the row is left NON-TERMINAL ('running'), so emit ONE explicit,
|
||||
// greppable ERROR so an operator can tell "survived a blip" from "gave up, run
|
||||
// held in memory until recovery" (the last warn alone says only "attempt 3/3").
|
||||
this.logger.error(
|
||||
`Run ${runId} (chat ${entry?.chatId ?? 'unknown'}) left NON-TERMINAL ` +
|
||||
`('running'): terminal write failed after ${
|
||||
AiChatRunService.FINALIZE_MAX_ATTEMPTS
|
||||
} attempts; entry retained in memory, recovery deferred to next settle / ` +
|
||||
`boot sweep`,
|
||||
lastError,
|
||||
);
|
||||
// RESTORE the claimed entry (and leave the run UNsettled) so a LATER settle
|
||||
// that arrives AFTER this restore MAY retry the terminal write — but that
|
||||
// in-process retry is NOT guaranteed (a concurrent settler caught in the retry
|
||||
// window above is consumed at the `active.delete` claim, and the no-streamText
|
||||
// path has no second settler at all). The UNCONDITIONAL backstop in every case
|
||||
// is the boot sweep on the next restart; the restored entry is bounded and
|
||||
// cleared on restart.
|
||||
if (entry) this.active.set(runId, entry);
|
||||
}
|
||||
|
||||
/** Small async backoff between terminal-write retries (F6). Isolated so it is
|
||||
* trivial to stub/fake-time in tests. */
|
||||
private delay(ms: number): Promise<void> {
|
||||
return new Promise((resolve) => setTimeout(resolve, ms));
|
||||
}
|
||||
|
||||
/**
|
||||
* Request an EXPLICIT stop of a run (the user pressed Stop). This is the ONLY
|
||||
* thing that aborts a run — distinct from a browser disconnect, which leaves
|
||||
* the run going. Aborts the in-process controller FIRST (the only thing that
|
||||
* actually stops the run, if this replica owns it), then makes a best-effort
|
||||
* attempt to stamp `stop_requested_at` — that audit write stamps only while the
|
||||
* row is active and may be skipped on a DB error or lost to the finalize race,
|
||||
* which is acceptable since the row still settles as 'aborted'. Returns true
|
||||
* when a stop took effect (row marked and/or controller aborted), false when
|
||||
* there was nothing active to stop.
|
||||
*/
|
||||
async requestStop(runId: string, workspaceId: string): Promise<boolean> {
|
||||
const entry = this.active.get(runId);
|
||||
if (entry) {
|
||||
// Abort the live turn FIRST -> streamText onAbort fires -> the partial is
|
||||
// persisted (#183) and finalizeRun settles the row as 'aborted'. This is
|
||||
// the ONLY thing that aborts a run, so it MUST NOT be hostage to the audit
|
||||
// write below: a transient failure on `markStopRequested` (pool exhaustion,
|
||||
// deadlock, dropped connection) must never leave the run executing despite
|
||||
// an explicit Stop. At worst only the `stop_requested_at` timestamp is lost.
|
||||
entry.controller.abort();
|
||||
}
|
||||
// Record `stop_requested_at` (best-effort). A transient DB failure here is
|
||||
// logged and treated as `marked = false`; the abort above already took
|
||||
// effect, so we never rethrow and skip stopping the run. Note: because
|
||||
// markStopRequested only stamps while the row is active, aborting first means
|
||||
// even a healthy write can lose the race against the resulting finalize and
|
||||
// skip the stamp — acceptable, as the row still settles as 'aborted' and only
|
||||
// this audit timestamp may be lost.
|
||||
let marked: unknown;
|
||||
try {
|
||||
marked = await this.runRepo.markStopRequested(runId, workspaceId);
|
||||
} catch (err) {
|
||||
marked = undefined;
|
||||
this.logger.warn(
|
||||
`requestStop: markStopRequested failed for run ${runId} ` +
|
||||
`(stop_requested_at not recorded); abort already issued: ` +
|
||||
`${err instanceof Error ? err.message : String(err)}`,
|
||||
);
|
||||
}
|
||||
return Boolean(marked) || Boolean(entry);
|
||||
}
|
||||
|
||||
/** Latest persisted run for a chat — the reconnect target (an in-flight or
|
||||
* finished run). Pure read-through to the repo. */
|
||||
getLatestForChat(
|
||||
chatId: string,
|
||||
workspaceId: string,
|
||||
): Promise<AiChatRun | undefined> {
|
||||
return this.runRepo.findLatestByChat(chatId, workspaceId);
|
||||
}
|
||||
|
||||
/** Fetch a run by id (workspace-scoped). Used to resolve + ownership-check an
|
||||
* explicit stop targeting a runId. */
|
||||
getRun(runId: string, workspaceId: string): Promise<AiChatRun | undefined> {
|
||||
return this.runRepo.findById(runId, workspaceId);
|
||||
}
|
||||
|
||||
/** The active run on a chat, if any (used to reject a concurrent start with a
|
||||
* clean 409 before committing to the stream). */
|
||||
getActiveForChat(
|
||||
chatId: string,
|
||||
workspaceId: string,
|
||||
): Promise<AiChatRun | undefined> {
|
||||
return this.runRepo.findActiveByChat(chatId, workspaceId);
|
||||
}
|
||||
|
||||
/** Test/diagnostic seam: whether this replica is holding a live controller for
|
||||
* the run. */
|
||||
isLocallyActive(runId: string): boolean {
|
||||
return this.active.has(runId);
|
||||
}
|
||||
}
|
||||
@@ -25,6 +25,7 @@ describe('AiChatController.boundChat', () => {
|
||||
};
|
||||
const controller = new AiChatController(
|
||||
{} as never,
|
||||
{} as never, // aiChatRunService
|
||||
aiChatRepo as never,
|
||||
{} as never,
|
||||
{} as never,
|
||||
|
||||
@@ -53,6 +53,7 @@ describe('AiChatController.export', () => {
|
||||
};
|
||||
const controller = new AiChatController(
|
||||
{} as never,
|
||||
{} as never, // aiChatRunService
|
||||
aiChatRepo as never,
|
||||
aiChatMessageRepo as never,
|
||||
{} as never,
|
||||
|
||||
@@ -0,0 +1,164 @@
|
||||
import { BadRequestException, ForbiddenException } from '@nestjs/common';
|
||||
import { AiChatController } from './ai-chat.controller';
|
||||
import type { User, Workspace } from '@docmost/db/types/entity.types';
|
||||
|
||||
/**
|
||||
* Wiring spec for the #184 run-reconnect / run-stop endpoints
|
||||
* (`POST /ai-chat/run` and `POST /ai-chat/stop`). Both are OWNER-gated via
|
||||
* assertOwnedChat (the requesting user must own the chat) and NOT flag-gated.
|
||||
* Exercised with hand-rolled mocks — no Nest graph, no DB. The controller's
|
||||
* constructor order is (aiChatService, aiChatRunService, aiChatRepo,
|
||||
* aiChatMessageRepo, aiTranscription).
|
||||
*/
|
||||
describe('AiChatController run endpoints (#184)', () => {
|
||||
const user = { id: 'u1' } as User;
|
||||
const workspace = { id: 'ws1' } as Workspace;
|
||||
|
||||
function makeController(opts: {
|
||||
chat?: unknown; // what aiChatRepo.findById returns (owner-gate)
|
||||
run?: unknown; // getLatestForChat / getRun result
|
||||
activeRun?: unknown; // getActiveForChat result
|
||||
message?: unknown; // aiChatMessageRepo.findById result
|
||||
stopped?: boolean; // requestStop result
|
||||
}) {
|
||||
const aiChatRunService = {
|
||||
getLatestForChat: jest.fn().mockResolvedValue(opts.run),
|
||||
getRun: jest.fn().mockResolvedValue(opts.run),
|
||||
getActiveForChat: jest.fn().mockResolvedValue(opts.activeRun),
|
||||
requestStop: jest.fn().mockResolvedValue(opts.stopped ?? false),
|
||||
};
|
||||
const aiChatRepo = {
|
||||
findById: jest.fn().mockResolvedValue(opts.chat),
|
||||
};
|
||||
const aiChatMessageRepo = {
|
||||
findById: jest.fn().mockResolvedValue(opts.message),
|
||||
};
|
||||
const controller = new AiChatController(
|
||||
{} as never, // aiChatService
|
||||
aiChatRunService as never,
|
||||
aiChatRepo as never,
|
||||
aiChatMessageRepo as never,
|
||||
{} as never, // aiTranscription
|
||||
{} as never, // pageRepo
|
||||
);
|
||||
return { controller, aiChatRunService, aiChatRepo, aiChatMessageRepo };
|
||||
}
|
||||
|
||||
describe('POST /ai-chat/run (getRun)', () => {
|
||||
it('owner-gates: a chat the user does not own throws ForbiddenException', async () => {
|
||||
const { controller, aiChatRunService } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'someone-else' },
|
||||
});
|
||||
await expect(
|
||||
controller.getRun({ chatId: 'c1' }, user, workspace),
|
||||
).rejects.toBeInstanceOf(ForbiddenException);
|
||||
// It must NOT reach the run lookup once the owner-gate fails.
|
||||
expect(aiChatRunService.getLatestForChat).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('returns { run: null, message: null } when the chat has never had a run', async () => {
|
||||
const { controller, aiChatRunService } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'u1' },
|
||||
run: undefined,
|
||||
});
|
||||
const res = await controller.getRun({ chatId: 'c1' }, user, workspace);
|
||||
expect(res).toEqual({ run: null, message: null });
|
||||
expect(aiChatRunService.getLatestForChat).toHaveBeenCalledWith(
|
||||
'c1',
|
||||
'ws1',
|
||||
);
|
||||
});
|
||||
|
||||
it('returns the run and its projected assistant message', async () => {
|
||||
const run = { id: 'run-1', chatId: 'c1', assistantMessageId: 'm1' };
|
||||
const message = { id: 'm1', role: 'assistant' };
|
||||
const { controller, aiChatMessageRepo } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'u1' },
|
||||
run,
|
||||
message,
|
||||
});
|
||||
const res = await controller.getRun({ chatId: 'c1' }, user, workspace);
|
||||
expect(res).toEqual({ run, message });
|
||||
expect(aiChatMessageRepo.findById).toHaveBeenCalledWith('m1', 'ws1');
|
||||
});
|
||||
|
||||
it('returns message: null when the run has no linked assistant message', async () => {
|
||||
const run = { id: 'run-1', chatId: 'c1', assistantMessageId: null };
|
||||
const { controller, aiChatMessageRepo } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'u1' },
|
||||
run,
|
||||
});
|
||||
const res = await controller.getRun({ chatId: 'c1' }, user, workspace);
|
||||
expect(res).toEqual({ run, message: null });
|
||||
expect(aiChatMessageRepo.findById).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('POST /ai-chat/stop (stopRun)', () => {
|
||||
it('throws BadRequestException when neither runId nor chatId is given', async () => {
|
||||
const { controller } = makeController({});
|
||||
await expect(
|
||||
controller.stopRun({}, user, workspace),
|
||||
).rejects.toBeInstanceOf(BadRequestException);
|
||||
});
|
||||
|
||||
it('stops by runId: owner-gates via the run’s chat, then requests the stop', async () => {
|
||||
const { controller, aiChatRunService, aiChatRepo } = makeController({
|
||||
run: { id: 'run-1', chatId: 'c1' },
|
||||
chat: { id: 'c1', creatorId: 'u1' },
|
||||
stopped: true,
|
||||
});
|
||||
const res = await controller.stopRun({ runId: 'run-1' }, user, workspace);
|
||||
expect(res).toEqual({ stopped: true });
|
||||
expect(aiChatRunService.getRun).toHaveBeenCalledWith('run-1', 'ws1');
|
||||
expect(aiChatRepo.findById).toHaveBeenCalledWith('c1', 'ws1');
|
||||
expect(aiChatRunService.requestStop).toHaveBeenCalledWith('run-1', 'ws1');
|
||||
});
|
||||
|
||||
it('stops by runId: a foreign run’s chat throws ForbiddenException (no stop)', async () => {
|
||||
const { controller, aiChatRunService } = makeController({
|
||||
run: { id: 'run-1', chatId: 'c1' },
|
||||
chat: { id: 'c1', creatorId: 'someone-else' },
|
||||
});
|
||||
await expect(
|
||||
controller.stopRun({ runId: 'run-1' }, user, workspace),
|
||||
).rejects.toBeInstanceOf(ForbiddenException);
|
||||
expect(aiChatRunService.requestStop).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('stops by runId: an unknown run reports { stopped: false }', async () => {
|
||||
const { controller, aiChatRunService } = makeController({
|
||||
run: undefined,
|
||||
});
|
||||
const res = await controller.stopRun({ runId: 'gone' }, user, workspace);
|
||||
expect(res).toEqual({ stopped: false });
|
||||
expect(aiChatRunService.requestStop).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('stops by chatId: owner-gates, resolves the active run, requests the stop', async () => {
|
||||
const { controller, aiChatRunService, aiChatRepo } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'u1' },
|
||||
activeRun: { id: 'run-9' },
|
||||
stopped: true,
|
||||
});
|
||||
const res = await controller.stopRun({ chatId: 'c1' }, user, workspace);
|
||||
expect(res).toEqual({ stopped: true });
|
||||
expect(aiChatRepo.findById).toHaveBeenCalledWith('c1', 'ws1');
|
||||
expect(aiChatRunService.getActiveForChat).toHaveBeenCalledWith(
|
||||
'c1',
|
||||
'ws1',
|
||||
);
|
||||
expect(aiChatRunService.requestStop).toHaveBeenCalledWith('run-9', 'ws1');
|
||||
});
|
||||
|
||||
it('stops by chatId: reports { stopped: false } when no run is active', async () => {
|
||||
const { controller, aiChatRunService } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'u1' },
|
||||
activeRun: undefined,
|
||||
});
|
||||
const res = await controller.stopRun({ chatId: 'c1' }, user, workspace);
|
||||
expect(res).toEqual({ stopped: false });
|
||||
expect(aiChatRunService.requestStop).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -1,6 +1,7 @@
|
||||
import {
|
||||
BadRequestException,
|
||||
Body,
|
||||
ConflictException,
|
||||
Controller,
|
||||
ForbiddenException,
|
||||
HttpCode,
|
||||
@@ -20,7 +21,13 @@ import { JwtAuthGuard } from '../../common/guards/jwt-auth.guard';
|
||||
import { AuthUser } from '../../common/decorators/auth-user.decorator';
|
||||
import { AuthWorkspace } from '../../common/decorators/auth-workspace.decorator';
|
||||
import { SkipTransform } from '../../common/decorators/skip-transform.decorator';
|
||||
import { AiChat, User, Workspace } from '@docmost/db/types/entity.types';
|
||||
import {
|
||||
AiChat,
|
||||
AiChatMessage,
|
||||
AiChatRun,
|
||||
User,
|
||||
Workspace,
|
||||
} from '@docmost/db/types/entity.types';
|
||||
import { PaginationOptions } from '@docmost/db/pagination/pagination-options';
|
||||
import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo';
|
||||
import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
|
||||
@@ -28,7 +35,12 @@ import { PageRepo } from '@docmost/db/repos/page/page.repo';
|
||||
import { UserThrottlerGuard } from '../../integrations/throttle/user-throttler.guard';
|
||||
import { AI_CHAT_THROTTLER } from '../../integrations/throttle/throttler-names';
|
||||
import { FileInterceptor } from '../../common/interceptors/file.interceptor';
|
||||
import { AiChatService, AiChatStreamBody } from './ai-chat.service';
|
||||
import {
|
||||
AiChatRunHooks,
|
||||
AiChatService,
|
||||
AiChatStreamBody,
|
||||
} from './ai-chat.service';
|
||||
import { AiChatRunService } from './ai-chat-run.service';
|
||||
import { AiTranscriptionService } from './ai-transcription.service';
|
||||
import {
|
||||
BoundChatDto,
|
||||
@@ -36,7 +48,9 @@ import {
|
||||
ExportChatDto,
|
||||
GeneratePageTitleDto,
|
||||
GetChatMessagesDto,
|
||||
GetRunDto,
|
||||
RenameChatDto,
|
||||
StopRunDto,
|
||||
} from './dto/ai-chat.dto';
|
||||
import { describeProviderError } from '../../integrations/ai/ai-error.util';
|
||||
import { buildChatMarkdown } from './chat-markdown.util';
|
||||
@@ -53,6 +67,7 @@ export class AiChatController {
|
||||
|
||||
constructor(
|
||||
private readonly aiChatService: AiChatService,
|
||||
private readonly aiChatRunService: AiChatRunService,
|
||||
private readonly aiChatRepo: AiChatRepo,
|
||||
private readonly aiChatMessageRepo: AiChatMessageRepo,
|
||||
private readonly aiTranscription: AiTranscriptionService,
|
||||
@@ -149,6 +164,75 @@ export class AiChatController {
|
||||
return { markdown };
|
||||
}
|
||||
|
||||
/**
|
||||
* Reconnect to the latest run of a chat (#184 phase 1). Returns the run's
|
||||
* persisted lifecycle state ({ status, error, stepCount, timings, ... }) plus
|
||||
* the assistant message it projects (the partial/final output) — the DB is the
|
||||
* source of truth, so this works for an in-flight run (the browser dropped, the
|
||||
* run kept going) and a finished one alike. Owner-gated via assertOwnedChat.
|
||||
* `{ run: null }` when the chat has never had a run.
|
||||
*/
|
||||
@HttpCode(HttpStatus.OK)
|
||||
@Post('run')
|
||||
async getRun(
|
||||
@Body() dto: GetRunDto,
|
||||
@AuthUser() user: User,
|
||||
@AuthWorkspace() workspace: Workspace,
|
||||
): Promise<{ run: AiChatRun | null; message: AiChatMessage | null }> {
|
||||
await this.assertOwnedChat(dto.chatId, user, workspace);
|
||||
const run = await this.aiChatRunService.getLatestForChat(
|
||||
dto.chatId,
|
||||
workspace.id,
|
||||
);
|
||||
if (!run) return { run: null, message: null };
|
||||
const message = run.assistantMessageId
|
||||
? await this.aiChatMessageRepo.findById(
|
||||
run.assistantMessageId,
|
||||
workspace.id,
|
||||
)
|
||||
: undefined;
|
||||
return { run, message: message ?? null };
|
||||
}
|
||||
|
||||
/**
|
||||
* Explicitly STOP an agent run (#184 phase 1) — the user pressed Stop. This is
|
||||
* the ONLY thing that ends a detached run; a browser disconnect deliberately
|
||||
* does not. Target by `runId` (from the streamed start metadata) or by `chatId`
|
||||
* (stop whatever run is active on it). Owner-gated. Returns
|
||||
* `{ stopped }` — false when there was nothing active to stop.
|
||||
*/
|
||||
@HttpCode(HttpStatus.OK)
|
||||
@Post('stop')
|
||||
async stopRun(
|
||||
@Body() dto: StopRunDto,
|
||||
@AuthUser() user: User,
|
||||
@AuthWorkspace() workspace: Workspace,
|
||||
): Promise<{ stopped: boolean }> {
|
||||
let runId = dto.runId;
|
||||
if (!runId && !dto.chatId) {
|
||||
throw new BadRequestException('runId or chatId is required');
|
||||
}
|
||||
if (runId) {
|
||||
// Resolve the run to its chat and owner-gate via that chat.
|
||||
const run = await this.aiChatRunService.getRun(runId, workspace.id);
|
||||
if (!run) return { stopped: false };
|
||||
await this.assertOwnedChat(run.chatId, user, workspace);
|
||||
} else {
|
||||
await this.assertOwnedChat(dto.chatId!, user, workspace);
|
||||
const active = await this.aiChatRunService.getActiveForChat(
|
||||
dto.chatId!,
|
||||
workspace.id,
|
||||
);
|
||||
if (!active) return { stopped: false };
|
||||
runId = active.id;
|
||||
}
|
||||
const stopped = await this.aiChatRunService.requestStop(
|
||||
runId,
|
||||
workspace.id,
|
||||
);
|
||||
return { stopped };
|
||||
}
|
||||
|
||||
/** Rename a chat. */
|
||||
@HttpCode(HttpStatus.OK)
|
||||
@Post('rename')
|
||||
@@ -200,11 +284,20 @@ export class AiChatController {
|
||||
@AuthWorkspace() workspace: Workspace,
|
||||
): Promise<void> {
|
||||
// A7 gate: the workspace must have AI chat explicitly enabled.
|
||||
const settings = (workspace.settings ?? {}) as { ai?: { chat?: boolean } };
|
||||
const settings = (workspace.settings ?? {}) as {
|
||||
ai?: { chat?: boolean; autonomousRuns?: boolean };
|
||||
};
|
||||
if (settings.ai?.chat !== true) {
|
||||
throw new ForbiddenException('AI chat is disabled');
|
||||
}
|
||||
|
||||
// #184 phase 1 flag: when ON, the turn becomes a detached, durable RUN — its
|
||||
// lifecycle is tracked in ai_chat_runs, a browser disconnect no longer aborts
|
||||
// it, and only an explicit /ai-chat/stop ends it. When OFF (the default) the
|
||||
// turn is socket-bound exactly as before, so existing deployments are
|
||||
// unaffected.
|
||||
const autonomousRuns = settings.ai?.autonomousRuns === true;
|
||||
|
||||
const sessionId = (req.raw as { sessionId?: string }).sessionId;
|
||||
if (!sessionId) {
|
||||
// The chat requires an interactive session to mint loopback tokens
|
||||
@@ -228,6 +321,58 @@ export class AiChatController {
|
||||
// HttpException) instead of breaking mid-stream.
|
||||
const model = await this.aiChatService.getChatModel(workspace.id, role);
|
||||
|
||||
// #184: one active run per chat. For an EXISTING chat reject a concurrent
|
||||
// start with a clean 409 BEFORE hijack (the common double-submit / second-tab
|
||||
// case), so the user gets JSON, not a mid-stream error. A brand-new chat
|
||||
// (no chatId) cannot have a prior run, and the DB partial unique index is the
|
||||
// backstop against any race that slips past this check.
|
||||
if (autonomousRuns && body.chatId) {
|
||||
const active = await this.aiChatRunService.getActiveForChat(
|
||||
body.chatId,
|
||||
workspace.id,
|
||||
);
|
||||
if (active) {
|
||||
throw new ConflictException({
|
||||
message: 'An agent run is already in progress for this chat',
|
||||
code: 'A_RUN_ALREADY_ACTIVE',
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Run-lifecycle hooks (#184), only when the flag is on. They wrap the turn in
|
||||
// a durable run whose abort is governed by the run (explicit stop), persist
|
||||
// its progress, and settle its terminal status — see AiChatRunService.
|
||||
const runHooks: AiChatRunHooks | undefined = autonomousRuns
|
||||
? {
|
||||
begin: (chatId) =>
|
||||
this.aiChatRunService.beginRun({
|
||||
chatId,
|
||||
workspaceId: workspace.id,
|
||||
userId: user.id,
|
||||
trigger: 'user',
|
||||
}),
|
||||
onAssistantSeeded: (runId, messageId) =>
|
||||
this.aiChatRunService.linkAssistantMessage(
|
||||
runId,
|
||||
workspace.id,
|
||||
messageId,
|
||||
),
|
||||
onStep: (runId, stepCount) =>
|
||||
void this.aiChatRunService.recordStep(
|
||||
runId,
|
||||
workspace.id,
|
||||
stepCount,
|
||||
),
|
||||
onSettled: (runId, status, error) =>
|
||||
this.aiChatRunService.finalizeRun(
|
||||
runId,
|
||||
workspace.id,
|
||||
status,
|
||||
error,
|
||||
),
|
||||
}
|
||||
: undefined;
|
||||
|
||||
// Abort the agent loop when the client disconnects. `close` also fires on
|
||||
// normal completion, so only abort when the response has not finished
|
||||
// writing (a genuine disconnect). `once` fires at most once and self-removes;
|
||||
@@ -242,18 +387,44 @@ export class AiChatController {
|
||||
// A genuine disconnect leaves the response unfinished (unlike a normal
|
||||
// completion, which also fires `close`). Such a drop — e.g. a reverse
|
||||
// proxy cutting the SSE mid-answer — is otherwise invisible server-side,
|
||||
// so log it here before aborting the agent loop.
|
||||
// so log it here.
|
||||
if (!res.raw.writableEnded) {
|
||||
this.logger.warn(
|
||||
`AI chat stream: client disconnected before completion; aborting turn ` +
|
||||
`(elapsed=${Date.now() - reqStartedAt}ms since request received)`,
|
||||
);
|
||||
controller.abort();
|
||||
if (autonomousRuns) {
|
||||
// #184: the turn is a DETACHED run. A disconnect must NOT abort it —
|
||||
// the run keeps executing and persisting server-side; the client
|
||||
// reconnects via /ai-chat/run (or re-stops via /ai-chat/stop). Log only.
|
||||
this.logger.log(
|
||||
`AI chat stream: client disconnected; run continues server-side ` +
|
||||
`(elapsed=${Date.now() - reqStartedAt}ms since request received)`,
|
||||
);
|
||||
} else {
|
||||
this.logger.warn(
|
||||
`AI chat stream: client disconnected before completion; aborting turn ` +
|
||||
`(elapsed=${Date.now() - reqStartedAt}ms since request received)`,
|
||||
);
|
||||
controller.abort();
|
||||
}
|
||||
}
|
||||
};
|
||||
req.raw.once('close', onClose);
|
||||
res.raw.once('finish', () => req.raw.off('close', onClose));
|
||||
|
||||
// #184: in detached mode the turn is NOT aborted on disconnect, so the SDK's
|
||||
// pipe keeps writing to a socket the client may have dropped — for the rest of
|
||||
// the (continuing) run. A write to the dead socket can emit an 'error' on the
|
||||
// raw response; without a listener that surfaces as an unhandled error event.
|
||||
// Swallow it (the run continues server-side regardless). Legacy mode aborts on
|
||||
// disconnect, so it does not need this and keeps its exact prior behavior.
|
||||
if (autonomousRuns) {
|
||||
res.raw.on('error', (err) => {
|
||||
this.logger.debug(
|
||||
`AI chat detached stream: post-disconnect socket error swallowed: ${
|
||||
err instanceof Error ? err.message : String(err)
|
||||
}`,
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
// Commit to streaming: hijack so Fastify stops managing the response and
|
||||
// the AI SDK can write the UI-message stream directly to the Node socket.
|
||||
res.hijack();
|
||||
@@ -268,15 +439,32 @@ export class AiChatController {
|
||||
signal: controller.signal,
|
||||
model,
|
||||
role,
|
||||
// #184: present only when the flag is on; wraps the turn in a durable run.
|
||||
runHooks,
|
||||
});
|
||||
} catch (err) {
|
||||
// Any failure AFTER hijack can no longer send a clean JSON error, so emit
|
||||
// a minimal error on the raw socket if nothing has been written yet.
|
||||
this.logger.error('AI chat stream failed', err as Error);
|
||||
// Any failure AFTER hijack can no longer go through Nest's exception
|
||||
// filter, so emit the error on the raw socket if nothing has been written
|
||||
// yet. The lost-the-race 409 (RunAlreadyActiveError -> ConflictException)
|
||||
// is raised by stream() BEFORE it writes a byte, so headers are still
|
||||
// unsent here: honor the HttpException's real status + body (a clean 409),
|
||||
// not a blanket 500. Everything else stays a 500.
|
||||
const isHttp = err instanceof HttpException;
|
||||
if (!isHttp) {
|
||||
this.logger.error('AI chat stream failed', err as Error);
|
||||
}
|
||||
if (!res.raw.headersSent) {
|
||||
res.raw.statusCode = 500;
|
||||
const status = isHttp ? err.getStatus() : 500;
|
||||
const payload = isHttp
|
||||
? err.getResponse()
|
||||
: { error: 'Internal server error' };
|
||||
res.raw.statusCode = status;
|
||||
res.raw.setHeader('Content-Type', 'application/json');
|
||||
res.raw.end(JSON.stringify({ error: 'Internal server error' }));
|
||||
res.raw.end(
|
||||
JSON.stringify(
|
||||
typeof payload === 'string' ? { message: payload } : payload,
|
||||
),
|
||||
);
|
||||
} else if (!res.raw.writableEnded) {
|
||||
res.raw.end();
|
||||
}
|
||||
|
||||
@@ -57,6 +57,7 @@ describe('AiChatController.generatePageTitle', () => {
|
||||
const aiChatService = { generatePageTitle: generate };
|
||||
const controller = new AiChatController(
|
||||
aiChatService as never,
|
||||
{} as never, // aiChatRunService
|
||||
{} as never,
|
||||
{} as never,
|
||||
{} as never,
|
||||
|
||||
@@ -3,6 +3,7 @@ import { AiModule } from '../../integrations/ai/ai.module';
|
||||
import { TokenModule } from '../auth/token.module';
|
||||
import { AiChatController } from './ai-chat.controller';
|
||||
import { AiChatService } from './ai-chat.service';
|
||||
import { AiChatRunService } from './ai-chat-run.service';
|
||||
import { AiTranscriptionService } from './ai-transcription.service';
|
||||
import { AiChatToolsService } from './tools/ai-chat-tools.service';
|
||||
import { EmbeddingModule } from './embedding/embedding.module';
|
||||
@@ -42,6 +43,7 @@ import { PublicShareChatToolsService } from './tools/public-share-chat-tools.ser
|
||||
controllers: [AiChatController, PublicShareChatController],
|
||||
providers: [
|
||||
AiChatService,
|
||||
AiChatRunService,
|
||||
AiTranscriptionService,
|
||||
AiChatToolsService,
|
||||
PublicShareChatService,
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
import { Logger } from '@nestjs/common';
|
||||
import { AiChatService } from './ai-chat.service';
|
||||
import { AiChatService, AiChatRunHooks } from './ai-chat.service';
|
||||
import { AiChatRunService } from './ai-chat-run.service';
|
||||
import type { User, Workspace } from '@docmost/db/types/entity.types';
|
||||
|
||||
/**
|
||||
* Lifecycle unit tests for AiChatService.onModuleInit (#183 crash-recovery
|
||||
@@ -61,3 +63,99 @@ describe('AiChatService.onModuleInit (startup sweep)', () => {
|
||||
expect(String(warnSpy.mock.calls[0][0])).toContain('db unavailable');
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* #184 CRITICAL run-lifecycle safety net (review fix). A transient failure
|
||||
* AFTER a successful beginRun but BEFORE streamText's terminal callbacks own the
|
||||
* lifecycle must STILL settle the run — otherwise the run row is stuck 'running'
|
||||
* forever (sweepRunning only runs at startup) and the partial unique index + the
|
||||
* controller pre-check 409 every future turn in that chat until a restart. Here
|
||||
* we model the very first bare await after beginRun (the user-message insert)
|
||||
* throwing, wiring the run hooks to a REAL AiChatRunService (mock repo) exactly
|
||||
* as the controller does, and assert the run is settled to 'error' and its
|
||||
* in-memory entry dropped (so a follow-up turn would NOT be 409'd).
|
||||
*/
|
||||
describe('AiChatService.stream run-lifecycle safety net (#184)', () => {
|
||||
const user = { id: 'u1' } as User;
|
||||
const workspace = { id: 'ws1' } as Workspace;
|
||||
|
||||
afterEach(() => jest.restoreAllMocks());
|
||||
|
||||
it('an exception after beginRun settles the run to error and drops the in-memory entry', async () => {
|
||||
jest.spyOn(Logger.prototype, 'error').mockImplementation(() => undefined);
|
||||
|
||||
// Real run service over a mock repo, so finalizeRun's in-memory bookkeeping
|
||||
// (active.delete) is exercised for real.
|
||||
const runRepo = {
|
||||
insert: jest.fn().mockResolvedValue({ id: 'run-1', status: 'running' }),
|
||||
update: jest.fn().mockResolvedValue({ id: 'run-1' }),
|
||||
};
|
||||
const runService = new AiChatRunService(runRepo as never, { isCloud: () => false } as never);
|
||||
|
||||
// The user-message insert (the first bare await after beginRun) throws.
|
||||
const aiChatMessageRepo = {
|
||||
insert: jest.fn().mockRejectedValue(new Error('insert boom')),
|
||||
};
|
||||
const aiChatRepo = {
|
||||
// Existing chat -> chatId stays, no new-chat insert path.
|
||||
findById: jest.fn().mockResolvedValue({ id: 'chat-1', creatorId: 'u1' }),
|
||||
};
|
||||
|
||||
const service = new AiChatService(
|
||||
{} as never, // ai
|
||||
aiChatRepo as never,
|
||||
aiChatMessageRepo as never,
|
||||
{} as never, // aiChatPageSnapshotRepo
|
||||
{} as never, // aiSettings
|
||||
{} as never, // tools
|
||||
{} as never, // mcpClients
|
||||
{} as never, // aiAgentRoleRepo
|
||||
{} as never, // pageRepo
|
||||
{} as never, // pageAccess
|
||||
{} as never, // environment
|
||||
);
|
||||
|
||||
const runHooks: AiChatRunHooks = {
|
||||
begin: (chatId) =>
|
||||
runService.beginRun({
|
||||
chatId,
|
||||
workspaceId: workspace.id,
|
||||
userId: user.id,
|
||||
trigger: 'user',
|
||||
}),
|
||||
onSettled: (runId, status, error) =>
|
||||
runService.finalizeRun(runId, workspace.id, status, error),
|
||||
};
|
||||
|
||||
await expect(
|
||||
service.stream({
|
||||
user,
|
||||
workspace,
|
||||
sessionId: 'sess',
|
||||
body: {
|
||||
chatId: 'chat-1',
|
||||
messages: [
|
||||
{ id: 'm', role: 'user', parts: [{ type: 'text', text: 'hi' }] },
|
||||
],
|
||||
},
|
||||
res: {} as never,
|
||||
signal: new AbortController().signal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
runHooks,
|
||||
}),
|
||||
).rejects.toThrow('insert boom');
|
||||
|
||||
// The run was begun...
|
||||
expect(runRepo.insert).toHaveBeenCalledTimes(1);
|
||||
// ...then settled to a terminal FAILED status by the safety net...
|
||||
expect(runRepo.update).toHaveBeenCalledTimes(1);
|
||||
expect(runRepo.update).toHaveBeenCalledWith(
|
||||
'run-1',
|
||||
'ws1',
|
||||
expect.objectContaining({ status: 'failed' }),
|
||||
);
|
||||
// ...and the in-memory entry is gone, so a follow-up turn is NOT 409'd.
|
||||
expect(runService.isLocallyActive('run-1')).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -0,0 +1,489 @@
|
||||
import { ConflictException, Logger } from '@nestjs/common';
|
||||
|
||||
// Mock the AI SDK so we can PROVE no provider call is made for the turn we are
|
||||
// about to reject. The race rejection happens at runHooks.begin(), long before
|
||||
// any streamText/generateText, so these never resolve a real model.
|
||||
jest.mock('ai', () => ({
|
||||
streamText: jest.fn(),
|
||||
generateText: jest.fn(),
|
||||
convertToModelMessages: jest.fn(() => []),
|
||||
stepCountIs: jest.fn(() => () => false),
|
||||
}));
|
||||
|
||||
import { streamText, generateText } from 'ai';
|
||||
import { AiChatService } from './ai-chat.service';
|
||||
import { RunAlreadyActiveError } from './ai-chat-run.service';
|
||||
|
||||
/**
|
||||
* Race-closure coverage for the "one active run per chat" guard (#184).
|
||||
*
|
||||
* THE BUG: two simultaneous POST /ai-chat/stream on the same chat both pass the
|
||||
* controller's cheap pre-check (TOCTOU), so the loser's run-row INSERT hits the
|
||||
* partial unique index. Previously that 23505 was SWALLOWED and the second turn
|
||||
* streamed UNTRACKED (no runId, not stoppable). THE FIX: beginRun surfaces a
|
||||
* RunAlreadyActiveError and stream() turns it into a 409 BEFORE any AI call —
|
||||
* the second turn never runs.
|
||||
*/
|
||||
describe('AiChatService.stream — concurrent-run race rejection (#184)', () => {
|
||||
const streamTextMock = streamText as unknown as jest.Mock;
|
||||
const generateTextMock = generateText as unknown as jest.Mock;
|
||||
|
||||
beforeEach(() => {
|
||||
streamTextMock.mockReset();
|
||||
generateTextMock.mockReset();
|
||||
});
|
||||
|
||||
// Minimal service whose only reachable deps before begin() are aiChatRepo
|
||||
// (resolve the existing chat) — everything past begin must remain untouched.
|
||||
function makeService(beginImpl: () => Promise<unknown>) {
|
||||
const aiChatMessageRepo = { insert: jest.fn() };
|
||||
const aiChatRepo = {
|
||||
// An existing chat: stream keeps the supplied chatId and skips creation.
|
||||
findById: jest.fn(async () => ({ id: 'chat-1', workspaceId: 'ws-1' })),
|
||||
insert: jest.fn(),
|
||||
};
|
||||
const svc = new AiChatService(
|
||||
{} as never, // ai
|
||||
aiChatRepo as never,
|
||||
aiChatMessageRepo as never,
|
||||
{} as never, // aiChatPageSnapshotRepo
|
||||
{} as never, // aiSettings
|
||||
{} as never, // tools
|
||||
{} as never, // mcpClients
|
||||
{} as never, // aiAgentRoleRepo
|
||||
{} as never, // pageRepo
|
||||
{} as never, // pageAccess
|
||||
{ isAiChatDeferredToolsEnabled: () => false } as never, // environment
|
||||
);
|
||||
const begin = jest.fn(beginImpl);
|
||||
return { svc, begin, aiChatRepo, aiChatMessageRepo };
|
||||
}
|
||||
|
||||
const baseArgs = (begin: jest.Mock) => ({
|
||||
user: { id: 'user-1' } as never,
|
||||
workspace: { id: 'ws-1' } as never,
|
||||
sessionId: 'sess-1',
|
||||
body: { chatId: 'chat-1', messages: [] } as never,
|
||||
res: { raw: {} } as never,
|
||||
signal: new AbortController().signal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
runHooks: {
|
||||
begin,
|
||||
onAssistantSeeded: jest.fn(),
|
||||
onStep: jest.fn(),
|
||||
onSettled: jest.fn(),
|
||||
} as never,
|
||||
});
|
||||
|
||||
it('rejects the racer with a 409 ConflictException BEFORE any AI call, and never persists an untracked turn', async () => {
|
||||
// begin loses the unique-index race -> RunAlreadyActiveError.
|
||||
const { svc, begin, aiChatMessageRepo } = makeService(() => {
|
||||
throw new RunAlreadyActiveError('chat-1');
|
||||
});
|
||||
|
||||
const promise = svc.stream(baseArgs(begin));
|
||||
|
||||
await expect(promise).rejects.toBeInstanceOf(ConflictException);
|
||||
await promise.catch((err: ConflictException) => {
|
||||
expect(err.getStatus()).toBe(409);
|
||||
expect((err.getResponse() as { code?: string }).code).toBe(
|
||||
'A_RUN_ALREADY_ACTIVE',
|
||||
);
|
||||
});
|
||||
|
||||
// The decisive assertions: the rejected racer spent NO tokens and left NO
|
||||
// untracked turn behind.
|
||||
expect(begin).toHaveBeenCalledTimes(1);
|
||||
expect(streamTextMock).not.toHaveBeenCalled();
|
||||
expect(generateTextMock).not.toHaveBeenCalled();
|
||||
expect(aiChatMessageRepo.insert).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* F3 — the LOAD-BEARING run-detach wiring: `effectiveSignal = handle.signal`
|
||||
* after runHooks.begin, then `abortSignal: effectiveSignal` passed to streamText.
|
||||
* That single line is what makes a run survive a browser disconnect (the agent
|
||||
* loop's abort is governed by the RUN's signal, not the socket): a regression to
|
||||
* the socket-bound signal would still pass every other test green while silently
|
||||
* breaking Stop + durability. These two tests pin the exact signal streamText
|
||||
* consumes on both paths.
|
||||
*/
|
||||
describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
|
||||
const streamTextMock = streamText as unknown as jest.Mock;
|
||||
|
||||
// A streamText result stub: the post-call drain + pipe are no-ops here; we only
|
||||
// care WHICH abortSignal streamText was handed.
|
||||
function makeStreamResult() {
|
||||
return {
|
||||
consumeStream: jest.fn(),
|
||||
pipeUIMessageStreamToResponse: jest.fn(),
|
||||
};
|
||||
}
|
||||
|
||||
// A raw-response stub sufficient for the post-streamText wiring
|
||||
// (stripStreamingHopByHopHeaders binds writeHead; startSseHeartbeat registers
|
||||
// close/finish listeners; flushHeaders is belt-and-braces).
|
||||
function makeRes() {
|
||||
return {
|
||||
raw: {
|
||||
writeHead: jest.fn(),
|
||||
write: jest.fn(),
|
||||
once: jest.fn(),
|
||||
on: jest.fn(),
|
||||
flushHeaders: jest.fn(),
|
||||
writableEnded: false,
|
||||
destroyed: false,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
// Wire only the deps reached on the way to streamText: resolve the existing
|
||||
// chat, persist the user + seed the assistant row, load (empty) history, the
|
||||
// admin settings, an empty external toolset + Docmost toolset.
|
||||
function makeService() {
|
||||
const aiChatRepo = {
|
||||
findById: jest.fn(async () => ({ id: 'chat-1', workspaceId: 'ws-1' })),
|
||||
insert: jest.fn(),
|
||||
};
|
||||
const aiChatMessageRepo = {
|
||||
insert: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
findAllByChat: jest.fn(async () => []),
|
||||
update: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
};
|
||||
const aiSettings = { resolve: jest.fn(async () => ({})) };
|
||||
const tools = { forUser: jest.fn(async () => ({})) };
|
||||
const mcpClients = {
|
||||
toolsFor: jest.fn(async () => ({
|
||||
tools: {},
|
||||
clients: [],
|
||||
outcomes: [],
|
||||
instructions: [],
|
||||
})),
|
||||
};
|
||||
const svc = new AiChatService(
|
||||
{} as never, // ai
|
||||
aiChatRepo as never,
|
||||
aiChatMessageRepo as never,
|
||||
{} as never, // aiChatPageSnapshotRepo
|
||||
aiSettings as never,
|
||||
tools as never,
|
||||
mcpClients as never,
|
||||
{} as never, // aiAgentRoleRepo
|
||||
{} as never, // pageRepo (openPage undefined -> never touched)
|
||||
{} as never, // pageAccess
|
||||
{ isAiChatDeferredToolsEnabled: () => false } as never, // environment
|
||||
);
|
||||
return { svc };
|
||||
}
|
||||
|
||||
const body = {
|
||||
chatId: 'chat-1',
|
||||
messages: [
|
||||
{ id: 'm1', role: 'user', parts: [{ type: 'text', text: 'hi' }] },
|
||||
],
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
streamTextMock.mockReset();
|
||||
streamTextMock.mockImplementation(() => makeStreamResult());
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'log')
|
||||
.mockImplementation(() => undefined as never);
|
||||
});
|
||||
|
||||
afterEach(() => jest.restoreAllMocks());
|
||||
|
||||
it('happy path (run-wrapped): streamText is driven with abortSignal === handle.signal (the RUN signal, NOT the socket)', async () => {
|
||||
const { svc } = makeService();
|
||||
const runController = new AbortController();
|
||||
const runSignal = runController.signal;
|
||||
const socketSignal = new AbortController().signal;
|
||||
|
||||
const begin = jest.fn(async () => ({ runId: 'run-1', signal: runSignal }));
|
||||
await svc.stream({
|
||||
user: { id: 'user-1' } as never,
|
||||
workspace: { id: 'ws-1' } as never,
|
||||
sessionId: 'sess-1',
|
||||
body: body as never,
|
||||
res: makeRes() as never,
|
||||
signal: socketSignal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
runHooks: {
|
||||
begin,
|
||||
onAssistantSeeded: jest.fn(),
|
||||
onStep: jest.fn(),
|
||||
onSettled: jest.fn(),
|
||||
} as never,
|
||||
});
|
||||
|
||||
expect(begin).toHaveBeenCalledTimes(1);
|
||||
expect(streamTextMock).toHaveBeenCalledTimes(1);
|
||||
// THE assertion: the agent loop's abort is wired to the RUN, so a browser
|
||||
// disconnect (which aborts only `socketSignal`) cannot end the turn.
|
||||
expect(streamTextMock.mock.calls[0][0].abortSignal).toBe(runSignal);
|
||||
expect(streamTextMock.mock.calls[0][0].abortSignal).not.toBe(socketSignal);
|
||||
});
|
||||
|
||||
it('legacy path (no runHooks): streamText is driven with the SOCKET signal', async () => {
|
||||
const { svc } = makeService();
|
||||
const socketSignal = new AbortController().signal;
|
||||
|
||||
await svc.stream({
|
||||
user: { id: 'user-1' } as never,
|
||||
workspace: { id: 'ws-1' } as never,
|
||||
sessionId: 'sess-1',
|
||||
body: body as never,
|
||||
res: makeRes() as never,
|
||||
signal: socketSignal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
// No runHooks -> the turn stays socket-bound (flag off / default).
|
||||
});
|
||||
|
||||
expect(streamTextMock).toHaveBeenCalledTimes(1);
|
||||
expect(streamTextMock.mock.calls[0][0].abortSignal).toBe(socketSignal);
|
||||
});
|
||||
|
||||
/**
|
||||
* F9 — streamText's TERMINAL callbacks carry the #184 run lifecycle:
|
||||
* onStepFinish -> runHooks.onStep(runId, stepCount)
|
||||
* onFinish -> runHooks.onSettled(runId, 'completed') (dominant path)
|
||||
* onAbort -> runHooks.onSettled(runId, 'aborted')
|
||||
* onError -> runHooks.onSettled(runId, 'error', cause)
|
||||
* makeStreamResult() ignores the streamText options, so these callbacks never
|
||||
* fire on their own — a regression in this wiring (esp. the success path) would
|
||||
* strand the run with NO test catching it. Here we CAPTURE the options streamText
|
||||
* was handed and invoke each callback with the real wiring, asserting the run
|
||||
* hooks fire with the right args.
|
||||
*/
|
||||
// Drive stream() to the point streamText is called, capturing the options object
|
||||
// (which carries onStepFinish/onFinish/onError/onAbort) and the run hooks.
|
||||
async function captureStreamCallbacks() {
|
||||
const { svc } = makeService();
|
||||
let capturedOpts: any;
|
||||
streamTextMock.mockImplementation((opts: any) => {
|
||||
capturedOpts = opts;
|
||||
return makeStreamResult();
|
||||
});
|
||||
const runHooks = {
|
||||
begin: jest.fn(async () => ({
|
||||
runId: 'run-1',
|
||||
signal: new AbortController().signal,
|
||||
})),
|
||||
onAssistantSeeded: jest.fn(),
|
||||
onStep: jest.fn(),
|
||||
onSettled: jest.fn(),
|
||||
};
|
||||
await svc.stream({
|
||||
user: { id: 'user-1' } as never,
|
||||
workspace: { id: 'ws-1' } as never,
|
||||
sessionId: 'sess-1',
|
||||
body: body as never,
|
||||
res: makeRes() as never,
|
||||
signal: new AbortController().signal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
runHooks: runHooks as never,
|
||||
});
|
||||
expect(capturedOpts).toBeDefined();
|
||||
return { capturedOpts, runHooks };
|
||||
}
|
||||
|
||||
it('F9: onStepFinish bumps the run step count, onFinish settles the run "completed" (the dominant autonomous-run path)', async () => {
|
||||
const { capturedOpts, runHooks } = await captureStreamCallbacks();
|
||||
|
||||
// A finished step -> onStep(runId, finishedStepCount).
|
||||
capturedOpts.onStepFinish({ text: 'step one', toolCalls: [], content: [] });
|
||||
expect(runHooks.onStep).toHaveBeenCalledWith('run-1', 1);
|
||||
capturedOpts.onStepFinish({ text: 'step two', toolCalls: [], content: [] });
|
||||
expect(runHooks.onStep).toHaveBeenLastCalledWith('run-1', 2);
|
||||
|
||||
// The success terminal callback settles the run.
|
||||
await capturedOpts.onFinish({
|
||||
text: 'done',
|
||||
finishReason: 'stop',
|
||||
totalUsage: {},
|
||||
usage: {},
|
||||
steps: [],
|
||||
});
|
||||
expect(runHooks.onSettled).toHaveBeenCalledWith('run-1', 'completed');
|
||||
});
|
||||
|
||||
it('F9: onAbort settles the run "aborted"', async () => {
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'warn')
|
||||
.mockImplementation(() => undefined as never);
|
||||
const { capturedOpts, runHooks } = await captureStreamCallbacks();
|
||||
|
||||
await capturedOpts.onAbort({ steps: [] });
|
||||
expect(runHooks.onSettled).toHaveBeenCalledWith('run-1', 'aborted');
|
||||
});
|
||||
|
||||
it('F9: onError settles the run "error" carrying the provider cause', async () => {
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'error')
|
||||
.mockImplementation(() => undefined as never);
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'warn')
|
||||
.mockImplementation(() => undefined as never);
|
||||
const { capturedOpts, runHooks } = await captureStreamCallbacks();
|
||||
|
||||
await capturedOpts.onError({ error: new Error('provider exploded') });
|
||||
expect(runHooks.onSettled).toHaveBeenCalledWith(
|
||||
'run-1',
|
||||
'error',
|
||||
expect.stringContaining('provider exploded'),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* F14 — the begin-failure RESILIENCE branch (the `else` of the run-race guard).
|
||||
*
|
||||
* stream() wraps runHooks.begin in try/catch with TWO branches:
|
||||
* - RunAlreadyActiveError -> 409 ConflictException (pinned above).
|
||||
* - ANY OTHER begin failure -> SWALLOW + continue UNTRACKED on the socket signal
|
||||
* (legacy fallback): it logs "...streaming without run tracking", leaves
|
||||
* `effectiveSignal = signal` (runId undefined) and serves the turn anyway.
|
||||
*
|
||||
* The contract: a transient beginRun failure (e.g. a non-unique DB error inserting
|
||||
* the run row) must STILL serve the user's turn — it must NOT re-throw and must NOT
|
||||
* be misclassified as a 409. A regression that re-threw here would break EVERY turn
|
||||
* on a begin failure with nothing to catch it. This branch is otherwise undriven by
|
||||
* any spec, so it is pinned here SEPARATELY from the 409 path: a plain begin error
|
||||
* proceeds to streamText with the SOCKET signal and still persists the user turn.
|
||||
*/
|
||||
describe('AiChatService.stream — begin-failure resilience / legacy fallback (#184 F14)', () => {
|
||||
const streamTextMock = streamText as unknown as jest.Mock;
|
||||
|
||||
function makeStreamResult() {
|
||||
return {
|
||||
consumeStream: jest.fn(),
|
||||
pipeUIMessageStreamToResponse: jest.fn(),
|
||||
};
|
||||
}
|
||||
|
||||
function makeRes() {
|
||||
return {
|
||||
raw: {
|
||||
writeHead: jest.fn(),
|
||||
write: jest.fn(),
|
||||
once: jest.fn(),
|
||||
on: jest.fn(),
|
||||
flushHeaders: jest.fn(),
|
||||
writableEnded: false,
|
||||
destroyed: false,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
// Same harness as the F3 abortSignal block, but it also exposes
|
||||
// aiChatMessageRepo so we can assert the user turn IS persisted (the turn really
|
||||
// streamed) despite begin() blowing up.
|
||||
function makeService() {
|
||||
const aiChatRepo = {
|
||||
findById: jest.fn(async () => ({ id: 'chat-1', workspaceId: 'ws-1' })),
|
||||
insert: jest.fn(),
|
||||
};
|
||||
const aiChatMessageRepo = {
|
||||
insert: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
findAllByChat: jest.fn(async () => []),
|
||||
update: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
};
|
||||
const aiSettings = { resolve: jest.fn(async () => ({})) };
|
||||
const tools = { forUser: jest.fn(async () => ({})) };
|
||||
const mcpClients = {
|
||||
toolsFor: jest.fn(async () => ({
|
||||
tools: {},
|
||||
clients: [],
|
||||
outcomes: [],
|
||||
instructions: [],
|
||||
})),
|
||||
};
|
||||
const svc = new AiChatService(
|
||||
{} as never, // ai
|
||||
aiChatRepo as never,
|
||||
aiChatMessageRepo as never,
|
||||
{} as never, // aiChatPageSnapshotRepo
|
||||
aiSettings as never,
|
||||
tools as never,
|
||||
mcpClients as never,
|
||||
{} as never, // aiAgentRoleRepo
|
||||
{} as never, // pageRepo
|
||||
{} as never, // pageAccess
|
||||
{ isAiChatDeferredToolsEnabled: () => false } as never, // environment
|
||||
);
|
||||
return { svc, aiChatMessageRepo };
|
||||
}
|
||||
|
||||
const body = {
|
||||
chatId: 'chat-1',
|
||||
messages: [
|
||||
{ id: 'm1', role: 'user', parts: [{ type: 'text', text: 'hi' }] },
|
||||
],
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
streamTextMock.mockReset();
|
||||
streamTextMock.mockImplementation(() => makeStreamResult());
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'log')
|
||||
.mockImplementation(() => undefined as never);
|
||||
});
|
||||
|
||||
afterEach(() => jest.restoreAllMocks());
|
||||
|
||||
it('a PLAIN begin() failure (NOT RunAlreadyActiveError) does NOT 409 — it swallows, logs, and streams the turn UNTRACKED on the socket signal', async () => {
|
||||
const errorSpy = jest
|
||||
.spyOn(Logger.prototype, 'error')
|
||||
.mockImplementation(() => undefined as never);
|
||||
|
||||
const { svc, aiChatMessageRepo } = makeService();
|
||||
const socketSignal = new AbortController().signal;
|
||||
|
||||
// A transient, NON-race begin failure (e.g. a non-unique DB error inserting
|
||||
// the run row). This is the `else` branch of the begin try/catch.
|
||||
const begin = jest.fn(async () => {
|
||||
throw new Error('insert failed');
|
||||
});
|
||||
|
||||
const promise = svc.stream({
|
||||
user: { id: 'user-1' } as never,
|
||||
workspace: { id: 'ws-1' } as never,
|
||||
sessionId: 'sess-1',
|
||||
body: body as never,
|
||||
res: makeRes() as never,
|
||||
signal: socketSignal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
runHooks: {
|
||||
begin,
|
||||
onAssistantSeeded: jest.fn(),
|
||||
onStep: jest.fn(),
|
||||
onSettled: jest.fn(),
|
||||
} as never,
|
||||
});
|
||||
|
||||
// The turn proceeds: NO throw at all (in particular NOT a 409).
|
||||
await expect(promise).resolves.toBeUndefined();
|
||||
|
||||
expect(begin).toHaveBeenCalledTimes(1);
|
||||
|
||||
// The resilience branch logged the legacy-fallback warning.
|
||||
expect(errorSpy).toHaveBeenCalledWith(
|
||||
expect.stringContaining('streaming without run tracking'),
|
||||
expect.anything(),
|
||||
);
|
||||
|
||||
// The turn really streamed: the user message was persisted and streamText ran.
|
||||
expect(aiChatMessageRepo.insert).toHaveBeenCalled();
|
||||
expect(streamTextMock).toHaveBeenCalledTimes(1);
|
||||
|
||||
// The decisive wiring: with no run handle, the fallback uses the SOCKET signal
|
||||
// (effectiveSignal = signal, runId undefined) — not a run-bound signal.
|
||||
expect(streamTextMock.mock.calls[0][0].abortSignal).toBe(socketSignal);
|
||||
});
|
||||
});
|
||||
@@ -453,6 +453,12 @@ describe('chatStreamMetadata', () => {
|
||||
});
|
||||
});
|
||||
|
||||
it('attaches the runId on the start part when a run wraps the turn (#184)', () => {
|
||||
expect(
|
||||
chatStreamMetadata({ type: 'start' }, 'chat-1', undefined, 'run-1'),
|
||||
).toEqual({ chatId: 'chat-1', runId: 'run-1' });
|
||||
});
|
||||
|
||||
it('returns the CUMULATIVE step usage passed in for the finish-step part', () => {
|
||||
// finish-step usage is per-step in v6; the caller accumulates and passes the
|
||||
// running sum, which this just wraps.
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -43,6 +43,30 @@ export class BoundChatDto {
|
||||
pageId: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Reconnect to the latest run of a chat (#184): fetch its persisted lifecycle
|
||||
* state (and the assistant message it projects) for an in-flight or finished run.
|
||||
*/
|
||||
export class GetRunDto {
|
||||
@IsString()
|
||||
chatId: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Explicitly STOP an agent run (#184): the user pressed Stop — distinct from a
|
||||
* browser disconnect, which never stops a run. Either the run id (preferred, from
|
||||
* the streamed start metadata) or the chat id (stop whatever run is active on it).
|
||||
*/
|
||||
export class StopRunDto {
|
||||
@IsOptional()
|
||||
@IsString()
|
||||
runId?: string;
|
||||
|
||||
@IsOptional()
|
||||
@IsString()
|
||||
chatId?: string;
|
||||
}
|
||||
|
||||
/** Export a chat to Markdown (#183). `lang` localizes the few fixed
|
||||
* role/tool-action labels; defaults to English server-side. */
|
||||
export class ExportChatDto {
|
||||
|
||||
@@ -55,6 +55,14 @@ export class UpdateWorkspaceDto extends PartialType(CreateWorkspaceDto) {
|
||||
@IsBoolean()
|
||||
aiDictationStreaming: boolean;
|
||||
|
||||
// #184: detached/autonomous agent runs (settings.ai.autonomousRuns). When on, a
|
||||
// chat turn becomes a server-side RUN that survives a browser disconnect; only
|
||||
// an explicit /ai-chat/stop ends it. Off by default; single-instance-only in
|
||||
// phase 1 (see AiChatRunService.warnIfMultiInstance / AGENTS.md).
|
||||
@IsOptional()
|
||||
@IsBoolean()
|
||||
autonomousRuns: boolean;
|
||||
|
||||
// Workspace master toggle that enables/disables the HTML embed block type.
|
||||
// Persisted at settings.htmlEmbed. ABSENT/false => OFF (default). The block
|
||||
// itself renders in a sandboxed iframe, so this is a feature switch, not a
|
||||
|
||||
@@ -526,6 +526,20 @@ export class WorkspaceService {
|
||||
);
|
||||
}
|
||||
|
||||
if (typeof updateWorkspaceDto.autonomousRuns !== 'undefined') {
|
||||
const prev = settingsBefore?.ai?.autonomousRuns ?? false;
|
||||
if (prev !== updateWorkspaceDto.autonomousRuns) {
|
||||
before.autonomousRuns = prev;
|
||||
after.autonomousRuns = updateWorkspaceDto.autonomousRuns;
|
||||
}
|
||||
await this.workspaceRepo.updateAiSettings(
|
||||
workspaceId,
|
||||
'autonomousRuns',
|
||||
updateWorkspaceDto.autonomousRuns,
|
||||
trx,
|
||||
);
|
||||
}
|
||||
|
||||
if (typeof updateWorkspaceDto.htmlEmbed !== 'undefined') {
|
||||
const prev = settingsBefore?.htmlEmbed ?? false;
|
||||
if (prev !== updateWorkspaceDto.htmlEmbed) {
|
||||
@@ -579,6 +593,7 @@ export class WorkspaceService {
|
||||
delete updateWorkspaceDto.aiChat;
|
||||
delete updateWorkspaceDto.aiDictation;
|
||||
delete updateWorkspaceDto.aiDictationStreaming;
|
||||
delete updateWorkspaceDto.autonomousRuns;
|
||||
delete updateWorkspaceDto.htmlEmbed;
|
||||
delete updateWorkspaceDto.trackerHead;
|
||||
delete updateWorkspaceDto.aiPublicShareAssistant;
|
||||
|
||||
Reference in New Issue
Block a user