Compare commits
33 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 5b17503df7 | |||
| 03eafa6c68 | |||
| a42f1ead48 | |||
| b7a3ec227d | |||
| 846341d7d4 | |||
| 32e10ca6d3 | |||
| 74d212cfd3 | |||
| 9a6389133b | |||
| 0b8497e496 | |||
| 433252bdb1 | |||
| 6f81067f4d | |||
| 55e8f61b3c | |||
| ab133fa0d0 | |||
| d42ca8dc57 | |||
| 25a31e6c0d | |||
| 3550bfa411 | |||
| a0ecf21cb5 | |||
| d81781aa27 | |||
| daca5ce8d6 | |||
| daeeb1f3f2 | |||
| a919c79cb2 | |||
| 798a81abfe | |||
| 216499d57b | |||
| bb5abf29a6 | |||
| 0ddeaadeee | |||
| bc632f9d56 | |||
| d3a049d176 | |||
| d738780370 | |||
| daf728676f | |||
| 1d8e3444f4 | |||
| c50f5b66bb | |||
| 51d44b6061 | |||
| 6247585b66 |
@@ -251,6 +251,16 @@ MCP_DOCMOST_PASSWORD=
|
||||
# Default 120000 (2 min).
|
||||
# AI_MCP_CALL_TIMEOUT_MS=120000
|
||||
|
||||
# Kill-switch for the agent API-key feature (#501). Default ON when unset — a
|
||||
# deploy that never sets it must NOT silently kill every agent. STRICT parse:
|
||||
# only the literals `true` / `false` are accepted; a typo like `=0`/`=off`/`=False`
|
||||
# FAILS AT BOOT by design (never silently read as "enabled"), so the switch is
|
||||
# guaranteed to actually flip when an operator flips it during an incident. When
|
||||
# set to `false`: all api-key auth is DENIED (every api-key token is rejected) and
|
||||
# the api-key management endpoints return 404. The resolved state is logged at boot
|
||||
# (`API keys: ENABLED/DISABLED (API_KEYS_ENABLED=...)`) so it is verifiable per deploy.
|
||||
# API_KEYS_ENABLED=true
|
||||
|
||||
# Max JSON/urlencoded request body size (bytes). Fastify's 1 MiB default is too
|
||||
# small for a long AI-chat research turn: the client resends the FULL message
|
||||
# history (every tool call + search result) on each turn, so a deep conversation's
|
||||
@@ -302,6 +312,16 @@ 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,
|
||||
|
||||
@@ -226,6 +226,13 @@ jobs:
|
||||
- name: Build mcp
|
||||
run: pnpm --filter @docmost/mcp build
|
||||
|
||||
# apps/server imports @docmost/token-estimate at runtime (history-budget.ts,
|
||||
# #490); its dist/ is gitignored and `test:e2e` type-checks + runs the code,
|
||||
# so build it here or tsc fails with TS2307 Cannot find module
|
||||
# '@docmost/token-estimate' (mirrors the editor-ext / mcp build steps above).
|
||||
- name: Build token-estimate
|
||||
run: pnpm --filter @docmost/token-estimate build
|
||||
|
||||
- name: Run migrations
|
||||
run: pnpm --filter ./apps/server migration:latest
|
||||
|
||||
|
||||
@@ -159,6 +159,14 @@ jobs:
|
||||
- name: Build prosemirror-markdown
|
||||
run: pnpm --filter @docmost/prosemirror-markdown build
|
||||
|
||||
# @docmost/token-estimate is a shared workspace package the client vitest
|
||||
# suite resolves via its dist build (main: ./dist/index.js); dist/ is
|
||||
# gitignored and `pnpm -r test` does NOT honour nx `dependsOn: ^build`, so
|
||||
# build it before the recursive test run or the client suite fails with
|
||||
# "Failed to resolve import '@docmost/token-estimate'" (#490).
|
||||
- name: Build token-estimate
|
||||
run: pnpm --filter @docmost/token-estimate build
|
||||
|
||||
- name: Run unit tests
|
||||
run: pnpm -r test
|
||||
|
||||
|
||||
@@ -135,6 +135,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
new resync path re-reads the live anchor so the suggestion applies against the
|
||||
current text, and orphaned anchors (whose marked run was deleted) are
|
||||
reconciled rather than left blocking. (#496)
|
||||
- **Save intentional page versions.** Press `Cmd/Ctrl+S` (or use the page menu)
|
||||
to save a named version of a page. The history panel now distinguishes
|
||||
intentional versions (a "Saved" / "Agent version" badge) from automatic
|
||||
snapshots, dims autosaves, and offers an "Only versions" filter. Automatic
|
||||
snapshots switched from a fixed interval to a trailing idle-flush with a
|
||||
max-wait ceiling, and a boundary snapshot is pinned whenever the editing source
|
||||
changes (e.g. a person's edits followed by the AI agent). (#370)
|
||||
|
||||
- **Place several images side by side in a row.** A new "Inline (side by
|
||||
side)" alignment mode in the image bubble menu renders consecutive inline
|
||||
@@ -385,6 +392,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
`@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)
|
||||
- **Cyrillic (and any non-ASCII) draw.io labels no longer turn into mojibake
|
||||
when a diagram is opened in the draw.io editor.** Agent-created diagrams
|
||||
(`drawioCreate`) and Confluence-imported diagrams stored their model in the
|
||||
SVG's `content=` attribute as base64; the draw.io editor decodes that via
|
||||
Latin-1 `atob` (no UTF-8 step), so every non-ASCII char (e.g. `Старт-бит`,
|
||||
`ё`, `—`) split into garbage and the editor's autosave then persisted the
|
||||
corrupted model, breaking the page preview too. Both write paths
|
||||
(`buildDrawioSvg`, the import service's `createDrawioSvg`) now write `content=`
|
||||
as XML-entity-escaped mxfile XML — draw.io's own native form, decoded by the
|
||||
DOM as UTF-8 — so labels open intact. The decoder reads both the new
|
||||
entity-encoded form and the old base64 form, so existing diagrams still open.
|
||||
*Healing pre-fix diagrams:* only a diagram that still holds its original
|
||||
(correct-UTF-8) base64 — i.e. one not yet opened/autosaved in the draw.io
|
||||
editor — can be repaired in place by `drawioGet` → `drawioUpdate` with the
|
||||
same XML (rewrites the attachment in the new form); no migration script is
|
||||
needed. A diagram that was already opened in the editor persisted the
|
||||
mojibake at rest, so `drawioGet` reads the already-corrupted text and
|
||||
`drawioUpdate` faithfully rewrites it — that text is lost and is not
|
||||
recoverable by a rewrite. (#507)
|
||||
- **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
|
||||
|
||||
@@ -59,6 +59,14 @@ COPY --from=builder /app/packages/mcp/data /app/packages/mcp/data
|
||||
COPY --from=builder /app/packages/prosemirror-markdown/build /app/packages/prosemirror-markdown/build
|
||||
COPY --from=builder /app/packages/prosemirror-markdown/package.json /app/packages/prosemirror-markdown/package.json
|
||||
|
||||
# apps/server imports @docmost/token-estimate (workspace:*) at runtime
|
||||
# (history-budget.ts, #490). tsc emits only dist/ and dist/ is gitignored, so the
|
||||
# prod install would resolve a broken workspace symlink and the server would die
|
||||
# with ERR_MODULE_NOT_FOUND on the first history-budget call. Ship the built
|
||||
# package + its manifest, mirroring prosemirror-markdown above.
|
||||
COPY --from=builder /app/packages/token-estimate/dist /app/packages/token-estimate/dist
|
||||
COPY --from=builder /app/packages/token-estimate/package.json /app/packages/token-estimate/package.json
|
||||
|
||||
# Copy root package files
|
||||
COPY --from=builder /app/package.json /app/package.json
|
||||
COPY --from=builder /app/pnpm*.yaml /app/
|
||||
|
||||
@@ -206,6 +206,137 @@ start the new migrations apply on top of your existing schema (`CREATE EXTENSION
|
||||
existing pages are indexed on their next edit. pgvector is still required for the migration to
|
||||
apply at all.
|
||||
|
||||
## Local embeddings server
|
||||
|
||||
The AI agent's semantic (RAG) search needs an **embeddings model**. Instead of paying a cloud
|
||||
provider (e.g. OpenAI `text-embedding-3-*`) to embed every page, you can run a small open-weights
|
||||
model yourself with Hugging Face
|
||||
[Text Embeddings Inference](https://github.com/huggingface/text-embeddings-inference) (TEI), which
|
||||
serves an OpenAI-compatible `/v1/embeddings` endpoint. `intfloat/multilingual-e5-small` is a good
|
||||
default: multilingual, 384-dim, and comfortable on CPU (~1–2 GB RAM, 1–2 vCPU). Point Gitmost at it
|
||||
under **Workspace settings → AI → Embeddings**.
|
||||
|
||||
### Option A — local (same Docker network as Gitmost)
|
||||
|
||||
Run TEI as a container on the network Gitmost is already on. The port is never published, so the
|
||||
endpoint stays internal and needs no authentication.
|
||||
|
||||
```yaml
|
||||
services:
|
||||
embeddings:
|
||||
image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.9 # pin version; use a cuda-* tag for GPU
|
||||
container_name: embeddings
|
||||
restart: unless-stopped
|
||||
networks:
|
||||
- gitmost_net # same network Gitmost is on
|
||||
command:
|
||||
- "--model-id"
|
||||
- "intfloat/multilingual-e5-small"
|
||||
- "--auto-truncate" # clamp over-long inputs instead of returning 413
|
||||
volumes:
|
||||
- tei-models:/data # weights are downloaded once and cached here
|
||||
|
||||
networks:
|
||||
gitmost_net:
|
||||
external: true # the network Gitmost already uses
|
||||
|
||||
volumes:
|
||||
tei-models:
|
||||
```
|
||||
|
||||
Gitmost settings (**Workspace settings → AI → Embeddings**):
|
||||
|
||||
| Field | Value |
|
||||
|-------------------|-----------------------------------|
|
||||
| Model | `intfloat/multilingual-e5-small` |
|
||||
| Base URL | `http://embeddings:80/v1/` |
|
||||
| Embedding API key | — (leave empty) |
|
||||
|
||||
> `embeddings` is the container name — Gitmost resolves it over DNS inside the Docker network.
|
||||
> The port is not published, so the endpoint is reachable only by containers on that network and
|
||||
> no authorization is required.
|
||||
|
||||
### Option B — separate host (public via Traefik + Let's Encrypt)
|
||||
|
||||
This assumes the host already runs Traefik with an ACME resolver (the example below uses
|
||||
`letsEncrypt`, the `websecure` entrypoint and a shared `docker_main_net` network). Replace the
|
||||
domain / network / resolver with your own.
|
||||
|
||||
**DNS:** add an A record `embeddings.example.com` → the IP of your Traefik host (same
|
||||
challenge / port 80 as the rest of your sites).
|
||||
|
||||
```yaml
|
||||
services:
|
||||
embeddings:
|
||||
image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.9 # pin version; cuda-* tag for GPU
|
||||
container_name: embeddings
|
||||
restart: unless-stopped
|
||||
networks:
|
||||
- docker_main_net # the network Traefik is attached to
|
||||
command:
|
||||
- "--model-id"
|
||||
- "intfloat/multilingual-e5-small"
|
||||
- "--auto-truncate"
|
||||
- "--api-key"
|
||||
- "sk-emb-REPLACE_WITH_YOUR_KEY"
|
||||
volumes:
|
||||
- tei-models:/data
|
||||
labels:
|
||||
traefik.enable: "true"
|
||||
traefik.http.routers.embeddings.rule: "Host(`embeddings.example.com`)"
|
||||
traefik.http.routers.embeddings.entrypoints: "websecure"
|
||||
traefik.http.routers.embeddings.tls: "true"
|
||||
traefik.http.routers.embeddings.tls.certresolver: "letsEncrypt"
|
||||
traefik.http.routers.embeddings.service: "embeddings"
|
||||
traefik.http.services.embeddings.loadbalancer.server.port: "80"
|
||||
# TEI enforces the Bearer key itself; Traefik only rate-limits to protect the CPU
|
||||
traefik.http.routers.embeddings.middlewares: "embeddings-rl"
|
||||
traefik.http.middlewares.embeddings-rl.ratelimit.average: "20"
|
||||
traefik.http.middlewares.embeddings-rl.ratelimit.burst: "40"
|
||||
traefik.http.middlewares.embeddings-rl.ratelimit.period: "1s"
|
||||
|
||||
networks:
|
||||
docker_main_net:
|
||||
external: true
|
||||
|
||||
volumes:
|
||||
tei-models:
|
||||
```
|
||||
|
||||
Gitmost settings (**Workspace settings → AI → Embeddings**):
|
||||
|
||||
| Field | Value |
|
||||
|-------------------|---------------------------------------|
|
||||
| Model | `intfloat/multilingual-e5-small` |
|
||||
| Base URL | `https://embeddings.example.com/v1/` |
|
||||
| Embedding API key | your `sk-emb-…` |
|
||||
|
||||
Check it from outside:
|
||||
|
||||
```bash
|
||||
curl -s https://embeddings.example.com/v1/embeddings \
|
||||
-H "Authorization: Bearer sk-emb-REPLACE_WITH_YOUR_KEY" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model":"intfloat/multilingual-e5-small","input":"query: hello"}' \
|
||||
| python3 -c 'import sys,json;print("dims:",len(json.load(sys.stdin)["data"][0]["embedding"]))'
|
||||
# -> dims: 384
|
||||
```
|
||||
|
||||
### Embeddings server notes
|
||||
|
||||
- **Vector dimension is 384.** If this Gitmost was previously embedded with a different model
|
||||
(e.g. `text-embedding-3-large` = 3072-dim), the old pgvector rows won't match the new dimension —
|
||||
clear the existing embeddings / re-index before switching. Gitmost only compares vectors of the
|
||||
same dimension, so mixed-dimension rows are silently ignored rather than searched.
|
||||
- **First start downloads the weights** (hundreds of MB) from `huggingface.co` into the
|
||||
`tei-models` volume; every start after that reads from the volume.
|
||||
- **Pin the version.** Pin the image, and optionally the model: add `--revision <commit-sha>` to
|
||||
`command` (the sha is on the model's page on Hugging Face).
|
||||
- **Air-gapped / no egress:** seed the `tei-models` volume ahead of time and add
|
||||
`environment: [HF_HUB_OFFLINE=1]`.
|
||||
- **GPU:** use the cuda tag of the same release (e.g.
|
||||
`ghcr.io/huggingface/text-embeddings-inference:cuda-1.9`) and start the container with `gpus: all`.
|
||||
|
||||
## Features
|
||||
|
||||
- Real-time collaboration
|
||||
|
||||
+131
@@ -193,6 +193,137 @@ dump/restore, существующий каталог данных переис
|
||||
> неизменным и бэкапьте вместе с базой данных.
|
||||
|
||||
|
||||
## Локальный сервер эмбеддингов
|
||||
|
||||
Семантическому (RAG) поиску AI-агента нужна **модель эмбеддингов**. Вместо оплаты облачного
|
||||
провайдера (например, OpenAI `text-embedding-3-*`) за эмбеддинг каждой страницы можно запустить
|
||||
небольшую open-weights модель у себя через Hugging Face
|
||||
[Text Embeddings Inference](https://github.com/huggingface/text-embeddings-inference) (TEI) — он
|
||||
отдаёт OpenAI-совместимый эндпоинт `/v1/embeddings`. Хороший дефолт — `intfloat/multilingual-e5-small`:
|
||||
многоязычная, 384-мерная, комфортно работает на CPU (~1–2 ГБ RAM, 1–2 vCPU). Пропишите её в
|
||||
**Настройки воркспейса → AI → Эмбеддинги**.
|
||||
|
||||
### Вариант A — локально (та же Docker-сеть, что и Gitmost)
|
||||
|
||||
Запустите TEI контейнером в той же сети, где уже работает Gitmost. Порт наружу не публикуется,
|
||||
поэтому эндпоинт остаётся внутренним и не требует авторизации.
|
||||
|
||||
```yaml
|
||||
services:
|
||||
embeddings:
|
||||
image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.9 # pin version; use a cuda-* tag for GPU
|
||||
container_name: embeddings
|
||||
restart: unless-stopped
|
||||
networks:
|
||||
- gitmost_net # same network Gitmost is on
|
||||
command:
|
||||
- "--model-id"
|
||||
- "intfloat/multilingual-e5-small"
|
||||
- "--auto-truncate" # clamp over-long inputs instead of returning 413
|
||||
volumes:
|
||||
- tei-models:/data # weights are downloaded once and cached here
|
||||
|
||||
networks:
|
||||
gitmost_net:
|
||||
external: true # the network Gitmost already uses
|
||||
|
||||
volumes:
|
||||
tei-models:
|
||||
```
|
||||
|
||||
Настройки Gitmost (**Настройки воркспейса → AI → Эмбеддинги**):
|
||||
|
||||
| Поле | Значение |
|
||||
|-------------------|-----------------------------------|
|
||||
| Model | `intfloat/multilingual-e5-small` |
|
||||
| Base URL | `http://embeddings:80/v1/` |
|
||||
| Embedding API key | — (оставить пустым) |
|
||||
|
||||
> `embeddings` — имя контейнера, Gitmost резолвит его по DNS внутри Docker-сети.
|
||||
> Наружу порт не публикуется, эндпоинт доступен только контейнерам этой сети, поэтому
|
||||
> авторизация не нужна.
|
||||
|
||||
### Вариант B — на отдельном хосте (наружу через Traefik + Let's Encrypt)
|
||||
|
||||
Предполагается, что на хосте уже есть Traefik с ACME-резолвером (в примере ниже — `letsEncrypt`,
|
||||
entrypoint `websecure`, общая сеть `docker_main_net`). Замените домен / сеть / резолвер на свои.
|
||||
|
||||
**DNS:** заведите A-запись `embeddings.example.com` → IP хоста с Traefik (тот же challenge / порт 80,
|
||||
что и у остальных сайтов).
|
||||
|
||||
```yaml
|
||||
services:
|
||||
embeddings:
|
||||
image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.9 # pin version; cuda-* tag for GPU
|
||||
container_name: embeddings
|
||||
restart: unless-stopped
|
||||
networks:
|
||||
- docker_main_net # the network Traefik is attached to
|
||||
command:
|
||||
- "--model-id"
|
||||
- "intfloat/multilingual-e5-small"
|
||||
- "--auto-truncate"
|
||||
- "--api-key"
|
||||
- "sk-emb-REPLACE_WITH_YOUR_KEY"
|
||||
volumes:
|
||||
- tei-models:/data
|
||||
labels:
|
||||
traefik.enable: "true"
|
||||
traefik.http.routers.embeddings.rule: "Host(`embeddings.example.com`)"
|
||||
traefik.http.routers.embeddings.entrypoints: "websecure"
|
||||
traefik.http.routers.embeddings.tls: "true"
|
||||
traefik.http.routers.embeddings.tls.certresolver: "letsEncrypt"
|
||||
traefik.http.routers.embeddings.service: "embeddings"
|
||||
traefik.http.services.embeddings.loadbalancer.server.port: "80"
|
||||
# TEI enforces the Bearer key itself; Traefik only rate-limits to protect the CPU
|
||||
traefik.http.routers.embeddings.middlewares: "embeddings-rl"
|
||||
traefik.http.middlewares.embeddings-rl.ratelimit.average: "20"
|
||||
traefik.http.middlewares.embeddings-rl.ratelimit.burst: "40"
|
||||
traefik.http.middlewares.embeddings-rl.ratelimit.period: "1s"
|
||||
|
||||
networks:
|
||||
docker_main_net:
|
||||
external: true
|
||||
|
||||
volumes:
|
||||
tei-models:
|
||||
```
|
||||
|
||||
Настройки Gitmost (**Настройки воркспейса → AI → Эмбеддинги**):
|
||||
|
||||
| Поле | Значение |
|
||||
|-------------------|---------------------------------------|
|
||||
| Model | `intfloat/multilingual-e5-small` |
|
||||
| Base URL | `https://embeddings.example.com/v1/` |
|
||||
| Embedding API key | ваш `sk-emb-…` |
|
||||
|
||||
Проверка снаружи:
|
||||
|
||||
```bash
|
||||
curl -s https://embeddings.example.com/v1/embeddings \
|
||||
-H "Authorization: Bearer sk-emb-REPLACE_WITH_YOUR_KEY" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model":"intfloat/multilingual-e5-small","input":"query: hello"}' \
|
||||
| python3 -c 'import sys,json;print("dims:",len(json.load(sys.stdin)["data"][0]["embedding"]))'
|
||||
# -> dims: 384
|
||||
```
|
||||
|
||||
### Заметки про сервер эмбеддингов
|
||||
|
||||
- **Размерность вектора — 384.** Если раньше этот Gitmost эмбеддился другой моделью
|
||||
(например, `text-embedding-3-large` = 3072-dim), старые строки в pgvector не совпадут по
|
||||
размерности — очистите существующие эмбеддинги / переиндексируйте перед переключением. Gitmost
|
||||
сравнивает только вектора одной размерности, поэтому строки другой размерности не участвуют в
|
||||
поиске, а не ломают его.
|
||||
- **Первый старт тянет веса** (сотни МБ) с `huggingface.co` в том `tei-models`; дальше — из тома.
|
||||
- **Пин версии.** Пиньте образ, а при желании и модель: добавьте в `command` `--revision <commit-sha>`
|
||||
(sha берётся со страницы модели на Hugging Face).
|
||||
- **Без egress (air-gapped):** засейте том `tei-models` заранее и добавьте
|
||||
`environment: [HF_HUB_OFFLINE=1]`.
|
||||
- **GPU:** возьмите cuda-тег того же релиза (например,
|
||||
`ghcr.io/huggingface/text-embeddings-inference:cuda-1.9`) и запустите контейнер с `gpus: all`.
|
||||
|
||||
|
||||
## Возможности
|
||||
|
||||
- Совместная работа в реальном времени
|
||||
|
||||
@@ -1418,5 +1418,14 @@
|
||||
"The commented text changed since this suggestion was made; it was not applied.": "The commented text changed since this suggestion was made; it was not applied.",
|
||||
"Dismiss": "Dismiss",
|
||||
"Suggestion dismissed": "Suggestion dismissed",
|
||||
"Failed to dismiss suggestion": "Failed to dismiss suggestion"
|
||||
"Failed to dismiss suggestion": "Failed to dismiss suggestion",
|
||||
"Save version": "Save version",
|
||||
"Ctrl+S": "Ctrl+S",
|
||||
"Version saved": "Version saved",
|
||||
"Already saved as the latest version": "Already saved as the latest version",
|
||||
"Agent version": "Agent version",
|
||||
"Boundary": "Boundary",
|
||||
"Autosave": "Autosave",
|
||||
"Only versions": "Only versions",
|
||||
"No saved versions yet.": "No saved versions yet."
|
||||
}
|
||||
|
||||
@@ -1433,5 +1433,14 @@
|
||||
"The commented text changed since this suggestion was made; it was not applied.": "Прокомментированный текст изменился после создания предложения; оно не было применено.",
|
||||
"Dismiss": "Не применять",
|
||||
"Suggestion dismissed": "Предложение отклонено",
|
||||
"Failed to dismiss suggestion": "Не удалось отклонить предложение"
|
||||
"Failed to dismiss suggestion": "Не удалось отклонить предложение",
|
||||
"Save version": "Сохранить версию",
|
||||
"Ctrl+S": "Ctrl+S",
|
||||
"Version saved": "Версия сохранена",
|
||||
"Already saved as the latest version": "Уже сохранено как последняя версия",
|
||||
"Agent version": "Версия агента",
|
||||
"Boundary": "Граница",
|
||||
"Autosave": "Автосейв",
|
||||
"Only versions": "Только версии",
|
||||
"No saved versions yet.": "Пока нет сохранённых версий."
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
@@ -269,17 +272,64 @@ export default function AiChatWindow() {
|
||||
const { data: messageRows, isLoading: messagesLoading } =
|
||||
useAiChatMessagesQuery(
|
||||
activeChatId ?? undefined,
|
||||
// DELIBERATELY DUMB: poll every 2.5s WHILE ARMED, otherwise off. NO error
|
||||
// checks (TanStack resets fetchFailureCount each fetch; the poll must survive
|
||||
// a server restart), NO tail checks, NO cap here — the settled/stalled/idle-cap
|
||||
// semantics all live in ChatThread's FSM, which disarms via onResumeFallback.
|
||||
() => (degradedPoll === true ? 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,
|
||||
);
|
||||
|
||||
// #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(() => {
|
||||
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
|
||||
// resume attempt would only ever 204; gating ChatThread's resume on it avoids a
|
||||
|
||||
@@ -172,9 +172,18 @@ function resetState() {
|
||||
h.state.getRun.mockResolvedValue({ run: null, message: null });
|
||||
}
|
||||
|
||||
// #491: the streaming tail carries a persisted step frontier (metadata.stepsPersisted),
|
||||
// which the tail-only attach reads as `n` in `?anchor=<id>&n=<n>`. Seeded WHOLE now.
|
||||
const streamingTail = () => [
|
||||
row("u1", "user", undefined, "hi"),
|
||||
row("a1", "assistant", "streaming", "partial"),
|
||||
{
|
||||
id: "a1",
|
||||
role: "assistant",
|
||||
content: "partial",
|
||||
status: "streaming",
|
||||
createdAt: "2026-01-01T00:00:00Z",
|
||||
metadata: { stepsPersisted: 2 },
|
||||
} as IAiChatMessageRow,
|
||||
];
|
||||
const settledTail = () => [
|
||||
row("u1", "user", undefined, "hi"),
|
||||
@@ -335,20 +344,24 @@ describe("ChatThread — send now", () => {
|
||||
expect(screen.getAllByLabelText("Remove queued message")).toHaveLength(1);
|
||||
});
|
||||
|
||||
it("Stop then a REAL network-drop finish exits to idle (honor-in-stopping), NOT a false reconnect", () => {
|
||||
it("Stop then a REAL network-drop finish exits to idle (honor-in-stopping), NOT a false reconnect", async () => {
|
||||
// Regression for the disconnect-first reorder: on the STOP path, even a drop-
|
||||
// form finish { isError:true, isDisconnect:true } arriving in `stopping` must be
|
||||
// HONORED (reducer) and exit to idle — it must NOT enter the reconnect ladder.
|
||||
startLocalStreamWithRun(); // live local stream, autonomous
|
||||
fireEvent.click(screen.getByLabelText("Stop")); // STOP_REQUESTED -> stopping
|
||||
h.state.error = { message: "Failed to fetch" };
|
||||
act(() => {
|
||||
// #491: the disconnect re-seeds from persist (async getRun) before dispatching
|
||||
// FINISH_DISCONNECT, which the reducer HONORS in `stopping` -> idle. Flush it.
|
||||
await act(async () => {
|
||||
h.state.onFinish?.({
|
||||
message: { id: "a1", role: "assistant", parts: [] },
|
||||
isAbort: false,
|
||||
isDisconnect: true,
|
||||
isError: true,
|
||||
});
|
||||
await Promise.resolve();
|
||||
await Promise.resolve();
|
||||
});
|
||||
expect(screen.queryByText(/reconnecting/i)).toBeNull();
|
||||
});
|
||||
@@ -803,19 +816,24 @@ describe("ChatThread — resume (attach) machinery", () => {
|
||||
expect(h.state.resumeStream).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("strips the streaming tail from the seed, keeps a user tail whole", () => {
|
||||
it("#491 tail-only: seeds the streaming tail WHOLE (no strip), keeps a user tail whole", () => {
|
||||
renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail() });
|
||||
expect(h.state.seededMessages).toHaveLength(1);
|
||||
// MUTATION-VERIFY: re-introduce the seed-strip and this goes red — the streaming
|
||||
// tail (steps 0..N-1) MUST be seeded so the SDK continuation appends the tail to
|
||||
// the RIGHT message. Both rows (user + assistant) are seeded.
|
||||
expect(h.state.seededMessages).toHaveLength(2);
|
||||
cleanup();
|
||||
resetState();
|
||||
renderThread({ autonomousRunsEnabled: true, initialRows: userTail() });
|
||||
expect(h.state.seededMessages).toHaveLength(1);
|
||||
});
|
||||
|
||||
it("builds the attach URL with expect=live&anchor only for a stripped streaming tail", () => {
|
||||
it("#491 tail-only: builds the attach URL with ?anchor=&n= from the persisted step frontier", () => {
|
||||
renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail() });
|
||||
// n=2 comes from a1's metadata.stepsPersisted (MUTATION-VERIFY: hardcode n=0 and
|
||||
// this fails). No `expect=live` param anymore.
|
||||
expect(h.state.transport!.prepareReconnectToStreamRequest!().api).toBe(
|
||||
"/api/ai-chat/runs/c1/stream?expect=live&anchor=a1",
|
||||
"/api/ai-chat/runs/c1/stream?anchor=a1&n=2",
|
||||
);
|
||||
cleanup();
|
||||
resetState();
|
||||
@@ -839,39 +857,41 @@ describe("ChatThread — resume (attach) machinery", () => {
|
||||
});
|
||||
}
|
||||
|
||||
it("204 on a streaming tail: restore + invalidate + onResumeFallback(true)", async () => {
|
||||
it("204 on a streaming tail: NO restore (row kept) + invalidate + onResumeFallback(true)", async () => {
|
||||
const { onResumeFallback, invalidateSpy } = renderThread({
|
||||
autonomousRunsEnabled: true,
|
||||
initialRows: streamingTail(),
|
||||
});
|
||||
await attachFetch({ status: 204, ok: false });
|
||||
expect(h.state.setMessages).toHaveBeenCalledTimes(1); // restore
|
||||
// #491 tail-only: the anchor row was never stripped, so there is NOTHING to
|
||||
// restore. MUTATION-VERIFY: re-add a restore setMessages here and it goes red.
|
||||
expect(h.state.setMessages).not.toHaveBeenCalled();
|
||||
expect(invalidateSpy).toHaveBeenCalledWith({
|
||||
queryKey: ["ai-chat-messages", "c1"],
|
||||
});
|
||||
expect(onResumeFallback).toHaveBeenCalledWith(true);
|
||||
});
|
||||
|
||||
it("F7 restart-survival: a 500 attach failure restores the row AND arms the poll", async () => {
|
||||
it("F7 restart-survival: a 500 attach failure arms the poll WITHOUT a restore", async () => {
|
||||
const { onResumeFallback, invalidateSpy } = renderThread({
|
||||
autonomousRunsEnabled: true,
|
||||
initialRows: streamingTail(),
|
||||
});
|
||||
await attachFetch({ status: 500, ok: false });
|
||||
expect(h.state.setMessages).toHaveBeenCalledTimes(1);
|
||||
expect(h.state.setMessages).not.toHaveBeenCalled();
|
||||
expect(invalidateSpy).toHaveBeenCalledWith({
|
||||
queryKey: ["ai-chat-messages", "c1"],
|
||||
});
|
||||
expect(onResumeFallback).toHaveBeenCalledWith(true);
|
||||
});
|
||||
|
||||
it("F7 restart-survival: a network throw restores the row AND arms the poll", async () => {
|
||||
it("F7 restart-survival: a network throw arms the poll WITHOUT a restore", async () => {
|
||||
const { onResumeFallback, invalidateSpy } = renderThread({
|
||||
autonomousRunsEnabled: true,
|
||||
initialRows: streamingTail(),
|
||||
});
|
||||
await attachFetch(new Error("network down"), true);
|
||||
expect(h.state.setMessages).toHaveBeenCalledTimes(1);
|
||||
expect(h.state.setMessages).not.toHaveBeenCalled();
|
||||
expect(invalidateSpy).toHaveBeenCalledWith({
|
||||
queryKey: ["ai-chat-messages", "c1"],
|
||||
});
|
||||
@@ -931,7 +951,7 @@ describe("ChatThread — resume (attach) machinery", () => {
|
||||
expect(h.state.sendMessage).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("an empty resumed message (starved replay) restores the row AND arms the poll", () => {
|
||||
it("an empty resumed message (starved replay) arms the poll WITHOUT a restore", () => {
|
||||
h.state.status = "ready";
|
||||
const { onResumeFallback } = renderThread({
|
||||
autonomousRunsEnabled: true,
|
||||
@@ -947,7 +967,9 @@ describe("ChatThread — resume (attach) machinery", () => {
|
||||
isError: false,
|
||||
});
|
||||
});
|
||||
expect(h.state.setMessages).toHaveBeenCalledTimes(1); // restore
|
||||
// #491 tail-only: the seeded steps 0..N-1 are still on screen (the SDK
|
||||
// continuation never wiped them), so there is nothing to restore — just poll.
|
||||
expect(h.state.setMessages).not.toHaveBeenCalled();
|
||||
expect(onResumeFallback).toHaveBeenCalledWith(true); // arm
|
||||
});
|
||||
|
||||
@@ -995,24 +1017,41 @@ describe("ChatThread — live reconnect + stalled", () => {
|
||||
cleanup();
|
||||
});
|
||||
|
||||
// #491: the authoritative PERSISTED assistant row `getRun` projects on a local
|
||||
// disconnect — the re-seed source. Its metadata.stepsPersisted becomes `n`.
|
||||
const persistedAnchor = (steps = 3) => ({
|
||||
run: { id: "run-1", status: "running" },
|
||||
message: {
|
||||
id: "a2",
|
||||
role: "assistant",
|
||||
content: "persisted 0..N-1",
|
||||
status: "streaming",
|
||||
createdAt: "2026-01-01T00:00:00Z",
|
||||
metadata: { stepsPersisted: steps },
|
||||
},
|
||||
});
|
||||
|
||||
// A REAL live SSE drop. ai@6.0.207 emits BOTH { isError:true, isDisconnect:true }
|
||||
// for a network TypeError AND sets useChat `error` — NOT the { isError:false,
|
||||
// error:null } form the old tests fed. This is the form browser QA hit; with the
|
||||
// buggy isError-first routing OR without the errorView render-gate these tests go
|
||||
// red (a real drop surfaces the terminal error banner, masking the reconnect
|
||||
// ladder). MUTATION-VERIFY of disconnect-first + the errorView phase-gate.
|
||||
function disconnect(message: unknown = liveMsg) {
|
||||
// for a network TypeError AND sets useChat `error`. #491: an autonomous local drop
|
||||
// now RE-SEEDS from persist (async getRun) BEFORE entering the reconnect ladder, so
|
||||
// this helper is async and flushes the getRun microtask before returning.
|
||||
async function disconnect(message: unknown = liveMsg) {
|
||||
h.state.error = { message: "Failed to fetch" }; // the SDK sets error on the drop
|
||||
act(() => {
|
||||
await act(async () => {
|
||||
h.state.onFinish?.({
|
||||
message,
|
||||
isAbort: false,
|
||||
isDisconnect: true,
|
||||
isError: true,
|
||||
});
|
||||
// Flush the getRun().then re-seed + the deferred FINISH_DISCONNECT dispatch.
|
||||
await Promise.resolve();
|
||||
await Promise.resolve();
|
||||
});
|
||||
}
|
||||
function renderLive() {
|
||||
// The persisted-anchor read the local disconnect performs to re-seed from persist.
|
||||
h.state.getRun.mockResolvedValue(persistedAnchor());
|
||||
const view = renderThread({
|
||||
autonomousRunsEnabled: true,
|
||||
initialRows: settledTail(),
|
||||
@@ -1032,35 +1071,80 @@ describe("ChatThread — live reconnect + stalled", () => {
|
||||
});
|
||||
}
|
||||
|
||||
it("a live disconnect starts a backoff reconnect (banner + resumeStream after backoff)", () => {
|
||||
it("#491: a live disconnect RE-SEEDS from persist, then backs off to reconnect with ?anchor=&n=", async () => {
|
||||
renderLive();
|
||||
disconnect();
|
||||
await disconnect();
|
||||
// The re-seed read the authoritative persisted row and replaced the live partial.
|
||||
// MUTATION-VERIFY: skip the getRun re-seed (send `n` off the live message) and the
|
||||
// n below no longer matches the PERSISTED stepsPersisted.
|
||||
expect(h.state.getRun).toHaveBeenCalledWith("c1");
|
||||
expect(h.state.setMessages).toHaveBeenCalled(); // re-seeded the store from persist
|
||||
expect(screen.getByText(/reconnecting/i)).toBeTruthy();
|
||||
expect(h.state.resumeStream).not.toHaveBeenCalled();
|
||||
advanceToAttempt(1);
|
||||
expect(h.state.resumeStream).toHaveBeenCalledTimes(1);
|
||||
// n=3 is the PERSISTED row's stepsPersisted (from getRun), NOT the live store.
|
||||
expect(h.state.transport!.prepareReconnectToStreamRequest!().api).toBe(
|
||||
"/api/ai-chat/runs/c1/stream?expect=live&anchor=a2",
|
||||
"/api/ai-chat/runs/c1/stream?anchor=a2&n=3",
|
||||
);
|
||||
});
|
||||
|
||||
it("#488 (browser QA): the reconnect banner is SHOWN, not masked by the residual useChat error", () => {
|
||||
it("#491 regression (#137/#161 dup): getRun REJECT on a live disconnect drops the live partial + nulls the anchor", async () => {
|
||||
// The re-seed source (getRun) FAILS — a flaky-network blip (SSE + getRun both
|
||||
// fail, network recovers in ~1s). The OLD .catch just re-entered the ladder with
|
||||
// NO re-seed and NO filter, so the reconnect could tail-apply the registry's
|
||||
// frames onto the live partial that ALREADY has those steps -> duplicated text.
|
||||
renderLive();
|
||||
h.state.getRun.mockReset();
|
||||
h.state.getRun.mockRejectedValue(new Error("network"));
|
||||
await disconnect(); // live partial = liveMsg (id "a2")
|
||||
expect(h.state.getRun).toHaveBeenCalledWith("c1");
|
||||
// THE GUARANTEE: on the getRun failure the live partial (a2) is FILTERED from the
|
||||
// store, so the reconnect can never tail-apply already-present steps onto it.
|
||||
// MUTATION-VERIFY: revert the .catch fix (enterReconnect only, no filter) and no
|
||||
// setMessages call removes a2 -> this reddens.
|
||||
const removedLivePartial = (
|
||||
h.state.setMessages as unknown as {
|
||||
mock: { calls: [unknown][] };
|
||||
}
|
||||
).mock.calls.some(([updater]) => {
|
||||
if (typeof updater !== "function") return false;
|
||||
const out = (updater as (p: { id: string }[]) => { id: string }[])([
|
||||
{ id: "a2" },
|
||||
{ id: "u1" },
|
||||
]);
|
||||
return !out.some((m) => m.id === "a2");
|
||||
});
|
||||
expect(removedLivePartial).toBe(true);
|
||||
expect(screen.getByText(/reconnecting/i)).toBeTruthy();
|
||||
advanceToAttempt(1);
|
||||
expect(h.state.resumeStream).toHaveBeenCalledTimes(1);
|
||||
// Anchor was nulled -> replay-from-start (no params) / 204 -> poll; never a stale
|
||||
// ?anchor=&n= over the live partial.
|
||||
expect(h.state.transport!.prepareReconnectToStreamRequest!().api).toBe(
|
||||
"/api/ai-chat/runs/c1/stream",
|
||||
);
|
||||
});
|
||||
|
||||
it("#488 (browser QA): the reconnect banner is SHOWN, not masked by the residual useChat error", async () => {
|
||||
// The drop sets useChat `error` (real SDK), and the terminal errorView describes
|
||||
// it ("Lost connection to the server"). The FSM phase-gate must let the
|
||||
// `reconnecting` banner WIN over that residual error. MUTATION-VERIFY: revert the
|
||||
// errorView phase-gate (show errorView whenever error is set) and the terminal
|
||||
// banner masks "reconnecting…" -> red.
|
||||
renderLive();
|
||||
disconnect();
|
||||
await disconnect();
|
||||
expect(h.state.error).not.toBeNull(); // the SDK error IS set during recovery
|
||||
expect(screen.getByText(/reconnecting/i)).toBeTruthy();
|
||||
// The terminal "Lost connection… reload" banner must NOT be showing.
|
||||
expect(screen.queryByText(/reload and try again/i)).toBeNull();
|
||||
});
|
||||
|
||||
it("#488 commit 2: a disconnect BEFORE the first assistant frame reconnects with NO anchor", () => {
|
||||
it("#488 commit 2: a disconnect BEFORE the first assistant frame reconnects with NO anchor", async () => {
|
||||
renderLive();
|
||||
disconnect(null); // no assistant message yet (pre-first-frame break)
|
||||
// No persisted assistant row for a pre-first-frame break -> no anchor.
|
||||
h.state.getRun.mockResolvedValue({ run: null, message: null });
|
||||
await disconnect(null); // no assistant message yet (pre-first-frame break)
|
||||
expect(screen.getByText(/reconnecting/i)).toBeTruthy();
|
||||
expect(
|
||||
screen.queryByText("Connection lost — the answer was interrupted."),
|
||||
@@ -1074,7 +1158,7 @@ describe("ChatThread — live reconnect + stalled", () => {
|
||||
|
||||
it("a live re-attach (2xx) clears the reconnect banner", async () => {
|
||||
renderLive();
|
||||
disconnect();
|
||||
await disconnect();
|
||||
advanceToAttempt(1);
|
||||
await reconnect({ status: 200, ok: true });
|
||||
expect(screen.queryByText(/reconnecting/i)).toBeNull();
|
||||
@@ -1082,7 +1166,7 @@ describe("ChatThread — live reconnect + stalled", () => {
|
||||
|
||||
it("a 204 arms the degraded poll and backs off to the next attempt", async () => {
|
||||
const { onResumeFallback } = renderLive();
|
||||
disconnect();
|
||||
await disconnect();
|
||||
advanceToAttempt(1);
|
||||
expect(h.state.resumeStream).toHaveBeenCalledTimes(1);
|
||||
await reconnect({ status: 204, ok: false });
|
||||
@@ -1094,7 +1178,7 @@ describe("ChatThread — live reconnect + stalled", () => {
|
||||
|
||||
it("exhausts the attempt limit into a manual Retry, which restarts the sequence", async () => {
|
||||
renderLive();
|
||||
disconnect();
|
||||
await disconnect();
|
||||
for (let n = 1; n <= 5; n++) {
|
||||
advanceToAttempt(n);
|
||||
expect(h.state.resumeStream).toHaveBeenCalledTimes(n);
|
||||
@@ -1112,22 +1196,23 @@ describe("ChatThread — live reconnect + stalled", () => {
|
||||
it("#488 commit 3: two breaks in a row produce two reconnect cycles", async () => {
|
||||
renderLive();
|
||||
// First break -> reconnect -> re-attach live.
|
||||
disconnect();
|
||||
await disconnect();
|
||||
advanceToAttempt(1);
|
||||
expect(h.state.resumeStream).toHaveBeenCalledTimes(1);
|
||||
await reconnect({ status: 200, ok: true });
|
||||
expect(screen.queryByText(/reconnecting/i)).toBeNull();
|
||||
// The re-attached observer stream drops AGAIN -> a SECOND reconnect cycle
|
||||
// (the old one-shot !wasResumed gate sent this to silent poll).
|
||||
disconnect();
|
||||
// The re-attached observer (live-follow) stream drops AGAIN -> a SECOND reconnect
|
||||
// cycle. #491: this too re-seeds from persist before re-attaching (never tail-
|
||||
// applies over the live-follow partial).
|
||||
await disconnect();
|
||||
expect(screen.getByText(/reconnecting/i)).toBeTruthy();
|
||||
advanceToAttempt(1);
|
||||
expect(h.state.resumeStream).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it("does NOT reconnect when autonomous runs are disabled", () => {
|
||||
it("does NOT reconnect when autonomous runs are disabled", async () => {
|
||||
renderThread({ autonomousRunsEnabled: false, initialRows: settledTail() });
|
||||
disconnect();
|
||||
await disconnect();
|
||||
expect(screen.queryByText(/reconnecting/i)).toBeNull();
|
||||
expect(
|
||||
screen.getByText("Connection lost — the answer was interrupted."),
|
||||
@@ -1138,7 +1223,7 @@ describe("ChatThread — live reconnect + stalled", () => {
|
||||
|
||||
it("#488 commit 4a: the poll idle cap surfaces a stalled banner + Retry (not silent)", async () => {
|
||||
renderLive();
|
||||
disconnect();
|
||||
await disconnect();
|
||||
advanceToAttempt(1);
|
||||
await reconnect({ status: 204, ok: false }); // arms the poll (reconnecting)
|
||||
// No activity for the whole idle cap -> stalled.
|
||||
|
||||
@@ -42,7 +42,7 @@ import { assistantMessageHasVisibleContent } from "@/features/ai-chat/utils/mess
|
||||
import {
|
||||
isStreamingTail,
|
||||
isSettledAssistantTail,
|
||||
seedRows,
|
||||
stepsPersistedOf,
|
||||
mergeById,
|
||||
} from "@/features/ai-chat/utils/resume-helpers.ts";
|
||||
import { getRun } from "@/features/ai-chat/services/ai-chat-service.ts";
|
||||
@@ -266,25 +266,33 @@ export default function ChatThread({
|
||||
// is NOT one of the lifecycle flags the FSM replaced.
|
||||
const mountedRef = useRef(true);
|
||||
|
||||
// attachStrategy DATA (behind the resumeStream effect; #491 swaps it to tail-only
|
||||
// WITHOUT touching the FSM). The controller is effect-owned (aborted in cleanup,
|
||||
// I5). `stripRef`/`strippedRowRef` are the current full-replay+strip anchor.
|
||||
// attachStrategy DATA (behind the resumeStream effect; #491 tail-only, WITHOUT
|
||||
// touching the FSM). The controller is effect-owned (aborted in cleanup, I5).
|
||||
// `anchorRef` is the PERSISTED assistant row that pins the run (server invariant
|
||||
// 6) and its persisted step frontier N: it feeds `?anchor=<id>&n=<stepsPersisted>`
|
||||
// so the tail-only attach returns frames for steps >= N (the seed carries 0..N-1).
|
||||
// It is NOT a "stripped" row — the seed keeps every row (tail-only replaces the
|
||||
// old full-replay+strip). Null when there is no streaming/active tail to resume.
|
||||
const attachAbortRef = useRef<AbortController | null>(null);
|
||||
const stripRef = useRef(chatId !== null && isStreamingTail(initialRows ?? []));
|
||||
const strippedRowRef = useRef<IAiChatMessageRow | null>(
|
||||
stripRef.current ? (initialRows ?? [])[initialRows!.length - 1] : null,
|
||||
const anchorRef = useRef<{ id: string; stepsPersisted: number } | null>(
|
||||
(() => {
|
||||
if (chatId === null || !isStreamingTail(initialRows ?? [])) return null;
|
||||
const rows = initialRows ?? [];
|
||||
const tail = rows[rows.length - 1];
|
||||
return { id: tail.id, stepsPersisted: stepsPersistedOf(tail) };
|
||||
})(),
|
||||
);
|
||||
// Effect-owned backoff timers (not lifecycle flags): the reconnect ladder and the
|
||||
// stalled inactivity cap. Cleared by the cancelReconnect effect / the cap effect.
|
||||
const reconnectTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
const idleCapTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
|
||||
// #491 tail-only: seed EVERY persisted row unchanged (no strip). The streaming
|
||||
// tail holds steps 0..N-1; the run-stream registry's tail (steps >= N) is APPENDED
|
||||
// to it by the SDK continuation (readUIMessageStream({ message })), so it must be
|
||||
// present in the store for the attach to continue the RIGHT message.
|
||||
const initialMessages = useMemo<UIMessage[]>(
|
||||
() =>
|
||||
seedRows(
|
||||
initialRows ?? [],
|
||||
stripRef.current && autonomousRunsEnabled === true,
|
||||
).map(rowToUiMessage),
|
||||
() => (initialRows ?? []).map(rowToUiMessage),
|
||||
[initialRows],
|
||||
);
|
||||
|
||||
@@ -335,21 +343,16 @@ export default function ChatThread({
|
||||
(eff: RunEffect, epoch: number) => {
|
||||
switch (eff.type) {
|
||||
case "resumeStream": {
|
||||
// The attach GET. Stamp the outcome's generation (I1). A reconnect
|
||||
// attempt filters the pinned live row from the store first (the mount
|
||||
// seed already stripped it), so the live replay's text-start rebuilds it
|
||||
// without duplicating parts (#430).
|
||||
// The attach GET. Stamp the outcome's generation (I1). #491 tail-only: the
|
||||
// store already holds EXACTLY the persisted steps 0..N-1 (the mount seed IS
|
||||
// persist; a reconnect was re-seeded from persist BEFORE FINISH_DISCONNECT
|
||||
// scheduled it — see the onFinish disconnect handler), so there is nothing
|
||||
// to filter here: the SDK continues that seeded message, appending the tail
|
||||
// (steps >= N) without duplicating the pre-drop partial step.
|
||||
pendingAttachEpochRef.current = epoch;
|
||||
// The resumed stream's onFinish is stamped with THIS attach generation
|
||||
// (F1), so a superseded attempt's late finish is dropped.
|
||||
turnEpochRef.current = epoch;
|
||||
if (machineRef.current.phase.name === "reconnecting") {
|
||||
const anchor = strippedRowRef.current;
|
||||
if (anchor)
|
||||
setMessagesRef.current?.((prev) =>
|
||||
prev.filter((m) => m.id !== anchor.id),
|
||||
);
|
||||
}
|
||||
void resumeStreamRef.current?.();
|
||||
break;
|
||||
}
|
||||
@@ -464,18 +467,23 @@ export default function ChatThread({
|
||||
new DefaultChatTransport<UIMessage>({
|
||||
api: "/api/ai-chat/stream",
|
||||
credentials: "include",
|
||||
prepareReconnectToStreamRequest: () => ({
|
||||
// Build the attach URL from the REAL chat id. ?expect=live&anchor=<row id>
|
||||
// only when a streaming tail was stripped: expect=live opts into a
|
||||
// finished-retained replay (safe only because the row is stripped and the
|
||||
// replay rebuilds it), and the anchor pins the replay to OUR run — a
|
||||
// mismatching (newer) run 204s into the restore+poll path instead.
|
||||
api: `/api/ai-chat/runs/${chatIdRef.current}/stream${
|
||||
stripRef.current
|
||||
? `?expect=live&anchor=${strippedRowRef.current!.id}`
|
||||
: ""
|
||||
}`,
|
||||
}),
|
||||
prepareReconnectToStreamRequest: () => {
|
||||
// #491 tail-only attach URL. When there is an anchor (a streaming/active
|
||||
// tail to resume) build `?anchor=<assistantRowId>&n=<stepsPersisted>`: the
|
||||
// server returns the TAIL — a synthetic `start` frame + frames for steps
|
||||
// >= n, then live — which the SDK continuation appends to the seeded row.
|
||||
// The server 204s (-> restore-noop + poll) when it cannot cover the
|
||||
// frontier (overflow/rotation gap) or the anchor mismatches (a newer run).
|
||||
// No anchor (a user tail / pre-first-frame break) => no params.
|
||||
const anchor = anchorRef.current;
|
||||
return {
|
||||
api: `/api/ai-chat/runs/${chatIdRef.current}/stream${
|
||||
anchor
|
||||
? `?anchor=${anchor.id}&n=${anchor.stepsPersisted}`
|
||||
: ""
|
||||
}`,
|
||||
};
|
||||
},
|
||||
fetch: async (input: RequestInfo | URL, init: RequestInit = {}) => {
|
||||
if ((init.method ?? "GET") !== "GET") {
|
||||
// Send path (POST). #488 commit 5: NO client 409 retry ladder anymore
|
||||
@@ -562,8 +570,9 @@ export default function ChatThread({
|
||||
|
||||
// Attach GET outcome -> FSM event. The epoch guard replaces BOTH the one-shot
|
||||
// 204 guard (noStreamHandledRef) and the unmount gate: a stale/superseded or
|
||||
// post-DISPOSE outcome is dropped (I1). For a NONE outcome the attachStrategy
|
||||
// recovery (restore the stripped row + invalidate for a fresh poll) runs first.
|
||||
// post-DISPOSE outcome is dropped (I1). #491 tail-only: on a NONE outcome there is
|
||||
// NOTHING to restore — the anchor row was never stripped from the view (the seed
|
||||
// keeps it) — so we only invalidate for a fresh poll + dispatch the FSM event.
|
||||
const handleAttachOutcome = useCallback(
|
||||
(ep: number, wasReconnecting: boolean, live: boolean) => {
|
||||
if (ep !== epochRef.current) return; // stale generation — drop
|
||||
@@ -575,10 +584,6 @@ export default function ChatThread({
|
||||
);
|
||||
return;
|
||||
}
|
||||
if (strippedRowRef.current)
|
||||
setMessagesRef.current?.((prev) =>
|
||||
mergeById(prev, rowToUiMessage(strippedRowRef.current!)),
|
||||
);
|
||||
queryClient.invalidateQueries({
|
||||
queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatIdRef.current),
|
||||
});
|
||||
@@ -661,56 +666,31 @@ export default function ChatThread({
|
||||
// keeps executing server-side — must win; only a NON-disconnect error (a
|
||||
// provider 500, `{ isError:true, isDisconnect:false }`) is terminal.
|
||||
if (isDisconnect) {
|
||||
if (wasObserver) {
|
||||
// A resumed/attached OBSERVER stream dropped. Recover via the degraded
|
||||
// poll (restore the stripped row only when there is no visible content;
|
||||
// never clobber a fuller on-screen tail, invariant 9). The FSM decides
|
||||
// reconnect-vs-poll from liveFollow (a live-follow drop reconnects again,
|
||||
// #488 commit 3; a mount-resume drop polls).
|
||||
if (mountedRef.current) {
|
||||
const hasVisible = msgHasVisible;
|
||||
if (!hasVisible && strippedRowRef.current)
|
||||
setMessages((prev) =>
|
||||
mergeById(prev, rowToUiMessage(strippedRowRef.current!)),
|
||||
);
|
||||
queryClient.invalidateQueries({
|
||||
queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatIdRef.current),
|
||||
});
|
||||
dispatch({
|
||||
type: "FINISH_DISCONNECT",
|
||||
hasVisibleContent: hasVisible,
|
||||
epoch: stampEpoch,
|
||||
});
|
||||
}
|
||||
if (!mountedRef.current) {
|
||||
setStopNotice(null);
|
||||
return;
|
||||
}
|
||||
// A LOCAL live turn dropped. #488 commit 2: recover by the RUN-FACT, not by
|
||||
// the presence of an assistant message — a setup-phase break (before the
|
||||
// first frame) still leaves a detached run writing to pages. In autonomous
|
||||
// mode a run is active for the whole turn, so seed the run-fact from the
|
||||
// start-metadata runId when known, else a sentinel (the attach GET goes by
|
||||
// chatId, not runId). Pin the assistant row as the strip/anchor when present.
|
||||
if (autonomousRunsEnabled === true && mountedRef.current) {
|
||||
const hasAnchor =
|
||||
message?.role === "assistant" && typeof message.id === "string";
|
||||
if (hasAnchor) {
|
||||
strippedRowRef.current = {
|
||||
id: message.id,
|
||||
role: "assistant",
|
||||
content: "",
|
||||
status: "streaming",
|
||||
createdAt: new Date().toISOString(),
|
||||
metadata: { parts: message.parts },
|
||||
};
|
||||
stripRef.current = true;
|
||||
} else {
|
||||
strippedRowRef.current = null;
|
||||
stripRef.current = false;
|
||||
}
|
||||
// No detached run to recover (legacy, non-autonomous): a plain disconnect —
|
||||
// terminal notice, no reconnect. (An observer only exists in autonomous mode,
|
||||
// so this is always a local turn.)
|
||||
if (autonomousRunsEnabled !== true) {
|
||||
dispatch({
|
||||
type: "RUN_FACT",
|
||||
runFact: { runId: extractRunId(message) ?? "pending" },
|
||||
type: "FINISH_DISCONNECT",
|
||||
hasVisibleContent: false,
|
||||
epoch: stampEpoch,
|
||||
});
|
||||
setStopNotice("disconnect");
|
||||
return;
|
||||
}
|
||||
// A mount-resume OBSERVER (one-shot resume, NOT live-follow) drop falls to
|
||||
// the degraded POLL, which merges by id — it does NOT attach, so there is
|
||||
// nothing to re-seed. #491 tail-only: the anchor row was never removed from
|
||||
// the view (the seed keeps it; the continuation only APPENDED), so nothing to
|
||||
// restore either. The FSM routes this to `polling` (ownership observer,
|
||||
// !liveFollow).
|
||||
if (wasObserver && !machineRef.current.ctx.liveFollow) {
|
||||
queryClient.invalidateQueries({
|
||||
queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatIdRef.current),
|
||||
});
|
||||
dispatch({
|
||||
type: "FINISH_DISCONNECT",
|
||||
@@ -718,14 +698,92 @@ export default function ChatThread({
|
||||
epoch: stampEpoch,
|
||||
});
|
||||
setStopNotice(null);
|
||||
} else {
|
||||
return;
|
||||
}
|
||||
// We will (re-)ENTER THE RECONNECT LADDER (an attach): a LOCAL live turn's
|
||||
// first drop, OR a live-follow observer's SUBSEQUENT drop (#488 commit 3).
|
||||
// #488 commit 2: recover by the RUN-FACT, not by the presence of an assistant
|
||||
// message — a setup-phase break still leaves a detached run writing to pages.
|
||||
//
|
||||
// #491 tail-only (THE crux): the live store holds a PARTIAL step that is AHEAD
|
||||
// of the persisted boundary; tail-applying the reconnect's step frames over it
|
||||
// would DUPLICATE that partial step. So entering reconnecting is ALWAYS via a
|
||||
// RE-SEED FROM PERSIST — never the live store. Fetch the authoritative
|
||||
// persisted assistant row (`getRun` returns the projected `message`), replace
|
||||
// the live partial by id (mergeById -> the store now holds EXACTLY steps
|
||||
// 0..N-1), and set the anchor to `{ id, n = stepsPersisted }`. Only AFTER the
|
||||
// re-seed is applied do we enter the ladder (FINISH_DISCONNECT schedules the
|
||||
// backoff) — so the attach can never tail-apply over the live partial.
|
||||
const cid = chatIdRef.current;
|
||||
// The live-message runId is the run-fact source (the attach GET keys on
|
||||
// chatId, so a sentinel still recovers a setup-phase break).
|
||||
const runId = extractRunId(message ?? undefined) ?? "pending";
|
||||
const enterReconnect = (fact: string): void => {
|
||||
if (!mountedRef.current) return;
|
||||
// Epoch-stamp the run-fact too (I1): the getRun rtt widens the
|
||||
// onFinish->dispatch window, so a concurrent SEND_LOCAL during it must be
|
||||
// able to drop this stale RUN_FACT (else it clobbers the new turn's
|
||||
// runFact.runId). Consistent with the postRun RUN_FACT stamp.
|
||||
dispatch({ type: "RUN_FACT", runFact: { runId: fact }, epoch: stampEpoch });
|
||||
dispatch({
|
||||
type: "FINISH_DISCONNECT",
|
||||
hasVisibleContent: false,
|
||||
hasVisibleContent: msgHasVisible,
|
||||
epoch: stampEpoch,
|
||||
});
|
||||
setStopNotice("disconnect");
|
||||
};
|
||||
// Restore the STRUCTURAL guarantee that the live partial is never the
|
||||
// tail-apply base: drop the live partial from the store by id and null the
|
||||
// anchor, so the reconnect replays from step 0 into a CLEAN store (a full
|
||||
// rebuild) or, past any rotation, 204s -> degraded poll. Used on BOTH the
|
||||
// no-persisted-row and getRun-FAILURE paths — after this there is no path
|
||||
// where the attach tail-applies frames onto a row that already has them
|
||||
// (the #137/#161 duplication class).
|
||||
const dropLivePartialAndReplayFromStart = (): void => {
|
||||
if (message?.role === "assistant" && typeof message.id === "string") {
|
||||
const liveId = message.id;
|
||||
setMessagesRef.current?.((prev) =>
|
||||
prev.filter((m) => m.id !== liveId),
|
||||
);
|
||||
}
|
||||
anchorRef.current = null;
|
||||
};
|
||||
if (cid) {
|
||||
void getRun(cid)
|
||||
.then((res) => {
|
||||
if (!mountedRef.current) return;
|
||||
const persisted = res.message;
|
||||
if (persisted && persisted.role === "assistant") {
|
||||
anchorRef.current = {
|
||||
id: persisted.id,
|
||||
stepsPersisted: stepsPersistedOf(persisted),
|
||||
};
|
||||
// Replace the live partial with the persisted row IN PLACE by id —
|
||||
// the re-seed from persist. The attach's tail (steps >= N) then
|
||||
// appends to a store holding EXACTLY steps 0..N-1: no duplication.
|
||||
setMessages((prev) => mergeById(prev, rowToUiMessage(persisted)));
|
||||
} else {
|
||||
// No persisted assistant row (pre-first-frame break): drop the live
|
||||
// partial + replay from start (no anchor/n) so nothing is duplicated.
|
||||
dropLivePartialAndReplayFromStart();
|
||||
}
|
||||
enterReconnect(res.run?.id ?? runId);
|
||||
})
|
||||
.catch(() => {
|
||||
if (!mountedRef.current) return;
|
||||
// Persist read FAILED: we cannot re-seed from fresh persist, and a
|
||||
// stale mount-time anchor over the live partial would tail-apply
|
||||
// already-present steps -> duplication (a flaky-network blip:
|
||||
// SSE + getRun both fail, network recovers in ~1s, the registry still
|
||||
// covers from the mount frontier). Restore the removed-filter guarantee
|
||||
// instead: drop the live partial + replay from start / 204 -> poll.
|
||||
dropLivePartialAndReplayFromStart();
|
||||
enterReconnect(runId);
|
||||
});
|
||||
} else {
|
||||
dropLivePartialAndReplayFromStart();
|
||||
enterReconnect(runId);
|
||||
}
|
||||
setStopNotice(null);
|
||||
return;
|
||||
}
|
||||
// A NON-disconnect stream error (a provider 500 etc.) -> terminal error banner.
|
||||
@@ -746,11 +804,10 @@ export default function ChatThread({
|
||||
if (mountedRef.current) {
|
||||
const hasVisible = msgHasVisible;
|
||||
if (!hasVisible) {
|
||||
// Starved replay: restore the stripped row + poll to the real terminal.
|
||||
if (strippedRowRef.current)
|
||||
setMessages((prev) =>
|
||||
mergeById(prev, rowToUiMessage(strippedRowRef.current!)),
|
||||
);
|
||||
// Starved replay (the tail carried no new steps). #491 tail-only: the
|
||||
// seeded steps 0..N-1 are still on screen (the SDK continuation never
|
||||
// wiped them — `start` does not reset parts), so there is nothing to
|
||||
// restore; just poll to the real terminal.
|
||||
queryClient.invalidateQueries({
|
||||
queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatIdRef.current),
|
||||
});
|
||||
@@ -863,12 +920,12 @@ export default function ChatThread({
|
||||
const tail = rows[rows.length - 1];
|
||||
if (!tail || tail.role !== "assistant") return;
|
||||
setMessages((prev) => mergeById(prev, rowToUiMessage(tail)));
|
||||
// Anchor-mismatch coherence: a restored stripped row A that a DIFFERENT run's
|
||||
// row B has replaced as the tail would linger as an orphan — settle A from
|
||||
// fresh history so no phantom row survives.
|
||||
const stripped = strippedRowRef.current;
|
||||
if (stripped && stripped.id !== tail.id) {
|
||||
const historical = rows.find((r) => r.id === stripped.id);
|
||||
// Anchor-mismatch coherence: if a DIFFERENT run's row B has replaced our anchor
|
||||
// row A as the tail, A would linger as an orphan — reconcile A by id from FRESH
|
||||
// PERSISTED history (not the pinned live row) so no phantom row survives.
|
||||
const anchor = anchorRef.current;
|
||||
if (anchor && anchor.id !== tail.id) {
|
||||
const historical = rows.find((r) => r.id === anchor.id);
|
||||
if (historical)
|
||||
setMessages((prev) => mergeById(prev, rowToUiMessage(historical)));
|
||||
}
|
||||
|
||||
@@ -57,6 +57,31 @@ 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
|
||||
|
||||
@@ -48,6 +48,7 @@ Legend: **†** = command-transition (bumps `epoch`, I1). Effects in `[…]`.
|
||||
| `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) |
|
||||
@@ -121,8 +122,7 @@ holds. **Pending column: empty.**
|
||||
| 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 | `stripRef` | **data** (attachStrategy) | strip+replay detail; the `resumeStream` effect reads it |
|
||||
| 15 | `strippedRowRef` | **data** (attachStrategy) | the anchor row |
|
||||
| 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 |
|
||||
@@ -151,8 +151,12 @@ message. Sources, in the order they update `ctx.runFact`:
|
||||
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 (future resume-stack iteration #491):** the delta will carry the run field;
|
||||
until then the poll drives to a terminal ROW, dispatched as `POLL_TERMINAL`.
|
||||
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.
|
||||
@@ -178,6 +182,9 @@ Pessimism rule: a stale-but-positive fact PERMITS entering recovery (attach); th
|
||||
/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** (strip+replay today) is behind the `resumeStream` effect; the
|
||||
resume-stack iteration (#491) swaps it to tail-only WITHOUT touching the FSM.
|
||||
- **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.
|
||||
|
||||
@@ -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`.
|
||||
|
||||
@@ -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" });
|
||||
});
|
||||
});
|
||||
@@ -3,11 +3,20 @@ import { atom } from "jotai";
|
||||
// import would drag the whole @tiptap/core engine into the eager graph of every
|
||||
// shell component that reads one of these atoms.
|
||||
import type { Editor } from "@tiptap/core";
|
||||
import type { HocuspocusProvider } from "@hocuspocus/provider";
|
||||
import { PageEditMode } from "@/features/user/types/user.types.ts";
|
||||
import type { DictationUnavailableReason } from "@/features/dictation/dictation-status";
|
||||
|
||||
export const pageEditorAtom = atom<Editor | null>(null);
|
||||
|
||||
// #370 — the active page's collab provider, published by the page editor so the
|
||||
// header menu can emit the "save-version" stateless signal (Cmd+S / button).
|
||||
// Null when the page is read-only / collab isn't connected. A typed initial
|
||||
// value (rather than an explicit generic) keeps jotai's overload resolution on
|
||||
// the writable PrimitiveAtom branch.
|
||||
const initialCollabProvider: HocuspocusProvider | null = null;
|
||||
export const collabProviderAtom = atom(initialCollabProvider);
|
||||
|
||||
export const titleEditorAtom = atom<Editor | null>(null);
|
||||
|
||||
export const readOnlyEditorAtom = atom<Editor | null>(null);
|
||||
|
||||
@@ -31,11 +31,18 @@ import { useAtom, useAtomValue, useSetAtom } from "jotai";
|
||||
import useCollaborationUrl from "@/features/editor/hooks/use-collaboration-url";
|
||||
import { currentUserAtom } from "@/features/user/atoms/current-user-atom";
|
||||
import {
|
||||
collabProviderAtom,
|
||||
currentPageEditModeAtom,
|
||||
dictationAvailabilityAtom,
|
||||
pageEditorAtom,
|
||||
yjsConnectionStatusAtom,
|
||||
} from "@/features/editor/atoms/editor-atoms";
|
||||
import { notifications } from "@mantine/notifications";
|
||||
import {
|
||||
VERSION_SAVED_MESSAGE_TYPE,
|
||||
type VersionSavedMessage,
|
||||
saveVersionPending,
|
||||
} from "@/features/page-history/version-messages";
|
||||
import { asideStateAtom } from "@/components/layouts/global/hooks/atoms/sidebar-atom";
|
||||
import {
|
||||
activeCommentIdAtom,
|
||||
@@ -124,6 +131,7 @@ export default function PageEditor({
|
||||
|
||||
const [currentUser] = useAtom(currentUserAtom);
|
||||
const [, setEditor] = useAtom(pageEditorAtom);
|
||||
const setCollabProvider = useSetAtom(collabProviderAtom);
|
||||
const [, setAsideState] = useAtom(asideStateAtom);
|
||||
const [, setActiveCommentId] = useAtom(activeCommentIdAtom);
|
||||
const [showCommentPopup, setShowCommentPopup] = useAtom(showCommentPopupAtom);
|
||||
@@ -181,6 +189,24 @@ export default function PageEditor({
|
||||
const onStatelessHandler = ({ payload }: onStatelessParameters) => {
|
||||
try {
|
||||
const message = JSON.parse(payload);
|
||||
// #370 — a version was saved somewhere; live-refresh the history panel
|
||||
// on every client. Only the client that pressed Save (tracked by the
|
||||
// module-level flag) shows the confirmation toast.
|
||||
if (message?.type === VERSION_SAVED_MESSAGE_TYPE) {
|
||||
const versionMsg = message as VersionSavedMessage;
|
||||
queryClient.invalidateQueries({
|
||||
queryKey: ["page-history-list"],
|
||||
});
|
||||
if (saveVersionPending.current) {
|
||||
saveVersionPending.current = false;
|
||||
notifications.show({
|
||||
message: versionMsg.alreadySaved
|
||||
? t("Already saved as the latest version")
|
||||
: t("Version saved"),
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (message?.type !== "page.updated" || !message.updatedAt) return;
|
||||
const pageData = queryClient.getQueryData<IPage>(["pages", slugId]);
|
||||
if (pageData) {
|
||||
@@ -238,12 +264,16 @@ export default function PageEditor({
|
||||
|
||||
local.on("synced", onLocalSyncedHandler);
|
||||
providersRef.current = { socket, local, remote };
|
||||
// #370 — publish the provider so the header menu can emit save-version.
|
||||
setCollabProvider(remote);
|
||||
setProvidersReady(true);
|
||||
} else {
|
||||
setCollabProvider(providersRef.current.remote);
|
||||
setProvidersReady(true);
|
||||
}
|
||||
// Only destroy on final unmount
|
||||
return () => {
|
||||
setCollabProvider(null);
|
||||
providersRef.current?.socket.destroy();
|
||||
providersRef.current?.remote.destroy();
|
||||
providersRef.current?.local.destroy();
|
||||
|
||||
@@ -1,4 +1,11 @@
|
||||
import { Text, Group, UnstyledButton, Avatar, Tooltip } from "@mantine/core";
|
||||
import {
|
||||
Text,
|
||||
Group,
|
||||
UnstyledButton,
|
||||
Avatar,
|
||||
Tooltip,
|
||||
Badge,
|
||||
} from "@mantine/core";
|
||||
import { CustomAvatar } from "@/components/ui/custom-avatar.tsx";
|
||||
import { AgentAvatarStack } from "@/components/ui/agent-avatar-stack.tsx";
|
||||
import { formattedDate } from "@/lib/time";
|
||||
@@ -7,36 +14,59 @@ import clsx from "clsx";
|
||||
import { IPageHistory } from "@/features/page-history/types/page.types";
|
||||
import { memo, useCallback } from "react";
|
||||
import { useSetAtom } from "jotai";
|
||||
import { useTranslation } from "react-i18next";
|
||||
import { historyAtoms } from "@/features/page-history/atoms/history-atoms.ts";
|
||||
|
||||
const MAX_VISIBLE_AVATARS = 5;
|
||||
|
||||
/**
|
||||
* #370 — map a snapshot's intentionality tier to its badge. `version: true`
|
||||
* marks the intentional points (manual / agent); autosaves (boundary / idle /
|
||||
* legacy null) are non-versions and get dimmed in the list.
|
||||
*/
|
||||
type HistoryKindMeta = { labelKey: string; color: string; version: boolean };
|
||||
export function historyKindMeta(kind?: string | null): HistoryKindMeta {
|
||||
switch (kind) {
|
||||
case "manual":
|
||||
return { labelKey: "Saved", color: "blue", version: true };
|
||||
case "agent":
|
||||
return { labelKey: "Agent version", color: "violet", version: true };
|
||||
case "boundary":
|
||||
return { labelKey: "Boundary", color: "gray", version: false };
|
||||
default: // "idle" | null | undefined (legacy autosave)
|
||||
return { labelKey: "Autosave", color: "gray", version: false };
|
||||
}
|
||||
}
|
||||
|
||||
interface HistoryItemProps {
|
||||
historyItem: IPageHistory;
|
||||
index: number;
|
||||
onSelect: (id: string, index: number) => void;
|
||||
onHover?: (id: string, index: number) => void;
|
||||
// The previous snapshot for diff/restore is resolved by id from the FULL list
|
||||
// in the parent (resolvePrevSnapshotId), so the item only needs to report its
|
||||
// own id — never a list index (which would be the filtered-view index).
|
||||
onSelect: (id: string) => void;
|
||||
onHover?: (id: string) => void;
|
||||
onHoverEnd?: () => void;
|
||||
isActive: boolean;
|
||||
}
|
||||
|
||||
const HistoryItem = memo(function HistoryItem({
|
||||
historyItem,
|
||||
index,
|
||||
onSelect,
|
||||
onHover,
|
||||
onHoverEnd,
|
||||
isActive,
|
||||
}: HistoryItemProps) {
|
||||
const setHistoryModalOpen = useSetAtom(historyAtoms);
|
||||
const { t } = useTranslation();
|
||||
const kindMeta = historyKindMeta(historyItem.kind);
|
||||
|
||||
const handleClick = useCallback(() => {
|
||||
onSelect(historyItem.id, index);
|
||||
}, [onSelect, historyItem.id, index]);
|
||||
onSelect(historyItem.id);
|
||||
}, [onSelect, historyItem.id]);
|
||||
|
||||
const handleMouseEnter = useCallback(() => {
|
||||
onHover?.(historyItem.id, index);
|
||||
}, [onHover, historyItem.id, index]);
|
||||
onHover?.(historyItem.id);
|
||||
}, [onHover, historyItem.id]);
|
||||
|
||||
const contributors = historyItem.contributors;
|
||||
const hasContributors = contributors && contributors.length > 0;
|
||||
@@ -49,8 +79,20 @@ const HistoryItem = memo(function HistoryItem({
|
||||
onMouseEnter={handleMouseEnter}
|
||||
onMouseLeave={onHoverEnd}
|
||||
className={clsx(classes.history, { [classes.active]: isActive })}
|
||||
// #370 — dim autosnapshots so intentional versions stand out.
|
||||
style={{ opacity: kindMeta.version ? 1 : 0.55 }}
|
||||
>
|
||||
<Text size="sm">{formattedDate(new Date(historyItem.createdAt))}</Text>
|
||||
<Group gap={6} wrap="nowrap" justify="space-between">
|
||||
<Text size="sm">{formattedDate(new Date(historyItem.createdAt))}</Text>
|
||||
<Badge
|
||||
size="xs"
|
||||
radius="sm"
|
||||
variant={kindMeta.version ? "filled" : "light"}
|
||||
color={kindMeta.color}
|
||||
>
|
||||
{t(kindMeta.labelKey)}
|
||||
</Badge>
|
||||
</Group>
|
||||
|
||||
<Group gap={6} wrap="nowrap" mt={4}>
|
||||
{hasContributors ? (
|
||||
|
||||
@@ -2,14 +2,16 @@ import {
|
||||
usePageHistoryListQuery,
|
||||
prefetchPageHistory,
|
||||
} from "@/features/page-history/queries/page-history-query";
|
||||
import HistoryItem from "@/features/page-history/components/history-item";
|
||||
import HistoryItem, {
|
||||
historyKindMeta,
|
||||
} from "@/features/page-history/components/history-item";
|
||||
import {
|
||||
activeHistoryIdAtom,
|
||||
activeHistoryPrevIdAtom,
|
||||
historyAtoms,
|
||||
} from "@/features/page-history/atoms/history-atoms";
|
||||
import { useAtom, useSetAtom } from "jotai";
|
||||
import { useCallback, useEffect, useMemo, useRef } from "react";
|
||||
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
import {
|
||||
Button,
|
||||
ScrollArea,
|
||||
@@ -17,9 +19,12 @@ import {
|
||||
Divider,
|
||||
Loader,
|
||||
Center,
|
||||
Switch,
|
||||
Text,
|
||||
} from "@mantine/core";
|
||||
import { useTranslation } from "react-i18next";
|
||||
import { useHistoryRestore } from "@/features/page-history/hooks";
|
||||
import { resolvePrevSnapshotId } from "@/features/page-history/utils/resolve-prev-snapshot";
|
||||
|
||||
const PREFETCH_DELAY_MS = 150;
|
||||
|
||||
@@ -47,6 +52,22 @@ function HistoryList({ pageId }: Props) {
|
||||
[pageHistoryData],
|
||||
);
|
||||
|
||||
// #370 — "only versions" filter: hide autosnapshots (idle/boundary/legacy
|
||||
// null), keep only intentional points (manual/agent). Filtering is over the
|
||||
// already-loaded pages; the diff/restore still targets the true previous
|
||||
// snapshot, so items carry their index within the FULL list.
|
||||
const [onlyVersions, setOnlyVersions] = useState(false);
|
||||
// Reuse historyKindMeta().version — the SAME predicate the badge (HistoryItem)
|
||||
// uses to mark intentional points — so the "Only versions" filter and the badge
|
||||
// can never drift apart when a future intentional kind is added.
|
||||
const visibleItems = useMemo(
|
||||
() =>
|
||||
onlyVersions
|
||||
? historyItems.filter((item) => historyKindMeta(item.kind).version)
|
||||
: historyItems,
|
||||
[historyItems, onlyVersions],
|
||||
);
|
||||
|
||||
const loadMoreRef = useRef<HTMLDivElement>(null);
|
||||
const prefetchTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
|
||||
@@ -60,11 +81,13 @@ function HistoryList({ pageId }: Props) {
|
||||
}, []);
|
||||
|
||||
const handleHover = useCallback(
|
||||
(historyId: string, index: number) => {
|
||||
(historyId: string) => {
|
||||
clearPrefetchTimeout();
|
||||
prefetchTimeoutRef.current = setTimeout(() => {
|
||||
prefetchPageHistory(historyId);
|
||||
const prevId = historyItems[index + 1]?.id;
|
||||
// The true previous snapshot in the FULL list (not the previous visible
|
||||
// one under the "only versions" filter).
|
||||
const prevId = resolvePrevSnapshotId(historyItems, historyId);
|
||||
if (prevId) {
|
||||
prefetchPageHistory(prevId);
|
||||
}
|
||||
@@ -78,9 +101,11 @@ function HistoryList({ pageId }: Props) {
|
||||
}, [clearPrefetchTimeout]);
|
||||
|
||||
const handleSelect = useCallback(
|
||||
(id: string, index: number) => {
|
||||
(id: string) => {
|
||||
setActiveHistoryId(id);
|
||||
setActiveHistoryPrevId(historyItems[index + 1]?.id ?? "");
|
||||
// Baseline = true previous snapshot in the FULL list, so the "only
|
||||
// versions" filter never diffs/restores against the wrong item.
|
||||
setActiveHistoryPrevId(resolvePrevSnapshotId(historyItems, id));
|
||||
},
|
||||
[historyItems, setActiveHistoryId, setActiveHistoryPrevId],
|
||||
);
|
||||
@@ -128,12 +153,27 @@ function HistoryList({ pageId }: Props) {
|
||||
|
||||
return (
|
||||
<div>
|
||||
<Group px="xs" py={6} justify="flex-end">
|
||||
<Switch
|
||||
size="xs"
|
||||
checked={onlyVersions}
|
||||
onChange={(e) => setOnlyVersions(e.currentTarget.checked)}
|
||||
label={t("Only versions")}
|
||||
/>
|
||||
</Group>
|
||||
|
||||
<ScrollArea h={620} w="100%" type="scroll" scrollbarSize={5}>
|
||||
{historyItems.map((historyItem, index) => (
|
||||
{onlyVersions && visibleItems.length === 0 && (
|
||||
<Center py="md">
|
||||
<Text size="sm" c="dimmed">
|
||||
{t("No saved versions yet.")}
|
||||
</Text>
|
||||
</Center>
|
||||
)}
|
||||
{visibleItems.map((historyItem) => (
|
||||
<HistoryItem
|
||||
key={historyItem.id}
|
||||
historyItem={historyItem}
|
||||
index={index}
|
||||
onSelect={handleSelect}
|
||||
onHover={handleHover}
|
||||
onHoverEnd={clearPrefetchTimeout}
|
||||
|
||||
@@ -24,6 +24,10 @@ export interface IPageHistory {
|
||||
updatedAt: string;
|
||||
lastUpdatedBy: IPageHistoryUser;
|
||||
contributors?: IPageHistoryUser[];
|
||||
// #370 — intentionality tier: 'manual'/'agent' are versions (intentional
|
||||
// points), 'idle'/'boundary' are autosnapshots; null/undefined = legacy
|
||||
// autosave. Derived server-side, drives the history badge + "versions" filter.
|
||||
kind?: "manual" | "agent" | "idle" | "boundary" | null;
|
||||
// Provenance markers copied off the page row when the snapshot was saved.
|
||||
// `'agent'` marks a version written by the AI agent; `lastUpdatedAiChatId`
|
||||
// (when present) deep-links to the chat that produced the edit.
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { resolvePrevSnapshotId } from "./resolve-prev-snapshot";
|
||||
|
||||
// #370 F4 — the risky client path: with the "only versions" filter active, diff
|
||||
// and restore must still baseline against the TRUE previous snapshot in the FULL
|
||||
// list, never the previous VISIBLE version (which would skip the autosnapshots
|
||||
// between two versions). These pin that the resolution is by FULL-list order.
|
||||
describe("resolvePrevSnapshotId", () => {
|
||||
// Newest-first, as the history list stores it: a version, then two autosaves,
|
||||
// then an older version.
|
||||
const full = [
|
||||
{ id: "v2", kind: "manual" },
|
||||
{ id: "a2", kind: "idle" },
|
||||
{ id: "a1", kind: "boundary" },
|
||||
{ id: "v1", kind: "manual" },
|
||||
{ id: "a0", kind: null },
|
||||
];
|
||||
|
||||
it("returns the immediate FULL-list successor, not the previous visible version", () => {
|
||||
// Selecting v2 while filtered to versions-only must baseline against a2 (the
|
||||
// real chronological predecessor), NOT v1 (the previous visible version).
|
||||
expect(resolvePrevSnapshotId(full, "v2")).toBe("a2");
|
||||
});
|
||||
|
||||
it("resolves an autosnapshot's predecessor by full-list order", () => {
|
||||
expect(resolvePrevSnapshotId(full, "a1")).toBe("v1");
|
||||
});
|
||||
|
||||
it("returns '' for the oldest item (no predecessor)", () => {
|
||||
expect(resolvePrevSnapshotId(full, "a0")).toBe("");
|
||||
});
|
||||
|
||||
it("returns '' for an id not in the list", () => {
|
||||
expect(resolvePrevSnapshotId(full, "missing")).toBe("");
|
||||
});
|
||||
|
||||
it("does not depend on a filtered subset — same result whatever is visible", () => {
|
||||
// The helper only ever sees the full list; a filtered view cannot change the
|
||||
// baseline it computes.
|
||||
expect(resolvePrevSnapshotId(full, "v1")).toBe("a0");
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,22 @@
|
||||
/**
|
||||
* #370 — resolve the TRUE previous snapshot for a history item.
|
||||
*
|
||||
* The history panel can be filtered to "only versions" (manual/agent), but diff
|
||||
* and restore must always compare against the immediately-preceding snapshot in
|
||||
* the FULL, unfiltered list — NOT the previous VISIBLE item. Comparing against
|
||||
* the previous visible version would silently skip the autosnapshots between two
|
||||
* versions and diff/restore the wrong baseline.
|
||||
*
|
||||
* Given the full (newest-first) list and an item id, this returns the id of the
|
||||
* item right after it in the full list (its chronological predecessor), or "" if
|
||||
* it is the oldest / not found. Pure and list-order-preserving so it can be unit
|
||||
* tested without mounting the component.
|
||||
*/
|
||||
export function resolvePrevSnapshotId(
|
||||
fullItems: ReadonlyArray<{ id: string }>,
|
||||
id: string,
|
||||
): string {
|
||||
const index = fullItems.findIndex((item) => item.id === id);
|
||||
if (index === -1) return "";
|
||||
return fullItems[index + 1]?.id ?? "";
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
/**
|
||||
* #370 — page-version stateless wire formats. Kept in one place so the client
|
||||
* emitter (Save hotkey / button) and the client listener (page-editor) agree
|
||||
* with the server (PersistenceExtension) on the message shapes.
|
||||
*/
|
||||
|
||||
/** Client → server: "save a version now". The server derives the tier
|
||||
* (manual/agent) from the signed connection actor, never from this payload. */
|
||||
export const SAVE_VERSION_MESSAGE_TYPE = "save-version";
|
||||
|
||||
/** Server → all clients: a version was saved (or promoted / already existed). */
|
||||
export const VERSION_SAVED_MESSAGE_TYPE = "version.saved";
|
||||
|
||||
export interface VersionSavedMessage {
|
||||
type: typeof VERSION_SAVED_MESSAGE_TYPE;
|
||||
historyId: string;
|
||||
kind: "manual" | "agent";
|
||||
/** True when the latest snapshot was already a manual version (a no-op save). */
|
||||
alreadySaved: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Cross-component coordination flag so only the client that pressed Save shows
|
||||
* the confirmation toast, while every other client silently refreshes its
|
||||
* history panel on the broadcast. A module-level ref avoids stale-closure
|
||||
* pitfalls in the editor's long-lived stateless handler.
|
||||
*/
|
||||
export const saveVersionPending = { current: false };
|
||||
@@ -3,6 +3,7 @@ import {
|
||||
IconArrowRight,
|
||||
IconArrowsHorizontal,
|
||||
IconClockHour4,
|
||||
IconDeviceFloppy,
|
||||
IconDots,
|
||||
IconEye,
|
||||
IconEyeOff,
|
||||
@@ -17,7 +18,7 @@ import {
|
||||
IconTrash,
|
||||
IconWifiOff,
|
||||
} from "@tabler/icons-react";
|
||||
import React, { useEffect, useRef, useState } from "react";
|
||||
import React, { useCallback, useEffect, useRef, useState } from "react";
|
||||
import { useAsideTriggerProps } from "@/hooks/use-toggle-aside.tsx";
|
||||
import { useAtom, useAtomValue } from "jotai";
|
||||
import { historyAtoms } from "@/features/page-history/atoms/history-atoms.ts";
|
||||
@@ -39,9 +40,14 @@ import { Trans, useTranslation } from "react-i18next";
|
||||
import ExportModal from "@/components/common/export-modal";
|
||||
import { convertProseMirrorToMarkdown } from "@docmost/prosemirror-markdown/browser";
|
||||
import {
|
||||
collabProviderAtom,
|
||||
pageEditorAtom,
|
||||
yjsConnectionStatusAtom,
|
||||
} from "@/features/editor/atoms/editor-atoms.ts";
|
||||
import {
|
||||
SAVE_VERSION_MESSAGE_TYPE,
|
||||
saveVersionPending,
|
||||
} from "@/features/page-history/version-messages.ts";
|
||||
import { formattedDate } from "@/lib/time.ts";
|
||||
import { PageEditModeToggle } from "@/features/user/components/page-state-pref.tsx";
|
||||
import MovePageModal from "@/features/page/components/move-page-modal.tsx";
|
||||
@@ -72,9 +78,34 @@ export default function PageHeaderMenu({ readOnly }: PageHeaderMenuProps) {
|
||||
});
|
||||
const isDeleted = !!page?.deletedAt;
|
||||
const [workspace] = useAtom(workspaceAtom);
|
||||
const collabProvider = useAtomValue(collabProviderAtom);
|
||||
// Community public-sharing entry point (replaces the removed EE PageShareModal)
|
||||
const workspaceSharingDisabled = workspace?.settings?.sharing?.disabled === true;
|
||||
|
||||
// #370 — explicit "save a version" (Cmd+S / Save button). One path for the
|
||||
// human; the server derives the tier from the signed actor. Readers can't save
|
||||
// (the button is hidden and the collab connection is read-only server-side).
|
||||
const handleSaveVersion = useCallback(() => {
|
||||
if (readOnly || !collabProvider) return;
|
||||
// Flag this client as the initiator so only it shows the confirmation toast;
|
||||
// a safety timeout clears it if no broadcast comes back (e.g. offline).
|
||||
saveVersionPending.current = true;
|
||||
window.setTimeout(() => {
|
||||
saveVersionPending.current = false;
|
||||
}, 5000);
|
||||
collabProvider.sendStateless(
|
||||
JSON.stringify({ type: SAVE_VERSION_MESSAGE_TYPE }),
|
||||
);
|
||||
}, [readOnly, collabProvider]);
|
||||
|
||||
// mod+S must also block the browser's "Save page" dialog. `triggerOnContent-
|
||||
// Editable` + empty ignore-list so it fires while typing in the editor/title.
|
||||
useHotkeys(
|
||||
[["mod+S", handleSaveVersion, { preventDefault: true }]],
|
||||
[],
|
||||
true,
|
||||
);
|
||||
|
||||
useHotkeys(
|
||||
[
|
||||
[
|
||||
@@ -133,15 +164,16 @@ export default function PageHeaderMenu({ readOnly }: PageHeaderMenuProps) {
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
|
||||
<PageActionMenu readOnly={readOnly} />
|
||||
<PageActionMenu readOnly={readOnly} onSaveVersion={handleSaveVersion} />
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
interface PageActionMenuProps {
|
||||
readOnly?: boolean;
|
||||
onSaveVersion?: () => void;
|
||||
}
|
||||
function PageActionMenu({ readOnly }: PageActionMenuProps) {
|
||||
function PageActionMenu({ readOnly, onSaveVersion }: PageActionMenuProps) {
|
||||
const { t } = useTranslation();
|
||||
const [, setHistoryModalOpen] = useAtom(historyAtoms);
|
||||
const clipboard = useClipboard({ timeout: 500 });
|
||||
@@ -303,6 +335,20 @@ function PageActionMenu({ readOnly }: PageActionMenuProps) {
|
||||
</Group>
|
||||
</Menu.Item>
|
||||
|
||||
{!readOnly && (
|
||||
<Menu.Item
|
||||
leftSection={<IconDeviceFloppy size={16} />}
|
||||
onClick={onSaveVersion}
|
||||
rightSection={
|
||||
<Text size="xs" c="dimmed">
|
||||
{t("Ctrl+S")}
|
||||
</Text>
|
||||
}
|
||||
>
|
||||
{t("Save version")}
|
||||
</Menu.Item>
|
||||
)}
|
||||
|
||||
<Menu.Item
|
||||
leftSection={<IconHistory size={16} />}
|
||||
onClick={openHistoryModal}
|
||||
|
||||
@@ -16,6 +16,7 @@ import { TransclusionService } from '../core/page/transclusion/transclusion.serv
|
||||
import { TransclusionModule } from '../core/page/transclusion/transclusion.module';
|
||||
import { StorageModule } from '../integrations/storage/storage.module';
|
||||
import { EnvironmentModule } from '../integrations/environment/environment.module';
|
||||
import { ApiKeyModule } from '../core/api-key/api-key.module';
|
||||
|
||||
@Module({
|
||||
providers: [
|
||||
@@ -31,6 +32,7 @@ import { EnvironmentModule } from '../integrations/environment/environment.modul
|
||||
exports: [CollaborationGateway],
|
||||
imports: [
|
||||
TokenModule,
|
||||
ApiKeyModule,
|
||||
WatcherModule,
|
||||
StorageModule.forRootAsync({
|
||||
imports: [EnvironmentModule],
|
||||
|
||||
@@ -1,10 +1,36 @@
|
||||
export const HISTORY_INTERVAL = 5 * 60 * 1000;
|
||||
export const HISTORY_FAST_INTERVAL = 60 * 1000;
|
||||
export const HISTORY_FAST_THRESHOLD = 5 * 60 * 1000;
|
||||
|
||||
// #348 — debounce window for the per-page RAG re-embed job. Repeated saves
|
||||
// within this window collapse to a single delayed job (coalesced by a stable
|
||||
// jobId), so active editing does not pile up expensive re-embeds (external API
|
||||
// + page_embeddings rewrite, concurrency 1). The worker reads the CURRENT page
|
||||
// state at run time, so the last content within the window wins.
|
||||
export const EMBED_DEBOUNCE_MS = 30 * 1000;
|
||||
|
||||
/**
|
||||
* #370 — page-history intentionality tiers. Domain of `page_history.kind`.
|
||||
* - 'manual' / 'agent' → Tier 1 versions (intentional points)
|
||||
* - 'idle' / 'boundary' → Tier 0 autosnapshots (safety net)
|
||||
* A legacy `null` kind is treated as an autosave.
|
||||
*/
|
||||
export type PageHistoryKind = 'manual' | 'agent' | 'idle' | 'boundary';
|
||||
|
||||
/**
|
||||
* #370 — trailing idle-flush windows. A page's pending idle snapshot is
|
||||
* re-armed on every store and fires this long after edits go quiet, so a burst
|
||||
* of edits collapses into a single autosnapshot instead of one-per-store. Human
|
||||
* sessions are noisier and less risky, so they flush less often than the agent.
|
||||
*/
|
||||
export const IDLE_INTERVAL_USER = 60 * 60 * 1000; // 60m
|
||||
export const IDLE_INTERVAL_AGENT = 15 * 60 * 1000; // 15m
|
||||
|
||||
/**
|
||||
* #370 — max-wait ceiling for the idle flush. Pure trailing debounce starves the
|
||||
* safety net: hocuspocus stores at least every ~45s, so a CONTINUOUS editing
|
||||
* session would re-arm the trailing timer forever and never take an idle
|
||||
* snapshot until edits finally go quiet (up to IDLE_INTERVAL_USER = 60m). This
|
||||
* ceiling bounds the actual wait from the FIRST edit of a burst, so an idle
|
||||
* snapshot fires at least this often during a long unbroken session — restoring
|
||||
* a recovery point cadence closer to the old heuristic without one-per-store
|
||||
* noise. Mirrors hocuspocus's own maxDebounce idea.
|
||||
*/
|
||||
export const IDLE_MAX_WAIT_USER = 10 * 60 * 1000; // 10m
|
||||
export const IDLE_MAX_WAIT_AGENT = 5 * 60 * 1000; // 5m
|
||||
|
||||
@@ -52,6 +52,7 @@ describe('AuthenticationExtension.onAuthenticate', () => {
|
||||
let pageRepo: { findById: jest.Mock };
|
||||
let spaceMemberRepo: { getUserSpaceRoles: jest.Mock };
|
||||
let pagePermissionRepo: { canUserEditPage: jest.Mock };
|
||||
let apiKeyService: { validate: jest.Mock };
|
||||
|
||||
// Build the hocuspocus onAuthenticate payload. connectionConfig.readOnly
|
||||
// starts false; the extension flips it to true on a read-only downgrade.
|
||||
@@ -79,12 +80,15 @@ describe('AuthenticationExtension.onAuthenticate', () => {
|
||||
}),
|
||||
};
|
||||
|
||||
apiKeyService = { validate: jest.fn().mockResolvedValue({ user: {}, workspace: {} }) };
|
||||
|
||||
ext = new AuthenticationExtension(
|
||||
tokenService as any,
|
||||
userRepo as any,
|
||||
pageRepo as any,
|
||||
spaceMemberRepo as any,
|
||||
pagePermissionRepo as any,
|
||||
apiKeyService as any,
|
||||
);
|
||||
// Silence the extension's logger (it warns/debugs on denial branches).
|
||||
jest.spyOn(ext['logger'], 'warn').mockImplementation(() => undefined);
|
||||
@@ -231,4 +235,73 @@ describe('AuthenticationExtension.onAuthenticate', () => {
|
||||
// No internal ai_chats row for an MCP/service-account collab edit → null.
|
||||
expect(ctx.aiChatId).toBeNull();
|
||||
});
|
||||
|
||||
// --- #501: api-key laundering guard (fail-closed discriminator) ----------
|
||||
describe('api-key laundering guard', () => {
|
||||
it('api_key principal → row-checks the key on connect (valid key proceeds)', async () => {
|
||||
tokenService.verifyJwt.mockResolvedValue(
|
||||
buildJwt({ principal: 'api_key', apiKeyId: 'key-1' }),
|
||||
);
|
||||
const data = buildData();
|
||||
await ext.onAuthenticate(data as any);
|
||||
|
||||
expect(apiKeyService.validate).toHaveBeenCalledTimes(1);
|
||||
expect(apiKeyService.validate).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ apiKeyId: 'key-1', type: JwtType.API_KEY }),
|
||||
);
|
||||
});
|
||||
|
||||
it('REVOKED api_key → Unauthorized on connect, BEFORE any page/user lookup', async () => {
|
||||
tokenService.verifyJwt.mockResolvedValue(
|
||||
buildJwt({ principal: 'api_key', apiKeyId: 'key-1' }),
|
||||
);
|
||||
// The shared validator denies a revoked key.
|
||||
apiKeyService.validate.mockRejectedValue(new UnauthorizedException());
|
||||
|
||||
await expect(ext.onAuthenticate(buildData() as any)).rejects.toThrow(
|
||||
UnauthorizedException,
|
||||
);
|
||||
// No new collab connection: the key check gates before page access.
|
||||
expect(pageRepo.findById).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('api_key principal missing apiKeyId → Unauthorized (malformed)', async () => {
|
||||
tokenService.verifyJwt.mockResolvedValue(buildJwt({ principal: 'api_key' }));
|
||||
await expect(ext.onAuthenticate(buildData() as any)).rejects.toThrow(
|
||||
UnauthorizedException,
|
||||
);
|
||||
expect(apiKeyService.validate).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('session principal → NO api-key check (session-backed, incl. internal agent)', async () => {
|
||||
tokenService.verifyJwt.mockResolvedValue(buildJwt({ principal: 'session' }));
|
||||
await ext.onAuthenticate(buildData() as any);
|
||||
expect(apiKeyService.validate).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('claimless token WITHIN the grace window → trusted (legacy pre-rollout)', async () => {
|
||||
// Default rolloutAt = now, so we are inside the grace window.
|
||||
tokenService.verifyJwt.mockResolvedValue(buildJwt()); // no principal
|
||||
await expect(ext.onAuthenticate(buildData() as any)).resolves.toBeDefined();
|
||||
expect(apiKeyService.validate).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('claimless token AFTER the grace window → Unauthorized (fail-closed)', async () => {
|
||||
// Move the rollout reference far into the past so the grace has elapsed.
|
||||
(ext as any).rolloutAt = Date.now() - 25 * 60 * 60 * 1000;
|
||||
tokenService.verifyJwt.mockResolvedValue(buildJwt()); // no principal
|
||||
await expect(ext.onAuthenticate(buildData() as any)).rejects.toThrow(
|
||||
UnauthorizedException,
|
||||
);
|
||||
});
|
||||
|
||||
it('infra error from the api-key row-check propagates (not masked)', async () => {
|
||||
tokenService.verifyJwt.mockResolvedValue(
|
||||
buildJwt({ principal: 'api_key', apiKeyId: 'key-1' }),
|
||||
);
|
||||
const boom = new Error('db down');
|
||||
apiKeyService.validate.mockRejectedValue(boom);
|
||||
await expect(ext.onAuthenticate(buildData() as any)).rejects.toBe(boom);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -14,20 +14,37 @@ import { findHighestUserSpaceRole } from '@docmost/db/repos/space/utils';
|
||||
import { SpaceRole } from '../../common/helpers/types/permission';
|
||||
import { isUserDisabled } from '../../common/helpers';
|
||||
import { getPageId } from '../collaboration.util';
|
||||
import { JwtCollabPayload, JwtType } from '../../core/auth/dto/jwt-payload';
|
||||
import {
|
||||
JwtApiKeyPayload,
|
||||
JwtCollabPayload,
|
||||
JwtType,
|
||||
} from '../../core/auth/dto/jwt-payload';
|
||||
import { resolveProvenance } from '../../common/decorators/auth-provenance.decorator';
|
||||
import { observeCollabAuth } from '../../integrations/metrics/metrics.registry';
|
||||
import { ApiKeyService } from '../../core/api-key/api-key.service';
|
||||
|
||||
// Max lifetime of a collab token (generateCollabToken uses expiresIn '24h'). Used
|
||||
// as the rollout grace window below: once this long has elapsed since this
|
||||
// process started serving the #501 code, every STILL-VALID collab token was
|
||||
// necessarily minted post-rollout and MUST carry the `principal` discriminator,
|
||||
// so a claimless one is a bug and is rejected (fail-closed) rather than trusted.
|
||||
const COLLAB_TOKEN_GRACE_MS = 24 * 60 * 60 * 1000;
|
||||
|
||||
@Injectable()
|
||||
export class AuthenticationExtension implements Extension {
|
||||
private readonly logger = new Logger(AuthenticationExtension.name);
|
||||
|
||||
// Reference instant for the claimless-rejection grace window. Overridable so a
|
||||
// unit test can drive the pre-/post-grace boundary without wall-clock waits.
|
||||
protected rolloutAt = Date.now();
|
||||
|
||||
constructor(
|
||||
private tokenService: TokenService,
|
||||
private userRepo: UserRepo,
|
||||
private pageRepo: PageRepo,
|
||||
private readonly spaceMemberRepo: SpaceMemberRepo,
|
||||
private readonly pagePermissionRepo: PagePermissionRepo,
|
||||
private readonly apiKeyService: ApiKeyService,
|
||||
) {}
|
||||
|
||||
async onAuthenticate(data: onAuthenticatePayload) {
|
||||
@@ -54,6 +71,36 @@ export class AuthenticationExtension implements Extension {
|
||||
throw new UnauthorizedException('Invalid collab token');
|
||||
}
|
||||
|
||||
// #501 — fail-closed api-key laundering guard. A collab token minted by an
|
||||
// api-key principal carries principal='api_key' + apiKeyId; re-check the key
|
||||
// on connect so a REVOKED key gets NO new collab connections (a collab token
|
||||
// outlives its 24h, but a revoked key can no longer open fresh ones). An
|
||||
// api-key token missing its apiKeyId is malformed → reject. A claimless token
|
||||
// (no recognized principal) is trusted only DURING the rollout grace window
|
||||
// (a legacy pre-rollout session token, which api keys could never mint);
|
||||
// once the grace has elapsed every valid token must carry the discriminator,
|
||||
// so a claimless one is a bug and is rejected (not silently trusted for 24h).
|
||||
const principal = jwtPayload.principal;
|
||||
if (principal === 'api_key') {
|
||||
if (!jwtPayload.apiKeyId) {
|
||||
throw new UnauthorizedException();
|
||||
}
|
||||
// Row-check via the SHARED validator: throws Unauthorized on a revoked/
|
||||
// expired/disabled key; an infra error propagates (not masked). No new
|
||||
// connection for a dead key.
|
||||
await this.apiKeyService.validate({
|
||||
sub: jwtPayload.sub,
|
||||
workspaceId: jwtPayload.workspaceId,
|
||||
apiKeyId: jwtPayload.apiKeyId,
|
||||
type: JwtType.API_KEY,
|
||||
} as JwtApiKeyPayload);
|
||||
} else if (principal !== 'session') {
|
||||
// Unrecognized/absent discriminator: reject once past the grace window.
|
||||
if (Date.now() - this.rolloutAt >= COLLAB_TOKEN_GRACE_MS) {
|
||||
throw new UnauthorizedException();
|
||||
}
|
||||
}
|
||||
|
||||
const userId = jwtPayload.sub;
|
||||
const workspaceId = jwtPayload.workspaceId;
|
||||
|
||||
|
||||
@@ -1,84 +1,93 @@
|
||||
import { computeHistoryJob, resolveSource } from './persistence.extension';
|
||||
import {
|
||||
computeHistoryJob,
|
||||
resolveSource,
|
||||
} from './persistence.extension';
|
||||
import {
|
||||
HISTORY_FAST_INTERVAL,
|
||||
HISTORY_FAST_THRESHOLD,
|
||||
HISTORY_INTERVAL,
|
||||
IDLE_INTERVAL_AGENT,
|
||||
IDLE_INTERVAL_USER,
|
||||
IDLE_MAX_WAIT_AGENT,
|
||||
IDLE_MAX_WAIT_USER,
|
||||
} from '../constants';
|
||||
|
||||
// A fixed clock + fixed createdAt make pageAge deterministic.
|
||||
const NOW = 1_700_000_000_000;
|
||||
const PAGE_ID = '550e8400-e29b-41d4-a716-446655440000';
|
||||
|
||||
// Build a minimal page whose age (NOW - createdAt) is exactly `ageMs`.
|
||||
const pageAged = (ageMs: number) => ({
|
||||
id: PAGE_ID,
|
||||
createdAt: new Date(NOW - ageMs),
|
||||
});
|
||||
const page = { id: PAGE_ID };
|
||||
|
||||
describe('computeHistoryJob', () => {
|
||||
it('agent edit → delay MUST be 0 and job id is source-keyed', () => {
|
||||
// INVARIANT (§15 H2 / persistence.extension): the agent delay MUST stay 0.
|
||||
// The worker re-reads the page row at run time, so any non-zero delay risks
|
||||
// snapshotting content a later human edit has already overwritten. This is
|
||||
// the load-bearing assertion of this spec — do not relax it.
|
||||
const { jobId, delay } = computeHistoryJob(pageAged(0), 'agent', NOW);
|
||||
expect(delay).toBe(0);
|
||||
expect(jobId).toBe(`${PAGE_ID}-agent`);
|
||||
});
|
||||
|
||||
it('agent edit on an OLD page is still delay 0 (age never applies to agents)', () => {
|
||||
// Even when the page is far older than the fast threshold, the agent path
|
||||
// must short-circuit to 0 — age-based debounce is a human-only concern.
|
||||
const { jobId, delay } = computeHistoryJob(
|
||||
pageAged(HISTORY_FAST_THRESHOLD + 60_000),
|
||||
'agent',
|
||||
NOW,
|
||||
);
|
||||
expect(delay).toBe(0);
|
||||
expect(jobId).toBe(`${PAGE_ID}-agent`);
|
||||
});
|
||||
|
||||
it('human edit on a YOUNG page (age < threshold) → fast interval, bare job id', () => {
|
||||
const { jobId, delay } = computeHistoryJob(
|
||||
pageAged(HISTORY_FAST_THRESHOLD - 1),
|
||||
'user',
|
||||
NOW,
|
||||
);
|
||||
expect(delay).toBe(HISTORY_FAST_INTERVAL);
|
||||
describe('computeHistoryJob (#370 — shared trailing idle pipeline)', () => {
|
||||
it('human edit → user idle window, bare page.id job', () => {
|
||||
// Humans and the agent now share ONE idle job per page (jobId = page.id).
|
||||
// The agent's old delay=0 fast path is GONE — intentional agent points now
|
||||
// arrive via the explicit save-version signal, not a zero-delay snapshot.
|
||||
const { jobId, delay } = computeHistoryJob(page, 'user');
|
||||
expect(delay).toBe(IDLE_INTERVAL_USER);
|
||||
expect(jobId).toBe(PAGE_ID);
|
||||
});
|
||||
|
||||
it('human edit on an OLD page (age > threshold) → standard interval', () => {
|
||||
const { jobId, delay } = computeHistoryJob(
|
||||
pageAged(HISTORY_FAST_THRESHOLD + 1),
|
||||
'user',
|
||||
NOW,
|
||||
);
|
||||
expect(delay).toBe(HISTORY_INTERVAL);
|
||||
it('agent edit → agent idle window (shorter), still the bare page.id job', () => {
|
||||
const { jobId, delay } = computeHistoryJob(page, 'agent');
|
||||
expect(delay).toBe(IDLE_INTERVAL_AGENT);
|
||||
// No `-agent` suffix anymore: the agent joins the common idle pipeline.
|
||||
expect(jobId).toBe(PAGE_ID);
|
||||
});
|
||||
|
||||
it('boundary: pageAge EXACTLY === threshold takes the slow branch (the `<` is strict)', () => {
|
||||
// Off-by-one guard: the condition is `pageAge < HISTORY_FAST_THRESHOLD`, so
|
||||
// an age of exactly the threshold is NOT "fast" — it must use HISTORY_INTERVAL.
|
||||
const { delay } = computeHistoryJob(
|
||||
pageAged(HISTORY_FAST_THRESHOLD),
|
||||
'user',
|
||||
NOW,
|
||||
);
|
||||
expect(delay).toBe(HISTORY_INTERVAL);
|
||||
it('agent flushes sooner than a human', () => {
|
||||
expect(IDLE_INTERVAL_AGENT).toBeLessThan(IDLE_INTERVAL_USER);
|
||||
});
|
||||
|
||||
it('treats any non-"agent" source string as human', () => {
|
||||
// resolveSource only ever yields 'agent' | 'user', but guard the contract:
|
||||
// the agent branch keys strictly on === 'agent'.
|
||||
const { jobId, delay } = computeHistoryJob(pageAged(0), 'user', NOW);
|
||||
expect(delay).toBe(HISTORY_FAST_INTERVAL);
|
||||
it('treats any non-"agent" source string as human (keys strictly on === agent)', () => {
|
||||
const { jobId, delay } = computeHistoryJob(page, 'user');
|
||||
expect(delay).toBe(IDLE_INTERVAL_USER);
|
||||
expect(jobId).toBe(PAGE_ID);
|
||||
});
|
||||
|
||||
// #370 review round-1 WARNING: the max-wait ceiling prevents autosnapshot
|
||||
// starvation during a continuous editing session (the trailing timer would
|
||||
// otherwise re-arm forever and never fire).
|
||||
describe('max-wait ceiling', () => {
|
||||
const T0 = 1_000_000; // arbitrary fixed epoch for deterministic tests
|
||||
|
||||
it('once a burst is armed, delay clamps to the remaining max-wait budget', () => {
|
||||
// 1 minute into the burst the USER interval (60m) far exceeds the remaining
|
||||
// max-wait budget (10m - 1m = 9m), so the delay is clamped DOWN to that
|
||||
// remaining budget — the full interval is NOT used once a ceiling applies.
|
||||
const { delay } = computeHistoryJob(page, 'user', T0, T0 + 60_000);
|
||||
expect(delay).toBe(IDLE_MAX_WAIT_USER - 60_000);
|
||||
});
|
||||
|
||||
it('never waits longer than the max-wait budget from the burst start', () => {
|
||||
// A store arriving right at the ceiling → delay 0 (fire promptly).
|
||||
const { delay } = computeHistoryJob(
|
||||
page,
|
||||
'user',
|
||||
T0,
|
||||
T0 + IDLE_MAX_WAIT_USER,
|
||||
);
|
||||
expect(delay).toBe(0);
|
||||
});
|
||||
|
||||
it('past the ceiling never returns a negative delay', () => {
|
||||
const { delay } = computeHistoryJob(
|
||||
page,
|
||||
'user',
|
||||
T0,
|
||||
T0 + IDLE_MAX_WAIT_USER + 5 * 60_000,
|
||||
);
|
||||
expect(delay).toBe(0);
|
||||
});
|
||||
|
||||
it('the agent ceiling is shorter than the user ceiling', () => {
|
||||
expect(IDLE_MAX_WAIT_AGENT).toBeLessThan(IDLE_MAX_WAIT_USER);
|
||||
const { delay } = computeHistoryJob(
|
||||
page,
|
||||
'agent',
|
||||
T0,
|
||||
T0 + IDLE_MAX_WAIT_AGENT,
|
||||
);
|
||||
expect(delay).toBe(0);
|
||||
});
|
||||
|
||||
it('without a burstStart there is no ceiling (backward-compatible)', () => {
|
||||
expect(computeHistoryJob(page, 'user').delay).toBe(IDLE_INTERVAL_USER);
|
||||
expect(computeHistoryJob(page, 'agent').delay).toBe(IDLE_INTERVAL_AGENT);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('resolveSource (truth table)', () => {
|
||||
|
||||
@@ -40,11 +40,12 @@ describe('PersistenceExtension.onStoreDocument — Approach-A boundary snapshot'
|
||||
let pageHistoryRepo: {
|
||||
saveHistory: jest.Mock;
|
||||
findPageLastHistory: jest.Mock;
|
||||
updateHistoryKind: jest.Mock;
|
||||
};
|
||||
let aiQueue: { add: jest.Mock };
|
||||
let historyQueue: { add: jest.Mock };
|
||||
let historyQueue: { add: jest.Mock; remove: jest.Mock };
|
||||
let notificationQueue: { add: jest.Mock };
|
||||
let collabHistory: { addContributors: jest.Mock };
|
||||
let collabHistory: { addContributors: jest.Mock; popContributors: jest.Mock };
|
||||
let transclusionService: {
|
||||
syncPageTransclusions: jest.Mock;
|
||||
syncPageReferences: jest.Mock;
|
||||
@@ -93,13 +94,22 @@ describe('PersistenceExtension.onStoreDocument — Approach-A boundary snapshot'
|
||||
pageHistoryRepo = {
|
||||
saveHistory: jest.fn().mockImplementation(async () => {
|
||||
callOrder.push('saveHistory');
|
||||
return { id: 'history-1' };
|
||||
}),
|
||||
findPageLastHistory: jest.fn().mockResolvedValue(null),
|
||||
updateHistoryKind: jest.fn().mockResolvedValue(undefined),
|
||||
};
|
||||
aiQueue = { add: jest.fn().mockResolvedValue(undefined) };
|
||||
historyQueue = { add: jest.fn().mockResolvedValue(undefined) };
|
||||
historyQueue = {
|
||||
add: jest.fn().mockResolvedValue(undefined),
|
||||
// #370 — enqueuePageHistory now removes any pending idle job before re-adding.
|
||||
remove: jest.fn().mockResolvedValue(undefined),
|
||||
};
|
||||
notificationQueue = { add: jest.fn().mockResolvedValue(undefined) };
|
||||
collabHistory = { addContributors: jest.fn().mockResolvedValue(undefined) };
|
||||
collabHistory = {
|
||||
addContributors: jest.fn().mockResolvedValue(undefined),
|
||||
popContributors: jest.fn().mockResolvedValue([]),
|
||||
};
|
||||
transclusionService = {
|
||||
syncPageTransclusions: jest.fn().mockResolvedValue(undefined),
|
||||
syncPageReferences: jest.fn().mockResolvedValue(undefined),
|
||||
@@ -165,6 +175,50 @@ describe('PersistenceExtension.onStoreDocument — Approach-A boundary snapshot'
|
||||
expect(pageRepo.updatePage.mock.calls[0][0].lastUpdatedSource).toBe('user');
|
||||
});
|
||||
|
||||
// #370 review round-1 SUGGESTION: the boundary was GENERALIZED from a
|
||||
// user→agent special-case to ANY lastUpdatedSource transition. These pin the
|
||||
// generalized behaviour it was rebuilt for.
|
||||
describe('generalized boundary — any source transition', () => {
|
||||
// Same persisted page but with an explicit prior source.
|
||||
const pageWithPriorSource = (prior: string | null) => ({
|
||||
...persistedHumanPage('NEW CONTENT'),
|
||||
lastUpdatedSource: prior,
|
||||
});
|
||||
|
||||
it('agent→user transition fires the boundary (pins the prior agent revision)', async () => {
|
||||
const document = ydocFor(doc('NEW CONTENT'));
|
||||
pageRepo.findById.mockResolvedValue(pageWithPriorSource('agent'));
|
||||
pageHistoryRepo.findPageLastHistory.mockResolvedValue(null);
|
||||
|
||||
await ext.onStoreDocument(buildData(document, 'user') as any);
|
||||
|
||||
expect(pageHistoryRepo.saveHistory).toHaveBeenCalledTimes(1);
|
||||
expect(callOrder).toEqual(['saveHistory', 'updatePage']);
|
||||
expect(pageRepo.updatePage.mock.calls[0][0].lastUpdatedSource).toBe('user');
|
||||
});
|
||||
|
||||
it('git→user transition fires the boundary (git-sync overwrite is a source change)', async () => {
|
||||
const document = ydocFor(doc('NEW CONTENT'));
|
||||
pageRepo.findById.mockResolvedValue(pageWithPriorSource('git'));
|
||||
pageHistoryRepo.findPageLastHistory.mockResolvedValue(null);
|
||||
|
||||
await ext.onStoreDocument(buildData(document, 'user') as any);
|
||||
|
||||
expect(pageHistoryRepo.saveHistory).toHaveBeenCalledTimes(1);
|
||||
expect(callOrder).toEqual(['saveHistory', 'updatePage']);
|
||||
});
|
||||
|
||||
it('a null prior source (first-ever edit) does NOT fire the boundary', async () => {
|
||||
const document = ydocFor(doc('NEW CONTENT'));
|
||||
pageRepo.findById.mockResolvedValue(pageWithPriorSource(null));
|
||||
|
||||
await ext.onStoreDocument(buildData(document, 'agent') as any);
|
||||
|
||||
expect(pageHistoryRepo.saveHistory).not.toHaveBeenCalled();
|
||||
expect(pageRepo.updatePage).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
|
||||
it('idempotency: unchanged content → no updatePage, no history, no queues', async () => {
|
||||
// The Y.Doc content equals the persisted content deeply → early skip.
|
||||
// A Y.Doc round-trip normalizes attrs (e.g. paragraph indent), so derive
|
||||
@@ -479,4 +533,231 @@ describe('PersistenceExtension.onStoreDocument — Approach-A boundary snapshot'
|
||||
// Contributors keyed by the UUID so they match the PAGE_HISTORY job (page.id).
|
||||
expect(collabHistory.addContributors.mock.calls[0][0]).toBe(PAGE_ID);
|
||||
});
|
||||
|
||||
// #370 — explicit save-version (Cmd+S / agent save tool) over the stateless
|
||||
// seam. The tier is derived from the SIGNED connection actor, the store path
|
||||
// is reused, and promote-not-dup avoids duplicating heavy content rows.
|
||||
describe('save-version (#370)', () => {
|
||||
const emitSave = (document: any, actor: 'user' | 'agent') =>
|
||||
ext.onStateless({
|
||||
connection: {
|
||||
readOnly: false,
|
||||
context: { user: { id: USER_ID, name: 'Alice' }, actor },
|
||||
} as any,
|
||||
documentName: `page.${PAGE_ID}`,
|
||||
document: document as any,
|
||||
payload: JSON.stringify({ type: 'save-version' }),
|
||||
} as any);
|
||||
|
||||
// findById returns a page whose content already equals the live doc, so the
|
||||
// store path is a no-op and we isolate the versioning decision.
|
||||
const pageMatchingDoc = (document: any) => ({
|
||||
...persistedHumanPage('IGNORED'),
|
||||
content: TiptapTransformer.fromYdoc(document, 'default'),
|
||||
});
|
||||
|
||||
it('human save with no prior snapshot → writes a manual version + broadcasts', async () => {
|
||||
const document = ydocFor(doc('VERSION ME'));
|
||||
pageRepo.findById.mockResolvedValue(pageMatchingDoc(document));
|
||||
pageHistoryRepo.findPageLastHistory.mockResolvedValue(null);
|
||||
|
||||
await emitSave(document, 'user');
|
||||
|
||||
expect(pageHistoryRepo.saveHistory).toHaveBeenCalledTimes(1);
|
||||
expect(pageHistoryRepo.saveHistory.mock.calls[0][1]).toEqual(
|
||||
expect.objectContaining({ kind: 'manual' }),
|
||||
);
|
||||
// The pending idle autosnapshot is cancelled by the explicit version.
|
||||
expect(historyQueue.remove).toHaveBeenCalledWith(PAGE_ID);
|
||||
const msg = JSON.parse(
|
||||
(document as any).broadcastStateless.mock.calls[(document as any).broadcastStateless.mock.calls.length - 1][0],
|
||||
);
|
||||
expect(msg).toMatchObject({
|
||||
type: 'version.saved',
|
||||
kind: 'manual',
|
||||
alreadySaved: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('agent save derives kind=agent from the signed actor', async () => {
|
||||
const document = ydocFor(doc('AGENT VERSION'));
|
||||
pageRepo.findById.mockResolvedValue(pageMatchingDoc(document));
|
||||
pageHistoryRepo.findPageLastHistory.mockResolvedValue(null);
|
||||
|
||||
await emitSave(document, 'agent');
|
||||
|
||||
expect(pageHistoryRepo.saveHistory.mock.calls[pageHistoryRepo.saveHistory.mock.calls.length - 1][1]).toEqual(
|
||||
expect.objectContaining({ kind: 'agent' }),
|
||||
);
|
||||
});
|
||||
|
||||
it('promote-not-dup: latest snapshot is an autosave with identical content → upgrades in place', async () => {
|
||||
const document = ydocFor(doc('SAME'));
|
||||
const page = pageMatchingDoc(document);
|
||||
pageRepo.findById.mockResolvedValue(page);
|
||||
pageHistoryRepo.findPageLastHistory.mockResolvedValue({
|
||||
id: 'auto-1',
|
||||
content: page.content,
|
||||
kind: 'idle',
|
||||
});
|
||||
|
||||
await emitSave(document, 'user');
|
||||
|
||||
// No heavy new content row — the existing autosave is promoted to manual.
|
||||
expect(pageHistoryRepo.updateHistoryKind).toHaveBeenCalledWith(
|
||||
'auto-1',
|
||||
'manual',
|
||||
expect.anything(),
|
||||
);
|
||||
expect(pageHistoryRepo.saveHistory).not.toHaveBeenCalled();
|
||||
const msg = JSON.parse(
|
||||
(document as any).broadcastStateless.mock.calls[(document as any).broadcastStateless.mock.calls.length - 1][0],
|
||||
);
|
||||
expect(msg).toMatchObject({ historyId: 'auto-1', alreadySaved: false });
|
||||
});
|
||||
|
||||
it('no-op when the latest snapshot is already a manual version of this content', async () => {
|
||||
const document = ydocFor(doc('ALREADY SAVED'));
|
||||
const page = pageMatchingDoc(document);
|
||||
pageRepo.findById.mockResolvedValue(page);
|
||||
pageHistoryRepo.findPageLastHistory.mockResolvedValue({
|
||||
id: 'ver-1',
|
||||
content: page.content,
|
||||
kind: 'manual',
|
||||
});
|
||||
|
||||
await emitSave(document, 'user');
|
||||
|
||||
expect(pageHistoryRepo.updateHistoryKind).not.toHaveBeenCalled();
|
||||
expect(pageHistoryRepo.saveHistory).not.toHaveBeenCalled();
|
||||
const msg = JSON.parse(
|
||||
(document as any).broadcastStateless.mock.calls[(document as any).broadcastStateless.mock.calls.length - 1][0],
|
||||
);
|
||||
expect(msg).toMatchObject({ alreadySaved: true, kind: 'manual' });
|
||||
});
|
||||
|
||||
it('a read-only connection cannot save a version', async () => {
|
||||
const document = ydocFor(doc('READER'));
|
||||
pageRepo.findById.mockResolvedValue(pageMatchingDoc(document));
|
||||
|
||||
await ext.onStateless({
|
||||
connection: {
|
||||
readOnly: true,
|
||||
context: { user: { id: USER_ID }, actor: 'user' },
|
||||
} as any,
|
||||
documentName: `page.${PAGE_ID}`,
|
||||
document: document as any,
|
||||
payload: JSON.stringify({ type: 'save-version' }),
|
||||
} as any);
|
||||
|
||||
expect(pageHistoryRepo.saveHistory).not.toHaveBeenCalled();
|
||||
expect(pageHistoryRepo.updateHistoryKind).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
// #370 F8-twin — a COMMIT abort (serialization/deadlock/conn-drop) rejects
|
||||
// OUTSIDE the tx callback, AFTER the destructive popContributors (SPOP) and
|
||||
// saveHistory ran but the INSERT rolled back. onStateless has no retry, so
|
||||
// the outer catch MUST re-add (SADD) the popped set or attribution is lost
|
||||
// irrecoverably. MUTATION: drop the outer catch → addContributors is never
|
||||
// called → this reddens.
|
||||
it('restores popped contributors when the commit aborts after the callback', async () => {
|
||||
const document = ydocFor(doc('VERSION ME'));
|
||||
pageRepo.findById.mockResolvedValue(pageMatchingDoc(document));
|
||||
// No matching snapshot → fresh version branch → pops contributors.
|
||||
pageHistoryRepo.findPageLastHistory.mockResolvedValue(null);
|
||||
collabHistory.popContributors.mockResolvedValue(['u1', 'u2']);
|
||||
|
||||
// A db whose commit REJECTS after the callback body resolved: the SPOP and
|
||||
// saveHistory already ran, then the tx aborts. onStoreDocument's flush uses
|
||||
// the same db but its content matches (no-op branch) and its own retry loop
|
||||
// swallows the throw, so only the versioning tx exercises the restore.
|
||||
const commitFailingDb = {
|
||||
transaction: () => ({
|
||||
execute: async (fn: (trx: any) => Promise<any>) => {
|
||||
await fn(trxStub);
|
||||
throw new Error('commit aborted (serialization_failure)');
|
||||
},
|
||||
}),
|
||||
};
|
||||
const ext2 = new PersistenceExtension(
|
||||
pageRepo as any,
|
||||
pageHistoryRepo as any,
|
||||
commitFailingDb as any,
|
||||
aiQueue as any,
|
||||
historyQueue as any,
|
||||
notificationQueue as any,
|
||||
collabHistory as any,
|
||||
transclusionService as any,
|
||||
);
|
||||
jest.spyOn(ext2['logger'], 'debug').mockImplementation(() => undefined);
|
||||
jest.spyOn(ext2['logger'], 'warn').mockImplementation(() => undefined);
|
||||
jest.spyOn(ext2['logger'], 'error').mockImplementation(() => undefined);
|
||||
|
||||
await expect(
|
||||
ext2.onStateless({
|
||||
connection: {
|
||||
readOnly: false,
|
||||
context: { user: { id: USER_ID, name: 'Alice' }, actor: 'user' },
|
||||
} as any,
|
||||
documentName: `page.${PAGE_ID}`,
|
||||
document: document as any,
|
||||
payload: JSON.stringify({ type: 'save-version' }),
|
||||
} as any),
|
||||
).rejects.toThrow();
|
||||
|
||||
// Attribution preserved: the popped set is SADD-restored, keyed by the page
|
||||
// UUID it was popped under.
|
||||
expect(collabHistory.addContributors).toHaveBeenCalledWith(PAGE_ID, [
|
||||
'u1',
|
||||
'u2',
|
||||
]);
|
||||
});
|
||||
|
||||
// #370 #260 — for a `page.<slugId>` document the idle job is armed under the
|
||||
// page UUID (computeHistoryJob's jobId = page.id), so the supersede-remove
|
||||
// must target page.id, not the raw slugId doc-name id, or it silently misses.
|
||||
it('cancels the superseded idle job by the page UUID for a slugId doc', async () => {
|
||||
const SLUG = 'slug-1'; // persistedHumanPage.slugId
|
||||
const document = ydocFor(doc('VERSION ME'));
|
||||
pageRepo.findById.mockResolvedValue(pageMatchingDoc(document));
|
||||
pageHistoryRepo.findPageLastHistory.mockResolvedValue(null);
|
||||
|
||||
await ext.onStateless({
|
||||
connection: {
|
||||
readOnly: false,
|
||||
context: { user: { id: USER_ID, name: 'Alice' }, actor: 'user' },
|
||||
} as any,
|
||||
documentName: `page.${SLUG}`,
|
||||
document: document as any,
|
||||
payload: JSON.stringify({ type: 'save-version' }),
|
||||
} as any);
|
||||
|
||||
// remove() keyed by the UUID (the real jobId), never the slugId.
|
||||
expect(historyQueue.remove).toHaveBeenCalledWith(PAGE_ID);
|
||||
expect(historyQueue.remove).not.toHaveBeenCalledWith(SLUG);
|
||||
});
|
||||
});
|
||||
|
||||
// #370 — the in-memory idle-burst marker must be dropped on doc unload (like
|
||||
// its sibling per-document maps) or it grows unbounded for every page that was
|
||||
// edited but never manually saved. MUTATION: drop the afterUnloadDocument
|
||||
// delete → the entry survives → this reddens.
|
||||
describe('idleBurstStart housekeeping', () => {
|
||||
it('afterUnloadDocument clears the idle-burst marker armed by a store', async () => {
|
||||
const document = ydocFor(doc('EDIT'));
|
||||
pageRepo.findById.mockResolvedValue(persistedHumanPage('EDIT'));
|
||||
|
||||
await ext.onStoreDocument(buildData(document, 'user') as any);
|
||||
|
||||
const map = ext['idleBurstStart'] as Map<string, number>;
|
||||
// Keyed by documentName (buildData uses `page.${PAGE_ID}`).
|
||||
expect(map.has(`page.${PAGE_ID}`)).toBe(true);
|
||||
|
||||
await ext.afterUnloadDocument({
|
||||
documentName: `page.${PAGE_ID}`,
|
||||
} as any);
|
||||
|
||||
expect(map.has(`page.${PAGE_ID}`)).toBe(false);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -37,9 +37,11 @@ import { Page } from '@docmost/db/types/entity.types';
|
||||
import { CollabHistoryService } from '../services/collab-history.service';
|
||||
import {
|
||||
EMBED_DEBOUNCE_MS,
|
||||
HISTORY_FAST_INTERVAL,
|
||||
HISTORY_FAST_THRESHOLD,
|
||||
HISTORY_INTERVAL,
|
||||
IDLE_INTERVAL_AGENT,
|
||||
IDLE_INTERVAL_USER,
|
||||
IDLE_MAX_WAIT_AGENT,
|
||||
IDLE_MAX_WAIT_USER,
|
||||
PageHistoryKind,
|
||||
} from '../constants';
|
||||
import { TransclusionService } from '../../core/page/transclusion/transclusion.service';
|
||||
import {
|
||||
@@ -56,6 +58,16 @@ import { hasTransclusionFamilyNodes } from '../../core/page/transclusion/utils/t
|
||||
*/
|
||||
export const INTENTIONAL_CLEAR_MESSAGE_TYPE = 'intentional-clear';
|
||||
|
||||
/**
|
||||
* #370 — wire format of the client→server "save a version" signal. Sent by the
|
||||
* human (Cmd+S / Save button) and by the agent's explicit save tool over the
|
||||
* SAME stateless channel. The intentionality tier ('manual' vs 'agent') is
|
||||
* derived SERVER-SIDE from the signed connection actor, never from this
|
||||
* payload, so a version's type is unforgeable. The document is taken from the
|
||||
* connection (not the payload), so the signal cannot be aimed at another page.
|
||||
*/
|
||||
export const SAVE_VERSION_MESSAGE_TYPE = 'save-version';
|
||||
|
||||
/**
|
||||
* #251 — how long an intentional-clear signal stays "pending" before it is
|
||||
* ignored. The signal is set on the clearing keystroke but consumed by the
|
||||
@@ -92,35 +104,39 @@ export function resolveSource(
|
||||
}
|
||||
|
||||
/**
|
||||
* Compute the BullMQ job id + delay for a page-history snapshot job. Pure so
|
||||
* the data-loss-sensitive timing arithmetic is unit-testable; `now` is injected
|
||||
* (caller passes `Date.now()`) for determinism.
|
||||
* #370 — compute the BullMQ job id + delay for a page's trailing idle-flush
|
||||
* autosnapshot. Pure so the timing is unit-testable.
|
||||
*
|
||||
* - Agent edits: delay 0 and a source-keyed job id `${page.id}-agent`. The
|
||||
* delay MUST stay 0 — the worker re-reads the page row at run time, so any
|
||||
* delay risks reading content a later human edit has already overwritten
|
||||
* (mis-tagged snapshot). 0 minimizes that window. The `-agent` suffix keeps
|
||||
* the job from coalescing with the bare-page.id human job.
|
||||
* - Human edits: age-based debounce so rapid human edits coalesce into one
|
||||
* snapshot; job id is the bare `page.id`.
|
||||
*
|
||||
* BullMQ forbids ':' in custom job ids (Redis key separator), so '-' is used;
|
||||
* page.id is a UUID, so `${page.id}-agent` cannot collide with a human job.
|
||||
* Both humans and the agent now share ONE idle pipeline (the agent's old
|
||||
* `delay=0` fast path is gone — intentional agent points arrive via the
|
||||
* explicit save-version signal instead). The job id is the bare `page.id`, so a
|
||||
* page has at most one pending idle job; the caller removes-and-re-adds it on
|
||||
* every store to keep it debounced to the trailing edge of an edit burst. The
|
||||
* window differs by source only: the agent flushes sooner than a human.
|
||||
*/
|
||||
export function computeHistoryJob(
|
||||
page: Pick<Page, 'id' | 'createdAt'>,
|
||||
page: Pick<Page, 'id'>,
|
||||
source: string,
|
||||
now: number,
|
||||
// Epoch ms of the FIRST edit in the current burst (when the pending idle job
|
||||
// was first armed). Used to enforce the max-wait ceiling so a continuous
|
||||
// editing session cannot re-arm the trailing timer forever. `now` is injectable
|
||||
// for tests; both default to a live clock / no ceiling when omitted.
|
||||
burstStart?: number,
|
||||
now: number = Date.now(),
|
||||
): { jobId: string; delay: number } {
|
||||
const isAgent = source === 'agent';
|
||||
const pageAge = now - new Date(page.createdAt).getTime();
|
||||
const delay = isAgent
|
||||
? 0
|
||||
: pageAge < HISTORY_FAST_THRESHOLD
|
||||
? HISTORY_FAST_INTERVAL
|
||||
: HISTORY_INTERVAL;
|
||||
const jobId = isAgent ? `${page.id}-agent` : page.id;
|
||||
return { jobId, delay };
|
||||
const interval = isAgent ? IDLE_INTERVAL_AGENT : IDLE_INTERVAL_USER;
|
||||
const maxWait = isAgent ? IDLE_MAX_WAIT_AGENT : IDLE_MAX_WAIT_USER;
|
||||
|
||||
let delay = interval;
|
||||
if (burstStart !== undefined) {
|
||||
// Time already elapsed since the burst's first edit; the snapshot must fire
|
||||
// no later than `maxWait` after that, so shrink the trailing delay to the
|
||||
// remaining budget (never negative, so BullMQ fires it promptly).
|
||||
const remaining = burstStart + maxWait - now;
|
||||
delay = Math.max(0, Math.min(interval, remaining));
|
||||
}
|
||||
return { jobId: page.id, delay };
|
||||
}
|
||||
|
||||
@Injectable()
|
||||
@@ -132,6 +148,28 @@ export class PersistenceExtension implements Extension {
|
||||
// coalescing window" per document and OR it across all edits in the window,
|
||||
// so the snapshot is marked 'agent' regardless of who wrote last.
|
||||
private agentTouched: Map<string, boolean> = new Map();
|
||||
// #370 — epoch ms of the FIRST edit in the current idle-flush burst. Keyed by
|
||||
// documentName (like its sibling per-document maps above), NOT by page.id, so
|
||||
// it can be cleaned in afterUnloadDocument alongside `contributors` /
|
||||
// `agentTouched` / `intentionalClear` when the doc unloads — otherwise any page
|
||||
// that was edited but never manually saved (the common case) would keep its
|
||||
// entry forever and the Map would grow unbounded in this long-lived process.
|
||||
// Set when the pending idle job is first armed (empty entry), read to enforce
|
||||
// the max-wait ceiling in computeHistoryJob, and cleared on doc unload or when
|
||||
// a manual save cancels the idle job so the next burst starts a fresh window.
|
||||
//
|
||||
// Single-process assumption (like `contributors` / `agentTouched` above): this
|
||||
// lives only in THIS collab process's memory. A restart, or a page's ownership
|
||||
// moving to another node, loses the burst-start marker. Consequence: a burst
|
||||
// that spans the restart looks like a fresh burst to the surviving process, so
|
||||
// its max-wait ceiling is re-anchored to the first post-restart edit — a single
|
||||
// continuous session straddling a restart can therefore wait up to ~2× the cap
|
||||
// for its idle snapshot (once for the lost pre-restart window, once for the new
|
||||
// one). Bounded and benign (it only DELAYS a safety-net autosnapshot; manual
|
||||
// saves are unaffected and the next quiet period always flushes), but the
|
||||
// assumption and its consequence are recorded here so no one mistakes the
|
||||
// in-memory marker for a durable, cross-process guarantee.
|
||||
private idleBurstStart: Map<string, number> = new Map();
|
||||
// #251 — per-document "intentional clear pending" flags. Keyed by
|
||||
// documentName, value = expiry timestamp (ms). Set by onStateless when the
|
||||
// client reports a deliberate clear; consumed once by the next
|
||||
@@ -363,20 +401,19 @@ export class PersistenceExtension implements Extension {
|
||||
//this.logger.debug('Contributors error:' + err?.['message']);
|
||||
}
|
||||
|
||||
// Approach A — boundary snapshot before the agent's first edit.
|
||||
// When this store is the agent's and the page's currently persisted
|
||||
// state was authored by a human, pin that human state as its own
|
||||
// history version BEFORE the agent overwrites it. `page` still holds
|
||||
// the OLD content/provenance here, so saveHistory(page) captures the
|
||||
// pre-agent state tagged 'user'. The agent's new content is
|
||||
// snapshotted later by the debounced PAGE_HISTORY job ('agent'). Skip
|
||||
// if the prior state is already agent-authored (boundary already
|
||||
// pinned on the user->agent transition), if the page is effectively
|
||||
// empty, or if the latest existing snapshot already equals this human
|
||||
// state (avoid duplicates).
|
||||
// #370 — boundary snapshot on ANY source transition. When the store
|
||||
// flips the page's provenance (user↔agent↔git), pin the OUTGOING
|
||||
// state as its own history version BEFORE the incoming source
|
||||
// overwrites it. `page` still holds the OLD content/provenance here,
|
||||
// so saveHistory(page) captures the pre-transition state tagged with
|
||||
// its own source, kind='boundary'. The incoming content is snapshotted
|
||||
// later by the debounced idle job. Skip if the page is effectively
|
||||
// empty or if the latest existing snapshot already equals this state
|
||||
// (the shared isDeepStrictEqual gate — avoids duplicates). Generalizing
|
||||
// beyond the old user→agent special-case also covers git-sync for free.
|
||||
if (
|
||||
lastUpdatedSource === 'agent' &&
|
||||
page.lastUpdatedSource !== 'agent'
|
||||
page.lastUpdatedSource &&
|
||||
page.lastUpdatedSource !== lastUpdatedSource
|
||||
) {
|
||||
// pageHistory.pageId is uuid-typed; use page.id (never the doc-name
|
||||
// slugId) so a `page.<slugId>` doc cannot throw 22P02 here (#260).
|
||||
@@ -384,15 +421,13 @@ export class PersistenceExtension implements Extension {
|
||||
page.id,
|
||||
{ includeContent: true, trx },
|
||||
);
|
||||
const humanBaselineMissing =
|
||||
const baselineMissing =
|
||||
!lastHistory ||
|
||||
!isDeepStrictEqual(lastHistory.content, page.content);
|
||||
if (
|
||||
!isEmptyParagraphDoc(page.content as any) &&
|
||||
humanBaselineMissing
|
||||
) {
|
||||
if (!isEmptyParagraphDoc(page.content as any) && baselineMissing) {
|
||||
await this.pageHistoryRepo.saveHistory(page, {
|
||||
contributorIds: page.contributorIds ?? undefined,
|
||||
kind: 'boundary',
|
||||
trx,
|
||||
});
|
||||
}
|
||||
@@ -522,7 +557,7 @@ export class PersistenceExtension implements Extension {
|
||||
{ jobId: `embed-${page.id}`, delay: EMBED_DEBOUNCE_MS },
|
||||
);
|
||||
|
||||
await this.enqueuePageHistory(page, lastUpdatedSource);
|
||||
await this.enqueuePageHistory(page, documentName, lastUpdatedSource);
|
||||
}
|
||||
|
||||
// #402 — report the serialized size for the store histogram's size_bucket.
|
||||
@@ -554,6 +589,14 @@ export class PersistenceExtension implements Extension {
|
||||
return; // unrelated / malformed stateless message
|
||||
}
|
||||
|
||||
// #370 — explicit "save a version" (human Cmd+S / agent save tool). Edit
|
||||
// rights are already enforced by the readOnly reject above (a reader can't
|
||||
// create a version), exactly as intentional-clear requires.
|
||||
if (message?.type === SAVE_VERSION_MESSAGE_TYPE) {
|
||||
await this.handleSaveVersion(data);
|
||||
return;
|
||||
}
|
||||
|
||||
if (message?.type !== INTENTIONAL_CLEAR_MESSAGE_TYPE) return;
|
||||
|
||||
this.intentionalClear.set(
|
||||
@@ -562,6 +605,160 @@ export class PersistenceExtension implements Extension {
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* #370 — persist an intentional version from the live in-memory ydoc.
|
||||
*
|
||||
* One stateless path serves BOTH the human and the agent; the tier is derived
|
||||
* SERVER-SIDE from the signed connection actor ('agent' → 'agent', anything
|
||||
* else → 'manual'), so the version type cannot be spoofed by the client. We
|
||||
* take the fresh ydoc from the collab process memory and run it through the
|
||||
* EXISTING store path first (so pages.content/ydoc reflect the exact content
|
||||
* being versioned — a REST endpoint would race the up-to-10s-stale page row),
|
||||
* then snapshot it into page_history with the intentional kind.
|
||||
*
|
||||
* Promote-not-dup: if the latest history row already holds this exact content
|
||||
* and it is an autosave (idle/boundary/legacy-null), upgrade its kind in place
|
||||
* instead of duplicating a heavy content row; if it is already 'manual', it is
|
||||
* a no-op (the client shows an "already saved" toast). Otherwise a fresh
|
||||
* version row is written, popping the aggregated contributors from Redis.
|
||||
*/
|
||||
private async handleSaveVersion(data: onStatelessPayload): Promise<void> {
|
||||
const { connection, document, documentName } = data;
|
||||
const context = connection?.context;
|
||||
const pageId = getPageId(documentName);
|
||||
// Unforgeable: 'agent' only for a signed agent connection, else 'manual'.
|
||||
const kind: PageHistoryKind =
|
||||
context?.actor === 'agent' ? 'agent' : 'manual';
|
||||
|
||||
// Flush the live ydoc through the normal store path so the page row + ydoc
|
||||
// hold exactly what we are about to version (also fires the idle enqueue we
|
||||
// supersede below, plus any source-transition boundary). onStoreDocument
|
||||
// only needs document/documentName/context.
|
||||
await this.onStoreDocument({
|
||||
document,
|
||||
documentName,
|
||||
context,
|
||||
} as onStoreDocumentPayload);
|
||||
|
||||
let result:
|
||||
| { historyId: string; kind: PageHistoryKind; alreadySaved: boolean }
|
||||
| undefined;
|
||||
|
||||
// #370 F8-twin — the contributor set popped from Redis (destructive SPOP)
|
||||
// must be restored if the version row does not durably land. The inner
|
||||
// try/catch below only covers a throw INSIDE the callback; but executeTx
|
||||
// COMMITS after the callback, so a commit-abort (serialization/deadlock/
|
||||
// connection drop — the transient class the epic retries in the processor)
|
||||
// rejects OUTSIDE the callback, after saveHistory already ran and the SPOP
|
||||
// already happened, while the INSERT rolls back. onStateless does NOT retry,
|
||||
// so an unrestored pop is a one-shot irrecoverable attribution loss (the
|
||||
// processor got exactly this fix: poppedForRestore + an outer catch). We
|
||||
// track the popped set here (keyed by the page UUID it was popped by — never
|
||||
// the doc-name id, which may be a slugId, #260) and restore it in the outer
|
||||
// catch. addContributors is an idempotent Redis SADD, so a double-restore is
|
||||
// harmless. versionedPageId is also reused below to remove the superseded
|
||||
// idle job by its real jobId (page.id).
|
||||
let poppedForRestore: string[] = [];
|
||||
let versionedPageId: string | undefined;
|
||||
|
||||
try {
|
||||
await executeTx(this.db, async (trx) => {
|
||||
const page = await this.pageRepo.findById(pageId, {
|
||||
withLock: true,
|
||||
includeContent: true,
|
||||
trx,
|
||||
});
|
||||
if (!page) return;
|
||||
versionedPageId = page.id;
|
||||
// Never version an effectively-empty page (mirrors the processor's
|
||||
// first-history guard); there is nothing intentional to pin.
|
||||
if (isEmptyParagraphDoc(page.content as any)) return;
|
||||
|
||||
const lastHistory = await this.pageHistoryRepo.findPageLastHistory(
|
||||
page.id,
|
||||
{ includeContent: true, trx },
|
||||
);
|
||||
|
||||
if (
|
||||
lastHistory &&
|
||||
isDeepStrictEqual(lastHistory.content, page.content)
|
||||
) {
|
||||
// Content is already snapshotted. Promote-not-dup.
|
||||
if (lastHistory.kind === 'manual') {
|
||||
result = {
|
||||
historyId: lastHistory.id,
|
||||
kind: 'manual',
|
||||
alreadySaved: true,
|
||||
};
|
||||
return;
|
||||
}
|
||||
await this.pageHistoryRepo.updateHistoryKind(
|
||||
lastHistory.id,
|
||||
kind,
|
||||
trx,
|
||||
);
|
||||
result = { historyId: lastHistory.id, kind, alreadySaved: false };
|
||||
return;
|
||||
}
|
||||
|
||||
// Fresh version row. Pop the contributors aggregated since the last
|
||||
// snapshot (SPOP); restore them if the write fails so they aren't lost.
|
||||
const contributorIds = await this.collabHistory.popContributors(
|
||||
page.id,
|
||||
);
|
||||
poppedForRestore = contributorIds;
|
||||
try {
|
||||
const saved = await this.pageHistoryRepo.saveHistory(page, {
|
||||
contributorIds,
|
||||
kind,
|
||||
trx,
|
||||
});
|
||||
result = { historyId: saved.id, kind, alreadySaved: false };
|
||||
} catch (err) {
|
||||
await this.collabHistory.addContributors(page.id, contributorIds);
|
||||
poppedForRestore = [];
|
||||
throw err;
|
||||
}
|
||||
});
|
||||
} catch (err) {
|
||||
// A throw here means the tx did NOT commit (callback threw, or the commit
|
||||
// itself failed and rolled back). If we popped contributors and the inner
|
||||
// catch did not already restore them, restore now so attribution is not
|
||||
// lost — onStateless has no retry to recover it. Restore by the page UUID
|
||||
// the pop was keyed under (versionedPageId is always set before the pop).
|
||||
if (poppedForRestore.length && versionedPageId) {
|
||||
await this.collabHistory.addContributors(
|
||||
versionedPageId,
|
||||
poppedForRestore,
|
||||
);
|
||||
}
|
||||
throw err;
|
||||
}
|
||||
|
||||
// Housekeeping: this explicit version supersedes the page's pending idle
|
||||
// autosnapshot, so cancel it and end the current idle burst so the next edit
|
||||
// starts a fresh max-wait window. Remove the idle job by its REAL jobId
|
||||
// (page.id UUID — computeHistoryJob arms it under page.id), not the raw
|
||||
// doc-name id which may be a slugId for a `page.<slugId>` doc (#260), or the
|
||||
// remove silently misses. The burst marker is keyed by documentName (like its
|
||||
// sibling per-document maps), and is also cleaned in afterUnloadDocument.
|
||||
if (versionedPageId) {
|
||||
await this.historyQueue.remove(versionedPageId).catch(() => undefined);
|
||||
}
|
||||
this.idleBurstStart.delete(documentName);
|
||||
|
||||
if (result) {
|
||||
document.broadcastStateless(
|
||||
JSON.stringify({
|
||||
type: 'version.saved',
|
||||
historyId: result.historyId,
|
||||
kind: result.kind,
|
||||
alreadySaved: result.alreadySaved,
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async onChange(data: onChangePayload) {
|
||||
const documentName = data.documentName;
|
||||
const userId = data.context?.user?.id;
|
||||
@@ -586,6 +783,10 @@ export class PersistenceExtension implements Extension {
|
||||
this.contributors.delete(documentName);
|
||||
this.agentTouched.delete(documentName);
|
||||
this.intentionalClear.delete(documentName);
|
||||
// #370 — drop the idle-burst marker with the other per-document maps so it
|
||||
// cannot accumulate across the process lifetime for never-manually-saved
|
||||
// pages. The pending idle job (if any) is a self-expiring BullMQ delayed job.
|
||||
this.idleBurstStart.delete(documentName);
|
||||
}
|
||||
|
||||
private consumeContributors(documentName: string): string[] {
|
||||
@@ -617,19 +818,80 @@ export class PersistenceExtension implements Extension {
|
||||
|
||||
private async enqueuePageHistory(
|
||||
page: Page,
|
||||
documentName: string,
|
||||
lastUpdatedSource: string,
|
||||
): Promise<void> {
|
||||
// Job id + delay arithmetic lives in the pure `computeHistoryJob` (see its
|
||||
// doc comment for the agent-delay-0 / age-based-debounce invariants).
|
||||
// #370 — trailing idle debounce with a max-wait ceiling. One pending idle
|
||||
// job per page (jobId = page.id); on every store we remove the pending
|
||||
// delayed job and re-add it, so the snapshot lands `delay` after edits go
|
||||
// quiet rather than once per store (precedent: workspace.service.ts).
|
||||
// remove() on a delayed job simply deletes it (0 if absent, no throw); if the
|
||||
// job is already ACTIVE and the remove is a no-op, the add still de-dups and
|
||||
// the processor's isDeepStrictEqual gate collapses the duplicate content.
|
||||
//
|
||||
// The FIRST arm of a burst records `burstStart`; computeHistoryJob shrinks
|
||||
// the delay to the remaining max-wait budget from that point, so a continuous
|
||||
// session cannot re-arm the trailing timer forever and starve the snapshot.
|
||||
// A burst marker older than THIS TIER's max-wait means the previous idle job
|
||||
// has already fired — start a fresh window instead of firing immediately on
|
||||
// the next edit. Must use the SAME source-specific max-wait computeHistoryJob
|
||||
// uses (agent 5m / user 10m): a hardcoded USER ceiling would leave an agent
|
||||
// burst's marker stale for 5..10m, forcing delay=0 on every store in that
|
||||
// window and writing one idle row per store — exactly the per-store bloat the
|
||||
// debounce exists to prevent, on the continuous-agent path.
|
||||
const maxWait =
|
||||
lastUpdatedSource === 'agent' ? IDLE_MAX_WAIT_AGENT : IDLE_MAX_WAIT_USER;
|
||||
const now = Date.now();
|
||||
// Keyed by documentName (see the map declaration) so afterUnloadDocument can
|
||||
// clean it; the queue jobId stays page.id (computeHistoryJob) as required.
|
||||
let burstStart = this.idleBurstStart.get(documentName);
|
||||
if (burstStart === undefined || now - burstStart >= maxWait) {
|
||||
burstStart = now;
|
||||
this.idleBurstStart.set(documentName, burstStart);
|
||||
}
|
||||
|
||||
const { jobId, delay } = computeHistoryJob(
|
||||
page,
|
||||
lastUpdatedSource,
|
||||
Date.now(),
|
||||
burstStart,
|
||||
now,
|
||||
);
|
||||
|
||||
// remove-then-add trailing-debounce idiom, and its ONE race. We delete the
|
||||
// pending delayed job and re-add it under the same jobId so the timer resets
|
||||
// to the trailing edge of the burst. The race is the small window between
|
||||
// these two awaits: if the delayed job's `delay` elapses in that gap it goes
|
||||
// ACTIVE, and then:
|
||||
// - remove() on an active/locked job is a no-op (BullMQ won't yank a job a
|
||||
// worker holds), and our `.catch(() => undefined)` swallows that too; and
|
||||
// - add() with a jobId that already exists (the now-active job's id) is
|
||||
// DROPPED by BullMQ — a duplicate add is a no-op.
|
||||
// So this store fails to re-arm the trailing job: the just-fired snapshot
|
||||
// captured content up to the moment it went active, and THIS edit is left
|
||||
// without a pending trailing job. It is bounded and self-healing — the NEXT
|
||||
// store re-arms a fresh delayed job (the id is free again once the active job
|
||||
// completes / removeOnComplete frees it), and the processor's
|
||||
// isDeepStrictEqual gate collapses any content-identical duplicate. The only
|
||||
// uncovered case is when the racing store was the LAST in the session: the
|
||||
// tail edits made after the job went active get NO trailing snapshot until
|
||||
// the next edit re-arms one. That is an acceptable safety-net gap (a manual
|
||||
// Save, a source-transition boundary, or simply the next edit all still cover
|
||||
// it), which is why the reviewer accepts documenting it here rather than
|
||||
// adding a post-add "did the add actually arm a job?" re-check.
|
||||
//
|
||||
// NOTE — do NOT "unify" this with the neighbouring embed-debounce idiom
|
||||
// (aiQueue.add of PAGE_CONTENT_UPDATED above): that one uses a STABLE jobId
|
||||
// and NO remove(), relying purely on BullMQ coalescing a repeated add under
|
||||
// the same id, because a re-embed only needs to eventually run once on the
|
||||
// latest content and re-anchoring its delay on every keystroke is undesirable.
|
||||
// THIS idiom deliberately removes-then-adds precisely to PUSH the delay back
|
||||
// to the trailing edge on every store (a true debounce), which coalescing
|
||||
// alone cannot do. Collapsing them would silently change the history cadence.
|
||||
await this.historyQueue.remove(jobId).catch(() => undefined);
|
||||
|
||||
await this.historyQueue.add(
|
||||
QueueJob.PAGE_HISTORY,
|
||||
{ pageId: page.id } as IPageHistoryJob,
|
||||
{ pageId: page.id, kind: 'idle' } as IPageHistoryJob,
|
||||
{ jobId, delay },
|
||||
);
|
||||
}
|
||||
|
||||
@@ -66,6 +66,15 @@ describe('HistoryProcessor.process', () => {
|
||||
notificationQueue = { add: jest.fn().mockResolvedValue(undefined) };
|
||||
generalQueue = { add: jest.fn().mockResolvedValue(undefined) };
|
||||
|
||||
// #370 F3 — the processor now serializes its find+save under a page-row lock
|
||||
// via executeTx. A db whose transaction().execute(fn) runs fn with a trx stub
|
||||
// drives the real executeTx() helper without a database.
|
||||
const db = {
|
||||
transaction: () => ({
|
||||
execute: (fn: (trx: any) => Promise<any>) => fn({ __trx: true }),
|
||||
}),
|
||||
};
|
||||
|
||||
// WorkerHost's constructor reads `this.worker`; passing repos positionally
|
||||
// matches the constructor and avoids the Nest DI container.
|
||||
proc = new HistoryProcessor(
|
||||
@@ -73,6 +82,7 @@ describe('HistoryProcessor.process', () => {
|
||||
pageRepo as any,
|
||||
collabHistory as any,
|
||||
watcherService as any,
|
||||
db as any,
|
||||
notificationQueue as any,
|
||||
generalQueue as any,
|
||||
);
|
||||
@@ -126,15 +136,26 @@ describe('HistoryProcessor.process', () => {
|
||||
await proc.process(buildJob());
|
||||
|
||||
expect(collabHistory.popContributors).toHaveBeenCalledWith(PAGE_ID);
|
||||
// #370 F3/F9 — the snapshot decision runs under a page-row lock. Pin the lock
|
||||
// structurally so a refactor that drops withLock/trx (silently reintroducing
|
||||
// the TOCTOU double-insert) turns this red. The tx stub is { __trx: true }.
|
||||
expect(pageRepo.findById).toHaveBeenCalledWith(
|
||||
PAGE_ID,
|
||||
expect.objectContaining({ withLock: true, trx: { __trx: true } }),
|
||||
);
|
||||
// #370 F7 — addPageWatchers MUST receive the trx, or its FK-check runs on a
|
||||
// separate connection and self-deadlocks against our FOR UPDATE. Asserting
|
||||
// the trx arg here is exactly what would have caught that regression.
|
||||
expect(watcherService.addPageWatchers).toHaveBeenCalledWith(
|
||||
['u1', 'u2'],
|
||||
PAGE_ID,
|
||||
SPACE_ID,
|
||||
WORKSPACE_ID,
|
||||
{ __trx: true },
|
||||
);
|
||||
expect(pageHistoryRepo.saveHistory).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ id: PAGE_ID }),
|
||||
{ contributorIds: ['u1', 'u2'] },
|
||||
{ contributorIds: ['u1', 'u2'], kind: 'idle', trx: { __trx: true } },
|
||||
);
|
||||
expect(generalQueue.add).toHaveBeenCalledWith(
|
||||
QueueJob.PAGE_BACKLINKS,
|
||||
@@ -186,6 +207,48 @@ describe('HistoryProcessor.process', () => {
|
||||
]);
|
||||
});
|
||||
|
||||
it('COMMIT failure (throw outside the tx callback) → contributors RESTORED', async () => {
|
||||
// #370 F8 — a commit-time failure throws OUTSIDE the callback, so the inner
|
||||
// try/catch does not run; the outer catch must restore the popped set (else a
|
||||
// BullMQ retry writes an unattributed version). Use a db whose execute() runs
|
||||
// the callback THEN throws, simulating a commit abort.
|
||||
pageHistoryRepo.findPageLastHistory.mockResolvedValue({
|
||||
content: { type: 'doc', content: [] },
|
||||
});
|
||||
const commitFail = {
|
||||
transaction: () => ({
|
||||
execute: async (fn: (trx: any) => Promise<any>) => {
|
||||
await fn({ __trx: true }); // callback succeeds (saveHistory ok)
|
||||
throw new Error('commit aborted'); // ...but the COMMIT fails
|
||||
},
|
||||
}),
|
||||
};
|
||||
const procCommitFail = new HistoryProcessor(
|
||||
pageHistoryRepo as any,
|
||||
pageRepo as any,
|
||||
collabHistory as any,
|
||||
watcherService as any,
|
||||
commitFail as any,
|
||||
notificationQueue as any,
|
||||
generalQueue as any,
|
||||
);
|
||||
jest
|
||||
.spyOn(procCommitFail['logger'], 'error')
|
||||
.mockImplementation(() => undefined);
|
||||
|
||||
await expect(procCommitFail.process(buildJob())).rejects.toThrow(
|
||||
'commit aborted',
|
||||
);
|
||||
// The inner catch did NOT run (save succeeded), so only the outer catch can
|
||||
// restore — assert it did.
|
||||
expect(collabHistory.addContributors).toHaveBeenCalledWith(PAGE_ID, [
|
||||
'u1',
|
||||
'u2',
|
||||
]);
|
||||
// And the post-snapshot queue work must NOT have run (we rethrew).
|
||||
expect(generalQueue.add).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('backlinks + notification queue failures are swallowed (history still committed)', async () => {
|
||||
pageHistoryRepo.findPageLastHistory.mockResolvedValue({
|
||||
content: { type: 'doc', content: [] },
|
||||
|
||||
@@ -19,6 +19,9 @@ import { isDeepStrictEqual } from 'node:util';
|
||||
import { CollabHistoryService } from '../services/collab-history.service';
|
||||
import { WatcherService } from '../../core/watcher/watcher.service';
|
||||
import { isEmptyParagraphDoc } from '../collaboration.util';
|
||||
import { InjectKysely } from 'nestjs-kysely';
|
||||
import { KyselyDB } from '@docmost/db/types/kysely.types';
|
||||
import { executeTx } from '@docmost/db/utils';
|
||||
|
||||
@Processor(QueueName.HISTORY_QUEUE)
|
||||
export class HistoryProcessor extends WorkerHost implements OnModuleDestroy {
|
||||
@@ -29,6 +32,7 @@ export class HistoryProcessor extends WorkerHost implements OnModuleDestroy {
|
||||
private readonly pageRepo: PageRepo,
|
||||
private readonly collabHistory: CollabHistoryService,
|
||||
private readonly watcherService: WatcherService,
|
||||
@InjectKysely() private readonly db: KyselyDB,
|
||||
@InjectQueue(QueueName.NOTIFICATION_QUEUE) private notificationQueue: Queue,
|
||||
@InjectQueue(QueueName.GENERAL_QUEUE) private generalQueue: Queue,
|
||||
) {
|
||||
@@ -41,6 +45,9 @@ export class HistoryProcessor extends WorkerHost implements OnModuleDestroy {
|
||||
try {
|
||||
const { pageId } = job.data;
|
||||
|
||||
// Read the page WITHOUT a lock first, only to bail early on the two cheap
|
||||
// no-write cases (page gone / empty first snapshot) without opening a
|
||||
// transaction. The authoritative check-then-write happens locked below.
|
||||
const page = await this.pageRepo.findById(pageId, {
|
||||
includeContent: true,
|
||||
});
|
||||
@@ -51,40 +58,109 @@ export class HistoryProcessor extends WorkerHost implements OnModuleDestroy {
|
||||
return;
|
||||
}
|
||||
|
||||
const lastHistory = await this.pageHistoryRepo.findPageLastHistory(
|
||||
pageId,
|
||||
{ includeContent: true },
|
||||
);
|
||||
// #370 F3 — the snapshot decision (findPageLastHistory → saveHistory) must
|
||||
// be serialized against manual-save/boundary writers, which run under a
|
||||
// page-row lock in onStoreDocument. Without it, this processor and a
|
||||
// concurrent manual-save each read the same lastHistory (MVCC), both see
|
||||
// content != lastHistory, and both insert — producing two page_history rows
|
||||
// with IDENTICAL content (one 'idle', one 'manual'), defeating
|
||||
// promote-not-dup and the version-vs-autosave split. Taking the same
|
||||
// page-row lock makes the second writer observe the first's committed row so
|
||||
// the isDeepStrictEqual gate collapses the duplicate. Only the read+write
|
||||
// is transacted; the post-snapshot queue work stays outside.
|
||||
let contributorIds: string[] = [];
|
||||
let snapshotWritten = false;
|
||||
let lastHistoryContent: unknown;
|
||||
// #370 F8 — the contributor set popped from Redis (destructive SPOP) must be
|
||||
// restored if the snapshot does not durably land. The inner try/catch only
|
||||
// covers a throw INSIDE the callback; a COMMIT failure (connection drop,
|
||||
// serialization/deadlock abort on commit — the transient class the epic
|
||||
// already retries) throws OUTSIDE it, rolling the snapshot back while the
|
||||
// pop is already gone. We track the popped set here and restore it in the
|
||||
// outer catch so a BullMQ retry re-attributes the version. addContributors
|
||||
// is an idempotent Redis SADD, so a double-restore is harmless.
|
||||
let poppedForRestore: string[] = [];
|
||||
|
||||
if (!lastHistory && isEmptyParagraphDoc(page.content as any)) {
|
||||
this.logger.debug(
|
||||
`Skipping first history for page ${pageId}: empty content`,
|
||||
);
|
||||
await this.collabHistory.clearContributors(pageId);
|
||||
try {
|
||||
await executeTx(this.db, async (trx) => {
|
||||
const lockedPage = await this.pageRepo.findById(pageId, {
|
||||
includeContent: true,
|
||||
withLock: true,
|
||||
trx,
|
||||
});
|
||||
if (!lockedPage) return;
|
||||
|
||||
const lastHistory = await this.pageHistoryRepo.findPageLastHistory(
|
||||
pageId,
|
||||
{ includeContent: true, trx },
|
||||
);
|
||||
lastHistoryContent = lastHistory?.content;
|
||||
|
||||
if (!lastHistory && isEmptyParagraphDoc(lockedPage.content as any)) {
|
||||
this.logger.debug(
|
||||
`Skipping first history for page ${pageId}: empty content`,
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
if (
|
||||
lastHistory &&
|
||||
isDeepStrictEqual(lastHistory.content, lockedPage.content)
|
||||
) {
|
||||
return; // already snapshotted at this content — nothing to write
|
||||
}
|
||||
|
||||
contributorIds = await this.collabHistory.popContributors(pageId);
|
||||
poppedForRestore = contributorIds;
|
||||
try {
|
||||
// Pass `trx` so the watcher insert's FK check (FOR KEY SHARE on
|
||||
// pages[pageId]) runs on the SAME connection that already holds the
|
||||
// FOR UPDATE lock from findById — otherwise it takes the FK lock on a
|
||||
// separate pool connection and self-deadlocks against our own tx.
|
||||
await this.watcherService.addPageWatchers(
|
||||
contributorIds,
|
||||
pageId,
|
||||
lockedPage.spaceId,
|
||||
lockedPage.workspaceId,
|
||||
trx,
|
||||
);
|
||||
|
||||
// #370 — every job on this queue is a trailing idle-flush autosnapshot.
|
||||
await this.pageHistoryRepo.saveHistory(lockedPage, {
|
||||
contributorIds,
|
||||
kind: job.data.kind ?? 'idle',
|
||||
trx,
|
||||
});
|
||||
snapshotWritten = true;
|
||||
this.logger.debug(`History created for page: ${pageId}`);
|
||||
} catch (err) {
|
||||
await this.collabHistory.addContributors(pageId, contributorIds);
|
||||
poppedForRestore = [];
|
||||
throw err;
|
||||
}
|
||||
});
|
||||
} catch (err) {
|
||||
// A throw here means the tx did NOT commit (callback threw, or the commit
|
||||
// itself failed and rolled back). If we popped contributors and the inner
|
||||
// catch did not already restore them, restore now so the retry keeps
|
||||
// attribution. snapshotWritten is irrelevant: it is set before commit, so
|
||||
// it can be true even when the commit rolled the snapshot back.
|
||||
if (poppedForRestore.length) {
|
||||
await this.collabHistory.addContributors(pageId, poppedForRestore);
|
||||
}
|
||||
throw err;
|
||||
}
|
||||
|
||||
// No snapshot written (page vanished / empty-first / unchanged content) →
|
||||
// clear the contributor set for the skip cases and stop.
|
||||
if (!snapshotWritten) {
|
||||
if (!lastHistoryContent && isEmptyParagraphDoc(page.content as any)) {
|
||||
await this.collabHistory.clearContributors(pageId);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (
|
||||
!lastHistory ||
|
||||
!isDeepStrictEqual(lastHistory.content, page.content)
|
||||
) {
|
||||
const contributorIds = await this.collabHistory.popContributors(pageId);
|
||||
|
||||
try {
|
||||
await this.watcherService.addPageWatchers(
|
||||
contributorIds,
|
||||
pageId,
|
||||
page.spaceId,
|
||||
page.workspaceId,
|
||||
);
|
||||
|
||||
await this.pageHistoryRepo.saveHistory(page, { contributorIds });
|
||||
this.logger.debug(`History created for page: ${pageId}`);
|
||||
} catch (err) {
|
||||
await this.collabHistory.addContributors(pageId, contributorIds);
|
||||
throw err;
|
||||
}
|
||||
|
||||
{
|
||||
const mentions = extractMentions(page.content);
|
||||
const pageMentions = extractPageMentions(mentions);
|
||||
const internalLinkSlugIds = extractInternalLinkSlugIds(page.content);
|
||||
@@ -102,7 +178,7 @@ export class HistoryProcessor extends WorkerHost implements OnModuleDestroy {
|
||||
);
|
||||
});
|
||||
|
||||
if (contributorIds.length > 0 && lastHistory?.content) {
|
||||
if (contributorIds.length > 0 && lastHistoryContent) {
|
||||
await this.notificationQueue
|
||||
.add(QueueJob.PAGE_UPDATED, {
|
||||
pageId,
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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',
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -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();
|
||||
|
||||
@@ -189,10 +189,11 @@ export function stepBudgetWarning(stepNumber: number): string {
|
||||
//
|
||||
// `system` is the in-scope system prompt; we CONCATENATE so the original
|
||||
// persona/context is preserved — a bare `system` override would REPLACE the
|
||||
// whole system prompt for the step. `activatedTools` is PER-TURN mutable state
|
||||
// owned by the streaming loop (a closure Set grown by loadTools); it is passed
|
||||
// in (not module-global, not persisted) so this stays a pure function of its
|
||||
// arguments.
|
||||
// whole system prompt for the step. `activatedTools` is a closure Set grown by
|
||||
// loadTools and owned by the streaming loop; the caller seeds it from and
|
||||
// persists it to the chat's metadata across turns (#490), but this function only
|
||||
// READS the Set it is handed, so it stays a pure function of its arguments (not
|
||||
// module-global).
|
||||
//
|
||||
// NOTE: at AI SDK v7 the per-step `system` field is renamed to `instructions`.
|
||||
// On v6 (`^6.0.134`) `system` is the correct field — adjust when bumping.
|
||||
@@ -1410,10 +1411,11 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
const baseTools = { ...external.tools, ...docmostTools };
|
||||
|
||||
// Deferred tool loading state (#332), scoped to THIS streaming loop:
|
||||
// - `activatedTools` is per-TURN mutable state — a fresh closure Set created
|
||||
// per streamText call, NOT module-global and NOT persisted, so a new turn
|
||||
// starts cold. loadTools.execute adds to it; prepareAgentStep reads it to
|
||||
// widen `activeTools` on the NEXT step.
|
||||
// - `activatedTools` is a fresh closure Set per streamText call (not
|
||||
// module-global), SEEDED from the chat's persisted metadata.activatedTools
|
||||
// (#490, just below) so activation carries across turns. loadTools.execute
|
||||
// adds to it; prepareAgentStep reads it to widen `activeTools` on the NEXT
|
||||
// step; turn end persists it back.
|
||||
// - `validDeferredNames` = every tool that is NOT core (the in-app deferred
|
||||
// tools + ALL external MCP tools), computed from the ACTUAL toolset so an
|
||||
// external tool is loadable by its namespaced name. loadTools rejects any
|
||||
@@ -1543,30 +1545,39 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
// Per-step (non-terminal) update: persist the finished steps the moment a
|
||||
// step ends. Tolerant — a failed update is logged and swallowed so it never
|
||||
// throws into the stream. Keeps status 'streaming'.
|
||||
const updateStreaming = async (): Promise<void> => {
|
||||
if (!assistantId) return;
|
||||
//
|
||||
// #491: it now SIGNALS its outcome — the persisted `stepsPersisted` count on
|
||||
// a CONFIRMED write, or null when it was skipped/failed. The caller rotates
|
||||
// the run-stream registry ring ONLY on a non-null return (a confirmed
|
||||
// persist), so a failed persist never rotates away a step nobody has (the
|
||||
// classic inversion bug); a failure just makes the ring cover more.
|
||||
const updateStreaming = async (): Promise<number | null> => {
|
||||
if (!assistantId) return null;
|
||||
// Cheap short-circuit once the turn is finalized (see `finalized` below).
|
||||
// The AUTHORITATIVE guard is `onlyIfStreaming` on the UPDATE: a late
|
||||
// fire-and-forget step update could still be in flight on another pool
|
||||
// connection when finalize runs, so the SQL `WHERE status='streaming'`
|
||||
// (not this flag) is what prevents it clobbering the terminal row.
|
||||
if (finalized) return;
|
||||
if (finalized) return null;
|
||||
// Build the flush ONCE so the returned count is EXACTLY the persisted
|
||||
// `stepsPersisted` (both derive from capturedSteps.length at this instant).
|
||||
const flushed = flushAssistant(capturedSteps, '', 'streaming', {
|
||||
pageChanged,
|
||||
partsCache,
|
||||
});
|
||||
const stepsPersisted = flushed.metadata.stepsPersisted as number;
|
||||
try {
|
||||
await this.aiChatMessageRepo.update(
|
||||
assistantId,
|
||||
workspace.id,
|
||||
flushAssistant(capturedSteps, '', 'streaming', {
|
||||
pageChanged,
|
||||
partsCache,
|
||||
}),
|
||||
{ onlyIfStreaming: true },
|
||||
);
|
||||
await this.aiChatMessageRepo.update(assistantId, workspace.id, flushed, {
|
||||
onlyIfStreaming: true,
|
||||
});
|
||||
return stepsPersisted;
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`Failed to update streaming assistant row: ${
|
||||
err instanceof Error ? err.message : 'unknown error'
|
||||
}`,
|
||||
);
|
||||
return null;
|
||||
}
|
||||
};
|
||||
|
||||
@@ -1742,7 +1753,24 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
// this point still recovers the step. Not awaited here (never block the
|
||||
// stream), but SERIALIZED via stepUpdateChain so the writes commit in
|
||||
// step order; updateStreaming is error-tolerant (logs + swallows).
|
||||
stepUpdateChain = stepUpdateChain.then(() => updateStreaming());
|
||||
// #491: on a CONFIRMED persist, rotate the run-stream registry ring to
|
||||
// drop the now-on-disk steps (stamp < stepsPersisted). Gated on the
|
||||
// resumable flag (same as open/bind) and identity-checked in the
|
||||
// registry; a null return (skipped/failed) rotates NOTHING (auto-safe).
|
||||
stepUpdateChain = stepUpdateChain.then(async () => {
|
||||
const persisted = await updateStreaming();
|
||||
if (
|
||||
persisted != null &&
|
||||
runId &&
|
||||
this.environment?.isAiChatResumableStreamEnabled?.()
|
||||
) {
|
||||
this.streamRegistry?.confirmPersistedStep(
|
||||
chatId,
|
||||
runId,
|
||||
persisted,
|
||||
);
|
||||
}
|
||||
});
|
||||
// #184: persist the run's progress (finished-step count). Fire-and-
|
||||
// forget; the hook swallows its own errors.
|
||||
if (runId) runHooks?.onStep?.(runId, capturedSteps.length);
|
||||
@@ -2891,6 +2919,18 @@ export function flushAssistant(
|
||||
// `parts`). Old rows have no marker and the legacy { output } shape; a
|
||||
// dual-shape query branches on this. Old rows are deliberately NOT migrated.
|
||||
toolTraceVersion: 2,
|
||||
// #491 STEP MARKER: the number of FINISHED steps whose parts are in THIS row,
|
||||
// written by the SAME flush that builds `parts` (atomically — they are both
|
||||
// derived from `finished`, so the marker can NEVER disagree with the persisted
|
||||
// parts). This is the step-alignment anchor the resume stack builds on:
|
||||
// - the registry rotates its retention ring only on a CONFIRMED persist of
|
||||
// step N (commit 3);
|
||||
// - attach slices the tail at "step > N" from the client's persisted seed.
|
||||
// It is NOT `run.stepCount`: recordStep is fire-and-forget and NOT atomic with
|
||||
// the parts write, so stepCount could race ahead of the persisted parts
|
||||
// (seed↔marker drift). The in-progress trailing text (an error/abort partial,
|
||||
// or a mid-stream flush) is NOT a finished step and is excluded from the count.
|
||||
stepsPersisted: finished.length,
|
||||
};
|
||||
// finishReason: prefer an explicit one; else derive a sensible value from the
|
||||
// terminal status (so onError/onAbort records keep their historical reason).
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
@@ -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,193 @@
|
||||
import { ForbiddenException, NotFoundException } from '@nestjs/common';
|
||||
import { ApiKeyController } from './api-key.controller';
|
||||
import {
|
||||
WorkspaceCaslAction,
|
||||
WorkspaceCaslSubject,
|
||||
} from '../casl/interfaces/workspace-ability.type';
|
||||
|
||||
/**
|
||||
* Authorization contract for the /api-keys management surface.
|
||||
*
|
||||
* - A token cannot manage tokens (an api_key PRINCIPAL is 403 on every method)
|
||||
* — GitHub-PAT semantics closing post-revocation laundering.
|
||||
* - admin (CASL Manage on API) sees/revokes all workspace keys; a member only
|
||||
* their own.
|
||||
* - kill-switch OFF -> the surface 404s (looks like the feature does not exist).
|
||||
*/
|
||||
|
||||
function makeController(over: any = {}) {
|
||||
const apiKeyService = {
|
||||
create: jest.fn().mockResolvedValue({
|
||||
token: 'tok',
|
||||
key: { id: 'k1', name: 'n', expiresAt: null, createdAt: new Date() },
|
||||
}),
|
||||
revoke: jest.fn().mockResolvedValue(undefined),
|
||||
...(over.apiKeyService ?? {}),
|
||||
};
|
||||
const apiKeyRepo = {
|
||||
findAllInWorkspace: jest.fn().mockResolvedValue([]),
|
||||
findByCreator: jest.fn().mockResolvedValue([]),
|
||||
findById: jest.fn(),
|
||||
...(over.apiKeyRepo ?? {}),
|
||||
};
|
||||
const workspaceAbility = {
|
||||
createForUser: jest.fn().mockReturnValue({
|
||||
can: (_a: any, _s: any) => over.canManage ?? false,
|
||||
}),
|
||||
...(over.workspaceAbility ?? {}),
|
||||
};
|
||||
const environmentService = {
|
||||
isApiKeysEnabled: jest.fn().mockReturnValue(over.enabled ?? true),
|
||||
...(over.environmentService ?? {}),
|
||||
};
|
||||
const auditService = { log: jest.fn() };
|
||||
const controller = new ApiKeyController(
|
||||
apiKeyService as any,
|
||||
apiKeyRepo as any,
|
||||
workspaceAbility as any,
|
||||
environmentService as any,
|
||||
auditService as any,
|
||||
);
|
||||
return {
|
||||
controller,
|
||||
apiKeyService,
|
||||
apiKeyRepo,
|
||||
workspaceAbility,
|
||||
environmentService,
|
||||
auditService,
|
||||
};
|
||||
}
|
||||
|
||||
const user = { id: 'u-1', workspaceId: 'ws-1' } as any;
|
||||
const workspace = { id: 'ws-1' } as any;
|
||||
const reqAccess = () => ({ raw: {}, ip: '1.2.3.4', socket: {} }) as any;
|
||||
const reqApiKey = () =>
|
||||
({ raw: { authType: 'api_key', apiKeyId: 'k-self' }, ip: '1.2.3.4', socket: {} }) as any;
|
||||
|
||||
describe('ApiKeyController — a token cannot manage tokens', () => {
|
||||
it('403 on create for an api_key principal', async () => {
|
||||
const { controller, apiKeyService } = makeController();
|
||||
await expect(
|
||||
controller.create({ name: 'x' } as any, user, reqApiKey()),
|
||||
).rejects.toBeInstanceOf(ForbiddenException);
|
||||
expect(apiKeyService.create).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('403 on list for an api_key principal', async () => {
|
||||
const { controller } = makeController();
|
||||
await expect(
|
||||
controller.list(user, workspace, reqApiKey()),
|
||||
).rejects.toBeInstanceOf(ForbiddenException);
|
||||
});
|
||||
|
||||
it('403 on revoke for an api_key principal', async () => {
|
||||
const { controller } = makeController();
|
||||
await expect(
|
||||
controller.revoke({ id: 'k1' } as any, user, workspace, reqApiKey()),
|
||||
).rejects.toBeInstanceOf(ForbiddenException);
|
||||
});
|
||||
});
|
||||
|
||||
describe('ApiKeyController — kill-switch OFF → 404', () => {
|
||||
it('create/list/revoke all 404 when disabled', async () => {
|
||||
const { controller } = makeController({ enabled: false });
|
||||
await expect(
|
||||
controller.create({ name: 'x' } as any, user, reqAccess()),
|
||||
).rejects.toBeInstanceOf(NotFoundException);
|
||||
await expect(
|
||||
controller.list(user, workspace, reqAccess()),
|
||||
).rejects.toBeInstanceOf(NotFoundException);
|
||||
await expect(
|
||||
controller.revoke({ id: 'k1' } as any, user, workspace, reqAccess()),
|
||||
).rejects.toBeInstanceOf(NotFoundException);
|
||||
});
|
||||
});
|
||||
|
||||
describe('ApiKeyController — list scoping', () => {
|
||||
it('admin (Manage on API) lists ALL workspace keys', async () => {
|
||||
const { controller, apiKeyRepo } = makeController({ canManage: true });
|
||||
await controller.list(user, workspace, reqAccess());
|
||||
expect(apiKeyRepo.findAllInWorkspace).toHaveBeenCalledWith('ws-1');
|
||||
expect(apiKeyRepo.findByCreator).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('member lists ONLY their own keys', async () => {
|
||||
const { controller, apiKeyRepo } = makeController({ canManage: false });
|
||||
await controller.list(user, workspace, reqAccess());
|
||||
expect(apiKeyRepo.findByCreator).toHaveBeenCalledWith('u-1', 'ws-1');
|
||||
expect(apiKeyRepo.findAllInWorkspace).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('ApiKeyController — revoke scoping', () => {
|
||||
it("member cannot revoke another user's key (403)", async () => {
|
||||
const { controller, apiKeyRepo, apiKeyService } = makeController({
|
||||
canManage: false,
|
||||
});
|
||||
apiKeyRepo.findById.mockResolvedValue({
|
||||
id: 'k1',
|
||||
creatorId: 'someone-else',
|
||||
workspaceId: 'ws-1',
|
||||
});
|
||||
await expect(
|
||||
controller.revoke({ id: 'k1' } as any, user, workspace, reqAccess()),
|
||||
).rejects.toBeInstanceOf(ForbiddenException);
|
||||
expect(apiKeyService.revoke).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('member CAN revoke their own key', async () => {
|
||||
const { controller, apiKeyRepo, apiKeyService } = makeController({
|
||||
canManage: false,
|
||||
});
|
||||
apiKeyRepo.findById.mockResolvedValue({
|
||||
id: 'k1',
|
||||
creatorId: 'u-1',
|
||||
workspaceId: 'ws-1',
|
||||
});
|
||||
await controller.revoke({ id: 'k1' } as any, user, workspace, reqAccess());
|
||||
expect(apiKeyService.revoke).toHaveBeenCalledWith('k1', 'ws-1');
|
||||
});
|
||||
|
||||
it("admin CAN revoke another user's key", async () => {
|
||||
const { controller, apiKeyRepo, apiKeyService } = makeController({
|
||||
canManage: true,
|
||||
});
|
||||
apiKeyRepo.findById.mockResolvedValue({
|
||||
id: 'k1',
|
||||
creatorId: 'someone-else',
|
||||
workspaceId: 'ws-1',
|
||||
});
|
||||
await controller.revoke({ id: 'k1' } as any, user, workspace, reqAccess());
|
||||
expect(apiKeyService.revoke).toHaveBeenCalledWith('k1', 'ws-1');
|
||||
});
|
||||
|
||||
it('404 when the key does not exist', async () => {
|
||||
const { controller, apiKeyRepo } = makeController({ canManage: true });
|
||||
apiKeyRepo.findById.mockResolvedValue(undefined);
|
||||
await expect(
|
||||
controller.revoke({ id: 'k1' } as any, user, workspace, reqAccess()),
|
||||
).rejects.toBeInstanceOf(NotFoundException);
|
||||
});
|
||||
});
|
||||
|
||||
describe('ApiKeyController — create emits audit + returns token once', () => {
|
||||
it('logs API_KEY_CREATED and returns the token + computed expiry', async () => {
|
||||
const { controller, auditService } = makeController();
|
||||
const res = await controller.create(
|
||||
{ name: 'ci' } as any,
|
||||
user,
|
||||
reqAccess(),
|
||||
);
|
||||
expect(res.token).toBe('tok');
|
||||
expect(res.apiKey).toMatchObject({ id: 'k1', name: 'n' });
|
||||
expect(auditService.log).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ event: 'api_key.created' }),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
// Sanity: the CASL subject used is the workspace API subject.
|
||||
it('uses WorkspaceCaslSubject.API with the Manage action', () => {
|
||||
expect(WorkspaceCaslSubject.API).toBe('api_key');
|
||||
expect(WorkspaceCaslAction.Manage).toBeDefined();
|
||||
});
|
||||
@@ -0,0 +1,201 @@
|
||||
import {
|
||||
Body,
|
||||
Controller,
|
||||
ForbiddenException,
|
||||
HttpCode,
|
||||
HttpStatus,
|
||||
Inject,
|
||||
Logger,
|
||||
NotFoundException,
|
||||
Post,
|
||||
Req,
|
||||
UseGuards,
|
||||
} from '@nestjs/common';
|
||||
import { Throttle } from '@nestjs/throttler';
|
||||
import { FastifyRequest } from 'fastify';
|
||||
import { ApiKeyService } from './api-key.service';
|
||||
import { ApiKeyRepo } from '@docmost/db/repos/api-key/api-key.repo';
|
||||
import { CreateApiKeyDto } from './dto/create-api-key.dto';
|
||||
import { RevokeApiKeyDto } from './dto/revoke-api-key.dto';
|
||||
import { AuthUser } from '../../common/decorators/auth-user.decorator';
|
||||
import { AuthWorkspace } from '../../common/decorators/auth-workspace.decorator';
|
||||
import { JwtAuthGuard } from '../../common/guards/jwt-auth.guard';
|
||||
import { User, Workspace } from '@docmost/db/types/entity.types';
|
||||
import WorkspaceAbilityFactory from '../casl/abilities/workspace-ability.factory';
|
||||
import {
|
||||
WorkspaceCaslAction,
|
||||
WorkspaceCaslSubject,
|
||||
} from '../casl/interfaces/workspace-ability.type';
|
||||
import { AuditEvent, AuditResource } from '../../common/events/audit-events';
|
||||
import {
|
||||
AUDIT_SERVICE,
|
||||
IAuditService,
|
||||
} from '../../integrations/audit/audit.service';
|
||||
import { EnvironmentService } from '../../integrations/environment/environment.service';
|
||||
import { UserThrottlerGuard } from '../../integrations/throttle/user-throttler.guard';
|
||||
import { AUTH_THROTTLER } from '../../integrations/throttle/throttler-names';
|
||||
|
||||
@UseGuards(JwtAuthGuard)
|
||||
@Controller('api-keys')
|
||||
export class ApiKeyController {
|
||||
private readonly logger = new Logger(ApiKeyController.name);
|
||||
|
||||
constructor(
|
||||
private readonly apiKeyService: ApiKeyService,
|
||||
private readonly apiKeyRepo: ApiKeyRepo,
|
||||
private readonly workspaceAbility: WorkspaceAbilityFactory,
|
||||
private readonly environmentService: EnvironmentService,
|
||||
@Inject(AUDIT_SERVICE) private readonly auditService: IAuditService,
|
||||
) {}
|
||||
|
||||
// Kill-switch OFF: the issuance surface disappears entirely (404), not a 403 —
|
||||
// it must look like the feature does not exist.
|
||||
private assertEnabled(): void {
|
||||
if (!this.environmentService.isApiKeysEnabled()) {
|
||||
throw new NotFoundException();
|
||||
}
|
||||
}
|
||||
|
||||
// A token cannot manage tokens (GitHub-PAT semantics): an api-key principal is
|
||||
// refused on the management surface, so a leaked key cannot mint a replacement
|
||||
// or revoke the keys that would lock it out (closes post-revocation laundering).
|
||||
private rejectApiKeyPrincipal(req: FastifyRequest): void {
|
||||
if ((req.raw as { authType?: string }).authType === 'api_key') {
|
||||
throw new ForbiddenException('API keys cannot manage API keys');
|
||||
}
|
||||
}
|
||||
|
||||
private clientIp(req: FastifyRequest): string {
|
||||
return req.ip ?? req.socket?.remoteAddress ?? 'unknown';
|
||||
}
|
||||
|
||||
@HttpCode(HttpStatus.OK)
|
||||
@UseGuards(JwtAuthGuard, UserThrottlerGuard)
|
||||
@Throttle({ [AUTH_THROTTLER]: { limit: 10, ttl: 60_000 } })
|
||||
@Post('create')
|
||||
async create(
|
||||
@Body() dto: CreateApiKeyDto,
|
||||
@AuthUser() user: User,
|
||||
@Req() req: FastifyRequest,
|
||||
) {
|
||||
this.assertEnabled();
|
||||
this.rejectApiKeyPrincipal(req);
|
||||
|
||||
// undefined -> service default (1 year); null -> unlimited; string -> Date.
|
||||
const expiresAt =
|
||||
dto.expiresAt === undefined
|
||||
? undefined
|
||||
: dto.expiresAt === null
|
||||
? null
|
||||
: new Date(dto.expiresAt);
|
||||
|
||||
const { token, key } = await this.apiKeyService.create(
|
||||
user,
|
||||
dto.name,
|
||||
expiresAt,
|
||||
);
|
||||
|
||||
this.auditService.log({
|
||||
event: AuditEvent.API_KEY_CREATED,
|
||||
resourceType: AuditResource.API_KEY,
|
||||
resourceId: key.id,
|
||||
});
|
||||
// The durable trail lives in container logs (AUDIT_SERVICE is a Noop in this
|
||||
// build). No token material — the JWT is only ever returned in the response.
|
||||
this.logger.log(
|
||||
`API key created: id=${key.id} name=${JSON.stringify(
|
||||
key.name,
|
||||
)} actor=${user.id} expiresAt=${
|
||||
key.expiresAt ? new Date(key.expiresAt).toISOString() : 'never'
|
||||
} ip=${this.clientIp(req)}`,
|
||||
);
|
||||
|
||||
// Return the token ONCE (never retrievable again) and the computed expiry so
|
||||
// the caller/UI can surface "expires <date>" (the year-default time-bomb
|
||||
// early-warning).
|
||||
return {
|
||||
token,
|
||||
apiKey: {
|
||||
id: key.id,
|
||||
name: key.name,
|
||||
expiresAt: key.expiresAt,
|
||||
createdAt: key.createdAt,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
@HttpCode(HttpStatus.OK)
|
||||
@Post('list')
|
||||
async list(
|
||||
@AuthUser() user: User,
|
||||
@AuthWorkspace() workspace: Workspace,
|
||||
@Req() req: FastifyRequest,
|
||||
) {
|
||||
this.assertEnabled();
|
||||
this.rejectApiKeyPrincipal(req);
|
||||
|
||||
const ability = this.workspaceAbility.createForUser(user, workspace);
|
||||
// Admin (CASL Manage on API): every key in the workspace, attributed to its
|
||||
// creator (closes "a leaked key of a departed employee is invisible"). A
|
||||
// member sees only their own.
|
||||
const keys = ability.can(
|
||||
WorkspaceCaslAction.Manage,
|
||||
WorkspaceCaslSubject.API,
|
||||
)
|
||||
? await this.apiKeyRepo.findAllInWorkspace(workspace.id)
|
||||
: await this.apiKeyRepo.findByCreator(user.id, workspace.id);
|
||||
|
||||
return keys.map((k) => ({
|
||||
id: k.id,
|
||||
name: k.name,
|
||||
expiresAt: k.expiresAt,
|
||||
lastUsedAt: k.lastUsedAt,
|
||||
createdAt: k.createdAt,
|
||||
creator: k.creator,
|
||||
}));
|
||||
}
|
||||
|
||||
@HttpCode(HttpStatus.OK)
|
||||
@Post('revoke')
|
||||
async revoke(
|
||||
@Body() dto: RevokeApiKeyDto,
|
||||
@AuthUser() user: User,
|
||||
@AuthWorkspace() workspace: Workspace,
|
||||
@Req() req: FastifyRequest,
|
||||
) {
|
||||
this.assertEnabled();
|
||||
this.rejectApiKeyPrincipal(req);
|
||||
|
||||
const key = await this.apiKeyRepo.findById(dto.id, workspace.id);
|
||||
// Uniform 404 for a missing/already-revoked key regardless of who asks — no
|
||||
// existence oracle for keys the caller may not own.
|
||||
if (!key) {
|
||||
throw new NotFoundException();
|
||||
}
|
||||
|
||||
const ability = this.workspaceAbility.createForUser(user, workspace);
|
||||
const canManageAll = ability.can(
|
||||
WorkspaceCaslAction.Manage,
|
||||
WorkspaceCaslSubject.API,
|
||||
);
|
||||
// Admin may revoke any key in the workspace; a member only their own.
|
||||
if (!canManageAll && key.creatorId !== user.id) {
|
||||
throw new ForbiddenException();
|
||||
}
|
||||
|
||||
await this.apiKeyService.revoke(key.id, workspace.id);
|
||||
|
||||
this.auditService.log({
|
||||
event: AuditEvent.API_KEY_DELETED,
|
||||
resourceType: AuditResource.API_KEY,
|
||||
resourceId: key.id,
|
||||
});
|
||||
this.logger.log(
|
||||
`API key revoked: id=${key.id} name=${JSON.stringify(
|
||||
key.name,
|
||||
)} actor=${user.id} ip=${this.clientIp(req)}`,
|
||||
);
|
||||
|
||||
return { success: true };
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
import { Module } from '@nestjs/common';
|
||||
import { ApiKeyService } from './api-key.service';
|
||||
import { ApiKeyController } from './api-key.controller';
|
||||
import { TokenModule } from '../auth/token.module';
|
||||
|
||||
// Core (non-EE) API-key feature: issuance REST endpoints + the shared validator
|
||||
// consumed by jwt.strategy (REST) and McpService (the /mcp Bearer router).
|
||||
// DatabaseModule (global) provides ApiKeyRepo/UserRepo/WorkspaceRepo; CaslModule
|
||||
// (global) provides WorkspaceAbilityFactory; TokenModule provides TokenService
|
||||
// (the no-exp api-key signer). ApiKeyService is exported so AuthModule (for
|
||||
// jwt.strategy) and McpModule (for the /mcp router) can inject it directly,
|
||||
// replacing the absent EE `ee/api-key` dynamic require.
|
||||
@Module({
|
||||
imports: [TokenModule],
|
||||
controllers: [ApiKeyController],
|
||||
providers: [ApiKeyService],
|
||||
exports: [ApiKeyService],
|
||||
})
|
||||
export class ApiKeyModule {}
|
||||
@@ -0,0 +1,299 @@
|
||||
import { UnauthorizedException } from '@nestjs/common';
|
||||
import { ApiKeyService } from './api-key.service';
|
||||
import { JwtType } from '../auth/dto/jwt-payload';
|
||||
|
||||
/**
|
||||
* Security contract for ApiKeyService.validate — the single validator shared by
|
||||
* jwt.strategy (REST) and the /mcp Bearer router.
|
||||
*
|
||||
* Invariants under test:
|
||||
* - Anti-enumeration: every DEFINITE deny (missing/revoked/expired row, disabled
|
||||
* user, workspace mismatch, kill-switch off) throws the SAME bare
|
||||
* UnauthorizedException — an agent cannot distinguish expired from revoked.
|
||||
* - deny-on-decision / 5xx-on-infra: an UNEXPECTED (infra) error PROPAGATES as
|
||||
* itself (→ 5xx), never masked as a 401.
|
||||
* - Expiry is read from the ROW, never a JWT exp claim.
|
||||
* - No validate cache (each call re-reads the row → immediate revocation).
|
||||
*/
|
||||
|
||||
const APP_SECRET = 'secret';
|
||||
|
||||
function makeDeps(over: Partial<Record<string, any>> = {}) {
|
||||
const apiKeyRepo = {
|
||||
findById: jest.fn(),
|
||||
insert: jest.fn(),
|
||||
softDelete: jest.fn(),
|
||||
touchLastUsed: jest.fn().mockResolvedValue(undefined),
|
||||
...(over.apiKeyRepo ?? {}),
|
||||
};
|
||||
const userRepo = {
|
||||
findById: jest.fn(),
|
||||
...(over.userRepo ?? {}),
|
||||
};
|
||||
const workspaceRepo = {
|
||||
findById: jest.fn().mockResolvedValue({ id: 'ws-1' }),
|
||||
...(over.workspaceRepo ?? {}),
|
||||
};
|
||||
const tokenService = {
|
||||
generateApiToken: jest.fn().mockResolvedValue('minted.jwt.token'),
|
||||
...(over.tokenService ?? {}),
|
||||
};
|
||||
const environmentService = {
|
||||
isApiKeysEnabled: jest.fn().mockReturnValue(true),
|
||||
getApiKeysEnabledRaw: jest.fn().mockReturnValue(undefined),
|
||||
getAppSecret: jest.fn().mockReturnValue(APP_SECRET),
|
||||
...(over.environmentService ?? {}),
|
||||
};
|
||||
const service = new (ApiKeyService as unknown as new (...a: any[]) => ApiKeyService)(
|
||||
apiKeyRepo,
|
||||
userRepo,
|
||||
workspaceRepo,
|
||||
tokenService,
|
||||
environmentService,
|
||||
);
|
||||
return { service, apiKeyRepo, userRepo, workspaceRepo, tokenService, environmentService };
|
||||
}
|
||||
|
||||
const payload = (over: Record<string, any> = {}) => ({
|
||||
sub: 'u-1',
|
||||
workspaceId: 'ws-1',
|
||||
apiKeyId: 'key-1',
|
||||
type: JwtType.API_KEY,
|
||||
...over,
|
||||
});
|
||||
|
||||
const activeRow = (over: Record<string, any> = {}) => ({
|
||||
id: 'key-1',
|
||||
name: 'k',
|
||||
creatorId: 'u-1',
|
||||
workspaceId: 'ws-1',
|
||||
expiresAt: null,
|
||||
lastUsedAt: null,
|
||||
deletedAt: null,
|
||||
...over,
|
||||
});
|
||||
|
||||
const activeUser = (over: Record<string, any> = {}) => ({
|
||||
id: 'u-1',
|
||||
workspaceId: 'ws-1',
|
||||
deactivatedAt: null,
|
||||
deletedAt: null,
|
||||
isAgent: false,
|
||||
...over,
|
||||
});
|
||||
|
||||
describe('ApiKeyService.validate', () => {
|
||||
it('returns { user, workspace } for a valid active key', async () => {
|
||||
const { service, apiKeyRepo, userRepo } = makeDeps();
|
||||
apiKeyRepo.findById.mockResolvedValue(activeRow());
|
||||
userRepo.findById.mockResolvedValue(activeUser());
|
||||
|
||||
const res = await service.validate(payload() as any);
|
||||
|
||||
expect(res.user).toMatchObject({ id: 'u-1' });
|
||||
expect(res.workspace).toMatchObject({ id: 'ws-1' });
|
||||
// Loads isAgent so downstream provenance does not silently degrade.
|
||||
expect(userRepo.findById).toHaveBeenCalledWith('u-1', 'ws-1', {
|
||||
includeIsAgent: true,
|
||||
});
|
||||
});
|
||||
|
||||
// --- Anti-enumeration: every deny is the SAME bare 401 -------------------
|
||||
const denyCases: Array<[string, (d: ReturnType<typeof makeDeps>) => void]> = [
|
||||
[
|
||||
'missing/orphaned row',
|
||||
(d) => d.apiKeyRepo.findById.mockResolvedValue(undefined),
|
||||
],
|
||||
[
|
||||
'revoked (soft-deleted → findById returns undefined)',
|
||||
(d) => d.apiKeyRepo.findById.mockResolvedValue(undefined),
|
||||
],
|
||||
[
|
||||
'expired (expires_at in the past)',
|
||||
(d) => {
|
||||
d.apiKeyRepo.findById.mockResolvedValue(
|
||||
activeRow({ expiresAt: new Date(Date.now() - 1000) }),
|
||||
);
|
||||
d.userRepo.findById.mockResolvedValue(activeUser());
|
||||
},
|
||||
],
|
||||
[
|
||||
'disabled user',
|
||||
(d) => {
|
||||
d.apiKeyRepo.findById.mockResolvedValue(activeRow());
|
||||
d.userRepo.findById.mockResolvedValue(
|
||||
activeUser({ deactivatedAt: new Date() }),
|
||||
);
|
||||
},
|
||||
],
|
||||
[
|
||||
'user not found',
|
||||
(d) => {
|
||||
d.apiKeyRepo.findById.mockResolvedValue(activeRow());
|
||||
d.userRepo.findById.mockResolvedValue(undefined);
|
||||
},
|
||||
],
|
||||
[
|
||||
'creator/sub mismatch',
|
||||
(d) => {
|
||||
d.apiKeyRepo.findById.mockResolvedValue(activeRow({ creatorId: 'other' }));
|
||||
d.userRepo.findById.mockResolvedValue(activeUser());
|
||||
},
|
||||
],
|
||||
[
|
||||
'workspace row missing',
|
||||
(d) => {
|
||||
d.apiKeyRepo.findById.mockResolvedValue(activeRow());
|
||||
d.userRepo.findById.mockResolvedValue(activeUser());
|
||||
d.workspaceRepo.findById.mockResolvedValue(undefined);
|
||||
},
|
||||
],
|
||||
[
|
||||
'kill-switch OFF',
|
||||
(d) => d.environmentService.isApiKeysEnabled.mockReturnValue(false),
|
||||
],
|
||||
[
|
||||
'malformed payload (missing apiKeyId)',
|
||||
() => undefined,
|
||||
],
|
||||
];
|
||||
|
||||
it.each(denyCases)(
|
||||
'throws a bare UnauthorizedException with NO message for: %s',
|
||||
async (label, arrange) => {
|
||||
const deps = makeDeps();
|
||||
arrange(deps);
|
||||
const p =
|
||||
label === 'malformed payload (missing apiKeyId)'
|
||||
? (payload({ apiKeyId: undefined }) as any)
|
||||
: (payload() as any);
|
||||
|
||||
const err = await deps.service.validate(p).catch((e) => e);
|
||||
expect(err).toBeInstanceOf(UnauthorizedException);
|
||||
// Anti-enumeration: identical, class-less message for every failure mode.
|
||||
expect(err.message).toBe('Unauthorized');
|
||||
},
|
||||
);
|
||||
|
||||
it('kill-switch OFF denies BEFORE any DB read', async () => {
|
||||
const { service, apiKeyRepo, environmentService } = makeDeps();
|
||||
environmentService.isApiKeysEnabled.mockReturnValue(false);
|
||||
|
||||
await expect(service.validate(payload() as any)).rejects.toBeInstanceOf(
|
||||
UnauthorizedException,
|
||||
);
|
||||
expect(apiKeyRepo.findById).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
// --- Infra error propagates (5xx), NOT masked as 401 ---------------------
|
||||
it('PROPAGATES an infra (DB) error instead of masking it as 401', async () => {
|
||||
const { service, apiKeyRepo } = makeDeps();
|
||||
const boom = new Error('connection terminated');
|
||||
apiKeyRepo.findById.mockRejectedValue(boom);
|
||||
|
||||
await expect(service.validate(payload() as any)).rejects.toBe(boom);
|
||||
});
|
||||
|
||||
// --- No cache: revocation is immediate on the next call ------------------
|
||||
it('re-reads the row on every call (no validate cache → immediate revocation)', async () => {
|
||||
const { service, apiKeyRepo, userRepo } = makeDeps();
|
||||
apiKeyRepo.findById.mockResolvedValue(activeRow());
|
||||
userRepo.findById.mockResolvedValue(activeUser());
|
||||
await service.validate(payload() as any);
|
||||
|
||||
// Now the key is revoked (row invisible) — the very next call denies.
|
||||
apiKeyRepo.findById.mockResolvedValue(undefined);
|
||||
await expect(service.validate(payload() as any)).rejects.toBeInstanceOf(
|
||||
UnauthorizedException,
|
||||
);
|
||||
expect(apiKeyRepo.findById).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
// --- last_used_at is throttled + fire-and-forget -------------------------
|
||||
it('touches last_used_at when stale and skips when fresh', async () => {
|
||||
const { service, apiKeyRepo, userRepo } = makeDeps();
|
||||
userRepo.findById.mockResolvedValue(activeUser());
|
||||
|
||||
// Stale (never used) -> touch.
|
||||
apiKeyRepo.findById.mockResolvedValue(activeRow({ lastUsedAt: null }));
|
||||
await service.validate(payload() as any);
|
||||
expect(apiKeyRepo.touchLastUsed).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Fresh (used 1 minute ago) -> skip.
|
||||
apiKeyRepo.touchLastUsed.mockClear();
|
||||
apiKeyRepo.findById.mockResolvedValue(
|
||||
activeRow({ lastUsedAt: new Date(Date.now() - 60_000) }),
|
||||
);
|
||||
await service.validate(payload() as any);
|
||||
expect(apiKeyRepo.touchLastUsed).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('a failing last_used_at touch never fails the request', async () => {
|
||||
const { service, apiKeyRepo, userRepo } = makeDeps();
|
||||
apiKeyRepo.findById.mockResolvedValue(activeRow({ lastUsedAt: null }));
|
||||
userRepo.findById.mockResolvedValue(activeUser());
|
||||
apiKeyRepo.touchLastUsed.mockRejectedValue(new Error('write failed'));
|
||||
|
||||
await expect(service.validate(payload() as any)).resolves.toMatchObject({
|
||||
user: { id: 'u-1' },
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('ApiKeyService.create (mint-then-insert, no exp)', () => {
|
||||
it('mints the JWT BEFORE inserting the row and returns the token once', async () => {
|
||||
const { service, apiKeyRepo, tokenService } = makeDeps();
|
||||
const order: string[] = [];
|
||||
tokenService.generateApiToken.mockImplementation(async () => {
|
||||
order.push('mint');
|
||||
return 'tok';
|
||||
});
|
||||
apiKeyRepo.insert.mockImplementation(async (row: any) => {
|
||||
order.push('insert');
|
||||
return { ...row, createdAt: new Date() };
|
||||
});
|
||||
|
||||
const user = { id: 'u-1', workspaceId: 'ws-1' } as any;
|
||||
const res = await service.create(user, 'my key');
|
||||
|
||||
expect(order).toEqual(['mint', 'insert']);
|
||||
expect(res.token).toBe('tok');
|
||||
// The id is generated before minting and reused for the row.
|
||||
const mintArg = tokenService.generateApiToken.mock.calls[0][0];
|
||||
const insertArg = apiKeyRepo.insert.mock.calls[0][0];
|
||||
expect(mintArg.apiKeyId).toBe(insertArg.id);
|
||||
});
|
||||
|
||||
it('applies the 1-year default when expiresAt is undefined', async () => {
|
||||
const { service, apiKeyRepo } = makeDeps();
|
||||
apiKeyRepo.insert.mockImplementation(async (row: any) => row);
|
||||
const user = { id: 'u-1', workspaceId: 'ws-1' } as any;
|
||||
|
||||
await service.create(user, 'k');
|
||||
|
||||
const insertArg = apiKeyRepo.insert.mock.calls[0][0];
|
||||
const days =
|
||||
(insertArg.expiresAt.getTime() - Date.now()) / (24 * 60 * 60 * 1000);
|
||||
expect(days).toBeGreaterThan(364);
|
||||
expect(days).toBeLessThan(367);
|
||||
});
|
||||
|
||||
it('honors an explicit null (unlimited)', async () => {
|
||||
const { service, apiKeyRepo } = makeDeps();
|
||||
apiKeyRepo.insert.mockImplementation(async (row: any) => row);
|
||||
const user = { id: 'u-1', workspaceId: 'ws-1' } as any;
|
||||
|
||||
await service.create(user, 'k', null);
|
||||
|
||||
expect(apiKeyRepo.insert.mock.calls[0][0].expiresAt).toBeNull();
|
||||
});
|
||||
|
||||
it('does NOT insert a row when minting fails (mint-then-insert → inert)', async () => {
|
||||
const { service, apiKeyRepo, tokenService } = makeDeps();
|
||||
tokenService.generateApiToken.mockRejectedValue(new Error('mint failed'));
|
||||
const user = { id: 'u-1', workspaceId: 'ws-1' } as any;
|
||||
|
||||
await expect(service.create(user, 'k')).rejects.toThrow('mint failed');
|
||||
expect(apiKeyRepo.insert).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,191 @@
|
||||
import {
|
||||
Injectable,
|
||||
Logger,
|
||||
OnModuleInit,
|
||||
UnauthorizedException,
|
||||
} from '@nestjs/common';
|
||||
import { v7 as uuid7 } from 'uuid';
|
||||
import { ApiKeyRepo } from '@docmost/db/repos/api-key/api-key.repo';
|
||||
import { UserRepo } from '@docmost/db/repos/user/user.repo';
|
||||
import { WorkspaceRepo } from '@docmost/db/repos/workspace/workspace.repo';
|
||||
import { TokenService } from '../auth/services/token.service';
|
||||
import { EnvironmentService } from '../../integrations/environment/environment.service';
|
||||
import { JwtApiKeyPayload } from '../auth/dto/jwt-payload';
|
||||
import { ApiKey, User, Workspace } from '@docmost/db/types/entity.types';
|
||||
import { isUserDisabled } from '../../common/helpers';
|
||||
|
||||
// Default lifetime for a new key when the caller does not specify one: 1 year.
|
||||
// The owner runs a homelab where agents live for years; forcing rotation is
|
||||
// operational pain, so an explicit `null` (unlimited) is also allowed.
|
||||
const DEFAULT_LIFETIME_MS = 365 * 24 * 60 * 60 * 1000;
|
||||
|
||||
// last_used_at is a best-effort forensics stamp, not an access record; we only
|
||||
// refresh it when it is older than this to avoid a write on every request. The
|
||||
// 1h resolution is a deliberate constant (forensics granularity, not accounting).
|
||||
const LAST_USED_THROTTLE_MS = 60 * 60 * 1000;
|
||||
|
||||
/**
|
||||
* Core API-key lifecycle service. Owns minting (create), the single validator
|
||||
* shared by BOTH the REST jwt.strategy path and the /mcp Bearer path (validate),
|
||||
* and revocation (revoke). The `api_keys` ROW is the sole source of truth for a
|
||||
* key's lifetime and revocation — never the JWT, which carries no `exp` claim.
|
||||
*/
|
||||
@Injectable()
|
||||
export class ApiKeyService implements OnModuleInit {
|
||||
private readonly logger = new Logger(ApiKeyService.name);
|
||||
|
||||
constructor(
|
||||
private readonly apiKeyRepo: ApiKeyRepo,
|
||||
private readonly userRepo: UserRepo,
|
||||
private readonly workspaceRepo: WorkspaceRepo,
|
||||
private readonly tokenService: TokenService,
|
||||
private readonly environmentService: EnvironmentService,
|
||||
) {}
|
||||
|
||||
onModuleInit() {
|
||||
// Boot log so the kill-switch state after each deploy is verifiable in logs.
|
||||
const enabled = this.environmentService.isApiKeysEnabled();
|
||||
const raw = this.environmentService.getApiKeysEnabledRaw();
|
||||
this.logger.log(
|
||||
`API keys: ${enabled ? 'ENABLED' : 'DISABLED'} (API_KEYS_ENABLED=${
|
||||
raw ?? 'unset'
|
||||
})`,
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Mint a new key for `user`. mint-then-insert ordering (R1):
|
||||
* 1. generate the id first — it must be in the JWT payload before the row.
|
||||
* 2. mint the JWT (no `exp` claim). A mint failure aborts before any row is
|
||||
* written (inert), so a half-created key cannot exist.
|
||||
* 3. insert the row last. A lost response leaves an orphaned row that is
|
||||
* visible in `list` and self-heals (the user revokes it).
|
||||
* The token is returned ONCE and never stored — the JWT is self-contained, so
|
||||
* no token material lives in the table.
|
||||
*
|
||||
* `expiresAt`: `undefined` -> default 1 year; `null` -> unlimited (explicit);
|
||||
* a Date -> that instant (a past date is rejected at the DTO layer).
|
||||
*/
|
||||
async create(
|
||||
user: User,
|
||||
name: string,
|
||||
expiresAt?: Date | null,
|
||||
): Promise<{ token: string; key: ApiKey }> {
|
||||
const resolvedExpiresAt =
|
||||
expiresAt === undefined
|
||||
? new Date(Date.now() + DEFAULT_LIFETIME_MS)
|
||||
: expiresAt;
|
||||
|
||||
const apiKeyId = uuid7();
|
||||
|
||||
const token = await this.tokenService.generateApiToken({
|
||||
apiKeyId,
|
||||
user,
|
||||
workspaceId: user.workspaceId,
|
||||
});
|
||||
|
||||
const key = await this.apiKeyRepo.insert({
|
||||
id: apiKeyId,
|
||||
name,
|
||||
creatorId: user.id,
|
||||
workspaceId: user.workspaceId,
|
||||
expiresAt: resolvedExpiresAt,
|
||||
});
|
||||
|
||||
return { token, key };
|
||||
}
|
||||
|
||||
/**
|
||||
* The single validator for an api-key principal, shared by jwt.strategy and the
|
||||
* /mcp Bearer router. Returns `{ user, workspace }` (the same shape the access
|
||||
* path returns) so the AuthUser/AuthWorkspace decorators and MCP identity work
|
||||
* unchanged.
|
||||
*
|
||||
* Failure semantics (R4, anti-enumeration): a DEFINITE negative fact — feature
|
||||
* disabled, missing/revoked/expired row, workspace mismatch, disabled user —
|
||||
* throws a bare `UnauthorizedException` (a single generic 401 for every case;
|
||||
* an agent cannot distinguish expired from revoked, and its reaction is
|
||||
* identical). An UNEXPECTED (infra) error is NOT caught here: it propagates so
|
||||
* the surface returns 5xx, never a masked 401 (deny-on-decision / 5xx-on-infra).
|
||||
* There is NO validate cache: it is 2–3 PK lookups (~1ms), two orders of
|
||||
* magnitude cheaper than the bcrypt it replaces; a cache would only add a
|
||||
* Date-serialization trap and a revocation lag. Revocation is immediate.
|
||||
*/
|
||||
async validate(
|
||||
payload: JwtApiKeyPayload,
|
||||
): Promise<{ user: User; workspace: Workspace }> {
|
||||
// Kill-switch OFF: deny unconditionally (same generic 401). Note the
|
||||
// endpoints additionally 404 at the controller; here we deny the token.
|
||||
if (!this.environmentService.isApiKeysEnabled()) {
|
||||
throw new UnauthorizedException();
|
||||
}
|
||||
|
||||
if (!payload?.apiKeyId || !payload?.sub || !payload?.workspaceId) {
|
||||
throw new UnauthorizedException();
|
||||
}
|
||||
|
||||
const row = await this.apiKeyRepo.findById(
|
||||
payload.apiKeyId,
|
||||
payload.workspaceId,
|
||||
);
|
||||
// Absent row = revoked (soft-deleted, invisible to findById), orphaned
|
||||
// (creator/workspace cascade-deleted), or never existed. All terminal deny.
|
||||
if (!row) {
|
||||
throw new UnauthorizedException();
|
||||
}
|
||||
|
||||
// Expiry is read from the ROW, never an `exp` JWT claim.
|
||||
if (row.expiresAt && row.expiresAt.getTime() <= Date.now()) {
|
||||
throw new UnauthorizedException();
|
||||
}
|
||||
|
||||
const user = await this.userRepo.findById(
|
||||
payload.sub,
|
||||
payload.workspaceId,
|
||||
{ includeIsAgent: true },
|
||||
);
|
||||
if (!user || isUserDisabled(user)) {
|
||||
throw new UnauthorizedException();
|
||||
}
|
||||
|
||||
// The key acts only as its creator (defence in depth against a token whose
|
||||
// signed `sub` ever drifted from the row's owner).
|
||||
if (row.creatorId !== user.id) {
|
||||
throw new UnauthorizedException();
|
||||
}
|
||||
|
||||
const workspace = await this.workspaceRepo.findById(payload.workspaceId);
|
||||
if (!workspace) {
|
||||
throw new UnauthorizedException();
|
||||
}
|
||||
|
||||
// Best-effort, throttled, fire-and-forget forensics stamp AFTER all checks.
|
||||
this.touchLastUsed(row);
|
||||
|
||||
return { user, workspace };
|
||||
}
|
||||
|
||||
/**
|
||||
* Revoke (soft-delete) a key. Authorization/ownership is decided by the caller
|
||||
* (the controller, via CASL); this only performs the terminal write. Idempotent:
|
||||
* a second revoke is a no-op (the row is already invisible).
|
||||
*/
|
||||
async revoke(id: string, workspaceId: string): Promise<void> {
|
||||
await this.apiKeyRepo.softDelete(id, workspaceId);
|
||||
}
|
||||
|
||||
// Throttled best-effort last_used_at bump: skip if it was touched within the
|
||||
// window; otherwise fire-and-forget so a stamp write never fails or slows the
|
||||
// request (mirrors SessionActivityService.trackActivity).
|
||||
private touchLastUsed(row: ApiKey): void {
|
||||
const last = row.lastUsedAt ? new Date(row.lastUsedAt).getTime() : 0;
|
||||
if (Date.now() - last < LAST_USED_THROTTLE_MS) return;
|
||||
void this.apiKeyRepo.touchLastUsed(row.id).catch((err) => {
|
||||
this.logger.warn(
|
||||
`Failed to update api_key last_used_at for ${row.id}: ${
|
||||
(err as Error)?.message ?? err
|
||||
}`,
|
||||
);
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
import {
|
||||
IsDateString,
|
||||
IsNotEmpty,
|
||||
IsOptional,
|
||||
IsString,
|
||||
MaxLength,
|
||||
registerDecorator,
|
||||
ValidationOptions,
|
||||
} from 'class-validator';
|
||||
|
||||
/**
|
||||
* `expiresAt` must be strictly in the future. Skips validation for `null`
|
||||
* (explicit "unlimited") and `undefined` (server applies the 1-year default),
|
||||
* so only an actually-supplied date is range-checked. Rejecting a PAST date at
|
||||
* the DTO layer means a caller cannot mint an already-dead key.
|
||||
*/
|
||||
function IsFutureDateString(options?: ValidationOptions) {
|
||||
return function (object: object, propertyName: string) {
|
||||
registerDecorator({
|
||||
name: 'isFutureDateString',
|
||||
target: object.constructor,
|
||||
propertyName,
|
||||
options: {
|
||||
message: 'expiresAt must be a date in the future',
|
||||
...options,
|
||||
},
|
||||
validator: {
|
||||
validate(value: unknown) {
|
||||
if (value === null || value === undefined) return true;
|
||||
if (typeof value !== 'string') return false;
|
||||
const t = Date.parse(value);
|
||||
return !Number.isNaN(t) && t > Date.now();
|
||||
},
|
||||
},
|
||||
});
|
||||
};
|
||||
}
|
||||
|
||||
export class CreateApiKeyDto {
|
||||
@IsString()
|
||||
@IsNotEmpty()
|
||||
@MaxLength(255)
|
||||
name: string;
|
||||
|
||||
// undefined -> default 1 year (applied server-side); null -> unlimited
|
||||
// (explicit); an ISO date string -> that instant, which must be in the future.
|
||||
@IsOptional()
|
||||
@IsDateString()
|
||||
@IsFutureDateString()
|
||||
expiresAt?: string | null;
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
import { IsUUID } from 'class-validator';
|
||||
|
||||
export class RevokeApiKeyDto {
|
||||
@IsUUID()
|
||||
id: string;
|
||||
}
|
||||
@@ -20,3 +20,71 @@ describe('AuthController', () => {
|
||||
expect(controller).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
// The collab-token handler is the ARM-SEAM of the #501 anti-laundering defense:
|
||||
// it derives the api-key origin args ({ apiKeyId }) from the SIGNED-derived
|
||||
// `req.raw` fields (stamped by jwt.strategy) and threads them into the collab
|
||||
// token mint, forcing principal='api_key' so a later key revoke rejects NEW
|
||||
// collab connections. A regression here (reading `body`, dropping `apiKeyId`,
|
||||
// or inverting the ternary) would silently mint api-key requests as
|
||||
// principal='session' and reopen the laundering hole with a green suite. These
|
||||
// tests pin the EXACT args the controller passes to authService.getCollabToken.
|
||||
describe('AuthController.collabToken arms the api-key origin (#501)', () => {
|
||||
let controller: AuthController;
|
||||
let getCollabToken: jest.Mock;
|
||||
|
||||
const user = { id: 'u-1', workspaceId: 'ws-1' } as any;
|
||||
const workspace = { id: 'ws-1' } as any;
|
||||
|
||||
beforeEach(() => {
|
||||
getCollabToken = jest.fn().mockResolvedValue({ token: 'ct' });
|
||||
controller = new AuthController(
|
||||
{ getCollabToken } as any, // authService
|
||||
{} as any, // sessionService
|
||||
{} as any, // environmentService
|
||||
{} as any, // moduleRef
|
||||
{} as any, // auditService
|
||||
);
|
||||
});
|
||||
|
||||
it('threads { apiKeyId } when req.raw marks an api_key principal (ARMED)', async () => {
|
||||
const req = { raw: { authType: 'api_key', apiKeyId: 'k-1' } } as any;
|
||||
|
||||
await controller.collabToken(user, workspace, req);
|
||||
|
||||
expect(getCollabToken).toHaveBeenCalledTimes(1);
|
||||
expect(getCollabToken).toHaveBeenCalledWith(user, 'ws-1', {
|
||||
apiKeyId: 'k-1',
|
||||
});
|
||||
});
|
||||
|
||||
it('passes undefined for a normal session request (NOT armed)', async () => {
|
||||
const req = { raw: {} } as any;
|
||||
|
||||
await controller.collabToken(user, workspace, req);
|
||||
|
||||
expect(getCollabToken).toHaveBeenCalledTimes(1);
|
||||
expect(getCollabToken).toHaveBeenCalledWith(user, 'ws-1', undefined);
|
||||
});
|
||||
|
||||
it('does NOT arm when api_key authType lacks an apiKeyId', async () => {
|
||||
const req = { raw: { authType: 'api_key' } } as any;
|
||||
|
||||
await controller.collabToken(user, workspace, req);
|
||||
|
||||
expect(getCollabToken).toHaveBeenCalledWith(user, 'ws-1', undefined);
|
||||
});
|
||||
|
||||
it('reads the SIGNED req.raw, never a spoofable client body', async () => {
|
||||
// A client-supplied body claiming api_key must be ignored: only the
|
||||
// jwt.strategy-stamped req.raw can arm the api-key origin.
|
||||
const req = {
|
||||
raw: {},
|
||||
body: { authType: 'api_key', apiKeyId: 'attacker' },
|
||||
} as any;
|
||||
|
||||
await controller.collabToken(user, workspace, req);
|
||||
|
||||
expect(getCollabToken).toHaveBeenCalledWith(user, 'ws-1', undefined);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -207,8 +207,20 @@ export class AuthController {
|
||||
async collabToken(
|
||||
@AuthUser() user: User,
|
||||
@AuthWorkspace() workspace: Workspace,
|
||||
@Req() req: FastifyRequest,
|
||||
) {
|
||||
return this.authService.getCollabToken(user, workspace.id);
|
||||
// Thread the api-key origin (#501): when the requester authenticated with an
|
||||
// api key (jwt.strategy stamped req.raw.authType/apiKeyId), the minted collab
|
||||
// token carries principal='api_key' + apiKeyId so a later revoke of the key
|
||||
// rejects NEW collab connections. A normal session request mints a
|
||||
// principal='session' token. Reading the SIGNED-derived req.raw fields (never
|
||||
// a client body) keeps it unspoofable.
|
||||
const raw = req.raw as { authType?: string; apiKeyId?: string };
|
||||
const apiKey =
|
||||
raw.authType === 'api_key' && raw.apiKeyId
|
||||
? { apiKeyId: raw.apiKeyId }
|
||||
: undefined;
|
||||
return this.authService.getCollabToken(user, workspace.id, apiKey);
|
||||
}
|
||||
|
||||
@SkipThrottle({ [AUTH_THROTTLER]: true })
|
||||
|
||||
@@ -5,9 +5,13 @@ import { JwtStrategy } from './strategies/jwt.strategy';
|
||||
import { WorkspaceModule } from '../workspace/workspace.module';
|
||||
import { SignupService } from './services/signup.service';
|
||||
import { TokenModule } from './token.module';
|
||||
import { ApiKeyModule } from '../api-key/api-key.module';
|
||||
|
||||
@Module({
|
||||
imports: [TokenModule, WorkspaceModule],
|
||||
// ApiKeyModule supplies ApiKeyService, injected into JwtStrategy so an
|
||||
// api_key Bearer/cookie token is validated directly (replacing the absent EE
|
||||
// `ee/api-key` dynamic require).
|
||||
imports: [TokenModule, WorkspaceModule, ApiKeyModule],
|
||||
controllers: [AuthController],
|
||||
providers: [AuthService, SignupService, JwtStrategy],
|
||||
exports: [SignupService, AuthService],
|
||||
|
||||
@@ -33,6 +33,17 @@ export type JwtPayload = {
|
||||
aiChatId?: string | null;
|
||||
};
|
||||
|
||||
// The AUTH-PRINCIPAL kind behind a collab token — distinct from `actor`
|
||||
// (provenance). Stamped into EVERY newly-minted collab token (#501): 'session'
|
||||
// for a normal user/session (incl. the internal AI agent, which is session-
|
||||
// backed), 'api_key' when the token was minted by an api-key principal (an
|
||||
// external MCP agent). The discriminator keys on the token's ORIGIN, so the
|
||||
// collab seam can re-check a revoked api key on connect and reject a claimless
|
||||
// token after the rollout grace window. NOT keyed on `actor:'agent'` — the
|
||||
// internal agent is 'agent' but session-backed, so it must stay on the no-check
|
||||
// path.
|
||||
export type CollabPrincipal = 'session' | 'api_key';
|
||||
|
||||
export type JwtCollabPayload = {
|
||||
sub: string;
|
||||
workspaceId: string;
|
||||
@@ -44,6 +55,13 @@ export type JwtCollabPayload = {
|
||||
// Nullable: an external MCP agent has no internal ai_chats row, so it carries
|
||||
// an 'agent' actor with a null aiChatId.
|
||||
aiChatId?: string | null;
|
||||
// Auth-principal discriminator (#501). Present on every post-rollout token;
|
||||
// its absence on a still-valid token past the grace window is treated as an
|
||||
// error, not trust (fail-closed).
|
||||
principal?: CollabPrincipal;
|
||||
// Only when principal === 'api_key': the minting key's id, so the collab seam
|
||||
// can row-check (and reject) a revoked key on connect.
|
||||
apiKeyId?: string;
|
||||
};
|
||||
|
||||
export type JwtExchangePayload = {
|
||||
|
||||
@@ -375,10 +375,20 @@ export class AuthService {
|
||||
}
|
||||
}
|
||||
|
||||
async getCollabToken(user: User, workspaceId: string) {
|
||||
async getCollabToken(
|
||||
user: User,
|
||||
workspaceId: string,
|
||||
// Origin of the request minting this collab token (#501). When the caller is
|
||||
// an api-key principal, its apiKeyId is threaded into the token so the collab
|
||||
// seam can re-check the key on connect (closing api-key -> long-lived-collab
|
||||
// laundering). Absent for a normal session/human request.
|
||||
apiKey?: { apiKeyId: string },
|
||||
) {
|
||||
const token = await this.tokenService.generateCollabToken(
|
||||
user,
|
||||
workspaceId,
|
||||
undefined,
|
||||
apiKey,
|
||||
);
|
||||
return { token };
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import { ForbiddenException, UnauthorizedException } from '@nestjs/common';
|
||||
import * as jwt from 'jsonwebtoken';
|
||||
import { TokenService } from './token.service';
|
||||
import { JwtType } from '../dto/jwt-payload';
|
||||
|
||||
@@ -213,4 +214,175 @@ describe('TokenService.generateCollabToken', () => {
|
||||
aiChatId: 'chat-456',
|
||||
});
|
||||
});
|
||||
|
||||
// #501 fail-closed discriminator: EVERY collab token carries a principal.
|
||||
it("defaults principal to 'session' with NO apiKeyId (normal/internal-agent path)", async () => {
|
||||
const { service, jwtService } = makeTokenService();
|
||||
await service.generateCollabToken(makeUser() as never, 'ws-1');
|
||||
const [payload] = jwtService.sign.mock.calls[0];
|
||||
expect(payload.principal).toBe('session');
|
||||
expect(payload).not.toHaveProperty('apiKeyId');
|
||||
});
|
||||
|
||||
it("the internal agent (provenance, NO apiKey) still gets principal='session'", async () => {
|
||||
const { service, jwtService } = makeTokenService();
|
||||
await service.generateCollabToken(
|
||||
makeUser() as never,
|
||||
'ws-1',
|
||||
{ actor: 'agent', aiChatId: 'chat-1' },
|
||||
);
|
||||
const [payload] = jwtService.sign.mock.calls[0];
|
||||
// Keyed on api-key ORIGIN, not actor: an is_agent session token is 'session'.
|
||||
expect(payload.principal).toBe('session');
|
||||
expect(payload).not.toHaveProperty('apiKeyId');
|
||||
});
|
||||
|
||||
it("stamps principal='api_key' + apiKeyId when minted by an api-key principal", async () => {
|
||||
const { service, jwtService } = makeTokenService();
|
||||
await service.generateCollabToken(
|
||||
makeUser() as never,
|
||||
'ws-1',
|
||||
undefined,
|
||||
{ apiKeyId: 'key-9' },
|
||||
);
|
||||
const [payload] = jwtService.sign.mock.calls[0];
|
||||
expect(payload.principal).toBe('api_key');
|
||||
expect(payload.apiKeyId).toBe('key-9');
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* API-key token minting MUST carry NO `exp` claim — the ONLY source of truth for
|
||||
* a key's lifetime/revocation is its `api_keys` row, checked on every request.
|
||||
*
|
||||
* This is a LIVE-bug regression: the shared JwtService is registered with a
|
||||
* global `signOptions.expiresIn` (default '90d') that merges into every sign(),
|
||||
* so an api-key minted through it silently gets exp=now+90d and an "unlimited"
|
||||
* key dies in 90 days. TokenService.generateApiToken mints through a dedicated
|
||||
* no-expiry signer instead. These tests construct the REAL signer (a real secret
|
||||
* via the stubbed EnvironmentService) and decode the produced JWT to assert the
|
||||
* observable property: no `exp`.
|
||||
*/
|
||||
describe('TokenService.generateApiToken (no exp claim ever)', () => {
|
||||
const APP_SECRET_LOCAL = 'apikey-secret';
|
||||
|
||||
function makeRealSignerService() {
|
||||
// Give the SHARED jwtService a global expiresIn so a regression (minting
|
||||
// through it) would show up as an exp claim — the exact live bug.
|
||||
const { JwtService } = require('@nestjs/jwt');
|
||||
const sharedJwt = new JwtService({
|
||||
secret: APP_SECRET_LOCAL,
|
||||
signOptions: { expiresIn: '90d', issuer: 'Docmost' },
|
||||
});
|
||||
const environmentService = {
|
||||
getAppSecret: () => APP_SECRET_LOCAL,
|
||||
};
|
||||
const service = new (TokenService as unknown as new (
|
||||
...args: unknown[]
|
||||
) => TokenService)(sharedJwt, environmentService);
|
||||
return { service };
|
||||
}
|
||||
|
||||
const user = makeUser({ id: 'svc-1', workspaceId: 'ws-1' });
|
||||
|
||||
it('mints an api-key JWT with NO exp claim and issuer Docmost', async () => {
|
||||
const { service } = makeRealSignerService();
|
||||
|
||||
const token = await service.generateApiToken({
|
||||
apiKeyId: 'key-1',
|
||||
user: user as never,
|
||||
workspaceId: 'ws-1',
|
||||
});
|
||||
|
||||
const decoded = jwt.decode(token) as Record<string, unknown>;
|
||||
// The observable security property: no expiry lives in the JWT.
|
||||
expect(decoded.exp).toBeUndefined();
|
||||
expect(decoded).toMatchObject({
|
||||
sub: 'svc-1',
|
||||
apiKeyId: 'key-1',
|
||||
workspaceId: 'ws-1',
|
||||
type: JwtType.API_KEY,
|
||||
iss: 'Docmost',
|
||||
});
|
||||
});
|
||||
|
||||
it('demonstrates the live bug it guards: the SHARED signer WOULD add exp', () => {
|
||||
const { JwtService } = require('@nestjs/jwt');
|
||||
const sharedJwt = new JwtService({
|
||||
secret: APP_SECRET_LOCAL,
|
||||
signOptions: { expiresIn: '90d', issuer: 'Docmost' },
|
||||
});
|
||||
// Even with empty per-call options the global expiresIn merges in.
|
||||
const leaky = sharedJwt.sign({ sub: 'x', type: JwtType.API_KEY }, {});
|
||||
expect((jwt.decode(leaky) as Record<string, unknown>).exp).toBeDefined();
|
||||
});
|
||||
|
||||
it('refuses to mint for a disabled user', async () => {
|
||||
const { service } = makeRealSignerService();
|
||||
await expect(
|
||||
service.generateApiToken({
|
||||
apiKeyId: 'key-1',
|
||||
user: makeUser({ deactivatedAt: new Date() }) as never,
|
||||
workspaceId: 'ws-1',
|
||||
}),
|
||||
).rejects.toBeInstanceOf(ForbiddenException);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* verifyJwtOneOf is the type-routing primitive: verify the signature once and
|
||||
* assert the token type is on an explicit allowlist. It must NOT degrade into a
|
||||
* "return whatever type" helper — a token whose type is off the allowlist is
|
||||
* rejected with the same generic error as a single-type mismatch.
|
||||
*/
|
||||
describe('TokenService.verifyJwtOneOf (allowlist type-routing)', () => {
|
||||
it('returns the payload when the type is on the allowlist', async () => {
|
||||
const verifyAsync = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ type: JwtType.API_KEY, sub: 'u-1' });
|
||||
const { service } = makeTokenService({ verifyAsync });
|
||||
|
||||
const payload = await service.verifyJwtOneOf(token123(), [
|
||||
JwtType.ACCESS,
|
||||
JwtType.API_KEY,
|
||||
]);
|
||||
|
||||
expect(payload).toMatchObject({ type: JwtType.API_KEY, sub: 'u-1' });
|
||||
expect(verifyAsync).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('accepts the OTHER allowed type too', async () => {
|
||||
const verifyAsync = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ type: JwtType.ACCESS, sub: 'u-1' });
|
||||
const { service } = makeTokenService({ verifyAsync });
|
||||
|
||||
await expect(
|
||||
service.verifyJwtOneOf(token123(), [JwtType.ACCESS, JwtType.API_KEY]),
|
||||
).resolves.toMatchObject({ type: JwtType.ACCESS });
|
||||
});
|
||||
|
||||
it('rejects a token whose type is OFF the allowlist (confused-deputy guard)', async () => {
|
||||
const verifyAsync = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ type: JwtType.COLLAB, sub: 'u-1' });
|
||||
const { service } = makeTokenService({ verifyAsync });
|
||||
|
||||
await expect(
|
||||
service.verifyJwtOneOf(token123(), [JwtType.ACCESS, JwtType.API_KEY]),
|
||||
).rejects.toBeInstanceOf(UnauthorizedException);
|
||||
});
|
||||
|
||||
it('verifies the signature exactly ONCE', async () => {
|
||||
const verifyAsync = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ type: JwtType.ACCESS });
|
||||
const { service } = makeTokenService({ verifyAsync });
|
||||
await service.verifyJwtOneOf(token123(), [JwtType.ACCESS]);
|
||||
expect(verifyAsync).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
|
||||
function token123(): string {
|
||||
return 'a.b.c';
|
||||
}
|
||||
|
||||
@@ -4,7 +4,6 @@ import {
|
||||
UnauthorizedException,
|
||||
} from '@nestjs/common';
|
||||
import { JwtService } from '@nestjs/jwt';
|
||||
import type { StringValue } from 'ms';
|
||||
import { EnvironmentService } from '../../../integrations/environment/environment.service';
|
||||
import {
|
||||
JwtApiKeyPayload,
|
||||
@@ -62,6 +61,12 @@ export class TokenService {
|
||||
// token carries no actor/aiChatId and is treated as 'user' downstream.
|
||||
// aiChatId is nullable for an external agent with no internal ai_chats row.
|
||||
provenance?: { actor: 'agent'; aiChatId: string | null },
|
||||
// Optional api-key origin (#501). When the collab token is minted by an
|
||||
// api-key principal (an external MCP agent), the caller passes the key id so
|
||||
// the token carries principal='api_key' + apiKeyId and the collab seam can
|
||||
// re-check the key on connect. Absent -> principal='session' (a normal
|
||||
// user/session, including the internal session-backed AI agent).
|
||||
apiKey?: { apiKeyId: string },
|
||||
): Promise<string> {
|
||||
if (isUserDisabled(user)) {
|
||||
throw new ForbiddenException();
|
||||
@@ -71,6 +76,10 @@ export class TokenService {
|
||||
sub: user.id,
|
||||
workspaceId,
|
||||
type: JwtType.COLLAB,
|
||||
// Fail-closed discriminator on EVERY minted token: 'api_key' when minted by
|
||||
// an api-key principal, else 'session'.
|
||||
principal: apiKey ? 'api_key' : 'session',
|
||||
...(apiKey ? { apiKeyId: apiKey.apiKeyId } : {}),
|
||||
...(provenance
|
||||
? { actor: provenance.actor, aiChatId: provenance.aiChatId }
|
||||
: {}),
|
||||
@@ -123,9 +132,8 @@ export class TokenService {
|
||||
apiKeyId: string;
|
||||
user: User;
|
||||
workspaceId: string;
|
||||
expiresIn?: StringValue | number;
|
||||
}): Promise<string> {
|
||||
const { apiKeyId, user, workspaceId, expiresIn } = opts;
|
||||
const { apiKeyId, user, workspaceId } = opts;
|
||||
if (isUserDisabled(user)) {
|
||||
throw new ForbiddenException();
|
||||
}
|
||||
@@ -137,7 +145,32 @@ export class TokenService {
|
||||
type: JwtType.API_KEY,
|
||||
};
|
||||
|
||||
return this.jwtService.sign(payload, expiresIn ? { expiresIn } : {});
|
||||
// API-key tokens carry NO `exp` claim EVER — the ONLY source of truth for a
|
||||
// key's lifetime and revocation is its `api_keys` row (checked on every
|
||||
// request), not the JWT. This CANNOT use `this.jwtService`: TokenModule
|
||||
// registers it with a global `signOptions.expiresIn` (JWT_TOKEN_EXPIRES_IN,
|
||||
// default '90d'), which merges into EVERY sign() call — even `sign(payload,
|
||||
// {})` — and `{ expiresIn: undefined }` THROWS rather than stripping it
|
||||
// (verified empirically). So an "unlimited" key minted through the shared
|
||||
// signer would silently get exp=now+90d and die in 90 days regardless of its
|
||||
// row. We mint through a dedicated no-expiry signer, re-stamping only
|
||||
// `issuer: 'Docmost'` for claim parity with the shared signer.
|
||||
return this.apiKeyJwtService().sign(payload);
|
||||
}
|
||||
|
||||
// Lazily-built JWT signer for API-key tokens: same APP_SECRET, same 'Docmost'
|
||||
// issuer, but WITHOUT the global `expiresIn` — so minted API-key tokens have no
|
||||
// `exp` claim. Built once and cached. Verification still goes through the
|
||||
// shared verifier (same secret); `verifyAsync` does not require an `exp`.
|
||||
private _apiKeyJwtService?: JwtService;
|
||||
private apiKeyJwtService(): JwtService {
|
||||
if (!this._apiKeyJwtService) {
|
||||
this._apiKeyJwtService = new JwtService({
|
||||
secret: this.environmentService.getAppSecret(),
|
||||
signOptions: { issuer: 'Docmost' },
|
||||
});
|
||||
}
|
||||
return this._apiKeyJwtService;
|
||||
}
|
||||
|
||||
async generatePdfRenderToken(
|
||||
@@ -177,4 +210,31 @@ export class TokenService {
|
||||
|
||||
return payload;
|
||||
}
|
||||
|
||||
/**
|
||||
* Verify a token's signature ONCE and assert its `type` is one of `allowed`.
|
||||
*
|
||||
* This is the type-routing primitive for surfaces that legitimately accept
|
||||
* more than one token type on the same Bearer slot (the /mcp Bearer path
|
||||
* accepts both an ACCESS and an API_KEY token). It is deliberately NOT a
|
||||
* "verify-and-return-whatever-type" helper — that would be a reusable
|
||||
* confused-deputy footgun (any caller could then feed an attachment/collab
|
||||
* token where an access token is expected). An explicit allowlist preserves
|
||||
* the type-pinning property of `verifyJwt`: a token whose `type` is not in the
|
||||
* allowlist is rejected with the SAME generic error as a type mismatch, and
|
||||
* the signature is verified exactly once (no double-verify).
|
||||
*/
|
||||
async verifyJwtOneOf(token: string, allowed: JwtType[]) {
|
||||
const payload = await this.jwtService.verifyAsync(token, {
|
||||
secret: this.environmentService.getAppSecret(),
|
||||
});
|
||||
|
||||
if (!allowed.includes(payload.type)) {
|
||||
throw new UnauthorizedException(
|
||||
'Invalid JWT token. Token type does not match.',
|
||||
);
|
||||
}
|
||||
|
||||
return payload;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -22,7 +22,8 @@ describe('JwtStrategy — provenance derivation', () => {
|
||||
const userSessionRepo: any = { findActiveById: jest.fn() };
|
||||
const sessionActivityService: any = { trackActivity: jest.fn() };
|
||||
const environmentService: any = { getAppSecret: () => 'test-secret' };
|
||||
const moduleRef: any = {};
|
||||
// ACCESS-path tests never touch the api-key seam; a bare stub suffices.
|
||||
const apiKeyService: any = { validate: jest.fn() };
|
||||
|
||||
const strategy = new JwtStrategy(
|
||||
userRepo,
|
||||
@@ -30,7 +31,7 @@ describe('JwtStrategy — provenance derivation', () => {
|
||||
userSessionRepo,
|
||||
sessionActivityService,
|
||||
environmentService,
|
||||
moduleRef,
|
||||
apiKeyService,
|
||||
);
|
||||
return { strategy, userRepo };
|
||||
}
|
||||
@@ -122,25 +123,29 @@ describe('JwtStrategy — provenance derivation', () => {
|
||||
});
|
||||
|
||||
/**
|
||||
* Provenance derivation on the API-KEY path (jwt.strategy.validateApiKey, #486).
|
||||
* Provenance derivation on the API-KEY path (jwt.strategy.validateApiKey, #486
|
||||
* + #501).
|
||||
*
|
||||
* 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',
|
||||
* SERVER-SIDE user returned by ApiKeyService.validate: 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.
|
||||
* #501 wires the CORE ApiKeyService (the EE `ee/api-key` module is absent in the
|
||||
* fork) directly into the strategy — no dynamic require. The strategy also stamps
|
||||
* `req.raw.authType='api_key'` + `req.raw.apiKeyId` for the "a token cannot manage
|
||||
* tokens" guard on the /api-keys surface.
|
||||
*/
|
||||
describe('JwtStrategy — API-key provenance derivation (#486)', () => {
|
||||
function makeApiKeyStrategy(validateApiKeyImpl: (p: any) => Promise<any>) {
|
||||
describe('JwtStrategy — API-key provenance derivation (#486/#501)', () => {
|
||||
function makeApiKeyStrategy(validateImpl: (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 validate = jest.fn(validateImpl);
|
||||
const apiKeyService: any = { validate };
|
||||
|
||||
const strategy = new JwtStrategy(
|
||||
userRepo,
|
||||
@@ -148,14 +153,9 @@ describe('JwtStrategy — API-key provenance derivation (#486)', () => {
|
||||
userSessionRepo,
|
||||
sessionActivityService,
|
||||
environmentService,
|
||||
moduleRef,
|
||||
apiKeyService,
|
||||
);
|
||||
// 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 };
|
||||
return { strategy, validate };
|
||||
}
|
||||
|
||||
const makeReq = () => ({ raw: {} as Record<string, any> });
|
||||
@@ -166,22 +166,23 @@ describe('JwtStrategy — API-key provenance derivation (#486)', () => {
|
||||
type: JwtType.API_KEY,
|
||||
});
|
||||
|
||||
it("stamps actor='agent' for an is_agent API key (from the validated user)", async () => {
|
||||
it("stamps actor='agent' + authType/apiKeyId for an is_agent API key", async () => {
|
||||
const validated = {
|
||||
user: { id: 'svc-1', isAgent: true },
|
||||
workspace: { id: 'ws-1' },
|
||||
};
|
||||
const { strategy, validateApiKey } = makeApiKeyStrategy(
|
||||
async () => validated,
|
||||
);
|
||||
const { strategy, validate } = makeApiKeyStrategy(async () => validated);
|
||||
const req = makeReq();
|
||||
|
||||
const result = await strategy.validate(req, apiKeyPayload() as any);
|
||||
|
||||
expect(validateApiKey).toHaveBeenCalledTimes(1);
|
||||
expect(validate).toHaveBeenCalledTimes(1);
|
||||
expect(req.raw.actor).toBe('agent');
|
||||
// API keys carry no internal ai_chats row -> null.
|
||||
expect(req.raw.aiChatId).toBeNull();
|
||||
// Principal-kind markers for the management-surface guard.
|
||||
expect(req.raw.authType).toBe('api_key');
|
||||
expect(req.raw.apiKeyId).toBe('key-1');
|
||||
// The validated auth object is returned unchanged (req.user shape preserved).
|
||||
expect(result).toBe(validated);
|
||||
});
|
||||
@@ -197,25 +198,19 @@ describe('JwtStrategy — API-key provenance derivation (#486)', () => {
|
||||
|
||||
expect(req.raw.actor).toBe('user');
|
||||
expect(req.raw.aiChatId).toBeNull();
|
||||
expect(req.raw.authType).toBe('api_key');
|
||||
});
|
||||
|
||||
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);
|
||||
it('propagates a validate() rejection and stamps nothing', async () => {
|
||||
const { strategy } = makeApiKeyStrategy(async () => {
|
||||
throw new UnauthorizedException();
|
||||
});
|
||||
const req = makeReq();
|
||||
|
||||
await expect(
|
||||
strategy.validate(req, apiKeyPayload() as any),
|
||||
).rejects.toThrow(UnauthorizedException);
|
||||
expect(req.raw.actor).toBeUndefined();
|
||||
expect(req.raw.authType).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import { Injectable, Logger, UnauthorizedException } from '@nestjs/common';
|
||||
import { Injectable, UnauthorizedException } from '@nestjs/common';
|
||||
import { PassportStrategy } from '@nestjs/passport';
|
||||
import { Strategy } from 'passport-jwt';
|
||||
import { EnvironmentService } from '../../../integrations/environment/environment.service';
|
||||
@@ -9,20 +9,18 @@ import { UserSessionRepo } from '@docmost/db/repos/session/user-session.repo';
|
||||
import { SessionActivityService } from '../../session/session-activity.service';
|
||||
import { FastifyRequest } from 'fastify';
|
||||
import { extractBearerTokenFromHeader, isUserDisabled } from '../../../common/helpers';
|
||||
import { ModuleRef } from '@nestjs/core';
|
||||
import { resolveProvenance } from '../../../common/decorators/auth-provenance.decorator';
|
||||
import { ApiKeyService } from '../../api-key/api-key.service';
|
||||
|
||||
@Injectable()
|
||||
export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
|
||||
private logger = new Logger('JwtStrategy');
|
||||
|
||||
constructor(
|
||||
private userRepo: UserRepo,
|
||||
private workspaceRepo: WorkspaceRepo,
|
||||
private userSessionRepo: UserSessionRepo,
|
||||
private sessionActivityService: SessionActivityService,
|
||||
private readonly environmentService: EnvironmentService,
|
||||
private moduleRef: ModuleRef,
|
||||
private readonly apiKeyService: ApiKeyService,
|
||||
) {
|
||||
super({
|
||||
jwtFromRequest: (req: FastifyRequest) => {
|
||||
@@ -102,12 +100,17 @@ export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
|
||||
}
|
||||
|
||||
private async validateApiKey(req: any, payload: JwtApiKeyPayload) {
|
||||
const apiKeyService = this.resolveApiKeyService();
|
||||
if (!apiKeyService) {
|
||||
throw new UnauthorizedException('Enterprise API Key module missing');
|
||||
}
|
||||
// The fork ships the core `ApiKeyService` (the EE `ee/api-key` module is
|
||||
// absent). `validate` throws a bare UnauthorizedException on any definite
|
||||
// deny (missing/revoked/expired row, disabled user, kill-switch off) and
|
||||
// propagates infra errors (→ 5xx) rather than masking them as a 401.
|
||||
const result = await this.apiKeyService.validate(payload);
|
||||
|
||||
const result = await apiKeyService.validateApiKey(payload);
|
||||
// Stamp the principal kind + key id so the /api-keys management surface can
|
||||
// enforce "a token cannot manage tokens". Done in this branch because it
|
||||
// returns before the shared ACCESS-path stamping below.
|
||||
req.raw.authType = 'api_key';
|
||||
req.raw.apiKeyId = payload.apiKeyId;
|
||||
|
||||
// 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
|
||||
@@ -119,32 +122,10 @@ export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
|
||||
// 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);
|
||||
const provenance = resolveProvenance(result.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
|
||||
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',
|
||||
);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,6 +6,7 @@ import {
|
||||
} from '@nestjs/common';
|
||||
import { UserModule } from './user/user.module';
|
||||
import { AuthModule } from './auth/auth.module';
|
||||
import { ApiKeyModule } from './api-key/api-key.module';
|
||||
import { WorkspaceModule } from './workspace/workspace.module';
|
||||
import { PageModule } from './page/page.module';
|
||||
import { AttachmentModule } from './attachment/attachment.module';
|
||||
@@ -29,6 +30,7 @@ import { ClsMiddleware } from 'nestjs-cls';
|
||||
imports: [
|
||||
UserModule,
|
||||
AuthModule,
|
||||
ApiKeyModule,
|
||||
WorkspaceModule,
|
||||
PageModule,
|
||||
AttachmentModule,
|
||||
|
||||
@@ -818,6 +818,11 @@ export class PageController {
|
||||
throw new NotFoundException('Page not found');
|
||||
}
|
||||
|
||||
// Target-only validateCanView is intentional: getPageBreadCrumbs returns
|
||||
// the full ancestor chain WITHOUT per-ancestor permission filtering. Safe
|
||||
// because page restrictions inherit down the tree, so any ancestor the
|
||||
// caller could not view would already hide the target here — see the
|
||||
// getPageBreadCrumbs docstring / #471.
|
||||
await this.pageAccessService.validateCanView(page, user);
|
||||
|
||||
return this.pageService.getPageBreadCrumbs(page.id);
|
||||
|
||||
@@ -1071,6 +1071,29 @@ export class PageService {
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Walk the ancestor chain of `childPageId` up to the space root, filtered
|
||||
* ONLY by `deletedAt` (+ MAX_PAGE_TREE_DEPTH) — WITHOUT per-ancestor
|
||||
* permission filtering. Callers that expose this to a user (the
|
||||
* `/breadcrumbs` endpoint) validate `validateCanView` on the TARGET page
|
||||
* only, then return the whole chain of ancestor titles (#471).
|
||||
*
|
||||
* This is safe — NOT a title leak — because page restrictions inherit DOWN
|
||||
* the tree: to view a page the caller must hold permission on EVERY
|
||||
* restricted ancestor (`validateCanView` -> `canUserAccessPage` checks the
|
||||
* full ancestor chain — see page-access.service.ts / page-permission.repo.ts
|
||||
* `canUserEditPage`). A restricted ancestor the caller may not see would
|
||||
* therefore already hide the TARGET page itself, so every ancestor reachable
|
||||
* here is one the caller is already entitled to view (content stays gated
|
||||
* regardless — getPage/getNode re-check permissions).
|
||||
*
|
||||
* Note the guarantee is the narrow "may view a descendant => may view its
|
||||
* ancestors", NOT "space membership sees every page" — restricted subtrees do
|
||||
* hide pages from members. Per-ancestor permission filtering here was
|
||||
* considered and declined as redundant given the inheritance invariant above
|
||||
* (#471). The same chain feeds the web-UI breadcrumb bar under identical CASL
|
||||
* scope.
|
||||
*/
|
||||
async getPageBreadCrumbs(childPageId: string, trx?: KyselyTransaction) {
|
||||
const ancestors = await dbOrTx(this.db, trx)
|
||||
.withRecursive('page_ancestors', (db) =>
|
||||
|
||||
@@ -37,6 +37,7 @@ import { AiProviderCredentialsRepo } from '@docmost/db/repos/ai-chat/ai-provider
|
||||
import { AiMcpServerRepo } from '@docmost/db/repos/ai-chat/ai-mcp-server.repo';
|
||||
import { AiAgentRoleRepo } from '@docmost/db/repos/ai-agent-roles/ai-agent-roles.repo';
|
||||
import { PageEmbeddingRepo } from '@docmost/db/repos/ai-chat/page-embedding.repo';
|
||||
import { ApiKeyRepo } from '@docmost/db/repos/api-key/api-key.repo';
|
||||
import { PageListener } from '@docmost/db/listeners/page.listener';
|
||||
import { PostgresJSDialect } from 'kysely-postgres-js';
|
||||
import * as postgres from 'postgres';
|
||||
@@ -129,6 +130,7 @@ import { firstSqlToken } from '../integrations/metrics/metrics.constants';
|
||||
AiMcpServerRepo,
|
||||
AiAgentRoleRepo,
|
||||
PageEmbeddingRepo,
|
||||
ApiKeyRepo,
|
||||
PageListener,
|
||||
],
|
||||
exports: [
|
||||
@@ -164,6 +166,7 @@ import { firstSqlToken } from '../integrations/metrics/metrics.constants';
|
||||
AiMcpServerRepo,
|
||||
AiAgentRoleRepo,
|
||||
PageEmbeddingRepo,
|
||||
ApiKeyRepo,
|
||||
],
|
||||
})
|
||||
export class DatabaseModule implements OnApplicationBootstrap {
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
import { type Kysely } from 'kysely';
|
||||
|
||||
/**
|
||||
* #370 — page-versioning intentionality tier on a history snapshot.
|
||||
*
|
||||
* Adds `page_history.kind`, the three-tier "how intentional was this snapshot"
|
||||
* marker that lets versions (intentional points) be told apart from autosaves:
|
||||
* - 'manual' — a human explicitly saved a version (Cmd+S / Save button)
|
||||
* - 'agent' — the AI agent explicitly saved a version
|
||||
* - 'idle' — trailing idle-flush autosnapshot (safety net)
|
||||
* - 'boundary' — autosnapshot pinned on a source transition (user↔agent↔git)
|
||||
*
|
||||
* Nullable with NO default (mirrors last_updated_source in the agent-provenance
|
||||
* migration): legacy rows predate the marker and read back as `null`, which the
|
||||
* client renders as a plain autosave. Stored as a short varchar to stay
|
||||
* forward-compatible without an enum migration.
|
||||
*/
|
||||
export async function up(db: Kysely<any>): Promise<void> {
|
||||
await db.schema
|
||||
.alterTable('page_history')
|
||||
.addColumn('kind', 'varchar(20)', (col) => col)
|
||||
.execute();
|
||||
}
|
||||
|
||||
export async function down(db: Kysely<any>): Promise<void> {
|
||||
await db.schema.alterTable('page_history').dropColumn('kind').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);
|
||||
},
|
||||
);
|
||||
});
|
||||
@@ -25,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);
|
||||
@@ -139,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,
|
||||
@@ -172,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
|
||||
@@ -214,7 +300,9 @@ export class AiChatMessageRepo {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
return db
|
||||
.updateTable('aiChatMessages')
|
||||
.set({ ...(patch as Record<string, unknown>), updatedAt: new Date() })
|
||||
// #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) =>
|
||||
@@ -249,7 +337,9 @@ export class AiChatMessageRepo {
|
||||
.set({
|
||||
status,
|
||||
metadata: sql`coalesce(metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`,
|
||||
updatedAt: new Date(),
|
||||
// #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)
|
||||
@@ -307,7 +397,9 @@ export class AiChatMessageRepo {
|
||||
.set({
|
||||
status: 'aborted',
|
||||
metadata: sql`coalesce(m.metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`,
|
||||
updatedAt: new Date(),
|
||||
// #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)
|
||||
@@ -351,7 +443,8 @@ export class AiChatMessageRepo {
|
||||
.set({
|
||||
status: 'aborted',
|
||||
metadata: sql`coalesce(metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`,
|
||||
updatedAt: new Date(),
|
||||
// #491: DB-clock stamp (see `update`). Staleness WHERE stays app-clock.
|
||||
updatedAt: sql`now()`,
|
||||
})
|
||||
.where('status', '=', 'streaming')
|
||||
.where('updatedAt', '<', staleBefore)
|
||||
|
||||
@@ -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,7 +136,11 @@ 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)
|
||||
@@ -162,14 +166,15 @@ export class AiChatRunRepo {
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<AiChatRun | undefined> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
const now = new Date();
|
||||
return db
|
||||
.updateTable('aiChatRuns')
|
||||
.set({
|
||||
status: patch.status,
|
||||
error: patch.error,
|
||||
finishedAt: now,
|
||||
updatedAt: now,
|
||||
// #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)
|
||||
@@ -192,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[])
|
||||
@@ -249,13 +255,14 @@ export class AiChatRunRepo {
|
||||
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[]);
|
||||
@@ -263,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();
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
import { Injectable } from '@nestjs/common';
|
||||
import { InjectKysely } from 'nestjs-kysely';
|
||||
import { KyselyDB, KyselyTransaction } from '@docmost/db/types/kysely.types';
|
||||
import { DB, Users } from '@docmost/db/types/db';
|
||||
import { dbOrTx } from '@docmost/db/utils';
|
||||
import { ApiKey, InsertableApiKey } from '@docmost/db/types/entity.types';
|
||||
import { jsonObjectFrom } from 'kysely/helpers/postgres';
|
||||
import { ExpressionBuilder } from 'kysely';
|
||||
|
||||
// Public-facing shape: the api_keys row plus a compact creator attribution used
|
||||
// by the admin list (the workspace-wide view attributes each key to its author).
|
||||
export type ApiKeyWithCreator = ApiKey & {
|
||||
creator: Pick<Users, 'id' | 'name' | 'email' | 'avatarUrl'> | null;
|
||||
};
|
||||
|
||||
@Injectable()
|
||||
export class ApiKeyRepo {
|
||||
constructor(@InjectKysely() private readonly db: KyselyDB) {}
|
||||
|
||||
// Compact creator attribution embedded in the admin list rows. A nested
|
||||
// subquery (jsonObjectFrom) keeps it one round-trip; the token material is
|
||||
// never stored, so nothing sensitive is joined in.
|
||||
private withCreator(eb: ExpressionBuilder<DB, 'apiKeys'>) {
|
||||
return jsonObjectFrom(
|
||||
eb
|
||||
.selectFrom('users')
|
||||
.select(['users.id', 'users.name', 'users.email', 'users.avatarUrl'])
|
||||
.whereRef('users.id', '=', 'apiKeys.creatorId'),
|
||||
).as('creator');
|
||||
}
|
||||
|
||||
async findById(
|
||||
id: string,
|
||||
workspaceId: string,
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<ApiKey | undefined> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
return db
|
||||
.selectFrom('apiKeys')
|
||||
.selectAll()
|
||||
.where('id', '=', id)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
// Convention (soft-delete): a revoked key is invisible to reads.
|
||||
.where('deletedAt', 'is', null)
|
||||
.executeTakeFirst();
|
||||
}
|
||||
|
||||
async findByCreator(
|
||||
creatorId: string,
|
||||
workspaceId: string,
|
||||
): Promise<ApiKeyWithCreator[]> {
|
||||
return this.db
|
||||
.selectFrom('apiKeys')
|
||||
.selectAll('apiKeys')
|
||||
.select((eb) => this.withCreator(eb))
|
||||
.where('creatorId', '=', creatorId)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.where('deletedAt', 'is', null)
|
||||
.orderBy('createdAt', 'desc')
|
||||
.execute() as unknown as Promise<ApiKeyWithCreator[]>;
|
||||
}
|
||||
|
||||
async findAllInWorkspace(
|
||||
workspaceId: string,
|
||||
): Promise<ApiKeyWithCreator[]> {
|
||||
return this.db
|
||||
.selectFrom('apiKeys')
|
||||
.selectAll('apiKeys')
|
||||
.select((eb) => this.withCreator(eb))
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.where('deletedAt', 'is', null)
|
||||
.orderBy('createdAt', 'desc')
|
||||
.execute() as unknown as Promise<ApiKeyWithCreator[]>;
|
||||
}
|
||||
|
||||
async insert(
|
||||
insertable: InsertableApiKey,
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<ApiKey> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
return db
|
||||
.insertInto('apiKeys')
|
||||
.values(insertable)
|
||||
.returningAll()
|
||||
.executeTakeFirst();
|
||||
}
|
||||
|
||||
// Soft-delete = revoke. Terminal: the row stays for forensics but is invisible
|
||||
// to every read above, so validate() denies it immediately (no grace window).
|
||||
async softDelete(id: string, workspaceId: string): Promise<void> {
|
||||
await this.db
|
||||
.updateTable('apiKeys')
|
||||
.set({ deletedAt: new Date(), updatedAt: new Date() })
|
||||
.where('id', '=', id)
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.where('deletedAt', 'is', null)
|
||||
.execute();
|
||||
}
|
||||
|
||||
// Throttled best-effort forensics stamp. Not on the critical path: callers
|
||||
// fire-and-forget and swallow errors, so it never fails a request.
|
||||
async touchLastUsed(id: string): Promise<void> {
|
||||
await this.db
|
||||
.updateTable('apiKeys')
|
||||
.set({ lastUsedAt: new Date() })
|
||||
.where('id', '=', id)
|
||||
.where('deletedAt', 'is', null)
|
||||
.execute();
|
||||
}
|
||||
}
|
||||
@@ -13,6 +13,7 @@ import { jsonArrayFrom, jsonObjectFrom } from 'kysely/helpers/postgres';
|
||||
import { ExpressionBuilder, sql } from 'kysely';
|
||||
import { DB } from '@docmost/db/types/db';
|
||||
import { resolveAgentProvenance } from '../agent-provenance';
|
||||
import { PageHistoryKind } from '../../../collaboration/constants';
|
||||
|
||||
/**
|
||||
* Role-resolution subquery for a page-history row's bound AI chat (#300). Joins
|
||||
@@ -46,6 +47,9 @@ export class PageHistoryRepo {
|
||||
'lastUpdatedById',
|
||||
'lastUpdatedSource',
|
||||
'lastUpdatedAiChatId',
|
||||
// #370 — intentionality tier ('manual' | 'agent' | 'idle' | 'boundary');
|
||||
// null on legacy rows (= autosave). Selected so callers can read/promote it.
|
||||
'kind',
|
||||
'contributorIds',
|
||||
'spaceId',
|
||||
'workspaceId',
|
||||
@@ -85,9 +89,15 @@ export class PageHistoryRepo {
|
||||
|
||||
async saveHistory(
|
||||
page: Page,
|
||||
opts?: { contributorIds?: string[]; trx?: KyselyTransaction },
|
||||
): Promise<void> {
|
||||
await this.insertPageHistory(
|
||||
opts?: {
|
||||
contributorIds?: string[];
|
||||
// #370 — intentionality tier for this snapshot. Omitted → null (legacy
|
||||
// autosave semantics). Callers derive it server-side, never from a client.
|
||||
kind?: PageHistoryKind;
|
||||
trx?: KyselyTransaction;
|
||||
},
|
||||
): Promise<PageHistory> {
|
||||
return await this.insertPageHistory(
|
||||
{
|
||||
pageId: page.id,
|
||||
slugId: page.slugId,
|
||||
@@ -99,6 +109,7 @@ export class PageHistoryRepo {
|
||||
// Copy the provenance marker off the page row, as for lastUpdatedById.
|
||||
lastUpdatedSource: page.lastUpdatedSource,
|
||||
lastUpdatedAiChatId: page.lastUpdatedAiChatId,
|
||||
kind: opts?.kind ?? null,
|
||||
contributorIds: opts?.contributorIds,
|
||||
spaceId: page.spaceId,
|
||||
workspaceId: page.workspaceId,
|
||||
@@ -107,6 +118,25 @@ export class PageHistoryRepo {
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* #370 — promote an existing snapshot's intentionality tier in place. Used by
|
||||
* the manual-save "promote-not-dup" path: when the latest history row already
|
||||
* holds the exact content being versioned, we upgrade its `kind` instead of
|
||||
* duplicating a heavy content row.
|
||||
*/
|
||||
async updateHistoryKind(
|
||||
pageHistoryId: string,
|
||||
kind: PageHistoryKind,
|
||||
trx?: KyselyTransaction,
|
||||
): Promise<void> {
|
||||
const db = dbOrTx(this.db, trx);
|
||||
await db
|
||||
.updateTable('pageHistory')
|
||||
.set({ kind })
|
||||
.where('id', '=', pageHistoryId)
|
||||
.execute();
|
||||
}
|
||||
|
||||
async findPageHistoryByPageId(pageId: string, pagination: PaginationOptions) {
|
||||
const query = this.db
|
||||
.selectFrom('pageHistory')
|
||||
|
||||
+1
@@ -280,6 +280,7 @@ export interface PageHistory {
|
||||
createdAt: Generated<Timestamp>;
|
||||
icon: string | null;
|
||||
id: Generated<string>;
|
||||
kind: string | null;
|
||||
lastUpdatedAiChatId: string | null;
|
||||
lastUpdatedById: string | null;
|
||||
lastUpdatedSource: string | null;
|
||||
|
||||
@@ -70,6 +70,23 @@ export class EnvironmentService {
|
||||
return this.configService.get<string>('JWT_TOKEN_EXPIRES_IN', '90d');
|
||||
}
|
||||
|
||||
// Kill-switch for the agent API-key feature. Default ON (unset -> enabled): a
|
||||
// deploy that never sets the variable must NOT silently kill every agent. The
|
||||
// parse is STRICT — the value is validated by environment.validation to be
|
||||
// exactly 'true' or 'false' (or absent), so a typo like `=0`/`=off`/`=False`
|
||||
// fails at boot rather than being read as "enabled" (which would leave an
|
||||
// operator who yanked the switch during an incident with it still on). When
|
||||
// OFF: validate() denies every api-key token and the issuance endpoints 404.
|
||||
isApiKeysEnabled(): boolean {
|
||||
return this.configService.get<string>('API_KEYS_ENABLED', 'true') !== 'false';
|
||||
}
|
||||
|
||||
// Raw value (or undefined) for the boot log, so the state after each deploy is
|
||||
// verifiable in container logs: `API keys: ENABLED/DISABLED (API_KEYS_ENABLED=...)`.
|
||||
getApiKeysEnabledRaw(): string | undefined {
|
||||
return this.configService.get<string>('API_KEYS_ENABLED');
|
||||
}
|
||||
|
||||
getCookieExpiresIn(): Date {
|
||||
const expiresInStr = this.getJwtTokenExpiresIn();
|
||||
let msUntilExpiry: number;
|
||||
|
||||
@@ -103,6 +103,15 @@ export class EnvironmentVariables {
|
||||
@IsString()
|
||||
TYPESENSE_LOCALE: string;
|
||||
|
||||
// Agent API-key kill-switch. Optional (absent -> default ON). STRICT: only the
|
||||
// literals 'true'/'false' are accepted, so `=0`/`=off`/`=False` fail at boot
|
||||
// instead of being silently read as "enabled" — the switch must actually flip
|
||||
// when an operator flips it. See EnvironmentService.isApiKeysEnabled.
|
||||
@IsOptional()
|
||||
@IsIn(['true', 'false'])
|
||||
@IsString()
|
||||
API_KEYS_ENABLED: string;
|
||||
|
||||
@IsOptional()
|
||||
@ValidateIf((obj) => obj.AI_DRIVER)
|
||||
@IsIn(['openai', 'openai-compatible', 'gemini', 'ollama'])
|
||||
|
||||
+131
@@ -0,0 +1,131 @@
|
||||
// p-limit and @sindresorhus/slugify are ESM-only and not in jest's transform
|
||||
// allowlist; both are irrelevant to createDrawioSvg (a pure fs + string method),
|
||||
// so they are mocked out to keep the module graph loadable under ts-jest.
|
||||
jest.mock('p-limit', () => ({
|
||||
__esModule: true,
|
||||
default: () => (fn: () => unknown) => fn(),
|
||||
}));
|
||||
jest.mock('@sindresorhus/slugify', () => ({
|
||||
__esModule: true,
|
||||
default: (input: string) => String(input),
|
||||
}));
|
||||
|
||||
import { promises as fs } from 'fs';
|
||||
import * as os from 'os';
|
||||
import * as path from 'path';
|
||||
import { ImportAttachmentService } from './import-attachment.service';
|
||||
|
||||
/**
|
||||
* Unit test for ImportAttachmentService.createDrawioSvg (issue #507).
|
||||
*
|
||||
* The Confluence import wraps a `.drawio` file into a `.drawio.svg` attachment.
|
||||
* The `content=` payload MUST be the mxfile XML entity-escaped (draw.io's native
|
||||
* form), NOT base64 — draw.io's editor decodes a base64 content= via Latin-1
|
||||
* atob, mangling every non-ASCII char into mojibake. createDrawioSvg touches no
|
||||
* injected dependency, so the service is built with placeholder deps.
|
||||
*/
|
||||
describe('ImportAttachmentService.createDrawioSvg (#507)', () => {
|
||||
const service = new ImportAttachmentService(
|
||||
{} as any,
|
||||
{} as any,
|
||||
{} as any,
|
||||
);
|
||||
const call = (p: string): Promise<Buffer> =>
|
||||
(service as any).createDrawioSvg(p);
|
||||
|
||||
let tmpDir: string;
|
||||
beforeAll(async () => {
|
||||
tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), 'drawio-507-'));
|
||||
});
|
||||
afterAll(async () => {
|
||||
await fs.rm(tmpDir, { recursive: true, force: true });
|
||||
});
|
||||
|
||||
const writeDrawio = async (name: string, xml: string): Promise<string> => {
|
||||
const p = path.join(tmpDir, name);
|
||||
await fs.writeFile(p, xml, 'utf-8');
|
||||
return p;
|
||||
};
|
||||
|
||||
it('writes content= as entity-encoded XML, not base64', async () => {
|
||||
const drawio =
|
||||
'<mxfile host="Confluence"><diagram name="Схема — ёж">' +
|
||||
'<mxGraphModel><root><mxCell id="0"/>' +
|
||||
'<mxCell id="2" value="Старт-бит" vertex="1" parent="0"/>' +
|
||||
'</root></mxGraphModel></diagram></mxfile>';
|
||||
const p = await writeDrawio('cyrillic.drawio', drawio);
|
||||
|
||||
const svg = (await call(p)).toString('utf-8');
|
||||
const content = /content="([^"]*)"/.exec(svg)?.[1];
|
||||
expect(content).toBeDefined();
|
||||
|
||||
// Entity-encoded XML form, starting with <mxfile — never a base64 blob.
|
||||
expect(content).toMatch(/^<mxfile/);
|
||||
expect(content).toContain('<');
|
||||
// Non-ASCII survives as raw UTF-8, with no Latin-1 mojibake.
|
||||
expect(content).toContain('Старт-бит');
|
||||
expect(content).toContain('Схема — ёж');
|
||||
expect(content).not.toContain('Ð');
|
||||
|
||||
// Decoding the attribute (un-escaping) yields the original drawio file.
|
||||
const decoded = content!
|
||||
.replace(/</g, '<')
|
||||
.replace(/>/g, '>')
|
||||
.replace(/"/g, '"')
|
||||
.replace(/&/g, '&');
|
||||
expect(decoded).toBe(drawio);
|
||||
});
|
||||
|
||||
it('encodes literal tab/newline/CR as numeric char-refs, not literal control chars (#507 F1)', async () => {
|
||||
// A literal tab/newline/CR inside the mxfile XML would be collapsed to a
|
||||
// single space by XML attribute-value normalization when the draw.io editor
|
||||
// reads content=, silently flattening multi-line labels and tab-bearing
|
||||
// values. They must be emitted as numeric char-refs instead.
|
||||
const drawio =
|
||||
'<mxfile><diagram name="p">' +
|
||||
'<mxGraphModel><root><mxCell id="0"/>' +
|
||||
'<mxCell id="2" value="col1\tcol2" style="html=1;\nshadow=0" vertex="1" parent="0"/>' +
|
||||
'<mxCell id="3" value="Line1\nLine2\rLine3" vertex="1" parent="0"/>' +
|
||||
'</root></mxGraphModel></diagram></mxfile>';
|
||||
const p = await writeDrawio('ctrl.drawio', drawio);
|
||||
|
||||
const svg = (await call(p)).toString('utf-8');
|
||||
const content = /content="([^"]*)"/.exec(svg)?.[1];
|
||||
expect(content).toBeDefined();
|
||||
// No literal control chars survive in the attribute value.
|
||||
expect(content).not.toMatch(/[\t\n\r]/);
|
||||
// They round-trip as numeric char-refs.
|
||||
expect(content).toContain('	');
|
||||
expect(content).toContain('
');
|
||||
expect(content).toContain('
');
|
||||
// Decoding (char-refs back to literal, entities back) recovers the file.
|
||||
const decoded = content!
|
||||
.replace(/</g, '<')
|
||||
.replace(/>/g, '>')
|
||||
.replace(/"/g, '"')
|
||||
.replace(/'/g, "'")
|
||||
.replace(/	/gi, '\t')
|
||||
.replace(/
/gi, '\n')
|
||||
.replace(/
/gi, '\r')
|
||||
.replace(/&/g, '&');
|
||||
expect(decoded).toBe(drawio);
|
||||
});
|
||||
|
||||
it('escapes XML metacharacters in the drawio payload', async () => {
|
||||
const drawio = '<mxfile><diagram name="a & b">"q" <x></diagram></mxfile>';
|
||||
const p = await writeDrawio('meta.drawio', drawio);
|
||||
|
||||
const svg = (await call(p)).toString('utf-8');
|
||||
const content = /content="([^"]*)"/.exec(svg)?.[1];
|
||||
expect(content).toBeDefined();
|
||||
// The attribute value must contain no bare `<`, `>` or `"` that would break
|
||||
// out of the content="..." attribute or the SVG element.
|
||||
expect(content).not.toMatch(/[<>"]/);
|
||||
const decoded = content!
|
||||
.replace(/</g, '<')
|
||||
.replace(/>/g, '>')
|
||||
.replace(/"/g, '"')
|
||||
.replace(/&/g, '&');
|
||||
expect(decoded).toBe(drawio);
|
||||
});
|
||||
});
|
||||
@@ -8,6 +8,7 @@ import { createReadStream } from 'node:fs';
|
||||
import { promises as fs } from 'fs';
|
||||
import { Readable } from 'stream';
|
||||
import { getMimeType, sanitizeFileName } from '../../../common/helpers';
|
||||
import { htmlEscape } from '../../../common/helpers/html-escaper';
|
||||
import { v7 } from 'uuid';
|
||||
import { FileTask } from '@docmost/db/types/entity.types';
|
||||
import { getAttachmentFolderPath } from '../../../core/attachment/attachment.utils';
|
||||
@@ -849,7 +850,12 @@ export class ImportAttachmentService {
|
||||
): Promise<Buffer> {
|
||||
try {
|
||||
const drawioContent = await fs.readFile(drawioPath, 'utf-8');
|
||||
const drawioBase64 = Buffer.from(drawioContent).toString('base64');
|
||||
// Write the mxfile XML XML-entity-escaped (draw.io's native content= form),
|
||||
// NOT base64. draw.io's editor decodes a base64 content= via Latin-1 atob
|
||||
// (no UTF-8 step), turning every non-ASCII char (Cyrillic, ё, —) into
|
||||
// mojibake; the entity-encoded form is decoded by the DOM as UTF-8 and
|
||||
// opens intact. Docmost's own decoder reads both forms.
|
||||
const drawioEscaped = this.xmlEscapeContent(drawioContent);
|
||||
|
||||
let imageElement = '';
|
||||
// If we have a PNG, include it in the SVG
|
||||
@@ -875,7 +881,7 @@ export class ImportAttachmentService {
|
||||
width="600"
|
||||
height="400"
|
||||
viewBox="0 0 600 400"
|
||||
content="${drawioBase64}">${imageElement}</svg>`;
|
||||
content="${drawioEscaped}">${imageElement}</svg>`;
|
||||
|
||||
return Buffer.from(svgContent, 'utf-8');
|
||||
} catch (error) {
|
||||
@@ -884,6 +890,24 @@ export class ImportAttachmentService {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Escape a string so it is safe as the value of a double-quoted XML attribute
|
||||
* (the `content=` payload of a `.drawio.svg`). The shared `htmlEscape` covers
|
||||
* `& < > " '` (a strict superset of what this attribute needs; the extra `'`
|
||||
* escape is harmless in a `"`-delimited value). On top of that, the numeric
|
||||
* char-refs for tab/newline/CR are required: a literal tab/newline/CR inside
|
||||
* an attribute value is collapsed to a single space by XML attribute-value
|
||||
* normalization on DOM read (both our decoder and the real draw.io editor),
|
||||
* silently flattening multi-line labels and tab-bearing values. Char-refs
|
||||
* survive that normalization (#507).
|
||||
*/
|
||||
private xmlEscapeContent(s: string): string {
|
||||
return htmlEscape(s)
|
||||
.replace(/\t/g, '	')
|
||||
.replace(/\n/g, '
')
|
||||
.replace(/\r/g, '
');
|
||||
}
|
||||
|
||||
private async uploadWithRetry(opts: {
|
||||
abs: string;
|
||||
storageFilePath: string;
|
||||
|
||||
@@ -304,37 +304,102 @@ export function clientIp(req: ClientIpRequest): string {
|
||||
return 'unknown';
|
||||
}
|
||||
|
||||
// Minimal structural shape of the TokenService.verifyJwt method we depend on,
|
||||
// so this module never imports the concrete TokenService (heavy graph).
|
||||
export interface AccessJwtVerifier {
|
||||
verifyJwt: (
|
||||
token: string,
|
||||
type: JwtType,
|
||||
) => Promise<{
|
||||
sub?: string;
|
||||
email?: string;
|
||||
workspaceId?: string;
|
||||
sessionId?: string;
|
||||
}>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Bind a TokenService-like verifier into a one-arg `verifyJwt(token)` that
|
||||
* ALWAYS enforces `JwtType.ACCESS`. This is the single place where the /mcp
|
||||
* Bearer path pins the token type: a Bearer access token must be verified AS an
|
||||
* access token (not refresh/exchange/collab/etc.), so the type literal is fixed
|
||||
* here rather than at the call site. McpService.verifyMcpBearer delegates to
|
||||
* this, keeping the `JwtType.ACCESS` choice testable without the heavy graph.
|
||||
*/
|
||||
export function bindAccessJwtVerifier(
|
||||
tokenService: AccessJwtVerifier,
|
||||
): (token: string) => Promise<{
|
||||
// The decoded payload shared by the /mcp Bearer allowlist. Carries the `type`
|
||||
// discriminator and the API-key `apiKeyId`, on top of the access-token fields.
|
||||
export interface McpBearerPayload {
|
||||
type?: JwtType;
|
||||
sub?: string;
|
||||
email?: string;
|
||||
workspaceId?: string;
|
||||
sessionId?: string;
|
||||
}> {
|
||||
return (token: string) => tokenService.verifyJwt(token, JwtType.ACCESS);
|
||||
apiKeyId?: string;
|
||||
}
|
||||
|
||||
// Minimal structural shape of the TokenService.verifyJwtOneOf method.
|
||||
export interface OneOfJwtVerifier {
|
||||
verifyJwtOneOf: (
|
||||
token: string,
|
||||
allowed: JwtType[],
|
||||
) => Promise<McpBearerPayload>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Bind a TokenService-like verifier into a one-arg `verifyJwtOneOf(token)` that
|
||||
* pins the /mcp Bearer ALLOWLIST to exactly {ACCESS, API_KEY}. This is the single
|
||||
* place the /mcp Bearer path pins the token type: the /mcp Bearer slot
|
||||
* legitimately accepts either an ACCESS token (a
|
||||
* human's session token) OR an API_KEY token (an agent's key), but NOTHING else
|
||||
* (collab/exchange/attachment/etc. are rejected with the generic type error).
|
||||
* The allowlist is fixed here rather than at the call site, and the signature is
|
||||
* verified exactly once (see verifyMcpBearer).
|
||||
*/
|
||||
export function bindMcpBearerVerifier(
|
||||
tokenService: OneOfJwtVerifier,
|
||||
): (token: string) => Promise<McpBearerPayload> {
|
||||
return (token: string) =>
|
||||
tokenService.verifyJwtOneOf(token, [JwtType.ACCESS, JwtType.API_KEY]);
|
||||
}
|
||||
|
||||
// Deps for the /mcp Bearer router. `verifyJwtOneOf` is the one-arg verifier bound
|
||||
// above (allowlist {ACCESS, API_KEY}); the ACCESS-specific revocation/disabled
|
||||
// deps mirror BearerVerifyDeps; `validateApiKey` is the SHARED api-key row-check.
|
||||
export interface McpBearerDeps
|
||||
extends Omit<BearerVerifyDeps, 'verifyJwt'> {
|
||||
verifyJwtOneOf: (token: string) => Promise<McpBearerPayload>;
|
||||
// Row-check for an API_KEY principal — the SAME validator REST uses. Throws
|
||||
// UnauthorizedException on a definite deny; PROPAGATES an infra error (→ 5xx),
|
||||
// never masking it as a 401. Not a login attempt: the Basic limiter is not
|
||||
// involved on this path.
|
||||
validateApiKey: (payload: McpBearerPayload) => Promise<unknown>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Verify a /mcp Bearer token that may be an ACCESS token OR an API_KEY token, and
|
||||
* route by type. The signature is verified EXACTLY ONCE (verifyJwtOneOf); the
|
||||
* result is reused so the ACCESS branch does not re-verify.
|
||||
*
|
||||
* - API_KEY -> bind to THIS instance's workspace FIRST (a token for another
|
||||
* workspace is rejected), THEN run the shared `validateApiKey` row-check.
|
||||
* No session/limiter involvement (an API key is not a login).
|
||||
* - ACCESS -> the unchanged `verifyBearerAccess` (session-active + not-disabled
|
||||
* checks), fed a closure over the already-verified payload so "verify once"
|
||||
* and "helper unchanged" coexist.
|
||||
*
|
||||
* Throws UnauthorizedException on any auth failure (uniform generic message — no
|
||||
* enumeration of why); propagates an infra error from `validateApiKey` as itself.
|
||||
*/
|
||||
export async function verifyMcpBearer(
|
||||
token: string,
|
||||
deps: McpBearerDeps,
|
||||
): Promise<{ sub?: string; email?: string }> {
|
||||
const generic = 'Invalid or expired token';
|
||||
const payload = await deps.verifyJwtOneOf(token);
|
||||
|
||||
if (payload.type === JwtType.API_KEY) {
|
||||
if (!payload.sub || !payload.workspaceId) {
|
||||
throw new UnauthorizedException(generic);
|
||||
}
|
||||
// Instance-binding (mirrors verifyBearerAccess): reject an API_KEY token
|
||||
// minted for a different workspace before touching the DB.
|
||||
if (
|
||||
deps.expectedWorkspaceId &&
|
||||
payload.workspaceId !== deps.expectedWorkspaceId
|
||||
) {
|
||||
throw new UnauthorizedException(generic);
|
||||
}
|
||||
// Shared row-check. A definite deny throws Unauthorized; an infra error
|
||||
// propagates (→ 5xx), which the caller must NOT convert to a 401.
|
||||
await deps.validateApiKey(payload);
|
||||
return { sub: payload.sub };
|
||||
}
|
||||
|
||||
// ACCESS: reuse verifyBearerAccess WITHOUT re-verifying the signature.
|
||||
return verifyBearerAccess(token, {
|
||||
verifyJwt: async () => payload,
|
||||
expectedWorkspaceId: deps.expectedWorkspaceId,
|
||||
findUser: deps.findUser,
|
||||
findActiveSession: deps.findActiveSession,
|
||||
});
|
||||
}
|
||||
|
||||
// Minimal shapes for the Bearer revocation/disabled check. Kept structural so
|
||||
@@ -728,18 +793,24 @@ export async function resolveMcpSessionConfig(
|
||||
};
|
||||
}
|
||||
|
||||
// --- 2) fallback A: Bearer access-JWT (user-supplied token) ---
|
||||
// --- 2) fallback A: Bearer JWT (user-supplied ACCESS or agent API_KEY) ---
|
||||
const bearer = extractBearer(authHeader);
|
||||
if (bearer) {
|
||||
let payload: { sub?: string; email?: string };
|
||||
try {
|
||||
payload = await deps.verifyAccessJwt(bearer);
|
||||
} catch (err) {
|
||||
const message =
|
||||
err instanceof Error && err.message
|
||||
? err.message
|
||||
: 'Invalid or expired token';
|
||||
throw new UnauthorizedException(message);
|
||||
// Anti-enumeration (Bearer leg): EVERY auth failure surfaces the SAME
|
||||
// generic 401 — expired/revoked/wrong-type/unknown are indistinguishable
|
||||
// to the caller (its reaction is identical either way). But an UNEXPECTED
|
||||
// (infra) error is NOT an auth verdict: rethrow it AS ITSELF so the surface
|
||||
// maps it to 5xx (mapAuthResultToResponse), never masking a DB/Redis
|
||||
// outage as a bad token. verifyMcpBearer throws UnauthorizedException on a
|
||||
// definite deny and lets an infra error from validateApiKey propagate.
|
||||
if (err instanceof UnauthorizedException) {
|
||||
throw new UnauthorizedException('Invalid or expired token');
|
||||
}
|
||||
throw err;
|
||||
}
|
||||
return {
|
||||
config: { apiUrl, getToken: async () => bearer },
|
||||
|
||||
@@ -115,6 +115,7 @@ function makeService(opts: {
|
||||
undefined as never, // userRepo
|
||||
undefined as never, // userSessionRepo
|
||||
moduleRef as never, // moduleRef (read by the MFA branch)
|
||||
undefined as never, // apiKeyService (unused by the login-gate path)
|
||||
undefined as never, // sandboxStore (unused by the login-gate path)
|
||||
);
|
||||
// Stop the constructor's unref'd sweep timer leaking across tests.
|
||||
|
||||
@@ -4,13 +4,16 @@ import { McpService } from './mcp.service';
|
||||
import { DatabaseModule } from '@docmost/db/database.module';
|
||||
import { AuthModule } from '../../core/auth/auth.module';
|
||||
import { TokenModule } from '../../core/auth/token.module';
|
||||
import { ApiKeyModule } from '../../core/api-key/api-key.module';
|
||||
|
||||
// Community MCP feature: the server itself serves the Model Context Protocol
|
||||
// over HTTP at /mcp. DatabaseModule (global) provides WorkspaceRepo. AuthModule
|
||||
// supplies AuthService (per-user HTTP-Basic login validation) and TokenModule
|
||||
// supplies TokenService (Bearer access-JWT verification for the token fallback).
|
||||
// supplies TokenService (Bearer JWT verification for the token path). ApiKeyModule
|
||||
// supplies ApiKeyService (the shared api-key row-check for the API_KEY Bearer
|
||||
// branch, so an agent authenticates with a key instead of the bcrypt Basic path).
|
||||
@Module({
|
||||
imports: [DatabaseModule, AuthModule, TokenModule],
|
||||
imports: [DatabaseModule, AuthModule, TokenModule, ApiKeyModule],
|
||||
controllers: [McpController],
|
||||
providers: [McpService],
|
||||
})
|
||||
|
||||
@@ -6,9 +6,10 @@ import {
|
||||
isCredentialsFailure,
|
||||
isInitializeRequestBody,
|
||||
verifyBearerAccess,
|
||||
verifyMcpBearer,
|
||||
bindMcpBearerVerifier,
|
||||
sharedTokenMatches,
|
||||
clientIp,
|
||||
bindAccessJwtVerifier,
|
||||
extractBearer,
|
||||
decideBasicGate,
|
||||
mapAuthResultToResponse,
|
||||
@@ -524,13 +525,28 @@ describe('resolveMcpSessionConfig', () => {
|
||||
expect(resolved.identity).toBe('bearer:user-9');
|
||||
});
|
||||
|
||||
it('Bearer invalid -> specific 401 from verifyAccessJwt', async () => {
|
||||
it('Bearer invalid -> UNIFORM generic 401 (anti-enumeration, reason not leaked)', async () => {
|
||||
// #501: the Bearer leg no longer surfaces the specific reason. Whether the
|
||||
// token is expired, revoked, wrong-type or unknown, the caller sees ONE bare
|
||||
// 'Invalid or expired token' — an agent's reaction is identical, and a leaked
|
||||
// class would be an enumeration oracle.
|
||||
const verifyAccessJwt = jest
|
||||
.fn()
|
||||
.mockRejectedValue(new UnauthorizedException('jwt expired'));
|
||||
await expect(
|
||||
resolveMcpSessionConfig('Bearer expired', makeDeps({ verifyAccessJwt })),
|
||||
).rejects.toThrow('jwt expired');
|
||||
).rejects.toThrow('Invalid or expired token');
|
||||
});
|
||||
|
||||
it('Bearer INFRA error -> propagates (NOT masked as 401)', async () => {
|
||||
// A non-UnauthorizedException (e.g. a DB outage in the api-key row-check) is
|
||||
// not an auth verdict: it must propagate so the surface maps it to 5xx.
|
||||
const verifyAccessJwt = jest
|
||||
.fn()
|
||||
.mockRejectedValue(new Error('connection terminated'));
|
||||
await expect(
|
||||
resolveMcpSessionConfig('Bearer x', makeDeps({ verifyAccessJwt })),
|
||||
).rejects.toThrow('connection terminated');
|
||||
});
|
||||
|
||||
it('no creds + env service account configured -> service-account config', async () => {
|
||||
@@ -1001,48 +1017,104 @@ describe('clientIp (XFF-fallback precedence, item 5)', () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe('bindAccessJwtVerifier enforces JwtType.ACCESS (item 3)', () => {
|
||||
it('calls TokenService.verifyJwt with JwtType.ACCESS as the second argument', async () => {
|
||||
// Mock TokenService: assert the type literal is pinned to ACCESS so swapping
|
||||
// to REFRESH (or omitting the type) breaks this test.
|
||||
const verifyJwt = jest
|
||||
describe('bindMcpBearerVerifier pins the {ACCESS, API_KEY} allowlist (#501)', () => {
|
||||
it('calls verifyJwtOneOf with exactly [ACCESS, API_KEY]', async () => {
|
||||
const verifyJwtOneOf = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ sub: 'user-1', workspaceId: 'ws-1' });
|
||||
const verify = bindAccessJwtVerifier({ verifyJwt });
|
||||
.mockResolvedValue({ type: JwtType.API_KEY, sub: 'u-1' });
|
||||
await bindMcpBearerVerifier({ verifyJwtOneOf })('the.jwt');
|
||||
expect(verifyJwtOneOf).toHaveBeenCalledWith('the.jwt', [
|
||||
JwtType.ACCESS,
|
||||
JwtType.API_KEY,
|
||||
]);
|
||||
// Pin the concrete enum values too.
|
||||
expect(verifyJwtOneOf.mock.calls[0][1]).toEqual(['access', 'api_key']);
|
||||
});
|
||||
});
|
||||
|
||||
await verify('the.access.jwt');
|
||||
|
||||
expect(verifyJwt).toHaveBeenCalledTimes(1);
|
||||
expect(verifyJwt).toHaveBeenCalledWith('the.access.jwt', JwtType.ACCESS);
|
||||
// Pin the real enum value too, so renaming/repointing the enum member is caught.
|
||||
expect(verifyJwt.mock.calls[0][1]).toBe('access');
|
||||
describe('verifyMcpBearer routes by token type (#501)', () => {
|
||||
const accessDeps = (over: any = {}) => ({
|
||||
verifyJwtOneOf: jest.fn(),
|
||||
expectedWorkspaceId: 'ws-1',
|
||||
findUser: jest.fn().mockResolvedValue({ deactivatedAt: null }),
|
||||
findActiveSession: jest
|
||||
.fn()
|
||||
.mockResolvedValue({ userId: 'u-1', workspaceId: 'ws-1' }),
|
||||
validateApiKey: jest.fn(),
|
||||
...over,
|
||||
});
|
||||
|
||||
it('passes through the verified payload', async () => {
|
||||
const payload = { sub: 'user-9', email: 'u@e.com', workspaceId: 'ws-1' };
|
||||
const verifyJwt = jest.fn().mockResolvedValue(payload);
|
||||
await expect(
|
||||
bindAccessJwtVerifier({ verifyJwt })('t'),
|
||||
).resolves.toBe(payload);
|
||||
it('API_KEY -> row-checks via validateApiKey and does NOT touch session/limiter', async () => {
|
||||
const deps = accessDeps({
|
||||
verifyJwtOneOf: jest.fn().mockResolvedValue({
|
||||
type: JwtType.API_KEY,
|
||||
sub: 'svc-1',
|
||||
workspaceId: 'ws-1',
|
||||
apiKeyId: 'k-1',
|
||||
}),
|
||||
validateApiKey: jest.fn().mockResolvedValue({ user: { id: 'svc-1' } }),
|
||||
});
|
||||
const res = await verifyMcpBearer('tok', deps);
|
||||
expect(deps.validateApiKey).toHaveBeenCalledTimes(1);
|
||||
// No session lookup on the api-key path (not a login).
|
||||
expect(deps.findActiveSession).not.toHaveBeenCalled();
|
||||
expect(res).toEqual({ sub: 'svc-1' });
|
||||
});
|
||||
|
||||
// The Bearer revocation/disabled checks (verifyBearerAccess) are covered above;
|
||||
// this binds the ACCESS-type enforcement that verifyMcpBearer wires in.
|
||||
it('feeds verifyBearerAccess so the whole Bearer chain enforces ACCESS', async () => {
|
||||
const verifyJwt = jest.fn().mockResolvedValue({
|
||||
sub: 'user-1',
|
||||
it('API_KEY for ANOTHER workspace -> rejected before the row-check', async () => {
|
||||
const deps = accessDeps({
|
||||
verifyJwtOneOf: jest.fn().mockResolvedValue({
|
||||
type: JwtType.API_KEY,
|
||||
sub: 'svc-1',
|
||||
workspaceId: 'ws-OTHER',
|
||||
apiKeyId: 'k-1',
|
||||
}),
|
||||
validateApiKey: jest.fn(),
|
||||
});
|
||||
await expect(verifyMcpBearer('tok', deps)).rejects.toBeInstanceOf(
|
||||
UnauthorizedException,
|
||||
);
|
||||
expect(deps.validateApiKey).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('API_KEY infra error from validateApiKey PROPAGATES (not masked)', async () => {
|
||||
const boom = new Error('db down');
|
||||
const deps = accessDeps({
|
||||
verifyJwtOneOf: jest.fn().mockResolvedValue({
|
||||
type: JwtType.API_KEY,
|
||||
sub: 'svc-1',
|
||||
workspaceId: 'ws-1',
|
||||
apiKeyId: 'k-1',
|
||||
}),
|
||||
validateApiKey: jest.fn().mockRejectedValue(boom),
|
||||
});
|
||||
await expect(verifyMcpBearer('tok', deps)).rejects.toBe(boom);
|
||||
});
|
||||
|
||||
it('ACCESS -> runs the session/disabled checks and does NOT call validateApiKey', async () => {
|
||||
const deps = accessDeps({
|
||||
verifyJwtOneOf: jest.fn().mockResolvedValue({
|
||||
type: JwtType.ACCESS,
|
||||
sub: 'u-1',
|
||||
workspaceId: 'ws-1',
|
||||
sessionId: 'sess-1',
|
||||
email: 'u@e.com',
|
||||
}),
|
||||
});
|
||||
const res = await verifyMcpBearer('tok', deps);
|
||||
expect(deps.findActiveSession).toHaveBeenCalledWith('sess-1');
|
||||
expect(deps.validateApiKey).not.toHaveBeenCalled();
|
||||
expect(res).toEqual({ sub: 'u-1', email: 'u@e.com' });
|
||||
});
|
||||
|
||||
it('verifies the signature exactly ONCE (single verifyJwtOneOf, no re-verify)', async () => {
|
||||
const verifyJwtOneOf = jest.fn().mockResolvedValue({
|
||||
type: JwtType.ACCESS,
|
||||
sub: 'u-1',
|
||||
workspaceId: 'ws-1',
|
||||
sessionId: 'sess-1',
|
||||
});
|
||||
const res = await verifyBearerAccess('t', {
|
||||
verifyJwt: bindAccessJwtVerifier({ verifyJwt }),
|
||||
findUser: jest.fn().mockResolvedValue({ deactivatedAt: null }),
|
||||
findActiveSession: jest
|
||||
.fn()
|
||||
.mockResolvedValue({ userId: 'user-1', workspaceId: 'ws-1' }),
|
||||
});
|
||||
expect(verifyJwt).toHaveBeenCalledWith('t', JwtType.ACCESS);
|
||||
expect(res).toEqual({ sub: 'user-1', email: undefined });
|
||||
await verifyMcpBearer('tok', accessDeps({ verifyJwtOneOf }));
|
||||
expect(verifyJwtOneOf).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1198,6 +1270,7 @@ describe('McpService.onModuleDestroy — CollabSession teardown (#486)', () => {
|
||||
{} as any,
|
||||
{} as any,
|
||||
{} as any,
|
||||
{} as any,
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -14,16 +14,17 @@ import { UserSessionRepo } from '@docmost/db/repos/session/user-session.repo';
|
||||
import { AuthService } from '../../core/auth/services/auth.service';
|
||||
import { TokenService } from '../../core/auth/services/token.service';
|
||||
import { validateSsoEnforcement } from '../../core/auth/auth.util';
|
||||
import { JwtPayload } from '../../core/auth/dto/jwt-payload';
|
||||
import { JwtApiKeyPayload } from '../../core/auth/dto/jwt-payload';
|
||||
import { Workspace } from '@docmost/db/types/entity.types';
|
||||
import { ApiKeyService } from '../../core/api-key/api-key.service';
|
||||
import {
|
||||
FailedLoginLimiter,
|
||||
resolveMcpSessionConfig,
|
||||
verifyBearerAccess,
|
||||
verifyMcpBearer,
|
||||
isInitializeRequestBody,
|
||||
sharedTokenMatches,
|
||||
clientIp,
|
||||
bindAccessJwtVerifier,
|
||||
bindMcpBearerVerifier,
|
||||
decideBasicGate,
|
||||
mapAuthResultToResponse,
|
||||
DocmostMcpConfig,
|
||||
@@ -105,6 +106,10 @@ export class McpService implements OnModuleDestroy {
|
||||
private readonly userRepo: UserRepo,
|
||||
private readonly userSessionRepo: UserSessionRepo,
|
||||
private readonly moduleRef: ModuleRef,
|
||||
// Shared api-key row-check for the /mcp API_KEY Bearer branch (same validator
|
||||
// REST uses). Also lets an agent authenticate to /mcp with an api key instead
|
||||
// of the bcrypt Basic path, so parallel reads stop starving the limiter.
|
||||
private readonly apiKeyService: ApiKeyService,
|
||||
// Shared singleton in-RAM blob store backing the stash tool.
|
||||
private readonly sandboxStore: SandboxStore,
|
||||
) {
|
||||
@@ -194,37 +199,41 @@ export class McpService implements OnModuleDestroy {
|
||||
}
|
||||
}
|
||||
|
||||
// Bearer access-JWT verification for the /mcp token fallback. verifyJwt only
|
||||
// checks signature/exp/type, but a logged-out (revoked) or disabled user can
|
||||
// still hold an unexpired access JWT. JwtStrategy additionally checks the
|
||||
// session is active and the user is not disabled; we mirror those exact checks
|
||||
// here so the MCP Bearer path is not weaker than the normal cookie/header path.
|
||||
// Bearer verification for the /mcp token path. The Bearer slot accepts EITHER
|
||||
// an ACCESS token (a human session token) OR an API_KEY token (an agent's key)
|
||||
// — the allowlist is pinned in bindMcpBearerVerifier. An ACCESS token is
|
||||
// checked exactly as JwtStrategy does (signature/exp/type + session-active +
|
||||
// not-disabled), so the MCP path is not weaker than the cookie/header path. An
|
||||
// API_KEY token is HMAC-verified (microseconds) then row-checked via the shared
|
||||
// ApiKeyService.validate — NOT a login attempt, so the Basic bcrypt path and
|
||||
// its anti-brute-force limiter are never touched (the parallel-reads fix).
|
||||
private async verifyMcpBearer(
|
||||
token: string,
|
||||
): Promise<{ sub?: string; email?: string }> {
|
||||
// Resolve THIS instance's workspace so verifyBearerAccess can bind the
|
||||
// token's `workspaceId` claim to it (mirrors JwtStrategy). The community
|
||||
// build is single-workspace (findFirst), so this is the default workspace
|
||||
// and the check is a no-op here; it only rejects a foreign-workspace token
|
||||
// in a multi-workspace deployment. Undefined (no workspace configured) means
|
||||
// no check — the credentials path would already have failed with no
|
||||
// workspace, and an undefined here keeps the helper a no-op rather than
|
||||
// rejecting every token.
|
||||
// Resolve THIS instance's workspace so the router can bind the token's
|
||||
// `workspaceId` claim to it (mirrors JwtStrategy). The community build is
|
||||
// single-workspace (findFirst), so this is the default workspace and the
|
||||
// check is a no-op here; it only rejects a foreign-workspace token in a
|
||||
// multi-workspace deployment. Undefined (no workspace configured) means no
|
||||
// check — the credentials path would already have failed with no workspace.
|
||||
const instanceWorkspace = await this.workspaceRepo.findFirst();
|
||||
// The revocation/disabled decision logic lives in the framework-free
|
||||
// verifyBearerAccess helper (unit-testable without the heavy auth graph);
|
||||
// this method only wires in the concrete TokenService + repos.
|
||||
return verifyBearerAccess(token, {
|
||||
// The JwtType.ACCESS enforcement lives in bindAccessJwtVerifier (a pure,
|
||||
// testable seam) so the type literal cannot silently drift to REFRESH.
|
||||
verifyJwt: bindAccessJwtVerifier(this.tokenService) as (
|
||||
t: string,
|
||||
) => Promise<JwtPayload>,
|
||||
// The type-routing + revocation/disabled decision logic lives in the
|
||||
// framework-free verifyMcpBearer helper (unit-testable without the heavy auth
|
||||
// graph); this method only wires in the concrete TokenService + repos + the
|
||||
// shared api-key validator.
|
||||
return verifyMcpBearer(token, {
|
||||
// The {ACCESS, API_KEY} allowlist enforcement lives in bindMcpBearerVerifier
|
||||
// (a pure, testable seam) so the type set cannot silently drift.
|
||||
verifyJwtOneOf: bindMcpBearerVerifier(this.tokenService),
|
||||
expectedWorkspaceId: instanceWorkspace?.id,
|
||||
findUser: (sub, workspaceId) =>
|
||||
this.userRepo.findById(sub, workspaceId),
|
||||
findActiveSession: (sessionId) =>
|
||||
this.userSessionRepo.findActiveById(sessionId),
|
||||
// Shared with REST: a definite deny throws Unauthorized, an infra error
|
||||
// propagates (→ 5xx). The /mcp bearer catch must preserve that distinction.
|
||||
validateApiKey: (payload) =>
|
||||
this.apiKeyService.validate(payload as JwtApiKeyPayload),
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
@@ -20,6 +20,10 @@ export interface IStripeSeatsSyncJob {
|
||||
|
||||
export interface IPageHistoryJob {
|
||||
pageId: string;
|
||||
// #370 — intentionality tier the worker stamps on the snapshot. All jobs on
|
||||
// this queue are trailing idle-flush autosnapshots, so this is 'idle' (absent
|
||||
// → treated as 'idle' by the processor).
|
||||
kind?: 'idle';
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -30,11 +30,13 @@ import {
|
||||
* tees the SSE frames into it via `consumeSseStream` while stamping the DB row id
|
||||
* via `generateMessageId` (both gated on runId + the resumable flag).
|
||||
*
|
||||
* Proven here: a finished run's replay is the full frame sequence incl `[DONE]`
|
||||
* with `start.messageId` == the seeded DB row id; the anchor check (invariant 6);
|
||||
* an attach opened BEFORE the first frame follows the live stream from frame 0; an
|
||||
* explicit stop surfaces `{"type":"abort"}` + `[DONE]` + end to the subscriber;
|
||||
* and the legacy (non-run) path tees nothing.
|
||||
* Proven here (tail-only #491): a finished run attached at its persisted frontier
|
||||
* N_final delivers only the TAIL past N (a synthetic `start` carrying the run-fact
|
||||
* + the terminal `finish`/`[DONE]`) — the step content below N lives in the seeded
|
||||
* DB row, NOT the ring; the anchor check (invariant 6); an attach opened BEFORE the
|
||||
* first frame follows the live stream from frame 0; an explicit stop surfaces
|
||||
* `{"type":"abort"}` + `[DONE]` + end to the subscriber; and the legacy (non-run)
|
||||
* path tees nothing.
|
||||
*/
|
||||
|
||||
const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));
|
||||
@@ -133,14 +135,16 @@ function liveSink(): {
|
||||
};
|
||||
}
|
||||
|
||||
// The SSE `start` frame carries the message id; pull it out of a `data: {...}`.
|
||||
function parseStartMessageId(frames: string[]): string | undefined {
|
||||
// Parse the first `start` frame's JSON out of a `data: {...}` sequence.
|
||||
function parseStartFrame(
|
||||
frames: string[],
|
||||
): { messageId?: string; messageMetadata?: any } | undefined {
|
||||
for (const f of frames) {
|
||||
const m = /^data: (\{.*\})\s*$/m.exec(f.trim());
|
||||
if (!m) continue;
|
||||
try {
|
||||
const json = JSON.parse(m[1]);
|
||||
if (json.type === 'start') return json.messageId;
|
||||
if (json.type === 'start') return json;
|
||||
} catch {
|
||||
/* not this frame */
|
||||
}
|
||||
@@ -270,7 +274,7 @@ describe('AiChatService run-stream attach [integration]', () => {
|
||||
await destroyTestDb();
|
||||
});
|
||||
|
||||
it('run-wrapped: replay is the full frame sequence incl [DONE], start.messageId == the seeded DB row id', async () => {
|
||||
it('run-wrapped, tail-only: a finished run at N_final delivers the run-fact start + finish/[DONE]; the step content lives in the seeded row', async () => {
|
||||
const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
|
||||
const registry = new AiChatStreamRegistryService();
|
||||
const runService = new AiChatRunService(runRepo, {
|
||||
@@ -300,27 +304,37 @@ describe('AiChatService run-stream attach [integration]', () => {
|
||||
);
|
||||
});
|
||||
const rowId = await assistantRowId(chatId);
|
||||
// The client reads its persisted step frontier N from the seeded row.
|
||||
const row: any = await msgRepo.findById(rowId, workspaceId);
|
||||
const nFinal = row.metadata.stepsPersisted as number;
|
||||
expect(nFinal).toBe(1); // a single finished step
|
||||
// The step content is in the SEEDED row (parts/content), not the ring.
|
||||
expect(JSON.stringify(row.metadata.parts)).toContain('Hello');
|
||||
|
||||
// Finished-run replay with expect=live + the correct anchor.
|
||||
// Attach at N_final with the correct anchor: the tail past step 1 is just
|
||||
// the terminal frames; step 0's 'Hello' is BELOW the frontier (seeded).
|
||||
const sink = liveSink();
|
||||
const att = await registry.attach(chatId, true, rowId, sink.cb);
|
||||
const att = await registry.attach(chatId, rowId, nFinal, sink.cb);
|
||||
expect(att).not.toBeNull();
|
||||
expect(att!.finished).toBe(true);
|
||||
// The tee captured frames (consumeSseStream was wired).
|
||||
expect(att!.replay.length).toBeGreaterThan(0);
|
||||
// generateMessageId stamped the DB row id onto the streamed start frame.
|
||||
expect(parseStartMessageId(att!.replay)).toBe(rowId);
|
||||
// The full sequence includes the streamed text and the terminal marker.
|
||||
const joined = att!.replay.join('');
|
||||
expect(joined).toContain('Hello');
|
||||
// The synthetic start frame carries the run-fact (runId/chatId), the source
|
||||
// of the run-fact on re-attach.
|
||||
const start = parseStartFrame(att!.replay);
|
||||
expect(start?.messageMetadata).toMatchObject({
|
||||
runId: box.runId,
|
||||
chatId,
|
||||
});
|
||||
// The terminal marker is delivered so the client's SDK closes the stream.
|
||||
expect(att!.replay.some((f) => f.includes('[DONE]'))).toBe(true);
|
||||
// 'Hello' (step 0, below the frontier) is NOT re-streamed — it is seeded.
|
||||
expect(att!.replay.some((f) => f.includes('Hello'))).toBe(false);
|
||||
} finally {
|
||||
registry.onModuleDestroy();
|
||||
await cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
it('anchor mismatch with expect=live returns null (invariant 6)', async () => {
|
||||
it('anchor mismatch returns null (invariant 6)', async () => {
|
||||
const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
|
||||
const registry = new AiChatStreamRegistryService();
|
||||
const runService = new AiChatRunService(runRepo, {
|
||||
@@ -347,7 +361,7 @@ describe('AiChatService run-stream attach [integration]', () => {
|
||||
const sink = liveSink();
|
||||
// A foreign anchor must NOT replay this run's transcript.
|
||||
expect(
|
||||
await registry.attach(chatId, true, 'a-different-run-row', sink.cb),
|
||||
await registry.attach(chatId, 'a-different-run-row', 1, sink.cb),
|
||||
).toBeNull();
|
||||
} finally {
|
||||
registry.onModuleDestroy();
|
||||
@@ -376,8 +390,11 @@ describe('AiChatService run-stream attach [integration]', () => {
|
||||
try {
|
||||
// Attach while the entry exists (opened at begin) but before any frame.
|
||||
const sink = liveSink();
|
||||
const att = (await registry.attach(chatId, false, undefined, sink.cb))!;
|
||||
expect(att.replay).toEqual([]); // nothing streamed yet -> replay from 0
|
||||
const att = (await registry.attach(chatId, undefined, 0, sink.cb))!;
|
||||
// Nothing streamed yet -> the tail is just the synthetic start frame; the
|
||||
// whole live stream (start..DONE) follows via onFrame after start().
|
||||
expect(att.replay).toHaveLength(1);
|
||||
expect(att.replay[0]).toContain('"type":"start"');
|
||||
att.start(); // go live (drains nothing, then follows)
|
||||
|
||||
// Now emit the whole turn.
|
||||
@@ -448,7 +465,7 @@ describe('AiChatService run-stream attach [integration]', () => {
|
||||
});
|
||||
try {
|
||||
const sink = liveSink();
|
||||
const att = (await registry.attach(chatId, false, undefined, sink.cb))!;
|
||||
const att = (await registry.attach(chatId, undefined, 0, sink.cb))!;
|
||||
att.start();
|
||||
|
||||
// Give streamText a beat to begin consuming the partial output.
|
||||
@@ -526,7 +543,9 @@ describe('AiChatService run-stream attach [integration]', () => {
|
||||
expect(entry).toBeDefined();
|
||||
expect(entry.finished).toBe(true);
|
||||
const sink = liveSink();
|
||||
expect(await registry.attach(chatId, false, undefined, sink.cb)).toBeNull();
|
||||
// Finished with an EMPTY ring (aborted before any frame) -> null -> the
|
||||
// client degrades to poll instead of hanging on an empty stream.
|
||||
expect(await registry.attach(chatId, undefined, 0, sink.cb)).toBeNull();
|
||||
} finally {
|
||||
registry.onModuleDestroy();
|
||||
await cleanup();
|
||||
@@ -556,8 +575,8 @@ describe('AiChatService run-stream attach [integration]', () => {
|
||||
});
|
||||
const sink = liveSink();
|
||||
// No entry was ever opened; attach always yields null.
|
||||
expect(await registry.attach(chatId, false, undefined, sink.cb)).toBeNull();
|
||||
expect(await registry.attach(chatId, true, 'anything', sink.cb)).toBeNull();
|
||||
expect(await registry.attach(chatId, undefined, 0, sink.cb)).toBeNull();
|
||||
expect(await registry.attach(chatId, 'anything', 1, sink.cb)).toBeNull();
|
||||
} finally {
|
||||
registry.onModuleDestroy();
|
||||
await cleanup();
|
||||
|
||||
@@ -322,20 +322,21 @@ describe('AiChatService.stream [integration]', () => {
|
||||
});
|
||||
|
||||
/**
|
||||
* #332 deferred tool loading, the ON path. The riskiest property is that the
|
||||
* per-turn `activatedTools` Set is created FRESH inside each stream() call, so a
|
||||
* tool a previous turn activated via loadTools is NOT still active when the next
|
||||
* turn starts — the new turn begins "cold" (CORE + loadTools only). The unit
|
||||
* tests only exercise pure prepareAgentStep with hand-fed Sets; this pins the
|
||||
* real wiring end-to-end (loadTools.execute -> activatedTools -> prepareStep ->
|
||||
* per-step activeTools) against the real streamText loop, and proves there is no
|
||||
* cross-turn leak. We drive a MockLanguageModelV3 whose step 1 calls
|
||||
* loadTools(['createPage']) and assert, via the model's recorded per-step
|
||||
* CallOptions.tools (the AI SDK filters the provider tool list by activeTools),
|
||||
* that the deferred tool becomes active on the SAME turn's next step but NOT on a
|
||||
* fresh turn's first step.
|
||||
* #332 + #490 deferred tool loading, the ON path. Turn 1 starts COLD (CORE +
|
||||
* loadTools only) and activates a deferred tool via loadTools; that activation
|
||||
* is PERSISTED into the chat's metadata.activatedTools (#490) so the NEXT turn
|
||||
* SEEDS from it and the tool is active from the fresh turn's FIRST step — the
|
||||
* model never re-runs loadTools to re-activate the same tool. The unit tests
|
||||
* only exercise pure prepareAgentStep with hand-fed Sets; this pins the real
|
||||
* wiring end-to-end (loadTools.execute -> activatedTools -> persist -> next-turn
|
||||
* seed -> prepareStep -> per-step activeTools) against the real streamText loop.
|
||||
* We drive a MockLanguageModelV3 whose step 1 calls loadTools(['createPage'])
|
||||
* and assert, via the model's recorded per-step CallOptions.tools (the AI SDK
|
||||
* filters the provider tool list by activeTools), that the deferred tool becomes
|
||||
* active on the SAME turn's next step AND, seeded from metadata, on the next
|
||||
* turn's first step.
|
||||
*/
|
||||
describe('deferred tool loading ON — per-turn activation, no leak (#332)', () => {
|
||||
describe('deferred tool loading ON — cross-turn activation persistence (#332 + #490)', () => {
|
||||
// A stub deferred (non-core) tool the agent can activate. Its execute is never
|
||||
// called — the model only needs to SEE it become active — but it must be a
|
||||
// valid AI-SDK tool so the SDK includes it in a step's tool list once active.
|
||||
@@ -451,7 +452,7 @@ describe('AiChatService.stream [integration]', () => {
|
||||
} as any);
|
||||
}
|
||||
|
||||
it('activates a deferred tool for the SAME turn, and a NEW turn starts cold (no leak)', async () => {
|
||||
it('activates a deferred tool for the SAME turn, and a NEW turn SEEDS it from persisted chat metadata (#490)', async () => {
|
||||
const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
|
||||
|
||||
// --- Turn 1: loadTools(createPage) on step 1, then answer on step 2. ---
|
||||
@@ -474,7 +475,7 @@ describe('AiChatService.stream [integration]', () => {
|
||||
// Step 2 of the SAME turn sees the just-activated deferred tool.
|
||||
expect(step2Tools).toContain('createPage');
|
||||
|
||||
// --- Turn 2 on the SAME chat: must start cold again. ---
|
||||
// --- Turn 2 on the SAME chat: seeds the persisted activation (#490). ---
|
||||
const model2 = new MockLanguageModelV3({
|
||||
doStream: async () => ({ stream: successStream() }),
|
||||
} as any);
|
||||
@@ -485,9 +486,10 @@ describe('AiChatService.stream [integration]', () => {
|
||||
|
||||
const nextTurnFirstStep = toolNames(model2.doStreamCalls[0]);
|
||||
expect(nextTurnFirstStep).toContain('loadTools');
|
||||
// The activated set is per-turn: the prior turn's createPage did NOT leak,
|
||||
// so the fresh turn's first step sees it deferred again.
|
||||
expect(nextTurnFirstStep).not.toContain('createPage');
|
||||
// #490: activation PERSISTS across turns — turn 1 wrote createPage into the
|
||||
// chat's metadata.activatedTools, so the next turn seeds from it and the
|
||||
// deferred tool is active from the FIRST step (no need to re-run loadTools).
|
||||
expect(nextTurnFirstStep).toContain('createPage');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -0,0 +1,162 @@
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import { Queue, Worker } from 'bullmq';
|
||||
import { PersistenceExtension } from '../../src/collaboration/extensions/persistence.extension';
|
||||
|
||||
/**
|
||||
* #370 — integration property of the idle-snapshot pipeline against REAL BullMQ.
|
||||
*
|
||||
* This is deliberately NOT a unit test of computeHistoryJob (that lives in
|
||||
* compute-history-job.spec.ts). The point here is the OBSERVABLE end-to-end
|
||||
* behaviour of the production `enqueuePageHistory` remove-then-add debounce
|
||||
* driving a real Redis-backed delayed queue + worker (the #431→#439 class: a
|
||||
* locally-correct function whose queue/timer property was never exercised):
|
||||
*
|
||||
* - a CONTINUOUS burst of stores lasting several caps yields periodic idle
|
||||
* snapshots — at least one per max-wait cap, NOT one-per-store; and
|
||||
* - an INTERMITTENT burst (a few stores, then quiet) yields exactly ONE
|
||||
* trailing snapshot.
|
||||
*
|
||||
* We shrink the idle windows to milliseconds (jest.mock of collaboration
|
||||
* constants) so real BullMQ delayed jobs actually promote within the test —
|
||||
* fake timers cannot advance Redis's own delayed-set clock, so the intervals
|
||||
* must be real but tiny. The production method under test is called verbatim.
|
||||
*/
|
||||
|
||||
// NOTE: jest.mock is hoisted above the module's const initializers, so its
|
||||
// factory cannot close over MAX_WAIT_MS/INTERVAL_MS — the literals are inlined
|
||||
// here and MUST stay in sync with the consts below (a single source of truth is
|
||||
// impossible across the hoist boundary).
|
||||
jest.mock('../../src/collaboration/constants', () => {
|
||||
const actual = jest.requireActual('../../src/collaboration/constants');
|
||||
return {
|
||||
...actual,
|
||||
IDLE_MAX_WAIT_USER: 300,
|
||||
IDLE_MAX_WAIT_AGENT: 300,
|
||||
IDLE_INTERVAL_USER: 1000,
|
||||
IDLE_INTERVAL_AGENT: 1000,
|
||||
};
|
||||
});
|
||||
|
||||
// Mirrors the mocked IDLE_MAX_WAIT_* above (IDLE_INTERVAL_* is 1000 > this, so
|
||||
// the max-wait ceiling is what actually governs the trailing delay).
|
||||
const MAX_WAIT_MS = 300;
|
||||
|
||||
const REDIS_CONNECTION = {
|
||||
host: process.env.TEST_REDIS_HOST ?? '127.0.0.1',
|
||||
port: Number(process.env.TEST_REDIS_PORT ?? 6379),
|
||||
};
|
||||
|
||||
const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));
|
||||
|
||||
describe('#370 idle-snapshot pipeline (real BullMQ)', () => {
|
||||
let queue: Queue;
|
||||
let worker: Worker;
|
||||
let extension: PersistenceExtension;
|
||||
// Every processed snapshot, tagged by pageId so the two scenarios stay isolated.
|
||||
const processed: Array<{ pageId: string; kind: string; at: number }> = [];
|
||||
const queueName = `history-idle-int-${randomUUID()}`;
|
||||
|
||||
beforeAll(async () => {
|
||||
queue = new Queue(queueName, {
|
||||
connection: REDIS_CONNECTION,
|
||||
// Mirror the production default (BullModule.forRoot removeOnComplete): the
|
||||
// enqueue idiom relies on the jobId being freed once a job completes so the
|
||||
// next burst can re-arm the same id.
|
||||
defaultJobOptions: { removeOnComplete: true, removeOnFail: true },
|
||||
});
|
||||
await queue.waitUntilReady();
|
||||
|
||||
worker = new Worker(
|
||||
queueName,
|
||||
async (job) => {
|
||||
processed.push({
|
||||
pageId: job.data?.pageId,
|
||||
kind: job.data?.kind,
|
||||
at: Date.now(),
|
||||
});
|
||||
},
|
||||
{ connection: REDIS_CONNECTION },
|
||||
);
|
||||
await worker.waitUntilReady();
|
||||
|
||||
// Construct the real extension; only historyQueue (5th ctor arg) and the
|
||||
// internal idleBurstStart map are exercised by enqueuePageHistory, so the
|
||||
// other collaborators can be null — the constructor only assigns fields.
|
||||
extension = new PersistenceExtension(
|
||||
null as any, // pageRepo
|
||||
null as any, // pageHistoryRepo
|
||||
null as any, // db
|
||||
null as any, // aiQueue
|
||||
queue as any, // historyQueue
|
||||
null as any, // notificationQueue
|
||||
null as any, // collabHistory
|
||||
null as any, // transclusionService
|
||||
);
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
// Force-close and fully drain so no BullMQ background activity (delayed-set
|
||||
// polling, blocking BRPOPLPUSH) bleeds into later suites in this single
|
||||
// shared jest worker (maxWorkers: 1).
|
||||
await worker?.close(true).catch(() => undefined);
|
||||
await queue?.obliterate({ force: true }).catch(() => undefined);
|
||||
await queue?.close();
|
||||
// Let the redis sockets settle before the next suite starts.
|
||||
await sleep(150);
|
||||
});
|
||||
|
||||
const arm = (pageId: string) =>
|
||||
(extension as any).enqueuePageHistory({ id: pageId }, 'user');
|
||||
|
||||
it('continuous burst over several caps → periodic idle snapshots (≥1 per cap, not one-per-store)', async () => {
|
||||
const pageId = randomUUID();
|
||||
const runMs = 6 * MAX_WAIT_MS; // ~6 caps of unbroken editing
|
||||
// Store cadence that does NOT evenly divide the cap: real hocuspocus stores
|
||||
// are not aligned to cap boundaries, so a boundary job promotes in the gap
|
||||
// before the next store's remove(). A cap-aligned cadence would instead land
|
||||
// a store exactly on every boundary and lose the snapshot to the documented
|
||||
// remove-vs-active race — an artefact of the test clock, not the pipeline.
|
||||
const stepMs = 70;
|
||||
const stores = Math.floor(runMs / stepMs);
|
||||
|
||||
const start = Date.now();
|
||||
let count = 0;
|
||||
while (Date.now() - start < runMs) {
|
||||
await arm(pageId);
|
||||
count++;
|
||||
await sleep(stepMs);
|
||||
}
|
||||
// Let the final armed job flush.
|
||||
await sleep(2 * MAX_WAIT_MS);
|
||||
|
||||
const snaps = processed.filter((p) => p.pageId === pageId);
|
||||
|
||||
// Every autosnapshot is an idle-kind row.
|
||||
expect(snaps.every((s) => s.kind === 'idle')).toBe(true);
|
||||
// Periodic: at least one per cap over a multi-cap burst (lower-bounded loosely
|
||||
// to stay robust; the property is "fires at least every cap", not a single
|
||||
// trailing snapshot).
|
||||
expect(snaps.length).toBeGreaterThanOrEqual(3);
|
||||
// But NOT one-per-store: ~`stores` stores were issued; the debounce must
|
||||
// collapse them to a small multiple of the cap count, nowhere near per-store.
|
||||
expect(snaps.length).toBeLessThanOrEqual(Math.ceil(stores / 2));
|
||||
});
|
||||
|
||||
it('intermittent burst (a few stores, then quiet) → exactly ONE trailing snapshot', async () => {
|
||||
const pageId = randomUUID();
|
||||
|
||||
// A short burst well within a single cap window, then silence.
|
||||
await arm(pageId);
|
||||
await sleep(40);
|
||||
await arm(pageId);
|
||||
await sleep(40);
|
||||
await arm(pageId);
|
||||
|
||||
// Wait comfortably past the cap so the single pending trailing job fires.
|
||||
await sleep(4 * MAX_WAIT_MS);
|
||||
|
||||
const snaps = processed.filter((p) => p.pageId === pageId);
|
||||
expect(snaps).toHaveLength(1);
|
||||
expect(snaps[0].kind).toBe('idle');
|
||||
});
|
||||
});
|
||||
@@ -387,7 +387,7 @@ bucketByDay(sessions, tz):
|
||||
факты ниже — ground truth, можно дозапросить файлы через gitea MCP по указанному SHA):
|
||||
|
||||
- `page_history.kind` — `varchar(20)`, NULLABLE, БЕЗ дефолта (migration
|
||||
`20260705T120000-page-history-kind.ts`). Домен: `manual`/`agent`/`idle`/`boundary`;
|
||||
`20260707T120000-page-history-kind.ts`). Домен: `manual`/`agent`/`idle`/`boundary`;
|
||||
legacy `null` = автосейв (`collaboration/constants.ts`, `PageHistoryKind`).
|
||||
- `kind` УЖЕ включён в `PageHistoryRepo.baseFields` (`page-history.repo.ts`) — читается всеми
|
||||
выборками истории. `saveHistory({kind})` и `updateHistoryKind(id, kind)` существуют.
|
||||
|
||||
@@ -178,10 +178,11 @@ function sliceModel(xml: string): string | null {
|
||||
// --- decode chain ----------------------------------------------------------
|
||||
|
||||
/**
|
||||
* Read the `content=` attribute out of a `.drawio.svg` string. Docmost stores a
|
||||
* base64 payload there (createDrawioSvg); draw.io's own SVG export may store the
|
||||
* XML entity-encoded instead. The DOM decodes entities for us, so the caller
|
||||
* only has to distinguish "starts with '<'" (raw XML) from base64.
|
||||
* Read the `content=` attribute out of a `.drawio.svg` string. Docmost writes
|
||||
* the mxfile XML entity-encoded there (buildDrawioSvg / createDrawioSvg), which
|
||||
* is also how draw.io's own SVG export stores it; older attachments stored a
|
||||
* base64 payload instead. The DOM decodes entities for us, so the caller only
|
||||
* has to distinguish "starts with '<'" (raw XML) from base64.
|
||||
*/
|
||||
export function extractContentAttr(svg: string): string {
|
||||
const { doc, error } = parseXml(svg);
|
||||
@@ -195,12 +196,23 @@ export function extractContentAttr(svg: string): string {
|
||||
// value itself never contains a double-quote (base64 / entity-encoded XML).
|
||||
const m = /content="([^"]*)"/.exec(svg);
|
||||
if (m) {
|
||||
// Decode the handful of XML entities a raw regex would leave encoded.
|
||||
// Decode the handful of XML entities a raw regex would leave encoded. The
|
||||
// numeric char-refs for tab/newline/CR MUST be decoded here too: the DOM
|
||||
// path above turns them back into the literal control chars, so this
|
||||
// regex fallback has to agree or the two decode paths diverge (#507).
|
||||
// `&` is decoded last so an escaped `&#x9;` reads back as the
|
||||
// literal text `	`, not a tab.
|
||||
return m[1]
|
||||
.replace(/</g, "<")
|
||||
.replace(/>/g, ">")
|
||||
.replace(/"/g, '"')
|
||||
.replace(/'/g, "'")
|
||||
.replace(/	/gi, "\t")
|
||||
.replace(/	/g, "\t")
|
||||
.replace(/
/gi, "\n")
|
||||
.replace(/ /g, "\n")
|
||||
.replace(/
/gi, "\r")
|
||||
.replace(/ /g, "\r")
|
||||
.replace(/&/g, "&");
|
||||
}
|
||||
throw new Error("drawio: SVG has no content= attribute to decode");
|
||||
@@ -307,9 +319,16 @@ export function encodeDrawioFile(modelXml: string, title = "Page-1"): string {
|
||||
/**
|
||||
* Build the `diagram.drawio.svg` attachment. Mirrors the import service's
|
||||
* createDrawioSvg contract exactly:
|
||||
* <svg xmlns=… xmlns:xlink=… content="${base64(drawioFile)}">${inner}</svg>
|
||||
* <svg xmlns=… xmlns:xlink=… content="${xmlEscape(drawioFile)}">${inner}</svg>
|
||||
* plus width/height/viewBox from the diagram bounding box and the schematic
|
||||
* preview as the visible children (`inner`).
|
||||
*
|
||||
* The `content=` value is the mxfile XML XML-entity-escaped (draw.io's own
|
||||
* native form), NOT base64. draw.io's editor decodes a base64 content= via
|
||||
* Latin-1 atob (no UTF-8 step), turning every non-ASCII char (e.g. Cyrillic,
|
||||
* ё, —) into mojibake; the entity-encoded form is decoded by the DOM as UTF-8
|
||||
* and opens intact. Our decoder (decodeDrawioSvg) reads both forms, so old
|
||||
* base64 attachments still round-trip.
|
||||
*/
|
||||
export function buildDrawioSvg(
|
||||
modelXml: string,
|
||||
@@ -318,14 +337,14 @@ export function buildDrawioSvg(
|
||||
title = "Page-1",
|
||||
): string {
|
||||
const file = encodeDrawioFile(modelXml, title);
|
||||
const base64 = Buffer.from(file, "utf-8").toString("base64");
|
||||
const content = xmlEscape(file);
|
||||
const w = Math.max(1, Math.round(bbox.width));
|
||||
const h = Math.max(1, Math.round(bbox.height));
|
||||
return (
|
||||
`<svg xmlns="http://www.w3.org/2000/svg" ` +
|
||||
`xmlns:xlink="http://www.w3.org/1999/xlink" ` +
|
||||
`width="${w}" height="${h}" viewBox="0 0 ${w} ${h}" ` +
|
||||
`content="${base64}">${inner}</svg>`
|
||||
`content="${content}">${inner}</svg>`
|
||||
);
|
||||
}
|
||||
|
||||
@@ -334,7 +353,15 @@ function xmlEscape(s: string): string {
|
||||
.replace(/&/g, "&")
|
||||
.replace(/</g, "<")
|
||||
.replace(/>/g, ">")
|
||||
.replace(/"/g, """);
|
||||
.replace(/"/g, """)
|
||||
// A literal tab/newline/CR inside an attribute value is collapsed to a
|
||||
// single space by XML attribute-value normalization on DOM read (both jsdom
|
||||
// here and the real draw.io editor), silently flattening multi-line labels
|
||||
// and tab-bearing values. Numeric char-refs survive that normalization, so
|
||||
// emit them the way draw.io's own native export does (#507).
|
||||
.replace(/\t/g, "	")
|
||||
.replace(/\n/g, "
")
|
||||
.replace(/\r/g, "
");
|
||||
}
|
||||
|
||||
// --- normalization + hash --------------------------------------------------
|
||||
|
||||
@@ -51,6 +51,14 @@ const hasRule = (issues, rule, cellId) =>
|
||||
(i) => i.rule === rule && (cellId === undefined || i.cellId === cellId),
|
||||
);
|
||||
|
||||
// Reverse the attribute-value XML escaping used in the `content=` payload.
|
||||
const unescapeAttr = (s) =>
|
||||
s
|
||||
.replace(/</g, "<")
|
||||
.replace(/>/g, ">")
|
||||
.replace(/"/g, '"')
|
||||
.replace(/&/g, "&");
|
||||
|
||||
// --- style parsing ---------------------------------------------------------
|
||||
|
||||
test("parseStyle: base stylename + key=value pairs", () => {
|
||||
@@ -335,16 +343,20 @@ test("encode/build: a title with < > \" & round-trips without corrupting the SVG
|
||||
|
||||
// The outer content="..." attribute must not be broken by the title: the raw
|
||||
// title metacharacters never appear literally in the SVG markup (they are
|
||||
// base64-encoded inside content=, and escaped inside the file XML).
|
||||
// entity-escaped inside content=, doubly so where the file XML already escaped
|
||||
// them inside name="...").
|
||||
const contentMatch = /content="([^"]*)"/.exec(svg);
|
||||
assert.ok(contentMatch, "SVG has a single well-formed content= attribute");
|
||||
|
||||
// The diagram model still decodes losslessly despite the exotic title.
|
||||
assert.equal(decodeDrawioSvg(svg), model);
|
||||
|
||||
// content= is now entity-encoded XML (draw.io's native form), never base64.
|
||||
assert.match(contentMatch[1], /^<mxfile/);
|
||||
|
||||
// The file XML is well-formed: the title lives in name="..." as escaped
|
||||
// entities, so unescaping recovers the original title byte-for-byte.
|
||||
const fileXml = Buffer.from(contentMatch[1], "base64").toString("utf-8");
|
||||
const fileXml = unescapeAttr(contentMatch[1]);
|
||||
const nameMatch = /<diagram id="[^"]*" name="([^"]*)">/.exec(fileXml);
|
||||
assert.ok(nameMatch, "the diagram name attribute is intact and quote-safe");
|
||||
const decodedTitle = nameMatch[1]
|
||||
@@ -358,3 +370,112 @@ test("encode/build: a title with < > \" & round-trips without corrupting the SVG
|
||||
const file = encodeDrawioFile(model, title);
|
||||
assert.match(file, /name="A < B > C " D & E">/);
|
||||
});
|
||||
|
||||
// --- #507: content= is entity-encoded XML, never base64 --------------------
|
||||
|
||||
// A model whose cell values carry Cyrillic, ё and an em dash — exactly the
|
||||
// characters draw.io's Latin-1 atob mangles when content= is base64.
|
||||
const CYRILLIC_MODEL =
|
||||
"<mxGraphModel><root>" +
|
||||
'<mxCell id="0"/><mxCell id="1" parent="0"/>' +
|
||||
'<mxCell id="2" value="Старт-бит — ёж" style="rounded=1;html=1;" vertex="1" parent="1">' +
|
||||
'<mxGeometry x="100" y="100" width="120" height="60" as="geometry"/></mxCell>' +
|
||||
"</root></mxGraphModel>";
|
||||
|
||||
test("#507: buildDrawioSvg writes content= as entity-encoded XML (not base64)", () => {
|
||||
const model = normalizeXml(CYRILLIC_MODEL);
|
||||
const svg = buildDrawioSvg(model, "<g/>", { width: 200, height: 120 }, "Диаграмма");
|
||||
|
||||
const contentMatch = /content="([^"]*)"/.exec(svg);
|
||||
assert.ok(contentMatch, "SVG has a single well-formed content= attribute");
|
||||
const content = contentMatch[1];
|
||||
|
||||
// The content= value is the entity-encoded mxfile XML — starts with `<mxfile`.
|
||||
assert.match(content, /^<mxfile/, "content= is entity-encoded mxfile XML");
|
||||
// It must NOT be a base64 blob: base64 has no XML entities and no literal `<`.
|
||||
assert.ok(content.includes("<"), "content= carries XML entities, not base64");
|
||||
|
||||
// Cyrillic / ё / — survive verbatim in the attribute (raw UTF-8, not atob-mangled).
|
||||
assert.ok(content.includes("Старт-бит — ёж"), "non-ASCII value is raw UTF-8 in content=");
|
||||
assert.ok(content.includes("Диаграмма"), "non-ASCII title is raw UTF-8 in content=");
|
||||
// The mojibake that base64+atob would have produced must be absent.
|
||||
assert.ok(!content.includes("Ð"), "no Latin-1 mojibake in content=");
|
||||
});
|
||||
|
||||
test("#507: Cyrillic model round-trips byte-stable through buildDrawioSvg -> decodeDrawioSvg", () => {
|
||||
const model = normalizeXml(CYRILLIC_MODEL);
|
||||
const svg = buildDrawioSvg(model, "<g/>", { width: 200, height: 120 }, "Заголовок — ё");
|
||||
assert.equal(decodeDrawioSvg(svg), model);
|
||||
});
|
||||
|
||||
test("#507 back-compat: an OLD base64-form .drawio.svg still decodes losslessly", () => {
|
||||
const model = normalizeXml(CYRILLIC_MODEL);
|
||||
// Reproduce the pre-fix write path: encodeDrawioFile -> base64 in content=.
|
||||
const file = encodeDrawioFile(model, "Старая диаграмма");
|
||||
const base64 = Buffer.from(file, "utf-8").toString("base64");
|
||||
const svg =
|
||||
`<svg xmlns="http://www.w3.org/2000/svg" content="${base64}"><g/></svg>`;
|
||||
// No XML entities, purely base64 alphabet — this is the legacy form.
|
||||
assert.ok(!base64.includes("<") && !base64.includes("&"));
|
||||
assert.equal(decodeDrawioSvg(svg), model);
|
||||
});
|
||||
|
||||
test("#507 negative: non-ASCII in a cell value AND in the title round-trip clean", () => {
|
||||
const model = normalizeXml(CYRILLIC_MODEL);
|
||||
const title = "Тест — ёмкость № 5";
|
||||
const svg = buildDrawioSvg(model, "<g/>", { width: 200, height: 120 }, title);
|
||||
|
||||
// Model recovered byte-for-byte.
|
||||
assert.equal(decodeDrawioSvg(svg), model);
|
||||
|
||||
// Title lands in the <diagram name="..."> of the decoded file XML, intact.
|
||||
const content = /content="([^"]*)"/.exec(svg)[1];
|
||||
const fileXml = unescapeAttr(content);
|
||||
const nameMatch = /<diagram id="[^"]*" name="([^"]*)">/.exec(fileXml);
|
||||
assert.ok(nameMatch, "diagram name attribute present");
|
||||
assert.equal(unescapeAttr(nameMatch[1]), title);
|
||||
});
|
||||
|
||||
// --- #507 F1: literal tab/newline/CR in an attribute value survive the
|
||||
// content= round-trip. The whole mxfile XML lives in one content="..." attr;
|
||||
// XML attribute-value normalization collapses a LITERAL tab/newline/CR to a
|
||||
// single space on DOM read, so the escape must emit numeric char-refs instead.
|
||||
const CTRL_MODEL =
|
||||
"<mxGraphModel><root>" +
|
||||
'<mxCell id="0"/><mxCell id="1" parent="0"/>' +
|
||||
'<mxCell id="2" value="col1\tcol2" style="html=1;\nshadow=0" vertex="1" parent="1">' +
|
||||
'<mxGeometry x="10" y="10" width="120" height="60" as="geometry"/></mxCell>' +
|
||||
'<mxCell id="3" value="Line1\nLine2\rLine3" vertex="1" parent="1">' +
|
||||
'<mxGeometry x="10" y="100" width="120" height="60" as="geometry"/></mxCell>' +
|
||||
"</root></mxGraphModel>";
|
||||
|
||||
test("#507 F1: literal tab/newline/CR in a value round-trip byte-stable (DOM decode)", () => {
|
||||
const svg = buildDrawioSvg(CTRL_MODEL, "<g/>", { width: 200, height: 200 });
|
||||
const content = /content="([^"]*)"/.exec(svg)[1];
|
||||
// The tab/newline/CR must be emitted as numeric char-refs, never as literal
|
||||
// control chars (which DOM attribute-value normalization would eat).
|
||||
assert.ok(
|
||||
!/[\t\n\r]/.test(content),
|
||||
"no literal tab/newline/CR survive in the content= attribute",
|
||||
);
|
||||
assert.ok(
|
||||
content.includes("	") &&
|
||||
content.includes("
") &&
|
||||
content.includes("
"),
|
||||
"tab/newline/CR are emitted as numeric char-refs",
|
||||
);
|
||||
// Full DOM decode recovers the model byte-for-byte, control chars intact.
|
||||
assert.equal(decodeDrawioSvg(svg), CTRL_MODEL);
|
||||
});
|
||||
|
||||
test("#507 F1: regex fallback decodes tab/newline/CR char-refs (agrees with DOM path)", () => {
|
||||
const svg = buildDrawioSvg(CTRL_MODEL, "<g/>", { width: 200, height: 200 });
|
||||
const content = /content="([^"]*)"/.exec(svg)[1];
|
||||
// A bare `&` makes the SVG wrapper malformed, forcing extractContentAttr onto
|
||||
// its regex fallback branch. That branch must decode the tab/newline/CR
|
||||
// char-refs exactly like the DOM path, or the two decoders diverge.
|
||||
const malformedSvg = `<svg content="${content}">&</svg>`;
|
||||
assert.equal(decodeDrawioSvg(malformedSvg), CTRL_MODEL);
|
||||
// The well-formed (DOM) path yields the identical result.
|
||||
assert.equal(decodeDrawioSvg(svg), CTRL_MODEL);
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user