Compare commits

...

5 Commits

Author SHA1 Message Date
agent_coder 3d8bc655eb docs(mcp): точный коммент toolAbortSignal — set-and-leave, не restore (#487, ревью nit)
Внутреннее ревью: docstring поля/метода toolAbortSignal утверждал «restores the
prior value on unwind», но wrapInAppToolWithCap намеренно НЕ восстанавливает
сигнал (set-and-leave) — именно это заставляет брошенного проигравшего race
бросить на следующем safe-point; корректность держится на свежем клиенте на ход +
перезаписи следующим вызовом. Комментарии приведены в соответствие с механизмом.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 05:34:46 +03:00
agent_coder b8aac9635b feat(ai-chat): ретраи finalizeAssistant + owner-write приоритет + двусторонний reconcile (#487)
Раньше finalizeAssistant ставил finalized=true ДО записи и не ретраил → один
неудачный UPDATE = строка вечно 'streaming'; свип был boot-only; run-свип
безусловный — асимметрия «run succeeded / message streaming навсегда».

- finalizeAssistant: bounded-ретраи; once-гейт закрывается ТОЛЬКО после успешной
  записи; возвращает ok. Правило owner-write: терминальная запись owner'а
  условна на status='streaming' OR metadata.finalizeFailed (repo.finalizeOwner) —
  перетирает reconcile-штамп, но не проставленный терминал. ВСЕ status-only
  штампы reconcile (stampTerminalIfStreaming, sweepStreaming) пишут строго
  onlyIfStreaming И мёржат metadata.finalizeFailed:true (иначе поздний owner-write
  не перетрёт).
- Порядок: попытка message-finalize → ран финализируется ВСЕГДА; при провале
  message onFinish помечает ран 'error' (не 'completed'). Ран не гейтится на
  message.
- Периодический reconcile-джоб (setInterval, env-tunable) клаузами по порядку:
  (a) пере-драйв зомби; (b) message streaming + ран терминален → штамп по статусу
  рана (succeeded-ран + зависшая строка → 'aborted'+finalizeFailed, НЕ
  'completed'-empty); (c) run running + НЕТ entry И НЕТ zombie + staleness →
  aborted (гейт «нет entry» первичный, staleness от last-progress updatedAt,
  X=max(2×per-call cap,15мин)+boot-warn); (d) message streaming + возраст>X + нет
  активной run-строки → aborted (двойной гейт). isInterruptResume исключает
  finalizeFailed-строки. Оппортунистический одно-чатовый reconcile при старте хода
  (best-effort, не фейлит ход). sweepStreaming boot-only → периодический.

Тесты (реальная БД): owner finalizeOwner чистит finalizeFailed; штамп не перетирает
терминал; поздний owner-write перетирает aborted-штамп; клаузы b/c/d (+живой entry
не трогать, двойной гейт d); «убить БД на finish → после восстановления ни строка,
ни ран не застряли». Юниты: finalizeFailed исключает interrupt-resume;
reconcileStaleRuns «нет entry».

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 05:15:46 +03:00
agent_coder d72e101216 feat(ai-chat): единый гейт конкурентности + серверный supersede (CAS) (#487)
Раньше в legacy-режиме (дефолт!) гейта конкурентности НЕ было вообще — проверка
409 сидела внутри if (autonomousRuns && chatId): два таба = два параллельных
стрима на один чат (интерливинг истории, падения convertToModelMessages).

Серверная часть (клиентская лестница ретраев — отдельной клиентской итерацией
FSM, см. ниже):
- Run-строка ОБЯЗАТЕЛЬНА в ОБОИХ режимах: legacy тоже проходит beginRun/
  finalizeRun (гейт партиального уникального индекса + runId в start-метаданных).
  Различие режимов сведено к семантике abort: legacy onClose при дисконнекте
  зовёт requestStop(runId) вместо аборта сокет-сигнала (которого streamText уже
  не потребляет). Второй таб в ЛЮБОМ режиме → 409 A_RUN_ALREADY_ACTIVE.
- Supersede CAS POST /stream { supersede: { runId: X } } (AiChatRunService
  .supersede): валидация X.chatId===body.chatId (иначе 400 SUPERSEDE_INVALID);
  нет активного рана → degrade в обычный send; активный ≠ X → 409
  SUPERSEDE_TARGET_MISMATCH + текущий runId; активный = X → requestStop →
  awaitSettled (таймаут W=10c) → «записано/сдался» (сдался → settleZombie
  применяет intended условным UPDATE) → ready; таймаут → 409 SUPERSEDE_TIMEOUT,
  ничего не персистится. W обоснован race-on-abort'ом коммита 1; DB-brownout →
  TIMEOUT штатен, W не увеличивать (env-tunable).
- Задокументированы ограничения: нет квиесценции сайд-эффектов (в промпт нового
  рана добавлена строка SUPERSEDE_NOTE «предыдущий ран прерван, его последние
  операции могли примениться с задержкой»); кража слота между освобождением и
  beginRun (→ MISMATCH с новым runId, бэкстоп — уникальный индекс).

Тесты: два параллельных старта (оба режима) → строго 409 второму; supersede при
живом долгом ТУЛЕ (не UPDATE-задержке) settle'ится быстро; все CAS-ветки; дубль
supersede-POST → degrade; зомби-путь через settleZombie; HTTP-маппинг ветвей.
Кейс «beginRun ok → insert user-msg упал → ран settled, слот свободен» покрыт
lifecycle-спекой против реального safety-net catch.

КЛИЕНТ (deferred, точно flagged): удаление лестницы SUPERSEDE_RETRY_DELAYS_MS/
isRunAlreadyActive/supersedeRetryRef и переход на supersede-в-body требуют
адопции runId из start-метаданных (сейчас клиент runId не трекает вовсе) — это
и есть та «итерация клиентского FSM», которую данный серверный контракт
разблокирует. Не трогаю фрагильный клиентский FSM без возможности E2E-валидации
в браузере, чтобы не внести непроверяемую регрессию.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 05:15:46 +03:00
agent_coder 89fcc60782 feat(ai-chat): зомби-механика finalizeRun + settledPromise + условные терминальные UPDATE (#487)
Раньше give-up-путь finalizeRun ВОССТАНАВЛИВАЛ entry (зомби неотличим от живого
рана), а терминальный runRepo.update был БЕЗУСЛОВНЫМ (последний писатель перетирал
терминальный статус).

- Все терминальные UPDATE ранов теперь УСЛОВНЫЕ: новый repo.finalizeIfActive
  пишет только пока строка pending|running (зеркально onlyIfStreaming у сообщений)
  → двойной settle схлопывается в benign no-op, терминальный статус не перетереть.
- При исчерпании ретраев finalizeRun НЕ восстанавливает entry, а оставляет
  ЗОМБИ-запись { terminalWriteFailed, intended:{status,error} } в отдельной map;
  settleZombie пере-драйвит intended условным UPDATE (зовётся reconcile/supersede/
  boot sweep). Зомби удаляется после успешного UPDATE или обнаружения строки уже
  терминальной.
- Per-run settledPromise в ОТДЕЛЬНОЙ map runId→deferred (переживает active.delete),
  создаётся в beginRun, резолвится ровно один раз исходом (записано/сдался); поздний
  подписчик через ранее взятую ссылку получает резолвленный; peekSettled: live
  deferred → зомби-синтез → undefined (читать строку). Обе map bounded.
- Задокументирована потеря (single-process): рестарт до пере-драйва → boot sweep
  пишет 'aborted' поверх фактического intended — неустранимо в phase 1.

Тесты: юниты give-up-пути (зомби вместо restore + notifier terminalWriteFailed +
settleZombie), двойного settle, позднего подписчика; интеграционные против реальной
Б, что условный finalizeIfActive не перетирает терминал и двойной settle схлопывается.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 05:15:46 +03:00
agent_coder 917c406489 fix(ai-chat): in-app тулы — race-on-abort + safe-points + per-call cap (#487)
In-app тул-обёртки отбрасывали второй аргумент options с abortSignal — это был
единственный класс тулов без отмены и без wall-clock cap. Контент-мутации идут
через collab-WS (mutatePageContent), не через axios, поэтому «прокинуть signal в
axios» главный write-путь не покрывал.

- Переиспользован race-паттерн wrapToolWithCallTimeout: каждый in-app тул гонится
  против композитного сигнала (Stop + per-call cap); на abort — реджект
  немедленно, проигравший промис отбрасывается (латентность мс, не сетевой
  teardown — от этого зависит таймаут supersede в коммите 3).
- Safe-point проверки сигнала между последовательными вызовами paginateAll и
  пре-коммитная проверка в mutatePageContent (+ реентрантный близнец) через
  DocmostClient.setToolAbortSignal, который обёртка публикует перед каждым вызовом.
- Per-call cap покрывает весь вызов, env-tunable (AI_CHAT_INAPP_TOOL_CALL_CAP_MS,
  дефолт 2 мин).
- Задокументировано ограничение (#487): abort между вызовами многовызовного
  write-тула оставляет частично применённую операцию — отмена гарантирует «новый
  вызов не стартует», не «запись не доехала».

Тесты (честное свойство «после Stop не стартует новый HTTP/WS-вызов»): юнит
wrapInAppToolWithCap (реджект-на-abort, отсутствие новых вызовов, per-call cap,
публикация сигнала) + mock-HTTP тест реального paginateAll safe-point.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 05:15:46 +03:00
19 changed files with 2519 additions and 238 deletions
@@ -43,6 +43,9 @@ function makeRepo(overrides: Record<string, jest.Mock> = {}) {
workspaceId: v.workspaceId,
})),
update: jest.fn(async () => ({ id: 'run-1' })),
// #487: terminal finalize now goes through the CONDITIONAL write. Default
// returns a truthy row (the run WAS active -> this call wrote it).
finalizeIfActive: jest.fn(async () => ({ id: 'run-1', status: 'succeeded' })),
markStopRequested: jest.fn(async () => ({ id: 'run-1' })),
findActiveByChat: jest.fn(async () => undefined),
findLatestByChat: jest.fn(async () => undefined),
@@ -336,14 +339,12 @@ describe('AiChatRunService run lifecycle', () => {
await svc.finalizeRun('run-1', 'ws-1', 'error', 'provider blew up');
expect(svc.isLocallyActive('run-1')).toBe(false);
expect(repo.update).toHaveBeenCalledWith(
// #487: the terminal write is CONDITIONAL (finalizeIfActive); finishedAt is
// stamped inside the repo method, so the service passes just status + error.
expect(repo.finalizeIfActive).toHaveBeenCalledWith(
'run-1',
'ws-1',
expect.objectContaining({
status: 'failed',
error: 'provider blew up',
finishedAt: expect.any(Date),
}),
expect.objectContaining({ status: 'failed', error: 'provider blew up' }),
);
});
@@ -366,8 +367,8 @@ describe('AiChatRunService run lifecycle', () => {
// 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(
expect(repo.finalizeIfActive).toHaveBeenCalledTimes(1);
expect(repo.finalizeIfActive).toHaveBeenCalledWith(
'run-1',
'ws-1',
expect.objectContaining({ status: 'failed', error: 'first' }),
@@ -389,8 +390,8 @@ describe('AiChatRunService run lifecycle', () => {
const updateGate = new Promise((res) => {
resolveUpdate = res;
});
const update = jest.fn(() => updateGate);
const repo = makeRepo({ update });
const finalizeIfActive = jest.fn(() => updateGate);
const repo = makeRepo({ finalizeIfActive });
const svc = new AiChatRunService(repo as never, makeEnv() as never);
await svc.beginRun({
chatId: 'chat-1',
@@ -399,23 +400,23 @@ describe('AiChatRunService run lifecycle', () => {
});
// 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.
// claims the entry (active.delete) and awaits the write; the second, started
// in the same macrotask, finds the entry already gone and returns at the claim
// WITHOUT ever writing.
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);
expect(finalizeIfActive).toHaveBeenCalledTimes(1);
// Let the single in-flight update land; both calls resolve cleanly.
resolveUpdate({ id: 'run-1' });
resolveUpdate({ id: 'run-1', status: 'succeeded' });
await Promise.all([p1, p2]);
expect(update).toHaveBeenCalledTimes(1);
expect(finalizeIfActive).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(
expect(finalizeIfActive).toHaveBeenCalledWith(
'run-1',
'ws-1',
expect.objectContaining({ status: 'succeeded' }),
@@ -431,10 +432,10 @@ describe('AiChatRunService run lifecycle', () => {
// 409s until a restart. The fix updates FIRST and retries.
let calls = 0;
const repo = makeRepo({
update: jest.fn(async () => {
finalizeIfActive: jest.fn(async () => {
calls += 1;
if (calls === 1) throw new Error('deadlock detected');
return { id: 'run-1' };
return { id: 'run-1', status: 'succeeded' };
}),
});
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined);
@@ -447,26 +448,29 @@ describe('AiChatRunService run lifecycle', () => {
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'.
// The retry landed the terminal write: the entry is dropped (slot freed), no
// zombie left, and the row carries the real terminal status.
expect(svc.isLocallyActive('run-1')).toBe(false);
expect(repo.update).toHaveBeenCalledTimes(2);
expect(repo.update).toHaveBeenLastCalledWith(
expect(svc.hasZombie('run-1')).toBe(false);
expect(repo.finalizeIfActive).toHaveBeenCalledTimes(2);
expect(repo.finalizeIfActive).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 () => {
it('#487 give-up: if the terminal write keeps failing, finalizeRun leaves a ZOMBIE (does NOT restore the entry) and settleZombie re-drives it', 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.
// #487 changes the give-up behaviour: the entry is NOT restored (a restored
// entry is indistinguishable from a live run). Instead a ZOMBIE record holds
// the intended terminal status, and a re-drive (settleZombie — called by the
// reconcile / supersede / opportunistic paths) applies it later.
let healthy = false;
const repo = makeRepo({
update: jest.fn(async () => {
finalizeIfActive: jest.fn(async () => {
if (!healthy) throw new Error('pool exhausted');
return { id: 'run-1' };
return { id: 'run-1', status: 'succeeded' };
}),
});
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined);
@@ -480,35 +484,83 @@ describe('AiChatRunService run lifecycle', () => {
userId: 'user-1',
});
// First settle: every bounded attempt fails -> entry retained, NOT settled.
// First settle: every bounded attempt fails -> ZOMBIE, entry NOT restored.
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.
expect(svc.isLocallyActive('run-1')).toBe(false); // NOT a live entry
expect(svc.hasZombie('run-1')).toBe(true);
expect(svc.zombieRunIds()).toContain('run-1');
// The give-up emits ONE explicit, greppable ERROR mentioning the zombie.
const gaveUp = errorSpy.mock.calls.some(
(c) =>
/NON-TERMINAL/.test(String(c[0])) &&
/ZOMBIE/.test(String(c[0])) &&
/run-1/.test(String(c[0])) &&
/chat-1/.test(String(c[0])),
);
expect(gaveUp).toBe(true);
// The settle notifier resolved as terminalWriteFailed (a subscriber learns the
// slot still needs the intended status applied).
const outcome = await svc.peekSettled('run-1');
expect(outcome).toEqual({
status: 'succeeded',
error: null,
terminalWriteFailed: true,
});
// The DB recovers; a later settle now succeeds and frees the slot.
// The DB recovers; a re-drive settles the zombie via the conditional UPDATE.
healthy = true;
await svc.finalizeRun('run-1', 'ws-1', 'completed');
expect(svc.isLocallyActive('run-1')).toBe(false);
expect(repo.update).toHaveBeenLastCalledWith(
const redriven = await svc.settleZombie('run-1');
expect(redriven).toBe(true);
expect(svc.hasZombie('run-1')).toBe(false);
expect(repo.finalizeIfActive).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;
// A later finalizeRun is idempotent (row already terminal): it no-ops at the
// once-gate, never re-writing.
const callsBefore = repo.finalizeIfActive.mock.calls.length;
await svc.finalizeRun('run-1', 'ws-1', 'error', 'late');
expect(repo.update).toHaveBeenCalledTimes(callsBefore);
expect(repo.finalizeIfActive).toHaveBeenCalledTimes(callsBefore);
});
it('#487 double-settle collapses to a benign no-op (conditional write; notifier resolves once)', async () => {
// A second concurrent settle is stopped at the synchronous active.delete
// claim, so the terminal write runs exactly once and the notifier resolves
// exactly once with the FIRST settler's outcome.
const repo = makeRepo();
const svc = new AiChatRunService(repo as never, makeEnv() as never);
await svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'u1' });
await svc.finalizeRun('run-1', 'ws-1', 'aborted');
await svc.finalizeRun('run-1', 'ws-1', 'error', 'late'); // no-op
expect(repo.finalizeIfActive).toHaveBeenCalledTimes(1);
const outcome = await svc.peekSettled('run-1');
// peekSettled after resolve+delete falls through (notifier dropped, no zombie)
// -> undefined; the FIRST settler already resolved any earlier subscriber.
expect(outcome).toBeUndefined();
});
it('#487 late settledPromise subscriber gets the resolved outcome', async () => {
const repo = makeRepo();
const svc = new AiChatRunService(repo as never, makeEnv() as never);
await svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'u1' });
// Subscribe BEFORE settle: hold the promise reference (as supersede does).
const early = svc.peekSettled('run-1');
expect(early).toBeDefined();
await svc.finalizeRun('run-1', 'ws-1', 'completed');
// The reference grabbed before settle resolves with the written outcome, even
// though the notifier was dropped from the map on resolve (bounded).
await expect(early).resolves.toEqual({
status: 'succeeded',
error: null,
terminalWriteFailed: false,
});
});
it('recordStep / linkAssistantMessage are best-effort: a repo failure is swallowed', async () => {
@@ -525,3 +577,197 @@ describe('AiChatRunService run lifecycle', () => {
).resolves.toBeUndefined();
});
});
describe('#487 AiChatRunService.supersede (CAS)', () => {
const chat = 'chat-1';
const ws = 'ws-1';
it('degrade: no active run on the chat -> caller sends a normal turn', async () => {
const repo = makeRepo({
findById: jest.fn(async () => undefined),
findActiveByChat: jest.fn(async () => undefined),
});
const svc = new AiChatRunService(repo as never, makeEnv() as never);
expect(await svc.supersede(chat, 'run-x', ws)).toEqual({ kind: 'degrade' });
});
it('invalid: the target run belongs to a DIFFERENT chat -> 400', async () => {
const repo = makeRepo({
findById: jest.fn(async () => ({
id: 'run-x',
chatId: 'other-chat',
workspaceId: ws,
})),
});
const svc = new AiChatRunService(repo as never, makeEnv() as never);
expect(await svc.supersede(chat, 'run-x', ws)).toEqual({ kind: 'invalid' });
});
it('mismatch: a DIFFERENT run is active than the one targeted -> current runId', async () => {
const repo = makeRepo({
findById: jest.fn(async () => ({ id: 'run-x', chatId: chat, workspaceId: ws })),
findActiveByChat: jest.fn(async () => ({
id: 'run-live',
chatId: chat,
workspaceId: ws,
status: 'running',
})),
});
const svc = new AiChatRunService(repo as never, makeEnv() as never);
expect(await svc.supersede(chat, 'run-x', ws)).toEqual({
kind: 'mismatch',
activeRunId: 'run-live',
});
});
it('ready: the target IS active -> stop it, await its (fast) settle, free the slot', async () => {
// Simulate a live long TOOL (NOT a slow UPDATE): the run stays active until an
// explicit Stop unwinds it; commit-1's race makes that settle land quickly.
// The abort listener stands in for streamText's onAbort -> finalizeRun.
const repo = makeRepo({
findById: jest.fn(async () => ({
id: 'run-1',
chatId: chat,
workspaceId: ws,
status: 'aborted',
error: null,
})),
findActiveByChat: jest.fn(async () => ({
id: 'run-1',
chatId: chat,
workspaceId: ws,
status: 'running',
})),
});
const svc = new AiChatRunService(repo as never, makeEnv() as never);
const handle = await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' });
handle.signal.addEventListener('abort', () => {
void svc.finalizeRun('run-1', ws, 'aborted');
});
// supersede: getRun -> getActiveByChat(==target) -> requestStop -> the abort
// listener settles the run -> awaitSettled resolves -> ready.
expect(await svc.supersede(chat, 'run-1', ws, 10_000)).toEqual({
kind: 'ready',
});
expect(handle.signal.aborted).toBe(true); // Stop reached the run
});
it('timeout: the target never settles within W -> 409 SUPERSEDE_TIMEOUT (nothing persisted)', async () => {
const repo = makeRepo({
findById: jest.fn(async () => ({ id: 'run-1', chatId: chat, workspaceId: ws })),
findActiveByChat: jest.fn(async () => ({
id: 'run-1',
chatId: chat,
workspaceId: ws,
status: 'running',
})),
});
const svc = new AiChatRunService(repo as never, makeEnv() as never);
await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' });
// Do NOT settle the run: a tiny W elapses -> timeout.
const result = await svc.supersede(chat, 'run-1', ws, 30);
expect(result).toEqual({ kind: 'timeout' });
});
it('ready then a DUPLICATE supersede POST degrades (the run is already gone)', async () => {
let active: unknown = {
id: 'run-1',
chatId: chat,
workspaceId: ws,
status: 'running',
};
const repo = makeRepo({
findById: jest.fn(async () => ({
id: 'run-1',
chatId: chat,
workspaceId: ws,
status: 'aborted',
error: null,
})),
findActiveByChat: jest.fn(async () => active),
finalizeIfActive: jest.fn(async () => {
active = undefined; // settling frees the active slot
return { id: 'run-1', status: 'aborted' };
}),
});
const svc = new AiChatRunService(repo as never, makeEnv() as never);
const handle = await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' });
handle.signal.addEventListener('abort', () => {
void svc.finalizeRun('run-1', ws, 'aborted');
});
expect(await svc.supersede(chat, 'run-1', ws, 10_000)).toEqual({
kind: 'ready',
});
// The duplicate POST for the same target now finds no active run -> degrade.
expect(await svc.supersede(chat, 'run-1', ws)).toEqual({ kind: 'degrade' });
});
it('reconcileStaleRuns: aborts a stale run with NO entry/zombie; NEVER touches a live entry', async () => {
const finalizeIfActive = jest.fn(async () => ({ id: 'x', status: 'aborted' }));
const repo = makeRepo({
insert: jest.fn(async (v: any) => ({
id: 'live-1',
status: 'running',
chatId: v.chatId,
workspaceId: v.workspaceId,
})),
finalizeIfActive,
findStaleActive: jest.fn(async () => [
{ id: 'orphan-1', workspaceId: ws, chatId: 'c-orphan' },
{ id: 'live-1', workspaceId: ws, chatId: 'c-live' },
]),
});
const svc = new AiChatRunService(repo as never, makeEnv() as never);
// A LIVE run this replica owns (in the `active` map).
await svc.beginRun({ chatId: 'c-live', workspaceId: ws, userId: 'u1' });
expect(svc.isLocallyActive('live-1')).toBe(true);
const aborted = await svc.reconcileStaleRuns(15 * 60 * 1000);
expect(aborted).toBe(1);
// The orphan (no entry) was aborted; the live entry was NEVER passed to the DB.
expect(finalizeIfActive).toHaveBeenCalledTimes(1);
expect(finalizeIfActive).toHaveBeenCalledWith(
'orphan-1',
ws,
expect.objectContaining({ status: 'aborted' }),
);
expect(svc.isLocallyActive('live-1')).toBe(true);
});
it('gave-up zombie: supersede applies the intended status (settleZombie) then is ready', async () => {
let healthy = false;
let active: unknown = {
id: 'run-1',
chatId: chat,
workspaceId: ws,
status: 'running',
};
const repo = makeRepo({
findById: jest.fn(async () => ({ id: 'run-1', chatId: chat, workspaceId: ws })),
findActiveByChat: jest.fn(async () => active),
finalizeIfActive: jest.fn(async () => {
if (!healthy) throw new Error('db down');
active = undefined;
return { id: 'run-1', status: 'aborted' };
}),
});
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined);
jest.spyOn(Logger.prototype, 'error').mockImplementation(() => undefined);
const svc = new AiChatRunService(repo as never, makeEnv() as never);
await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' });
// The run's terminal write gives up -> zombie (row still 'running').
await svc.finalizeRun('run-1', ws, 'aborted');
expect(svc.hasZombie('run-1')).toBe(true);
// The DB recovers; supersede awaits the (already-resolved, terminalWriteFailed)
// settle, then settleZombie applies the intended status -> ready.
healthy = true;
expect(await svc.supersede(chat, 'run-1', ws, 10_000)).toEqual({
kind: 'ready',
});
expect(svc.hasZombie('run-1')).toBe(false);
});
});
@@ -34,6 +34,88 @@ export class RunAlreadyActiveError extends Error {
export type TurnTerminalStatus = 'completed' | 'error' | 'aborted';
export type RunTerminalStatus = 'succeeded' | 'failed' | 'aborted';
/** The terminal run statuses — the row is done once it reads one of these. */
export const RUN_TERMINAL_STATUSES: readonly RunTerminalStatus[] = [
'succeeded',
'failed',
'aborted',
];
/** Whether a persisted run status is terminal (settled). */
export function isRunTerminal(status: string | null | undefined): boolean {
return (
status === 'succeeded' || status === 'failed' || status === 'aborted'
);
}
/**
* #487: the outcome a run's {@link AiChatRunService.finalizeRun} settled with.
* `terminalWriteFailed` = the terminal write GAVE UP after the bounded retry, so
* the row is still non-terminal ('running') and a ZOMBIE record holds the
* `intended` status for a later re-drive (reconcile / supersede / boot sweep). A
* subscriber (supersede, #487 commit 3) uses this to decide whether the slot is
* genuinely free or must first have the intended status applied.
*/
export interface RunSettleOutcome {
status: RunTerminalStatus;
error: string | null;
terminalWriteFailed: boolean;
}
/**
* #487: how long a supersede waits for the target run to settle after Stop before
* it degrades to `SUPERSEDE_TIMEOUT`. W=10s is generous under a HEALTHY DB: commit
* 1's race-on-abort makes an in-app tool abort->settle in ms/hundreds of ms, so a
* live run releases its slot well within the window. Under a DB brownout the
* timeout is normal (the write cannot land); W must NOT be raised to paper
* over a slow DB — a SUPERSEDE_TIMEOUT is the honest signal (nothing persisted,
* the composer keeps the user's text). Env-tunable for ops, default 10s.
*/
export const SUPERSEDE_SETTLE_TIMEOUT_MS = (() => {
const raw = Number(process.env.AI_CHAT_SUPERSEDE_TIMEOUT_MS);
return Number.isFinite(raw) && raw > 0 ? raw : 10_000;
})();
/**
* #487: the result of the supersede CAS ({@link AiChatRunService.supersede}).
* - `degrade` : no active run on the chat (it ended between click and POST) —
* the caller sends a NORMAL turn (NOT a mismatch);
* - `invalid` : the target runId belongs to a DIFFERENT chat (malformed CAS 400);
* - `mismatch` : a DIFFERENT run is active than the one the client targeted —
* 409 SUPERSEDE_TARGET_MISMATCH carrying the current `activeRunId`
* (the client does NOT auto-retry);
* - `timeout` : the target did not settle within W — 409 SUPERSEDE_TIMEOUT,
* nothing persisted;
* - `ready` : the target was stopped AND settled (or its zombie's intended was
* applied) — the slot is free; the caller may beginRun the new run.
*/
export type SupersedeResult =
| { kind: 'degrade' }
| { kind: 'invalid' }
| { kind: 'mismatch'; activeRunId: string }
| { kind: 'timeout' }
| { kind: 'ready' };
/** A one-shot settle notifier (#487): `resolve` is called EXACTLY ONCE. */
interface Deferred<T> {
promise: Promise<T>;
resolve: (value: T) => void;
}
/**
* #487: a run whose terminal write GAVE UP (every bounded attempt failed). The
* row is stranded non-terminal ('running'); this record is the ONLY thing that
* distinguishes it from a live run, and carries the `intended` terminal status so
* a re-drive can apply it via the conditional UPDATE. Process-local (phase-1
* single-process assumption): a restart drops it, and the boot sweep then writes
* 'aborted' over the intended — a documented loss (see finalizeRun).
*/
interface ZombieRun {
workspaceId: string;
chatId: string;
intended: { status: RunTerminalStatus; error: string | null };
}
export function mapTurnStatusToRun(
status: TurnTerminalStatus,
): RunTerminalStatus {
@@ -101,6 +183,22 @@ export class AiChatRunService implements OnModuleInit {
// uptime — negligible in phase 1's single process.
private readonly settled = new Set<string>();
// #487 runId -> one-shot settle notifier. Kept in a SEPARATE map from `active`
// ON PURPOSE: it must OUTLIVE the `active.delete` claim inside finalizeRun (the
// claim frees the slot the instant finalize starts), so a subscriber can still
// await the outcome after the entry is gone. Created in beginRun, resolved
// EXACTLY ONCE in finalizeRun, then removed (bounded). Absence => this replica
// has no live notifier: a subscriber falls back to the zombie map, then to the
// row (see peekSettled). Process-local (phase-1 single-process assumption).
private readonly settledPromises = new Map<string, Deferred<RunSettleOutcome>>();
// #487 runId -> ZOMBIE record: a run whose terminal write gave up (row stranded
// non-terminal). BOUNDED — an entry is added only on give-up and removed on a
// successful re-drive (settleZombie) or when the row is found already terminal;
// a process restart clears it (and the boot sweep settles the stranded row).
// Process-local (phase-1 single-process assumption).
private readonly zombies = new Map<string, ZombieRun>();
// 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
@@ -224,6 +322,10 @@ export class AiChatRunService implements OnModuleInit {
chatId: args.chatId,
workspaceId: args.workspaceId,
});
// #487: arm the one-shot settle notifier BEFORE returning, so a subscriber
// that races in immediately after begin always finds a promise to await. It
// is resolved exactly once when the run settles (or gives up).
this.settledPromises.set(run.id, this.makeDeferred<RunSettleOutcome>());
return { runId: run.id, signal: controller.signal };
}
@@ -263,47 +365,43 @@ export class AiChatRunService implements OnModuleInit {
}
/**
* 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.
* Finalize a run to its terminal status (succeeded / failed / aborted) via a
* CONDITIONAL UPDATE, stamping finishedAt + any error. Atomically safe against a
* concurrent settle AND robust against a transient terminal-write failure.
*
* 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.
* claim 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.
*
* 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.
* ALL TERMINAL WRITES ARE CONDITIONAL (#487): `finalizeIfActive` only flips a
* row still in pending|running (mirror of the assistant message's
* `onlyIfStreaming`). So even a settle that DID reach the UPDATE (e.g. a
* reconcile stamp racing an owner finalize) can never clobber a terminal status
* — the loser matches nothing and is a benign no-op. `active.delete` is the
* fast, in-process gate; the conditional WHERE is the authoritative one.
*
* 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.
* ZOMBIE ON GIVE-UP (#487): if every bounded attempt THROWS (the DB is down for
* the whole finalize), we do NOT restore the entry. The row is stranded
* non-terminal ('running'); we record a ZOMBIE `{ terminalWriteFailed, intended
* }` (the ONLY thing distinguishing this dead run from a live one) and resolve
* the settle notifier with `terminalWriteFailed: true`. A restore would make the
* zombie indistinguishable from a live run to every reader; instead a re-drive
* (settleZombie, called by the periodic reconcile / supersede / opportunistic
* paths) applies the intended status later via the same conditional UPDATE.
*
* DOCUMENTED LOSS (#487, single-process phase 1): if the process RESTARTS before
* a zombie is re-driven, the in-memory zombie map is gone and the boot sweep
* (unconditional) writes 'aborted' over the ACTUAL intended status. This is
* unavoidable while the run lifecycle is single-process — there is no durable
* record of `intended`; a cross-process durable intent is deferred to phase 2.
*
* IDEMPOTENT: the settle notifier resolves EXACTLY ONCE; a second settle is
* stopped at `settled.has` or the `active.delete` claim, so a double-settle
* collapses to a single write and can never double-resolve or clobber the row.
*/
async finalizeRun(
runId: string,
@@ -314,13 +412,17 @@ export class AiChatRunService implements OnModuleInit {
// ---- 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.
// Capture the entry BEFORE the delete for the give-up log context.
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;
const status = mapTurnStatusToRun(turnStatus);
const err = error ?? null;
const chatId = entry?.chatId ?? 'unknown';
let lastError: unknown;
for (
let attempt = 1;
@@ -328,47 +430,294 @@ export class AiChatRunService implements OnModuleInit {
attempt++
) {
try {
await this.runRepo.update(runId, workspaceId, {
status: mapTurnStatusToRun(turnStatus),
finishedAt: new Date(),
error: error ?? null,
const row = await this.runRepo.finalizeIfActive(runId, workspaceId, {
status,
error: err,
});
// Terminal write landed: arm the once-gate. The entry is already gone
// (claimed above); we do NOT restore it. The slot is now free.
// No throw => the row is now terminal (we wrote it, or it was ALREADY
// terminal — another writer won the conditional UPDATE, a benign no-op).
this.settled.add(runId);
this.zombies.delete(runId);
// Resolve with the persisted outcome: our status when WE wrote it, else
// the row's real terminal status (re-read on the already-terminal path so
// a subscriber never sees a status we did not actually persist).
const outcome: RunSettleOutcome = row
? { status, error: err, terminalWriteFailed: false }
: await this.readTerminalOutcome(runId, workspaceId, status, err);
this.resolveSettled(runId, outcome);
return;
} catch (err) {
lastError = err;
} catch (err2) {
lastError = err2;
this.logger.warn(
`Failed to finalize run ${runId} (attempt ${attempt}/${
AiChatRunService.FINALIZE_MAX_ATTEMPTS
}): ${err instanceof Error ? err.message : 'unknown error'}`,
}): ${err2 instanceof Error ? err2.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").
// Every attempt threw: GIVE UP. The row is stranded non-terminal ('running').
// Do NOT restore the entry (a restored entry is indistinguishable from a live
// run); leave a ZOMBIE record instead, and resolve the notifier as
// terminalWriteFailed so a subscriber knows the slot still needs the intended
// status applied. One explicit, greppable ERROR so an operator can tell a
// give-up from a per-attempt blip.
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`,
`Run ${runId} (chat ${chatId}) left NON-TERMINAL ('running'): terminal ` +
`write failed after ${AiChatRunService.FINALIZE_MAX_ATTEMPTS} attempts; ` +
`ZOMBIE recorded (intended '${status}'), recovery deferred to reconcile / ` +
`supersede / 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);
this.zombies.set(runId, {
workspaceId,
chatId,
intended: { status, error: err },
});
this.resolveSettled(runId, { status, error: err, terminalWriteFailed: true });
}
/**
* #487: re-drive a zombie run's intended terminal write (the conditional
* UPDATE). Called by the periodic reconcile (commit 4), an opportunistic
* single-chat reconcile, and supersede (commit 3). On success — the row is now
* terminal (written OR found already terminal) — the zombie is cleared and the
* once-gate armed; on another failure the zombie is kept for a later retry.
* Returns true when the row is now terminal. Best-effort; never throws.
*/
async settleZombie(runId: string): Promise<boolean> {
const z = this.zombies.get(runId);
if (!z) return false;
try {
await this.runRepo.finalizeIfActive(runId, z.workspaceId, {
status: z.intended.status,
error: z.intended.error,
});
this.zombies.delete(runId);
this.settled.add(runId);
return true;
} catch (err) {
this.logger.warn(
`Re-drive of zombie run ${runId} (chat ${z.chatId}) failed; will retry ` +
`later: ${err instanceof Error ? err.message : 'unknown error'}`,
);
return false;
}
}
/**
* #487 reconcile clause (c): abort runs the DB still shows active (pending|
* running) but that this replica does NOT own — NO live entry AND NO zombie —
* and that have been UNTOUCHED past `staleMs` (from last-progress `updated_at`,
* NOT startedAt, so a legit long marathon is never a candidate). "No entry" is
* the PRIMARY gate: a live entry (an actively-executing run on this replica) is
* NEVER aborted, whatever its age. Returns the number aborted. Best-effort —
* never throws (a periodic-job failure must not crash the process).
*/
async reconcileStaleRuns(staleMs: number): Promise<number> {
let candidates: Array<{ id: string; workspaceId: string; chatId: string }>;
try {
candidates = await this.runRepo.findStaleActive(staleMs);
} catch (err) {
this.logger.warn(
`Reconcile (stale runs) query failed: ${
err instanceof Error ? err.message : 'unknown error'
}`,
);
return 0;
}
let aborted = 0;
for (const c of candidates) {
// PRIMARY gate: never touch a live entry, and never race a zombie we are
// already re-driving (settleZombie owns those).
if (this.active.has(c.id) || this.zombies.has(c.id)) continue;
try {
const row = await this.runRepo.finalizeIfActive(c.id, c.workspaceId, {
status: 'aborted',
error: 'Run aborted by reconcile: no live runner (stale).',
});
if (row) {
aborted += 1;
this.settled.add(c.id);
}
} catch (err) {
this.logger.warn(
`Reconcile abort of stale run ${c.id} failed: ${
err instanceof Error ? err.message : 'unknown error'
}`,
);
}
}
return aborted;
}
/**
* #487: the run's settle outcome as seen by THIS replica, or undefined when it
* has no record (the caller then reads the row — the DB is the source of truth).
* A LIVE deferred (still settling, or resolved-but-not-yet-consumed) wins; a
* ZOMBIE synthesizes the give-up outcome. A subscriber (supersede) races this
* against a timeout.
*/
peekSettled(runId: string): Promise<RunSettleOutcome> | undefined {
const d = this.settledPromises.get(runId);
if (d) return d.promise;
const z = this.zombies.get(runId);
if (z) {
return Promise.resolve({
status: z.intended.status,
error: z.intended.error,
terminalWriteFailed: true,
});
}
return undefined;
}
/**
* #487: await a run's settle outcome, bounded by `timeoutMs`. Returns the
* outcome on settle, or undefined on TIMEOUT (or when this replica has no record
* of the run and its row is not terminal). Uses the LIVE settle notifier / the
* zombie synth when present; else reads the row (the DB is the source of truth
* once the in-memory record is gone). The subscriber (supersede) grabs this
* right after Stop; commit 1's race makes the settle land in ms on a healthy DB.
*/
async awaitSettled(
runId: string,
workspaceId: string,
timeoutMs: number,
): Promise<RunSettleOutcome | undefined> {
const pending = this.peekSettled(runId);
if (pending) {
let timer: ReturnType<typeof setTimeout> | undefined;
const timeout = new Promise<undefined>((resolve) => {
timer = setTimeout(() => resolve(undefined), timeoutMs);
timer.unref?.();
});
try {
return await Promise.race([pending, timeout]);
} finally {
if (timer) clearTimeout(timer);
}
}
// No live notifier and no zombie: read the row (already settled-and-written,
// or unknown here). A terminal row is an outcome; anything else -> undefined.
const row = await this.runRepo.findById(runId, workspaceId);
if (row && isRunTerminal(row.status)) {
return {
status: row.status as RunTerminalStatus,
error: row.error ?? null,
terminalWriteFailed: false,
};
}
return undefined;
}
/**
* #487: the SERVER supersede CAS for `POST /stream { supersede: { runId: X } }`.
* Atomically transitions "X is the chat's active run" -> "X is stopped, settled,
* slot free" so the caller can start a replacement run. See {@link
* SupersedeResult} for the branch semantics.
*
* On a `ready` result the caller MUST still go through the normal beginRun gate
* (the partial unique index) — between the slot freeing here and beginRun a
* neighbouring tab's ordinary POST can win the slot (documented SLOT-THEFT: the
* loser then gets a MISMATCH carrying the NEW runId). There is also NO side-
* effect quiescence: an in-flight write of the stopped run may still land AFTER
* the new run starts (commit 1 stops the NEXT call, not one already committing),
* so the caller adds a prompt note to the new run.
*/
async supersede(
chatId: string,
targetRunId: string,
workspaceId: string,
timeoutMs: number = SUPERSEDE_SETTLE_TIMEOUT_MS,
): Promise<SupersedeResult> {
// Validate the target belongs to THIS chat (a CAS targeting another chat's run
// is malformed -> 400). A missing row is NOT invalid: the run may have ended
// and been pruned; the active-run check below decides degrade vs mismatch.
const target = await this.getRun(targetRunId, workspaceId);
if (target && target.chatId !== chatId) return { kind: 'invalid' };
const active = await this.getActiveForChat(chatId, workspaceId);
// No active run: it ended between the client's click and this POST — this is a
// DEGRADE to a normal send, NOT a mismatch (the user's intent still holds).
if (!active) return { kind: 'degrade' };
// A DIFFERENT run is active than the one the client saw -> mismatch. The
// client does not auto-retry; it surfaces the new runId.
if (active.id !== targetRunId) {
return { kind: 'mismatch', activeRunId: active.id };
}
// The target IS active: stop it, then await its settle within W.
await this.requestStop(targetRunId, workspaceId);
const outcome = await this.awaitSettled(targetRunId, workspaceId, timeoutMs);
if (!outcome) return { kind: 'timeout' };
// Gave up (terminal write failed): apply the intended status via the
// conditional UPDATE so the slot actually frees. If that ALSO fails, the row
// is still stranded -> treat as a timeout (nothing persisted for the new run).
if (outcome.terminalWriteFailed) {
const settled = await this.settleZombie(targetRunId);
if (!settled) return { kind: 'timeout' };
}
return { kind: 'ready' };
}
/** #487 test/diagnostic seam: whether a give-up zombie is held for this run. */
hasZombie(runId: string): boolean {
return this.zombies.has(runId);
}
/** #487: every zombie runId held on this replica (reconcile clause a, commit 4). */
zombieRunIds(): string[] {
return [...this.zombies.keys()];
}
/** #487: create a one-shot deferred (resolve captured for a later single call). */
private makeDeferred<T>(): Deferred<T> {
let resolve!: (value: T) => void;
const promise = new Promise<T>((r) => {
resolve = r;
});
return { promise, resolve };
}
/** #487: resolve a run's settle notifier EXACTLY ONCE, then drop it (bounded).
* A subscriber that already grabbed the promise still resolves; a later one
* falls back to the zombie map / the row (see peekSettled). */
private resolveSettled(runId: string, outcome: RunSettleOutcome): void {
const d = this.settledPromises.get(runId);
if (!d) return;
this.settledPromises.delete(runId);
d.resolve(outcome);
}
/** #487: read the persisted terminal outcome when the conditional finalize was a
* no-op (the row was already terminal). Falls back to the intended status when
* the read fails or the row is unexpectedly missing/non-terminal. */
private async readTerminalOutcome(
runId: string,
workspaceId: string,
fallbackStatus: RunTerminalStatus,
fallbackError: string | null,
): Promise<RunSettleOutcome> {
try {
const row = await this.runRepo.findById(runId, workspaceId);
if (row && isRunTerminal(row.status)) {
return {
status: row.status as RunTerminalStatus,
error: row.error ?? null,
terminalWriteFailed: false,
};
}
} catch {
// Fall through to the intended status — best-effort only.
}
return {
status: fallbackStatus,
error: fallbackError,
terminalWriteFailed: false,
};
}
/** Small async backoff between terminal-write retries (F6). Isolated so it is
@@ -115,7 +115,7 @@ describe('finalizeAssistant dispatch (planFinalizeAssistant + applyFinalize)', (
// Drive the SAME applyFinalize the service calls (no duplicated logic).
async function dispatchFinalize(
repo: { insert: jest.Mock; update: jest.Mock },
repo: { insert: jest.Mock; finalizeOwner: jest.Mock },
assistantId: string | undefined,
flushed: AssistantFlush,
): Promise<void> {
@@ -135,21 +135,22 @@ describe('finalizeAssistant dispatch (planFinalizeAssistant + applyFinalize)', (
expect(planFinalizeAssistant(undefined)).toEqual({ kind: 'insert' });
});
it('(a) upfront insert succeeded -> finalize UPDATEs the row by id', async () => {
const repo = { insert: jest.fn(), update: jest.fn() };
it('(a) upfront insert succeeded -> finalize CONDITIONALLY updates the row by id (#487 owner-write)', async () => {
const repo = { insert: jest.fn(), finalizeOwner: jest.fn() };
const flushed = flushAssistant([], 'final answer', 'completed', {
finishReason: 'stop',
});
await dispatchFinalize(repo, 'a1', flushed);
expect(repo.update).toHaveBeenCalledWith('a1', workspaceId, flushed);
// #487: the owner write is the CONDITIONAL finalizeOwner, not a raw update.
expect(repo.finalizeOwner).toHaveBeenCalledWith('a1', workspaceId, flushed);
expect(repo.insert).not.toHaveBeenCalled();
});
it('(b) upfront insert failed -> finalize INSERTs the terminal payload', async () => {
const repo = { insert: jest.fn(), update: jest.fn() };
const repo = { insert: jest.fn(), finalizeOwner: jest.fn() };
const flushed = flushAssistant([], 'partial', 'error', { error: 'boom' });
await dispatchFinalize(repo, undefined, flushed);
expect(repo.update).not.toHaveBeenCalled();
expect(repo.finalizeOwner).not.toHaveBeenCalled();
expect(repo.insert).toHaveBeenCalledTimes(1);
const arg = repo.insert.mock.calls[0][0];
// The fallback insert carries the terminal content/status/metadata.
@@ -0,0 +1,221 @@
import {
BadRequestException,
ConflictException,
HttpException,
} from '@nestjs/common';
import { AiChatController } from './ai-chat.controller';
import type { User, Workspace } from '@docmost/db/types/entity.types';
/**
* #487 commit 3 — the single concurrency GATE (both modes) + the server supersede
* CAS, at the controller boundary. The gate + CAS run BEFORE res.hijack(), so a
* rejected concurrent start / a CAS branch returns clean JSON (an HttpException
* the controller's post-hijack catch re-serializes). These assert the OBSERVABLE
* HTTP contract against the real controller + a stubbed run service.
*/
describe('#487 AiChatController.stream — gate + supersede', () => {
const user = { id: 'u1' } as User;
function wsWith(autonomousRuns: boolean): Workspace {
return {
id: 'ws1',
settings: { ai: { chat: true, autonomousRuns } },
} as unknown as Workspace;
}
function makeReqRes(body: Record<string, unknown>) {
const req = {
raw: { sessionId: 'sess', once: jest.fn(), destroyed: false },
body,
};
const res = {
raw: {
writableEnded: false,
headersSent: false,
on: jest.fn(),
once: jest.fn(),
setHeader: jest.fn(),
end: jest.fn(),
statusCode: 200,
flushHeaders: jest.fn(),
},
hijack: jest.fn(),
status: jest.fn().mockReturnThis(),
send: jest.fn(),
};
return { req, res };
}
function makeController(runServiceOverrides: Record<string, jest.Mock>) {
const aiChatService = {
resolveRoleForRequest: jest.fn().mockResolvedValue(null),
getChatModel: jest.fn().mockResolvedValue({}),
stream: jest.fn().mockResolvedValue(undefined),
};
const aiChatRunService = {
getActiveForChat: jest.fn().mockResolvedValue(undefined),
supersede: jest.fn(),
beginRun: jest.fn().mockResolvedValue({
runId: 'run-new',
signal: new AbortController().signal,
}),
linkAssistantMessage: jest.fn(),
recordStep: jest.fn(),
finalizeRun: jest.fn(),
requestStop: jest.fn(),
...runServiceOverrides,
};
const controller = new AiChatController(
aiChatService as never,
aiChatRunService as never,
{} as never, // aiChatRepo
{} as never, // aiChatMessageRepo
{} as never, // aiTranscription
{} as never, // pageRepo
);
return { controller, aiChatService, aiChatRunService };
}
const codeOf = (err: unknown) =>
(((err as HttpException).getResponse() as Record<string, unknown>) ?? {})
.code;
describe('single concurrency gate — BOTH modes reject the second tab with 409', () => {
for (const autonomousRuns of [true, false]) {
it(`rejects a concurrent start with 409 A_RUN_ALREADY_ACTIVE (autonomousRuns=${autonomousRuns})`, async () => {
const { controller, aiChatRunService } = makeController({
getActiveForChat: jest
.fn()
.mockResolvedValue({ id: 'run-live', chatId: 'c1' }),
});
const { req, res } = makeReqRes({ chatId: 'c1' });
let thrown: unknown;
try {
await controller.stream(
req as never,
res as never,
user,
wsWith(autonomousRuns),
);
} catch (e) {
thrown = e;
}
expect(thrown).toBeInstanceOf(ConflictException);
expect((thrown as HttpException).getStatus()).toBe(409);
expect(codeOf(thrown)).toBe('A_RUN_ALREADY_ACTIVE');
// Rejected BEFORE committing to the stream (no hijack, no service.stream).
expect(res.hijack).not.toHaveBeenCalled();
expect(aiChatRunService.getActiveForChat).toHaveBeenCalledWith(
'c1',
'ws1',
);
});
}
});
it('supersede MISMATCH -> 409 SUPERSEDE_TARGET_MISMATCH carrying the current runId', async () => {
const { controller } = makeController({
supersede: jest
.fn()
.mockResolvedValue({ kind: 'mismatch', activeRunId: 'run-other' }),
});
const { req, res } = makeReqRes({
chatId: 'c1',
supersede: { runId: 'run-x' },
});
let thrown: unknown;
try {
await controller.stream(req as never, res as never, user, wsWith(true));
} catch (e) {
thrown = e;
}
expect(thrown).toBeInstanceOf(ConflictException);
expect(codeOf(thrown)).toBe('SUPERSEDE_TARGET_MISMATCH');
expect(
((thrown as HttpException).getResponse() as Record<string, unknown>)
.activeRunId,
).toBe('run-other');
expect(res.hijack).not.toHaveBeenCalled();
});
it('supersede TIMEOUT -> 409 SUPERSEDE_TIMEOUT, nothing streamed', async () => {
const { controller } = makeController({
supersede: jest.fn().mockResolvedValue({ kind: 'timeout' }),
});
const { req, res } = makeReqRes({
chatId: 'c1',
supersede: { runId: 'run-x' },
});
let thrown: unknown;
try {
await controller.stream(req as never, res as never, user, wsWith(false));
} catch (e) {
thrown = e;
}
expect(thrown).toBeInstanceOf(ConflictException);
expect(codeOf(thrown)).toBe('SUPERSEDE_TIMEOUT');
expect(res.hijack).not.toHaveBeenCalled();
});
it('supersede INVALID (target on another chat) -> 400 SUPERSEDE_INVALID', async () => {
const { controller } = makeController({
supersede: jest.fn().mockResolvedValue({ kind: 'invalid' }),
});
const { req, res } = makeReqRes({
chatId: 'c1',
supersede: { runId: 'run-x' },
});
let thrown: unknown;
try {
await controller.stream(req as never, res as never, user, wsWith(true));
} catch (e) {
thrown = e;
}
expect(thrown).toBeInstanceOf(BadRequestException);
expect(codeOf(thrown)).toBe('SUPERSEDE_INVALID');
});
it('supersede without chatId -> 400 SUPERSEDE_INVALID', async () => {
const { controller, aiChatRunService } = makeController({});
const { req, res } = makeReqRes({ supersede: { runId: 'run-x' } });
let thrown: unknown;
try {
await controller.stream(req as never, res as never, user, wsWith(true));
} catch (e) {
thrown = e;
}
expect(thrown).toBeInstanceOf(BadRequestException);
expect(codeOf(thrown)).toBe('SUPERSEDE_INVALID');
expect(aiChatRunService.supersede).not.toHaveBeenCalled();
});
it('supersede READY -> proceeds to stream with superseded=true', async () => {
const { controller, aiChatService } = makeController({
supersede: jest.fn().mockResolvedValue({ kind: 'ready' }),
getActiveForChat: jest.fn().mockResolvedValue(undefined), // slot free after CAS
});
const { req, res } = makeReqRes({
chatId: 'c1',
supersede: { runId: 'run-x' },
});
await controller.stream(req as never, res as never, user, wsWith(true));
expect(res.hijack).toHaveBeenCalled();
expect(aiChatService.stream).toHaveBeenCalledTimes(1);
expect(aiChatService.stream.mock.calls[0][0].superseded).toBe(true);
// The run hooks are always present now (both modes).
expect(aiChatService.stream.mock.calls[0][0].runHooks).toBeDefined();
});
it('supersede DEGRADE -> proceeds to a normal send (superseded=false)', async () => {
const { controller, aiChatService } = makeController({
supersede: jest.fn().mockResolvedValue({ kind: 'degrade' }),
});
const { req, res } = makeReqRes({
chatId: 'c1',
supersede: { runId: 'run-x' },
});
await controller.stream(req as never, res as never, user, wsWith(false));
expect(aiChatService.stream).toHaveBeenCalledTimes(1);
expect(aiChatService.stream.mock.calls[0][0].superseded).toBe(false);
});
});
@@ -432,12 +432,66 @@ 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) {
// #487: server-side supersede CAS ("interrupt and send now"). When the client
// asks to replace a live run, atomically STOP it and wait for it to settle
// before this turn claims the slot. Runs BEFORE hijack so every branch returns
// clean JSON (the client keeps the composer text on a 409). See
// AiChatRunService.supersede for the branch semantics.
let superseded = false;
const supersedeRunId = body.supersede?.runId;
if (supersedeRunId) {
if (!body.chatId) {
throw new BadRequestException({
message: 'supersede requires chatId',
code: 'SUPERSEDE_INVALID',
});
}
const result = await this.aiChatRunService.supersede(
body.chatId,
supersedeRunId,
workspace.id,
);
switch (result.kind) {
case 'invalid':
throw new BadRequestException({
message: 'The run to supersede does not belong to this chat',
code: 'SUPERSEDE_INVALID',
});
case 'mismatch':
// A DIFFERENT run is active than the one the client targeted. Surface
// the CURRENT runId; the client does NOT auto-retry (a stale CAS).
throw new ConflictException({
message: 'A different agent run is now active on this chat',
code: 'SUPERSEDE_TARGET_MISMATCH',
activeRunId: result.activeRunId,
});
case 'timeout':
// The target did not settle within W — nothing was persisted, the
// composer keeps the text. NOT a rollback: the stop is already issued.
throw new ConflictException({
message:
'The previous run did not stop in time; nothing was sent — please try again',
code: 'SUPERSEDE_TIMEOUT',
});
case 'ready':
// The target stopped and settled: the slot is free. Prompt the new run
// that the old run's last operations may still be applying.
superseded = true;
break;
case 'degrade':
// The run already ended between click and POST — send normally.
break;
}
}
// #487: one active run per chat — ENFORCED IN BOTH MODES now (legacy mode used
// to have NO gate, so two tabs streamed two parallel turns on one chat, which
// interleaved history and crashed convertToModelMessages). Reject a concurrent
// start with a clean pre-hijack 409 (double-submit / second-tab). A brand-new
// chat (no chatId) cannot have a prior run, and the DB partial unique index in
// beginRun is the authoritative backstop for any race that slips past here
// (including a slot stolen between a supersede release and beginRun).
if (body.chatId) {
const active = await this.aiChatRunService.getActiveForChat(
body.chatId,
workspace.id,
@@ -446,107 +500,94 @@ export class AiChatController {
throw new ConflictException({
message: 'An agent run is already in progress for this chat',
code: 'A_RUN_ALREADY_ACTIVE',
activeRunId: active.id,
});
}
}
// 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: async (chatId) => {
const handle = await this.aiChatRunService.beginRun({
chatId,
workspaceId: workspace.id,
userId: user.id,
trigger: 'user',
});
// #184 phase 1.5: register the run-stream entry at BEGIN (before any
// frame) so a tab that attaches in the begin->seed window finds an
// entry to wait on. Gated on AI_CHAT_RESUMABLE_STREAM: with the flag
// off nothing is registered and attach always 204s.
if (
handle?.runId &&
this.environment?.isAiChatResumableStreamEnabled?.()
) {
this.streamRegistry?.open(chatId, handle.runId);
}
return handle;
},
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,
),
// #487: the turn is ALWAYS a first-class RUN now (both modes). The mode
// difference is only the abort semantics on a browser disconnect (onClose
// below). currentRunId is captured at begin so a legacy disconnect can stop
// the run through its stop lever.
let currentRunId: string | undefined;
const runHooks: AiChatRunHooks = {
begin: async (chatId) => {
const handle = await this.aiChatRunService.beginRun({
chatId,
workspaceId: workspace.id,
userId: user.id,
trigger: 'user',
});
currentRunId = handle?.runId;
// #184 phase 1.5: register the run-stream entry at BEGIN (before any
// frame) so a tab that attaches in the begin->seed window finds an entry
// to wait on. Gated on AI_CHAT_RESUMABLE_STREAM.
if (
handle?.runId &&
this.environment?.isAiChatResumableStreamEnabled?.()
) {
this.streamRegistry?.open(chatId, handle.runId);
}
: undefined;
return handle;
},
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),
};
// 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;
// we also drop it on response `finish` so it never lingers after the stream
// completes normally (the AI SDK pipes the response fire-and-forget, so we
// cannot simply remove it once `stream()` returns).
// Handle a client disconnect. `close` also fires on normal completion, so only
// act when the response has not finished writing (a genuine disconnect). `once`
// fires at most once and self-removes; we also drop it on response `finish`.
// DIAGNOSTIC (Safari stream-drop investigation) — temporary: wall-clock at
// which a Safari disconnect is observed, measured from request receipt.
const reqStartedAt = Date.now();
const controller = new AbortController();
const onClose = (): void => {
// 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.
if (!res.raw.writableEnded) {
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.
// #184: a DETACHED run — a disconnect must NOT stop 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 {
// #487: legacy — a disconnect ENDS the turn, but the turn is now a RUN,
// so stop it through the run's stop lever (requestStop). streamText no
// longer consumes the socket signal (effectiveSignal is the run signal),
// so aborting `controller` would do nothing; requestStop aborts the run.
this.logger.warn(
`AI chat stream: client disconnected before completion; aborting turn ` +
`(elapsed=${Date.now() - reqStartedAt}ms since request received)`,
`AI chat stream: client disconnected before completion; stopping the ` +
`run (elapsed=${Date.now() - reqStartedAt}ms since request received)`,
);
controller.abort();
if (currentRunId) {
void this.aiChatRunService.requestStop(currentRunId, workspace.id);
}
}
}
};
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)
}`,
);
});
}
// #184/#487: the run/pipe can outlive the socket in BOTH modes now (autonomous
// keeps going; legacy keeps going until requestStop's abort unwinds the turn).
// The SDK's pipe may then write to a dropped socket and emit an 'error' on the
// raw response — swallow it so it never surfaces as an unhandled error event.
res.raw.on('error', (err) => {
this.logger.debug(
`AI chat 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.
@@ -562,8 +603,10 @@ export class AiChatController {
signal: controller.signal,
model,
role,
// #184: present only when the flag is on; wraps the turn in a durable run.
// #487: the turn is always run-wrapped now (both modes).
runHooks,
// #487: warn the new run that a superseded run's last ops may still apply.
superseded,
});
} catch (err) {
// Any failure AFTER hijack can no longer go through Nest's exception
@@ -101,6 +101,22 @@ const INTERRUPT_NOTE =
'assume your previous response was complete, and do not silently restart the ' +
'partial work — build on it or follow the new instruction.';
/**
* #487: injected on a turn started by SUPERSEDING a previous run (the user hit
* "interrupt and send now" while a run was live). The previous run was Stopped,
* but there is NO side-effect quiescence — a write it had already committed, or
* one committing at the moment of Stop, may land with a small delay AFTER this new
* run starts. So the model is told its picture of the page/state may be a beat
* stale and to re-read before assuming an edit did or did not apply.
*/
const SUPERSEDE_NOTE =
'NOTE: A previous agent run in this conversation was just interrupted so this ' +
'new turn could start. That run was stopped, but any operation it had already ' +
'begun (e.g. a page edit) may still be applied with a short delay. Do not ' +
'assume the document/state is exactly as the interrupted run left it — if you ' +
'need to rely on the current content, RE-READ it with the page tools before ' +
'acting rather than trusting a cached view.';
/**
* Injected on a turn where the open page was hand-edited by the user (or anyone
* else) AFTER the agent's previous response ended (#274). The server takes a
@@ -203,6 +219,14 @@ export interface BuildSystemPromptInput {
* (partial) answer was cut off by the user's new message.
*/
interrupted?: boolean;
/**
* #487: true when THIS turn was started by superseding a still-live previous run
* ("interrupt and send now"). Adds SUPERSEDE_NOTE so the model knows the previous
* run's last operations may still be applying and to re-read state it depends on.
* Distinct from `interrupted` (which is about a PARTIAL prior answer in history);
* both can be set together. Self-clears — set only for the superseding turn.
*/
superseded?: boolean;
/**
* Set only when the open page was edited by the user AFTER the agent's previous
* turn ended (#274), confirmed server-side by diffing the current page against
@@ -311,6 +335,7 @@ export function buildSystemPrompt({
openedPage,
mcpInstructions,
interrupted,
superseded,
pageChanged,
deferredToolsEnabled,
toolCatalog,
@@ -360,6 +385,13 @@ export function buildSystemPrompt({
context += `\n${INTERRUPT_NOTE}`;
}
// Supersede note (#487): present only for a turn that stopped and replaced a
// still-live previous run — warns the model the previous run's last operations
// may still be applying (no side-effect quiescence).
if (superseded) {
context += `\n${SUPERSEDE_NOTE}`;
}
// Per-turn page-change note (#274). Added to the context section (inside the
// safety sandwich), present only when the server detected that the open page
// was edited by the user since the agent's last turn ended. The diff content is
@@ -89,6 +89,11 @@ describe('AiChatService.stream run-lifecycle safety net (#184)', () => {
const runRepo = {
insert: jest.fn().mockResolvedValue({ id: 'run-1', status: 'running' }),
update: jest.fn().mockResolvedValue({ id: 'run-1' }),
// #487: the terminal settle now goes through the CONDITIONAL write.
finalizeIfActive: jest
.fn()
.mockResolvedValue({ id: 'run-1', status: 'failed' }),
findById: jest.fn().mockResolvedValue(undefined),
};
const runService = new AiChatRunService(runRepo as never, { isCloud: () => false } as never);
@@ -148,9 +153,10 @@ describe('AiChatService.stream run-lifecycle safety net (#184)', () => {
// 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(
// ...then settled to a terminal FAILED status by the safety net (via the
// #487 conditional write)...
expect(runRepo.finalizeIfActive).toHaveBeenCalledTimes(1);
expect(runRepo.finalizeIfActive).toHaveBeenCalledWith(
'run-1',
'ws1',
expect.objectContaining({ status: 'failed' }),
@@ -155,6 +155,8 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
insert: jest.fn(async () => ({ id: 'msg-1' })),
findAllByChat: jest.fn(async () => []),
update: jest.fn(async () => ({ id: 'msg-1' })),
finalizeOwner: jest.fn(async () => ({ id: 'msg-1' })),
findStreamingWithTerminalRun: jest.fn(async () => []),
};
const aiSettings = { resolve: jest.fn(async () => ({})) };
const tools = { forUser: jest.fn(async () => ({})) };
@@ -332,7 +334,13 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
usage: {},
steps: [],
});
expect(runHooks.onSettled).toHaveBeenCalledWith('run-1', 'completed');
// #487: onFinish passes the (undefined) error slot so a message-finalize
// failure could error-mark the run; on the success path it is undefined.
expect(runHooks.onSettled).toHaveBeenCalledWith(
'run-1',
'completed',
undefined,
);
});
it('F9: onAbort settles the run "aborted"', async () => {
@@ -415,6 +423,8 @@ describe('AiChatService.stream — begin-failure fails the turn (#184 F14 / #486
insert: jest.fn(async () => ({ id: 'msg-1' })),
findAllByChat: jest.fn(async () => []),
update: jest.fn(async () => ({ id: 'msg-1' })),
finalizeOwner: jest.fn(async () => ({ id: 'msg-1' })),
findStreamingWithTerminalRun: jest.fn(async () => []),
};
const aiSettings = { resolve: jest.fn(async () => ({})) };
const tools = { forUser: jest.fn(async () => ({})) };
@@ -1315,8 +1315,12 @@ describe('AiChatService page-change lifecycle (#274)', () => {
describe('isInterruptResume', () => {
// history tail is the just-inserted user row; [len-2] is the previous turn.
const withPrev = (
prev: { role: string; status?: string | null } | null,
): Array<{ role: string; status?: string | null }> =>
prev: {
role: string;
status?: string | null;
metadata?: unknown;
} | null,
): Array<{ role: string; status?: string | null; metadata?: unknown }> =>
prev
? [prev, { role: 'user', status: null }]
: [{ role: 'user', status: null }];
@@ -1357,6 +1361,33 @@ describe('isInterruptResume', () => {
it('false when there is no preceding turn (only the new user row)', () => {
expect(isInterruptResume(withPrev(null), true)).toBe(false);
});
it('#487 EXCLUDES a reconcile stamp (finalizeFailed) — not a genuine interruption', () => {
// A row a reconcile settled to 'aborted' carries metadata.finalizeFailed. It
// must NOT be treated as an interrupt-resume (that would inject a false
// "you were interrupted" note), even though its status is 'aborted'.
expect(
isInterruptResume(
withPrev({
role: 'assistant',
status: 'aborted',
metadata: { finalizeFailed: true },
}),
true,
),
).toBe(false);
// A genuine abort (no finalizeFailed) still counts.
expect(
isInterruptResume(
withPrev({
role: 'assistant',
status: 'aborted',
metadata: { parts: [] },
}),
true,
),
).toBe(true);
});
});
/**
@@ -1419,6 +1450,9 @@ describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', ()
insert: jest.fn(async () => ({ id: 'msg-1' })),
findAllByChat: jest.fn(async () => []),
update: jest.fn(async () => ({ id: 'msg-1' })),
// #487: the terminal owner-write + the opportunistic reconcile query.
finalizeOwner: jest.fn(async () => ({ id: 'msg-1' })),
findStreamingWithTerminalRun: jest.fn(async () => []),
};
const aiSettings = { resolve: jest.fn(async () => ({})) };
const tools = { forUser: jest.fn(async () => ({})) };
@@ -1623,6 +1657,19 @@ describe('AiChatService.stream — token-degeneration reaction (#444)', () => {
return { id };
},
),
// #487: the terminal owner-write records into the SAME `updated` recorder so
// assertions on the terminal 'completed'/'error'/'aborted' write still hold.
finalizeOwner: jest.fn(
async (
id: string,
workspaceId: string,
patch: Record<string, unknown>,
) => {
updated.push({ id, workspaceId, patch });
return { id };
},
),
findStreamingWithTerminalRun: jest.fn(async () => []),
};
const aiSettings = { resolve: jest.fn(async () => ({})) };
const tools = { forUser: jest.fn(async () => ({})) };
+284 -35
View File
@@ -3,6 +3,7 @@ import {
ForbiddenException,
Injectable,
Logger,
OnModuleDestroy,
OnModuleInit,
ServiceUnavailableException,
} from '@nestjs/common';
@@ -42,7 +43,11 @@ import {
makeLoadToolsTool,
buildExternalToolCatalog,
} from './tools/tool-tiers';
import { RunAlreadyActiveError } from './ai-chat-run.service';
import {
RunAlreadyActiveError,
AiChatRunService,
} from './ai-chat-run.service';
import { inAppToolCallCapMs } from './tools/ai-chat-tools.service';
import { computePageChange } from './page-change/page-change.util';
import {
sanitizeSelection,
@@ -250,15 +255,23 @@ export function cleanGeneratedTitle(text: string): string {
* partial output is already in history thanks to the step-granular write path).
*/
export function isInterruptResume(
history: Array<{ role: string; status?: string | null }>,
history: Array<{
role: string;
status?: string | null;
metadata?: unknown;
}>,
clientInterrupted: boolean | undefined,
): boolean {
if (clientInterrupted !== true) return false;
const prev = history[history.length - 2];
return (
prev?.role === 'assistant' &&
(prev.status === 'aborted' || prev.status === 'streaming')
);
if (prev?.role !== 'assistant') return false;
// #487: a reconcile STAMP (metadata.finalizeFailed) is NOT a genuine user
// interruption — the previous turn's process died and a reconcile settled the
// row as 'aborted'. Treating it as an interrupt-resume would inject a false
// "you were interrupted" note. Exclude any finalizeFailed row.
const meta = prev.metadata as { finalizeFailed?: unknown } | null | undefined;
if (meta && meta.finalizeFailed === true) return false;
return prev.status === 'aborted' || prev.status === 'streaming';
}
/**
@@ -379,6 +392,14 @@ export interface AiChatStreamBody {
// it against persisted history (`isInterruptResume`) before injecting the
// interrupt note, so a spoofed/stale flag on an ordinary turn is ignored.
interrupted?: boolean;
// #487: server-side supersede CAS. When present, this POST asks the server to
// STOP the run `supersede.runId` (which the client saw as the chat's active run)
// and, once it has settled, start THIS turn in its place. The server validates
// the target against the chat and answers 400 (wrong chat) / 409
// SUPERSEDE_TARGET_MISMATCH / 409 SUPERSEDE_TIMEOUT, or proceeds normally
// (degrade / ready). Absent => an ordinary send (rejected with 409
// A_RUN_ALREADY_ACTIVE if a run is already active on the chat).
supersede?: { runId?: string } | null;
// useChat sends the full UIMessage list; the last one is the new user turn.
messages?: UIMessage[];
}
@@ -428,6 +449,11 @@ export interface AiChatStreamArgs {
// chat row (existing chat) or the request body (new chat). null => universal
// assistant. Carried here so the turn never re-loads it.
role: AiAgentRole | null;
// #487: true when this turn was started by SUPERSEDING a still-live previous run
// (the controller ran the supersede CAS to a `ready` result). Adds the
// SUPERSEDE_NOTE to the system prompt (the previous run's last ops may still be
// applying — no side-effect quiescence). Absent on an ordinary send.
superseded?: boolean;
}
/**
@@ -444,7 +470,7 @@ export interface AiChatStreamArgs {
* can be rebuilt for `convertToModelMessages`.
*/
@Injectable()
export class AiChatService implements OnModuleInit {
export class AiChatService implements OnModuleInit, OnModuleDestroy {
private readonly logger = new Logger(AiChatService.name);
constructor(
@@ -465,8 +491,17 @@ export class AiChatService implements OnModuleInit {
// constructions (int-specs) compile unchanged; Nest always injects the real
// provider in production. Only ever touched on the run-wrapped + flag-on path.
private readonly streamRegistry?: AiChatStreamRegistryService,
// #487: the run lifecycle service, for the periodic + opportunistic reconcile
// (zombie re-drive + stale-run abort). OPTIONAL so positional test
// constructions compile unchanged; Nest always injects the real singleton, so
// reconcile sees the SAME in-memory active/zombie maps the runner mutates.
private readonly aiChatRunService?: AiChatRunService,
) {}
// #487: periodic reconcile timer (single-process phase 1). Started in
// onModuleInit, cleared in onModuleDestroy.
private reconcileTimer?: ReturnType<typeof setInterval>;
/**
* Crash-recovery sweep on server start (#183): any assistant row left in the
* 'streaming' state is the relic of a turn whose process died before it
@@ -491,6 +526,158 @@ export class AiChatService implements OnModuleInit {
}`,
);
}
// #487: start the PERIODIC reconcile (was boot-only). It heals both directions
// of the run<->message lifecycle asymmetry that a boot sweep alone left to the
// NEXT restart. Single-process phase 1: the in-memory active/zombie maps are
// authoritative, so "no live entry" is a safe primary gate.
const staleMs = this.reconcileStalenessMs();
// boot-warn if the per-call cap is configured so high the derived staleness is
// unusually long (a stale run then lingers longer before reconcile aborts it).
if (staleMs > 30 * 60 * 1000) {
this.logger.warn(
`#487 reconcile staleness is ${Math.round(staleMs / 60000)}min ` +
`(derived from max(2 x per-call cap, 15min)); a per-call cap this high ` +
`delays stale-run recovery. Review AI_CHAT_INAPP_TOOL_CALL_CAP_MS.`,
);
}
const intervalMs = this.reconcileIntervalMs();
this.reconcileTimer = setInterval(() => {
void this.reconcile().catch((err) => {
this.logger.warn(
`Periodic reconcile failed: ${
err instanceof Error ? err.message : 'unknown error'
}`,
);
});
}, intervalMs);
this.reconcileTimer.unref?.();
}
/** #487: stop the periodic reconcile timer on shutdown. */
onModuleDestroy(): void {
if (this.reconcileTimer) {
clearInterval(this.reconcileTimer);
this.reconcileTimer = undefined;
}
}
/**
* #487: reconcile staleness threshold X a run/message is only a "no live
* runner" abort candidate once UNTOUCHED past this. Derived as
* max(2 x per-call cap, 15min): 2x the longest legitimate single tool call plus
* a floor, so a marathon turn making steady progress (updatedAt bumped each
* step) is never swept.
*/
private reconcileStalenessMs(): number {
return Math.max(2 * inAppToolCallCapMs(), 15 * 60 * 1000);
}
/** #487: how often the periodic reconcile runs (env-tunable, default 2min). */
private reconcileIntervalMs(): number {
const raw = Number(process.env.AI_CHAT_RECONCILE_INTERVAL_MS);
return Number.isFinite(raw) && raw > 0 ? raw : 2 * 60 * 1000;
}
/**
* #487: the periodic BIDIRECTIONAL reconcile. Runs the clauses IN ORDER; each is
* best-effort (a failure of one never blocks the others). Single-process phase 1
* the run service's in-memory maps are authoritative for "live entry".
*
* (a) re-drive ZOMBIE runs (a terminal write that gave up) apply the intended
* status via the conditional UPDATE;
* (b) message 'streaming' + its RUN terminal -> stamp the message by the run's
* status (succeeded-run + stuck row -> 'aborted'+finalizeFailed, NOT
* 'completed' with empty parts the final text lived only in the dead
* process's memory, a documented loss);
* (c) run active + NO live entry + NO zombie + stale -> aborted (the run
* service applies the "no entry" primary gate + last-progress staleness);
* (d) message 'streaming' + age>X + NO active run on the chat -> aborted
* (historical-row safety, double-gated).
*/
async reconcile(): Promise<void> {
const staleMs = this.reconcileStalenessMs();
// (a) zombie re-drive.
if (this.aiChatRunService) {
for (const runId of this.aiChatRunService.zombieRunIds()) {
try {
await this.aiChatRunService.settleZombie(runId);
} catch (err) {
this.logger.warn(
`Reconcile (a) zombie ${runId} re-drive failed: ${
err instanceof Error ? err.message : 'unknown error'
}`,
);
}
}
}
// (b) message streaming + run terminal -> stamp message by run status.
try {
const stuck = await this.aiChatMessageRepo.findStreamingWithTerminalRun();
for (const s of stuck) {
// succeeded-run -> 'aborted' (NOT 'completed'-empty); failed -> 'error';
// aborted -> 'aborted'. All via the finalizeFailed stamp.
const status = s.runStatus === 'failed' ? 'error' : 'aborted';
await this.aiChatMessageRepo.stampTerminalIfStreaming(
s.messageId,
s.workspaceId,
status,
);
}
} catch (err) {
this.logger.warn(
`Reconcile (b) message<-run failed: ${
err instanceof Error ? err.message : 'unknown error'
}`,
);
}
// (c) stale active run with no live runner -> aborted.
if (this.aiChatRunService) {
try {
await this.aiChatRunService.reconcileStaleRuns(staleMs);
} catch (err) {
this.logger.warn(
`Reconcile (c) stale-run abort failed: ${
err instanceof Error ? err.message : 'unknown error'
}`,
);
}
}
// (d) historical streaming row, no active run on the chat, stale -> aborted.
try {
await this.aiChatMessageRepo.sweepStreamingWithoutActiveRun(staleMs);
} catch (err) {
this.logger.warn(
`Reconcile (d) historical-row sweep failed: ${
err instanceof Error ? err.message : 'unknown error'
}`,
);
}
}
/**
* #487: OPPORTUNISTIC single-chat reconcile at the start of a turn (beginRun /
* supersede path), so a user who returns to a chat with a stuck streaming row
* (its run already terminal) sees it settled WITHOUT waiting for the periodic
* job. Best-effort a failure NEVER fails the turn (swallowed by the caller).
*/
async reconcileChat(chatId: string, workspaceId: string): Promise<void> {
const stuck = await this.aiChatMessageRepo.findStreamingWithTerminalRun(50, {
chatId,
workspaceId,
});
for (const s of stuck) {
const status = s.runStatus === 'failed' ? 'error' : 'aborted';
await this.aiChatMessageRepo.stampTerminalIfStreaming(
s.messageId,
s.workspaceId,
status,
);
}
}
/**
@@ -727,6 +914,7 @@ export class AiChatService implements OnModuleInit {
model,
role,
runHooks,
superseded,
}: AiChatStreamArgs): Promise<void> {
// Resolve / create the chat. A new chat is created when no valid chatId is
// supplied or the supplied one does not belong to this workspace.
@@ -828,6 +1016,20 @@ export class AiChatService implements OnModuleInit {
}
}
// #487: opportunistic single-chat reconcile — settle any streaming row on this
// chat whose run is already terminal BEFORE this turn's history load, so the
// user never waits on the periodic job and the new turn's model history is not
// polluted by a stuck 'streaming' row. Best-effort: it must NEVER fail the turn.
try {
await this.reconcileChat(chatId, workspace.id);
} catch (err) {
this.logger.debug(
`Opportunistic reconcile for chat ${chatId} failed (ignored): ${
err instanceof Error ? err.message : 'unknown error'
}`,
);
}
try {
// Extract the incoming user turn (the last user message from useChat).
const incoming = lastUserMessage(body.messages);
@@ -1033,6 +1235,9 @@ export class AiChatService implements OnModuleInit {
// History-confirmed interrupt-resume flag (#198): adds the interrupt note
// so the model treats the partial answer above as cut off, not finished.
interrupted,
// #487: this turn superseded a still-live run — warn the model the
// previous run's last ops may still be applying (no quiescence).
superseded,
// Detected between-turns human edit to the open page (#274): adds the
// page_changed note + unified diff so the agent doesn't overwrite it.
pageChanged,
@@ -1187,29 +1392,59 @@ export class AiChatService implements OnModuleInit {
// callbacks — mirroring the pre-#183 persist-at-most-once guard for the
// TERMINAL status (the row may be updated many times with 'streaming' before
// this fires once).
// #487: the once-gate closes ONLY AFTER a successful write, and the write is
// BOUNDED-RETRIED. Previously `finalized` was set BEFORE the write and never
// retried, so a single failed UPDATE stranded the row 'streaming' forever
// (the boot-only sweep was the only recovery). Now a transient blip is ridden
// out in place; a total give-up leaves the gate OPEN and logs, and the
// periodic reconcile (clauses b/d) later settles the row. Returns whether the
// terminal write LANDED, so the caller can error-mark the RUN on a message
// failure (the run is finalized regardless — never gated on the message).
let finalized = false;
const FINALIZE_MSG_MAX_ATTEMPTS = 3;
const finalizeAssistant = async (
flushed: AssistantFlush,
): Promise<void> => {
if (finalized) return;
finalized = true;
): Promise<boolean> => {
if (finalized) return true;
const plan = planFinalizeAssistant(assistantId);
try {
// Shared dispatch (see applyFinalize): UPDATE the upfront row, or — when
// the upfront insert failed (kind 'insert') — INSERT the terminal row as
// the only safety against losing the turn entirely.
await applyFinalize(
this.aiChatMessageRepo,
plan,
{ chatId, workspaceId: workspace.id, userId: user.id },
flushed,
);
} catch (err) {
this.logger.error(
`Failed to finalize assistant message (kind=${plan.kind})`,
err as Error,
);
let lastError: unknown;
for (let attempt = 1; attempt <= FINALIZE_MSG_MAX_ATTEMPTS; attempt++) {
try {
// Shared dispatch (see applyFinalize): conditionally UPDATE the upfront
// row (owner-write priority), or — when the upfront insert failed (kind
// 'insert') — INSERT the terminal row as the only safety against losing
// the turn entirely.
await applyFinalize(
this.aiChatMessageRepo,
plan,
{ chatId, workspaceId: workspace.id, userId: user.id },
flushed,
);
finalized = true; // gate closes ONLY after a successful write
return true;
} catch (err) {
lastError = err;
this.logger.warn(
`Assistant message finalize attempt ${attempt}/${FINALIZE_MSG_MAX_ATTEMPTS} ` +
`failed (kind=${plan.kind}): ${
err instanceof Error ? err.message : 'unknown error'
}`,
);
if (attempt < FINALIZE_MSG_MAX_ATTEMPTS) {
await new Promise((r) => setTimeout(r, 50 * attempt));
}
}
}
// Gave up: leave the gate OPEN (no in-process second settler exists — the
// terminal callbacks are mutually exclusive) and log. The periodic reconcile
// settles the stranded row; a late owner-write is impossible for this turn,
// so the reconcile stamp (aborted+finalizeFailed) is the final state.
this.logger.error(
`Assistant message finalize GAVE UP after ${FINALIZE_MSG_MAX_ATTEMPTS} ` +
`attempts (row left 'streaming', chat ${chatId}); reconcile will settle it`,
lastError as Error,
);
return false;
};
// DIAGNOSTIC (Safari stream-drop investigation) — temporary. Measure
@@ -1354,7 +1589,7 @@ export class AiChatService implements OnModuleInit {
const stepExhausted = steps.length >= MAX_AGENT_STEPS;
const emptyTurnMarker =
!producedText && stepExhausted ? STEP_LIMIT_NO_ANSWER_MARKER : '';
await finalizeAssistant(
const msgOk = await finalizeAssistant(
flushAssistant(steps as StepLike[], emptyTurnMarker, 'completed', {
finishReason: finishReason as string,
usage: totalUsage as StreamUsage,
@@ -1368,9 +1603,19 @@ export class AiChatService implements OnModuleInit {
pageChanged,
}),
);
// #184: settle the RUN as succeeded (best-effort, after the projection
// is finalized above).
if (runId) await runHooks?.onSettled?.(runId, 'completed');
// #184/#487: the RUN is finalized ALWAYS (never gated on the message).
// If the message finalize GAVE UP, error-mark the run so the asymmetry
// "run succeeded / message streaming forever" cannot arise; the
// periodic reconcile then settles the stuck message from this run.
if (runId) {
await runHooks?.onSettled?.(
runId,
msgOk ? 'completed' : 'error',
msgOk
? undefined
: 'Assistant message could not be persisted (finalize failed).',
);
}
// Lifecycle: release the external MCP clients leased for this turn.
await closeExternalClients();
@@ -2111,7 +2356,10 @@ export function planFinalizeAssistant(
* a test mock both satisfy it). */
export interface FinalizeRepo {
insert(insertable: Record<string, unknown>): Promise<unknown>;
update(
// #487: the OWNER terminal write is CONDITIONAL (status='streaming' OR
// metadata.finalizeFailed) so the owner overwrites a reconcile stamp but never
// an already-proper terminal row (owner-write priority).
finalizeOwner(
id: string,
workspaceId: string,
patch: AssistantFlush,
@@ -2120,10 +2368,11 @@ export interface FinalizeRepo {
/**
* Apply a finalize `plan` to the repo with the terminal `flushed` payload (#183):
* UPDATE the upfront row, or INSERT a fresh terminal row as the fallback when the
* upfront insert failed. The SINGLE dispatch shared by the service's
* finalizeAssistant and its test, so the test exercises the real path instead of
* a copy (#186 review). Pure of error handling the caller wraps it.
* conditionally UPDATE the upfront row (owner-write priority, #487), or INSERT a
* fresh terminal row as the fallback when the upfront insert failed. The SINGLE
* dispatch shared by the service's finalizeAssistant and its test, so the test
* exercises the real path instead of a copy (#186 review). Pure of error
* handling the caller wraps it (and RETRIES it, #487).
*/
export async function applyFinalize(
repo: FinalizeRepo,
@@ -2132,7 +2381,7 @@ export async function applyFinalize(
flushed: AssistantFlush,
): Promise<void> {
if (plan.kind === 'update') {
await repo.update(plan.id, base.workspaceId, flushed);
await repo.finalizeOwner(plan.id, base.workspaceId, flushed);
return;
}
await repo.insert({
@@ -0,0 +1,160 @@
import {
wrapInAppToolWithCap,
inAppToolCallCapMs,
type ToolAbortSignalSink,
} from './ai-chat-tools.service';
import type { Tool, ToolCallOptions } from 'ai';
/**
* #487 commit 1 in-app tool race-on-abort + safe-points + per-call cap.
*
* Tests assert the HONEST observable property the spec names "after Stop, NO
* new HTTP/WS call STARTS; an already-started single call may take either
* outcome" against the REAL wrapper mechanism (the composite abort signal it
* publishes on the client + the RACE it runs), NOT a timing-dependent proxy like
* "the write didn't land".
*/
// A minimal stand-in for the client's `toolAbortSignal` field. In production the
// wrapper publishes the composite here and the client's paginateAll /
// mutatePageContent safe-points read it; the fake "tool" below reads it the same
// way, so this exercises the real contract without a live DB / collab socket.
class FakeClient implements ToolAbortSignalSink {
private signal: AbortSignal | null = null;
setToolAbortSignal(signal: AbortSignal | null): void {
this.signal = signal;
}
getToolAbortSignal(): AbortSignal | null {
return this.signal;
}
}
// A ToolCallOptions with just the field the wrapper reads (abortSignal). The AI
// SDK passes a fuller object; the wrapper only spreads it and reads abortSignal.
const opts = (abortSignal?: AbortSignal): ToolCallOptions =>
({ toolCallId: 't1', messages: [], abortSignal }) as unknown as ToolCallOptions;
const tick = (ms = 5) => new Promise((r) => setTimeout(r, ms));
describe('#487 wrapInAppToolWithCap — race-on-abort + safe-points', () => {
it('after Stop, no NEW simulated call starts (multi-call tool)', async () => {
const client = new FakeClient();
const started: number[] = [];
// A multi-call tool that mirrors paginateAll: it consults the client signal
// at a safe-point BEFORE starting each simulated network call.
const multiCall: Tool = {
execute: (async (_args: unknown) => {
for (let i = 0; i < 6; i++) {
// Safe-point: exactly what paginateAll / mutatePageContent do.
client.getToolAbortSignal()?.throwIfAborted();
started.push(i);
await tick(10);
}
return 'done';
}) as unknown as Tool['execute'],
} as Tool;
const wrapped = wrapInAppToolWithCap(multiCall, client, 10_000);
const ac = new AbortController();
const call = (
wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise<unknown>
)({}, opts(ac.signal));
// Let one or two calls start, then Stop.
await tick(12);
ac.abort(new Error('user stop'));
await expect(call).rejects.toThrow(); // wrapper rejects promptly
const startedAtStop = started.length;
// Give the abandoned loser ample time; its next safe-point must throw because
// the (aborted) composite is still published on the client.
await tick(60);
expect(started.length).toBe(startedAtStop);
// It must NOT have run the whole sequence (that would mean Stop did nothing).
expect(started.length).toBeLessThan(6);
});
it('rejects immediately on Stop even if the call never settles (discard loser)', async () => {
const client = new FakeClient();
let settled = false;
const hang: Tool = {
execute: (async () => {
await new Promise(() => undefined); // never resolves
settled = true;
}) as unknown as Tool['execute'],
} as Tool;
const wrapped = wrapInAppToolWithCap(hang, client, 10_000);
const ac = new AbortController();
const call = (
wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise<unknown>
)({}, opts(ac.signal));
await tick(5);
ac.abort();
await expect(call).rejects.toThrow();
expect(settled).toBe(false);
});
it('per-call cap rejects a hung call with no Stop signal', async () => {
const client = new FakeClient();
const hang: Tool = {
execute: (async () => {
await new Promise(() => undefined);
}) as unknown as Tool['execute'],
} as Tool;
// Tiny cap; no options.abortSignal at all (composite = cap only).
const wrapped = wrapInAppToolWithCap(hang, client, 20);
const start = Date.now();
await expect(
(wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise<unknown>)(
{},
opts(undefined),
),
).rejects.toThrow(/per-call cap/);
expect(Date.now() - start).toBeLessThan(2000);
});
it('publishes a composite signal on the client for the duration of the call', async () => {
const client = new FakeClient();
let seenDuringCall: AbortSignal | null = null;
const probe: Tool = {
execute: (async () => {
seenDuringCall = client.getToolAbortSignal();
return 'ok';
}) as unknown as Tool['execute'],
} as Tool;
const wrapped = wrapInAppToolWithCap(probe, client, 10_000);
const ac = new AbortController();
await (
wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise<unknown>
)({}, opts(ac.signal));
expect(seenDuringCall).not.toBeNull();
// The published composite must reflect the turn's Stop signal.
ac.abort();
expect((seenDuringCall as unknown as AbortSignal).aborted).toBe(true);
});
it('a completed call returns its raw result unchanged', async () => {
const client = new FakeClient();
const ok: Tool = {
execute: (async () => ({ items: [1, 2, 3] })) as unknown as Tool['execute'],
} as Tool;
const wrapped = wrapInAppToolWithCap(ok, client, 10_000);
const res = await (
wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise<unknown>
)({}, opts(new AbortController().signal));
expect(res).toEqual({ items: [1, 2, 3] });
});
it('cap is env-tunable with a 2-minute default', () => {
const prev = process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS;
delete process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS;
expect(inAppToolCallCapMs()).toBe(120_000);
process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS = '5000';
expect(inAppToolCallCapMs()).toBe(5000);
process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS = 'not-a-number';
expect(inAppToolCallCapMs()).toBe(120_000);
if (prev === undefined) delete process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS;
else process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS = prev;
});
});
@@ -1,5 +1,5 @@
import { Injectable, Logger } from '@nestjs/common';
import { tool, type Tool } from 'ai';
import { tool, type Tool, type ToolCallOptions } from 'ai';
import { z } from 'zod';
import { User } from '@docmost/db/types/entity.types';
import { TokenService } from '../../auth/services/token.service';
@@ -159,6 +159,129 @@ function __assertClientCallContract(client: DocmostClientLike): void {
* existing service-account `/mcp` path already calls loopback successfully, so
* this works for single-workspace self-host.
*/
/**
* #487: wall-clock cap for a SINGLE in-app tool call, env-tunable via
* `AI_CHAT_INAPP_TOOL_CALL_CAP_MS`. Bounds a read tool that would otherwise
* paginate for minutes and a content write whose collab commit hangs, and is the
* per-call CAP half of the composite abort signal every in-app tool is wrapped
* with (the other half is the turn's Stop signal). Default 2 minutes: generous
* for a legitimate long read/write, tight enough that a stuck call cannot pin the
* turn. The reconcile staleness floor (#487 commit 4) is derived as
* max(2 x this cap, 15 min), so keep this well under that.
*/
export function inAppToolCallCapMs(): number {
const raw = Number(process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS);
return Number.isFinite(raw) && raw > 0 ? raw : 120_000;
}
/** #487: the composite signal's reason as an Error (informative thrown value). */
function inAppAbortReason(signal: AbortSignal): Error {
const r = signal.reason;
return r instanceof Error
? r
: new Error(typeof r === 'string' ? r : 'In-app tool call aborted');
}
/**
* The client surface {@link wrapInAppToolWithCap} drives (#487). Both methods are
* OPTIONAL: the real loopback DocmostClient implements them (so a Stop/cap reaches
* its pagination / pre-commit safe-points), but a client that omits them still
* gets the OUTER guarantee the race rejects on abort regardless. This keeps the
* wrapper decoupled from the exact client shape (unit-test doubles need not stub
* the plumbing).
*/
export interface ToolAbortSignalSink {
setToolAbortSignal?(signal: AbortSignal | null): void;
getToolAbortSignal?(): AbortSignal | null;
}
/**
* #487: wrap an in-app tool so a Stop (the turn's `options.abortSignal`) OR the
* per-call wall-clock cap REJECTS the call immediately, and so that SAME
* composite signal reaches the client's pagination / pre-commit safe-points (via
* `client.setToolAbortSignal`) making a Stop stop the NEXT HTTP/WS call from
* starting.
*
* Reuses the RACE pattern of `wrapToolWithCallTimeout` (mcp-clients.service.ts):
* the call is raced against the composite signal, so on abort we reject in the
* SAME tick and DISCARD the loser promise. Its network / collab teardown latency
* therefore never blocks the turn the supersede timeout W=10s (#487 commit 3)
* relies on this abort->settle latency being milliseconds, not a socket teardown.
* Awaiting the client's own signal-into-write path alone would NOT satisfy this
* (the loser could still be tearing down a collab socket).
*
* The composite is SET on the client at entry and deliberately NOT restored on
* unwind: after this wrapper rejects on abort, the ABANDONED loser promise keeps
* running, and its safe-points read the client field leaving the (aborted)
* composite there is exactly what makes the loser's NEXT call throw and stop. The
* next in-app tool call overwrites the field with its own fresh composite before
* any of its safe-points run, so a stale settled signal never leaks forward.
* SINGLE-WRITER by phase-1 assumption see DocmostClientContext.toolAbortSignal
* for the parallel-call caveat (#487).
*
* KNOWN LIMITATION (#487): a write tool that issues SEVERAL sequential collab
* commits can be aborted BETWEEN commits, leaving a partially-applied operation.
* Cancel guarantees "no NEW call starts", NOT "the write didn't land".
*/
export function wrapInAppToolWithCap(
toolDef: Tool,
client: ToolAbortSignalSink,
capMs: number,
): Tool {
const original = toolDef.execute;
if (typeof original !== 'function') return toolDef;
const execute = async (args: unknown, options: ToolCallOptions) => {
const capController = new AbortController();
const timer = setTimeout(() => {
capController.abort(
new Error(`In-app tool call exceeded the ${capMs}ms per-call cap`),
);
}, capMs);
timer.unref?.();
const composite = options?.abortSignal
? AbortSignal.any([options.abortSignal, capController.signal])
: capController.signal;
// Reject the MOMENT the composite fires, independent of whether `original`
// ever settles (a hung collab write / read would otherwise pin the turn). The
// losing `original` is left pending; Promise.race attaches a rejection
// handler to both inputs so a late rejection is never unhandled.
const aborted = new Promise<never>((_, reject) => {
const fail = () => reject(inAppAbortReason(composite));
if (composite.aborted) fail();
else composite.addEventListener('abort', fail, { once: true });
});
// Publish the composite so the client's pagination / pre-commit safe-points
// observe it (see the "not restored on unwind" rationale above). Guarded: a
// client without the plumbing still gets the OUTER race guarantee below.
client.setToolAbortSignal?.(composite);
try {
return await Promise.race([
(original as (a: unknown, o: ToolCallOptions) => Promise<unknown>)(
args,
{ ...options, abortSignal: composite },
),
aborted,
]);
} finally {
clearTimeout(timer);
}
};
return { ...toolDef, execute } as unknown as Tool;
}
/** #487: apply {@link wrapInAppToolWithCap} to every tool in a set. */
export function wrapInAppToolsWithCap(
tools: Record<string, Tool>,
client: ToolAbortSignalSink,
capMs: number,
): Record<string, Tool> {
const out: Record<string, Tool> = {};
for (const [name, t] of Object.entries(tools)) {
out[name] = wrapInAppToolWithCap(t, client, capMs);
}
return out;
}
@Injectable()
export class AiChatToolsService {
private readonly logger = new Logger(AiChatToolsService.name);
@@ -186,7 +309,12 @@ export class AiChatToolsService {
sessionId: string,
workspaceId: string,
aiChatId: string,
): Promise<DocmostClientLike> {
// #487: the returned client also carries the tool-cancellation plumbing
// (setToolAbortSignal/getToolAbortSignal). These are host plumbing, NOT part
// of the tool-execute surface (DocmostClientMethod), so they are surfaced here
// as an intersection rather than by widening that Pick — keeping the
// positional-call drift-guard (#446) scoped to the actual tool methods.
): Promise<DocmostClientLike & ToolAbortSignalSink> {
const apiUrl =
process.env.MCP_DOCMOST_API_URL ||
`http://127.0.0.1:${process.env.PORT || 3000}/api`;
@@ -630,7 +758,15 @@ export class AiChatToolsService {
// dependency and reuses the CASL enforcement already on `client`. When the
// loaded package predates #417 (factory undefined) or the loader is mocked in
// a unit test, signalling is a pure no-op and results are byte-identical.
if (!createCommentSignalTracker) return tools;
// #487: wrap every in-app tool with the race-on-abort + per-call cap guard so
// a Stop / cap rejects immediately AND reaches the client's write/pagination
// safe-points. Applied as the OUTERMOST wrapper (over the comment-signal
// wrapper below) so the race governs the whole call. The client carries the
// per-call composite signal via setToolAbortSignal.
const capMs = inAppToolCallCapMs();
if (!createCommentSignalTracker) {
return wrapInAppToolsWithCap(tools, client, capMs);
}
const tracker = createCommentSignalTracker({
probe: async (pageId: string, sinceMs: number) => {
@@ -659,7 +795,11 @@ export class AiChatToolsService {
},
});
return wrapToolsWithCommentSignal(tools, tracker);
return wrapInAppToolsWithCap(
wrapToolsWithCommentSignal(tools, tracker),
client,
capMs,
);
}
}
@@ -1,5 +1,6 @@
import { Injectable, Logger } from '@nestjs/common';
import { InjectKysely } from 'nestjs-kysely';
import { sql } from 'kysely';
import { KyselyDB, KyselyTransaction } from '../../types/kysely.types';
import { dbOrTx } from '../../utils';
import {
@@ -188,6 +189,144 @@ export class AiChatMessageRepo {
return query.returning(this.baseFields).executeTakeFirst();
}
/**
* #487 OWNER terminal write the streamText terminal callback's finalize. Like
* `update` but CONDITIONAL on `status='streaming' OR metadata.finalizeFailed`:
* the owner writes its real content EITHER when the row is still streaming (the
* normal case) OR when a reconcile stamp already flipped it to a terminal status
* but marked `finalizeFailed:true` the owner's real content OVERWRITES that
* placeholder stamp (owner-write priority, #487). A row that is properly terminal
* (no finalizeFailed) is left untouched (undefined) idempotent. The `patch`
* carries the real metadata WITHOUT finalizeFailed, so a successful write CLEARS
* the flag. Returns the updated row, or undefined when nothing matched.
*/
async finalizeOwner(
id: string,
workspaceId: string,
patch: Partial<{
content: string | null;
toolCalls: unknown;
metadata: unknown;
status: string | null;
}>,
trx?: KyselyTransaction,
): Promise<AiChatMessage | undefined> {
const db = dbOrTx(this.db, trx);
return db
.updateTable('aiChatMessages')
.set({ ...(patch as Record<string, unknown>), updatedAt: new Date() })
.where('id', '=', id)
.where('workspaceId', '=', workspaceId)
.where((eb) =>
eb.or([
eb('status', '=', 'streaming'),
eb(sql<string>`(metadata->>'finalizeFailed')`, '=', 'true'),
]),
)
.returning(this.baseFields)
.executeTakeFirst();
}
/**
* #487 RECONCILE status-only stamp settle a stuck 'streaming' row to a
* terminal status WITHOUT the owner's real content (which lived only in the
* dead process's memory — a documented loss). CONDITIONAL on `status='streaming'`
* (never touches an already-terminal row) AND it MERGES `finalizeFailed:true`
* into metadata (preserving the partial `parts` already persisted) so a LATER
* owner-write (finalizeOwner) can still OVERWRITE this placeholder with real
* content, and so `isInterruptResume` can EXCLUDE this row (a reconcile stamp is
* not a genuine user interruption). Returns the updated row, or undefined.
*/
async stampTerminalIfStreaming(
id: string,
workspaceId: string,
status: 'aborted' | 'error' | 'completed',
trx?: KyselyTransaction,
): Promise<AiChatMessage | undefined> {
const db = dbOrTx(this.db, trx);
return db
.updateTable('aiChatMessages')
.set({
status,
metadata: sql`coalesce(metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`,
updatedAt: new Date(),
})
.where('id', '=', id)
.where('workspaceId', '=', workspaceId)
.where('status', '=', 'streaming')
.returning(this.baseFields)
.executeTakeFirst();
}
/**
* #487 reconcile clause (b): streaming assistant rows whose linked RUN has
* already reached a terminal status an asymmetry ("run settled / message
* streaming forever") the periodic reconcile heals by stamping the message.
* Returns the message id + its run's terminal status, bounded.
*/
async findStreamingWithTerminalRun(
limit = 200,
// #487: scope to ONE chat for the opportunistic per-turn reconcile (removes
// reconcile latency from the user-visible path); omit for the periodic sweep.
chat?: { chatId: string; workspaceId: string },
): Promise<
Array<{ messageId: string; workspaceId: string; runStatus: string }>
> {
let query = this.db
.selectFrom('aiChatMessages as m')
.innerJoin('aiChatRuns as r', 'r.assistantMessageId', 'm.id')
.select([
'm.id as messageId',
'm.workspaceId as workspaceId',
'r.status as runStatus',
])
.where('m.status', '=', 'streaming')
.where('r.status', 'in', ['succeeded', 'failed', 'aborted']);
if (chat) {
query = query
.where('m.chatId', '=', chat.chatId)
.where('m.workspaceId', '=', chat.workspaceId);
}
return query.limit(limit).execute();
}
/**
* #487 reconcile clause (d) historical-row safety: streaming rows older than
* `staleMs` whose chat has NO active run row (double-gated). Settle them to
* 'aborted' + finalizeFailed (so a late owner-write could still overwrite).
* Returns the count. Used ONLY by the periodic reconcile, never at boot.
*/
async sweepStreamingWithoutActiveRun(
staleMs: number,
trx?: KyselyTransaction,
): Promise<number> {
const db = dbOrTx(this.db, trx);
const staleBefore = new Date(Date.now() - staleMs);
const rows = await db
.updateTable('aiChatMessages as m')
.set({
status: 'aborted',
metadata: sql`coalesce(m.metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`,
updatedAt: new Date(),
})
.where('m.status', '=', 'streaming')
.where('m.updatedAt', '<', staleBefore)
.where((eb) =>
eb.not(
eb.exists(
eb
.selectFrom('aiChatRuns as r')
.select('r.id')
.whereRef('r.chatId', '=', 'm.chatId')
.where('r.status', 'in', ['pending', 'running']),
),
),
)
.returning('m.id')
.execute();
return rows.length;
}
/**
* Crash-recovery sweep (#183): flip every assistant row still left in the
* 'streaming' state (a turn that died mid-write before reaching a terminal
@@ -200,13 +339,20 @@ export class AiChatMessageRepo {
* step, so an actively-streaming row never matches; this prevents a fresh
* replica's boot-sweep from aborting a turn another replica is still streaming
* in a multi-instance deploy.
*
* #487: the sweep now ALSO marks `finalizeFailed:true` so a late owner-write can
* overwrite this placeholder with real content (owner-write priority).
*/
async sweepStreaming(trx?: KyselyTransaction): Promise<number> {
const db = dbOrTx(this.db, trx);
const staleBefore = new Date(Date.now() - SWEEP_STREAMING_STALE_MS);
const rows = await db
.updateTable('aiChatMessages')
.set({ status: 'aborted', updatedAt: new Date() })
.set({
status: 'aborted',
metadata: sql`coalesce(metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`,
updatedAt: new Date(),
})
.where('status', '=', 'streaming')
.where('updatedAt', '<', staleBefore)
.returning('id')
@@ -143,6 +143,41 @@ export class AiChatRunRepo {
.executeTakeFirst();
}
/**
* #487: CONDITIONAL terminal finalize flip a run to a terminal status and
* stamp `finished_at` ONLY while it is still active (pending|running), mirroring
* the assistant message's `onlyIfStreaming` guard. A double-settle (a late or
* second writer, a supersede applying a zombie's intended, a reconcile stamp)
* matches NOTHING once the row is terminal and is a benign no-op so a terminal
* status can never be clobbered by a later writer (last-writer-wins is gone).
*
* Returns the updated row when it WAS active (this call wrote it), else
* undefined (the row was already terminal another writer won). The caller
* distinguishes the two to resolve the correct settle outcome.
*/
async finalizeIfActive(
id: string,
workspaceId: string,
patch: { status: string; error: string | null },
trx?: KyselyTransaction,
): Promise<AiChatRun | undefined> {
const db = dbOrTx(this.db, trx);
const now = new Date();
return db
.updateTable('aiChatRuns')
.set({
status: patch.status,
error: patch.error,
finishedAt: now,
updatedAt: now,
})
.where('id', '=', id)
.where('workspaceId', '=', workspaceId)
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[])
.returning(this.baseFields)
.executeTakeFirst();
}
/**
* Mark an EXPLICIT stop request on an active run (distinct from a browser
* disconnect, which never stops a run). Stamps `stop_requested_at` ONLY while
@@ -184,6 +219,31 @@ export class AiChatRunRepo {
* sweeps only runs UNTOUCHED past the window. Phase 1 is single-process, so the
* boot path supplies no window.
*/
/**
* #487 reconcile clause (c): active (pending|running) runs UNTOUCHED past
* `staleMs` candidates for "no live runner" abort. Staleness is measured from
* `updated_at` (the LAST-PROGRESS timestamp recordStep bumps it), NOT
* `started_at`, so a legitimate long-running marathon (1125 min of steady
* progress) is never a candidate. The caller filters these against its in-memory
* `active` / zombie maps ("no entry" is the PRIMARY gate a live entry is never
* aborted) before settling any of them. Bounded.
*/
async findStaleActive(
staleMs: number,
limit = 200,
trx?: KyselyTransaction,
): Promise<Array<{ id: string; workspaceId: string; chatId: string }>> {
const db = dbOrTx(this.db, trx);
const staleBefore = new Date(Date.now() - staleMs);
return db
.selectFrom('aiChatRuns')
.select(['id', 'workspaceId', 'chatId'])
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[])
.where('updatedAt', '<', staleBefore)
.limit(limit)
.execute();
}
async sweepRunning(
opts: { staleMs?: number } = {},
trx?: KyselyTransaction,
@@ -0,0 +1,305 @@
import { Kysely } from 'kysely';
import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
import { AiChatRunRepo } from '@docmost/db/repos/ai-chat/ai-chat-run.repo';
import { AiChatRunService } from '../../src/core/ai-chat/ai-chat-run.service';
import {
getTestDb,
destroyTestDb,
createWorkspace,
createUser,
createChat,
createMessage,
} from './db';
/**
* #487 commit 4 bidirectional reconcile + owner-write priority, real SQL.
*
* Proves the OBSERVABLE recovery properties against docmost_test:
* - the CONDITIONAL owner-write beats a reconcile stamp, and a stamp never
* clobbers a proper terminal row;
* - a LATE owner-finalize with real content OVERWRITES a reconcile 'aborted'
* stamp (finalizeFailed);
* - each reconcile clause (b message<-run, c stale-run, d historical row) settles
* the stuck row/run, and a LIVE run entry is never touched;
* - the "kill DB on finish" recovery: after the DB comes back, neither the
* message row nor the run row stays stuck.
*/
describe('#487 reconcile + owner-write priority [integration]', () => {
let db: Kysely<any>;
let messageRepo: AiChatMessageRepo;
let runRepo: AiChatRunRepo;
let runService: AiChatRunService;
let workspaceId: string;
let userId: string;
beforeAll(async () => {
db = getTestDb();
messageRepo = new AiChatMessageRepo(db as any);
runRepo = new AiChatRunRepo(db as any);
runService = new AiChatRunService(runRepo, { isCloud: () => false } as never);
workspaceId = (await createWorkspace(db)).id;
userId = (await createUser(db, workspaceId)).id;
});
afterAll(async () => {
await destroyTestDb();
});
const newChat = async () =>
(await createChat(db, { workspaceId, creatorId: userId })).id;
const metaOf = async (id: string): Promise<Record<string, unknown> | null> => {
const row = await messageRepo.findById(id, workspaceId);
return (row?.metadata as Record<string, unknown> | null) ?? null;
};
it('owner finalizeOwner writes a streaming row and CLEARS finalizeFailed', async () => {
const chatId = await newChat();
const m = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
status: 'streaming',
metadata: { parts: [] },
});
const wrote = await messageRepo.finalizeOwner(m.id, workspaceId, {
content: 'final answer',
status: 'completed',
metadata: { parts: [{ type: 'text', text: 'final answer' }] },
} as never);
expect(wrote!.status).toBe('completed');
expect((await metaOf(m.id))?.finalizeFailed).toBeUndefined();
});
it('a reconcile stamp NEVER clobbers a proper terminal row (finalizeOwner is a no-op there)', async () => {
const chatId = await newChat();
const m = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
status: 'completed',
content: 'real',
metadata: { parts: [] },
});
// The reconcile stamp is onlyIfStreaming -> no-op on a completed row.
const stamped = await messageRepo.stampTerminalIfStreaming(
m.id,
workspaceId,
'aborted',
);
expect(stamped).toBeUndefined();
expect((await messageRepo.findById(m.id, workspaceId))!.status).toBe(
'completed',
);
});
it('LATE owner-finalize with real content OVERWRITES a reconcile aborted stamp', async () => {
const chatId = await newChat();
const m = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
status: 'streaming',
metadata: { parts: [{ type: 'text', text: 'partial' }] },
});
// Reconcile stamps it aborted + finalizeFailed (final text lived only in mem).
const stamped = await messageRepo.stampTerminalIfStreaming(
m.id,
workspaceId,
'aborted',
);
expect(stamped!.status).toBe('aborted');
expect((await metaOf(m.id))?.finalizeFailed).toBe(true);
// A LATE owner-write (finalizeFailed=true satisfies the OR) overwrites it with
// real content, clearing the flag — owner-write priority.
const wrote = await messageRepo.finalizeOwner(m.id, workspaceId, {
content: 'the real final answer',
status: 'completed',
metadata: { parts: [{ type: 'text', text: 'the real final answer' }] },
} as never);
expect(wrote!.status).toBe('completed');
expect(wrote!.content).toBe('the real final answer');
expect((await metaOf(m.id))?.finalizeFailed).toBeUndefined();
});
it('clause (c): a stale active run with NO live entry -> aborted; a LIVE entry is untouched', async () => {
// Stale run, NOT owned by this replica (no entry) -> reconcile aborts it.
const staleChat = await newChat();
const stale = await runRepo.insert({
chatId: staleChat,
workspaceId,
createdBy: userId,
status: 'running',
});
await db
.updateTable('aiChatRuns')
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
.where('id', '=', stale.id)
.execute();
// A live run OWNED by this replica (beginRun registers an in-memory entry),
// ALSO backdated stale — the "no entry" primary gate must protect it.
const liveChat = await newChat();
const live = await runService.beginRun({
chatId: liveChat,
workspaceId,
userId,
});
await db
.updateTable('aiChatRuns')
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
.where('id', '=', live.runId)
.execute();
const aborted = await runService.reconcileStaleRuns(15 * 60 * 1000);
expect(aborted).toBeGreaterThanOrEqual(1);
expect((await runRepo.findById(stale.id, workspaceId))!.status).toBe(
'aborted',
);
// The live entry is NEVER aborted, however stale its row looks.
expect((await runRepo.findById(live.runId, workspaceId))!.status).toBe(
'running',
);
expect(runService.isLocallyActive(live.runId)).toBe(true);
// cleanup the live run
await runService.finalizeRun(live.runId, workspaceId, 'aborted');
});
it('clause (b): a streaming message whose RUN is terminal is stamped by run status (succeeded -> aborted, NOT completed-empty)', async () => {
const chatId = await newChat();
const msg = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
status: 'streaming',
metadata: { parts: [] },
});
// A SUCCEEDED run linked to the still-streaming message (the asymmetry).
const run = await runRepo.insert({
chatId,
workspaceId,
createdBy: userId,
status: 'running',
assistantMessageId: msg.id,
});
await runRepo.finalizeIfActive(run.id, workspaceId, {
status: 'succeeded',
error: null,
});
const stuck = await messageRepo.findStreamingWithTerminalRun();
const mine = stuck.find((s) => s.messageId === msg.id);
expect(mine?.runStatus).toBe('succeeded');
// Reconcile clause (b): succeeded run -> message 'aborted' (NOT 'completed'),
// the final text lived only in memory (documented loss), +finalizeFailed.
const status = mine!.runStatus === 'failed' ? 'error' : 'aborted';
await messageRepo.stampTerminalIfStreaming(msg.id, workspaceId, status);
const row = await messageRepo.findById(msg.id, workspaceId);
expect(row!.status).toBe('aborted');
expect((row!.metadata as Record<string, unknown>).finalizeFailed).toBe(true);
});
it('clause (d): a stale streaming row with NO active run on the chat -> aborted+finalizeFailed', async () => {
const chatId = await newChat();
const msg = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
status: 'streaming',
metadata: { parts: [] },
});
await db
.updateTable('aiChatMessages')
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
.where('id', '=', msg.id)
.execute();
const swept = await messageRepo.sweepStreamingWithoutActiveRun(
15 * 60 * 1000,
);
expect(swept).toBeGreaterThanOrEqual(1);
const row = await messageRepo.findById(msg.id, workspaceId);
expect(row!.status).toBe('aborted');
expect((row!.metadata as Record<string, unknown>).finalizeFailed).toBe(true);
});
it('clause (d) is DOUBLE-GATED: a stale streaming row WITH an active run on the chat is left alone', async () => {
const chatId = await newChat();
const msg = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
status: 'streaming',
metadata: { parts: [] },
});
await db
.updateTable('aiChatMessages')
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
.where('id', '=', msg.id)
.execute();
// An ACTIVE run on the same chat -> clause (d) must NOT touch the message.
const run = await runRepo.insert({
chatId,
workspaceId,
createdBy: userId,
status: 'running',
});
await messageRepo.sweepStreamingWithoutActiveRun(15 * 60 * 1000);
expect((await messageRepo.findById(msg.id, workspaceId))!.status).toBe(
'streaming',
);
await runRepo.finalizeIfActive(run.id, workspaceId, {
status: 'aborted',
error: null,
});
});
it('"kill DB on finish" recovery: after the DB is back, reconcile leaves NEITHER the row nor the run stuck', async () => {
// Simulate a process that seeded the assistant row + run, then died before
// finalizing EITHER (a mid-turn crash): a streaming message + a running run,
// both stale, with no in-memory entry (fresh service = fresh maps).
const chatId = await newChat();
const msg = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
status: 'streaming',
metadata: { parts: [{ type: 'text', text: 'partial' }] },
});
const run = await runRepo.insert({
chatId,
workspaceId,
createdBy: userId,
status: 'running',
assistantMessageId: msg.id,
});
await db
.updateTable('aiChatRuns')
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
.where('id', '=', run.id)
.execute();
await db
.updateTable('aiChatMessages')
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
.where('id', '=', msg.id)
.execute();
// Reconcile (as the periodic job would): (c) aborts the orphan run, then
// (b) settles the message from the now-terminal run.
await runService.reconcileStaleRuns(15 * 60 * 1000);
const stuck = await messageRepo.findStreamingWithTerminalRun();
for (const s of stuck) {
const status = s.runStatus === 'failed' ? 'error' : 'aborted';
await messageRepo.stampTerminalIfStreaming(s.messageId, s.workspaceId, status);
}
// Neither is stuck: the run is terminal AND the message is terminal.
expect((await runRepo.findById(run.id, workspaceId))!.status).toBe('aborted');
const row = await messageRepo.findById(msg.id, workspaceId);
expect(row!.status).toBe('aborted');
expect((row!.metadata as Record<string, unknown>).finalizeFailed).toBe(true);
});
});
@@ -281,6 +281,52 @@ describe('AiChatRun durable lifecycle [integration]', () => {
});
});
it('#487 finalizeIfActive is CONDITIONAL: a late terminal write cannot clobber the settled status (real SQL)', async () => {
const c = (await createChat(db, { workspaceId, creatorId: userId })).id;
const run = await runRepo.insert({
chatId: c,
workspaceId,
createdBy: userId,
status: 'running',
});
// First terminal write: the run IS active, so it flips + returns the row.
const first = await runRepo.finalizeIfActive(run.id, workspaceId, {
status: 'succeeded',
error: null,
});
expect(first!.status).toBe('succeeded');
expect(first!.finishedAt).toBeTruthy();
// A late/second writer tries to flip it to 'aborted' — the WHERE status IN
// ('pending','running') guard matches NOTHING now, so it is a benign no-op.
const second = await runRepo.finalizeIfActive(run.id, workspaceId, {
status: 'aborted',
error: 'late clobber attempt',
});
expect(second).toBeUndefined();
// The persisted terminal status is UNCHANGED — last-writer-wins is gone.
const row = await runRepo.findById(run.id, workspaceId);
expect(row!.status).toBe('succeeded');
expect(row!.error).toBeNull();
});
it('#487 double-settle through the service collapses to one write at the SQL gate', async () => {
const c = (await createChat(db, { workspaceId, creatorId: userId })).id;
const handle = await service.beginRun({ chatId: c, workspaceId, userId });
// First settle writes 'aborted' via the conditional write.
await service.finalizeRun(handle.runId, workspaceId, 'aborted');
// A late safety-net settle to 'error' is a no-op (row already terminal).
await service.finalizeRun(handle.runId, workspaceId, 'error', 'late');
const row = await runRepo.findById(handle.runId, workspaceId);
expect(row!.status).toBe('aborted');
expect(service.isLocallyActive(handle.runId)).toBe(false);
expect(service.hasZombie(handle.runId)).toBe(false);
});
it('sweepRunning() with NO args (boot sweep / variant C) aborts even a FRESH running run', async () => {
// F1/DECISION C at the SQL level: the unconditional boot sweep has NO
// staleness window, so a run updated just now (a fast restart) is settled too
+63 -2
View File
@@ -162,6 +162,43 @@ export abstract class DocmostClientContext {
// bypassed on a forced refresh (the 401/403 reauth path). null = no token yet.
protected collabTokenCache: { token: string; mintedAt: number } | null = null;
// #487: an OPTIONAL abort signal the in-app tool host sets before each tool
// call (a composite of the turn's Stop signal + a per-call wall-clock cap). It
// is checked at safe-points BETWEEN the sequential HTTP calls of a paginated
// read (paginateAll) and just before the atomic collab commit of a write (the
// mutatePage/replacePage/mutateLiveContentUnlocked seams), so a Stop / cap
// stops the NEXT network call from STARTING. An already-started single call may
// still land — a documented limitation (#487).
//
// SINGLE-WRITER by phase-1 assumption: exactly one DocmostClient is built per
// turn and shared by every tool call; the host sets this per call and does NOT
// restore the prior value on unwind (set-and-leave) — a fresh client per turn
// plus overwrite-by-the-next-call keeps it correct, and leaving a settled
// call's signal in place is what makes a discarded race-loser throw on its
// next safe-point. If the model emits PARALLEL in-app
// tool calls they share this one field, so the per-call CAP of one call is not
// guaranteed to bound another's in-flight pagination — but every composite the
// host sets carries the SAME turn Stop signal, so a Stop still aborts whichever
// signal is current. #487.
protected toolAbortSignal: AbortSignal | null = null;
/**
* #487: set (or clear with null) the in-app tool abort signal governing the
* NEXT client call's safe-points. The host wraps each in-app tool call: it sets
* the composite (Stop + per-call cap) here before invoking the tool and leaves
* it in place afterwards (set-and-leave, NOT restored) the next call
* overwrites it, and a fresh client is built per turn. Public so the
* server-side tool wrapper can reach it; harmless (a no-op) when never set.
*/
public setToolAbortSignal(signal: AbortSignal | null): void {
this.toolAbortSignal = signal;
}
/** #487: the abort signal currently governing this client's safe-points. */
public getToolAbortSignal(): AbortSignal | null {
return this.toolAbortSignal;
}
// Two construction forms:
// - new DocmostClient(config) // discriminated union (current)
// - new DocmostClient(baseURL, email, password) // legacy positional creds
@@ -563,6 +600,10 @@ export abstract class DocmostClientContext {
this.onMetricFn?.("collab_connect_timeouts_total", 1),
});
try {
// #487 PRE-COMMIT safe-point (reentrant twin of mutatePageContent): a
// Stop/cap after acquiring the session but before the atomic write skips
// this commit. Same limitation applies (stops the NEXT commit only).
this.toolAbortSignal?.throwIfAborted();
return await session.mutate(transform);
} catch (e) {
// Drop the session on any failure so the next call reconnects fresh.
@@ -594,6 +635,11 @@ export abstract class DocmostClientContext {
let truncated = false;
for (let page = 0; page < MAX_PAGES; page++) {
// #487 safe-point: a Stop (or the in-app tool per-call cap) that fires
// BETWEEN sequential page fetches must stop the NEXT request from starting
// — a read tool that would otherwise paginate for minutes is interrupted
// here. throwIfAborted() rejects with the signal's reason.
this.toolAbortSignal?.throwIfAborted();
const payload: Record<string, any> = {
...basePayload,
limit: clampedLimit,
@@ -701,7 +747,15 @@ export abstract class DocmostClientContext {
// #486: on a rejected collab-WS handshake, invalidate + refresh the token and
// retry the write once (symmetric to the HTTP-401 reauth path).
return this.writeWithCollabAuthRetry(collabToken, (token) =>
mutatePageContent(pageUuid, token, apiUrl, transform),
// #487: thread the in-app tool signal to mutatePageContent's pre-commit
// safe-point so a Stop/cap during the connect/lock window skips the write.
mutatePageContent(
pageUuid,
token,
apiUrl,
transform,
this.toolAbortSignal ?? undefined,
),
);
}
@@ -725,7 +779,14 @@ export abstract class DocmostClientContext {
// #486: on a rejected collab-WS handshake, invalidate + refresh the token and
// retry the write once (symmetric to the HTTP-401 reauth path).
return this.writeWithCollabAuthRetry(collabToken, (token) =>
replacePageContent(pageUuid, doc, token, apiUrl),
// #487: same pre-commit safe-point as mutatePage, for full-document writes.
replacePageContent(
pageUuid,
doc,
token,
apiUrl,
this.toolAbortSignal ?? undefined,
),
);
}
+16
View File
@@ -254,6 +254,12 @@ export async function mutatePageContent(
collabToken: string,
baseUrl: string,
transform: (liveDoc: any) => any | null,
// #487: optional abort signal carrying the turn's Stop + the in-app tool
// per-call cap. Checked as the PRE-COMMIT safe-point below (after the session
// is acquired, immediately before the atomic read->write), so a Stop that
// arrives during the connect/lock window stops THIS write from landing. See the
// limitation note at the check.
signal?: AbortSignal,
): Promise<MutationResult> {
return withPageLock(pageId, async () => {
if (process.env.DEBUG) {
@@ -266,6 +272,13 @@ export async function mutatePageContent(
const session = await acquireCollabSession(pageId, collabToken, baseUrl);
try {
// #487 PRE-COMMIT safe-point: if the turn was Stopped (or the in-app tool
// per-call cap fired) after we acquired the collab session but before the
// atomic write, throw NOW so this commit never runs. KNOWN LIMITATION
// (#487): this only stops THIS commit — a write tool that already committed
// an EARLIER call this turn leaves that op applied. Cancel guarantees "no
// NEW commit starts", NOT "the write didn't land".
signal?.throwIfAborted();
return await session.mutate(transform);
} catch (e) {
// Drop the session on any failure so the next call reconnects fresh (this
@@ -291,6 +304,8 @@ export async function replacePageContent(
prosemirrorDoc: any,
collabToken: string,
baseUrl: string,
// #487: threaded straight to mutatePageContent's pre-commit safe-point.
signal?: AbortSignal,
): Promise<MutationResult> {
// Fail fast on a bad document instead of deferring the failure into the
// collaboration write (where TiptapTransformer.toYdoc(undefined) used to
@@ -307,6 +322,7 @@ export async function replacePageContent(
collabToken,
baseUrl,
() => prosemirrorDoc,
signal,
);
}
@@ -0,0 +1,143 @@
// #487 commit 1 — the in-app tool cancellation safe-point inside paginateAll.
//
// The in-app tool host sets a composite abort signal on the client
// (setToolAbortSignal) before each tool call; paginateAll checks it at a
// safe-point BEFORE every sequential page fetch, so a Stop that lands mid-read
// stops the NEXT HTTP request from STARTING (a read tool can no longer paginate
// for minutes past a Stop). This pins the HONEST observable property against the
// REAL client + a real HTTP server: "after Stop, no NEW request starts".
import { test, after } from "node:test";
import assert from "node:assert/strict";
import http from "node:http";
import { DocmostClient } from "../../build/client.js";
function readBody(req) {
return new Promise((resolve) => {
let raw = "";
req.on("data", (c) => (raw += c));
req.on("end", () => resolve(raw));
});
}
function sendJson(res, status, obj, extra = {}) {
res.writeHead(status, { "Content-Type": "application/json", ...extra });
res.end(JSON.stringify(obj));
}
const openServers = [];
async function spawn(handler) {
const server = await new Promise((resolve) => {
const s = http.createServer(handler);
s.listen(0, "127.0.0.1", () => resolve(s));
});
openServers.push(server);
const { port } = server.address();
return { baseURL: `http://127.0.0.1:${port}/api` };
}
after(async () => {
await Promise.all(openServers.map((s) => new Promise((r) => s.close(r))));
});
function handleLogin(req, res) {
if (req.url === "/api/auth/login") {
sendJson(res, 200, { success: true }, {
"Set-Cookie": "authToken=t; Path=/; HttpOnly",
});
return true;
}
return false;
}
// A Stop that lands DURING pagination: the server aborts the client signal as it
// serves page 1 (more pages remain). The loop's next safe-point must throw before
// the page-2 request is sent.
test("paginateAll stops the NEXT request when the signal aborts mid-pagination", async () => {
let requests = 0;
const ac = new AbortController();
const { baseURL } = await spawn(async (req, res) => {
await readBody(req);
if (handleLogin(req, res)) return;
if (req.url === "/api/spaces") {
requests++;
// Simulate a user Stop that lands while page 1 is in flight.
if (requests === 1) ac.abort(new Error("user stop"));
sendJson(res, 200, {
success: true,
data: {
items: [{ id: `p${requests}` }],
meta: { hasNextPage: true, nextCursor: `c${requests}` },
},
});
return;
}
sendJson(res, 404, {});
});
const client = new DocmostClient(baseURL, "user@example.com", "pw");
client.setToolAbortSignal(ac.signal);
await assert.rejects(
() => client.paginateAll("/spaces", {}),
/user stop/,
"the aborted safe-point rejects with the signal's reason",
);
assert.equal(requests, 1, "page 2 never started after the Stop");
});
// A Stop that is already in effect before the read starts: zero requests fire.
test("paginateAll starts no request when the signal is already aborted", async () => {
let requests = 0;
const { baseURL } = await spawn(async (req, res) => {
await readBody(req);
if (handleLogin(req, res)) return;
if (req.url === "/api/spaces") {
requests++;
sendJson(res, 200, {
success: true,
data: { items: [], meta: { hasNextPage: false, nextCursor: null } },
});
return;
}
sendJson(res, 404, {});
});
const client = new DocmostClient(baseURL, "user@example.com", "pw");
// Warm the auth so ensureAuthenticated does not itself POST after the abort.
await client.ensureAuthenticated();
const ac = new AbortController();
ac.abort(new Error("already stopped"));
client.setToolAbortSignal(ac.signal);
await assert.rejects(() => client.paginateAll("/spaces", {}), /already stopped/);
assert.equal(requests, 0, "no /spaces request started once already aborted");
});
// Without a tool signal (default), pagination is unaffected — the safe-point is a
// pure no-op, so pre-#487 behaviour is byte-identical.
test("paginateAll is unaffected when no tool signal is set", async () => {
let requests = 0;
const PAGES = {
"": { items: [{ id: "a" }], nextCursor: "c1" },
c1: { items: [{ id: "b" }], nextCursor: null },
};
const { baseURL } = await spawn(async (req, res) => {
const raw = await readBody(req);
if (handleLogin(req, res)) return;
if (req.url === "/api/spaces") {
requests++;
const body = JSON.parse(raw || "{}");
const page = PAGES[body.cursor ?? ""] ?? { items: [], nextCursor: null };
sendJson(res, 200, {
success: true,
data: {
items: page.items,
meta: { hasNextPage: page.nextCursor != null, nextCursor: page.nextCursor },
},
});
return;
}
sendJson(res, 404, {});
});
const client = new DocmostClient(baseURL, "user@example.com", "pw");
const all = await client.paginateAll("/spaces", {});
assert.equal(requests, 2, "both pages fetched with no signal set");
assert.deepEqual(all.map((p) => p.id), ["a", "b"]);
});