refactor(review): address PR #186 re-review (approve-with-comments)

Approve-with-comments re-review; no blockers. All 7 actionable points (8 is a
forward-looking architecture note — recommendation A, keep as-is):

1. chat-markdown.util spec: restore parity coverage of the removed client spec —
   tool error state (+ errorText), unknown-tool fallback (`Ran tool <name>` en /
   `Выполнил инструмент <name>` ru), and the circular-output stringify catch.
2. findAllByChat row cap is now testable (injectable limit) + an int-spec proves
   truncation on a modest volume.
3. Stability: the per-step durability updates are SERIALIZED via a promise chain
   (stepUpdateChain) so they commit in step order — onlyIfStreaming already
   closed the finalize race, this closes inter-step ordering.
4. findAllByChat keeps the NEWEST messages on truncation (order DESC + reverse,
   like findRecent) and logs a warning with chatId, instead of silently dropping
   the newest tail.
5. The LABELS parity comment already references the real path (tool-parts.tsx /
   toolLabelKey) — confirmed accurate.
6. Removed the redundant 'off-by-one boundary' test (strict subset of the two
   adjacent prepareAgentStep cases).
7. Extracted the terminal-finalize dispatch into a shared `applyFinalize`, used
   by BOTH the service's finalizeAssistant and its test — the test now exercises
   the real path, not a copy, so a production drift fails it.

Verified: server build + 325 ai-chat unit + 6 integration; prettier clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
claude code agent 227
2026-06-25 12:28:35 +03:00
parent ea61c96a7c
commit aa7a115f66
7 changed files with 212 additions and 62 deletions

View File

