Compare commits
37 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 9f0e880e0b | |||
| d35956b9e9 | |||
| c0e39a395c | |||
| c144abcb83 | |||
| 71517552b7 | |||
| 3f217f4835 | |||
| 791a707eb6 | |||
| 60205398bb | |||
| 38a09c5ca1 | |||
| 1b05224b27 | |||
| b50b32bf64 | |||
| 1d704d4ec5 | |||
| 2cb07ef7fb | |||
| 9ba97b6d2b | |||
| 0b9de2c25e | |||
| 123cba7de1 | |||
| 001eb1c923 | |||
| 002c931c6b | |||
| 307ad49324 | |||
| 1031ed1ae8 | |||
| 3b6634c6be | |||
| fdc37de3e8 | |||
| 4a750c1e7f | |||
| ea7c4d7cd2 | |||
| 255024fbe2 | |||
| b934fb2292 | |||
| 14da295185 | |||
| de8f9c804c | |||
| 8ebdfe156f | |||
| 31176d5f93 | |||
| 5f0c49d47c | |||
| b2e5b51420 | |||
| 26d9a70a1c | |||
| e1ebd79f30 | |||
| fc83ea28ab | |||
| ee15278c8f | |||
| 09ab92eccf |
+55
-3
@@ -225,11 +225,26 @@ MCP_DOCMOST_PASSWORD=
|
||||
|
||||
# Silence timeout (ms) for EXTERNAL-MCP transport ONLY (not the chat provider).
|
||||
# Tighter than AI_STREAM_TIMEOUT_MS so a byte-silent/hung MCP server is broken in
|
||||
# ~1 min instead of 15. Note it also cuts a legitimately long but byte-silent
|
||||
# single tool call (a slow crawl that emits nothing until done) and an SSE
|
||||
# transport idling >1 min BETWEEN tool calls. Default 60000 (1 min).
|
||||
# ~1 min instead of 15. It cuts a legitimately long but byte-silent single tool
|
||||
# call (a slow crawl that emits nothing until done) on the HTTP (streamable)
|
||||
# transport, which opens a fresh request per call. The SSE transport — one
|
||||
# long-lived body across many calls — is NO LONGER governed by this timeout
|
||||
# (as of #489): its idle-BETWEEN-calls window has its own, raised bodyTimeout,
|
||||
# AI_MCP_SSE_BODY_TIMEOUT_MS below. Default 60000 (1 min).
|
||||
# AI_MCP_STREAM_TIMEOUT_MS=60000
|
||||
|
||||
# bodyTimeout (ms) for the EXTERNAL-MCP SSE transport ONLY (#489). The SSE
|
||||
# transport holds ONE response body open across many tool calls, so undici's
|
||||
# bodyTimeout (time between body bytes) counts the LEGITIMATE silence BETWEEN the
|
||||
# model's tool calls, not just a hung single call. At the tight 1-min silence
|
||||
# timeout above, a normal >1-min gap between calls would break the SSE socket and
|
||||
# the cache would serve a dead client until TTL — so the SSE transport gets its
|
||||
# OWN, RAISED bodyTimeout. A single stuck call is still bounded by the per-call
|
||||
# cap (AI_MCP_CALL_TIMEOUT_MS), and a socket that does break is healed by the
|
||||
# in-run transport-error retry. The HTTP (streamable) transport keeps the tight
|
||||
# timeout. Default 600000 (10 min).
|
||||
# AI_MCP_SSE_BODY_TIMEOUT_MS=600000
|
||||
|
||||
# Total wall-clock cap (ms) for ONE external MCP tool call (app-level, not
|
||||
# transport). Aborts a tool that keeps the socket warm (SSE heartbeats / trickle)
|
||||
# but never returns a result — which the silence timeout above never breaks.
|
||||
@@ -288,6 +303,29 @@ MCP_DOCMOST_PASSWORD=
|
||||
# registry is process-local).
|
||||
# AI_CHAT_RESUMABLE_STREAM=false
|
||||
|
||||
# --- Run lifecycle tunables (#487) ---
|
||||
# These govern the universal run machinery (every turn is now a first-class run,
|
||||
# both modes) and rarely need changing.
|
||||
#
|
||||
# How long a server-side SUPERSEDE ("interrupt and send now") waits for the target
|
||||
# run to settle after issuing Stop before it degrades to a 409 SUPERSEDE_TIMEOUT
|
||||
# (nothing sent, the composer keeps the user's text). 10s is generous under a
|
||||
# healthy DB; do NOT raise it to paper over a slow DB — a SUPERSEDE_TIMEOUT is the
|
||||
# honest signal. Default 10000 (10s).
|
||||
# AI_CHAT_SUPERSEDE_TIMEOUT_MS=10000
|
||||
#
|
||||
# How often the periodic bidirectional reconcile job runs (heals runs/messages
|
||||
# left dangling by a crash or a lost terminal write). Default 120000 (2 min).
|
||||
# AI_CHAT_RECONCILE_INTERVAL_MS=120000
|
||||
#
|
||||
# Wall-clock cap for a SINGLE in-app tool call (a long paginated read, or a content
|
||||
# write whose collab commit hangs) — the per-call half of the composite abort
|
||||
# signal every in-app tool is wrapped with (the other half is the turn's Stop).
|
||||
# The reconcile staleness floor is derived as max(2 x this cap, 15min), so a very
|
||||
# high value delays stale-run recovery (the server boot-warns above 30min). Default
|
||||
# 120000 (2 min).
|
||||
# AI_CHAT_INAPP_TOOL_CALL_CAP_MS=120000
|
||||
|
||||
# --- 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
|
||||
@@ -335,6 +373,20 @@ MCP_DOCMOST_PASSWORD=
|
||||
# VictoriaMetrics/Prometheus reaching it as <host>:<port>/metrics.
|
||||
# METRICS_PORT=9464
|
||||
#
|
||||
# METRICS_BIND — interface the /metrics listener binds to. DEFAULT 127.0.0.1
|
||||
# (loopback only), so the unauthenticated endpoint is NOT exposed on all
|
||||
# interfaces. If the scraper runs in a SEPARATE container and reaches this as
|
||||
# docmost:9464, set METRICS_BIND=0.0.0.0 — but then also set METRICS_TOKEN
|
||||
# and/or keep the port on a private network, since /metrics is otherwise open.
|
||||
# METRICS_BIND=127.0.0.1
|
||||
#
|
||||
# METRICS_TOKEN — optional Bearer token guarding /metrics. When set, every
|
||||
# scrape MUST send `Authorization: Bearer <token>` (others get 401). Configure
|
||||
# the scraper with the same bearer token (e.g. VictoriaMetrics/vmagent
|
||||
# `bearer_token`, Prometheus `authorization.credentials`). Leave unset only
|
||||
# when the endpoint is bound to loopback or an otherwise-trusted network.
|
||||
# METRICS_TOKEN=
|
||||
#
|
||||
# 2) CLIENT_TELEMETRY_ENABLED — the public client perf-telemetry sink.
|
||||
# OFF by default. When true, the unauthenticated POST /api/telemetry/vitals
|
||||
# endpoint is registered and browsers collect + send web-vitals / editor
|
||||
|
||||
@@ -29,6 +29,10 @@ packages/mcp/build/
|
||||
# is a build artifact like build/ — never committed, always fresh.
|
||||
packages/mcp/src/registry-stamp.generated.ts
|
||||
|
||||
# token-estimate compiled output (#490; built in CI/Docker via `pnpm build` /
|
||||
# the server `pretest`, never committed, so src/ and prod can never diverge).
|
||||
packages/token-estimate/dist/
|
||||
|
||||
# Logs
|
||||
logs
|
||||
*.log
|
||||
|
||||
@@ -455,7 +455,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.
|
||||
- `core/ai-chat/ai-chat-run.service.ts` + `ai_chat_runs` — **every agent turn is now a first-class server-side RUN** (`#184`, universalized in `#487`): its lifecycle is tracked in `ai_chat_runs` in **both** modes, and the single-active-run-per-chat concurrency gate is enforced universally (a legacy second tab now gets a clean `409 A_RUN_ALREADY_ACTIVE` instead of a second parallel stream that interleaved history). The per-workspace `settings.ai.autonomousRuns` flag (off by default) **no longer gates whether a turn is a run** — it now controls **only the browser-disconnect semantics**: when ON the run is *detached* (a disconnect leaves it executing server-side; only an explicit `POST /ai-chat/stop` ends it, and a client reconnects/live-follows via `POST /ai-chat/run`); when OFF (legacy) a disconnect ends the turn by stopping its run via the run's stop lever. `#487` also adds a server-side **supersede** CAS ("interrupt and send now") to `POST /ai-chat/stream` (`supersede: { runId }`): it atomically stops the chat's currently-active run and waits for it to settle before the new turn claims the slot, returning `SUPERSEDE_INVALID` / `SUPERSEDE_TARGET_MISMATCH` / `SUPERSEDE_TIMEOUT` on the non-proceed branches. **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:
|
||||
@@ -470,7 +470,7 @@ Vite SPA. Code is organized by feature under `apps/client/src/features/*` (mirro
|
||||
- **Errors must never be swallowed or shown as generic messages.** Every caught error MUST (1) be logged in full to the console/logger — error name, message, stack, `cause`, and (for HTTP/provider failures) the status code and response body — and (2) be surfaced to the user with a *specific, human-readable explanation of what actually went wrong*, never a bare generic string like "Something went wrong" / "Could not start recording" / "Transcription failed". Include the real reason (the underlying error/provider message) in the user-facing text. On the server, wrap third-party/provider failures with `describeProviderError` (or equivalent) and rethrow as a meaningful HTTP status + message — never let them collapse into an opaque 500. On the client, `console.error(<context>, err)` the raw error AND show the extracted reason (e.g. `err.response?.data?.message`, or the error `name: message`) in the notification.
|
||||
- The version string shown in the UI comes from `APP_VERSION` (CI/Docker) or `git describe --tags --always` (local), resolved in `vite.config.ts` — not from `package.json`.
|
||||
- Server TS config is permissive (`noImplicitAny: false`, `strictNullChecks: false`, `no-explicit-any` lint disabled). Follow the existing relaxed style rather than tightening types broadly.
|
||||
- Dependency versions are heavily pinned via `pnpm.overrides` and `pnpm.patchedDependencies` (`scimmy`, `yjs`, `ai`) in the root `package.json`. Don't bump pinned/patched deps casually; the patches and overrides exist for compatibility/security reasons. The `ai@6.0.134` patch disables the SDK's O(n²) cumulative `partialOutput` accumulation when no output strategy is requested (server heap OOM on long agent runs, #184; tripwire test: `apps/server/src/integrations/ai/ai-sdk-partial-output.patch.spec.ts`) — it MUST be re-created via `pnpm patch` when bumping `ai`.
|
||||
- Dependency versions are heavily pinned via `pnpm.overrides` and `pnpm.patchedDependencies` (`scimmy`, `yjs`, `ai`) in the root `package.json`. Don't bump pinned/patched deps casually; the patches and overrides exist for compatibility/security reasons. The `ai@6.0.134` patch carries TWO independent server fixes, each with its own tripwire test: (1) it disables the SDK's O(n²) cumulative `partialOutput` accumulation when no output strategy is requested (server heap OOM on long agent runs, #184; tripwire: `apps/server/src/integrations/ai/ai-sdk-partial-output.patch.spec.ts`); (2) it fixes `writeToServerResponse`'s drain-hang — the loop awaited only `"drain"` under backpressure, so a mid-write client disconnect parked the pipe forever and leaked the reader/buffers until restart; it now races `"drain"` against `"close"`/`"error"`, cancels the reader on disconnect, and swallows the fire-and-forget read rejection (#486; tripwire: `apps/server/src/integrations/ai/ai-sdk-drain-hang.patch.spec.ts`). Both tripwires assert BOTH installed dist builds carry their patch marker. The patch MUST be re-created via `pnpm patch` when bumping `ai`.
|
||||
- **The MCP tool inventory in `SERVER_INSTRUCTIONS` is GENERATED from the registry** (`packages/mcp/src/server-instructions.ts`: `buildToolInventory()` over `SHARED_TOOL_SPECS`) and spliced into the hand-written routing prose (`ROUTING_PROSE`). So adding/renaming/removing a **shared** spec in `packages/mcp/src/tool-specs.ts` auto-updates the `<tool_inventory>` — no manual `SERVER_INSTRUCTIONS` edit needed. Only an **inline** MCP-only tool (those registered via `server.registerTool(...)` in `index.ts`, not through the registry) needs a one-line entry in `INLINE_MCP_INVENTORY`. Enforced by `packages/mcp/test/unit/tool-inventory.test.mjs`, which fails when a registered tool is missing from the generated inventory (there is no `EXCEPTIONS` opt-out anymore — every tool must appear). Update `ROUTING_PROSE` when a tool's *intent guidance* (when-to-use) changes. `packages/mcp/build/` is gitignored and rebuilt in CI/Docker via `pnpm build` (same convention as `git-sync`/`prosemirror-markdown`) — never commit it; rebuild locally after editing to run the tests.
|
||||
|
||||
## CI / release
|
||||
|
||||
+102
@@ -115,6 +115,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
the old ProseMirror-JSON output. Released together with the `#411`/`#412`
|
||||
breaking window so external configs break exactly once. (#413)
|
||||
|
||||
- **The Prometheus `/metrics` listener now binds to `127.0.0.1` (loopback) by
|
||||
default instead of `0.0.0.0` (all interfaces).** This closes an unauthenticated
|
||||
endpoint that was previously reachable on every interface. **DEPLOY MIGRATION —
|
||||
cross-container scraping breaks silently otherwise:** if your scraper runs in a
|
||||
SEPARATE container and reaches the app as `docmost:9464` (the exact topology the
|
||||
old `0.0.0.0` hardcode served), you MUST now set `METRICS_BIND=0.0.0.0` — and,
|
||||
because that re-exposes the endpoint, also set `METRICS_TOKEN=<secret>` and
|
||||
configure the scraper with a matching Bearer token. Without `METRICS_BIND`, the
|
||||
scraper can no longer connect and metrics go dark with no error. See the
|
||||
`METRICS_BIND` / `METRICS_TOKEN` block in `.env.example` for the migration.
|
||||
Same-host (loopback) scrapers need no change. (#486)
|
||||
|
||||
### Added
|
||||
|
||||
- **Place several images side by side in a row.** A new "Inline (side by
|
||||
@@ -190,6 +202,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
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)
|
||||
- **Server-side "interrupt and send now" (supersede) for AI chat.** `POST
|
||||
/ai-chat/stream` now accepts a `supersede: { runId }` field: when the user sends
|
||||
a new message while a run is active, the server atomically stops that run and
|
||||
waits for it to settle before the new turn claims the chat's single run slot,
|
||||
instead of the send being rejected as concurrent. The compare-and-set surfaces
|
||||
three codes on its non-proceed branches — `SUPERSEDE_INVALID` (the targeted run
|
||||
is malformed / belongs to another chat), `SUPERSEDE_TARGET_MISMATCH` (a
|
||||
different run is now active; carries the current `activeRunId`), and
|
||||
`SUPERSEDE_TIMEOUT` (the previous run did not stop within the settle window, so
|
||||
nothing was sent and the composer keeps the text). Tunable via
|
||||
`AI_CHAT_SUPERSEDE_TIMEOUT_MS` (default 10s). (#487)
|
||||
- **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
|
||||
@@ -270,6 +293,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
### Changed
|
||||
|
||||
- **Every AI-chat turn is now a first-class server-side run, and one run per chat
|
||||
is enforced in both modes.** The run machinery from `#184` was universalized: a
|
||||
turn is tracked in `ai_chat_runs` and gated by the single-active-run-per-chat
|
||||
index regardless of the `settings.ai.autonomousRuns` flag. **Behavior change:**
|
||||
a second tab (or a double-submit) that starts a turn while one is already active
|
||||
on the chat is now rejected up front with `409 A_RUN_ALREADY_ACTIVE` (carrying
|
||||
the `activeRunId`); previously, on the legacy path, it opened a second parallel
|
||||
stream on the same chat that interleaved history. The `autonomousRuns` flag no
|
||||
longer controls whether a turn is a run — it now governs **only** the
|
||||
browser-disconnect semantics (ON = detached/survives a disconnect; OFF = a
|
||||
disconnect stops the run). (#487)
|
||||
- **Client markdown paste/copy and AI-chat rendering now go through the canonical
|
||||
converter.** Pasting markdown into the editor, "Copy as markdown", the AI title
|
||||
generator, and the AI-chat markdown renderer all now use
|
||||
@@ -302,6 +336,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
### Fixed
|
||||
|
||||
- **A chat with one malformed message part no longer 500s on every turn, and a
|
||||
failed send no longer duplicates the user's message.** Incoming client parts
|
||||
are now whitelisted to `text` (a forged tool-result part can no longer reach
|
||||
the persisted history or the model context), and the turn is converted BEFORE
|
||||
the user row is inserted, so a mid-flight failure cannot leave a duplicate
|
||||
user row that a retry then compounds. A single part that still fails to convert
|
||||
degrades to a `[tool context omitted]` marker on that one row instead of
|
||||
bricking the whole chat. (#489)
|
||||
- **A transport drop to an external MCP server now heals within the same turn.**
|
||||
On an undici transport error, a read-only MCP tool reconnects its server and
|
||||
retries once within the run; a write is never auto-retried (it may already have
|
||||
applied). One flapping server no longer nulls the shared client cache, so other
|
||||
servers' cached clients are untouched. The SSE transport also gets a raised
|
||||
body-timeout so a legitimate >1-min idle between the model's tool calls no
|
||||
longer breaks a long-lived SSE socket (new `AI_MCP_SSE_BODY_TIMEOUT_MS`, default
|
||||
10 min; see `.env.example`). (#489)
|
||||
|
||||
- **The server no longer runs out of heap during long autonomous agent runs.** A
|
||||
new pnpm patch on `ai@6.0.134` stops the SDK from building a cumulative
|
||||
snapshot of the ENTIRE turn text on every streamed text-delta when no output
|
||||
@@ -310,6 +361,39 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
`tee()` branch of the stream result — a ~20-step, ~28k-chunk agent run
|
||||
retained ~1.7 GB and OOM'd the 2 GB JS heap. Streaming granularity is
|
||||
unchanged; the patch must be re-created if `ai` is ever bumped. (#184)
|
||||
|
||||
- **The server no longer leaks a hung stream pipe on every mid-run client
|
||||
disconnect.** The same `ai@6.0.134` pnpm patch now also fixes the SDK's
|
||||
`writeToServerResponse`, which awaited only a `"drain"` event under
|
||||
backpressure: when a client disconnected mid-write the socket never drained, so
|
||||
the write loop parked forever, `response.end()` was unreachable, and the stream
|
||||
reader plus buffered chunks were pinned until process restart (every mid-run
|
||||
disconnect in autonomous mode leaked one). The patch races `"drain"` against
|
||||
`"close"`/`"error"`, cancels the reader and ends the response on disconnect, and
|
||||
swallows the fire-and-forget read rejection instead of crashing on an
|
||||
unhandledRejection. (#486)
|
||||
|
||||
- **A failed autonomous agent-run start no longer becomes an unstoppable ghost
|
||||
run.** When `beginRun` failed for a transient reason (e.g. a DB-pool blip),
|
||||
the turn previously continued with NO run row — invisible to `/stop`, not
|
||||
aborted on disconnect, and able to slip a second run past the one-run-per-chat
|
||||
gate, leaving an unstoppable run until restart. The turn now fails fast with an
|
||||
honest `503 A_RUN_BEGIN_FAILED` before the first byte (no orphan state), and the
|
||||
client shows a "temporary — please try again" message instead of a misleading
|
||||
"provider not configured". (#486)
|
||||
|
||||
- **A pathological draw.io graph can no longer wedge the whole server.** The ELK
|
||||
auto-layout (`layout:"elk"`) ran elkjs synchronously on the main event loop, so
|
||||
a graph at the node/edge cap blocked ALL HTTP/SSE/loopback traffic while it
|
||||
churned — and the old `setTimeout` "timeout" could never fire because the same
|
||||
thread was blocked. Layout now runs in a worker thread with the timeout enforced
|
||||
by `worker.terminate()`; the main loop stays responsive. (#486)
|
||||
|
||||
- **The `/health` Redis probe no longer leaks a client on every tick while Redis
|
||||
is down.** It built a new `ioredis` client per probe and disconnected it only on
|
||||
success, so during an outage each health tick added another forever-reconnecting
|
||||
client (an unbounded handle leak). A single long-lived probe client is now
|
||||
reused and closed on shutdown. (#486)
|
||||
- **Internal links in exported Markdown no longer lose their visible text.** A
|
||||
link whose target page name had no file extension (e.g. a bare title) was
|
||||
collapsed to empty text during export, producing an unclickable, label-less
|
||||
@@ -386,6 +470,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
share); any other value now returns the generic "not found" instead of
|
||||
serving the page. (#218)
|
||||
|
||||
- **Tool and provider error text no longer leaks to anonymous readers in the
|
||||
public-share AI chat.** A failing tool's raw error (which could carry an
|
||||
internal page title or a stack fragment) and a provider error (which bundles the
|
||||
provider `statusCode` and response body — potentially the internal baseUrl or
|
||||
model name) were streamed verbatim to the anonymous reader over SSE. Errors are
|
||||
now sanitized at the source: the share toolset collapses any unclassified tool
|
||||
error to a safe generic string (safe, classified tool messages still pass
|
||||
through for the model's self-correction), and the anonymous stream `onError`
|
||||
maps provider failures to a fixed set of neutral strings — the full detail goes
|
||||
only to the server log. A UI render gate is layered on top. (closes #394)
|
||||
|
||||
- **The Prometheus `/metrics` endpoint can now require Bearer authentication and
|
||||
is loopback-bound by default.** Previously it listened on all interfaces with no
|
||||
auth. Setting `METRICS_TOKEN` requires every scrape to present
|
||||
`Authorization: Bearer <token>` (compared in constant time), and the listener
|
||||
defaults to `127.0.0.1` (see the Breaking Changes entry for the cross-container
|
||||
migration). (#486)
|
||||
|
||||
## [0.94.0] - 2026-06-26
|
||||
|
||||
This release makes AI chat durable and fast: assistant turns are persisted to
|
||||
|
||||
@@ -22,6 +22,7 @@
|
||||
"@casl/react": "5.0.1",
|
||||
"@docmost/editor-ext": "workspace:*",
|
||||
"@docmost/prosemirror-markdown": "workspace:*",
|
||||
"@docmost/token-estimate": "workspace:*",
|
||||
"@excalidraw/excalidraw": "0.18.0-3a5ef40",
|
||||
"@mantine/core": "8.3.18",
|
||||
"@mantine/dates": "8.3.18",
|
||||
|
||||
@@ -203,6 +203,52 @@ describe("ChatThread — send now (#198)", () => {
|
||||
});
|
||||
});
|
||||
|
||||
// #486: the final onFinish -> flushNext() must be gated on the live-mount flag.
|
||||
// A clean onFinish can land AFTER the thread unmounts (New-chat / chat-switch
|
||||
// mid-stream — the async attach/resume settles late); flushing then dequeues and
|
||||
// re-POSTs a queued message from an abandoned thread (a "ghost" send).
|
||||
describe("ChatThread — onFinish flush gated on mount (#486)", () => {
|
||||
beforeEach(resetState);
|
||||
afterEach(cleanup);
|
||||
|
||||
it("a clean onFinish WHILE MOUNTED flushes the queued message (control)", () => {
|
||||
renderThread();
|
||||
fireEvent.click(screen.getByTestId("queue-btn")); // enqueue "queued text"
|
||||
expect(h.state.sendMessage).not.toHaveBeenCalled();
|
||||
|
||||
act(() => {
|
||||
h.state.onFinish?.({
|
||||
message: { id: "a", role: "assistant", parts: [] },
|
||||
isAbort: false,
|
||||
isDisconnect: false,
|
||||
isError: false,
|
||||
});
|
||||
});
|
||||
// Mounted: the queue flushes normally.
|
||||
expect(h.state.sendMessage).toHaveBeenCalledWith({ text: "queued text" });
|
||||
});
|
||||
|
||||
it("a clean onFinish AFTER unmount does NOT flush (no ghost send)", () => {
|
||||
const { unmount } = renderThread();
|
||||
fireEvent.click(screen.getByTestId("queue-btn")); // enqueue "queued text"
|
||||
h.state.sendMessage.mockClear();
|
||||
|
||||
// Chat switched away mid-stream: the streamer unmounts...
|
||||
unmount();
|
||||
// ...and a late, clean onFinish lands on the abandoned thread.
|
||||
act(() => {
|
||||
h.state.onFinish?.({
|
||||
message: { id: "a", role: "assistant", parts: [] },
|
||||
isAbort: false,
|
||||
isDisconnect: false,
|
||||
isError: false,
|
||||
});
|
||||
});
|
||||
// Gated on mountedRef: NOTHING is sent from the dead thread.
|
||||
expect(h.state.sendMessage).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
// #396: in autonomous mode a live sendNow must additionally request the
|
||||
// AUTHORITATIVE server stop of the detached run (a local abort is only a client
|
||||
// disconnect the server ignores) and arm a bounded 409 retry so the re-POST
|
||||
|
||||
@@ -659,7 +659,13 @@ export default function ChatThread({
|
||||
return;
|
||||
}
|
||||
if (isAbort || isDisconnect || isError) return;
|
||||
flushNext();
|
||||
// Gate the final flush on the live-mount flag (#486): a clean onFinish can
|
||||
// land AFTER this thread unmounted (a New-chat / chat-switch mid-stream —
|
||||
// the async attach/resume settles late). Flushing then dequeues and POSTs a
|
||||
// queued message from an abandoned thread — a "ghost" send / ghost chat.
|
||||
// Every other queue side effect already guards on mountedRef; this last one
|
||||
// was the gap.
|
||||
if (mountedRef.current) flushNext();
|
||||
},
|
||||
// `onError` runs in addition to `onFinish` (which ai@6 also calls on error).
|
||||
// Log the raw failure here for devtools; the UI shows a friendly classified
|
||||
|
||||
@@ -47,6 +47,13 @@ interface MessageItemProps {
|
||||
* agent's raw query/argument text.
|
||||
*/
|
||||
showInput?: boolean;
|
||||
/**
|
||||
* Forwarded to ToolCallCard: whether a failed tool card renders its raw
|
||||
* errorText. Defaults to true (internal chat). The public share passes false so
|
||||
* internal detail in a tool error is never painted (belt to the server-side
|
||||
* byte sanitization).
|
||||
*/
|
||||
showErrors?: boolean;
|
||||
/**
|
||||
* Neutralize internal/relative markdown links in the rendered answer (drop
|
||||
* their href so they become inert text). Defaults to false (internal chat,
|
||||
@@ -125,6 +132,7 @@ function MessageItem({
|
||||
message,
|
||||
showCitations = true,
|
||||
showInput = true,
|
||||
showErrors = true,
|
||||
neutralizeInternalLinks = false,
|
||||
assistantName,
|
||||
turnStreaming = false,
|
||||
@@ -219,6 +227,7 @@ function MessageItem({
|
||||
part={part as unknown as ToolUiPart}
|
||||
showCitations={showCitations}
|
||||
showInput={showInput}
|
||||
showErrors={showErrors}
|
||||
/>
|
||||
);
|
||||
}
|
||||
@@ -284,6 +293,7 @@ export function arePropsEqual(
|
||||
prev.signature === next.signature &&
|
||||
prev.showCitations === next.showCitations &&
|
||||
prev.showInput === next.showInput &&
|
||||
prev.showErrors === next.showErrors &&
|
||||
prev.neutralizeInternalLinks === next.neutralizeInternalLinks &&
|
||||
prev.assistantName === next.assistantName &&
|
||||
// The turn-end flip re-renders every row once (cheap, terminal event) —
|
||||
|
||||
@@ -32,6 +32,12 @@ interface MessageListProps {
|
||||
* doesn't see the agent's raw query/argument text.
|
||||
*/
|
||||
showInput?: boolean;
|
||||
/**
|
||||
* Forwarded to MessageItem -> ToolCallCard: whether a failed tool card renders
|
||||
* its raw errorText. Defaults to true (internal chat). The public share passes
|
||||
* false so internal detail in a tool error is never painted.
|
||||
*/
|
||||
showErrors?: boolean;
|
||||
/**
|
||||
* Forwarded to MessageItem: neutralize internal/relative markdown links in
|
||||
* the rendered answers (drop their href so they render as inert text).
|
||||
@@ -127,6 +133,7 @@ export default function MessageList({
|
||||
emptyState,
|
||||
showCitations = true,
|
||||
showInput = true,
|
||||
showErrors = true,
|
||||
neutralizeInternalLinks = false,
|
||||
assistantName,
|
||||
}: MessageListProps) {
|
||||
@@ -217,6 +224,7 @@ export default function MessageList({
|
||||
signature={messageSignature(message)}
|
||||
showCitations={showCitations}
|
||||
showInput={showInput}
|
||||
showErrors={showErrors}
|
||||
neutralizeInternalLinks={neutralizeInternalLinks}
|
||||
assistantName={assistantName}
|
||||
// Turn-level liveness, gated to the TAIL row: only the tail message
|
||||
|
||||
@@ -30,6 +30,16 @@ interface ToolCallCardProps {
|
||||
* the extra summary line, leaving the card (the action log) intact.
|
||||
*/
|
||||
showInput?: boolean;
|
||||
/**
|
||||
* Whether to render the tool's raw errorText on a failed call. Defaults to true
|
||||
* (the internal chat, where the operator may debug). The public share passes
|
||||
* false: a tool error string can carry internal detail (an internal page title,
|
||||
* a stack fragment, a provider message). This is the RENDER gate only — the
|
||||
* authoritative fix also sanitizes the bytes server-side (see
|
||||
* PublicShareChatToolsService.forShare), so a share reader never receives raw
|
||||
* error text over the wire, not just never sees it painted (#394).
|
||||
*/
|
||||
showErrors?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -41,6 +51,7 @@ export default function ToolCallCard({
|
||||
part,
|
||||
showCitations = true,
|
||||
showInput = true,
|
||||
showErrors = true,
|
||||
}: ToolCallCardProps) {
|
||||
const { t } = useTranslation();
|
||||
const toolName = getToolName(part);
|
||||
@@ -74,7 +85,7 @@ export default function ToolCallCard({
|
||||
</Text>
|
||||
)}
|
||||
|
||||
{state === "error" && part.errorText && (
|
||||
{state === "error" && showErrors && part.errorText && (
|
||||
<Text size="xs" c="red" mt={2}>
|
||||
{part.errorText}
|
||||
</Text>
|
||||
|
||||
@@ -6,10 +6,13 @@ describe("estimateTokens", () => {
|
||||
expect(estimateTokens("")).toBe(0);
|
||||
});
|
||||
|
||||
it("ceils chars/4 so any non-empty text is at least 1 token", () => {
|
||||
// #490: migrated onto the shared @docmost/token-estimate module (chars/2.5, up
|
||||
// from the old client-only chars/4) so the client counter and the server replay
|
||||
// budgeter can never diverge.
|
||||
it("ceils chars/2.5 so any non-empty text is at least 1 token", () => {
|
||||
expect(estimateTokens("a")).toBe(1);
|
||||
expect(estimateTokens("abcd")).toBe(1);
|
||||
expect(estimateTokens("abcde")).toBe(2);
|
||||
expect(estimateTokens("12345678")).toBe(2);
|
||||
expect(estimateTokens("ab")).toBe(1);
|
||||
expect(estimateTokens("abcde")).toBe(2); // 5 / 2.5 = 2
|
||||
expect(estimateTokens("x".repeat(10))).toBe(4); // 10 / 2.5 = 4
|
||||
});
|
||||
});
|
||||
|
||||
@@ -2,18 +2,10 @@
|
||||
* Rough client-side token estimation for AI-chat UI affordances.
|
||||
*
|
||||
* No provider streams exact per-token usage mid-stream, so any in-flight figure
|
||||
* is a CLIENT ESTIMATE (chars/≈4 heuristic). Pure + unit-testable: it never runs
|
||||
* a real BPE tokenizer (that would be O(n²) on the hot path, bloat the bundle,
|
||||
* and be wrong for Gemini/Ollama anyway). Used by the in-body reasoning counter
|
||||
* ("Thinking · N tokens").
|
||||
* is a CLIENT ESTIMATE. This re-exports the SHARED estimator from
|
||||
* `@docmost/token-estimate` (chars/2.5) so the in-body counter and the server's
|
||||
* replay budgeter use the SAME heuristic — two divergent estimators would mean
|
||||
* "the badge shows 60%" while "the budgeter already trimmed" (#490). Used by the
|
||||
* in-body reasoning counter ("Thinking · N tokens").
|
||||
*/
|
||||
|
||||
/**
|
||||
* Rough token estimate for a piece of text using the standard chars/≈4 heuristic.
|
||||
* Returns 0 for empty/whitespace-free-of-content input, and ceils so any
|
||||
* non-empty text counts as at least one token.
|
||||
*/
|
||||
export function estimateTokens(text: string): number {
|
||||
if (!text) return 0;
|
||||
return Math.ceil(text.length / 4);
|
||||
}
|
||||
export { estimateTokens } from "@docmost/token-estimate";
|
||||
|
||||
@@ -23,6 +23,25 @@ describe("describeChatError", () => {
|
||||
});
|
||||
});
|
||||
|
||||
it("classifies an A_RUN_BEGIN_FAILED 503 as a temporary run-start failure, NOT provider-not-configured (#486)", () => {
|
||||
// The FULL real body the server writes for a beginRun failure: a
|
||||
// ServiceUnavailableException(object) whose response is serialized verbatim
|
||||
// onto the raw socket, self-describing statusCode 503 + the run-start code.
|
||||
const body =
|
||||
'{"message":"Could not start the agent run. This is usually temporary — please try again.","code":"A_RUN_BEGIN_FAILED","statusCode":503}';
|
||||
expect(describeChatError(body, t)).toEqual({
|
||||
title: "Could not start the run",
|
||||
detail:
|
||||
"The agent run could not be started. This is usually temporary — please try again.",
|
||||
});
|
||||
// ORDER GUARD: even though the body ALSO carries statusCode 503 (which the
|
||||
// generic branch matches), the A_RUN_BEGIN_FAILED branch runs first, so it is
|
||||
// never mislabeled "AI provider not configured".
|
||||
expect(describeChatError(body, t).title).not.toBe(
|
||||
"AI provider not configured",
|
||||
);
|
||||
});
|
||||
|
||||
it("classifies a dropped connection (ECONNRESET) as a lost-connection error", () => {
|
||||
expect(
|
||||
describeChatError("Cannot connect to API: read ECONNRESET", t).title,
|
||||
|
||||
@@ -24,6 +24,21 @@ export function describeChatError(
|
||||
): ChatErrorView {
|
||||
const msg = message ?? "";
|
||||
|
||||
// Our own "could not start the run" gate (A_RUN_BEGIN_FAILED, #486): a 503
|
||||
// whose body carries this code is a TEMPORARY server-side failure while
|
||||
// starting the run (e.g. a DB-pool blip), NOT an unconfigured provider. It MUST
|
||||
// be matched STRICTLY BEFORE the generic 503 branch below, which would
|
||||
// otherwise mislabel it "The AI provider is not configured" and tell the user
|
||||
// to call an admin instead of just retrying.
|
||||
if (/"code"\s*:\s*"A_RUN_BEGIN_FAILED"/.test(msg)) {
|
||||
return {
|
||||
title: t("Could not start the run"),
|
||||
detail: t(
|
||||
"The agent run could not be started. This is usually temporary — please try again.",
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
if (/"statusCode"\s*:\s*403\b/.test(msg)) {
|
||||
return {
|
||||
title: t("AI chat is disabled"),
|
||||
|
||||
@@ -168,6 +168,10 @@ export default function ShareAiWidget({
|
||||
// Anonymous reader: suppress the tool-argument summary line so the
|
||||
// agent's raw query/argument text isn't shown on the public share.
|
||||
showInput={false}
|
||||
// Anonymous reader: never paint a tool's raw errorText (it can carry
|
||||
// internal detail). This is the render gate; the bytes are also
|
||||
// sanitized server-side in PublicShareChatToolsService.forShare (#394).
|
||||
showErrors={false}
|
||||
// Anonymous reader: neutralize internal/relative links in the
|
||||
// assistant's markdown so internal UUIDs/auth-gated routes don't
|
||||
// leak as clickable links (external http(s) links are kept).
|
||||
|
||||
@@ -23,7 +23,7 @@
|
||||
"migration:reset": "tsx src/database/migrate.ts down-to NO_MIGRATIONS",
|
||||
"migration:codegen": "kysely-codegen --dialect=postgres --camel-case --env-file=../../.env --out-file=./src/database/types/db.d.ts",
|
||||
"lint": "eslint \"{src,apps,libs,test}/**/*.ts\" --fix",
|
||||
"pretest": "pnpm --filter @docmost/editor-ext build && pnpm --filter @docmost/prosemirror-markdown build",
|
||||
"pretest": "pnpm --filter @docmost/editor-ext build && pnpm --filter @docmost/prosemirror-markdown build && pnpm --filter @docmost/token-estimate build",
|
||||
"test": "jest",
|
||||
"test:int": "jest --config test/jest-integration.json",
|
||||
"test:watch": "jest --watch",
|
||||
@@ -44,6 +44,7 @@
|
||||
"@docmost/mcp": "workspace:*",
|
||||
"@docmost/pdf-inspector": "1.9.6",
|
||||
"@docmost/prosemirror-markdown": "workspace:*",
|
||||
"@docmost/token-estimate": "workspace:*",
|
||||
"@fastify/compress": "^9.0.0",
|
||||
"@fastify/cookie": "^11.0.2",
|
||||
"@fastify/multipart": "^10.0.0",
|
||||
@@ -206,6 +207,7 @@
|
||||
"^@docmost/db/(.*)$": "<rootDir>/database/$1",
|
||||
"^@docmost/transactional/(.*)$": "<rootDir>/integrations/transactional/$1",
|
||||
"^@docmost/ee/(.*)$": "<rootDir>/ee/$1",
|
||||
"^@docmost/token-estimate$": "<rootDir>/../../../packages/token-estimate/src/index.ts",
|
||||
"^src/(.*)$": "<rootDir>/$1",
|
||||
"^@tiptap/react$": "<rootDir>/../test/stubs/tiptap-react.js"
|
||||
}
|
||||
|
||||
@@ -43,6 +43,9 @@ function makeRepo(overrides: Record<string, jest.Mock> = {}) {
|
||||
workspaceId: v.workspaceId,
|
||||
})),
|
||||
update: jest.fn(async () => ({ id: 'run-1' })),
|
||||
// #487: terminal finalize now goes through the CONDITIONAL write. Default
|
||||
// returns a truthy row (the run WAS active -> this call wrote it).
|
||||
finalizeIfActive: jest.fn(async () => ({ id: 'run-1', status: 'succeeded' })),
|
||||
markStopRequested: jest.fn(async () => ({ id: 'run-1' })),
|
||||
findActiveByChat: jest.fn(async () => undefined),
|
||||
findLatestByChat: jest.fn(async () => undefined),
|
||||
@@ -336,14 +339,12 @@ describe('AiChatRunService run lifecycle', () => {
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'error', 'provider blew up');
|
||||
|
||||
expect(svc.isLocallyActive('run-1')).toBe(false);
|
||||
expect(repo.update).toHaveBeenCalledWith(
|
||||
// #487: the terminal write is CONDITIONAL (finalizeIfActive); finishedAt is
|
||||
// stamped inside the repo method, so the service passes just status + error.
|
||||
expect(repo.finalizeIfActive).toHaveBeenCalledWith(
|
||||
'run-1',
|
||||
'ws-1',
|
||||
expect.objectContaining({
|
||||
status: 'failed',
|
||||
error: 'provider blew up',
|
||||
finishedAt: expect.any(Date),
|
||||
}),
|
||||
expect.objectContaining({ status: 'failed', error: 'provider blew up' }),
|
||||
);
|
||||
});
|
||||
|
||||
@@ -366,8 +367,8 @@ describe('AiChatRunService run lifecycle', () => {
|
||||
// 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(
|
||||
expect(repo.finalizeIfActive).toHaveBeenCalledTimes(1);
|
||||
expect(repo.finalizeIfActive).toHaveBeenCalledWith(
|
||||
'run-1',
|
||||
'ws-1',
|
||||
expect.objectContaining({ status: 'failed', error: 'first' }),
|
||||
@@ -389,8 +390,8 @@ describe('AiChatRunService run lifecycle', () => {
|
||||
const updateGate = new Promise((res) => {
|
||||
resolveUpdate = res;
|
||||
});
|
||||
const update = jest.fn(() => updateGate);
|
||||
const repo = makeRepo({ update });
|
||||
const finalizeIfActive = jest.fn(() => updateGate);
|
||||
const repo = makeRepo({ finalizeIfActive });
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await svc.beginRun({
|
||||
chatId: 'chat-1',
|
||||
@@ -399,23 +400,23 @@ describe('AiChatRunService run lifecycle', () => {
|
||||
});
|
||||
|
||||
// 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.
|
||||
// claims the entry (active.delete) and awaits the write; the second, started
|
||||
// in the same macrotask, finds the entry already gone and returns at the claim
|
||||
// WITHOUT ever writing.
|
||||
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);
|
||||
expect(finalizeIfActive).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Let the single in-flight update land; both calls resolve cleanly.
|
||||
resolveUpdate({ id: 'run-1' });
|
||||
resolveUpdate({ id: 'run-1', status: 'succeeded' });
|
||||
await Promise.all([p1, p2]);
|
||||
|
||||
expect(update).toHaveBeenCalledTimes(1);
|
||||
expect(finalizeIfActive).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(
|
||||
expect(finalizeIfActive).toHaveBeenCalledWith(
|
||||
'run-1',
|
||||
'ws-1',
|
||||
expect.objectContaining({ status: 'succeeded' }),
|
||||
@@ -431,10 +432,10 @@ describe('AiChatRunService run lifecycle', () => {
|
||||
// 409s until a restart. The fix updates FIRST and retries.
|
||||
let calls = 0;
|
||||
const repo = makeRepo({
|
||||
update: jest.fn(async () => {
|
||||
finalizeIfActive: jest.fn(async () => {
|
||||
calls += 1;
|
||||
if (calls === 1) throw new Error('deadlock detected');
|
||||
return { id: 'run-1' };
|
||||
return { id: 'run-1', status: 'succeeded' };
|
||||
}),
|
||||
});
|
||||
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined);
|
||||
@@ -447,26 +448,29 @@ describe('AiChatRunService run lifecycle', () => {
|
||||
|
||||
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'.
|
||||
// The retry landed the terminal write: the entry is dropped (slot freed), no
|
||||
// zombie left, and the row carries the real terminal status.
|
||||
expect(svc.isLocallyActive('run-1')).toBe(false);
|
||||
expect(repo.update).toHaveBeenCalledTimes(2);
|
||||
expect(repo.update).toHaveBeenLastCalledWith(
|
||||
expect(svc.hasZombie('run-1')).toBe(false);
|
||||
expect(repo.finalizeIfActive).toHaveBeenCalledTimes(2);
|
||||
expect(repo.finalizeIfActive).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 () => {
|
||||
it('#487 give-up: if the terminal write keeps failing, finalizeRun leaves a ZOMBIE (does NOT restore the entry) and settleZombie re-drives it', 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.
|
||||
// #487 changes the give-up behaviour: the entry is NOT restored (a restored
|
||||
// entry is indistinguishable from a live run). Instead a ZOMBIE record holds
|
||||
// the intended terminal status, and a re-drive (settleZombie — called by the
|
||||
// reconcile / supersede / opportunistic paths) applies it later.
|
||||
let healthy = false;
|
||||
const repo = makeRepo({
|
||||
update: jest.fn(async () => {
|
||||
finalizeIfActive: jest.fn(async () => {
|
||||
if (!healthy) throw new Error('pool exhausted');
|
||||
return { id: 'run-1' };
|
||||
return { id: 'run-1', status: 'succeeded' };
|
||||
}),
|
||||
});
|
||||
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined);
|
||||
@@ -480,35 +484,83 @@ describe('AiChatRunService run lifecycle', () => {
|
||||
userId: 'user-1',
|
||||
});
|
||||
|
||||
// First settle: every bounded attempt fails -> entry retained, NOT settled.
|
||||
// First settle: every bounded attempt fails -> ZOMBIE, entry NOT restored.
|
||||
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.
|
||||
expect(svc.isLocallyActive('run-1')).toBe(false); // NOT a live entry
|
||||
expect(svc.hasZombie('run-1')).toBe(true);
|
||||
expect(svc.zombieRunIds()).toContain('run-1');
|
||||
// The give-up emits ONE explicit, greppable ERROR mentioning the zombie.
|
||||
const gaveUp = errorSpy.mock.calls.some(
|
||||
(c) =>
|
||||
/NON-TERMINAL/.test(String(c[0])) &&
|
||||
/ZOMBIE/.test(String(c[0])) &&
|
||||
/run-1/.test(String(c[0])) &&
|
||||
/chat-1/.test(String(c[0])),
|
||||
);
|
||||
expect(gaveUp).toBe(true);
|
||||
// The settle notifier resolved as terminalWriteFailed (a subscriber learns the
|
||||
// slot still needs the intended status applied).
|
||||
const outcome = await svc.peekSettled('run-1');
|
||||
expect(outcome).toEqual({
|
||||
status: 'succeeded',
|
||||
error: null,
|
||||
terminalWriteFailed: true,
|
||||
});
|
||||
|
||||
// The DB recovers; a later settle now succeeds and frees the slot.
|
||||
// The DB recovers; a re-drive settles the zombie via the conditional UPDATE.
|
||||
healthy = true;
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'completed');
|
||||
expect(svc.isLocallyActive('run-1')).toBe(false);
|
||||
expect(repo.update).toHaveBeenLastCalledWith(
|
||||
const redriven = await svc.settleZombie('run-1');
|
||||
expect(redriven).toBe(true);
|
||||
expect(svc.hasZombie('run-1')).toBe(false);
|
||||
expect(repo.finalizeIfActive).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;
|
||||
// A later finalizeRun is idempotent (row already terminal): it no-ops at the
|
||||
// once-gate, never re-writing.
|
||||
const callsBefore = repo.finalizeIfActive.mock.calls.length;
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'error', 'late');
|
||||
expect(repo.update).toHaveBeenCalledTimes(callsBefore);
|
||||
expect(repo.finalizeIfActive).toHaveBeenCalledTimes(callsBefore);
|
||||
});
|
||||
|
||||
it('#487 double-settle collapses to a benign no-op (conditional write; notifier resolves once)', async () => {
|
||||
// A second concurrent settle is stopped at the synchronous active.delete
|
||||
// claim, so the terminal write runs exactly once and the notifier resolves
|
||||
// exactly once with the FIRST settler's outcome.
|
||||
const repo = makeRepo();
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'u1' });
|
||||
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'aborted');
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'error', 'late'); // no-op
|
||||
|
||||
expect(repo.finalizeIfActive).toHaveBeenCalledTimes(1);
|
||||
const outcome = await svc.peekSettled('run-1');
|
||||
// peekSettled after resolve+delete falls through (notifier dropped, no zombie)
|
||||
// -> undefined; the FIRST settler already resolved any earlier subscriber.
|
||||
expect(outcome).toBeUndefined();
|
||||
});
|
||||
|
||||
it('#487 late settledPromise subscriber gets the resolved outcome', async () => {
|
||||
const repo = makeRepo();
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'u1' });
|
||||
|
||||
// Subscribe BEFORE settle: hold the promise reference (as supersede does).
|
||||
const early = svc.peekSettled('run-1');
|
||||
expect(early).toBeDefined();
|
||||
|
||||
await svc.finalizeRun('run-1', 'ws-1', 'completed');
|
||||
|
||||
// The reference grabbed before settle resolves with the written outcome, even
|
||||
// though the notifier was dropped from the map on resolve (bounded).
|
||||
await expect(early).resolves.toEqual({
|
||||
status: 'succeeded',
|
||||
error: null,
|
||||
terminalWriteFailed: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('recordStep / linkAssistantMessage are best-effort: a repo failure is swallowed', async () => {
|
||||
@@ -525,3 +577,197 @@ describe('AiChatRunService run lifecycle', () => {
|
||||
).resolves.toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('#487 AiChatRunService.supersede (CAS)', () => {
|
||||
const chat = 'chat-1';
|
||||
const ws = 'ws-1';
|
||||
|
||||
it('degrade: no active run on the chat -> caller sends a normal turn', async () => {
|
||||
const repo = makeRepo({
|
||||
findById: jest.fn(async () => undefined),
|
||||
findActiveByChat: jest.fn(async () => undefined),
|
||||
});
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
expect(await svc.supersede(chat, 'run-x', ws)).toEqual({ kind: 'degrade' });
|
||||
});
|
||||
|
||||
it('invalid: the target run belongs to a DIFFERENT chat -> 400', async () => {
|
||||
const repo = makeRepo({
|
||||
findById: jest.fn(async () => ({
|
||||
id: 'run-x',
|
||||
chatId: 'other-chat',
|
||||
workspaceId: ws,
|
||||
})),
|
||||
});
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
expect(await svc.supersede(chat, 'run-x', ws)).toEqual({ kind: 'invalid' });
|
||||
});
|
||||
|
||||
it('mismatch: a DIFFERENT run is active than the one targeted -> current runId', async () => {
|
||||
const repo = makeRepo({
|
||||
findById: jest.fn(async () => ({ id: 'run-x', chatId: chat, workspaceId: ws })),
|
||||
findActiveByChat: jest.fn(async () => ({
|
||||
id: 'run-live',
|
||||
chatId: chat,
|
||||
workspaceId: ws,
|
||||
status: 'running',
|
||||
})),
|
||||
});
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
expect(await svc.supersede(chat, 'run-x', ws)).toEqual({
|
||||
kind: 'mismatch',
|
||||
activeRunId: 'run-live',
|
||||
});
|
||||
});
|
||||
|
||||
it('ready: the target IS active -> stop it, await its (fast) settle, free the slot', async () => {
|
||||
// Simulate a live long TOOL (NOT a slow UPDATE): the run stays active until an
|
||||
// explicit Stop unwinds it; commit-1's race makes that settle land quickly.
|
||||
// The abort listener stands in for streamText's onAbort -> finalizeRun.
|
||||
const repo = makeRepo({
|
||||
findById: jest.fn(async () => ({
|
||||
id: 'run-1',
|
||||
chatId: chat,
|
||||
workspaceId: ws,
|
||||
status: 'aborted',
|
||||
error: null,
|
||||
})),
|
||||
findActiveByChat: jest.fn(async () => ({
|
||||
id: 'run-1',
|
||||
chatId: chat,
|
||||
workspaceId: ws,
|
||||
status: 'running',
|
||||
})),
|
||||
});
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
const handle = await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' });
|
||||
handle.signal.addEventListener('abort', () => {
|
||||
void svc.finalizeRun('run-1', ws, 'aborted');
|
||||
});
|
||||
|
||||
// supersede: getRun -> getActiveByChat(==target) -> requestStop -> the abort
|
||||
// listener settles the run -> awaitSettled resolves -> ready.
|
||||
expect(await svc.supersede(chat, 'run-1', ws, 10_000)).toEqual({
|
||||
kind: 'ready',
|
||||
});
|
||||
expect(handle.signal.aborted).toBe(true); // Stop reached the run
|
||||
});
|
||||
|
||||
it('timeout: the target never settles within W -> 409 SUPERSEDE_TIMEOUT (nothing persisted)', async () => {
|
||||
const repo = makeRepo({
|
||||
findById: jest.fn(async () => ({ id: 'run-1', chatId: chat, workspaceId: ws })),
|
||||
findActiveByChat: jest.fn(async () => ({
|
||||
id: 'run-1',
|
||||
chatId: chat,
|
||||
workspaceId: ws,
|
||||
status: 'running',
|
||||
})),
|
||||
});
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' });
|
||||
// Do NOT settle the run: a tiny W elapses -> timeout.
|
||||
const result = await svc.supersede(chat, 'run-1', ws, 30);
|
||||
expect(result).toEqual({ kind: 'timeout' });
|
||||
});
|
||||
|
||||
it('ready then a DUPLICATE supersede POST degrades (the run is already gone)', async () => {
|
||||
let active: unknown = {
|
||||
id: 'run-1',
|
||||
chatId: chat,
|
||||
workspaceId: ws,
|
||||
status: 'running',
|
||||
};
|
||||
const repo = makeRepo({
|
||||
findById: jest.fn(async () => ({
|
||||
id: 'run-1',
|
||||
chatId: chat,
|
||||
workspaceId: ws,
|
||||
status: 'aborted',
|
||||
error: null,
|
||||
})),
|
||||
findActiveByChat: jest.fn(async () => active),
|
||||
finalizeIfActive: jest.fn(async () => {
|
||||
active = undefined; // settling frees the active slot
|
||||
return { id: 'run-1', status: 'aborted' };
|
||||
}),
|
||||
});
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
const handle = await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' });
|
||||
handle.signal.addEventListener('abort', () => {
|
||||
void svc.finalizeRun('run-1', ws, 'aborted');
|
||||
});
|
||||
|
||||
expect(await svc.supersede(chat, 'run-1', ws, 10_000)).toEqual({
|
||||
kind: 'ready',
|
||||
});
|
||||
// The duplicate POST for the same target now finds no active run -> degrade.
|
||||
expect(await svc.supersede(chat, 'run-1', ws)).toEqual({ kind: 'degrade' });
|
||||
});
|
||||
|
||||
it('reconcileStaleRuns: aborts a stale run with NO entry/zombie; NEVER touches a live entry', async () => {
|
||||
const finalizeIfActive = jest.fn(async () => ({ id: 'x', status: 'aborted' }));
|
||||
const repo = makeRepo({
|
||||
insert: jest.fn(async (v: any) => ({
|
||||
id: 'live-1',
|
||||
status: 'running',
|
||||
chatId: v.chatId,
|
||||
workspaceId: v.workspaceId,
|
||||
})),
|
||||
finalizeIfActive,
|
||||
findStaleActive: jest.fn(async () => [
|
||||
{ id: 'orphan-1', workspaceId: ws, chatId: 'c-orphan' },
|
||||
{ id: 'live-1', workspaceId: ws, chatId: 'c-live' },
|
||||
]),
|
||||
});
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
// A LIVE run this replica owns (in the `active` map).
|
||||
await svc.beginRun({ chatId: 'c-live', workspaceId: ws, userId: 'u1' });
|
||||
expect(svc.isLocallyActive('live-1')).toBe(true);
|
||||
|
||||
const aborted = await svc.reconcileStaleRuns(15 * 60 * 1000);
|
||||
expect(aborted).toBe(1);
|
||||
// The orphan (no entry) was aborted; the live entry was NEVER passed to the DB.
|
||||
expect(finalizeIfActive).toHaveBeenCalledTimes(1);
|
||||
expect(finalizeIfActive).toHaveBeenCalledWith(
|
||||
'orphan-1',
|
||||
ws,
|
||||
expect.objectContaining({ status: 'aborted' }),
|
||||
);
|
||||
expect(svc.isLocallyActive('live-1')).toBe(true);
|
||||
});
|
||||
|
||||
it('gave-up zombie: supersede applies the intended status (settleZombie) then is ready', async () => {
|
||||
let healthy = false;
|
||||
let active: unknown = {
|
||||
id: 'run-1',
|
||||
chatId: chat,
|
||||
workspaceId: ws,
|
||||
status: 'running',
|
||||
};
|
||||
const repo = makeRepo({
|
||||
findById: jest.fn(async () => ({ id: 'run-1', chatId: chat, workspaceId: ws })),
|
||||
findActiveByChat: jest.fn(async () => active),
|
||||
finalizeIfActive: jest.fn(async () => {
|
||||
if (!healthy) throw new Error('db down');
|
||||
active = undefined;
|
||||
return { id: 'run-1', status: 'aborted' };
|
||||
}),
|
||||
});
|
||||
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined);
|
||||
jest.spyOn(Logger.prototype, 'error').mockImplementation(() => undefined);
|
||||
const svc = new AiChatRunService(repo as never, makeEnv() as never);
|
||||
await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' });
|
||||
|
||||
// The run's terminal write gives up -> zombie (row still 'running').
|
||||
await svc.finalizeRun('run-1', ws, 'aborted');
|
||||
expect(svc.hasZombie('run-1')).toBe(true);
|
||||
|
||||
// The DB recovers; supersede awaits the (already-resolved, terminalWriteFailed)
|
||||
// settle, then settleZombie applies the intended status -> ready.
|
||||
healthy = true;
|
||||
expect(await svc.supersede(chat, 'run-1', ws, 10_000)).toEqual({
|
||||
kind: 'ready',
|
||||
});
|
||||
expect(svc.hasZombie('run-1')).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -34,6 +34,88 @@ export class RunAlreadyActiveError extends Error {
|
||||
export type TurnTerminalStatus = 'completed' | 'error' | 'aborted';
|
||||
export type RunTerminalStatus = 'succeeded' | 'failed' | 'aborted';
|
||||
|
||||
/** The terminal run statuses — the row is done once it reads one of these. */
|
||||
export const RUN_TERMINAL_STATUSES: readonly RunTerminalStatus[] = [
|
||||
'succeeded',
|
||||
'failed',
|
||||
'aborted',
|
||||
];
|
||||
|
||||
/** Whether a persisted run status is terminal (settled). */
|
||||
export function isRunTerminal(status: string | null | undefined): boolean {
|
||||
return (
|
||||
status === 'succeeded' || status === 'failed' || status === 'aborted'
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* #487: the outcome a run's {@link AiChatRunService.finalizeRun} settled with.
|
||||
* `terminalWriteFailed` = the terminal write GAVE UP after the bounded retry, so
|
||||
* the row is still non-terminal ('running') and a ZOMBIE record holds the
|
||||
* `intended` status for a later re-drive (reconcile / supersede / boot sweep). A
|
||||
* subscriber (supersede, #487 commit 3) uses this to decide whether the slot is
|
||||
* genuinely free or must first have the intended status applied.
|
||||
*/
|
||||
export interface RunSettleOutcome {
|
||||
status: RunTerminalStatus;
|
||||
error: string | null;
|
||||
terminalWriteFailed: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* #487: how long a supersede waits for the target run to settle after Stop before
|
||||
* it degrades to `SUPERSEDE_TIMEOUT`. W=10s is generous under a HEALTHY DB: commit
|
||||
* 1's race-on-abort makes an in-app tool abort->settle in ms/hundreds of ms, so a
|
||||
* live run releases its slot well within the window. Under a DB brownout the
|
||||
* timeout is normal (the write cannot land); W must NOT be raised to paper
|
||||
* over a slow DB — a SUPERSEDE_TIMEOUT is the honest signal (nothing persisted,
|
||||
* the composer keeps the user's text). Env-tunable for ops, default 10s.
|
||||
*/
|
||||
export const SUPERSEDE_SETTLE_TIMEOUT_MS = (() => {
|
||||
const raw = Number(process.env.AI_CHAT_SUPERSEDE_TIMEOUT_MS);
|
||||
return Number.isFinite(raw) && raw > 0 ? raw : 10_000;
|
||||
})();
|
||||
|
||||
/**
|
||||
* #487: the result of the supersede CAS ({@link AiChatRunService.supersede}).
|
||||
* - `degrade` : no active run on the chat (it ended between click and POST) —
|
||||
* the caller sends a NORMAL turn (NOT a mismatch);
|
||||
* - `invalid` : the target runId belongs to a DIFFERENT chat (malformed CAS 400);
|
||||
* - `mismatch` : a DIFFERENT run is active than the one the client targeted —
|
||||
* 409 SUPERSEDE_TARGET_MISMATCH carrying the current `activeRunId`
|
||||
* (the client does NOT auto-retry);
|
||||
* - `timeout` : the target did not settle within W — 409 SUPERSEDE_TIMEOUT,
|
||||
* nothing persisted;
|
||||
* - `ready` : the target was stopped AND settled (or its zombie's intended was
|
||||
* applied) — the slot is free; the caller may beginRun the new run.
|
||||
*/
|
||||
export type SupersedeResult =
|
||||
| { kind: 'degrade' }
|
||||
| { kind: 'invalid' }
|
||||
| { kind: 'mismatch'; activeRunId: string }
|
||||
| { kind: 'timeout' }
|
||||
| { kind: 'ready' };
|
||||
|
||||
/** A one-shot settle notifier (#487): `resolve` is called EXACTLY ONCE. */
|
||||
interface Deferred<T> {
|
||||
promise: Promise<T>;
|
||||
resolve: (value: T) => void;
|
||||
}
|
||||
|
||||
/**
|
||||
* #487: a run whose terminal write GAVE UP (every bounded attempt failed). The
|
||||
* row is stranded non-terminal ('running'); this record is the ONLY thing that
|
||||
* distinguishes it from a live run, and carries the `intended` terminal status so
|
||||
* a re-drive can apply it via the conditional UPDATE. Process-local (phase-1
|
||||
* single-process assumption): a restart drops it, and the boot sweep then writes
|
||||
* 'aborted' over the intended — a documented loss (see finalizeRun).
|
||||
*/
|
||||
interface ZombieRun {
|
||||
workspaceId: string;
|
||||
chatId: string;
|
||||
intended: { status: RunTerminalStatus; error: string | null };
|
||||
}
|
||||
|
||||
export function mapTurnStatusToRun(
|
||||
status: TurnTerminalStatus,
|
||||
): RunTerminalStatus {
|
||||
@@ -101,6 +183,22 @@ export class AiChatRunService implements OnModuleInit {
|
||||
// uptime — negligible in phase 1's single process.
|
||||
private readonly settled = new Set<string>();
|
||||
|
||||
// #487 runId -> one-shot settle notifier. Kept in a SEPARATE map from `active`
|
||||
// ON PURPOSE: it must OUTLIVE the `active.delete` claim inside finalizeRun (the
|
||||
// claim frees the slot the instant finalize starts), so a subscriber can still
|
||||
// await the outcome after the entry is gone. Created in beginRun, resolved
|
||||
// EXACTLY ONCE in finalizeRun, then removed (bounded). Absence => this replica
|
||||
// has no live notifier: a subscriber falls back to the zombie map, then to the
|
||||
// row (see peekSettled). Process-local (phase-1 single-process assumption).
|
||||
private readonly settledPromises = new Map<string, Deferred<RunSettleOutcome>>();
|
||||
|
||||
// #487 runId -> ZOMBIE record: a run whose terminal write gave up (row stranded
|
||||
// non-terminal). BOUNDED — an entry is added only on give-up and removed on a
|
||||
// successful re-drive (settleZombie) or when the row is found already terminal;
|
||||
// a process restart clears it (and the boot sweep settles the stranded row).
|
||||
// Process-local (phase-1 single-process assumption).
|
||||
private readonly zombies = new Map<string, ZombieRun>();
|
||||
|
||||
// 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
|
||||
@@ -224,6 +322,10 @@ export class AiChatRunService implements OnModuleInit {
|
||||
chatId: args.chatId,
|
||||
workspaceId: args.workspaceId,
|
||||
});
|
||||
// #487: arm the one-shot settle notifier BEFORE returning, so a subscriber
|
||||
// that races in immediately after begin always finds a promise to await. It
|
||||
// is resolved exactly once when the run settles (or gives up).
|
||||
this.settledPromises.set(run.id, this.makeDeferred<RunSettleOutcome>());
|
||||
return { runId: run.id, signal: controller.signal };
|
||||
}
|
||||
|
||||
@@ -263,47 +365,43 @@ export class AiChatRunService implements OnModuleInit {
|
||||
}
|
||||
|
||||
/**
|
||||
* 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.
|
||||
* Finalize a run to its terminal status (succeeded / failed / aborted) via a
|
||||
* CONDITIONAL UPDATE, stamping finishedAt + any error. Atomically safe against a
|
||||
* concurrent settle AND robust against a transient terminal-write failure.
|
||||
*
|
||||
* 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.
|
||||
* claim 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.
|
||||
*
|
||||
* 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.
|
||||
* ALL TERMINAL WRITES ARE CONDITIONAL (#487): `finalizeIfActive` only flips a
|
||||
* row still in pending|running (mirror of the assistant message's
|
||||
* `onlyIfStreaming`). So even a settle that DID reach the UPDATE (e.g. a
|
||||
* reconcile stamp racing an owner finalize) can never clobber a terminal status
|
||||
* — the loser matches nothing and is a benign no-op. `active.delete` is the
|
||||
* fast, in-process gate; the conditional WHERE is the authoritative one.
|
||||
*
|
||||
* 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.
|
||||
* ZOMBIE ON GIVE-UP (#487): if every bounded attempt THROWS (the DB is down for
|
||||
* the whole finalize), we do NOT restore the entry. The row is stranded
|
||||
* non-terminal ('running'); we record a ZOMBIE `{ terminalWriteFailed, intended
|
||||
* }` (the ONLY thing distinguishing this dead run from a live one) and resolve
|
||||
* the settle notifier with `terminalWriteFailed: true`. A restore would make the
|
||||
* zombie indistinguishable from a live run to every reader; instead a re-drive
|
||||
* (settleZombie, called by the periodic reconcile / supersede / opportunistic
|
||||
* paths) applies the intended status later via the same conditional UPDATE.
|
||||
*
|
||||
* DOCUMENTED LOSS (#487, single-process phase 1): if the process RESTARTS before
|
||||
* a zombie is re-driven, the in-memory zombie map is gone and the boot sweep
|
||||
* (unconditional) writes 'aborted' over the ACTUAL intended status. This is
|
||||
* unavoidable while the run lifecycle is single-process — there is no durable
|
||||
* record of `intended`; a cross-process durable intent is deferred to phase 2.
|
||||
*
|
||||
* IDEMPOTENT: the settle notifier resolves EXACTLY ONCE; a second settle is
|
||||
* stopped at `settled.has` or the `active.delete` claim, so a double-settle
|
||||
* collapses to a single write and can never double-resolve or clobber the row.
|
||||
*/
|
||||
async finalizeRun(
|
||||
runId: string,
|
||||
@@ -314,13 +412,17 @@ export class AiChatRunService implements OnModuleInit {
|
||||
// ---- 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.
|
||||
// Capture the entry BEFORE the delete for the give-up log context.
|
||||
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;
|
||||
|
||||
const status = mapTurnStatusToRun(turnStatus);
|
||||
const err = error ?? null;
|
||||
const chatId = entry?.chatId ?? 'unknown';
|
||||
|
||||
let lastError: unknown;
|
||||
for (
|
||||
let attempt = 1;
|
||||
@@ -328,47 +430,294 @@ export class AiChatRunService implements OnModuleInit {
|
||||
attempt++
|
||||
) {
|
||||
try {
|
||||
await this.runRepo.update(runId, workspaceId, {
|
||||
status: mapTurnStatusToRun(turnStatus),
|
||||
finishedAt: new Date(),
|
||||
error: error ?? null,
|
||||
const row = await this.runRepo.finalizeIfActive(runId, workspaceId, {
|
||||
status,
|
||||
error: err,
|
||||
});
|
||||
// Terminal write landed: arm the once-gate. The entry is already gone
|
||||
// (claimed above); we do NOT restore it. The slot is now free.
|
||||
// No throw => the row is now terminal (we wrote it, or it was ALREADY
|
||||
// terminal — another writer won the conditional UPDATE, a benign no-op).
|
||||
this.settled.add(runId);
|
||||
this.zombies.delete(runId);
|
||||
// Resolve with the persisted outcome: our status when WE wrote it, else
|
||||
// the row's real terminal status (re-read on the already-terminal path so
|
||||
// a subscriber never sees a status we did not actually persist).
|
||||
const outcome: RunSettleOutcome = row
|
||||
? { status, error: err, terminalWriteFailed: false }
|
||||
: await this.readTerminalOutcome(runId, workspaceId, status, err);
|
||||
this.resolveSettled(runId, outcome);
|
||||
return;
|
||||
} catch (err) {
|
||||
lastError = err;
|
||||
} catch (err2) {
|
||||
lastError = err2;
|
||||
this.logger.warn(
|
||||
`Failed to finalize run ${runId} (attempt ${attempt}/${
|
||||
AiChatRunService.FINALIZE_MAX_ATTEMPTS
|
||||
}): ${err instanceof Error ? err.message : 'unknown error'}`,
|
||||
}): ${err2 instanceof Error ? err2.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").
|
||||
// Every attempt threw: GIVE UP. The row is stranded non-terminal ('running').
|
||||
// Do NOT restore the entry (a restored entry is indistinguishable from a live
|
||||
// run); leave a ZOMBIE record instead, and resolve the notifier as
|
||||
// terminalWriteFailed so a subscriber knows the slot still needs the intended
|
||||
// status applied. One explicit, greppable ERROR so an operator can tell a
|
||||
// give-up from a per-attempt blip.
|
||||
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`,
|
||||
`Run ${runId} (chat ${chatId}) left NON-TERMINAL ('running'): terminal ` +
|
||||
`write failed after ${AiChatRunService.FINALIZE_MAX_ATTEMPTS} attempts; ` +
|
||||
`ZOMBIE recorded (intended '${status}'), recovery deferred to reconcile / ` +
|
||||
`supersede / 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);
|
||||
this.zombies.set(runId, {
|
||||
workspaceId,
|
||||
chatId,
|
||||
intended: { status, error: err },
|
||||
});
|
||||
this.resolveSettled(runId, { status, error: err, terminalWriteFailed: true });
|
||||
}
|
||||
|
||||
/**
|
||||
* #487: re-drive a zombie run's intended terminal write (the conditional
|
||||
* UPDATE). Called by the periodic reconcile (commit 4), an opportunistic
|
||||
* single-chat reconcile, and supersede (commit 3). On success — the row is now
|
||||
* terminal (written OR found already terminal) — the zombie is cleared and the
|
||||
* once-gate armed; on another failure the zombie is kept for a later retry.
|
||||
* Returns true when the row is now terminal. Best-effort; never throws.
|
||||
*/
|
||||
async settleZombie(runId: string): Promise<boolean> {
|
||||
const z = this.zombies.get(runId);
|
||||
if (!z) return false;
|
||||
try {
|
||||
await this.runRepo.finalizeIfActive(runId, z.workspaceId, {
|
||||
status: z.intended.status,
|
||||
error: z.intended.error,
|
||||
});
|
||||
this.zombies.delete(runId);
|
||||
this.settled.add(runId);
|
||||
return true;
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`Re-drive of zombie run ${runId} (chat ${z.chatId}) failed; will retry ` +
|
||||
`later: ${err instanceof Error ? err.message : 'unknown error'}`,
|
||||
);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* #487 reconcile clause (c): abort runs the DB still shows active (pending|
|
||||
* running) but that this replica does NOT own — NO live entry AND NO zombie —
|
||||
* and that have been UNTOUCHED past `staleMs` (from last-progress `updated_at`,
|
||||
* NOT startedAt, so a legit long marathon is never a candidate). "No entry" is
|
||||
* the PRIMARY gate: a live entry (an actively-executing run on this replica) is
|
||||
* NEVER aborted, whatever its age. Returns the number aborted. Best-effort —
|
||||
* never throws (a periodic-job failure must not crash the process).
|
||||
*/
|
||||
async reconcileStaleRuns(staleMs: number): Promise<number> {
|
||||
let candidates: Array<{ id: string; workspaceId: string; chatId: string }>;
|
||||
try {
|
||||
candidates = await this.runRepo.findStaleActive(staleMs);
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`Reconcile (stale runs) query failed: ${
|
||||
err instanceof Error ? err.message : 'unknown error'
|
||||
}`,
|
||||
);
|
||||
return 0;
|
||||
}
|
||||
let aborted = 0;
|
||||
for (const c of candidates) {
|
||||
// PRIMARY gate: never touch a live entry, and never race a zombie we are
|
||||
// already re-driving (settleZombie owns those).
|
||||
if (this.active.has(c.id) || this.zombies.has(c.id)) continue;
|
||||
try {
|
||||
const row = await this.runRepo.finalizeIfActive(c.id, c.workspaceId, {
|
||||
status: 'aborted',
|
||||
error: 'Run aborted by reconcile: no live runner (stale).',
|
||||
});
|
||||
if (row) {
|
||||
aborted += 1;
|
||||
this.settled.add(c.id);
|
||||
}
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`Reconcile abort of stale run ${c.id} failed: ${
|
||||
err instanceof Error ? err.message : 'unknown error'
|
||||
}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
return aborted;
|
||||
}
|
||||
|
||||
/**
|
||||
* #487: the run's settle outcome as seen by THIS replica, or undefined when it
|
||||
* has no record (the caller then reads the row — the DB is the source of truth).
|
||||
* A LIVE deferred (still settling, or resolved-but-not-yet-consumed) wins; a
|
||||
* ZOMBIE synthesizes the give-up outcome. A subscriber (supersede) races this
|
||||
* against a timeout.
|
||||
*/
|
||||
peekSettled(runId: string): Promise<RunSettleOutcome> | undefined {
|
||||
const d = this.settledPromises.get(runId);
|
||||
if (d) return d.promise;
|
||||
const z = this.zombies.get(runId);
|
||||
if (z) {
|
||||
return Promise.resolve({
|
||||
status: z.intended.status,
|
||||
error: z.intended.error,
|
||||
terminalWriteFailed: true,
|
||||
});
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* #487: await a run's settle outcome, bounded by `timeoutMs`. Returns the
|
||||
* outcome on settle, or undefined on TIMEOUT (or when this replica has no record
|
||||
* of the run and its row is not terminal). Uses the LIVE settle notifier / the
|
||||
* zombie synth when present; else reads the row (the DB is the source of truth
|
||||
* once the in-memory record is gone). The subscriber (supersede) grabs this
|
||||
* right after Stop; commit 1's race makes the settle land in ms on a healthy DB.
|
||||
*/
|
||||
async awaitSettled(
|
||||
runId: string,
|
||||
workspaceId: string,
|
||||
timeoutMs: number,
|
||||
): Promise<RunSettleOutcome | undefined> {
|
||||
const pending = this.peekSettled(runId);
|
||||
if (pending) {
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
const timeout = new Promise<undefined>((resolve) => {
|
||||
timer = setTimeout(() => resolve(undefined), timeoutMs);
|
||||
timer.unref?.();
|
||||
});
|
||||
try {
|
||||
return await Promise.race([pending, timeout]);
|
||||
} finally {
|
||||
if (timer) clearTimeout(timer);
|
||||
}
|
||||
}
|
||||
// No live notifier and no zombie: read the row (already settled-and-written,
|
||||
// or unknown here). A terminal row is an outcome; anything else -> undefined.
|
||||
const row = await this.runRepo.findById(runId, workspaceId);
|
||||
if (row && isRunTerminal(row.status)) {
|
||||
return {
|
||||
status: row.status as RunTerminalStatus,
|
||||
error: row.error ?? null,
|
||||
terminalWriteFailed: false,
|
||||
};
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* #487: the SERVER supersede CAS for `POST /stream { supersede: { runId: X } }`.
|
||||
* Atomically transitions "X is the chat's active run" -> "X is stopped, settled,
|
||||
* slot free" so the caller can start a replacement run. See {@link
|
||||
* SupersedeResult} for the branch semantics.
|
||||
*
|
||||
* On a `ready` result the caller MUST still go through the normal beginRun gate
|
||||
* (the partial unique index) — between the slot freeing here and beginRun a
|
||||
* neighbouring tab's ordinary POST can win the slot (documented SLOT-THEFT: the
|
||||
* loser then gets a MISMATCH carrying the NEW runId). There is also NO side-
|
||||
* effect quiescence: an in-flight write of the stopped run may still land AFTER
|
||||
* the new run starts (commit 1 stops the NEXT call, not one already committing),
|
||||
* so the caller adds a prompt note to the new run.
|
||||
*/
|
||||
async supersede(
|
||||
chatId: string,
|
||||
targetRunId: string,
|
||||
workspaceId: string,
|
||||
timeoutMs: number = SUPERSEDE_SETTLE_TIMEOUT_MS,
|
||||
): Promise<SupersedeResult> {
|
||||
// Validate the target belongs to THIS chat (a CAS targeting another chat's run
|
||||
// is malformed -> 400). A missing row is NOT invalid: the run may have ended
|
||||
// and been pruned; the active-run check below decides degrade vs mismatch.
|
||||
const target = await this.getRun(targetRunId, workspaceId);
|
||||
if (target && target.chatId !== chatId) return { kind: 'invalid' };
|
||||
|
||||
const active = await this.getActiveForChat(chatId, workspaceId);
|
||||
// No active run: it ended between the client's click and this POST — this is a
|
||||
// DEGRADE to a normal send, NOT a mismatch (the user's intent still holds).
|
||||
if (!active) return { kind: 'degrade' };
|
||||
// A DIFFERENT run is active than the one the client saw -> mismatch. The
|
||||
// client does not auto-retry; it surfaces the new runId.
|
||||
if (active.id !== targetRunId) {
|
||||
return { kind: 'mismatch', activeRunId: active.id };
|
||||
}
|
||||
|
||||
// The target IS active: stop it, then await its settle within W.
|
||||
await this.requestStop(targetRunId, workspaceId);
|
||||
const outcome = await this.awaitSettled(targetRunId, workspaceId, timeoutMs);
|
||||
if (!outcome) return { kind: 'timeout' };
|
||||
// Gave up (terminal write failed): apply the intended status via the
|
||||
// conditional UPDATE so the slot actually frees. If that ALSO fails, the row
|
||||
// is still stranded -> treat as a timeout (nothing persisted for the new run).
|
||||
if (outcome.terminalWriteFailed) {
|
||||
const settled = await this.settleZombie(targetRunId);
|
||||
if (!settled) return { kind: 'timeout' };
|
||||
}
|
||||
return { kind: 'ready' };
|
||||
}
|
||||
|
||||
/** #487 test/diagnostic seam: whether a give-up zombie is held for this run. */
|
||||
hasZombie(runId: string): boolean {
|
||||
return this.zombies.has(runId);
|
||||
}
|
||||
|
||||
/** #487: every zombie runId held on this replica (reconcile clause a, commit 4). */
|
||||
zombieRunIds(): string[] {
|
||||
return [...this.zombies.keys()];
|
||||
}
|
||||
|
||||
/** #487: create a one-shot deferred (resolve captured for a later single call). */
|
||||
private makeDeferred<T>(): Deferred<T> {
|
||||
let resolve!: (value: T) => void;
|
||||
const promise = new Promise<T>((r) => {
|
||||
resolve = r;
|
||||
});
|
||||
return { promise, resolve };
|
||||
}
|
||||
|
||||
/** #487: resolve a run's settle notifier EXACTLY ONCE, then drop it (bounded).
|
||||
* A subscriber that already grabbed the promise still resolves; a later one
|
||||
* falls back to the zombie map / the row (see peekSettled). */
|
||||
private resolveSettled(runId: string, outcome: RunSettleOutcome): void {
|
||||
const d = this.settledPromises.get(runId);
|
||||
if (!d) return;
|
||||
this.settledPromises.delete(runId);
|
||||
d.resolve(outcome);
|
||||
}
|
||||
|
||||
/** #487: read the persisted terminal outcome when the conditional finalize was a
|
||||
* no-op (the row was already terminal). Falls back to the intended status when
|
||||
* the read fails or the row is unexpectedly missing/non-terminal. */
|
||||
private async readTerminalOutcome(
|
||||
runId: string,
|
||||
workspaceId: string,
|
||||
fallbackStatus: RunTerminalStatus,
|
||||
fallbackError: string | null,
|
||||
): Promise<RunSettleOutcome> {
|
||||
try {
|
||||
const row = await this.runRepo.findById(runId, workspaceId);
|
||||
if (row && isRunTerminal(row.status)) {
|
||||
return {
|
||||
status: row.status as RunTerminalStatus,
|
||||
error: row.error ?? null,
|
||||
terminalWriteFailed: false,
|
||||
};
|
||||
}
|
||||
} catch {
|
||||
// Fall through to the intended status — best-effort only.
|
||||
}
|
||||
return {
|
||||
status: fallbackStatus,
|
||||
error: fallbackError,
|
||||
terminalWriteFailed: false,
|
||||
};
|
||||
}
|
||||
|
||||
/** Small async backoff between terminal-write retries (F6). Isolated so it is
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
import {
|
||||
ConflictException,
|
||||
Logger,
|
||||
ServiceUnavailableException,
|
||||
} from '@nestjs/common';
|
||||
import { AiChatService } from './ai-chat.service';
|
||||
import { RunAlreadyActiveError } from './ai-chat-run.service';
|
||||
|
||||
/**
|
||||
* Fail-fast guard for beginRun failures (#486, commit 4).
|
||||
*
|
||||
* When runHooks.begin() rejects for a reason OTHER than RunAlreadyActiveError
|
||||
* (e.g. a DB-pool blip), the turn must NOT continue untracked. The old code
|
||||
* logged and streamed anyway, leaving a run with NO run-row: in autonomous mode
|
||||
* nobody could abort it (/stop can't see it, disconnect doesn't abort it, and the
|
||||
* one-run gate would admit a SECOND run) — an unstoppable invisible run until
|
||||
* restart. The fix throws A_RUN_BEGIN_FAILED (503) BEFORE the first byte and
|
||||
* before the user row is persisted.
|
||||
*
|
||||
* We drive `stream()` directly on a prototype instance wired with only the
|
||||
* collaborators it touches before the throw, so the assertion is on the REAL
|
||||
* control flow, not a mock of it.
|
||||
*/
|
||||
describe('AiChatService beginRun failure (#486)', () => {
|
||||
function makeService(insertSpy: jest.Mock): AiChatService {
|
||||
// Bypass the (heavy) DI constructor: exercise the real stream() method on a
|
||||
// bare prototype instance with just the fields reached before the throw.
|
||||
// `any` because the private `logger` field makes a typed intersection collapse.
|
||||
const svc = Object.create(AiChatService.prototype);
|
||||
svc.aiChatRepo = {
|
||||
// Existing chat -> no insert path; chatId is kept as-is.
|
||||
findById: jest.fn().mockResolvedValue({ id: 'chat1' }),
|
||||
};
|
||||
svc.aiChatMessageRepo = { insert: insertSpy };
|
||||
svc.logger = new Logger('test');
|
||||
return svc as AiChatService;
|
||||
}
|
||||
|
||||
const baseArgs = () => {
|
||||
const write = jest.fn();
|
||||
const res = {
|
||||
raw: { write, writableEnded: false, headersSent: false },
|
||||
};
|
||||
return {
|
||||
user: { id: 'u1' } as never,
|
||||
workspace: { id: 'w1' } as never,
|
||||
sessionId: 's1',
|
||||
// openPage undefined -> resolveOpenPageContext returns null without any DB
|
||||
// call; chatId present -> the existing-chat path.
|
||||
body: { chatId: 'chat1', messages: [] } as never,
|
||||
res: res as never,
|
||||
signal: new AbortController().signal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
write,
|
||||
};
|
||||
};
|
||||
|
||||
it('throws A_RUN_BEGIN_FAILED (503) before the first byte and before persisting the user turn', async () => {
|
||||
const insertSpy = jest.fn();
|
||||
const svc = makeService(insertSpy);
|
||||
const { write, ...args } = baseArgs();
|
||||
|
||||
const runHooks = {
|
||||
begin: jest.fn().mockRejectedValue(new Error('DB pool exhausted')),
|
||||
} as never;
|
||||
|
||||
let caught: unknown;
|
||||
try {
|
||||
await svc.stream({ ...args, runHooks });
|
||||
} catch (e) {
|
||||
caught = e;
|
||||
}
|
||||
|
||||
expect(caught).toBeInstanceOf(ServiceUnavailableException);
|
||||
const http = caught as ServiceUnavailableException;
|
||||
expect(http.getStatus()).toBe(503);
|
||||
expect(http.getResponse()).toMatchObject({ code: 'A_RUN_BEGIN_FAILED' });
|
||||
|
||||
// Fail-fast: nothing was written to the socket and NO user message row was
|
||||
// persisted, so the turn left no orphan state to clean up.
|
||||
expect(write).not.toHaveBeenCalled();
|
||||
expect(insertSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('still maps a lost-the-race RunAlreadyActiveError to a 409, not A_RUN_BEGIN_FAILED', async () => {
|
||||
const insertSpy = jest.fn();
|
||||
const svc = makeService(insertSpy);
|
||||
const { write, ...args } = baseArgs();
|
||||
|
||||
const runHooks = {
|
||||
begin: jest.fn().mockRejectedValue(new RunAlreadyActiveError('chat1')),
|
||||
} as never;
|
||||
|
||||
let caught: unknown;
|
||||
try {
|
||||
await svc.stream({ ...args, runHooks });
|
||||
} catch (e) {
|
||||
caught = e;
|
||||
}
|
||||
|
||||
expect(caught).toBeInstanceOf(ConflictException);
|
||||
expect((caught as ConflictException).getResponse()).toMatchObject({
|
||||
code: 'A_RUN_ALREADY_ACTIVE',
|
||||
});
|
||||
expect(write).not.toHaveBeenCalled();
|
||||
expect(insertSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -115,7 +115,7 @@ describe('finalizeAssistant dispatch (planFinalizeAssistant + applyFinalize)', (
|
||||
|
||||
// Drive the SAME applyFinalize the service calls (no duplicated logic).
|
||||
async function dispatchFinalize(
|
||||
repo: { insert: jest.Mock; update: jest.Mock },
|
||||
repo: { insert: jest.Mock; finalizeOwner: jest.Mock },
|
||||
assistantId: string | undefined,
|
||||
flushed: AssistantFlush,
|
||||
): Promise<void> {
|
||||
@@ -135,21 +135,22 @@ describe('finalizeAssistant dispatch (planFinalizeAssistant + applyFinalize)', (
|
||||
expect(planFinalizeAssistant(undefined)).toEqual({ kind: 'insert' });
|
||||
});
|
||||
|
||||
it('(a) upfront insert succeeded -> finalize UPDATEs the row by id', async () => {
|
||||
const repo = { insert: jest.fn(), update: jest.fn() };
|
||||
it('(a) upfront insert succeeded -> finalize CONDITIONALLY updates the row by id (#487 owner-write)', async () => {
|
||||
const repo = { insert: jest.fn(), finalizeOwner: jest.fn() };
|
||||
const flushed = flushAssistant([], 'final answer', 'completed', {
|
||||
finishReason: 'stop',
|
||||
});
|
||||
await dispatchFinalize(repo, 'a1', flushed);
|
||||
expect(repo.update).toHaveBeenCalledWith('a1', workspaceId, flushed);
|
||||
// #487: the owner write is the CONDITIONAL finalizeOwner, not a raw update.
|
||||
expect(repo.finalizeOwner).toHaveBeenCalledWith('a1', workspaceId, flushed);
|
||||
expect(repo.insert).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('(b) upfront insert failed -> finalize INSERTs the terminal payload', async () => {
|
||||
const repo = { insert: jest.fn(), update: jest.fn() };
|
||||
const repo = { insert: jest.fn(), finalizeOwner: jest.fn() };
|
||||
const flushed = flushAssistant([], 'partial', 'error', { error: 'boom' });
|
||||
await dispatchFinalize(repo, undefined, flushed);
|
||||
expect(repo.update).not.toHaveBeenCalled();
|
||||
expect(repo.finalizeOwner).not.toHaveBeenCalled();
|
||||
expect(repo.insert).toHaveBeenCalledTimes(1);
|
||||
const arg = repo.insert.mock.calls[0][0];
|
||||
// The fallback insert carries the terminal content/status/metadata.
|
||||
|
||||
@@ -0,0 +1,279 @@
|
||||
import {
|
||||
BadRequestException,
|
||||
ConflictException,
|
||||
ForbiddenException,
|
||||
HttpException,
|
||||
} from '@nestjs/common';
|
||||
import { AiChatController } from './ai-chat.controller';
|
||||
import type { User, Workspace } from '@docmost/db/types/entity.types';
|
||||
|
||||
/**
|
||||
* #487 commit 3 — the single concurrency GATE (both modes) + the server supersede
|
||||
* CAS, at the controller boundary. The gate + CAS run BEFORE res.hijack(), so a
|
||||
* rejected concurrent start / a CAS branch returns clean JSON (an HttpException
|
||||
* the controller's post-hijack catch re-serializes). These assert the OBSERVABLE
|
||||
* HTTP contract against the real controller + a stubbed run service.
|
||||
*/
|
||||
describe('#487 AiChatController.stream — gate + supersede', () => {
|
||||
const user = { id: 'u1' } as User;
|
||||
|
||||
function wsWith(autonomousRuns: boolean): Workspace {
|
||||
return {
|
||||
id: 'ws1',
|
||||
settings: { ai: { chat: true, autonomousRuns } },
|
||||
} as unknown as Workspace;
|
||||
}
|
||||
|
||||
function makeReqRes(body: Record<string, unknown>) {
|
||||
const req = {
|
||||
raw: { sessionId: 'sess', once: jest.fn(), destroyed: false },
|
||||
body,
|
||||
};
|
||||
const res = {
|
||||
raw: {
|
||||
writableEnded: false,
|
||||
headersSent: false,
|
||||
on: jest.fn(),
|
||||
once: jest.fn(),
|
||||
setHeader: jest.fn(),
|
||||
end: jest.fn(),
|
||||
statusCode: 200,
|
||||
flushHeaders: jest.fn(),
|
||||
},
|
||||
hijack: jest.fn(),
|
||||
status: jest.fn().mockReturnThis(),
|
||||
send: jest.fn(),
|
||||
};
|
||||
return { req, res };
|
||||
}
|
||||
|
||||
function makeController(
|
||||
runServiceOverrides: Record<string, jest.Mock>,
|
||||
// The chat assertOwnedChat resolves. Default: a chat OWNED by `user` (u1), so
|
||||
// the ownership gate is transparent to the gate/CAS assertions below. Pass a
|
||||
// foreign-owner (or undefined) chat to exercise the #487 owner rejection.
|
||||
chat: { creatorId: string } | undefined = { creatorId: 'u1' },
|
||||
) {
|
||||
const aiChatService = {
|
||||
resolveRoleForRequest: jest.fn().mockResolvedValue(null),
|
||||
getChatModel: jest.fn().mockResolvedValue({}),
|
||||
stream: jest.fn().mockResolvedValue(undefined),
|
||||
};
|
||||
const aiChatRunService = {
|
||||
getActiveForChat: jest.fn().mockResolvedValue(undefined),
|
||||
supersede: jest.fn(),
|
||||
beginRun: jest.fn().mockResolvedValue({
|
||||
runId: 'run-new',
|
||||
signal: new AbortController().signal,
|
||||
}),
|
||||
linkAssistantMessage: jest.fn(),
|
||||
recordStep: jest.fn(),
|
||||
finalizeRun: jest.fn(),
|
||||
requestStop: jest.fn(),
|
||||
...runServiceOverrides,
|
||||
};
|
||||
const aiChatRepo = { findById: jest.fn().mockResolvedValue(chat) };
|
||||
const controller = new AiChatController(
|
||||
aiChatService as never,
|
||||
aiChatRunService as never,
|
||||
aiChatRepo as never, // aiChatRepo
|
||||
{} as never, // aiChatMessageRepo
|
||||
{} as never, // aiTranscription
|
||||
{} as never, // pageRepo
|
||||
);
|
||||
return { controller, aiChatService, aiChatRunService, aiChatRepo };
|
||||
}
|
||||
|
||||
const codeOf = (err: unknown) =>
|
||||
(((err as HttpException).getResponse() as Record<string, unknown>) ?? {})
|
||||
.code;
|
||||
|
||||
describe('single concurrency gate — BOTH modes reject the second tab with 409', () => {
|
||||
for (const autonomousRuns of [true, false]) {
|
||||
it(`rejects a concurrent start with 409 A_RUN_ALREADY_ACTIVE (autonomousRuns=${autonomousRuns})`, async () => {
|
||||
const { controller, aiChatRunService } = makeController({
|
||||
getActiveForChat: jest
|
||||
.fn()
|
||||
.mockResolvedValue({ id: 'run-live', chatId: 'c1' }),
|
||||
});
|
||||
const { req, res } = makeReqRes({ chatId: 'c1' });
|
||||
let thrown: unknown;
|
||||
try {
|
||||
await controller.stream(
|
||||
req as never,
|
||||
res as never,
|
||||
user,
|
||||
wsWith(autonomousRuns),
|
||||
);
|
||||
} catch (e) {
|
||||
thrown = e;
|
||||
}
|
||||
expect(thrown).toBeInstanceOf(ConflictException);
|
||||
expect((thrown as HttpException).getStatus()).toBe(409);
|
||||
expect(codeOf(thrown)).toBe('A_RUN_ALREADY_ACTIVE');
|
||||
// Rejected BEFORE committing to the stream (no hijack, no service.stream).
|
||||
expect(res.hijack).not.toHaveBeenCalled();
|
||||
expect(aiChatRunService.getActiveForChat).toHaveBeenCalledWith(
|
||||
'c1',
|
||||
'ws1',
|
||||
);
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
// #487 [security, F1]: stream() MUST owner-gate an existing chat exactly like its
|
||||
// six sibling endpoints, BEFORE the supersede CAS. Otherwise a same-workspace
|
||||
// non-owner could POST a supersede against another user's chat and (a) harvest
|
||||
// that user's active runId from the 409 SUPERSEDE_TARGET_MISMATCH body, then (b)
|
||||
// requestStop the foreign run. The gate must reject FIRST — no run lookup, no
|
||||
// supersede, no stop, no runId leak.
|
||||
describe('cross-user ownership gate (F1)', () => {
|
||||
it('a non-owner streaming against someone else\'s chat is rejected (403) with NO runId leak and NO foreign requestStop', async () => {
|
||||
// A live run exists on the victim's chat. Without the gate the supersede CAS
|
||||
// would run and (faithful to the run service) return a MISMATCH carrying the
|
||||
// victim's runId — the exact leak. With the gate it must never be reached.
|
||||
const getActiveForChat = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ id: 'run-victim', chatId: 'c-other' });
|
||||
const supersede = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ kind: 'mismatch', activeRunId: 'run-victim' });
|
||||
const requestStop = jest.fn();
|
||||
const { controller, aiChatService } = makeController(
|
||||
{ getActiveForChat, supersede, requestStop },
|
||||
{ creatorId: 'someone-else' }, // the chat is NOT owned by u1
|
||||
);
|
||||
const { req, res } = makeReqRes({
|
||||
chatId: 'c-other',
|
||||
supersede: { runId: 'guessed-uuid' },
|
||||
});
|
||||
let thrown: unknown;
|
||||
try {
|
||||
await controller.stream(req as never, res as never, user, wsWith(true));
|
||||
} catch (e) {
|
||||
thrown = e;
|
||||
}
|
||||
// Rejected by the ownership gate (403), the SAME shape the neighbors use.
|
||||
expect(thrown).toBeInstanceOf(ForbiddenException);
|
||||
expect((thrown as HttpException).getStatus()).toBe(403);
|
||||
// Crucially NOT a 409 that would carry activeRunId — no runId is leaked.
|
||||
const payload = JSON.stringify(
|
||||
(thrown as HttpException).getResponse() ?? {},
|
||||
);
|
||||
expect(payload).not.toContain('run-victim');
|
||||
expect(codeOf(thrown)).not.toBe('SUPERSEDE_TARGET_MISMATCH');
|
||||
// The gate short-circuits BEFORE any run machinery runs.
|
||||
expect(getActiveForChat).not.toHaveBeenCalled();
|
||||
expect(supersede).not.toHaveBeenCalled();
|
||||
expect(requestStop).not.toHaveBeenCalled();
|
||||
expect(aiChatService.stream).not.toHaveBeenCalled();
|
||||
expect(res.hijack).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
it('supersede MISMATCH -> 409 SUPERSEDE_TARGET_MISMATCH carrying the current runId', async () => {
|
||||
const { controller } = makeController({
|
||||
supersede: jest
|
||||
.fn()
|
||||
.mockResolvedValue({ kind: 'mismatch', activeRunId: 'run-other' }),
|
||||
});
|
||||
const { req, res } = makeReqRes({
|
||||
chatId: 'c1',
|
||||
supersede: { runId: 'run-x' },
|
||||
});
|
||||
let thrown: unknown;
|
||||
try {
|
||||
await controller.stream(req as never, res as never, user, wsWith(true));
|
||||
} catch (e) {
|
||||
thrown = e;
|
||||
}
|
||||
expect(thrown).toBeInstanceOf(ConflictException);
|
||||
expect(codeOf(thrown)).toBe('SUPERSEDE_TARGET_MISMATCH');
|
||||
expect(
|
||||
((thrown as HttpException).getResponse() as Record<string, unknown>)
|
||||
.activeRunId,
|
||||
).toBe('run-other');
|
||||
expect(res.hijack).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('supersede TIMEOUT -> 409 SUPERSEDE_TIMEOUT, nothing streamed', async () => {
|
||||
const { controller } = makeController({
|
||||
supersede: jest.fn().mockResolvedValue({ kind: 'timeout' }),
|
||||
});
|
||||
const { req, res } = makeReqRes({
|
||||
chatId: 'c1',
|
||||
supersede: { runId: 'run-x' },
|
||||
});
|
||||
let thrown: unknown;
|
||||
try {
|
||||
await controller.stream(req as never, res as never, user, wsWith(false));
|
||||
} catch (e) {
|
||||
thrown = e;
|
||||
}
|
||||
expect(thrown).toBeInstanceOf(ConflictException);
|
||||
expect(codeOf(thrown)).toBe('SUPERSEDE_TIMEOUT');
|
||||
expect(res.hijack).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('supersede INVALID (target on another chat) -> 400 SUPERSEDE_INVALID', async () => {
|
||||
const { controller } = makeController({
|
||||
supersede: jest.fn().mockResolvedValue({ kind: 'invalid' }),
|
||||
});
|
||||
const { req, res } = makeReqRes({
|
||||
chatId: 'c1',
|
||||
supersede: { runId: 'run-x' },
|
||||
});
|
||||
let thrown: unknown;
|
||||
try {
|
||||
await controller.stream(req as never, res as never, user, wsWith(true));
|
||||
} catch (e) {
|
||||
thrown = e;
|
||||
}
|
||||
expect(thrown).toBeInstanceOf(BadRequestException);
|
||||
expect(codeOf(thrown)).toBe('SUPERSEDE_INVALID');
|
||||
});
|
||||
|
||||
it('supersede without chatId -> 400 SUPERSEDE_INVALID', async () => {
|
||||
const { controller, aiChatRunService } = makeController({});
|
||||
const { req, res } = makeReqRes({ supersede: { runId: 'run-x' } });
|
||||
let thrown: unknown;
|
||||
try {
|
||||
await controller.stream(req as never, res as never, user, wsWith(true));
|
||||
} catch (e) {
|
||||
thrown = e;
|
||||
}
|
||||
expect(thrown).toBeInstanceOf(BadRequestException);
|
||||
expect(codeOf(thrown)).toBe('SUPERSEDE_INVALID');
|
||||
expect(aiChatRunService.supersede).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('supersede READY -> proceeds to stream with superseded=true', async () => {
|
||||
const { controller, aiChatService } = makeController({
|
||||
supersede: jest.fn().mockResolvedValue({ kind: 'ready' }),
|
||||
getActiveForChat: jest.fn().mockResolvedValue(undefined), // slot free after CAS
|
||||
});
|
||||
const { req, res } = makeReqRes({
|
||||
chatId: 'c1',
|
||||
supersede: { runId: 'run-x' },
|
||||
});
|
||||
await controller.stream(req as never, res as never, user, wsWith(true));
|
||||
expect(res.hijack).toHaveBeenCalled();
|
||||
expect(aiChatService.stream).toHaveBeenCalledTimes(1);
|
||||
expect(aiChatService.stream.mock.calls[0][0].superseded).toBe(true);
|
||||
// The run hooks are always present now (both modes).
|
||||
expect(aiChatService.stream.mock.calls[0][0].runHooks).toBeDefined();
|
||||
});
|
||||
|
||||
it('supersede DEGRADE -> proceeds to a normal send (superseded=false)', async () => {
|
||||
const { controller, aiChatService } = makeController({
|
||||
supersede: jest.fn().mockResolvedValue({ kind: 'degrade' }),
|
||||
});
|
||||
const { req, res } = makeReqRes({
|
||||
chatId: 'c1',
|
||||
supersede: { runId: 'run-x' },
|
||||
});
|
||||
await controller.stream(req as never, res as never, user, wsWith(false));
|
||||
expect(aiChatService.stream).toHaveBeenCalledTimes(1);
|
||||
expect(aiChatService.stream.mock.calls[0][0].superseded).toBe(false);
|
||||
});
|
||||
});
|
||||
@@ -418,6 +418,19 @@ export class AiChatController {
|
||||
|
||||
const body = (req.body ?? {}) as AiChatStreamBody;
|
||||
|
||||
// #487 [security]: gate cross-user access to an EXISTING chat BEFORE anything
|
||||
// reads its runs. Every sibling endpoint (getRun/stop/history/rename/delete/
|
||||
// attachRunStream) owner-checks the chat via assertOwnedChat; stream() must too.
|
||||
// Without this a same-workspace member who is NOT the chat owner could POST a
|
||||
// supersede against another user's chat and (a) harvest that user's active runId
|
||||
// out of the 409 SUPERSEDE_TARGET_MISMATCH body, then (b) requestStop the foreign
|
||||
// run. Gate on the chatId the client sent, when present — a brand-new chat (no
|
||||
// chatId) has no prior owner to check. Mirrors /stop's owner check (403 as the
|
||||
// neighbors do), and runs pre-hijack so it returns clean JSON.
|
||||
if (body.chatId) {
|
||||
await this.assertOwnedChat(body.chatId, user, workspace);
|
||||
}
|
||||
|
||||
// Resolve the agent role for this turn BEFORE hijack: existing chats read it
|
||||
// from ai_chats.role_id (authoritative), a new chat from body.roleId. The
|
||||
// role drives both the persona and the optional model override below.
|
||||
@@ -432,12 +445,66 @@ 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) {
|
||||
// #487: server-side supersede CAS ("interrupt and send now"). When the client
|
||||
// asks to replace a live run, atomically STOP it and wait for it to settle
|
||||
// before this turn claims the slot. Runs BEFORE hijack so every branch returns
|
||||
// clean JSON (the client keeps the composer text on a 409). See
|
||||
// AiChatRunService.supersede for the branch semantics.
|
||||
let superseded = false;
|
||||
const supersedeRunId = body.supersede?.runId;
|
||||
if (supersedeRunId) {
|
||||
if (!body.chatId) {
|
||||
throw new BadRequestException({
|
||||
message: 'supersede requires chatId',
|
||||
code: 'SUPERSEDE_INVALID',
|
||||
});
|
||||
}
|
||||
const result = await this.aiChatRunService.supersede(
|
||||
body.chatId,
|
||||
supersedeRunId,
|
||||
workspace.id,
|
||||
);
|
||||
switch (result.kind) {
|
||||
case 'invalid':
|
||||
throw new BadRequestException({
|
||||
message: 'The run to supersede does not belong to this chat',
|
||||
code: 'SUPERSEDE_INVALID',
|
||||
});
|
||||
case 'mismatch':
|
||||
// A DIFFERENT run is active than the one the client targeted. Surface
|
||||
// the CURRENT runId; the client does NOT auto-retry (a stale CAS).
|
||||
throw new ConflictException({
|
||||
message: 'A different agent run is now active on this chat',
|
||||
code: 'SUPERSEDE_TARGET_MISMATCH',
|
||||
activeRunId: result.activeRunId,
|
||||
});
|
||||
case 'timeout':
|
||||
// The target did not settle within W — nothing was persisted, the
|
||||
// composer keeps the text. NOT a rollback: the stop is already issued.
|
||||
throw new ConflictException({
|
||||
message:
|
||||
'The previous run did not stop in time; nothing was sent — please try again',
|
||||
code: 'SUPERSEDE_TIMEOUT',
|
||||
});
|
||||
case 'ready':
|
||||
// The target stopped and settled: the slot is free. Prompt the new run
|
||||
// that the old run's last operations may still be applying.
|
||||
superseded = true;
|
||||
break;
|
||||
case 'degrade':
|
||||
// The run already ended between click and POST — send normally.
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// #487: one active run per chat — ENFORCED IN BOTH MODES now (legacy mode used
|
||||
// to have NO gate, so two tabs streamed two parallel turns on one chat, which
|
||||
// interleaved history and crashed convertToModelMessages). Reject a concurrent
|
||||
// start with a clean pre-hijack 409 (double-submit / second-tab). A brand-new
|
||||
// chat (no chatId) cannot have a prior run, and the DB partial unique index in
|
||||
// beginRun is the authoritative backstop for any race that slips past here
|
||||
// (including a slot stolen between a supersede release and beginRun).
|
||||
if (body.chatId) {
|
||||
const active = await this.aiChatRunService.getActiveForChat(
|
||||
body.chatId,
|
||||
workspace.id,
|
||||
@@ -446,107 +513,94 @@ export class AiChatController {
|
||||
throw new ConflictException({
|
||||
message: 'An agent run is already in progress for this chat',
|
||||
code: 'A_RUN_ALREADY_ACTIVE',
|
||||
activeRunId: active.id,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// 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: async (chatId) => {
|
||||
const handle = await this.aiChatRunService.beginRun({
|
||||
chatId,
|
||||
workspaceId: workspace.id,
|
||||
userId: user.id,
|
||||
trigger: 'user',
|
||||
});
|
||||
// #184 phase 1.5: register the run-stream entry at BEGIN (before any
|
||||
// frame) so a tab that attaches in the begin->seed window finds an
|
||||
// entry to wait on. Gated on AI_CHAT_RESUMABLE_STREAM: with the flag
|
||||
// off nothing is registered and attach always 204s.
|
||||
if (
|
||||
handle?.runId &&
|
||||
this.environment?.isAiChatResumableStreamEnabled?.()
|
||||
) {
|
||||
this.streamRegistry?.open(chatId, handle.runId);
|
||||
}
|
||||
return handle;
|
||||
},
|
||||
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,
|
||||
),
|
||||
// #487: the turn is ALWAYS a first-class RUN now (both modes). The mode
|
||||
// difference is only the abort semantics on a browser disconnect (onClose
|
||||
// below). currentRunId is captured at begin so a legacy disconnect can stop
|
||||
// the run through its stop lever.
|
||||
let currentRunId: string | undefined;
|
||||
const runHooks: AiChatRunHooks = {
|
||||
begin: async (chatId) => {
|
||||
const handle = await this.aiChatRunService.beginRun({
|
||||
chatId,
|
||||
workspaceId: workspace.id,
|
||||
userId: user.id,
|
||||
trigger: 'user',
|
||||
});
|
||||
currentRunId = handle?.runId;
|
||||
// #184 phase 1.5: register the run-stream entry at BEGIN (before any
|
||||
// frame) so a tab that attaches in the begin->seed window finds an entry
|
||||
// to wait on. Gated on AI_CHAT_RESUMABLE_STREAM.
|
||||
if (
|
||||
handle?.runId &&
|
||||
this.environment?.isAiChatResumableStreamEnabled?.()
|
||||
) {
|
||||
this.streamRegistry?.open(chatId, handle.runId);
|
||||
}
|
||||
: undefined;
|
||||
return handle;
|
||||
},
|
||||
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),
|
||||
};
|
||||
|
||||
// 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;
|
||||
// we also drop it on response `finish` so it never lingers after the stream
|
||||
// completes normally (the AI SDK pipes the response fire-and-forget, so we
|
||||
// cannot simply remove it once `stream()` returns).
|
||||
// Handle a client disconnect. `close` also fires on normal completion, so only
|
||||
// act when the response has not finished writing (a genuine disconnect). `once`
|
||||
// fires at most once and self-removes; we also drop it on response `finish`.
|
||||
// DIAGNOSTIC (Safari stream-drop investigation) — temporary: wall-clock at
|
||||
// which a Safari disconnect is observed, measured from request receipt.
|
||||
const reqStartedAt = Date.now();
|
||||
const controller = new AbortController();
|
||||
const onClose = (): void => {
|
||||
// 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.
|
||||
if (!res.raw.writableEnded) {
|
||||
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.
|
||||
// #184: a DETACHED run — a disconnect must NOT stop 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 {
|
||||
// #487: legacy — a disconnect ENDS the turn, but the turn is now a RUN,
|
||||
// so stop it through the run's stop lever (requestStop). streamText no
|
||||
// longer consumes the socket signal (effectiveSignal is the run signal),
|
||||
// so aborting `controller` would do nothing; requestStop aborts the run.
|
||||
this.logger.warn(
|
||||
`AI chat stream: client disconnected before completion; aborting turn ` +
|
||||
`(elapsed=${Date.now() - reqStartedAt}ms since request received)`,
|
||||
`AI chat stream: client disconnected before completion; stopping the ` +
|
||||
`run (elapsed=${Date.now() - reqStartedAt}ms since request received)`,
|
||||
);
|
||||
controller.abort();
|
||||
if (currentRunId) {
|
||||
void this.aiChatRunService.requestStop(currentRunId, workspace.id);
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
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)
|
||||
}`,
|
||||
);
|
||||
});
|
||||
}
|
||||
// #184/#487: the run/pipe can outlive the socket in BOTH modes now (autonomous
|
||||
// keeps going; legacy keeps going until requestStop's abort unwinds the turn).
|
||||
// The SDK's pipe may then write to a dropped socket and emit an 'error' on the
|
||||
// raw response — swallow it so it never surfaces as an unhandled error event.
|
||||
res.raw.on('error', (err) => {
|
||||
this.logger.debug(
|
||||
`AI chat 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.
|
||||
@@ -562,8 +616,10 @@ export class AiChatController {
|
||||
signal: controller.signal,
|
||||
model,
|
||||
role,
|
||||
// #184: present only when the flag is on; wraps the turn in a durable run.
|
||||
// #487: the turn is always run-wrapped now (both modes).
|
||||
runHooks,
|
||||
// #487: warn the new run that a superseded run's last ops may still apply.
|
||||
superseded,
|
||||
});
|
||||
} catch (err) {
|
||||
// Any failure AFTER hijack can no longer go through Nest's exception
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
// #489 — client-parts validation + resilient history conversion.
|
||||
//
|
||||
// These unit tests exercise the two exported helpers against the REAL
|
||||
// `convertToModelMessages` from `ai` (NOT a mock): a genuinely malformed part
|
||||
// (a `null` element inside a parts array) makes the real converter throw
|
||||
// ("Cannot read properties of null"), which is the actual production
|
||||
// "bricked chat" mechanism this fix defends against. Asserting against the real
|
||||
// converter (rather than a mock-shaped error) is the whole point — a mock would
|
||||
// hide a version change in the converter's throw behaviour.
|
||||
import { convertToModelMessages, type UIMessage } from 'ai';
|
||||
import {
|
||||
sanitizeUserParts,
|
||||
convertHistoryResilient,
|
||||
TOOL_CONTEXT_OMITTED_MARKER,
|
||||
} from './ai-chat.service';
|
||||
|
||||
type Row = Omit<UIMessage, 'id'> & { id: string };
|
||||
|
||||
describe('sanitizeUserParts (#489, branch: validation on receipt)', () => {
|
||||
it('keeps whitelisted text parts unchanged', () => {
|
||||
const drops: string[] = [];
|
||||
const out = sanitizeUserParts(
|
||||
[
|
||||
{ type: 'text', text: 'a' },
|
||||
{ type: 'text', text: 'b' },
|
||||
] as UIMessage['parts'],
|
||||
(t) => drops.push(t),
|
||||
);
|
||||
expect(out).toEqual([
|
||||
{ type: 'text', text: 'a' },
|
||||
{ type: 'text', text: 'b' },
|
||||
]);
|
||||
expect(drops).toEqual([]);
|
||||
});
|
||||
|
||||
it('drops a non-text part (a tool-part in input-available) and reports its type', () => {
|
||||
const drops: string[] = [];
|
||||
const out = sanitizeUserParts(
|
||||
[
|
||||
{ type: 'text', text: 'hi' },
|
||||
{
|
||||
type: 'tool-getPage',
|
||||
toolCallId: 't1',
|
||||
state: 'input-available',
|
||||
input: { pageId: 'p' },
|
||||
},
|
||||
] as unknown as UIMessage['parts'],
|
||||
(t) => drops.push(t),
|
||||
);
|
||||
expect(out).toEqual([{ type: 'text', text: 'hi' }]);
|
||||
expect(drops).toEqual(['tool-getPage']);
|
||||
});
|
||||
|
||||
it('drops a null part (the shape that would poison convertToModelMessages)', () => {
|
||||
const drops: string[] = [];
|
||||
const out = sanitizeUserParts(
|
||||
[{ type: 'text', text: 'hi' }, null] as unknown as UIMessage['parts'],
|
||||
(t) => drops.push(t),
|
||||
);
|
||||
expect(out).toEqual([{ type: 'text', text: 'hi' }]);
|
||||
expect(drops).toEqual(['(unknown)']);
|
||||
});
|
||||
|
||||
it('returns undefined when nothing survives (so a null metadata is persisted)', () => {
|
||||
const out = sanitizeUserParts(
|
||||
[
|
||||
{ type: 'tool-x', toolCallId: 't', state: 'input-available' },
|
||||
] as unknown as UIMessage['parts'],
|
||||
() => undefined,
|
||||
);
|
||||
expect(out).toBeUndefined();
|
||||
});
|
||||
|
||||
it('returns undefined for a non-array input', () => {
|
||||
expect(
|
||||
sanitizeUserParts(undefined as unknown as UIMessage['parts'], () => undefined),
|
||||
).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('convertHistoryResilient (#489, branches: happy + per-row degradation)', () => {
|
||||
it('happy path: healthy history converts identically to convertToModelMessages, no degrade', async () => {
|
||||
const history: Row[] = [
|
||||
{ id: 'u1', role: 'user', parts: [{ type: 'text', text: 'hi' }] },
|
||||
{ id: 'a1', role: 'assistant', parts: [{ type: 'text', text: 'hello' }] },
|
||||
];
|
||||
const degrades: number[] = [];
|
||||
const out = await convertHistoryResilient(history, (i) => degrades.push(i));
|
||||
const expected = await convertToModelMessages(history as UIMessage[]);
|
||||
expect(out).toEqual(expected);
|
||||
expect(degrades).toEqual([]);
|
||||
});
|
||||
|
||||
it('REAL poison: a null part throws in the batch converter but is isolated and degraded to a marker', async () => {
|
||||
// Sanity: the real converter genuinely throws on this shape.
|
||||
const poisoned: Row = {
|
||||
id: 'a1',
|
||||
role: 'assistant',
|
||||
parts: [
|
||||
{ type: 'text', text: 'earlier answer' },
|
||||
null,
|
||||
] as unknown as UIMessage['parts'],
|
||||
};
|
||||
await expect(
|
||||
convertToModelMessages([poisoned as UIMessage]),
|
||||
).rejects.toThrow();
|
||||
|
||||
const history: Row[] = [
|
||||
{ id: 'u1', role: 'user', parts: [{ type: 'text', text: 'first' }] },
|
||||
poisoned,
|
||||
{ id: 'u2', role: 'user', parts: [{ type: 'text', text: 'second' }] },
|
||||
];
|
||||
const degrades: number[] = [];
|
||||
const out = await convertHistoryResilient(history, (i) => degrades.push(i));
|
||||
|
||||
// Only the poisoned row (index 1) is degraded.
|
||||
expect(degrades).toEqual([1]);
|
||||
// Healthy rows survive verbatim.
|
||||
const flat = JSON.stringify(out);
|
||||
expect(flat).toContain('first');
|
||||
expect(flat).toContain('second');
|
||||
// The degraded row carries its readable text AND the truncation marker so the
|
||||
// model sees that tool context was omitted (never a silent loss).
|
||||
expect(flat).toContain('earlier answer');
|
||||
expect(flat).toContain(TOOL_CONTEXT_OMITTED_MARKER);
|
||||
// The whole batch converted (3 model messages, none dropped).
|
||||
expect(out).toHaveLength(3);
|
||||
});
|
||||
|
||||
it('a fully-poisoned row (no readable text) still degrades to just the marker', async () => {
|
||||
const history: Row[] = [
|
||||
{
|
||||
id: 'a1',
|
||||
role: 'assistant',
|
||||
parts: [null] as unknown as UIMessage['parts'],
|
||||
},
|
||||
];
|
||||
const out = await convertHistoryResilient(history, () => undefined);
|
||||
expect(out).toHaveLength(1);
|
||||
expect(JSON.stringify(out)).toContain(TOOL_CONTEXT_OMITTED_MARKER);
|
||||
});
|
||||
});
|
||||
@@ -101,6 +101,22 @@ const INTERRUPT_NOTE =
|
||||
'assume your previous response was complete, and do not silently restart the ' +
|
||||
'partial work — build on it or follow the new instruction.';
|
||||
|
||||
/**
|
||||
* #487: injected on a turn started by SUPERSEDING a previous run (the user hit
|
||||
* "interrupt and send now" while a run was live). The previous run was Stopped,
|
||||
* but there is NO side-effect quiescence — a write it had already committed, or
|
||||
* one committing at the moment of Stop, may land with a small delay AFTER this new
|
||||
* run starts. So the model is told its picture of the page/state may be a beat
|
||||
* stale and to re-read before assuming an edit did or did not apply.
|
||||
*/
|
||||
const SUPERSEDE_NOTE =
|
||||
'NOTE: A previous agent run in this conversation was just interrupted so this ' +
|
||||
'new turn could start. That run was stopped, but any operation it had already ' +
|
||||
'begun (e.g. a page edit) may still be applied with a short delay. Do not ' +
|
||||
'assume the document/state is exactly as the interrupted run left it — if you ' +
|
||||
'need to rely on the current content, RE-READ it with the page tools before ' +
|
||||
'acting rather than trusting a cached view.';
|
||||
|
||||
/**
|
||||
* Injected on a turn where the open page was hand-edited by the user (or anyone
|
||||
* else) AFTER the agent's previous response ended (#274). The server takes a
|
||||
@@ -203,6 +219,14 @@ export interface BuildSystemPromptInput {
|
||||
* (partial) answer was cut off by the user's new message.
|
||||
*/
|
||||
interrupted?: boolean;
|
||||
/**
|
||||
* #487: true when THIS turn was started by superseding a still-live previous run
|
||||
* ("interrupt and send now"). Adds SUPERSEDE_NOTE so the model knows the previous
|
||||
* run's last operations may still be applying and to re-read state it depends on.
|
||||
* Distinct from `interrupted` (which is about a PARTIAL prior answer in history);
|
||||
* both can be set together. Self-clears — set only for the superseding turn.
|
||||
*/
|
||||
superseded?: boolean;
|
||||
/**
|
||||
* Set only when the open page was edited by the user AFTER the agent's previous
|
||||
* turn ended (#274), confirmed server-side by diffing the current page against
|
||||
@@ -311,6 +335,7 @@ export function buildSystemPrompt({
|
||||
openedPage,
|
||||
mcpInstructions,
|
||||
interrupted,
|
||||
superseded,
|
||||
pageChanged,
|
||||
deferredToolsEnabled,
|
||||
toolCatalog,
|
||||
@@ -360,6 +385,13 @@ export function buildSystemPrompt({
|
||||
context += `\n${INTERRUPT_NOTE}`;
|
||||
}
|
||||
|
||||
// Supersede note (#487): present only for a turn that stopped and replaced a
|
||||
// still-live previous run — warns the model the previous run's last operations
|
||||
// may still be applying (no side-effect quiescence).
|
||||
if (superseded) {
|
||||
context += `\n${SUPERSEDE_NOTE}`;
|
||||
}
|
||||
|
||||
// Per-turn page-change note (#274). Added to the context section (inside the
|
||||
// safety sandwich), present only when the server detected that the open page
|
||||
// was edited by the user since the agent's last turn ended. The diff content is
|
||||
|
||||
@@ -89,11 +89,22 @@ describe('AiChatService.stream run-lifecycle safety net (#184)', () => {
|
||||
const runRepo = {
|
||||
insert: jest.fn().mockResolvedValue({ id: 'run-1', status: 'running' }),
|
||||
update: jest.fn().mockResolvedValue({ id: 'run-1' }),
|
||||
// #487: the terminal settle now goes through the CONDITIONAL write.
|
||||
finalizeIfActive: jest
|
||||
.fn()
|
||||
.mockResolvedValue({ id: 'run-1', status: 'failed' }),
|
||||
findById: jest.fn().mockResolvedValue(undefined),
|
||||
};
|
||||
const runService = new AiChatRunService(runRepo as never, { isCloud: () => false } as never);
|
||||
|
||||
// The user-message insert (the first bare await after beginRun) throws.
|
||||
// The user-message insert throws. #489 runs the history load + convert BEFORE
|
||||
// the insert (convert-before-insert, so a retry cannot duplicate the user row),
|
||||
// so `findAllByChat` (a real repo method) is now called first — stub it to an
|
||||
// empty history so the flow reaches the insert. Both awaits are AFTER beginRun,
|
||||
// so the "exception after beginRun -> settled to error" invariant is unchanged;
|
||||
// the throw point simply moved from insert to a later insert after a no-op load.
|
||||
const aiChatMessageRepo = {
|
||||
findAllByChat: jest.fn().mockResolvedValue([]),
|
||||
insert: jest.fn().mockRejectedValue(new Error('insert boom')),
|
||||
};
|
||||
const aiChatRepo = {
|
||||
@@ -148,9 +159,10 @@ describe('AiChatService.stream run-lifecycle safety net (#184)', () => {
|
||||
|
||||
// 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(
|
||||
// ...then settled to a terminal FAILED status by the safety net (via the
|
||||
// #487 conditional write)...
|
||||
expect(runRepo.finalizeIfActive).toHaveBeenCalledTimes(1);
|
||||
expect(runRepo.finalizeIfActive).toHaveBeenCalledWith(
|
||||
'run-1',
|
||||
'ws1',
|
||||
expect.objectContaining({ status: 'failed' }),
|
||||
|
||||
@@ -1,4 +1,8 @@
|
||||
import { ConflictException, Logger } from '@nestjs/common';
|
||||
import {
|
||||
ConflictException,
|
||||
Logger,
|
||||
ServiceUnavailableException,
|
||||
} 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
|
||||
@@ -151,6 +155,8 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
|
||||
insert: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
findAllByChat: jest.fn(async () => []),
|
||||
update: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
finalizeOwner: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
findStreamingWithTerminalRun: jest.fn(async () => []),
|
||||
};
|
||||
const aiSettings = { resolve: jest.fn(async () => ({})) };
|
||||
const tools = { forUser: jest.fn(async () => ({})) };
|
||||
@@ -175,7 +181,7 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
|
||||
{} as never, // pageAccess
|
||||
{ isAiChatDeferredToolsEnabled: () => false, isAiChatFinalStepLockdownEnabled: () => false } as never, // environment
|
||||
);
|
||||
return { svc };
|
||||
return { svc, aiChatMessageRepo };
|
||||
}
|
||||
|
||||
const body = {
|
||||
@@ -281,7 +287,7 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
|
||||
// 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();
|
||||
const { svc, aiChatMessageRepo } = makeService();
|
||||
let capturedOpts: any;
|
||||
streamTextMock.mockImplementation((opts: any) => {
|
||||
capturedOpts = opts;
|
||||
@@ -308,7 +314,7 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
|
||||
runHooks: runHooks as never,
|
||||
});
|
||||
expect(capturedOpts).toBeDefined();
|
||||
return { capturedOpts, runHooks };
|
||||
return { capturedOpts, runHooks, aiChatMessageRepo };
|
||||
}
|
||||
|
||||
it('F9: onStepFinish bumps the run step count, onFinish settles the run "completed" (the dominant autonomous-run path)', async () => {
|
||||
@@ -328,7 +334,13 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
|
||||
usage: {},
|
||||
steps: [],
|
||||
});
|
||||
expect(runHooks.onSettled).toHaveBeenCalledWith('run-1', 'completed');
|
||||
// #487: onFinish passes the (undefined) error slot so a message-finalize
|
||||
// failure could error-mark the run; on the success path it is undefined.
|
||||
expect(runHooks.onSettled).toHaveBeenCalledWith(
|
||||
'run-1',
|
||||
'completed',
|
||||
undefined,
|
||||
);
|
||||
});
|
||||
|
||||
it('F9: onAbort settles the run "aborted"', async () => {
|
||||
@@ -357,25 +369,70 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
|
||||
expect.stringContaining('provider exploded'),
|
||||
);
|
||||
});
|
||||
|
||||
// #490 reactive branch: a provider CONTEXT-OVERFLOW 400 in onError is classified,
|
||||
// records a distinguishable cause, and stamps metadata.replayOverflow so the NEXT
|
||||
// turn's budgeter trims aggressively (the recovery that un-bricks the chat).
|
||||
it('#490: a context-overflow 400 stamps replayOverflow on the finalized row', async () => {
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'error')
|
||||
.mockImplementation(() => undefined as never);
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'warn')
|
||||
.mockImplementation(() => undefined as never);
|
||||
const { capturedOpts, aiChatMessageRepo } = await captureStreamCallbacks();
|
||||
|
||||
const overflow = Object.assign(new Error('too large'), {
|
||||
statusCode: 400,
|
||||
message:
|
||||
"This model's maximum context length is 128000 tokens. However, your messages resulted in 214000 tokens. Please reduce the length.",
|
||||
});
|
||||
await capturedOpts.onError({ error: overflow });
|
||||
|
||||
// The seed row exists (finalizeOwner is the owner-write path).
|
||||
expect(aiChatMessageRepo.finalizeOwner).toHaveBeenCalled();
|
||||
const calls = aiChatMessageRepo.finalizeOwner.mock.calls as any[][];
|
||||
const patch = calls[calls.length - 1][2] as {
|
||||
status: string;
|
||||
metadata: Record<string, unknown>;
|
||||
};
|
||||
expect(patch.status).toBe('error');
|
||||
expect(patch.metadata.replayOverflow).toBe(true);
|
||||
expect(patch.metadata.error).toContain('контекстное окно');
|
||||
});
|
||||
|
||||
it('#490: a non-overflow error does NOT stamp replayOverflow', async () => {
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'error')
|
||||
.mockImplementation(() => undefined as never);
|
||||
const { capturedOpts, aiChatMessageRepo } = await captureStreamCallbacks();
|
||||
await capturedOpts.onError({ error: new Error('network reset') });
|
||||
const calls = aiChatMessageRepo.finalizeOwner.mock.calls as any[][];
|
||||
const patch = calls[calls.length - 1][2] as {
|
||||
status: string;
|
||||
metadata: Record<string, unknown>;
|
||||
};
|
||||
expect('replayOverflow' in patch.metadata).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* F14 — the begin-failure RESILIENCE branch (the `else` of the run-race guard).
|
||||
* F14 — the begin-failure 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.
|
||||
* - ANY OTHER begin failure -> throw ServiceUnavailableException(A_RUN_BEGIN_FAILED)
|
||||
* BEFORE the first byte (#486, commit 4).
|
||||
*
|
||||
* 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.
|
||||
* POLICY CHANGE (#486): the OLD contract here was "SWALLOW + stream the turn
|
||||
* UNTRACKED on the socket signal". That was reversed: an untracked run is
|
||||
* invisible to /stop, is not aborted on disconnect, and slips past the one-run
|
||||
* gate — an unstoppable ghost run in autonomous mode. Now a plain begin failure
|
||||
* FAILS the turn fast with a 503 A_RUN_BEGIN_FAILED, before any user row is
|
||||
* persisted and before streamText runs. This case is INVERTED (not deleted) so
|
||||
* the "plain begin failure" path stays explicitly pinned under the new policy.
|
||||
*/
|
||||
describe('AiChatService.stream — begin-failure resilience / legacy fallback (#184 F14)', () => {
|
||||
describe('AiChatService.stream — begin-failure fails the turn (#184 F14 / #486)', () => {
|
||||
const streamTextMock = streamText as unknown as jest.Mock;
|
||||
|
||||
function makeStreamResult() {
|
||||
@@ -411,6 +468,8 @@ describe('AiChatService.stream — begin-failure resilience / legacy fallback (#
|
||||
insert: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
findAllByChat: jest.fn(async () => []),
|
||||
update: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
finalizeOwner: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
findStreamingWithTerminalRun: jest.fn(async () => []),
|
||||
};
|
||||
const aiSettings = { resolve: jest.fn(async () => ({})) };
|
||||
const tools = { forUser: jest.fn(async () => ({})) };
|
||||
@@ -455,7 +514,7 @@ describe('AiChatService.stream — begin-failure resilience / legacy fallback (#
|
||||
|
||||
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 () => {
|
||||
it('a PLAIN begin() failure (NOT RunAlreadyActiveError) FAILS the turn with a 503 A_RUN_BEGIN_FAILED before the first byte — NO untracked stream (#486)', async () => {
|
||||
const errorSpy = jest
|
||||
.spyOn(Logger.prototype, 'error')
|
||||
.mockImplementation(() => undefined as never);
|
||||
@@ -487,28 +546,26 @@ describe('AiChatService.stream — begin-failure resilience / legacy fallback (#
|
||||
} as never,
|
||||
});
|
||||
|
||||
// The turn proceeds: NO throw at all (in particular NOT a 409).
|
||||
await expect(promise).resolves.toBeUndefined();
|
||||
// NEW POLICY: the turn is REJECTED with a 503 A_RUN_BEGIN_FAILED (not a 409,
|
||||
// and NOT swallowed into an untracked stream).
|
||||
await expect(promise).rejects.toBeInstanceOf(ServiceUnavailableException);
|
||||
const err = (await promise.catch(
|
||||
(e) => e,
|
||||
)) as ServiceUnavailableException;
|
||||
expect(err.getStatus()).toBe(503);
|
||||
expect(err.getResponse()).toMatchObject({ code: 'A_RUN_BEGIN_FAILED' });
|
||||
|
||||
expect(begin).toHaveBeenCalledTimes(1);
|
||||
|
||||
// The resilience branch logged the legacy-fallback warning.
|
||||
// It logged the fail-the-turn line.
|
||||
expect(errorSpy).toHaveBeenCalledWith(
|
||||
expect.stringContaining('streaming without run tracking'),
|
||||
expect.stringContaining('failing the turn'),
|
||||
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. #444:
|
||||
// the signal is unioned with the degeneration controller via AbortSignal.any,
|
||||
// so assert the socket abort still reaches the turn rather than identity.
|
||||
const passed = streamTextMock.mock.calls[0][0].abortSignal as AbortSignal;
|
||||
expect(passed.aborted).toBe(false);
|
||||
socketController.abort();
|
||||
expect(passed.aborted).toBe(true);
|
||||
// Fail-fast: the turn NEVER streamed — no user row persisted, no streamText
|
||||
// call, so no orphan/untracked run was left behind.
|
||||
expect(aiChatMessageRepo.insert).not.toHaveBeenCalled();
|
||||
expect(streamTextMock).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -13,6 +13,7 @@ import {
|
||||
compactToolOutput,
|
||||
assistantParts,
|
||||
serializeSteps,
|
||||
type StepPartsCache,
|
||||
rowToUiMessage,
|
||||
prepareAgentStep,
|
||||
stepBudgetWarning,
|
||||
@@ -28,6 +29,9 @@ import {
|
||||
FINAL_STEP_NUDGE,
|
||||
STEP_LIMIT_NO_ANSWER_MARKER,
|
||||
OUTPUT_DEGENERATION_ERROR,
|
||||
lastAssistantContextTokens,
|
||||
lastAssistantReplayOverflow,
|
||||
seedActivatedTools,
|
||||
} from './ai-chat.service';
|
||||
import type { AiChatMessage, Workspace } from '@docmost/db/types/entity.types';
|
||||
import { buildSystemPrompt } from './ai-chat.prompt';
|
||||
@@ -114,6 +118,54 @@ describe('compactToolOutput', () => {
|
||||
describe('assistantParts', () => {
|
||||
type AnyPart = Record<string, unknown>;
|
||||
|
||||
// #490 memoization: assistantParts builds each step's parts once and caches
|
||||
// them by the step OBJECT's identity, so a mid-stream flush does not
|
||||
// re-stringify every prior step's (large) output. Observable property: with a
|
||||
// shared cache, the second call over the SAME step object returns the cached
|
||||
// (identical) part array even if the step's underlying output was swapped —
|
||||
// proving the work was memoized, not redone.
|
||||
it('memoizes a step by identity (shared cache => one build per step)', () => {
|
||||
const cache: StepPartsCache = new WeakMap();
|
||||
const step = {
|
||||
text: 'x',
|
||||
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: {} }],
|
||||
toolResults: [{ toolCallId: 'c1', toolName: 'getPage', output: { v: 1 } }],
|
||||
};
|
||||
const first = assistantParts([step], '', cache) as AnyPart[];
|
||||
expect((first.find((p) => p.type === 'tool-getPage')!.output as any).v).toBe(
|
||||
1,
|
||||
);
|
||||
// Swap the output for a NEW value; a re-build would pick it up, a cache hit
|
||||
// keeps the first result.
|
||||
step.toolResults[0] = {
|
||||
toolCallId: 'c1',
|
||||
toolName: 'getPage',
|
||||
output: { v: 2 },
|
||||
};
|
||||
const second = assistantParts([step], '', cache) as AnyPart[];
|
||||
expect((second.find((p) => p.type === 'tool-getPage')!.output as any).v).toBe(
|
||||
1,
|
||||
);
|
||||
// Same cached part objects are reused.
|
||||
expect(second.find((p) => p.type === 'tool-getPage')).toBe(
|
||||
first.find((p) => p.type === 'tool-getPage'),
|
||||
);
|
||||
});
|
||||
|
||||
it('without a cache, each call rebuilds (no stale memo)', () => {
|
||||
const step = {
|
||||
text: 'x',
|
||||
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: {} }],
|
||||
toolResults: [{ toolCallId: 'c1', toolName: 'getPage', output: { v: 1 } }],
|
||||
};
|
||||
const first = assistantParts([step], '') as AnyPart[];
|
||||
step.toolResults[0].output = { v: 2 };
|
||||
const second = assistantParts([step], '') as AnyPart[];
|
||||
expect((second.find((p) => p.type === 'tool-getPage')!.output as any).v).toBe(
|
||||
2,
|
||||
);
|
||||
});
|
||||
|
||||
it('emits output-available for a tool-call WITH a paired result', () => {
|
||||
const steps = [
|
||||
{
|
||||
@@ -231,61 +283,299 @@ describe('assistantParts', () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe('serializeSteps', () => {
|
||||
// #490 trace format v2: per call the trace stores { input } for the call and an
|
||||
// OUTCOME element — { ok: true } on success, { error, kind: 'thrown' } on a
|
||||
// thrown tool-error, { error, kind: 'interrupted' } on a mid-step abort. The tool
|
||||
// OUTPUT is no longer duplicated here (it lives once in metadata.parts).
|
||||
describe('serializeSteps (trace v2)', () => {
|
||||
it('returns null when there are no calls or results', () => {
|
||||
expect(serializeSteps([])).toBeNull();
|
||||
});
|
||||
|
||||
it('flattens calls and results into a compact trace', () => {
|
||||
it('pairs a successful call with an { ok: true } outcome and NO output', () => {
|
||||
const trace = serializeSteps([
|
||||
{
|
||||
toolCalls: [{ toolName: 'getPage', input: { id: 'p1' } }],
|
||||
toolResults: [{ toolName: 'getPage', output: { title: 'T' } }],
|
||||
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: { id: 'p1' } }],
|
||||
toolResults: [{ toolCallId: 'c1', toolName: 'getPage' }],
|
||||
},
|
||||
]) as Array<Record<string, unknown>>;
|
||||
expect(trace).toHaveLength(2);
|
||||
expect(trace[0]).toEqual({ toolName: 'getPage', input: { id: 'p1' } });
|
||||
expect(trace[1]).toEqual({ toolName: 'getPage', output: { title: 'T' } });
|
||||
expect(trace[1]).toEqual({ toolName: 'getPage', ok: true });
|
||||
// The output is NOT stored in the trace any more (dedup: it lives in parts).
|
||||
expect(trace.some((e) => 'output' in e)).toBe(false);
|
||||
});
|
||||
|
||||
it('records a THROWN tool failure (tool-error part) with its error message', () => {
|
||||
it('records a THROWN failure with { error, kind: "thrown" }', () => {
|
||||
const trace = serializeSteps([
|
||||
{
|
||||
toolCalls: [{ toolName: 'editPageText', input: { id: 'p1' } }],
|
||||
toolCalls: [
|
||||
{ toolCallId: 'c1', toolName: 'editPageText', input: { id: 'p1' } },
|
||||
],
|
||||
toolResults: [],
|
||||
content: [
|
||||
{
|
||||
type: 'tool-error',
|
||||
toolCallId: 'c1',
|
||||
toolName: 'editPageText',
|
||||
error: new Error('page is locked'),
|
||||
},
|
||||
],
|
||||
},
|
||||
]) as Array<Record<string, unknown>>;
|
||||
// The call element is followed by a paired error element (mirroring how a
|
||||
// successful result is appended), so the failure survives in the trace.
|
||||
expect(trace).toHaveLength(2);
|
||||
expect(trace[0]).toEqual({ toolName: 'editPageText', input: { id: 'p1' } });
|
||||
expect(trace[1]).toEqual({
|
||||
toolName: 'editPageText',
|
||||
error: 'page is locked',
|
||||
kind: 'thrown',
|
||||
});
|
||||
});
|
||||
|
||||
it('truncates a very long tool-error message to the tool-output limit', () => {
|
||||
it('marks an interrupted call (no result, no throw) with kind "interrupted"', () => {
|
||||
const trace = serializeSteps([
|
||||
{
|
||||
toolCalls: [
|
||||
{ toolCallId: 'c1', toolName: 'createComment', input: { x: 1 } },
|
||||
],
|
||||
toolResults: [],
|
||||
content: [],
|
||||
},
|
||||
]) as Array<Record<string, unknown>>;
|
||||
expect(trace).toHaveLength(2);
|
||||
expect(trace[1]).toEqual({
|
||||
toolName: 'createComment',
|
||||
error: 'Tool call did not complete.',
|
||||
kind: 'interrupted',
|
||||
});
|
||||
// Structurally distinct from a thrown hard-fail so it never inflates an
|
||||
// error-rate scan.
|
||||
expect((trace[1] as { kind: string }).kind).not.toBe('thrown');
|
||||
});
|
||||
|
||||
it('truncates a very long thrown-error message to the tool-output limit', () => {
|
||||
const long = 'x'.repeat(5000);
|
||||
const trace = serializeSteps([
|
||||
{
|
||||
toolCalls: [{ toolName: 'editPageText', input: {} }],
|
||||
toolCalls: [{ toolCallId: 'c1', toolName: 'editPageText', input: {} }],
|
||||
toolResults: [],
|
||||
content: [{ type: 'tool-error', toolName: 'editPageText', error: long }],
|
||||
content: [
|
||||
{
|
||||
type: 'tool-error',
|
||||
toolCallId: 'c1',
|
||||
toolName: 'editPageText',
|
||||
error: long,
|
||||
},
|
||||
],
|
||||
},
|
||||
]) as Array<Record<string, unknown>>;
|
||||
const errorText = trace[1].error as string;
|
||||
// Truncated (not the full 5000 chars) and carries the omission marker.
|
||||
expect(errorText.length).toBeLessThan(long.length);
|
||||
expect(errorText).toContain('chars omitted');
|
||||
});
|
||||
|
||||
it('pairs parallel calls in one step with their outcomes by id', () => {
|
||||
const trace = serializeSteps([
|
||||
{
|
||||
toolCalls: [
|
||||
{ toolCallId: 'a', toolName: 'getPage', input: {} },
|
||||
{ toolCallId: 'b', toolName: 'searchPages', input: {} },
|
||||
],
|
||||
toolResults: [{ toolCallId: 'b', toolName: 'searchPages' }],
|
||||
content: [
|
||||
{ type: 'tool-error', toolCallId: 'a', toolName: 'getPage', error: 'nope' },
|
||||
],
|
||||
},
|
||||
]) as Array<Record<string, unknown>>;
|
||||
// call a, outcome a (thrown), call b, outcome b (ok)
|
||||
expect(trace).toHaveLength(4);
|
||||
expect(trace[1]).toEqual({ toolName: 'getPage', error: 'nope', kind: 'thrown' });
|
||||
expect(trace[3]).toEqual({ toolName: 'searchPages', ok: true });
|
||||
});
|
||||
});
|
||||
|
||||
// #490: every assistant row flushAssistant writes carries the v2 era marker so a
|
||||
// dual-shape diagnostic query can branch on the trace shape without inspecting it.
|
||||
describe('toolTraceVersion era marker (#490)', () => {
|
||||
it('stamps metadata.toolTraceVersion = 2 on every flushed row', () => {
|
||||
const seed = flushAssistant([], '', 'streaming');
|
||||
expect(seed.metadata.toolTraceVersion).toBe(2);
|
||||
const done = flushAssistant(
|
||||
[
|
||||
{
|
||||
text: 'ok',
|
||||
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: {} }],
|
||||
toolResults: [{ toolCallId: 'c1', toolName: 'getPage' }],
|
||||
},
|
||||
],
|
||||
'',
|
||||
'completed',
|
||||
{ finishReason: 'stop' },
|
||||
);
|
||||
expect(done.metadata.toolTraceVersion).toBe(2);
|
||||
});
|
||||
});
|
||||
|
||||
// #490 replay-budget signal helpers over persisted history.
|
||||
describe('lastAssistantContextTokens', () => {
|
||||
const row = (
|
||||
role: string,
|
||||
metadata: Record<string, unknown> | null,
|
||||
): AiChatMessage => ({ role, metadata }) as unknown as AiChatMessage;
|
||||
|
||||
it('reads the most recent assistant turn contextTokens (provider fact)', () => {
|
||||
const hist = [
|
||||
row('user', null),
|
||||
row('assistant', { contextTokens: 12000 }),
|
||||
row('user', null),
|
||||
row('assistant', { contextTokens: 41000 }),
|
||||
];
|
||||
expect(lastAssistantContextTokens(hist)).toBe(41000);
|
||||
});
|
||||
|
||||
it('returns undefined when the last assistant turn recorded no usage', () => {
|
||||
const hist = [row('assistant', { error: 'boom' }), row('user', null)];
|
||||
expect(lastAssistantContextTokens(hist)).toBeUndefined();
|
||||
expect(lastAssistantContextTokens([])).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
// #490 snapshotOpenPage fast-path: skip the full Markdown export + upsert when a
|
||||
// snapshot already exists at the page's CURRENT version (same updated_at instant).
|
||||
describe('snapshotOpenPage fast-path (#490)', () => {
|
||||
function makeSvc(existingSnapshot: unknown, pageUpdatedAt: Date) {
|
||||
const exportPageMarkdown = jest.fn(async () => '# md');
|
||||
const upsert = jest.fn(async () => undefined);
|
||||
const findByChatPage = jest.fn(async () => existingSnapshot);
|
||||
const pageRepo = {
|
||||
findById: jest.fn(async () => ({
|
||||
id: 'p1',
|
||||
workspaceId: 'ws1',
|
||||
updatedAt: pageUpdatedAt,
|
||||
})),
|
||||
};
|
||||
const svc = new AiChatService(
|
||||
{} as never, // ai
|
||||
{} as never, // aiChatRepo
|
||||
{} as never, // aiChatMessageRepo
|
||||
{ findByChatPage, upsert } as never, // aiChatPageSnapshotRepo
|
||||
{} as never, // aiSettings
|
||||
{ exportPageMarkdown } as never, // tools
|
||||
{} as never, // mcpClients
|
||||
{} as never, // aiAgentRoleRepo
|
||||
pageRepo as never, // pageRepo
|
||||
{} as never, // pageAccess
|
||||
{} as never, // environment
|
||||
);
|
||||
return { svc, exportPageMarkdown, upsert, findByChatPage };
|
||||
}
|
||||
|
||||
const args = () =>
|
||||
[
|
||||
'chat1',
|
||||
'p1',
|
||||
{ id: 'ws1' } as never,
|
||||
{ id: 'u1' } as never,
|
||||
'sess',
|
||||
] as const;
|
||||
|
||||
it('skips export + upsert when the snapshot is already at this page version', async () => {
|
||||
const t = new Date('2026-07-07T10:00:00Z');
|
||||
const { svc, exportPageMarkdown, upsert } = makeSvc(
|
||||
{ pageUpdatedAt: t, contentMd: '# md' },
|
||||
t,
|
||||
);
|
||||
await (svc as unknown as { snapshotOpenPage: (...a: unknown[]) => Promise<void> })
|
||||
.snapshotOpenPage(...args());
|
||||
expect(exportPageMarkdown).not.toHaveBeenCalled();
|
||||
expect(upsert).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('exports + upserts when the page advanced since the snapshot', async () => {
|
||||
const { svc, exportPageMarkdown, upsert } = makeSvc(
|
||||
{ pageUpdatedAt: new Date('2026-07-07T10:00:00Z'), contentMd: 'old' },
|
||||
new Date('2026-07-07T11:00:00Z'),
|
||||
);
|
||||
await (svc as unknown as { snapshotOpenPage: (...a: unknown[]) => Promise<void> })
|
||||
.snapshotOpenPage(...args());
|
||||
expect(exportPageMarkdown).toHaveBeenCalledTimes(1);
|
||||
expect(upsert).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('seeds (exports + upserts) on the first turn (no snapshot yet)', async () => {
|
||||
const { svc, exportPageMarkdown, upsert } = makeSvc(
|
||||
undefined,
|
||||
new Date('2026-07-07T10:00:00Z'),
|
||||
);
|
||||
await (svc as unknown as { snapshotOpenPage: (...a: unknown[]) => Promise<void> })
|
||||
.snapshotOpenPage(...args());
|
||||
expect(exportPageMarkdown).toHaveBeenCalledTimes(1);
|
||||
expect(upsert).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
|
||||
// #490 deferred-tool activation persisted across turns.
|
||||
describe('seedActivatedTools', () => {
|
||||
const valid = new Set(['Search_web', 'getPageJson', 'diffPageVersions']);
|
||||
|
||||
it('seeds from persisted metadata, intersected with current valid names', () => {
|
||||
expect(
|
||||
seedActivatedTools(
|
||||
{ activatedTools: ['Search_web', 'getPageJson'] },
|
||||
valid,
|
||||
),
|
||||
).toEqual(['Search_web', 'getPageJson']);
|
||||
});
|
||||
|
||||
it('drops a stored tool that is no longer valid (allowlist/role changed)', () => {
|
||||
// 'Habr_publish' was activated before but is not in the current allowlist.
|
||||
expect(
|
||||
seedActivatedTools({ activatedTools: ['Search_web', 'Habr_publish'] }, valid),
|
||||
).toEqual(['Search_web']);
|
||||
});
|
||||
|
||||
it('is empty/robust for missing, non-array, or unknown-shaped metadata', () => {
|
||||
expect(seedActivatedTools(undefined, valid)).toEqual([]);
|
||||
expect(seedActivatedTools({}, valid)).toEqual([]);
|
||||
expect(seedActivatedTools({ activatedTools: 'nope' }, valid)).toEqual([]);
|
||||
expect(
|
||||
seedActivatedTools({ activatedTools: [1, 'getPageJson', null] }, valid),
|
||||
).toEqual(['getPageJson']);
|
||||
});
|
||||
|
||||
it('de-duplicates stored names', () => {
|
||||
expect(
|
||||
seedActivatedTools(
|
||||
{ activatedTools: ['getPageJson', 'getPageJson'] },
|
||||
valid,
|
||||
),
|
||||
).toEqual(['getPageJson']);
|
||||
});
|
||||
});
|
||||
|
||||
describe('lastAssistantReplayOverflow', () => {
|
||||
const row = (
|
||||
role: string,
|
||||
metadata: Record<string, unknown> | null,
|
||||
): AiChatMessage => ({ role, metadata }) as unknown as AiChatMessage;
|
||||
|
||||
it('is true only when the LAST assistant turn overflowed', () => {
|
||||
expect(
|
||||
lastAssistantReplayOverflow([
|
||||
row('assistant', { replayOverflow: true }),
|
||||
row('user', null),
|
||||
]),
|
||||
).toBe(true);
|
||||
// A recovered (later, non-overflow) assistant turn clears it.
|
||||
expect(
|
||||
lastAssistantReplayOverflow([
|
||||
row('assistant', { replayOverflow: true }),
|
||||
row('user', null),
|
||||
row('assistant', { contextTokens: 5 }),
|
||||
]),
|
||||
).toBe(false);
|
||||
expect(lastAssistantReplayOverflow([])).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('rowToUiMessage', () => {
|
||||
@@ -618,6 +908,23 @@ describe('flushAssistant', () => {
|
||||
expect(flushed.metadata.error).toBe('boom');
|
||||
});
|
||||
|
||||
// #490 observability: the replay budgeter's decision is stamped on the turn.
|
||||
it('records replayTrimmedToTokens + replayOverflow when provided', () => {
|
||||
const f = flushAssistant([], '', 'error', {
|
||||
error: 'ctx',
|
||||
replayTrimmedToTokens: 42_000,
|
||||
replayOverflow: true,
|
||||
});
|
||||
expect(f.metadata.replayTrimmedToTokens).toBe(42_000);
|
||||
expect(f.metadata.replayOverflow).toBe(true);
|
||||
});
|
||||
|
||||
it('omits the replay metadata when not provided', () => {
|
||||
const f = flushAssistant([], '', 'completed', { finishReason: 'stop' });
|
||||
expect('replayTrimmedToTokens' in f.metadata).toBe(false);
|
||||
expect('replayOverflow' in f.metadata).toBe(false);
|
||||
});
|
||||
|
||||
// #274 observability: the page-change diff the agent saw this turn is persisted
|
||||
// to metadata.pageChanged when a non-empty diff was injected, and omitted when
|
||||
// the diff is empty/whitespace or the arg is not supplied.
|
||||
@@ -1315,8 +1622,12 @@ describe('AiChatService page-change lifecycle (#274)', () => {
|
||||
describe('isInterruptResume', () => {
|
||||
// history tail is the just-inserted user row; [len-2] is the previous turn.
|
||||
const withPrev = (
|
||||
prev: { role: string; status?: string | null } | null,
|
||||
): Array<{ role: string; status?: string | null }> =>
|
||||
prev: {
|
||||
role: string;
|
||||
status?: string | null;
|
||||
metadata?: unknown;
|
||||
} | null,
|
||||
): Array<{ role: string; status?: string | null; metadata?: unknown }> =>
|
||||
prev
|
||||
? [prev, { role: 'user', status: null }]
|
||||
: [{ role: 'user', status: null }];
|
||||
@@ -1357,6 +1668,33 @@ describe('isInterruptResume', () => {
|
||||
it('false when there is no preceding turn (only the new user row)', () => {
|
||||
expect(isInterruptResume(withPrev(null), true)).toBe(false);
|
||||
});
|
||||
|
||||
it('#487 EXCLUDES a reconcile stamp (finalizeFailed) — not a genuine interruption', () => {
|
||||
// A row a reconcile settled to 'aborted' carries metadata.finalizeFailed. It
|
||||
// must NOT be treated as an interrupt-resume (that would inject a false
|
||||
// "you were interrupted" note), even though its status is 'aborted'.
|
||||
expect(
|
||||
isInterruptResume(
|
||||
withPrev({
|
||||
role: 'assistant',
|
||||
status: 'aborted',
|
||||
metadata: { finalizeFailed: true },
|
||||
}),
|
||||
true,
|
||||
),
|
||||
).toBe(false);
|
||||
// A genuine abort (no finalizeFailed) still counts.
|
||||
expect(
|
||||
isInterruptResume(
|
||||
withPrev({
|
||||
role: 'assistant',
|
||||
status: 'aborted',
|
||||
metadata: { parts: [] },
|
||||
}),
|
||||
true,
|
||||
),
|
||||
).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
@@ -1409,7 +1747,7 @@ describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', ()
|
||||
}
|
||||
|
||||
// Wire only the deps reached on the way to the pipe call, plus a spy registry.
|
||||
function makeService(opts: { resumable: boolean }) {
|
||||
function makeService(opts: { resumable: boolean; history?: unknown[] }) {
|
||||
const aiChatRepo = {
|
||||
findById: jest.fn(async () => ({ id: 'chat-1', workspaceId: 'ws-1' })),
|
||||
insert: jest.fn(),
|
||||
@@ -1417,8 +1755,11 @@ describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', ()
|
||||
const aiChatMessageRepo = {
|
||||
// Both the user insert and the assistant seed return the same row id.
|
||||
insert: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
findAllByChat: jest.fn(async () => []),
|
||||
findAllByChat: jest.fn(async () => opts.history ?? []),
|
||||
update: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
// #487: the terminal owner-write + the opportunistic reconcile query.
|
||||
finalizeOwner: jest.fn(async () => ({ id: 'msg-1' })),
|
||||
findStreamingWithTerminalRun: jest.fn(async () => []),
|
||||
};
|
||||
const aiSettings = { resolve: jest.fn(async () => ({})) };
|
||||
const tools = { forUser: jest.fn(async () => ({})) };
|
||||
@@ -1453,7 +1794,7 @@ describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', ()
|
||||
} as never,
|
||||
streamRegistry as never,
|
||||
);
|
||||
return { svc, streamRegistry };
|
||||
return { svc, streamRegistry, aiChatMessageRepo };
|
||||
}
|
||||
|
||||
const body = {
|
||||
@@ -1536,6 +1877,86 @@ describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', ()
|
||||
await expect(drive(svc, makeRunHooks())).rejects.toThrow('boom');
|
||||
expect(streamRegistry.abortEntry).toHaveBeenCalledWith('chat-1', 'run-1');
|
||||
});
|
||||
|
||||
// #489 REGRESSION (against the REAL convertToModelMessages — not mocked here):
|
||||
// a persisted history row whose parts contain a `null` element makes the real
|
||||
// convertToModelMessages THROW ("Cannot read properties of null"). Pre-fix that
|
||||
// 500-ed every turn forever and each retry appended a duplicate user row. The
|
||||
// fix converts BEFORE the insert and isolates the poisoned row per-row, degrading
|
||||
// it to text with a "[tool context omitted]" marker. Assert the turn still runs,
|
||||
// the marker reaches the model, and exactly ONE user row is inserted.
|
||||
it('#489: a poisoned OLD-history row keeps the chat working; the marker reaches the model; one user insert', async () => {
|
||||
const { svc, aiChatMessageRepo } = makeService({
|
||||
resumable: false,
|
||||
history: [
|
||||
{
|
||||
id: 'old-1',
|
||||
role: 'assistant',
|
||||
content: 'earlier answer',
|
||||
// A null part is the poison: rowToUiMessage keeps it (the array is
|
||||
// non-empty) and the real convertToModelMessages throws on it.
|
||||
metadata: { parts: [{ type: 'text', text: 'earlier answer' }, null] },
|
||||
status: 'completed',
|
||||
},
|
||||
],
|
||||
});
|
||||
// Must NOT throw — the poisoned row is degraded, not fatal.
|
||||
await drive(svc, makeRunHooks());
|
||||
expect(streamTextMock).toHaveBeenCalledTimes(1);
|
||||
const passedMessages = streamTextMock.mock.calls[0][0].messages;
|
||||
const serialized = JSON.stringify(passedMessages);
|
||||
// The model sees the truncation marker (silent tool-context loss is not ok)
|
||||
// AND the row's readable text is preserved alongside it.
|
||||
expect(serialized).toContain('[tool context omitted]');
|
||||
expect(serialized).toContain('earlier answer');
|
||||
// Exactly ONE user row inserted (no duplicate), inserted AFTER conversion.
|
||||
const userInserts = aiChatMessageRepo.insert.mock.calls
|
||||
.map((c: unknown[]) => c[0] as { role?: string })
|
||||
.filter((r) => r.role === 'user');
|
||||
expect(userInserts).toHaveLength(1);
|
||||
});
|
||||
|
||||
// #489: client-supplied non-text parts (a tool-part in `input-available`, the
|
||||
// exact "bricking" payload) are dropped ON RECEIPT — never persisted — so they
|
||||
// can never poison future turns. Only the text survives into metadata.parts.
|
||||
it('#489: a non-text client part is stripped before persist (only text survives)', async () => {
|
||||
const { svc, aiChatMessageRepo } = makeService({ resumable: false });
|
||||
await svc.stream({
|
||||
user: { id: 'u1' } as never,
|
||||
workspace: { id: 'ws-1' } as never,
|
||||
sessionId: 's1',
|
||||
body: {
|
||||
chatId: 'chat-1',
|
||||
messages: [
|
||||
{
|
||||
id: 'm1',
|
||||
role: 'user',
|
||||
parts: [
|
||||
{ type: 'text', text: 'hello' },
|
||||
// untrusted tool-part — must be dropped, never persisted
|
||||
{
|
||||
type: 'tool-getPage',
|
||||
toolCallId: 't1',
|
||||
state: 'input-available',
|
||||
input: { pageId: 'p' },
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
} as never,
|
||||
res: makeRes() as never,
|
||||
signal: new AbortController().signal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
runHooks: makeRunHooks() as never,
|
||||
});
|
||||
const userInsert = aiChatMessageRepo.insert.mock.calls
|
||||
.map((c: unknown[]) => c[0] as { role?: string; metadata?: unknown })
|
||||
.find((r) => r.role === 'user');
|
||||
const parts = (userInsert?.metadata as { parts?: Array<{ type: string }> })
|
||||
?.parts;
|
||||
expect(parts).toEqual([{ type: 'text', text: 'hello' }]);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
@@ -1623,6 +2044,19 @@ describe('AiChatService.stream — token-degeneration reaction (#444)', () => {
|
||||
return { id };
|
||||
},
|
||||
),
|
||||
// #487: the terminal owner-write records into the SAME `updated` recorder so
|
||||
// assertions on the terminal 'completed'/'error'/'aborted' write still hold.
|
||||
finalizeOwner: jest.fn(
|
||||
async (
|
||||
id: string,
|
||||
workspaceId: string,
|
||||
patch: Record<string, unknown>,
|
||||
) => {
|
||||
updated.push({ id, workspaceId, patch });
|
||||
return { id };
|
||||
},
|
||||
),
|
||||
findStreamingWithTerminalRun: jest.fn(async () => []),
|
||||
};
|
||||
const aiSettings = { resolve: jest.fn(async () => ({})) };
|
||||
const tools = { forUser: jest.fn(async () => ({})) };
|
||||
@@ -1882,3 +2316,148 @@ describe('AiChatService.stream — token-degeneration reaction (#444)', () => {
|
||||
expect(patch.content).not.toContain(STEP_LIMIT_NO_ANSWER_MARKER);
|
||||
});
|
||||
});
|
||||
|
||||
// #487 F3 — the reconcile() / reconcileChat() ORCHESTRATORS. The individual
|
||||
// clauses are exercised elsewhere; these pin the production orchestration the
|
||||
// per-clause specs do not: the clause ORDER, the per-clause try/catch ISOLATION
|
||||
// (one clause throwing must NOT abort the others), and reconcileChat() (which runs
|
||||
// at the start of every turn and was entirely uncovered).
|
||||
describe('AiChatService.reconcile / reconcileChat orchestrators (#487 F3)', () => {
|
||||
let warnSpy: jest.SpyInstance;
|
||||
beforeEach(() => {
|
||||
// Silence the intentional clause-failure warnings (kept out of test output).
|
||||
warnSpy = jest
|
||||
.spyOn(Logger.prototype, 'warn')
|
||||
.mockImplementation(() => undefined);
|
||||
});
|
||||
afterEach(() => {
|
||||
warnSpy.mockRestore();
|
||||
});
|
||||
|
||||
function makeService(opts: {
|
||||
messageRepo?: Record<string, jest.Mock>;
|
||||
runService?: Record<string, jest.Mock>;
|
||||
}) {
|
||||
const aiChatMessageRepo = {
|
||||
findStreamingWithTerminalRun: jest.fn(async () => []),
|
||||
stampTerminalIfStreaming: jest.fn(async () => undefined),
|
||||
sweepStreamingWithoutActiveRun: jest.fn(async () => 0),
|
||||
...(opts.messageRepo ?? {}),
|
||||
};
|
||||
const aiChatRunService = opts.runService
|
||||
? {
|
||||
zombieRunIds: jest.fn(() => []),
|
||||
settleZombie: jest.fn(async () => true),
|
||||
reconcileStaleRuns: jest.fn(async () => 0),
|
||||
...opts.runService,
|
||||
}
|
||||
: undefined;
|
||||
const svc = new AiChatService(
|
||||
{} as never, // ai
|
||||
{} as never, // aiChatRepo
|
||||
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
|
||||
{} as never, // streamRegistry
|
||||
aiChatRunService as never, // aiChatRunService (#487)
|
||||
);
|
||||
return { svc, aiChatMessageRepo, aiChatRunService };
|
||||
}
|
||||
|
||||
it('reconcile() fires all four clauses IN ORDER (a -> b -> c -> d)', async () => {
|
||||
const order: string[] = [];
|
||||
const { svc } = makeService({
|
||||
messageRepo: {
|
||||
findStreamingWithTerminalRun: jest.fn(async () => {
|
||||
order.push('b:find');
|
||||
return [
|
||||
{ messageId: 'm1', workspaceId: 'ws1', runStatus: 'succeeded' },
|
||||
];
|
||||
}),
|
||||
stampTerminalIfStreaming: jest.fn(async () => {
|
||||
order.push('b:stamp');
|
||||
}),
|
||||
sweepStreamingWithoutActiveRun: jest.fn(async () => {
|
||||
order.push('d');
|
||||
return 0;
|
||||
}),
|
||||
},
|
||||
runService: {
|
||||
zombieRunIds: jest.fn(() => ['z1']),
|
||||
settleZombie: jest.fn(async () => {
|
||||
order.push('a');
|
||||
return true;
|
||||
}),
|
||||
reconcileStaleRuns: jest.fn(async () => {
|
||||
order.push('c');
|
||||
return 0;
|
||||
}),
|
||||
},
|
||||
});
|
||||
|
||||
await svc.reconcile();
|
||||
|
||||
expect(order).toEqual(['a', 'b:find', 'b:stamp', 'c', 'd']);
|
||||
});
|
||||
|
||||
it('a clause that THROWS does not abort the remaining clauses (per-clause try/catch isolation)', async () => {
|
||||
const { svc, aiChatMessageRepo, aiChatRunService } = makeService({
|
||||
messageRepo: {
|
||||
// Clause (b) blows up mid-reconcile.
|
||||
findStreamingWithTerminalRun: jest.fn(async () => {
|
||||
throw new Error('clause b DB blip');
|
||||
}),
|
||||
},
|
||||
runService: {
|
||||
zombieRunIds: jest.fn(() => ['z1']),
|
||||
},
|
||||
});
|
||||
|
||||
// reconcile() must SETTLE (the clause-b failure is swallowed), not reject.
|
||||
await expect(svc.reconcile()).resolves.toBeUndefined();
|
||||
|
||||
// (a) ran before (b); crucially (c) and (d) STILL ran despite (b) throwing —
|
||||
// the property a missing try/catch would break. MUTATION-VERIFY: drop clause
|
||||
// (b)'s try/catch and this reddens (the throw propagates, skipping c + d).
|
||||
expect(aiChatRunService!.settleZombie).toHaveBeenCalled(); // (a)
|
||||
expect(aiChatRunService!.reconcileStaleRuns).toHaveBeenCalled(); // (c)
|
||||
expect(
|
||||
aiChatMessageRepo.sweepStreamingWithoutActiveRun,
|
||||
).toHaveBeenCalled(); // (d)
|
||||
});
|
||||
|
||||
it('reconcileChat() settles THIS chat\'s stuck streaming rows by their run status', async () => {
|
||||
const { svc, aiChatMessageRepo } = makeService({
|
||||
messageRepo: {
|
||||
findStreamingWithTerminalRun: jest.fn(async () => [
|
||||
{ messageId: 'm1', workspaceId: 'ws1', runStatus: 'failed' },
|
||||
{ messageId: 'm2', workspaceId: 'ws1', runStatus: 'succeeded' },
|
||||
]),
|
||||
},
|
||||
});
|
||||
|
||||
await svc.reconcileChat('chat-1', 'ws1');
|
||||
|
||||
// Scoped to THIS chat and bounded at 50 (the user-facing opportunistic path).
|
||||
expect(
|
||||
aiChatMessageRepo.findStreamingWithTerminalRun,
|
||||
).toHaveBeenCalledWith(50, { chatId: 'chat-1', workspaceId: 'ws1' });
|
||||
// failed-run -> 'error'; every other terminal status -> 'aborted'.
|
||||
expect(aiChatMessageRepo.stampTerminalIfStreaming).toHaveBeenCalledWith(
|
||||
'm1',
|
||||
'ws1',
|
||||
'error',
|
||||
);
|
||||
expect(aiChatMessageRepo.stampTerminalIfStreaming).toHaveBeenCalledWith(
|
||||
'm2',
|
||||
'ws1',
|
||||
'aborted',
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,209 @@
|
||||
import { randomBytes } from 'crypto';
|
||||
import { Client } from 'pg';
|
||||
import { flushAssistant, serializeSteps } from './ai-chat.service';
|
||||
|
||||
/**
|
||||
* #490 write-volume regression — an OBSERVABLE-PROPERTY test on a LIVE Postgres,
|
||||
* not "bytes through a mock repo" (a mock measures exactly the thing that does not
|
||||
* hurt). It drives a realistic 50-step run where each step returns a ~100 KB tool
|
||||
* output and, at every `onStepFinish`, UPDATEs the assistant row the way the
|
||||
* service does — then reads the REAL write volume via the `pg_current_wal_lsn()`
|
||||
* delta around the run.
|
||||
*
|
||||
* The property proven: v2 stores each tool OUTPUT only in `metadata.parts`, no
|
||||
* longer ALSO in the `tool_calls` trace. So:
|
||||
* 1. the trace (`tool_calls`) column's write volume is now O(Σ steps) — tiny,
|
||||
* linear outcome flags — vs the pre-#490 O(N²) that re-persisted every prior
|
||||
* output on every step; and
|
||||
* 2. the FULL-row write volume drops sharply (the duplicated output copy is gone).
|
||||
*
|
||||
* Connects to the local gitmost test Postgres (docker `gitmost-test-pg` on :5432);
|
||||
* SKIPS cleanly when that DB is not reachable so it never breaks a DB-less CI.
|
||||
*/
|
||||
const CONN =
|
||||
process.env.WAL_TEST_DATABASE_URL ??
|
||||
'postgresql://docmost:docmost_dev_pw@localhost:5432/docmost';
|
||||
|
||||
// A step whose tool output is ~100 KB (a page read), in the SDK StepLike shape.
|
||||
// The body is INCOMPRESSIBLE random text — a `'x'.repeat()` filler would TOAST-
|
||||
// compress to nothing and hide the real write volume (a page body does not).
|
||||
function makeStep(i: number, outputBytes = 100_000) {
|
||||
const body = randomBytes(Math.ceil(outputBytes * 0.75)).toString('base64');
|
||||
return {
|
||||
text: `step ${i} reasoning`,
|
||||
toolCalls: [{ toolCallId: `c${i}`, toolName: 'getPage', input: { id: `p${i}` } }],
|
||||
toolResults: [
|
||||
{
|
||||
toolCallId: `c${i}`,
|
||||
toolName: 'getPage',
|
||||
output: { id: `p${i}`, title: `Page ${i}`, body },
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
// The pre-#490 (v1) trace: outputs stored a SECOND time in `tool_calls`
|
||||
// (the duplication #490 removed). Mirrors the OLD serializeSteps shape.
|
||||
function v1Trace(steps: ReturnType<typeof makeStep>[]): unknown {
|
||||
const calls: unknown[] = [];
|
||||
for (const s of steps) {
|
||||
for (const c of s.toolCalls) calls.push({ toolName: c.toolName, input: c.input });
|
||||
for (const r of s.toolResults)
|
||||
calls.push({ toolName: r.toolName, output: r.output });
|
||||
}
|
||||
return calls;
|
||||
}
|
||||
|
||||
async function walDelta(
|
||||
client: Client,
|
||||
fn: () => Promise<void>,
|
||||
): Promise<number> {
|
||||
const before = (await client.query('SELECT pg_current_wal_lsn() AS l')).rows[0]
|
||||
.l as string;
|
||||
await fn();
|
||||
// NOTE: do NOT pg_switch_wal() here — a segment switch pads the LSN to the next
|
||||
// 16 MB boundary and would swamp the actual write delta. The raw LSN advances by
|
||||
// the bytes of WAL emitted, which is exactly what we want to measure.
|
||||
const after = (await client.query('SELECT pg_current_wal_lsn() AS l')).rows[0]
|
||||
.l as string;
|
||||
return Number(
|
||||
(await client.query('SELECT pg_wal_lsn_diff($1,$2) AS d', [after, before]))
|
||||
.rows[0].d,
|
||||
);
|
||||
}
|
||||
|
||||
describe('#490 write-volume on a live Postgres (pg_current_wal_lsn delta)', () => {
|
||||
let client: Client | undefined;
|
||||
let available = false;
|
||||
|
||||
beforeAll(async () => {
|
||||
try {
|
||||
client = new Client(CONN);
|
||||
await client.connect();
|
||||
await client.query('SELECT pg_current_wal_lsn()');
|
||||
available = true;
|
||||
} catch {
|
||||
available = false;
|
||||
client = undefined;
|
||||
}
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await client?.end().catch(() => undefined);
|
||||
});
|
||||
|
||||
const STEPS = 50;
|
||||
|
||||
it('v2 trace write volume is O(Σ steps) — a tiny fraction of the v1 duplicate', async () => {
|
||||
if (!available || !client) {
|
||||
console.warn('SKIP: gitmost-test-pg not reachable; skipping WAL test.');
|
||||
return;
|
||||
}
|
||||
const c = client;
|
||||
// Isolated table so we measure only the tool_calls (trace) column's writes.
|
||||
await c.query('DROP TABLE IF EXISTS _wal_trace');
|
||||
await c.query('CREATE TABLE _wal_trace(id int primary key, tool_calls jsonb)');
|
||||
await c.query("INSERT INTO _wal_trace VALUES (1, '[]'::jsonb)");
|
||||
|
||||
const steps: ReturnType<typeof makeStep>[] = [];
|
||||
|
||||
// v1: each step re-persists ALL prior outputs into the trace (the O(N²) churn).
|
||||
const v1 = await walDelta(c, async () => {
|
||||
const acc: ReturnType<typeof makeStep>[] = [];
|
||||
for (let i = 0; i < STEPS; i++) {
|
||||
acc.push(makeStep(i));
|
||||
await c.query('UPDATE _wal_trace SET tool_calls=$1 WHERE id=1', [
|
||||
JSON.stringify(v1Trace(acc)),
|
||||
]);
|
||||
}
|
||||
steps.push(...acc);
|
||||
});
|
||||
|
||||
await c.query("UPDATE _wal_trace SET tool_calls='[]'::jsonb WHERE id=1");
|
||||
|
||||
// v2: the REAL serializeSteps — outcome flags only, NO outputs.
|
||||
const v2 = await walDelta(c, async () => {
|
||||
const acc: ReturnType<typeof makeStep>[] = [];
|
||||
for (let i = 0; i < STEPS; i++) {
|
||||
acc.push(makeStep(i));
|
||||
await c.query('UPDATE _wal_trace SET tool_calls=$1 WHERE id=1', [
|
||||
JSON.stringify(serializeSteps(acc)),
|
||||
]);
|
||||
}
|
||||
});
|
||||
|
||||
await c.query('DROP TABLE IF EXISTS _wal_trace');
|
||||
|
||||
// eslint-disable-next-line no-console
|
||||
console.log(
|
||||
`[#490 WAL] trace column over ${STEPS} steps: v1=${(v1 / 1e6).toFixed(1)}MB ` +
|
||||
`v2=${(v2 / 1e6).toFixed(2)}MB (${(v1 / v2).toFixed(0)}x smaller)`,
|
||||
);
|
||||
|
||||
// The trace no longer carries outputs: v2 is a tiny fraction of v1's WAL.
|
||||
expect(v2).toBeLessThan(v1 * 0.1);
|
||||
// And v2's trace WAL is small in absolute terms — O(Σ steps) of flags, not
|
||||
// O(N² × output). 50 steps of ~40-byte flags is well under a few MB of WAL.
|
||||
expect(v2).toBeLessThan(5_000_000);
|
||||
// v1's duplicate alone is huge (≈ the 100 KB output re-written N² times).
|
||||
expect(v1).toBeGreaterThan(50_000_000);
|
||||
}, 120_000);
|
||||
|
||||
it('the full assistant row write drops sharply once the duplicate is gone', async () => {
|
||||
if (!available || !client) return;
|
||||
const c = client;
|
||||
await c.query('DROP TABLE IF EXISTS _wal_full');
|
||||
await c.query(
|
||||
'CREATE TABLE _wal_full(id int primary key, content text, tool_calls jsonb, metadata jsonb, status text)',
|
||||
);
|
||||
await c.query("INSERT INTO _wal_full VALUES (1, '', '[]'::jsonb, '{}'::jsonb, 'streaming')");
|
||||
|
||||
const writeRow = async (patch: {
|
||||
content: string;
|
||||
toolCalls: unknown;
|
||||
metadata: unknown;
|
||||
status: string;
|
||||
}) =>
|
||||
c.query(
|
||||
'UPDATE _wal_full SET content=$1, tool_calls=$2, metadata=$3, status=$4 WHERE id=1',
|
||||
[
|
||||
patch.content,
|
||||
JSON.stringify(patch.toolCalls ?? null),
|
||||
JSON.stringify(patch.metadata),
|
||||
patch.status,
|
||||
],
|
||||
);
|
||||
|
||||
// v2 (real flushAssistant): outputs live once, in metadata.parts.
|
||||
const v2 = await walDelta(c, async () => {
|
||||
const acc: ReturnType<typeof makeStep>[] = [];
|
||||
for (let i = 0; i < STEPS; i++) {
|
||||
acc.push(makeStep(i));
|
||||
await writeRow(flushAssistant(acc as never, '', 'streaming'));
|
||||
}
|
||||
});
|
||||
|
||||
await c.query("UPDATE _wal_full SET content='', tool_calls='[]'::jsonb, metadata='{}'::jsonb WHERE id=1");
|
||||
|
||||
// v1: same row PLUS the duplicated outputs in the trace column.
|
||||
const v1 = await walDelta(c, async () => {
|
||||
const acc: ReturnType<typeof makeStep>[] = [];
|
||||
for (let i = 0; i < STEPS; i++) {
|
||||
acc.push(makeStep(i));
|
||||
const f = flushAssistant(acc as never, '', 'streaming');
|
||||
await writeRow({ ...f, toolCalls: v1Trace(acc) });
|
||||
}
|
||||
});
|
||||
|
||||
await c.query('DROP TABLE IF EXISTS _wal_full');
|
||||
|
||||
// eslint-disable-next-line no-console
|
||||
console.log(
|
||||
`[#490 WAL] full row over ${STEPS} steps: v1=${(v1 / 1e6).toFixed(1)}MB ` +
|
||||
`v2=${(v2 / 1e6).toFixed(1)}MB (saved ${((1 - v2 / v1) * 100).toFixed(0)}%)`,
|
||||
);
|
||||
|
||||
// Removing the duplicated trace copy is a large, real write-volume reduction.
|
||||
expect(v2).toBeLessThan(v1 * 0.75);
|
||||
}, 120_000);
|
||||
});
|
||||
@@ -0,0 +1,261 @@
|
||||
import { errors } from 'undici';
|
||||
import {
|
||||
McpClientsService,
|
||||
isRetryableConnectError,
|
||||
} from './mcp-clients.service';
|
||||
|
||||
/**
|
||||
* #489 — external-MCP in-run transport recovery.
|
||||
*
|
||||
* The transport-error classification + retry gate are exercised against the REAL
|
||||
* undici error CLASSES prod throws (`errors.SocketError` / `errors.BodyTimeoutError`,
|
||||
* carrying the true `UND_ERR_*` codes and class names), wrapped EXACTLY as undici's
|
||||
* `fetch` wraps them — a `TypeError('fetch failed'|'terminated')` whose `.cause` is
|
||||
* the undici error. These are the real classes, not hand-rolled `{code:'...'}`
|
||||
* mocks: constructing the genuine class is what makes this a faithful test of the
|
||||
* prod predicate (epic root-cause #4 — a mock-shaped predicate would leave the
|
||||
* evict/retry path silently dead in production while CI stays green). We construct
|
||||
* rather than drive a live fetch because Jest's environment degrades the live-fetch
|
||||
* error to a generic `Error` cause (no undici code), which would NOT be the prod
|
||||
* shape.
|
||||
*/
|
||||
|
||||
/** A REAL undici socket reset, wrapped as fetch wraps it. */
|
||||
function realSocketResetError(): unknown {
|
||||
const err = new TypeError('fetch failed');
|
||||
(err as { cause?: unknown }).cause = new errors.SocketError('other side closed');
|
||||
return err;
|
||||
}
|
||||
|
||||
/** A REAL undici body timeout, wrapped as fetch wraps it. */
|
||||
function realBodyTimeoutError(): unknown {
|
||||
const err = new TypeError('terminated');
|
||||
(err as { cause?: unknown }).cause = new errors.BodyTimeoutError();
|
||||
return err;
|
||||
}
|
||||
|
||||
type FakeServer = {
|
||||
id: string;
|
||||
name: string;
|
||||
transport: 'http' | 'sse';
|
||||
url: string;
|
||||
headersEnc: string | null;
|
||||
toolAllowlist: string[] | null;
|
||||
instructions: string | null;
|
||||
};
|
||||
|
||||
const server = (over: Partial<FakeServer> = {}): FakeServer => ({
|
||||
id: 's1',
|
||||
name: 'srv',
|
||||
transport: 'http',
|
||||
url: 'http://example.test/mcp',
|
||||
headersEnc: null,
|
||||
toolAllowlist: null,
|
||||
instructions: null,
|
||||
...over,
|
||||
});
|
||||
|
||||
function buildService(servers: FakeServer[], trusted = false) {
|
||||
const repo = { listEnabled: jest.fn().mockResolvedValue(servers) };
|
||||
const service = new McpClientsService(repo as never, {} as never);
|
||||
// Seed a DETERMINISTIC write-class map so the retry gate is controlled here
|
||||
// (the production map loads from @docmost/mcp via a dynamic ESM import). getPage
|
||||
// is a read, patchNode is a write — the real classifications.
|
||||
(
|
||||
service as unknown as { writeClassMapPromise: Promise<unknown> }
|
||||
).writeClassMapPromise = Promise.resolve({
|
||||
getPage: 'readOnly',
|
||||
patchNode: 'write',
|
||||
});
|
||||
// The service only APPLIES that map to a TRUSTED internal Docmost server
|
||||
// (isInternalDocmostServer, really false for every third-party row). A retry
|
||||
// test needs a trusted server to exercise the readOnly-retry path at all, so it
|
||||
// passes trusted=true to model a Docmost-origin server; the third-party
|
||||
// double-apply test leaves it at the real value (false).
|
||||
if (trusted) {
|
||||
jest
|
||||
.spyOn(
|
||||
service as unknown as {
|
||||
isInternalDocmostServer: (s: FakeServer) => boolean;
|
||||
},
|
||||
'isInternalDocmostServer',
|
||||
)
|
||||
.mockReturnValue(true);
|
||||
}
|
||||
return { service, repo };
|
||||
}
|
||||
|
||||
/** Spy the private `connect` so each call yields a controlled fake client whose
|
||||
* single tool's execute is the supplied function. Returns the connect spy. */
|
||||
function stubConnect(
|
||||
service: McpClientsService,
|
||||
toolName: string,
|
||||
execs: Array<(...a: unknown[]) => Promise<unknown>>,
|
||||
) {
|
||||
let n = 0;
|
||||
return jest
|
||||
.spyOn(
|
||||
service as unknown as { connect: (s: FakeServer) => Promise<unknown> },
|
||||
'connect',
|
||||
)
|
||||
.mockImplementation(async () => {
|
||||
const exec = execs[Math.min(n, execs.length - 1)];
|
||||
n += 1;
|
||||
return {
|
||||
tools: async () => ({ [toolName]: { description: 'x', execute: exec } }),
|
||||
close: jest.fn().mockResolvedValue(undefined),
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
const opts = (abortSignal?: AbortSignal) =>
|
||||
({ toolCallId: 't', messages: [], abortSignal }) as never;
|
||||
|
||||
describe('isRetryableConnectError (#489, REAL error shapes)', () => {
|
||||
it('classifies a real undici socket reset and body timeout as retryable', async () => {
|
||||
const socketErr = await realSocketResetError();
|
||||
const bodyErr = await realBodyTimeoutError();
|
||||
expect(isRetryableConnectError(socketErr)).toBe(true);
|
||||
expect(isRetryableConnectError(bodyErr)).toBe(true);
|
||||
// Unwraps a wrapped cause chain (e.g. an MCPClientError around the socket err).
|
||||
const wrapped = new Error('mcp call failed');
|
||||
(wrapped as { cause?: unknown }).cause = socketErr;
|
||||
expect(isRetryableConnectError(wrapped)).toBe(true);
|
||||
});
|
||||
|
||||
it('does NOT classify an application-level error as a transport break', () => {
|
||||
expect(isRetryableConnectError(new Error('validation failed'))).toBe(false);
|
||||
expect(isRetryableConnectError({ name: 'HttpError', status: 400 })).toBe(false);
|
||||
expect(isRetryableConnectError(undefined)).toBe(false);
|
||||
expect(isRetryableConnectError('boom')).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('McpClientsService in-run transport recovery (#489)', () => {
|
||||
afterEach(() => jest.restoreAllMocks());
|
||||
|
||||
it('a readOnly tool whose transport breaks reconnects and retries WITHIN the same run', async () => {
|
||||
const realErr = await realSocketResetError();
|
||||
const { service } = buildService([server()], true);
|
||||
const first = jest.fn().mockRejectedValue(realErr);
|
||||
const second = jest.fn().mockResolvedValue({ ok: true });
|
||||
const connectSpy = stubConnect(service, 'getPage', [first, second]);
|
||||
|
||||
const toolset = await service.toolsFor('ws-1');
|
||||
const tool = toolset.tools['srv_getPage'];
|
||||
const result = await (tool.execute as (a: unknown, o: unknown) => Promise<unknown>)(
|
||||
{ pageId: 'p' },
|
||||
opts(),
|
||||
);
|
||||
|
||||
// The repeat call within the run got a LIVE client and succeeded.
|
||||
expect(result).toEqual({ ok: true });
|
||||
expect(first).toHaveBeenCalledTimes(1);
|
||||
expect(second).toHaveBeenCalledTimes(1);
|
||||
// Exactly one reconnect was minted (initial build connect + one recovery).
|
||||
expect(connectSpy).toHaveBeenCalledTimes(2);
|
||||
// The run accumulated BOTH leases (old + reconnected) — released together at end.
|
||||
expect(toolset.clients).toHaveLength(2);
|
||||
await Promise.all(toolset.clients.map((c) => c.close()));
|
||||
});
|
||||
|
||||
it('a WRITE tool does NOT auto-retry on a transport error (indeterminate)', async () => {
|
||||
const realErr = await realSocketResetError();
|
||||
const { service } = buildService([server()], true);
|
||||
const exec = jest.fn().mockRejectedValue(realErr);
|
||||
const connectSpy = stubConnect(service, 'patchNode', [exec]);
|
||||
|
||||
const toolset = await service.toolsFor('ws-2');
|
||||
const tool = toolset.tools['srv_patchNode'];
|
||||
await expect(
|
||||
(tool.execute as (a: unknown, o: unknown) => Promise<unknown>)(
|
||||
{ pageId: 'p' },
|
||||
opts(),
|
||||
),
|
||||
).rejects.toThrow(/MAY have already applied/);
|
||||
|
||||
// Called exactly once — NO blind retry (avoids double-apply, the #435 class).
|
||||
expect(exec).toHaveBeenCalledTimes(1);
|
||||
// No fresh connection was minted for a write.
|
||||
expect(connectSpy).toHaveBeenCalledTimes(1);
|
||||
await Promise.all(toolset.clients.map((c) => c.close()));
|
||||
});
|
||||
|
||||
it('does NOT retry (or reconnect) after the run is aborted (Stop)', async () => {
|
||||
const realErr = await realSocketResetError();
|
||||
const { service } = buildService([server()], true);
|
||||
const controller = new AbortController();
|
||||
// The transport error arrives, but the run was Stopped in the same tick.
|
||||
const first = jest.fn().mockImplementation(async () => {
|
||||
controller.abort();
|
||||
throw realErr;
|
||||
});
|
||||
const second = jest.fn().mockResolvedValue({ ok: true });
|
||||
const connectSpy = stubConnect(service, 'getPage', [first, second]);
|
||||
|
||||
const toolset = await service.toolsFor('ws-3');
|
||||
const tool = toolset.tools['srv_getPage'];
|
||||
await expect(
|
||||
(tool.execute as (a: unknown, o: unknown) => Promise<unknown>)(
|
||||
{ pageId: 'p' },
|
||||
opts(controller.signal),
|
||||
),
|
||||
).rejects.toBeDefined();
|
||||
|
||||
// getPage IS readOnly, but the Stop blocks the retry — no second call, no mint.
|
||||
expect(second).not.toHaveBeenCalled();
|
||||
expect(connectSpy).toHaveBeenCalledTimes(1);
|
||||
await Promise.all(toolset.clients.map((c) => c.close()));
|
||||
});
|
||||
|
||||
it('an app-level (non-transport) tool error is surfaced verbatim, never retried', async () => {
|
||||
const { service } = buildService([server()], true);
|
||||
const appErr = new Error('tool says: bad input');
|
||||
const exec = jest.fn().mockRejectedValue(appErr);
|
||||
const connectSpy = stubConnect(service, 'getPage', [exec]);
|
||||
|
||||
const toolset = await service.toolsFor('ws-4');
|
||||
const tool = toolset.tools['srv_getPage'];
|
||||
await expect(
|
||||
(tool.execute as (a: unknown, o: unknown) => Promise<unknown>)(
|
||||
{ pageId: 'p' },
|
||||
opts(),
|
||||
),
|
||||
).rejects.toThrow('tool says: bad input');
|
||||
expect(exec).toHaveBeenCalledTimes(1);
|
||||
expect(connectSpy).toHaveBeenCalledTimes(1); // no reconnect for an app error
|
||||
await Promise.all(toolset.clients.map((c) => c.close()));
|
||||
});
|
||||
|
||||
// #489 (review, MEDIUM) — the Docmost write-class map keys by DOCMOST tool
|
||||
// names; a THIRD-PARTY server may name a WRITE tool `getPage` (a Docmost read
|
||||
// name). It must NOT inherit readOnly and must NOT auto-retry on a transport
|
||||
// error — a blind retry of that write is a double-apply (the #435 class). Here
|
||||
// the server is UNTRUSTED (buildService default, isInternalDocmostServer=false),
|
||||
// so the map is not applied and `getPage` classifies as a write.
|
||||
//
|
||||
// MUTATION-VERIFY: forcing the server "trusted" (buildService(..., true)) makes
|
||||
// `getPage` inherit readOnly -> it WOULD reconnect+retry (connect twice) and the
|
||||
// assertions below fail — i.e. removing the trust scope re-opens the bug.
|
||||
it('a THIRD-PARTY WRITE tool named like a Docmost read does NOT auto-retry (no double-apply)', async () => {
|
||||
const realErr = await realSocketResetError();
|
||||
// Untrusted: default trusted=false — a real third-party server.
|
||||
const { service } = buildService([server()]);
|
||||
const exec = jest.fn().mockRejectedValue(realErr);
|
||||
const connectSpy = stubConnect(service, 'getPage', [exec, exec]);
|
||||
|
||||
const toolset = await service.toolsFor('ws-5');
|
||||
const tool = toolset.tools['srv_getPage'];
|
||||
await expect(
|
||||
(tool.execute as (a: unknown, o: unknown) => Promise<unknown>)(
|
||||
{ pageId: 'p' },
|
||||
opts(),
|
||||
),
|
||||
).rejects.toThrow(/MAY have already applied/);
|
||||
|
||||
// Exactly one call, NO reconnect — the name collision granted no readOnly-retry.
|
||||
expect(exec).toHaveBeenCalledTimes(1);
|
||||
expect(connectSpy).toHaveBeenCalledTimes(1);
|
||||
await Promise.all(toolset.clients.map((c) => c.close()));
|
||||
});
|
||||
});
|
||||
@@ -106,8 +106,11 @@ describe('McpClientsService.decryptHeaders', () => {
|
||||
|
||||
describe('McpClientsService.guardedFetch (SSRF per-request guard)', () => {
|
||||
// The bound guardedFetch closure lives on the instance as a private field.
|
||||
// #489 split it into per-transport HTTP/SSE bindings (they differ only in the
|
||||
// dispatcher's bodyTimeout); the SSRF guard is identical, so testing the HTTP
|
||||
// one is sufficient.
|
||||
const guardedFetchOf = (service: McpClientsService) =>
|
||||
(service as unknown as { guardedFetch: typeof fetch }).guardedFetch;
|
||||
(service as unknown as { guardedFetchHttp: typeof fetch }).guardedFetchHttp;
|
||||
|
||||
let fetchSpy: jest.SpiedFunction<typeof fetch>;
|
||||
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import { isIP } from 'node:net';
|
||||
import { lookup as dnsLookup, type LookupAddress } from 'node:dns';
|
||||
import { pathToFileURL } from 'node:url';
|
||||
import { Injectable, Logger } from '@nestjs/common';
|
||||
import { type Tool, type ToolCallOptions } from 'ai';
|
||||
import { createMCPClient } from '@ai-sdk/mcp';
|
||||
@@ -10,9 +11,29 @@ import {
|
||||
streamingDispatcherOptions,
|
||||
mcpStreamTimeoutMs,
|
||||
mcpCallTimeoutMs,
|
||||
mcpSseBodyTimeoutMs,
|
||||
} from '../../../integrations/ai/ai-streaming-fetch';
|
||||
import { SecretBoxService } from '../../../integrations/crypto/secret-box';
|
||||
import { isUrlAllowed, isIpAllowed } from './ssrf-guard';
|
||||
// TYPE-ONLY (erased at compile): @docmost/mcp is ESM-only and cannot be a runtime
|
||||
// `require()` from this commonjs module (same constraint as docmost-client.loader).
|
||||
// The write-class MAP is loaded lazily via the dynamic-import trick below.
|
||||
import type { ToolWriteClass } from '@docmost/mcp';
|
||||
|
||||
// TS(commonjs) downlevels a literal `import()` to `require()`, which cannot load
|
||||
// the ESM-only @docmost/mcp. Indirect through Function so the real dynamic
|
||||
// `import()` survives compilation (same trick as docmost-client.loader.ts).
|
||||
const esmImport = new Function(
|
||||
'specifier',
|
||||
'return import(specifier)',
|
||||
) as (specifier: string) => Promise<unknown>;
|
||||
|
||||
/** Local read-only predicate — avoids a value import of the ESM-only package.
|
||||
* Only a pure read is retry-safe after a transport break (a write is
|
||||
* indeterminate). Kept in lockstep with @docmost/mcp's isRetryableWriteClass. */
|
||||
function isReadOnlyWriteClass(writeClass: ToolWriteClass | undefined): boolean {
|
||||
return writeClass === 'readOnly';
|
||||
}
|
||||
|
||||
/** A closable external MCP client handle. */
|
||||
export interface Closable {
|
||||
@@ -81,12 +102,52 @@ const MAX_TOOL_NAME_LENGTH = 64;
|
||||
* close until the turn releases it, so a TTL expiry mid-turn never closes a
|
||||
* client a stream is still executing against.
|
||||
*/
|
||||
/**
|
||||
* Where a merged (namespaced) tool came from, so the per-run recovery wrapper
|
||||
* (#489) can, on a transport error, reconnect THAT server and re-resolve the SAME
|
||||
* underlying tool by its raw name. `writeClass` gates the single auto-retry (a
|
||||
* read is retry-safe; a write is indeterminate). `serverIndex` indexes the
|
||||
* entry's `servers` array (which server config to reconnect).
|
||||
*/
|
||||
interface ToolProvenance {
|
||||
serverIndex: number;
|
||||
rawName: string;
|
||||
writeClass: ToolWriteClass | undefined;
|
||||
}
|
||||
|
||||
/** A live reconnected server (its fresh client + raw call-timeout-wrapped tools). */
|
||||
interface RecoveredServerState {
|
||||
client: McpClient;
|
||||
tools: Record<string, Tool>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Per-run, per-server recovery binding (#489). `current` is the server's LIVE
|
||||
* target for this run: `null` means "use the ORIGINAL cached client/template";
|
||||
* a non-null value is a reconnected throwaway client all this server's tools now
|
||||
* call. `reconnecting` dedupes concurrent reconnects so only ONE fresh client is
|
||||
* minted per death (a losing concurrent call awaits it and retries on the SAME
|
||||
* new client — the CAS-by-identity rule).
|
||||
*/
|
||||
interface ServerBinding {
|
||||
current: RecoveredServerState | null;
|
||||
reconnecting?: Promise<RecoveredServerState>;
|
||||
}
|
||||
|
||||
interface CacheEntry {
|
||||
tools: Record<string, Tool>;
|
||||
clients: McpClient[];
|
||||
outcomes: ServerOutcome[];
|
||||
/** Prompt guidance for qualifying servers (see McpServerInstruction). */
|
||||
instructions: McpServerInstruction[];
|
||||
/**
|
||||
* The enabled server configs used to build this entry (#489), so the per-run
|
||||
* recovery wrapper can reconnect a specific server by index. Parallel to the
|
||||
* indices referenced by {@link toolMeta}.
|
||||
*/
|
||||
servers: AiMcpServer[];
|
||||
/** merged-tool-key -> provenance (#489), for the per-run recovery wrapper. */
|
||||
toolMeta: Record<string, ToolProvenance>;
|
||||
expiresAt: number;
|
||||
/** Active leases (turns currently using these clients). */
|
||||
refCount: number;
|
||||
@@ -120,20 +181,82 @@ export class McpClientsService {
|
||||
*/
|
||||
private readonly cache = new Map<string, Promise<CacheEntry>>();
|
||||
/**
|
||||
* A single shared SSRF-pinned dispatcher for ALL outbound external-MCP fetches.
|
||||
* Its custom connect.lookup runs per connection, so one instance safely guards
|
||||
* every server's connections (we never connect to an unvalidated IP).
|
||||
* SSRF-pinned dispatchers for outbound external-MCP fetches. Both use the SAME
|
||||
* custom connect.lookup (so every connection is IP-validated), but carry a
|
||||
* DIFFERENT `bodyTimeout` (#489): the HTTP (streamable) transport opens a fresh
|
||||
* request per call, so it keeps the tight silence timeout; the SSE transport
|
||||
* holds ONE long-lived body open across many calls, so a >1-min idle BETWEEN
|
||||
* calls is LEGITIMATE and must not break the socket — it gets a much larger
|
||||
* bodyTimeout. (headersTimeout stays tight on both.)
|
||||
*/
|
||||
private readonly dispatcher: Dispatcher = buildPinnedDispatcher();
|
||||
/** guardedFetch bound to the pinned dispatcher; reused by every transport. */
|
||||
private readonly guardedFetch: typeof fetch = (input, init) =>
|
||||
guardedFetch(this.dispatcher, input, init);
|
||||
private readonly dispatcherHttp: Dispatcher = buildPinnedDispatcher(
|
||||
mcpStreamTimeoutMs(),
|
||||
);
|
||||
private readonly dispatcherSse: Dispatcher = buildPinnedDispatcher(
|
||||
mcpSseBodyTimeoutMs(),
|
||||
);
|
||||
/** guardedFetch bound to each dispatcher; picked by transport type in connect(). */
|
||||
private readonly guardedFetchHttp: typeof fetch = (input, init) =>
|
||||
guardedFetch(this.dispatcherHttp, input, init);
|
||||
private readonly guardedFetchSse: typeof fetch = (input, init) =>
|
||||
guardedFetch(this.dispatcherSse, input, init);
|
||||
|
||||
/**
|
||||
* Memoized write-class map (#489), loaded lazily from @docmost/mcp via the
|
||||
* dynamic-import trick. Keyed by tool name (=== mcpName). A tool NOT in the map
|
||||
* (any third-party external MCP tool) classifies as `undefined` -> treated as a
|
||||
* write by the retry gate (the safe default: never blind-retry an unknown tool).
|
||||
* On any load failure the map is `{}` (every tool -> no auto-retry), so a
|
||||
* missing/older @docmost/mcp build only DISABLES retries, never mis-retries.
|
||||
*/
|
||||
private writeClassMapPromise: Promise<Record<string, ToolWriteClass>> | null =
|
||||
null;
|
||||
|
||||
constructor(
|
||||
private readonly repo: AiMcpServerRepo,
|
||||
private readonly secretBox: SecretBoxService,
|
||||
) {}
|
||||
|
||||
/**
|
||||
* Whether an external MCP server is the TRUSTED internal Docmost MCP server —
|
||||
* the only server whose tools may be classified by the Docmost write-class map
|
||||
* (#489 review). Today this is ALWAYS false: every `ai_mcp_servers` row is an
|
||||
* admin-configured THIRD-PARTY endpoint (there is no builtin/self flag, sentinel
|
||||
* URL, or synthetic server in this path — Docmost's OWN tools are exposed via the
|
||||
* separate in-app tools path, never through this external-MCP client). So no
|
||||
* third-party tool can inherit `readOnly` by a name collision with a Docmost read
|
||||
* tool, and none is ever auto-retried on a transport error (which would risk a
|
||||
* double-apply — the #435 class). Flip this (an explicit `kind`/`isBuiltin`
|
||||
* column, or a configured self-MCP URL) if a trusted internal server is ever
|
||||
* introduced. A method (not a free function) so it is a single, mockable seam.
|
||||
*/
|
||||
private isInternalDocmostServer(_server: AiMcpServer): boolean {
|
||||
return false;
|
||||
}
|
||||
|
||||
/** Lazily load + memoize the shared write-class map (see the field doc). */
|
||||
private getWriteClassMap(): Promise<Record<string, ToolWriteClass>> {
|
||||
if (!this.writeClassMapPromise) {
|
||||
this.writeClassMapPromise = (async () => {
|
||||
try {
|
||||
const entry = require.resolve('@docmost/mcp');
|
||||
const mod = (await esmImport(pathToFileURL(entry).href)) as {
|
||||
SHARED_TOOL_WRITE_CLASS?: Record<string, ToolWriteClass>;
|
||||
};
|
||||
return mod.SHARED_TOOL_WRITE_CLASS ?? {};
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`Could not load MCP write-class map (auto-retry disabled): ${shortError(
|
||||
err,
|
||||
)}`,
|
||||
);
|
||||
return {};
|
||||
}
|
||||
})();
|
||||
}
|
||||
return this.writeClassMapPromise;
|
||||
}
|
||||
|
||||
/**
|
||||
* Build (or reuse a cached) external toolset for a workspace. Returns the
|
||||
* merged tools, the open client handles to release, and per-server outcomes.
|
||||
@@ -162,11 +285,37 @@ export class McpClientsService {
|
||||
}
|
||||
},
|
||||
};
|
||||
// One release handle drives the whole leased entry; closing it releases all
|
||||
// underlying clients together (they share the same lease lifecycle).
|
||||
|
||||
// #489: the run accumulates a SET of leases — the primary cache lease PLUS any
|
||||
// throwaway client minted by an in-run transport-recovery reconnect. They are
|
||||
// NEVER released mid-run (releasing a swapped-out client while a concurrent
|
||||
// in-flight call still holds it would INDUCE a second failure); the caller
|
||||
// releases the WHOLE set together at turn-end. A recovery reconnect pushes its
|
||||
// lease onto this live array, which the consumer closes over.
|
||||
const leaseSet: Closable[] = [release];
|
||||
|
||||
// #489: per-RUN transport-recovery binding, one per server, SHARED by all of
|
||||
// that server's tools so a swap by one call is seen by the next (CAS by
|
||||
// identity). Kept per-run (here, not in the cached entry) because the binding
|
||||
// + lease-set state is per-run.
|
||||
const bindings = new Map<number, ServerBinding>();
|
||||
const capMs = mcpCallTimeoutMs();
|
||||
|
||||
// Wrap each cached tool with the recovery layer. On a transport error a
|
||||
// declared readOnly tool reconnects its server and retries ONCE; a write is
|
||||
// never blind-retried (indeterminate — may have applied before the reset). A
|
||||
// tool without provenance (a minimal stub entry in a test) passes through raw.
|
||||
const tools: Record<string, Tool> = {};
|
||||
for (const [key, tool] of Object.entries(entry.tools)) {
|
||||
const meta = entry.toolMeta?.[key];
|
||||
tools[key] = meta
|
||||
? this.wrapWithTransportRecovery(entry, meta, tool, leaseSet, bindings, capMs)
|
||||
: tool;
|
||||
}
|
||||
|
||||
return {
|
||||
tools: entry.tools,
|
||||
clients: [release],
|
||||
tools,
|
||||
clients: leaseSet,
|
||||
outcomes: entry.outcomes,
|
||||
instructions: entry.instructions,
|
||||
};
|
||||
@@ -254,6 +403,16 @@ export class McpClientsService {
|
||||
// Per-call total wall-clock cap, read once for this build (env-overridable).
|
||||
const callTimeoutMs = mcpCallTimeoutMs();
|
||||
const instructions: McpServerInstruction[] = [];
|
||||
// merged-key -> provenance for the per-run recovery wrapper (#489).
|
||||
const toolMeta: Record<string, ToolProvenance> = {};
|
||||
// Shared Docmost write-class map (#489) — classifies a tool by its raw name.
|
||||
// Loaded ONLY when at least one server is a TRUSTED internal Docmost server
|
||||
// (see isInternalDocmostServer): for third-party servers the map is never
|
||||
// applied (a name collision must not grant readOnly-retry), so we skip the
|
||||
// dynamic ESM load entirely in that (currently universal) case.
|
||||
const writeClassMap = servers.some((s) => this.isInternalDocmostServer(s))
|
||||
? await this.getWriteClassMap()
|
||||
: null;
|
||||
|
||||
// Per-server connect+tools result, still tagged with its server so the merge
|
||||
// below can be applied in the SAME order as `servers` (see the parallel note).
|
||||
@@ -327,11 +486,23 @@ export class McpClientsService {
|
||||
// against names already merged from earlier servers, so no external
|
||||
// tool is silently overwritten on collision. The returned count drives
|
||||
// whether this server's prompt guidance is included (≥1 tool merged).
|
||||
// #489 (review): the Docmost write-class map keys by DOCMOST tool names and
|
||||
// may ONLY be trusted for a server KNOWN to be the internal Docmost MCP
|
||||
// server. Every row here is an admin-configured THIRD-PARTY endpoint, so a
|
||||
// third-party WRITE tool that happens to be named like a Docmost read
|
||||
// (getPage, listPages, ...) must NOT inherit readOnly — that would auto-retry
|
||||
// a mutation on a transport error (double-apply, the #435 class). Gate the
|
||||
// map on the trust check; untrusted servers get writeClass=undefined -> the
|
||||
// recovery wrapper treats them as writes and never auto-retries.
|
||||
const trustWriteClass = this.isInternalDocmostServer(server);
|
||||
const merged = this.mergeNamespaced(
|
||||
tools,
|
||||
result.guarded,
|
||||
server.name,
|
||||
server.id,
|
||||
toolMeta,
|
||||
i,
|
||||
trustWriteClass ? writeClassMap : null,
|
||||
);
|
||||
outcomes.push({ name: server.name, ok: true });
|
||||
// Include this server's guidance ONLY when it actually contributed at
|
||||
@@ -353,6 +524,8 @@ export class McpClientsService {
|
||||
clients,
|
||||
outcomes,
|
||||
instructions,
|
||||
servers,
|
||||
toolMeta,
|
||||
expiresAt: Date.now() + CACHE_TTL_MS,
|
||||
refCount: 0,
|
||||
evicted: false,
|
||||
@@ -379,18 +552,33 @@ export class McpClientsService {
|
||||
picked: Record<string, Tool>,
|
||||
serverName: string,
|
||||
serverId: string,
|
||||
toolMeta: Record<string, ToolProvenance>,
|
||||
serverIndex: number,
|
||||
// The Docmost write-class map, or `null` for an UNTRUSTED (third-party)
|
||||
// server whose tools must all default to write (never auto-retried).
|
||||
writeClassMap: Record<string, ToolWriteClass> | null,
|
||||
): { count: number; prefix: string } {
|
||||
let count = 0;
|
||||
for (const [name, tool] of Object.entries(namespace(picked, serverName))) {
|
||||
let key = name;
|
||||
for (const { full, raw, tool } of namespace(picked, serverName)) {
|
||||
let key = full;
|
||||
if (key in target) {
|
||||
const original = key;
|
||||
key = disambiguate(name, serverId, (candidate) => candidate in target);
|
||||
key = disambiguate(full, serverId, (candidate) => candidate in target);
|
||||
this.logger.debug(
|
||||
`External MCP tool name "${original}" collided; renamed to "${key}"`,
|
||||
);
|
||||
}
|
||||
target[key] = tool;
|
||||
// Record provenance so the per-run recovery wrapper (#489) can reconnect
|
||||
// this tool's server and re-resolve it by its raw name. writeClass is set
|
||||
// ONLY from a TRUSTED (internal-Docmost) map; for a third-party server the
|
||||
// map is null -> writeClass stays undefined -> the wrapper treats the tool
|
||||
// as a write and never auto-retries it (no double-apply on name collision).
|
||||
toolMeta[key] = {
|
||||
serverIndex,
|
||||
rawName: raw,
|
||||
writeClass: writeClassMap ? writeClassMap[raw] : undefined,
|
||||
};
|
||||
count += 1;
|
||||
}
|
||||
return { count, prefix: namespacePrefix(serverName) };
|
||||
@@ -424,7 +612,10 @@ export class McpClientsService {
|
||||
// Defense in depth: re-validate the actual request host on EVERY fetch
|
||||
// AND pin the socket to a validated IP via the dispatcher's connect
|
||||
// lookup, closing the DNS-rebinding TOCTOU between check and connect.
|
||||
fetch: this.guardedFetch,
|
||||
// #489: the SSE transport uses the raised-bodyTimeout dispatcher (idle
|
||||
// between calls is legit); HTTP uses the tight one.
|
||||
fetch:
|
||||
transportType === 'sse' ? this.guardedFetchSse : this.guardedFetchHttp,
|
||||
},
|
||||
})) as unknown as McpClient;
|
||||
return client;
|
||||
@@ -505,6 +696,176 @@ export class McpClientsService {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Wrap one merged external tool with the per-run transport-recovery layer (#489).
|
||||
*
|
||||
* attempt 1 runs on the server's CURRENT binding (the cached client, or a client
|
||||
* a sibling tool already reconnected this run). On a REAL transport error
|
||||
* (undici/@ai-sdk socket/body-timeout shapes — {@link isRetryableConnectError},
|
||||
* NOT a mock) and ONLY for a declared readOnly tool, it reconnects the server
|
||||
* and retries EXACTLY ONCE on the fresh client; a write is surfaced as an
|
||||
* indeterminate error (it may have applied before the reset — never
|
||||
* blind-retried). A single per-call cap bounds BOTH attempts + the reconnect,
|
||||
* and the run's abort signal is checked before the retry AND before minting a
|
||||
* fresh connection (no connection is opened for a stopped run).
|
||||
*/
|
||||
private wrapWithTransportRecovery(
|
||||
entry: CacheEntry,
|
||||
meta: ToolProvenance,
|
||||
template: Tool,
|
||||
leaseSet: Closable[],
|
||||
bindings: Map<number, ServerBinding>,
|
||||
capMs: number,
|
||||
): Tool {
|
||||
const original = template.execute;
|
||||
if (typeof original !== 'function') return template;
|
||||
const service = this;
|
||||
const { serverIndex, rawName, writeClass } = meta;
|
||||
|
||||
let binding = bindings.get(serverIndex);
|
||||
if (!binding) {
|
||||
binding = { current: null };
|
||||
bindings.set(serverIndex, binding);
|
||||
}
|
||||
const boundBinding = binding;
|
||||
|
||||
const execute = async (args: unknown, options: ToolCallOptions) => {
|
||||
// The per-call cap governs the WHOLE sequence (attempt1 + reconnect +
|
||||
// attempt2). Compose it with the run's abort signal so a Stop or the cap
|
||||
// ends any awaited call — @ai-sdk/mcp does not settle on abort, so we RACE.
|
||||
const capController = new AbortController();
|
||||
const capTimer = setTimeout(() => {
|
||||
capController.abort(new Error(`MCP tool call timed out after ${capMs}ms`));
|
||||
}, capMs);
|
||||
capTimer.unref?.();
|
||||
const runSignal = options?.abortSignal;
|
||||
const composed = runSignal
|
||||
? AbortSignal.any([runSignal, capController.signal])
|
||||
: capController.signal;
|
||||
const stopped = () => runSignal?.aborted === true || capController.signal.aborted;
|
||||
|
||||
const callOn = async (
|
||||
exec: NonNullable<Tool['execute']>,
|
||||
): Promise<unknown> => {
|
||||
const aborted = new Promise<never>((_, reject) => {
|
||||
const fail = () => reject(abortReason(composed));
|
||||
if (composed.aborted) fail();
|
||||
else composed.addEventListener('abort', fail, { once: true });
|
||||
});
|
||||
return Promise.race([exec(args, { ...options, abortSignal: composed }), aborted]);
|
||||
};
|
||||
|
||||
const execFor = (
|
||||
state: RecoveredServerState | null,
|
||||
): NonNullable<Tool['execute']> | undefined =>
|
||||
state ? (state.tools[rawName]?.execute as NonNullable<Tool['execute']>) : original;
|
||||
|
||||
try {
|
||||
// Snapshot the target BEFORE the call so a swap by a concurrent call is
|
||||
// detected by identity in the catch.
|
||||
const attemptState = boundBinding.current;
|
||||
const attemptExec = execFor(attemptState);
|
||||
if (typeof attemptExec !== 'function') {
|
||||
throw new Error(`external MCP tool "${rawName}" is not callable`);
|
||||
}
|
||||
try {
|
||||
return await callOn(attemptExec);
|
||||
} catch (err) {
|
||||
// Never retry on a Stop or an exhausted cap.
|
||||
if (stopped()) throw err;
|
||||
// Only a genuine transport break is a recovery candidate.
|
||||
if (!isRetryableConnectError(err)) throw err;
|
||||
// A write tool is INDETERMINATE on a transport error (may have applied
|
||||
// before the reset) — surface that; do NOT auto-retry (double-apply is
|
||||
// the #435 incident class).
|
||||
if (!isReadOnlyWriteClass(writeClass)) {
|
||||
throw new Error(
|
||||
`external MCP tool "${rawName}" hit a transport error and MAY have already ` +
|
||||
`applied on the server — not retried automatically; verify state before ` +
|
||||
`retrying. (${shortError(err)})`,
|
||||
);
|
||||
}
|
||||
// Abort check BEFORE minting a fresh connection (no socket for a
|
||||
// stopped run). LIMITATION (#489, LOW): the reconnect's own connect is
|
||||
// bounded by CONNECT_TIMEOUT_MS but does NOT itself observe `composed`,
|
||||
// so a Stop that lands DURING the handshake is only honored at the next
|
||||
// `stopped()` gate (before the retry) — a bounded ≤5s late-abort window;
|
||||
// the throwaway client is closed at turn-end regardless. Threading
|
||||
// `composed` into the SHARED (CAS-deduped) reconnect is deliberately
|
||||
// avoided: it would let the first caller's abort tear down a reconnect a
|
||||
// concurrent still-live caller depends on.
|
||||
if (stopped()) throw err;
|
||||
// CAS-swap by IDENTITY: mint+swap only if nobody swapped since this
|
||||
// call's snapshot; a losing concurrent call awaits the same reconnect
|
||||
// and retries on the SAME fresh client.
|
||||
let target: RecoveredServerState;
|
||||
if (boundBinding.current === attemptState) {
|
||||
if (!boundBinding.reconnecting) {
|
||||
boundBinding.reconnecting = (async () => {
|
||||
const server = entry.servers[serverIndex];
|
||||
const fresh = await service.reconnectServer(server, capMs);
|
||||
leaseSet.push(fresh.lease); // accumulate; released at turn-end
|
||||
boundBinding.current = fresh.state;
|
||||
return fresh.state;
|
||||
})();
|
||||
// Clear the in-flight marker once it settles (success or failure) so
|
||||
// a LATER death of the new client can reconnect again.
|
||||
void boundBinding.reconnecting.then(
|
||||
() => (boundBinding.reconnecting = undefined),
|
||||
() => (boundBinding.reconnecting = undefined),
|
||||
);
|
||||
}
|
||||
target = await boundBinding.reconnecting;
|
||||
} else {
|
||||
target = boundBinding.current as RecoveredServerState;
|
||||
}
|
||||
// Abort check BEFORE the retry.
|
||||
if (stopped()) throw err;
|
||||
const retryExec = execFor(target);
|
||||
if (typeof retryExec !== 'function') throw err;
|
||||
return await callOn(retryExec);
|
||||
}
|
||||
} finally {
|
||||
clearTimeout(capTimer);
|
||||
}
|
||||
};
|
||||
return { ...template, execute } as unknown as Tool;
|
||||
}
|
||||
|
||||
/**
|
||||
* Reconnect ONE server for an in-run recovery (#489): open a fresh client and
|
||||
* list+wrap its tools. The throwaway client is NOT cached — it is owned by the
|
||||
* RUN via the returned lease (closed at turn-end), independent of the shared
|
||||
* cache entry (whose TTL rebuild heals future turns). On a failure the fresh
|
||||
* client is closed so its socket never leaks.
|
||||
*/
|
||||
private async reconnectServer(
|
||||
server: AiMcpServer,
|
||||
capMs: number,
|
||||
): Promise<{ state: RecoveredServerState; lease: Closable }> {
|
||||
const client = await this.connectWithTimeout(server, CONNECT_TIMEOUT_MS);
|
||||
let tools: Record<string, Tool>;
|
||||
try {
|
||||
const raw = await withTimeout(client.tools(), CONNECT_TIMEOUT_MS);
|
||||
const allow = server.toolAllowlist;
|
||||
const picked =
|
||||
Array.isArray(allow) && allow.length > 0 ? pick(raw, allow) : raw;
|
||||
tools = wrapToolsWithCallTimeout(picked, capMs);
|
||||
} catch (err) {
|
||||
void client.close().catch(() => undefined);
|
||||
throw err;
|
||||
}
|
||||
let released = false;
|
||||
const lease: Closable = {
|
||||
close: async () => {
|
||||
if (released) return;
|
||||
released = true;
|
||||
await client.close().catch(() => undefined);
|
||||
},
|
||||
};
|
||||
return { state: { client, tools }, lease };
|
||||
}
|
||||
|
||||
/** Mark an entry evicted; close its clients now if nothing is leasing them. */
|
||||
private evict(entry: CacheEntry): void {
|
||||
clearTimeout(entry.timer);
|
||||
@@ -554,22 +915,21 @@ export function validateResolvedAddresses(addrs: readonly LookupAddress[]): {
|
||||
* certificate validation still uses the real hostname (we never rewrite the URL
|
||||
* to an IP literal).
|
||||
*/
|
||||
function buildPinnedDispatcher(): Agent {
|
||||
// External-MCP traffic uses a DEDICATED, shorter silence timeout
|
||||
function buildPinnedDispatcher(bodyTimeoutMs: number): Agent {
|
||||
// External-MCP traffic uses a DEDICATED, shorter HEADERS silence timeout
|
||||
// (`AI_MCP_STREAM_TIMEOUT_MS`, default 1 min) — deliberately tighter than the
|
||||
// chat provider's 15-min `streamTimeoutMs()` — so a byte-silent/hung MCP
|
||||
// upstream is broken in ~1 min instead of 15. We keep the keep-alive options
|
||||
// from `streamingDispatcherOptions()` but OVERRIDE headers/body timeouts.
|
||||
// Accepted trade-off: a legitimately long but byte-silent single tool call,
|
||||
// and an SSE transport idling >1 min BETWEEN tool calls, are also cut here; the
|
||||
// per-call total cap (wrapToolsWithCallTimeout, `AI_MCP_CALL_TIMEOUT_MS`) is the
|
||||
// complementary guard for chatty-but-stuck calls that keep the socket warm yet
|
||||
// never return.
|
||||
const mcpSilenceMs = mcpStreamTimeoutMs();
|
||||
// from `streamingDispatcherOptions()` but OVERRIDE the timeouts. `bodyTimeout`
|
||||
// is passed in per-transport (#489): tight for HTTP (fresh request per call),
|
||||
// raised for SSE (one long-lived body across calls — idle BETWEEN calls is
|
||||
// legit). The per-call total cap (`AI_MCP_CALL_TIMEOUT_MS`) is the complementary
|
||||
// guard for chatty-but-stuck calls that keep the socket warm yet never return.
|
||||
const headersMs = mcpStreamTimeoutMs();
|
||||
return new Agent({
|
||||
...streamingDispatcherOptions(),
|
||||
headersTimeout: mcpSilenceMs,
|
||||
bodyTimeout: mcpSilenceMs,
|
||||
headersTimeout: headersMs,
|
||||
bodyTimeout: bodyTimeoutMs,
|
||||
connect: {
|
||||
lookup: (hostname, _options, callback) => {
|
||||
// Always resolve ALL addresses ourselves; do not trust the caller's
|
||||
@@ -669,18 +1029,22 @@ function pick(
|
||||
function namespace(
|
||||
tools: Record<string, Tool>,
|
||||
serverName: string,
|
||||
): Record<string, Tool> {
|
||||
): Array<{ full: string; raw: string; tool: Tool }> {
|
||||
const prefix = namespacePrefix(serverName);
|
||||
const out: Record<string, Tool> = {};
|
||||
const out: Array<{ full: string; raw: string; tool: Tool }> = [];
|
||||
const taken: Record<string, true> = {};
|
||||
for (const [name, t] of Object.entries(tools)) {
|
||||
const safe = sanitizeName(name);
|
||||
let full = capName(`${prefix}_${safe}`);
|
||||
// Duplicate names within ONE server can still collide after sanitize/
|
||||
// truncate — suffix-disambiguate so the second tool is not overwritten.
|
||||
if (full in out) {
|
||||
full = disambiguate(full, '', (candidate) => candidate in out);
|
||||
if (full in taken) {
|
||||
full = disambiguate(full, '', (candidate) => candidate in taken);
|
||||
}
|
||||
out[full] = t;
|
||||
taken[full] = true;
|
||||
// Keep the RAW (un-namespaced) name alongside the merged key so the per-run
|
||||
// recovery wrapper (#489) can re-resolve the same tool on a fresh client.
|
||||
out.push({ full, raw: name, tool: t });
|
||||
}
|
||||
return out;
|
||||
}
|
||||
@@ -804,6 +1168,69 @@ export function wrapToolWithCallTimeout(tool: Tool, ms: number): Tool {
|
||||
return { ...tool, execute } as unknown as Tool;
|
||||
}
|
||||
|
||||
/**
|
||||
* undici / Node network error CODES that mean the connection broke (not an
|
||||
* application-level error) — a transient transport failure a readOnly call may
|
||||
* safely retry after reconnecting. Matched against the REAL error shapes (#489):
|
||||
* a socket reset surfaces as `TypeError: fetch failed` whose `.cause` is an
|
||||
* undici `SocketError { code:'UND_ERR_SOCKET' }`; a body-timeout as
|
||||
* `TypeError: terminated` whose `.cause` is `BodyTimeoutError`. Classifying by
|
||||
* these real codes/names (not by mock errors) is essential — a mock-shaped
|
||||
* predicate would leave eviction silently dead in production while CI is green.
|
||||
*/
|
||||
const RETRYABLE_TRANSPORT_ERROR_CODES: ReadonlySet<string> = new Set([
|
||||
'ECONNRESET',
|
||||
'ECONNREFUSED',
|
||||
'ECONNABORTED',
|
||||
'EPIPE',
|
||||
'ETIMEDOUT',
|
||||
'EAI_AGAIN',
|
||||
'ENETUNREACH',
|
||||
'EHOSTUNREACH',
|
||||
'UND_ERR_SOCKET',
|
||||
'UND_ERR_CONNECT_TIMEOUT',
|
||||
'UND_ERR_HEADERS_TIMEOUT',
|
||||
'UND_ERR_BODY_TIMEOUT',
|
||||
'UND_ERR_CLOSED',
|
||||
'UND_ERR_DESTROYED',
|
||||
]);
|
||||
|
||||
/** undici error CLASS names for the same transport-break conditions. */
|
||||
const RETRYABLE_TRANSPORT_ERROR_NAMES: ReadonlySet<string> = new Set([
|
||||
'SocketError',
|
||||
'BodyTimeoutError',
|
||||
'HeadersTimeoutError',
|
||||
'ConnectTimeoutError',
|
||||
'ClientClosedError',
|
||||
'ClientDestroyedError',
|
||||
]);
|
||||
|
||||
/**
|
||||
* Whether `err` is a retryable TRANSPORT break (a broken socket / body timeout),
|
||||
* classified by the REAL undici/@ai-sdk error shapes (#489). undici surfaces a
|
||||
* reset as `TypeError('fetch failed'|'terminated')` with the real error in
|
||||
* `.cause`, and @ai-sdk/mcp may wrap it again in an `MCPClientError` (cause
|
||||
* chain), so we walk `.cause` (bounded depth) checking `.code` and `.name`. An
|
||||
* app-level tool error (a 4xx, a validation failure) is NOT retryable and returns
|
||||
* false — only a connection-level failure heals with a reconnect.
|
||||
*/
|
||||
export function isRetryableConnectError(err: unknown, depth = 0): boolean {
|
||||
if (!err || typeof err !== 'object' || depth > 6) return false;
|
||||
const e = err as {
|
||||
code?: unknown;
|
||||
name?: unknown;
|
||||
cause?: unknown;
|
||||
};
|
||||
if (typeof e.code === 'string' && RETRYABLE_TRANSPORT_ERROR_CODES.has(e.code)) {
|
||||
return true;
|
||||
}
|
||||
if (typeof e.name === 'string' && RETRYABLE_TRANSPORT_ERROR_NAMES.has(e.name)) {
|
||||
return true;
|
||||
}
|
||||
if (e.cause != null) return isRetryableConnectError(e.cause, depth + 1);
|
||||
return false;
|
||||
}
|
||||
|
||||
/** The signal's reason as an Error (informative thrown value on abort/timeout). */
|
||||
function abortReason(signal: AbortSignal): Error {
|
||||
const r = signal.reason;
|
||||
|
||||
@@ -0,0 +1,266 @@
|
||||
import type { ModelMessage } from 'ai';
|
||||
import {
|
||||
resolveReplayBudget,
|
||||
isContextOverflowError,
|
||||
estimateMessagesTokens,
|
||||
trimHistoryForReplay,
|
||||
REPLAY_BUDGET_DEFAULT_TOKENS,
|
||||
REPLAY_TRUNCATION_MARKER,
|
||||
REPLAY_TURN_COLLAPSED_MARKER,
|
||||
} from './history-budget';
|
||||
|
||||
describe('resolveReplayBudget', () => {
|
||||
it('uses min(default, 0.7 x window) for a configured window', () => {
|
||||
// 0.7 x 60k = 42k < 100k
|
||||
expect(resolveReplayBudget(60_000)).toEqual({
|
||||
thresholdTokens: 42_000,
|
||||
usedDefault: false,
|
||||
});
|
||||
// 0.7 x 1M = 700k, capped to the 100k default
|
||||
expect(resolveReplayBudget(1_000_000)).toEqual({
|
||||
thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS,
|
||||
usedDefault: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('accepts the raw ::text stored form', () => {
|
||||
expect(resolveReplayBudget('60000').thresholdTokens).toBe(42_000);
|
||||
});
|
||||
|
||||
// The crux (#490): a chat with NO context window configured must STILL be
|
||||
// budgeted — those are exactly the installs that hit terminal overflow.
|
||||
it('applies the flat default when the window is unset/empty', () => {
|
||||
expect(resolveReplayBudget(undefined)).toEqual({
|
||||
thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS,
|
||||
usedDefault: true,
|
||||
});
|
||||
expect(resolveReplayBudget('')).toEqual({
|
||||
thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS,
|
||||
usedDefault: true,
|
||||
});
|
||||
expect(resolveReplayBudget(' ')).toEqual({
|
||||
thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS,
|
||||
usedDefault: true,
|
||||
});
|
||||
});
|
||||
|
||||
it('treats an explicit 0 as the off-switch (distinct from unset)', () => {
|
||||
expect(resolveReplayBudget(0)).toEqual({
|
||||
thresholdTokens: null,
|
||||
usedDefault: false,
|
||||
});
|
||||
expect(resolveReplayBudget('0')).toEqual({
|
||||
thresholdTokens: null,
|
||||
usedDefault: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('falls back to the default on a negative/garbage value', () => {
|
||||
expect(resolveReplayBudget(-5).usedDefault).toBe(true);
|
||||
expect(resolveReplayBudget('abc').usedDefault).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('isContextOverflowError', () => {
|
||||
it('classifies a real provider 400 context-overflow shape', () => {
|
||||
// OpenAI-compatible shape.
|
||||
expect(
|
||||
isContextOverflowError({
|
||||
statusCode: 400,
|
||||
message:
|
||||
"This model's maximum context length is 128000 tokens. However, your messages resulted in 214000 tokens. Please reduce the length of the messages.",
|
||||
}),
|
||||
).toBe(true);
|
||||
// Anthropic-style wording.
|
||||
expect(
|
||||
isContextOverflowError({
|
||||
status: 400,
|
||||
message: 'prompt is too long: 250000 tokens > 200000 maximum',
|
||||
}),
|
||||
).toBe(true);
|
||||
// Nested body + string status.
|
||||
expect(
|
||||
isContextOverflowError({
|
||||
response: { status: '400' },
|
||||
message: 'input is too long for the requested model',
|
||||
}),
|
||||
).toBe(true);
|
||||
// Error instance with the cause carrying the body.
|
||||
const e = new Error('Bad request');
|
||||
(e as any).statusCode = 400;
|
||||
(e as any).cause = new Error('maximum context window exceeded');
|
||||
expect(isContextOverflowError(e)).toBe(true);
|
||||
});
|
||||
|
||||
it('does NOT classify unrelated 400s or auth/rate-limit errors', () => {
|
||||
expect(
|
||||
isContextOverflowError({ statusCode: 400, message: 'invalid tool schema' }),
|
||||
).toBe(false);
|
||||
expect(
|
||||
isContextOverflowError({
|
||||
statusCode: 429,
|
||||
message: 'context length exceeded but rate limited',
|
||||
}),
|
||||
).toBe(false);
|
||||
expect(isContextOverflowError({ statusCode: 500, message: 'server error' })).toBe(
|
||||
false,
|
||||
);
|
||||
expect(isContextOverflowError(undefined)).toBe(false);
|
||||
expect(isContextOverflowError('some random string')).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
// Helpers to build ModelMessage fixtures in the ai@6 shape.
|
||||
const userMsg = (text: string): ModelMessage =>
|
||||
({ role: 'user', content: [{ type: 'text', text }] }) as ModelMessage;
|
||||
const assistantMsg = (
|
||||
text: string,
|
||||
toolCallId?: string,
|
||||
toolName?: string,
|
||||
): ModelMessage =>
|
||||
({
|
||||
role: 'assistant',
|
||||
content: [
|
||||
{ type: 'text', text },
|
||||
...(toolCallId
|
||||
? [{ type: 'tool-call', toolCallId, toolName, input: {} }]
|
||||
: []),
|
||||
],
|
||||
}) as ModelMessage;
|
||||
const toolMsg = (
|
||||
toolCallId: string,
|
||||
toolName: string,
|
||||
value: unknown,
|
||||
): ModelMessage =>
|
||||
({
|
||||
role: 'tool',
|
||||
content: [
|
||||
{ type: 'tool-result', toolCallId, toolName, output: { type: 'json', value } },
|
||||
],
|
||||
}) as ModelMessage;
|
||||
|
||||
describe('trimHistoryForReplay', () => {
|
||||
it('null budget disables trimming (returns the same reference)', () => {
|
||||
const msgs = [userMsg('hi'), assistantMsg('yo')];
|
||||
const r = trimHistoryForReplay(msgs, null);
|
||||
expect(r.trimmed).toBe(false);
|
||||
expect(r.messages).toBe(msgs);
|
||||
});
|
||||
|
||||
it('leaves history under budget untouched (same reference)', () => {
|
||||
const msgs = [userMsg('hi'), assistantMsg('a short answer')];
|
||||
const r = trimHistoryForReplay(msgs, 100_000);
|
||||
expect(r.trimmed).toBe(false);
|
||||
expect(r.messages).toBe(msgs);
|
||||
});
|
||||
|
||||
it('truncates OLD tool outputs but keeps recent turns full', () => {
|
||||
const big = 'X'.repeat(40_000); // ~16k tokens on its own
|
||||
const msgs: ModelMessage[] = [];
|
||||
// 6 OLD turns (indices 0..5), each with a huge tool output.
|
||||
for (let i = 0; i < 6; i++) {
|
||||
msgs.push(userMsg(`old q${i}`));
|
||||
msgs.push(assistantMsg('looking', `c${i}`, 'getPage'));
|
||||
msgs.push(toolMsg(`c${i}`, 'getPage', { body: big }));
|
||||
msgs.push(assistantMsg(`old a${i}`));
|
||||
}
|
||||
// 3 small recent turns, then the CURRENT turn with its own huge tool output.
|
||||
// With REPLAY_KEEP_RECENT_TURNS=4 the last 4 user-turns stay full, so only
|
||||
// these small recent turns + the current big one are kept full; the 6 old
|
||||
// turns above fall in the trim region.
|
||||
for (let i = 0; i < 3; i++) {
|
||||
msgs.push(userMsg(`recent q${i}`));
|
||||
msgs.push(assistantMsg(`recent a${i}`));
|
||||
}
|
||||
msgs.push(userMsg('current q'));
|
||||
msgs.push(assistantMsg('looking', 'cR', 'getPage'));
|
||||
msgs.push(toolMsg('cR', 'getPage', { body: big }));
|
||||
msgs.push(assistantMsg('current a'));
|
||||
|
||||
// Budget large enough that phase-1 tool truncation alone brings it under.
|
||||
const r = trimHistoryForReplay(msgs, 30_000);
|
||||
expect(r.trimmed).toBe(true);
|
||||
const flat = JSON.stringify(r.messages);
|
||||
// The CURRENT turn's tool output survives in full.
|
||||
expect(flat).toContain(big);
|
||||
// Old outputs were truncated with the marker.
|
||||
expect(flat).toContain(REPLAY_TRUNCATION_MARKER);
|
||||
// Phase 1 sufficed: the oldest turns were NOT collapsed.
|
||||
expect(flat).not.toContain(REPLAY_TURN_COLLAPSED_MARKER);
|
||||
expect(estimateMessagesTokens(r.messages)).toBeLessThan(
|
||||
estimateMessagesTokens(msgs),
|
||||
);
|
||||
});
|
||||
|
||||
it('collapses the oldest turns when tool truncation is not enough', () => {
|
||||
// Many turns with LARGE assistant TEXT (not tool output) so phase 1 can't help.
|
||||
const bigText = 'слово '.repeat(8_000); // large Cyrillic text per turn
|
||||
const msgs: ModelMessage[] = [];
|
||||
for (let i = 0; i < 12; i++) {
|
||||
msgs.push(userMsg(`q${i}`));
|
||||
msgs.push(assistantMsg(bigText));
|
||||
}
|
||||
const r = trimHistoryForReplay(msgs, 30_000);
|
||||
expect(r.trimmed).toBe(true);
|
||||
// Oldest turns collapsed; result fits (best-effort) and is much smaller.
|
||||
expect(estimateMessagesTokens(r.messages)).toBeLessThan(
|
||||
estimateMessagesTokens(msgs),
|
||||
);
|
||||
// The LAST turn's text is preserved in full (recent turns stay full).
|
||||
expect(JSON.stringify(r.messages[r.messages.length - 1])).toContain(bigText);
|
||||
});
|
||||
|
||||
it('is deterministic / byte-stable for identical inputs', () => {
|
||||
const big = 'Y'.repeat(30_000);
|
||||
const build = (): ModelMessage[] => {
|
||||
const m: ModelMessage[] = [];
|
||||
for (let i = 0; i < 10; i++) {
|
||||
m.push(userMsg(`q${i}`));
|
||||
m.push(assistantMsg('t', `c${i}`, 'getPage'));
|
||||
m.push(toolMsg(`c${i}`, 'getPage', { body: big }));
|
||||
}
|
||||
return m;
|
||||
};
|
||||
const a = trimHistoryForReplay(build(), 15_000);
|
||||
const b = trimHistoryForReplay(build(), 15_000);
|
||||
expect(JSON.stringify(a.messages)).toBe(JSON.stringify(b.messages));
|
||||
});
|
||||
|
||||
it('never leaves an unpaired tool-call after collapsing (balanced history)', () => {
|
||||
const big = 'Z'.repeat(40_000);
|
||||
const msgs: ModelMessage[] = [];
|
||||
for (let i = 0; i < 10; i++) {
|
||||
msgs.push(userMsg(`q${i}`));
|
||||
msgs.push(assistantMsg('t', `c${i}`, 'getPage'));
|
||||
msgs.push(toolMsg(`c${i}`, 'getPage', { body: big }));
|
||||
}
|
||||
const r = trimHistoryForReplay(msgs, 8_000);
|
||||
// Count tool-call vs tool-result parts in the trimmed output.
|
||||
let calls = 0;
|
||||
let results = 0;
|
||||
for (const m of r.messages) {
|
||||
if (!Array.isArray(m.content)) continue;
|
||||
for (const p of m.content as Array<{ type?: string }>) {
|
||||
if (p.type === 'tool-call') calls++;
|
||||
if (p.type === 'tool-result' || p.type === 'tool-error') results++;
|
||||
}
|
||||
}
|
||||
// Every surviving tool-call has a surviving result (collapsing drops BOTH).
|
||||
expect(calls).toBe(results);
|
||||
// Collapsed turns carry the marker.
|
||||
expect(JSON.stringify(r.messages)).toContain(REPLAY_TURN_COLLAPSED_MARKER);
|
||||
});
|
||||
|
||||
it('respects the provider fact: under-budget contextTokens skips trimming', () => {
|
||||
const big = 'W'.repeat(60_000);
|
||||
const msgs = [
|
||||
userMsg('q'),
|
||||
assistantMsg('t', 'c1', 'getPage'),
|
||||
toolMsg('c1', 'getPage', { body: big }),
|
||||
];
|
||||
// char-estimate is high, but the provider says we are well under budget.
|
||||
const r = trimHistoryForReplay(msgs, 100_000, 5_000);
|
||||
expect(r.trimmed).toBe(false);
|
||||
expect(r.messages).toBe(msgs);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,356 @@
|
||||
/**
|
||||
* History-replay token budget (#490).
|
||||
*
|
||||
* The whole persisted conversation is replayed to the provider on EVERY turn, so
|
||||
* a long chat eventually exceeds the model's context window and the provider 400s
|
||||
* on every turn — terminally (the chat "bricks"). This module bounds the replayed
|
||||
* history at REPLAY TIME only: it never mutates what is persisted (the DB stays
|
||||
* the full record), and its output is a deterministic, byte-stable function of its
|
||||
* input so the trimmed prefix is identical turn to turn (provider prompt-cache
|
||||
* friendliness — real money on long chats).
|
||||
*
|
||||
* The PRIMARY signal is the provider's own fact: `metadata.contextTokens` from the
|
||||
* last turn. The chars-based {@link estimateTokens} (shared with the client) is
|
||||
* used only for the DELTA of not-yet-sent messages, to decide WHAT to trim, and as
|
||||
* the fallback for chats with no usage yet.
|
||||
*/
|
||||
import type { ModelMessage } from 'ai';
|
||||
import { estimateTokens } from '@docmost/token-estimate';
|
||||
|
||||
/** Flat default budget when no context window is configured (tokens). */
|
||||
export const REPLAY_BUDGET_DEFAULT_TOKENS = 100_000;
|
||||
/** Fraction of a configured context window used as the budget. */
|
||||
export const REPLAY_BUDGET_WINDOW_FRACTION = 0.7;
|
||||
/**
|
||||
* Fraction of the normal budget used for the REACTIVE re-trim after a provider
|
||||
* context-overflow 400 — the preventive estimate under-counted, so cut harder.
|
||||
*/
|
||||
export const REPLAY_AGGRESSIVE_FRACTION = 0.5;
|
||||
/**
|
||||
* Turns (a user message + its assistant/tool replies) kept FULL at the tail,
|
||||
* including the current one — never trimmed. Older turns are compacted first.
|
||||
*/
|
||||
export const REPLAY_KEEP_RECENT_TURNS = 4;
|
||||
/** Leading chars kept from a truncated old tool output. */
|
||||
export const REPLAY_TOOL_OUTPUT_HEAD = 800;
|
||||
/** Trailing chars kept from a truncated old tool output. */
|
||||
export const REPLAY_TOOL_OUTPUT_TAIL = 300;
|
||||
/** Marker inserted where an old tool output was truncated for replay. */
|
||||
export const REPLAY_TRUNCATION_MARKER =
|
||||
'[…truncated for replay; call the tool again to read the full output]';
|
||||
/** Marker for a whole old turn collapsed to its text. */
|
||||
export const REPLAY_TURN_COLLAPSED_MARKER =
|
||||
'[earlier tool activity omitted for replay]';
|
||||
|
||||
export interface ReplayBudget {
|
||||
/** Token threshold above which replay history is trimmed; `null` = OFF. */
|
||||
thresholdTokens: number | null;
|
||||
/** True when the flat default was used (no context window configured). */
|
||||
usedDefault: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve the replay budget from the RAW stored `chatContextWindow` (text/number).
|
||||
* - a positive value -> `min(default, floor(fraction × window))`
|
||||
* - explicit `0` -> OFF (admin opt-out; `null` threshold)
|
||||
* - unset/empty/invalid-> the flat default (still protects — the installations
|
||||
* that hit terminal overflow are exactly the ones that never set a window)
|
||||
*
|
||||
* Note the raw value is needed because the parsed `chatContextWindow` collapses
|
||||
* both `0` and unset to `undefined`, which would erase the explicit off-switch.
|
||||
*/
|
||||
export function resolveReplayBudget(rawContextWindow: unknown): ReplayBudget {
|
||||
let n: number | undefined;
|
||||
if (typeof rawContextWindow === 'number') {
|
||||
n = rawContextWindow;
|
||||
} else if (typeof rawContextWindow === 'string') {
|
||||
const t = rawContextWindow.trim();
|
||||
n = t === '' ? undefined : Number(t);
|
||||
}
|
||||
// Unset / empty / non-numeric / negative -> flat default (the protective case).
|
||||
if (n === undefined || !Number.isFinite(n) || n < 0) {
|
||||
return { thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS, usedDefault: true };
|
||||
}
|
||||
// Explicit 0 -> off-switch.
|
||||
if (n === 0) {
|
||||
return { thresholdTokens: null, usedDefault: false };
|
||||
}
|
||||
return {
|
||||
thresholdTokens: Math.min(
|
||||
REPLAY_BUDGET_DEFAULT_TOKENS,
|
||||
Math.floor(REPLAY_BUDGET_WINDOW_FRACTION * n),
|
||||
),
|
||||
usedDefault: false,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* True when a provider error is a CONTEXT-OVERFLOW rejection (the prompt exceeds
|
||||
* the model's window). Providers surface this as an HTTP 400 with a recognizable
|
||||
* message; match both the status and the message patterns robustly across
|
||||
* OpenAI-compatible / Anthropic / Gemini wordings, since the exact shape varies.
|
||||
*/
|
||||
export function isContextOverflowError(error: unknown): boolean {
|
||||
const status = extractStatus(error);
|
||||
const msg = extractMessage(error).toLowerCase();
|
||||
// Message patterns seen across providers for "prompt too long".
|
||||
const overflowPattern =
|
||||
/context (?:length|window)|maximum context|too many tokens|too large for|reduce the length|prompt is too long|input (?:is )?too long|exceeds? the (?:maximum )?(?:context|token)|maximum.*tokens|string too long/;
|
||||
if (!overflowPattern.test(msg)) return false;
|
||||
// A 400/413 with an overflow-shaped message is an overflow. Some providers
|
||||
// omit/rewrite the status, so accept the message match when the status is
|
||||
// unknown, but reject it for auth/rate-limit statuses that never mean overflow.
|
||||
if (status === 400 || status === 413) return true;
|
||||
if (status === 401 || status === 403 || status === 429) return false;
|
||||
return true;
|
||||
}
|
||||
|
||||
function extractStatus(error: unknown): number | undefined {
|
||||
if (!error || typeof error !== 'object') return undefined;
|
||||
const e = error as Record<string, unknown>;
|
||||
for (const k of ['statusCode', 'status']) {
|
||||
const v = e[k];
|
||||
if (typeof v === 'number') return v;
|
||||
if (typeof v === 'string' && /^\d+$/.test(v)) return Number(v);
|
||||
}
|
||||
// Nested (e.g. { response: { status } } / { cause: { statusCode } }).
|
||||
for (const k of ['response', 'cause', 'data']) {
|
||||
const nested = e[k];
|
||||
if (nested && typeof nested === 'object') {
|
||||
const s = extractStatus(nested);
|
||||
if (s !== undefined) return s;
|
||||
}
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function extractMessage(error: unknown): string {
|
||||
if (error == null) return '';
|
||||
if (typeof error === 'string') return error;
|
||||
if (error instanceof Error) {
|
||||
// Include nested causes (provider libs wrap the real body in `cause`).
|
||||
const cause = (error as { cause?: unknown }).cause;
|
||||
return `${error.message} ${cause ? extractMessage(cause) : ''}`;
|
||||
}
|
||||
if (typeof error === 'object') {
|
||||
const e = error as Record<string, unknown>;
|
||||
const parts: string[] = [];
|
||||
for (const k of ['message', 'error', 'body', 'responseBody', 'data']) {
|
||||
const v = e[k];
|
||||
if (typeof v === 'string') parts.push(v);
|
||||
else if (v && typeof v === 'object') parts.push(extractMessage(v));
|
||||
}
|
||||
return parts.join(' ');
|
||||
}
|
||||
return String(error);
|
||||
}
|
||||
|
||||
/** Rough token size of a ModelMessage array via the shared chars estimator. */
|
||||
export function estimateMessagesTokens(
|
||||
messages: ReadonlyArray<ModelMessage>,
|
||||
): number {
|
||||
let total = 0;
|
||||
for (const m of messages) {
|
||||
total += estimateTokens(serializeContent(m.content));
|
||||
}
|
||||
return total;
|
||||
}
|
||||
|
||||
function serializeContent(content: unknown): string {
|
||||
if (typeof content === 'string') return content;
|
||||
try {
|
||||
return JSON.stringify(content) ?? '';
|
||||
} catch {
|
||||
return '';
|
||||
}
|
||||
}
|
||||
|
||||
/** Deep JSON string of an arbitrary value, bounded so estimation never throws. */
|
||||
function stringifyValue(value: unknown): string {
|
||||
if (typeof value === 'string') return value;
|
||||
try {
|
||||
return JSON.stringify(value) ?? String(value);
|
||||
} catch {
|
||||
return String(value);
|
||||
}
|
||||
}
|
||||
|
||||
export interface TrimResult {
|
||||
messages: ModelMessage[];
|
||||
/** Whether any trimming was applied. */
|
||||
trimmed: boolean;
|
||||
/** Estimated tokens of the returned messages (chars-based). */
|
||||
estimatedTokens: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Bound the replayed history to `budgetTokens`, deterministically. Returns the
|
||||
* SAME array reference (no copy) when nothing needs trimming, so the common case
|
||||
* is free and byte-identical. Trimming order (spec #490):
|
||||
* 1. truncate OLD turns' tool outputs (head+tail + marker) — the bulk of the size
|
||||
* 2. mechanically collapse the OLDEST turns to their text (concatenation, no LLM)
|
||||
* 3. the current + last {@link REPLAY_KEEP_RECENT_TURNS} turns stay FULL
|
||||
*
|
||||
* `budgetTokens === null` disables trimming. `priorContextTokens` (the provider's
|
||||
* fact from last turn) short-circuits the decision: when it is known and already
|
||||
* under budget we skip trimming even if the char-estimate is higher (the provider
|
||||
* count is authoritative). The char-estimate drives WHAT to cut.
|
||||
*/
|
||||
export function trimHistoryForReplay(
|
||||
messages: ModelMessage[],
|
||||
budgetTokens: number | null,
|
||||
priorContextTokens?: number,
|
||||
): TrimResult {
|
||||
if (budgetTokens == null) {
|
||||
return { messages, trimmed: false, estimatedTokens: 0 };
|
||||
}
|
||||
const estimated = estimateMessagesTokens(messages);
|
||||
// Decision signal: prefer the provider's fact (last turn's contextTokens) plus
|
||||
// the estimated delta of the messages appended since; fall back to the pure
|
||||
// char-estimate for a chat with no usage yet.
|
||||
const projected =
|
||||
priorContextTokens != null
|
||||
? Math.max(priorContextTokens, estimated)
|
||||
: estimated;
|
||||
if (projected <= budgetTokens) {
|
||||
return { messages, trimmed: false, estimatedTokens: estimated };
|
||||
}
|
||||
|
||||
// The tail we always keep full: from the Nth-from-last user message onward.
|
||||
const boundary = recentBoundaryIndex(messages, REPLAY_KEEP_RECENT_TURNS);
|
||||
const tail = messages.slice(boundary);
|
||||
let head = messages.slice(0, boundary).map(cloneMessage);
|
||||
|
||||
// Phase 1: truncate old tool outputs.
|
||||
for (const m of head) {
|
||||
if (m.role === 'tool') truncateToolMessage(m);
|
||||
}
|
||||
let out = [...head, ...tail];
|
||||
let est = estimateMessagesTokens(out);
|
||||
if (est <= budgetTokens) {
|
||||
return { messages: out, trimmed: true, estimatedTokens: est };
|
||||
}
|
||||
|
||||
// Phase 2: collapse the oldest turns (in `head`) to their text, one at a time,
|
||||
// from the oldest, until we fit or the whole head is collapsed.
|
||||
const turns = splitTurns(head);
|
||||
const collapsed: ModelMessage[] = [];
|
||||
let i = 0;
|
||||
for (; i < turns.length; i++) {
|
||||
if (est <= budgetTokens) break;
|
||||
collapsed.push(...collapseTurn(turns[i]));
|
||||
// Re-estimate the whole prospective output.
|
||||
const remaining = turns.slice(i + 1).flat();
|
||||
out = [...collapsed, ...remaining, ...tail];
|
||||
est = estimateMessagesTokens(out);
|
||||
}
|
||||
// Include any turns we didn't need to collapse.
|
||||
const remaining = turns.slice(i).flat();
|
||||
out = [...collapsed, ...remaining, ...tail];
|
||||
est = estimateMessagesTokens(out);
|
||||
return { messages: out, trimmed: true, estimatedTokens: est };
|
||||
}
|
||||
|
||||
/** Index of the first message of the Nth-from-last user turn (0 if fewer). */
|
||||
function recentBoundaryIndex(
|
||||
messages: ReadonlyArray<ModelMessage>,
|
||||
keepTurns: number,
|
||||
): number {
|
||||
const userIdx: number[] = [];
|
||||
for (let i = 0; i < messages.length; i++) {
|
||||
if (messages[i].role === 'user') userIdx.push(i);
|
||||
}
|
||||
if (userIdx.length <= keepTurns) return 0;
|
||||
return userIdx[userIdx.length - keepTurns];
|
||||
}
|
||||
|
||||
/** Split a message list into turns; each turn starts at a `user` message. */
|
||||
function splitTurns(messages: ModelMessage[]): ModelMessage[][] {
|
||||
const turns: ModelMessage[][] = [];
|
||||
for (const m of messages) {
|
||||
if (m.role === 'user' || turns.length === 0) turns.push([m]);
|
||||
else turns[turns.length - 1].push(m);
|
||||
}
|
||||
return turns;
|
||||
}
|
||||
|
||||
/**
|
||||
* Collapse a whole turn to its plain text (mechanical concatenation, not an LLM
|
||||
* summary). Keeps the user message; replaces the assistant/tool messages with a
|
||||
* single assistant text message = the assistant's concatenated text + a marker
|
||||
* when tool activity was dropped. Dropping BOTH the tool-call and tool-result
|
||||
* parts together keeps the rebuilt history balanced (no unpaired calls).
|
||||
*/
|
||||
function collapseTurn(turn: ModelMessage[]): ModelMessage[] {
|
||||
const out: ModelMessage[] = [];
|
||||
let assistantText = '';
|
||||
let hadTools = false;
|
||||
for (const m of turn) {
|
||||
if (m.role === 'user') {
|
||||
out.push(m);
|
||||
} else if (m.role === 'assistant') {
|
||||
const { text, tools } = extractAssistantText(m.content);
|
||||
assistantText += text;
|
||||
hadTools = hadTools || tools;
|
||||
} else if (m.role === 'tool') {
|
||||
hadTools = true;
|
||||
} else {
|
||||
out.push(m);
|
||||
}
|
||||
}
|
||||
const note =
|
||||
(assistantText ? assistantText.trimEnd() : '') +
|
||||
(hadTools
|
||||
? `${assistantText ? '\n\n' : ''}${REPLAY_TURN_COLLAPSED_MARKER}`
|
||||
: '');
|
||||
if (note) out.push({ role: 'assistant', content: note } as ModelMessage);
|
||||
return out;
|
||||
}
|
||||
|
||||
function extractAssistantText(content: unknown): {
|
||||
text: string;
|
||||
tools: boolean;
|
||||
} {
|
||||
if (typeof content === 'string') return { text: content, tools: false };
|
||||
if (!Array.isArray(content)) return { text: '', tools: false };
|
||||
let text = '';
|
||||
let tools = false;
|
||||
for (const part of content) {
|
||||
const type = (part as { type?: string })?.type;
|
||||
if (type === 'text') text += (part as { text?: string }).text ?? '';
|
||||
else if (type === 'tool-call') tools = true;
|
||||
}
|
||||
return { text, tools };
|
||||
}
|
||||
|
||||
/** Truncate every tool-result output in a `tool` message to head+tail+marker. */
|
||||
function truncateToolMessage(message: ModelMessage): void {
|
||||
const content = message.content;
|
||||
if (!Array.isArray(content)) return;
|
||||
for (const part of content) {
|
||||
const p = part as { type?: string; output?: { type?: string; value?: unknown } };
|
||||
if (p.type !== 'tool-result' && p.type !== 'tool-error') continue;
|
||||
if (!p.output) continue;
|
||||
const raw = stringifyValue(p.output.value);
|
||||
const budget = REPLAY_TOOL_OUTPUT_HEAD + REPLAY_TOOL_OUTPUT_TAIL;
|
||||
if (raw.length <= budget + REPLAY_TRUNCATION_MARKER.length) continue;
|
||||
const truncated =
|
||||
raw.slice(0, REPLAY_TOOL_OUTPUT_HEAD) +
|
||||
`\n${REPLAY_TRUNCATION_MARKER}\n` +
|
||||
raw.slice(raw.length - REPLAY_TOOL_OUTPUT_TAIL);
|
||||
// Represent the shrunk output as a text output (a valid tool-result output).
|
||||
p.output = { type: 'text', value: truncated };
|
||||
}
|
||||
}
|
||||
|
||||
/** Shallow-ish clone so trimming never mutates the caller's (persisted-derived)
|
||||
* message objects — only the OLD region is cloned before it is edited. */
|
||||
function cloneMessage(m: ModelMessage): ModelMessage {
|
||||
if (typeof m.content === 'string') return { ...m };
|
||||
return {
|
||||
...m,
|
||||
content: (m.content as unknown[]).map((p) =>
|
||||
p && typeof p === 'object' ? { ...(p as object) } : p,
|
||||
),
|
||||
} as ModelMessage;
|
||||
}
|
||||
@@ -1,11 +1,24 @@
|
||||
import { Logger } from '@nestjs/common';
|
||||
import { streamText } from 'ai';
|
||||
import {
|
||||
hasRepeatedLineRun,
|
||||
hasPeriodicTail,
|
||||
isDegenerateOutput,
|
||||
truncateDegeneratedTail,
|
||||
shouldCheckDegeneration,
|
||||
DEGENERATION_CHECK_STEP,
|
||||
REPEATED_LINES_THRESHOLD,
|
||||
MIN_PERIOD_REPEATS,
|
||||
} from './output-degeneration';
|
||||
import { AiChatService } from './ai-chat.service';
|
||||
|
||||
// Mock ONLY streamText so we can capture the onChunk/onStepFinish callbacks the
|
||||
// service registers and drive them by hand; every other `ai` export the service
|
||||
// uses (convertToModelMessages, stepCountIs, …) stays real.
|
||||
jest.mock('ai', () => {
|
||||
const actual = jest.requireActual('ai');
|
||||
return { ...actual, streamText: jest.fn() };
|
||||
});
|
||||
|
||||
/**
|
||||
* Unit tests for the token-degeneration detector (#444) — the sole anti-babble
|
||||
@@ -180,3 +193,188 @@ describe('truncateDegeneratedTail', () => {
|
||||
expect(truncateDegeneratedTail(text)).toBe(text);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* Throttle + step-boundary reset (#486). The stream keeps a watermark
|
||||
* (`lastDegenerationCheckLen`) that is an OFFSET into the accumulated step text.
|
||||
* On a step boundary the accumulator resets to '', so the watermark MUST reset to
|
||||
* 0 too — otherwise the throttle goes silent for the whole next step. These tests
|
||||
* pin the pure decision AND the reset property that ai-chat.service.onStepFinish
|
||||
* now enforces.
|
||||
*/
|
||||
describe('shouldCheckDegeneration (throttle) + step-boundary reset (#486)', () => {
|
||||
it('fires once the text grows a full DEGENERATION_CHECK_STEP past the mark', () => {
|
||||
expect(shouldCheckDegeneration(DEGENERATION_CHECK_STEP, 0)).toBe(true);
|
||||
expect(shouldCheckDegeneration(DEGENERATION_CHECK_STEP - 1, 0)).toBe(false);
|
||||
expect(shouldCheckDegeneration(5000, 3000)).toBe(true); // grew 2000 since mark
|
||||
expect(shouldCheckDegeneration(4000, 3000)).toBe(false); // grew only 1000
|
||||
});
|
||||
|
||||
it('BUG (no reset): a stale large watermark silences the next step', () => {
|
||||
// End of a long step: the watermark sits at 5000. The step ends and the
|
||||
// accumulator resets to '' — but if the watermark is NOT reset, a fresh short
|
||||
// degenerate burst (length 2000) never triggers a check: 2000 - 5000 < STEP.
|
||||
const staleWatermark = 5000;
|
||||
const nextStepLen = DEGENERATION_CHECK_STEP; // a fresh 2KB burst
|
||||
expect(shouldCheckDegeneration(nextStepLen, staleWatermark)).toBe(false);
|
||||
});
|
||||
|
||||
it('FIX (reset to 0): the same short degenerate burst IS checked and detected', () => {
|
||||
// onStepFinish now zeroes the watermark, so the fresh burst re-arms the check.
|
||||
const resetWatermark = 0;
|
||||
const degenerateBurst = 'loadTools.\n'.repeat(300); // real degeneration
|
||||
expect(degenerateBurst.length).toBeGreaterThanOrEqual(DEGENERATION_CHECK_STEP);
|
||||
// The throttle now fires...
|
||||
expect(
|
||||
shouldCheckDegeneration(degenerateBurst.length, resetWatermark),
|
||||
).toBe(true);
|
||||
// ...and the detector catches the loop that would otherwise stream unchecked.
|
||||
expect(isDegenerateOutput(degenerateBurst)).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* BEHAVIOR guard for the ACTUAL fix (#486, ai-chat.service.onStepFinish resets
|
||||
* lastDegenerationCheckLen to 0). The pure tests above use a hard-coded
|
||||
* resetWatermark, so a REVERT of the real `lastDegenerationCheckLen = 0` line
|
||||
* would not redden any of them. This drives the REAL onChunk/onStepFinish
|
||||
* closures from stream() end to end and asserts the run is aborted when a fresh
|
||||
* degenerate burst arrives in the step AFTER a long clean step — which only
|
||||
* happens if the watermark was actually zeroed on the step boundary.
|
||||
*/
|
||||
describe('AiChatService: onStepFinish re-arms the degeneration watermark (#486)', () => {
|
||||
const streamTextMock = streamText as unknown as jest.Mock;
|
||||
|
||||
function makeRes() {
|
||||
return {
|
||||
raw: {
|
||||
writeHead: jest.fn(),
|
||||
write: jest.fn(),
|
||||
once: jest.fn(),
|
||||
on: jest.fn(),
|
||||
flushHeaders: jest.fn(),
|
||||
writableEnded: false,
|
||||
destroyed: false,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
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: [],
|
||||
})),
|
||||
};
|
||||
return 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,
|
||||
// Lockdown OFF -> the degeneration guard is the active anti-babble path.
|
||||
isAiChatFinalStepLockdownEnabled: () => false,
|
||||
} as never, // environment
|
||||
);
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
streamTextMock.mockReset();
|
||||
jest.spyOn(Logger.prototype, 'log').mockImplementation(() => undefined as never);
|
||||
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined as never);
|
||||
});
|
||||
|
||||
afterEach(() => jest.restoreAllMocks());
|
||||
|
||||
it('aborts on a fresh degenerate burst in the NEXT step (reverting the reset line reddens this)', async () => {
|
||||
let captured:
|
||||
| {
|
||||
onChunk?: (e: { chunk: { type: string; text: string } }) => void;
|
||||
onStepFinish?: (step: unknown) => void;
|
||||
abortSignal?: AbortSignal;
|
||||
}
|
||||
| undefined;
|
||||
streamTextMock.mockImplementation((opts: never) => {
|
||||
captured = opts;
|
||||
return {
|
||||
consumeStream: jest.fn(),
|
||||
pipeUIMessageStreamToResponse: jest.fn(),
|
||||
};
|
||||
});
|
||||
|
||||
const svc = makeService();
|
||||
await svc.stream({
|
||||
user: { id: 'user-1' } as never,
|
||||
workspace: { id: 'ws-1' } as never,
|
||||
sessionId: 'sess-1',
|
||||
body: {
|
||||
chatId: 'chat-1',
|
||||
messages: [
|
||||
{ id: 'm1', role: 'user', parts: [{ type: 'text', text: 'hi' }] },
|
||||
],
|
||||
} as never,
|
||||
res: makeRes() as never,
|
||||
signal: new AbortController().signal,
|
||||
model: {} as never,
|
||||
role: null,
|
||||
// No runHooks -> legacy path (socket signal), degeneration guard active.
|
||||
});
|
||||
|
||||
expect(streamTextMock).toHaveBeenCalledTimes(1);
|
||||
const onChunk = captured!.onChunk!;
|
||||
const onStepFinish = captured!.onStepFinish!;
|
||||
const abortSignal = captured!.abortSignal!;
|
||||
expect(abortSignal.aborted).toBe(false);
|
||||
|
||||
// STEP 1: a LONG, non-degenerate first step. Distinct lines never trip the
|
||||
// detector, but they advance the throttle watermark far past the burst size
|
||||
// that follows (to ~5x the step). This is the stale watermark that, WITHOUT
|
||||
// the reset, would silence step 2.
|
||||
let counter = 0;
|
||||
let accumulated = 0;
|
||||
while (accumulated < DEGENERATION_CHECK_STEP * 5) {
|
||||
const line = `unique clean line number ${counter++} with distinct words\n`;
|
||||
accumulated += line.length;
|
||||
onChunk({ chunk: { type: 'text-delta', text: line } });
|
||||
}
|
||||
expect(abortSignal.aborted).toBe(false); // clean step must not abort
|
||||
|
||||
// STEP BOUNDARY: the real onStepFinish resets inProgressText AND (the fix)
|
||||
// zeroes lastDegenerationCheckLen.
|
||||
onStepFinish({ text: 'a clean first step', toolCalls: [], toolResults: [] });
|
||||
|
||||
// STEP 2: a FRESH, short degenerate burst (~3.3KB). Its length is far below
|
||||
// the step-1 stale watermark (~10KB), so WITHOUT the reset the throttle stays
|
||||
// silent and this streams unchecked. WITH the reset (watermark 0) it re-arms,
|
||||
// the detector fires, and the run aborts.
|
||||
const burst = 'loadTools.\n'.repeat(300);
|
||||
expect(burst.length).toBeGreaterThanOrEqual(DEGENERATION_CHECK_STEP);
|
||||
expect(burst.length).toBeLessThan(DEGENERATION_CHECK_STEP * 5);
|
||||
onChunk({ chunk: { type: 'text-delta', text: burst } });
|
||||
|
||||
// The decisive assertion: the composed abortSignal (unioned with the
|
||||
// degeneration controller) is now aborted. Reverting `lastDegenerationCheckLen
|
||||
// = 0` in onStepFinish makes this stay false.
|
||||
expect(abortSignal.aborted).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -131,6 +131,32 @@ export function isDegenerateOutput(text: string): boolean {
|
||||
return hasRepeatedLineRun(text) || hasPeriodicTail(text);
|
||||
}
|
||||
|
||||
/**
|
||||
* How many bytes the in-progress text must grow before the (amortized) tail
|
||||
* heuristics are re-run. Shared with ai-chat.service so the throttle the stream
|
||||
* applies is the SAME one the unit test drives.
|
||||
*/
|
||||
export const DEGENERATION_CHECK_STEP = 2000;
|
||||
|
||||
/**
|
||||
* Throttle decision for the degeneration guard (#444/#486). Returns true when
|
||||
* the accumulated text has grown at least DEGENERATION_CHECK_STEP bytes past the
|
||||
* last-checked offset, so the pure rules only fire every ~2KB. Pure; the caller
|
||||
* updates its watermark to `textLen` when this returns true.
|
||||
*
|
||||
* The watermark is an offset INTO the accumulator, so when the accumulator is
|
||||
* reset to '' on a step boundary the caller MUST reset the watermark to 0 too
|
||||
* (#486). Otherwise `textLen - lastCheckLen` goes negative after the reset and
|
||||
* this returns false until a later step re-grows past the stale offset — a whole
|
||||
* degenerate step could stream unchecked.
|
||||
*/
|
||||
export function shouldCheckDegeneration(
|
||||
textLen: number,
|
||||
lastCheckLen: number,
|
||||
): boolean {
|
||||
return textLen - lastCheckLen >= DEGENERATION_CHECK_STEP;
|
||||
}
|
||||
|
||||
/**
|
||||
* Truncate a degenerated tail before persist so hundreds of KB of garbage never
|
||||
* reach the DB / replay (#444). Keeps everything up to and including the FIRST
|
||||
|
||||
@@ -0,0 +1,241 @@
|
||||
// Break the editor-ext import chain (share.service -> collaboration.util ->
|
||||
// @docmost/editor-ext -> @tiptap/core) that is unresolvable in this jest env and
|
||||
// pre-existingly breaks these specs. jsonToMarkdown is never reached in these
|
||||
// tests (the tools fail before rendering markdown).
|
||||
jest.mock('../../collaboration/collaboration.util', () => ({
|
||||
jsonToMarkdown: () => '',
|
||||
}));
|
||||
|
||||
import { Logger } from '@nestjs/common';
|
||||
import { MockLanguageModelV3, simulateReadableStream } from 'ai/test';
|
||||
import { PublicShareChatService } from './public-share-chat.service';
|
||||
import { PublicShareChatToolsService } from './tools/public-share-chat-tools.service';
|
||||
|
||||
/**
|
||||
* SECURITY integration guard for #394 (commit 5): a tool's or the provider's raw
|
||||
* error text must NOT leak to an anonymous public-share reader.
|
||||
*
|
||||
* The render gate (ToolCallCard showErrors=false) hides the text in the DOM but
|
||||
* NOT on the wire, so this test asserts on the RAW SSE BYTES the server writes —
|
||||
* exactly the channel the render gate masks. We drive the real
|
||||
* PublicShareChatService.stream() with a real share toolset (its underlying
|
||||
* services mocked to fail) and a mock model, then inspect every byte piped to the
|
||||
* fake socket.
|
||||
*/
|
||||
|
||||
// A minimal ServerResponse stand-in that records every written chunk.
|
||||
class FakeSocket {
|
||||
chunks: string[] = [];
|
||||
statusCode = 200;
|
||||
writableEnded = false;
|
||||
destroyed = false;
|
||||
headersSent = false;
|
||||
writeHead(): this {
|
||||
this.headersSent = true;
|
||||
return this;
|
||||
}
|
||||
setHeader(): void {}
|
||||
removeHeader(): void {}
|
||||
getHeader(): undefined {
|
||||
return undefined;
|
||||
}
|
||||
flushHeaders(): void {}
|
||||
write(chunk: unknown): boolean {
|
||||
this.chunks.push(
|
||||
typeof chunk === 'string' ? chunk : Buffer.from(chunk as never).toString('utf8'),
|
||||
);
|
||||
return true;
|
||||
}
|
||||
end(chunk?: unknown): void {
|
||||
if (chunk) this.write(chunk);
|
||||
this.writableEnded = true;
|
||||
}
|
||||
on(): this {
|
||||
return this;
|
||||
}
|
||||
once(): this {
|
||||
return this;
|
||||
}
|
||||
get body(): string {
|
||||
return this.chunks.join('');
|
||||
}
|
||||
}
|
||||
|
||||
/** Mock model that issues one getSharePage tool call, then finishes with text. */
|
||||
function toolCallingModel(): MockLanguageModelV3 {
|
||||
let call = 0;
|
||||
return new MockLanguageModelV3({
|
||||
doStream: async () => {
|
||||
call++;
|
||||
if (call === 1) {
|
||||
return {
|
||||
stream: simulateReadableStream({
|
||||
chunks: [
|
||||
{ type: 'stream-start' as const, warnings: [] },
|
||||
{ type: 'tool-input-start' as const, id: 't1', toolName: 'getSharePage' },
|
||||
{ type: 'tool-input-end' as const, id: 't1' },
|
||||
{
|
||||
type: 'tool-call' as const,
|
||||
toolCallId: 't1',
|
||||
toolName: 'getSharePage',
|
||||
input: '{"pageId":"secret-page"}',
|
||||
},
|
||||
{
|
||||
type: 'finish' as const,
|
||||
finishReason: { unified: 'tool-calls' as const, raw: 'tool_calls' },
|
||||
usage: {
|
||||
inputTokens: { total: 1, noCache: undefined, cacheRead: undefined, cacheWrite: undefined },
|
||||
outputTokens: { total: 1, text: 1, reasoning: undefined },
|
||||
},
|
||||
},
|
||||
],
|
||||
}),
|
||||
};
|
||||
}
|
||||
return {
|
||||
stream: simulateReadableStream({
|
||||
chunks: [
|
||||
{ type: 'stream-start' as const, warnings: [] },
|
||||
{ type: 'text-start' as const, id: '1' },
|
||||
{ type: 'text-delta' as const, id: '1', delta: 'Sorry.' },
|
||||
{ type: 'text-end' as const, id: '1' },
|
||||
{
|
||||
type: 'finish' as const,
|
||||
finishReason: { unified: 'stop' as const, raw: 'stop' },
|
||||
usage: {
|
||||
inputTokens: { total: 1, noCache: undefined, cacheRead: undefined, cacheWrite: undefined },
|
||||
outputTokens: { total: 1, text: 1, reasoning: undefined },
|
||||
},
|
||||
},
|
||||
],
|
||||
}),
|
||||
};
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
/** Mock model whose stream emits a provider error carrying an internal secret. */
|
||||
function providerErrorModel(secret: string): MockLanguageModelV3 {
|
||||
return new MockLanguageModelV3({
|
||||
doStream: async () => ({
|
||||
stream: simulateReadableStream({
|
||||
chunks: [
|
||||
{ type: 'stream-start' as const, warnings: [] },
|
||||
{
|
||||
type: 'error' as const,
|
||||
error: {
|
||||
statusCode: 503,
|
||||
message: 'Service Unavailable',
|
||||
responseBody: `upstream ${secret} model=internal-gpt`,
|
||||
},
|
||||
},
|
||||
],
|
||||
}),
|
||||
}),
|
||||
});
|
||||
}
|
||||
|
||||
function makeService(toolsService: PublicShareChatToolsService): {
|
||||
svc: PublicShareChatService;
|
||||
logSpy: jest.SpyInstance;
|
||||
} {
|
||||
const svc = Object.create(PublicShareChatService.prototype);
|
||||
const logger = new Logger('test');
|
||||
const logSpy = jest.spyOn(logger, 'error').mockImplementation(() => undefined);
|
||||
jest.spyOn(logger, 'warn').mockImplementation(() => undefined);
|
||||
svc.tools = toolsService;
|
||||
svc.logger = logger;
|
||||
svc.tokenBudget = { record: jest.fn().mockResolvedValue(undefined) };
|
||||
return { svc, logSpy };
|
||||
}
|
||||
|
||||
async function runStream(
|
||||
svc: PublicShareChatService,
|
||||
model: MockLanguageModelV3,
|
||||
): Promise<FakeSocket> {
|
||||
const socket = new FakeSocket();
|
||||
await svc.stream({
|
||||
workspaceId: 'ws1',
|
||||
shareId: 'share1',
|
||||
share: { id: 'share1', pageId: 'p1', sharedPage: { id: 'p1', title: 'Docs' } },
|
||||
openedPage: null,
|
||||
messages: [
|
||||
{ id: 'm1', role: 'user', parts: [{ type: 'text', text: 'read the page' }] } as never,
|
||||
],
|
||||
res: { raw: socket } as never,
|
||||
signal: new AbortController().signal,
|
||||
model: model as never,
|
||||
role: null,
|
||||
});
|
||||
// Let the piped stream drain fully.
|
||||
await new Promise((r) => setTimeout(r, 300));
|
||||
return socket;
|
||||
}
|
||||
|
||||
describe('public share chat error leak (#394)', () => {
|
||||
afterEach(() => jest.restoreAllMocks());
|
||||
|
||||
it('does NOT leak a tool\'s raw internal error to the SSE bytes (generic classified string instead)', async () => {
|
||||
const SECRET = 'INTERNAL_baseUrl_http://provider.internal:8080/v1';
|
||||
const shareService = {
|
||||
// The canonical boundary throws a RAW internal error (with a secret).
|
||||
resolveReadableSharePage: jest
|
||||
.fn()
|
||||
.mockRejectedValue(new Error(`db failed at ${SECRET} stack@line42`)),
|
||||
};
|
||||
const tools = new PublicShareChatToolsService(
|
||||
shareService as never,
|
||||
{} as never,
|
||||
{} as never,
|
||||
);
|
||||
const { svc } = makeService(tools);
|
||||
|
||||
const socket = await runStream(svc, toolCallingModel());
|
||||
|
||||
// The tool-output-error frame is present on the wire...
|
||||
expect(socket.body).toContain('tool-output-error');
|
||||
// ...but it carries ONLY the generic classified string — never the secret,
|
||||
// the raw driver message, or a stack fragment.
|
||||
expect(socket.body).toContain('The tool could not complete the request.');
|
||||
expect(socket.body).not.toContain(SECRET);
|
||||
expect(socket.body).not.toContain('stack@line42');
|
||||
expect(socket.body).not.toContain('db failed');
|
||||
});
|
||||
|
||||
it('passes a SAFE ShareToolError message (page not available) through to the bytes', async () => {
|
||||
const shareService = {
|
||||
// Not found in this share -> the tool throws the classified SAFE message.
|
||||
resolveReadableSharePage: jest.fn().mockResolvedValue(null),
|
||||
};
|
||||
const tools = new PublicShareChatToolsService(
|
||||
shareService as never,
|
||||
{} as never,
|
||||
{} as never,
|
||||
);
|
||||
const { svc } = makeService(tools);
|
||||
|
||||
const socket = await runStream(svc, toolCallingModel());
|
||||
expect(socket.body).toContain('tool-output-error');
|
||||
expect(socket.body).toContain('not available in this share');
|
||||
});
|
||||
|
||||
it('does NOT leak a provider error (statusCode + response body) to the SSE bytes', async () => {
|
||||
const SECRET = 'http://provider.internal:8080';
|
||||
const tools = new PublicShareChatToolsService(
|
||||
{} as never,
|
||||
{} as never,
|
||||
{} as never,
|
||||
);
|
||||
const { svc, logSpy } = makeService(tools);
|
||||
|
||||
const socket = await runStream(svc, providerErrorModel(SECRET));
|
||||
|
||||
// The anon sees a fixed classified string, not the provider body/baseUrl/model.
|
||||
expect(socket.body).toContain('temporarily unavailable');
|
||||
expect(socket.body).not.toContain(SECRET);
|
||||
expect(socket.body).not.toContain('internal-gpt');
|
||||
// The FULL provider detail is logged server-side only.
|
||||
const logged = logSpy.mock.calls.map((c) => String(c[0])).join('\n');
|
||||
expect(logged).toContain(SECRET);
|
||||
});
|
||||
});
|
||||
@@ -12,7 +12,10 @@ import { AiAgentRoleRepo } from '@docmost/db/repos/ai-agent-roles/ai-agent-roles
|
||||
import { AiAgentRole } from '@docmost/db/types/entity.types';
|
||||
import { AiService } from '../../integrations/ai/ai.service';
|
||||
import { AiSettingsService } from '../../integrations/ai/ai-settings.service';
|
||||
import { PublicShareChatToolsService } from './tools/public-share-chat-tools.service';
|
||||
import {
|
||||
PublicShareChatToolsService,
|
||||
ShareToolError,
|
||||
} from './tools/public-share-chat-tools.service';
|
||||
import { buildShareSystemPrompt } from './public-share-chat.prompt';
|
||||
import { roleModelOverride } from './roles/role-model-config';
|
||||
import {
|
||||
@@ -102,6 +105,30 @@ export function filterShareTranscript(messages: UIMessage[]): UIMessage[] {
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Fixed, classified strings an ANONYMOUS share reader may see when the assistant
|
||||
* stream fails (#394). These reveal NOTHING about the internal provider, its
|
||||
* baseUrl, the model name, or the raw response body — unlike describeProviderError
|
||||
* (which is for the server log / the authenticated operator only). We classify by
|
||||
* HTTP status where available so the reader still gets a useful hint (retry vs.
|
||||
* give up) without any internal detail.
|
||||
*/
|
||||
export function classifyAnonStreamError(error: unknown): string {
|
||||
const status =
|
||||
typeof error === 'object' && error !== null
|
||||
? (error as { statusCode?: number }).statusCode
|
||||
: undefined;
|
||||
if (status === 429) {
|
||||
return 'The assistant is receiving too many requests right now. Please try again shortly.';
|
||||
}
|
||||
if (typeof status === 'number' && status >= 500) {
|
||||
return 'The assistant is temporarily unavailable. Please try again.';
|
||||
}
|
||||
// Any other failure (including a bare connection error with no status): a
|
||||
// single neutral line. No provider identity, no config, no response body.
|
||||
return 'The assistant could not complete your request. Please try again.';
|
||||
}
|
||||
|
||||
/**
|
||||
* Anonymous, read-only AI assistant for a single PUBLIC share tree.
|
||||
*
|
||||
@@ -318,11 +345,28 @@ export class PublicShareChatService {
|
||||
result.pipeUIMessageStreamToResponse(res.raw, {
|
||||
headers: { 'X-Accel-Buffering': 'no' },
|
||||
onError: (error: unknown) => {
|
||||
// Reuse the shared formatter so provider error formatting stays
|
||||
// unified between the log line and the streamed error message — a
|
||||
// share reader sees 402/429/503 causes consistently with the
|
||||
// authenticated path.
|
||||
return describeProviderError(error, 'AI stream error');
|
||||
// SECURITY (#394): the string this returns is written verbatim into the
|
||||
// SSE error frame delivered to an ANONYMOUS reader (for a tool failure
|
||||
// it becomes the atomic `tool-output-error` frame's errorText; for a
|
||||
// stream/provider failure, the terminal error frame).
|
||||
//
|
||||
// A ShareToolError is already a classified, safe tool message (see
|
||||
// PublicShareChatToolsService.wrapToolErrors) — pass it through so the
|
||||
// reader still gets the useful "page not available in this share" hint.
|
||||
if (error instanceof ShareToolError) {
|
||||
return error.message;
|
||||
}
|
||||
// Anything else is a provider/stream error. describeProviderError
|
||||
// bundles the provider statusCode AND response body, which can carry the
|
||||
// internal baseUrl or model name — NEVER expose that to the public. Log
|
||||
// the full detail server-side only and return a fixed classified string.
|
||||
this.logger.error(
|
||||
`Public share chat pipe error: ${describeProviderError(
|
||||
error,
|
||||
'AI stream error',
|
||||
)}`,
|
||||
);
|
||||
return classifyAnonStreamError(error);
|
||||
},
|
||||
});
|
||||
|
||||
|
||||
@@ -808,7 +808,7 @@ describe('PublicShareChatToolsService share scoping', () => {
|
||||
};
|
||||
|
||||
await expect(getSharePage.execute({ pageId: 'p-outside' })).rejects.toThrow(
|
||||
/not part of this published share/i,
|
||||
/not available in this share/i,
|
||||
);
|
||||
// The tool delegated the resolve to the canonical boundary with the
|
||||
// forShare-scoped shareId, and returned NO content for a non-resolving page.
|
||||
@@ -841,7 +841,7 @@ describe('PublicShareChatToolsService share scoping', () => {
|
||||
|
||||
await expect(
|
||||
getSharePage.execute({ pageId: 'p-restricted' }),
|
||||
).rejects.toThrow(/not part of this published share/i);
|
||||
).rejects.toThrow(/not available in this share/i);
|
||||
// No content was ever sanitized/returned for the blocked page.
|
||||
expect(shareService.updatePublicAttachments).not.toHaveBeenCalled();
|
||||
});
|
||||
@@ -1003,7 +1003,7 @@ describe('public-share assistant boundary locks (red-team regression guards)', (
|
||||
};
|
||||
await expect(
|
||||
getSharePage.execute({ pageId: 'p-elsewhere' }),
|
||||
).rejects.toThrow(/not part of this published share/i);
|
||||
).rejects.toThrow(/not available in this share/i);
|
||||
// The forged share id is the scope the boundary re-derivation rejects against.
|
||||
expect(shareService.resolveReadableSharePage).toHaveBeenCalledWith(
|
||||
'FORGED-SHARE',
|
||||
|
||||
@@ -0,0 +1,160 @@
|
||||
import {
|
||||
wrapInAppToolWithCap,
|
||||
inAppToolCallCapMs,
|
||||
type ToolAbortSignalSink,
|
||||
} from './ai-chat-tools.service';
|
||||
import type { Tool, ToolCallOptions } from 'ai';
|
||||
|
||||
/**
|
||||
* #487 commit 1 — in-app tool race-on-abort + safe-points + per-call cap.
|
||||
*
|
||||
* Tests assert the HONEST observable property the spec names — "after Stop, NO
|
||||
* new HTTP/WS call STARTS; an already-started single call may take either
|
||||
* outcome" — against the REAL wrapper mechanism (the composite abort signal it
|
||||
* publishes on the client + the RACE it runs), NOT a timing-dependent proxy like
|
||||
* "the write didn't land".
|
||||
*/
|
||||
|
||||
// A minimal stand-in for the client's `toolAbortSignal` field. In production the
|
||||
// wrapper publishes the composite here and the client's paginateAll /
|
||||
// mutatePageContent safe-points read it; the fake "tool" below reads it the same
|
||||
// way, so this exercises the real contract without a live DB / collab socket.
|
||||
class FakeClient implements ToolAbortSignalSink {
|
||||
private signal: AbortSignal | null = null;
|
||||
setToolAbortSignal(signal: AbortSignal | null): void {
|
||||
this.signal = signal;
|
||||
}
|
||||
getToolAbortSignal(): AbortSignal | null {
|
||||
return this.signal;
|
||||
}
|
||||
}
|
||||
|
||||
// A ToolCallOptions with just the field the wrapper reads (abortSignal). The AI
|
||||
// SDK passes a fuller object; the wrapper only spreads it and reads abortSignal.
|
||||
const opts = (abortSignal?: AbortSignal): ToolCallOptions =>
|
||||
({ toolCallId: 't1', messages: [], abortSignal }) as unknown as ToolCallOptions;
|
||||
|
||||
const tick = (ms = 5) => new Promise((r) => setTimeout(r, ms));
|
||||
|
||||
describe('#487 wrapInAppToolWithCap — race-on-abort + safe-points', () => {
|
||||
it('after Stop, no NEW simulated call starts (multi-call tool)', async () => {
|
||||
const client = new FakeClient();
|
||||
const started: number[] = [];
|
||||
// A multi-call tool that mirrors paginateAll: it consults the client signal
|
||||
// at a safe-point BEFORE starting each simulated network call.
|
||||
const multiCall: Tool = {
|
||||
execute: (async (_args: unknown) => {
|
||||
for (let i = 0; i < 6; i++) {
|
||||
// Safe-point: exactly what paginateAll / mutatePageContent do.
|
||||
client.getToolAbortSignal()?.throwIfAborted();
|
||||
started.push(i);
|
||||
await tick(10);
|
||||
}
|
||||
return 'done';
|
||||
}) as unknown as Tool['execute'],
|
||||
} as Tool;
|
||||
|
||||
const wrapped = wrapInAppToolWithCap(multiCall, client, 10_000);
|
||||
const ac = new AbortController();
|
||||
const call = (
|
||||
wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise<unknown>
|
||||
)({}, opts(ac.signal));
|
||||
|
||||
// Let one or two calls start, then Stop.
|
||||
await tick(12);
|
||||
ac.abort(new Error('user stop'));
|
||||
|
||||
await expect(call).rejects.toThrow(); // wrapper rejects promptly
|
||||
const startedAtStop = started.length;
|
||||
|
||||
// Give the abandoned loser ample time; its next safe-point must throw because
|
||||
// the (aborted) composite is still published on the client.
|
||||
await tick(60);
|
||||
expect(started.length).toBe(startedAtStop);
|
||||
// It must NOT have run the whole sequence (that would mean Stop did nothing).
|
||||
expect(started.length).toBeLessThan(6);
|
||||
});
|
||||
|
||||
it('rejects immediately on Stop even if the call never settles (discard loser)', async () => {
|
||||
const client = new FakeClient();
|
||||
let settled = false;
|
||||
const hang: Tool = {
|
||||
execute: (async () => {
|
||||
await new Promise(() => undefined); // never resolves
|
||||
settled = true;
|
||||
}) as unknown as Tool['execute'],
|
||||
} as Tool;
|
||||
const wrapped = wrapInAppToolWithCap(hang, client, 10_000);
|
||||
const ac = new AbortController();
|
||||
const call = (
|
||||
wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise<unknown>
|
||||
)({}, opts(ac.signal));
|
||||
await tick(5);
|
||||
ac.abort();
|
||||
await expect(call).rejects.toThrow();
|
||||
expect(settled).toBe(false);
|
||||
});
|
||||
|
||||
it('per-call cap rejects a hung call with no Stop signal', async () => {
|
||||
const client = new FakeClient();
|
||||
const hang: Tool = {
|
||||
execute: (async () => {
|
||||
await new Promise(() => undefined);
|
||||
}) as unknown as Tool['execute'],
|
||||
} as Tool;
|
||||
// Tiny cap; no options.abortSignal at all (composite = cap only).
|
||||
const wrapped = wrapInAppToolWithCap(hang, client, 20);
|
||||
const start = Date.now();
|
||||
await expect(
|
||||
(wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise<unknown>)(
|
||||
{},
|
||||
opts(undefined),
|
||||
),
|
||||
).rejects.toThrow(/per-call cap/);
|
||||
expect(Date.now() - start).toBeLessThan(2000);
|
||||
});
|
||||
|
||||
it('publishes a composite signal on the client for the duration of the call', async () => {
|
||||
const client = new FakeClient();
|
||||
let seenDuringCall: AbortSignal | null = null;
|
||||
const probe: Tool = {
|
||||
execute: (async () => {
|
||||
seenDuringCall = client.getToolAbortSignal();
|
||||
return 'ok';
|
||||
}) as unknown as Tool['execute'],
|
||||
} as Tool;
|
||||
const wrapped = wrapInAppToolWithCap(probe, client, 10_000);
|
||||
const ac = new AbortController();
|
||||
await (
|
||||
wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise<unknown>
|
||||
)({}, opts(ac.signal));
|
||||
expect(seenDuringCall).not.toBeNull();
|
||||
// The published composite must reflect the turn's Stop signal.
|
||||
ac.abort();
|
||||
expect((seenDuringCall as unknown as AbortSignal).aborted).toBe(true);
|
||||
});
|
||||
|
||||
it('a completed call returns its raw result unchanged', async () => {
|
||||
const client = new FakeClient();
|
||||
const ok: Tool = {
|
||||
execute: (async () => ({ items: [1, 2, 3] })) as unknown as Tool['execute'],
|
||||
} as Tool;
|
||||
const wrapped = wrapInAppToolWithCap(ok, client, 10_000);
|
||||
const res = await (
|
||||
wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise<unknown>
|
||||
)({}, opts(new AbortController().signal));
|
||||
expect(res).toEqual({ items: [1, 2, 3] });
|
||||
});
|
||||
|
||||
it('cap is env-tunable with a 2-minute default', () => {
|
||||
const prev = process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS;
|
||||
delete process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS;
|
||||
expect(inAppToolCallCapMs()).toBe(120_000);
|
||||
process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS = '5000';
|
||||
expect(inAppToolCallCapMs()).toBe(5000);
|
||||
process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS = 'not-a-number';
|
||||
expect(inAppToolCallCapMs()).toBe(120_000);
|
||||
if (prev === undefined) delete process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS;
|
||||
else process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS = prev;
|
||||
});
|
||||
});
|
||||
@@ -1,5 +1,5 @@
|
||||
import { Injectable, Logger } from '@nestjs/common';
|
||||
import { tool, type Tool } from 'ai';
|
||||
import { tool, type Tool, type ToolCallOptions } from 'ai';
|
||||
import { z } from 'zod';
|
||||
import { User } from '@docmost/db/types/entity.types';
|
||||
import { TokenService } from '../../auth/services/token.service';
|
||||
@@ -159,6 +159,129 @@ function __assertClientCallContract(client: DocmostClientLike): void {
|
||||
* existing service-account `/mcp` path already calls loopback successfully, so
|
||||
* this works for single-workspace self-host.
|
||||
*/
|
||||
/**
|
||||
* #487: wall-clock cap for a SINGLE in-app tool call, env-tunable via
|
||||
* `AI_CHAT_INAPP_TOOL_CALL_CAP_MS`. Bounds a read tool that would otherwise
|
||||
* paginate for minutes and a content write whose collab commit hangs, and is the
|
||||
* per-call CAP half of the composite abort signal every in-app tool is wrapped
|
||||
* with (the other half is the turn's Stop signal). Default 2 minutes: generous
|
||||
* for a legitimate long read/write, tight enough that a stuck call cannot pin the
|
||||
* turn. The reconcile staleness floor (#487 commit 4) is derived as
|
||||
* max(2 x this cap, 15 min), so keep this well under that.
|
||||
*/
|
||||
export function inAppToolCallCapMs(): number {
|
||||
const raw = Number(process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS);
|
||||
return Number.isFinite(raw) && raw > 0 ? raw : 120_000;
|
||||
}
|
||||
|
||||
/** #487: the composite signal's reason as an Error (informative thrown value). */
|
||||
function inAppAbortReason(signal: AbortSignal): Error {
|
||||
const r = signal.reason;
|
||||
return r instanceof Error
|
||||
? r
|
||||
: new Error(typeof r === 'string' ? r : 'In-app tool call aborted');
|
||||
}
|
||||
|
||||
/**
|
||||
* The client surface {@link wrapInAppToolWithCap} drives (#487). Both methods are
|
||||
* OPTIONAL: the real loopback DocmostClient implements them (so a Stop/cap reaches
|
||||
* its pagination / pre-commit safe-points), but a client that omits them still
|
||||
* gets the OUTER guarantee — the race rejects on abort regardless. This keeps the
|
||||
* wrapper decoupled from the exact client shape (unit-test doubles need not stub
|
||||
* the plumbing).
|
||||
*/
|
||||
export interface ToolAbortSignalSink {
|
||||
setToolAbortSignal?(signal: AbortSignal | null): void;
|
||||
getToolAbortSignal?(): AbortSignal | null;
|
||||
}
|
||||
|
||||
/**
|
||||
* #487: wrap an in-app tool so a Stop (the turn's `options.abortSignal`) OR the
|
||||
* per-call wall-clock cap REJECTS the call immediately, and so that SAME
|
||||
* composite signal reaches the client's pagination / pre-commit safe-points (via
|
||||
* `client.setToolAbortSignal`) — making a Stop stop the NEXT HTTP/WS call from
|
||||
* starting.
|
||||
*
|
||||
* Reuses the RACE pattern of `wrapToolWithCallTimeout` (mcp-clients.service.ts):
|
||||
* the call is raced against the composite signal, so on abort we reject in the
|
||||
* SAME tick and DISCARD the loser promise. Its network / collab teardown latency
|
||||
* therefore never blocks the turn — the supersede timeout W=10s (#487 commit 3)
|
||||
* relies on this abort->settle latency being milliseconds, not a socket teardown.
|
||||
* Awaiting the client's own signal-into-write path alone would NOT satisfy this
|
||||
* (the loser could still be tearing down a collab socket).
|
||||
*
|
||||
* The composite is SET on the client at entry and deliberately NOT restored on
|
||||
* unwind: after this wrapper rejects on abort, the ABANDONED loser promise keeps
|
||||
* running, and its safe-points read the client field — leaving the (aborted)
|
||||
* composite there is exactly what makes the loser's NEXT call throw and stop. The
|
||||
* next in-app tool call overwrites the field with its own fresh composite before
|
||||
* any of its safe-points run, so a stale settled signal never leaks forward.
|
||||
* SINGLE-WRITER by phase-1 assumption — see DocmostClientContext.toolAbortSignal
|
||||
* for the parallel-call caveat (#487).
|
||||
*
|
||||
* KNOWN LIMITATION (#487): a write tool that issues SEVERAL sequential collab
|
||||
* commits can be aborted BETWEEN commits, leaving a partially-applied operation.
|
||||
* Cancel guarantees "no NEW call starts", NOT "the write didn't land".
|
||||
*/
|
||||
export function wrapInAppToolWithCap(
|
||||
toolDef: Tool,
|
||||
client: ToolAbortSignalSink,
|
||||
capMs: number,
|
||||
): Tool {
|
||||
const original = toolDef.execute;
|
||||
if (typeof original !== 'function') return toolDef;
|
||||
const execute = async (args: unknown, options: ToolCallOptions) => {
|
||||
const capController = new AbortController();
|
||||
const timer = setTimeout(() => {
|
||||
capController.abort(
|
||||
new Error(`In-app tool call exceeded the ${capMs}ms per-call cap`),
|
||||
);
|
||||
}, capMs);
|
||||
timer.unref?.();
|
||||
const composite = options?.abortSignal
|
||||
? AbortSignal.any([options.abortSignal, capController.signal])
|
||||
: capController.signal;
|
||||
// Reject the MOMENT the composite fires, independent of whether `original`
|
||||
// ever settles (a hung collab write / read would otherwise pin the turn). The
|
||||
// losing `original` is left pending; Promise.race attaches a rejection
|
||||
// handler to both inputs so a late rejection is never unhandled.
|
||||
const aborted = new Promise<never>((_, reject) => {
|
||||
const fail = () => reject(inAppAbortReason(composite));
|
||||
if (composite.aborted) fail();
|
||||
else composite.addEventListener('abort', fail, { once: true });
|
||||
});
|
||||
// Publish the composite so the client's pagination / pre-commit safe-points
|
||||
// observe it (see the "not restored on unwind" rationale above). Guarded: a
|
||||
// client without the plumbing still gets the OUTER race guarantee below.
|
||||
client.setToolAbortSignal?.(composite);
|
||||
try {
|
||||
return await Promise.race([
|
||||
(original as (a: unknown, o: ToolCallOptions) => Promise<unknown>)(
|
||||
args,
|
||||
{ ...options, abortSignal: composite },
|
||||
),
|
||||
aborted,
|
||||
]);
|
||||
} finally {
|
||||
clearTimeout(timer);
|
||||
}
|
||||
};
|
||||
return { ...toolDef, execute } as unknown as Tool;
|
||||
}
|
||||
|
||||
/** #487: apply {@link wrapInAppToolWithCap} to every tool in a set. */
|
||||
export function wrapInAppToolsWithCap(
|
||||
tools: Record<string, Tool>,
|
||||
client: ToolAbortSignalSink,
|
||||
capMs: number,
|
||||
): Record<string, Tool> {
|
||||
const out: Record<string, Tool> = {};
|
||||
for (const [name, t] of Object.entries(tools)) {
|
||||
out[name] = wrapInAppToolWithCap(t, client, capMs);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
@Injectable()
|
||||
export class AiChatToolsService {
|
||||
private readonly logger = new Logger(AiChatToolsService.name);
|
||||
@@ -186,7 +309,12 @@ export class AiChatToolsService {
|
||||
sessionId: string,
|
||||
workspaceId: string,
|
||||
aiChatId: string,
|
||||
): Promise<DocmostClientLike> {
|
||||
// #487: the returned client also carries the tool-cancellation plumbing
|
||||
// (setToolAbortSignal/getToolAbortSignal). These are host plumbing, NOT part
|
||||
// of the tool-execute surface (DocmostClientMethod), so they are surfaced here
|
||||
// as an intersection rather than by widening that Pick — keeping the
|
||||
// positional-call drift-guard (#446) scoped to the actual tool methods.
|
||||
): Promise<DocmostClientLike & ToolAbortSignalSink> {
|
||||
const apiUrl =
|
||||
process.env.MCP_DOCMOST_API_URL ||
|
||||
`http://127.0.0.1:${process.env.PORT || 3000}/api`;
|
||||
@@ -630,7 +758,15 @@ export class AiChatToolsService {
|
||||
// dependency and reuses the CASL enforcement already on `client`. When the
|
||||
// loaded package predates #417 (factory undefined) or the loader is mocked in
|
||||
// a unit test, signalling is a pure no-op and results are byte-identical.
|
||||
if (!createCommentSignalTracker) return tools;
|
||||
// #487: wrap every in-app tool with the race-on-abort + per-call cap guard so
|
||||
// a Stop / cap rejects immediately AND reaches the client's write/pagination
|
||||
// safe-points. Applied as the OUTERMOST wrapper (over the comment-signal
|
||||
// wrapper below) so the race governs the whole call. The client carries the
|
||||
// per-call composite signal via setToolAbortSignal.
|
||||
const capMs = inAppToolCallCapMs();
|
||||
if (!createCommentSignalTracker) {
|
||||
return wrapInAppToolsWithCap(tools, client, capMs);
|
||||
}
|
||||
|
||||
const tracker = createCommentSignalTracker({
|
||||
probe: async (pageId: string, sinceMs: number) => {
|
||||
@@ -659,7 +795,11 @@ export class AiChatToolsService {
|
||||
},
|
||||
});
|
||||
|
||||
return wrapToolsWithCommentSignal(tools, tracker);
|
||||
return wrapInAppToolsWithCap(
|
||||
wrapToolsWithCommentSignal(tools, tracker),
|
||||
client,
|
||||
capMs,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,7 +1,15 @@
|
||||
import { createHash } from 'node:crypto';
|
||||
import { mkdtempSync, mkdirSync, writeFileSync, rmSync } from 'node:fs';
|
||||
import {
|
||||
mkdtempSync,
|
||||
mkdirSync,
|
||||
writeFileSync,
|
||||
rmSync,
|
||||
readdirSync,
|
||||
statSync,
|
||||
readFileSync,
|
||||
} from 'node:fs';
|
||||
import { tmpdir } from 'node:os';
|
||||
import { join } from 'node:path';
|
||||
import { dirname, join, relative, sep } from 'node:path';
|
||||
|
||||
import { computeSrcRegistryStamp } from './docmost-client.loader';
|
||||
|
||||
@@ -30,10 +38,14 @@ function assertStaleGuard(
|
||||
}
|
||||
}
|
||||
|
||||
// Build a throwaway `<pkg>/build/index.js` + optional `<pkg>/src/tool-specs.ts`
|
||||
// layout so `computeSrcRegistryStamp(<pkg>/build/index.js)` resolves src the same
|
||||
// way the loader does (dirname(dirname(entry))/src/tool-specs.ts).
|
||||
function makeFakePackage(toolSpecsSource: string | null): {
|
||||
// Build a throwaway `<pkg>/build/index.js` + optional `<pkg>/src/` tree so
|
||||
// `computeSrcRegistryStamp(<pkg>/build/index.js)` resolves src the same way the
|
||||
// loader does (dirname(dirname(entry))/src). Since #486 the stamp hashes the WHOLE
|
||||
// src tree, so a fixture is a { relPath: content } map. A bare string is sugar for
|
||||
// a single `tool-specs.ts`; `null` means "no src tree" (the prod no-op path).
|
||||
function makeFakePackage(
|
||||
src: string | Record<string, string> | null,
|
||||
): {
|
||||
entry: string;
|
||||
cleanup: () => void;
|
||||
} {
|
||||
@@ -42,10 +54,15 @@ function makeFakePackage(toolSpecsSource: string | null): {
|
||||
mkdirSync(buildDir, { recursive: true });
|
||||
const entry = join(buildDir, 'index.js');
|
||||
writeFileSync(entry, '// fake @docmost/mcp build entry\n', 'utf8');
|
||||
if (toolSpecsSource !== null) {
|
||||
if (src !== null) {
|
||||
const files =
|
||||
typeof src === 'string' ? { 'tool-specs.ts': src } : src;
|
||||
const srcDir = join(root, 'src');
|
||||
mkdirSync(srcDir, { recursive: true });
|
||||
writeFileSync(join(srcDir, 'tool-specs.ts'), toolSpecsSource, 'utf8');
|
||||
for (const [rel, content] of Object.entries(files)) {
|
||||
const full = join(srcDir, rel);
|
||||
mkdirSync(dirname(full), { recursive: true });
|
||||
writeFileSync(full, content, 'utf8');
|
||||
}
|
||||
}
|
||||
return { entry, cleanup: () => rmSync(root, { recursive: true, force: true }) };
|
||||
}
|
||||
@@ -93,34 +110,109 @@ describe('computeSrcRegistryStamp (#447 stale-build guard)', () => {
|
||||
}
|
||||
});
|
||||
|
||||
// CROSS-IMPL EQUALITY (covers reviewer suggestion 2). The SAME fixed input and
|
||||
// #486 CORE (negative): an edit to a NON-tool-specs src file (client.ts) with a
|
||||
// rebuild NOT run must move the src stamp away from the built REGISTRY_STAMP, so
|
||||
// the loader's stale-check refuses. Under the old tool-specs.ts-only hash this
|
||||
// edit was invisible and a stale build/ served the old client.ts silently.
|
||||
it('a client.ts edit (no rebuild) moves the src stamp -> loader refuses (#486)', () => {
|
||||
// "Built" state: the package as it was compiled.
|
||||
const built = makeFakePackage({
|
||||
'tool-specs.ts': 'export const SPECS = 1;\n',
|
||||
'client.ts': "export const impl = 'v1';\n",
|
||||
});
|
||||
// "Dev edited src, forgot to rebuild": client.ts changed, tool-specs.ts not.
|
||||
const edited = makeFakePackage({
|
||||
'tool-specs.ts': 'export const SPECS = 1;\n',
|
||||
'client.ts': "export const impl = 'v2';\n",
|
||||
});
|
||||
try {
|
||||
const builtStamp = computeSrcRegistryStamp(built.entry);
|
||||
const editedStamp = computeSrcRegistryStamp(edited.entry);
|
||||
expect(builtStamp).not.toBeNull();
|
||||
expect(editedStamp).not.toBe(builtStamp);
|
||||
// build/ still carries builtStamp; src now hashes to editedStamp -> refuse.
|
||||
expect(() => assertStaleGuard(editedStamp, builtStamp as string)).toThrow(
|
||||
STALE_BUILD_MESSAGE,
|
||||
);
|
||||
} finally {
|
||||
built.cleanup();
|
||||
edited.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
// *.generated.ts is excluded (the codegen's own output — a fixed-point cycle
|
||||
// otherwise): its presence/content must not move the stamp.
|
||||
it('excludes *.generated.ts from the stamp', () => {
|
||||
const without = makeFakePackage({ 'tool-specs.ts': 'x\n' });
|
||||
const withGen = makeFakePackage({
|
||||
'tool-specs.ts': 'x\n',
|
||||
'registry-stamp.generated.ts': 'export const REGISTRY_STAMP = "abc";\n',
|
||||
});
|
||||
try {
|
||||
expect(computeSrcRegistryStamp(withGen.entry)).toBe(
|
||||
computeSrcRegistryStamp(without.entry),
|
||||
);
|
||||
} finally {
|
||||
without.cleanup();
|
||||
withGen.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
// CROSS-IMPL EQUALITY (covers reviewer suggestion 2). The SAME fixed tree and
|
||||
// EXPECTED hash are asserted in the mcp-side node test
|
||||
// (packages/mcp/test/unit/registry-stamp.test.mjs) against the codegen's
|
||||
// `computeRegistryStamp`. Asserting the SAME pair here against the loader's
|
||||
// `computeSrcRegistryStamp` proves both implementations normalize+hash
|
||||
// `computeSrcRegistryStamp` proves both implementations enumerate+normalize+hash
|
||||
// identically; a divergence in EITHER side reddens one of the two tests.
|
||||
it('matches the documented cross-impl hash for a fixed input', () => {
|
||||
const FIXED_INPUT = 'line1\r\nline2\n';
|
||||
const EXPECTED =
|
||||
'683376e290829b482c2655745caffa7a1dccfa10afaa62dac2b42dd6c68d0f83';
|
||||
const { entry, cleanup } = makeFakePackage(FIXED_INPUT);
|
||||
const CROSS_IMPL_TREE = {
|
||||
'tool-specs.ts': 'line1\r\nline2\n',
|
||||
'client/read.ts': 'export const R = 1;\n',
|
||||
'registry-stamp.generated.ts': 'export const REGISTRY_STAMP="ignored";\n',
|
||||
};
|
||||
const CROSS_IMPL_EXPECTED =
|
||||
'131c1b9e4e2f5a7d6cef91ca8df619822b442f52bc45ebd09474a4c1d6728616';
|
||||
|
||||
it('matches the documented cross-impl hash for a fixed tree', () => {
|
||||
const { entry, cleanup } = makeFakePackage(CROSS_IMPL_TREE);
|
||||
try {
|
||||
expect(computeSrcRegistryStamp(entry)).toBe(EXPECTED);
|
||||
expect(computeSrcRegistryStamp(entry)).toBe(CROSS_IMPL_EXPECTED);
|
||||
} finally {
|
||||
cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
it('the documented EXPECTED is the normalize+sha256 of the fixed input', () => {
|
||||
// Proves EXPECTED is not a magic constant but the documented computation.
|
||||
const FIXED_INPUT = 'line1\r\nline2\n';
|
||||
const normalized = FIXED_INPUT.replace(/\r\n/g, '\n').replace(/\n$/, '');
|
||||
const expected = createHash('sha256')
|
||||
.update(normalized, 'utf8')
|
||||
.digest('hex');
|
||||
const { entry, cleanup } = makeFakePackage(FIXED_INPUT);
|
||||
it('the documented EXPECTED is the enumerate+normalize+sha256 of the tree', () => {
|
||||
// Proves EXPECTED is not a magic constant but the documented computation — a
|
||||
// local re-implementation of the loader's tree walk.
|
||||
const { entry, cleanup } = makeFakePackage(CROSS_IMPL_TREE);
|
||||
try {
|
||||
expect(computeSrcRegistryStamp(entry)).toBe(expected);
|
||||
const srcDir = join(dirname(dirname(entry)), 'src');
|
||||
const collect = (dir: string): string[] => {
|
||||
const out: string[] = [];
|
||||
for (const e of readdirSync(dir)) {
|
||||
const f = join(dir, e);
|
||||
if (statSync(f).isDirectory()) out.push(...collect(f));
|
||||
else if (e.endsWith('.ts') && !e.endsWith('.generated.ts'))
|
||||
out.push(f);
|
||||
}
|
||||
return out;
|
||||
};
|
||||
const files = collect(srcDir)
|
||||
.map((abs) => ({ rel: relative(srcDir, abs).split(sep).join('/'), abs }))
|
||||
.sort((a, b) => (a.rel < b.rel ? -1 : a.rel > b.rel ? 1 : 0));
|
||||
const h = createHash('sha256');
|
||||
for (const { rel, abs } of files) {
|
||||
const n = readFileSync(abs, 'utf8')
|
||||
.replace(/\r\n/g, '\n')
|
||||
.replace(/\n$/, '');
|
||||
h.update(rel, 'utf8');
|
||||
h.update('\0', 'utf8');
|
||||
h.update(n, 'utf8');
|
||||
h.update('\0', 'utf8');
|
||||
}
|
||||
const localHash = h.digest('hex');
|
||||
expect(computeSrcRegistryStamp(entry)).toBe(localHash);
|
||||
expect(localHash).toBe(CROSS_IMPL_EXPECTED);
|
||||
} finally {
|
||||
cleanup();
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import { createHash } from 'node:crypto';
|
||||
import { existsSync, readFileSync } from 'node:fs';
|
||||
import { dirname, join } from 'node:path';
|
||||
import { existsSync, readdirSync, readFileSync, statSync } from 'node:fs';
|
||||
import { dirname, join, relative, sep } from 'node:path';
|
||||
import { pathToFileURL } from 'node:url';
|
||||
import type { DocmostClient, SharedToolSpec } from '@docmost/mcp';
|
||||
|
||||
@@ -191,33 +191,52 @@ interface DocmostMcpModule {
|
||||
* present. Returns the stamp string, or `null` when the source is absent (a prod
|
||||
* image ships only build/, no src/). MUST stay byte-for-byte identical to
|
||||
* packages/mcp/scripts/gen-registry-stamp.mjs's `computeRegistryStamp` so the
|
||||
* build-time and src-time hashes agree: same input file (src/tool-specs.ts), same
|
||||
* normalization (CRLF -> LF, strip a single trailing newline), same sha256.
|
||||
* build-time and src-time hashes agree: same file set (every src/**\/*.ts except
|
||||
* *.generated.ts), same POSIX-relative sort, same per-file normalization (CRLF ->
|
||||
* LF, strip a single trailing newline) with the same path+content framing, same
|
||||
* sha256. Hashing the WHOLE src tree (not just tool-specs.ts) is #486: an edit to
|
||||
* client.ts / a client/* module / comment-signal / drawio-* without a rebuild
|
||||
* must also be caught, otherwise build/ silently serves the old code.
|
||||
*
|
||||
* DEV vs PROD detection is by FILE EXISTENCE, not NODE_ENV: we resolve the
|
||||
* package's own directory from `require.resolve('@docmost/mcp')` (which points at
|
||||
* build/index.js) and look for ../src/tool-specs.ts next to it. In a dev/test
|
||||
* worktree that file exists; in a prod image (build/ only, src/ stripped) it does
|
||||
* not, so this returns null and the caller skips the check. Any error (ENOENT, a
|
||||
* bad resolve) is swallowed to null — the stale-check must NEVER break startup.
|
||||
* build/index.js) and look for ../src next to it. In a dev/test worktree that
|
||||
* directory exists; in a prod image (build/ only, src/ stripped) it does not, so
|
||||
* this returns null and the caller skips the check. Any error (ENOENT, a bad
|
||||
* resolve) is swallowed to null — the stale-check must NEVER break startup.
|
||||
*
|
||||
* Exported for unit testing (docmost-client.loader.spec.ts): the export keyword
|
||||
* is behaviourally a no-op — the module-internal caller `loadDocmostMcp` is
|
||||
* unaffected. The test drives the null (no-src) path and asserts this
|
||||
* normalize+sha256 stays identical to the codegen's `computeRegistryStamp`.
|
||||
* enumerate+normalize+sha256 stays identical to the codegen's
|
||||
* `computeRegistryStamp`.
|
||||
*/
|
||||
export function computeSrcRegistryStamp(packageEntry: string): string | null {
|
||||
try {
|
||||
// packageEntry is <pkg>/build/index.js; the source lives at <pkg>/src/.
|
||||
const toolSpecsPath = join(
|
||||
dirname(dirname(packageEntry)),
|
||||
'src',
|
||||
'tool-specs.ts',
|
||||
);
|
||||
if (!existsSync(toolSpecsPath)) return null; // prod: no src tree -> skip.
|
||||
const source = readFileSync(toolSpecsPath, 'utf8');
|
||||
const normalized = source.replace(/\r\n/g, '\n').replace(/\n$/, '');
|
||||
return createHash('sha256').update(normalized, 'utf8').digest('hex');
|
||||
const srcDir = join(dirname(dirname(packageEntry)), 'src');
|
||||
if (!existsSync(srcDir)) return null; // prod: no src tree -> skip.
|
||||
// Enumerate every src/**\/*.ts except the codegen's own *.generated.ts
|
||||
// output (including it would be a fixed-point cycle). Sort by POSIX-relative
|
||||
// path so ordering is platform-independent, then fold each file's relative
|
||||
// path + normalized content into one hash — identical to the codegen.
|
||||
const files = collectStampFiles(srcDir)
|
||||
.map((abs) => ({
|
||||
rel: relative(srcDir, abs).split(sep).join('/'),
|
||||
abs,
|
||||
}))
|
||||
.sort((a, b) => (a.rel < b.rel ? -1 : a.rel > b.rel ? 1 : 0));
|
||||
const hash = createHash('sha256');
|
||||
for (const { rel, abs } of files) {
|
||||
const normalized = readFileSync(abs, 'utf8')
|
||||
.replace(/\r\n/g, '\n')
|
||||
.replace(/\n$/, '');
|
||||
hash.update(rel, 'utf8');
|
||||
hash.update('\0', 'utf8');
|
||||
hash.update(normalized, 'utf8');
|
||||
hash.update('\0', 'utf8');
|
||||
}
|
||||
return hash.digest('hex');
|
||||
} catch {
|
||||
// Never let a resolution/read hiccup break server startup — treat as "no
|
||||
// src available" and skip the check (identical to the prod no-op path).
|
||||
@@ -225,6 +244,24 @@ export function computeSrcRegistryStamp(packageEntry: string): string | null {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Recursively enumerate every `*.ts` under `dir`, EXCLUDING `*.generated.ts`.
|
||||
* Mirror of the codegen's `collectStampFiles` (packages/mcp/scripts/
|
||||
* gen-registry-stamp.mjs) — keep the two walk/filter rules identical.
|
||||
*/
|
||||
function collectStampFiles(dir: string): string[] {
|
||||
const out: string[] = [];
|
||||
for (const entry of readdirSync(dir)) {
|
||||
const full = join(dir, entry);
|
||||
if (statSync(full).isDirectory()) {
|
||||
out.push(...collectStampFiles(full));
|
||||
} else if (entry.endsWith('.ts') && !entry.endsWith('.generated.ts')) {
|
||||
out.push(full);
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
// TS with module:commonjs downlevels a literal `import()` to `require()`, which
|
||||
// cannot load the ESM-only `@docmost/mcp` package. Indirect through Function so
|
||||
// the real dynamic `import()` survives compilation and can load ESM from
|
||||
|
||||
@@ -224,7 +224,7 @@ describe('PublicShareChatToolsService.forShare', () => {
|
||||
(tools.getSharePage as unknown as ToolExec).execute({
|
||||
pageId: 'page-1',
|
||||
}),
|
||||
).rejects.toThrow('That page is not part of this published share.');
|
||||
).rejects.toThrow('The requested page is not available in this share.');
|
||||
|
||||
// No content is ever fetched/returned for a non-resolving page.
|
||||
expect(shareService.updatePublicAttachments).not.toHaveBeenCalled();
|
||||
|
||||
@@ -7,6 +7,22 @@ import { PageRepo } from '@docmost/db/repos/page/page.repo';
|
||||
import { jsonToMarkdown } from '../../../collaboration/collaboration.util';
|
||||
import { modelFriendlyInput } from './model-friendly-input';
|
||||
|
||||
/**
|
||||
* A tool error whose message is DELIBERATELY safe to expose to an anonymous
|
||||
* share reader (and to the model, for self-correction). Every OTHER thrown error
|
||||
* is treated as internal and replaced with a generic string by `wrapToolErrors`,
|
||||
* so a raw exception message — an internal page title, a DB/stack fragment, a
|
||||
* driver detail — never rides the public UI stream (#394).
|
||||
*/
|
||||
export class ShareToolError extends Error {}
|
||||
|
||||
// The only two classified strings an anonymous reader may ever see from a tool
|
||||
// failure. The specific one keeps the model's self-correction useful ("try a
|
||||
// different page"); the generic one reveals nothing about the internal fault.
|
||||
const SHARE_TOOL_ERROR_NOT_AVAILABLE =
|
||||
'The requested page is not available in this share.';
|
||||
const SHARE_TOOL_ERROR_GENERIC = 'The tool could not complete the request.';
|
||||
|
||||
/**
|
||||
* Isolated, READ-ONLY toolset for the ANONYMOUS public-share assistant.
|
||||
*
|
||||
@@ -44,7 +60,7 @@ export class PublicShareChatToolsService {
|
||||
* are NO write tools, NO comments/history, NO cross-space or external tools.
|
||||
*/
|
||||
forShare(shareId: string, workspaceId: string): Record<string, Tool> {
|
||||
return {
|
||||
return this.wrapToolErrors({
|
||||
searchSharePages: tool({
|
||||
description:
|
||||
'Search the pages of THIS published documentation share for a ' +
|
||||
@@ -96,7 +112,7 @@ export class PublicShareChatToolsService {
|
||||
execute: async ({ pageId }) => {
|
||||
const id = (pageId ?? '').trim();
|
||||
if (!id) {
|
||||
throw new Error('A pageId is required.');
|
||||
throw new ShareToolError('A pageId is required.');
|
||||
}
|
||||
// Resolve via the SINGLE canonical share-access boundary: confirms the
|
||||
// page resolves to THIS share (recursive CTE up the tree, honouring
|
||||
@@ -112,7 +128,7 @@ export class PublicShareChatToolsService {
|
||||
workspaceId,
|
||||
);
|
||||
if (!resolved) {
|
||||
throw new Error('That page is not part of this published share.');
|
||||
throw new ShareToolError(SHARE_TOOL_ERROR_NOT_AVAILABLE);
|
||||
}
|
||||
const { page } = resolved;
|
||||
|
||||
@@ -193,6 +209,57 @@ export class PublicShareChatToolsService {
|
||||
}
|
||||
},
|
||||
}),
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Wrap every tool's `execute` so a THROWN error is sanitized in ONE place —
|
||||
* closing the byte leak, the render, and the model context at once (#394).
|
||||
*
|
||||
* The AI SDK surfaces a tool-execution throw as an atomic `tool-output-error`
|
||||
* frame on the v6 UI stream whose `errorText` is the thrown message; on the
|
||||
* public share that frame goes straight to an anonymous reader. Unwrapped, a
|
||||
* raw exception (an internal page title, a DB/stack fragment, a driver detail)
|
||||
* would ride that frame verbatim. Here we catch it, LOG the full detail
|
||||
* server-side only, and re-throw a CLASSIFIED, safe error: the tool's own
|
||||
* intentional ShareToolError messages pass through (they keep the model's
|
||||
* self-correction useful), everything else collapses to a generic string.
|
||||
*/
|
||||
private wrapToolErrors(
|
||||
tools: Record<string, Tool>,
|
||||
): Record<string, Tool> {
|
||||
const wrapped: Record<string, Tool> = {};
|
||||
for (const [name, t] of Object.entries(tools)) {
|
||||
const original = t.execute;
|
||||
if (typeof original !== 'function') {
|
||||
wrapped[name] = t;
|
||||
continue;
|
||||
}
|
||||
wrapped[name] = {
|
||||
...t,
|
||||
execute: async (args: unknown, options: unknown) => {
|
||||
try {
|
||||
return await (
|
||||
original as (a: unknown, o: unknown) => Promise<unknown>
|
||||
)(args, options);
|
||||
} catch (err) {
|
||||
const safe =
|
||||
err instanceof ShareToolError
|
||||
? err.message
|
||||
: SHARE_TOOL_ERROR_GENERIC;
|
||||
// Full detail to the server log ONLY — never to the anon.
|
||||
this.logger.warn(
|
||||
`Public share tool "${name}" failed: ${
|
||||
err instanceof Error ? err.message : String(err)
|
||||
}`,
|
||||
);
|
||||
// This safe string is ALL that rides the tool-output-error frame,
|
||||
// becomes model context, and could be rendered — one choke point.
|
||||
throw new ShareToolError(safe);
|
||||
}
|
||||
},
|
||||
} as Tool;
|
||||
}
|
||||
return wrapped;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -120,3 +120,102 @@ describe('JwtStrategy — provenance derivation', () => {
|
||||
expect(req.raw.actor).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* Provenance derivation on the API-KEY path (jwt.strategy.validateApiKey, #486).
|
||||
*
|
||||
* The access-token path stamped provenance; the API-key path returned early
|
||||
* WITHOUT it, so an is_agent API key's REST writes recorded no 'agent' marker.
|
||||
* The API-key payload carries no signed claim, so provenance is resolved from the
|
||||
* SERVER-SIDE user returned by ApiKeyService.validateApiKey: isAgent -> 'agent',
|
||||
* otherwise 'user'; aiChatId is always null (an API key has no ai_chats row).
|
||||
*
|
||||
* The enterprise ApiKeyService is not bundled in the OSS build, so the strategy
|
||||
* loads it through an overridable `resolveApiKeyService` seam that we stub here.
|
||||
*/
|
||||
describe('JwtStrategy — API-key provenance derivation (#486)', () => {
|
||||
function makeApiKeyStrategy(validateApiKeyImpl: (p: any) => Promise<any>) {
|
||||
const userRepo: any = { findById: jest.fn() };
|
||||
const workspaceRepo: any = { findById: jest.fn() };
|
||||
const userSessionRepo: any = { findActiveById: jest.fn() };
|
||||
const sessionActivityService: any = { trackActivity: jest.fn() };
|
||||
const environmentService: any = { getAppSecret: () => 'test-secret' };
|
||||
const moduleRef: any = {};
|
||||
|
||||
const strategy = new JwtStrategy(
|
||||
userRepo,
|
||||
workspaceRepo,
|
||||
userSessionRepo,
|
||||
sessionActivityService,
|
||||
environmentService,
|
||||
moduleRef,
|
||||
);
|
||||
// Stub the EE ApiKeyService seam (the real module is not in the OSS build).
|
||||
const validateApiKey = jest.fn(validateApiKeyImpl);
|
||||
jest
|
||||
.spyOn(strategy as any, 'resolveApiKeyService')
|
||||
.mockReturnValue({ validateApiKey });
|
||||
return { strategy, validateApiKey };
|
||||
}
|
||||
|
||||
const makeReq = () => ({ raw: {} as Record<string, any> });
|
||||
const apiKeyPayload = () => ({
|
||||
sub: 'svc-1',
|
||||
workspaceId: 'ws-1',
|
||||
apiKeyId: 'key-1',
|
||||
type: JwtType.API_KEY,
|
||||
});
|
||||
|
||||
it("stamps actor='agent' for an is_agent API key (from the validated user)", async () => {
|
||||
const validated = {
|
||||
user: { id: 'svc-1', isAgent: true },
|
||||
workspace: { id: 'ws-1' },
|
||||
};
|
||||
const { strategy, validateApiKey } = makeApiKeyStrategy(
|
||||
async () => validated,
|
||||
);
|
||||
const req = makeReq();
|
||||
|
||||
const result = await strategy.validate(req, apiKeyPayload() as any);
|
||||
|
||||
expect(validateApiKey).toHaveBeenCalledTimes(1);
|
||||
expect(req.raw.actor).toBe('agent');
|
||||
// API keys carry no internal ai_chats row -> null.
|
||||
expect(req.raw.aiChatId).toBeNull();
|
||||
// The validated auth object is returned unchanged (req.user shape preserved).
|
||||
expect(result).toBe(validated);
|
||||
});
|
||||
|
||||
it("stamps actor='user' for an ordinary (non-agent) API key", async () => {
|
||||
const { strategy } = makeApiKeyStrategy(async () => ({
|
||||
user: { id: 'u-1', isAgent: false },
|
||||
workspace: { id: 'ws-1' },
|
||||
}));
|
||||
const req = makeReq();
|
||||
|
||||
await strategy.validate(req, apiKeyPayload() as any);
|
||||
|
||||
expect(req.raw.actor).toBe('user');
|
||||
expect(req.raw.aiChatId).toBeNull();
|
||||
});
|
||||
|
||||
it('throws Unauthorized (and stamps nothing) when the EE module is missing', async () => {
|
||||
const userRepo: any = { findById: jest.fn() };
|
||||
const strategy = new JwtStrategy(
|
||||
userRepo,
|
||||
{ findById: jest.fn() } as any,
|
||||
{ findActiveById: jest.fn() } as any,
|
||||
{ trackActivity: jest.fn() } as any,
|
||||
{ getAppSecret: () => 'test-secret' } as any,
|
||||
{} as any,
|
||||
);
|
||||
// EE not bundled: the seam returns null.
|
||||
jest.spyOn(strategy as any, 'resolveApiKeyService').mockReturnValue(null);
|
||||
const req = makeReq();
|
||||
|
||||
await expect(
|
||||
strategy.validate(req, apiKeyPayload() as any),
|
||||
).rejects.toThrow(UnauthorizedException);
|
||||
expect(req.raw.actor).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -102,28 +102,49 @@ export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
|
||||
}
|
||||
|
||||
private async validateApiKey(req: any, payload: JwtApiKeyPayload) {
|
||||
let ApiKeyModule: any;
|
||||
let isApiKeyModuleReady = false;
|
||||
const apiKeyService = this.resolveApiKeyService();
|
||||
if (!apiKeyService) {
|
||||
throw new UnauthorizedException('Enterprise API Key module missing');
|
||||
}
|
||||
|
||||
const result = await apiKeyService.validateApiKey(payload);
|
||||
|
||||
// Stamp the agent-edit provenance for the API-KEY path too (#486). Unlike the
|
||||
// access-token path above, it CANNOT be resolved before this point: the
|
||||
// API-key payload carries no signed actor/aiChatId claim, and the user (with
|
||||
// its isAgent flag) is unknown until the key is validated. Claim semantics for
|
||||
// API keys: an is_agent API key (an agent service account) stamps 'agent' on
|
||||
// every REST write; an ordinary API key resolves to 'user'. An API key has no
|
||||
// internal ai_chats row, so aiChatId is always null. Derived from the
|
||||
// SERVER-SIDE user (never a client field), so an 'agent' badge is unspoofable
|
||||
// — mirroring the access-token path. Passing `null` for the claim means the
|
||||
// actor is decided solely by user.isAgent.
|
||||
const provenance = resolveProvenance((result as any)?.user, null);
|
||||
req.raw.actor = provenance.actor;
|
||||
req.raw.aiChatId = provenance.aiChatId;
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve the enterprise ApiKeyService, or `null` when the EE module is not
|
||||
* bundled in this build (community build). Extracted as an overridable seam so
|
||||
* the API-key provenance stamping can be unit-tested without the EE package
|
||||
* present (docmost is OSS + a separate EE bundle; `require` of the EE path
|
||||
* throws here). Any load/resolve error is treated as "module missing".
|
||||
*/
|
||||
protected resolveApiKeyService(): {
|
||||
validateApiKey: (payload: JwtApiKeyPayload) => Promise<unknown>;
|
||||
} | null {
|
||||
try {
|
||||
// eslint-disable-next-line @typescript-eslint/no-require-imports
|
||||
ApiKeyModule = require('./../../../ee/api-key/api-key.service');
|
||||
isApiKeyModuleReady = true;
|
||||
const ApiKeyModule = require('./../../../ee/api-key/api-key.service');
|
||||
return this.moduleRef.get(ApiKeyModule.ApiKeyService, { strict: false });
|
||||
} catch (err) {
|
||||
this.logger.debug(
|
||||
'API Key module requested but enterprise module not bundled in this build',
|
||||
);
|
||||
isApiKeyModuleReady = false;
|
||||
return null;
|
||||
}
|
||||
|
||||
if (isApiKeyModuleReady) {
|
||||
const ApiKeyService = this.moduleRef.get(ApiKeyModule.ApiKeyService, {
|
||||
strict: false,
|
||||
});
|
||||
|
||||
return ApiKeyService.validateApiKey(payload);
|
||||
}
|
||||
|
||||
throw new UnauthorizedException('Enterprise API Key module missing');
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
import { type Kysely, sql } from 'kysely';
|
||||
|
||||
export async function up(db: Kysely<any>): Promise<void> {
|
||||
// Chat-level metadata bag (#490). First use: the deferred-tool ACTIVATION set
|
||||
// (`activatedTools`) is persisted here so it survives across turns — previously
|
||||
// the set was reset every turn, forcing the model to re-run loadTools and pay a
|
||||
// fresh round-trip to re-activate the same tools each turn. On load the stored
|
||||
// set is intersected with the current valid deferred names, so an allowlist /
|
||||
// role change can never inject a now-nonexistent tool.
|
||||
//
|
||||
// jsonb, defaulted to '{}' so every row (incl. pre-migration ones, backfilled
|
||||
// by the default) is a readable object — the app never has to null-guard the
|
||||
// bag itself, only individual keys.
|
||||
await db.schema
|
||||
.alterTable('ai_chats')
|
||||
.addColumn('metadata', 'jsonb', (col) =>
|
||||
col.notNull().defaultTo(sql`'{}'::jsonb`),
|
||||
)
|
||||
.execute();
|
||||
}
|
||||
|
||||
export async function down(db: Kysely<any>): Promise<void> {
|
||||
await db.schema.alterTable('ai_chats').dropColumn('metadata').execute();
|
||||
}
|
||||
@@ -1,5 +1,6 @@
|
||||
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 {
|
||||
@@ -188,6 +189,144 @@ export class AiChatMessageRepo {
|
||||
return query.returning(this.baseFields).executeTakeFirst();
|
||||
}
|
||||
|
||||
/**
|
||||
* #487 OWNER terminal write — the streamText terminal callback's finalize. Like
|
||||
* `update` but CONDITIONAL on `status='streaming' OR metadata.finalizeFailed`:
|
||||
* the owner writes its real content EITHER when the row is still streaming (the
|
||||
* normal case) OR when a reconcile stamp already flipped it to a terminal status
|
||||
* but marked `finalizeFailed:true` — the owner's real content OVERWRITES that
|
||||
* placeholder stamp (owner-write priority, #487). A row that is properly terminal
|
||||
* (no finalizeFailed) is left untouched (undefined) — idempotent. The `patch`
|
||||
* carries the real metadata WITHOUT finalizeFailed, so a successful write CLEARS
|
||||
* the flag. Returns the updated row, or undefined when nothing matched.
|
||||
*/
|
||||
async finalizeOwner(
|
||||
id: string,
|
||||
workspaceId: string,
|
||||
patch: Partial<{
|
||||
content: string | null;
|
||||
toolCalls: unknown;
|
||||
metadata: unknown;
|
||||
status: string | null;
|
||||
}>,
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<AiChatMessage | undefined> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
return db
|
||||
.updateTable('aiChatMessages')
|
||||
.set({ ...(patch as Record<string, unknown>), updatedAt: new Date() })
|
||||
.where('id', '=', id)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.where((eb) =>
|
||||
eb.or([
|
||||
eb('status', '=', 'streaming'),
|
||||
eb(sql<string>`(metadata->>'finalizeFailed')`, '=', 'true'),
|
||||
]),
|
||||
)
|
||||
.returning(this.baseFields)
|
||||
.executeTakeFirst();
|
||||
}
|
||||
|
||||
/**
|
||||
* #487 RECONCILE status-only stamp — settle a stuck 'streaming' row to a
|
||||
* terminal status WITHOUT the owner's real content (which lived only in the
|
||||
* dead process's memory — a documented loss). CONDITIONAL on `status='streaming'`
|
||||
* (never touches an already-terminal row) AND it MERGES `finalizeFailed:true`
|
||||
* into metadata (preserving the partial `parts` already persisted) so a LATER
|
||||
* owner-write (finalizeOwner) can still OVERWRITE this placeholder with real
|
||||
* content, and so `isInterruptResume` can EXCLUDE this row (a reconcile stamp is
|
||||
* not a genuine user interruption). Returns the updated row, or undefined.
|
||||
*/
|
||||
async stampTerminalIfStreaming(
|
||||
id: string,
|
||||
workspaceId: string,
|
||||
status: 'aborted' | 'error' | 'completed',
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<AiChatMessage | undefined> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
return db
|
||||
.updateTable('aiChatMessages')
|
||||
.set({
|
||||
status,
|
||||
metadata: sql`coalesce(metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`,
|
||||
updatedAt: new Date(),
|
||||
})
|
||||
.where('id', '=', id)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.where('status', '=', 'streaming')
|
||||
.returning(this.baseFields)
|
||||
.executeTakeFirst();
|
||||
}
|
||||
|
||||
/**
|
||||
* #487 reconcile clause (b): streaming assistant rows whose linked RUN has
|
||||
* already reached a terminal status — an asymmetry ("run settled / message
|
||||
* streaming forever") the periodic reconcile heals by stamping the message.
|
||||
* Returns the message id + its run's terminal status, bounded.
|
||||
*/
|
||||
async findStreamingWithTerminalRun(
|
||||
limit = 200,
|
||||
// #487: scope to ONE chat for the opportunistic per-turn reconcile (removes
|
||||
// reconcile latency from the user-visible path); omit for the periodic sweep.
|
||||
chat?: { chatId: string; workspaceId: string },
|
||||
): Promise<
|
||||
Array<{ messageId: string; workspaceId: string; runStatus: string }>
|
||||
> {
|
||||
let query = this.db
|
||||
.selectFrom('aiChatMessages as m')
|
||||
.innerJoin('aiChatRuns as r', 'r.assistantMessageId', 'm.id')
|
||||
.select([
|
||||
'm.id as messageId',
|
||||
'm.workspaceId as workspaceId',
|
||||
'r.status as runStatus',
|
||||
])
|
||||
.where('m.status', '=', 'streaming')
|
||||
.where('r.status', 'in', ['succeeded', 'failed', 'aborted']);
|
||||
if (chat) {
|
||||
query = query
|
||||
.where('m.chatId', '=', chat.chatId)
|
||||
.where('m.workspaceId', '=', chat.workspaceId);
|
||||
}
|
||||
return query.limit(limit).execute();
|
||||
}
|
||||
|
||||
/**
|
||||
* #487 reconcile clause (d) — historical-row safety: streaming rows older than
|
||||
* `staleMs` whose chat has NO active run row (double-gated). Settle them to
|
||||
* 'aborted' + finalizeFailed (so a late owner-write could still overwrite).
|
||||
* Returns the count. Used ONLY by the periodic reconcile, never at boot.
|
||||
*/
|
||||
async sweepStreamingWithoutActiveRun(
|
||||
staleMs: number,
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<number> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
const staleBefore = new Date(Date.now() - staleMs);
|
||||
const rows = await db
|
||||
.updateTable('aiChatMessages as m')
|
||||
.set({
|
||||
status: 'aborted',
|
||||
metadata: sql`coalesce(m.metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`,
|
||||
updatedAt: new Date(),
|
||||
})
|
||||
.where('m.status', '=', 'streaming')
|
||||
.where('m.updatedAt', '<', staleBefore)
|
||||
.where((eb) =>
|
||||
eb.not(
|
||||
eb.exists(
|
||||
eb
|
||||
.selectFrom('aiChatRuns as r')
|
||||
.select('r.id')
|
||||
.whereRef('r.chatId', '=', 'm.chatId')
|
||||
.where('r.status', 'in', ['pending', 'running']),
|
||||
),
|
||||
),
|
||||
)
|
||||
.returning('m.id')
|
||||
.execute();
|
||||
return rows.length;
|
||||
}
|
||||
|
||||
/**
|
||||
* Crash-recovery sweep (#183): flip every assistant row still left in the
|
||||
* 'streaming' state (a turn that died mid-write before reaching a terminal
|
||||
@@ -200,13 +339,20 @@ export class AiChatMessageRepo {
|
||||
* step, so an actively-streaming row never matches; this prevents a fresh
|
||||
* replica's boot-sweep from aborting a turn another replica is still streaming
|
||||
* in a multi-instance deploy.
|
||||
*
|
||||
* #487: the sweep now ALSO marks `finalizeFailed:true` so a late owner-write can
|
||||
* overwrite this placeholder with real content (owner-write priority).
|
||||
*/
|
||||
async sweepStreaming(trx?: KyselyTransaction): Promise<number> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
const staleBefore = new Date(Date.now() - SWEEP_STREAMING_STALE_MS);
|
||||
const rows = await db
|
||||
.updateTable('aiChatMessages')
|
||||
.set({ status: 'aborted', updatedAt: new Date() })
|
||||
.set({
|
||||
status: 'aborted',
|
||||
metadata: sql`coalesce(metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`,
|
||||
updatedAt: new Date(),
|
||||
})
|
||||
.where('status', '=', 'streaming')
|
||||
.where('updatedAt', '<', staleBefore)
|
||||
.returning('id')
|
||||
|
||||
@@ -143,6 +143,41 @@ export class AiChatRunRepo {
|
||||
.executeTakeFirst();
|
||||
}
|
||||
|
||||
/**
|
||||
* #487: CONDITIONAL terminal finalize — flip a run to a terminal status and
|
||||
* stamp `finished_at` ONLY while it is still active (pending|running), mirroring
|
||||
* the assistant message's `onlyIfStreaming` guard. A double-settle (a late or
|
||||
* second writer, a supersede applying a zombie's intended, a reconcile stamp)
|
||||
* matches NOTHING once the row is terminal and is a benign no-op — so a terminal
|
||||
* status can never be clobbered by a later writer (last-writer-wins is gone).
|
||||
*
|
||||
* Returns the updated row when it WAS active (this call wrote it), else
|
||||
* undefined (the row was already terminal — another writer won). The caller
|
||||
* distinguishes the two to resolve the correct settle outcome.
|
||||
*/
|
||||
async finalizeIfActive(
|
||||
id: string,
|
||||
workspaceId: string,
|
||||
patch: { status: string; error: string | null },
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<AiChatRun | undefined> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
const now = new Date();
|
||||
return db
|
||||
.updateTable('aiChatRuns')
|
||||
.set({
|
||||
status: patch.status,
|
||||
error: patch.error,
|
||||
finishedAt: now,
|
||||
updatedAt: now,
|
||||
})
|
||||
.where('id', '=', id)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[])
|
||||
.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
|
||||
@@ -184,6 +219,31 @@ export class AiChatRunRepo {
|
||||
* sweeps only runs UNTOUCHED past the window. Phase 1 is single-process, so the
|
||||
* boot path supplies no window.
|
||||
*/
|
||||
/**
|
||||
* #487 reconcile clause (c): active (pending|running) runs UNTOUCHED past
|
||||
* `staleMs` — candidates for "no live runner" abort. Staleness is measured from
|
||||
* `updated_at` (the LAST-PROGRESS timestamp — recordStep bumps it), NOT
|
||||
* `started_at`, so a legitimate long-running marathon (11–25 min of steady
|
||||
* progress) is never a candidate. The caller filters these against its in-memory
|
||||
* `active` / zombie maps ("no entry" is the PRIMARY gate — a live entry is never
|
||||
* aborted) before settling any of them. Bounded.
|
||||
*/
|
||||
async findStaleActive(
|
||||
staleMs: number,
|
||||
limit = 200,
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<Array<{ id: string; workspaceId: string; chatId: string }>> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
const staleBefore = new Date(Date.now() - staleMs);
|
||||
return db
|
||||
.selectFrom('aiChatRuns')
|
||||
.select(['id', 'workspaceId', 'chatId'])
|
||||
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[])
|
||||
.where('updatedAt', '<', staleBefore)
|
||||
.limit(limit)
|
||||
.execute();
|
||||
}
|
||||
|
||||
async sweepRunning(
|
||||
opts: { staleMs?: number } = {},
|
||||
trx?: KyselyTransaction,
|
||||
|
||||
+3
@@ -606,6 +606,9 @@ export interface AiChats {
|
||||
// The document the chat was created in (open page at first message). NULL =>
|
||||
// started outside any document. ON DELETE SET NULL on the page FK.
|
||||
pageId: string | null;
|
||||
// Chat-level metadata bag (#490). jsonb, defaulted to '{}'. First key:
|
||||
// `activatedTools` — the deferred-tool activation set persisted across turns.
|
||||
metadata: Generated<Json>;
|
||||
createdAt: Generated<Timestamp>;
|
||||
updatedAt: Generated<Timestamp>;
|
||||
deletedAt: Timestamp | null;
|
||||
|
||||
@@ -0,0 +1,133 @@
|
||||
import { readFileSync } from 'fs';
|
||||
import { EventEmitter } from 'node:events';
|
||||
import { streamText } from 'ai';
|
||||
import { MockLanguageModelV3, simulateReadableStream } from 'ai/test';
|
||||
|
||||
/**
|
||||
* Regression tests for the writeToServerResponse drain-hang fix in
|
||||
* patches/ai@6.0.134.patch (#486, commit 6).
|
||||
*
|
||||
* Unpatched ai@6.0.134's writeToServerResponse awaits ONLY `once("drain")` when
|
||||
* response.write() returns false (backpressure). If the client disconnects
|
||||
* mid-write the socket never drains, so that await never resolves: the read loop
|
||||
* parks FOREVER, its `finally { response.end() }` is unreachable, and the stream
|
||||
* reader + buffered chunks are pinned until process restart. In autonomous mode
|
||||
* the run keeps producing output after the disconnect, so EVERY mid-run
|
||||
* disconnect leaks a hung pipe. The patch races drain against close/error, and on
|
||||
* a terminal socket event cancels the reader and breaks so `finally` always runs.
|
||||
*
|
||||
* This drives the REAL patched writeToServerResponse through the public
|
||||
* pipeUIMessageStreamToResponse API with a response that never drains and closes
|
||||
* mid-write — exactly the leak scenario.
|
||||
*/
|
||||
|
||||
/** A ServerResponse-like emitter whose first write() stalls (returns false) and
|
||||
* then "closes" like a disconnecting client — never firing 'drain'. */
|
||||
class DisconnectingResponse extends EventEmitter {
|
||||
ended = false;
|
||||
writeCount = 0;
|
||||
statusCode = 200;
|
||||
writableEnded = false;
|
||||
destroyed = false;
|
||||
writeHead(): this {
|
||||
return this;
|
||||
}
|
||||
setHeader(): void {}
|
||||
flushHeaders(): void {}
|
||||
write(): boolean {
|
||||
this.writeCount++;
|
||||
if (this.writeCount === 1) {
|
||||
// Simulate the client vanishing mid-write: backpressure (false) and then a
|
||||
// 'close' on the next tick, and CRUCIALLY never a 'drain'. Unpatched, the
|
||||
// loop would await drain forever here.
|
||||
setImmediate(() => this.emit('close'));
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
end(): void {
|
||||
this.ended = true;
|
||||
this.writableEnded = true;
|
||||
this.emit('finish');
|
||||
}
|
||||
}
|
||||
|
||||
function makeModel() {
|
||||
return new MockLanguageModelV3({
|
||||
doStream: async () => ({
|
||||
stream: simulateReadableStream({
|
||||
chunks: [
|
||||
{ type: 'stream-start' as const, warnings: [] },
|
||||
{ type: 'text-start' as const, id: '1' },
|
||||
{ type: 'text-delta' as const, id: '1', delta: 'hello ' },
|
||||
{ type: 'text-delta' as const, id: '1', delta: 'world' },
|
||||
{ type: 'text-end' as const, id: '1' },
|
||||
{
|
||||
type: 'finish' as const,
|
||||
finishReason: { unified: 'stop' as const, raw: 'stop' },
|
||||
usage: {
|
||||
inputTokens: { total: 1, noCache: undefined, cacheRead: undefined, cacheWrite: undefined },
|
||||
outputTokens: { total: 1, text: 1, reasoning: undefined },
|
||||
},
|
||||
},
|
||||
],
|
||||
}),
|
||||
}),
|
||||
});
|
||||
}
|
||||
|
||||
describe('ai@6.0.134 pnpm patch: writeToServerResponse drain-hang (#486)', () => {
|
||||
it('ends the response (does NOT hang) when the socket closes mid-write without draining', async () => {
|
||||
const result = streamText({ model: makeModel(), prompt: 'hi' });
|
||||
const res = new DisconnectingResponse();
|
||||
// Drain the SDK stream independently, like the production detached path.
|
||||
void result.consumeStream({ onError: () => undefined });
|
||||
result.pipeUIMessageStreamToResponse(res as never);
|
||||
|
||||
// TRIPWIRE: the patched loop exits on 'close' and runs finally -> end().
|
||||
// Unpatched, it awaits 'drain' forever and this never becomes true.
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
const started = Date.now();
|
||||
const poll = setInterval(() => {
|
||||
if (res.ended) {
|
||||
clearInterval(poll);
|
||||
resolve();
|
||||
} else if (Date.now() - started > 3000) {
|
||||
clearInterval(poll);
|
||||
reject(new Error('writeToServerResponse hung: response never ended'));
|
||||
}
|
||||
}, 20);
|
||||
});
|
||||
|
||||
expect(res.ended).toBe(true);
|
||||
});
|
||||
|
||||
it('does not emit an unhandledRejection when the fire-and-forget read() throws', async () => {
|
||||
// The patch swallows read()'s rejection (fire-and-forget) with a log instead
|
||||
// of letting it surface as a process-killing unhandledRejection.
|
||||
const rejections: unknown[] = [];
|
||||
const onUnhandled = (e: unknown) => rejections.push(e);
|
||||
process.on('unhandledRejection', onUnhandled);
|
||||
// Silence the patch's diagnostic console.error for the throwing read().
|
||||
const errSpy = jest.spyOn(console, 'error').mockImplementation(() => undefined);
|
||||
try {
|
||||
const result = streamText({ model: makeModel(), prompt: 'hi' });
|
||||
const res = new DisconnectingResponse();
|
||||
void result.consumeStream({ onError: () => undefined });
|
||||
result.pipeUIMessageStreamToResponse(res as never);
|
||||
await new Promise((r) => setTimeout(r, 300));
|
||||
} finally {
|
||||
process.off('unhandledRejection', onUnhandled);
|
||||
errSpy.mockRestore();
|
||||
}
|
||||
expect(rejections).toEqual([]);
|
||||
});
|
||||
|
||||
it('both installed dist builds (CJS and ESM) carry the #486 patch marker', () => {
|
||||
const cjsPath = require.resolve('ai');
|
||||
const mjsPath = cjsPath.replace(/index\.js$/, 'index.mjs');
|
||||
expect(cjsPath).toMatch(/index\.js$/);
|
||||
expect(readFileSync(cjsPath, 'utf8')).toContain('PATCH(docmost #486)');
|
||||
expect(readFileSync(mjsPath, 'utf8')).toContain('PATCH(docmost #486)');
|
||||
});
|
||||
});
|
||||
@@ -245,6 +245,9 @@ export class AiSettingsService {
|
||||
// Max context window for the chat header badge denominator. Stored as
|
||||
// ::text; 0/unset/invalid = no limit (undefined).
|
||||
chatContextWindow: parsePositiveInt(provider.chatContextWindow),
|
||||
// RAW stored value (#490): the replay budgeter reads this to distinguish an
|
||||
// explicit `0` (off-switch) from unset, which parsePositiveInt cannot.
|
||||
chatContextWindowRaw: provider.chatContextWindow,
|
||||
// Plain passthrough; getChatModel defaults unset to 'openai-compatible'.
|
||||
chatApiStyle: provider.chatApiStyle,
|
||||
// Cheap model id for the anonymous public-share assistant; reuses the chat
|
||||
|
||||
@@ -129,6 +129,12 @@ const DEFAULT_MCP_STREAM_TIMEOUT_MS = 60_000;
|
||||
/** Default total wall-clock cap for ONE external MCP tool call (2 min). */
|
||||
const DEFAULT_MCP_CALL_TIMEOUT_MS = 120_000;
|
||||
|
||||
/**
|
||||
* Default `bodyTimeout` for the EXTERNAL-MCP SSE transport (10 min) — #489.
|
||||
* Deliberately much LARGER than {@link DEFAULT_MCP_STREAM_TIMEOUT_MS}.
|
||||
*/
|
||||
const DEFAULT_MCP_SSE_BODY_TIMEOUT_MS = 600_000;
|
||||
|
||||
/**
|
||||
* SILENCE timeout (ms) for EXTERNAL-MCP transport ONLY. Override with
|
||||
* `AI_MCP_STREAM_TIMEOUT_MS`; a missing/invalid/non-positive value falls back to
|
||||
@@ -164,6 +170,26 @@ export function mcpCallTimeoutMs(): number {
|
||||
return positiveEnv('AI_MCP_CALL_TIMEOUT_MS', DEFAULT_MCP_CALL_TIMEOUT_MS);
|
||||
}
|
||||
|
||||
/**
|
||||
* `bodyTimeout` (ms) for the EXTERNAL-MCP **SSE** transport ONLY — #489. Override
|
||||
* with `AI_MCP_SSE_BODY_TIMEOUT_MS`; a missing/invalid/non-positive value falls
|
||||
* back to {@link DEFAULT_MCP_SSE_BODY_TIMEOUT_MS} (10 min).
|
||||
*
|
||||
* The SSE transport holds ONE long-lived response body open across many tool
|
||||
* calls, so undici's `bodyTimeout` (time between body bytes) counts the LEGITIMATE
|
||||
* silence BETWEEN calls, not just a hung single call. At the tight HTTP silence
|
||||
* timeout ({@link mcpStreamTimeoutMs}, 1 min) a normal >1-min gap between the
|
||||
* model's tool calls would break the SSE socket, and the cache would then serve a
|
||||
* dead client until TTL. So the SSE transport gets its OWN, RAISED bodyTimeout;
|
||||
* the per-call total cap ({@link mcpCallTimeoutMs}) still bounds a single stuck
|
||||
* call, and the app-level transport-error retry heals a socket that does break.
|
||||
* The HTTP (streamable) transport keeps the tight timeout — it opens a fresh
|
||||
* request per call, so idle-between-calls does not apply there.
|
||||
*/
|
||||
export function mcpSseBodyTimeoutMs(): number {
|
||||
return positiveEnv('AI_MCP_SSE_BODY_TIMEOUT_MS', DEFAULT_MCP_SSE_BODY_TIMEOUT_MS);
|
||||
}
|
||||
|
||||
/**
|
||||
* undici `Agent` options for streaming AI traffic — the (generous, finite)
|
||||
* silence timeouts plus the keep-alive recycle window. Shared by the chat
|
||||
|
||||
@@ -105,6 +105,10 @@ export interface ResolvedAiConfig extends Partial<AiProviderSettings> {
|
||||
// Max context window in tokens; surfaced to the chat header badge as the
|
||||
// "current / max" denominator. 0/unset = no limit.
|
||||
chatContextWindow?: number;
|
||||
// RAW stored context window (::text), BEFORE parsePositiveInt collapses `0` and
|
||||
// unset to `undefined`. The #490 replay budgeter needs the raw value to honor an
|
||||
// explicit `0` off-switch distinctly from "unset -> flat default".
|
||||
chatContextWindowRaw?: string | number;
|
||||
// Cheap model id for the public-share assistant; reuses the chat creds.
|
||||
publicShareChatModel?: string;
|
||||
// Agent-role id whose persona the public-share assistant adopts (empty/unset
|
||||
|
||||
@@ -0,0 +1,211 @@
|
||||
import type { HealthIndicatorService } from '@nestjs/terminus';
|
||||
import type { EnvironmentService } from '../environment/environment.service';
|
||||
|
||||
/**
|
||||
* Integration guard for the /health Redis-probe handle leak (#486, commit 2).
|
||||
*
|
||||
* The bug: `pingCheck` built `new Redis(...)` per call and only disconnected on
|
||||
* the SUCCESS path, so when Redis is DOWN every probe tick added ANOTHER
|
||||
* forever-reconnecting client — an unbounded handle/client leak for the duration
|
||||
* of the outage. The fix reuses ONE long-lived probe client.
|
||||
*
|
||||
* This is an OBSERVABLE-property test, not an assertion on a mocked return value:
|
||||
* we point the indicator at a REAL, refused TCP endpoint (a dead port) so ioredis
|
||||
* genuinely fails to connect, run many probes, and assert the number of live
|
||||
* Redis CLIENTS created stays at exactly ONE. `ioredis` is delegated to its real
|
||||
* implementation (requireActual) — only the constructor is wrapped to COUNT the
|
||||
* real clients it creates, which is precisely the leaking resource.
|
||||
*/
|
||||
import type { Redis } from 'ioredis';
|
||||
|
||||
const mockLiveClients: Redis[] = [];
|
||||
|
||||
/**
|
||||
* Fully tear a REAL ioredis client down so NO timer survives jest's 1s exit
|
||||
* window (this suite must exit cleanly WITHOUT forceExit; see #382).
|
||||
*
|
||||
* `connector.disconnect()` arms a ~12s "force-destroy the stream" `setTimeout`
|
||||
* that is cleared ONLY by the stream's 'close' event — but only when the
|
||||
* connector still holds a stream. Two problem cases:
|
||||
* - a LIVE/connecting socket: disconnect arms the timer and 'close' may lag
|
||||
* past jest's window, so we destroy the socket to make 'close' fire NOW;
|
||||
* - a client BETWEEN reconnect attempts to a dead port: the held socket is
|
||||
* ALREADY destroyed (its 'close' fired long ago), so disconnect would arm a
|
||||
* timer whose clearing 'close' can never come again. We drop that dead stream
|
||||
* reference BEFORE disconnect so the doomed timer is never armed.
|
||||
* `disconnect()` itself also clears ioredis' own reconnect backoff timer.
|
||||
*/
|
||||
type DrainableStream = { destroyed?: boolean; destroy?: () => void } | null;
|
||||
type DrainableClient = {
|
||||
removeAllListeners: (event: string) => void;
|
||||
disconnect: () => void;
|
||||
stream?: DrainableStream;
|
||||
connector?: { stream?: DrainableStream };
|
||||
};
|
||||
|
||||
async function drainClient(client: Redis): Promise<void> {
|
||||
if (!client || client.status === 'end') return;
|
||||
const c = client as unknown as DrainableClient;
|
||||
c.removeAllListeners('error');
|
||||
|
||||
// Drop an already-dead held socket so disconnect() can't arm a timer whose
|
||||
// clearing 'close' will never fire again.
|
||||
if (c.connector?.stream && c.connector.stream.destroyed) {
|
||||
c.connector.stream = null;
|
||||
}
|
||||
if (c.stream && c.stream.destroyed) {
|
||||
c.stream = null;
|
||||
}
|
||||
|
||||
await new Promise<void>((resolve) => {
|
||||
let done = false;
|
||||
const finish = () => {
|
||||
if (done) return;
|
||||
done = true;
|
||||
resolve();
|
||||
};
|
||||
client.once('end', finish);
|
||||
// reconnect=false (the default): stop the retry loop and close the socket.
|
||||
client.disconnect();
|
||||
// Force any still-live socket closed NOW so the connector's stream-destroy
|
||||
// timer clears inside jest's window instead of lagging behind a real 'close'.
|
||||
if (c.stream && !c.stream.destroyed) {
|
||||
c.stream.destroy?.();
|
||||
}
|
||||
// Fallback for a client with no live stream to emit 'end' (unref'd so it
|
||||
// can never itself hold the loop open).
|
||||
const fallback = setTimeout(finish, 500);
|
||||
(fallback as { unref?: () => void }).unref?.();
|
||||
});
|
||||
}
|
||||
|
||||
async function drainAll(): Promise<void> {
|
||||
await Promise.all(mockLiveClients.map((c) => drainClient(c)));
|
||||
}
|
||||
|
||||
jest.mock('ioredis', () => {
|
||||
const actual = jest.requireActual('ioredis');
|
||||
const RealRedis = actual.Redis ?? actual.default ?? actual;
|
||||
class CountingRedis extends RealRedis {
|
||||
constructor(...args: unknown[]) {
|
||||
super(...(args as []));
|
||||
mockLiveClients.push(this as never);
|
||||
}
|
||||
}
|
||||
return { ...actual, Redis: CountingRedis, default: CountingRedis };
|
||||
});
|
||||
|
||||
// Import AFTER the mock is registered so the class picks up the counting client.
|
||||
import { RedisHealthIndicator } from './redis.health';
|
||||
|
||||
describe('RedisHealthIndicator handle leak (#486)', () => {
|
||||
const indicatorService = {
|
||||
check: (key: string) => ({
|
||||
up: () => ({ [key]: { status: 'up' } }),
|
||||
down: (message: string) => ({ [key]: { status: 'down', message } }),
|
||||
}),
|
||||
} as unknown as HealthIndicatorService;
|
||||
|
||||
// A port with (almost certainly) nothing listening -> connection refused fast.
|
||||
const environmentService = {
|
||||
getRedisUrl: () => 'redis://127.0.0.1:6399/0',
|
||||
} as unknown as EnvironmentService;
|
||||
|
||||
let indicator: RedisHealthIndicator;
|
||||
|
||||
beforeEach(() => {
|
||||
mockLiveClients.length = 0;
|
||||
indicator = new RedisHealthIndicator(indicatorService, environmentService);
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
// Drain (destroy socket + AWAIT 'end') every client the test created FIRST,
|
||||
// so each is fully 'end' before onModuleDestroy's disconnect runs — that way
|
||||
// no ioredis reconnect / stream-destroy timer outlives jest's exit window.
|
||||
await drainAll();
|
||||
indicator.onModuleDestroy();
|
||||
});
|
||||
|
||||
it('creates exactly ONE Redis client across many probes while Redis is DOWN', async () => {
|
||||
const N = 8;
|
||||
for (let i = 0; i < N; i++) {
|
||||
const result = await indicator.pingCheck('redis');
|
||||
// Down endpoint -> every probe reports "down" (not an unhandled crash).
|
||||
expect(result.redis.status).toBe('down');
|
||||
}
|
||||
|
||||
// THE OBSERVABLE LEAK: on the buggy code this is N (a fresh, never-cleaned
|
||||
// reconnecting client per probe). The fix reuses one shared client.
|
||||
expect(mockLiveClients).toHaveLength(1);
|
||||
});
|
||||
|
||||
it('onModuleDestroy releases the probe client (a later probe builds a fresh one)', async () => {
|
||||
await indicator.pingCheck('redis');
|
||||
expect(mockLiveClients).toHaveLength(1);
|
||||
|
||||
indicator.onModuleDestroy();
|
||||
// A second destroy is a safe no-op (probeClient was nulled).
|
||||
indicator.onModuleDestroy();
|
||||
|
||||
// After shutdown the indicator lazily builds a NEW client on the next probe,
|
||||
// proving the old one was truly released rather than reused.
|
||||
await indicator.pingCheck('redis');
|
||||
expect(mockLiveClients).toHaveLength(2);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* Happy-path regression guard (#486, B2): the FIRST probe against a LIVE Redis
|
||||
* must report UP.
|
||||
*
|
||||
* With `lazyConnect: true` + `enableOfflineQueue: false`, a freshly-built client
|
||||
* is in the `wait` state and the socket opens lazily. If the very first `ping()`
|
||||
* is issued before an explicit `connect()`, ioredis rejects it instantly with
|
||||
* "Stream isn't writeable and enableOfflineQueue options is false" — a FALSE
|
||||
* DOWN even though Redis is alive. The fix opens the socket before the first
|
||||
* ping. This exercises a REAL ioredis client against a REAL TCP redis server
|
||||
* (not a mock), so a regression genuinely reddens it.
|
||||
*/
|
||||
describe('RedisHealthIndicator live Redis first-probe (#486, B2)', () => {
|
||||
const indicatorService = {
|
||||
check: (key: string) => ({
|
||||
up: () => ({ [key]: { status: 'up' } }),
|
||||
down: (message: string) => ({ [key]: { status: 'down', message } }),
|
||||
}),
|
||||
} as unknown as HealthIndicatorService;
|
||||
|
||||
// A REAL running redis (see the neighboring harness / CI env).
|
||||
const environmentService = {
|
||||
getRedisUrl: () => 'redis://127.0.0.1:6379/0',
|
||||
} as unknown as EnvironmentService;
|
||||
|
||||
let indicator: RedisHealthIndicator;
|
||||
|
||||
beforeEach(() => {
|
||||
mockLiveClients.length = 0;
|
||||
indicator = new RedisHealthIndicator(indicatorService, environmentService);
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
// Await full socket close of every live client (see drainClient) BEFORE
|
||||
// onModuleDestroy: a real, connected ioredis client MUST be drained to 'end'
|
||||
// or its stream-destroy timer keeps the jest worker alive past the 1s window.
|
||||
await drainAll();
|
||||
indicator.onModuleDestroy();
|
||||
});
|
||||
|
||||
it('reports UP on the FIRST probe against a live Redis', async () => {
|
||||
// The VERY FIRST probe — no warm-up ping — must be UP.
|
||||
const result = await indicator.pingCheck('redis');
|
||||
expect(result.redis.status).toBe('up');
|
||||
});
|
||||
|
||||
it('stays UP on a probe AFTER onModuleDestroy re-creates the client', async () => {
|
||||
await indicator.pingCheck('redis');
|
||||
indicator.onModuleDestroy();
|
||||
// The re-created client is again in `wait`; the first ping on it must still
|
||||
// open the socket (the false-DOWN also recurs on the post-destroy path).
|
||||
const result = await indicator.pingCheck('redis');
|
||||
expect(result.redis.status).toBe('up');
|
||||
});
|
||||
});
|
||||
@@ -2,33 +2,173 @@ import {
|
||||
HealthIndicatorResult,
|
||||
HealthIndicatorService,
|
||||
} from '@nestjs/terminus';
|
||||
import { Injectable, Logger } from '@nestjs/common';
|
||||
import { Injectable, Logger, OnModuleDestroy } from '@nestjs/common';
|
||||
import { EnvironmentService } from '../environment/environment.service';
|
||||
import { Redis } from 'ioredis';
|
||||
|
||||
@Injectable()
|
||||
export class RedisHealthIndicator {
|
||||
export class RedisHealthIndicator implements OnModuleDestroy {
|
||||
private readonly logger = new Logger(RedisHealthIndicator.name);
|
||||
|
||||
/**
|
||||
* ONE long-lived probe connection, reused across every /health tick. The old
|
||||
* code built `new Redis(...)` per call and only `disconnect()`d on the SUCCESS
|
||||
* path, so while Redis was DOWN every probe added a fresh, forever-reconnecting
|
||||
* client — a handle leak that grew without bound for as long as the outage (and
|
||||
* the health checker keeps polling) lasted. A single shared client keeps at most
|
||||
* ONE background reconnect loop regardless of how many probes run.
|
||||
*/
|
||||
private probeClient: Redis | null = null;
|
||||
|
||||
/**
|
||||
* How long the first-ping `connect()` may take before a probe gives up and
|
||||
* reports DOWN. A `connect()` against a truly-down Redis never settles on its
|
||||
* own (ioredis retries the socket indefinitely per its retryStrategy), so the
|
||||
* probe MUST bound it or the /health handler would hang. Kept short so a real
|
||||
* outage is reported fast; localhost/live Redis connects well within it.
|
||||
*/
|
||||
private static readonly CONNECT_TIMEOUT_MS = 2000;
|
||||
|
||||
/**
|
||||
* The single in-flight first-`connect()`, memoized so CONCURRENT probes share
|
||||
* it. k8s liveness+readiness hit /health in parallel on startup: without this,
|
||||
* probe A drives `connect()` (the client leaves the `wait` state) and probe B,
|
||||
* seeing a not-`wait`/not-`ready` client, would skip connect and fire `ping()`
|
||||
* at a still-opening socket → an instant FALSE DOWN. With the memo, B awaits
|
||||
* the SAME connect. Cleared once it settles so a later disconnect / re-create
|
||||
* starts a fresh connect.
|
||||
*/
|
||||
private connectingPromise: Promise<void> | null = null;
|
||||
|
||||
constructor(
|
||||
private readonly healthIndicatorService: HealthIndicatorService,
|
||||
private environmentService: EnvironmentService,
|
||||
) {}
|
||||
|
||||
private getProbeClient(): Redis {
|
||||
if (!this.probeClient) {
|
||||
this.probeClient = new Redis(this.environmentService.getRedisUrl(), {
|
||||
// Constructing must never throw or eagerly connect; the first ping opens
|
||||
// the socket. This lets us build the client once and reuse it.
|
||||
lazyConnect: true,
|
||||
// A health probe must fail FAST, not queue behind a stuck reconnect: one
|
||||
// retry per request, and no offline queue so a ping while disconnected
|
||||
// rejects immediately instead of buffering commands that pile up in RAM.
|
||||
maxRetriesPerRequest: 1,
|
||||
enableOfflineQueue: false,
|
||||
});
|
||||
// ioredis emits 'error' on every failed (re)connect; with no listener that
|
||||
// surfaces as an unhandled 'error' event and can crash the process. Swallow
|
||||
// it here — pingCheck already reports health — and log at debug so a Redis
|
||||
// outage does not flood the logs.
|
||||
this.probeClient.on('error', (err) => {
|
||||
this.logger.debug(
|
||||
`Redis probe connection error: ${
|
||||
err instanceof Error ? err.message : String(err)
|
||||
}`,
|
||||
);
|
||||
});
|
||||
}
|
||||
return this.probeClient;
|
||||
}
|
||||
|
||||
/**
|
||||
* Open the probe socket BEFORE the first ping. `lazyConnect: true` leaves a
|
||||
* freshly-built (or post-destroy re-built) client in the `wait` state: the
|
||||
* socket is NOT open yet, so with `enableOfflineQueue: false` the very first
|
||||
* `ping()` rejects instantly with "Stream isn't writeable and
|
||||
* enableOfflineQueue options is false" even when Redis is perfectly alive — a
|
||||
* false DOWN on the happy path. We drive `connect()` ONLY from `wait`; once
|
||||
* the client is connected, ioredis owns its own (re)connect loop and a ping
|
||||
* issued while it reconnects still fast-fails to a correct DOWN (offline queue
|
||||
* stays off). A failed/timed-out connect rejects → reported DOWN, which is the
|
||||
* right signal for a truly-down Redis.
|
||||
*/
|
||||
private ensureConnected(client: Redis): Promise<void> {
|
||||
// Already open — steady state, nothing to do.
|
||||
if (client.status === 'ready') return Promise.resolve();
|
||||
// A first-connect is already in flight (possibly started by a CONCURRENT
|
||||
// probe): await the SAME one instead of racing a second connect() (ioredis
|
||||
// throws "already connecting") or firing ping() at a not-yet-open socket.
|
||||
if (this.connectingPromise) return this.connectingPromise;
|
||||
// Only DRIVE connect() from the initial `wait` state (fresh / post-destroy
|
||||
// re-created client). In any other non-ready state ioredis already owns its
|
||||
// (re)connect loop; a ping there fast-fails to a correct DOWN, so we must not
|
||||
// start a competing connect.
|
||||
if (client.status !== 'wait') return Promise.resolve();
|
||||
|
||||
const promise = this.connectWithTimeout(client).finally(() => {
|
||||
// Clear only if still ours, so a later disconnect / re-create can connect
|
||||
// again. Whether it resolved or rejected, the memo has served its window.
|
||||
if (this.connectingPromise === promise) {
|
||||
this.connectingPromise = null;
|
||||
}
|
||||
});
|
||||
this.connectingPromise = promise;
|
||||
return promise;
|
||||
}
|
||||
|
||||
private connectWithTimeout(client: Redis): Promise<void> {
|
||||
return new Promise<void>((resolve, reject) => {
|
||||
let settled = false;
|
||||
const timer = setTimeout(() => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
reject(new Error('Redis probe connect timed out'));
|
||||
}, RedisHealthIndicator.CONNECT_TIMEOUT_MS);
|
||||
// Never let THIS timer alone keep the event loop (or a jest worker) alive;
|
||||
// it is cleared on settle anyway, this is belt-and-braces.
|
||||
timer.unref?.();
|
||||
// `.catch` is always attached, so a connect() that rejects AFTER we have
|
||||
// already timed out is handled here (guarded by `settled`) and never
|
||||
// surfaces as an unhandled rejection.
|
||||
client
|
||||
.connect()
|
||||
.then(() => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
clearTimeout(timer);
|
||||
resolve();
|
||||
})
|
||||
.catch((err) => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
clearTimeout(timer);
|
||||
reject(err);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
async pingCheck(key: string): Promise<HealthIndicatorResult> {
|
||||
const indicator = this.healthIndicatorService.check(key);
|
||||
|
||||
try {
|
||||
const redis = new Redis(this.environmentService.getRedisUrl(), {
|
||||
maxRetriesPerRequest: 15,
|
||||
});
|
||||
|
||||
const redis = this.getProbeClient();
|
||||
// Open the socket before the first ping (see ensureConnected); without
|
||||
// this the first probe after (re)creation falsely reports DOWN on a live
|
||||
// Redis because lazyConnect defers the connect past the first ping.
|
||||
await this.ensureConnected(redis);
|
||||
await redis.ping();
|
||||
redis.disconnect();
|
||||
return indicator.up();
|
||||
} catch (e) {
|
||||
this.logger.error(e);
|
||||
return indicator.down(`${key} is not available`);
|
||||
}
|
||||
}
|
||||
|
||||
onModuleDestroy(): void {
|
||||
if (this.probeClient) {
|
||||
// disconnect() (not quit()) tears the socket + reconnect loop down
|
||||
// immediately without waiting on a round-trip to a possibly-down server.
|
||||
// Do NOT removeAllListeners() with no event name — that would also strip
|
||||
// ioredis' OWN internal listeners and break its teardown; our 'error'
|
||||
// listener is harmless and dies with the dropped client reference.
|
||||
this.probeClient.disconnect();
|
||||
this.probeClient = null;
|
||||
}
|
||||
// Drop any in-flight first-connect memo so the NEXT client (lazily rebuilt on
|
||||
// the next probe) starts a fresh connect rather than awaiting a promise tied
|
||||
// to the client we just tore down.
|
||||
this.connectingPromise = null;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,6 +16,7 @@ import {
|
||||
} from './mcp-auth.helpers';
|
||||
import { JwtType } from '../../core/auth/dto/jwt-payload';
|
||||
import { CREDENTIALS_MISMATCH_MESSAGE } from '../../core/auth/auth.constants';
|
||||
import { McpService } from './mcp.service';
|
||||
|
||||
// The /mcp per-user auth decision logic is tested through the framework-free
|
||||
// `resolveMcpSessionConfig` helper that McpService delegates to. McpService
|
||||
@@ -1179,3 +1180,46 @@ describe('mapAuthResultToResponse (handle status/body mapping, refactor R2)', ()
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
// #486: onModuleDestroy must ALSO tear down the live loopback CollabSessions, not
|
||||
// just clear the sweep timer — otherwise the embedded MCP's collab sockets keep
|
||||
// docs pinned open on the collab server past process exit. The teardown goes
|
||||
// through an overridable seam (destroyAllMcpSessions) so it can be spied without
|
||||
// loading the ESM-only @docmost/mcp package.
|
||||
describe('McpService.onModuleDestroy — CollabSession teardown (#486)', () => {
|
||||
function makeService(): McpService {
|
||||
// The constructor only stores its deps and starts the (unref'd) sweep timer,
|
||||
// so bare stubs suffice. onModuleDestroy clears that timer, so no leak.
|
||||
return new McpService(
|
||||
{} as any,
|
||||
{} as any,
|
||||
{} as any,
|
||||
{} as any,
|
||||
{} as any,
|
||||
{} as any,
|
||||
{} as any,
|
||||
);
|
||||
}
|
||||
|
||||
it('destroys all sessions AND clears the sweep timer on shutdown', async () => {
|
||||
const svc = makeService();
|
||||
const destroy = jest.fn().mockResolvedValue(undefined);
|
||||
(svc as any).destroyAllMcpSessions = destroy;
|
||||
const clearSpy = jest.spyOn(global, 'clearInterval');
|
||||
|
||||
await svc.onModuleDestroy();
|
||||
|
||||
expect(destroy).toHaveBeenCalledTimes(1);
|
||||
expect(clearSpy).toHaveBeenCalledWith((svc as any).sweepTimer);
|
||||
clearSpy.mockRestore();
|
||||
});
|
||||
|
||||
it('swallows a teardown failure so shutdown never throws', async () => {
|
||||
const svc = makeService();
|
||||
(svc as any).destroyAllMcpSessions = jest
|
||||
.fn()
|
||||
.mockRejectedValue(new Error('collab teardown boom'));
|
||||
|
||||
await expect(svc.onModuleDestroy()).resolves.toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -119,10 +119,42 @@ export class McpService implements OnModuleDestroy {
|
||||
this.sweepTimer.unref?.();
|
||||
}
|
||||
|
||||
onModuleDestroy(): void {
|
||||
async onModuleDestroy(): Promise<void> {
|
||||
clearInterval(this.sweepTimer);
|
||||
// Tear down any live loopback CollabSession providers at shutdown (#486). The
|
||||
// embedded MCP (and the in-app AI agent) open Hocuspocus collab sockets against
|
||||
// THIS process; without an explicit teardown those sessions keep their docs
|
||||
// "open" on the collab server and hold providers/buffers until they idle out,
|
||||
// so a restart can race a doc still pinned by the dying worker. Best-effort:
|
||||
// any failure is logged, never allowed to break shutdown.
|
||||
try {
|
||||
await this.destroyAllMcpSessions();
|
||||
} catch (err) {
|
||||
this.logger.error(
|
||||
'MCP CollabSession teardown on shutdown failed',
|
||||
err as Error,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve @docmost/mcp's `destroyAllSessions` and invoke it (#486). The live
|
||||
* CollabSession registry is a module-level singleton in the ESM package, shared
|
||||
* by every entry (`.`/`./http`), so this tears down ALL sessions regardless of
|
||||
* which surface opened them. The module is already loaded whenever MCP was used;
|
||||
* if it was never loaded (or is absent) the import + no-op is harmless.
|
||||
*
|
||||
* Held as an overridable field so a unit test can spy the teardown without
|
||||
* loading the ESM-only package or standing up the DI graph.
|
||||
*/
|
||||
private destroyAllMcpSessions: () => Promise<void> = async () => {
|
||||
const entry = require.resolve('@docmost/mcp');
|
||||
const mod = (await esmImport(pathToFileURL(entry).href)) as {
|
||||
destroyAllSessions?: () => void;
|
||||
};
|
||||
mod.destroyAllSessions?.();
|
||||
};
|
||||
|
||||
// Service account the embedded MCP uses to talk back to this Docmost
|
||||
// instance over loopback REST + the collaboration WebSocket. Now OPTIONAL:
|
||||
// it is only a fallback when no per-user Basic/Bearer credentials are sent.
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
import { get as httpGet } from 'node:http';
|
||||
import { AddressInfo } from 'node:net';
|
||||
import { createServer } from 'node:http';
|
||||
|
||||
// Drive the metrics HTTP server without the load-time METRICS_PORT gate: mock the
|
||||
// registry so isMetricsEnabled()/getMetricsRegistry() are always satisfied. What
|
||||
// we assert is observed over a REAL socket (bind address, status codes), not on
|
||||
// the mock.
|
||||
jest.mock('./metrics.registry', () => ({
|
||||
isMetricsEnabled: () => true,
|
||||
getMetricsRegistry: () => ({
|
||||
metrics: async () => '# HELP up test\nup 1\n',
|
||||
contentType: 'text/plain; version=0.0.4',
|
||||
}),
|
||||
}));
|
||||
|
||||
import {
|
||||
startMetricsServer,
|
||||
closeMetricsServer,
|
||||
resolveMetricsBind,
|
||||
resolveMetricsToken,
|
||||
} from './metrics.server';
|
||||
|
||||
/** Find a free TCP port (the metrics server requires METRICS_PORT > 0). */
|
||||
function freePort(): Promise<number> {
|
||||
return new Promise((resolve, reject) => {
|
||||
const s = createServer();
|
||||
s.once('error', reject);
|
||||
s.listen(0, '127.0.0.1', () => {
|
||||
const p = (s.address() as AddressInfo).port;
|
||||
s.close(() => resolve(p));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** Minimal GET against 127.0.0.1:port with optional Authorization header. */
|
||||
function req(
|
||||
port: number,
|
||||
headers: Record<string, string> = {},
|
||||
): Promise<{ status: number; body: string }> {
|
||||
return new Promise((resolve, reject) => {
|
||||
const r = httpGet(
|
||||
{ host: '127.0.0.1', port, path: '/metrics', headers },
|
||||
(res) => {
|
||||
let body = '';
|
||||
res.on('data', (c) => (body += c));
|
||||
res.on('end', () =>
|
||||
resolve({ status: res.statusCode ?? 0, body }),
|
||||
);
|
||||
},
|
||||
);
|
||||
r.on('error', reject);
|
||||
});
|
||||
}
|
||||
|
||||
describe('metrics server bind + auth (#486)', () => {
|
||||
const saved = {
|
||||
bind: process.env.METRICS_BIND,
|
||||
token: process.env.METRICS_TOKEN,
|
||||
port: process.env.METRICS_PORT,
|
||||
};
|
||||
|
||||
afterEach(async () => {
|
||||
await closeMetricsServer();
|
||||
process.env.METRICS_BIND = saved.bind;
|
||||
process.env.METRICS_TOKEN = saved.token;
|
||||
process.env.METRICS_PORT = saved.port;
|
||||
delete process.env.METRICS_BIND;
|
||||
delete process.env.METRICS_TOKEN;
|
||||
});
|
||||
|
||||
describe('resolveMetricsBind', () => {
|
||||
it('defaults to loopback 127.0.0.1', () => {
|
||||
delete process.env.METRICS_BIND;
|
||||
expect(resolveMetricsBind()).toBe('127.0.0.1');
|
||||
});
|
||||
it('honours the METRICS_BIND override', () => {
|
||||
process.env.METRICS_BIND = '0.0.0.0';
|
||||
expect(resolveMetricsBind()).toBe('0.0.0.0');
|
||||
});
|
||||
it('treats a blank override as unset (loopback)', () => {
|
||||
process.env.METRICS_BIND = ' ';
|
||||
expect(resolveMetricsBind()).toBe('127.0.0.1');
|
||||
});
|
||||
});
|
||||
|
||||
describe('resolveMetricsToken', () => {
|
||||
it('is null when unset', () => {
|
||||
delete process.env.METRICS_TOKEN;
|
||||
expect(resolveMetricsToken()).toBeNull();
|
||||
});
|
||||
it('returns the trimmed token when set', () => {
|
||||
process.env.METRICS_TOKEN = ' s3cret ';
|
||||
expect(resolveMetricsToken()).toBe('s3cret');
|
||||
});
|
||||
});
|
||||
|
||||
it('binds to loopback by default and serves /metrics without auth when no token', async () => {
|
||||
delete process.env.METRICS_BIND;
|
||||
delete process.env.METRICS_TOKEN;
|
||||
const port = await freePort();
|
||||
process.env.METRICS_PORT = String(port);
|
||||
|
||||
const server = startMetricsServer();
|
||||
expect(server).not.toBeNull();
|
||||
await new Promise<void>((resolve) => {
|
||||
if (server!.listening) resolve();
|
||||
else server!.once('listening', () => resolve());
|
||||
});
|
||||
// OBSERVABLE: the listener bound to loopback, not 0.0.0.0.
|
||||
expect((server!.address() as AddressInfo).address).toBe('127.0.0.1');
|
||||
|
||||
const res = await req(port);
|
||||
expect(res.status).toBe(200);
|
||||
expect(res.body).toContain('up 1');
|
||||
});
|
||||
|
||||
it('rejects unauthenticated scrapes with 401 and accepts the exact Bearer token', async () => {
|
||||
delete process.env.METRICS_BIND;
|
||||
process.env.METRICS_TOKEN = 'topsecret';
|
||||
const port = await freePort();
|
||||
process.env.METRICS_PORT = String(port);
|
||||
|
||||
const server = startMetricsServer();
|
||||
expect(server).not.toBeNull();
|
||||
|
||||
// No auth -> 401.
|
||||
const noAuth = await req(port);
|
||||
expect(noAuth.status).toBe(401);
|
||||
|
||||
// Wrong token, DIFFERENT length -> 401 (short-circuits on the length guard).
|
||||
const wrong = await req(port, { authorization: 'Bearer nope' });
|
||||
expect(wrong.status).toBe(401);
|
||||
|
||||
// Wrong token, SAME length -> 401. This drives the timingSafeEqual compare
|
||||
// itself (the length guard passes: 'Bearer topsecreX' has the same length as
|
||||
// 'Bearer topsecret'). Pins the constant-time compare: a regression that made
|
||||
// it return true would let this equal-length wrong token through — the
|
||||
// different-length case above would NOT catch that.
|
||||
const sameLen = await req(port, { authorization: 'Bearer topsecreX' });
|
||||
expect(sameLen.status).toBe(401);
|
||||
|
||||
// Correct token -> 200 with the metrics body.
|
||||
const ok = await req(port, { authorization: 'Bearer topsecret' });
|
||||
expect(ok.status).toBe(200);
|
||||
expect(ok.body).toContain('up 1');
|
||||
});
|
||||
});
|
||||
@@ -1,7 +1,27 @@
|
||||
import { createServer, Server } from 'node:http';
|
||||
import { timingSafeEqual } from 'node:crypto';
|
||||
import { Logger } from '@nestjs/common';
|
||||
import { getMetricsRegistry, isMetricsEnabled } from './metrics.registry';
|
||||
|
||||
/**
|
||||
* Constant-time compare of the presented Authorization header against the
|
||||
* expected `Bearer <token>`. This is the ONLY auth layer for the metrics
|
||||
* endpoint, so a naive `!==` would leak the token byte-by-byte via timing.
|
||||
* timingSafeEqual requires equal-length buffers, so a length mismatch short-
|
||||
* circuits to "not equal" (its own length is not itself a useful oracle: the
|
||||
* expected string length is fixed by config, not secret-derived).
|
||||
*/
|
||||
function bearerMatches(
|
||||
presented: string | undefined,
|
||||
expected: string,
|
||||
): boolean {
|
||||
if (typeof presented !== 'string') return false;
|
||||
const a = Buffer.from(presented);
|
||||
const b = Buffer.from(expected);
|
||||
if (a.length !== b.length) return false;
|
||||
return timingSafeEqual(a, b);
|
||||
}
|
||||
|
||||
/**
|
||||
* Start the Prometheus scrape endpoint on a SEPARATE port, taken from
|
||||
* `METRICS_PORT`. There is NO default port: when `METRICS_PORT` is unset the
|
||||
@@ -16,6 +36,30 @@ import { getMetricsRegistry, isMetricsEnabled } from './metrics.registry';
|
||||
*/
|
||||
let metricsServer: Server | null = null;
|
||||
|
||||
/**
|
||||
* Interface the metrics endpoint binds to. Defaults to LOOPBACK (127.0.0.1) so
|
||||
* the unauthenticated `/metrics` surface is NOT exposed on all interfaces by
|
||||
* default — the old `0.0.0.0` bind put an auth-less endpoint on every interface.
|
||||
* Deployments where the scraper runs in a SEPARATE container (and reaches this as
|
||||
* `docmost:9464`) set `METRICS_BIND=0.0.0.0`, ideally together with METRICS_TOKEN
|
||||
* and/or a private network so the port is not world-readable.
|
||||
*/
|
||||
export function resolveMetricsBind(): string {
|
||||
const raw = (process.env.METRICS_BIND ?? '').trim();
|
||||
return raw.length > 0 ? raw : '127.0.0.1';
|
||||
}
|
||||
|
||||
/**
|
||||
* Optional Bearer token guarding `/metrics`. When `METRICS_TOKEN` is set, every
|
||||
* scrape must present `Authorization: Bearer <token>`; unset (default) leaves the
|
||||
* endpoint open (safe when bound to loopback / a trusted network). Returns the
|
||||
* trimmed token or null when unset/blank.
|
||||
*/
|
||||
export function resolveMetricsToken(): string | null {
|
||||
const raw = (process.env.METRICS_TOKEN ?? '').trim();
|
||||
return raw.length > 0 ? raw : null;
|
||||
}
|
||||
|
||||
export function startMetricsServer(): Server | null {
|
||||
if (!isMetricsEnabled()) return null;
|
||||
|
||||
@@ -31,8 +75,22 @@ export function startMetricsServer(): Server | null {
|
||||
return null;
|
||||
}
|
||||
|
||||
const bind = resolveMetricsBind();
|
||||
const token = resolveMetricsToken();
|
||||
|
||||
const server = createServer(async (req, res) => {
|
||||
if (req.method === 'GET' && req.url === '/metrics') {
|
||||
// Optional Bearer auth: reject scrapes without the exact token when one is
|
||||
// configured. This is the auth layer the old all-interfaces bind lacked.
|
||||
if (token) {
|
||||
const auth = req.headers['authorization'];
|
||||
if (!bearerMatches(auth, `Bearer ${token}`)) {
|
||||
res.statusCode = 401;
|
||||
res.setHeader('WWW-Authenticate', 'Bearer');
|
||||
res.end();
|
||||
return;
|
||||
}
|
||||
}
|
||||
try {
|
||||
const body = await register.metrics();
|
||||
res.setHeader('Content-Type', register.contentType);
|
||||
@@ -48,10 +106,14 @@ export function startMetricsServer(): Server | null {
|
||||
res.end();
|
||||
});
|
||||
|
||||
// Bind on all interfaces: the scraper (VictoriaMetrics) reaches this from
|
||||
// another container as docmost:9464. The port is not published to the host.
|
||||
server.listen(port, '0.0.0.0', () => {
|
||||
logger.log(`Metrics endpoint listening on :${port}/metrics`);
|
||||
// Bind to loopback by default so the auth-less endpoint is not exposed on all
|
||||
// interfaces. Set METRICS_BIND=0.0.0.0 (ideally with METRICS_TOKEN) when the
|
||||
// scraper runs in a separate container and reaches this as docmost:9464.
|
||||
server.listen(port, bind, () => {
|
||||
logger.log(
|
||||
`Metrics endpoint listening on ${bind}:${port}/metrics` +
|
||||
(token ? ' (Bearer auth required)' : ''),
|
||||
);
|
||||
});
|
||||
|
||||
server.on('error', (err) => {
|
||||
|
||||
@@ -31,9 +31,6 @@ export enum QueueJob {
|
||||
IMPORT_TASK = 'import-task',
|
||||
EXPORT_TASK = 'export-task',
|
||||
|
||||
SEARCH_REMOVE_PAGE = 'search-remove-page',
|
||||
SEARCH_REMOVE_ASSET = 'search-remove-attachment',
|
||||
SEARCH_REMOVE_FACE = 'search-remove-comment',
|
||||
TYPESENSE_FLUSH = 'typesense-flush',
|
||||
|
||||
PAGE_CREATED = 'page-created',
|
||||
|
||||
@@ -0,0 +1,305 @@
|
||||
import { Kysely } from 'kysely';
|
||||
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 { AiChatRunService } from '../../src/core/ai-chat/ai-chat-run.service';
|
||||
import {
|
||||
getTestDb,
|
||||
destroyTestDb,
|
||||
createWorkspace,
|
||||
createUser,
|
||||
createChat,
|
||||
createMessage,
|
||||
} from './db';
|
||||
|
||||
/**
|
||||
* #487 commit 4 — bidirectional reconcile + owner-write priority, real SQL.
|
||||
*
|
||||
* Proves the OBSERVABLE recovery properties against docmost_test:
|
||||
* - the CONDITIONAL owner-write beats a reconcile stamp, and a stamp never
|
||||
* clobbers a proper terminal row;
|
||||
* - a LATE owner-finalize with real content OVERWRITES a reconcile 'aborted'
|
||||
* stamp (finalizeFailed);
|
||||
* - each reconcile clause (b message<-run, c stale-run, d historical row) settles
|
||||
* the stuck row/run, and a LIVE run entry is never touched;
|
||||
* - the "kill DB on finish" recovery: after the DB comes back, neither the
|
||||
* message row nor the run row stays stuck.
|
||||
*/
|
||||
describe('#487 reconcile + owner-write priority [integration]', () => {
|
||||
let db: Kysely<any>;
|
||||
let messageRepo: AiChatMessageRepo;
|
||||
let runRepo: AiChatRunRepo;
|
||||
let runService: AiChatRunService;
|
||||
let workspaceId: string;
|
||||
let userId: string;
|
||||
|
||||
beforeAll(async () => {
|
||||
db = getTestDb();
|
||||
messageRepo = new AiChatMessageRepo(db as any);
|
||||
runRepo = new AiChatRunRepo(db as any);
|
||||
runService = new AiChatRunService(runRepo, { isCloud: () => false } as never);
|
||||
workspaceId = (await createWorkspace(db)).id;
|
||||
userId = (await createUser(db, workspaceId)).id;
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await destroyTestDb();
|
||||
});
|
||||
|
||||
const newChat = async () =>
|
||||
(await createChat(db, { workspaceId, creatorId: userId })).id;
|
||||
|
||||
const metaOf = async (id: string): Promise<Record<string, unknown> | null> => {
|
||||
const row = await messageRepo.findById(id, workspaceId);
|
||||
return (row?.metadata as Record<string, unknown> | null) ?? null;
|
||||
};
|
||||
|
||||
it('owner finalizeOwner writes a streaming row and CLEARS finalizeFailed', async () => {
|
||||
const chatId = await newChat();
|
||||
const m = await createMessage(db, {
|
||||
workspaceId,
|
||||
chatId,
|
||||
role: 'assistant',
|
||||
status: 'streaming',
|
||||
metadata: { parts: [] },
|
||||
});
|
||||
const wrote = await messageRepo.finalizeOwner(m.id, workspaceId, {
|
||||
content: 'final answer',
|
||||
status: 'completed',
|
||||
metadata: { parts: [{ type: 'text', text: 'final answer' }] },
|
||||
} as never);
|
||||
expect(wrote!.status).toBe('completed');
|
||||
expect((await metaOf(m.id))?.finalizeFailed).toBeUndefined();
|
||||
});
|
||||
|
||||
it('a reconcile stamp NEVER clobbers a proper terminal row (finalizeOwner is a no-op there)', async () => {
|
||||
const chatId = await newChat();
|
||||
const m = await createMessage(db, {
|
||||
workspaceId,
|
||||
chatId,
|
||||
role: 'assistant',
|
||||
status: 'completed',
|
||||
content: 'real',
|
||||
metadata: { parts: [] },
|
||||
});
|
||||
// The reconcile stamp is onlyIfStreaming -> no-op on a completed row.
|
||||
const stamped = await messageRepo.stampTerminalIfStreaming(
|
||||
m.id,
|
||||
workspaceId,
|
||||
'aborted',
|
||||
);
|
||||
expect(stamped).toBeUndefined();
|
||||
expect((await messageRepo.findById(m.id, workspaceId))!.status).toBe(
|
||||
'completed',
|
||||
);
|
||||
});
|
||||
|
||||
it('LATE owner-finalize with real content OVERWRITES a reconcile aborted stamp', async () => {
|
||||
const chatId = await newChat();
|
||||
const m = await createMessage(db, {
|
||||
workspaceId,
|
||||
chatId,
|
||||
role: 'assistant',
|
||||
status: 'streaming',
|
||||
metadata: { parts: [{ type: 'text', text: 'partial' }] },
|
||||
});
|
||||
// Reconcile stamps it aborted + finalizeFailed (final text lived only in mem).
|
||||
const stamped = await messageRepo.stampTerminalIfStreaming(
|
||||
m.id,
|
||||
workspaceId,
|
||||
'aborted',
|
||||
);
|
||||
expect(stamped!.status).toBe('aborted');
|
||||
expect((await metaOf(m.id))?.finalizeFailed).toBe(true);
|
||||
|
||||
// A LATE owner-write (finalizeFailed=true satisfies the OR) overwrites it with
|
||||
// real content, clearing the flag — owner-write priority.
|
||||
const wrote = await messageRepo.finalizeOwner(m.id, workspaceId, {
|
||||
content: 'the real final answer',
|
||||
status: 'completed',
|
||||
metadata: { parts: [{ type: 'text', text: 'the real final answer' }] },
|
||||
} as never);
|
||||
expect(wrote!.status).toBe('completed');
|
||||
expect(wrote!.content).toBe('the real final answer');
|
||||
expect((await metaOf(m.id))?.finalizeFailed).toBeUndefined();
|
||||
});
|
||||
|
||||
it('clause (c): a stale active run with NO live entry -> aborted; a LIVE entry is untouched', async () => {
|
||||
// Stale run, NOT owned by this replica (no entry) -> reconcile aborts it.
|
||||
const staleChat = await newChat();
|
||||
const stale = await runRepo.insert({
|
||||
chatId: staleChat,
|
||||
workspaceId,
|
||||
createdBy: userId,
|
||||
status: 'running',
|
||||
});
|
||||
await db
|
||||
.updateTable('aiChatRuns')
|
||||
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
|
||||
.where('id', '=', stale.id)
|
||||
.execute();
|
||||
|
||||
// A live run OWNED by this replica (beginRun registers an in-memory entry),
|
||||
// ALSO backdated stale — the "no entry" primary gate must protect it.
|
||||
const liveChat = await newChat();
|
||||
const live = await runService.beginRun({
|
||||
chatId: liveChat,
|
||||
workspaceId,
|
||||
userId,
|
||||
});
|
||||
await db
|
||||
.updateTable('aiChatRuns')
|
||||
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
|
||||
.where('id', '=', live.runId)
|
||||
.execute();
|
||||
|
||||
const aborted = await runService.reconcileStaleRuns(15 * 60 * 1000);
|
||||
expect(aborted).toBeGreaterThanOrEqual(1);
|
||||
expect((await runRepo.findById(stale.id, workspaceId))!.status).toBe(
|
||||
'aborted',
|
||||
);
|
||||
// The live entry is NEVER aborted, however stale its row looks.
|
||||
expect((await runRepo.findById(live.runId, workspaceId))!.status).toBe(
|
||||
'running',
|
||||
);
|
||||
expect(runService.isLocallyActive(live.runId)).toBe(true);
|
||||
|
||||
// cleanup the live run
|
||||
await runService.finalizeRun(live.runId, workspaceId, 'aborted');
|
||||
});
|
||||
|
||||
it('clause (b): a streaming message whose RUN is terminal is stamped by run status (succeeded -> aborted, NOT completed-empty)', async () => {
|
||||
const chatId = await newChat();
|
||||
const msg = await createMessage(db, {
|
||||
workspaceId,
|
||||
chatId,
|
||||
role: 'assistant',
|
||||
status: 'streaming',
|
||||
metadata: { parts: [] },
|
||||
});
|
||||
// A SUCCEEDED run linked to the still-streaming message (the asymmetry).
|
||||
const run = await runRepo.insert({
|
||||
chatId,
|
||||
workspaceId,
|
||||
createdBy: userId,
|
||||
status: 'running',
|
||||
assistantMessageId: msg.id,
|
||||
});
|
||||
await runRepo.finalizeIfActive(run.id, workspaceId, {
|
||||
status: 'succeeded',
|
||||
error: null,
|
||||
});
|
||||
|
||||
const stuck = await messageRepo.findStreamingWithTerminalRun();
|
||||
const mine = stuck.find((s) => s.messageId === msg.id);
|
||||
expect(mine?.runStatus).toBe('succeeded');
|
||||
// Reconcile clause (b): succeeded run -> message 'aborted' (NOT 'completed'),
|
||||
// the final text lived only in memory (documented loss), +finalizeFailed.
|
||||
const status = mine!.runStatus === 'failed' ? 'error' : 'aborted';
|
||||
await messageRepo.stampTerminalIfStreaming(msg.id, workspaceId, status);
|
||||
const row = await messageRepo.findById(msg.id, workspaceId);
|
||||
expect(row!.status).toBe('aborted');
|
||||
expect((row!.metadata as Record<string, unknown>).finalizeFailed).toBe(true);
|
||||
});
|
||||
|
||||
it('clause (d): a stale streaming row with NO active run on the chat -> aborted+finalizeFailed', async () => {
|
||||
const chatId = await newChat();
|
||||
const msg = await createMessage(db, {
|
||||
workspaceId,
|
||||
chatId,
|
||||
role: 'assistant',
|
||||
status: 'streaming',
|
||||
metadata: { parts: [] },
|
||||
});
|
||||
await db
|
||||
.updateTable('aiChatMessages')
|
||||
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
|
||||
.where('id', '=', msg.id)
|
||||
.execute();
|
||||
|
||||
const swept = await messageRepo.sweepStreamingWithoutActiveRun(
|
||||
15 * 60 * 1000,
|
||||
);
|
||||
expect(swept).toBeGreaterThanOrEqual(1);
|
||||
const row = await messageRepo.findById(msg.id, workspaceId);
|
||||
expect(row!.status).toBe('aborted');
|
||||
expect((row!.metadata as Record<string, unknown>).finalizeFailed).toBe(true);
|
||||
});
|
||||
|
||||
it('clause (d) is DOUBLE-GATED: a stale streaming row WITH an active run on the chat is left alone', async () => {
|
||||
const chatId = await newChat();
|
||||
const msg = await createMessage(db, {
|
||||
workspaceId,
|
||||
chatId,
|
||||
role: 'assistant',
|
||||
status: 'streaming',
|
||||
metadata: { parts: [] },
|
||||
});
|
||||
await db
|
||||
.updateTable('aiChatMessages')
|
||||
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
|
||||
.where('id', '=', msg.id)
|
||||
.execute();
|
||||
// An ACTIVE run on the same chat -> clause (d) must NOT touch the message.
|
||||
const run = await runRepo.insert({
|
||||
chatId,
|
||||
workspaceId,
|
||||
createdBy: userId,
|
||||
status: 'running',
|
||||
});
|
||||
|
||||
await messageRepo.sweepStreamingWithoutActiveRun(15 * 60 * 1000);
|
||||
expect((await messageRepo.findById(msg.id, workspaceId))!.status).toBe(
|
||||
'streaming',
|
||||
);
|
||||
await runRepo.finalizeIfActive(run.id, workspaceId, {
|
||||
status: 'aborted',
|
||||
error: null,
|
||||
});
|
||||
});
|
||||
|
||||
it('"kill DB on finish" recovery: after the DB is back, reconcile leaves NEITHER the row nor the run stuck', async () => {
|
||||
// Simulate a process that seeded the assistant row + run, then died before
|
||||
// finalizing EITHER (a mid-turn crash): a streaming message + a running run,
|
||||
// both stale, with no in-memory entry (fresh service = fresh maps).
|
||||
const chatId = await newChat();
|
||||
const msg = await createMessage(db, {
|
||||
workspaceId,
|
||||
chatId,
|
||||
role: 'assistant',
|
||||
status: 'streaming',
|
||||
metadata: { parts: [{ type: 'text', text: 'partial' }] },
|
||||
});
|
||||
const run = await runRepo.insert({
|
||||
chatId,
|
||||
workspaceId,
|
||||
createdBy: userId,
|
||||
status: 'running',
|
||||
assistantMessageId: msg.id,
|
||||
});
|
||||
await db
|
||||
.updateTable('aiChatRuns')
|
||||
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
|
||||
.where('id', '=', run.id)
|
||||
.execute();
|
||||
await db
|
||||
.updateTable('aiChatMessages')
|
||||
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
|
||||
.where('id', '=', msg.id)
|
||||
.execute();
|
||||
|
||||
// Reconcile (as the periodic job would): (c) aborts the orphan run, then
|
||||
// (b) settles the message from the now-terminal run.
|
||||
await runService.reconcileStaleRuns(15 * 60 * 1000);
|
||||
const stuck = await messageRepo.findStreamingWithTerminalRun();
|
||||
for (const s of stuck) {
|
||||
const status = s.runStatus === 'failed' ? 'error' : 'aborted';
|
||||
await messageRepo.stampTerminalIfStreaming(s.messageId, s.workspaceId, status);
|
||||
}
|
||||
|
||||
// Neither is stuck: the run is terminal AND the message is terminal.
|
||||
expect((await runRepo.findById(run.id, workspaceId))!.status).toBe('aborted');
|
||||
const row = await messageRepo.findById(msg.id, workspaceId);
|
||||
expect(row!.status).toBe('aborted');
|
||||
expect((row!.metadata as Record<string, unknown>).finalizeFailed).toBe(true);
|
||||
});
|
||||
});
|
||||
@@ -281,6 +281,52 @@ describe('AiChatRun durable lifecycle [integration]', () => {
|
||||
});
|
||||
});
|
||||
|
||||
it('#487 finalizeIfActive is CONDITIONAL: a late terminal write cannot clobber the settled status (real SQL)', async () => {
|
||||
const c = (await createChat(db, { workspaceId, creatorId: userId })).id;
|
||||
const run = await runRepo.insert({
|
||||
chatId: c,
|
||||
workspaceId,
|
||||
createdBy: userId,
|
||||
status: 'running',
|
||||
});
|
||||
|
||||
// First terminal write: the run IS active, so it flips + returns the row.
|
||||
const first = await runRepo.finalizeIfActive(run.id, workspaceId, {
|
||||
status: 'succeeded',
|
||||
error: null,
|
||||
});
|
||||
expect(first!.status).toBe('succeeded');
|
||||
expect(first!.finishedAt).toBeTruthy();
|
||||
|
||||
// A late/second writer tries to flip it to 'aborted' — the WHERE status IN
|
||||
// ('pending','running') guard matches NOTHING now, so it is a benign no-op.
|
||||
const second = await runRepo.finalizeIfActive(run.id, workspaceId, {
|
||||
status: 'aborted',
|
||||
error: 'late clobber attempt',
|
||||
});
|
||||
expect(second).toBeUndefined();
|
||||
|
||||
// The persisted terminal status is UNCHANGED — last-writer-wins is gone.
|
||||
const row = await runRepo.findById(run.id, workspaceId);
|
||||
expect(row!.status).toBe('succeeded');
|
||||
expect(row!.error).toBeNull();
|
||||
});
|
||||
|
||||
it('#487 double-settle through the service collapses to one write at the SQL gate', async () => {
|
||||
const c = (await createChat(db, { workspaceId, creatorId: userId })).id;
|
||||
const handle = await service.beginRun({ chatId: c, workspaceId, userId });
|
||||
|
||||
// First settle writes 'aborted' via the conditional write.
|
||||
await service.finalizeRun(handle.runId, workspaceId, 'aborted');
|
||||
// A late safety-net settle to 'error' is a no-op (row already terminal).
|
||||
await service.finalizeRun(handle.runId, workspaceId, 'error', 'late');
|
||||
|
||||
const row = await runRepo.findById(handle.runId, workspaceId);
|
||||
expect(row!.status).toBe('aborted');
|
||||
expect(service.isLocallyActive(handle.runId)).toBe(false);
|
||||
expect(service.hasZombie(handle.runId)).toBe(false);
|
||||
});
|
||||
|
||||
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
|
||||
|
||||
+174
-78
@@ -8,19 +8,35 @@ real pain (a "which tools fail most?" analysis that confidently answered
|
||||
|
||||
Read the **Gotchas** section before you trust any error count.
|
||||
|
||||
> **TWO ERAS — check the marker first.** The `tool_calls` shape changed in **#490
|
||||
> (trace v2)**. A row written by v2 carries `metadata.toolTraceVersion = 2`; older
|
||||
> rows have no such key. The two shapes store DIFFERENT things (v2 dropped the tool
|
||||
> OUTPUT from the trace), so **every query below is dual-shape** — branch on the
|
||||
> marker. **Never compare an aggregate or trend across the era boundary**: a metric
|
||||
> jump on the cut-over week is an artifact of the shape change, not a behavior
|
||||
> change.
|
||||
|
||||
## TL;DR
|
||||
|
||||
- Agent chats live in Postgres, DB `docmost`, tables `ai_chat_*`.
|
||||
- Each tool invocation is stored as **two** array elements (a `tool-call` part and
|
||||
a `tool-result` part), so naive counting double-counts.
|
||||
- **A tool that *throws* writes no result part.** Since the #407 fix its error is
|
||||
persisted as a dedicated `{toolName, error}` element in `tool_calls` (queryable +
|
||||
replayed to the model). **Rows written before #407 still drop it** — the error is
|
||||
nowhere in the DB and shows only in the live UI. So `isError` / `success=false`
|
||||
scans under-report by design, and pre-#407 thrown errors are invisible.
|
||||
- To find where agents fail: (1) soft-failure markers in `tool_calls`, (2) the new
|
||||
`error` field for thrown errors (new rows) / the orphan-gap proxy (old rows),
|
||||
(3) server logs / the live UI for full stack traces beyond the truncated message.
|
||||
- **Era marker:** `metadata.toolTraceVersion = 2` ⇒ v2 (#490) row; absent ⇒ legacy row.
|
||||
- Each tool invocation is stored as **two** consecutive array elements — a
|
||||
`tool-call` part then an OUTCOME part — so naive counting double-counts.
|
||||
- **v2 (#490):** outcome is `{toolName, ok: true}` on success, or
|
||||
`{toolName, error, kind: 'thrown'|'interrupted'}` on failure. The tool **OUTPUT
|
||||
is NOT in `tool_calls`** any more — it lives once in `metadata.parts` (this
|
||||
removed a hundreds-of-MB-per-run write duplication). Soft-failure analysis
|
||||
therefore reads `metadata.parts`, not `tool_calls`.
|
||||
- **legacy:** outcome is `{toolName, output}` on success; a **thrown** failure is
|
||||
a `{toolName, error}` element **only on rows after #407**, and is dropped
|
||||
entirely (silent orphan) on pre-#407 rows.
|
||||
- **A tool that *throws* writes no result part.** In v2 it is a
|
||||
`{error, kind:'thrown'}` element; an interrupted/aborted call is a distinct
|
||||
`{error, kind:'interrupted'}`. `isError`/`success=false` scans read the *output*
|
||||
and so under-report thrown failures in every era.
|
||||
- To find where agents fail: (1) soft-failure markers in `metadata.parts` outputs
|
||||
(v2) / `tool_calls` outputs (legacy), (2) the `error`/`kind` fields for thrown
|
||||
failures (v2 + post-#407), (3) server logs / the live UI for full stack traces.
|
||||
|
||||
## Where the data lives
|
||||
|
||||
@@ -53,33 +69,67 @@ are rows in `workspaces`, not separate deployments.
|
||||
separate `tool` role), `content` (text), `tool_calls` (jsonb array), `metadata`
|
||||
(jsonb, holds run `error` + rendered `parts`), `status`, `tsv` (full-text index).
|
||||
|
||||
## Era marker — check this before every query
|
||||
|
||||
```sql
|
||||
-- how many rows are in each era?
|
||||
SELECT COALESCE((metadata->>'toolTraceVersion'), 'legacy') AS era, count(*)
|
||||
FROM ai_chat_messages
|
||||
WHERE role = 'assistant' AND jsonb_typeof(tool_calls) = 'array'
|
||||
GROUP BY 1 ORDER BY 2 DESC;
|
||||
```
|
||||
|
||||
- `toolTraceVersion = '2'` → **v2** (#490): outcome flags, **no output in the trace**.
|
||||
- `NULL` (`'legacy'`) → pre-#490: outcome carries the tool `output` inline.
|
||||
|
||||
**Do not trend a metric across the cut-over.** The shape change alone shifts counts
|
||||
(e.g. "elements with `output`" collapses to zero for v2), so a week that straddles
|
||||
the boundary shows an artifact, not a behavior change. Segment by era, or restrict to
|
||||
one era, before comparing.
|
||||
|
||||
## How tool calls are stored — READ THIS
|
||||
|
||||
Tool calls are **not** one-object-per-call. Each logical invocation is split into
|
||||
two consecutive elements of the `tool_calls` array:
|
||||
two consecutive elements of the `tool_calls` array — a **call** then an **outcome**.
|
||||
The outcome shape is era-dependent:
|
||||
|
||||
```text
|
||||
index 0: { "toolName": "getPage", "input": { "pageId": "…" } } ← tool-call (has input, NO output)
|
||||
index 1: { "toolName": "getPage", "output": { … } } ← tool-result (has output, NO input)
|
||||
# v2 (#490) — metadata.toolTraceVersion = 2
|
||||
index 0: { "toolName":"getPage", "input":{...} } ← call (has input)
|
||||
index 1: { "toolName":"getPage", "ok":true } ← success (NO output here)
|
||||
or : { "toolName":"getPage", "error":"…", "kind":"thrown" } ← threw
|
||||
or : { "toolName":"getPage", "error":"…", "kind":"interrupted" } ← aborted mid-step
|
||||
|
||||
# legacy — no toolTraceVersion
|
||||
index 0: { "toolName":"getPage", "input":{...} } ← call (has input, NO output)
|
||||
index 1: { "toolName":"getPage", "output":{...} } ← success (has output)
|
||||
or : { "toolName":"getPage", "error":"…" } ← threw (post-#407 only)
|
||||
```
|
||||
|
||||
The keys that appear on an element are `toolName`, `input`, `output`, and — for a
|
||||
**thrown** failure on rows written after the #407 fix — `error` (the tool's error
|
||||
message; see the "Hard failures" section below). There is no `state`, no `errorText`,
|
||||
no `type`. On pre-#407 rows a thrown failure has NO paired result element at all
|
||||
(silent orphan). Consequences:
|
||||
The keys that can appear: `toolName`, `input` (call), and on the outcome — **v2:**
|
||||
`ok` **or** `error`+`kind`; **legacy:** `output` **or** (post-#407) `error`. There is
|
||||
no `state`, no `errorText`, no `type` in `tool_calls` (those live on `metadata.parts`).
|
||||
Consequences:
|
||||
|
||||
1. **Real invocation count = elements that have `output` or `error`.** Counting every
|
||||
element double-counts (you get ~2× and a spurious "~50% of every tool has no output").
|
||||
2. **Pairing:** a call = a `tool-call` part followed by its result part. A success
|
||||
carries `output`; a thrown failure (post-#407) carries `error` instead. Both carry
|
||||
`toolName`, so you can group by tool on either.
|
||||
1. **Real invocation count** — count the OUTCOME elements, not every element (else you
|
||||
double-count): **v2** = elements with `ok` or `error`; **legacy** = elements with
|
||||
`output` or `error`.
|
||||
2. **Pairing:** a call (`input`) is followed by its outcome. `toolName` is on both, so
|
||||
you can group by tool on either. In v2 the `kind` field separates a real hard-fail
|
||||
(`thrown`) from an aborted call (`interrupted`) — a distinction legacy rows cannot
|
||||
make (both are orphans; see below).
|
||||
3. **The tool OUTPUT is only in `metadata.parts` on v2 rows.** To inspect what a tool
|
||||
returned (soft-error markers, page bodies) on a v2 row, read the parts
|
||||
(`part->>'type' LIKE 'tool-%'`, `part->>'state' = 'output-available'`, `part->'output'`),
|
||||
not `tool_calls`.
|
||||
|
||||
## The two classes of failure (and which the DB can see)
|
||||
|
||||
### 1. Soft failures — tool RAN and returned an error-shaped result → PERSISTED ✅
|
||||
|
||||
These are visible in the `tool-result` `output`. The marker differs per tool:
|
||||
These are visible in the tool `output` — **on v2 rows in `metadata.parts`** (the
|
||||
`output-available` part's `output`), on **legacy rows in the `tool_calls` outcome
|
||||
element's `output`**. The marker differs per tool:
|
||||
|
||||
| Tool(s) | Error marker in `output` |
|
||||
| --- | --- |
|
||||
@@ -91,37 +141,32 @@ These are visible in the `tool-result` `output`. The marker differs per tool:
|
||||
Note `editPageText` returns `failed: []` on success — filtering on the *presence*
|
||||
of the key gives false positives; filter on **non-empty**.
|
||||
|
||||
### 2. Hard failures — tool THREW → NOW PERSISTED ✅ (since the #407 fix)
|
||||
### 2. Hard failures — tool THREW → PERSISTED ✅
|
||||
|
||||
When a tool throws (the classic one is `patchNode` / `insertNode` / `tableUpdateCell`
|
||||
→ `Failed to encode document to Yjs (fromJSON): Unknown node type: undefined`), the
|
||||
runtime still writes **no `tool-result` part** — the failure is an ai@6 `tool-error`
|
||||
content part instead. **Since the #407 fix, that error is persisted**: `serializeSteps`
|
||||
appends a dedicated element `{toolName, error: "<message>"}` right after the failed
|
||||
call, mirroring how a successful `{toolName, output}` element is appended. So a thrown
|
||||
error now leaves a queryable `error` field carrying its (truncated) reason, and the
|
||||
same real text is replayed to the model on the next turn (an `output-error` part with
|
||||
the real `errorText`, no longer the `'Tool call did not complete.'` placeholder).
|
||||
runtime writes **no `tool-result` part** — the failure is an ai@6 `tool-error` content
|
||||
part. How that lands in `tool_calls` depends on the era:
|
||||
|
||||
**Cutover caveat — old rows keep the old blind shape.** Rows written **before** this
|
||||
change have the two-part shape (`call` + `output` only) and simply **drop** thrown
|
||||
errors, leaving a silent **orphan** (a `call` with no `output` *and* no `error`). Rows
|
||||
written **after** the fix additionally carry the `error` element. So:
|
||||
- **v2 (#490):** a `{toolName, error, kind:'thrown'}` outcome element. An interrupted /
|
||||
aborted mid-step call is a **distinct** `{toolName, error:'Tool call did not
|
||||
complete.', kind:'interrupted'}` element — so you can tell a real hard-fail from an
|
||||
abort **directly, without the orphan heuristic**. Query `kind = 'thrown'`.
|
||||
- **post-#407 legacy:** a `{toolName, error}` element (no `kind`) right after the call.
|
||||
- **pre-#407 legacy:** the error is **dropped** — a silent **orphan** (a `call` with no
|
||||
`output` *and* no `error`).
|
||||
|
||||
- **New rows:** query the `error` field directly (see the hard-error query below) — no
|
||||
orphan heuristic needed for thrown failures.
|
||||
- **Old rows (pre-#407):** the only DB-side proxy is still an **orphan**: a `tool-call`
|
||||
part with no matching `tool-result` *and* no `error`. Orphans also appear when a run
|
||||
is **aborted** mid-flight (server restart), so a high-volume tool (`createComment`,
|
||||
`searchInPage`, `Search_web_search`) shows orphans from aborts, not real errors on
|
||||
old rows. Treat the orphan gap as an *upper bound*, and cross-check the tool: a gap on
|
||||
a structural editor (`patchNode`, `insertNode`, `updatePageJson`, `transformPage`) is
|
||||
almost certainly a thrown Yjs-encode error; a gap on `createComment` is mostly aborts.
|
||||
The same real error text is replayed to the model on the next turn (an `output-error`
|
||||
part with the real `errorText`, from `metadata.parts`), in every era.
|
||||
|
||||
A note on the aborted-call fallback: a call with **neither** a result **nor** a
|
||||
`tool-error` (genuinely interrupted mid-step) still replays with the
|
||||
`'Tool call did not complete.'` placeholder and persists as an orphan — that path is
|
||||
unchanged, and is distinct from a real thrown error, which now carries `error`.
|
||||
**Cutover caveat.** Only pre-#407 legacy rows need the orphan proxy: an orphan is a
|
||||
`tool-call` with no matching outcome. Orphans there also appear when a run is **aborted**
|
||||
mid-flight (server restart), so a high-volume tool (`createComment`, `searchInPage`,
|
||||
`Search_web_search`) shows orphans from aborts, not real errors. Treat the orphan gap as
|
||||
an *upper bound* and cross-check the tool: a gap on a structural editor (`patchNode`,
|
||||
`insertNode`, `updatePageJson`, `transformPage`) is almost certainly a thrown Yjs-encode
|
||||
error; a gap on `createComment` is mostly aborts. **On v2 rows this ambiguity is gone**
|
||||
— `kind` labels each outcome.
|
||||
|
||||
### 3. Run-level failures → `ai_chat_runs`
|
||||
|
||||
@@ -134,22 +179,34 @@ the wild: `Run interrupted by a server restart.` (aborts) and
|
||||
|
||||
Run all of these via `docker exec gitmost-postgresql psql -U docmost -d docmost -P pager=off -c "…"`.
|
||||
|
||||
**Real invocation count per tool** (result parts only — the correct denominator):
|
||||
**Real invocation count per tool** (outcome parts only — the correct denominator).
|
||||
Dual-shape: a v2 outcome has `ok` or `error`; a legacy outcome has `output` or `error`:
|
||||
|
||||
```sql
|
||||
SELECT elem->>'toolName' AS tool, count(*) AS calls
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
|
||||
WHERE jsonb_typeof(m.tool_calls) = 'array' AND elem ? 'output'
|
||||
WHERE jsonb_typeof(m.tool_calls) = 'array'
|
||||
AND (elem ? 'ok' OR elem ? 'output' OR elem ? 'error')
|
||||
GROUP BY 1 ORDER BY 2 DESC;
|
||||
```
|
||||
|
||||
**Soft errors per tool** (everything the DB can honestly see):
|
||||
**Soft errors per tool.** The soft-error marker lives in the tool OUTPUT — which on
|
||||
**v2 rows is in `metadata.parts`**, on **legacy rows is in the `tool_calls` outcome
|
||||
element**. This query UNIONs both eras, projecting each output as `o`:
|
||||
|
||||
```sql
|
||||
WITH res AS (
|
||||
-- v2 (#490): output is in metadata.parts (output-available tool parts)
|
||||
SELECT part->>'type' AS tool, part->'output' AS o
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.metadata->'parts') part
|
||||
WHERE (m.metadata->>'toolTraceVersion') = '2'
|
||||
AND part->>'type' LIKE 'tool-%' AND part->>'state' = 'output-available'
|
||||
UNION ALL
|
||||
-- legacy: output is inline in the tool_calls outcome element
|
||||
SELECT elem->>'toolName' AS tool, elem->'output' AS o
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
|
||||
WHERE jsonb_typeof(m.tool_calls) = 'array' AND elem ? 'output'
|
||||
WHERE (m.metadata->>'toolTraceVersion') IS NULL
|
||||
AND jsonb_typeof(m.tool_calls) = 'array' AND elem ? 'output'
|
||||
)
|
||||
SELECT tool, count(*) AS calls,
|
||||
sum(COALESCE(
|
||||
@@ -167,13 +224,23 @@ FROM res GROUP BY tool HAVING sum(COALESCE(
|
||||
ORDER BY soft_errors DESC;
|
||||
```
|
||||
|
||||
**`editPageText` failure reasons** (the most common real agent mistake — bad `find`):
|
||||
Note the v2 `tool` label is the part type (`tool-editPageText`); strip the `tool-`
|
||||
prefix if you join it against the legacy `toolName`.
|
||||
|
||||
**`editPageText` failure reasons** (the most common real agent mistake — bad `find`).
|
||||
Same dual-shape output source:
|
||||
|
||||
```sql
|
||||
WITH res AS (
|
||||
SELECT part->'output' AS o
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.metadata->'parts') part
|
||||
WHERE (m.metadata->>'toolTraceVersion') = '2'
|
||||
AND part->>'type' = 'tool-editPageText' AND part->>'state' = 'output-available'
|
||||
UNION ALL
|
||||
SELECT elem->'output' AS o
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
|
||||
WHERE jsonb_typeof(m.tool_calls) = 'array'
|
||||
WHERE (m.metadata->>'toolTraceVersion') IS NULL
|
||||
AND jsonb_typeof(m.tool_calls) = 'array'
|
||||
AND elem->>'toolName' = 'editPageText' AND elem ? 'output'
|
||||
)
|
||||
SELECT f->>'reason' AS reason, count(*)
|
||||
@@ -182,30 +249,43 @@ WHERE jsonb_typeof(o->'failed') = 'array'
|
||||
GROUP BY 1 ORDER BY 2 DESC;
|
||||
```
|
||||
|
||||
**Hard errors — persisted `error` field per tool (NEW rows, since #407)** — thrown
|
||||
tool failures now carry their real reason, so query them directly:
|
||||
**Hard errors — persisted `error` field per tool (v2 + post-#407 rows)** — thrown tool
|
||||
failures carry their real reason, so query them directly. On **v2** rows exclude the
|
||||
`interrupted` kind so an aborted call is not counted as a hard-fail:
|
||||
|
||||
```sql
|
||||
SELECT elem->>'toolName' AS tool, count(*) AS thrown_errors,
|
||||
min(elem->>'error') AS sample_error
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
|
||||
WHERE jsonb_typeof(m.tool_calls) = 'array' AND elem ? 'error'
|
||||
-- v2 rows label the kind; a legacy error element has no kind (count it).
|
||||
AND COALESCE(elem->>'kind', 'thrown') = 'thrown'
|
||||
GROUP BY 1 ORDER BY 2 DESC;
|
||||
```
|
||||
|
||||
Aborted mid-step calls on v2 rows are a distinct, directly countable population:
|
||||
|
||||
```sql
|
||||
SELECT elem->>'toolName' AS tool, count(*) AS interrupted
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
|
||||
WHERE jsonb_typeof(m.tool_calls) = 'array' AND elem->>'kind' = 'interrupted'
|
||||
GROUP BY 1 ORDER BY 2 DESC;
|
||||
```
|
||||
|
||||
**Hard-error proxy for OLD rows (pre-#407) — orphan gap per tool, WITH a spread column**
|
||||
(call parts minus result parts, plus how many distinct chats the gap is spread across).
|
||||
This covers rows written before thrown errors were persisted; on new rows a thrown
|
||||
failure now has its own `error` element (use the query above) and an orphan means only
|
||||
a genuinely aborted mid-step call:
|
||||
(call parts minus outcome parts, plus how many distinct chats the gap is spread across).
|
||||
This is needed ONLY for pre-#407 legacy rows (v2 and post-#407 rows carry the error /
|
||||
`kind` directly — use the queries above). The `WHERE` restricts to the legacy era so v2
|
||||
rows (where an `ok` outcome is not an `output`) never produce phantom orphans:
|
||||
|
||||
```sql
|
||||
WITH parts AS (
|
||||
SELECT m.chat_id, elem->>'toolName' AS tool,
|
||||
(elem ? 'input' AND NOT (elem ? 'output')) AS is_call,
|
||||
(elem ? 'output' OR elem ? 'error') AS is_result
|
||||
(elem ? 'input' AND NOT (elem ? 'output') AND NOT (elem ? 'ok')) AS is_call,
|
||||
(elem ? 'output' OR elem ? 'error' OR elem ? 'ok') AS is_result
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
|
||||
WHERE jsonb_typeof(m.tool_calls) = 'array' AND m.role = 'assistant'
|
||||
AND (m.metadata->>'toolTraceVersion') IS NULL
|
||||
),
|
||||
per_chat AS (
|
||||
SELECT tool, chat_id, sum(is_call::int) - sum(is_result::int) AS gap
|
||||
@@ -261,11 +341,21 @@ WHERE tsv @@ websearch_to_tsquery('english', 'some phrase') LIMIT 20;
|
||||
|
||||
## Don't blow up your context
|
||||
|
||||
A single `tool_calls` row can be **300–400 KB** (results embed full page content and
|
||||
search payloads). Never `SELECT tool_calls` (or `jsonb_pretty(tool_calls)`) raw.
|
||||
Always project just the keys you need and truncate:
|
||||
Tool outputs embed full page content and search payloads (hundreds of KB per row).
|
||||
On **legacy** rows they are in `tool_calls`; on **v2** rows they moved to
|
||||
`metadata->'parts'` (the `tool_calls` trace itself is now small). Never `SELECT
|
||||
tool_calls` / `metadata` (or `jsonb_pretty(...)`) raw — project just the keys you need
|
||||
and truncate:
|
||||
|
||||
```sql
|
||||
-- v2: outputs live in metadata.parts
|
||||
SELECT part->>'type',
|
||||
left(regexp_replace((part->'output')::text, '\s+', ' ', 'g'), 200)
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.metadata->'parts') part
|
||||
WHERE (m.metadata->>'toolTraceVersion') = '2'
|
||||
AND part->>'state' = 'output-available' LIMIT 5;
|
||||
|
||||
-- legacy: outputs live in tool_calls
|
||||
SELECT elem->>'toolName',
|
||||
left(regexp_replace((elem->'output')::text, '\s+', ' ', 'g'), 200)
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
|
||||
@@ -280,26 +370,32 @@ docker compose -p gitmost logs -f --tail=100 # whole stack
|
||||
```
|
||||
|
||||
Logging is `json-file`, `max-size=10m max-file=5` → ~50 MB retained, then rotated,
|
||||
and **wiped on container recreate**. Since the #407 fix, thrown-tool error text is
|
||||
**persisted in the `error` field** of `tool_calls` (see the hard-error query above), so
|
||||
you no longer depend on live logs for it. Logs/live UI remain useful for **pre-#407
|
||||
rows** (whose thrown errors were dropped) and for full stack traces beyond the
|
||||
truncated stored message. A per-tool `tool_calls_total{tool,status}` metric to
|
||||
VictoriaMetrics is still a possible future add for aggregate dashboards.
|
||||
and **wiped on container recreate**. Thrown-tool error text is **persisted** — in the
|
||||
`error` field of `tool_calls` (v2 `kind:'thrown'` / post-#407 legacy) — so you no longer
|
||||
depend on live logs for it. Logs/live UI remain useful for **pre-#407 rows** (whose
|
||||
thrown errors were dropped) and for full stack traces beyond the truncated stored
|
||||
message. A per-tool `tool_calls_total{tool,status}` metric to VictoriaMetrics is still a
|
||||
possible future add for aggregate dashboards.
|
||||
|
||||
## Gotchas checklist
|
||||
|
||||
- [ ] Counting every `tool_calls` element → **overcount**. Count `output` elements; add `error` elements for thrown failures (new rows), but don't count both as invocations.
|
||||
- [ ] `isError` / `success=false` ≈ 0 does **not** mean "no errors" — thrown errors are a separate `error` element (new rows) or dropped entirely (pre-#407 rows).
|
||||
- [ ] Thrown errors persist only on rows written **after the #407 fix** — pre-#407 rows still drop them (orphan only). Mind the cutover when trending over time.
|
||||
- [ ] **Check `metadata.toolTraceVersion` first.** v2 (`= 2`) has no output in `tool_calls`; legacy has it inline. Never trend a metric across the era boundary.
|
||||
- [ ] Counting every `tool_calls` element → **overcount**. Count OUTCOME elements — v2: `ok` or `error`; legacy: `output` or `error` — never both call+outcome as invocations.
|
||||
- [ ] `isError` / `success=false` ≈ 0 does **not** mean "no errors" — thrown errors are an `error` element (v2 `kind:'thrown'` / post-#407), not in the output.
|
||||
- [ ] **v2:** soft-error markers (the tool output) are in `metadata.parts`, NOT `tool_calls`. Legacy: they are in the `tool_calls` outcome `output`.
|
||||
- [ ] **v2:** `kind` splits a real hard-fail (`thrown`) from an aborted call (`interrupted`) directly — no orphan heuristic needed. The orphan gap is a pre-#407-legacy-only proxy.
|
||||
- [ ] `editPageText.failed` is `[]` on success — test for **non-empty**, not presence.
|
||||
- [ ] Orphan gap on OLD rows mixes thrown errors **and** aborted runs — split by tool. On NEW rows a thrown error is its own `error` element, so a gap ≈ aborted call.
|
||||
- [ ] `aborted` runs = server restarts, `failed` runs = provider overload — not agent mistakes.
|
||||
- [ ] Never dump a raw `tool_calls` cell — it can be hundreds of KB.
|
||||
- [ ] Logs are ephemeral (≤50 MB, wiped on recreate) — grab hard-error text live.
|
||||
- [ ] Never dump a raw `tool_calls` **or** `metadata.parts` cell — outputs are hundreds of KB.
|
||||
- [ ] Logs are ephemeral (≤50 MB, wiped on recreate) — grab pre-#407 hard-error text live.
|
||||
|
||||
## Snapshot (2026-07-07, illustrative — rerun the queries for current numbers)
|
||||
|
||||
> All rows in this snapshot predate #490, so they are **legacy-era** (outputs inline in
|
||||
> `tool_calls`, orphan proxy for thrown errors). Do not trend these numbers against v2
|
||||
> rows — segment by `toolTraceVersion` first.
|
||||
|
||||
|
||||
- 226 chats, 732 messages, 46 runs; ~4 400 real tool invocations.
|
||||
- Soft errors (persisted): `editPageText` 4/79 (bad/non-unique `find`) + 9 markdown-in-`find` warnings; `semanticSearch` 3/4 (`unavailable`); `Habr_update_draft_from_docmost` 1/2 (`doc` sent as object, not string).
|
||||
- Missing-result proxy, read WITH the spread column:
|
||||
|
||||
@@ -1,60 +1,100 @@
|
||||
// Codegen: emit src/registry-stamp.generated.ts with a REGISTRY_STAMP hash of
|
||||
// the tool-specs REGISTRY CONTENT, so a build/ vs src/ skew (issue #447) is
|
||||
// detectable at runtime.
|
||||
// the ENTIRE src/ tree, so a build/ vs src/ skew (issue #447) is detectable at
|
||||
// runtime for ANY source file — not just tool-specs.ts.
|
||||
//
|
||||
// WHY hash the raw source text (not extracted structured data):
|
||||
// SHARED_TOOL_SPECS carries `buildShape` functions (the input SCHEMAS) which are
|
||||
// NOT serializable. The input schema is exactly one of the things that MUST stay
|
||||
// in sync between build/ and src/, so we cannot drop it from the hash. Rather
|
||||
// than probe zod with a fragile shim to reconstruct the schema shape, we hash the
|
||||
// STABLE, deterministic source TEXT of tool-specs.ts. That text fully captures
|
||||
// every field that must stay in sync — mcpName, inAppKey, description, tier,
|
||||
// catalogLine AND the buildShape bodies (input schemas) — with zero probing
|
||||
// fragility. Any edit to a spec (a renamed tool, a reworded description, a
|
||||
// changed schema field) changes the text and therefore the stamp.
|
||||
// WHY hash the whole src tree (not just tool-specs.ts): the runtime tools are
|
||||
// assembled from far more than the spec registry — client.ts, the client/*
|
||||
// domain modules, comment-signal.ts and the drawio-* helpers all ship in build/
|
||||
// and are loaded by the in-app server. Hashing ONLY tool-specs.ts meant an edit
|
||||
// to any of those (e.g. a behavioural fix in client.ts) left the stamp unchanged,
|
||||
// so a stale build/ served the OLD code silently (issue #486). Hashing every
|
||||
// src/**/*.ts closes that gap: any source edit changes the stamp.
|
||||
//
|
||||
// DETERMINISM: the hash is computed over the file bytes with line endings
|
||||
// normalized to LF and a single trailing newline stripped, so a CRLF checkout or
|
||||
// an editor's trailing-newline habit cannot make build/ and src/ disagree. No
|
||||
// Date.now / randomness. The loader's dev-only stale-check (docmost-client.loader.ts)
|
||||
// re-runs THIS SAME normalization + sha256 over src/tool-specs.ts and compares to
|
||||
// the built REGISTRY_STAMP; the two must compute identically.
|
||||
// WHY hash the raw source text (not extracted structured data): the tool input
|
||||
// SCHEMAS live as `buildShape` functions which are NOT serializable, so we cannot
|
||||
// reduce them to structured data without a fragile zod shim. Hashing the STABLE,
|
||||
// deterministic source TEXT captures every field that must stay in sync with zero
|
||||
// probing fragility. Any edit to any source file changes the text → the stamp.
|
||||
//
|
||||
// DETERMINISM: files are enumerated recursively, filtered to *.ts EXCLUDING
|
||||
// *.generated.ts (the codegen's OWN output — including it would create a
|
||||
// fixed-point cycle), and sorted by their POSIX-normalized path relative to src/
|
||||
// so the order is platform-independent. Each file contributes its relative path
|
||||
// AND its content with line endings normalized to LF and a single trailing
|
||||
// newline stripped, so a CRLF checkout or an editor's trailing-newline habit
|
||||
// cannot make build/ and src/ disagree. No Date.now / randomness. The loader's
|
||||
// dev-only stale-check (docmost-client.loader.ts) re-runs THIS SAME enumeration +
|
||||
// normalization + sha256 and compares to the built REGISTRY_STAMP; the two must
|
||||
// compute identically.
|
||||
//
|
||||
// This script runs from the `build` and `pretest` npm scripts BEFORE tsc, so
|
||||
// build/ always carries a stamp derived from the tool-specs.ts that was compiled.
|
||||
// build/ always carries a stamp derived from the src/ tree that was compiled.
|
||||
|
||||
import { createHash } from 'node:crypto';
|
||||
import { readFileSync, writeFileSync } from 'node:fs';
|
||||
import { readdirSync, readFileSync, statSync, writeFileSync } from 'node:fs';
|
||||
import { fileURLToPath } from 'node:url';
|
||||
import { dirname, join } from 'node:path';
|
||||
import { dirname, join, relative, sep } from 'node:path';
|
||||
|
||||
const __dirname = dirname(fileURLToPath(import.meta.url));
|
||||
const SRC_DIR = join(__dirname, '..', 'src');
|
||||
const TOOL_SPECS_PATH = join(SRC_DIR, 'tool-specs.ts');
|
||||
const OUT_PATH = join(SRC_DIR, 'registry-stamp.generated.ts');
|
||||
|
||||
/**
|
||||
* Deterministic stamp of the tool-specs registry content. Kept as a plain
|
||||
* function (exported) so the algorithm has a single home; the loader duplicates
|
||||
* only the tiny normalize+sha256 steps because it lives in the CJS server build
|
||||
* and cannot import this ESM script. If you change the normalization here, mirror
|
||||
* it in apps/server/src/core/ai-chat/tools/docmost-client.loader.ts.
|
||||
* Recursively enumerate every `*.ts` file under `dir`, EXCLUDING the codegen's
|
||||
* own `*.generated.ts` output (a self-referential cycle otherwise). Returns
|
||||
* absolute paths, unsorted (the caller sorts by relative path for determinism).
|
||||
* Kept as a plain exported function so the algorithm has a single home; the
|
||||
* loader duplicates it because it lives in the CJS server build and cannot import
|
||||
* this ESM script. If you change the walk/filter here, mirror it in
|
||||
* apps/server/src/core/ai-chat/tools/docmost-client.loader.ts.
|
||||
*/
|
||||
export function computeRegistryStamp(toolSpecsSource) {
|
||||
const normalized = toolSpecsSource.replace(/\r\n/g, '\n').replace(/\n$/, '');
|
||||
return createHash('sha256').update(normalized, 'utf8').digest('hex');
|
||||
export function collectStampFiles(dir) {
|
||||
const out = [];
|
||||
for (const entry of readdirSync(dir)) {
|
||||
const full = join(dir, entry);
|
||||
if (statSync(full).isDirectory()) {
|
||||
out.push(...collectStampFiles(full));
|
||||
} else if (entry.endsWith('.ts') && !entry.endsWith('.generated.ts')) {
|
||||
out.push(full);
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/**
|
||||
* Deterministic stamp of the whole src/ tree. Enumerate + sort by POSIX-relative
|
||||
* path, then fold each file's relative path AND normalized content into one
|
||||
* sha256. MUST stay byte-for-byte identical to the loader's recompute.
|
||||
*/
|
||||
export function computeRegistryStamp(srcDir) {
|
||||
const files = collectStampFiles(srcDir)
|
||||
.map((abs) => ({
|
||||
rel: relative(srcDir, abs).split(sep).join('/'),
|
||||
abs,
|
||||
}))
|
||||
.sort((a, b) => (a.rel < b.rel ? -1 : a.rel > b.rel ? 1 : 0));
|
||||
const hash = createHash('sha256');
|
||||
for (const { rel, abs } of files) {
|
||||
const normalized = readFileSync(abs, 'utf8')
|
||||
.replace(/\r\n/g, '\n')
|
||||
.replace(/\n$/, '');
|
||||
hash.update(rel, 'utf8');
|
||||
hash.update('\0', 'utf8');
|
||||
hash.update(normalized, 'utf8');
|
||||
hash.update('\0', 'utf8');
|
||||
}
|
||||
return hash.digest('hex');
|
||||
}
|
||||
|
||||
function main() {
|
||||
const source = readFileSync(TOOL_SPECS_PATH, 'utf8');
|
||||
const stamp = computeRegistryStamp(source);
|
||||
const stamp = computeRegistryStamp(SRC_DIR);
|
||||
const out =
|
||||
'// AUTO-GENERATED by scripts/gen-registry-stamp.mjs — DO NOT EDIT BY HAND.\n' +
|
||||
'// A deterministic hash of src/tool-specs.ts content (tool names, descriptions,\n' +
|
||||
'// tiers, catalog lines and input schemas). Regenerated on every build/pretest\n' +
|
||||
'// so build/ always matches the compiled src. The in-app loader recomputes this\n' +
|
||||
'// from src and refuses to run on a mismatch (issue #447). This file is\n' +
|
||||
'// gitignored and produced by the build — see .gitignore.\n' +
|
||||
'// A deterministic hash of the whole src/ tree (every src/**/*.ts except\n' +
|
||||
'// *.generated.ts). Regenerated on every build/pretest so build/ always\n' +
|
||||
'// matches the compiled src. The in-app loader recomputes this from src and\n' +
|
||||
'// refuses to run on a mismatch (issue #447/#486). This file is gitignored\n' +
|
||||
'// and produced by the build — see .gitignore.\n' +
|
||||
`export const REGISTRY_STAMP = ${JSON.stringify(stamp)};\n`;
|
||||
writeFileSync(OUT_PATH, out, 'utf8');
|
||||
// eslint-disable-next-line no-console
|
||||
|
||||
@@ -59,6 +59,39 @@ import {
|
||||
mergeFootnoteDefinitions,
|
||||
} from "../lib/transforms.js";
|
||||
|
||||
// Max concurrent per-page comment fetches in checkNewComments (#490). The scan is
|
||||
// O(N) independent REST reads over the working set; running them one-at-a-time made
|
||||
// a large space linear in round-trips. A small cap parallelizes without hammering
|
||||
// the server (or exhausting sockets). 6 is a conservative middle of the 5–8 band.
|
||||
const COMMENT_SCAN_CONCURRENCY = 6;
|
||||
|
||||
/**
|
||||
* Map `items` through `fn` with at most `limit` in flight, preserving INPUT ORDER
|
||||
* in the returned array. A tiny bounded pool (no p-limit dependency): `limit`
|
||||
* workers pull the next index off a shared cursor until the list is drained.
|
||||
*/
|
||||
async function mapWithConcurrency<T, R>(
|
||||
items: readonly T[],
|
||||
limit: number,
|
||||
fn: (item: T, index: number) => Promise<R>,
|
||||
): Promise<R[]> {
|
||||
const results = new Array<R>(items.length);
|
||||
let cursor = 0;
|
||||
const worker = async (): Promise<void> => {
|
||||
for (;;) {
|
||||
const i = cursor++;
|
||||
if (i >= items.length) return;
|
||||
results[i] = await fn(items[i], i);
|
||||
}
|
||||
};
|
||||
const workers = Array.from(
|
||||
{ length: Math.max(1, Math.min(limit, items.length)) },
|
||||
() => worker(),
|
||||
);
|
||||
await Promise.all(workers);
|
||||
return results;
|
||||
}
|
||||
|
||||
// Public method surface of CommentsMixin (issue #450) — a NAMED type so the factory
|
||||
// return type is expressible in the emitted .d.ts (the anonymous mixin class
|
||||
// carries the base's protected shared state, which would otherwise trip TS4094).
|
||||
@@ -655,27 +688,32 @@ export function CommentsMixin<TBase extends GConstructor<DocmostClientContext>>(
|
||||
parentPageId,
|
||||
);
|
||||
|
||||
// 2. Fetch comments for each page, keep ones created after since
|
||||
const results: any[] = [];
|
||||
for (const page of pagesInScope) {
|
||||
try {
|
||||
// Full feed (incl. resolved): a "new comments since" scan reports all
|
||||
// recent activity; the active-only filter is scoped to listComments.
|
||||
const comments = (await this.listComments(page.id, true)).items;
|
||||
const newComments = comments.filter(
|
||||
(c: any) => new Date(c.createdAt) > sinceDate,
|
||||
);
|
||||
if (newComments.length > 0) {
|
||||
results.push({
|
||||
pageId: page.id,
|
||||
pageTitle: page.title,
|
||||
comments: newComments,
|
||||
});
|
||||
// 2. Fetch comments for each page, keep ones created after since. Runs with
|
||||
// bounded concurrency (#490) instead of one-at-a-time — the per-page reads are
|
||||
// independent, so a large working set no longer costs O(N) serial round-trips.
|
||||
// Order is preserved (mapWithConcurrency keeps input order), so the output is
|
||||
// deterministic regardless of which fetch finishes first.
|
||||
const perPage = await mapWithConcurrency(
|
||||
pagesInScope,
|
||||
COMMENT_SCAN_CONCURRENCY,
|
||||
async (page: any) => {
|
||||
try {
|
||||
// Full feed (incl. resolved): a "new comments since" scan reports all
|
||||
// recent activity; the active-only filter is scoped to listComments.
|
||||
const comments = (await this.listComments(page.id, true)).items;
|
||||
const newComments = comments.filter(
|
||||
(c: any) => new Date(c.createdAt) > sinceDate,
|
||||
);
|
||||
return newComments.length > 0
|
||||
? { pageId: page.id, pageTitle: page.title, comments: newComments }
|
||||
: null;
|
||||
} catch (e: any) {
|
||||
// Skip pages with errors (e.g. deleted between calls)
|
||||
return null;
|
||||
}
|
||||
} catch (e: any) {
|
||||
// Skip pages with errors (e.g. deleted between calls)
|
||||
}
|
||||
}
|
||||
},
|
||||
);
|
||||
const results: any[] = perPage.filter((r): r is any => r !== null);
|
||||
|
||||
const totalNewComments = results.reduce(
|
||||
(sum, r) => sum + r.comments.length,
|
||||
|
||||
@@ -19,7 +19,10 @@ import {
|
||||
assertYjsEncodable,
|
||||
MutationResult,
|
||||
} from "../lib/collaboration.js";
|
||||
import { acquireCollabSession } from "../lib/collab-session.js";
|
||||
import {
|
||||
acquireCollabSession,
|
||||
isCollabAuthFailedError,
|
||||
} from "../lib/collab-session.js";
|
||||
import { withPageLock, isUuid } from "../lib/page-lock.js";
|
||||
import { getCollabToken, performLogin } from "../lib/auth-utils.js";
|
||||
import { formatDocmostAxiosError } from "./errors.js";
|
||||
@@ -167,6 +170,43 @@ export abstract class DocmostClientContext {
|
||||
// cached conversion can never leak across identities. See getpage-cache.ts.
|
||||
protected getPageCache = new GetPageConversionCache();
|
||||
|
||||
// #487: an OPTIONAL abort signal the in-app tool host sets before each tool
|
||||
// call (a composite of the turn's Stop signal + a per-call wall-clock cap). It
|
||||
// is checked at safe-points BETWEEN the sequential HTTP calls of a paginated
|
||||
// read (paginateAll) and just before the atomic collab commit of a write (the
|
||||
// mutatePage/replacePage/mutateLiveContentUnlocked seams), so a Stop / cap
|
||||
// stops the NEXT network call from STARTING. An already-started single call may
|
||||
// still land — a documented limitation (#487).
|
||||
//
|
||||
// SINGLE-WRITER by phase-1 assumption: exactly one DocmostClient is built per
|
||||
// turn and shared by every tool call; the host sets this per call and does NOT
|
||||
// restore the prior value on unwind (set-and-leave) — a fresh client per turn
|
||||
// plus overwrite-by-the-next-call keeps it correct, and leaving a settled
|
||||
// call's signal in place is what makes a discarded race-loser throw on its
|
||||
// next safe-point. If the model emits PARALLEL in-app
|
||||
// tool calls they share this one field, so the per-call CAP of one call is not
|
||||
// guaranteed to bound another's in-flight pagination — but every composite the
|
||||
// host sets carries the SAME turn Stop signal, so a Stop still aborts whichever
|
||||
// signal is current. #487.
|
||||
protected toolAbortSignal: AbortSignal | null = null;
|
||||
|
||||
/**
|
||||
* #487: set (or clear with null) the in-app tool abort signal governing the
|
||||
* NEXT client call's safe-points. The host wraps each in-app tool call: it sets
|
||||
* the composite (Stop + per-call cap) here before invoking the tool and leaves
|
||||
* it in place afterwards (set-and-leave, NOT restored) — the next call
|
||||
* overwrites it, and a fresh client is built per turn. Public so the
|
||||
* server-side tool wrapper can reach it; harmless (a no-op) when never set.
|
||||
*/
|
||||
public setToolAbortSignal(signal: AbortSignal | null): void {
|
||||
this.toolAbortSignal = signal;
|
||||
}
|
||||
|
||||
/** #487: the abort signal currently governing this client's safe-points. */
|
||||
public getToolAbortSignal(): AbortSignal | null {
|
||||
return this.toolAbortSignal;
|
||||
}
|
||||
|
||||
// Two construction forms:
|
||||
// - new DocmostClient(config) // discriminated union (current)
|
||||
// - new DocmostClient(baseURL, email, password) // legacy positional creds
|
||||
@@ -492,6 +532,37 @@ export abstract class DocmostClientContext {
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Run a collab write and, on a Hocuspocus HANDSHAKE auth failure, self-heal
|
||||
* once (#486). Symmetric to the HTTP-401 path in getCollabTokenWithReauth: the
|
||||
* REST interceptor and login() already drop the cached collab token on a 401/
|
||||
* 403, but a rejected WEBSOCKET handshake left the stale token in the cache, so
|
||||
* every subsequent mutation kept re-presenting the same bad token for up to the
|
||||
* collab-token TTL (minutes) with no self-heal. Here, when the write rejects
|
||||
* with the tagged collab-auth error, we invalidate the cached token and retry
|
||||
* the write EXACTLY once with a force-refreshed token. Not a loop: a second
|
||||
* failure (or any non-auth error) propagates unchanged.
|
||||
*
|
||||
* `write` receives the token to use, so the retry can hand it a genuinely fresh
|
||||
* one rather than re-running with the same stale string.
|
||||
*/
|
||||
protected async writeWithCollabAuthRetry<T>(
|
||||
collabToken: string,
|
||||
write: (token: string) => Promise<T>,
|
||||
): Promise<T> {
|
||||
try {
|
||||
return await write(collabToken);
|
||||
} catch (e) {
|
||||
if (!isCollabAuthFailedError(e)) throw e;
|
||||
// The WS handshake rejected our token: drop it from the cache so it can't
|
||||
// be reused for the rest of the TTL, mint a fresh one (forceRefresh bypasses
|
||||
// the cache and re-invokes the provider/login), and retry the write once.
|
||||
this.collabTokenCache = null;
|
||||
const fresh = await this.getCollabTokenWithReauth(true);
|
||||
return await write(fresh);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Connect to the collaboration websocket, read the live doc, apply
|
||||
* `transform`, write the result, and wait for the server to persist it —
|
||||
@@ -526,19 +597,28 @@ export abstract class DocmostClientContext {
|
||||
// unsyncedChanges/connectionLost ack logic live in CollabSession.mutate,
|
||||
// preserved verbatim from the old inline machine (incl. the #152 structural
|
||||
// diff that keeps a live editor's cursor anchored).
|
||||
const session = await acquireCollabSession(pageId, collabToken, this.apiUrl, {
|
||||
// Only the actual 25s collab connect timeout emits this — the connect-vs-
|
||||
// unload signal; the other failure paths must NOT emit it.
|
||||
onConnectTimeout: () =>
|
||||
this.onMetricFn?.("collab_connect_timeouts_total", 1),
|
||||
// Wrap in the collab-auth self-heal (#486): a rejected WS handshake drops the
|
||||
// cached collab token and retries once with a fresh one (the retry passes the
|
||||
// refreshed token down to acquireCollabSession via `token`).
|
||||
return this.writeWithCollabAuthRetry(collabToken, async (token) => {
|
||||
const session = await acquireCollabSession(pageId, token, this.apiUrl, {
|
||||
// Only the actual 25s collab connect timeout emits this — the connect-vs-
|
||||
// unload signal; the other failure paths must NOT emit it.
|
||||
onConnectTimeout: () =>
|
||||
this.onMetricFn?.("collab_connect_timeouts_total", 1),
|
||||
});
|
||||
try {
|
||||
// #487 PRE-COMMIT safe-point (reentrant twin of mutatePageContent): a
|
||||
// Stop/cap after acquiring the session but before the atomic write skips
|
||||
// this commit. Same limitation applies (stops the NEXT commit only).
|
||||
this.toolAbortSignal?.throwIfAborted();
|
||||
return await session.mutate(transform);
|
||||
} catch (e) {
|
||||
// Drop the session on any failure so the next call reconnects fresh.
|
||||
session.destroy("mutate failed");
|
||||
throw e;
|
||||
}
|
||||
});
|
||||
try {
|
||||
return await session.mutate(transform);
|
||||
} catch (e) {
|
||||
// Drop the session on any failure so the next call reconnects fresh.
|
||||
session.destroy("mutate failed");
|
||||
throw e;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -563,6 +643,11 @@ export abstract class DocmostClientContext {
|
||||
let truncated = false;
|
||||
|
||||
for (let page = 0; page < MAX_PAGES; page++) {
|
||||
// #487 safe-point: a Stop (or the in-app tool per-call cap) that fires
|
||||
// BETWEEN sequential page fetches must stop the NEXT request from starting
|
||||
// — a read tool that would otherwise paginate for minutes is interrupted
|
||||
// here. throwIfAborted() rejects with the signal's reason.
|
||||
this.toolAbortSignal?.throwIfAborted();
|
||||
const payload: Record<string, any> = {
|
||||
...basePayload,
|
||||
limit: clampedLimit,
|
||||
@@ -667,7 +752,19 @@ export abstract class DocmostClientContext {
|
||||
transform: (doc: any) => any,
|
||||
): Promise<{ doc?: any; verify?: any }> {
|
||||
const pageUuid = await this.resolvePageId(pageId);
|
||||
return mutatePageContent(pageUuid, collabToken, apiUrl, transform);
|
||||
// #486: on a rejected collab-WS handshake, invalidate + refresh the token and
|
||||
// retry the write once (symmetric to the HTTP-401 reauth path).
|
||||
return this.writeWithCollabAuthRetry(collabToken, (token) =>
|
||||
// #487: thread the in-app tool signal to mutatePageContent's pre-commit
|
||||
// safe-point so a Stop/cap during the connect/lock window skips the write.
|
||||
mutatePageContent(
|
||||
pageUuid,
|
||||
token,
|
||||
apiUrl,
|
||||
transform,
|
||||
this.toolAbortSignal ?? undefined,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -687,7 +784,18 @@ export abstract class DocmostClientContext {
|
||||
apiUrl: string,
|
||||
): Promise<{ doc?: any; verify?: any }> {
|
||||
const pageUuid = await this.resolvePageId(pageId);
|
||||
return replacePageContent(pageUuid, doc, collabToken, apiUrl);
|
||||
// #486: on a rejected collab-WS handshake, invalidate + refresh the token and
|
||||
// retry the write once (symmetric to the HTTP-401 reauth path).
|
||||
return this.writeWithCollabAuthRetry(collabToken, (token) =>
|
||||
// #487: same pre-commit safe-point as mutatePage, for full-document writes.
|
||||
replacePageContent(
|
||||
pageUuid,
|
||||
doc,
|
||||
token,
|
||||
apiUrl,
|
||||
this.toolAbortSignal ?? undefined,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -127,8 +127,13 @@ export function ReadMixin<TBase extends GConstructor<DocmostClientContext>>(Base
|
||||
"listPages: tree mode requires a spaceId (a page tree is scoped to one space). Pass spaceId, or omit tree to get the recent-pages list.",
|
||||
);
|
||||
}
|
||||
const { pages } = await this.enumerateSpacePages(spaceId);
|
||||
return buildPageTree(pages);
|
||||
// #486: propagate `truncated` (same pattern as check_new_comments). The old
|
||||
// code dropped it, so a caller handed an INCOMPLETE tree (the stdio-fallback
|
||||
// BFS hit its node cap) had no way to know pages were missing. Return the
|
||||
// tree alongside the flag; the primary /pages/tree path is uncapped so this
|
||||
// is false there.
|
||||
const { pages, truncated } = await this.enumerateSpacePages(spaceId);
|
||||
return { tree: buildPageTree(pages), truncated };
|
||||
}
|
||||
|
||||
const clampedLimit = Math.max(1, Math.min(100, limit));
|
||||
|
||||
@@ -4,7 +4,6 @@ import { readFileSync } from "fs";
|
||||
import { fileURLToPath } from "url";
|
||||
import { dirname, join } from "path";
|
||||
import { DocmostClient, DocmostMcpConfig } from "./client.js";
|
||||
import { parseNodeArg } from "@docmost/prosemirror-markdown";
|
||||
import { searchShapes } from "./lib/drawio-shapes.js";
|
||||
import { getGuideSection } from "./lib/drawio-guide.js";
|
||||
import { SHARED_TOOL_SPECS, SharedToolSpec } from "./tool-specs.js";
|
||||
@@ -31,6 +30,13 @@ export { destroyAllSessions } from "./lib/collab-session.js";
|
||||
// internals directly; it goes through loadDocmostMcp()).
|
||||
export { SHARED_TOOL_SPECS } from "./tool-specs.js";
|
||||
export type { SharedToolSpec } from "./tool-specs.js";
|
||||
// #489 — write-class registry consumed by the in-app external-MCP retry gate.
|
||||
export {
|
||||
SHARED_TOOL_WRITE_CLASS,
|
||||
isRetryableWriteClass,
|
||||
assertEverySpecDeclaresWriteClass,
|
||||
} from "./tool-specs.js";
|
||||
export type { ToolWriteClass } from "./tool-specs.js";
|
||||
|
||||
// Re-export the build-time REGISTRY_STAMP (issue #447): a deterministic hash of
|
||||
// the tool-specs registry content, generated into src/registry-stamp.generated.ts
|
||||
|
||||
@@ -37,6 +37,25 @@ const CONNECT_TIMEOUT_MS = 25000;
|
||||
/** Time we wait for the server to acknowledge our write before giving up. */
|
||||
const PERSIST_TIMEOUT_MS = 20000;
|
||||
|
||||
/**
|
||||
* Marker property set on the Error thrown when the Hocuspocus handshake REJECTS
|
||||
* our collab token (onAuthenticationFailed). The client wraps content writes so
|
||||
* that on this specific failure it invalidates its cached collab token and
|
||||
* retries once with a fresh one — symmetric to the HTTP-401 reauth path (#486).
|
||||
* A plain message-match would be brittle; a tagged property is unambiguous and
|
||||
* survives teardown (which rejects pending ops with this SAME error object).
|
||||
*/
|
||||
const COLLAB_AUTH_FAILED_MARKER = "collabAuthFailed";
|
||||
|
||||
/** True when `e` is the tagged collab-WS auth-failure error (see marker above). */
|
||||
export function isCollabAuthFailedError(e: unknown): boolean {
|
||||
return !!(
|
||||
e &&
|
||||
typeof e === "object" &&
|
||||
(e as Record<string, unknown>)[COLLAB_AUTH_FAILED_MARKER] === true
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Tunables, read fresh from the environment on every acquire so tests (and a
|
||||
* live rollback) can change them without reloading the module. Mirrors how
|
||||
@@ -302,10 +321,13 @@ export class CollabSession {
|
||||
this.openResolve?.();
|
||||
},
|
||||
onAuthenticationFailed: () => {
|
||||
this.teardown(
|
||||
new Error("Authentication failed for collaboration connection"),
|
||||
true,
|
||||
);
|
||||
// Tag the error so the client can tell a REJECTED collab token apart
|
||||
// from a generic disconnect and invalidate + refresh it (#486).
|
||||
const err = new Error(
|
||||
"Authentication failed for collaboration connection",
|
||||
) as Error & { [COLLAB_AUTH_FAILED_MARKER]?: boolean };
|
||||
err[COLLAB_AUTH_FAILED_MARKER] = true;
|
||||
this.teardown(err, true);
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
@@ -254,6 +254,12 @@ export async function mutatePageContent(
|
||||
collabToken: string,
|
||||
baseUrl: string,
|
||||
transform: (liveDoc: any) => any | null,
|
||||
// #487: optional abort signal carrying the turn's Stop + the in-app tool
|
||||
// per-call cap. Checked as the PRE-COMMIT safe-point below (after the session
|
||||
// is acquired, immediately before the atomic read->write), so a Stop that
|
||||
// arrives during the connect/lock window stops THIS write from landing. See the
|
||||
// limitation note at the check.
|
||||
signal?: AbortSignal,
|
||||
): Promise<MutationResult> {
|
||||
return withPageLock(pageId, async () => {
|
||||
if (process.env.DEBUG) {
|
||||
@@ -266,6 +272,13 @@ export async function mutatePageContent(
|
||||
|
||||
const session = await acquireCollabSession(pageId, collabToken, baseUrl);
|
||||
try {
|
||||
// #487 PRE-COMMIT safe-point: if the turn was Stopped (or the in-app tool
|
||||
// per-call cap fired) after we acquired the collab session but before the
|
||||
// atomic write, throw NOW so this commit never runs. KNOWN LIMITATION
|
||||
// (#487): this only stops THIS commit — a write tool that already committed
|
||||
// an EARLIER call this turn leaves that op applied. Cancel guarantees "no
|
||||
// NEW commit starts", NOT "the write didn't land".
|
||||
signal?.throwIfAborted();
|
||||
return await session.mutate(transform);
|
||||
} catch (e) {
|
||||
// Drop the session on any failure so the next call reconnects fresh (this
|
||||
@@ -291,6 +304,8 @@ export async function replacePageContent(
|
||||
prosemirrorDoc: any,
|
||||
collabToken: string,
|
||||
baseUrl: string,
|
||||
// #487: threaded straight to mutatePageContent's pre-commit safe-point.
|
||||
signal?: AbortSignal,
|
||||
): Promise<MutationResult> {
|
||||
// Fail fast on a bad document instead of deferring the failure into the
|
||||
// collaboration write (where TiptapTransformer.toYdoc(undefined) used to
|
||||
@@ -307,6 +322,7 @@ export async function replacePageContent(
|
||||
collabToken,
|
||||
baseUrl,
|
||||
() => prosemirrorDoc,
|
||||
signal,
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
// exactly mxGraph's convention for a child of a container, so they map across
|
||||
// directly. Container sizes are computed by ELK; leaf sizes are preserved.
|
||||
|
||||
import ELK from "elkjs/lib/elk.bundled.js";
|
||||
import { Worker } from "node:worker_threads";
|
||||
import { JSDOM } from "jsdom";
|
||||
import { normalizeInput, parseCells, type DrawioCell } from "./drawio-xml.js";
|
||||
|
||||
@@ -18,22 +18,33 @@ import { normalizeInput, parseCells, type DrawioCell } from "./drawio-xml.js";
|
||||
const DEFAULT_W = 140;
|
||||
const DEFAULT_H = 60;
|
||||
|
||||
// DoS bounds for the in-process ELK layout. The mxGraph XML is LLM-supplied
|
||||
// (layout:"elk" in drawioCreate/drawioUpdate) and elkjs runs synchronously on
|
||||
// the MCP server's event loop, so an unbounded graph would block it for
|
||||
// seconds-to-minutes. A ~1MB XML (well under the stage-1 16MB cap) can carry
|
||||
// thousands of nodes. We cap the graph size and race the layout against a
|
||||
// wall-clock timeout; on either bound we fall back to the ORIGINAL model, the
|
||||
// same best-effort contract the catch already honours.
|
||||
// DoS bounds for the ELK layout. The mxGraph XML is LLM-supplied (layout:"elk"
|
||||
// in drawioCreate/drawioUpdate). elkjs' layout() returns a Promise but runs the
|
||||
// crossing-minimisation SYNCHRONOUSLY — it blocks whatever thread it runs on for
|
||||
// the whole pass. A ~1MB XML (well under the stage-1 16MB cap) can carry
|
||||
// thousands of nodes. We (a) cap the graph size before ever calling ELK and
|
||||
// (b) run the layout in a WORKER THREAD so the main event loop stays free, with
|
||||
// the wall-clock timeout enforced by terminating that worker. On either bound we
|
||||
// fall back to the ORIGINAL model, the same best-effort contract the catch honours.
|
||||
// - 500 nodes lays out in well under a second; beyond that ELK cost climbs
|
||||
// steeply, so refuse and leave the (already-valid) model untouched.
|
||||
// - Edges dominate the layered-crossing cost, so allow a bit more headroom
|
||||
// (1000) than nodes but still bound them.
|
||||
// - 5s is generous for any graph within the caps yet short enough that a
|
||||
// pathological input can never wedge the server.
|
||||
// - The timeout is a HARD kill of the worker thread — the only way to interrupt
|
||||
// synchronous JS. The in-process setTimeout race we used before was an
|
||||
// illusion: the timer could never fire while the SAME thread was blocked
|
||||
// inside elkjs, so it "protected" nothing. Now the timer runs on the main
|
||||
// thread while ELK runs on the worker, so it can actually fire and terminate.
|
||||
const ELK_MAX_NODES = 500;
|
||||
const ELK_MAX_EDGES = 1000;
|
||||
const ELK_TIMEOUT_MS = 5000;
|
||||
// Wall-clock ceiling for a single layout pass. Overridable for tests (a tiny
|
||||
// value forces the terminate-on-timeout path deterministically); a non-positive
|
||||
// or unparseable override falls back to the default.
|
||||
const ELK_TIMEOUT_DEFAULT_MS = 5000;
|
||||
function resolveElkTimeoutMs(): number {
|
||||
const raw = Number(process.env.DRAWIO_ELK_TIMEOUT_MS);
|
||||
return Number.isFinite(raw) && raw > 0 ? Math.floor(raw) : ELK_TIMEOUT_DEFAULT_MS;
|
||||
}
|
||||
|
||||
// Spacing is set >=150px on purpose so an ELK layout never trips the linter's
|
||||
// "gap between adjacent shapes < 150px" quality warning (acceptance #3).
|
||||
@@ -78,13 +89,57 @@ interface ElkGraph extends ElkNode {
|
||||
edges?: ElkEdge[];
|
||||
}
|
||||
|
||||
/**
|
||||
* Run one ELK layered layout on a worker thread and resolve with the laid-out
|
||||
* graph. The timeout is enforced by `worker.terminate()` — a HARD kill, which is
|
||||
* the only way to interrupt elkjs' synchronous crossing-minimisation once it has
|
||||
* started. Rejects on timeout, worker error, or an early exit; the caller treats
|
||||
* any rejection as "keep the original model" (best-effort layout).
|
||||
*/
|
||||
function layoutInWorker(graph: ElkGraph, timeoutMs: number): Promise<ElkGraph> {
|
||||
return new Promise((resolve, reject) => {
|
||||
const worker = new Worker(
|
||||
new URL("./drawio-layout.worker.js", import.meta.url),
|
||||
{ workerData: { graph } },
|
||||
);
|
||||
let settled = false;
|
||||
const finish = (fn: () => void) => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
clearTimeout(timer);
|
||||
// Always tear the worker down: on the happy path so it does not linger,
|
||||
// on timeout so the blocked synchronous ELK run is actually interrupted.
|
||||
void worker.terminate();
|
||||
fn();
|
||||
};
|
||||
const timer = setTimeout(
|
||||
() => finish(() => reject(new Error("ELK layout timed out"))),
|
||||
timeoutMs,
|
||||
);
|
||||
worker.once("message", (msg: { ok?: boolean; laid?: ElkGraph; error?: string }) => {
|
||||
finish(() =>
|
||||
msg?.ok
|
||||
? resolve(msg.laid as ElkGraph)
|
||||
: reject(new Error(msg?.error ?? "ELK layout failed")),
|
||||
);
|
||||
});
|
||||
worker.once("error", (err) => finish(() => reject(err)));
|
||||
worker.once("exit", (code) => {
|
||||
// A clean exit after we already settled is normal (terminate()); only an
|
||||
// unexpected early exit while still pending is a failure.
|
||||
if (settled) return;
|
||||
finish(() => reject(new Error(`ELK worker exited early (code ${code})`)));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Apply an ELK layered layout to a drawio input and return a full mxGraphModel
|
||||
* string with rewritten geometry. Accepts the same three input forms as
|
||||
* drawioCreate (a bare model, an <mxfile>, or a <mxCell> list). Async because
|
||||
* elkjs' layout() is promise-based. On any layout failure the ORIGINAL
|
||||
* (normalized) model is returned unchanged — layout is best-effort polish, never
|
||||
* a reason to fail the write.
|
||||
* the layout runs on a worker thread. On any layout failure (including a
|
||||
* terminate-on-timeout) the ORIGINAL (normalized) model is returned unchanged —
|
||||
* layout is best-effort polish, never a reason to fail the write.
|
||||
*/
|
||||
export async function applyElkLayout(inputXml: string): Promise<string> {
|
||||
const modelXml = normalizeInput(inputXml);
|
||||
@@ -150,26 +205,14 @@ export async function applyElkLayout(inputXml: string): Promise<string> {
|
||||
};
|
||||
|
||||
let laid: ElkGraph;
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
try {
|
||||
// elkjs ships a CJS default export whose interop shape varies across
|
||||
// module systems; resolve the real constructor at runtime, then cast (the
|
||||
// runtime call is verified — see the layout unit test).
|
||||
const Ctor: any = (ELK as any).default ?? ELK;
|
||||
const elk = new Ctor();
|
||||
// Race the layout against a wall-clock timeout so a graph that is under the
|
||||
// node/edge caps but still pathologically slow can never wedge the server.
|
||||
const timeout = new Promise<never>((_, reject) => {
|
||||
timer = setTimeout(
|
||||
() => reject(new Error("ELK layout timed out")),
|
||||
ELK_TIMEOUT_MS,
|
||||
);
|
||||
});
|
||||
laid = (await Promise.race([elk.layout(graph as any), timeout])) as ElkGraph;
|
||||
// Run the (synchronous-under-the-hood) ELK pass on a worker thread so the
|
||||
// main event loop is never blocked, and enforce the wall-clock ceiling by
|
||||
// terminating that worker on timeout. A graph under the node/edge caps but
|
||||
// still pathologically slow is hard-killed instead of wedging anything.
|
||||
laid = await layoutInWorker(graph, resolveElkTimeoutMs());
|
||||
} catch {
|
||||
return modelXml; // best-effort: keep the model as-is on timeout or ELK failure
|
||||
} finally {
|
||||
if (timer) clearTimeout(timer);
|
||||
}
|
||||
|
||||
// Collect computed geometry per node id (coords are parent-relative already).
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
// Worker-thread entry for the ELK layered layout (issue #486, commit 1).
|
||||
//
|
||||
// elkjs' layout() returns a Promise but runs the actual crossing-minimisation
|
||||
// SYNCHRONOUSLY — it blocks whatever thread it runs on for the whole pass. On
|
||||
// the in-app MCP host that thread used to be the main NestJS event loop, so a
|
||||
// pathological graph at the node/edge cap could wedge ALL HTTP/SSE/loopback
|
||||
// traffic while it churned. Running it HERE, on a dedicated worker thread, keeps
|
||||
// the main loop free; the parent enforces the wall-clock timeout by calling
|
||||
// `worker.terminate()` — the only way to interrupt synchronous JS — since the
|
||||
// in-process `setTimeout` race the parent used before could never fire while the
|
||||
// same thread was blocked inside elkjs.
|
||||
import { parentPort, workerData } from "node:worker_threads";
|
||||
import ELK from "elkjs/lib/elk.bundled.js";
|
||||
|
||||
interface WorkerInput {
|
||||
graph: unknown;
|
||||
}
|
||||
|
||||
const { graph } = (workerData ?? {}) as WorkerInput;
|
||||
|
||||
// elkjs ships a CJS default export whose interop shape varies across module
|
||||
// systems; resolve the real constructor at runtime (same as the parent did).
|
||||
const Ctor: any = (ELK as any).default ?? ELK;
|
||||
const elk = new Ctor();
|
||||
|
||||
elk
|
||||
.layout(graph as any)
|
||||
.then((laid: unknown) => {
|
||||
parentPort?.postMessage({ ok: true, laid });
|
||||
})
|
||||
.catch((err: unknown) => {
|
||||
parentPort?.postMessage({
|
||||
ok: false,
|
||||
error: err instanceof Error ? err.message : String(err),
|
||||
});
|
||||
});
|
||||
@@ -127,6 +127,18 @@ export interface SharedToolSpec {
|
||||
mcpName: string;
|
||||
/** camelCase key in the ai-SDK tools object (the in-app layer). */
|
||||
inAppKey: string;
|
||||
/**
|
||||
* Write-class of the tool (#489), declared on EVERY spec (a registration-time
|
||||
* assert enforces completeness; `satisfies Record<string, SharedToolSpec>`
|
||||
* makes it a compile error to omit). 'readOnly' = a pure read that mutates
|
||||
* NOTHING durable, so it is safe to auto-retry once after a transport break.
|
||||
* 'write' = anything that mutates a page/comment/share/diagram/etc — a
|
||||
* transport error is INDETERMINATE (the server may have applied it before the
|
||||
* connection reset), so it is NEVER blind-retried (a retry would double-apply,
|
||||
* the #435 incident class). Consumed by the external-MCP retry path
|
||||
* (mcp-clients.service.ts) to gate its single auto-retry.
|
||||
*/
|
||||
writeClass: 'readOnly' | 'write';
|
||||
/** Single canonical model-facing description used by both layers. */
|
||||
description: string;
|
||||
/**
|
||||
@@ -240,6 +252,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
getWorkspace: {
|
||||
mcpName: 'getWorkspace',
|
||||
inAppKey: 'getWorkspace',
|
||||
writeClass: 'readOnly',
|
||||
description: 'Fetch metadata about the current workspace (name, settings).',
|
||||
tier: 'core',
|
||||
catalogLine: 'getWorkspace — fetch current workspace metadata (name, settings).',
|
||||
@@ -249,6 +262,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
listSpaces: {
|
||||
mcpName: 'listSpaces',
|
||||
inAppKey: 'listSpaces',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
'List the spaces the current user can access. Returns the array of ' +
|
||||
'spaces (id, name, slug, ...).',
|
||||
@@ -260,6 +274,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
listShares: {
|
||||
mcpName: 'listShares',
|
||||
inAppKey: 'listShares',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
'List all public shares in the workspace with page titles and public URLs.',
|
||||
tier: 'deferred',
|
||||
@@ -272,6 +287,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
getPageJson: {
|
||||
mcpName: 'getPageJson',
|
||||
inAppKey: 'getPageJson',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
'Get page details with the raw ProseMirror JSON content (lossless: ' +
|
||||
'includes block ids, callouts, tables, link/image attributes) plus the ' +
|
||||
@@ -289,6 +305,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
getOutline: {
|
||||
mcpName: 'getOutline',
|
||||
inAppKey: 'getOutline',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
"Return a COMPACT outline of a page's top-level blocks ({index, type, " +
|
||||
'id, level, firstText}; tables add rows/cols/header; lists add item ' +
|
||||
@@ -309,6 +326,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
getNode: {
|
||||
mcpName: 'getNode',
|
||||
inAppKey: 'getNode',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
"Fetch a single block for editing. `nodeId` is a block id from the page " +
|
||||
'outline or page-JSON view (works for headings/paragraphs/callouts/images), OR ' +
|
||||
@@ -350,6 +368,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
searchInPage: {
|
||||
mcpName: 'searchInPage',
|
||||
inAppKey: 'searchInPage',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
'Find every occurrence of a string (or regex) INSIDE one page and get ' +
|
||||
'WHERE each is — instead of pulling blocks one-by-one with getNode. ' +
|
||||
@@ -413,6 +432,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
deleteNode: {
|
||||
mcpName: 'deleteNode',
|
||||
inAppKey: 'deleteNode',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Remove a single block by its attrs.id (from the page outline or ' +
|
||||
'page-JSON view) WITHOUT resending the whole document.',
|
||||
@@ -438,6 +458,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
patchNode: {
|
||||
mcpName: 'patchNode',
|
||||
inAppKey: 'patchNode',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Replace a single content block identified by its attrs.id, WITHOUT ' +
|
||||
'resending the whole document; the replacement keeps the same block id. ' +
|
||||
@@ -505,6 +526,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
insertNode: {
|
||||
mcpName: 'insertNode',
|
||||
inAppKey: 'insertNode',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Insert content before/after another block (by attrs.id or anchor text) ' +
|
||||
'or append it at the end (top level). For before/after you MUST provide ' +
|
||||
@@ -597,6 +619,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
sharePage: {
|
||||
mcpName: 'sharePage',
|
||||
inAppKey: 'sharePage',
|
||||
writeClass: 'write',
|
||||
// CANONICAL: merges the MCP copy's URL-format + idempotency detail with the
|
||||
// in-app copy's reversibility note; keeps the security framing both had.
|
||||
description:
|
||||
@@ -626,6 +649,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
unsharePage: {
|
||||
mcpName: 'unsharePage',
|
||||
inAppKey: 'unsharePage',
|
||||
writeClass: 'write',
|
||||
description: 'Remove the public share of a page (revokes the public URL).',
|
||||
tier: 'deferred',
|
||||
catalogLine: "unsharePage — revoke a page's public share (removes the public URL).",
|
||||
@@ -640,6 +664,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
diffPageVersions: {
|
||||
mcpName: 'diffPageVersions',
|
||||
inAppKey: 'diffPageVersions',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
'Diff two versions of a page and return a Docmost-equivalent change set ' +
|
||||
'(inserted/deleted text, integrity counts for images/links/tables/' +
|
||||
@@ -672,6 +697,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
listPageHistory: {
|
||||
mcpName: 'listPageHistory',
|
||||
inAppKey: 'listPageHistory',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
"List a page's saved versions (Docmost auto-snapshots on every save), " +
|
||||
'newest first, cursor-paginated. Returns { items, nextCursor }; each ' +
|
||||
@@ -693,6 +719,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
restorePageVersion: {
|
||||
mcpName: 'restorePageVersion',
|
||||
inAppKey: 'restorePageVersion',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Restore a page to a saved version: writes that version\'s content back ' +
|
||||
'as the page\'s current content (Docmost has no restore endpoint, so ' +
|
||||
@@ -713,6 +740,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
importPageMarkdown: {
|
||||
mcpName: 'importPageMarkdown',
|
||||
inAppKey: 'importPageMarkdown',
|
||||
writeClass: 'write',
|
||||
// IN-APP ONLY (issue #411): the external /mcp surface no longer exposes
|
||||
// importPageMarkdown — the registry loop in index.ts skips inAppOnly specs,
|
||||
// so this stays available to the in-app agent (round-tripping an EXPORTED
|
||||
@@ -742,6 +770,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
copyPageContent: {
|
||||
mcpName: 'copyPageContent',
|
||||
inAppKey: 'copyPageContent',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
"Replace targetPageId's content with a copy of sourcePageId's content, " +
|
||||
'entirely server-side — the document is NOT sent through the model. The ' +
|
||||
@@ -770,6 +799,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
editPageText: {
|
||||
mcpName: 'editPageText',
|
||||
inAppKey: 'editPageText',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
"Surgical find/replace inside a page's text, preserving all block " +
|
||||
'ids and marks. A find MAY cross bold/italic/link boundaries; the ' +
|
||||
@@ -819,6 +849,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
stashPage: {
|
||||
mcpName: 'stashPage',
|
||||
inAppKey: 'stashPage',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
'Serialize a whole page (the full ProseMirror JSON, as getPageJson ' +
|
||||
'returns) into an ephemeral in-memory blob and return ONLY a short ' +
|
||||
@@ -880,6 +911,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
getPage: {
|
||||
mcpName: 'getPage',
|
||||
inAppKey: 'getPage',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
'Fetch a single page as Markdown by its id. Returns the page title and ' +
|
||||
'its Markdown content. The converter is canonical (round-trips text and ' +
|
||||
@@ -919,13 +951,15 @@ export const SHARED_TOOL_SPECS = {
|
||||
listPages: {
|
||||
mcpName: 'listPages',
|
||||
inAppKey: 'listPages',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
'List the most recent pages (ordered by updatedAt, descending), ' +
|
||||
'optionally scoped to a single space. Returns a bounded list (default ' +
|
||||
'50, max 100) — use search for lookups in large spaces. tree:true (with ' +
|
||||
"spaceId) returns the space's full page hierarchy as a nested tree, but " +
|
||||
'is DEPRECATED — use getTree instead (leaner nodes, plus rootPageId / ' +
|
||||
'maxDepth).',
|
||||
"spaceId) returns { tree, truncated } — the space's full page hierarchy " +
|
||||
'as a nested tree, plus a `truncated` flag that is true when the tree was ' +
|
||||
'capped and is INCOMPLETE — but is DEPRECATED, use getTree instead ' +
|
||||
'(leaner nodes, plus rootPageId / maxDepth).',
|
||||
tier: 'core',
|
||||
catalogLine:
|
||||
"listPages — list recent pages (tree:true is deprecated; use getTree for the hierarchy).",
|
||||
@@ -964,6 +998,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
getTree: {
|
||||
mcpName: 'getTree',
|
||||
inAppKey: 'getTree',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
"Get a space's page hierarchy (or one subtree) as a nested tree in a " +
|
||||
'SINGLE request — completely and without loss. Each node is ' +
|
||||
@@ -1008,6 +1043,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
getPageContext: {
|
||||
mcpName: 'getPageContext',
|
||||
inAppKey: 'getPageContext',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
'Given a pageId, get its LOCATION and immediate surroundings (metadata ' +
|
||||
'only, no page content) in one call — answers "where am I / what is ' +
|
||||
@@ -1037,6 +1073,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
createPage: {
|
||||
mcpName: 'createPage',
|
||||
inAppKey: 'createPage',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Create a new page with a Markdown body in a space, optionally under a ' +
|
||||
'parent page (omit parentPageId to create at the space root). Returns ' +
|
||||
@@ -1087,6 +1124,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
movePage: {
|
||||
mcpName: 'movePage',
|
||||
inAppKey: 'movePage',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Move a page under a new parent page, or to the space root when no ' +
|
||||
'parent is given. Reversible: move it back at any time.',
|
||||
@@ -1180,6 +1218,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
renamePage: {
|
||||
mcpName: 'renamePage',
|
||||
inAppKey: 'renamePage',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Rename a page (change its title only; the body is untouched, never ' +
|
||||
'resent). Reversible: rename back at any time.',
|
||||
@@ -1201,6 +1240,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
deletePage: {
|
||||
mcpName: 'deletePage',
|
||||
inAppKey: 'deletePage',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Move a page to the trash — SOFT delete only: the page can be restored ' +
|
||||
'from trash and nothing is ever permanently deleted.',
|
||||
@@ -1234,6 +1274,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
updatePageJson: {
|
||||
mcpName: 'updatePageJson',
|
||||
inAppKey: 'updatePageJson',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
"Replace a page's content with a raw ProseMirror JSON document (lossless " +
|
||||
'write: preserves the block ids, callouts, tables and attributes you pass ' +
|
||||
@@ -1290,6 +1331,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
updatePageMarkdown: {
|
||||
mcpName: 'updatePageMarkdown',
|
||||
inAppKey: 'updatePageMarkdown',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
"Replace a page's body with new Markdown content (and optionally its " +
|
||||
'title). The whole body is re-imported from the markdown (block ids ' +
|
||||
@@ -1324,6 +1366,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
exportPageMarkdown: {
|
||||
mcpName: 'exportPageMarkdown',
|
||||
inAppKey: 'exportPageMarkdown',
|
||||
writeClass: 'readOnly',
|
||||
// CANONICAL: the MCP copy (a strict superset of the terse in-app wording).
|
||||
description:
|
||||
'Export a page to a single self-contained Docmost-flavoured Markdown ' +
|
||||
@@ -1372,6 +1415,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
createComment: {
|
||||
mcpName: 'createComment',
|
||||
inAppKey: 'createComment',
|
||||
writeClass: 'write',
|
||||
// CANONICAL: the in-app copy (the more-maintained one). It keeps the same
|
||||
// rules as the MCP copy — inline-only, top-level requires a `selection`, no
|
||||
// page-level comments, replies inherit the anchor, suggestedText must be
|
||||
@@ -1504,6 +1548,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
listComments: {
|
||||
mcpName: 'listComments',
|
||||
inAppKey: 'listComments',
|
||||
writeClass: 'readOnly',
|
||||
// CANONICAL: the two copies are near-identical; the MCP copy is the
|
||||
// superset (it keeps the "(pagination is handled internally)" note the
|
||||
// in-app copy dropped), so it is used verbatim.
|
||||
@@ -1531,6 +1576,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
resolveComment: {
|
||||
mcpName: 'resolveComment',
|
||||
inAppKey: 'resolveComment',
|
||||
writeClass: 'write',
|
||||
// CANONICAL: the MCP copy's richer wording, minus its reference
|
||||
// to `deleteComment` (a sibling tool that does NOT exist in the in-app
|
||||
// layer) — rephrased transport-neutrally per the registry convention.
|
||||
@@ -1571,6 +1617,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
checkNewComments: {
|
||||
mcpName: 'checkNewComments',
|
||||
inAppKey: 'checkNewComments',
|
||||
writeClass: 'readOnly',
|
||||
// CANONICAL: the MCP copy (the more detailed of the two). The MCP layer's
|
||||
// execute-side guard that rejects an unparseable `since` timestamp stays in
|
||||
// its execute body (per-layer logic), not in the shared schema.
|
||||
@@ -1650,6 +1697,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
tableInsertRow: {
|
||||
mcpName: 'tableInsertRow',
|
||||
inAppKey: 'tableInsertRow',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Insert a row of plain-text cells into a table. `table` is `#<index>` ' +
|
||||
'from the page outline, or a block id inside it. `cells` is the text per ' +
|
||||
@@ -1683,6 +1731,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
tableDeleteRow: {
|
||||
mcpName: 'tableDeleteRow',
|
||||
inAppKey: 'tableDeleteRow',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Delete the row at 0-based `index` from a table (`table` is `#<index>` ' +
|
||||
'from the page outline, or a block id inside it). Refuses to delete the ' +
|
||||
@@ -1706,6 +1755,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
tableUpdateCell: {
|
||||
mcpName: 'tableUpdateCell',
|
||||
inAppKey: 'tableUpdateCell',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Set the plain-text content of cell [row, col] (0-based) in a table ' +
|
||||
'(`table` is `#<index>` from the page outline, or a block id inside it). ' +
|
||||
@@ -1746,6 +1796,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
insertFootnote: {
|
||||
mcpName: 'insertFootnote',
|
||||
inAppKey: 'insertFootnote',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Insert an AUTHOR-INLINE footnote: you specify only WHERE (anchorText) ' +
|
||||
'and WHAT (text). The footnote marker is placed right after anchorText in ' +
|
||||
@@ -1783,6 +1834,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
insertImage: {
|
||||
mcpName: 'insertImage',
|
||||
inAppKey: 'insertImage',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Download an image from a web (http/https) URL and insert it into ' +
|
||||
'a page in one step. By default ' +
|
||||
@@ -1827,6 +1879,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
replaceImage: {
|
||||
mcpName: 'replaceImage',
|
||||
inAppKey: 'replaceImage',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Replace an existing image on a page with a new image fetched from a web ' +
|
||||
'(http/https) URL: uploads the new file as a NEW ' +
|
||||
@@ -1865,6 +1918,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
drawioGet: {
|
||||
mcpName: 'drawioGet',
|
||||
inAppKey: 'drawioGet',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
'Read a draw.io diagram on a page as mxGraph XML (default) or as its raw ' +
|
||||
'`.drawio.svg`. `node` is the drawio node\'s attrs.id (from getOutline / ' +
|
||||
@@ -1898,6 +1952,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
drawioCreate: {
|
||||
mcpName: 'drawioCreate',
|
||||
inAppKey: 'drawioCreate',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Create a draw.io diagram from mxGraph XML and insert it as a diagram ' +
|
||||
'block. `xml` is a bare `<mxGraphModel>` OR a list of `<mxCell>` elements ' +
|
||||
@@ -1968,6 +2023,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
drawioUpdate: {
|
||||
mcpName: 'drawioUpdate',
|
||||
inAppKey: 'drawioUpdate',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Replace a draw.io diagram\'s content with new mxGraph XML (same lint ' +
|
||||
'pipeline as drawioCreate). `baseHash` is MANDATORY: pass the hash from ' +
|
||||
@@ -2019,6 +2075,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
drawioEditCells: {
|
||||
mcpName: 'drawioEditCells',
|
||||
inAppKey: 'drawioEditCells',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Make TARGETED, id-based edits to an existing draw.io diagram instead of ' +
|
||||
'resending the whole XML (a full-XML diff is fragile — draw.io reorders ' +
|
||||
@@ -2077,6 +2134,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
drawioFromGraph: {
|
||||
mcpName: 'drawioFromGraph',
|
||||
inAppKey: 'drawioFromGraph',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Build a draw.io diagram from a SEMANTIC graph — you describe nodes, groups ' +
|
||||
'and edges by MEANING and the server picks every coordinate, color and icon ' +
|
||||
@@ -2201,6 +2259,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
drawioFromMermaid: {
|
||||
mcpName: 'drawioFromMermaid',
|
||||
inAppKey: 'drawioFromMermaid',
|
||||
writeClass: 'write',
|
||||
description:
|
||||
'Convert Mermaid `flowchart` text into an EDITABLE draw.io diagram (LLMs ' +
|
||||
'write Mermaid reliably). Best for STANDARD flowcharts/decision trees: ' +
|
||||
@@ -2247,6 +2306,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
drawioShapes: {
|
||||
mcpName: 'drawioShapes',
|
||||
inAppKey: 'drawioShapes',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
'Look up VERIFIED draw.io stencil style-strings so you never guess a ' +
|
||||
'`shape=mxgraph.*` name (a wrong name renders as an EMPTY BOX). Searches a ' +
|
||||
@@ -2293,6 +2353,7 @@ export const SHARED_TOOL_SPECS = {
|
||||
drawioGuide: {
|
||||
mcpName: 'drawioGuide',
|
||||
inAppKey: 'drawioGuide',
|
||||
writeClass: 'readOnly',
|
||||
description:
|
||||
'Progressive-disclosure draw.io authoring reference. Call with a `section` ' +
|
||||
'to pull one focused, <=4KB chapter instead of bloating context: ' +
|
||||
@@ -2322,3 +2383,53 @@ export const SHARED_TOOL_SPECS = {
|
||||
inlineBothHosts: true,
|
||||
},
|
||||
} satisfies Record<string, SharedToolSpec>;
|
||||
|
||||
// --- write-class registry (#489) ------------------------------------------
|
||||
|
||||
/** A tool's retry-safety class. 'readOnly' may be auto-retried once after a
|
||||
* transport break; 'write' is indeterminate and must never be blind-retried. */
|
||||
export type ToolWriteClass = 'readOnly' | 'write';
|
||||
|
||||
/**
|
||||
* Name → write-class map for the shared registry, keyed by mcpName (=== inAppKey).
|
||||
* The external-MCP retry path (mcp-clients.service.ts) looks a tool up here by its
|
||||
* RAW (un-namespaced) name to decide whether a transport failure may be retried.
|
||||
* A tool NOT in this map (a third-party external MCP tool) is treated as 'write'
|
||||
* by the consumer — the safe default (never blind-retry an unknown tool).
|
||||
*/
|
||||
export const SHARED_TOOL_WRITE_CLASS: Record<string, ToolWriteClass> =
|
||||
Object.fromEntries(
|
||||
Object.values(SHARED_TOOL_SPECS).map((spec) => [spec.mcpName, spec.writeClass]),
|
||||
);
|
||||
|
||||
/** Whether a write-class permits a single automatic retry after a transport
|
||||
* break. Only a pure read is retry-safe; everything mutating is indeterminate. */
|
||||
export function isRetryableWriteClass(
|
||||
writeClass: ToolWriteClass | undefined,
|
||||
): boolean {
|
||||
return writeClass === 'readOnly';
|
||||
}
|
||||
|
||||
/**
|
||||
* Registration-time assert (#489): EVERY spec must declare a valid write-class.
|
||||
* `satisfies Record<string, SharedToolSpec>` already makes an omission a compile
|
||||
* error, but this guards a raw/cast construction path and documents the invariant
|
||||
* at the point of use. Runs once on import — both hosts import this module, so
|
||||
* both get the check. Throws (fails startup) rather than silently mis-gating a
|
||||
* retry in production.
|
||||
*/
|
||||
export function assertEverySpecDeclaresWriteClass(): void {
|
||||
for (const [key, spec] of Object.entries(SHARED_TOOL_SPECS)) {
|
||||
const wc = (spec as SharedToolSpec).writeClass;
|
||||
if (wc !== 'readOnly' && wc !== 'write') {
|
||||
throw new Error(
|
||||
`tool-specs: spec "${key}" must declare writeClass ('readOnly' | 'write'), got ${JSON.stringify(
|
||||
wc,
|
||||
)}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Enforce at module load (registration time) on both hosts.
|
||||
assertEverySpecDeclaresWriteClass();
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
// Unit tests for the collab-token reset on a Hocuspocus WS auth failure (#486).
|
||||
//
|
||||
// Before this fix the cached collab token (#435) was dropped ONLY on an HTTP
|
||||
// 401/403 (the REST interceptor + login()); a rejected collab-WEBSOCKET handshake
|
||||
// left the stale token in the cache, so every subsequent mutation re-presented
|
||||
// the SAME bad token for up to the collab-token TTL (minutes) with no self-heal.
|
||||
//
|
||||
// The fix wraps collab writes in `writeWithCollabAuthRetry`: when the write
|
||||
// rejects with the tagged collab-auth error (collab-session.ts's
|
||||
// onAuthenticationFailed), it invalidates the cached token and retries the write
|
||||
// ONCE with a force-refreshed token — symmetric to the HTTP-401 path.
|
||||
//
|
||||
// writeWithCollabAuthRetry / getCollabTokenWithReauth are protected in TS but
|
||||
// plain methods on the compiled build, so the tests call them directly (same
|
||||
// convention as collab-token-cache.test.mjs).
|
||||
import { test, afterEach } from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
import { DocmostClient } from "../../build/client.js";
|
||||
|
||||
const ENV_KEY = "MCP_COLLAB_TOKEN_TTL_MS";
|
||||
afterEach(() => {
|
||||
delete process.env[ENV_KEY];
|
||||
});
|
||||
|
||||
// A counting provider that returns a distinct token each call so a cached
|
||||
// (reused) token is visibly the SAME string while a fresh mint is different.
|
||||
function countingProvider() {
|
||||
let n = 0;
|
||||
const fn = async () => {
|
||||
n++;
|
||||
return `provider-token-${n}`;
|
||||
};
|
||||
return {
|
||||
fn,
|
||||
get calls() {
|
||||
return n;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
// The tagged error collab-session.ts throws on a rejected WS handshake.
|
||||
function collabAuthError() {
|
||||
const err = new Error("Authentication failed for collaboration connection");
|
||||
err.collabAuthFailed = true;
|
||||
return err;
|
||||
}
|
||||
|
||||
test("a WS auth failure clears the cached token and retries the write with a FRESH one (#486)", async () => {
|
||||
process.env[ENV_KEY] = "300000"; // 5 min: the cache is warm across the burst.
|
||||
const p = countingProvider();
|
||||
const client = new DocmostClient({
|
||||
apiUrl: "http://127.0.0.1:1/api",
|
||||
getToken: async () => "access",
|
||||
getCollabToken: p.fn,
|
||||
});
|
||||
|
||||
// Warm the cache the way a real write would (mints provider-token-1).
|
||||
const initial = await client.getCollabTokenWithReauth();
|
||||
assert.equal(initial, "provider-token-1");
|
||||
assert.equal(p.calls, 1);
|
||||
|
||||
const tokensSeen = [];
|
||||
const write = async (token) => {
|
||||
tokensSeen.push(token);
|
||||
// The FIRST attempt (with the stale cached token) fails the WS handshake;
|
||||
// the retry (with a fresh token) succeeds.
|
||||
if (tokensSeen.length === 1) throw collabAuthError();
|
||||
return `written-with:${token}`;
|
||||
};
|
||||
|
||||
const result = await client.writeWithCollabAuthRetry(initial, write);
|
||||
|
||||
assert.equal(tokensSeen.length, 2, "write attempted exactly twice (one retry)");
|
||||
assert.equal(tokensSeen[0], "provider-token-1", "first attempt used the stale token");
|
||||
assert.equal(
|
||||
tokensSeen[1],
|
||||
"provider-token-2",
|
||||
"retry used a FRESH force-refreshed token, not the stale cached one",
|
||||
);
|
||||
assert.equal(result, "written-with:provider-token-2", "the retry's result wins");
|
||||
assert.equal(p.calls, 2, "exactly one extra mint for the retry — no loop");
|
||||
|
||||
// The cache now holds the fresh token, so a subsequent op reuses it (proving
|
||||
// the stale token was evicted and the fresh one cached, not re-minted).
|
||||
const next = await client.getCollabTokenWithReauth();
|
||||
assert.equal(next, "provider-token-2", "the fresh token replaced the stale cache");
|
||||
assert.equal(p.calls, 2, "served from cache — provider not re-invoked");
|
||||
});
|
||||
|
||||
test("a successful write is NOT retried and mints nothing extra", async () => {
|
||||
process.env[ENV_KEY] = "300000";
|
||||
const p = countingProvider();
|
||||
const client = new DocmostClient({
|
||||
apiUrl: "http://127.0.0.1:1/api",
|
||||
getToken: async () => "access",
|
||||
getCollabToken: p.fn,
|
||||
});
|
||||
|
||||
const initial = await client.getCollabTokenWithReauth(); // provider-token-1
|
||||
let attempts = 0;
|
||||
const result = await client.writeWithCollabAuthRetry(initial, async (token) => {
|
||||
attempts++;
|
||||
return `ok:${token}`;
|
||||
});
|
||||
|
||||
assert.equal(attempts, 1, "no retry on success");
|
||||
assert.equal(result, "ok:provider-token-1");
|
||||
assert.equal(p.calls, 1, "no extra mint");
|
||||
});
|
||||
|
||||
test("a NON-auth write error propagates unchanged (no reset, no retry)", async () => {
|
||||
process.env[ENV_KEY] = "300000";
|
||||
const p = countingProvider();
|
||||
const client = new DocmostClient({
|
||||
apiUrl: "http://127.0.0.1:1/api",
|
||||
getToken: async () => "access",
|
||||
getCollabToken: p.fn,
|
||||
});
|
||||
|
||||
const initial = await client.getCollabTokenWithReauth(); // provider-token-1
|
||||
let attempts = 0;
|
||||
await assert.rejects(
|
||||
client.writeWithCollabAuthRetry(initial, async () => {
|
||||
attempts++;
|
||||
throw new Error("collab connection closed before persist"); // NOT tagged.
|
||||
}),
|
||||
/closed before persist/,
|
||||
);
|
||||
|
||||
assert.equal(attempts, 1, "a non-auth error is not retried");
|
||||
assert.equal(p.calls, 1, "the cache is untouched -> no fresh mint");
|
||||
|
||||
// Cache still holds the original token (was never invalidated).
|
||||
const still = await client.getCollabTokenWithReauth();
|
||||
assert.equal(still, "provider-token-1");
|
||||
assert.equal(p.calls, 1);
|
||||
});
|
||||
@@ -0,0 +1,55 @@
|
||||
// Unit test: listPages tree mode must propagate the `truncated` flag (#486).
|
||||
//
|
||||
// enumerateSpacePages returns { pages, truncated } — truncated is true ONLY when
|
||||
// the stdio-fallback BFS hit its node cap (the primary /pages/tree path is
|
||||
// uncapped). The old tree-mode listPages destructured only `pages` and returned a
|
||||
// bare tree, dropping `truncated`, so a caller handed an INCOMPLETE tree had no
|
||||
// way to know pages were missing. The fix returns { tree, truncated } (same
|
||||
// pattern check_new_comments uses).
|
||||
//
|
||||
// Reaching the real cap (MAX_NODES = 10000) in a mock is impractical, so we stub
|
||||
// enumerateSpacePages directly to assert the flag is threaded through verbatim.
|
||||
import { test } from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
import { DocmostClient } from "../../build/client.js";
|
||||
|
||||
function stubClient() {
|
||||
const client = new DocmostClient({
|
||||
apiUrl: "http://127.0.0.1:1/api",
|
||||
getToken: async () => "access",
|
||||
});
|
||||
// No network: the tree path only calls ensureAuthenticated + enumerateSpacePages.
|
||||
client.ensureAuthenticated = async () => {};
|
||||
return client;
|
||||
}
|
||||
|
||||
const onePage = [{ id: "r1", title: "Root", parentPageId: null }];
|
||||
|
||||
test("tree mode carries truncated:true when the enumeration truncated (#486)", async () => {
|
||||
const client = stubClient();
|
||||
client.enumerateSpacePages = async () => ({ pages: onePage, truncated: true });
|
||||
|
||||
const res = await client.listPages("space-1", 50, true);
|
||||
|
||||
assert.equal(res.truncated, true, "the truncated flag is threaded through");
|
||||
assert.ok(Array.isArray(res.tree), "the built tree rides alongside the flag");
|
||||
assert.equal(res.tree[0].id, "r1");
|
||||
});
|
||||
|
||||
test("tree mode carries truncated:false for a complete enumeration", async () => {
|
||||
const client = stubClient();
|
||||
client.enumerateSpacePages = async () => ({ pages: onePage, truncated: false });
|
||||
|
||||
const res = await client.listPages("space-1", 50, true);
|
||||
|
||||
assert.equal(res.truncated, false);
|
||||
assert.equal(res.tree[0].id, "r1");
|
||||
});
|
||||
|
||||
test("tree mode still requires a spaceId", async () => {
|
||||
const client = stubClient();
|
||||
await assert.rejects(
|
||||
client.listPages(undefined, 50, true),
|
||||
/tree mode requires a spaceId/,
|
||||
);
|
||||
});
|
||||
@@ -0,0 +1,143 @@
|
||||
// #487 commit 1 — the in-app tool cancellation safe-point inside paginateAll.
|
||||
//
|
||||
// The in-app tool host sets a composite abort signal on the client
|
||||
// (setToolAbortSignal) before each tool call; paginateAll checks it at a
|
||||
// safe-point BEFORE every sequential page fetch, so a Stop that lands mid-read
|
||||
// stops the NEXT HTTP request from STARTING (a read tool can no longer paginate
|
||||
// for minutes past a Stop). This pins the HONEST observable property against the
|
||||
// REAL client + a real HTTP server: "after Stop, no NEW request starts".
|
||||
import { test, after } from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
import http from "node:http";
|
||||
import { DocmostClient } from "../../build/client.js";
|
||||
|
||||
function readBody(req) {
|
||||
return new Promise((resolve) => {
|
||||
let raw = "";
|
||||
req.on("data", (c) => (raw += c));
|
||||
req.on("end", () => resolve(raw));
|
||||
});
|
||||
}
|
||||
function sendJson(res, status, obj, extra = {}) {
|
||||
res.writeHead(status, { "Content-Type": "application/json", ...extra });
|
||||
res.end(JSON.stringify(obj));
|
||||
}
|
||||
const openServers = [];
|
||||
async function spawn(handler) {
|
||||
const server = await new Promise((resolve) => {
|
||||
const s = http.createServer(handler);
|
||||
s.listen(0, "127.0.0.1", () => resolve(s));
|
||||
});
|
||||
openServers.push(server);
|
||||
const { port } = server.address();
|
||||
return { baseURL: `http://127.0.0.1:${port}/api` };
|
||||
}
|
||||
after(async () => {
|
||||
await Promise.all(openServers.map((s) => new Promise((r) => s.close(r))));
|
||||
});
|
||||
function handleLogin(req, res) {
|
||||
if (req.url === "/api/auth/login") {
|
||||
sendJson(res, 200, { success: true }, {
|
||||
"Set-Cookie": "authToken=t; Path=/; HttpOnly",
|
||||
});
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
// A Stop that lands DURING pagination: the server aborts the client signal as it
|
||||
// serves page 1 (more pages remain). The loop's next safe-point must throw before
|
||||
// the page-2 request is sent.
|
||||
test("paginateAll stops the NEXT request when the signal aborts mid-pagination", async () => {
|
||||
let requests = 0;
|
||||
const ac = new AbortController();
|
||||
const { baseURL } = await spawn(async (req, res) => {
|
||||
await readBody(req);
|
||||
if (handleLogin(req, res)) return;
|
||||
if (req.url === "/api/spaces") {
|
||||
requests++;
|
||||
// Simulate a user Stop that lands while page 1 is in flight.
|
||||
if (requests === 1) ac.abort(new Error("user stop"));
|
||||
sendJson(res, 200, {
|
||||
success: true,
|
||||
data: {
|
||||
items: [{ id: `p${requests}` }],
|
||||
meta: { hasNextPage: true, nextCursor: `c${requests}` },
|
||||
},
|
||||
});
|
||||
return;
|
||||
}
|
||||
sendJson(res, 404, {});
|
||||
});
|
||||
|
||||
const client = new DocmostClient(baseURL, "user@example.com", "pw");
|
||||
client.setToolAbortSignal(ac.signal);
|
||||
|
||||
await assert.rejects(
|
||||
() => client.paginateAll("/spaces", {}),
|
||||
/user stop/,
|
||||
"the aborted safe-point rejects with the signal's reason",
|
||||
);
|
||||
assert.equal(requests, 1, "page 2 never started after the Stop");
|
||||
});
|
||||
|
||||
// A Stop that is already in effect before the read starts: zero requests fire.
|
||||
test("paginateAll starts no request when the signal is already aborted", async () => {
|
||||
let requests = 0;
|
||||
const { baseURL } = await spawn(async (req, res) => {
|
||||
await readBody(req);
|
||||
if (handleLogin(req, res)) return;
|
||||
if (req.url === "/api/spaces") {
|
||||
requests++;
|
||||
sendJson(res, 200, {
|
||||
success: true,
|
||||
data: { items: [], meta: { hasNextPage: false, nextCursor: null } },
|
||||
});
|
||||
return;
|
||||
}
|
||||
sendJson(res, 404, {});
|
||||
});
|
||||
|
||||
const client = new DocmostClient(baseURL, "user@example.com", "pw");
|
||||
// Warm the auth so ensureAuthenticated does not itself POST after the abort.
|
||||
await client.ensureAuthenticated();
|
||||
const ac = new AbortController();
|
||||
ac.abort(new Error("already stopped"));
|
||||
client.setToolAbortSignal(ac.signal);
|
||||
|
||||
await assert.rejects(() => client.paginateAll("/spaces", {}), /already stopped/);
|
||||
assert.equal(requests, 0, "no /spaces request started once already aborted");
|
||||
});
|
||||
|
||||
// Without a tool signal (default), pagination is unaffected — the safe-point is a
|
||||
// pure no-op, so pre-#487 behaviour is byte-identical.
|
||||
test("paginateAll is unaffected when no tool signal is set", async () => {
|
||||
let requests = 0;
|
||||
const PAGES = {
|
||||
"": { items: [{ id: "a" }], nextCursor: "c1" },
|
||||
c1: { items: [{ id: "b" }], nextCursor: null },
|
||||
};
|
||||
const { baseURL } = await spawn(async (req, res) => {
|
||||
const raw = await readBody(req);
|
||||
if (handleLogin(req, res)) return;
|
||||
if (req.url === "/api/spaces") {
|
||||
requests++;
|
||||
const body = JSON.parse(raw || "{}");
|
||||
const page = PAGES[body.cursor ?? ""] ?? { items: [], nextCursor: null };
|
||||
sendJson(res, 200, {
|
||||
success: true,
|
||||
data: {
|
||||
items: page.items,
|
||||
meta: { hasNextPage: page.nextCursor != null, nextCursor: page.nextCursor },
|
||||
},
|
||||
});
|
||||
return;
|
||||
}
|
||||
sendJson(res, 404, {});
|
||||
});
|
||||
|
||||
const client = new DocmostClient(baseURL, "user@example.com", "pw");
|
||||
const all = await client.paginateAll("/spaces", {});
|
||||
assert.equal(requests, 2, "both pages fetched with no signal set");
|
||||
assert.deepEqual(all.map((p) => p.id), ["a", "b"]);
|
||||
});
|
||||
@@ -184,12 +184,14 @@ test("enumerateSpacePages (via listPages tree) uses one /pages/tree request", as
|
||||
});
|
||||
|
||||
const client = new DocmostClient(baseURL, "user@example.com", "pw");
|
||||
// listPages tree:true -> enumerateSpacePages(spaceId) -> buildPageTree.
|
||||
const tree = await client.listPages("space-1", 50, true);
|
||||
// listPages tree:true -> enumerateSpacePages(spaceId) -> { tree, truncated }.
|
||||
const { tree, truncated } = await client.listPages("space-1", 50, true);
|
||||
|
||||
assert.equal(treeRequests, 1, "exactly one /pages/tree request for the space");
|
||||
assert.equal(sidebarRequests, 0, "no per-node sidebar BFS requests");
|
||||
assert.deepEqual(treeBody, { spaceId: "space-1" }, "space scope posts spaceId only");
|
||||
// The uncapped /pages/tree path is never truncated (#486).
|
||||
assert.equal(truncated, false, "primary /pages/tree path is not truncated");
|
||||
// buildPageTree nests c1 under r1; two roots at the top level.
|
||||
assert.equal(tree.length, 2, "two root nodes");
|
||||
const r1 = tree.find((n) => n.id === "r1");
|
||||
@@ -249,7 +251,7 @@ test("enumerateSpacePages falls back to the cursor BFS on /pages/tree 404", asyn
|
||||
});
|
||||
|
||||
const client = new DocmostClient(baseURL, "user@example.com", "pw");
|
||||
const tree = await client.listPages("space-1", 50, true);
|
||||
const { tree, truncated } = await client.listPages("space-1", 50, true);
|
||||
|
||||
assert.ok(treeRequests >= 1, "the tree endpoint was attempted first");
|
||||
assert.deepEqual(
|
||||
@@ -257,6 +259,8 @@ test("enumerateSpacePages falls back to the cursor BFS on /pages/tree 404", asyn
|
||||
["<root>", "r1"],
|
||||
"fell back to the sidebar BFS: roots then the root's children",
|
||||
);
|
||||
// Small fallback walk well under the node cap -> not truncated (#486).
|
||||
assert.equal(truncated, false, "fallback BFS below the cap is not truncated");
|
||||
assert.equal(tree.length, 1, "one root in the built tree");
|
||||
assert.equal(tree[0].children[0].id, "c1", "leaf nested via the BFS");
|
||||
});
|
||||
@@ -438,3 +442,71 @@ test("checkNewComments subtree includes the root without a separate getPageRaw",
|
||||
assert.equal(result.checkedPages, 2, "root + one descendant scanned");
|
||||
assert.equal(result.totalNewComments, 1, "the root's fresh comment found");
|
||||
});
|
||||
|
||||
// -----------------------------------------------------------------------------
|
||||
// 6) checkNewComments parallelism (#490): the per-page comment fetches run with
|
||||
// bounded concurrency (not one-at-a-time), and the results still preserve the
|
||||
// page order deterministically regardless of which fetch finishes first.
|
||||
// -----------------------------------------------------------------------------
|
||||
test("checkNewComments fetches pages concurrently (bounded) and preserves order", async () => {
|
||||
// A subtree with 12 descendants so the scan has plenty to parallelize.
|
||||
const NODES = [{ id: "parent", title: "Parent", parentPageId: null, hasChildren: true }];
|
||||
for (let i = 0; i < 12; i++) {
|
||||
NODES.push({ id: `k${i}`, title: `Kid ${i}`, parentPageId: "parent", hasChildren: false });
|
||||
}
|
||||
|
||||
let inFlight = 0;
|
||||
let maxInFlight = 0;
|
||||
|
||||
const { baseURL } = await spawn(async (req, res) => {
|
||||
const raw = await readBody(req);
|
||||
if (handleLogin(req, res)) return;
|
||||
if (req.url === "/api/pages/tree") {
|
||||
sendJson(res, 200, { success: true, data: { items: NODES } });
|
||||
return;
|
||||
}
|
||||
if (req.url === "/api/comments") {
|
||||
const body = JSON.parse(raw || "{}");
|
||||
inFlight++;
|
||||
maxInFlight = Math.max(maxInFlight, inFlight);
|
||||
// Hold the response briefly so concurrent fetches actually overlap.
|
||||
setTimeout(() => {
|
||||
inFlight--;
|
||||
// Every page carries one fresh comment so ordering is observable.
|
||||
sendJson(res, 200, {
|
||||
success: true,
|
||||
data: {
|
||||
items: [
|
||||
{ id: `c-${body.pageId}`, createdAt: "2030-01-01T00:00:00.000Z", content: null },
|
||||
],
|
||||
meta: { nextCursor: null },
|
||||
},
|
||||
});
|
||||
}, 25);
|
||||
return;
|
||||
}
|
||||
sendJson(res, 404, {});
|
||||
});
|
||||
|
||||
const client = new DocmostClient(baseURL, "user@example.com", "pw");
|
||||
const result = await client.checkNewComments(
|
||||
"space-1",
|
||||
"2020-01-01T00:00:00.000Z",
|
||||
"parent",
|
||||
);
|
||||
|
||||
// 13 pages (parent + 12 kids) were scanned; each had a fresh comment.
|
||||
assert.equal(result.checkedPages, 13, "all pages scanned");
|
||||
assert.equal(result.totalNewComments, 13, "one fresh comment per page");
|
||||
// Parallelism: more than one request was in flight at once, but never above the
|
||||
// cap (6). A serial implementation would show maxInFlight === 1.
|
||||
assert.ok(maxInFlight > 1, `expected concurrent fetches, saw max ${maxInFlight}`);
|
||||
assert.ok(maxInFlight <= 6, `concurrency must be bounded, saw ${maxInFlight}`);
|
||||
// Deterministic order: results follow the page-enumeration order (parent first).
|
||||
assert.equal(result.comments[0].pageId, "parent", "results preserve page order");
|
||||
assert.deepEqual(
|
||||
result.comments.map((r) => r.pageId),
|
||||
["parent", ...Array.from({ length: 12 }, (_, i) => `k${i}`)],
|
||||
"result order matches the enumeration order regardless of finish order",
|
||||
);
|
||||
});
|
||||
|
||||
@@ -0,0 +1,164 @@
|
||||
// #487 F4 — the WRITE-side cancellation safe-point.
|
||||
//
|
||||
// Every content-mutating collab write (collaboration.mutatePageContent and the
|
||||
// reentrant twin client.mutateLiveContentUnlocked used by replaceImage) checks
|
||||
// the in-app tool abort signal at a PRE-COMMIT safe-point — after the collab
|
||||
// session is acquired but immediately BEFORE the atomic read->write
|
||||
// (session.mutate). So a Stop (or the per-call cap) that lands during the
|
||||
// connect/lock window stops THIS write from landing: no new commit starts once
|
||||
// aborted. paginate-abort-safepoint.test.mjs pins the READ half; this pins the
|
||||
// integrity-critical WRITE half — remove the `throwIfAborted()` and the transform
|
||||
// would run and the doc would be mutated past a Stop.
|
||||
//
|
||||
// There is no collab server in the unit env, so we swap the provider factory
|
||||
// (__setCollabProviderFactory) for a fake that reports an immediate successful
|
||||
// sync. That makes acquireCollabSession SUCCEED and hand back a live, ready
|
||||
// session, so the ONLY thing standing between the call and session.mutate is the
|
||||
// safe-point under test. The transform is instrumented to prove it never runs.
|
||||
import { test, afterEach } from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
import { mutatePageContent } from "../../build/lib/collaboration.js";
|
||||
import {
|
||||
__setCollabProviderFactory,
|
||||
destroyAllSessions,
|
||||
} from "../../build/lib/collab-session.js";
|
||||
import { DocmostClient } from "../../build/client.js";
|
||||
|
||||
const BASE_URL = "http://127.0.0.1:1/api";
|
||||
// mutatePageContent locks via withPageLock, which demands a canonical page UUID
|
||||
// (resolve-then-lock invariant, #260/#449); the unlocked twin does not.
|
||||
const PAGE_UUID = "11111111-1111-4111-8111-111111111111";
|
||||
|
||||
// A fake HocuspocusProvider that immediately reports a successful initial sync so
|
||||
// CollabSession.open() resolves to a ready session, and stays "synced" with zero
|
||||
// unsynced changes so a reached session.mutate() would resolve at once. It speaks
|
||||
// only the tiny CollabProviderLike surface the session depends on.
|
||||
function syncedProviderFactory(config) {
|
||||
// Fire the initial-sync callback so open() settles as ready.
|
||||
config.onSynced();
|
||||
return {
|
||||
synced: true,
|
||||
unsyncedChanges: 0,
|
||||
destroy() {},
|
||||
on() {},
|
||||
off() {},
|
||||
};
|
||||
}
|
||||
|
||||
// Disable the session cache so every acquire opens (and the failure path destroys)
|
||||
// its own ephemeral session — no cross-test session reuse.
|
||||
process.env.MCP_COLLAB_SESSION_IDLE_MS = "0";
|
||||
|
||||
afterEach(() => {
|
||||
__setCollabProviderFactory(null); // restore the real factory
|
||||
destroyAllSessions();
|
||||
});
|
||||
|
||||
// --- collaboration.mutatePageContent (the page-locked write path) -------------
|
||||
|
||||
test("mutatePageContent rejects at the pre-commit safe-point BEFORE session.mutate when the signal is already aborted", async () => {
|
||||
__setCollabProviderFactory(syncedProviderFactory);
|
||||
let transformCalls = 0;
|
||||
const ac = new AbortController();
|
||||
ac.abort(new Error("user stop"));
|
||||
|
||||
await assert.rejects(
|
||||
() =>
|
||||
mutatePageContent(
|
||||
PAGE_UUID,
|
||||
"collab-jwt",
|
||||
BASE_URL,
|
||||
(liveDoc) => {
|
||||
transformCalls++;
|
||||
return liveDoc;
|
||||
},
|
||||
ac.signal,
|
||||
),
|
||||
/user stop/,
|
||||
"the aborted safe-point rejects with the signal's reason before committing",
|
||||
);
|
||||
assert.equal(
|
||||
transformCalls,
|
||||
0,
|
||||
"the transform (and therefore session.mutate) must NEVER run once aborted",
|
||||
);
|
||||
});
|
||||
|
||||
test("mutatePageContent (control) DOES reach session.mutate and invoke the transform when the signal is live", async () => {
|
||||
__setCollabProviderFactory(syncedProviderFactory);
|
||||
let transformCalls = 0;
|
||||
const ac = new AbortController(); // never aborted
|
||||
|
||||
const result = await mutatePageContent(
|
||||
PAGE_UUID,
|
||||
"collab-jwt",
|
||||
BASE_URL,
|
||||
(liveDoc) => {
|
||||
transformCalls++;
|
||||
return null; // null -> no-op write; still proves the transform was invoked
|
||||
},
|
||||
ac.signal,
|
||||
);
|
||||
assert.equal(
|
||||
transformCalls,
|
||||
1,
|
||||
"with a live signal the safe-point is a no-op and session.mutate runs the transform",
|
||||
);
|
||||
assert.ok(result && result.verify, "a MutationResult is returned");
|
||||
});
|
||||
|
||||
// --- client.mutateLiveContentUnlocked (the reentrant twin, replaceImage) ------
|
||||
|
||||
test("mutateLiveContentUnlocked rejects at the pre-commit safe-point BEFORE session.mutate when the tool signal is already aborted", async () => {
|
||||
__setCollabProviderFactory(syncedProviderFactory);
|
||||
const client = new DocmostClient({
|
||||
apiUrl: BASE_URL,
|
||||
getToken: async () => "access",
|
||||
getCollabToken: async () => "collab-jwt",
|
||||
});
|
||||
let transformCalls = 0;
|
||||
const ac = new AbortController();
|
||||
ac.abort(new Error("cap fired"));
|
||||
client.setToolAbortSignal(ac.signal);
|
||||
|
||||
await assert.rejects(
|
||||
() =>
|
||||
client.mutateLiveContentUnlocked("page-1", "collab-jwt", (liveDoc) => {
|
||||
transformCalls++;
|
||||
return liveDoc;
|
||||
}),
|
||||
/cap fired/,
|
||||
"the aborted safe-point rejects with the signal's reason before committing",
|
||||
);
|
||||
assert.equal(
|
||||
transformCalls,
|
||||
0,
|
||||
"the transform (and therefore session.mutate) must NEVER run once aborted",
|
||||
);
|
||||
});
|
||||
|
||||
test("mutateLiveContentUnlocked (control) DOES reach session.mutate and invoke the transform when the tool signal is live", async () => {
|
||||
__setCollabProviderFactory(syncedProviderFactory);
|
||||
const client = new DocmostClient({
|
||||
apiUrl: BASE_URL,
|
||||
getToken: async () => "access",
|
||||
getCollabToken: async () => "collab-jwt",
|
||||
});
|
||||
let transformCalls = 0;
|
||||
client.setToolAbortSignal(new AbortController().signal); // live
|
||||
|
||||
const result = await client.mutateLiveContentUnlocked(
|
||||
"page-1",
|
||||
"collab-jwt",
|
||||
(liveDoc) => {
|
||||
transformCalls++;
|
||||
return null;
|
||||
},
|
||||
);
|
||||
assert.equal(
|
||||
transformCalls,
|
||||
1,
|
||||
"with a live signal the safe-point is a no-op and session.mutate runs the transform",
|
||||
);
|
||||
assert.ok(result && result.verify, "a MutationResult is returned");
|
||||
});
|
||||
@@ -101,6 +101,88 @@ test("DoS guard: a graph over the node cap is returned unchanged, quickly", asyn
|
||||
assert.ok(dt < 2000, `cap path should be fast, took ${dt}ms`);
|
||||
});
|
||||
|
||||
/** Build a layered DAG near the caps: `n` vertices, up to ~2 edges each into the
|
||||
* next layer of `layerSize`. Used as a real worst-case graph for the benchmark. */
|
||||
function layeredGraph(n, layerSize) {
|
||||
let cells = "";
|
||||
for (let i = 2; i < 2 + n; i++) {
|
||||
cells +=
|
||||
`<mxCell id="${i}" value="N${i}" style="rounded=1;html=1;" vertex="1" parent="1">` +
|
||||
`<mxGeometry x="10" y="10" width="120" height="60" as="geometry"/></mxCell>`;
|
||||
}
|
||||
let ei = 0;
|
||||
for (let i = 2; i < 2 + n; i++) {
|
||||
for (const off of [layerSize, layerSize + 1]) {
|
||||
const t = i + off;
|
||||
if (t < 2 + n) cells += `<mxCell id="e${ei++}" edge="1" parent="1" source="${i}" target="${t}"><mxGeometry relative="1" as="geometry"/></mxCell>`;
|
||||
}
|
||||
}
|
||||
return (
|
||||
'<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>' +
|
||||
cells +
|
||||
"</root></mxGraphModel>"
|
||||
);
|
||||
}
|
||||
|
||||
test("terminate-on-timeout: a layout that exceeds the wall-clock ceiling is hard-killed and the original model is returned (#486)", async () => {
|
||||
// A 1ms ceiling fires before the worker can even finish loading elkjs, so the
|
||||
// parent must terminate() the worker and fall back to the ORIGINAL model. On
|
||||
// the OLD in-process race this timer could never fire while the SAME thread was
|
||||
// blocked inside elkjs — the fallback path was unreachable; here it works.
|
||||
const prev = process.env.DRAWIO_ELK_TIMEOUT_MS;
|
||||
process.env.DRAWIO_ELK_TIMEOUT_MS = "1";
|
||||
try {
|
||||
const model = layeredGraph(400, 20);
|
||||
const t0 = Date.now();
|
||||
const laid = await applyElkLayout(model);
|
||||
const dt = Date.now() - t0;
|
||||
// Original geometry is preserved verbatim: every vertex is still stacked at
|
||||
// (10,10), proving NO ELK coordinates were applied (the pass was killed).
|
||||
const verts = parseCells(laid).filter((c) => c.vertex);
|
||||
assert.equal(verts.length, 400, "all vertices survived the fallback");
|
||||
for (const v of verts) {
|
||||
assert.equal(v.geometry.x, 10, "x untouched -> layout was terminated");
|
||||
assert.equal(v.geometry.y, 10, "y untouched -> layout was terminated");
|
||||
}
|
||||
// The kill is prompt: terminate() returns the call well under the natural
|
||||
// layout time for a 400-node graph.
|
||||
assert.ok(dt < 2000, `terminate path should be prompt, took ${dt}ms`);
|
||||
} finally {
|
||||
if (prev === undefined) delete process.env.DRAWIO_ELK_TIMEOUT_MS;
|
||||
else process.env.DRAWIO_ELK_TIMEOUT_MS = prev;
|
||||
}
|
||||
});
|
||||
|
||||
test("benchmark guard: a worst-case graph AT the cap lays out without wedging the main event loop (#486)", async () => {
|
||||
// ~500 nodes / ~1000 edges — a real worst case at the node/edge caps. The
|
||||
// layout runs on a WORKER thread, so the MAIN event loop must stay responsive
|
||||
// throughout: a timer scheduled on the main thread keeps firing while ELK
|
||||
// churns. On the OLD synchronous-on-main-thread code this counter would be
|
||||
// pinned at 0 for the whole layout (event loop wedged) — exactly the prod fire.
|
||||
const model = layeredGraph(500, 20);
|
||||
let mainLoopTicks = 0;
|
||||
const iv = setInterval(() => {
|
||||
mainLoopTicks++;
|
||||
}, 2);
|
||||
const t0 = Date.now();
|
||||
const laid = await applyElkLayout(model);
|
||||
const dt = Date.now() - t0;
|
||||
clearInterval(iv);
|
||||
|
||||
assert.ok(
|
||||
mainLoopTicks > 0,
|
||||
"main event loop must stay responsive while ELK runs on the worker",
|
||||
);
|
||||
// Benchmark guard: the worst-case graph actually LAYS OUT within the default
|
||||
// ceiling (it did not fall back). At least one vertex moved off the stack.
|
||||
const verts = parseCells(laid).filter((c) => c.vertex);
|
||||
assert.equal(verts.length, 500, "all vertices survived");
|
||||
const moved = verts.some((v) => v.geometry.x !== 10 || v.geometry.y !== 10);
|
||||
assert.ok(moved, "layout was applied (did not time out / fall back)");
|
||||
// Sanity ceiling well under the 5s wall-clock timeout.
|
||||
assert.ok(dt < 5000, `worst-case layout should be under the ceiling, took ${dt}ms`);
|
||||
});
|
||||
|
||||
test("layout is best-effort: an empty/degenerate model is returned intact", async () => {
|
||||
const model =
|
||||
'<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/></root></mxGraphModel>';
|
||||
|
||||
@@ -1,101 +1,220 @@
|
||||
import { test } from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
import { createHash } from "node:crypto";
|
||||
import { readFileSync } from "node:fs";
|
||||
import {
|
||||
mkdtempSync,
|
||||
mkdirSync,
|
||||
writeFileSync,
|
||||
rmSync,
|
||||
readdirSync,
|
||||
statSync,
|
||||
readFileSync,
|
||||
} from "node:fs";
|
||||
import { tmpdir } from "node:os";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import { dirname, join } from "node:path";
|
||||
import { dirname, join, relative, sep } from "node:path";
|
||||
import { createHash } from "node:crypto";
|
||||
|
||||
import { computeRegistryStamp } from "../../scripts/gen-registry-stamp.mjs";
|
||||
import { REGISTRY_STAMP } from "../../build/index.js";
|
||||
|
||||
// Guard tests for the build/src-skew stamp (issue #447). The codegen script
|
||||
// exports `computeRegistryStamp(sourceText)` — a sha256 over normalized source
|
||||
// text (CRLF->LF, single trailing newline stripped). The in-app loader
|
||||
// (apps/server/.../docmost-client.loader.ts) DUPLICATES that normalize+sha256 to
|
||||
// recompute the stamp from src and refuse a stale build. These tests pin the
|
||||
// algorithm's behaviour AND assert the built stamp matches the current src, so a
|
||||
// stale generated file OR a normalize divergence reddens.
|
||||
// Guard tests for the build/src-skew stamp (issues #447/#486). The codegen script
|
||||
// exports `computeRegistryStamp(srcDir)` — a sha256 over the WHOLE src/ tree
|
||||
// (every src/**/*.ts EXCEPT *.generated.ts), each file folded in as its
|
||||
// POSIX-relative path + its normalized content (CRLF->LF, single trailing newline
|
||||
// stripped). Hashing the whole tree (not just tool-specs.ts) is #486: an edit to
|
||||
// client.ts / a client/* module without a rebuild must ALSO redden. The in-app
|
||||
// loader (apps/server/.../docmost-client.loader.ts) DUPLICATES this enumerate+
|
||||
// normalize+sha256 to refuse a stale build. These tests pin the algorithm and
|
||||
// assert the built stamp matches the current src.
|
||||
|
||||
const __dirname = dirname(fileURLToPath(import.meta.url));
|
||||
const TOOL_SPECS_PATH = join(__dirname, "..", "..", "src", "tool-specs.ts");
|
||||
const SRC_DIR = join(__dirname, "..", "..", "src");
|
||||
|
||||
test("computeRegistryStamp is deterministic: same input -> same hash", () => {
|
||||
const input = "export const X = 1;\nexport const Y = 2;\n";
|
||||
assert.equal(computeRegistryStamp(input), computeRegistryStamp(input));
|
||||
// Build a throwaway src/ tree from a { relPath: content } map and return its dir.
|
||||
function makeSrcTree(files) {
|
||||
const root = mkdtempSync(join(tmpdir(), "mcp-stamp-tree-"));
|
||||
const src = join(root, "src");
|
||||
for (const [rel, content] of Object.entries(files)) {
|
||||
const full = join(src, rel);
|
||||
mkdirSync(dirname(full), { recursive: true });
|
||||
writeFileSync(full, content, "utf8");
|
||||
}
|
||||
return { src, cleanup: () => rmSync(root, { recursive: true, force: true }) };
|
||||
}
|
||||
|
||||
test("computeRegistryStamp is deterministic: same tree -> same hash", () => {
|
||||
const a = makeSrcTree({ "tool-specs.ts": "export const X = 1;\n" });
|
||||
const b = makeSrcTree({ "tool-specs.ts": "export const X = 1;\n" });
|
||||
try {
|
||||
assert.equal(computeRegistryStamp(a.src), computeRegistryStamp(b.src));
|
||||
} finally {
|
||||
a.cleanup();
|
||||
b.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
test("computeRegistryStamp returns a 64-char lowercase hex sha256", () => {
|
||||
const stamp = computeRegistryStamp("anything");
|
||||
assert.match(stamp, /^[0-9a-f]{64}$/);
|
||||
const t = makeSrcTree({ "tool-specs.ts": "anything\n" });
|
||||
try {
|
||||
assert.match(computeRegistryStamp(t.src), /^[0-9a-f]{64}$/);
|
||||
} finally {
|
||||
t.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
test("normalizes CRLF vs LF: the same content hashes equal", () => {
|
||||
const lf = "line1\nline2\nline3";
|
||||
const crlf = "line1\r\nline2\r\nline3";
|
||||
assert.equal(computeRegistryStamp(crlf), computeRegistryStamp(lf));
|
||||
// #486 CORE: an edit to a NON-tool-specs source file (client.ts) must change the
|
||||
// stamp. Under the old single-file (tool-specs.ts only) hash this edit was
|
||||
// invisible and a stale build/ served the old client.ts silently.
|
||||
test("editing client.ts (not tool-specs.ts) changes the stamp (#486)", () => {
|
||||
const before = makeSrcTree({
|
||||
"tool-specs.ts": "export const SPECS = 1;\n",
|
||||
"client.ts": "export const impl = 'v1';\n",
|
||||
});
|
||||
const after = makeSrcTree({
|
||||
"tool-specs.ts": "export const SPECS = 1;\n",
|
||||
"client.ts": "export const impl = 'v2';\n",
|
||||
});
|
||||
try {
|
||||
assert.notEqual(
|
||||
computeRegistryStamp(before.src),
|
||||
computeRegistryStamp(after.src),
|
||||
"a client.ts edit with an unchanged tool-specs.ts must move the stamp",
|
||||
);
|
||||
} finally {
|
||||
before.cleanup();
|
||||
after.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
test("normalizes a trailing newline: with/without a final \\n hashes equal", () => {
|
||||
const noTrailing = "alpha\nbeta";
|
||||
const trailing = "alpha\nbeta\n";
|
||||
assert.equal(computeRegistryStamp(trailing), computeRegistryStamp(noTrailing));
|
||||
test("editing a nested client/* module changes the stamp", () => {
|
||||
const before = makeSrcTree({
|
||||
"tool-specs.ts": "x\n",
|
||||
"client/read.ts": "export const READ = 1;\n",
|
||||
});
|
||||
const after = makeSrcTree({
|
||||
"tool-specs.ts": "x\n",
|
||||
"client/read.ts": "export const READ = 2;\n",
|
||||
});
|
||||
try {
|
||||
assert.notEqual(
|
||||
computeRegistryStamp(before.src),
|
||||
computeRegistryStamp(after.src),
|
||||
);
|
||||
} finally {
|
||||
before.cleanup();
|
||||
after.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
test("a CRLF checkout WITH a trailing CRLF still hashes equal to bare LF", () => {
|
||||
// A worst-case Windows checkout: CRLF line endings + a trailing CRLF. Both the
|
||||
// \r\n->\n replace and the trailing-newline strip must apply for parity.
|
||||
const bare = "alpha\nbeta";
|
||||
const crlfTrailing = "alpha\r\nbeta\r\n";
|
||||
assert.equal(
|
||||
computeRegistryStamp(crlfTrailing),
|
||||
computeRegistryStamp(bare),
|
||||
);
|
||||
// *.generated.ts is EXCLUDED (else the codegen's own output is a fixed-point
|
||||
// cycle): adding/removing/changing it must not move the stamp.
|
||||
test("*.generated.ts is excluded from the stamp", () => {
|
||||
const without = makeSrcTree({ "tool-specs.ts": "x\n" });
|
||||
const withGen = makeSrcTree({
|
||||
"tool-specs.ts": "x\n",
|
||||
"registry-stamp.generated.ts": 'export const REGISTRY_STAMP = "abc";\n',
|
||||
});
|
||||
try {
|
||||
assert.equal(
|
||||
computeRegistryStamp(without.src),
|
||||
computeRegistryStamp(withGen.src),
|
||||
"a *.generated.ts file must not affect the stamp",
|
||||
);
|
||||
} finally {
|
||||
without.cleanup();
|
||||
withGen.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
test("a real content change hashes differently", () => {
|
||||
const before = "export const description = 'search a page';\n";
|
||||
const after = "export const description = 'search a PAGE';\n";
|
||||
assert.notEqual(computeRegistryStamp(before), computeRegistryStamp(after));
|
||||
test("a CRLF checkout WITH trailing CRLF hashes equal to bare LF", () => {
|
||||
const bare = makeSrcTree({ "tool-specs.ts": "alpha\nbeta" });
|
||||
const crlfTrailing = makeSrcTree({ "tool-specs.ts": "alpha\r\nbeta\r\n" });
|
||||
try {
|
||||
assert.equal(
|
||||
computeRegistryStamp(crlfTrailing.src),
|
||||
computeRegistryStamp(bare.src),
|
||||
);
|
||||
} finally {
|
||||
bare.cleanup();
|
||||
crlfTrailing.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
// Only a SINGLE trailing newline is stripped — a second blank line is content and
|
||||
// must change the hash. This pins the exact `/\n$/` semantics the loader mirrors.
|
||||
// Only a SINGLE trailing newline is stripped — a second blank line is content.
|
||||
test("only ONE trailing newline is stripped (two differ from one)", () => {
|
||||
assert.notEqual(
|
||||
computeRegistryStamp("x\n"),
|
||||
computeRegistryStamp("x\n\n"),
|
||||
);
|
||||
const one = makeSrcTree({ "tool-specs.ts": "x\n" });
|
||||
const two = makeSrcTree({ "tool-specs.ts": "x\n\n" });
|
||||
try {
|
||||
assert.notEqual(
|
||||
computeRegistryStamp(one.src),
|
||||
computeRegistryStamp(two.src),
|
||||
);
|
||||
} finally {
|
||||
one.cleanup();
|
||||
two.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
// Cross-impl equality against a fixed, documented input. The SAME literal input
|
||||
// and expected hash are asserted in the server-side jest test
|
||||
// (docmost-client.loader.spec.ts). If either side's normalize+sha256 ever
|
||||
// diverges, one of the two tests reddens. Input exercises BOTH normalize steps.
|
||||
test("fixed-input hash matches the documented cross-impl value", () => {
|
||||
const FIXED_INPUT = "line1\r\nline2\n";
|
||||
const EXPECTED =
|
||||
"683376e290829b482c2655745caffa7a1dccfa10afaa62dac2b42dd6c68d0f83";
|
||||
assert.equal(computeRegistryStamp(FIXED_INPUT), EXPECTED);
|
||||
// Cross-impl equality against a fixed, documented tree. The SAME literal tree and
|
||||
// expected hash are asserted in the server-side jest test
|
||||
// (docmost-client.loader.spec.ts). If either side's enumerate+normalize+sha256
|
||||
// ever diverges, one of the two tests reddens. The tree exercises: a nested file,
|
||||
// BOTH normalize steps (tool-specs.ts uses CRLF + trailing \n) and the
|
||||
// *.generated.ts exclusion.
|
||||
const CROSS_IMPL_TREE = {
|
||||
"tool-specs.ts": "line1\r\nline2\n",
|
||||
"client/read.ts": "export const R = 1;\n",
|
||||
"registry-stamp.generated.ts": 'export const REGISTRY_STAMP="ignored";\n',
|
||||
};
|
||||
const CROSS_IMPL_EXPECTED =
|
||||
"131c1b9e4e2f5a7d6cef91ca8df619822b442f52bc45ebd09474a4c1d6728616";
|
||||
|
||||
test("fixed-tree hash matches the documented cross-impl value", () => {
|
||||
const t = makeSrcTree(CROSS_IMPL_TREE);
|
||||
try {
|
||||
assert.equal(computeRegistryStamp(t.src), CROSS_IMPL_EXPECTED);
|
||||
} finally {
|
||||
t.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
// DESYNC GUARD (covers reviewer suggestion 2). Recompute the stamp from the
|
||||
// actual src/tool-specs.ts and assert it equals the REGISTRY_STAMP baked into the
|
||||
// freshly-built build/index.js. This reddens if the generated file is stale OR if
|
||||
// the codegen normalize ever diverges from what produced the built stamp.
|
||||
test("built REGISTRY_STAMP equals the stamp recomputed from src/tool-specs.ts", () => {
|
||||
const source = readFileSync(TOOL_SPECS_PATH, "utf8");
|
||||
assert.equal(computeRegistryStamp(source), REGISTRY_STAMP);
|
||||
// Sanity: the EXPECTED constant is not a magic value but the documented
|
||||
// enumerate+normalize+sha256 of CROSS_IMPL_TREE (a local re-implementation).
|
||||
test("the documented EXPECTED is the enumerate+normalize+sha256 of the tree", () => {
|
||||
const t = makeSrcTree(CROSS_IMPL_TREE);
|
||||
try {
|
||||
const collect = (dir) => {
|
||||
const out = [];
|
||||
for (const e of readdirSync(dir)) {
|
||||
const f = join(dir, e);
|
||||
if (statSync(f).isDirectory()) out.push(...collect(f));
|
||||
else if (e.endsWith(".ts") && !e.endsWith(".generated.ts")) out.push(f);
|
||||
}
|
||||
return out;
|
||||
};
|
||||
const files = collect(t.src)
|
||||
.map((abs) => ({ rel: relative(t.src, abs).split(sep).join("/"), abs }))
|
||||
.sort((a, b) => (a.rel < b.rel ? -1 : a.rel > b.rel ? 1 : 0));
|
||||
const h = createHash("sha256");
|
||||
for (const { rel, abs } of files) {
|
||||
const n = readFileSync(abs, "utf8")
|
||||
.replace(/\r\n/g, "\n")
|
||||
.replace(/\n$/, "");
|
||||
h.update(rel, "utf8");
|
||||
h.update("\0", "utf8");
|
||||
h.update(n, "utf8");
|
||||
h.update("\0", "utf8");
|
||||
}
|
||||
assert.equal(h.digest("hex"), CROSS_IMPL_EXPECTED);
|
||||
} finally {
|
||||
t.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
// Sanity: the fixed-input helper computes the SAME way the codegen does, proving
|
||||
// the EXPECTED constant above is not an arbitrary magic value but the documented
|
||||
// normalize+sha256 of FIXED_INPUT. Belt-and-braces so a bad EXPECTED can't hide a
|
||||
// real regression.
|
||||
test("the documented EXPECTED constant is the normalize+sha256 of FIXED_INPUT", () => {
|
||||
const FIXED_INPUT = "line1\r\nline2\n";
|
||||
const normalized = FIXED_INPUT.replace(/\r\n/g, "\n").replace(/\n$/, "");
|
||||
const expected = createHash("sha256")
|
||||
.update(normalized, "utf8")
|
||||
.digest("hex");
|
||||
assert.equal(computeRegistryStamp(FIXED_INPUT), expected);
|
||||
// DESYNC GUARD. Recompute the stamp from the REAL src/ tree and assert it equals
|
||||
// the REGISTRY_STAMP baked into the freshly-built build/index.js. This reddens if
|
||||
// the generated file is stale OR if the codegen ever diverges from what produced
|
||||
// the built stamp.
|
||||
test("built REGISTRY_STAMP equals the stamp recomputed from src/", () => {
|
||||
assert.equal(computeRegistryStamp(SRC_DIR), REGISTRY_STAMP);
|
||||
});
|
||||
|
||||
@@ -2,7 +2,12 @@ import { test } from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
import { z } from "zod";
|
||||
|
||||
import { SHARED_TOOL_SPECS } from "../../build/tool-specs.js";
|
||||
import {
|
||||
SHARED_TOOL_SPECS,
|
||||
SHARED_TOOL_WRITE_CLASS,
|
||||
isRetryableWriteClass,
|
||||
assertEverySpecDeclaresWriteClass,
|
||||
} from "../../build/tool-specs.js";
|
||||
|
||||
// The shared registry is consumed by BOTH the zod-v3 MCP server and the zod-v4
|
||||
// in-app AI-SDK service, so every spec must carry the cross-layer wiring
|
||||
@@ -43,6 +48,41 @@ test("mcpName and inAppKey are each unique across the registry", () => {
|
||||
}
|
||||
});
|
||||
|
||||
// #489 — every spec must declare its write-class so the external-MCP retry path
|
||||
// can gate a single auto-retry ONLY on a pure read (a blind retry of a write =
|
||||
// double-apply). The declaration is enforced at registration time.
|
||||
test("#489: every spec declares a valid writeClass ('readOnly' | 'write')", () => {
|
||||
for (const [key, spec] of Object.entries(SHARED_TOOL_SPECS)) {
|
||||
assert.ok(
|
||||
spec.writeClass === "readOnly" || spec.writeClass === "write",
|
||||
`${key}: missing/invalid writeClass: ${JSON.stringify(spec.writeClass)}`,
|
||||
);
|
||||
}
|
||||
// The registration-time assert must not throw for the shipped registry.
|
||||
assert.doesNotThrow(() => assertEverySpecDeclaresWriteClass());
|
||||
});
|
||||
|
||||
test("#489: SHARED_TOOL_WRITE_CLASS maps every mcpName to its class; helper gates on readOnly", () => {
|
||||
const specs = Object.values(SHARED_TOOL_SPECS);
|
||||
assert.equal(Object.keys(SHARED_TOOL_WRITE_CLASS).length, specs.length);
|
||||
for (const spec of specs) {
|
||||
assert.equal(SHARED_TOOL_WRITE_CLASS[spec.mcpName], spec.writeClass);
|
||||
}
|
||||
// Only a readOnly tool is retry-eligible; a write tool and an unknown tool are not.
|
||||
assert.equal(isRetryableWriteClass("readOnly"), true);
|
||||
assert.equal(isRetryableWriteClass("write"), false);
|
||||
assert.equal(isRetryableWriteClass(undefined), false);
|
||||
});
|
||||
|
||||
test("#489: representative reads are readOnly and representative writes are write", () => {
|
||||
for (const name of ["getPage", "getTree", "searchInPage", "listComments"]) {
|
||||
assert.equal(SHARED_TOOL_SPECS[name].writeClass, "readOnly", `${name} should be readOnly`);
|
||||
}
|
||||
for (const name of ["patchNode", "createPage", "deletePage", "createComment", "drawioCreate"]) {
|
||||
assert.equal(SHARED_TOOL_SPECS[name].writeClass, "write", `${name} should be write`);
|
||||
}
|
||||
});
|
||||
|
||||
test("buildShape (when present) returns a usable ZodRawShape with a real zod", () => {
|
||||
for (const [key, spec] of Object.entries(SHARED_TOOL_SPECS)) {
|
||||
if (!spec.buildShape) continue;
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
{
|
||||
"name": "@docmost/token-estimate",
|
||||
"version": "0.1.0",
|
||||
"description": "Shared, provider-agnostic token estimator (chars/2.5) used by the AI-chat client counter and the server history-replay budgeter, so the two never diverge.",
|
||||
"private": true,
|
||||
"main": "./dist/index.js",
|
||||
"types": "./dist/index.d.ts",
|
||||
"scripts": {
|
||||
"build": "tsc",
|
||||
"watch": "tsc --watch",
|
||||
"test": "vitest run",
|
||||
"test:watch": "vitest"
|
||||
},
|
||||
"license": "MIT",
|
||||
"devDependencies": {
|
||||
"typescript": "^5.0.0",
|
||||
"vitest": "4.1.6"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { estimateTokens, CHARS_PER_TOKEN } from './index';
|
||||
|
||||
describe('estimateTokens (shared chars/2.5)', () => {
|
||||
it('returns 0 for empty / nullish input', () => {
|
||||
expect(estimateTokens('')).toBe(0);
|
||||
expect(estimateTokens(null)).toBe(0);
|
||||
expect(estimateTokens(undefined)).toBe(0);
|
||||
});
|
||||
|
||||
it('uses the chars/2.5 ratio, ceiled', () => {
|
||||
expect(CHARS_PER_TOKEN).toBe(2.5);
|
||||
// 5 chars / 2.5 = 2
|
||||
expect(estimateTokens('abcde')).toBe(2);
|
||||
// any non-empty string is at least 1 token (ceil)
|
||||
expect(estimateTokens('a')).toBe(1);
|
||||
// 100 chars / 2.5 = 40
|
||||
expect(estimateTokens('x'.repeat(100))).toBe(40);
|
||||
});
|
||||
|
||||
it('counts Cyrillic ~2x higher than the old chars/4 rule (no undercount)', () => {
|
||||
const cyr = 'привет мир как дела'; // 19 chars
|
||||
expect(estimateTokens(cyr)).toBe(Math.ceil(19 / 2.5)); // 8
|
||||
expect(estimateTokens(cyr)).toBeGreaterThan(Math.ceil(19 / 4)); // > 5
|
||||
});
|
||||
|
||||
it('is deterministic / byte-stable (same input => same output)', () => {
|
||||
const s = 'the quick brown fox';
|
||||
expect(estimateTokens(s)).toBe(estimateTokens(s));
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,35 @@
|
||||
/**
|
||||
* Shared, provider-agnostic token estimator (#490).
|
||||
*
|
||||
* No provider exposes an exact tokenizer we can afford to run on the hot path (a
|
||||
* real BPE pass is O(n²)-ish, bloats the client bundle, and is wrong for
|
||||
* Gemini/Ollama anyway), so both the client's in-body counter AND the server's
|
||||
* history-replay budgeter use this ONE cheap chars-based heuristic. Keeping it in
|
||||
* a single shared module is deliberate: two independent estimators drift, and then
|
||||
* "the badge shows 60%" while "the budgeter already trimmed" — the exact confusion
|
||||
* this package prevents.
|
||||
*
|
||||
* Ratio: **chars / 2.5**. Most content here is Cyrillic, where a token is ~2.5
|
||||
* characters; the common English `chars/4` rule of thumb UNDER-counts Cyrillic by
|
||||
* ~2×, which for a budget check is the dangerous direction (it lets the context
|
||||
* overflow). 2.5 slightly over-estimates pure English/code, which is the SAFE
|
||||
* direction for a budget. This is an estimate, never an exact count — the
|
||||
* authoritative figure is always the provider's reported usage; the estimate is
|
||||
* for UI affordances, the delta of not-yet-sent messages, and deciding what to
|
||||
* trim.
|
||||
*/
|
||||
|
||||
/** Characters per token for the shared estimate. See the module comment. */
|
||||
export const CHARS_PER_TOKEN = 2.5;
|
||||
|
||||
/**
|
||||
* Rough token estimate for a piece of text (chars / {@link CHARS_PER_TOKEN}).
|
||||
* Returns 0 for empty/nullish input, and ceils so any non-empty text counts as at
|
||||
* least one token. Pure and deterministic (byte-stable), so the same text always
|
||||
* yields the same estimate — which the server budgeter relies on to keep replay
|
||||
* trimming stable turn to turn (provider prompt-cache friendliness).
|
||||
*/
|
||||
export function estimateTokens(text: string | null | undefined): number {
|
||||
if (!text) return 0;
|
||||
return Math.ceil(text.length / CHARS_PER_TOKEN);
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"target": "ES2022",
|
||||
"module": "CommonJS",
|
||||
"moduleResolution": "Node",
|
||||
"outDir": "./dist",
|
||||
"rootDir": "./src",
|
||||
"strict": true,
|
||||
"esModuleInterop": true,
|
||||
"skipLibCheck": true,
|
||||
"forceConsistentCasingInFileNames": true,
|
||||
"declaration": true
|
||||
},
|
||||
"include": ["src/**/*"],
|
||||
"exclude": ["src/**/*.test.ts"]
|
||||
}
|
||||
+114
-6
@@ -1,8 +1,62 @@
|
||||
diff --git a/dist/index.js b/dist/index.js
|
||||
index ae447a12f7823ec0a00837ee9f0eb809a610d5f8..a3402b2c2d021ef432cfa76e35d370073d525135 100644
|
||||
index ae447a12f7823ec0a00837ee9f0eb809a610d5f8..210b5a9009e5cf1537cb2007c524007c1a01253c 100644
|
||||
--- a/dist/index.js
|
||||
+++ b/dist/index.js
|
||||
@@ -6578,9 +6578,19 @@ function createOutputTransformStream(output) {
|
||||
@@ -5036,9 +5036,40 @@ function writeToServerResponse({
|
||||
break;
|
||||
const canContinue = response.write(value);
|
||||
if (!canContinue) {
|
||||
- await new Promise((resolve3) => {
|
||||
- response.once("drain", resolve3);
|
||||
+ // PATCH(docmost #486): race "drain" against "close"/"error". The
|
||||
+ // original awaited ONLY "drain", so a client that disconnected mid-write
|
||||
+ // (the socket never drains) parked this loop FOREVER: the finally never
|
||||
+ // ran, response.end() was unreachable, and the reader + buffered chunks
|
||||
+ // were held until process restart. On close/error we cancel the reader
|
||||
+ // and break so the finally always runs (safe for detached runs:
|
||||
+ // consumeStream drains the SDK stream independently). Listener hygiene:
|
||||
+ // all three once-listeners are removed on the first settle, so they
|
||||
+ // cannot pile up one-per-stall.
|
||||
+ const closed = await new Promise((resolve3) => {
|
||||
+ function finish(isClosed) {
|
||||
+ response.removeListener("drain", onDrain);
|
||||
+ response.removeListener("close", onClose);
|
||||
+ response.removeListener("error", onError);
|
||||
+ resolve3(isClosed);
|
||||
+ }
|
||||
+ function onDrain() {
|
||||
+ finish(false);
|
||||
+ }
|
||||
+ function onClose() {
|
||||
+ finish(true);
|
||||
+ }
|
||||
+ function onError() {
|
||||
+ finish(true);
|
||||
+ }
|
||||
+ response.once("drain", onDrain);
|
||||
+ response.once("close", onClose);
|
||||
+ response.once("error", onError);
|
||||
});
|
||||
+ if (closed) {
|
||||
+ await reader.cancel().catch(() => {
|
||||
+ });
|
||||
+ break;
|
||||
+ }
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
@@ -5047,7 +5078,9 @@ function writeToServerResponse({
|
||||
response.end();
|
||||
}
|
||||
};
|
||||
- read();
|
||||
+ read().catch((error) => {
|
||||
+ console.error("ai writeToServerResponse read() failed:", error);
|
||||
+ });
|
||||
}
|
||||
|
||||
// src/text-stream/pipe-text-stream-to-response.ts
|
||||
@@ -6578,9 +6611,19 @@ function createOutputTransformStream(output) {
|
||||
controller.enqueue({ part: chunk, partialOutput: void 0 });
|
||||
return;
|
||||
}
|
||||
@@ -23,7 +77,7 @@ index ae447a12f7823ec0a00837ee9f0eb809a610d5f8..a3402b2c2d021ef432cfa76e35d37007
|
||||
const result = await output.parsePartialOutput({ text: text2 });
|
||||
if (result !== void 0) {
|
||||
const currentJson = JSON.stringify(result.partial);
|
||||
@@ -6959,7 +6969,7 @@ var DefaultStreamTextResult = class {
|
||||
@@ -6959,7 +7002,7 @@ var DefaultStreamTextResult = class {
|
||||
})
|
||||
);
|
||||
}
|
||||
@@ -33,10 +87,64 @@ index ae447a12f7823ec0a00837ee9f0eb809a610d5f8..a3402b2c2d021ef432cfa76e35d37007
|
||||
maxRetries: maxRetriesArg,
|
||||
abortSignal
|
||||
diff --git a/dist/index.mjs b/dist/index.mjs
|
||||
index 663875332e3f9a9bd167c25583c515876f42951b..b840b0502c9894df983e0154805abb80e70e6331 100644
|
||||
index 663875332e3f9a9bd167c25583c515876f42951b..a51514a390d22811a407edd8b703e21793586cc8 100644
|
||||
--- a/dist/index.mjs
|
||||
+++ b/dist/index.mjs
|
||||
@@ -6501,9 +6501,19 @@ function createOutputTransformStream(output) {
|
||||
@@ -4957,9 +4957,40 @@ function writeToServerResponse({
|
||||
break;
|
||||
const canContinue = response.write(value);
|
||||
if (!canContinue) {
|
||||
- await new Promise((resolve3) => {
|
||||
- response.once("drain", resolve3);
|
||||
+ // PATCH(docmost #486): race "drain" against "close"/"error". The
|
||||
+ // original awaited ONLY "drain", so a client that disconnected mid-write
|
||||
+ // (the socket never drains) parked this loop FOREVER: the finally never
|
||||
+ // ran, response.end() was unreachable, and the reader + buffered chunks
|
||||
+ // were held until process restart. On close/error we cancel the reader
|
||||
+ // and break so the finally always runs (safe for detached runs:
|
||||
+ // consumeStream drains the SDK stream independently). Listener hygiene:
|
||||
+ // all three once-listeners are removed on the first settle, so they
|
||||
+ // cannot pile up one-per-stall.
|
||||
+ const closed = await new Promise((resolve3) => {
|
||||
+ function finish(isClosed) {
|
||||
+ response.removeListener("drain", onDrain);
|
||||
+ response.removeListener("close", onClose);
|
||||
+ response.removeListener("error", onError);
|
||||
+ resolve3(isClosed);
|
||||
+ }
|
||||
+ function onDrain() {
|
||||
+ finish(false);
|
||||
+ }
|
||||
+ function onClose() {
|
||||
+ finish(true);
|
||||
+ }
|
||||
+ function onError() {
|
||||
+ finish(true);
|
||||
+ }
|
||||
+ response.once("drain", onDrain);
|
||||
+ response.once("close", onClose);
|
||||
+ response.once("error", onError);
|
||||
});
|
||||
+ if (closed) {
|
||||
+ await reader.cancel().catch(() => {
|
||||
+ });
|
||||
+ break;
|
||||
+ }
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
@@ -4968,7 +4999,9 @@ function writeToServerResponse({
|
||||
response.end();
|
||||
}
|
||||
};
|
||||
- read();
|
||||
+ read().catch((error) => {
|
||||
+ console.error("ai writeToServerResponse read() failed:", error);
|
||||
+ });
|
||||
}
|
||||
|
||||
// src/text-stream/pipe-text-stream-to-response.ts
|
||||
@@ -6501,9 +6534,19 @@ function createOutputTransformStream(output) {
|
||||
controller.enqueue({ part: chunk, partialOutput: void 0 });
|
||||
return;
|
||||
}
|
||||
@@ -57,7 +165,7 @@ index 663875332e3f9a9bd167c25583c515876f42951b..b840b0502c9894df983e0154805abb80
|
||||
const result = await output.parsePartialOutput({ text: text2 });
|
||||
if (result !== void 0) {
|
||||
const currentJson = JSON.stringify(result.partial);
|
||||
@@ -6882,7 +6892,7 @@ var DefaultStreamTextResult = class {
|
||||
@@ -6882,7 +6925,7 @@ var DefaultStreamTextResult = class {
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
Generated
+21
-6
@@ -48,7 +48,7 @@ patchedDependencies:
|
||||
hash: d8dc66e5ec3b9d23a876b979f493b6aa901fd2d965be54729495da5136296a42
|
||||
path: patches/@hocuspocus__server@3.4.4.patch
|
||||
ai@6.0.134:
|
||||
hash: f60bfc3357e01e1f3978c6c40fdd65aeb33fefaad7179cde8676465b6c5ff4d9
|
||||
hash: e8c599b3963eb01b9ed1481683b7b795ce94137aa1a0c951917c20c8b870299b
|
||||
path: patches/ai@6.0.134.patch
|
||||
scimmy@1.3.5:
|
||||
hash: 775d80f86830b2c5dd1a250c9802c10f8fc3da3c7898373de5aa0c23993d1673
|
||||
@@ -287,6 +287,9 @@ importers:
|
||||
'@docmost/prosemirror-markdown':
|
||||
specifier: workspace:*
|
||||
version: link:../../packages/prosemirror-markdown
|
||||
'@docmost/token-estimate':
|
||||
specifier: workspace:*
|
||||
version: link:../../packages/token-estimate
|
||||
'@excalidraw/excalidraw':
|
||||
specifier: 0.18.0-3a5ef40
|
||||
version: 0.18.0-3a5ef40(@types/react-dom@18.3.1)(@types/react@18.3.12)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)
|
||||
@@ -561,6 +564,9 @@ importers:
|
||||
'@docmost/prosemirror-markdown':
|
||||
specifier: workspace:*
|
||||
version: link:../../packages/prosemirror-markdown
|
||||
'@docmost/token-estimate':
|
||||
specifier: workspace:*
|
||||
version: link:../../packages/token-estimate
|
||||
'@fastify/compress':
|
||||
specifier: ^9.0.0
|
||||
version: 9.0.0
|
||||
@@ -644,10 +650,10 @@ importers:
|
||||
version: 8.3.0(socket.io-adapter@2.5.4)
|
||||
ai:
|
||||
specifier: ^6.0.134
|
||||
version: 6.0.134(patch_hash=f60bfc3357e01e1f3978c6c40fdd65aeb33fefaad7179cde8676465b6c5ff4d9)(zod@4.3.6)
|
||||
version: 6.0.134(patch_hash=e8c599b3963eb01b9ed1481683b7b795ce94137aa1a0c951917c20c8b870299b)(zod@4.3.6)
|
||||
ai-sdk-ollama:
|
||||
specifier: ^3.8.1
|
||||
version: 3.8.1(ai@6.0.134(patch_hash=f60bfc3357e01e1f3978c6c40fdd65aeb33fefaad7179cde8676465b6c5ff4d9)(zod@4.3.6))(zod@4.3.6)
|
||||
version: 3.8.1(ai@6.0.134(patch_hash=e8c599b3963eb01b9ed1481683b7b795ce94137aa1a0c951917c20c8b870299b)(zod@4.3.6))(zod@4.3.6)
|
||||
bcrypt:
|
||||
specifier: ^6.0.0
|
||||
version: 6.0.0
|
||||
@@ -1148,6 +1154,15 @@ importers:
|
||||
specifier: 4.1.6
|
||||
version: 4.1.6(@opentelemetry/api@1.9.0)(@types/node@20.19.43)(@vitest/coverage-v8@4.1.6)(happy-dom@20.8.9)(jsdom@25.0.0)(vite@8.0.5(@types/node@20.19.43)(esbuild@0.28.0)(jiti@2.4.2)(less@4.2.0)(sugarss@5.0.1(postcss@8.5.14))(terser@5.39.0)(tsx@4.21.0)(yaml@2.8.3))
|
||||
|
||||
packages/token-estimate:
|
||||
devDependencies:
|
||||
typescript:
|
||||
specifier: ^5.0.0
|
||||
version: 5.9.3
|
||||
vitest:
|
||||
specifier: 4.1.6
|
||||
version: 4.1.6(@opentelemetry/api@1.9.0)(@types/node@25.5.0)(@vitest/coverage-v8@4.1.6)(happy-dom@20.8.9)(jsdom@27.4.0(@noble/hashes@2.0.1))(vite@8.0.5(@types/node@25.5.0)(esbuild@0.28.0)(jiti@2.4.2)(less@4.2.0)(sugarss@5.0.1(postcss@8.5.14))(terser@5.39.0)(tsx@4.21.0)(yaml@2.8.3))
|
||||
|
||||
packages:
|
||||
|
||||
'@aashutoshrathi/word-wrap@1.2.6':
|
||||
@@ -16455,17 +16470,17 @@ snapshots:
|
||||
|
||||
agent-base@7.1.4: {}
|
||||
|
||||
ai-sdk-ollama@3.8.1(ai@6.0.134(patch_hash=f60bfc3357e01e1f3978c6c40fdd65aeb33fefaad7179cde8676465b6c5ff4d9)(zod@4.3.6))(zod@4.3.6):
|
||||
ai-sdk-ollama@3.8.1(ai@6.0.134(patch_hash=e8c599b3963eb01b9ed1481683b7b795ce94137aa1a0c951917c20c8b870299b)(zod@4.3.6))(zod@4.3.6):
|
||||
dependencies:
|
||||
'@ai-sdk/provider': 3.0.8
|
||||
'@ai-sdk/provider-utils': 4.0.21(zod@4.3.6)
|
||||
ai: 6.0.134(patch_hash=f60bfc3357e01e1f3978c6c40fdd65aeb33fefaad7179cde8676465b6c5ff4d9)(zod@4.3.6)
|
||||
ai: 6.0.134(patch_hash=e8c599b3963eb01b9ed1481683b7b795ce94137aa1a0c951917c20c8b870299b)(zod@4.3.6)
|
||||
jsonrepair: 3.13.3
|
||||
ollama: 0.6.3
|
||||
transitivePeerDependencies:
|
||||
- zod
|
||||
|
||||
ai@6.0.134(patch_hash=f60bfc3357e01e1f3978c6c40fdd65aeb33fefaad7179cde8676465b6c5ff4d9)(zod@4.3.6):
|
||||
ai@6.0.134(patch_hash=e8c599b3963eb01b9ed1481683b7b795ce94137aa1a0c951917c20c8b870299b)(zod@4.3.6):
|
||||
dependencies:
|
||||
'@ai-sdk/gateway': 3.0.77(zod@4.3.6)
|
||||
'@ai-sdk/provider': 3.0.8
|
||||
|
||||
Reference in New Issue
Block a user