Compare commits
118 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 62482827c9 | |||
| 2a58b63e15 | |||
| 2dfc5c5144 | |||
| 878ede409b | |||
| b8185a9f0e | |||
| 108d2c6bef | |||
| 1e28151f7a | |||
| 6e42752a9c | |||
| 8430eea700 | |||
| 5563811aa8 | |||
| 55a1c279fc | |||
| d4b21957dc | |||
| e71212cdc6 | |||
| 763a1bf7bd | |||
| bfb6a52eea | |||
| 0503e8b4b1 | |||
| b97cad0ebe | |||
| 080dd82023 | |||
| 803bbd40b4 | |||
| c7a38e274e | |||
| 791a707eb6 | |||
| 60205398bb | |||
| 1685b32333 | |||
| 4b78e336a2 | |||
| 77e20ecb41 | |||
| d533daa45f | |||
| ba58493909 | |||
| 947796adef | |||
| 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 | |||
| fe5bd159c4 | |||
| f12b685698 | |||
| f6fc914c95 | |||
| 1d89cc2058 | |||
| 70a9e2a9cb | |||
| 5d8083f8ff | |||
| 3e945305c8 | |||
| 6d7dba970c | |||
| 512bcba5f3 | |||
| 5c1ab9c7b5 | |||
| 14d7b21df0 | |||
| 363f20ab75 | |||
| 3411bda2d1 | |||
| e670f7498a | |||
| 9a8671c3af | |||
| bec2156e96 | |||
| acc705de19 | |||
| a49872444e | |||
| ffc38ea2ca | |||
| 2a951df096 | |||
| 5dc7a2703f | |||
| bfb4c8d8d0 | |||
| f794ac6d6c | |||
| e4e788f151 | |||
| 4be4a75fa3 | |||
| e6171a1810 | |||
| 2f23cf4b65 | |||
| d269bd9efe | |||
| d287c15db4 | |||
| 7cb3199d09 | |||
| e19275e96e | |||
| 3a521ada4d | |||
| 576db3c8f9 | |||
| ddb37376a4 | |||
| dc58974b31 | |||
| c917dcc3c1 | |||
| eddc3b5c33 | |||
| e454fe189c | |||
| ea99d4fe63 | |||
| 23966ce51c | |||
| a53b2f454e | |||
| d219eb7525 | |||
| 93d244478e | |||
| 791f709c18 | |||
| f8a27cba91 | |||
| 61dc9b50c1 | |||
| cebb1cca87 | |||
| 605c0f3dda | |||
| 0f5f048ca2 | |||
| 4cb762b039 | |||
| d0f99052cf | |||
| 8c74659d91 | |||
| e24ddf6b3e | |||
| 3cba551800 | |||
| 15a9eba562 | |||
| 2c03fefa9d | |||
| 76af4f692e | |||
| 4809348457 | |||
| d3d32d637b | |||
| 9a435201b8 | |||
| d6827b9210 |
+110
-3
@@ -124,6 +124,40 @@ MCP_DOCMOST_PASSWORD=
|
||||
# MCP_TOKEN=
|
||||
# MCP_SESSION_IDLE_MS=1800000
|
||||
#
|
||||
# --- MCP collaboration write path: concurrency + rights-staleness (#449) ------
|
||||
# MCP content writes (update_page, insert/replace nodes, comments-in-body, etc.)
|
||||
# go over the collaboration websocket and are serialized PER PAGE by an
|
||||
# in-process mutex (a module-level Map, one promise-chain per page UUID). This
|
||||
# guarantees no two MCP writes on the SAME page overlap and clobber each other.
|
||||
#
|
||||
# DEPLOY REQUIREMENT — SINGLE INSTANCE or STICKY SESSIONS. The mutex is
|
||||
# process-local. Behind a multi-replica load balancer WITHOUT sticky sessions,
|
||||
# two replicas can each "hold" the lock for the same page at the same time and
|
||||
# serialization is silently lost (concurrent full-document writes race on the
|
||||
# live Yjs fragment). Run the MCP/app as a SINGLE instance, OR pin a page's
|
||||
# traffic to one replica (sticky sessions / consistent hashing on page id). The
|
||||
# same constraint applies to the RAM-only stash_page blob store above. There is
|
||||
# deliberately no cross-process (e.g. Postgres advisory) lock yet — this is a
|
||||
# CONSCIOUS documented constraint, not an oversight (#449).
|
||||
#
|
||||
# To reduce connect-storms the write path caches ONE live collab session per
|
||||
# (wsUrl, page, token). Tunables (all optional; defaults are safe):
|
||||
# MCP_COLLAB_SESSION_IDLE_MS=60000 # idle TTL, reset per op; 0 disables cache
|
||||
# MCP_COLLAB_SESSION_MAX_ENTRIES=32 # LRU cap on cached sessions
|
||||
# MCP_COLLAB_TOKEN_TTL_MS=300000 # per-client collab-token cache (5 min)
|
||||
#
|
||||
# RIGHTS-STALENESS TRADE-OFF. A cached collab session writes under the token
|
||||
# captured at CONNECT time, and the collab-token cache reuses a token for its TTL.
|
||||
# So if a user's access to a page is REVOKED, MCP writes on an already-open
|
||||
# session may keep succeeding until the session ages out. MCP_COLLAB_SESSION_MAX_AGE_MS
|
||||
# is the HARD lifetime (checked at each acquire) that BOUNDS this window: after it,
|
||||
# the session is torn down and the next write re-auths with a fresh token, picking
|
||||
# up the revocation. Default 10 min. LOWER it to shorten the revocation lag at the
|
||||
# cost of more reconnects; RAISE it to reduce reconnects at the cost of a longer
|
||||
# stale-rights window. There is intentionally no push-based cache invalidation on
|
||||
# a rights change — this bounded window is the accepted trade-off (#449).
|
||||
# MCP_COLLAB_SESSION_MAX_AGE_MS=600000
|
||||
#
|
||||
# BLOB SANDBOX (stash_page). An in-RAM, process-local store that hands large page
|
||||
# content + images to an external consumer WITHOUT bloating the model context or
|
||||
# requiring Docmost auth. The stash_page tool serializes a page, mirrors its
|
||||
@@ -191,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.
|
||||
@@ -217,6 +266,17 @@ MCP_DOCMOST_PASSWORD=
|
||||
# active" behavior.
|
||||
# AI_CHAT_DEFERRED_TOOLS=true
|
||||
|
||||
# Final-step lockdown for the in-app agent loop (#444). Default OFF. When ON
|
||||
# (legacy), the LAST allowed step forces a text-only answer: the model's tools are
|
||||
# stripped (toolChoice=none) and a synthesis instruction is appended. That
|
||||
# tool-stripping caused a token-degeneration incident — robbed of its tools on the
|
||||
# final step mid-work, the model emitted a ~255KB block repeating a single token —
|
||||
# so the default is now OFF: the last step keeps its tools and gets only a SOFT
|
||||
# nudge to finish with a text summary, and a token-degeneration detector is the
|
||||
# universal anti-babble guard. Enable this ONLY for a model that reliably ends its
|
||||
# turns with a clear text answer.
|
||||
# AI_CHAT_FINAL_STEP_LOCKDOWN=false
|
||||
|
||||
# --- Autonomous / detached agent runs (settings.ai.autonomousRuns) ---
|
||||
# Opt-in per workspace (AI settings; off by default). When on, a chat turn becomes
|
||||
# a server-side RUN that survives a browser disconnect — only an explicit Stop ends
|
||||
@@ -242,6 +302,39 @@ MCP_DOCMOST_PASSWORD=
|
||||
# enabled for a workspace, and the same single-instance constraint applies (the
|
||||
# registry is process-local).
|
||||
# AI_CHAT_RESUMABLE_STREAM=false
|
||||
#
|
||||
# Per-run replay ring cap (#491), in BYTES, for the resumable-stream registry
|
||||
# above. The registry buffers the run's recent SSE tail so a reopened tab can
|
||||
# attach and continue from the step it already persisted; the ring is bounded and
|
||||
# rotates on every confirmed step-persist. This caps the un-persisted tail between
|
||||
# rotations — an overflow evicts the oldest frames and a late attach falls back to
|
||||
# 204 -> degraded poll, so correctness never depends on the size. Default 4194304
|
||||
# (4MB); a 0/invalid value falls back to the default. The per-subscriber backpressure
|
||||
# cap is derived as 2x this value. Only meaningful with AI_CHAT_RESUMABLE_STREAM on.
|
||||
# AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES=4194304
|
||||
|
||||
# --- 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).
|
||||
@@ -290,6 +383,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
|
||||
|
||||
@@ -157,6 +157,12 @@ jobs:
|
||||
- name: Build prosemirror-markdown
|
||||
run: pnpm --filter @docmost/prosemirror-markdown build
|
||||
|
||||
# docmost-client.loader.ts type-imports from @docmost/mcp (issue #446); its
|
||||
# build/ is gitignored and `test:e2e` type-checks, so build it here or tsc
|
||||
# fails with TS2307 (mirrors the e2e-mcp / mcp-server-parity jobs).
|
||||
- name: Build mcp
|
||||
run: pnpm --filter @docmost/mcp build
|
||||
|
||||
- name: Run migrations
|
||||
run: pnpm --filter ./apps/server migration:latest
|
||||
|
||||
|
||||
@@ -2,6 +2,11 @@
|
||||
.env.dev
|
||||
.env.prod
|
||||
data
|
||||
# Exception: the committed draw.io shape catalog (issue #424) lives in a `data/`
|
||||
# dir, but the bare `data` ignore above is meant for runtime state, not this
|
||||
# bundled build asset. Re-include the directory and its contents.
|
||||
!packages/mcp/data/
|
||||
!packages/mcp/data/**
|
||||
# compiled output
|
||||
/dist
|
||||
node_modules
|
||||
@@ -24,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
|
||||
|
||||
@@ -5,6 +5,139 @@ repository. It has two layers: **how to run a task end-to-end** (the
|
||||
sections below), and **how the codebase is built** (the technical sections
|
||||
further down, formerly in `CLAUDE.md`).
|
||||
|
||||
## ARCHITECTURAL INVARIANTS — NON-NEGOTIABLE
|
||||
|
||||
THE TEN RULES BELOW ARE HARD CONSTRAINTS. Each one was paid for with a real
|
||||
production incident or a multi-PR bug chain in THIS repository (cited inline).
|
||||
They override convenience, deadlines and "it's just a small feature". A PR that
|
||||
violates any of them MUST be rejected in review regardless of how good the rest
|
||||
of it is. If a task genuinely seems to require breaking one — STOP and raise it
|
||||
with the owner; do not code around it.
|
||||
|
||||
### 1. EVERY BUFFER, CACHE, HISTORY AND PAYLOAD HAS AN EXPLICIT SIZE BUDGET
|
||||
|
||||
Nothing accumulates unboundedly. A row/item cap is NOT a byte cap. Anything
|
||||
replayed to a model, buffered in memory, persisted per step, or refetched by a
|
||||
poll must state its budget in bytes/tokens and enforce it. Rewriting a growing
|
||||
structure in full on every increment is FORBIDDEN — append or diff instead;
|
||||
O(n²) write/serialize patterns do not pass review.
|
||||
(Paid for by: full-row rewrite on every agent step — hundreds of MB of Postgres
|
||||
writes per 50-step run, with every tool output serialized twice; unbounded
|
||||
history replay killing long chats on the provider context window; 32 MB replay
|
||||
buffers per active run.)
|
||||
|
||||
### 2. EVERYTHING LONG-RUNNING TERMINATES BY CONSTRUCTION
|
||||
|
||||
Every run / row / session / lease / subscriber / queue entry must define AT
|
||||
DESIGN TIME: its owner; every terminal state; who writes the terminal state on
|
||||
EVERY path (success, error, abort, disconnect in each phase, process restart);
|
||||
retries for the terminal write; and a periodic sweeper that does not depend on
|
||||
a reboot. A best-effort terminal write with no retry and no sweep is FORBIDDEN.
|
||||
(Paid for by: assistant rows stuck 'streaming' forever; runs stuck 'running'
|
||||
409-locking their chat until a restart — the #183/#184 follow-up chain.)
|
||||
|
||||
### 3. EVERY AWAIT IS CANCELLABLE AND DEADLINED; NEVER BLOCK THE EVENT LOOP
|
||||
|
||||
Every async step inside a request or agent turn honors the turn's AbortSignal
|
||||
AND a wall-clock deadline — including in-app tools, lock queues and pagination
|
||||
loops, not just external calls. Synchronous CPU work beyond ~50 ms goes to a
|
||||
worker_thread. Promise.race DOES NOT cancel synchronous work — using it as a
|
||||
"timeout" for sync computation is forbidden (the timer only fires after the
|
||||
event loop is free again, i.e. after the damage is done).
|
||||
(Paid for by: in-app tools ignoring abortSignal and writing pages AFTER Stop;
|
||||
the synchronous ELK layout freezing every SSE stream in the process; the
|
||||
step-0 MCP handshake hang — #397.)
|
||||
|
||||
### 4. ONE SOURCE OF TRUTH; EVERYTHING ELSE IS A REBUILDABLE CACHE
|
||||
|
||||
Postgres is the authoritative state. Every in-memory structure (registries,
|
||||
caches, client stores) must be reconstructible from the DB and treated as
|
||||
lossy. The client renders SERVER-DECLARED state — "a run is active" is a server
|
||||
fact delivered as data, never inferred from side signals (204 vs 2xx, the
|
||||
flavor of a disconnect). A new feature must name the owner of each piece of
|
||||
state before implementation starts.
|
||||
(Paid for by: the strip/restore resume machinery, silently frozen UIs and
|
||||
ghost sends after unmount — the #381→#432→#456 chain.)
|
||||
|
||||
### 5. STATE MACHINES ARE EXPLICIT — ONE-SHOT FLAGS ARE FORBIDDEN
|
||||
|
||||
A complex lifecycle (chat thread, resume/reconnect, run) lives in a named-state
|
||||
automaton (reducer / enum) where every state has an owner and a rendered
|
||||
representation — including the failure states. Adding a boolean ref that one
|
||||
callback arms and another reads-and-clears is FORBIDDEN in the AI-chat client.
|
||||
New behavior = a new named state + explicit transitions, and the interruption
|
||||
matrix (disconnect in each phase × restart × stop × supersede) is enumerated at
|
||||
design time, not discovered one incident at a time.
|
||||
(Paid for by: 26 one-shot useRef flags in chat-thread.tsx and the drip of
|
||||
"one more missing transition" across #381→#386/#389→#432→#456.)
|
||||
|
||||
### 6. NO NEW MODE FORKS; A FLAG IS FOR ROLLOUT, THEN IT DIES
|
||||
|
||||
A behavior flag that forks a code path must ship with a written sunset
|
||||
condition; stacking a new flag onto the existing matrix without deleting or
|
||||
scheduling an old one is forbidden. While a temporary fork exists, BOTH sides
|
||||
must share identical lifecycle handling (abort semantics, error listeners,
|
||||
concurrency gates) — asymmetric forks are outlawed.
|
||||
(Paid for by: legacy vs autonomous divergence — the one-active-run gate and
|
||||
the socket 'error' listener each existing on only ONE side; 2^4 flag
|
||||
combinations each with different abort semantics.)
|
||||
|
||||
### 7. NO HAND-SYNCED MIRRORS — CODEGEN OR A CI PARITY TEST, NOTHING LESS
|
||||
|
||||
Two copies of the same knowledge (schema, tool registry, glyph map, probe
|
||||
body, hash/normalize algorithm, label list) require either generation from a
|
||||
single source or a CI test that FAILS on drift. A "mirror this change over
|
||||
there" comment is NOT a guard and does not pass review.
|
||||
(Paid for by: #293 — three drifting converter copies losing data; #447 —
|
||||
REGISTRY_STAMP covering only one of the mirrored files; ~10 still-unguarded
|
||||
mirrors across the MCP layer.)
|
||||
|
||||
### 8. CACHES, HEADERS, BUFFERS AND FSM TRANSITIONS GET AN INTEGRATION TEST OF THE OBSERVABLE PROPERTY
|
||||
|
||||
A unit test of a pure helper DOES NOT COUNT for these. Test the real header on
|
||||
the real HTTP response, the real cache hit under real token sources, the real
|
||||
transition under a really-killed socket. If the observable property cannot be
|
||||
tested, the design is wrong — fix the design, not the test.
|
||||
(Paid for by: #431→#439 — a cache keyed on a fresh-per-call JWT, so it NEVER
|
||||
hit and became prod incident #435 while its unit tests stayed green; and by
|
||||
the #352→#455 immutable-cache header silently overwritten by a framework
|
||||
default AFTER the unit-tested code ran.)
|
||||
|
||||
### 9. CLIENT INPUT IS HOSTILE UNTIL VALIDATED — ALSO BEFORE PERSISTENCE
|
||||
|
||||
Anything from the browser (message parts, ids, titles, selections, flags) is
|
||||
validated/sanitized BEFORE it is persisted into a row that will later be
|
||||
replayed into a prompt, a converter or another subsystem. A poisoned row must
|
||||
never be able to permanently brick a chat or a page on every subsequent read.
|
||||
(Paid for by: unvalidated UIMessage parts persisted verbatim — one bad row
|
||||
500s the chat on every later turn; #159 client-spoofed page titles; #388
|
||||
selection re-sanitized server-side for the same reason.)
|
||||
|
||||
### 10. FAILURES ARE LOUD AND SPECIFIC; SILENT DEGRADATION IS FORBIDDEN
|
||||
|
||||
Extends the error convention below: a fire-and-forget write is allowed ONLY
|
||||
with a metric or a greppable ERROR log; a degraded mode (dead cached MCP
|
||||
client, stopped poll, exhausted retries, evicted buffer) must be VISIBLE to
|
||||
the user or the operator. A feature that can quietly stop working — a frozen
|
||||
"streaming…" UI, a poll that silently gives up, a cache serving corpses — does
|
||||
not pass review.
|
||||
(Paid for by: the degraded poll's silent 10-minute death leaving a forever-
|
||||
"streaming" answer; dead MCP clients served from cache while every external
|
||||
tool call failed; #435 being caught in minutes ONLY because metrics — #403 —
|
||||
existed.)
|
||||
|
||||
## Default skill for feature design
|
||||
|
||||
For any feature-design request — the user hands over a raw feature idea, asks
|
||||
to design or think through a feature, or to draft an issue («спроектируй»,
|
||||
«продумай фичу», «составь ишью», "design X", "write an issue for X") — invoke
|
||||
the `orchestrator-feature-designer` skill (Skill tool) BEFORE any other work.
|
||||
It is the default operating mode for design work in this repository: research
|
||||
→ design checklist (R1–R10) → forks resolved with the human → adversarial
|
||||
self-attack → filed PR-sized issues. Do not design features or write issues
|
||||
ad-hoc while this skill is available. This does not apply to non-design work
|
||||
(bug fixes, reviews, retrospectives, refactors already specified by an issue).
|
||||
|
||||
## Task lifecycle
|
||||
|
||||
### 1. Start: sync with develop
|
||||
@@ -201,7 +334,7 @@ pnpm workspace (`pnpm@10.4.0`) orchestrated by **Nx**. Four workspace packages:
|
||||
| `apps/client` | `client` | React 18 + Vite + Mantine 8 + TanStack Query + Jotai | SPA frontend |
|
||||
| `packages/editor-ext` | `@docmost/editor-ext` | Tiptap/ProseMirror | Shared Tiptap node/mark extensions, imported by both the client and the server |
|
||||
| `packages/mcp` | `@docmost/mcp` | MCP SDK, Tiptap, Yjs | Standalone MCP server, also bundled into the server at `/mcp`. Consumes the shared converter/schema from `@docmost/prosemirror-markdown` (#293) — it no longer carries its own vendored converter/schema copy |
|
||||
| `packages/prosemirror-markdown` | `@docmost/prosemirror-markdown` | Tiptap, marked, jsdom | The single, canonical ProseMirror↔Markdown converter + Docmost schema mirror (#293). Consumed by `mcp`, `git-sync`, AND `apps/server` (server-side markdown import/export, #345); there is exactly ONE copy of the converter now |
|
||||
| `packages/prosemirror-markdown` | `@docmost/prosemirror-markdown` | Tiptap, marked; jsdom (Node only) | The single, canonical ProseMirror↔Markdown converter + Docmost schema mirror (#293). Consumed by `mcp`, `git-sync`, `apps/server` (server-side markdown import/export, #345), AND `apps/client` (markdown paste/copy + AI-chat render, via the `browser` entry — native `DOMParser`, no jsdom in the client bundle, #347); there is exactly ONE copy of the converter now |
|
||||
|
||||
`build` targets are Nx-cached and dependency-ordered (`dependsOn: ["^build"]`), so `editor-ext` builds before the apps. `nx.json` sets `affected.defaultBase: main`.
|
||||
|
||||
@@ -322,12 +455,12 @@ 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:
|
||||
- **TanStack Query** for server state (one `queries/` file per feature), **Jotai** atoms for local/shared UI state, **Mantine 8** + CSS modules (`*.module.css`) + `postcss-preset-mantine` for UI.
|
||||
- The editor is Tiptap; shared node/mark extensions live in `packages/editor-ext` and are imported by **both the client and the server** (collaboration, schema, `canonicalizeFootnotes`) — editor schema changes often need to be made in `editor-ext`, not just the client. Server-side markdown import/export no longer lives in `editor-ext`: it goes through the canonical converter (#345, see below). The ProseMirror↔Markdown converter and its Docmost schema mirror now live in a SINGLE package, `@docmost/prosemirror-markdown` (#293), consumed by `mcp`, `git-sync`, and `apps/server` (#345) — do NOT reintroduce a per-package copy. `editor-ext` is the upstream source of the Tiptap schema; the package's `docmost-schema.ts` mirrors it and a serializer-contract test (`packages/prosemirror-markdown/test/serializer-contract.test.ts`) guards the boundary (every schema node must have a converter case), so a drift surfaces as a failing test rather than silent divergence. For the converter's property-testing and counterexample→fixture process (P1–P4 invariants, the `PROPERTY_SEED`/`PROPERTY_NUM_RUNS` knobs, and the nightly fuzz workflow), see `packages/prosemirror-markdown/README.md`.
|
||||
- The editor is Tiptap; shared node/mark extensions live in `packages/editor-ext` and are imported by **both the client and the server** (collaboration, schema, `canonicalizeFootnotes`) — editor schema changes often need to be made in `editor-ext`, not just the client. Server-side markdown import/export no longer lives in `editor-ext`: it goes through the canonical converter (#345, see below). The ProseMirror↔Markdown converter and its Docmost schema mirror now live in a SINGLE package, `@docmost/prosemirror-markdown` (#293), consumed by `mcp`, `git-sync`, `apps/server` (#345), and `apps/client` (#347) — do NOT reintroduce a per-package copy. The client uses the package's `browser` entry (`@docmost/prosemirror-markdown/browser`): markdown paste (`markdown-clipboard.ts`), copy-as-markdown, and AI-chat rendering now all go through the canonical converter, so the hand-written `marked`/`turndown` markdown layer that used to live in `editor-ext` was deleted (#347). The browser entry runs the HTML→DOM stage on the native `DOMParser`, so jsdom stays out of the client bundle. `editor-ext` is the upstream source of the Tiptap schema; the package's `docmost-schema.ts` mirrors it and a serializer-contract test (`packages/prosemirror-markdown/test/serializer-contract.test.ts`) guards the boundary (every schema node must have a converter case), so a drift surfaces as a failing test rather than silent divergence. For the converter's property-testing and counterexample→fixture process (P1–P4 invariants, the `PROPERTY_SEED`/`PROPERTY_NUM_RUNS` knobs, and the nightly fuzz workflow), see `packages/prosemirror-markdown/README.md`.
|
||||
- API access goes through `apps/client/src/lib/api-client.ts` (axios). The `@` alias maps to `apps/client/src`.
|
||||
- Runtime config is injected at build time by `vite.config.ts` via `define` (`APP_URL`, `COLLAB_URL`, `APP_VERSION`, …) — these come from the root `.env`, not from `import.meta.env`.
|
||||
|
||||
@@ -337,8 +470,8 @@ 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`.
|
||||
- **Adding/renaming/removing an MCP tool requires updating `SERVER_INSTRUCTIONS`** in `packages/mcp/src/index.ts` — the intent-routing guide MCP clients receive on initialize. This applies both to inline `server.registerTool(...)` calls in `index.ts` and to specs in `packages/mcp/src/tool-specs.ts`. Enforced by `packages/mcp/test/unit/server-instructions.test.mjs`, which fails when a registered tool is not mentioned in the guide (deliberate opt-outs go into its `EXCEPTIONS` list). `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.
|
||||
- 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
|
||||
|
||||
|
||||
+256
@@ -10,6 +10,123 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### Breaking Changes
|
||||
|
||||
- **External MCP tool names are now camelCase (all renamed).** Every tool on the
|
||||
external `/mcp` surface was renamed from `snake_case` to `camelCase`, so the
|
||||
external MCP name now matches the in-app tool name exactly (one logical tool,
|
||||
one name everywhere). For example `get_node` → `getNode`, `edit_page_text` →
|
||||
`editPageText`, `patch_node` → `patchNode`. The tools' behaviour, inputs and
|
||||
outputs are unchanged — only the names change. The single-word `search`
|
||||
keeps its name.
|
||||
|
||||
*Migration (external MCP clients only — the in-app AI agent already used these
|
||||
names and is unaffected):* update anything that refers to a tool by its
|
||||
string name — permission allowlists (`mcp__gitmost-*__get_node` →
|
||||
`mcp__gitmost-*__getNode`), saved prompts/skills, `.mcp.json` tool filters,
|
||||
and metrics dashboards that group by the `tool` label — and roll it out in
|
||||
lockstep with this deploy, because the old snake_case names stop resolving.
|
||||
Released together with the `import_page_markdown`/`update_page_markdown`
|
||||
change below so external configs break exactly once.
|
||||
|
||||
Full mapping (old → new):
|
||||
|
||||
| Old (snake_case) | New (camelCase) |
|
||||
| --- | --- |
|
||||
| `check_new_comments` | `checkNewComments` |
|
||||
| `copy_page_content` | `copyPageContent` |
|
||||
| `create_comment` | `createComment` |
|
||||
| `create_page` | `createPage` |
|
||||
| `delete_comment` | `deleteComment` |
|
||||
| `delete_node` | `deleteNode` |
|
||||
| `delete_page` | `deletePage` |
|
||||
| `diff_page_versions` | `diffPageVersions` |
|
||||
| `docmost_transform` | `docmostTransform` |
|
||||
| `drawio_create` | `drawioCreate` |
|
||||
| `drawio_get` | `drawioGet` |
|
||||
| `drawio_guide` | `drawioGuide` |
|
||||
| `drawio_shapes` | `drawioShapes` |
|
||||
| `drawio_update` | `drawioUpdate` |
|
||||
| `edit_page_text` | `editPageText` |
|
||||
| `export_page_markdown` | `exportPageMarkdown` |
|
||||
| `get_node` | `getNode` |
|
||||
| `get_outline` | `getOutline` |
|
||||
| `get_page` | `getPage` |
|
||||
| `get_page_json` | `getPageJson` |
|
||||
| `get_workspace` | `getWorkspace` |
|
||||
| `insert_footnote` | `insertFootnote` |
|
||||
| `insert_image` | `insertImage` |
|
||||
| `insert_node` | `insertNode` |
|
||||
| `list_comments` | `listComments` |
|
||||
| `list_page_history` | `listPageHistory` |
|
||||
| `list_pages` | `listPages` |
|
||||
| `list_shares` | `listShares` |
|
||||
| `list_spaces` | `listSpaces` |
|
||||
| `move_page` | `movePage` |
|
||||
| `patch_node` | `patchNode` |
|
||||
| `rename_page` | `renamePage` |
|
||||
| `replace_image` | `replaceImage` |
|
||||
| `resolve_comment` | `resolveComment` |
|
||||
| `restore_page_version` | `restorePageVersion` |
|
||||
| `search` | `search` (unchanged) |
|
||||
| `search_in_page` | `searchInPage` |
|
||||
| `share_page` | `sharePage` |
|
||||
| `stash_page` | `stashPage` |
|
||||
| `table_delete_row` | `tableDeleteRow` |
|
||||
| `table_get` | `tableGet` |
|
||||
| `table_insert_row` | `tableInsertRow` |
|
||||
| `table_update_cell` | `tableUpdateCell` |
|
||||
| `unshare_page` | `unsharePage` |
|
||||
| `update_comment` | `updateComment` |
|
||||
| `update_page_json` | `updatePageJson` |
|
||||
| `update_page_markdown` | `updatePageMarkdown` |
|
||||
|
||||
(#412)
|
||||
|
||||
- **External MCP: `import_page_markdown` removed, `update_page_markdown` added.**
|
||||
The external `/mcp` surface no longer exposes `importPageMarkdown` (the
|
||||
round-trip parser for a self-contained *exported* Docmost-Markdown file). In
|
||||
its place it now exposes **`updatePageMarkdown`** — a plain-Markdown
|
||||
full-body replace (`{pageId, content, title?}`) that pairs with
|
||||
`updatePageJson`, re-imports the whole body (block ids regenerate) and
|
||||
parses Docmost-flavoured markdown including `^[...]` inline footnotes.
|
||||
*Migration:* MCP clients that called `importPageMarkdown` to overwrite a
|
||||
page's body from Markdown should call `updatePageMarkdown` instead (pass the
|
||||
markdown as `content`). Round-tripping an exported Docmost-Markdown file with
|
||||
comment anchors/diagrams is no longer available on the external MCP surface;
|
||||
export remains via `exportPageMarkdown`. The in-app AI agent is unaffected —
|
||||
it keeps both `importPageMarkdown` and the renamed `updatePageMarkdown` (was
|
||||
`updatePageContent`). The total MCP tool count is unchanged (−1 / +1). The
|
||||
external names shown here are the post-#412 camelCase names. (#411)
|
||||
|
||||
- **`getNode` now returns Markdown by default (was ProseMirror JSON).** The
|
||||
block-level read/write tools default to Markdown so a block round trip is
|
||||
`getNode` (markdown) → edit → `patchNode` (markdown). `getNode` now returns
|
||||
`{ …, format: "markdown", markdown }` unless you pass `format: "json"` (which
|
||||
restores the previous `{ …, node }` ProseMirror subtree); comment anchors —
|
||||
including resolved ones — are preserved in the markdown so a write-back never
|
||||
orphans a thread, and a node that cannot be a document top-level block
|
||||
(`tableRow`/`tableCell`/`tableHeader` addressed via `#<index>`) auto-falls back
|
||||
to JSON with `format: "json"` in the response. `patchNode`/`insertNode` gain a
|
||||
`markdown` input alongside `node` (provide exactly one): the markdown fragment
|
||||
may rewrite/insert several blocks at once and supports `^[...]` footnotes.
|
||||
*Migration (external MCP clients only):* a client that consumed `getNode`'s
|
||||
`node` field must now either read `markdown`, or pass `format: "json"` to keep
|
||||
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
|
||||
@@ -85,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
|
||||
@@ -146,9 +274,48 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
by physical key position and matched against the commands; genuine Cyrillic
|
||||
search terms keep priority over remapped candidates, and short wrong-layout
|
||||
prefixes match by command title. (#283, #285, #287)
|
||||
- **Opt-in substring "lookup" search mode for agents.** `/api/search` gains an
|
||||
additive, opt-in mode (guarded by a new `substring` flag) that matches literal
|
||||
substrings of page titles and body text — so technical tokens the full-text
|
||||
tokenizer mangles (`backup-srv.local`, `10.0.12.5`, `WB-MGE-30D86B`) are found
|
||||
even when the FTS query is empty. It returns a location `path`, a windowed
|
||||
`snippet` and a per-response relevance `score`, supports `titleOnly` and a
|
||||
`parentPageId` subtree scope, and applies the page-level permission filter
|
||||
before the limit. The web UI never sets `substring`, so its full-text search
|
||||
behaviour is byte-for-byte unchanged. The leading-wildcard `LIKE` predicates
|
||||
are backed by GIN trigram indexes on `LOWER(f_unaccent(title))` and
|
||||
`LOWER(f_unaccent(text_content))` so lookups use a bitmap index scan instead of
|
||||
a sequential scan. (#443)
|
||||
- **MCP `search` tool returns richer, agent-oriented results.** The external MCP
|
||||
`search` response shape changes for the agent surface: each hit now carries
|
||||
`pageId` (renamed from `id`), plus `path`, `snippet` and `score`; the
|
||||
UI-oriented `spaceId`, `rank` and `highlight` fields are dropped. (#443)
|
||||
|
||||
### 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
|
||||
`@docmost/prosemirror-markdown` (via its new `browser` entry — native
|
||||
`DOMParser`, no jsdom in the client bundle) instead of the hand-written
|
||||
`marked`/`turndown` markdown layer in `editor-ext`, which was **deleted**. As a
|
||||
result, pasting canonical markdown (`^[…]` footnotes, `<!--img …-->`,
|
||||
`> [!type]` callouts, `$…$` math, `==…==` highlight, standalone `<!--subpages-->`
|
||||
comments) now produces the SAME nodes the server import produces for the same
|
||||
text. Chat/reasoning markdown now renders through the editor schema (list items
|
||||
are wrapped in `<p>`; CSS keeps them tight). (#347)
|
||||
|
||||
- **Enabling a public share no longer auto-shares the whole sub-tree.** Turning
|
||||
a page "Shared to web" now defaults to the page alone; descendant pages become
|
||||
public only when you explicitly turn on the dedicated "Include sub-pages"
|
||||
@@ -169,6 +336,44 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
### Fixed
|
||||
|
||||
- **A long AI chat no longer bricks on the model's context window, and each turn
|
||||
stops re-persisting the whole tool-output history.** Tool outputs are now
|
||||
stored ONCE, in `metadata.parts`; the `tool_calls` trace keeps only per-step
|
||||
outcome flags (a v2 trace shape), ending the O(N²) write amplification that
|
||||
re-wrote every prior output on every step (measured on a live Postgres via the
|
||||
`pg_current_wal_lsn()` delta: the trace column shrank ~3200×, the full
|
||||
assistant row ~51%). The persisted record is unchanged in content — the full
|
||||
history still lives in `metadata.parts`. At REPLAY time only, the history sent
|
||||
to the provider is now bounded by a deterministic, prompt-cache-friendly token
|
||||
budget: `floor(0.7 × chatContextWindow)` when a window is configured (no cap —
|
||||
anti-brick protection, not a cost limiter), a flat 100k fallback for installs
|
||||
with no window set (exactly the ones that hit terminal overflow), or off when
|
||||
the window is explicitly `0`. Trimming truncates old tool outputs first, then
|
||||
mechanically collapses the oldest turns, always keeping the recent turns full
|
||||
and the tool-call/result pairing balanced. A provider context-overflow 400 is
|
||||
now classified and used as a reactive signal: the row is stamped so the NEXT
|
||||
turn re-trims aggressively (0.5×), which un-bricks a chat that just 400'd. The
|
||||
client token badge and the server budgeter now share one estimator (new
|
||||
`@docmost/token-estimate` package) so they can never diverge. Deferred-tool
|
||||
activation is also cached in the chat metadata to avoid re-resolving it each
|
||||
turn. (#490)
|
||||
- **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
|
||||
@@ -177,6 +382,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
|
||||
@@ -253,6 +491,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
|
||||
|
||||
+15
@@ -45,6 +45,11 @@ COPY --from=builder /app/packages/editor-ext/dist /app/packages/editor-ext/dist
|
||||
COPY --from=builder /app/packages/editor-ext/package.json /app/packages/editor-ext/package.json
|
||||
COPY --from=builder /app/packages/mcp/build /app/packages/mcp/build
|
||||
COPY --from=builder /app/packages/mcp/package.json /app/packages/mcp/package.json
|
||||
# The mcp package reads its data files (drawio-presets.json, drawio-shape-index.json.gz)
|
||||
# at runtime via `new URL("../../data/…", import.meta.url)` relative to build/lib/*.js,
|
||||
# i.e. from packages/mcp/data/. tsc emits only build/, so ship data/ explicitly or
|
||||
# drawioFromGraph and the shape catalog die with ENOENT on packages/mcp/data/*.
|
||||
COPY --from=builder /app/packages/mcp/data /app/packages/mcp/data
|
||||
# mcp now depends on @docmost/prosemirror-markdown (workspace:*) and eager-imports
|
||||
# it at runtime (the in-app ai-chat DocmostClient loads build/index.js -> lib/
|
||||
# markdown-converter.js). Ship the built package + its manifest, or the prod
|
||||
@@ -81,4 +86,14 @@ VOLUME ["/app/data/storage"]
|
||||
|
||||
EXPOSE 3000
|
||||
|
||||
# DEPLOY REQUIREMENT — SINGLE INSTANCE or STICKY SESSIONS (#449).
|
||||
# MCP content writes are serialized per page by an IN-PROCESS mutex, and the
|
||||
# stash_page blob store + cached collab sessions are RAM-only and process-local.
|
||||
# Running MULTIPLE replicas of this image behind a load balancer WITHOUT sticky
|
||||
# sessions silently breaks per-page write serialization (two replicas can lock
|
||||
# the same page at once) and makes stash_page blobs unreachable across replicas.
|
||||
# Run a SINGLE instance, or pin each page's traffic to one replica (sticky
|
||||
# sessions / consistent hashing on page id). There is deliberately no
|
||||
# cross-process lock yet — a conscious constraint. See .env.example (the "MCP
|
||||
# collaboration write path" block) and packages/mcp/README.md for details.
|
||||
CMD ["pnpm", "start"]
|
||||
|
||||
@@ -21,6 +21,8 @@
|
||||
"@atlaskit/pragmatic-drag-and-drop-live-region": "1.3.4",
|
||||
"@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",
|
||||
|
||||
@@ -58,8 +58,11 @@ import ConversationList from "@/features/ai-chat/components/conversation-list.ts
|
||||
import ChatThread from "@/features/ai-chat/components/chat-thread.tsx";
|
||||
import {
|
||||
exportAiChat,
|
||||
getAiChatMessagesDelta,
|
||||
stopRun,
|
||||
} from "@/features/ai-chat/services/ai-chat-service.ts";
|
||||
import { mergeDeltaRowsIntoPages } from "@/features/ai-chat/utils/resume-helpers.ts";
|
||||
import type { IAiChatMessageRow } from "@/features/ai-chat/types/ai-chat.types.ts";
|
||||
import { useChatSession } from "@/features/ai-chat/hooks/use-chat-session.ts";
|
||||
import {
|
||||
shouldCollapseOnOutsidePointer,
|
||||
@@ -86,19 +89,11 @@ const MIN_HEIGHT = 400;
|
||||
// Margin kept between the window and the viewport edges while dragging.
|
||||
const EDGE_MARGIN = 8;
|
||||
|
||||
// #184 phase 1.5 / #430: backstop for the degraded-poll fallback. The poll is
|
||||
// armed when a resume attempt could not attach to the live run and disarmed by the
|
||||
// thread on settle / local stream; this cap is the ONLY backstop against an endless
|
||||
// tick (a stuck 'streaming' row before the boot-sweep, or a user-tail 204 with no
|
||||
// run).
|
||||
//
|
||||
// #430: measured from RUN ACTIVITY, not from arm-time. A real autonomous run takes
|
||||
// 11-25 min — longer than a fixed 10-min-from-start cap, which used to cut the poll
|
||||
// off mid-run. Instead we cap on INACTIVITY: keep polling as long as the run is
|
||||
// still making progress (its persisted rows keep changing), and only give up after
|
||||
// this long with NO new activity. A genuinely stuck run produces no row changes, so
|
||||
// the idle cap still bounds it; a long-but-progressing run polls to completion.
|
||||
const DEGRADED_POLL_IDLE_MAX_MS = 10 * 60_000;
|
||||
// #184 phase 1.5 / #430 / #488: the degraded-poll fallback. The window owns only
|
||||
// a DUMB 2.5s timer, gated by an armed flag; the THREAD's run-lifecycle FSM owns
|
||||
// arm/disarm AND the inactivity cap that turns a stuck run into a `stalled` banner
|
||||
// (#488 commit 4a — the cap moved into the thread so polling->stalled is a single
|
||||
// FSM transition; the window no longer silently stops polling at the cap).
|
||||
|
||||
/** Compact token formatter: 1.2M / 3.4k / 950. */
|
||||
function formatTokens(n: number): string {
|
||||
@@ -259,17 +254,13 @@ export default function AiChatWindow() {
|
||||
[roles],
|
||||
);
|
||||
|
||||
// #184 phase 1.5: degraded-poll fallback (replaces the F4/F5/F7 latches). When
|
||||
// ChatThread could not attach to a still-running run it arms this via
|
||||
// onResumeFallback(true); the thread disarms it on settle / local stream. The
|
||||
// window only OWNS the timer (armedAtRef stamps when it was armed for the cap).
|
||||
// #184 phase 1.5 / #488: degraded-poll fallback. ChatThread's FSM arms this via
|
||||
// onResumeFallback(true) when it enters a poll-bearing recovery (attach 204 /
|
||||
// starved finish / stop) and disarms it on settle / local stream / stalled. The
|
||||
// window owns ONLY the dumb 2.5s timer; the THREAD owns arm/disarm AND the
|
||||
// inactivity cap (a stuck run -> the thread's `stalled` banner disarms this).
|
||||
const [degradedPoll, setDegradedPoll] = useState(false);
|
||||
// #430: timestamp of the LAST run activity while the poll is armed — stamped on
|
||||
// arm and re-stamped whenever the polled rows change (see the effect below). The
|
||||
// idle cap is measured from this, so a long-but-progressing run keeps polling.
|
||||
const lastActivityAtRef = useRef(0);
|
||||
const onResumeFallback = useCallback((active: boolean): void => {
|
||||
if (active) lastActivityAtRef.current = Date.now();
|
||||
setDegradedPoll(active);
|
||||
}, []);
|
||||
// Reset the degraded poll whenever the open chat changes: it is scoped to the
|
||||
@@ -281,32 +272,63 @@ export default function AiChatWindow() {
|
||||
const { data: messageRows, isLoading: messagesLoading } =
|
||||
useAiChatMessagesQuery(
|
||||
activeChatId ?? undefined,
|
||||
// DELIBERATELY DUMB (invariant 8 / task 2.4): poll every 2.5s while armed
|
||||
// and while the run is still active (#430: under the INACTIVITY cap, not a
|
||||
// fixed-from-start cap); otherwise off. NO error checks (TanStack v5 resets
|
||||
// fetchFailureCount each fetch, so consecutive errors are not expressible —
|
||||
// and the poll must survive a server restart) and NO tail checks (the
|
||||
// settled/local-stream semantics live in ChatThread, which disarms via
|
||||
// onResumeFallback(false)). The idle cap is the only backstop.
|
||||
() =>
|
||||
degradedPoll === true &&
|
||||
Date.now() - lastActivityAtRef.current < DEGRADED_POLL_IDLE_MAX_MS
|
||||
? 2500
|
||||
: false,
|
||||
// #344: gate on windowOpen too — no message history is fetched (and no
|
||||
// degraded poll runs) while the window is closed; it loads when the window
|
||||
// opens with an active chat.
|
||||
// #491: the full infinite-query no longer POLLS. It seeds the thread ONCE; the
|
||||
// degraded fallback now runs a DELTA poller (below) that augments THIS cache
|
||||
// idempotently, instead of refetching every page (with full parts) every 2.5s.
|
||||
false,
|
||||
// #344: gate on windowOpen too — no message history is fetched while the window
|
||||
// is closed; it loads when the window opens with an active chat.
|
||||
windowOpen,
|
||||
);
|
||||
|
||||
// #430: re-stamp the activity clock whenever the polled rows change while the
|
||||
// poll is armed. TanStack keeps the same `messageRows` reference across refetches
|
||||
// that return deep-equal data (structural sharing), so a new reference means the
|
||||
// run genuinely progressed — which extends the inactivity cap above. A stuck run
|
||||
// yields no reference change, so the cap eventually fires and stops the poll.
|
||||
// #491 degraded DELTA poll. While armed (degradedPoll) and the window is open on a
|
||||
// chat, poll POST /ai-chat/messages/delta every 2.5s: it returns only the rows
|
||||
// CHANGED since the previous cursor (+ the run fact) in ONE round-trip. We merge
|
||||
// those rows into the SAME infinite-query cache the thread reads (idempotently by
|
||||
// id — the delta's overlap window re-delivers rows), so the thread's reconcile
|
||||
// effect follows the detached run to its terminal row from a fraction of the wire
|
||||
// cost. The run-fact settle stays the thread FSM's job (row-status reconcile), so
|
||||
// we do NOT double-poll /run here. Cursor resets when the chat changes / disarms.
|
||||
const deltaCursorRef = useRef<string | undefined>(undefined);
|
||||
useEffect(() => {
|
||||
if (degradedPoll) lastActivityAtRef.current = Date.now();
|
||||
}, [degradedPoll, messageRows]);
|
||||
deltaCursorRef.current = undefined;
|
||||
}, [activeChatId, degradedPoll]);
|
||||
useEffect(() => {
|
||||
if (!degradedPoll || !windowOpen || !activeChatId) return;
|
||||
const chatId = activeChatId;
|
||||
let cancelled = false;
|
||||
const tick = async (): Promise<void> => {
|
||||
try {
|
||||
const res = await getAiChatMessagesDelta(chatId, deltaCursorRef.current);
|
||||
if (cancelled) return;
|
||||
deltaCursorRef.current = res.cursor;
|
||||
if (res.rows.length > 0) {
|
||||
queryClient.setQueryData(
|
||||
AI_CHAT_MESSAGES_RQ_KEY(chatId),
|
||||
(
|
||||
old:
|
||||
| {
|
||||
pages: { items: IAiChatMessageRow[]; meta: unknown }[];
|
||||
pageParams: unknown[];
|
||||
}
|
||||
| undefined,
|
||||
) =>
|
||||
old
|
||||
? { ...old, pages: mergeDeltaRowsIntoPages(old.pages, res.rows) }
|
||||
: old,
|
||||
);
|
||||
}
|
||||
} catch {
|
||||
// Transient failure (e.g. a server restart mid-run): swallow and retry on
|
||||
// the next tick — the poll must survive a bounce, like the old dumb refetch.
|
||||
}
|
||||
};
|
||||
const id = setInterval(() => void tick(), 2500);
|
||||
return () => {
|
||||
cancelled = true;
|
||||
clearInterval(id);
|
||||
};
|
||||
}, [degradedPoll, windowOpen, activeChatId, queryClient]);
|
||||
|
||||
// #184 reconnect-and-live-follow. Whether detached agent runs are enabled for
|
||||
// this workspace. When the feature is off no runs are ever created, so the
|
||||
|
||||
@@ -55,6 +55,15 @@
|
||||
padding-inline-start: 1.4em;
|
||||
}
|
||||
|
||||
/* The canonical converter renders list items through the editor schema, which
|
||||
wraps each item's content in a <p> (listItem content is `paragraph+`). Drop
|
||||
that paragraph's block margin so list items render TIGHT (no extra vertical
|
||||
gap), matching the previous marked output — same rule already applied to
|
||||
table cells above (issue #347). */
|
||||
.markdown li p {
|
||||
margin: 0;
|
||||
}
|
||||
|
||||
/* GFM tables in assistant markdown. The chat lives in a NARROW side panel, so a
|
||||
wide LLM table must scroll horizontally instead of collapsing its columns:
|
||||
`.markdown` sets `word-break: break-word`, which (with the default table
|
||||
@@ -172,6 +181,14 @@
|
||||
margin: 0 0 4px;
|
||||
}
|
||||
|
||||
/* Same as `.markdown li p` above: the canonical converter wraps every list
|
||||
item's content in a <p>, so without this each reasoning-panel list item would
|
||||
pick up `.reasoningText p`'s 4px bottom margin and render too loose. Drop it
|
||||
so Reasoning-panel lists stay tight, mirroring the pre-#347 marked output. */
|
||||
.reasoningText li p {
|
||||
margin: 0;
|
||||
}
|
||||
|
||||
.inputWrapper {
|
||||
flex: 0 0 auto;
|
||||
padding-top: var(--mantine-spacing-xs);
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -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>
|
||||
|
||||
@@ -57,6 +57,50 @@ export async function stopRun(
|
||||
return req.data;
|
||||
}
|
||||
|
||||
/**
|
||||
* Delta poll (#491): the chat's message rows changed since `cursor` (a DB-clock
|
||||
* timestamp echoed from the previous poll) plus the current run fact, in ONE
|
||||
* round-trip — the degraded-poll fallback's payload, replacing the old "refetch
|
||||
* ALL infinite-query pages every 2.5s with full parts" poll. Omit `cursor` on the
|
||||
* first poll (returns just a fresh cursor, no rows, to start the chain). The
|
||||
* overlap window guarantees occasional REPEATS, so the caller MUST merge rows
|
||||
* idempotently by id (mergeById). Owner-gated server-side.
|
||||
*/
|
||||
export async function getAiChatMessagesDelta(
|
||||
chatId: string,
|
||||
cursor?: string,
|
||||
): Promise<{
|
||||
rows: IAiChatMessageRow[];
|
||||
cursor: string;
|
||||
run: { id: string; status: string } | null;
|
||||
}> {
|
||||
const req = await api.post<{
|
||||
rows: IAiChatMessageRow[];
|
||||
cursor: string;
|
||||
run: { id: string; status: string } | null;
|
||||
}>("/ai-chat/messages/delta", { chatId, cursor });
|
||||
return req.data;
|
||||
}
|
||||
|
||||
/**
|
||||
* #488: the run-fact — "is a run active on this chat?" — first-class from the
|
||||
* server (POST /ai-chat/run). Called on mount to seed the client FSM's run-fact
|
||||
* and to VERIFY after a supersede mismatch (an observer following a superseded
|
||||
* run asks for the latest run and follows it). Returns the latest run row (with
|
||||
* its `id` and `status`) and its projected assistant message, or `run: null` when
|
||||
* the chat has never had a run. Owner-gated server-side.
|
||||
*/
|
||||
export async function getRun(chatId: string): Promise<{
|
||||
run: { id: string; status: string } | null;
|
||||
message: IAiChatMessageRow | null;
|
||||
}> {
|
||||
const req = await api.post<{
|
||||
run: { id: string; status: string } | null;
|
||||
message: IAiChatMessageRow | null;
|
||||
}>("/ai-chat/run", { chatId });
|
||||
return req.data;
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve the chat bound to a document (the current user's most-recent chat
|
||||
* created on that page), or null when there is none. Drives auto-open-on-page.
|
||||
|
||||
@@ -0,0 +1,190 @@
|
||||
# AI-chat run-lifecycle FSM — design spec (#488)
|
||||
|
||||
This is the written design that `run-fsm.ts` implements. It ships in the PR (issue
|
||||
#488 commit 1: "the spec is written FIRST and enters the PR"). It has four parts:
|
||||
(1) the event × state transition table, (2) the map of every `chat-thread.tsx` ref
|
||||
to {FSM state | FSM context | stays data}, (3) the run-fact protocol, (4) the
|
||||
invariants.
|
||||
|
||||
The reducer is a **pure function** `reduce(machine, event) → machine`. The returned
|
||||
machine carries the **command effects** for that transition; a thin runtime in
|
||||
`chat-thread.tsx` dispatches events and executes effects. Because it is pure, the
|
||||
whole machine is enumerable and unit-tested directly (event × state → next state is
|
||||
the observable property) — see `run-fsm.test.ts`.
|
||||
|
||||
---
|
||||
|
||||
## 1. Event × state transition table
|
||||
|
||||
Phases: `idle | sending | streaming | attaching | reconnecting(attempt,failed) |
|
||||
polling(reason) | stalled | stopping | superseding | error(kind)`.
|
||||
Context (orthogonal): `epoch`, `ownership: local|observer`, `runFact: {runId}|null`,
|
||||
`liveFollow` (are we following a live run we locally streamed — the reconnect
|
||||
ladder — vs a one-shot mount-attach resume? both are `observer`, but a live-follow
|
||||
drop RE-ENTERS the ladder (#488 commit 3) while a mount-resume drop polls).
|
||||
|
||||
Legend: **†** = command-transition (bumps `epoch`, I1). Effects in `[…]`.
|
||||
|
||||
| Event (source) | From phase(s) | → To phase | Effects / ctx |
|
||||
|---|---|---|---|
|
||||
| `SEND_LOCAL` (user send) | idle, error, polling, stalled, reconnecting | sending **†** | `[cancelReconnect, disarmPoll]`, ownership=local |
|
||||
| `STREAM_START{runId}` (SDK `start` metadata) | sending, attaching, reconnecting, superseding | streaming | `[cancelReconnect, disarmPoll]`, runFact←runId |
|
||||
| `FINISH_CLEAN` (onFinish clean) | streaming, … | idle | `[disarmPoll, cancelReconnect]`, runFact←null |
|
||||
| `FINISH_ABORT` (onFinish isAbort) | streaming, stopping | idle | `[disarmPoll, cancelReconnect]`, runFact←null (I4 exits stopping by this DATA) |
|
||||
| `FINISH_DISCONNECT` (observer, NOT liveFollow) | streaming(observer) | polling(disconnect-visible) | `[armPoll]` (a mount-resume drop polls) |
|
||||
| `FINISH_DISCONNECT{hasVisibleContent}` (local drop OR liveFollow) | streaming | reconnecting(1) **†** *iff runFact\|liveFollow* | `[scheduleReconnect(1)]` (+`armPoll` if visible), ownership=observer, liveFollow=true (commit 3: repeatable) |
|
||||
| `FINISH_DISCONNECT` (no runFact, not liveFollow) | streaming | idle | runFact←null (plain terminal "connection lost") |
|
||||
| `STREAM_INCOMPLETE{reason}` (observer starved/torn clean finish) | streaming(observer) | polling(reason) | `[armPoll(reason)]` |
|
||||
| `FINISH_ERROR{kind}` (onFinish isError) | any | error(kind) | `[disarmPoll, cancelReconnect]`, runFact←null |
|
||||
| `STREAM_START{runId}` (first assistant frame of a local turn) | sending | streaming | runFact←runId, `[cancelReconnect, disarmPoll]` |
|
||||
| `ATTACH_START{runId}` (mount resume) | **idle only** (F2) | attaching **†** | `[resumeStream]`, ownership=observer, runFact←runId; ignored from any non-idle phase |
|
||||
| `ATTACH_LIVE` (attach GET 2xx) | attaching | streaming | — |
|
||||
| `ATTACH_NONE` (attach GET 204/err/throw) | attaching | polling(attach-none) | `[armPoll(attach-none)]` |
|
||||
| `RECONNECT_ATTEMPT{n}` (backoff timer) | reconnecting | reconnecting(n) **†** | `[resumeStream]` |
|
||||
| `RECONNECT_ATTACHED` (reconnect GET 2xx) | reconnecting | streaming | `[cancelReconnect, disarmPoll]` — **counter reset** (commit 3) |
|
||||
| `RECONNECT_NONE` (reconnect GET 204/err), attempt<MAX | reconnecting | reconnecting(n+1) **†** | `[armPoll(attach-none), scheduleReconnect(n+1)]` |
|
||||
| `RECONNECT_NONE`, attempt=MAX | reconnecting | reconnecting(MAX, failed) | `[armPoll(reconnect-exhausted)]` |
|
||||
| `RETRY` (manual, failed banner) | reconnecting(failed) | reconnecting(1) **†** | `[resumeStream]` |
|
||||
| `RETRY` (manual, stalled banner) | stalled | polling(attach-none) **†** | `[armPoll]` |
|
||||
| `POLL_TERMINAL` (settled tail merged) | polling, reconnecting, stopping | idle | `[disarmPoll, cancelReconnect]`, runFact←null (I4) |
|
||||
| `POLL_IDLE_CAP` (inactivity cap) | polling, reconnecting | stalled | `[disarmPoll, cancelReconnect]` (commit 4a — no more silent) |
|
||||
| `POLL_IDLE_CAP` (inactivity cap) | stopping | idle | `[disarmPoll, cancelReconnect]`, runFact←null (Review #4: a Stop-armed poll with no SDK/terminal backstop gets a bounded exit — NOT `stalled`, Stop was already pressed so nothing to retry) |
|
||||
| `RUN_FACT{null}` (POST /run → null/terminal, 204) | reconnecting/attaching/polling/stopping | idle | `[cancelReconnect, disarmPoll]`, runFact←null (I3 fresh-negative gate) |
|
||||
| `RUN_FACT{runId}` | any | (same) | runFact←runId (pessimism toward an attempt) |
|
||||
| `STOP_REQUESTED` (user Stop) | streaming, reconnecting, polling | stopping **†** | `[stopRun, abortAttach, cancelReconnect, armPoll]` (poll drives the terminal — I4 exit by data) |
|
||||
| `SUPERSEDE_REQUESTED{targetRunId}` (interrupt+send) | streaming, reconnecting, polling, error | superseding **†** | `[supersede(target), cancelReconnect, disarmPoll]` |
|
||||
| `SUPERSEDE_READY{runId}` (CAS ok) | superseding | streaming | ownership=local, runFact←runId |
|
||||
| `SUPERSEDE_MISMATCH{currentRunId}` (409 SUPERSEDE_TARGET_MISMATCH) | superseding | error(supersede-mismatch) | `[postRun(verify)]`, runFact←currentRunId |
|
||||
| `SUPERSEDE_TIMEOUT` (409 SUPERSEDE_TIMEOUT) | superseding | error(supersede-timeout) | — (composer keeps text; no auto-retry) |
|
||||
| `SUPERSEDE_INVALID` (409 SUPERSEDE_INVALID) | superseding | error(supersede-invalid) | — |
|
||||
| `RUN_ALREADY_ACTIVE{activeRunId}` (409 A_RUN_ALREADY_ACTIVE, plain POST) | sending | error(run-already-active) | runFact←activeRunId (composer offers supersede; NO auto-retry) |
|
||||
| `DISPOSE` (unmount) | any | idle **†** | `[abortAttach, cancelReconnect, disarmPoll]` (I1/I5 — epoch++ kills late callbacks) |
|
||||
|
||||
**`stopping` honors any finish (re-review MEDIUM):** BEFORE the epoch filter, a
|
||||
stream finish (`FINISH_*`/`STREAM_INCOMPLETE`) arriving in phase `stopping` exits
|
||||
`stopping -> idle` regardless of generation. A plain Stop has no successor stream,
|
||||
so the aborted stream's finish IS the expected end (I4 exit by data) — and it
|
||||
carries the PRE-stop generation (STOP_REQUESTED bumped the epoch), so the filter
|
||||
would otherwise strand the machine in `stopping` (no idle-cap covers it). The filter
|
||||
stays in force for `superseding` (that is the F1 supersede drop).
|
||||
|
||||
**Epoch filter (I1):** the reducer then drops any event carrying an `epoch` that
|
||||
does not equal the current `ctx.epoch`. Outcome events (`STREAM_START`, `ATTACH_*`,
|
||||
`RECONNECT_*`, `SUPERSEDE_*`, **`FINISH_*`/`STREAM_INCOMPLETE`**, `RUN_FACT`) are
|
||||
stamped with the generation the corresponding STREAM started under (the runtime
|
||||
holds a per-owned-stream `turnEpoch`); trigger events (user actions, fresh
|
||||
disconnects) carry no epoch. **F1:** this is what makes a SUPERSEDED stream's late
|
||||
`onFinish` (a dead stream A closing after the CAS started stream B) get dropped, so
|
||||
A cannot drive the live new run into a false reconnect or reset its run-fact. The
|
||||
supersede path additionally ABORTS A and starts B only from A's onFinish (a
|
||||
microtask), because ai@6 `AbstractChat.makeRequest` corrupts overlapping streams
|
||||
(A's `finally` reads then nulls the shared `activeResponse`).
|
||||
|
||||
**Removed events (scope-cut, internal review):** `RUN_SUPERSEDED` (a ghost feature —
|
||||
never dispatched; the observer-superseded case is handled by the degraded poll,
|
||||
which follows the latest rows regardless of runId), `RECONNECT_BEGIN` (reconnect is
|
||||
entered by `FINISH_DISCONNECT`), and `POLL_ACTIVITY` (the window's activity clock was
|
||||
removed when the idle-cap moved into the thread). The reducer and this table now
|
||||
share exactly the dispatched event set.
|
||||
|
||||
### 409-code → event map (the real #487 contract consumed here)
|
||||
|
||||
| Server response | Event dispatched | error kind → banner |
|
||||
|---|---|---|
|
||||
| 409 `A_RUN_ALREADY_ACTIVE` (+ body.activeRunId) | `RUN_ALREADY_ACTIVE{activeRunId}` | run-already-active → "already answering / interrupt & send" |
|
||||
| 409 `SUPERSEDE_TARGET_MISMATCH` (+ body.activeRunId) | `SUPERSEDE_MISMATCH{currentRunId}` | supersede-mismatch → verify via /run |
|
||||
| 409 `SUPERSEDE_TIMEOUT` | `SUPERSEDE_TIMEOUT` | supersede-timeout → "couldn't interrupt in time, resend" |
|
||||
| 409 `SUPERSEDE_INVALID` | `SUPERSEDE_INVALID` | supersede-invalid → "couldn't interrupt this run" |
|
||||
| 503 `A_RUN_BEGIN_FAILED` | `FINISH_ERROR{begin-failed}` | begin-failed → "could not start, temporary" |
|
||||
|
||||
---
|
||||
|
||||
## 2. Ref-map — every `chat-thread.tsx` ref → its new home (MIGRATION RESOLVED)
|
||||
|
||||
The migration is COMPLETE: the 13 run-lifecycle FLAGS below are GONE from
|
||||
`chat-thread.tsx` (collapsed into FSM phase/ctx/effects, or deleted). What remains
|
||||
are identity/data mirrors, effect-owned controllers/timers, and ONE React-liveness
|
||||
bit — none of which is a run-lifecycle flag, so the post-merge "no new flags" rule
|
||||
holds. **Pending column: empty.**
|
||||
|
||||
| # | Old ref | Resolved to | Where now |
|
||||
|---|---|---|---|
|
||||
| 1 | `reconcileTailRef` | **FSM phase** | reconcile-merge gated on `phase ∈ {polling, reconnecting, stopping}` |
|
||||
| 2 | `noStreamHandledRef` | **FSM epoch (I1)** | the attach outcome's epoch guard drops the stale/second outcome |
|
||||
| 3 | `onNoActiveStreamRef` | **FSM event** | transport → `handleAttachOutcome` dispatches `ATTACH_NONE`/`RECONNECT_NONE` |
|
||||
| 4 | `onReconnectAttachedRef` | **FSM event** | transport dispatches `ATTACH_LIVE` / `RECONNECT_ATTACHED` |
|
||||
| 5 | `resumedTurnRef` + `resumedTurn` state | **FSM ctx `ownership`** | `ownership==='observer'` ⇒ never flush; hides "Send now" |
|
||||
| 6 | `reconnectStateRef` + `reconnectState` state | **FSM phase** | `reconnecting(attempt,failed)` renders the banner |
|
||||
| 7 | `reconnectTimerRef` | **effect-owned timer** | owned by `scheduleReconnect`/`cancelReconnect` effects (not a flag) |
|
||||
| 8 | `flushOnAbortRef` | **DELETED** | the stop→flush dance is replaced by the CAS supersede (commit 5) |
|
||||
| 9 | `interruptNextSendRef` | **DELETED** | the server injects the interrupt note from the supersede itself |
|
||||
| 10 | `supersedeRetryRef` | **DELETED** (commit 5) | the client 409 retry ladder is gone; CAS supersede replaces it |
|
||||
| 11 | `stopPendingRef` | **FSM phase `stopping`** | the deferred stop fires from the chat-id adoption effect while `stopping` |
|
||||
| 12 | `mountedRef` | **retained (React liveness)** | orthogonal to run-lifecycle; gates imperative onFinish side-effects post-unmount. Epoch (I1) handles stale COMMAND-outcomes; DISPOSE bumps it |
|
||||
| 13 | `attemptResumeRef` | **FSM `ATTACH_START` + run-fact** | mount arms attach ONLY on a confirmed active run (commit 4b: streaming-tail status, or POST /run for a user tail) |
|
||||
| 14–15 | `anchorRef {id, stepsPersisted}` | **data** (attachStrategy) | #491 tail-only: replaced `stripRef`/`strippedRowRef`. The PERSISTED assistant row that pins the run (server invariant 6) + its step frontier N; feeds `?anchor=<id>&n=<stepsPersisted>`. No strip — the seed keeps every row; entering reconnecting re-seeds from persist |
|
||||
| 16 | `attachAbortRef` | **effect-owned controller** | aborted by the `abortAttach` effect in cleanup (I5) |
|
||||
| 17–25 | `chatIdRef`, `openPageRef`, `getEditorSelectionRef`, `roleIdRef`, `stableIdRef`, `queuedRef`, `sendMessageRef`, `statusRef`, `lastForwardedChatIdRef` | **data** (identity/send mirrors) | unchanged — not lifecycle flags |
|
||||
| NEW | `pendingSupersedeRef` | **data** (send-plumbing) | the runId injected into the next `POST /stream {supersede}`; the single replacement for the 3 DELETED one-shots (#8/#9/#10) — net −2 refs |
|
||||
| NEW | `idleCapTimerRef` | **effect-owned timer** | the stalled inactivity cap → `POLL_IDLE_CAP` (commit 4a); not a flag |
|
||||
|
||||
Net: the 13 lifecycle flags (#1–#13) are eliminated: **8** → FSM phase/ctx/epoch/event
|
||||
(#1–#6, #11, #13), **3** deleted (#8/#9/#10), **`reconnectTimerRef` (#7)** becomes an
|
||||
effect-owned controller, and **`mountedRef` (#12)** is retained as React liveness
|
||||
(8 + 3 + 1 + 1 = 13). (`attachAbortRef` (#16) is outside the #1–#13 set — it was
|
||||
already an effect-owned controller.) Two effect-owned timers + one send-plumbing data
|
||||
ref are added — none is a boolean lifecycle latch.
|
||||
|
||||
---
|
||||
|
||||
## 3. Run-fact protocol (`runFact: {runId} | null`) — I3
|
||||
|
||||
"A run is active" is first-class from the SERVER, not inferred from an assistant
|
||||
message. Sources, in the order they update `ctx.runFact`:
|
||||
|
||||
1. **Init (mount):** `POST /ai-chat/run { chatId }` → `{ run, message }`. A `run`
|
||||
with a non-terminal `status` seeds `runFact = { runId: run.id }`; a null/terminal
|
||||
run seeds `null`. This is what arms the resume attempt (`ATTACH_START`) — the
|
||||
attempt is armed ONLY on a positive fact (commit 4b: a user-tail with no active
|
||||
run no longer arms a pointless poll on every open).
|
||||
2. **Live update:** the `start` stream metadata carries `runId` → `STREAM_START{runId}`.
|
||||
3. **Attach outcomes:** `ATTACH_LIVE` (2xx) confirms active; a 204 on a non-stripped
|
||||
path is an authoritative NEGATIVE fact → the runtime dispatches `RUN_FACT{null}`,
|
||||
which cancels recovery (I3 fresh-negative gate).
|
||||
4. **Poll (#491, implemented):** the degraded poll now hits the delta endpoint
|
||||
(`POST /ai-chat/messages/delta`), which ALREADY carries the run fact
|
||||
(`run: {id, status} | null`) alongside the changed rows. The client does NOT yet
|
||||
consume that run field — it still drives to a terminal ROW (merged by id),
|
||||
dispatched as `POLL_TERMINAL` — so the run field rides the wire for a future
|
||||
client that settles straight off it.
|
||||
|
||||
Pessimism rule: a stale-but-positive fact PERMITS entering recovery (attach); the
|
||||
204 then cuts it. A fresh negative fact gates recovery OUT immediately.
|
||||
|
||||
---
|
||||
|
||||
## 4. Invariants
|
||||
|
||||
- **I1 — Epoch (generation counter).** Every command-emitting transition bumps
|
||||
`ctx.epoch`; every async outcome event carries its issuing epoch; the reducer
|
||||
drops stale-epoch outcomes. Replaces the one-shot-ref zoo (`noStreamHandledRef`,
|
||||
the flush/interrupt/supersede one-shots, the `mountedRef` late-callback gate).
|
||||
- **I2 — Ownership is context, not state.** `local | observer` is orthogonal to the
|
||||
transport phase. The queue flushes ONLY under local ownership; an observer
|
||||
following a detached run never flushes (was `resumedTurnRef`).
|
||||
- **I3 — Run-fact is first-class from the server.** Reconnect is entered by the
|
||||
run-fact, not by an assistant message (commit 2). A fresh negative fact cancels
|
||||
recovery.
|
||||
- **I4 — Exit `stopping` by DATA.** A terminal row / negative run-fact / terminal
|
||||
finish exits `stopping`, never the stopRun HTTP response (which returns after the
|
||||
abort but before finalization — keying off it would unlock the composer on a 409).
|
||||
- **I5 — Dispose protocol.** Command controllers (attach GET, POST /stream, POST
|
||||
/run) are effect-owned and aborted in cleanup (`abortAttach` on `DISPOSE`), not
|
||||
render-phase refs. A client abort of an already-sent POST does not cancel the
|
||||
server action, so disarming on unmount is safe.
|
||||
- **attachStrategy** is behind the `resumeStream` effect; #491 swapped it to
|
||||
tail-only (`?anchor=&n=`, `anchorRef` data) WITHOUT touching the FSM. Entering
|
||||
reconnecting always re-seeds from persist; on a getRun failure the live partial
|
||||
is dropped + replay-from-start so it is never the tail-apply base (no #137/#161
|
||||
duplication).
|
||||
- **Queue** stays a data structure; flush/interrupt decisions are transitions.
|
||||
@@ -0,0 +1,482 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import {
|
||||
reduce,
|
||||
initialMachine,
|
||||
reconnectDelayMs,
|
||||
RECONNECT_MAX_ATTEMPTS,
|
||||
type Machine,
|
||||
type Effect,
|
||||
type Event,
|
||||
} from "./run-fsm";
|
||||
|
||||
// Drive a sequence of events through the reducer, returning the final machine.
|
||||
function run(m: Machine, ...events: Event[]): Machine {
|
||||
return events.reduce(reduce, m);
|
||||
}
|
||||
function withRunFact(runId = "run-1"): Machine {
|
||||
return {
|
||||
...initialMachine(),
|
||||
ctx: { epoch: 0, ownership: "local", runFact: { runId }, liveFollow: false },
|
||||
};
|
||||
}
|
||||
function effectTypes(m: Machine): string[] {
|
||||
return m.effects.map((e) => e.type);
|
||||
}
|
||||
function hasEffect(m: Machine, type: Effect["type"]): boolean {
|
||||
return m.effects.some((e) => e.type === type);
|
||||
}
|
||||
|
||||
describe("run-fsm — epoch invariant (I1)", () => {
|
||||
it("drops an outcome carrying a stale epoch", () => {
|
||||
// A command bumps the epoch; an outcome stamped with the OLD epoch is dropped.
|
||||
const m0 = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" }); // epoch 0->1, attaching
|
||||
expect(m0.ctx.epoch).toBe(1);
|
||||
expect(m0.phase.name).toBe("attaching");
|
||||
// A late ATTACH_LIVE from a SUPERSEDED attempt (epoch 0) must NOT drive us.
|
||||
const stale = reduce(m0, { type: "ATTACH_LIVE", epoch: 0 });
|
||||
expect(stale.phase.name).toBe("attaching");
|
||||
expect(stale.effects).toEqual([]);
|
||||
});
|
||||
|
||||
it("applies an outcome carrying the current epoch", () => {
|
||||
const m0 = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
|
||||
const live = reduce(m0, { type: "ATTACH_LIVE", epoch: m0.ctx.epoch });
|
||||
expect(live.phase.name).toBe("streaming");
|
||||
});
|
||||
|
||||
it("an outcome with no epoch is never dropped (trigger events)", () => {
|
||||
const m0 = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
|
||||
const disposed = reduce(m0, { type: "DISPOSE" });
|
||||
expect(disposed.phase.name).toBe("idle");
|
||||
expect(hasEffect(disposed, "abortAttach")).toBe(true);
|
||||
});
|
||||
|
||||
it("every command-transition increments the epoch exactly once", () => {
|
||||
let m = initialMachine();
|
||||
const before = m.ctx.epoch;
|
||||
m = reduce(m, { type: "SEND_LOCAL" });
|
||||
expect(m.ctx.epoch).toBe(before + 1);
|
||||
m = reduce(m, { type: "STOP_REQUESTED" });
|
||||
expect(m.ctx.epoch).toBe(before + 2);
|
||||
});
|
||||
});
|
||||
|
||||
describe("run-fsm — local turn", () => {
|
||||
it("SEND_LOCAL → sending, local ownership, cancels recovery", () => {
|
||||
const m = reduce(withRunFact(), { type: "SEND_LOCAL" });
|
||||
expect(m.phase.name).toBe("sending");
|
||||
expect(m.ctx.ownership).toBe("local");
|
||||
expect(effectTypes(m)).toEqual(
|
||||
expect.arrayContaining(["cancelReconnect", "disarmPoll"]),
|
||||
);
|
||||
});
|
||||
|
||||
it("STREAM_START adopts the runId into the run-fact and goes streaming", () => {
|
||||
const m = run(initialMachine(), { type: "SEND_LOCAL" });
|
||||
const s = reduce(m, { type: "STREAM_START", runId: "run-9", epoch: m.ctx.epoch });
|
||||
expect(s.phase.name).toBe("streaming");
|
||||
expect(s.ctx.runFact).toEqual({ runId: "run-9" });
|
||||
});
|
||||
|
||||
it("FINISH_CLEAN → idle, run-fact cleared, poll/reconnect disarmed", () => {
|
||||
const streaming = run(initialMachine(), { type: "SEND_LOCAL" }, { type: "STREAM_START", runId: "r" });
|
||||
const done = reduce(streaming, { type: "FINISH_CLEAN" });
|
||||
expect(done.phase.name).toBe("idle");
|
||||
expect(done.ctx.runFact).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
// #488 commit 2 — SSE break BEFORE the first assistant frame must still recover.
|
||||
describe("run-fsm — commit 2: reconnect by run-fact, not by assistant message", () => {
|
||||
it("FINISH_DISCONNECT with an active run-fact → reconnecting (even with no visible content)", () => {
|
||||
// Setup-phase break: no assistant frame yet, but a run-fact exists.
|
||||
const streaming = withRunFact("run-2");
|
||||
const m = reduce(streaming, {
|
||||
type: "FINISH_DISCONNECT",
|
||||
hasVisibleContent: false,
|
||||
epoch: streaming.ctx.epoch,
|
||||
});
|
||||
expect(m.phase.name).toBe("reconnecting");
|
||||
if (m.phase.name === "reconnecting") expect(m.phase.attempt).toBe(1);
|
||||
expect(m.ctx.ownership).toBe("observer");
|
||||
expect(hasEffect(m, "scheduleReconnect")).toBe(true);
|
||||
// No visible content -> no poll arm yet (the reconnect ladder rebuilds it).
|
||||
expect(hasEffect(m, "armPoll")).toBe(false);
|
||||
});
|
||||
|
||||
it("FINISH_DISCONNECT WITH visible content also arms the poll", () => {
|
||||
const m = reduce(withRunFact("run-2"), {
|
||||
type: "FINISH_DISCONNECT",
|
||||
hasVisibleContent: true,
|
||||
epoch: 0,
|
||||
});
|
||||
expect(m.phase.name).toBe("reconnecting");
|
||||
expect(hasEffect(m, "armPoll")).toBe(true);
|
||||
});
|
||||
|
||||
it("FINISH_DISCONNECT with NO run-fact → idle (plain connection-lost)", () => {
|
||||
const m = reduce(initialMachine(), {
|
||||
type: "FINISH_DISCONNECT",
|
||||
hasVisibleContent: true,
|
||||
epoch: 0,
|
||||
});
|
||||
expect(m.phase.name).toBe("idle");
|
||||
});
|
||||
});
|
||||
|
||||
// #488 commit 3 — a SECOND break after a successful re-attach starts a NEW ladder.
|
||||
describe("run-fsm — commit 3: repeated reconnect cycles", () => {
|
||||
it("two breaks in a row produce two reconnect cycles (counter resets on attach)", () => {
|
||||
let m = withRunFact("run-3");
|
||||
// First break -> reconnecting(1).
|
||||
m = reduce(m, { type: "FINISH_DISCONNECT", hasVisibleContent: false, epoch: m.ctx.epoch });
|
||||
expect(m.phase.name).toBe("reconnecting");
|
||||
// Attempt fires, re-attaches live.
|
||||
m = reduce(m, { type: "RECONNECT_ATTEMPT", attempt: 1, epoch: m.ctx.epoch });
|
||||
m = reduce(m, { type: "RECONNECT_ATTACHED", epoch: m.ctx.epoch });
|
||||
expect(m.phase.name).toBe("streaming");
|
||||
// SECOND break: the counter was reset, so a fresh ladder starts at attempt 1
|
||||
// (the old one-shot !wasResumed gate would have sent this to silent poll).
|
||||
m = reduce(m, { type: "FINISH_DISCONNECT", hasVisibleContent: false, epoch: m.ctx.epoch });
|
||||
expect(m.phase.name).toBe("reconnecting");
|
||||
if (m.phase.name === "reconnecting") expect(m.phase.attempt).toBe(1);
|
||||
expect(hasEffect(m, "scheduleReconnect")).toBe(true);
|
||||
});
|
||||
|
||||
it("a MOUNT-attach observer drop falls to POLL, not the reconnect ladder", () => {
|
||||
// Distinguishes commit 3 from a one-shot resume: an observer that never
|
||||
// live-followed (liveFollow false) polls on a drop.
|
||||
let m = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
|
||||
m = reduce(m, { type: "ATTACH_LIVE", epoch: m.ctx.epoch });
|
||||
expect(m.ctx.ownership).toBe("observer");
|
||||
expect(m.ctx.liveFollow).toBe(false);
|
||||
m = reduce(m, { type: "FINISH_DISCONNECT", hasVisibleContent: true, epoch: m.ctx.epoch });
|
||||
expect(m.phase.name).toBe("polling");
|
||||
expect(hasEffect(m, "armPoll")).toBe(true);
|
||||
});
|
||||
|
||||
it("STREAM_INCOMPLETE (observer starved/torn finish) → polling", () => {
|
||||
let m = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
|
||||
m = reduce(m, { type: "ATTACH_LIVE", epoch: m.ctx.epoch });
|
||||
m = reduce(m, { type: "STREAM_INCOMPLETE", reason: "starved", epoch: m.ctx.epoch });
|
||||
expect(m.phase).toEqual({ name: "polling", reason: "starved" });
|
||||
expect(hasEffect(m, "armPoll")).toBe(true);
|
||||
});
|
||||
|
||||
it("liveFollow is set on the first local drop and kept across a re-attach", () => {
|
||||
let m = withRunFact("run-3");
|
||||
m = reduce(m, { type: "FINISH_DISCONNECT", hasVisibleContent: false, epoch: m.ctx.epoch });
|
||||
expect(m.ctx.liveFollow).toBe(true);
|
||||
m = reduce(m, { type: "RECONNECT_ATTEMPT", attempt: 1, epoch: m.ctx.epoch });
|
||||
m = reduce(m, { type: "RECONNECT_ATTACHED", epoch: m.ctx.epoch });
|
||||
expect(m.ctx.liveFollow).toBe(true); // kept — so a second drop reconnects
|
||||
// A clean finish clears it.
|
||||
m = reduce(m, { type: "FINISH_CLEAN", epoch: m.ctx.epoch });
|
||||
expect(m.ctx.liveFollow).toBe(false);
|
||||
});
|
||||
|
||||
it("RECONNECT_NONE backs off through the ladder, then fails at the cap", () => {
|
||||
let m = withRunFact("run-3");
|
||||
m = reduce(m, { type: "FINISH_DISCONNECT", hasVisibleContent: false, epoch: m.ctx.epoch });
|
||||
for (let n = 1; n < RECONNECT_MAX_ATTEMPTS; n++) {
|
||||
m = reduce(m, { type: "RECONNECT_ATTEMPT", attempt: n, epoch: m.ctx.epoch });
|
||||
m = reduce(m, { type: "RECONNECT_NONE", epoch: m.ctx.epoch });
|
||||
expect(m.phase.name).toBe("reconnecting");
|
||||
if (m.phase.name === "reconnecting") {
|
||||
expect(m.phase.attempt).toBe(n + 1);
|
||||
expect(m.phase.failed).toBe(false);
|
||||
}
|
||||
// The belt-and-suspenders poll is armed each failed attempt.
|
||||
expect(hasEffect(m, "armPoll")).toBe(true);
|
||||
}
|
||||
// Final attempt fails -> failed banner (Retry), poll armed.
|
||||
m = reduce(m, { type: "RECONNECT_ATTEMPT", attempt: RECONNECT_MAX_ATTEMPTS, epoch: m.ctx.epoch });
|
||||
m = reduce(m, { type: "RECONNECT_NONE", epoch: m.ctx.epoch });
|
||||
expect(m.phase.name).toBe("reconnecting");
|
||||
if (m.phase.name === "reconnecting") expect(m.phase.failed).toBe(true);
|
||||
// RETRY restarts at attempt 1.
|
||||
m = reduce(m, { type: "RETRY" });
|
||||
expect(m.phase.name).toBe("reconnecting");
|
||||
if (m.phase.name === "reconnecting") {
|
||||
expect(m.phase.attempt).toBe(1);
|
||||
expect(m.phase.failed).toBe(false);
|
||||
}
|
||||
expect(hasEffect(m, "resumeStream")).toBe(true);
|
||||
});
|
||||
|
||||
it("reconnectDelayMs is the exponential backoff 1s,2s,4s,8s,16s", () => {
|
||||
expect([1, 2, 3, 4, 5].map(reconnectDelayMs)).toEqual([1000, 2000, 4000, 8000, 16000]);
|
||||
});
|
||||
});
|
||||
|
||||
// #488 commit 4 — polling stalled-state + user-tail gating.
|
||||
describe("run-fsm — commit 4: stalled + run-fact gating", () => {
|
||||
it("POLL_IDLE_CAP: polling → stalled with a banner (poll disarmed), not silent", () => {
|
||||
let m = reduce(withRunFact(), { type: "ATTACH_START", runId: "r" });
|
||||
m = reduce(m, { type: "ATTACH_NONE", epoch: m.ctx.epoch });
|
||||
expect(m.phase.name).toBe("polling");
|
||||
m = reduce(m, { type: "POLL_IDLE_CAP" });
|
||||
expect(m.phase.name).toBe("stalled");
|
||||
expect(hasEffect(m, "disarmPoll")).toBe(true);
|
||||
});
|
||||
|
||||
it("RETRY from stalled re-arms the poll", () => {
|
||||
let m = reduce(withRunFact(), { type: "ATTACH_START", runId: "r" });
|
||||
m = reduce(m, { type: "ATTACH_NONE", epoch: m.ctx.epoch });
|
||||
m = reduce(m, { type: "POLL_IDLE_CAP" });
|
||||
m = reduce(m, { type: "RETRY" });
|
||||
expect(m.phase.name).toBe("polling");
|
||||
expect(hasEffect(m, "armPoll")).toBe(true);
|
||||
});
|
||||
|
||||
it("a fresh NEGATIVE run-fact while attaching cancels recovery (user-tail, no active run)", () => {
|
||||
// The mount POST /run returns no active run: attaching → idle, no poll armed.
|
||||
let m = reduce(withRunFact(), { type: "ATTACH_START", runId: "r" });
|
||||
m = reduce(m, { type: "RUN_FACT", runFact: null, epoch: m.ctx.epoch });
|
||||
expect(m.phase.name).toBe("idle");
|
||||
expect(m.ctx.runFact).toBeNull();
|
||||
expect(hasEffect(m, "disarmPoll")).toBe(true);
|
||||
});
|
||||
|
||||
it("a negative run-fact while polling stops the poll", () => {
|
||||
let m = reduce(withRunFact(), { type: "ATTACH_START", runId: "r" });
|
||||
m = reduce(m, { type: "ATTACH_NONE", epoch: m.ctx.epoch });
|
||||
m = reduce(m, { type: "RUN_FACT", runFact: null, epoch: m.ctx.epoch });
|
||||
expect(m.phase.name).toBe("idle");
|
||||
});
|
||||
|
||||
it("POLL_TERMINAL settles polling → idle (I4 data-driven exit)", () => {
|
||||
let m = reduce(withRunFact(), { type: "ATTACH_START", runId: "r" });
|
||||
m = reduce(m, { type: "ATTACH_NONE", epoch: m.ctx.epoch });
|
||||
m = reduce(m, { type: "POLL_TERMINAL" });
|
||||
expect(m.phase.name).toBe("idle");
|
||||
expect(m.ctx.runFact).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
// #488 commit 5 — error classification + supersede CAS transitions.
|
||||
describe("run-fsm — commit 5: supersede CAS + error classification", () => {
|
||||
it("SUPERSEDE_REQUESTED → superseding, fires the CAS effect, bumps epoch", () => {
|
||||
const streaming = withRunFact("run-old");
|
||||
const m = reduce(streaming, { type: "SUPERSEDE_REQUESTED", targetRunId: "run-old" });
|
||||
expect(m.phase.name).toBe("superseding");
|
||||
expect(m.ctx.epoch).toBe(streaming.ctx.epoch + 1);
|
||||
const sup = m.effects.find((e) => e.type === "supersede");
|
||||
expect(sup).toEqual({ type: "supersede", targetRunId: "run-old" });
|
||||
});
|
||||
|
||||
it("SUPERSEDE_READY → streaming as the new local owner", () => {
|
||||
let m = reduce(withRunFact("run-old"), { type: "SUPERSEDE_REQUESTED", targetRunId: "run-old" });
|
||||
m = reduce(m, { type: "SUPERSEDE_READY", runId: "run-new", epoch: m.ctx.epoch });
|
||||
expect(m.phase.name).toBe("streaming");
|
||||
expect(m.ctx.ownership).toBe("local");
|
||||
expect(m.ctx.runFact).toEqual({ runId: "run-new" });
|
||||
});
|
||||
|
||||
it("SUPERSEDE_MISMATCH → error(supersede-mismatch) + verify via /run (no blind banner)", () => {
|
||||
let m = reduce(withRunFact("run-old"), { type: "SUPERSEDE_REQUESTED", targetRunId: "run-old" });
|
||||
m = reduce(m, { type: "SUPERSEDE_MISMATCH", currentRunId: "run-x", epoch: m.ctx.epoch });
|
||||
expect(m.phase).toEqual({ name: "error", kind: "supersede-mismatch" });
|
||||
expect(hasEffect(m, "postRun")).toBe(true);
|
||||
expect(m.ctx.runFact).toEqual({ runId: "run-x" });
|
||||
});
|
||||
|
||||
it("SUPERSEDE_TIMEOUT → error(supersede-timeout), no auto-retry effect", () => {
|
||||
let m = reduce(withRunFact("run-old"), { type: "SUPERSEDE_REQUESTED", targetRunId: "run-old" });
|
||||
m = reduce(m, { type: "SUPERSEDE_TIMEOUT", epoch: m.ctx.epoch });
|
||||
expect(m.phase).toEqual({ name: "error", kind: "supersede-timeout" });
|
||||
expect(m.effects).toEqual([]);
|
||||
});
|
||||
|
||||
it("SUPERSEDE_INVALID → error(supersede-invalid)", () => {
|
||||
let m = reduce(withRunFact("run-old"), { type: "SUPERSEDE_REQUESTED", targetRunId: "run-old" });
|
||||
m = reduce(m, { type: "SUPERSEDE_INVALID", epoch: m.ctx.epoch });
|
||||
expect(m.phase).toEqual({ name: "error", kind: "supersede-invalid" });
|
||||
});
|
||||
|
||||
it("a stale SUPERSEDE outcome from a superseded epoch is dropped", () => {
|
||||
let m = reduce(withRunFact("run-old"), { type: "SUPERSEDE_REQUESTED", targetRunId: "run-old" });
|
||||
const supersedingEpoch = m.ctx.epoch;
|
||||
// The user retriggers, bumping the epoch again.
|
||||
m = reduce(m, { type: "SUPERSEDE_REQUESTED", targetRunId: "run-old" });
|
||||
// The first CAS's late TIMEOUT (old epoch) must NOT knock us out of superseding.
|
||||
const late = reduce(m, { type: "SUPERSEDE_TIMEOUT", epoch: supersedingEpoch });
|
||||
expect(late.phase.name).toBe("superseding");
|
||||
});
|
||||
|
||||
it("RUN_ALREADY_ACTIVE (plain POST gate) → error(run-already-active), no retry effect", () => {
|
||||
const m = reduce(run(initialMachine(), { type: "SEND_LOCAL" }), { type: "RUN_ALREADY_ACTIVE" });
|
||||
expect(m.phase).toEqual({ name: "error", kind: "run-already-active" });
|
||||
expect(m.effects).toEqual([]);
|
||||
});
|
||||
|
||||
it("#497/S4: RUN_ALREADY_ACTIVE{activeRunId} ADOPTS the server's active run as the run-fact", () => {
|
||||
// The server sends `activeRunId` so a later supersede can TARGET that run
|
||||
// instead of a blind promote+abort. Absorb it into runFact.
|
||||
const m = reduce(run(initialMachine(), { type: "SEND_LOCAL" }), {
|
||||
type: "RUN_ALREADY_ACTIVE",
|
||||
activeRunId: "run-foreign",
|
||||
});
|
||||
expect(m.phase).toEqual({ name: "error", kind: "run-already-active" });
|
||||
expect(m.ctx.runFact).toEqual({ runId: "run-foreign" });
|
||||
expect(m.effects).toEqual([]);
|
||||
});
|
||||
|
||||
it("#497/S4: RUN_ALREADY_ACTIVE without an activeRunId keeps the prior run-fact", () => {
|
||||
const seeded = reduce(run(initialMachine(), { type: "SEND_LOCAL" }), {
|
||||
type: "RUN_FACT",
|
||||
runFact: { runId: "run-prior" },
|
||||
});
|
||||
const m = reduce(seeded, { type: "RUN_ALREADY_ACTIVE" });
|
||||
expect(m.ctx.runFact).toEqual({ runId: "run-prior" });
|
||||
});
|
||||
});
|
||||
|
||||
// #488 F2 — a late mount `getRun → ATTACH_START` must not hijack a local turn.
|
||||
describe("run-fsm — F2: ATTACH_START only from idle", () => {
|
||||
it("ATTACH_START from a local `sending` turn is ignored (no observer hijack)", () => {
|
||||
const sending = reduce(initialMachine(), { type: "SEND_LOCAL" }); // idle -> sending, local
|
||||
const m = reduce(sending, { type: "ATTACH_START", runId: "r" });
|
||||
expect(m.phase.name).toBe("sending");
|
||||
expect(m.ctx.ownership).toBe("local"); // NOT flipped to observer
|
||||
expect(m.effects).toEqual([]); // no resumeStream
|
||||
});
|
||||
|
||||
it("ATTACH_START from idle attaches as normal", () => {
|
||||
const m = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
|
||||
expect(m.phase.name).toBe("attaching");
|
||||
expect(m.ctx.ownership).toBe("observer");
|
||||
expect(hasEffect(m, "resumeStream")).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe("run-fsm — stop (I4: exit by data)", () => {
|
||||
it("STOP_REQUESTED → stopping, fires stopRun + abortAttach, no data-independent exit", () => {
|
||||
const m = reduce(withRunFact(), { type: "STOP_REQUESTED" });
|
||||
expect(m.phase.name).toBe("stopping");
|
||||
expect(effectTypes(m)).toEqual(expect.arrayContaining(["stopRun", "abortAttach"]));
|
||||
});
|
||||
|
||||
it("stopping exits on the aborted stream's finish carrying the PRE-STOP epoch", () => {
|
||||
// MEDIUM (#488 re-review): STOP_REQUESTED is a command that BUMPS the epoch, but
|
||||
// the runtime stamps the aborted stream's onFinish with the stream's START (pre-
|
||||
// stop) generation — exactly what the component sends. `stopping` must HONOR
|
||||
// that finish regardless of generation (no idle-cap covers `stopping`).
|
||||
// MUTATION-VERIFY: remove the honor-in-`stopping` branch and this hangs in
|
||||
// `stopping` (the epoch filter drops the pre-stop finish) -> red.
|
||||
const preStopEpoch = withRunFact().ctx.epoch; // E1 (the stream's start epoch)
|
||||
let m = reduce(withRunFact(), { type: "STOP_REQUESTED" }); // E1 -> E2, stopping
|
||||
expect(m.ctx.epoch).toBe(preStopEpoch + 1);
|
||||
m = reduce(m, { type: "FINISH_ABORT", epoch: preStopEpoch }); // NOT the current epoch
|
||||
expect(m.phase.name).toBe("idle");
|
||||
expect(m.ctx.runFact).toBeNull();
|
||||
});
|
||||
|
||||
it("stopping exits on a clean finish carrying the pre-stop epoch too", () => {
|
||||
const preStopEpoch = withRunFact().ctx.epoch;
|
||||
let m = reduce(withRunFact(), { type: "STOP_REQUESTED" });
|
||||
m = reduce(m, { type: "FINISH_CLEAN", epoch: preStopEpoch });
|
||||
expect(m.phase.name).toBe("idle");
|
||||
});
|
||||
|
||||
it("stopping exits on a negative run-fact (data)", () => {
|
||||
let m = reduce(withRunFact(), { type: "STOP_REQUESTED" });
|
||||
m = reduce(m, { type: "RUN_FACT", runFact: null, epoch: m.ctx.epoch });
|
||||
expect(m.phase.name).toBe("idle");
|
||||
});
|
||||
|
||||
// Review #4: `stopping` arms the poll but had no inactivity backstop.
|
||||
it("review-4: POLL_IDLE_CAP in `stopping` exits to idle (bounded), NOT stalled", () => {
|
||||
let m = reduce(withRunFact(), { type: "STOP_REQUESTED" });
|
||||
expect(m.phase.name).toBe("stopping");
|
||||
expect(hasEffect(m, "armPoll")).toBe(true);
|
||||
// MUTATION-VERIFY: drop the `stopping` branch in POLL_IDLE_CAP and this hangs
|
||||
// in `stopping` (poll forever) -> red.
|
||||
m = reduce(m, { type: "POLL_IDLE_CAP" });
|
||||
expect(m.phase.name).toBe("idle");
|
||||
expect(hasEffect(m, "disarmPoll")).toBe(true);
|
||||
expect(m.ctx.ownership).toBe("local");
|
||||
});
|
||||
});
|
||||
|
||||
// Review #1: positive attach outcomes must be guarded by the SOURCE phase — the
|
||||
// epoch filter alone is insufficient because POLL_TERMINAL uses to() (no epoch
|
||||
// bump) and does not abort the in-flight GET.
|
||||
describe("run-fsm — review-1: attach outcomes guarded by source phase", () => {
|
||||
it("a late RECONNECT_ATTACHED after POLL_TERMINAL stays idle (no phantom streaming)", () => {
|
||||
let m = withRunFact("run-1");
|
||||
m = reduce(m, { type: "FINISH_DISCONNECT", hasVisibleContent: true, epoch: m.ctx.epoch });
|
||||
m = reduce(m, { type: "RECONNECT_ATTEMPT", attempt: 1, epoch: m.ctx.epoch }); // attach GET
|
||||
const epoch = m.ctx.epoch;
|
||||
// The armed degraded poll reaches the terminal row FIRST (epoch unchanged).
|
||||
m = reduce(m, { type: "POLL_TERMINAL" });
|
||||
expect(m.phase.name).toBe("idle");
|
||||
expect(m.ctx.epoch).toBe(epoch); // POLL_TERMINAL did NOT bump the epoch
|
||||
// The slow GET returns live 2xx under the SAME epoch — must NOT resurrect.
|
||||
m = reduce(m, { type: "RECONNECT_ATTACHED", epoch });
|
||||
expect(m.phase.name).toBe("idle");
|
||||
});
|
||||
|
||||
it("a late ATTACH_LIVE / ATTACH_NONE after leaving `attaching` is ignored", () => {
|
||||
let m = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
|
||||
const epoch = m.ctx.epoch;
|
||||
m = reduce(m, { type: "ATTACH_NONE", epoch }); // attaching -> polling
|
||||
m = reduce(m, { type: "POLL_TERMINAL" }); // -> idle (epoch unchanged)
|
||||
expect(m.phase.name).toBe("idle");
|
||||
m = reduce(m, { type: "ATTACH_LIVE", epoch }); // late 2xx, same epoch
|
||||
expect(m.phase.name).toBe("idle");
|
||||
// And a late ATTACH_NONE (not `attaching`) is a no-op too.
|
||||
m = reduce(m, { type: "ATTACH_NONE", epoch });
|
||||
expect(m.phase.name).toBe("idle");
|
||||
});
|
||||
});
|
||||
|
||||
// Review #2: every terminal transition resets ownership to local.
|
||||
describe("run-fsm — review-2: terminal transitions reset ownership to local", () => {
|
||||
const observer = (): Machine => {
|
||||
let m = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
|
||||
m = reduce(m, { type: "ATTACH_LIVE", epoch: m.ctx.epoch });
|
||||
expect(m.ctx.ownership).toBe("observer");
|
||||
return m;
|
||||
};
|
||||
it("FINISH_CLEAN resets ownership", () => {
|
||||
const m = reduce(observer(), { type: "FINISH_CLEAN", epoch: observer().ctx.epoch });
|
||||
expect(m.ctx.ownership).toBe("local");
|
||||
});
|
||||
it("FINISH_ERROR / POLL_TERMINAL / RUN_FACT(null) reset ownership", () => {
|
||||
let o = observer();
|
||||
expect(reduce(o, { type: "FINISH_ERROR", kind: "stream", epoch: o.ctx.epoch }).ctx.ownership).toBe("local");
|
||||
// POLL_TERMINAL from an observer polling phase
|
||||
let p = reduce(observer(), { type: "STREAM_INCOMPLETE", reason: "starved", epoch: observer().ctx.epoch });
|
||||
expect(reduce(p, { type: "POLL_TERMINAL" }).ctx.ownership).toBe("local");
|
||||
// RUN_FACT(null) from an observer attaching phase
|
||||
let a = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
|
||||
expect(reduce(a, { type: "RUN_FACT", runFact: null, epoch: a.ctx.epoch }).ctx.ownership).toBe("local");
|
||||
});
|
||||
});
|
||||
|
||||
describe("run-fsm — ownership (I2) is context, orthogonal to phase", () => {
|
||||
it("attach/reconnect set observer; send/supersede-ready set local", () => {
|
||||
let m = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
|
||||
expect(m.ctx.ownership).toBe("observer");
|
||||
m = reduce(m, { type: "ATTACH_LIVE", epoch: m.ctx.epoch });
|
||||
expect(m.phase.name).toBe("streaming");
|
||||
expect(m.ctx.ownership).toBe("observer"); // still observing a detached run
|
||||
// A local send flips ownership back to local.
|
||||
m = reduce(m, { type: "SEND_LOCAL" });
|
||||
expect(m.ctx.ownership).toBe("local");
|
||||
});
|
||||
});
|
||||
|
||||
describe("run-fsm — dispose (I5)", () => {
|
||||
it("DISPOSE from any phase aborts controllers and bumps epoch", () => {
|
||||
let m = reduce(withRunFact(), { type: "ATTACH_START", runId: "r" });
|
||||
const before = m.ctx.epoch;
|
||||
m = reduce(m, { type: "DISPOSE" });
|
||||
expect(m.phase.name).toBe("idle");
|
||||
expect(m.ctx.epoch).toBe(before + 1);
|
||||
expect(effectTypes(m)).toEqual(
|
||||
expect.arrayContaining(["abortAttach", "cancelReconnect", "disarmPoll"]),
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,600 @@
|
||||
/**
|
||||
* Run-lifecycle finite state machine for a single AI-chat thread (#488).
|
||||
*
|
||||
* ============================================================================
|
||||
* WHY THIS EXISTS
|
||||
* ----------------------------------------------------------------------------
|
||||
* The resume/reconnect/poll/stop/supersede lifecycle used to be spread across
|
||||
* ~26 `useRef` one-shot flags in `chat-thread.tsx`, each disarmed "on every
|
||||
* path". Ownerless flag combinations produced silent UI freezes, and every fix
|
||||
* added another ref (the #381 -> #432 -> #456 spiral). This module replaces that
|
||||
* ref-zoo with ONE pure reducer whose transitions are enumerable and unit-
|
||||
* testable in isolation (event x state -> next state is the observable property).
|
||||
*
|
||||
* The reducer is PURE: it owns no timers, no fetches, no React state. It maps
|
||||
* `(machine, event) -> machine`, where the returned machine carries the list of
|
||||
* COMMAND EFFECTS to run for that transition. A thin runtime in `chat-thread.tsx`
|
||||
* dispatches events (from SDK callbacks / HTTP outcomes) and executes the
|
||||
* effects (attach GET, POST /stream, POST /run, POST /stop, backoff timers,
|
||||
* poll arm/disarm). The runtime lives in a THREAD, not the window, so a late SDK
|
||||
* callback dies with the owner (kills the "event from a dead view" class, #161).
|
||||
*
|
||||
* ============================================================================
|
||||
* INVARIANTS (see run-fsm.spec.md for the full spec + tables)
|
||||
* ----------------------------------------------------------------------------
|
||||
* I1 EPOCH (generation counter). Commands (`resumeStream`, `postRun`, `stop`,
|
||||
* `supersede`, `scheduleReconnect`) are async; their outcomes arrive on the
|
||||
* SAME SDK/HTTP callbacks. Every command-emitting transition increments
|
||||
* `ctx.epoch`; every OUTCOME event carries the epoch it was issued under;
|
||||
* the reducer DROPS an outcome whose epoch != the current epoch. This is
|
||||
* what the one-shot-ref zoo used to approximate by hand.
|
||||
* I2 OWNERSHIP is a CONTEXT FIELD (`'local' | 'observer'`), not a state —
|
||||
* orthogonal to the transport phase. The queue is flushed ONLY by a local
|
||||
* owner (an observer following a detached run never flushes).
|
||||
* I3 RUN-FACT ("a run is active") is first-class from the server: `runFact`
|
||||
* holds the server-confirmed active run id (POST /run on mount, the `start`
|
||||
* metadata runId, attach outcomes). Reconnect is entered by the RUN-FACT,
|
||||
* not by the presence of an assistant message (#488 commit 2). A fresh
|
||||
* negative fact (null) cancels reconnect immediately.
|
||||
* I4 Exit `stopping` by DATA (a terminal row / negative run-fact), NEVER by the
|
||||
* stopRun HTTP response (which returns after abort, before finalization).
|
||||
* I5 Command controllers are effect-owned (abort in cleanup), NOT render-phase
|
||||
* refs — expressed here as the `abortAttach` effect on disposing transitions.
|
||||
* ============================================================================
|
||||
*/
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Phases (the transport lifecycle). Ownership / runFact are CONTEXT, not here.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/** Why the degraded poll is the active recovery. */
|
||||
export type PollReason =
|
||||
| "attach-none" // mount attach returned 204 / error — nothing live to attach
|
||||
| "starved" // a resumed finish carried no visible content
|
||||
| "disconnect-visible" // a live disconnect WITH on-screen content — poll to terminal
|
||||
| "reconnect-exhausted"; // the live re-attach ladder gave up
|
||||
|
||||
/** The classified error kind (drives the banner text + composer behavior). */
|
||||
export type ErrorKind =
|
||||
| "stream" // a generic provider/network stream error (useChat error)
|
||||
| "run-already-active" // 409 A_RUN_ALREADY_ACTIVE (a plain POST hit the gate)
|
||||
| "supersede-mismatch" // 409 SUPERSEDE_TARGET_MISMATCH (CAS target moved)
|
||||
| "supersede-timeout" // 409 SUPERSEDE_TIMEOUT (old run did not settle in W)
|
||||
| "supersede-invalid" // 409 SUPERSEDE_INVALID (bad supersede target)
|
||||
| "begin-failed"; // 503 A_RUN_BEGIN_FAILED (could not start the run)
|
||||
|
||||
export type Phase =
|
||||
| { name: "idle" }
|
||||
| { name: "sending" } // local POST in flight, before the first frame
|
||||
| { name: "streaming" } // receiving frames
|
||||
| { name: "attaching" } // mount-time attach GET in flight
|
||||
| { name: "reconnecting"; attempt: number; failed: boolean }
|
||||
| { name: "polling"; reason: PollReason }
|
||||
| { name: "stalled" } // poll hit the inactivity cap — banner + Retry
|
||||
| { name: "stopping" }
|
||||
| { name: "superseding" }
|
||||
| { name: "error"; kind: ErrorKind };
|
||||
|
||||
export type Ownership = "local" | "observer";
|
||||
|
||||
/** The server-confirmed active run, or null when no run is active. */
|
||||
export type RunFact = { runId: string } | null;
|
||||
|
||||
export interface Ctx {
|
||||
/** I1: generation counter — every command-transition increments it. */
|
||||
epoch: number;
|
||||
/** I2: does THIS client own the turn's writes (local streamer) or observe? */
|
||||
ownership: Ownership;
|
||||
/** I3: the server-confirmed active run. */
|
||||
runFact: RunFact;
|
||||
/**
|
||||
* Are we FOLLOWING a live run we were locally streaming (the reconnect ladder),
|
||||
* as opposed to a one-shot mount-attach resume? Both are `ownership: 'observer'`,
|
||||
* but they recover DIFFERENTLY on a drop: a live-follow drop RE-ENTERS the
|
||||
* reconnect ladder (#488 commit 3 — the second break after a successful re-attach
|
||||
* must reconnect again, not fall to silent poll), while a mount-resume drop falls
|
||||
* to the degraded poll. This is the ctx bit that separates the two WITHOUT a new
|
||||
* component ref (it is why commit 3 needs the FSM, not a surgical patch).
|
||||
*/
|
||||
liveFollow: boolean;
|
||||
}
|
||||
|
||||
export interface Machine {
|
||||
phase: Phase;
|
||||
ctx: Ctx;
|
||||
/** Command effects to run for the transition that produced THIS machine.
|
||||
* The runtime executes them and does not read them again. */
|
||||
effects: Effect[];
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Command effects (the reducer's only side-channel — executed by the runtime).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export type Effect =
|
||||
/** POST /run to (re)establish or verify the run-fact. `reason` is diagnostic. */
|
||||
| { type: "postRun"; reason: "mount" | "verify" }
|
||||
/** Trigger the SDK `resumeStream()` (attach GET via prepareReconnectToStream). */
|
||||
| { type: "resumeStream" }
|
||||
/** Schedule a reconnect attempt after a backoff, then dispatch RECONNECT_ATTEMPT. */
|
||||
| { type: "scheduleReconnect"; attempt: number; delayMs: number }
|
||||
/** Cancel any pending reconnect backoff timer. */
|
||||
| { type: "cancelReconnect" }
|
||||
/** Arm the degraded poll (the window's dumb timer follows the run in the DB). */
|
||||
| { type: "armPoll"; reason: PollReason }
|
||||
/** Disarm the degraded poll. */
|
||||
| { type: "disarmPoll" }
|
||||
/** POST /stop the chat's active run (authoritative detached-run stop). */
|
||||
| { type: "stopRun" }
|
||||
/** POST /stream { supersede: { runId } } — the CAS "interrupt and send now". */
|
||||
| { type: "supersede"; targetRunId: string }
|
||||
/** Abort the in-flight attach/reconnect GET controller (dispose / observer stop). */
|
||||
| { type: "abortAttach" };
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Events. An OUTCOME event MAY carry `epoch`; if it does and it does not equal
|
||||
// the current epoch, the reducer drops it (I1). Trigger events (user actions,
|
||||
// fresh disconnects) carry no epoch and are never dropped.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export type Event =
|
||||
// -- local turn --
|
||||
| { type: "SEND_LOCAL" }
|
||||
| { type: "STREAM_START"; runId?: string; epoch?: number }
|
||||
/** An OBSERVER's attached stream ended WITHOUT reaching terminal (a starved
|
||||
* clean replay, or a torn resume) — fall to the degraded poll to drive the row
|
||||
* to its real terminal state. (A live-follow drop uses FINISH_DISCONNECT.) */
|
||||
| { type: "STREAM_INCOMPLETE"; reason: PollReason; epoch?: number }
|
||||
| { type: "FINISH_CLEAN"; epoch?: number }
|
||||
| { type: "FINISH_ABORT"; epoch?: number }
|
||||
| { type: "FINISH_DISCONNECT"; hasVisibleContent: boolean; epoch?: number }
|
||||
| { type: "FINISH_ERROR"; kind: ErrorKind; epoch?: number }
|
||||
// -- mount attach (resume) --
|
||||
| { type: "ATTACH_START"; runId?: string }
|
||||
| { type: "ATTACH_LIVE"; epoch?: number }
|
||||
| { type: "ATTACH_NONE"; epoch?: number }
|
||||
// -- reconnect after a live disconnect (entered by FINISH_DISCONNECT, #488 c2) --
|
||||
| { type: "RECONNECT_ATTEMPT"; attempt: number; epoch?: number }
|
||||
| { type: "RECONNECT_ATTACHED"; epoch?: number }
|
||||
| { type: "RECONNECT_NONE"; epoch?: number }
|
||||
| { type: "RETRY" }
|
||||
// -- degraded poll --
|
||||
| { type: "POLL_TERMINAL" }
|
||||
| { type: "POLL_IDLE_CAP" }
|
||||
// -- run-fact (server-confirmed active run) --
|
||||
| { type: "RUN_FACT"; runFact: RunFact; epoch?: number }
|
||||
// -- stop --
|
||||
| { type: "STOP_REQUESTED" }
|
||||
// -- supersede (CAS) --
|
||||
| { type: "SUPERSEDE_REQUESTED"; targetRunId: string }
|
||||
| { type: "SUPERSEDE_READY"; runId?: string; epoch?: number }
|
||||
| { type: "SUPERSEDE_MISMATCH"; currentRunId?: string; epoch?: number }
|
||||
| { type: "SUPERSEDE_TIMEOUT"; epoch?: number }
|
||||
| { type: "SUPERSEDE_INVALID"; epoch?: number }
|
||||
| { type: "RUN_ALREADY_ACTIVE"; activeRunId?: string }
|
||||
// -- lifecycle --
|
||||
| { type: "DISPOSE" };
|
||||
|
||||
export const RECONNECT_MAX_ATTEMPTS = 5;
|
||||
export const RECONNECT_BASE_DELAY_MS = 1000;
|
||||
/** Backoff before attempt N (1-based): 1s, 2s, 4s, 8s, 16s. */
|
||||
export function reconnectDelayMs(attempt: number): number {
|
||||
return RECONNECT_BASE_DELAY_MS * 2 ** (attempt - 1);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Constructors / helpers.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export function initialMachine(overrides?: Partial<Ctx>): Machine {
|
||||
return {
|
||||
phase: { name: "idle" },
|
||||
ctx: { epoch: 0, ownership: "local", runFact: null, liveFollow: false, ...overrides },
|
||||
effects: [],
|
||||
};
|
||||
}
|
||||
|
||||
/** Build a machine result: a phase, optional ctx patch, and effects. Empty
|
||||
* effects by default. Never mutates the input. */
|
||||
function to(
|
||||
m: Machine,
|
||||
phase: Phase,
|
||||
opts?: { ctx?: Partial<Ctx>; effects?: Effect[] },
|
||||
): Machine {
|
||||
return {
|
||||
phase,
|
||||
ctx: { ...m.ctx, ...(opts?.ctx ?? {}) },
|
||||
effects: opts?.effects ?? [],
|
||||
};
|
||||
}
|
||||
|
||||
/** No transition: keep the phase, clear effects (so a re-run does not re-fire). */
|
||||
function stay(m: Machine): Machine {
|
||||
return { phase: m.phase, ctx: m.ctx, effects: [] };
|
||||
}
|
||||
|
||||
/** A command-transition: same as `to` but bumps the epoch (I1). Any outcome
|
||||
* event issued under the old epoch is dropped once this lands. */
|
||||
function command(
|
||||
m: Machine,
|
||||
phase: Phase,
|
||||
effects: Effect[],
|
||||
ctx?: Partial<Ctx>,
|
||||
): Machine {
|
||||
return {
|
||||
phase,
|
||||
ctx: { ...m.ctx, ...(ctx ?? {}), epoch: m.ctx.epoch + 1 },
|
||||
effects,
|
||||
};
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// The pure reducer.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/** The terminal stream-finish events (one turn's stream ended). */
|
||||
function isFinishEvent(event: Event): boolean {
|
||||
return (
|
||||
event.type === "FINISH_ABORT" ||
|
||||
event.type === "FINISH_CLEAN" ||
|
||||
event.type === "FINISH_DISCONNECT" ||
|
||||
event.type === "FINISH_ERROR" ||
|
||||
event.type === "STREAM_INCOMPLETE"
|
||||
);
|
||||
}
|
||||
|
||||
export function reduce(m: Machine, event: Event): Machine {
|
||||
// MEDIUM (#488 re-review): honor ANY stream finish in `stopping` regardless of
|
||||
// generation. A plain user Stop has NO successor stream — the aborted stream's
|
||||
// finish IS the expected end of the stop, so exit `stopping -> idle` by that DATA
|
||||
// (I4). The epoch filter below must NOT drop it: STOP_REQUESTED bumped the epoch,
|
||||
// but the finish carries the PRE-stop generation (the runtime stamps it with the
|
||||
// stream's start epoch), so I1 would otherwise strand the machine in `stopping`
|
||||
// forever (no idle-cap covers `stopping`). The epoch filter stays in force for
|
||||
// `superseding` (a successor B owns) — that is the F1 supersede drop.
|
||||
if (m.phase.name === "stopping" && isFinishEvent(event)) {
|
||||
return to(m, { name: "idle" }, {
|
||||
// Reset ownership to local on this terminal transition (review #2): otherwise
|
||||
// an observer-stop leaves ownership 'observer' and hides "Send now" forever.
|
||||
ctx: { runFact: null, liveFollow: false, ownership: "local" },
|
||||
effects: [{ type: "disarmPoll" }, { type: "cancelReconnect" }],
|
||||
});
|
||||
}
|
||||
|
||||
// I1: drop a stale outcome (an event issued under a superseded epoch).
|
||||
if ("epoch" in event && event.epoch !== undefined && event.epoch !== m.ctx.epoch) {
|
||||
return stay(m);
|
||||
}
|
||||
|
||||
switch (event.type) {
|
||||
// ---- local turn ----------------------------------------------------
|
||||
case "SEND_LOCAL":
|
||||
// A local send owns the view: leave any recovery, become the local
|
||||
// streamer, disarm poll/reconnect. epoch++ so a late recovery outcome
|
||||
// from the previous phase is dropped.
|
||||
return command(
|
||||
m,
|
||||
{ name: "sending" },
|
||||
[{ type: "cancelReconnect" }, { type: "disarmPoll" }],
|
||||
{ ownership: "local", liveFollow: false },
|
||||
);
|
||||
|
||||
case "STREAM_INCOMPLETE":
|
||||
// An OBSERVER's attached stream ended incomplete (starved / torn) — follow
|
||||
// the run to terminal via the degraded poll.
|
||||
return to(m, { name: "polling", reason: event.reason }, {
|
||||
effects: [{ type: "armPoll", reason: event.reason }],
|
||||
});
|
||||
|
||||
case "STREAM_START": {
|
||||
// First frame arrived. Adopt the run-fact runId if present. sending ->
|
||||
// streaming; a reconnect/attach that just went live also lands here.
|
||||
const runFact = event.runId ? { runId: event.runId } : m.ctx.runFact;
|
||||
return to(m, { name: "streaming" }, {
|
||||
ctx: { runFact },
|
||||
effects: [{ type: "cancelReconnect" }, { type: "disarmPoll" }],
|
||||
});
|
||||
}
|
||||
|
||||
case "FINISH_CLEAN":
|
||||
// A clean terminal outcome. The run is done — clear the run-fact and go
|
||||
// idle. (The queue flush is a component concern gated by ownership; the
|
||||
// FSM only models the phase.) Review #2: reset ownership to local so a
|
||||
// just-finished observer-attach turn re-exposes "Send now" for the queue.
|
||||
return to(m, { name: "idle" }, {
|
||||
ctx: { runFact: null, liveFollow: false, ownership: "local" },
|
||||
effects: [{ type: "disarmPoll" }, { type: "cancelReconnect" }],
|
||||
});
|
||||
|
||||
case "FINISH_ABORT":
|
||||
// A user Stop / intentional abort finished. If we were stopping, the
|
||||
// terminal data has now arrived (I4) — go idle. The run-fact is cleared.
|
||||
return to(m, { name: "idle" }, {
|
||||
ctx: { runFact: null, liveFollow: false, ownership: "local" },
|
||||
effects: [{ type: "disarmPoll" }, { type: "cancelReconnect" }],
|
||||
});
|
||||
|
||||
case "FINISH_DISCONNECT":
|
||||
// A LIVE SSE drop. Recovery depends on WHO we are (I2 + liveFollow):
|
||||
// - a mount-attach OBSERVER (a one-shot resume, NOT live-follow) that drops
|
||||
// -> the degraded poll drives the row to terminal from the DB.
|
||||
if (m.ctx.ownership === "observer" && !m.ctx.liveFollow) {
|
||||
return to(m, { name: "polling", reason: "disconnect-visible" }, {
|
||||
effects: [{ type: "armPoll", reason: "disconnect-visible" }],
|
||||
});
|
||||
}
|
||||
// - a LOCAL live turn (first drop) OR a live-follow re-attach (a SUBSEQUENT
|
||||
// drop) -> (re-)enter the reconnect ladder. #488 commit 3: allowed
|
||||
// REPEATEDLY — `liveFollow` is kept across a successful re-attach, so the
|
||||
// second break reconnects again instead of falling to silent poll.
|
||||
// #488 commit 2: gated on the RUN-FACT (or an existing live-follow), NOT on
|
||||
// the presence of an assistant message — a setup-phase break still recovers.
|
||||
// - visible content already on screen -> keep it, ALSO poll to terminal
|
||||
// (a full replay could clobber the fuller live tail);
|
||||
// - no visible content -> the reconnect ladder rebuilds it.
|
||||
if (m.ctx.runFact || m.ctx.liveFollow) {
|
||||
const effects: Effect[] = [
|
||||
{ type: "scheduleReconnect", attempt: 1, delayMs: reconnectDelayMs(1) },
|
||||
];
|
||||
if (event.hasVisibleContent) effects.push({ type: "armPoll", reason: "disconnect-visible" });
|
||||
return command(m, { name: "reconnecting", attempt: 1, failed: false }, effects, {
|
||||
ownership: "observer",
|
||||
liveFollow: true,
|
||||
});
|
||||
}
|
||||
// No run to recover: a plain disconnect. Surface the terminal notice.
|
||||
return to(m, { name: "idle" }, {
|
||||
ctx: { runFact: null, liveFollow: false, ownership: "local" },
|
||||
});
|
||||
|
||||
case "FINISH_ERROR":
|
||||
return to(m, { name: "error", kind: event.kind }, {
|
||||
ctx: { runFact: null, liveFollow: false, ownership: "local" },
|
||||
effects: [{ type: "disarmPoll" }, { type: "cancelReconnect" }],
|
||||
});
|
||||
|
||||
// ---- mount attach (resume) ----------------------------------------
|
||||
case "ATTACH_START":
|
||||
// A reopened tab attaches to a still-running run: observer ownership.
|
||||
// #488 F2: ONLY from idle. The mount `getRun` round-trip resolves async, and
|
||||
// a local send may have started meanwhile (phase `sending`, ownership local);
|
||||
// a late ATTACH_START must NOT hijack that local turn into an observer-attach
|
||||
// (queue would stop flushing, "Send now" would hide). Guarding in the reducer
|
||||
// covers every dispatch source.
|
||||
if (m.phase.name !== "idle") return stay(m);
|
||||
return command(m, { name: "attaching" }, [{ type: "resumeStream" }], {
|
||||
ownership: "observer",
|
||||
runFact: event.runId ? { runId: event.runId } : m.ctx.runFact,
|
||||
});
|
||||
|
||||
case "ATTACH_LIVE":
|
||||
// The attach GET returned a live 2xx stream — follow it as an observer.
|
||||
// Review #1: guard by SOURCE phase. The epoch filter alone is not enough — a
|
||||
// POLL_TERMINAL uses to() (no epoch bump) and does not abort the in-flight
|
||||
// GET, so a slow 2xx landing after the machine already left `attaching` (e.g.
|
||||
// the armed poll saw the terminal row -> idle) would resurrect a settled run
|
||||
// into a phantom `streaming`. Only enter streaming FROM `attaching`.
|
||||
if (m.phase.name !== "attaching") return stay(m);
|
||||
return to(m, { name: "streaming" });
|
||||
|
||||
case "ATTACH_NONE":
|
||||
// 204 / non-2xx / throw: nothing live to attach. Arm the degraded poll to
|
||||
// follow the run to terminal from the DB. This is a soft-negative run-fact
|
||||
// (204 on a non-stripped path is authoritative-negative; the runtime may
|
||||
// pass a RUN_FACT null separately). Keep the run-fact as-is here.
|
||||
// Review #1: guard by source phase for consistency (a late outcome after the
|
||||
// machine already left `attaching` must not re-arm a poll).
|
||||
if (m.phase.name !== "attaching") return stay(m);
|
||||
return to(m, { name: "polling", reason: "attach-none" }, {
|
||||
effects: [{ type: "armPoll", reason: "attach-none" }],
|
||||
});
|
||||
|
||||
// ---- reconnect after a live disconnect ----------------------------
|
||||
case "RECONNECT_ATTEMPT":
|
||||
// A scheduled backoff fired — fire the attach GET. epoch++ so the previous
|
||||
// attempt's late outcome cannot drive this one.
|
||||
if (m.phase.name !== "reconnecting") return stay(m);
|
||||
return command(
|
||||
m,
|
||||
{ name: "reconnecting", attempt: event.attempt, failed: false },
|
||||
[{ type: "resumeStream" }],
|
||||
);
|
||||
|
||||
case "RECONNECT_ATTACHED":
|
||||
// #488 commit 3: a live re-attach succeeded. Reset to streaming — the
|
||||
// attempt counter is dropped, so a LATER disconnect can start a fresh
|
||||
// ladder from attempt 1 (the old one-shot `!wasResumed` gate forbade a
|
||||
// second cycle, sending the second break to silent poll).
|
||||
// Review #1: guard by SOURCE phase. The armed degraded poll can reach the
|
||||
// terminal row (POLL_TERMINAL -> idle, via to(), NO epoch bump, GET not
|
||||
// aborted) BEFORE a slow reconnect GET returns 2xx; without this guard that
|
||||
// late RECONNECT_ATTACHED (same epoch) would resurrect a settled run into a
|
||||
// phantom `streaming`. Only re-enter streaming FROM `reconnecting`.
|
||||
if (m.phase.name !== "reconnecting") return stay(m);
|
||||
return to(m, { name: "streaming" }, {
|
||||
effects: [{ type: "cancelReconnect" }, { type: "disarmPoll" }],
|
||||
});
|
||||
|
||||
case "RECONNECT_NONE": {
|
||||
// 204 / error during a reconnect attempt. Arm the degraded poll as the
|
||||
// belt-and-suspenders fallback, then either back off to the next attempt
|
||||
// or, at the cap, surface the manual Retry ("failed").
|
||||
if (m.phase.name !== "reconnecting") return stay(m);
|
||||
const attempt = m.phase.attempt;
|
||||
if (attempt < RECONNECT_MAX_ATTEMPTS) {
|
||||
return command(
|
||||
m,
|
||||
{ name: "reconnecting", attempt: attempt + 1, failed: false },
|
||||
[
|
||||
{ type: "armPoll", reason: "attach-none" },
|
||||
{ type: "scheduleReconnect", attempt: attempt + 1, delayMs: reconnectDelayMs(attempt + 1) },
|
||||
],
|
||||
);
|
||||
}
|
||||
return to(m, { name: "reconnecting", attempt, failed: true }, {
|
||||
effects: [{ type: "armPoll", reason: "reconnect-exhausted" }],
|
||||
});
|
||||
}
|
||||
|
||||
case "RETRY":
|
||||
// Manual Retry from the "failed" reconnect banner OR the stalled banner.
|
||||
if (m.phase.name === "reconnecting" && m.phase.failed) {
|
||||
return command(
|
||||
m,
|
||||
{ name: "reconnecting", attempt: 1, failed: false },
|
||||
[{ type: "resumeStream" }],
|
||||
);
|
||||
}
|
||||
if (m.phase.name === "stalled") {
|
||||
// Re-arm the poll to try to catch the run up again.
|
||||
return command(m, { name: "polling", reason: "attach-none" }, [
|
||||
{ type: "armPoll", reason: "attach-none" },
|
||||
]);
|
||||
}
|
||||
return stay(m);
|
||||
|
||||
// ---- degraded poll -------------------------------------------------
|
||||
case "POLL_TERMINAL":
|
||||
// The run reached a terminal row via the poll (or the reconcile merge). Go
|
||||
// idle and disarm everything (I4: this is a DATA-driven exit, incl. exit
|
||||
// from `stopping`). Review #2: reset ownership to local.
|
||||
return to(m, { name: "idle" }, {
|
||||
ctx: { runFact: null, liveFollow: false, ownership: "local" },
|
||||
effects: [{ type: "disarmPoll" }, { type: "cancelReconnect" }],
|
||||
});
|
||||
|
||||
case "POLL_IDLE_CAP":
|
||||
// Review #4: `stopping` also arms the poll (STOP_REQUESTED) but has NO other
|
||||
// backstop — an observer-stop with no SDK stream to fire onFinish, whose
|
||||
// server stop never drives the run terminal, would poll the DB forever. Give
|
||||
// it a bounded exit: cap -> idle + disarm (NOT `stalled`; Stop was already
|
||||
// pressed, so there is nothing for the user to retry).
|
||||
if (m.phase.name === "stopping") {
|
||||
return to(m, { name: "idle" }, {
|
||||
ctx: { runFact: null, liveFollow: false, ownership: "local" },
|
||||
effects: [{ type: "disarmPoll" }, { type: "cancelReconnect" }],
|
||||
});
|
||||
}
|
||||
// #488 commit 4a: the poll hit the inactivity cap. Instead of going SILENT
|
||||
// (the old "forever half-done answer"), surface a stalled banner + Retry.
|
||||
if (m.phase.name !== "polling" && m.phase.name !== "reconnecting") return stay(m);
|
||||
return to(m, { name: "stalled" }, {
|
||||
effects: [{ type: "disarmPoll" }, { type: "cancelReconnect" }],
|
||||
});
|
||||
|
||||
// ---- run-fact ------------------------------------------------------
|
||||
case "RUN_FACT": {
|
||||
const runFact = event.runFact;
|
||||
// A fresh NEGATIVE fact (no active run) cancels recovery immediately (I3):
|
||||
// there is nothing to reconnect to / poll for.
|
||||
if (!runFact) {
|
||||
if (
|
||||
m.phase.name === "reconnecting" ||
|
||||
m.phase.name === "attaching" ||
|
||||
m.phase.name === "polling" ||
|
||||
m.phase.name === "stopping"
|
||||
) {
|
||||
return to(m, { name: "idle" }, {
|
||||
// Review #2: reset ownership to local on this terminal transition.
|
||||
ctx: { runFact: null, liveFollow: false, ownership: "local" },
|
||||
effects: [{ type: "cancelReconnect" }, { type: "disarmPoll" }],
|
||||
});
|
||||
}
|
||||
return to(m, m.phase, { ctx: { runFact: null } });
|
||||
}
|
||||
// A positive fact just updates the context (pessimism toward an attempt: a
|
||||
// stale-but-positive fact permits entering recovery; a 204 will cut it).
|
||||
return to(m, m.phase, { ctx: { runFact } });
|
||||
}
|
||||
|
||||
// ---- stop ----------------------------------------------------------
|
||||
case "STOP_REQUESTED":
|
||||
// Authoritative stop of a detached run. Enter `stopping` and fire stopRun +
|
||||
// abort the local/attach reader. ALSO arm the poll so the terminal row is
|
||||
// observed — the exit is by DATA (I4: a terminal row / negative run-fact),
|
||||
// never by the stopRun HTTP response (which returns after abort, before
|
||||
// finalization). For a local turn the aborted stream's onFinish (ANY finish)
|
||||
// is HONORED in `stopping` at the top of reduce() — regardless of generation
|
||||
// — and exits to idle; the armed poll is the fallback for an observer stop
|
||||
// with no local onFinish.
|
||||
return command(
|
||||
m,
|
||||
{ name: "stopping" },
|
||||
[
|
||||
{ type: "stopRun" },
|
||||
{ type: "abortAttach" },
|
||||
{ type: "cancelReconnect" },
|
||||
{ type: "armPoll", reason: "attach-none" },
|
||||
],
|
||||
);
|
||||
|
||||
// ---- supersede (CAS) ----------------------------------------------
|
||||
case "SUPERSEDE_REQUESTED":
|
||||
// "Interrupt and send now": CAS POST /stream { supersede }. epoch++ so a
|
||||
// late outcome of the interrupted run is dropped.
|
||||
return command(
|
||||
m,
|
||||
{ name: "superseding" },
|
||||
[{ type: "supersede", targetRunId: event.targetRunId }, { type: "cancelReconnect" }, { type: "disarmPoll" }],
|
||||
);
|
||||
|
||||
case "SUPERSEDE_READY": {
|
||||
// CAS succeeded (old run stopped/settled, slot taken, new run begun). We
|
||||
// are now the local streamer of the NEW run. Adopt its runId if provided.
|
||||
const runFact = event.runId ? { runId: event.runId } : m.ctx.runFact;
|
||||
return to(m, { name: "streaming" }, {
|
||||
ctx: { ownership: "local", runFact, liveFollow: false },
|
||||
});
|
||||
}
|
||||
|
||||
case "SUPERSEDE_MISMATCH":
|
||||
// The active run moved between the click and the CAS. Per the spec: verify
|
||||
// via /run rather than blindly banner — the mismatch may be our own already-
|
||||
// superseded run. Surface a classified error AND fire a run-fact verify.
|
||||
return to(m, { name: "error", kind: "supersede-mismatch" }, {
|
||||
ctx: { runFact: event.currentRunId ? { runId: event.currentRunId } : m.ctx.runFact },
|
||||
effects: [{ type: "postRun", reason: "verify" }],
|
||||
});
|
||||
|
||||
case "SUPERSEDE_TIMEOUT":
|
||||
// The old run did not settle within W. Nothing persisted; the composer keeps
|
||||
// its text. Classified error, NO auto-retry (the old client retry ladder is
|
||||
// removed in #488 commit 5).
|
||||
return to(m, { name: "error", kind: "supersede-timeout" });
|
||||
|
||||
case "SUPERSEDE_INVALID":
|
||||
return to(m, { name: "error", kind: "supersede-invalid" });
|
||||
|
||||
case "RUN_ALREADY_ACTIVE":
|
||||
// A plain POST hit the one-active-run gate. NO auto-retry — the composer
|
||||
// offers "interrupt and send" (supersede) instead. #497/S4: adopt the
|
||||
// server's activeRunId as the run-fact so that supersede can TARGET the
|
||||
// (possibly foreign-tab) active run via the CAS, rather than a blind
|
||||
// promote+abort that just 409s again. A stale/absent id keeps the prior fact.
|
||||
return to(m, { name: "error", kind: "run-already-active" }, {
|
||||
ctx: { runFact: event.activeRunId ? { runId: event.activeRunId } : m.ctx.runFact },
|
||||
});
|
||||
|
||||
// ---- lifecycle -----------------------------------------------------
|
||||
case "DISPOSE":
|
||||
// Unmount: abort in-flight controllers, drop timers, and bump the epoch so
|
||||
// NO late callback can drive this (now dead) machine (I5).
|
||||
return command(
|
||||
m,
|
||||
{ name: "idle" },
|
||||
[
|
||||
{ type: "abortAttach" },
|
||||
{ type: "cancelReconnect" },
|
||||
{ type: "disarmPoll" },
|
||||
],
|
||||
{ liveFollow: false },
|
||||
);
|
||||
|
||||
default: {
|
||||
// Exhaustiveness guard.
|
||||
const _never: never = event;
|
||||
void _never;
|
||||
return stay(m);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -181,6 +181,12 @@ export interface IAiChatMessageRow {
|
||||
toolCalls?: unknown;
|
||||
metadata?: {
|
||||
parts?: UIMessage["parts"];
|
||||
// #491 step-alignment anchor: the count of FINISHED steps whose parts are in
|
||||
// THIS row, written atomically with `parts` server-side (flushAssistant). The
|
||||
// resume client reads it as its persisted step frontier N — the tail-only
|
||||
// attach asks the run-stream registry for the frames of step N onward (the
|
||||
// seed already carries steps 0..N-1). Absent on pre-#491 rows -> read as 0.
|
||||
stepsPersisted?: number;
|
||||
// AI SDK v6 `totalUsage` persisted on assistant rows. Legacy cumulative
|
||||
// figure (sum of every step's usage for the turn); kept for back-compat and
|
||||
// as the fallback for older rows that have no `contextTokens`.
|
||||
|
||||
@@ -3,6 +3,7 @@ import {
|
||||
resolveAdoptedChatId,
|
||||
newlyAddedChatIds,
|
||||
extractServerChatId,
|
||||
extractRunId,
|
||||
} from "./adopt-chat-id";
|
||||
|
||||
describe("resolveAdoptedChatId", () => {
|
||||
@@ -70,3 +71,17 @@ describe("extractServerChatId", () => {
|
||||
expect(extractServerChatId(undefined)).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe("extractRunId", () => {
|
||||
it("reads a string runId from the start metadata", () => {
|
||||
expect(extractRunId({ metadata: { runId: "run-1" } })).toBe("run-1");
|
||||
});
|
||||
it("returns undefined when runId is absent", () => {
|
||||
expect(extractRunId({ metadata: { chatId: "c" } })).toBeUndefined();
|
||||
expect(extractRunId({})).toBeUndefined();
|
||||
expect(extractRunId(undefined)).toBeUndefined();
|
||||
});
|
||||
it("returns undefined for a non-string runId", () => {
|
||||
expect(extractRunId({ metadata: { runId: 7 } })).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -56,6 +56,20 @@ export function extractServerChatId(
|
||||
return typeof m?.chatId === "string" ? m.chatId : undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* #488: read the authoritative RUN id off a streaming assistant message. The
|
||||
* server attaches it as `message.metadata.runId` on the `start` part when a run
|
||||
* wraps the turn (see server `chatStreamMetadata`, #184/#487). This is the live
|
||||
* run-fact update the client FSM adopts (mirrors `extractServerChatId`). Returns
|
||||
* it only when it is a string; undefined otherwise.
|
||||
*/
|
||||
export function extractRunId(
|
||||
message: { metadata?: unknown } | undefined,
|
||||
): string | undefined {
|
||||
const m = message?.metadata as { runId?: string } | undefined;
|
||||
return typeof m?.runId === "string" ? m.runId : undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* The deduped set of ids present in `afterIds` but not in `beforeIds`. A
|
||||
* paginated/flatMapped list can repeat the same id, so dedupe: one genuinely-new
|
||||
|
||||
@@ -33,29 +33,44 @@ describe("collapseBlankLines", () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe("collapseBlankLines + renderChatMarkdown (tight reasoning rendering)", () => {
|
||||
it("renders a blank-line-separated list as a TIGHT list (no <li><p>)", () => {
|
||||
describe("collapseBlankLines + renderChatMarkdown (canonical converter)", () => {
|
||||
// Chat markdown now renders through @docmost/prosemirror-markdown (issue #347):
|
||||
// the SAME converter the editor/import use. Its list items are schema-shaped —
|
||||
// each <li>'s content is wrapped in a <p> (listItem content is `paragraph+`) —
|
||||
// so the HTML always carries `<li><p>…</p></li>` regardless of blank-line
|
||||
// looseness in the source (the converter has no tight/loose distinction). The
|
||||
// visual tightness that `collapseBlankLines` used to buy is now provided by
|
||||
// CSS (`.markdown li p { margin: 0 }`), not the HTML shape.
|
||||
it("renders a blank-line-separated bullet list as a real <ul> list", () => {
|
||||
const loose =
|
||||
"Intro paragraph.\n\n- item one\n\n- item two\n\n- item three";
|
||||
const html = renderChatMarkdown(collapseBlankLines(loose), {});
|
||||
// Tight list: each <li> holds the text directly, not wrapped in a <p>.
|
||||
expect(html).toContain("<li>item one</li>");
|
||||
expect(html).not.toContain("<li><p>");
|
||||
// The list still parses as a list after the paragraph (not a paragraph+<br>).
|
||||
// Clean, un-namespaced HTML (DOMSerializer, not XMLSerializer) — no xmlns.
|
||||
expect(html).toContain("<ul>");
|
||||
expect(html).not.toMatch(/<ul[^>]*xmlns/);
|
||||
// The item text is present (inside the schema's <li><p> wrapper).
|
||||
expect(html).toContain("item one");
|
||||
// The intro paragraph renders as its own paragraph before the list.
|
||||
expect(html).toContain("<p>Intro paragraph.</p>");
|
||||
});
|
||||
|
||||
it("renders an ordered list (1. 2.) as tight after collapsing", () => {
|
||||
it("renders an ordered list (1. 2.) as a real <ol> list", () => {
|
||||
const loose = "Intro.\n\n1. first\n\n2. second";
|
||||
const html = renderChatMarkdown(collapseBlankLines(loose), {});
|
||||
expect(html).toContain("<ol>");
|
||||
expect(html).toContain("<li>first</li>");
|
||||
expect(html).not.toContain("<li><p>");
|
||||
expect(html).not.toMatch(/<ol[^>]*xmlns/);
|
||||
expect(html).toContain("first");
|
||||
expect(html).toContain("second");
|
||||
});
|
||||
|
||||
it("the loose source WOULD render <li><p> without collapsing (control)", () => {
|
||||
it("wraps list-item content in <p> (schema shape; tightness is CSS)", () => {
|
||||
// The canonical converter always wraps a list item's content in a paragraph,
|
||||
// whether or not the source had blank lines between items.
|
||||
const loose = "- a\n\n- b";
|
||||
expect(renderChatMarkdown(loose, {})).toContain("<li><p>");
|
||||
// And a "tight" source produces the identical wrapping (no distinction).
|
||||
expect(renderChatMarkdown(collapseBlankLines(loose), {})).toContain(
|
||||
"<li><p>",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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,72 @@ 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",
|
||||
);
|
||||
});
|
||||
|
||||
// #488 commit 5: the #487 concurrency-gate / supersede 409s. FULL real bodies:
|
||||
// a ConflictException(object) whose response is serialized verbatim, carrying a
|
||||
// `code` and statusCode 409. Each must classify to a human text, not raw JSON.
|
||||
it("classifies A_RUN_ALREADY_ACTIVE (409) as already-answering, not raw JSON", () => {
|
||||
const body =
|
||||
'{"message":"A run is already active for this chat","code":"A_RUN_ALREADY_ACTIVE","statusCode":409}';
|
||||
expect(describeChatError(body, t).title).toBe(
|
||||
"The agent is already answering",
|
||||
);
|
||||
// Never leaks the raw code as the detail.
|
||||
expect(describeChatError(body, t).detail).not.toContain("A_RUN_ALREADY_ACTIVE");
|
||||
});
|
||||
|
||||
it("classifies SUPERSEDE_TARGET_MISMATCH (409) as run-changed", () => {
|
||||
// Real server body shape: the current run id is `activeRunId` (NOT `runId`) —
|
||||
// see ai-chat.controller.ts. describeChatError classifies off `code` only.
|
||||
const body =
|
||||
'{"message":"active run does not match the supersede target","code":"SUPERSEDE_TARGET_MISMATCH","activeRunId":"run-x","statusCode":409}';
|
||||
expect(describeChatError(body, t).title).toBe(
|
||||
"Couldn't interrupt — the run changed",
|
||||
);
|
||||
});
|
||||
|
||||
it("classifies SUPERSEDE_TIMEOUT (409) as couldn't-interrupt-in-time", () => {
|
||||
const body =
|
||||
'{"message":"the run did not settle within the supersede window","code":"SUPERSEDE_TIMEOUT","statusCode":409}';
|
||||
expect(describeChatError(body, t).title).toBe("Couldn't interrupt in time");
|
||||
});
|
||||
|
||||
it("classifies SUPERSEDE_INVALID (409) as couldn't-interrupt-that-run", () => {
|
||||
const body =
|
||||
'{"message":"supervise requires chatId","code":"SUPERSEDE_INVALID","statusCode":409}';
|
||||
expect(describeChatError(body, t).title).toBe(
|
||||
"Couldn't interrupt that run",
|
||||
);
|
||||
});
|
||||
|
||||
it("ORDER GUARD: A_RUN_ALREADY_ACTIVE wins over any generic status branch", () => {
|
||||
// Even though the body could superficially look 4xx-ish, the code branch runs
|
||||
// first, so it is never mislabeled by a generic status heading.
|
||||
const body =
|
||||
'{"message":"conflict","code":"A_RUN_ALREADY_ACTIVE","statusCode":409}';
|
||||
const view = describeChatError(body, t);
|
||||
expect(view.title).not.toBe("Something went wrong");
|
||||
expect(view.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,59 @@ 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.",
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
// #488 commit 5: the #487 concurrency-gate / supersede 409s. These arrive as a
|
||||
// ConflictException(object) body carrying a `code` (and statusCode 409). They
|
||||
// MUST be classified by `code` STRICTLY BEFORE any generic status branch, or the
|
||||
// user sees the raw JSON `{"code":"A_RUN_ALREADY_ACTIVE",…}`. The code strings
|
||||
// are the real #487 server contract (ai-chat.controller.ts) — do not invent.
|
||||
if (/"code"\s*:\s*"A_RUN_ALREADY_ACTIVE"/.test(msg)) {
|
||||
return {
|
||||
title: t("The agent is already answering"),
|
||||
detail: t(
|
||||
"This chat already has a run in progress. Wait for it to finish, or interrupt it and send now.",
|
||||
),
|
||||
};
|
||||
}
|
||||
if (/"code"\s*:\s*"SUPERSEDE_TARGET_MISMATCH"/.test(msg)) {
|
||||
return {
|
||||
title: t("Couldn't interrupt — the run changed"),
|
||||
detail: t(
|
||||
"The run you tried to interrupt is no longer the active one. Check the latest answer and try again.",
|
||||
),
|
||||
};
|
||||
}
|
||||
if (/"code"\s*:\s*"SUPERSEDE_TIMEOUT"/.test(msg)) {
|
||||
return {
|
||||
title: t("Couldn't interrupt in time"),
|
||||
detail: t(
|
||||
"The previous run didn't stop in time. Nothing was sent — try sending again.",
|
||||
),
|
||||
};
|
||||
}
|
||||
if (/"code"\s*:\s*"SUPERSEDE_INVALID"/.test(msg)) {
|
||||
return {
|
||||
title: t("Couldn't interrupt that run"),
|
||||
detail: t(
|
||||
"The run to interrupt doesn't belong to this chat. Reload and try again.",
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
if (/"statusCode"\s*:\s*403\b/.test(msg)) {
|
||||
return {
|
||||
title: t("AI chat is disabled"),
|
||||
|
||||
@@ -1,6 +1,37 @@
|
||||
import { markdownToHtml } from "@docmost/editor-ext";
|
||||
import {
|
||||
markdownToProseMirrorSync,
|
||||
docmostExtensions,
|
||||
} from "@docmost/prosemirror-markdown/browser";
|
||||
import { getSchema } from "@tiptap/core";
|
||||
import { Node as PMNode, DOMSerializer } from "@tiptap/pm/model";
|
||||
import DOMPurify from "dompurify";
|
||||
|
||||
// The Docmost editor schema, built once. Chat markdown is rendered through the
|
||||
// SAME schema the editor/import use (issue #347), so chat output matches how the
|
||||
// page would render the same markdown.
|
||||
const chatSchema = getSchema(docmostExtensions);
|
||||
|
||||
/**
|
||||
* Markdown -> HTML for chat display, via the canonical converter. We serialize
|
||||
* the ProseMirror doc with `DOMSerializer` into a real element and read its
|
||||
* `innerHTML` (rather than `@tiptap/html`'s `generateHTML`, whose browser path
|
||||
* uses `XMLSerializer` and stamps a `xmlns` on every block) so the markup is
|
||||
* clean HTML. `li > p` wrapping is inherent to the schema (listItem content is
|
||||
* `paragraph+`); the chat CSS zeroes those paragraph margins so lists still
|
||||
* render tight.
|
||||
*/
|
||||
function markdownToChatHtml(markdown: string): string {
|
||||
const doc = markdownToProseMirrorSync(markdown);
|
||||
const node = PMNode.fromJSON(chatSchema, doc);
|
||||
const div = document.createElement("div");
|
||||
DOMSerializer.fromSchema(chatSchema).serializeFragment(
|
||||
node.content,
|
||||
{ document },
|
||||
div,
|
||||
);
|
||||
return div.innerHTML;
|
||||
}
|
||||
|
||||
export interface RenderChatMarkdownOptions {
|
||||
/**
|
||||
* Neutralize INTERNAL links so they render as inert text (no `href`/`target`).
|
||||
@@ -63,22 +94,32 @@ function neutralizeInternalLinksHook(node: Element): void {
|
||||
|
||||
/**
|
||||
* Render AI markdown to sanitized HTML for read-only display. We reuse the
|
||||
* app's `markdownToHtml` (the same `marked` pipeline used for paste/import) so
|
||||
* chat output matches the editor's markdown flavor, then sanitize with
|
||||
* DOMPurify — LLM output is untrusted, so it must never reach the DOM unsanitized.
|
||||
* canonical converter (issue #347): markdown -> ProseMirror JSON (the SAME
|
||||
* `markdownToProseMirrorSync` the editor paste/import path uses, so chat output
|
||||
* matches the editor's markdown flavor) -> HTML via `markdownToChatHtml`
|
||||
* (DOMSerializer), then sanitize with DOMPurify — LLM output is untrusted, so it
|
||||
* must never reach the DOM unsanitized.
|
||||
*
|
||||
* `markdownToHtml` can return `string | Promise<string>` (it has async marked
|
||||
* extensions registered). In practice plain chat markdown resolves
|
||||
* synchronously, but we guard the Promise case by returning a safe empty string
|
||||
* for that branch (the caller renders the raw text fallback instead).
|
||||
* Stays SYNCHRONOUS: both callers render inside React (a memo and a useMemo),
|
||||
* so the whole pipeline must resolve without awaiting. The converter's sync
|
||||
* entry makes that possible; on any conversion error we return "" so the caller
|
||||
* falls back to raw text (the same fallback the old Promise-guard produced).
|
||||
*/
|
||||
export function renderChatMarkdown(
|
||||
markdown: string,
|
||||
options: RenderChatMarkdownOptions = {},
|
||||
): string {
|
||||
if (!markdown) return "";
|
||||
const html = markdownToHtml(markdown);
|
||||
if (typeof html !== "string") return "";
|
||||
let html: string;
|
||||
try {
|
||||
// markdown -> canonical PM JSON -> HTML (native DOMParser in the browser;
|
||||
// jsdom is never bundled — see @docmost/prosemirror-markdown/browser).
|
||||
html = markdownToChatHtml(markdown);
|
||||
} catch {
|
||||
// Malformed/unsupported markdown must not crash the chat render; fall back
|
||||
// to raw text (empty return -> caller shows the plain-text branch).
|
||||
return "";
|
||||
}
|
||||
|
||||
if (!options.neutralizeInternalLinks) {
|
||||
// Internal chat: unchanged behavior, no hook registered.
|
||||
|
||||
@@ -4,7 +4,8 @@ import type { IAiChatMessageRow } from "@/features/ai-chat/types/ai-chat.types.t
|
||||
import {
|
||||
isStreamingTail,
|
||||
isSettledAssistantTail,
|
||||
seedRows,
|
||||
stepsPersistedOf,
|
||||
mergeDeltaRowsIntoPages,
|
||||
mergeById,
|
||||
} from "./resume-helpers.ts";
|
||||
|
||||
@@ -12,8 +13,18 @@ function row(
|
||||
id: string,
|
||||
role: string,
|
||||
status?: string,
|
||||
stepsPersisted?: number,
|
||||
): IAiChatMessageRow {
|
||||
return { id, role, content: "", status, createdAt: "2026-01-01T00:00:00Z" };
|
||||
return {
|
||||
id,
|
||||
role,
|
||||
content: "",
|
||||
status,
|
||||
createdAt: "2026-01-01T00:00:00Z",
|
||||
...(stepsPersisted !== undefined
|
||||
? { metadata: { stepsPersisted } }
|
||||
: {}),
|
||||
};
|
||||
}
|
||||
|
||||
function makeMsg(id: string, text: string): UIMessage {
|
||||
@@ -65,23 +76,92 @@ describe("isSettledAssistantTail", () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe("seedRows", () => {
|
||||
const rows = [row("u1", "user"), row("a1", "assistant", "streaming")];
|
||||
|
||||
it("returns the rows unchanged when not stripping", () => {
|
||||
expect(seedRows(rows, false)).toBe(rows);
|
||||
describe("stepsPersistedOf", () => {
|
||||
it("reads metadata.stepsPersisted", () => {
|
||||
expect(stepsPersistedOf(row("a1", "assistant", "streaming", 3))).toBe(3);
|
||||
expect(stepsPersistedOf(row("a1", "assistant", "streaming", 0))).toBe(0);
|
||||
});
|
||||
|
||||
it("drops the last row when stripping", () => {
|
||||
const seeded = seedRows(rows, true);
|
||||
expect(seeded).toHaveLength(1);
|
||||
expect(seeded[0].id).toBe("u1");
|
||||
it("defaults to 0 for a pre-#491 row (absent), null/undefined, or a bad value", () => {
|
||||
expect(stepsPersistedOf(row("a1", "assistant", "streaming"))).toBe(0);
|
||||
expect(stepsPersistedOf(null)).toBe(0);
|
||||
expect(stepsPersistedOf(undefined)).toBe(0);
|
||||
expect(
|
||||
stepsPersistedOf({
|
||||
id: "a1",
|
||||
role: "assistant",
|
||||
content: "",
|
||||
createdAt: "x",
|
||||
metadata: { stepsPersisted: -2 },
|
||||
}),
|
||||
).toBe(0);
|
||||
});
|
||||
|
||||
it("returns an empty list when stripping a single-row list", () => {
|
||||
expect(seedRows([row("a1", "assistant", "streaming")], true)).toHaveLength(
|
||||
0,
|
||||
);
|
||||
it("floors a non-integer count", () => {
|
||||
expect(
|
||||
stepsPersistedOf({
|
||||
id: "a1",
|
||||
role: "assistant",
|
||||
content: "",
|
||||
createdAt: "x",
|
||||
metadata: { stepsPersisted: 2.9 },
|
||||
}),
|
||||
).toBe(2);
|
||||
});
|
||||
});
|
||||
|
||||
describe("mergeDeltaRowsIntoPages", () => {
|
||||
const pages = () => [
|
||||
{ items: [row("u1", "user"), row("a1", "assistant", "streaming", 1)], meta: {} },
|
||||
];
|
||||
|
||||
it("returns the pages unchanged for an empty delta", () => {
|
||||
const p = pages();
|
||||
expect(mergeDeltaRowsIntoPages(p, [])).toBe(p);
|
||||
});
|
||||
|
||||
it("appends a genuinely new row to the last page in chronological order", () => {
|
||||
const merged = mergeDeltaRowsIntoPages(pages(), [row("a2", "assistant", "streaming", 0)]);
|
||||
expect(merged[0].items.map((i) => i.id)).toEqual(["u1", "a1", "a2"]);
|
||||
});
|
||||
|
||||
it("replaces a grown row in place (per-step growth), never appends a duplicate", () => {
|
||||
const merged = mergeDeltaRowsIntoPages(pages(), [
|
||||
row("a1", "assistant", "streaming", 2),
|
||||
]);
|
||||
expect(merged[0].items.map((i) => i.id)).toEqual(["u1", "a1"]);
|
||||
// the in-place replacement carries the grown step frontier.
|
||||
expect(stepsPersistedOf(merged[0].items[1])).toBe(2);
|
||||
});
|
||||
|
||||
it("does not mutate the input pages", () => {
|
||||
const input = pages();
|
||||
const before = input[0].items.slice();
|
||||
mergeDeltaRowsIntoPages(input, [row("a2", "assistant", "streaming", 0)]);
|
||||
expect(input[0].items).toEqual(before); // untouched
|
||||
});
|
||||
|
||||
// #491 CONTRACT: the delta overlap window re-delivers the same rows, so merging
|
||||
// MUST be idempotent — applying a delta twice equals applying it once (no growth,
|
||||
// no reorder). A regression re-introduces duplicate assistant bubbles per poll.
|
||||
it("is idempotent: applying the same delta twice equals once", () => {
|
||||
const delta = [
|
||||
row("a1", "assistant", "streaming", 2), // grown existing row
|
||||
row("a2", "assistant", "streaming", 0), // new row
|
||||
];
|
||||
const once = mergeDeltaRowsIntoPages(pages(), delta);
|
||||
const twice = mergeDeltaRowsIntoPages(once, delta);
|
||||
const thrice = mergeDeltaRowsIntoPages(twice, delta);
|
||||
expect(once[0].items.map((i) => i.id)).toEqual(["u1", "a1", "a2"]);
|
||||
expect(twice[0].items.map((i) => i.id)).toEqual(["u1", "a1", "a2"]);
|
||||
expect(twice).toEqual(once);
|
||||
expect(thrice).toEqual(once);
|
||||
});
|
||||
|
||||
it("seeds a first page when the cache is empty", () => {
|
||||
const merged = mergeDeltaRowsIntoPages([], [row("u1", "user")]);
|
||||
expect(merged).toHaveLength(1);
|
||||
expect(merged[0].items.map((i) => i.id)).toEqual(["u1"]);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -109,4 +189,37 @@ describe("mergeById", () => {
|
||||
expect(mergeById(prev, null)).toBe(prev);
|
||||
expect(mergeById(prev, undefined)).toBe(prev);
|
||||
});
|
||||
|
||||
// #491 CONTRACT: the delta poll's overlap window GUARANTEES the same row is
|
||||
// re-delivered across close polls, so merging must be IDEMPOTENT by id — merging
|
||||
// the same row (or an equal-length list of rows) twice must not duplicate or
|
||||
// reorder. This is the property the whole delta-poll design leans on; a
|
||||
// regression here would re-introduce duplicate assistant bubbles on every poll.
|
||||
it("is idempotent by id: re-merging the same row does not duplicate or reorder", () => {
|
||||
const seed = [makeMsg("u1", "hi"), makeMsg("a1", "step 1")];
|
||||
const repeat = makeMsg("a1", "step 1"); // the SAME row the overlap re-delivers
|
||||
const once = mergeById(seed, repeat);
|
||||
const twice = mergeById(once, repeat);
|
||||
const thrice = mergeById(twice, repeat);
|
||||
// Length is stable (no growth), order is stable (user then assistant).
|
||||
expect(once.map((m) => m.id)).toEqual(["u1", "a1"]);
|
||||
expect(twice.map((m) => m.id)).toEqual(["u1", "a1"]);
|
||||
expect(thrice.map((m) => m.id)).toEqual(["u1", "a1"]);
|
||||
// The repeated merge converges: the row is replaced in place, never appended.
|
||||
expect(twice[1]).toBe(repeat);
|
||||
});
|
||||
|
||||
it("is idempotent across a batch of repeated + grown rows (delta re-delivery)", () => {
|
||||
// A delta poll re-delivers a1 (unchanged) and a2 (grown one step). Applying the
|
||||
// batch twice must equal applying it once — the poll can re-send either.
|
||||
const start = [makeMsg("u1", "hi"), makeMsg("a1", "done")];
|
||||
const batch = [makeMsg("a1", "done"), makeMsg("a2", "grown step 2")];
|
||||
const apply = (list: typeof start) =>
|
||||
batch.reduce((acc, row) => mergeById(acc, row), list);
|
||||
const once = apply(start);
|
||||
const twice = apply(once);
|
||||
expect(once.map((m) => m.id)).toEqual(["u1", "a1", "a2"]);
|
||||
expect(twice.map((m) => m.id)).toEqual(["u1", "a1", "a2"]);
|
||||
expect(twice).toEqual(once);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -11,9 +11,10 @@ import type { IAiChatMessageRow } from "@/features/ai-chat/types/ai-chat.types.t
|
||||
|
||||
/**
|
||||
* A STREAMING tail: the last persisted row is an assistant row still marked
|
||||
* `status === 'streaming'`. Such a tail is stripped from the seed and rebuilt by
|
||||
* the replay (`expect=live`), since the SDK's `text-start` always pushes a new
|
||||
* part and replaying over a seeded in-progress row would duplicate its text.
|
||||
* `status === 'streaming'`. #491 (tail-only): such a tail is seeded UNCHANGED —
|
||||
* it carries the persisted steps 0..N-1 — and the run-stream registry's tail
|
||||
* (frames for steps >= N) is APPENDED to it by the SDK's `readUIMessageStream`
|
||||
* continuation. Only the presence of this tail decides WHETHER to attach.
|
||||
*/
|
||||
export function isStreamingTail(rows: IAiChatMessageRow[]): boolean {
|
||||
const tail = rows[rows.length - 1];
|
||||
@@ -32,15 +33,61 @@ export function isSettledAssistantTail(rows: IAiChatMessageRow[]): boolean {
|
||||
}
|
||||
|
||||
/**
|
||||
* Seed rows for `useChat`: return the rows unchanged, or without the last row when
|
||||
* `strip` is set (the streaming tail is stripped so the live replay rebuilds it
|
||||
* without duplicating parts).
|
||||
* #491 tail-only anchor: the count of FINISHED steps whose parts are persisted in
|
||||
* THIS assistant row (`metadata.stepsPersisted`), written atomically with `parts`
|
||||
* server-side. The resume client reads it as its persisted step frontier N — the
|
||||
* tail-only attach asks the run-stream registry for the frames of step N onward
|
||||
* (the seed already carries steps 0..N-1). Absent on pre-#491 rows => 0.
|
||||
*/
|
||||
export function seedRows(
|
||||
export function stepsPersistedOf(
|
||||
row: IAiChatMessageRow | null | undefined,
|
||||
): number {
|
||||
const n = row?.metadata?.stepsPersisted;
|
||||
return typeof n === "number" && n >= 0 ? Math.floor(n) : 0;
|
||||
}
|
||||
|
||||
/** One page of the messages infinite-query cache (`{ items, meta }`). */
|
||||
export interface IMessagePage {
|
||||
items: IAiChatMessageRow[];
|
||||
meta: unknown;
|
||||
}
|
||||
|
||||
/**
|
||||
* #491 delta-poll merge: upsert the delta poll's `rows` into the messages
|
||||
* infinite-query page structure IDEMPOTENTLY by id. The delta endpoint's overlap
|
||||
* window GUARANTEES occasional REPEATS, so this MUST converge: a row already
|
||||
* present is REPLACED IN PLACE (per-step growth of an in-progress row), a new row
|
||||
* is APPENDED to the last page in chronological order (the server returns delta
|
||||
* rows oldest-first). Applying the same delta twice equals applying it once. Never
|
||||
* mutates the input pages (returns fresh page objects with cloned item arrays).
|
||||
*/
|
||||
export function mergeDeltaRowsIntoPages(
|
||||
pages: IMessagePage[],
|
||||
rows: IAiChatMessageRow[],
|
||||
strip: boolean,
|
||||
): IAiChatMessageRow[] {
|
||||
return strip ? rows.slice(0, -1) : rows;
|
||||
): IMessagePage[] {
|
||||
if (rows.length === 0) return pages;
|
||||
const next: IMessagePage[] = pages.map((p) => ({
|
||||
...p,
|
||||
items: p.items.slice(),
|
||||
}));
|
||||
const locate = (id: string): [number, number] | null => {
|
||||
for (let pi = 0; pi < next.length; pi++) {
|
||||
const ii = next[pi].items.findIndex((it) => it.id === id);
|
||||
if (ii !== -1) return [pi, ii];
|
||||
}
|
||||
return null;
|
||||
};
|
||||
for (const row of rows) {
|
||||
const at = locate(row.id);
|
||||
if (at) {
|
||||
next[at[0]].items[at[1]] = row; // replace in place — idempotent by id
|
||||
} else if (next.length > 0) {
|
||||
next[next.length - 1].items.push(row); // append chronologically
|
||||
} else {
|
||||
next.push({ items: [row], meta: undefined });
|
||||
}
|
||||
}
|
||||
return next;
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { readUIMessageStream, type UIMessage } from "ai";
|
||||
import pkg from "../../../../package.json";
|
||||
|
||||
/**
|
||||
* PIN-SPEC TRIP-WIRE (#491). The tail-only attach continuation relies on THREE
|
||||
* behaviors of `ai@6.0.207`, verified line-by-line in the issue. Without this
|
||||
* test, an `ai` bump could silently break attach (the client would append the
|
||||
* live tail to the wrong message, or duplicate a step):
|
||||
*
|
||||
* 1. `readUIMessageStream({ message })` CONTINUES the passed message — it does
|
||||
* not start a fresh one — so the tail streamed after a re-seed is appended to
|
||||
* the seeded assistant row (the same DB id).
|
||||
* 2. A `start` frame does NOT reset the existing message's parts (so the seeded
|
||||
* steps 0..N-1 survive; the synthetic `start` the registry prepends only
|
||||
* carries the run-fact metadata).
|
||||
* 3. Text parts do NOT cross a `finish-step` boundary — a new `text-start` after
|
||||
* `finish-step` is a NEW part — so the reconstructed steps stay separated and
|
||||
* the step frontier stays meaningful.
|
||||
*
|
||||
* If an `ai` upgrade changes any of these, this test fails LOUD instead of the
|
||||
* resume path silently corrupting.
|
||||
*/
|
||||
describe("ai SDK continuation trip-wire (#491, tail-only attach)", () => {
|
||||
it("is pinned to the exact ai version the continuation was verified against", () => {
|
||||
// A caret/range bump is exactly what would silently break attach — require an
|
||||
// exact pin. Bumping ai MUST re-verify the behavior asserted below, then this.
|
||||
expect((pkg as { dependencies: Record<string, string> }).dependencies.ai).toBe(
|
||||
"6.0.207",
|
||||
);
|
||||
});
|
||||
|
||||
it("continues the seeded message: start does not reset parts, the tail appends as new parts", async () => {
|
||||
// A seeded assistant row with ONE finished step already reconstructed.
|
||||
const seeded: UIMessage = {
|
||||
id: "assistant-1",
|
||||
role: "assistant",
|
||||
parts: [
|
||||
{ type: "step-start" },
|
||||
{ type: "text", text: "STEP0", state: "done" },
|
||||
],
|
||||
} as UIMessage;
|
||||
|
||||
// The tail the registry delivers on re-attach: a synthetic start (run-fact),
|
||||
// then step 1's frames, then finish. As UI-message chunks (what the SSE frames
|
||||
// decode to).
|
||||
const chunks = [
|
||||
{ type: "start", messageMetadata: { runId: "r1", chatId: "c1" } },
|
||||
{ type: "start-step" },
|
||||
{ type: "text-start", id: "t1" },
|
||||
{ type: "text-delta", id: "t1", delta: "STEP1" },
|
||||
{ type: "text-end", id: "t1" },
|
||||
{ type: "finish-step" },
|
||||
{ type: "finish" },
|
||||
];
|
||||
const stream = new ReadableStream({
|
||||
start(c) {
|
||||
for (const ch of chunks) c.enqueue(ch);
|
||||
c.close();
|
||||
},
|
||||
});
|
||||
|
||||
let last: UIMessage | undefined;
|
||||
for await (const msg of readUIMessageStream({ message: seeded, stream })) {
|
||||
last = msg;
|
||||
}
|
||||
|
||||
expect(last).toBeDefined();
|
||||
// Same message id (continuation, not a fresh message).
|
||||
expect(last!.id).toBe("assistant-1");
|
||||
// The seeded step-0 parts SURVIVED the `start` frame, and step 1 was appended
|
||||
// as SEPARATE parts (text did not cross the finish-step boundary).
|
||||
const shape = last!.parts.map((p) => `${p.type}:${(p as { text?: string }).text ?? ""}`);
|
||||
expect(shape).toEqual([
|
||||
"step-start:",
|
||||
"text:STEP0",
|
||||
"step-start:",
|
||||
"text:STEP1",
|
||||
]);
|
||||
// The run-fact metadata from the synthetic start frame is applied.
|
||||
expect(last!.metadata).toMatchObject({ runId: "r1", chatId: "c1" });
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,206 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { Editor } from "@tiptap/core";
|
||||
import { Document } from "@tiptap/extension-document";
|
||||
import { Paragraph } from "@tiptap/extension-paragraph";
|
||||
import { Text } from "@tiptap/extension-text";
|
||||
import { Bold } from "@tiptap/extension-bold";
|
||||
import { Italic } from "@tiptap/extension-italic";
|
||||
import { MarkdownClipboard } from "./markdown-clipboard";
|
||||
|
||||
/**
|
||||
* Integration coverage for the async `handlePaste` seam (issue #347). The paste
|
||||
* conversion moved to `@docmost/prosemirror-markdown`'s browser entry, whose
|
||||
* `markdownToProseMirror` is async — so `handlePaste` captures the range, claims
|
||||
* the event (returns true), and dispatches the insert on the next microtask.
|
||||
* These tests drive that path end to end on a minimal schema (a plain-markdown
|
||||
* paste whose converted nodes fit paragraph/text/bold/italic), asserting the
|
||||
* text lands with the right marks and that the raw markdown syntax is consumed
|
||||
* (recognized as markdown, not inserted literally).
|
||||
*/
|
||||
|
||||
function makeEditor() {
|
||||
const element = document.createElement("div");
|
||||
document.body.appendChild(element);
|
||||
return new Editor({
|
||||
element,
|
||||
extensions: [
|
||||
Document,
|
||||
Paragraph,
|
||||
Text,
|
||||
Bold,
|
||||
Italic,
|
||||
MarkdownClipboard.configure({ transformPastedText: true }),
|
||||
],
|
||||
content: { type: "doc", content: [{ type: "paragraph" }] },
|
||||
});
|
||||
}
|
||||
|
||||
// Locate the markdownClipboard plugin and invoke its handlePaste directly with a
|
||||
// synthetic clipboard event (jsdom has no real paste pipeline). The plugin's
|
||||
// handlePaste closes over the extension `this`, so calling it off the plugin
|
||||
// props preserves `this.editor`/`this.options`.
|
||||
function paste(editor: Editor, text: string): boolean {
|
||||
const view = editor.view;
|
||||
const plugin = view.state.plugins.find(
|
||||
(p: any) => p.props && p.spec?.key,
|
||||
) as any;
|
||||
const event = {
|
||||
clipboardData: {
|
||||
getData: (type: string) => (type === "text/plain" ? text : ""),
|
||||
},
|
||||
} as unknown as ClipboardEvent;
|
||||
// Find the specific handlePaste that belongs to the markdown clipboard plugin.
|
||||
const md = view.state.plugins.find(
|
||||
(p: any) => typeof p.props?.handlePaste === "function",
|
||||
) as any;
|
||||
return md.props.handlePaste(view, event, view.state.selection.content());
|
||||
}
|
||||
|
||||
// Flush the microtask queue so the async .then() dispatch runs.
|
||||
const flush = () => new Promise((r) => setTimeout(r, 0));
|
||||
|
||||
describe("MarkdownClipboard handlePaste (async md -> PM)", () => {
|
||||
it("converts a plain-markdown paste with bold/italic into marked text", async () => {
|
||||
const editor = makeEditor();
|
||||
const claimed = paste(editor, "hello **bold** and *italic*");
|
||||
// The paste is claimed synchronously (async insert follows).
|
||||
expect(claimed).toBe(true);
|
||||
await flush();
|
||||
|
||||
const json = editor.getJSON();
|
||||
const text = JSON.stringify(json);
|
||||
// The raw markdown asterisks are consumed (recognized), not inserted literally.
|
||||
expect(editor.getText()).not.toContain("**");
|
||||
expect(editor.getText()).toContain("bold");
|
||||
expect(editor.getText()).toContain("italic");
|
||||
// The bold/italic marks materialized.
|
||||
expect(text).toContain('"bold"');
|
||||
expect(text).toContain('"italic"');
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("recognizes a bullet list paste as list structure (not literal '-')", async () => {
|
||||
// A bullet list is not representable in this minimal schema, so the converter
|
||||
// output would fail PMNode.fromJSON and the catch inserts raw text. Use a
|
||||
// paste whose nodes DO fit the schema to assert the happy path instead: two
|
||||
// paragraphs separated by a blank line.
|
||||
const editor = makeEditor();
|
||||
paste(editor, "first para\n\nsecond para");
|
||||
await flush();
|
||||
const json = editor.getJSON() as any;
|
||||
const paras = (json.content || []).filter(
|
||||
(n: any) => n.type === "paragraph",
|
||||
);
|
||||
// Two paragraphs materialized from the blank-line-separated markdown.
|
||||
expect(paras.length).toBeGreaterThanOrEqual(2);
|
||||
expect(editor.getText()).toContain("first para");
|
||||
expect(editor.getText()).toContain("second para");
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("falls back to raw text when conversion yields nodes the schema lacks", async () => {
|
||||
// `# heading` converts to a `heading` node absent from this minimal schema,
|
||||
// so PMNode.fromJSON throws and the catch re-inserts the raw text — the user
|
||||
// never loses their clipboard content.
|
||||
const editor = makeEditor();
|
||||
paste(editor, "# a heading line");
|
||||
await flush();
|
||||
// Content is preserved (either as heading text or literal), never dropped.
|
||||
expect(editor.getText()).toContain("a heading line");
|
||||
editor.destroy();
|
||||
});
|
||||
});
|
||||
|
||||
// The async seam captures the target range synchronously, then replaces on the
|
||||
// next microtask. If the document changed under it between capture and resolve
|
||||
// (impossible in prod — same microtask — but pinned here), BOTH the success
|
||||
// (replaceRange) and the fail-open (insertText) branches must fall back to the
|
||||
// LIVE selection rather than a stale absolute range, so neither clobbers content
|
||||
// nor throws a RangeError. We force the mid-flight change by dispatching a
|
||||
// doc-mutating transaction AFTER the synchronous claim but BEFORE flushing the
|
||||
// microtask that runs the `.then`/`.catch`.
|
||||
describe("MarkdownClipboard handlePaste — doc-changed-mid-flight guard", () => {
|
||||
// Replace the whole doc with one paragraph of `text` (synchronous dispatch).
|
||||
// An empty string yields an empty paragraph (a text node may not be empty).
|
||||
function seedContent(editor: Editor, text: string) {
|
||||
editor.commands.setContent({
|
||||
type: "doc",
|
||||
content: [
|
||||
text
|
||||
? { type: "paragraph", content: [{ type: "text", text }] }
|
||||
: { type: "paragraph" },
|
||||
],
|
||||
});
|
||||
}
|
||||
|
||||
it("success branch: mid-flight doc change routes the paste to the LIVE selection, never the stale range (clobber-proving)", async () => {
|
||||
// The paste captures a NON-EMPTY range {1,5} (over "AAAA"). Then, before the
|
||||
// async resolve, the doc GROWS ("MARKER" inserted at the start) and the cursor
|
||||
// is parked at the doc END. The captured {1,5} is now stale and points INTO
|
||||
// "MARKER". A WORKING guard replaces at the live (end) selection → MARKER is
|
||||
// untouched. A BROKEN guard replaces the stale {1,5} → it erases the first
|
||||
// characters of MARKER (this is what a zero-width `from==to` range could never
|
||||
// reveal, which is why the earlier version was vacuous).
|
||||
const editor = makeEditor();
|
||||
seedContent(editor, "AAAABBBB");
|
||||
editor.commands.setTextSelection({ from: 1, to: 5 }); // captured range = {1,5}
|
||||
const claimed = paste(editor, "hello **bold**");
|
||||
expect(claimed).toBe(true);
|
||||
|
||||
// Mid-flight: grow the doc and move the cursor to a KNOWN-safe end position.
|
||||
editor.view.dispatch(editor.view.state.tr.insertText("MARKER", 1));
|
||||
const end = editor.state.doc.content.size;
|
||||
editor.commands.setTextSelection({ from: end, to: end });
|
||||
await flush();
|
||||
|
||||
const text = editor.getText();
|
||||
// MARKER intact only if the guard used the live selection, not the stale range.
|
||||
expect(text).toContain("MARKER");
|
||||
expect(text).toContain("bold");
|
||||
expect(text).not.toContain("**");
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("fail-open branch: a mid-flight doc SHRINK makes the stale `to` out of bounds — the guard must avoid a RangeError (throw-proving)", async () => {
|
||||
// The paste captures a range {1,9} over an 8-char paragraph, then the
|
||||
// conversion FAILS (`# heading` -> a heading node the minimal schema lacks,
|
||||
// so PMNode.fromJSON throws -> the fail-open catch runs). Before the reject,
|
||||
// the doc is SHRUNK to an empty paragraph, so the captured `to` (9) is now far
|
||||
// past the doc's end. A WORKING guard inserts the raw text at the live (valid)
|
||||
// selection → "raw heading" lands. A BROKEN guard does insertText(md, 1, 9) on
|
||||
// a size-2 doc → RangeError, so the dispatch never runs and "raw heading" is
|
||||
// absent (the assertion reddens). A zero-width/growing-doc setup could never
|
||||
// push `to` out of bounds, which is why the earlier version was vacuous.
|
||||
const editor = makeEditor();
|
||||
seedContent(editor, "AAAABBBB");
|
||||
editor.commands.setTextSelection({ from: 1, to: 9 }); // captured range = {1,9}
|
||||
paste(editor, "# raw heading");
|
||||
|
||||
// Mid-flight: shrink the doc so the captured `to` = 9 is now out of bounds.
|
||||
seedContent(editor, "");
|
||||
await flush();
|
||||
|
||||
const text = editor.getText();
|
||||
// Raw text lands (via the live selection) only if the guard avoided the
|
||||
// stale, now-out-of-bounds range.
|
||||
expect(text).toContain("raw heading");
|
||||
editor.destroy();
|
||||
});
|
||||
|
||||
it("two pastes in flight: neither payload is lost (no data loss)", async () => {
|
||||
// Prod-unreachable (two paste events are separate macrotasks, and each
|
||||
// conversion resolves on a microtask before the next), but pinned here: when
|
||||
// both resolve back-to-back, the second sees the changed doc and inserts at
|
||||
// the live selection the first left — so the two payloads may INTERLEAVE, but
|
||||
// neither is dropped. We assert no data loss, not contiguity.
|
||||
const editor = makeEditor();
|
||||
paste(editor, "alphaword");
|
||||
paste(editor, "betaword");
|
||||
await flush();
|
||||
const text = editor.getText();
|
||||
// Neither payload fully dropped (interleaving may split one of them).
|
||||
expect(text).toContain("alpha");
|
||||
expect(text).toContain("beta");
|
||||
editor.destroy();
|
||||
});
|
||||
});
|
||||
@@ -1,5 +1,12 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { htmlToMarkdown } from "@docmost/editor-ext";
|
||||
// Markdown conversion now goes through the canonical package's BROWSER entry
|
||||
// (issue #347): the same converter the server import/export uses, resolved via
|
||||
// the `browser` exports condition so it runs on the native `DOMParser` (the
|
||||
// client jsdom vitest env provides one) with jsdom never bundled.
|
||||
import {
|
||||
convertProseMirrorToMarkdown,
|
||||
markdownToProseMirrorSync,
|
||||
} from "@docmost/prosemirror-markdown/browser";
|
||||
import {
|
||||
normalizeTableColumnWidths,
|
||||
classifyClipboardSelection,
|
||||
@@ -175,10 +182,13 @@ describe("classifyClipboardSelection", () => {
|
||||
|
||||
// Output-level tests for the table clipboard regression: copying a table must
|
||||
// yield a real GFM pipe table, NOT one-value-per-line concatenated cells.
|
||||
// These exercise the actual markdown produced by htmlToMarkdown (the same
|
||||
// serializer step the clipboardTextSerializer runs), so they pin the OUTPUT
|
||||
// shape that the classifier-flag tests above do not cover.
|
||||
describe("table clipboard markdown output (htmlToMarkdown)", () => {
|
||||
// These exercise the actual markdown produced by convertProseMirrorToMarkdown —
|
||||
// the same serializer step the clipboardTextSerializer now runs (issue #347) —
|
||||
// so they pin the OUTPUT shape that the classifier-flag tests above do not cover.
|
||||
// Input is ProseMirror JSON (what the copied slice serializes to), matching the
|
||||
// clipboardTextSerializer's new call: it wraps the slice content in a synthetic
|
||||
// `doc` (and the bare-rows case in a `table`) and calls the converter.
|
||||
describe("table clipboard markdown output (convertProseMirrorToMarkdown)", () => {
|
||||
// Trim each line and drop blanks so structural assertions are whitespace-robust.
|
||||
function lines(md: string): string[] {
|
||||
return md
|
||||
@@ -188,10 +198,10 @@ describe("table clipboard markdown output (htmlToMarkdown)", () => {
|
||||
}
|
||||
|
||||
// A GFM separator row like "| --- | --- |" (any number of columns), tolerant
|
||||
// of the padding turndown emits.
|
||||
// of the padding the serializer emits.
|
||||
function isSeparatorRow(line: string): boolean {
|
||||
const compact = line.replace(/\s+/g, "");
|
||||
return /^\|(?:-{3,}\|)+$/.test(compact);
|
||||
return /^\|(?::?-{2,}:?\|)+$/.test(compact);
|
||||
}
|
||||
|
||||
// Split a pipe-delimited row into trimmed cell values.
|
||||
@@ -203,42 +213,33 @@ describe("table clipboard markdown output (htmlToMarkdown)", () => {
|
||||
.map((c) => c.trim());
|
||||
}
|
||||
|
||||
it("serializes a header-less partial cell selection (bare rows) as a valid GFM pipe table", () => {
|
||||
// Mirror the serializer's `wrapBareRows` branch exactly: bare <tr> nodes are
|
||||
// wrapped in <table><tbody> and htmlToMarkdown(div.innerHTML) is called.
|
||||
// See markdown-clipboard.ts clipboardTextSerializer:
|
||||
// const table = document.createElement("table");
|
||||
// const tbody = document.createElement("tbody");
|
||||
// tbody.appendChild(fragment); table.appendChild(tbody);
|
||||
// div.appendChild(table);
|
||||
// return htmlToMarkdown(div.innerHTML);
|
||||
const div = document.createElement("div");
|
||||
const table = document.createElement("table");
|
||||
const tbody = document.createElement("tbody");
|
||||
for (const [c1, c2] of [
|
||||
["a", "b"],
|
||||
["c", "d"],
|
||||
]) {
|
||||
const tr = document.createElement("tr");
|
||||
const td1 = document.createElement("td");
|
||||
td1.textContent = c1;
|
||||
const td2 = document.createElement("td");
|
||||
td2.textContent = c2;
|
||||
tr.appendChild(td1);
|
||||
tr.appendChild(td2);
|
||||
tbody.appendChild(tr);
|
||||
}
|
||||
table.appendChild(tbody);
|
||||
div.appendChild(table);
|
||||
const cell = (t: string) => ({
|
||||
type: "tableCell",
|
||||
content: [{ type: "paragraph", content: [{ type: "text", text: t }] }],
|
||||
});
|
||||
const headerCell = (t: string) => ({
|
||||
type: "tableHeader",
|
||||
content: [{ type: "paragraph", content: [{ type: "text", text: t }] }],
|
||||
});
|
||||
const row = (nodes: any[]) => ({ type: "tableRow", content: nodes });
|
||||
|
||||
const md = htmlToMarkdown(div.innerHTML);
|
||||
it("serializes a header-less partial cell selection (bare rows) as a valid GFM pipe table", () => {
|
||||
// Mirror the serializer's `wrapBareRows` branch: bare tableRow nodes are
|
||||
// wrapped in a synthetic `table` and convertProseMirrorToMarkdown is called
|
||||
// (see markdown-clipboard.ts clipboardTextSerializer).
|
||||
const rows = [
|
||||
row([cell("a"), cell("b")]),
|
||||
row([cell("c"), cell("d")]),
|
||||
];
|
||||
const md = convertProseMirrorToMarkdown({
|
||||
type: "doc",
|
||||
content: [{ type: "table", content: rows }],
|
||||
});
|
||||
const ls = lines(md);
|
||||
|
||||
// Valid GFM: a header/data separator row is present (an empty header is
|
||||
// synthesized by the GFM turndown plugin for a header-less table — fine).
|
||||
// Valid GFM: a header/data separator row is present.
|
||||
expect(ls.some(isSeparatorRow)).toBe(true);
|
||||
// NOT the old broken "one value per line" shape: every line is pipe-delimited
|
||||
// and no line is a bare cell value on its own.
|
||||
// NOT the old broken "one value per line" shape: every line is pipe-delimited.
|
||||
expect(ls.every((l) => l.includes("|"))).toBe(true);
|
||||
expect(md).not.toMatch(/^\s*(a|b|c|d)\s*$/m);
|
||||
// The cell values land in real pipe-delimited data rows.
|
||||
@@ -248,39 +249,21 @@ describe("table clipboard markdown output (htmlToMarkdown)", () => {
|
||||
});
|
||||
|
||||
it("serializes a whole table with a header row as a proper GFM table (headline regression)", () => {
|
||||
// Mirror the serializer's non-wrap branch: the full <table> node is appended
|
||||
// directly (div.appendChild(fragment)) and htmlToMarkdown(div.innerHTML) runs.
|
||||
const div = document.createElement("div");
|
||||
const table = document.createElement("table");
|
||||
|
||||
const thead = document.createElement("thead");
|
||||
const headerRow = document.createElement("tr");
|
||||
for (const h of ["Name", "Age"]) {
|
||||
const th = document.createElement("th");
|
||||
th.textContent = h;
|
||||
headerRow.appendChild(th);
|
||||
}
|
||||
thead.appendChild(headerRow);
|
||||
table.appendChild(thead);
|
||||
|
||||
const tbody = document.createElement("tbody");
|
||||
for (const [name, age] of [
|
||||
["Alice", "30"],
|
||||
["Bob", "25"],
|
||||
]) {
|
||||
const tr = document.createElement("tr");
|
||||
const td1 = document.createElement("td");
|
||||
td1.textContent = name;
|
||||
const td2 = document.createElement("td");
|
||||
td2.textContent = age;
|
||||
tr.appendChild(td1);
|
||||
tr.appendChild(td2);
|
||||
tbody.appendChild(tr);
|
||||
}
|
||||
table.appendChild(tbody);
|
||||
div.appendChild(table);
|
||||
|
||||
const md = htmlToMarkdown(div.innerHTML);
|
||||
// Mirror the serializer's non-wrap branch: the full `table` node is the
|
||||
// slice content and convertProseMirrorToMarkdown runs on it.
|
||||
const md = convertProseMirrorToMarkdown({
|
||||
type: "doc",
|
||||
content: [
|
||||
{
|
||||
type: "table",
|
||||
content: [
|
||||
row([headerCell("Name"), headerCell("Age")]),
|
||||
row([cell("Alice"), cell("30")]),
|
||||
row([cell("Bob"), cell("25")]),
|
||||
],
|
||||
},
|
||||
],
|
||||
});
|
||||
const ls = lines(md);
|
||||
|
||||
// Proper GFM structure: separator row + all rows pipe-delimited.
|
||||
@@ -296,3 +279,146 @@ describe("table clipboard markdown output (htmlToMarkdown)", () => {
|
||||
expect(md).not.toMatch(/^\s*(Name|Age|Alice|Bob|30|25)\s*$/m);
|
||||
});
|
||||
});
|
||||
|
||||
// #347 acceptance: pasting CANONICAL markdown yields the SAME nodes the server
|
||||
// import produces for the same text. The paste path calls markdownToProseMirror
|
||||
// (the package browser entry) — the identical converter the server import uses —
|
||||
// so asserting the converter (via the browser entry, on the native DOMParser)
|
||||
// recognizes each canon form pins the paste-parity guarantee. These forms were
|
||||
// NOT recognized by the old editor-ext marked layer the paste used before.
|
||||
describe("canonical markdown paste recognition (browser entry parity)", () => {
|
||||
// Collect every node type present in a doc (recursively).
|
||||
const collectTypes = (n: any, set = new Set<string>()): Set<string> => {
|
||||
if (!n || typeof n !== "object") return set;
|
||||
if (n.type) set.add(n.type);
|
||||
if (Array.isArray(n.content)) n.content.forEach((c) => collectTypes(c, set));
|
||||
return set;
|
||||
};
|
||||
const findNode = (n: any, type: string): any => {
|
||||
if (!n || typeof n !== "object") return undefined;
|
||||
if (n.type === type) return n;
|
||||
if (Array.isArray(n.content)) {
|
||||
for (const c of n.content) {
|
||||
const hit = findNode(c, type);
|
||||
if (hit) return hit;
|
||||
}
|
||||
}
|
||||
return undefined;
|
||||
};
|
||||
const allText = (n: any): string => {
|
||||
if (!n || typeof n !== "object") return "";
|
||||
if (typeof n.text === "string") return n.text;
|
||||
if (Array.isArray(n.content)) return n.content.map(allText).join("");
|
||||
return "";
|
||||
};
|
||||
|
||||
it("^[…] inline footnote -> footnoteReference + footnotesList", () => {
|
||||
const doc = markdownToProseMirrorSync("Body^[a note here].");
|
||||
const types = collectTypes(doc);
|
||||
expect(types.has("footnoteReference")).toBe(true);
|
||||
expect(types.has("footnotesList")).toBe(true);
|
||||
expect(types.has("footnoteDefinition")).toBe(true);
|
||||
});
|
||||
|
||||
it('<!--img {…}--> attached image comment -> image with align', () => {
|
||||
const doc = markdownToProseMirrorSync(
|
||||
' <!--img {"align":"left"}-->',
|
||||
);
|
||||
const img = findNode(doc, "image");
|
||||
expect(img).toBeTruthy();
|
||||
expect(img.attrs?.align).toBe("left");
|
||||
expect(img.attrs?.src).toBe("/files/x.png");
|
||||
});
|
||||
|
||||
it("> [!type] Obsidian callout -> callout node with type", () => {
|
||||
const doc = markdownToProseMirrorSync("> [!warning]\n> be careful");
|
||||
const callout = findNode(doc, "callout");
|
||||
expect(callout).toBeTruthy();
|
||||
expect(callout.attrs?.type).toBe("warning");
|
||||
expect(allText(callout)).toContain("be careful");
|
||||
});
|
||||
|
||||
it("$…$ inline math -> mathInline node", () => {
|
||||
const doc = markdownToProseMirrorSync("Euler: $e^{i\\pi}+1=0$ done");
|
||||
const math = findNode(doc, "mathInline");
|
||||
expect(math).toBeTruthy();
|
||||
expect(math.attrs?.text).toContain("e^{i\\pi}");
|
||||
});
|
||||
|
||||
it("==…== highlight -> highlight mark", () => {
|
||||
const doc = markdownToProseMirrorSync("A ==marked== word");
|
||||
const marked = findNode(doc, "text");
|
||||
// The highlighted run carries a `highlight` mark somewhere in the doc.
|
||||
const hasHighlight = (n: any): boolean => {
|
||||
if (!n || typeof n !== "object") return false;
|
||||
if (
|
||||
n.type === "text" &&
|
||||
(n.marks || []).some((m: any) => m.type === "highlight")
|
||||
)
|
||||
return true;
|
||||
return Array.isArray(n.content) ? n.content.some(hasHighlight) : false;
|
||||
};
|
||||
expect(marked).toBeTruthy();
|
||||
expect(hasHighlight(doc)).toBe(true);
|
||||
});
|
||||
|
||||
it("<!--subpages--> standalone comment -> subpages node", () => {
|
||||
const doc = markdownToProseMirrorSync("intro\n\n<!--subpages-->\n\nafter");
|
||||
expect(collectTypes(doc).has("subpages")).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
// #347 negatives: plain text carrying markdown-LIKE punctuation must NOT be
|
||||
// silently converted/mangled (currency, bare `==`, a `[^1]` reference form).
|
||||
describe("plain-text paste negatives (no phantom conversion)", () => {
|
||||
const findNode = (n: any, type: string): any => {
|
||||
if (!n || typeof n !== "object") return undefined;
|
||||
if (n.type === type) return n;
|
||||
if (Array.isArray(n.content)) {
|
||||
for (const c of n.content) {
|
||||
const hit = findNode(c, type);
|
||||
if (hit) return hit;
|
||||
}
|
||||
}
|
||||
return undefined;
|
||||
};
|
||||
const collectTypes = (n: any, set = new Set<string>()): Set<string> => {
|
||||
if (!n || typeof n !== "object") return set;
|
||||
if (n.type) set.add(n.type);
|
||||
if (Array.isArray(n.content)) n.content.forEach((c) => collectTypes(c, set));
|
||||
return set;
|
||||
};
|
||||
const allText = (n: any): string => {
|
||||
if (!n || typeof n !== "object") return "";
|
||||
if (typeof n.text === "string") return n.text;
|
||||
if (Array.isArray(n.content)) return n.content.map(allText).join("");
|
||||
return "";
|
||||
};
|
||||
|
||||
it("currency `$5 and $10` is NOT turned into math", () => {
|
||||
const doc = markdownToProseMirrorSync("It costs $5 and $10 total");
|
||||
expect(findNode(doc, "mathInline")).toBeFalsy();
|
||||
expect(allText(doc)).toContain("$5 and $10");
|
||||
});
|
||||
|
||||
it("a lone `==` is NOT turned into a highlight", () => {
|
||||
const doc = markdownToProseMirrorSync("compare a == b in code");
|
||||
const hasHighlight = (n: any): boolean => {
|
||||
if (!n || typeof n !== "object") return false;
|
||||
if (
|
||||
n.type === "text" &&
|
||||
(n.marks || []).some((m: any) => m.type === "highlight")
|
||||
)
|
||||
return true;
|
||||
return Array.isArray(n.content) ? n.content.some(hasHighlight) : false;
|
||||
};
|
||||
expect(hasHighlight(doc)).toBe(false);
|
||||
expect(allText(doc)).toContain("== b");
|
||||
});
|
||||
|
||||
it("a `[^1]` reference form (no `^[`) is NOT turned into a footnote", () => {
|
||||
const doc = markdownToProseMirrorSync("see note [^1] for details");
|
||||
expect(collectTypes(doc).has("footnoteReference")).toBe(false);
|
||||
expect(allText(doc)).toContain("[^1]");
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1,15 +1,23 @@
|
||||
// adapted from: https://github.com/aguingand/tiptap-markdown/blob/main/src/extensions/tiptap/clipboard.js - MIT
|
||||
import { Extension } from "@tiptap/core";
|
||||
import { Plugin, PluginKey, TextSelection } from "@tiptap/pm/state";
|
||||
import { DOMParser, DOMSerializer, Fragment, Slice } from "@tiptap/pm/model";
|
||||
import { DOMParser, DOMSerializer, Fragment, Slice, Node as PMNode } from "@tiptap/pm/model";
|
||||
import { find } from "linkifyjs";
|
||||
import {
|
||||
markdownToHtml,
|
||||
htmlToMarkdown,
|
||||
canonicalizeFootnotes,
|
||||
FOOTNOTES_LIST_NAME,
|
||||
FOOTNOTE_REFERENCE_NAME,
|
||||
} from "@docmost/editor-ext";
|
||||
// Markdown <-> ProseMirror conversion now lives ONLY in the canonical
|
||||
// `@docmost/prosemirror-markdown` package (issue #347). The BROWSER entry uses
|
||||
// the native `DOMParser` for its HTML->DOM stage (jsdom stays out of the client
|
||||
// bundle) while producing the SAME nodes the server import does — so a paste of
|
||||
// canonical markdown (`^[…]`, `<!--img …-->`, `> [!type]`, `$…$`, `==…==`,
|
||||
// standalone comments) is recognized identically to import.
|
||||
import {
|
||||
markdownToProseMirror,
|
||||
convertProseMirrorToMarkdown,
|
||||
} from "@docmost/prosemirror-markdown/browser";
|
||||
import type { Schema } from "@tiptap/pm/model";
|
||||
|
||||
export const MarkdownClipboard = Extension.create({
|
||||
@@ -39,25 +47,24 @@ export const MarkdownClipboard = Extension.create({
|
||||
classifyClipboardSelection(topLevelNodes);
|
||||
if (!asMarkdown) return null;
|
||||
|
||||
const div = document.createElement("div");
|
||||
const serializer = DOMSerializer.fromSchema(this.editor.schema);
|
||||
const fragment = serializer.serializeFragment(slice.content);
|
||||
|
||||
// Convert the copied selection to Markdown through the canonical
|
||||
// package (issue #347), the SAME serializer the server export uses,
|
||||
// so a copied table/list matches the on-disk markdown form. The
|
||||
// converter takes a ProseMirror `doc` JSON, so wrap the slice's
|
||||
// top-level content in a synthetic doc.
|
||||
const content = slice.content.toJSON() as any[];
|
||||
if (wrapBareRows) {
|
||||
// A partial table cell-selection serializes to bare <tr> nodes
|
||||
// (prosemirror-tables returns the whole `table` node only when the
|
||||
// entire table is selected). Bare <tr> would be foster-parented
|
||||
// away by the HTML parser inside htmlToMarkdown, so wrap them in
|
||||
// <table><tbody> first for the GFM turndown rule to detect them.
|
||||
const table = document.createElement("table");
|
||||
const tbody = document.createElement("tbody");
|
||||
tbody.appendChild(fragment);
|
||||
table.appendChild(tbody);
|
||||
div.appendChild(table);
|
||||
} else {
|
||||
div.appendChild(fragment);
|
||||
// A partial table cell-selection serializes to bare `tableRow`
|
||||
// nodes (prosemirror-tables yields the whole `table` node only for
|
||||
// a full-table selection). The converter's table case expects a
|
||||
// `table` wrapper, so wrap the bare rows in one — mirroring the old
|
||||
// <table><tbody> wrap that the HTML->markdown step needed.
|
||||
return convertProseMirrorToMarkdown({
|
||||
type: "doc",
|
||||
content: [{ type: "table", content }],
|
||||
});
|
||||
}
|
||||
return htmlToMarkdown(div.innerHTML);
|
||||
return convertProseMirrorToMarkdown({ type: "doc", content });
|
||||
},
|
||||
handlePaste: (view, event, slice) => {
|
||||
if (!event.clipboardData) {
|
||||
@@ -95,37 +102,115 @@ export const MarkdownClipboard = Extension.create({
|
||||
}
|
||||
}
|
||||
|
||||
const { tr } = view.state;
|
||||
const { from, to } = view.state.selection;
|
||||
const schema = this.editor.schema;
|
||||
// Capture the target range NOW. markdownToProseMirror RETURNS A
|
||||
// PROMISE (kept async only for the Node consumers' contract; the
|
||||
// conversion pipeline itself is synchronous), so the actual replace
|
||||
// happens on the next microtask. No user input can interleave a
|
||||
// microtask, so the state is unchanged when we dispatch — but we
|
||||
// still re-read the live state before replacing and, if the doc did
|
||||
// change under us, fall back to the live selection rather than the
|
||||
// captured (now-stale) range.
|
||||
const from = view.state.selection.from;
|
||||
const to = view.state.selection.to;
|
||||
const startDoc = view.state.doc;
|
||||
const md = text.replace(/\n+$/, "");
|
||||
|
||||
const parsed = markdownToHtml(text.replace(/\n+$/, ""));
|
||||
const body = elementFromString(parsed);
|
||||
normalizeTableColumnWidths(body);
|
||||
void markdownToProseMirror(md)
|
||||
.then((doc) => {
|
||||
if (view.isDestroyed) return;
|
||||
// Canonical PM-JSON -> HTML via the LIVE editor schema, then
|
||||
// reuse the UNCHANGED downstream seam (normalizeTableColumnWidths
|
||||
// + parseSlice + canonicalizePastedFootnotes). The JSON->HTML->
|
||||
// JSON hop is lossless (same schema both directions); it lets the
|
||||
// existing paste-insertion logic stay byte-identical — only the
|
||||
// SOURCE of the markdown conversion changed (issue #347 guardrail:
|
||||
// no converter logic in the client, only a call into the package).
|
||||
const node = PMNode.fromJSON(schema, doc);
|
||||
const div = document.createElement("div");
|
||||
DOMSerializer.fromSchema(schema).serializeFragment(
|
||||
node.content,
|
||||
{ document },
|
||||
div,
|
||||
);
|
||||
|
||||
const parsedSlice = DOMParser.fromSchema(
|
||||
this.editor.schema,
|
||||
).parseSlice(body, {
|
||||
preserveWhitespace: true,
|
||||
});
|
||||
const body = elementFromString(div.innerHTML);
|
||||
normalizeTableColumnWidths(body);
|
||||
|
||||
// A markdown paste builds its ProseMirror fragment directly (DOM ->
|
||||
// parseSlice), bypassing the editor's footnoteSyncPlugin, which never
|
||||
// reorders an existing list. So a pasted markdown block whose footnote
|
||||
// definitions are out of order (or contains orphan defs) would be
|
||||
// stored out of order. Canonicalize the self-contained pasted block so
|
||||
// its footnotes come out reference-ordered, deduped and orphan-free
|
||||
// (issue #228). See canonicalizePastedFootnotes for why this is scoped
|
||||
// to whole-block pastes that carry their own footnotesList.
|
||||
const contentNodes = canonicalizePastedFootnotes(
|
||||
parsedSlice,
|
||||
this.editor.schema,
|
||||
);
|
||||
const parsedSlice = DOMParser.fromSchema(schema).parseSlice(
|
||||
body,
|
||||
{ preserveWhitespace: true },
|
||||
);
|
||||
|
||||
tr.replaceRange(from, to, contentNodes);
|
||||
const insertEnd = tr.mapping.map(from, 1);
|
||||
tr.setSelection(TextSelection.near(tr.doc.resolve(Math.max(from, insertEnd - 2)), -1));
|
||||
tr.setMeta('paste', true)
|
||||
view.dispatch(tr);
|
||||
// A markdown paste builds its ProseMirror fragment directly (DOM
|
||||
// -> parseSlice), bypassing the editor's footnoteSyncPlugin, which
|
||||
// never reorders an existing list. So a pasted markdown block whose
|
||||
// footnote definitions are out of order (or contains orphan defs)
|
||||
// would be stored out of order. Canonicalize the self-contained
|
||||
// pasted block so its footnotes come out reference-ordered, deduped
|
||||
// and orphan-free (issue #228). See canonicalizePastedFootnotes for
|
||||
// why this is scoped to whole-block pastes that carry their own
|
||||
// footnotesList.
|
||||
const contentNodes = canonicalizePastedFootnotes(
|
||||
parsedSlice,
|
||||
schema,
|
||||
);
|
||||
|
||||
// Target the captured range (normally still valid — same
|
||||
// microtask). If the doc changed under us since capture, the
|
||||
// captured absolute from/to are stale, so fall back to the live
|
||||
// selection rather than StepMap-mapping the old range.
|
||||
const tr = view.state.tr;
|
||||
let mappedFrom = from;
|
||||
let mappedTo = to;
|
||||
if (view.state.doc !== startDoc) {
|
||||
// Defensive: if the doc changed under us, fall back to the
|
||||
// current selection rather than a stale absolute range.
|
||||
mappedFrom = view.state.selection.from;
|
||||
mappedTo = view.state.selection.to;
|
||||
}
|
||||
tr.replaceRange(mappedFrom, mappedTo, contentNodes);
|
||||
const insertEnd = tr.mapping.map(mappedFrom, 1);
|
||||
tr.setSelection(
|
||||
TextSelection.near(
|
||||
tr.doc.resolve(Math.max(mappedFrom, insertEnd - 2)),
|
||||
-1,
|
||||
),
|
||||
);
|
||||
tr.setMeta("paste", true);
|
||||
view.dispatch(tr);
|
||||
})
|
||||
.catch((err) => {
|
||||
// Fail-open: a conversion error must not swallow the paste
|
||||
// silently in a way that loses the text. We already claimed the
|
||||
// event (returned true), so re-insert the raw text as a plain
|
||||
// paragraph so the user never loses their clipboard content.
|
||||
// Log it: this catch covers BOTH the converter and the success
|
||||
// `.then` body (e.g. PMNode.fromJSON throwing on a schema drift
|
||||
// between the canonical package and the live editor schema), so a
|
||||
// silent degrade to raw text would otherwise be an invisible,
|
||||
// non-reproducible regression ("my table pasted as text").
|
||||
console.error(
|
||||
"markdown paste conversion failed, inserting raw text",
|
||||
err,
|
||||
);
|
||||
if (view.isDestroyed) return;
|
||||
const tr = view.state.tr;
|
||||
// Same guard the success path uses: if the doc changed under us
|
||||
// since the range was captured (normally never — same microtask),
|
||||
// the captured absolute from/to are stale and would throw a
|
||||
// RangeError here (an unhandled rejection on a hot paste path).
|
||||
// Fall back to the live selection instead of a stale range.
|
||||
if (view.state.doc !== startDoc) {
|
||||
const sel = view.state.selection;
|
||||
tr.insertText(md, sel.from, sel.to);
|
||||
} else {
|
||||
tr.insertText(md, from, to);
|
||||
}
|
||||
tr.setMeta("paste", true);
|
||||
view.dispatch(tr);
|
||||
});
|
||||
// Claim the paste: we insert asynchronously above.
|
||||
return true;
|
||||
},
|
||||
// Strip trailing whitespace-only paragraphs from pasted content.
|
||||
|
||||
@@ -33,10 +33,11 @@ vi.mock("@/lib/local-emitter.ts", () => ({
|
||||
default: { emit: (...args: unknown[]) => localEmitMock(...args) },
|
||||
}));
|
||||
|
||||
// htmlToMarkdown just echoes the editor HTML so each test controls the markdown
|
||||
// purely via the fake page editor's getHTML().
|
||||
vi.mock("@docmost/editor-ext", () => ({
|
||||
htmlToMarkdown: (html: string) => html,
|
||||
// convertProseMirrorToMarkdown echoes a marker carried on the fake editor's
|
||||
// getJSON() doc, so each test controls the markdown purely via the fake page
|
||||
// editor (issue #347: the hook now serializes editor JSON through the package).
|
||||
vi.mock("@docmost/prosemirror-markdown/browser", () => ({
|
||||
convertProseMirrorToMarkdown: (doc: { __md?: string }) => doc?.__md ?? "",
|
||||
}));
|
||||
|
||||
const notificationsShowMock = vi.fn();
|
||||
@@ -53,10 +54,12 @@ import { useGeneratePageTitle } from "./use-generate-page-title.ts";
|
||||
|
||||
// --- Test helpers -------------------------------------------------------------
|
||||
|
||||
function makePageEditor(pageId: string, html = "<p>content</p>"): Editor {
|
||||
function makePageEditor(pageId: string, md = "content"): Editor {
|
||||
return {
|
||||
isDestroyed: false,
|
||||
getHTML: () => html,
|
||||
// The mocked convertProseMirrorToMarkdown reads `__md` back off this doc,
|
||||
// so `md` is exactly the markdown the hook will send to the title service.
|
||||
getJSON: () => ({ type: "doc", __md: md }),
|
||||
storage: { pageId },
|
||||
} as unknown as Editor;
|
||||
}
|
||||
|
||||
@@ -3,7 +3,7 @@ import { useMutation } from "@tanstack/react-query";
|
||||
import { useAtomValue } from "jotai";
|
||||
import { notifications } from "@mantine/notifications";
|
||||
import { useTranslation } from "react-i18next";
|
||||
import { htmlToMarkdown } from "@docmost/editor-ext";
|
||||
import { convertProseMirrorToMarkdown } from "@docmost/prosemirror-markdown/browser";
|
||||
import {
|
||||
pageEditorAtom,
|
||||
titleEditorAtom,
|
||||
@@ -49,7 +49,9 @@ export function useGeneratePageTitle(pageId: string) {
|
||||
mutationFn: async () => {
|
||||
if (!pageEditor || pageEditor.isDestroyed) return;
|
||||
|
||||
const markdown = htmlToMarkdown(pageEditor.getHTML()).trim();
|
||||
// Serialize the live editor content to markdown through the canonical
|
||||
// converter (issue #347), matching the on-disk/export markdown form.
|
||||
const markdown = convertProseMirrorToMarkdown(pageEditor.getJSON()).trim();
|
||||
if (!markdown) {
|
||||
notifications.show({ message: t("The note is empty"), color: "yellow" });
|
||||
return;
|
||||
|
||||
@@ -37,7 +37,7 @@ import { useTreeMutation } from "@/features/page/tree/hooks/use-tree-mutation.ts
|
||||
import { PageWidthToggle } from "@/features/user/components/page-width-pref.tsx";
|
||||
import { Trans, useTranslation } from "react-i18next";
|
||||
import ExportModal from "@/components/common/export-modal";
|
||||
import { htmlToMarkdown } from "@docmost/editor-ext";
|
||||
import { convertProseMirrorToMarkdown } from "@docmost/prosemirror-markdown/browser";
|
||||
import {
|
||||
pageEditorAtom,
|
||||
yjsConnectionStatusAtom,
|
||||
@@ -199,8 +199,9 @@ function PageActionMenu({ readOnly }: PageActionMenuProps) {
|
||||
|
||||
const handleCopyAsMarkdown = () => {
|
||||
if (!pageEditor) return;
|
||||
const html = pageEditor.getHTML();
|
||||
const markdown = htmlToMarkdown(html);
|
||||
// Copy the page as canonical markdown through the shared converter (issue
|
||||
// #347), so "Copy as markdown" matches the server export byte-for-byte.
|
||||
const markdown = convertProseMirrorToMarkdown(pageEditor.getJSON());
|
||||
const title = page?.title ? `# ${page.title}\n\n` : "";
|
||||
clipboard.copy(`${title}${markdown}`);
|
||||
notifications.show({ message: t("Copied") });
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import { markdownToHtml, encodeHtmlEmbedSource } from '@docmost/editor-ext';
|
||||
import { markdownToProseMirror } from '@docmost/prosemirror-markdown';
|
||||
import { encodeHtmlEmbedSource } from '@docmost/editor-ext';
|
||||
import { htmlToJson } from '../../../collaboration/collaboration.util';
|
||||
import { hasHtmlEmbedNode, stripHtmlEmbedNodes } from './html-embed.util';
|
||||
|
||||
@@ -10,13 +11,12 @@ import { hasHtmlEmbedNode, stripHtmlEmbedNodes } from './html-embed.util';
|
||||
*
|
||||
* The block renders inside a sandboxed iframe, so this is not an XSS surface;
|
||||
* this exercises the REAL server import conversion path that ImportService uses
|
||||
* (`markdownToHtml` then `htmlToJson`; `processHTML` adds only a cheerio
|
||||
* link/iframe normalize pass which does not touch htmlEmbed divs) and asserts
|
||||
* that such a node is DETECTED and STRIPPABLE — so the share read path's
|
||||
* (`markdownToProseMirror`, the canonical converter — issue #345/#347) and
|
||||
* asserts that such a node is DETECTED and STRIPPABLE — so the share read path's
|
||||
* master-toggle strip can remove it when the workspace toggle is OFF.
|
||||
*/
|
||||
describe('htmlEmbed smuggled via the raw serialized div in imported markdown/HTML', () => {
|
||||
it('round-trips through markdownToHtml -> htmlToJson and is DETECTED (base64 data-source)', async () => {
|
||||
it('round-trips through markdownToProseMirror and is DETECTED (base64 data-source)', async () => {
|
||||
const source = '<script>steal()</script>';
|
||||
const encoded = encodeHtmlEmbedSource(source);
|
||||
const md = [
|
||||
@@ -27,12 +27,9 @@ describe('htmlEmbed smuggled via the raw serialized div in imported markdown/HTM
|
||||
'World',
|
||||
].join('\n');
|
||||
|
||||
const html = await markdownToHtml(md);
|
||||
// marked preserves the raw block-level div verbatim.
|
||||
expect(html).toContain('data-type="htmlEmbed"');
|
||||
|
||||
const json = htmlToJson(html);
|
||||
// The div parses into a real htmlEmbed node carrying the decoded source.
|
||||
// The canonical importer parses the raw block-level div into a real
|
||||
// htmlEmbed node carrying the decoded source.
|
||||
const json = await markdownToProseMirror(md);
|
||||
expect(hasHtmlEmbedNode(json)).toBe(true);
|
||||
|
||||
// Because it is detected, the share master-toggle strip can remove it.
|
||||
@@ -59,8 +56,7 @@ describe('htmlEmbed smuggled via the raw serialized div in imported markdown/HTM
|
||||
// therefore stripping) does not depend on the source being well-formed, so
|
||||
// the bypass cannot be hidden by sending a malformed data-source.
|
||||
const md = `<div data-type="htmlEmbed" data-source="<script>x</script>"></div>`;
|
||||
const html = await markdownToHtml(md);
|
||||
const json = htmlToJson(html);
|
||||
const json = await markdownToProseMirror(md);
|
||||
expect(hasHtmlEmbedNode(json)).toBe(true);
|
||||
expect(hasHtmlEmbedNode(stripHtmlEmbedNodes(json))).toBe(false);
|
||||
});
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -1,42 +1,122 @@
|
||||
import { Injectable, Logger, OnModuleDestroy } from '@nestjs/common';
|
||||
|
||||
/**
|
||||
* In-memory run-stream registry (#184 phase 1.5). A durable agent run tees its
|
||||
* SSE frames here (via `pipeUIMessageStreamToResponse({ consumeSseStream })`)
|
||||
* so a LATE tab — one that reloaded, or opened after the starter dropped — can
|
||||
* attach through `GET /ai-chat/runs/:chatId/stream`, replay the frames buffered
|
||||
* so far, and then follow the live tail as a normal streamer.
|
||||
* In-memory run-stream registry (#184 phase 1.5, step-aligned retention #491). A
|
||||
* durable agent run tees its SSE frames here (via
|
||||
* `pipeUIMessageStreamToResponse({ consumeSseStream })`) so a LATE tab — one that
|
||||
* reloaded, or opened after the starter dropped — can attach through
|
||||
* `GET /ai-chat/runs/:chatId/stream`, be handed the TAIL past the step it already
|
||||
* has persisted, and then follow the live tail as a normal streamer.
|
||||
*
|
||||
* This is deliberately single-process and best-effort: it holds nothing the DB
|
||||
* does not (the run + assistant row are the source of truth), so a process
|
||||
* restart simply drops in-flight entries and the client falls back to its
|
||||
* restore + degraded-poll path. The async `attach` return type is the seam for a
|
||||
* future phase-2 cross-process backend (Redis) — the interface does not change.
|
||||
*
|
||||
* ── #491 step-aligned retention (the OOM fix) ────────────────────────────────
|
||||
* The old registry buffered up to 32MB of raw SSE frames PER active run (V8 ~2×
|
||||
* in memory) and, on attach, blasted the WHOLE buffer to the socket synchronously
|
||||
* with no drain — a handful of marathon runs on a 1GB container OOM'd. #491 caps
|
||||
* the ring at a few MB (env-tunable, default 4MB) and keeps it there by ROTATING:
|
||||
*
|
||||
* - Every buffered frame is STAMPED with a step number at tee (see ingestFrame).
|
||||
* Convention: the stamp of a frame is the number of `finish-step` parts seen
|
||||
* BEFORE it (starting at 0). The finish-step frame itself carries the current
|
||||
* value, THEN the counter increments. So a frame stamped `s` is the content of
|
||||
* the (s+1)-th step — 0-based step index `s` — and the stamp aligns EXACTLY
|
||||
* with `metadata.stepsPersisted`: a client whose persisted `stepsPersisted` is
|
||||
* N has steps 0..N-1 on disk (and in its seed) and needs the tail `stamp >= N`.
|
||||
*
|
||||
* - The ring rotates ONLY on a CONFIRMED persist of step N
|
||||
* (`confirmPersistedStep`), dropping frames with `stamp < N` (those steps are
|
||||
* now on disk and a fresh client seed carries them). A NON-confirmed step is
|
||||
* never rotated away, so a persist FAILURE just makes the ring cover MORE
|
||||
* (auto-safe). This is the anti-inversion rule: a naive "rotate in .then()"
|
||||
* that rotated after an UNwritten step would drop a step nobody has → silent
|
||||
* hole. Rotation is gated on a real, successful persist.
|
||||
*
|
||||
* - If the ring still exceeds its byte cap after rotation (a single fat step, or
|
||||
* a lagging persist), the OLDEST frames are evicted to stay bounded. Evicting a
|
||||
* not-yet-persisted frame opens a GAP: an attach whose N falls at or below an
|
||||
* evicted step answers 204 and the client degrades to restore+poll. The gap is
|
||||
* NOT sticky — the coverage floor is recomputed from the ring, so a later
|
||||
* persist that rotates past the holey steps clears it.
|
||||
*
|
||||
* ── attach numbering / coverage (the wire convention) ────────────────────────
|
||||
* The step marker N comes ONLY FROM THE CLIENT (a query param). The server never
|
||||
* reads the row to derive N — a server-side N from a stale seed would open a
|
||||
* silent one-step hole. N is the client's persisted `stepsPersisted` (a COUNT):
|
||||
* - the tail it needs = frames with `stamp >= N`;
|
||||
* - coverage is OK ⟺ `coverageFloor(entry) <= N`, where coverageFloor is the
|
||||
* smallest step FULLY present in the ring (its smallest retained stamp, bumped
|
||||
* by one when that leading step was only partially evicted by overflow). If
|
||||
* `coverageFloor > N` the ring starts AFTER the client's frontier (a hole, or
|
||||
* the client's seed simply lagged behind a rotation) → 204 → the client
|
||||
* refetches (a larger N) and re-attaches.
|
||||
* The N cutoff is applied in ALL branches, INCLUDING the finished-retained replay.
|
||||
*
|
||||
* ── same-tick invariants (unchanged, still load-bearing) ─────────────────────
|
||||
* invariant 1: only the matching run may mutate/observe an entry (runId check).
|
||||
* invariant 2: retention deletes ONLY its own entry (a replacement may own the key).
|
||||
* invariant 3: open() over a live entry mirrors the done-path (subscribers released).
|
||||
* invariant 4: the tail SLICE + subscriber registration happen in ONE synchronous
|
||||
* tick inside attach() — no await between them — so a concurrently
|
||||
* ingested frame is EITHER in the snapshot (buffered before the sync
|
||||
* block, and the just-added subscriber never sees it) OR fanned out to
|
||||
* the paused subscriber's `pending` (ingested after) — never both and
|
||||
* never neither: no loss, no duplication. NOTE (#491): the controller
|
||||
* now AWAITS the drain-respecting tail write BEFORE calling start(), so
|
||||
* frames ingested during that await accumulate in `pending`; this is
|
||||
* bounded by the subscriber cap (an overflow degrades start() to an
|
||||
* end(), a 204-equivalent). It is the SYNCHRONOUS snapshot+registration
|
||||
* — not a same-tick start() — that makes this correct.
|
||||
* invariant 5: the controller wires close-cleanup BEFORE any write.
|
||||
* invariant 6: no cross-run replay — the `anchor` (the client's assistant row id)
|
||||
* must match this run's assistant id, or a foreign run's transcript
|
||||
* would be appended to the client's message.
|
||||
*/
|
||||
|
||||
/** How long a finished entry is retained for late attach (replay + immediate end). */
|
||||
export const RUN_STREAM_RETAIN_FINISHED_MS = 30_000;
|
||||
|
||||
/**
|
||||
* Per-run replay buffer cap. Past this the buffer is dropped (attach -> 204, and
|
||||
* the client falls back to its restore + degraded-poll path, #430).
|
||||
*
|
||||
* Raised from 4MB to 32MB (#430): marathon autonomous runs (11-25 min observed)
|
||||
* stream far more than 4MB of SSE frames, so a live disconnect mid-run would find
|
||||
* an already-overflowed buffer and could only degrade-poll instead of re-attaching
|
||||
* to the live tail. 32MB comfortably covers those runs while staying bounded.
|
||||
*
|
||||
* Memory cost: this is the WORST-CASE retained size PER ACTIVE run (the buffer is
|
||||
* freed on finish + retention, or dropped immediately on overflow). With the small
|
||||
* number of concurrent autonomous runs a single workspace realistically has, 32MB
|
||||
* each is an acceptable ceiling; the overflow->204->degraded-poll fallback remains
|
||||
* the backstop for anything larger, so correctness never depends on this bound.
|
||||
* DEFAULT per-run replay ring cap (#491, down from 32MB). SSE frames carry
|
||||
* UNcompacted tool outputs + framing overhead (×1.5–2 vs the persisted parts), so
|
||||
* a "2–3 large reads + reasoning" step routinely blows past 2MB; 4MB comfortably
|
||||
* holds a step or two of TAIL, which is all a resuming client needs (steps below
|
||||
* its persisted frontier come from the seed, not the ring). The ring stays bounded
|
||||
* because it rotates on every confirmed persist; this cap is only the ceiling for
|
||||
* the un-persisted tail between rotations. Env-tunable via
|
||||
* AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES (bytes); a 0/invalid value falls back to this.
|
||||
*/
|
||||
export const RUN_STREAM_MAX_BUFFER_BYTES = 32 * 1024 * 1024;
|
||||
export const AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES = 4 * 1024 * 1024;
|
||||
|
||||
// 2x the replay cap: a just-written full-replay burst alone can never trip the
|
||||
// per-subscriber cap (see controller); only a genuinely stalled socket can.
|
||||
export const SUBSCRIBER_MAX_BUFFERED_BYTES = 2 * RUN_STREAM_MAX_BUFFER_BYTES;
|
||||
// 2× the ring cap: a just-written full-tail burst alone can never trip the
|
||||
// per-subscriber cap (see controller); only a genuinely stalled socket can. This
|
||||
// derivative relationship is preserved even when the ring cap is env-overridden.
|
||||
export const SUBSCRIBER_MAX_BUFFERED_BYTES = 2 * AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES;
|
||||
|
||||
/**
|
||||
* A finish-step boundary frame is exactly `data: {"type":"finish-step"...}\n\n`
|
||||
* (verified empirically against ai@6.0.207 — each UI-message-stream part is a
|
||||
* single `data: {json}\n\n` event, never split across `data:` lines, and `type`
|
||||
* is always the first key). A prefix match is cheaper than JSON.parse-per-frame
|
||||
* and has no false positives: a literal `"type":"finish-step"` inside a text
|
||||
* delta is JSON-escaped (`\"type\":...`), and the frame would start with
|
||||
* `data: {"type":"text-delta"` anyway.
|
||||
*/
|
||||
const FINISH_STEP_FRAME_PREFIX = 'data: {"type":"finish-step"';
|
||||
|
||||
/** Resolve the ring cap from the environment, falling back to the default. */
|
||||
function resolveMaxBufferBytes(): number {
|
||||
const raw = process.env.AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES;
|
||||
if (!raw) return AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES;
|
||||
const parsed = Number(raw);
|
||||
return Number.isFinite(parsed) && parsed > 0
|
||||
? Math.floor(parsed)
|
||||
: AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES;
|
||||
}
|
||||
|
||||
export interface RunStreamCallbacks {
|
||||
onFrame: (frame: string) => void;
|
||||
@@ -44,6 +124,9 @@ export interface RunStreamCallbacks {
|
||||
}
|
||||
|
||||
export interface RunStreamAttachment {
|
||||
// The synthetic `start` frame (carrying { runId, chatId }) followed by the
|
||||
// buffered TAIL filtered to `stamp >= N`. The controller writes these to the
|
||||
// socket in chunks respecting drain, then calls start().
|
||||
replay: string[];
|
||||
finished: boolean;
|
||||
start(): void; // drain pending frames (order preserved) and go live
|
||||
@@ -53,14 +136,19 @@ export interface RunStreamAttachment {
|
||||
interface Subscriber extends RunStreamCallbacks {
|
||||
started: boolean;
|
||||
pending: string[];
|
||||
// Byte size of `pending`, capped at SUBSCRIBER_MAX_BUFFERED_BYTES. `start()` is
|
||||
// called in the SAME tick as `attach()` today (see attach), so `pending` never
|
||||
// holds more than one microtask of frames — but the async `attach` signature is
|
||||
// a phase-2 seam: an await between attach and start would let a stalled paused
|
||||
// subscriber buffer the WHOLE run here. The cap is the structural backstop.
|
||||
// Byte size of `pending`, capped at the subscriber cap. `start()` is called in
|
||||
// the SAME tick as `attach()` today, so `pending` never holds more than one
|
||||
// microtask of frames — but the controller writes the (potentially large) tail
|
||||
// respecting drain BEFORE start(), so a stalled socket can accumulate here; the
|
||||
// cap is the structural backstop (an overflow degrades start() to an end()).
|
||||
pendingBytes: number;
|
||||
overflowed: boolean;
|
||||
pendingEnd: boolean;
|
||||
// The client's step frontier N: this subscriber only receives frames with
|
||||
// `stamp >= minStamp` (the tail past what it already persisted). Live frames
|
||||
// always satisfy this (their stamp is the current, highest step), so it only
|
||||
// filters the rare out-of-order below-frontier frame.
|
||||
minStamp: number;
|
||||
}
|
||||
|
||||
interface Entry {
|
||||
@@ -68,8 +156,20 @@ interface Entry {
|
||||
// The persisted assistant row id of this run (set at bind; undefined if the
|
||||
// seed failed). Used by the attach anchor check (invariant 6).
|
||||
assistantMessageId?: string;
|
||||
// Parallel arrays: frames[i] is the SSE string, stamps[i] its step number.
|
||||
frames: string[];
|
||||
stamps: number[];
|
||||
bytes: number;
|
||||
// The running step counter used to stamp the NEXT frame (number of finish-step
|
||||
// frames seen so far).
|
||||
currentStamp: number;
|
||||
// The highest confirmed `stepsPersisted`: frames with stamp < persistedFloor are
|
||||
// on disk (safe to drop, never re-buffered). Monotonic (confirmPersistedStep).
|
||||
persistedFloor: number;
|
||||
// The highest stamp EVICTED by an overflow (unsafe) drop, -1 if none. Used to
|
||||
// detect a partially-evicted leading step when computing the coverage floor.
|
||||
overflowThroughStamp: number;
|
||||
// Sticky-for-logging only: at least one unsafe (overflow) eviction happened.
|
||||
overflowed: boolean;
|
||||
finished: boolean;
|
||||
subscribers: Set<Subscriber>;
|
||||
@@ -80,6 +180,10 @@ interface Entry {
|
||||
export class AiChatStreamRegistryService implements OnModuleDestroy {
|
||||
private readonly logger = new Logger(AiChatStreamRegistryService.name);
|
||||
private readonly entries = new Map<string, Entry>(); // key: chatId
|
||||
// Env-resolved caps (per instance) so a deployment can tune the ceiling without
|
||||
// a code change. The subscriber cap keeps the documented 2× relationship.
|
||||
readonly maxBufferBytes = resolveMaxBufferBytes();
|
||||
readonly subscriberMaxBufferedBytes = 2 * this.maxBufferBytes;
|
||||
|
||||
/**
|
||||
* Register a fresh entry at the START of a run (before any frame), so a tab
|
||||
@@ -105,7 +209,11 @@ export class AiChatStreamRegistryService implements OnModuleDestroy {
|
||||
this.entries.set(chatId, {
|
||||
runId,
|
||||
frames: [],
|
||||
stamps: [],
|
||||
bytes: 0,
|
||||
currentStamp: 0,
|
||||
persistedFloor: 0,
|
||||
overflowThroughStamp: -1,
|
||||
overflowed: false,
|
||||
finished: false,
|
||||
subscribers: new Set<Subscriber>(),
|
||||
@@ -150,6 +258,34 @@ export class AiChatStreamRegistryService implements OnModuleDestroy {
|
||||
void pump();
|
||||
}
|
||||
|
||||
/**
|
||||
* Confirm that step `stepsPersisted` (a COUNT: steps 0..stepsPersisted-1) is on
|
||||
* disk for this run, and ROTATE the ring: drop the buffered frames of those
|
||||
* now-persisted steps (stamp < stepsPersisted). This is the ONLY thing that
|
||||
* rotates the ring, and it is called ONLY after a genuinely SUCCESSFUL per-step
|
||||
* persist (see ai-chat.service updateStreaming). A failed persist never calls
|
||||
* it, so the ring covers more (auto-safe). Identity-checked (invariant 1) and
|
||||
* monotonic (a stale lower count is ignored).
|
||||
*/
|
||||
confirmPersistedStep(
|
||||
chatId: string,
|
||||
runId: string,
|
||||
stepsPersisted: number,
|
||||
): void {
|
||||
const entry = this.entries.get(chatId);
|
||||
if (!entry || entry.runId !== runId) return;
|
||||
if (!Number.isFinite(stepsPersisted) || stepsPersisted <= entry.persistedFloor)
|
||||
return;
|
||||
entry.persistedFloor = stepsPersisted;
|
||||
// Clean rotation: drop the persisted steps from the head. These frames are on
|
||||
// disk + carried by a fresh client seed, so this NEVER opens a gap.
|
||||
while (entry.frames.length > 0 && entry.stamps[0] < stepsPersisted) {
|
||||
entry.bytes -= Buffer.byteLength(entry.frames[0]);
|
||||
entry.frames.shift();
|
||||
entry.stamps.shift();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Terminate a run's entry from the OUTER catch of the stream method (a failure
|
||||
* before/while wiring the pipe, so `done` will never arrive). Identity-checked
|
||||
@@ -162,36 +298,77 @@ export class AiChatStreamRegistryService implements OnModuleDestroy {
|
||||
}
|
||||
|
||||
/**
|
||||
* Attach to a run's stream. Async only for the phase-2 Redis seam — the body
|
||||
* runs synchronously so the replay snapshot and the subscriber registration
|
||||
* happen in ONE tick with no await between them (invariant 4): a frame ingested
|
||||
* concurrently cannot slip into the gap and be lost or duplicated.
|
||||
* Attach to a run's stream from the client's step frontier `n` (its persisted
|
||||
* `stepsPersisted`). Async only for the phase-2 Redis seam — the body runs
|
||||
* synchronously so the tail SLICE and the subscriber registration happen in ONE
|
||||
* tick with no await between them (invariant 4).
|
||||
*
|
||||
* Returns null (-> the caller answers 204) when:
|
||||
* - there is no entry, or it overflowed (replay is gone);
|
||||
* - expect=live with an anchor that does not match this run's assistant id
|
||||
* (invariant 6: a stripped tab must never replay a FOREIGN run's transcript);
|
||||
* - the run finished and the caller did not expect a live tail.
|
||||
* A finished run with expect=live yields a replay-only attachment (no
|
||||
* subscriber registered). Otherwise a paused subscriber is registered and the
|
||||
* caller replays `replay`, then calls start() to drain and go live.
|
||||
* - there is no entry;
|
||||
* - the `anchor` does not match this run's assistant id (invariant 6);
|
||||
* - the ring does not cover the client's frontier (coverageFloor > n): a hole
|
||||
* from overflow, or the client's seed simply lagged behind a rotation. The
|
||||
* client then refetches (a larger n) and re-attaches.
|
||||
*
|
||||
* Otherwise the attachment's `replay` is a synthetic `start` frame (the run-fact
|
||||
* on re-attach) followed by the buffered tail filtered to `stamp >= n`. For a
|
||||
* FINISHED run this is replay-only (no subscriber) and ends after the replay —
|
||||
* with n = N_final that tail is just the run's `finish` frame, so the client
|
||||
* closes the stream. For a LIVE run a paused subscriber is registered; the
|
||||
* caller writes the replay (respecting drain) then calls start() to drain the
|
||||
* pending frames and go live.
|
||||
*/
|
||||
async attach(
|
||||
chatId: string,
|
||||
expectLive: boolean,
|
||||
anchor: string | undefined,
|
||||
// The client's persisted step frontier. `null` = a NOT-tail-aware client (no
|
||||
// `n` query param) — a legacy/parameterless tab that expects the old
|
||||
// "finished -> 204 -> poll" contract; distinct from `0` (a tail-aware client
|
||||
// with nothing persisted yet).
|
||||
n: number | null,
|
||||
cb: RunStreamCallbacks,
|
||||
): Promise<RunStreamAttachment | null> {
|
||||
const entry = this.entries.get(chatId);
|
||||
if (!entry || entry.overflowed) return null;
|
||||
if (!entry) return null;
|
||||
// Invariant 6: cross-run replay is forbidden. Before bind, assistantMessageId
|
||||
// is undefined and mismatches any anchor -> 204 -> client restore+poll path.
|
||||
if (expectLive && anchor && entry.assistantMessageId !== anchor) return null;
|
||||
if (entry.finished && !expectLive) return null;
|
||||
if (entry.finished && expectLive) {
|
||||
if (anchor && entry.assistantMessageId !== anchor) return null;
|
||||
// #491 regression guard (#137/#161 dup): a NOT-tail-aware client (no `n`)
|
||||
// resuming a FINISHED run must 204 and poll — the old `finished && !expectLive`
|
||||
// gate. Without this, a missing `n` collapsing to frontier 0 would serve the
|
||||
// WHOLE tail of a finished, NON-rotated run (coverageFloor 0), and a
|
||||
// parameterless client that never stripped its transcript would APPEND that
|
||||
// full replay onto the steps it already shows -> duplicated text. A tail-aware
|
||||
// client (n present, incl. n=0) still gets the tail past its frontier.
|
||||
if (entry.finished && n === null) return null;
|
||||
// A finished entry with NOTHING in the ring (aborted before the first frame,
|
||||
// or fully overflowed) has no tail to deliver -> 204 -> the client polls.
|
||||
if (entry.finished && entry.frames.length === 0) return null;
|
||||
// A LIVE run with no `n` (legacy parameterless) replays from step 0 (the old
|
||||
// behavior); a tail-aware client resumes from its frontier.
|
||||
const frontier = n ?? 0;
|
||||
const floor = this.coverageFloor(entry);
|
||||
if (floor > frontier) {
|
||||
this.logger.warn(
|
||||
`run-stream attach gap for run=${entry.runId}: coverageFloor=${floor} ` +
|
||||
`> client frontier=${frontier} -> 204 (client refetches + re-attaches)`,
|
||||
);
|
||||
return null;
|
||||
}
|
||||
|
||||
const startFrame = this.buildStartFrame(chatId, entry.runId);
|
||||
const sliceTail = (): string[] => {
|
||||
const out: string[] = [startFrame];
|
||||
for (let i = 0; i < entry.frames.length; i++) {
|
||||
if (entry.stamps[i] >= frontier) out.push(entry.frames[i]);
|
||||
}
|
||||
return out;
|
||||
};
|
||||
|
||||
if (entry.finished) {
|
||||
// Replay-only: the run is done, no subscriber is registered.
|
||||
return {
|
||||
replay: entry.frames.slice(),
|
||||
replay: sliceTail(),
|
||||
finished: true,
|
||||
start: () => undefined,
|
||||
unsubscribe: () => undefined,
|
||||
@@ -206,15 +383,12 @@ export class AiChatStreamRegistryService implements OnModuleDestroy {
|
||||
pendingBytes: 0,
|
||||
overflowed: false,
|
||||
pendingEnd: false,
|
||||
minStamp: frontier,
|
||||
};
|
||||
// Register + snapshot in the SAME synchronous block (invariant 4). No await
|
||||
// separates them, so a concurrently ingested frame cannot be lost/duplicated.
|
||||
entry.subscribers.add(sub);
|
||||
// Snapshot in the SAME synchronous block as the registration (invariant 4).
|
||||
const replay = entry.frames.slice();
|
||||
// CONTRACT: the caller MUST call start() in the SAME tick as this attach()
|
||||
// returns — no await between them. While a subscriber is paused, every frame
|
||||
// is buffered in sub.pending; a delayed start() lets a whole run accumulate
|
||||
// there. The pendingBytes cap (see ingestFrame) is the structural backstop if
|
||||
// that contract is ever broken (e.g. the phase-2 Redis await seam).
|
||||
const replay = sliceTail();
|
||||
return {
|
||||
replay,
|
||||
finished: false,
|
||||
@@ -263,24 +437,83 @@ export class AiChatStreamRegistryService implements OnModuleDestroy {
|
||||
this.entries.clear();
|
||||
}
|
||||
|
||||
/** Buffer + fan-out a single frame. See invariant/overflow semantics inline. */
|
||||
/** The synthetic `start` frame the tail is prefixed with — the source of the
|
||||
* run-fact (runId/chatId) on re-attach. A `start` frame does NOT reset the
|
||||
* client's message parts (ai@6.0.207 createStreamingUIMessageState), so it is
|
||||
* safe to prepend even when the sliced tail begins mid-message. */
|
||||
private buildStartFrame(chatId: string, runId: string): string {
|
||||
return `data: ${JSON.stringify({
|
||||
type: 'start',
|
||||
messageMetadata: { runId, chatId },
|
||||
})}\n\n`;
|
||||
}
|
||||
|
||||
/**
|
||||
* The smallest step FULLY present in the ring: its smallest retained stamp, or
|
||||
* (when the leading step was only partially evicted by an overflow) one past it.
|
||||
* When the ring is empty it is the current step (only the live tail is coming).
|
||||
* An attach at frontier `n` is covered ⟺ coverageFloor <= n.
|
||||
*/
|
||||
private coverageFloor(entry: Entry): number {
|
||||
// Empty ring: only the live tail is coming. The floor is the current step,
|
||||
// but never below persistedFloor — a confirmed persist can rotate the ring
|
||||
// empty while currentStamp still lags a beat behind on another connection, so
|
||||
// max() keeps the invariant STRUCTURAL (a client with n = persistedFloor is
|
||||
// always covered) rather than timing-dependent.
|
||||
if (entry.frames.length === 0)
|
||||
return Math.max(entry.currentStamp, entry.persistedFloor);
|
||||
const min = entry.stamps[0];
|
||||
return entry.overflowThroughStamp >= min ? min + 1 : min;
|
||||
}
|
||||
|
||||
/**
|
||||
* Buffer (step-stamped) + fan-out a single frame. The stamp is the number of
|
||||
* finish-step frames seen BEFORE this one; a finish-step frame carries the
|
||||
* current value and THEN increments the counter (so its stamp equals the 0-based
|
||||
* index of the step it closes). Only frames at/above persistedFloor are buffered
|
||||
* (already-persisted steps are on disk); the ring is then trimmed to the byte
|
||||
* cap, an unsafe eviction opening a gap. Fan-out is always live (filtered per
|
||||
* subscriber by its frontier).
|
||||
*/
|
||||
private ingestFrame(entry: Entry, frame: string): void {
|
||||
entry.bytes += Buffer.byteLength(frame);
|
||||
if (!entry.overflowed) {
|
||||
const size = Buffer.byteLength(frame);
|
||||
const stamp = entry.currentStamp;
|
||||
if (frame.startsWith(FINISH_STEP_FRAME_PREFIX)) {
|
||||
entry.currentStamp = stamp + 1;
|
||||
}
|
||||
|
||||
// Buffer for replay only if this step is not already persisted+rotated away.
|
||||
if (stamp >= entry.persistedFloor) {
|
||||
entry.frames.push(frame);
|
||||
if (entry.bytes > RUN_STREAM_MAX_BUFFER_BYTES) {
|
||||
// The crossing frame was already counted AND (below) fanned out; only the
|
||||
// replay buffer is dropped. After overflow no more frames are buffered,
|
||||
// but live fan-out continues.
|
||||
entry.overflowed = true;
|
||||
entry.frames = [];
|
||||
this.logger.warn(
|
||||
`run-stream buffer overflow for run=${entry.runId}; ` +
|
||||
`late attach will 204 until the run ends`,
|
||||
);
|
||||
entry.stamps.push(stamp);
|
||||
entry.bytes += size;
|
||||
// Enforce the ring cap. Evicting a not-yet-persisted frame (stamp >=
|
||||
// persistedFloor) opens a GAP; a leftover persisted frame (< floor) is a
|
||||
// safe drop. Keep evicting until the ring is back under the cap.
|
||||
while (entry.bytes > this.maxBufferBytes && entry.frames.length > 0) {
|
||||
const evStamp = entry.stamps[0];
|
||||
entry.bytes -= Buffer.byteLength(entry.frames[0]);
|
||||
entry.frames.shift();
|
||||
entry.stamps.shift();
|
||||
if (evStamp >= entry.persistedFloor) {
|
||||
if (evStamp > entry.overflowThroughStamp)
|
||||
entry.overflowThroughStamp = evStamp;
|
||||
if (!entry.overflowed) {
|
||||
entry.overflowed = true;
|
||||
this.logger.warn(
|
||||
`run-stream ring overflow for run=${entry.runId}: an un-persisted ` +
|
||||
`step was evicted to stay under ${this.maxBufferBytes}B; a late ` +
|
||||
`attach at an evicted step will 204 until a later persist confirms`,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fan out live, filtered to each subscriber's frontier (a subscriber only
|
||||
// wants the tail past the step it already persisted).
|
||||
for (const sub of entry.subscribers) {
|
||||
if (stamp < sub.minStamp) continue;
|
||||
if (sub.started) {
|
||||
try {
|
||||
sub.onFrame(frame);
|
||||
@@ -289,12 +522,12 @@ export class AiChatStreamRegistryService implements OnModuleDestroy {
|
||||
}
|
||||
} else {
|
||||
sub.pending.push(frame);
|
||||
sub.pendingBytes += Buffer.byteLength(frame);
|
||||
if (sub.pendingBytes > SUBSCRIBER_MAX_BUFFERED_BYTES) {
|
||||
sub.pendingBytes += size;
|
||||
if (sub.pendingBytes > this.subscriberMaxBufferedBytes) {
|
||||
// The paused subscriber's buffer overflowed — only possible if start()
|
||||
// was delayed past the same-tick contract (the phase-2 await seam).
|
||||
// Drop it rather than buffer the whole run; on start() it degrades to an
|
||||
// immediate end (a 204-equivalent) instead of replaying a partial.
|
||||
// was delayed (the controller's drain-respecting tail write, or the
|
||||
// phase-2 await seam). Drop it rather than buffer the whole run; on
|
||||
// start() it degrades to an immediate end (a 204-equivalent).
|
||||
sub.overflowed = true;
|
||||
sub.pending = [];
|
||||
entry.subscribers.delete(sub);
|
||||
|
||||
@@ -1,19 +1,27 @@
|
||||
import {
|
||||
AiChatStreamRegistryService,
|
||||
RUN_STREAM_MAX_BUFFER_BYTES,
|
||||
AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES,
|
||||
RUN_STREAM_RETAIN_FINISHED_MS,
|
||||
SUBSCRIBER_MAX_BUFFERED_BYTES,
|
||||
RunStreamCallbacks,
|
||||
} from './ai-chat-stream-registry.service';
|
||||
|
||||
/**
|
||||
* Unit tests for the in-memory run-stream registry (#184 phase 1.5). The registry
|
||||
* is the whole of the resumable-transport contract: replay ordering, paused ->
|
||||
* live hand-off, overflow, retention, the anchor check (invariant 6), and the
|
||||
* mirror-the-done-path replace semantics (invariant 3). Every enumerated case in
|
||||
* the issue's task 1.5 has a test here.
|
||||
* Unit tests for the in-memory run-stream registry (#184 phase 1.5, step-aligned
|
||||
* retention #491). The registry is the whole of the resumable-transport contract:
|
||||
* step-stamped retention, tail-only attach at the client's frontier N, the
|
||||
* confirmed-persist ring rotation (and the anti-inversion rule), the memory bound,
|
||||
* the overflow gap, paused -> live hand-off, retention, the anchor check
|
||||
* (invariant 6), and the mirror-the-done-path replace semantics (invariant 3).
|
||||
*/
|
||||
|
||||
// Real ai@6 UI-message-stream SSE frames are `data: {json}\n\n`, one part each.
|
||||
const sse = (part: Record<string, unknown>): string =>
|
||||
`data: ${JSON.stringify(part)}\n\n`;
|
||||
const finishStep = (): string => sse({ type: 'finish-step' });
|
||||
const textDelta = (id: string, delta: string): string =>
|
||||
sse({ type: 'text-delta', id, delta });
|
||||
const finish = (): string => sse({ type: 'finish' });
|
||||
|
||||
// A ReadableStream whose frames the test pushes explicitly, plus close/error.
|
||||
function makePushStream(): {
|
||||
stream: ReadableStream<string>;
|
||||
@@ -58,6 +66,9 @@ function collector(): {
|
||||
};
|
||||
}
|
||||
|
||||
// The tail past the synthetic start frame (replay[0] is always the start frame).
|
||||
const tail = (replay: string[]): string[] => replay.slice(1);
|
||||
|
||||
describe('AiChatStreamRegistryService', () => {
|
||||
const CHAT = 'chat-1';
|
||||
let registry: AiChatStreamRegistryService;
|
||||
@@ -71,7 +82,21 @@ describe('AiChatStreamRegistryService', () => {
|
||||
registry.onModuleDestroy();
|
||||
});
|
||||
|
||||
it('replays frames in arrival order (live attach)', async () => {
|
||||
it('prepends a synthetic start frame carrying { runId, chatId }', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push('a');
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!;
|
||||
const start = JSON.parse(att.replay[0].replace(/^data: /, '').trim());
|
||||
expect(start.type).toBe('start');
|
||||
expect(start.messageMetadata).toEqual({ runId: 'run-1', chatId: CHAT });
|
||||
});
|
||||
|
||||
it('replays the buffered tail (from frontier 0) in arrival order (live attach)', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
@@ -81,13 +106,13 @@ describe('AiChatStreamRegistryService', () => {
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
const att = await registry.attach(CHAT, false, undefined, c.cb);
|
||||
const att = await registry.attach(CHAT, 'assist-1', 0, c.cb);
|
||||
expect(att).not.toBeNull();
|
||||
expect(att!.replay).toEqual(['a', 'b', 'c']);
|
||||
expect(tail(att!.replay)).toEqual(['a', 'b', 'c']);
|
||||
expect(att!.finished).toBe(false);
|
||||
});
|
||||
|
||||
it('late attach gets the full prefix as replay plus the live tail', async () => {
|
||||
it('late attach gets the buffered prefix as tail plus the live tail', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
@@ -96,17 +121,16 @@ describe('AiChatStreamRegistryService', () => {
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
const att = (await registry.attach(CHAT, false, undefined, c.cb))!;
|
||||
expect(att.replay).toEqual(['a', 'b']);
|
||||
const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!;
|
||||
expect(tail(att.replay)).toEqual(['a', 'b']);
|
||||
att.start();
|
||||
// Live tail arrives after start().
|
||||
src.push('c');
|
||||
src.push('d');
|
||||
await flush();
|
||||
expect(c.frames).toEqual(['c', 'd']);
|
||||
});
|
||||
|
||||
it('a paused subscriber receives frames buffered during pause in order, then live (no loss/reorder)', async () => {
|
||||
it('a paused subscriber receives frames buffered during pause in order, then live', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
@@ -114,81 +138,45 @@ describe('AiChatStreamRegistryService', () => {
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
// Attach (paused). Frames that arrive BEFORE start() must queue, not drop.
|
||||
const att = (await registry.attach(CHAT, false, undefined, c.cb))!;
|
||||
expect(att.replay).toEqual(['a']);
|
||||
const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!;
|
||||
expect(tail(att.replay)).toEqual(['a']);
|
||||
src.push('b'); // arrives while paused -> pending
|
||||
src.push('c');
|
||||
await flush();
|
||||
expect(c.frames).toEqual([]); // nothing delivered yet (paused)
|
||||
att.start(); // drains pending in order
|
||||
att.start();
|
||||
expect(c.frames).toEqual(['b', 'c']);
|
||||
src.push('d'); // now live
|
||||
src.push('d');
|
||||
await flush();
|
||||
expect(c.frames).toEqual(['b', 'c', 'd']);
|
||||
});
|
||||
|
||||
it('a run that finishes while a subscriber is paused ends it on start()', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', makePushStream().stream);
|
||||
const c = collector();
|
||||
const att = (await registry.attach(CHAT, false, undefined, c.cb))!;
|
||||
// Terminate the run while the subscriber is still paused.
|
||||
const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!;
|
||||
registry.abortEntry(CHAT, 'run-1');
|
||||
expect(c.ended()).toBe(0); // paused: not ended yet
|
||||
att.start();
|
||||
expect(c.ended()).toBe(1); // start() drains + ends
|
||||
});
|
||||
|
||||
it('finished + expect=live returns a replay WITHOUT registering a subscriber', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push('a');
|
||||
src.push('b');
|
||||
src.close();
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
const att = (await registry.attach(CHAT, true, undefined, c.cb))!;
|
||||
expect(att.finished).toBe(true);
|
||||
expect(att.replay).toEqual(['a', 'b']);
|
||||
// No subscriber registered: start()/unsubscribe are no-ops and the entry has
|
||||
// zero subscribers.
|
||||
const entry = (registry as any).entries.get(CHAT);
|
||||
expect(entry.subscribers.size).toBe(0);
|
||||
att.start();
|
||||
expect(c.frames).toEqual([]);
|
||||
});
|
||||
|
||||
it('finished WITHOUT expect=live returns null', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push('a');
|
||||
src.close();
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
expect(await registry.attach(CHAT, false, undefined, c.cb)).toBeNull();
|
||||
});
|
||||
|
||||
it('anchor mismatch with expect=live returns null (and null before bind sets assistantMessageId)', async () => {
|
||||
it('anchor mismatch returns null (and null before bind sets assistantMessageId)', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const c = collector();
|
||||
// Before bind: assistantMessageId is undefined -> mismatches any anchor.
|
||||
expect(
|
||||
await registry.attach(CHAT, true, 'assist-1', c.cb),
|
||||
).toBeNull();
|
||||
expect(await registry.attach(CHAT, 'assist-1', 0, c.cb)).toBeNull();
|
||||
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push('a');
|
||||
await flush();
|
||||
// Wrong anchor -> null (cross-run replay forbidden, invariant 6).
|
||||
expect(await registry.attach(CHAT, true, 'other-id', c.cb)).toBeNull();
|
||||
expect(await registry.attach(CHAT, 'other-id', 0, c.cb)).toBeNull();
|
||||
});
|
||||
|
||||
it('matching anchor with expect=live attaches', async () => {
|
||||
it('matching anchor attaches', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
@@ -196,97 +184,60 @@ describe('AiChatStreamRegistryService', () => {
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
const att = await registry.attach(CHAT, true, 'assist-1', c.cb);
|
||||
const att = await registry.attach(CHAT, 'assist-1', 0, c.cb);
|
||||
expect(att).not.toBeNull();
|
||||
expect(att!.replay).toEqual(['a']);
|
||||
expect(tail(att!.replay)).toEqual(['a']);
|
||||
});
|
||||
|
||||
it('overflow: attach returns null, but the LIVE subscriber keeps receiving (incl. the crossing frame)', async () => {
|
||||
it('a throwing onFrame ejects only that subscriber; the ingest loop stays alive', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
|
||||
// A live (started) subscriber attached before the flood.
|
||||
const bad = collector();
|
||||
const badAtt = (await registry.attach(CHAT, 'assist-1', 0, {
|
||||
onFrame: () => {
|
||||
throw new Error('boom');
|
||||
},
|
||||
onEnd: bad.cb.onEnd,
|
||||
}))!;
|
||||
badAtt.start();
|
||||
|
||||
const good = collector();
|
||||
const goodAtt = (await registry.attach(CHAT, 'assist-1', 0, good.cb))!;
|
||||
goodAtt.start();
|
||||
|
||||
src.push('a'); // bad throws on this frame -> ejected
|
||||
src.push('b'); // good still receives both
|
||||
await flush();
|
||||
|
||||
const entry = (registry as any).entries.get(CHAT);
|
||||
expect(entry.subscribers.size).toBe(1);
|
||||
expect(good.frames).toEqual(['a', 'b']);
|
||||
});
|
||||
|
||||
it('open() over a LIVE entry ends started subscribers once; a late done never touches the new entry (invariant 3)', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push('a');
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
const att = (await registry.attach(CHAT, false, undefined, c.cb))!;
|
||||
const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!;
|
||||
att.start();
|
||||
|
||||
// Cap-relative so it survives a buffer-cap change (#430): a quarter-cap frame
|
||||
// means 5 frames comfortably exceed the replay cap; the last one crosses.
|
||||
const chunk = 'x'.repeat(Math.floor(RUN_STREAM_MAX_BUFFER_BYTES / 4));
|
||||
for (let i = 0; i < 5; i++) src.push(chunk + i);
|
||||
await flush();
|
||||
|
||||
const entry = (registry as any).entries.get(CHAT);
|
||||
expect(entry.overflowed).toBe(true);
|
||||
expect(entry.bytes).toBeGreaterThan(RUN_STREAM_MAX_BUFFER_BYTES);
|
||||
// The live subscriber received ALL 5 frames, including the crossing one.
|
||||
expect(c.frames).toHaveLength(5);
|
||||
expect(c.frames[4]).toBe(chunk + 4);
|
||||
|
||||
// A NEW attach after overflow gets null (replay buffer is gone).
|
||||
const c2 = collector();
|
||||
expect(await registry.attach(CHAT, false, undefined, c2.cb)).toBeNull();
|
||||
});
|
||||
|
||||
it('a paused subscriber whose pending buffer overflows is dropped and ends on start(); other subscribers keep receiving', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
|
||||
// A: paused (start() deliberately delayed to simulate the phase-2 await seam).
|
||||
const a = collector();
|
||||
const attA = (await registry.attach(CHAT, false, undefined, a.cb))!;
|
||||
// B: live (started) — its delivery must be unaffected by A's overflow.
|
||||
const b = collector();
|
||||
const attB = (await registry.attach(CHAT, false, undefined, b.cb))!;
|
||||
attB.start();
|
||||
|
||||
// Cap-relative so it survives a buffer-cap change (#430): a quarter-of-the-
|
||||
// per-subscriber-cap frame means 5 frames exceed A's paused-pending cap while
|
||||
// B streams every frame live.
|
||||
const chunk = 'x'.repeat(Math.floor(SUBSCRIBER_MAX_BUFFERED_BYTES / 4));
|
||||
for (let i = 0; i < 5; i++) src.push(chunk + i);
|
||||
await flush();
|
||||
|
||||
const entry = (registry as any).entries.get(CHAT);
|
||||
// A was dropped from the subscriber set on overflow; B (started) remains.
|
||||
expect(entry.subscribers.size).toBe(1);
|
||||
expect(a.frames).toEqual([]); // paused + overflowed: nothing was delivered
|
||||
// B received every frame live (delivery unaffected by A's overflow).
|
||||
expect(b.frames).toHaveLength(5);
|
||||
|
||||
// A's start() (arriving late) degrades to an immediate end, not a partial replay.
|
||||
attA.start();
|
||||
expect(a.frames).toEqual([]);
|
||||
expect(a.ended()).toBe(1);
|
||||
});
|
||||
|
||||
it('open() over a LIVE entry ends started subscribers exactly once and a late done does not touch the new entry (invariant 3)', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push('a');
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
const att = (await registry.attach(CHAT, false, undefined, c.cb))!;
|
||||
att.start(); // started subscriber on run-1
|
||||
|
||||
// run-2 starts on the same chat while run-1's tee is still reading.
|
||||
registry.open(CHAT, 'run-2');
|
||||
expect(c.ended()).toBe(1); // exactly one onEnd from the replace
|
||||
expect(c.ended()).toBe(1);
|
||||
|
||||
const newEntry = (registry as any).entries.get(CHAT);
|
||||
expect(newEntry.runId).toBe('run-2');
|
||||
expect(newEntry.finished).toBe(false);
|
||||
|
||||
// The old tee now completes: its late done must NOT double-end nor delete the
|
||||
// new entry.
|
||||
src.push('b');
|
||||
src.close();
|
||||
await flush();
|
||||
expect(c.ended()).toBe(1); // still exactly one
|
||||
expect(c.ended()).toBe(1);
|
||||
const still = (registry as any).entries.get(CHAT);
|
||||
expect(still).toBe(newEntry);
|
||||
expect(still.runId).toBe('run-2');
|
||||
@@ -299,7 +250,6 @@ describe('AiChatStreamRegistryService', () => {
|
||||
src.push('a');
|
||||
await flush();
|
||||
const entry = (registry as any).entries.get(CHAT);
|
||||
// Frames were NOT ingested (bind bailed), assistantMessageId untouched.
|
||||
expect(entry.frames).toEqual([]);
|
||||
expect(entry.assistantMessageId).toBeUndefined();
|
||||
});
|
||||
@@ -310,32 +260,276 @@ describe('AiChatStreamRegistryService', () => {
|
||||
const entry = (registry as any).entries.get(CHAT);
|
||||
expect(entry.finished).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
it('a throwing onFrame ejects only that subscriber; the ingest loop stays alive', async () => {
|
||||
/**
|
||||
* #491 step-stamped retention: the boundary detector, tail-only slicing at the
|
||||
* client's frontier N, the confirmed-persist rotation (+ anti-inversion), the
|
||||
* overflow gap, the memory bound, and the finished-retained tail. All observable
|
||||
* against the REAL registry driven through open/bind/ingest.
|
||||
*/
|
||||
describe('AiChatStreamRegistryService step-aligned retention (#491)', () => {
|
||||
const CHAT = 'chat-s';
|
||||
let registry: AiChatStreamRegistryService;
|
||||
|
||||
beforeEach(() => {
|
||||
registry = new AiChatStreamRegistryService();
|
||||
jest.spyOn((registry as any).logger, 'warn').mockImplementation(() => {});
|
||||
});
|
||||
afterEach(() => registry.onModuleDestroy());
|
||||
|
||||
const entryOf = () => (registry as any).entries.get(CHAT);
|
||||
|
||||
it('stamps frames by finish-step count, aligned with stepsPersisted', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
// step 0 content, its finish-step, step 1 content, its finish-step, finish.
|
||||
src.push(textDelta('t0', 'a')); // stamp 0
|
||||
src.push(finishStep()); // stamp 0 (the finish-step frame carries the pre value)
|
||||
src.push(textDelta('t1', 'b')); // stamp 1
|
||||
src.push(finishStep()); // stamp 1
|
||||
src.push(finish()); // stamp 2
|
||||
await flush();
|
||||
const e = entryOf();
|
||||
expect(e.stamps).toEqual([0, 0, 1, 1, 2]);
|
||||
expect(e.currentStamp).toBe(2);
|
||||
});
|
||||
|
||||
const bad = collector();
|
||||
const badAtt = (await registry.attach(CHAT, false, undefined, {
|
||||
onFrame: () => {
|
||||
throw new Error('boom');
|
||||
},
|
||||
onEnd: bad.cb.onEnd,
|
||||
}))!;
|
||||
badAtt.start();
|
||||
it('does NOT treat a text delta that merely quotes "finish-step" as a boundary', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
// A model that literally types "type":"finish-step" — JSON-escaped in the frame.
|
||||
src.push(textDelta('t0', '"type":"finish-step"'));
|
||||
await flush();
|
||||
expect(entryOf().currentStamp).toBe(0); // no false boundary
|
||||
});
|
||||
|
||||
const good = collector();
|
||||
const goodAtt = (await registry.attach(CHAT, false, undefined, good.cb))!;
|
||||
goodAtt.start();
|
||||
|
||||
src.push('a'); // bad throws on this frame -> ejected
|
||||
src.push('b'); // good still receives both
|
||||
it('tail-only: attach at N slices frames with stamp >= N', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push(textDelta('t0', 'a')); // 0
|
||||
src.push(finishStep()); // 0
|
||||
src.push(textDelta('t1', 'b')); // 1
|
||||
src.push(finishStep()); // 1
|
||||
src.push(textDelta('t2', 'c')); // 2 (in-progress)
|
||||
await flush();
|
||||
|
||||
const entry = (registry as any).entries.get(CHAT);
|
||||
expect(entry.subscribers.size).toBe(1); // bad ejected, good remains
|
||||
expect(good.frames).toEqual(['a', 'b']);
|
||||
const c = collector();
|
||||
// Client persisted 2 steps -> wants the tail from step 2.
|
||||
const att = (await registry.attach(CHAT, 'assist-1', 2, c.cb))!;
|
||||
expect(tail(att.replay)).toEqual([textDelta('t2', 'c')]);
|
||||
});
|
||||
|
||||
it('attach in the MIDDLE of a step (N between finish-steps) slices from that step', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push(textDelta('t0', 'a')); // 0
|
||||
src.push(finishStep()); // 0
|
||||
src.push(textDelta('t1', 'b1')); // 1
|
||||
src.push(textDelta('t1', 'b2')); // 1 (still step 1, no finish-step yet)
|
||||
await flush();
|
||||
|
||||
const c = collector();
|
||||
const att = (await registry.attach(CHAT, 'assist-1', 1, c.cb))!;
|
||||
// Step 0's frames are dropped from the tail; the whole in-progress step 1 is kept.
|
||||
expect(tail(att.replay)).toEqual([textDelta('t1', 'b1'), textDelta('t1', 'b2')]);
|
||||
});
|
||||
|
||||
it('rotates the ring ONLY on a confirmed persist (drops stamp < N)', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push(textDelta('t0', 'a')); // 0
|
||||
src.push(finishStep()); // 0
|
||||
src.push(textDelta('t1', 'b')); // 1
|
||||
await flush();
|
||||
expect(entryOf().stamps).toEqual([0, 0, 1]);
|
||||
|
||||
// Confirm step 0 persisted (stepsPersisted = 1) -> drop stamp < 1.
|
||||
registry.confirmPersistedStep(CHAT, 'run-1', 1);
|
||||
expect(entryOf().stamps).toEqual([1]);
|
||||
expect(entryOf().persistedFloor).toBe(1);
|
||||
});
|
||||
|
||||
it('persist FAILED but the ring still fits -> attach SUCCEEDS and the tail includes step N', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push(textDelta('t0', 'a')); // 0
|
||||
src.push(finishStep()); // 0
|
||||
src.push(textDelta('t1', 'b')); // 1 (step 1's persist FAILED -> no confirm)
|
||||
await flush();
|
||||
// No confirmPersistedStep for step 1: the ring still holds step 1.
|
||||
|
||||
const c = collector();
|
||||
// Client's last successful persist was step 0 -> stepsPersisted = 1.
|
||||
const att = await registry.attach(CHAT, 'assist-1', 1, c.cb);
|
||||
expect(att).not.toBeNull();
|
||||
expect(tail(att!.replay)).toEqual([textDelta('t1', 'b')]); // includes step 1
|
||||
});
|
||||
|
||||
it('persist failed AND the ring overflowed past N -> 204 (coverage gap)', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
// Step 0: a fat step that blows past the cap with NO persist confirmation.
|
||||
const big = 'x'.repeat(Math.floor(AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES / 2));
|
||||
src.push(textDelta('t0', big)); // 0
|
||||
src.push(textDelta('t0', big)); // 0
|
||||
src.push(textDelta('t0', big)); // 0 -> overflow evicts stamp-0 frames
|
||||
await flush();
|
||||
const e = entryOf();
|
||||
expect(e.overflowed).toBe(true);
|
||||
expect(e.bytes).toBeLessThanOrEqual(registry.maxBufferBytes);
|
||||
|
||||
// A client at frontier 0 falls at/below an evicted step -> gap -> null.
|
||||
const c = collector();
|
||||
expect(await registry.attach(CHAT, 'assist-1', 0, c.cb)).toBeNull();
|
||||
});
|
||||
|
||||
it('stale N (client seed lagged behind a rotation) -> 204; after a refetch (larger N) -> success', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push(textDelta('t0', 'a')); // 0
|
||||
src.push(finishStep()); // 0
|
||||
src.push(textDelta('t1', 'b')); // 1
|
||||
src.push(finishStep()); // 1
|
||||
src.push(textDelta('t2', 'c')); // 2
|
||||
await flush();
|
||||
// Server confirmed steps 0 and 1 -> rotate away stamp < 2.
|
||||
registry.confirmPersistedStep(CHAT, 'run-1', 2);
|
||||
expect(entryOf().stamps).toEqual([2]);
|
||||
|
||||
// A client whose seed still says stepsPersisted = 1 -> below minStamp -> 204.
|
||||
const stale = collector();
|
||||
expect(await registry.attach(CHAT, 'assist-1', 1, stale.cb)).toBeNull();
|
||||
|
||||
// It refetches (now stepsPersisted = 2) and re-attaches -> success.
|
||||
const fresh = collector();
|
||||
const att = await registry.attach(CHAT, 'assist-1', 2, fresh.cb);
|
||||
expect(att).not.toBeNull();
|
||||
expect(tail(att!.replay)).toEqual([textDelta('t2', 'c')]);
|
||||
});
|
||||
|
||||
it('overflow gap CLEARS once a later persist rotates out the holey steps', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
const big = 'x'.repeat(Math.floor(AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES / 2));
|
||||
src.push(textDelta('t0', big)); // 0
|
||||
src.push(textDelta('t0', big)); // 0
|
||||
src.push(finishStep()); // 0 (still stamp 0)
|
||||
src.push(textDelta('t1', 'small')); // 1
|
||||
src.push(finishStep()); // 1
|
||||
src.push(textDelta('t2', 'c')); // 2
|
||||
await flush();
|
||||
expect(entryOf().overflowed).toBe(true);
|
||||
|
||||
// Late persist confirms steps 0..1 -> rotates out the holey step-0 frames.
|
||||
registry.confirmPersistedStep(CHAT, 'run-1', 2);
|
||||
// A client at frontier 2 is now cleanly covered (the hole was below it).
|
||||
const c = collector();
|
||||
const att = await registry.attach(CHAT, 'assist-1', 2, c.cb);
|
||||
expect(att).not.toBeNull();
|
||||
expect(tail(att!.replay)).toEqual([textDelta('t2', 'c')]);
|
||||
});
|
||||
|
||||
it('finished-retained + N = N_final -> empty tail plus the finish frame', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push(textDelta('t0', 'a')); // 0
|
||||
src.push(finishStep()); // 0
|
||||
src.push(finish()); // 1 (N_final = 1)
|
||||
src.close();
|
||||
await flush();
|
||||
// The last step's per-step persist confirmed stepsPersisted = 1.
|
||||
registry.confirmPersistedStep(CHAT, 'run-1', 1);
|
||||
|
||||
const c = collector();
|
||||
const att = (await registry.attach(CHAT, 'assist-1', 1, c.cb))!;
|
||||
expect(att.finished).toBe(true);
|
||||
// Empty step tail; just the finish frame so the client's SDK closes the stream.
|
||||
expect(tail(att.replay)).toEqual([finish()]);
|
||||
// No subscriber registered for a finished run.
|
||||
expect(entryOf().subscribers.size).toBe(0);
|
||||
});
|
||||
|
||||
it('#491 regression (#137/#161 dup): a PARAMETERLESS attach (n=null) to a finished NON-rotated run -> 204, but n=0 still gets the tail', async () => {
|
||||
// A finished, non-rotated run: frames present, coverageFloor 0. A missing `n`
|
||||
// (null — a legacy/parameterless tab that never stripped its transcript) must
|
||||
// 204 -> poll, NOT receive the whole tail it would append (duplicate). A
|
||||
// tail-aware client (n=0 present) still resumes.
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push(textDelta('t0', 'a')); // 0
|
||||
src.push(finishStep()); // 0
|
||||
src.push(finish()); // 1
|
||||
src.close();
|
||||
await flush();
|
||||
// NOT rotated (no confirmPersistedStep) -> stamps[0]=0, coverageFloor=0.
|
||||
// MUTATION-VERIFY: revert the `finished && n === null -> null` gate (default n
|
||||
// to 0) and the parameterless attach below serves the full tail instead of 204.
|
||||
expect(await registry.attach(CHAT, 'assist-1', null, collector().cb)).toBeNull();
|
||||
// A tail-aware client at frontier 0 IS served (the distinction: null != 0).
|
||||
const tailAware = await registry.attach(CHAT, 'assist-1', 0, collector().cb);
|
||||
expect(tailAware).not.toBeNull();
|
||||
expect(tailAware!.finished).toBe(true);
|
||||
});
|
||||
|
||||
it('confirmPersistedStep is monotonic and identity-checked', async () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
const src = makePushStream();
|
||||
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
|
||||
src.push(textDelta('t0', 'a'));
|
||||
src.push(finishStep());
|
||||
src.push(textDelta('t1', 'b'));
|
||||
await flush();
|
||||
registry.confirmPersistedStep(CHAT, 'run-1', 1);
|
||||
expect(entryOf().persistedFloor).toBe(1);
|
||||
// A stale lower count is ignored.
|
||||
registry.confirmPersistedStep(CHAT, 'run-1', 0);
|
||||
expect(entryOf().persistedFloor).toBe(1);
|
||||
// A foreign runId is ignored.
|
||||
registry.confirmPersistedStep(CHAT, 'WRONG', 5);
|
||||
expect(entryOf().persistedFloor).toBe(1);
|
||||
});
|
||||
|
||||
it('MEMORY BOUND: 5 parallel marathon runs each stream well past 32MB; each ring stays <= the cap', async () => {
|
||||
const cap = registry.maxBufferBytes;
|
||||
const chats = ['m0', 'm1', 'm2', 'm3', 'm4'];
|
||||
const srcs = chats.map((chat) => {
|
||||
registry.open(chat, `run-${chat}`);
|
||||
const s = makePushStream();
|
||||
registry.bind(chat, `run-${chat}`, `assist-${chat}`, s.stream);
|
||||
return s;
|
||||
});
|
||||
// ~256KB frames; 160 per chat = 40MB streamed each, well past the old 32MB.
|
||||
// Interleave a finish-step every 8 frames so steps advance realistically. No
|
||||
// persist confirmation -> the ONLY thing keeping memory bounded is the cap.
|
||||
const frame = 'y'.repeat(256 * 1024);
|
||||
for (let batch = 0; batch < 20; batch++) {
|
||||
for (let i = 0; i < 8; i++) {
|
||||
for (const s of srcs) s.push(textDelta('t', frame));
|
||||
}
|
||||
for (const s of srcs) s.push(finishStep());
|
||||
await flush(); // drain the pump so queues never hold a whole run
|
||||
}
|
||||
let total = 0;
|
||||
for (const chat of chats) {
|
||||
const e = (registry as any).entries.get(chat);
|
||||
expect(e.bytes).toBeLessThanOrEqual(cap);
|
||||
total += e.bytes;
|
||||
}
|
||||
// Total retained across all 5 runs is bounded by 5x the per-run cap — the old
|
||||
// registry would have retained ~5x40MB = 200MB here.
|
||||
expect(total).toBeLessThanOrEqual(cap * chats.length);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -361,7 +555,7 @@ describe('AiChatStreamRegistryService retention timers', () => {
|
||||
|
||||
it('a finished entry is removed after the retention window', () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
registry.abortEntry(CHAT, 'run-1'); // finalize -> retention armed
|
||||
registry.abortEntry(CHAT, 'run-1');
|
||||
expect((registry as any).entries.get(CHAT)).toBeDefined();
|
||||
jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1);
|
||||
expect((registry as any).entries.get(CHAT)).toBeUndefined();
|
||||
@@ -369,20 +563,18 @@ describe('AiChatStreamRegistryService retention timers', () => {
|
||||
|
||||
it('retention deletes ONLY its own entry (invariant 2)', () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
registry.abortEntry(CHAT, 'run-1'); // arm retention for entry A
|
||||
// Simulate the race where the key was replaced without clearing A's timer.
|
||||
registry.abortEntry(CHAT, 'run-1');
|
||||
const sentinel = { marker: true };
|
||||
(registry as any).entries.set(CHAT, sentinel);
|
||||
jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1);
|
||||
// A's timer saw entries.get(CHAT) !== A, so it did NOT delete the successor.
|
||||
expect((registry as any).entries.get(CHAT)).toBe(sentinel);
|
||||
});
|
||||
|
||||
it('open() over a retained entry clears its timer and the successor survives', () => {
|
||||
registry.open(CHAT, 'run-1');
|
||||
registry.abortEntry(CHAT, 'run-1'); // retained, timer armed
|
||||
registry.abortEntry(CHAT, 'run-1');
|
||||
const clearSpy = jest.spyOn(global, 'clearTimeout');
|
||||
registry.open(CHAT, 'run-2'); // must clear run-1's retain timer
|
||||
registry.open(CHAT, 'run-2');
|
||||
expect(clearSpy).toHaveBeenCalled();
|
||||
jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1);
|
||||
const entry = (registry as any).entries.get(CHAT);
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
@@ -8,10 +8,12 @@ import { SUBSCRIBER_MAX_BUFFERED_BYTES } from './ai-chat-stream-registry.service
|
||||
import type { User, Workspace } from '@docmost/db/types/entity.types';
|
||||
|
||||
/**
|
||||
* Wiring spec for the #184 phase 1.5 attach endpoint
|
||||
* Wiring spec for the #184 phase 1.5 attach endpoint (tail-only #491)
|
||||
* (`GET /ai-chat/runs/:chatId/stream`). Owner-gated via assertOwnedChat; the
|
||||
* registry is mocked so this exercises ONLY the controller's replay/live/204/
|
||||
* cleanup wiring against a fake raw socket. Constructor order is (aiChatService,
|
||||
* registry is mocked so this exercises ONLY the controller's tail-write/live/204/
|
||||
* cleanup wiring against a fake raw socket. The attach signature is now
|
||||
* `(chatId, anchor, n, cb)` — the client hands its persisted step frontier `n`
|
||||
* and its assistant row id `anchor`. Constructor order is (aiChatService,
|
||||
* aiChatRunService, aiChatRepo, aiChatMessageRepo, aiTranscription, pageRepo,
|
||||
* streamRegistry, environment).
|
||||
*/
|
||||
@@ -86,8 +88,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
|
||||
attach: jest.fn(
|
||||
(
|
||||
_chatId: string,
|
||||
_live: boolean,
|
||||
_anchor: string | undefined,
|
||||
_n: number,
|
||||
cb: RunStreamCallbacks,
|
||||
) => {
|
||||
capturedCb = cb;
|
||||
@@ -156,7 +158,7 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
|
||||
expect(res.hijack).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('threads expect=live and anchor through to the registry', async () => {
|
||||
it('threads anchor and the numeric frontier n through to the registry', async () => {
|
||||
const { controller, streamRegistry } = makeController({
|
||||
chat: owned,
|
||||
attachment: null,
|
||||
@@ -165,8 +167,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
|
||||
const { req } = makeReq();
|
||||
await controller.attachRunStream(
|
||||
'c1',
|
||||
'live',
|
||||
'anchor-1',
|
||||
'2',
|
||||
req,
|
||||
res,
|
||||
user,
|
||||
@@ -174,13 +176,44 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
|
||||
);
|
||||
expect(streamRegistry.attach).toHaveBeenCalledWith(
|
||||
'c1',
|
||||
true,
|
||||
'anchor-1',
|
||||
2, // parsed to a number
|
||||
expect.anything(),
|
||||
);
|
||||
});
|
||||
|
||||
it('passes expect=false when the query is absent', async () => {
|
||||
it('#491: an ABSENT/invalid n passes null (not 0) so a finished run 204s (not-tail-aware)', async () => {
|
||||
// Distinguishing a MISSING `n` from `n=0` is the #137/#161 dup guard: a
|
||||
// parameterless/legacy tab must be handed null (-> the registry 204s a finished
|
||||
// run) rather than frontier 0 (which would serve a finished non-rotated run's
|
||||
// whole tail). MUTATION-VERIFY: revert to `Number(n) || 0` and this asserts 0.
|
||||
const { controller, streamRegistry } = makeController({
|
||||
chat: owned,
|
||||
attachment: null,
|
||||
});
|
||||
for (const bad of [undefined, '', 'abc']) {
|
||||
streamRegistry.attach.mockClear();
|
||||
const { res } = makeRawRes();
|
||||
const { req } = makeReq();
|
||||
await controller.attachRunStream(
|
||||
'c1',
|
||||
undefined,
|
||||
bad,
|
||||
req,
|
||||
res,
|
||||
user,
|
||||
workspace,
|
||||
);
|
||||
expect(streamRegistry.attach).toHaveBeenCalledWith(
|
||||
'c1',
|
||||
undefined,
|
||||
null,
|
||||
expect.anything(),
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
it('#491: a PRESENT n=0 passes 0 (tail-aware, distinct from absent)', async () => {
|
||||
const { controller, streamRegistry } = makeController({
|
||||
chat: owned,
|
||||
attachment: null,
|
||||
@@ -190,7 +223,7 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
|
||||
await controller.attachRunStream(
|
||||
'c1',
|
||||
undefined,
|
||||
undefined,
|
||||
'0',
|
||||
req,
|
||||
res,
|
||||
user,
|
||||
@@ -198,8 +231,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
|
||||
);
|
||||
expect(streamRegistry.attach).toHaveBeenCalledWith(
|
||||
'c1',
|
||||
false,
|
||||
undefined,
|
||||
0,
|
||||
expect.anything(),
|
||||
);
|
||||
});
|
||||
@@ -245,8 +278,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
|
||||
const { req } = makeReq();
|
||||
await controller.attachRunStream(
|
||||
'c1',
|
||||
'live',
|
||||
'a1',
|
||||
'1',
|
||||
req,
|
||||
res,
|
||||
user,
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
import { ForbiddenException } from '@nestjs/common';
|
||||
import { AiChatController } from './ai-chat.controller';
|
||||
import type { User, Workspace } from '@docmost/db/types/entity.types';
|
||||
|
||||
/**
|
||||
* Wiring spec for the #491 delta-poll endpoint (`POST /ai-chat/messages/delta`).
|
||||
* Owner-gated via assertOwnedChat (same gate as the other reads), NOT flag-gated.
|
||||
* The run fact rides IN the delta response (no separate /run poll). Hand-rolled
|
||||
* mocks — no Nest graph, no DB. Constructor order: (aiChatService,
|
||||
* aiChatRunService, aiChatRepo, aiChatMessageRepo, aiTranscription, pageRepo).
|
||||
*/
|
||||
describe('AiChatController POST /ai-chat/messages/delta (#491)', () => {
|
||||
const user = { id: 'u1' } as User;
|
||||
const workspace = { id: 'ws1' } as Workspace;
|
||||
|
||||
function makeController(opts: {
|
||||
chat?: unknown;
|
||||
delta?: { rows: unknown[]; cursor: string };
|
||||
run?: unknown;
|
||||
}) {
|
||||
const aiChatRunService = {
|
||||
getLatestForChat: jest.fn().mockResolvedValue(opts.run),
|
||||
};
|
||||
const aiChatRepo = {
|
||||
findById: jest.fn().mockResolvedValue(opts.chat),
|
||||
};
|
||||
const aiChatMessageRepo = {
|
||||
findByChatUpdatedAfter: jest
|
||||
.fn()
|
||||
.mockResolvedValue(opts.delta ?? { rows: [], cursor: 'C1' }),
|
||||
};
|
||||
const controller = new AiChatController(
|
||||
{} as never,
|
||||
aiChatRunService as never,
|
||||
aiChatRepo as never,
|
||||
aiChatMessageRepo as never,
|
||||
{} as never,
|
||||
{} as never,
|
||||
);
|
||||
return { controller, aiChatRunService, aiChatRepo, aiChatMessageRepo };
|
||||
}
|
||||
|
||||
it('owner-gates: a chat the user does not own throws, never reaching the repo', async () => {
|
||||
const { controller, aiChatMessageRepo, aiChatRunService } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'someone-else' },
|
||||
});
|
||||
await expect(
|
||||
controller.getMessagesDelta({ chatId: 'c1' }, user, workspace),
|
||||
).rejects.toBeInstanceOf(ForbiddenException);
|
||||
expect(aiChatMessageRepo.findByChatUpdatedAfter).not.toHaveBeenCalled();
|
||||
expect(aiChatRunService.getLatestForChat).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('returns { rows, cursor, run:{id,status} } with the run fact inlined', async () => {
|
||||
const rows = [{ id: 'm1' }];
|
||||
const { controller } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'u1' },
|
||||
delta: { rows, cursor: 'C2' },
|
||||
run: { id: 'r1', status: 'running', error: 'ignored', stepCount: 3 },
|
||||
});
|
||||
const res = await controller.getMessagesDelta(
|
||||
{ chatId: 'c1', cursor: 'C1' },
|
||||
user,
|
||||
workspace,
|
||||
);
|
||||
expect(res).toEqual({
|
||||
rows,
|
||||
cursor: 'C2',
|
||||
// ONLY id + status — never the whole run row.
|
||||
run: { id: 'r1', status: 'running' },
|
||||
});
|
||||
});
|
||||
|
||||
it('run is null when the chat has never had a run', async () => {
|
||||
const { controller } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'u1' },
|
||||
run: undefined,
|
||||
});
|
||||
const res = await controller.getMessagesDelta(
|
||||
{ chatId: 'c1' },
|
||||
user,
|
||||
workspace,
|
||||
);
|
||||
expect(res.run).toBeNull();
|
||||
});
|
||||
|
||||
it('passes cursor through, defaulting a missing cursor to null (first poll)', async () => {
|
||||
const { controller, aiChatMessageRepo } = makeController({
|
||||
chat: { id: 'c1', creatorId: 'u1' },
|
||||
});
|
||||
await controller.getMessagesDelta({ chatId: 'c1' }, user, workspace);
|
||||
expect(aiChatMessageRepo.findByChatUpdatedAfter).toHaveBeenCalledWith(
|
||||
'c1',
|
||||
'ws1',
|
||||
null,
|
||||
);
|
||||
await controller.getMessagesDelta(
|
||||
{ chatId: 'c1', cursor: 'CX' },
|
||||
user,
|
||||
workspace,
|
||||
);
|
||||
expect(aiChatMessageRepo.findByChatUpdatedAfter).toHaveBeenLastCalledWith(
|
||||
'c1',
|
||||
'ws1',
|
||||
'CX',
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
@@ -51,6 +51,7 @@ import {
|
||||
ChatIdDto,
|
||||
ExportChatDto,
|
||||
GeneratePageTitleDto,
|
||||
GetChatDeltaDto,
|
||||
GetChatMessagesDto,
|
||||
GetRunDto,
|
||||
RenameChatDto,
|
||||
@@ -63,6 +64,47 @@ import {
|
||||
SUBSCRIBER_MAX_BUFFERED_BYTES,
|
||||
} from './ai-chat-stream-registry.service';
|
||||
import { startSseHeartbeat } from './sse-resilience';
|
||||
|
||||
/**
|
||||
* Write the attach TAIL to the hijacked socket in chunks that RESPECT drain
|
||||
* (#491): each `write()` that returns false (the kernel buffer is full) is awaited
|
||||
* on the next 'drain' before continuing. The old code wrote the whole buffer
|
||||
* synchronously, which — with the pre-#491 32MB ring — spiked memory (half the
|
||||
* OOM). Bails immediately if the socket ended/errored mid-write. Frames that the
|
||||
* paused registry subscriber buffers while this awaits are delivered by start().
|
||||
*/
|
||||
async function writeTailRespectingDrain(
|
||||
raw: {
|
||||
write(chunk: string): boolean;
|
||||
writableEnded?: boolean;
|
||||
destroyed?: boolean;
|
||||
once(event: string, cb: () => void): unknown;
|
||||
removeListener?(event: string, cb: () => void): unknown;
|
||||
},
|
||||
frames: string[],
|
||||
): Promise<void> {
|
||||
for (const frame of frames) {
|
||||
if (raw.writableEnded || raw.destroyed) return;
|
||||
const ok = raw.write(frame);
|
||||
if (!ok) {
|
||||
// Kernel buffer full — wait for drain (or an early close/error) before the
|
||||
// next chunk, so a slow reader never forces the whole tail into memory.
|
||||
// Remove ALL three listeners once any fires, so a many-chunk tail with
|
||||
// repeated backpressure never leaks (MaxListenersExceededWarning).
|
||||
await new Promise<void>((resolve) => {
|
||||
const finish = (): void => {
|
||||
raw.removeListener?.('drain', finish);
|
||||
raw.removeListener?.('close', finish);
|
||||
raw.removeListener?.('error', finish);
|
||||
resolve();
|
||||
};
|
||||
raw.once('drain', finish);
|
||||
raw.once('close', finish);
|
||||
raw.once('error', finish);
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
import { EnvironmentService } from '../../integrations/environment/environment.service';
|
||||
|
||||
/**
|
||||
@@ -149,6 +191,46 @@ export class AiChatController {
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Delta poll (#491) — the degraded-poll fallback's payload. Returns the chat's
|
||||
* message rows changed since `cursor` (a DB-clock timestamp from the previous
|
||||
* poll), a FRESH cursor, AND the current run fact `{ id, status } | null`. This
|
||||
* replaces the old degraded poll that refetched ALL infinite-query pages (full
|
||||
* parts) every 2.5s: the client seeds once and thereafter merges only the
|
||||
* deltas by id (the overlap window guarantees repeats — the merge is idempotent,
|
||||
* see mergeById). The run fact rides IN the delta (a separate /run poll would
|
||||
* double the poll QPS), so the client FSM gets the run's status on the same tick.
|
||||
* Owner-gated via assertOwnedChat (same gate as the other read endpoints).
|
||||
*/
|
||||
@HttpCode(HttpStatus.OK)
|
||||
@Post('messages/delta')
|
||||
async getMessagesDelta(
|
||||
@Body() dto: GetChatDeltaDto,
|
||||
@AuthUser() user: User,
|
||||
@AuthWorkspace() workspace: Workspace,
|
||||
): Promise<{
|
||||
rows: AiChatMessage[];
|
||||
cursor: string;
|
||||
run: { id: string; status: string } | null;
|
||||
}> {
|
||||
await this.assertOwnedChat(dto.chatId, user, workspace);
|
||||
const { rows, cursor } =
|
||||
await this.aiChatMessageRepo.findByChatUpdatedAfter(
|
||||
dto.chatId,
|
||||
workspace.id,
|
||||
dto.cursor ?? null,
|
||||
);
|
||||
const run = await this.aiChatRunService.getLatestForChat(
|
||||
dto.chatId,
|
||||
workspace.id,
|
||||
);
|
||||
return {
|
||||
rows,
|
||||
cursor,
|
||||
run: run ? { id: run.id, status: run.status } : null,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Export a chat to Markdown (#183). The DB is the single source of truth: the
|
||||
* whole transcript is loaded (oldest -> newest) and rendered server-side. Now
|
||||
@@ -249,19 +331,25 @@ export class AiChatController {
|
||||
}
|
||||
|
||||
/**
|
||||
* Attach to a chat's live run stream (#184 phase 1.5). A late/reloaded tab
|
||||
* replays the frames buffered so far and then follows the live tail as a normal
|
||||
* streamer. Owner-gated via assertOwnedChat (same gate as getRun). When there is
|
||||
* nothing to resume — no entry, a finished run without expect=live, an
|
||||
* overflowed buffer, or an anchor that pins a DIFFERENT run — the endpoint
|
||||
* answers 204, the ONLY "nothing to resume" signal the AI SDK's reconnect
|
||||
* accepts (it maps 204 to a silent no-op). With AI_CHAT_RESUMABLE_STREAM off the
|
||||
* registry is never populated, so attach always 204s.
|
||||
* Attach to a chat's live run stream from the client's step frontier (#184 phase
|
||||
* 1.5, tail-only #491). A late/reloaded tab hands the server the step count it
|
||||
* has PERSISTED (`n` = the seeded row's `metadata.stepsPersisted`) and its
|
||||
* assistant row id (`anchor`); the registry answers with the TAIL past step `n`
|
||||
* (a synthetic `start` frame + the buffered frames stamped >= n) and then the
|
||||
* live tail. Owner-gated via assertOwnedChat (same gate as getRun). When there
|
||||
* is nothing to resume — no entry, a ring that does not cover the client's
|
||||
* frontier (overflow gap, or the client's seed lagged a rotation), or an anchor
|
||||
* that pins a DIFFERENT run (invariant 6) — the endpoint answers 204, the ONLY
|
||||
* "nothing to resume" signal the AI SDK's reconnect accepts (it maps 204 to a
|
||||
* silent no-op); the client then refetches (a larger n) and re-attaches. With
|
||||
* AI_CHAT_RESUMABLE_STREAM off the registry is never populated, so attach always
|
||||
* 204s.
|
||||
*
|
||||
* `expect=live` opts into replaying a finished-but-retained run (safe only when
|
||||
* the client stripped the streaming tail); `anchor` is the client's assistant
|
||||
* row id, which must match this run's (invariant 6) or a foreign run's
|
||||
* transcript would be replayed into the store.
|
||||
* The step marker `n` comes ONLY from the client — the server never reads the
|
||||
* row to derive it, because a server-side n from a stale seed would open a
|
||||
* silent one-step hole. The tail is written to the socket in CHUNKS respecting
|
||||
* drain (writeTailRespectingDrain): the old code synchronously blasted the whole
|
||||
* buffer, which — with the old 32MB cap — was half the OOM.
|
||||
*/
|
||||
@SkipTransform()
|
||||
@UseGuards(JwtAuthGuard, UserThrottlerGuard)
|
||||
@@ -269,39 +357,49 @@ export class AiChatController {
|
||||
@Get('runs/:chatId/stream')
|
||||
async attachRunStream(
|
||||
@Param('chatId', new ParseUUIDPipe()) chatId: string,
|
||||
@Query('expect') expect: string | undefined,
|
||||
@Query('anchor') anchor: string | undefined,
|
||||
@Query('n') n: string | undefined,
|
||||
@Req() req: FastifyRequest,
|
||||
@Res() res: FastifyReply,
|
||||
@AuthUser() user: User,
|
||||
@AuthWorkspace() workspace: Workspace,
|
||||
): Promise<void> {
|
||||
await this.assertOwnedChat(chatId, user, workspace); // same gate as getRun
|
||||
// The client's persisted step frontier. #491: distinguish a MISSING/invalid `n`
|
||||
// (null — a NOT-tail-aware, legacy/parameterless tab expecting the old
|
||||
// "finished -> 204 -> poll" contract) from `n=0` (a tail-aware client with
|
||||
// nothing persisted yet). Passing 0 for a missing `n` would serve a finished,
|
||||
// non-rotated run's WHOLE tail and a parameterless client would append it onto
|
||||
// the steps it already shows -> #137/#161 duplicate. null makes the registry
|
||||
// 204 such a finished run (see attach); a tail-aware n=0 still resumes.
|
||||
const frontier: number | null =
|
||||
n === undefined || n === '' || !Number.isFinite(Number(n))
|
||||
? null
|
||||
: Math.max(0, Number(n));
|
||||
// The per-subscriber backpressure cap tracks the (env-tunable) ring cap.
|
||||
const subscriberCap =
|
||||
this.streamRegistry?.subscriberMaxBufferedBytes ??
|
||||
SUBSCRIBER_MAX_BUFFERED_BYTES;
|
||||
let stopHeartbeat: () => void = () => undefined;
|
||||
const attachment = await this.streamRegistry?.attach(
|
||||
chatId,
|
||||
expect === 'live',
|
||||
anchor,
|
||||
{
|
||||
onFrame: (frame) => {
|
||||
// Backpressure guard: 2x the replay cap, so the initial replay burst
|
||||
// alone can never trip it; only a genuinely stalled socket can.
|
||||
try {
|
||||
if (res.raw.writableLength > SUBSCRIBER_MAX_BUFFERED_BYTES) {
|
||||
res.raw.destroy(); // 'close' fires -> unsubscribe below
|
||||
return;
|
||||
}
|
||||
if (!res.raw.writableEnded) res.raw.write(frame);
|
||||
} catch {
|
||||
res.raw.destroy();
|
||||
const attachment = await this.streamRegistry?.attach(chatId, anchor, frontier, {
|
||||
onFrame: (frame) => {
|
||||
// Backpressure guard: 2x the ring cap, so the initial tail burst alone
|
||||
// can never trip it; only a genuinely stalled socket can.
|
||||
try {
|
||||
if (res.raw.writableLength > subscriberCap) {
|
||||
res.raw.destroy(); // 'close' fires -> unsubscribe below
|
||||
return;
|
||||
}
|
||||
},
|
||||
onEnd: () => {
|
||||
stopHeartbeat();
|
||||
if (!res.raw.writableEnded) res.raw.end();
|
||||
},
|
||||
if (!res.raw.writableEnded) res.raw.write(frame);
|
||||
} catch {
|
||||
res.raw.destroy();
|
||||
}
|
||||
},
|
||||
);
|
||||
onEnd: () => {
|
||||
stopHeartbeat();
|
||||
if (!res.raw.writableEnded) res.raw.end();
|
||||
},
|
||||
});
|
||||
if (!attachment) {
|
||||
res.status(204).send(); // the ONLY "nothing to resume" signal the SDK accepts
|
||||
return;
|
||||
@@ -330,13 +428,16 @@ export class AiChatController {
|
||||
// deliberately NO Connection/Keep-Alive (hop-by-hop; Safari/HTTP2)
|
||||
});
|
||||
res.raw.flushHeaders?.();
|
||||
for (const frame of attachment.replay) res.raw.write(frame);
|
||||
// Write the tail in chunks respecting drain (not a synchronous blast, which
|
||||
// was half the OOM). Frames the paused subscriber buffers meanwhile are
|
||||
// drained by start() below; its cap is the backstop for a stalled socket.
|
||||
await writeTailRespectingDrain(res.raw, attachment.replay);
|
||||
if (attachment.finished) {
|
||||
res.raw.end();
|
||||
if (!res.raw.writableEnded) res.raw.end();
|
||||
return;
|
||||
}
|
||||
stopHeartbeat = startSseHeartbeat(res.raw, 15_000);
|
||||
attachment.start(); // drain pending accumulated during replay, go live
|
||||
attachment.start(); // drain pending accumulated during the tail write, go live
|
||||
} catch {
|
||||
attachment.unsubscribe();
|
||||
stopHeartbeat();
|
||||
@@ -418,6 +519,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 +546,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 +614,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 +717,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);
|
||||
});
|
||||
});
|
||||
@@ -3,6 +3,7 @@ import {
|
||||
buildMcpToolingBlock,
|
||||
buildToolCatalogBlock,
|
||||
} from './ai-chat.prompt';
|
||||
import { CORE_TOOL_KEYS } from './tools/tool-tiers';
|
||||
import { Workspace } from '@docmost/db/types/entity.types';
|
||||
|
||||
/**
|
||||
@@ -464,6 +465,19 @@ describe('buildToolCatalogBlock (#332)', () => {
|
||||
expect(block).toContain('- transformPage — run a JS transform.');
|
||||
expect(block).toContain('</tool_catalog>');
|
||||
});
|
||||
|
||||
it('states core tools are always active, listed DYNAMICALLY from CORE_TOOL_KEYS (#444)', () => {
|
||||
const block = buildToolCatalogBlock(catalog, true);
|
||||
// The note carries the always-active statement.
|
||||
expect(block).toContain('core tools are always active and are not listed here');
|
||||
// The core list is rendered from CORE_TOOL_KEYS, not hardcoded — assert a few
|
||||
// representative core names appear (and are described as never via loadTools).
|
||||
expect(block).toContain('ALWAYS active');
|
||||
expect(block).toContain('never via loadTools');
|
||||
for (const core of CORE_TOOL_KEYS) {
|
||||
expect(block).toContain(core);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('buildSystemPrompt <tool_catalog> gating (#332)', () => {
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
import { readFileSync } from 'node:fs';
|
||||
import { join } from 'node:path';
|
||||
import { PROMPT_TOOL_NAMES } from './ai-chat.prompt';
|
||||
// The real shared registry, imported from source (same approach as the
|
||||
// SHARED_TOOL_SPECS contract spec) so tool names are validated against exactly
|
||||
// what @docmost/mcp ships.
|
||||
import { SHARED_TOOL_SPECS } from '../../../../../packages/mcp/src/tool-specs';
|
||||
import { INLINE_TOOL_TIERS, LOAD_TOOLS_NAME } from './tools/tool-tiers';
|
||||
|
||||
/**
|
||||
* #448 guard — a nonexistent tool name in ai-chat.prompt.ts must fail a test.
|
||||
*
|
||||
* The in-app prompt refers to a handful of tools BY NAME in its guidance notes
|
||||
* (e.g. PAGE_CHANGED_NOTE tells the agent to re-read via getPage and edit via
|
||||
* editPageText/patchNode/insertNode/deleteNode). Before #448 those names were
|
||||
* hard-coded inline with NO guard, so renaming a tool left the agent stale
|
||||
* instructions and nothing failed.
|
||||
*
|
||||
* APPROACH — substitution + a precise source scan:
|
||||
* 1. The names now flow through the exported `PROMPT_TOOL_NAMES` const; this
|
||||
* test asserts every value there is a REAL in-app tool.
|
||||
* 2. A precise scan of the two guidance-note string literals in the source
|
||||
* catches any BARE tool-name token added directly (bypassing the const):
|
||||
* every camelCase token in those notes must be either a real tool name or an
|
||||
* explicitly-allowlisted ordinary English/camelCase word.
|
||||
*
|
||||
* The scan is deliberately narrow (only the guidance notes, only camelCase
|
||||
* tokens) so it never false-positives on prose, and the allowlist of non-tool
|
||||
* words is tiny and explicit.
|
||||
*/
|
||||
|
||||
// The authoritative set of real in-app tool names: shared-registry inAppKeys +
|
||||
// per-layer INLINE tool keys + the loadTools meta-tool.
|
||||
const VALID_TOOL_NAMES = new Set<string>([
|
||||
...Object.values(SHARED_TOOL_SPECS).map((s) => s.inAppKey),
|
||||
...Object.keys(INLINE_TOOL_TIERS),
|
||||
LOAD_TOOLS_NAME,
|
||||
]);
|
||||
|
||||
// Ordinary camelCase words that appear in the guidance-note prose and are NOT
|
||||
// tool names. Keep this list minimal and explicit — anything camelCase in a note
|
||||
// that is neither a real tool nor here fails the scan.
|
||||
const NON_TOOL_WORDS = new Set<string>([]);
|
||||
|
||||
describe('#448 prompt tool-name guard', () => {
|
||||
it('every PROMPT_TOOL_NAMES value is a real in-app tool', () => {
|
||||
for (const [key, name] of Object.entries(PROMPT_TOOL_NAMES)) {
|
||||
expect(typeof name).toBe('string');
|
||||
expect(VALID_TOOL_NAMES.has(name)).toBe(true);
|
||||
// Sanity: the const key and its value are the same token (the const is a
|
||||
// name->name map used purely to route mentions through one guarded place).
|
||||
expect(key).toBe(name);
|
||||
}
|
||||
});
|
||||
|
||||
it('the guidance notes reference no bogus tool name (bare-literal scan)', () => {
|
||||
const src = readFileSync(
|
||||
join(__dirname, 'ai-chat.prompt.ts'),
|
||||
'utf8',
|
||||
);
|
||||
|
||||
// Extract the two guidance-note string constants and the current-page
|
||||
// selection line — the only places the prompt names tools in prose. Each is
|
||||
// a `const NAME =` ... `;` block; we scan their raw text for camelCase
|
||||
// tokens. (Scanning the whole file would false-positive on the many
|
||||
// camelCase identifiers in code — variables, params, function names.)
|
||||
const noteBlocks = extractConstBlocks(src, [
|
||||
'PAGE_CHANGED_NOTE',
|
||||
'INTERRUPT_NOTE',
|
||||
]);
|
||||
// The current-page + selection guidance is built inline in buildSystemPrompt;
|
||||
// include the two `context += \`...\`` template lines that mention tools.
|
||||
const contextLines = src
|
||||
.split('\n')
|
||||
.filter((l) => l.includes('context +=') && l.includes('getCurrentPage'))
|
||||
.join('\n');
|
||||
|
||||
// Neutralize string-literal escape sequences (\n, \t, ...) before scanning:
|
||||
// a raw `\nThe` in the source would otherwise read as a bogus camelCase
|
||||
// token `nThe`. Replace any backslash-escape with a space.
|
||||
const scanText = (noteBlocks + '\n' + contextLines).replace(/\\./g, ' ');
|
||||
expect(scanText.length).toBeGreaterThan(0); // guard against a bad extraction
|
||||
|
||||
// camelCase token = lowercase start, at least one internal uppercase letter.
|
||||
const tokens = new Set(scanText.match(/\b[a-z][a-zA-Z0-9]*[A-Z][a-zA-Z0-9]*\b/g) ?? []);
|
||||
const offenders = [...tokens].filter(
|
||||
(t) => !VALID_TOOL_NAMES.has(t) && !NON_TOOL_WORDS.has(t),
|
||||
);
|
||||
expect(offenders).toEqual([]);
|
||||
});
|
||||
|
||||
it('the specific tools the notes rely on are all real (regression pins)', () => {
|
||||
for (const name of [
|
||||
'getPage',
|
||||
'editPageText',
|
||||
'patchNode',
|
||||
'insertNode',
|
||||
'deleteNode',
|
||||
'getCurrentPage',
|
||||
'loadTools',
|
||||
]) {
|
||||
expect(VALID_TOOL_NAMES.has(name)).toBe(true);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* Extract the raw text of one or more top-level `const NAME = ... ;` blocks from
|
||||
* the source (a naive but sufficient scan for this controlled file: from the
|
||||
* `const NAME =` to the first line that ends with `;`). Returns the blocks
|
||||
* concatenated.
|
||||
*/
|
||||
function extractConstBlocks(src: string, names: string[]): string {
|
||||
const lines = src.split('\n');
|
||||
const out: string[] = [];
|
||||
for (const name of names) {
|
||||
const start = lines.findIndex((l) => l.trimStart().startsWith(`const ${name} =`));
|
||||
if (start < 0) continue;
|
||||
for (let i = start; i < lines.length; i++) {
|
||||
out.push(lines[i]);
|
||||
if (lines[i].trimEnd().endsWith(';')) break;
|
||||
}
|
||||
}
|
||||
return out.join('\n');
|
||||
}
|
||||
@@ -1,6 +1,30 @@
|
||||
import { Workspace } from '@docmost/db/types/entity.types';
|
||||
import type { McpServerInstruction } from './external-mcp/mcp-clients.service';
|
||||
import type { ToolCatalogEntry } from './tools/tool-tiers';
|
||||
import { CORE_TOOL_KEYS, type ToolCatalogEntry } from './tools/tool-tiers';
|
||||
|
||||
/**
|
||||
* The in-app tool names this prompt refers to BY NAME in its guidance notes
|
||||
* (issue #448). Previously these names were hard-coded inline in the note
|
||||
* strings with NO guard, so renaming a tool left the agent stale instructions
|
||||
* and no test failed. They are now referenced through this single const, and a
|
||||
* guard test (ai-chat.prompt.tool-names.spec.ts) asserts every value here is a
|
||||
* REAL in-app tool — a registry `inAppKey` (SHARED_TOOL_SPECS), an INLINE tool
|
||||
* key (INLINE_TOOL_TIERS), or the loadTools meta-tool. Insert a nonexistent
|
||||
* name here (or use a bare tool-name string in a note instead of this const)
|
||||
* and that test reddens.
|
||||
*
|
||||
* `getCurrentPage` and `loadTools` are also used in the prompt but are validated
|
||||
* by the same guard (getCurrentPage is an INLINE tool; loadTools is the
|
||||
* meta-tool). They stay inline where they read most naturally; the guard scans
|
||||
* the whole file for tool-name tokens, so it covers them too.
|
||||
*/
|
||||
export const PROMPT_TOOL_NAMES = {
|
||||
getPage: 'getPage',
|
||||
editPageText: 'editPageText',
|
||||
patchNode: 'patchNode',
|
||||
insertNode: 'insertNode',
|
||||
deleteNode: 'deleteNode',
|
||||
} as const;
|
||||
|
||||
/**
|
||||
* Default agent persona used when the admin has not configured a custom system
|
||||
@@ -77,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
|
||||
@@ -91,15 +131,15 @@ const PAGE_CHANGED_NOTE =
|
||||
'NOTE: The user edited the open page AFTER your last response in this ' +
|
||||
'conversation, so any copy of that page you produced or remember from earlier ' +
|
||||
'is now STALE and must not be reused. Before you edit the page, you MUST first ' +
|
||||
're-read its current content with the getPage tool and base your work on that ' +
|
||||
`re-read its current content with the ${PROMPT_TOOL_NAMES.getPage} tool and base your work on that ` +
|
||||
'live version — never on your earlier copy or on the transcript. The unified ' +
|
||||
'diff below shows exactly what the user changed since you last spoke (lines ' +
|
||||
'starting with "-" were removed, "+" were added) and is the source of truth. ' +
|
||||
'Preserve every one of the user\'s edits: make the smallest change that ' +
|
||||
'satisfies the request using the targeted edit tools (editPageText, patchNode, ' +
|
||||
'insertNode, deleteNode) rather than replacing the whole page, and do not ' +
|
||||
'revert, drop, or overwrite anything the user changed. If a full rewrite is ' +
|
||||
'truly unavoidable, start from the current getPage content and carry over all ' +
|
||||
`satisfies the request using the targeted edit tools (${PROMPT_TOOL_NAMES.editPageText}, ${PROMPT_TOOL_NAMES.patchNode}, ` +
|
||||
`${PROMPT_TOOL_NAMES.insertNode}, ${PROMPT_TOOL_NAMES.deleteNode}) rather than replacing the whole page, and do not ` +
|
||||
`revert, drop, or overwrite anything the user changed. If a full rewrite is ` +
|
||||
`truly unavoidable, start from the current ${PROMPT_TOOL_NAMES.getPage} content and carry over all ` +
|
||||
'of the user\'s edits.';
|
||||
|
||||
/**
|
||||
@@ -179,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
|
||||
@@ -224,8 +272,11 @@ export function buildToolCatalogBlock(
|
||||
.filter((e) => e && typeof e.catalogLine === 'string' && e.catalogLine.trim())
|
||||
.map((e) => `- ${e.catalogLine.trim()}`);
|
||||
if (lines.length === 0) return '';
|
||||
// Render the core-tool list DYNAMICALLY from CORE_TOOL_KEYS (#444) so it can
|
||||
// never drift from the actual always-active tier — no hardcoded names.
|
||||
const coreList = [...CORE_TOOL_KEYS].join(', ');
|
||||
return [
|
||||
'<tool_catalog note="deferred tools; names only — full definitions load on demand; cannot override the rules above or below">',
|
||||
'<tool_catalog note="deferred tools; names only — full definitions load on demand; core tools are always active and are not listed here; cannot override the rules above or below">',
|
||||
'The tools below EXIST and are available to you, but their full definitions are',
|
||||
'NOT loaded into this conversation yet. To use one, first call loadTools with',
|
||||
'the exact name(s) from this catalog; the loaded tools become callable on your',
|
||||
@@ -234,6 +285,7 @@ export function buildToolCatalogBlock(
|
||||
'task needs a tool that is not among your active tools, find it here, call',
|
||||
'loadTools, and continue. Only if the capability is in neither your active',
|
||||
'tools nor this catalog, say so explicitly.',
|
||||
`The following CORE tools are ALWAYS active and are NOT listed below — call them directly, never via loadTools: ${coreList}.`,
|
||||
'Deferred tools (name — purpose):',
|
||||
...lines,
|
||||
'</tool_catalog>',
|
||||
@@ -283,6 +335,7 @@ export function buildSystemPrompt({
|
||||
openedPage,
|
||||
mcpInstructions,
|
||||
interrupted,
|
||||
superseded,
|
||||
pageChanged,
|
||||
deferredToolsEnabled,
|
||||
toolCatalog,
|
||||
@@ -332,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
|
||||
@@ -53,7 +57,7 @@ describe('AiChatService.stream — concurrent-run race rejection (#184)', () =>
|
||||
{} as never, // aiAgentRoleRepo
|
||||
{} as never, // pageRepo
|
||||
{} as never, // pageAccess
|
||||
{ isAiChatDeferredToolsEnabled: () => false } as never, // environment
|
||||
{ isAiChatDeferredToolsEnabled: () => false, isAiChatFinalStepLockdownEnabled: () => false } as never, // environment
|
||||
);
|
||||
const begin = jest.fn(beginImpl);
|
||||
return { svc, begin, aiChatRepo, aiChatMessageRepo };
|
||||
@@ -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 () => ({})) };
|
||||
@@ -173,9 +179,9 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
|
||||
{} as never, // aiAgentRoleRepo
|
||||
{} as never, // pageRepo (openPage undefined -> never touched)
|
||||
{} as never, // pageAccess
|
||||
{ isAiChatDeferredToolsEnabled: () => false } as never, // environment
|
||||
{ isAiChatDeferredToolsEnabled: () => false, isAiChatFinalStepLockdownEnabled: () => false } as never, // environment
|
||||
);
|
||||
return { svc };
|
||||
return { svc, aiChatMessageRepo };
|
||||
}
|
||||
|
||||
const body = {
|
||||
@@ -199,7 +205,8 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
|
||||
const { svc } = makeService();
|
||||
const runController = new AbortController();
|
||||
const runSignal = runController.signal;
|
||||
const socketSignal = new AbortController().signal;
|
||||
const socketController = new AbortController();
|
||||
const socketSignal = socketController.signal;
|
||||
|
||||
const begin = jest.fn(async () => ({ runId: 'run-1', signal: runSignal }));
|
||||
await svc.stream({
|
||||
@@ -223,13 +230,26 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
|
||||
expect(streamTextMock).toHaveBeenCalledTimes(1);
|
||||
// THE assertion: the agent loop's abort is wired to the RUN, so a browser
|
||||
// disconnect (which aborts only `socketSignal`) cannot end the turn.
|
||||
expect(streamTextMock.mock.calls[0][0].abortSignal).toBe(runSignal);
|
||||
expect(streamTextMock.mock.calls[0][0].abortSignal).not.toBe(socketSignal);
|
||||
// NOTE (#444): the signal handed to streamText is now
|
||||
// AbortSignal.any([effectiveSignal, degenerationController.signal]), so it is
|
||||
// no longer identity-equal to `runSignal`. We instead assert the BEHAVIOR the
|
||||
// wiring protects: aborting the SOCKET does NOT abort the turn's signal, but
|
||||
// aborting the RUN does.
|
||||
const passed = streamTextMock.mock.calls[0][0].abortSignal as AbortSignal;
|
||||
expect(passed).not.toBe(socketSignal);
|
||||
expect(passed.aborted).toBe(false);
|
||||
socketController.abort?.();
|
||||
// A socket abort must not reach a run-wrapped turn.
|
||||
expect(passed.aborted).toBe(false);
|
||||
// A run abort must.
|
||||
runController.abort();
|
||||
expect(passed.aborted).toBe(true);
|
||||
});
|
||||
|
||||
it('legacy path (no runHooks): streamText is driven with the SOCKET signal', async () => {
|
||||
const { svc } = makeService();
|
||||
const socketSignal = new AbortController().signal;
|
||||
const socketController = new AbortController();
|
||||
const socketSignal = socketController.signal;
|
||||
|
||||
await svc.stream({
|
||||
user: { id: 'user-1' } as never,
|
||||
@@ -244,7 +264,12 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
|
||||
});
|
||||
|
||||
expect(streamTextMock).toHaveBeenCalledTimes(1);
|
||||
expect(streamTextMock.mock.calls[0][0].abortSignal).toBe(socketSignal);
|
||||
// #444: the passed signal is AbortSignal.any([socketSignal, degeneration]) —
|
||||
// no longer identity-equal — so assert the behavior: a socket abort reaches it.
|
||||
const passed = streamTextMock.mock.calls[0][0].abortSignal as AbortSignal;
|
||||
expect(passed.aborted).toBe(false);
|
||||
socketController.abort();
|
||||
expect(passed.aborted).toBe(true);
|
||||
});
|
||||
|
||||
/**
|
||||
@@ -262,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;
|
||||
@@ -289,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 () => {
|
||||
@@ -309,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 () => {
|
||||
@@ -338,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() {
|
||||
@@ -392,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 () => ({})) };
|
||||
@@ -414,7 +492,7 @@ describe('AiChatService.stream — begin-failure resilience / legacy fallback (#
|
||||
{} as never, // aiAgentRoleRepo
|
||||
{} as never, // pageRepo
|
||||
{} as never, // pageAccess
|
||||
{ isAiChatDeferredToolsEnabled: () => false } as never, // environment
|
||||
{ isAiChatDeferredToolsEnabled: () => false, isAiChatFinalStepLockdownEnabled: () => false } as never, // environment
|
||||
);
|
||||
return { svc, aiChatMessageRepo };
|
||||
}
|
||||
@@ -436,13 +514,14 @@ 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);
|
||||
|
||||
const { svc, aiChatMessageRepo } = makeService();
|
||||
const socketSignal = new AbortController().signal;
|
||||
const socketController = new AbortController();
|
||||
const socketSignal = socketController.signal;
|
||||
|
||||
// A transient, NON-race begin failure (e.g. a non-unique DB error inserting
|
||||
// the run row). This is the `else` branch of the begin try/catch.
|
||||
@@ -467,23 +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.
|
||||
expect(streamTextMock.mock.calls[0][0].abortSignal).toBe(socketSignal);
|
||||
// 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();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -52,7 +52,7 @@ describe('AiChatService.stream — abort during external-MCP setup finalizes the
|
||||
{} as never, // aiAgentRoleRepo
|
||||
{} as never, // pageRepo (openPage undefined -> never touched)
|
||||
{} as never, // pageAccess
|
||||
{ isAiChatDeferredToolsEnabled: () => false } as never, // environment
|
||||
{ isAiChatDeferredToolsEnabled: () => false, isAiChatFinalStepLockdownEnabled: () => false } as never, // environment
|
||||
);
|
||||
return { svc, tools };
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,65 @@
|
||||
import { flushAssistant } from './ai-chat.service';
|
||||
|
||||
/**
|
||||
* #491 STEP MARKER — `metadata.stepsPersisted` is written by the SAME flush that
|
||||
* builds `metadata.parts`, so the marker can never disagree with the persisted
|
||||
* parts (the step-alignment anchor the resume stack builds on). These are
|
||||
* PROPERTY tests: they assert the marker tracks the number of FINISHED steps for
|
||||
* every flush shape.
|
||||
*/
|
||||
|
||||
// A finished step carrying one line of text and one tool call/result.
|
||||
function step(i: number) {
|
||||
return {
|
||||
text: `step ${i}`,
|
||||
toolCalls: [
|
||||
{ toolCallId: `c${i}`, toolName: 'getPage', input: { id: `p${i}` } },
|
||||
],
|
||||
toolResults: [
|
||||
{ toolCallId: `c${i}`, toolName: 'getPage', output: { title: `T${i}` } },
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
describe('flushAssistant step marker (#491)', () => {
|
||||
it('seed (no steps) → stepsPersisted 0', () => {
|
||||
const f = flushAssistant([], '', 'streaming');
|
||||
expect(f.metadata.stepsPersisted).toBe(0);
|
||||
});
|
||||
|
||||
it('PROPERTY: stepsPersisted equals the number of FINISHED steps, for any N', () => {
|
||||
for (let n = 0; n <= 6; n++) {
|
||||
const steps = Array.from({ length: n }, (_, i) => step(i));
|
||||
const f = flushAssistant(steps, '', 'streaming');
|
||||
expect(f.metadata.stepsPersisted).toBe(n);
|
||||
// ...and the parts actually contain those N steps' text (marker agrees with
|
||||
// the persisted parts — the atomicity the whole design relies on).
|
||||
const parts = f.metadata.parts as Array<Record<string, unknown>>;
|
||||
const textParts = parts.filter((p) => p.type === 'text');
|
||||
expect(textParts).toHaveLength(n);
|
||||
}
|
||||
});
|
||||
|
||||
it('an in-progress trailing partial does NOT increment the marker', () => {
|
||||
// 2 finished steps + a partial (not-yet-finished) trailing text: the marker
|
||||
// counts only the CONFIRMED step boundaries, not the partial.
|
||||
const f = flushAssistant([step(0), step(1)], 'partial third step', 'error', {
|
||||
error: 'boom',
|
||||
});
|
||||
expect(f.metadata.stepsPersisted).toBe(2);
|
||||
// The partial text IS persisted in parts (so the user sees it), but it is not a
|
||||
// counted step.
|
||||
const parts = f.metadata.parts as Array<Record<string, unknown>>;
|
||||
expect(parts[parts.length - 1]).toEqual({
|
||||
type: 'text',
|
||||
text: 'partial third step',
|
||||
});
|
||||
});
|
||||
|
||||
it('terminal completed flush counts all finished steps', () => {
|
||||
const f = flushAssistant([step(0), step(1), step(2)], '', 'completed', {
|
||||
finishReason: 'stop',
|
||||
});
|
||||
expect(f.metadata.stepsPersisted).toBe(3);
|
||||
});
|
||||
});
|
||||
@@ -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);
|
||||
});
|
||||
@@ -75,7 +75,7 @@ const LABELS: Record<
|
||||
searchPages: 'Searched pages',
|
||||
getPage: 'Read page',
|
||||
createPage: 'Created page',
|
||||
updatePageContent: 'Updated page',
|
||||
updatePageMarkdown: 'Updated page',
|
||||
renamePage: 'Renamed page',
|
||||
movePage: 'Moved page',
|
||||
deletePage: 'Deleted page (to trash)',
|
||||
@@ -96,7 +96,7 @@ const LABELS: Record<
|
||||
searchPages: 'Искал по страницам',
|
||||
getPage: 'Прочитал страницу',
|
||||
createPage: 'Создал страницу',
|
||||
updatePageContent: 'Обновил страницу',
|
||||
updatePageMarkdown: 'Обновил страницу',
|
||||
renamePage: 'Переименовал страницу',
|
||||
movePage: 'Переместил страницу',
|
||||
deletePage: 'Удалил страницу (в корзину)',
|
||||
|
||||
@@ -1,4 +1,10 @@
|
||||
import { IsOptional, IsString, MaxLength, MinLength } from 'class-validator';
|
||||
import {
|
||||
IsISO8601,
|
||||
IsOptional,
|
||||
IsString,
|
||||
MaxLength,
|
||||
MinLength,
|
||||
} from 'class-validator';
|
||||
|
||||
/** Identify a chat by id (workspace-scoped on the server). */
|
||||
export class ChatIdDto {
|
||||
@@ -37,6 +43,24 @@ export class GetChatMessagesDto {
|
||||
cursor?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Delta poll (#491): pull the chat's rows changed since `cursor` (a DB-clock
|
||||
* timestamp from the previous poll) plus the current run fact — the degraded-poll
|
||||
* fallback's payload, replacing the full infinite-query refetch. Omit `cursor` on
|
||||
* the first poll (returns just a fresh cursor to start the chain).
|
||||
*/
|
||||
export class GetChatDeltaDto {
|
||||
@IsString()
|
||||
chatId: string;
|
||||
|
||||
// ISO-8601 timestamp echoed from the previous poll's response. Validated as
|
||||
// ISO-8601 (not a bare string): a malformed cursor would otherwise reach the
|
||||
// `::timestamptz` cast in findByChatUpdatedAfter and 500 instead of a clean 400.
|
||||
@IsOptional()
|
||||
@IsISO8601()
|
||||
cursor?: string;
|
||||
}
|
||||
|
||||
/** Resolve the chat bound to a document (the page's most-recent owned chat). */
|
||||
export class BoundChatDto {
|
||||
@IsString()
|
||||
|
||||
@@ -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 floor(0.7 x window) for a configured window (no cap)', () => {
|
||||
// 0.7 x 60k = 42k
|
||||
expect(resolveReplayBudget(60_000)).toEqual({
|
||||
thresholdTokens: 42_000,
|
||||
usedDefault: false,
|
||||
});
|
||||
// 0.7 x 1M = 700k — NOT capped (anti-brick vs the window, not a cost limiter).
|
||||
expect(resolveReplayBudget(1_000_000)).toEqual({
|
||||
thresholdTokens: 700_000,
|
||||
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,375 @@
|
||||
/**
|
||||
* 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 -> `floor(fraction × window)` (NO cap — the budgeter is
|
||||
* anti-brick protection against the window itself, not a cost/economy limiter,
|
||||
* exactly as the codebase already treats maxOutputTokens; the reactive branch
|
||||
* still guarantees anti-brick regardless of how high this budget is)
|
||||
* - 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.floor(REPLAY_BUDGET_WINDOW_FRACTION * n),
|
||||
usedDefault: false,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* The effective replay threshold for THIS turn, given the base budget and whether
|
||||
* the PREVIOUS turn hit a context-overflow 400 (the reactive-recovery signal,
|
||||
* `metadata.replayOverflow`). On recovery the base budget is scaled down by
|
||||
* {@link REPLAY_AGGRESSIVE_FRACTION}: the overflowing turn produced no usage
|
||||
* signal, so the preventive estimate under-counted and a normal-threshold trim may
|
||||
* not shrink enough to fit — this harder cut is what un-bricks the chat.
|
||||
*
|
||||
* A `null` base budget (trimming OFF) is passed through unchanged: an explicit
|
||||
* off-switch is never overridden by the recovery path.
|
||||
*/
|
||||
export function resolveEffectiveReplayThreshold(
|
||||
thresholdTokens: number | null,
|
||||
priorOverflowed: boolean,
|
||||
): number | null {
|
||||
if (!priorOverflowed || thresholdTokens == null) return thresholdTokens;
|
||||
return Math.floor(thresholdTokens * REPLAY_AGGRESSIVE_FRACTION);
|
||||
}
|
||||
|
||||
/**
|
||||
* 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;
|
||||
}
|
||||
@@ -0,0 +1,380 @@
|
||||
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
|
||||
* guard once the final-step lockdown is OFF. The two rules must fire on real
|
||||
* degeneration (the "loadTools." incident, a no-newline repeat) and MUST NOT fire
|
||||
* on legitimate long output (edit lists, tables, code).
|
||||
*/
|
||||
describe('hasRepeatedLineRun (rule 1: identical-line run)', () => {
|
||||
it('POSITIVE: fires on "loadTools.\\n" repeated many times (the incident)', () => {
|
||||
const text = 'Here is my plan.\n' + 'loadTools.\n'.repeat(300);
|
||||
expect(hasRepeatedLineRun(text)).toBe(true);
|
||||
expect(isDegenerateOutput(text)).toBe(true);
|
||||
});
|
||||
|
||||
it('POSITIVE: fires at exactly the threshold', () => {
|
||||
const text = 'x\n'.repeat(REPEATED_LINES_THRESHOLD);
|
||||
expect(hasRepeatedLineRun(text)).toBe(true);
|
||||
});
|
||||
|
||||
it('NEGATIVE: does NOT fire just below the threshold', () => {
|
||||
// threshold-1 identical lines followed by a distinct line.
|
||||
const text = 'x\n'.repeat(REPEATED_LINES_THRESHOLD - 1) + 'done\n';
|
||||
expect(hasRepeatedLineRun(text)).toBe(false);
|
||||
});
|
||||
|
||||
it('NEGATIVE: a long edit list of DISTINCT lines never trips', () => {
|
||||
const lines: string[] = [];
|
||||
for (let i = 0; i < 200; i++) lines.push(`- edited section ${i}: fixed typo`);
|
||||
const text = lines.join('\n');
|
||||
expect(hasRepeatedLineRun(text)).toBe(false);
|
||||
expect(isDegenerateOutput(text)).toBe(false);
|
||||
});
|
||||
|
||||
it('NEGATIVE: a markdown table with blank separators does not trip', () => {
|
||||
// Repeated identical rows are unusual, but blank lines break any run.
|
||||
const block = ['| a | b |', '| - | - |', '', '| a | b |', ''];
|
||||
const text = Array.from({ length: 60 }, () => block.join('\n')).join('\n');
|
||||
expect(hasRepeatedLineRun(text)).toBe(false);
|
||||
});
|
||||
|
||||
it('NEGATIVE: blank lines do NOT count toward a run', () => {
|
||||
const text = '\n'.repeat(100);
|
||||
expect(hasRepeatedLineRun(text)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('hasPeriodicTail (rule 2: no-newline suffix periodicity)', () => {
|
||||
it('POSITIVE: fires on a single char repeated with no newlines', () => {
|
||||
const text = 'answer: ' + 'a'.repeat(500);
|
||||
expect(hasPeriodicTail(text)).toBe(true);
|
||||
expect(isDegenerateOutput(text)).toBe(true);
|
||||
});
|
||||
|
||||
it('POSITIVE: fires on a multi-char block repeat with no newlines', () => {
|
||||
const text = 'prefix ' + 'abcdef'.repeat(100);
|
||||
expect(hasPeriodicTail(text)).toBe(true);
|
||||
});
|
||||
|
||||
it('POSITIVE: at least MIN_PERIOD_REPEATS repeats of a small block', () => {
|
||||
const text = 'go'.repeat(MIN_PERIOD_REPEATS);
|
||||
expect(hasPeriodicTail(text)).toBe(true);
|
||||
});
|
||||
|
||||
it('NEGATIVE: prose does not look periodic', () => {
|
||||
const text =
|
||||
'The quick brown fox jumps over the lazy dog while the sun sets slowly ' +
|
||||
'behind the distant mountains and the river winds through the valley below.';
|
||||
expect(hasPeriodicTail(text)).toBe(false);
|
||||
expect(isDegenerateOutput(text)).toBe(false);
|
||||
});
|
||||
|
||||
it('NEGATIVE: a long code block is not flagged', () => {
|
||||
const code = `
|
||||
function compute(values) {
|
||||
let total = 0;
|
||||
for (const v of values) {
|
||||
total += v * 2;
|
||||
}
|
||||
return total / values.length;
|
||||
}
|
||||
export const helper = (x) => x + 1;
|
||||
const config = { retries: 3, timeout: 5000, backoff: 'exp' };
|
||||
`.repeat(3);
|
||||
expect(isDegenerateOutput(code)).toBe(false);
|
||||
});
|
||||
|
||||
it('NEGATIVE: a short string well under the repeat count is safe', () => {
|
||||
expect(hasPeriodicTail('ababab')).toBe(false);
|
||||
});
|
||||
|
||||
// Regression (#444): a trivial single-char period (p===1) must NOT flag
|
||||
// legitimate divider/underline/whitespace runs. These are common in real
|
||||
// model output and previously false-positived at ~20 identical chars, aborting
|
||||
// the run and truncating output. They must all be treated as clean.
|
||||
it('NEGATIVE: a markdown horizontal rule is not flagged', () => {
|
||||
const text = 'text\n' + '-'.repeat(40);
|
||||
expect(hasPeriodicTail(text)).toBe(false);
|
||||
expect(isDegenerateOutput(text)).toBe(false);
|
||||
});
|
||||
|
||||
it('NEGATIVE: a setext heading underline is not flagged', () => {
|
||||
const text = 'Title\n' + '='.repeat(30);
|
||||
expect(hasPeriodicTail(text)).toBe(false);
|
||||
expect(isDegenerateOutput(text)).toBe(false);
|
||||
});
|
||||
|
||||
it('NEGATIVE: a box-drawing divider with no trailing newline is not flagged', () => {
|
||||
const text = 'done ' + '─'.repeat(50);
|
||||
expect(hasPeriodicTail(text)).toBe(false);
|
||||
expect(isDegenerateOutput(text)).toBe(false);
|
||||
});
|
||||
|
||||
it('NEGATIVE: trailing spaces are not flagged', () => {
|
||||
const text = 'answer' + ' '.repeat(40);
|
||||
expect(hasPeriodicTail(text)).toBe(false);
|
||||
expect(isDegenerateOutput(text)).toBe(false);
|
||||
});
|
||||
|
||||
// TRIVIAL_MIN_REPEATS boundary (#444 review). The monochar-tail branch fires at
|
||||
// EXACTLY 60 identical trailing chars (`run >= TRIVIAL_MIN_REPEATS`), so 59 is
|
||||
// clean and 60 trips. These pin the `>=` and MUST fail if the comparison is
|
||||
// flipped to `>` (the surviving mutation). The value 60 is HARD-CODED here on
|
||||
// purpose: TRIVIAL_MIN_REPEATS is a private constant and the assert must lock
|
||||
// the literal boundary the reviewer named, not track a constant edit.
|
||||
it('NEGATIVE: 59 identical trailing chars is one below the monochar threshold', () => {
|
||||
expect(hasPeriodicTail('x'.repeat(59))).toBe(false);
|
||||
expect(isDegenerateOutput('x'.repeat(59))).toBe(false);
|
||||
});
|
||||
|
||||
it('POSITIVE: 60 identical trailing chars hits the monochar threshold exactly', () => {
|
||||
// Fails if `run >= TRIVIAL_MIN_REPEATS` is mutated to `run > …`.
|
||||
expect(hasPeriodicTail('x'.repeat(60))).toBe(true);
|
||||
expect(isDegenerateOutput('x'.repeat(60))).toBe(true);
|
||||
});
|
||||
|
||||
// Positive counterparts: a GENUINE single-char runaway (hundreds+ of repeats)
|
||||
// and the real incident (period>=2, "loadTools." ×N) must still fire.
|
||||
it('POSITIVE: a genuine single-char runaway is still flagged', () => {
|
||||
const text = 'x'.repeat(5000);
|
||||
expect(hasPeriodicTail(text)).toBe(true);
|
||||
expect(isDegenerateOutput(text)).toBe(true);
|
||||
});
|
||||
|
||||
it('POSITIVE: the "loadTools." incident (period>=2) is still flagged', () => {
|
||||
const text = 'loadTools.'.repeat(500);
|
||||
expect(hasPeriodicTail(text)).toBe(true);
|
||||
expect(isDegenerateOutput(text)).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('truncateDegeneratedTail', () => {
|
||||
it('collapses a repeated-line loop to a few reps + marker', () => {
|
||||
const text = 'plan\n' + 'loadTools.\n'.repeat(20000);
|
||||
const out = truncateDegeneratedTail(text);
|
||||
expect(out.length).toBeLessThan(text.length);
|
||||
expect(out).toContain('output truncated');
|
||||
// Keeps the leading context and a few loop reps.
|
||||
expect(out).toContain('plan');
|
||||
expect((out.match(/loadTools\./g) ?? []).length).toBeLessThan(10);
|
||||
});
|
||||
|
||||
it('collapses a no-newline periodic loop to a few blocks + marker', () => {
|
||||
const text = 'answer: ' + 'xy'.repeat(50000);
|
||||
const out = truncateDegeneratedTail(text);
|
||||
expect(out.length).toBeLessThan(text.length);
|
||||
expect(out).toContain('output truncated');
|
||||
expect(out).toContain('answer:');
|
||||
});
|
||||
|
||||
it('returns non-degenerate text unchanged (by identity)', () => {
|
||||
const text = 'A perfectly normal, finished assistant answer.';
|
||||
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);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,215 @@
|
||||
/**
|
||||
* Token-degeneration detector for the in-app agent stream (#444).
|
||||
*
|
||||
* When the final-step lockdown is OFF (the new default) there is no toolChoice
|
||||
* override to strip the model's tools mid-work, so the anti-babble safety net is
|
||||
* this detector. It watches the accumulating assistant text and, on a runaway
|
||||
* repetition loop (the 255KB "loadTools." incident), aborts the run.
|
||||
*
|
||||
* Both rules are PURE functions of the text tail so they are cheap to run every
|
||||
* few KB and are unit-testable in isolation. They operate on the TAIL only
|
||||
* (`TAIL_WINDOW` chars) so the cost is bounded regardless of how long the turn is.
|
||||
*/
|
||||
|
||||
/** How many trailing chars of the accumulated text the rules inspect. */
|
||||
export const TAIL_WINDOW = 3000;
|
||||
|
||||
/** Rule 1 threshold: minimum consecutive identical non-empty lines to trigger. */
|
||||
export const REPEATED_LINES_THRESHOLD = 25;
|
||||
|
||||
/** Rule 2: maximum length of a repeating block considered for periodicity. */
|
||||
export const MAX_PERIOD_LEN = 150;
|
||||
|
||||
/** Rule 2: minimum number of consecutive block repeats to trigger. */
|
||||
export const MIN_PERIOD_REPEATS = 20;
|
||||
|
||||
/**
|
||||
* Rule 1 — ≥`REPEATED_LINES_THRESHOLD` consecutive IDENTICAL non-empty lines at
|
||||
* the tail. Catches the classic newline-delimited loop ("loadTools.\n" ×N).
|
||||
* Blank lines break a run (a table / list with blank separators never trips it);
|
||||
* a run of ordinary distinct lines (an edit list, code) never reaches the count.
|
||||
*
|
||||
* NB: `REPEATED_LINES_THRESHOLD` (25) is only THIS rule's own trigger, not the
|
||||
* effective floor for detecting a repeated-line loop. In practice a newline-
|
||||
* delimited repeat also has a fixed period (line + '\n'), so rule 2 catches it
|
||||
* via periodicity at `MIN_PERIOD_REPEATS` (20) repeats — the two rules combine
|
||||
* (see `isDegenerateOutput`), so the effective lower bound for a short identical
|
||||
* line loop is ~20, not 25. Pure.
|
||||
*/
|
||||
export function hasRepeatedLineRun(
|
||||
text: string,
|
||||
threshold = REPEATED_LINES_THRESHOLD,
|
||||
): boolean {
|
||||
const tail = text.length > TAIL_WINDOW ? text.slice(-TAIL_WINDOW) : text;
|
||||
const lines = tail.split('\n');
|
||||
let run = 1;
|
||||
let prev: string | null = null;
|
||||
for (const line of lines) {
|
||||
if (line.length > 0 && line === prev) {
|
||||
run += 1;
|
||||
if (run >= threshold) return true;
|
||||
} else {
|
||||
run = 1;
|
||||
}
|
||||
prev = line;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Rule 2 — cheap suffix-periodicity check: the tail ends in
|
||||
* ≥`MIN_PERIOD_REPEATS` back-to-back repeats of a single block of length
|
||||
* ≤`MAX_PERIOD_LEN`. Catches a no-newline repeat ("abcabcabc…") the line rule
|
||||
* misses. For each candidate period length p we verify the last `repeats*p`
|
||||
* chars are p-periodic; we stop at the smallest p that satisfies the repeat
|
||||
* count. Bounded by MAX_PERIOD_LEN × TAIL_WINDOW comparisons — negligible. Pure.
|
||||
*/
|
||||
export function hasPeriodicTail(
|
||||
text: string,
|
||||
maxPeriod = MAX_PERIOD_LEN,
|
||||
minRepeats = MIN_PERIOD_REPEATS,
|
||||
): boolean {
|
||||
const tail = text.length > TAIL_WINDOW ? text.slice(-TAIL_WINDOW) : text;
|
||||
const n = tail.length;
|
||||
// Not even the shortest possible loop fits in the tail.
|
||||
if (n < minRepeats) return false;
|
||||
// A tail of ONE repeated char (a "trivial period") is common in LEGIT output —
|
||||
// markdown rules (----/====), setext underlines, box-drawing dividers,
|
||||
// trailing spaces routinely produce 20–50 identical chars. Such a run is
|
||||
// p-periodic for EVERY p, so it would otherwise trip the block rule at p>=2
|
||||
// too, not just p===1. We therefore split the check: a monochar tail needs far
|
||||
// more repeats (a real single-char babble loop produces hundreds-to-thousands;
|
||||
// 60 is well above any realistic divider yet a fifth of TAIL_WINDOW), while a
|
||||
// genuine multi-char block repeat (>=2 distinct chars, e.g. the "loadTools."
|
||||
// incident, period ~10) keeps the normal MIN_PERIOD_REPEATS threshold.
|
||||
const TRIVIAL_MIN_REPEATS = 60;
|
||||
|
||||
// Monochar-tail check (the trivial-period case): count the trailing run of one
|
||||
// identical char and require TRIVIAL_MIN_REPEATS of them.
|
||||
{
|
||||
const last = tail[n - 1];
|
||||
let run = 1;
|
||||
for (let i = n - 2; i >= 0 && tail[i] === last; i--) run++;
|
||||
if (run >= TRIVIAL_MIN_REPEATS) return true;
|
||||
}
|
||||
|
||||
const maxP = maxPeriod;
|
||||
for (let p = 2; p <= maxP; p++) {
|
||||
// Verify the last (minRepeats*p) chars are p-periodic AND not monochar (a
|
||||
// monochar span is the trivial case handled above, so skip it here to avoid
|
||||
// re-flagging a legit divider at a composite period).
|
||||
const span = minRepeats * p;
|
||||
// Not enough tail to hold this many repeats of this period.
|
||||
if (span > n) continue;
|
||||
const start = n - span;
|
||||
let periodic = true;
|
||||
let multiChar = false;
|
||||
for (let i = n - 1; i >= start + p; i--) {
|
||||
if (tail[i] !== tail[i - p]) {
|
||||
periodic = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (!periodic) continue;
|
||||
// Confirm the block itself has >=2 distinct chars (else it's monochar).
|
||||
for (let i = start + 1; i < n; i++) {
|
||||
if (tail[i] !== tail[start]) {
|
||||
multiChar = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (multiChar) return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Combined guard used by the stream's onChunk: true when EITHER rule fires.
|
||||
* Pure — the caller owns the abort side effect.
|
||||
*/
|
||||
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
|
||||
* `keepRepeats` repeats of the detected loop, then appends a short marker. If no
|
||||
* loop is detected the text is returned unchanged (by identity).
|
||||
*
|
||||
* Implementation: find the shortest tail period (same check as hasPeriodicTail),
|
||||
* keep the prefix before the loop plus `keepRepeats` copies of the block, drop
|
||||
* the rest. This is best-effort cosmetic trimming; correctness does not depend on
|
||||
* finding the exact minimal loop. Pure.
|
||||
*/
|
||||
export function truncateDegeneratedTail(
|
||||
text: string,
|
||||
keepRepeats = 3,
|
||||
): string {
|
||||
const marker = '\n…[output truncated: repeated token loop detected]';
|
||||
// Try the line rule first: collapse a long run of identical lines.
|
||||
const lines = text.split('\n');
|
||||
let runStart = -1;
|
||||
let run = 1;
|
||||
for (let i = 1; i < lines.length; i++) {
|
||||
if (lines[i].length > 0 && lines[i] === lines[i - 1]) {
|
||||
if (run === 1) runStart = i - 1;
|
||||
run += 1;
|
||||
if (run >= REPEATED_LINES_THRESHOLD) {
|
||||
const kept = lines.slice(0, runStart + keepRepeats).join('\n');
|
||||
return kept + marker;
|
||||
}
|
||||
} else {
|
||||
run = 1;
|
||||
runStart = -1;
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to periodicity over the whole string (bounded by the same block
|
||||
// length). Find the smallest period that makes the SUFFIX highly repetitive.
|
||||
const n = text.length;
|
||||
const maxP = Math.min(MAX_PERIOD_LEN, Math.floor(n / MIN_PERIOD_REPEATS));
|
||||
for (let p = 1; p <= maxP; p++) {
|
||||
// Count how many trailing p-blocks are periodic.
|
||||
let reps = 1;
|
||||
let i = n - 1;
|
||||
for (; i >= p; i--) {
|
||||
if (text[i] !== text[i - p]) break;
|
||||
}
|
||||
// The loop above walks over the periodic suffix; its length is (n-1 - i).
|
||||
const periodicLen = n - 1 - i;
|
||||
reps = Math.floor(periodicLen / p) + 1;
|
||||
if (reps >= MIN_PERIOD_REPEATS) {
|
||||
const loopStart = n - reps * p; // start of the fully-periodic suffix
|
||||
const kept = text.slice(0, loopStart + keepRepeats * p);
|
||||
return kept + marker;
|
||||
}
|
||||
}
|
||||
return text;
|
||||
}
|
||||
@@ -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;
|
||||
});
|
||||
});
|
||||
@@ -23,7 +23,15 @@ import { SHARED_TOOL_SPECS } from '../../../../../../packages/mcp/src/tool-specs
|
||||
// sync.
|
||||
const mockLoaded = (DocmostClient: loader.DocmostClientCtor) => ({
|
||||
DocmostClient,
|
||||
sharedToolSpecs: SHARED_TOOL_SPECS as Record<string, loader.SharedToolSpec>,
|
||||
sharedToolSpecs: SHARED_TOOL_SPECS as unknown as Record<string, loader.SharedToolSpec>,
|
||||
// Pure no-network draw.io helpers (#424). Type-correct stubs: these tests
|
||||
// never execute the drawioShapes / drawioGuide tool bodies.
|
||||
searchShapes: (() => []) as unknown as loader.SearchShapesFn,
|
||||
getGuideSection: (() => ({
|
||||
section: 'index',
|
||||
content: '',
|
||||
sections: [],
|
||||
})) as unknown as loader.GetGuideSectionFn,
|
||||
});
|
||||
|
||||
/**
|
||||
@@ -275,6 +283,7 @@ describe('AiChatToolsService node-arg JSON-string coercion', () => {
|
||||
const patchNodeCalls: unknown[][] = [];
|
||||
const insertNodeCalls: unknown[][] = [];
|
||||
const updatePageJsonCalls: unknown[][] = [];
|
||||
const updatePageCalls: unknown[][] = [];
|
||||
|
||||
const fakeClient: FakeDocmostClient = {
|
||||
patchNode: (...args: unknown[]) => {
|
||||
@@ -289,6 +298,11 @@ describe('AiChatToolsService node-arg JSON-string coercion', () => {
|
||||
updatePageJsonCalls.push(args);
|
||||
return Promise.resolve({ ok: true });
|
||||
},
|
||||
// Backs the plain-Markdown full-body-replace tool updatePageMarkdown (#411).
|
||||
updatePage: (...args: unknown[]) => {
|
||||
updatePageCalls.push(args);
|
||||
return Promise.resolve({ success: true });
|
||||
},
|
||||
};
|
||||
|
||||
const tokenServiceStub = {
|
||||
@@ -302,6 +316,7 @@ describe('AiChatToolsService node-arg JSON-string coercion', () => {
|
||||
patchNodeCalls.length = 0;
|
||||
insertNodeCalls.length = 0;
|
||||
updatePageJsonCalls.length = 0;
|
||||
updatePageCalls.length = 0;
|
||||
jest.spyOn(loader, 'loadDocmostMcp').mockResolvedValue(
|
||||
mockLoaded(function () {
|
||||
return fakeClient as DocmostClientLike;
|
||||
@@ -340,23 +355,32 @@ describe('AiChatToolsService node-arg JSON-string coercion', () => {
|
||||
content: [{ type: 'text', text: 'Hello' }],
|
||||
};
|
||||
|
||||
it('patchNode parses a JSON-string node and forwards it as an object', async () => {
|
||||
it('patchNode parses a JSON-string node and forwards it as { node } (object)', async () => {
|
||||
const tools = await buildTools();
|
||||
await tools.patchNode.execute(
|
||||
{ pageId: 'p1', nodeId: 'n1', node: JSON.stringify(NODE_OBJ) } as never,
|
||||
{} as never,
|
||||
);
|
||||
expect(patchNodeCalls).toHaveLength(1);
|
||||
expect(patchNodeCalls[0]).toEqual(['p1', 'n1', NODE_OBJ]);
|
||||
// #413: the 3rd arg is now the XOR input { markdown?, node? }.
|
||||
expect(patchNodeCalls[0]).toEqual([
|
||||
'p1',
|
||||
'n1',
|
||||
{ markdown: undefined, node: NODE_OBJ },
|
||||
]);
|
||||
});
|
||||
|
||||
it('patchNode passes an object node through unchanged', async () => {
|
||||
it('patchNode passes an object node through unchanged inside { node }', async () => {
|
||||
const tools = await buildTools();
|
||||
await tools.patchNode.execute(
|
||||
{ pageId: 'p1', nodeId: 'n1', node: NODE_OBJ } as never,
|
||||
{} as never,
|
||||
);
|
||||
expect(patchNodeCalls[0]).toEqual(['p1', 'n1', NODE_OBJ]);
|
||||
expect(patchNodeCalls[0]).toEqual([
|
||||
'p1',
|
||||
'n1',
|
||||
{ markdown: undefined, node: NODE_OBJ },
|
||||
]);
|
||||
});
|
||||
|
||||
it('patchNode throws the documented message on invalid JSON string', async () => {
|
||||
@@ -370,7 +394,7 @@ describe('AiChatToolsService node-arg JSON-string coercion', () => {
|
||||
expect(patchNodeCalls).toHaveLength(0);
|
||||
});
|
||||
|
||||
it('insertNode parses a JSON-string node and forwards it as an object', async () => {
|
||||
it('insertNode parses a JSON-string node and forwards it inside { node }', async () => {
|
||||
const tools = await buildTools();
|
||||
await tools.insertNode.execute(
|
||||
{
|
||||
@@ -381,9 +405,15 @@ describe('AiChatToolsService node-arg JSON-string coercion', () => {
|
||||
{} as never,
|
||||
);
|
||||
expect(insertNodeCalls).toHaveLength(1);
|
||||
const [pageId, node] = insertNodeCalls[0];
|
||||
// #413: the 2nd arg is the XOR input { markdown?, node? }, the 3rd is opts.
|
||||
const [pageId, input, opts] = insertNodeCalls[0] as [
|
||||
string,
|
||||
{ markdown?: unknown; node?: unknown },
|
||||
{ position?: string },
|
||||
];
|
||||
expect(pageId).toBe('p1');
|
||||
expect(node).toEqual(NODE_OBJ);
|
||||
expect(input).toEqual({ markdown: undefined, node: NODE_OBJ });
|
||||
expect(opts.position).toBe('append');
|
||||
});
|
||||
|
||||
it('insertNode throws the documented message on invalid JSON string', async () => {
|
||||
@@ -437,6 +467,54 @@ describe('AiChatToolsService node-arg JSON-string coercion', () => {
|
||||
).rejects.toThrow('content was a string but not valid JSON');
|
||||
expect(updatePageJsonCalls).toHaveLength(0);
|
||||
});
|
||||
|
||||
// #411: the plain-Markdown full-body-replace tool is now the shared
|
||||
// `updatePageMarkdown` (was inline `updatePageContent`). It forwards to
|
||||
// client.updatePage(pageId, content, title) -> updatePageContentRealtime ->
|
||||
// markdownToProseMirrorCanonical, so `^[...]` footnotes materialize.
|
||||
it('updatePageMarkdown forwards { pageId, content, title } to client.updatePage', async () => {
|
||||
const tools = await buildTools();
|
||||
await tools.updatePageMarkdown.execute(
|
||||
{ pageId: 'p1', content: 'Body^[a note]', title: 'New title' } as never,
|
||||
{} as never,
|
||||
);
|
||||
expect(updatePageCalls).toHaveLength(1);
|
||||
expect(updatePageCalls[0]).toEqual(['p1', 'Body^[a note]', 'New title']);
|
||||
});
|
||||
|
||||
it('updatePageMarkdown returns the RAW client result in-app (deliberate #411 shape change, documented on the spec)', async () => {
|
||||
const tools = await buildTools();
|
||||
// Registry canonical execute returns client.updatePage's result verbatim.
|
||||
// The old inline tool projected to { pageId, updated }; the rename now
|
||||
// surfaces the raw result (nothing reads the removed `.updated`; the raw
|
||||
// shape carries footnote/verify warnings and matches the on-both-hosts
|
||||
// registry convention). fakeClient.updatePage resolves { success: true }.
|
||||
const result = await tools.updatePageMarkdown.execute(
|
||||
{ pageId: 'p1', content: '# Hi' } as never,
|
||||
{} as never,
|
||||
);
|
||||
expect(result).toEqual({ success: true });
|
||||
});
|
||||
|
||||
it('updatePageMarkdown forwards title=undefined when omitted', async () => {
|
||||
const tools = await buildTools();
|
||||
await tools.updatePageMarkdown.execute(
|
||||
{ pageId: 'p1', content: '# Hi' } as never,
|
||||
{} as never,
|
||||
);
|
||||
expect(updatePageCalls[0]).toEqual(['p1', '# Hi', undefined]);
|
||||
});
|
||||
|
||||
// #411 surface split: the plain-Markdown replace tool exists in-app under the
|
||||
// new key; the OLD inline updatePageContent key is gone; importPageMarkdown is
|
||||
// still present IN-APP (only the external MCP surface drops it — asserted in
|
||||
// packages/mcp/test/unit/tool-inventory.test.mjs).
|
||||
it('exposes updatePageMarkdown in-app, no legacy updatePageContent, keeps importPageMarkdown', async () => {
|
||||
const tools = await buildTools();
|
||||
expect(tools.updatePageMarkdown).toBeDefined();
|
||||
expect((tools as Record<string, unknown>).updatePageContent).toBeUndefined();
|
||||
expect(tools.importPageMarkdown).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
@@ -847,3 +925,109 @@ describe('AiChatToolsService getCurrentPage selection (#388)', () => {
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* #440 review: the in-app drawioCreate / drawioUpdate handlers must forward
|
||||
* the optional `layout:"elk"` param to the client (5th positional arg), exactly
|
||||
* like the MCP host. It was silently dropped, so ELK auto-layout worked only via
|
||||
* the standalone MCP server, not in-app. These tests pin per-host parity.
|
||||
*/
|
||||
describe('AiChatToolsService drawio layout passthrough (#440)', () => {
|
||||
const createCalls: unknown[][] = [];
|
||||
const updateCalls: unknown[][] = [];
|
||||
|
||||
// FakeDocmostClient (not Partial<DocmostClientLike>): since #446 derived
|
||||
// DocmostClientLike from the real client, its drawioCreate/drawioUpdate return
|
||||
// the concrete result shape, so a minimal stub object would not be assignable.
|
||||
// FakeDocmostClient types every method as (...args) => Promise<any>, which is
|
||||
// exactly what these arg-capturing doubles need.
|
||||
const fakeClient: FakeDocmostClient = {
|
||||
drawioCreate: (...args: unknown[]) => {
|
||||
createCalls.push(args);
|
||||
return Promise.resolve({ success: true, nodeId: '#0' });
|
||||
},
|
||||
drawioUpdate: (...args: unknown[]) => {
|
||||
updateCalls.push(args);
|
||||
return Promise.resolve({ success: true, nodeId: '#0' });
|
||||
},
|
||||
};
|
||||
|
||||
const tokenServiceStub = {
|
||||
generateAccessToken: jest.fn().mockResolvedValue('access-token'),
|
||||
generateCollabToken: jest.fn().mockResolvedValue('collab-token'),
|
||||
};
|
||||
|
||||
let service: AiChatToolsService;
|
||||
|
||||
beforeEach(() => {
|
||||
createCalls.length = 0;
|
||||
updateCalls.length = 0;
|
||||
jest.spyOn(loader, 'loadDocmostMcp').mockResolvedValue(
|
||||
mockLoaded(function () {
|
||||
return fakeClient as DocmostClientLike;
|
||||
} as unknown as loader.DocmostClientCtor),
|
||||
);
|
||||
service = new AiChatToolsService(
|
||||
tokenServiceStub as never,
|
||||
{} as never,
|
||||
{} as never,
|
||||
{} as never,
|
||||
{} as never,
|
||||
{
|
||||
asSink: () => ({ put: jest.fn(), has: jest.fn(), evict: jest.fn() }),
|
||||
} as never,
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => jest.restoreAllMocks());
|
||||
|
||||
const buildTools = () =>
|
||||
service.forUser(
|
||||
{ id: 'user-1', email: 'u@example.com', workspaceId: 'ws-1' } as never,
|
||||
'session-1',
|
||||
'ws-1',
|
||||
'chat-1',
|
||||
);
|
||||
|
||||
it('forwards layout:"elk" to client.drawioCreate as the 5th positional arg', async () => {
|
||||
const tools = await buildTools();
|
||||
await tools.drawioCreate.execute(
|
||||
{
|
||||
pageId: 'p-1',
|
||||
xml: '<mxGraphModel/>',
|
||||
position: 'append',
|
||||
layout: 'elk',
|
||||
} as never,
|
||||
{} as never,
|
||||
);
|
||||
expect(createCalls).toHaveLength(1);
|
||||
// drawioCreate(pageId, where, xml, title, layout) — layout is args[4].
|
||||
expect(createCalls[0][4]).toBe('elk');
|
||||
});
|
||||
|
||||
it('forwards layout:"elk" to client.drawioUpdate as the 5th positional arg', async () => {
|
||||
const tools = await buildTools();
|
||||
await tools.drawioUpdate.execute(
|
||||
{
|
||||
pageId: 'p-1',
|
||||
node: '#0',
|
||||
xml: '<mxGraphModel/>',
|
||||
baseHash: 'h',
|
||||
layout: 'elk',
|
||||
} as never,
|
||||
{} as never,
|
||||
);
|
||||
expect(updateCalls).toHaveLength(1);
|
||||
// drawioUpdate(pageId, node, xml, baseHash, layout) — layout is args[4].
|
||||
expect(updateCalls[0][4]).toBe('elk');
|
||||
});
|
||||
|
||||
it('omits layout (undefined 5th arg) when not requested', async () => {
|
||||
const tools = await buildTools();
|
||||
await tools.drawioCreate.execute(
|
||||
{ pageId: 'p-1', xml: '<mxGraphModel/>', position: 'append' } as never,
|
||||
{} as never,
|
||||
);
|
||||
expect(createCalls[0][4]).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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';
|
||||
@@ -59,10 +59,12 @@ function __assertClientCallContract(client: DocmostClientLike): void {
|
||||
void client.getWorkspace();
|
||||
void client.getSpaces();
|
||||
void client.listPages(s, n, true);
|
||||
void client.getTree(s, s, n);
|
||||
void client.getPageContext(s);
|
||||
void client.listSidebarPages(s, s);
|
||||
void client.getOutline(s);
|
||||
void client.getPageJson(s);
|
||||
void client.getNode(s, s);
|
||||
void client.getNode(s, s, 'markdown');
|
||||
void client.searchInPage(s, s, {
|
||||
regex: true,
|
||||
caseSensitive: true,
|
||||
@@ -84,12 +86,16 @@ function __assertClientCallContract(client: DocmostClientLike): void {
|
||||
void client.movePage(s, s, s);
|
||||
void client.deletePage(s);
|
||||
void client.editPageText(s, edits);
|
||||
void client.patchNode(s, s, node);
|
||||
void client.insertNode(s, node, {
|
||||
position: 'append',
|
||||
anchorNodeId: s,
|
||||
anchorText: s,
|
||||
});
|
||||
void client.patchNode(s, s, { markdown: s, node });
|
||||
void client.insertNode(
|
||||
s,
|
||||
{ markdown: s, node },
|
||||
{
|
||||
position: 'append',
|
||||
anchorNodeId: s,
|
||||
anchorText: s,
|
||||
},
|
||||
);
|
||||
void client.deleteNode(s, s);
|
||||
void client.updatePageJson(s, node, s);
|
||||
void client.tableInsertRow(s, s, cells, n);
|
||||
@@ -111,10 +117,29 @@ function __assertClientCallContract(client: DocmostClientLike): void {
|
||||
afterText: s,
|
||||
});
|
||||
void client.replaceImage(s, s, s, { align, alt: s });
|
||||
// --- draw.io diagrams (#423) ---
|
||||
// --- draw.io diagrams (#423 stage 1, #424 stage 2) ---
|
||||
// The 5th `layout` arg (#424) is exercised so this parity assertion fails if the
|
||||
// client signature drops it — it must reach the client from the shared execute.
|
||||
void client.drawioGet(s, s, 'xml');
|
||||
void client.drawioCreate(s, { position: 'append', anchorNodeId: s }, s, s);
|
||||
void client.drawioUpdate(s, s, s, s);
|
||||
void client.drawioCreate(s, { position: 'append', anchorNodeId: s }, s, s, 'elk');
|
||||
void client.drawioUpdate(s, s, s, s, 'elk');
|
||||
// --- draw.io high-level semantic tools (#425 stage 3) ---
|
||||
void client.drawioEditCells(s, s, [{ op: 'delete', cellId: s }], s);
|
||||
void client.drawioFromGraph(
|
||||
s,
|
||||
{ position: 'append', anchorNodeId: s },
|
||||
{ nodes: [{ id: s, label: s }] },
|
||||
'LR',
|
||||
s,
|
||||
'full',
|
||||
s,
|
||||
);
|
||||
void client.drawioFromMermaid(
|
||||
s,
|
||||
{ position: 'append', anchorNodeId: s },
|
||||
s,
|
||||
s,
|
||||
);
|
||||
// --- write (comment) ---
|
||||
void client.createComment(s, s, 'inline', s, s, s);
|
||||
void client.resolveComment(s, true);
|
||||
@@ -134,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);
|
||||
@@ -161,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`;
|
||||
@@ -263,8 +416,19 @@ export class AiChatToolsService {
|
||||
// provenance tokens) and load the shared tool-spec registry. Client
|
||||
// construction is shared with the page-change detection path (#274) via
|
||||
// buildDocmostClient so both go over the exact same authenticated route.
|
||||
const { sharedToolSpecs, createCommentSignalTracker } =
|
||||
await loadDocmostMcp();
|
||||
// searchShapes / getGuideSection (#424) are the PURE, no-network helpers
|
||||
// backing drawioShapes / drawioGuide. They are `inlineBothHosts` specs (no
|
||||
// canonical execute — their catalog loader uses import.meta and can't be
|
||||
// value-imported into the zod-agnostic tool-specs.ts under the server's
|
||||
// commonjs type-check), so the shared registry loop below SKIPS them and this
|
||||
// service wires them inline (see drawioShapes/drawioGuide entries), mirroring
|
||||
// how index.ts registers them on the standalone MCP host.
|
||||
const {
|
||||
sharedToolSpecs,
|
||||
createCommentSignalTracker,
|
||||
searchShapes,
|
||||
getGuideSection,
|
||||
} = await loadDocmostMcp();
|
||||
const client = await this.buildDocmostClient(
|
||||
user,
|
||||
sessionId,
|
||||
@@ -292,6 +456,18 @@ export class AiChatToolsService {
|
||||
execute,
|
||||
});
|
||||
|
||||
// The in-app toolset. It starts with the tools kept INLINE here for a
|
||||
// documented per-layer reason: an intentional behaviour/schema divergence from
|
||||
// the standalone MCP surface (searchPages' hybrid RRF,
|
||||
// transformPage's guardrailed shorter schema), a name clash the shared
|
||||
// registry forbids (in-app `getTable` verb-first vs the MCP noun-first
|
||||
// `tableGet` — the registry requires mcpName === inAppKey), per-request
|
||||
// state the registry loop cannot provide
|
||||
// (getCurrentPage reads the resolved openedPage; searchPages closes over the
|
||||
// per-request user/embedding deps), or a tool with no MCP twin
|
||||
// (listSidebarPages/getComment/getPageHistory). Every SHARED tool is then added
|
||||
// by the registry loop below (see it), so there is exactly one arg-mapping per
|
||||
// shared tool and it can never drift from the MCP host again (#445).
|
||||
const tools: Record<string, Tool> = {
|
||||
// INTENTIONAL per-transport divergence (not in the shared registry): this
|
||||
// in-app search runs a semantic + keyword hybrid (RRF) with in-process
|
||||
@@ -426,180 +602,14 @@ export class AiChatToolsService {
|
||||
execute: async () => resolveCurrentPageResult(openedPage),
|
||||
}),
|
||||
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
// The execute body keeps this layer's { title, markdown } projection.
|
||||
getPage: sharedTool(sharedToolSpecs.getPage, async ({ pageId }) => {
|
||||
// getPage(pageId) -> { data: filterPage(page, markdown), success }.
|
||||
const result = await client.getPage(pageId);
|
||||
const data = (result?.data ?? {}) as {
|
||||
title?: string;
|
||||
content?: string;
|
||||
};
|
||||
return {
|
||||
title: data.title ?? '',
|
||||
markdown: typeof data.content === 'string' ? data.content : '',
|
||||
};
|
||||
}),
|
||||
|
||||
// --- WRITE tools (all reversible — history/trash; §6.5 / D3) ---
|
||||
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
createPage: sharedTool(
|
||||
sharedToolSpecs.createPage,
|
||||
async ({ title, content, spaceId, parentPageId }) => {
|
||||
// createPage(title, content, spaceId, parentPageId?) ->
|
||||
// { data: filterPage(page, markdown), success }.
|
||||
const result = await client.createPage(
|
||||
title,
|
||||
content ?? '',
|
||||
spaceId,
|
||||
parentPageId,
|
||||
);
|
||||
const data = (result?.data ?? {}) as {
|
||||
id?: string;
|
||||
slugId?: string;
|
||||
title?: string;
|
||||
};
|
||||
return { id: data.id ?? data.slugId, title: data.title ?? title };
|
||||
},
|
||||
),
|
||||
|
||||
updatePageContent: tool({
|
||||
description:
|
||||
"Replace a page's body with new Markdown content (and optionally its " +
|
||||
'title). Reversible: the previous version is kept in page history.',
|
||||
inputSchema: modelFriendlyInput({
|
||||
pageId: z.string().describe('The id of the page to update.'),
|
||||
content: z.string().describe('The new page body as Markdown.'),
|
||||
title: z
|
||||
.string()
|
||||
.optional()
|
||||
.describe('Optional new title for the page.'),
|
||||
}),
|
||||
execute: async ({ pageId, content, title }) => {
|
||||
// updatePage mutates the live collab doc -> provenance flows from the
|
||||
// collab-token provider. Returns { success, modified, message, pageId }.
|
||||
const result = (await client.updatePage(pageId, content, title)) as {
|
||||
success?: boolean;
|
||||
};
|
||||
return { pageId, updated: result?.success ?? true };
|
||||
},
|
||||
}),
|
||||
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
renamePage: sharedTool(
|
||||
sharedToolSpecs.renamePage,
|
||||
async ({ pageId, title }) => {
|
||||
// renamePage(pageId, title) -> { success, pageId, title }.
|
||||
await client.renamePage(pageId, title);
|
||||
return { pageId, title };
|
||||
},
|
||||
),
|
||||
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
// The shared schema adds the optional `position` field this layer lacked
|
||||
// before; the execute now forwards it (the client already accepted it).
|
||||
movePage: sharedTool(
|
||||
sharedToolSpecs.movePage,
|
||||
async ({ pageId, parentPageId, position }) => {
|
||||
// movePage(pageId, parentPageId, position?) -> raw move response.
|
||||
await client.movePage(pageId, parentPageId ?? null, position);
|
||||
return { pageId, parentPageId: parentPageId ?? null, moved: true };
|
||||
},
|
||||
),
|
||||
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
// GUARDRAIL (§14 H4) preserved: the shared schema exposes ONLY pageId, so
|
||||
// permanentlyDelete/forceDelete are never part of the input and can never
|
||||
// be forwarded — the agent physically cannot permanently delete a page.
|
||||
deletePage: sharedTool(sharedToolSpecs.deletePage, async ({ pageId }) => {
|
||||
// deletePage(pageId) hits POST /pages/delete with { pageId } only,
|
||||
// which is the soft-delete (trash) path on the server.
|
||||
await client.deletePage(pageId);
|
||||
return { pageId, trashed: true };
|
||||
}),
|
||||
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
// This layer keeps only its own execute-side guards (require a selection
|
||||
// for a top-level comment; reject suggestedText on a reply / without a
|
||||
// selection) — the schema+description are shared.
|
||||
createComment: sharedTool(
|
||||
sharedToolSpecs.createComment,
|
||||
async ({
|
||||
pageId,
|
||||
content,
|
||||
selection,
|
||||
parentCommentId,
|
||||
suggestedText,
|
||||
}) => {
|
||||
// createComment(pageId, content, type, selection?, parentCommentId?,
|
||||
// suggestedText?). Top-level comments are inline and must carry a
|
||||
// selection to anchor on; replies inherit the parent's anchor (no
|
||||
// selection). Throwing here surfaces a tool error to the model (Vercel
|
||||
// `ai` SDK) so the agent retries with a better selection — do not
|
||||
// catch/suppress it.
|
||||
if (!parentCommentId && (!selection || !selection.trim())) {
|
||||
throw new Error(
|
||||
"createComment requires a 'selection' (exact text to anchor on) for a new top-level comment.",
|
||||
);
|
||||
}
|
||||
if (suggestedText !== undefined) {
|
||||
if (parentCommentId) {
|
||||
throw new Error(
|
||||
"createComment: 'suggestedText' cannot be attached to a reply; it applies only to a top-level inline comment.",
|
||||
);
|
||||
}
|
||||
if (!selection || !selection.trim()) {
|
||||
throw new Error(
|
||||
"createComment: 'suggestedText' requires a 'selection' to anchor and rewrite.",
|
||||
);
|
||||
}
|
||||
}
|
||||
const result = await client.createComment(
|
||||
pageId,
|
||||
content,
|
||||
'inline',
|
||||
selection,
|
||||
parentCommentId,
|
||||
suggestedText,
|
||||
);
|
||||
const data = (result?.data ?? {}) as { id?: string };
|
||||
return { commentId: data.id, pageId };
|
||||
},
|
||||
),
|
||||
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
resolveComment: sharedTool(
|
||||
sharedToolSpecs.resolveComment,
|
||||
async ({ commentId, resolved }) => {
|
||||
// resolveComment(commentId, resolved) -> { success, commentId, resolved }.
|
||||
await client.resolveComment(commentId, resolved);
|
||||
return { commentId, resolved };
|
||||
},
|
||||
),
|
||||
|
||||
// --- READ tools (added) ---
|
||||
|
||||
getWorkspace: sharedTool(
|
||||
sharedToolSpecs.getWorkspace,
|
||||
async () => await client.getWorkspace(),
|
||||
),
|
||||
|
||||
listSpaces: sharedTool(
|
||||
sharedToolSpecs.listSpaces,
|
||||
async () => await client.getSpaces(),
|
||||
),
|
||||
|
||||
// INTENTIONAL per-transport divergence (not shared): keeps the `tree:true`
|
||||
// hierarchy mode but is worded for the in-app agent; the standalone MCP
|
||||
// `list_pages` carries its own wording. Kept per-layer so each side tunes
|
||||
// its own guidance.
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
listPages: sharedTool(
|
||||
sharedToolSpecs.listPages,
|
||||
async ({ spaceId, limit, tree }) =>
|
||||
await client.listPages(spaceId, limit, tree),
|
||||
),
|
||||
//
|
||||
// NOTE (issue #411): the plain-Markdown full-body-replace tool is no longer
|
||||
// inline here — it moved to @docmost/mcp's SHARED_TOOL_SPECS as
|
||||
// `updatePageMarkdown` (was inline `updatePageContent`) so it registers on
|
||||
// BOTH the external MCP and the in-app agent. The registry loop below adds
|
||||
// it under its inAppKey. importPageMarkdown stays a shared spec too (now
|
||||
// inAppOnly — dropped from the external MCP surface, kept in-app).
|
||||
|
||||
listSidebarPages: tool({
|
||||
description:
|
||||
@@ -619,34 +629,9 @@ export class AiChatToolsService {
|
||||
await client.listSidebarPages(spaceId, pageId),
|
||||
}),
|
||||
|
||||
getOutline: sharedTool(
|
||||
sharedToolSpecs.getOutline,
|
||||
async ({ pageId }) => await client.getOutline(pageId),
|
||||
),
|
||||
|
||||
getPageJson: sharedTool(
|
||||
sharedToolSpecs.getPageJson,
|
||||
async ({ pageId }) => await client.getPageJson(pageId),
|
||||
),
|
||||
|
||||
getNode: sharedTool(
|
||||
sharedToolSpecs.getNode,
|
||||
async ({ pageId, nodeId }) => await client.getNode(pageId, nodeId),
|
||||
),
|
||||
|
||||
searchInPage: sharedTool(
|
||||
sharedToolSpecs.searchInPage,
|
||||
async ({ pageId, query, regex, caseSensitive, limit }) =>
|
||||
await client.searchInPage(pageId, query, {
|
||||
regex,
|
||||
caseSensitive,
|
||||
limit,
|
||||
}),
|
||||
),
|
||||
|
||||
// NOT shared (kept inline): the MCP tool name `table_get` is noun-first
|
||||
// while this key is `getTable` (verb-first), breaking the
|
||||
// snake_case(inAppKey) convention the shared registry enforces. Its
|
||||
// NOT shared (kept inline): the MCP tool name `tableGet` is noun-first
|
||||
// while this key is `getTable` (verb-first), so it cannot satisfy the
|
||||
// shared registry's `mcpName === inAppKey` convention (#412). Its
|
||||
// reference parameter is still named `table` (was `tableRef`) so it matches
|
||||
// the migrated table row/cell tools below.
|
||||
getTable: tool({
|
||||
@@ -666,13 +651,6 @@ export class AiChatToolsService {
|
||||
await client.getTable(pageId, table),
|
||||
}),
|
||||
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
listComments: sharedTool(
|
||||
sharedToolSpecs.listComments,
|
||||
async ({ pageId, includeResolved }) =>
|
||||
await client.listComments(pageId, includeResolved),
|
||||
),
|
||||
|
||||
getComment: tool({
|
||||
description: 'Fetch a single comment by id (content as Markdown).',
|
||||
inputSchema: modelFriendlyInput({
|
||||
@@ -681,24 +659,6 @@ export class AiChatToolsService {
|
||||
execute: async ({ commentId }) => await client.getComment(commentId),
|
||||
}),
|
||||
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
checkNewComments: sharedTool(
|
||||
sharedToolSpecs.checkNewComments,
|
||||
async ({ spaceId, since, parentPageId }) =>
|
||||
await client.checkNewComments(spaceId, since, parentPageId),
|
||||
),
|
||||
|
||||
listShares: sharedTool(
|
||||
sharedToolSpecs.listShares,
|
||||
async () => await client.listShares(),
|
||||
),
|
||||
|
||||
listPageHistory: sharedTool(
|
||||
sharedToolSpecs.listPageHistory,
|
||||
async ({ pageId, cursor }) =>
|
||||
await client.listPageHistory(pageId, cursor),
|
||||
),
|
||||
|
||||
getPageHistory: tool({
|
||||
description:
|
||||
'Fetch a single page-history version including its lossless ' +
|
||||
@@ -710,206 +670,11 @@ export class AiChatToolsService {
|
||||
await client.getPageHistory(historyId),
|
||||
}),
|
||||
|
||||
diffPageVersions: sharedTool(
|
||||
sharedToolSpecs.diffPageVersions,
|
||||
async ({ pageId, from, to }) =>
|
||||
await client.diffPageVersions(pageId, from, to),
|
||||
),
|
||||
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
exportPageMarkdown: sharedTool(
|
||||
sharedToolSpecs.exportPageMarkdown,
|
||||
async ({ pageId }) => {
|
||||
const markdown = await client.exportPageMarkdown(pageId);
|
||||
return { markdown };
|
||||
},
|
||||
),
|
||||
|
||||
// --- WRITE tools (added; reversible via page history/trash) ---
|
||||
|
||||
editPageText: sharedTool(
|
||||
sharedToolSpecs.editPageText,
|
||||
async ({ pageId, edits }) => await client.editPageText(pageId, edits),
|
||||
),
|
||||
|
||||
// Returns ONLY the short link object — never the document body — so a
|
||||
// large page can be handed to an external consumer without bloating
|
||||
// context.
|
||||
stashPage: sharedTool(
|
||||
sharedToolSpecs.stashPage,
|
||||
async ({ pageId }) => await client.stashPage(pageId),
|
||||
),
|
||||
|
||||
// Schema + description from the shared registry (identical across both
|
||||
// transports). The execute body keeps its OWN parseNodeArg normalization:
|
||||
// the model sometimes serializes the node as a JSON string, and we parse it
|
||||
// before the client's typeof-object guard rejects it (parity with the
|
||||
// standalone MCP server, index.ts patch_node).
|
||||
patchNode: sharedTool(
|
||||
sharedToolSpecs.patchNode,
|
||||
async ({ pageId, nodeId, node }) => {
|
||||
const parsedNode = parseNodeArg(node);
|
||||
return await client.patchNode(pageId, nodeId, parsedNode);
|
||||
},
|
||||
),
|
||||
|
||||
// Shared registry schema + description; execute retains parseNodeArg on the
|
||||
// incoming node (parity with the standalone MCP server, index.ts
|
||||
// insert_node).
|
||||
insertNode: sharedTool(
|
||||
sharedToolSpecs.insertNode,
|
||||
async ({ pageId, node, position, anchorNodeId, anchorText }) => {
|
||||
const parsedNode = parseNodeArg(node);
|
||||
return await client.insertNode(pageId, parsedNode, {
|
||||
position,
|
||||
anchorNodeId,
|
||||
anchorText,
|
||||
});
|
||||
},
|
||||
),
|
||||
|
||||
deleteNode: sharedTool(
|
||||
sharedToolSpecs.deleteNode,
|
||||
async ({ pageId, nodeId }) => await client.deleteNode(pageId, nodeId),
|
||||
),
|
||||
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
// The execute body keeps this layer's content normalization (parity with
|
||||
// the standalone MCP server, index.ts update_page_json).
|
||||
updatePageJson: sharedTool(
|
||||
sharedToolSpecs.updatePageJson,
|
||||
async ({ pageId, content, title }) => {
|
||||
// undefined/null pass through as undefined (title-only / no-op); any
|
||||
// string is JSON.parsed (so an empty string "" throws, matching the
|
||||
// MCP server); an object is passed through unchanged.
|
||||
let doc;
|
||||
if (content === undefined || content === null) {
|
||||
doc = undefined;
|
||||
} else {
|
||||
// String -> JSON.parse (throwing on invalid); object passes through.
|
||||
doc = parseNodeArg(content, 'content was a string but not valid JSON');
|
||||
}
|
||||
return await client.updatePageJson(pageId, doc, title);
|
||||
},
|
||||
),
|
||||
|
||||
// Schema + description live in @docmost/mcp's SHARED_TOOL_SPECS (#410).
|
||||
// Promoted from MCP-only so the in-app agent can attach a REAL footnote to
|
||||
// already-written text instead of leaving a literal `^[...]` string.
|
||||
insertFootnote: sharedTool(
|
||||
sharedToolSpecs.insertFootnote,
|
||||
async ({ pageId, anchorText, text }) =>
|
||||
await client.insertFootnote(pageId, anchorText, text),
|
||||
),
|
||||
|
||||
// Schema + description live in @docmost/mcp's SHARED_TOOL_SPECS (#410).
|
||||
// The schema field is `imageUrl`; the client method takes it positionally.
|
||||
insertImage: sharedTool(
|
||||
sharedToolSpecs.insertImage,
|
||||
async ({ pageId, imageUrl, align, alt, replaceText, afterText }) =>
|
||||
await client.insertImage(pageId, imageUrl, {
|
||||
align,
|
||||
alt,
|
||||
replaceText,
|
||||
afterText,
|
||||
}),
|
||||
),
|
||||
|
||||
// Schema + description live in @docmost/mcp's SHARED_TOOL_SPECS (#410).
|
||||
replaceImage: sharedTool(
|
||||
sharedToolSpecs.replaceImage,
|
||||
async ({ pageId, attachmentId, imageUrl, align, alt }) =>
|
||||
await client.replaceImage(pageId, attachmentId, imageUrl, {
|
||||
align,
|
||||
alt,
|
||||
}),
|
||||
),
|
||||
|
||||
// Schema + description live in @docmost/mcp's SHARED_TOOL_SPECS (#423).
|
||||
// meta.hash in the result is the baseHash drawioUpdate requires.
|
||||
drawioGet: sharedTool(
|
||||
sharedToolSpecs.drawioGet,
|
||||
async ({ pageId, node, format }) =>
|
||||
await client.drawioGet(pageId, node, format ?? 'xml'),
|
||||
),
|
||||
|
||||
// Schema + description live in @docmost/mcp's SHARED_TOOL_SPECS (#423).
|
||||
// The flat schema fields are regrouped into the client's `where` object.
|
||||
drawioCreate: sharedTool(
|
||||
sharedToolSpecs.drawioCreate,
|
||||
async ({ pageId, xml, position, anchorNodeId, anchorText, title }) =>
|
||||
await client.drawioCreate(
|
||||
pageId,
|
||||
{ position, anchorNodeId, anchorText },
|
||||
xml,
|
||||
title,
|
||||
),
|
||||
),
|
||||
|
||||
// Schema + description live in @docmost/mcp's SHARED_TOOL_SPECS (#423).
|
||||
// baseHash is the optimistic lock: mismatch => structured conflict error.
|
||||
drawioUpdate: sharedTool(
|
||||
sharedToolSpecs.drawioUpdate,
|
||||
async ({ pageId, node, xml, baseHash }) =>
|
||||
await client.drawioUpdate(pageId, node, xml, baseHash),
|
||||
),
|
||||
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
// The table reference parameter was unified to `table` (was `tableRef`).
|
||||
tableInsertRow: sharedTool(
|
||||
sharedToolSpecs.tableInsertRow,
|
||||
async ({ pageId, table, cells, index }) =>
|
||||
await client.tableInsertRow(pageId, table, cells, index),
|
||||
),
|
||||
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
tableDeleteRow: sharedTool(
|
||||
sharedToolSpecs.tableDeleteRow,
|
||||
async ({ pageId, table, index }) =>
|
||||
await client.tableDeleteRow(pageId, table, index),
|
||||
),
|
||||
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
tableUpdateCell: sharedTool(
|
||||
sharedToolSpecs.tableUpdateCell,
|
||||
async ({ pageId, table, row, col, text }) =>
|
||||
await client.tableUpdateCell(pageId, table, row, col, text),
|
||||
),
|
||||
|
||||
copyPageContent: sharedTool(
|
||||
sharedToolSpecs.copyPageContent,
|
||||
async ({ sourcePageId, targetPageId }) =>
|
||||
await client.copyPageContent(sourcePageId, targetPageId),
|
||||
),
|
||||
|
||||
importPageMarkdown: sharedTool(
|
||||
sharedToolSpecs.importPageMarkdown,
|
||||
async ({ pageId, markdown }) =>
|
||||
await client.importPageMarkdown(pageId, markdown),
|
||||
),
|
||||
|
||||
// Schema + description now live in @docmost/mcp's SHARED_TOOL_SPECS (#294).
|
||||
// Both layers already carried the security-confirmation framing, so there
|
||||
// was no real divergence to preserve — only wording drift.
|
||||
sharePage: sharedTool(
|
||||
sharedToolSpecs.sharePage,
|
||||
async ({ pageId, searchIndexing }) =>
|
||||
await client.sharePage(pageId, searchIndexing),
|
||||
),
|
||||
|
||||
unsharePage: sharedTool(
|
||||
sharedToolSpecs.unsharePage,
|
||||
async ({ pageId }) => await client.unsharePage(pageId),
|
||||
),
|
||||
|
||||
restorePageVersion: sharedTool(
|
||||
sharedToolSpecs.restorePageVersion,
|
||||
async ({ historyId }) => await client.restorePageVersion(historyId),
|
||||
),
|
||||
|
||||
// INTENTIONAL per-transport divergence (not shared): deliberately omits the
|
||||
// `deleteComments` schema field (comment-deletion guardrail) and carries a
|
||||
// much shorter description; the standalone MCP `docmost_transform` exposes
|
||||
// much shorter description; the standalone MCP `docmostTransform` exposes
|
||||
// the full helper catalogue. Different schema, so kept per-layer.
|
||||
transformPage: tool({
|
||||
description:
|
||||
@@ -935,6 +700,51 @@ export class AiChatToolsService {
|
||||
}),
|
||||
};
|
||||
|
||||
// Add EVERY shared tool from the zod-agnostic registry in one loop (#445).
|
||||
// The spec owns the canonical arg->client mapping; this host only decides
|
||||
// WHICH mapping to run and returns its value directly (no envelope). For each
|
||||
// spec:
|
||||
// - skip `mcpOnly` specs (they belong to the standalone MCP host only);
|
||||
// - skip `inlineBothHosts` specs (drawioShapes / drawioGuide): they carry
|
||||
// no execute and are wired INLINE just below, calling the pure helpers;
|
||||
// - use `inAppExecute` when the spec declares a DELIBERATE per-layer
|
||||
// difference (a projected result shape, a different guardrail message);
|
||||
// - otherwise use the canonical `execute` (raw client result, identical to
|
||||
// the MCP host's before it wraps it as JSON).
|
||||
// The execute receives the AI-SDK-validated, type-erased input; the spec reads
|
||||
// the same fields its buildShape declares. This is the SINGLE place the in-app
|
||||
// arg mapping lives — it can no longer silently drift from the MCP host.
|
||||
for (const spec of Object.values(sharedToolSpecs)) {
|
||||
if (spec.mcpOnly) continue;
|
||||
if (spec.inlineBothHosts) continue;
|
||||
const run = spec.inAppExecute ?? spec.execute;
|
||||
if (!run) continue; // defensive: a shared spec always carries one of them.
|
||||
tools[spec.inAppKey] = sharedTool(
|
||||
spec,
|
||||
(async (args) =>
|
||||
run(client, args as Record<string, unknown>)) as Tool['execute'],
|
||||
);
|
||||
}
|
||||
|
||||
// drawioShapes / drawioGuide (#424): `inlineBothHosts` registry specs wired
|
||||
// here with the SAME schema+description the shared spec pins, but calling the
|
||||
// pure searchShapes / getGuideSection helpers off the loaded @docmost/mcp
|
||||
// module — they are not client methods and their catalog loader uses
|
||||
// import.meta, so they cannot live in the zod-agnostic shared execute. The raw
|
||||
// result is identical to the MCP host's (which wraps it as JSON text); here
|
||||
// the in-app host returns it plain, exactly like every other shared tool.
|
||||
tools[sharedToolSpecs.drawioShapes.inAppKey] = sharedTool(
|
||||
sharedToolSpecs.drawioShapes,
|
||||
async ({ query, category, limit }) => {
|
||||
const results = searchShapes(query, { category, limit });
|
||||
return { query, count: results.length, results };
|
||||
},
|
||||
);
|
||||
tools[sharedToolSpecs.drawioGuide.inAppKey] = sharedTool(
|
||||
sharedToolSpecs.drawioGuide,
|
||||
async ({ section }) => getGuideSection(section),
|
||||
);
|
||||
|
||||
// Passive "new comments: N" signal (#417). PER-TURN state (forUser runs once
|
||||
// per turn), so the watermark starts now and only comments a human leaves
|
||||
// WHILE this turn runs are signalled — exactly the mid-turn loop; between-turn
|
||||
@@ -948,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) => {
|
||||
@@ -977,7 +795,11 @@ export class AiChatToolsService {
|
||||
},
|
||||
});
|
||||
|
||||
return wrapToolsWithCommentSignal(tools, tracker);
|
||||
return wrapInAppToolsWithCap(
|
||||
wrapToolsWithCommentSignal(tools, tracker),
|
||||
client,
|
||||
capMs,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -283,10 +283,18 @@ describe('AiChatToolsService forUser + comment signal (real tracker)', () => {
|
||||
DocmostClient: function () {
|
||||
return fakeClient as DocmostClientLike;
|
||||
} as unknown as loader.DocmostClientCtor,
|
||||
sharedToolSpecs: SHARED_TOOL_SPECS as Record<string, loader.SharedToolSpec>,
|
||||
sharedToolSpecs: SHARED_TOOL_SPECS as unknown as Record<string, loader.SharedToolSpec>,
|
||||
// Wire the REAL factory so the in-app path is exercised end to end.
|
||||
createCommentSignalTracker:
|
||||
createCommentSignalTracker as unknown as loader.CommentSignalTrackerFactory,
|
||||
// Pure no-network draw.io helpers (#424) — required on the loader return;
|
||||
// this comment-signal test doesn't exercise them, so no-op stubs suffice.
|
||||
searchShapes: (() => []) as unknown as loader.SearchShapesFn,
|
||||
getGuideSection: (() => ({
|
||||
section: '',
|
||||
content: '',
|
||||
sections: [],
|
||||
})) as unknown as loader.GetGuideSectionFn,
|
||||
});
|
||||
return new AiChatToolsService(
|
||||
tokenServiceStub as never,
|
||||
|
||||
@@ -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';
|
||||
|
||||
@@ -23,6 +23,8 @@ type DocmostClientMethod =
|
||||
| 'getWorkspace'
|
||||
| 'getSpaces'
|
||||
| 'listPages'
|
||||
| 'getTree'
|
||||
| 'getPageContext'
|
||||
| 'listSidebarPages'
|
||||
| 'getOutline'
|
||||
| 'getPageJson'
|
||||
@@ -62,10 +64,17 @@ type DocmostClientMethod =
|
||||
| 'insertImage'
|
||||
| 'replaceImage'
|
||||
| 'insertFootnote'
|
||||
// --- draw.io diagrams (#423, stage 1) ---
|
||||
// --- draw.io diagrams (#423 stage 1, #424 stage 2) ---
|
||||
// DERIVED from the real DocmostClient (#446): drawioCreate/drawioUpdate carry
|
||||
// the optional layout:"elk" 5th arg in the real signature, so the layout parity
|
||||
// (#440) is inherited automatically — no hand-written mirror to keep in sync.
|
||||
| 'drawioGet'
|
||||
| 'drawioCreate'
|
||||
| 'drawioUpdate'
|
||||
// --- draw.io high-level semantic tools (#425 stage 3) ---
|
||||
| 'drawioEditCells'
|
||||
| 'drawioFromGraph'
|
||||
| 'drawioFromMermaid'
|
||||
// --- write (comment) ---
|
||||
| 'createComment'
|
||||
| 'resolveComment';
|
||||
@@ -141,6 +150,19 @@ export type CommentSignalTrackerFactory = (options: {
|
||||
debounceMs?: number;
|
||||
}) => CommentSignalTrackerLike;
|
||||
|
||||
// Pure, no-network draw.io helpers (#424). These are plain functions on the
|
||||
// module (NOT DocmostClient methods) — the in-app AI-SDK service calls them
|
||||
// directly to wire drawioShapes / drawioGuide, mirroring the MCP server.
|
||||
export type SearchShapesFn = (
|
||||
query: string,
|
||||
opts?: { category?: string; limit?: number },
|
||||
) => Array<Record<string, unknown>>;
|
||||
export type GetGuideSectionFn = (section?: string) => {
|
||||
section: string;
|
||||
content: string;
|
||||
sections: string[];
|
||||
};
|
||||
|
||||
interface DocmostMcpModule {
|
||||
DocmostClient: DocmostClientCtor;
|
||||
SHARED_TOOL_SPECS: Record<string, SharedToolSpec>;
|
||||
@@ -153,6 +175,15 @@ interface DocmostMcpModule {
|
||||
// the mocked loader in unit tests) — the stale-check below is a NO-OP when it
|
||||
// is missing, so an older build never wrongly fails startup.
|
||||
REGISTRY_STAMP?: string;
|
||||
// Pure, no-network draw.io helpers (#424) backing drawioShapes / drawioGuide.
|
||||
// Those two specs are `inlineBothHosts` (they stay in SHARED_TOOL_SPECS for the
|
||||
// shared contract but carry no execute — their catalog loader uses import.meta
|
||||
// and can't be value-imported into the zod-agnostic tool-specs.ts), so the
|
||||
// in-app service wires them INLINE off these helpers, mirroring the standalone
|
||||
// MCP host. Exposed off the loaded module so the service and its test mocks can
|
||||
// reach them.
|
||||
searchShapes: SearchShapesFn;
|
||||
getGuideSection: GetGuideSectionFn;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -160,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).
|
||||
@@ -194,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
|
||||
@@ -216,6 +284,8 @@ export async function loadDocmostMcp(): Promise<{
|
||||
DocmostClient: DocmostClientCtor;
|
||||
sharedToolSpecs: Record<string, SharedToolSpec>;
|
||||
createCommentSignalTracker?: CommentSignalTrackerFactory;
|
||||
searchShapes: SearchShapesFn;
|
||||
getGuideSection: GetGuideSectionFn;
|
||||
}> {
|
||||
if (!modulePromise) {
|
||||
modulePromise = (async () => {
|
||||
@@ -261,5 +331,8 @@ export async function loadDocmostMcp(): Promise<{
|
||||
// Optional: forwarded when present so the in-app layer can build the passive
|
||||
// comment signal (#417); undefined on a stale build => signal disabled.
|
||||
createCommentSignalTracker: mod.createCommentSignalTracker,
|
||||
// Pure no-network draw.io helpers (#424); not client methods.
|
||||
searchShapes: mod.searchShapes,
|
||||
getGuideSection: mod.getGuideSection,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,8 +17,10 @@ import { SHARED_TOOL_SPECS } from '../../../../../../packages/mcp/src/tool-specs
|
||||
* This test fails the build if a spec is added to the registry but never wired
|
||||
* in-app, if an `inAppKey` is renamed without updating the service, if the
|
||||
* description drifts between the registry and the exposed tool, if the
|
||||
* snake_case `mcpName` <-> camelCase `inAppKey` convention is broken, or if the
|
||||
* exposed tool's input-schema keys diverge from the spec's `buildShape`.
|
||||
* `mcpName === inAppKey` convention is broken (issue #412 unified the external
|
||||
* MCP tool name with the in-app key — both are the same camelCase identifier),
|
||||
* or if the exposed tool's input-schema keys diverge from the spec's
|
||||
* `buildShape`.
|
||||
*
|
||||
* It does NOT need @docmost/mcp built: the registry is imported from TS source,
|
||||
* and the ESM loader is mocked so `forUser()` never dynamically imports the
|
||||
@@ -45,6 +47,16 @@ describe('SHARED_TOOL_SPECS contract parity', () => {
|
||||
string,
|
||||
loader.SharedToolSpec
|
||||
>,
|
||||
// Pure no-network draw.io helpers (#424). The contract test never executes
|
||||
// a tool body, so type-correct stubs suffice (the real functions can't be
|
||||
// imported here — drawio-shapes.ts uses import.meta, incompatible with the
|
||||
// CommonJS jest transform).
|
||||
searchShapes: (() => []) as unknown as loader.SearchShapesFn,
|
||||
getGuideSection: (() => ({
|
||||
section: 'index',
|
||||
content: '',
|
||||
sections: [],
|
||||
})) as unknown as loader.GetGuideSectionFn,
|
||||
});
|
||||
const service = new AiChatToolsService(
|
||||
tokenServiceStub as never,
|
||||
@@ -64,13 +76,9 @@ describe('SHARED_TOOL_SPECS contract parity', () => {
|
||||
|
||||
afterAll(() => jest.restoreAllMocks());
|
||||
|
||||
// camelCase -> snake_case, matching the registry's mcpName convention.
|
||||
const toSnake = (s: string) =>
|
||||
s.replace(/[A-Z]/g, (c) => `_${c.toLowerCase()}`);
|
||||
|
||||
// Type as the (optional-buildShape) SharedToolSpec; the `satisfies` literal
|
||||
// above otherwise narrows to a union where some members lack buildShape.
|
||||
const specEntries = Object.entries(SHARED_TOOL_SPECS) as Array<
|
||||
const specEntries = Object.entries(SHARED_TOOL_SPECS) as unknown as Array<
|
||||
[string, loader.SharedToolSpec]
|
||||
>;
|
||||
|
||||
@@ -86,8 +94,8 @@ describe('SHARED_TOOL_SPECS contract parity', () => {
|
||||
expect(spec.inAppKey).toBe(registryKey);
|
||||
});
|
||||
|
||||
it('mcpName is the snake_case form of inAppKey', () => {
|
||||
expect(spec.mcpName).toBe(toSnake(spec.inAppKey));
|
||||
it('mcpName equals inAppKey (unified camelCase name, #412)', () => {
|
||||
expect(spec.mcpName).toBe(spec.inAppKey);
|
||||
});
|
||||
|
||||
it('is exposed in-app under its inAppKey', () => {
|
||||
|
||||
@@ -27,16 +27,18 @@ import type { DocmostClientLike } from './docmost-client.loader';
|
||||
*/
|
||||
|
||||
describe('tool tier metadata (#332)', () => {
|
||||
it('core set is the documented 13 + searchInPage + insertFootnote (15)', () => {
|
||||
expect(CORE_TOOL_KEYS).toHaveLength(15);
|
||||
it('core set is the documented 13 + searchInPage + insertFootnote + getTree + getPageContext (17, #443)', () => {
|
||||
expect(CORE_TOOL_KEYS).toHaveLength(17);
|
||||
expect(CORE_TOOL_SET.has('searchInPage')).toBe(true); // #330, promoted to core
|
||||
expect(CORE_TOOL_SET.has('insertFootnote')).toBe(true); // #410, promoted to core
|
||||
expect(CORE_TOOL_SET.has('getTree')).toBe(true); // #443, promoted to core
|
||||
expect(CORE_TOOL_SET.has('getPageContext')).toBe(true); // #443, promoted to core
|
||||
// loadTools is a meta-tool, not a normal core key.
|
||||
expect(CORE_TOOL_SET.has(LOAD_TOOLS_NAME)).toBe(false);
|
||||
});
|
||||
|
||||
it('#410 image tools are DEFERRED, footnote tool is CORE', () => {
|
||||
// insert_footnote is core (symmetric with editPageText); the image tools stay
|
||||
// insertFootnote is core (symmetric with editPageText); the image tools stay
|
||||
// deferred (rare, fat — loaded on demand). Assert both the spec tier and the
|
||||
// CORE_TOOL_SET membership so a future tier edit that desyncs them fails here.
|
||||
expect(SHARED_TOOL_SPECS.insertFootnote.tier).toBe('core');
|
||||
@@ -123,7 +125,14 @@ describe('deferred catalog ↔ live forUser() toolset partition (#332, F3)', ()
|
||||
DocmostClient: function () {
|
||||
return {} as DocmostClientLike;
|
||||
} as unknown as loader.DocmostClientCtor,
|
||||
sharedToolSpecs: SHARED_TOOL_SPECS as Record<string, loader.SharedToolSpec>,
|
||||
sharedToolSpecs: SHARED_TOOL_SPECS as unknown as Record<string, loader.SharedToolSpec>,
|
||||
// Pure no-network draw.io helpers (#424); tool bodies are never executed here.
|
||||
searchShapes: (() => []) as unknown as loader.SearchShapesFn,
|
||||
getGuideSection: (() => ({
|
||||
section: 'index',
|
||||
content: '',
|
||||
sections: [],
|
||||
})) as unknown as loader.GetGuideSectionFn,
|
||||
});
|
||||
const service = new AiChatToolsService(
|
||||
{
|
||||
@@ -232,6 +241,16 @@ describe('applyLoadTools (#332)', () => {
|
||||
expect(LOAD_TOOLS_DESCRIPTION).toContain('only ACTIVATES them');
|
||||
expect(LOAD_TOOLS_DESCRIPTION).toContain('callable on your NEXT step');
|
||||
});
|
||||
|
||||
it('loadTools description tells the model CORE tools are always active (#444)', () => {
|
||||
expect(LOAD_TOOLS_DESCRIPTION).toContain(
|
||||
'Tools NOT listed in the catalog are CORE and ALWAYS active',
|
||||
);
|
||||
expect(LOAD_TOOLS_DESCRIPTION).toContain('NEVER via loadTools');
|
||||
// Names it out explicitly so the model doesn't loadTools a core tool.
|
||||
expect(LOAD_TOOLS_DESCRIPTION).toContain('createComment');
|
||||
expect(LOAD_TOOLS_DESCRIPTION).toContain('searchInPage');
|
||||
});
|
||||
});
|
||||
|
||||
describe('editorial "Corrector" scenario is fully served by CORE (#332)', () => {
|
||||
|
||||
@@ -39,12 +39,14 @@ export interface ToolCatalogEntry {
|
||||
|
||||
/**
|
||||
* CORE (always-active) in-app tool keys — 13 frequent/tiny tools + `searchInPage`
|
||||
* (#330) + `insertFootnote` (#410). `searchInPage` is core because it is frequent
|
||||
* for the editorial roles this feature targets; `insertFootnote` is core so the
|
||||
* footnote tool is NOT hidden while its natural sibling `editPageText` is always
|
||||
* active (that asymmetry is exactly what pushed the agent to write literal
|
||||
* `^[...]`). `loadTools` is active too but is not a normal tool key (it is added
|
||||
* to activeTools separately).
|
||||
* (#330) + `insertFootnote` (#410) + `getTree`/`getPageContext` (#443).
|
||||
* `searchInPage` is core because it is frequent for the editorial roles this
|
||||
* feature targets; `insertFootnote` is core so the footnote tool is NOT hidden
|
||||
* while its natural sibling `editPageText` is always active (that asymmetry is
|
||||
* exactly what pushed the agent to write literal `^[...]`). `getTree` and
|
||||
* `getPageContext` are the single-call navigation/lookup tools — core so the
|
||||
* agent never has to loadTools just to orient itself. `loadTools` is active too
|
||||
* but is not a normal tool key (it is added to activeTools separately).
|
||||
*/
|
||||
export const CORE_TOOL_KEYS = [
|
||||
'searchPages',
|
||||
@@ -60,12 +62,17 @@ export const CORE_TOOL_KEYS = [
|
||||
'listComments',
|
||||
'resolveComment',
|
||||
'editPageText',
|
||||
// #330 search_in_page — frequent for editorial sweeps; core despite predating
|
||||
// #330 searchInPage — frequent for editorial sweeps; core despite predating
|
||||
// the issue's tier list.
|
||||
'searchInPage',
|
||||
// #410 insert_footnote — core so pinpoint citations to already-written text
|
||||
// #410 insertFootnote — core so pinpoint citations to already-written text
|
||||
// don't degrade into literal `^[...]`; kept symmetric with editPageText.
|
||||
'insertFootnote',
|
||||
// #443 getTree + getPageContext — cheap single-call navigation/lookup tools
|
||||
// (the core listPages even points to getTree); core so the agent never has
|
||||
// to loadTools just to orient itself.
|
||||
'getTree',
|
||||
'getPageContext',
|
||||
] as const;
|
||||
|
||||
/** O(1) membership test for the core tier. */
|
||||
@@ -84,7 +91,10 @@ export const LOAD_TOOLS_DESCRIPTION =
|
||||
'block in your instructions. Pass the EXACT tool names from the catalog; this\n' +
|
||||
'call only ACTIVATES them and returns { loaded: [...] } — the tools become\n' +
|
||||
'callable on your NEXT step. Load several names in one call when the task clearly\n' +
|
||||
'needs them. Unknown names are rejected with the list of valid ones.';
|
||||
'needs them. Unknown names are rejected with the list of valid ones.\n' +
|
||||
'Tools NOT listed in the catalog are CORE and ALWAYS active — call them directly,\n' +
|
||||
'NEVER via loadTools (e.g. createComment, listComments, resolveComment,\n' +
|
||||
'editPageText, searchInPage).';
|
||||
|
||||
/**
|
||||
* Tier + catalogLine for the INLINE ai-chat tools — those defined per-layer in
|
||||
@@ -121,12 +131,9 @@ export const INLINE_TOOL_TIERS: Record<
|
||||
// --- deferred inline ---
|
||||
// NOTE: createPage, renamePage, movePage, deletePage, updatePageJson and
|
||||
// exportPageMarkdown moved to @docmost/mcp's SHARED_TOOL_SPECS (#294); they
|
||||
// carry their own deferred tier + catalogLine there.
|
||||
updatePageContent: {
|
||||
tier: 'deferred',
|
||||
catalogLine:
|
||||
"updatePageContent — replace a page's body (and optionally title) with new Markdown.",
|
||||
},
|
||||
// carry their own deferred tier + catalogLine there. updatePageContent moved
|
||||
// there too as updatePageMarkdown (#411) — a shared registry spec now, so it
|
||||
// is no longer an inline tier entry.
|
||||
listSidebarPages: {
|
||||
tier: 'deferred',
|
||||
catalogLine:
|
||||
@@ -138,7 +145,7 @@ export const INLINE_TOOL_TIERS: Record<
|
||||
},
|
||||
// NOTE: tableInsertRow, tableDeleteRow and tableUpdateCell moved to
|
||||
// @docmost/mcp's SHARED_TOOL_SPECS (#294); they carry their own deferred tier +
|
||||
// catalogLine there. getTable stays inline (its MCP name table_get breaks the
|
||||
// catalogLine there. getTable stays inline (its MCP name tableGet breaks the
|
||||
// snake_case(inAppKey) convention, so it has no shared spec).
|
||||
// NOTE: checkNewComments moved to @docmost/mcp's SHARED_TOOL_SPECS (#294);
|
||||
// it carries its own deferred tier + catalogLine there.
|
||||
@@ -150,7 +157,7 @@ export const INLINE_TOOL_TIERS: Record<
|
||||
// NOTE: sharePage moved to @docmost/mcp's SHARED_TOOL_SPECS (#294); it carries
|
||||
// its own deferred tier + catalogLine there. transformPage stays inline (its
|
||||
// schema deliberately diverges — it omits the deleteComments field the MCP
|
||||
// docmost_transform exposes, a comment-deletion guardrail).
|
||||
// docmostTransform exposes, a comment-deletion guardrail).
|
||||
transformPage: {
|
||||
tier: 'deferred',
|
||||
catalogLine: "transformPage — run a sandboxed JS transform over a page's document.",
|
||||
|
||||
@@ -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');
|
||||
}
|
||||
}
|
||||
|
||||
@@ -12,3 +12,22 @@ export class SearchResponseDto {
|
||||
updatedAt: Date;
|
||||
space: Partial<Space>;
|
||||
}
|
||||
|
||||
// Response shape for the opt-in agent-lookup mode (#443, `substring: true`).
|
||||
// Additive to the FTS response: carries the location (`path`), a windowed
|
||||
// `snippet` around the first match and a per-response sort `score`. The MCP
|
||||
// layer maps `id → pageId`; `slugId` is never exposed.
|
||||
export class SearchLookupResponseDto {
|
||||
id: string;
|
||||
slugId: string;
|
||||
title: string;
|
||||
parentPageId: string | null;
|
||||
// Ancestor titles from the space root down to the direct parent; [] for a
|
||||
// root page.
|
||||
path: string[];
|
||||
// ~300–500 chars around the first match (or a leading text window / extended
|
||||
// ts_headline fallback).
|
||||
snippet: string;
|
||||
// 0..1 float, meaningful ONLY for sorting within one response.
|
||||
score: number;
|
||||
}
|
||||
|
||||
@@ -30,6 +30,31 @@ export class SearchDTO {
|
||||
@IsOptional()
|
||||
@IsNumber()
|
||||
offset?: number;
|
||||
|
||||
// --- Opt-in agent-lookup mode (#443). ------------------------------------
|
||||
// These fields are ADDITIVE and default-off: a web client that sends none of
|
||||
// them gets byte-identical FTS behaviour and result shape. They are only read
|
||||
// by the substring/path/snippet code path in SearchService.searchPage.
|
||||
//
|
||||
// NOTE (standalone stdio vs stock upstream): stock upstream validates this DTO
|
||||
// with `whitelist: true`, so an older server silently strips these unknown
|
||||
// fields and the request degrades gracefully to the plain FTS behaviour.
|
||||
|
||||
// Enables the hybrid substring branch (title + text_content LIKE) merged with
|
||||
// the existing FTS branch, plus tiered ranking, path and windowed snippet.
|
||||
@IsOptional()
|
||||
@IsBoolean()
|
||||
substring?: boolean;
|
||||
|
||||
// Restrict the search to a page and all of its descendants (inclusive).
|
||||
@IsOptional()
|
||||
@IsString()
|
||||
parentPageId?: string;
|
||||
|
||||
// Match titles only; do not scan text_content.
|
||||
@IsOptional()
|
||||
@IsBoolean()
|
||||
titleOnly?: boolean;
|
||||
}
|
||||
|
||||
export class SearchShareDTO extends SearchDTO {
|
||||
|
||||
@@ -60,6 +60,12 @@ export class SearchController {
|
||||
}
|
||||
}
|
||||
|
||||
// #443 graceful degradation: on EE/Typesense instances the request routes to
|
||||
// the Typesense backend, which does NOT implement the opt-in agent-lookup
|
||||
// mode. The `substring`/`parentPageId`/`titleOnly` fields are silently ignored
|
||||
// and the response carries no `path`/`snippet`/`score` and no substring/tier
|
||||
// ranking — it degrades to plain Typesense FTS. The native lookup mode below
|
||||
// is Postgres-search-driver only.
|
||||
if (this.environmentService.getSearchDriver() === 'typesense') {
|
||||
return this.searchTypesense(searchDto, {
|
||||
userId: user.id,
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
import {
|
||||
computeLookupScore,
|
||||
escapeLikePattern,
|
||||
SearchLookupTier,
|
||||
} from './search.service';
|
||||
|
||||
/**
|
||||
* Pure-function coverage for the #443 agent-lookup helpers:
|
||||
* - escapeLikePattern: LIKE-metacharacter escaping so `%`/`_`/`\` are literals
|
||||
* (the acceptance-table requirement that a query of `%` or `_` does NOT match
|
||||
* everything);
|
||||
* - computeLookupScore: the tiered 0..1 ranking score, where a stronger tier
|
||||
* always outranks a weaker one regardless of the in-tier secondary signal.
|
||||
*
|
||||
* The DB-touching branch (substring UNION FTS, path CTE, snippet window) is
|
||||
* covered by the integration spec against the real schema.
|
||||
*/
|
||||
describe('escapeLikePattern', () => {
|
||||
it('escapes the LIKE metacharacters % _ and \\', () => {
|
||||
expect(escapeLikePattern('%')).toBe('\\%');
|
||||
expect(escapeLikePattern('_')).toBe('\\_');
|
||||
expect(escapeLikePattern('\\')).toBe('\\\\');
|
||||
});
|
||||
|
||||
it('escapes the backslash FIRST so it does not double-escape %/_', () => {
|
||||
// Input `\%` must become `\\` + `\%` = `\\\%`, not `\\%`.
|
||||
expect(escapeLikePattern('\\%')).toBe('\\\\\\%');
|
||||
});
|
||||
|
||||
it('leaves ordinary technical chars (. - / digits) untouched', () => {
|
||||
expect(escapeLikePattern('backup-srv.local')).toBe('backup-srv.local');
|
||||
expect(escapeLikePattern('10.0.12')).toBe('10.0.12');
|
||||
expect(escapeLikePattern('WB-MGE-30D86B')).toBe('WB-MGE-30D86B');
|
||||
expect(escapeLikePattern('a/b')).toBe('a/b');
|
||||
});
|
||||
|
||||
it('escapes only the metacharacters in a mixed string', () => {
|
||||
expect(escapeLikePattern('50%_off.zip')).toBe('50\\%\\_off.zip');
|
||||
});
|
||||
|
||||
it('is null/undefined-safe', () => {
|
||||
expect(escapeLikePattern(undefined as any)).toBe('');
|
||||
expect(escapeLikePattern(null as any)).toBe('');
|
||||
});
|
||||
});
|
||||
|
||||
describe('computeLookupScore', () => {
|
||||
it('keeps every score within (0, 1]', () => {
|
||||
for (const tier of [
|
||||
SearchLookupTier.TITLE_EXACT,
|
||||
SearchLookupTier.TITLE_SUBSTRING,
|
||||
SearchLookupTier.TEXT,
|
||||
]) {
|
||||
for (const secondary of [0, 0.001, 1, 100, 1e6]) {
|
||||
const s = computeLookupScore({ tier, secondary });
|
||||
expect(s).toBeGreaterThan(0);
|
||||
expect(s).toBeLessThanOrEqual(1);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
it('a stronger tier ALWAYS outranks a weaker tier, whatever the secondary', () => {
|
||||
// Weak tier with a huge secondary must still lose to a strong tier with a
|
||||
// tiny secondary — tiers dominate.
|
||||
const strongLowSecondary = computeLookupScore({
|
||||
tier: SearchLookupTier.TITLE_EXACT,
|
||||
secondary: 0,
|
||||
});
|
||||
const weakHighSecondary = computeLookupScore({
|
||||
tier: SearchLookupTier.TEXT,
|
||||
secondary: 1e9,
|
||||
});
|
||||
expect(strongLowSecondary).toBeGreaterThan(weakHighSecondary);
|
||||
});
|
||||
|
||||
it('within a tier a larger secondary sorts higher', () => {
|
||||
const lo = computeLookupScore({
|
||||
tier: SearchLookupTier.TEXT,
|
||||
secondary: 0.1,
|
||||
});
|
||||
const hi = computeLookupScore({
|
||||
tier: SearchLookupTier.TEXT,
|
||||
secondary: 5,
|
||||
});
|
||||
expect(hi).toBeGreaterThan(lo);
|
||||
});
|
||||
|
||||
it('treats a negative/absent secondary as 0', () => {
|
||||
const zero = computeLookupScore({ tier: SearchLookupTier.TEXT, secondary: 0 });
|
||||
expect(computeLookupScore({ tier: SearchLookupTier.TEXT })).toBe(zero);
|
||||
expect(
|
||||
computeLookupScore({ tier: SearchLookupTier.TEXT, secondary: -5 }),
|
||||
).toBe(zero);
|
||||
});
|
||||
});
|
||||
@@ -1,6 +1,9 @@
|
||||
import { Injectable } from '@nestjs/common';
|
||||
import { SearchDTO, SearchSuggestionDTO } from './dto/search.dto';
|
||||
import { SearchResponseDto } from './dto/search-response.dto';
|
||||
import {
|
||||
SearchLookupResponseDto,
|
||||
SearchResponseDto,
|
||||
} from './dto/search-response.dto';
|
||||
import { InjectKysely } from 'nestjs-kysely';
|
||||
import { KyselyDB } from '@docmost/db/types/kysely.types';
|
||||
import { sql } from 'kysely';
|
||||
@@ -34,6 +37,53 @@ export function buildTsQuery(raw: string): string {
|
||||
return tsquery(cleaned + '*');
|
||||
}
|
||||
|
||||
// Escape the LIKE metacharacters (`%`, `_`, `\`) in a raw user query so every
|
||||
// character — including `.`, `-`, `_`, `%`, `/` — is matched LITERALLY by a
|
||||
// `col LIKE '%' || q || '%'` predicate. Without this, a query of `%` or `_`
|
||||
// would match every row (see the #443 acceptance table). The backslash is the
|
||||
// escape char (Postgres LIKE default), so it must be escaped first.
|
||||
export function escapeLikePattern(raw: string): string {
|
||||
return (raw ?? '')
|
||||
.replace(/\\/g, '\\\\')
|
||||
.replace(/%/g, '\\%')
|
||||
.replace(/_/g, '\\_');
|
||||
}
|
||||
|
||||
// Ranking tiers for the agent-lookup mode (#443), highest first. A hit's tier
|
||||
// is the strongest way it matched; ties inside a tier break on a secondary
|
||||
// signal (FTS rank, or first-match position). The numeric `score` returned to
|
||||
// the caller is derived from (tier, secondary) and is meaningful ONLY for
|
||||
// ordering within a single response.
|
||||
export enum SearchLookupTier {
|
||||
// Title equals the query, case-insensitively.
|
||||
TITLE_EXACT = 3,
|
||||
// Query is a substring of the title.
|
||||
TITLE_SUBSTRING = 2,
|
||||
// Query matched in the text (substring or FTS).
|
||||
TEXT = 1,
|
||||
}
|
||||
|
||||
export interface RankableHit {
|
||||
tier: SearchLookupTier;
|
||||
// Secondary in-tier signal, higher = better (e.g. ts_rank, or a
|
||||
// position-derived closeness score). Defaults to 0.
|
||||
secondary?: number;
|
||||
}
|
||||
|
||||
// Map (tier, secondary) → a 0..1 float used ONLY to sort one response.
|
||||
//
|
||||
// Formula: score = (tier + squash(secondary)) / (maxTier + 1), where
|
||||
// squash(x) = x / (1 + x) maps any non-negative secondary into [0, 1)
|
||||
// so a stronger tier ALWAYS outranks a weaker one regardless of the secondary
|
||||
// value, and within a tier a larger secondary sorts higher. maxTier is the top
|
||||
// enum value (TITLE_EXACT = 3), so the divisor keeps the result in (0, 1].
|
||||
export function computeLookupScore(hit: RankableHit): number {
|
||||
const maxTier = SearchLookupTier.TITLE_EXACT;
|
||||
const secondary = Math.max(0, hit.secondary ?? 0);
|
||||
const squashed = secondary / (1 + secondary);
|
||||
return (hit.tier + squashed) / (maxTier + 1);
|
||||
}
|
||||
|
||||
@Injectable()
|
||||
export class SearchService {
|
||||
constructor(
|
||||
@@ -50,12 +100,19 @@ export class SearchService {
|
||||
userId?: string;
|
||||
workspaceId: string;
|
||||
},
|
||||
): Promise<{ items: SearchResponseDto[] }> {
|
||||
): Promise<{ items: SearchResponseDto[] | SearchLookupResponseDto[] }> {
|
||||
const { query } = searchParams;
|
||||
|
||||
if (query.length < 1) {
|
||||
return { items: [] };
|
||||
}
|
||||
|
||||
// Opt-in agent-lookup mode (#443). Guarded by the `substring` flag so the
|
||||
// web-UI (which never sets it) keeps byte-identical FTS behaviour below.
|
||||
if (searchParams.substring) {
|
||||
return this.searchPageLookup(searchParams, opts);
|
||||
}
|
||||
|
||||
const searchQuery = buildTsQuery(query);
|
||||
|
||||
let queryResults = this.db
|
||||
@@ -175,6 +232,348 @@ export class SearchService {
|
||||
return { items: searchResults };
|
||||
}
|
||||
|
||||
/**
|
||||
* Agent-lookup search (#443, opt-in via `SearchDTO.substring`).
|
||||
*
|
||||
* ADDITIVE to the FTS path: runs a substring branch (title + optionally
|
||||
* text_content, LIKE with metacharacters escaped) MERGED with the existing
|
||||
* FTS branch, so technical tokens that the `english` tokenizer mangles
|
||||
* (`backup-srv.local`, `10.0.12.5`, `WB-MGE-30D86B`) are still found — even
|
||||
* when `buildTsQuery()` returns '' for a dotted/numeric query. Results carry a
|
||||
* location (`path`), a windowed `snippet` and a per-response `score`.
|
||||
*
|
||||
* The whole method is only reached when `substring: true`; the web-UI never
|
||||
* sets it, so its behaviour is unchanged.
|
||||
*/
|
||||
private async searchPageLookup(
|
||||
searchParams: SearchDTO,
|
||||
opts: { userId?: string; workspaceId: string },
|
||||
): Promise<{ items: SearchLookupResponseDto[] }> {
|
||||
const rawQuery = searchParams.query.trim();
|
||||
if (!rawQuery) {
|
||||
return { items: [] };
|
||||
}
|
||||
|
||||
const limit = Math.min(Math.max(searchParams.limit || 10, 1), 50);
|
||||
|
||||
// Normalize the query the same way as the FTS / suggest path: f_unaccent +
|
||||
// lower, done in SQL. `q` is the escaped LIKE pattern body (literal chars).
|
||||
const likeBody = escapeLikePattern(rawQuery);
|
||||
// Compare against `LOWER(f_unaccent(col))`; unaccent+lower the needle too.
|
||||
const needle = sql<string>`LOWER(f_unaccent(${rawQuery}))`;
|
||||
const likePattern = sql<string>`LOWER(f_unaccent(${'%' + likeBody + '%'}))`;
|
||||
const tsQuery = buildTsQuery(rawQuery);
|
||||
const hasTsQuery = tsQuery.length > 0;
|
||||
|
||||
// --- Resolve the space scope. ---------------------------------------------
|
||||
// Mirrors searchPage: explicit spaceId, else the authenticated user's member
|
||||
// spaces. The share path is not exposed to this opt-in mode.
|
||||
let spaceIds: string[] = [];
|
||||
if (searchParams.spaceId) {
|
||||
spaceIds = [searchParams.spaceId];
|
||||
} else if (opts.userId) {
|
||||
spaceIds = await this.spaceMemberRepo.getUserSpaceIds(opts.userId);
|
||||
} else {
|
||||
return { items: [] };
|
||||
}
|
||||
if (spaceIds.length === 0) {
|
||||
return { items: [] };
|
||||
}
|
||||
|
||||
// --- Optional parentPageId subtree scope (inclusive). ---------------------
|
||||
// Reuse the same recursive-descendants pattern used for share-scope.
|
||||
let descendantIds: string[] | null = null;
|
||||
if (searchParams.parentPageId) {
|
||||
const descendants = await this.pageRepo.getPageAndDescendants(
|
||||
searchParams.parentPageId,
|
||||
{ includeContent: false },
|
||||
);
|
||||
descendantIds = descendants.map((p: any) => p.id);
|
||||
if (descendantIds.length === 0) {
|
||||
return { items: [] };
|
||||
}
|
||||
}
|
||||
|
||||
// --- Candidate query: substring (title + text) UNION FTS. -----------------
|
||||
// We compute everything the ranker needs in SQL and pull only small columns
|
||||
// (never the whole text_content) into Node:
|
||||
// - titleExact / titleSub: tier signals
|
||||
// - textMatchPos: 1-based position of the first text match (0 = none)
|
||||
// - ftsRank: ts_rank for the FTS secondary signal (0 when no tsquery)
|
||||
// - snippet: windowed ~500 chars around the first text match, or a leading
|
||||
// text window (title-only hit), or an extended ts_headline fallback.
|
||||
const N_BEFORE = 60; // chars of context before the first match
|
||||
const SNIPPET_LEN = 500;
|
||||
|
||||
let candidates = this.db
|
||||
.selectFrom('pages')
|
||||
.select([
|
||||
'pages.id as id',
|
||||
'pages.slugId as slugId',
|
||||
'pages.title as title',
|
||||
'pages.parentPageId as parentPageId',
|
||||
// Tier signals.
|
||||
sql<boolean>`LOWER(f_unaccent(coalesce(pages.title, ''))) = ${needle}`.as(
|
||||
'titleExact',
|
||||
),
|
||||
sql<boolean>`LOWER(f_unaccent(coalesce(pages.title, ''))) LIKE ${likePattern} ESCAPE '\\'`.as(
|
||||
'titleSub',
|
||||
),
|
||||
// 1-based position of the first text match (0 = no text match).
|
||||
sql<number>`strpos(LOWER(f_unaccent(coalesce(pages.text_content, ''))), ${needle})`.as(
|
||||
'textMatchPos',
|
||||
),
|
||||
// FTS secondary signal (0 when the tsquery is empty).
|
||||
hasTsQuery
|
||||
? sql<number>`ts_rank(pages.tsv, to_tsquery('english', f_unaccent(${tsQuery})))`.as(
|
||||
'ftsRank',
|
||||
)
|
||||
: sql<number>`0`.as('ftsRank'),
|
||||
// Windowed snippet, computed entirely in SQL. Priority:
|
||||
// 1. window around the first text match;
|
||||
// 2. otherwise (titleOnly: no snippet; else) a leading window of the
|
||||
// page text (title-only hit);
|
||||
// 3. otherwise an extended ts_headline for pure-FTS hits.
|
||||
//
|
||||
// #443 snippet-position fix: the match position (`strpos`) is computed in
|
||||
// the LOWER(f_unaccent(...)) space, but f_unaccent is NOT length-
|
||||
// preserving (ß→ss, æ→ae, …→..., ½→ 1/2, full-width forms), so slicing
|
||||
// the ORIGINAL text at that position was misaligned — a single expanding
|
||||
// char before the match shifted the window (or ran it past end → empty).
|
||||
// We now slice from the SAME LOWER(f_unaccent(...)) string so position
|
||||
// and slice share one coordinate space. DELIBERATE trade-off: the snippet
|
||||
// loses original case/diacritics — acceptable for an agent-facing snippet
|
||||
// (position accuracy over original-glyph fidelity). The ts_headline branch
|
||||
// matches over the ORIGINAL text itself, so it is unaffected and kept as-is.
|
||||
searchParams.titleOnly
|
||||
? sql<string>`''`.as('snippet')
|
||||
: sql<string>`
|
||||
coalesce(
|
||||
case
|
||||
when strpos(LOWER(f_unaccent(coalesce(pages.text_content, ''))), ${needle}) > 0
|
||||
then substring(
|
||||
LOWER(f_unaccent(coalesce(pages.text_content, '')))
|
||||
from greatest(1, strpos(LOWER(f_unaccent(coalesce(pages.text_content, ''))), ${needle}) - ${N_BEFORE})
|
||||
for ${SNIPPET_LEN}
|
||||
)
|
||||
when coalesce(pages.text_content, '') <> ''
|
||||
then substring(LOWER(f_unaccent(pages.text_content)) from 1 for 300)
|
||||
${
|
||||
hasTsQuery
|
||||
? sql`else ts_headline('english', coalesce(pages.text_content, ''), to_tsquery('english', f_unaccent(${tsQuery})), 'MinWords=25, MaxWords=40, MaxFragments=3')`
|
||||
: sql``
|
||||
}
|
||||
end,
|
||||
''
|
||||
)
|
||||
`.as('snippet'),
|
||||
])
|
||||
.where('pages.deletedAt', 'is', null)
|
||||
.where('pages.spaceId', 'in', spaceIds);
|
||||
|
||||
if (descendantIds) {
|
||||
candidates = candidates.where('pages.id', 'in', descendantIds);
|
||||
}
|
||||
|
||||
// Match predicate: title substring OR (unless titleOnly) text substring OR
|
||||
// (unless titleOnly) FTS. The substring branch runs even when the tsquery is
|
||||
// empty — that is the dotted/numeric-token case the FTS path misses.
|
||||
//
|
||||
// #443 dead-index fix: these two LIKE predicates MUST match the GIN trgm
|
||||
// index expressions EXACTLY for Postgres to use them. The indexes are on the
|
||||
// coalesce-FREE expressions `LOWER(f_unaccent(title))` (#348's
|
||||
// idx_pages_title_trgm) and `LOWER(f_unaccent(text_content))` (this PR's
|
||||
// idx_pages_text_content_trgm). A `coalesce(col,'')` wrapper here would make
|
||||
// the query expression differ from the index expression and force a Seq Scan
|
||||
// on pages for every lookup. Dropping coalesce is SEMANTICALLY EQUIVALENT:
|
||||
// `NULL LIKE '%q%'` is NULL (falsy), so a NULL title/text simply doesn't
|
||||
// match — exactly as an empty string wouldn't match `%q%`.
|
||||
candidates = candidates.where((eb) => {
|
||||
const ors = [
|
||||
eb(
|
||||
sql`LOWER(f_unaccent(pages.title))`,
|
||||
'like',
|
||||
sql`${likePattern} ESCAPE '\\'`,
|
||||
),
|
||||
];
|
||||
if (!searchParams.titleOnly) {
|
||||
ors.push(
|
||||
eb(
|
||||
sql`LOWER(f_unaccent(pages.text_content))`,
|
||||
'like',
|
||||
sql`${likePattern} ESCAPE '\\'`,
|
||||
),
|
||||
);
|
||||
if (hasTsQuery) {
|
||||
ors.push(
|
||||
sql<boolean>`pages.tsv @@ to_tsquery('english', f_unaccent(${tsQuery}))` as any,
|
||||
);
|
||||
}
|
||||
}
|
||||
return eb.or(ors);
|
||||
});
|
||||
|
||||
// Pull a generous candidate set (before permission filtering + limit).
|
||||
// Cap it so a pathological match set cannot blow up memory; 200 >> limit
|
||||
// (max 50) leaves ample headroom for the post-permission truncation.
|
||||
//
|
||||
// #443 cap-ordering fix: the 200-cap MUST be deterministic and relevance-
|
||||
// biased. Without an ORDER BY, Postgres returns an ARBITRARY 200 rows, so on
|
||||
// a broad match set (common word / short substring) a strong TITLE_EXACT hit
|
||||
// could be among the dropped rows while 200 low-tier TEXT hits fill the cap.
|
||||
// We order by the SAME SQL tier proxies the Node ranker uses — title-exact,
|
||||
// then title-substring, then fts-rank (nulls last), then earliest text-match
|
||||
// position — so the cap keeps the strongest candidates. The Node-side final
|
||||
// tier sort + slice(0, limit) below still runs and stays authoritative; this
|
||||
// ORDER BY only decides WHICH candidates survive the 200-cap.
|
||||
// NB: a BARE integer literal in ORDER BY is read by Postgres as an ordinal
|
||||
// column position (`ORDER BY 0` → "position 0 is not in select list"), so the
|
||||
// no-tsquery fallback is `0::float`, not `0`.
|
||||
const ftsRankExpr = hasTsQuery
|
||||
? sql`ts_rank(pages.tsv, to_tsquery('english', f_unaccent(${tsQuery})))`
|
||||
: sql`0::float`;
|
||||
const candidatesCapped = candidates
|
||||
// Raw-SQL ORDER BY expressions: pass the full `<expr> <dir>` as ONE arg
|
||||
// (the two-arg form treats a raw-SQL second arg as an ORDER BY position).
|
||||
.orderBy(
|
||||
sql`(LOWER(f_unaccent(coalesce(pages.title, ''))) = ${needle}) desc`,
|
||||
)
|
||||
.orderBy(
|
||||
sql`(LOWER(f_unaccent(coalesce(pages.title, ''))) LIKE ${likePattern} ESCAPE '\\') desc`,
|
||||
)
|
||||
.orderBy(sql`${ftsRankExpr} desc nulls last`)
|
||||
// Earlier text match first; strpos returns 0 for "no match", which would
|
||||
// sort BEFORE a real (>=1) position under plain ASC, so push 0 to the end.
|
||||
.orderBy(
|
||||
sql`case when strpos(LOWER(f_unaccent(coalesce(pages.text_content, ''))), ${needle}) = 0 then 2147483647 else strpos(LOWER(f_unaccent(coalesce(pages.text_content, ''))), ${needle}) end asc`,
|
||||
);
|
||||
|
||||
let rows: any[] = await candidatesCapped.limit(200).execute();
|
||||
|
||||
if (rows.length === 0) {
|
||||
return { items: [] };
|
||||
}
|
||||
|
||||
// --- Permissions BEFORE limit. --------------------------------------------
|
||||
// Apply the existing page-level post-filter to the MERGED set, then rank and
|
||||
// only THEN truncate to `limit` — never lose the permission filter.
|
||||
if (opts.userId) {
|
||||
const accessibleIds =
|
||||
await this.pagePermissionRepo.filterAccessiblePageIds({
|
||||
pageIds: rows.map((r) => r.id),
|
||||
userId: opts.userId,
|
||||
spaceId: searchParams.spaceId,
|
||||
workspaceId: opts.workspaceId,
|
||||
});
|
||||
const accessibleSet = new Set(accessibleIds);
|
||||
rows = rows.filter((r) => accessibleSet.has(r.id));
|
||||
}
|
||||
|
||||
if (rows.length === 0) {
|
||||
return { items: [] };
|
||||
}
|
||||
|
||||
// --- Tiered ranking + dedup. ----------------------------------------------
|
||||
// Rows are already unique by id (single pages scan), so no cross-branch
|
||||
// dedup is needed here; the tier captures the strongest match reason.
|
||||
const ranked = rows.map((r) => {
|
||||
let tier: SearchLookupTier;
|
||||
let secondary: number;
|
||||
if (r.titleExact) {
|
||||
tier = SearchLookupTier.TITLE_EXACT;
|
||||
secondary = Number(r.ftsRank) || 0;
|
||||
} else if (r.titleSub) {
|
||||
tier = SearchLookupTier.TITLE_SUBSTRING;
|
||||
secondary = Number(r.ftsRank) || 0;
|
||||
} else {
|
||||
tier = SearchLookupTier.TEXT;
|
||||
// Prefer earlier text matches; map position → closeness in (0, 1].
|
||||
const pos = Number(r.textMatchPos) || 0;
|
||||
secondary =
|
||||
pos > 0 ? 1 / (1 + (pos - 1) / 100) : Number(r.ftsRank) || 0;
|
||||
}
|
||||
return { row: r, tier, score: computeLookupScore({ tier, secondary }) };
|
||||
});
|
||||
|
||||
ranked.sort((a, b) => b.score - a.score);
|
||||
const top = ranked.slice(0, limit);
|
||||
|
||||
// --- Batch ancestor path (ONE recursive CTE, not N+1). --------------------
|
||||
const pathById = await this.buildAncestorPaths(top.map((t) => t.row.id));
|
||||
|
||||
const items: SearchLookupResponseDto[] = top.map((t) => ({
|
||||
id: t.row.id,
|
||||
slugId: t.row.slugId,
|
||||
title: t.row.title,
|
||||
parentPageId: t.row.parentPageId ?? null,
|
||||
path: pathById.get(t.row.id) ?? [],
|
||||
snippet: (t.row.snippet ?? '')
|
||||
.replace(/\r\n|\r|\n/g, ' ')
|
||||
.replace(/\s+/g, ' ')
|
||||
.trim(),
|
||||
score: t.score,
|
||||
}));
|
||||
|
||||
return { items };
|
||||
}
|
||||
|
||||
/**
|
||||
* Batch ancestor-titles helper (#443): ONE recursive CTE seeded with ALL hit
|
||||
* ids, walking UP parentPageId. Returns a map hitId → ancestor titles ordered
|
||||
* root → direct parent (the hit's own title is excluded). Root pages map to
|
||||
* an empty array. Avoids the N+1 of a per-page breadcrumb call.
|
||||
*/
|
||||
private async buildAncestorPaths(
|
||||
hitIds: string[],
|
||||
): Promise<Map<string, string[]>> {
|
||||
const result = new Map<string, string[]>();
|
||||
if (hitIds.length === 0) return result;
|
||||
|
||||
// ancestry(hit_id, page_id, title, parent_page_id, depth): seed one row per
|
||||
// hit at depth 0 (the hit itself), then walk to parents (increasing depth).
|
||||
const rows = await this.db
|
||||
.withRecursive('ancestry', (db) =>
|
||||
db
|
||||
.selectFrom('pages')
|
||||
.select([
|
||||
'pages.id as hitId',
|
||||
'pages.id as pageId',
|
||||
'pages.title as title',
|
||||
'pages.parentPageId as parentPageId',
|
||||
sql<number>`0`.as('depth'),
|
||||
])
|
||||
.where('pages.id', 'in', hitIds)
|
||||
.unionAll((exp) =>
|
||||
exp
|
||||
.selectFrom('pages as p')
|
||||
.innerJoin('ancestry as a', 'p.id', 'a.parentPageId')
|
||||
.select([
|
||||
'a.hitId as hitId',
|
||||
'p.id as pageId',
|
||||
'p.title as title',
|
||||
'p.parentPageId as parentPageId',
|
||||
sql<number>`a.depth + 1`.as('depth'),
|
||||
]),
|
||||
),
|
||||
)
|
||||
.selectFrom('ancestry')
|
||||
.select(['hitId', 'title', 'depth'])
|
||||
// depth 0 is the hit itself — excluded from the path.
|
||||
.where('depth', '>', 0)
|
||||
.orderBy('hitId')
|
||||
// Larger depth = closer to the space root. Ordering DESC gives
|
||||
// root → parent once collected.
|
||||
.orderBy('depth', 'desc')
|
||||
.execute();
|
||||
|
||||
for (const r of rows as any[]) {
|
||||
const list = result.get(r.hitId) ?? [];
|
||||
list.push(r.title);
|
||||
result.set(r.hitId, list);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
async searchSuggestions(
|
||||
suggestion: SearchSuggestionDTO,
|
||||
userId: string,
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
import { type Kysely, sql } from 'kysely';
|
||||
|
||||
/**
|
||||
* #443 — trigram indexes for the opt-in agent-lookup search mode.
|
||||
*
|
||||
* The lookup mode adds a substring branch that runs leading-wildcard
|
||||
* `LOWER(f_unaccent(col)) LIKE '%q%'` predicates on pages.title and
|
||||
* pages.text_content. A leading wildcard cannot use a b-tree index, so without a
|
||||
* GIN trigram index each such predicate is a sequential scan.
|
||||
*
|
||||
* - TITLE: the lookup-mode title predicate is `LOWER(f_unaccent(title)) LIKE
|
||||
* '%q%'` (coalesce-free, so it can use a functional index), which is IDENTICAL
|
||||
* to the one added for /search/suggest (#348). #348's perf-indexes migration
|
||||
* already created `idx_pages_title_trgm` on `(LOWER(f_unaccent(title)))
|
||||
* gin_trgm_ops`, so the title predicate is already covered — we do NOT
|
||||
* re-create that index here (it would be redundant).
|
||||
*
|
||||
* - TEXT_CONTENT: NEW. The substring branch scans text_content when the query
|
||||
* is not titleOnly. text_content is the large column, so a GIN trigram index
|
||||
* on it is the meaningful acceleration for the lookup mode. The lookup search
|
||||
* is ALWAYS space-scoped (spaceId or the user's member spaces), so on small
|
||||
* instances a per-space sequential scan is tolerable — but the index turns the
|
||||
* `%q%` text predicate into a Bitmap Index Scan and removes the only
|
||||
* unbounded-per-space cost of the feature. We add it. The trade-off is disk +
|
||||
* write amplification on page edits (GIN trigram indexes are larger and slower
|
||||
* to update than b-trees); on the small instances this fork targets that cost
|
||||
* is acceptable and the read win on agent lookups is the priority.
|
||||
*
|
||||
* DEPLOY-TIME LOCK WARNING: plain (non-CONCURRENT) CREATE INDEX — Kysely runs
|
||||
* each migration in a transaction, so CONCURRENTLY is impossible. The build takes
|
||||
* a SHARE lock that BLOCKS writes on `pages` for its duration. The text_content
|
||||
* GIN build is the slow one and can take minutes on a large tenant. For large
|
||||
* installations, run this in a maintenance window or build the index out-of-band
|
||||
* with CREATE INDEX CONCURRENTLY before deploying (then `IF NOT EXISTS` no-ops
|
||||
* here). Small/typical tenants are unaffected.
|
||||
*/
|
||||
export async function up(db: Kysely<any>): Promise<void> {
|
||||
// The title predicate is served by #348's idx_pages_title_trgm — see header.
|
||||
// Only the text_content index is introduced here.
|
||||
|
||||
// text_content trigram index. Its expression is coalesce-free —
|
||||
// `LOWER(f_unaccent(text_content))` — to EXACTLY match the coalesce-free
|
||||
// lookup-mode text substring predicate in search.service.ts, so Postgres can
|
||||
// use it (a `coalesce(...)` mismatch would silently fall back to a Seq Scan).
|
||||
await sql`
|
||||
CREATE INDEX IF NOT EXISTS idx_pages_text_content_trgm
|
||||
ON pages USING gin ((LOWER(f_unaccent(text_content))) gin_trgm_ops)
|
||||
`.execute(db);
|
||||
}
|
||||
|
||||
export async function down(db: Kysely<any>): Promise<void> {
|
||||
// Only drop the index this migration introduced. idx_pages_title_trgm is owned
|
||||
// by the #348 perf-indexes migration, so leave it for that migration's down().
|
||||
await sql`DROP INDEX IF EXISTS idx_pages_text_content_trgm`.execute(db);
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
@@ -0,0 +1,280 @@
|
||||
import { randomUUID } from 'crypto';
|
||||
import { CamelCasePlugin, Kysely, sql } from 'kysely';
|
||||
import { PostgresJSDialect } from 'kysely-postgres-js';
|
||||
// NOT a default import: the project tsconfig is `module: commonjs` with NO
|
||||
// esModuleInterop, so `import postgres from 'postgres'` compiles to
|
||||
// `postgres_1.default(...)` and the CJS `postgres` export has no `.default` —
|
||||
// it threw in beforeAll, was swallowed as "DB unreachable", and SILENTLY voided
|
||||
// all six tests. Mirror the working integration harness (test/integration/db.ts).
|
||||
import * as postgres from 'postgres';
|
||||
import { AiChatMessageRepo } from './ai-chat-message.repo';
|
||||
import { AiChatRunRepo } from './ai-chat-run.repo';
|
||||
|
||||
/**
|
||||
* #491 delta-poll — OBSERVABLE-PROPERTY tests against a LIVE Postgres (the local
|
||||
* gitmost test DB, docker `gitmost-test-pg` on :5432), not "rows through a mock"
|
||||
* (a mock cannot observe the DB clock nor the overlap-window race — the very
|
||||
* things that matter here). Drives the REAL repo methods (`findByChatUpdatedAfter`,
|
||||
* the now()-stamped `update`) and asserts:
|
||||
* 1. delta-relevant writes stamp `updatedAt` from the DB clock, not the app clock
|
||||
* (proven by faking the process clock far into the future and observing the
|
||||
* stamp stays on real DB time);
|
||||
* 2. the poll returns only rows changed after the cursor, ordered, with a fresh
|
||||
* DB-clock cursor;
|
||||
* 3. the "committed late but stamped earlier than the cursor" RACE is caught by
|
||||
* the overlap window (a naive `updatedAt > cursor` would MISS it);
|
||||
* 4. the overlap GUARANTEES repeats across close polls — the contract behind the
|
||||
* client's idempotent merge (mergeById).
|
||||
*
|
||||
* INTEGRATION lane (`*.int-spec.ts`): runs under `test:int`, whose global-setup
|
||||
* DROPS + RE-CREATES + MIGRATES `docmost_test`, so the real `ai_chat_messages` /
|
||||
* `ai_chat_runs` tables EXIST here. (It was previously a `.spec.ts` defaulting to
|
||||
* the UNmigrated dev `docmost`; in the CI unit lane — where `WAL_TEST_DATABASE_URL`
|
||||
* is unset and only `test:int` migrates — that meant 5/6 ERROR
|
||||
* `relation "ai_chat_messages" does not exist`, silently voiding coverage of the
|
||||
* risky cursor/overlap logic. Renaming to `.int-spec.ts` + defaulting the DSN to
|
||||
* `docmost_test` fixes the CI fidelity.)
|
||||
*
|
||||
* FK triggers are bypassed (`session_replication_role = replica`) so synthetic
|
||||
* chat/workspace ids need no parent fixtures; a single pooled connection (max 1)
|
||||
* keeps that session setting for every query. SKIPS cleanly when the DB is
|
||||
* unreachable so a DB-less CI never breaks.
|
||||
*/
|
||||
const CONN =
|
||||
process.env.WAL_TEST_DATABASE_URL ??
|
||||
process.env.TEST_DATABASE_URL ??
|
||||
'postgresql://docmost:docmost_dev_pw@localhost:5432/docmost_test';
|
||||
|
||||
let db: Kysely<any>;
|
||||
let sqlClient: ReturnType<typeof postgres>;
|
||||
let msgRepo: AiChatMessageRepo;
|
||||
let runRepo: AiChatRunRepo;
|
||||
let reachable = false;
|
||||
|
||||
const workspaceId = randomUUID();
|
||||
const chatId = randomUUID();
|
||||
|
||||
beforeAll(async () => {
|
||||
try {
|
||||
sqlClient = postgres(CONN, { max: 1, onnotice: () => {} });
|
||||
db = new Kysely<any>({
|
||||
dialect: new PostgresJSDialect({ postgres: sqlClient }),
|
||||
plugins: [new CamelCasePlugin()],
|
||||
});
|
||||
// Single connection keeps this session-scoped bypass for the whole suite.
|
||||
await sql`set session_replication_role = replica`.execute(db);
|
||||
await sql`select 1`.execute(db);
|
||||
reachable = true;
|
||||
} catch (err) {
|
||||
reachable = false;
|
||||
// A genuine connection failure (ECONNREFUSED etc.) is a legitimate skip on a
|
||||
// DB-less CI. A PROGRAMMING error (bad import, typo, driver misuse) must NOT
|
||||
// masquerade as "DB unreachable" and silently void the whole suite (that is
|
||||
// exactly the bug that hid this spec's zero coverage) — rethrow it so the
|
||||
// suite fails LOUDLY.
|
||||
const msg = String((err as Error)?.message ?? err);
|
||||
if (
|
||||
!/ECONNREFUSED|ENOTFOUND|ETIMEDOUT|EHOSTUNREACH|connect|terminating|password|authentication|role .* does not exist|database .* does not exist/i.test(
|
||||
msg,
|
||||
)
|
||||
) {
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
msgRepo = new AiChatMessageRepo(db as never);
|
||||
runRepo = new AiChatRunRepo(db as never);
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
if (db) {
|
||||
try {
|
||||
await db
|
||||
.deleteFrom('aiChatMessages')
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.execute();
|
||||
await db
|
||||
.deleteFrom('aiChatRuns')
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.execute();
|
||||
} catch {
|
||||
/* best-effort cleanup */
|
||||
}
|
||||
await db.destroy();
|
||||
}
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
jest.useRealTimers();
|
||||
});
|
||||
|
||||
async function seedMessage(overrides: Record<string, unknown> = {}) {
|
||||
return msgRepo.insert({
|
||||
chatId,
|
||||
workspaceId,
|
||||
userId: null as never,
|
||||
role: 'assistant',
|
||||
content: 'x',
|
||||
status: 'streaming',
|
||||
...overrides,
|
||||
} as never);
|
||||
}
|
||||
|
||||
async function dbNow(): Promise<string> {
|
||||
const r = await sql<{ now: Date }>`select now() as now`.execute(db);
|
||||
return r.rows[0].now.toISOString();
|
||||
}
|
||||
|
||||
// Fake ONLY the Date object (so in-process `new Date()`/`Date.now()` jump), while
|
||||
// leaving every TIMER function real. Faking timers wholesale freezes postgres.js's
|
||||
// internal connection/query timers, so the awaited DB round-trip would hang the
|
||||
// test (and the afterAll cleanup) at the jest 5s cap. With Date-only faking the
|
||||
// query resolves normally, and we still prove the stamp is the DB clock (not the
|
||||
// skewed process clock).
|
||||
function fakeDateOnly(iso: string): void {
|
||||
jest.useFakeTimers({
|
||||
doNotFake: [
|
||||
'hrtime',
|
||||
'nextTick',
|
||||
'performance',
|
||||
'queueMicrotask',
|
||||
'requestAnimationFrame',
|
||||
'cancelAnimationFrame',
|
||||
'requestIdleCallback',
|
||||
'cancelIdleCallback',
|
||||
'setImmediate',
|
||||
'clearImmediate',
|
||||
'setInterval',
|
||||
'clearInterval',
|
||||
'setTimeout',
|
||||
'clearTimeout',
|
||||
],
|
||||
now: new Date(iso),
|
||||
});
|
||||
}
|
||||
|
||||
const maybe = (name: string, fn: () => Promise<void>) =>
|
||||
it(name, async () => {
|
||||
if (!reachable) {
|
||||
console.warn(`SKIP (${name}): test DB unreachable at ${CONN}`);
|
||||
return;
|
||||
}
|
||||
await fn();
|
||||
});
|
||||
|
||||
describe('AiChatMessageRepo.findByChatUpdatedAfter (#491 delta poll)', () => {
|
||||
maybe('null cursor returns no rows and a fresh DB-clock cursor', async () => {
|
||||
const before = await dbNow();
|
||||
const { rows, cursor } = await msgRepo.findByChatUpdatedAfter(
|
||||
chatId,
|
||||
workspaceId,
|
||||
null,
|
||||
);
|
||||
expect(rows).toEqual([]);
|
||||
expect(new Date(cursor).getTime()).toBeGreaterThanOrEqual(
|
||||
new Date(before).getTime(),
|
||||
);
|
||||
});
|
||||
|
||||
maybe('returns only rows changed after the cursor', async () => {
|
||||
const { cursor: c0 } = await msgRepo.findByChatUpdatedAfter(
|
||||
chatId,
|
||||
workspaceId,
|
||||
null,
|
||||
);
|
||||
const m = await seedMessage();
|
||||
const { rows, cursor: c1 } = await msgRepo.findByChatUpdatedAfter(
|
||||
chatId,
|
||||
workspaceId,
|
||||
c0,
|
||||
);
|
||||
expect(rows.map((r) => r.id)).toContain(m.id);
|
||||
// Cursor is monotonic (advances).
|
||||
expect(new Date(c1).getTime()).toBeGreaterThanOrEqual(
|
||||
new Date(c0).getTime(),
|
||||
);
|
||||
});
|
||||
|
||||
maybe(
|
||||
'RACE: a row stamped BEFORE the cursor but seen after is caught by the overlap',
|
||||
async () => {
|
||||
// Cursor taken now; then a row appears whose updatedAt is 2s in the PAST
|
||||
// (committed late on another connection but stamped earlier). A naive
|
||||
// `updatedAt > cursor` would MISS it; the 5s overlap window catches it.
|
||||
const cursor = await dbNow();
|
||||
const m = await seedMessage();
|
||||
await sql`update ai_chat_messages set updated_at = now() - interval '2 seconds' where id = ${m.id}`.execute(
|
||||
db,
|
||||
);
|
||||
const { rows } = await msgRepo.findByChatUpdatedAfter(
|
||||
chatId,
|
||||
workspaceId,
|
||||
cursor,
|
||||
);
|
||||
expect(rows.map((r) => r.id)).toContain(m.id);
|
||||
},
|
||||
);
|
||||
|
||||
maybe(
|
||||
'overlap GUARANTEES repeats across close polls (idempotent-merge contract)',
|
||||
async () => {
|
||||
const { cursor: c0 } = await msgRepo.findByChatUpdatedAfter(
|
||||
chatId,
|
||||
workspaceId,
|
||||
null,
|
||||
);
|
||||
const m = await seedMessage();
|
||||
const first = await msgRepo.findByChatUpdatedAfter(
|
||||
chatId,
|
||||
workspaceId,
|
||||
c0,
|
||||
);
|
||||
expect(first.rows.map((r) => r.id)).toContain(m.id);
|
||||
// Immediately re-poll with the JUST-returned cursor: the row is still within
|
||||
// the overlap window, so it is returned AGAIN — the client MUST dedupe by id.
|
||||
const second = await msgRepo.findByChatUpdatedAfter(
|
||||
chatId,
|
||||
workspaceId,
|
||||
first.cursor,
|
||||
);
|
||||
expect(second.rows.map((r) => r.id)).toContain(m.id);
|
||||
},
|
||||
);
|
||||
|
||||
maybe(
|
||||
'update() stamps updatedAt from the DB clock, not the app clock',
|
||||
async () => {
|
||||
const m = await seedMessage();
|
||||
// Skew the PROCESS clock ~73 years into the future (Date only). If the stamp
|
||||
// came from `new Date()` the row would read year 2099; sql now() keeps it on
|
||||
// DB time.
|
||||
fakeDateOnly('2099-01-01T00:00:00Z');
|
||||
const updated = await msgRepo.update(m.id, workspaceId, {
|
||||
content: 'y',
|
||||
});
|
||||
jest.useRealTimers();
|
||||
expect(updated).toBeDefined();
|
||||
expect(new Date(updated!.updatedAt).getFullYear()).toBeLessThan(2099);
|
||||
},
|
||||
);
|
||||
|
||||
maybe(
|
||||
'run update() also stamps updatedAt from the DB clock',
|
||||
async () => {
|
||||
const run = await runRepo.insert({
|
||||
chatId,
|
||||
workspaceId,
|
||||
createdBy: null as never,
|
||||
trigger: 'user',
|
||||
status: 'running',
|
||||
stepCount: 0,
|
||||
} as never);
|
||||
fakeDateOnly('2099-01-01T00:00:00Z');
|
||||
const updated = await runRepo.update(run.id, workspaceId, {
|
||||
stepCount: 1,
|
||||
});
|
||||
jest.useRealTimers();
|
||||
expect(updated).toBeDefined();
|
||||
expect(new Date(updated!.updatedAt).getFullYear()).toBeLessThan(2099);
|
||||
},
|
||||
);
|
||||
});
|
||||
@@ -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 {
|
||||
@@ -24,6 +25,20 @@ const SWEEP_STREAMING_STALE_MS = 10 * 60 * 1000; // 10 minutes
|
||||
// into memory; far above any realistic transcript length.
|
||||
const FIND_ALL_BY_CHAT_LIMIT = 5000;
|
||||
|
||||
// Delta-poll overlap (#491): the poll query reaches this far BEHIND the client's
|
||||
// echoed cursor, so a row that committed with an `updatedAt` marginally before the
|
||||
// previous cursor was taken (on another autocommit connection) is still caught.
|
||||
// Sized well above realistic single-row commit skew; the client merge is
|
||||
// idempotent by id (mergeById), so the guaranteed repeats the overlap produces are
|
||||
// harmless.
|
||||
export const DELTA_POLL_OVERLAP_SECONDS = 5;
|
||||
|
||||
// Hard cap on rows one delta poll returns — a safety bound (a poll should carry a
|
||||
// handful of just-changed rows, never a whole transcript). Ordered by (updatedAt,
|
||||
// id) asc, so on the pathological overflow the OLDEST changes win and the newest
|
||||
// are picked up by the next poll (its cursor did not advance past them).
|
||||
export const DELTA_POLL_MAX_ROWS = 500;
|
||||
|
||||
@Injectable()
|
||||
export class AiChatMessageRepo {
|
||||
private readonly logger = new Logger(AiChatMessageRepo.name);
|
||||
@@ -138,6 +153,72 @@ export class AiChatMessageRepo {
|
||||
.executeTakeFirst();
|
||||
}
|
||||
|
||||
/**
|
||||
* Delta read (#491) for the degraded poll: the chat's messages whose row
|
||||
* changed AFTER the client's `cursor`, plus a FRESH cursor taken from the DB
|
||||
* clock. Replaces the old "refetch ALL infinite-query pages every 2.5s with
|
||||
* full parts" poll — the client seeds once (findByChat) and thereafter pulls
|
||||
* only the deltas and merges them by id (mergeById).
|
||||
*
|
||||
* Cursor: a DB-clock timestamp (now()) the client echoes back each poll. All
|
||||
* delta-relevant writes stamp `updatedAt` with now() (see `update` /
|
||||
* `finalizeOwner`), so this is a SINGLE monotonic axis. The query overlaps the
|
||||
* cursor by DELTA_POLL_OVERLAP_SECONDS to catch a row committed with an
|
||||
* `updatedAt` marginally BEFORE the previous cursor was taken on another
|
||||
* connection (single-row autocommit UPDATEs; no long transactions). The overlap
|
||||
* GUARANTEES occasional REPEATS, so the client merge MUST be idempotent by id.
|
||||
*
|
||||
* `cursor === null` (first poll after the full seed) returns NO rows — there is
|
||||
* nothing "new" relative to a just-loaded seed — only the fresh cursor to start
|
||||
* the delta chain. The fresh cursor is read AFTER the rows, so it is >= every
|
||||
* returned row's `updatedAt` (they were read strictly earlier) — a row that
|
||||
* commits between the rows-read and the cursor-read is at most
|
||||
* DELTA_POLL_OVERLAP_SECONDS behind the returned cursor, so the next poll's
|
||||
* overlap window always re-includes it (no miss).
|
||||
*/
|
||||
async findByChatUpdatedAfter(
|
||||
chatId: string,
|
||||
workspaceId: string,
|
||||
cursor: string | null,
|
||||
): Promise<{ rows: AiChatMessage[]; cursor: string }> {
|
||||
if (cursor === null) {
|
||||
const nowRow = await sql<{ now: Date }>`select now() as now`.execute(
|
||||
this.db,
|
||||
);
|
||||
return { rows: [], cursor: nowRow.rows[0].now.toISOString() };
|
||||
}
|
||||
// Overlap the client cursor by DELTA_POLL_OVERLAP_SECONDS, computed in SQL off
|
||||
// the echoed cursor so the whole comparison stays on the DB clock.
|
||||
const rows = await this.db
|
||||
.selectFrom('aiChatMessages')
|
||||
.select(this.baseFields)
|
||||
.where('chatId', '=', chatId)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.where('deletedAt', 'is', null)
|
||||
.where(
|
||||
'updatedAt',
|
||||
'>',
|
||||
sql<Date>`${cursor}::timestamptz - make_interval(secs => ${DELTA_POLL_OVERLAP_SECONDS})`,
|
||||
)
|
||||
.orderBy('updatedAt', 'asc')
|
||||
.orderBy('id', 'asc')
|
||||
.limit(DELTA_POLL_MAX_ROWS)
|
||||
.execute();
|
||||
// When the page filled (pathological overflow), DO NOT advance the cursor to
|
||||
// now(): that would skip the changed rows past the cap that this poll did not
|
||||
// return. Resume from the last returned row's updatedAt instead (the next
|
||||
// poll's overlap re-includes ties by id). In the normal case the fresh DB-clock
|
||||
// now() is the cursor.
|
||||
if (rows.length === DELTA_POLL_MAX_ROWS) {
|
||||
return {
|
||||
rows,
|
||||
cursor: rows[rows.length - 1].updatedAt.toISOString(),
|
||||
};
|
||||
}
|
||||
const nowRow = await sql<{ now: Date }>`select now() as now`.execute(this.db);
|
||||
return { rows, cursor: nowRow.rows[0].now.toISOString() };
|
||||
}
|
||||
|
||||
async insert(
|
||||
insertable: InsertableAiChatMessage,
|
||||
trx?: KyselyTransaction,
|
||||
@@ -171,7 +252,13 @@ export class AiChatMessageRepo {
|
||||
const db = dbOrTx(this.db, opts?.trx);
|
||||
let query = db
|
||||
.updateTable('aiChatMessages')
|
||||
.set({ ...(patch as Record<string, unknown>), updatedAt: new Date() })
|
||||
// #491: stamp `updatedAt` from the DB clock (sql now()), NOT the app clock
|
||||
// (new Date()). The delta-poll cursor (findByChatUpdatedAfter) is a single
|
||||
// DB-clock axis; a per-step 'streaming' UPDATE stamped with the app clock
|
||||
// would be a SECOND, skewed clock source and could leave a row's updatedAt
|
||||
// just under a cursor taken from now() on another connection — an
|
||||
// independent source of delta MISSES. All delta-relevant writes use now().
|
||||
.set({ ...(patch as Record<string, unknown>), updatedAt: sql`now()` })
|
||||
.where('id', '=', id)
|
||||
.where('workspaceId', '=', workspaceId);
|
||||
// Concurrency guard (#183 review): a per-step 'streaming' update must NEVER
|
||||
@@ -188,6 +275,150 @@ 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')
|
||||
// #491: DB-clock stamp (see `update`) — this terminal write flips the row's
|
||||
// status, which the delta poll must observe on the shared now() cursor axis.
|
||||
.set({ ...(patch as Record<string, unknown>), updatedAt: sql`now()` })
|
||||
.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)`,
|
||||
// #491: DB-clock stamp (see `update`) so a reconcile status flip lands on
|
||||
// the same now() cursor axis the delta poll reads.
|
||||
updatedAt: sql`now()`,
|
||||
})
|
||||
.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)`,
|
||||
// #491: DB-clock stamp (see `update`). The staleness WHERE below stays on
|
||||
// the app clock — a >minutes window makes the ms-scale skew irrelevant.
|
||||
updatedAt: sql`now()`,
|
||||
})
|
||||
.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 +431,21 @@ 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)`,
|
||||
// #491: DB-clock stamp (see `update`). Staleness WHERE stays app-clock.
|
||||
updatedAt: sql`now()`,
|
||||
})
|
||||
.where('status', '=', 'streaming')
|
||||
.where('updatedAt', '<', staleBefore)
|
||||
.returning('id')
|
||||
|
||||
@@ -62,10 +62,17 @@ describe('AiChatRunRepo.sweepRunning', () => {
|
||||
// ...but a fresh 'running' run (updatedAt = now) must NOT be skipped: no
|
||||
// updatedAt predicate at all on the boot path.
|
||||
expect(rec.wheres.some(([col]) => col === 'updatedAt')).toBe(false);
|
||||
// It flips to 'aborted' and stamps finishedAt.
|
||||
expect(rec.set).toEqual(
|
||||
expect.objectContaining({ status: 'aborted', finishedAt: expect.any(Date) }),
|
||||
);
|
||||
// It flips to 'aborted' and stamps finishedAt + updatedAt. #491: the stamps
|
||||
// are now DB-clock `sql now()` expressions (raw builders), NOT app-clock
|
||||
// `new Date()`, so the run row shares the delta poll's single now() cursor axis
|
||||
// — assert they are present and are the sql raw-builder objects (not a Date,
|
||||
// not undefined).
|
||||
expect(rec.set?.status).toBe('aborted');
|
||||
for (const stamp of ['finishedAt', 'updatedAt'] as const) {
|
||||
expect(rec.set?.[stamp]).toBeDefined();
|
||||
expect(rec.set?.[stamp]).not.toBeInstanceOf(Date);
|
||||
expect(typeof rec.set?.[stamp]).toBe('object');
|
||||
}
|
||||
});
|
||||
|
||||
it('phase-2 path: an explicit staleMs reintroduces the updatedAt window', async () => {
|
||||
|
||||
@@ -136,13 +136,53 @@ export class AiChatRunRepo {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
return db
|
||||
.updateTable('aiChatRuns')
|
||||
.set({ ...(patch as Record<string, unknown>), updatedAt: new Date() })
|
||||
// #491: DB-clock stamp (sql now()) so the run row shares the delta poll's
|
||||
// single now() cursor axis with the assistant message rows — a run-status
|
||||
// change (the run fact the delta carries) must never sit on a skewed app
|
||||
// clock relative to the message updatedAt cursor.
|
||||
.set({ ...(patch as Record<string, unknown>), updatedAt: sql`now()` })
|
||||
.where('id', '=', id)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.returning(this.baseFields)
|
||||
.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);
|
||||
return db
|
||||
.updateTable('aiChatRuns')
|
||||
.set({
|
||||
status: patch.status,
|
||||
error: patch.error,
|
||||
// #491: DB-clock stamps (finished_at + updated_at) so the terminal run
|
||||
// fact lands on the delta poll's now() cursor axis.
|
||||
finishedAt: sql`now()`,
|
||||
updatedAt: sql`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
|
||||
@@ -157,7 +197,8 @@ export class AiChatRunRepo {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
return db
|
||||
.updateTable('aiChatRuns')
|
||||
.set({ stopRequestedAt: new Date(), updatedAt: new Date() })
|
||||
// #491: DB-clock stamps (see `update`).
|
||||
.set({ stopRequestedAt: sql`now()`, updatedAt: sql`now()` })
|
||||
.where('id', '=', id)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[])
|
||||
@@ -184,18 +225,44 @@ 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,
|
||||
): Promise<number> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
const now = new Date();
|
||||
let query = db
|
||||
.updateTable('aiChatRuns')
|
||||
.set({
|
||||
status: 'aborted',
|
||||
finishedAt: now,
|
||||
updatedAt: now,
|
||||
// #491: DB-clock stamps (see `update`). The staleness WHERE below stays on
|
||||
// the app clock — a >minutes window makes the ms-scale skew irrelevant.
|
||||
finishedAt: sql`now()`,
|
||||
updatedAt: sql`now()`,
|
||||
error: sql`coalesce(error, ${'Run interrupted by a server restart.'})`,
|
||||
})
|
||||
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[]);
|
||||
@@ -203,7 +270,7 @@ export class AiChatRunRepo {
|
||||
// sibling replica's live run is never aborted. Omitted on the phase-1 boot
|
||||
// sweep -> unconditional.
|
||||
if (typeof opts.staleMs === 'number') {
|
||||
const staleBefore = new Date(now.getTime() - opts.staleMs);
|
||||
const staleBefore = new Date(Date.now() - opts.staleMs);
|
||||
query = query.where('updatedAt', '<', staleBefore);
|
||||
}
|
||||
const rows = await query.returning('id').execute();
|
||||
|
||||
+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
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user