feat(ai-chat): durable detached agent runs — phase 1 (#184/#234)

Squashed for a clean rebase onto develop (was 19 commits; the reviewer approved
the net diff at fb246080). Detaches an agent run from the HTTP request/browser
window: a run is a first-class lifecycle object (ai_chat_runs), a browser
disconnect no longer kills it, a concurrent-run insert-gate prevents double runs,
and a reopened chat live-follows a still-running run via a polled observer merge.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
agent_coder
2026-07-04 23:20:21 +03:00
parent 5336f06d10
commit 68899a2c2e
37 changed files with 4535 additions and 605 deletions
@@ -31,6 +31,7 @@ import { FavoriteRepo } from '@docmost/db/repos/favorite/favorite.repo';
import { TemplateRepo } from '@docmost/db/repos/template/template.repo';
import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo';
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 { AiChatPageSnapshotRepo } from '@docmost/db/repos/ai-chat/ai-chat-page-snapshot.repo';
import { AiProviderCredentialsRepo } from '@docmost/db/repos/ai-chat/ai-provider-credentials.repo';
import { AiMcpServerRepo } from '@docmost/db/repos/ai-chat/ai-mcp-server.repo';
@@ -105,6 +106,7 @@ import { normalizePostgresUrl } from '../common/helpers';
TemplateRepo,
AiChatRepo,
AiChatMessageRepo,
AiChatRunRepo,
AiChatPageSnapshotRepo,
AiProviderCredentialsRepo,
AiMcpServerRepo,
@@ -139,6 +141,7 @@ import { normalizePostgresUrl } from '../common/helpers';
TemplateRepo,
AiChatRepo,
AiChatMessageRepo,
AiChatRunRepo,
AiChatPageSnapshotRepo,
AiProviderCredentialsRepo,
AiMcpServerRepo,
@@ -0,0 +1,106 @@
import { type Kysely, sql } from 'kysely';
/**
* `ai_chat_runs` — the agent RUN as a first-class, server-side lifecycle object
* (#184 phase 1: autonomous agent runs detached from the browser window).
*
* Until now an agent turn lived ONLY as long as the HTTP request was open
* (`res.hijack()` in ai-chat.controller.ts); a browser disconnect aborted it.
* This table makes a turn a persistent object the server owns: it is created
* when a run starts (inserted directly as 'running' in phase 1 — 'pending' is
* only this column's default + a reserved value, never written by code yet) and
* advances to succeeded|failed|aborted, surviving the subscriber (browser) going
* away when it settles. The DB is the source of
* truth — a later client reconnects/sees the result by reading this row plus the
* assistant message it projects (`assistant_message_id`).
*
* The assistant message row (#183 step-granular durability) is the PROJECTION of
* a run's output; this row is the run's LIFECYCLE. They are linked by
* `assistant_message_id` (SET NULL if the message is later pruned).
*
* `status` : 'pending' | 'running' | 'succeeded' | 'failed' | 'aborted'.
* `trigger` : 'user' | 'autostart' | 'schedule' | 'api' | 'continue' — only
* 'user' is produced in phase 1; the others are reserved for the
* autonomy triggers deferred to phase 2 so they need no later
* migration.
*
* ONE ACTIVE RUN PER CHAT is enforced by a partial unique index on `chat_id`
* WHERE status IN ('pending','running'): an autonomous run and a user run can
* never trample each other on the same chat. Settled runs (succeeded/failed/
* aborted) are excluded from the index so a chat can accumulate any number of
* historical runs.
*/
export async function up(db: Kysely<any>): Promise<void> {
await db.schema
.createTable('ai_chat_runs')
.ifNotExists()
.addColumn('id', 'uuid', (col) =>
col.primaryKey().defaultTo(sql`gen_uuid_v7()`),
)
.addColumn('chat_id', 'uuid', (col) =>
col.references('ai_chats.id').onDelete('cascade').notNull(),
)
.addColumn('workspace_id', 'uuid', (col) =>
col.references('workspaces.id').onDelete('cascade').notNull(),
)
// The human who triggered the run (audit). SET NULL on user deletion so the
// run history outlives its author; NULL is also the natural value for a
// future system/cron/api trigger with no human actor.
.addColumn('created_by', 'uuid', (col) =>
col.references('users.id').onDelete('set null'),
)
// The assistant message this run materializes (the #183 projection). SET NULL
// if that message row is later deleted; nullable because the run row is
// created a moment BEFORE the assistant row is seeded.
.addColumn('assistant_message_id', 'uuid', (col) =>
col.references('ai_chat_messages.id').onDelete('set null'),
)
.addColumn('trigger', 'varchar(20)', (col) =>
col.notNull().defaultTo('user'),
)
.addColumn('status', 'varchar(20)', (col) =>
col.notNull().defaultTo('pending'),
)
// Terminal error message for a failed run (provider/transport cause),
// mirroring the assistant message's metadata.error.
.addColumn('error', 'text', (col) => col)
// Number of agent steps finished so far (kept monotonic with the projection).
.addColumn('step_count', 'integer', (col) => col.notNull().defaultTo(0))
// Set when an EXPLICIT user stop is requested (distinct from a mere browser
// disconnect, which never stops a run). The runner aborts the turn and the
// run settles as 'aborted'.
.addColumn('stop_requested_at', 'timestamptz', (col) => col)
.addColumn('started_at', 'timestamptz', (col) => col)
.addColumn('finished_at', 'timestamptz', (col) => col)
.addColumn('created_at', 'timestamptz', (col) =>
col.notNull().defaultTo(sql`now()`),
)
.addColumn('updated_at', 'timestamptz', (col) =>
col.notNull().defaultTo(sql`now()`),
)
.execute();
// Reconnect / "latest run for this chat" reads hit chat_id first.
await db.schema
.createIndex('ai_chat_runs_chat_id_idx')
.ifNotExists()
.on('ai_chat_runs')
.column('chat_id')
.execute();
// One ACTIVE run per chat (advisory at the DB level): a second pending/running
// run on the same chat is rejected, so a user turn and an autonomous turn can
// never race on the same chat. Partial so settled runs do not collide.
await db.schema
.createIndex('ai_chat_runs_one_active_per_chat')
.ifNotExists()
.on('ai_chat_runs')
.column('chat_id')
.unique()
.where(sql.ref('status'), 'in', sql`('pending','running')`)
.execute();
}
export async function down(db: Kysely<any>): Promise<void> {
await db.schema.dropTable('ai_chat_runs').execute();
}
@@ -121,6 +121,23 @@ export class AiChatMessageRepo {
return rows.reverse();
}
/** Fetch a single message by id + workspace (e.g. a run's projection row for
* the #184 reconnect read). Returns undefined when nothing matches. */
async findById(
id: string,
workspaceId: string,
trx?: KyselyTransaction,
): Promise<AiChatMessage | undefined> {
const db = dbOrTx(this.db, trx);
return db
.selectFrom('aiChatMessages')
.select(this.baseFields)
.where('id', '=', id)
.where('workspaceId', '=', workspaceId)
.where('deletedAt', 'is', null)
.executeTakeFirst();
}
async insert(
insertable: InsertableAiChatMessage,
trx?: KyselyTransaction,
@@ -0,0 +1,82 @@
import { AiChatRunRepo, SWEEP_RUN_STALE_MS } from './ai-chat-run.repo';
import type { KyselyDB } from '../../types/kysely.types';
/**
* Unit coverage for AiChatRunRepo.sweepRunning over a chainable builder mock (no
* live DB). The F1 invariant under test (DECISION C): the BOOT sweep is
* UNCONDITIONAL — it adds NO `updatedAt <` predicate, so a fresh 'running' run
* (updatedAt = now) IS settled rather than skipped by a staleness window. The
* window is added ONLY when an explicit `staleMs` is supplied (the future phase-2
* multi-instance timer sweep). We assert the EXACT predicates the spec mandates.
*/
describe('AiChatRunRepo.sweepRunning', () => {
type Recorded = {
table?: string;
set?: Record<string, unknown>;
wheres: Array<[string, string, unknown]>;
returning?: string;
};
function makeDb(swept: Array<{ id: string }>): {
db: KyselyDB;
rec: Recorded;
} {
const rec: Recorded = { wheres: [] };
const builder: Record<string, unknown> = {};
builder.set = (v: Record<string, unknown>) => {
rec.set = v;
return builder;
};
builder.where = (col: string, op: string, val: unknown) => {
rec.wheres.push([col, op, val]);
return builder;
};
builder.returning = (col: string) => {
rec.returning = col;
return builder;
};
builder.execute = () => Promise.resolve(swept);
const db = {
updateTable: (table: string) => {
rec.table = table;
return builder;
},
} as unknown as KyselyDB;
return { db, rec };
}
it('F1: the boot sweep (no staleMs) is UNCONDITIONAL — only a status filter, NO updatedAt window', async () => {
const { db, rec } = makeDb([{ id: 'r1' }, { id: 'r2' }]);
const repo = new AiChatRunRepo(db);
const swept = await repo.sweepRunning();
expect(swept).toBe(2);
expect(rec.table).toBe('aiChatRuns');
// The status filter is always present...
expect(rec.wheres).toContainEqual([
'status',
'in',
expect.arrayContaining(['pending', 'running']),
]);
// ...but a fresh 'running' run (updatedAt = now) must NOT be skipped: no
// updatedAt predicate at all on the boot path.
expect(rec.wheres.some(([col]) => col === 'updatedAt')).toBe(false);
// It flips to 'aborted' and stamps finishedAt.
expect(rec.set).toEqual(
expect.objectContaining({ status: 'aborted', finishedAt: expect.any(Date) }),
);
});
it('phase-2 path: an explicit staleMs reintroduces the updatedAt window', async () => {
const { db, rec } = makeDb([]);
const repo = new AiChatRunRepo(db);
await repo.sweepRunning({ staleMs: SWEEP_RUN_STALE_MS });
const updatedAtWhere = rec.wheres.find(([col]) => col === 'updatedAt');
expect(updatedAtWhere).toBeDefined();
expect(updatedAtWhere![1]).toBe('<');
expect(updatedAtWhere![2]).toBeInstanceOf(Date);
});
});
@@ -0,0 +1,212 @@
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 {
AiChatRun,
InsertableAiChatRun,
} from '@docmost/db/types/entity.types';
// Statuses that count as "the run is still live" (an autonomous and a user run
// must never both be live on one chat — enforced by the partial unique index and
// checked here for friendly 409s before the insert races the constraint).
export const ACTIVE_RUN_STATUSES = ['pending', 'running'] as const;
// Crash-recovery sweep recency threshold (mirrors AiChatMessageRepo.sweepStreaming,
// #183): when a staleness window is supplied, a 'running'/'pending' run is only
// swept to 'aborted' once it has been UNTOUCHED for this long, so a sibling
// replica's boot-sweep can never abort a run another replica is actively
// executing. The runner bumps `updatedAt` on every step, so a live run never
// matches. PHASE 1 is single-process and the boot sweep passes NO window (every
// dangling run is settled unconditionally — see sweepRunning / F1). This constant
// is the window to reintroduce for the phase-2 multi-instance timer sweep.
export const SWEEP_RUN_STALE_MS = 10 * 60 * 1000; // 10 minutes
/**
* Repository for `ai_chat_runs` (#184 phase 1): the agent run as a first-class,
* server-side lifecycle object detached from the HTTP request. The run row is the
* point a client subscribes/reconnects to (by `id` or by chat); the assistant
* message it links to (`assistantMessageId`) is the #183 projection of its output.
*/
@Injectable()
export class AiChatRunRepo {
private readonly logger = new Logger(AiChatRunRepo.name);
private baseFields: Array<keyof AiChatRun> = [
'id',
'chatId',
'workspaceId',
'createdBy',
'assistantMessageId',
'trigger',
'status',
'error',
'stepCount',
'stopRequestedAt',
'startedAt',
'finishedAt',
'createdAt',
'updatedAt',
];
constructor(@InjectKysely() private readonly db: KyselyDB) {}
async insert(
insertable: InsertableAiChatRun,
trx?: KyselyTransaction,
): Promise<AiChatRun> {
const db = dbOrTx(this.db, trx);
return db
.insertInto('aiChatRuns')
.values(insertable)
.returning(this.baseFields)
.executeTakeFirst();
}
async findById(
id: string,
workspaceId: string,
trx?: KyselyTransaction,
): Promise<AiChatRun | undefined> {
const db = dbOrTx(this.db, trx);
return db
.selectFrom('aiChatRuns')
.select(this.baseFields)
.where('id', '=', id)
.where('workspaceId', '=', workspaceId)
.executeTakeFirst();
}
/** The currently-active (pending|running) run for a chat, if any. At most one
* exists thanks to the partial unique index. */
async findActiveByChat(
chatId: string,
workspaceId: string,
trx?: KyselyTransaction,
): Promise<AiChatRun | undefined> {
const db = dbOrTx(this.db, trx);
return db
.selectFrom('aiChatRuns')
.select(this.baseFields)
.where('chatId', '=', chatId)
.where('workspaceId', '=', workspaceId)
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[])
.executeTakeFirst();
}
/** The most-recent run for a chat (active or settled) — the reconnect target. */
async findLatestByChat(
chatId: string,
workspaceId: string,
trx?: KyselyTransaction,
): Promise<AiChatRun | undefined> {
const db = dbOrTx(this.db, trx);
return db
.selectFrom('aiChatRuns')
.select(this.baseFields)
.where('chatId', '=', chatId)
.where('workspaceId', '=', workspaceId)
.orderBy('createdAt', 'desc')
.orderBy('id', 'desc')
.limit(1)
.executeTakeFirst();
}
/**
* Patch a run by id + workspace; always bumps `updatedAt`. Used for every
* lifecycle transition (mark running, link the assistant message, bump
* step_count, finalize succeeded/failed/aborted). Returns the updated row or
* undefined when nothing matched (e.g. a foreign workspace).
*/
async update(
id: string,
workspaceId: string,
patch: Partial<{
status: string;
error: string | null;
stepCount: number;
assistantMessageId: string | null;
stopRequestedAt: Date | null;
startedAt: Date | null;
finishedAt: Date | null;
}>,
trx?: KyselyTransaction,
): Promise<AiChatRun | undefined> {
const db = dbOrTx(this.db, trx);
return db
.updateTable('aiChatRuns')
.set({ ...(patch as Record<string, unknown>), updatedAt: new Date() })
.where('id', '=', id)
.where('workspaceId', '=', workspaceId)
.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
* the run is still active, so a late stop on an already-settled run is a no-op.
* Returns the row when a stop was recorded, else undefined (nothing active).
*/
async markStopRequested(
id: string,
workspaceId: string,
trx?: KyselyTransaction,
): Promise<AiChatRun | undefined> {
const db = dbOrTx(this.db, trx);
return db
.updateTable('aiChatRuns')
.set({ stopRequestedAt: new Date(), updatedAt: new Date() })
.where('id', '=', id)
.where('workspaceId', '=', workspaceId)
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[])
.returning(this.baseFields)
.executeTakeFirst();
}
/**
* Crash-recovery sweep (mirrors AiChatMessageRepo.sweepStreaming): flip every
* run still left pending/running — a run whose process died before reaching a
* terminal status — to 'aborted', stamping `finished_at`. Returns the number
* swept. Workspace-wide on purpose (a crash can dangle runs in any workspace).
*
* F1 (DECISION C): the BOOT sweep is UNCONDITIONAL — it passes no `staleMs`, so
* EVERY dangling run is settled regardless of how recently it was touched. On a
* fresh single-process boot any pending|running run is definitionally hung (no
* runner is alive to own it), so a fast restart (deploy/OOM within minutes of
* the last step) no longer leaves a run stuck 'running' forever — which would
* make the one-active-run gate 409 every future turn in that chat.
*
* The optional `staleMs` window is reintroduced ONLY for the future phase-2
* multi-instance timer sweep (see {@link SWEEP_RUN_STALE_MS}): there a booting
* replica must NOT abort a run another replica is actively executing, so it
* sweeps only runs UNTOUCHED past the window. Phase 1 is single-process, so the
* boot path supplies no window.
*/
async sweepRunning(
opts: { staleMs?: number } = {},
trx?: KyselyTransaction,
): Promise<number> {
const db = dbOrTx(this.db, trx);
const now = new Date();
let query = db
.updateTable('aiChatRuns')
.set({
status: 'aborted',
finishedAt: now,
updatedAt: now,
error: sql`coalesce(error, ${'Run interrupted by a server restart.'})`,
})
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[]);
// Multi-instance (phase 2) only: skip runs touched within the window so a
// sibling replica's live run is never aborted. Omitted on the phase-1 boot
// sweep -> unconditional.
if (typeof opts.staleMs === 'number') {
const staleBefore = new Date(now.getTime() - opts.staleMs);
query = query.where('updatedAt', '<', staleBefore);
}
const rows = await query.returning('id').execute();
return rows.length;
}
}
+30
View File
@@ -647,6 +647,35 @@ export interface AiChatMessages {
deletedAt: Timestamp | null;
}
// The agent RUN as a first-class server-side lifecycle object (#184 phase 1).
// Mirrors migration 20260627T130000-ai-chat-runs.ts. A run is created when an
// agent turn starts and survives the browser disconnecting; the DB is the source
// of truth a later client reconnects to. `assistantMessageId` links to the #183
// projection row (the assistant message this run materializes).
export interface AiChatRuns {
id: Generated<string>;
chatId: string;
workspaceId: string;
// SET NULL on user deletion (the run history outlives its author); also NULL
// for a future non-human trigger (cron/api).
createdBy: string | null;
// The assistant message this run materializes; SET NULL if it is pruned.
assistantMessageId: string | null;
// 'user' | 'autostart' | 'schedule' | 'api' | 'continue' (only 'user' is
// produced in phase 1; the rest are reserved for the deferred autonomy triggers).
trigger: Generated<string>;
// 'pending' | 'running' | 'succeeded' | 'failed' | 'aborted'.
status: Generated<string>;
error: string | null;
stepCount: Generated<number>;
// Set when an EXPLICIT user stop is requested (distinct from a disconnect).
stopRequestedAt: Timestamp | null;
startedAt: Timestamp | null;
finishedAt: Timestamp | null;
createdAt: Generated<Timestamp>;
updatedAt: Generated<Timestamp>;
}
// Per-(chat,page) snapshot of the open page's Markdown at the END of the agent's
// previous turn (#274). Mirrors migration 20260702T120000-ai-chat-page-snapshot.ts.
// The next turn diffs the CURRENT Markdown against `contentMd` to surface edits a
@@ -683,6 +712,7 @@ export interface DB {
aiAgentRoles: AiAgentRoles;
aiChats: AiChats;
aiChatMessages: AiChatMessages;
aiChatRuns: AiChatRuns;
aiChatPageSnapshots: AiChatPageSnapshots;
apiKeys: ApiKeys;
attachments: Attachments;
+15 -7
View File
@@ -3,6 +3,7 @@ import {
AiAgentRoles,
AiChats,
AiChatMessages,
AiChatRuns,
AiChatPageSnapshots,
Attachments,
Comments,
@@ -56,10 +57,12 @@ export type UpdatableAiChat = Updateable<Omit<AiChats, 'id'>>;
// full-text search. It is omitted from the public type so it never leaks
// into HTTP responses or the chat history fed to the language model.
export type AiChatMessage = Omit<Selectable<AiChatMessages>, 'tsv'>;
export type InsertableAiChatMessage = Omit<
Insertable<AiChatMessages>,
'tsv'
>;
export type InsertableAiChatMessage = Omit<Insertable<AiChatMessages>, 'tsv'>;
// AI Chat Run (#184 phase 1): the agent run as a first-class lifecycle object,
// detached from the HTTP request / browser window.
export type AiChatRun = Selectable<AiChatRuns>;
export type InsertableAiChatRun = Insertable<AiChatRuns>;
// AI Chat Page Snapshot (#274): per-(chat,page) Markdown snapshot taken at the
// end of the agent's previous turn, diffed against the current page next turn to
@@ -214,11 +217,14 @@ export type UpdatableFavorite = Updateable<Omit<Favorites, 'id'>>;
// Page Transclusion
export type PageTransclusion = Selectable<PageTransclusions>;
export type InsertablePageTransclusion = Insertable<PageTransclusions>;
export type UpdatablePageTransclusion = Updateable<Omit<PageTransclusions, 'id'>>;
export type UpdatablePageTransclusion = Updateable<
Omit<PageTransclusions, 'id'>
>;
// Page Transclusion Reference
export type PageTransclusionReference = Selectable<PageTransclusionReferences>;
export type InsertablePageTransclusionReference = Insertable<PageTransclusionReferences>;
export type InsertablePageTransclusionReference =
Insertable<PageTransclusionReferences>;
export type UpdatablePageTransclusionReference = Updateable<
Omit<PageTransclusionReferences, 'id'>
>;
@@ -288,7 +294,9 @@ export type UpdatablePagePermission = Updateable<Omit<_PagePermissions, 'id'>>;
// Page Verification
export type PageVerification = Selectable<_PageVerifications>;
export type InsertablePageVerification = Insertable<_PageVerifications>;
export type UpdatablePageVerification = Updateable<Omit<_PageVerifications, 'id'>>;
export type UpdatablePageVerification = Updateable<
Omit<_PageVerifications, 'id'>
>;
// Page Verifier
export type PageVerifier = Selectable<_PageVerifiers>;