@@ -2,6 +2,7 @@ import { ForbiddenException } from '@nestjs/common';
import { AiChatController } from './ai-chat.controller';
import {
planFinalizeAssistant,
applyFinalize,
flushAssistant,
type AssistantFlush,
} from './ai-chat.service';
@@ -103,35 +104,25 @@ describe('AiChatController.export', () => {
* (assistantId is undefined) finalize falls back to INSERTing the terminal row
* so the turn is not lost — the only safety against losing the turn entirely.
*
* `planFinalizeAssistant` is the pure decision; this also drives a tiny harness
* that mirrors the service's `finalizeAssistant` repo dispatch over a mock repo,
* proving both branches issue the right call with the terminal payload.
* `planFinalizeAssistant` is the pure decision; `applyFinalize` is the REAL
* dispatch the service uses, exercised here over a mock repo (not a copy of the
* logic) so a production drift would fail the test (#186 review).
*/
describe('finalizeAssistant dispatch (planFinalizeAssistant)', () => {
describe('finalizeAssistant dispatch (planFinalizeAssistant + applyFinalize)', () => {
const workspaceId = 'ws1';
// Mirror of the service's finalize repo-dispatch over the plan: UPDATE when an
// upfront row exists, else INSERT the terminal row.
// Drive the SAME applyFinalize the service calls (no duplicated logic).
async function dispatchFinalize(
repo: { insert: jest.Mock; update: jest.Mock },
assistantId: string | undefined,
flushed: AssistantFlush,
): Promise<void> {
const plan = planFinalizeAssistant(assistantId);
if (plan.kind === 'insert') {
await repo.insert({
chatId: 'c1',
workspaceId,
userId: 'u1',
role: 'assistant',
content: flushed.content,
toolCalls: flushed.toolCalls ?? null,
metadata: flushed.metadata,
status: flushed.status,
});
} else {
await repo.update(plan.id, workspaceId, flushed);
}
await applyFinalize(
repo,
planFinalizeAssistant(assistantId),
{ chatId: 'c1', workspaceId, userId: 'u1' },
flushed,
);
}
it('plan: update when the upfront insert returned an id', () => {

View File

@@ -229,15 +229,6 @@ describe('prepareAgentStep', () => {
// The synthesis instruction is appended.
expect(result?.system).toContain(FINAL_STEP_INSTRUCTION);
});
it('pins the off-by-one boundary (MAX-2 is not final, MAX-1 is)', () => {
// Boundary expressed via the constant, not a hardcoded 18/19, so the test
// tracks MAX_AGENT_STEPS if the cap ever changes.
expect(prepareAgentStep(MAX_AGENT_STEPS - 2, 'SYS')).toBeUndefined();
const atBoundary = prepareAgentStep(MAX_AGENT_STEPS - 1, 'SYS');
expect(atBoundary).toBeDefined();
expect(atBoundary?.toolChoice).toBe('none');
});
});
/**

View File

@@ -445,6 +445,13 @@ export class AiChatService implements OnModuleInit {
}
};
// Serialize the per-step updates (#183 review): onStepFinish fires them
// without await, so two could otherwise commit out of order on different pool
// connections (step N landing after N+1). Chaining each onto the previous
// keeps the persisted row monotonic with step order; each link short-circuits
// on `finalized`, so a tail of late updates is cheap.
let stepUpdateChain: Promise<void> = Promise.resolve();
// Terminal finalize: write the completed/error/aborted row exactly once
// across the (mutually-exclusive, at-most-once) onFinish/onError/onAbort
// callbacks — mirroring the pre-#183 persist-at-most-once guard for the
@@ -457,32 +464,21 @@ export class AiChatService implements OnModuleInit {
if (finalized) return;
finalized = true;
const plan = planFinalizeAssistant(assistantId);
if (plan.kind === 'insert') {
// The upfront insert failed: fall back to inserting the terminal row so
// the turn is not lost entirely.
try {
await this.aiChatMessageRepo.insert({
chatId,
workspaceId: workspace.id,
userId: user.id,
role: 'assistant',
content: flushed.content,
toolCalls: (flushed.toolCalls ?? null) as never,
metadata: flushed.metadata as never,
status: flushed.status,
});
} catch (err) {
this.logger.error(
'Failed to persist terminal assistant message',
err as Error,
);
}
return;
}
try {
await this.aiChatMessageRepo.update(plan.id, workspace.id, flushed);
// 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', err as Error);
this.logger.error(
`Failed to finalize assistant message (kind=${plan.kind})`,
err as Error,
);
}
};
@@ -536,9 +532,10 @@ export class AiChatService implements OnModuleInit {
inProgressText = '';
// Step-granular durability (#183): persist this finished step (its text +
// tool calls + tool RESULTS) the moment it ends, so a process death after
// this point still recovers the step. Fire-and-forget but error-tolerant
// (updateStreaming logs + swallows) — never throw into the stream.
void updateStreaming();
// this point still recovers the step. Not awaited here (never block the
// stream), but SERIALIZED via stepUpdateChain so the writes commit in
// step order; updateStreaming is error-tolerant (logs + swallows).
stepUpdateChain = stepUpdateChain.then(() => updateStreaming());
},
onFinish: async ({ text, finishReason, totalUsage, usage, steps }) => {
// DIAGNOSTIC (Safari stream-drop investigation) — temporary: success
@@ -1116,6 +1113,46 @@ export function planFinalizeAssistant(
return assistantId ? { kind: 'update', id: assistantId } : { kind: 'insert' };
}
/** The repo surface the terminal finalize needs (structural — the real repo and
* a test mock both satisfy it). */
export interface FinalizeRepo {
insert(insertable: Record<string, unknown>): Promise<unknown>;
update(
id: string,
workspaceId: string,
patch: AssistantFlush,
): Promise<unknown>;
}
/**
* 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.
*/
export async function applyFinalize(
repo: FinalizeRepo,
plan: { kind: 'update'; id: string } | { kind: 'insert' },
base: { chatId: string; workspaceId: string; userId: string },
flushed: AssistantFlush,
): Promise<void> {
if (plan.kind === 'update') {
await repo.update(plan.id, base.workspaceId, flushed);
return;
}
await repo.insert({
chatId: base.chatId,
workspaceId: base.workspaceId,
userId: base.userId,
role: 'assistant',
content: flushed.content,
toolCalls: flushed.toolCalls ?? null,
metadata: flushed.metadata,
status: flushed.status,
});
}
/**
* PURE assistant-row builder (#183 step-granular durability). Given the turn's
* accumulated steps + the in-progress (not-yet-finished) text + the lifecycle

View File

@@ -122,6 +122,80 @@ describe('buildChatMarkdown (server) — structure', () => {
expect(md).toContain('"title": "Hello"');
});
// #186 re-review pt 1: restore the parity coverage of the removed client spec —
// error state, unknown-tool fallback (en + ru), and the circular-stringify catch.
it('renders a tool part in the error state with its errorText', () => {
const md = buildChatMarkdown({
title: 'T',
chatId: 'c',
rows: [
row({
role: 'assistant',
metadata: {
parts: [
{
type: 'tool-getPage',
state: 'output-error',
input: { id: 'p1' },
errorText: 'page not found',
},
],
} as never,
}),
],
});
expect(md).toContain('**Tool: Read page** (`getPage`) — error');
expect(md).toContain('**Error:** page not found');
});
it('falls back to "Ran tool <name>" for an unknown tool (en) and the ru variant', () => {
const parts = [
{
type: 'tool-mysteryTool',
state: 'output-available',
output: { ok: 1 },
},
];
const en = buildChatMarkdown({
title: 'T',
chatId: 'c',
rows: [row({ role: 'assistant', metadata: { parts } as never })],
});
expect(en).toContain('**Tool: Ran tool mysteryTool** (`mysteryTool`)');
const ru = buildChatMarkdown({
title: 'T',
chatId: 'c',
lang: 'ru',
rows: [row({ role: 'assistant', metadata: { parts } as never })],
});
expect(ru).toContain('Выполнил инструмент mysteryTool');
});
it('does not throw on a circular tool output (falls back to String)', () => {
const circular: Record<string, unknown> = {};
circular.self = circular;
expect(() =>
buildChatMarkdown({
title: 'T',
chatId: 'c',
rows: [
row({
role: 'assistant',
metadata: {
parts: [
{
type: 'tool-getPage',
state: 'output-available',
output: circular,
},
],
} as never,
}),
],
}),
).not.toThrow();
});
it('emits a token footer + total when usage is present', () => {
const md = buildChatMarkdown({
title: 'T',

View File

@@ -1,4 +1,4 @@
import { Injectable } from '@nestjs/common';
import { Injectable, Logger } from '@nestjs/common';
import { InjectKysely } from 'nestjs-kysely';
import { KyselyDB, KyselyTransaction } from '../../types/kysely.types';
import { dbOrTx } from '../../utils';
@@ -25,6 +25,8 @@ const FIND_ALL_BY_CHAT_LIMIT = 5000;
@Injectable()
export class AiChatMessageRepo {
private readonly logger = new Logger(AiChatMessageRepo.name);
constructor(@InjectKysely() private readonly db: KyselyDB) {}
// The `tsv` column is a trigger-maintained tsvector used only for
@@ -87,17 +89,32 @@ export class AiChatMessageRepo {
async findAllByChat(
chatId: string,
workspaceId: string,
// Injectable for tests so truncation can be exercised on a modest volume.
limit: number = FIND_ALL_BY_CHAT_LIMIT,
): Promise<AiChatMessage[]> {
return this.db
// Fetch newest-first (+1 to DETECT truncation), so on overflow we keep the
// NEWEST `limit` messages — the recent conversation matters most for an
// export — rather than silently dropping the tail (#183 review). Reverse back
// to chronological for rendering, like findRecent.
const rows = await this.db
.selectFrom('aiChatMessages')
.select(this.baseFields)
.where('chatId', '=', chatId)
.where('workspaceId', '=', workspaceId)
.where('deletedAt', 'is', null)
.orderBy('createdAt', 'asc')
.orderBy('id', 'asc')
.limit(FIND_ALL_BY_CHAT_LIMIT)
.orderBy('createdAt', 'desc')
.orderBy('id', 'desc')
.limit(limit + 1)
.execute();
if (rows.length > limit) {
rows.length = limit; // keep the newest `limit` (rows are newest-first here)
this.logger.warn(
`Chat ${chatId} export truncated to the newest ${limit} messages ` +
`(older messages omitted).`,
);
}
return rows.reverse();
}
// Load the most RECENT `limit` messages for a chat and return them in

View File

@@ -231,4 +231,40 @@ describe('AiChatMessageRepo.update + sweepStreaming [integration]', () => {
// ...while the stale one alongside it was swept to 'aborted'.
expect(byId.get(stale.id)!.status).toBe('aborted');
});
it('findAllByChat caps the result, keeping the NEWEST messages in order (#183 review)', async () => {
// A dedicated chat so the cap test is independent of the rows above.
const cappedChat = (
await createChat(db, { workspaceId, creatorId: userId })
).id;
const base = Date.now();
// Three messages at strictly increasing timestamps.
await createMessage(db, {
workspaceId,
chatId: cappedChat,
content: 'm1-oldest',
createdAt: new Date(base),
});
await createMessage(db, {
workspaceId,
chatId: cappedChat,
content: 'm2',
createdAt: new Date(base + 1000),
});
await createMessage(db, {
workspaceId,
chatId: cappedChat,
content: 'm3-newest',
createdAt: new Date(base + 2000),
});
// Cap of 2 -> the OLDEST message is dropped; the newest two stay, in
// chronological order (oldest -> newest).
const capped = await repo.findAllByChat(cappedChat, workspaceId, 2);
expect(capped.map((r) => r.content)).toEqual(['m2', 'm3-newest']);
// Without a cap (well above the row count) all three come back in order.
const all = await repo.findAllByChat(cappedChat, workspaceId, 100);
expect(all.map((r) => r.content)).toEqual(['m1-oldest', 'm2', 'm3-newest']);
});
});

View File

@@ -238,6 +238,9 @@ export async function createMessage(
content?: string | null;
status?: string | null;
metadata?: unknown;
// Explicit timestamp so a test can control message ORDER (the default DB
// now() can tie within a millisecond, and the v4 id is not time-ordered).
createdAt?: Date;
},
): Promise<{ id: string }> {
const id = randomUUID();
@@ -252,6 +255,7 @@ export async function createMessage(
content: args.content ?? null,
status: args.status ?? null,
metadata: (args.metadata ?? null) as any,
...(args.createdAt ? { createdAt: args.createdAt } : {}),
})
.returning(['id'])
.executeTakeFirstOrThrow();