Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 agent_coder
|
ae3dfd8de6 | ||
|
agent_coder
|
c96fafc4ad | ||
|
agent_coder
|
2f8c5d9a98 |
@@ -27,6 +27,7 @@ vi.mock("@/features/ai-chat/utils/markdown.ts", async () => {
|
||||
|
||||
import MessageItem from "./message-item";
|
||||
import { messageSignature } from "@/features/ai-chat/utils/message-signature.ts";
|
||||
import { splitPlainChunks } from "./streaming-plain-text";
|
||||
|
||||
// matchMedia (read by MantineProvider) is stubbed globally in vitest.setup.ts.
|
||||
|
||||
@@ -114,3 +115,89 @@ describe("MessageItem markdown memoization", () => {
|
||||
expect(queryByText("streamed answer")).not.toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
// PERF SMOKE (#492): the whole point of the incremental streaming render is that
|
||||
// the ANSWER path costs O(number of markdown blocks), NOT O(number of throttled
|
||||
// ~20Hz ticks). Pre-#492 the finalized MarkdownPart re-parsed the WHOLE growing
|
||||
// answer on every delta — a synthetic ~100 KB stream measured 394 renderChatMarkdown
|
||||
// calls (one per tick). With the incremental render each STABILIZED block is parsed
|
||||
// exactly once (memoized in MarkdownChunk) and the live tail is cheap plain text, so
|
||||
// the call count collapses to ~= the block count regardless of tick granularity.
|
||||
describe("MessageItem streaming answer render is O(blocks), not O(ticks)", () => {
|
||||
// ~100 KB answer. Each section is a heading + a paragraph — TWO blank-line
|
||||
// delimited markdown blocks — so the safe-cut block count is ~2× the section
|
||||
// count. The perf claim is about the BLOCK count (the memoization granularity),
|
||||
// measured directly with splitPlainChunks below, not the section count.
|
||||
const buildAnswer = () => {
|
||||
const SECTIONS = 100;
|
||||
const paragraphs: string[] = [];
|
||||
for (let i = 0; i < SECTIONS; i++) {
|
||||
paragraphs.push(`## Section ${i}\n\n` + "lorem ipsum dolor ".repeat(55));
|
||||
}
|
||||
const full = paragraphs.join("\n\n");
|
||||
// The number of memoized markdown blocks the incremental render splits into
|
||||
// (all but the live tail are parsed once each).
|
||||
return { full, blocks: splitPlainChunks(full).length };
|
||||
};
|
||||
|
||||
const streamMsg = (text: string, state: "streaming" | "done"): UIMessage =>
|
||||
({
|
||||
id: "m1",
|
||||
role: "assistant",
|
||||
parts: [{ type: "text", text, state }],
|
||||
}) as UIMessage;
|
||||
|
||||
it("parses each block ~once over a 100KB stream (≈blocks, ≪ ticks)", () => {
|
||||
renderChatMarkdownSpy.mockClear();
|
||||
const { full, blocks } = buildAnswer();
|
||||
const CHUNK = 128; // a realistic ~20Hz throttled delta size
|
||||
const ticks = Math.ceil(full.length / CHUNK);
|
||||
|
||||
let msg = streamMsg(full.slice(0, CHUNK), "streaming");
|
||||
const { rerender } = render(
|
||||
<MantineProvider>
|
||||
<MessageItem
|
||||
message={msg}
|
||||
signature={messageSignature(msg)}
|
||||
turnStreaming
|
||||
/>
|
||||
</MantineProvider>,
|
||||
);
|
||||
for (let end = 2 * CHUNK; end < full.length; end += CHUNK) {
|
||||
msg = streamMsg(full.slice(0, end), "streaming");
|
||||
rerender(
|
||||
<MantineProvider>
|
||||
<MessageItem
|
||||
message={msg}
|
||||
signature={messageSignature(msg)}
|
||||
turnStreaming
|
||||
/>
|
||||
</MantineProvider>,
|
||||
);
|
||||
}
|
||||
// Finalize: the streaming→done flip renders the whole answer through ONE
|
||||
// canonical pass (visual parity), so the finished DOM matches the pre-#492
|
||||
// output. This is the single extra parse on top of the per-block ones.
|
||||
const done = streamMsg(full, "done");
|
||||
rerender(
|
||||
<MantineProvider>
|
||||
<MessageItem message={done} signature={messageSignature(done)} />
|
||||
</MantineProvider>,
|
||||
);
|
||||
|
||||
const calls = renderChatMarkdownSpy.mock.calls.length;
|
||||
// Sanity: the stream really had far more ticks than blocks (else the test is
|
||||
// vacuous — the point is that calls scale with blocks, not ticks).
|
||||
expect(ticks).toBeGreaterThan(blocks * 3);
|
||||
// O(blocks): each stabilized block parsed once + the single final whole-text
|
||||
// parse. A small constant absorbs the finalize render and the live-tail block;
|
||||
// the load-bearing claim is the bound below.
|
||||
expect(calls).toBeLessThanOrEqual(blocks + 2);
|
||||
// ≪ ticks — and, non-vacuously, the blocks WERE parsed (not skipped entirely).
|
||||
expect(calls).toBeLessThan(ticks / 3);
|
||||
expect(calls).toBeGreaterThan(blocks / 2);
|
||||
// MUTATION-VERIFY (documented, not run here): dropping the `memo()` wrapper on
|
||||
// MarkdownChunk (so every stable block re-parses each tick) drives `calls`
|
||||
// toward `ticks` (~394), reddening both upper-bound assertions above.
|
||||
});
|
||||
});
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { render } from "@testing-library/react";
|
||||
import { MantineProvider } from "@mantine/core";
|
||||
import type { UIMessage } from "@ai-sdk/react";
|
||||
|
||||
// Stub react-i18next (the component reads `useTranslation`). Mirrors the other
|
||||
// message-item specs.
|
||||
vi.mock("react-i18next", () => ({
|
||||
useTranslation: () => ({ t: (key: string) => key }),
|
||||
}));
|
||||
|
||||
import MessageItem from "./message-item";
|
||||
import { messageSignature } from "@/features/ai-chat/utils/message-signature.ts";
|
||||
// The REAL canonical renderer (NOT the spy the memo test installs): this file
|
||||
// exercises the actual markdown output so the visual-regression assertions below
|
||||
// compare against genuine HTML (incl. the schema's `<li><p>` wrappers).
|
||||
import { renderChatMarkdown } from "@/features/ai-chat/utils/markdown.ts";
|
||||
import classes from "./ai-chat.module.css";
|
||||
|
||||
const msg = (
|
||||
parts: UIMessage["parts"],
|
||||
extra?: Partial<UIMessage>,
|
||||
): UIMessage =>
|
||||
({ id: "m1", role: "assistant", parts, ...extra }) as UIMessage;
|
||||
|
||||
const renderRow = (message: UIMessage, turnStreaming = false) =>
|
||||
render(
|
||||
<MantineProvider>
|
||||
<MessageItem
|
||||
message={message}
|
||||
signature={messageSignature(message)}
|
||||
turnStreaming={turnStreaming}
|
||||
/>
|
||||
</MantineProvider>,
|
||||
);
|
||||
|
||||
// A rich multi-block answer that exercises headings, a list (the `<li><p>` case
|
||||
// the scoped CSS tightens), inline emphasis, and multiple paragraphs.
|
||||
const ANSWER = [
|
||||
"# Заголовок",
|
||||
"",
|
||||
"Первый абзац с **жирным** и `кодом`.",
|
||||
"",
|
||||
"- пункт один",
|
||||
"- пункт два",
|
||||
"",
|
||||
"Второй абзац.",
|
||||
].join("\n");
|
||||
|
||||
describe("MessageItem final render — visual parity with the canonical pipeline", () => {
|
||||
it("a finalized text part renders exactly renderChatMarkdown(text)", () => {
|
||||
const { container } = renderRow(
|
||||
msg([{ type: "text", text: ANSWER, state: "done" }]),
|
||||
);
|
||||
const block = container.querySelector(`.${classes.markdown}`);
|
||||
expect(block).not.toBeNull();
|
||||
// Byte-for-byte the canonical output (the SAME whole-text pass the pre-#492
|
||||
// MarkdownPart produced), including `<li><p>…</p></li>` wrappers.
|
||||
expect(block!.innerHTML).toBe(renderChatMarkdown(ANSWER, {}));
|
||||
// The list wrapper is really present (guards against a vacuous empty render).
|
||||
expect(container.querySelectorAll("li p").length).toBe(2);
|
||||
});
|
||||
|
||||
it("the streaming incremental view CONVERGES to the canonical render on finish", () => {
|
||||
// Mount mid-stream (live tail) — the DOM here is the incremental view.
|
||||
const { container, rerender } = render(
|
||||
<MantineProvider>
|
||||
<MessageItem
|
||||
message={msg([{ type: "text", text: ANSWER, state: "streaming" }])}
|
||||
signature={messageSignature(
|
||||
msg([{ type: "text", text: ANSWER, state: "streaming" }]),
|
||||
)}
|
||||
turnStreaming
|
||||
/>
|
||||
</MantineProvider>,
|
||||
);
|
||||
// Finish the turn: state flips to done AND the turn is no longer streaming.
|
||||
const done = msg([{ type: "text", text: ANSWER, state: "done" }]);
|
||||
rerender(
|
||||
<MantineProvider>
|
||||
<MessageItem message={done} signature={messageSignature(done)} />
|
||||
</MantineProvider>,
|
||||
);
|
||||
// After finish there is exactly ONE canonical markdown container whose HTML is
|
||||
// the whole-text render — identical to the non-streaming path above.
|
||||
const blocks = container.querySelectorAll(`.${classes.markdown}`);
|
||||
expect(blocks.length).toBe(1);
|
||||
expect(blocks[0].innerHTML).toBe(renderChatMarkdown(ANSWER, {}));
|
||||
});
|
||||
|
||||
it("neutralizeInternalLinks is honored on the finalized render", () => {
|
||||
const linkAnswer = "См. [страницу](/p/abc).";
|
||||
const { container } = render(
|
||||
<MantineProvider>
|
||||
<MessageItem
|
||||
message={msg([{ type: "text", text: linkAnswer, state: "done" }])}
|
||||
signature={messageSignature(
|
||||
msg([{ type: "text", text: linkAnswer, state: "done" }]),
|
||||
)}
|
||||
neutralizeInternalLinks
|
||||
/>
|
||||
</MantineProvider>,
|
||||
);
|
||||
const block = container.querySelector(`.${classes.markdown}`);
|
||||
expect(block!.innerHTML).toBe(
|
||||
renderChatMarkdown(linkAnswer, { neutralizeInternalLinks: true }),
|
||||
);
|
||||
// The internal link was made inert (no href) by the neutralization flag.
|
||||
const a = container.querySelector("a");
|
||||
expect(a?.hasAttribute("href")).toBe(false);
|
||||
});
|
||||
});
|
||||
@@ -4,6 +4,7 @@ import { useTranslation } from "react-i18next";
|
||||
import type { UIMessage } from "@ai-sdk/react";
|
||||
import ToolCallCard from "@/features/ai-chat/components/tool-call-card.tsx";
|
||||
import ReasoningBlock from "@/features/ai-chat/components/reasoning-block.tsx";
|
||||
import { StreamingMarkdownText } from "@/features/ai-chat/components/streaming-markdown-text.tsx";
|
||||
import ChatErrorAlert from "@/features/ai-chat/components/chat-error-alert.tsx";
|
||||
import ChatStoppedNotice from "@/features/ai-chat/components/chat-stopped-notice.tsx";
|
||||
import { ToolUiPart, isToolPart } from "@/features/ai-chat/utils/tool-parts.tsx";
|
||||
@@ -86,17 +87,39 @@ interface MessageItemProps {
|
||||
* One assistant text part rendered as sanitized markdown. Memoized on its inputs
|
||||
* so a finalized text part is NOT re-parsed on every streamed delta: during a
|
||||
* turn only the actively-growing tail part changes its `text`, so every earlier
|
||||
* part hits the memo and skips the expensive marked + DOMPurify pass. Props are
|
||||
* primitives, so React.memo's default shallow compare is exactly right (the
|
||||
* `text` string is compared by value).
|
||||
* part hits the memo and skips the expensive canonical parse + DOMPurify pass.
|
||||
* Props are primitives, so React.memo's default shallow compare is exactly right
|
||||
* (the `text` string is compared by value).
|
||||
*
|
||||
* Streaming gate (#492) — mirrors ReasoningBlock:
|
||||
* - `streaming` (this is the live, actively-growing tail part of an in-flight
|
||||
* turn): render incrementally via StreamingMarkdownText — the stabilized blocks
|
||||
* go through the canonical pipeline (each parsed ONCE, memoized) and only the
|
||||
* live tail is cheap plain text. This makes the per-tick cost O(new blocks),
|
||||
* not the pre-#492 O(ticks) whole-answer re-parse on every ~20Hz delta.
|
||||
* - finalized (the common case, and the turn-end flip): render the WHOLE text
|
||||
* through ONE canonical pass — byte-identical to the pre-#492 output (visual
|
||||
* parity). The row re-renders on the streaming→done flip because
|
||||
* `messageSignature` tracks each part's `state` (and `turnStreaming` flips at
|
||||
* turn end), so the incremental view always converges to this single render.
|
||||
*/
|
||||
const MarkdownPart = memo(function MarkdownPart({
|
||||
text,
|
||||
neutralizeInternalLinks,
|
||||
streaming,
|
||||
}: {
|
||||
text: string;
|
||||
neutralizeInternalLinks: boolean;
|
||||
streaming: boolean;
|
||||
}) {
|
||||
if (streaming) {
|
||||
return (
|
||||
<StreamingMarkdownText
|
||||
text={text}
|
||||
neutralizeInternalLinks={neutralizeInternalLinks}
|
||||
/>
|
||||
);
|
||||
}
|
||||
const html = renderChatMarkdown(text, { neutralizeInternalLinks });
|
||||
if (html) {
|
||||
return (
|
||||
@@ -179,47 +202,10 @@ function MessageItem({
|
||||
{resolveAssistantName(assistantName) ?? t("AI agent")}
|
||||
</Text>
|
||||
{message.parts.map((part, index) => {
|
||||
if (part.type === "reasoning") {
|
||||
// Reasoning ("thinking") -> a collapsible block with its own token
|
||||
// count. Empty/whitespace reasoning with no authoritative count carries
|
||||
// nothing to show, so skip it (avoids an empty 0-token block).
|
||||
const text = (part as { text?: string }).text ?? "";
|
||||
if (!text.trim() && !(reasoningTokens && reasoningTokens > 0))
|
||||
return null;
|
||||
// Absent state (persisted rows) and "done" both mean finalized.
|
||||
// `messageSignature` already includes each part's `state`, so the
|
||||
// streaming→done flip changes the row signature and re-renders this
|
||||
// row — which is what lets ReasoningBlock switch from chunked plain
|
||||
// text to its one-time markdown parse (see reasoning-block.tsx).
|
||||
// ALSO require the turn to be live: a part stranded at
|
||||
// `state:"streaming"` after the turn ended (no `reasoning-end` — see
|
||||
// the `turnStreaming` prop doc) must still finalize and parse.
|
||||
const streaming =
|
||||
turnStreaming && (part as { state?: string }).state === "streaming";
|
||||
return (
|
||||
<ReasoningBlock
|
||||
key={index}
|
||||
text={text}
|
||||
tokens={reasoningTokens}
|
||||
streaming={streaming}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
if (part.type === "text") {
|
||||
// Skip empty/whitespace-only text parts (a streaming message often
|
||||
// starts with an empty text part before the first token arrives); the
|
||||
// typing indicator covers that gap until real content streams in.
|
||||
if (!part.text.trim()) return null;
|
||||
return (
|
||||
<MarkdownPart
|
||||
key={index}
|
||||
text={part.text}
|
||||
neutralizeInternalLinks={neutralizeInternalLinks}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
// Tool parts (`tool-*` / `dynamic-tool`) are template-literal kinds, so
|
||||
// they cannot be a `switch` case; the runtime guard handles them, and the
|
||||
// switch below covers every CLOSED (literal-typed) part kind with a
|
||||
// compile-time exhaustiveness check in its default.
|
||||
if (isToolPart(part.type)) {
|
||||
return (
|
||||
<ToolCallCard
|
||||
@@ -232,7 +218,76 @@ function MessageItem({
|
||||
);
|
||||
}
|
||||
|
||||
return null;
|
||||
switch (part.type) {
|
||||
case "reasoning": {
|
||||
// Reasoning ("thinking") -> a collapsible block with its own token
|
||||
// count. Empty/whitespace reasoning with no authoritative count
|
||||
// carries nothing to show, so skip it (avoids an empty 0-token block).
|
||||
const text = part.text ?? "";
|
||||
if (!text.trim() && !(reasoningTokens && reasoningTokens > 0))
|
||||
return null;
|
||||
// Absent state (persisted rows) and "done" both mean finalized.
|
||||
// `messageSignature` already includes each part's `state`, so the
|
||||
// streaming→done flip changes the row signature and re-renders this
|
||||
// row — which is what lets ReasoningBlock switch from chunked plain
|
||||
// text to its one-time markdown parse (see reasoning-block.tsx).
|
||||
// ALSO require the turn to be live: a part stranded at
|
||||
// `state:"streaming"` after the turn ended (no `reasoning-end` — see
|
||||
// the `turnStreaming` prop doc) must still finalize and parse.
|
||||
const streaming = turnStreaming && part.state === "streaming";
|
||||
return (
|
||||
<ReasoningBlock
|
||||
key={index}
|
||||
text={text}
|
||||
tokens={reasoningTokens}
|
||||
streaming={streaming}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
case "text": {
|
||||
// Skip empty/whitespace-only text parts (a streaming message often
|
||||
// starts with an empty text part before the first token arrives); the
|
||||
// typing indicator covers that gap until real content streams in.
|
||||
if (!part.text.trim()) return null;
|
||||
// The live, actively-growing tail part of the in-flight turn renders
|
||||
// incrementally (see MarkdownPart); a finalized part (persisted, or
|
||||
// the turn-end flip) renders the whole text through one canonical
|
||||
// pass. Same liveness rule as the reasoning branch above.
|
||||
const streaming = turnStreaming && part.state === "streaming";
|
||||
return (
|
||||
<MarkdownPart
|
||||
key={index}
|
||||
text={part.text}
|
||||
neutralizeInternalLinks={neutralizeInternalLinks}
|
||||
streaming={streaming}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
case "source-url":
|
||||
case "source-document":
|
||||
case "file":
|
||||
case "step-start":
|
||||
// Not surfaced in the chat bubble (v1) — same as the pre-#492 default.
|
||||
return null;
|
||||
|
||||
default: {
|
||||
// Compile-time exhaustiveness over the CLOSED union members: every
|
||||
// literal-typed part kind is handled above, so the only kinds that
|
||||
// can reach here are the OPEN template-literal ones (`tool-*` — caught
|
||||
// by the guard at runtime — and `data-*`) plus `dynamic-tool`. Adding
|
||||
// a NEW closed part kind to UIMessagePart makes this assignment fail
|
||||
// to compile, forcing it to be handled instead of silently ignored
|
||||
// (this replaces the pre-#492 fall-through `return null` + WARNING).
|
||||
const _exhaustive:
|
||||
| `tool-${string}`
|
||||
| "dynamic-tool"
|
||||
| `data-${string}` = part.type;
|
||||
void _exhaustive;
|
||||
return null;
|
||||
}
|
||||
}
|
||||
})}
|
||||
{/* A persisted turn error (server stored it in metadata.error). Rendered
|
||||
here so it survives a thread remount and shows in reopened history. */}
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
import { memo, useMemo } from "react";
|
||||
import { splitPlainChunks } from "@/features/ai-chat/components/streaming-plain-text.tsx";
|
||||
import { renderChatMarkdown } from "@/features/ai-chat/utils/markdown.ts";
|
||||
import classes from "@/features/ai-chat/components/ai-chat.module.css";
|
||||
|
||||
/**
|
||||
* One STABILIZED markdown block, rendered through the canonical pipeline and
|
||||
* memoized on its string prop. During streaming only the TAIL chunk grows (the
|
||||
* `splitPlainChunks` append-only invariant guarantees every earlier chunk is
|
||||
* byte-identical across deltas), so React skips every stable block and each one
|
||||
* is parsed by `renderChatMarkdown` EXACTLY ONCE — turning the pre-#492
|
||||
* "re-parse the whole accumulated answer on every ~20Hz tick" (O(ticks)) into
|
||||
* O(number of blocks). The markup is DOMPurify-sanitized inside renderChatMarkdown
|
||||
* before it reaches `dangerouslySetInnerHTML`.
|
||||
*
|
||||
* NOTE (transient streaming-only artifact): a safe cut is a blank-line boundary,
|
||||
* so a construct that legitimately contains a blank line (e.g. a fenced code block
|
||||
* with an empty line) can be split across chunks and render oddly WHILE it is still
|
||||
* streaming. This is cosmetic and self-heals: the moment the part finalizes,
|
||||
* MarkdownPart renders the WHOLE text through one canonical pass (visual parity
|
||||
* with the pre-#492 output). The reasoning path makes the same trade (plain text
|
||||
* while streaming, one markdown parse at the end).
|
||||
*/
|
||||
const MarkdownChunk = memo(function MarkdownChunk({
|
||||
text,
|
||||
neutralizeInternalLinks,
|
||||
}: {
|
||||
text: string;
|
||||
neutralizeInternalLinks: boolean;
|
||||
}) {
|
||||
const html = renderChatMarkdown(text, { neutralizeInternalLinks });
|
||||
if (html) {
|
||||
return (
|
||||
<div
|
||||
className={classes.markdown}
|
||||
// Sanitized by renderChatMarkdown (DOMPurify) before insertion.
|
||||
dangerouslySetInnerHTML={{ __html: html }}
|
||||
/>
|
||||
);
|
||||
}
|
||||
// Malformed/unsupported markdown could not render synchronously: raw text.
|
||||
return (
|
||||
<div className={classes.markdown} style={{ whiteSpace: "pre-wrap" }}>
|
||||
{text}
|
||||
</div>
|
||||
);
|
||||
});
|
||||
|
||||
/**
|
||||
* The cheap streaming-time stand-in for the finalized answer's one-time markdown
|
||||
* parse (see MarkdownPart in message-item.tsx). Mirrors StreamingPlainText's
|
||||
* chunked-memo pattern but renders the STABILIZED prefix as real markdown (each
|
||||
* block parsed once, memoized) and only the LIVE tail as flat plain text — so the
|
||||
* user sees formatted output for everything up to the last safe cut, and the not-
|
||||
* yet-stable tail (which markdown-parsing every tick would make O(ticks)) stays a
|
||||
* single cheap escaped text node until it stabilizes into a new block.
|
||||
*
|
||||
* `splitPlainChunks` yields chunks where, under append-only growth, every chunk
|
||||
* except the LAST is immutable; the last chunk is the live tail. Index keys are
|
||||
* therefore stable (a given index never changes to a different chunk's content).
|
||||
*/
|
||||
export function StreamingMarkdownText({
|
||||
text,
|
||||
neutralizeInternalLinks,
|
||||
}: {
|
||||
text: string;
|
||||
neutralizeInternalLinks: boolean;
|
||||
}) {
|
||||
const chunks = useMemo(() => splitPlainChunks(text), [text]);
|
||||
return (
|
||||
<>
|
||||
{chunks.map((chunk, index) =>
|
||||
index < chunks.length - 1 ? (
|
||||
<MarkdownChunk
|
||||
key={index}
|
||||
text={chunk}
|
||||
neutralizeInternalLinks={neutralizeInternalLinks}
|
||||
/>
|
||||
) : (
|
||||
// The live tail: flat, React-escaped plain text (no markdown parse, no
|
||||
// sanitizer, no innerHTML). `pre-wrap` preserves its newlines; trailing
|
||||
// separator newlines are dropped at display time so the block gap comes
|
||||
// from the markdown margins, not a doubled empty line (mirrors
|
||||
// PlainChunk in streaming-plain-text.tsx).
|
||||
<div
|
||||
key={index}
|
||||
className={classes.markdown}
|
||||
style={{ whiteSpace: "pre-wrap" }}
|
||||
>
|
||||
{chunk.replace(/\n+$/, "")}
|
||||
</div>
|
||||
),
|
||||
)}
|
||||
</>
|
||||
);
|
||||
}
|
||||
@@ -35,6 +35,7 @@ import {
|
||||
import { PaginationOptions } from '@docmost/db/pagination/pagination-options';
|
||||
import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo';
|
||||
import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
|
||||
import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo';
|
||||
import { PageRepo } from '@docmost/db/repos/page/page.repo';
|
||||
import { UserThrottlerGuard } from '../../integrations/throttle/user-throttler.guard';
|
||||
import { AI_CHAT_THROTTLER } from '../../integrations/throttle/throttler-names';
|
||||
@@ -43,6 +44,8 @@ import {
|
||||
AiChatRunHooks,
|
||||
AiChatService,
|
||||
AiChatStreamBody,
|
||||
rowHasInlineParts,
|
||||
hydrateAssistantParts,
|
||||
} from './ai-chat.service';
|
||||
import { AiChatRunService } from './ai-chat-run.service';
|
||||
import { AiTranscriptionService } from './ai-transcription.service';
|
||||
@@ -129,8 +132,39 @@ export class AiChatController {
|
||||
// production. Only touched on the resumable-stream (flag-on) path.
|
||||
private readonly streamRegistry?: AiChatStreamRegistryService,
|
||||
private readonly environment?: EnvironmentService,
|
||||
// #492: reconstruct a #492 mid-run record's parts from the steps table before
|
||||
// returning rows to the client / export. OPTIONAL so positional controller
|
||||
// specs compile unchanged; when absent, hydration is skipped (old-era rows
|
||||
// already carry inline parts, so nothing to reconstruct).
|
||||
private readonly aiChatRunStepRepo?: AiChatRunStepRepo,
|
||||
) {}
|
||||
|
||||
/**
|
||||
* Reconstruct parts for any assistant rows that don't carry them INLINE — a
|
||||
* #492 mid-run record whose per-step parts live in `ai_chat_run_steps` (the
|
||||
* append-persist backend). Every FINISHED row (old-era + #492) and every old-era
|
||||
* streaming snapshot already has inline `metadata.parts`, so the common path
|
||||
* fetches NOTHING and returns the rows untouched; only an actively-streaming
|
||||
* new-style row triggers the batch step fetch. Consumers (seed/poll/export) read
|
||||
* `metadata.parts` off the returned rows exactly as before — the era switch is
|
||||
* invisible to them (reconstructRunParts contract).
|
||||
*/
|
||||
private async withReconstructedParts(
|
||||
rows: AiChatMessage[],
|
||||
workspaceId: string,
|
||||
): Promise<AiChatMessage[]> {
|
||||
if (!this.aiChatRunStepRepo) return rows;
|
||||
const needy = rows.filter(
|
||||
(r) => r.role === 'assistant' && !rowHasInlineParts(r),
|
||||
);
|
||||
if (needy.length === 0) return rows;
|
||||
const stepsByMessage = await this.aiChatRunStepRepo.findByMessageIds(
|
||||
needy.map((r) => r.id),
|
||||
workspaceId,
|
||||
);
|
||||
return hydrateAssistantParts(rows, stepsByMessage);
|
||||
}
|
||||
|
||||
/** List the requesting user's chats in this workspace (paginated). */
|
||||
@HttpCode(HttpStatus.OK)
|
||||
@Post('chats')
|
||||
@@ -184,11 +218,17 @@ export class AiChatController {
|
||||
@AuthWorkspace() workspace: Workspace,
|
||||
) {
|
||||
await this.assertOwnedChat(dto.chatId, user, workspace);
|
||||
return this.aiChatMessageRepo.findByChat(
|
||||
const page = await this.aiChatMessageRepo.findByChat(
|
||||
dto.chatId,
|
||||
workspace.id,
|
||||
pagination,
|
||||
);
|
||||
// #492: reconstruct parts for any active new-style row so the client seed sees
|
||||
// `metadata.parts` unchanged (a no-op for the finished rows that fill a page).
|
||||
return {
|
||||
...page,
|
||||
items: await this.withReconstructedParts(page.items, workspace.id),
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -225,7 +265,10 @@ export class AiChatController {
|
||||
workspace.id,
|
||||
);
|
||||
return {
|
||||
rows,
|
||||
// #492: the delta of an actively-streaming new-style row carries its parts
|
||||
// reconstructed from the steps table, so the degraded poll shows persisted
|
||||
// progress exactly as the pre-#492 full-row snapshot did.
|
||||
rows: await this.withReconstructedParts(rows, workspace.id),
|
||||
cursor,
|
||||
run: run ? { id: run.id, status: run.status } : null,
|
||||
};
|
||||
@@ -247,8 +290,10 @@ export class AiChatController {
|
||||
@AuthWorkspace() workspace: Workspace,
|
||||
): Promise<{ markdown: string }> {
|
||||
const chat = await this.assertOwnedChat(dto.chatId, user, workspace);
|
||||
const rows = await this.aiChatMessageRepo.findAllByChat(
|
||||
dto.chatId,
|
||||
const rows = await this.withReconstructedParts(
|
||||
await this.aiChatMessageRepo.findAllByChat(dto.chatId, workspace.id),
|
||||
// #492: an interrupted-but-still-active turn exports its persisted steps
|
||||
// (reconstructed from the steps table) just like the pre-#492 full row did.
|
||||
workspace.id,
|
||||
);
|
||||
const markdown = buildChatMarkdown({
|
||||
@@ -288,7 +333,13 @@ export class AiChatController {
|
||||
workspace.id,
|
||||
)
|
||||
: undefined;
|
||||
return { run, message: message ?? null };
|
||||
// #492: reconnect to an IN-FLIGHT run reconstructs the projection row's parts
|
||||
// from the steps table (the row itself carries only the step marker mid-run);
|
||||
// a finished run's row already has inline parts, so this is a no-op.
|
||||
const [hydrated] = message
|
||||
? await this.withReconstructedParts([message], workspace.id)
|
||||
: [undefined];
|
||||
return { run, message: hydrated ?? null };
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -22,6 +22,7 @@ import { AiSettingsService } from '../../integrations/ai/ai-settings.service';
|
||||
import { describeProviderError } from '../../integrations/ai/ai-error.util';
|
||||
import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo';
|
||||
import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
|
||||
import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo';
|
||||
import { AiChatPageSnapshotRepo } from '@docmost/db/repos/ai-chat/ai-chat-page-snapshot.repo';
|
||||
import { AiAgentRoleRepo } from '@docmost/db/repos/ai-agent-roles/ai-agent-roles.repo';
|
||||
import { PageRepo } from '@docmost/db/repos/page/page.repo';
|
||||
@@ -518,6 +519,12 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
// 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,
|
||||
// #492 append-persist: per-step INSERT into the lightweight steps table (the
|
||||
// O(Σ steps) replacement for the O(n²) full-row `metadata.parts` rewrite).
|
||||
// OPTIONAL so existing positional constructions (int-specs) compile unchanged;
|
||||
// Nest injects the real singleton. When ABSENT the per-step path falls back to
|
||||
// the pre-#492 full-row flush (no regression, only no WAL win).
|
||||
private readonly aiChatRunStepRepo?: AiChatRunStepRepo,
|
||||
) {}
|
||||
|
||||
// #487: periodic reconcile timer (single-process phase 1). Started in
|
||||
@@ -1559,17 +1566,57 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
// connection when finalize runs, so the SQL `WHERE status='streaming'`
|
||||
// (not this flag) is what prevents it clobbering the terminal row.
|
||||
if (finalized) return null;
|
||||
// Build the flush ONCE so the returned count is EXACTLY the persisted
|
||||
// `stepsPersisted` (both derive from capturedSteps.length at this instant).
|
||||
const flushed = flushAssistant(capturedSteps, '', 'streaming', {
|
||||
pageChanged,
|
||||
partsCache,
|
||||
});
|
||||
const stepsPersisted = flushed.metadata.stepsPersisted as number;
|
||||
// The count derives from capturedSteps.length at THIS instant, so the
|
||||
// returned value is EXACTLY the persisted `stepsPersisted` the ring rotates
|
||||
// on (whether we take the append-persist path or the legacy fallback).
|
||||
const stepsPersisted = capturedSteps.length;
|
||||
try {
|
||||
await this.aiChatMessageRepo.update(assistantId, workspace.id, flushed, {
|
||||
onlyIfStreaming: true,
|
||||
});
|
||||
if (this.aiChatRunStepRepo) {
|
||||
// #492 APPEND-PERSIST: write only THIS finished step's parts to the
|
||||
// steps table (O(step) WAL), then bump the row's CHEAP step marker —
|
||||
// NO growing `metadata.parts` blob (that O(n²) full-row rewrite is
|
||||
// exactly what this removes). The full `metadata.parts` is assembled
|
||||
// once at finalize; a mid-run resume seed is reconstructed from the
|
||||
// step rows (reconstructRunParts). The INSERT is idempotent
|
||||
// (ON CONFLICT DO NOTHING), so a re-fired step never doubles the parts.
|
||||
const index = stepsPersisted - 1;
|
||||
if (index >= 0) {
|
||||
const stepParts = assistantParts(
|
||||
[capturedSteps[index]],
|
||||
'',
|
||||
partsCache,
|
||||
);
|
||||
await this.aiChatRunStepRepo.insertStep(
|
||||
assistantId,
|
||||
workspace.id,
|
||||
index,
|
||||
stepParts,
|
||||
);
|
||||
}
|
||||
// Marker UPDATE: advance stepsPersisted + keep the toolTrace era marker
|
||||
// (bumps updatedAt so the delta poll observes the step, and carries the
|
||||
// frontier a resuming client attaches from). Scoped onlyIfStreaming so a
|
||||
// late marker never clobbers the terminal finalize.
|
||||
await this.aiChatMessageRepo.update(
|
||||
assistantId,
|
||||
workspace.id,
|
||||
{ metadata: stepMarkerMetadata(stepsPersisted) },
|
||||
{ onlyIfStreaming: true },
|
||||
);
|
||||
} else {
|
||||
// Legacy fallback (no steps table wired — positional test builds): the
|
||||
// pre-#492 full-row flush, so parts still land inline on the row.
|
||||
const flushed = flushAssistant(capturedSteps, '', 'streaming', {
|
||||
pageChanged,
|
||||
partsCache,
|
||||
});
|
||||
await this.aiChatMessageRepo.update(
|
||||
assistantId,
|
||||
workspace.id,
|
||||
flushed,
|
||||
{ onlyIfStreaming: true },
|
||||
);
|
||||
}
|
||||
return stepsPersisted;
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
@@ -2749,6 +2796,122 @@ export function rowToUiMessage(row: AiChatMessage): Omit<UIMessage, 'id'> & {
|
||||
return { id: row.id, role, parts: parts as UIMessage['parts'] };
|
||||
}
|
||||
|
||||
/**
|
||||
* Cheap step-marker metadata for the #492 per-step UPDATE. Advances
|
||||
* `stepsPersisted` (the resume attach frontier) and keeps the `toolTraceVersion`
|
||||
* era marker, WITHOUT the growing `parts` blob (those live in the steps table
|
||||
* now; the full `metadata.parts` is assembled once at finalize by flushAssistant).
|
||||
* `parts: []` is kept for shape stability — it reads as an empty inline-parts row,
|
||||
* which is exactly the discriminator that routes reconstruction to the steps table.
|
||||
*/
|
||||
export function stepMarkerMetadata(
|
||||
stepsPersisted: number,
|
||||
): Record<string, unknown> {
|
||||
return { parts: [], toolTraceVersion: 2, stepsPersisted };
|
||||
}
|
||||
|
||||
/**
|
||||
* Whether an assistant row already carries its full UI parts INLINE on the row
|
||||
* (`metadata.parts`). TRUE for every FINISHED row — old-era rows AND #492 rows,
|
||||
* whose full parts are assembled once at finalize — and for old-era streaming
|
||||
* snapshots (the pre-#492 per-step full-row flush). FALSE for a #492 MID-RUN
|
||||
* record, whose per-step parts live in the `ai_chat_run_steps` table. This is the
|
||||
* era discriminator the reconstruct seam branches on — no schema flag needed.
|
||||
*/
|
||||
export function rowHasInlineParts(row: { metadata?: unknown }): boolean {
|
||||
const meta = (row.metadata ?? {}) as { parts?: unknown };
|
||||
return Array.isArray(meta.parts) && meta.parts.length > 0;
|
||||
}
|
||||
|
||||
/**
|
||||
* Concatenate persisted per-step parts (in `stepIndex` order) into the turn's UI
|
||||
* parts (#492). Reproduces EXACTLY what flushAssistant → assistantParts would have
|
||||
* written to `metadata.parts` for those finished steps, since each step row stored
|
||||
* `assistantParts([step])` at persist time.
|
||||
*/
|
||||
export function assembleStepParts(
|
||||
stepRows: ReadonlyArray<{ stepIndex: number; parts: unknown }>,
|
||||
): UIMessage['parts'] {
|
||||
const parts: Array<Record<string, unknown>> = [];
|
||||
for (const step of [...stepRows].sort((a, b) => a.stepIndex - b.stepIndex)) {
|
||||
if (Array.isArray(step.parts)) {
|
||||
parts.push(...(step.parts as Array<Record<string, unknown>>));
|
||||
}
|
||||
}
|
||||
return parts as UIMessage['parts'];
|
||||
}
|
||||
|
||||
/**
|
||||
* reconstructRunParts (#492) — the single backend-switch seam. Given an assistant
|
||||
* ROW and its persisted step rows, return the turn's UI `parts` + the persisted
|
||||
* step count, reading from the ROW when it already carries inline parts (old-era
|
||||
* records AND every finished record) and from the STEPS TABLE otherwise (a #492
|
||||
* mid-run record). The higher-level consumers (attach seed, delta poll, export)
|
||||
* route their row→parts through this / {@link hydrateAssistantParts}, so old and
|
||||
* new records reconstruct identically WITHOUT the consumers branching on the era.
|
||||
*/
|
||||
export function reconstructRunParts(
|
||||
row: { metadata?: unknown; content?: string | null },
|
||||
stepRows: ReadonlyArray<{ stepIndex: number; parts: unknown }>,
|
||||
): { parts: UIMessage['parts']; stepsPersisted: number } {
|
||||
if (rowHasInlineParts(row)) {
|
||||
const meta = row.metadata as {
|
||||
parts: UIMessage['parts'];
|
||||
stepsPersisted?: number;
|
||||
};
|
||||
return {
|
||||
parts: meta.parts,
|
||||
stepsPersisted:
|
||||
typeof meta.stepsPersisted === 'number'
|
||||
? meta.stepsPersisted
|
||||
: stepRows.length,
|
||||
};
|
||||
}
|
||||
if (stepRows.length > 0) {
|
||||
return {
|
||||
parts: assembleStepParts(stepRows),
|
||||
stepsPersisted: stepRows.length,
|
||||
};
|
||||
}
|
||||
// No inline parts and no step rows: an old-era seed / empty streaming row. Fall
|
||||
// back to a single text part from `content` (mirrors rowToUiMessage).
|
||||
return {
|
||||
parts: textPart(row.content ?? '') as UIMessage['parts'],
|
||||
stepsPersisted: 0,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Fill each assistant row's `metadata.parts` from its step rows when the row does
|
||||
* not already carry them inline (a #492 mid-run record), so a consumer that reads
|
||||
* `metadata.parts` off the RAW row (the client seed/poll, the Markdown export)
|
||||
* sees the reconstructed parts with NO change to itself. Rows that already have
|
||||
* inline parts (old-era + finished) and non-assistant rows pass through untouched.
|
||||
* Pure: returns new row objects, never mutates the inputs.
|
||||
*/
|
||||
export function hydrateAssistantParts<
|
||||
T extends { id: string; role?: string; metadata?: unknown },
|
||||
>(
|
||||
rows: ReadonlyArray<T>,
|
||||
stepsByMessage: Map<
|
||||
string,
|
||||
ReadonlyArray<{ stepIndex: number; parts: unknown }>
|
||||
>,
|
||||
): T[] {
|
||||
return rows.map((row) => {
|
||||
if (row.role !== 'assistant' || rowHasInlineParts(row)) return row;
|
||||
const steps = stepsByMessage.get(row.id);
|
||||
if (!steps || steps.length === 0) return row;
|
||||
return {
|
||||
...row,
|
||||
metadata: {
|
||||
...((row.metadata ?? {}) as Record<string, unknown>),
|
||||
parts: assembleStepParts(steps),
|
||||
},
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* The persisted-row patch shape produced by {@link flushAssistant}. It is the
|
||||
* SAME shape the assistant repo insert/update consume (content + toolCalls +
|
||||
|
||||
@@ -32,6 +32,7 @@ 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 { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.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';
|
||||
@@ -125,6 +126,7 @@ import { firstSqlToken } from '../integrations/metrics/metrics.constants';
|
||||
AiChatRepo,
|
||||
AiChatMessageRepo,
|
||||
AiChatRunRepo,
|
||||
AiChatRunStepRepo,
|
||||
AiChatPageSnapshotRepo,
|
||||
AiProviderCredentialsRepo,
|
||||
AiMcpServerRepo,
|
||||
@@ -161,6 +163,7 @@ import { firstSqlToken } from '../integrations/metrics/metrics.constants';
|
||||
AiChatRepo,
|
||||
AiChatMessageRepo,
|
||||
AiChatRunRepo,
|
||||
AiChatRunStepRepo,
|
||||
AiChatPageSnapshotRepo,
|
||||
AiProviderCredentialsRepo,
|
||||
AiMcpServerRepo,
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
import { type Kysely, sql } from 'kysely';
|
||||
|
||||
/**
|
||||
* `ai_chat_run_steps` — append-only per-step persistence for an assistant turn
|
||||
* (#492 wave C). Each finished agent step's UI `parts` (its text part + a part
|
||||
* per tool call, WITH the tool output) is INSERTed as its own lightweight row the
|
||||
* moment the step ends, instead of REWRITING the whole assistant row's growing
|
||||
* `metadata.parts` jsonb on every `onStepFinish`.
|
||||
*
|
||||
* WHY a separate table + INSERT (not a jsonb `||` append on the message row): a
|
||||
* Postgres jsonb UPDATE rewrites the ENTIRE TOASTed row version under MVCC, so
|
||||
* re-persisting a growing `metadata.parts` on every step is O(n²) write volume
|
||||
* (a 50-step run with ~100 KB tool outputs wrote hundreds of MB of WAL / dead
|
||||
* tuples per turn, hammering autovacuum). `||` would only shave the network
|
||||
* payload — the WAL/TOAST rewrite harm remains. An INSERT into a per-step table
|
||||
* writes ONLY that step's bytes, so the per-turn write volume is O(Σ steps).
|
||||
*
|
||||
* The full `metadata.parts` on the message row is assembled ONCE at finalize (the
|
||||
* terminal completed/error/aborted write). Mid-run, a resuming client's seed is
|
||||
* reconstructed by concatenating these step rows in `step_index` order — which
|
||||
* reproduces exactly what the old per-step full-row rewrite persisted. Records
|
||||
* written the OLD way (full `metadata.parts` on the row, no step rows) still
|
||||
* reconstruct from the row unchanged; the two eras are distinguished by whether
|
||||
* the row already carries non-empty `metadata.parts` (see reconstructRunParts /
|
||||
* assembleStepParts in ai-chat.service.ts).
|
||||
*
|
||||
* ON DELETE CASCADE on `message_id`: the step rows are a derived projection of the
|
||||
* assistant message; they must vanish with it (or with its workspace).
|
||||
*/
|
||||
export async function up(db: Kysely<any>): Promise<void> {
|
||||
await db.schema
|
||||
.createTable('ai_chat_run_steps')
|
||||
.ifNotExists()
|
||||
.addColumn('id', 'uuid', (col) =>
|
||||
col.primaryKey().defaultTo(sql`gen_uuid_v7()`),
|
||||
)
|
||||
// The assistant message row this step belongs to (the #183 projection). The
|
||||
// step rows are a derived, per-step slice of that message, so they cascade.
|
||||
.addColumn('message_id', 'uuid', (col) =>
|
||||
col.references('ai_chat_messages.id').onDelete('cascade').notNull(),
|
||||
)
|
||||
.addColumn('workspace_id', 'uuid', (col) =>
|
||||
col.references('workspaces.id').onDelete('cascade').notNull(),
|
||||
)
|
||||
// 0-based index of the finished step within the turn. Ordering key for
|
||||
// reconstruction; unique per message (idempotent step re-persist).
|
||||
.addColumn('step_index', 'integer', (col) => col.notNull())
|
||||
// The step's UI parts (text part + a `tool-*` part per call, WITH output).
|
||||
// Concatenated in step order to rebuild the turn's `metadata.parts`.
|
||||
.addColumn('parts', 'jsonb', (col) => col.notNull())
|
||||
.addColumn('created_at', 'timestamptz', (col) =>
|
||||
col.notNull().defaultTo(sql`now()`),
|
||||
)
|
||||
.execute();
|
||||
|
||||
// Idempotent per-step persist: a retried INSERT of the same (message, step)
|
||||
// is a no-op (the service uses ON CONFLICT DO NOTHING). This also serves the
|
||||
// reconstruction read (WHERE message_id ORDER BY step_index).
|
||||
await db.schema
|
||||
.createIndex('ai_chat_run_steps_message_step_uidx')
|
||||
.ifNotExists()
|
||||
.on('ai_chat_run_steps')
|
||||
.columns(['message_id', 'step_index'])
|
||||
.unique()
|
||||
.execute();
|
||||
}
|
||||
|
||||
export async function down(db: Kysely<any>): Promise<void> {
|
||||
await db.schema.dropTable('ai_chat_run_steps').ifExists().execute();
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
import { Injectable } from '@nestjs/common';
|
||||
import { InjectKysely } from 'nestjs-kysely';
|
||||
import { sql } from 'kysely';
|
||||
import { KyselyDB, KyselyTransaction } from '../../types/kysely.types';
|
||||
import { dbOrTx } from '../../utils';
|
||||
import { AiChatRunStep } from '@docmost/db/types/entity.types';
|
||||
|
||||
/**
|
||||
* Append-only per-step persistence for an assistant turn (#492). Each finished
|
||||
* agent step's UI `parts` (its text part + a `tool-*` part per call, WITH the
|
||||
* tool output) is INSERTed as its own lightweight row the moment the step ends —
|
||||
* instead of REWRITING the assistant row's growing `metadata.parts` jsonb on every
|
||||
* `onStepFinish` (a Postgres jsonb UPDATE rewrites the whole TOASTed row version
|
||||
* under MVCC, so that was O(n²) WAL/dead-tuple churn per turn).
|
||||
*
|
||||
* The full `metadata.parts` on the message row is assembled ONCE at finalize;
|
||||
* mid-run, a resuming client's seed is rebuilt from these rows in `stepIndex`
|
||||
* order (see `assembleStepParts` / the reconstruct seam in ai-chat.service.ts).
|
||||
* Every method is workspace-scoped as defense-in-depth.
|
||||
*/
|
||||
@Injectable()
|
||||
export class AiChatRunStepRepo {
|
||||
constructor(@InjectKysely() private readonly db: KyselyDB) {}
|
||||
|
||||
/**
|
||||
* Append one finished step's parts. Idempotent: a retried persist of the SAME
|
||||
* (message, stepIndex) is a no-op via ON CONFLICT DO NOTHING — the per-step
|
||||
* writes are fired fire-and-forget + serialized, and a duplicate must never
|
||||
* throw into the stream or double the parts. Returns whether a NEW row landed
|
||||
* (false = the step was already persisted).
|
||||
*/
|
||||
async insertStep(
|
||||
messageId: string,
|
||||
workspaceId: string,
|
||||
stepIndex: number,
|
||||
parts: unknown,
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<boolean> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
const inserted = await db
|
||||
.insertInto('aiChatRunSteps')
|
||||
.values({
|
||||
messageId,
|
||||
workspaceId,
|
||||
stepIndex,
|
||||
// jsonb column: cast through never (same pattern as the message repo).
|
||||
parts: parts as never,
|
||||
})
|
||||
.onConflict((oc) => oc.columns(['messageId', 'stepIndex']).doNothing())
|
||||
.returning('id')
|
||||
.executeTakeFirst();
|
||||
return inserted !== undefined;
|
||||
}
|
||||
|
||||
/** All persisted steps for ONE assistant message, in step order. */
|
||||
async findByMessage(
|
||||
messageId: string,
|
||||
workspaceId: string,
|
||||
): Promise<AiChatRunStep[]> {
|
||||
return this.db
|
||||
.selectFrom('aiChatRunSteps')
|
||||
.selectAll('aiChatRunSteps')
|
||||
.where('messageId', '=', messageId)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.orderBy('stepIndex', 'asc')
|
||||
.execute();
|
||||
}
|
||||
|
||||
/**
|
||||
* All persisted steps for a SET of assistant messages, grouped by messageId
|
||||
* (each group in step order). One query for the batch — the hydration seam
|
||||
* (getMessages / delta / export) calls this only for the rows that actually
|
||||
* need reconstruction (an active new-style row whose `metadata.parts` is still
|
||||
* empty), which is usually none, so this is skipped on the common path.
|
||||
*/
|
||||
async findByMessageIds(
|
||||
messageIds: string[],
|
||||
workspaceId: string,
|
||||
): Promise<Map<string, AiChatRunStep[]>> {
|
||||
const byMessage = new Map<string, AiChatRunStep[]>();
|
||||
if (messageIds.length === 0) return byMessage;
|
||||
const rows = await this.db
|
||||
.selectFrom('aiChatRunSteps')
|
||||
.selectAll('aiChatRunSteps')
|
||||
.where('messageId', 'in', messageIds)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.orderBy('stepIndex', 'asc')
|
||||
.execute();
|
||||
for (const row of rows) {
|
||||
const list = byMessage.get(row.messageId);
|
||||
if (list) list.push(row);
|
||||
else byMessage.set(row.messageId, [row]);
|
||||
}
|
||||
return byMessage;
|
||||
}
|
||||
|
||||
/**
|
||||
* How many steps are persisted for a message (its step-marker floor). Exposed
|
||||
* for the reconstruct contract (`reconstructRunParts → { parts, stepsPersisted }`)
|
||||
* so a caller can align a resume attach without materializing every step's parts.
|
||||
*/
|
||||
async countByMessage(messageId: string, workspaceId: string): Promise<number> {
|
||||
const row = await this.db
|
||||
.selectFrom('aiChatRunSteps')
|
||||
.select(sql<number>`count(*)::int`.as('n'))
|
||||
.where('messageId', '=', messageId)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.executeTakeFirst();
|
||||
return row?.n ?? 0;
|
||||
}
|
||||
}
|
||||
+17
@@ -692,6 +692,22 @@ export interface AiChatRuns {
|
||||
updatedAt: Generated<Timestamp>;
|
||||
}
|
||||
|
||||
// Append-only per-step persistence for an assistant turn (#492). Mirrors
|
||||
// migration 20260708T120000-ai-chat-run-steps.ts. Each finished agent step's UI
|
||||
// `parts` are INSERTed as their own row (instead of rewriting the message row's
|
||||
// growing `metadata.parts` jsonb every step — an O(n²) WAL/TOAST churn). The full
|
||||
// `metadata.parts` is assembled once at finalize; mid-run a resuming client's seed
|
||||
// is rebuilt by concatenating these rows in `stepIndex` order. Cascades with the
|
||||
// assistant message row it projects.
|
||||
export interface AiChatRunSteps {
|
||||
id: Generated<string>;
|
||||
messageId: string;
|
||||
workspaceId: string;
|
||||
stepIndex: number;
|
||||
parts: Json;
|
||||
createdAt: 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
|
||||
@@ -729,6 +745,7 @@ export interface DB {
|
||||
aiChats: AiChats;
|
||||
aiChatMessages: AiChatMessages;
|
||||
aiChatRuns: AiChatRuns;
|
||||
aiChatRunSteps: AiChatRunSteps;
|
||||
aiChatPageSnapshots: AiChatPageSnapshots;
|
||||
apiKeys: ApiKeys;
|
||||
attachments: Attachments;
|
||||
|
||||
@@ -4,6 +4,7 @@ import {
|
||||
AiChats,
|
||||
AiChatMessages,
|
||||
AiChatRuns,
|
||||
AiChatRunSteps,
|
||||
AiChatPageSnapshots,
|
||||
Attachments,
|
||||
Comments,
|
||||
@@ -64,6 +65,12 @@ export type InsertableAiChatMessage = Omit<Insertable<AiChatMessages>, 'tsv'>;
|
||||
export type AiChatRun = Selectable<AiChatRuns>;
|
||||
export type InsertableAiChatRun = Insertable<AiChatRuns>;
|
||||
|
||||
// AI Chat Run Step (#492): append-only per-step parts persistence. Each finished
|
||||
// agent step's UI parts are stored as their own row; the full turn's parts are
|
||||
// assembled from these (in stepIndex order) for a mid-run resume seed.
|
||||
export type AiChatRunStep = Selectable<AiChatRunSteps>;
|
||||
export type InsertableAiChatRunStep = Insertable<AiChatRunSteps>;
|
||||
|
||||
// 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
|
||||
// detect human edits made between turns.
|
||||
|
||||
@@ -0,0 +1,173 @@
|
||||
import { randomBytes } from 'crypto';
|
||||
import { Kysely, sql } from 'kysely';
|
||||
import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo';
|
||||
import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
|
||||
import {
|
||||
assistantParts,
|
||||
flushAssistant,
|
||||
stepMarkerMetadata,
|
||||
} from '../../src/core/ai-chat/ai-chat.service';
|
||||
import {
|
||||
getTestDb,
|
||||
destroyTestDb,
|
||||
createWorkspace,
|
||||
createUser,
|
||||
createChat,
|
||||
} from './db';
|
||||
|
||||
/**
|
||||
* #492 append-persist — WRITE-VOLUME regression on a LIVE Postgres, measured via
|
||||
* the `pg_current_wal_lsn()` delta around a realistic multi-step run driven through
|
||||
* the REAL repos (not a mock — a mock cannot observe MVCC/TOAST rewrite volume, the
|
||||
* whole point). Proves the core claim:
|
||||
*
|
||||
* NEW (per-step INSERT into ai_chat_run_steps + a CHEAP step-marker UPDATE on the
|
||||
* message row) writes O(Σ steps) of WAL — each step writes only its own bytes.
|
||||
*
|
||||
* OLD (the pre-#492 full-row rewrite: re-persist the GROWING metadata.parts on
|
||||
* every onStepFinish) writes O(n²) — step k rewrites the whole TOASTed jsonb of
|
||||
* all k prior outputs.
|
||||
*
|
||||
* The OLD path here IS the reverted behavior, so this doubles as the mutation
|
||||
* check: swapping the new path back to `flushAssistant` full-row UPDATEs reddens
|
||||
* the assertion (OLD is many times larger).
|
||||
*/
|
||||
type Step = {
|
||||
text: string;
|
||||
toolCalls: Array<{ toolCallId: string; toolName: string; input: unknown }>;
|
||||
toolResults: Array<{ toolCallId: string; toolName: string; output: unknown }>;
|
||||
};
|
||||
|
||||
// ~100 KB INCOMPRESSIBLE output per step (a page read). Random base64 so TOAST
|
||||
// cannot compress it away and hide the real write volume.
|
||||
function makeStep(i: number, outputBytes = 100_000): Step {
|
||||
const body = randomBytes(Math.ceil(outputBytes * 0.75)).toString('base64');
|
||||
return {
|
||||
text: `step ${i} reasoning`,
|
||||
toolCalls: [
|
||||
{ toolCallId: `c${i}`, toolName: 'getPage', input: { id: `p${i}` } },
|
||||
],
|
||||
toolResults: [
|
||||
{
|
||||
toolCallId: `c${i}`,
|
||||
toolName: 'getPage',
|
||||
output: { id: `p${i}`, title: `Page ${i}`, body },
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
async function walDelta(
|
||||
db: Kysely<any>,
|
||||
fn: () => Promise<void>,
|
||||
): Promise<number> {
|
||||
const before = (
|
||||
await sql<{ l: string }>`select pg_current_wal_lsn() as l`.execute(db)
|
||||
).rows[0].l;
|
||||
await fn();
|
||||
// NOTE: no pg_switch_wal() — a segment switch pads the LSN to the next 16 MB
|
||||
// boundary and would swamp the delta. The raw LSN advances by the WAL bytes.
|
||||
const after = (
|
||||
await sql<{ l: string }>`select pg_current_wal_lsn() as l`.execute(db)
|
||||
).rows[0].l;
|
||||
return Number(
|
||||
(
|
||||
await sql<{
|
||||
d: string;
|
||||
}>`select pg_wal_lsn_diff(${after}::pg_lsn, ${before}::pg_lsn) as d`.execute(
|
||||
db,
|
||||
)
|
||||
).rows[0].d,
|
||||
);
|
||||
}
|
||||
|
||||
describe('#492 append-persist write volume (pg_current_wal_lsn delta) [integration]', () => {
|
||||
let db: Kysely<any>;
|
||||
let stepRepo: AiChatRunStepRepo;
|
||||
let msgRepo: AiChatMessageRepo;
|
||||
let workspaceId: string;
|
||||
let userId: string;
|
||||
let chatId: string;
|
||||
|
||||
beforeAll(async () => {
|
||||
db = getTestDb();
|
||||
stepRepo = new AiChatRunStepRepo(db as any);
|
||||
msgRepo = new AiChatMessageRepo(db as any);
|
||||
workspaceId = (await createWorkspace(db)).id;
|
||||
userId = (await createUser(db, workspaceId)).id;
|
||||
chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await destroyTestDb();
|
||||
});
|
||||
|
||||
const seedRow = () =>
|
||||
msgRepo.insert({
|
||||
chatId,
|
||||
workspaceId,
|
||||
userId,
|
||||
role: 'assistant',
|
||||
content: '',
|
||||
status: 'streaming',
|
||||
metadata: stepMarkerMetadata(0) as never,
|
||||
});
|
||||
|
||||
const STEPS = 40;
|
||||
|
||||
it('NEW per-step INSERT is O(Σ steps); OLD full-row rewrite is O(n²)', async () => {
|
||||
const steps: Step[] = [];
|
||||
for (let i = 0; i < STEPS; i++) steps.push(makeStep(i));
|
||||
|
||||
// NEW: per-step INSERT of THIS step's parts + a cheap marker UPDATE.
|
||||
const newRow = await seedRow();
|
||||
const newWal = await walDelta(db, async () => {
|
||||
for (let i = 0; i < STEPS; i++) {
|
||||
await stepRepo.insertStep(
|
||||
newRow.id,
|
||||
workspaceId,
|
||||
i,
|
||||
assistantParts([steps[i]], ''),
|
||||
);
|
||||
await msgRepo.update(
|
||||
newRow.id,
|
||||
workspaceId,
|
||||
{ metadata: stepMarkerMetadata(i + 1) },
|
||||
{ onlyIfStreaming: true },
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
// OLD (the pre-#492 revert): re-persist the GROWING metadata.parts on the
|
||||
// message row on every step.
|
||||
const oldRow = await seedRow();
|
||||
const oldWal = await walDelta(db, async () => {
|
||||
const acc: Step[] = [];
|
||||
for (let i = 0; i < STEPS; i++) {
|
||||
acc.push(steps[i]);
|
||||
await msgRepo.update(
|
||||
oldRow.id,
|
||||
workspaceId,
|
||||
flushAssistant(acc as never, '', 'streaming'),
|
||||
{ onlyIfStreaming: true },
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
// eslint-disable-next-line no-console
|
||||
console.log(
|
||||
`[#492 WAL] ${STEPS} steps ×100KB: new=${(newWal / 1e6).toFixed(1)}MB ` +
|
||||
`old=${(oldWal / 1e6).toFixed(1)}MB (${(oldWal / newWal).toFixed(
|
||||
1,
|
||||
)}x smaller)`,
|
||||
);
|
||||
|
||||
// O(Σ steps): ~STEPS × (100KB output + marker) of WAL. 40 × ~100KB parts plus
|
||||
// 40 tiny markers is a few tens of MB at most — bounded, linear in step count.
|
||||
expect(newWal).toBeLessThan(30_000_000);
|
||||
// O(n²): step k rewrites ~k × 100KB. Σ over 40 steps ≈ 80+ MB — far larger.
|
||||
expect(oldWal).toBeGreaterThan(30_000_000);
|
||||
// The load-bearing claim: the new path writes a small FRACTION of the old.
|
||||
expect(newWal).toBeLessThan(oldWal * 0.35);
|
||||
}, 120_000);
|
||||
});
|
||||
@@ -0,0 +1,164 @@
|
||||
import { randomBytes } from 'crypto';
|
||||
import { Kysely } from 'kysely';
|
||||
import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo';
|
||||
import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
|
||||
import {
|
||||
assistantParts,
|
||||
reconstructRunParts,
|
||||
hydrateAssistantParts,
|
||||
stepMarkerMetadata,
|
||||
rowHasInlineParts,
|
||||
} from '../../src/core/ai-chat/ai-chat.service';
|
||||
import {
|
||||
getTestDb,
|
||||
destroyTestDb,
|
||||
createWorkspace,
|
||||
createUser,
|
||||
createChat,
|
||||
} from './db';
|
||||
|
||||
/**
|
||||
* #492 append-persist — the reconstruct CONTRACT on a live Postgres. Proves that a
|
||||
* turn persisted the NEW way (per-step rows in `ai_chat_run_steps`, only a step
|
||||
* marker on the message row) reconstructs to the SAME UI parts as a turn persisted
|
||||
* the OLD way (full `metadata.parts` inline on the row, no step rows) — so the
|
||||
* era-switch is invisible to attach / delta-poll / export. Real repos + real jsonb
|
||||
* roundtrip, not a mock (a mock cannot prove the parts survive the jsonb column
|
||||
* byte-identical).
|
||||
*/
|
||||
type Step = {
|
||||
text: string;
|
||||
toolCalls: Array<{ toolCallId: string; toolName: string; input: unknown }>;
|
||||
toolResults: Array<{ toolCallId: string; toolName: string; output: unknown }>;
|
||||
};
|
||||
|
||||
// A realistic step: some text + a getPage tool call whose ~100 KB body is
|
||||
// INCOMPRESSIBLE random base64 (a 'x'.repeat filler would TOAST away and hide the
|
||||
// real bytes). Under MAX_TOOL_OUTPUT_BYTES (200 KB) it is stored uncompacted.
|
||||
function makeStep(i: number, outputBytes = 4_000): Step {
|
||||
const body = randomBytes(Math.ceil(outputBytes * 0.75)).toString('base64');
|
||||
return {
|
||||
text: `step ${i} text`,
|
||||
toolCalls: [
|
||||
{ toolCallId: `c${i}`, toolName: 'getPage', input: { id: `p${i}` } },
|
||||
],
|
||||
toolResults: [
|
||||
{
|
||||
toolCallId: `c${i}`,
|
||||
toolName: 'getPage',
|
||||
output: { id: `p${i}`, title: `Page ${i}`, body },
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
describe('AiChatRunStepRepo + reconstruct contract [integration]', () => {
|
||||
let db: Kysely<any>;
|
||||
let stepRepo: AiChatRunStepRepo;
|
||||
let msgRepo: AiChatMessageRepo;
|
||||
let workspaceId: string;
|
||||
let userId: string;
|
||||
let chatId: string;
|
||||
|
||||
beforeAll(async () => {
|
||||
db = getTestDb();
|
||||
stepRepo = new AiChatRunStepRepo(db as any);
|
||||
msgRepo = new AiChatMessageRepo(db as any);
|
||||
workspaceId = (await createWorkspace(db)).id;
|
||||
userId = (await createUser(db, workspaceId)).id;
|
||||
chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await destroyTestDb();
|
||||
});
|
||||
|
||||
const seedRow = (metadata: unknown, status: string) =>
|
||||
msgRepo.insert({
|
||||
chatId,
|
||||
workspaceId,
|
||||
userId,
|
||||
role: 'assistant',
|
||||
content: '',
|
||||
status,
|
||||
metadata: metadata as never,
|
||||
});
|
||||
|
||||
it('insertStep is idempotent per (message, stepIndex) and reads back in order', async () => {
|
||||
const row = await seedRow(stepMarkerMetadata(0), 'streaming');
|
||||
const parts0 = assistantParts([makeStep(0)], '');
|
||||
const parts1 = assistantParts([makeStep(1)], '');
|
||||
|
||||
expect(await stepRepo.insertStep(row.id, workspaceId, 0, parts0)).toBe(true);
|
||||
expect(await stepRepo.insertStep(row.id, workspaceId, 1, parts1)).toBe(true);
|
||||
// A retried persist of the SAME step is a no-op (ON CONFLICT DO NOTHING).
|
||||
expect(await stepRepo.insertStep(row.id, workspaceId, 0, parts0)).toBe(
|
||||
false,
|
||||
);
|
||||
|
||||
const steps = await stepRepo.findByMessage(row.id, workspaceId);
|
||||
expect(steps.map((s) => s.stepIndex)).toEqual([0, 1]);
|
||||
expect(await stepRepo.countByMessage(row.id, workspaceId)).toBe(2);
|
||||
|
||||
// Batch fetch groups by message id in step order.
|
||||
const map = await stepRepo.findByMessageIds([row.id], workspaceId);
|
||||
expect(map.get(row.id)!.map((s) => s.stepIndex)).toEqual([0, 1]);
|
||||
});
|
||||
|
||||
it('a NEW-style (step-table) run reconstructs identically to an OLD-style (inline) run', async () => {
|
||||
const steps = [makeStep(10), makeStep(11)];
|
||||
// The inline parts the OLD full-row flush would have written.
|
||||
const fullParts = assistantParts(steps, '');
|
||||
|
||||
// OLD-style record: full parts inline on the row, NO step rows.
|
||||
const oldRow = await seedRow(
|
||||
{ parts: fullParts, toolTraceVersion: 2, stepsPersisted: 2 },
|
||||
'completed',
|
||||
);
|
||||
|
||||
// NEW-style record: only a step marker on the row + per-step rows.
|
||||
const newRow = await seedRow(stepMarkerMetadata(2), 'streaming');
|
||||
for (let i = 0; i < steps.length; i++) {
|
||||
await stepRepo.insertStep(
|
||||
newRow.id,
|
||||
workspaceId,
|
||||
i,
|
||||
assistantParts([steps[i]], ''),
|
||||
);
|
||||
}
|
||||
|
||||
// Re-read both from the DB (proves the jsonb roundtrip).
|
||||
const oldFetched = await msgRepo.findById(oldRow.id, workspaceId);
|
||||
const newFetched = await msgRepo.findById(newRow.id, workspaceId);
|
||||
const oldSteps = await stepRepo.findByMessage(oldRow.id, workspaceId);
|
||||
const newSteps = await stepRepo.findByMessage(newRow.id, workspaceId);
|
||||
|
||||
// The discriminator: the old row carries inline parts, the new one does not.
|
||||
expect(rowHasInlineParts(oldFetched!)).toBe(true);
|
||||
expect(rowHasInlineParts(newFetched!)).toBe(false);
|
||||
expect(oldSteps).toHaveLength(0);
|
||||
expect(newSteps).toHaveLength(2);
|
||||
|
||||
const oldRecon = reconstructRunParts(oldFetched!, oldSteps);
|
||||
const newRecon = reconstructRunParts(newFetched!, newSteps);
|
||||
|
||||
// Both reconstruct to the SAME parts + step count — the era is invisible.
|
||||
expect(newRecon.parts).toEqual(fullParts);
|
||||
expect(oldRecon.parts).toEqual(fullParts);
|
||||
expect(newRecon.parts).toEqual(oldRecon.parts);
|
||||
expect(newRecon.stepsPersisted).toBe(2);
|
||||
expect(oldRecon.stepsPersisted).toBe(2);
|
||||
|
||||
// hydrateAssistantParts fills the new row's metadata.parts to match the old
|
||||
// row's inline parts — so a consumer reading `metadata.parts` off the raw row
|
||||
// (the client seed/poll, export) is unchanged across the era.
|
||||
const map = await stepRepo.findByMessageIds([newRow.id], workspaceId);
|
||||
const [hydrated] = hydrateAssistantParts([newFetched!], map);
|
||||
expect((hydrated.metadata as { parts: unknown }).parts).toEqual(fullParts);
|
||||
// A row that already has inline parts passes through untouched (same ref-shape).
|
||||
const [oldPassThrough] = hydrateAssistantParts([oldFetched!], map);
|
||||
expect((oldPassThrough.metadata as { parts: unknown }).parts).toEqual(
|
||||
fullParts,
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,63 @@
|
||||
import { Kysely, sql } from 'kysely';
|
||||
import {
|
||||
up,
|
||||
down,
|
||||
} from '../../src/database/migrations/20260708T120000-ai-chat-run-steps';
|
||||
import { getTestDb, destroyTestDb } from './db';
|
||||
|
||||
/**
|
||||
* #492 migration up/down roundtrip on a LIVE Postgres. global-setup already
|
||||
* migrated docmost_test to latest (so the table exists at start); this drives the
|
||||
* migration's own down()/up() and asserts the table presence toggles, then leaves
|
||||
* it PRESENT (up) so the shared test DB is intact for any spec that runs after.
|
||||
*/
|
||||
async function tableExists(db: Kysely<any>): Promise<boolean> {
|
||||
const row = (
|
||||
await sql<{ t: string | null }>`select to_regclass('ai_chat_run_steps') as t`.execute(
|
||||
db,
|
||||
)
|
||||
).rows[0];
|
||||
return row.t !== null;
|
||||
}
|
||||
|
||||
async function uniqueIndexExists(db: Kysely<any>): Promise<boolean> {
|
||||
const row = (
|
||||
await sql<{
|
||||
t: string | null;
|
||||
}>`select to_regclass('ai_chat_run_steps_message_step_uidx') as t`.execute(db)
|
||||
).rows[0];
|
||||
return row.t !== null;
|
||||
}
|
||||
|
||||
describe('20260708 ai_chat_run_steps migration roundtrip [integration]', () => {
|
||||
let db: Kysely<any>;
|
||||
|
||||
beforeAll(() => {
|
||||
db = getTestDb();
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
// Belt-and-suspenders: guarantee the table is present for later specs even if
|
||||
// an assertion threw mid-roundtrip.
|
||||
if (!(await tableExists(db))) await up(db);
|
||||
await destroyTestDb();
|
||||
});
|
||||
|
||||
it('down() drops the table+index and up() recreates them (idempotent)', async () => {
|
||||
// Starts applied (global-setup migrated to latest).
|
||||
expect(await tableExists(db)).toBe(true);
|
||||
expect(await uniqueIndexExists(db)).toBe(true);
|
||||
|
||||
await down(db);
|
||||
expect(await tableExists(db)).toBe(false);
|
||||
expect(await uniqueIndexExists(db)).toBe(false);
|
||||
|
||||
await up(db);
|
||||
expect(await tableExists(db)).toBe(true);
|
||||
expect(await uniqueIndexExists(db)).toBe(true);
|
||||
|
||||
// up() is idempotent (ifNotExists) — a second run is a harmless no-op.
|
||||
await expect(up(db)).resolves.not.toThrow();
|
||||
expect(await tableExists(db)).toBe(true);
|
||||
});
|
||||
});
|
||||
@@ -31,11 +31,7 @@ import { TransformsMixin, type ITransformsMixin } from "./client/transforms.js";
|
||||
// existing importer (index.ts, http.ts, stdio.ts, the in-app host) keeps working
|
||||
// with ZERO changes.
|
||||
export type { DocmostMcpConfig, SandboxPut } from "./client/context.js";
|
||||
export {
|
||||
formatDocmostAxiosError,
|
||||
assertFullUuid,
|
||||
formatSpaceNotAccessible,
|
||||
} from "./client/errors.js";
|
||||
export { formatDocmostAxiosError, assertFullUuid } from "./client/errors.js";
|
||||
|
||||
// Branded canonical page-identity type (#435): the internal page UUID is a
|
||||
// distinct nominal type so an unresolved raw/slug string can't be swapped into
|
||||
|
||||
@@ -725,15 +725,10 @@ export function CommentsMixin<TBase extends GConstructor<DocmostClientContext>>(
|
||||
// The subtree scope (parentPageId given) already INCLUDES the root node
|
||||
// itself: /pages/tree seeds getPageAndDescendants with id = parentPageId, so
|
||||
// no separate getPageRaw fetch for the parent is needed.
|
||||
// #534: the enumerateSpacePages seed (`/pages/tree`) 404s for a bad or
|
||||
// inaccessible spaceId; wrap it so that 404 becomes an actionable "spaceId
|
||||
// not accessible" hint instead of the opaque "Space permissions not found".
|
||||
// Only the whole enumeration is wrapped (the only 404 source here) — see the
|
||||
// wrap-allowlist invariant on withSpaceAccessDiagnostics.
|
||||
const { pages: pagesInScope, truncated } =
|
||||
await this.withSpaceAccessDiagnostics(spaceId, "checkNewComments", () =>
|
||||
this.enumerateSpacePages(spaceId, parentPageId),
|
||||
);
|
||||
const { pages: pagesInScope, truncated } = await this.enumerateSpacePages(
|
||||
spaceId,
|
||||
parentPageId,
|
||||
);
|
||||
|
||||
// 2. Fetch comments for each page, keep ones created after since. Runs with
|
||||
// bounded concurrency (#490) instead of one-at-a-time — the per-page reads are
|
||||
|
||||
@@ -26,10 +26,7 @@ import {
|
||||
import { withPageLock, isUuid } from "../lib/page-lock.js";
|
||||
import type { PageId } from "../lib/page-id.js";
|
||||
import { getCollabToken, performLogin } from "../lib/auth-utils.js";
|
||||
import {
|
||||
formatDocmostAxiosError,
|
||||
formatSpaceNotAccessible,
|
||||
} from "./errors.js";
|
||||
import { formatDocmostAxiosError } from "./errors.js";
|
||||
import { GetPageConversionCache } from "./getpage-cache.js";
|
||||
|
||||
// A generic mixin base constructor (issue #450). Each domain mixin is a factory
|
||||
@@ -120,36 +117,6 @@ function readCollabTokenTtlMs(): number {
|
||||
return Number.isFinite(raw) ? Math.max(0, raw) : 5 * 60 * 1000;
|
||||
}
|
||||
|
||||
/**
|
||||
* Accessible-space index cache TTL in milliseconds (issue #534). Read fresh from
|
||||
* the environment on every access — mirroring readCollabTokenTtlMs above — so a
|
||||
* test or a live rollback can change it without reloading the module.
|
||||
*
|
||||
* The index (see getAccessibleSpaceIndex) is fetched ONLY on the enrich-on-404
|
||||
* slow path to turn an opaque "Space permissions not found" 404 into a factual
|
||||
* "spaceId X is not among your accessible spaces" hint; a short TTL keeps a burst
|
||||
* of failing tool calls from re-sweeping /spaces each time while never widening
|
||||
* the permission-staleness window meaningfully. Default 60s. An EXPLICIT 0 (or
|
||||
* negative) DISABLES the cache (exact fetch-per-enrichment). Unset/unparseable
|
||||
* (NaN) falls back to the 60s default with the cache ON.
|
||||
*/
|
||||
function readSpacesCacheTtlMs(): number {
|
||||
const raw = parseInt(process.env.MCP_SPACES_CACHE_TTL_MS ?? "", 10);
|
||||
return Number.isFinite(raw) ? Math.max(0, raw) : 60000;
|
||||
}
|
||||
|
||||
/**
|
||||
* The set of spaces the current token can see, plus a `complete` flag that is
|
||||
* false when the /spaces listing was truncated at the pagination ceiling. Used
|
||||
* by the enrich-on-404 diagnostics: an authoritative membership test is only
|
||||
* possible when `complete` is true (see withSpaceAccessDiagnostics).
|
||||
*/
|
||||
export type AccessibleSpaceIndex = {
|
||||
ids: Set<string>;
|
||||
spaces: { id: string; name: string }[];
|
||||
complete: boolean;
|
||||
};
|
||||
|
||||
export abstract class DocmostClientContext {
|
||||
protected client: AxiosInstance;
|
||||
protected token: string | null = null;
|
||||
@@ -197,27 +164,6 @@ 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;
|
||||
|
||||
// Accessible-space index cache + single-flight (issue #534). TWO separate
|
||||
// fields, mirroring loginPromise (in-flight dedup) vs collabTokenCache
|
||||
// (persistent value):
|
||||
// - spaceIndexCache: the last SUCCESSFULLY-FETCHED, COMPLETE index plus the
|
||||
// wall-clock time it was fetched. Written ONLY from a resolved /spaces
|
||||
// sweep whose result was complete (a truncated list is never cached, since
|
||||
// it cannot answer "is this spaceId missing?"). Per-instance (a
|
||||
// DocmostClient is built per user / per chat) so it can never leak across
|
||||
// identities; invalidated on every identity change exactly like
|
||||
// collabTokenCache (login() + the 401/403 reauth interceptor).
|
||||
// - spaceIndexInFlight: dedups concurrent enrich-on-404 fetches into ONE
|
||||
// /spaces sweep. CRITICAL INVARIANT: this promise is nulled in `.finally`
|
||||
// on BOTH resolve AND reject — a rejected/settled promise is NEVER
|
||||
// memoized, so a transient /spaces blip during one failed tool call cannot
|
||||
// poison the diagnostics for the rest of the session.
|
||||
protected spaceIndexCache: {
|
||||
index: AccessibleSpaceIndex;
|
||||
fetchedAt: number;
|
||||
} | null = null;
|
||||
protected spaceIndexInFlight: Promise<AccessibleSpaceIndex> | null = null;
|
||||
|
||||
// Content-addressed conversion cache for getPage (issue #479). Keyed on
|
||||
// (canonical pageId, updatedAt, optionsHash) -> the converted Markdown, so a
|
||||
// re-read of an UNCHANGED page skips the expensive convertProseMirrorToMarkdown
|
||||
@@ -335,9 +281,6 @@ export abstract class DocmostClientContext {
|
||||
// keep serving a collab token minted under the old one.
|
||||
this.token = null;
|
||||
this.collabTokenCache = null;
|
||||
// #534: a new identity/login must not keep serving a space index
|
||||
// computed under the old token (same reasoning as collabTokenCache).
|
||||
this.spaceIndexCache = null;
|
||||
delete this.client.defaults.headers.common["Authorization"];
|
||||
try {
|
||||
await this.login();
|
||||
@@ -470,8 +413,6 @@ export abstract class DocmostClientContext {
|
||||
// Identity (re)established: drop any collab token minted under a
|
||||
// previous identity so the #435 cache can never outlive it.
|
||||
this.collabTokenCache = null;
|
||||
// #534: likewise drop the accessible-space index of the old identity.
|
||||
this.spaceIndexCache = null;
|
||||
this.client.defaults.headers.common["Authorization"] =
|
||||
`Bearer ${token}`;
|
||||
})
|
||||
@@ -689,34 +630,13 @@ export abstract class DocmostClientContext {
|
||||
}
|
||||
|
||||
/**
|
||||
* Generic pagination handler for Docmost API endpoints. Thin wrapper over
|
||||
* paginateAllWithMeta that discards the `truncated` flag — the historical
|
||||
* contract every caller (getSpaces, etc.) relies on. Callers that need to KNOW
|
||||
* whether the result set was complete (e.g. #534's getAccessibleSpaceIndex,
|
||||
* which must not assert "spaceId missing" against a truncated list) call
|
||||
* paginateAllWithMeta directly.
|
||||
* Generic pagination handler for Docmost API endpoints
|
||||
*/
|
||||
async paginateAll<T = any>(
|
||||
endpoint: string,
|
||||
basePayload: Record<string, any> = {},
|
||||
limit: number = 100,
|
||||
): Promise<T[]> {
|
||||
return (await this.paginateAllWithMeta<T>(endpoint, basePayload, limit))
|
||||
.items;
|
||||
}
|
||||
|
||||
/**
|
||||
* Generic pagination handler that ALSO surfaces whether the result was
|
||||
* truncated at the MAX_PAGES ceiling. `paginateAll` swallows this flag (it only
|
||||
* warns); callers that must distinguish "complete listing" from "gave up at the
|
||||
* cap" use this overload. `truncated` is true iff the loop stopped at the
|
||||
* ceiling while the server still reported more pages.
|
||||
*/
|
||||
async paginateAllWithMeta<T = any>(
|
||||
endpoint: string,
|
||||
basePayload: Record<string, any> = {},
|
||||
limit: number = 100,
|
||||
): Promise<{ items: T[]; truncated: boolean }> {
|
||||
await this.ensureAuthenticated();
|
||||
|
||||
const clampedLimit = Math.max(1, Math.min(100, limit));
|
||||
@@ -777,137 +697,7 @@ export abstract class DocmostClientContext {
|
||||
);
|
||||
}
|
||||
|
||||
return { items: allItems, truncated };
|
||||
}
|
||||
|
||||
/**
|
||||
* The set of spaces the current token can access (issue #534), fetched from the
|
||||
* single source of truth — the `/spaces` listing — with a per-instance
|
||||
* short-TTL cache and single-flight dedup. Used ONLY by the enrich-on-404 slow
|
||||
* path (withSpaceAccessDiagnostics), so the happy path incurs ZERO extra
|
||||
* requests.
|
||||
*
|
||||
* `complete` is `!truncated`: it is false when the /spaces listing was cut at
|
||||
* the pagination ceiling. A truncated index can never authoritatively answer
|
||||
* "is this spaceId missing?", so only a complete result is cached AND only a
|
||||
* complete result is allowed to drive the "not accessible" rewrite.
|
||||
*
|
||||
* Cache/single-flight discipline (see the spaceIndexCache / spaceIndexInFlight
|
||||
* field docs):
|
||||
* - serve a fresh, complete cached index without any request;
|
||||
* - otherwise collapse concurrent callers onto ONE in-flight /spaces sweep;
|
||||
* - write the persistent cache ONLY from a resolved, complete fetch;
|
||||
* - null the in-flight promise on BOTH resolve and reject (never memoize a
|
||||
* rejected promise — a transient /spaces failure must be retried fresh).
|
||||
*/
|
||||
async getAccessibleSpaceIndex(): Promise<AccessibleSpaceIndex> {
|
||||
const ttl = readSpacesCacheTtlMs();
|
||||
|
||||
// Fast path: a still-fresh, complete cached index needs no request at all.
|
||||
if (
|
||||
ttl > 0 &&
|
||||
this.spaceIndexCache &&
|
||||
Date.now() - this.spaceIndexCache.fetchedAt < ttl
|
||||
) {
|
||||
return this.spaceIndexCache.index;
|
||||
}
|
||||
|
||||
// Single-flight: a concurrent enrichment joins the in-flight sweep instead of
|
||||
// issuing its own. (A settled/rejected promise is never left here — see the
|
||||
// `.finally` below — so this only ever joins a genuinely in-progress fetch.)
|
||||
if (this.spaceIndexInFlight) return this.spaceIndexInFlight;
|
||||
|
||||
const fetchPromise = (async (): Promise<AccessibleSpaceIndex> => {
|
||||
const { items, truncated } = await this.paginateAllWithMeta("/spaces", {});
|
||||
const spaces = items.map((s: any) => ({
|
||||
id: s?.id,
|
||||
name: s?.name,
|
||||
}));
|
||||
return {
|
||||
ids: new Set(spaces.map((s) => s.id)),
|
||||
spaces,
|
||||
complete: !truncated,
|
||||
};
|
||||
})();
|
||||
|
||||
this.spaceIndexInFlight = fetchPromise
|
||||
.then((index) => {
|
||||
// Cache ONLY a complete result, and only while the cache is enabled.
|
||||
if (ttl > 0 && index.complete) {
|
||||
this.spaceIndexCache = { index, fetchedAt: Date.now() };
|
||||
}
|
||||
return index;
|
||||
})
|
||||
.finally(() => {
|
||||
// CRITICAL (#534): clear the in-flight slot on BOTH resolve and reject.
|
||||
// Nulling on reject too means a transient /spaces error is retried by the
|
||||
// NEXT enrichment with a fresh fetch, never re-serving the rejection.
|
||||
this.spaceIndexInFlight = null;
|
||||
});
|
||||
|
||||
return this.spaceIndexInFlight;
|
||||
}
|
||||
|
||||
/**
|
||||
* Wrap a client method whose 404 means "the supplied spaceId is not accessible"
|
||||
* and, ONLY on that 404, replace the opaque server text ("Space permissions not
|
||||
* found") with a factual, actionable message naming the spaceId and the spaces
|
||||
* the token can actually see (issue #534). A HINT layered on top of the
|
||||
* backend, which stays authoritative — so it FAILS OPEN on ANY uncertainty:
|
||||
* every branch below that is not a confident "this spaceId is genuinely
|
||||
* missing" rethrows the ORIGINAL server error unchanged. The happy path returns
|
||||
* fn()'s value with zero extra requests.
|
||||
*
|
||||
* WRAP-ALLOWLIST INVARIANT (load-bearing — read before wrapping a new method):
|
||||
* among the currently wrapped tools a 404 comes ONLY from the spaceId
|
||||
* membership / space-permissions check — their pageId / rootPageId /
|
||||
* parentPageId branches resolve to 403 or 200, NEVER 404. If a future change
|
||||
* adds a `NotFoundException` to `/pages/tree`, `/pages/recent`,
|
||||
* `/pages/sidebar-pages` or `/search` (e.g. "page not found"), this enrichment
|
||||
* would MISATTRIBUTE that 404 to the spaceId. Re-audit the wrapped call before
|
||||
* relying on this, and only wrap paths where the sole 404 cause is the space.
|
||||
*/
|
||||
protected async withSpaceAccessDiagnostics<T>(
|
||||
spaceId: string,
|
||||
mcpName: string,
|
||||
fn: () => Promise<T>,
|
||||
): Promise<T> {
|
||||
try {
|
||||
return await fn();
|
||||
} catch (e) {
|
||||
// Abort/cap wins FIRST and is detected by the SIGNAL FLAG, not e.name: a
|
||||
// per-call cap may be an AbortSignal.timeout() (reason name "TimeoutError")
|
||||
// or a custom reason, so `e.name === 'AbortError'` is NOT reliable (#534
|
||||
// hole B). A stopped/capped turn must propagate its reason, never trigger a
|
||||
// /spaces sweep or a rewrite.
|
||||
if (this.toolAbortSignal?.aborted) throw e;
|
||||
|
||||
// Only a 404 is enrichable; any other status/shape is a different failure.
|
||||
if (!(axios.isAxiosError(e) && e.response?.status === 404)) throw e;
|
||||
|
||||
let idx: AccessibleSpaceIndex;
|
||||
try {
|
||||
idx = await this.getAccessibleSpaceIndex();
|
||||
} catch (fetchErr) {
|
||||
// The /spaces sweep itself failed. If we were aborted mid-sweep,
|
||||
// propagate the abort reason; otherwise FAIL OPEN with the ORIGINAL
|
||||
// server error rather than a misleading "not found".
|
||||
if (this.toolAbortSignal?.aborted) throw fetchErr;
|
||||
if (process.env.DEBUG) {
|
||||
console.error("space-diag: /spaces fetch failed:", fetchErr);
|
||||
}
|
||||
throw e;
|
||||
}
|
||||
|
||||
// Fail open when the listing is incomplete (can't assert "missing") or when
|
||||
// the spaceId IS present (the 404 is about something else, not the space).
|
||||
if (!idx.complete) throw e;
|
||||
if (idx.ids.has(spaceId)) throw e;
|
||||
|
||||
// Confident: the spaceId is well-formed but not among the accessible
|
||||
// spaces. Replace the opaque server text with the actionable fact.
|
||||
throw new Error(formatSpaceNotAccessible(mcpName, spaceId, idx.spaces));
|
||||
}
|
||||
return allItems;
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -46,60 +46,6 @@ export function assertFullUuid(
|
||||
}
|
||||
}
|
||||
|
||||
// Max number of accessible spaces to enumerate inline in the "space not
|
||||
// accessible" message (issue #534) before collapsing the rest into a "(+N ещё)"
|
||||
// tail, so a workspace with many spaces cannot blow up the model context.
|
||||
const SPACE_LIST_CAP = 10;
|
||||
|
||||
/**
|
||||
* Compose the model-facing "spaceId is not accessible" message (issue #534).
|
||||
* This is FACT text about the supplied spaceId that REPLACES the opaque server
|
||||
* string ("Space permissions not found") on the enrich-on-404 path — it names
|
||||
* the exact bad id, lists the spaces the token can actually see (id + name, so
|
||||
* the agent can copy the right id verbatim), and points at `listSpaces`.
|
||||
*
|
||||
* Deliberately Russian: like the other agent-facing tool guidance in this repo,
|
||||
* this is the message the acting agent reads to self-correct.
|
||||
*/
|
||||
export function formatSpaceNotAccessible(
|
||||
mcpName: string,
|
||||
spaceId: string,
|
||||
spaces: { id: string; name: string }[],
|
||||
): string {
|
||||
// No accessible spaces at all — a distinct diagnosis (token has no space
|
||||
// access), not "you picked the wrong one from this list".
|
||||
if (!Array.isArray(spaces) || spaces.length === 0) {
|
||||
return `${mcpName}: spaceId "${spaceId}" недоступен, и доступных тебе спейсов нет — проверь доступ токена / вызови listSpaces.`;
|
||||
}
|
||||
|
||||
const shown = spaces.slice(0, SPACE_LIST_CAP);
|
||||
const remaining = spaces.length - shown.length;
|
||||
const tail =
|
||||
remaining > 0 ? ` (+${remaining} ещё, см. listSpaces)` : "";
|
||||
|
||||
// Cap the whole message at ERROR_MESSAGE_CAP (same budget as
|
||||
// formatDocmostAxiosError) so ~10 long space names cannot blow up the model
|
||||
// context. Truncate ONLY the interpolated space LIST — the fixed prefix (which
|
||||
// carries the bad spaceId) and the fixed suffix (the "…из listSpaces"
|
||||
// instruction) are always kept intact, so the actionable parts survive even
|
||||
// when the list is trimmed.
|
||||
const prefix = `${mcpName}: spaceId "${spaceId}" не найден среди доступных тебе спейсов. Доступные: `;
|
||||
const suffix = ` — скопируй нужный id дословно из listSpaces.`;
|
||||
let listed = shown.map((s) => `${s.id} (${s.name})`).join(", ") + tail;
|
||||
const budget = ERROR_MESSAGE_CAP - prefix.length - suffix.length;
|
||||
if (listed.length > budget) {
|
||||
listed = listed.slice(0, Math.max(0, budget - 1)) + "…";
|
||||
}
|
||||
const message = `${prefix}${listed}${suffix}`;
|
||||
// Unconditional final backstop (mirrors formatDocmostAxiosError): the list
|
||||
// cap above assumes a well-formed prefix, but a pathologically long
|
||||
// agent-supplied spaceId lives in the prefix and would otherwise blow past the
|
||||
// budget. Cap the WHOLE message so nothing bloats the model context.
|
||||
return message.length > ERROR_MESSAGE_CAP
|
||||
? message.slice(0, ERROR_MESSAGE_CAP - 1) + "…"
|
||||
: message;
|
||||
}
|
||||
|
||||
// Keep ONLY the pathname of a request (no host, no query string, no fragment)
|
||||
// so the message never leaks a host or query params. Resolves a relative
|
||||
// config.url against config.baseURL, then discards everything but the path.
|
||||
|
||||
@@ -132,29 +132,17 @@ export function ReadMixin<TBase extends GConstructor<DocmostClientContext>>(Base
|
||||
// BFS hit its node cap) had no way to know pages were missing. Return the
|
||||
// tree alongside the flag; the primary /pages/tree path is uncapped so this
|
||||
// is false there.
|
||||
// #534: spaceId is required here; wrap so a bad-spaceId 404 (from the
|
||||
// /pages/tree seed inside enumerateSpacePages) becomes an actionable hint.
|
||||
return this.withSpaceAccessDiagnostics(spaceId, "listPages", async () => {
|
||||
const { pages, truncated } = await this.enumerateSpacePages(spaceId);
|
||||
return { tree: buildPageTree(pages), truncated };
|
||||
});
|
||||
const { pages, truncated } = await this.enumerateSpacePages(spaceId);
|
||||
return { tree: buildPageTree(pages), truncated };
|
||||
}
|
||||
|
||||
const clampedLimit = Math.max(1, Math.min(100, limit));
|
||||
const payload: Record<string, any> = { limit: clampedLimit, page: 1 };
|
||||
if (spaceId) payload.spaceId = spaceId;
|
||||
// #534: only the WITH-spaceId recent path can 404 on space access; wrap it so
|
||||
// that 404 is rewritten. Without a spaceId there is no space to diagnose, so
|
||||
// the wrapper is inert (run the request directly).
|
||||
const runRecent = async () => {
|
||||
const response = await this.client.post("/pages/recent", payload);
|
||||
const data = response.data;
|
||||
const items = data.data?.items || data.items || [];
|
||||
return items.map((page: any) => filterPage(page));
|
||||
};
|
||||
return spaceId
|
||||
? this.withSpaceAccessDiagnostics(spaceId, "listPages", runRecent)
|
||||
: runRecent();
|
||||
const response = await this.client.post("/pages/recent", payload);
|
||||
const data = response.data;
|
||||
const items = data.data?.items || data.items || [];
|
||||
return items.map((page: any) => filterPage(page));
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -186,16 +174,8 @@ export function ReadMixin<TBase extends GConstructor<DocmostClientContext>>(Base
|
||||
"getTree: spaceId is required (a page tree is scoped to one space).",
|
||||
);
|
||||
}
|
||||
// #534: the 404 for a bad/inaccessible spaceId surfaces from the
|
||||
// `/pages/tree` seeding step inside enumerateSpacePages (which is NOT in a
|
||||
// try/catch of its own for that case) — wrap the whole body so it is caught
|
||||
// and rewritten into an actionable "spaceId not accessible" message. Only the
|
||||
// space membership check can 404 here (rootPageId 403/200), see the
|
||||
// wrap-allowlist invariant on withSpaceAccessDiagnostics.
|
||||
return this.withSpaceAccessDiagnostics(spaceId, "getTree", async () => {
|
||||
const { pages } = await this.enumerateSpacePages(spaceId, rootPageId);
|
||||
return buildPageTree(pages, { shape: "getTree", maxDepth });
|
||||
});
|
||||
const { pages } = await this.enumerateSpacePages(spaceId, rootPageId);
|
||||
return buildPageTree(pages, { shape: "getTree", maxDepth });
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -720,27 +700,17 @@ export function ReadMixin<TBase extends GConstructor<DocmostClientContext>>(Base
|
||||
if (limit !== undefined) {
|
||||
payload.limit = Math.max(1, Math.min(50, limit));
|
||||
}
|
||||
const response = await this.client.post("/search", payload);
|
||||
|
||||
const runSearch = async () => {
|
||||
const response = await this.client.post("/search", payload);
|
||||
// Normalize both response shapes: bare array and paginated { items: [...] }
|
||||
const data = response.data?.data;
|
||||
const items = Array.isArray(data) ? data : data?.items || [];
|
||||
const filteredItems = items.map((item: any) => filterSearchResult(item));
|
||||
|
||||
// Normalize both response shapes: bare array and paginated { items: [...] }
|
||||
const data = response.data?.data;
|
||||
const items = Array.isArray(data) ? data : data?.items || [];
|
||||
const filteredItems = items.map((item: any) => filterSearchResult(item));
|
||||
|
||||
return {
|
||||
items: filteredItems,
|
||||
success: response.data?.success || false,
|
||||
};
|
||||
return {
|
||||
items: filteredItems,
|
||||
success: response.data?.success || false,
|
||||
};
|
||||
|
||||
// #534: a search scoped to a spaceId 404s when that space is inaccessible;
|
||||
// wrap only that case so the 404 becomes an actionable hint. A workspace-wide
|
||||
// search (no spaceId) has no space to diagnose — run it directly (inert).
|
||||
return spaceId
|
||||
? this.withSpaceAccessDiagnostics(spaceId, "search", runSearch)
|
||||
: runSearch();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -1,423 +0,0 @@
|
||||
// Issue #534: enrich-on-404 space-access diagnostics.
|
||||
//
|
||||
// When a tool is handed a well-formed but non-existent/inaccessible spaceId, the
|
||||
// server answers the space-permissions check with an opaque 404 ("Space
|
||||
// permissions not found"). The client wrapper (withSpaceAccessDiagnostics) turns
|
||||
// ONLY that 404 into an actionable message naming the bad spaceId and the spaces
|
||||
// the token can actually see — while FAILING OPEN (rethrowing the original
|
||||
// server error unchanged) on every source of uncertainty.
|
||||
//
|
||||
// These tests drive the assembled DocmostClient with its inner seams
|
||||
// (enumerateSpacePages / client.post / paginateAllWithMeta) stubbed at runtime,
|
||||
// mirroring the stub-client style of error-diagnostics.test.mjs. Only criterion
|
||||
// 9 (createPage is NOT wrapped) uses a real offline http server, because
|
||||
// createPage's 404 arrives over a bare-axios multipart path.
|
||||
import { test, after } from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
import http from "node:http";
|
||||
import axios, { AxiosError } from "axios";
|
||||
import {
|
||||
DocmostClient,
|
||||
formatSpaceNotAccessible,
|
||||
} from "../../build/client.js";
|
||||
|
||||
// Two accessible spaces (id + name is all getAccessibleSpaceIndex maps/uses).
|
||||
const SPACES = [
|
||||
{ id: "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa", name: "Engineering" },
|
||||
{ id: "bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb", name: "Design" },
|
||||
];
|
||||
// A well-formed UUID that is NOT among the accessible spaces.
|
||||
const BAD = "99999999-9999-4999-8999-999999999999";
|
||||
|
||||
// Build an AxiosError shaped exactly as the response interceptor would hand it
|
||||
// on a 404 — status readable, axios.isAxiosError() true.
|
||||
function makeAxiosErr(status, url = "/pages/tree", message = "boom") {
|
||||
const config = { method: "post", url, baseURL: "http://host.example/api" };
|
||||
const response = {
|
||||
status,
|
||||
statusText: String(status),
|
||||
data: { message },
|
||||
headers: {},
|
||||
config,
|
||||
};
|
||||
return new AxiosError(
|
||||
`Request failed with status code ${status}`,
|
||||
"ERR_BAD_REQUEST",
|
||||
config,
|
||||
{},
|
||||
response,
|
||||
);
|
||||
}
|
||||
function make404(url = "/pages/tree") {
|
||||
return makeAxiosErr(404, url, "Space permissions not found");
|
||||
}
|
||||
|
||||
// A client whose token is pre-set (so ensureAuthenticated never hits the
|
||||
// network) and whose /spaces sweep is a counted stub. Individual tests override
|
||||
// enumerateSpacePages / client.post to shape the method-under-test's outcome.
|
||||
function makeClient({ spaces = SPACES, truncated = false } = {}) {
|
||||
const c = new DocmostClient("http://127.0.0.1:1/api", "u@example.com", "pw");
|
||||
c.token = "t";
|
||||
c.client.defaults.headers.common["Authorization"] = "Bearer t";
|
||||
c._spacesFetches = 0;
|
||||
c.paginateAllWithMeta = async (endpoint) => {
|
||||
if (endpoint === "/spaces") {
|
||||
c._spacesFetches++;
|
||||
return { items: spaces, truncated };
|
||||
}
|
||||
throw new Error(`unexpected paginate endpoint ${endpoint}`);
|
||||
};
|
||||
return c;
|
||||
}
|
||||
|
||||
// --- Criterion 1: bad spaceId on getTree -> actionable rewrite --------------
|
||||
test("getTree with a non-existent spaceId is rewritten into an actionable hint", async () => {
|
||||
const c = makeClient();
|
||||
c.enumerateSpacePages = async () => {
|
||||
throw make404();
|
||||
};
|
||||
|
||||
await assert.rejects(
|
||||
() => c.getTree(BAD),
|
||||
(e) => {
|
||||
assert.ok(e.message.includes(BAD), "names the passed spaceId");
|
||||
// At least one valid "id (name)" pair.
|
||||
assert.ok(
|
||||
e.message.includes(`${SPACES[0].id} (${SPACES[0].name})`),
|
||||
"lists an accessible id (name)",
|
||||
);
|
||||
assert.ok(e.message.includes("listSpaces"), "points at listSpaces");
|
||||
assert.ok(
|
||||
!e.message.includes("Space permissions not found"),
|
||||
"the opaque server text is replaced",
|
||||
);
|
||||
return true;
|
||||
},
|
||||
);
|
||||
assert.equal(c._spacesFetches, 1, "one /spaces sweep on the enrichment path");
|
||||
});
|
||||
|
||||
// --- Criterion 2: happy path -> no /spaces request --------------------------
|
||||
test("getTree with a valid spaceId returns the tree and makes NO /spaces request", async () => {
|
||||
const c = makeClient();
|
||||
// Spy client.post to prove no /spaces POST is issued on the happy path.
|
||||
const posted = [];
|
||||
c.client.post = async (url) => {
|
||||
posted.push(url);
|
||||
throw new Error(`unexpected post ${url}`);
|
||||
};
|
||||
c.enumerateSpacePages = async () => ({ pages: [], truncated: false });
|
||||
|
||||
const res = await c.getTree(SPACES[0].id);
|
||||
assert.ok(Array.isArray(res), "tree returned as before");
|
||||
assert.equal(c._spacesFetches, 0, "no /spaces sweep on the happy path");
|
||||
assert.ok(
|
||||
!posted.includes("/spaces"),
|
||||
"no /spaces POST on the happy path",
|
||||
);
|
||||
});
|
||||
|
||||
// --- Criterion 3: short-TTL cache + refetch after expiry --------------------
|
||||
test("two bad getTree within TTL share ONE /spaces fetch; a refetch happens after TTL", async () => {
|
||||
const c = makeClient();
|
||||
c.enumerateSpacePages = async () => {
|
||||
throw make404();
|
||||
};
|
||||
|
||||
await assert.rejects(() => c.getTree(BAD), /не найден среди/);
|
||||
await assert.rejects(() => c.getTree(BAD), /не найден среди/);
|
||||
assert.equal(c._spacesFetches, 1, "second call served from cache");
|
||||
|
||||
// Age the cache past the default 60s TTL -> exactly one refetch.
|
||||
assert.ok(c.spaceIndexCache, "a complete result was cached");
|
||||
c.spaceIndexCache.fetchedAt = Date.now() - 10 * 60 * 1000;
|
||||
await assert.rejects(() => c.getTree(BAD), /не найден среди/);
|
||||
assert.equal(c._spacesFetches, 2, "exactly one refetch after TTL");
|
||||
});
|
||||
|
||||
// --- Criterion 4: no spaceId -> wrapper inert -------------------------------
|
||||
test("listPages WITHOUT spaceId is inert (raw error propagates, no /spaces sweep)", async () => {
|
||||
const c = makeClient();
|
||||
c.client.post = async () => {
|
||||
throw make404("/pages/recent");
|
||||
};
|
||||
await assert.rejects(
|
||||
() => c.listPages(),
|
||||
(e) => {
|
||||
assert.equal(e.response?.status, 404, "raw server 404 propagates");
|
||||
assert.ok(!/не найден среди/.test(e.message ?? ""), "not rewritten");
|
||||
return true;
|
||||
},
|
||||
);
|
||||
assert.equal(c._spacesFetches, 0, "no /spaces sweep without a spaceId");
|
||||
});
|
||||
|
||||
test("search WITHOUT spaceId is inert (raw error propagates, no /spaces sweep)", async () => {
|
||||
const c = makeClient();
|
||||
c.client.post = async () => {
|
||||
throw make404("/search");
|
||||
};
|
||||
await assert.rejects(
|
||||
() => c.search("query"),
|
||||
(e) => {
|
||||
assert.equal(e.response?.status, 404, "raw server 404 propagates");
|
||||
assert.ok(!/не найден среди/.test(e.message ?? ""), "not rewritten");
|
||||
return true;
|
||||
},
|
||||
);
|
||||
assert.equal(c._spacesFetches, 0, "no /spaces sweep without a spaceId");
|
||||
});
|
||||
|
||||
// --- Fail-open: a NON-404 error is never enriched (regression guard) --------
|
||||
test("a non-404 error (500) on a wrapped call propagates unchanged, with NO /spaces sweep", async () => {
|
||||
const c = makeClient();
|
||||
// A real server failure — must surface as-is, never be swallowed by the
|
||||
// enrichment path or reformatted into a "not found among your spaces" message.
|
||||
c.enumerateSpacePages = async () => {
|
||||
throw makeAxiosErr(500, "/pages/tree", "Internal Server Error");
|
||||
};
|
||||
await assert.rejects(
|
||||
() => c.getTree(BAD),
|
||||
(e) => {
|
||||
assert.equal(e.response?.status, 500, "the original 500 propagates");
|
||||
assert.ok(
|
||||
!/не найден среди/.test(e.message ?? ""),
|
||||
"a non-404 is NOT rewritten",
|
||||
);
|
||||
return true;
|
||||
},
|
||||
);
|
||||
assert.equal(c._spacesFetches, 0, "no /spaces sweep for a non-404");
|
||||
});
|
||||
|
||||
// --- Fail-open: 404 when the spaceId IS accessible -> not about the space ----
|
||||
test("a 404 when the spaceId IS in the accessible index fails open (the 404 is about something else)", async () => {
|
||||
const c = makeClient();
|
||||
// The wrapped call 404s, but the spaceId is genuinely accessible — the 404
|
||||
// must be about some OTHER resource, so the original error propagates and is
|
||||
// NOT falsely rewritten to "spaceId not found among your spaces".
|
||||
c.enumerateSpacePages = async () => {
|
||||
throw make404();
|
||||
};
|
||||
await assert.rejects(
|
||||
() => c.getTree(SPACES[0].id),
|
||||
(e) => {
|
||||
assert.equal(e.response?.status, 404, "original 404 preserved");
|
||||
assert.ok(
|
||||
!/не найден среди/.test(e.message ?? ""),
|
||||
"no false 'not found' when the space is accessible",
|
||||
);
|
||||
return true;
|
||||
},
|
||||
);
|
||||
assert.equal(c._spacesFetches, 1, "the index WAS consulted to make this call");
|
||||
});
|
||||
|
||||
// --- Criterion 5: incomplete listing -> fail open ---------------------------
|
||||
test("a truncated /spaces listing (!complete) fails OPEN with the original 404", async () => {
|
||||
const c = makeClient({ truncated: true });
|
||||
c.enumerateSpacePages = async () => {
|
||||
throw make404();
|
||||
};
|
||||
await assert.rejects(
|
||||
() => c.getTree(BAD),
|
||||
(e) => {
|
||||
assert.equal(e.response?.status, 404, "original server error preserved");
|
||||
assert.ok(!/не найден среди/.test(e.message ?? ""), "no false rewrite");
|
||||
return true;
|
||||
},
|
||||
);
|
||||
assert.equal(c.spaceIndexCache, null, "a truncated result is never cached");
|
||||
});
|
||||
|
||||
// --- Criterion 6: /spaces fetch fails -> fail open; in-flight nulled on reject
|
||||
test("a /spaces fetch failure fails OPEN, logs under DEBUG, and never memoizes the rejected in-flight promise", async () => {
|
||||
const c = makeClient();
|
||||
let fetchCount = 0;
|
||||
c.paginateAllWithMeta = async () => {
|
||||
fetchCount++;
|
||||
throw new Error("network boom");
|
||||
};
|
||||
c.enumerateSpacePages = async () => {
|
||||
throw make404();
|
||||
};
|
||||
|
||||
const prevDebug = process.env.DEBUG;
|
||||
process.env.DEBUG = "1";
|
||||
const errs = [];
|
||||
const origErr = console.error;
|
||||
console.error = (...a) => errs.push(a.map(String).join(" "));
|
||||
try {
|
||||
await assert.rejects(
|
||||
() => c.getTree(BAD),
|
||||
(e) => {
|
||||
assert.equal(e.response?.status, 404, "original 404 rethrown");
|
||||
return true;
|
||||
},
|
||||
);
|
||||
} finally {
|
||||
console.error = origErr;
|
||||
if (prevDebug === undefined) delete process.env.DEBUG;
|
||||
else process.env.DEBUG = prevDebug;
|
||||
}
|
||||
|
||||
assert.ok(
|
||||
errs.some((l) => l.includes("space-diag: /spaces fetch failed")),
|
||||
"a DEBUG stderr line is emitted",
|
||||
);
|
||||
assert.equal(c.spaceIndexInFlight, null, "in-flight promise nulled on reject");
|
||||
|
||||
// The rejected in-flight promise must NOT be reused: a second enrichment does
|
||||
// a genuinely fresh fetch (fetchCount increments to 2).
|
||||
await assert.rejects(
|
||||
() => c.getTree(BAD),
|
||||
(e) => e.response?.status === 404,
|
||||
);
|
||||
assert.equal(fetchCount, 2, "fresh fetch — rejected promise not memoized");
|
||||
});
|
||||
|
||||
// --- Criterion 7: zero accessible spaces ------------------------------------
|
||||
test("zero accessible spaces yields the dedicated 'no accessible spaces' message", async () => {
|
||||
const c = makeClient({ spaces: [] });
|
||||
c.enumerateSpacePages = async () => {
|
||||
throw make404();
|
||||
};
|
||||
await assert.rejects(
|
||||
() => c.getTree(BAD),
|
||||
(e) => {
|
||||
assert.ok(
|
||||
e.message.includes("доступных тебе спейсов нет"),
|
||||
"the zero-spaces branch fired",
|
||||
);
|
||||
assert.ok(e.message.includes(BAD), "still names the bad spaceId");
|
||||
return true;
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
// --- Criterion 8: abort/cap during enrichment -------------------------------
|
||||
test("an aborted signal (custom/TimeoutError reason) propagates its reason, not a rewrite, and skips the sweep", async () => {
|
||||
const c = makeClient();
|
||||
const ac = new AbortController();
|
||||
const reason = new Error("per-call cap exceeded");
|
||||
reason.name = "TimeoutError"; // NOT 'AbortError' — hole B
|
||||
ac.abort(reason);
|
||||
c.setToolAbortSignal(ac.signal);
|
||||
// Simulate paginateAll's throwIfAborted(): the inner op rejects with the
|
||||
// signal's reason (not an AxiosError).
|
||||
c.enumerateSpacePages = async () => {
|
||||
throw ac.signal.reason;
|
||||
};
|
||||
|
||||
await assert.rejects(
|
||||
() => c.getTree(BAD),
|
||||
(e) => {
|
||||
assert.equal(e, reason, "the abort/cap reason itself propagates");
|
||||
assert.equal(e.name, "TimeoutError");
|
||||
assert.ok(!/не найден среди/.test(e.message ?? ""), "no rewrite");
|
||||
return true;
|
||||
},
|
||||
);
|
||||
assert.equal(c._spacesFetches, 0, "abort short-circuits before any sweep");
|
||||
});
|
||||
|
||||
// --- Criterion 9: createPage is NOT wrapped (real offline server) -----------
|
||||
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 = [];
|
||||
function spawn(handler) {
|
||||
return new Promise((resolve) => {
|
||||
const server = http.createServer(handler);
|
||||
server.listen(0, "127.0.0.1", () => {
|
||||
openServers.push(server);
|
||||
const { port } = server.address();
|
||||
resolve({ baseURL: `http://127.0.0.1:${port}/api` });
|
||||
});
|
||||
});
|
||||
}
|
||||
after(async () => {
|
||||
await Promise.all(openServers.map((s) => new Promise((r) => s.close(r))));
|
||||
});
|
||||
|
||||
test("createPage with an inaccessible spaceId returns the server error as-is (NOT wrapped)", async () => {
|
||||
const { baseURL } = await spawn(async (req, res) => {
|
||||
await readBody(req);
|
||||
if (req.url === "/api/auth/login") {
|
||||
sendJson(res, 200, { success: true }, {
|
||||
"Set-Cookie": "authToken=t; Path=/; HttpOnly",
|
||||
});
|
||||
return;
|
||||
}
|
||||
if (req.url === "/api/pages/import") {
|
||||
// The space-permissions 404 createPage would see for a bad spaceId.
|
||||
sendJson(res, 404, { message: "Space permissions not found" });
|
||||
return;
|
||||
}
|
||||
sendJson(res, 404, {});
|
||||
});
|
||||
|
||||
const client = new DocmostClient(baseURL, "u@example.com", "pw");
|
||||
// Prove the diagnostics path is never entered from createPage.
|
||||
let idxCalls = 0;
|
||||
const realIdx = client.getAccessibleSpaceIndex.bind(client);
|
||||
client.getAccessibleSpaceIndex = async () => {
|
||||
idxCalls++;
|
||||
return realIdx();
|
||||
};
|
||||
|
||||
await assert.rejects(
|
||||
() => client.createPage("Title", "body", BAD),
|
||||
(e) => {
|
||||
assert.equal(e.response?.status, 404, "the raw server 404 surfaces");
|
||||
assert.ok(
|
||||
!/не найден среди/.test(e.message ?? ""),
|
||||
"createPage's 404 is NOT rewritten (multi-cause path)",
|
||||
);
|
||||
return true;
|
||||
},
|
||||
);
|
||||
assert.equal(idxCalls, 0, "createPage never invokes the space diagnostics");
|
||||
});
|
||||
|
||||
// --- formatSpaceNotAccessible unit shape ------------------------------------
|
||||
test("formatSpaceNotAccessible caps the inline list at 10 and appends a (+N ещё) tail", () => {
|
||||
// Short ids/names so the whole message stays under the length cap and the full
|
||||
// list-cap behaviour (first 10 shown, rest collapsed) is observable intact.
|
||||
const many = Array.from({ length: 13 }, (_, i) => ({
|
||||
id: `s${i}`,
|
||||
name: `${i}`,
|
||||
}));
|
||||
const msg = formatSpaceNotAccessible("getTree", BAD, many);
|
||||
assert.ok(msg.includes("s0 (0)"));
|
||||
assert.ok(msg.includes("s9 (9)"), "10th entry (index 9) is shown");
|
||||
assert.ok(!msg.includes("s10 (10)"), "the 11th is collapsed");
|
||||
assert.ok(msg.includes("(+3 ещё, см. listSpaces)"), "tail counts the remainder");
|
||||
assert.ok(msg.length <= 300, `short-name message stays under the cap (${msg.length})`);
|
||||
});
|
||||
|
||||
test("formatSpaceNotAccessible caps the assembled message at ERROR_MESSAGE_CAP (300)", () => {
|
||||
// 10 spaces with long names would, uncapped, produce a message several times
|
||||
// over the 300-char budget. The cap must keep it compact while still carrying
|
||||
// the bad spaceId and the listSpaces pointer.
|
||||
const longName = "X".repeat(120);
|
||||
const spaces = Array.from({ length: 10 }, (_, i) => ({
|
||||
id: `id-${i}`,
|
||||
name: `${longName}-${i}`,
|
||||
}));
|
||||
const msg = formatSpaceNotAccessible("getTree", BAD, spaces);
|
||||
assert.ok(
|
||||
msg.length <= 300,
|
||||
`message must be <= 300 chars, got ${msg.length}`,
|
||||
);
|
||||
assert.ok(msg.includes(BAD), "the bad spaceId survives the cap");
|
||||
assert.ok(msg.includes("listSpaces"), "the listSpaces pointer survives the cap");
|
||||
});
|
||||
Reference in New Issue
Block a user