Compare commits
5 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 8b36294b7b | |||
| 0fa2f9fb91 | |||
| e89ac627dd | |||
| f665f6fdd2 | |||
| 68899a2c2e |
@@ -209,6 +209,20 @@ MCP_DOCMOST_PASSWORD=
|
||||
# active" behavior.
|
||||
# AI_CHAT_DEFERRED_TOOLS=true
|
||||
|
||||
# --- Autonomous / detached agent runs (settings.ai.autonomousRuns) ---
|
||||
# Opt-in per workspace (AI settings; off by default). When on, a chat turn becomes
|
||||
# a server-side RUN that survives a browser disconnect — only an explicit Stop ends
|
||||
# it, and a client reconnects/live-follows the run.
|
||||
#
|
||||
# DEPLOY CONSTRAINT — SINGLE-INSTANCE ONLY in phase 1: Stop and the in-process
|
||||
# AbortController that backs it are process-local, so a Stop only aborts a run
|
||||
# executing on the SAME replica that owns it (cross-instance pub/sub stop is phase
|
||||
# 2 and not yet reliable). Do NOT enable autonomousRuns on a horizontally-scaled
|
||||
# deployment (multiple replicas behind a load balancer, or Docmost cloud
|
||||
# CLOUD=true) — run a single instance instead. The server logs a startup WARNING
|
||||
# when it detects a multi-instance deployment (CLOUD=true) so the constraint is
|
||||
# visible, and a startup sweep settles any run left dangling by a restart.
|
||||
|
||||
# --- Anonymous public-share AI assistant ---
|
||||
# Opt-in per workspace (AI settings -> "public share assistant"; off by default).
|
||||
# When enabled, anonymous visitors of a published share can ask an AI about that
|
||||
|
||||
@@ -279,6 +279,7 @@ The API server is a Fastify app with a global `/api` prefix (`main.ts` excludes
|
||||
- `core/ai-chat/tools/` — the agent's ~40 read+write tools. Every tool runs under the **calling user's** CASL permissions via a per-user loopback access token (`docmost-client.loader.ts`), so the agent can never exceed what the user could do. Only **reversible** operations are exposed (page history + trash; no permanent delete). Agent edits get an "AI agent" provenance badge in page history (`20260616T130000-agent-provenance` migration).
|
||||
- `core/ai-chat/embedding/` — RAG indexer + a BullMQ consumer on `AI_QUEUE` that embeds pages into `page_embeddings` (vector search), complementing Postgres full-text search. Pages are (re)indexed on edit; `AI_EMBEDDING_TIMEOUT_MS` bounds a hung embeddings endpoint.
|
||||
- `core/ai-chat/external-mcp/` — admins can attach external MCP servers (e.g. Tavily) to give the agent web access. **`ssrf-guard.ts` validates outbound MCP URLs against SSRF** — keep that guard in the path when touching external-MCP connection logic.
|
||||
- `core/ai-chat/ai-chat-run.service.ts` + `ai_chat_runs` — **detached/autonomous agent runs** (`#184`), behind the per-workspace `settings.ai.autonomousRuns` flag (off by default). When on, a turn becomes a server-side RUN that survives a browser disconnect; only an explicit `POST /ai-chat/stop` ends it, and a client reconnects/live-follows via `POST /ai-chat/run`. **DEPLOY CONSTRAINT — single-instance only in phase 1:** Stop and the AbortController that backs it are process-local, so a Stop only aborts a run executing on the **same** replica that owns it (cross-instance pub/sub stop is phase 2). Do **not** enable `autonomousRuns` on a horizontally-scaled deployment (multiple replicas behind a load balancer, or Docmost cloud `CLOUD=true`) — run a single instance instead. The server logs a startup WARNING when it detects a multi-instance deployment (`CLOUD=true`) so the constraint is visible. The startup sweep settles any run left dangling by a restart.
|
||||
|
||||
### Client structure
|
||||
Vite SPA. Code is organized by feature under `apps/client/src/features/*` (mirrors the server domains: `page`, `space`, `comment`, `ai-chat`, `editor`, …). Conventions:
|
||||
|
||||
@@ -72,6 +72,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
append/prepend fragments, nor to COMMENT bodies — a comment may legitimately
|
||||
contain a standalone footnote definition, which canonicalization would drop.
|
||||
(#228)
|
||||
- **Detached, autonomous agent runs that survive a browser disconnect.** When the
|
||||
new `settings.ai.autonomousRuns` workspace flag is on (off by default), an
|
||||
AI-chat turn becomes a first-class, server-side RUN tracked in a new
|
||||
`ai_chat_runs` table instead of a socket-bound stream: closing the tab or
|
||||
losing the connection no longer aborts the turn — it keeps executing and
|
||||
persisting server-side, and only an explicit Stop ends it. A client can
|
||||
reconnect and live-follow (or stop) an in-flight run via `POST /ai-chat/run`
|
||||
(resolve the latest run + its assistant message for a chat) and
|
||||
`POST /ai-chat/stop` (stop by `runId` or `chatId`). A partial unique index
|
||||
enforces one active run per chat, and a startup sweep settles any run left
|
||||
dangling by a restart. Phase 1 is single-instance-only (cross-instance Stop is
|
||||
not yet reliable); the server warns at startup on a horizontally-scaled
|
||||
deployment. (#184)
|
||||
- **Out-of-band page transfer via an in-RAM blob sandbox (`stash_page`).** A
|
||||
new MCP tool serializes a whole page (its full ProseMirror JSON, with every
|
||||
internal image/file mirrored) into an ephemeral in-RAM blob and returns only
|
||||
|
||||
@@ -19,7 +19,7 @@ import {
|
||||
IconPlus,
|
||||
IconX,
|
||||
} from "@tabler/icons-react";
|
||||
import { useAtom, useSetAtom } from "jotai";
|
||||
import { useAtom, useAtomValue, useSetAtom } from "jotai";
|
||||
import { useLocation, useMatch } from "react-router-dom";
|
||||
import { useTranslation } from "react-i18next";
|
||||
import { useQueryClient } from "@tanstack/react-query";
|
||||
@@ -41,13 +41,24 @@ import { extractPageSlugId } from "@/lib";
|
||||
import {
|
||||
AI_CHATS_RQ_KEY,
|
||||
AI_CHAT_MESSAGES_RQ_KEY,
|
||||
AI_CHAT_RUN_RQ_KEY,
|
||||
useAiChatMessagesQuery,
|
||||
useAiChatRunQuery,
|
||||
useAiChatsQuery,
|
||||
useAiRolesQuery,
|
||||
} from "@/features/ai-chat/queries/ai-chat-query.ts";
|
||||
import {
|
||||
shouldClearLatchOnQueryError,
|
||||
shouldClearStoppingLatch,
|
||||
shouldObserveRun,
|
||||
} from "@/features/ai-chat/utils/run-polling.ts";
|
||||
import { workspaceAtom } from "@/features/user/atoms/current-user-atom";
|
||||
import ConversationList from "@/features/ai-chat/components/conversation-list.tsx";
|
||||
import ChatThread from "@/features/ai-chat/components/chat-thread.tsx";
|
||||
import { exportAiChat } from "@/features/ai-chat/services/ai-chat-service.ts";
|
||||
import {
|
||||
exportAiChat,
|
||||
stopRun,
|
||||
} from "@/features/ai-chat/services/ai-chat-service.ts";
|
||||
import { useChatSession } from "@/features/ai-chat/hooks/use-chat-session.ts";
|
||||
import {
|
||||
shouldCollapseOnOutsidePointer,
|
||||
@@ -234,6 +245,147 @@ export default function AiChatWindow() {
|
||||
const { data: messageRows, isLoading: messagesLoading } =
|
||||
useAiChatMessagesQuery(activeChatId ?? undefined);
|
||||
|
||||
// #184 reconnect-and-live-follow. Whether detached agent runs are enabled for
|
||||
// this workspace. The reconnect endpoint itself is NOT flag-gated server-side
|
||||
// (it is only owner-gated and returns `{ run: null }` when the chat has no
|
||||
// run); but when the feature is off no runs are ever created, so polling it
|
||||
// would always come back empty — we gate it off here to avoid pointless polls.
|
||||
const workspace = useAtomValue(workspaceAtom);
|
||||
const autonomousRunsEnabled =
|
||||
workspace?.settings?.ai?.autonomousRuns === true;
|
||||
|
||||
// Whether THIS tab is the one actively streaming the open chat's run locally
|
||||
// (it started the run here and holds the SSE). Reported up from ChatThread. We
|
||||
// are the STREAMER while true and a passive OBSERVER while false — the basis of
|
||||
// the observer-vs-streamer detection. Reset to false by the fresh ChatThread's
|
||||
// mount effect on every chat switch.
|
||||
const [localStreaming, setLocalStreaming] = useState(false);
|
||||
const onStreamingChange = useCallback((streaming: boolean) => {
|
||||
setLocalStreaming(streaming);
|
||||
}, []);
|
||||
|
||||
// #184 Stop wiring. While a detached run is being stopped we SUPPRESS the
|
||||
// observer merge so the stopping run's still-persisting output does not
|
||||
// re-stream back into view between the moment the user pressed Stop and the run
|
||||
// actually settling as 'aborted' server-side. Polling itself keeps running (so
|
||||
// the terminal transition is still detected) — only the visual merge is gated.
|
||||
// Cleared when the run is observed terminal (below) or the chat is switched.
|
||||
const [stoppingRun, setStoppingRun] = useState(false);
|
||||
// Reset the stopping latch whenever the open chat changes: it is scoped to the
|
||||
// run of the previously-open chat.
|
||||
useEffect(() => {
|
||||
setStoppingRun(false);
|
||||
}, [activeChatId]);
|
||||
|
||||
// Authoritative stop of the open chat's detached run (the Stop button in
|
||||
// autonomous mode). Latch "stopping" first (suppresses the re-stream flash),
|
||||
// then request the server stop — the ONLY thing that ends a detached run; a mere
|
||||
// local SSE abort is a client disconnect the server ignores. On failure we
|
||||
// release the latch so the observer resumes (better to show the live run than to
|
||||
// freeze the view) and surface the error.
|
||||
const handleServerStop = useCallback(
|
||||
(chatId: string): void => {
|
||||
setStoppingRun(true);
|
||||
// #234 F4: drop the PREVIOUS turn's run from the cache so `run` becomes null
|
||||
// until the CURRENT turn's run is fetched fresh. Without this, once the local
|
||||
// stream aborts (localStreaming -> false) the run query re-enables and
|
||||
// react-query SYNCHRONOUSLY returns the still-cached prior terminal run; the
|
||||
// terminal effect would then clear the stopping latch against that STALE run
|
||||
// before the current turn's (still-running, detached, growing) run is ever
|
||||
// observed — re-opening the observer merge and flashing the growing output
|
||||
// over the frozen row. With the cache cleared the terminal effect's
|
||||
// `if (!run) return` holds the latch until the current run itself is observed
|
||||
// terminal (see shouldClearStoppingLatch).
|
||||
queryClient.removeQueries({ queryKey: AI_CHAT_RUN_RQ_KEY(chatId) });
|
||||
void stopRun(chatId).catch(() => {
|
||||
setStoppingRun(false);
|
||||
notifications.show({
|
||||
message: t("Failed to stop the run"),
|
||||
color: "red",
|
||||
});
|
||||
});
|
||||
},
|
||||
[t, queryClient],
|
||||
);
|
||||
|
||||
// Poll the latest run of the open chat ONLY when we are a passive observer:
|
||||
// feature on, a chat is open, and we are NOT the local streamer (the streamer
|
||||
// already has the live SSE — polling/merging too would double-render). The
|
||||
// query's own status-keyed refetchInterval stops once the run is terminal.
|
||||
const { data: runData, isError: runQueryFailed } = useAiChatRunQuery(
|
||||
activeChatId ?? undefined,
|
||||
autonomousRunsEnabled && !localStreaming,
|
||||
);
|
||||
const run = runData?.run ?? null;
|
||||
|
||||
// Safety net (#234 F4 review): after handleServerStop clears the run cache,
|
||||
// `run` is null until the current turn's run is fetched fresh, and the terminal
|
||||
// effect below holds the latch via `if (!run) return`. If that refetch instead
|
||||
// ERRORS PERMANENTLY (the GET-run keeps failing) while we are no longer the
|
||||
// streamer, the run stays null, its status-keyed refetchInterval is off, and
|
||||
// nothing would ever observe a terminal run — freezing the view with the
|
||||
// observer merge suppressed. Release the latch on that error so the live view
|
||||
// resumes rather than stays stuck (the local stopRun may already have succeeded
|
||||
// independently).
|
||||
//
|
||||
// #234 F7: this must NOT fire on a TRANSIENT error while `run` is still an
|
||||
// ACTIVE held run. In TanStack Query v5 (retry:false) the query's `data` is
|
||||
// RETAINED on error, so `runQueryFailed` can be true while `run` is still
|
||||
// pending/running — releasing then would re-open the observer merge and flash
|
||||
// the growing detached run over the frozen row (the very flash F4 prevents). The
|
||||
// decision is the pure, unit-tested `shouldClearLatchOnQueryError`, which gates
|
||||
// on the run NOT being active: it cures only the genuine permanent-null-freeze
|
||||
// (`run === null`) and never releases against an active run.
|
||||
useEffect(() => {
|
||||
if (
|
||||
shouldClearLatchOnQueryError({
|
||||
stoppingRun,
|
||||
isLocalStreaming: localStreaming,
|
||||
runQueryFailed,
|
||||
run,
|
||||
})
|
||||
)
|
||||
setStoppingRun(false);
|
||||
}, [stoppingRun, localStreaming, runQueryFailed, run]);
|
||||
// The run's incrementally-persisted assistant message to merge into the thread,
|
||||
// but only while we are an observer (never when we are the streamer — guards
|
||||
// against a stale poll fighting the live stream). Includes a terminal run so the
|
||||
// final persisted output is shown on reopen.
|
||||
const observedRow =
|
||||
shouldObserveRun(run, localStreaming) && !stoppingRun
|
||||
? (runData?.message ?? null)
|
||||
: null;
|
||||
|
||||
// When the observed run reaches a terminal status, do a final messages refetch
|
||||
// so the persisted final state (token/context badge, export source) is shown,
|
||||
// then the query's refetchInterval has already stopped polling. Deduped per run
|
||||
// id so it fires exactly once per run, not on every subsequent poll-less render.
|
||||
const finalizedRunIdRef = useRef<string | null>(null);
|
||||
useEffect(() => {
|
||||
if (!run || !activeChatId) return;
|
||||
if (run.status === "pending" || run.status === "running") {
|
||||
// Active again (a new run) — re-arm so its terminal transition fires once.
|
||||
finalizedRunIdRef.current = null;
|
||||
return;
|
||||
}
|
||||
// Terminal: a stop we requested has landed (or the run finished on its own),
|
||||
// so release the stopping latch — the observer merge can now show the final
|
||||
// persisted (aborted/finished) output without any live re-stream. The decision
|
||||
// is the pure, unit-tested `shouldClearStoppingLatch` (run-polling.ts): release
|
||||
// ONLY when we requested a stop, this tab is no longer the streamer, AND the
|
||||
// CURRENT run is terminal. The #234 F4 cache removal in handleServerStop makes
|
||||
// `run` null (this branch's `if (!run) return` above holds) until the current
|
||||
// turn's run is fetched fresh, so the latch can never clear against a stale
|
||||
// cached run.
|
||||
if (shouldClearStoppingLatch({ stoppingRun, run, isLocalStreaming: localStreaming }))
|
||||
setStoppingRun(false);
|
||||
if (finalizedRunIdRef.current === run.id) return;
|
||||
finalizedRunIdRef.current = run.id;
|
||||
queryClient.invalidateQueries({
|
||||
queryKey: AI_CHAT_MESSAGES_RQ_KEY(activeChatId),
|
||||
});
|
||||
}, [run, activeChatId, queryClient, stoppingRun, localStreaming]);
|
||||
|
||||
// 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"
|
||||
@@ -882,6 +1034,18 @@ export default function AiChatWindow() {
|
||||
assistantName={currentRole?.name}
|
||||
onTurnFinished={onTurnFinished}
|
||||
onServerChatId={onServerChatId}
|
||||
// #184: live-follow a still-running run when we reopened the chat as
|
||||
// a passive observer; null when there is nothing to observe or this
|
||||
// tab is the streamer. onStreamingChange lets the window stop polling
|
||||
// while we are the streamer.
|
||||
observedRow={observedRow}
|
||||
onStreamingChange={onStreamingChange}
|
||||
// #184: in autonomous mode the Stop button must hit the authoritative
|
||||
// server stop (a local SSE abort is a client disconnect the server
|
||||
// ignores). onServerStop also arms the "stopping" latch above so the
|
||||
// stopped run's output does not re-stream via the observer merge.
|
||||
autonomousRunsEnabled={autonomousRunsEnabled}
|
||||
onServerStop={handleServerStop}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
|
||||
@@ -11,6 +11,7 @@ const h = vi.hoisted(() => ({
|
||||
onFinish: null as null | ((arg: Record<string, unknown>) => void),
|
||||
sendMessage: vi.fn(),
|
||||
stop: vi.fn(),
|
||||
setMessages: vi.fn(),
|
||||
transport: null as null | {
|
||||
prepareSendMessagesRequest: (arg: {
|
||||
messages: unknown[];
|
||||
@@ -30,6 +31,8 @@ vi.mock("@ai-sdk/react", () => ({
|
||||
status: h.state.status,
|
||||
stop: h.state.stop,
|
||||
error: null,
|
||||
// #184: ChatThread reads setMessages to merge a polled observer run.
|
||||
setMessages: h.state.setMessages,
|
||||
};
|
||||
},
|
||||
}));
|
||||
@@ -228,3 +231,56 @@ describe("ChatThread — turn-end decision (onFinish)", () => {
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
// #184 passive-observer merge: when reconnecting to a still-running run, the
|
||||
// parent feeds the polled run message via `observedRow`; ChatThread merges it via
|
||||
// setMessages — but ONLY when this tab is NOT itself streaming (the streamer's
|
||||
// SSE owns the view, so a stale observedRow must never overwrite it).
|
||||
describe("ChatThread — observer run merge (#184)", () => {
|
||||
beforeEach(() => {
|
||||
h.state.onFinish = null;
|
||||
h.state.setMessages.mockReset();
|
||||
});
|
||||
|
||||
const observedRow = {
|
||||
id: "a-run",
|
||||
role: "assistant",
|
||||
content: "step 1\nstep 2",
|
||||
metadata: {
|
||||
parts: [{ type: "text", text: "step 1\nstep 2" }],
|
||||
},
|
||||
createdAt: "2026-01-01T00:00:00Z",
|
||||
} as const;
|
||||
|
||||
function renderObserver(status: string) {
|
||||
h.state.status = status;
|
||||
render(
|
||||
<MantineProvider>
|
||||
<ChatThread
|
||||
chatId="c1"
|
||||
initialRows={[]}
|
||||
onTurnFinished={vi.fn()}
|
||||
observedRow={observedRow as never}
|
||||
/>
|
||||
</MantineProvider>,
|
||||
);
|
||||
}
|
||||
|
||||
it("merges the polled run message when this tab is a passive observer", () => {
|
||||
renderObserver("ready");
|
||||
expect(h.state.setMessages).toHaveBeenCalledTimes(1);
|
||||
// The updater replaces/append the observed assistant row by id.
|
||||
const updater = h.state.setMessages.mock.calls[0][0] as (
|
||||
prev: { id: string; parts: { text: string }[] }[],
|
||||
) => { id: string; parts: { text: string }[] }[];
|
||||
const merged = updater([{ id: "u1", parts: [{ text: "hi" }] }]);
|
||||
expect(merged).toHaveLength(2);
|
||||
expect(merged[1].id).toBe("a-run");
|
||||
expect(merged[1].parts[0].text).toBe("step 1\nstep 2");
|
||||
});
|
||||
|
||||
it("does NOT merge while THIS tab is the streamer (no double-render)", () => {
|
||||
renderObserver("streaming");
|
||||
expect(h.state.setMessages).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -24,6 +24,7 @@ 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 { mergeObservedMessage } from "@/features/ai-chat/utils/run-polling.ts";
|
||||
import {
|
||||
dequeue,
|
||||
enqueueMessage,
|
||||
@@ -86,6 +87,29 @@ interface ChatThreadProps {
|
||||
* Copy/export button available mid-stream). Distinct from onTurnFinished,
|
||||
* which fires only at the terminal outcome. */
|
||||
onServerChatId?: (serverChatId?: string) => void;
|
||||
/** #184 reconnect-and-live-follow. When THIS tab reopened a chat whose agent
|
||||
* run is still going (it is a PASSIVE OBSERVER — it did not start the run here),
|
||||
* the parent polls the reconnect endpoint and feeds the run's incrementally-
|
||||
* persisted assistant message here; we merge it into the live list so new
|
||||
* steps/tool-calls appear as they are persisted. Null when there is nothing to
|
||||
* observe (no run, feature off, or this tab IS the streamer). The merge is
|
||||
* ADDITIONALLY guarded by our own `isStreaming`, so a stale value can never
|
||||
* fight the local stream when we are the streamer. */
|
||||
observedRow?: IAiChatMessageRow | null;
|
||||
/** Report this tab's live streaming status up to the parent, so it can stop
|
||||
* polling the run while WE are the active streamer (the SSE owns the view) and
|
||||
* resume once we go idle. Called from an effect on every transition. */
|
||||
onStreamingChange?: (streaming: boolean) => void;
|
||||
/** #184: whether detached/autonomous agent runs are enabled for this workspace.
|
||||
* When true the Stop button must additionally hit the AUTHORITATIVE server stop
|
||||
* (via onServerStop) — aborting only the local SSE is just a client disconnect,
|
||||
* which the server deliberately ignores, so the detached run would keep going. */
|
||||
autonomousRunsEnabled?: boolean;
|
||||
/** #184: request the server-side stop of this chat's active run (the parent owns
|
||||
* the endpoint call + the "stopping" latch that keeps observer-polling from
|
||||
* immediately re-streaming the stopping run's output). Called with the resolved
|
||||
* chat id when the user presses Stop in autonomous mode. */
|
||||
onServerStop?: (chatId: string) => void;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -131,6 +155,10 @@ export default function ChatThread({
|
||||
assistantName,
|
||||
onTurnFinished,
|
||||
onServerChatId,
|
||||
observedRow,
|
||||
onStreamingChange,
|
||||
autonomousRunsEnabled,
|
||||
onServerStop,
|
||||
}: ChatThreadProps) {
|
||||
const { t } = useTranslation();
|
||||
|
||||
@@ -216,6 +244,16 @@ export default function ChatThread({
|
||||
const flushOnAbortRef = useRef(false);
|
||||
const interruptNextSendRef = useRef(false);
|
||||
|
||||
// #234 F5: the user pressed Stop while streaming a BRAND-NEW chat whose server
|
||||
// chat id has not been adopted yet (the `start` chunk carrying it hadn't landed
|
||||
// when Stop was pressed). A local SSE abort alone does NOT stop the DETACHED
|
||||
// autonomous run — it keeps burning tokens and WRITING TO PAGES — so we cannot
|
||||
// just no-op. We latch the stop as PENDING and fire the authoritative server
|
||||
// stop the moment onServerChatId adopts the id (below). Read-and-cleared there;
|
||||
// also defused on every new turn start so it can never fire against a later,
|
||||
// unrelated turn's run.
|
||||
const stopPendingRef = useRef(false);
|
||||
|
||||
// FIFO dequeue + send the next queued message (no-op when the queue is empty).
|
||||
// Returns whether a message was actually sent, so callers can tell an empty
|
||||
// dequeue (nothing to flush) from a real send.
|
||||
@@ -274,7 +312,7 @@ export default function ChatThread({
|
||||
[],
|
||||
);
|
||||
|
||||
const { messages, sendMessage, status, stop, error } = useChat({
|
||||
const { messages, sendMessage, status, stop, error, setMessages } = useChat({
|
||||
// Stable per-mount key. Existing chats use their real id; new chats use a
|
||||
// generated client id (never `undefined`) so the store is NOT re-created on
|
||||
// every render mid-stream (see `chatStoreId` above).
|
||||
@@ -365,7 +403,14 @@ export default function ChatThread({
|
||||
return;
|
||||
lastForwardedChatIdRef.current = serverChatId;
|
||||
onServerChatId(serverChatId);
|
||||
}, [messages, onServerChatId]);
|
||||
// #234 F5: if Stop was pressed before the id was known, the authoritative
|
||||
// server stop was deferred to this adoption point — fire it now with the
|
||||
// just-adopted id. One-shot (read-and-clear) so it can't fire twice.
|
||||
if (stopPendingRef.current) {
|
||||
stopPendingRef.current = false;
|
||||
onServerStop?.(serverChatId);
|
||||
}
|
||||
}, [messages, onServerChatId, onServerStop]);
|
||||
|
||||
// Live "turn was interrupted" marker for the CURRENT session. The red error
|
||||
// banner (driven by `error`) covers the error case; this covers an aborted
|
||||
@@ -378,6 +423,27 @@ export default function ChatThread({
|
||||
|
||||
const isStreaming = status === "submitted" || status === "streaming";
|
||||
|
||||
// #184: report our live streaming status up so the parent stops polling the run
|
||||
// while WE are the streamer (the SSE owns the view) and resumes once we go idle.
|
||||
// Effect (not render) so it never updates parent state during our own render;
|
||||
// fires on mount with `false`, which also re-syncs the parent after a chat
|
||||
// switch remounts this thread (a fresh mount is idle until the user sends).
|
||||
useEffect(() => {
|
||||
onStreamingChange?.(isStreaming);
|
||||
}, [isStreaming, onStreamingChange]);
|
||||
|
||||
// #184 passive-observer merge: when the parent feeds a polled run message (we
|
||||
// reopened a chat whose run is still going and did NOT start it here), merge it
|
||||
// into the live list so new steps/tool-calls appear as they are persisted. Hard-
|
||||
// gated by `!isStreaming`: if THIS tab is actually the streamer, the local SSE
|
||||
// owns the view and a stale observedRow must never overwrite it. `observedRow`
|
||||
// is a stable per-poll object, so this runs once per poll, not per render.
|
||||
useEffect(() => {
|
||||
if (isStreaming || !observedRow) return;
|
||||
const observed = rowToUiMessage(observedRow);
|
||||
setMessages((prev) => mergeObservedMessage(prev, observed));
|
||||
}, [observedRow, isStreaming, setMessages]);
|
||||
|
||||
// "Send now" on a queued message: interrupt the current turn and immediately
|
||||
// send THIS message, keeping the agent's partial output. Other queued messages
|
||||
// stay queued and flush normally after the new turn. Reuses the existing
|
||||
@@ -409,6 +475,40 @@ export default function ChatThread({
|
||||
[setQueue, stop],
|
||||
);
|
||||
|
||||
// Stop the current turn. ALWAYS abort the local SSE (`stop()`) so the composer
|
||||
// returns to idle immediately. In AUTONOMOUS mode the turn is a DETACHED run:
|
||||
// aborting the local SSE is only a client disconnect, which the server ignores,
|
||||
// so the run would keep executing — we ADDITIONALLY request the authoritative
|
||||
// server-side stop (the parent owns that call + the "stopping" latch that keeps
|
||||
// observer-polling from re-streaming the stopping run's output). The chat id is
|
||||
// read live from chatIdRef (adopted early at the stream's `start` chunk); if it
|
||||
// is not known yet — a brand-new chat in the first moment of its first turn —
|
||||
// only the local abort happens (there is no server-side run handle to stop yet).
|
||||
const handleStop = useCallback(() => {
|
||||
stop();
|
||||
if (!autonomousRunsEnabled) return;
|
||||
if (chatIdRef.current) {
|
||||
onServerStop?.(chatIdRef.current);
|
||||
} else {
|
||||
// #234 F5: no chat id yet (brand-new chat in the first moment of its first
|
||||
// turn, before the `start` chunk adopted the id). Latch the stop as pending;
|
||||
// the onServerChatId adoption effect fires the deferred server stop as soon
|
||||
// as the id appears, so the detached run is still authoritatively stopped
|
||||
// instead of left running by a silent local-only abort.
|
||||
//
|
||||
// KNOWN LIMITATION (#234 F5 review): `stop()` above has already aborted the
|
||||
// local SSE reader. In the rare sub-window where Stop is pressed while still
|
||||
// `submitted` (request sent, not one chunk read yet), that abort can cancel
|
||||
// the reader BEFORE the `start` chunk is applied to `messages`, so the
|
||||
// adoption effect never runs and this pending stop never fires. The detached
|
||||
// run then keeps going for that turn. This is not a regression (the pre-fix
|
||||
// behavior sent no server stop at all); closing it fully would require
|
||||
// deferring the local abort until adoption, which is riskier and out of scope
|
||||
// for this fix. Documented so a future change can address the abort-ordering.
|
||||
stopPendingRef.current = true;
|
||||
}
|
||||
}, [stop, autonomousRunsEnabled, onServerStop]);
|
||||
|
||||
// Clear the stopped marker as soon as a new turn begins streaming, and drop any
|
||||
// stale "Send now" interrupt flags. On the legit interrupt path both refs are
|
||||
// already consumed synchronously (onFinish + prepareSendMessagesRequest) before
|
||||
@@ -420,6 +520,11 @@ export default function ChatThread({
|
||||
setStopNotice(null);
|
||||
flushOnAbortRef.current = false;
|
||||
interruptNextSendRef.current = false;
|
||||
// #234 F5: a new turn is starting — drop any pending deferred-stop from a
|
||||
// previous turn that never adopted an id, so it can never fire against this
|
||||
// (or a later) unrelated turn's run. A deferred stop for the CURRENT turn is
|
||||
// set AFTER this effect (on the Stop click), so this does not clobber it.
|
||||
stopPendingRef.current = false;
|
||||
}
|
||||
}, [isStreaming]);
|
||||
|
||||
@@ -539,7 +644,7 @@ export default function ChatThread({
|
||||
<ChatInput
|
||||
onSend={(text) => sendMessage({ text })}
|
||||
onQueue={enqueue}
|
||||
onStop={stop}
|
||||
onStop={handleStop}
|
||||
isStreaming={isStreaming}
|
||||
/>
|
||||
</Stack>
|
||||
|
||||
@@ -12,6 +12,7 @@ import {
|
||||
deleteAiChat,
|
||||
deleteAiRole,
|
||||
getAiChatMessages,
|
||||
getAiChatRun,
|
||||
getAiChats,
|
||||
getAiRoleCatalog,
|
||||
getAiRoleCatalogBundle,
|
||||
@@ -24,6 +25,7 @@ import {
|
||||
import {
|
||||
IAiChat,
|
||||
IAiChatMessageRow,
|
||||
IAiChatRunResponse,
|
||||
IAiRole,
|
||||
IAiRoleCatalog,
|
||||
IAiRoleCatalogBundle,
|
||||
@@ -34,6 +36,7 @@ import {
|
||||
IAiRoleUpdateFromCatalogResult,
|
||||
} from "@/features/ai-chat/types/ai-chat.types.ts";
|
||||
import { IPagination } from "@/lib/types.ts";
|
||||
import { runPollInterval } from "@/features/ai-chat/utils/run-polling.ts";
|
||||
|
||||
export const AI_CHATS_RQ_KEY = ["ai-chats"];
|
||||
export const AI_ROLES_RQ_KEY = ["ai-roles"];
|
||||
@@ -51,16 +54,18 @@ export const AI_CHAT_MESSAGES_RQ_KEY = (chatId: string) => [
|
||||
"ai-chat-messages",
|
||||
chatId,
|
||||
];
|
||||
export const AI_CHAT_RUN_RQ_KEY = (chatId: string) => ["ai-chat-run", chatId];
|
||||
|
||||
/** Paginated list of the current user's chats (auto-loads further pages). */
|
||||
export function useAiChatsQuery() {
|
||||
const query = useInfiniteQuery({
|
||||
queryKey: AI_CHATS_RQ_KEY,
|
||||
queryFn: ({ pageParam }) =>
|
||||
getAiChats({ cursor: pageParam, limit: 50 }),
|
||||
queryFn: ({ pageParam }) => getAiChats({ cursor: pageParam, limit: 50 }),
|
||||
initialPageParam: undefined as string | undefined,
|
||||
getNextPageParam: (lastPage) =>
|
||||
lastPage.meta.hasNextPage ? (lastPage.meta.nextCursor ?? undefined) : undefined,
|
||||
lastPage.meta.hasNextPage
|
||||
? (lastPage.meta.nextCursor ?? undefined)
|
||||
: undefined,
|
||||
});
|
||||
|
||||
const data = useMemo<IPagination<IAiChat> | undefined>(() => {
|
||||
@@ -90,7 +95,9 @@ export function useAiChatMessagesQuery(chatId: string | undefined) {
|
||||
getAiChatMessages({ chatId: chatId as string, cursor: pageParam }),
|
||||
initialPageParam: undefined as string | undefined,
|
||||
getNextPageParam: (lastPage) =>
|
||||
lastPage.meta.hasNextPage ? (lastPage.meta.nextCursor ?? undefined) : undefined,
|
||||
lastPage.meta.hasNextPage
|
||||
? (lastPage.meta.nextCursor ?? undefined)
|
||||
: undefined,
|
||||
enabled: !!chatId,
|
||||
});
|
||||
|
||||
@@ -131,6 +138,34 @@ export function useAiChatMessagesQuery(chatId: string | undefined) {
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Reconnect to a chat's latest agent run and LIVE-FOLLOW it (#184). While the run
|
||||
* is active the query re-polls every {@link runPollInterval} ms (driven off the
|
||||
* fetched `run.status`, the same status-keyed refetchInterval pattern as the
|
||||
* embeddings reindex polling); once the run reaches a terminal status — or there
|
||||
* is no run — the interval returns `false` and polling stops on its own. Polling
|
||||
* is thus naturally bounded by the run terminating; no separate timeout cap.
|
||||
*
|
||||
* `enabled` gates the whole thing: callers pass `false` when the autonomous-runs
|
||||
* feature is off (the endpoint is NOT flag-gated server-side, but with the feature
|
||||
* off the chat has no runs, so polling would only ever return `{ run: null }`) OR
|
||||
* when THIS tab is the one actively streaming the run (the live SSE owns the view,
|
||||
* so we must not also poll/merge). The global `retry: false` means a failed fetch
|
||||
* leaves `data` undefined, so refetchInterval(undefined run) returns false — a
|
||||
* failed fetch can never spin a tight loop.
|
||||
*/
|
||||
export function useAiChatRunQuery(
|
||||
chatId: string | undefined,
|
||||
enabled: boolean,
|
||||
) {
|
||||
return useQuery<IAiChatRunResponse, Error>({
|
||||
queryKey: AI_CHAT_RUN_RQ_KEY(chatId ?? ""),
|
||||
queryFn: () => getAiChatRun(chatId as string),
|
||||
enabled: !!chatId && enabled,
|
||||
refetchInterval: (query) => runPollInterval(query.state.data?.run),
|
||||
});
|
||||
}
|
||||
|
||||
export function useRenameAiChatMutation() {
|
||||
const queryClient = useQueryClient();
|
||||
const { t } = useTranslation();
|
||||
@@ -280,11 +315,14 @@ export function useImportAiRolesFromCatalogMutation() {
|
||||
mutationFn: (payload) => importAiRolesFromCatalog(payload),
|
||||
onSuccess: (result) => {
|
||||
notifications.show({
|
||||
message: t("Imported {{created}}, renamed {{renamed}}, skipped {{skipped}}", {
|
||||
created: result.created,
|
||||
renamed: result.renamed,
|
||||
skipped: result.skipped,
|
||||
}),
|
||||
message: t(
|
||||
"Imported {{created}}, renamed {{renamed}}, skipped {{skipped}}",
|
||||
{
|
||||
created: result.created,
|
||||
renamed: result.renamed,
|
||||
skipped: result.skipped,
|
||||
},
|
||||
),
|
||||
});
|
||||
// Surface partial failures (e.g. unique-name races) as a red warning.
|
||||
if (result.errors.length > 0) {
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
import React from "react";
|
||||
import { renderHook, waitFor } from "@testing-library/react";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import type { IAiChatRunResponse } from "@/features/ai-chat/types/ai-chat.types.ts";
|
||||
|
||||
// react-i18next is pulled in transitively by ai-chat-query.ts (the mutation hooks
|
||||
// use it); stub it so the module imports cleanly in this hook test.
|
||||
vi.mock("react-i18next", () => ({
|
||||
useTranslation: () => ({ t: (key: string) => key }),
|
||||
}));
|
||||
|
||||
vi.mock("@mantine/notifications", () => ({
|
||||
notifications: { show: vi.fn() },
|
||||
}));
|
||||
|
||||
// Mock the whole service module; only getAiChatRun is exercised here, but the
|
||||
// other named exports must exist so ai-chat-query.ts imports resolve.
|
||||
vi.mock("@/features/ai-chat/services/ai-chat-service.ts", () => ({
|
||||
getAiChatRun: vi.fn(),
|
||||
getAiChatMessages: vi.fn(),
|
||||
getAiChats: vi.fn(),
|
||||
getAiRoleCatalog: vi.fn(),
|
||||
getAiRoleCatalogBundle: vi.fn(),
|
||||
getAiRoles: vi.fn(),
|
||||
importAiRolesFromCatalog: vi.fn(),
|
||||
createAiRole: vi.fn(),
|
||||
deleteAiChat: vi.fn(),
|
||||
deleteAiRole: vi.fn(),
|
||||
renameAiChat: vi.fn(),
|
||||
updateAiRole: vi.fn(),
|
||||
updateAiRoleFromCatalog: vi.fn(),
|
||||
}));
|
||||
|
||||
import { getAiChatRun } from "@/features/ai-chat/services/ai-chat-service.ts";
|
||||
import { useAiChatRunQuery } from "@/features/ai-chat/queries/ai-chat-query.ts";
|
||||
|
||||
function createWrapper() {
|
||||
const queryClient = new QueryClient({
|
||||
defaultOptions: { queries: { retry: false } },
|
||||
});
|
||||
return function Wrapper({ children }: { children: React.ReactNode }) {
|
||||
return (
|
||||
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
|
||||
);
|
||||
};
|
||||
}
|
||||
|
||||
const runningResponse: IAiChatRunResponse = {
|
||||
run: { id: "run-1", chatId: "c1", status: "running" },
|
||||
message: {
|
||||
id: "a1",
|
||||
role: "assistant",
|
||||
content: "working...",
|
||||
createdAt: "2026-01-01T00:00:00Z",
|
||||
},
|
||||
};
|
||||
|
||||
describe("useAiChatRunQuery — enable gating", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("fetches the run when enabled (passive observer, feature on)", async () => {
|
||||
vi.mocked(getAiChatRun).mockResolvedValue(runningResponse);
|
||||
const { result } = renderHook(() => useAiChatRunQuery("c1", true), {
|
||||
wrapper: createWrapper(),
|
||||
});
|
||||
await waitFor(() => expect(result.current.isSuccess).toBe(true));
|
||||
expect(getAiChatRun).toHaveBeenCalledWith("c1");
|
||||
expect(result.current.data?.run?.status).toBe("running");
|
||||
});
|
||||
|
||||
it("does NOT fetch when disabled (this tab is the streamer / feature off)", async () => {
|
||||
vi.mocked(getAiChatRun).mockResolvedValue(runningResponse);
|
||||
renderHook(() => useAiChatRunQuery("c1", false), {
|
||||
wrapper: createWrapper(),
|
||||
});
|
||||
// Give any errant fetch a chance to fire, then assert none did.
|
||||
await new Promise((r) => setTimeout(r, 20));
|
||||
expect(getAiChatRun).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("does NOT fetch when there is no chat id", async () => {
|
||||
vi.mocked(getAiChatRun).mockResolvedValue(runningResponse);
|
||||
renderHook(() => useAiChatRunQuery(undefined, true), {
|
||||
wrapper: createWrapper(),
|
||||
});
|
||||
await new Promise((r) => setTimeout(r, 20));
|
||||
expect(getAiChatRun).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -5,6 +5,7 @@ import {
|
||||
IAiChatListParams,
|
||||
IAiChatMessageRow,
|
||||
IAiChatMessagesParams,
|
||||
IAiChatRunResponse,
|
||||
IAiRole,
|
||||
IAiRoleCatalog,
|
||||
IAiRoleCatalogBundle,
|
||||
@@ -42,6 +43,38 @@ export async function getAiChatMessages(
|
||||
return req.data;
|
||||
}
|
||||
|
||||
/**
|
||||
* Reconnect to the latest agent run of a chat (#184). Returns the run's
|
||||
* persisted lifecycle state and the assistant message it materializes (the
|
||||
* partial output while the run is in-flight, the final output once it finished).
|
||||
* The DB is the source of truth, so this works for an in-flight run (the browser
|
||||
* dropped, the run kept going) and a finished one alike; `{ run: null }` when the
|
||||
* chat has never had a run. Owner-gated server-side (the requesting user must own
|
||||
* the chat); it is NOT flag-gated — when the feature is off the chat simply has no
|
||||
* runs, so the endpoint returns `{ run: null }`.
|
||||
*/
|
||||
export async function getAiChatRun(
|
||||
chatId: string,
|
||||
): Promise<IAiChatRunResponse> {
|
||||
const req = await api.post<IAiChatRunResponse>("/ai-chat/run", { chatId });
|
||||
return req.data;
|
||||
}
|
||||
|
||||
/**
|
||||
* Explicitly STOP the active agent run of a chat (#184). This is the ONLY thing
|
||||
* that ends a DETACHED run — a mere browser disconnect (aborting the local SSE)
|
||||
* is deliberately ignored server-side, so the client must call this to actually
|
||||
* stop an autonomous run. Targeted by `chatId` (the server resolves whatever run
|
||||
* is active on it); owner-gated server-side. Returns `{ stopped }` — false when
|
||||
* there was nothing active to stop.
|
||||
*/
|
||||
export async function stopRun(
|
||||
chatId: string,
|
||||
): Promise<{ stopped: boolean }> {
|
||||
const req = await api.post<{ stopped: boolean }>("/ai-chat/stop", { chatId });
|
||||
return req.data;
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve the chat bound to a document (the current user's most-recent chat
|
||||
* created on that page), or null when there is none. Drives auto-open-on-page.
|
||||
|
||||
@@ -200,6 +200,38 @@ export interface IAiChatMessageRow {
|
||||
createdAt: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* A persisted agent-run row (#184), mirroring the `ai_chat_runs` fields the
|
||||
* client reads from `POST /ai-chat/run`. Only `status` is load-bearing for the
|
||||
* reconnect-and-live-update UX (it drives the poll cadence); the rest are carried
|
||||
* for display/diagnostics. The DB is the source of truth, so this resolves for an
|
||||
* in-flight run (the browser dropped, the run kept going) and a finished one.
|
||||
*/
|
||||
export interface IAiChatRun {
|
||||
id: string;
|
||||
chatId: string;
|
||||
// 'pending' | 'running' | 'succeeded' | 'failed' | 'aborted'. The first two are
|
||||
// ACTIVE (keep polling); the rest are TERMINAL (stop polling).
|
||||
status: "pending" | "running" | "succeeded" | "failed" | "aborted" | string;
|
||||
error?: string | null;
|
||||
stepCount?: number;
|
||||
assistantMessageId?: string | null;
|
||||
startedAt?: string | null;
|
||||
finishedAt?: string | null;
|
||||
createdAt?: string;
|
||||
updatedAt?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Response of `POST /ai-chat/run` (#184): the latest run of a chat and the
|
||||
* assistant message it materializes (the partial/final output, projected from the
|
||||
* persisted rows). Both are `null` when the chat has never had a run.
|
||||
*/
|
||||
export interface IAiChatRunResponse {
|
||||
run: IAiChatRun | null;
|
||||
message: IAiChatMessageRow | null;
|
||||
}
|
||||
|
||||
export interface IAiChatListParams extends QueryParams {}
|
||||
|
||||
export interface IAiChatMessagesParams {
|
||||
|
||||
@@ -0,0 +1,303 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import type { UIMessage } from "@ai-sdk/react";
|
||||
import type { IAiChatRun } from "@/features/ai-chat/types/ai-chat.types.ts";
|
||||
import {
|
||||
RUN_POLL_INTERVAL_MS,
|
||||
isRunActive,
|
||||
runPollInterval,
|
||||
shouldObserveRun,
|
||||
shouldClearStoppingLatch,
|
||||
shouldClearLatchOnQueryError,
|
||||
mergeObservedMessage,
|
||||
} from "./run-polling.ts";
|
||||
|
||||
function makeRun(status: string): IAiChatRun {
|
||||
return { id: "run-1", chatId: "c1", status };
|
||||
}
|
||||
|
||||
function makeMsg(id: string, text: string): UIMessage {
|
||||
return {
|
||||
id,
|
||||
role: "assistant",
|
||||
parts: [{ type: "text", text }],
|
||||
} as UIMessage;
|
||||
}
|
||||
|
||||
describe("isRunActive", () => {
|
||||
it("treats pending and running as active", () => {
|
||||
expect(isRunActive(makeRun("pending"))).toBe(true);
|
||||
expect(isRunActive(makeRun("running"))).toBe(true);
|
||||
});
|
||||
|
||||
it("treats terminal / unknown / nullish as not active", () => {
|
||||
expect(isRunActive(makeRun("succeeded"))).toBe(false);
|
||||
expect(isRunActive(makeRun("failed"))).toBe(false);
|
||||
expect(isRunActive(makeRun("aborted"))).toBe(false);
|
||||
expect(isRunActive(makeRun("weird-future-status"))).toBe(false);
|
||||
expect(isRunActive(null)).toBe(false);
|
||||
expect(isRunActive(undefined)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("runPollInterval (the refetchInterval helper)", () => {
|
||||
it("returns 2000ms while the run is pending/running", () => {
|
||||
expect(runPollInterval(makeRun("pending"))).toBe(RUN_POLL_INTERVAL_MS);
|
||||
expect(runPollInterval(makeRun("running"))).toBe(RUN_POLL_INTERVAL_MS);
|
||||
expect(RUN_POLL_INTERVAL_MS).toBe(2000);
|
||||
});
|
||||
|
||||
it("returns false (stop polling) once the run is terminal", () => {
|
||||
expect(runPollInterval(makeRun("succeeded"))).toBe(false);
|
||||
expect(runPollInterval(makeRun("failed"))).toBe(false);
|
||||
expect(runPollInterval(makeRun("aborted"))).toBe(false);
|
||||
});
|
||||
|
||||
it("returns false (no polling) when there is no run", () => {
|
||||
expect(runPollInterval(null)).toBe(false);
|
||||
expect(runPollInterval(undefined)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("shouldObserveRun (observer-vs-streamer decision)", () => {
|
||||
it("observes an active run when this tab is NOT the local streamer", () => {
|
||||
expect(shouldObserveRun(makeRun("running"), false)).toBe(true);
|
||||
expect(shouldObserveRun(makeRun("pending"), false)).toBe(true);
|
||||
});
|
||||
|
||||
it("observes a terminal run too (so the final output shows on reopen)", () => {
|
||||
expect(shouldObserveRun(makeRun("succeeded"), false)).toBe(true);
|
||||
});
|
||||
|
||||
it("does NOT observe when this tab IS the streamer (no double-render)", () => {
|
||||
expect(shouldObserveRun(makeRun("running"), true)).toBe(false);
|
||||
expect(shouldObserveRun(makeRun("succeeded"), true)).toBe(false);
|
||||
});
|
||||
|
||||
it("does NOT observe when there is no run", () => {
|
||||
expect(shouldObserveRun(null, false)).toBe(false);
|
||||
expect(shouldObserveRun(undefined, false)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("shouldClearStoppingLatch (#234 latch-release decision)", () => {
|
||||
// The one case the latch SHOULD clear: we requested a stop, we are the passive
|
||||
// observer (not streaming), and the CURRENT run is terminal.
|
||||
it("clears only when stopping, observing, and the run is terminal", () => {
|
||||
expect(
|
||||
shouldClearStoppingLatch({
|
||||
stoppingRun: true,
|
||||
run: makeRun("aborted"),
|
||||
isLocalStreaming: false,
|
||||
}),
|
||||
).toBe(true);
|
||||
expect(
|
||||
shouldClearStoppingLatch({
|
||||
stoppingRun: true,
|
||||
run: makeRun("succeeded"),
|
||||
isLocalStreaming: false,
|
||||
}),
|
||||
).toBe(true);
|
||||
expect(
|
||||
shouldClearStoppingLatch({
|
||||
stoppingRun: true,
|
||||
run: makeRun("failed"),
|
||||
isLocalStreaming: false,
|
||||
}),
|
||||
).toBe(true);
|
||||
});
|
||||
|
||||
// Round-3 regression: clearing while THIS tab is still the local streamer would
|
||||
// re-open the flash for the current turn the moment we switch to observer role.
|
||||
// A predicate lacking the streaming gate would (wrongly) return true here.
|
||||
it("does NOT clear while this tab is the local streamer", () => {
|
||||
expect(
|
||||
shouldClearStoppingLatch({
|
||||
stoppingRun: true,
|
||||
run: makeRun("aborted"),
|
||||
isLocalStreaming: true,
|
||||
}),
|
||||
).toBe(false);
|
||||
expect(
|
||||
shouldClearStoppingLatch({
|
||||
stoppingRun: true,
|
||||
run: makeRun("succeeded"),
|
||||
isLocalStreaming: true,
|
||||
}),
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
// The detached run keeps growing after a local abort — while it is still
|
||||
// active the latch MUST hold so the observer merge stays suppressed.
|
||||
it("does NOT clear while the run is still active", () => {
|
||||
expect(
|
||||
shouldClearStoppingLatch({
|
||||
stoppingRun: true,
|
||||
run: makeRun("running"),
|
||||
isLocalStreaming: false,
|
||||
}),
|
||||
).toBe(false);
|
||||
expect(
|
||||
shouldClearStoppingLatch({
|
||||
stoppingRun: true,
|
||||
run: makeRun("pending"),
|
||||
isLocalStreaming: false,
|
||||
}),
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
// #234 F4: on Stop the stale PREVIOUS-turn run is removed from the cache, so the
|
||||
// observed `run` is null until the current turn's run is fetched fresh. A null
|
||||
// run HOLDS the latch — it can never clear against the just-removed stale run,
|
||||
// only against the current turn's own terminal run once observed.
|
||||
it("does NOT clear against a removed/absent run (F4 stale-run guard)", () => {
|
||||
expect(
|
||||
shouldClearStoppingLatch({
|
||||
stoppingRun: true,
|
||||
run: null,
|
||||
isLocalStreaming: false,
|
||||
}),
|
||||
).toBe(false);
|
||||
expect(
|
||||
shouldClearStoppingLatch({
|
||||
stoppingRun: true,
|
||||
run: undefined,
|
||||
isLocalStreaming: false,
|
||||
}),
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
it("does NOT clear when no stop was requested", () => {
|
||||
expect(
|
||||
shouldClearStoppingLatch({
|
||||
stoppingRun: false,
|
||||
run: makeRun("aborted"),
|
||||
isLocalStreaming: false,
|
||||
}),
|
||||
).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("shouldClearLatchOnQueryError (#234 F7 error-safety-net decision)", () => {
|
||||
// This guards the REAL anti-flash decision the component's run-query-error
|
||||
// safety-net effect uses (ai-chat-window.tsx wires the effect to THIS helper,
|
||||
// not a copy — so the test is non-vacuous vs the live code).
|
||||
|
||||
// (b) The F7 hole: a TRANSIENT run-query error while `run` is STILL ACTIVE must
|
||||
// NOT clear the latch. TanStack Query v5 retains `data` on error, so
|
||||
// runQueryFailed can be true while the held run is still pending/running.
|
||||
// Against the PRE-F7 condition (without `!isRunActive(run)`) this would return
|
||||
// true — so this assertion fails on the buggy code (non-vacuous).
|
||||
it("does NOT clear on a transient error while the run is still ACTIVE (F7)", () => {
|
||||
expect(
|
||||
shouldClearLatchOnQueryError({
|
||||
stoppingRun: true,
|
||||
isLocalStreaming: false,
|
||||
runQueryFailed: true,
|
||||
run: makeRun("running"),
|
||||
}),
|
||||
).toBe(false);
|
||||
expect(
|
||||
shouldClearLatchOnQueryError({
|
||||
stoppingRun: true,
|
||||
isLocalStreaming: false,
|
||||
runQueryFailed: true,
|
||||
run: makeRun("pending"),
|
||||
}),
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
// (a) The genuine permanent-null-freeze: run cache cleared by removeQueries +
|
||||
// the refetch keeps ERRORING, so `run === null`. This is the ONLY case the
|
||||
// safety-net exists to cure — it MUST clear so the frozen view resumes.
|
||||
it("clears on a permanent error when the run is null (permanent-null-freeze)", () => {
|
||||
expect(
|
||||
shouldClearLatchOnQueryError({
|
||||
stoppingRun: true,
|
||||
isLocalStreaming: false,
|
||||
runQueryFailed: true,
|
||||
run: null,
|
||||
}),
|
||||
).toBe(true);
|
||||
expect(
|
||||
shouldClearLatchOnQueryError({
|
||||
stoppingRun: true,
|
||||
isLocalStreaming: false,
|
||||
runQueryFailed: true,
|
||||
run: undefined,
|
||||
}),
|
||||
).toBe(true);
|
||||
});
|
||||
|
||||
// A TERMINAL run also satisfies `!isRunActive`; clearing then is harmless — the
|
||||
// terminal effect (shouldClearStoppingLatch) already clears for a terminal run,
|
||||
// so this only ever agrees with it. Asserted so the (c) reasoning is pinned.
|
||||
it("clears on an error when the run is terminal (harmless, agrees with terminal effect)", () => {
|
||||
expect(
|
||||
shouldClearLatchOnQueryError({
|
||||
stoppingRun: true,
|
||||
isLocalStreaming: false,
|
||||
runQueryFailed: true,
|
||||
run: makeRun("aborted"),
|
||||
}),
|
||||
).toBe(true);
|
||||
});
|
||||
|
||||
it("does NOT clear without an actual query error", () => {
|
||||
expect(
|
||||
shouldClearLatchOnQueryError({
|
||||
stoppingRun: true,
|
||||
isLocalStreaming: false,
|
||||
runQueryFailed: false,
|
||||
run: null,
|
||||
}),
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
it("does NOT clear while this tab is the local streamer", () => {
|
||||
expect(
|
||||
shouldClearLatchOnQueryError({
|
||||
stoppingRun: true,
|
||||
isLocalStreaming: true,
|
||||
runQueryFailed: true,
|
||||
run: null,
|
||||
}),
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
it("does NOT clear when no stop was requested", () => {
|
||||
expect(
|
||||
shouldClearLatchOnQueryError({
|
||||
stoppingRun: false,
|
||||
isLocalStreaming: false,
|
||||
runQueryFailed: true,
|
||||
run: null,
|
||||
}),
|
||||
).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("mergeObservedMessage", () => {
|
||||
it("replaces the message with the same id in place (per-step growth)", () => {
|
||||
const prev = [makeMsg("u1", "hi"), makeMsg("a1", "step 1")];
|
||||
const observed = makeMsg("a1", "step 1\nstep 2");
|
||||
const next = mergeObservedMessage(prev, observed);
|
||||
expect(next).toHaveLength(2);
|
||||
expect(next[1]).toBe(observed);
|
||||
expect(next[0]).toBe(prev[0]); // untouched
|
||||
expect(next).not.toBe(prev); // new array (never mutates input)
|
||||
});
|
||||
|
||||
it("appends when the observed message is not yet present", () => {
|
||||
const prev = [makeMsg("u1", "hi")];
|
||||
const observed = makeMsg("a1", "first token");
|
||||
const next = mergeObservedMessage(prev, observed);
|
||||
expect(next).toHaveLength(2);
|
||||
expect(next[1]).toBe(observed);
|
||||
});
|
||||
|
||||
it("returns the original list unchanged when there is nothing to merge", () => {
|
||||
const prev = [makeMsg("u1", "hi")];
|
||||
expect(mergeObservedMessage(prev, null)).toBe(prev);
|
||||
expect(mergeObservedMessage(prev, undefined)).toBe(prev);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,151 @@
|
||||
import type { UIMessage } from "@ai-sdk/react";
|
||||
import type { IAiChatRun } from "@/features/ai-chat/types/ai-chat.types.ts";
|
||||
|
||||
/**
|
||||
* Reconnect-and-live-follow helpers (#184). When a chat is reopened while its
|
||||
* agent run is STILL going, this tab is a PASSIVE OBSERVER: it did not start the
|
||||
* run here (no local SSE stream), so it catches up by POLLING the reconnect
|
||||
* endpoint (`POST /ai-chat/run`) and merging the run's incrementally-persisted
|
||||
* assistant message into the rendered thread. These are the small pure decisions
|
||||
* that machinery hangs off, extracted so they can be unit-tested in isolation
|
||||
* (mirrors how reindex polling / editor-sync-state are tested).
|
||||
*/
|
||||
|
||||
/** How often to re-poll the reconnect endpoint while a run is ACTIVE. */
|
||||
export const RUN_POLL_INTERVAL_MS = 2000;
|
||||
|
||||
// 'pending' and 'running' are the two ACTIVE statuses; 'succeeded' | 'failed' |
|
||||
// 'aborted' are TERMINAL (and any unknown future status is treated as terminal,
|
||||
// so a stale/odd value never polls forever).
|
||||
const ACTIVE_STATUSES = new Set(["pending", "running"]);
|
||||
|
||||
/** Whether a run is still going (worth polling / merging live updates from). */
|
||||
export function isRunActive(run: IAiChatRun | null | undefined): boolean {
|
||||
return !!run && ACTIVE_STATUSES.has(run.status);
|
||||
}
|
||||
|
||||
/**
|
||||
* The TanStack Query `refetchInterval` value for the run query: poll every
|
||||
* {@link RUN_POLL_INTERVAL_MS} while the run is active, and `false` (stop) once
|
||||
* it is terminal or there is no run. Polling is thus naturally bounded by the run
|
||||
* reaching a terminal status — no separate timeout cap is needed.
|
||||
*/
|
||||
export function runPollInterval(
|
||||
run: IAiChatRun | null | undefined,
|
||||
): number | false {
|
||||
return isRunActive(run) ? RUN_POLL_INTERVAL_MS : false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Observer-vs-streamer decision. We render the polled run message (catch up +
|
||||
* keep advancing) ONLY when this tab is a passive observer: there IS a run AND
|
||||
* this tab is NOT the one locally streaming it (we reconnected, we didn't start
|
||||
* it here). When this tab is the streamer, the live SSE stream owns the view, so
|
||||
* we neither poll nor merge — avoiding a double-render fight. Terminal runs still
|
||||
* merge (so the final persisted output is shown on reopen); the poll itself is
|
||||
* stopped separately by {@link runPollInterval}.
|
||||
*/
|
||||
export function shouldObserveRun(
|
||||
run: IAiChatRun | null | undefined,
|
||||
localStreaming: boolean,
|
||||
): boolean {
|
||||
return !!run && !localStreaming;
|
||||
}
|
||||
|
||||
/**
|
||||
* Should the "stopping" latch — which suppresses the observer re-stream flash
|
||||
* after the user pressed Stop — be RELEASED now? All three must hold:
|
||||
* - `stoppingRun`: we actually requested a stop (otherwise nothing to release);
|
||||
* - `!isLocalStreaming`: this tab is NOT the local streamer. While we are the
|
||||
* streamer the run query is disabled, so the observed `run` is not the run we
|
||||
* are following — releasing the latch then would re-open the flash for the
|
||||
* current turn the instant we switch to observer role;
|
||||
* - the observed `run` EXISTS and has reached a TERMINAL status.
|
||||
*
|
||||
* The null / still-active `run` case is the #234 F4 invariant. On Stop the stale
|
||||
* PREVIOUS-turn run is removed from the query cache (`removeQueries`), so `run`
|
||||
* is null until the CURRENT turn's run is re-fetched fresh; a null or active run
|
||||
* therefore HOLDS the latch, so it can only ever clear against the current turn's
|
||||
* OWN terminal run — never a stale cached one. (The cache removal itself is
|
||||
* integration-level in AiChatWindow; this predicate encodes the decision given
|
||||
* whatever run is currently observed, and a stale terminal run is
|
||||
* indistinguishable from a current terminal run at the predicate level — hence
|
||||
* the cache removal is what guarantees only the current run is ever passed here.)
|
||||
*/
|
||||
export function shouldClearStoppingLatch(args: {
|
||||
stoppingRun: boolean;
|
||||
run: IAiChatRun | null | undefined;
|
||||
isLocalStreaming: boolean;
|
||||
}): boolean {
|
||||
const { stoppingRun, run, isLocalStreaming } = args;
|
||||
if (!stoppingRun || isLocalStreaming) return false;
|
||||
return !!run && !isRunActive(run);
|
||||
}
|
||||
|
||||
/**
|
||||
* Should the "stopping" latch be RELEASED by the run-query ERROR safety-net?
|
||||
* (#234 F7 — a NEW path of the same re-stream flash the F4 latch exists to
|
||||
* prevent.) After Stop, `handleServerStop` clears the run cache; the terminal
|
||||
* effect then holds the latch via `if (!run) return` until the CURRENT turn's run
|
||||
* is fetched fresh. If that refetch instead ERRORS permanently, `run` stays null,
|
||||
* its status-keyed refetchInterval is off, and nothing would ever observe a
|
||||
* terminal run — freezing the view with the observer merge suppressed. This
|
||||
* safety-net cures ONLY that genuine permanent-null-freeze.
|
||||
*
|
||||
* All four must hold:
|
||||
* - `stoppingRun`: we actually requested a stop (otherwise nothing to release);
|
||||
* - `!isLocalStreaming`: this tab is NOT the local streamer (same reason as
|
||||
* {@link shouldClearStoppingLatch});
|
||||
* - `runQueryFailed`: the run query is in its error state (TanStack Query v5 with
|
||||
* retry:false — isError);
|
||||
* - `!isRunActive(run)`: the observed `run` is NOT an active (pending/running)
|
||||
* held run. This is the F7 gate. In TanStack Query v5 the query's `data` is
|
||||
* RETAINED on error, so `runQueryFailed` can be true while `run` is STILL an
|
||||
* ACTIVE run (a single transient GET-run failure in the window between Stop and
|
||||
* settle). Without this gate a transient error would release the latch early —
|
||||
* re-opening the observer merge and flashing the growing detached run over the
|
||||
* frozen row (exactly the F4 flash). Gating on the run NOT being active means we
|
||||
* only ever cure the permanent-null-freeze (`run === null`, so
|
||||
* `isRunActive(null)` is false), never release against an active run.
|
||||
*
|
||||
* (A terminal `run` also satisfies `!isRunActive(run)`; clearing then is harmless
|
||||
* — the terminal effect's {@link shouldClearStoppingLatch} already clears the
|
||||
* latch for a terminal run, so this only ever agrees with it, never conflicts.)
|
||||
*
|
||||
* INVARIANT (do not break): clearing the latch on the `run === null` branch is safe
|
||||
* ONLY because the run query's `refetchInterval` (see {@link runPollInterval}) stops
|
||||
* polling when the data is empty — so after we clear on null+error there is no
|
||||
* subsequent auto-poll that could return a still-active detached run and re-open the
|
||||
* merge. If `refetchInterval` is ever changed to keep polling on `run === null`/on
|
||||
* error, this null-branch clear would re-open the F7 flash through the null path.
|
||||
* Do not change the run query's refetchInterval without re-checking this path.
|
||||
*/
|
||||
export function shouldClearLatchOnQueryError(args: {
|
||||
stoppingRun: boolean;
|
||||
isLocalStreaming: boolean;
|
||||
runQueryFailed: boolean;
|
||||
run: IAiChatRun | null | undefined;
|
||||
}): boolean {
|
||||
const { stoppingRun, isLocalStreaming, runQueryFailed, run } = args;
|
||||
return (
|
||||
stoppingRun && !isLocalStreaming && runQueryFailed && !isRunActive(run)
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Merge an observed assistant message into the rendered list: replace the message
|
||||
* with the same id in place (the in-progress assistant row is already seeded from
|
||||
* history, so per-step growth replaces it), or append it when absent. Returns a
|
||||
* new array; the input is never mutated.
|
||||
*/
|
||||
export function mergeObservedMessage(
|
||||
messages: UIMessage[],
|
||||
observed: UIMessage | null | undefined,
|
||||
): UIMessage[] {
|
||||
if (!observed) return messages;
|
||||
const idx = messages.findIndex((m) => m.id === observed.id);
|
||||
if (idx === -1) return [...messages, observed];
|
||||
const next = messages.slice();
|
||||
next[idx] = observed;
|
||||
return next;
|
||||
}
|
||||
@@ -45,6 +45,7 @@ import {
|
||||
TiptapPdf,
|
||||
PageBreak,
|
||||
SearchAndReplace,
|
||||
MultiCursor,
|
||||
Mention,
|
||||
TableDndExtension,
|
||||
TableHandleCommandsExtension,
|
||||
@@ -447,6 +448,10 @@ export const mainExtensions = [
|
||||
};
|
||||
},
|
||||
}).configure(),
|
||||
// Multi-cursor editing (MVP / Variant A): select-all-occurrences + type into
|
||||
// all at once. Does not depend on collaboration, so it lives in mainExtensions
|
||||
// (available in both the plain and collaborative editors).
|
||||
MultiCursor,
|
||||
Columns,
|
||||
Column,
|
||||
AutoJoiner.configure({
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
@import "./core.css";
|
||||
@import "./collaboration.css";
|
||||
@import "./multi-cursor.css";
|
||||
@import "./task-list.css";
|
||||
@import "./placeholder.css";
|
||||
@import "./drag-handle.css";
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
/*
|
||||
* Multi-cursor (issue #196). Deliberately DISTINCT from the collaboration
|
||||
* carets (collaboration.css) so a user never confuses their own multi-cursors
|
||||
* with a co-author's caret: solid accent-blue carets + a translucent blue
|
||||
* range highlight, versus the thin dark collaboration caret with a name label.
|
||||
*/
|
||||
|
||||
/* A secondary caret rendered as a Decoration.widget at each cursor position. */
|
||||
.multi-cursor__caret {
|
||||
position: relative;
|
||||
display: inline-block;
|
||||
width: 0;
|
||||
height: 1em;
|
||||
vertical-align: text-bottom;
|
||||
pointer-events: none;
|
||||
}
|
||||
|
||||
.multi-cursor__caret::after {
|
||||
content: "";
|
||||
position: absolute;
|
||||
left: -1px;
|
||||
top: 0;
|
||||
bottom: 0;
|
||||
width: 2px;
|
||||
background: #2b6cb0;
|
||||
animation: multi-cursor-blink 1s steps(1) infinite;
|
||||
}
|
||||
|
||||
/* Optional label class reserved for future per-cursor annotations. */
|
||||
.multi-cursor__label {
|
||||
position: absolute;
|
||||
top: -1.4em;
|
||||
left: -1px;
|
||||
font-size: 0.7rem;
|
||||
line-height: normal;
|
||||
padding: 0.05rem 0.25rem;
|
||||
border-radius: 3px 3px 3px 0;
|
||||
background: #2b6cb0;
|
||||
color: #fff;
|
||||
white-space: nowrap;
|
||||
user-select: none;
|
||||
pointer-events: none;
|
||||
}
|
||||
|
||||
/* Inline highlight for a multi-cursor RANGE (from < to). */
|
||||
.multi-cursor__selection {
|
||||
background: rgba(43, 108, 176, 0.28);
|
||||
border-radius: 2px;
|
||||
}
|
||||
|
||||
@keyframes multi-cursor-blink {
|
||||
0%,
|
||||
50% {
|
||||
opacity: 1;
|
||||
}
|
||||
51%,
|
||||
100% {
|
||||
opacity: 0;
|
||||
}
|
||||
}
|
||||
+62
@@ -394,6 +394,10 @@ export default function AiProviderSettings() {
|
||||
useState<boolean>(
|
||||
workspace?.settings?.ai?.publicShareAssistant ?? false,
|
||||
);
|
||||
// #184: detached/autonomous agent runs (settings.ai.autonomousRuns).
|
||||
const [autonomousRunsEnabled, setAutonomousRunsEnabled] = useState<boolean>(
|
||||
workspace?.settings?.ai?.autonomousRuns ?? false,
|
||||
);
|
||||
const [chatToggleLoading, setChatToggleLoading] = useState(false);
|
||||
const [searchToggleLoading, setSearchToggleLoading] = useState(false);
|
||||
const [dictationToggleLoading, setDictationToggleLoading] = useState(false);
|
||||
@@ -403,6 +407,8 @@ export default function AiProviderSettings() {
|
||||
publicShareAssistantToggleLoading,
|
||||
setPublicShareAssistantToggleLoading,
|
||||
] = useState(false);
|
||||
const [autonomousRunsToggleLoading, setAutonomousRunsToggleLoading] =
|
||||
useState(false);
|
||||
|
||||
// Whether a key is currently stored server-side (drives the placeholder).
|
||||
const [hasApiKey, setHasApiKey] = useState(false);
|
||||
@@ -730,6 +736,37 @@ export default function AiProviderSettings() {
|
||||
}
|
||||
}
|
||||
|
||||
// Optimistic toggle for detached/autonomous agent runs
|
||||
// (settings.ai.autonomousRuns). When on, a chat turn becomes a server-side run
|
||||
// that survives a browser disconnect and can be reconnected to / live-followed;
|
||||
// only an explicit Stop ends it. Off by default; single-instance-only in phase 1.
|
||||
async function handleToggleAutonomousRuns(value: boolean) {
|
||||
setAutonomousRunsToggleLoading(true);
|
||||
const previous = autonomousRunsEnabled;
|
||||
setAutonomousRunsEnabled(value);
|
||||
try {
|
||||
const updated = await updateWorkspace({ autonomousRuns: value });
|
||||
setWorkspace({
|
||||
...updated,
|
||||
settings: {
|
||||
...updated.settings,
|
||||
ai: { ...updated.settings?.ai, autonomousRuns: value },
|
||||
},
|
||||
});
|
||||
notifications.show({ message: t("Updated successfully") });
|
||||
} catch (err) {
|
||||
setAutonomousRunsEnabled(previous);
|
||||
const message = (err as { response?: { data?: { message?: string } } })
|
||||
?.response?.data?.message;
|
||||
notifications.show({
|
||||
message: message ?? t("Failed to update data"),
|
||||
color: "red",
|
||||
});
|
||||
} finally {
|
||||
setAutonomousRunsToggleLoading(false);
|
||||
}
|
||||
}
|
||||
|
||||
// Admins only — match the previous behavior.
|
||||
if (!isAdmin) {
|
||||
return (
|
||||
@@ -960,6 +997,31 @@ export default function AiProviderSettings() {
|
||||
{...form.getInputProps("publicShareAssistantRoleId")}
|
||||
/>
|
||||
|
||||
{/* Detached/autonomous agent runs: a chat turn becomes a server-side run
|
||||
that survives a browser disconnect; only an explicit Stop ends it.
|
||||
Single-instance-only in phase 1. */}
|
||||
<Group justify="space-between" align="center" wrap="nowrap" mt="md">
|
||||
<Stack gap={0}>
|
||||
<Text fw={600} size="sm">
|
||||
{t("Autonomous agent runs")}
|
||||
</Text>
|
||||
<Text size="xs" c="dimmed">
|
||||
{t(
|
||||
"Keep an agent turn running server-side even if the browser disconnects; reconnect and follow it on reopen. Single-instance deployments only.",
|
||||
)}
|
||||
</Text>
|
||||
</Stack>
|
||||
<Switch
|
||||
label={t("Enabled")}
|
||||
labelPosition="left"
|
||||
checked={autonomousRunsEnabled}
|
||||
disabled={autonomousRunsToggleLoading}
|
||||
onChange={(e) =>
|
||||
handleToggleAutonomousRuns(e.currentTarget.checked)
|
||||
}
|
||||
/>
|
||||
</Group>
|
||||
|
||||
<Group mt="md" align="center">
|
||||
<Button
|
||||
variant="default"
|
||||
|
||||
@@ -26,6 +26,9 @@ export interface IWorkspace {
|
||||
aiDictation?: boolean;
|
||||
aiDictationStreaming?: boolean;
|
||||
aiPublicShareAssistant?: boolean;
|
||||
// Write-only field for updateWorkspace({ autonomousRuns }). Read state lives at
|
||||
// settings.ai.autonomousRuns.
|
||||
autonomousRuns?: boolean;
|
||||
trashRetentionDays?: number;
|
||||
// Default lifetime (HOURS) for new temporary notes; frozen per-note at creation.
|
||||
temporaryNoteHours?: number;
|
||||
@@ -65,6 +68,9 @@ export interface IWorkspaceAiSettings {
|
||||
dictation?: boolean;
|
||||
dictationStreaming?: boolean;
|
||||
publicShareAssistant?: boolean;
|
||||
// #184: detached agent runs (a run survives a browser disconnect and can be
|
||||
// reconnected to / live-followed on reopen). Gates the run-reconnect polling.
|
||||
autonomousRuns?: boolean;
|
||||
}
|
||||
|
||||
export interface IWorkspaceSharingSettings {
|
||||
|
||||
@@ -1,10 +1,3 @@
|
||||
export const HISTORY_INTERVAL = 5 * 60 * 1000;
|
||||
export const HISTORY_FAST_INTERVAL = 60 * 1000;
|
||||
export const HISTORY_FAST_THRESHOLD = 5 * 60 * 1000;
|
||||
|
||||
// #348 — debounce window for the per-page RAG re-embed job. Repeated saves
|
||||
// within this window collapse to a single delayed job (coalesced by a stable
|
||||
// jobId), so active editing does not pile up expensive re-embeds (external API
|
||||
// + page_embeddings rewrite, concurrency 1). The worker reads the CURRENT page
|
||||
// state at run time, so the last content within the window wins.
|
||||
export const EMBED_DEBOUNCE_MS = 30 * 1000;
|
||||
|
||||
@@ -431,17 +431,7 @@ describe('PersistenceExtension.onStoreDocument — Approach-A boundary snapshot'
|
||||
it('uses the canonical page.id (not the slugId doc name) for post-store side effects (#260)', async () => {
|
||||
const SLUG = 'slug-1'; // persistedHumanPage.slugId; findById resolves it
|
||||
const document = ydocFor(doc('NEW AGENT CONTENT'));
|
||||
// #348 — the transclusion sync now runs only when the new OR the previously
|
||||
// persisted content carries a transclusion-family node. Give the persisted
|
||||
// (old) content a pageEmbed so the sync path is exercised and the #260
|
||||
// UUID-vs-slugId contract asserted below is still verified.
|
||||
pageRepo.findById.mockResolvedValue({
|
||||
...persistedHumanPage('NEW AGENT CONTENT'),
|
||||
content: {
|
||||
type: 'doc',
|
||||
content: [{ type: 'pageEmbed', attrs: { sourcePageId: 'src-1' } }],
|
||||
},
|
||||
});
|
||||
pageRepo.findById.mockResolvedValue(persistedHumanPage('NEW AGENT CONTENT'));
|
||||
pageHistoryRepo.findPageLastHistory.mockResolvedValue(null);
|
||||
|
||||
// A `page.<slugId>` document name (the bug's smoking gun), agent store over
|
||||
|
||||
@@ -36,13 +36,11 @@ import {
|
||||
import { Page } from '@docmost/db/types/entity.types';
|
||||
import { CollabHistoryService } from '../services/collab-history.service';
|
||||
import {
|
||||
EMBED_DEBOUNCE_MS,
|
||||
HISTORY_FAST_INTERVAL,
|
||||
HISTORY_FAST_THRESHOLD,
|
||||
HISTORY_INTERVAL,
|
||||
} from '../constants';
|
||||
import { TransclusionService } from '../../core/page/transclusion/transclusion.service';
|
||||
import { hasTransclusionFamilyNodes } from '../../core/page/transclusion/utils/transclusion-prosemirror.util';
|
||||
import { observeCollabStore } from '../../integrations/metrics/metrics.registry';
|
||||
|
||||
/**
|
||||
@@ -417,18 +415,7 @@ export class PersistenceExtension implements Extension {
|
||||
// Use the canonical page UUID (page.id), not the doc-name id, which may be
|
||||
// a slugId for a `page.<slugId>` doc (#260). The transclusion/reference
|
||||
// syncs write uuid-typed columns, so a slugId here threw Postgres 22P02.
|
||||
//
|
||||
// #348 — skip the three sync SELECTs when neither the new content nor the
|
||||
// previously-persisted content has any transclusion/reference/pageEmbed
|
||||
// node: nothing to insert, and (the DB mirrors the old content) nothing to
|
||||
// delete. Whenever either side has one, run the idempotent sync exactly as
|
||||
// before so removals are still reconciled.
|
||||
if (
|
||||
hasTransclusionFamilyNodes(tiptapJson) ||
|
||||
hasTransclusionFamilyNodes(page.content)
|
||||
) {
|
||||
await this.syncTransclusion(page.id, page.workspaceId, tiptapJson);
|
||||
}
|
||||
await this.syncTransclusion(page.id, page.workspaceId, tiptapJson);
|
||||
}
|
||||
|
||||
if (page) {
|
||||
@@ -444,17 +431,7 @@ export class PersistenceExtension implements Extension {
|
||||
(m) => m.entityId,
|
||||
);
|
||||
|
||||
// #348 — only enqueue when the mentioned-user set actually GAINED a member.
|
||||
// The processor (processPageMention) already no-ops when every current
|
||||
// mention was present before (newMentions.length === 0), so skipping the
|
||||
// enqueue in that case is behavior-identical and avoids piling up no-op jobs
|
||||
// on every save of a page that merely CONTAINS (unchanged) mentions.
|
||||
const oldMentionedUserIdSet = new Set(oldMentionedUserIds);
|
||||
const hasNewMentionedUser = userMentions.some(
|
||||
(m) => !oldMentionedUserIdSet.has(m.entityId),
|
||||
);
|
||||
|
||||
if (hasNewMentionedUser) {
|
||||
if (userMentions.length > 0) {
|
||||
await this.notificationQueue.add(QueueJob.PAGE_MENTION_NOTIFICATION, {
|
||||
userMentions: userMentions.map((m) => ({
|
||||
userId: m.entityId,
|
||||
@@ -469,23 +446,12 @@ export class PersistenceExtension implements Extension {
|
||||
} as IPageMentionNotificationJob);
|
||||
}
|
||||
|
||||
await this.aiQueue.add(
|
||||
QueueJob.PAGE_CONTENT_UPDATED,
|
||||
{
|
||||
// Canonical UUID: the embedding reindex resolves pages by uuid, so a
|
||||
// slugId here threw Postgres 22P02 invalid-uuid (#260).
|
||||
pageIds: [page.id],
|
||||
workspaceId: page.workspaceId,
|
||||
},
|
||||
// #348 — coalesce re-embeds during active editing. A stable per-page
|
||||
// jobId + delay means repeated saves within EMBED_DEBOUNCE_MS collapse
|
||||
// to one delayed job instead of one expensive re-embed per save. The
|
||||
// worker reads the current page state at run time, so last content wins.
|
||||
// BullMQ forbids ':' in custom job ids (Redis key separator), so '-' is
|
||||
// used; page.id is a UUID, so the id is unique per page. removeOnComplete
|
||||
// (queue.module) frees the id after each run so the next window re-arms.
|
||||
{ jobId: `embed-${page.id}`, delay: EMBED_DEBOUNCE_MS },
|
||||
);
|
||||
await this.aiQueue.add(QueueJob.PAGE_CONTENT_UPDATED, {
|
||||
// Canonical UUID: the embedding reindex resolves pages by uuid, so a
|
||||
// slugId here threw Postgres 22P02 invalid-uuid (#260).
|
||||
pageIds: [page.id],
|
||||
workspaceId: page.workspaceId,
|
||||
});
|
||||
|
||||
await this.enqueuePageHistory(page, lastUpdatedSource);
|
||||
}
|
||||
|
||||
@@ -220,13 +220,6 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
|
||||
};
|
||||
|
||||
async maintainLock(documentName: string) {
|
||||
// #348 — clear any existing timer for this document before installing a new
|
||||
// one. Without this, a second maintainLock for the same document (a
|
||||
// reload-without-unload) overwrites this.locks[documentName] and leaks the
|
||||
// previous interval, which keeps firing SET forever with no way to clear it.
|
||||
if (this.locks[documentName]) {
|
||||
clearInterval(this.locks[documentName]);
|
||||
}
|
||||
this.locks[documentName] = setInterval(() => {
|
||||
this.pub.set(
|
||||
this.getKey(documentName),
|
||||
|
||||
@@ -4,21 +4,8 @@ export const CacheKey = {
|
||||
`perm:space-roles:${userId}:${spaceId}`,
|
||||
PAGE_CAN_EDIT: (userId: string, pageId: string) =>
|
||||
`perm:can-edit:${userId}:${pageId}`,
|
||||
// #348 — DomainMiddleware workspace resolution. Self-hosted resolves the single
|
||||
// workspace (constant key); cloud resolves by the request subdomain (lowercased
|
||||
// to match the case-insensitive `LOWER(hostname)` lookup). Every WorkspaceRepo
|
||||
// mutator busts these, so staleness is bounded by both explicit invalidation and
|
||||
// the short TTL below.
|
||||
WORKSPACE_SELF_HOSTED: 'workspace:self-hosted',
|
||||
WORKSPACE_BY_HOST: (subdomain: string) =>
|
||||
`workspace:byhost:${subdomain.toLowerCase()}`,
|
||||
};
|
||||
|
||||
// Permission caches dedupe repeated checks within and across short request bursts.
|
||||
// 5s keeps staleness on revocations bounded.
|
||||
export const PERMISSION_CACHE_TTL_MS = 5_000;
|
||||
|
||||
// #348 — workspace row changes rarely; a short TTL bounds staleness of
|
||||
// security-relevant fields (enforceSso/enforceMfa/status) even if an explicit
|
||||
// bust is ever missed, while still removing the per-request workspace query.
|
||||
export const WORKSPACE_CACHE_TTL_MS = 15_000;
|
||||
|
||||
@@ -1,42 +1,13 @@
|
||||
import { Inject, Injectable, NestMiddleware } from '@nestjs/common';
|
||||
import { Injectable, NestMiddleware, NotFoundException } from '@nestjs/common';
|
||||
import { FastifyRequest, FastifyReply } from 'fastify';
|
||||
import { CACHE_MANAGER } from '@nestjs/cache-manager';
|
||||
import { Cache } from 'cache-manager';
|
||||
import { EnvironmentService } from '../../integrations/environment/environment.service';
|
||||
import { WorkspaceRepo } from '@docmost/db/repos/workspace/workspace.repo';
|
||||
import { Workspace } from '@docmost/db/types/entity.types';
|
||||
import { withCache } from '../helpers/with-cache';
|
||||
import { CacheKey, WORKSPACE_CACHE_TTL_MS } from '../helpers/cache-keys';
|
||||
|
||||
// #348 — timestamptz columns on the workspace row. The cache store (Keyv/Redis)
|
||||
// JSON-serializes values, so a cached workspace comes back with these fields as
|
||||
// ISO strings. Reviving them to Date keeps the cached path byte-identical to the
|
||||
// direct DB path (postgres.js returns Date), so nothing downstream can observe a
|
||||
// cache hit vs miss. Idempotent: `new Date(date)` on an already-Date value is a
|
||||
// no-op-equivalent. Keep in sync with the workspace timestamptz columns.
|
||||
const WORKSPACE_DATE_FIELDS: Array<keyof Workspace> = [
|
||||
'createdAt',
|
||||
'updatedAt',
|
||||
'deletedAt',
|
||||
'trialEndAt',
|
||||
];
|
||||
|
||||
function reviveWorkspaceDates(workspace: Workspace): Workspace {
|
||||
for (const field of WORKSPACE_DATE_FIELDS) {
|
||||
const value = workspace[field];
|
||||
if (value != null) {
|
||||
(workspace as any)[field] = new Date(value as any);
|
||||
}
|
||||
}
|
||||
return workspace;
|
||||
}
|
||||
|
||||
@Injectable()
|
||||
export class DomainMiddleware implements NestMiddleware {
|
||||
constructor(
|
||||
private workspaceRepo: WorkspaceRepo,
|
||||
private environmentService: EnvironmentService,
|
||||
@Inject(CACHE_MANAGER) private readonly cacheManager: Cache,
|
||||
) {}
|
||||
async use(
|
||||
req: FastifyRequest['raw'],
|
||||
@@ -44,21 +15,13 @@ export class DomainMiddleware implements NestMiddleware {
|
||||
next: () => void,
|
||||
) {
|
||||
if (this.environmentService.isSelfHosted()) {
|
||||
// #348 — cache the single-workspace lookup that runs on every request.
|
||||
// Invalidated by every WorkspaceRepo mutator (see bustWorkspaceCache).
|
||||
const workspace = await withCache(
|
||||
this.cacheManager,
|
||||
CacheKey.WORKSPACE_SELF_HOSTED,
|
||||
WORKSPACE_CACHE_TTL_MS,
|
||||
() => this.workspaceRepo.findFirst(),
|
||||
);
|
||||
const workspace = await this.workspaceRepo.findFirst();
|
||||
if (!workspace) {
|
||||
//throw new NotFoundException('Workspace not found');
|
||||
(req as any).workspaceId = null;
|
||||
return next();
|
||||
}
|
||||
|
||||
reviveWorkspaceDates(workspace);
|
||||
// TODO: unify
|
||||
(req as any).workspaceId = workspace.id;
|
||||
(req as any).workspace = workspace;
|
||||
@@ -66,21 +29,13 @@ export class DomainMiddleware implements NestMiddleware {
|
||||
const header = req.headers.host;
|
||||
const subdomain = header.split('.')[0];
|
||||
|
||||
// #348 — cache per-subdomain workspace resolution. Keyed by subdomain (the
|
||||
// hostname column); busted per hostname by every WorkspaceRepo mutator.
|
||||
const workspace = await withCache(
|
||||
this.cacheManager,
|
||||
CacheKey.WORKSPACE_BY_HOST(subdomain),
|
||||
WORKSPACE_CACHE_TTL_MS,
|
||||
() => this.workspaceRepo.findByHostname(subdomain),
|
||||
);
|
||||
const workspace = await this.workspaceRepo.findByHostname(subdomain);
|
||||
|
||||
if (!workspace) {
|
||||
(req as any).workspaceId = null;
|
||||
return next();
|
||||
}
|
||||
|
||||
reviveWorkspaceDates(workspace);
|
||||
(req as any).workspaceId = workspace.id;
|
||||
(req as any).workspace = workspace;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,527 @@
|
||||
import { Logger } from '@nestjs/common';
|
||||
import {
|
||||
AiChatRunService,
|
||||
RunAlreadyActiveError,
|
||||
ONE_ACTIVE_RUN_PER_CHAT_INDEX,
|
||||
mapTurnStatusToRun,
|
||||
} from './ai-chat-run.service';
|
||||
|
||||
/** Shape a Postgres unique-violation the way the postgres.js driver surfaces it:
|
||||
* SQLSTATE 23505 + the offending index in `constraint_name`. */
|
||||
function uniqueViolation(constraintName: string): Error & {
|
||||
code: string;
|
||||
constraint_name: string;
|
||||
} {
|
||||
return Object.assign(
|
||||
new Error('duplicate key value violates unique constraint'),
|
||||
{
|
||||
code: '23505',
|
||||
constraint_name: constraintName,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Unit coverage for the #184 phase-1 run lifecycle (AiChatRunService) with a
|
||||
* hand-rolled mock repo — no Nest graph, no DB. The invariant under test is the
|
||||
* one that makes a run "autonomous": a run keeps going when its SUBSCRIBER (the
|
||||
* browser) detaches, and ONLY an explicit stop aborts it. We assert that at the
|
||||
* abort-signal level (the signal the agent loop actually consumes).
|
||||
*/
|
||||
|
||||
/** Minimal EnvironmentService stub. Single-instance (CLOUD unset) by default. */
|
||||
function makeEnv(isCloud = false) {
|
||||
return { isCloud: () => isCloud };
|
||||
}
|
||||
|
||||
function makeRepo(overrides: Record<string, jest.Mock> = {}) {
|
||||
return {
|
||||
insert: jest.fn(async (v: any) => ({
|
||||
id: 'run-1',
|
||||
status: v.status ?? 'running',
|
||||
chatId: v.chatId,
|
||||
workspaceId: v.workspaceId,
|
||||
})),
|
||||
update: jest.fn(async () => ({ id: 'run-1' })),
|
||||
markStopRequested: jest.fn(async () => ({ id: 'run-1' })),
|
||||
findActiveByChat: jest.fn(async () => undefined),
|
||||
findLatestByChat: jest.fn(async () => undefined),
|
||||
findById: jest.fn(async () => undefined),
|
||||
sweepRunning: jest.fn(async () => 0),
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
describe('mapTurnStatusToRun', () => {
|
||||
it('maps the turn terminal status to the run terminal status', () => {
|
||||
expect(mapTurnStatusToRun('completed')).toBe('succeeded');
|
||||
expect(mapTurnStatusToRun('error')).toBe('failed');
|
||||
expect(mapTurnStatusToRun('aborted')).toBe('aborted');
|
||||
});
|
||||
});
|
||||
|
||||
describe('AiChatRunService.onModuleInit (startup sweep)', () => {
|
||||
afterEach(() => jest.restoreAllMocks());
|
||||
|
||||
it('calls sweepRunning and resolves; logs when > 0', async () => {
|
||||
const repo = makeRepo({ sweepRunning: jest.fn(async () => 2) });
|
||||
const logSpy = jest
|
||||
.spyOn(Logger.prototype, 'log')
|
||||
.mockImplementation(() => undefined);
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await expect(svc.onModuleInit()).resolves.toBeUndefined();
|
||||
expect(repo.sweepRunning).toHaveBeenCalledTimes(1);
|
||||
expect(logSpy).toHaveBeenCalledTimes(1);
|
||||
expect(String(logSpy.mock.calls[0][0])).toContain('2');
|
||||
});
|
||||
|
||||
it('a sweep failure is swallowed (never blocks startup)', async () => {
|
||||
const repo = makeRepo({
|
||||
sweepRunning: jest.fn(async () => {
|
||||
throw new Error('db down');
|
||||
}),
|
||||
});
|
||||
const warnSpy = jest
|
||||
.spyOn(Logger.prototype, 'warn')
|
||||
.mockImplementation(() => undefined);
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await expect(svc.onModuleInit()).resolves.toBeUndefined();
|
||||
// The first warn is the sweep failure (the multi-instance warn never fires
|
||||
// single-instance), so the message is the db error.
|
||||
expect(String(warnSpy.mock.calls[0][0])).toContain('db down');
|
||||
});
|
||||
|
||||
it('F1 (DECISION C): the boot sweep is UNCONDITIONAL — sweepRunning is called with NO staleness window, so a fresh running run (updatedAt = now) is settled, not skipped', async () => {
|
||||
// The bug: a fast restart (deploy/OOM within minutes of the last step) left a
|
||||
// run stuck 'running' under the old 10-min window, 409ing every later turn in
|
||||
// the chat. The fix settles ALL pending|running on boot. We assert the service
|
||||
// invokes sweepRunning with no `staleMs` (the unconditional path); the repo's
|
||||
// own spec proves no-window => no updatedAt filter.
|
||||
const repo = makeRepo({ sweepRunning: jest.fn(async () => 1) });
|
||||
jest.spyOn(Logger.prototype, 'log').mockImplementation(() => undefined);
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await svc.onModuleInit();
|
||||
expect(repo.sweepRunning).toHaveBeenCalledTimes(1);
|
||||
const callArgs = repo.sweepRunning.mock.calls[0] as unknown[];
|
||||
const firstArg = callArgs[0] as { staleMs?: number } | undefined;
|
||||
// Either no opts at all, or opts without a staleMs window => unconditional.
|
||||
expect(firstArg?.staleMs).toBeUndefined();
|
||||
});
|
||||
|
||||
it('F2 (DECISION A): warns at startup that autonomousRuns is single-instance-only when a horizontally-scaled deployment (CLOUD) is detected', async () => {
|
||||
const repo = makeRepo();
|
||||
const warnSpy = jest
|
||||
.spyOn(Logger.prototype, 'warn')
|
||||
.mockImplementation(() => undefined);
|
||||
const svc = new AiChatRunService(repo as never, makeEnv(true) as never);
|
||||
await svc.onModuleInit();
|
||||
const warned = warnSpy.mock.calls.some((c) =>
|
||||
/single-instance-only/i.test(String(c[0])),
|
||||
);
|
||||
expect(warned).toBe(true);
|
||||
});
|
||||
|
||||
it('F2: does NOT warn about multi-instance on a single-instance (CLOUD unset) deployment', async () => {
|
||||
const repo = makeRepo();
|
||||
const warnSpy = jest
|
||||
.spyOn(Logger.prototype, 'warn')
|
||||
.mockImplementation(() => undefined);
|
||||
const svc = new AiChatRunService(repo as never, makeEnv(false) as never);
|
||||
await svc.onModuleInit();
|
||||
const warned = warnSpy.mock.calls.some((c) =>
|
||||
/single-instance-only/i.test(String(c[0])),
|
||||
);
|
||||
expect(warned).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('AiChatRunService run lifecycle', () => {
|
||||
it('beginRun inserts a running row and registers a live abort controller', async () => {
|
||||
const repo = makeRepo();
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
const handle = await svc.beginRun({
|
||||
chatId: 'chat-1',
|
||||
workspaceId: 'ws-1',
|
||||
userId: 'user-1',
|
||||
});
|
||||
expect(repo.insert).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
chatId: 'chat-1',
|
||||
workspaceId: 'ws-1',
|
||||
createdBy: 'user-1',
|
||||
status: 'running',
|
||||
trigger: 'user',
|
||||
}),
|
||||
);
|
||||
expect(handle.runId).toBe('run-1');
|
||||
expect(handle.signal.aborted).toBe(false);
|
||||
expect(svc.isLocallyActive('run-1')).toBe(true);
|
||||
});
|
||||
|
||||
it('beginRun REJECTS the racer: a 23505 on the one-active-per-chat index throws RunAlreadyActiveError (not swallowed) and registers no controller', async () => {
|
||||
// The race: the controller's cheap pre-check passed for BOTH concurrent
|
||||
// turns, so the loser's INSERT hits the partial unique index. That rejection
|
||||
// is the authoritative gate — it must surface, not be swallowed into an
|
||||
// untracked turn.
|
||||
const repo = makeRepo({
|
||||
insert: jest.fn(async () => {
|
||||
throw uniqueViolation(ONE_ACTIVE_RUN_PER_CHAT_INDEX);
|
||||
}),
|
||||
});
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await expect(
|
||||
svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'user-1' }),
|
||||
).rejects.toBeInstanceOf(RunAlreadyActiveError);
|
||||
// No controller leaked for a rejected start.
|
||||
expect(svc.isLocallyActive('run-1')).toBe(false);
|
||||
});
|
||||
|
||||
it('beginRun does NOT mask an unrelated unique violation as already-active', async () => {
|
||||
// A 23505 on some OTHER constraint is a real bug, not the race — it must
|
||||
// propagate unchanged so it is never silently treated as "already active".
|
||||
const other = uniqueViolation('ai_chat_runs_pkey');
|
||||
const repo = makeRepo({
|
||||
insert: jest.fn(async () => {
|
||||
throw other;
|
||||
}),
|
||||
});
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await expect(
|
||||
svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'user-1' }),
|
||||
).rejects.toBe(other);
|
||||
});
|
||||
|
||||
it('beginRun propagates a non-unique insert failure unchanged', async () => {
|
||||
const boom = new Error('connection reset');
|
||||
const repo = makeRepo({
|
||||
insert: jest.fn(async () => {
|
||||
throw boom;
|
||||
}),
|
||||
});
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await expect(
|
||||
svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'user-1' }),
|
||||
).rejects.toBe(boom);
|
||||
});
|
||||
|
||||
it('two concurrent begins on one chat: exactly one wins, the other is rejected as already-active', async () => {
|
||||
// Integration-style: model the DB partial unique index with a one-shot slot.
|
||||
// The first insert claims it; the second hits a 23505 on the active index.
|
||||
let slotTaken = false;
|
||||
const repo = makeRepo({
|
||||
insert: jest.fn(async (v: any) => {
|
||||
if (slotTaken) throw uniqueViolation(ONE_ACTIVE_RUN_PER_CHAT_INDEX);
|
||||
slotTaken = true;
|
||||
return { id: 'run-win', status: v.status, chatId: v.chatId };
|
||||
}),
|
||||
});
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
const results = await Promise.allSettled([
|
||||
svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'user-1' }),
|
||||
svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'user-1' }),
|
||||
]);
|
||||
const fulfilled = results.filter((r) => r.status === 'fulfilled');
|
||||
const rejected = results.filter((r) => r.status === 'rejected');
|
||||
expect(fulfilled).toHaveLength(1);
|
||||
expect(rejected).toHaveLength(1);
|
||||
expect((rejected[0] as PromiseRejectedResult).reason).toBeInstanceOf(
|
||||
RunAlreadyActiveError,
|
||||
);
|
||||
// Exactly the winner is locally active.
|
||||
expect(svc.isLocallyActive('run-win')).toBe(true);
|
||||
});
|
||||
|
||||
it('a SUBSCRIBER detaching does NOT abort the run (only an explicit stop does)', async () => {
|
||||
const repo = makeRepo();
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
const handle = await svc.beginRun({
|
||||
chatId: 'chat-1',
|
||||
workspaceId: 'ws-1',
|
||||
userId: 'user-1',
|
||||
});
|
||||
// Model a browser disconnect: nothing in the run service is told to stop.
|
||||
// The signal the agent loop consumes must stay un-aborted and the run stays
|
||||
// locally active — i.e. it keeps running server-side.
|
||||
expect(handle.signal.aborted).toBe(false);
|
||||
expect(svc.isLocallyActive('run-1')).toBe(true);
|
||||
// markStopRequested was never called by a mere detach.
|
||||
expect(repo.markStopRequested).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('requestStop aborts the live controller, marks the row, and reports true', async () => {
|
||||
const repo = makeRepo();
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
const handle = await svc.beginRun({
|
||||
chatId: 'chat-1',
|
||||
workspaceId: 'ws-1',
|
||||
userId: 'user-1',
|
||||
});
|
||||
const aborted = jest.fn();
|
||||
handle.signal.addEventListener('abort', aborted);
|
||||
|
||||
const result = await svc.requestStop('run-1', 'ws-1');
|
||||
|
||||
expect(result).toBe(true);
|
||||
expect(handle.signal.aborted).toBe(true);
|
||||
expect(aborted).toHaveBeenCalledTimes(1);
|
||||
expect(repo.markStopRequested).toHaveBeenCalledWith('run-1', 'ws-1');
|
||||
});
|
||||
|
||||
it('requestStop on a run this replica does NOT hold still marks the row (true)', async () => {
|
||||
// e.g. after a restart, or a sibling replica owns the controller. The row is
|
||||
// marked so the owning replica/sweep settles it; we report a stop took effect.
|
||||
const repo = makeRepo({
|
||||
markStopRequested: jest.fn(async () => ({ id: 'run-9' })),
|
||||
});
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
const result = await svc.requestStop('run-9', 'ws-1');
|
||||
expect(result).toBe(true);
|
||||
expect(svc.isLocallyActive('run-9')).toBe(false);
|
||||
});
|
||||
|
||||
it('requestStop still aborts the live controller when markStopRequested rejects (transient DB error)', async () => {
|
||||
// F15: the in-memory abort is the ONLY thing that stops a run and must not be
|
||||
// hostage to the audit write of stop_requested_at. A transient failure on
|
||||
// markStopRequested must NOT prevent abort() nor make requestStop throw.
|
||||
const warnSpy = jest
|
||||
.spyOn(Logger.prototype, 'warn')
|
||||
.mockImplementation(() => undefined);
|
||||
const repo = makeRepo({
|
||||
markStopRequested: jest.fn(async () => {
|
||||
throw new Error('pool exhausted');
|
||||
}),
|
||||
});
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
const handle = await svc.beginRun({
|
||||
chatId: 'chat-1',
|
||||
workspaceId: 'ws-1',
|
||||
userId: 'user-1',
|
||||
});
|
||||
const aborted = jest.fn();
|
||||
handle.signal.addEventListener('abort', aborted);
|
||||
|
||||
// Does NOT throw despite the DB write rejecting.
|
||||
const result = await svc.requestStop('run-1', 'ws-1');
|
||||
|
||||
// The live turn was aborted even though the audit write failed...
|
||||
expect(handle.signal.aborted).toBe(true);
|
||||
expect(aborted).toHaveBeenCalledTimes(1);
|
||||
expect(repo.markStopRequested).toHaveBeenCalledWith('run-1', 'ws-1');
|
||||
// ...the catch branch logged the swallowed failure...
|
||||
expect(warnSpy).toHaveBeenCalledTimes(1);
|
||||
// ...and a stop is reported as having taken effect (the entry existed).
|
||||
expect(result).toBe(true);
|
||||
warnSpy.mockRestore();
|
||||
});
|
||||
|
||||
it('requestStop on an already-settled run (nothing active) reports false', async () => {
|
||||
const repo = makeRepo({
|
||||
markStopRequested: jest.fn(async () => undefined),
|
||||
});
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
const result = await svc.requestStop('run-done', 'ws-1');
|
||||
expect(result).toBe(false);
|
||||
});
|
||||
|
||||
it('finalizeRun settles the row to the mapped status with finishedAt and drops the in-memory entry', async () => {
|
||||
const repo = makeRepo();
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await svc.beginRun({
|
||||
chatId: 'chat-1',
|
||||
workspaceId: 'ws-1',
|
||||
userId: 'user-1',
|
||||
});
|
||||
expect(svc.isLocallyActive('run-1')).toBe(true);
|
||||
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'error', 'provider blew up');
|
||||
|
||||
expect(svc.isLocallyActive('run-1')).toBe(false);
|
||||
expect(repo.update).toHaveBeenCalledWith(
|
||||
'run-1',
|
||||
'ws-1',
|
||||
expect.objectContaining({
|
||||
status: 'failed',
|
||||
error: 'provider blew up',
|
||||
finishedAt: expect.any(Date),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('finalizeRun is IDEMPOTENT: a second settle no-ops (single terminal write)', async () => {
|
||||
// The #184 review fix: AiChatService.stream wraps the turn in a safety-net
|
||||
// catch that settles a failed turn AND streamText's terminal callback may
|
||||
// also settle — both routes call finalizeRun. Only the FIRST may write the
|
||||
// terminal row; the second must no-op so a late settle can never clobber the
|
||||
// real terminal status or double-write the row.
|
||||
const repo = makeRepo();
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await svc.beginRun({
|
||||
chatId: 'chat-1',
|
||||
workspaceId: 'ws-1',
|
||||
userId: 'user-1',
|
||||
});
|
||||
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'error', 'first');
|
||||
expect(svc.isLocallyActive('run-1')).toBe(false);
|
||||
// A second settle (e.g. a streamText callback firing after the catch) no-ops.
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'completed', undefined);
|
||||
|
||||
expect(repo.update).toHaveBeenCalledTimes(1);
|
||||
expect(repo.update).toHaveBeenCalledWith(
|
||||
'run-1',
|
||||
'ws-1',
|
||||
expect.objectContaining({ status: 'failed', error: 'first' }),
|
||||
);
|
||||
});
|
||||
|
||||
it('CONCURRENCY: two simultaneous finalizeRun on the same run write the terminal row EXACTLY ONCE (the 2nd caller exits synchronously at the atomic claim)', async () => {
|
||||
// The CRITICAL race: AiChatService.stream's safety-net catch settles the turn
|
||||
// to 'error' while a streamText terminal callback also settles it — both call
|
||||
// finalizeRun for the SAME runId. The once-gate must close ATOMICALLY: a
|
||||
// `settled.has` check alone is read BEFORE the awaited UPDATE, so both callers
|
||||
// would pass it and BOTH write the row (last-write-wins clobber + double
|
||||
// write). The fix claims the run with a SYNCHRONOUS `active.delete` before any
|
||||
// await, so the second caller returns in the same tick, before the UPDATE.
|
||||
//
|
||||
// We force the two calls to overlap by making `update` return a promise we
|
||||
// resolve only AFTER both finalizeRun calls have run their synchronous bodies.
|
||||
let resolveUpdate!: (v: unknown) => void;
|
||||
const updateGate = new Promise((res) => {
|
||||
resolveUpdate = res;
|
||||
});
|
||||
const update = jest.fn(() => updateGate);
|
||||
const repo = makeRepo({ update });
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await svc.beginRun({
|
||||
chatId: 'chat-1',
|
||||
workspaceId: 'ws-1',
|
||||
userId: 'user-1',
|
||||
});
|
||||
|
||||
// Fire both before the (pending) update resolves. The first synchronously
|
||||
// claims the entry (active.delete) and awaits update; the second, started in
|
||||
// the same macrotask, finds the entry already gone and returns at the claim
|
||||
// WITHOUT ever calling update.
|
||||
const p1 = svc.finalizeRun('run-1', 'ws-1', 'completed');
|
||||
const p2 = svc.finalizeRun('run-1', 'ws-1', 'error', 'safety-net');
|
||||
|
||||
// The decisive assertion: exactly one caller reached the terminal UPDATE.
|
||||
expect(update).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Let the single in-flight update land; both calls resolve cleanly.
|
||||
resolveUpdate({ id: 'run-1' });
|
||||
await Promise.all([p1, p2]);
|
||||
|
||||
expect(update).toHaveBeenCalledTimes(1);
|
||||
// The winner is the FIRST caller ('completed' -> 'succeeded'); the late
|
||||
// 'error' settle never wrote, so it could not clobber the real status.
|
||||
expect(update).toHaveBeenCalledWith(
|
||||
'run-1',
|
||||
'ws-1',
|
||||
expect.objectContaining({ status: 'succeeded' }),
|
||||
);
|
||||
expect(svc.isLocallyActive('run-1')).toBe(false);
|
||||
});
|
||||
|
||||
it('F6: a TRANSIENT terminal-write failure is ridden out by the bounded retry — the run is settled, not stranded', async () => {
|
||||
// The bug: finalizeRun used to DROP the in-memory entry BEFORE the terminal
|
||||
// UPDATE, then only warn-log a failure. A single transient blip (pool
|
||||
// exhaustion / deadlock / connection hiccup) on that PK UPDATE left the row
|
||||
// 'running' with nothing left to recover it -> every later turn in that chat
|
||||
// 409s until a restart. The fix updates FIRST and retries.
|
||||
let calls = 0;
|
||||
const repo = makeRepo({
|
||||
update: jest.fn(async () => {
|
||||
calls += 1;
|
||||
if (calls === 1) throw new Error('deadlock detected');
|
||||
return { id: 'run-1' };
|
||||
}),
|
||||
});
|
||||
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined);
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await svc.beginRun({
|
||||
chatId: 'chat-1',
|
||||
workspaceId: 'ws-1',
|
||||
userId: 'user-1',
|
||||
});
|
||||
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'completed');
|
||||
|
||||
// The retry landed the terminal write: the entry is dropped (slot freed) and
|
||||
// the row carries the real terminal status — NOT stranded at 'running'.
|
||||
expect(svc.isLocallyActive('run-1')).toBe(false);
|
||||
expect(repo.update).toHaveBeenCalledTimes(2);
|
||||
expect(repo.update).toHaveBeenLastCalledWith(
|
||||
'run-1',
|
||||
'ws-1',
|
||||
expect.objectContaining({ status: 'succeeded' }),
|
||||
);
|
||||
});
|
||||
|
||||
it('F6: if the terminal write keeps failing, the entry is RETAINED and a LATER settle completes it (chat not permanently 409d)', async () => {
|
||||
// Worst case: the DB is down for the whole first finalize (all attempts fail).
|
||||
// The run must NOT be silently lost — the entry stays so a subsequent settle
|
||||
// (a streamText callback, requestStop -> onAbort, or a future sweep) can retry.
|
||||
let healthy = false;
|
||||
const repo = makeRepo({
|
||||
update: jest.fn(async () => {
|
||||
if (!healthy) throw new Error('pool exhausted');
|
||||
return { id: 'run-1' };
|
||||
}),
|
||||
});
|
||||
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined);
|
||||
const errorSpy = jest
|
||||
.spyOn(Logger.prototype, 'error')
|
||||
.mockImplementation(() => undefined);
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await svc.beginRun({
|
||||
chatId: 'chat-1',
|
||||
workspaceId: 'ws-1',
|
||||
userId: 'user-1',
|
||||
});
|
||||
|
||||
// First settle: every bounded attempt fails -> entry retained, NOT settled.
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'completed');
|
||||
expect(svc.isLocallyActive('run-1')).toBe(true);
|
||||
// F12: the give-up emits ONE explicit, greppable ERROR (run + chat context)
|
||||
// so an operator can tell "gave up, run held in memory" from a per-attempt
|
||||
// blip — distinct from the per-attempt warns.
|
||||
const gaveUp = errorSpy.mock.calls.some(
|
||||
(c) =>
|
||||
/NON-TERMINAL/.test(String(c[0])) &&
|
||||
/run-1/.test(String(c[0])) &&
|
||||
/chat-1/.test(String(c[0])),
|
||||
);
|
||||
expect(gaveUp).toBe(true);
|
||||
|
||||
// The DB recovers; a later settle now succeeds and frees the slot.
|
||||
healthy = true;
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'completed');
|
||||
expect(svc.isLocallyActive('run-1')).toBe(false);
|
||||
expect(repo.update).toHaveBeenLastCalledWith(
|
||||
'run-1',
|
||||
'ws-1',
|
||||
expect.objectContaining({ status: 'succeeded' }),
|
||||
);
|
||||
|
||||
// And it is now idempotent: a further settle no-ops (terminal row already
|
||||
// written), so a double-settle can never clobber the real status.
|
||||
const callsBefore = repo.update.mock.calls.length;
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'error', 'late');
|
||||
expect(repo.update).toHaveBeenCalledTimes(callsBefore);
|
||||
});
|
||||
|
||||
it('recordStep / linkAssistantMessage are best-effort: a repo failure is swallowed', async () => {
|
||||
const repo = makeRepo({
|
||||
update: jest.fn(async () => {
|
||||
throw new Error('transient');
|
||||
}),
|
||||
});
|
||||
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined);
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await expect(svc.recordStep('run-1', 'ws-1', 3)).resolves.toBeUndefined();
|
||||
await expect(
|
||||
svc.linkAssistantMessage('run-1', 'ws-1', 'msg-1'),
|
||||
).resolves.toBeUndefined();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,452 @@
|
||||
import { Injectable, Logger, OnModuleInit } from '@nestjs/common';
|
||||
import { AiChatRunRepo } from '@docmost/db/repos/ai-chat/ai-chat-run.repo';
|
||||
import { AiChatRun } from '@docmost/db/types/entity.types';
|
||||
import { isUniqueViolation, violatedConstraint } from '@docmost/db/utils';
|
||||
import { EnvironmentService } from '../../integrations/environment/environment.service';
|
||||
|
||||
/** Name of the partial unique index enforcing "one active run per chat" (see the
|
||||
* ai_chat_runs migration). A 23505 on THIS constraint is the race-safe signal
|
||||
* that a concurrent turn already owns the chat — distinct from any other unique
|
||||
* collision, which must NOT be silently treated as "already active". */
|
||||
export const ONE_ACTIVE_RUN_PER_CHAT_INDEX = 'ai_chat_runs_one_active_per_chat';
|
||||
|
||||
/**
|
||||
* Thrown by {@link AiChatRunService.beginRun} when the run-row INSERT loses the
|
||||
* race for a chat's single active slot (the partial unique index rejects it with
|
||||
* a 23505). This is the AUTHORITATIVE concurrency gate: the controller's cheap
|
||||
* pre-check is only a fast-path, and a request that slips past it must NOT run
|
||||
* untracked. The caller (AiChatService.stream) translates this into a 409 and
|
||||
* aborts the turn BEFORE any AI/provider call.
|
||||
*/
|
||||
export class RunAlreadyActiveError extends Error {
|
||||
constructor(public readonly chatId: string) {
|
||||
super(`An agent run is already in progress for chat ${chatId}`);
|
||||
this.name = 'RunAlreadyActiveError';
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* The terminal status of a TURN (the #183 assistant-row lifecycle) maps onto the
|
||||
* terminal status of a RUN (#184). A turn that completed -> the run succeeded; a
|
||||
* turn that errored -> the run failed; a turn aborted (explicit user stop) -> the
|
||||
* run aborted. Pure + unit-testable.
|
||||
*/
|
||||
export type TurnTerminalStatus = 'completed' | 'error' | 'aborted';
|
||||
export type RunTerminalStatus = 'succeeded' | 'failed' | 'aborted';
|
||||
|
||||
export function mapTurnStatusToRun(
|
||||
status: TurnTerminalStatus,
|
||||
): RunTerminalStatus {
|
||||
switch (status) {
|
||||
case 'completed':
|
||||
return 'succeeded';
|
||||
case 'error':
|
||||
return 'failed';
|
||||
case 'aborted':
|
||||
return 'aborted';
|
||||
}
|
||||
}
|
||||
|
||||
/** An in-flight run held in process memory: its AbortController is the ONLY thing
|
||||
* that can stop the turn (an explicit user stop), independent of the browser
|
||||
* socket. A mere disconnect never touches it, so the run keeps going. */
|
||||
interface ActiveRun {
|
||||
controller: AbortController;
|
||||
chatId: string;
|
||||
workspaceId: string;
|
||||
}
|
||||
|
||||
/** The live handle the streaming path drives a run through (returned by
|
||||
* {@link AiChatRunService.beginRun}). The `signal` governs the agent loop's
|
||||
* abort — wired to the run, NOT to the HTTP socket. */
|
||||
export interface RunHandle {
|
||||
runId: string;
|
||||
signal: AbortSignal;
|
||||
}
|
||||
|
||||
/**
|
||||
* AiChatRunService (#184 phase 1) — owns the agent RUN as a first-class,
|
||||
* server-side lifecycle object detached from the HTTP request / browser window.
|
||||
*
|
||||
* Responsibilities:
|
||||
* - create a run row when a turn starts (inserted directly as 'running'; the
|
||||
* 'pending' status is only the column default + a reserved value, never
|
||||
* written by code in phase 1) and register an in-memory AbortController for it
|
||||
* (the explicit-stop lever);
|
||||
* - finalize the run row (succeeded / failed / aborted) and unregister it;
|
||||
* - service an EXPLICIT user stop (`requestStop`) — the ONLY thing that aborts a
|
||||
* run; a browser disconnect deliberately does NOT;
|
||||
* - crash-recovery sweep of dangling runs on startup.
|
||||
*
|
||||
* The agent loop itself still runs in AiChatService.stream (reusing #183's
|
||||
* step-granular durable write path, `consumeStream` already drains it independent
|
||||
* of the socket); this service only wraps it in a durable lifecycle and an
|
||||
* abort handle that outlives the subscriber.
|
||||
*/
|
||||
@Injectable()
|
||||
export class AiChatRunService implements OnModuleInit {
|
||||
private readonly logger = new Logger(AiChatRunService.name);
|
||||
|
||||
// runId -> ActiveRun. Process-local on purpose (phase 1 is single-process /
|
||||
// in-memory transport; a cross-process BullMQ runner + Redis stop-signal is
|
||||
// deferred to phase 2). A stop for a runId not in this map (e.g. after a
|
||||
// restart) still records `stop_requested_at` on the row.
|
||||
private readonly active = new Map<string, ActiveRun>();
|
||||
|
||||
// runIds whose TERMINAL row write has SUCCEEDED — the idempotency once-gate
|
||||
// (F6). A finalize must short-circuit only AFTER the terminal write has landed,
|
||||
// NOT merely after the in-memory entry was dropped: a transient UPDATE failure
|
||||
// has to stay retryable, so "already settled" means "row already terminal", not
|
||||
// "entry already gone". Grows by one short UUID per finished run over process
|
||||
// uptime — negligible in phase 1's single process.
|
||||
private readonly settled = new Set<string>();
|
||||
|
||||
// Bounded retry for the terminal write (F6): a single PK UPDATE can fail
|
||||
// transiently under many fire-and-forget writes (pool exhaustion, deadlock, a
|
||||
// brief connection blip). Riding out that blip in-place matters because the
|
||||
// dominant success path (streamText onFinish) settles exactly ONCE — if that
|
||||
// write is dropped and never retried, the row is stranded 'running' and the
|
||||
// one-active-run gate 409s every future turn in the chat until a restart (no
|
||||
// periodic sweep in phase 1).
|
||||
private static readonly FINALIZE_MAX_ATTEMPTS = 3;
|
||||
private static readonly FINALIZE_RETRY_BASE_MS = 50;
|
||||
|
||||
constructor(
|
||||
private readonly runRepo: AiChatRunRepo,
|
||||
private readonly environment: EnvironmentService,
|
||||
) {}
|
||||
|
||||
/**
|
||||
* Crash-recovery sweep on server start: settle EVERY run still left
|
||||
* pending/running to 'aborted' (F1 / DECISION C). The boot sweep is
|
||||
* UNCONDITIONAL — no staleness window — because phase 1 is single-process: on a
|
||||
* fresh boot any pending|running run is definitionally hung (no live runner owns
|
||||
* it), so even a fast restart (deploy/OOM within minutes of the last step) can
|
||||
* no longer leave a run stuck 'running' forever (which would make the
|
||||
* one-active-run gate 409 every future turn in that chat). The staleness window
|
||||
* is reintroduced only for the phase-2 multi-instance timer sweep, where a
|
||||
* booting replica must not abort a run another replica is actively executing.
|
||||
* Best-effort — a sweep failure is logged but MUST NOT block startup (mirrors
|
||||
* AiChatService.onModuleInit for #183).
|
||||
*/
|
||||
async onModuleInit(): Promise<void> {
|
||||
this.warnIfMultiInstance();
|
||||
try {
|
||||
// No `staleMs`: unconditional boot sweep (F1). See AiChatRunRepo.sweepRunning.
|
||||
const swept = await this.runRepo.sweepRunning();
|
||||
if (swept > 0) {
|
||||
this.logger.log(
|
||||
`Startup sweep: marked ${swept} dangling agent run(s) as 'aborted'.`,
|
||||
);
|
||||
}
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`Startup sweep of dangling runs failed: ${
|
||||
err instanceof Error ? err.message : 'unknown error'
|
||||
}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* F2 (DECISION A): autonomous runs are SINGLE-INSTANCE-ONLY in phase 1. An
|
||||
* explicit Stop, and the in-memory AbortController that backs it, are
|
||||
* process-local: a Stop only aborts the live turn if it lands on the SAME
|
||||
* replica that owns the run (it still stamps `stop_requested_at` cross-instance,
|
||||
* but nothing reads that flag during an active run yet). Cross-instance pub/sub
|
||||
* stop is phase 2. So if the deployment is horizontally scaled, warn loudly at
|
||||
* startup that a Stop may not reach a run executing on another replica.
|
||||
*
|
||||
* DETECTION: this codebase always wires the socket.io Redis adapter (REDIS_URL
|
||||
* is mandatory), so the adapter alone is NOT a horizontal-scaling signal. The
|
||||
* authoritative signal the codebase has is `CLOUD=true` (EnvironmentService
|
||||
* .isCloud()), the Docmost-cloud multi-replica deployment. We warn whenever that
|
||||
* is set, because any workspace could enable settings.ai.autonomousRuns. A
|
||||
* self-hosted operator running multiple replicas behind a load balancer is also
|
||||
* multi-instance; the deploy docs (.env.example / AGENTS.md) spell out the
|
||||
* single-instance constraint for that case.
|
||||
*/
|
||||
private warnIfMultiInstance(): void {
|
||||
if (this.environment.isCloud()) {
|
||||
this.logger.warn(
|
||||
'Autonomous agent runs (settings.ai.autonomousRuns) are SINGLE-INSTANCE-ONLY ' +
|
||||
'in phase 1: a horizontally-scaled deployment was detected (CLOUD=true). ' +
|
||||
'An explicit Stop only aborts a run executing on the same replica that owns ' +
|
||||
'it (cross-instance Stop is not yet reliable — phase 2). Run a single ' +
|
||||
'instance if you enable autonomousRuns, or keep the flag off.',
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Start a run for a turn: insert the run row (status 'running', startedAt now),
|
||||
* register a fresh AbortController for it, and return a {@link RunHandle} whose
|
||||
* `signal` the agent loop uses. The DB partial unique index guarantees at most
|
||||
* one active run per chat — a second concurrent start on the same chat REJECTS
|
||||
* at the insert (a 23505 on {@link ONE_ACTIVE_RUN_PER_CHAT_INDEX}). That
|
||||
* rejection is the AUTHORITATIVE race gate: it is surfaced as a distinct
|
||||
* {@link RunAlreadyActiveError} (NOT swallowed), so the caller turns it into a
|
||||
* 409 and never streams an untracked turn. The controller is registered AFTER a
|
||||
* successful insert so a rejected start leaks nothing.
|
||||
*/
|
||||
async beginRun(args: {
|
||||
chatId: string;
|
||||
workspaceId: string;
|
||||
userId: string;
|
||||
trigger?: string;
|
||||
}): Promise<RunHandle> {
|
||||
let run: AiChatRun;
|
||||
try {
|
||||
run = await this.runRepo.insert({
|
||||
chatId: args.chatId,
|
||||
workspaceId: args.workspaceId,
|
||||
createdBy: args.userId,
|
||||
trigger: args.trigger ?? 'user',
|
||||
status: 'running',
|
||||
startedAt: new Date(),
|
||||
});
|
||||
} catch (err) {
|
||||
// The race backstop: a concurrent turn already holds this chat's single
|
||||
// active slot, so the partial unique index rejected our insert. Surface a
|
||||
// distinct signal — the caller MUST reject this turn (409), not run it
|
||||
// untracked. Any OTHER error propagates unchanged.
|
||||
if (
|
||||
isUniqueViolation(err) &&
|
||||
violatedConstraint(err) === ONE_ACTIVE_RUN_PER_CHAT_INDEX
|
||||
) {
|
||||
throw new RunAlreadyActiveError(args.chatId);
|
||||
}
|
||||
throw err;
|
||||
}
|
||||
const controller = new AbortController();
|
||||
this.active.set(run.id, {
|
||||
controller,
|
||||
chatId: args.chatId,
|
||||
workspaceId: args.workspaceId,
|
||||
});
|
||||
return { runId: run.id, signal: controller.signal };
|
||||
}
|
||||
|
||||
/** Link the assistant message (the #183 projection) to its run. Best-effort. */
|
||||
async linkAssistantMessage(
|
||||
runId: string,
|
||||
workspaceId: string,
|
||||
assistantMessageId: string,
|
||||
): Promise<void> {
|
||||
try {
|
||||
await this.runRepo.update(runId, workspaceId, { assistantMessageId });
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`Failed to link assistant message to run ${runId}: ${
|
||||
err instanceof Error ? err.message : 'unknown error'
|
||||
}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/** Persist progress: bump the run's finished-step count. Best-effort (never
|
||||
* blocks or breaks the stream). */
|
||||
async recordStep(
|
||||
runId: string,
|
||||
workspaceId: string,
|
||||
stepCount: number,
|
||||
): Promise<void> {
|
||||
try {
|
||||
await this.runRepo.update(runId, workspaceId, { stepCount });
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`Failed to record step for run ${runId}: ${
|
||||
err instanceof Error ? err.message : 'unknown error'
|
||||
}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Finalize a run to its terminal status (succeeded / failed / aborted),
|
||||
* stamping finishedAt + any error. Best-effort, but ROBUST against a transient
|
||||
* terminal-write failure (F6) AND atomically safe against a concurrent settle.
|
||||
*
|
||||
* ATOMIC ONCE-CLAIM (the gate must close in ONE synchronous tick): two
|
||||
* finalizeRun calls for the SAME run can race — the documented real path is
|
||||
* AiChatService.stream's safety-net catch settling the turn to 'error' while a
|
||||
* streamText terminal callback (onFinish/onAbort/onError) ALSO settles it. The
|
||||
* `settled.has` check alone is NOT a gate: it is read BEFORE the awaited UPDATE,
|
||||
* so two callers can both see `false` and both write the row (last-write-wins
|
||||
* clobbers the real terminal status, and the bounded retry only widens that
|
||||
* window). The claim therefore happens via `active.delete`, a SYNCHRONOUS
|
||||
* check-and-clear with NO await between the gate and the entry removal: the
|
||||
* second concurrent caller finds the entry already gone and returns in the same
|
||||
* tick, before any UPDATE. The transition "nobody is finalizing" -> "I am
|
||||
* finalizing" is thus a single atomic step.
|
||||
*
|
||||
* ORDER MATTERS (F6): once we own the claim, the terminal UPDATE happens FIRST;
|
||||
* only once it SUCCEEDS do we record the run as settled. If the UPDATE fails on
|
||||
* every bounded attempt we RESTORE the in-memory entry, leave the run UNsettled,
|
||||
* and emit an ERROR signal that the row is left non-terminal 'running' (which
|
||||
* would 409 every future turn in the chat until recovery). An in-process retry
|
||||
* by a LATER settle is only POSSIBLE, never guaranteed: it needs (a) the entry
|
||||
* to have been restored at the give-up path AND (b) a fresh settler to arrive
|
||||
* AFTER that restore. A concurrent settler that arrives DURING the retry window
|
||||
* — while the entry is deleted for backoff and not yet restored — is consumed at
|
||||
* the synchronous `active.delete` claim (it finds nothing to delete and returns
|
||||
* a no-op), so it does NOT become an in-process retrier. The NO-streamText path
|
||||
* (the turn threw before streamText was wired, so ONLY the safety-net ever
|
||||
* settles) likewise has no second in-process settler at all. The UNCONDITIONAL
|
||||
* backstop in every case is the boot sweep on the next restart (phase 1 has no
|
||||
* periodic in-process sweep); the retained entry is bounded (cleared on restart)
|
||||
* and harmless meanwhile.
|
||||
*
|
||||
* IDEMPOTENT on SUCCESS (#184 review): the terminal write happens AT MOST ONCE
|
||||
* per run. After a successful write the once-gate keys off {@link settled} (the
|
||||
* terminal row already written) so a settle arriving AFTER the entry was already
|
||||
* dropped-and-settled returns early; a settle racing the in-flight write is
|
||||
* stopped earlier still, by the `active.delete` claim. Either way a genuine
|
||||
* double-settle collapses to a single write and a late settle can never clobber
|
||||
* the real terminal status or double-write the row.
|
||||
*/
|
||||
async finalizeRun(
|
||||
runId: string,
|
||||
workspaceId: string,
|
||||
turnStatus: TurnTerminalStatus,
|
||||
error?: string,
|
||||
): Promise<void> {
|
||||
// ---- Atomic once-claim (synchronous; NO await before the gate closes) ----
|
||||
// Already terminally written -> idempotent no-op.
|
||||
if (this.settled.has(runId)) return;
|
||||
// Capture the entry BEFORE the delete so a total-failure path can restore it.
|
||||
const entry = this.active.get(runId);
|
||||
// SYNCHRONOUS check-and-clear: the FIRST caller deletes (claims) the entry;
|
||||
// any concurrent SECOND caller finds nothing to delete and returns HERE, in
|
||||
// the same tick, before any await — so it can never reach the UPDATE.
|
||||
if (!this.active.delete(runId)) return;
|
||||
|
||||
let lastError: unknown;
|
||||
for (
|
||||
let attempt = 1;
|
||||
attempt <= AiChatRunService.FINALIZE_MAX_ATTEMPTS;
|
||||
attempt++
|
||||
) {
|
||||
try {
|
||||
await this.runRepo.update(runId, workspaceId, {
|
||||
status: mapTurnStatusToRun(turnStatus),
|
||||
finishedAt: new Date(),
|
||||
error: error ?? null,
|
||||
});
|
||||
// Terminal write landed: arm the once-gate. The entry is already gone
|
||||
// (claimed above); we do NOT restore it. The slot is now free.
|
||||
this.settled.add(runId);
|
||||
return;
|
||||
} catch (err) {
|
||||
lastError = err;
|
||||
this.logger.warn(
|
||||
`Failed to finalize run ${runId} (attempt ${attempt}/${
|
||||
AiChatRunService.FINALIZE_MAX_ATTEMPTS
|
||||
}): ${err instanceof Error ? err.message : 'unknown error'}`,
|
||||
);
|
||||
if (attempt < AiChatRunService.FINALIZE_MAX_ATTEMPTS) {
|
||||
await this.delay(AiChatRunService.FINALIZE_RETRY_BASE_MS * attempt);
|
||||
}
|
||||
}
|
||||
}
|
||||
// Every attempt failed: this is a give-up, materially worse than a per-attempt
|
||||
// blip — the row is left NON-TERMINAL ('running'), so emit ONE explicit,
|
||||
// greppable ERROR so an operator can tell "survived a blip" from "gave up, run
|
||||
// held in memory until recovery" (the last warn alone says only "attempt 3/3").
|
||||
this.logger.error(
|
||||
`Run ${runId} (chat ${entry?.chatId ?? 'unknown'}) left NON-TERMINAL ` +
|
||||
`('running'): terminal write failed after ${
|
||||
AiChatRunService.FINALIZE_MAX_ATTEMPTS
|
||||
} attempts; entry retained in memory, recovery deferred to next settle / ` +
|
||||
`boot sweep`,
|
||||
lastError,
|
||||
);
|
||||
// RESTORE the claimed entry (and leave the run UNsettled) so a LATER settle
|
||||
// that arrives AFTER this restore MAY retry the terminal write — but that
|
||||
// in-process retry is NOT guaranteed (a concurrent settler caught in the retry
|
||||
// window above is consumed at the `active.delete` claim, and the no-streamText
|
||||
// path has no second settler at all). The UNCONDITIONAL backstop in every case
|
||||
// is the boot sweep on the next restart; the restored entry is bounded and
|
||||
// cleared on restart.
|
||||
if (entry) this.active.set(runId, entry);
|
||||
}
|
||||
|
||||
/** Small async backoff between terminal-write retries (F6). Isolated so it is
|
||||
* trivial to stub/fake-time in tests. */
|
||||
private delay(ms: number): Promise<void> {
|
||||
return new Promise((resolve) => setTimeout(resolve, ms));
|
||||
}
|
||||
|
||||
/**
|
||||
* Request an EXPLICIT stop of a run (the user pressed Stop). This is the ONLY
|
||||
* thing that aborts a run — distinct from a browser disconnect, which leaves
|
||||
* the run going. Aborts the in-process controller FIRST (the only thing that
|
||||
* actually stops the run, if this replica owns it), then makes a best-effort
|
||||
* attempt to stamp `stop_requested_at` — that audit write stamps only while the
|
||||
* row is active and may be skipped on a DB error or lost to the finalize race,
|
||||
* which is acceptable since the row still settles as 'aborted'. Returns true
|
||||
* when a stop took effect (row marked and/or controller aborted), false when
|
||||
* there was nothing active to stop.
|
||||
*/
|
||||
async requestStop(runId: string, workspaceId: string): Promise<boolean> {
|
||||
const entry = this.active.get(runId);
|
||||
if (entry) {
|
||||
// Abort the live turn FIRST -> streamText onAbort fires -> the partial is
|
||||
// persisted (#183) and finalizeRun settles the row as 'aborted'. This is
|
||||
// the ONLY thing that aborts a run, so it MUST NOT be hostage to the audit
|
||||
// write below: a transient failure on `markStopRequested` (pool exhaustion,
|
||||
// deadlock, dropped connection) must never leave the run executing despite
|
||||
// an explicit Stop. At worst only the `stop_requested_at` timestamp is lost.
|
||||
entry.controller.abort();
|
||||
}
|
||||
// Record `stop_requested_at` (best-effort). A transient DB failure here is
|
||||
// logged and treated as `marked = false`; the abort above already took
|
||||
// effect, so we never rethrow and skip stopping the run. Note: because
|
||||
// markStopRequested only stamps while the row is active, aborting first means
|
||||
// even a healthy write can lose the race against the resulting finalize and
|
||||
// skip the stamp — acceptable, as the row still settles as 'aborted' and only
|
||||
// this audit timestamp may be lost.
|
||||
let marked: unknown;
|
||||
try {
|
||||
marked = await this.runRepo.markStopRequested(runId, workspaceId);
|
||||
} catch (err) {
|
||||
marked = undefined;
|
||||
this.logger.warn(
|
||||
`requestStop: markStopRequested failed for run ${runId} ` +
|
||||
`(stop_requested_at not recorded); abort already issued: ` +
|
||||
`${err instanceof Error ? err.message : String(err)}`,
|
||||
);
|
||||
}
|
||||
return Boolean(marked) || Boolean(entry);
|
||||
}
|
||||
|
||||
/** Latest persisted run for a chat — the reconnect target (an in-flight or
|
||||
* finished run). Pure read-through to the repo. */
|
||||
getLatestForChat(
|
||||
chatId: string,
|
||||
workspaceId: string,
|
||||
): Promise<AiChatRun | undefined> {
|
||||
return this.runRepo.findLatestByChat(chatId, workspaceId);
|
||||
}
|
||||
|
||||
/** Fetch a run by id (workspace-scoped). Used to resolve + ownership-check an
|
||||
* explicit stop targeting a runId. */
|
||||
getRun(runId: string, workspaceId: string): Promise<AiChatRun | undefined> {
|
||||
return this.runRepo.findById(runId, workspaceId);
|
||||
}
|
||||
|
||||
/** The active run on a chat, if any (used to reject a concurrent start with a
|
||||
* clean 409 before committing to the stream). */
|
||||
getActiveForChat(
|
||||
chatId: string,
|
||||
workspaceId: string,
|
||||
): Promise<AiChatRun | undefined> {
|
||||
return this.runRepo.findActiveByChat(chatId, workspaceId);
|
||||
}
|
||||
|
||||
/** Test/diagnostic seam: whether this replica is holding a live controller for
|
||||
* the run. */
|
||||
isLocallyActive(runId: string): boolean {
|
||||
return this.active.has(runId);
|
||||
}
|
||||
}
|
||||
@@ -25,6 +25,7 @@ describe('AiChatController.boundChat', () => {
|
||||
};
|
||||
const controller = new AiChatController(
|
||||
{} as never,
|
||||
{} as never, // aiChatRunService
|
||||
aiChatRepo as never,
|
||||
{} as never,
|
||||
{} as never,
|
||||
|
||||
@@ -53,6 +53,7 @@ describe('AiChatController.export', () => {
|
||||
};
|
||||
const controller = new AiChatController(
|
||||
{} as never,
|
||||
{} as never, // aiChatRunService
|
||||
aiChatRepo as never,
|
||||
aiChatMessageRepo as never,
|
||||
{} as never,
|
||||
|
||||
@@ -0,0 +1,164 @@
|
||||
import { BadRequestException, ForbiddenException } from '@nestjs/common';
|
||||
import { AiChatController } from './ai-chat.controller';
|
||||
import type { User, Workspace } from '@docmost/db/types/entity.types';
|
||||
|
||||
/**
|
||||
* Wiring spec for the #184 run-reconnect / run-stop endpoints
|
||||
* (`POST /ai-chat/run` and `POST /ai-chat/stop`). Both are OWNER-gated via
|
||||
* assertOwnedChat (the requesting user must own the chat) and NOT flag-gated.
|
||||
* Exercised with hand-rolled mocks — no Nest graph, no DB. The controller's
|
||||
* constructor order is (aiChatService, aiChatRunService, aiChatRepo,
|
||||
* aiChatMessageRepo, aiTranscription).
|
||||
*/
|
||||
describe('AiChatController run endpoints (#184)', () => {
|
||||
const user = { id: 'u1' } as User;
|
||||
const workspace = { id: 'ws1' } as Workspace;
|
||||
|
||||
function makeController(opts: {
|
||||
chat?: unknown; // what aiChatRepo.findById returns (owner-gate)
|
||||
run?: unknown; // getLatestForChat / getRun result
|
||||
activeRun?: unknown; // getActiveForChat result
|
||||
message?: unknown; // aiChatMessageRepo.findById result
|
||||
stopped?: boolean; // requestStop result
|
||||
}) {
|
||||
const aiChatRunService = {
|
||||
getLatestForChat: jest.fn().mockResolvedValue(opts.run),
|
||||
getRun: jest.fn().mockResolvedValue(opts.run),
|
||||
getActiveForChat: jest.fn().mockResolvedValue(opts.activeRun),
|
||||
requestStop: jest.fn().mockResolvedValue(opts.stopped ?? false),
|
||||
};
|
||||
const aiChatRepo = {
|
||||
findById: jest.fn().mockResolvedValue(opts.chat),
|
||||
};
|
||||
const aiChatMessageRepo = {
|
||||
findById: jest.fn().mockResolvedValue(opts.message),
|
||||
};
|
||||
const controller = new AiChatController(
|
||||
{} as never, // aiChatService
|
||||
aiChatRunService as never,
|
||||
aiChatRepo as never,
|
||||
aiChatMessageRepo as never,
|
||||
{} as never, // aiTranscription
|
||||
{} as never, // pageRepo
|
||||
);
|
||||
return { controller, aiChatRunService, aiChatRepo, aiChatMessageRepo };
|
||||
}
|
||||
|
||||
describe('POST /ai-chat/run (getRun)', () => {
|
||||
it('owner-gates: a chat the user does not own throws ForbiddenException', async () => {
|
||||
const { controller, aiChatRunService } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'someone-else' },
|
||||
});
|
||||
await expect(
|
||||
controller.getRun({ chatId: 'c1' }, user, workspace),
|
||||
).rejects.toBeInstanceOf(ForbiddenException);
|
||||
// It must NOT reach the run lookup once the owner-gate fails.
|
||||
expect(aiChatRunService.getLatestForChat).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('returns { run: null, message: null } when the chat has never had a run', async () => {
|
||||
const { controller, aiChatRunService } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'u1' },
|
||||
run: undefined,
|
||||
});
|
||||
const res = await controller.getRun({ chatId: 'c1' }, user, workspace);
|
||||
expect(res).toEqual({ run: null, message: null });
|
||||
expect(aiChatRunService.getLatestForChat).toHaveBeenCalledWith(
|
||||
'c1',
|
||||
'ws1',
|
||||
);
|
||||
});
|
||||
|
||||
it('returns the run and its projected assistant message', async () => {
|
||||
const run = { id: 'run-1', chatId: 'c1', assistantMessageId: 'm1' };
|
||||
const message = { id: 'm1', role: 'assistant' };
|
||||
const { controller, aiChatMessageRepo } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'u1' },
|
||||
run,
|
||||
message,
|
||||
});
|
||||
const res = await controller.getRun({ chatId: 'c1' }, user, workspace);
|
||||
expect(res).toEqual({ run, message });
|
||||
expect(aiChatMessageRepo.findById).toHaveBeenCalledWith('m1', 'ws1');
|
||||
});
|
||||
|
||||
it('returns message: null when the run has no linked assistant message', async () => {
|
||||
const run = { id: 'run-1', chatId: 'c1', assistantMessageId: null };
|
||||
const { controller, aiChatMessageRepo } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'u1' },
|
||||
run,
|
||||
});
|
||||
const res = await controller.getRun({ chatId: 'c1' }, user, workspace);
|
||||
expect(res).toEqual({ run, message: null });
|
||||
expect(aiChatMessageRepo.findById).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('POST /ai-chat/stop (stopRun)', () => {
|
||||
it('throws BadRequestException when neither runId nor chatId is given', async () => {
|
||||
const { controller } = makeController({});
|
||||
await expect(
|
||||
controller.stopRun({}, user, workspace),
|
||||
).rejects.toBeInstanceOf(BadRequestException);
|
||||
});
|
||||
|
||||
it('stops by runId: owner-gates via the run’s chat, then requests the stop', async () => {
|
||||
const { controller, aiChatRunService, aiChatRepo } = makeController({
|
||||
run: { id: 'run-1', chatId: 'c1' },
|
||||
chat: { id: 'c1', creatorId: 'u1' },
|
||||
stopped: true,
|
||||
});
|
||||
const res = await controller.stopRun({ runId: 'run-1' }, user, workspace);
|
||||
expect(res).toEqual({ stopped: true });
|
||||
expect(aiChatRunService.getRun).toHaveBeenCalledWith('run-1', 'ws1');
|
||||
expect(aiChatRepo.findById).toHaveBeenCalledWith('c1', 'ws1');
|
||||
expect(aiChatRunService.requestStop).toHaveBeenCalledWith('run-1', 'ws1');
|
||||
});
|
||||
|
||||
it('stops by runId: a foreign run’s chat throws ForbiddenException (no stop)', async () => {
|
||||
const { controller, aiChatRunService } = makeController({
|
||||
run: { id: 'run-1', chatId: 'c1' },
|
||||
chat: { id: 'c1', creatorId: 'someone-else' },
|
||||
});
|
||||
await expect(
|
||||
controller.stopRun({ runId: 'run-1' }, user, workspace),
|
||||
).rejects.toBeInstanceOf(ForbiddenException);
|
||||
expect(aiChatRunService.requestStop).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('stops by runId: an unknown run reports { stopped: false }', async () => {
|
||||
const { controller, aiChatRunService } = makeController({
|
||||
run: undefined,
|
||||
});
|
||||
const res = await controller.stopRun({ runId: 'gone' }, user, workspace);
|
||||
expect(res).toEqual({ stopped: false });
|
||||
expect(aiChatRunService.requestStop).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('stops by chatId: owner-gates, resolves the active run, requests the stop', async () => {
|
||||
const { controller, aiChatRunService, aiChatRepo } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'u1' },
|
||||
activeRun: { id: 'run-9' },
|
||||
stopped: true,
|
||||
});
|
||||
const res = await controller.stopRun({ chatId: 'c1' }, user, workspace);
|
||||
expect(res).toEqual({ stopped: true });
|
||||
expect(aiChatRepo.findById).toHaveBeenCalledWith('c1', 'ws1');
|
||||
expect(aiChatRunService.getActiveForChat).toHaveBeenCalledWith(
|
||||
'c1',
|
||||
'ws1',
|
||||
);
|
||||
expect(aiChatRunService.requestStop).toHaveBeenCalledWith('run-9', 'ws1');
|
||||
});
|
||||
|
||||
it('stops by chatId: reports { stopped: false } when no run is active', async () => {
|
||||
const { controller, aiChatRunService } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'u1' },
|
||||
activeRun: undefined,
|
||||
});
|
||||
const res = await controller.stopRun({ chatId: 'c1' }, user, workspace);
|
||||
expect(res).toEqual({ stopped: false });
|
||||
expect(aiChatRunService.requestStop).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -1,6 +1,7 @@
|
||||
import {
|
||||
BadRequestException,
|
||||
Body,
|
||||
ConflictException,
|
||||
Controller,
|
||||
ForbiddenException,
|
||||
HttpCode,
|
||||
@@ -20,7 +21,13 @@ import { JwtAuthGuard } from '../../common/guards/jwt-auth.guard';
|
||||
import { AuthUser } from '../../common/decorators/auth-user.decorator';
|
||||
import { AuthWorkspace } from '../../common/decorators/auth-workspace.decorator';
|
||||
import { SkipTransform } from '../../common/decorators/skip-transform.decorator';
|
||||
import { AiChat, User, Workspace } from '@docmost/db/types/entity.types';
|
||||
import {
|
||||
AiChat,
|
||||
AiChatMessage,
|
||||
AiChatRun,
|
||||
User,
|
||||
Workspace,
|
||||
} from '@docmost/db/types/entity.types';
|
||||
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';
|
||||
@@ -28,7 +35,12 @@ 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';
|
||||
import { FileInterceptor } from '../../common/interceptors/file.interceptor';
|
||||
import { AiChatService, AiChatStreamBody } from './ai-chat.service';
|
||||
import {
|
||||
AiChatRunHooks,
|
||||
AiChatService,
|
||||
AiChatStreamBody,
|
||||
} from './ai-chat.service';
|
||||
import { AiChatRunService } from './ai-chat-run.service';
|
||||
import { AiTranscriptionService } from './ai-transcription.service';
|
||||
import {
|
||||
BoundChatDto,
|
||||
@@ -36,7 +48,9 @@ import {
|
||||
ExportChatDto,
|
||||
GeneratePageTitleDto,
|
||||
GetChatMessagesDto,
|
||||
GetRunDto,
|
||||
RenameChatDto,
|
||||
StopRunDto,
|
||||
} from './dto/ai-chat.dto';
|
||||
import { describeProviderError } from '../../integrations/ai/ai-error.util';
|
||||
import { buildChatMarkdown } from './chat-markdown.util';
|
||||
@@ -53,6 +67,7 @@ export class AiChatController {
|
||||
|
||||
constructor(
|
||||
private readonly aiChatService: AiChatService,
|
||||
private readonly aiChatRunService: AiChatRunService,
|
||||
private readonly aiChatRepo: AiChatRepo,
|
||||
private readonly aiChatMessageRepo: AiChatMessageRepo,
|
||||
private readonly aiTranscription: AiTranscriptionService,
|
||||
@@ -149,6 +164,75 @@ export class AiChatController {
|
||||
return { markdown };
|
||||
}
|
||||
|
||||
/**
|
||||
* Reconnect to the latest run of a chat (#184 phase 1). Returns the run's
|
||||
* persisted lifecycle state ({ status, error, stepCount, timings, ... }) plus
|
||||
* the assistant message it projects (the partial/final output) — the DB is the
|
||||
* source of truth, so this works for an in-flight run (the browser dropped, the
|
||||
* run kept going) and a finished one alike. Owner-gated via assertOwnedChat.
|
||||
* `{ run: null }` when the chat has never had a run.
|
||||
*/
|
||||
@HttpCode(HttpStatus.OK)
|
||||
@Post('run')
|
||||
async getRun(
|
||||
@Body() dto: GetRunDto,
|
||||
@AuthUser() user: User,
|
||||
@AuthWorkspace() workspace: Workspace,
|
||||
): Promise<{ run: AiChatRun | null; message: AiChatMessage | null }> {
|
||||
await this.assertOwnedChat(dto.chatId, user, workspace);
|
||||
const run = await this.aiChatRunService.getLatestForChat(
|
||||
dto.chatId,
|
||||
workspace.id,
|
||||
);
|
||||
if (!run) return { run: null, message: null };
|
||||
const message = run.assistantMessageId
|
||||
? await this.aiChatMessageRepo.findById(
|
||||
run.assistantMessageId,
|
||||
workspace.id,
|
||||
)
|
||||
: undefined;
|
||||
return { run, message: message ?? null };
|
||||
}
|
||||
|
||||
/**
|
||||
* Explicitly STOP an agent run (#184 phase 1) — the user pressed Stop. This is
|
||||
* the ONLY thing that ends a detached run; a browser disconnect deliberately
|
||||
* does not. Target by `runId` (from the streamed start metadata) or by `chatId`
|
||||
* (stop whatever run is active on it). Owner-gated. Returns
|
||||
* `{ stopped }` — false when there was nothing active to stop.
|
||||
*/
|
||||
@HttpCode(HttpStatus.OK)
|
||||
@Post('stop')
|
||||
async stopRun(
|
||||
@Body() dto: StopRunDto,
|
||||
@AuthUser() user: User,
|
||||
@AuthWorkspace() workspace: Workspace,
|
||||
): Promise<{ stopped: boolean }> {
|
||||
let runId = dto.runId;
|
||||
if (!runId && !dto.chatId) {
|
||||
throw new BadRequestException('runId or chatId is required');
|
||||
}
|
||||
if (runId) {
|
||||
// Resolve the run to its chat and owner-gate via that chat.
|
||||
const run = await this.aiChatRunService.getRun(runId, workspace.id);
|
||||
if (!run) return { stopped: false };
|
||||
await this.assertOwnedChat(run.chatId, user, workspace);
|
||||
} else {
|
||||
await this.assertOwnedChat(dto.chatId!, user, workspace);
|
||||
const active = await this.aiChatRunService.getActiveForChat(
|
||||
dto.chatId!,
|
||||
workspace.id,
|
||||
);
|
||||
if (!active) return { stopped: false };
|
||||
runId = active.id;
|
||||
}
|
||||
const stopped = await this.aiChatRunService.requestStop(
|
||||
runId,
|
||||
workspace.id,
|
||||
);
|
||||
return { stopped };
|
||||
}
|
||||
|
||||
/** Rename a chat. */
|
||||
@HttpCode(HttpStatus.OK)
|
||||
@Post('rename')
|
||||
@@ -200,11 +284,20 @@ export class AiChatController {
|
||||
@AuthWorkspace() workspace: Workspace,
|
||||
): Promise<void> {
|
||||
// A7 gate: the workspace must have AI chat explicitly enabled.
|
||||
const settings = (workspace.settings ?? {}) as { ai?: { chat?: boolean } };
|
||||
const settings = (workspace.settings ?? {}) as {
|
||||
ai?: { chat?: boolean; autonomousRuns?: boolean };
|
||||
};
|
||||
if (settings.ai?.chat !== true) {
|
||||
throw new ForbiddenException('AI chat is disabled');
|
||||
}
|
||||
|
||||
// #184 phase 1 flag: when ON, the turn becomes a detached, durable RUN — its
|
||||
// lifecycle is tracked in ai_chat_runs, a browser disconnect no longer aborts
|
||||
// it, and only an explicit /ai-chat/stop ends it. When OFF (the default) the
|
||||
// turn is socket-bound exactly as before, so existing deployments are
|
||||
// unaffected.
|
||||
const autonomousRuns = settings.ai?.autonomousRuns === true;
|
||||
|
||||
const sessionId = (req.raw as { sessionId?: string }).sessionId;
|
||||
if (!sessionId) {
|
||||
// The chat requires an interactive session to mint loopback tokens
|
||||
@@ -228,6 +321,58 @@ export class AiChatController {
|
||||
// HttpException) instead of breaking mid-stream.
|
||||
const model = await this.aiChatService.getChatModel(workspace.id, role);
|
||||
|
||||
// #184: one active run per chat. For an EXISTING chat reject a concurrent
|
||||
// start with a clean 409 BEFORE hijack (the common double-submit / second-tab
|
||||
// case), so the user gets JSON, not a mid-stream error. A brand-new chat
|
||||
// (no chatId) cannot have a prior run, and the DB partial unique index is the
|
||||
// backstop against any race that slips past this check.
|
||||
if (autonomousRuns && body.chatId) {
|
||||
const active = await this.aiChatRunService.getActiveForChat(
|
||||
body.chatId,
|
||||
workspace.id,
|
||||
);
|
||||
if (active) {
|
||||
throw new ConflictException({
|
||||
message: 'An agent run is already in progress for this chat',
|
||||
code: 'A_RUN_ALREADY_ACTIVE',
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Run-lifecycle hooks (#184), only when the flag is on. They wrap the turn in
|
||||
// a durable run whose abort is governed by the run (explicit stop), persist
|
||||
// its progress, and settle its terminal status — see AiChatRunService.
|
||||
const runHooks: AiChatRunHooks | undefined = autonomousRuns
|
||||
? {
|
||||
begin: (chatId) =>
|
||||
this.aiChatRunService.beginRun({
|
||||
chatId,
|
||||
workspaceId: workspace.id,
|
||||
userId: user.id,
|
||||
trigger: 'user',
|
||||
}),
|
||||
onAssistantSeeded: (runId, messageId) =>
|
||||
this.aiChatRunService.linkAssistantMessage(
|
||||
runId,
|
||||
workspace.id,
|
||||
messageId,
|
||||
),
|
||||
onStep: (runId, stepCount) =>
|
||||
void this.aiChatRunService.recordStep(
|
||||
runId,
|
||||
workspace.id,
|
||||
stepCount,
|
||||
),
|
||||
onSettled: (runId, status, error) =>
|
||||
this.aiChatRunService.finalizeRun(
|
||||
runId,
|
||||
workspace.id,
|
||||
status,
|
||||
error,
|
||||
),
|
||||
}
|
||||
: undefined;
|
||||
|
||||
// Abort the agent loop when the client disconnects. `close` also fires on
|
||||
// normal completion, so only abort when the response has not finished
|
||||
// writing (a genuine disconnect). `once` fires at most once and self-removes;
|
||||
@@ -242,18 +387,44 @@ export class AiChatController {
|
||||
// A genuine disconnect leaves the response unfinished (unlike a normal
|
||||
// completion, which also fires `close`). Such a drop — e.g. a reverse
|
||||
// proxy cutting the SSE mid-answer — is otherwise invisible server-side,
|
||||
// so log it here before aborting the agent loop.
|
||||
// so log it here.
|
||||
if (!res.raw.writableEnded) {
|
||||
this.logger.warn(
|
||||
`AI chat stream: client disconnected before completion; aborting turn ` +
|
||||
`(elapsed=${Date.now() - reqStartedAt}ms since request received)`,
|
||||
);
|
||||
controller.abort();
|
||||
if (autonomousRuns) {
|
||||
// #184: the turn is a DETACHED run. A disconnect must NOT abort it —
|
||||
// the run keeps executing and persisting server-side; the client
|
||||
// reconnects via /ai-chat/run (or re-stops via /ai-chat/stop). Log only.
|
||||
this.logger.log(
|
||||
`AI chat stream: client disconnected; run continues server-side ` +
|
||||
`(elapsed=${Date.now() - reqStartedAt}ms since request received)`,
|
||||
);
|
||||
} else {
|
||||
this.logger.warn(
|
||||
`AI chat stream: client disconnected before completion; aborting turn ` +
|
||||
`(elapsed=${Date.now() - reqStartedAt}ms since request received)`,
|
||||
);
|
||||
controller.abort();
|
||||
}
|
||||
}
|
||||
};
|
||||
req.raw.once('close', onClose);
|
||||
res.raw.once('finish', () => req.raw.off('close', onClose));
|
||||
|
||||
// #184: in detached mode the turn is NOT aborted on disconnect, so the SDK's
|
||||
// pipe keeps writing to a socket the client may have dropped — for the rest of
|
||||
// the (continuing) run. A write to the dead socket can emit an 'error' on the
|
||||
// raw response; without a listener that surfaces as an unhandled error event.
|
||||
// Swallow it (the run continues server-side regardless). Legacy mode aborts on
|
||||
// disconnect, so it does not need this and keeps its exact prior behavior.
|
||||
if (autonomousRuns) {
|
||||
res.raw.on('error', (err) => {
|
||||
this.logger.debug(
|
||||
`AI chat detached stream: post-disconnect socket error swallowed: ${
|
||||
err instanceof Error ? err.message : String(err)
|
||||
}`,
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
// Commit to streaming: hijack so Fastify stops managing the response and
|
||||
// the AI SDK can write the UI-message stream directly to the Node socket.
|
||||
res.hijack();
|
||||
@@ -268,15 +439,32 @@ export class AiChatController {
|
||||
signal: controller.signal,
|
||||
model,
|
||||
role,
|
||||
// #184: present only when the flag is on; wraps the turn in a durable run.
|
||||
runHooks,
|
||||
});
|
||||
} catch (err) {
|
||||
// Any failure AFTER hijack can no longer send a clean JSON error, so emit
|
||||
// a minimal error on the raw socket if nothing has been written yet.
|
||||
this.logger.error('AI chat stream failed', err as Error);
|
||||
// Any failure AFTER hijack can no longer go through Nest's exception
|
||||
// filter, so emit the error on the raw socket if nothing has been written
|
||||
// yet. The lost-the-race 409 (RunAlreadyActiveError -> ConflictException)
|
||||
// is raised by stream() BEFORE it writes a byte, so headers are still
|
||||
// unsent here: honor the HttpException's real status + body (a clean 409),
|
||||
// not a blanket 500. Everything else stays a 500.
|
||||
const isHttp = err instanceof HttpException;
|
||||
if (!isHttp) {
|
||||
this.logger.error('AI chat stream failed', err as Error);
|
||||
}
|
||||
if (!res.raw.headersSent) {
|
||||
res.raw.statusCode = 500;
|
||||
const status = isHttp ? err.getStatus() : 500;
|
||||
const payload = isHttp
|
||||
? err.getResponse()
|
||||
: { error: 'Internal server error' };
|
||||
res.raw.statusCode = status;
|
||||
res.raw.setHeader('Content-Type', 'application/json');
|
||||
res.raw.end(JSON.stringify({ error: 'Internal server error' }));
|
||||
res.raw.end(
|
||||
JSON.stringify(
|
||||
typeof payload === 'string' ? { message: payload } : payload,
|
||||
),
|
||||
);
|
||||
} else if (!res.raw.writableEnded) {
|
||||
res.raw.end();
|
||||
}
|
||||
|
||||
@@ -57,6 +57,7 @@ describe('AiChatController.generatePageTitle', () => {
|
||||
const aiChatService = { generatePageTitle: generate };
|
||||
const controller = new AiChatController(
|
||||
aiChatService as never,
|
||||
{} as never, // aiChatRunService
|
||||
{} as never,
|
||||
{} as never,
|
||||
{} as never,
|
||||
|
||||
@@ -3,6 +3,7 @@ import { AiModule } from '../../integrations/ai/ai.module';
|
||||
import { TokenModule } from '../auth/token.module';
|
||||
import { AiChatController } from './ai-chat.controller';
|
||||
import { AiChatService } from './ai-chat.service';
|
||||
import { AiChatRunService } from './ai-chat-run.service';
|
||||
import { AiTranscriptionService } from './ai-transcription.service';
|
||||
import { AiChatToolsService } from './tools/ai-chat-tools.service';
|
||||
import { EmbeddingModule } from './embedding/embedding.module';
|
||||
@@ -42,6 +43,7 @@ import { PublicShareChatToolsService } from './tools/public-share-chat-tools.ser
|
||||
controllers: [AiChatController, PublicShareChatController],
|
||||
providers: [
|
||||
AiChatService,
|
||||
AiChatRunService,
|
||||
AiTranscriptionService,
|
||||
AiChatToolsService,
|
||||
PublicShareChatService,
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
import { Logger } from '@nestjs/common';
|
||||
import { AiChatService } from './ai-chat.service';
|
||||
import { AiChatService, AiChatRunHooks } from './ai-chat.service';
|
||||
import { AiChatRunService } from './ai-chat-run.service';
|
||||
import type { User, Workspace } from '@docmost/db/types/entity.types';
|
||||
|
||||
/**
|
||||
* Lifecycle unit tests for AiChatService.onModuleInit (#183 crash-recovery
|
||||
@@ -61,3 +63,99 @@ describe('AiChatService.onModuleInit (startup sweep)', () => {
|
||||
expect(String(warnSpy.mock.calls[0][0])).toContain('db unavailable');
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* #184 CRITICAL run-lifecycle safety net (review fix). A transient failure
|
||||
* AFTER a successful beginRun but BEFORE streamText's terminal callbacks own the
|
||||
* lifecycle must STILL settle the run — otherwise the run row is stuck 'running'
|
||||
* forever (sweepRunning only runs at startup) and the partial unique index + the
|
||||
* controller pre-check 409 every future turn in that chat until a restart. Here
|
||||
* we model the very first bare await after beginRun (the user-message insert)
|
||||
* throwing, wiring the run hooks to a REAL AiChatRunService (mock repo) exactly
|
||||
* as the controller does, and assert the run is settled to 'error' and its
|
||||
* in-memory entry dropped (so a follow-up turn would NOT be 409'd).
|
||||
*/
|
||||
describe('AiChatService.stream run-lifecycle safety net (#184)', () => {
|
||||
const user = { id: 'u1' } as User;
|
||||
const workspace = { id: 'ws1' } as Workspace;
|
||||
|
||||
afterEach(() => jest.restoreAllMocks());
|
||||
|
||||
it('an exception after beginRun settles the run to error and drops the in-memory entry', async () => {
|
||||
jest.spyOn(Logger.prototype, 'error').mockImplementation(() => undefined);
|
||||
|
||||
// Real run service over a mock repo, so finalizeRun's in-memory bookkeeping
|
||||
// (active.delete) is exercised for real.
|
||||
const runRepo = {
|
||||
insert: jest.fn().mockResolvedValue({ id: 'run-1', status: 'running' }),
|
||||
update: jest.fn().mockResolvedValue({ id: 'run-1' }),
|
||||
};
|
||||
const runService = new AiChatRunService(runRepo as never, { isCloud: () => false } as never);
|
||||
|
||||
// The user-message insert (the first bare await after beginRun) throws.
|
||||
const aiChatMessageRepo = {
|
||||
insert: jest.fn().mockRejectedValue(new Error('insert boom')),
|
||||
};
|
||||
const aiChatRepo = {
|
||||
// Existing chat -> chatId stays, no new-chat insert path.
|
||||
findById: jest.fn().mockResolvedValue({ id: 'chat-1', creatorId: 'u1' }),
|
||||
};
|
||||
|
||||
const service = new AiChatService(
|
||||
{} as never, // ai
|
||||
aiChatRepo as never,
|
||||
aiChatMessageRepo as never,
|
||||
{} as never, // aiChatPageSnapshotRepo
|
||||
{} as never, // aiSettings
|
||||
{} as never, // tools
|
||||
{} as never, // mcpClients
|
||||
{} as never, // aiAgentRoleRepo
|
||||
{} as never, // pageRepo
|
||||
{} as never, // pageAccess
|
||||
{} as never, // environment
|
||||
);
|
||||
|
||||
const runHooks: AiChatRunHooks = {
|
||||
begin: (chatId) =>
|
||||
runService.beginRun({
|
||||
chatId,
|
||||
workspaceId: workspace.id,
|
||||
userId: user.id,
|
||||
trigger: 'user',
|
||||
}),
|
||||
onSettled: (runId, status, error) =>
|
||||
runService.finalizeRun(runId, workspace.id, status, error),
|
||||
};
|
||||
|
||||
await expect(
|
||||
service.stream({
|
||||
user,
|
||||
workspace,
|
||||
sessionId: 'sess',
|
||||
body: {
|
||||
chatId: 'chat-1',
|
||||
messages: [
|
||||
{ id: 'm', role: 'user', parts: [{ type: 'text', text: 'hi' }] },
|
||||
],
|
||||
},
|
||||
res: {} as never,
|
||||
signal: new AbortController().signal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
runHooks,
|
||||
}),
|
||||
).rejects.toThrow('insert boom');
|
||||
|
||||
// The run was begun...
|
||||
expect(runRepo.insert).toHaveBeenCalledTimes(1);
|
||||
// ...then settled to a terminal FAILED status by the safety net...
|
||||
expect(runRepo.update).toHaveBeenCalledTimes(1);
|
||||
expect(runRepo.update).toHaveBeenCalledWith(
|
||||
'run-1',
|
||||
'ws1',
|
||||
expect.objectContaining({ status: 'failed' }),
|
||||
);
|
||||
// ...and the in-memory entry is gone, so a follow-up turn is NOT 409'd.
|
||||
expect(runService.isLocallyActive('run-1')).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -0,0 +1,489 @@
|
||||
import { ConflictException, Logger } from '@nestjs/common';
|
||||
|
||||
// Mock the AI SDK so we can PROVE no provider call is made for the turn we are
|
||||
// about to reject. The race rejection happens at runHooks.begin(), long before
|
||||
// any streamText/generateText, so these never resolve a real model.
|
||||
jest.mock('ai', () => ({
|
||||
streamText: jest.fn(),
|
||||
generateText: jest.fn(),
|
||||
convertToModelMessages: jest.fn(() => []),
|
||||
stepCountIs: jest.fn(() => () => false),
|
||||
}));
|
||||
|
||||
import { streamText, generateText } from 'ai';
|
||||
import { AiChatService } from './ai-chat.service';
|
||||
import { RunAlreadyActiveError } from './ai-chat-run.service';
|
||||
|
||||
/**
|
||||
* Race-closure coverage for the "one active run per chat" guard (#184).
|
||||
*
|
||||
* THE BUG: two simultaneous POST /ai-chat/stream on the same chat both pass the
|
||||
* controller's cheap pre-check (TOCTOU), so the loser's run-row INSERT hits the
|
||||
* partial unique index. Previously that 23505 was SWALLOWED and the second turn
|
||||
* streamed UNTRACKED (no runId, not stoppable). THE FIX: beginRun surfaces a
|
||||
* RunAlreadyActiveError and stream() turns it into a 409 BEFORE any AI call —
|
||||
* the second turn never runs.
|
||||
*/
|
||||
describe('AiChatService.stream — concurrent-run race rejection (#184)', () => {
|
||||
const streamTextMock = streamText as unknown as jest.Mock;
|
||||
const generateTextMock = generateText as unknown as jest.Mock;
|
||||
|
||||
beforeEach(() => {
|
||||
streamTextMock.mockReset();
|
||||
generateTextMock.mockReset();
|
||||
});
|
||||
|
||||
// Minimal service whose only reachable deps before begin() are aiChatRepo
|
||||
// (resolve the existing chat) — everything past begin must remain untouched.
|
||||
function makeService(beginImpl: () => Promise<unknown>) {
|
||||
const aiChatMessageRepo = { insert: jest.fn() };
|
||||
const aiChatRepo = {
|
||||
// An existing chat: stream keeps the supplied chatId and skips creation.
|
||||
findById: jest.fn(async () => ({ id: 'chat-1', workspaceId: 'ws-1' })),
|
||||
insert: jest.fn(),
|
||||
};
|
||||
const svc = new AiChatService(
|
||||
{} as never, // ai
|
||||
aiChatRepo as never,
|
||||
aiChatMessageRepo as never,
|
||||
{} as never, // aiChatPageSnapshotRepo
|
||||
{} as never, // aiSettings
|
||||
{} as never, // tools
|
||||
{} as never, // mcpClients
|
||||
{} as never, // aiAgentRoleRepo
|
||||
{} as never, // pageRepo
|
||||
{} as never, // pageAccess
|
||||
{ isAiChatDeferredToolsEnabled: () => false } as never, // environment
|
||||
);
|
||||
const begin = jest.fn(beginImpl);
|
||||
return { svc, begin, aiChatRepo, aiChatMessageRepo };
|
||||
}
|
||||
|
||||
const baseArgs = (begin: jest.Mock) => ({
|
||||
user: { id: 'user-1' } as never,
|
||||
workspace: { id: 'ws-1' } as never,
|
||||
sessionId: 'sess-1',
|
||||
body: { chatId: 'chat-1', messages: [] } as never,
|
||||
res: { raw: {} } as never,
|
||||
signal: new AbortController().signal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
runHooks: {
|
||||
begin,
|
||||
onAssistantSeeded: jest.fn(),
|
||||
onStep: jest.fn(),
|
||||
onSettled: jest.fn(),
|
||||
} as never,
|
||||
});
|
||||
|
||||
it('rejects the racer with a 409 ConflictException BEFORE any AI call, and never persists an untracked turn', async () => {
|
||||
// begin loses the unique-index race -> RunAlreadyActiveError.
|
||||
const { svc, begin, aiChatMessageRepo } = makeService(() => {
|
||||
throw new RunAlreadyActiveError('chat-1');
|
||||
});
|
||||
|
||||
const promise = svc.stream(baseArgs(begin));
|
||||
|
||||
await expect(promise).rejects.toBeInstanceOf(ConflictException);
|
||||
await promise.catch((err: ConflictException) => {
|
||||
expect(err.getStatus()).toBe(409);
|
||||
expect((err.getResponse() as { code?: string }).code).toBe(
|
||||
'A_RUN_ALREADY_ACTIVE',
|
||||
);
|
||||
});
|
||||
|
||||
// The decisive assertions: the rejected racer spent NO tokens and left NO
|
||||
// untracked turn behind.
|
||||
expect(begin).toHaveBeenCalledTimes(1);
|
||||
expect(streamTextMock).not.toHaveBeenCalled();
|
||||
expect(generateTextMock).not.toHaveBeenCalled();
|
||||
expect(aiChatMessageRepo.insert).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* F3 — the LOAD-BEARING run-detach wiring: `effectiveSignal = handle.signal`
|
||||
* after runHooks.begin, then `abortSignal: effectiveSignal` passed to streamText.
|
||||
* That single line is what makes a run survive a browser disconnect (the agent
|
||||
* loop's abort is governed by the RUN's signal, not the socket): a regression to
|
||||
* the socket-bound signal would still pass every other test green while silently
|
||||
* breaking Stop + durability. These two tests pin the exact signal streamText
|
||||
* consumes on both paths.
|
||||
*/
|
||||
describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
|
||||
const streamTextMock = streamText as unknown as jest.Mock;
|
||||
|
||||
// A streamText result stub: the post-call drain + pipe are no-ops here; we only
|
||||
// care WHICH abortSignal streamText was handed.
|
||||
function makeStreamResult() {
|
||||
return {
|
||||
consumeStream: jest.fn(),
|
||||
pipeUIMessageStreamToResponse: jest.fn(),
|
||||
};
|
||||
}
|
||||
|
||||
// A raw-response stub sufficient for the post-streamText wiring
|
||||
// (stripStreamingHopByHopHeaders binds writeHead; startSseHeartbeat registers
|
||||
// close/finish listeners; flushHeaders is belt-and-braces).
|
||||
function makeRes() {
|
||||
return {
|
||||
raw: {
|
||||
writeHead: jest.fn(),
|
||||
write: jest.fn(),
|
||||
once: jest.fn(),
|
||||
on: jest.fn(),
|
||||
flushHeaders: jest.fn(),
|
||||
writableEnded: false,
|
||||
destroyed: false,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
// Wire only the deps reached on the way to streamText: resolve the existing
|
||||
// chat, persist the user + seed the assistant row, load (empty) history, the
|
||||
// admin settings, an empty external toolset + Docmost toolset.
|
||||
function makeService() {
|
||||
const aiChatRepo = {
|
||||
findById: jest.fn(async () => ({ id: 'chat-1', workspaceId: 'ws-1' })),
|
||||
insert: jest.fn(),
|
||||
};
|
||||
const aiChatMessageRepo = {
|
||||
insert: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
findAllByChat: jest.fn(async () => []),
|
||||
update: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
};
|
||||
const aiSettings = { resolve: jest.fn(async () => ({})) };
|
||||
const tools = { forUser: jest.fn(async () => ({})) };
|
||||
const mcpClients = {
|
||||
toolsFor: jest.fn(async () => ({
|
||||
tools: {},
|
||||
clients: [],
|
||||
outcomes: [],
|
||||
instructions: [],
|
||||
})),
|
||||
};
|
||||
const svc = new AiChatService(
|
||||
{} as never, // ai
|
||||
aiChatRepo as never,
|
||||
aiChatMessageRepo as never,
|
||||
{} as never, // aiChatPageSnapshotRepo
|
||||
aiSettings as never,
|
||||
tools as never,
|
||||
mcpClients as never,
|
||||
{} as never, // aiAgentRoleRepo
|
||||
{} as never, // pageRepo (openPage undefined -> never touched)
|
||||
{} as never, // pageAccess
|
||||
{ isAiChatDeferredToolsEnabled: () => false } as never, // environment
|
||||
);
|
||||
return { svc };
|
||||
}
|
||||
|
||||
const body = {
|
||||
chatId: 'chat-1',
|
||||
messages: [
|
||||
{ id: 'm1', role: 'user', parts: [{ type: 'text', text: 'hi' }] },
|
||||
],
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
streamTextMock.mockReset();
|
||||
streamTextMock.mockImplementation(() => makeStreamResult());
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'log')
|
||||
.mockImplementation(() => undefined as never);
|
||||
});
|
||||
|
||||
afterEach(() => jest.restoreAllMocks());
|
||||
|
||||
it('happy path (run-wrapped): streamText is driven with abortSignal === handle.signal (the RUN signal, NOT the socket)', async () => {
|
||||
const { svc } = makeService();
|
||||
const runController = new AbortController();
|
||||
const runSignal = runController.signal;
|
||||
const socketSignal = new AbortController().signal;
|
||||
|
||||
const begin = jest.fn(async () => ({ runId: 'run-1', signal: runSignal }));
|
||||
await svc.stream({
|
||||
user: { id: 'user-1' } as never,
|
||||
workspace: { id: 'ws-1' } as never,
|
||||
sessionId: 'sess-1',
|
||||
body: body as never,
|
||||
res: makeRes() as never,
|
||||
signal: socketSignal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
runHooks: {
|
||||
begin,
|
||||
onAssistantSeeded: jest.fn(),
|
||||
onStep: jest.fn(),
|
||||
onSettled: jest.fn(),
|
||||
} as never,
|
||||
});
|
||||
|
||||
expect(begin).toHaveBeenCalledTimes(1);
|
||||
expect(streamTextMock).toHaveBeenCalledTimes(1);
|
||||
// THE assertion: the agent loop's abort is wired to the RUN, so a browser
|
||||
// disconnect (which aborts only `socketSignal`) cannot end the turn.
|
||||
expect(streamTextMock.mock.calls[0][0].abortSignal).toBe(runSignal);
|
||||
expect(streamTextMock.mock.calls[0][0].abortSignal).not.toBe(socketSignal);
|
||||
});
|
||||
|
||||
it('legacy path (no runHooks): streamText is driven with the SOCKET signal', async () => {
|
||||
const { svc } = makeService();
|
||||
const socketSignal = new AbortController().signal;
|
||||
|
||||
await svc.stream({
|
||||
user: { id: 'user-1' } as never,
|
||||
workspace: { id: 'ws-1' } as never,
|
||||
sessionId: 'sess-1',
|
||||
body: body as never,
|
||||
res: makeRes() as never,
|
||||
signal: socketSignal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
// No runHooks -> the turn stays socket-bound (flag off / default).
|
||||
});
|
||||
|
||||
expect(streamTextMock).toHaveBeenCalledTimes(1);
|
||||
expect(streamTextMock.mock.calls[0][0].abortSignal).toBe(socketSignal);
|
||||
});
|
||||
|
||||
/**
|
||||
* F9 — streamText's TERMINAL callbacks carry the #184 run lifecycle:
|
||||
* onStepFinish -> runHooks.onStep(runId, stepCount)
|
||||
* onFinish -> runHooks.onSettled(runId, 'completed') (dominant path)
|
||||
* onAbort -> runHooks.onSettled(runId, 'aborted')
|
||||
* onError -> runHooks.onSettled(runId, 'error', cause)
|
||||
* makeStreamResult() ignores the streamText options, so these callbacks never
|
||||
* fire on their own — a regression in this wiring (esp. the success path) would
|
||||
* strand the run with NO test catching it. Here we CAPTURE the options streamText
|
||||
* was handed and invoke each callback with the real wiring, asserting the run
|
||||
* hooks fire with the right args.
|
||||
*/
|
||||
// Drive stream() to the point streamText is called, capturing the options object
|
||||
// (which carries onStepFinish/onFinish/onError/onAbort) and the run hooks.
|
||||
async function captureStreamCallbacks() {
|
||||
const { svc } = makeService();
|
||||
let capturedOpts: any;
|
||||
streamTextMock.mockImplementation((opts: any) => {
|
||||
capturedOpts = opts;
|
||||
return makeStreamResult();
|
||||
});
|
||||
const runHooks = {
|
||||
begin: jest.fn(async () => ({
|
||||
runId: 'run-1',
|
||||
signal: new AbortController().signal,
|
||||
})),
|
||||
onAssistantSeeded: jest.fn(),
|
||||
onStep: jest.fn(),
|
||||
onSettled: jest.fn(),
|
||||
};
|
||||
await svc.stream({
|
||||
user: { id: 'user-1' } as never,
|
||||
workspace: { id: 'ws-1' } as never,
|
||||
sessionId: 'sess-1',
|
||||
body: body as never,
|
||||
res: makeRes() as never,
|
||||
signal: new AbortController().signal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
runHooks: runHooks as never,
|
||||
});
|
||||
expect(capturedOpts).toBeDefined();
|
||||
return { capturedOpts, runHooks };
|
||||
}
|
||||
|
||||
it('F9: onStepFinish bumps the run step count, onFinish settles the run "completed" (the dominant autonomous-run path)', async () => {
|
||||
const { capturedOpts, runHooks } = await captureStreamCallbacks();
|
||||
|
||||
// A finished step -> onStep(runId, finishedStepCount).
|
||||
capturedOpts.onStepFinish({ text: 'step one', toolCalls: [], content: [] });
|
||||
expect(runHooks.onStep).toHaveBeenCalledWith('run-1', 1);
|
||||
capturedOpts.onStepFinish({ text: 'step two', toolCalls: [], content: [] });
|
||||
expect(runHooks.onStep).toHaveBeenLastCalledWith('run-1', 2);
|
||||
|
||||
// The success terminal callback settles the run.
|
||||
await capturedOpts.onFinish({
|
||||
text: 'done',
|
||||
finishReason: 'stop',
|
||||
totalUsage: {},
|
||||
usage: {},
|
||||
steps: [],
|
||||
});
|
||||
expect(runHooks.onSettled).toHaveBeenCalledWith('run-1', 'completed');
|
||||
});
|
||||
|
||||
it('F9: onAbort settles the run "aborted"', async () => {
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'warn')
|
||||
.mockImplementation(() => undefined as never);
|
||||
const { capturedOpts, runHooks } = await captureStreamCallbacks();
|
||||
|
||||
await capturedOpts.onAbort({ steps: [] });
|
||||
expect(runHooks.onSettled).toHaveBeenCalledWith('run-1', 'aborted');
|
||||
});
|
||||
|
||||
it('F9: onError settles the run "error" carrying the provider cause', async () => {
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'error')
|
||||
.mockImplementation(() => undefined as never);
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'warn')
|
||||
.mockImplementation(() => undefined as never);
|
||||
const { capturedOpts, runHooks } = await captureStreamCallbacks();
|
||||
|
||||
await capturedOpts.onError({ error: new Error('provider exploded') });
|
||||
expect(runHooks.onSettled).toHaveBeenCalledWith(
|
||||
'run-1',
|
||||
'error',
|
||||
expect.stringContaining('provider exploded'),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* F14 — the begin-failure RESILIENCE branch (the `else` of the run-race guard).
|
||||
*
|
||||
* stream() wraps runHooks.begin in try/catch with TWO branches:
|
||||
* - RunAlreadyActiveError -> 409 ConflictException (pinned above).
|
||||
* - ANY OTHER begin failure -> SWALLOW + continue UNTRACKED on the socket signal
|
||||
* (legacy fallback): it logs "...streaming without run tracking", leaves
|
||||
* `effectiveSignal = signal` (runId undefined) and serves the turn anyway.
|
||||
*
|
||||
* The contract: a transient beginRun failure (e.g. a non-unique DB error inserting
|
||||
* the run row) must STILL serve the user's turn — it must NOT re-throw and must NOT
|
||||
* be misclassified as a 409. A regression that re-threw here would break EVERY turn
|
||||
* on a begin failure with nothing to catch it. This branch is otherwise undriven by
|
||||
* any spec, so it is pinned here SEPARATELY from the 409 path: a plain begin error
|
||||
* proceeds to streamText with the SOCKET signal and still persists the user turn.
|
||||
*/
|
||||
describe('AiChatService.stream — begin-failure resilience / legacy fallback (#184 F14)', () => {
|
||||
const streamTextMock = streamText as unknown as jest.Mock;
|
||||
|
||||
function makeStreamResult() {
|
||||
return {
|
||||
consumeStream: jest.fn(),
|
||||
pipeUIMessageStreamToResponse: jest.fn(),
|
||||
};
|
||||
}
|
||||
|
||||
function makeRes() {
|
||||
return {
|
||||
raw: {
|
||||
writeHead: jest.fn(),
|
||||
write: jest.fn(),
|
||||
once: jest.fn(),
|
||||
on: jest.fn(),
|
||||
flushHeaders: jest.fn(),
|
||||
writableEnded: false,
|
||||
destroyed: false,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
// Same harness as the F3 abortSignal block, but it also exposes
|
||||
// aiChatMessageRepo so we can assert the user turn IS persisted (the turn really
|
||||
// streamed) despite begin() blowing up.
|
||||
function makeService() {
|
||||
const aiChatRepo = {
|
||||
findById: jest.fn(async () => ({ id: 'chat-1', workspaceId: 'ws-1' })),
|
||||
insert: jest.fn(),
|
||||
};
|
||||
const aiChatMessageRepo = {
|
||||
insert: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
findAllByChat: jest.fn(async () => []),
|
||||
update: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
};
|
||||
const aiSettings = { resolve: jest.fn(async () => ({})) };
|
||||
const tools = { forUser: jest.fn(async () => ({})) };
|
||||
const mcpClients = {
|
||||
toolsFor: jest.fn(async () => ({
|
||||
tools: {},
|
||||
clients: [],
|
||||
outcomes: [],
|
||||
instructions: [],
|
||||
})),
|
||||
};
|
||||
const svc = new AiChatService(
|
||||
{} as never, // ai
|
||||
aiChatRepo as never,
|
||||
aiChatMessageRepo as never,
|
||||
{} as never, // aiChatPageSnapshotRepo
|
||||
aiSettings as never,
|
||||
tools as never,
|
||||
mcpClients as never,
|
||||
{} as never, // aiAgentRoleRepo
|
||||
{} as never, // pageRepo
|
||||
{} as never, // pageAccess
|
||||
{ isAiChatDeferredToolsEnabled: () => false } as never, // environment
|
||||
);
|
||||
return { svc, aiChatMessageRepo };
|
||||
}
|
||||
|
||||
const body = {
|
||||
chatId: 'chat-1',
|
||||
messages: [
|
||||
{ id: 'm1', role: 'user', parts: [{ type: 'text', text: 'hi' }] },
|
||||
],
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
streamTextMock.mockReset();
|
||||
streamTextMock.mockImplementation(() => makeStreamResult());
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'log')
|
||||
.mockImplementation(() => undefined as never);
|
||||
});
|
||||
|
||||
afterEach(() => jest.restoreAllMocks());
|
||||
|
||||
it('a PLAIN begin() failure (NOT RunAlreadyActiveError) does NOT 409 — it swallows, logs, and streams the turn UNTRACKED on the socket signal', async () => {
|
||||
const errorSpy = jest
|
||||
.spyOn(Logger.prototype, 'error')
|
||||
.mockImplementation(() => undefined as never);
|
||||
|
||||
const { svc, aiChatMessageRepo } = makeService();
|
||||
const socketSignal = new AbortController().signal;
|
||||
|
||||
// A transient, NON-race begin failure (e.g. a non-unique DB error inserting
|
||||
// the run row). This is the `else` branch of the begin try/catch.
|
||||
const begin = jest.fn(async () => {
|
||||
throw new Error('insert failed');
|
||||
});
|
||||
|
||||
const promise = svc.stream({
|
||||
user: { id: 'user-1' } as never,
|
||||
workspace: { id: 'ws-1' } as never,
|
||||
sessionId: 'sess-1',
|
||||
body: body as never,
|
||||
res: makeRes() as never,
|
||||
signal: socketSignal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
runHooks: {
|
||||
begin,
|
||||
onAssistantSeeded: jest.fn(),
|
||||
onStep: jest.fn(),
|
||||
onSettled: jest.fn(),
|
||||
} as never,
|
||||
});
|
||||
|
||||
// The turn proceeds: NO throw at all (in particular NOT a 409).
|
||||
await expect(promise).resolves.toBeUndefined();
|
||||
|
||||
expect(begin).toHaveBeenCalledTimes(1);
|
||||
|
||||
// The resilience branch logged the legacy-fallback warning.
|
||||
expect(errorSpy).toHaveBeenCalledWith(
|
||||
expect.stringContaining('streaming without run tracking'),
|
||||
expect.anything(),
|
||||
);
|
||||
|
||||
// The turn really streamed: the user message was persisted and streamText ran.
|
||||
expect(aiChatMessageRepo.insert).toHaveBeenCalled();
|
||||
expect(streamTextMock).toHaveBeenCalledTimes(1);
|
||||
|
||||
// The decisive wiring: with no run handle, the fallback uses the SOCKET signal
|
||||
// (effectiveSignal = signal, runId undefined) — not a run-bound signal.
|
||||
expect(streamTextMock.mock.calls[0][0].abortSignal).toBe(socketSignal);
|
||||
});
|
||||
});
|
||||
@@ -453,6 +453,12 @@ describe('chatStreamMetadata', () => {
|
||||
});
|
||||
});
|
||||
|
||||
it('attaches the runId on the start part when a run wraps the turn (#184)', () => {
|
||||
expect(
|
||||
chatStreamMetadata({ type: 'start' }, 'chat-1', undefined, 'run-1'),
|
||||
).toEqual({ chatId: 'chat-1', runId: 'run-1' });
|
||||
});
|
||||
|
||||
it('returns the CUMULATIVE step usage passed in for the finish-step part', () => {
|
||||
// finish-step usage is per-step in v6; the caller accumulates and passes the
|
||||
// running sum, which this just wraps.
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -43,6 +43,30 @@ export class BoundChatDto {
|
||||
pageId: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Reconnect to the latest run of a chat (#184): fetch its persisted lifecycle
|
||||
* state (and the assistant message it projects) for an in-flight or finished run.
|
||||
*/
|
||||
export class GetRunDto {
|
||||
@IsString()
|
||||
chatId: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Explicitly STOP an agent run (#184): the user pressed Stop — distinct from a
|
||||
* browser disconnect, which never stops a run. Either the run id (preferred, from
|
||||
* the streamed start metadata) or the chat id (stop whatever run is active on it).
|
||||
*/
|
||||
export class StopRunDto {
|
||||
@IsOptional()
|
||||
@IsString()
|
||||
runId?: string;
|
||||
|
||||
@IsOptional()
|
||||
@IsString()
|
||||
chatId?: string;
|
||||
}
|
||||
|
||||
/** Export a chat to Markdown (#183). `lang` localizes the few fixed
|
||||
* role/tool-action labels; defaults to English server-side. */
|
||||
export class ExportChatDto {
|
||||
|
||||
@@ -51,21 +51,7 @@ export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
|
||||
throw new UnauthorizedException();
|
||||
}
|
||||
|
||||
// #348 — reuse the workspace DomainMiddleware already loaded for this request
|
||||
// instead of re-querying it. `validate()` above has confirmed
|
||||
// `req.raw.workspaceId === payload.workspaceId` (or that it is unset), and the
|
||||
// middleware sets `req.raw.workspace` alongside `req.raw.workspaceId` from the
|
||||
// SAME workspace row, so when the ids match this is that row. NOTE it is the
|
||||
// middleware's `selectAll` object (a superset of the fallback `findById` base
|
||||
// fields — it also carries licenseKey/auditRetentionDays); that is harmless
|
||||
// here because every consumer reads this workspace via the AuthWorkspace
|
||||
// decorator, which already preferred `req.raw.workspace` (the selectAll object)
|
||||
// over `req.user.workspace` before this change. Fall back to the query if the
|
||||
// middleware did not populate it (a path that bypasses DomainMiddleware).
|
||||
const workspace =
|
||||
req.raw.workspace && req.raw.workspaceId === payload.workspaceId
|
||||
? req.raw.workspace
|
||||
: await this.workspaceRepo.findById(payload.workspaceId);
|
||||
const workspace = await this.workspaceRepo.findById(payload.workspaceId);
|
||||
|
||||
if (!workspace) {
|
||||
throw new UnauthorizedException();
|
||||
|
||||
@@ -38,8 +38,6 @@ export class FavoriteService {
|
||||
await this.pagePermissionRepo.filterAccessiblePageIds({
|
||||
pageIds: result.items,
|
||||
userId,
|
||||
// #348 — favorites load at app-start; enable the workspace short-circuit.
|
||||
workspaceId,
|
||||
});
|
||||
const accessibleSet = new Set(accessibleIds);
|
||||
result.items = result.items.filter((id) => accessibleSet.has(id));
|
||||
@@ -127,8 +125,6 @@ export class FavoriteService {
|
||||
await this.pagePermissionRepo.filterAccessiblePageIds({
|
||||
pageIds,
|
||||
userId,
|
||||
// #348 — workspace-level short-circuit for the favorites list.
|
||||
workspaceId,
|
||||
});
|
||||
accessiblePageSet = new Set(accessibleIds);
|
||||
}
|
||||
|
||||
@@ -23,12 +23,7 @@ export class NotificationController {
|
||||
@Body() dto: ListNotificationsDto,
|
||||
@AuthUser() user: User,
|
||||
) {
|
||||
return this.notificationService.findByUserId(
|
||||
user.id,
|
||||
dto,
|
||||
dto.type,
|
||||
user.workspaceId,
|
||||
);
|
||||
return this.notificationService.findByUserId(user.id, dto, dto.type);
|
||||
}
|
||||
|
||||
@HttpCode(HttpStatus.OK)
|
||||
|
||||
@@ -45,7 +45,6 @@ export class NotificationService {
|
||||
userId: string,
|
||||
pagination: PaginationOptions,
|
||||
type: NotificationTab = 'all',
|
||||
workspaceId?: string | null,
|
||||
) {
|
||||
const result = await this.notificationRepo.findByUserId(
|
||||
userId,
|
||||
@@ -62,8 +61,6 @@ export class NotificationService {
|
||||
await this.pagePermissionRepo.filterAccessiblePageIds({
|
||||
pageIds,
|
||||
userId,
|
||||
// #348 — notifications list; enable the workspace short-circuit.
|
||||
workspaceId,
|
||||
});
|
||||
const accessibleSet = new Set(accessiblePageIds);
|
||||
|
||||
|
||||
@@ -446,11 +446,7 @@ export class PageController {
|
||||
);
|
||||
}
|
||||
|
||||
return this.pageService.getRecentPages(
|
||||
user.id,
|
||||
pagination,
|
||||
user.workspaceId,
|
||||
);
|
||||
return this.pageService.getRecentPages(user.id, pagination);
|
||||
}
|
||||
|
||||
@HttpCode(HttpStatus.OK)
|
||||
@@ -473,13 +469,7 @@ export class PageController {
|
||||
}
|
||||
}
|
||||
|
||||
return this.pageService.getCreatedByPages(
|
||||
targetUserId,
|
||||
user.id,
|
||||
pagination,
|
||||
dto.spaceId,
|
||||
user.workspaceId,
|
||||
);
|
||||
return this.pageService.getCreatedByPages(targetUserId, user.id, pagination, dto.spaceId);
|
||||
}
|
||||
|
||||
@HttpCode(HttpStatus.OK)
|
||||
|
||||
@@ -1163,7 +1163,6 @@ export class PageService {
|
||||
async getRecentPages(
|
||||
userId: string,
|
||||
pagination: PaginationOptions,
|
||||
workspaceId?: string | null,
|
||||
): Promise<CursorPaginationResult<Page>> {
|
||||
const result = await this.pageRepo.getRecentPages(userId, pagination);
|
||||
|
||||
@@ -1173,8 +1172,6 @@ export class PageService {
|
||||
await this.pagePermissionRepo.filterAccessiblePageIds({
|
||||
pageIds,
|
||||
userId,
|
||||
// #348 — cross-space "recent"; enable the workspace short-circuit.
|
||||
workspaceId,
|
||||
});
|
||||
const accessibleSet = new Set(accessibleIds);
|
||||
result.items = result.items.filter((p) => accessibleSet.has(p.id));
|
||||
@@ -1188,7 +1185,6 @@ export class PageService {
|
||||
requestingUserId: string,
|
||||
pagination: PaginationOptions,
|
||||
spaceId?: string,
|
||||
workspaceId?: string | null,
|
||||
): Promise<CursorPaginationResult<Page>> {
|
||||
const result = await this.pageRepo.getCreatedByPages(
|
||||
creatorId,
|
||||
@@ -1203,9 +1199,6 @@ export class PageService {
|
||||
await this.pagePermissionRepo.filterAccessiblePageIds({
|
||||
pageIds,
|
||||
userId: requestingUserId,
|
||||
spaceId,
|
||||
// #348 — enable the workspace short-circuit when not space-scoped.
|
||||
workspaceId,
|
||||
});
|
||||
const accessibleSet = new Set(accessibleIds);
|
||||
result.items = result.items.filter((p) => accessibleSet.has(p.id));
|
||||
|
||||
@@ -93,41 +93,6 @@ function collectNodes<T>(
|
||||
return Array.from(byKey.values());
|
||||
}
|
||||
|
||||
/**
|
||||
* #348 — cheap early-exit probe: does this doc contain ANY node the transclusion
|
||||
* syncs care about (`transclusionSource` / `transclusionReference` / `pageEmbed`)?
|
||||
* Lets the collab store skip the three sync SELECTs when neither the previous nor
|
||||
* the new content has any such node — there is nothing to insert, and (since the
|
||||
* DB mirrors the previously-persisted content) nothing to delete. Walks once and
|
||||
* short-circuits on the first match; uses the same depth ceiling as the
|
||||
* collectors. Deliberately does NOT skip `transclusionSource` subtrees: it only
|
||||
* answers "any node present?", so descending everywhere is strictly conservative
|
||||
* (it can never wrongly report "none").
|
||||
*/
|
||||
export function hasTransclusionFamilyNodes(doc: unknown): boolean {
|
||||
const visit = (node: any, depth: number): boolean => {
|
||||
if (!node || typeof node !== 'object') return false;
|
||||
if (depth > MAX_PM_WALK_DEPTH) return false;
|
||||
|
||||
if (
|
||||
node.type === TRANSCLUSION_TYPE ||
|
||||
node.type === REFERENCE_TYPE ||
|
||||
node.type === PAGE_EMBED_TYPE
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
|
||||
if (Array.isArray(node.content)) {
|
||||
for (const child of node.content) {
|
||||
if (visit(child, depth + 1)) return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
};
|
||||
|
||||
return visit(doc, 0);
|
||||
}
|
||||
|
||||
/**
|
||||
* Walks a ProseMirror JSON document and returns one snapshot per top-level
|
||||
* `transclusion` node. Does not recurse into transclusions (schema disallows
|
||||
|
||||
@@ -155,8 +155,6 @@ export class SearchService {
|
||||
pageIds,
|
||||
userId: opts.userId,
|
||||
spaceId: searchParams.spaceId,
|
||||
// #348 — enables the workspace-level short-circuit when not space-scoped.
|
||||
workspaceId: opts.workspaceId,
|
||||
});
|
||||
const accessibleSet = new Set(accessibleIds);
|
||||
results = results.filter((r: any) => accessibleSet.has(r.id));
|
||||
@@ -268,8 +266,6 @@ export class SearchService {
|
||||
await this.pagePermissionRepo.filterAccessiblePageIds({
|
||||
pageIds,
|
||||
userId,
|
||||
// #348 — workspace-level short-circuit for the suggest path.
|
||||
workspaceId,
|
||||
});
|
||||
const accessibleSet = new Set(accessibleIds);
|
||||
pages = pages.filter((p) => accessibleSet.has(p.id));
|
||||
|
||||
@@ -55,6 +55,14 @@ export class UpdateWorkspaceDto extends PartialType(CreateWorkspaceDto) {
|
||||
@IsBoolean()
|
||||
aiDictationStreaming: boolean;
|
||||
|
||||
// #184: detached/autonomous agent runs (settings.ai.autonomousRuns). When on, a
|
||||
// chat turn becomes a server-side RUN that survives a browser disconnect; only
|
||||
// an explicit /ai-chat/stop ends it. Off by default; single-instance-only in
|
||||
// phase 1 (see AiChatRunService.warnIfMultiInstance / AGENTS.md).
|
||||
@IsOptional()
|
||||
@IsBoolean()
|
||||
autonomousRuns: boolean;
|
||||
|
||||
// Workspace master toggle that enables/disables the HTML embed block type.
|
||||
// Persisted at settings.htmlEmbed. ABSENT/false => OFF (default). The block
|
||||
// itself renders in a sandboxed iframe, so this is a feature switch, not a
|
||||
|
||||
@@ -526,6 +526,20 @@ export class WorkspaceService {
|
||||
);
|
||||
}
|
||||
|
||||
if (typeof updateWorkspaceDto.autonomousRuns !== 'undefined') {
|
||||
const prev = settingsBefore?.ai?.autonomousRuns ?? false;
|
||||
if (prev !== updateWorkspaceDto.autonomousRuns) {
|
||||
before.autonomousRuns = prev;
|
||||
after.autonomousRuns = updateWorkspaceDto.autonomousRuns;
|
||||
}
|
||||
await this.workspaceRepo.updateAiSettings(
|
||||
workspaceId,
|
||||
'autonomousRuns',
|
||||
updateWorkspaceDto.autonomousRuns,
|
||||
trx,
|
||||
);
|
||||
}
|
||||
|
||||
if (typeof updateWorkspaceDto.htmlEmbed !== 'undefined') {
|
||||
const prev = settingsBefore?.htmlEmbed ?? false;
|
||||
if (prev !== updateWorkspaceDto.htmlEmbed) {
|
||||
@@ -579,6 +593,7 @@ export class WorkspaceService {
|
||||
delete updateWorkspaceDto.aiChat;
|
||||
delete updateWorkspaceDto.aiDictation;
|
||||
delete updateWorkspaceDto.aiDictationStreaming;
|
||||
delete updateWorkspaceDto.autonomousRuns;
|
||||
delete updateWorkspaceDto.htmlEmbed;
|
||||
delete updateWorkspaceDto.trackerHead;
|
||||
delete updateWorkspaceDto.aiPublicShareAssistant;
|
||||
|
||||
@@ -31,6 +31,7 @@ import { FavoriteRepo } from '@docmost/db/repos/favorite/favorite.repo';
|
||||
import { TemplateRepo } from '@docmost/db/repos/template/template.repo';
|
||||
import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo';
|
||||
import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
|
||||
import { AiChatRunRepo } from '@docmost/db/repos/ai-chat/ai-chat-run.repo';
|
||||
import { AiChatPageSnapshotRepo } from '@docmost/db/repos/ai-chat/ai-chat-page-snapshot.repo';
|
||||
import { AiProviderCredentialsRepo } from '@docmost/db/repos/ai-chat/ai-provider-credentials.repo';
|
||||
import { AiMcpServerRepo } from '@docmost/db/repos/ai-chat/ai-mcp-server.repo';
|
||||
@@ -122,6 +123,7 @@ import { firstSqlToken } from '../integrations/metrics/metrics.constants';
|
||||
TemplateRepo,
|
||||
AiChatRepo,
|
||||
AiChatMessageRepo,
|
||||
AiChatRunRepo,
|
||||
AiChatPageSnapshotRepo,
|
||||
AiProviderCredentialsRepo,
|
||||
AiMcpServerRepo,
|
||||
@@ -156,6 +158,7 @@ import { firstSqlToken } from '../integrations/metrics/metrics.constants';
|
||||
TemplateRepo,
|
||||
AiChatRepo,
|
||||
AiChatMessageRepo,
|
||||
AiChatRunRepo,
|
||||
AiChatPageSnapshotRepo,
|
||||
AiProviderCredentialsRepo,
|
||||
AiMcpServerRepo,
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
import { type Kysely, sql } from 'kysely';
|
||||
|
||||
/**
|
||||
* `ai_chat_runs` — the agent RUN as a first-class, server-side lifecycle object
|
||||
* (#184 phase 1: autonomous agent runs detached from the browser window).
|
||||
*
|
||||
* Until now an agent turn lived ONLY as long as the HTTP request was open
|
||||
* (`res.hijack()` in ai-chat.controller.ts); a browser disconnect aborted it.
|
||||
* This table makes a turn a persistent object the server owns: it is created
|
||||
* when a run starts (inserted directly as 'running' in phase 1 — 'pending' is
|
||||
* only this column's default + a reserved value, never written by code yet) and
|
||||
* advances to succeeded|failed|aborted, surviving the subscriber (browser) going
|
||||
* away when it settles. The DB is the source of
|
||||
* truth — a later client reconnects/sees the result by reading this row plus the
|
||||
* assistant message it projects (`assistant_message_id`).
|
||||
*
|
||||
* The assistant message row (#183 step-granular durability) is the PROJECTION of
|
||||
* a run's output; this row is the run's LIFECYCLE. They are linked by
|
||||
* `assistant_message_id` (SET NULL if the message is later pruned).
|
||||
*
|
||||
* `status` : 'pending' | 'running' | 'succeeded' | 'failed' | 'aborted'.
|
||||
* `trigger` : 'user' | 'autostart' | 'schedule' | 'api' | 'continue' — only
|
||||
* 'user' is produced in phase 1; the others are reserved for the
|
||||
* autonomy triggers deferred to phase 2 so they need no later
|
||||
* migration.
|
||||
*
|
||||
* ONE ACTIVE RUN PER CHAT is enforced by a partial unique index on `chat_id`
|
||||
* WHERE status IN ('pending','running'): an autonomous run and a user run can
|
||||
* never trample each other on the same chat. Settled runs (succeeded/failed/
|
||||
* aborted) are excluded from the index so a chat can accumulate any number of
|
||||
* historical runs.
|
||||
*/
|
||||
export async function up(db: Kysely<any>): Promise<void> {
|
||||
await db.schema
|
||||
.createTable('ai_chat_runs')
|
||||
.ifNotExists()
|
||||
.addColumn('id', 'uuid', (col) =>
|
||||
col.primaryKey().defaultTo(sql`gen_uuid_v7()`),
|
||||
)
|
||||
.addColumn('chat_id', 'uuid', (col) =>
|
||||
col.references('ai_chats.id').onDelete('cascade').notNull(),
|
||||
)
|
||||
.addColumn('workspace_id', 'uuid', (col) =>
|
||||
col.references('workspaces.id').onDelete('cascade').notNull(),
|
||||
)
|
||||
// The human who triggered the run (audit). SET NULL on user deletion so the
|
||||
// run history outlives its author; NULL is also the natural value for a
|
||||
// future system/cron/api trigger with no human actor.
|
||||
.addColumn('created_by', 'uuid', (col) =>
|
||||
col.references('users.id').onDelete('set null'),
|
||||
)
|
||||
// The assistant message this run materializes (the #183 projection). SET NULL
|
||||
// if that message row is later deleted; nullable because the run row is
|
||||
// created a moment BEFORE the assistant row is seeded.
|
||||
.addColumn('assistant_message_id', 'uuid', (col) =>
|
||||
col.references('ai_chat_messages.id').onDelete('set null'),
|
||||
)
|
||||
.addColumn('trigger', 'varchar(20)', (col) =>
|
||||
col.notNull().defaultTo('user'),
|
||||
)
|
||||
.addColumn('status', 'varchar(20)', (col) =>
|
||||
col.notNull().defaultTo('pending'),
|
||||
)
|
||||
// Terminal error message for a failed run (provider/transport cause),
|
||||
// mirroring the assistant message's metadata.error.
|
||||
.addColumn('error', 'text', (col) => col)
|
||||
// Number of agent steps finished so far (kept monotonic with the projection).
|
||||
.addColumn('step_count', 'integer', (col) => col.notNull().defaultTo(0))
|
||||
// Set when an EXPLICIT user stop is requested (distinct from a mere browser
|
||||
// disconnect, which never stops a run). The runner aborts the turn and the
|
||||
// run settles as 'aborted'.
|
||||
.addColumn('stop_requested_at', 'timestamptz', (col) => col)
|
||||
.addColumn('started_at', 'timestamptz', (col) => col)
|
||||
.addColumn('finished_at', 'timestamptz', (col) => col)
|
||||
.addColumn('created_at', 'timestamptz', (col) =>
|
||||
col.notNull().defaultTo(sql`now()`),
|
||||
)
|
||||
.addColumn('updated_at', 'timestamptz', (col) =>
|
||||
col.notNull().defaultTo(sql`now()`),
|
||||
)
|
||||
.execute();
|
||||
|
||||
// Reconnect / "latest run for this chat" reads hit chat_id first.
|
||||
await db.schema
|
||||
.createIndex('ai_chat_runs_chat_id_idx')
|
||||
.ifNotExists()
|
||||
.on('ai_chat_runs')
|
||||
.column('chat_id')
|
||||
.execute();
|
||||
|
||||
// One ACTIVE run per chat (advisory at the DB level): a second pending/running
|
||||
// run on the same chat is rejected, so a user turn and an autonomous turn can
|
||||
// never race on the same chat. Partial so settled runs do not collide.
|
||||
await db.schema
|
||||
.createIndex('ai_chat_runs_one_active_per_chat')
|
||||
.ifNotExists()
|
||||
.on('ai_chat_runs')
|
||||
.column('chat_id')
|
||||
.unique()
|
||||
.where(sql.ref('status'), 'in', sql`('pending','running')`)
|
||||
.execute();
|
||||
}
|
||||
|
||||
export async function down(db: Kysely<any>): Promise<void> {
|
||||
await db.schema.dropTable('ai_chat_runs').execute();
|
||||
}
|
||||
@@ -1,118 +0,0 @@
|
||||
import { type Kysely, sql } from 'kysely';
|
||||
|
||||
/**
|
||||
* #348 — targeted hot-path indexes.
|
||||
*
|
||||
* 1. GIN trigram indexes for `/search/suggest`. That endpoint runs a
|
||||
* leading-wildcard `LOWER(f_unaccent(col)) LIKE '%q%'` per keystroke, which
|
||||
* is a sequential scan without a trigram index. The index EXPRESSIONS below
|
||||
* are `LOWER(f_unaccent(title|name))`, matching the predicates in
|
||||
* search.service.ts exactly so the planner uses them (verified with EXPLAIN:
|
||||
* the suggest predicate resolves to a Bitmap Index Scan on these indexes).
|
||||
*
|
||||
* IMMUTABLE-wrapper fix (required for the index to build): `f_unaccent` was
|
||||
* defined as `SELECT unaccent('unaccent', $1)` (the two-arg, dictionary-named
|
||||
* unaccent). That body CANNOT be used in an index expression: when Postgres
|
||||
* inlines the IMMUTABLE SQL wrapper while building the index it fails to
|
||||
* resolve the two-arg call (`function unaccent(unknown, text) does not exist`,
|
||||
* the `'unaccent'` literal loses its regdictionary coercion). The single-arg
|
||||
* `unaccent($1)` is the same operation (the default text-search dictionary IS
|
||||
* `unaccent`; verified byte-equal on accented samples), and — crucially —
|
||||
* SCHEMA-QUALIFIED as `public.unaccent($1)` it inlines cleanly, so the index
|
||||
* builds. We therefore `CREATE OR REPLACE` `f_unaccent` to the qualified
|
||||
* single-arg body. This is output-identical for every existing caller (the
|
||||
* tsvector trigger, the main `tsv @@` search, and the suggest LIKE), so no
|
||||
* reindex/backfill is needed; `down()` restores the original two-arg body.
|
||||
* (The `unaccent` extension is installed in `public` in this codebase, which
|
||||
* is why `public.unaccent` is the correct qualification.)
|
||||
*
|
||||
* 2. Composite indexes for two ORDER-BY-only-on-id queries that currently sort
|
||||
* on top of a created_at index:
|
||||
* - page_history: `findPageHistoryByPageId` does WHERE page_id ORDER BY id
|
||||
* DESC, but only `(page_id, created_at DESC)` exists → extra sort.
|
||||
* - comments: `findPageComments` does WHERE page_id ORDER BY id ASC, but only
|
||||
* `(page_id)` exists → extra sort.
|
||||
*
|
||||
* DEPLOY-TIME LOCK WARNING: these are plain (non-CONCURRENT) CREATE INDEX
|
||||
* statements — CONCURRENTLY is impossible because Kysely runs each migration in a
|
||||
* transaction. They take a SHARE lock that BLOCKS writes (INSERT/UPDATE/DELETE) on
|
||||
* pages/users/groups/comments/page_history for the duration of the build. The two
|
||||
* GIN trigram builds on pages.title / users.name are the slow ones and can take
|
||||
* minutes on a large tenant → a write-outage window during the deploy migration.
|
||||
* For large installations, run this migration in a maintenance window, or build
|
||||
* the trigram indexes out-of-band with CREATE INDEX CONCURRENTLY before deploying
|
||||
* (then this migration's `IF NOT EXISTS` is a no-op). Small/typical tenants are
|
||||
* unaffected.
|
||||
*/
|
||||
export async function up(db: Kysely<any>): Promise<void> {
|
||||
// Index-compatible, output-identical redefinition of f_unaccent (see header).
|
||||
await sql`
|
||||
CREATE OR REPLACE FUNCTION f_unaccent(text)
|
||||
RETURNS text
|
||||
LANGUAGE sql
|
||||
IMMUTABLE PARALLEL SAFE STRICT
|
||||
AS $func$
|
||||
SELECT public.unaccent($1);
|
||||
$func$
|
||||
`.execute(db);
|
||||
|
||||
// Search-suggest trigram indexes. Expressions match search.service.ts.
|
||||
await sql`
|
||||
CREATE INDEX IF NOT EXISTS idx_pages_title_trgm
|
||||
ON pages USING gin ((LOWER(f_unaccent(title))) gin_trgm_ops)
|
||||
`.execute(db);
|
||||
|
||||
await sql`
|
||||
CREATE INDEX IF NOT EXISTS idx_users_name_trgm
|
||||
ON users USING gin ((LOWER(f_unaccent(name))) gin_trgm_ops)
|
||||
`.execute(db);
|
||||
|
||||
await sql`
|
||||
CREATE INDEX IF NOT EXISTS idx_groups_name_trgm
|
||||
ON groups USING gin ((LOWER(f_unaccent(name))) gin_trgm_ops)
|
||||
`.execute(db);
|
||||
|
||||
// page_history: WHERE page_id ORDER BY id DESC (findPageHistoryByPageId).
|
||||
await sql`
|
||||
CREATE INDEX IF NOT EXISTS idx_page_history_page_id
|
||||
ON page_history (page_id, id DESC)
|
||||
`.execute(db);
|
||||
|
||||
// comments: WHERE page_id ORDER BY id ASC (findPageComments).
|
||||
await sql`
|
||||
CREATE INDEX IF NOT EXISTS idx_comments_page_id_id
|
||||
ON comments (page_id, id)
|
||||
`.execute(db);
|
||||
|
||||
// page_access(workspace_id): #348 made hasRestrictedPagesInWorkspace uncached
|
||||
// (F1 fix), so `EXISTS(SELECT 1 FROM page_access WHERE workspace_id=?)` now runs
|
||||
// per-request on every whole-workspace list endpoint (global search + suggest,
|
||||
// favorites, notifications, recent, created-by). page_access only had a
|
||||
// space_id index → that EXISTS was a seq scan in the common zero-restriction
|
||||
// case. This index makes it an index-only existence probe.
|
||||
await sql`
|
||||
CREATE INDEX IF NOT EXISTS idx_page_access_workspace_id
|
||||
ON page_access (workspace_id)
|
||||
`.execute(db);
|
||||
}
|
||||
|
||||
export async function down(db: Kysely<any>): Promise<void> {
|
||||
// Drop the expression indexes before restoring the function body.
|
||||
await sql`DROP INDEX IF EXISTS idx_pages_title_trgm`.execute(db);
|
||||
await sql`DROP INDEX IF EXISTS idx_users_name_trgm`.execute(db);
|
||||
await sql`DROP INDEX IF EXISTS idx_groups_name_trgm`.execute(db);
|
||||
await sql`DROP INDEX IF EXISTS idx_page_history_page_id`.execute(db);
|
||||
await sql`DROP INDEX IF EXISTS idx_comments_page_id_id`.execute(db);
|
||||
await sql`DROP INDEX IF EXISTS idx_page_access_workspace_id`.execute(db);
|
||||
|
||||
// Restore the original two-arg (dictionary-named) f_unaccent body.
|
||||
await sql`
|
||||
CREATE OR REPLACE FUNCTION f_unaccent(text)
|
||||
RETURNS text
|
||||
LANGUAGE sql
|
||||
IMMUTABLE PARALLEL SAFE STRICT
|
||||
AS $func$
|
||||
SELECT unaccent('unaccent', $1);
|
||||
$func$
|
||||
`.execute(db);
|
||||
}
|
||||
@@ -121,6 +121,23 @@ export class AiChatMessageRepo {
|
||||
return rows.reverse();
|
||||
}
|
||||
|
||||
/** Fetch a single message by id + workspace (e.g. a run's projection row for
|
||||
* the #184 reconnect read). Returns undefined when nothing matches. */
|
||||
async findById(
|
||||
id: string,
|
||||
workspaceId: string,
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<AiChatMessage | undefined> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
return db
|
||||
.selectFrom('aiChatMessages')
|
||||
.select(this.baseFields)
|
||||
.where('id', '=', id)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.where('deletedAt', 'is', null)
|
||||
.executeTakeFirst();
|
||||
}
|
||||
|
||||
async insert(
|
||||
insertable: InsertableAiChatMessage,
|
||||
trx?: KyselyTransaction,
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
import { AiChatRunRepo, SWEEP_RUN_STALE_MS } from './ai-chat-run.repo';
|
||||
import type { KyselyDB } from '../../types/kysely.types';
|
||||
|
||||
/**
|
||||
* Unit coverage for AiChatRunRepo.sweepRunning over a chainable builder mock (no
|
||||
* live DB). The F1 invariant under test (DECISION C): the BOOT sweep is
|
||||
* UNCONDITIONAL — it adds NO `updatedAt <` predicate, so a fresh 'running' run
|
||||
* (updatedAt = now) IS settled rather than skipped by a staleness window. The
|
||||
* window is added ONLY when an explicit `staleMs` is supplied (the future phase-2
|
||||
* multi-instance timer sweep). We assert the EXACT predicates the spec mandates.
|
||||
*/
|
||||
describe('AiChatRunRepo.sweepRunning', () => {
|
||||
type Recorded = {
|
||||
table?: string;
|
||||
set?: Record<string, unknown>;
|
||||
wheres: Array<[string, string, unknown]>;
|
||||
returning?: string;
|
||||
};
|
||||
|
||||
function makeDb(swept: Array<{ id: string }>): {
|
||||
db: KyselyDB;
|
||||
rec: Recorded;
|
||||
} {
|
||||
const rec: Recorded = { wheres: [] };
|
||||
const builder: Record<string, unknown> = {};
|
||||
builder.set = (v: Record<string, unknown>) => {
|
||||
rec.set = v;
|
||||
return builder;
|
||||
};
|
||||
builder.where = (col: string, op: string, val: unknown) => {
|
||||
rec.wheres.push([col, op, val]);
|
||||
return builder;
|
||||
};
|
||||
builder.returning = (col: string) => {
|
||||
rec.returning = col;
|
||||
return builder;
|
||||
};
|
||||
builder.execute = () => Promise.resolve(swept);
|
||||
const db = {
|
||||
updateTable: (table: string) => {
|
||||
rec.table = table;
|
||||
return builder;
|
||||
},
|
||||
} as unknown as KyselyDB;
|
||||
return { db, rec };
|
||||
}
|
||||
|
||||
it('F1: the boot sweep (no staleMs) is UNCONDITIONAL — only a status filter, NO updatedAt window', async () => {
|
||||
const { db, rec } = makeDb([{ id: 'r1' }, { id: 'r2' }]);
|
||||
const repo = new AiChatRunRepo(db);
|
||||
|
||||
const swept = await repo.sweepRunning();
|
||||
|
||||
expect(swept).toBe(2);
|
||||
expect(rec.table).toBe('aiChatRuns');
|
||||
// The status filter is always present...
|
||||
expect(rec.wheres).toContainEqual([
|
||||
'status',
|
||||
'in',
|
||||
expect.arrayContaining(['pending', 'running']),
|
||||
]);
|
||||
// ...but a fresh 'running' run (updatedAt = now) must NOT be skipped: no
|
||||
// updatedAt predicate at all on the boot path.
|
||||
expect(rec.wheres.some(([col]) => col === 'updatedAt')).toBe(false);
|
||||
// It flips to 'aborted' and stamps finishedAt.
|
||||
expect(rec.set).toEqual(
|
||||
expect.objectContaining({ status: 'aborted', finishedAt: expect.any(Date) }),
|
||||
);
|
||||
});
|
||||
|
||||
it('phase-2 path: an explicit staleMs reintroduces the updatedAt window', async () => {
|
||||
const { db, rec } = makeDb([]);
|
||||
const repo = new AiChatRunRepo(db);
|
||||
|
||||
await repo.sweepRunning({ staleMs: SWEEP_RUN_STALE_MS });
|
||||
|
||||
const updatedAtWhere = rec.wheres.find(([col]) => col === 'updatedAt');
|
||||
expect(updatedAtWhere).toBeDefined();
|
||||
expect(updatedAtWhere![1]).toBe('<');
|
||||
expect(updatedAtWhere![2]).toBeInstanceOf(Date);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,212 @@
|
||||
import { Injectable, Logger } from '@nestjs/common';
|
||||
import { InjectKysely } from 'nestjs-kysely';
|
||||
import { sql } from 'kysely';
|
||||
import { KyselyDB, KyselyTransaction } from '../../types/kysely.types';
|
||||
import { dbOrTx } from '../../utils';
|
||||
import {
|
||||
AiChatRun,
|
||||
InsertableAiChatRun,
|
||||
} from '@docmost/db/types/entity.types';
|
||||
|
||||
// Statuses that count as "the run is still live" (an autonomous and a user run
|
||||
// must never both be live on one chat — enforced by the partial unique index and
|
||||
// checked here for friendly 409s before the insert races the constraint).
|
||||
export const ACTIVE_RUN_STATUSES = ['pending', 'running'] as const;
|
||||
|
||||
// Crash-recovery sweep recency threshold (mirrors AiChatMessageRepo.sweepStreaming,
|
||||
// #183): when a staleness window is supplied, a 'running'/'pending' run is only
|
||||
// swept to 'aborted' once it has been UNTOUCHED for this long, so a sibling
|
||||
// replica's boot-sweep can never abort a run another replica is actively
|
||||
// executing. The runner bumps `updatedAt` on every step, so a live run never
|
||||
// matches. PHASE 1 is single-process and the boot sweep passes NO window (every
|
||||
// dangling run is settled unconditionally — see sweepRunning / F1). This constant
|
||||
// is the window to reintroduce for the phase-2 multi-instance timer sweep.
|
||||
export const SWEEP_RUN_STALE_MS = 10 * 60 * 1000; // 10 minutes
|
||||
|
||||
/**
|
||||
* Repository for `ai_chat_runs` (#184 phase 1): the agent run as a first-class,
|
||||
* server-side lifecycle object detached from the HTTP request. The run row is the
|
||||
* point a client subscribes/reconnects to (by `id` or by chat); the assistant
|
||||
* message it links to (`assistantMessageId`) is the #183 projection of its output.
|
||||
*/
|
||||
@Injectable()
|
||||
export class AiChatRunRepo {
|
||||
private readonly logger = new Logger(AiChatRunRepo.name);
|
||||
|
||||
private baseFields: Array<keyof AiChatRun> = [
|
||||
'id',
|
||||
'chatId',
|
||||
'workspaceId',
|
||||
'createdBy',
|
||||
'assistantMessageId',
|
||||
'trigger',
|
||||
'status',
|
||||
'error',
|
||||
'stepCount',
|
||||
'stopRequestedAt',
|
||||
'startedAt',
|
||||
'finishedAt',
|
||||
'createdAt',
|
||||
'updatedAt',
|
||||
];
|
||||
|
||||
constructor(@InjectKysely() private readonly db: KyselyDB) {}
|
||||
|
||||
async insert(
|
||||
insertable: InsertableAiChatRun,
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<AiChatRun> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
return db
|
||||
.insertInto('aiChatRuns')
|
||||
.values(insertable)
|
||||
.returning(this.baseFields)
|
||||
.executeTakeFirst();
|
||||
}
|
||||
|
||||
async findById(
|
||||
id: string,
|
||||
workspaceId: string,
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<AiChatRun | undefined> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
return db
|
||||
.selectFrom('aiChatRuns')
|
||||
.select(this.baseFields)
|
||||
.where('id', '=', id)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.executeTakeFirst();
|
||||
}
|
||||
|
||||
/** The currently-active (pending|running) run for a chat, if any. At most one
|
||||
* exists thanks to the partial unique index. */
|
||||
async findActiveByChat(
|
||||
chatId: string,
|
||||
workspaceId: string,
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<AiChatRun | undefined> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
return db
|
||||
.selectFrom('aiChatRuns')
|
||||
.select(this.baseFields)
|
||||
.where('chatId', '=', chatId)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[])
|
||||
.executeTakeFirst();
|
||||
}
|
||||
|
||||
/** The most-recent run for a chat (active or settled) — the reconnect target. */
|
||||
async findLatestByChat(
|
||||
chatId: string,
|
||||
workspaceId: string,
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<AiChatRun | undefined> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
return db
|
||||
.selectFrom('aiChatRuns')
|
||||
.select(this.baseFields)
|
||||
.where('chatId', '=', chatId)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.orderBy('createdAt', 'desc')
|
||||
.orderBy('id', 'desc')
|
||||
.limit(1)
|
||||
.executeTakeFirst();
|
||||
}
|
||||
|
||||
/**
|
||||
* Patch a run by id + workspace; always bumps `updatedAt`. Used for every
|
||||
* lifecycle transition (mark running, link the assistant message, bump
|
||||
* step_count, finalize succeeded/failed/aborted). Returns the updated row or
|
||||
* undefined when nothing matched (e.g. a foreign workspace).
|
||||
*/
|
||||
async update(
|
||||
id: string,
|
||||
workspaceId: string,
|
||||
patch: Partial<{
|
||||
status: string;
|
||||
error: string | null;
|
||||
stepCount: number;
|
||||
assistantMessageId: string | null;
|
||||
stopRequestedAt: Date | null;
|
||||
startedAt: Date | null;
|
||||
finishedAt: Date | null;
|
||||
}>,
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<AiChatRun | undefined> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
return db
|
||||
.updateTable('aiChatRuns')
|
||||
.set({ ...(patch as Record<string, unknown>), updatedAt: new Date() })
|
||||
.where('id', '=', id)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.returning(this.baseFields)
|
||||
.executeTakeFirst();
|
||||
}
|
||||
|
||||
/**
|
||||
* Mark an EXPLICIT stop request on an active run (distinct from a browser
|
||||
* disconnect, which never stops a run). Stamps `stop_requested_at` ONLY while
|
||||
* the run is still active, so a late stop on an already-settled run is a no-op.
|
||||
* Returns the row when a stop was recorded, else undefined (nothing active).
|
||||
*/
|
||||
async markStopRequested(
|
||||
id: string,
|
||||
workspaceId: string,
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<AiChatRun | undefined> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
return db
|
||||
.updateTable('aiChatRuns')
|
||||
.set({ stopRequestedAt: new Date(), updatedAt: new Date() })
|
||||
.where('id', '=', id)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[])
|
||||
.returning(this.baseFields)
|
||||
.executeTakeFirst();
|
||||
}
|
||||
|
||||
/**
|
||||
* Crash-recovery sweep (mirrors AiChatMessageRepo.sweepStreaming): flip every
|
||||
* run still left pending/running — a run whose process died before reaching a
|
||||
* terminal status — to 'aborted', stamping `finished_at`. Returns the number
|
||||
* swept. Workspace-wide on purpose (a crash can dangle runs in any workspace).
|
||||
*
|
||||
* F1 (DECISION C): the BOOT sweep is UNCONDITIONAL — it passes no `staleMs`, so
|
||||
* EVERY dangling run is settled regardless of how recently it was touched. On a
|
||||
* fresh single-process boot any pending|running run is definitionally hung (no
|
||||
* runner is alive to own it), so a fast restart (deploy/OOM within minutes of
|
||||
* the last step) no longer leaves a run stuck 'running' forever — which would
|
||||
* make the one-active-run gate 409 every future turn in that chat.
|
||||
*
|
||||
* The optional `staleMs` window is reintroduced ONLY for the future phase-2
|
||||
* multi-instance timer sweep (see {@link SWEEP_RUN_STALE_MS}): there a booting
|
||||
* replica must NOT abort a run another replica is actively executing, so it
|
||||
* sweeps only runs UNTOUCHED past the window. Phase 1 is single-process, so the
|
||||
* boot path supplies no window.
|
||||
*/
|
||||
async sweepRunning(
|
||||
opts: { staleMs?: number } = {},
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<number> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
const now = new Date();
|
||||
let query = db
|
||||
.updateTable('aiChatRuns')
|
||||
.set({
|
||||
status: 'aborted',
|
||||
finishedAt: now,
|
||||
updatedAt: now,
|
||||
error: sql`coalesce(error, ${'Run interrupted by a server restart.'})`,
|
||||
})
|
||||
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[]);
|
||||
// Multi-instance (phase 2) only: skip runs touched within the window so a
|
||||
// sibling replica's live run is never aborted. Omitted on the phase-1 boot
|
||||
// sweep -> unconditional.
|
||||
if (typeof opts.staleMs === 'number') {
|
||||
const staleBefore = new Date(now.getTime() - opts.staleMs);
|
||||
query = query.where('updatedAt', '<', staleBefore);
|
||||
}
|
||||
const rows = await query.returning('id').execute();
|
||||
return rows.length;
|
||||
}
|
||||
}
|
||||
@@ -657,9 +657,8 @@ export class PagePermissionRepo {
|
||||
pageIds: string[];
|
||||
userId: string;
|
||||
spaceId?: string;
|
||||
workspaceId?: string | null;
|
||||
}): Promise<string[]> {
|
||||
const { pageIds, userId, spaceId, workspaceId } = opts;
|
||||
const { pageIds, userId, spaceId } = opts;
|
||||
if (pageIds.length === 0) return [];
|
||||
|
||||
if (spaceId) {
|
||||
@@ -667,17 +666,6 @@ export class PagePermissionRepo {
|
||||
if (!hasRestrictions) {
|
||||
return pageIds;
|
||||
}
|
||||
} else if (workspaceId) {
|
||||
// #348 — whole-workspace callers (no spaceId: favorites, notifications,
|
||||
// recent, created-by, global search) skip the recursive-ancestor CTE + anti
|
||||
// -join entirely when the workspace has ZERO restricted pages. When any
|
||||
// restriction DOES exist, fall through to the identical CTE below, so
|
||||
// behavior is unchanged whenever restrictions are present.
|
||||
const hasRestrictions =
|
||||
await this.hasRestrictedPagesInWorkspace(workspaceId);
|
||||
if (!hasRestrictions) {
|
||||
return pageIds;
|
||||
}
|
||||
}
|
||||
|
||||
const results = await this.db
|
||||
@@ -915,39 +903,6 @@ export class PagePermissionRepo {
|
||||
return Boolean(result?.exists);
|
||||
}
|
||||
|
||||
/**
|
||||
* Workspace-level analogue of hasRestrictedPagesInSpace: does ANY page in the
|
||||
* whole workspace carry a restriction? Lets whole-workspace access filters
|
||||
* short-circuit the recursive-ancestor CTE when nothing is restricted at all.
|
||||
*
|
||||
* UNCACHED (like the sibling hasRestrictedPagesInSpace) — a single cheap
|
||||
* `EXISTS(pageAccess WHERE workspaceId=?)` per call. This is an ACCESS-CONTROL
|
||||
* gate on whole-workspace list endpoints, so it must never go stale: caching it
|
||||
* (even 5s) reintroduced a leak the space-path never had — a concurrent
|
||||
* whole-workspace read in the insert->commit window of the FIRST restricted page
|
||||
* could re-populate `false` under withCache (read-then-set, no del-during-read
|
||||
* guard) and override the insert bust, leaking that page to unauthorized users
|
||||
* for up to the TTL (#348 review F1). An uncached EXISTS removes both the
|
||||
* cache/DB asymmetry with hasRestrictedPagesInSpace and that race; the space
|
||||
* path already accepts this exact per-call cost.
|
||||
*/
|
||||
async hasRestrictedPagesInWorkspace(workspaceId: string): Promise<boolean> {
|
||||
const result = await this.db
|
||||
.selectNoFrom((eb) =>
|
||||
eb
|
||||
.exists(
|
||||
eb
|
||||
.selectFrom('pageAccess')
|
||||
.select(sql`1`.as('one'))
|
||||
.where('pageAccess.workspaceId', '=', workspaceId),
|
||||
)
|
||||
.as('exists'),
|
||||
)
|
||||
.executeTakeFirst();
|
||||
|
||||
return Boolean(result?.exists);
|
||||
}
|
||||
|
||||
/**
|
||||
* Given a list of parent page IDs, return which ones have at least one accessible child.
|
||||
* Efficient batch query for sidebar hasChildren calculation.
|
||||
|
||||
@@ -581,9 +581,6 @@ export class PageRepo {
|
||||
const query = this.db
|
||||
.selectFrom('pages')
|
||||
.select(this.baseFields)
|
||||
// NOTE: `content` IS needed here — the trash UI reads page.content to render
|
||||
// the deleted-page preview modal (trash.tsx handlePageClick ->
|
||||
// TrashPageContentModal pageContent). Do NOT drop it (see #348 review F3).
|
||||
.select('content')
|
||||
.select((eb) => this.withSpace(eb))
|
||||
.select((eb) => this.withDeletedBy(eb))
|
||||
|
||||
@@ -1,6 +1,4 @@
|
||||
import { Inject, Injectable } from '@nestjs/common';
|
||||
import { CACHE_MANAGER } from '@nestjs/cache-manager';
|
||||
import { Cache } from 'cache-manager';
|
||||
import { Injectable } from '@nestjs/common';
|
||||
import { InjectKysely } from 'nestjs-kysely';
|
||||
import { KyselyDB, KyselyTransaction } from '../../types/kysely.types';
|
||||
import { dbOrTx } from '../../utils';
|
||||
@@ -11,7 +9,6 @@ import {
|
||||
} from '@docmost/db/types/entity.types';
|
||||
import { ExpressionBuilder, sql } from 'kysely';
|
||||
import { DB, Workspaces } from '@docmost/db/types/db';
|
||||
import { CacheKey } from '../../../common/helpers/cache-keys';
|
||||
|
||||
/**
|
||||
* Writable `settings.ai.provider` keys, enforced at this generic SQL layer. This
|
||||
@@ -64,34 +61,7 @@ export class WorkspaceRepo {
|
||||
'temporaryNoteHours',
|
||||
'isScimEnabled',
|
||||
];
|
||||
constructor(
|
||||
@InjectKysely() private readonly db: KyselyDB,
|
||||
@Inject(CACHE_MANAGER) private readonly cacheManager: Cache,
|
||||
) {}
|
||||
|
||||
/**
|
||||
* #348 — bust the DomainMiddleware workspace caches after any workspace write.
|
||||
* Deletes BOTH the self-hosted (constant) key and the cloud per-hostname key so
|
||||
* a single implementation covers either deployment mode (the irrelevant key is a
|
||||
* harmless no-op). Best-effort: a cache error must never fail the write, and a
|
||||
* missed bust is bounded by WORKSPACE_CACHE_TTL_MS. Note: a hostname RENAME only
|
||||
* busts the NEW hostname's key (the row returned here carries the new hostname);
|
||||
* the old key expires via TTL.
|
||||
*/
|
||||
private async bustWorkspaceCache(
|
||||
workspace?: Pick<Workspace, 'hostname'> | undefined,
|
||||
): Promise<void> {
|
||||
try {
|
||||
await this.cacheManager.del(CacheKey.WORKSPACE_SELF_HOSTED);
|
||||
if (workspace?.hostname) {
|
||||
await this.cacheManager.del(
|
||||
CacheKey.WORKSPACE_BY_HOST(workspace.hostname),
|
||||
);
|
||||
}
|
||||
} catch {
|
||||
// cache is best-effort; TTL is the backstop
|
||||
}
|
||||
}
|
||||
constructor(@InjectKysely() private readonly db: KyselyDB) {}
|
||||
|
||||
async findById(
|
||||
workspaceId: string,
|
||||
@@ -174,14 +144,12 @@ export class WorkspaceRepo {
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<Workspace> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
const workspace = await db
|
||||
return db
|
||||
.updateTable('workspaces')
|
||||
.set({ ...updatableWorkspace, updatedAt: new Date() })
|
||||
.where('id', '=', workspaceId)
|
||||
.returning(this.baseFields)
|
||||
.executeTakeFirst();
|
||||
await this.bustWorkspaceCache(workspace);
|
||||
return workspace;
|
||||
}
|
||||
|
||||
async insertWorkspace(
|
||||
@@ -189,14 +157,11 @@ export class WorkspaceRepo {
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<Workspace> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
const workspace = await db
|
||||
return db
|
||||
.insertInto('workspaces')
|
||||
.values(insertableWorkspace)
|
||||
.returning(this.baseFields)
|
||||
.executeTakeFirst();
|
||||
// Bust the cached "not found" so a fresh install / new tenant is seen at once.
|
||||
await this.bustWorkspaceCache(workspace);
|
||||
return workspace;
|
||||
}
|
||||
|
||||
async count(): Promise<number> {
|
||||
@@ -238,7 +203,7 @@ export class WorkspaceRepo {
|
||||
trx?: KyselyTransaction,
|
||||
) {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
const workspace = await db
|
||||
return db
|
||||
.updateTable('workspaces')
|
||||
.set({
|
||||
settings: sql`COALESCE(settings, '{}'::jsonb)
|
||||
@@ -249,8 +214,6 @@ export class WorkspaceRepo {
|
||||
.where('id', '=', workspaceId)
|
||||
.returning(this.baseFields)
|
||||
.executeTakeFirst();
|
||||
await this.bustWorkspaceCache(workspace);
|
||||
return workspace;
|
||||
}
|
||||
|
||||
async updateAiSettings(
|
||||
@@ -260,7 +223,7 @@ export class WorkspaceRepo {
|
||||
trx?: KyselyTransaction,
|
||||
) {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
const workspace = await db
|
||||
return db
|
||||
.updateTable('workspaces')
|
||||
.set({
|
||||
settings: sql`COALESCE(settings, '{}'::jsonb)
|
||||
@@ -271,8 +234,6 @@ export class WorkspaceRepo {
|
||||
.where('id', '=', workspaceId)
|
||||
.returning(this.baseFields)
|
||||
.executeTakeFirst();
|
||||
await this.bustWorkspaceCache(workspace);
|
||||
return workspace;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -311,7 +272,7 @@ export class WorkspaceRepo {
|
||||
entries.flatMap(([k, v]) => [sql.lit(k), sql`${v}::text`]),
|
||||
)})`
|
||||
: sql`'{}'::jsonb`;
|
||||
const workspace = await db
|
||||
return db
|
||||
.updateTable('workspaces')
|
||||
.set({
|
||||
settings: sql`COALESCE(settings, '{}'::jsonb) || jsonb_build_object(
|
||||
@@ -326,8 +287,6 @@ export class WorkspaceRepo {
|
||||
.where('id', '=', workspaceId)
|
||||
.returning(this.baseFields)
|
||||
.executeTakeFirst();
|
||||
await this.bustWorkspaceCache(workspace);
|
||||
return workspace;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -344,7 +303,7 @@ export class WorkspaceRepo {
|
||||
trx?: KyselyTransaction,
|
||||
) {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
const workspace = await db
|
||||
return db
|
||||
.updateTable('workspaces')
|
||||
.set({
|
||||
settings: sql`COALESCE(settings, '{}'::jsonb)
|
||||
@@ -354,8 +313,6 @@ export class WorkspaceRepo {
|
||||
.where('id', '=', workspaceId)
|
||||
.returning(this.baseFields)
|
||||
.executeTakeFirst();
|
||||
await this.bustWorkspaceCache(workspace);
|
||||
return workspace;
|
||||
}
|
||||
|
||||
async updateSharingSettings(
|
||||
@@ -365,7 +322,7 @@ export class WorkspaceRepo {
|
||||
trx?: KyselyTransaction,
|
||||
) {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
const workspace = await db
|
||||
return db
|
||||
.updateTable('workspaces')
|
||||
.set({
|
||||
settings: sql`COALESCE(settings, '{}'::jsonb)
|
||||
@@ -376,8 +333,6 @@ export class WorkspaceRepo {
|
||||
.where('id', '=', workspaceId)
|
||||
.returning(this.baseFields)
|
||||
.executeTakeFirst();
|
||||
await this.bustWorkspaceCache(workspace);
|
||||
return workspace;
|
||||
}
|
||||
|
||||
async updateTemplateSettings(
|
||||
@@ -387,7 +342,7 @@ export class WorkspaceRepo {
|
||||
trx?: KyselyTransaction,
|
||||
) {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
const workspace = await db
|
||||
return db
|
||||
.updateTable('workspaces')
|
||||
.set({
|
||||
settings: sql`COALESCE(settings, '{}'::jsonb)
|
||||
@@ -398,8 +353,6 @@ export class WorkspaceRepo {
|
||||
.where('id', '=', workspaceId)
|
||||
.returning(this.baseFields)
|
||||
.executeTakeFirst();
|
||||
await this.bustWorkspaceCache(workspace);
|
||||
return workspace;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
+30
@@ -659,6 +659,35 @@ export interface AiChatMessages {
|
||||
deletedAt: Timestamp | null;
|
||||
}
|
||||
|
||||
// The agent RUN as a first-class server-side lifecycle object (#184 phase 1).
|
||||
// Mirrors migration 20260704T130000-ai-chat-runs.ts. A run is created when an
|
||||
// agent turn starts and survives the browser disconnecting; the DB is the source
|
||||
// of truth a later client reconnects to. `assistantMessageId` links to the #183
|
||||
// projection row (the assistant message this run materializes).
|
||||
export interface AiChatRuns {
|
||||
id: Generated<string>;
|
||||
chatId: string;
|
||||
workspaceId: string;
|
||||
// SET NULL on user deletion (the run history outlives its author); also NULL
|
||||
// for a future non-human trigger (cron/api).
|
||||
createdBy: string | null;
|
||||
// The assistant message this run materializes; SET NULL if it is pruned.
|
||||
assistantMessageId: string | null;
|
||||
// 'user' | 'autostart' | 'schedule' | 'api' | 'continue' (only 'user' is
|
||||
// produced in phase 1; the rest are reserved for the deferred autonomy triggers).
|
||||
trigger: Generated<string>;
|
||||
// 'pending' | 'running' | 'succeeded' | 'failed' | 'aborted'.
|
||||
status: Generated<string>;
|
||||
error: string | null;
|
||||
stepCount: Generated<number>;
|
||||
// Set when an EXPLICIT user stop is requested (distinct from a disconnect).
|
||||
stopRequestedAt: Timestamp | null;
|
||||
startedAt: Timestamp | null;
|
||||
finishedAt: Timestamp | null;
|
||||
createdAt: Generated<Timestamp>;
|
||||
updatedAt: Generated<Timestamp>;
|
||||
}
|
||||
|
||||
// Per-(chat,page) snapshot of the open page's Markdown at the END of the agent's
|
||||
// previous turn (#274). Mirrors migration 20260702T120000-ai-chat-page-snapshot.ts.
|
||||
// The next turn diffs the CURRENT Markdown against `contentMd` to surface edits a
|
||||
@@ -695,6 +724,7 @@ export interface DB {
|
||||
aiAgentRoles: AiAgentRoles;
|
||||
aiChats: AiChats;
|
||||
aiChatMessages: AiChatMessages;
|
||||
aiChatRuns: AiChatRuns;
|
||||
aiChatPageSnapshots: AiChatPageSnapshots;
|
||||
apiKeys: ApiKeys;
|
||||
attachments: Attachments;
|
||||
|
||||
@@ -3,6 +3,7 @@ import {
|
||||
AiAgentRoles,
|
||||
AiChats,
|
||||
AiChatMessages,
|
||||
AiChatRuns,
|
||||
AiChatPageSnapshots,
|
||||
Attachments,
|
||||
Comments,
|
||||
@@ -56,10 +57,12 @@ export type UpdatableAiChat = Updateable<Omit<AiChats, 'id'>>;
|
||||
// full-text search. It is omitted from the public type so it never leaks
|
||||
// into HTTP responses or the chat history fed to the language model.
|
||||
export type AiChatMessage = Omit<Selectable<AiChatMessages>, 'tsv'>;
|
||||
export type InsertableAiChatMessage = Omit<
|
||||
Insertable<AiChatMessages>,
|
||||
'tsv'
|
||||
>;
|
||||
export type InsertableAiChatMessage = Omit<Insertable<AiChatMessages>, 'tsv'>;
|
||||
|
||||
// AI Chat Run (#184 phase 1): the agent run as a first-class lifecycle object,
|
||||
// detached from the HTTP request / browser window.
|
||||
export type AiChatRun = Selectable<AiChatRuns>;
|
||||
export type InsertableAiChatRun = Insertable<AiChatRuns>;
|
||||
|
||||
// AI Chat Page Snapshot (#274): per-(chat,page) Markdown snapshot taken at the
|
||||
// end of the agent's previous turn, diffed against the current page next turn to
|
||||
@@ -214,11 +217,14 @@ export type UpdatableFavorite = Updateable<Omit<Favorites, 'id'>>;
|
||||
// Page Transclusion
|
||||
export type PageTransclusion = Selectable<PageTransclusions>;
|
||||
export type InsertablePageTransclusion = Insertable<PageTransclusions>;
|
||||
export type UpdatablePageTransclusion = Updateable<Omit<PageTransclusions, 'id'>>;
|
||||
export type UpdatablePageTransclusion = Updateable<
|
||||
Omit<PageTransclusions, 'id'>
|
||||
>;
|
||||
|
||||
// Page Transclusion Reference
|
||||
export type PageTransclusionReference = Selectable<PageTransclusionReferences>;
|
||||
export type InsertablePageTransclusionReference = Insertable<PageTransclusionReferences>;
|
||||
export type InsertablePageTransclusionReference =
|
||||
Insertable<PageTransclusionReferences>;
|
||||
export type UpdatablePageTransclusionReference = Updateable<
|
||||
Omit<PageTransclusionReferences, 'id'>
|
||||
>;
|
||||
@@ -288,7 +294,9 @@ export type UpdatablePagePermission = Updateable<Omit<_PagePermissions, 'id'>>;
|
||||
// Page Verification
|
||||
export type PageVerification = Selectable<_PageVerifications>;
|
||||
export type InsertablePageVerification = Insertable<_PageVerifications>;
|
||||
export type UpdatablePageVerification = Updateable<Omit<_PageVerifications, 'id'>>;
|
||||
export type UpdatablePageVerification = Updateable<
|
||||
Omit<_PageVerifications, 'id'>
|
||||
>;
|
||||
|
||||
// Page Verifier
|
||||
export type PageVerifier = Selectable<_PageVerifiers>;
|
||||
|
||||
@@ -0,0 +1,304 @@
|
||||
import { Kysely } from 'kysely';
|
||||
import {
|
||||
AiChatRunRepo,
|
||||
SWEEP_RUN_STALE_MS,
|
||||
} from '@docmost/db/repos/ai-chat/ai-chat-run.repo';
|
||||
import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
|
||||
import { AiChatRunService } from '../../src/core/ai-chat/ai-chat-run.service';
|
||||
import {
|
||||
getTestDb,
|
||||
destroyTestDb,
|
||||
createWorkspace,
|
||||
createUser,
|
||||
createChat,
|
||||
} from './db';
|
||||
|
||||
/**
|
||||
* Integration coverage for the #184 phase-1 durable agent run: real SQL against
|
||||
* docmost_test. Proves the core invariant primitives — a run is a first-class
|
||||
* lifecycle row, at most one is active per chat, a detached run's progress
|
||||
* survives with NO subscriber, an explicit stop settles it as aborted, a
|
||||
* reconnect read returns the persisted state, and a crash sweep recovers
|
||||
* dangling runs.
|
||||
*/
|
||||
describe('AiChatRun durable lifecycle [integration]', () => {
|
||||
let db: Kysely<any>;
|
||||
let runRepo: AiChatRunRepo;
|
||||
let messageRepo: AiChatMessageRepo;
|
||||
let service: AiChatRunService;
|
||||
let workspaceId: string;
|
||||
let otherWorkspaceId: string;
|
||||
let userId: string;
|
||||
let chatId: string;
|
||||
|
||||
beforeAll(async () => {
|
||||
db = getTestDb();
|
||||
runRepo = new AiChatRunRepo(db as any);
|
||||
messageRepo = new AiChatMessageRepo(db as any);
|
||||
// Boot-sweep isn't triggered here; the isCloud stub is all the service needs
|
||||
// for these direct-call integration cases (F7).
|
||||
service = new AiChatRunService(runRepo, { isCloud: () => false } as never);
|
||||
workspaceId = (await createWorkspace(db)).id;
|
||||
otherWorkspaceId = (await createWorkspace(db)).id;
|
||||
userId = (await createUser(db, workspaceId)).id;
|
||||
chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await destroyTestDb();
|
||||
});
|
||||
|
||||
// Each test that creates an active run settles it (or uses its own chat) so the
|
||||
// partial unique index does not bleed across tests.
|
||||
|
||||
it('insert + findById round-trips a run row, defaulting status/trigger', async () => {
|
||||
const run = await runRepo.insert({
|
||||
chatId,
|
||||
workspaceId,
|
||||
createdBy: userId,
|
||||
});
|
||||
expect(run.status).toBe('pending');
|
||||
expect(run.trigger).toBe('user');
|
||||
expect(run.stepCount).toBe(0);
|
||||
|
||||
const found = await runRepo.findById(run.id, workspaceId);
|
||||
expect(found!.id).toBe(run.id);
|
||||
// Workspace-scoped: a foreign workspace sees nothing.
|
||||
expect(await runRepo.findById(run.id, otherWorkspaceId)).toBeUndefined();
|
||||
|
||||
// settle so it does not occupy the active slot
|
||||
await runRepo.update(run.id, workspaceId, {
|
||||
status: 'succeeded',
|
||||
finishedAt: new Date(),
|
||||
});
|
||||
});
|
||||
|
||||
it('enforces ONE ACTIVE run per chat (partial unique index rejects a second)', async () => {
|
||||
const activeChat = (
|
||||
await createChat(db, { workspaceId, creatorId: userId })
|
||||
).id;
|
||||
const first = await runRepo.insert({
|
||||
chatId: activeChat,
|
||||
workspaceId,
|
||||
createdBy: userId,
|
||||
status: 'running',
|
||||
});
|
||||
// A second pending/running run on the SAME chat must be rejected by the DB.
|
||||
await expect(
|
||||
runRepo.insert({
|
||||
chatId: activeChat,
|
||||
workspaceId,
|
||||
createdBy: userId,
|
||||
status: 'running',
|
||||
}),
|
||||
).rejects.toThrow();
|
||||
|
||||
// findActiveByChat returns exactly the one active run.
|
||||
const active = await runRepo.findActiveByChat(activeChat, workspaceId);
|
||||
expect(active!.id).toBe(first.id);
|
||||
|
||||
// Once it settles, the slot frees and a new run may start.
|
||||
await runRepo.update(first.id, workspaceId, {
|
||||
status: 'succeeded',
|
||||
finishedAt: new Date(),
|
||||
});
|
||||
expect(
|
||||
await runRepo.findActiveByChat(activeChat, workspaceId),
|
||||
).toBeUndefined();
|
||||
const second = await runRepo.insert({
|
||||
chatId: activeChat,
|
||||
workspaceId,
|
||||
createdBy: userId,
|
||||
status: 'running',
|
||||
});
|
||||
expect(second.id).not.toBe(first.id);
|
||||
await runRepo.update(second.id, workspaceId, {
|
||||
status: 'aborted',
|
||||
finishedAt: new Date(),
|
||||
});
|
||||
});
|
||||
|
||||
it('DETACHED run: persists + finalizes succeeded with NO subscriber, reconnect returns state', async () => {
|
||||
// A dedicated chat so the active-run slot is clean.
|
||||
const runChat = (
|
||||
await createChat(db, { workspaceId, creatorId: userId })
|
||||
).id;
|
||||
|
||||
// beginRun = the runner starts the turn (registers an in-memory controller).
|
||||
const handle = await service.beginRun({
|
||||
chatId: runChat,
|
||||
workspaceId,
|
||||
userId,
|
||||
});
|
||||
expect(handle.signal.aborted).toBe(false);
|
||||
expect(service.isLocallyActive(handle.runId)).toBe(true);
|
||||
|
||||
// The assistant projection row (#183) is seeded + linked.
|
||||
const seeded = await messageRepo.insert({
|
||||
chatId: runChat,
|
||||
workspaceId,
|
||||
userId,
|
||||
role: 'assistant',
|
||||
content: '',
|
||||
status: 'streaming',
|
||||
metadata: { parts: [] } as never,
|
||||
});
|
||||
await service.linkAssistantMessage(handle.runId, workspaceId, seeded.id);
|
||||
|
||||
// Progress is persisted as steps finish — NO HTTP socket involved here at all.
|
||||
await service.recordStep(handle.runId, workspaceId, 1);
|
||||
await messageRepo.update(seeded.id, workspaceId, {
|
||||
content: 'partial work',
|
||||
metadata: { parts: [{ type: 'text', text: 'partial work' }] },
|
||||
});
|
||||
|
||||
// The turn completes; finalize the projection then the run.
|
||||
await messageRepo.update(seeded.id, workspaceId, {
|
||||
content: 'final answer',
|
||||
status: 'completed',
|
||||
});
|
||||
await service.finalizeRun(handle.runId, workspaceId, 'completed');
|
||||
|
||||
expect(service.isLocallyActive(handle.runId)).toBe(false);
|
||||
|
||||
// Reconnect: the latest run for the chat + its projected message, from the DB.
|
||||
const run = await service.getLatestForChat(runChat, workspaceId);
|
||||
expect(run!.status).toBe('succeeded');
|
||||
expect(run!.stepCount).toBe(1);
|
||||
expect(run!.assistantMessageId).toBe(seeded.id);
|
||||
expect(run!.finishedAt).toBeTruthy();
|
||||
const message = await messageRepo.findById(seeded.id, workspaceId);
|
||||
expect(message!.status).toBe('completed');
|
||||
expect(message!.content).toBe('final answer');
|
||||
});
|
||||
|
||||
it('EXPLICIT stop aborts the run signal, marks the row, and settles as aborted', async () => {
|
||||
const runChat = (
|
||||
await createChat(db, { workspaceId, creatorId: userId })
|
||||
).id;
|
||||
const handle = await service.beginRun({
|
||||
chatId: runChat,
|
||||
workspaceId,
|
||||
userId,
|
||||
});
|
||||
|
||||
// User presses Stop.
|
||||
const stopped = await service.requestStop(handle.runId, workspaceId);
|
||||
expect(stopped).toBe(true);
|
||||
expect(handle.signal.aborted).toBe(true);
|
||||
|
||||
// The row carries the stop request (distinct from a disconnect, which would
|
||||
// leave stop_requested_at NULL).
|
||||
const afterStop = await runRepo.findById(handle.runId, workspaceId);
|
||||
expect(afterStop!.stopRequestedAt).toBeTruthy();
|
||||
|
||||
// The terminal callback (onAbort) settles the run.
|
||||
await service.finalizeRun(handle.runId, workspaceId, 'aborted');
|
||||
const run = await service.getLatestForChat(runChat, workspaceId);
|
||||
expect(run!.status).toBe('aborted');
|
||||
});
|
||||
|
||||
it('markStopRequested is a no-op on an already-settled run (returns undefined)', async () => {
|
||||
const runChat = (
|
||||
await createChat(db, { workspaceId, creatorId: userId })
|
||||
).id;
|
||||
const run = await runRepo.insert({
|
||||
chatId: runChat,
|
||||
workspaceId,
|
||||
createdBy: userId,
|
||||
status: 'running',
|
||||
});
|
||||
await runRepo.update(run.id, workspaceId, {
|
||||
status: 'succeeded',
|
||||
finishedAt: new Date(),
|
||||
});
|
||||
const marked = await runRepo.markStopRequested(run.id, workspaceId);
|
||||
expect(marked).toBeUndefined();
|
||||
});
|
||||
|
||||
it('sweepRunning aborts STALE dangling runs but not fresh or settled ones', async () => {
|
||||
const sweepChat1 = (
|
||||
await createChat(db, { workspaceId, creatorId: userId })
|
||||
).id;
|
||||
const sweepChat2 = (
|
||||
await createChat(db, { workspaceId, creatorId: userId })
|
||||
).id;
|
||||
const sweepChat3 = (
|
||||
await createChat(db, { workspaceId, creatorId: userId })
|
||||
).id;
|
||||
|
||||
const stale = await runRepo.insert({
|
||||
chatId: sweepChat1,
|
||||
workspaceId,
|
||||
createdBy: userId,
|
||||
status: 'running',
|
||||
});
|
||||
const fresh = await runRepo.insert({
|
||||
chatId: sweepChat2,
|
||||
workspaceId,
|
||||
createdBy: userId,
|
||||
status: 'running',
|
||||
});
|
||||
const settled = await runRepo.insert({
|
||||
chatId: sweepChat3,
|
||||
workspaceId,
|
||||
createdBy: userId,
|
||||
status: 'running',
|
||||
});
|
||||
await runRepo.update(settled.id, workspaceId, {
|
||||
status: 'succeeded',
|
||||
finishedAt: new Date(),
|
||||
});
|
||||
// Backdate the stale run's updatedAt past the 10-minute staleness window.
|
||||
await db
|
||||
.updateTable('aiChatRuns')
|
||||
.set({ updatedAt: new Date(Date.now() - 20 * 60 * 1000) })
|
||||
.where('id', '=', stale.id)
|
||||
.execute();
|
||||
|
||||
// WINDOWED sweep (phase-2 multi-instance timer path): only runs older than the
|
||||
// staleness window are aborted, so a sibling replica's fresh run survives. The
|
||||
// no-arg boot sweep (variant C) is unconditional — covered separately below.
|
||||
const swept = await runRepo.sweepRunning({ staleMs: SWEEP_RUN_STALE_MS });
|
||||
expect(swept).toBeGreaterThanOrEqual(1);
|
||||
|
||||
expect((await runRepo.findById(stale.id, workspaceId))!.status).toBe(
|
||||
'aborted',
|
||||
);
|
||||
// Fresh (recently-updated) running run survives the WINDOWED sweep — a sibling
|
||||
// replica may still be executing it.
|
||||
expect((await runRepo.findById(fresh.id, workspaceId))!.status).toBe(
|
||||
'running',
|
||||
);
|
||||
expect((await runRepo.findById(settled.id, workspaceId))!.status).toBe(
|
||||
'succeeded',
|
||||
);
|
||||
|
||||
// cleanup active fresh run
|
||||
await runRepo.update(fresh.id, workspaceId, {
|
||||
status: 'aborted',
|
||||
finishedAt: new Date(),
|
||||
});
|
||||
});
|
||||
|
||||
it('sweepRunning() with NO args (boot sweep / variant C) aborts even a FRESH running run', async () => {
|
||||
// F1/DECISION C at the SQL level: the unconditional boot sweep has NO
|
||||
// staleness window, so a run updated just now (a fast restart) is settled too
|
||||
// — otherwise it would stay 'running' forever and 409 every future turn.
|
||||
const bootChat = (
|
||||
await createChat(db, { workspaceId, creatorId: userId })
|
||||
).id;
|
||||
const fresh = await runRepo.insert({
|
||||
chatId: bootChat,
|
||||
workspaceId,
|
||||
createdBy: userId,
|
||||
status: 'running',
|
||||
});
|
||||
// updatedAt = now (fresh, untouched). The no-arg sweep settles it anyway.
|
||||
const swept = await runRepo.sweepRunning();
|
||||
expect(swept).toBeGreaterThanOrEqual(1);
|
||||
expect((await runRepo.findById(fresh.id, workspaceId))!.status).toBe(
|
||||
'aborted',
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -1,113 +0,0 @@
|
||||
import { Kysely } from 'kysely';
|
||||
import { PagePermissionRepo } from '@docmost/db/repos/page/page-permission.repo';
|
||||
import { GroupRepo } from '@docmost/db/repos/group/group.repo';
|
||||
import {
|
||||
getTestDb,
|
||||
destroyTestDb,
|
||||
createWorkspace,
|
||||
createSpace,
|
||||
createUser,
|
||||
createPage,
|
||||
} from './db';
|
||||
|
||||
/**
|
||||
* #348 — the whole-workspace access-filter short-circuit is an ACCESS-CONTROL
|
||||
* path, so it must produce the SAME result as the full recursive-ancestor CTE.
|
||||
*
|
||||
* filterAccessiblePageIds({ workspaceId }) (no spaceId — the favorites /
|
||||
* notifications / recent / created-by / global-search callers) skips the CTE only
|
||||
* when the workspace has ZERO restricted pages. A page is "restricted &
|
||||
* inaccessible" when it (or an ancestor) has a `pageAccess` row and the user has
|
||||
* no matching `pagePermissions`. Driven against real Postgres, asserts:
|
||||
* 1. zero restrictions -> short-circuit returns the full input set;
|
||||
* 2. a restriction present -> the CTE runs and drops the page the user can't
|
||||
* reach while keeping the reachable ones (behavior unchanged);
|
||||
* 3. inserting the FIRST pageAccess flips hasRestrictedPagesInWorkspace
|
||||
* false -> true immediately (the 0->1 transition — now uncached, no stale
|
||||
* window, review F1); it is scoped per workspace.
|
||||
*/
|
||||
describe('#348 filterAccessiblePageIds workspace short-circuit (real PG)', () => {
|
||||
let db: Kysely<any>;
|
||||
let repo: PagePermissionRepo;
|
||||
let workspaceId: string;
|
||||
let otherWorkspaceId: string;
|
||||
let userId: string;
|
||||
let spaceId: string;
|
||||
|
||||
beforeAll(async () => {
|
||||
db = getTestDb();
|
||||
// hasRestrictedPagesInWorkspace is now uncached, and no other cached
|
||||
// permission path is exercised here, so a no-op cache stub suffices.
|
||||
const cacheStub = {
|
||||
get: async () => undefined,
|
||||
set: async () => undefined,
|
||||
del: async () => undefined,
|
||||
} as never;
|
||||
repo = new PagePermissionRepo(db, new GroupRepo(db), cacheStub);
|
||||
|
||||
const ws = await createWorkspace(db);
|
||||
workspaceId = ws.id;
|
||||
const other = await createWorkspace(db);
|
||||
otherWorkspaceId = other.id;
|
||||
const user = await createUser(db, workspaceId);
|
||||
userId = user.id;
|
||||
const space = await createSpace(db, workspaceId);
|
||||
spaceId = space.id;
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await destroyTestDb();
|
||||
});
|
||||
|
||||
it('zero restrictions: short-circuit returns the full input set', async () => {
|
||||
const p1 = await createPage(db, { workspaceId, spaceId });
|
||||
const p2 = await createPage(db, { workspaceId, spaceId });
|
||||
|
||||
expect(await repo.hasRestrictedPagesInWorkspace(workspaceId)).toBe(false);
|
||||
|
||||
const ids = [p1.id, p2.id];
|
||||
const filtered = await repo.filterAccessiblePageIds({
|
||||
pageIds: ids,
|
||||
userId,
|
||||
workspaceId,
|
||||
});
|
||||
expect(new Set(filtered)).toEqual(new Set(ids));
|
||||
});
|
||||
|
||||
it('a restriction present: filters out the page the user cannot reach', async () => {
|
||||
const openPage = await createPage(db, { workspaceId, spaceId });
|
||||
const restrictedPage = await createPage(db, { workspaceId, spaceId });
|
||||
|
||||
// Add a pageAccess row on restrictedPage with NO matching pagePermissions for
|
||||
// `userId` → the CTE anti-join marks it inaccessible for this user.
|
||||
await db
|
||||
.insertInto('pageAccess')
|
||||
.values({
|
||||
pageId: restrictedPage.id,
|
||||
workspaceId,
|
||||
spaceId,
|
||||
accessLevel: 'read',
|
||||
creatorId: userId,
|
||||
})
|
||||
.execute();
|
||||
|
||||
// 0->1 transition is reflected immediately (uncached).
|
||||
expect(await repo.hasRestrictedPagesInWorkspace(workspaceId)).toBe(true);
|
||||
|
||||
const filtered = await repo.filterAccessiblePageIds({
|
||||
pageIds: [openPage.id, restrictedPage.id],
|
||||
userId,
|
||||
workspaceId,
|
||||
});
|
||||
expect(filtered).toContain(openPage.id);
|
||||
expect(filtered).not.toContain(restrictedPage.id);
|
||||
});
|
||||
|
||||
it('hasRestrictedPagesInWorkspace is scoped per workspace', async () => {
|
||||
// The other workspace has no pageAccess rows → still false, unaffected by the
|
||||
// restriction added above in `workspaceId`.
|
||||
expect(await repo.hasRestrictedPagesInWorkspace(otherWorkspaceId)).toBe(
|
||||
false,
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -1,25 +1,7 @@
|
||||
import { Kysely } from 'kysely';
|
||||
import { WorkspaceRepo } from '@docmost/db/repos/workspace/workspace.repo';
|
||||
import { CacheKey } from 'src/common/helpers/cache-keys';
|
||||
import { getTestDb, destroyTestDb, createWorkspace } from './db';
|
||||
|
||||
// A minimal Map-backed cache double with a working `del` (the previous `{}` stub
|
||||
// made bustWorkspaceCache's `del` throw into its own try/catch, so the #348
|
||||
// invalidation was never actually exercised — review F6).
|
||||
function makeCacheDouble() {
|
||||
const store = new Map<string, unknown>();
|
||||
return {
|
||||
store,
|
||||
get: async (k: string) => store.get(k),
|
||||
set: async (k: string, v: unknown) => {
|
||||
store.set(k, v);
|
||||
},
|
||||
del: async (k: string) => {
|
||||
store.delete(k);
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* A — WorkspaceRepo.updateSetting jsonb-MERGE (the html-embed kill-switch
|
||||
* write-half). Setting a single top-level key must NOT clobber sibling
|
||||
@@ -33,9 +15,7 @@ describe('WorkspaceRepo.updateSetting (jsonb merge) [integration]', () => {
|
||||
beforeAll(() => {
|
||||
db = getTestDb();
|
||||
// Repos are plain classes taking @InjectKysely() db — instantiate directly.
|
||||
// 2nd arg is CACHE_MANAGER (used only to bust the #348 workspace cache); a
|
||||
// stub is fine here since bustWorkspaceCache is best-effort (try/catch).
|
||||
repo = new WorkspaceRepo(db as any, {} as any);
|
||||
repo = new WorkspaceRepo(db as any);
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
@@ -78,62 +58,3 @@ describe('WorkspaceRepo.updateSetting (jsonb merge) [integration]', () => {
|
||||
expect(updated.settings).toEqual({ htmlEmbed: false });
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* #348 F6 — the DomainMiddleware workspace cache (WORKSPACE_SELF_HOSTED /
|
||||
* WORKSPACE_BY_HOST, 15s TTL) caches security-relevant fields (enforceSso/
|
||||
* enforceMfa/status). Its correctness rests entirely on bustWorkspaceCache being
|
||||
* called from every mutator. This exercises the real invalidation with a working
|
||||
* cache double (not the {} stub, whose del throws-and-swallows): warm the cache
|
||||
* like DomainMiddleware, mutate, and assert the busted key is gone so a stale
|
||||
* workspace row can't outlive the mutation.
|
||||
*/
|
||||
describe('WorkspaceRepo bustWorkspaceCache invalidation [integration]', () => {
|
||||
let db: Kysely<any>;
|
||||
|
||||
beforeAll(() => {
|
||||
db = getTestDb();
|
||||
});
|
||||
afterAll(async () => {
|
||||
await destroyTestDb();
|
||||
});
|
||||
|
||||
it('updateSetting busts the self-hosted workspace cache key', async () => {
|
||||
const cache = makeCacheDouble();
|
||||
const repo = new WorkspaceRepo(db as any, cache as any);
|
||||
const ws = await createWorkspace(db, { settings: {} });
|
||||
|
||||
// Warm the cache as DomainMiddleware would (self-hosted key).
|
||||
cache.store.set(CacheKey.WORKSPACE_SELF_HOSTED, ws);
|
||||
expect(cache.store.has(CacheKey.WORKSPACE_SELF_HOSTED)).toBe(true);
|
||||
|
||||
await repo.updateSetting(ws.id, 'htmlEmbed', true);
|
||||
|
||||
// The mutation must have invalidated the cached row.
|
||||
expect(cache.store.has(CacheKey.WORKSPACE_SELF_HOSTED)).toBe(false);
|
||||
});
|
||||
|
||||
it('updateSharingSettings busts the by-host workspace cache key too', async () => {
|
||||
const cache = makeCacheDouble();
|
||||
const repo = new WorkspaceRepo(db as any, cache as any);
|
||||
const ws = await createWorkspace(db, { settings: {} });
|
||||
// createWorkspace assigns a unique hostname; read it back for the by-host key.
|
||||
const { hostname } = await db
|
||||
.selectFrom('workspaces')
|
||||
.select(['hostname'])
|
||||
.where('id', '=', ws.id)
|
||||
.executeTakeFirstOrThrow();
|
||||
|
||||
// Warm BOTH keys (self-hosted + by-host); the by-host bust needs the row's
|
||||
// hostname, which the mutator returns from the DB.
|
||||
cache.store.set(CacheKey.WORKSPACE_SELF_HOSTED, ws);
|
||||
cache.store.set(CacheKey.WORKSPACE_BY_HOST(hostname as string), ws);
|
||||
|
||||
await repo.updateSharingSettings(ws.id, 'allowInvite', true);
|
||||
|
||||
expect(cache.store.has(CacheKey.WORKSPACE_SELF_HOSTED)).toBe(false);
|
||||
expect(cache.store.has(CacheKey.WORKSPACE_BY_HOST(hostname as string))).toBe(
|
||||
false,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -20,6 +20,7 @@ export * from "./lib/html-embed/html-embed";
|
||||
export * from "./lib/mention";
|
||||
export * from "./lib/markdown";
|
||||
export * from "./lib/search-and-replace";
|
||||
export * from "./lib/multi-cursor";
|
||||
export * from "./lib/embed-provider";
|
||||
export * from "./lib/subpages";
|
||||
export * from "./lib/transclusion";
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
import { MultiCursor } from "./multi-cursor";
|
||||
export * from "./multi-cursor";
|
||||
export default MultiCursor;
|
||||
@@ -0,0 +1,453 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { Editor } from "@tiptap/core";
|
||||
import { Document } from "@tiptap/extension-document";
|
||||
import { Paragraph } from "@tiptap/extension-paragraph";
|
||||
import { Text } from "@tiptap/extension-text";
|
||||
import { Bold } from "@tiptap/extension-bold";
|
||||
import { Node as PMNode } from "@tiptap/pm/model";
|
||||
import { MultiCursor, multiCursorPluginKey, MAX_CURSORS } from "./multi-cursor";
|
||||
import { findOccurrences } from "../search-and-replace/find-occurrences";
|
||||
|
||||
const extensions = [Document, Paragraph, Text, Bold, MultiCursor];
|
||||
|
||||
function makeEditor(content?: any) {
|
||||
return new Editor({
|
||||
extensions,
|
||||
content: content ?? { type: "doc", content: [{ type: "paragraph" }] },
|
||||
});
|
||||
}
|
||||
|
||||
function doc(...paragraphs: string[]) {
|
||||
return {
|
||||
type: "doc",
|
||||
content: paragraphs.map((text) => ({
|
||||
type: "paragraph",
|
||||
content: text ? [{ type: "text", text }] : [],
|
||||
})),
|
||||
};
|
||||
}
|
||||
|
||||
function paraTexts(d: PMNode): string[] {
|
||||
const out: string[] = [];
|
||||
d.forEach((node) => {
|
||||
if (node.type.name === "paragraph") out.push(node.textContent);
|
||||
});
|
||||
return out;
|
||||
}
|
||||
|
||||
function cursors(editor: Editor) {
|
||||
return multiCursorPluginKey.getState(editor.state)!.cursors;
|
||||
}
|
||||
|
||||
// Simulate typing a character through the real handleTextInput routing (the
|
||||
// browser path). someMethod-equivalent: dispatch a DOM-ish text input by calling
|
||||
// the view's input handler directly.
|
||||
function typeText(editor: Editor, text: string) {
|
||||
const { from, to } = editor.state.selection;
|
||||
// props.handleTextInput is what ProseMirror calls on beforeinput/keypress.
|
||||
const handled = editor.view.someProp(
|
||||
"handleTextInput",
|
||||
(fn) => fn(editor.view, from, to, text) || false,
|
||||
);
|
||||
if (!handled) {
|
||||
// Fall back to a normal insertion (no active multi-cursor set).
|
||||
editor.view.dispatch(editor.state.tr.insertText(text, from, to));
|
||||
}
|
||||
}
|
||||
|
||||
function pressKey(editor: Editor, key: string) {
|
||||
editor.view.someProp("handleKeyDown", (fn) =>
|
||||
fn(editor.view, new KeyboardEvent("keydown", { key })),
|
||||
);
|
||||
}
|
||||
|
||||
describe("multi-cursor: selectAllOccurrences", () => {
|
||||
it("finds EVERY occurrence of a repeated word under the cursor", () => {
|
||||
const editor = makeEditor(doc("foo bar foo baz foo"));
|
||||
// Cursor inside the first "foo".
|
||||
editor.commands.setTextSelection(2);
|
||||
expect(editor.commands.selectAllOccurrences()).toBe(true);
|
||||
|
||||
const cs = cursors(editor);
|
||||
expect(cs.length).toBe(3);
|
||||
// Every cursor spans a "foo".
|
||||
for (const c of cs) {
|
||||
expect(editor.state.doc.textBetween(c.from, c.to)).toBe("foo");
|
||||
}
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("uses the current non-empty selection as the term", () => {
|
||||
const editor = makeEditor(doc("ab abc ab abcd ab"));
|
||||
// Select the first "ab".
|
||||
editor.commands.setTextSelection({ from: 1, to: 3 });
|
||||
expect(editor.state.doc.textBetween(1, 3)).toBe("ab");
|
||||
editor.commands.selectAllOccurrences();
|
||||
// Literal substring match (selection is not whole-word), so every "ab"
|
||||
// including those inside "abc"/"abcd" is matched: 5 total.
|
||||
const cs = cursors(editor);
|
||||
expect(cs.length).toBe(5);
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("whole-word matching from a word cursor does not match substrings", () => {
|
||||
const editor = makeEditor(doc("cat category cat scatter cat"));
|
||||
editor.commands.setTextSelection(2); // inside first "cat"
|
||||
editor.commands.selectAllOccurrences();
|
||||
// Only the three standalone "cat" words, not "category"/"scatter".
|
||||
expect(cursors(editor).length).toBe(3);
|
||||
editor.destroy();
|
||||
});
|
||||
});
|
||||
|
||||
describe("multi-cursor: mass typing (single transaction)", () => {
|
||||
it("types text into N carets at once", () => {
|
||||
const editor = makeEditor(doc("foo foo foo"));
|
||||
editor.commands.setTextSelection(2);
|
||||
editor.commands.selectAllOccurrences();
|
||||
expect(cursors(editor).length).toBe(3);
|
||||
|
||||
// Typing replaces each selected "foo" with "X".
|
||||
typeText(editor, "X");
|
||||
expect(paraTexts(editor.state.doc)).toEqual(["X X X"]);
|
||||
|
||||
// The cursors are now carets right after each inserted "X".
|
||||
const cs = cursors(editor);
|
||||
expect(cs.length).toBe(3);
|
||||
for (const c of cs) expect(c.from).toBe(c.to);
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("continues typing at the resulting carets (append semantics)", () => {
|
||||
const editor = makeEditor(doc("a a a"));
|
||||
editor.commands.setTextSelection(1);
|
||||
editor.commands.selectAllOccurrences();
|
||||
typeText(editor, "b"); // each "a" -> "b"
|
||||
typeText(editor, "c"); // append at each caret -> "bc"
|
||||
expect(paraTexts(editor.state.doc)).toEqual(["bc bc bc"]);
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("applies the whole multi-edit in a SINGLE transaction (one undo step)", () => {
|
||||
// "One Cmd/Ctrl+Z undoes the whole multi-edit" holds iff the N edits land in
|
||||
// ONE transaction (history groups by transaction). @tiptap/extension-history
|
||||
// is not a dependency here, so rather than exercise undo we assert the
|
||||
// property that guarantees it: typing into N cursors is exactly ONE dispatch.
|
||||
const editor = makeEditor(doc("foo foo foo"));
|
||||
editor.commands.setTextSelection(2);
|
||||
editor.commands.selectAllOccurrences();
|
||||
expect(cursors(editor).length).toBe(3);
|
||||
|
||||
const orig = editor.view.dispatch.bind(editor.view);
|
||||
let dispatches = 0;
|
||||
editor.view.dispatch = (tr) => {
|
||||
dispatches += 1;
|
||||
return orig(tr);
|
||||
};
|
||||
typeText(editor, "Z");
|
||||
editor.view.dispatch = orig;
|
||||
|
||||
expect(dispatches).toBe(1); // all three edits share one transaction
|
||||
expect(paraTexts(editor.state.doc)).toEqual(["Z Z Z"]);
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("off-by-one guard: reverse-order iteration keeps every position valid", () => {
|
||||
// If the mass edit iterated FORWARD, inserting at an earlier cursor would
|
||||
// shift every later cursor and corrupt the result. Different-length
|
||||
// replacement makes such a bug visible.
|
||||
const editor = makeEditor(doc("x x x x"));
|
||||
editor.commands.setTextSelection(1);
|
||||
editor.commands.selectAllOccurrences();
|
||||
expect(cursors(editor).length).toBe(4);
|
||||
typeText(editor, "LONG");
|
||||
expect(paraTexts(editor.state.doc)).toEqual(["LONG LONG LONG LONG"]);
|
||||
editor.destroy();
|
||||
});
|
||||
});
|
||||
|
||||
describe("multi-cursor: mass Backspace / Delete", () => {
|
||||
it("Backspace removes one char before each caret", () => {
|
||||
const editor = makeEditor(doc("foo foo foo"));
|
||||
editor.commands.setTextSelection(2);
|
||||
editor.commands.selectAllOccurrences();
|
||||
// Collapse selections to carets at the END of each "foo" by typing then
|
||||
// removing is complex; instead type to convert ranges into carets first.
|
||||
typeText(editor, "ab"); // each "foo" -> "ab", carets after "ab"
|
||||
expect(paraTexts(editor.state.doc)).toEqual(["ab ab ab"]);
|
||||
pressKey(editor, "Backspace"); // remove the trailing "b" at each caret
|
||||
expect(paraTexts(editor.state.doc)).toEqual(["a a a"]);
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("Delete removes one char after each caret", () => {
|
||||
const editor = makeEditor(doc("fooX fooX"));
|
||||
// Literal (selection) match of "foo" -> both occurrences inside "fooX".
|
||||
editor.commands.setTextSelection({ from: 1, to: 4 }); // first "foo"
|
||||
editor.commands.selectAllOccurrences();
|
||||
expect(cursors(editor).length).toBe(2);
|
||||
typeText(editor, "foo"); // rewrite "foo", carets now sit before each "X"
|
||||
expect(paraTexts(editor.state.doc)).toEqual(["fooX fooX"]);
|
||||
pressKey(editor, "Delete"); // remove the "X" after each caret
|
||||
expect(paraTexts(editor.state.doc)).toEqual(["foo foo"]);
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("Backspace at a block-start caret is a no-op for that cursor", () => {
|
||||
const editor = makeEditor(doc("ab", "ab"));
|
||||
// Select both "ab" then convert to carets at start by replacing with "".
|
||||
editor.commands.setTextSelection({ from: 1, to: 3 }); // first "ab"
|
||||
editor.commands.selectAllOccurrences();
|
||||
// Move carets to block start: type "" is not possible; instead delete range.
|
||||
pressKey(editor, "Backspace"); // deletes each selected "ab"
|
||||
expect(paraTexts(editor.state.doc)).toEqual(["", ""]);
|
||||
// Carets are now at each block start; another Backspace must not throw and
|
||||
// must not merge blocks (still two empty paragraphs).
|
||||
pressKey(editor, "Backspace");
|
||||
expect(paraTexts(editor.state.doc)).toEqual(["", ""]);
|
||||
editor.destroy();
|
||||
});
|
||||
});
|
||||
|
||||
describe("multi-cursor: addNextOccurrence (Cmd/Ctrl+D)", () => {
|
||||
it("first press selects the current word, next press adds the next", () => {
|
||||
const editor = makeEditor(doc("go go go"));
|
||||
editor.commands.setTextSelection(2); // inside first "go"
|
||||
editor.commands.addNextOccurrence();
|
||||
expect(cursors(editor).length).toBe(1);
|
||||
editor.commands.addNextOccurrence();
|
||||
expect(cursors(editor).length).toBe(2);
|
||||
editor.commands.addNextOccurrence();
|
||||
expect(cursors(editor).length).toBe(3);
|
||||
// Nothing left to add — stays at 3.
|
||||
editor.commands.addNextOccurrence();
|
||||
expect(cursors(editor).length).toBe(3);
|
||||
for (const c of cursors(editor)) {
|
||||
expect(editor.state.doc.textBetween(c.from, c.to)).toBe("go");
|
||||
}
|
||||
editor.destroy();
|
||||
});
|
||||
});
|
||||
|
||||
describe("multi-cursor: position remapping", () => {
|
||||
it("remaps cursors after a LOCAL edit before them", () => {
|
||||
const editor = makeEditor(doc("foo foo"));
|
||||
editor.commands.setTextSelection(2);
|
||||
editor.commands.selectAllOccurrences();
|
||||
const before = cursors(editor).map((c) => ({ ...c }));
|
||||
|
||||
// Insert unrelated text at the very start (pos 1), shifting everything +5.
|
||||
editor.view.dispatch(editor.state.tr.insertText("HELLO", 1));
|
||||
|
||||
const after = cursors(editor);
|
||||
expect(after.length).toBe(before.length);
|
||||
for (let i = 0; i < after.length; i += 1) {
|
||||
expect(after[i].from).toBe(before[i].from + 5);
|
||||
expect(after[i].to).toBe(before[i].to + 5);
|
||||
// And they still point at "foo".
|
||||
expect(editor.state.doc.textBetween(after[i].from, after[i].to)).toBe(
|
||||
"foo",
|
||||
);
|
||||
}
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("remaps cursors after a simulated REMOTE edit (ordinary transaction)", () => {
|
||||
const editor = makeEditor(doc("foo bar foo"));
|
||||
editor.commands.setTextSelection(2);
|
||||
editor.commands.selectAllOccurrences();
|
||||
const before = cursors(editor).map((c) => ({ ...c }));
|
||||
expect(before.length).toBe(2);
|
||||
|
||||
// y-prosemirror applies remote changes as ordinary transactions. Emulate a
|
||||
// remote insertion between the two "foo"s (inside "bar", pos 6) with a tr
|
||||
// that carries NO multi-cursor meta — exactly like a collaborator's edit.
|
||||
const tr = editor.state.tr.insertText("ZZ", 6);
|
||||
editor.view.dispatch(tr);
|
||||
|
||||
const after = cursors(editor);
|
||||
// The first "foo" (before the insertion) is unchanged; the second shifts +2.
|
||||
expect(after[0].from).toBe(before[0].from);
|
||||
expect(after[1].from).toBe(before[1].from + 2);
|
||||
for (const c of after) {
|
||||
expect(editor.state.doc.textBetween(c.from, c.to)).toBe("foo");
|
||||
}
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("a REMOTE delete UNDER a cursor collapses it to a caret (not drop), leaving others intact", () => {
|
||||
// The riskiest remap path: a collaborator deletes the very text one cursor
|
||||
// spans. Both edges map with assoc +1 and there is no drop logic, so the
|
||||
// deleted-over cursor CONTRACT is: it collapses to a zero-width caret at the
|
||||
// deletion point (from === to) and STAYS in the set — it is not removed.
|
||||
// Untouched cursors keep spanning their occurrence. Pinning this makes the
|
||||
// collapse-not-drop choice explicit (review #372 F2).
|
||||
const editor = makeEditor(doc("foo bar foo"));
|
||||
editor.commands.setTextSelection(2);
|
||||
editor.commands.selectAllOccurrences();
|
||||
const before = cursors(editor).map((c) => ({ ...c }));
|
||||
expect(before.length).toBe(2);
|
||||
|
||||
// Remote (no multi-cursor meta) delete of the FIRST "foo" range.
|
||||
const tr = editor.state.tr.delete(before[0].from, before[0].to);
|
||||
editor.view.dispatch(tr);
|
||||
|
||||
const after = cursors(editor);
|
||||
// Still two cursors — the deleted-over one is NOT dropped.
|
||||
expect(after.length).toBe(2);
|
||||
// The first collapsed to a caret at the deletion point.
|
||||
expect(after[0].from).toBe(after[0].to);
|
||||
expect(after[0].from).toBe(before[0].from);
|
||||
// The second still spans "foo" (shifted left by the 3 removed chars).
|
||||
expect(after[1].from).toBe(before[1].from - 3);
|
||||
expect(editor.state.doc.textBetween(after[1].from, after[1].to)).toBe("foo");
|
||||
// Sanity: the document now reads " bar foo".
|
||||
expect(paraTexts(editor.state.doc)).toEqual([" bar foo"]);
|
||||
editor.destroy();
|
||||
});
|
||||
});
|
||||
|
||||
describe("multi-cursor: collapse / exit", () => {
|
||||
it("exitMultiCursor clears the set", () => {
|
||||
const editor = makeEditor(doc("foo foo"));
|
||||
editor.commands.setTextSelection(2);
|
||||
editor.commands.selectAllOccurrences();
|
||||
expect(cursors(editor).length).toBe(2);
|
||||
editor.commands.exitMultiCursor();
|
||||
expect(cursors(editor).length).toBe(0);
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("an arrow key collapses the set", () => {
|
||||
const editor = makeEditor(doc("foo foo"));
|
||||
editor.commands.setTextSelection(2);
|
||||
editor.commands.selectAllOccurrences();
|
||||
expect(cursors(editor).length).toBe(2);
|
||||
pressKey(editor, "ArrowRight");
|
||||
expect(cursors(editor).length).toBe(0);
|
||||
editor.destroy();
|
||||
});
|
||||
});
|
||||
|
||||
describe("multi-cursor: collapse on composition / mousedown", () => {
|
||||
// Invoke a plugin handleDOMEvents handler through the real prop plumbing.
|
||||
function fireDOM(editor: Editor, name: string): void {
|
||||
editor.view.someProp("handleDOMEvents", (handlers: any) => {
|
||||
const h = handlers && handlers[name];
|
||||
if (h) h(editor.view, new Event(name));
|
||||
return false;
|
||||
});
|
||||
}
|
||||
|
||||
it("collapses the set on compositionstart (IME) — MVP does not multi-IME", () => {
|
||||
const editor = makeEditor(doc("foo foo"));
|
||||
editor.commands.setTextSelection(2);
|
||||
editor.commands.selectAllOccurrences();
|
||||
expect(cursors(editor).length).toBe(2);
|
||||
fireDOM(editor, "compositionstart");
|
||||
expect(cursors(editor).length).toBe(0);
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("collapses the set on a plain mousedown (VS Code behaviour)", () => {
|
||||
const editor = makeEditor(doc("foo foo"));
|
||||
editor.commands.setTextSelection(2);
|
||||
editor.commands.selectAllOccurrences();
|
||||
expect(cursors(editor).length).toBe(2);
|
||||
fireDOM(editor, "mousedown");
|
||||
expect(cursors(editor).length).toBe(0);
|
||||
editor.destroy();
|
||||
});
|
||||
});
|
||||
|
||||
describe("multi-cursor: hard cap", () => {
|
||||
it("never activates more than MAX_CURSORS cursors", () => {
|
||||
const many = new Array(MAX_CURSORS + 20).fill("w").join(" ");
|
||||
const editor = makeEditor(doc(many));
|
||||
editor.commands.setTextSelection(2);
|
||||
editor.commands.selectAllOccurrences();
|
||||
expect(cursors(editor).length).toBe(MAX_CURSORS);
|
||||
editor.destroy();
|
||||
});
|
||||
});
|
||||
|
||||
describe("multi-cursor: marks are carried across a mass edit", () => {
|
||||
it("preserves marks spanning each replaced range", () => {
|
||||
const editor = makeEditor({
|
||||
type: "doc",
|
||||
content: [
|
||||
{
|
||||
type: "paragraph",
|
||||
content: [
|
||||
{ type: "text", text: "a " },
|
||||
{ type: "text", marks: [{ type: "bold" }], text: "key" },
|
||||
{ type: "text", text: " b " },
|
||||
{ type: "text", marks: [{ type: "bold" }], text: "key" },
|
||||
],
|
||||
},
|
||||
],
|
||||
});
|
||||
editor.commands.setTextSelection(3); // inside first bold "key"
|
||||
editor.commands.selectAllOccurrences();
|
||||
expect(cursors(editor).length).toBe(2);
|
||||
typeText(editor, "NEW");
|
||||
|
||||
// Both replacements keep the bold mark.
|
||||
let boldRuns = 0;
|
||||
editor.state.doc.descendants((node) => {
|
||||
if (
|
||||
node.isText &&
|
||||
node.text === "NEW" &&
|
||||
node.marks.some((m) => m.type.name === "bold")
|
||||
) {
|
||||
boldRuns += 1;
|
||||
}
|
||||
});
|
||||
expect(boldRuns).toBe(2);
|
||||
editor.destroy();
|
||||
});
|
||||
});
|
||||
|
||||
// The extracted find-occurrences util must return the SAME occurrences that the
|
||||
// old inline walk produced (and that search-and-replace still relies on).
|
||||
describe("find-occurrences util", () => {
|
||||
it("finds all matches of a literal regex across text nodes", () => {
|
||||
const editor = makeEditor(doc("foo foofoo foo"));
|
||||
const results = findOccurrences(editor.state.doc, /foo/gu);
|
||||
// 4 occurrences: two standalone + two inside "foofoo".
|
||||
expect(results.length).toBe(4);
|
||||
for (const r of results) {
|
||||
expect(editor.state.doc.textBetween(r.from, r.to)).toBe("foo");
|
||||
}
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("ignores whitespace-only matches and empty regex", () => {
|
||||
const editor = makeEditor(doc("a b c"));
|
||||
expect(findOccurrences(editor.state.doc, null as any).length).toBe(0);
|
||||
// A whitespace regex yields no results (matches are trimmed away).
|
||||
expect(findOccurrences(editor.state.doc, /\s/gu).length).toBe(0);
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("finds a match spanning two differently-marked contiguous text nodes", () => {
|
||||
const editor = makeEditor({
|
||||
type: "doc",
|
||||
content: [
|
||||
{
|
||||
type: "paragraph",
|
||||
content: [
|
||||
{ type: "text", text: "wo" },
|
||||
{ type: "text", marks: [{ type: "bold" }], text: "rd" },
|
||||
],
|
||||
},
|
||||
],
|
||||
});
|
||||
const results = findOccurrences(editor.state.doc, /word/gu);
|
||||
expect(results.length).toBe(1);
|
||||
expect(editor.state.doc.textBetween(results[0].from, results[0].to)).toBe(
|
||||
"word",
|
||||
);
|
||||
editor.destroy();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,545 @@
|
||||
import { Extension, Range } from "@tiptap/core";
|
||||
import { Decoration, DecorationSet, EditorView } from "@tiptap/pm/view";
|
||||
import {
|
||||
Plugin,
|
||||
PluginKey,
|
||||
TextSelection,
|
||||
type EditorState,
|
||||
} from "@tiptap/pm/state";
|
||||
import { Mark } from "@tiptap/pm/model";
|
||||
import { findOccurrences } from "../search-and-replace/find-occurrences";
|
||||
|
||||
/**
|
||||
* Multi-cursor editing — MVP (issue #196, "Variant A").
|
||||
*
|
||||
* VS Code-style multi-cursor limited to "select all occurrences of a word (or
|
||||
* the current selection) and type into all of them at once", built ON TOP OF
|
||||
* the search-and-replace mass-transaction machinery:
|
||||
*
|
||||
* - Cmd/Ctrl+Shift+L (selectAllOccurrences): the word under the cursor (or the
|
||||
* current non-empty selection) -> ALL its occurrences become active cursors.
|
||||
* - Cmd/Ctrl+D (addNextOccurrence): add the NEXT occurrence of the term.
|
||||
* - Typing / Backspace / Delete apply to EVERY active cursor in ONE
|
||||
* transaction (so a single Cmd/Ctrl+Z undoes the whole multi-edit).
|
||||
* - Esc (exitMultiCursor): collapse back to a single cursor.
|
||||
*
|
||||
* The single-transaction, reverse-order edit mechanic mirrors `replaceAll` in
|
||||
* search-and-replace.ts: we iterate cursors from the END of the document to the
|
||||
* START so an earlier edit never invalidates a later position, carrying the
|
||||
* marks that span each range.
|
||||
*
|
||||
* CONSCIOUS v1 OUT-OF-SCOPE BOUNDARIES (these are "Variant B", deliberately NOT
|
||||
* built here):
|
||||
* - Alt+Click arbitrary carets and Alt+drag column selection.
|
||||
* - Cmd/Ctrl+Alt+Up/Down "add cursor on the adjacent line".
|
||||
* - Simultaneous IME / composition input into multiple positions — on
|
||||
* `compositionstart` we collapse back to a single cursor.
|
||||
* - Cursors spanning different schema nodes in one edit.
|
||||
*
|
||||
* NOT out of scope, but worth stating precisely: there is NO schema-aware or
|
||||
* structural cursor. Occurrences are found by a plain text-node walk
|
||||
* (`findOccurrences`), so a term that appears inside a table cell, code block or
|
||||
* callout DOES get a cursor there and IS edited — as plain text, exactly like
|
||||
* `replaceAll`. There is no special table/code handling; the per-cursor try/catch
|
||||
* only SKIPS a cursor whose edit would violate the schema (never applied
|
||||
* half-way), it does not exclude those node types from matching.
|
||||
*/
|
||||
|
||||
interface MultiCursorState {
|
||||
// Each active cursor: a caret when from === to, a range when from < to.
|
||||
cursors: Range[];
|
||||
}
|
||||
|
||||
export const multiCursorPluginKey = new PluginKey<MultiCursorState>(
|
||||
"multiCursor",
|
||||
);
|
||||
|
||||
// Hard safety cap on simultaneously-active cursors — stop adding past it.
|
||||
export const MAX_CURSORS = 100;
|
||||
|
||||
export interface MultiCursorStorage {
|
||||
// Whether the active term matches whole words only. Set to true when the set
|
||||
// was seeded from a bare cursor (word under caret), false when seeded from an
|
||||
// explicit selection (literal substring match, like VS Code). Remembered so
|
||||
// addNextOccurrence keeps matching the same way as selectAllOccurrences.
|
||||
wholeWord: boolean;
|
||||
}
|
||||
|
||||
declare module "@tiptap/core" {
|
||||
interface Storage {
|
||||
multiCursor: MultiCursorStorage;
|
||||
}
|
||||
interface Commands<ReturnType> {
|
||||
multiCursor: {
|
||||
/** Select all occurrences of the word/selection as active cursors. */
|
||||
selectAllOccurrences: () => ReturnType;
|
||||
/** Add the next occurrence of the current term to the cursor set. */
|
||||
addNextOccurrence: () => ReturnType;
|
||||
/** Collapse the multi-cursor set back to a single cursor. */
|
||||
exitMultiCursor: () => ReturnType;
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Term helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function escapeRegExp(s: string): string {
|
||||
return s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
|
||||
}
|
||||
|
||||
// A "word" is a run of letters/numbers/underscore; those get whole-word
|
||||
// matching (\b…\b) so a term never matches inside a larger word. Anything else
|
||||
// (punctuation, phrases) is matched literally. Case-sensitive, like VS Code.
|
||||
function isWordTerm(s: string): boolean {
|
||||
return /^[\p{L}\p{N}_]+$/u.test(s);
|
||||
}
|
||||
|
||||
// wholeWord uses \b…\b so the term never matches inside a larger word; it only
|
||||
// applies to word-like terms (a term containing punctuation cannot be
|
||||
// whole-word-bounded meaningfully). Otherwise the term is matched literally.
|
||||
function buildTermRegex(term: string, wholeWord: boolean): RegExp {
|
||||
const esc = escapeRegExp(term);
|
||||
return wholeWord && isWordTerm(term)
|
||||
? new RegExp(`\\b${esc}\\b`, "gu")
|
||||
: new RegExp(esc, "gu");
|
||||
}
|
||||
|
||||
// Word under a position: returns the exact { from, to } range and its text, or
|
||||
// null if the position is not inside a word in a textblock.
|
||||
function getWordAt(
|
||||
state: EditorState,
|
||||
pos: number,
|
||||
): { from: number; to: number; text: string } | null {
|
||||
const $pos = state.doc.resolve(pos);
|
||||
const parent = $pos.parent;
|
||||
if (!parent.isTextblock) return null;
|
||||
|
||||
const text = parent.textContent;
|
||||
const offset = $pos.parentOffset;
|
||||
const start = $pos.start();
|
||||
const wordRe = /[\p{L}\p{N}_]+/gu;
|
||||
|
||||
let m: RegExpExecArray | null;
|
||||
while ((m = wordRe.exec(text)) !== null) {
|
||||
const s = m.index;
|
||||
const e = m.index + m[0].length;
|
||||
if (offset >= s && offset <= e) {
|
||||
return { from: start + s, to: start + e, text: m[0] };
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Plugin-state access
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function getCursors(state: EditorState): Range[] {
|
||||
const st = multiCursorPluginKey.getState(state);
|
||||
return st ? st.cursors : [];
|
||||
}
|
||||
|
||||
function setCursors(view: EditorView, cursors: Range[]): void {
|
||||
view.dispatch(view.state.tr.setMeta(multiCursorPluginKey, cursors));
|
||||
}
|
||||
|
||||
function collapse(view: EditorView): void {
|
||||
setCursors(view, []);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// The single-transaction, reverse-order mass edit (mirrors replaceAll)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
interface EditOp {
|
||||
from: number;
|
||||
to: number;
|
||||
// Text to insert at `from` after deleting [from, to); "" for a pure delete.
|
||||
text: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Apply one edit per cursor in ONE transaction. Ops are processed from the END
|
||||
* of the document to the START so an earlier edit never shifts a later position
|
||||
* (mirrors `replaceAll`). Each cursor is wrapped independently: a schema
|
||||
* violation SKIPS that one cursor instead of throwing away the whole
|
||||
* transaction, so the document is never left half-applied.
|
||||
*
|
||||
* After building the transaction the new cursor positions are recomputed by
|
||||
* mapping each op's original anchor through `tr.mapping` (which also remaps any
|
||||
* concurrent changes), so carets land right after their inserted text.
|
||||
*/
|
||||
function dispatchMassEdit(view: EditorView, ops: EditOp[]): boolean {
|
||||
if (!ops.length) return false;
|
||||
|
||||
const { state } = view;
|
||||
const tr = state.tr;
|
||||
const schema = state.schema;
|
||||
|
||||
// Ascending by `from`; iterate reverse so earlier positions stay valid.
|
||||
const sorted = [...ops].sort((a, b) => a.from - b.from);
|
||||
const appliedLen: number[] = new Array(sorted.length).fill(0);
|
||||
|
||||
for (let i = sorted.length - 1; i >= 0; i -= 1) {
|
||||
const { from, to, text } = sorted[i];
|
||||
try {
|
||||
let marks: readonly Mark[] = [];
|
||||
if (text) {
|
||||
if (to > from) {
|
||||
// Carry all marks spanning the replaced range.
|
||||
const set = new Set<Mark>();
|
||||
tr.doc.nodesBetween(from, to, (node) => {
|
||||
if (node.isText && node.marks) {
|
||||
node.marks.forEach((mk) => set.add(mk));
|
||||
}
|
||||
});
|
||||
marks = Array.from(set);
|
||||
} else {
|
||||
// Caret: continue the marks active at the insertion point.
|
||||
marks = state.storedMarks || state.doc.resolve(from).marks();
|
||||
}
|
||||
}
|
||||
|
||||
// ONE atomic step per cursor: replaceWith covers both insert (from === to)
|
||||
// and replace (to > from); a pure delete (empty text) uses delete. This
|
||||
// can never leave a cursor half-applied (deleted but not re-inserted) the
|
||||
// way a separate delete-then-insert pair could if the insert step threw.
|
||||
if (text) {
|
||||
tr.replaceWith(from, to, schema.text(text, marks as Mark[]));
|
||||
} else if (to > from) {
|
||||
tr.delete(from, to);
|
||||
}
|
||||
|
||||
appliedLen[i] = text.length;
|
||||
} catch {
|
||||
// Per-cursor backstop (text-only MVP): drop this cursor's edit, keep the
|
||||
// rest of the transaction intact.
|
||||
appliedLen[i] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
if (!tr.docChanged) return false;
|
||||
|
||||
// Recompute cursor carets from the ORIGINAL op anchors through the full map.
|
||||
const newCursors: Range[] = sorted.map((op, i) => {
|
||||
const start = tr.mapping.map(op.from, -1);
|
||||
const caret = start + appliedLen[i];
|
||||
return { from: caret, to: caret };
|
||||
});
|
||||
|
||||
tr.setMeta(multiCursorPluginKey, newCursors);
|
||||
|
||||
// Park the native selection on the last caret so the browser draws exactly
|
||||
// one real caret; the rest are our decoration widgets.
|
||||
const last = newCursors[newCursors.length - 1];
|
||||
tr.setSelection(TextSelection.create(tr.doc, last.from));
|
||||
|
||||
view.dispatch(tr);
|
||||
return true;
|
||||
}
|
||||
|
||||
function buildDeleteOps(
|
||||
state: EditorState,
|
||||
cursors: Range[],
|
||||
forward: boolean,
|
||||
): EditOp[] {
|
||||
return cursors.map((c) => {
|
||||
// A selected range: Backspace/Delete removes the whole range.
|
||||
if (c.to > c.from) return { from: c.from, to: c.to, text: "" };
|
||||
|
||||
const $pos = state.doc.resolve(c.from);
|
||||
if (forward) {
|
||||
// Delete: at the end of a textblock there is nothing to remove (a no-op;
|
||||
// MVP does not merge blocks across a multi-cursor set).
|
||||
if ($pos.parentOffset >= $pos.parent.content.size) {
|
||||
return { from: c.from, to: c.from, text: "" };
|
||||
}
|
||||
return { from: c.from, to: c.from + 1, text: "" };
|
||||
}
|
||||
// Backspace: at the start of a textblock there is nothing to remove.
|
||||
if ($pos.parentOffset <= 0) {
|
||||
return { from: c.from, to: c.from, text: "" };
|
||||
}
|
||||
return { from: c.from - 1, to: c.from, text: "" };
|
||||
});
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Extension
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export const MultiCursor = Extension.create<unknown, MultiCursorStorage>({
|
||||
name: "multiCursor",
|
||||
|
||||
addStorage() {
|
||||
return { wholeWord: true };
|
||||
},
|
||||
|
||||
addCommands() {
|
||||
return {
|
||||
selectAllOccurrences:
|
||||
() =>
|
||||
({ editor, state, tr, dispatch }) => {
|
||||
let term: string;
|
||||
// A bare cursor expands to the whole word; an explicit selection is
|
||||
// matched literally (VS Code semantics).
|
||||
const wholeWord = state.selection.empty;
|
||||
if (wholeWord) {
|
||||
const word = getWordAt(state, state.selection.from);
|
||||
if (!word) return false;
|
||||
term = word.text;
|
||||
} else {
|
||||
term = state.doc.textBetween(
|
||||
state.selection.from,
|
||||
state.selection.to,
|
||||
);
|
||||
}
|
||||
if (!term.trim()) return false;
|
||||
editor.storage.multiCursor.wholeWord = wholeWord;
|
||||
|
||||
const results = findOccurrences(
|
||||
state.doc,
|
||||
buildTermRegex(term, wholeWord),
|
||||
).slice(0, MAX_CURSORS);
|
||||
if (!results.length) return false;
|
||||
|
||||
if (dispatch) {
|
||||
tr.setMeta(multiCursorPluginKey, results);
|
||||
const last = results[results.length - 1];
|
||||
tr.setSelection(TextSelection.create(tr.doc, last.from, last.to));
|
||||
dispatch(tr);
|
||||
}
|
||||
return true;
|
||||
},
|
||||
|
||||
addNextOccurrence:
|
||||
() =>
|
||||
({ editor, state, tr, dispatch }) => {
|
||||
const existing = getCursors(state);
|
||||
let cursors: Range[];
|
||||
|
||||
if (!existing.length) {
|
||||
// First press: turn the current word/selection into the one cursor.
|
||||
let range: Range;
|
||||
const wholeWord = state.selection.empty;
|
||||
if (wholeWord) {
|
||||
const word = getWordAt(state, state.selection.from);
|
||||
if (!word) return false;
|
||||
range = { from: word.from, to: word.to };
|
||||
} else {
|
||||
range = { from: state.selection.from, to: state.selection.to };
|
||||
}
|
||||
editor.storage.multiCursor.wholeWord = wholeWord;
|
||||
cursors = [range];
|
||||
} else {
|
||||
// Subsequent press: add the next unselected occurrence of the term,
|
||||
// matched the SAME way (whole-word vs literal) the set was seeded.
|
||||
if (existing.length >= MAX_CURSORS) return true;
|
||||
|
||||
const first = existing[0];
|
||||
const term = state.doc.textBetween(first.from, first.to);
|
||||
if (!term.trim()) return false;
|
||||
|
||||
const results = findOccurrences(
|
||||
state.doc,
|
||||
buildTermRegex(term, editor.storage.multiCursor.wholeWord),
|
||||
);
|
||||
const keys = new Set(existing.map((c) => `${c.from}:${c.to}`));
|
||||
const notSelected = results.filter(
|
||||
(r) => !keys.has(`${r.from}:${r.to}`),
|
||||
);
|
||||
if (!notSelected.length) return true; // all occurrences selected
|
||||
|
||||
const maxTo = Math.max(...existing.map((c) => c.to));
|
||||
const next =
|
||||
notSelected.find((r) => r.from >= maxTo) || notSelected[0];
|
||||
cursors = [...existing, next];
|
||||
}
|
||||
|
||||
if (dispatch) {
|
||||
tr.setMeta(multiCursorPluginKey, cursors);
|
||||
const last = cursors[cursors.length - 1];
|
||||
tr.setSelection(TextSelection.create(tr.doc, last.from, last.to));
|
||||
dispatch(tr);
|
||||
}
|
||||
return true;
|
||||
},
|
||||
|
||||
exitMultiCursor:
|
||||
() =>
|
||||
({ tr, dispatch }) => {
|
||||
if (dispatch) {
|
||||
tr.setMeta(multiCursorPluginKey, []);
|
||||
dispatch(tr);
|
||||
}
|
||||
return true;
|
||||
},
|
||||
};
|
||||
},
|
||||
|
||||
addKeyboardShortcuts() {
|
||||
return {
|
||||
"Mod-Shift-l": () => {
|
||||
this.editor.commands.selectAllOccurrences();
|
||||
// Always consume so the browser's default is prevented.
|
||||
return true;
|
||||
},
|
||||
"Mod-d": () => {
|
||||
this.editor.commands.addNextOccurrence();
|
||||
// Consume unconditionally to prevent the browser's Cmd/Ctrl+D bookmark.
|
||||
return true;
|
||||
},
|
||||
Escape: () => {
|
||||
// Only swallow Escape while a multi-cursor set is active; otherwise let
|
||||
// Escape keep its other behaviours (e.g. closing dialogs).
|
||||
if (!getCursors(this.editor.state).length) return false;
|
||||
return this.editor.commands.exitMultiCursor();
|
||||
},
|
||||
};
|
||||
},
|
||||
|
||||
addProseMirrorPlugins() {
|
||||
return [
|
||||
new Plugin<MultiCursorState>({
|
||||
key: multiCursorPluginKey,
|
||||
|
||||
state: {
|
||||
init: () => ({ cursors: [] }),
|
||||
apply(tr, value): MultiCursorState {
|
||||
// A command (or a mass edit) can set/clear the cursor set directly.
|
||||
// Its cursors are already in the post-transaction coordinate space,
|
||||
// so they take priority over remapping.
|
||||
const meta = tr.getMeta(multiCursorPluginKey) as
|
||||
| Range[]
|
||||
| undefined;
|
||||
if (meta !== undefined) {
|
||||
return { cursors: meta.slice(0, MAX_CURSORS) };
|
||||
}
|
||||
|
||||
if (!value.cursors.length) return value;
|
||||
|
||||
// Remap surviving cursors across ANY doc change — this covers both
|
||||
// local edits and REMOTE Yjs edits (y-prosemirror applies remote
|
||||
// changes as ordinary transactions, so mapping them here keeps every
|
||||
// multi-cursor correctly positioned without special-casing collab).
|
||||
if (tr.docChanged) {
|
||||
// Map both edges with the SAME association (+1) so content
|
||||
// inserted at a boundary shifts the whole cursor right and a caret
|
||||
// (from === to) can never invert into a range.
|
||||
const cursors = value.cursors.map((c) => ({
|
||||
from: tr.mapping.map(c.from, 1),
|
||||
to: tr.mapping.map(c.to, 1),
|
||||
}));
|
||||
return { cursors };
|
||||
}
|
||||
|
||||
return value;
|
||||
},
|
||||
},
|
||||
|
||||
props: {
|
||||
decorations(state) {
|
||||
const st = multiCursorPluginKey.getState(state);
|
||||
if (!st || !st.cursors.length) return DecorationSet.empty;
|
||||
|
||||
const decorations: Decoration[] = [];
|
||||
st.cursors.forEach((c, i) => {
|
||||
if (c.from === c.to) {
|
||||
decorations.push(
|
||||
Decoration.widget(
|
||||
c.from,
|
||||
() => {
|
||||
const el = document.createElement("span");
|
||||
el.className = "multi-cursor__caret";
|
||||
return el;
|
||||
},
|
||||
{ side: 0, key: `mc-caret-${i}` },
|
||||
),
|
||||
);
|
||||
} else {
|
||||
decorations.push(
|
||||
Decoration.inline(c.from, c.to, {
|
||||
class: "multi-cursor__selection",
|
||||
}),
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
return DecorationSet.create(state.doc, decorations);
|
||||
},
|
||||
|
||||
handleTextInput(view, _from, _to, text) {
|
||||
const cursors = getCursors(view.state);
|
||||
if (!cursors.length) return false;
|
||||
|
||||
// Insert `text` at EVERY cursor in one transaction. Returning true
|
||||
// prevents ProseMirror's own single-position insert at the native
|
||||
// selection, so there is no double-insert there.
|
||||
const ops = cursors.map((c) => ({
|
||||
from: c.from,
|
||||
to: c.to,
|
||||
text,
|
||||
}));
|
||||
return dispatchMassEdit(view, ops);
|
||||
},
|
||||
|
||||
handleKeyDown(view, event) {
|
||||
const cursors = getCursors(view.state);
|
||||
if (!cursors.length) return false;
|
||||
|
||||
if (event.key === "Backspace") {
|
||||
dispatchMassEdit(view, buildDeleteOps(view.state, cursors, false));
|
||||
return true;
|
||||
}
|
||||
if (event.key === "Delete") {
|
||||
dispatchMassEdit(view, buildDeleteOps(view.state, cursors, true));
|
||||
return true;
|
||||
}
|
||||
|
||||
// Let modifier combinations (our own shortcuts, copy, etc.) through
|
||||
// WITHOUT collapsing the set.
|
||||
if (event.metaKey || event.ctrlKey || event.altKey) return false;
|
||||
|
||||
// Navigation / block keys collapse back to a single cursor, then let
|
||||
// ProseMirror handle the movement on the native selection.
|
||||
const COLLAPSE_KEYS = [
|
||||
"ArrowLeft",
|
||||
"ArrowRight",
|
||||
"ArrowUp",
|
||||
"ArrowDown",
|
||||
"Home",
|
||||
"End",
|
||||
"PageUp",
|
||||
"PageDown",
|
||||
"Enter",
|
||||
"Tab",
|
||||
];
|
||||
if (COLLAPSE_KEYS.includes(event.key)) {
|
||||
collapse(view);
|
||||
return false;
|
||||
}
|
||||
|
||||
return false;
|
||||
},
|
||||
|
||||
handleDOMEvents: {
|
||||
// A plain click exits multi-cursor (VS Code behaviour).
|
||||
mousedown: (view) => {
|
||||
if (getCursors(view.state).length) collapse(view);
|
||||
return false;
|
||||
},
|
||||
// MVP does not drive multi-position IME — collapse on composition.
|
||||
compositionstart: (view) => {
|
||||
if (getCursors(view.state).length) collapse(view);
|
||||
return false;
|
||||
},
|
||||
},
|
||||
},
|
||||
}),
|
||||
];
|
||||
},
|
||||
});
|
||||
|
||||
export default MultiCursor;
|
||||
@@ -0,0 +1,69 @@
|
||||
import { Range } from "@tiptap/core";
|
||||
import { Node as PMNode } from "@tiptap/pm/model";
|
||||
|
||||
interface TextNodesWithPosition {
|
||||
text: string;
|
||||
pos: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Shared "find all occurrences of a term in the doc" primitive.
|
||||
*
|
||||
* Walks every text node of the document and returns each regex match as a
|
||||
* `{ from, to }` range. Contiguous text nodes (which may differ only by marks)
|
||||
* are concatenated into a single run, so a match that spans e.g. "wo" + bold
|
||||
* "rd" is still found; runs are split by any non-text node, so a match never
|
||||
* crosses a node boundary. Whitespace-only matches are ignored.
|
||||
*
|
||||
* This is used by BOTH search-and-replace (highlight/replace) and multi-cursor
|
||||
* (turn occurrences into active cursors) so the two stay behaviourally in sync.
|
||||
* Extracted verbatim from the original `processSearches` walk.
|
||||
*/
|
||||
export function findOccurrences(doc: PMNode, searchTerm: RegExp): Range[] {
|
||||
const results: Range[] = [];
|
||||
|
||||
if (!searchTerm) return results;
|
||||
|
||||
let textNodesWithPosition: TextNodesWithPosition[] = [];
|
||||
let index = 0;
|
||||
|
||||
doc?.descendants((node, pos) => {
|
||||
if (node.isText) {
|
||||
if (textNodesWithPosition[index]) {
|
||||
textNodesWithPosition[index] = {
|
||||
text: textNodesWithPosition[index].text + node.text,
|
||||
pos: textNodesWithPosition[index].pos,
|
||||
};
|
||||
} else {
|
||||
textNodesWithPosition[index] = {
|
||||
text: `${node.text}`,
|
||||
pos,
|
||||
};
|
||||
}
|
||||
} else {
|
||||
index += 1;
|
||||
}
|
||||
});
|
||||
|
||||
textNodesWithPosition = textNodesWithPosition.filter(Boolean);
|
||||
|
||||
for (const element of textNodesWithPosition) {
|
||||
const { text, pos } = element;
|
||||
const matches = Array.from(text.matchAll(searchTerm)).filter(
|
||||
([matchText]) => matchText.trim(),
|
||||
);
|
||||
|
||||
for (const m of matches) {
|
||||
if (m[0] === "") break;
|
||||
|
||||
if (m.index !== undefined) {
|
||||
results.push({
|
||||
from: pos + m.index,
|
||||
to: pos + m.index + m[0].length,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return results;
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
import { SearchAndReplace } from './search-and-replace'
|
||||
export * from './search-and-replace'
|
||||
export * from './find-occurrences'
|
||||
export default SearchAndReplace
|
||||
@@ -29,6 +29,7 @@ import {
|
||||
type Transaction,
|
||||
} from "@tiptap/pm/state";
|
||||
import { Node as PMNode, Mark } from "@tiptap/pm/model";
|
||||
import { findOccurrences } from "./find-occurrences";
|
||||
|
||||
declare module "@tiptap/core" {
|
||||
interface Storage {
|
||||
@@ -76,11 +77,6 @@ declare module "@tiptap/core" {
|
||||
}
|
||||
}
|
||||
|
||||
interface TextNodesWithPosition {
|
||||
text: string;
|
||||
pos: number;
|
||||
}
|
||||
|
||||
const getRegex = (
|
||||
s: string,
|
||||
disableRegex: boolean,
|
||||
@@ -104,10 +100,6 @@ function processSearches(
|
||||
resultIndex: number,
|
||||
): ProcessedSearches {
|
||||
const decorations: Decoration[] = [];
|
||||
const results: Range[] = [];
|
||||
|
||||
let textNodesWithPosition: TextNodesWithPosition[] = [];
|
||||
let index = 0;
|
||||
|
||||
if (!searchTerm) {
|
||||
return {
|
||||
@@ -116,43 +108,8 @@ function processSearches(
|
||||
};
|
||||
}
|
||||
|
||||
doc?.descendants((node, pos) => {
|
||||
if (node.isText) {
|
||||
if (textNodesWithPosition[index]) {
|
||||
textNodesWithPosition[index] = {
|
||||
text: textNodesWithPosition[index].text + node.text,
|
||||
pos: textNodesWithPosition[index].pos,
|
||||
};
|
||||
} else {
|
||||
textNodesWithPosition[index] = {
|
||||
text: `${node.text}`,
|
||||
pos,
|
||||
};
|
||||
}
|
||||
} else {
|
||||
index += 1;
|
||||
}
|
||||
});
|
||||
|
||||
textNodesWithPosition = textNodesWithPosition.filter(Boolean);
|
||||
|
||||
for (const element of textNodesWithPosition) {
|
||||
const { text, pos } = element;
|
||||
const matches = Array.from(text.matchAll(searchTerm)).filter(
|
||||
([matchText]) => matchText.trim(),
|
||||
);
|
||||
|
||||
for (const m of matches) {
|
||||
if (m[0] === "") break;
|
||||
|
||||
if (m.index !== undefined) {
|
||||
results.push({
|
||||
from: pos + m.index,
|
||||
to: pos + m.index + m[0].length,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
// Shared find-all-occurrences primitive (also used by multi-cursor).
|
||||
const results: Range[] = findOccurrences(doc, searchTerm);
|
||||
|
||||
for (let i = 0; i < results.length; i += 1) {
|
||||
const r = results[i];
|
||||
|
||||
Reference in New Issue
Block a user