Compare commits

..

4 Commits

Author SHA1 Message Date
agent_coder 46bb55dbd1 Merge remote-tracking branch 'gitea/develop' into fix/534-mcp-spaceid 2026-07-12 05:58:05 +03:00
agent_coder 8b34a428f4 fix(mcp): unconditional final ERROR_MESSAGE_CAP backstop in formatSpaceNotAccessible (#536 review)
The list-only cap assumed a well-formed prefix; a pathologically long
agent-supplied spaceId (unvalidated for length) lives in the prefix and bypassed
the cap. Add the same unconditional final slice formatDocmostAxiosError uses, so
the WHOLE message is bounded regardless of spaceId length. No-op for normal
inputs (36-char UUID) where the list cap already keeps it <= 300.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 05:58:04 +03:00
agent_coder e236782260 fix(mcp): cap #534 space-not-accessible message + cover fail-open guards
Review follow-up (#536).

F1 (test coverage): add two tests that pin the wrapper's fail-open guards, which
previously survived mutation:
- a non-404 error (500) on a wrapped call with a spaceId propagates unchanged and
  triggers NO /spaces sweep (a real 5xx/403 must never be swallowed/reformatted);
- a 404 when the spaceId IS in the accessible index fails open (the 404 is about
  another resource), so it is not falsely rewritten to "not found among spaces".
Both mutation-verified: forcing the non-404 condition to false reddens the first;
removing the id-present guard reddens the second.

F2 (conventions): formatSpaceNotAccessible now honours ERROR_MESSAGE_CAP (300),
the same budget formatDocmostAxiosError enforces. Only the interpolated space
list is truncated (with an ellipsis); the fixed prefix (bad spaceId) and suffix
(listSpaces pointer) are always kept, so the actionable parts survive. Test: 10
long space names -> message <= 300 and still contains the spaceId + listSpaces.
Adjusted the existing list-cap test to short ids/names so it exercises the 10-item
cap + "(+N ещё)" tail below the length cap.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 05:49:31 +03:00
agent_coder 575125a5dc fix(mcp): enrich bad/inaccessible spaceId 404 into an actionable error (#534)
A well-formed but non-existent/inaccessible spaceId made the space-permissions
check answer with an opaque 404 ("Space permissions not found"), which the agent
could not self-correct from. Add an enrich-on-404 wrapper at the client-method
level (Variant A — no backend change) that, ONLY on that 404, replaces the server
text with a factual message naming the bad spaceId and the spaces the token can
actually see, pointing at listSpaces.

Mechanism (context.ts):
- getAccessibleSpaceIndex(): the token's accessible spaces from the single source
  of truth (/spaces), with a per-instance short-TTL cache (MCP_SPACES_CACHE_TTL_MS,
  default 60s; 0 disables) and single-flight dedup. Only a COMPLETE (untruncated)
  result is cached and only a complete result may drive the rewrite. The in-flight
  promise is nulled on BOTH resolve and reject so a transient /spaces blip is never
  memoized. Cache invalidated on every identity change, mirroring collabTokenCache
  (login() + the 401/403 reauth interceptor).
- paginateAllWithMeta(): surfaces the `truncated` flag paginateAll swallows;
  paginateAll now delegates to it (unchanged contract, getSpaces untouched).
- withSpaceAccessDiagnostics(spaceId, mcpName, fn): abort/cap wins FIRST (by the
  toolAbortSignal flag, NOT e.name — a cap may be TimeoutError/custom); only a 404
  is enrichable; FAILS OPEN (rethrows the original server error) on every source of
  uncertainty (fetch failed / !complete / spaceId present / aborted).

Wrapped only the paths where the sole 404 cause is the space membership check:
getTree, listPages (tree + recent-with-spaceId), search (with spaceId),
checkNewComments. createPage/getPageContext are deliberately NOT wrapped.
formatSpaceNotAccessible (errors.ts) composes the message (<=10 spaces + "(+N ещё)"
tail; distinct zero-spaces variant).

Tests: 11 new unit tests (stub client) covering all 9 acceptance criteria +
message shape; full mcp unit suite green (701 tests).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 02:20:12 +03:00
22 changed files with 812 additions and 1850 deletions
@@ -27,7 +27,6 @@ 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.
@@ -115,89 +114,3 @@ 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.
});
});
@@ -1,112 +0,0 @@
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,7 +4,6 @@ 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";
@@ -87,39 +86,17 @@ 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 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.
* 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).
*/
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 (
@@ -202,10 +179,47 @@ function MessageItem({
{resolveAssistantName(assistantName) ?? t("AI agent")}
</Text>
{message.parts.map((part, index) => {
// 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 (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}
/>
);
}
if (isToolPart(part.type)) {
return (
<ToolCallCard
@@ -218,76 +232,7 @@ function MessageItem({
);
}
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;
}
}
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. */}
@@ -1,96 +0,0 @@
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,7 +35,6 @@ 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';
@@ -44,8 +43,6 @@ import {
AiChatRunHooks,
AiChatService,
AiChatStreamBody,
rowHasInlineParts,
hydrateAssistantParts,
} from './ai-chat.service';
import { AiChatRunService } from './ai-chat-run.service';
import { AiTranscriptionService } from './ai-transcription.service';
@@ -132,39 +129,8 @@ 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')
@@ -218,17 +184,11 @@ export class AiChatController {
@AuthWorkspace() workspace: Workspace,
) {
await this.assertOwnedChat(dto.chatId, user, workspace);
const page = await this.aiChatMessageRepo.findByChat(
return 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),
};
}
/**
@@ -265,10 +225,7 @@ export class AiChatController {
workspace.id,
);
return {
// #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),
rows,
cursor,
run: run ? { id: run.id, status: run.status } : null,
};
@@ -290,10 +247,8 @@ export class AiChatController {
@AuthWorkspace() workspace: Workspace,
): Promise<{ markdown: string }> {
const chat = await this.assertOwnedChat(dto.chatId, user, workspace);
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.
const rows = await this.aiChatMessageRepo.findAllByChat(
dto.chatId,
workspace.id,
);
const markdown = buildChatMarkdown({
@@ -333,13 +288,7 @@ export class AiChatController {
workspace.id,
)
: undefined;
// #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 };
return { run, message: message ?? null };
}
/**
+12 -203
View File
@@ -22,7 +22,6 @@ 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';
@@ -519,12 +518,6 @@ 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
@@ -1121,34 +1114,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
chatId,
workspace.id,
);
// #492: HYDRATE needy assistant rows from the steps table BEFORE the replay
// map. A #492 mid-run assistant row carries only a step marker
// (metadata.parts:[]); its real per-step parts live in `ai_chat_run_steps`.
// The graceful terminal callbacks (onFinish/onError/onAbort -> flushAssistant)
// assemble the full inline parts, so a normally-ended turn already has them.
// But a HARD crash mid-run (SIGKILL/OOM) fires NO terminal callback, so the
// row stays parts:[]; without this, rowToUiMessage falls back to an empty
// text part and the partial tool-calls/results/text — durable in the steps
// table — would DROP OUT of the model's replay context (regressing #183
// step-granular durability for the model consumer). Mirrors the controller's
// withReconstructedParts EXACTLY (same needy predicate + hydration helper).
// Guarded on the optional repo: absent (positional test builds) degrades to
// the current behavior rather than crashing.
let replayHistory = oldHistory;
if (this.aiChatRunStepRepo) {
const needy = oldHistory.filter(
(r) => r.role === 'assistant' && !rowHasInlineParts(r),
);
if (needy.length > 0) {
const stepsByMessage = await this.aiChatRunStepRepo.findByMessageIds(
needy.map((r) => r.id),
workspace.id,
);
replayHistory = hydrateAssistantParts(oldHistory, stepsByMessage);
}
}
const uiMessages: Array<Omit<UIMessage, 'id'> & { id: string }> = [
...replayHistory.map(rowToUiMessage),
...oldHistory.map(rowToUiMessage),
{
id: 'pending-user',
role: 'user',
@@ -1187,9 +1154,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// hint — confirm it against the persisted history (the preceding assistant
// turn must really be aborted/streaming) so a spoofed flag cannot inject the
// interrupt note onto an ordinary turn. The partial output the model needs is
// already in `messages`: a #492 mid-run row's per-step parts live only in the
// `ai_chat_run_steps` table and were hydrated into the replay history above,
// so the aborted assistant turn replays WITH its partial parts intact.
// already in `messages` (the aborted assistant row replays via findRecent).
// Append the new user turn (shape-only) so index -2 is the prior assistant.
const interrupted = isInterruptResume(
[...oldHistory, { role: 'user', status: null, metadata: null }],
@@ -1594,57 +1559,17 @@ 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;
// 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;
// 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;
try {
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 },
);
}
await this.aiChatMessageRepo.update(assistantId, workspace.id, flushed, {
onlyIfStreaming: true,
});
return stepsPersisted;
} catch (err) {
this.logger.warn(
@@ -2824,122 +2749,6 @@ 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,7 +32,6 @@ 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';
@@ -126,7 +125,6 @@ import { firstSqlToken } from '../integrations/metrics/metrics.constants';
AiChatRepo,
AiChatMessageRepo,
AiChatRunRepo,
AiChatRunStepRepo,
AiChatPageSnapshotRepo,
AiProviderCredentialsRepo,
AiMcpServerRepo,
@@ -163,7 +161,6 @@ import { firstSqlToken } from '../integrations/metrics/metrics.constants';
AiChatRepo,
AiChatMessageRepo,
AiChatRunRepo,
AiChatRunStepRepo,
AiChatPageSnapshotRepo,
AiProviderCredentialsRepo,
AiMcpServerRepo,
@@ -1,70 +0,0 @@
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();
}
@@ -1,95 +0,0 @@
import { Injectable } from '@nestjs/common';
import { InjectKysely } from 'nestjs-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;
}
}
-17
View File
@@ -692,22 +692,6 @@ 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
@@ -745,7 +729,6 @@ export interface DB {
aiChats: AiChats;
aiChatMessages: AiChatMessages;
aiChatRuns: AiChatRuns;
aiChatRunSteps: AiChatRunSteps;
aiChatPageSnapshots: AiChatPageSnapshots;
apiKeys: ApiKeys;
attachments: Attachments;
@@ -4,7 +4,6 @@ import {
AiChats,
AiChatMessages,
AiChatRuns,
AiChatRunSteps,
AiChatPageSnapshots,
Attachments,
Comments,
@@ -65,12 +64,6 @@ 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.
@@ -1,412 +0,0 @@
import * as http from 'node:http';
import { Kysely } from 'kysely';
import { tool } from 'ai';
import { z } from 'zod';
import { MockLanguageModelV3, convertArrayToReadableStream } from 'ai/test';
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 {
AiChatService,
assembleStepParts,
assistantParts,
rowHasInlineParts,
stepMarkerMetadata,
} from 'src/core/ai-chat/ai-chat.service';
import {
getTestDb,
destroyTestDb,
createWorkspace,
createUser,
createChat,
createMessage,
} from './db';
/**
* #492 append-persist the REAL onStep WRITE path (F2) and the model-REPLAY
* hydration path (F1), driven through `AiChatService.stream` against a LIVE
* Postgres with a REAL `AiChatRunStepRepo` INJECTED. The existing append-persist
* int-specs hand-roll the insert+marker cycle via the repos directly and build
* the service with `aiChatRunStepRepo: undefined` (only the legacy-fallback branch
* is covered), so an off-by-one on `stepsPersisted-1`, a wrong `capturedSteps`
* slice, or a broken marker payload would pass all of them. These tests exercise
* the actual `updateStreaming` append-persist branch end to end.
*
* The seam is the injected `model` (a seeded `MockLanguageModelV3` from `ai/test`)
* plus a REAL Node `ServerResponse` as the hijacked socket mirrors
* ai-chat-stream.int-spec.ts.
*/
const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));
async function waitFor(
cond: () => Promise<boolean> | boolean,
{ timeoutMs = 15_000, stepMs = 25 } = {},
): Promise<void> {
const start = Date.now();
while (Date.now() - start < timeoutMs) {
if (await cond()) return;
await sleep(stepMs);
}
throw new Error('waitFor: condition not met within timeout');
}
// A real Node ServerResponse wired to a live socket (identical helper to the
// stream int-spec) so the SDK's pipe/heartbeat writes behave as in prod.
function makeRealResponse(): Promise<{
res: http.ServerResponse;
cleanup: () => Promise<void>;
}> {
return new Promise((resolve) => {
const server = http.createServer((_req, res) => {
resolve({
res,
cleanup: () =>
new Promise<void>((done) => {
try {
if (!res.writableEnded) res.end();
} catch {
/* socket already gone */
}
server.close(() => done());
}),
});
});
server.listen(0, () => {
const port = (server.address() as any).port;
const creq = http.request({ port, method: 'GET' }, (cres) => {
cres.resume();
});
creq.on('error', () => undefined);
creq.end();
});
});
}
// Stream parts for a normal, successful single-step turn.
function successStream() {
return convertArrayToReadableStream([
{ type: 'stream-start', warnings: [] },
{ type: 'text-start', id: 't1' },
{ type: 'text-delta', id: 't1', delta: 'Hello' },
{ type: 'text-delta', id: 't1', delta: ' there' },
{ type: 'text-end', id: 't1' },
{
type: 'finish',
finishReason: 'stop',
usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 },
},
] as any);
}
// A THREE-step turn: steps 0 and 1 each emit text + an `echo` tool call (the SDK
// runs the tool and continues); step 2 answers and stops. Three steps is
// deliberate: the LAST finished step's append-persist write races the terminal
// finalize (which writes the full inline parts anyway, so a lost last-step row is
// by design), but the NON-final steps 0 and 1 always drain to the steps table
// before finalize — so those are what the test asserts on deterministically.
function threeStepModel(): MockLanguageModelV3 {
let step = 0;
const toolStep = (i: number) => ({
stream: convertArrayToReadableStream([
{ type: 'stream-start', warnings: [] },
{ type: 'text-start', id: `s${i}` },
{ type: 'text-delta', id: `s${i}`, delta: `step ${i} ` },
{ type: 'text-end', id: `s${i}` },
{
type: 'tool-call',
toolCallId: `c${i}`,
toolName: 'echo',
input: JSON.stringify({ msg: `m${i}` }),
},
{
type: 'finish',
finishReason: 'tool-calls',
usage: { inputTokens: 5, outputTokens: 3, totalTokens: 8 },
},
] as any),
});
return new MockLanguageModelV3({
doStream: async () => {
const n = step++;
// Realistic inter-step latency. A real model spends seconds per step, so the
// fire-and-forget per-step write chain drains to the steps table BETWEEN
// steps; the mock otherwise collapses all steps into microseconds and the
// terminal finalize wins the race before any but the first step persists.
if (n > 0) await sleep(200);
if (n < 2) return toolStep(n);
return {
stream: convertArrayToReadableStream([
{ type: 'stream-start', warnings: [] },
{ type: 'text-start', id: 's2' },
{ type: 'text-delta', id: 's2', delta: 'final answer' },
{ type: 'text-end', id: 's2' },
{
type: 'finish',
finishReason: 'stop',
usage: { inputTokens: 6, outputTokens: 4, totalTokens: 10 },
},
] as any),
};
},
} as any);
}
describe('#492 append-persist service paths [integration]', () => {
let db: Kysely<any>;
let aiChatRepo: AiChatRepo;
let msgRepo: AiChatMessageRepo;
let stepRepo: AiChatRunStepRepo;
let workspaceId: string;
let userId: string;
let closeCalls: number;
const mcpClients = {
toolsFor: async () => ({
tools: {},
clients: [
{
close: async () => {
closeCalls += 1;
},
},
],
outcomes: [],
instructions: [],
}),
};
// Build the service WITH a REAL AiChatRunStepRepo injected (the property under
// test) — unlike the legacy-fallback harness that passes it as undefined.
const echoTool = tool({
description: 'echo the message back',
inputSchema: z.object({ msg: z.string() }),
execute: async ({ msg }) => ({ echoed: msg }),
});
function buildService(): AiChatService {
return new AiChatService(
{ getChatModel: async () => null } as any,
aiChatRepo,
msgRepo,
{} as any, // aiChatPageSnapshotRepo
{ resolve: async () => null } as any, // aiSettings
{ forUser: async () => ({ echo: echoTool }) } as any, // tools
mcpClients as any,
{} as any, // aiAgentRoleRepo
{} as any, // pageRepo
{} as any, // pageAccess
{
isAiChatDeferredToolsEnabled: () => false,
isAiChatFinalStepLockdownEnabled: () => false,
} as any, // environment (deferred OFF -> all tools active every step)
undefined, // streamRegistry
undefined, // aiChatRunService
stepRepo, // #492 aiChatRunStepRepo — the append-persist backend
);
}
function userUiMessage(text: string) {
return {
id: `u-${Math.random()}`,
role: 'user',
parts: [{ type: 'text', text }],
};
}
async function runStream(opts: {
model: MockLanguageModelV3;
chatId: string;
body: any;
}): Promise<void> {
closeCalls = 0;
const service = buildService();
const { res, cleanup } = await makeRealResponse();
try {
await service.stream({
user: { id: userId, workspaceId } as any,
workspace: { id: workspaceId, name: 'WS' } as any,
sessionId: 'sess-1',
body: opts.body,
res: { raw: res } as any,
signal: new AbortController().signal,
model: opts.model as any,
role: null,
} as any);
await waitFor(async () => {
const rows = await msgRepo.findAllByChat(opts.chatId, workspaceId);
return rows.some(
(r) =>
r.role === 'assistant' &&
['completed', 'error', 'aborted'].includes(r.status as string),
);
});
await waitFor(() => closeCalls > 0, { timeoutMs: 5_000 });
} finally {
await cleanup();
}
}
beforeAll(async () => {
db = getTestDb();
aiChatRepo = new AiChatRepo(db as any);
msgRepo = new AiChatMessageRepo(db as any);
stepRepo = new AiChatRunStepRepo(db as any);
workspaceId = (await createWorkspace(db)).id;
userId = (await createUser(db, workspaceId)).id;
});
afterAll(async () => {
await destroyTestDb();
});
// --- F2: the real onStep append-persist WRITE branch -----------------------
it('drives steps through the real onStep path: per-step rows + marker match a single-row flush', async () => {
const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
const model = threeStepModel();
// Capture the mid-run step-marker UPDATEs the append-persist branch writes on
// the assistant row (a { parts: [], toolTraceVersion, stepsPersisted } patch).
const updateSpy = jest.spyOn(msgRepo, 'update');
try {
await runStream({
model,
chatId,
body: { chatId, messages: [userUiMessage('call the tool then answer')] },
});
const rows = await msgRepo.findAllByChat(chatId, workspaceId);
const assistant = rows.find((r) => r.role === 'assistant')!;
expect(assistant).toBeDefined();
expect(assistant.status).toBe('completed');
// The turn finalizes with the FULL inline parts assembled by a single-row
// flush (assistantParts over every step) — the baseline the per-step slices
// must reproduce.
expect(rowHasInlineParts(assistant)).toBe(true);
const finalParts = (assistant.metadata as { parts: any[] }).parts;
// The two NON-final finished steps each landed their own row, in stepIndex
// order. (The fire-and-forget write chain drains before the next step, so
// poll until both are on disk; the LAST step's write may lose the finalize
// race, which is by design — its parts are already in `finalParts`.)
await waitFor(async () => {
const s = await stepRepo.findByMessage(assistant.id, workspaceId);
return s.length >= 2;
});
const steps = await stepRepo.findByMessage(assistant.id, workspaceId);
expect(steps[0].stepIndex).toBe(0);
expect(steps[1].stepIndex).toBe(1);
// Each per-step row carries a NON-trivial slice: this step's text part + its
// paired tool part (guards a mutation that persists empty/whole-turn parts).
const s0 = steps[0].parts as any[];
expect(s0).toContainEqual({ type: 'text', text: 'step 0 ' });
expect(s0.some((p) => p.type === 'tool-echo')).toBe(true);
// The per-step slices are EXACTLY the corresponding prefix of the single-row
// flush: assembleStepParts([step0, step1]) === finalParts[0 .. len0+len1].
// This is what an off-by-one on `stepsPersisted-1` (a wrong `capturedSteps`
// slice) or a shifted stepIndex breaks — the prefix no longer aligns.
const prefixLen =
(steps[0].parts as any[]).length + (steps[1].parts as any[]).length;
expect(assembleStepParts([steps[0], steps[1]] as any)).toEqual(
finalParts.slice(0, prefixLen),
);
// The mid-run step markers advanced 1 -> 2 -> ... (the resume frontier), each
// a shape-stable empty-parts marker equal to a single-row flush's marker.
const markerCounts = updateSpy.mock.calls
.map((c) => (c[2] as any)?.metadata)
.filter(
(m) =>
m &&
Array.isArray(m.parts) &&
m.parts.length === 0 &&
typeof m.stepsPersisted === 'number',
)
.map((m) => m.stepsPersisted);
// Monotonic from 1, covering at least the two non-final steps.
expect(markerCounts.slice(0, 2)).toEqual([1, 2]);
expect(
updateSpy.mock.calls
.map((c) => (c[2] as any)?.metadata)
.find((m) => m && m.stepsPersisted === 2),
).toEqual(stepMarkerMetadata(2));
} finally {
updateSpy.mockRestore();
}
}, 60_000);
// --- F1: model-REPLAY hydrates a hard-crashed mid-run turn from the steps table
it('replays a hard-crashed mid-run turn WITH its partial steps hydrated from the steps table', async () => {
const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
// Prior turn: a genuine user question...
await createMessage(db, {
workspaceId,
chatId,
userId,
role: 'user',
content: 'What is in the design doc?',
createdAt: new Date(Date.now() - 3000),
});
// ...and an assistant row that a HARD crash (SIGKILL/OOM) left mid-run: only a
// step marker on the row (metadata.parts:[] , content:''), NO terminal
// callback ever fired, so its real parts live ONLY in ai_chat_run_steps.
const crashed = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
content: '',
status: 'aborted',
metadata: stepMarkerMetadata(1),
createdAt: new Date(Date.now() - 2000),
});
// The durable partial step: some reasoning text + a completed getPage tool
// call (input + output), exactly what #183 step-granular durability preserves.
await stepRepo.insertStep(
crashed.id,
workspaceId,
0,
assistantParts(
[
{
text: 'HYDRATED_PARTIAL_STEP the doc says',
toolCalls: [
{ toolCallId: 'g1', toolName: 'getPage', input: { id: 'p1' } },
],
toolResults: [
{
toolCallId: 'g1',
toolName: 'getPage',
output: { id: 'p1', body: 'PARTIAL_TOOL_OUTPUT budget section' },
},
],
} as any,
],
'',
),
);
// The NEXT turn: the model just answers. The service must REPLAY the crashed
// assistant turn with its partial parts hydrated from the steps table.
const model = new MockLanguageModelV3({
doStream: async () => ({ stream: successStream() }),
} as any);
await runStream({
model,
chatId,
body: { chatId, messages: [userUiMessage('Continue please')] },
});
expect(model.doStreamCalls.length).toBeGreaterThan(0);
const prompt = JSON.stringify(model.doStreamCalls[0].prompt);
// The partial step's TEXT reached the model context (it would be an empty text
// part without hydration — rowToUiMessage falls back to `content:''`).
expect(prompt).toContain('HYDRATED_PARTIAL_STEP');
// The partial TOOL RESULT survived too (durable in the steps table, replayed).
expect(prompt).toContain('PARTIAL_TOOL_OUTPUT');
// The genuine prior user turn is present as well (sanity: real history replay).
expect(prompt).toContain('What is in the design doc?');
}, 60_000);
});
@@ -1,173 +0,0 @@
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);
});
@@ -1,169 +0,0 @@
import { Kysely } from 'kysely';
import { AiChatController } from 'src/core/ai-chat/ai-chat.controller';
import {
assembleStepParts,
assistantParts,
stepMarkerMetadata,
} from 'src/core/ai-chat/ai-chat.service';
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 type { User, Workspace } from '@docmost/db/types/entity.types';
import {
getTestDb,
destroyTestDb,
createWorkspace,
createUser,
createChat,
createMessage,
} from './db';
/**
* #492 controller hydration (crash-before-finalize RESUME) on a LIVE Postgres.
* `AiChatController.withReconstructedParts` is wired into getMessages/delta/export/
* run, but `aiChatRunStepRepo` is OPTIONAL and every controller unit spec passes it
* as `undefined`, so the hydration branch early-returns and NEVER executes in those
* tests. This drives the real read path a mid-run streaming row (marker only,
* empty inline parts) PLUS its `ai_chat_run_steps` rows through getMessages WITH
* the repo present, exercising the `role==='assistant' && !rowHasInlineParts`
* needy predicate, the workspace-scoped batch step fetch, and the endpoint binding.
*/
describe('#492 controller hydration read path [integration]', () => {
let db: Kysely<any>;
let aiChatRepo: AiChatRepo;
let msgRepo: AiChatMessageRepo;
let stepRepo: AiChatRunStepRepo;
let workspaceId: string;
let otherWorkspaceId: string;
let userId: string;
// Build the controller WITH a real AiChatRunStepRepo injected (position 9), the
// seam the unit specs leave undefined. Only the read-path deps are real.
function buildController(): AiChatController {
return new AiChatController(
{} as any, // aiChatService
{} as any, // aiChatRunService
aiChatRepo,
msgRepo,
{} as any, // aiTranscription
{} as any, // pageRepo
undefined, // streamRegistry
undefined, // environment
stepRepo, // #492 aiChatRunStepRepo
);
}
beforeAll(async () => {
db = getTestDb();
aiChatRepo = new AiChatRepo(db as any);
msgRepo = new AiChatMessageRepo(db as any);
stepRepo = new AiChatRunStepRepo(db as any);
workspaceId = (await createWorkspace(db)).id;
otherWorkspaceId = (await createWorkspace(db)).id;
userId = (await createUser(db, workspaceId)).id;
});
afterAll(async () => {
await destroyTestDb();
});
it('getMessages reconstructs a mid-run row from the steps table (finished rows untouched)', async () => {
const chatId = (
await createChat(db, { workspaceId, creatorId: userId })
).id;
const user = { id: userId } as User;
const workspace = { id: workspaceId } as Workspace;
// A prior FINISHED assistant row that already carries inline parts — the needy
// predicate must SKIP it (no step fetch), returned untouched.
const finishedParts = assistantParts(
[{ text: 'done earlier', toolCalls: [], toolResults: [] } as any],
'',
);
await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
content: 'done earlier',
status: 'completed',
metadata: { parts: finishedParts, toolTraceVersion: 2, stepsPersisted: 1 },
createdAt: new Date(Date.now() - 3000),
});
// The mid-run row a crash-before-finalize left behind: a step marker only
// (parts:[] , content:''), status 'streaming'. Its real parts live ONLY in the
// steps table.
const midRun = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
content: '',
status: 'streaming',
metadata: stepMarkerMetadata(2),
createdAt: new Date(Date.now() - 1000),
});
const step0 = assistantParts(
[
{
text: 'reasoning about the page',
toolCalls: [
{ toolCallId: 'g1', toolName: 'getPage', input: { id: 'p1' } },
],
toolResults: [
{ toolCallId: 'g1', toolName: 'getPage', output: { id: 'p1', body: 'B' } },
],
} as any,
],
'',
);
const step1 = assistantParts(
[{ text: 'partial synthesis so far', toolCalls: [], toolResults: [] } as any],
'',
);
await stepRepo.insertStep(midRun.id, workspaceId, 0, step0);
await stepRepo.insertStep(midRun.id, workspaceId, 1, step1);
// Workspace-scoping guard: a step row for the SAME message id under a DIFFERENT
// workspace must NEVER leak into this workspace's reconstruction.
await stepRepo.insertStep(midRun.id, otherWorkspaceId, 99, [
{ type: 'text', text: 'FOREIGN_WORKSPACE_LEAK' },
]);
const res = await buildController().getMessages(
{ chatId } as any,
{ limit: 50 } as any,
user,
workspace,
);
const items = res.items as any[];
const finished = items.find((r) => r.status === 'completed');
const reconstructed = items.find((r) => r.id === midRun.id);
// The finished row passed through with its inline parts unchanged.
expect(finished.metadata.parts).toEqual(finishedParts);
// The mid-run row's parts were reconstructed from the two step rows, in order,
// exactly as assembleStepParts concatenates them — the client seed sees the
// persisted progress with no change to itself.
const expected = assembleStepParts([
{ stepIndex: 0, parts: step0 },
{ stepIndex: 1, parts: step1 },
] as any);
expect(reconstructed.metadata.parts).toEqual(expected);
// The foreign-workspace step row did NOT leak in.
expect(JSON.stringify(reconstructed.metadata.parts)).not.toContain(
'FOREIGN_WORKSPACE_LEAK',
);
// Sanity: reconstruction produced real content (text + the paired tool part +
// the second step's text), not an empty fallback.
expect(reconstructed.metadata.parts).toContainEqual({
type: 'text',
text: 'reasoning about the page',
});
expect(
(reconstructed.metadata.parts as any[]).some((p) => p.type === 'tool-getPage'),
).toBe(true);
}, 60_000);
});
@@ -1,163 +0,0 @@
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]);
// 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,
);
});
});
@@ -1,63 +0,0 @@
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);
});
});
+5 -1
View File
@@ -31,7 +31,11 @@ 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 } from "./client/errors.js";
export {
formatDocmostAxiosError,
assertFullUuid,
formatSpaceNotAccessible,
} 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
+9 -4
View File
@@ -725,10 +725,15 @@ 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.
const { pages: pagesInScope, truncated } = await this.enumerateSpacePages(
spaceId,
parentPageId,
);
// #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),
);
// 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
+213 -3
View File
@@ -26,7 +26,10 @@ 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 } from "./errors.js";
import {
formatDocmostAxiosError,
formatSpaceNotAccessible,
} from "./errors.js";
import { GetPageConversionCache } from "./getpage-cache.js";
// A generic mixin base constructor (issue #450). Each domain mixin is a factory
@@ -117,6 +120,36 @@ 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;
@@ -164,6 +197,27 @@ 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
@@ -281,6 +335,9 @@ 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();
@@ -413,6 +470,8 @@ 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}`;
})
@@ -630,13 +689,34 @@ export abstract class DocmostClientContext {
}
/**
* Generic pagination handler for Docmost API endpoints
* 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.
*/
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));
@@ -697,7 +777,137 @@ export abstract class DocmostClientContext {
);
}
return allItems;
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));
}
}
+54
View File
@@ -46,6 +46,60 @@ 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.
+46 -16
View File
@@ -132,17 +132,29 @@ 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.
const { pages, truncated } = await this.enumerateSpacePages(spaceId);
return { tree: buildPageTree(pages), truncated };
// #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 clampedLimit = Math.max(1, Math.min(100, limit));
const payload: Record<string, any> = { limit: clampedLimit, page: 1 };
if (spaceId) payload.spaceId = spaceId;
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));
// #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();
}
/**
@@ -174,8 +186,16 @@ export function ReadMixin<TBase extends GConstructor<DocmostClientContext>>(Base
"getTree: spaceId is required (a page tree is scoped to one space).",
);
}
const { pages } = await this.enumerateSpacePages(spaceId, rootPageId);
return buildPageTree(pages, { shape: "getTree", maxDepth });
// #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 });
});
}
/**
@@ -700,17 +720,27 @@ 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);
// 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));
const runSearch = async () => {
const response = await this.client.post("/search", payload);
return {
items: filteredItems,
success: response.data?.success || false,
// 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,
};
};
// #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();
}
}
@@ -0,0 +1,423 @@
// 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");
});