Compare commits

..

2 Commits

Author SHA1 Message Date
claude code agent 227
88199703fe fix(ai): store chatContextWindow as a JSON number, not a ::text string
chatContextWindow (#189) is the first numeric provider field routed
through WorkspaceRepo.updateAiProviderSettings, whose patch builder cast
every value as `${v}::text`. The DTO validates it as @IsInt(), so a JS
number 200000 was stored as the JSON STRING "200000". The client guards
require `typeof === "number"` (ai-chat-window.tsx, context-badge.tsx),
so the `/ max` badge denominator never rendered and the whole feature
silently no-opped.

Branch the jsonb_build_object value cast by JS runtime type: numbers ->
::numeric (real JSON number), booleans -> ::boolean, everything else ->
::text (unchanged for the existing string fields). This is the root fix
(store as a real number) rather than coercing on read, so every reader
sees the correct type.

Add a DB round-trip int-spec asserting
jsonb_typeof(settings->'ai'->'provider'->'chatContextWindow') = 'number'
and that the value re-reads as the number 200000, including the
partial-merge path. CHANGELOG: Added entry for the chatContextWindow
setting and a Changed entry for the badge's new "used / max" meaning.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-26 17:19:34 +03:00
claude code agent 227
d88fe4cde7 feat(ai-chat): context badge shows current/max (#189)
The header badge in the floating AI-chat window flipped meaning between
states (a live per-turn token counter while streaming vs. the context
size at rest), which made it "reset to 1" on each prompt and confused
users. Make it consistently show the current context size, with the
model's context window as an optional "/ max" denominator.

The max comes from a new admin-set AI setting (chatContextWindow, in
tokens) — provider-independent and always exact. The server stamps it
onto the assistant message metadata (maxContextTokens) next to
contextTokens, so the client reads both from the last row with no
client-side model resolution (survives shares / future per-role models).

- server: chatContextWindow in AiProviderSettings/keys/masked/resolved,
  DTO (@IsInt @Min(0)), settings-service resolve/getMasked, repo parity
  allowlist; flushAssistant writes metadata.maxContextTokens when > 0.
- client: ContextBadge component (extracted, shows "current [/ max]",
  no live mode); removed the liveTurnTokens header path + dead util fn;
  Context-window NumberInput in AI settings; i18n strings.
- live "Thinking · N tokens" feedback in the chat body is unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-26 06:27:45 +03:00
28 changed files with 522 additions and 878 deletions

View File

@@ -43,6 +43,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
OpenRouter, etc.; `openai` uses the official provider (real-OpenAI
reasoning-model request shaping). Chosen explicitly rather than inferred from
the base URL, since a custom URL can front real OpenAI too. (#175, #177)
- **AI chat "Context window (tokens)" setting (`chatContextWindow`).** A new
admin field in AI settings that records the chat model's context-window size.
When set (> 0) it becomes the denominator of the header context-badge, which
now reads "used / max"; `0`/empty clears the limit and the badge shows only
the current context as before. There is no provider-independent way to read a
model's window automatically, so it is an explicit workspace-level value.
(#189)
- **Per-MCP-server instructions in the agent prompt.** Each external MCP server
now has an admin-authored `instructions` field ("how/when to use this server's
tools") that is injected into the agent's system prompt next to that server's
@@ -52,14 +59,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- **Footnote multi-backlinks.** A footnote referenced more than once now shows a
back-link per reference (↩ a b c …), each scrolling to its own occurrence, like
Pandoc/Wikipedia; a single-reference footnote keeps the plain ↩. (#168)
- **Model-friendly AI-chat tool-input errors.** When the model calls an in-app
AI tool with bad arguments, the validation failure is now a concise,
human-readable message that NAMES each offending parameter (by its dotted
path) and appends a fixed retry hint ("include every REQUIRED parameter…, do
not drop ids like `pageId`"), instead of the raw zod text. This nudges the
model to re-issue the call correctly — particularly in parallel tool-call
batches where it tends to drop a repeated id. The required/optional contract
and unknown-key stripping are unchanged. (#190)
### Changed
@@ -69,6 +68,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
model's reasoning out of the box. An endpoint that is real OpenAI behind a
custom base URL should set the new `chatApiStyle` "Protocol" to `openai`. (#177)
- **AI chat header context-badge now shows "used / max".** When an admin sets
the new `chatContextWindow`, the badge displays the current context size over
the configured window (e.g. `120k / 200k`) instead of switching to a live
per-turn token counter during streaming. With no window configured the badge
keeps showing just the current context. (#189)
- **Footnotes now reuse (Pandoc semantics).** Multiple `[^a]` references to the
same id are ONE footnote — one number, one definition, several back-references
— instead of being renamed to `a__2`, `a__3`. Duplicate `[^a]:` definitions are
@@ -100,14 +105,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
no longer froze on the previous step's authoritative usage; the current step's
estimate is combined per-component with `max`, so the count rises smoothly and
never jumps backwards. (#163)
- **Concurrent page moves can no longer lose a subtree to a cycle.** Two
opposing re-parents racing each other (A: X under Y, B: Y under X) could each
pass a cycle check built from a stale snapshot and commit a cycle, orphaning a
subtree. A genuine re-parent under a concrete parent now serializes: it locks
the moved page and the destination parent `FOR UPDATE` in a canonical
(UUID-sorted) order — so opposing moves can't deadlock — and re-runs the cycle
check INSIDE the transaction against the now-committed state. Same-parent
reorders and moves to root keep the lock-free path. (#159)
## [0.93.0] - 2026-06-21

View File

@@ -1168,7 +1168,10 @@
"Built-in assistant persona": "Built-in assistant persona",
"Minimize": "Minimize",
"Current context size": "Current context size",
"Tokens generated this turn": "Tokens generated this turn",
"Context size / model limit": "Context size / model limit",
"Context window (tokens)": "Context window (tokens)",
"Shows used / total in the chat header badge; empty hides the total.": "Shows used / total in the chat header badge; empty hides the total.",
"e.g. 200000": "e.g. 200000",
"AI agent": "AI agent",
"Take a look at the current document": "Take a look at the current document",
"AI agent is typing…": "AI agent is typing…",

View File

@@ -705,7 +705,10 @@
"Copy chat": "Копировать чат",
"Created successfully": "Успешно создано",
"Current context size": "Текущий размер контекста",
"Tokens generated this turn": "Токенов сгенерировано за ход",
"Context size / model limit": "Размер контекста / лимит модели",
"Context window (tokens)": "Размер окна контекста (токены)",
"Shows used / total in the chat header badge; empty hides the total.": "Показывает использовано/всего в шапке чата; пусто — скрыть лимит.",
"e.g. 200000": "напр. 200000",
"Delete this chat?": "Удалить этот чат?",
"Deleted successfully": "Успешно удалено",
"Edited by AI agent on behalf of {{name}}": "Отредактировано AI-агентом от имени {{name}}",

View File

@@ -6,7 +6,7 @@ import {
useRef,
useState,
} from "react";
import { Group, Loader, Tooltip } from "@mantine/core";
import { Group, Loader } from "@mantine/core";
import {
IconArrowsDiagonal,
IconCheck,
@@ -39,6 +39,7 @@ import {
} from "@/features/ai-chat/queries/ai-chat-query.ts";
import ConversationList from "@/features/ai-chat/components/conversation-list.tsx";
import ChatThread from "@/features/ai-chat/components/chat-thread.tsx";
import { ContextBadge } from "@/features/ai-chat/components/context-badge.tsx";
import { exportAiChat } from "@/features/ai-chat/services/ai-chat-service.ts";
import { useChatSession } from "@/features/ai-chat/hooks/use-chat-session.ts";
import {
@@ -60,13 +61,6 @@ const MIN_HEIGHT = 400;
// Margin kept between the window and the viewport edges while dragging.
const EDGE_MARGIN = 8;
/** Compact token formatter: 1.2M / 3.4k / 950. */
function formatTokens(n: number): string {
if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
if (n >= 1_000) return `${(n / 1_000).toFixed(1)}k`;
return String(n);
}
// Compute the initial top-right placement at the default size, fitted to the
// current viewport. Reads `window` only when called (inside an effect).
function computeInitialGeom() {
@@ -161,12 +155,6 @@ export default function AiChatWindow() {
const { data: messageRows, isLoading: messagesLoading } =
useAiChatMessagesQuery(activeChatId ?? undefined);
// Live turn-token total (reasoning + output) for the in-flight turn, pushed up
// (THROTTLED to ~8 Hz inside ChatThread) so the header badge ticks mid-stream.
// `null` means no turn is in flight -> the badge falls back to the persisted
// context size below.
const [liveTurnTokens, setLiveTurnTokens] = useState<number | null>(null);
// The page the user is currently viewing. AiChatWindow lives in a pathless
// parent layout route, so useParams() can't see :pageSlug. Match the full
// pathname against the authenticated page route instead so "the current page"
@@ -306,6 +294,21 @@ export default function AiChatWindow() {
return 0;
}, [activeChatId, messageRows]);
// The model's context-window size (badge denominator), read from the most
// recent assistant row that carries it. Admin-configured in AI settings and
// stamped onto the turn server-side, so it travels with the message metadata —
// no client-side model resolution, and it survives public shares / per-role
// models automatically. 0 (no limit configured, or older rows) → the badge
// hides the denominator and shows only the current context size.
const maxContextTokens = useMemo(() => {
if (!activeChatId || !messageRows) return 0;
for (let i = messageRows.length - 1; i >= 0; i--) {
const max = messageRows[i].metadata?.maxContextTokens;
if (typeof max === "number" && max > 0) return max;
}
return 0;
}, [activeChatId, messageRows]);
// On (re)open, settle the geometry before paint (useLayoutEffect → no
// first-frame jump): compute an initial top-right placement the first time,
// and re-clamp an existing geometry to the current viewport on later opens
@@ -495,23 +498,14 @@ export default function AiChatWindow() {
)}
<div style={{ flex: 1, display: "flex", justifyContent: "center" }}>
{/* While a turn streams, show the LIVE turn-token count (ticks ~8 Hz);
once it finishes, fall back to the persisted context size. Require
> 0 so the very first emit (an empty tail message, count 0) does not
flash a "0" badge before any token streams in (#151 review). */}
{liveTurnTokens !== null && liveTurnTokens > 0 ? (
<Tooltip label={t("Tokens generated this turn")} withArrow>
<span className={classes.badge}>
{formatTokens(liveTurnTokens)}
</span>
</Tooltip>
) : contextTokens > 0 ? (
<Tooltip label={t("Current context size")} withArrow>
<span className={classes.badge}>
{formatTokens(contextTokens)}
</span>
</Tooltip>
) : null}
{/* Context badge: always "current / max" context size (or just current
when no model limit is configured). It no longer flips to a live
per-turn generation counter mid-stream — that live feedback lives in
the chat body's "Thinking · N tokens" block. */}
<ContextBadge
contextTokens={contextTokens}
maxContextTokens={maxContextTokens}
/>
</div>
<div style={{ display: "flex", alignItems: "center", gap: 1 }}>
@@ -634,7 +628,6 @@ export default function AiChatWindow() {
assistantName={currentRole?.name}
onTurnFinished={onTurnFinished}
onServerChatId={onServerChatId}
onLiveTurnTokens={setLiveTurnTokens}
/>
)}
</div>

View File

@@ -20,7 +20,6 @@ import {
} from "@/features/ai-chat/utils/role-launch.ts";
import { describeChatError } from "@/features/ai-chat/utils/error-message.ts";
import { extractServerChatId } from "@/features/ai-chat/utils/adopt-chat-id.ts";
import { liveTurnTokens } from "@/features/ai-chat/utils/count-stream-tokens.ts";
import {
dequeue,
enqueueMessage,
@@ -67,12 +66,6 @@ interface ChatThreadProps {
* Copy/export button available mid-stream). Distinct from onTurnFinished,
* which fires only at the terminal outcome. */
onServerChatId?: (serverChatId?: string) => void;
/** Reports the live turn-token total (reasoning + output) for the in-flight
* turn so the parent can show a header badge that ticks mid-stream. THROTTLED
* here (~8 Hz) so the parent re-renders a handful of times a second, not on
* every streamed delta. Called with `null` when no turn is in flight (the
* parent then reverts the badge to the persisted context size). */
onLiveTurnTokens?: (tokens: number | null) => void;
}
/**
@@ -117,7 +110,6 @@ export default function ChatThread({
assistantName,
onTurnFinished,
onServerChatId,
onLiveTurnTokens,
}: ChatThreadProps) {
const { t } = useTranslation();
@@ -328,53 +320,6 @@ export default function ChatThread({
// the SAME on-screen banner text can be mirrored into the export (issue #160).
const errorView = error ? describeChatError(error.message ?? "", t) : null;
// Report the live turn-token total to the parent header badge, THROTTLED to
// ~8 Hz so the parent re-renders a few times a second instead of on every
// streamed delta. The tail assistant message's reasoning+output (estimate while
// streaming, authoritative once a step reports usage) is the live figure. When
// the turn ends we emit a final exact value, then `null` so the parent reverts
// the badge to the persisted context size.
const lastEmitRef = useRef(0);
const emitTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
useEffect(() => {
if (!onLiveTurnTokens) return;
if (!isStreaming) {
// Turn ended (or never started): clear any pending throttle and revert.
if (emitTimerRef.current) {
clearTimeout(emitTimerRef.current);
emitTimerRef.current = null;
}
lastEmitRef.current = 0;
onLiveTurnTokens(null);
return;
}
const tail = messages[messages.length - 1];
const live = tail?.role === "assistant" ? liveTurnTokens(tail) : null;
const total = live ? live.reasoning + live.output : 0;
const now = Date.now();
const MIN_INTERVAL = 120; // ms (~8 Hz)
const elapsed = now - lastEmitRef.current;
if (elapsed >= MIN_INTERVAL) {
lastEmitRef.current = now;
onLiveTurnTokens(total);
} else if (!emitTimerRef.current) {
// Schedule a trailing emit so the FINAL value of a burst is not dropped.
emitTimerRef.current = setTimeout(() => {
emitTimerRef.current = null;
lastEmitRef.current = Date.now();
onLiveTurnTokens(total);
}, MIN_INTERVAL - elapsed);
}
}, [messages, isStreaming, onLiveTurnTokens]);
// Clear any pending throttle timer on unmount (chat switch via `key`) so a
// trailing emit can't fire into a torn-down thread's parent.
useEffect(() => {
return () => {
if (emitTimerRef.current) clearTimeout(emitTimerRef.current);
};
}, []);
// A role was picked with autoStart=false: the role is bound but NOTHING was
// sent, so chatId stays null and the empty state would keep showing the cards.
// This flag hides the cards and reveals the composer (with the role indicated)

View File

@@ -0,0 +1,69 @@
import { describe, it, expect } from "vitest";
import { render, screen, fireEvent } from "@testing-library/react";
import { MantineProvider } from "@mantine/core";
import { ContextBadge, formatTokens } from "./context-badge";
// matchMedia (read by MantineProvider) is stubbed globally in vitest.setup.ts.
// Without an I18nextProvider, `t(key)` returns the key verbatim, so tooltip
// labels assert against their English source strings.
function renderBadge(props: {
contextTokens: number;
maxContextTokens?: number;
}) {
return render(
<MantineProvider>
<ContextBadge {...props} />
</MantineProvider>,
);
}
describe("formatTokens", () => {
it("formats with k / M suffixes", () => {
expect(formatTokens(572)).toBe("572");
expect(formatTokens(200_000)).toBe("200.0k");
expect(formatTokens(1_500_000)).toBe("1.5M");
});
});
describe("ContextBadge", () => {
it("shows `current / max` when a limit is configured", () => {
renderBadge({ contextTokens: 572, maxContextTokens: 200_000 });
expect(screen.getByText("572 / 200.0k")).toBeDefined();
});
it("shows only the current size when no limit is configured", () => {
renderBadge({ contextTokens: 572, maxContextTokens: 0 });
expect(screen.getByText("572")).toBeDefined();
// No denominator rendered.
expect(screen.queryByText(/\//)).toBeNull();
});
it("treats an undefined limit as no limit", () => {
renderBadge({ contextTokens: 1234 });
expect(screen.getByText("1.2k")).toBeDefined();
expect(screen.queryByText(/\//)).toBeNull();
});
it("renders nothing until there is a current context size", () => {
const { container } = renderBadge({
contextTokens: 0,
maxContextTokens: 200_000,
});
expect(container.querySelector("span")).toBeNull();
});
it("never flips to a live per-turn counter (no live mode); shows context as-is even above max", () => {
// `current > max` (estimate drift / smaller-model role) is shown unclamped.
renderBadge({ contextTokens: 210_000, maxContextTokens: 200_000 });
expect(screen.getByText("210.0k / 200.0k")).toBeDefined();
});
it("exposes the limit tooltip label on hover", async () => {
renderBadge({ contextTokens: 572, maxContextTokens: 200_000 });
fireEvent.mouseEnter(screen.getByText("572 / 200.0k"));
expect(
await screen.findByText("Context size / model limit"),
).toBeDefined();
});
});

View File

@@ -0,0 +1,61 @@
import { Tooltip } from "@mantine/core";
import { useTranslation } from "react-i18next";
import classes from "@/features/ai-chat/components/ai-chat-window.module.css";
/** Compact token formatter: 1.2M / 3.4k / 950. */
export function formatTokens(n: number): string {
if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
if (n >= 1_000) return `${(n / 1_000).toFixed(1)}k`;
return String(n);
}
interface ContextBadgeProps {
// Current context size for the active chat (tokens occupied in the model's
// window). 0 = unknown → nothing is rendered.
contextTokens: number;
// The model's context-window size (tokens), from AI settings. 0/undefined =
// no limit known → only the current size is shown (no denominator).
maxContextTokens?: number;
}
/**
* Header badge that ALWAYS shows the current context size, and — when the model's
* context-window size is configured — appends "/ max" so the badge reads
* "current / max" (e.g. `572 / 200k`). This is a single, stable meaning: unlike
* the previous design it never flips to a live per-turn generation counter while
* streaming (that live feedback lives in the chat body's "Thinking · N tokens").
*
* No limit configured (or older history rows without it) → the denominator is
* hidden and the badge shows the current size only, matching the prior at-rest
* behaviour. `context > max` (estimate drift, or a role on a smaller model) is
* shown as-is, without clamping.
*/
export function ContextBadge({
contextTokens,
maxContextTokens,
}: ContextBadgeProps) {
const { t } = useTranslation();
// Nothing to show until the first persisted context figure exists.
if (!(contextTokens > 0)) return null;
const hasMax = typeof maxContextTokens === "number" && maxContextTokens > 0;
const label = hasMax
? `${formatTokens(contextTokens)} / ${formatTokens(maxContextTokens)}`
: formatTokens(contextTokens);
return (
<Tooltip
label={
hasMax
? t("Context size / model limit")
: t("Current context size")
}
withArrow
>
<span className={classes.badge}>{label}</span>
</Tooltip>
);
}
export default ContextBadge;

View File

@@ -113,9 +113,14 @@ export interface IAiChatMessageRow {
};
// Current context size for the turn = final-step (input+output) tokens, i.e.
// how much the conversation occupies in the model's context window after this
// turn. Distinct from `usage` (legacy cumulative totalUsage). Shown in the
// floating window's header badge.
// turn. Distinct from `usage` (legacy cumulative totalUsage). Shown as the
// numerator of the floating window's "current / max" header badge.
contextTokens?: number;
// The model's context-window size (tokens), admin-configured in AI settings
// and stamped onto the turn server-side. The denominator of the header badge.
// Absent/0 (older rows, or no limit configured) → the badge hides the
// denominator and shows only the current context size (`contextTokens`).
maxContextTokens?: number;
// Set on an assistant row whose turn ended in a provider/stream error; the
// raw provider error text (e.g. "402: ...") for inline display in the thread.
error?: string;

View File

@@ -1,17 +1,5 @@
import { describe, expect, it } from "vitest";
import type { UIMessage } from "@ai-sdk/react";
import {
estimateTokens,
liveTurnTokens,
} from "@/features/ai-chat/utils/count-stream-tokens.ts";
const msg = (parts: unknown[], metadata?: unknown): UIMessage =>
({
id: Math.random().toString(),
role: "assistant",
parts,
metadata,
}) as UIMessage;
import { estimateTokens } from "@/features/ai-chat/utils/count-stream-tokens.ts";
describe("estimateTokens", () => {
it("returns 0 for the empty string", () => {
@@ -25,147 +13,3 @@ describe("estimateTokens", () => {
expect(estimateTokens("12345678")).toBe(2);
});
});
describe("liveTurnTokens — estimate path", () => {
it("is all zeros for an undefined message", () => {
expect(liveTurnTokens(undefined)).toEqual({
reasoning: 0,
output: 0,
authoritative: false,
});
});
it("is all zeros for a parts-less message", () => {
expect(liveTurnTokens({ id: "x", role: "assistant" } as UIMessage)).toEqual({
reasoning: 0,
output: 0,
authoritative: false,
});
});
it("estimates output from text parts", () => {
// 8 chars -> 2 tokens.
const r = liveTurnTokens(msg([{ type: "text", text: "12345678" }]));
expect(r).toEqual({ reasoning: 0, output: 2, authoritative: false });
});
it("estimates reasoning from reasoning parts (kept separate from output)", () => {
const r = liveTurnTokens(
msg([
{ type: "reasoning", text: "12345678" },
{ type: "text", text: "abcd" },
]),
);
expect(r).toEqual({ reasoning: 2, output: 1, authoritative: false });
});
it("accumulates across multiple text + reasoning parts (multi-step)", () => {
const r = liveTurnTokens(
msg([
{ type: "reasoning", text: "abcd" }, // 1
{ type: "text", text: "abcd" }, // 1
{ type: "tool-getPage", state: "output-available" }, // ignored
{ type: "reasoning", text: "abcd" }, // 1
{ type: "text", text: "abcdefgh" }, // 2
]),
);
expect(r).toEqual({ reasoning: 2, output: 3, authoritative: false });
});
it("ignores non text/reasoning parts (tools, step-start)", () => {
const r = liveTurnTokens(
msg([
{ type: "step-start" },
{ type: "tool-getPage", state: "input-available" },
]),
);
expect(r).toEqual({ reasoning: 0, output: 0, authoritative: false });
});
});
describe("liveTurnTokens — authoritative path", () => {
it("returns authoritative usage verbatim, splitting reasoning out of output", () => {
// outputTokens INCLUDES reasoning in the AI SDK shape -> answer = 100 - 30.
const r = liveTurnTokens(
msg([{ type: "text", text: "estimate would be tiny" }], {
usage: { inputTokens: 500, outputTokens: 100, reasoningTokens: 30 },
}),
);
expect(r).toEqual({ reasoning: 30, output: 70, authoritative: true });
});
it("treats missing reasoningTokens as 0 and keeps full output", () => {
const r = liveTurnTokens(
msg([{ type: "text", text: "x" }], {
usage: { inputTokens: 10, outputTokens: 42 },
}),
);
expect(r).toEqual({ reasoning: 0, output: 42, authoritative: true });
});
it("never returns a negative output when reasoning exceeds reported output", () => {
const r = liveTurnTokens(
msg([], { usage: { outputTokens: 10, reasoningTokens: 40 } }),
);
expect(r).toEqual({ reasoning: 40, output: 0, authoritative: true });
});
it("falls back to the estimate when metadata has no usage object", () => {
const r = liveTurnTokens(
msg([{ type: "text", text: "abcd" }], { chatId: "c1" }),
);
expect(r).toEqual({ reasoning: 0, output: 1, authoritative: false });
});
});
describe("liveTurnTokens — combined authoritative + estimate (#163)", () => {
it("ticks the in-flight step above the completed-steps authoritative base", () => {
// The authoritative usage is the sum over COMPLETED steps (step 1). The
// CURRENT step is streaming and its text is NOT in `usage` yet, but it IS in
// the parts -> the running estimate must push the live figure above the base
// so the badge keeps growing between step boundaries.
const longText = "x".repeat(800); // 800 chars -> 200 est output tokens
const r = liveTurnTokens(
msg([{ type: "text", text: longText }], {
usage: { inputTokens: 500, outputTokens: 40 }, // step-1 base: 40 output
}),
);
// max(authOutput=40, estOutput=200) = 200 -> the counter ticks, not frozen.
expect(r.output).toBe(200);
expect(r.authoritative).toBe(true);
});
it("ticks reasoning of the in-flight step above the authoritative reasoning base", () => {
const longReasoning = "r".repeat(400); // 400 chars -> 100 est reasoning
const r = liveTurnTokens(
msg([{ type: "reasoning", text: longReasoning }], {
usage: { inputTokens: 100, outputTokens: 20, reasoningTokens: 20 },
}),
);
// reasoning: max(20, 100) = 100 ; output: max(max(0,20-20)=0, 0) = 0.
expect(r.reasoning).toBe(100);
expect(r.output).toBe(0);
expect(r.authoritative).toBe(true);
});
it("snaps to the authoritative figure once it exceeds the rough estimate", () => {
// Short on-screen text (estimate tiny) but a large authoritative output:
// the exact figure wins at the boundary (the counter never under-reports).
const r = liveTurnTokens(
msg([{ type: "text", text: "abcd" }], {
usage: { inputTokens: 10, outputTokens: 5000 },
}),
);
expect(r.output).toBe(5000);
});
it("is monotonic: max never drops below the authoritative base when the estimate is smaller", () => {
// Mirrors the legacy 'verbatim' tests: estimate < authoritative -> unchanged.
const r = liveTurnTokens(
msg([{ type: "text", text: "tiny" }], {
usage: { inputTokens: 500, outputTokens: 100, reasoningTokens: 30 },
}),
);
expect(r).toEqual({ reasoning: 30, output: 70, authoritative: true });
});
});

View File

@@ -1,18 +1,16 @@
import type { UIMessage } from "@ai-sdk/react";
/**
* Live token counting for a streaming AI-chat turn — split into REASONING
* (thinking) and OUTPUT (answer) tokens, mirroring how Claude Code shows
* `Thinking… · 60 tokens` next to its thinking indicator.
* Live token ESTIMATION for a streaming AI-chat turn.
*
* No provider streams exact per-token usage mid-stream, so the live number is a
* CLIENT ESTIMATE (chars/≈4 heuristic) that is reconciled to AUTHORITATIVE usage
* once the server attaches it on a step/turn boundary (see the server's
* `chatStreamMetadata` + the client's read of `message.metadata.usage`). When
* authoritative usage is present we return it verbatim (the number "jumps to
* exact"); otherwise we return the running estimate. Pure + unit-testable: it
* never runs a real BPE tokenizer (that would be O(n²) on the hot path, bloat the
* CLIENT ESTIMATE (chars/≈4 heuristic). It powers the chat body's
* `Thinking… · N tokens` indicator (see `ReasoningBlock`), which reconciles to
* the authoritative server usage once it lands. Pure + unit-testable: it never
* runs a real BPE tokenizer (that would be O(n²) on the hot path, bloat the
* bundle, and be wrong for Gemini/Ollama anyway).
*
* The former header-badge `liveTurnTokens()` split was removed with #189 (the
* header badge now shows the stable "current / max" context size, not a live
* per-turn counter); the live feedback remains in `ReasoningBlock`.
*/
/**
@@ -24,90 +22,3 @@ export function estimateTokens(text: string): number {
if (!text) return 0;
return Math.ceil(text.length / 4);
}
/** Authoritative per-step/turn usage the server attaches to message metadata. */
export interface AuthoritativeUsage {
inputTokens?: number;
outputTokens?: number;
totalTokens?: number;
reasoningTokens?: number;
}
/** Live token split for a turn's tail (streaming) assistant message. */
export interface LiveTurnTokens {
/** Thinking/reasoning tokens (estimate, or authoritative when available). */
reasoning: number;
/** Answer/output tokens (estimate, or authoritative when available). */
output: number;
/** True when the numbers come from authoritative server usage, not estimate. */
authoritative: boolean;
}
/** Read the authoritative usage off a UIMessage's metadata, if the server set it. */
function metadataUsage(message: UIMessage): AuthoritativeUsage | undefined {
const meta = message?.metadata as
| { usage?: AuthoritativeUsage }
| undefined;
const usage = meta?.usage;
if (!usage || typeof usage !== "object") return undefined;
return usage;
}
/**
* Token split for the given (streaming) assistant message.
*
* COMBINES the authoritative server usage with the running text estimate so the
* counter ticks in real time AND lands exact. The server only attaches
* `metadata.usage` at a step/turn boundary (`finish-step`/`finish`) and it is
* CUMULATIVE over COMPLETED steps — it does NOT yet include the in-flight step.
* So a multi-step turn that returned the authoritative figure verbatim would
* FREEZE between boundaries and jump in steps (issue #163).
*
* Instead we always compute the running ESTIMATE (chars/≈4 over the message's
* `reasoning`/`text` parts, which grows on every streamed delta) and take the
* per-component MAX of the authoritative base and the estimate:
* - between boundaries the estimate of the in-flight step ticks the number up;
* - at a boundary the authoritative figure snaps it to exact;
* - because the server's usage is cumulative and we only ever take the max, the
* number is MONOTONIC — it never drops.
*
* Providers that don't stream reasoning text still surface a reasoning count once
* the authoritative usage arrives (`max(reasoningTokens, 0)`); on the pure
* estimate path (no usage yet) such a turn shows `reasoning: 0` until then.
*/
export function liveTurnTokens(message: UIMessage | undefined): LiveTurnTokens {
if (!message) return { reasoning: 0, output: 0, authoritative: false };
// Running ESTIMATE over every reasoning/text part — grows on each delta. This
// includes the IN-FLIGHT step, which the authoritative usage does not cover yet.
let estReasoning = 0;
let estOutput = 0;
for (const part of message.parts ?? []) {
if (part.type === "reasoning") {
estReasoning += estimateTokens((part as { text?: string }).text ?? "");
} else if (part.type === "text") {
estOutput += estimateTokens((part as { text?: string }).text ?? "");
}
}
const usage = metadataUsage(message);
if (!usage) {
// No authoritative usage streamed yet: the estimate IS the live figure.
return { reasoning: estReasoning, output: estOutput, authoritative: false };
}
// Authoritative sum over COMPLETED steps. `outputTokens` already INCLUDES
// reasoning in the AI SDK usage shape, so subtract it out for the "answer"
// figure (never go negative if a provider reports them inconsistently).
const authReasoning = usage.reasoningTokens ?? 0;
const authOutput = Math.max(0, (usage.outputTokens ?? 0) - authReasoning);
// Per-component max: the in-flight step's estimate ticks above the completed-
// steps base between boundaries, and the authoritative figure wins once it
// exceeds the (rough) estimate at the next boundary. Monotonic by construction.
return {
reasoning: Math.max(authReasoning, estReasoning),
output: Math.max(authOutput, estOutput),
authoritative: true,
};
}

View File

@@ -7,6 +7,7 @@ import {
Button,
Group,
Modal,
NumberInput,
Paper,
PasswordInput,
Select,
@@ -85,6 +86,9 @@ const formSchema = z.object({
chatModel: z.string(),
// Chat provider implementation (reasoning surfacing). Default openai-compatible.
chatApiStyle: z.enum(["openai-compatible", "openai"]),
// Model context-window size (tokens) shown as the chat header badge's "max".
// Empty string = no limit (NumberInput emits "" when cleared).
chatContextWindow: z.union([z.number(), z.literal("")]),
// Cheap model id for the anonymous public-share assistant; empty = use chatModel.
publicShareChatModel: z.string(),
// Agent-role id whose persona the public-share assistant adopts; empty =
@@ -312,6 +316,7 @@ export default function AiProviderSettings() {
initialValues: {
chatModel: "",
chatApiStyle: "openai-compatible" as ChatApiStyle,
chatContextWindow: "" as number | "",
publicShareChatModel: "",
publicShareAssistantRoleId: "",
embeddingModel: "",
@@ -335,6 +340,10 @@ export default function AiProviderSettings() {
form.setValues({
chatModel: settings.chatModel ?? "",
chatApiStyle: settings.chatApiStyle ?? "openai-compatible",
// 0/unset = no limit → show an empty field (not a literal "0").
chatContextWindow: settings.chatContextWindow
? settings.chatContextWindow
: "",
publicShareChatModel: settings.publicShareChatModel ?? "",
publicShareAssistantRoleId: settings.publicShareAssistantRoleId ?? "",
embeddingModel: settings.embeddingModel ?? "",
@@ -365,6 +374,11 @@ export default function AiProviderSettings() {
driver: "openai",
chatModel: values.chatModel,
chatApiStyle: values.chatApiStyle,
// Empty → 0, which clears the limit server-side (badge shows current only).
chatContextWindow:
typeof values.chatContextWindow === "number"
? values.chatContextWindow
: 0,
// Cheap model id for the anonymous public-share assistant; empty falls
// back to chatModel server-side.
publicShareChatModel: values.publicShareChatModel,
@@ -785,6 +799,22 @@ export default function AiProviderSettings() {
{...form.getInputProps("chatApiStyle")}
/>
<NumberInput
mt="sm"
label={t("Context window (tokens)")}
description={t(
"Shows used / total in the chat header badge; empty hides the total.",
)}
placeholder={t("e.g. 200000")}
min={0}
step={1000}
allowDecimal={false}
allowNegative={false}
thousandSeparator=" "
disabled={isLoading}
{...form.getInputProps("chatContextWindow")}
/>
{/* Anonymous public-share assistant: a single master toggle + an
optional cheaper model id. Reuses this card's driver/URL/key. */}
<Group justify="space-between" align="center" wrap="nowrap" mt="md">

View File

@@ -23,6 +23,9 @@ export interface IAiSettings {
driver?: AiDriver;
chatModel?: string;
chatApiStyle?: ChatApiStyle;
// Chat model context-window size (tokens); shown as the "max" in the chat
// header context badge. 0/unset = no limit (badge shows the current size only).
chatContextWindow?: number;
// Cheap model id for the anonymous public-share assistant; empty = chatModel.
publicShareChatModel?: string;
// Agent-role id whose persona the public-share assistant adopts; empty =
@@ -57,6 +60,8 @@ export interface IAiSettingsUpdate {
driver?: AiDriver;
chatModel?: string;
chatApiStyle?: ChatApiStyle;
// Chat model context-window size (tokens); 0 clears the limit.
chatContextWindow?: number;
publicShareChatModel?: string;
// Agent-role id whose persona the public-share assistant adopts; empty =
// built-in locked persona.

View File

@@ -292,6 +292,26 @@ describe('flushAssistant', () => {
expect(f.metadata.contextTokens).toBe(15);
});
it('completed: writes maxContextTokens when the model limit is > 0', () => {
const f = flushAssistant([toolStep], '', 'completed', {
contextTokens: 15,
maxContextTokens: 200_000,
});
expect(f.metadata.maxContextTokens).toBe(200_000);
});
it('omits maxContextTokens when the limit is unset or 0', () => {
const unset = flushAssistant([toolStep], '', 'completed', {
contextTokens: 15,
});
expect('maxContextTokens' in unset.metadata).toBe(false);
const zero = flushAssistant([toolStep], '', 'completed', {
contextTokens: 15,
maxContextTokens: 0,
});
expect('maxContextTokens' in zero.metadata).toBe(false);
});
it('error: records the error and a derived finishReason', () => {
const f = flushAssistant([], 'partial answer', 'error', { error: 'boom' });
expect(f.status).toBe('error');

View File

@@ -616,6 +616,9 @@ export class AiChatService implements OnModuleInit {
contextTokens:
(usage?.inputTokens ?? 0) + (usage?.outputTokens ?? 0) ||
undefined,
// Admin-configured context-window size for this model (badge max).
// Resolved once per turn above; written to metadata only when > 0.
maxContextTokens: resolved?.chatContextWindow,
}),
);
// Lifecycle: release the external MCP clients leased for this turn.
@@ -1223,6 +1226,10 @@ export function flushAssistant(
finishReason?: string;
usage?: ChatStreamUsage | StreamUsage | undefined;
contextTokens?: number;
// Admin-configured context-window size (tokens) for this turn's model; the
// denominator of the client's "current / max" header badge. Written only
// when > 0 (0/unset = no limit known → the badge shows current only).
maxContextTokens?: number;
error?: string;
},
): AssistantFlush {
@@ -1253,6 +1260,9 @@ export function flushAssistant(
normalizeStreamUsage(extra.usage as StreamUsage) ?? extra.usage;
}
if (extra?.contextTokens) metadata.contextTokens = extra.contextTokens;
if (extra?.maxContextTokens && extra.maxContextTokens > 0) {
metadata.maxContextTokens = extra.maxContextTokens;
}
if (extra?.error) metadata.error = extra.error;
return {

View File

@@ -120,26 +120,21 @@ describe('AiChatToolsService deletePage guardrail (H4)', () => {
const tools = await buildTools();
const deletePage = tools.deletePage;
// inputSchema is now an AI SDK `Schema` (not a raw zod object). Its
// `validate` runs the same zod safeParse and forwards the STRIPPED data, so
// a permanent/force flag is never part of the validated input the SDK then
// hands to execute.
// The Zod input schema only allows `pageId`; parsing strips/ignores extra
// keys, so a permanent/force flag is never part of the validated input.
const schema = (deletePage as unknown as { inputSchema: unknown })
.inputSchema as {
validate: (
v: unknown,
) => Promise<{ success: boolean; value?: Record<string, unknown> }>;
parse: (v: unknown) => Record<string, unknown>;
};
const result = await schema.validate({
const parsed = schema.parse({
pageId: 'page-789',
permanentlyDelete: true,
forceDelete: true,
});
expect(result.success).toBe(true);
expect(result.value).toHaveProperty('pageId', 'page-789');
expect(result.value).not.toHaveProperty('permanentlyDelete');
expect(result.value).not.toHaveProperty('forceDelete');
expect(parsed).toHaveProperty('pageId', 'page-789');
expect(parsed).not.toHaveProperty('permanentlyDelete');
expect(parsed).not.toHaveProperty('forceDelete');
});
});
@@ -212,25 +207,21 @@ describe('AiChatToolsService expanded toolset guardrails', () => {
const tools = await buildTools();
const transformPage = tools.transformPage;
// inputSchema is now an AI SDK `Schema`; its `validate` runs the same zod
// safeParse, which only allows pageId/transformJs/dryRun and strips unknown
// keys — so deleteComments can never reach the client.
// The Zod input schema only allows pageId/transformJs/dryRun; parsing
// strips unknown keys, so deleteComments can never reach the client.
const schema = (transformPage as unknown as { inputSchema: unknown })
.inputSchema as {
validate: (
v: unknown,
) => Promise<{ success: boolean; value?: Record<string, unknown> }>;
parse: (v: unknown) => Record<string, unknown>;
};
const result = await schema.validate({
const parsed = schema.parse({
pageId: 'p',
transformJs: '(d)=>d',
dryRun: true,
deleteComments: true,
});
expect(result.success).toBe(true);
expect(result.value).toHaveProperty('pageId', 'p');
expect(result.value).not.toHaveProperty('deleteComments');
expect(parsed).toHaveProperty('pageId', 'p');
expect(parsed).not.toHaveProperty('deleteComments');
});
});

View File

@@ -15,7 +15,6 @@ import {
} from './docmost-client.loader';
import { resolveCurrentPageResult } from './current-page.util';
import { parseNodeArg } from './parse-node-arg';
import { modelFriendlyInput } from './model-friendly-input';
/**
* Per-user, per-request adapter that exposes Docmost READ operations to the
@@ -103,9 +102,9 @@ export class AiChatToolsService {
): Tool =>
tool({
description: spec.description,
inputSchema: modelFriendlyInput(
spec.buildShape ? (spec.buildShape(z) as z.ZodRawShape) : {},
),
inputSchema: spec.buildShape
? z.object(spec.buildShape(z) as z.ZodRawShape)
: z.object({}),
execute,
});
@@ -119,7 +118,7 @@ export class AiChatToolsService {
'and entities), not a full sentence. If the first results look weak ' +
'or incomplete, search again with different wording or synonyms ' +
'before answering.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
query: z.string().describe('The search query.'),
limit: z
.number()
@@ -228,7 +227,7 @@ export class AiChatToolsService {
'"the current page", or "here" refers to. Returns the page id and title, ' +
'or null if the user is not currently on a page. Call this first whenever ' +
'the user refers to the current page without giving an explicit id.',
inputSchema: modelFriendlyInput({}),
inputSchema: z.object({}),
execute: async () => resolveCurrentPageResult(openedPage),
}),
@@ -236,7 +235,7 @@ export class AiChatToolsService {
description:
'Fetch a single page as Markdown by its page id. Returns the page ' +
'title and its Markdown content.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id (or slugId) of the page.'),
}),
execute: async ({ pageId }) => {
@@ -260,7 +259,7 @@ export class AiChatToolsService {
'Create a new page with a Markdown body in a space, optionally under ' +
'a parent page. Returns the new page id and title. Reversible: a page ' +
'can be moved to trash later.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
title: z.string().describe('The title of the new page.'),
content: z
.string()
@@ -295,7 +294,7 @@ export class AiChatToolsService {
description:
"Replace a page's body with new Markdown content (and optionally its " +
'title). Reversible: the previous version is kept in page history.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page to update.'),
content: z.string().describe('The new page body as Markdown.'),
title: z
@@ -317,7 +316,7 @@ export class AiChatToolsService {
description:
"Rename a page (change its title only; the body is untouched). " +
'Reversible: rename back at any time.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page to rename.'),
title: z.string().describe('The new title.'),
}),
@@ -332,7 +331,7 @@ export class AiChatToolsService {
description:
'Move a page under a new parent page, or to the space root when no ' +
'parent is given. Reversible: move it back at any time.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page to move.'),
parentPageId: z
.string()
@@ -354,7 +353,7 @@ export class AiChatToolsService {
description:
'Move a page to the trash (SOFT delete only — fully reversible; the ' +
'page can be restored from trash). This NEVER permanently deletes.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page to move to trash.'),
}),
// GUARDRAIL (§14 H4): the only field ever passed to the client is
@@ -380,7 +379,7 @@ export class AiChatToolsService {
'"selection not found" error, retry with a corrected EXACT selection ' +
'copied verbatim from a single paragraph/block. Reversible via the ' +
'comment UI.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page to comment on.'),
content: z.string().describe('The comment body as Markdown.'),
selection: z
@@ -429,7 +428,7 @@ export class AiChatToolsService {
description:
'Resolve or reopen a top-level comment thread (reversible — toggle ' +
'the resolved flag). Only top-level comments can be resolved.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
commentId: z
.string()
.describe('The id of the top-level comment to resolve/reopen.'),
@@ -461,7 +460,7 @@ export class AiChatToolsService {
'List the most recent pages, optionally scoped to a single space. ' +
'Returns a bounded list (default 50, max 100). Pass tree:true (with ' +
"spaceId) to instead get the space's full page hierarchy as a nested tree.",
inputSchema: modelFriendlyInput({
inputSchema: z.object({
spaceId: z
.string()
.optional()
@@ -489,7 +488,7 @@ export class AiChatToolsService {
'List sidebar pages for a space. With no pageId, returns the ' +
"space's ROOT pages; with a pageId, returns that page's direct " +
'CHILDREN.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
spaceId: z.string().describe('The id of the space.'),
pageId: z
.string()
@@ -521,7 +520,7 @@ export class AiChatToolsService {
description:
'Read a table as a matrix of cell texts (plus a parallel cellIds ' +
'matrix so cells can be addressed for rich edits).',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page.'),
tableRef: z
.string()
@@ -537,7 +536,7 @@ export class AiChatToolsService {
listComments: tool({
description:
'List all comments on a page (content as Markdown).',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page.'),
}),
execute: async ({ pageId }) => await client.listComments(pageId),
@@ -545,7 +544,7 @@ export class AiChatToolsService {
getComment: tool({
description: 'Fetch a single comment by id (content as Markdown).',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
commentId: z.string().describe('The id of the comment.'),
}),
execute: async ({ commentId }) => await client.getComment(commentId),
@@ -555,7 +554,7 @@ export class AiChatToolsService {
description:
'Find new comments across a space (optionally scoped to a subtree) ' +
'created after a given timestamp.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
spaceId: z.string().describe('The id of the space to scan.'),
since: z
.string()
@@ -587,7 +586,7 @@ export class AiChatToolsService {
description:
'Fetch a single page-history version including its lossless ' +
'ProseMirror content.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
historyId: z.string().describe('The id of the history version.'),
}),
execute: async ({ historyId }) =>
@@ -605,7 +604,7 @@ export class AiChatToolsService {
'Export a page to a single self-contained Docmost-flavoured ' +
'Markdown file (meta + body + comment threads). Lossless round-trip ' +
'with importPageMarkdown.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page to export.'),
}),
execute: async ({ pageId }) => {
@@ -631,7 +630,7 @@ export class AiChatToolsService {
'{"type":"text","text":"x","marks":[{"type":"bold"}]}. The node arg ' +
'may be a JSON object or a JSON string (both accepted). Reversible: ' +
'the previous version is kept in page history.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page.'),
nodeId: z
.string()
@@ -664,7 +663,7 @@ export class AiChatToolsService {
'{"type":"text","text":"x","marks":[{"type":"bold"}]}. The node arg ' +
'may be a JSON object or a JSON string (both accepted). Reversible ' +
'via page history.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page.'),
node: z
.any()
@@ -723,7 +722,7 @@ export class AiChatToolsService {
'object or a JSON string (both accepted). Omit content for a ' +
'title-only update. Reversible: the previous version is kept in page ' +
'history.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page to update.'),
content: z
.any()
@@ -754,7 +753,7 @@ export class AiChatToolsService {
description:
'Insert a row of plain-text cells into a table. Reversible via ' +
'page history.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page.'),
tableRef: z
.string()
@@ -773,7 +772,7 @@ export class AiChatToolsService {
tableDeleteRow: tool({
description:
'Delete a table row at a 0-based index. Reversible via page history.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page.'),
tableRef: z
.string()
@@ -788,7 +787,7 @@ export class AiChatToolsService {
description:
'Set the plain-text content of a table cell at [row, col] (0-based). ' +
'Reversible via page history.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page.'),
tableRef: z
.string()
@@ -818,7 +817,7 @@ export class AiChatToolsService {
'Make a page PUBLICLY accessible and return its public URL. ' +
'Reversible via unsharePage. Only share when the user explicitly ' +
'asked, since this exposes the page to anyone with the link.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page to share.'),
searchIndexing: z
.boolean()
@@ -845,7 +844,7 @@ export class AiChatToolsService {
"page's ProseMirror document for complex/scripted rewrites. dryRun " +
'(default true) previews a diff WITHOUT writing; set dryRun:false to ' +
'apply. Reversible: applying creates a new page-history snapshot.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z.string().describe('The id of the page to transform.'),
transformJs: z
.string()

View File

@@ -1,112 +0,0 @@
import { z } from 'zod';
import { modelFriendlyInput } from './model-friendly-input';
/**
* Unit tests for the model-friendly input wrapper (issue #190): validation
* failures must report a human-readable, parameter-naming message (not the raw
* zod text), successful validation must strip unknown keys (preserving the
* strip guardrails), and the JSON schema handed to the model must keep the
* required/optional contract and field descriptions intact.
*/
describe('modelFriendlyInput', () => {
// A representative shape: a required id + description, plus an optional field.
const shape = {
pageId: z.string().describe('The id of the page to comment on.'),
content: z.string(),
limit: z.number().int().optional(),
};
// The AI SDK `Schema` exposes a `validate` callback and a `jsonSchema` field;
// type them loosely for the test.
type SchemaLike = {
validate?: (
v: unknown,
) =>
| { success: boolean; value?: Record<string, unknown>; error?: Error }
| PromiseLike<{
success: boolean;
value?: Record<string, unknown>;
error?: Error;
}>;
jsonSchema: unknown;
};
it('reports a model-friendly error naming the missing REQUIRED param + retry hint', async () => {
const schema = modelFriendlyInput(shape) as unknown as SchemaLike;
// Drop the required `pageId` (the parallel-batch failure mode).
const result = await schema.validate!({ content: 'hi' });
expect(result.success).toBe(false);
const message = result.error?.message ?? '';
// Names the offending parameter by name.
expect(message).toContain('pageId');
// Carries the fixed actionable retry hint.
expect(message).toContain('Include every REQUIRED parameter and retry');
expect(message).toContain('do not drop ids like "pageId"');
// It must NOT be the bare raw zod text alone — our wrapper prefix is present.
expect(message).toContain('Invalid tool input');
});
it('accepts valid input and STRIPS unknown keys (keeps declared ones)', async () => {
const schema = modelFriendlyInput(shape) as unknown as SchemaLike;
const result = await schema.validate!({
pageId: 'p-1',
content: 'hello',
// An extra unknown key a (compromised) model might emit.
permanentlyDelete: true,
});
expect(result.success).toBe(true);
expect(result.value).toEqual({ pageId: 'p-1', content: 'hello' });
expect(result.value).not.toHaveProperty('permanentlyDelete');
});
it('produces a draft-07 JSON schema that preserves required + descriptions', async () => {
const schema = modelFriendlyInput(shape) as unknown as SchemaLike;
// jsonSchema may be a value or a promise; await either way.
const json = (await Promise.resolve(schema.jsonSchema)) as {
required?: string[];
properties?: Record<string, { description?: string }>;
};
// Required contract preserved: pageId + content required, limit optional.
expect(json.required).toEqual(expect.arrayContaining(['pageId', 'content']));
expect(json.required).not.toContain('limit');
// Field description preserved.
expect(json.properties?.pageId?.description).toBe(
'The id of the page to comment on.',
);
});
it('de-duplicates a parameter that produces MULTIPLE issues on the same path', async () => {
// A single field can fail several zod checks at once (here min-length AND a
// regex), yielding two issues with the SAME path. The friendly message must
// name that parameter only once (the `seen` dedup branch).
const multiIssueShape = {
code: z
.string()
.min(5)
.regex(/^[0-9]+$/),
};
const schema = modelFriendlyInput(
multiIssueShape,
) as unknown as SchemaLike;
// "ab" violates BOTH the min(5) and the digit-only regex.
const result = await schema.validate!({ code: 'ab' });
expect(result.success).toBe(false);
const message = result.error?.message ?? '';
// The parameter name appears exactly once despite two underlying issues.
const occurrences = message.split('parameter "code"').length - 1;
expect(occurrences).toBe(1);
});
it('handles a root-level type error with a "(root)" parameter name', async () => {
const schema = modelFriendlyInput(shape) as unknown as SchemaLike;
// Passing a non-object yields an issue with an empty path.
const result = await schema.validate!('not an object');
expect(result.success).toBe(false);
expect(result.error?.message).toContain('(root)');
});
});

View File

@@ -1,72 +0,0 @@
import { jsonSchema, type JSONSchema7, type Schema } from 'ai';
import { z } from 'zod';
// Centralized input-schema wrapper for in-app AI tools. The JSON schema handed
// to the model is derived from the same zod shape (so `required`/`description`/
// constraints are unchanged), but validation failures are reported with a
// human-readable message that NAMES the offending parameter(s) and asks the
// model to retry with every required field — instead of the raw zod text. This
// matters for parallel tool-call batches where the model tends to drop a
// repeated id like `pageId`.
// Fixed, actionable hint appended to every validation error. Kept as a constant
// so the message stays deterministic and the spec can assert on it verbatim.
const RETRY_HINT =
'Include every REQUIRED parameter and retry; when issuing parallel tool ' +
'calls, do not drop ids like "pageId".';
/**
* Turn a zod validation error into a concise, model-friendly message that names
* each offending parameter (by its dotted path; the root object is "(root)"),
* gives a short reason, and ends with the fixed retry hint. Repeated parameter
* names are de-duplicated and the output is deterministic.
*/
export function formatIssues(error: z.ZodError): string {
const seen = new Set<string>();
const parts: string[] = [];
for (const issue of error.issues) {
const name =
Array.isArray(issue.path) && issue.path.length > 0
? issue.path.join('.')
: '(root)';
if (seen.has(name)) continue;
seen.add(name);
// Prefer zod's own message (e.g. "Invalid input: expected string, received
// undefined"); fall back to a generic reason when it is missing.
const reason = issue.message ? issue.message : 'missing or invalid';
parts.push(`parameter "${name}": ${reason}`);
}
const summary = parts.length > 0 ? parts.join('; ') : 'invalid tool input';
return `Invalid tool input — ${summary}. ${RETRY_HINT}`;
}
/**
* Build an AI SDK `Schema` from a zod raw shape. The JSON schema exposed to the
* model is derived from the zod object (preserving `required`, `description`,
* and field constraints), so the required/optional contract is UNCHANGED. On a
* validation failure we return a model-friendly error (see `formatIssues`); on
* success we return the PARSED data, which has unknown keys stripped by zod —
* this preserves the existing strip guardrails (e.g. deletePage never forwards
* permanentlyDelete/forceDelete; transformPage never forwards deleteComments).
*/
export function modelFriendlyInput<Shape extends z.ZodRawShape>(
shape: Shape,
): Schema<z.infer<z.ZodObject<Shape>>> {
const object = z.object(shape);
// draft-07 JSON schema for the model (keeps required/description/constraints).
const schema = z.toJSONSchema(object, { target: 'draft-7' }) as JSONSchema7;
return jsonSchema<z.infer<typeof object>>(schema, {
validate: (value: unknown) => {
const result = object.safeParse(value);
if (result.success) {
// Return the PARSED (unknown-key-stripped) data so the SDK forwards a
// clean object to execute — preserves the existing strip guardrails.
return { success: true as const, value: result.data };
}
return {
success: false as const,
error: new Error(formatIssues(result.error)),
};
},
});
}

View File

@@ -5,7 +5,6 @@ import { ShareService } from '../../share/share.service';
import { SearchService } from '../../search/search.service';
import { PageRepo } from '@docmost/db/repos/page/page.repo';
import { jsonToMarkdown } from '../../../collaboration/collaboration.util';
import { modelFriendlyInput } from './model-friendly-input';
/**
* Isolated, READ-ONLY toolset for the ANONYMOUS public-share assistant.
@@ -53,7 +52,7 @@ export class PublicShareChatToolsService {
'(key terms and entities), not a full sentence. If the first ' +
'results look weak, search again with different wording before ' +
'answering. Only pages inside this share are ever returned.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
query: z.string().describe('The search query.'),
limit: z
.number()
@@ -88,7 +87,7 @@ export class PublicShareChatToolsService {
'Markdown, by its page id. Returns the page title and its Markdown ' +
'content. Only pages inside this share can be read; reading any ' +
'other page fails.',
inputSchema: modelFriendlyInput({
inputSchema: z.object({
pageId: z
.string()
.describe('The id (or slugId) of a page within this share.'),
@@ -143,7 +142,7 @@ export class PublicShareChatToolsService {
'List the pages (titles + ids) that make up THIS published ' +
'documentation share, so you can orient yourself before reading or ' +
'searching. Only pages inside this share are listed.',
inputSchema: modelFriendlyInput({}),
inputSchema: z.object({}),
execute: async () => {
// Reuse the same share-tree logic the public /shares/tree route uses:
// it validates the share + workspace, excludes restricted subtrees,

View File

@@ -3,22 +3,6 @@ import { PageService } from './page.service';
import { MovePageDto } from '../dto/move-page.dto';
import { Page } from '@docmost/db/types/entity.types';
// A permissive chainable Proxy stands in for the locked Kysely trx so the
// FOR-UPDATE lock query chains inside executeTx(this.db, ...) resolve. Shared by
// every spec that drives a transactional write (movePage cycle guard, movePage
// provenance, movePageToSpace).
const makeChain = () => {
const c: any = new Proxy(function () {}, {
get: (_t, p) =>
p === 'then'
? undefined
: p === 'execute' || p === 'executeTakeFirst'
? () => Promise.resolve([])
: () => c,
});
return c;
};
// Direct instantiation with stub deps. The Test.createTestingModule form failed
// to resolve the @InjectKysely()/@InjectQueue() tokens at compile(), and this
// smoke test only needs the service to construct.
@@ -55,14 +39,12 @@ describe('PageService', () => {
// Build a PageService whose pageRepo (findById/updatePage) and own
// getPageBreadCrumbs are mockable, while every other collaborator stays a
// bare stub. We only need to drive the three cycle-guard branches, so we
// mock minimally rather than standing up the whole DI graph. The trx stub
// comes from the shared module-level `makeChain` helper.
// mock minimally rather than standing up the whole DI graph.
const makeService = (overrides?: {
breadcrumbs?: Array<{ id: string }>;
}) => {
const pageRepo = {
// Destination parent lookup: a valid, non-deleted, same-space page. Also
// serves the FOR-UPDATE lock reads inside the transaction.
// Destination parent lookup: a valid, non-deleted, same-space page.
findById: jest.fn().mockResolvedValue({
id: 'dest-parent',
deletedAt: null,
@@ -75,19 +57,11 @@ describe('PageService', () => {
const eventEmitter = { emit: jest.fn() };
// Re-parenting under a concrete parent now runs through
// executeTx(this.db, ...), which calls db.transaction().execute(fn). The
// trxStub is the value handed to the callback (the locked transaction).
const trxStub = makeChain();
const db = {
transaction: jest.fn(() => ({ execute: (fn: any) => fn(trxStub) })),
};
const svc = new PageService(
pageRepo as any, // pageRepo
{} as any, // pagePermissionRepo
{} as any, // attachmentRepo
db as any, // db
{} as any, // db
{} as any, // storageService
{} as any, // attachmentQueue
{} as any, // aiQueue
@@ -105,7 +79,7 @@ describe('PageService', () => {
.spyOn(svc, 'getPageBreadCrumbs')
.mockResolvedValue((overrides?.breadcrumbs ?? []) as any);
return { svc, pageRepo, eventEmitter, trxStub, db };
return { svc, pageRepo, eventEmitter };
};
// movePage takes `movedPage` as a param. Keep its parentPageId distinct from
@@ -172,65 +146,6 @@ describe('PageService', () => {
await expect(svc.movePage(dto, makeMovedPage())).resolves.not.toThrow();
expect(pageRepo.updatePage).toHaveBeenCalledTimes(1);
});
it('serializes a legitimate re-parent under FOR UPDATE in canonical lock order (#159 #9)', async () => {
// Destination's ancestor chain does NOT contain the moved page -> no cycle.
const { svc, pageRepo, trxStub } = makeService({
breadcrumbs: [{ id: 'dest-parent' }, { id: 'root' }],
});
const getBreadcrumbsSpy = jest.spyOn(svc, 'getPageBreadCrumbs');
const dto: MovePageDto = {
pageId: 'page-1',
position: VALID_POSITION,
parentPageId: 'dest-parent',
};
await expect(svc.movePage(dto, makeMovedPage())).resolves.not.toThrow();
// Both rows are locked FOR UPDATE inside the transaction: findById is
// called with { withLock: true, trx: <the locked tx> } for the moved page
// and the destination parent.
const lockCalls = pageRepo.findById.mock.calls.filter(
(c: any[]) => c[1]?.withLock === true,
);
expect(lockCalls).toHaveLength(2);
for (const call of lockCalls) {
expect(call[1].withLock).toBe(true);
expect(call[1].trx).toBe(trxStub);
}
// Locks are acquired in a canonical (id-sorted) order so the two opposing
// moves serialize without deadlocking.
const lockedIds = lockCalls.map((c: any[]) => c[0]);
expect(lockedIds).toEqual(['page-1', 'dest-parent'].sort());
// The cycle re-check runs inside the locked transaction (trx passed).
expect(getBreadcrumbsSpy).toHaveBeenCalledWith('dest-parent', trxStub);
// The update is written inside the same transaction (trx is the 3rd arg).
expect(pageRepo.updatePage).toHaveBeenCalledTimes(1);
expect(pageRepo.updatePage.mock.calls[0][2]).toBe(trxStub);
});
it('moves a page to root WITHOUT a transaction (no cycle possible)', async () => {
// A move-to-root (parentPageId === null) can never create a cycle, so it
// takes the unlocked else-branch: updatePage runs with NO trx and the
// db.transaction() serialization path is skipped entirely.
const { svc, pageRepo, db } = makeService();
const dto: MovePageDto = {
pageId: 'page-1',
position: VALID_POSITION,
parentPageId: null,
};
await expect(svc.movePage(dto, makeMovedPage())).resolves.not.toThrow();
// No FOR-UPDATE serialization: the transaction was never opened.
expect(db.transaction).not.toHaveBeenCalled();
// The update is written outside any transaction (3rd arg is undefined).
expect(pageRepo.updatePage).toHaveBeenCalledTimes(1);
expect(pageRepo.updatePage.mock.calls[0][2]).toBeUndefined();
});
});
describe('agent provenance stamping (#143)', () => {
@@ -344,8 +259,6 @@ describe('PageService', () => {
describe('movePage() → updatePage', () => {
const VALID_POSITION = 'a0';
// Re-parenting under a concrete parent runs through executeTx(this.db, ...);
// the shared `makeChain` helper stands in for the locked Kysely trx.
const run = async (provenance: any) => {
const pageRepo = {
findById: jest.fn().mockResolvedValue({
@@ -355,12 +268,9 @@ describe('PageService', () => {
}),
updatePage: jest.fn().mockResolvedValue({ numUpdatedRows: 1n }),
};
const trxStub = makeChain();
const svc = makeSvc({
pageRepo,
db: {
transaction: () => ({ execute: (fn: any) => fn(trxStub) }),
} as any,
db: {} as any,
});
// Legitimate move: destination ancestors do NOT include the moved page.
jest
@@ -406,9 +316,20 @@ describe('PageService', () => {
describe('movePageToSpace() → root-page updatePage', () => {
// movePageToSpace runs its writes inside executeTx(this.db, cb), which
// calls this.db.transaction().execute(fn => fn(trx)). The shared
// `makeChain` helper stands in for the Kysely trx so arbitrary chains
// resolve.
// calls this.db.transaction().execute(fn => fn(trx)). A permissive
// chainable Proxy stands in for the Kysely trx so arbitrary chains resolve.
const makeChain = () => {
const c: any = new Proxy(function () {}, {
get: (_t, p) =>
p === 'then'
? undefined
: p === 'execute' || p === 'executeTakeFirst'
? () => Promise.resolve([])
: () => c,
});
return c;
};
const run = async (provenance: any) => {
const trxStub = makeChain();
const db = {

View File

@@ -15,13 +15,13 @@ import {
executeWithCursorPagination,
} from '@docmost/db/pagination/cursor-pagination';
import { InjectKysely } from 'nestjs-kysely';
import { KyselyDB, KyselyTransaction } from '@docmost/db/types/kysely.types';
import { KyselyDB } from '@docmost/db/types/kysely.types';
import { generateJitteredKeyBetween } from 'fractional-indexing-jittered';
import { MovePageDto } from '../dto/move-page.dto';
import { shapeSidebarPagesTree } from './sidebar-pages-tree.util';
import { generateSlugId } from '../../../common/helpers';
import { getPageTitle } from '../../../common/helpers';
import { executeTx, dbOrTx } from '@docmost/db/utils';
import { executeTx } from '@docmost/db/utils';
import { AttachmentRepo } from '@docmost/db/repos/attachment/attachment.repo';
import { v7 as uuid7 } from 'uuid';
import {
@@ -915,53 +915,34 @@ export class PageService {
}
}
// Cheap self-move guard (no DB) — keep before the transaction.
if (dto.parentPageId && dto.parentPageId === dto.pageId) {
throw new BadRequestException('Cannot move a page into its own subtree');
// Server-side cycle guard: a page may not be moved into itself or into any
// page within its own subtree. Without this, an MCP/REST/agent caller (or a
// fast drag racing the client check) could persist a cycle and broadcast it.
// Only relevant when re-parenting under a concrete parent; moving to root
// (parentPageId null/undefined) can never create a cycle.
if (dto.parentPageId) {
if (dto.parentPageId === dto.pageId) {
throw new BadRequestException('Cannot move a page into its own subtree');
}
// Walk the destination parent's ancestor chain (reusing the breadcrumb
// ancestor CTE). If the page being moved appears among those ancestors,
// the destination lives inside the moved page's subtree -> cycle.
const destAncestors = await this.getPageBreadCrumbs(dto.parentPageId);
if (destAncestors.some((ancestor) => ancestor.id === dto.pageId)) {
throw new BadRequestException('Cannot move a page into its own subtree');
}
}
const updateValues = {
position: dto.position,
parentPageId: parentPageId,
// Agent-edit provenance: annotate the source on an agent move. A normal
// user request leaves the existing source value unchanged.
...agentSourceFields(provenance, 'lastUpdatedSource', 'lastUpdatedAiChatId'),
};
let updateResult;
if (typeof parentPageId === 'string') {
// Genuine re-parent under a concrete parent: the ONLY path that can create
// a cycle. Two opposing moves (A: X under Y, B: Y under X) racing each
// other could each pass a cycle check built from a stale snapshot and
// persist a cycle (#159 finding #9). Serialize them: lock the moved page
// and the destination parent FOR UPDATE in a canonical (id-sorted) order
// so they cannot deadlock, then run the cycle check INSIDE the transaction
// against the now-committed state.
updateResult = await executeTx(this.db, async (trx) => {
// Both opposing moves touch the same two rows {pageId, parentPageId};
// a fixed lock order forces one to wait for the other to commit. Lock by
// canonical UUIDs — `dto.pageId` can be a slugId (MovePageDto.pageId is a
// bare @IsString), so two opposing moves passing slugIds could sort into
// different lock orders and deadlock (AB-BA). `movedPage.id` is the
// resolved row UUID, matching `parentPageId`.
const lockIds = [movedPage.id, parentPageId].sort();
for (const id of lockIds) {
await this.pageRepo.findById(id, { withLock: true, trx });
}
// Re-read the destination's ancestor chain within the locked tx: it now
// reflects any concurrent re-parent that committed before we got the lock.
const destAncestors = await this.getPageBreadCrumbs(parentPageId, trx);
if (destAncestors.some((ancestor) => ancestor.id === dto.pageId)) {
throw new BadRequestException(
'Cannot move a page into its own subtree',
);
}
return this.pageRepo.updatePage(updateValues, dto.pageId, trx);
});
} else {
// Same-parent reorder or move-to-root: no cycle possible, no lock needed.
updateResult = await this.pageRepo.updatePage(updateValues, dto.pageId);
}
const updateResult = await this.pageRepo.updatePage(
{
position: dto.position,
parentPageId: parentPageId,
// Agent-edit provenance: annotate the source on an agent move. A normal
// user request leaves the existing source value unchanged.
...agentSourceFields(provenance, 'lastUpdatedSource', 'lastUpdatedAiChatId'),
},
dto.pageId,
);
// Guard against a phantom broadcast: if the row was concurrently deleted or
// otherwise not updated, skip the PAGE_MOVED event so we don't replay a move
@@ -1000,8 +981,8 @@ export class PageService {
});
}
async getPageBreadCrumbs(childPageId: string, trx?: KyselyTransaction) {
const ancestors = await dbOrTx(this.db, trx)
async getPageBreadCrumbs(childPageId: string) {
const ancestors = await this.db
.withRecursive('page_ancestors', (db) =>
db
.selectFrom('pages')

View File

@@ -21,6 +21,7 @@ export const AI_PROVIDER_SETTINGS_ALLOWED: readonly string[] = [
'driver',
'chatModel',
'chatApiStyle',
'chatContextWindow',
'embeddingModel',
'baseUrl',
'embeddingBaseUrl',
@@ -255,11 +256,17 @@ export class WorkspaceRepo {
): Promise<Workspace> {
const db = dbOrTx(this.db, trx);
// Assemble the provider object IN SQL. Keys are fixed provider field names
// (sql.lit -> inlined literals, no injection); values are bound params cast
// to ::text — postgres.js sends bound params untyped, and jsonb_build_object's
// value args are polymorphic ("any"), so without the explicit ::text cast
// Postgres throws "could not determine data type of parameter $1". The result
// is a real jsonb object, never a double-encoded string. The CASE self-heals
// (sql.lit -> inlined literals, no injection); values are bound params with
// an explicit cast — postgres.js sends bound params untyped, and
// jsonb_build_object's value args are polymorphic ("any"), so without the
// cast Postgres throws "could not determine data type of parameter $1". The
// cast is branched by the JS runtime type so the value lands in jsonb with
// the matching JSON type: a number stays a JSON number (e.g.
// chatContextWindow → `{"chatContextWindow":200000}`, jsonb_typeof 'number'),
// a boolean a JSON boolean, everything else a JSON string. A plain `::text`
// for all would store a numeric field as the JSON STRING `"200000"`, which
// the client's `typeof === "number"` guards reject. The result is a real
// jsonb object, never a double-encoded string. The CASE self-heals
// workspaces whose settings.ai.provider was previously corrupted into an
// array/string.
const entries = Object.entries(provider).filter(
@@ -267,7 +274,14 @@ export class WorkspaceRepo {
);
const patch = entries.length
? sql`jsonb_build_object(${sql.join(
entries.flatMap(([k, v]) => [sql.lit(k), sql`${v}::text`]),
entries.flatMap(([k, v]) => [
sql.lit(k),
typeof v === 'number'
? sql`${v}::numeric`
: typeof v === 'boolean'
? sql`${v}::boolean`
: sql`${v}::text`,
]),
)})`
: sql`'{}'::jsonb`;
return db

View File

@@ -41,3 +41,35 @@ describe('UpdateAiSettingsDto.chatApiStyle', () => {
expect(errs.find((e) => e.property === 'chatApiStyle')).toBeUndefined();
});
});
/** DTO validation for chatContextWindow (@IsOptional @IsInt @Min(0)). */
describe('UpdateAiSettingsDto.chatContextWindow', () => {
const errorsFor = async (chatContextWindow: unknown) =>
validate(plainToInstance(UpdateAiSettingsDto, { chatContextWindow }));
it('accepts a non-negative integer (incl. 0 = clear the limit)', async () => {
for (const v of [0, 200000]) {
const errs = await errorsFor(v);
expect(
errs.find((e) => e.property === 'chatContextWindow'),
).toBeUndefined();
}
});
it('rejects a negative value', async () => {
const errs = await errorsFor(-1);
expect(errs.find((e) => e.property === 'chatContextWindow')).toBeDefined();
});
it('rejects a non-integer value', async () => {
const errs = await errorsFor(1.5);
expect(errs.find((e) => e.property === 'chatContextWindow')).toBeDefined();
});
it('accepts the field being omitted (optional)', async () => {
const errs = await validate(plainToInstance(UpdateAiSettingsDto, {}));
expect(
errs.find((e) => e.property === 'chatContextWindow'),
).toBeUndefined();
});
});

View File

@@ -27,6 +27,8 @@ export interface UpdateAiSettingsInput {
driver?: AiDriver;
chatModel?: string;
chatApiStyle?: ChatApiStyle;
// Chat context-window size (tokens); 0/empty clears the limit.
chatContextWindow?: number;
embeddingModel?: string;
baseUrl?: string;
embeddingBaseUrl?: string;
@@ -162,6 +164,8 @@ export class AiSettingsService {
chatModel: provider.chatModel,
// Plain passthrough; getChatModel defaults unset to 'openai-compatible'.
chatApiStyle: provider.chatApiStyle,
// Admin-configured context-window size; 0/unset = no limit (badge denominator).
chatContextWindow: provider.chatContextWindow,
// Cheap model id for the anonymous public-share assistant; reuses the chat
// driver/baseUrl/apiKey. Empty/unset → callers fall back to chatModel.
publicShareChatModel: provider.publicShareChatModel,
@@ -244,6 +248,7 @@ export class AiSettingsService {
driver: provider.driver,
chatModel: provider.chatModel,
chatApiStyle: provider.chatApiStyle,
chatContextWindow: provider.chatContextWindow,
embeddingModel: provider.embeddingModel,
baseUrl: provider.baseUrl,
embeddingBaseUrl: provider.embeddingBaseUrl,

View File

@@ -35,6 +35,13 @@ export interface AiProviderSettings {
// Chat provider implementation for the `openai` driver. Unset → defaults to
// 'openai-compatible' (so reasoning is surfaced by default). See ChatApiStyle.
chatApiStyle?: ChatApiStyle;
// Admin-configured chat model context-window size, in tokens. There is no
// provider-independent way to discover this (OpenAI's /v1/models usually omits
// it, Gemini/Ollama/OpenRouter each expose it differently), so it is entered
// manually. Surfaced to the chat client (via assistant message metadata) as the
// denominator of the header "current / max" context badge. Empty/0 = no limit
// known → the badge shows only the current context size.
chatContextWindow?: number;
embeddingModel?: string;
baseUrl?: string;
// Embedding-specific base URL. Falls back to `baseUrl` when empty/unset.
@@ -73,6 +80,7 @@ export const PROVIDER_SETTINGS_KEYS = [
'driver',
'chatModel',
'chatApiStyle',
'chatContextWindow',
'embeddingModel',
'baseUrl',
'embeddingBaseUrl',
@@ -98,6 +106,10 @@ export const PROVIDER_SETTINGS_KEYS = [
export interface ResolvedAiConfig extends Partial<AiProviderSettings> {
driver?: AiDriver;
chatModel?: string;
// Admin-configured chat context-window size (tokens); 0/unset = no limit. Used
// as the header context-badge denominator. Re-declared for parity with the
// explicit fields above.
chatContextWindow?: number;
// Cheap model id for the public-share assistant; reuses the chat creds.
publicShareChatModel?: string;
// Agent-role id whose persona the public-share assistant adopts (empty/unset
@@ -117,6 +129,8 @@ export interface MaskedAiSettings {
driver?: AiDriver;
chatModel?: string;
chatApiStyle?: ChatApiStyle;
// Admin-configured chat context-window size (tokens); 0/unset = no limit.
chatContextWindow?: number;
embeddingModel?: string;
baseUrl?: string;
embeddingBaseUrl?: string;

View File

@@ -1,4 +1,4 @@
import { IsIn, IsOptional, IsString } from 'class-validator';
import { IsIn, IsInt, IsOptional, IsString, Min } from 'class-validator';
import {
AI_DRIVERS,
AiDriver,
@@ -29,6 +29,13 @@ export class UpdateAiSettingsDto {
@IsIn(CHAT_API_STYLES)
chatApiStyle?: ChatApiStyle;
// Chat model context-window size in tokens (header context-badge denominator).
// 0 (or empty) clears the limit so the badge shows only the current context.
@IsOptional()
@IsInt()
@Min(0)
chatContextWindow?: number;
@IsOptional()
@IsString()
embeddingModel?: string;

View File

@@ -1,122 +0,0 @@
import { Kysely } from 'kysely';
import { BadRequestException } from '@nestjs/common';
import { PageRepo } from '@docmost/db/repos/page/page.repo';
import { PageService } from '../../src/core/page/services/page.service';
import {
getTestDb,
destroyTestDb,
createWorkspace,
createSpace,
createPage,
} from './db';
/**
* #159 finding #9 — concurrent opposing page moves must not create a cycle.
*
* User A drags X under Y while user B simultaneously drags Y under X. Before the
* fix, the cycle guard (a breadcrumb/ancestor read) and the parent-update write
* were NOT in a transaction, so both moves ran against a stale snapshot, both
* passed their cycle check, and both committed -> X.parent=Y AND Y.parent=X: a
* cycle with no path to root, which breaks the recursive ancestor CTEs and makes
* both subtrees vanish from the sidebar.
*
* The fix wraps the cycle check + update in one READ COMMITTED transaction and
* locks the two involved rows FOR UPDATE in a canonical (id-sorted) order, then
* re-runs the cycle check inside the lock. One move wins; the loser sees the
* committed re-parent and trips the "own subtree" guard.
*
* NOTE: this runs against real Postgres in CI (the integration suite). There is
* no local Postgres in dev, so it is not expected to run locally.
*/
describe('movePage concurrent opposing moves do not create a cycle [integration]', () => {
let db: Kysely<any>;
let workspaceId: string;
let spaceId: string;
beforeAll(async () => {
db = getTestDb();
workspaceId = (await createWorkspace(db)).id;
spaceId = (await createSpace(db, workspaceId)).id;
});
afterAll(async () => {
await destroyTestDb();
});
it('lets exactly one opposing move win and never persists a cycle', async () => {
// Real collaborators against the test DB. movePage only touches db-backed
// methods on pageRepo plus the db itself and the event emitter (stubbed).
const pageRepo = new PageRepo(
db as any,
{} as any,
{ emit: () => undefined } as any,
);
const svc = new PageService(
pageRepo as any, // pageRepo
{} as any, // pagePermissionRepo
{} as any, // attachmentRepo
db as any, // db
{} as any, // storageService
{} as any, // attachmentQueue
{} as any, // aiQueue
{} as any, // generalQueue
{ emit: () => undefined } as any, // eventEmitter
{} as any, // collaborationGateway
{} as any, // watcherService
{} as any, // transclusionService
);
// Seed R (root), X and Y as children of R.
const root = await createPage(db, { workspaceId, spaceId, title: 'R' });
const pageX = await createPage(db, { workspaceId, spaceId, title: 'X' });
const pageY = await createPage(db, { workspaceId, spaceId, title: 'Y' });
// createPage does not set parentPageId; wire X.parent=R and Y.parent=R and
// give each a valid fractional-index position.
await db
.updateTable('pages')
.set({ parentPageId: root.id, position: 'a0' })
.where('id', 'in', [pageX.id, pageY.id])
.execute();
// The Page snapshots movePage receives (parentPageId must be R so each move
// is a genuine re-parent rather than a same-parent reorder).
const movedX = await pageRepo.findById(pageX.id);
const movedY = await pageRepo.findById(pageY.id);
// Two opposing moves racing: A moves X under Y, B moves Y under X.
const results = await Promise.allSettled([
svc.movePage(
{ pageId: pageX.id, parentPageId: pageY.id, position: 'a1' },
movedX,
),
svc.movePage(
{ pageId: pageY.id, parentPageId: pageX.id, position: 'a1' },
movedY,
),
]);
// Exactly one fulfilled and one rejected; the rejection is the cycle guard.
const fulfilled = results.filter((r) => r.status === 'fulfilled');
const rejected = results.filter((r) => r.status === 'rejected');
expect(fulfilled).toHaveLength(1);
expect(rejected).toHaveLength(1);
const reason = (rejected[0] as PromiseRejectedResult).reason;
expect(reason).toBeInstanceOf(BadRequestException);
expect(String(reason?.message)).toContain('own subtree');
// No cycle persisted: re-fetch X and Y and assert NOT (X->Y AND Y->X).
const xNow = await pageRepo.findById(pageX.id);
const yNow = await pageRepo.findById(pageY.id);
const isCycle =
xNow.parentPageId === pageY.id && yNow.parentPageId === pageX.id;
expect(isCycle).toBe(false);
// Exactly one re-parent took effect: one of X/Y still points at R, the other
// points at its sibling.
const xWon = xNow.parentPageId === pageY.id && yNow.parentPageId === root.id;
const yWon = yNow.parentPageId === pageX.id && xNow.parentPageId === root.id;
expect(xWon || yWon).toBe(true);
});
});

View File

@@ -0,0 +1,91 @@
import { Kysely, sql } from 'kysely';
import { WorkspaceRepo } from '@docmost/db/repos/workspace/workspace.repo';
import { getTestDb, destroyTestDb, createWorkspace } from './db';
/**
* WorkspaceRepo.updateAiProviderSettings numeric round-trip (#189, #213).
*
* `chatContextWindow` is the first NUMERIC provider field routed through this
* generic SQL layer. The patch builder must cast a JS number so it lands in
* jsonb as a JSON NUMBER, not the JSON STRING `"200000"` — the client guards
* (`typeof === "number"`) reject a string, silently killing the `/ max` badge
* denominator. A plain `::text` cast (the prior code) regressed exactly this.
* These specs are real SQL and assert both the JS value type and the on-disk
* `jsonb_typeof`.
*/
describe('WorkspaceRepo.updateAiProviderSettings (numeric round-trip) [integration]', () => {
let db: Kysely<any>;
let repo: WorkspaceRepo;
beforeAll(() => {
db = getTestDb();
repo = new WorkspaceRepo(db as any);
});
afterAll(async () => {
await destroyTestDb();
});
it('stores chatContextWindow as a JSON number (not a "200000" string)', async () => {
const ws = await createWorkspace(db, { settings: undefined });
const updated = await repo.updateAiProviderSettings(ws.id, {
driver: 'openai',
chatModel: 'gpt-4o',
chatContextWindow: 200000,
});
// Returned row: the number survives as a real JS number, alongside the
// string fields which stay strings.
const provider = (updated.settings as any)?.ai?.provider;
expect(provider.chatContextWindow).toBe(200000);
expect(typeof provider.chatContextWindow).toBe('number');
expect(provider.driver).toBe('openai');
expect(provider.chatModel).toBe('gpt-4o');
// On disk: the jsonb value is typed 'number' (the must-fix assertion), and
// sibling string fields are typed 'string'.
const typed = await db
.selectFrom('workspaces')
.select([
sql<string>`jsonb_typeof(settings->'ai'->'provider'->'chatContextWindow')`.as(
'windowType',
),
sql<string>`jsonb_typeof(settings->'ai'->'provider'->'chatModel')`.as(
'modelType',
),
])
.where('id', '=', ws.id)
.executeTakeFirstOrThrow();
expect(typed.windowType).toBe('number');
expect(typed.modelType).toBe('string');
});
it('re-reads chatContextWindow as a number after a partial-merge update', async () => {
const ws = await createWorkspace(db, {
settings: { ai: { provider: { driver: 'openai', chatModel: 'x' } } },
});
// Merge in only the numeric field; siblings must be preserved and the value
// must still be a JSON number, not a string.
await repo.updateAiProviderSettings(ws.id, { chatContextWindow: 128000 });
const row = await db
.selectFrom('workspaces')
.select([
'settings',
sql<string>`jsonb_typeof(settings->'ai'->'provider'->'chatContextWindow')`.as(
'windowType',
),
])
.where('id', '=', ws.id)
.executeTakeFirstOrThrow();
expect(row.windowType).toBe('number');
const provider = (row.settings as any)?.ai?.provider;
expect(provider.chatContextWindow).toBe(128000);
expect(provider.driver).toBe('openai');
expect(provider.chatModel).toBe('x');
});
});