From 9c1def15de797722737c8ba9da44429175014f65 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 3 Jul 2026 11:34:57 +0800 Subject: [PATCH 01/51] feat: add page citation asset support --- .gitignore | 1 + eslint.config.mjs | 1 + next.config.ts | 3 + package.json | 7 +- pnpm-lock.yaml | 672 +++++++++++++++++- pnpm-workspace.yaml | 4 + scripts/prepare-effect.sh | 13 + src/agent-harness/ledger.ts | 8 + src/agent-harness/types.ts | 4 + src/components/chat-message-list.test.ts | 50 ++ src/components/chat-message-list.tsx | 74 +- src/domains/chat/chat-citation-persistence.ts | 1 + src/domains/chat/citations.test.ts | 19 +- src/domains/chat/citations.ts | 6 +- src/domains/chat/index.ts | 28 +- src/domains/chat/page-citation-assets.test.ts | 216 ++++++ src/domains/chat/page-citation-assets.ts | 381 ++++++++++ src/domains/chat/types.ts | 2 +- src/domains/chat/view.ts | 4 + .../sources/page-citation-assets.test.ts | 82 +++ src/domains/sources/page-citation-assets.ts | 69 ++ .../source-reconcile-route-workflow.test.ts | 78 ++ .../source-reconcile-route-workflow.ts | 98 +++ .../sources/source-reconcile-workflow.test.ts | 32 + .../sources/source-reconcile-workflow.ts | 33 + src/domains/workspace/integration.test.ts | 2 + src/domains/workspace/persistence.test.ts | 2 + src/integrations/knowhere-sdk-storage.test.ts | 320 +++++++++ src/integrations/knowhere-sdk-storage.ts | 294 ++++++++ 29 files changed, 2440 insertions(+), 64 deletions(-) create mode 100755 scripts/prepare-effect.sh create mode 100644 src/domains/chat/page-citation-assets.test.ts create mode 100644 src/domains/chat/page-citation-assets.ts create mode 100644 src/domains/sources/page-citation-assets.test.ts create mode 100644 src/domains/sources/page-citation-assets.ts create mode 100644 src/integrations/knowhere-sdk-storage.test.ts create mode 100644 src/integrations/knowhere-sdk-storage.ts diff --git a/.gitignore b/.gitignore index 74889db..2a2a9f1 100644 --- a/.gitignore +++ b/.gitignore @@ -25,6 +25,7 @@ # misc .DS_Store *.pem +.repos/effect # debug npm-debug.log* diff --git a/eslint.config.mjs b/eslint.config.mjs index ac9d10f..6ed5798 100644 --- a/eslint.config.mjs +++ b/eslint.config.mjs @@ -13,6 +13,7 @@ const eslintConfig = defineConfig([ "build/**", "playwright-report/**", "test-results/**", + ".repos/effect/**", "next-env.d.ts", ]), ]); diff --git a/next.config.ts b/next.config.ts index e294c70..b73f520 100644 --- a/next.config.ts +++ b/next.config.ts @@ -7,6 +7,9 @@ const nextConfig: NextConfig = { "pg", "@neondatabase/serverless", "postgres", + "@ontos-ai/knowhere-sdk", + "@napi-rs/canvas", + "piscina", ], allowedDevOrigins: [ "127.0.0.1", diff --git a/package.json b/package.json index 73b1e26..c998de4 100644 --- a/package.json +++ b/package.json @@ -2,11 +2,14 @@ "name": "knowhere-notebook", "version": "0.1.0", "private": true, + "engines": { + "node": ">=22.13.0" + }, "scripts": { "dev": "pnpm run upstash:dev & next dev", "build": "next build", "typecheck": "tsc --noEmit", - "prepare": "effect-language-service patch", + "prepare": "./scripts/prepare-effect.sh && effect-language-service patch", "start": "next start", "lint": "eslint", "test": "vitest run", @@ -23,6 +26,7 @@ "@ai-sdk/react": "^3.0.177", "@antv/chart-visualization-skills": "0.1.3", "@effect/platform": "^0.96.1", + "@napi-rs/canvas": "^1.0.2", "@neondatabase/serverless": "^1.1.0", "@ontos-ai/knowhere-sdk": "^2.0.0", "@radix-ui/react-alert-dialog": "^1.1.15", @@ -51,6 +55,7 @@ "next": "16.2.4", "next-themes": "^0.4.6", "pdfjs-dist": "5.4.296", + "piscina": "^5.2.0", "postgres": "^3.4.9", "posthog-js": "^1.386.8", "react": "19.2.4", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 1c4223f..978e912 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -4,6 +4,9 @@ settings: autoInstallPeers: true excludeLinksFromLockfile: false +overrides: + '@ontos-ai/knowhere-sdk': file:../knowhere-node-sdk + importers: .: @@ -17,12 +20,15 @@ importers: '@effect/platform': specifier: ^0.96.1 version: 0.96.1(effect@3.21.2) + '@napi-rs/canvas': + specifier: ^1.0.2 + version: 1.0.2 '@neondatabase/serverless': specifier: ^1.1.0 version: 1.1.0 '@ontos-ai/knowhere-sdk': - specifier: ^2.0.0 - version: 2.0.0 + specifier: file:../knowhere-node-sdk + version: file:../knowhere-node-sdk(ioredis@5.11.1) '@radix-ui/react-alert-dialog': specifier: ^1.1.15 version: 1.1.15(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4) @@ -101,6 +107,9 @@ importers: pdfjs-dist: specifier: 5.4.296 version: 5.4.296 + piscina: + specifier: ^5.2.0 + version: 5.2.0 postgres: specifier: ^3.4.9 version: 3.4.9 @@ -167,7 +176,7 @@ importers: version: 19.2.3(@types/react@19.2.14) '@vitejs/plugin-react': specifier: ^6.0.1 - version: 6.0.1(babel-plugin-react-compiler@1.0.0)(vite@8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0)) + version: 6.0.1(babel-plugin-react-compiler@1.0.0)(vite@8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0)(yaml@2.9.0)) babel-plugin-react-compiler: specifier: ^1.0.0 version: 1.0.0 @@ -194,7 +203,7 @@ importers: version: 6.0.3 vitest: specifier: ^4.1.5 - version: 4.1.5(@opentelemetry/api@1.9.0)(@types/node@20.19.39)(jsdom@29.1.1(@noble/hashes@1.8.0))(msw@2.14.3(@types/node@20.19.39)(typescript@6.0.3))(vite@8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0)) + version: 4.1.5(@opentelemetry/api@1.9.0)(@types/node@20.19.39)(jsdom@29.1.1(@noble/hashes@1.8.0))(msw@2.14.3(@types/node@20.19.39)(typescript@6.0.3))(vite@8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0)(yaml@2.9.0)) packages: @@ -433,6 +442,19 @@ packages: resolution: {integrity: sha512-EXnJjIy6zQ3nUO/MZ+ynWUb8B895KZPotd1++oTs9JjDkplwM7cb6zo8Zq2zU6piwq+KflO7amXbEfj1UMpHkw==} hasBin: true + '@effect/platform-node-shared@4.0.0-beta.93': + resolution: {integrity: sha512-XUqZ2u5GglBqY8q2jj4Q7GjN5K/enedk8auZM9rY/l5a/myaQTrQp3QnvpIK4/Yg0WFjLGuctGPMKWRk3OLIrA==} + engines: {node: '>=18.0.0'} + peerDependencies: + effect: ^4.0.0-beta.93 + + '@effect/platform-node@4.0.0-beta.93': + resolution: {integrity: sha512-QagsCGR0ZOXaCQqS5qGR2mcDng4LiP2bYhiiX1D6UC8cT9vsusVVOHiJWn8CupeDx+yVnPcu81QmA/SDt6GM1w==} + engines: {node: '>=18.0.0'} + peerDependencies: + effect: ^4.0.0-beta.93 + ioredis: ^5.7.0 + '@effect/platform@0.96.1': resolution: {integrity: sha512-cjB1QZZYEP8JXCFNGvBLVi0T6YUBQTmOVEUA3SDbiQ6RUO+p6CE3eyD2vMWmrz5nE8yY5QSAuOV9v0boEcUv+A==} peerDependencies: @@ -1175,6 +1197,9 @@ packages: '@types/node': optional: true + '@ioredis/commands@1.10.0': + resolution: {integrity: sha512-UmeW7z4LfctwoQ5wkhVzgq8tXkreED2xZGpX+Bg+zA+WJFZCT6c062AfCK/Dfk81xZnnwdhJCUMkitihRaoC2Q==} + '@jridgewell/gen-mapping@0.3.13': resolution: {integrity: sha512-2kkt/7niJ6MgEPxF0bYdQ6etZaA+fQvDcLKckhy1yIQOzaoKjBBjSj63/aLVjYE3qhRt5dvM+uUyfCg6UKCBbA==} @@ -1206,31 +1231,61 @@ packages: cpu: [arm64] os: [darwin] + '@msgpackr-extract/msgpackr-extract-darwin-arm64@3.0.4': + resolution: {integrity: sha512-LCkGo6JDfaBhgST7UpPWgNgLINpcpabaHfyz5OBx75nUYxBsaEPxjnyNjWpeb/xBup/682QnBfRBy2/LvPutZQ==} + cpu: [arm64] + os: [darwin] + '@msgpackr-extract/msgpackr-extract-darwin-x64@3.0.3': resolution: {integrity: sha512-mdzd3AVzYKuUmiWOQ8GNhl64/IoFGol569zNRdkLReh6LRLHOXxU4U8eq0JwaD8iFHdVGqSy4IjFL4reoWCDFw==} cpu: [x64] os: [darwin] + '@msgpackr-extract/msgpackr-extract-darwin-x64@3.0.4': + resolution: {integrity: sha512-zExlW9zUJKZH/tOtVMttwjKa4Xm/3KcNjnE3dPN92uCktwavMxpgCA3MoJK/DOnTWsQgo224OaST27/mPNAf+w==} + cpu: [x64] + os: [darwin] + '@msgpackr-extract/msgpackr-extract-linux-arm64@3.0.3': resolution: {integrity: sha512-YxQL+ax0XqBJDZiKimS2XQaf+2wDGVa1enVRGzEvLLVFeqa5kx2bWbtcSXgsxjQB7nRqqIGFIcLteF/sHeVtQg==} cpu: [arm64] os: [linux] + '@msgpackr-extract/msgpackr-extract-linux-arm64@3.0.4': + resolution: {integrity: sha512-dgX0P/9wGPJeHFBG+ZmhgE6bmtMt7NP5CRBGyyktpopdk/mW4POnrpQsSLtKI1dwpc+pPLuXHDh6vvskyQE/sw==} + cpu: [arm64] + os: [linux] + '@msgpackr-extract/msgpackr-extract-linux-arm@3.0.3': resolution: {integrity: sha512-fg0uy/dG/nZEXfYilKoRe7yALaNmHoYeIoJuJ7KJ+YyU2bvY8vPv27f7UKhGRpY6euFYqEVhxCFZgAUNQBM3nw==} cpu: [arm] os: [linux] + '@msgpackr-extract/msgpackr-extract-linux-arm@3.0.4': + resolution: {integrity: sha512-Tg3yX65f5GbtXLkrYEHE5oibZG9epyYWas7FogTTEJeDEF9JlXJzKgXaNhT3UXlTOeA+AfZpYZYZ0uPj7Cfquw==} + cpu: [arm] + os: [linux] + '@msgpackr-extract/msgpackr-extract-linux-x64@3.0.3': resolution: {integrity: sha512-cvwNfbP07pKUfq1uH+S6KJ7dT9K8WOE4ZiAcsrSes+UY55E/0jLYc+vq+DO7jlmqRb5zAggExKm0H7O/CBaesg==} cpu: [x64] os: [linux] + '@msgpackr-extract/msgpackr-extract-linux-x64@3.0.4': + resolution: {integrity: sha512-8TNXMEjJc3QEy7R/x1INhgiU+XakDAFUzBhaz7+Rbrs8NH5UQeHQxxmzsSBJGyV6I1jW79undiQm8tOI+D+8FQ==} + cpu: [x64] + os: [linux] + '@msgpackr-extract/msgpackr-extract-win32-x64@3.0.3': resolution: {integrity: sha512-x0fWaQtYp4E6sktbsdAqnehxDgEc/VwM7uLsRCYWaiGu0ykYdZPiS8zCWdnjHwyiumousxfBm4SO31eXqwEZhQ==} cpu: [x64] os: [win32] + '@msgpackr-extract/msgpackr-extract-win32-x64@3.0.4': + resolution: {integrity: sha512-CmCXPQrkbwExx3j946/PtHWHbYJiCRBRDl4BlkRQcJB/YOwQxJRTpoo7aTsortjgoJ1x7opzTSxn7C+ASSLVjQ==} + cpu: [x64] + os: [win32] + '@mswjs/interceptors@0.41.8': resolution: {integrity: sha512-pRLMNKTSGRoLq+KnEB/7OY5vijw1XmcheAAOiv6pj7W1FG32kAGqj1C/RK/cqxRGr1Fh+zBi8sDur8kj3EQv6A==} engines: {node: '>=18'} @@ -1241,24 +1296,48 @@ packages: cpu: [arm64] os: [android] + '@napi-rs/canvas-android-arm64@1.0.2': + resolution: {integrity: sha512-IMXKVQod0ol4vt3gmClUfXz4JAgHYESGPCUqmH3lQxBoL0K/2greJaQE1HVBVxWWFKfLc4OLZVdxg7kXVyXv+g==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [android] + '@napi-rs/canvas-darwin-arm64@0.1.100': resolution: {integrity: sha512-2PcswRaC7Ly645DGt88///zuFDhJxJYdKAs1uU3mfk1atYkXufgcgLfBpk6Tm12nCQBaNt1wpybuPZ4qOhTo8A==} engines: {node: '>= 10'} cpu: [arm64] os: [darwin] + '@napi-rs/canvas-darwin-arm64@1.0.2': + resolution: {integrity: sha512-Sc8tPi6cF+5lqOzCCKFALJHhDiRwyMzTPYm3bbhdXsOunU0lQO5f05ucyOzN2r55I23Hg5bsjH63uSCvWp3EgQ==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [darwin] + '@napi-rs/canvas-darwin-x64@0.1.100': resolution: {integrity: sha512-ePNZtj7pNIva/siZMg+HmbeozkIjqUIYdoymH8HaA3qK7LfzFN4WMBM8G6HQ9ZC+H3+Dnn5pqtiXpgLykaPOhw==} engines: {node: '>= 10'} cpu: [x64] os: [darwin] + '@napi-rs/canvas-darwin-x64@1.0.2': + resolution: {integrity: sha512-niDXZ9LhKB1zLrUdYB64RHQFDGz9rr0eGx061qtJJU3U20EMMIx28ADF5fVYbhtOgkWQrBjFicfaye1yM0U62A==} + engines: {node: '>= 10'} + cpu: [x64] + os: [darwin] + '@napi-rs/canvas-linux-arm-gnueabihf@0.1.100': resolution: {integrity: sha512-d5cDB48oWFGU8/XPhUOFAlySgb/VAu7D+s8fi55K1Pcfg8aPplHWqMgibhVLU8ky7Pyg/fuiVLz4Nf3JrSTuUA==} engines: {node: '>= 10'} cpu: [arm] os: [linux] + '@napi-rs/canvas-linux-arm-gnueabihf@1.0.2': + resolution: {integrity: sha512-sgatQL9JxGRH/Amzcvu0P3t8Am3duou74CisfuJ41Dwt8cWy723z/9KZ8LlgmxfypEwEZxSTNFJtU8d281lmhQ==} + engines: {node: '>= 10'} + cpu: [arm] + os: [linux] + '@napi-rs/canvas-linux-arm64-gnu@0.1.100': resolution: {integrity: sha512-rDxgxRu69RvDlX/bh9o22DxLsGr8EqsNgotL9+RwQE1S0b0cqeatqsw6aW45mukm0B42DIAaAacKaYQ8cqS1nw==} engines: {node: '>= 10'} @@ -1266,6 +1345,13 @@ packages: os: [linux] libc: [glibc] + '@napi-rs/canvas-linux-arm64-gnu@1.0.2': + resolution: {integrity: sha512-dgKuX0peF3xwY6ZF5QxGS4wbfDqpoFAJYXiLSp+guZKARQUKMkRqZSDrXKj7nfrec3UCMzC0PFCPte0ES98AiA==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [linux] + libc: [glibc] + '@napi-rs/canvas-linux-arm64-musl@0.1.100': resolution: {integrity: sha512-K3mDW66N+xT2/V439u1alFANiBUjdEx2gLiNYnCmUsva5jZMxWTjafBYwTzYK+EMFMHrUoabuU+T1BIP5CgbYQ==} engines: {node: '>= 10'} @@ -1273,6 +1359,13 @@ packages: os: [linux] libc: [musl] + '@napi-rs/canvas-linux-arm64-musl@1.0.2': + resolution: {integrity: sha512-qwROoDIC9upfvDoRLuPn2aNg9CGW1x0Ygr4k2Or+8paA9d0qBLwk87U+g8KQpoOviKoPoiwl97kvBYuYD7qZoA==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [linux] + libc: [musl] + '@napi-rs/canvas-linux-riscv64-gnu@0.1.100': resolution: {integrity: sha512-mooqUBTIsccZpnoQC4NgrC1v6C1vof39etLNMnBwCY+p0gajWJvAHLGQ6g/gGyS5YrpDW+GefSN4+Cvcr08UWw==} engines: {node: '>= 10'} @@ -1280,6 +1373,13 @@ packages: os: [linux] libc: [glibc] + '@napi-rs/canvas-linux-riscv64-gnu@1.0.2': + resolution: {integrity: sha512-fXRjnPihdnbO6qy1QQOgxAonb68A0TCEG7rj1x7v7rxNElsE8EVIKIEUTvyDtU+sthYSbX+8e7g3oZiLGnOmxw==} + engines: {node: '>= 10'} + cpu: [riscv64] + os: [linux] + libc: [glibc] + '@napi-rs/canvas-linux-x64-gnu@0.1.100': resolution: {integrity: sha512-1eCvkDCazm7FFhsT7DfGOdSaHgZVK3bt/dSBl5EWHOWmnz+I7j8tPseJqqD81NF+MH21jKUK4wQSDjN0mdhnTg==} engines: {node: '>= 10'} @@ -1287,6 +1387,13 @@ packages: os: [linux] libc: [glibc] + '@napi-rs/canvas-linux-x64-gnu@1.0.2': + resolution: {integrity: sha512-nPR97DXhbWIAy7yazF3jc06kEPMqYMLmPzFOVNlwKPfIoSChnI+x7dc0hTLaihz3jxrjL6j4BbA7earxfx4X3g==} + engines: {node: '>= 10'} + cpu: [x64] + os: [linux] + libc: [glibc] + '@napi-rs/canvas-linux-x64-musl@0.1.100': resolution: {integrity: sha512-20arT6lnI19S68qNlii73TSEDbECNgzMz2EpldC1V3mZFuRkeujXkcebRk0LRJe9SEUAooYiLokfMViY8IX7yA==} engines: {node: '>= 10'} @@ -1294,22 +1401,158 @@ packages: os: [linux] libc: [musl] + '@napi-rs/canvas-linux-x64-musl@1.0.2': + resolution: {integrity: sha512-l7zZY5+jL5qnBZtDz7CoBtY6p7EkHu422g/0zWwrOrzIwWyWxZFRfZZORY1UG7YApymPLx+UbOkN206xXn/c1Q==} + engines: {node: '>= 10'} + cpu: [x64] + os: [linux] + libc: [musl] + '@napi-rs/canvas-win32-arm64-msvc@0.1.100': resolution: {integrity: sha512-DZFFT1wIAg37LJw37yhMRFfjATd3vTQzjZ1Yki8u2vhO6Hi5VE6BVaGQ1aaDu7xb4iMErz+9EOwjpS7xcxFeBw==} engines: {node: '>= 10'} cpu: [arm64] os: [win32] + '@napi-rs/canvas-win32-arm64-msvc@1.0.2': + resolution: {integrity: sha512-yE0koHCFF4PIbMc2o2SEALhnipz7WBISh5glLvQiomtIoCcW0np3H4Lw93ceJAfJttTTeIIWFbwH84F7EVzjMQ==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [win32] + '@napi-rs/canvas-win32-x64-msvc@0.1.100': resolution: {integrity: sha512-MyT1j3mHC2+Lu4pBi9mKyMJhtP6U7k7EldY7sj/uS5gJA65gTXt8MefJQXLJo5d/vZbuWmfxzkEUNc/urV3pHA==} engines: {node: '>= 10'} cpu: [x64] os: [win32] + '@napi-rs/canvas-win32-x64-msvc@1.0.2': + resolution: {integrity: sha512-okU8/t2foV6C31n0GtvEMbfD5rOFc70+/6xUNME9Guld29sgSOIGUEDScAWFlcP3k5TYQRl9TNkwJEEjh15w8A==} + engines: {node: '>= 10'} + cpu: [x64] + os: [win32] + '@napi-rs/canvas@0.1.100': resolution: {integrity: sha512-xglYA6q3XO5P3BNJYxVZ1IV7DLVjp1Py6nwag88YntrS+3vKHyYcMqXVS4ZztJmwz2uGvz1FWhI/4LgbR5uQDA==} engines: {node: '>= 10'} + '@napi-rs/canvas@1.0.2': + resolution: {integrity: sha512-EYEqlMYaCbpZDz+IgDH5xp9MTd3ui4dmGqbQYryhMLnSRxrhHKq5KQWHHKxFUcEP4Hp8/BWgvqXocX4j7iSbOQ==} + engines: {node: '>= 10'} + + '@napi-rs/nice-android-arm-eabi@1.1.1': + resolution: {integrity: sha512-kjirL3N6TnRPv5iuHw36wnucNqXAO46dzK9oPb0wj076R5Xm8PfUVA9nAFB5ZNMmfJQJVKACAPd/Z2KYMppthw==} + engines: {node: '>= 10'} + cpu: [arm] + os: [android] + + '@napi-rs/nice-android-arm64@1.1.1': + resolution: {integrity: sha512-blG0i7dXgbInN5urONoUCNf+DUEAavRffrO7fZSeoRMJc5qD+BJeNcpr54msPF6qfDD6kzs9AQJogZvT2KD5nw==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [android] + + '@napi-rs/nice-darwin-arm64@1.1.1': + resolution: {integrity: sha512-s/E7w45NaLqTGuOjC2p96pct4jRfo61xb9bU1unM/MJ/RFkKlJyJDx7OJI/O0ll/hrfpqKopuAFDV8yo0hfT7A==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [darwin] + + '@napi-rs/nice-darwin-x64@1.1.1': + resolution: {integrity: sha512-dGoEBnVpsdcC+oHHmW1LRK5eiyzLwdgNQq3BmZIav+9/5WTZwBYX7r5ZkQC07Nxd3KHOCkgbHSh4wPkH1N1LiQ==} + engines: {node: '>= 10'} + cpu: [x64] + os: [darwin] + + '@napi-rs/nice-freebsd-x64@1.1.1': + resolution: {integrity: sha512-kHv4kEHAylMYmlNwcQcDtXjklYp4FCf0b05E+0h6nDHsZ+F0bDe04U/tXNOqrx5CmIAth4vwfkjjUmp4c4JktQ==} + engines: {node: '>= 10'} + cpu: [x64] + os: [freebsd] + + '@napi-rs/nice-linux-arm-gnueabihf@1.1.1': + resolution: {integrity: sha512-E1t7K0efyKXZDoZg1LzCOLxgolxV58HCkaEkEvIYQx12ht2pa8hoBo+4OB3qh7e+QiBlp1SRf+voWUZFxyhyqg==} + engines: {node: '>= 10'} + cpu: [arm] + os: [linux] + + '@napi-rs/nice-linux-arm64-gnu@1.1.1': + resolution: {integrity: sha512-CIKLA12DTIZlmTaaKhQP88R3Xao+gyJxNWEn04wZwC2wmRapNnxCUZkVwggInMJvtVElA+D4ZzOU5sX4jV+SmQ==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [linux] + libc: [glibc] + + '@napi-rs/nice-linux-arm64-musl@1.1.1': + resolution: {integrity: sha512-+2Rzdb3nTIYZ0YJF43qf2twhqOCkiSrHx2Pg6DJaCPYhhaxbLcdlV8hCRMHghQ+EtZQWGNcS2xF4KxBhSGeutg==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [linux] + libc: [musl] + + '@napi-rs/nice-linux-ppc64-gnu@1.1.1': + resolution: {integrity: sha512-4FS8oc0GeHpwvv4tKciKkw3Y4jKsL7FRhaOeiPei0X9T4Jd619wHNe4xCLmN2EMgZoeGg+Q7GY7BsvwKpL22Tg==} + engines: {node: '>= 10'} + cpu: [ppc64] + os: [linux] + libc: [glibc] + + '@napi-rs/nice-linux-riscv64-gnu@1.1.1': + resolution: {integrity: sha512-HU0nw9uD4FO/oGCCk409tCi5IzIZpH2agE6nN4fqpwVlCn5BOq0MS1dXGjXaG17JaAvrlpV5ZeyZwSon10XOXw==} + engines: {node: '>= 10'} + cpu: [riscv64] + os: [linux] + libc: [glibc] + + '@napi-rs/nice-linux-s390x-gnu@1.1.1': + resolution: {integrity: sha512-2YqKJWWl24EwrX0DzCQgPLKQBxYDdBxOHot1KWEq7aY2uYeX+Uvtv4I8xFVVygJDgf6/92h9N3Y43WPx8+PAgQ==} + engines: {node: '>= 10'} + cpu: [s390x] + os: [linux] + libc: [glibc] + + '@napi-rs/nice-linux-x64-gnu@1.1.1': + resolution: {integrity: sha512-/gaNz3R92t+dcrfCw/96pDopcmec7oCcAQ3l/M+Zxr82KT4DljD37CpgrnXV+pJC263JkW572pdbP3hP+KjcIg==} + engines: {node: '>= 10'} + cpu: [x64] + os: [linux] + libc: [glibc] + + '@napi-rs/nice-linux-x64-musl@1.1.1': + resolution: {integrity: sha512-xScCGnyj/oppsNPMnevsBe3pvNaoK7FGvMjT35riz9YdhB2WtTG47ZlbxtOLpjeO9SqqQ2J2igCmz6IJOD5JYw==} + engines: {node: '>= 10'} + cpu: [x64] + os: [linux] + libc: [musl] + + '@napi-rs/nice-openharmony-arm64@1.1.1': + resolution: {integrity: sha512-6uJPRVwVCLDeoOaNyeiW0gp2kFIM4r7PL2MczdZQHkFi9gVlgm+Vn+V6nTWRcu856mJ2WjYJiumEajfSm7arPQ==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [openharmony] + + '@napi-rs/nice-win32-arm64-msvc@1.1.1': + resolution: {integrity: sha512-uoTb4eAvM5B2aj/z8j+Nv8OttPf2m+HVx3UjA5jcFxASvNhQriyCQF1OB1lHL43ZhW+VwZlgvjmP5qF3+59atA==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [win32] + + '@napi-rs/nice-win32-ia32-msvc@1.1.1': + resolution: {integrity: sha512-CNQqlQT9MwuCsg1Vd/oKXiuH+TcsSPJmlAFc5frFyX/KkOh0UpBLEj7aoY656d5UKZQMQFP7vJNa1DNUNORvug==} + engines: {node: '>= 10'} + cpu: [ia32] + os: [win32] + + '@napi-rs/nice-win32-x64-msvc@1.1.1': + resolution: {integrity: sha512-vB+4G/jBQCAh0jelMTY3+kgFy00Hlx2f2/1zjMoH821IbplbWZOkLiTYXQkygNTzQJTq5cvwBDgn2ppHD+bglQ==} + engines: {node: '>= 10'} + cpu: [x64] + os: [win32] + + '@napi-rs/nice@1.1.1': + resolution: {integrity: sha512-xJIPs+bYuc9ASBl+cvGsKbGrJmS6fAKaSZCnT0lhahT5rhA2VVy9/EcIgd2JhtEuFOJNx7UHNn/qiTPTY4nrQw==} + engines: {node: '>= 10'} + '@napi-rs/wasm-runtime@0.2.12': resolution: {integrity: sha512-ZVWUcfwY4E/yPitQJl481FjFo3K22D6qF0DuFH6Y/nbnE11GY5uguDxZMGXPQ8WQ0128MXQD7TnfHyK4oWoIJQ==} @@ -1409,9 +1652,9 @@ packages: resolution: {integrity: sha512-nn5ozdjYQpUCZlWGuxcJY/KpxkWQs4DcbMCmKojjyrYDEAGy4Ce19NN4v5MduafTwJlbKc99UA8YhSVqq9yPZA==} engines: {node: '>=12.4.0'} - '@ontos-ai/knowhere-sdk@2.0.0': - resolution: {integrity: sha512-Zhx8mn/8pYyhBaO9BOdQYsVZybGXyq15YctjHSU8TwgNjJEzX5lEiBEkp0kIGFPh3S+/7X3rZGakefEPf5G2tQ==} - engines: {node: '>=20.19.0', npm: '>=10.0.0', pnpm: '>=9.0.0'} + '@ontos-ai/knowhere-sdk@file:../knowhere-node-sdk': + resolution: {directory: ../knowhere-node-sdk, type: directory} + engines: {node: '>=22.13.0', npm: '>=10.0.0', pnpm: '>=9.0.0'} '@open-draft/deferred-promise@2.2.0': resolution: {integrity: sha512-CecwLWx3rhxVQF6V4bAgPS5t+So2sTbPgAzafKkVizyi7tlwpcFpdFqq+wqF2OwNBmqFuu6tOyouTuxgpMfzmA==} @@ -2174,6 +2417,9 @@ packages: '@types/validate-npm-package-name@4.0.2': resolution: {integrity: sha512-lrpDziQipxCEeK5kWxvljWYhUvOiB2A9izZd9B2AFarYAkqZshb4lPbRs7zKEic6eGtH8V/2qJW+dPp9OtF6bw==} + '@types/ws@8.18.1': + resolution: {integrity: sha512-ThVF6DCVhA8kUGy+aazFQ4kXQ7E1Ty7A3ypFOe0IcJV8O/M511G99AW24irKrW56Wt44yG9+ij8FaqoBGkuBXg==} + '@typescript-eslint/eslint-plugin@8.59.2': resolution: {integrity: sha512-j/bwmkBvHUtPNxzuWe5z6BEk3q54YRyGlBXkSsmfoih7zNrBvl5A9A98anlp/7JbyZcWIJ8KXo/3Tq/DjFLtuQ==} engines: {node: ^18.18.0 || ^20.9.0 || >=21.1.0} @@ -2415,6 +2661,10 @@ packages: engines: {node: '>=0.4.0'} hasBin: true + agent-base@6.0.2: + resolution: {integrity: sha512-RZNwNclF7+MS/8bDg70amg32dyeZGZxiDuQmZxKLAlQjr3jGyLx+4Kkk58UO7D2QdgFIQCovuSuZESne6RG6XQ==} + engines: {node: '>= 6.0.0'} + agent-base@7.1.4: resolution: {integrity: sha512-MnA+YT8fwfJPgBx3m60MNqakm30XOkyIoH1y6huTQvC0PwZG7ki8NacLBcrPbNoo8vEZy7Jpuk7+jMO+CUovTQ==} engines: {node: '>= 14'} @@ -2533,8 +2783,8 @@ packages: resolution: {integrity: sha512-KunSNx+TVpkAw/6ULfhnx+HWRecjqZGTOyquAoWHYLRSdK1tB5Ihce1ZW+UY3fj33bYAFWPu7W/GRSmmrCGuxA==} engines: {node: '>=4'} - axios@1.16.0: - resolution: {integrity: sha512-6hp5CwvTPlN2A31g5dxnwAX0orzM7pmCRDLnZSX772mv8WDqICwFjowHuPs04Mc8deIld1+ejhtaMn5vp6b+1w==} + axios@1.18.1: + resolution: {integrity: sha512-3nTvFlvpn9Zu/RkHUqtc7/+al4UpRW5az71ap5zccp6e8RAYEzhMTecX8Dz1wWDYrPpUoB1HAQEGEAEvUr7S9g==} axobject-query@4.1.0: resolution: {integrity: sha512-qIj0G9wZbMGNLjLmg1PT6v2mE9AH2zlnADJD/2tC6E00hgmhUOfEB6greHPAfLRSufHqROIUTkw6E+M3lH0PTQ==} @@ -2670,6 +2920,10 @@ packages: resolution: {integrity: sha512-eYm0QWBtUrBWZWG0d386OGAw16Z995PiOVo2B7bjWSbHedGl5e0ZWaq65kOGgUSNesEIDkB9ISbTg/JK9dhCZA==} engines: {node: '>=6'} + cluster-key-slot@1.1.1: + resolution: {integrity: sha512-rwHwUfXL40Chm1r08yrhU3qpUvdVlgkKNeyeGPOxnW8/SyVDvgRaed/Uz54AqWNaTCAThlj6QAs3TZcKI0xDEw==} + engines: {node: '>=0.10.0'} + code-block-writer@13.0.3: resolution: {integrity: sha512-Oofo0pq3IKnsFtuHqSF7TqBfr71aeyZDVJ0HpmqB7FBM2qEigL0iPONSCZSO9pE9dZTAxANe5XHG9Uy0YMv8cg==} @@ -2852,6 +3106,10 @@ packages: resolution: {integrity: sha512-ZySD7Nf91aLB0RxL4KGrKHBXl7Eds1DAmEdcoVawXnLD7SDhpNgtuII2aAkg7a7QS41jxPSZ17p4VdGnMHk3MQ==} engines: {node: '>=0.4.0'} + denque@2.1.0: + resolution: {integrity: sha512-HVQE3AAb/pxF8fQAoiqpvg9i3evqug3hoiwakOyZAwJm+6vZehbkYXZ0l4JxS+I3QxM97v5aaRNhj8v5oBhekw==} + engines: {node: '>=0.10'} + depd@2.0.0: resolution: {integrity: sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==} engines: {node: '>= 0.8'} @@ -3007,6 +3265,9 @@ packages: effect@3.21.2: resolution: {integrity: sha512-rXd2FGDM8KdjSIrc+mqEELo7ScW7xTVxEf1iInmPSpIde9/nyGuFM710cjTo7/EreGXiUX2MOonPpprbz2XHCg==} + effect@4.0.0-beta.93: + resolution: {integrity: sha512-wNS5MKFa3C42uBfIDik2oJ78lhpoYz2hN4oBR0229BeeDCIrkg/FiOvoiPGdCVlWa7MEKxEL5I0f8AILVHSD9A==} + electron-to-chromium@1.5.351: resolution: {integrity: sha512-9D7Iqx8RImSvCnOsj86rCH6eQjZFQoM04Jn6HnZVM0Nu/G58/gmKYQ1d12MZTbjQbQSTGI8nwEy07ErsA2slLA==} @@ -3279,6 +3540,10 @@ packages: resolution: {integrity: sha512-h5+1OzzfCC3Ef7VbtKdcv7zsstUQwUDlYpUTvjeUsJAssPgLn7QzbboPtL5ro04Mq0rPOsMzl7q5hIbRs2wD1A==} engines: {node: '>=8.0.0'} + fast-check@4.8.0: + resolution: {integrity: sha512-GOJ158CUMnN6cSahsv4+ExARvIDuzzinFjkp0E9WtiBa5zcVeLozVkWaE4IzFcc+Y48Wp1EDlUZsXRyAztQcSg==} + engines: {node: '>=12.17.0'} + fast-deep-equal@3.1.3: resolution: {integrity: sha512-f3qQ9oQy9j2AhBe/H9VC91wLmKBCCU/gDOnKNAYG5hswO7BLKj09Hc5HYNz9cGI++xlpDCIgDaitVs03ATR84Q==} @@ -3370,8 +3635,8 @@ packages: resolution: {integrity: sha512-dKx12eRCVIzqCxFGplyFKJMPvLEWgmNtUrpTiJIR5u97zEhRG8ySrtboPHZXx7daLxQVrl643cTzbab2tkQjxg==} engines: {node: '>= 0.4'} - form-data@4.0.5: - resolution: {integrity: sha512-8RipRLol37bNs2bhoV67fiTEvdTrbMUYcFTiy3+wuuOnUog2QBHCZWXDRijWQfAkhBj2Uf5UnVaiWwA5vdd82w==} + form-data@4.0.6: + resolution: {integrity: sha512-vKatAh4SlVfgbv+YtmhiRjhEMJsYpsG1Y2rMQtR+SVSbytsSD1YGzDIcrAJmdFec88u/+VoGmxnl+80gL1tRCQ==} engines: {node: '>= 6'} formdata-polyfill@4.0.10: @@ -3522,6 +3787,10 @@ packages: resolution: {integrity: sha512-ej4AhfhfL2Q2zpMmLo7U1Uv9+PyhIZpgQLGT1F9miIGmiCJIoCgSmczFdrc97mWT4kVY72KA+WnnhJ5pghSvSg==} engines: {node: '>= 0.4'} + hasown@2.0.4: + resolution: {integrity: sha512-T2UbfbBEF32wiepXIsMlTW9+dDYC6wMh/t/vYA4tuOMKqWz/n3vr1NFSxQiyP+zk2mXsoMA/i/7qV6LKut1t1A==} + engines: {node: '>= 0.4'} + hast-util-to-jsx-runtime@2.3.6: resolution: {integrity: sha512-zl6s8LwNyo1P9uw+XJGvZtdFF1GdAkOg8ujOw+4Pyb76874fLps4ueHXDhXWdk6YHQ6OgUtinliG7RsYvCbbBg==} @@ -3552,6 +3821,10 @@ packages: resolution: {integrity: sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==} engines: {node: '>= 0.8'} + https-proxy-agent@5.0.1: + resolution: {integrity: sha512-dFcAjpTQFgoLMzC2VwU+C/CbS7uRL0lWmxDITmqm7C+7F0Odmj6s9l6alZc6AELXhrnggM2CeWSXHGOdX2YtwA==} + engines: {node: '>= 6'} + https-proxy-agent@7.0.6: resolution: {integrity: sha512-vK9P5/iUfdl95AI+JVyUuIcVtd4ofvtrOr3HNtM2yxC9bnMbEdp3x01OhQNnjb8IJYi38VlTE3mBXwcfvywuSw==} engines: {node: '>= 14'} @@ -3590,6 +3863,10 @@ packages: inherits@2.0.4: resolution: {integrity: sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==} + ini@7.0.0: + resolution: {integrity: sha512-ifK0CgjALofS5bkrcTy4RaQ9Vx2Knf/eLeIO+NaswQEpH1UblrtTSCIvN71qQDMq0PeQ/SSPojvEJp9vvvfr+w==} + engines: {node: ^22.22.2 || ^24.15.0 || >=26.0.0} + inline-style-parser@0.2.7: resolution: {integrity: sha512-Nb2ctOyNR8DqQoR0OwRG95uNWIC0C1lCgf5Naz5H6Ji72KZ8OcFZLz2P5sNgwlyoJ8Yif11oMuYs5pBQa86csA==} @@ -3597,6 +3874,10 @@ packages: resolution: {integrity: sha512-4gd7VpWNQNB4UKKCFFVcp1AVv+FMOgs9NKzjHKusc8jTMhd5eL1NqQqOpE0KzMds804/yHlglp3uxgluOqAPLw==} engines: {node: '>= 0.4'} + ioredis@5.11.1: + resolution: {integrity: sha512-ehuGcf94bQXhfagULNXrJdfnWO38v070jxSx/qE87Kjzmu2fU7ro5EFAb+OPituLqgfyuQaym5DlrNydW2sJ9A==} + engines: {node: '>=12.22.0'} + ip-address@10.1.0: resolution: {integrity: sha512-XXADHxXmvT9+CRxhXg56LJovE+bmWnEWB78LB83VZTprKTmaC5QfruXocxzTZ2Kl0DNwKuBdlIhjL8LeY8Sf8Q==} engines: {node: '>= 12'} @@ -3901,6 +4182,9 @@ packages: resolution: {integrity: sha512-o+NO+8WrRiQEE4/7nwRJhN1HWpVmJm511pBHUxPLtp0BUISzlBplORYSmTclCnJvQq2tKu/sgl3xVpkc7ZWuQQ==} engines: {node: '>=6'} + kubernetes-types@1.30.0: + resolution: {integrity: sha512-Dew1okvhM/SQcIa2rcgujNndZwU8VnSapDgdxlYoB84ZlpAD43U6KLAFqYo17ykSFGHNPrg0qry0bP+GJd9v7Q==} + language-subtag-registry@0.3.23: resolution: {integrity: sha512-0K65Lea881pHotoGEa5gDlMxt3pctLi2RplBb7Ezh4rRdLEOtgi7n4EwK9lamnUCkKBqaeKRVebTq6BAxSkpXQ==} @@ -4225,6 +4509,11 @@ packages: resolution: {integrity: sha512-Lbgzdk0h4juoQ9fCKXW4by0UJqj+nOOrI9MJ1sSj4nI8aI2eo1qmvQEie4VD1glsS250n15LsWsYtCugiStS5A==} engines: {node: '>=18'} + mime@4.1.0: + resolution: {integrity: sha512-X5ju04+cAzsojXKes0B/S4tcYtFAJ6tTMuSPBEn9CPGlrWr8Fiw7qYeLT0XyH80HSoAoqWCaz+MWKh22P7G1cw==} + engines: {node: '>=16'} + hasBin: true + mimic-fn@2.1.0: resolution: {integrity: sha512-OqbOk5oEQeAZ8WXWydlu9HJjz9WVdEIvamMCcXmuqUYjTknH/sqsWvhQ3vgwKFRR1HpjvNBKQ37nbJgYzGqGcg==} engines: {node: '>=6'} @@ -4250,9 +4539,16 @@ packages: resolution: {integrity: sha512-P0efT1C9jIdVRefqjzOQ9Xml57zpOXnIuS+csaB4MdZbTdmGDLo8XhzBG1N7aO11gKDDkJvBLULeFTo46wwreA==} hasBin: true + msgpackr-extract@3.0.4: + resolution: {integrity: sha512-4kmO/MdyUIkLIvTPr8VHLil4AtoKIoniWPIEk5+CDy0xnWC84azhSFmuJ7PxZdsYtiP5kEeQsORAVIeMgxT+Hw==} + hasBin: true + msgpackr@1.11.12: resolution: {integrity: sha512-RBdJ1Un7yGlXWajrkxcSa93nvQ0w4zBf60c0yYv7YtBelP8H2FA7XsfBbMHtXKXUMUxH7zV3Zuozh+kUQWhHvg==} + msgpackr@2.0.4: + resolution: {integrity: sha512-o1C5KRmuRt+apqMr1HuGSqWStZoRBUpEsCsl15uM9VdAF1qHLtvMOU2En747EnTyEl6c4pzPewRMFF31s1CNbA==} + msw@2.14.3: resolution: {integrity: sha512-kk8G5cocVlJ4wsKMGZegn2H6XLOEKjbA+nSJE2354e/SRp4mDicCHUYnMXpymzVcVDCs+GUAsmNqSn+yHv4T2A==} engines: {node: '>=18'} @@ -4490,6 +4786,10 @@ packages: resolution: {integrity: sha512-DlOzet0HO7OEnmUmB6wWGJrrdvbyJKftI1bhMitK7O2N8W2gc757yyYBbINy9IDafXAV9wmKr9t7xsTaNKRG5Q==} engines: {node: '>=20.16.0 || >=22.3.0'} + pdfjs-dist@6.1.200: + resolution: {integrity: sha512-o8MolyzirkkLrcdsae/HEOiIcXWI7DS5zGpvqW8xTC2YUsW30rltFw2bDGvw/fskUdEMrQm2br68jzDS5BH2vw==} + engines: {node: '>=22.13.0 || >=24'} + picocolors@1.1.1: resolution: {integrity: sha512-xceH2snhtb5M9liqDsmEw56le376mTZkEX/jEb/RxNFyegNul7eNslCXP9FDj/Lcu0X8KEyMceP2ntpaHrDEVA==} @@ -4501,6 +4801,10 @@ packages: resolution: {integrity: sha512-QP88BAKvMam/3NxH6vj2o21R6MjxZUAd6nlwAS/pnGvN9IVLocLHxGYIzFhg6fUQ+5th6P4dv4eW9jX3DSIj7A==} engines: {node: '>=12'} + piscina@5.2.0: + resolution: {integrity: sha512-DszUCKeVN/5G5QKo6jAVHL8fmKnkJvQ0ACiVgY7YGCq3TUB2oznAOayvZPIAdEThvhczkXR+qm3IHsNXpFCYfA==} + engines: {node: '>=20.x'} + pkce-challenge@5.0.1: resolution: {integrity: sha512-wQ0b/W4Fr01qtpHlqSqspcj3EhBvimsdh0KlHhH8HRZnMsEa0ea2fTULOXOS9ccQr3om+GcGRk4e+isrZWV8qQ==} engines: {node: '>=16.20.0'} @@ -4585,6 +4889,9 @@ packages: pure-rand@6.1.0: resolution: {integrity: sha512-bVWawvoZoBYpp6yIoQtQXHZjmz35RSVHnUOTefl8Vcjr8snTPY1wnpSPMWekcFwbxI6gtmT7rSYPFvz71ldiOA==} + pure-rand@8.4.1: + resolution: {integrity: sha512-c58R2+SPFcSIPXoU834QN/KPDDOSd8sXcSrqf6e83Me6Rrp1EYkxukkjXMVrKvKaADs1SOyNkWdfvLf6zY8qLQ==} + qs@6.15.1: resolution: {integrity: sha512-6YHEFRL9mfgcAvql/XhwTvf5jKcOiiupt2FiJxHkiX1z4j7WL8J/jRHYLluORvc1XxB5rV20KoeK00gVJamspg==} engines: {node: '>=0.6'} @@ -4671,6 +4978,14 @@ packages: resolution: {integrity: sha512-YTUo+Flmw4ZXiWfQKGcwwc11KnoRAYgzAE2E7mXKCjSviTKShtxBsN6YUUBB2gtaBzKzeKunxhUwNHQuRryhWA==} engines: {node: '>= 4'} + redis-errors@1.2.0: + resolution: {integrity: sha512-1qny3OExCf0UvUV/5wpYKf2YwPcOqXzkwKKSmKHiE6ZMQs5heeE/c8eXK+PNllPvmjgAbfnsbpkGZWy8cBpn9w==} + engines: {node: '>=4'} + + redis-parser@3.0.0: + resolution: {integrity: sha512-DJnGAeenTdpMEH6uAJRK/uiyEIH9WVsUmoLwzudwGJUwZPp80PDBWPHXSAGNPwNvIXAbe7MSUB1zQFugFml66A==} + engines: {node: '>=4'} + reflect.getprototypeof@1.0.10: resolution: {integrity: sha512-00o4I+DVrefhv+nX0ulyi3biSHCPDe+yLv5o/p6d/UVlirijB8E16FtfwSAi4g3tcqrQ4lRAqQSoFEZJehYEcw==} engines: {node: '>= 0.4'} @@ -4880,6 +5195,9 @@ packages: stackback@0.0.2: resolution: {integrity: sha512-1XMJE5fQo1jGH6Y/7ebnwPOBEkIEnT4QF32d5R1+VXdXveM0IBMJt8zfaxX1P3QhVwrYe+576+jkANtSS2mBbw==} + standard-as-callback@2.1.0: + resolution: {integrity: sha512-qoRRSyROncaz1z0mvYqIE4lCd9p2R90i6GxW3uZv5ucSu8tU7B5HXUP1gG8pVZsYNVaXjk8ClXHPttLyxAL48A==} + statuses@2.0.2: resolution: {integrity: sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==} engines: {node: '>= 0.8'} @@ -5053,6 +5371,10 @@ packages: resolution: {integrity: sha512-o5sSPKEkg/DIQNmH43V0/uerLrpzVedkUh8tGNvaeXpfpuwjKenlSox/2O/BTlZUtEe+JG7s5YhEz608PlAHRA==} engines: {node: '>=0.6'} + toml@4.1.2: + resolution: {integrity: sha512-m0vXfHODcw3gk+KONAOlVQ5yNHc3yS3B1ybM3HS1vqDoS0RWTDDVBVVTYi8hH0k+2OM1vmo9fb1WX9EVqjqfHA==} + engines: {node: '>=20'} + tough-cookie@6.0.1: resolution: {integrity: sha512-LktZQb3IeoUWB9lqR5EWTHgW/VTITCXg4D21M+lvybRVdylLrRMnqaIONLVb5mav8vM19m44HIcGq4qASeu2Qw==} engines: {node: '>=16'} @@ -5152,6 +5474,10 @@ packages: resolution: {integrity: sha512-xXnp4kTyor2Zq+J1FfPI6Eq3ew5h6Vl0F/8d9XU5zZQf1tX9s2Su1/3PiMmUANFULpmksxkClamIZcaUqryHsQ==} engines: {node: '>=20.18.1'} + undici@8.6.0: + resolution: {integrity: sha512-l2FlC6I510GawyEd1qgcE/okihKrzy+BRTEBlu6T0fdbM9m5yxtIH5Oa3ysRsH0zC4EhmWUEaSDsy2QngBeRlw==} + engines: {node: '>=22.19.0'} + unicorn-magic@0.3.0: resolution: {integrity: sha512-+QBBXBCvifc56fsbuxZQ6Sic3wqqc3WWaqxs58gvJrcOuN83HGTCwz3oS5phzU9LthRNE9VrJCFCLUgHeeFnfA==} engines: {node: '>=18'} @@ -5225,6 +5551,10 @@ packages: util-deprecate@1.0.2: resolution: {integrity: sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw==} + uuid@14.0.1: + resolution: {integrity: sha512-6ZxzVpzDXDa3bJWaHilVayA+BH/1zmxCJoVgvmqJnid/gPoKHxUrS/aC/T6LGQtNHT+XHG9fXPJB4d+IrU30Ew==} + hasBin: true + validate-npm-package-name@7.0.2: resolution: {integrity: sha512-hVDIBwsRruT73PbK7uP5ebUt+ezEtCmzZz3F59BSr2F6OVFnJ/6h8liuvdLrQ88Xmnk6/+xGGuq+pG9WwTuy3A==} engines: {node: ^20.17.0 || >=22.9.0} @@ -5391,6 +5721,18 @@ packages: wrappy@1.0.2: resolution: {integrity: sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ==} + ws@8.21.0: + resolution: {integrity: sha512-Vsp28b7DRcimFQvrqu2Wek3z1iYxDCWqHYB8Qsnk/S4RfaCQzPGPyBNuVjJV3cd6UiKtUtp6sNM77gWvzcCH+g==} + engines: {node: '>=10.0.0'} + peerDependencies: + bufferutil: ^4.0.1 + utf-8-validate: '>=5.0.2' + peerDependenciesMeta: + bufferutil: + optional: true + utf-8-validate: + optional: true + wsl-utils@0.3.1: resolution: {integrity: sha512-g/eziiSUNBSsdDJtCLB8bdYEUMj4jR7AGeUo96p/3dTafgjHhpF4RiCFPiRILwjQoDXx5MqkBr4fwWtR3Ky4Wg==} engines: {node: '>=20'} @@ -5413,6 +5755,11 @@ packages: yallist@3.1.1: resolution: {integrity: sha512-a4UGQaWPH59mOXUYnAG2ewncQS4i4F43Tv3JoAM+s2VDAmS9NsK8GpDMLrCHPksFT7h3K6TOoUNn2pb7RoXx4g==} + yaml@2.9.0: + resolution: {integrity: sha512-2AvhNX3mb8zd6Zy7INTtSpl1F15HW6Wnqj0srWlkKLcpYl/gMIMJiyuGq2KeI2YFxUPjdlB+3Lc10seMLtL4cA==} + engines: {node: '>= 14.6'} + hasBin: true + yargs-parser@21.1.1: resolution: {integrity: sha512-tVpsJW7DdjecAiFpbIB1e3qxIQsE6NoPc5/eTdrbbIC4h0LVsWhnoa3g+m2HclBIujHzsxZ4VJVA+GUuc2/LBw==} engines: {node: '>=12'} @@ -5747,6 +6094,26 @@ snapshots: '@effect/language-service@0.85.1': {} + '@effect/platform-node-shared@4.0.0-beta.93(effect@4.0.0-beta.93)': + dependencies: + '@types/ws': 8.18.1 + effect: 4.0.0-beta.93 + ws: 8.21.0 + transitivePeerDependencies: + - bufferutil + - utf-8-validate + + '@effect/platform-node@4.0.0-beta.93(effect@4.0.0-beta.93)(ioredis@5.11.1)': + dependencies: + '@effect/platform-node-shared': 4.0.0-beta.93(effect@4.0.0-beta.93) + effect: 4.0.0-beta.93 + ioredis: 5.11.1 + mime: 4.1.0 + undici: 8.6.0 + transitivePeerDependencies: + - bufferutil + - utf-8-validate + '@effect/platform@0.96.1(effect@3.21.2)': dependencies: effect: 3.21.2 @@ -6213,6 +6580,8 @@ snapshots: optionalDependencies: '@types/node': 20.19.39 + '@ioredis/commands@1.10.0': {} + '@jridgewell/gen-mapping@0.3.13': dependencies: '@jridgewell/sourcemap-codec': 1.5.5 @@ -6257,21 +6626,39 @@ snapshots: '@msgpackr-extract/msgpackr-extract-darwin-arm64@3.0.3': optional: true + '@msgpackr-extract/msgpackr-extract-darwin-arm64@3.0.4': + optional: true + '@msgpackr-extract/msgpackr-extract-darwin-x64@3.0.3': optional: true + '@msgpackr-extract/msgpackr-extract-darwin-x64@3.0.4': + optional: true + '@msgpackr-extract/msgpackr-extract-linux-arm64@3.0.3': optional: true + '@msgpackr-extract/msgpackr-extract-linux-arm64@3.0.4': + optional: true + '@msgpackr-extract/msgpackr-extract-linux-arm@3.0.3': optional: true + '@msgpackr-extract/msgpackr-extract-linux-arm@3.0.4': + optional: true + '@msgpackr-extract/msgpackr-extract-linux-x64@3.0.3': optional: true + '@msgpackr-extract/msgpackr-extract-linux-x64@3.0.4': + optional: true + '@msgpackr-extract/msgpackr-extract-win32-x64@3.0.3': optional: true + '@msgpackr-extract/msgpackr-extract-win32-x64@3.0.4': + optional: true + '@mswjs/interceptors@0.41.8': dependencies: '@open-draft/deferred-promise': 2.2.0 @@ -6284,36 +6671,69 @@ snapshots: '@napi-rs/canvas-android-arm64@0.1.100': optional: true + '@napi-rs/canvas-android-arm64@1.0.2': + optional: true + '@napi-rs/canvas-darwin-arm64@0.1.100': optional: true + '@napi-rs/canvas-darwin-arm64@1.0.2': + optional: true + '@napi-rs/canvas-darwin-x64@0.1.100': optional: true + '@napi-rs/canvas-darwin-x64@1.0.2': + optional: true + '@napi-rs/canvas-linux-arm-gnueabihf@0.1.100': optional: true + '@napi-rs/canvas-linux-arm-gnueabihf@1.0.2': + optional: true + '@napi-rs/canvas-linux-arm64-gnu@0.1.100': optional: true + '@napi-rs/canvas-linux-arm64-gnu@1.0.2': + optional: true + '@napi-rs/canvas-linux-arm64-musl@0.1.100': optional: true + '@napi-rs/canvas-linux-arm64-musl@1.0.2': + optional: true + '@napi-rs/canvas-linux-riscv64-gnu@0.1.100': optional: true + '@napi-rs/canvas-linux-riscv64-gnu@1.0.2': + optional: true + '@napi-rs/canvas-linux-x64-gnu@0.1.100': optional: true + '@napi-rs/canvas-linux-x64-gnu@1.0.2': + optional: true + '@napi-rs/canvas-linux-x64-musl@0.1.100': optional: true + '@napi-rs/canvas-linux-x64-musl@1.0.2': + optional: true + '@napi-rs/canvas-win32-arm64-msvc@0.1.100': optional: true + '@napi-rs/canvas-win32-arm64-msvc@1.0.2': + optional: true + '@napi-rs/canvas-win32-x64-msvc@0.1.100': optional: true + '@napi-rs/canvas-win32-x64-msvc@1.0.2': + optional: true + '@napi-rs/canvas@0.1.100': optionalDependencies: '@napi-rs/canvas-android-arm64': 0.1.100 @@ -6329,6 +6749,92 @@ snapshots: '@napi-rs/canvas-win32-x64-msvc': 0.1.100 optional: true + '@napi-rs/canvas@1.0.2': + optionalDependencies: + '@napi-rs/canvas-android-arm64': 1.0.2 + '@napi-rs/canvas-darwin-arm64': 1.0.2 + '@napi-rs/canvas-darwin-x64': 1.0.2 + '@napi-rs/canvas-linux-arm-gnueabihf': 1.0.2 + '@napi-rs/canvas-linux-arm64-gnu': 1.0.2 + '@napi-rs/canvas-linux-arm64-musl': 1.0.2 + '@napi-rs/canvas-linux-riscv64-gnu': 1.0.2 + '@napi-rs/canvas-linux-x64-gnu': 1.0.2 + '@napi-rs/canvas-linux-x64-musl': 1.0.2 + '@napi-rs/canvas-win32-arm64-msvc': 1.0.2 + '@napi-rs/canvas-win32-x64-msvc': 1.0.2 + + '@napi-rs/nice-android-arm-eabi@1.1.1': + optional: true + + '@napi-rs/nice-android-arm64@1.1.1': + optional: true + + '@napi-rs/nice-darwin-arm64@1.1.1': + optional: true + + '@napi-rs/nice-darwin-x64@1.1.1': + optional: true + + '@napi-rs/nice-freebsd-x64@1.1.1': + optional: true + + '@napi-rs/nice-linux-arm-gnueabihf@1.1.1': + optional: true + + '@napi-rs/nice-linux-arm64-gnu@1.1.1': + optional: true + + '@napi-rs/nice-linux-arm64-musl@1.1.1': + optional: true + + '@napi-rs/nice-linux-ppc64-gnu@1.1.1': + optional: true + + '@napi-rs/nice-linux-riscv64-gnu@1.1.1': + optional: true + + '@napi-rs/nice-linux-s390x-gnu@1.1.1': + optional: true + + '@napi-rs/nice-linux-x64-gnu@1.1.1': + optional: true + + '@napi-rs/nice-linux-x64-musl@1.1.1': + optional: true + + '@napi-rs/nice-openharmony-arm64@1.1.1': + optional: true + + '@napi-rs/nice-win32-arm64-msvc@1.1.1': + optional: true + + '@napi-rs/nice-win32-ia32-msvc@1.1.1': + optional: true + + '@napi-rs/nice-win32-x64-msvc@1.1.1': + optional: true + + '@napi-rs/nice@1.1.1': + optionalDependencies: + '@napi-rs/nice-android-arm-eabi': 1.1.1 + '@napi-rs/nice-android-arm64': 1.1.1 + '@napi-rs/nice-darwin-arm64': 1.1.1 + '@napi-rs/nice-darwin-x64': 1.1.1 + '@napi-rs/nice-freebsd-x64': 1.1.1 + '@napi-rs/nice-linux-arm-gnueabihf': 1.1.1 + '@napi-rs/nice-linux-arm64-gnu': 1.1.1 + '@napi-rs/nice-linux-arm64-musl': 1.1.1 + '@napi-rs/nice-linux-ppc64-gnu': 1.1.1 + '@napi-rs/nice-linux-riscv64-gnu': 1.1.1 + '@napi-rs/nice-linux-s390x-gnu': 1.1.1 + '@napi-rs/nice-linux-x64-gnu': 1.1.1 + '@napi-rs/nice-linux-x64-musl': 1.1.1 + '@napi-rs/nice-openharmony-arm64': 1.1.1 + '@napi-rs/nice-win32-arm64-msvc': 1.1.1 + '@napi-rs/nice-win32-ia32-msvc': 1.1.1 + '@napi-rs/nice-win32-x64-msvc': 1.1.1 + optional: true + '@napi-rs/wasm-runtime@0.2.12': dependencies: '@emnapi/core': 1.10.0 @@ -6397,12 +6903,21 @@ snapshots: '@nolyfill/is-core-module@1.0.39': {} - '@ontos-ai/knowhere-sdk@2.0.0': + '@ontos-ai/knowhere-sdk@file:../knowhere-node-sdk(ioredis@5.11.1)': dependencies: - axios: 1.16.0 + '@effect/platform-node': 4.0.0-beta.93(effect@4.0.0-beta.93)(ioredis@5.11.1) + '@napi-rs/canvas': 1.0.2 + axios: 1.18.1 + effect: 4.0.0-beta.93 jszip: 3.10.1 + pdfjs-dist: 6.1.200 + piscina: 5.2.0 transitivePeerDependencies: + - bufferutil - debug + - ioredis + - supports-color + - utf-8-validate '@open-draft/deferred-promise@2.2.0': {} @@ -7065,6 +7580,10 @@ snapshots: '@types/validate-npm-package-name@4.0.2': {} + '@types/ws@8.18.1': + dependencies: + '@types/node': 20.19.39 + '@typescript-eslint/eslint-plugin@8.59.2(@typescript-eslint/parser@8.59.2(eslint@9.39.4(jiti@2.7.0))(typescript@6.0.3))(eslint@9.39.4(jiti@2.7.0))(typescript@6.0.3)': dependencies: '@eslint-community/regexpp': 4.12.2 @@ -7238,10 +7757,10 @@ snapshots: '@vercel/oidc@3.2.0': {} - '@vitejs/plugin-react@6.0.1(babel-plugin-react-compiler@1.0.0)(vite@8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0))': + '@vitejs/plugin-react@6.0.1(babel-plugin-react-compiler@1.0.0)(vite@8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0)(yaml@2.9.0))': dependencies: '@rolldown/pluginutils': 1.0.0-rc.7 - vite: 8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0) + vite: 8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0)(yaml@2.9.0) optionalDependencies: babel-plugin-react-compiler: 1.0.0 @@ -7254,14 +7773,14 @@ snapshots: chai: 6.2.2 tinyrainbow: 3.1.0 - '@vitest/mocker@4.1.5(msw@2.14.3(@types/node@20.19.39)(typescript@6.0.3))(vite@8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0))': + '@vitest/mocker@4.1.5(msw@2.14.3(@types/node@20.19.39)(typescript@6.0.3))(vite@8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0)(yaml@2.9.0))': dependencies: '@vitest/spy': 4.1.5 estree-walker: 3.0.3 magic-string: 0.30.21 optionalDependencies: msw: 2.14.3(@types/node@20.19.39)(typescript@6.0.3) - vite: 8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0) + vite: 8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0)(yaml@2.9.0) '@vitest/pretty-format@4.1.5': dependencies: @@ -7300,6 +7819,12 @@ snapshots: acorn@8.16.0: {} + agent-base@6.0.2: + dependencies: + debug: 4.4.3 + transitivePeerDependencies: + - supports-color + agent-base@7.1.4: {} ai@6.0.175(zod@4.4.3): @@ -7443,13 +7968,15 @@ snapshots: axe-core@4.11.4: {} - axios@1.16.0: + axios@1.18.1: dependencies: follow-redirects: 1.16.0 - form-data: 4.0.5 + form-data: 4.0.6 + https-proxy-agent: 5.0.1 proxy-from-env: 2.1.0 transitivePeerDependencies: - debug + - supports-color axobject-query@4.1.0: {} @@ -7578,6 +8105,8 @@ snapshots: clsx@2.1.1: {} + cluster-key-slot@1.1.1: {} + code-block-writer@13.0.3: {} color-convert@2.0.1: @@ -7721,6 +8250,8 @@ snapshots: delayed-stream@1.0.0: {} + denque@2.1.0: {} + depd@2.0.0: {} dequal@2.0.3: {} @@ -7790,6 +8321,19 @@ snapshots: '@standard-schema/spec': 1.1.0 fast-check: 3.23.2 + effect@4.0.0-beta.93: + dependencies: + '@standard-schema/spec': 1.1.0 + fast-check: 4.8.0 + find-my-way-ts: 0.1.6 + ini: 7.0.0 + kubernetes-types: 1.30.0 + msgpackr: 2.0.4 + multipasta: 0.2.7 + toml: 4.1.2 + uuid: 14.0.1 + yaml: 2.9.0 + electron-to-chromium@1.5.351: {} emoji-regex@10.6.0: {} @@ -8305,6 +8849,10 @@ snapshots: dependencies: pure-rand: 6.1.0 + fast-check@4.8.0: + dependencies: + pure-rand: 8.4.1 + fast-deep-equal@3.1.3: {} fast-glob@3.3.1: @@ -8397,12 +8945,12 @@ snapshots: dependencies: is-callable: 1.2.7 - form-data@4.0.5: + form-data@4.0.6: dependencies: asynckit: 0.4.0 combined-stream: 1.0.8 es-set-tostringtag: 2.1.0 - hasown: 2.0.3 + hasown: 2.0.4 mime-types: 2.1.35 formdata-polyfill@4.0.10: @@ -8539,6 +9087,10 @@ snapshots: dependencies: function-bind: 1.1.2 + hasown@2.0.4: + dependencies: + function-bind: 1.1.2 + hast-util-to-jsx-runtime@2.3.6: dependencies: '@types/estree': 1.0.8 @@ -8592,6 +9144,13 @@ snapshots: statuses: 2.0.2 toidentifier: 1.0.1 + https-proxy-agent@5.0.1: + dependencies: + agent-base: 6.0.2 + debug: 4.4.3 + transitivePeerDependencies: + - supports-color + https-proxy-agent@7.0.6: dependencies: agent-base: 7.1.4 @@ -8622,6 +9181,8 @@ snapshots: inherits@2.0.4: {} + ini@7.0.0: {} + inline-style-parser@0.2.7: {} internal-slot@1.1.0: @@ -8630,6 +9191,18 @@ snapshots: hasown: 2.0.3 side-channel: 1.1.0 + ioredis@5.11.1: + dependencies: + '@ioredis/commands': 1.10.0 + cluster-key-slot: 1.1.1 + debug: 4.4.3 + denque: 2.1.0 + redis-errors: 1.2.0 + redis-parser: 3.0.0 + standard-as-callback: 2.1.0 + transitivePeerDependencies: + - supports-color + ip-address@10.1.0: {} ipaddr.js@1.9.1: {} @@ -8911,6 +9484,8 @@ snapshots: kleur@4.1.5: {} + kubernetes-types@1.30.0: {} + language-subtag-registry@0.3.23: {} language-tags@1.0.9: @@ -9412,6 +9987,8 @@ snapshots: dependencies: mime-db: 1.54.0 + mime@4.1.0: {} + mimic-fn@2.1.0: {} mimic-function@5.0.1: {} @@ -9440,10 +10017,26 @@ snapshots: '@msgpackr-extract/msgpackr-extract-win32-x64': 3.0.3 optional: true + msgpackr-extract@3.0.4: + dependencies: + node-gyp-build-optional-packages: 5.2.2 + optionalDependencies: + '@msgpackr-extract/msgpackr-extract-darwin-arm64': 3.0.4 + '@msgpackr-extract/msgpackr-extract-darwin-x64': 3.0.4 + '@msgpackr-extract/msgpackr-extract-linux-arm': 3.0.4 + '@msgpackr-extract/msgpackr-extract-linux-arm64': 3.0.4 + '@msgpackr-extract/msgpackr-extract-linux-x64': 3.0.4 + '@msgpackr-extract/msgpackr-extract-win32-x64': 3.0.4 + optional: true + msgpackr@1.11.12: optionalDependencies: msgpackr-extract: 3.0.3 + msgpackr@2.0.4: + optionalDependencies: + msgpackr-extract: 3.0.4 + msw@2.14.3(@types/node@20.19.39)(typescript@6.0.3): dependencies: '@inquirer/confirm': 6.0.12(@types/node@20.19.39) @@ -9709,12 +10302,20 @@ snapshots: optionalDependencies: '@napi-rs/canvas': 0.1.100 + pdfjs-dist@6.1.200: + optionalDependencies: + '@napi-rs/canvas': 1.0.2 + picocolors@1.1.1: {} picomatch@2.3.2: {} picomatch@4.0.4: {} + piscina@5.2.0: + optionalDependencies: + '@napi-rs/nice': 1.1.1 + pkce-challenge@5.0.1: {} playwright-core@1.58.2: {} @@ -9799,6 +10400,8 @@ snapshots: pure-rand@6.1.0: {} + pure-rand@8.4.1: {} + qs@6.15.1: dependencies: side-channel: 1.1.0 @@ -9905,6 +10508,12 @@ snapshots: tiny-invariant: 1.3.3 tslib: 2.8.1 + redis-errors@1.2.0: {} + + redis-parser@3.0.0: + dependencies: + redis-errors: 1.2.0 + reflect.getprototypeof@1.0.10: dependencies: call-bind: 1.0.9 @@ -10251,6 +10860,8 @@ snapshots: stackback@0.0.2: {} + standard-as-callback@2.1.0: {} + statuses@2.0.2: {} std-env@4.1.0: {} @@ -10423,6 +11034,8 @@ snapshots: toidentifier@1.0.1: {} + toml@4.1.2: {} + tough-cookie@6.0.1: dependencies: tldts: 7.0.30 @@ -10543,6 +11156,8 @@ snapshots: undici@7.25.0: {} + undici@8.6.0: {} + unicorn-magic@0.3.0: {} unified@11.0.5: @@ -10639,6 +11254,8 @@ snapshots: util-deprecate@1.0.2: {} + uuid@14.0.1: {} + validate-npm-package-name@7.0.2: {} vary@1.1.2: {} @@ -10653,7 +11270,7 @@ snapshots: '@types/unist': 3.0.3 vfile-message: 4.0.3 - vite@8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0): + vite@8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0)(yaml@2.9.0): dependencies: lightningcss: 1.32.0 picomatch: 4.0.4 @@ -10666,11 +11283,12 @@ snapshots: fsevents: 2.3.3 jiti: 2.7.0 tsx: 4.21.0 + yaml: 2.9.0 - vitest@4.1.5(@opentelemetry/api@1.9.0)(@types/node@20.19.39)(jsdom@29.1.1(@noble/hashes@1.8.0))(msw@2.14.3(@types/node@20.19.39)(typescript@6.0.3))(vite@8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0)): + vitest@4.1.5(@opentelemetry/api@1.9.0)(@types/node@20.19.39)(jsdom@29.1.1(@noble/hashes@1.8.0))(msw@2.14.3(@types/node@20.19.39)(typescript@6.0.3))(vite@8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0)(yaml@2.9.0)): dependencies: '@vitest/expect': 4.1.5 - '@vitest/mocker': 4.1.5(msw@2.14.3(@types/node@20.19.39)(typescript@6.0.3))(vite@8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0)) + '@vitest/mocker': 4.1.5(msw@2.14.3(@types/node@20.19.39)(typescript@6.0.3))(vite@8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0)(yaml@2.9.0)) '@vitest/pretty-format': 4.1.5 '@vitest/runner': 4.1.5 '@vitest/snapshot': 4.1.5 @@ -10687,7 +11305,7 @@ snapshots: tinyexec: 1.1.2 tinyglobby: 0.2.16 tinyrainbow: 3.1.0 - vite: 8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0) + vite: 8.0.10(@types/node@20.19.39)(esbuild@0.27.7)(jiti@2.7.0)(tsx@4.21.0)(yaml@2.9.0) why-is-node-running: 2.3.0 optionalDependencies: '@opentelemetry/api': 1.9.0 @@ -10784,6 +11402,8 @@ snapshots: wrappy@1.0.2: {} + ws@8.21.0: {} + wsl-utils@0.3.1: dependencies: is-wsl: 3.1.1 @@ -10799,6 +11419,8 @@ snapshots: yallist@3.1.1: {} + yaml@2.9.0: {} + yargs-parser@21.1.1: {} yargs@17.7.2: diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index d989c76..a49c9d9 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -1,5 +1,9 @@ ignoredBuiltDependencies: - sharp - unrs-resolver + minimumReleaseAgeExclude: - '@ontos-ai/knowhere-sdk@2.0.0' + +overrides: + '@ontos-ai/knowhere-sdk': file:../knowhere-node-sdk diff --git a/scripts/prepare-effect.sh b/scripts/prepare-effect.sh new file mode 100755 index 0000000..24a9d31 --- /dev/null +++ b/scripts/prepare-effect.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env sh + +set -eu + +repo_dir=".repos/effect" +repo_url="https://github.com/Effect-TS/effect-smol" + +if [ -d "$repo_dir/.git" ]; then + exit 0 +fi + +mkdir -p ".repos" +git clone "$repo_url" "$repo_dir" diff --git a/src/agent-harness/ledger.ts b/src/agent-harness/ledger.ts index da197e8..103f9c0 100644 --- a/src/agent-harness/ledger.ts +++ b/src/agent-harness/ledger.ts @@ -68,10 +68,14 @@ export function createEvidenceLedger() { chunk: { ref: `r${retrievalIndex}:referenced:${index + 1}`, kind: "referenced_chunk", + chunkId: chunk.chunkId, content, contentPreview: content, chunkType: chunk.chunkType, score: null, + sourceChunkPath: chunk.sourceChunkPath, + filePath: chunk.filePath, + metadata: chunk.metadata, source: { documentId: chunk.documentId, sourceFileName: null, @@ -146,6 +150,10 @@ function addChunkFromResult(input: { contentPreview: buildContentPreview(input.result.content), chunkType: input.result.chunkType, score: input.result.score, + chunkId: input.result.chunkId, + sourceChunkPath: input.result.sourceChunkPath, + filePath: input.result.filePath, + metadata: input.result.metadata, source: { documentId: input.result.source.documentId, sourceFileName: input.result.source.sourceFileName, diff --git a/src/agent-harness/types.ts b/src/agent-harness/types.ts index 1a1cce9..c559623 100644 --- a/src/agent-harness/types.ts +++ b/src/agent-harness/types.ts @@ -94,10 +94,14 @@ export type RetrievalCapability = { export type EvidenceChunk = { readonly ref: string readonly kind: "result" | "referenced_chunk" + readonly chunkId?: string readonly content: string readonly contentPreview: string readonly chunkType: string readonly score: number | null + readonly sourceChunkPath?: string | null + readonly filePath?: string | null + readonly metadata?: Readonly> readonly source: { readonly documentId?: string | null readonly sourceFileName?: string | null diff --git a/src/components/chat-message-list.test.ts b/src/components/chat-message-list.test.ts index 8d58630..f94478e 100644 --- a/src/components/chat-message-list.test.ts +++ b/src/components/chat-message-list.test.ts @@ -129,6 +129,56 @@ describe("ChatMessageList", () => { ); }); + it("renders a separate page image link without replacing source focus", async () => { + const user = userEvent.setup(); + const onCitationClick = vi.fn(); + + render( + React.createElement(ChatMessageList, { + messages: [ + { + id: "assistant_1", + role: "assistant", + content: "The referenced page discusses revenue.", + citations: [ + { + chunkType: "page", + score: 0.9, + pageCitationAssetUrl: + "https://blob.example/pages/page-000004.png", + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "Page 4", + }, + }, + ], + }, + ], + onCitationClick, + }), + ); + + const citationButton = screen.getByRole("button", { + name: "Open source report.pdf", + }); + const pageImageLink = screen.getByRole("link", { + name: "Open page image for report.pdf", + }); + + expect(pageImageLink.getAttribute("href")).toBe( + "https://blob.example/pages/page-000004.png", + ); + + await user.click(citationButton); + expect(onCitationClick).toHaveBeenCalledWith( + expect.objectContaining({ + pageCitationAssetUrl: "https://blob.example/pages/page-000004.png", + }), + "assistant_1:0", + ); + }); + it("removes description-only source labels without changing other markdown whitespace", () => { render( React.createElement(ChatMessageList, { diff --git a/src/components/chat-message-list.tsx b/src/components/chat-message-list.tsx index e5e6460..531c71f 100644 --- a/src/components/chat-message-list.tsx +++ b/src/components/chat-message-list.tsx @@ -2,7 +2,7 @@ import { type CSSProperties, type ReactElement } from "react"; import { type VirtualItem } from "@tanstack/react-virtual"; -import { ImageIcon, MessageCircle } from "lucide-react"; +import { ExternalLink, ImageIcon, MessageCircle } from "lucide-react"; import ReactMarkdown, { defaultUrlTransform, type Components, @@ -633,28 +633,60 @@ function CitationChip({ citationId: string, ) => void; }): ReactElement { + const pageCitationAssetUrl = getTrimmedCitationField( + citation.pageCitationAssetUrl, + ); + return ( - - - + + - {label} - - - - {tooltipLabel} - - + {tooltipLabel} + + + {pageCitationAssetUrl ? ( + + + + + + + + Open page image + + + ) : null} + ); } diff --git a/src/domains/chat/chat-citation-persistence.ts b/src/domains/chat/chat-citation-persistence.ts index bf2a1a1..54e89de 100644 --- a/src/domains/chat/chat-citation-persistence.ts +++ b/src/domains/chat/chat-citation-persistence.ts @@ -86,6 +86,7 @@ function toCitationView( chunkType: citation.chunkType, score: citation.score, assetUrl: citation.assetUrl, + pageCitationAssetUrl: citation.pageCitationAssetUrl, description: "description" in citation ? citation.description : undefined, source: { documentId: citation.source.documentId, diff --git a/src/domains/chat/citations.test.ts b/src/domains/chat/citations.test.ts index ce681ec..b6daa0a 100644 --- a/src/domains/chat/citations.test.ts +++ b/src/domains/chat/citations.test.ts @@ -1,7 +1,7 @@ import { describe, expect, it } from "vitest" -import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" import { toChatCitationViews } from "./citations" +import type { PageCitationAssetRetrievalResult } from "./page-citation-assets" describe("toChatCitationViews", () => { it("extracts first citation description for each generated source label", () => { @@ -31,11 +31,24 @@ describe("toChatCitationViews", () => { { ...secondResult, description: "margin expansion" }, ]) }) + + it("preserves page citation asset URLs", () => { + const result = makeRetrievalResult({ + chunkType: "page", + pageCitationAssetUrl: "https://blob.example/pages/page-000004.png", + }) + + const citations = toChatCitationViews([result], "Grounded answer.") + + expect(citations[0]?.pageCitationAssetUrl).toBe( + "https://blob.example/pages/page-000004.png", + ) + }) }) function makeRetrievalResult( - overrides: Partial = {}, -): RetrievalResult { + overrides: Partial = {}, +): PageCitationAssetRetrievalResult { return { content: "Grounding content", chunkType: "text", diff --git a/src/domains/chat/citations.ts b/src/domains/chat/citations.ts index e010bdc..b6b4038 100644 --- a/src/domains/chat/citations.ts +++ b/src/domains/chat/citations.ts @@ -2,9 +2,10 @@ import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" import type { ChatCitationView } from "@/domains/chat/types" +import type { PageCitationAssetRetrievalResult } from "./page-citation-assets" export function toChatCitationViews( - results: readonly RetrievalResult[], + results: readonly PageCitationAssetRetrievalResult[], answer: string, ): ChatCitationView[] { const descriptionsBySourceNumber = getCitationDescriptions(answer) @@ -16,6 +17,9 @@ export function toChatCitationViews( chunkType: result.chunkType, score: result.score, ...(result.assetUrl ? { assetUrl: result.assetUrl } : {}), + ...(result.pageCitationAssetUrl + ? { pageCitationAssetUrl: result.pageCitationAssetUrl } + : {}), ...(description ? { description } : {}), source: { documentId: result.source.documentId, diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index 2d7245d..b0a95d9 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -38,6 +38,7 @@ import { enrichRetrievalResultsWithAssetUrls, removeRetrievedMediaAssetUrls, } from "./media-assets" +import { enrichRetrievalResultsWithPageCitationAssetUrls } from "./page-citation-assets" const DEFAULT_TOP_K = 8 const MAX_AGENTIC_TOP_K = 12 @@ -262,7 +263,12 @@ export const answerQuestionWithRetrieval = ( hardenedArtifacts: hardenedMedia.artifacts, }), }) - const citationResults = hardenedMedia.results + const citationResults = yield* Effect.tryPromise(() => + enrichRetrievalResultsWithPageCitationAssetUrls({ + results: hardenedMedia.results, + sources: input.sources, + }), + ) const displayArtifacts = hardenedMedia.artifacts ?? [] logger.info("chat-agent: answer complete", { answerLength: answer.length, @@ -792,17 +798,7 @@ function mapManifestCitationsToResults( resolveChunkForAssetRef(citation.ref, assetsByRef, chunksByRef) if (!chunk) continue - const retrievalResult: RetrievalResult = { - content: chunk.content, - chunkType: chunk.chunkType, - score: chunk.score, - ...(chunk.assetUrl ? { assetUrl: chunk.assetUrl } : {}), - source: { - documentId: chunk.source.documentId ?? undefined, - sourceFileName: chunk.source.sourceFileName ?? undefined, - sectionPath: chunk.source.sectionPath ?? undefined, - }, - } + const retrievalResult = toRetrievalResultFromEvidenceChunk(chunk) const key = getRetrievalResultKey(retrievalResult) if (seenKeys.has(key)) continue @@ -889,10 +885,14 @@ function toRetrievalResultFromEvidenceChunk( chunk: EvidenceChunk, ): RetrievalResult { return { + ...(chunk.chunkId ? { chunkId: chunk.chunkId } : {}), content: chunk.content, chunkType: chunk.chunkType, score: chunk.score, ...(chunk.assetUrl ? { assetUrl: chunk.assetUrl } : {}), + ...(chunk.sourceChunkPath ? { sourceChunkPath: chunk.sourceChunkPath } : {}), + ...(chunk.filePath ? { filePath: chunk.filePath } : {}), + ...(chunk.metadata ? { metadata: chunk.metadata } : {}), source: { documentId: chunk.source.documentId ?? undefined, sourceFileName: chunk.source.sourceFileName ?? undefined, @@ -921,10 +921,14 @@ function collectRetrievalResults( for (const result of [ ...response.results, ...response.referencedChunks.map((chunk): RetrievalResult => ({ + chunkId: chunk.chunkId, content: "", chunkType: chunk.chunkType, score: null, ...(chunk.assetUrl ? { assetUrl: chunk.assetUrl } : {}), + ...(chunk.sourceChunkPath ? { sourceChunkPath: chunk.sourceChunkPath } : {}), + ...(chunk.filePath ? { filePath: chunk.filePath } : {}), + ...(chunk.metadata ? { metadata: chunk.metadata } : {}), source: { documentId: chunk.documentId, sourceFileName: sourceTitlesByDocumentId.get(chunk.documentId), diff --git a/src/domains/chat/page-citation-assets.test.ts b/src/domains/chat/page-citation-assets.test.ts new file mode 100644 index 0000000..c76b385 --- /dev/null +++ b/src/domains/chat/page-citation-assets.test.ts @@ -0,0 +1,216 @@ +import { describe, expect, it } from "vitest" +import type { + KnowhereSdkStorage, + KnowhereSdkStorageHead, + KnowhereSdkStorageObject, + KnowhereSdkStorageReadResult, + KnowhereSdkStorageWriteResult, + RetrievalResult, +} from "@ontos-ai/knowhere-sdk" + +import type { Source } from "@/infrastructure/db/schema" +import { enrichRetrievalResultsWithPageCitationAssetUrls } from "./page-citation-assets" + +class MemorySdkStorage implements KnowhereSdkStorage { + private readonly objects = new Map() + + constructor(objects: Readonly>) { + for (const [key, value] of Object.entries(objects)) { + this.objects.set(key, new TextEncoder().encode(value)) + } + } + + headObject(): Promise { + return Promise.resolve(null) + } + + writeObject( + input: KnowhereSdkStorageObject, + ): Promise { + if (input.body instanceof Uint8Array) { + this.objects.set(input.key, new Uint8Array(input.body)) + } + return Promise.resolve({ key: input.key }) + } + + readObject(key: string): Promise { + const body = this.objects.get(key) + return Promise.resolve(body ? { body } : null) + } + + getObjectUrl(key: string): Promise { + return Promise.resolve(`https://blob.example/${key}`) + } +} + +describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { + it("adds a page citation asset URL for page results with matching metadata", async () => { + const storage = new MemorySdkStorage({ + "page-citation-assets/documents/doc_1/current.json": JSON.stringify({ + version: 1, + documentId: "doc_1", + jobId: "job_1", + variant: "default", + indexKey: + "page-citation-assets/documents/doc_1/jobs/job_1/variants/default/index.json", + updatedAt: "2026-07-03T00:00:00.000Z", + }), + "page-citation-assets/documents/doc_1/jobs/job_1/variants/default/index.json": + JSON.stringify({ + version: 1, + documentId: "doc_1", + jobId: "job_1", + variant: "default", + generatedAt: "2026-07-03T00:00:00.000Z", + assets: [ + { + pageNum: 4, + key: "page-citation-assets/documents/doc_1/jobs/job_1/variants/default/scale-1/page-000004.png", + assetUrl: "https://blob.example/pages/page-000004.png", + mimeType: "image/png", + width: 1200, + height: 1600, + source: "client-rendered-pdf-page", + variant: "default", + }, + ], + }), + }) + + const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ + results: [ + makeRetrievalResult({ + chunkType: "page", + metadata: { pageNums: [4] }, + }), + ], + sources: [makeSource()], + storage, + }) + + expect(result?.pageCitationAssetUrl).toBe( + "https://blob.example/pages/page-000004.png", + ) + }) + + it("uses SDK-enriched page asset metadata before reading the index", async () => { + const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ + results: [ + makeRetrievalResult({ + chunkType: "page", + metadata: { + pageAssets: [ + { + pageNum: 2, + key: "page-key", + assetUrl: "https://blob.example/pages/page-000002.png", + }, + ], + }, + }), + ], + sources: [makeSource()], + storage: new MemorySdkStorage({}), + }) + + expect(result?.pageCitationAssetUrl).toBe( + "https://blob.example/pages/page-000002.png", + ) + }) + + it("chooses the SDK-enriched asset matching the citation page metadata", async () => { + const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ + results: [ + makeRetrievalResult({ + chunkType: "page", + metadata: { + pageNums: [4], + pageAssets: [ + { + pageNum: 2, + key: "page-2-key", + assetUrl: "https://blob.example/pages/page-000002.png", + }, + { + pageNum: 4, + key: "page-4-key", + assetUrl: "https://blob.example/pages/page-000004.png", + }, + ], + }, + }), + ], + sources: [makeSource()], + storage: new MemorySdkStorage({}), + }) + + expect(result?.pageCitationAssetUrl).toBe( + "https://blob.example/pages/page-000004.png", + ) + }) + + it("leaves non-page results unchanged even when they have page metadata", async () => { + const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ + results: [ + makeRetrievalResult({ + chunkType: "text", + metadata: { + pageNums: [4], + pageAssets: [ + { + pageNum: 4, + key: "page-key", + assetUrl: "https://blob.example/pages/page-000004.png", + }, + ], + }, + }), + ], + sources: [makeSource()], + storage: new MemorySdkStorage({}), + }) + + expect(result?.pageCitationAssetUrl).toBeUndefined() + }) +}) + +function makeRetrievalResult( + overrides: Partial = {}, +): RetrievalResult { + return { + chunkId: "chunk_page_4", + content: "Page four summary", + chunkType: "page", + score: 0.8, + metadata: { pageNums: [4] }, + source: { + documentId: "doc_1", + sourceFileName: "source.pdf", + sectionPath: "Page 4", + }, + ...overrides, + } +} + +function makeSource(overrides: Partial = {}): Source { + return { + id: "source_1", + workspaceId: "workspace_1", + title: "source.pdf", + mimeType: "application/pdf", + sizeBytes: 1, + status: "ready", + failureReason: null, + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: null, + originalBlobUrl: null, + demoKey: null, + createdAt: new Date("2026-07-03T00:00:00.000Z"), + updatedAt: new Date("2026-07-03T00:00:00.000Z"), + deletedAt: null, + ...overrides, + } +} diff --git a/src/domains/chat/page-citation-assets.ts b/src/domains/chat/page-citation-assets.ts new file mode 100644 index 0000000..dd30a8d --- /dev/null +++ b/src/domains/chat/page-citation-assets.ts @@ -0,0 +1,381 @@ +import "server-only" + +import { createHash } from "node:crypto" +import type { + KnowhereSdkStorage, + PageCitationAsset, + PageCitationAssetCurrentIndex, + PageCitationAssetIndex, + RetrievalResult, +} from "@ontos-ai/knowhere-sdk" + +import type { Source } from "@/infrastructure/db/schema" +import { createVercelBlobKnowhereSdkStorage } from "@/integrations/knowhere-sdk-storage" + +export type PageCitationAssetRetrievalResult = RetrievalResult & { + readonly pageCitationAssetUrl?: string +} + +type EnrichRetrievalResultsWithPageCitationAssetUrlsInput = { + readonly results: readonly RetrievalResult[] + readonly sources: readonly Source[] + readonly storage?: KnowhereSdkStorage +} + +type PageCitationAssetCandidate = { + readonly pageNum: number + readonly assetUrl?: string + readonly key?: string +} + +const defaultVariant = "default" + +export async function enrichRetrievalResultsWithPageCitationAssetUrls({ + results, + sources, + storage = createVercelBlobKnowhereSdkStorage(), +}: EnrichRetrievalResultsWithPageCitationAssetUrlsInput): Promise< + PageCitationAssetRetrievalResult[] +> { + if (results.length === 0) return [] + + const sourcesByDocumentId = new Map( + sources.flatMap((source): readonly [string, Source][] => + source.knowhereDocumentId ? [[source.knowhereDocumentId, source]] : [], + ), + ) + const indexByKey = new Map>() + + return Promise.all( + results.map((result) => + enrichRetrievalResultWithPageCitationAssetUrl({ + result, + storage, + source: result.source.documentId + ? sourcesByDocumentId.get(result.source.documentId) + : undefined, + indexByKey, + }), + ), + ) +} + +async function enrichRetrievalResultWithPageCitationAssetUrl(input: { + readonly result: RetrievalResult + readonly source?: Source + readonly storage: KnowhereSdkStorage + readonly indexByKey: Map> +}): Promise { + if (!isPageResult(input.result)) return input.result + + const pageNumbers = getPageNumbers(input.result.metadata) + const directAsset = getDirectPageCitationAsset(input.result, pageNumbers) + if (directAsset?.assetUrl) { + return { + ...input.result, + pageCitationAssetUrl: directAsset.assetUrl, + } + } + + if (pageNumbers.length === 0) return input.result + + const documentId = getTrimmedString(input.result.source.documentId) + if (!documentId) return input.result + + const index = await resolvePageCitationAssetIndex({ + storage: input.storage, + documentId, + jobId: getResultJobId(input.result, input.source), + indexByKey: input.indexByKey, + }) + if (!index) return input.result + + const asset = index.assets.find((candidate) => + pageNumbers.includes(candidate.pageNum), + ) + if (!asset) return input.result + + const assetUrl = + getTrimmedString(asset.assetUrl) ?? + (asset.key && input.storage.getObjectUrl + ? await input.storage.getObjectUrl(asset.key).catch(() => null) + : null) + if (!assetUrl) return input.result + + return { + ...input.result, + pageCitationAssetUrl: assetUrl, + } +} + +function isPageResult(result: RetrievalResult): boolean { + return result.chunkType.toLowerCase() === "page" +} + +function getDirectPageCitationAsset( + result: RetrievalResult, + pageNumbers: readonly number[], +): PageCitationAssetCandidate | null { + const candidates = [ + isRecord(result) ? result.pageAssets : undefined, + result.metadata?.pageAssets, + ].flatMap(parsePageCitationAssetCandidates) + + if (pageNumbers.length > 0) { + const matchingCandidates = candidates.filter((candidate) => + pageNumbers.includes(candidate.pageNum), + ) + const matchingAssetWithUrl = matchingCandidates.find( + (candidate) => candidate.assetUrl, + ) + if (matchingAssetWithUrl) return matchingAssetWithUrl + if (matchingCandidates[0]) return matchingCandidates[0] + } + + return candidates.find((candidate) => candidate.assetUrl) ?? candidates[0] ?? null +} + +function parsePageCitationAssetCandidates( + value: unknown, +): readonly PageCitationAssetCandidate[] { + if (!Array.isArray(value)) return [] + + return value.flatMap((item): PageCitationAssetCandidate[] => { + if (!isRecord(item)) return [] + const pageNum = getPositiveInteger(item.pageNum) + if (!pageNum) return [] + + return [ + { + pageNum, + assetUrl: getTrimmedString(item.assetUrl) ?? undefined, + key: getTrimmedString(item.key) ?? undefined, + }, + ] + }) +} + +async function resolvePageCitationAssetIndex(input: { + readonly storage: KnowhereSdkStorage + readonly documentId: string + readonly jobId: string | null + readonly indexByKey: Map> +}): Promise { + if (!input.storage.readObject) return null + + if (input.jobId) { + const jobIndexKey = createPageCitationAssetIndexKey( + input.documentId, + input.jobId, + ) + const jobIndex = await readCachedIndex( + input.storage, + jobIndexKey, + input.indexByKey, + ) + if (jobIndex) return jobIndex + } + + const current = await readCurrentIndexPointer(input.storage, input.documentId) + if (!current) return null + + return readCachedIndex(input.storage, current.indexKey, input.indexByKey) +} + +async function readCachedIndex( + storage: KnowhereSdkStorage, + key: string, + cache: Map>, +): Promise { + let cached = cache.get(key) + if (!cached) { + cached = readPageCitationAssetIndex(storage, key) + cache.set(key, cached) + } + return cached +} + +async function readCurrentIndexPointer( + storage: KnowhereSdkStorage, + documentId: string, +): Promise { + const key = createPageCitationAssetCurrentIndexKey(documentId) + const parsed = await readJsonObject(storage, key) + if (!isRecord(parsed)) return null + + const indexKey = getTrimmedString(parsed.indexKey) + const jobId = getTrimmedString(parsed.jobId) + const variant = getTrimmedString(parsed.variant) + if (!indexKey || !jobId || !variant) return null + + return { + version: 1, + documentId, + jobId, + variant, + indexKey, + namespace: getTrimmedString(parsed.namespace) ?? undefined, + jobResultId: getTrimmedString(parsed.jobResultId) ?? undefined, + updatedAt: getTrimmedString(parsed.updatedAt) ?? "", + } +} + +async function readPageCitationAssetIndex( + storage: KnowhereSdkStorage, + key: string, +): Promise { + const parsed = await readJsonObject(storage, key) + if (!isRecord(parsed) || !Array.isArray(parsed.assets)) return null + + const assets = parsed.assets.flatMap((asset): PageCitationAsset[] => { + if (!isRecord(asset)) return [] + const pageNum = getPositiveInteger(asset.pageNum) + const keyValue = getTrimmedString(asset.key) + const mimeType = getPageCitationMimeType(asset.mimeType) + const width = getPositiveInteger(asset.width) + const height = getPositiveInteger(asset.height) + if (!pageNum || !keyValue || !mimeType || !width || !height) return [] + + return [ + { + pageNum, + key: keyValue, + assetUrl: getTrimmedString(asset.assetUrl) ?? undefined, + mimeType, + width, + height, + source: "client-rendered-pdf-page", + variant: getTrimmedString(asset.variant) ?? defaultVariant, + }, + ] + }) + if (assets.length === 0) return null + + return { + version: 1, + documentId: getTrimmedString(parsed.documentId) ?? "", + jobId: getTrimmedString(parsed.jobId) ?? "", + variant: getTrimmedString(parsed.variant) ?? defaultVariant, + generatedAt: getTrimmedString(parsed.generatedAt) ?? "", + namespace: getTrimmedString(parsed.namespace) ?? undefined, + jobResultId: getTrimmedString(parsed.jobResultId) ?? undefined, + assets, + } +} + +async function readJsonObject( + storage: KnowhereSdkStorage, + key: string, +): Promise { + if (!storage.readObject) return null + + const object = await storage.readObject(key).catch(() => null) + if (!object) return null + + try { + const parsed: unknown = JSON.parse(new TextDecoder().decode(object.body)) + return parsed + } catch { + return null + } +} + +function getResultJobId( + result: RetrievalResult, + source: Source | undefined, +): string | null { + const metadata = result.metadata + return ( + getRecordString(metadata, "jobId") ?? + getRecordString(metadata, "job_id") ?? + getTrimmedString(source?.knowhereJobId) ?? + null + ) +} + +function getPageNumbers( + metadata: Readonly> | undefined, +): readonly number[] { + if (!metadata) return [] + + const values = [metadata.pageNums, metadata.page_nums, metadata.pageNum] + const pageNumbers = new Set() + + for (const value of values) { + if (Array.isArray(value)) { + for (const item of value) { + const pageNum = getPositiveInteger(item) + if (pageNum) pageNumbers.add(pageNum) + } + continue + } + + const pageNum = getPositiveInteger(value) + if (pageNum) pageNumbers.add(pageNum) + } + + return [...pageNumbers].sort((left, right) => left - right) +} + +function createPageCitationAssetIndexKey( + documentId: string, + jobId: string, +): string { + return [ + "page-citation-assets", + "documents", + toSafeKeySegment(documentId), + "jobs", + toSafeKeySegment(jobId), + "variants", + defaultVariant, + "index.json", + ].join("/") +} + +function createPageCitationAssetCurrentIndexKey(documentId: string): string { + return [ + "page-citation-assets", + "documents", + toSafeKeySegment(documentId), + "current.json", + ].join("/") +} + +function toSafeKeySegment(value: string): string { + const normalized = value.replace(/[^A-Za-z0-9._-]+/g, "_").replace(/^_+|_+$/g, "") + if (normalized.length > 0 && normalized === value) return normalized + + const hash = createHash("sha256").update(value).digest("hex").slice(0, 12) + return `${(normalized || "value").slice(0, 48)}-${hash}` +} + +function getPageCitationMimeType(value: unknown): PageCitationAsset["mimeType"] | null { + return value === "image/png" || value === "image/jpeg" ? value : null +} + +function getRecordString( + record: Readonly> | undefined, + key: string, +): string | null { + if (!record) return null + return getTrimmedString(record[key]) +} + +function getTrimmedString(value: unknown): string | null { + if (typeof value !== "string") return null + const trimmed = value.trim() + return trimmed.length > 0 ? trimmed : null +} + +function getPositiveInteger(value: unknown): number | null { + return typeof value === "number" && + Number.isSafeInteger(value) && + value > 0 + ? value + : null +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null +} diff --git a/src/domains/chat/types.ts b/src/domains/chat/types.ts index f3f987d..a0469e0 100644 --- a/src/domains/chat/types.ts +++ b/src/domains/chat/types.ts @@ -1,12 +1,12 @@ /** * Chat citation / retrieval hit. Mirrors RetrievalResult from the SDK. - * No chunkId here; retrieval does not expose one. */ export type RetrievalResultView = { readonly content: string readonly chunkType: string readonly score: number | null readonly assetUrl?: string + readonly pageCitationAssetUrl?: string readonly source: { readonly documentId?: string | null readonly sourceFileName?: string | null diff --git a/src/domains/chat/view.ts b/src/domains/chat/view.ts index 9afc4f6..ed39295 100644 --- a/src/domains/chat/view.ts +++ b/src/domains/chat/view.ts @@ -51,6 +51,7 @@ function toPersistedCitationViews(value: unknown): ChatCitationView[] | undefine chunkType: getString(item.chunkType) ?? "text", score: getNumber(item.score) ?? 0, assetUrl: getString(item.assetUrl), + pageCitationAssetUrl: getString(item.pageCitationAssetUrl), description: getString(item.description), source: { documentId: getString(item.source.documentId), @@ -81,6 +82,9 @@ function toPersistedArtifactViews(value: unknown): ChatArtifactView[] | undefine chunkType: getString(item.citation.chunkType) ?? "text", score: getNumber(item.citation.score) ?? 0, assetUrl: getString(item.citation.assetUrl), + pageCitationAssetUrl: getString( + item.citation.pageCitationAssetUrl, + ), description: getString(item.citation.description), source: { documentId: getString(item.citation.source.documentId), diff --git a/src/domains/sources/page-citation-assets.test.ts b/src/domains/sources/page-citation-assets.test.ts new file mode 100644 index 0000000..d748929 --- /dev/null +++ b/src/domains/sources/page-citation-assets.test.ts @@ -0,0 +1,82 @@ +import { describe, expect, it, vi } from "vitest" +import type { + KnowhereSdkStorage, + KnowhereSdkStorageHead, + KnowhereSdkStorageObject, + KnowhereSdkStorageReadResult, + KnowhereSdkStorageWriteResult, + PageRenderer, + RenderedPage, +} from "@ontos-ai/knowhere-sdk" + +import { prepareSourcePageCitationAssets } from "./page-citation-assets" + +class MemoryStorage implements KnowhereSdkStorage { + headObject(): Promise { + return Promise.resolve(null) + } + + writeObject( + input: KnowhereSdkStorageObject, + ): Promise { + return Promise.resolve({ key: input.key }) + } + + readObject(): Promise { + return Promise.resolve(null) + } +} + +describe("prepareSourcePageCitationAssets", () => { + it("calls SDK cacheJobResult with tolerant bounded page citation options", async () => { + const storage = new MemoryStorage() + const renderer: PageRenderer & { close: () => Promise } = { + renderPage(): Promise { + return Promise.resolve({ + body: new Uint8Array([1]), + mimeType: "image/png", + width: 1, + height: 1, + }) + }, + close: vi.fn(async () => undefined), + } + const cacheJobResult = vi.fn(async () => ({ + pageCitationAssetWarnings: [ + { + code: renderLimitExceededCode, + message: "limit reached", + }, + ], + })) + + const result = await prepareSourcePageCitationAssets({ + client: { + knowledge: { + cacheJobResult, + }, + }, + sourceId: "source_1", + jobId: "job_1", + documentId: "doc_1", + storage, + renderer, + maxPagesToRenderPerRun: 3, + }) + + expect(cacheJobResult).toHaveBeenCalledWith({ + jobId: "job_1", + localDocumentId: "source_1", + pageCitationAssets: { + storage, + renderer, + strict: false, + maxPagesToRenderPerRun: 3, + }, + }) + expect(result.warnings).toHaveLength(1) + expect(renderer.close).toHaveBeenCalled() + }) +}) + +const renderLimitExceededCode = "render_limit_exceeded" as const diff --git a/src/domains/sources/page-citation-assets.ts b/src/domains/sources/page-citation-assets.ts new file mode 100644 index 0000000..7349bc4 --- /dev/null +++ b/src/domains/sources/page-citation-assets.ts @@ -0,0 +1,69 @@ +import "server-only" + +import { createPdfJsPageRenderer } from "@ontos-ai/knowhere-sdk/page-renderer-pdfjs" +import type { + KnowhereSdkStorage, + PageCitationAssetWarning, + PageCitationAssetsOptions, + PageRenderer, +} from "@ontos-ai/knowhere-sdk" + +import { createVercelBlobKnowhereSdkStorage } from "@/integrations/knowhere-sdk-storage" + +type SourcePageCitationAssetClient = { + readonly knowledge: { + readonly cacheJobResult: (input: { + readonly jobId: string + readonly localDocumentId: string + readonly pageCitationAssets: PageCitationAssetsOptions + }) => Promise<{ + readonly pageCitationAssetWarnings?: readonly PageCitationAssetWarning[] + }> + } +} + +export type PrepareSourcePageCitationAssetsInput = { + readonly client: SourcePageCitationAssetClient + readonly sourceId: string + readonly jobId: string + readonly documentId: string + readonly storage?: KnowhereSdkStorage + readonly renderer?: PageRenderer & { close?: () => Promise } + readonly maxPagesToRenderPerRun?: number +} + +export type PrepareSourcePageCitationAssetsResult = { + readonly warnings: readonly PageCitationAssetWarning[] +} + +const defaultMaxPagesToRenderPerRun = 25 + +export async function prepareSourcePageCitationAssets({ + client, + sourceId, + jobId, + storage = createVercelBlobKnowhereSdkStorage(), + renderer, + maxPagesToRenderPerRun = defaultMaxPagesToRenderPerRun, +}: PrepareSourcePageCitationAssetsInput): Promise { + const pageRenderer = renderer ?? createPdfJsPageRenderer({ maxThreads: 1 }) + + try { + const result = await client.knowledge.cacheJobResult({ + jobId, + localDocumentId: sourceId, + pageCitationAssets: { + storage, + renderer: pageRenderer, + strict: false, + maxPagesToRenderPerRun, + }, + }) + + return { + warnings: result.pageCitationAssetWarnings ?? [], + } + } finally { + await pageRenderer.close?.().catch(() => undefined) + } +} diff --git a/src/domains/sources/source-reconcile-route-workflow.test.ts b/src/domains/sources/source-reconcile-route-workflow.test.ts index 0841c54..a9e8a17 100644 --- a/src/domains/sources/source-reconcile-route-workflow.test.ts +++ b/src/domains/sources/source-reconcile-route-workflow.test.ts @@ -5,16 +5,23 @@ const mocks = vi.hoisted(() => ({ loggerInfo: vi.fn(), loggerWarn: vi.fn(), makeKnowhereClient: vi.fn(), + markSourceFailedAfterReconciliation: vi.fn(), markFailed: vi.fn(), markSourceReadyAfterReconciliation: vi.fn(), pollSourceReconciliation: vi.fn(), + prepareSourcePageCitationAssets: vi.fn(), })) vi.mock("@/domains/sources/source-reconcile-workflow", () => ({ + markSourceFailedAfterReconciliation: mocks.markSourceFailedAfterReconciliation, markSourceReadyAfterReconciliation: mocks.markSourceReadyAfterReconciliation, pollSourceReconciliation: mocks.pollSourceReconciliation, })) +vi.mock("@/domains/sources/page-citation-assets", () => ({ + prepareSourcePageCitationAssets: mocks.prepareSourcePageCitationAssets, +})) + vi.mock("@/domains/sources/workflow-runtime", () => ({ sourceWorkflowRuntime: { markFailed: mocks.markFailed, @@ -76,6 +83,16 @@ describe("sourceReconcileRouteWorkflow", () => { mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ status: "ready", }) + mocks.prepareSourcePageCitationAssets.mockResolvedValue({ + warnings: [ + { + code: "render_limit_exceeded", + message: "One page was skipped.", + documentId: "doc_1", + jobId: "job_1", + }, + ], + }) try { await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ @@ -95,9 +112,70 @@ describe("sourceReconcileRouteWorkflow", () => { sourceId: "source_1", documentId: "doc_1", }) + expect(mocks.prepareSourcePageCitationAssets).toHaveBeenCalledWith({ + client, + sourceId: "source_1", + jobId: "job_1", + documentId: "doc_1", + }) + expect( + mocks.prepareSourcePageCitationAssets.mock.invocationCallOrder[0], + ).toBeLessThan( + mocks.markSourceReadyAfterReconciliation.mock.invocationCallOrder[0]!, + ) + expect(mocks.loggerWarn).toHaveBeenCalledWith( + "workflow: page citation asset warning", + expect.objectContaining({ + sourceId: "source_1", + jobId: "job_1", + documentId: "doc_1", + }), + ) expect(continuations).toEqual([]) }) + it("marks the source failed when page citation asset preparation throws", async () => { + const context = createWorkflowContext() + const restore = + sourceReconcileRouteWorkflow.setContinuationTriggerForTesting( + async () => undefined, + ) + const client = { jobs: {}, knowledge: {} } + mocks.makeKnowhereClient.mockReturnValue(client) + mocks.pollSourceReconciliation.mockResolvedValue({ + kind: "ready-to-prepare", + jobId: "job_1", + documentId: "doc_1", + }) + mocks.prepareSourcePageCitationAssets.mockRejectedValue( + new Error("renderer setup failed"), + ) + mocks.markSourceFailedAfterReconciliation.mockResolvedValue({ + status: "failed", + }) + + try { + await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ + context, + payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + }), + }) + } finally { + restore() + } + + expect(mocks.markSourceFailedAfterReconciliation).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "source_1", + reason: + "Page citation asset preparation failed: renderer setup failed", + }) + expect(mocks.markSourceReadyAfterReconciliation).not.toHaveBeenCalled() + }) + it("triggers a fresh poll run when Knowhere is still running after the segment budget", async () => { const context = createWorkflowContext() const continuations: ContinuationTriggerInput[] = [] diff --git a/src/domains/sources/source-reconcile-route-workflow.ts b/src/domains/sources/source-reconcile-route-workflow.ts index a1e732a..f195250 100644 --- a/src/domains/sources/source-reconcile-route-workflow.ts +++ b/src/domains/sources/source-reconcile-route-workflow.ts @@ -3,11 +3,13 @@ import "server-only" import { Client, type WorkflowContext } from "@upstash/workflow" import { + markSourceFailedAfterReconciliation, markSourceReadyAfterReconciliation, pollSourceReconciliation, } from "@/domains/sources/source-reconcile-workflow" import { makeKnowhereClient } from "@/integrations/knowhere" import { logger } from "@/lib/logger" +import { prepareSourcePageCitationAssets } from "./page-citation-assets" import { sourceWorkflowRuntime } from "./workflow-runtime" type ReconcilePayload = { @@ -39,6 +41,19 @@ type ContinuationTriggerInput = { readonly workflowRunId: string } +type PageCitationAssetPreparationStepResult = + | { + readonly ok: true + readonly warnings: Awaited< + ReturnType + >["warnings"] + } + | { + readonly ok: false + readonly reason: string + readonly error: string + } + const maxPollAttempts = 25 const initialDelaySeconds = 3 const maxDelaySeconds = 30 @@ -118,6 +133,16 @@ async function runPollAndMirrorWorkflow(input: { return } + const preparation = await runPageCitationAssetPreparation({ + context, + workspaceId, + sourceId, + client, + jobId: jobToPrepare.jobId, + documentId: jobToPrepare.documentId, + }) + if (!preparation.ok) return + const ready = await context.run("source-ready", async () => markSourceReadyAfterReconciliation({ workspaceId, @@ -132,6 +157,79 @@ async function runPollAndMirrorWorkflow(input: { }) } +async function runPageCitationAssetPreparation(input: { + readonly context: ReconcileWorkflowContext + readonly workspaceId: string + readonly sourceId: string + readonly client: ReturnType + readonly jobId: string + readonly documentId: string +}): Promise<{ readonly ok: boolean }> { + const result = await input.context.run( + "prepare-page-citation-assets", + async (): Promise => { + try { + const preparation = await prepareSourcePageCitationAssets({ + client: input.client, + sourceId: input.sourceId, + jobId: input.jobId, + documentId: input.documentId, + }) + + return { + ok: true, + warnings: preparation.warnings, + } + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error) + return { + ok: false, + reason: `Page citation asset preparation failed: ${getSafeFailureReason( + errorMessage, + )}`, + error: errorMessage, + } + } + }, + ) + + if (!result.ok) { + const failed = await input.context.run( + "source-failed-page-citation-assets", + async () => + markSourceFailedAfterReconciliation({ + workspaceId: input.workspaceId, + sourceId: input.sourceId, + reason: result.reason, + }), + ) + logger.error("workflow: page citation asset preparation failed", { + sourceId: input.sourceId, + jobId: input.jobId, + documentId: input.documentId, + status: failed.status, + error: result.error, + }) + return { ok: false } + } + + for (const warning of result.warnings) { + logger.warn("workflow: page citation asset warning", { + sourceId: input.sourceId, + jobId: input.jobId, + documentId: input.documentId, + warning, + }) + } + logger.info("workflow: page citation assets prepared", { + sourceId: input.sourceId, + jobId: input.jobId, + documentId: input.documentId, + warningCount: result.warnings.length, + }) + return { ok: true } +} + function normalizeReconcilePayload( payload: ReconcilePayload, ): NormalizedReconcilePayload { diff --git a/src/domains/sources/source-reconcile-workflow.test.ts b/src/domains/sources/source-reconcile-workflow.test.ts index 3f0632e..f1907ef 100644 --- a/src/domains/sources/source-reconcile-workflow.test.ts +++ b/src/domains/sources/source-reconcile-workflow.test.ts @@ -3,6 +3,7 @@ import type { JobResult } from "@ontos-ai/knowhere-sdk" import type { Source, Workspace } from "@/infrastructure/db/schema" import { + markSourceFailedAfterReconciliation, markSourceReadyAfterReconciliation, pollSourceReconciliation, } from "./source-reconcile-workflow" @@ -197,3 +198,34 @@ describe("markSourceReadyAfterReconciliation", () => { ) }) }) + +describe("markSourceFailedAfterReconciliation", () => { + it("marks failed and cleans staged uploads after fatal preparation errors", async () => { + const source = makeSource({ + stagedBlobPathname: "source-uploads/upload_1/document.pdf", + }) + const repository = createRepository(source) + const blobStore = { + deleteStagedSourceBlob: vi.fn(async () => undefined), + } + + const result = await markSourceFailedAfterReconciliation({ + workspaceId: workspace.id, + sourceId: "source_1", + reason: "Page citation asset preparation failed.", + repository, + blobStore, + }) + + expect(result).toEqual({ status: "failed" }) + expect(repository.markFailed).toHaveBeenCalledWith( + workspace.id, + "source_1", + "Page citation asset preparation failed.", + "parsing", + ) + expect(blobStore.deleteStagedSourceBlob).toHaveBeenCalledWith( + "source-uploads/upload_1/document.pdf", + ) + }) +}) diff --git a/src/domains/sources/source-reconcile-workflow.ts b/src/domains/sources/source-reconcile-workflow.ts index 0355350..4571038 100644 --- a/src/domains/sources/source-reconcile-workflow.ts +++ b/src/domains/sources/source-reconcile-workflow.ts @@ -55,6 +55,14 @@ export type MarkSourceReadyAfterReconciliationInput = { readonly blobStore?: SourceReconcileWorkflowBlobStore } +export type MarkSourceFailedAfterReconciliationInput = { + readonly workspaceId: string + readonly sourceId: string + readonly reason: string + readonly repository?: SourceReconcileWorkflowRepository + readonly blobStore?: SourceReconcileWorkflowBlobStore +} + export type PollSourceReconciliationResult = | { readonly kind: "waiting" @@ -75,6 +83,10 @@ export type MarkSourceReadyAfterReconciliationResult = { readonly status: string } +export type MarkSourceFailedAfterReconciliationResult = { + readonly status: string +} + export async function pollSourceReconciliation({ workspaceId, sourceId, @@ -165,6 +177,27 @@ export async function markSourceReadyAfterReconciliation({ return { status: readySource.status } } +export async function markSourceFailedAfterReconciliation({ + workspaceId, + sourceId, + reason, + repository = sourceWorkflowRuntime, + blobStore = vercelBlobStore, +}: MarkSourceFailedAfterReconciliationInput): Promise { + const source = await repository.findInWorkspace(workspaceId, sourceId) + if (!source) return { status: "gone" } + if (source.status !== "parsing") return { status: source.status } + + await failSourceAndCleanup({ + workspaceId, + source, + reason, + repository, + blobStore, + }) + return { status: "failed" } +} + async function failSourceAndCleanup(input: { readonly workspaceId: string readonly source: Source diff --git a/src/domains/workspace/integration.test.ts b/src/domains/workspace/integration.test.ts index 797cdbf..9d27083 100644 --- a/src/domains/workspace/integration.test.ts +++ b/src/domains/workspace/integration.test.ts @@ -331,6 +331,7 @@ describeIfDb("workspace helpers — integration", () => { chunkType: "text", score: 0.91, assetUrl: "https://assets.example/doc.pdf", + pageCitationAssetUrl: "https://assets.example/page-4.png", description: "intro summary", source: { documentId: "doc_123", @@ -353,6 +354,7 @@ describeIfDb("workspace helpers — integration", () => { chunkType: "text", score: 0.91, assetUrl: "https://assets.example/doc.pdf", + pageCitationAssetUrl: "https://assets.example/page-4.png", description: "intro summary", source: { documentId: "doc_123", diff --git a/src/domains/workspace/persistence.test.ts b/src/domains/workspace/persistence.test.ts index 47534d1..1c9be69 100644 --- a/src/domains/workspace/persistence.test.ts +++ b/src/domains/workspace/persistence.test.ts @@ -207,6 +207,7 @@ describe("chatRepository", () => { chunkType: "text", score: 0.99, assetUrl: "https://assets.example/doc.pdf", + pageCitationAssetUrl: "https://assets.example/page-4.png", description: "intro", source: { documentId: "doc_1", @@ -225,6 +226,7 @@ describe("chatRepository", () => { chunkType: "text", score: 0.99, assetUrl: "https://assets.example/doc.pdf", + pageCitationAssetUrl: "https://assets.example/page-4.png", description: "intro", source: { documentId: "doc_1", diff --git a/src/integrations/knowhere-sdk-storage.test.ts b/src/integrations/knowhere-sdk-storage.test.ts new file mode 100644 index 0000000..c5a19b6 --- /dev/null +++ b/src/integrations/knowhere-sdk-storage.test.ts @@ -0,0 +1,320 @@ +import { BlobNotFoundError, type HeadBlobResult } from "@vercel/blob" +import { describe, expect, it } from "vitest" + +import { + createVercelBlobKnowhereSdkStorage, + type VercelBlobStore, +} from "./knowhere-sdk-storage" + +type StoredBlob = { + readonly pathname: string + readonly body: Uint8Array + readonly contentType?: string + readonly url: string +} + +class MemoryVercelBlobStore implements VercelBlobStore { + readonly objects = new Map() + + async put( + pathname: string, + body: Parameters[1], + options: Parameters[2], + ): Promise<{ readonly pathname: string; readonly url: string }> { + const url = `https://blob.example/${pathname}` + this.objects.set(pathname, { + pathname, + body: await readBody(body), + contentType: options.contentType, + url, + }) + return { pathname, url } + } + + async head(pathname: string): Promise { + const object = this.objects.get(pathname) + if (!object) throw new BlobNotFoundError() + + return { + size: object.body.byteLength, + uploadedAt: new Date("2026-07-03T00:00:00.000Z"), + pathname, + contentType: object.contentType ?? "application/octet-stream", + contentDisposition: "", + url: object.url, + downloadUrl: `${object.url}?download=1`, + cacheControl: "", + etag: `etag:${pathname}`, + } + } + + async get( + pathname: string, + ): Promise>> { + const object = this.objects.get(pathname) + if (!object) return null + + return { + statusCode: 200, + stream: toReadableStream(object.body), + headers: new Headers(), + blob: { + url: object.url, + downloadUrl: `${object.url}?download=1`, + pathname, + contentDisposition: "", + cacheControl: "", + uploadedAt: new Date("2026-07-03T00:00:00.000Z"), + etag: `etag:${pathname}`, + contentType: object.contentType ?? "application/octet-stream", + size: object.body.byteLength, + }, + } + } + + async list(input: { + readonly prefix: string + readonly limit: number + readonly cursor?: string + }): Promise>> { + const start = input.cursor ? Number(input.cursor) : 0 + const matches = [...this.objects.values()] + .filter((object) => object.pathname.startsWith(input.prefix)) + .sort((left, right) => left.pathname.localeCompare(right.pathname)) + const page = matches.slice(start, start + input.limit) + const nextIndex = start + page.length + const hasMore = nextIndex < matches.length + + return { + blobs: page.map((object) => ({ + url: object.url, + downloadUrl: `${object.url}?download=1`, + pathname: object.pathname, + size: object.body.byteLength, + uploadedAt: new Date("2026-07-03T00:00:00.000Z"), + etag: `etag:${object.pathname}`, + })), + hasMore, + ...(hasMore ? { cursor: String(nextIndex) } : {}), + } + } + + async del(pathname: string | readonly string[]): Promise { + const pathnames = typeof pathname === "string" ? [pathname] : [...pathname] + for (const item of pathnames) { + this.objects.delete(item) + } + } +} + +describe("createVercelBlobKnowhereSdkStorage", () => { + it("rejects unsafe storage keys", async () => { + const storage = createVercelBlobKnowhereSdkStorage( + new MemoryVercelBlobStore(), + ) + + for (const key of [ + "", + "/absolute/path.png", + "nested//empty.png", + "nested/../escape.png", + "nested/./dot.png", + "nested\\windows.png", + "nested/\0/null.png", + ]) { + await expect( + storage.writeObject({ + key, + body: new Uint8Array([1]), + contentType: "image/png", + }), + ).rejects.toThrow("Invalid Knowhere SDK storage key") + } + }) + + it("writes objects with metadata sidecars and returns blob URLs", async () => { + const blobStore = new MemoryVercelBlobStore() + const storage = createVercelBlobKnowhereSdkStorage(blobStore) + + const result = await storage.writeObject({ + key: "page-citation-assets/documents/doc_1/page.png", + body: new Uint8Array([1, 2, 3]), + contentType: "image/png", + metadata: { + width: "1200", + height: "1600", + mimeType: "image/png", + }, + }) + + expect(result).toEqual({ + key: "page-citation-assets/documents/doc_1/page.png", + url: "https://blob.example/page-citation-assets/documents/doc_1/page.png", + }) + expect( + blobStore.objects.has("page-citation-assets/documents/doc_1/page.png"), + ).toBe(true) + const metadataSidecar = blobStore.objects.get( + "page-citation-assets/documents/doc_1/page.png.metadata.json", + ) + expect(metadataSidecar).toBeDefined() + if (!metadataSidecar) throw new Error("Expected metadata sidecar.") + + expect(JSON.parse(decodeBody(metadataSidecar.body))).toEqual({ + metadata: { + width: "1200", + height: "1600", + mimeType: "image/png", + }, + }) + }) + + it("combines blob head data with sidecar metadata", async () => { + const blobStore = new MemoryVercelBlobStore() + const storage = createVercelBlobKnowhereSdkStorage(blobStore) + + await storage.writeObject({ + key: "page-citation-assets/documents/doc_1/page.png", + body: new Uint8Array([1, 2, 3]), + contentType: "image/png", + metadata: { + width: "1200", + height: "1600", + }, + }) + + await expect( + storage.headObject("page-citation-assets/documents/doc_1/page.png"), + ).resolves.toEqual({ + key: "page-citation-assets/documents/doc_1/page.png", + contentType: "image/png", + contentLength: 3, + metadata: { + width: "1200", + height: "1600", + }, + }) + }) + + it("reads objects, resolves URLs, deletes objects, and deletes prefixes", async () => { + const blobStore = new MemoryVercelBlobStore() + const storage = createVercelBlobKnowhereSdkStorage(blobStore) + + await storage.writeObject({ + key: "prefix/one.png", + body: new Uint8Array([1]), + contentType: "image/png", + }) + await storage.writeObject({ + key: "prefix/two.png", + body: new Uint8Array([2]), + contentType: "image/png", + }) + await storage.writeObject({ + key: "prefix-sibling/two.png", + body: new Uint8Array([4]), + contentType: "image/png", + }) + await storage.writeObject({ + key: "other/three.png", + body: new Uint8Array([3]), + contentType: "image/png", + }) + + const readObject = storage.readObject + const getObjectUrl = storage.getObjectUrl + const deleteObject = storage.deleteObject + const deletePrefix = storage.deletePrefix + expect(readObject).toBeDefined() + expect(getObjectUrl).toBeDefined() + expect(deleteObject).toBeDefined() + expect(deletePrefix).toBeDefined() + if (!readObject || !getObjectUrl || !deleteObject || !deletePrefix) { + throw new Error("Expected full storage adapter methods.") + } + + await expect(storage.readObject?.("prefix/one.png")).resolves.toMatchObject({ + body: new Uint8Array([1]), + contentType: "image/png", + }) + await expect(storage.getObjectUrl?.("prefix/one.png")).resolves.toBe( + "https://blob.example/prefix/one.png", + ) + + await storage.deleteObject?.("prefix/one.png") + expect(blobStore.objects.has("prefix/one.png")).toBe(false) + expect(blobStore.objects.has("prefix/one.png.metadata.json")).toBe(false) + + await storage.deletePrefix?.("prefix") + expect([...blobStore.objects.keys()].sort()).toEqual([ + "other/three.png", + "other/three.png.metadata.json", + "prefix-sibling/two.png", + "prefix-sibling/two.png.metadata.json", + ]) + }) + + it("accepts trailing slashes when deleting prefixes", async () => { + const blobStore = new MemoryVercelBlobStore() + const storage = createVercelBlobKnowhereSdkStorage(blobStore) + + await storage.writeObject({ + key: "prefix/one.png", + body: new Uint8Array([1]), + contentType: "image/png", + }) + + await storage.deletePrefix?.("prefix/") + + expect(blobStore.objects.size).toBe(0) + }) +}) + +async function readBody( + body: Parameters[1], +): Promise { + if (typeof body === "string") return new TextEncoder().encode(body) + if (body instanceof Uint8Array) return new Uint8Array(body) + if (body instanceof Blob) return new Uint8Array(await body.arrayBuffer()) + return readStream(body) +} + +async function readStream( + stream: ReadableStream, +): Promise { + const reader = stream.getReader() + const chunks: Uint8Array[] = [] + let totalLength = 0 + + try { + while (true) { + const result = await reader.read() + if (result.done) break + chunks.push(result.value) + totalLength += result.value.byteLength + } + } finally { + reader.releaseLock() + } + + const body = new Uint8Array(totalLength) + let offset = 0 + for (const chunk of chunks) { + body.set(chunk, offset) + offset += chunk.byteLength + } + return body +} + +function toReadableStream(body: Uint8Array): ReadableStream { + return new ReadableStream({ + start(controller) { + controller.enqueue(body) + controller.close() + }, + }) +} + +function decodeBody(body: Uint8Array): string { + return new TextDecoder().decode(body) +} diff --git a/src/integrations/knowhere-sdk-storage.ts b/src/integrations/knowhere-sdk-storage.ts new file mode 100644 index 0000000..387bc21 --- /dev/null +++ b/src/integrations/knowhere-sdk-storage.ts @@ -0,0 +1,294 @@ +import "server-only" + +import { Buffer } from "node:buffer" +import { + BlobNotFoundError, + del, + get, + head, + list, + put, + type GetBlobResult, + type HeadBlobResult, + type ListBlobResult, + type PutBlobResult, +} from "@vercel/blob" +import type { + KnowhereSdkStorage, + KnowhereSdkStorageBody, + KnowhereSdkStorageHead, + KnowhereSdkStorageObject, + KnowhereSdkStorageReadResult, + KnowhereSdkStorageWriteResult, +} from "@ontos-ai/knowhere-sdk" + +type VercelBlobBody = Buffer | Blob | ReadableStream + +export type VercelBlobStore = { + readonly put: ( + pathname: string, + body: VercelBlobBody | string, + options: { + readonly access: "public" + readonly addRandomSuffix: false + readonly allowOverwrite: true + readonly contentType?: string + }, + ) => Promise> + readonly head: (pathname: string) => Promise + readonly get: ( + pathname: string, + options: { readonly access: "public" }, + ) => Promise + readonly list: (options: { + readonly prefix: string + readonly limit: number + readonly cursor?: string + }) => Promise + readonly del: (pathname: string | readonly string[]) => Promise +} + +type StoredMetadataSidecar = { + readonly metadata?: Readonly> +} + +const metadataSuffix = ".metadata.json" +const deletePrefixPageSize = 1000 + +const vercelBlobStore: VercelBlobStore = { + put, + head, + get, + list, + del: (pathname) => + del(typeof pathname === "string" ? pathname : [...pathname]), +} + +class VercelBlobKnowhereSdkStorage implements KnowhereSdkStorage { + private readonly blobStore: VercelBlobStore + + constructor(blobStore: VercelBlobStore = vercelBlobStore) { + this.blobStore = blobStore + } + + async headObject(key: string): Promise { + validateBlobStorageKey(key) + + try { + const object = await this.blobStore.head(key) + const metadata = await this.readMetadata(key) + return { + key, + contentType: object.contentType, + contentLength: object.size, + metadata, + } + } catch (error) { + if (isBlobNotFoundError(error)) return null + throw error + } + } + + async writeObject( + input: KnowhereSdkStorageObject, + ): Promise { + validateBlobStorageKey(input.key) + + const object = await this.blobStore.put( + input.key, + toVercelBlobBody(input.body), + { + access: "public", + addRandomSuffix: false, + allowOverwrite: true, + contentType: input.contentType, + }, + ) + await this.writeMetadata(input.key, input.metadata ?? {}) + + return { + key: object.pathname, + url: object.url, + } + } + + async getObjectUrl(key: string): Promise { + validateBlobStorageKey(key) + + try { + const object = await this.blobStore.head(key) + return object.url + } catch (error) { + if (isBlobNotFoundError(error)) return null + throw error + } + } + + async readObject(key: string): Promise { + validateBlobStorageKey(key) + + const object = await this.blobStore.get(key, { access: "public" }) + if (!object || object.statusCode !== 200) return null + + return { + body: await readStream(object.stream), + contentType: object.blob.contentType, + metadata: await this.readMetadata(key), + } + } + + async deleteObject(key: string): Promise { + validateBlobStorageKey(key) + await this.deleteIgnoringMissing([key, getMetadataKey(key)]) + } + + async deletePrefix(prefix: string): Promise { + validateBlobStoragePrefix(prefix) + const blobPrefix = prefix.endsWith("/") ? prefix : `${prefix}/` + + let cursor: string | undefined + do { + const page = await this.blobStore.list({ + prefix: blobPrefix, + limit: deletePrefixPageSize, + ...(cursor ? { cursor } : {}), + }) + const pathnames = page.blobs.map((blob) => blob.pathname) + if (pathnames.length > 0) await this.deleteIgnoringMissing(pathnames) + cursor = page.cursor + if (page.hasMore && !cursor) break + } while (cursor) + } + + private async readMetadata( + key: string, + ): Promise>> { + const metadataObject = await this.blobStore.get(getMetadataKey(key), { + access: "public", + }) + if (!metadataObject || metadataObject.statusCode !== 200) return {} + + const body = await readStream(metadataObject.stream) + const parsed: unknown = JSON.parse(new TextDecoder().decode(body)) + return parseStoredMetadataSidecar(parsed).metadata ?? {} + } + + private async writeMetadata( + key: string, + metadata: Readonly>, + ): Promise { + const sidecar: StoredMetadataSidecar = { metadata } + await this.blobStore.put( + getMetadataKey(key), + JSON.stringify(sidecar, null, 2), + { + access: "public", + addRandomSuffix: false, + allowOverwrite: true, + contentType: "application/json", + }, + ) + } + + private async deleteIgnoringMissing( + pathnames: readonly string[], + ): Promise { + try { + await this.blobStore.del(pathnames) + } catch (error) { + if (!isBlobNotFoundError(error)) throw error + } + } +} + +export function createVercelBlobKnowhereSdkStorage( + blobStore?: VercelBlobStore, +): KnowhereSdkStorage { + return new VercelBlobKnowhereSdkStorage(blobStore) +} + +function validateBlobStorageKey(key: string): void { + if (!isSafeRelativePosixPath(key)) { + throw new Error(`Invalid Knowhere SDK storage key: ${key}`) + } +} + +function validateBlobStoragePrefix(prefix: string): void { + const normalizedPrefix = prefix.endsWith("/") ? prefix.slice(0, -1) : prefix + if (!isSafeRelativePosixPath(normalizedPrefix)) { + throw new Error(`Invalid Knowhere SDK storage prefix: ${prefix}`) + } +} + +function isSafeRelativePosixPath(value: string): boolean { + if ( + value.length === 0 || + value.startsWith("/") || + value.includes("\\") || + value.includes("\0") + ) { + return false + } + + return value + .split("/") + .every((segment) => segment.length > 0 && segment !== "." && segment !== "..") +} + +function getMetadataKey(key: string): string { + return `${key}${metadataSuffix}` +} + +function toVercelBlobBody(body: KnowhereSdkStorageBody): VercelBlobBody { + if (body instanceof Uint8Array) return Buffer.from(body) + return body +} + +async function readStream( + stream: ReadableStream, +): Promise { + const reader = stream.getReader() + const chunks: Uint8Array[] = [] + let totalLength = 0 + + try { + while (true) { + const result = await reader.read() + if (result.done) break + chunks.push(result.value) + totalLength += result.value.byteLength + } + } finally { + reader.releaseLock() + } + + const body = new Uint8Array(totalLength) + let offset = 0 + for (const chunk of chunks) { + body.set(chunk, offset) + offset += chunk.byteLength + } + return body +} + +function parseStoredMetadataSidecar(value: unknown): StoredMetadataSidecar { + if (!isRecord(value)) return {} + const metadata = value.metadata + if (!isRecord(metadata)) return {} + + return { + metadata: Object.fromEntries( + Object.entries(metadata).filter( + (entry): entry is [string, string] => typeof entry[1] === "string", + ), + ), + } +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null +} + +function isBlobNotFoundError(error: unknown): boolean { + return error instanceof BlobNotFoundError +} From a44e5227651da7cf3b8c28b3287bd3de9a1a3a01 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 3 Jul 2026 17:40:37 +0800 Subject: [PATCH 02/51] Consume server page citation assets --- src/domains/chat/page-citation-assets.test.ts | 144 +++----- src/domains/chat/page-citation-assets.ts | 245 +------------- .../sources/page-citation-assets.test.ts | 82 ----- src/domains/sources/page-citation-assets.ts | 69 ---- .../source-reconcile-route-workflow.test.ts | 80 +---- .../source-reconcile-route-workflow.ts | 98 ------ src/integrations/knowhere-sdk-storage.test.ts | 320 ------------------ src/integrations/knowhere-sdk-storage.ts | 294 ---------------- 8 files changed, 55 insertions(+), 1277 deletions(-) delete mode 100644 src/domains/sources/page-citation-assets.test.ts delete mode 100644 src/domains/sources/page-citation-assets.ts delete mode 100644 src/integrations/knowhere-sdk-storage.test.ts delete mode 100644 src/integrations/knowhere-sdk-storage.ts diff --git a/src/domains/chat/page-citation-assets.test.ts b/src/domains/chat/page-citation-assets.test.ts index c76b385..f7a1812 100644 --- a/src/domains/chat/page-citation-assets.test.ts +++ b/src/domains/chat/page-citation-assets.test.ts @@ -1,124 +1,90 @@ import { describe, expect, it } from "vitest" -import type { - KnowhereSdkStorage, - KnowhereSdkStorageHead, - KnowhereSdkStorageObject, - KnowhereSdkStorageReadResult, - KnowhereSdkStorageWriteResult, - RetrievalResult, -} from "@ontos-ai/knowhere-sdk" +import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" import { enrichRetrievalResultsWithPageCitationAssetUrls } from "./page-citation-assets" -class MemorySdkStorage implements KnowhereSdkStorage { - private readonly objects = new Map() - - constructor(objects: Readonly>) { - for (const [key, value] of Object.entries(objects)) { - this.objects.set(key, new TextEncoder().encode(value)) - } - } - - headObject(): Promise { - return Promise.resolve(null) - } - - writeObject( - input: KnowhereSdkStorageObject, - ): Promise { - if (input.body instanceof Uint8Array) { - this.objects.set(input.key, new Uint8Array(input.body)) - } - return Promise.resolve({ key: input.key }) - } - - readObject(key: string): Promise { - const body = this.objects.get(key) - return Promise.resolve(body ? { body } : null) - } - - getObjectUrl(key: string): Promise { - return Promise.resolve(`https://blob.example/${key}`) - } -} - describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { - it("adds a page citation asset URL for page results with matching metadata", async () => { - const storage = new MemorySdkStorage({ - "page-citation-assets/documents/doc_1/current.json": JSON.stringify({ - version: 1, - documentId: "doc_1", - jobId: "job_1", - variant: "default", - indexKey: - "page-citation-assets/documents/doc_1/jobs/job_1/variants/default/index.json", - updatedAt: "2026-07-03T00:00:00.000Z", - }), - "page-citation-assets/documents/doc_1/jobs/job_1/variants/default/index.json": - JSON.stringify({ - version: 1, - documentId: "doc_1", - jobId: "job_1", - variant: "default", - generatedAt: "2026-07-03T00:00:00.000Z", - assets: [ - { - pageNum: 4, - key: "page-citation-assets/documents/doc_1/jobs/job_1/variants/default/scale-1/page-000004.png", - assetUrl: "https://blob.example/pages/page-000004.png", - mimeType: "image/png", - width: 1200, - height: 1600, - source: "client-rendered-pdf-page", - variant: "default", - }, - ], + it("uses server-provided page asset URLs from result metadata", async () => { + const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ + results: [ + makeRetrievalResult({ + chunkType: "page", + metadata: { + pageNums: [2], + pageAssets: [ + { + pageNum: 2, + artifactRef: "page_citation_assets/page-2.png", + assetUrl: "https://assets.example/pages/page-2.png", + }, + ], + }, }), + ], + sources: [makeSource()], }) + expect(result?.pageCitationAssetUrl).toBe( + "https://assets.example/pages/page-2.png", + ) + }) + + it("uses server-provided top-level page assets", async () => { const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ results: [ makeRetrievalResult({ chunkType: "page", - metadata: { pageNums: [4] }, - }), + pageAssets: [ + { + pageNum: 4, + artifactRef: "page_citation_assets/page-4.png", + assetUrl: "https://assets.example/pages/page-4.png", + contentType: "image/png", + source: "knowhere-rendered-page-citation-source", + }, + ], + } as Partial), ], sources: [makeSource()], - storage, }) expect(result?.pageCitationAssetUrl).toBe( - "https://blob.example/pages/page-000004.png", + "https://assets.example/pages/page-4.png", ) }) - it("uses SDK-enriched page asset metadata before reading the index", async () => { + it("chooses the asset matching the citation page metadata", async () => { const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ results: [ makeRetrievalResult({ chunkType: "page", metadata: { + pageNums: [4], pageAssets: [ { pageNum: 2, - key: "page-key", - assetUrl: "https://blob.example/pages/page-000002.png", + artifactRef: "page_citation_assets/page-2.png", + assetUrl: "https://assets.example/pages/page-2.png", + }, + { + pageNum: 4, + artifactRef: "page_citation_assets/page-4.png", + assetUrl: "https://assets.example/pages/page-4.png", }, ], }, }), ], sources: [makeSource()], - storage: new MemorySdkStorage({}), }) expect(result?.pageCitationAssetUrl).toBe( - "https://blob.example/pages/page-000002.png", + "https://assets.example/pages/page-4.png", ) }) - it("chooses the SDK-enriched asset matching the citation page metadata", async () => { + it("does not invent a link when the server only provides an artifact ref", async () => { const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ results: [ makeRetrievalResult({ @@ -126,27 +92,18 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { metadata: { pageNums: [4], pageAssets: [ - { - pageNum: 2, - key: "page-2-key", - assetUrl: "https://blob.example/pages/page-000002.png", - }, { pageNum: 4, - key: "page-4-key", - assetUrl: "https://blob.example/pages/page-000004.png", + artifactRef: "page_citation_assets/page-4.png", }, ], }, }), ], sources: [makeSource()], - storage: new MemorySdkStorage({}), }) - expect(result?.pageCitationAssetUrl).toBe( - "https://blob.example/pages/page-000004.png", - ) + expect(result?.pageCitationAssetUrl).toBeUndefined() }) it("leaves non-page results unchanged even when they have page metadata", async () => { @@ -159,15 +116,14 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { pageAssets: [ { pageNum: 4, - key: "page-key", - assetUrl: "https://blob.example/pages/page-000004.png", + artifactRef: "page_citation_assets/page-4.png", + assetUrl: "https://assets.example/pages/page-4.png", }, ], }, }), ], sources: [makeSource()], - storage: new MemorySdkStorage({}), }) expect(result?.pageCitationAssetUrl).toBeUndefined() diff --git a/src/domains/chat/page-citation-assets.ts b/src/domains/chat/page-citation-assets.ts index dd30a8d..dccdd88 100644 --- a/src/domains/chat/page-citation-assets.ts +++ b/src/domains/chat/page-citation-assets.ts @@ -1,16 +1,8 @@ import "server-only" -import { createHash } from "node:crypto" -import type { - KnowhereSdkStorage, - PageCitationAsset, - PageCitationAssetCurrentIndex, - PageCitationAssetIndex, - RetrievalResult, -} from "@ontos-ai/knowhere-sdk" +import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" -import { createVercelBlobKnowhereSdkStorage } from "@/integrations/knowhere-sdk-storage" export type PageCitationAssetRetrievalResult = RetrievalResult & { readonly pageCitationAssetUrl?: string @@ -19,42 +11,26 @@ export type PageCitationAssetRetrievalResult = RetrievalResult & { type EnrichRetrievalResultsWithPageCitationAssetUrlsInput = { readonly results: readonly RetrievalResult[] readonly sources: readonly Source[] - readonly storage?: KnowhereSdkStorage } type PageCitationAssetCandidate = { readonly pageNum: number readonly assetUrl?: string - readonly key?: string } -const defaultVariant = "default" - export async function enrichRetrievalResultsWithPageCitationAssetUrls({ results, sources, - storage = createVercelBlobKnowhereSdkStorage(), }: EnrichRetrievalResultsWithPageCitationAssetUrlsInput): Promise< PageCitationAssetRetrievalResult[] > { + void sources if (results.length === 0) return [] - const sourcesByDocumentId = new Map( - sources.flatMap((source): readonly [string, Source][] => - source.knowhereDocumentId ? [[source.knowhereDocumentId, source]] : [], - ), - ) - const indexByKey = new Map>() - return Promise.all( results.map((result) => enrichRetrievalResultWithPageCitationAssetUrl({ result, - storage, - source: result.source.documentId - ? sourcesByDocumentId.get(result.source.documentId) - : undefined, - indexByKey, }), ), ) @@ -62,9 +38,6 @@ export async function enrichRetrievalResultsWithPageCitationAssetUrls({ async function enrichRetrievalResultWithPageCitationAssetUrl(input: { readonly result: RetrievalResult - readonly source?: Source - readonly storage: KnowhereSdkStorage - readonly indexByKey: Map> }): Promise { if (!isPageResult(input.result)) return input.result @@ -77,35 +50,7 @@ async function enrichRetrievalResultWithPageCitationAssetUrl(input: { } } - if (pageNumbers.length === 0) return input.result - - const documentId = getTrimmedString(input.result.source.documentId) - if (!documentId) return input.result - - const index = await resolvePageCitationAssetIndex({ - storage: input.storage, - documentId, - jobId: getResultJobId(input.result, input.source), - indexByKey: input.indexByKey, - }) - if (!index) return input.result - - const asset = index.assets.find((candidate) => - pageNumbers.includes(candidate.pageNum), - ) - if (!asset) return input.result - - const assetUrl = - getTrimmedString(asset.assetUrl) ?? - (asset.key && input.storage.getObjectUrl - ? await input.storage.getObjectUrl(asset.key).catch(() => null) - : null) - if (!assetUrl) return input.result - - return { - ...input.result, - pageCitationAssetUrl: assetUrl, - } + return input.result } function isPageResult(result: RetrievalResult): boolean { @@ -149,148 +94,9 @@ function parsePageCitationAssetCandidates( { pageNum, assetUrl: getTrimmedString(item.assetUrl) ?? undefined, - key: getTrimmedString(item.key) ?? undefined, - }, - ] - }) -} - -async function resolvePageCitationAssetIndex(input: { - readonly storage: KnowhereSdkStorage - readonly documentId: string - readonly jobId: string | null - readonly indexByKey: Map> -}): Promise { - if (!input.storage.readObject) return null - - if (input.jobId) { - const jobIndexKey = createPageCitationAssetIndexKey( - input.documentId, - input.jobId, - ) - const jobIndex = await readCachedIndex( - input.storage, - jobIndexKey, - input.indexByKey, - ) - if (jobIndex) return jobIndex - } - - const current = await readCurrentIndexPointer(input.storage, input.documentId) - if (!current) return null - - return readCachedIndex(input.storage, current.indexKey, input.indexByKey) -} - -async function readCachedIndex( - storage: KnowhereSdkStorage, - key: string, - cache: Map>, -): Promise { - let cached = cache.get(key) - if (!cached) { - cached = readPageCitationAssetIndex(storage, key) - cache.set(key, cached) - } - return cached -} - -async function readCurrentIndexPointer( - storage: KnowhereSdkStorage, - documentId: string, -): Promise { - const key = createPageCitationAssetCurrentIndexKey(documentId) - const parsed = await readJsonObject(storage, key) - if (!isRecord(parsed)) return null - - const indexKey = getTrimmedString(parsed.indexKey) - const jobId = getTrimmedString(parsed.jobId) - const variant = getTrimmedString(parsed.variant) - if (!indexKey || !jobId || !variant) return null - - return { - version: 1, - documentId, - jobId, - variant, - indexKey, - namespace: getTrimmedString(parsed.namespace) ?? undefined, - jobResultId: getTrimmedString(parsed.jobResultId) ?? undefined, - updatedAt: getTrimmedString(parsed.updatedAt) ?? "", - } -} - -async function readPageCitationAssetIndex( - storage: KnowhereSdkStorage, - key: string, -): Promise { - const parsed = await readJsonObject(storage, key) - if (!isRecord(parsed) || !Array.isArray(parsed.assets)) return null - - const assets = parsed.assets.flatMap((asset): PageCitationAsset[] => { - if (!isRecord(asset)) return [] - const pageNum = getPositiveInteger(asset.pageNum) - const keyValue = getTrimmedString(asset.key) - const mimeType = getPageCitationMimeType(asset.mimeType) - const width = getPositiveInteger(asset.width) - const height = getPositiveInteger(asset.height) - if (!pageNum || !keyValue || !mimeType || !width || !height) return [] - - return [ - { - pageNum, - key: keyValue, - assetUrl: getTrimmedString(asset.assetUrl) ?? undefined, - mimeType, - width, - height, - source: "client-rendered-pdf-page", - variant: getTrimmedString(asset.variant) ?? defaultVariant, }, ] }) - if (assets.length === 0) return null - - return { - version: 1, - documentId: getTrimmedString(parsed.documentId) ?? "", - jobId: getTrimmedString(parsed.jobId) ?? "", - variant: getTrimmedString(parsed.variant) ?? defaultVariant, - generatedAt: getTrimmedString(parsed.generatedAt) ?? "", - namespace: getTrimmedString(parsed.namespace) ?? undefined, - jobResultId: getTrimmedString(parsed.jobResultId) ?? undefined, - assets, - } -} - -async function readJsonObject( - storage: KnowhereSdkStorage, - key: string, -): Promise { - if (!storage.readObject) return null - - const object = await storage.readObject(key).catch(() => null) - if (!object) return null - - try { - const parsed: unknown = JSON.parse(new TextDecoder().decode(object.body)) - return parsed - } catch { - return null - } -} - -function getResultJobId( - result: RetrievalResult, - source: Source | undefined, -): string | null { - const metadata = result.metadata - return ( - getRecordString(metadata, "jobId") ?? - getRecordString(metadata, "job_id") ?? - getTrimmedString(source?.knowhereJobId) ?? - null - ) } function getPageNumbers( @@ -317,51 +123,6 @@ function getPageNumbers( return [...pageNumbers].sort((left, right) => left - right) } -function createPageCitationAssetIndexKey( - documentId: string, - jobId: string, -): string { - return [ - "page-citation-assets", - "documents", - toSafeKeySegment(documentId), - "jobs", - toSafeKeySegment(jobId), - "variants", - defaultVariant, - "index.json", - ].join("/") -} - -function createPageCitationAssetCurrentIndexKey(documentId: string): string { - return [ - "page-citation-assets", - "documents", - toSafeKeySegment(documentId), - "current.json", - ].join("/") -} - -function toSafeKeySegment(value: string): string { - const normalized = value.replace(/[^A-Za-z0-9._-]+/g, "_").replace(/^_+|_+$/g, "") - if (normalized.length > 0 && normalized === value) return normalized - - const hash = createHash("sha256").update(value).digest("hex").slice(0, 12) - return `${(normalized || "value").slice(0, 48)}-${hash}` -} - -function getPageCitationMimeType(value: unknown): PageCitationAsset["mimeType"] | null { - return value === "image/png" || value === "image/jpeg" ? value : null -} - -function getRecordString( - record: Readonly> | undefined, - key: string, -): string | null { - if (!record) return null - return getTrimmedString(record[key]) -} - function getTrimmedString(value: unknown): string | null { if (typeof value !== "string") return null const trimmed = value.trim() diff --git a/src/domains/sources/page-citation-assets.test.ts b/src/domains/sources/page-citation-assets.test.ts deleted file mode 100644 index d748929..0000000 --- a/src/domains/sources/page-citation-assets.test.ts +++ /dev/null @@ -1,82 +0,0 @@ -import { describe, expect, it, vi } from "vitest" -import type { - KnowhereSdkStorage, - KnowhereSdkStorageHead, - KnowhereSdkStorageObject, - KnowhereSdkStorageReadResult, - KnowhereSdkStorageWriteResult, - PageRenderer, - RenderedPage, -} from "@ontos-ai/knowhere-sdk" - -import { prepareSourcePageCitationAssets } from "./page-citation-assets" - -class MemoryStorage implements KnowhereSdkStorage { - headObject(): Promise { - return Promise.resolve(null) - } - - writeObject( - input: KnowhereSdkStorageObject, - ): Promise { - return Promise.resolve({ key: input.key }) - } - - readObject(): Promise { - return Promise.resolve(null) - } -} - -describe("prepareSourcePageCitationAssets", () => { - it("calls SDK cacheJobResult with tolerant bounded page citation options", async () => { - const storage = new MemoryStorage() - const renderer: PageRenderer & { close: () => Promise } = { - renderPage(): Promise { - return Promise.resolve({ - body: new Uint8Array([1]), - mimeType: "image/png", - width: 1, - height: 1, - }) - }, - close: vi.fn(async () => undefined), - } - const cacheJobResult = vi.fn(async () => ({ - pageCitationAssetWarnings: [ - { - code: renderLimitExceededCode, - message: "limit reached", - }, - ], - })) - - const result = await prepareSourcePageCitationAssets({ - client: { - knowledge: { - cacheJobResult, - }, - }, - sourceId: "source_1", - jobId: "job_1", - documentId: "doc_1", - storage, - renderer, - maxPagesToRenderPerRun: 3, - }) - - expect(cacheJobResult).toHaveBeenCalledWith({ - jobId: "job_1", - localDocumentId: "source_1", - pageCitationAssets: { - storage, - renderer, - strict: false, - maxPagesToRenderPerRun: 3, - }, - }) - expect(result.warnings).toHaveLength(1) - expect(renderer.close).toHaveBeenCalled() - }) -}) - -const renderLimitExceededCode = "render_limit_exceeded" as const diff --git a/src/domains/sources/page-citation-assets.ts b/src/domains/sources/page-citation-assets.ts deleted file mode 100644 index 7349bc4..0000000 --- a/src/domains/sources/page-citation-assets.ts +++ /dev/null @@ -1,69 +0,0 @@ -import "server-only" - -import { createPdfJsPageRenderer } from "@ontos-ai/knowhere-sdk/page-renderer-pdfjs" -import type { - KnowhereSdkStorage, - PageCitationAssetWarning, - PageCitationAssetsOptions, - PageRenderer, -} from "@ontos-ai/knowhere-sdk" - -import { createVercelBlobKnowhereSdkStorage } from "@/integrations/knowhere-sdk-storage" - -type SourcePageCitationAssetClient = { - readonly knowledge: { - readonly cacheJobResult: (input: { - readonly jobId: string - readonly localDocumentId: string - readonly pageCitationAssets: PageCitationAssetsOptions - }) => Promise<{ - readonly pageCitationAssetWarnings?: readonly PageCitationAssetWarning[] - }> - } -} - -export type PrepareSourcePageCitationAssetsInput = { - readonly client: SourcePageCitationAssetClient - readonly sourceId: string - readonly jobId: string - readonly documentId: string - readonly storage?: KnowhereSdkStorage - readonly renderer?: PageRenderer & { close?: () => Promise } - readonly maxPagesToRenderPerRun?: number -} - -export type PrepareSourcePageCitationAssetsResult = { - readonly warnings: readonly PageCitationAssetWarning[] -} - -const defaultMaxPagesToRenderPerRun = 25 - -export async function prepareSourcePageCitationAssets({ - client, - sourceId, - jobId, - storage = createVercelBlobKnowhereSdkStorage(), - renderer, - maxPagesToRenderPerRun = defaultMaxPagesToRenderPerRun, -}: PrepareSourcePageCitationAssetsInput): Promise { - const pageRenderer = renderer ?? createPdfJsPageRenderer({ maxThreads: 1 }) - - try { - const result = await client.knowledge.cacheJobResult({ - jobId, - localDocumentId: sourceId, - pageCitationAssets: { - storage, - renderer: pageRenderer, - strict: false, - maxPagesToRenderPerRun, - }, - }) - - return { - warnings: result.pageCitationAssetWarnings ?? [], - } - } finally { - await pageRenderer.close?.().catch(() => undefined) - } -} diff --git a/src/domains/sources/source-reconcile-route-workflow.test.ts b/src/domains/sources/source-reconcile-route-workflow.test.ts index a9e8a17..14e9b03 100644 --- a/src/domains/sources/source-reconcile-route-workflow.test.ts +++ b/src/domains/sources/source-reconcile-route-workflow.test.ts @@ -5,23 +5,16 @@ const mocks = vi.hoisted(() => ({ loggerInfo: vi.fn(), loggerWarn: vi.fn(), makeKnowhereClient: vi.fn(), - markSourceFailedAfterReconciliation: vi.fn(), markFailed: vi.fn(), markSourceReadyAfterReconciliation: vi.fn(), pollSourceReconciliation: vi.fn(), - prepareSourcePageCitationAssets: vi.fn(), })) vi.mock("@/domains/sources/source-reconcile-workflow", () => ({ - markSourceFailedAfterReconciliation: mocks.markSourceFailedAfterReconciliation, markSourceReadyAfterReconciliation: mocks.markSourceReadyAfterReconciliation, pollSourceReconciliation: mocks.pollSourceReconciliation, })) -vi.mock("@/domains/sources/page-citation-assets", () => ({ - prepareSourcePageCitationAssets: mocks.prepareSourcePageCitationAssets, -})) - vi.mock("@/domains/sources/workflow-runtime", () => ({ sourceWorkflowRuntime: { markFailed: mocks.markFailed, @@ -83,16 +76,6 @@ describe("sourceReconcileRouteWorkflow", () => { mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ status: "ready", }) - mocks.prepareSourcePageCitationAssets.mockResolvedValue({ - warnings: [ - { - code: "render_limit_exceeded", - message: "One page was skipped.", - documentId: "doc_1", - jobId: "job_1", - }, - ], - }) try { await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ @@ -112,70 +95,11 @@ describe("sourceReconcileRouteWorkflow", () => { sourceId: "source_1", documentId: "doc_1", }) - expect(mocks.prepareSourcePageCitationAssets).toHaveBeenCalledWith({ - client, - sourceId: "source_1", - jobId: "job_1", - documentId: "doc_1", - }) - expect( - mocks.prepareSourcePageCitationAssets.mock.invocationCallOrder[0], - ).toBeLessThan( - mocks.markSourceReadyAfterReconciliation.mock.invocationCallOrder[0]!, - ) - expect(mocks.loggerWarn).toHaveBeenCalledWith( - "workflow: page citation asset warning", - expect.objectContaining({ - sourceId: "source_1", - jobId: "job_1", - documentId: "doc_1", - }), - ) + expect(client).toEqual({ jobs: {} }) + expect(mocks.loggerWarn).not.toHaveBeenCalled() expect(continuations).toEqual([]) }) - it("marks the source failed when page citation asset preparation throws", async () => { - const context = createWorkflowContext() - const restore = - sourceReconcileRouteWorkflow.setContinuationTriggerForTesting( - async () => undefined, - ) - const client = { jobs: {}, knowledge: {} } - mocks.makeKnowhereClient.mockReturnValue(client) - mocks.pollSourceReconciliation.mockResolvedValue({ - kind: "ready-to-prepare", - jobId: "job_1", - documentId: "doc_1", - }) - mocks.prepareSourcePageCitationAssets.mockRejectedValue( - new Error("renderer setup failed"), - ) - mocks.markSourceFailedAfterReconciliation.mockResolvedValue({ - status: "failed", - }) - - try { - await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ - context, - payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ - workspaceId: "workspace_1", - sourceId: "source_1", - apiKey: "jwt_1", - }), - }) - } finally { - restore() - } - - expect(mocks.markSourceFailedAfterReconciliation).toHaveBeenCalledWith({ - workspaceId: "workspace_1", - sourceId: "source_1", - reason: - "Page citation asset preparation failed: renderer setup failed", - }) - expect(mocks.markSourceReadyAfterReconciliation).not.toHaveBeenCalled() - }) - it("triggers a fresh poll run when Knowhere is still running after the segment budget", async () => { const context = createWorkflowContext() const continuations: ContinuationTriggerInput[] = [] diff --git a/src/domains/sources/source-reconcile-route-workflow.ts b/src/domains/sources/source-reconcile-route-workflow.ts index f195250..a1e732a 100644 --- a/src/domains/sources/source-reconcile-route-workflow.ts +++ b/src/domains/sources/source-reconcile-route-workflow.ts @@ -3,13 +3,11 @@ import "server-only" import { Client, type WorkflowContext } from "@upstash/workflow" import { - markSourceFailedAfterReconciliation, markSourceReadyAfterReconciliation, pollSourceReconciliation, } from "@/domains/sources/source-reconcile-workflow" import { makeKnowhereClient } from "@/integrations/knowhere" import { logger } from "@/lib/logger" -import { prepareSourcePageCitationAssets } from "./page-citation-assets" import { sourceWorkflowRuntime } from "./workflow-runtime" type ReconcilePayload = { @@ -41,19 +39,6 @@ type ContinuationTriggerInput = { readonly workflowRunId: string } -type PageCitationAssetPreparationStepResult = - | { - readonly ok: true - readonly warnings: Awaited< - ReturnType - >["warnings"] - } - | { - readonly ok: false - readonly reason: string - readonly error: string - } - const maxPollAttempts = 25 const initialDelaySeconds = 3 const maxDelaySeconds = 30 @@ -133,16 +118,6 @@ async function runPollAndMirrorWorkflow(input: { return } - const preparation = await runPageCitationAssetPreparation({ - context, - workspaceId, - sourceId, - client, - jobId: jobToPrepare.jobId, - documentId: jobToPrepare.documentId, - }) - if (!preparation.ok) return - const ready = await context.run("source-ready", async () => markSourceReadyAfterReconciliation({ workspaceId, @@ -157,79 +132,6 @@ async function runPollAndMirrorWorkflow(input: { }) } -async function runPageCitationAssetPreparation(input: { - readonly context: ReconcileWorkflowContext - readonly workspaceId: string - readonly sourceId: string - readonly client: ReturnType - readonly jobId: string - readonly documentId: string -}): Promise<{ readonly ok: boolean }> { - const result = await input.context.run( - "prepare-page-citation-assets", - async (): Promise => { - try { - const preparation = await prepareSourcePageCitationAssets({ - client: input.client, - sourceId: input.sourceId, - jobId: input.jobId, - documentId: input.documentId, - }) - - return { - ok: true, - warnings: preparation.warnings, - } - } catch (error) { - const errorMessage = error instanceof Error ? error.message : String(error) - return { - ok: false, - reason: `Page citation asset preparation failed: ${getSafeFailureReason( - errorMessage, - )}`, - error: errorMessage, - } - } - }, - ) - - if (!result.ok) { - const failed = await input.context.run( - "source-failed-page-citation-assets", - async () => - markSourceFailedAfterReconciliation({ - workspaceId: input.workspaceId, - sourceId: input.sourceId, - reason: result.reason, - }), - ) - logger.error("workflow: page citation asset preparation failed", { - sourceId: input.sourceId, - jobId: input.jobId, - documentId: input.documentId, - status: failed.status, - error: result.error, - }) - return { ok: false } - } - - for (const warning of result.warnings) { - logger.warn("workflow: page citation asset warning", { - sourceId: input.sourceId, - jobId: input.jobId, - documentId: input.documentId, - warning, - }) - } - logger.info("workflow: page citation assets prepared", { - sourceId: input.sourceId, - jobId: input.jobId, - documentId: input.documentId, - warningCount: result.warnings.length, - }) - return { ok: true } -} - function normalizeReconcilePayload( payload: ReconcilePayload, ): NormalizedReconcilePayload { diff --git a/src/integrations/knowhere-sdk-storage.test.ts b/src/integrations/knowhere-sdk-storage.test.ts deleted file mode 100644 index c5a19b6..0000000 --- a/src/integrations/knowhere-sdk-storage.test.ts +++ /dev/null @@ -1,320 +0,0 @@ -import { BlobNotFoundError, type HeadBlobResult } from "@vercel/blob" -import { describe, expect, it } from "vitest" - -import { - createVercelBlobKnowhereSdkStorage, - type VercelBlobStore, -} from "./knowhere-sdk-storage" - -type StoredBlob = { - readonly pathname: string - readonly body: Uint8Array - readonly contentType?: string - readonly url: string -} - -class MemoryVercelBlobStore implements VercelBlobStore { - readonly objects = new Map() - - async put( - pathname: string, - body: Parameters[1], - options: Parameters[2], - ): Promise<{ readonly pathname: string; readonly url: string }> { - const url = `https://blob.example/${pathname}` - this.objects.set(pathname, { - pathname, - body: await readBody(body), - contentType: options.contentType, - url, - }) - return { pathname, url } - } - - async head(pathname: string): Promise { - const object = this.objects.get(pathname) - if (!object) throw new BlobNotFoundError() - - return { - size: object.body.byteLength, - uploadedAt: new Date("2026-07-03T00:00:00.000Z"), - pathname, - contentType: object.contentType ?? "application/octet-stream", - contentDisposition: "", - url: object.url, - downloadUrl: `${object.url}?download=1`, - cacheControl: "", - etag: `etag:${pathname}`, - } - } - - async get( - pathname: string, - ): Promise>> { - const object = this.objects.get(pathname) - if (!object) return null - - return { - statusCode: 200, - stream: toReadableStream(object.body), - headers: new Headers(), - blob: { - url: object.url, - downloadUrl: `${object.url}?download=1`, - pathname, - contentDisposition: "", - cacheControl: "", - uploadedAt: new Date("2026-07-03T00:00:00.000Z"), - etag: `etag:${pathname}`, - contentType: object.contentType ?? "application/octet-stream", - size: object.body.byteLength, - }, - } - } - - async list(input: { - readonly prefix: string - readonly limit: number - readonly cursor?: string - }): Promise>> { - const start = input.cursor ? Number(input.cursor) : 0 - const matches = [...this.objects.values()] - .filter((object) => object.pathname.startsWith(input.prefix)) - .sort((left, right) => left.pathname.localeCompare(right.pathname)) - const page = matches.slice(start, start + input.limit) - const nextIndex = start + page.length - const hasMore = nextIndex < matches.length - - return { - blobs: page.map((object) => ({ - url: object.url, - downloadUrl: `${object.url}?download=1`, - pathname: object.pathname, - size: object.body.byteLength, - uploadedAt: new Date("2026-07-03T00:00:00.000Z"), - etag: `etag:${object.pathname}`, - })), - hasMore, - ...(hasMore ? { cursor: String(nextIndex) } : {}), - } - } - - async del(pathname: string | readonly string[]): Promise { - const pathnames = typeof pathname === "string" ? [pathname] : [...pathname] - for (const item of pathnames) { - this.objects.delete(item) - } - } -} - -describe("createVercelBlobKnowhereSdkStorage", () => { - it("rejects unsafe storage keys", async () => { - const storage = createVercelBlobKnowhereSdkStorage( - new MemoryVercelBlobStore(), - ) - - for (const key of [ - "", - "/absolute/path.png", - "nested//empty.png", - "nested/../escape.png", - "nested/./dot.png", - "nested\\windows.png", - "nested/\0/null.png", - ]) { - await expect( - storage.writeObject({ - key, - body: new Uint8Array([1]), - contentType: "image/png", - }), - ).rejects.toThrow("Invalid Knowhere SDK storage key") - } - }) - - it("writes objects with metadata sidecars and returns blob URLs", async () => { - const blobStore = new MemoryVercelBlobStore() - const storage = createVercelBlobKnowhereSdkStorage(blobStore) - - const result = await storage.writeObject({ - key: "page-citation-assets/documents/doc_1/page.png", - body: new Uint8Array([1, 2, 3]), - contentType: "image/png", - metadata: { - width: "1200", - height: "1600", - mimeType: "image/png", - }, - }) - - expect(result).toEqual({ - key: "page-citation-assets/documents/doc_1/page.png", - url: "https://blob.example/page-citation-assets/documents/doc_1/page.png", - }) - expect( - blobStore.objects.has("page-citation-assets/documents/doc_1/page.png"), - ).toBe(true) - const metadataSidecar = blobStore.objects.get( - "page-citation-assets/documents/doc_1/page.png.metadata.json", - ) - expect(metadataSidecar).toBeDefined() - if (!metadataSidecar) throw new Error("Expected metadata sidecar.") - - expect(JSON.parse(decodeBody(metadataSidecar.body))).toEqual({ - metadata: { - width: "1200", - height: "1600", - mimeType: "image/png", - }, - }) - }) - - it("combines blob head data with sidecar metadata", async () => { - const blobStore = new MemoryVercelBlobStore() - const storage = createVercelBlobKnowhereSdkStorage(blobStore) - - await storage.writeObject({ - key: "page-citation-assets/documents/doc_1/page.png", - body: new Uint8Array([1, 2, 3]), - contentType: "image/png", - metadata: { - width: "1200", - height: "1600", - }, - }) - - await expect( - storage.headObject("page-citation-assets/documents/doc_1/page.png"), - ).resolves.toEqual({ - key: "page-citation-assets/documents/doc_1/page.png", - contentType: "image/png", - contentLength: 3, - metadata: { - width: "1200", - height: "1600", - }, - }) - }) - - it("reads objects, resolves URLs, deletes objects, and deletes prefixes", async () => { - const blobStore = new MemoryVercelBlobStore() - const storage = createVercelBlobKnowhereSdkStorage(blobStore) - - await storage.writeObject({ - key: "prefix/one.png", - body: new Uint8Array([1]), - contentType: "image/png", - }) - await storage.writeObject({ - key: "prefix/two.png", - body: new Uint8Array([2]), - contentType: "image/png", - }) - await storage.writeObject({ - key: "prefix-sibling/two.png", - body: new Uint8Array([4]), - contentType: "image/png", - }) - await storage.writeObject({ - key: "other/three.png", - body: new Uint8Array([3]), - contentType: "image/png", - }) - - const readObject = storage.readObject - const getObjectUrl = storage.getObjectUrl - const deleteObject = storage.deleteObject - const deletePrefix = storage.deletePrefix - expect(readObject).toBeDefined() - expect(getObjectUrl).toBeDefined() - expect(deleteObject).toBeDefined() - expect(deletePrefix).toBeDefined() - if (!readObject || !getObjectUrl || !deleteObject || !deletePrefix) { - throw new Error("Expected full storage adapter methods.") - } - - await expect(storage.readObject?.("prefix/one.png")).resolves.toMatchObject({ - body: new Uint8Array([1]), - contentType: "image/png", - }) - await expect(storage.getObjectUrl?.("prefix/one.png")).resolves.toBe( - "https://blob.example/prefix/one.png", - ) - - await storage.deleteObject?.("prefix/one.png") - expect(blobStore.objects.has("prefix/one.png")).toBe(false) - expect(blobStore.objects.has("prefix/one.png.metadata.json")).toBe(false) - - await storage.deletePrefix?.("prefix") - expect([...blobStore.objects.keys()].sort()).toEqual([ - "other/three.png", - "other/three.png.metadata.json", - "prefix-sibling/two.png", - "prefix-sibling/two.png.metadata.json", - ]) - }) - - it("accepts trailing slashes when deleting prefixes", async () => { - const blobStore = new MemoryVercelBlobStore() - const storage = createVercelBlobKnowhereSdkStorage(blobStore) - - await storage.writeObject({ - key: "prefix/one.png", - body: new Uint8Array([1]), - contentType: "image/png", - }) - - await storage.deletePrefix?.("prefix/") - - expect(blobStore.objects.size).toBe(0) - }) -}) - -async function readBody( - body: Parameters[1], -): Promise { - if (typeof body === "string") return new TextEncoder().encode(body) - if (body instanceof Uint8Array) return new Uint8Array(body) - if (body instanceof Blob) return new Uint8Array(await body.arrayBuffer()) - return readStream(body) -} - -async function readStream( - stream: ReadableStream, -): Promise { - const reader = stream.getReader() - const chunks: Uint8Array[] = [] - let totalLength = 0 - - try { - while (true) { - const result = await reader.read() - if (result.done) break - chunks.push(result.value) - totalLength += result.value.byteLength - } - } finally { - reader.releaseLock() - } - - const body = new Uint8Array(totalLength) - let offset = 0 - for (const chunk of chunks) { - body.set(chunk, offset) - offset += chunk.byteLength - } - return body -} - -function toReadableStream(body: Uint8Array): ReadableStream { - return new ReadableStream({ - start(controller) { - controller.enqueue(body) - controller.close() - }, - }) -} - -function decodeBody(body: Uint8Array): string { - return new TextDecoder().decode(body) -} diff --git a/src/integrations/knowhere-sdk-storage.ts b/src/integrations/knowhere-sdk-storage.ts deleted file mode 100644 index 387bc21..0000000 --- a/src/integrations/knowhere-sdk-storage.ts +++ /dev/null @@ -1,294 +0,0 @@ -import "server-only" - -import { Buffer } from "node:buffer" -import { - BlobNotFoundError, - del, - get, - head, - list, - put, - type GetBlobResult, - type HeadBlobResult, - type ListBlobResult, - type PutBlobResult, -} from "@vercel/blob" -import type { - KnowhereSdkStorage, - KnowhereSdkStorageBody, - KnowhereSdkStorageHead, - KnowhereSdkStorageObject, - KnowhereSdkStorageReadResult, - KnowhereSdkStorageWriteResult, -} from "@ontos-ai/knowhere-sdk" - -type VercelBlobBody = Buffer | Blob | ReadableStream - -export type VercelBlobStore = { - readonly put: ( - pathname: string, - body: VercelBlobBody | string, - options: { - readonly access: "public" - readonly addRandomSuffix: false - readonly allowOverwrite: true - readonly contentType?: string - }, - ) => Promise> - readonly head: (pathname: string) => Promise - readonly get: ( - pathname: string, - options: { readonly access: "public" }, - ) => Promise - readonly list: (options: { - readonly prefix: string - readonly limit: number - readonly cursor?: string - }) => Promise - readonly del: (pathname: string | readonly string[]) => Promise -} - -type StoredMetadataSidecar = { - readonly metadata?: Readonly> -} - -const metadataSuffix = ".metadata.json" -const deletePrefixPageSize = 1000 - -const vercelBlobStore: VercelBlobStore = { - put, - head, - get, - list, - del: (pathname) => - del(typeof pathname === "string" ? pathname : [...pathname]), -} - -class VercelBlobKnowhereSdkStorage implements KnowhereSdkStorage { - private readonly blobStore: VercelBlobStore - - constructor(blobStore: VercelBlobStore = vercelBlobStore) { - this.blobStore = blobStore - } - - async headObject(key: string): Promise { - validateBlobStorageKey(key) - - try { - const object = await this.blobStore.head(key) - const metadata = await this.readMetadata(key) - return { - key, - contentType: object.contentType, - contentLength: object.size, - metadata, - } - } catch (error) { - if (isBlobNotFoundError(error)) return null - throw error - } - } - - async writeObject( - input: KnowhereSdkStorageObject, - ): Promise { - validateBlobStorageKey(input.key) - - const object = await this.blobStore.put( - input.key, - toVercelBlobBody(input.body), - { - access: "public", - addRandomSuffix: false, - allowOverwrite: true, - contentType: input.contentType, - }, - ) - await this.writeMetadata(input.key, input.metadata ?? {}) - - return { - key: object.pathname, - url: object.url, - } - } - - async getObjectUrl(key: string): Promise { - validateBlobStorageKey(key) - - try { - const object = await this.blobStore.head(key) - return object.url - } catch (error) { - if (isBlobNotFoundError(error)) return null - throw error - } - } - - async readObject(key: string): Promise { - validateBlobStorageKey(key) - - const object = await this.blobStore.get(key, { access: "public" }) - if (!object || object.statusCode !== 200) return null - - return { - body: await readStream(object.stream), - contentType: object.blob.contentType, - metadata: await this.readMetadata(key), - } - } - - async deleteObject(key: string): Promise { - validateBlobStorageKey(key) - await this.deleteIgnoringMissing([key, getMetadataKey(key)]) - } - - async deletePrefix(prefix: string): Promise { - validateBlobStoragePrefix(prefix) - const blobPrefix = prefix.endsWith("/") ? prefix : `${prefix}/` - - let cursor: string | undefined - do { - const page = await this.blobStore.list({ - prefix: blobPrefix, - limit: deletePrefixPageSize, - ...(cursor ? { cursor } : {}), - }) - const pathnames = page.blobs.map((blob) => blob.pathname) - if (pathnames.length > 0) await this.deleteIgnoringMissing(pathnames) - cursor = page.cursor - if (page.hasMore && !cursor) break - } while (cursor) - } - - private async readMetadata( - key: string, - ): Promise>> { - const metadataObject = await this.blobStore.get(getMetadataKey(key), { - access: "public", - }) - if (!metadataObject || metadataObject.statusCode !== 200) return {} - - const body = await readStream(metadataObject.stream) - const parsed: unknown = JSON.parse(new TextDecoder().decode(body)) - return parseStoredMetadataSidecar(parsed).metadata ?? {} - } - - private async writeMetadata( - key: string, - metadata: Readonly>, - ): Promise { - const sidecar: StoredMetadataSidecar = { metadata } - await this.blobStore.put( - getMetadataKey(key), - JSON.stringify(sidecar, null, 2), - { - access: "public", - addRandomSuffix: false, - allowOverwrite: true, - contentType: "application/json", - }, - ) - } - - private async deleteIgnoringMissing( - pathnames: readonly string[], - ): Promise { - try { - await this.blobStore.del(pathnames) - } catch (error) { - if (!isBlobNotFoundError(error)) throw error - } - } -} - -export function createVercelBlobKnowhereSdkStorage( - blobStore?: VercelBlobStore, -): KnowhereSdkStorage { - return new VercelBlobKnowhereSdkStorage(blobStore) -} - -function validateBlobStorageKey(key: string): void { - if (!isSafeRelativePosixPath(key)) { - throw new Error(`Invalid Knowhere SDK storage key: ${key}`) - } -} - -function validateBlobStoragePrefix(prefix: string): void { - const normalizedPrefix = prefix.endsWith("/") ? prefix.slice(0, -1) : prefix - if (!isSafeRelativePosixPath(normalizedPrefix)) { - throw new Error(`Invalid Knowhere SDK storage prefix: ${prefix}`) - } -} - -function isSafeRelativePosixPath(value: string): boolean { - if ( - value.length === 0 || - value.startsWith("/") || - value.includes("\\") || - value.includes("\0") - ) { - return false - } - - return value - .split("/") - .every((segment) => segment.length > 0 && segment !== "." && segment !== "..") -} - -function getMetadataKey(key: string): string { - return `${key}${metadataSuffix}` -} - -function toVercelBlobBody(body: KnowhereSdkStorageBody): VercelBlobBody { - if (body instanceof Uint8Array) return Buffer.from(body) - return body -} - -async function readStream( - stream: ReadableStream, -): Promise { - const reader = stream.getReader() - const chunks: Uint8Array[] = [] - let totalLength = 0 - - try { - while (true) { - const result = await reader.read() - if (result.done) break - chunks.push(result.value) - totalLength += result.value.byteLength - } - } finally { - reader.releaseLock() - } - - const body = new Uint8Array(totalLength) - let offset = 0 - for (const chunk of chunks) { - body.set(chunk, offset) - offset += chunk.byteLength - } - return body -} - -function parseStoredMetadataSidecar(value: unknown): StoredMetadataSidecar { - if (!isRecord(value)) return {} - const metadata = value.metadata - if (!isRecord(metadata)) return {} - - return { - metadata: Object.fromEntries( - Object.entries(metadata).filter( - (entry): entry is [string, string] => typeof entry[1] === "string", - ), - ), - } -} - -function isRecord(value: unknown): value is Record { - return typeof value === "object" && value !== null -} - -function isBlobNotFoundError(error: unknown): boolean { - return error instanceof BlobNotFoundError -} From ce8357316b78b530270ed0047f6cdb663bd12af5 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 3 Jul 2026 18:51:32 +0800 Subject: [PATCH 03/51] Mirror page citation assets to Blob --- src/domains/chat/index.test.ts | 109 +++++++++++++++++- src/domains/chat/index.ts | 57 +++++++-- .../chat/media-asset-hardening.test.ts | 99 +++++++++++++++- src/domains/chat/media-asset-hardening.ts | 104 ++++++++++++----- 4 files changed, 326 insertions(+), 43 deletions(-) diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 0c8cfd6..9daadc3 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -9,7 +9,10 @@ import { generateAgenticOutputManifest, parseChatRequestBody, } from "." -import type { HardenMediaAssetUrlsInput } from "./media-asset-hardening" +import type { + HardenableRetrievalResult, + HardenMediaAssetUrlsInput, +} from "./media-asset-hardening" import type { Source } from "@/infrastructure/db/schema" import type { ChatArtifactView } from "@/domains/chat/types" @@ -665,6 +668,110 @@ describe("answerQuestionWithRetrieval", () => { expect(answer.artifacts?.[0]?.citation?.assetUrl).toBe(hardenedAssetUrl); }); + it("hardens page citation asset URLs before returning citations", async () => { + const rawPageAssetUrl = + "https://knowhere-storage.example/results/job_1/page_citation_assets/page-4.png?AWSAccessKeyId=test"; + const hardenedPageAssetUrl = + "https://blob.example/workspaces/workspace_1/chat-assets/source-source_pages/page-4.png"; + const result = makeRetrievalResult({ + chunkType: "page", + metadata: { + pageNums: [4], + pageAssets: [ + { + pageNum: 4, + artifactRef: "page_citation_assets/page-4.png", + assetUrl: rawPageAssetUrl, + contentType: "image/png", + }, + ], + }, + source: { + documentId: "doc_pages", + sourceFileName: "document-generated.pdf", + sectionPath: "Page 4", + }, + }); + const retrieval = { + query: vi.fn().mockResolvedValue({ + results: [result], + evidenceText: "Page four evidence.", + referencedChunks: [], + namespace: "notebook-workspace", + query: "page four evidence", + routerUsed: "workflow_single_step", + answerText: null, + }), + }; + const generateAnswer = vi.fn(async ({ searchSources }) => { + await searchSources({ query: "page four evidence" }); + return makeHarnessRunResult(`This page has the answer. ${rawPageAssetUrl}`); + }); + const hardenMediaAssetUrls = vi.fn( + async ({ + results, + artifacts, + }: HardenMediaAssetUrlsInput): Promise<{ + results: HardenableRetrievalResult[] + artifacts?: ChatArtifactView[] + }> => ({ + results: results.map( + (candidate): HardenableRetrievalResult => ({ + ...candidate, + pageCitationAssetUrl: + candidate.pageCitationAssetUrl === rawPageAssetUrl + ? hardenedPageAssetUrl + : candidate.pageCitationAssetUrl, + }), + ), + ...(artifacts ? { artifacts: [...artifacts] } : {}), + }), + ); + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "What is on page four?", + namespace: "notebook-workspace", + sources: [ + makeSource({ + id: "source_pages", + title: "deck.pdf", + knowhereDocumentId: "doc_pages", + }), + ], + excludedSourceIds: [], + retrieval, + generateAnswer, + hardenMediaAssetUrls, + messages: [], + }), + ); + + expect(hardenMediaAssetUrls).toHaveBeenCalledWith({ + results: [ + expect.objectContaining({ + pageCitationAssetUrl: rawPageAssetUrl, + source: expect.objectContaining({ + sourceFileName: "deck.pdf", + }), + }), + ], + artifacts: undefined, + }); + expect(answer.answer).toBe("This page has the answer."); + expect(answer.answer).not.toContain("knowhere-storage.example"); + expect(answer.citations).toEqual([ + expect.objectContaining({ + chunkType: "page", + pageCitationAssetUrl: hardenedPageAssetUrl, + source: expect.objectContaining({ + sourceFileName: "deck.pdf", + }), + }), + ]); + expect(answer.citations[0]?.pageCitationAssetUrl).not.toBe(rawPageAssetUrl); + }); + it("returns only harness-selected artifacts when retrieval has extra media candidates", async () => { const frontAssetUrl = "https://blob.example/images/id-front.jpg"; const backAssetUrl = "https://blob.example/images/id-back.jpg"; diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index b0a95d9..ec19628 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -39,6 +39,7 @@ import { removeRetrievedMediaAssetUrls, } from "./media-assets" import { enrichRetrievalResultsWithPageCitationAssetUrls } from "./page-citation-assets" +import type { HardenableRetrievalResult } from "./media-asset-hardening" const DEFAULT_TOP_K = 8 const MAX_AGENTIC_TOP_K = 12 @@ -245,11 +246,17 @@ export const answerQuestionWithRetrieval = ( evidenceText: formatRetrievalEvidenceText(retrievalResponses), }), ) + const pageCitationResults = yield* Effect.tryPromise(() => + enrichRetrievalResultsWithPageCitationAssetUrls({ + results: enrichedResults, + sources: input.sources, + }), + ) const artifacts = toChatArtifactViewsFromHarness(generatedAnswer, input.sources) const hardenedMedia = yield* Effect.tryPromise(() => hardenAnswerMediaAssetUrls({ input, - results: enrichedResults, + results: pageCitationResults, artifacts, }), ) @@ -258,26 +265,21 @@ export const answerQuestionWithRetrieval = ( results: getGeneratedAnswerSanitizerResults({ rawResults, enrichedResults, + pageCitationResults, hardenedResults: hardenedMedia.results, artifacts, hardenedArtifacts: hardenedMedia.artifacts, }), }) - const citationResults = yield* Effect.tryPromise(() => - enrichRetrievalResultsWithPageCitationAssetUrls({ - results: hardenedMedia.results, - sources: input.sources, - }), - ) const displayArtifacts = hardenedMedia.artifacts ?? [] logger.info("chat-agent: answer complete", { answerLength: answer.length, - citationCount: citationResults.length, + citationCount: hardenedMedia.results.length, artifactCount: displayArtifacts.length, }) return { answer, - citations: toChatCitationViews(citationResults, answer), + citations: toChatCitationViews(hardenedMedia.results, answer), artifacts: displayArtifacts, } }) @@ -422,7 +424,7 @@ function normalizeHarnessSource( type AnswerMediaAssetHardeningInput = { readonly input: AnswerQuestionInput - readonly results: readonly RetrievalResult[] + readonly results: readonly HardenableRetrievalResult[] readonly artifacts?: readonly ChatArtifactView[] } @@ -431,7 +433,7 @@ async function hardenAnswerMediaAssetUrls({ results, artifacts, }: AnswerMediaAssetHardeningInput): Promise<{ - readonly results: RetrievalResult[] + readonly results: HardenableRetrievalResult[] readonly artifacts?: ChatArtifactView[] }> { if (!input.hardenMediaAssetUrls) { @@ -462,7 +464,8 @@ async function hardenAnswerMediaAssetUrls({ type GeneratedAnswerSanitizerResultsInput = { readonly rawResults: readonly RetrievalResult[] readonly enrichedResults: readonly RetrievalResult[] - readonly hardenedResults: readonly RetrievalResult[] + readonly pageCitationResults: readonly HardenableRetrievalResult[] + readonly hardenedResults: readonly HardenableRetrievalResult[] readonly artifacts?: readonly ChatArtifactView[] readonly hardenedArtifacts?: readonly ChatArtifactView[] } @@ -470,6 +473,7 @@ type GeneratedAnswerSanitizerResultsInput = { function getGeneratedAnswerSanitizerResults({ rawResults, enrichedResults, + pageCitationResults, hardenedResults, artifacts, hardenedArtifacts, @@ -477,12 +481,32 @@ function getGeneratedAnswerSanitizerResults({ return [ ...rawResults, ...enrichedResults, + ...toPageCitationSanitizerResults(pageCitationResults), ...hardenedResults, + ...toPageCitationSanitizerResults(hardenedResults), ...toArtifactSanitizerResults(artifacts), ...toArtifactSanitizerResults(hardenedArtifacts), ] } +function toPageCitationSanitizerResults( + results: readonly HardenableRetrievalResult[], +): RetrievalResult[] { + return results.flatMap((result): RetrievalResult[] => { + if (!result.pageCitationAssetUrl) return [] + + return [ + { + content: result.content, + chunkType: result.chunkType, + score: result.score, + assetUrl: result.pageCitationAssetUrl, + source: result.source, + }, + ] + }) +} + function toArtifactSanitizerResults( artifacts: readonly ChatArtifactView[] | undefined, ): RetrievalResult[] { @@ -506,6 +530,15 @@ function toArtifactSanitizerResults( }), ) } + if (artifact.citation?.pageCitationAssetUrl) { + results.push( + toArtifactSanitizerResult({ + assetUrl: artifact.citation.pageCitationAssetUrl, + artifact, + citation: artifact.citation, + }), + ) + } return results }) } diff --git a/src/domains/chat/media-asset-hardening.test.ts b/src/domains/chat/media-asset-hardening.test.ts index 70cf562..fe1f3df 100644 --- a/src/domains/chat/media-asset-hardening.test.ts +++ b/src/domains/chat/media-asset-hardening.test.ts @@ -1,11 +1,11 @@ import { afterEach, describe, expect, it, vi } from "vitest" -import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" import { hardenChatMediaAssetUrls, type ChatMediaAssetBlobStore, type FetchChatMediaAsset, + type HardenableRetrievalResult, } from "./media-asset-hardening" const loggerMock = vi.hoisted(() => ({ @@ -73,6 +73,54 @@ describe("hardenChatMediaAssetUrls", () => { ) }) + it("copies server page citation asset URLs into Notebook chat assets", async () => { + const rawPageAssetUrl = + "https://knowhere-storage.example/results/job_1/page_citation_assets/page-4.png?AWSAccessKeyId=test&Signature=secret" + const blobPageAssetUrl = + "https://blob.example/workspaces/workspace_1/chat-assets/source-source_1/page-4.png" + const blobStore = makeBlobStore(blobPageAssetUrl) + const fetchAsset = makeFetchAsset("page-image-bytes", "image/png") + + const result = await hardenChatMediaAssetUrls({ + workspaceId: "workspace_1", + sources: [ + makeSource({ + id: "source_1", + knowhereDocumentId: "doc_pages", + }), + ], + results: [ + makeRetrievalResult({ + chunkType: "page", + pageCitationAssetUrl: rawPageAssetUrl, + source: { + documentId: "doc_pages", + sourceFileName: "deck.pdf", + sectionPath: "Page 4", + }, + }), + ], + blobStore, + fetchAsset, + }) + + expect(fetchAsset).toHaveBeenCalledWith(rawPageAssetUrl) + expect(blobStore.put).toHaveBeenCalledWith( + expect.stringMatching( + /^workspaces\/workspace_1\/chat-assets\/source-source_1\/[a-f0-9]{24}-page-4\.png$/, + ), + expect.any(Buffer), + { + access: "public", + allowOverwrite: true, + contentType: "image/png", + multipart: true, + }, + ) + expect(result.results[0]?.pageCitationAssetUrl).toBe(blobPageAssetUrl) + expect(result.results[0]?.assetUrl).toBeUndefined() + }) + it("uses an existing parsed asset URL before fetching the upstream URL", async () => { const rawAssetUrl = "https://knowhere-storage.example/results/job_1/images/id-front.jpg?AWSAccessKeyId=test" @@ -234,6 +282,51 @@ describe("hardenChatMediaAssetUrls", () => { expect(artifact?.assetUrl).toBe(blobAssetUrl) expect(artifact?.citation?.assetUrl).toBe(blobAssetUrl) }) + + it("rewrites nested citation page citation asset URLs", async () => { + const rawPageAssetUrl = + "https://knowhere-storage.example/results/job_1/page_citation_assets/page-8.png?AWSAccessKeyId=test" + const blobPageAssetUrl = + "https://blob.example/workspaces/workspace_1/chat-assets/source-source_pages/page-8.png" + const blobStore = makeBlobStore(blobPageAssetUrl) + const fetchAsset = makeFetchAsset("page-eight-image", "image/png") + + const result = await hardenChatMediaAssetUrls({ + workspaceId: "workspace_1", + sources: [ + makeSource({ + id: "source_pages", + knowhereDocumentId: "doc_pages", + }), + ], + results: [], + artifacts: [ + { + type: "image", + ref: "asset:r1:result:1", + assetUrl: + "https://blob.example/workspaces/workspace_1/chat-assets/source-source_pages/figure.png", + label: "figure", + citation: { + chunkType: "page", + score: 0.9, + pageCitationAssetUrl: rawPageAssetUrl, + source: { + documentId: "doc_pages", + sourceFileName: "deck.pdf", + sectionPath: "Page 8", + }, + }, + }, + ], + blobStore, + fetchAsset, + }) + + const [artifact] = result.artifacts ?? [] + expect(fetchAsset).toHaveBeenCalledWith(rawPageAssetUrl) + expect(artifact?.citation?.pageCitationAssetUrl).toBe(blobPageAssetUrl) + }) }) function makeFetchAsset( @@ -257,8 +350,8 @@ function makeBlobStore(url: string): ChatMediaAssetBlobStore { } function makeRetrievalResult( - overrides: Partial = {}, -): RetrievalResult { + overrides: Partial = {}, +): HardenableRetrievalResult { return { content: "Asset evidence", chunkType: "text", diff --git a/src/domains/chat/media-asset-hardening.ts b/src/domains/chat/media-asset-hardening.ts index e0a813e..4161587 100644 --- a/src/domains/chat/media-asset-hardening.ts +++ b/src/domains/chat/media-asset-hardening.ts @@ -13,13 +13,17 @@ import { logger } from "@/lib/logger" import type { LoadSourceAssetUrls } from "./media-assets" import { resolveAssetUrlFromReferenceText } from "./media-assets" +export type HardenableRetrievalResult = RetrievalResult & { + readonly pageCitationAssetUrl?: string +} + export type HardenMediaAssetUrlsInput = { - readonly results: readonly RetrievalResult[] + readonly results: readonly HardenableRetrievalResult[] readonly artifacts?: readonly ChatArtifactView[] } export type HardenMediaAssetUrlsResult = { - readonly results: RetrievalResult[] + readonly results: HardenableRetrievalResult[] readonly artifacts?: ChatArtifactView[] } @@ -121,7 +125,7 @@ export async function hardenChatMediaAssetUrls({ } const hardenedResults = await Promise.all( - results.map((result): Promise => + results.map((result): Promise => hardenRetrievalResult(result, context), ), ) @@ -140,25 +144,48 @@ export async function hardenChatMediaAssetUrls({ } async function hardenRetrievalResult( - result: RetrievalResult, + result: HardenableRetrievalResult, context: HardeningContext, -): Promise { +): Promise { const assetUrl = getTrimmedString(result.assetUrl) - if (!assetUrl) return result - - const hardenedAssetUrl = await hardenAssetUrl( - { - assetUrl, - source: result.source, - content: result.content, - }, - context, + const pageCitationAssetUrl = getTrimmedString(result.pageCitationAssetUrl) + if (!assetUrl && !pageCitationAssetUrl) return result + + const hardenedAssetUrl = assetUrl + ? await hardenAssetUrl( + { + assetUrl, + source: result.source, + content: result.content, + }, + context, + ) + : undefined + const hardenedPageCitationAssetUrl = pageCitationAssetUrl + ? await hardenAssetUrl( + { + assetUrl: pageCitationAssetUrl, + source: result.source, + content: result.content, + }, + context, + ) + : undefined + const hasAssetUrlChange = Boolean( + assetUrl && hardenedAssetUrl !== result.assetUrl, + ) + const hasPageCitationAssetUrlChange = Boolean( + pageCitationAssetUrl && + hardenedPageCitationAssetUrl !== result.pageCitationAssetUrl, ) - if (hardenedAssetUrl === result.assetUrl) return result + if (!hasAssetUrlChange && !hasPageCitationAssetUrlChange) return result return { ...result, - assetUrl: hardenedAssetUrl, + ...(hardenedAssetUrl ? { assetUrl: hardenedAssetUrl } : {}), + ...(hardenedPageCitationAssetUrl + ? { pageCitationAssetUrl: hardenedPageCitationAssetUrl } + : {}), } } @@ -200,21 +227,44 @@ async function hardenCitation( context: HardeningContext, ): Promise { const assetUrl = getTrimmedString(citation.assetUrl) - if (!assetUrl) return citation - - const hardenedAssetUrl = await hardenAssetUrl( - { - assetUrl, - source: citation.source, - content: citation.content, - }, - context, + const pageCitationAssetUrl = getTrimmedString(citation.pageCitationAssetUrl) + if (!assetUrl && !pageCitationAssetUrl) return citation + + const hardenedAssetUrl = assetUrl + ? await hardenAssetUrl( + { + assetUrl, + source: citation.source, + content: citation.content, + }, + context, + ) + : undefined + const hardenedPageCitationAssetUrl = pageCitationAssetUrl + ? await hardenAssetUrl( + { + assetUrl: pageCitationAssetUrl, + source: citation.source, + content: citation.content, + }, + context, + ) + : undefined + const hasAssetUrlChange = Boolean( + assetUrl && hardenedAssetUrl !== citation.assetUrl, + ) + const hasPageCitationAssetUrlChange = Boolean( + pageCitationAssetUrl && + hardenedPageCitationAssetUrl !== citation.pageCitationAssetUrl, ) - if (hardenedAssetUrl === citation.assetUrl) return citation + if (!hasAssetUrlChange && !hasPageCitationAssetUrlChange) return citation return { ...citation, - assetUrl: hardenedAssetUrl, + ...(hardenedAssetUrl ? { assetUrl: hardenedAssetUrl } : {}), + ...(hardenedPageCitationAssetUrl + ? { pageCitationAssetUrl: hardenedPageCitationAssetUrl } + : {}), } } From 109bf3660a9f8fa1527940467a3145192cc51556 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 3 Jul 2026 21:14:10 +0800 Subject: [PATCH 04/51] Read page citation assets from metadata --- src/domains/chat/index.test.ts | 99 +++++++++++++++++++ src/domains/chat/page-citation-assets.test.ts | 24 ----- src/domains/chat/page-citation-assets.ts | 5 +- 3 files changed, 100 insertions(+), 28 deletions(-) diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 9daadc3..459add9 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -772,6 +772,105 @@ describe("answerQuestionWithRetrieval", () => { expect(answer.citations[0]?.pageCitationAssetUrl).not.toBe(rawPageAssetUrl); }); + it("hardens page citation asset URLs from referenced chunk metadata", async () => { + const rawPageAssetUrl = + "https://knowhere-storage.example/results/job_1/page_citation_assets/page-6.png?AWSAccessKeyId=test"; + const hardenedPageAssetUrl = + "https://blob.example/workspaces/workspace_1/chat-assets/source-source_pages/page-6.png"; + const retrieval = { + query: vi.fn().mockResolvedValue({ + results: [], + evidenceText: "Page six evidence.", + referencedChunks: [ + { + chunkId: "chunk_page_6", + documentId: "doc_pages", + chunkType: "page", + sectionPath: "Page 6", + filePath: null, + jobId: "job_1", + assetUrl: rawPageAssetUrl, + metadata: { + pageNums: [6], + pageAssets: [ + { + pageNum: 6, + artifactRef: "page_citation_assets/page-6.png", + assetUrl: rawPageAssetUrl, + contentType: "image/png", + }, + ], + }, + }, + ], + namespace: "notebook-workspace", + query: "page six evidence", + routerUsed: "workflow_single_step", + answerText: null, + }), + }; + const generateAnswer = vi.fn(async ({ searchSources }) => { + await searchSources({ query: "page six evidence" }); + return makeHarnessRunResult("This page has referenced evidence."); + }); + const hardenMediaAssetUrls = vi.fn( + async ({ + results, + artifacts, + }: HardenMediaAssetUrlsInput): Promise<{ + results: HardenableRetrievalResult[] + artifacts?: ChatArtifactView[] + }> => ({ + results: results.map( + (candidate): HardenableRetrievalResult => ({ + ...candidate, + pageCitationAssetUrl: + candidate.pageCitationAssetUrl === rawPageAssetUrl + ? hardenedPageAssetUrl + : candidate.pageCitationAssetUrl, + }), + ), + ...(artifacts ? { artifacts: [...artifacts] } : {}), + }), + ); + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "What is on page six?", + namespace: "notebook-workspace", + sources: [ + makeSource({ + id: "source_pages", + title: "deck.pdf", + knowhereDocumentId: "doc_pages", + }), + ], + excludedSourceIds: [], + retrieval, + generateAnswer, + hardenMediaAssetUrls, + messages: [], + }), + ); + + expect(hardenMediaAssetUrls).toHaveBeenCalledWith({ + results: [ + expect.objectContaining({ + metadata: expect.objectContaining({ + pageAssets: [ + expect.objectContaining({ + assetUrl: rawPageAssetUrl, + }), + ], + }), + pageCitationAssetUrl: rawPageAssetUrl, + }), + ], + artifacts: undefined, + }); + expect(answer.citations[0]?.pageCitationAssetUrl).toBe(hardenedPageAssetUrl); + }); + it("returns only harness-selected artifacts when retrieval has extra media candidates", async () => { const frontAssetUrl = "https://blob.example/images/id-front.jpg"; const backAssetUrl = "https://blob.example/images/id-back.jpg"; diff --git a/src/domains/chat/page-citation-assets.test.ts b/src/domains/chat/page-citation-assets.test.ts index f7a1812..8a251cf 100644 --- a/src/domains/chat/page-citation-assets.test.ts +++ b/src/domains/chat/page-citation-assets.test.ts @@ -30,30 +30,6 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { ) }) - it("uses server-provided top-level page assets", async () => { - const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ - results: [ - makeRetrievalResult({ - chunkType: "page", - pageAssets: [ - { - pageNum: 4, - artifactRef: "page_citation_assets/page-4.png", - assetUrl: "https://assets.example/pages/page-4.png", - contentType: "image/png", - source: "knowhere-rendered-page-citation-source", - }, - ], - } as Partial), - ], - sources: [makeSource()], - }) - - expect(result?.pageCitationAssetUrl).toBe( - "https://assets.example/pages/page-4.png", - ) - }) - it("chooses the asset matching the citation page metadata", async () => { const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ results: [ diff --git a/src/domains/chat/page-citation-assets.ts b/src/domains/chat/page-citation-assets.ts index dccdd88..4f29a50 100644 --- a/src/domains/chat/page-citation-assets.ts +++ b/src/domains/chat/page-citation-assets.ts @@ -61,10 +61,7 @@ function getDirectPageCitationAsset( result: RetrievalResult, pageNumbers: readonly number[], ): PageCitationAssetCandidate | null { - const candidates = [ - isRecord(result) ? result.pageAssets : undefined, - result.metadata?.pageAssets, - ].flatMap(parsePageCitationAssetCandidates) + const candidates = parsePageCitationAssetCandidates(result.metadata?.pageAssets) if (pageNumbers.length > 0) { const matchingCandidates = candidates.filter((candidate) => From 3377032d6c7e5c414d35bb1b00640717481bc3d5 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 4 Jul 2026 00:08:39 +0800 Subject: [PATCH 05/51] Sync page citation assets to Blob --- pnpm-lock.yaml | 265 +----------------- src/domains/chat/index.ts | 1 + src/domains/chat/page-citation-assets.test.ts | 36 ++- src/domains/chat/page-citation-assets.ts | 76 ++++- .../sources/parse-result-storage-adapter.ts | 99 +++++++ .../source-reconcile-route-workflow.test.ts | 58 +++- .../source-reconcile-route-workflow.ts | 28 ++ 7 files changed, 297 insertions(+), 266 deletions(-) create mode 100644 src/domains/sources/parse-result-storage-adapter.ts diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 978e912..94e3f05 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -28,7 +28,7 @@ importers: version: 1.1.0 '@ontos-ai/knowhere-sdk': specifier: file:../knowhere-node-sdk - version: file:../knowhere-node-sdk(ioredis@5.11.1) + version: file:../knowhere-node-sdk '@radix-ui/react-alert-dialog': specifier: ^1.1.15 version: 1.1.15(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4) @@ -442,19 +442,6 @@ packages: resolution: {integrity: sha512-EXnJjIy6zQ3nUO/MZ+ynWUb8B895KZPotd1++oTs9JjDkplwM7cb6zo8Zq2zU6piwq+KflO7amXbEfj1UMpHkw==} hasBin: true - '@effect/platform-node-shared@4.0.0-beta.93': - resolution: {integrity: sha512-XUqZ2u5GglBqY8q2jj4Q7GjN5K/enedk8auZM9rY/l5a/myaQTrQp3QnvpIK4/Yg0WFjLGuctGPMKWRk3OLIrA==} - engines: {node: '>=18.0.0'} - peerDependencies: - effect: ^4.0.0-beta.93 - - '@effect/platform-node@4.0.0-beta.93': - resolution: {integrity: sha512-QagsCGR0ZOXaCQqS5qGR2mcDng4LiP2bYhiiX1D6UC8cT9vsusVVOHiJWn8CupeDx+yVnPcu81QmA/SDt6GM1w==} - engines: {node: '>=18.0.0'} - peerDependencies: - effect: ^4.0.0-beta.93 - ioredis: ^5.7.0 - '@effect/platform@0.96.1': resolution: {integrity: sha512-cjB1QZZYEP8JXCFNGvBLVi0T6YUBQTmOVEUA3SDbiQ6RUO+p6CE3eyD2vMWmrz5nE8yY5QSAuOV9v0boEcUv+A==} peerDependencies: @@ -1197,9 +1184,6 @@ packages: '@types/node': optional: true - '@ioredis/commands@1.10.0': - resolution: {integrity: sha512-UmeW7z4LfctwoQ5wkhVzgq8tXkreED2xZGpX+Bg+zA+WJFZCT6c062AfCK/Dfk81xZnnwdhJCUMkitihRaoC2Q==} - '@jridgewell/gen-mapping@0.3.13': resolution: {integrity: sha512-2kkt/7niJ6MgEPxF0bYdQ6etZaA+fQvDcLKckhy1yIQOzaoKjBBjSj63/aLVjYE3qhRt5dvM+uUyfCg6UKCBbA==} @@ -1231,61 +1215,31 @@ packages: cpu: [arm64] os: [darwin] - '@msgpackr-extract/msgpackr-extract-darwin-arm64@3.0.4': - resolution: {integrity: sha512-LCkGo6JDfaBhgST7UpPWgNgLINpcpabaHfyz5OBx75nUYxBsaEPxjnyNjWpeb/xBup/682QnBfRBy2/LvPutZQ==} - cpu: [arm64] - os: [darwin] - '@msgpackr-extract/msgpackr-extract-darwin-x64@3.0.3': resolution: {integrity: sha512-mdzd3AVzYKuUmiWOQ8GNhl64/IoFGol569zNRdkLReh6LRLHOXxU4U8eq0JwaD8iFHdVGqSy4IjFL4reoWCDFw==} cpu: [x64] os: [darwin] - '@msgpackr-extract/msgpackr-extract-darwin-x64@3.0.4': - resolution: {integrity: sha512-zExlW9zUJKZH/tOtVMttwjKa4Xm/3KcNjnE3dPN92uCktwavMxpgCA3MoJK/DOnTWsQgo224OaST27/mPNAf+w==} - cpu: [x64] - os: [darwin] - '@msgpackr-extract/msgpackr-extract-linux-arm64@3.0.3': resolution: {integrity: sha512-YxQL+ax0XqBJDZiKimS2XQaf+2wDGVa1enVRGzEvLLVFeqa5kx2bWbtcSXgsxjQB7nRqqIGFIcLteF/sHeVtQg==} cpu: [arm64] os: [linux] - '@msgpackr-extract/msgpackr-extract-linux-arm64@3.0.4': - resolution: {integrity: sha512-dgX0P/9wGPJeHFBG+ZmhgE6bmtMt7NP5CRBGyyktpopdk/mW4POnrpQsSLtKI1dwpc+pPLuXHDh6vvskyQE/sw==} - cpu: [arm64] - os: [linux] - '@msgpackr-extract/msgpackr-extract-linux-arm@3.0.3': resolution: {integrity: sha512-fg0uy/dG/nZEXfYilKoRe7yALaNmHoYeIoJuJ7KJ+YyU2bvY8vPv27f7UKhGRpY6euFYqEVhxCFZgAUNQBM3nw==} cpu: [arm] os: [linux] - '@msgpackr-extract/msgpackr-extract-linux-arm@3.0.4': - resolution: {integrity: sha512-Tg3yX65f5GbtXLkrYEHE5oibZG9epyYWas7FogTTEJeDEF9JlXJzKgXaNhT3UXlTOeA+AfZpYZYZ0uPj7Cfquw==} - cpu: [arm] - os: [linux] - '@msgpackr-extract/msgpackr-extract-linux-x64@3.0.3': resolution: {integrity: sha512-cvwNfbP07pKUfq1uH+S6KJ7dT9K8WOE4ZiAcsrSes+UY55E/0jLYc+vq+DO7jlmqRb5zAggExKm0H7O/CBaesg==} cpu: [x64] os: [linux] - '@msgpackr-extract/msgpackr-extract-linux-x64@3.0.4': - resolution: {integrity: sha512-8TNXMEjJc3QEy7R/x1INhgiU+XakDAFUzBhaz7+Rbrs8NH5UQeHQxxmzsSBJGyV6I1jW79undiQm8tOI+D+8FQ==} - cpu: [x64] - os: [linux] - '@msgpackr-extract/msgpackr-extract-win32-x64@3.0.3': resolution: {integrity: sha512-x0fWaQtYp4E6sktbsdAqnehxDgEc/VwM7uLsRCYWaiGu0ykYdZPiS8zCWdnjHwyiumousxfBm4SO31eXqwEZhQ==} cpu: [x64] os: [win32] - '@msgpackr-extract/msgpackr-extract-win32-x64@3.0.4': - resolution: {integrity: sha512-CmCXPQrkbwExx3j946/PtHWHbYJiCRBRDl4BlkRQcJB/YOwQxJRTpoo7aTsortjgoJ1x7opzTSxn7C+ASSLVjQ==} - cpu: [x64] - os: [win32] - '@mswjs/interceptors@0.41.8': resolution: {integrity: sha512-pRLMNKTSGRoLq+KnEB/7OY5vijw1XmcheAAOiv6pj7W1FG32kAGqj1C/RK/cqxRGr1Fh+zBi8sDur8kj3EQv6A==} engines: {node: '>=18'} @@ -2417,9 +2371,6 @@ packages: '@types/validate-npm-package-name@4.0.2': resolution: {integrity: sha512-lrpDziQipxCEeK5kWxvljWYhUvOiB2A9izZd9B2AFarYAkqZshb4lPbRs7zKEic6eGtH8V/2qJW+dPp9OtF6bw==} - '@types/ws@8.18.1': - resolution: {integrity: sha512-ThVF6DCVhA8kUGy+aazFQ4kXQ7E1Ty7A3ypFOe0IcJV8O/M511G99AW24irKrW56Wt44yG9+ij8FaqoBGkuBXg==} - '@typescript-eslint/eslint-plugin@8.59.2': resolution: {integrity: sha512-j/bwmkBvHUtPNxzuWe5z6BEk3q54YRyGlBXkSsmfoih7zNrBvl5A9A98anlp/7JbyZcWIJ8KXo/3Tq/DjFLtuQ==} engines: {node: ^18.18.0 || ^20.9.0 || >=21.1.0} @@ -2920,10 +2871,6 @@ packages: resolution: {integrity: sha512-eYm0QWBtUrBWZWG0d386OGAw16Z995PiOVo2B7bjWSbHedGl5e0ZWaq65kOGgUSNesEIDkB9ISbTg/JK9dhCZA==} engines: {node: '>=6'} - cluster-key-slot@1.1.1: - resolution: {integrity: sha512-rwHwUfXL40Chm1r08yrhU3qpUvdVlgkKNeyeGPOxnW8/SyVDvgRaed/Uz54AqWNaTCAThlj6QAs3TZcKI0xDEw==} - engines: {node: '>=0.10.0'} - code-block-writer@13.0.3: resolution: {integrity: sha512-Oofo0pq3IKnsFtuHqSF7TqBfr71aeyZDVJ0HpmqB7FBM2qEigL0iPONSCZSO9pE9dZTAxANe5XHG9Uy0YMv8cg==} @@ -3106,10 +3053,6 @@ packages: resolution: {integrity: sha512-ZySD7Nf91aLB0RxL4KGrKHBXl7Eds1DAmEdcoVawXnLD7SDhpNgtuII2aAkg7a7QS41jxPSZ17p4VdGnMHk3MQ==} engines: {node: '>=0.4.0'} - denque@2.1.0: - resolution: {integrity: sha512-HVQE3AAb/pxF8fQAoiqpvg9i3evqug3hoiwakOyZAwJm+6vZehbkYXZ0l4JxS+I3QxM97v5aaRNhj8v5oBhekw==} - engines: {node: '>=0.10'} - depd@2.0.0: resolution: {integrity: sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==} engines: {node: '>= 0.8'} @@ -3265,9 +3208,6 @@ packages: effect@3.21.2: resolution: {integrity: sha512-rXd2FGDM8KdjSIrc+mqEELo7ScW7xTVxEf1iInmPSpIde9/nyGuFM710cjTo7/EreGXiUX2MOonPpprbz2XHCg==} - effect@4.0.0-beta.93: - resolution: {integrity: sha512-wNS5MKFa3C42uBfIDik2oJ78lhpoYz2hN4oBR0229BeeDCIrkg/FiOvoiPGdCVlWa7MEKxEL5I0f8AILVHSD9A==} - electron-to-chromium@1.5.351: resolution: {integrity: sha512-9D7Iqx8RImSvCnOsj86rCH6eQjZFQoM04Jn6HnZVM0Nu/G58/gmKYQ1d12MZTbjQbQSTGI8nwEy07ErsA2slLA==} @@ -3540,10 +3480,6 @@ packages: resolution: {integrity: sha512-h5+1OzzfCC3Ef7VbtKdcv7zsstUQwUDlYpUTvjeUsJAssPgLn7QzbboPtL5ro04Mq0rPOsMzl7q5hIbRs2wD1A==} engines: {node: '>=8.0.0'} - fast-check@4.8.0: - resolution: {integrity: sha512-GOJ158CUMnN6cSahsv4+ExARvIDuzzinFjkp0E9WtiBa5zcVeLozVkWaE4IzFcc+Y48Wp1EDlUZsXRyAztQcSg==} - engines: {node: '>=12.17.0'} - fast-deep-equal@3.1.3: resolution: {integrity: sha512-f3qQ9oQy9j2AhBe/H9VC91wLmKBCCU/gDOnKNAYG5hswO7BLKj09Hc5HYNz9cGI++xlpDCIgDaitVs03ATR84Q==} @@ -3863,10 +3799,6 @@ packages: inherits@2.0.4: resolution: {integrity: sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==} - ini@7.0.0: - resolution: {integrity: sha512-ifK0CgjALofS5bkrcTy4RaQ9Vx2Knf/eLeIO+NaswQEpH1UblrtTSCIvN71qQDMq0PeQ/SSPojvEJp9vvvfr+w==} - engines: {node: ^22.22.2 || ^24.15.0 || >=26.0.0} - inline-style-parser@0.2.7: resolution: {integrity: sha512-Nb2ctOyNR8DqQoR0OwRG95uNWIC0C1lCgf5Naz5H6Ji72KZ8OcFZLz2P5sNgwlyoJ8Yif11oMuYs5pBQa86csA==} @@ -3874,10 +3806,6 @@ packages: resolution: {integrity: sha512-4gd7VpWNQNB4UKKCFFVcp1AVv+FMOgs9NKzjHKusc8jTMhd5eL1NqQqOpE0KzMds804/yHlglp3uxgluOqAPLw==} engines: {node: '>= 0.4'} - ioredis@5.11.1: - resolution: {integrity: sha512-ehuGcf94bQXhfagULNXrJdfnWO38v070jxSx/qE87Kjzmu2fU7ro5EFAb+OPituLqgfyuQaym5DlrNydW2sJ9A==} - engines: {node: '>=12.22.0'} - ip-address@10.1.0: resolution: {integrity: sha512-XXADHxXmvT9+CRxhXg56LJovE+bmWnEWB78LB83VZTprKTmaC5QfruXocxzTZ2Kl0DNwKuBdlIhjL8LeY8Sf8Q==} engines: {node: '>= 12'} @@ -4182,9 +4110,6 @@ packages: resolution: {integrity: sha512-o+NO+8WrRiQEE4/7nwRJhN1HWpVmJm511pBHUxPLtp0BUISzlBplORYSmTclCnJvQq2tKu/sgl3xVpkc7ZWuQQ==} engines: {node: '>=6'} - kubernetes-types@1.30.0: - resolution: {integrity: sha512-Dew1okvhM/SQcIa2rcgujNndZwU8VnSapDgdxlYoB84ZlpAD43U6KLAFqYo17ykSFGHNPrg0qry0bP+GJd9v7Q==} - language-subtag-registry@0.3.23: resolution: {integrity: sha512-0K65Lea881pHotoGEa5gDlMxt3pctLi2RplBb7Ezh4rRdLEOtgi7n4EwK9lamnUCkKBqaeKRVebTq6BAxSkpXQ==} @@ -4509,11 +4434,6 @@ packages: resolution: {integrity: sha512-Lbgzdk0h4juoQ9fCKXW4by0UJqj+nOOrI9MJ1sSj4nI8aI2eo1qmvQEie4VD1glsS250n15LsWsYtCugiStS5A==} engines: {node: '>=18'} - mime@4.1.0: - resolution: {integrity: sha512-X5ju04+cAzsojXKes0B/S4tcYtFAJ6tTMuSPBEn9CPGlrWr8Fiw7qYeLT0XyH80HSoAoqWCaz+MWKh22P7G1cw==} - engines: {node: '>=16'} - hasBin: true - mimic-fn@2.1.0: resolution: {integrity: sha512-OqbOk5oEQeAZ8WXWydlu9HJjz9WVdEIvamMCcXmuqUYjTknH/sqsWvhQ3vgwKFRR1HpjvNBKQ37nbJgYzGqGcg==} engines: {node: '>=6'} @@ -4539,16 +4459,9 @@ packages: resolution: {integrity: sha512-P0efT1C9jIdVRefqjzOQ9Xml57zpOXnIuS+csaB4MdZbTdmGDLo8XhzBG1N7aO11gKDDkJvBLULeFTo46wwreA==} hasBin: true - msgpackr-extract@3.0.4: - resolution: {integrity: sha512-4kmO/MdyUIkLIvTPr8VHLil4AtoKIoniWPIEk5+CDy0xnWC84azhSFmuJ7PxZdsYtiP5kEeQsORAVIeMgxT+Hw==} - hasBin: true - msgpackr@1.11.12: resolution: {integrity: sha512-RBdJ1Un7yGlXWajrkxcSa93nvQ0w4zBf60c0yYv7YtBelP8H2FA7XsfBbMHtXKXUMUxH7zV3Zuozh+kUQWhHvg==} - msgpackr@2.0.4: - resolution: {integrity: sha512-o1C5KRmuRt+apqMr1HuGSqWStZoRBUpEsCsl15uM9VdAF1qHLtvMOU2En747EnTyEl6c4pzPewRMFF31s1CNbA==} - msw@2.14.3: resolution: {integrity: sha512-kk8G5cocVlJ4wsKMGZegn2H6XLOEKjbA+nSJE2354e/SRp4mDicCHUYnMXpymzVcVDCs+GUAsmNqSn+yHv4T2A==} engines: {node: '>=18'} @@ -4786,10 +4699,6 @@ packages: resolution: {integrity: sha512-DlOzet0HO7OEnmUmB6wWGJrrdvbyJKftI1bhMitK7O2N8W2gc757yyYBbINy9IDafXAV9wmKr9t7xsTaNKRG5Q==} engines: {node: '>=20.16.0 || >=22.3.0'} - pdfjs-dist@6.1.200: - resolution: {integrity: sha512-o8MolyzirkkLrcdsae/HEOiIcXWI7DS5zGpvqW8xTC2YUsW30rltFw2bDGvw/fskUdEMrQm2br68jzDS5BH2vw==} - engines: {node: '>=22.13.0 || >=24'} - picocolors@1.1.1: resolution: {integrity: sha512-xceH2snhtb5M9liqDsmEw56le376mTZkEX/jEb/RxNFyegNul7eNslCXP9FDj/Lcu0X8KEyMceP2ntpaHrDEVA==} @@ -4889,9 +4798,6 @@ packages: pure-rand@6.1.0: resolution: {integrity: sha512-bVWawvoZoBYpp6yIoQtQXHZjmz35RSVHnUOTefl8Vcjr8snTPY1wnpSPMWekcFwbxI6gtmT7rSYPFvz71ldiOA==} - pure-rand@8.4.1: - resolution: {integrity: sha512-c58R2+SPFcSIPXoU834QN/KPDDOSd8sXcSrqf6e83Me6Rrp1EYkxukkjXMVrKvKaADs1SOyNkWdfvLf6zY8qLQ==} - qs@6.15.1: resolution: {integrity: sha512-6YHEFRL9mfgcAvql/XhwTvf5jKcOiiupt2FiJxHkiX1z4j7WL8J/jRHYLluORvc1XxB5rV20KoeK00gVJamspg==} engines: {node: '>=0.6'} @@ -4978,14 +4884,6 @@ packages: resolution: {integrity: sha512-YTUo+Flmw4ZXiWfQKGcwwc11KnoRAYgzAE2E7mXKCjSviTKShtxBsN6YUUBB2gtaBzKzeKunxhUwNHQuRryhWA==} engines: {node: '>= 4'} - redis-errors@1.2.0: - resolution: {integrity: sha512-1qny3OExCf0UvUV/5wpYKf2YwPcOqXzkwKKSmKHiE6ZMQs5heeE/c8eXK+PNllPvmjgAbfnsbpkGZWy8cBpn9w==} - engines: {node: '>=4'} - - redis-parser@3.0.0: - resolution: {integrity: sha512-DJnGAeenTdpMEH6uAJRK/uiyEIH9WVsUmoLwzudwGJUwZPp80PDBWPHXSAGNPwNvIXAbe7MSUB1zQFugFml66A==} - engines: {node: '>=4'} - reflect.getprototypeof@1.0.10: resolution: {integrity: sha512-00o4I+DVrefhv+nX0ulyi3biSHCPDe+yLv5o/p6d/UVlirijB8E16FtfwSAi4g3tcqrQ4lRAqQSoFEZJehYEcw==} engines: {node: '>= 0.4'} @@ -5195,9 +5093,6 @@ packages: stackback@0.0.2: resolution: {integrity: sha512-1XMJE5fQo1jGH6Y/7ebnwPOBEkIEnT4QF32d5R1+VXdXveM0IBMJt8zfaxX1P3QhVwrYe+576+jkANtSS2mBbw==} - standard-as-callback@2.1.0: - resolution: {integrity: sha512-qoRRSyROncaz1z0mvYqIE4lCd9p2R90i6GxW3uZv5ucSu8tU7B5HXUP1gG8pVZsYNVaXjk8ClXHPttLyxAL48A==} - statuses@2.0.2: resolution: {integrity: sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==} engines: {node: '>= 0.8'} @@ -5371,10 +5266,6 @@ packages: resolution: {integrity: sha512-o5sSPKEkg/DIQNmH43V0/uerLrpzVedkUh8tGNvaeXpfpuwjKenlSox/2O/BTlZUtEe+JG7s5YhEz608PlAHRA==} engines: {node: '>=0.6'} - toml@4.1.2: - resolution: {integrity: sha512-m0vXfHODcw3gk+KONAOlVQ5yNHc3yS3B1ybM3HS1vqDoS0RWTDDVBVVTYi8hH0k+2OM1vmo9fb1WX9EVqjqfHA==} - engines: {node: '>=20'} - tough-cookie@6.0.1: resolution: {integrity: sha512-LktZQb3IeoUWB9lqR5EWTHgW/VTITCXg4D21M+lvybRVdylLrRMnqaIONLVb5mav8vM19m44HIcGq4qASeu2Qw==} engines: {node: '>=16'} @@ -5474,10 +5365,6 @@ packages: resolution: {integrity: sha512-xXnp4kTyor2Zq+J1FfPI6Eq3ew5h6Vl0F/8d9XU5zZQf1tX9s2Su1/3PiMmUANFULpmksxkClamIZcaUqryHsQ==} engines: {node: '>=20.18.1'} - undici@8.6.0: - resolution: {integrity: sha512-l2FlC6I510GawyEd1qgcE/okihKrzy+BRTEBlu6T0fdbM9m5yxtIH5Oa3ysRsH0zC4EhmWUEaSDsy2QngBeRlw==} - engines: {node: '>=22.19.0'} - unicorn-magic@0.3.0: resolution: {integrity: sha512-+QBBXBCvifc56fsbuxZQ6Sic3wqqc3WWaqxs58gvJrcOuN83HGTCwz3oS5phzU9LthRNE9VrJCFCLUgHeeFnfA==} engines: {node: '>=18'} @@ -5551,10 +5438,6 @@ packages: util-deprecate@1.0.2: resolution: {integrity: sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw==} - uuid@14.0.1: - resolution: {integrity: sha512-6ZxzVpzDXDa3bJWaHilVayA+BH/1zmxCJoVgvmqJnid/gPoKHxUrS/aC/T6LGQtNHT+XHG9fXPJB4d+IrU30Ew==} - hasBin: true - validate-npm-package-name@7.0.2: resolution: {integrity: sha512-hVDIBwsRruT73PbK7uP5ebUt+ezEtCmzZz3F59BSr2F6OVFnJ/6h8liuvdLrQ88Xmnk6/+xGGuq+pG9WwTuy3A==} engines: {node: ^20.17.0 || >=22.9.0} @@ -5721,18 +5604,6 @@ packages: wrappy@1.0.2: resolution: {integrity: sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ==} - ws@8.21.0: - resolution: {integrity: sha512-Vsp28b7DRcimFQvrqu2Wek3z1iYxDCWqHYB8Qsnk/S4RfaCQzPGPyBNuVjJV3cd6UiKtUtp6sNM77gWvzcCH+g==} - engines: {node: '>=10.0.0'} - peerDependencies: - bufferutil: ^4.0.1 - utf-8-validate: '>=5.0.2' - peerDependenciesMeta: - bufferutil: - optional: true - utf-8-validate: - optional: true - wsl-utils@0.3.1: resolution: {integrity: sha512-g/eziiSUNBSsdDJtCLB8bdYEUMj4jR7AGeUo96p/3dTafgjHhpF4RiCFPiRILwjQoDXx5MqkBr4fwWtR3Ky4Wg==} engines: {node: '>=20'} @@ -6094,26 +5965,6 @@ snapshots: '@effect/language-service@0.85.1': {} - '@effect/platform-node-shared@4.0.0-beta.93(effect@4.0.0-beta.93)': - dependencies: - '@types/ws': 8.18.1 - effect: 4.0.0-beta.93 - ws: 8.21.0 - transitivePeerDependencies: - - bufferutil - - utf-8-validate - - '@effect/platform-node@4.0.0-beta.93(effect@4.0.0-beta.93)(ioredis@5.11.1)': - dependencies: - '@effect/platform-node-shared': 4.0.0-beta.93(effect@4.0.0-beta.93) - effect: 4.0.0-beta.93 - ioredis: 5.11.1 - mime: 4.1.0 - undici: 8.6.0 - transitivePeerDependencies: - - bufferutil - - utf-8-validate - '@effect/platform@0.96.1(effect@3.21.2)': dependencies: effect: 3.21.2 @@ -6580,8 +6431,6 @@ snapshots: optionalDependencies: '@types/node': 20.19.39 - '@ioredis/commands@1.10.0': {} - '@jridgewell/gen-mapping@0.3.13': dependencies: '@jridgewell/sourcemap-codec': 1.5.5 @@ -6626,39 +6475,21 @@ snapshots: '@msgpackr-extract/msgpackr-extract-darwin-arm64@3.0.3': optional: true - '@msgpackr-extract/msgpackr-extract-darwin-arm64@3.0.4': - optional: true - '@msgpackr-extract/msgpackr-extract-darwin-x64@3.0.3': optional: true - '@msgpackr-extract/msgpackr-extract-darwin-x64@3.0.4': - optional: true - '@msgpackr-extract/msgpackr-extract-linux-arm64@3.0.3': optional: true - '@msgpackr-extract/msgpackr-extract-linux-arm64@3.0.4': - optional: true - '@msgpackr-extract/msgpackr-extract-linux-arm@3.0.3': optional: true - '@msgpackr-extract/msgpackr-extract-linux-arm@3.0.4': - optional: true - '@msgpackr-extract/msgpackr-extract-linux-x64@3.0.3': optional: true - '@msgpackr-extract/msgpackr-extract-linux-x64@3.0.4': - optional: true - '@msgpackr-extract/msgpackr-extract-win32-x64@3.0.3': optional: true - '@msgpackr-extract/msgpackr-extract-win32-x64@3.0.4': - optional: true - '@mswjs/interceptors@0.41.8': dependencies: '@open-draft/deferred-promise': 2.2.0 @@ -6903,21 +6734,13 @@ snapshots: '@nolyfill/is-core-module@1.0.39': {} - '@ontos-ai/knowhere-sdk@file:../knowhere-node-sdk(ioredis@5.11.1)': + '@ontos-ai/knowhere-sdk@file:../knowhere-node-sdk': dependencies: - '@effect/platform-node': 4.0.0-beta.93(effect@4.0.0-beta.93)(ioredis@5.11.1) - '@napi-rs/canvas': 1.0.2 axios: 1.18.1 - effect: 4.0.0-beta.93 jszip: 3.10.1 - pdfjs-dist: 6.1.200 - piscina: 5.2.0 transitivePeerDependencies: - - bufferutil - debug - - ioredis - supports-color - - utf-8-validate '@open-draft/deferred-promise@2.2.0': {} @@ -7580,10 +7403,6 @@ snapshots: '@types/validate-npm-package-name@4.0.2': {} - '@types/ws@8.18.1': - dependencies: - '@types/node': 20.19.39 - '@typescript-eslint/eslint-plugin@8.59.2(@typescript-eslint/parser@8.59.2(eslint@9.39.4(jiti@2.7.0))(typescript@6.0.3))(eslint@9.39.4(jiti@2.7.0))(typescript@6.0.3)': dependencies: '@eslint-community/regexpp': 4.12.2 @@ -8105,8 +7924,6 @@ snapshots: clsx@2.1.1: {} - cluster-key-slot@1.1.1: {} - code-block-writer@13.0.3: {} color-convert@2.0.1: @@ -8250,8 +8067,6 @@ snapshots: delayed-stream@1.0.0: {} - denque@2.1.0: {} - depd@2.0.0: {} dequal@2.0.3: {} @@ -8321,19 +8136,6 @@ snapshots: '@standard-schema/spec': 1.1.0 fast-check: 3.23.2 - effect@4.0.0-beta.93: - dependencies: - '@standard-schema/spec': 1.1.0 - fast-check: 4.8.0 - find-my-way-ts: 0.1.6 - ini: 7.0.0 - kubernetes-types: 1.30.0 - msgpackr: 2.0.4 - multipasta: 0.2.7 - toml: 4.1.2 - uuid: 14.0.1 - yaml: 2.9.0 - electron-to-chromium@1.5.351: {} emoji-regex@10.6.0: {} @@ -8849,10 +8651,6 @@ snapshots: dependencies: pure-rand: 6.1.0 - fast-check@4.8.0: - dependencies: - pure-rand: 8.4.1 - fast-deep-equal@3.1.3: {} fast-glob@3.3.1: @@ -9181,8 +8979,6 @@ snapshots: inherits@2.0.4: {} - ini@7.0.0: {} - inline-style-parser@0.2.7: {} internal-slot@1.1.0: @@ -9191,18 +8987,6 @@ snapshots: hasown: 2.0.3 side-channel: 1.1.0 - ioredis@5.11.1: - dependencies: - '@ioredis/commands': 1.10.0 - cluster-key-slot: 1.1.1 - debug: 4.4.3 - denque: 2.1.0 - redis-errors: 1.2.0 - redis-parser: 3.0.0 - standard-as-callback: 2.1.0 - transitivePeerDependencies: - - supports-color - ip-address@10.1.0: {} ipaddr.js@1.9.1: {} @@ -9484,8 +9268,6 @@ snapshots: kleur@4.1.5: {} - kubernetes-types@1.30.0: {} - language-subtag-registry@0.3.23: {} language-tags@1.0.9: @@ -9987,8 +9769,6 @@ snapshots: dependencies: mime-db: 1.54.0 - mime@4.1.0: {} - mimic-fn@2.1.0: {} mimic-function@5.0.1: {} @@ -10017,26 +9797,10 @@ snapshots: '@msgpackr-extract/msgpackr-extract-win32-x64': 3.0.3 optional: true - msgpackr-extract@3.0.4: - dependencies: - node-gyp-build-optional-packages: 5.2.2 - optionalDependencies: - '@msgpackr-extract/msgpackr-extract-darwin-arm64': 3.0.4 - '@msgpackr-extract/msgpackr-extract-darwin-x64': 3.0.4 - '@msgpackr-extract/msgpackr-extract-linux-arm': 3.0.4 - '@msgpackr-extract/msgpackr-extract-linux-arm64': 3.0.4 - '@msgpackr-extract/msgpackr-extract-linux-x64': 3.0.4 - '@msgpackr-extract/msgpackr-extract-win32-x64': 3.0.4 - optional: true - msgpackr@1.11.12: optionalDependencies: msgpackr-extract: 3.0.3 - msgpackr@2.0.4: - optionalDependencies: - msgpackr-extract: 3.0.4 - msw@2.14.3(@types/node@20.19.39)(typescript@6.0.3): dependencies: '@inquirer/confirm': 6.0.12(@types/node@20.19.39) @@ -10302,10 +10066,6 @@ snapshots: optionalDependencies: '@napi-rs/canvas': 0.1.100 - pdfjs-dist@6.1.200: - optionalDependencies: - '@napi-rs/canvas': 1.0.2 - picocolors@1.1.1: {} picomatch@2.3.2: {} @@ -10400,8 +10160,6 @@ snapshots: pure-rand@6.1.0: {} - pure-rand@8.4.1: {} - qs@6.15.1: dependencies: side-channel: 1.1.0 @@ -10508,12 +10266,6 @@ snapshots: tiny-invariant: 1.3.3 tslib: 2.8.1 - redis-errors@1.2.0: {} - - redis-parser@3.0.0: - dependencies: - redis-errors: 1.2.0 - reflect.getprototypeof@1.0.10: dependencies: call-bind: 1.0.9 @@ -10860,8 +10612,6 @@ snapshots: stackback@0.0.2: {} - standard-as-callback@2.1.0: {} - statuses@2.0.2: {} std-env@4.1.0: {} @@ -11034,8 +10784,6 @@ snapshots: toidentifier@1.0.1: {} - toml@4.1.2: {} - tough-cookie@6.0.1: dependencies: tldts: 7.0.30 @@ -11156,8 +10904,6 @@ snapshots: undici@7.25.0: {} - undici@8.6.0: {} - unicorn-magic@0.3.0: {} unified@11.0.5: @@ -11254,8 +11000,6 @@ snapshots: util-deprecate@1.0.2: {} - uuid@14.0.1: {} - validate-npm-package-name@7.0.2: {} vary@1.1.2: {} @@ -11402,8 +11146,6 @@ snapshots: wrappy@1.0.2: {} - ws@8.21.0: {} - wsl-utils@0.3.1: dependencies: is-wsl: 3.1.1 @@ -11419,7 +11161,8 @@ snapshots: yallist@3.1.1: {} - yaml@2.9.0: {} + yaml@2.9.0: + optional: true yargs-parser@21.1.1: {} diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index ec19628..bbf95b9 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -250,6 +250,7 @@ export const answerQuestionWithRetrieval = ( enrichRetrievalResultsWithPageCitationAssetUrls({ results: enrichedResults, sources: input.sources, + loadSourceAssetUrls: input.loadSourceAssetUrls, }), ) const artifacts = toChatArtifactViewsFromHarness(generatedAnswer, input.sources) diff --git a/src/domains/chat/page-citation-assets.test.ts b/src/domains/chat/page-citation-assets.test.ts index 8a251cf..6d0ee01 100644 --- a/src/domains/chat/page-citation-assets.test.ts +++ b/src/domains/chat/page-citation-assets.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it } from "vitest" +import { describe, expect, it, vi } from "vitest" import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" @@ -30,6 +30,40 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { ) }) + it("prefers stored Blob URLs over server-provided page asset URLs", async () => { + const loadSourceAssetUrls = vi.fn().mockResolvedValue({ + "page_citation_assets/page-2.png": + "https://blob.example/page_citation_assets/page-2.png", + }) + + const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ + results: [ + makeRetrievalResult({ + chunkType: "page", + metadata: { + pageNums: [2], + pageAssets: [ + { + pageNum: 2, + artifactRef: "page_citation_assets/page-2.png", + assetUrl: "https://assets.example/pages/page-2.png", + }, + ], + }, + }), + ], + sources: [makeSource()], + loadSourceAssetUrls, + }) + + expect(loadSourceAssetUrls).toHaveBeenCalledWith( + expect.objectContaining({ id: "source_1" }), + ) + expect(result?.pageCitationAssetUrl).toBe( + "https://blob.example/page_citation_assets/page-2.png", + ) + }) + it("chooses the asset matching the citation page metadata", async () => { const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ results: [ diff --git a/src/domains/chat/page-citation-assets.ts b/src/domains/chat/page-citation-assets.ts index 4f29a50..dd1d34d 100644 --- a/src/domains/chat/page-citation-assets.ts +++ b/src/domains/chat/page-citation-assets.ts @@ -3,6 +3,7 @@ import "server-only" import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" +import type { LoadSourceAssetUrls } from "./media-assets" export type PageCitationAssetRetrievalResult = RetrievalResult & { readonly pageCitationAssetUrl?: string @@ -11,26 +12,41 @@ export type PageCitationAssetRetrievalResult = RetrievalResult & { type EnrichRetrievalResultsWithPageCitationAssetUrlsInput = { readonly results: readonly RetrievalResult[] readonly sources: readonly Source[] + readonly loadSourceAssetUrls?: LoadSourceAssetUrls } type PageCitationAssetCandidate = { readonly pageNum: number + readonly artifactRef?: string readonly assetUrl?: string } export async function enrichRetrievalResultsWithPageCitationAssetUrls({ results, sources, + loadSourceAssetUrls, }: EnrichRetrievalResultsWithPageCitationAssetUrlsInput): Promise< PageCitationAssetRetrievalResult[] > { - void sources if (results.length === 0) return [] + const sourcesByDocumentId = new Map( + sources.flatMap((source): readonly [string, Source][] => + source.knowhereDocumentId ? [[source.knowhereDocumentId, source]] : [], + ), + ) + const assetUrlsBySourceId = new Map< + string, + Promise>> + >() + return Promise.all( results.map((result) => enrichRetrievalResultWithPageCitationAssetUrl({ result, + sourcesByDocumentId, + loadSourceAssetUrls, + assetUrlsBySourceId, }), ), ) @@ -38,11 +54,30 @@ export async function enrichRetrievalResultsWithPageCitationAssetUrls({ async function enrichRetrievalResultWithPageCitationAssetUrl(input: { readonly result: RetrievalResult + readonly sourcesByDocumentId: ReadonlyMap + readonly loadSourceAssetUrls?: LoadSourceAssetUrls + readonly assetUrlsBySourceId: Map< + string, + Promise>> + > }): Promise { if (!isPageResult(input.result)) return input.result const pageNumbers = getPageNumbers(input.result.metadata) const directAsset = getDirectPageCitationAsset(input.result, pageNumbers) + const sourceAssetUrl = await getStoredPageCitationAssetUrl({ + result: input.result, + directAsset, + sourcesByDocumentId: input.sourcesByDocumentId, + loadSourceAssetUrls: input.loadSourceAssetUrls, + assetUrlsBySourceId: input.assetUrlsBySourceId, + }) + if (sourceAssetUrl) { + return { + ...input.result, + pageCitationAssetUrl: sourceAssetUrl, + } + } if (directAsset?.assetUrl) { return { ...input.result, @@ -53,6 +88,31 @@ async function enrichRetrievalResultWithPageCitationAssetUrl(input: { return input.result } +async function getStoredPageCitationAssetUrl(input: { + readonly result: RetrievalResult + readonly directAsset: PageCitationAssetCandidate | null + readonly sourcesByDocumentId: ReadonlyMap + readonly loadSourceAssetUrls?: LoadSourceAssetUrls + readonly assetUrlsBySourceId: Map< + string, + Promise>> + > +}): Promise { + const artifactRef = getTrimmedString(input.directAsset?.artifactRef) + const documentId = getTrimmedString(input.result.source.documentId) + if (!artifactRef || !documentId || !input.loadSourceAssetUrls) return null + + const source = input.sourcesByDocumentId.get(documentId) + if (!source) return null + + const assetUrls = await getCachedSourceAssetUrls( + source, + input.loadSourceAssetUrls, + input.assetUrlsBySourceId, + ) + return getTrimmedString(assetUrls[artifactRef]) +} + function isPageResult(result: RetrievalResult): boolean { return result.chunkType.toLowerCase() === "page" } @@ -90,12 +150,26 @@ function parsePageCitationAssetCandidates( return [ { pageNum, + artifactRef: getTrimmedString(item.artifactRef) ?? undefined, assetUrl: getTrimmedString(item.assetUrl) ?? undefined, }, ] }) } +async function getCachedSourceAssetUrls( + source: Source, + loadSourceAssetUrls: LoadSourceAssetUrls, + cache: Map>>>, +): Promise>> { + let cached = cache.get(source.id) + if (!cached) { + cached = loadSourceAssetUrls(source).catch(() => ({})) + cache.set(source.id, cached) + } + return cached +} + function getPageNumbers( metadata: Readonly> | undefined, ): readonly number[] { diff --git a/src/domains/sources/parse-result-storage-adapter.ts b/src/domains/sources/parse-result-storage-adapter.ts new file mode 100644 index 0000000..908dd34 --- /dev/null +++ b/src/domains/sources/parse-result-storage-adapter.ts @@ -0,0 +1,99 @@ +import "server-only" + +import { head, put, BlobNotFoundError } from "@vercel/blob" +import type { + KnowhereAssetStorageObject, + KnowhereAssetStorageOptions, +} from "@ontos-ai/knowhere-sdk" + +export type ParsedResultAssetStorageInput = { + readonly workspaceId: string + readonly sourceId: string +} + +export type ParsedResultAssetIndex = { + readonly workspaceId: string + readonly sourceId: string + readonly assetUrlsByFilePath: Readonly> + readonly updatedAt: string +} + +const parsedResultDirectoryName = "parsed-result" +const parsedResultAssetIndexFileName = "asset-index.json" + +export function createParsedResultStorageAdapter({ + workspaceId, + sourceId, +}: ParsedResultAssetStorageInput): KnowhereAssetStorageOptions { + return { + adapter: { + async headObject(key) { + try { + const blob = await head(key) + return { + key, + contentType: blob.contentType, + contentLength: blob.size, + url: blob.url, + } + } catch (error) { + if (error instanceof BlobNotFoundError) return null + throw error + } + }, + async getObjectUrl(key) { + try { + return (await head(key)).url + } catch (error) { + if (error instanceof BlobNotFoundError) return null + throw error + } + }, + async writeObject(input: KnowhereAssetStorageObject) { + const blob = await put(input.key, Buffer.from(input.body), { + access: "public", + addRandomSuffix: false, + allowOverwrite: true, + contentType: input.contentType, + multipart: true, + }) + + return { + key: blob.pathname, + url: blob.url, + } + }, + }, + keyPrefix: getParsedResultBlobPrefix(workspaceId, sourceId), + skipExisting: true, + } +} + +export async function writeParsedResultAssetIndex(input: { + readonly workspaceId: string + readonly sourceId: string + readonly assetUrlsByFilePath: Readonly> +}): Promise { + const index: ParsedResultAssetIndex = { + workspaceId: input.workspaceId, + sourceId: input.sourceId, + assetUrlsByFilePath: input.assetUrlsByFilePath, + updatedAt: new Date().toISOString(), + } + const pathname = `${getParsedResultBlobPrefix( + input.workspaceId, + input.sourceId, + )}/${parsedResultAssetIndexFileName}` + const blob = await put(pathname, JSON.stringify(index, null, 2), { + access: "public", + addRandomSuffix: false, + allowOverwrite: true, + contentType: "application/json; charset=utf-8", + }) + + return blob.url +} + +function getParsedResultBlobPrefix(workspaceId: string, sourceId: string): string { + return `workspaces/${workspaceId}/sources/${sourceId}/${parsedResultDirectoryName}` +} diff --git a/src/domains/sources/source-reconcile-route-workflow.test.ts b/src/domains/sources/source-reconcile-route-workflow.test.ts index 14e9b03..170941f 100644 --- a/src/domains/sources/source-reconcile-route-workflow.test.ts +++ b/src/domains/sources/source-reconcile-route-workflow.test.ts @@ -1,6 +1,9 @@ import { afterEach, describe, expect, it, vi } from "vitest" const mocks = vi.hoisted(() => ({ + createParsedResultStorageAdapter: vi.fn(), + saveParseResult: vi.fn(), + writeParsedResultAssetIndex: vi.fn(), loggerError: vi.fn(), loggerInfo: vi.fn(), loggerWarn: vi.fn(), @@ -18,9 +21,15 @@ vi.mock("@/domains/sources/source-reconcile-workflow", () => ({ vi.mock("@/domains/sources/workflow-runtime", () => ({ sourceWorkflowRuntime: { markFailed: mocks.markFailed, + saveParseResult: mocks.saveParseResult, }, })) +vi.mock("./parse-result-storage-adapter", () => ({ + createParsedResultStorageAdapter: mocks.createParsedResultStorageAdapter, + writeParsedResultAssetIndex: mocks.writeParsedResultAssetIndex, +})) + vi.mock("@/integrations/knowhere", () => ({ makeKnowhereClient: mocks.makeKnowhereClient, })) @@ -57,7 +66,7 @@ describe("sourceReconcileRouteWorkflow", () => { }) }) - it("marks the source ready when Knowhere publishes a document id", async () => { + it("caches parsed assets before marking the source ready", async () => { const context = createWorkflowContext() const continuations: ContinuationTriggerInput[] = [] const restore = @@ -66,8 +75,25 @@ describe("sourceReconcileRouteWorkflow", () => { continuations.push(input) }, ) - const client = { jobs: {} } + const storageAdapter = { + adapter: { + writeObject: vi.fn(), + }, + keyPrefix: "workspaces/workspace_1/sources/source_1/parsed-result", + } + const cacheJobResult = vi.fn().mockResolvedValue({ + assetUrlsByFilePath: { + "page_citation_assets/page-1.png": + "https://blob.example/page_citation_assets/page-1.png", + }, + }) + const client = { jobs: {}, knowledge: { cacheJobResult } } + mocks.createParsedResultStorageAdapter.mockReturnValue(storageAdapter) mocks.makeKnowhereClient.mockReturnValue(client) + mocks.writeParsedResultAssetIndex.mockResolvedValue( + "https://blob.example/asset-index.json", + ) + mocks.saveParseResult.mockResolvedValue({ id: "parse_result_1" }) mocks.pollSourceReconciliation.mockResolvedValue({ kind: "ready-to-prepare", jobId: "job_1", @@ -90,12 +116,38 @@ describe("sourceReconcileRouteWorkflow", () => { restore() } + expect(mocks.createParsedResultStorageAdapter).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "source_1", + }) + expect(cacheJobResult).toHaveBeenCalledWith({ + jobId: "job_1", + storageAdapter, + }) + expect(mocks.writeParsedResultAssetIndex).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "source_1", + assetUrlsByFilePath: { + "page_citation_assets/page-1.png": + "https://blob.example/page_citation_assets/page-1.png", + }, + }) + expect(mocks.saveParseResult).toHaveBeenCalledWith( + "workspace_1", + "source_1", + { + resultBlobUrl: "https://blob.example/asset-index.json", + assetUrlsByFilePath: { + "page_citation_assets/page-1.png": + "https://blob.example/page_citation_assets/page-1.png", + }, + }, + ) expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ workspaceId: "workspace_1", sourceId: "source_1", documentId: "doc_1", }) - expect(client).toEqual({ jobs: {} }) expect(mocks.loggerWarn).not.toHaveBeenCalled() expect(continuations).toEqual([]) }) diff --git a/src/domains/sources/source-reconcile-route-workflow.ts b/src/domains/sources/source-reconcile-route-workflow.ts index a1e732a..cc03d40 100644 --- a/src/domains/sources/source-reconcile-route-workflow.ts +++ b/src/domains/sources/source-reconcile-route-workflow.ts @@ -8,6 +8,10 @@ import { } from "@/domains/sources/source-reconcile-workflow" import { makeKnowhereClient } from "@/integrations/knowhere" import { logger } from "@/lib/logger" +import { + createParsedResultStorageAdapter, + writeParsedResultAssetIndex, +} from "./parse-result-storage-adapter" import { sourceWorkflowRuntime } from "./workflow-runtime" type ReconcilePayload = { @@ -118,6 +122,30 @@ async function runPollAndMirrorWorkflow(input: { return } + const assetUrlsByFilePath = await context.run("cache-job-result-assets", async () => { + const cachedResult = await client.knowledge.cacheJobResult({ + jobId: jobToPrepare.jobId, + storageAdapter: createParsedResultStorageAdapter({ + workspaceId, + sourceId, + }), + }) + return cachedResult.assetUrlsByFilePath ?? {} + }) + const resultBlobUrl = await context.run("write-parse-result-asset-index", async () => + writeParsedResultAssetIndex({ + workspaceId, + sourceId, + assetUrlsByFilePath, + }), + ) + await context.run("save-parse-result-assets", async () => + sourceWorkflowRuntime.saveParseResult(workspaceId, sourceId, { + resultBlobUrl, + assetUrlsByFilePath, + }), + ) + const ready = await context.run("source-ready", async () => markSourceReadyAfterReconciliation({ workspaceId, From 907083e5bf4785b610539a71af9e4dd738ca19dd Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 4 Jul 2026 01:23:08 +0800 Subject: [PATCH 06/51] Read parsed chunks from Blob snapshot --- drizzle/0010_minor_piledriver.sql | 2 + drizzle/meta/0010_snapshot.json | 728 ++++++++++++++++++ drizzle/meta/_journal.json | 7 + .../sources/[sourceId]/chunks/route.test.ts | 184 ++++- src/domains/chat/index.test.ts | 20 +- src/domains/chat/page-citation-assets.test.ts | 18 +- src/domains/chat/page-citation-assets.ts | 12 +- src/domains/chunks/server.test.ts | 123 +++ src/domains/chunks/server.ts | 176 ++++- .../sources/parse-result-storage-adapter.ts | 172 ++++- src/domains/sources/repository.ts | 3 + src/domains/sources/route-chunks.ts | 65 ++ src/domains/sources/route-dependencies.ts | 1 + src/domains/sources/route-types.ts | 9 + src/domains/sources/service.ts | 5 + .../sources/source-parse-result-repository.ts | 35 + .../source-reconcile-route-workflow.test.ts | 65 +- .../source-reconcile-route-workflow.ts | 40 +- src/domains/sources/workflow-runtime.ts | 18 + src/infrastructure/db/schema.ts | 9 +- 20 files changed, 1608 insertions(+), 84 deletions(-) create mode 100644 drizzle/0010_minor_piledriver.sql create mode 100644 drizzle/meta/0010_snapshot.json diff --git a/drizzle/0010_minor_piledriver.sql b/drizzle/0010_minor_piledriver.sql new file mode 100644 index 0000000..973029b --- /dev/null +++ b/drizzle/0010_minor_piledriver.sql @@ -0,0 +1,2 @@ +ALTER TABLE "source_parse_results" ADD COLUMN "snapshot_manifest_url" text;--> statement-breakpoint +ALTER TABLE "source_parse_results" ADD COLUMN "snapshot_manifest_key" text; \ No newline at end of file diff --git a/drizzle/meta/0010_snapshot.json b/drizzle/meta/0010_snapshot.json new file mode 100644 index 0000000..b9b5a4f --- /dev/null +++ b/drizzle/meta/0010_snapshot.json @@ -0,0 +1,728 @@ +{ + "id": "cc41c815-4e52-4d01-ac2e-abbc791331b7", + "prevId": "c59ee2f8-6b59-4a2c-98a3-d71a67388d47", + "version": "7", + "dialect": "postgresql", + "tables": { + "public.chat_messages": { + "name": "chat_messages", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "thread_id": { + "name": "thread_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "role": { + "name": "role", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "content": { + "name": "content", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "citations": { + "name": "citations", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "artifacts": { + "name": "artifacts", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "chat_messages_thread_created_idx": { + "name": "chat_messages_thread_created_idx", + "columns": [ + { + "expression": "thread_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "created_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "chat_messages_thread_id_chat_threads_id_fk": { + "name": "chat_messages_thread_id_chat_threads_id_fk", + "tableFrom": "chat_messages", + "tableTo": "chat_threads", + "columnsFrom": [ + "thread_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.chat_threads": { + "name": "chat_threads", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "workspace_id": { + "name": "workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "title": { + "name": "title", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "demo_key": { + "name": "demo_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + } + }, + "indexes": { + "chat_threads_workspace_updated_idx": { + "name": "chat_threads_workspace_updated_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "updated_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "where": "deleted_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + }, + "chat_threads_workspace_demo_key_idx": { + "name": "chat_threads_workspace_demo_key_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "demo_key", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "chat_threads_workspace_id_workspaces_id_fk": { + "name": "chat_threads_workspace_id_workspaces_id_fk", + "tableFrom": "chat_threads", + "tableTo": "workspaces", + "columnsFrom": [ + "workspace_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.demo_source_visibilities": { + "name": "demo_source_visibilities", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "workspace_id": { + "name": "workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "demo_source_id": { + "name": "demo_source_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "hidden_at": { + "name": "hidden_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "demo_source_visibilities_workspace_source_idx": { + "name": "demo_source_visibilities_workspace_source_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "demo_source_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "demo_source_visibilities_workspace_idx": { + "name": "demo_source_visibilities_workspace_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "demo_source_visibilities_workspace_id_workspaces_id_fk": { + "name": "demo_source_visibilities_workspace_id_workspaces_id_fk", + "tableFrom": "demo_source_visibilities", + "tableTo": "workspaces", + "columnsFrom": [ + "workspace_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.source_parse_results": { + "name": "source_parse_results", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "source_id": { + "name": "source_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "result_blob_url": { + "name": "result_blob_url", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "snapshot_manifest_url": { + "name": "snapshot_manifest_url", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "snapshot_manifest_key": { + "name": "snapshot_manifest_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "asset_urls": { + "name": "asset_urls", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "source_parse_results_source_id_idx": { + "name": "source_parse_results_source_id_idx", + "columns": [ + { + "expression": "source_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "source_parse_results_source_id_sources_id_fk": { + "name": "source_parse_results_source_id_sources_id_fk", + "tableFrom": "source_parse_results", + "tableTo": "sources", + "columnsFrom": [ + "source_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": { + "source_parse_results_source_id_unique": { + "name": "source_parse_results_source_id_unique", + "nullsNotDistinct": false, + "columns": [ + "source_id" + ] + } + }, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.sources": { + "name": "sources", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "workspace_id": { + "name": "workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "title": { + "name": "title", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "mime_type": { + "name": "mime_type", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "size_bytes": { + "name": "size_bytes", + "type": "bigint", + "primaryKey": false, + "notNull": true + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "failure_reason": { + "name": "failure_reason", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "knowhere_job_id": { + "name": "knowhere_job_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "knowhere_document_id": { + "name": "knowhere_document_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "staged_blob_pathname": { + "name": "staged_blob_pathname", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "staged_blob_url": { + "name": "staged_blob_url", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "original_blob_pathname": { + "name": "original_blob_pathname", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "original_blob_url": { + "name": "original_blob_url", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "demo_key": { + "name": "demo_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + } + }, + "indexes": { + "sources_workspace_created_idx": { + "name": "sources_workspace_created_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "created_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "where": "deleted_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + }, + "sources_workspace_status_idx": { + "name": "sources_workspace_status_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "status", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "sources_workspace_demo_key_idx": { + "name": "sources_workspace_demo_key_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "demo_key", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "sources_workspace_document_idx": { + "name": "sources_workspace_document_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "knowhere_document_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "where": "knowhere_document_id IS NOT NULL AND deleted_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "sources_workspace_id_workspaces_id_fk": { + "name": "sources_workspace_id_workspaces_id_fk", + "tableFrom": "sources", + "tableTo": "workspaces", + "columnsFrom": [ + "workspace_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.workspaces": { + "name": "workspaces", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "namespace": { + "name": "namespace", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "workspaces_user_id_idx": { + "name": "workspaces_user_id_idx", + "columns": [ + { + "expression": "user_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": { + "workspaces_user_id_unique": { + "name": "workspaces_user_id_unique", + "nullsNotDistinct": false, + "columns": [ + "user_id" + ] + }, + "workspaces_namespace_unique": { + "name": "workspaces_namespace_unique", + "nullsNotDistinct": false, + "columns": [ + "namespace" + ] + } + }, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + } + }, + "enums": {}, + "schemas": {}, + "sequences": {}, + "roles": {}, + "policies": {}, + "views": {}, + "_meta": { + "columns": {}, + "schemas": {}, + "tables": {} + } +} \ No newline at end of file diff --git a/drizzle/meta/_journal.json b/drizzle/meta/_journal.json index 568b77f..7174911 100644 --- a/drizzle/meta/_journal.json +++ b/drizzle/meta/_journal.json @@ -71,6 +71,13 @@ "when": 1782472980935, "tag": "0009_true_lila_cheney", "breakpoints": true + }, + { + "idx": 10, + "version": "7", + "when": 1783098230819, + "tag": "0010_minor_piledriver", + "breakpoints": true } ] } \ No newline at end of file diff --git a/src/app/api/sources/[sourceId]/chunks/route.test.ts b/src/app/api/sources/[sourceId]/chunks/route.test.ts index 524a657..c41319c 100644 --- a/src/app/api/sources/[sourceId]/chunks/route.test.ts +++ b/src/app/api/sources/[sourceId]/chunks/route.test.ts @@ -11,6 +11,7 @@ const mocks = vi.hoisted(() => ({ findSourceInWorkspace: vi.fn(), getCurrentUser: vi.fn(), getSourceParseAssetUrls: vi.fn(), + getSourceParseSnapshotMetadata: vi.fn(), localizeRemoteDocument: vi.fn(), makeKnowhereClient: vi.fn(), requireUser: vi.fn(), @@ -51,6 +52,7 @@ vi.mock("@/domains/sources/service", () => ({ sourceService: { findInWorkspace: mocks.findSourceInWorkspace, getParseAssetUrls: mocks.getSourceParseAssetUrls, + getParseSnapshotMetadata: mocks.getSourceParseSnapshotMetadata, localizeRemoteDocument: mocks.localizeRemoteDocument, updateSourceRevisionKey: mocks.updateSourceRevisionKey, }, @@ -71,6 +73,16 @@ describe("GET /api/sources/[sourceId]/chunks", () => { mocks.blobPut.mockImplementation(async (pathname: string) => ({ url: `https://blob.example/${pathname}`, })) + mocks.getSourceParseAssetUrls.mockResolvedValue({}) + mocks.getSourceParseSnapshotMetadata.mockResolvedValue({ + resultBlobUrl: + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + snapshotManifestUrl: + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + assetUrlsByFilePath: {}, + }) mocks.updateSourceRevisionKey.mockResolvedValue(null) }) @@ -429,30 +441,10 @@ describe("GET /api/sources/[sourceId]/chunks", () => { } }) - it("loads authenticated workspace chunks without probing the demo endpoint first", async () => { + it("loads authenticated workspace chunks from the parsed Blob snapshot", async () => { const knowhereClient = { documents: { - listChunks: vi.fn(async () => ({ - chunks: [ - { - id: "dchk_1", - chunkId: "parser_1", - chunkType: "text", - content: "Workspace chunk", - sectionPath: "Summary", - sourceChunkPath: "Default_Root/notes.pdf/Summary", - filePath: null, - metadata: {}, - sortOrder: 0, - }, - ], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - })), + listChunks: vi.fn(), }, } mocks.getCurrentUser.mockResolvedValue({ @@ -487,7 +479,81 @@ describe("GET /api/sources/[sourceId]/chunks", () => { }) mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") mocks.makeKnowhereClient.mockReturnValue(knowhereClient) - mocks.getSourceParseAssetUrls.mockResolvedValue({}) + mocks.getSourceParseAssetUrls.mockResolvedValue({ + "images/chart.png": "https://blob.example/images/chart.png", + }) + mocks.getSourceParseSnapshotMetadata.mockResolvedValue({ + resultBlobUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000002/parsed-result/manifest/current.json", + snapshotManifestUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000002/parsed-result/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000002/parsed-result/manifest/current.json", + assetUrlsByFilePath: { + "images/chart.png": "https://blob.example/images/chart.png", + }, + }) + mocks.blobGet.mockImplementation(async (pathname: string) => { + if (pathname.endsWith("/manifest/current.json")) { + return { + statusCode: 200, + stream: createTextStream( + JSON.stringify({ + version: 1, + kind: "knowhere-parsed-result-snapshot", + jobId: "job_1", + documentId: "doc_1", + sourceFileName: "notes.pdf", + totalChunks: 1, + chunkPageSize: 1, + chunkPages: [ + { + page: 1, + pageSize: 1, + chunkCount: 1, + key: "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000002/parsed-result/chunks/page-1.json", + }, + ], + assetUrlsByFilePath: { + "images/chart.png": "https://blob.example/images/chart.png", + }, + createdAt: "2026-07-03T00:00:00.000Z", + }), + ), + } + } + if (pathname.endsWith("/chunks/page-1.json")) { + return { + statusCode: 200, + stream: createTextStream( + JSON.stringify({ + version: 1, + jobId: "job_1", + documentId: "doc_1", + sourceFileName: "notes.pdf", + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + chunks: [ + { + id: "dchk_1", + chunkId: "parser_1", + chunkType: "image", + content: "Workspace chunk", + sectionPath: "Summary", + sourceChunkPath: "Default_Root/notes.pdf/Summary", + filePath: "images/chart.png", + metadata: {}, + sortOrder: 1, + }, + ], + }), + ), + } + } + return null + }) const response = await GET( new NextRequest( @@ -502,6 +568,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { chunkId: "dchk_1", parserChunkId: "parser_1", documentId: "doc_1", + assetUrl: "https://blob.example/images/chart.png", sourceTitle: "notes.pdf", }, ], @@ -513,11 +580,68 @@ describe("GET /api/sources/[sourceId]/chunks", () => { }) expect(response.status).toBe(200) expect(mocks.fetchDemoChunkPage).not.toHaveBeenCalled() - expect(knowhereClient.documents.listChunks).toHaveBeenCalledWith("doc_1", { - page: 1, - pageSize: 1, - includeAssetUrls: true, + expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled() + }) + + it("returns processing when a ready workspace source has no complete parsed snapshot", async () => { + const knowhereClient = { + documents: { + listChunks: vi.fn(), + }, + } + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue({ + id: "00000000-0000-0000-0000-000000000002", + workspaceId: "workspace_1", + title: "notes.pdf", + mimeType: "application/pdf", + sizeBytes: 1024, + status: "ready", + failureReason: null, + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: null, + originalBlobUrl: null, + demoKey: null, + createdAt: new Date("2026-05-10T00:00:00.000Z"), + updatedAt: new Date("2026-05-10T00:00:00.000Z"), + deletedAt: null, + }) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.makeKnowhereClient.mockReturnValue(knowhereClient) + mocks.getSourceParseSnapshotMetadata.mockResolvedValue(null) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/chunks?page=1&pageSize=1", + ), + { params: Promise.resolve({ sourceId: "00000000-0000-0000-0000-000000000002" }) }, + ) + + await expect(response.json()).resolves.toEqual({ + chunks: [], + pagination: { + page: 1, + pageSize: 1, + total: 0, + totalPages: 0, + }, + message: "Source parsed snapshot is still being prepared.", }) + expect(response.status).toBe(202) + expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled() }) it("materializes a remote source id on open before loading chunks", async () => { @@ -650,3 +774,9 @@ describe("GET /api/sources/[sourceId]/chunks", () => { ) }) }) + +function createTextStream(text: string): ReadableStream { + const stream = new Response(text).body + if (!stream) throw new Error("Response body stream was not created.") + return stream +} diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 459add9..71604f6 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -671,6 +671,8 @@ describe("answerQuestionWithRetrieval", () => { it("hardens page citation asset URLs before returning citations", async () => { const rawPageAssetUrl = "https://knowhere-storage.example/results/job_1/page_citation_assets/page-4.png?AWSAccessKeyId=test"; + const storedPageAssetUrl = + "https://blob.example/workspaces/workspace_1/sources/source_pages/parsed-result/page_citation_assets/page-4.png"; const hardenedPageAssetUrl = "https://blob.example/workspaces/workspace_1/chat-assets/source-source_pages/page-4.png"; const result = makeRetrievalResult({ @@ -705,7 +707,7 @@ describe("answerQuestionWithRetrieval", () => { }; const generateAnswer = vi.fn(async ({ searchSources }) => { await searchSources({ query: "page four evidence" }); - return makeHarnessRunResult(`This page has the answer. ${rawPageAssetUrl}`); + return makeHarnessRunResult(`This page has the answer. ${storedPageAssetUrl}`); }); const hardenMediaAssetUrls = vi.fn( async ({ @@ -720,6 +722,8 @@ describe("answerQuestionWithRetrieval", () => { ...candidate, pageCitationAssetUrl: candidate.pageCitationAssetUrl === rawPageAssetUrl + ? hardenedPageAssetUrl + : candidate.pageCitationAssetUrl === storedPageAssetUrl ? hardenedPageAssetUrl : candidate.pageCitationAssetUrl, }), @@ -743,6 +747,9 @@ describe("answerQuestionWithRetrieval", () => { retrieval, generateAnswer, hardenMediaAssetUrls, + loadSourceAssetUrls: vi.fn(async () => ({ + "page_citation_assets/page-4.png": storedPageAssetUrl, + })), messages: [], }), ); @@ -750,7 +757,7 @@ describe("answerQuestionWithRetrieval", () => { expect(hardenMediaAssetUrls).toHaveBeenCalledWith({ results: [ expect.objectContaining({ - pageCitationAssetUrl: rawPageAssetUrl, + pageCitationAssetUrl: storedPageAssetUrl, source: expect.objectContaining({ sourceFileName: "deck.pdf", }), @@ -775,6 +782,8 @@ describe("answerQuestionWithRetrieval", () => { it("hardens page citation asset URLs from referenced chunk metadata", async () => { const rawPageAssetUrl = "https://knowhere-storage.example/results/job_1/page_citation_assets/page-6.png?AWSAccessKeyId=test"; + const storedPageAssetUrl = + "https://blob.example/workspaces/workspace_1/sources/source_pages/parsed-result/page_citation_assets/page-6.png"; const hardenedPageAssetUrl = "https://blob.example/workspaces/workspace_1/chat-assets/source-source_pages/page-6.png"; const retrieval = { @@ -826,6 +835,8 @@ describe("answerQuestionWithRetrieval", () => { ...candidate, pageCitationAssetUrl: candidate.pageCitationAssetUrl === rawPageAssetUrl + ? hardenedPageAssetUrl + : candidate.pageCitationAssetUrl === storedPageAssetUrl ? hardenedPageAssetUrl : candidate.pageCitationAssetUrl, }), @@ -849,6 +860,9 @@ describe("answerQuestionWithRetrieval", () => { retrieval, generateAnswer, hardenMediaAssetUrls, + loadSourceAssetUrls: vi.fn(async () => ({ + "page_citation_assets/page-6.png": storedPageAssetUrl, + })), messages: [], }), ); @@ -863,7 +877,7 @@ describe("answerQuestionWithRetrieval", () => { }), ], }), - pageCitationAssetUrl: rawPageAssetUrl, + pageCitationAssetUrl: storedPageAssetUrl, }), ], artifacts: undefined, diff --git a/src/domains/chat/page-citation-assets.test.ts b/src/domains/chat/page-citation-assets.test.ts index 6d0ee01..ed62c66 100644 --- a/src/domains/chat/page-citation-assets.test.ts +++ b/src/domains/chat/page-citation-assets.test.ts @@ -5,7 +5,7 @@ import type { Source } from "@/infrastructure/db/schema" import { enrichRetrievalResultsWithPageCitationAssetUrls } from "./page-citation-assets" describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { - it("uses server-provided page asset URLs from result metadata", async () => { + it("does not use direct server-provided page asset URLs from result metadata", async () => { const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ results: [ makeRetrievalResult({ @@ -25,12 +25,10 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { sources: [makeSource()], }) - expect(result?.pageCitationAssetUrl).toBe( - "https://assets.example/pages/page-2.png", - ) + expect(result?.pageCitationAssetUrl).toBeUndefined() }) - it("prefers stored Blob URLs over server-provided page asset URLs", async () => { + it("uses stored Blob URLs for page citation assets", async () => { const loadSourceAssetUrls = vi.fn().mockResolvedValue({ "page_citation_assets/page-2.png": "https://blob.example/page_citation_assets/page-2.png", @@ -64,7 +62,12 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { ) }) - it("chooses the asset matching the citation page metadata", async () => { + it("chooses the stored asset matching the citation page metadata", async () => { + const loadSourceAssetUrls = vi.fn().mockResolvedValue({ + "page_citation_assets/page-4.png": + "https://blob.example/page_citation_assets/page-4.png", + }) + const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ results: [ makeRetrievalResult({ @@ -87,10 +90,11 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { }), ], sources: [makeSource()], + loadSourceAssetUrls, }) expect(result?.pageCitationAssetUrl).toBe( - "https://assets.example/pages/page-4.png", + "https://blob.example/page_citation_assets/page-4.png", ) }) diff --git a/src/domains/chat/page-citation-assets.ts b/src/domains/chat/page-citation-assets.ts index dd1d34d..77b8d53 100644 --- a/src/domains/chat/page-citation-assets.ts +++ b/src/domains/chat/page-citation-assets.ts @@ -78,12 +78,6 @@ async function enrichRetrievalResultWithPageCitationAssetUrl(input: { pageCitationAssetUrl: sourceAssetUrl, } } - if (directAsset?.assetUrl) { - return { - ...input.result, - pageCitationAssetUrl: directAsset.assetUrl, - } - } return input.result } @@ -127,14 +121,10 @@ function getDirectPageCitationAsset( const matchingCandidates = candidates.filter((candidate) => pageNumbers.includes(candidate.pageNum), ) - const matchingAssetWithUrl = matchingCandidates.find( - (candidate) => candidate.assetUrl, - ) - if (matchingAssetWithUrl) return matchingAssetWithUrl if (matchingCandidates[0]) return matchingCandidates[0] } - return candidates.find((candidate) => candidate.assetUrl) ?? candidates[0] ?? null + return candidates[0] ?? null } function parsePageCitationAssetCandidates( diff --git a/src/domains/chunks/server.test.ts b/src/domains/chunks/server.test.ts index f58f5ea..2442b53 100644 --- a/src/domains/chunks/server.test.ts +++ b/src/domains/chunks/server.test.ts @@ -9,6 +9,107 @@ import { } from "./server" describe("server chunk cache", () => { + it("loads ready source chunks from the parsed Blob snapshot without calling Knowhere listChunks", async () => { + const listChunks = vi.fn(async () => ({ + documentId: "doc_1", + chunks: [], + pagination: { + page: 1, + pageSize: 1, + total: 0, + totalPages: 0, + }, + })) + const snapshotStore = createSnapshotStore({ + "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json": + JSON.stringify({ + version: 1, + kind: "knowhere-parsed-result-snapshot", + jobId: "job_1", + documentId: "doc_1", + sourceFileName: "notes.pdf", + totalChunks: 1, + chunkPageSize: 1, + chunkPages: [ + { + page: 1, + pageSize: 1, + chunkCount: 1, + key: "workspaces/workspace_1/sources/source_1/parsed-result/chunks/page-1.json", + }, + ], + assetUrlsByFilePath: { + "images/chart.png": "https://blob.example/images/chart.png", + }, + createdAt: "2026-07-03T00:00:00.000Z", + }), + "workspaces/workspace_1/sources/source_1/parsed-result/chunks/page-1.json": + JSON.stringify({ + version: 1, + jobId: "job_1", + documentId: "doc_1", + sourceFileName: "notes.pdf", + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + chunks: [ + { + id: "chunk_image_1", + chunkId: "parser_image_1", + chunkType: "image", + content: "Chart summary", + sectionPath: "Revenue", + sourceChunkPath: "notes.pdf/Revenue", + filePath: "images/chart.png", + sortOrder: 1, + metadata: { summary: "Chart summary" }, + }, + ], + }), + }) + + const page = await Effect.runPromise( + loadChunkPageForSource( + makeSource({ status: "ready", knowhereJobId: "job_1" }), + { documents: { listChunks } }, + { page: 1, pageSize: 1 }, + { + snapshot: { + resultBlobUrl: + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + snapshotManifestUrl: + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + assetUrlsByFilePath: {}, + }, + snapshotStore, + workspaceId: "workspace_1", + }, + ), + ) + + expect(listChunks).not.toHaveBeenCalled() + expect(page).toEqual({ + chunks: [ + expect.objectContaining({ + chunkId: "chunk_image_1", + parserChunkId: "parser_image_1", + type: "image", + assetUrl: "https://blob.example/images/chart.png", + sourceTitle: "notes.pdf", + }), + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + }) + }) + it("returns upstream chunks on a visible cache miss and warms mirrored assets in the background", async () => { const warmTasks: Array<() => Promise> = [] const cacheStore = createCacheStore() @@ -355,6 +456,13 @@ type TestCacheStore = { readonly putMock: TestCachePutMock } +type TestSnapshotStore = { + readonly get: ( + pathname: string, + options: { readonly access: "public" }, + ) => Promise +} + type TestCacheGetMock = Mock< ( pathname: string, @@ -391,6 +499,21 @@ function createCacheStore(overrides: Partial<{ } } +function createSnapshotStore( + entriesByPathname: Readonly>, +): TestSnapshotStore { + return { + get: vi.fn(async (pathname: string) => { + const entry = entriesByPathname[pathname] + if (!entry) return null + return { + statusCode: 200 as const, + stream: createTextStream(entry), + } + }), + } +} + function makeDocumentChunk( overrides: Partial = {}, ): DocumentChunk { diff --git a/src/domains/chunks/server.ts b/src/domains/chunks/server.ts index a28ee3b..64b7fe4 100644 --- a/src/domains/chunks/server.ts +++ b/src/domains/chunks/server.ts @@ -7,7 +7,6 @@ import { after } from "next/server" import { Effect } from "effect" import type { DocumentChunk, - DocumentChunkListResponse, } from "@ontos-ai/knowhere-sdk" import { @@ -18,6 +17,13 @@ import { type ChunkPageParams, type LoadChunksOptions, } from "@/domains/chunks" +import { + readParsedResultSnapshotChunkPage, + readParsedResultSnapshotManifest, + type ParsedResultSnapshotBlobStore, + type ParsedResultSnapshotChunk, + type ParsedResultSnapshotManifest, +} from "@/domains/sources/parse-result-storage-adapter" import type { ParsedChunkView } from "@/domains/chunks/types" import type { Source } from "@/infrastructure/db/schema" import { logger } from "@/lib/logger" @@ -59,6 +65,8 @@ type ChunkPageWarmScheduler = (task: () => Promise) => void type ServerLoadChunksOptions = LoadChunksOptions & { readonly workspaceId?: string + readonly snapshot?: SourceParsedSnapshotReadModel | null + readonly snapshotStore?: ParsedResultSnapshotBlobStore readonly cacheStore?: ChunkPageBlobStore readonly fetchAsset?: FetchChunkAsset readonly mode?: ChunkPageMode @@ -86,6 +94,13 @@ type MirrorableChunkAsset = { readonly filePath?: string } +export type SourceParsedSnapshotReadModel = { + readonly resultBlobUrl: string + readonly snapshotManifestUrl?: string | null + readonly snapshotManifestKey?: string | null + readonly assetUrlsByFilePath: Readonly> +} + const documentChunkPageSize = 200 const visibleChunkPageMode: ChunkPageMode = "visible" const structureChunkPageMode: ChunkPageMode = "structure" @@ -124,6 +139,17 @@ export const loadChunksForSource = ( Effect.gen(function* () { if (source.status !== "ready" || !source.knowhereDocumentId) return [] + if (options.snapshot) { + const snapshotChunks = yield* Effect.promise(() => + loadAllSnapshotChunks({ + source, + snapshot: options.snapshot!, + blobStore: options.snapshotStore, + }), + ) + return resolveChunkConnectionTargets(snapshotChunks) + } + const chunks: ParsedChunkView[] = [] let page = 1 let totalPages = 1 @@ -156,6 +182,17 @@ export const loadChunkPageForSource = ( return emptyPage } + if (options.snapshot) { + return yield* Effect.promise(() => + loadSnapshotChunkPage({ + source, + params, + snapshot: options.snapshot!, + blobStore: options.snapshotStore, + }), + ) + } + const mode = options.mode ?? visibleChunkPageMode const workspaceId = options.workspaceId ?? source.workspaceId const cacheStore = options.cacheStore ?? defaultBlobStore @@ -236,6 +273,138 @@ export const loadChunkPageForSource = ( return chunkPage }) +async function loadAllSnapshotChunks(input: { + readonly source: Source + readonly snapshot: SourceParsedSnapshotReadModel + readonly blobStore?: ParsedResultSnapshotBlobStore +}): Promise { + const manifest = await loadSnapshotManifest(input) + if (!manifest) return [] + + const pages = await Promise.all( + manifest.chunkPages.map((page) => + readParsedResultSnapshotChunkPage({ + pageKey: page.key, + blobStore: input.blobStore, + }), + ), + ) + const chunks = pages.flatMap((page) => page?.chunks ?? []) + return chunks.map((chunk) => + toParsedChunkView( + toDocumentChunkFromSnapshot(chunk), + input.source.title, + input.source.knowhereDocumentId ?? undefined, + { + assetUrlsByFilePath: { + ...manifest.assetUrlsByFilePath, + ...input.snapshot.assetUrlsByFilePath, + }, + }, + ), + ) +} + +async function loadSnapshotChunkPage(input: { + readonly source: Source + readonly params: ChunkPageParams + readonly snapshot: SourceParsedSnapshotReadModel + readonly blobStore?: ParsedResultSnapshotBlobStore +}): Promise { + const manifest = await loadSnapshotManifest(input) + if (!manifest) return createEmptyChunkPage(input.params) + + const pageReference = manifest.chunkPages.find( + (page) => page.page === input.params.page && page.pageSize === input.params.pageSize, + ) + if (!pageReference) { + return createEmptyChunkPage({ + page: input.params.page, + pageSize: input.params.pageSize, + }) + } + + const page = await readParsedResultSnapshotChunkPage({ + pageKey: pageReference.key, + blobStore: input.blobStore, + }) + if (!page) return createEmptyChunkPage(input.params) + + const assetUrlsByFilePath = { + ...manifest.assetUrlsByFilePath, + ...input.snapshot.assetUrlsByFilePath, + } + const chunks = page.chunks.map((chunk) => + toParsedChunkView( + toDocumentChunkFromSnapshot(chunk), + input.source.title, + input.source.knowhereDocumentId ?? undefined, + { assetUrlsByFilePath }, + ), + ) + + return { + chunks, + pagination: { + page: page.page, + pageSize: page.pageSize, + total: page.total, + totalPages: page.totalPages, + }, + } +} + +async function loadSnapshotManifest(input: { + readonly source: Source + readonly snapshot: SourceParsedSnapshotReadModel + readonly blobStore?: ParsedResultSnapshotBlobStore +}): Promise { + const manifestKey = input.snapshot.snapshotManifestKey + if (!manifestKey) { + logger.warn("chunks: ready source missing parsed snapshot manifest", { + sourceId: input.source.id, + documentId: input.source.knowhereDocumentId, + }) + return null + } + + const manifest = await readParsedResultSnapshotManifest({ + workspaceId: input.source.workspaceId, + sourceId: input.source.id, + manifestKey, + blobStore: input.blobStore, + }) + if (!manifest) { + logger.warn("chunks: parsed snapshot manifest could not be read", { + sourceId: input.source.id, + documentId: input.source.knowhereDocumentId, + manifestKey, + }) + } + return manifest +} + +function toDocumentChunkFromSnapshot(chunk: ParsedResultSnapshotChunk): DocumentChunk { + return { + id: chunk.id, + chunkId: chunk.chunkId, + chunkType: normalizeDocumentChunkType(chunk.chunkType), + content: chunk.content, + sectionId: null, + sectionPath: chunk.sectionPath ?? null, + sourceChunkPath: chunk.sourceChunkPath, + filePath: chunk.filePath ?? null, + sortOrder: chunk.sortOrder, + metadata: chunk.metadata, + assetUrl: chunk.assetUrl ?? null, + } +} + +function normalizeDocumentChunkType(value: string): DocumentChunk["chunkType"] { + if (value === "image" || value === "table" || value === "page") return value + return "text" +} + export async function warmChunkPageCache( input: WarmChunkPageCacheInput, ): Promise { @@ -653,7 +822,10 @@ function getMirroredAssetContentType( } function getRevisionKey( - response: Pick, + response: { + readonly jobId?: string | null + readonly jobResultId?: string | null + }, source: Source, ): string | null { return ( diff --git a/src/domains/sources/parse-result-storage-adapter.ts b/src/domains/sources/parse-result-storage-adapter.ts index 908dd34..d195b8f 100644 --- a/src/domains/sources/parse-result-storage-adapter.ts +++ b/src/domains/sources/parse-result-storage-adapter.ts @@ -1,6 +1,6 @@ import "server-only" -import { head, put, BlobNotFoundError } from "@vercel/blob" +import { get, head, put, BlobNotFoundError } from "@vercel/blob" import type { KnowhereAssetStorageObject, KnowhereAssetStorageOptions, @@ -18,8 +18,56 @@ export type ParsedResultAssetIndex = { readonly updatedAt: string } +export type ParsedResultSnapshotChunk = { + readonly id: string + readonly chunkId: string + readonly chunkType: string + readonly contentSource?: string + readonly content: string + readonly sectionPath?: string + readonly sourceChunkPath: string + readonly filePath?: string + readonly sortOrder: number + readonly metadata: Record + readonly assetUrl?: string +} + +export type ParsedResultSnapshotChunkPage = { + readonly version: 1 + readonly jobId: string + readonly documentId?: string + readonly namespace?: string + readonly sourceFileName: string + readonly page: number + readonly pageSize: number + readonly total: number + readonly totalPages: number + readonly chunks: readonly ParsedResultSnapshotChunk[] +} + +export type ParsedResultSnapshotManifest = { + readonly version: 1 + readonly kind: "knowhere-parsed-result-snapshot" + readonly jobId: string + readonly documentId?: string + readonly namespace?: string + readonly sourceFileName: string + readonly totalChunks: number + readonly chunkPageSize: number + readonly chunkPages: readonly { + readonly page: number + readonly pageSize: number + readonly chunkCount: number + readonly key: string + readonly url?: string + }[] + readonly assetUrlsByFilePath: Readonly> + readonly createdAt: string +} + const parsedResultDirectoryName = "parsed-result" const parsedResultAssetIndexFileName = "asset-index.json" +const parsedResultSnapshotManifestPath = "manifest/current.json" export function createParsedResultStorageAdapter({ workspaceId, @@ -69,6 +117,42 @@ export function createParsedResultStorageAdapter({ } } +export function getParsedResultSnapshotManifestKey({ + workspaceId, + sourceId, +}: ParsedResultAssetStorageInput): string { + return `${getParsedResultBlobPrefix( + workspaceId, + sourceId, + )}/${parsedResultSnapshotManifestPath}` +} + +export async function readParsedResultSnapshotManifest(input: { + readonly workspaceId: string + readonly sourceId: string + readonly manifestKey?: string | null + readonly blobStore?: ParsedResultSnapshotBlobStore +}): Promise { + const key = + input.manifestKey ?? + getParsedResultSnapshotManifestKey({ + workspaceId: input.workspaceId, + sourceId: input.sourceId, + }) + const text = await readBlobText(key, input.blobStore ?? vercelSnapshotBlobStore) + if (!text) return null + return parseSnapshotManifest(text) +} + +export async function readParsedResultSnapshotChunkPage(input: { + readonly pageKey: string + readonly blobStore?: ParsedResultSnapshotBlobStore +}): Promise { + const text = await readBlobText(input.pageKey, input.blobStore ?? vercelSnapshotBlobStore) + if (!text) return null + return parseSnapshotChunkPage(text) +} + export async function writeParsedResultAssetIndex(input: { readonly workspaceId: string readonly sourceId: string @@ -97,3 +181,89 @@ export async function writeParsedResultAssetIndex(input: { function getParsedResultBlobPrefix(workspaceId: string, sourceId: string): string { return `workspaces/${workspaceId}/sources/${sourceId}/${parsedResultDirectoryName}` } + +type ParsedResultSnapshotBlobGetResult = + | { + readonly statusCode: 200 + readonly stream: ReadableStream + } + | { + readonly statusCode: 304 + readonly stream: null + } + +export type ParsedResultSnapshotBlobStore = { + readonly get: ( + pathname: string, + options: { readonly access: "public" }, + ) => Promise +} + +const vercelSnapshotBlobStore: ParsedResultSnapshotBlobStore = { + get: (pathname, options) => get(pathname, options), +} + +async function readBlobText( + key: string, + blobStore: ParsedResultSnapshotBlobStore, +): Promise { + try { + const result = await blobStore.get(key, { access: "public" }) + if (!result || result.statusCode !== 200) return null + return new Response(result.stream).text() + } catch (error) { + if (error instanceof BlobNotFoundError) return null + throw error + } +} + +function parseSnapshotManifest(text: string): ParsedResultSnapshotManifest | null { + try { + const value: unknown = JSON.parse(text) + return isSnapshotManifest(value) ? value : null + } catch { + return null + } +} + +function parseSnapshotChunkPage(text: string): ParsedResultSnapshotChunkPage | null { + try { + const value: unknown = JSON.parse(text) + return isSnapshotChunkPage(value) ? value : null + } catch { + return null + } +} + +function isSnapshotManifest(value: unknown): value is ParsedResultSnapshotManifest { + if (!isRecord(value)) return false + return ( + value["kind"] === "knowhere-parsed-result-snapshot" && + value["version"] === 1 && + typeof value["jobId"] === "string" && + typeof value["sourceFileName"] === "string" && + typeof value["totalChunks"] === "number" && + typeof value["chunkPageSize"] === "number" && + Array.isArray(value["chunkPages"]) && + isRecord(value["assetUrlsByFilePath"]) && + typeof value["createdAt"] === "string" + ) +} + +function isSnapshotChunkPage(value: unknown): value is ParsedResultSnapshotChunkPage { + if (!isRecord(value)) return false + return ( + value["version"] === 1 && + typeof value["jobId"] === "string" && + typeof value["sourceFileName"] === "string" && + typeof value["page"] === "number" && + typeof value["pageSize"] === "number" && + typeof value["total"] === "number" && + typeof value["totalPages"] === "number" && + Array.isArray(value["chunks"]) + ) +} + +function isRecord(value: unknown): value is Readonly> { + return typeof value === "object" && value !== null && !Array.isArray(value) +} diff --git a/src/domains/sources/repository.ts b/src/domains/sources/repository.ts index c7c4b6f..4d9c061 100644 --- a/src/domains/sources/repository.ts +++ b/src/domains/sources/repository.ts @@ -21,6 +21,7 @@ type SourceRepository = { readonly saveParseResultEffect: typeof sourceParseResultRepository.saveParseResultEffect readonly mergeParseAssetUrlsEffect: typeof sourceParseResultRepository.mergeParseAssetUrlsEffect readonly getParseResultProgressEffect: typeof sourceParseResultRepository.getParseResultProgressEffect + readonly getParseSnapshotMetadataEffect: typeof sourceParseResultRepository.getParseSnapshotMetadataEffect readonly getParseAssetUrlsEffect: typeof sourceParseResultRepository.getParseAssetUrlsEffect } @@ -45,5 +46,7 @@ export const sourceRepository: SourceRepository = { sourceParseResultRepository.mergeParseAssetUrlsEffect, getParseResultProgressEffect: sourceParseResultRepository.getParseResultProgressEffect, + getParseSnapshotMetadataEffect: + sourceParseResultRepository.getParseSnapshotMetadataEffect, getParseAssetUrlsEffect: sourceParseResultRepository.getParseAssetUrlsEffect, } diff --git a/src/domains/sources/route-chunks.ts b/src/domains/sources/route-chunks.ts index 733ccc1..fdbe0b1 100644 --- a/src/domains/sources/route-chunks.ts +++ b/src/domains/sources/route-chunks.ts @@ -88,8 +88,16 @@ const loadSourceChunksEffect = ( const client = yield* Effect.tryPromise(() => getClientForWorkspace(workspace.id, input.cookieHeader, deps), ) + const snapshot = yield* Effect.tryPromise(() => + deps.sourceService.getParseSnapshotMetadata(workspace.id, source.id), + ) + if (!isCompleteSnapshot(snapshot)) { + return sourceSnapshotProcessing(input) + } + if (input.shouldLoadAll) { const chunks = yield* deps.loadChunksForSource(source, client, { + snapshot, workspaceId: workspace.id, onRevisionKey: async (revisionKey) => { await deps.sourceService.updateSourceRevisionKey( @@ -111,6 +119,7 @@ const loadSourceChunksEffect = ( input.pageParams, { assetUrlsByFilePath, + snapshot, workspaceId: workspace.id, onRevisionKey: async (revisionKey) => { await deps.sourceService.updateSourceRevisionKey( @@ -294,4 +303,60 @@ function sourceNotFound(): JsonRouteResult<{ readonly message: string }> { return routeResult.error(404, "Source not found.") } +function sourceSnapshotProcessing( + input: LoadSourceChunksInput, +): JsonRouteResult<{ + readonly chunks: [] + readonly pagination?: { + readonly page: number + readonly pageSize: number + readonly total: 0 + readonly totalPages: 0 + } + readonly message: string +}> { + if (input.shouldLoadAll) { + return routeResult.ok( + { + chunks: [], + message: "Source parsed snapshot is still being prepared.", + }, + 202, + ) + } + + return routeResult.ok( + { + chunks: [], + pagination: { + page: input.pageParams.page, + pageSize: input.pageParams.pageSize, + total: 0, + totalPages: 0, + }, + message: "Source parsed snapshot is still being prepared.", + }, + 202, + ) +} + +function isCompleteSnapshot( + snapshot: + | { + readonly snapshotManifestKey?: string | null + readonly snapshotManifestUrl?: string | null + } + | null, +): snapshot is { + readonly snapshotManifestKey: string + readonly snapshotManifestUrl: string +} { + return ( + typeof snapshot?.snapshotManifestKey === "string" && + snapshot.snapshotManifestKey.length > 0 && + typeof snapshot.snapshotManifestUrl === "string" && + snapshot.snapshotManifestUrl.length > 0 + ) +} + export { createRouteChunks } diff --git a/src/domains/sources/route-dependencies.ts b/src/domains/sources/route-dependencies.ts index 25313ce..d48ede2 100644 --- a/src/domains/sources/route-dependencies.ts +++ b/src/domains/sources/route-dependencies.ts @@ -46,6 +46,7 @@ const defaultDependencies: SourceRouteServiceDependencies = { sourceService: { findInWorkspace: defaultSourceService.findInWorkspace, getParseAssetUrls: defaultSourceService.getParseAssetUrls, + getParseSnapshotMetadata: defaultSourceService.getParseSnapshotMetadata, hideDemoSource: defaultSourceService.hideDemoSource, listHiddenDemoSourceIds: defaultSourceService.listHiddenDemoSourceIds, localizeRemoteDocument: defaultSourceService.localizeRemoteDocument, diff --git a/src/domains/sources/route-types.ts b/src/domains/sources/route-types.ts index 7303ec3..79d0635 100644 --- a/src/domains/sources/route-types.ts +++ b/src/domains/sources/route-types.ts @@ -183,6 +183,15 @@ type SourceWorkflowService = { workspaceId: string, sourceId: string, ) => Promise>> + readonly getParseSnapshotMetadata: ( + workspaceId: string, + sourceId: string, + ) => Promise<{ + readonly resultBlobUrl: string + readonly snapshotManifestUrl?: string | null + readonly snapshotManifestKey?: string | null + readonly assetUrlsByFilePath: Readonly> + } | null> readonly hideDemoSource: ( workspaceId: string, demoSourceId: string, diff --git a/src/domains/sources/service.ts b/src/domains/sources/service.ts index 0d7c7f6..57701e6 100644 --- a/src/domains/sources/service.ts +++ b/src/domains/sources/service.ts @@ -21,6 +21,10 @@ type SourceService = { workspaceId: string, sourceId: string, ) => Promise>> + readonly getParseSnapshotMetadata: ( + workspaceId: string, + sourceId: string, + ) => ReturnType readonly listForWorkspace: (workspaceId: string) => Promise readonly localizeRemoteDocument: ( workspaceId: string, @@ -106,6 +110,7 @@ const retrySourceToKnowhere: SourceService["retrySourceToKnowhere"] = ( export const sourceService: SourceService = { findInWorkspace: sourceWorkflowRuntime.findInWorkspace, getParseAssetUrls: sourceWorkflowRuntime.getParseAssetUrls, + getParseSnapshotMetadata: sourceWorkflowRuntime.getParseSnapshotMetadata, hideDemoSource: sourceWorkflowRuntime.hideDemoSource, listHiddenDemoSourceIds: sourceWorkflowRuntime.listHiddenDemoSourceIds, listForWorkspace: sourceWorkflowRuntime.listForWorkspace, diff --git a/src/domains/sources/source-parse-result-repository.ts b/src/domains/sources/source-parse-result-repository.ts index 8a5711e..c3e9859 100644 --- a/src/domains/sources/source-parse-result-repository.ts +++ b/src/domains/sources/source-parse-result-repository.ts @@ -12,11 +12,22 @@ import { sourceRowRepository } from "./source-row-repository" type SaveSourceParseResultInput = { readonly resultBlobUrl: string + readonly snapshotManifestUrl?: string + readonly snapshotManifestKey?: string readonly assetUrlsByFilePath: Readonly> } type SourceParseResultProgress = { readonly resultBlobUrl: string + readonly snapshotManifestUrl?: string | null + readonly snapshotManifestKey?: string | null + readonly assetUrlsByFilePath: Readonly> +} + +export type SourceParseSnapshotMetadata = { + readonly resultBlobUrl: string + readonly snapshotManifestUrl?: string | null + readonly snapshotManifestKey?: string | null readonly assetUrlsByFilePath: Readonly> } @@ -35,6 +46,10 @@ type SourceParseResultRepository = { workspaceId: string, sourceId: string, ) => Effect.Effect + readonly getParseSnapshotMetadataEffect: ( + workspaceId: string, + sourceId: string, + ) => Effect.Effect readonly getParseAssetUrlsEffect: ( workspaceId: string, sourceId: string, @@ -62,12 +77,16 @@ const saveParseResultEffect: SourceParseResultRepository["saveParseResultEffect" .values({ sourceId, resultBlobUrl: input.resultBlobUrl, + snapshotManifestUrl: input.snapshotManifestUrl, + snapshotManifestKey: input.snapshotManifestKey, assetUrls: input.assetUrlsByFilePath, }) .onConflictDoUpdate({ target: sourceParseResults.sourceId, set: { resultBlobUrl: input.resultBlobUrl, + snapshotManifestUrl: input.snapshotManifestUrl, + snapshotManifestKey: input.snapshotManifestKey, assetUrls: input.assetUrlsByFilePath, updatedAt: sql`now()`, }, @@ -93,12 +112,16 @@ const mergeParseAssetUrlsEffect: SourceParseResultRepository["mergeParseAssetUrl .values({ sourceId, resultBlobUrl: input.resultBlobUrl, + snapshotManifestUrl: input.snapshotManifestUrl, + snapshotManifestKey: input.snapshotManifestKey, assetUrls: input.assetUrlsByFilePath, }) .onConflictDoUpdate({ target: sourceParseResults.sourceId, set: { resultBlobUrl: input.resultBlobUrl, + snapshotManifestUrl: input.snapshotManifestUrl, + snapshotManifestKey: input.snapshotManifestKey, assetUrls: buildAtomicAssetUrlsMergeSql( input.assetUrlsByFilePath, ), @@ -124,6 +147,8 @@ const getParseResultProgressEffect: SourceParseResultRepository["getParseResultP db .select({ resultBlobUrl: sourceParseResults.resultBlobUrl, + snapshotManifestUrl: sourceParseResults.snapshotManifestUrl, + snapshotManifestKey: sourceParseResults.snapshotManifestKey, assetUrls: sourceParseResults.assetUrls, }) .from(sourceParseResults) @@ -135,6 +160,8 @@ const getParseResultProgressEffect: SourceParseResultRepository["getParseResultP return { resultBlobUrl: progress.resultBlobUrl, + snapshotManifestUrl: progress.snapshotManifestUrl, + snapshotManifestKey: progress.snapshotManifestKey, assetUrlsByFilePath: progress.assetUrls, } }) @@ -159,9 +186,17 @@ const getParseAssetUrlsEffect: SourceParseResultRepository["getParseAssetUrlsEff return row[0]?.assetUrls ?? {} }) +const getParseSnapshotMetadataEffect: SourceParseResultRepository["getParseSnapshotMetadataEffect"] = + (workspaceId: string, sourceId: string) => + Effect.gen(function* () { + const progress = yield* getParseResultProgressEffect(workspaceId, sourceId) + return progress + }) + export const sourceParseResultRepository: SourceParseResultRepository = { saveParseResultEffect, mergeParseAssetUrlsEffect, getParseResultProgressEffect, + getParseSnapshotMetadataEffect, getParseAssetUrlsEffect, } diff --git a/src/domains/sources/source-reconcile-route-workflow.test.ts b/src/domains/sources/source-reconcile-route-workflow.test.ts index 170941f..4821229 100644 --- a/src/domains/sources/source-reconcile-route-workflow.test.ts +++ b/src/domains/sources/source-reconcile-route-workflow.test.ts @@ -3,7 +3,6 @@ import { afterEach, describe, expect, it, vi } from "vitest" const mocks = vi.hoisted(() => ({ createParsedResultStorageAdapter: vi.fn(), saveParseResult: vi.fn(), - writeParsedResultAssetIndex: vi.fn(), loggerError: vi.fn(), loggerInfo: vi.fn(), loggerWarn: vi.fn(), @@ -27,7 +26,6 @@ vi.mock("@/domains/sources/workflow-runtime", () => ({ vi.mock("./parse-result-storage-adapter", () => ({ createParsedResultStorageAdapter: mocks.createParsedResultStorageAdapter, - writeParsedResultAssetIndex: mocks.writeParsedResultAssetIndex, })) vi.mock("@/integrations/knowhere", () => ({ @@ -66,7 +64,7 @@ describe("sourceReconcileRouteWorkflow", () => { }) }) - it("caches parsed assets before marking the source ready", async () => { + it("writes a parsed snapshot before marking the source ready", async () => { const context = createWorkflowContext() const continuations: ContinuationTriggerInput[] = [] const restore = @@ -86,13 +84,16 @@ describe("sourceReconcileRouteWorkflow", () => { "page_citation_assets/page-1.png": "https://blob.example/page_citation_assets/page-1.png", }, + parsedSnapshot: { + manifestKey: + "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + manifestUrl: + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + }, }) const client = { jobs: {}, knowledge: { cacheJobResult } } mocks.createParsedResultStorageAdapter.mockReturnValue(storageAdapter) mocks.makeKnowhereClient.mockReturnValue(client) - mocks.writeParsedResultAssetIndex.mockResolvedValue( - "https://blob.example/asset-index.json", - ) mocks.saveParseResult.mockResolvedValue({ id: "parse_result_1" }) mocks.pollSourceReconciliation.mockResolvedValue({ kind: "ready-to-prepare", @@ -124,19 +125,16 @@ describe("sourceReconcileRouteWorkflow", () => { jobId: "job_1", storageAdapter, }) - expect(mocks.writeParsedResultAssetIndex).toHaveBeenCalledWith({ - workspaceId: "workspace_1", - sourceId: "source_1", - assetUrlsByFilePath: { - "page_citation_assets/page-1.png": - "https://blob.example/page_citation_assets/page-1.png", - }, - }) expect(mocks.saveParseResult).toHaveBeenCalledWith( "workspace_1", "source_1", { - resultBlobUrl: "https://blob.example/asset-index.json", + resultBlobUrl: + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + snapshotManifestUrl: + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", assetUrlsByFilePath: { "page_citation_assets/page-1.png": "https://blob.example/page_citation_assets/page-1.png", @@ -152,6 +150,43 @@ describe("sourceReconcileRouteWorkflow", () => { expect(continuations).toEqual([]) }) + it("does not mark ready when snapshot manifest storage is missing", async () => { + const context = createWorkflowContext() + const storageAdapter = { + adapter: { + writeObject: vi.fn(), + }, + keyPrefix: "workspaces/workspace_1/sources/source_1/parsed-result", + } + const cacheJobResult = vi.fn().mockResolvedValue({ + assetUrlsByFilePath: {}, + }) + mocks.createParsedResultStorageAdapter.mockReturnValue(storageAdapter) + mocks.makeKnowhereClient.mockReturnValue({ + jobs: {}, + knowledge: { cacheJobResult }, + }) + mocks.pollSourceReconciliation.mockResolvedValue({ + kind: "ready-to-prepare", + jobId: "job_1", + documentId: "doc_1", + }) + + await expect( + sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ + context, + payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + }), + }), + ).rejects.toThrow("Parsed result snapshot was not written") + + expect(mocks.saveParseResult).not.toHaveBeenCalled() + expect(mocks.markSourceReadyAfterReconciliation).not.toHaveBeenCalled() + }) + it("triggers a fresh poll run when Knowhere is still running after the segment budget", async () => { const context = createWorkflowContext() const continuations: ContinuationTriggerInput[] = [] diff --git a/src/domains/sources/source-reconcile-route-workflow.ts b/src/domains/sources/source-reconcile-route-workflow.ts index cc03d40..cb28c20 100644 --- a/src/domains/sources/source-reconcile-route-workflow.ts +++ b/src/domains/sources/source-reconcile-route-workflow.ts @@ -10,7 +10,6 @@ import { makeKnowhereClient } from "@/integrations/knowhere" import { logger } from "@/lib/logger" import { createParsedResultStorageAdapter, - writeParsedResultAssetIndex, } from "./parse-result-storage-adapter" import { sourceWorkflowRuntime } from "./workflow-runtime" @@ -24,6 +23,14 @@ type ReconcilePayload = { type ReconcilePhase = "poll-and-ready" | "poll-and-mirror" | "asset-batches" +type ParsedSnapshotResponse = { + readonly assetUrlsByFilePath?: Readonly> + readonly parsedSnapshot?: { + readonly manifestKey: string + readonly manifestUrl?: string + } +} + type NormalizedReconcilePayload = { readonly workspaceId: string readonly sourceId: string @@ -122,27 +129,32 @@ async function runPollAndMirrorWorkflow(input: { return } - const assetUrlsByFilePath = await context.run("cache-job-result-assets", async () => { - const cachedResult = await client.knowledge.cacheJobResult({ + const snapshot = await context.run("cache-job-result-snapshot", async () => { + const cachedResult = (await client.knowledge.cacheJobResult({ jobId: jobToPrepare.jobId, storageAdapter: createParsedResultStorageAdapter({ workspaceId, sourceId, }), - }) - return cachedResult.assetUrlsByFilePath ?? {} + })) as ParsedSnapshotResponse + const manifest = cachedResult.parsedSnapshot + if (!manifest?.manifestUrl || !manifest.manifestKey) { + throw new Error( + "Parsed result snapshot was not written; refusing to mark source ready.", + ) + } + return { + assetUrlsByFilePath: cachedResult.assetUrlsByFilePath ?? {}, + snapshotManifestUrl: manifest.manifestUrl, + snapshotManifestKey: manifest.manifestKey, + } }) - const resultBlobUrl = await context.run("write-parse-result-asset-index", async () => - writeParsedResultAssetIndex({ - workspaceId, - sourceId, - assetUrlsByFilePath, - }), - ) await context.run("save-parse-result-assets", async () => sourceWorkflowRuntime.saveParseResult(workspaceId, sourceId, { - resultBlobUrl, - assetUrlsByFilePath, + resultBlobUrl: snapshot.snapshotManifestUrl, + snapshotManifestUrl: snapshot.snapshotManifestUrl, + snapshotManifestKey: snapshot.snapshotManifestKey, + assetUrlsByFilePath: snapshot.assetUrlsByFilePath, }), ) diff --git a/src/domains/sources/workflow-runtime.ts b/src/domains/sources/workflow-runtime.ts index bbec158..8527fab 100644 --- a/src/domains/sources/workflow-runtime.ts +++ b/src/domains/sources/workflow-runtime.ts @@ -62,6 +62,17 @@ type SourceWorkflowRuntime = UploadRepositoryRuntime & { sourceId: string, ) => Promise<{ readonly resultBlobUrl: string + readonly snapshotManifestUrl?: string | null + readonly snapshotManifestKey?: string | null + readonly assetUrlsByFilePath: Readonly> + } | null> + readonly getParseSnapshotMetadata: ( + workspaceId: string, + sourceId: string, + ) => Promise<{ + readonly resultBlobUrl: string + readonly snapshotManifestUrl?: string | null + readonly snapshotManifestKey?: string | null readonly assetUrlsByFilePath: Readonly> } | null> readonly listForWorkspace: (workspaceId: string) => Promise @@ -248,6 +259,12 @@ const getParseAssetUrls: SourceWorkflowRuntime["getParseAssetUrls"] = ( sourceRepository.getParseAssetUrlsEffect(workspaceId, sourceId), ) +const getParseSnapshotMetadata: SourceWorkflowRuntime["getParseSnapshotMetadata"] = + (workspaceId: string, sourceId: string) => + databaseRuntime.runPromise( + sourceRepository.getParseSnapshotMetadataEffect(workspaceId, sourceId), + ) + function createUploadRepository( runtime: UploadRepositoryRuntime = sourceWorkflowRuntime, ): UploadSourceRepository { @@ -287,6 +304,7 @@ export const sourceWorkflowRuntime: SourceWorkflowRuntime = { findInWorkspace, getParseAssetUrls, getParseResultProgress, + getParseSnapshotMetadata, hideDemoSource, listForWorkspace, listHiddenDemoSourceIds, diff --git a/src/infrastructure/db/schema.ts b/src/infrastructure/db/schema.ts index 720a303..5e62a2c 100644 --- a/src/infrastructure/db/schema.ts +++ b/src/infrastructure/db/schema.ts @@ -161,10 +161,9 @@ export type NewDemoSourceVisibility = typeof demoSourceVisibilities.$inferInsert /** * Notebook-owned parse-result artifact index for one source. * - * Knowhere's chunk list currently may omit media asset URLs, while parsed chunk - * metadata still points to ZIP-relative files like `images/image-1.jpg`. - * This table stores the Notebook Blob copy of the result ZIP plus a - * file-path-to-public-URL map for those extracted parsed artifacts. + * Blob is the Notebook-owned read model for parsed chunks after source + * reconciliation completes. This row stores the current parsed snapshot + * manifest and the file-path-to-public-URL map for parsed media artifacts. */ export const sourceParseResults = pgTable( "source_parse_results", @@ -175,6 +174,8 @@ export const sourceParseResults = pgTable( .references(() => sources.id, { onDelete: "cascade" }) .unique(), resultBlobUrl: text("result_blob_url").notNull(), + snapshotManifestUrl: text("snapshot_manifest_url"), + snapshotManifestKey: text("snapshot_manifest_key"), assetUrls: jsonb("asset_urls") .$type>>() .notNull(), From 6b30c48d7cebe1baa6593181613ad4f1765b1bf0 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 4 Jul 2026 03:24:37 +0800 Subject: [PATCH 07/51] Sync remote parsed snapshots into Notebook Blob --- .../sources/[sourceId]/chunks/route.test.ts | 361 ++++++++++++++++-- .../remote-parsed-snapshot-sync.test.ts | 316 +++++++++++++++ .../sources/remote-parsed-snapshot-sync.ts | 210 ++++++++++ src/domains/sources/route-chunks.ts | 30 +- src/domains/sources/route-dependencies.ts | 1 + src/domains/sources/route-types.ts | 7 + src/domains/sources/service.ts | 3 + 7 files changed, 897 insertions(+), 31 deletions(-) create mode 100644 src/domains/sources/remote-parsed-snapshot-sync.test.ts create mode 100644 src/domains/sources/remote-parsed-snapshot-sync.ts diff --git a/src/app/api/sources/[sourceId]/chunks/route.test.ts b/src/app/api/sources/[sourceId]/chunks/route.test.ts index c41319c..f9b9cbd 100644 --- a/src/app/api/sources/[sourceId]/chunks/route.test.ts +++ b/src/app/api/sources/[sourceId]/chunks/route.test.ts @@ -15,6 +15,7 @@ const mocks = vi.hoisted(() => ({ localizeRemoteDocument: vi.fn(), makeKnowhereClient: vi.fn(), requireUser: vi.fn(), + syncRemoteParsedSnapshot: vi.fn(), updateSourceRevisionKey: vi.fn(), })) @@ -54,6 +55,7 @@ vi.mock("@/domains/sources/service", () => ({ getParseAssetUrls: mocks.getSourceParseAssetUrls, getParseSnapshotMetadata: mocks.getSourceParseSnapshotMetadata, localizeRemoteDocument: mocks.localizeRemoteDocument, + syncRemoteParsedSnapshot: mocks.syncRemoteParsedSnapshot, updateSourceRevisionKey: mocks.updateSourceRevisionKey, }, })) @@ -84,6 +86,15 @@ describe("GET /api/sources/[sourceId]/chunks", () => { assetUrlsByFilePath: {}, }) mocks.updateSourceRevisionKey.mockResolvedValue(null) + mocks.syncRemoteParsedSnapshot.mockResolvedValue({ + resultBlobUrl: + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + snapshotManifestUrl: + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + assetUrlsByFilePath: {}, + }) }) it("serves API-owned demo chunks for anonymous canonical demo sources", async () => { @@ -644,7 +655,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled() }) - it("materializes a remote source id on open before loading chunks", async () => { + it("materializes a remote source id and loads chunks from the parsed Blob snapshot", async () => { const knowhereClient = { documents: { list: vi.fn(async () => ({ @@ -661,29 +672,10 @@ describe("GET /api/sources/[sourceId]/chunks", () => { }, ], })), - listChunks: vi.fn(async () => ({ - documentId: "doc_remote", - jobResultId: "job_result_1", - chunks: [ - { - id: "dchk_remote", - chunkId: "parser_remote", - chunkType: "text", - content: "Remote chunk", - sectionPath: "Summary", - sourceChunkPath: "Default_Root/remote.pdf/Summary", - filePath: null, - metadata: {}, - sortOrder: 0, - }, - ], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - })), + listChunks: vi.fn(), + }, + knowledge: { + cacheJobResult: vi.fn(), }, } mocks.getCurrentUser.mockResolvedValue({ @@ -700,7 +692,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { mocks.fetchDemoChunkPage.mockRejectedValue(new Error("not a demo")) mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") mocks.makeKnowhereClient.mockReturnValue(knowhereClient) - mocks.localizeRemoteDocument.mockResolvedValue({ + const localizedSource = { id: "00000000-0000-0000-0000-000000000009", workspaceId: "workspace_1", title: "remote.pdf", @@ -718,6 +710,79 @@ describe("GET /api/sources/[sourceId]/chunks", () => { createdAt: new Date("2026-05-10T00:00:00.000Z"), updatedAt: new Date("2026-05-10T00:00:00.000Z"), deletedAt: null, + } + mocks.localizeRemoteDocument.mockResolvedValue(localizedSource) + mocks.syncRemoteParsedSnapshot.mockResolvedValue({ + resultBlobUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + snapshotManifestUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + assetUrlsByFilePath: { + "pages/page-1.png": "https://blob.example/pages/page-1.png", + }, + }) + mocks.blobGet.mockImplementation(async (pathname: string) => { + if (pathname.endsWith("/manifest/current.json")) { + return { + statusCode: 200, + stream: createTextStream( + JSON.stringify({ + version: 1, + kind: "knowhere-parsed-result-snapshot", + jobId: "job_result_1", + documentId: "doc_remote", + sourceFileName: "remote.pdf", + totalChunks: 1, + chunkPageSize: 1, + chunkPages: [ + { + page: 1, + pageSize: 1, + chunkCount: 1, + key: "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/chunks/page-1.json", + }, + ], + assetUrlsByFilePath: { + "pages/page-1.png": "https://blob.example/pages/page-1.png", + }, + createdAt: "2026-07-03T00:00:00.000Z", + }), + ), + } + } + if (pathname.endsWith("/chunks/page-1.json")) { + return { + statusCode: 200, + stream: createTextStream( + JSON.stringify({ + version: 1, + jobId: "job_result_1", + documentId: "doc_remote", + sourceFileName: "remote.pdf", + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + chunks: [ + { + id: "dchk_remote", + chunkId: "parser_remote", + chunkType: "page", + content: "Remote chunk", + sectionPath: "Summary", + sourceChunkPath: "Default_Root/remote.pdf/Summary", + filePath: "pages/page-1.png", + metadata: {}, + sortOrder: 0, + }, + ], + }), + ), + } + } + return null }) const response = await GET( @@ -737,6 +802,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { chunkId: "dchk_remote", parserChunkId: "parser_remote", documentId: "doc_remote", + assetUrl: "https://blob.example/pages/page-1.png", sourceTitle: "remote.pdf", }, ], @@ -764,14 +830,253 @@ describe("GET /api/sources/[sourceId]/chunks", () => { revisionKey: "job_result_1", }, ) - expect(knowhereClient.documents.listChunks).toHaveBeenCalledWith( - "doc_remote", + expect(mocks.syncRemoteParsedSnapshot).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + source: localizedSource, + client: { + documents: knowhereClient.documents, + knowledge: knowhereClient.knowledge, + }, + }) + expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled() + expect(knowhereClient.knowledge.cacheJobResult).not.toHaveBeenCalled() + }) + + it("returns processing for a remote source when the parsed Blob snapshot is incomplete", async () => { + const knowhereClient = { + documents: { + list: vi.fn(async () => ({ + documents: [ + { + documentId: "doc_remote", + namespace: "default", + status: "active", + currentJobResultId: "job_result_1", + sourceFileName: "remote.pdf", + documentMetadata: { + mimeType: "application/pdf", + }, + }, + ], + })), + listChunks: vi.fn(), + }, + knowledge: { + cacheJobResult: vi.fn(), + }, + } + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.fetchDemoChunkPage.mockRejectedValue(new Error("not a demo")) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.makeKnowhereClient.mockReturnValue(knowhereClient) + mocks.localizeRemoteDocument.mockResolvedValue({ + id: "00000000-0000-0000-0000-000000000009", + workspaceId: "workspace_1", + title: "remote.pdf", + mimeType: "application/pdf", + sizeBytes: 0, + status: "ready", + failureReason: null, + knowhereJobId: "job_result_1", + knowhereDocumentId: "doc_remote", + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: null, + originalBlobUrl: null, + demoKey: null, + createdAt: new Date("2026-05-10T00:00:00.000Z"), + updatedAt: new Date("2026-05-10T00:00:00.000Z"), + deletedAt: null, + }) + mocks.syncRemoteParsedSnapshot.mockResolvedValue(null) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/knowhere-doc:default:doc_remote/chunks?page=1&pageSize=1", + ), { + params: Promise.resolve({ + sourceId: "knowhere-doc:default:doc_remote", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + chunks: [], + pagination: { page: 1, pageSize: 1, - includeAssetUrls: true, + total: 0, + totalPages: 0, + }, + message: "Source parsed snapshot is still being prepared.", + }) + expect(response.status).toBe(202) + expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled() + }) + + it("loads remote snapshot chunks after an interrupted sync left the source parsing", async () => { + const knowhereClient = { + documents: { + list: vi.fn(async () => ({ + documents: [ + { + documentId: "doc_remote", + namespace: "default", + status: "active", + currentJobResultId: "job_result_1", + sourceFileName: "remote.pdf", + documentMetadata: { + mimeType: "application/pdf", + }, + }, + ], + })), + listChunks: vi.fn(), + }, + knowledge: { + cacheJobResult: vi.fn(), + }, + } + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.fetchDemoChunkPage.mockRejectedValue(new Error("not a demo")) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.makeKnowhereClient.mockReturnValue(knowhereClient) + mocks.localizeRemoteDocument.mockResolvedValue({ + id: "00000000-0000-0000-0000-000000000009", + workspaceId: "workspace_1", + title: "remote.pdf", + mimeType: "application/pdf", + sizeBytes: 0, + status: "parsing", + failureReason: null, + knowhereJobId: "job_result_1", + knowhereDocumentId: "doc_remote", + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: null, + originalBlobUrl: null, + demoKey: null, + createdAt: new Date("2026-05-10T00:00:00.000Z"), + updatedAt: new Date("2026-05-10T00:00:00.000Z"), + deletedAt: null, + }) + mocks.syncRemoteParsedSnapshot.mockResolvedValue({ + resultBlobUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + snapshotManifestUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + assetUrlsByFilePath: {}, + }) + mocks.blobGet.mockImplementation(async (pathname: string) => { + if (pathname.endsWith("/manifest/current.json")) { + return { + statusCode: 200, + stream: createTextStream( + JSON.stringify({ + version: 1, + kind: "knowhere-parsed-result-snapshot", + jobId: "job_result_1", + documentId: "doc_remote", + sourceFileName: "remote.pdf", + totalChunks: 1, + chunkPageSize: 1, + chunkPages: [ + { + page: 1, + pageSize: 1, + chunkCount: 1, + key: "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/chunks/page-1.json", + }, + ], + assetUrlsByFilePath: {}, + createdAt: "2026-07-03T00:00:00.000Z", + }), + ), + } + } + if (pathname.endsWith("/chunks/page-1.json")) { + return { + statusCode: 200, + stream: createTextStream( + JSON.stringify({ + version: 1, + jobId: "job_result_1", + documentId: "doc_remote", + sourceFileName: "remote.pdf", + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + chunks: [ + { + id: "dchk_remote", + chunkId: "parser_remote", + chunkType: "text", + content: "Remote chunk", + sectionPath: "Summary", + sourceChunkPath: "Default_Root/remote.pdf/Summary", + metadata: {}, + sortOrder: 0, + }, + ], + }), + ), + } + } + return null + }) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/knowhere-doc:default:doc_remote/chunks?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "knowhere-doc:default:doc_remote", + }), }, ) + + await expect(response.json()).resolves.toMatchObject({ + chunks: [ + { + chunkId: "dchk_remote", + parserChunkId: "parser_remote", + documentId: "doc_remote", + sourceTitle: "remote.pdf", + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + }, + }) + expect(response.status).toBe(200) + expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled() }) }) diff --git a/src/domains/sources/remote-parsed-snapshot-sync.test.ts b/src/domains/sources/remote-parsed-snapshot-sync.test.ts new file mode 100644 index 0000000..8b0d7f9 --- /dev/null +++ b/src/domains/sources/remote-parsed-snapshot-sync.test.ts @@ -0,0 +1,316 @@ +import { afterEach, describe, expect, it, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + createParsedResultStorageAdapter: vi.fn(), + loggerInfo: vi.fn(), +})) + +vi.mock("./parse-result-storage-adapter", () => ({ + createParsedResultStorageAdapter: mocks.createParsedResultStorageAdapter, +})) + +vi.mock("@/lib/logger", () => ({ + logger: { + info: mocks.loggerInfo, + }, +})) + +import type { Source, SourceParseResult } from "@/infrastructure/db/schema" +import { syncRemoteParsedSnapshot } from "./remote-parsed-snapshot-sync" + +describe("syncRemoteParsedSnapshot", () => { + afterEach(() => { + vi.clearAllMocks() + }) + + it("caches an already-parsed remote document into the Notebook Blob snapshot", async () => { + const storageAdapter = { + adapter: { + writeObject: vi.fn(), + }, + keyPrefix: + "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result", + } + const listChunks = vi.fn(async () => ({ + jobResultId: "job_result_1", + })) + const cacheJobResult = vi.fn(async () => ({ + assetUrlsByFilePath: { + "pages/page-1.png": "https://blob.example/pages/page-1.png", + }, + parsedSnapshot: { + manifestKey: + "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + manifestUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + }, + })) + const repository = createRepository() + mocks.createParsedResultStorageAdapter.mockReturnValue(storageAdapter) + + const snapshot = await syncRemoteParsedSnapshot({ + workspaceId: "workspace_1", + source: makeSource({ + knowhereJobId: null, + }), + client: { + documents: { + listChunks, + }, + knowledge: { + cacheJobResult, + }, + }, + repository, + }) + + expect(listChunks).toHaveBeenCalledWith("doc_remote", { + page: 1, + pageSize: 1, + includeAssetUrls: false, + }) + expect(mocks.createParsedResultStorageAdapter).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "00000000-0000-0000-0000-000000000009", + }) + expect(cacheJobResult).toHaveBeenCalledWith({ + jobId: "job_result_1", + storageAdapter, + }) + expect(repository.markParsing).toHaveBeenCalledWith( + "workspace_1", + "00000000-0000-0000-0000-000000000009", + "job_result_1", + "doc_remote", + ) + expect(repository.saveParseResult).toHaveBeenCalledWith( + "workspace_1", + "00000000-0000-0000-0000-000000000009", + { + resultBlobUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + snapshotManifestUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + assetUrlsByFilePath: { + "pages/page-1.png": "https://blob.example/pages/page-1.png", + }, + }, + ) + expect(repository.markReady).toHaveBeenCalledWith( + "workspace_1", + "00000000-0000-0000-0000-000000000009", + "doc_remote", + ) + expect(snapshot).toEqual({ + resultBlobUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + snapshotManifestUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + assetUrlsByFilePath: { + "pages/page-1.png": "https://blob.example/pages/page-1.png", + }, + }) + }) + + it("returns the existing complete snapshot without calling Knowhere", async () => { + const existingSnapshot = { + resultBlobUrl: + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + snapshotManifestUrl: + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + assetUrlsByFilePath: {}, + } + const repository = createRepository({ + getParseSnapshotMetadata: vi.fn(async () => existingSnapshot), + }) + const listChunks = vi.fn() + const cacheJobResult = vi.fn() + + const snapshot = await syncRemoteParsedSnapshot({ + workspaceId: "workspace_1", + source: makeSource(), + client: { + documents: { + listChunks, + }, + knowledge: { + cacheJobResult, + }, + }, + repository, + }) + + expect(snapshot).toEqual(existingSnapshot) + expect(listChunks).not.toHaveBeenCalled() + expect(cacheJobResult).not.toHaveBeenCalled() + expect(repository.markParsing).not.toHaveBeenCalled() + expect(repository.markReady).not.toHaveBeenCalled() + }) + + it("marks a parsing source ready when a complete snapshot already exists", async () => { + const existingSnapshot = { + resultBlobUrl: + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + snapshotManifestUrl: + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + assetUrlsByFilePath: {}, + } + const repository = createRepository({ + getParseSnapshotMetadata: vi.fn(async () => existingSnapshot), + }) + const listChunks = vi.fn() + const cacheJobResult = vi.fn() + + const snapshot = await syncRemoteParsedSnapshot({ + workspaceId: "workspace_1", + source: makeSource({ + status: "parsing", + knowhereDocumentId: "doc_remote", + }), + client: { + documents: { + listChunks, + }, + knowledge: { + cacheJobResult, + }, + }, + repository, + }) + + expect(snapshot).toEqual(existingSnapshot) + expect(listChunks).not.toHaveBeenCalled() + expect(cacheJobResult).not.toHaveBeenCalled() + expect(repository.markReady).toHaveBeenCalledWith( + "workspace_1", + "00000000-0000-0000-0000-000000000009", + "doc_remote", + ) + }) + + it("retries an interrupted remote snapshot sync from a parsing source row", async () => { + const storageAdapter = { + adapter: { + writeObject: vi.fn(), + }, + keyPrefix: + "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result", + } + const listChunks = vi.fn() + const cacheJobResult = vi.fn(async () => ({ + assetUrlsByFilePath: {}, + parsedSnapshot: { + manifestKey: + "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + manifestUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + }, + })) + const repository = createRepository() + mocks.createParsedResultStorageAdapter.mockReturnValue(storageAdapter) + + await syncRemoteParsedSnapshot({ + workspaceId: "workspace_1", + source: makeSource({ + status: "parsing", + knowhereJobId: "job_result_1", + }), + client: { + documents: { + listChunks, + }, + knowledge: { + cacheJobResult, + }, + }, + repository, + }) + + expect(listChunks).not.toHaveBeenCalled() + expect(cacheJobResult).toHaveBeenCalledWith({ + jobId: "job_result_1", + storageAdapter, + }) + expect(repository.markReady).toHaveBeenCalledWith( + "workspace_1", + "00000000-0000-0000-0000-000000000009", + "doc_remote", + ) + }) +}) + +type SyncRemoteParsedSnapshotInput = Parameters< + typeof syncRemoteParsedSnapshot +>[0] + +type TestRemoteParsedSnapshotRepository = NonNullable< + SyncRemoteParsedSnapshotInput["repository"] +> + +function createRepository( + overrides: Partial = {}, +): TestRemoteParsedSnapshotRepository { + return { + ...createRepositoryShape(), + ...overrides, + } +} + +function createRepositoryShape(): TestRemoteParsedSnapshotRepository { + return { + getParseSnapshotMetadata: vi.fn(async () => null), + markParsing: vi.fn(async () => makeSource({ status: "parsing" })), + markReady: vi.fn(async () => makeSource()), + saveParseResult: vi.fn(async () => makeSourceParseResult()), + } +} + +function makeSource(overrides: Partial = {}): Source { + return { + id: "00000000-0000-0000-0000-000000000009", + workspaceId: "workspace_1", + title: "remote.pdf", + mimeType: "application/pdf", + sizeBytes: 0, + status: "ready", + failureReason: null, + knowhereJobId: "job_result_1", + knowhereDocumentId: "doc_remote", + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: null, + originalBlobUrl: null, + demoKey: null, + createdAt: new Date("2026-05-10T00:00:00.000Z"), + updatedAt: new Date("2026-05-10T00:00:00.000Z"), + deletedAt: null, + ...overrides, + } +} + +function makeSourceParseResult( + overrides: Partial = {}, +): SourceParseResult { + return { + id: "parse_result_1", + sourceId: "00000000-0000-0000-0000-000000000009", + resultBlobUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + snapshotManifestUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", + assetUrls: {}, + createdAt: new Date("2026-05-10T00:00:00.000Z"), + updatedAt: new Date("2026-05-10T00:00:00.000Z"), + ...overrides, + } +} diff --git a/src/domains/sources/remote-parsed-snapshot-sync.ts b/src/domains/sources/remote-parsed-snapshot-sync.ts new file mode 100644 index 0000000..fb7edac --- /dev/null +++ b/src/domains/sources/remote-parsed-snapshot-sync.ts @@ -0,0 +1,210 @@ +import "server-only" + +import { Effect } from "effect" + +import type { Source } from "@/infrastructure/db/schema" +import { logger } from "@/lib/logger" +import { createParsedResultStorageAdapter } from "./parse-result-storage-adapter" +import { sourceWorkflowRuntime } from "./workflow-runtime" + +type RemoteParsedSnapshotSource = Pick< + Source, + "id" | "workspaceId" | "knowhereDocumentId" | "knowhereJobId" | "status" +> + +type RemoteParsedSnapshotClient = { + readonly documents: { + listChunks( + documentId: string, + params: { + readonly page: number + readonly pageSize: number + readonly includeAssetUrls: boolean + }, + ): Promise<{ + readonly jobId?: string | null + readonly jobResultId?: string | null + }> + } + readonly knowledge: { + cacheJobResult(params: { + readonly jobId: string + readonly storageAdapter: ReturnType + }): Promise + } +} + +type RemoteParsedSnapshotResponse = { + readonly assetUrlsByFilePath?: Readonly> + readonly parsedSnapshot?: { + readonly manifestKey: string + readonly manifestUrl?: string + } +} + +type RemoteParsedSnapshotRepository = { + readonly getParseSnapshotMetadata: typeof sourceWorkflowRuntime.getParseSnapshotMetadata + readonly markParsing: typeof sourceWorkflowRuntime.markParsing + readonly markReady: typeof sourceWorkflowRuntime.markReady + readonly saveParseResult: typeof sourceWorkflowRuntime.saveParseResult +} + +type SyncRemoteParsedSnapshotInput = { + readonly workspaceId: string + readonly source: RemoteParsedSnapshotSource + readonly client: RemoteParsedSnapshotClient + readonly repository?: RemoteParsedSnapshotRepository +} + +type RemoteParsedSnapshotReadModel = { + readonly resultBlobUrl: string + readonly snapshotManifestUrl?: string | null + readonly snapshotManifestKey?: string | null + readonly assetUrlsByFilePath: Readonly> +} + +const jobProbePage = 1 +const jobProbePageSize = 1 + +const syncRemoteParsedSnapshotEffect = Effect.fn( + "syncRemoteParsedSnapshot", +)(function* ({ + workspaceId, + source, + client, + repository = sourceWorkflowRuntime, +}: SyncRemoteParsedSnapshotInput) { + const existingSnapshot = yield* Effect.tryPromise(() => + repository.getParseSnapshotMetadata(workspaceId, source.id), + ) + if (isCompleteSnapshot(existingSnapshot)) { + if (source.status === "parsing" && source.knowhereDocumentId) { + yield* Effect.tryPromise(() => + repository.markReady(workspaceId, source.id, source.knowhereDocumentId!), + ) + } + return existingSnapshot + } + + if (source.status !== "ready" && source.status !== "parsing") { + return existingSnapshot + } + + const documentId = source.knowhereDocumentId + if (!documentId) { + return yield* Effect.die( + new Error("Remote source is missing a Knowhere document id."), + ) + } + + const jobId = yield* resolveRemoteDocumentJobId({ + client, + documentId, + source, + }) + yield* Effect.tryPromise(() => + repository.markParsing(workspaceId, source.id, jobId, documentId), + ) + + const snapshot = yield* Effect.tryPromise(async () => { + const cachedResult = await client.knowledge.cacheJobResult({ + jobId, + storageAdapter: createParsedResultStorageAdapter({ + workspaceId, + sourceId: source.id, + }), + }) + const manifest = cachedResult.parsedSnapshot + if (!manifest?.manifestUrl || !manifest.manifestKey) { + throw new Error( + "Remote parsed snapshot was not written; refusing to mark source ready.", + ) + } + return { + assetUrlsByFilePath: cachedResult.assetUrlsByFilePath ?? {}, + snapshotManifestUrl: manifest.manifestUrl, + snapshotManifestKey: manifest.manifestKey, + } + }) + + const savedSnapshot = yield* Effect.tryPromise(() => + repository.saveParseResult(workspaceId, source.id, { + resultBlobUrl: snapshot.snapshotManifestUrl, + snapshotManifestUrl: snapshot.snapshotManifestUrl, + snapshotManifestKey: snapshot.snapshotManifestKey, + assetUrlsByFilePath: snapshot.assetUrlsByFilePath, + }), + ) + if (!savedSnapshot) { + return yield* Effect.die( + new Error("Remote parsed snapshot could not be saved for the source."), + ) + } + + yield* Effect.tryPromise(() => + repository.markReady(workspaceId, source.id, documentId), + ) + + logger.info("sources: remote parsed snapshot synced", { + sourceId: source.id, + documentId, + jobId, + snapshotManifestKey: snapshot.snapshotManifestKey, + }) + + return { + resultBlobUrl: snapshot.snapshotManifestUrl, + snapshotManifestUrl: snapshot.snapshotManifestUrl, + snapshotManifestKey: snapshot.snapshotManifestKey, + assetUrlsByFilePath: snapshot.assetUrlsByFilePath, + } +}) + +export async function syncRemoteParsedSnapshot( + input: SyncRemoteParsedSnapshotInput, +): Promise { + return Effect.runPromise(syncRemoteParsedSnapshotEffect(input)) +} + +const resolveRemoteDocumentJobId = Effect.fn("resolveRemoteDocumentJobId")( + function* ({ + client, + documentId, + source, + }: { + readonly client: RemoteParsedSnapshotClient + readonly documentId: string + readonly source: RemoteParsedSnapshotSource + }) { + if (source.knowhereJobId) return source.knowhereJobId + + const response = yield* Effect.tryPromise(() => + client.documents.listChunks(documentId, { + page: jobProbePage, + pageSize: jobProbePageSize, + includeAssetUrls: false, + }), + ) + const jobId = response.jobId ?? response.jobResultId + if (!jobId) { + return yield* Effect.die( + new Error("Remote document chunk metadata did not include a job id."), + ) + } + return jobId + }, +) + +function isCompleteSnapshot( + snapshot: RemoteParsedSnapshotReadModel | null, +): snapshot is RemoteParsedSnapshotReadModel & { + readonly snapshotManifestKey: string + readonly snapshotManifestUrl: string +} { + return ( + typeof snapshot?.snapshotManifestKey === "string" && + snapshot.snapshotManifestKey.length > 0 && + typeof snapshot.snapshotManifestUrl === "string" && + snapshot.snapshotManifestUrl.length > 0 + ) +} diff --git a/src/domains/sources/route-chunks.ts b/src/domains/sources/route-chunks.ts index fdbe0b1..ff9a174 100644 --- a/src/domains/sources/route-chunks.ts +++ b/src/domains/sources/route-chunks.ts @@ -94,7 +94,6 @@ const loadSourceChunksEffect = ( if (!isCompleteSnapshot(snapshot)) { return sourceSnapshotProcessing(input) } - if (input.shouldLoadAll) { const chunks = yield* deps.loadChunksForSource(source, client, { snapshot, @@ -168,9 +167,32 @@ const loadRemoteChunkPageEffect = ( revisionKey: remoteDocument.revisionKey ?? null, }), ) + const knowledge = client.knowledge + if (!knowledge) { + throw new Error("Knowhere client does not support parsed snapshot caching.") + } + const snapshot = yield* Effect.tryPromise(() => + deps.sourceService.syncRemoteParsedSnapshot({ + workspaceId: workspace.id, + source, + client: { + documents: client.documents, + knowledge, + }, + }), + ) + if (!isCompleteSnapshot(snapshot)) { + return sourceSnapshotProcessing(input) + } + const readySource = { + ...source, + status: "ready" as const, + knowhereDocumentId: source.knowhereDocumentId ?? remoteDocument.documentId, + } if (input.shouldLoadAll) { - const chunks = yield* deps.loadChunksForSource(source, client, { + const chunks = yield* deps.loadChunksForSource(readySource, client, { + snapshot, workspaceId: workspace.id, onRevisionKey: async (revisionKey) => { await deps.sourceService.updateSourceRevisionKey( @@ -184,10 +206,12 @@ const loadRemoteChunkPageEffect = ( } const chunkPage = yield* deps.loadChunkPageForSource( - source, + readySource, client, input.pageParams, { + assetUrlsByFilePath: snapshot.assetUrlsByFilePath, + snapshot, workspaceId: workspace.id, onRevisionKey: async (revisionKey) => { await deps.sourceService.updateSourceRevisionKey( diff --git a/src/domains/sources/route-dependencies.ts b/src/domains/sources/route-dependencies.ts index d48ede2..cdc5751 100644 --- a/src/domains/sources/route-dependencies.ts +++ b/src/domains/sources/route-dependencies.ts @@ -55,6 +55,7 @@ const defaultDependencies: SourceRouteServiceDependencies = { upsertMaterializedDemoSource: defaultSourceService.upsertMaterializedDemoSource, retrySourceToKnowhere: defaultSourceService.retrySourceToKnowhere, + syncRemoteParsedSnapshot: defaultSourceService.syncRemoteParsedSnapshot, uploadSourceBlobToKnowhere: defaultSourceService.uploadSourceBlobToKnowhere, uploadSourceToKnowhere: defaultSourceService.uploadSourceToKnowhere, }, diff --git a/src/domains/sources/route-types.ts b/src/domains/sources/route-types.ts index 79d0635..a8deed9 100644 --- a/src/domains/sources/route-types.ts +++ b/src/domains/sources/route-types.ts @@ -18,8 +18,13 @@ import type { import type { RouteResult } from "@/lib/route-result" import type { SourceBlobUploadInput } from "./blob-upload" import type { sourceViewOptionsBySourceId } from "./counts" +import type { syncRemoteParsedSnapshot } from "./remote-parsed-snapshot-sync" import type { UploadKnowhereClient } from "./upload" +type SyncRemoteParsedSnapshotInput = Parameters< + typeof syncRemoteParsedSnapshot +>[0] + type SourceRouteKnowhereClient = UploadKnowhereClient & ChunkKnowhereClient & { readonly documents: ChunkKnowhereClient["documents"] & { @@ -54,6 +59,7 @@ type SourceRouteKnowhereClient = UploadKnowhereClient & }> archive(documentId: string): Promise } + readonly knowledge?: SyncRemoteParsedSnapshotInput["client"]["knowledge"] } type SourceUploadRequest = @@ -171,6 +177,7 @@ type SourceWorkflowService = { source: Source, knowhere: UploadKnowhereClient, ) => Promise + readonly syncRemoteParsedSnapshot: typeof syncRemoteParsedSnapshot readonly findInWorkspace: ( workspaceId: string, sourceId: string, diff --git a/src/domains/sources/service.ts b/src/domains/sources/service.ts index 57701e6..2e25596 100644 --- a/src/domains/sources/service.ts +++ b/src/domains/sources/service.ts @@ -9,6 +9,7 @@ import { uploadSourceBlobToKnowhereEffect, uploadSourceToKnowhereEffect, } from "./upload" +import { syncRemoteParsedSnapshot } from "./remote-parsed-snapshot-sync" import { retrySourceToKnowhereEffect } from "./retry" import { sourceWorkflowRuntime } from "./workflow-runtime" @@ -65,6 +66,7 @@ type SourceService = { source: Source, knowhere: UploadKnowhereClient, ) => Promise + readonly syncRemoteParsedSnapshot: typeof syncRemoteParsedSnapshot } const uploadSourceToKnowhere: SourceService["uploadSourceToKnowhere"] = ( @@ -122,4 +124,5 @@ export const sourceService: SourceService = { uploadSourceToKnowhere, uploadSourceBlobToKnowhere, retrySourceToKnowhere, + syncRemoteParsedSnapshot, } From 95a879698315e26f2336389df45f347325281627 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 4 Jul 2026 15:50:01 +0800 Subject: [PATCH 08/51] Use server-safe parsed snapshot loads --- .../sources/[sourceId]/chunks/route.test.ts | 8 +++--- .../remote-parsed-snapshot-sync.test.ts | 26 +++++++++---------- .../sources/remote-parsed-snapshot-sync.ts | 8 +++--- .../source-reconcile-route-workflow.test.ts | 10 +++---- .../source-reconcile-route-workflow.ts | 8 +++--- 5 files changed, 30 insertions(+), 30 deletions(-) diff --git a/src/app/api/sources/[sourceId]/chunks/route.test.ts b/src/app/api/sources/[sourceId]/chunks/route.test.ts index f9b9cbd..f00d531 100644 --- a/src/app/api/sources/[sourceId]/chunks/route.test.ts +++ b/src/app/api/sources/[sourceId]/chunks/route.test.ts @@ -675,7 +675,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { listChunks: vi.fn(), }, knowledge: { - cacheJobResult: vi.fn(), + loadJobResult: vi.fn(), }, } mocks.getCurrentUser.mockResolvedValue({ @@ -839,7 +839,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { }, }) expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled() - expect(knowhereClient.knowledge.cacheJobResult).not.toHaveBeenCalled() + expect(knowhereClient.knowledge.loadJobResult).not.toHaveBeenCalled() }) it("returns processing for a remote source when the parsed Blob snapshot is incomplete", async () => { @@ -862,7 +862,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { listChunks: vi.fn(), }, knowledge: { - cacheJobResult: vi.fn(), + loadJobResult: vi.fn(), }, } mocks.getCurrentUser.mockResolvedValue({ @@ -945,7 +945,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { listChunks: vi.fn(), }, knowledge: { - cacheJobResult: vi.fn(), + loadJobResult: vi.fn(), }, } mocks.getCurrentUser.mockResolvedValue({ diff --git a/src/domains/sources/remote-parsed-snapshot-sync.test.ts b/src/domains/sources/remote-parsed-snapshot-sync.test.ts index 8b0d7f9..6a91582 100644 --- a/src/domains/sources/remote-parsed-snapshot-sync.test.ts +++ b/src/domains/sources/remote-parsed-snapshot-sync.test.ts @@ -23,7 +23,7 @@ describe("syncRemoteParsedSnapshot", () => { vi.clearAllMocks() }) - it("caches an already-parsed remote document into the Notebook Blob snapshot", async () => { + it("loads an already-parsed remote document into the Notebook Blob snapshot", async () => { const storageAdapter = { adapter: { writeObject: vi.fn(), @@ -34,7 +34,7 @@ describe("syncRemoteParsedSnapshot", () => { const listChunks = vi.fn(async () => ({ jobResultId: "job_result_1", })) - const cacheJobResult = vi.fn(async () => ({ + const loadJobResult = vi.fn(async () => ({ assetUrlsByFilePath: { "pages/page-1.png": "https://blob.example/pages/page-1.png", }, @@ -58,7 +58,7 @@ describe("syncRemoteParsedSnapshot", () => { listChunks, }, knowledge: { - cacheJobResult, + loadJobResult, }, }, repository, @@ -73,7 +73,7 @@ describe("syncRemoteParsedSnapshot", () => { workspaceId: "workspace_1", sourceId: "00000000-0000-0000-0000-000000000009", }) - expect(cacheJobResult).toHaveBeenCalledWith({ + expect(loadJobResult).toHaveBeenCalledWith({ jobId: "job_result_1", storageAdapter, }) @@ -130,7 +130,7 @@ describe("syncRemoteParsedSnapshot", () => { getParseSnapshotMetadata: vi.fn(async () => existingSnapshot), }) const listChunks = vi.fn() - const cacheJobResult = vi.fn() + const loadJobResult = vi.fn() const snapshot = await syncRemoteParsedSnapshot({ workspaceId: "workspace_1", @@ -140,7 +140,7 @@ describe("syncRemoteParsedSnapshot", () => { listChunks, }, knowledge: { - cacheJobResult, + loadJobResult, }, }, repository, @@ -148,7 +148,7 @@ describe("syncRemoteParsedSnapshot", () => { expect(snapshot).toEqual(existingSnapshot) expect(listChunks).not.toHaveBeenCalled() - expect(cacheJobResult).not.toHaveBeenCalled() + expect(loadJobResult).not.toHaveBeenCalled() expect(repository.markParsing).not.toHaveBeenCalled() expect(repository.markReady).not.toHaveBeenCalled() }) @@ -167,7 +167,7 @@ describe("syncRemoteParsedSnapshot", () => { getParseSnapshotMetadata: vi.fn(async () => existingSnapshot), }) const listChunks = vi.fn() - const cacheJobResult = vi.fn() + const loadJobResult = vi.fn() const snapshot = await syncRemoteParsedSnapshot({ workspaceId: "workspace_1", @@ -180,7 +180,7 @@ describe("syncRemoteParsedSnapshot", () => { listChunks, }, knowledge: { - cacheJobResult, + loadJobResult, }, }, repository, @@ -188,7 +188,7 @@ describe("syncRemoteParsedSnapshot", () => { expect(snapshot).toEqual(existingSnapshot) expect(listChunks).not.toHaveBeenCalled() - expect(cacheJobResult).not.toHaveBeenCalled() + expect(loadJobResult).not.toHaveBeenCalled() expect(repository.markReady).toHaveBeenCalledWith( "workspace_1", "00000000-0000-0000-0000-000000000009", @@ -205,7 +205,7 @@ describe("syncRemoteParsedSnapshot", () => { "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result", } const listChunks = vi.fn() - const cacheJobResult = vi.fn(async () => ({ + const loadJobResult = vi.fn(async () => ({ assetUrlsByFilePath: {}, parsedSnapshot: { manifestKey: @@ -228,14 +228,14 @@ describe("syncRemoteParsedSnapshot", () => { listChunks, }, knowledge: { - cacheJobResult, + loadJobResult, }, }, repository, }) expect(listChunks).not.toHaveBeenCalled() - expect(cacheJobResult).toHaveBeenCalledWith({ + expect(loadJobResult).toHaveBeenCalledWith({ jobId: "job_result_1", storageAdapter, }) diff --git a/src/domains/sources/remote-parsed-snapshot-sync.ts b/src/domains/sources/remote-parsed-snapshot-sync.ts index fb7edac..3e939c4 100644 --- a/src/domains/sources/remote-parsed-snapshot-sync.ts +++ b/src/domains/sources/remote-parsed-snapshot-sync.ts @@ -27,7 +27,7 @@ type RemoteParsedSnapshotClient = { }> } readonly knowledge: { - cacheJobResult(params: { + loadJobResult(params: { readonly jobId: string readonly storageAdapter: ReturnType }): Promise @@ -107,21 +107,21 @@ const syncRemoteParsedSnapshotEffect = Effect.fn( ) const snapshot = yield* Effect.tryPromise(async () => { - const cachedResult = await client.knowledge.cacheJobResult({ + const loadedResult = await client.knowledge.loadJobResult({ jobId, storageAdapter: createParsedResultStorageAdapter({ workspaceId, sourceId: source.id, }), }) - const manifest = cachedResult.parsedSnapshot + const manifest = loadedResult.parsedSnapshot if (!manifest?.manifestUrl || !manifest.manifestKey) { throw new Error( "Remote parsed snapshot was not written; refusing to mark source ready.", ) } return { - assetUrlsByFilePath: cachedResult.assetUrlsByFilePath ?? {}, + assetUrlsByFilePath: loadedResult.assetUrlsByFilePath ?? {}, snapshotManifestUrl: manifest.manifestUrl, snapshotManifestKey: manifest.manifestKey, } diff --git a/src/domains/sources/source-reconcile-route-workflow.test.ts b/src/domains/sources/source-reconcile-route-workflow.test.ts index 4821229..ea68cbe 100644 --- a/src/domains/sources/source-reconcile-route-workflow.test.ts +++ b/src/domains/sources/source-reconcile-route-workflow.test.ts @@ -79,7 +79,7 @@ describe("sourceReconcileRouteWorkflow", () => { }, keyPrefix: "workspaces/workspace_1/sources/source_1/parsed-result", } - const cacheJobResult = vi.fn().mockResolvedValue({ + const loadJobResult = vi.fn().mockResolvedValue({ assetUrlsByFilePath: { "page_citation_assets/page-1.png": "https://blob.example/page_citation_assets/page-1.png", @@ -91,7 +91,7 @@ describe("sourceReconcileRouteWorkflow", () => { "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", }, }) - const client = { jobs: {}, knowledge: { cacheJobResult } } + const client = { jobs: {}, knowledge: { loadJobResult } } mocks.createParsedResultStorageAdapter.mockReturnValue(storageAdapter) mocks.makeKnowhereClient.mockReturnValue(client) mocks.saveParseResult.mockResolvedValue({ id: "parse_result_1" }) @@ -121,7 +121,7 @@ describe("sourceReconcileRouteWorkflow", () => { workspaceId: "workspace_1", sourceId: "source_1", }) - expect(cacheJobResult).toHaveBeenCalledWith({ + expect(loadJobResult).toHaveBeenCalledWith({ jobId: "job_1", storageAdapter, }) @@ -158,13 +158,13 @@ describe("sourceReconcileRouteWorkflow", () => { }, keyPrefix: "workspaces/workspace_1/sources/source_1/parsed-result", } - const cacheJobResult = vi.fn().mockResolvedValue({ + const loadJobResult = vi.fn().mockResolvedValue({ assetUrlsByFilePath: {}, }) mocks.createParsedResultStorageAdapter.mockReturnValue(storageAdapter) mocks.makeKnowhereClient.mockReturnValue({ jobs: {}, - knowledge: { cacheJobResult }, + knowledge: { loadJobResult }, }) mocks.pollSourceReconciliation.mockResolvedValue({ kind: "ready-to-prepare", diff --git a/src/domains/sources/source-reconcile-route-workflow.ts b/src/domains/sources/source-reconcile-route-workflow.ts index cb28c20..16f9054 100644 --- a/src/domains/sources/source-reconcile-route-workflow.ts +++ b/src/domains/sources/source-reconcile-route-workflow.ts @@ -129,22 +129,22 @@ async function runPollAndMirrorWorkflow(input: { return } - const snapshot = await context.run("cache-job-result-snapshot", async () => { - const cachedResult = (await client.knowledge.cacheJobResult({ + const snapshot = await context.run("load-job-result-snapshot", async () => { + const loadedResult = (await client.knowledge.loadJobResult({ jobId: jobToPrepare.jobId, storageAdapter: createParsedResultStorageAdapter({ workspaceId, sourceId, }), })) as ParsedSnapshotResponse - const manifest = cachedResult.parsedSnapshot + const manifest = loadedResult.parsedSnapshot if (!manifest?.manifestUrl || !manifest.manifestKey) { throw new Error( "Parsed result snapshot was not written; refusing to mark source ready.", ) } return { - assetUrlsByFilePath: cachedResult.assetUrlsByFilePath ?? {}, + assetUrlsByFilePath: loadedResult.assetUrlsByFilePath ?? {}, snapshotManifestUrl: manifest.manifestUrl, snapshotManifestKey: manifest.manifestKey, } From 0304aeee051ff36f5bf1cd572237475668f3bc74 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 4 Jul 2026 17:22:26 +0800 Subject: [PATCH 09/51] Sync legacy parsed sources before chunk reads --- .../sources/[sourceId]/chunks/route.test.ts | 23 +++ src/components/chunks-panel.tsx | 14 +- src/components/source-row.tsx | 7 +- src/components/workspace-citation-focus.ts | 3 + .../workspace-selected-chunks.test.ts | 32 +++ src/components/workspace-selected-chunks.ts | 30 ++- src/components/workspace-shell-layout.tsx | 3 + src/components/workspace-shell.tsx | 1 + src/domains/chat/route-answer.ts | 28 ++- src/domains/chat/route-service.test.ts | 63 ++++++ .../demo/workspace-source-resolution.ts | 2 +- src/domains/sources/counts.test.ts | 70 ++++++- src/domains/sources/counts.ts | 80 ++++++-- src/domains/sources/route-chunks.ts | 28 ++- src/domains/sources/route-dependencies.ts | 2 + src/domains/sources/route-listing.ts | 8 +- src/domains/sources/route-service.test.ts | 184 ++++++++++++++++++ src/domains/sources/route-types.ts | 10 + src/domains/sources/service.ts | 60 ++++++ src/domains/sources/types.ts | 2 +- src/domains/workspace/client.test.ts | 27 +++ src/domains/workspace/client.ts | 4 + src/domains/workspace/initial-state.ts | 8 +- src/infrastructure/db/schema.ts | 10 +- 24 files changed, 645 insertions(+), 54 deletions(-) diff --git a/src/app/api/sources/[sourceId]/chunks/route.test.ts b/src/app/api/sources/[sourceId]/chunks/route.test.ts index f00d531..911646a 100644 --- a/src/app/api/sources/[sourceId]/chunks/route.test.ts +++ b/src/app/api/sources/[sourceId]/chunks/route.test.ts @@ -9,6 +9,7 @@ const mocks = vi.hoisted(() => ({ ensureWorkspace: vi.fn(), fetchDemoChunkPage: vi.fn(), findSourceInWorkspace: vi.fn(), + ensureParsedSnapshotForRead: vi.fn(), getCurrentUser: vi.fn(), getSourceParseAssetUrls: vi.fn(), getSourceParseSnapshotMetadata: vi.fn(), @@ -51,6 +52,7 @@ vi.mock("@vercel/blob", () => ({ vi.mock("@/domains/sources/service", () => ({ sourceService: { + ensureParsedSnapshotForRead: mocks.ensureParsedSnapshotForRead, findInWorkspace: mocks.findSourceInWorkspace, getParseAssetUrls: mocks.getSourceParseAssetUrls, getParseSnapshotMetadata: mocks.getSourceParseSnapshotMetadata, @@ -85,6 +87,15 @@ describe("GET /api/sources/[sourceId]/chunks", () => { "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", assetUrlsByFilePath: {}, }) + mocks.ensureParsedSnapshotForRead.mockResolvedValue({ + resultBlobUrl: + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + snapshotManifestUrl: + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + assetUrlsByFilePath: {}, + }) mocks.updateSourceRevisionKey.mockResolvedValue(null) mocks.syncRemoteParsedSnapshot.mockResolvedValue({ resultBlobUrl: @@ -504,6 +515,17 @@ describe("GET /api/sources/[sourceId]/chunks", () => { "images/chart.png": "https://blob.example/images/chart.png", }, }) + mocks.ensureParsedSnapshotForRead.mockResolvedValue({ + resultBlobUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000002/parsed-result/manifest/current.json", + snapshotManifestUrl: + "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000002/parsed-result/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000002/parsed-result/manifest/current.json", + assetUrlsByFilePath: { + "images/chart.png": "https://blob.example/images/chart.png", + }, + }) mocks.blobGet.mockImplementation(async (pathname: string) => { if (pathname.endsWith("/manifest/current.json")) { return { @@ -633,6 +655,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") mocks.makeKnowhereClient.mockReturnValue(knowhereClient) mocks.getSourceParseSnapshotMetadata.mockResolvedValue(null) + mocks.ensureParsedSnapshotForRead.mockResolvedValue(null) const response = await GET( new NextRequest( diff --git a/src/components/chunks-panel.tsx b/src/components/chunks-panel.tsx index 898ca39..78cd170 100644 --- a/src/components/chunks-panel.tsx +++ b/src/components/chunks-panel.tsx @@ -57,6 +57,7 @@ export type ChunksPanelProps = { isLoadingMore?: boolean; isLoadingAllChunks?: boolean; hasMoreChunks?: boolean; + processingMessage?: string | null; onLoadMore?: () => void; onLoadAllChunks?: () => void; onLoginClick?: () => void; @@ -84,6 +85,7 @@ export function ChunksPanel({ isLoadingMore = false, isLoadingAllChunks = false, hasMoreChunks = false, + processingMessage = null, onLoadMore, onLoadAllChunks, onLoginClick, @@ -360,7 +362,7 @@ export function ChunksPanel({ className="mx-auto flex w-[90%] min-w-0 max-w-[1600px] flex-col items-center p-3 sm:p-6" > {isLoading ? ( - + ) : chunks.length === 0 ? ( selectedSource ? ( @@ -1138,13 +1140,19 @@ function EmptyChunks(): ReactNode { ); } -function LoadingChunks(): ReactNode { +function LoadingChunks({ + message, +}: { + readonly message?: string | null; +}): ReactNode { return (
-

Loading parsed chunks...

+

+ {message ?? "Loading parsed chunks..."} +

); } diff --git a/src/components/source-row.tsx b/src/components/source-row.tsx index abfe8cb..e964c93 100644 --- a/src/components/source-row.tsx +++ b/src/components/source-row.tsx @@ -102,7 +102,7 @@ export function SourceRow({ }`} > {isReady - ? `${getReadySourceLabel(source)} · ${source.chunkCount ?? 0} chunks` + ? getReadySourceStatusText(source) : source.status === "parsing" ? "Preparing" : source.status === "uploading" @@ -197,6 +197,11 @@ function getReadySourceLabel(source: SourceView): string { return "Processed"; } +function getReadySourceStatusText(source: SourceView): string { + if (typeof source.chunkCount !== "number") return getReadySourceLabel(source); + return `${getReadySourceLabel(source)} · ${source.chunkCount} chunks`; +} + function fileIconTint(title: string): { bg: string; fg: string } { const ext = title.split(".").pop()?.toLowerCase(); switch (ext) { diff --git a/src/components/workspace-citation-focus.ts b/src/components/workspace-citation-focus.ts index b0d0f66..3485dd2 100644 --- a/src/components/workspace-citation-focus.ts +++ b/src/components/workspace-citation-focus.ts @@ -43,6 +43,7 @@ type WorkspaceCitationFocus = { readonly requestChunkFocus: (chunkId: string | null) => void readonly isSelectedChunksLoading: boolean readonly isSelectedChunksLoadingMore: boolean + readonly selectedChunksMessage: string | null readonly selectedChunks: ParsedChunkView[] readonly selectedSource: SourceView | undefined } @@ -80,6 +81,7 @@ export function useWorkspaceCitationFocus({ handleLoadMoreChunks, isSelectedChunksLoading, isSelectedChunksLoadingMore, + selectedChunksMessage, selectedChunks, selectedSource, } = useWorkspaceSelectedChunks({ @@ -277,6 +279,7 @@ export function useWorkspaceCitationFocus({ pendingCitationId, prefetchedChunksBySourceId, requestChunkFocus, + selectedChunksMessage, selectedChunks, selectedSource, } diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index 2772c85..fa08df5 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -129,6 +129,38 @@ describe("useWorkspaceSelectedChunks", () => { }); }); + it("treats a processing chunk page as a loading state", async () => { + fetchChunkPageMock.mockResolvedValue({ + chunks: [], + isProcessing: true, + message: "Source parsed snapshot is still being prepared.", + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }); + + const { result } = renderHook( + () => + useWorkspaceSelectedChunks({ + selectedSourceId: "source_1", + sources: [readySource], + prefetchedChunksBySourceId: {}, + }), + { wrapper: createSWRWrapper }, + ); + + await waitFor(() => + expect(result.current.selectedChunksMessage).toBe( + "Source parsed snapshot is still being prepared.", + ), + ); + expect(result.current.isSelectedChunksLoading).toBe(true); + expect(result.current.selectedChunks).toEqual([]); + }); + it("returns an empty chunk list when no source is selected", () => { const { result } = renderHook( () => diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index 4ec9f97..6d79175 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -24,6 +24,7 @@ type WorkspaceSelectedChunks = { readonly handleLoadMoreChunks: () => void readonly isSelectedChunksLoading: boolean readonly isSelectedChunksLoadingMore: boolean + readonly selectedChunksMessage: string | null readonly selectedChunks: ParsedChunkView[] readonly selectedSource: SourceView | undefined } @@ -57,8 +58,11 @@ export function useWorkspaceSelectedChunks({ { revalidateIfStale: false, keepPreviousData: false, + refreshInterval: (pages: readonly SourceChunksResponse[] | undefined) => + hasProcessingChunkPage(pages) ? 2_000 : 0, }, ) + const selectedChunksMessage = getSelectedChunksMessage(selectedChunkPages) const pagedSelectedChunks = useMemo( () => resolveChunkConnectionTargets( @@ -90,10 +94,11 @@ export function useWorkspaceSelectedChunks({ typeof selectedChunkPages[selectedChunkPageCount - 1] === "undefined", ) const isSelectedChunksLoading = - selectedChunkSourceId !== null && - !prefetchedSelectedChunks && - !selectedChunkPages && - isChunksLoading + Boolean(selectedChunksMessage) || + (selectedChunkSourceId !== null && + !prefetchedSelectedChunks && + !selectedChunkPages && + isChunksLoading) function handleLoadMoreChunks(): void { if (!hasMoreSelectedChunks || isSelectedChunksLoadingMore) return @@ -105,6 +110,7 @@ export function useWorkspaceSelectedChunks({ handleLoadMoreChunks, isSelectedChunksLoading, isSelectedChunksLoadingMore, + selectedChunksMessage, selectedChunks, selectedSource, } @@ -118,6 +124,22 @@ function fetchChunksByKey([ return workspaceClient.fetchChunkPage(sourceId, page) } +function hasProcessingChunkPage( + pages: readonly SourceChunksResponse[] | undefined, +): boolean { + return pages?.some((page) => page.isProcessing) ?? false +} + +function getSelectedChunksMessage( + pages: readonly SourceChunksResponse[] | undefined, +): string | null { + const page = pages?.find( + (candidate) => + candidate.isProcessing && typeof candidate.message === "string", + ) + return page?.message ?? null +} + function mergeVisibleChunkAssetUrls( chunks: readonly ParsedChunkView[], visibleChunks: readonly ParsedChunkView[], diff --git a/src/components/workspace-shell-layout.tsx b/src/components/workspace-shell-layout.tsx index 235eb07..3b9e997 100644 --- a/src/components/workspace-shell-layout.tsx +++ b/src/components/workspace-shell-layout.tsx @@ -80,6 +80,7 @@ export type WorkspaceShellLayoutProps = { readonly pendingCitationId: string | null readonly readySourceCount: number readonly selectedChunks: readonly ParsedChunkView[] + readonly selectedChunksMessage: string | null readonly selectedSourceFile: SourceOriginalFileView | null readonly selectedSourceId: string | null readonly selectedSourceTitle: string | null @@ -271,6 +272,7 @@ export function WorkspaceShellLayout( isLoading={props.isSelectedChunksLoading} isLoadingAllChunks={props.isSelectedAllChunksLoading} isLoadingMore={props.isSelectedChunksLoadingMore} + processingMessage={props.selectedChunksMessage} hasMoreChunks={props.hasMoreSelectedChunks} onLoadAllChunks={props.onLoadAllChunks} onLoadMore={props.onLoadMoreChunks} @@ -412,6 +414,7 @@ export function WorkspaceShellLayout( isLoading={props.isSelectedChunksLoading} isLoadingAllChunks={props.isSelectedAllChunksLoading} isLoadingMore={props.isSelectedChunksLoadingMore} + processingMessage={props.selectedChunksMessage} hasMoreChunks={props.hasMoreSelectedChunks} onLoadAllChunks={props.onLoadAllChunks} onLoadMore={props.onLoadMoreChunks} diff --git a/src/components/workspace-shell.tsx b/src/components/workspace-shell.tsx index 254c91e..4dd8b9c 100644 --- a/src/components/workspace-shell.tsx +++ b/src/components/workspace-shell.tsx @@ -236,6 +236,7 @@ function WorkspaceShellContent({ pendingCitationId={citationFocus.pendingCitationId} readySourceCount={sourceWorkflow.readySourceCount} selectedChunks={citationFocus.selectedChunks} + selectedChunksMessage={citationFocus.selectedChunksMessage} selectedSourceFile={citationFocus.selectedSource?.originalFile ?? null} selectedSourceId={sourceWorkflow.selectedSourceId} selectedSourceTitle={selectedSourceTitle} diff --git a/src/domains/chat/route-answer.ts b/src/domains/chat/route-answer.ts index ea75943..3df55e0 100644 --- a/src/domains/chat/route-answer.ts +++ b/src/domains/chat/route-answer.ts @@ -68,8 +68,32 @@ const answerChatEffect = (input: AnswerChatInput) => apiKey, }), ) - const loadSourceAssetUrls = (source: (typeof sources)[number]) => - sourceService.getParseAssetUrls(workspace.id, source.id) + const loadSourceAssetUrls = async ( + source: (typeof sources)[number], + ): Promise>> => { + try { + const snapshot = await sourceService.ensureParsedSnapshotForRead({ + workspaceId: workspace.id, + source, + client: client.knowledge + ? { + documents: client.documents, + knowledge: client.knowledge, + } + : null, + }) + if (snapshot) return snapshot.assetUrlsByFilePath + } catch (error) { + logger.warn("chat: parsed snapshot sync for assets failed", { + workspaceId: workspace.id, + sourceId: source.id, + documentId: source.knowhereDocumentId, + error: summarizeUnknownError(error), + }) + } + + return sourceService.getParseAssetUrls(workspace.id, source.id) + } const result: Either.Either = yield* Effect.tryPromise(() => diff --git a/src/domains/chat/route-service.test.ts b/src/domains/chat/route-service.test.ts index 4358e8a..189e862 100644 --- a/src/domains/chat/route-service.test.ts +++ b/src/domains/chat/route-service.test.ts @@ -73,6 +73,7 @@ vi.mock("@/lib/logger", () => ({ import { chatAnswerRouteService } from "./route-answer" import { chatThreadRouteService } from "./route-threads" +import { sourceService } from "@/domains/sources/service" describe("chat route services", () => { beforeEach(() => { @@ -140,6 +141,68 @@ describe("chat route services", () => { ) }) + it("ensures a legacy ready source snapshot before chat loads citation assets", async () => { + const workspace = makeWorkspace() + const documents = { listChunks: vi.fn() } + const knowledge = { loadJobResult: vi.fn() } + const client = { documents, knowledge, retrieval: { query: vi.fn() } } + const readySource = makeSource({ + status: "ready", + knowhereDocumentId: "doc_legacy", + knowhereJobId: null, + }) + const snapshot = { + resultBlobUrl: "https://blob.example/manifest/current.json", + snapshotManifestUrl: "https://blob.example/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", + assetUrlsByFilePath: { + "pages/page-1.png": "https://blob.example/pages/page-1.png", + }, + } + const ensureParsedSnapshotForRead = vi + .spyOn(sourceService, "ensureParsedSnapshotForRead") + .mockResolvedValue(snapshot) + mocks.getAuthenticatedWithClient.mockResolvedValue({ + user: { id: "user_1" }, + workspace, + apiKey: "jwt_123", + client, + }) + mocks.listSourcesForWorkspace.mockResolvedValue([readySource]) + mocks.handleChatTurn.mockImplementation( + async (input: { + readonly loadSourceAssetUrls?: ( + source: Source, + ) => Promise>> + }) => { + const assetUrls = await input.loadSourceAssetUrls?.(readySource) + expect(assetUrls).toEqual(snapshot.assetUrlsByFilePath) + return Either.right({ + threadId: "thread_1", + messages: [ + { id: "message_user", role: "user", content: "Show the page" }, + { id: "message_assistant", role: "assistant", content: "Answer" }, + ], + }) + }, + ) + + const result = await chatAnswerRouteService.answerChat({ + body: { message: "Show the page" }, + }) + + expect(result.status).toBe(200) + expect(ensureParsedSnapshotForRead).toHaveBeenCalledWith({ + workspaceId: workspace.id, + source: readySource, + client: { + documents, + knowledge, + }, + }) + }) + it("triggers background reconciliation for parsing sources without blocking chat", async () => { const workspace = makeWorkspace() const client = { retrieval: { query: vi.fn() } } diff --git a/src/domains/demo/workspace-source-resolution.ts b/src/domains/demo/workspace-source-resolution.ts index 39abc5a..d0918af 100644 --- a/src/domains/demo/workspace-source-resolution.ts +++ b/src/domains/demo/workspace-source-resolution.ts @@ -39,7 +39,7 @@ export function resolveWorkspaceDemoSources( } } -export function getWorkspaceSourcesNeedingKnowhereChunkCount( +export function getWorkspaceSourcesNeedingChunkCount( sources: readonly Source[], ): Source[] { return sources.filter((source) => !source.demoKey) diff --git a/src/domains/sources/counts.test.ts b/src/domains/sources/counts.test.ts index 22cdda9..e520eff 100644 --- a/src/domains/sources/counts.test.ts +++ b/src/domains/sources/counts.test.ts @@ -29,13 +29,36 @@ function makeSource(overrides: Partial = {}): Source { } describe("countChunksBySourceId", () => { - it("counts chunks only for ready sources with a Knowhere document id", async () => { - const listChunks = vi.fn().mockResolvedValue({ - pagination: { total: 12 }, - }) + it("counts ready source chunks from the parsed snapshot manifest", async () => { + const listChunks = vi.fn() const mockClient = { documents: { listChunks }, } as unknown as Knowhere + const repository = { + getParseSnapshotMetadata: vi.fn(async (workspaceId: string, sourceId: string) => + sourceId === "ready" + ? { + resultBlobUrl: "https://blob.example/manifest/current.json", + snapshotManifestUrl: "https://blob.example/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/ready/parsed-result/manifest/current.json", + assetUrlsByFilePath: {}, + } + : null, + ), + } + const readSnapshotManifest = vi.fn(async () => ({ + version: 1 as const, + kind: "knowhere-parsed-result-snapshot" as const, + jobId: "job_1", + documentId: "doc_ready", + sourceFileName: "notes.pdf", + totalChunks: 12, + chunkPageSize: 50, + chunkPages: [], + assetUrlsByFilePath: {}, + createdAt: "2026-07-04T00:00:00.000Z", + })) const { countChunksBySourceId } = await import("./counts") @@ -47,22 +70,44 @@ describe("countChunksBySourceId", () => { makeSource({ id: "missing-doc", knowhereDocumentId: null }), ], mockClient, + { + repository, + readSnapshotManifest, + }, ), ) - expect(listChunks).toHaveBeenCalledOnce() - expect(listChunks).toHaveBeenCalledWith("doc_ready", { - page: 1, - pageSize: 1, + expect(listChunks).not.toHaveBeenCalled() + expect(repository.getParseSnapshotMetadata).toHaveBeenCalledWith( + "workspace_1", + "ready", + ) + expect(readSnapshotManifest).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "ready", + manifestKey: + "workspaces/workspace_1/sources/ready/parsed-result/manifest/current.json", }) expect(counts).toEqual(new Map([["ready", 12]])) }) - it("skips a source count when Knowhere chunks lookup fails", async () => { - const listChunks = vi.fn().mockRejectedValue(new Error("temporary outage")) + it("skips a source count when snapshot manifest lookup fails", async () => { + const listChunks = vi.fn() const mockClient = { documents: { listChunks }, } as unknown as Knowhere + const repository = { + getParseSnapshotMetadata: vi.fn(async () => ({ + resultBlobUrl: "https://blob.example/manifest/current.json", + snapshotManifestUrl: "https://blob.example/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/ready/parsed-result/manifest/current.json", + assetUrlsByFilePath: {}, + })), + } + const readSnapshotManifest = vi.fn(async () => { + throw new Error("temporary outage") + }) const { countChunksBySourceId } = await import("./counts") @@ -70,10 +115,15 @@ describe("countChunksBySourceId", () => { countChunksBySourceId( [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], mockClient, + { + repository, + readSnapshotManifest, + }, ), ) expect(counts.size).toBe(0) + expect(listChunks).not.toHaveBeenCalled() }) it("does not count materialized demo sources through their copied document id", async () => { diff --git a/src/domains/sources/counts.ts b/src/domains/sources/counts.ts index 310a694..8f934bb 100644 --- a/src/domains/sources/counts.ts +++ b/src/domains/sources/counts.ts @@ -1,15 +1,32 @@ import "server-only" -import { Effect, Either } from "effect" +import { Effect } from "effect" import type Knowhere from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" +import { + readParsedResultSnapshotManifest, + type ParsedResultSnapshotManifest, +} from "./parse-result-storage-adapter" +import { sourceWorkflowRuntime } from "./workflow-runtime" + +type ChunkCountRepository = Pick< + typeof sourceWorkflowRuntime, + "getParseSnapshotMetadata" +> + +type ChunkCountOptions = { + readonly readSnapshotManifest?: typeof readParsedResultSnapshotManifest + readonly repository?: ChunkCountRepository +} export const countChunksBySourceId = ( sources: readonly Source[], client: Knowhere, + options: ChunkCountOptions = {}, ) => Effect.gen(function* () { + void client const readySources = sources.filter( (source) => !source.demoKey && @@ -21,23 +38,18 @@ export const countChunksBySourceId = ( const entries = yield* Effect.all( readySources.map((source) => Effect.gen(function* () { - const documentId = source.knowhereDocumentId - if (!documentId) return [source.id, undefined] as const - - const result = yield* Effect.either( - Effect.tryPromise(() => - client.documents.listChunks(documentId, { - page: 1, - pageSize: 1, - }), + const manifest = yield* Effect.tryPromise(() => + loadSourceSnapshotManifest(source, options), + ).pipe( + Effect.catchAll(() => + Effect.succeed(null), ), ) - - if (Either.isLeft(result)) return [source.id, undefined] as const + if (!manifest) return [source.id, undefined] as const return [ source.id, - result.right.pagination.total, + manifest.totalChunks, ] as const }), ), @@ -55,9 +67,10 @@ export const countChunksBySourceId = ( export const sourceViewOptionsBySourceId = ( sources: readonly Source[], client: Knowhere, + options: ChunkCountOptions = {}, ) => Effect.gen(function* () { - const counts = yield* countChunksBySourceId(sources, client) + const counts = yield* countChunksBySourceId(sources, client, options) return new Map( sources.map((source) => [ source.id, @@ -65,3 +78,42 @@ export const sourceViewOptionsBySourceId = ( ]), ) }) + +async function loadSourceSnapshotManifest( + source: Source, + options: ChunkCountOptions, +): Promise { + const repository = options.repository ?? sourceWorkflowRuntime + const readSnapshotManifest = + options.readSnapshotManifest ?? readParsedResultSnapshotManifest + const snapshot = await repository.getParseSnapshotMetadata( + source.workspaceId, + source.id, + ) + if (!isCompleteSnapshot(snapshot)) return null + + return readSnapshotManifest({ + workspaceId: source.workspaceId, + sourceId: source.id, + manifestKey: snapshot.snapshotManifestKey, + }) +} + +function isCompleteSnapshot( + snapshot: + | { + readonly snapshotManifestKey?: string | null + readonly snapshotManifestUrl?: string | null + } + | null, +): snapshot is { + readonly snapshotManifestKey: string + readonly snapshotManifestUrl: string +} { + return ( + typeof snapshot?.snapshotManifestKey === "string" && + snapshot.snapshotManifestKey.length > 0 && + typeof snapshot.snapshotManifestUrl === "string" && + snapshot.snapshotManifestUrl.length > 0 + ) +} diff --git a/src/domains/sources/route-chunks.ts b/src/domains/sources/route-chunks.ts index ff9a174..b8fbca2 100644 --- a/src/domains/sources/route-chunks.ts +++ b/src/domains/sources/route-chunks.ts @@ -89,7 +89,28 @@ const loadSourceChunksEffect = ( getClientForWorkspace(workspace.id, input.cookieHeader, deps), ) const snapshot = yield* Effect.tryPromise(() => - deps.sourceService.getParseSnapshotMetadata(workspace.id, source.id), + deps.sourceService.ensureParsedSnapshotForRead({ + workspaceId: workspace.id, + source, + client: client.knowledge + ? { + documents: client.documents, + knowledge: client.knowledge, + } + : null, + }), + ).pipe( + Effect.catchAll((error) => + Effect.sync(() => { + logger.warn("sources: parsed snapshot sync for read failed", { + workspaceId: workspace.id, + sourceId: source.id, + documentId: source.knowhereDocumentId, + error: getErrorMessage(error), + }) + return null + }), + ), ) if (!isCompleteSnapshot(snapshot)) { return sourceSnapshotProcessing(input) @@ -109,15 +130,12 @@ const loadSourceChunksEffect = ( return routeResult.ok({ chunks }) } - const assetUrlsByFilePath = yield* Effect.tryPromise(() => - deps.sourceService.getParseAssetUrls(workspace.id, source.id), - ) const chunkPage = yield* deps.loadChunkPageForSource( source, client, input.pageParams, { - assetUrlsByFilePath, + assetUrlsByFilePath: snapshot.assetUrlsByFilePath, snapshot, workspaceId: workspace.id, onRevisionKey: async (revisionKey) => { diff --git a/src/domains/sources/route-dependencies.ts b/src/domains/sources/route-dependencies.ts index cdc5751..af81a3c 100644 --- a/src/domains/sources/route-dependencies.ts +++ b/src/domains/sources/route-dependencies.ts @@ -44,6 +44,8 @@ const defaultDependencies: SourceRouteServiceDependencies = { ), requireUser, sourceService: { + ensureParsedSnapshotForRead: + defaultSourceService.ensureParsedSnapshotForRead, findInWorkspace: defaultSourceService.findInWorkspace, getParseAssetUrls: defaultSourceService.getParseAssetUrls, getParseSnapshotMetadata: defaultSourceService.getParseSnapshotMetadata, diff --git a/src/domains/sources/route-listing.ts b/src/domains/sources/route-listing.ts index 9476285..35c0c72 100644 --- a/src/domains/sources/route-listing.ts +++ b/src/domains/sources/route-listing.ts @@ -3,7 +3,7 @@ import { Effect } from "effect" import { demoView } from "@/domains/demo/view" import { getMaterializedDemoSourceViewOptionsBySourceId, - getWorkspaceSourcesNeedingKnowhereChunkCount, + getWorkspaceSourcesNeedingChunkCount, resolveWorkspaceDemoSources, } from "@/domains/demo/workspace-source-resolution" import { routeResult } from "@/lib/route-result" @@ -96,8 +96,8 @@ const listSourcesEffect = ( client, localSources: demoSourceResolution.workspaceSources, }) - const sourcesNeedingKnowhereChunkCount = - getWorkspaceSourcesNeedingKnowhereChunkCount(workspaceSources) + const sourcesNeedingChunkCount = + getWorkspaceSourcesNeedingChunkCount(workspaceSources) const materializedDemoSourceOptions = getMaterializedDemoSourceViewOptionsBySourceId(workspaceSources, catalog) yield* Effect.sync(() => @@ -111,7 +111,7 @@ const listSourcesEffect = ( }), ) const sourceOptions = yield* deps.getSourceViewOptionsBySourceId( - sourcesNeedingKnowhereChunkCount, + sourcesNeedingChunkCount, client, ) const hiddenDemoSourceIds = new Set( diff --git a/src/domains/sources/route-service.test.ts b/src/domains/sources/route-service.test.ts index 3d11b96..ca37390 100644 --- a/src/domains/sources/route-service.test.ts +++ b/src/domains/sources/route-service.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it, vi } from "vitest"; import { Effect } from "effect"; import type { Job } from "@ontos-ai/knowhere-sdk"; +import type { ChunkPage } from "@/domains/chunks"; import type { Source, Workspace } from "@/infrastructure/db/schema"; import { createRouteListing } from "./route-listing"; import { createSourceRouteService } from "./route-service"; @@ -119,6 +120,189 @@ describe("source route service", () => { expect(listHiddenDemoSourceIds).toHaveBeenCalledWith(workspace.id); }); + it("syncs a ready legacy source snapshot before loading displayed chunks", async () => { + const readySource: Source = { + ...source, + id: "00000000-0000-4000-8000-000000000001", + knowhereJobId: null, + knowhereDocumentId: "doc_legacy", + status: "ready", + }; + const snapshot = { + resultBlobUrl: "https://blob.example/manifest/current.json", + snapshotManifestUrl: "https://blob.example/manifest/current.json", + snapshotManifestKey: + "workspaces/workspace_1/sources/00000000-0000-4000-8000-000000000001/parsed-result/manifest/current.json", + assetUrlsByFilePath: { + "images/chart.png": "https://blob.example/images/chart.png", + }, + }; + const listChunks = vi.fn(async () => { + throw new Error("display path must not call Knowhere listChunks"); + }); + const knowledge = { + loadJobResult: vi.fn(), + }; + const knowhereClient = { + documents: { + archive: vi.fn(async () => undefined), + listChunks, + }, + jobs: { + create: vi.fn(), + get: vi.fn(), + upload: vi.fn(), + }, + knowledge, + }; + const loadedPage: ChunkPage = { + chunks: [ + { + chunkId: "chunk_1", + type: "text", + content: "Chunk body", + sourceTitle: "notes.pdf", + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 1, + totalPages: 1, + }, + }; + const loadChunkPageForSource = vi.fn(() => Effect.succeed(loadedPage)); + const syncRemoteParsedSnapshot = vi.fn(async () => snapshot); + const service = createSourceRouteService({ + ensureApiKeyForWorkspace: vi.fn(async () => "jwt_123"), + ensureWorkspace: vi.fn(async () => workspace), + getCurrentUser: vi.fn(async () => ({ + id: "user_1", + email: null, + name: null, + })), + loadChunkPageForSource, + makeKnowhereClient: vi.fn(() => knowhereClient), + sourceService: { + findInWorkspace: vi.fn(async () => readySource), + ensureParsedSnapshotForRead: syncRemoteParsedSnapshot, + updateSourceRevisionKey: vi.fn(async () => readySource), + }, + }); + + const result = await service.loadSourceChunks({ + cookieHeader: "session=abc", + sourceId: readySource.id, + shouldLoadAll: false, + pageParams: { page: 1, pageSize: 50 }, + }); + + expect(result).toEqual({ + status: 200, + body: { + chunks: [ + { + chunkId: "chunk_1", + type: "text", + content: "Chunk body", + sourceTitle: "notes.pdf", + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 1, + totalPages: 1, + }, + }, + }); + expect(syncRemoteParsedSnapshot).toHaveBeenCalledWith({ + workspaceId: workspace.id, + source: readySource, + client: { + documents: knowhereClient.documents, + knowledge, + }, + }); + expect(loadChunkPageForSource).toHaveBeenCalledWith( + readySource, + knowhereClient, + { page: 1, pageSize: 50 }, + expect.objectContaining({ + assetUrlsByFilePath: snapshot.assetUrlsByFilePath, + snapshot, + workspaceId: workspace.id, + }), + ); + expect(listChunks).not.toHaveBeenCalled(); + }); + + it("returns processing when a ready legacy source snapshot is not synced yet", async () => { + const readySource: Source = { + ...source, + id: "00000000-0000-4000-8000-000000000002", + knowhereDocumentId: "doc_legacy", + knowhereJobId: null, + status: "ready", + }; + const knowledge = { + loadJobResult: vi.fn(), + }; + const knowhereClient = { + documents: { + archive: vi.fn(async () => undefined), + listChunks: vi.fn(), + }, + jobs: { + create: vi.fn(), + get: vi.fn(), + upload: vi.fn(), + }, + knowledge, + }; + const loadChunkPageForSource = vi.fn(); + const service = createSourceRouteService({ + ensureApiKeyForWorkspace: vi.fn(async () => "jwt_123"), + ensureWorkspace: vi.fn(async () => workspace), + getCurrentUser: vi.fn(async () => ({ + id: "user_1", + email: null, + name: null, + })), + loadChunkPageForSource, + makeKnowhereClient: vi.fn(() => knowhereClient), + sourceService: { + findInWorkspace: vi.fn(async () => readySource), + ensureParsedSnapshotForRead: vi.fn(async () => { + throw new Error("snapshot import still running"); + }), + }, + }); + + const result = await service.loadSourceChunks({ + cookieHeader: "session=abc", + sourceId: readySource.id, + shouldLoadAll: false, + pageParams: { page: 1, pageSize: 50 }, + }); + + expect(result).toEqual({ + status: 202, + body: { + chunks: [], + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + message: "Source parsed snapshot is still being prepared.", + }, + }); + expect(loadChunkPageForSource).not.toHaveBeenCalled(); + expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled(); + }); + it("lists shared default and legacy namespace documents as lightweight remote sources", async () => { const localReadySource: Source = { ...source, diff --git a/src/domains/sources/route-types.ts b/src/domains/sources/route-types.ts index a8deed9..a32cad2 100644 --- a/src/domains/sources/route-types.ts +++ b/src/domains/sources/route-types.ts @@ -162,6 +162,16 @@ type SourceRouteService = { } type SourceWorkflowService = { + readonly ensureParsedSnapshotForRead: (input: { + readonly workspaceId: string + readonly source: Source + readonly client?: SyncRemoteParsedSnapshotInput["client"] | null + }) => Promise<{ + readonly resultBlobUrl: string + readonly snapshotManifestUrl?: string | null + readonly snapshotManifestKey?: string | null + readonly assetUrlsByFilePath: Readonly> + } | null> readonly uploadSourceToKnowhere: ( workspace: Workspace, file: File, diff --git a/src/domains/sources/service.ts b/src/domains/sources/service.ts index 2e25596..11e9199 100644 --- a/src/domains/sources/service.ts +++ b/src/domains/sources/service.ts @@ -14,6 +14,9 @@ import { retrySourceToKnowhereEffect } from "./retry" import { sourceWorkflowRuntime } from "./workflow-runtime" type SourceService = { + readonly ensureParsedSnapshotForRead: ( + input: EnsureParsedSnapshotForReadInput, + ) => Promise readonly findInWorkspace: ( workspaceId: string, sourceId: string, @@ -69,6 +72,19 @@ type SourceService = { readonly syncRemoteParsedSnapshot: typeof syncRemoteParsedSnapshot } +type EnsureParsedSnapshotForReadInput = { + readonly workspaceId: string + readonly source: Source + readonly client?: Parameters[0]["client"] | null +} + +type EnsureParsedSnapshotForReadResult = { + readonly resultBlobUrl: string + readonly snapshotManifestUrl?: string | null + readonly snapshotManifestKey?: string | null + readonly assetUrlsByFilePath: Readonly> +} + const uploadSourceToKnowhere: SourceService["uploadSourceToKnowhere"] = ( workspace: Workspace, file: File, @@ -109,7 +125,25 @@ const retrySourceToKnowhere: SourceService["retrySourceToKnowhere"] = ( }), ) +const ensureParsedSnapshotForRead: SourceService["ensureParsedSnapshotForRead"] = + async ({ workspaceId, source, client }: EnsureParsedSnapshotForReadInput) => { + const existingSnapshot = + await sourceWorkflowRuntime.getParseSnapshotMetadata(workspaceId, source.id) + if (isCompleteParsedSnapshot(existingSnapshot)) return existingSnapshot + + if (!canSyncParsedSnapshotForRead(source, client)) { + return existingSnapshot + } + + return syncRemoteParsedSnapshot({ + workspaceId, + source, + client, + }) + } + export const sourceService: SourceService = { + ensureParsedSnapshotForRead, findInWorkspace: sourceWorkflowRuntime.findInWorkspace, getParseAssetUrls: sourceWorkflowRuntime.getParseAssetUrls, getParseSnapshotMetadata: sourceWorkflowRuntime.getParseSnapshotMetadata, @@ -126,3 +160,29 @@ export const sourceService: SourceService = { retrySourceToKnowhere, syncRemoteParsedSnapshot, } + +function canSyncParsedSnapshotForRead( + source: Source, + client: EnsureParsedSnapshotForReadInput["client"], +): client is Parameters[0]["client"] { + return ( + Boolean(client) && + (source.status === "ready" || source.status === "parsing") && + typeof source.knowhereDocumentId === "string" && + source.knowhereDocumentId.length > 0 + ) +} + +function isCompleteParsedSnapshot( + snapshot: EnsureParsedSnapshotForReadResult | null, +): snapshot is EnsureParsedSnapshotForReadResult & { + readonly snapshotManifestKey: string + readonly snapshotManifestUrl: string +} { + return ( + typeof snapshot?.snapshotManifestKey === "string" && + snapshot.snapshotManifestKey.length > 0 && + typeof snapshot.snapshotManifestUrl === "string" && + snapshot.snapshotManifestUrl.length > 0 + ) +} diff --git a/src/domains/sources/types.ts b/src/domains/sources/types.ts index e1abe9e..8479f87 100644 --- a/src/domains/sources/types.ts +++ b/src/domains/sources/types.ts @@ -48,7 +48,7 @@ export type SourceView = { readonly originalFile?: SourceOriginalFileView /** Official Library metadata when this row is an API-owned catalog item. */ readonly officialLibrary?: SourceOfficialLibraryView - /** Count from the Knowhere chunks API, not a local aggregate. */ + /** Count from the Notebook parsed snapshot manifest when available. */ readonly chunkCount?: number /** User opt-out for this query session. Drives excludeDocumentIds. */ readonly excludedFromQuery?: boolean diff --git a/src/domains/workspace/client.test.ts b/src/domains/workspace/client.test.ts index 42e8c65..bc19bba 100644 --- a/src/domains/workspace/client.test.ts +++ b/src/domains/workspace/client.test.ts @@ -63,6 +63,33 @@ describe("workspaceClient", () => { }) }) + it("preserves source chunk processing messages", async () => { + mockRouteClient.getJson.mockResolvedValue({ + chunks: [], + message: "Source parsed snapshot is still being prepared.", + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }) + + const page = await workspaceClient.fetchChunkPage("source_1", 1) + + expect(page).toEqual({ + chunks: [], + isProcessing: true, + message: "Source parsed snapshot is still being prepared.", + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }) + }) + it("throws materialization route errors instead of treating them as empty sources", async () => { mockRouteClient.postJsonWithStatus.mockResolvedValue({ status: 502, diff --git a/src/domains/workspace/client.ts b/src/domains/workspace/client.ts index 9a7b3cd..a650b6f 100644 --- a/src/domains/workspace/client.ts +++ b/src/domains/workspace/client.ts @@ -24,6 +24,8 @@ const workspaceClientConfig = { type SourceChunksResponse = { chunks?: ParsedChunkView[] + isProcessing?: boolean + message?: string pagination?: { page: number pageSize: number @@ -128,6 +130,8 @@ async function fetchChunkPage( return { chunks: Array.isArray(body.chunks) ? body.chunks : [], + ...(typeof body.message === "string" ? { message: body.message } : {}), + ...(typeof body.message === "string" ? { isProcessing: true } : {}), pagination: body.pagination, } } diff --git a/src/domains/workspace/initial-state.ts b/src/domains/workspace/initial-state.ts index 368d43b..851a34b 100644 --- a/src/domains/workspace/initial-state.ts +++ b/src/domains/workspace/initial-state.ts @@ -7,7 +7,7 @@ import type { ParsedChunkView } from "@/domains/chunks/types" import { demoView } from "@/domains/demo/view" import { getMaterializedDemoSourceViewOptionsBySourceId, - getWorkspaceSourcesNeedingKnowhereChunkCount, + getWorkspaceSourcesNeedingChunkCount, resolveWorkspaceDemoSources, } from "@/domains/demo/workspace-source-resolution" import { chatThreadService } from "@/domains/chat/thread-service" @@ -331,8 +331,8 @@ export const loadWorkspaceShellInitialStateEffect = ( localSources: demoSourceResolution.workspaceSources, }), ) - const sourcesNeedingKnowhereChunkCount = - getWorkspaceSourcesNeedingKnowhereChunkCount(workspaceSources) + const sourcesNeedingChunkCount = + getWorkspaceSourcesNeedingChunkCount(workspaceSources) const materializedDemoSourceOptions = getMaterializedDemoSourceViewOptionsBySourceId( workspaceSources, @@ -354,7 +354,7 @@ export const loadWorkspaceShellInitialStateEffect = ( operation: "sourceViewOptionsBySourceId", }, deps.sourceViewOptionsBySourceId( - sourcesNeedingKnowhereChunkCount, + sourcesNeedingChunkCount, client, ), ) diff --git a/src/infrastructure/db/schema.ts b/src/infrastructure/db/schema.ts index 5e62a2c..8c909a9 100644 --- a/src/infrastructure/db/schema.ts +++ b/src/infrastructure/db/schema.ts @@ -17,8 +17,8 @@ import { * - Postgres stores only metadata, status, Knowhere IDs, and chat * threads/messages. * - It does NOT store file bytes or chunk copies in Postgres. Original - * uploads and parsed media artifacts live in Blob storage; chunks are - * fetched on demand from Knowhere's chunks API. + * uploads and parsed-source snapshots live in Blob storage; retrieval + * stays upstream in Knowhere. * * Soft delete: * - Every user-visible resource has a nullable `deleted_at` timestamp. @@ -67,9 +67,9 @@ export type NewWorkspace = typeof workspaces.$inferInsert; * - `status` — lifecycle: uploading | parsing | ready | failed * - `failure_reason` — human-readable error text when status=failed * - `knowhere_job_id` — set once the parse job is created - * - `knowhere_document_id` — set when parsing completes; the sole handle - * used to fetch chunks and to exclude a source - * from a retrieval query + * - `knowhere_document_id` — set when parsing completes; used to import + * parsed snapshots and to exclude a source from a + * retrieval query * - `original_blob_*` — public Blob pointer for the original upload preview * and download path * - `staged_blob_*` — legacy temporary Blob staging pointer retained for From 1269fb22fe74df4a0b7e9a6ba844bc91bc6509f8 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 5 Jul 2026 02:18:47 +0800 Subject: [PATCH 10/51] Move parsed chunk reads and sync into the SDK storage model Replace the hand-rolled Blob snapshot machinery with the SDK's ParsedDocumentStorage: a Vercel-Blob adapter, a QStash-backed resumable sync scheduler + /api/sources/parsed-sync route, and chunk reads routed through knowledge.readChunks/getDocumentOutline (assetUrlPolicy:"durable"). Parse flow syncs to storage before markReady and stages sync failures as failure_stage="storage_sync"; retry resumes sync without reparsing. Chat citation hardening is unified onto the durable read path. Adds migration 0011 (revision_key/sync_status/sync_error on source_parse_results, failure_stage on sources). Deletes chunks/server, remote-parsed-snapshot-sync, and parse-result-storage-adapter. Co-Authored-By: Claude Opus 4.8 --- drizzle/0011_rare_deadpool.sql | 5 + drizzle/meta/0011_snapshot.json | 752 +++++++++++++++ drizzle/meta/_journal.json | 7 + .../demo-sources/materialize/route.test.ts | 1 + .../sources/[sourceId]/chunks/route.test.ts | 682 +++---------- src/app/api/sources/parsed-sync/route.ts | 27 + src/app/api/sources/route.test.ts | 1 + src/domains/chat/index.test.ts | 1 + .../chat/media-asset-hardening.test.ts | 283 ++---- src/domains/chat/media-asset-hardening.ts | 470 +++------ src/domains/chat/media-assets.test.ts | 1 + src/domains/chat/page-citation-assets.test.ts | 1 + src/domains/chat/route-answer.ts | 27 +- src/domains/chat/route-service.test.ts | 83 +- src/domains/chat/service.test.ts | 1 + src/domains/chunks/index.test.ts | 1 + src/domains/chunks/index.ts | 31 +- src/domains/chunks/read.test.ts | 123 +++ src/domains/chunks/read.ts | 142 +++ src/domains/chunks/server.test.ts | 563 ----------- src/domains/chunks/server.ts | 911 ------------------ src/domains/sources/counts.test.ts | 78 +- src/domains/sources/counts.ts | 90 +- .../sources/parse-result-storage-adapter.ts | 269 ------ .../parsed-document-blob-storage.test.ts | 223 +++++ .../sources/parsed-document-blob-storage.ts | 291 ++++++ .../sources/parsed-document-sync-scheduler.ts | 154 +++ .../parsed-sync-route-workflow.test.ts | 201 ++++ .../sources/parsed-sync-route-workflow.ts | 226 +++++ src/domains/sources/reconcile.test.ts | 1 + .../remote-parsed-snapshot-sync.test.ts | 316 ------ .../sources/remote-parsed-snapshot-sync.ts | 210 ---- src/domains/sources/repository.ts | 2 + src/domains/sources/retry.test.ts | 2 + src/domains/sources/route-chunks.ts | 194 ++-- src/domains/sources/route-dependencies.ts | 52 +- src/domains/sources/route-retry.test.ts | 107 ++ src/domains/sources/route-retry.ts | 73 +- src/domains/sources/route-service.test.ts | 157 +-- src/domains/sources/route-types.ts | 41 - src/domains/sources/service.ts | 79 -- .../sources/source-parse-result-repository.ts | 85 +- .../source-reconcile-route-workflow.test.ts | 173 ++-- .../source-reconcile-route-workflow.ts | 121 ++- .../sources/source-reconcile-workflow.test.ts | 1 + src/domains/sources/source-row-repository.ts | 6 + src/domains/sources/upload.test.ts | 1 + src/domains/sources/view.test.ts | 1 + src/domains/sources/workflow-runtime.test.ts | 1 + src/domains/sources/workflow-runtime.ts | 44 +- src/domains/workspace/initial-state.test.ts | 1 + src/domains/workspace/integration.test.ts | 2 +- src/infrastructure/db/schema.ts | 22 +- src/integrations/knowhere.ts | 53 + 54 files changed, 3276 insertions(+), 4114 deletions(-) create mode 100644 drizzle/0011_rare_deadpool.sql create mode 100644 drizzle/meta/0011_snapshot.json create mode 100644 src/app/api/sources/parsed-sync/route.ts create mode 100644 src/domains/chunks/read.test.ts create mode 100644 src/domains/chunks/read.ts delete mode 100644 src/domains/chunks/server.test.ts delete mode 100644 src/domains/chunks/server.ts delete mode 100644 src/domains/sources/parse-result-storage-adapter.ts create mode 100644 src/domains/sources/parsed-document-blob-storage.test.ts create mode 100644 src/domains/sources/parsed-document-blob-storage.ts create mode 100644 src/domains/sources/parsed-document-sync-scheduler.ts create mode 100644 src/domains/sources/parsed-sync-route-workflow.test.ts create mode 100644 src/domains/sources/parsed-sync-route-workflow.ts delete mode 100644 src/domains/sources/remote-parsed-snapshot-sync.test.ts delete mode 100644 src/domains/sources/remote-parsed-snapshot-sync.ts create mode 100644 src/domains/sources/route-retry.test.ts diff --git a/drizzle/0011_rare_deadpool.sql b/drizzle/0011_rare_deadpool.sql new file mode 100644 index 0000000..0cf5000 --- /dev/null +++ b/drizzle/0011_rare_deadpool.sql @@ -0,0 +1,5 @@ +ALTER TABLE "source_parse_results" ALTER COLUMN "result_blob_url" DROP NOT NULL;--> statement-breakpoint +ALTER TABLE "source_parse_results" ADD COLUMN "revision_key" text;--> statement-breakpoint +ALTER TABLE "source_parse_results" ADD COLUMN "sync_status" text;--> statement-breakpoint +ALTER TABLE "source_parse_results" ADD COLUMN "sync_error" text;--> statement-breakpoint +ALTER TABLE "sources" ADD COLUMN "failure_stage" text; \ No newline at end of file diff --git a/drizzle/meta/0011_snapshot.json b/drizzle/meta/0011_snapshot.json new file mode 100644 index 0000000..0f8448c --- /dev/null +++ b/drizzle/meta/0011_snapshot.json @@ -0,0 +1,752 @@ +{ + "id": "bab84445-8eb2-4d15-856e-b2f59edcaf47", + "prevId": "cc41c815-4e52-4d01-ac2e-abbc791331b7", + "version": "7", + "dialect": "postgresql", + "tables": { + "public.chat_messages": { + "name": "chat_messages", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "thread_id": { + "name": "thread_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "role": { + "name": "role", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "content": { + "name": "content", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "citations": { + "name": "citations", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "artifacts": { + "name": "artifacts", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "chat_messages_thread_created_idx": { + "name": "chat_messages_thread_created_idx", + "columns": [ + { + "expression": "thread_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "created_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "chat_messages_thread_id_chat_threads_id_fk": { + "name": "chat_messages_thread_id_chat_threads_id_fk", + "tableFrom": "chat_messages", + "tableTo": "chat_threads", + "columnsFrom": [ + "thread_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.chat_threads": { + "name": "chat_threads", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "workspace_id": { + "name": "workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "title": { + "name": "title", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "demo_key": { + "name": "demo_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + } + }, + "indexes": { + "chat_threads_workspace_updated_idx": { + "name": "chat_threads_workspace_updated_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "updated_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "where": "deleted_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + }, + "chat_threads_workspace_demo_key_idx": { + "name": "chat_threads_workspace_demo_key_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "demo_key", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "chat_threads_workspace_id_workspaces_id_fk": { + "name": "chat_threads_workspace_id_workspaces_id_fk", + "tableFrom": "chat_threads", + "tableTo": "workspaces", + "columnsFrom": [ + "workspace_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.demo_source_visibilities": { + "name": "demo_source_visibilities", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "workspace_id": { + "name": "workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "demo_source_id": { + "name": "demo_source_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "hidden_at": { + "name": "hidden_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "demo_source_visibilities_workspace_source_idx": { + "name": "demo_source_visibilities_workspace_source_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "demo_source_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "demo_source_visibilities_workspace_idx": { + "name": "demo_source_visibilities_workspace_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "demo_source_visibilities_workspace_id_workspaces_id_fk": { + "name": "demo_source_visibilities_workspace_id_workspaces_id_fk", + "tableFrom": "demo_source_visibilities", + "tableTo": "workspaces", + "columnsFrom": [ + "workspace_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.source_parse_results": { + "name": "source_parse_results", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "source_id": { + "name": "source_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "result_blob_url": { + "name": "result_blob_url", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "snapshot_manifest_url": { + "name": "snapshot_manifest_url", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "snapshot_manifest_key": { + "name": "snapshot_manifest_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "revision_key": { + "name": "revision_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "sync_status": { + "name": "sync_status", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "sync_error": { + "name": "sync_error", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "asset_urls": { + "name": "asset_urls", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "source_parse_results_source_id_idx": { + "name": "source_parse_results_source_id_idx", + "columns": [ + { + "expression": "source_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "source_parse_results_source_id_sources_id_fk": { + "name": "source_parse_results_source_id_sources_id_fk", + "tableFrom": "source_parse_results", + "tableTo": "sources", + "columnsFrom": [ + "source_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": { + "source_parse_results_source_id_unique": { + "name": "source_parse_results_source_id_unique", + "nullsNotDistinct": false, + "columns": [ + "source_id" + ] + } + }, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.sources": { + "name": "sources", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "workspace_id": { + "name": "workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "title": { + "name": "title", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "mime_type": { + "name": "mime_type", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "size_bytes": { + "name": "size_bytes", + "type": "bigint", + "primaryKey": false, + "notNull": true + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "failure_reason": { + "name": "failure_reason", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "failure_stage": { + "name": "failure_stage", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "knowhere_job_id": { + "name": "knowhere_job_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "knowhere_document_id": { + "name": "knowhere_document_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "staged_blob_pathname": { + "name": "staged_blob_pathname", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "staged_blob_url": { + "name": "staged_blob_url", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "original_blob_pathname": { + "name": "original_blob_pathname", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "original_blob_url": { + "name": "original_blob_url", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "demo_key": { + "name": "demo_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + } + }, + "indexes": { + "sources_workspace_created_idx": { + "name": "sources_workspace_created_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "created_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "where": "deleted_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + }, + "sources_workspace_status_idx": { + "name": "sources_workspace_status_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "status", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "sources_workspace_demo_key_idx": { + "name": "sources_workspace_demo_key_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "demo_key", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "sources_workspace_document_idx": { + "name": "sources_workspace_document_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "knowhere_document_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "where": "knowhere_document_id IS NOT NULL AND deleted_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "sources_workspace_id_workspaces_id_fk": { + "name": "sources_workspace_id_workspaces_id_fk", + "tableFrom": "sources", + "tableTo": "workspaces", + "columnsFrom": [ + "workspace_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.workspaces": { + "name": "workspaces", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "namespace": { + "name": "namespace", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "workspaces_user_id_idx": { + "name": "workspaces_user_id_idx", + "columns": [ + { + "expression": "user_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": { + "workspaces_user_id_unique": { + "name": "workspaces_user_id_unique", + "nullsNotDistinct": false, + "columns": [ + "user_id" + ] + }, + "workspaces_namespace_unique": { + "name": "workspaces_namespace_unique", + "nullsNotDistinct": false, + "columns": [ + "namespace" + ] + } + }, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + } + }, + "enums": {}, + "schemas": {}, + "sequences": {}, + "roles": {}, + "policies": {}, + "views": {}, + "_meta": { + "columns": {}, + "schemas": {}, + "tables": {} + } +} \ No newline at end of file diff --git a/drizzle/meta/_journal.json b/drizzle/meta/_journal.json index 7174911..fc814c6 100644 --- a/drizzle/meta/_journal.json +++ b/drizzle/meta/_journal.json @@ -78,6 +78,13 @@ "when": 1783098230819, "tag": "0010_minor_piledriver", "breakpoints": true + }, + { + "idx": 11, + "version": "7", + "when": 1783175614191, + "tag": "0011_rare_deadpool", + "breakpoints": true } ] } \ No newline at end of file diff --git a/src/app/api/demo-sources/materialize/route.test.ts b/src/app/api/demo-sources/materialize/route.test.ts index 9b41725..0818dca 100644 --- a/src/app/api/demo-sources/materialize/route.test.ts +++ b/src/app/api/demo-sources/materialize/route.test.ts @@ -232,6 +232,7 @@ function makeSource( sizeBytes: 5648867, status: "ready", failureReason: null, + failureStage: null, knowhereJobId: null, knowhereDocumentId: "doc_user_copy", stagedBlobPathname: null, diff --git a/src/app/api/sources/[sourceId]/chunks/route.test.ts b/src/app/api/sources/[sourceId]/chunks/route.test.ts index 911646a..85fb1bd 100644 --- a/src/app/api/sources/[sourceId]/chunks/route.test.ts +++ b/src/app/api/sources/[sourceId]/chunks/route.test.ts @@ -2,22 +2,17 @@ import { NextRequest } from "next/server" import { beforeEach, describe, expect, it, vi } from "vitest" const mocks = vi.hoisted(() => ({ - blobGet: vi.fn(), - blobPut: vi.fn(), deleteBlob: vi.fn(), ensureApiKeyForWorkspace: vi.fn(), ensureWorkspace: vi.fn(), fetchDemoChunkPage: vi.fn(), findSourceInWorkspace: vi.fn(), - ensureParsedSnapshotForRead: vi.fn(), getCurrentUser: vi.fn(), - getSourceParseAssetUrls: vi.fn(), - getSourceParseSnapshotMetadata: vi.fn(), localizeRemoteDocument: vi.fn(), makeKnowhereClient: vi.fn(), + makeKnowhereClientWithParsedStorage: vi.fn(), + readChunks: vi.fn(), requireUser: vi.fn(), - syncRemoteParsedSnapshot: vi.fn(), - updateSourceRevisionKey: vi.fn(), })) vi.mock("next/headers", () => ({ @@ -42,23 +37,18 @@ vi.mock("@/infrastructure/auth", () => ({ vi.mock("@/integrations/knowhere", () => ({ makeKnowhereClient: mocks.makeKnowhereClient, + makeKnowhereClientWithParsedStorage: + mocks.makeKnowhereClientWithParsedStorage, })) vi.mock("@vercel/blob", () => ({ del: mocks.deleteBlob, - get: mocks.blobGet, - put: mocks.blobPut, })) vi.mock("@/domains/sources/service", () => ({ sourceService: { - ensureParsedSnapshotForRead: mocks.ensureParsedSnapshotForRead, findInWorkspace: mocks.findSourceInWorkspace, - getParseAssetUrls: mocks.getSourceParseAssetUrls, - getParseSnapshotMetadata: mocks.getSourceParseSnapshotMetadata, localizeRemoteDocument: mocks.localizeRemoteDocument, - syncRemoteParsedSnapshot: mocks.syncRemoteParsedSnapshot, - updateSourceRevisionKey: mocks.updateSourceRevisionKey, }, })) @@ -73,38 +63,9 @@ import { GET } from "./route" describe("GET /api/sources/[sourceId]/chunks", () => { beforeEach(() => { vi.clearAllMocks() - mocks.blobGet.mockResolvedValue(null) - mocks.blobPut.mockImplementation(async (pathname: string) => ({ - url: `https://blob.example/${pathname}`, - })) - mocks.getSourceParseAssetUrls.mockResolvedValue({}) - mocks.getSourceParseSnapshotMetadata.mockResolvedValue({ - resultBlobUrl: - "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - snapshotManifestUrl: - "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - assetUrlsByFilePath: {}, - }) - mocks.ensureParsedSnapshotForRead.mockResolvedValue({ - resultBlobUrl: - "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - snapshotManifestUrl: - "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - assetUrlsByFilePath: {}, - }) - mocks.updateSourceRevisionKey.mockResolvedValue(null) - mocks.syncRemoteParsedSnapshot.mockResolvedValue({ - resultBlobUrl: - "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - snapshotManifestUrl: - "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - assetUrlsByFilePath: {}, + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: { documents: { listChunks: vi.fn() } }, + knowledge: { readChunks: mocks.readChunks }, }) }) @@ -166,7 +127,6 @@ describe("GET /api/sources/[sourceId]/chunks", () => { }) expect(mocks.ensureApiKeyForWorkspace).not.toHaveBeenCalled() expect(mocks.makeKnowhereClient).not.toHaveBeenCalled() - expect(mocks.getSourceParseAssetUrls).not.toHaveBeenCalled() }) it("loads every API-owned demo chunk page for full anonymous chunk requests", async () => { @@ -321,7 +281,6 @@ describe("GET /api/sources/[sourceId]/chunks", () => { expect(mocks.findSourceInWorkspace).not.toHaveBeenCalled() expect(mocks.ensureApiKeyForWorkspace).not.toHaveBeenCalled() expect(mocks.makeKnowhereClient).not.toHaveBeenCalled() - expect(mocks.getSourceParseAssetUrls).not.toHaveBeenCalled() }) it("serves demo chunks for authenticated materialized demo sources", async () => { @@ -411,7 +370,6 @@ describe("GET /api/sources/[sourceId]/chunks", () => { }) expect(mocks.ensureApiKeyForWorkspace).not.toHaveBeenCalled() expect(mocks.makeKnowhereClient).not.toHaveBeenCalled() - expect(mocks.getSourceParseAssetUrls).not.toHaveBeenCalled() }) it("logs the demo chunk load failure before returning 404", async () => { @@ -463,12 +421,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { } }) - it("loads authenticated workspace chunks from the parsed Blob snapshot", async () => { - const knowhereClient = { - documents: { - listChunks: vi.fn(), - }, - } + it("loads authenticated workspace chunks through the SDK durable read", async () => { mocks.getCurrentUser.mockResolvedValue({ id: "user_1", email: null, @@ -480,128 +433,55 @@ describe("GET /api/sources/[sourceId]/chunks", () => { namespace: "notebook-workspace_1", createdAt: new Date("2026-05-10T00:00:00.000Z"), }) - mocks.findSourceInWorkspace.mockResolvedValue({ - id: "00000000-0000-0000-0000-000000000002", - workspaceId: "workspace_1", - title: "notes.pdf", - mimeType: "application/pdf", - sizeBytes: 1024, - status: "ready", - failureReason: null, - knowhereJobId: "job_1", - knowhereDocumentId: "doc_1", - stagedBlobPathname: null, - stagedBlobUrl: null, - originalBlobPathname: null, - originalBlobUrl: null, - demoKey: null, - createdAt: new Date("2026-05-10T00:00:00.000Z"), - updatedAt: new Date("2026-05-10T00:00:00.000Z"), - deletedAt: null, - }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + }), + ) mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") - mocks.makeKnowhereClient.mockReturnValue(knowhereClient) - mocks.getSourceParseAssetUrls.mockResolvedValue({ - "images/chart.png": "https://blob.example/images/chart.png", - }) - mocks.getSourceParseSnapshotMetadata.mockResolvedValue({ - resultBlobUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000002/parsed-result/manifest/current.json", - snapshotManifestUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000002/parsed-result/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000002/parsed-result/manifest/current.json", - assetUrlsByFilePath: { - "images/chart.png": "https://blob.example/images/chart.png", - }, - }) - mocks.ensureParsedSnapshotForRead.mockResolvedValue({ - resultBlobUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000002/parsed-result/manifest/current.json", - snapshotManifestUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000002/parsed-result/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000002/parsed-result/manifest/current.json", - assetUrlsByFilePath: { - "images/chart.png": "https://blob.example/images/chart.png", - }, - }) - mocks.blobGet.mockImplementation(async (pathname: string) => { - if (pathname.endsWith("/manifest/current.json")) { - return { - statusCode: 200, - stream: createTextStream( - JSON.stringify({ - version: 1, - kind: "knowhere-parsed-result-snapshot", - jobId: "job_1", - documentId: "doc_1", - sourceFileName: "notes.pdf", - totalChunks: 1, - chunkPageSize: 1, - chunkPages: [ - { - page: 1, - pageSize: 1, - chunkCount: 1, - key: "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000002/parsed-result/chunks/page-1.json", - }, - ], - assetUrlsByFilePath: { - "images/chart.png": "https://blob.example/images/chart.png", - }, - createdAt: "2026-07-03T00:00:00.000Z", - }), - ), - } - } - if (pathname.endsWith("/chunks/page-1.json")) { - return { - statusCode: 200, - stream: createTextStream( - JSON.stringify({ - version: 1, - jobId: "job_1", - documentId: "doc_1", - sourceFileName: "notes.pdf", - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - chunks: [ - { - id: "dchk_1", - chunkId: "parser_1", - chunkType: "image", - content: "Workspace chunk", - sectionPath: "Summary", - sourceChunkPath: "Default_Root/notes.pdf/Summary", - filePath: "images/chart.png", - metadata: {}, - sortOrder: 1, - }, - ], - }), - ), - } - } - return null + mocks.readChunks.mockResolvedValue({ + document: { localDocumentId: "doc_1" }, + chunks: [ + { + position: 1, + chunkId: "parser_1", + chunkType: "image", + content: "Workspace chunk", + readableContent: "Workspace chunk", + sectionPath: "Summary", + sourceChunkPath: "Summary", + filePath: "images/chart.png", + assetUrl: + "https://fake.public.blob.vercel-storage.com/images/chart.png", + metadata: {}, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, }) const response = await GET( new NextRequest( "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/chunks?page=1&pageSize=1", ), - { params: Promise.resolve({ sourceId: "00000000-0000-0000-0000-000000000002" }) }, + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, ) await expect(response.json()).resolves.toMatchObject({ chunks: [ { - chunkId: "dchk_1", parserChunkId: "parser_1", documentId: "doc_1", - assetUrl: "https://blob.example/images/chart.png", + assetUrl: + "https://fake.public.blob.vercel-storage.com/images/chart.png", sourceTitle: "notes.pdf", }, ], @@ -613,15 +493,16 @@ describe("GET /api/sources/[sourceId]/chunks", () => { }) expect(response.status).toBe(200) expect(mocks.fetchDemoChunkPage).not.toHaveBeenCalled() - expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled() + expect(mocks.readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "job_1", + page: 1, + pageSize: 1, + assetUrlPolicy: "durable", + }) }) - it("returns processing when a ready workspace source has no complete parsed snapshot", async () => { - const knowhereClient = { - documents: { - listChunks: vi.fn(), - }, - } + it("returns processing when a workspace source is not ready", async () => { mocks.getCurrentUser.mockResolvedValue({ id: "user_1", email: null, @@ -633,35 +514,24 @@ describe("GET /api/sources/[sourceId]/chunks", () => { namespace: "notebook-workspace_1", createdAt: new Date("2026-05-10T00:00:00.000Z"), }) - mocks.findSourceInWorkspace.mockResolvedValue({ - id: "00000000-0000-0000-0000-000000000002", - workspaceId: "workspace_1", - title: "notes.pdf", - mimeType: "application/pdf", - sizeBytes: 1024, - status: "ready", - failureReason: null, - knowhereJobId: "job_1", - knowhereDocumentId: "doc_1", - stagedBlobPathname: null, - stagedBlobUrl: null, - originalBlobPathname: null, - originalBlobUrl: null, - demoKey: null, - createdAt: new Date("2026-05-10T00:00:00.000Z"), - updatedAt: new Date("2026-05-10T00:00:00.000Z"), - deletedAt: null, - }) - mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") - mocks.makeKnowhereClient.mockReturnValue(knowhereClient) - mocks.getSourceParseSnapshotMetadata.mockResolvedValue(null) - mocks.ensureParsedSnapshotForRead.mockResolvedValue(null) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + status: "parsing", + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + }), + ) const response = await GET( new NextRequest( "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/chunks?page=1&pageSize=1", ), - { params: Promise.resolve({ sourceId: "00000000-0000-0000-0000-000000000002" }) }, + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, ) await expect(response.json()).resolves.toEqual({ @@ -672,13 +542,13 @@ describe("GET /api/sources/[sourceId]/chunks", () => { total: 0, totalPages: 0, }, - message: "Source parsed snapshot is still being prepared.", + message: "Source is still being prepared.", }) expect(response.status).toBe(202) - expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled() + expect(mocks.readChunks).not.toHaveBeenCalled() }) - it("materializes a remote source id and loads chunks from the parsed Blob snapshot", async () => { + it("materializes a remote source id and reads chunks through the SDK", async () => { const knowhereClient = { documents: { list: vi.fn(async () => ({ @@ -697,9 +567,6 @@ describe("GET /api/sources/[sourceId]/chunks", () => { })), listChunks: vi.fn(), }, - knowledge: { - loadJobResult: vi.fn(), - }, } mocks.getCurrentUser.mockResolvedValue({ id: "user_1", @@ -715,97 +582,35 @@ describe("GET /api/sources/[sourceId]/chunks", () => { mocks.fetchDemoChunkPage.mockRejectedValue(new Error("not a demo")) mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") mocks.makeKnowhereClient.mockReturnValue(knowhereClient) - const localizedSource = { - id: "00000000-0000-0000-0000-000000000009", - workspaceId: "workspace_1", - title: "remote.pdf", - mimeType: "application/pdf", - sizeBytes: 0, - status: "ready", - failureReason: null, - knowhereJobId: "job_result_1", - knowhereDocumentId: "doc_remote", - stagedBlobPathname: null, - stagedBlobUrl: null, - originalBlobPathname: null, - originalBlobUrl: null, - demoKey: null, - createdAt: new Date("2026-05-10T00:00:00.000Z"), - updatedAt: new Date("2026-05-10T00:00:00.000Z"), - deletedAt: null, - } - mocks.localizeRemoteDocument.mockResolvedValue(localizedSource) - mocks.syncRemoteParsedSnapshot.mockResolvedValue({ - resultBlobUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - snapshotManifestUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - assetUrlsByFilePath: { - "pages/page-1.png": "https://blob.example/pages/page-1.png", - }, - }) - mocks.blobGet.mockImplementation(async (pathname: string) => { - if (pathname.endsWith("/manifest/current.json")) { - return { - statusCode: 200, - stream: createTextStream( - JSON.stringify({ - version: 1, - kind: "knowhere-parsed-result-snapshot", - jobId: "job_result_1", - documentId: "doc_remote", - sourceFileName: "remote.pdf", - totalChunks: 1, - chunkPageSize: 1, - chunkPages: [ - { - page: 1, - pageSize: 1, - chunkCount: 1, - key: "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/chunks/page-1.json", - }, - ], - assetUrlsByFilePath: { - "pages/page-1.png": "https://blob.example/pages/page-1.png", - }, - createdAt: "2026-07-03T00:00:00.000Z", - }), - ), - } - } - if (pathname.endsWith("/chunks/page-1.json")) { - return { - statusCode: 200, - stream: createTextStream( - JSON.stringify({ - version: 1, - jobId: "job_result_1", - documentId: "doc_remote", - sourceFileName: "remote.pdf", - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - chunks: [ - { - id: "dchk_remote", - chunkId: "parser_remote", - chunkType: "page", - content: "Remote chunk", - sectionPath: "Summary", - sourceChunkPath: "Default_Root/remote.pdf/Summary", - filePath: "pages/page-1.png", - metadata: {}, - sortOrder: 0, - }, - ], - }), - ), - } - } - return null + mocks.localizeRemoteDocument.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000009", + title: "remote.pdf", + knowhereJobId: "job_result_1", + knowhereDocumentId: "doc_remote", + }), + ) + mocks.readChunks.mockResolvedValue({ + document: { localDocumentId: "doc_remote" }, + chunks: [ + { + position: 1, + chunkId: "parser_remote", + chunkType: "page", + content: "Remote chunk", + readableContent: "Remote chunk", + sectionPath: "Summary", + sourceChunkPath: "Summary", + filePath: "pages/page-1.png", + assetUrl: + "https://fake.public.blob.vercel-storage.com/pages/page-1.png", + metadata: {}, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, }) const response = await GET( @@ -822,10 +627,10 @@ describe("GET /api/sources/[sourceId]/chunks", () => { await expect(response.json()).resolves.toMatchObject({ chunks: [ { - chunkId: "dchk_remote", parserChunkId: "parser_remote", documentId: "doc_remote", - assetUrl: "https://blob.example/pages/page-1.png", + assetUrl: + "https://fake.public.blob.vercel-storage.com/pages/page-1.png", sourceTitle: "remote.pdf", }, ], @@ -841,270 +646,45 @@ describe("GET /api/sources/[sourceId]/chunks", () => { "workspace_1", "session=abc", ) - expect(mocks.localizeRemoteDocument).toHaveBeenCalledWith( - "workspace_1", - { - documentId: "doc_remote", - namespace: "default", - status: "ready", - title: "remote.pdf", - mimeType: "application/pdf", - sizeBytes: undefined, - revisionKey: "job_result_1", - }, - ) - expect(mocks.syncRemoteParsedSnapshot).toHaveBeenCalledWith({ - workspaceId: "workspace_1", - source: localizedSource, - client: { - documents: knowhereClient.documents, - knowledge: knowhereClient.knowledge, - }, - }) - expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled() - expect(knowhereClient.knowledge.loadJobResult).not.toHaveBeenCalled() - }) - - it("returns processing for a remote source when the parsed Blob snapshot is incomplete", async () => { - const knowhereClient = { - documents: { - list: vi.fn(async () => ({ - documents: [ - { - documentId: "doc_remote", - namespace: "default", - status: "active", - currentJobResultId: "job_result_1", - sourceFileName: "remote.pdf", - documentMetadata: { - mimeType: "application/pdf", - }, - }, - ], - })), - listChunks: vi.fn(), - }, - knowledge: { - loadJobResult: vi.fn(), - }, - } - mocks.getCurrentUser.mockResolvedValue({ - id: "user_1", - email: null, - name: null, - }) - mocks.ensureWorkspace.mockResolvedValue({ - id: "workspace_1", - userId: "user_1", - namespace: "notebook-workspace_1", - createdAt: new Date("2026-05-10T00:00:00.000Z"), - }) - mocks.fetchDemoChunkPage.mockRejectedValue(new Error("not a demo")) - mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") - mocks.makeKnowhereClient.mockReturnValue(knowhereClient) - mocks.localizeRemoteDocument.mockResolvedValue({ - id: "00000000-0000-0000-0000-000000000009", - workspaceId: "workspace_1", - title: "remote.pdf", - mimeType: "application/pdf", - sizeBytes: 0, + expect(mocks.localizeRemoteDocument).toHaveBeenCalledWith("workspace_1", { + documentId: "doc_remote", + namespace: "default", status: "ready", - failureReason: null, - knowhereJobId: "job_result_1", - knowhereDocumentId: "doc_remote", - stagedBlobPathname: null, - stagedBlobUrl: null, - originalBlobPathname: null, - originalBlobUrl: null, - demoKey: null, - createdAt: new Date("2026-05-10T00:00:00.000Z"), - updatedAt: new Date("2026-05-10T00:00:00.000Z"), - deletedAt: null, - }) - mocks.syncRemoteParsedSnapshot.mockResolvedValue(null) - - const response = await GET( - new NextRequest( - "http://localhost:3001/api/sources/knowhere-doc:default:doc_remote/chunks?page=1&pageSize=1", - ), - { - params: Promise.resolve({ - sourceId: "knowhere-doc:default:doc_remote", - }), - }, - ) - - await expect(response.json()).resolves.toEqual({ - chunks: [], - pagination: { - page: 1, - pageSize: 1, - total: 0, - totalPages: 0, - }, - message: "Source parsed snapshot is still being prepared.", - }) - expect(response.status).toBe(202) - expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled() - }) - - it("loads remote snapshot chunks after an interrupted sync left the source parsing", async () => { - const knowhereClient = { - documents: { - list: vi.fn(async () => ({ - documents: [ - { - documentId: "doc_remote", - namespace: "default", - status: "active", - currentJobResultId: "job_result_1", - sourceFileName: "remote.pdf", - documentMetadata: { - mimeType: "application/pdf", - }, - }, - ], - })), - listChunks: vi.fn(), - }, - knowledge: { - loadJobResult: vi.fn(), - }, - } - mocks.getCurrentUser.mockResolvedValue({ - id: "user_1", - email: null, - name: null, - }) - mocks.ensureWorkspace.mockResolvedValue({ - id: "workspace_1", - userId: "user_1", - namespace: "notebook-workspace_1", - createdAt: new Date("2026-05-10T00:00:00.000Z"), - }) - mocks.fetchDemoChunkPage.mockRejectedValue(new Error("not a demo")) - mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") - mocks.makeKnowhereClient.mockReturnValue(knowhereClient) - mocks.localizeRemoteDocument.mockResolvedValue({ - id: "00000000-0000-0000-0000-000000000009", - workspaceId: "workspace_1", title: "remote.pdf", mimeType: "application/pdf", - sizeBytes: 0, - status: "parsing", - failureReason: null, - knowhereJobId: "job_result_1", - knowhereDocumentId: "doc_remote", - stagedBlobPathname: null, - stagedBlobUrl: null, - originalBlobPathname: null, - originalBlobUrl: null, - demoKey: null, - createdAt: new Date("2026-05-10T00:00:00.000Z"), - updatedAt: new Date("2026-05-10T00:00:00.000Z"), - deletedAt: null, + sizeBytes: undefined, + revisionKey: "job_result_1", }) - mocks.syncRemoteParsedSnapshot.mockResolvedValue({ - resultBlobUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - snapshotManifestUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - assetUrlsByFilePath: {}, - }) - mocks.blobGet.mockImplementation(async (pathname: string) => { - if (pathname.endsWith("/manifest/current.json")) { - return { - statusCode: 200, - stream: createTextStream( - JSON.stringify({ - version: 1, - kind: "knowhere-parsed-result-snapshot", - jobId: "job_result_1", - documentId: "doc_remote", - sourceFileName: "remote.pdf", - totalChunks: 1, - chunkPageSize: 1, - chunkPages: [ - { - page: 1, - pageSize: 1, - chunkCount: 1, - key: "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/chunks/page-1.json", - }, - ], - assetUrlsByFilePath: {}, - createdAt: "2026-07-03T00:00:00.000Z", - }), - ), - } - } - if (pathname.endsWith("/chunks/page-1.json")) { - return { - statusCode: 200, - stream: createTextStream( - JSON.stringify({ - version: 1, - jobId: "job_result_1", - documentId: "doc_remote", - sourceFileName: "remote.pdf", - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - chunks: [ - { - id: "dchk_remote", - chunkId: "parser_remote", - chunkType: "text", - content: "Remote chunk", - sectionPath: "Summary", - sourceChunkPath: "Default_Root/remote.pdf/Summary", - metadata: {}, - sortOrder: 0, - }, - ], - }), - ), - } - } - return null - }) - - const response = await GET( - new NextRequest( - "http://localhost:3001/api/sources/knowhere-doc:default:doc_remote/chunks?page=1&pageSize=1", - ), - { - params: Promise.resolve({ - sourceId: "knowhere-doc:default:doc_remote", - }), - }, - ) - - await expect(response.json()).resolves.toMatchObject({ - chunks: [ - { - chunkId: "dchk_remote", - parserChunkId: "parser_remote", - documentId: "doc_remote", - sourceTitle: "remote.pdf", - }, - ], - pagination: { - page: 1, - pageSize: 1, - total: 1, - }, + expect(mocks.readChunks).toHaveBeenCalledWith({ + documentId: "doc_remote", + revisionKey: "job_result_1", + page: 1, + pageSize: 1, + assetUrlPolicy: "durable", }) - expect(response.status).toBe(200) - expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled() }) }) -function createTextStream(text: string): ReadableStream { - const stream = new Response(text).body - if (!stream) throw new Error("Response body stream was not created.") - return stream +function makeReadySource(overrides: Record) { + return { + id: "00000000-0000-0000-0000-000000000002", + workspaceId: "workspace_1", + title: "notes.pdf", + mimeType: "application/pdf", + sizeBytes: 1024, + status: "ready", + failureReason: null, + failureStage: null, + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: null, + originalBlobUrl: null, + demoKey: null, + createdAt: new Date("2026-05-10T00:00:00.000Z"), + updatedAt: new Date("2026-05-10T00:00:00.000Z"), + deletedAt: null, + ...overrides, + } } diff --git a/src/app/api/sources/parsed-sync/route.ts b/src/app/api/sources/parsed-sync/route.ts new file mode 100644 index 0000000..5272176 --- /dev/null +++ b/src/app/api/sources/parsed-sync/route.ts @@ -0,0 +1,27 @@ +import { serve } from "@upstash/workflow/nextjs" + +import { parsedSyncRouteWorkflow } from "@/domains/sources/parsed-sync-route-workflow" + +type ParsedSyncPayload = Parameters< + typeof parsedSyncRouteWorkflow.normalizeParsedSyncPayload +>[0] + +export const { POST } = serve( + async (context) => { + const payload = parsedSyncRouteWorkflow.normalizeParsedSyncPayload( + context.requestPayload, + ) + await parsedSyncRouteWorkflow.runParsedSyncWorkflow({ + context, + payload, + }) + }, + { + failureFunction: async ({ context, failResponse }) => { + await parsedSyncRouteWorkflow.markSyncFailedAfterWorkflowFailure( + context.requestPayload, + failResponse, + ) + }, + }, +) diff --git a/src/app/api/sources/route.test.ts b/src/app/api/sources/route.test.ts index a78639b..a7b10c6 100644 --- a/src/app/api/sources/route.test.ts +++ b/src/app/api/sources/route.test.ts @@ -67,6 +67,7 @@ const source: Source = { sizeBytes: 5, status: "parsing", failureReason: null, + failureStage: null, knowhereJobId: "job_1", knowhereDocumentId: null, stagedBlobPathname: null, diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 71604f6..1b51180 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -1993,6 +1993,7 @@ function makeSource(overrides: Partial = {}): Source { sizeBytes: 100, status: "ready", failureReason: null, + failureStage: null, knowhereJobId: "job_123", knowhereDocumentId: "doc_included", stagedBlobPathname: null, diff --git a/src/domains/chat/media-asset-hardening.test.ts b/src/domains/chat/media-asset-hardening.test.ts index fe1f3df..d3c87b2 100644 --- a/src/domains/chat/media-asset-hardening.test.ts +++ b/src/domains/chat/media-asset-hardening.test.ts @@ -3,8 +3,7 @@ import { afterEach, describe, expect, it, vi } from "vitest" import type { Source } from "@/infrastructure/db/schema" import { hardenChatMediaAssetUrls, - type ChatMediaAssetBlobStore, - type FetchChatMediaAsset, + isNotebookOwnedAssetUrl, type HardenableRetrievalResult, } from "./media-asset-hardening" @@ -20,116 +19,42 @@ vi.mock("@/lib/logger", () => ({ afterEach(() => { vi.clearAllMocks() - delete process.env.KNOWHERE_BASE_URL }) describe("hardenChatMediaAssetUrls", () => { - it("copies upstream absolute asset URLs into Notebook chat assets", async () => { - const rawAssetUrl = - "https://knowhere-storage.example/results/job_1/images/image-6-%E6%83%85%E6%84%9F%E5%88%86%E7%B1%BB%E6%A8%A1%E5%9E%8B.jpg?AWSAccessKeyId=test&Signature=secret" - const blobStore = makeBlobStore( - "https://blob.example/workspaces/workspace_1/chat-assets/source-source_1/image-6.jpg", - ) - const fetchAsset = makeFetchAsset("image-bytes", "image/jpeg") + it("keeps an already Notebook-owned asset URL without loading the asset map", async () => { + const ownedUrl = + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_1/rev_1/assets/images/a.png" + const loadSourceAssetUrls = vi.fn(async () => ({})) const result = await hardenChatMediaAssetUrls({ workspaceId: "workspace_1", - sources: [ - makeSource({ - id: "source_1", - knowhereDocumentId: "doc_model", - }), - ], + sources: [makeSource({ id: "source_1", knowhereDocumentId: "doc_1" })], results: [ makeRetrievalResult({ chunkType: "image", - assetUrl: rawAssetUrl, + assetUrl: ownedUrl, source: { - documentId: "doc_model", + documentId: "doc_1", sourceFileName: "model.pdf", sectionPath: "Root", }, }), ], - blobStore, - fetchAsset, - }) - - expect(fetchAsset).toHaveBeenCalledWith(rawAssetUrl) - expect(blobStore.put).toHaveBeenCalledWith( - expect.stringMatching( - /^workspaces\/workspace_1\/chat-assets\/source-source_1\/[a-f0-9]{24}-image-6\.jpg$/, - ), - expect.any(Buffer), - { - access: "public", - allowOverwrite: true, - contentType: "image/jpeg", - multipart: true, - }, - ) - expect(result.results[0]?.assetUrl).toBe( - "https://blob.example/workspaces/workspace_1/chat-assets/source-source_1/image-6.jpg", - ) - }) - - it("copies server page citation asset URLs into Notebook chat assets", async () => { - const rawPageAssetUrl = - "https://knowhere-storage.example/results/job_1/page_citation_assets/page-4.png?AWSAccessKeyId=test&Signature=secret" - const blobPageAssetUrl = - "https://blob.example/workspaces/workspace_1/chat-assets/source-source_1/page-4.png" - const blobStore = makeBlobStore(blobPageAssetUrl) - const fetchAsset = makeFetchAsset("page-image-bytes", "image/png") - - const result = await hardenChatMediaAssetUrls({ - workspaceId: "workspace_1", - sources: [ - makeSource({ - id: "source_1", - knowhereDocumentId: "doc_pages", - }), - ], - results: [ - makeRetrievalResult({ - chunkType: "page", - pageCitationAssetUrl: rawPageAssetUrl, - source: { - documentId: "doc_pages", - sourceFileName: "deck.pdf", - sectionPath: "Page 4", - }, - }), - ], - blobStore, - fetchAsset, + loadSourceAssetUrls, }) - expect(fetchAsset).toHaveBeenCalledWith(rawPageAssetUrl) - expect(blobStore.put).toHaveBeenCalledWith( - expect.stringMatching( - /^workspaces\/workspace_1\/chat-assets\/source-source_1\/[a-f0-9]{24}-page-4\.png$/, - ), - expect.any(Buffer), - { - access: "public", - allowOverwrite: true, - contentType: "image/png", - multipart: true, - }, - ) - expect(result.results[0]?.pageCitationAssetUrl).toBe(blobPageAssetUrl) - expect(result.results[0]?.assetUrl).toBeUndefined() + expect(result.results[0]?.assetUrl).toBe(ownedUrl) + expect(loadSourceAssetUrls).not.toHaveBeenCalled() }) - it("uses an existing parsed asset URL before fetching the upstream URL", async () => { + it("resolves a raw asset URL to the durable parsed asset URL", async () => { const rawAssetUrl = "https://knowhere-storage.example/results/job_1/images/id-front.jpg?AWSAccessKeyId=test" - const parsedAssetUrl = - "https://blob.example/workspaces/workspace_1/sources/source_identity/parsed-result/images/id-front.jpg" - const blobStore = makeBlobStore("https://blob.example/should-not-upload.jpg") - const fetchAsset = makeFetchAsset("should-not-fetch", "image/jpeg") + const durableUrl = + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/assets/images/id-front.jpg" const loadSourceAssetUrls = vi.fn().mockResolvedValue({ - "images/id-front.jpg": parsedAssetUrl, + "images/id-front.jpg": durableUrl, }) const result = await hardenChatMediaAssetUrls({ @@ -152,99 +77,49 @@ describe("hardenChatMediaAssetUrls", () => { }), ], loadSourceAssetUrls, - blobStore, - fetchAsset, }) expect(loadSourceAssetUrls).toHaveBeenCalledWith( expect.objectContaining({ id: "source_identity" }), ) - expect(fetchAsset).not.toHaveBeenCalled() - expect(blobStore.put).not.toHaveBeenCalled() - expect(result.results[0]?.assetUrl).toBe(parsedAssetUrl) + expect(result.results[0]?.assetUrl).toBe(durableUrl) }) - it("fetches demo asset routes from the upstream demo API", async () => { - process.env.KNOWHERE_BASE_URL = "https://knowhere.example" - const demoAssetUrl = - "/api/demo-sources/demo_source_1/assets/images/demo%20chart.png" - const blobStore = makeBlobStore( - "https://blob.example/workspaces/workspace_1/chat-assets/demo-demo_source_1/demo-chart.png", - ) - const fetchAsset = makeFetchAsset("demo-image", "image/png") - - const result = await hardenChatMediaAssetUrls({ - workspaceId: "workspace_1", - sources: [], - results: [ - makeRetrievalResult({ - chunkType: "image", - assetUrl: demoAssetUrl, - source: { - documentId: "demo_doc", - sourceFileName: "demo.pdf", - sectionPath: "images/demo chart.png", - }, - }), - ], - blobStore, - fetchAsset, - }) - - expect(fetchAsset).toHaveBeenCalledWith( - "https://knowhere.example/api/v1/demo/sources/demo_source_1/assets/images/demo%20chart.png", - ) - expect(fetchAsset).not.toHaveBeenCalledWith(demoAssetUrl) - expect(blobStore.put).toHaveBeenCalledWith( - expect.stringContaining("/chat-assets/demo-demo_source_1/"), - expect.any(Buffer), - expect.objectContaining({ contentType: "image/png" }), - ) - expect(result.results[0]?.assetUrl).toBe( - "https://blob.example/workspaces/workspace_1/chat-assets/demo-demo_source_1/demo-chart.png", - ) - }) - - it("falls back to the raw URL when hardening fails", async () => { + it("omits an asset URL that cannot be resolved to a durable URL", async () => { const rawAssetUrl = "https://knowhere-storage.example/results/job_1/tables/table-1.html?AWSAccessKeyId=test" - const blobStore = makeBlobStore("https://blob.example/should-not-exist.html") - const fetchAsset: FetchChatMediaAsset = vi - .fn() - .mockRejectedValue(new Error("expired URL")) + const loadSourceAssetUrls = vi.fn().mockResolvedValue({}) const result = await hardenChatMediaAssetUrls({ workspaceId: "workspace_1", - sources: [], + sources: [ + makeSource({ id: "source_1", knowhereDocumentId: "doc_1" }), + ], results: [ makeRetrievalResult({ chunkType: "table", assetUrl: rawAssetUrl, + source: { + documentId: "doc_1", + sourceFileName: "source.pdf", + sectionPath: "tables/table-1.html", + }, }), ], - blobStore, - fetchAsset, + loadSourceAssetUrls, }) - expect(result.results[0]?.assetUrl).toBe(rawAssetUrl) - expect(blobStore.put).not.toHaveBeenCalled() - expect(loggerMock.warn).toHaveBeenCalledWith( - "chat-agent: media asset hardening failed; keeping raw URL", - expect.objectContaining({ - assetUrl: - "https://knowhere-storage.example/results/job_1/tables/table-1.html", - error: "expired URL", - }), - ) + expect(result.results[0]?.assetUrl).toBeUndefined() }) - it("rewrites artifact asset URLs and nested citation asset URLs", async () => { + it("resolves an artifact asset URL and its nested citation URL", async () => { const rawAssetUrl = "https://knowhere-storage.example/results/job_1/images/front.jpg?AWSAccessKeyId=test" - const blobAssetUrl = - "https://blob.example/workspaces/workspace_1/chat-assets/source-source_identity/front.jpg" - const blobStore = makeBlobStore(blobAssetUrl) - const fetchAsset = makeFetchAsset("front-image", "image/jpeg") + const durableUrl = + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/assets/images/front.jpg" + const loadSourceAssetUrls = vi.fn().mockResolvedValue({ + "images/front.jpg": durableUrl, + }) const result = await hardenChatMediaAssetUrls({ workspaceId: "workspace_1", @@ -273,82 +148,35 @@ describe("hardenChatMediaAssetUrls", () => { }, }, ], - blobStore, - fetchAsset, + loadSourceAssetUrls, }) const [artifact] = result.artifacts ?? [] - expect(fetchAsset).toHaveBeenCalledTimes(1) - expect(artifact?.assetUrl).toBe(blobAssetUrl) - expect(artifact?.citation?.assetUrl).toBe(blobAssetUrl) + expect(artifact?.assetUrl).toBe(durableUrl) + expect(artifact?.citation?.assetUrl).toBe(durableUrl) }) +}) - it("rewrites nested citation page citation asset URLs", async () => { - const rawPageAssetUrl = - "https://knowhere-storage.example/results/job_1/page_citation_assets/page-8.png?AWSAccessKeyId=test" - const blobPageAssetUrl = - "https://blob.example/workspaces/workspace_1/chat-assets/source-source_pages/page-8.png" - const blobStore = makeBlobStore(blobPageAssetUrl) - const fetchAsset = makeFetchAsset("page-eight-image", "image/png") - - const result = await hardenChatMediaAssetUrls({ - workspaceId: "workspace_1", - sources: [ - makeSource({ - id: "source_pages", - knowhereDocumentId: "doc_pages", - }), - ], - results: [], - artifacts: [ - { - type: "image", - ref: "asset:r1:result:1", - assetUrl: - "https://blob.example/workspaces/workspace_1/chat-assets/source-source_pages/figure.png", - label: "figure", - citation: { - chunkType: "page", - score: 0.9, - pageCitationAssetUrl: rawPageAssetUrl, - source: { - documentId: "doc_pages", - sourceFileName: "deck.pdf", - sectionPath: "Page 8", - }, - }, - }, - ], - blobStore, - fetchAsset, - }) - - const [artifact] = result.artifacts ?? [] - expect(fetchAsset).toHaveBeenCalledWith(rawPageAssetUrl) - expect(artifact?.citation?.pageCitationAssetUrl).toBe(blobPageAssetUrl) +describe("isNotebookOwnedAssetUrl", () => { + it("treats Vercel Blob hosts and parsed-document paths as owned", () => { + expect( + isNotebookOwnedAssetUrl( + "https://fake.public.blob.vercel-storage.com/x/y.png", + ), + ).toBe(true) + expect( + isNotebookOwnedAssetUrl( + "https://cdn.example/workspaces/w/parsed-documents/d/r/assets/a.png", + ), + ).toBe(true) + expect( + isNotebookOwnedAssetUrl( + "https://knowhere-storage.example/results/job_1/images/a.png?sig=x", + ), + ).toBe(false) }) }) -function makeFetchAsset( - body: string, - contentType: string, -): FetchChatMediaAsset { - return vi.fn().mockResolvedValue( - new Response(Buffer.from(body), { - status: 200, - headers: { - "content-type": contentType, - }, - }), - ) -} - -function makeBlobStore(url: string): ChatMediaAssetBlobStore { - return { - put: vi.fn().mockResolvedValue({ url }), - } -} - function makeRetrievalResult( overrides: Partial = {}, ): HardenableRetrievalResult { @@ -374,6 +202,7 @@ function makeSource(overrides: Partial = {}): Source { sizeBytes: 100, status: "ready", failureReason: null, + failureStage: null, knowhereJobId: "job_123", knowhereDocumentId: "doc_1", stagedBlobPathname: null, diff --git a/src/domains/chat/media-asset-hardening.ts b/src/domains/chat/media-asset-hardening.ts index 4161587..f768515 100644 --- a/src/domains/chat/media-asset-hardening.ts +++ b/src/domains/chat/media-asset-hardening.ts @@ -1,6 +1,3 @@ -import path from "node:path" -import { createHash } from "node:crypto" -import { put } from "@vercel/blob" import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" import type { @@ -8,7 +5,6 @@ import type { ChatCitationView, } from "@/domains/chat/types" import type { Source } from "@/infrastructure/db/schema" -import { knowhereDemoApi } from "@/integrations/knowhere-demo" import { logger } from "@/lib/logger" import type { LoadSourceAssetUrls } from "./media-assets" import { resolveAssetUrlFromReferenceText } from "./media-assets" @@ -31,30 +27,11 @@ export type HardenMediaAssetUrls = ( input: HardenMediaAssetUrlsInput, ) => Promise -export type ChatMediaAssetBlobStore = { - readonly put: ( - pathname: string, - body: Buffer, - options: ChatMediaAssetBlobPutOptions, - ) => Promise<{ readonly url: string }> -} - -export type ChatMediaAssetBlobPutOptions = { - readonly access?: "public" - readonly allowOverwrite?: boolean - readonly contentType: string - readonly multipart?: boolean -} - -export type FetchChatMediaAsset = (url: string) => Promise - export type HardenChatMediaAssetUrlsForWorkspaceInput = HardenMediaAssetUrlsInput & { readonly workspaceId: string readonly sources: readonly Source[] - readonly loadSourceAssetUrls?: LoadSourceAssetUrls - readonly blobStore?: ChatMediaAssetBlobStore - readonly fetchAsset?: FetchChatMediaAsset + readonly loadSourceAssetUrls: LoadSourceAssetUrls } type AssetReferenceSource = ChatCitationView["source"] @@ -65,63 +42,38 @@ type AssetUrlReference = { readonly content?: string } -type AssetFetchRequest = { - readonly fetchUrl: string - readonly canonicalKey: string - readonly sourceSegment: string - readonly suggestedFileName: string -} - type HardeningContext = { - readonly workspaceId: string readonly sourcesByDocumentId: ReadonlyMap - readonly loadSourceAssetUrls?: LoadSourceAssetUrls + readonly loadSourceAssetUrls: LoadSourceAssetUrls readonly assetUrlsBySourceId: Map< string, Promise>> > - readonly hardenedAssetUrlByKey: Map> - readonly blobStore: ChatMediaAssetBlobStore - readonly fetchAsset: FetchChatMediaAsset } -type DemoAssetRoute = { - readonly demoSourceId: string - readonly encodedAssetPath: string - readonly decodedAssetPath: string -} - -const chatAssetsDirectoryName = "chat-assets" const parsedResultDirectoryName = "parsed-result" -const fallbackContentType = "application/octet-stream" -const defaultFetchAsset: FetchChatMediaAsset = (url) => fetch(url) -const defaultBlobStore: ChatMediaAssetBlobStore = { - put: (pathname, body, options) => - put(pathname, body, { - access: options.access ?? "public", - allowOverwrite: options.allowOverwrite, - contentType: options.contentType, - multipart: options.multipart, - }), -} +const chatAssetsDirectoryName = "chat-assets" +/** + * Resolve chat citation/media asset URLs to durable Notebook Blob URLs. The + * single hardening path is the SDK's `assetUrlPolicy: "durable"` read that + * `loadSourceAssetUrls` performs; here we only map a retrieval result's + * reference text to the durable URL that read produced. + * + * An asset that is already Notebook-owned is kept as-is. An asset that cannot + * be resolved to a durable URL is omitted rather than exposing a presigned + * Knowhere URL to the client. + */ export async function hardenChatMediaAssetUrls({ results, artifacts, - workspaceId, sources, loadSourceAssetUrls, - blobStore = defaultBlobStore, - fetchAsset = defaultFetchAsset, }: HardenChatMediaAssetUrlsForWorkspaceInput): Promise { const context: HardeningContext = { - workspaceId, sourcesByDocumentId: createSourcesByDocumentId(sources), loadSourceAssetUrls, assetUrlsBySourceId: new Map(), - hardenedAssetUrlByKey: new Map(), - blobStore, - fetchAsset, } const hardenedResults = await Promise.all( @@ -152,17 +104,13 @@ async function hardenRetrievalResult( if (!assetUrl && !pageCitationAssetUrl) return result const hardenedAssetUrl = assetUrl - ? await hardenAssetUrl( - { - assetUrl, - source: result.source, - content: result.content, - }, + ? await resolveDurableAssetUrl( + { assetUrl, source: result.source, content: result.content }, context, ) : undefined const hardenedPageCitationAssetUrl = pageCitationAssetUrl - ? await hardenAssetUrl( + ? await resolveDurableAssetUrl( { assetUrl: pageCitationAssetUrl, source: result.source, @@ -171,22 +119,13 @@ async function hardenRetrievalResult( context, ) : undefined - const hasAssetUrlChange = Boolean( - assetUrl && hardenedAssetUrl !== result.assetUrl, - ) - const hasPageCitationAssetUrlChange = Boolean( - pageCitationAssetUrl && - hardenedPageCitationAssetUrl !== result.pageCitationAssetUrl, - ) - if (!hasAssetUrlChange && !hasPageCitationAssetUrlChange) return result - return { - ...result, - ...(hardenedAssetUrl ? { assetUrl: hardenedAssetUrl } : {}), - ...(hardenedPageCitationAssetUrl - ? { pageCitationAssetUrl: hardenedPageCitationAssetUrl } - : {}), - } + return applyAssetUrls(result, { + hadAssetUrl: Boolean(assetUrl), + hadPageCitationAssetUrl: Boolean(pageCitationAssetUrl), + assetUrl: hardenedAssetUrl, + pageCitationAssetUrl: hardenedPageCitationAssetUrl, + }) } async function hardenArtifact( @@ -197,27 +136,26 @@ async function hardenArtifact( ? await hardenCitation(artifact.citation, context) : undefined const assetUrl = getTrimmedString(artifact.assetUrl) - if (!assetUrl) { - return citation && citation !== artifact.citation - ? { ...artifact, citation } - : artifact - } + const hardenedAssetUrl = assetUrl + ? await resolveDurableAssetUrl( + { + assetUrl, + source: artifact.citation?.source, + content: artifact.label, + }, + context, + ) + : undefined - const hardenedAssetUrl = await hardenAssetUrl( - { - assetUrl, - source: artifact.citation?.source, - content: artifact.label, - }, - context, - ) - const hasAssetUrlChange = hardenedAssetUrl !== artifact.assetUrl - const hasCitationChange = citation && citation !== artifact.citation - if (!hasAssetUrlChange && !hasCitationChange) return artifact + const citationChanged = citation && citation !== artifact.citation + const assetUrlChanged = assetUrl + ? hardenedAssetUrl !== artifact.assetUrl + : false + if (!citationChanged && !assetUrlChanged) return artifact return { ...artifact, - assetUrl: hardenedAssetUrl, + ...(assetUrl ? { assetUrl: hardenedAssetUrl } : {}), ...(citation ? { citation } : {}), } } @@ -231,17 +169,13 @@ async function hardenCitation( if (!assetUrl && !pageCitationAssetUrl) return citation const hardenedAssetUrl = assetUrl - ? await hardenAssetUrl( - { - assetUrl, - source: citation.source, - content: citation.content, - }, + ? await resolveDurableAssetUrl( + { assetUrl, source: citation.source, content: citation.content }, context, ) : undefined const hardenedPageCitationAssetUrl = pageCitationAssetUrl - ? await hardenAssetUrl( + ? await resolveDurableAssetUrl( { assetUrl: pageCitationAssetUrl, source: citation.source, @@ -250,70 +184,34 @@ async function hardenCitation( context, ) : undefined - const hasAssetUrlChange = Boolean( - assetUrl && hardenedAssetUrl !== citation.assetUrl, - ) - const hasPageCitationAssetUrlChange = Boolean( - pageCitationAssetUrl && - hardenedPageCitationAssetUrl !== citation.pageCitationAssetUrl, - ) - if (!hasAssetUrlChange && !hasPageCitationAssetUrlChange) return citation - return { - ...citation, - ...(hardenedAssetUrl ? { assetUrl: hardenedAssetUrl } : {}), - ...(hardenedPageCitationAssetUrl - ? { pageCitationAssetUrl: hardenedPageCitationAssetUrl } - : {}), - } + return applyAssetUrls(citation, { + hadAssetUrl: Boolean(assetUrl), + hadPageCitationAssetUrl: Boolean(pageCitationAssetUrl), + assetUrl: hardenedAssetUrl, + pageCitationAssetUrl: hardenedPageCitationAssetUrl, + }) } -async function hardenAssetUrl( +/** + * Return a durable Notebook-owned URL for a reference: keep already-owned URLs, + * otherwise resolve against the source's durable parsed asset map. Returns + * `undefined` when no durable URL is available so callers omit the URL rather + * than leak a presigned Knowhere URL. + */ +async function resolveDurableAssetUrl( reference: AssetUrlReference, context: HardeningContext, -): Promise { +): Promise { if (isNotebookOwnedAssetUrl(reference.assetUrl)) { return reference.assetUrl } - const parsedAssetUrl = await resolveParsedAssetUrl(reference, context) - if (parsedAssetUrl) return parsedAssetUrl - - const fetchRequest = resolveAssetFetchRequest(reference.assetUrl) - if (!fetchRequest) return reference.assetUrl - - const source = resolveSourceForReference(reference, context) - const sourceSegment = source - ? `source-${toSafePathSegment(source.id)}` - : fetchRequest.sourceSegment - const hardeningKey = [ - context.workspaceId, - source?.id ?? reference.source?.documentId ?? "", - fetchRequest.canonicalKey, - ].join("\0") - const cached = context.hardenedAssetUrlByKey.get(hardeningKey) - if (cached) return cached - - const hardenedAssetUrl = copyAssetToBlob({ - reference, - fetchRequest, - context, - sourceSegment, - hardeningKey, - }) - context.hardenedAssetUrlByKey.set(hardeningKey, hardenedAssetUrl) - return hardenedAssetUrl -} - -async function resolveParsedAssetUrl( - reference: AssetUrlReference, - context: HardeningContext, -): Promise { const source = resolveSourceForReference(reference, context) - if (!source || !context.loadSourceAssetUrls) return null + if (!source) return undefined const assetUrlsByFilePath = await getCachedSourceAssetUrls(source, context) - return resolveAssetUrlFromReferenceText({ + const durableUrl = resolveAssetUrlFromReferenceText({ values: [ reference.source?.sectionPath, reference.content, @@ -321,6 +219,46 @@ async function resolveParsedAssetUrl( ], assetUrlsByFilePath, }) + return durableUrl ?? undefined +} + +function applyAssetUrls< + T extends { + readonly assetUrl?: string | null + readonly pageCitationAssetUrl?: string | null + }, +>( + value: T, + hardened: { + readonly hadAssetUrl: boolean + readonly hadPageCitationAssetUrl: boolean + readonly assetUrl: string | undefined + readonly pageCitationAssetUrl: string | undefined + }, +): T { + const assetUrlChanged = + hardened.hadAssetUrl && hardened.assetUrl !== value.assetUrl + const pageCitationChanged = + hardened.hadPageCitationAssetUrl && + hardened.pageCitationAssetUrl !== value.pageCitationAssetUrl + if (!assetUrlChanged && !pageCitationChanged) return value + + const next: Record = { ...value } + if (hardened.hadAssetUrl) { + if (hardened.assetUrl) { + next["assetUrl"] = hardened.assetUrl + } else { + delete next["assetUrl"] + } + } + if (hardened.hadPageCitationAssetUrl) { + if (hardened.pageCitationAssetUrl) { + next["pageCitationAssetUrl"] = hardened.pageCitationAssetUrl + } else { + delete next["pageCitationAssetUrl"] + } + } + return next as T } async function getCachedSourceAssetUrls( @@ -330,116 +268,25 @@ async function getCachedSourceAssetUrls( const cached = context.assetUrlsBySourceId.get(source.id) if (cached) return cached - const loaded = context.loadSourceAssetUrls - ? context.loadSourceAssetUrls(source).catch((error: unknown) => { - logger.warn("chat-agent: failed to load parsed asset map", { - sourceId: source.id, - error: formatUnknownError(error), - }) - return {} + const loaded = context + .loadSourceAssetUrls(source) + .catch((error: unknown) => { + logger.warn("chat: failed to load durable parsed asset map", { + sourceId: source.id, + error: formatUnknownError(error), }) - : Promise.resolve({}) + return {} + }) context.assetUrlsBySourceId.set(source.id, loaded) return loaded } -async function copyAssetToBlob(input: { - readonly reference: AssetUrlReference - readonly fetchRequest: AssetFetchRequest - readonly context: HardeningContext - readonly sourceSegment: string - readonly hardeningKey: string -}): Promise { - try { - const response = await input.context.fetchAsset(input.fetchRequest.fetchUrl) - if (!response.ok) { - logger.warn("chat-agent: media asset hardening fetch failed", { - assetUrl: redactAssetUrl(input.reference.assetUrl), - status: response.status, - }) - return input.reference.assetUrl - } - - const body = Buffer.from(await response.arrayBuffer()) - const contentType = normalizeContentType( - response.headers.get("content-type"), - input.fetchRequest.suggestedFileName, - ) - const blobPathname = getChatAssetBlobPathname({ - workspaceId: input.context.workspaceId, - sourceSegment: input.sourceSegment, - hardeningKey: input.hardeningKey, - suggestedFileName: input.fetchRequest.suggestedFileName, - contentType, - }) - const blob = await input.context.blobStore.put(blobPathname, body, { - access: "public", - allowOverwrite: true, - contentType, - multipart: true, - }) - return blob.url - } catch (error) { - logger.warn("chat-agent: media asset hardening failed; keeping raw URL", { - assetUrl: redactAssetUrl(input.reference.assetUrl), - error: formatUnknownError(error), - }) - return input.reference.assetUrl - } -} - -function resolveAssetFetchRequest(assetUrl: string): AssetFetchRequest | null { - const demoAsset = parseDemoAssetRoute(assetUrl) - if (demoAsset) { - return { - fetchUrl: knowhereDemoApi.resolveApiURL( - `/api/v1/demo/sources/${encodeURIComponent( - demoAsset.demoSourceId, - )}/assets/${demoAsset.encodedAssetPath}`, - ), - canonicalKey: `demo:${demoAsset.demoSourceId}:${demoAsset.decodedAssetPath}`, - sourceSegment: `demo-${toSafePathSegment(demoAsset.demoSourceId)}`, - suggestedFileName: getPathBasename(demoAsset.decodedAssetPath), - } - } - - const absoluteUrl = parseAbsoluteHttpUrl(assetUrl) - if (!absoluteUrl) return null - - return { - fetchUrl: assetUrl, - canonicalKey: `url:${absoluteUrl.origin}${absoluteUrl.pathname}`, - sourceSegment: `external-${hashText(absoluteUrl.origin).slice(0, 16)}`, - suggestedFileName: getPathBasename(absoluteUrl.pathname), - } -} - -function parseDemoAssetRoute(assetUrl: string): DemoAssetRoute | null { - const pathname = getAssetUrlPathname(assetUrl) - const match = /^\/api\/demo-sources\/([^/]+)\/assets\/(.+)$/.exec(pathname) - const encodedDemoSourceId = match?.[1] - const encodedAssetPath = match?.[2] - if (!encodedDemoSourceId || !encodedAssetPath) return null - - const demoSourceId = decodeUrlComponent(encodedDemoSourceId) - const assetPathSegments = encodedAssetPath - .split("/") - .map(decodeUrlComponent) - .filter((segment): boolean => segment.length > 0) - if (!demoSourceId || assetPathSegments.length === 0) return null - - return { - demoSourceId, - encodedAssetPath: assetPathSegments.map(encodeURIComponent).join("/"), - decodedAssetPath: assetPathSegments.join("/"), - } -} - -function isNotebookOwnedAssetUrl(assetUrl: string): boolean { +export function isNotebookOwnedAssetUrl(assetUrl: string): boolean { const pathname = getAssetUrlPathname(assetUrl).toLowerCase() if ( pathname.includes(`/${parsedResultDirectoryName}/`) || - pathname.includes(`/${chatAssetsDirectoryName}/`) + pathname.includes(`/${chatAssetsDirectoryName}/`) || + pathname.includes("/parsed-documents/") ) { return true } @@ -449,80 +296,6 @@ function isNotebookOwnedAssetUrl(assetUrl: string): boolean { return hostname?.endsWith(".blob.vercel-storage.com") === true } -function getChatAssetBlobPathname(input: { - readonly workspaceId: string - readonly sourceSegment: string - readonly hardeningKey: string - readonly suggestedFileName: string - readonly contentType: string -}): string { - const hash = hashText(input.hardeningKey).slice(0, 24) - const fileName = toSafeFileName(input.suggestedFileName, input.contentType) - return [ - "workspaces", - toSafePathSegment(input.workspaceId), - chatAssetsDirectoryName, - input.sourceSegment, - `${hash}-${fileName}`, - ].join("/") -} - -function normalizeContentType( - value: string | null, - fileName: string, -): string { - const normalized = value?.replace(/\s+/g, " ").trim() - if (normalized) return normalized - return getContentTypeForPath(fileName) -} - -function getContentTypeForPath(filePath: string): string { - const extension = path.extname(filePath).toLowerCase() - if (extension === ".jpg" || extension === ".jpeg") return "image/jpeg" - if (extension === ".png") return "image/png" - if (extension === ".gif") return "image/gif" - if (extension === ".webp") return "image/webp" - if (extension === ".svg") return "image/svg+xml" - if (extension === ".html" || extension === ".htm") { - return "text/html; charset=utf-8" - } - if (extension === ".csv") return "text/csv; charset=utf-8" - if (extension === ".pdf") return "application/pdf" - return fallbackContentType -} - -function getExtensionForContentType(contentType: string): string { - const normalized = contentType.split(";")[0]?.trim().toLowerCase() - if (normalized === "image/jpeg") return ".jpg" - if (normalized === "image/png") return ".png" - if (normalized === "image/gif") return ".gif" - if (normalized === "image/webp") return ".webp" - if (normalized === "image/svg+xml") return ".svg" - if (normalized === "text/html") return ".html" - if (normalized === "text/csv") return ".csv" - if (normalized === "application/pdf") return ".pdf" - return ".bin" -} - -function toSafeFileName(fileName: string, contentType: string): string { - const extension = getSafeFileExtension(fileName, contentType) - const stem = path.basename(fileName, path.extname(fileName)) - const safeStem = toSafePathSegment(stem) - return `${safeStem}${extension}` -} - -function getSafeFileExtension(fileName: string, contentType: string): string { - const extension = path.extname(fileName).toLowerCase() - if (/^\.[a-z0-9]{1,12}$/.test(extension)) return extension - return getExtensionForContentType(contentType) -} - -function getPathBasename(value: string): string { - const decodedPath = decodeUrlComponent(value) - const basename = decodedPath.replaceAll("\\", "/").split("/").pop() - return basename && basename.trim().length > 0 ? basename : "asset" -} - function getAssetUrlPathname(assetUrl: string): string { try { return new URL(assetUrl, "http://notebook.local").pathname @@ -558,33 +331,6 @@ function createSourcesByDocumentId( ) } -function toSafePathSegment(value: string): string { - const decoded = decodeUrlComponent(value) - const normalized = decoded - .replace(/[^A-Za-z0-9._-]+/g, "-") - .replace(/^-+|-+$/g, "") - .slice(0, 80) - return normalized || hashText(value).slice(0, 16) -} - -function hashText(value: string): string { - return createHash("sha256").update(value).digest("hex") -} - -function decodeUrlComponent(value: string): string { - try { - return decodeURIComponent(value) - } catch { - return value - } -} - -function redactAssetUrl(assetUrl: string): string { - const absoluteUrl = parseAbsoluteHttpUrl(assetUrl) - if (absoluteUrl) return `${absoluteUrl.origin}${absoluteUrl.pathname}` - return getAssetUrlPathname(assetUrl) -} - function formatUnknownError(error: unknown): string { if (error instanceof Error) return error.message return String(error) diff --git a/src/domains/chat/media-assets.test.ts b/src/domains/chat/media-assets.test.ts index f79338b..2fc84cb 100644 --- a/src/domains/chat/media-assets.test.ts +++ b/src/domains/chat/media-assets.test.ts @@ -298,6 +298,7 @@ function makeSource(overrides: Partial = {}): Source { sizeBytes: 100, status: "ready", failureReason: null, + failureStage: null, knowhereJobId: "job_1", knowhereDocumentId: "doc_1", stagedBlobPathname: null, diff --git a/src/domains/chat/page-citation-assets.test.ts b/src/domains/chat/page-citation-assets.test.ts index ed62c66..d99dbcb 100644 --- a/src/domains/chat/page-citation-assets.test.ts +++ b/src/domains/chat/page-citation-assets.test.ts @@ -171,6 +171,7 @@ function makeSource(overrides: Partial = {}): Source { sizeBytes: 1, status: "ready", failureReason: null, + failureStage: null, knowhereJobId: "job_1", knowhereDocumentId: "doc_1", stagedBlobPathname: null, diff --git a/src/domains/chat/route-answer.ts b/src/domains/chat/route-answer.ts index 3df55e0..9582b77 100644 --- a/src/domains/chat/route-answer.ts +++ b/src/domains/chat/route-answer.ts @@ -11,9 +11,10 @@ import { type ChatTurnValue, } from "@/domains/chat/service" import { chatTurnPersistence } from "@/domains/chat/chat-turn-persistence" +import { readSourceAssetUrls } from "@/domains/chunks/read" import { startBackgroundReconciliation } from "@/domains/sources/background-reconcile" -import { sourceService } from "@/domains/sources/service" import { sourceWorkflowRuntime } from "@/domains/sources/workflow-runtime" +import { makeKnowhereClientWithParsedStorage } from "@/integrations/knowhere" import { notebookRequestContext } from "@/domains/workspace/request-context" import type { Source } from "@/infrastructure/db/schema" import { isAuthError } from "@/integrations/dashboard/api-key-service" @@ -68,31 +69,29 @@ const answerChatEffect = (input: AnswerChatInput) => apiKey, }), ) + const { knowledge } = makeKnowhereClientWithParsedStorage(apiKey, { + workspaceId: workspace.id, + }) const loadSourceAssetUrls = async ( source: (typeof sources)[number], ): Promise>> => { + if (source.status !== "ready" || !source.knowhereDocumentId) return {} + try { - const snapshot = await sourceService.ensureParsedSnapshotForRead({ - workspaceId: workspace.id, - source, - client: client.knowledge - ? { - documents: client.documents, - knowledge: client.knowledge, - } - : null, + return await readSourceAssetUrls({ + knowledge, + documentId: source.knowhereDocumentId, + revisionKey: source.knowhereJobId, }) - if (snapshot) return snapshot.assetUrlsByFilePath } catch (error) { - logger.warn("chat: parsed snapshot sync for assets failed", { + logger.warn("chat: durable parsed asset read failed", { workspaceId: workspace.id, sourceId: source.id, documentId: source.knowhereDocumentId, error: summarizeUnknownError(error), }) + return {} } - - return sourceService.getParseAssetUrls(workspace.id, source.id) } const result: Either.Either = diff --git a/src/domains/chat/route-service.test.ts b/src/domains/chat/route-service.test.ts index 189e862..f4df92a 100644 --- a/src/domains/chat/route-service.test.ts +++ b/src/domains/chat/route-service.test.ts @@ -18,6 +18,8 @@ const mocks = vi.hoisted(() => ({ loggerInfo: vi.fn(), loggerWarn: vi.fn(), listSourcesForWorkspace: vi.fn(), + makeKnowhereClientWithParsedStorage: vi.fn(), + readChunks: vi.fn(), softDeleteChatThread: vi.fn(), startBackgroundReconciliation: vi.fn(), })) @@ -51,6 +53,11 @@ vi.mock("@/domains/workspace/request-context", () => ({ }, })) +vi.mock("@/integrations/knowhere", () => ({ + makeKnowhereClientWithParsedStorage: + mocks.makeKnowhereClientWithParsedStorage, +})) + vi.mock("@/domains/chat/thread-service", () => ({ chatThreadService: { appendMessage: mocks.appendMessageToThread, @@ -73,11 +80,22 @@ vi.mock("@/lib/logger", () => ({ import { chatAnswerRouteService } from "./route-answer" import { chatThreadRouteService } from "./route-threads" -import { sourceService } from "@/domains/sources/service" describe("chat route services", () => { beforeEach(() => { vi.clearAllMocks() + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: { documents: { listChunks: vi.fn() } }, + knowledge: { readChunks: mocks.readChunks }, + }) + mocks.readChunks.mockResolvedValue({ + document: { localDocumentId: "doc" }, + chunks: [], + page: 1, + pageSize: 200, + totalChunks: 0, + totalPages: 1, + }) }) it("orchestrates a chat turn from request body to response body", async () => { @@ -141,28 +159,37 @@ describe("chat route services", () => { ) }) - it("ensures a legacy ready source snapshot before chat loads citation assets", async () => { + it("builds durable citation asset URLs from SDK reads for a ready source", async () => { const workspace = makeWorkspace() - const documents = { listChunks: vi.fn() } - const knowledge = { loadJobResult: vi.fn() } - const client = { documents, knowledge, retrieval: { query: vi.fn() } } + const client = { retrieval: { query: vi.fn() } } const readySource = makeSource({ status: "ready", knowhereDocumentId: "doc_legacy", - knowhereJobId: null, + knowhereJobId: "job_1", + }) + const durableUrl = + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_legacy/job_1/assets/pages/page-1.png" + mocks.readChunks.mockResolvedValue({ + document: { localDocumentId: "doc_legacy" }, + chunks: [ + { + position: 1, + chunkId: "c1", + chunkType: "page", + content: "Page", + readableContent: "Page", + sectionPath: "Page 1", + sourceChunkPath: "Page 1", + filePath: "pages/page-1.png", + assetUrl: durableUrl, + metadata: {}, + }, + ], + page: 1, + pageSize: 200, + totalChunks: 1, + totalPages: 1, }) - const snapshot = { - resultBlobUrl: "https://blob.example/manifest/current.json", - snapshotManifestUrl: "https://blob.example/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - assetUrlsByFilePath: { - "pages/page-1.png": "https://blob.example/pages/page-1.png", - }, - } - const ensureParsedSnapshotForRead = vi - .spyOn(sourceService, "ensureParsedSnapshotForRead") - .mockResolvedValue(snapshot) mocks.getAuthenticatedWithClient.mockResolvedValue({ user: { id: "user_1" }, workspace, @@ -177,7 +204,7 @@ describe("chat route services", () => { ) => Promise>> }) => { const assetUrls = await input.loadSourceAssetUrls?.(readySource) - expect(assetUrls).toEqual(snapshot.assetUrlsByFilePath) + expect(assetUrls).toEqual({ "pages/page-1.png": durableUrl }) return Either.right({ threadId: "thread_1", messages: [ @@ -193,13 +220,16 @@ describe("chat route services", () => { }) expect(result.status).toBe(200) - expect(ensureParsedSnapshotForRead).toHaveBeenCalledWith({ - workspaceId: workspace.id, - source: readySource, - client: { - documents, - knowledge, - }, + expect(mocks.makeKnowhereClientWithParsedStorage).toHaveBeenCalledWith( + "jwt_123", + { workspaceId: workspace.id }, + ) + expect(mocks.readChunks).toHaveBeenCalledWith({ + documentId: "doc_legacy", + revisionKey: "job_1", + page: 1, + pageSize: 200, + assetUrlPolicy: "durable", }) }) @@ -444,6 +474,7 @@ function makeSource(overrides: Partial = {}): Source { sizeBytes: 100, status: "ready", failureReason: null, + failureStage: null, knowhereJobId: "job_123", knowhereDocumentId: "doc_1", stagedBlobPathname: null, diff --git a/src/domains/chat/service.test.ts b/src/domains/chat/service.test.ts index e931f04..7c0ce92 100644 --- a/src/domains/chat/service.test.ts +++ b/src/domains/chat/service.test.ts @@ -270,6 +270,7 @@ function makeSource(overrides: Partial = {}): Source { sizeBytes: 100, status: "ready", failureReason: null, + failureStage: null, knowhereJobId: "job_123", knowhereDocumentId: "doc_included", stagedBlobPathname: null, diff --git a/src/domains/chunks/index.test.ts b/src/domains/chunks/index.test.ts index 0dfc22f..316bfd6 100644 --- a/src/domains/chunks/index.test.ts +++ b/src/domains/chunks/index.test.ts @@ -578,6 +578,7 @@ function makeSource(overrides: Partial = {}): Source { sizeBytes: 100, status: "ready", failureReason: null, + failureStage: null, knowhereJobId: "job_123", knowhereDocumentId: "doc_123", stagedBlobPathname: null, diff --git a/src/domains/chunks/index.ts b/src/domains/chunks/index.ts index 8a285d7..5e1b25a 100644 --- a/src/domains/chunks/index.ts +++ b/src/domains/chunks/index.ts @@ -1,5 +1,5 @@ import { Effect } from "effect" -import type { DocumentChunk } from "@ontos-ai/knowhere-sdk" +import type { DocumentChunk, KnowledgeReadChunk } from "@ontos-ai/knowhere-sdk" import { parsedChunkNormalization } from "./normalization" import type { Source } from "@/infrastructure/db/schema" @@ -179,6 +179,35 @@ export function toParsedChunkView( }) } +/** + * Map an SDK `KnowledgeReadChunk` (from `knowledge.readChunks`) to the view + * model. Asset URLs on durable reads are already hardened SDK-side, so no + * `assetUrlsByFilePath` remap is needed here. + */ +export function toParsedChunkViewFromReadChunk( + chunk: KnowledgeReadChunk, + sourceTitle: string, + documentId?: string, +): ParsedChunkView { + return parsedChunkNormalization.createParsedChunkView({ + chunkId: chunk.chunkId, + documentId, + parserChunkId: chunk.chunkId, + sectionPath: chunk.sectionPath, + chunkType: chunk.chunkType, + contentSource: chunk.contentSource, + content: chunk.content, + metadata: chunk.metadata, + filePathCandidates: [ + chunk.filePath, + chunk.metadata["filePath"], + chunk.metadata["file_path"], + ], + assetUrl: chunk.assetUrl, + sourceTitle, + }) +} + export function resolveCitationChunk( citation: ChatCitationView, chunks: readonly ParsedChunkView[], diff --git a/src/domains/chunks/read.test.ts b/src/domains/chunks/read.test.ts new file mode 100644 index 0000000..1d0786d --- /dev/null +++ b/src/domains/chunks/read.test.ts @@ -0,0 +1,123 @@ +import { describe, expect, it, vi } from "vitest" +import type { Knowledge } from "@ontos-ai/knowhere-sdk" + +import { readAllSourceChunks, readSourceChunkPage } from "./read" + +function makeReadChunk(overrides: Record = {}) { + return { + position: 1, + chunkId: "parser_1", + chunkType: "text", + content: "Body", + readableContent: "Body", + sectionPath: "Summary", + sourceChunkPath: "Summary", + filePath: undefined, + assetUrl: undefined, + metadata: {}, + ...overrides, + } +} + +describe("readSourceChunkPage", () => { + it("reads a durable page and maps chunks to the view model", async () => { + const readChunks = vi.fn(async () => ({ + document: { localDocumentId: "doc_1" }, + chunks: [ + makeReadChunk({ + chunkType: "image", + filePath: "images/a.png", + assetUrl: "https://blob.example/images/a.png", + }), + ], + page: 2, + pageSize: 50, + totalChunks: 120, + totalPages: 3, + })) + const knowledge = { readChunks } as unknown as Knowledge + + const result = await readSourceChunkPage({ + knowledge, + source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, + params: { page: 2, pageSize: 50 }, + }) + + expect(readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "rev_1", + page: 2, + pageSize: 50, + assetUrlPolicy: "durable", + }) + expect(result.pagination).toEqual({ + page: 2, + pageSize: 50, + total: 120, + totalPages: 3, + }) + expect(result.chunks[0]).toMatchObject({ + parserChunkId: "parser_1", + documentId: "doc_1", + assetUrl: "https://blob.example/images/a.png", + sourceTitle: "notes.pdf", + }) + }) + + it("omits revisionKey when the source has none", async () => { + const readChunks = vi.fn(async () => ({ + document: { localDocumentId: "doc_1" }, + chunks: [], + page: 1, + pageSize: 50, + totalChunks: 0, + totalPages: 1, + })) + const knowledge = { readChunks } as unknown as Knowledge + + await readSourceChunkPage({ + knowledge, + source: { documentId: "doc_1", title: "notes.pdf", revisionKey: null }, + params: { page: 1, pageSize: 50 }, + }) + + expect(readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + page: 1, + pageSize: 50, + assetUrlPolicy: "durable", + }) + }) +}) + +describe("readAllSourceChunks", () => { + it("pages the SDK to exhaustion", async () => { + const readChunks = vi + .fn() + .mockResolvedValueOnce({ + document: { localDocumentId: "doc_1" }, + chunks: [makeReadChunk({ chunkId: "c1" })], + page: 1, + pageSize: 200, + totalChunks: 2, + totalPages: 2, + }) + .mockResolvedValueOnce({ + document: { localDocumentId: "doc_1" }, + chunks: [makeReadChunk({ chunkId: "c2" })], + page: 2, + pageSize: 200, + totalChunks: 2, + totalPages: 2, + }) + const knowledge = { readChunks } as unknown as Knowledge + + const chunks = await readAllSourceChunks({ + knowledge, + source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, + }) + + expect(readChunks).toHaveBeenCalledTimes(2) + expect(chunks.map((chunk) => chunk.parserChunkId)).toEqual(["c1", "c2"]) + }) +}) diff --git a/src/domains/chunks/read.ts b/src/domains/chunks/read.ts new file mode 100644 index 0000000..f375d48 --- /dev/null +++ b/src/domains/chunks/read.ts @@ -0,0 +1,142 @@ +import "server-only" + +import type { Knowledge } from "@ontos-ai/knowhere-sdk" + +import { toParsedChunkViewFromReadChunk, type ChunkPage, type ChunkPageParams } from "@/domains/chunks" +import type { ParsedChunkView } from "@/domains/chunks/types" + +const loadAllPageSize = 200 + +type ReadableSource = { + readonly documentId: string + readonly title: string + readonly revisionKey?: string | null +} + +/** + * Read a single display page of parsed chunks through the SDK. The SDK serves + * from configured Blob storage when fresh and falls back to Knowhere remote + * otherwise, hardening visible asset URLs into durable Blob URLs + * (`assetUrlPolicy: "durable"`) and scheduling a background sync on a miss. + */ +export async function readSourceChunkPage(input: { + readonly knowledge: Knowledge + readonly source: ReadableSource + readonly params: ChunkPageParams +}): Promise { + const response = await input.knowledge.readChunks({ + documentId: input.source.documentId, + ...(input.source.revisionKey ? { revisionKey: input.source.revisionKey } : {}), + page: input.params.page, + pageSize: input.params.pageSize, + assetUrlPolicy: "durable", + }) + + const chunks = response.chunks.map((chunk) => + toParsedChunkViewFromReadChunk(chunk, input.source.title, input.source.documentId), + ) + + return { + chunks, + pagination: { + page: response.page ?? input.params.page, + pageSize: response.pageSize ?? input.params.pageSize, + total: response.totalChunks ?? chunks.length, + totalPages: + response.totalPages ?? + Math.max(1, Math.ceil(chunks.length / input.params.pageSize)), + }, + } +} + +/** + * Read every parsed chunk for a source by paging the SDK to exhaustion. Used by + * the tree view and load-all display mode. + */ +export async function readAllSourceChunks(input: { + readonly knowledge: Knowledge + readonly source: ReadableSource +}): Promise { + const chunks: ParsedChunkView[] = [] + let page = 1 + let totalPages = 1 + + do { + const response = await input.knowledge.readChunks({ + documentId: input.source.documentId, + ...(input.source.revisionKey + ? { revisionKey: input.source.revisionKey } + : {}), + page, + pageSize: loadAllPageSize, + assetUrlPolicy: "durable", + }) + for (const chunk of response.chunks) { + chunks.push( + toParsedChunkViewFromReadChunk( + chunk, + input.source.title, + input.source.documentId, + ), + ) + } + totalPages = Math.max(1, response.totalPages ?? 1) + page += 1 + } while (page <= totalPages) + + return chunks +} + +/** + * Build a `filePath -> durable Blob URL` map for a source by paging durable + * reads to exhaustion. This is the single asset-hardening path for chat: the + * SDK writes any missing asset into Blob during the durable read and returns + * the durable URL, which we index by both the chunk file path and any + * `metadata.pageAssets[].artifactRef`. + */ +export async function readSourceAssetUrls(input: { + readonly knowledge: Knowledge + readonly documentId: string + readonly revisionKey?: string | null +}): Promise>> { + const assetUrlsByFilePath: Record = {} + let page = 1 + let totalPages = 1 + + do { + const response = await input.knowledge.readChunks({ + documentId: input.documentId, + ...(input.revisionKey ? { revisionKey: input.revisionKey } : {}), + page, + pageSize: loadAllPageSize, + assetUrlPolicy: "durable", + }) + for (const chunk of response.chunks) { + if (chunk.filePath && chunk.assetUrl) { + assetUrlsByFilePath[chunk.filePath] = chunk.assetUrl + } + collectPageAssetUrls(chunk.metadata, assetUrlsByFilePath) + } + totalPages = Math.max(1, response.totalPages ?? 1) + page += 1 + } while (page <= totalPages) + + return assetUrlsByFilePath +} + +function collectPageAssetUrls( + metadata: Record, + target: Record, +): void { + const pageAssets = metadata["pageAssets"] + if (!Array.isArray(pageAssets)) return + for (const pageAsset of pageAssets) { + if (typeof pageAsset !== "object" || pageAsset === null) continue + const record = pageAsset as Record + const artifactRef = record["artifactRef"] + const assetUrl = record["assetUrl"] + if (typeof artifactRef === "string" && typeof assetUrl === "string") { + target[artifactRef] = assetUrl + } + } +} diff --git a/src/domains/chunks/server.test.ts b/src/domains/chunks/server.test.ts deleted file mode 100644 index 2442b53..0000000 --- a/src/domains/chunks/server.test.ts +++ /dev/null @@ -1,563 +0,0 @@ -import { describe, expect, it, vi, type Mock } from "vitest" -import { Effect } from "effect" -import type { DocumentChunk } from "@ontos-ai/knowhere-sdk" - -import type { Source } from "@/infrastructure/db/schema" -import { - loadChunkPageForSource, - loadChunksForSource, -} from "./server" - -describe("server chunk cache", () => { - it("loads ready source chunks from the parsed Blob snapshot without calling Knowhere listChunks", async () => { - const listChunks = vi.fn(async () => ({ - documentId: "doc_1", - chunks: [], - pagination: { - page: 1, - pageSize: 1, - total: 0, - totalPages: 0, - }, - })) - const snapshotStore = createSnapshotStore({ - "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json": - JSON.stringify({ - version: 1, - kind: "knowhere-parsed-result-snapshot", - jobId: "job_1", - documentId: "doc_1", - sourceFileName: "notes.pdf", - totalChunks: 1, - chunkPageSize: 1, - chunkPages: [ - { - page: 1, - pageSize: 1, - chunkCount: 1, - key: "workspaces/workspace_1/sources/source_1/parsed-result/chunks/page-1.json", - }, - ], - assetUrlsByFilePath: { - "images/chart.png": "https://blob.example/images/chart.png", - }, - createdAt: "2026-07-03T00:00:00.000Z", - }), - "workspaces/workspace_1/sources/source_1/parsed-result/chunks/page-1.json": - JSON.stringify({ - version: 1, - jobId: "job_1", - documentId: "doc_1", - sourceFileName: "notes.pdf", - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - chunks: [ - { - id: "chunk_image_1", - chunkId: "parser_image_1", - chunkType: "image", - content: "Chart summary", - sectionPath: "Revenue", - sourceChunkPath: "notes.pdf/Revenue", - filePath: "images/chart.png", - sortOrder: 1, - metadata: { summary: "Chart summary" }, - }, - ], - }), - }) - - const page = await Effect.runPromise( - loadChunkPageForSource( - makeSource({ status: "ready", knowhereJobId: "job_1" }), - { documents: { listChunks } }, - { page: 1, pageSize: 1 }, - { - snapshot: { - resultBlobUrl: - "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - snapshotManifestUrl: - "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - assetUrlsByFilePath: {}, - }, - snapshotStore, - workspaceId: "workspace_1", - }, - ), - ) - - expect(listChunks).not.toHaveBeenCalled() - expect(page).toEqual({ - chunks: [ - expect.objectContaining({ - chunkId: "chunk_image_1", - parserChunkId: "parser_image_1", - type: "image", - assetUrl: "https://blob.example/images/chart.png", - sourceTitle: "notes.pdf", - }), - ], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - }) - }) - - it("returns upstream chunks on a visible cache miss and warms mirrored assets in the background", async () => { - const warmTasks: Array<() => Promise> = [] - const cacheStore = createCacheStore() - const listChunks = vi.fn(async () => ({ - documentId: "doc_1", - jobResultId: "revision_1", - chunks: [ - makeDocumentChunk({ - id: "image_1", - chunkId: "parser_image_1", - chunkType: "image", - filePath: "images/image-1.png", - assetUrl: "https://knowhere.example/assets/image-1.png", - }), - ], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - })) - const fetchAsset = vi.fn(async () => - new Response("image-body", { - headers: { "content-type": "image/png" }, - }), - ) - - const page = await Effect.runPromise( - loadChunkPageForSource( - makeSource({ knowhereJobId: "revision_1" }), - { documents: { listChunks } }, - { page: 1, pageSize: 1 }, - { - cacheStore, - fetchAsset, - scheduleWarm: (task) => warmTasks.push(task), - workspaceId: "workspace_1", - }, - ), - ) - - expect(page.chunks[0]?.assetUrl).toBe( - "https://knowhere.example/assets/image-1.png", - ) - expect(cacheStore.putMock).not.toHaveBeenCalled() - expect(warmTasks).toHaveLength(1) - - await warmTasks[0]?.() - - expect(fetchAsset).toHaveBeenCalledWith( - "https://knowhere.example/assets/image-1.png", - ) - expect(cacheStore.putMock).toHaveBeenCalledWith( - expect.stringContaining("/chunk-assets/revision_1/"), - Buffer.from("image-body"), - expect.objectContaining({ contentType: "image/png" }), - ) - const cachedPagePut = cacheStore.putMock.mock.calls.find( - ([pathname]) => - typeof pathname === "string" && pathname.endsWith(".json"), - ) - expect(cachedPagePut).toBeDefined() - const cachedPage = JSON.parse(String(cachedPagePut?.[1])) as { - readonly chunks: readonly { readonly assetUrl?: string }[] - } - expect(cachedPage.chunks[0]?.assetUrl).toContain( - "https://blob.example/workspaces/workspace_1/sources/source_1/chunk-assets/revision_1/", - ) - }) - - it("returns a cached visible page after verifying the current Knowhere revision", async () => { - const cachedPage = { - chunks: [ - { - chunkId: "image_1", - documentId: "doc_1", - type: "image", - content: "", - assetUrl: "https://blob.example/image-1.png", - sourceTitle: "notes.pdf", - }, - ], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - } - const listChunks = vi.fn(async () => ({ - documentId: "doc_1", - jobResultId: "revision_1", - chunks: [], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - })) - const cacheStore = createCacheStore({ - get: vi.fn(async () => ({ - statusCode: 200, - stream: createTextStream(JSON.stringify(cachedPage)), - })), - }) - - const page = await Effect.runPromise( - loadChunkPageForSource( - makeSource({ knowhereJobId: "revision_1" }), - { documents: { listChunks } }, - { page: 1, pageSize: 1 }, - { - cacheStore, - workspaceId: "workspace_1", - }, - ), - ) - - expect(page).toEqual(cachedPage) - expect(listChunks).toHaveBeenCalledWith("doc_1", { - page: 1, - pageSize: 1, - includeAssetUrls: false, - }) - expect(cacheStore.getMock).toHaveBeenCalledWith( - expect.stringContaining("/revision_1/visible/page-1-size-1.json"), - { access: "public" }, - ) - }) - - it("ignores old cached pages when Knowhere reports a new job id", async () => { - const warmTasks: Array<() => Promise> = [] - const staleCachedPage = { - chunks: [ - { - chunkId: "stale_1", - documentId: "doc_1", - type: "text", - content: "Old chunk", - sourceTitle: "notes.pdf", - }, - ], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - } - const cacheStore = createCacheStore({ - get: vi.fn(async (pathname: string) => - pathname.includes("/job_old/") - ? { - statusCode: 200, - stream: createTextStream(JSON.stringify(staleCachedPage)), - } - : null, - ), - }) - const listChunks = vi.fn(async ( - _documentId: string, - params: { readonly includeAssetUrls: boolean }, - ) => ({ - documentId: "doc_1", - jobId: "job_new", - chunks: params.includeAssetUrls - ? [ - makeDocumentChunk({ - id: "text_new", - chunkId: "parser_text_new", - content: "New chunk", - }), - ] - : [], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - })) - const onRevisionKey = vi.fn(async () => undefined) - - const page = await Effect.runPromise( - loadChunkPageForSource( - makeSource({ knowhereJobId: "job_old" }), - { documents: { listChunks } }, - { page: 1, pageSize: 1 }, - { - cacheStore, - onRevisionKey, - scheduleWarm: (task) => warmTasks.push(task), - workspaceId: "workspace_1", - }, - ), - ) - - expect(page.chunks[0]?.content).toBe("New chunk") - expect(cacheStore.getMock).toHaveBeenCalledWith( - expect.stringContaining("/job_new/visible/page-1-size-1.json"), - { access: "public" }, - ) - expect(cacheStore.getMock).not.toHaveBeenCalledWith( - expect.stringContaining("/job_old/visible/page-1-size-1.json"), - expect.anything(), - ) - expect(onRevisionKey).toHaveBeenCalledWith("job_new") - expect(warmTasks).toHaveLength(1) - }) - - it("uses structure-only chunk loading for full-tree requests", async () => { - const warmTasks: Array<() => Promise> = [] - const cacheStore = createCacheStore() - const listChunks = vi.fn(async () => ({ - documentId: "doc_1", - jobResultId: "revision_1", - chunks: [ - makeDocumentChunk({ - id: "text_1", - chunkId: "parser_text_1", - content: "A text chunk", - }), - ], - pagination: { - page: 1, - pageSize: 200, - total: 1, - totalPages: 1, - }, - })) - const fetchAsset = vi.fn() - - const chunks = await Effect.runPromise( - loadChunksForSource( - makeSource({ knowhereJobId: "revision_1" }), - { documents: { listChunks } }, - { - cacheStore, - fetchAsset, - scheduleWarm: (task) => warmTasks.push(task), - workspaceId: "workspace_1", - }, - ), - ) - - expect(chunks).toMatchObject([{ chunkId: "text_1" }]) - expect(listChunks).toHaveBeenCalledWith("doc_1", { - page: 1, - pageSize: 200, - includeAssetUrls: false, - }) - expect(fetchAsset).not.toHaveBeenCalled() - expect(warmTasks).toHaveLength(1) - await warmTasks[0]?.() - expect(cacheStore.putMock).toHaveBeenCalledWith( - expect.stringContaining("/structure/page-1-size-200.json"), - expect.any(String), - expect.objectContaining({ contentType: "application/json; charset=utf-8" }), - ) - }) - - it("caches media chunks without previews when Knowhere has no upstream asset URL", async () => { - const warmTasks: Array<() => Promise> = [] - const cacheStore = createCacheStore() - const listChunks = vi.fn(async () => ({ - documentId: "doc_1", - jobResultId: "revision_1", - chunks: [ - makeDocumentChunk({ - id: "image_1", - chunkId: "parser_image_1", - chunkType: "image", - filePath: "images/image-1.png", - assetUrl: null, - }), - ], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - })) - const fetchAsset = vi.fn() - - const page = await Effect.runPromise( - loadChunkPageForSource( - makeSource({ knowhereJobId: "revision_1" }), - { documents: { listChunks } }, - { page: 1, pageSize: 1 }, - { - cacheStore, - fetchAsset, - scheduleWarm: (task) => warmTasks.push(task), - workspaceId: "workspace_1", - }, - ), - ) - - expect(page.chunks[0]?.assetUrl).toBeUndefined() - await warmTasks[0]?.() - expect(fetchAsset).not.toHaveBeenCalled() - const cachedPagePut = cacheStore.putMock.mock.calls.find( - ([pathname]) => - typeof pathname === "string" && pathname.endsWith(".json"), - ) - const cachedPage = JSON.parse(String(cachedPagePut?.[1])) as { - readonly chunks: readonly { readonly assetUrl?: string }[] - } - expect(cachedPage.chunks[0]?.assetUrl).toBeUndefined() - }) -}) - -type TestCacheGetResult = - | { - readonly statusCode: 200 - readonly stream: ReadableStream - } - | { - readonly statusCode: 304 - readonly stream: null - } - -type TestCachePutOptions = { - readonly access: "public" - readonly allowOverwrite: boolean - readonly contentType: string - readonly multipart?: boolean -} - -type TestCacheStore = { - readonly get: ( - pathname: string, - options: { readonly access: "public" }, - ) => Promise - readonly put: ( - pathname: string, - body: string | Buffer, - options: TestCachePutOptions, - ) => Promise<{ readonly url: string }> - readonly getMock: TestCacheGetMock - readonly putMock: TestCachePutMock -} - -type TestSnapshotStore = { - readonly get: ( - pathname: string, - options: { readonly access: "public" }, - ) => Promise -} - -type TestCacheGetMock = Mock< - ( - pathname: string, - options: { readonly access: "public" }, - ) => Promise -> - -type TestCachePutMock = Mock< - ( - pathname: string, - body: string | Buffer, - options: TestCachePutOptions, - ) => Promise<{ readonly url: string }> -> - -function createCacheStore(overrides: Partial<{ - readonly get: TestCacheGetMock - readonly put: TestCachePutMock -}> = {}): TestCacheStore { - const getMock = - overrides.get ?? - vi.fn(async () => null) - const putMock = - overrides.put ?? - vi.fn(async (pathname: string) => ({ - url: `https://blob.example/${pathname}`, - })) - - return { - get: (pathname, options) => getMock(pathname, options), - put: (pathname, body, options) => putMock(pathname, body, options), - getMock, - putMock, - } -} - -function createSnapshotStore( - entriesByPathname: Readonly>, -): TestSnapshotStore { - return { - get: vi.fn(async (pathname: string) => { - const entry = entriesByPathname[pathname] - if (!entry) return null - return { - statusCode: 200 as const, - stream: createTextStream(entry), - } - }), - } -} - -function makeDocumentChunk( - overrides: Partial = {}, -): DocumentChunk { - return { - id: "document_chunk_1", - chunkId: "parser_chunk_1", - chunkType: "text", - content: "Chunk content", - sectionId: null, - sectionPath: null, - sourceChunkPath: null, - filePath: null, - sortOrder: 1, - metadata: {}, - assetUrl: null, - ...overrides, - } -} - -function makeSource(overrides: Partial = {}): Source { - return { - id: "source_1", - workspaceId: "workspace_1", - title: "notes.pdf", - mimeType: "application/pdf", - sizeBytes: 100, - status: "ready", - failureReason: null, - knowhereJobId: "revision_1", - knowhereDocumentId: "doc_1", - stagedBlobPathname: null, - stagedBlobUrl: null, - originalBlobPathname: null, - originalBlobUrl: null, - demoKey: null, - createdAt: new Date("2026-05-06T00:00:00Z"), - updatedAt: new Date("2026-05-06T00:00:00Z"), - deletedAt: null, - ...overrides, - } -} - -function createTextStream(text: string): ReadableStream { - const stream = new Response(text).body - if (!stream) throw new Error("Response body stream was not created.") - return stream -} diff --git a/src/domains/chunks/server.ts b/src/domains/chunks/server.ts deleted file mode 100644 index 64b7fe4..0000000 --- a/src/domains/chunks/server.ts +++ /dev/null @@ -1,911 +0,0 @@ -import "server-only" - -import path from "node:path" -import { createHash } from "node:crypto" -import { get as getBlob, put } from "@vercel/blob" -import { after } from "next/server" -import { Effect } from "effect" -import type { - DocumentChunk, -} from "@ontos-ai/knowhere-sdk" - -import { - resolveChunkConnectionTargets, - toParsedChunkView, - type ChunkKnowhereClient, - type ChunkPage, - type ChunkPageParams, - type LoadChunksOptions, -} from "@/domains/chunks" -import { - readParsedResultSnapshotChunkPage, - readParsedResultSnapshotManifest, - type ParsedResultSnapshotBlobStore, - type ParsedResultSnapshotChunk, - type ParsedResultSnapshotManifest, -} from "@/domains/sources/parse-result-storage-adapter" -import type { ParsedChunkView } from "@/domains/chunks/types" -import type { Source } from "@/infrastructure/db/schema" -import { logger } from "@/lib/logger" - -type ChunkPageMode = "visible" | "structure" - -type ChunkPageBlobGetResult = - | { - readonly statusCode: 200 - readonly stream: ReadableStream - } - | { - readonly statusCode: 304 - readonly stream: null - } - -type ChunkPageBlobPutOptions = { - readonly access: "public" - readonly allowOverwrite: boolean - readonly contentType: string - readonly multipart?: boolean -} - -type ChunkPageBlobStore = { - readonly get: ( - pathname: string, - options: { readonly access: "public" }, - ) => Promise - readonly put: ( - pathname: string, - body: string | Buffer, - options: ChunkPageBlobPutOptions, - ) => Promise<{ readonly url: string }> -} - -type FetchChunkAsset = (assetUrl: string) => Promise - -type ChunkPageWarmScheduler = (task: () => Promise) => void - -type ServerLoadChunksOptions = LoadChunksOptions & { - readonly workspaceId?: string - readonly snapshot?: SourceParsedSnapshotReadModel | null - readonly snapshotStore?: ParsedResultSnapshotBlobStore - readonly cacheStore?: ChunkPageBlobStore - readonly fetchAsset?: FetchChunkAsset - readonly mode?: ChunkPageMode - readonly onRevisionKey?: (revisionKey: string) => Promise - readonly scheduleWarm?: ChunkPageWarmScheduler -} - -type WarmChunkPageCacheInput = { - readonly source: Source - readonly client: ChunkKnowhereClient - readonly params: ChunkPageParams - readonly revisionKey: string - readonly workspaceId: string - readonly cacheStore: ChunkPageBlobStore - readonly fetchAsset: FetchChunkAsset - readonly scheduleWarm: ChunkPageWarmScheduler - readonly startAssetIndex?: number - readonly mirroredAssetUrlsByOriginalUrl?: Readonly> -} - -type MirrorableChunkAsset = { - readonly chunkId: string - readonly chunkType: "image" | "table" - readonly assetUrl: string - readonly filePath?: string -} - -export type SourceParsedSnapshotReadModel = { - readonly resultBlobUrl: string - readonly snapshotManifestUrl?: string | null - readonly snapshotManifestKey?: string | null - readonly assetUrlsByFilePath: Readonly> -} - -const documentChunkPageSize = 200 -const visibleChunkPageMode: ChunkPageMode = "visible" -const structureChunkPageMode: ChunkPageMode = "structure" -const maximumMirroredAssetsPerWarmStep = 50 -const maximumWarmStepDurationMs = 45_000 -const assetMirrorConcurrency = 10 - -const defaultBlobStore: ChunkPageBlobStore = { - get: (pathname, options) => getBlob(pathname, options), - put: (pathname, body, options) => - put(pathname, body, { - access: options.access, - allowOverwrite: options.allowOverwrite, - contentType: options.contentType, - multipart: options.multipart, - }), -} - -const defaultFetchAsset: FetchChunkAsset = (assetUrl: string) => fetch(assetUrl) - -const defaultScheduleWarm: ChunkPageWarmScheduler = ( - task: () => Promise, -) => { - try { - after(task) - } catch { - void task() - } -} - -export const loadChunksForSource = ( - source: Source, - client: ChunkKnowhereClient, - options: ServerLoadChunksOptions = {}, -) => - Effect.gen(function* () { - if (source.status !== "ready" || !source.knowhereDocumentId) return [] - - if (options.snapshot) { - const snapshotChunks = yield* Effect.promise(() => - loadAllSnapshotChunks({ - source, - snapshot: options.snapshot!, - blobStore: options.snapshotStore, - }), - ) - return resolveChunkConnectionTargets(snapshotChunks) - } - - const chunks: ParsedChunkView[] = [] - let page = 1 - let totalPages = 1 - - do { - const chunkPage = yield* loadChunkPageForSource(source, client, { - page, - pageSize: documentChunkPageSize, - }, { - ...options, - mode: structureChunkPageMode, - }) - chunks.push(...chunkPage.chunks) - totalPages = chunkPage.pagination.totalPages - page += 1 - } while (page <= totalPages) - - return resolveChunkConnectionTargets(chunks) - }) - -export const loadChunkPageForSource = ( - source: Source, - client: ChunkKnowhereClient, - params: ChunkPageParams, - options: ServerLoadChunksOptions = {}, -) => - Effect.gen(function* () { - const emptyPage = createEmptyChunkPage(params) - if (source.status !== "ready" || !source.knowhereDocumentId) { - return emptyPage - } - - if (options.snapshot) { - return yield* Effect.promise(() => - loadSnapshotChunkPage({ - source, - params, - snapshot: options.snapshot!, - blobStore: options.snapshotStore, - }), - ) - } - - const mode = options.mode ?? visibleChunkPageMode - const workspaceId = options.workspaceId ?? source.workspaceId - const cacheStore = options.cacheStore ?? defaultBlobStore - const includeAssetUrls = mode === visibleChunkPageMode - const revisionProbeResponse = yield* Effect.promise(() => - client.documents.listChunks(source.knowhereDocumentId!, { - page: params.page, - pageSize: params.pageSize, - includeAssetUrls: false, - }), - ) - const probeRevisionKey = getRevisionKey(revisionProbeResponse, source) - if (probeRevisionKey) { - scheduleRevisionKeyUpdate(source, probeRevisionKey, options.onRevisionKey) - const cachedPage = yield* Effect.promise(() => - readCachedChunkPage({ - cacheStore, - documentId: source.knowhereDocumentId!, - mode, - params, - revisionKey: probeRevisionKey, - workspaceId, - }), - ) - if (cachedPage) return cachedPage - } - - const response = includeAssetUrls - ? yield* Effect.promise(() => - client.documents.listChunks(source.knowhereDocumentId!, { - page: params.page, - pageSize: params.pageSize, - includeAssetUrls, - }), - ) - : revisionProbeResponse - const revisionKey = getRevisionKey(response, source) ?? probeRevisionKey - if (revisionKey && revisionKey !== probeRevisionKey) { - scheduleRevisionKeyUpdate(source, revisionKey, options.onRevisionKey) - } - - const chunkPage = createChunkPageFromResponse({ - response, - source, - params, - options: - mode === visibleChunkPageMode - ? { assetUrlsByFilePath: options.assetUrlsByFilePath } - : {}, - }) - - if (revisionKey) { - if (mode === visibleChunkPageMode) { - scheduleChunkPageWarm({ - source, - client, - params, - revisionKey, - workspaceId, - cacheStore, - fetchAsset: options.fetchAsset ?? defaultFetchAsset, - scheduleWarm: options.scheduleWarm ?? defaultScheduleWarm, - }) - } else { - scheduleStructurePageCacheWrite({ - cacheStore, - chunkPage, - documentId: source.knowhereDocumentId, - mode, - params, - revisionKey, - scheduleWarm: options.scheduleWarm ?? defaultScheduleWarm, - workspaceId, - }) - } - } - - return chunkPage - }) - -async function loadAllSnapshotChunks(input: { - readonly source: Source - readonly snapshot: SourceParsedSnapshotReadModel - readonly blobStore?: ParsedResultSnapshotBlobStore -}): Promise { - const manifest = await loadSnapshotManifest(input) - if (!manifest) return [] - - const pages = await Promise.all( - manifest.chunkPages.map((page) => - readParsedResultSnapshotChunkPage({ - pageKey: page.key, - blobStore: input.blobStore, - }), - ), - ) - const chunks = pages.flatMap((page) => page?.chunks ?? []) - return chunks.map((chunk) => - toParsedChunkView( - toDocumentChunkFromSnapshot(chunk), - input.source.title, - input.source.knowhereDocumentId ?? undefined, - { - assetUrlsByFilePath: { - ...manifest.assetUrlsByFilePath, - ...input.snapshot.assetUrlsByFilePath, - }, - }, - ), - ) -} - -async function loadSnapshotChunkPage(input: { - readonly source: Source - readonly params: ChunkPageParams - readonly snapshot: SourceParsedSnapshotReadModel - readonly blobStore?: ParsedResultSnapshotBlobStore -}): Promise { - const manifest = await loadSnapshotManifest(input) - if (!manifest) return createEmptyChunkPage(input.params) - - const pageReference = manifest.chunkPages.find( - (page) => page.page === input.params.page && page.pageSize === input.params.pageSize, - ) - if (!pageReference) { - return createEmptyChunkPage({ - page: input.params.page, - pageSize: input.params.pageSize, - }) - } - - const page = await readParsedResultSnapshotChunkPage({ - pageKey: pageReference.key, - blobStore: input.blobStore, - }) - if (!page) return createEmptyChunkPage(input.params) - - const assetUrlsByFilePath = { - ...manifest.assetUrlsByFilePath, - ...input.snapshot.assetUrlsByFilePath, - } - const chunks = page.chunks.map((chunk) => - toParsedChunkView( - toDocumentChunkFromSnapshot(chunk), - input.source.title, - input.source.knowhereDocumentId ?? undefined, - { assetUrlsByFilePath }, - ), - ) - - return { - chunks, - pagination: { - page: page.page, - pageSize: page.pageSize, - total: page.total, - totalPages: page.totalPages, - }, - } -} - -async function loadSnapshotManifest(input: { - readonly source: Source - readonly snapshot: SourceParsedSnapshotReadModel - readonly blobStore?: ParsedResultSnapshotBlobStore -}): Promise { - const manifestKey = input.snapshot.snapshotManifestKey - if (!manifestKey) { - logger.warn("chunks: ready source missing parsed snapshot manifest", { - sourceId: input.source.id, - documentId: input.source.knowhereDocumentId, - }) - return null - } - - const manifest = await readParsedResultSnapshotManifest({ - workspaceId: input.source.workspaceId, - sourceId: input.source.id, - manifestKey, - blobStore: input.blobStore, - }) - if (!manifest) { - logger.warn("chunks: parsed snapshot manifest could not be read", { - sourceId: input.source.id, - documentId: input.source.knowhereDocumentId, - manifestKey, - }) - } - return manifest -} - -function toDocumentChunkFromSnapshot(chunk: ParsedResultSnapshotChunk): DocumentChunk { - return { - id: chunk.id, - chunkId: chunk.chunkId, - chunkType: normalizeDocumentChunkType(chunk.chunkType), - content: chunk.content, - sectionId: null, - sectionPath: chunk.sectionPath ?? null, - sourceChunkPath: chunk.sourceChunkPath, - filePath: chunk.filePath ?? null, - sortOrder: chunk.sortOrder, - metadata: chunk.metadata, - assetUrl: chunk.assetUrl ?? null, - } -} - -function normalizeDocumentChunkType(value: string): DocumentChunk["chunkType"] { - if (value === "image" || value === "table" || value === "page") return value - return "text" -} - -export async function warmChunkPageCache( - input: WarmChunkPageCacheInput, -): Promise { - const response = await input.client.documents.listChunks( - input.source.knowhereDocumentId!, - { - page: input.params.page, - pageSize: input.params.pageSize, - includeAssetUrls: true, - }, - ) - const assets = collectMirrorableChunkAssets(response.chunks) - const startAssetIndex = input.startAssetIndex ?? 0 - const mirroredAssetUrls = new Map( - Object.entries(input.mirroredAssetUrlsByOriginalUrl ?? {}), - ) - - const stepStartedAt = Date.now() - const assetBatch = assets.slice( - startAssetIndex, - startAssetIndex + maximumMirroredAssetsPerWarmStep, - ) - const mirroredBatch = await mirrorChunkAssets({ - assets: assetBatch, - cacheStore: input.cacheStore, - fetchAsset: input.fetchAsset, - revisionKey: input.revisionKey, - source: input.source, - }) - for (const mirroredAsset of mirroredBatch) { - mirroredAssetUrls.set(mirroredAsset.assetUrl, mirroredAsset.blobUrl) - } - - const nextAssetIndex = startAssetIndex + assetBatch.length - if ( - nextAssetIndex < assets.length && - Date.now() - stepStartedAt < maximumWarmStepDurationMs - ) { - input.scheduleWarm(() => - warmChunkPageCache({ - ...input, - startAssetIndex: nextAssetIndex, - mirroredAssetUrlsByOriginalUrl: Object.fromEntries(mirroredAssetUrls), - }), - ) - return - } - - if (nextAssetIndex < assets.length) { - input.scheduleWarm(() => - warmChunkPageCache({ - ...input, - startAssetIndex: nextAssetIndex, - mirroredAssetUrlsByOriginalUrl: Object.fromEntries(mirroredAssetUrls), - }), - ) - return - } - - const rewrittenChunks = response.chunks.map((chunk) => - rewriteChunkAssetUrl(chunk, mirroredAssetUrls), - ) - const chunkPage = createChunkPageFromResponse({ - response: { - ...response, - chunks: rewrittenChunks, - }, - source: input.source, - params: input.params, - options: {}, - }) - - await writeCachedChunkPage({ - cacheStore: input.cacheStore, - chunkPage, - documentId: input.source.knowhereDocumentId!, - mode: visibleChunkPageMode, - params: input.params, - revisionKey: input.revisionKey, - workspaceId: input.workspaceId, - }) -} - -function scheduleChunkPageWarm(input: WarmChunkPageCacheInput): void { - input.scheduleWarm(async () => { - try { - await warmChunkPageCache(input) - } catch (error) { - logger.warn("chunks: chunk page cache warm failed", { - sourceId: input.source.id, - documentId: input.source.knowhereDocumentId, - page: input.params.page, - pageSize: input.params.pageSize, - revisionKey: input.revisionKey, - error: getErrorMessage(error), - }) - } - }) -} - -function scheduleStructurePageCacheWrite(input: { - readonly cacheStore: ChunkPageBlobStore - readonly chunkPage: ChunkPage - readonly documentId: string - readonly mode: ChunkPageMode - readonly params: ChunkPageParams - readonly revisionKey: string - readonly scheduleWarm: ChunkPageWarmScheduler - readonly workspaceId: string -}): void { - input.scheduleWarm(async () => { - try { - await writeCachedChunkPage(input) - } catch (error) { - logger.warn("chunks: structure chunk page cache write failed", { - documentId: input.documentId, - page: input.params.page, - pageSize: input.params.pageSize, - revisionKey: input.revisionKey, - error: getErrorMessage(error), - }) - } - }) -} - -async function readCachedChunkPage(input: { - readonly cacheStore: ChunkPageBlobStore - readonly documentId: string - readonly mode: ChunkPageMode - readonly params: ChunkPageParams - readonly revisionKey: string - readonly workspaceId: string -}): Promise { - const pathname = getChunkPageCachePathname(input) - const result = await input.cacheStore.get(pathname, { access: "public" }) - if (!result || result.statusCode !== 200) return null - - const text = await new Response(result.stream).text() - return parseCachedChunkPage(text) -} - -async function writeCachedChunkPage(input: { - readonly cacheStore: ChunkPageBlobStore - readonly chunkPage: ChunkPage - readonly documentId: string - readonly mode: ChunkPageMode - readonly params: ChunkPageParams - readonly revisionKey: string - readonly workspaceId: string -}): Promise { - await input.cacheStore.put( - getChunkPageCachePathname(input), - JSON.stringify(input.chunkPage), - { - access: "public", - allowOverwrite: true, - contentType: "application/json; charset=utf-8", - }, - ) -} - -function createChunkPageFromResponse(input: { - readonly response: { - readonly chunks: readonly DocumentChunk[] - readonly pagination?: { - readonly page?: number - readonly pageSize?: number - readonly total?: number - readonly totalPages?: number - } - } - readonly source: Source - readonly params: ChunkPageParams - readonly options: LoadChunksOptions -}): ChunkPage { - const chunks = input.response.chunks.map((chunk) => - toParsedChunkView( - chunk, - input.source.title, - input.source.knowhereDocumentId ?? undefined, - input.options, - ), - ) - - return { - chunks, - pagination: { - page: getFiniteNonNegativeNumber( - input.response.pagination?.page, - input.params.page, - ), - pageSize: getFiniteNonNegativeNumber( - input.response.pagination?.pageSize, - input.params.pageSize, - ), - total: getFiniteNonNegativeNumber( - input.response.pagination?.total, - chunks.length, - ), - totalPages: getFiniteNonNegativeNumber( - input.response.pagination?.totalPages, - Math.ceil(chunks.length / input.params.pageSize), - ), - }, - } -} - -function createEmptyChunkPage(params: ChunkPageParams): ChunkPage { - return { - chunks: [], - pagination: { - page: params.page, - pageSize: params.pageSize, - total: 0, - totalPages: 0, - }, - } -} - -function parseCachedChunkPage(text: string): ChunkPage | null { - try { - const value: unknown = JSON.parse(text) - if (!isRecord(value)) return null - const chunks = value["chunks"] - const pagination = value["pagination"] - if (!Array.isArray(chunks) || !isRecord(pagination)) return null - const page = getNumber(pagination["page"]) - const pageSize = getNumber(pagination["pageSize"]) - const total = getNumber(pagination["total"]) - const totalPages = getNumber(pagination["totalPages"]) - if ( - page === undefined || - pageSize === undefined || - total === undefined || - totalPages === undefined - ) { - return null - } - - return { - chunks: chunks.filter(isParsedChunkView), - pagination: { - page, - pageSize, - total, - totalPages, - }, - } - } catch { - return null - } -} - -async function mirrorChunkAssets(input: { - readonly assets: readonly MirrorableChunkAsset[] - readonly cacheStore: ChunkPageBlobStore - readonly fetchAsset: FetchChunkAsset - readonly revisionKey: string - readonly source: Source -}): Promise< - readonly { - readonly assetUrl: string - readonly blobUrl: string - }[] -> { - return mapWithConcurrency( - input.assets, - assetMirrorConcurrency, - async (asset) => { - const response = await input.fetchAsset(asset.assetUrl) - if (!response.ok) { - throw new Error( - `Chunk asset fetch failed with status ${response.status}.`, - ) - } - - const body = Buffer.from(await response.arrayBuffer()) - const blob = await input.cacheStore.put( - getMirroredAssetPathname({ - asset, - revisionKey: input.revisionKey, - source: input.source, - }), - body, - { - access: "public", - allowOverwrite: true, - contentType: getMirroredAssetContentType(asset, response), - multipart: true, - }, - ) - return { - assetUrl: asset.assetUrl, - blobUrl: blob.url, - } - }, - ) -} - -async function mapWithConcurrency( - inputs: readonly Input[], - concurrency: number, - mapInput: (input: Input) => Promise, -): Promise { - const results: Array = [] - let nextIndex = 0 - const workerCount = Math.min(concurrency, inputs.length) - - async function runWorker(): Promise { - while (nextIndex < inputs.length) { - const index = nextIndex - nextIndex += 1 - const input = inputs[index] - if (input === undefined) return - results[index] = await mapInput(input) - } - } - - await Promise.all( - Array.from({ length: workerCount }, () => runWorker()), - ) - - return results.filter( - (result): result is Output => result !== undefined, - ) -} - -function collectMirrorableChunkAssets( - chunks: readonly DocumentChunk[], -): readonly MirrorableChunkAsset[] { - return chunks.flatMap((chunk): readonly MirrorableChunkAsset[] => { - if (chunk.chunkType !== "image" && chunk.chunkType !== "table") return [] - if (typeof chunk.assetUrl !== "string" || chunk.assetUrl.length === 0) { - return [] - } - - return [ - { - chunkId: chunk.id, - chunkType: chunk.chunkType, - assetUrl: chunk.assetUrl, - ...(chunk.filePath ? { filePath: chunk.filePath } : {}), - }, - ] - }) -} - -function rewriteChunkAssetUrl( - chunk: DocumentChunk, - mirroredAssetUrls: ReadonlyMap, -): DocumentChunk { - if (typeof chunk.assetUrl !== "string") return chunk - const mirroredAssetUrl = mirroredAssetUrls.get(chunk.assetUrl) - return mirroredAssetUrl ? { ...chunk, assetUrl: mirroredAssetUrl } : chunk -} - -function getChunkPageCachePathname(input: { - readonly documentId: string - readonly mode: ChunkPageMode - readonly params: ChunkPageParams - readonly revisionKey: string - readonly workspaceId: string -}): string { - return [ - "workspaces", - encodePathSegment(input.workspaceId), - "chunk-pages", - encodePathSegment(input.documentId), - encodePathSegment(input.revisionKey), - input.mode, - `page-${input.params.page}-size-${input.params.pageSize}.json`, - ].join("/") -} - -function getMirroredAssetPathname(input: { - readonly asset: MirrorableChunkAsset - readonly revisionKey: string - readonly source: Source -}): string { - const basename = getMirroredAssetBasename(input.asset) - return [ - "workspaces", - encodePathSegment(input.source.workspaceId), - "sources", - encodePathSegment(input.source.id), - "chunk-assets", - encodePathSegment(input.revisionKey), - `${hashValue(input.asset.assetUrl)}-${basename}`, - ].join("/") -} - -function getMirroredAssetBasename(asset: MirrorableChunkAsset): string { - const candidate = - asset.filePath ?? - getUrlPathname(asset.assetUrl).split("/").filter(Boolean).at(-1) ?? - `${asset.chunkId}.bin` - return candidate.replace(/[^a-zA-Z0-9._-]+/g, "-") -} - -function getMirroredAssetContentType( - asset: MirrorableChunkAsset, - response: Response, -): string { - const headerContentType = response.headers.get("content-type") - if (headerContentType) return headerContentType - if (asset.chunkType === "table") return "text/html; charset=utf-8" - - const extension = path.extname(asset.filePath ?? getUrlPathname(asset.assetUrl)) - .toLowerCase() - if (extension === ".jpg" || extension === ".jpeg") return "image/jpeg" - if (extension === ".png") return "image/png" - if (extension === ".webp") return "image/webp" - if (extension === ".gif") return "image/gif" - return "application/octet-stream" -} - -function getRevisionKey( - response: { - readonly jobId?: string | null - readonly jobResultId?: string | null - }, - source: Source, -): string | null { - return ( - getNonEmptyString(response.jobId) ?? - getNonEmptyString(response.jobResultId) ?? - getFallbackRevisionKey(source) - ) -} - -function getFallbackRevisionKey(source: Source): string | null { - return ( - getNonEmptyString(source.knowhereJobId) ?? - getNonEmptyString(source.knowhereDocumentId) - ) -} - -function scheduleRevisionKeyUpdate( - source: Source, - revisionKey: string, - onRevisionKey: ((revisionKey: string) => Promise) | undefined, -): void { - if (!onRevisionKey || source.knowhereJobId === revisionKey) return - - void onRevisionKey(revisionKey).catch((error: unknown) => { - logger.warn("chunks: source revision key update failed", { - sourceId: source.id, - documentId: source.knowhereDocumentId, - revisionKey, - error: getErrorMessage(error), - }) - }) -} - -function getFiniteNonNegativeNumber( - value: number | undefined, - fallback: number, -): number { - return typeof value === "number" && Number.isFinite(value) && value >= 0 - ? value - : fallback -} - -function isParsedChunkView(value: unknown): value is ParsedChunkView { - if (!isRecord(value)) return false - return ( - typeof value["chunkId"] === "string" && - typeof value["type"] === "string" && - typeof value["content"] === "string" && - typeof value["sourceTitle"] === "string" - ) -} - -function getNumber(value: unknown): number | undefined { - return typeof value === "number" && Number.isFinite(value) ? value : undefined -} - -function getNonEmptyString(value: unknown): string | null { - return typeof value === "string" && value.length > 0 ? value : null -} - -function isRecord(value: unknown): value is Readonly> { - return typeof value === "object" && value !== null && !Array.isArray(value) -} - -function getUrlPathname(value: string): string { - try { - return new URL(value).pathname - } catch { - return value.split("?")[0] ?? value - } -} - -function encodePathSegment(value: string): string { - return encodeURIComponent(value) -} - -function hashValue(value: string): string { - return createHash("sha256").update(value).digest("hex").slice(0, 16) -} - -function getErrorMessage(error: unknown): string { - return error instanceof Error ? error.message : String(error) -} diff --git a/src/domains/sources/counts.test.ts b/src/domains/sources/counts.test.ts index e520eff..22f0b8c 100644 --- a/src/domains/sources/counts.test.ts +++ b/src/domains/sources/counts.test.ts @@ -14,6 +14,7 @@ function makeSource(overrides: Partial = {}): Source { sizeBytes: 1, status: "ready", failureReason: null, + failureStage: null, knowhereJobId: "job_1", knowhereDocumentId: "doc_1", stagedBlobPathname: null, @@ -29,36 +30,11 @@ function makeSource(overrides: Partial = {}): Source { } describe("countChunksBySourceId", () => { - it("counts ready source chunks from the parsed snapshot manifest", async () => { - const listChunks = vi.fn() + it("counts ready source chunks from the document total", async () => { + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) const mockClient = { documents: { listChunks }, } as unknown as Knowhere - const repository = { - getParseSnapshotMetadata: vi.fn(async (workspaceId: string, sourceId: string) => - sourceId === "ready" - ? { - resultBlobUrl: "https://blob.example/manifest/current.json", - snapshotManifestUrl: "https://blob.example/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/ready/parsed-result/manifest/current.json", - assetUrlsByFilePath: {}, - } - : null, - ), - } - const readSnapshotManifest = vi.fn(async () => ({ - version: 1 as const, - kind: "knowhere-parsed-result-snapshot" as const, - jobId: "job_1", - documentId: "doc_ready", - sourceFileName: "notes.pdf", - totalChunks: 12, - chunkPageSize: 50, - chunkPages: [], - assetUrlsByFilePath: {}, - createdAt: "2026-07-04T00:00:00.000Z", - })) const { countChunksBySourceId } = await import("./counts") @@ -66,48 +42,32 @@ describe("countChunksBySourceId", () => { countChunksBySourceId( [ makeSource({ id: "ready", knowhereDocumentId: "doc_ready" }), - makeSource({ id: "parsing", status: "parsing", knowhereDocumentId: null }), + makeSource({ + id: "parsing", + status: "parsing", + knowhereDocumentId: null, + }), makeSource({ id: "missing-doc", knowhereDocumentId: null }), ], mockClient, - { - repository, - readSnapshotManifest, - }, ), ) - expect(listChunks).not.toHaveBeenCalled() - expect(repository.getParseSnapshotMetadata).toHaveBeenCalledWith( - "workspace_1", - "ready", - ) - expect(readSnapshotManifest).toHaveBeenCalledWith({ - workspaceId: "workspace_1", - sourceId: "ready", - manifestKey: - "workspaces/workspace_1/sources/ready/parsed-result/manifest/current.json", + expect(listChunks).toHaveBeenCalledTimes(1) + expect(listChunks).toHaveBeenCalledWith("doc_ready", { + page: 1, + pageSize: 1, }) expect(counts).toEqual(new Map([["ready", 12]])) }) - it("skips a source count when snapshot manifest lookup fails", async () => { - const listChunks = vi.fn() + it("skips a source count when the document total lookup fails", async () => { + const listChunks = vi.fn(async () => { + throw new Error("temporary outage") + }) const mockClient = { documents: { listChunks }, } as unknown as Knowhere - const repository = { - getParseSnapshotMetadata: vi.fn(async () => ({ - resultBlobUrl: "https://blob.example/manifest/current.json", - snapshotManifestUrl: "https://blob.example/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/ready/parsed-result/manifest/current.json", - assetUrlsByFilePath: {}, - })), - } - const readSnapshotManifest = vi.fn(async () => { - throw new Error("temporary outage") - }) const { countChunksBySourceId } = await import("./counts") @@ -115,15 +75,11 @@ describe("countChunksBySourceId", () => { countChunksBySourceId( [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], mockClient, - { - repository, - readSnapshotManifest, - }, ), ) expect(counts.size).toBe(0) - expect(listChunks).not.toHaveBeenCalled() + expect(listChunks).toHaveBeenCalledTimes(1) }) it("does not count materialized demo sources through their copied document id", async () => { diff --git a/src/domains/sources/counts.ts b/src/domains/sources/counts.ts index 8f934bb..e6e0c9b 100644 --- a/src/domains/sources/counts.ts +++ b/src/domains/sources/counts.ts @@ -4,29 +4,24 @@ import { Effect } from "effect" import type Knowhere from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" -import { - readParsedResultSnapshotManifest, - type ParsedResultSnapshotManifest, -} from "./parse-result-storage-adapter" -import { sourceWorkflowRuntime } from "./workflow-runtime" -type ChunkCountRepository = Pick< - typeof sourceWorkflowRuntime, - "getParseSnapshotMetadata" -> - -type ChunkCountOptions = { - readonly readSnapshotManifest?: typeof readParsedResultSnapshotManifest - readonly repository?: ChunkCountRepository +type CountChunksClient = { + readonly documents: { + listChunks( + documentId: string, + params: { readonly page: number; readonly pageSize: number }, + ): Promise<{ + readonly pagination?: { readonly total?: number } + }> + } } export const countChunksBySourceId = ( sources: readonly Source[], client: Knowhere, - options: ChunkCountOptions = {}, ) => Effect.gen(function* () { - void client + const countClient = client as unknown as CountChunksClient const readySources = sources.filter( (source) => !source.demoKey && @@ -38,19 +33,10 @@ export const countChunksBySourceId = ( const entries = yield* Effect.all( readySources.map((source) => Effect.gen(function* () { - const manifest = yield* Effect.tryPromise(() => - loadSourceSnapshotManifest(source, options), - ).pipe( - Effect.catchAll(() => - Effect.succeed(null), - ), - ) - if (!manifest) return [source.id, undefined] as const - - return [ - source.id, - manifest.totalChunks, - ] as const + const total = yield* Effect.tryPromise(() => + loadSourceChunkCount(countClient, source.knowhereDocumentId!), + ).pipe(Effect.catchAll(() => Effect.succeed(undefined))) + return [source.id, total] as const }), ), { concurrency: "unbounded" }, @@ -67,10 +53,9 @@ export const countChunksBySourceId = ( export const sourceViewOptionsBySourceId = ( sources: readonly Source[], client: Knowhere, - options: ChunkCountOptions = {}, ) => Effect.gen(function* () { - const counts = yield* countChunksBySourceId(sources, client, options) + const counts = yield* countChunksBySourceId(sources, client) return new Map( sources.map((source) => [ source.id, @@ -79,41 +64,14 @@ export const sourceViewOptionsBySourceId = ( ) }) -async function loadSourceSnapshotManifest( - source: Source, - options: ChunkCountOptions, -): Promise { - const repository = options.repository ?? sourceWorkflowRuntime - const readSnapshotManifest = - options.readSnapshotManifest ?? readParsedResultSnapshotManifest - const snapshot = await repository.getParseSnapshotMetadata( - source.workspaceId, - source.id, - ) - if (!isCompleteSnapshot(snapshot)) return null - - return readSnapshotManifest({ - workspaceId: source.workspaceId, - sourceId: source.id, - manifestKey: snapshot.snapshotManifestKey, +async function loadSourceChunkCount( + client: CountChunksClient, + documentId: string, +): Promise { + const response = await client.documents.listChunks(documentId, { + page: 1, + pageSize: 1, }) -} - -function isCompleteSnapshot( - snapshot: - | { - readonly snapshotManifestKey?: string | null - readonly snapshotManifestUrl?: string | null - } - | null, -): snapshot is { - readonly snapshotManifestKey: string - readonly snapshotManifestUrl: string -} { - return ( - typeof snapshot?.snapshotManifestKey === "string" && - snapshot.snapshotManifestKey.length > 0 && - typeof snapshot.snapshotManifestUrl === "string" && - snapshot.snapshotManifestUrl.length > 0 - ) + const total = response.pagination?.total + return typeof total === "number" && Number.isFinite(total) ? total : undefined } diff --git a/src/domains/sources/parse-result-storage-adapter.ts b/src/domains/sources/parse-result-storage-adapter.ts deleted file mode 100644 index d195b8f..0000000 --- a/src/domains/sources/parse-result-storage-adapter.ts +++ /dev/null @@ -1,269 +0,0 @@ -import "server-only" - -import { get, head, put, BlobNotFoundError } from "@vercel/blob" -import type { - KnowhereAssetStorageObject, - KnowhereAssetStorageOptions, -} from "@ontos-ai/knowhere-sdk" - -export type ParsedResultAssetStorageInput = { - readonly workspaceId: string - readonly sourceId: string -} - -export type ParsedResultAssetIndex = { - readonly workspaceId: string - readonly sourceId: string - readonly assetUrlsByFilePath: Readonly> - readonly updatedAt: string -} - -export type ParsedResultSnapshotChunk = { - readonly id: string - readonly chunkId: string - readonly chunkType: string - readonly contentSource?: string - readonly content: string - readonly sectionPath?: string - readonly sourceChunkPath: string - readonly filePath?: string - readonly sortOrder: number - readonly metadata: Record - readonly assetUrl?: string -} - -export type ParsedResultSnapshotChunkPage = { - readonly version: 1 - readonly jobId: string - readonly documentId?: string - readonly namespace?: string - readonly sourceFileName: string - readonly page: number - readonly pageSize: number - readonly total: number - readonly totalPages: number - readonly chunks: readonly ParsedResultSnapshotChunk[] -} - -export type ParsedResultSnapshotManifest = { - readonly version: 1 - readonly kind: "knowhere-parsed-result-snapshot" - readonly jobId: string - readonly documentId?: string - readonly namespace?: string - readonly sourceFileName: string - readonly totalChunks: number - readonly chunkPageSize: number - readonly chunkPages: readonly { - readonly page: number - readonly pageSize: number - readonly chunkCount: number - readonly key: string - readonly url?: string - }[] - readonly assetUrlsByFilePath: Readonly> - readonly createdAt: string -} - -const parsedResultDirectoryName = "parsed-result" -const parsedResultAssetIndexFileName = "asset-index.json" -const parsedResultSnapshotManifestPath = "manifest/current.json" - -export function createParsedResultStorageAdapter({ - workspaceId, - sourceId, -}: ParsedResultAssetStorageInput): KnowhereAssetStorageOptions { - return { - adapter: { - async headObject(key) { - try { - const blob = await head(key) - return { - key, - contentType: blob.contentType, - contentLength: blob.size, - url: blob.url, - } - } catch (error) { - if (error instanceof BlobNotFoundError) return null - throw error - } - }, - async getObjectUrl(key) { - try { - return (await head(key)).url - } catch (error) { - if (error instanceof BlobNotFoundError) return null - throw error - } - }, - async writeObject(input: KnowhereAssetStorageObject) { - const blob = await put(input.key, Buffer.from(input.body), { - access: "public", - addRandomSuffix: false, - allowOverwrite: true, - contentType: input.contentType, - multipart: true, - }) - - return { - key: blob.pathname, - url: blob.url, - } - }, - }, - keyPrefix: getParsedResultBlobPrefix(workspaceId, sourceId), - skipExisting: true, - } -} - -export function getParsedResultSnapshotManifestKey({ - workspaceId, - sourceId, -}: ParsedResultAssetStorageInput): string { - return `${getParsedResultBlobPrefix( - workspaceId, - sourceId, - )}/${parsedResultSnapshotManifestPath}` -} - -export async function readParsedResultSnapshotManifest(input: { - readonly workspaceId: string - readonly sourceId: string - readonly manifestKey?: string | null - readonly blobStore?: ParsedResultSnapshotBlobStore -}): Promise { - const key = - input.manifestKey ?? - getParsedResultSnapshotManifestKey({ - workspaceId: input.workspaceId, - sourceId: input.sourceId, - }) - const text = await readBlobText(key, input.blobStore ?? vercelSnapshotBlobStore) - if (!text) return null - return parseSnapshotManifest(text) -} - -export async function readParsedResultSnapshotChunkPage(input: { - readonly pageKey: string - readonly blobStore?: ParsedResultSnapshotBlobStore -}): Promise { - const text = await readBlobText(input.pageKey, input.blobStore ?? vercelSnapshotBlobStore) - if (!text) return null - return parseSnapshotChunkPage(text) -} - -export async function writeParsedResultAssetIndex(input: { - readonly workspaceId: string - readonly sourceId: string - readonly assetUrlsByFilePath: Readonly> -}): Promise { - const index: ParsedResultAssetIndex = { - workspaceId: input.workspaceId, - sourceId: input.sourceId, - assetUrlsByFilePath: input.assetUrlsByFilePath, - updatedAt: new Date().toISOString(), - } - const pathname = `${getParsedResultBlobPrefix( - input.workspaceId, - input.sourceId, - )}/${parsedResultAssetIndexFileName}` - const blob = await put(pathname, JSON.stringify(index, null, 2), { - access: "public", - addRandomSuffix: false, - allowOverwrite: true, - contentType: "application/json; charset=utf-8", - }) - - return blob.url -} - -function getParsedResultBlobPrefix(workspaceId: string, sourceId: string): string { - return `workspaces/${workspaceId}/sources/${sourceId}/${parsedResultDirectoryName}` -} - -type ParsedResultSnapshotBlobGetResult = - | { - readonly statusCode: 200 - readonly stream: ReadableStream - } - | { - readonly statusCode: 304 - readonly stream: null - } - -export type ParsedResultSnapshotBlobStore = { - readonly get: ( - pathname: string, - options: { readonly access: "public" }, - ) => Promise -} - -const vercelSnapshotBlobStore: ParsedResultSnapshotBlobStore = { - get: (pathname, options) => get(pathname, options), -} - -async function readBlobText( - key: string, - blobStore: ParsedResultSnapshotBlobStore, -): Promise { - try { - const result = await blobStore.get(key, { access: "public" }) - if (!result || result.statusCode !== 200) return null - return new Response(result.stream).text() - } catch (error) { - if (error instanceof BlobNotFoundError) return null - throw error - } -} - -function parseSnapshotManifest(text: string): ParsedResultSnapshotManifest | null { - try { - const value: unknown = JSON.parse(text) - return isSnapshotManifest(value) ? value : null - } catch { - return null - } -} - -function parseSnapshotChunkPage(text: string): ParsedResultSnapshotChunkPage | null { - try { - const value: unknown = JSON.parse(text) - return isSnapshotChunkPage(value) ? value : null - } catch { - return null - } -} - -function isSnapshotManifest(value: unknown): value is ParsedResultSnapshotManifest { - if (!isRecord(value)) return false - return ( - value["kind"] === "knowhere-parsed-result-snapshot" && - value["version"] === 1 && - typeof value["jobId"] === "string" && - typeof value["sourceFileName"] === "string" && - typeof value["totalChunks"] === "number" && - typeof value["chunkPageSize"] === "number" && - Array.isArray(value["chunkPages"]) && - isRecord(value["assetUrlsByFilePath"]) && - typeof value["createdAt"] === "string" - ) -} - -function isSnapshotChunkPage(value: unknown): value is ParsedResultSnapshotChunkPage { - if (!isRecord(value)) return false - return ( - value["version"] === 1 && - typeof value["jobId"] === "string" && - typeof value["sourceFileName"] === "string" && - typeof value["page"] === "number" && - typeof value["pageSize"] === "number" && - typeof value["total"] === "number" && - typeof value["totalPages"] === "number" && - Array.isArray(value["chunks"]) - ) -} - -function isRecord(value: unknown): value is Readonly> { - return typeof value === "object" && value !== null && !Array.isArray(value) -} diff --git a/src/domains/sources/parsed-document-blob-storage.test.ts b/src/domains/sources/parsed-document-blob-storage.test.ts new file mode 100644 index 0000000..8a3c427 --- /dev/null +++ b/src/domains/sources/parsed-document-blob-storage.test.ts @@ -0,0 +1,223 @@ +import { describe, expect, it } from "vitest" +import type { + KnowhereParsedSnapshotChunkPage, + KnowhereParsedSnapshotManifest, + ParsedDocumentSyncProgress, +} from "@ontos-ai/knowhere-sdk" + +import { + BlobParsedDocumentStorage, + type ParsedDocumentBlobStore, +} from "./parsed-document-blob-storage" + +type StoredBlob = { + readonly body: string | Buffer + readonly contentType: string +} + +function createFakeBlobStore(): { + readonly store: ParsedDocumentBlobStore + readonly objects: Map +} { + const objects = new Map() + const toUrl = (pathname: string) => + `https://fake.public.blob.vercel-storage.com/${pathname}` + + const store: ParsedDocumentBlobStore = { + get: async (pathname) => { + const object = objects.get(pathname) + if (!object) return null + const body = + typeof object.body === "string" + ? object.body + : new Uint8Array(object.body) + return { + statusCode: 200, + stream: new Response(body).body as ReadableStream, + } + }, + put: async (pathname, body, options) => { + objects.set(pathname, { body, contentType: options.contentType }) + return { url: toUrl(pathname), pathname } + }, + head: async (pathname) => { + const object = objects.get(pathname) + return object ? { url: toUrl(pathname) } : null + }, + del: async (pathname) => { + objects.delete(pathname) + }, + } + + return { store, objects } +} + +const documentId = "doc_123" +const revisionKey = "job_result_1" + +const manifest: KnowhereParsedSnapshotManifest = { + version: 1, + kind: "knowhere-parsed-result-snapshot", + jobId: "job_1", + revisionKey, + documentId, + sourceFileName: "example.pdf", + totalChunks: 2, + chunkPageSize: 200, + chunkPages: [{ page: 1, pageSize: 200, chunkCount: 2, key: "chunks/page-1.json" }], + assetUrlsByFilePath: {}, + createdAt: "2026-07-04T00:00:00.000Z", +} + +const chunkPage: KnowhereParsedSnapshotChunkPage = { + version: 1, + jobId: "job_1", + revisionKey, + documentId, + sourceFileName: "example.pdf", + page: 1, + pageSize: 200, + total: 2, + totalPages: 1, + chunks: [ + { + id: "c1", + chunkId: "c1", + chunkType: "text", + content: "hello", + sourceChunkPath: "example.pdf", + sortOrder: 0, + metadata: {}, + }, + ], +} + +describe("BlobParsedDocumentStorage", () => { + it("round-trips a manifest keyed by workspace/document/revision", async () => { + const { store, objects } = createFakeBlobStore() + const storage = new BlobParsedDocumentStorage({ + workspaceId: "ws_1", + blobStore: store, + }) + + await storage.writeManifest({ documentId, revisionKey, manifest }) + + expect([...objects.keys()]).toContain( + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/manifest/current.json", + ) + const read = await storage.readManifest({ documentId, revisionKey }) + expect(read).toEqual(manifest) + }) + + it("round-trips a chunk page", async () => { + const { store } = createFakeBlobStore() + const storage = new BlobParsedDocumentStorage({ + workspaceId: "ws_1", + blobStore: store, + }) + + await storage.writeChunkPage({ documentId, revisionKey, page: chunkPage }) + const read = await storage.readChunkPage({ documentId, revisionKey, page: 1 }) + expect(read).toEqual(chunkPage) + }) + + it("round-trips sync progress", async () => { + const { store } = createFakeBlobStore() + const storage = new BlobParsedDocumentStorage({ + workspaceId: "ws_1", + blobStore: store, + }) + const progress: ParsedDocumentSyncProgress = { + documentId, + revisionKey, + nextChunkPage: 3, + nextAssetIndex: 0, + status: "running", + updatedAt: "2026-07-04T00:00:00.000Z", + } + + await storage.writeSyncProgress(progress) + const read = await storage.readSyncProgress({ documentId, revisionKey }) + expect(read).toEqual(progress) + }) + + it("writes an asset and resolves a durable URL", async () => { + const { store } = createFakeBlobStore() + const storage = new BlobParsedDocumentStorage({ + workspaceId: "ws_1", + blobStore: store, + }) + + const written = await storage.writeAsset({ + documentId, + revisionKey, + sourcePath: "images/fig-1.png", + body: new Uint8Array([1, 2, 3]), + contentType: "image/png", + }) + expect(written.url).toContain( + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/assets/images/fig-1.png", + ) + + const url = await storage.getAssetUrl({ + documentId, + revisionKey, + sourcePath: "images/fig-1.png", + }) + expect(url).toBe(written.url) + }) + + it("returns null for a missing manifest, chunk page, progress, and asset", async () => { + const { store } = createFakeBlobStore() + const storage = new BlobParsedDocumentStorage({ + workspaceId: "ws_1", + blobStore: store, + }) + + expect(await storage.readManifest({ documentId, revisionKey })).toBeNull() + expect( + await storage.readChunkPage({ documentId, revisionKey, page: 9 }), + ).toBeNull() + expect(await storage.readSyncProgress({ documentId, revisionKey })).toBeNull() + expect( + await storage.getAssetUrl({ documentId, revisionKey, sourcePath: "x.png" }), + ).toBeNull() + }) + + it("isolates artifacts by revision key so a stale revision does not read", async () => { + const { store } = createFakeBlobStore() + const storage = new BlobParsedDocumentStorage({ + workspaceId: "ws_1", + blobStore: store, + }) + + await storage.writeManifest({ documentId, revisionKey, manifest }) + const staleRead = await storage.readManifest({ + documentId, + revisionKey: "job_result_2", + }) + expect(staleRead).toBeNull() + }) + + it("rejects traversal in the document id, revision key, and asset path", async () => { + const { store } = createFakeBlobStore() + const storage = new BlobParsedDocumentStorage({ + workspaceId: "ws_1", + blobStore: store, + }) + + await expect( + storage.readManifest({ documentId: "../escape", revisionKey }), + ).rejects.toThrow(/Invalid parsed storage segment/) + await expect( + storage.readManifest({ documentId, revisionKey: "a/b" }), + ).rejects.toThrow(/Invalid parsed storage segment/) + await expect( + storage.getAssetUrl({ + documentId, + revisionKey, + sourcePath: "../../etc/passwd", + }), + ).rejects.toThrow(/Invalid parsed storage path/) + }) +}) diff --git a/src/domains/sources/parsed-document-blob-storage.ts b/src/domains/sources/parsed-document-blob-storage.ts new file mode 100644 index 0000000..23a93a0 --- /dev/null +++ b/src/domains/sources/parsed-document-blob-storage.ts @@ -0,0 +1,291 @@ +import "server-only" + +import { del, get, head, put, BlobNotFoundError } from "@vercel/blob" +import type { + KnowhereParsedSnapshotChunkPage, + KnowhereParsedSnapshotManifest, + ParsedDocumentStorage, + ParsedDocumentStorageAsset, + ParsedDocumentStorageAssetParams, + ParsedDocumentStorageChunkPageParams, + ParsedDocumentStorageDocument, + ParsedDocumentStorageManifestParams, + ParsedDocumentSyncProgress, +} from "@ontos-ai/knowhere-sdk" + +/** + * Vercel Blob backed implementation of the SDK `ParsedDocumentStorage` + * interface. The SDK only ever calls this with `{ documentId, revisionKey }`, + * so the adapter is constructed with the enclosing workspace and derives all + * blob keys from documentId + revisionKey. Keying every artifact under the + * revision key gives the SDK freshness contract for free: a manifest read at a + * revision path always carries a matching `revisionKey`. + * + * The `blobStore` seam mirrors the pattern used elsewhere in the sources + * domain so tests can inject a fake store instead of hitting Vercel Blob. + */ + +const parsedDocumentsDirectoryName = "parsed-documents" +const manifestStoragePath = "manifest/current.json" +const syncProgressStoragePath = "sync-progress.json" +const jsonContentType = "application/json; charset=utf-8" + +type BlobGetResult = + | { + readonly statusCode: 200 + readonly stream: ReadableStream + } + | { + readonly statusCode: 304 + readonly stream: null + } + +type BlobPutResult = { + readonly url: string + readonly pathname: string +} + +type BlobHeadResult = { + readonly url: string +} + +/** + * Minimal blob operations the adapter depends on. The default binds directly to + * `@vercel/blob`; tests provide an in-memory stand-in. + */ +export type ParsedDocumentBlobStore = { + readonly get: ( + pathname: string, + options: { readonly access: "public" }, + ) => Promise + readonly put: ( + pathname: string, + body: string | Buffer, + options: { + readonly access: "public" + readonly allowOverwrite: boolean + readonly contentType: string + readonly multipart?: boolean + }, + ) => Promise + readonly head: (pathname: string) => Promise + readonly del: (pathname: string) => Promise +} + +export type BlobParsedDocumentStorageInput = { + readonly workspaceId: string + readonly blobStore?: ParsedDocumentBlobStore +} + +const vercelBlobStore: ParsedDocumentBlobStore = { + get: (pathname, options) => get(pathname, options), + put: async (pathname, body, options) => { + const blob = await put(pathname, body, { + access: options.access, + addRandomSuffix: false, + allowOverwrite: options.allowOverwrite, + contentType: options.contentType, + multipart: options.multipart, + }) + return { url: blob.url, pathname: blob.pathname } + }, + head: async (pathname) => { + try { + const blob = await head(pathname) + return { url: blob.url } + } catch (error) { + if (error instanceof BlobNotFoundError) return null + throw error + } + }, + del: (pathname) => del(pathname), +} + +export class BlobParsedDocumentStorage implements ParsedDocumentStorage { + private readonly workspaceId: string + private readonly blobStore: ParsedDocumentBlobStore + + constructor(input: BlobParsedDocumentStorageInput) { + this.workspaceId = input.workspaceId + this.blobStore = input.blobStore ?? vercelBlobStore + } + + async readManifest( + params: ParsedDocumentStorageManifestParams, + ): Promise { + return this.readJson( + this.getManifestKey(params.documentId, params.revisionKey), + ) + } + + async writeManifest(params: { + readonly documentId: string + readonly revisionKey: string + readonly manifest: KnowhereParsedSnapshotManifest + }): Promise { + await this.writeJson( + this.getManifestKey(params.documentId, params.revisionKey), + params.manifest, + ) + } + + async readChunkPage( + params: ParsedDocumentStorageChunkPageParams, + ): Promise { + // chunkType filtering happens SDK-side after read; storage returns the full page. + return this.readJson( + this.getChunkPageKey(params.documentId, params.revisionKey, params.page), + ) + } + + async writeChunkPage(params: { + readonly documentId: string + readonly revisionKey: string + readonly page: KnowhereParsedSnapshotChunkPage + }): Promise { + await this.writeJson( + this.getChunkPageKey( + params.documentId, + params.revisionKey, + params.page.page, + ), + params.page, + ) + } + + async writeAsset( + params: ParsedDocumentStorageDocument & ParsedDocumentStorageAsset, + ): Promise<{ readonly sourcePath: string; readonly url?: string }> { + const blob = await this.blobStore.put( + this.getAssetKey(params.documentId, params.revisionKey, params.sourcePath), + Buffer.from(params.body), + { + access: "public", + allowOverwrite: true, + contentType: params.contentType, + multipart: true, + }, + ) + return { sourcePath: params.sourcePath, url: blob.url } + } + + async getAssetUrl( + params: ParsedDocumentStorageAssetParams, + ): Promise { + const result = await this.blobStore.head( + this.getAssetKey(params.documentId, params.revisionKey, params.sourcePath), + ) + return result?.url ?? null + } + + async readSyncProgress( + params: ParsedDocumentStorageDocument, + ): Promise { + return this.readJson( + this.getSyncProgressKey(params.documentId, params.revisionKey), + ) + } + + async writeSyncProgress(params: ParsedDocumentSyncProgress): Promise { + await this.writeJson( + this.getSyncProgressKey(params.documentId, params.revisionKey), + params, + ) + } + + private getRevisionPrefix(documentId: string, revisionKey: string): string { + return [ + "workspaces", + normalizePathSegment(this.workspaceId), + parsedDocumentsDirectoryName, + normalizePathSegment(documentId), + normalizePathSegment(revisionKey), + ].join("/") + } + + private getManifestKey(documentId: string, revisionKey: string): string { + return `${this.getRevisionPrefix(documentId, revisionKey)}/${manifestStoragePath}` + } + + private getChunkPageKey( + documentId: string, + revisionKey: string, + page: number, + ): string { + return `${this.getRevisionPrefix(documentId, revisionKey)}/chunks/page-${page}.json` + } + + private getSyncProgressKey(documentId: string, revisionKey: string): string { + return `${this.getRevisionPrefix(documentId, revisionKey)}/${syncProgressStoragePath}` + } + + private getAssetKey( + documentId: string, + revisionKey: string, + sourcePath: string, + ): string { + return `${this.getRevisionPrefix(documentId, revisionKey)}/assets/${normalizeRelativeStoragePath(sourcePath)}` + } + + private async readJson(key: string): Promise { + const text = await this.readBlobText(key) + if (text === null) return null + try { + return JSON.parse(text) as T + } catch { + return null + } + } + + private async writeJson(key: string, value: unknown): Promise { + await this.blobStore.put(key, JSON.stringify(value), { + access: "public", + allowOverwrite: true, + contentType: jsonContentType, + }) + } + + private async readBlobText(key: string): Promise { + try { + const result = await this.blobStore.get(key, { access: "public" }) + if (!result || result.statusCode !== 200) return null + return await new Response(result.stream).text() + } catch (error) { + if (error instanceof BlobNotFoundError) return null + throw error + } + } +} + +/** + * Reject path segments that could traverse outside the intended prefix. Mirrors + * the SDK `DiskParsedDocumentStorage` guard so blob keys stay well-formed. + */ +function normalizePathSegment(value: string): string { + if ( + value.length === 0 || + value.includes("\0") || + value.includes("/") || + value.includes("\\") || + value === "." || + value === ".." + ) { + throw new Error(`Invalid parsed storage segment: ${value}`) + } + return value +} + +function normalizeRelativeStoragePath(value: string): string { + const normalized = value.replaceAll("\\", "/").replace(/^\.\/+/, "") + if ( + normalized.length === 0 || + normalized.includes("\0") || + normalized.startsWith("/") || + normalized + .split("/") + .some((part) => part.length === 0 || part === "." || part === "..") + ) { + throw new Error(`Invalid parsed storage path: ${value}`) + } + return normalized +} diff --git a/src/domains/sources/parsed-document-sync-scheduler.ts b/src/domains/sources/parsed-document-sync-scheduler.ts new file mode 100644 index 0000000..5d2957d --- /dev/null +++ b/src/domains/sources/parsed-document-sync-scheduler.ts @@ -0,0 +1,154 @@ +import "server-only" + +import { Client } from "@upstash/workflow" +import type { ParsedDocumentSyncScheduler } from "@ontos-ai/knowhere-sdk" + +import { logger } from "@/lib/logger" + +/** + * A `ParsedDocumentSyncScheduler` whose `schedule` enqueues a durable QStash + * continuation instead of running the (non-serializable) task closure inline. + * The SDK schedules a background sync on every storage read-miss; the closure + * captures a per-request client we cannot serialize across serverless + * invocations, so we discard it and trigger `/api/sources/parsed-sync`, which + * rebuilds a parsed-storage client and loops `syncParsedDocument` to completion. + * + * The scheduler is pre-bound to a single `{ workspaceId, sourceId, documentId }` + * because the SDK invokes `schedule(task)` with no arguments. `revisionKey` is + * intentionally NOT part of the bound identity — the sync route resolves the + * current revision itself so a scheduler bound before the revision is known + * still enqueues correctly. + */ + +export type ParsedSyncTrigger = (input: { + readonly url: string + readonly body: ParsedSyncPayload + readonly workflowRunId: string +}) => Promise + +export type ParsedSyncPayload = { + readonly workspaceId: string + readonly sourceId: string + readonly documentId: string + readonly apiKey: string + readonly revisionKey?: string + readonly segmentIndex?: number +} + +export type CreateParsedDocumentSyncSchedulerInput = { + readonly workspaceId: string + readonly sourceId: string + readonly documentId: string + readonly apiKey: string + readonly revisionKey?: string + readonly trigger?: ParsedSyncTrigger +} + +// Bounded per-document guard: the SDK may schedule on every read-miss within a +// single request. The cooldown collapses duplicate same-process enqueues +// without permanently blocking future syncs. +const triggerCooldownMs: number = 60_000 +const lastTriggeredAtByKey: Map = new Map() + +function resolveBaseURL(): string { + return process.env.NOTEBOOK_PUBLIC_URL ?? "http://localhost:3000" +} + +export function getParsedSyncWorkflowRunId(input: { + readonly documentId: string + readonly revisionKey: string + readonly segmentIndex: number +}): string { + return `${input.documentId}-sync-${input.revisionKey}-${input.segmentIndex}` +} + +const defaultTrigger: ParsedSyncTrigger = async (input) => { + const token = process.env.QSTASH_TOKEN + if (!token) { + throw new Error("QSTASH_TOKEN is required to schedule parsed document sync.") + } + await new Client({ token }).trigger({ + url: input.url, + body: input.body, + workflowRunId: input.workflowRunId, + retries: 3, + }) +} + +export function getParsedSyncUrl(): string { + return `${resolveBaseURL()}/api/sources/parsed-sync` +} + +/** + * Enqueue a durable parsed-document sync from outside the SDK read path — used + * by the reconcile workflow to hand off to the resumable parsed-sync route, and + * by retry to resume a failed storage sync. Starts at segment 0. + */ +export async function enqueueParsedDocumentSync( + input: { + readonly workspaceId: string + readonly sourceId: string + readonly documentId: string + readonly apiKey: string + readonly revisionKey?: string + }, + trigger: ParsedSyncTrigger = defaultTrigger, +): Promise { + await trigger({ + url: getParsedSyncUrl(), + body: { ...input, segmentIndex: 0 }, + workflowRunId: getParsedSyncWorkflowRunId({ + documentId: input.documentId, + revisionKey: input.revisionKey ?? "initial", + segmentIndex: 0, + }), + }) +} + +export function createParsedDocumentSyncScheduler( + input: CreateParsedDocumentSyncSchedulerInput, +): ParsedDocumentSyncScheduler { + const trigger = input.trigger ?? defaultTrigger + const cooldownKey = `${input.documentId}|${input.revisionKey ?? "unknown"}` + + return { + schedule: () => { + const now = Date.now() + const lastTriggeredAt = lastTriggeredAtByKey.get(cooldownKey) + if ( + lastTriggeredAt !== undefined && + now - lastTriggeredAt < triggerCooldownMs + ) { + return + } + lastTriggeredAtByKey.set(cooldownKey, now) + + const url = `${resolveBaseURL()}/api/sources/parsed-sync` + const workflowRunId = getParsedSyncWorkflowRunId({ + documentId: input.documentId, + revisionKey: input.revisionKey ?? "initial", + segmentIndex: 0, + }) + void trigger({ + url, + body: { + workspaceId: input.workspaceId, + sourceId: input.sourceId, + documentId: input.documentId, + apiKey: input.apiKey, + revisionKey: input.revisionKey, + segmentIndex: 0, + }, + workflowRunId, + }).catch((error: unknown) => { + lastTriggeredAtByKey.delete(cooldownKey) + logger.error("parsed-sync-scheduler: failed to enqueue sync", { + workspaceId: input.workspaceId, + sourceId: input.sourceId, + documentId: input.documentId, + error: error instanceof Error ? error.message : String(error), + }) + }) + }, + } +} diff --git a/src/domains/sources/parsed-sync-route-workflow.test.ts b/src/domains/sources/parsed-sync-route-workflow.test.ts new file mode 100644 index 0000000..555c186 --- /dev/null +++ b/src/domains/sources/parsed-sync-route-workflow.test.ts @@ -0,0 +1,201 @@ +import { afterEach, describe, expect, it, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + makeKnowhereClientWithParsedStorage: vi.fn(), + updateSyncStatus: vi.fn(), + findInWorkspace: vi.fn(), + markFailed: vi.fn(), + markReady: vi.fn(), + markSourceReadyAfterReconciliation: vi.fn(), + loggerInfo: vi.fn(), + loggerError: vi.fn(), +})) + +vi.mock("@/integrations/knowhere", () => ({ + makeKnowhereClientWithParsedStorage: mocks.makeKnowhereClientWithParsedStorage, +})) + +vi.mock("./workflow-runtime", () => ({ + sourceWorkflowRuntime: { + updateSyncStatus: mocks.updateSyncStatus, + findInWorkspace: mocks.findInWorkspace, + markFailed: mocks.markFailed, + markReady: mocks.markReady, + }, +})) + +vi.mock("./source-reconcile-workflow", () => ({ + markSourceReadyAfterReconciliation: mocks.markSourceReadyAfterReconciliation, +})) + +vi.mock("@/lib/logger", () => ({ + logger: { info: mocks.loggerInfo, error: mocks.loggerError }, +})) + +import { parsedSyncRouteWorkflow } from "./parsed-sync-route-workflow" + +type RunStep = (id: string, task: () => Promise | T) => Promise + +function createContext(overrides: { url?: string } = {}) { + const run: RunStep = async (_id, task) => task() + return { + run, + url: overrides.url ?? "https://notebook.example/api/sources/parsed-sync", + } +} + +const basePayload = { + workspaceId: "workspace_1", + sourceId: "source_1", + documentId: "doc_1", + apiKey: "key_1", + segmentIndex: 0, +} + +describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { + afterEach(() => { + vi.clearAllMocks() + }) + + it("marks the source ready when sync completes in one segment", async () => { + const syncParsedDocument = vi.fn(async () => ({ + documentId: "doc_1", + revisionKey: "rev_1", + completed: true, + })) + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: {}, + knowledge: { syncParsedDocument }, + }) + mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ + status: "ready", + }) + + await parsedSyncRouteWorkflow.runParsedSyncWorkflow({ + context: createContext(), + payload: basePayload, + }) + + expect(syncParsedDocument).toHaveBeenCalledTimes(1) + expect(syncParsedDocument).toHaveBeenCalledWith({ documentId: "doc_1" }) + expect(mocks.updateSyncStatus).toHaveBeenCalledWith( + "workspace_1", + "source_1", + { revisionKey: "rev_1", syncStatus: "completed" }, + ) + expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "source_1", + documentId: "doc_1", + }) + }) + + it("triggers a continuation and does not mark ready when sync is incomplete", async () => { + const syncParsedDocument = vi.fn(async () => ({ + documentId: "doc_1", + revisionKey: "rev_1", + completed: false, + })) + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: {}, + knowledge: { syncParsedDocument }, + }) + const triggered: Array<{ workflowRunId: string; segmentIndex?: number }> = [] + const restore = parsedSyncRouteWorkflow.setContinuationTriggerForTesting( + async (input) => { + triggered.push({ + workflowRunId: input.workflowRunId, + segmentIndex: input.payload.segmentIndex, + }) + }, + ) + + try { + await parsedSyncRouteWorkflow.runParsedSyncWorkflow({ + context: createContext(), + payload: basePayload, + }) + } finally { + restore() + } + + expect(mocks.markSourceReadyAfterReconciliation).not.toHaveBeenCalled() + expect(triggered).toHaveLength(1) + expect(triggered[0]?.segmentIndex).toBe(1) + expect(triggered[0]?.workflowRunId).toBe("doc_1-sync-rev_1-1") + expect(mocks.updateSyncStatus).toHaveBeenLastCalledWith( + "workspace_1", + "source_1", + { revisionKey: "rev_1", syncStatus: "running" }, + ) + }) + + it("passes an explicit revisionKey into syncParsedDocument when provided", async () => { + const syncParsedDocument = vi.fn(async () => ({ + documentId: "doc_1", + revisionKey: "rev_9", + completed: true, + })) + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: {}, + knowledge: { syncParsedDocument }, + }) + mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ + status: "ready", + }) + + await parsedSyncRouteWorkflow.runParsedSyncWorkflow({ + context: createContext(), + payload: { ...basePayload, revisionKey: "rev_9" }, + }) + + expect(syncParsedDocument).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "rev_9", + }) + }) +}) + +describe("parsedSyncRouteWorkflow.markSyncFailedAfterWorkflowFailure", () => { + afterEach(() => { + vi.clearAllMocks() + }) + + it("fails a parsing source with failure_stage storage_sync", async () => { + mocks.findInWorkspace.mockResolvedValue({ id: "source_1", status: "parsing" }) + + await parsedSyncRouteWorkflow.markSyncFailedAfterWorkflowFailure( + { ...basePayload, revisionKey: "rev_1" }, + "boom", + ) + + expect(mocks.updateSyncStatus).toHaveBeenCalledWith( + "workspace_1", + "source_1", + { revisionKey: "rev_1", syncStatus: "failed", syncError: "boom" }, + ) + expect(mocks.markFailed).toHaveBeenCalledWith( + "workspace_1", + "source_1", + expect.stringContaining("storage sync failed"), + "parsing", + "storage_sync", + ) + }) + + it("does not fail an already-ready source, only records sync_status", async () => { + mocks.findInWorkspace.mockResolvedValue({ id: "source_1", status: "ready" }) + + await parsedSyncRouteWorkflow.markSyncFailedAfterWorkflowFailure( + basePayload, + "boom", + ) + + expect(mocks.updateSyncStatus).toHaveBeenCalledWith( + "workspace_1", + "source_1", + { revisionKey: undefined, syncStatus: "failed", syncError: "boom" }, + ) + expect(mocks.markFailed).not.toHaveBeenCalled() + }) +}) diff --git a/src/domains/sources/parsed-sync-route-workflow.ts b/src/domains/sources/parsed-sync-route-workflow.ts new file mode 100644 index 0000000..c33ace8 --- /dev/null +++ b/src/domains/sources/parsed-sync-route-workflow.ts @@ -0,0 +1,226 @@ +import "server-only" + +import { Client, type WorkflowContext } from "@upstash/workflow" +import type { KnowledgeSyncParsedDocumentResponse } from "@ontos-ai/knowhere-sdk" + +import { makeKnowhereClientWithParsedStorage } from "@/integrations/knowhere" +import { logger } from "@/lib/logger" +import { + getParsedSyncWorkflowRunId, + type ParsedSyncPayload, +} from "./parsed-document-sync-scheduler" +import { markSourceReadyAfterReconciliation } from "./source-reconcile-workflow" +import { sourceWorkflowRuntime } from "./workflow-runtime" + +type ParsedSyncWorkflowContext = Pick< + WorkflowContext, + "run" | "url" +> + +type NormalizedParsedSyncPayload = { + readonly workspaceId: string + readonly sourceId: string + readonly documentId: string + readonly apiKey: string + readonly revisionKey?: string + readonly segmentIndex: number +} + +type ContinuationTriggerInput = { + readonly url: string + readonly payload: ParsedSyncPayload + readonly workflowRunId: string +} + +// Sync steps per workflow segment. Each `syncParsedDocument` call is bounded by +// the SDK limits (pages + deadline); this caps how many bounded steps we run in +// one serverless invocation before handing off to a fresh continuation. +const maxSyncStepsPerSegment = 4 + +let triggerContinuation: typeof triggerParsedSyncContinuation = + triggerParsedSyncContinuation + +function normalizeParsedSyncPayload( + payload: ParsedSyncPayload, +): NormalizedParsedSyncPayload { + return { + workspaceId: payload.workspaceId, + sourceId: payload.sourceId, + documentId: payload.documentId, + apiKey: payload.apiKey, + revisionKey: payload.revisionKey, + segmentIndex: + typeof payload.segmentIndex === "number" && + Number.isInteger(payload.segmentIndex) && + payload.segmentIndex >= 0 + ? payload.segmentIndex + : 0, + } +} + +async function runParsedSyncWorkflow(input: { + readonly context: ParsedSyncWorkflowContext + readonly payload: NormalizedParsedSyncPayload +}): Promise { + const { context, payload } = input + const { workspaceId, sourceId, documentId, apiKey } = payload + const { knowledge } = makeKnowhereClientWithParsedStorage(apiKey, { + workspaceId, + }) + + let revisionKey = payload.revisionKey + let completed = false + + for (let step = 0; step < maxSyncStepsPerSegment; step++) { + const result: KnowledgeSyncParsedDocumentResponse = await context.run( + `sync-${payload.segmentIndex}-${step}`, + async () => + knowledge.syncParsedDocument({ + documentId, + ...(revisionKey ? { revisionKey } : {}), + }), + ) + revisionKey = result.revisionKey + + await context.run(`record-progress-${payload.segmentIndex}-${step}`, () => + sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { + revisionKey, + syncStatus: result.completed ? "completed" : "running", + }), + ) + + if (result.completed) { + completed = true + break + } + } + + if (!completed) { + const nextSegmentIndex = payload.segmentIndex + 1 + await context.run( + `trigger-sync-continuation-${nextSegmentIndex}`, + async () => + triggerContinuation({ + url: context.url, + payload: { + workspaceId, + sourceId, + documentId, + apiKey, + revisionKey, + segmentIndex: nextSegmentIndex, + }, + workflowRunId: getParsedSyncWorkflowRunId({ + documentId, + revisionKey: revisionKey ?? "initial", + segmentIndex: nextSegmentIndex, + }), + }), + ) + logger.info("parsed-sync: continuation triggered", { + sourceId, + documentId, + segmentIndex: nextSegmentIndex, + }) + return + } + + // A parsing source is only readied once its parsed snapshot is fully synced. + // For an already-ready source (read-miss backfill) markReady is a no-op guard. + const ready = await context.run("source-ready", async () => + markSourceReadyAfterReconciliation({ + workspaceId, + sourceId, + documentId, + }), + ) + logger.info("parsed-sync: parsed document sync finished", { + sourceId, + documentId, + revisionKey, + status: ready.status, + }) +} + +async function triggerParsedSyncContinuation( + input: ContinuationTriggerInput, +): Promise { + const token = process.env.QSTASH_TOKEN + if (!token) { + throw new Error("QSTASH_TOKEN is required to continue parsed document sync.") + } + + await new Client({ token }).trigger({ + url: input.url, + body: input.payload, + workflowRunId: input.workflowRunId, + retries: 3, + }) +} + +async function markSyncFailedAfterWorkflowFailure( + payload: ParsedSyncPayload, + failResponse: string, +): Promise { + const normalized = normalizeParsedSyncPayload(payload) + const reason = getSafeFailureReason(failResponse) + + // Record the storage-sync failure. A source still `parsing` is failed with + // failure_stage=storage_sync so a retry resumes sync without reparsing; an + // already-ready source is left ready (it still serves via remote fallback), + // only its sync_status is marked failed for observability. + await sourceWorkflowRuntime.updateSyncStatus( + normalized.workspaceId, + normalized.sourceId, + { + revisionKey: normalized.revisionKey, + syncStatus: "failed", + syncError: reason, + }, + ) + + const source = await sourceWorkflowRuntime.findInWorkspace( + normalized.workspaceId, + normalized.sourceId, + ) + if (source?.status === "parsing") { + await sourceWorkflowRuntime.markFailed( + normalized.workspaceId, + normalized.sourceId, + `Parsed document storage sync failed: ${reason}`, + "parsing", + "storage_sync", + ) + } + + logger.error("parsed-sync: marked sync failed after workflow failure", { + workspaceId: normalized.workspaceId, + sourceId: normalized.sourceId, + documentId: normalized.documentId, + segmentIndex: normalized.segmentIndex, + sourceStatus: source?.status, + }) +} + +function getSafeFailureReason(value: string): string { + const normalized = value.replace(/\s+/g, " ").trim() + if (normalized.length === 0) return "retry attempts were exhausted." + return normalized.slice(0, 500) +} + +function setContinuationTriggerForTesting( + trigger: typeof triggerParsedSyncContinuation, +): () => void { + const previous = triggerContinuation + triggerContinuation = trigger + return () => { + triggerContinuation = previous + } +} + +export const parsedSyncRouteWorkflow = { + markSyncFailedAfterWorkflowFailure, + normalizeParsedSyncPayload, + runParsedSyncWorkflow, + setContinuationTriggerForTesting, +} diff --git a/src/domains/sources/reconcile.test.ts b/src/domains/sources/reconcile.test.ts index 4c90172..5be0cea 100644 --- a/src/domains/sources/reconcile.test.ts +++ b/src/domains/sources/reconcile.test.ts @@ -20,6 +20,7 @@ function makeSource(overrides: Partial): Source { sizeBytes: 1, status: "parsing", failureReason: null, + failureStage: null, knowhereJobId: "job_1", knowhereDocumentId: null, stagedBlobPathname: null, diff --git a/src/domains/sources/remote-parsed-snapshot-sync.test.ts b/src/domains/sources/remote-parsed-snapshot-sync.test.ts deleted file mode 100644 index 6a91582..0000000 --- a/src/domains/sources/remote-parsed-snapshot-sync.test.ts +++ /dev/null @@ -1,316 +0,0 @@ -import { afterEach, describe, expect, it, vi } from "vitest" - -const mocks = vi.hoisted(() => ({ - createParsedResultStorageAdapter: vi.fn(), - loggerInfo: vi.fn(), -})) - -vi.mock("./parse-result-storage-adapter", () => ({ - createParsedResultStorageAdapter: mocks.createParsedResultStorageAdapter, -})) - -vi.mock("@/lib/logger", () => ({ - logger: { - info: mocks.loggerInfo, - }, -})) - -import type { Source, SourceParseResult } from "@/infrastructure/db/schema" -import { syncRemoteParsedSnapshot } from "./remote-parsed-snapshot-sync" - -describe("syncRemoteParsedSnapshot", () => { - afterEach(() => { - vi.clearAllMocks() - }) - - it("loads an already-parsed remote document into the Notebook Blob snapshot", async () => { - const storageAdapter = { - adapter: { - writeObject: vi.fn(), - }, - keyPrefix: - "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result", - } - const listChunks = vi.fn(async () => ({ - jobResultId: "job_result_1", - })) - const loadJobResult = vi.fn(async () => ({ - assetUrlsByFilePath: { - "pages/page-1.png": "https://blob.example/pages/page-1.png", - }, - parsedSnapshot: { - manifestKey: - "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - manifestUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - }, - })) - const repository = createRepository() - mocks.createParsedResultStorageAdapter.mockReturnValue(storageAdapter) - - const snapshot = await syncRemoteParsedSnapshot({ - workspaceId: "workspace_1", - source: makeSource({ - knowhereJobId: null, - }), - client: { - documents: { - listChunks, - }, - knowledge: { - loadJobResult, - }, - }, - repository, - }) - - expect(listChunks).toHaveBeenCalledWith("doc_remote", { - page: 1, - pageSize: 1, - includeAssetUrls: false, - }) - expect(mocks.createParsedResultStorageAdapter).toHaveBeenCalledWith({ - workspaceId: "workspace_1", - sourceId: "00000000-0000-0000-0000-000000000009", - }) - expect(loadJobResult).toHaveBeenCalledWith({ - jobId: "job_result_1", - storageAdapter, - }) - expect(repository.markParsing).toHaveBeenCalledWith( - "workspace_1", - "00000000-0000-0000-0000-000000000009", - "job_result_1", - "doc_remote", - ) - expect(repository.saveParseResult).toHaveBeenCalledWith( - "workspace_1", - "00000000-0000-0000-0000-000000000009", - { - resultBlobUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - snapshotManifestUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - assetUrlsByFilePath: { - "pages/page-1.png": "https://blob.example/pages/page-1.png", - }, - }, - ) - expect(repository.markReady).toHaveBeenCalledWith( - "workspace_1", - "00000000-0000-0000-0000-000000000009", - "doc_remote", - ) - expect(snapshot).toEqual({ - resultBlobUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - snapshotManifestUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - assetUrlsByFilePath: { - "pages/page-1.png": "https://blob.example/pages/page-1.png", - }, - }) - }) - - it("returns the existing complete snapshot without calling Knowhere", async () => { - const existingSnapshot = { - resultBlobUrl: - "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - snapshotManifestUrl: - "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - assetUrlsByFilePath: {}, - } - const repository = createRepository({ - getParseSnapshotMetadata: vi.fn(async () => existingSnapshot), - }) - const listChunks = vi.fn() - const loadJobResult = vi.fn() - - const snapshot = await syncRemoteParsedSnapshot({ - workspaceId: "workspace_1", - source: makeSource(), - client: { - documents: { - listChunks, - }, - knowledge: { - loadJobResult, - }, - }, - repository, - }) - - expect(snapshot).toEqual(existingSnapshot) - expect(listChunks).not.toHaveBeenCalled() - expect(loadJobResult).not.toHaveBeenCalled() - expect(repository.markParsing).not.toHaveBeenCalled() - expect(repository.markReady).not.toHaveBeenCalled() - }) - - it("marks a parsing source ready when a complete snapshot already exists", async () => { - const existingSnapshot = { - resultBlobUrl: - "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - snapshotManifestUrl: - "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - assetUrlsByFilePath: {}, - } - const repository = createRepository({ - getParseSnapshotMetadata: vi.fn(async () => existingSnapshot), - }) - const listChunks = vi.fn() - const loadJobResult = vi.fn() - - const snapshot = await syncRemoteParsedSnapshot({ - workspaceId: "workspace_1", - source: makeSource({ - status: "parsing", - knowhereDocumentId: "doc_remote", - }), - client: { - documents: { - listChunks, - }, - knowledge: { - loadJobResult, - }, - }, - repository, - }) - - expect(snapshot).toEqual(existingSnapshot) - expect(listChunks).not.toHaveBeenCalled() - expect(loadJobResult).not.toHaveBeenCalled() - expect(repository.markReady).toHaveBeenCalledWith( - "workspace_1", - "00000000-0000-0000-0000-000000000009", - "doc_remote", - ) - }) - - it("retries an interrupted remote snapshot sync from a parsing source row", async () => { - const storageAdapter = { - adapter: { - writeObject: vi.fn(), - }, - keyPrefix: - "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result", - } - const listChunks = vi.fn() - const loadJobResult = vi.fn(async () => ({ - assetUrlsByFilePath: {}, - parsedSnapshot: { - manifestKey: - "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - manifestUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - }, - })) - const repository = createRepository() - mocks.createParsedResultStorageAdapter.mockReturnValue(storageAdapter) - - await syncRemoteParsedSnapshot({ - workspaceId: "workspace_1", - source: makeSource({ - status: "parsing", - knowhereJobId: "job_result_1", - }), - client: { - documents: { - listChunks, - }, - knowledge: { - loadJobResult, - }, - }, - repository, - }) - - expect(listChunks).not.toHaveBeenCalled() - expect(loadJobResult).toHaveBeenCalledWith({ - jobId: "job_result_1", - storageAdapter, - }) - expect(repository.markReady).toHaveBeenCalledWith( - "workspace_1", - "00000000-0000-0000-0000-000000000009", - "doc_remote", - ) - }) -}) - -type SyncRemoteParsedSnapshotInput = Parameters< - typeof syncRemoteParsedSnapshot ->[0] - -type TestRemoteParsedSnapshotRepository = NonNullable< - SyncRemoteParsedSnapshotInput["repository"] -> - -function createRepository( - overrides: Partial = {}, -): TestRemoteParsedSnapshotRepository { - return { - ...createRepositoryShape(), - ...overrides, - } -} - -function createRepositoryShape(): TestRemoteParsedSnapshotRepository { - return { - getParseSnapshotMetadata: vi.fn(async () => null), - markParsing: vi.fn(async () => makeSource({ status: "parsing" })), - markReady: vi.fn(async () => makeSource()), - saveParseResult: vi.fn(async () => makeSourceParseResult()), - } -} - -function makeSource(overrides: Partial = {}): Source { - return { - id: "00000000-0000-0000-0000-000000000009", - workspaceId: "workspace_1", - title: "remote.pdf", - mimeType: "application/pdf", - sizeBytes: 0, - status: "ready", - failureReason: null, - knowhereJobId: "job_result_1", - knowhereDocumentId: "doc_remote", - stagedBlobPathname: null, - stagedBlobUrl: null, - originalBlobPathname: null, - originalBlobUrl: null, - demoKey: null, - createdAt: new Date("2026-05-10T00:00:00.000Z"), - updatedAt: new Date("2026-05-10T00:00:00.000Z"), - deletedAt: null, - ...overrides, - } -} - -function makeSourceParseResult( - overrides: Partial = {}, -): SourceParseResult { - return { - id: "parse_result_1", - sourceId: "00000000-0000-0000-0000-000000000009", - resultBlobUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - snapshotManifestUrl: - "https://blob.example/workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/00000000-0000-0000-0000-000000000009/parsed-result/manifest/current.json", - assetUrls: {}, - createdAt: new Date("2026-05-10T00:00:00.000Z"), - updatedAt: new Date("2026-05-10T00:00:00.000Z"), - ...overrides, - } -} diff --git a/src/domains/sources/remote-parsed-snapshot-sync.ts b/src/domains/sources/remote-parsed-snapshot-sync.ts deleted file mode 100644 index 3e939c4..0000000 --- a/src/domains/sources/remote-parsed-snapshot-sync.ts +++ /dev/null @@ -1,210 +0,0 @@ -import "server-only" - -import { Effect } from "effect" - -import type { Source } from "@/infrastructure/db/schema" -import { logger } from "@/lib/logger" -import { createParsedResultStorageAdapter } from "./parse-result-storage-adapter" -import { sourceWorkflowRuntime } from "./workflow-runtime" - -type RemoteParsedSnapshotSource = Pick< - Source, - "id" | "workspaceId" | "knowhereDocumentId" | "knowhereJobId" | "status" -> - -type RemoteParsedSnapshotClient = { - readonly documents: { - listChunks( - documentId: string, - params: { - readonly page: number - readonly pageSize: number - readonly includeAssetUrls: boolean - }, - ): Promise<{ - readonly jobId?: string | null - readonly jobResultId?: string | null - }> - } - readonly knowledge: { - loadJobResult(params: { - readonly jobId: string - readonly storageAdapter: ReturnType - }): Promise - } -} - -type RemoteParsedSnapshotResponse = { - readonly assetUrlsByFilePath?: Readonly> - readonly parsedSnapshot?: { - readonly manifestKey: string - readonly manifestUrl?: string - } -} - -type RemoteParsedSnapshotRepository = { - readonly getParseSnapshotMetadata: typeof sourceWorkflowRuntime.getParseSnapshotMetadata - readonly markParsing: typeof sourceWorkflowRuntime.markParsing - readonly markReady: typeof sourceWorkflowRuntime.markReady - readonly saveParseResult: typeof sourceWorkflowRuntime.saveParseResult -} - -type SyncRemoteParsedSnapshotInput = { - readonly workspaceId: string - readonly source: RemoteParsedSnapshotSource - readonly client: RemoteParsedSnapshotClient - readonly repository?: RemoteParsedSnapshotRepository -} - -type RemoteParsedSnapshotReadModel = { - readonly resultBlobUrl: string - readonly snapshotManifestUrl?: string | null - readonly snapshotManifestKey?: string | null - readonly assetUrlsByFilePath: Readonly> -} - -const jobProbePage = 1 -const jobProbePageSize = 1 - -const syncRemoteParsedSnapshotEffect = Effect.fn( - "syncRemoteParsedSnapshot", -)(function* ({ - workspaceId, - source, - client, - repository = sourceWorkflowRuntime, -}: SyncRemoteParsedSnapshotInput) { - const existingSnapshot = yield* Effect.tryPromise(() => - repository.getParseSnapshotMetadata(workspaceId, source.id), - ) - if (isCompleteSnapshot(existingSnapshot)) { - if (source.status === "parsing" && source.knowhereDocumentId) { - yield* Effect.tryPromise(() => - repository.markReady(workspaceId, source.id, source.knowhereDocumentId!), - ) - } - return existingSnapshot - } - - if (source.status !== "ready" && source.status !== "parsing") { - return existingSnapshot - } - - const documentId = source.knowhereDocumentId - if (!documentId) { - return yield* Effect.die( - new Error("Remote source is missing a Knowhere document id."), - ) - } - - const jobId = yield* resolveRemoteDocumentJobId({ - client, - documentId, - source, - }) - yield* Effect.tryPromise(() => - repository.markParsing(workspaceId, source.id, jobId, documentId), - ) - - const snapshot = yield* Effect.tryPromise(async () => { - const loadedResult = await client.knowledge.loadJobResult({ - jobId, - storageAdapter: createParsedResultStorageAdapter({ - workspaceId, - sourceId: source.id, - }), - }) - const manifest = loadedResult.parsedSnapshot - if (!manifest?.manifestUrl || !manifest.manifestKey) { - throw new Error( - "Remote parsed snapshot was not written; refusing to mark source ready.", - ) - } - return { - assetUrlsByFilePath: loadedResult.assetUrlsByFilePath ?? {}, - snapshotManifestUrl: manifest.manifestUrl, - snapshotManifestKey: manifest.manifestKey, - } - }) - - const savedSnapshot = yield* Effect.tryPromise(() => - repository.saveParseResult(workspaceId, source.id, { - resultBlobUrl: snapshot.snapshotManifestUrl, - snapshotManifestUrl: snapshot.snapshotManifestUrl, - snapshotManifestKey: snapshot.snapshotManifestKey, - assetUrlsByFilePath: snapshot.assetUrlsByFilePath, - }), - ) - if (!savedSnapshot) { - return yield* Effect.die( - new Error("Remote parsed snapshot could not be saved for the source."), - ) - } - - yield* Effect.tryPromise(() => - repository.markReady(workspaceId, source.id, documentId), - ) - - logger.info("sources: remote parsed snapshot synced", { - sourceId: source.id, - documentId, - jobId, - snapshotManifestKey: snapshot.snapshotManifestKey, - }) - - return { - resultBlobUrl: snapshot.snapshotManifestUrl, - snapshotManifestUrl: snapshot.snapshotManifestUrl, - snapshotManifestKey: snapshot.snapshotManifestKey, - assetUrlsByFilePath: snapshot.assetUrlsByFilePath, - } -}) - -export async function syncRemoteParsedSnapshot( - input: SyncRemoteParsedSnapshotInput, -): Promise { - return Effect.runPromise(syncRemoteParsedSnapshotEffect(input)) -} - -const resolveRemoteDocumentJobId = Effect.fn("resolveRemoteDocumentJobId")( - function* ({ - client, - documentId, - source, - }: { - readonly client: RemoteParsedSnapshotClient - readonly documentId: string - readonly source: RemoteParsedSnapshotSource - }) { - if (source.knowhereJobId) return source.knowhereJobId - - const response = yield* Effect.tryPromise(() => - client.documents.listChunks(documentId, { - page: jobProbePage, - pageSize: jobProbePageSize, - includeAssetUrls: false, - }), - ) - const jobId = response.jobId ?? response.jobResultId - if (!jobId) { - return yield* Effect.die( - new Error("Remote document chunk metadata did not include a job id."), - ) - } - return jobId - }, -) - -function isCompleteSnapshot( - snapshot: RemoteParsedSnapshotReadModel | null, -): snapshot is RemoteParsedSnapshotReadModel & { - readonly snapshotManifestKey: string - readonly snapshotManifestUrl: string -} { - return ( - typeof snapshot?.snapshotManifestKey === "string" && - snapshot.snapshotManifestKey.length > 0 && - typeof snapshot.snapshotManifestUrl === "string" && - snapshot.snapshotManifestUrl.length > 0 - ) -} diff --git a/src/domains/sources/repository.ts b/src/domains/sources/repository.ts index 4d9c061..3c6d01c 100644 --- a/src/domains/sources/repository.ts +++ b/src/domains/sources/repository.ts @@ -23,6 +23,7 @@ type SourceRepository = { readonly getParseResultProgressEffect: typeof sourceParseResultRepository.getParseResultProgressEffect readonly getParseSnapshotMetadataEffect: typeof sourceParseResultRepository.getParseSnapshotMetadataEffect readonly getParseAssetUrlsEffect: typeof sourceParseResultRepository.getParseAssetUrlsEffect + readonly updateSyncStatusEffect: typeof sourceParseResultRepository.updateSyncStatusEffect } export const sourceRepository: SourceRepository = { @@ -49,4 +50,5 @@ export const sourceRepository: SourceRepository = { getParseSnapshotMetadataEffect: sourceParseResultRepository.getParseSnapshotMetadataEffect, getParseAssetUrlsEffect: sourceParseResultRepository.getParseAssetUrlsEffect, + updateSyncStatusEffect: sourceParseResultRepository.updateSyncStatusEffect, } diff --git a/src/domains/sources/retry.test.ts b/src/domains/sources/retry.test.ts index 86c03a1..19eb565 100644 --- a/src/domains/sources/retry.test.ts +++ b/src/domains/sources/retry.test.ts @@ -18,6 +18,7 @@ describe("retrySourceToKnowhereEffect", () => { const parsingSource = makeSource({ status: "parsing", failureReason: null, + failureStage: null, knowhereJobId: "job_retry", knowhereDocumentId: "doc_retry", }) @@ -123,6 +124,7 @@ function makeSource(overrides: Partial = {}): Source { sizeBytes: 5, status: "failed", failureReason: "Knowhere upload failed.", + failureStage: null, knowhereJobId: null, knowhereDocumentId: null, stagedBlobPathname: null, diff --git a/src/domains/sources/route-chunks.ts b/src/domains/sources/route-chunks.ts index b8fbca2..7d97920 100644 --- a/src/domains/sources/route-chunks.ts +++ b/src/domains/sources/route-chunks.ts @@ -1,6 +1,8 @@ import { Effect } from "effect" import { demoView } from "@/domains/demo/view" +import { readAllSourceChunks, readSourceChunkPage } from "@/domains/chunks/read" +import { resolveChunkConnectionTargets } from "@/domains/chunks" import type { DemoChunkPage } from "@/integrations/knowhere-demo" import { logger } from "@/lib/logger" import { routeResult } from "@/lib/route-result" @@ -8,7 +10,10 @@ import { decodeRemoteSourceId, findRemoteLibraryDocumentBySourceId, } from "./remote-library" -import { getClientForWorkspace } from "./route-dependencies" +import { + getClientForWorkspace, + getKnowledgeForSource, +} from "./route-dependencies" import { sourceRowRepository } from "./source-row-repository" import type { JsonRouteResult, @@ -23,8 +28,6 @@ type RouteChunksDependencies = Pick< | "ensureApiKeyForWorkspace" | "ensureWorkspace" | "getCurrentUser" - | "loadChunkPageForSource" - | "loadChunksForSource" | "makeKnowhereClient" | "sourceService" > @@ -85,67 +88,42 @@ const loadSourceChunksEffect = ( return demoResult ?? sourceNotFound() } - const client = yield* Effect.tryPromise(() => - getClientForWorkspace(workspace.id, input.cookieHeader, deps), - ) - const snapshot = yield* Effect.tryPromise(() => - deps.sourceService.ensureParsedSnapshotForRead({ - workspaceId: workspace.id, - source, - client: client.knowledge - ? { - documents: client.documents, - knowledge: client.knowledge, - } - : null, - }), - ).pipe( - Effect.catchAll((error) => - Effect.sync(() => { - logger.warn("sources: parsed snapshot sync for read failed", { - workspaceId: workspace.id, - sourceId: source.id, - documentId: source.knowhereDocumentId, - error: getErrorMessage(error), - }) - return null - }), - ), - ) - if (!isCompleteSnapshot(snapshot)) { + // Reads never return zero chunks for a ready remote document; the SDK falls + // back to Knowhere when Blob storage is missing or stale. A source that is + // not yet ready has no published document to read. + if (source.status !== "ready" || !source.knowhereDocumentId) { return sourceSnapshotProcessing(input) } + + const apiKey = yield* Effect.tryPromise(() => + deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), + ) + const knowledge = getKnowledgeForSource({ + apiKey, + workspaceId: workspace.id, + sourceId: source.id, + documentId: source.knowhereDocumentId, + revisionKey: source.knowhereJobId, + }) + const readableSource = { + documentId: source.knowhereDocumentId, + title: source.title, + revisionKey: source.knowhereJobId, + } + if (input.shouldLoadAll) { - const chunks = yield* deps.loadChunksForSource(source, client, { - snapshot, - workspaceId: workspace.id, - onRevisionKey: async (revisionKey) => { - await deps.sourceService.updateSourceRevisionKey( - workspace.id, - source.id, - revisionKey, - ) - }, - }) - return routeResult.ok({ chunks }) + const chunks = yield* Effect.tryPromise(() => + readAllSourceChunks({ knowledge, source: readableSource }), + ) + return routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }) } - const chunkPage = yield* deps.loadChunkPageForSource( - source, - client, - input.pageParams, - { - assetUrlsByFilePath: snapshot.assetUrlsByFilePath, - snapshot, - workspaceId: workspace.id, - onRevisionKey: async (revisionKey) => { - await deps.sourceService.updateSourceRevisionKey( - workspace.id, - source.id, - revisionKey, - ) - }, - }, + const chunkPage = yield* Effect.tryPromise(() => + readSourceChunkPage({ + knowledge, + source: readableSource, + params: input.pageParams, + }), ) return routeResult.ok(chunkPage) }) @@ -163,6 +141,9 @@ const loadRemoteChunkPageEffect = ( const workspace = yield* Effect.tryPromise(() => deps.ensureWorkspace(user.id), ) + const apiKey = yield* Effect.tryPromise(() => + deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), + ) const client = yield* Effect.tryPromise(() => getClientForWorkspace(workspace.id, input.cookieHeader, deps), ) @@ -185,60 +166,34 @@ const loadRemoteChunkPageEffect = ( revisionKey: remoteDocument.revisionKey ?? null, }), ) - const knowledge = client.knowledge - if (!knowledge) { - throw new Error("Knowhere client does not support parsed snapshot caching.") - } - const snapshot = yield* Effect.tryPromise(() => - deps.sourceService.syncRemoteParsedSnapshot({ - workspaceId: workspace.id, - source, - client: { - documents: client.documents, - knowledge, - }, - }), - ) - if (!isCompleteSnapshot(snapshot)) { - return sourceSnapshotProcessing(input) - } - const readySource = { - ...source, - status: "ready" as const, - knowhereDocumentId: source.knowhereDocumentId ?? remoteDocument.documentId, + const documentId = source.knowhereDocumentId ?? remoteDocument.documentId + + const knowledge = getKnowledgeForSource({ + apiKey, + workspaceId: workspace.id, + sourceId: source.id, + documentId, + revisionKey: source.knowhereJobId ?? remoteDocument.revisionKey ?? null, + }) + const readableSource = { + documentId, + title: source.title, + revisionKey: source.knowhereJobId ?? remoteDocument.revisionKey ?? null, } if (input.shouldLoadAll) { - const chunks = yield* deps.loadChunksForSource(readySource, client, { - snapshot, - workspaceId: workspace.id, - onRevisionKey: async (revisionKey) => { - await deps.sourceService.updateSourceRevisionKey( - workspace.id, - source.id, - revisionKey, - ) - }, - }) - return routeResult.ok({ chunks }) + const chunks = yield* Effect.tryPromise(() => + readAllSourceChunks({ knowledge, source: readableSource }), + ) + return routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }) } - const chunkPage = yield* deps.loadChunkPageForSource( - readySource, - client, - input.pageParams, - { - assetUrlsByFilePath: snapshot.assetUrlsByFilePath, - snapshot, - workspaceId: workspace.id, - onRevisionKey: async (revisionKey) => { - await deps.sourceService.updateSourceRevisionKey( - workspace.id, - source.id, - revisionKey, - ) - }, - }, + const chunkPage = yield* Effect.tryPromise(() => + readSourceChunkPage({ + knowledge, + source: readableSource, + params: input.pageParams, + }), ) return routeResult.ok(chunkPage) }) @@ -361,7 +316,7 @@ function sourceSnapshotProcessing( return routeResult.ok( { chunks: [], - message: "Source parsed snapshot is still being prepared.", + message: "Source is still being prepared.", }, 202, ) @@ -376,29 +331,10 @@ function sourceSnapshotProcessing( total: 0, totalPages: 0, }, - message: "Source parsed snapshot is still being prepared.", + message: "Source is still being prepared.", }, 202, ) } -function isCompleteSnapshot( - snapshot: - | { - readonly snapshotManifestKey?: string | null - readonly snapshotManifestUrl?: string | null - } - | null, -): snapshot is { - readonly snapshotManifestKey: string - readonly snapshotManifestUrl: string -} { - return ( - typeof snapshot?.snapshotManifestKey === "string" && - snapshot.snapshotManifestKey.length > 0 && - typeof snapshot.snapshotManifestUrl === "string" && - snapshot.snapshotManifestUrl.length > 0 - ) -} - export { createRouteChunks } diff --git a/src/domains/sources/route-dependencies.ts b/src/domains/sources/route-dependencies.ts index af81a3c..a39635f 100644 --- a/src/domains/sources/route-dependencies.ts +++ b/src/domains/sources/route-dependencies.ts @@ -1,17 +1,18 @@ import "server-only" import { del } from "@vercel/blob" +import type { Knowledge } from "@ontos-ai/knowhere-sdk" -import { - loadChunkPageForSource, - loadChunksForSource, -} from "@/domains/chunks/server" import { ensureApiKeyForWorkspace } from "@/integrations/dashboard/api-key-service" -import { makeKnowhereClient as makeDefaultKnowhereClient } from "@/integrations/knowhere" +import { + makeKnowhereClient as makeDefaultKnowhereClient, + makeKnowhereClientWithParsedStorage, +} from "@/integrations/knowhere" import { knowhereDemoApi } from "@/integrations/knowhere-demo" import { getCurrentUser, requireUser } from "@/infrastructure/auth" import { workspaceService } from "@/domains/workspace/service" import { sourceViewOptionsBySourceId as getDefaultSourceViewOptionsBySourceId } from "./counts" +import { createParsedDocumentSyncScheduler } from "./parsed-document-sync-scheduler" import { reconcileSourcesForWorkspace as reconcileDefaultSourcesForWorkspace } from "./reconcile" import { sourceWorkflowRuntime } from "./workflow-runtime" import { sourceService as defaultSourceService } from "./service" @@ -32,8 +33,6 @@ const defaultDependencies: SourceRouteServiceDependencies = { sources, client as ReturnType, ), - loadChunkPageForSource, - loadChunksForSource, makeKnowhereClient: (apiKey: string) => makeDefaultKnowhereClient(apiKey) as SourceRouteKnowhereClient, listSourcesForWorkspace: sourceWorkflowRuntime.listForWorkspace, @@ -44,20 +43,14 @@ const defaultDependencies: SourceRouteServiceDependencies = { ), requireUser, sourceService: { - ensureParsedSnapshotForRead: - defaultSourceService.ensureParsedSnapshotForRead, findInWorkspace: defaultSourceService.findInWorkspace, - getParseAssetUrls: defaultSourceService.getParseAssetUrls, - getParseSnapshotMetadata: defaultSourceService.getParseSnapshotMetadata, hideDemoSource: defaultSourceService.hideDemoSource, listHiddenDemoSourceIds: defaultSourceService.listHiddenDemoSourceIds, localizeRemoteDocument: defaultSourceService.localizeRemoteDocument, - updateSourceRevisionKey: defaultSourceService.updateSourceRevisionKey, softDelete: defaultSourceService.softDelete, upsertMaterializedDemoSource: defaultSourceService.upsertMaterializedDemoSource, retrySourceToKnowhere: defaultSourceService.retrySourceToKnowhere, - syncRemoteParsedSnapshot: defaultSourceService.syncRemoteParsedSnapshot, uploadSourceBlobToKnowhere: defaultSourceService.uploadSourceBlobToKnowhere, uploadSourceToKnowhere: defaultSourceService.uploadSourceToKnowhere, }, @@ -92,4 +85,35 @@ async function getClientForWorkspace( return deps.makeKnowhereClient(apiKey) } -export { createSourceRouteDependencies, getClientForWorkspace } +/** + * Build a `Knowledge` configured with Vercel-Blob parsed storage plus a + * QStash-backed background-sync scheduler bound to this source's document. + * Chunk reads go through the returned `knowledge`; on a storage miss the SDK + * serves from Knowhere remote and the scheduler enqueues a durable backfill. + */ +function getKnowledgeForSource(input: { + readonly apiKey: string + readonly workspaceId: string + readonly sourceId: string + readonly documentId: string + readonly revisionKey?: string | null +}): Knowledge { + const scheduler = createParsedDocumentSyncScheduler({ + workspaceId: input.workspaceId, + sourceId: input.sourceId, + documentId: input.documentId, + apiKey: input.apiKey, + revisionKey: input.revisionKey ?? undefined, + }) + const { knowledge } = makeKnowhereClientWithParsedStorage(input.apiKey, { + workspaceId: input.workspaceId, + scheduler, + }) + return knowledge +} + +export { + createSourceRouteDependencies, + getClientForWorkspace, + getKnowledgeForSource, +} diff --git a/src/domains/sources/route-retry.test.ts b/src/domains/sources/route-retry.test.ts new file mode 100644 index 0000000..fc1316c --- /dev/null +++ b/src/domains/sources/route-retry.test.ts @@ -0,0 +1,107 @@ +import { describe, expect, it, vi } from "vitest" + +import { createRouteRetry } from "./route-retry" +import type { Source, Workspace } from "@/infrastructure/db/schema" + +const workspace: Workspace = { + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), +} + +function makeSource(overrides: Partial = {}): Source { + return { + id: "source_1", + workspaceId: "workspace_1", + title: "notes.pdf", + mimeType: "application/pdf", + sizeBytes: 5, + status: "failed", + failureReason: "boom", + failureStage: null, + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: null, + originalBlobUrl: null, + demoKey: null, + createdAt: new Date("2026-05-10T00:00:00.000Z"), + updatedAt: new Date("2026-05-10T00:00:00.000Z"), + deletedAt: null, + ...overrides, + } +} + +function baseDeps() { + return { + ensureApiKeyForWorkspace: vi.fn(async () => "jwt_123"), + ensureWorkspace: vi.fn(async () => workspace), + makeKnowhereClient: vi.fn(), + requireUser: vi.fn(async () => ({ id: "user_1", email: null, name: null })), + sourceService: { + retrySourceToKnowhere: vi.fn(), + }, + } +} + +describe("createRouteRetry", () => { + it("resumes parsed sync for a storage_sync failure without reparsing", async () => { + const storageSyncFailed = makeSource({ failureStage: "storage_sync" }) + const resumedSource = makeSource({ status: "parsing", failureStage: null }) + const resumeParsedSync = vi.fn(async () => resumedSource) + const deps = baseDeps() + const retry = createRouteRetry({ + ...deps, + resumeParsedSync, + sourceService: { + findInWorkspace: vi.fn(async () => storageSyncFailed), + retrySourceToKnowhere: deps.sourceService.retrySourceToKnowhere, + }, + } as unknown as Parameters[0]) + + const result = await retry.retrySource({ + cookieHeader: "session=abc", + sourceId: "source_1", + }) + + expect(result.status).toBe(200) + expect(resumeParsedSync).toHaveBeenCalledWith({ + workspace, + source: storageSyncFailed, + apiKey: "jwt_123", + }) + expect(deps.sourceService.retrySourceToKnowhere).not.toHaveBeenCalled() + }) + + it("reparses a plain parse failure with a saved original blob", async () => { + const parseFailed = makeSource({ + failureStage: "parse", + originalBlobUrl: "https://blob.example/source-uploads/u/document.pdf", + originalBlobPathname: "source-uploads/u/document.pdf", + }) + const parsingSource = makeSource({ status: "parsing", failureStage: null }) + const resumeParsedSync = vi.fn() + const retrySourceToKnowhere = vi.fn(async () => parsingSource) + const deps = baseDeps() + const retry = createRouteRetry({ + ...deps, + resumeParsedSync, + makeKnowhereClient: vi.fn(() => ({})), + sourceService: { + findInWorkspace: vi.fn(async () => parseFailed), + retrySourceToKnowhere, + }, + } as unknown as Parameters[0]) + + const result = await retry.retrySource({ + cookieHeader: "session=abc", + sourceId: "source_1", + }) + + expect(result.status).toBe(200) + expect(resumeParsedSync).not.toHaveBeenCalled() + expect(retrySourceToKnowhere).toHaveBeenCalled() + }) +}) diff --git a/src/domains/sources/route-retry.ts b/src/domains/sources/route-retry.ts index 68464bb..094c91c 100644 --- a/src/domains/sources/route-retry.ts +++ b/src/domains/sources/route-retry.ts @@ -1,8 +1,12 @@ import { Effect } from "effect" +import { logger } from "@/lib/logger" import { routeResult } from "@/lib/route-result" import { startBackgroundReconciliation } from "./background-reconcile" +import { enqueueParsedDocumentSync } from "./parsed-document-sync-scheduler" +import { sourceWorkflowRuntime } from "./workflow-runtime" import { toSourceView } from "./view" +import type { Source, Workspace } from "@/infrastructure/db/schema" import type { JsonRouteResult, RetrySourceBody, @@ -17,7 +21,9 @@ type RouteRetryDependencies = Pick< | "makeKnowhereClient" | "requireUser" | "sourceService" -> +> & { + readonly resumeParsedSync?: typeof resumeParsedSync +} type RouteRetry = { readonly retrySource: ( @@ -51,6 +57,23 @@ const retrySourceEffect = ( if (source.status !== "failed") { return routeResult.error(409, "Only failed sources can be retried.") } + + // A storage-sync failure keeps the parsed Knowhere document intact — resume + // the parsed-document sync from the stored revision instead of reparsing. + if (source.failureStage === "storage_sync" && source.knowhereDocumentId) { + const apiKey = yield* Effect.tryPromise(() => + deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), + ) + const resumedSource = yield* Effect.tryPromise(() => + (deps.resumeParsedSync ?? resumeParsedSync)({ + workspace, + source, + apiKey, + }), + ) + return routeResult.ok({ source: toSourceView(resumedSource) }) + } + if (!source.originalBlobUrl || !source.originalBlobPathname) { return routeResult.error( 409, @@ -75,4 +98,50 @@ const retrySourceEffect = ( return routeResult.ok({ source: toSourceView(retriedSource) }) }) -export { createRouteRetry } +/** + * Resume a parsed-document storage sync that failed after a successful parse. + * Moves the source back to `parsing` (clearing the failure stage), resets the + * sync status, and re-enqueues the resumable parsed-sync workflow from the + * stored document + revision — no reparse. + */ +async function resumeParsedSync(input: { + readonly workspace: Workspace + readonly source: Source + readonly apiKey: string +}): Promise { + const { workspace, source, apiKey } = input + const documentId = source.knowhereDocumentId + if (!documentId) return source + + const revisionKey = source.knowhereJobId ?? undefined + const parsingSource = await sourceWorkflowRuntime.markParsing( + workspace.id, + source.id, + revisionKey ?? documentId, + documentId, + "failed", + ) + if (!parsingSource) return source + + await sourceWorkflowRuntime.updateSyncStatus(workspace.id, source.id, { + revisionKey, + syncStatus: "running", + syncError: null, + }) + await enqueueParsedDocumentSync({ + workspaceId: workspace.id, + sourceId: source.id, + documentId, + apiKey, + revisionKey, + }) + logger.info("sources: resumed parsed document storage sync on retry", { + workspaceId: workspace.id, + sourceId: source.id, + documentId, + revisionKey, + }) + return parsingSource +} + +export { createRouteRetry, resumeParsedSync } diff --git a/src/domains/sources/route-service.test.ts b/src/domains/sources/route-service.test.ts index ca37390..26e60cb 100644 --- a/src/domains/sources/route-service.test.ts +++ b/src/domains/sources/route-service.test.ts @@ -2,7 +2,6 @@ import { describe, expect, it, vi } from "vitest"; import { Effect } from "effect"; import type { Job } from "@ontos-ai/knowhere-sdk"; -import type { ChunkPage } from "@/domains/chunks"; import type { Source, Workspace } from "@/infrastructure/db/schema"; import { createRouteListing } from "./route-listing"; import { createSourceRouteService } from "./route-service"; @@ -23,6 +22,7 @@ const source: Source = { sizeBytes: 5, status: "parsing", failureReason: null, + failureStage: null, knowhereJobId: "job_1", knowhereDocumentId: null, stagedBlobPathname: null, @@ -120,147 +120,15 @@ describe("source route service", () => { expect(listHiddenDemoSourceIds).toHaveBeenCalledWith(workspace.id); }); - it("syncs a ready legacy source snapshot before loading displayed chunks", async () => { - const readySource: Source = { - ...source, - id: "00000000-0000-4000-8000-000000000001", - knowhereJobId: null, - knowhereDocumentId: "doc_legacy", - status: "ready", - }; - const snapshot = { - resultBlobUrl: "https://blob.example/manifest/current.json", - snapshotManifestUrl: "https://blob.example/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/00000000-0000-4000-8000-000000000001/parsed-result/manifest/current.json", - assetUrlsByFilePath: { - "images/chart.png": "https://blob.example/images/chart.png", - }, - }; - const listChunks = vi.fn(async () => { - throw new Error("display path must not call Knowhere listChunks"); - }); - const knowledge = { - loadJobResult: vi.fn(), - }; - const knowhereClient = { - documents: { - archive: vi.fn(async () => undefined), - listChunks, - }, - jobs: { - create: vi.fn(), - get: vi.fn(), - upload: vi.fn(), - }, - knowledge, - }; - const loadedPage: ChunkPage = { - chunks: [ - { - chunkId: "chunk_1", - type: "text", - content: "Chunk body", - sourceTitle: "notes.pdf", - }, - ], - pagination: { - page: 1, - pageSize: 50, - total: 1, - totalPages: 1, - }, - }; - const loadChunkPageForSource = vi.fn(() => Effect.succeed(loadedPage)); - const syncRemoteParsedSnapshot = vi.fn(async () => snapshot); - const service = createSourceRouteService({ - ensureApiKeyForWorkspace: vi.fn(async () => "jwt_123"), - ensureWorkspace: vi.fn(async () => workspace), - getCurrentUser: vi.fn(async () => ({ - id: "user_1", - email: null, - name: null, - })), - loadChunkPageForSource, - makeKnowhereClient: vi.fn(() => knowhereClient), - sourceService: { - findInWorkspace: vi.fn(async () => readySource), - ensureParsedSnapshotForRead: syncRemoteParsedSnapshot, - updateSourceRevisionKey: vi.fn(async () => readySource), - }, - }); - - const result = await service.loadSourceChunks({ - cookieHeader: "session=abc", - sourceId: readySource.id, - shouldLoadAll: false, - pageParams: { page: 1, pageSize: 50 }, - }); - - expect(result).toEqual({ - status: 200, - body: { - chunks: [ - { - chunkId: "chunk_1", - type: "text", - content: "Chunk body", - sourceTitle: "notes.pdf", - }, - ], - pagination: { - page: 1, - pageSize: 50, - total: 1, - totalPages: 1, - }, - }, - }); - expect(syncRemoteParsedSnapshot).toHaveBeenCalledWith({ - workspaceId: workspace.id, - source: readySource, - client: { - documents: knowhereClient.documents, - knowledge, - }, - }); - expect(loadChunkPageForSource).toHaveBeenCalledWith( - readySource, - knowhereClient, - { page: 1, pageSize: 50 }, - expect.objectContaining({ - assetUrlsByFilePath: snapshot.assetUrlsByFilePath, - snapshot, - workspaceId: workspace.id, - }), - ); - expect(listChunks).not.toHaveBeenCalled(); - }); - - it("returns processing when a ready legacy source snapshot is not synced yet", async () => { - const readySource: Source = { + it("returns processing without building a client when a source is not ready", async () => { + const parsingSource: Source = { ...source, id: "00000000-0000-4000-8000-000000000002", knowhereDocumentId: "doc_legacy", knowhereJobId: null, - status: "ready", - }; - const knowledge = { - loadJobResult: vi.fn(), - }; - const knowhereClient = { - documents: { - archive: vi.fn(async () => undefined), - listChunks: vi.fn(), - }, - jobs: { - create: vi.fn(), - get: vi.fn(), - upload: vi.fn(), - }, - knowledge, + status: "parsing", }; - const loadChunkPageForSource = vi.fn(); + const makeKnowhereClient = vi.fn(); const service = createSourceRouteService({ ensureApiKeyForWorkspace: vi.fn(async () => "jwt_123"), ensureWorkspace: vi.fn(async () => workspace), @@ -269,19 +137,15 @@ describe("source route service", () => { email: null, name: null, })), - loadChunkPageForSource, - makeKnowhereClient: vi.fn(() => knowhereClient), + makeKnowhereClient, sourceService: { - findInWorkspace: vi.fn(async () => readySource), - ensureParsedSnapshotForRead: vi.fn(async () => { - throw new Error("snapshot import still running"); - }), + findInWorkspace: vi.fn(async () => parsingSource), }, }); const result = await service.loadSourceChunks({ cookieHeader: "session=abc", - sourceId: readySource.id, + sourceId: parsingSource.id, shouldLoadAll: false, pageParams: { page: 1, pageSize: 50 }, }); @@ -296,11 +160,9 @@ describe("source route service", () => { total: 0, totalPages: 0, }, - message: "Source parsed snapshot is still being prepared.", + message: "Source is still being prepared.", }, }); - expect(loadChunkPageForSource).not.toHaveBeenCalled(); - expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled(); }); it("lists shared default and legacy namespace documents as lightweight remote sources", async () => { @@ -985,6 +847,7 @@ describe("source route service", () => { ...failedSource, status: "parsing", failureReason: null, + failureStage: null, knowhereJobId: "job_retry", }; const knowhereClient = { diff --git a/src/domains/sources/route-types.ts b/src/domains/sources/route-types.ts index a32cad2..96bd80c 100644 --- a/src/domains/sources/route-types.ts +++ b/src/domains/sources/route-types.ts @@ -3,10 +3,6 @@ import type { ChunkPage, ChunkPageParams, } from "@/domains/chunks" -import type { - loadChunkPageForSource, - loadChunksForSource, -} from "@/domains/chunks/server" import type { ParsedChunkView } from "@/domains/chunks/types" import type { SourceStatus, SourceView } from "@/domains/sources/types" import type { AuthUser } from "@/infrastructure/auth" @@ -18,13 +14,8 @@ import type { import type { RouteResult } from "@/lib/route-result" import type { SourceBlobUploadInput } from "./blob-upload" import type { sourceViewOptionsBySourceId } from "./counts" -import type { syncRemoteParsedSnapshot } from "./remote-parsed-snapshot-sync" import type { UploadKnowhereClient } from "./upload" -type SyncRemoteParsedSnapshotInput = Parameters< - typeof syncRemoteParsedSnapshot ->[0] - type SourceRouteKnowhereClient = UploadKnowhereClient & ChunkKnowhereClient & { readonly documents: ChunkKnowhereClient["documents"] & { @@ -59,7 +50,6 @@ type SourceRouteKnowhereClient = UploadKnowhereClient & }> archive(documentId: string): Promise } - readonly knowledge?: SyncRemoteParsedSnapshotInput["client"]["knowledge"] } type SourceUploadRequest = @@ -162,16 +152,6 @@ type SourceRouteService = { } type SourceWorkflowService = { - readonly ensureParsedSnapshotForRead: (input: { - readonly workspaceId: string - readonly source: Source - readonly client?: SyncRemoteParsedSnapshotInput["client"] | null - }) => Promise<{ - readonly resultBlobUrl: string - readonly snapshotManifestUrl?: string | null - readonly snapshotManifestKey?: string | null - readonly assetUrlsByFilePath: Readonly> - } | null> readonly uploadSourceToKnowhere: ( workspace: Workspace, file: File, @@ -187,7 +167,6 @@ type SourceWorkflowService = { source: Source, knowhere: UploadKnowhereClient, ) => Promise - readonly syncRemoteParsedSnapshot: typeof syncRemoteParsedSnapshot readonly findInWorkspace: ( workspaceId: string, sourceId: string, @@ -196,19 +175,6 @@ type SourceWorkflowService = { workspaceId: string, sourceId: string, ) => Promise - readonly getParseAssetUrls: ( - workspaceId: string, - sourceId: string, - ) => Promise>> - readonly getParseSnapshotMetadata: ( - workspaceId: string, - sourceId: string, - ) => Promise<{ - readonly resultBlobUrl: string - readonly snapshotManifestUrl?: string | null - readonly snapshotManifestKey?: string | null - readonly assetUrlsByFilePath: Readonly> - } | null> readonly hideDemoSource: ( workspaceId: string, demoSourceId: string, @@ -226,11 +192,6 @@ type SourceWorkflowService = { readonly revisionKey?: string | null }, ) => Promise - readonly updateSourceRevisionKey: ( - workspaceId: string, - sourceId: string, - revisionKey: string, - ) => Promise readonly upsertMaterializedDemoSource: ( workspaceId: string, input: { @@ -266,8 +227,6 @@ type SourceRouteServiceDependencies = { sources: readonly Source[], client: SourceRouteKnowhereClient, ) => ReturnType - readonly loadChunkPageForSource: typeof loadChunkPageForSource - readonly loadChunksForSource: typeof loadChunksForSource readonly makeKnowhereClient: (apiKey: string) => SourceRouteKnowhereClient readonly listSourcesForWorkspace: (workspaceId: string) => Promise readonly reconcileSourcesForWorkspace: ( diff --git a/src/domains/sources/service.ts b/src/domains/sources/service.ts index 11e9199..4fcbfd5 100644 --- a/src/domains/sources/service.ts +++ b/src/domains/sources/service.ts @@ -9,36 +9,19 @@ import { uploadSourceBlobToKnowhereEffect, uploadSourceToKnowhereEffect, } from "./upload" -import { syncRemoteParsedSnapshot } from "./remote-parsed-snapshot-sync" import { retrySourceToKnowhereEffect } from "./retry" import { sourceWorkflowRuntime } from "./workflow-runtime" type SourceService = { - readonly ensureParsedSnapshotForRead: ( - input: EnsureParsedSnapshotForReadInput, - ) => Promise readonly findInWorkspace: ( workspaceId: string, sourceId: string, ) => Promise - readonly getParseAssetUrls: ( - workspaceId: string, - sourceId: string, - ) => Promise>> - readonly getParseSnapshotMetadata: ( - workspaceId: string, - sourceId: string, - ) => ReturnType readonly listForWorkspace: (workspaceId: string) => Promise readonly localizeRemoteDocument: ( workspaceId: string, input: Parameters[1], ) => Promise - readonly updateSourceRevisionKey: ( - workspaceId: string, - sourceId: string, - revisionKey: string, - ) => Promise readonly listHiddenDemoSourceIds: (workspaceId: string) => Promise readonly hideDemoSource: ( workspaceId: string, @@ -69,20 +52,6 @@ type SourceService = { source: Source, knowhere: UploadKnowhereClient, ) => Promise - readonly syncRemoteParsedSnapshot: typeof syncRemoteParsedSnapshot -} - -type EnsureParsedSnapshotForReadInput = { - readonly workspaceId: string - readonly source: Source - readonly client?: Parameters[0]["client"] | null -} - -type EnsureParsedSnapshotForReadResult = { - readonly resultBlobUrl: string - readonly snapshotManifestUrl?: string | null - readonly snapshotManifestKey?: string | null - readonly assetUrlsByFilePath: Readonly> } const uploadSourceToKnowhere: SourceService["uploadSourceToKnowhere"] = ( @@ -125,64 +94,16 @@ const retrySourceToKnowhere: SourceService["retrySourceToKnowhere"] = ( }), ) -const ensureParsedSnapshotForRead: SourceService["ensureParsedSnapshotForRead"] = - async ({ workspaceId, source, client }: EnsureParsedSnapshotForReadInput) => { - const existingSnapshot = - await sourceWorkflowRuntime.getParseSnapshotMetadata(workspaceId, source.id) - if (isCompleteParsedSnapshot(existingSnapshot)) return existingSnapshot - - if (!canSyncParsedSnapshotForRead(source, client)) { - return existingSnapshot - } - - return syncRemoteParsedSnapshot({ - workspaceId, - source, - client, - }) - } - export const sourceService: SourceService = { - ensureParsedSnapshotForRead, findInWorkspace: sourceWorkflowRuntime.findInWorkspace, - getParseAssetUrls: sourceWorkflowRuntime.getParseAssetUrls, - getParseSnapshotMetadata: sourceWorkflowRuntime.getParseSnapshotMetadata, hideDemoSource: sourceWorkflowRuntime.hideDemoSource, listHiddenDemoSourceIds: sourceWorkflowRuntime.listHiddenDemoSourceIds, listForWorkspace: sourceWorkflowRuntime.listForWorkspace, localizeRemoteDocument: sourceWorkflowRuntime.localizeRemoteDocument, - updateSourceRevisionKey: sourceWorkflowRuntime.updateRevisionKey, softDelete: sourceWorkflowRuntime.softDelete, upsertMaterializedDemoSource: sourceWorkflowRuntime.upsertMaterializedDemoSource, uploadSourceToKnowhere, uploadSourceBlobToKnowhere, retrySourceToKnowhere, - syncRemoteParsedSnapshot, -} - -function canSyncParsedSnapshotForRead( - source: Source, - client: EnsureParsedSnapshotForReadInput["client"], -): client is Parameters[0]["client"] { - return ( - Boolean(client) && - (source.status === "ready" || source.status === "parsing") && - typeof source.knowhereDocumentId === "string" && - source.knowhereDocumentId.length > 0 - ) -} - -function isCompleteParsedSnapshot( - snapshot: EnsureParsedSnapshotForReadResult | null, -): snapshot is EnsureParsedSnapshotForReadResult & { - readonly snapshotManifestKey: string - readonly snapshotManifestUrl: string -} { - return ( - typeof snapshot?.snapshotManifestKey === "string" && - snapshot.snapshotManifestKey.length > 0 && - typeof snapshot.snapshotManifestUrl === "string" && - snapshot.snapshotManifestUrl.length > 0 - ) } diff --git a/src/domains/sources/source-parse-result-repository.ts b/src/domains/sources/source-parse-result-repository.ts index c3e9859..6ddc1e2 100644 --- a/src/domains/sources/source-parse-result-repository.ts +++ b/src/domains/sources/source-parse-result-repository.ts @@ -14,13 +14,25 @@ type SaveSourceParseResultInput = { readonly resultBlobUrl: string readonly snapshotManifestUrl?: string readonly snapshotManifestKey?: string + readonly revisionKey?: string + readonly syncStatus?: SourceParseSyncStatus + readonly syncError?: string | null readonly assetUrlsByFilePath: Readonly> } +export type SourceParseSyncStatus = + | "pending" + | "running" + | "completed" + | "failed" + type SourceParseResultProgress = { readonly resultBlobUrl: string readonly snapshotManifestUrl?: string | null readonly snapshotManifestKey?: string | null + readonly revisionKey?: string | null + readonly syncStatus?: string | null + readonly syncError?: string | null readonly assetUrlsByFilePath: Readonly> } @@ -28,9 +40,18 @@ export type SourceParseSnapshotMetadata = { readonly resultBlobUrl: string readonly snapshotManifestUrl?: string | null readonly snapshotManifestKey?: string | null + readonly revisionKey?: string | null + readonly syncStatus?: string | null + readonly syncError?: string | null readonly assetUrlsByFilePath: Readonly> } +type UpdateSyncStatusInput = { + readonly revisionKey?: string + readonly syncStatus: SourceParseSyncStatus + readonly syncError?: string | null +} + type SourceParseResultRepository = { readonly saveParseResultEffect: ( workspaceId: string, @@ -54,6 +75,11 @@ type SourceParseResultRepository = { workspaceId: string, sourceId: string, ) => Effect.Effect>, never, DbClient> + readonly updateSyncStatusEffect: ( + workspaceId: string, + sourceId: string, + input: UpdateSyncStatusInput, + ) => Effect.Effect } export function buildAtomicAssetUrlsMergeSql( @@ -79,6 +105,9 @@ const saveParseResultEffect: SourceParseResultRepository["saveParseResultEffect" resultBlobUrl: input.resultBlobUrl, snapshotManifestUrl: input.snapshotManifestUrl, snapshotManifestKey: input.snapshotManifestKey, + revisionKey: input.revisionKey, + syncStatus: input.syncStatus, + syncError: input.syncError, assetUrls: input.assetUrlsByFilePath, }) .onConflictDoUpdate({ @@ -87,6 +116,9 @@ const saveParseResultEffect: SourceParseResultRepository["saveParseResultEffect" resultBlobUrl: input.resultBlobUrl, snapshotManifestUrl: input.snapshotManifestUrl, snapshotManifestKey: input.snapshotManifestKey, + revisionKey: input.revisionKey, + syncStatus: input.syncStatus, + syncError: input.syncError, assetUrls: input.assetUrlsByFilePath, updatedAt: sql`now()`, }, @@ -114,6 +146,9 @@ const mergeParseAssetUrlsEffect: SourceParseResultRepository["mergeParseAssetUrl resultBlobUrl: input.resultBlobUrl, snapshotManifestUrl: input.snapshotManifestUrl, snapshotManifestKey: input.snapshotManifestKey, + revisionKey: input.revisionKey, + syncStatus: input.syncStatus, + syncError: input.syncError, assetUrls: input.assetUrlsByFilePath, }) .onConflictDoUpdate({ @@ -122,6 +157,9 @@ const mergeParseAssetUrlsEffect: SourceParseResultRepository["mergeParseAssetUrl resultBlobUrl: input.resultBlobUrl, snapshotManifestUrl: input.snapshotManifestUrl, snapshotManifestKey: input.snapshotManifestKey, + revisionKey: input.revisionKey, + syncStatus: input.syncStatus, + syncError: input.syncError, assetUrls: buildAtomicAssetUrlsMergeSql( input.assetUrlsByFilePath, ), @@ -149,6 +187,9 @@ const getParseResultProgressEffect: SourceParseResultRepository["getParseResultP resultBlobUrl: sourceParseResults.resultBlobUrl, snapshotManifestUrl: sourceParseResults.snapshotManifestUrl, snapshotManifestKey: sourceParseResults.snapshotManifestKey, + revisionKey: sourceParseResults.revisionKey, + syncStatus: sourceParseResults.syncStatus, + syncError: sourceParseResults.syncError, assetUrls: sourceParseResults.assetUrls, }) .from(sourceParseResults) @@ -159,13 +200,54 @@ const getParseResultProgressEffect: SourceParseResultRepository["getParseResultP if (!progress) return null return { - resultBlobUrl: progress.resultBlobUrl, + resultBlobUrl: progress.resultBlobUrl ?? "", snapshotManifestUrl: progress.snapshotManifestUrl, snapshotManifestKey: progress.snapshotManifestKey, + revisionKey: progress.revisionKey, + syncStatus: progress.syncStatus, + syncError: progress.syncError, assetUrlsByFilePath: progress.assetUrls, } }) +const updateSyncStatusEffect: SourceParseResultRepository["updateSyncStatusEffect"] = + (workspaceId: string, sourceId: string, input: UpdateSyncStatusInput) => + Effect.gen(function* () { + const db = yield* DbClient + const source = yield* Effect.promise(() => + sourceRowRepository.findInWorkspaceWithDb(db, workspaceId, sourceId), + ) + if (!source) return null + + // Upsert: a read-miss backfill may set sync status before any parsed + // snapshot row exists, so insert a bare row when one is not present yet. + const [result] = yield* Effect.promise(() => + db + .insert(sourceParseResults) + .values({ + sourceId, + revisionKey: input.revisionKey, + syncStatus: input.syncStatus, + syncError: input.syncError ?? null, + assetUrls: {}, + }) + .onConflictDoUpdate({ + target: sourceParseResults.sourceId, + set: { + ...(input.revisionKey !== undefined + ? { revisionKey: input.revisionKey } + : {}), + syncStatus: input.syncStatus, + syncError: input.syncError ?? null, + updatedAt: sql`now()`, + }, + }) + .returning(), + ) + + return result ?? null + }) + const getParseAssetUrlsEffect: SourceParseResultRepository["getParseAssetUrlsEffect"] = (workspaceId: string, sourceId: string) => Effect.gen(function* () { @@ -199,4 +281,5 @@ export const sourceParseResultRepository: SourceParseResultRepository = { getParseResultProgressEffect, getParseSnapshotMetadataEffect, getParseAssetUrlsEffect, + updateSyncStatusEffect, } diff --git a/src/domains/sources/source-reconcile-route-workflow.test.ts b/src/domains/sources/source-reconcile-route-workflow.test.ts index ea68cbe..52f9c6b 100644 --- a/src/domains/sources/source-reconcile-route-workflow.test.ts +++ b/src/domains/sources/source-reconcile-route-workflow.test.ts @@ -1,13 +1,13 @@ import { afterEach, describe, expect, it, vi } from "vitest" const mocks = vi.hoisted(() => ({ - createParsedResultStorageAdapter: vi.fn(), - saveParseResult: vi.fn(), + enqueueParsedDocumentSync: vi.fn(), + updateSyncStatus: vi.fn(), + markFailed: vi.fn(), loggerError: vi.fn(), loggerInfo: vi.fn(), loggerWarn: vi.fn(), - makeKnowhereClient: vi.fn(), - markFailed: vi.fn(), + makeKnowhereClientWithParsedStorage: vi.fn(), markSourceReadyAfterReconciliation: vi.fn(), pollSourceReconciliation: vi.fn(), })) @@ -20,16 +20,17 @@ vi.mock("@/domains/sources/source-reconcile-workflow", () => ({ vi.mock("@/domains/sources/workflow-runtime", () => ({ sourceWorkflowRuntime: { markFailed: mocks.markFailed, - saveParseResult: mocks.saveParseResult, + updateSyncStatus: mocks.updateSyncStatus, }, })) -vi.mock("./parse-result-storage-adapter", () => ({ - createParsedResultStorageAdapter: mocks.createParsedResultStorageAdapter, +vi.mock("./parsed-document-sync-scheduler", () => ({ + enqueueParsedDocumentSync: mocks.enqueueParsedDocumentSync, })) vi.mock("@/integrations/knowhere", () => ({ - makeKnowhereClient: mocks.makeKnowhereClient, + makeKnowhereClientWithParsedStorage: + mocks.makeKnowhereClientWithParsedStorage, })) vi.mock("@/lib/logger", () => ({ @@ -42,6 +43,30 @@ vi.mock("@/lib/logger", () => ({ import { sourceReconcileRouteWorkflow } from "./source-reconcile-route-workflow" +function createClient(overrides: { + syncParsedDocument?: ReturnType + jobResultId?: string + jobId?: string +}) { + const listChunks = vi.fn(async () => ({ + jobResultId: overrides.jobResultId ?? "rev_1", + jobId: overrides.jobId ?? "job_1", + })) + return { + client: { jobs: {}, documents: { listChunks } }, + knowledge: { + syncParsedDocument: + overrides.syncParsedDocument ?? + vi.fn(async () => ({ + documentId: "doc_1", + revisionKey: "rev_1", + completed: true, + })), + }, + listChunks, + } +} + describe("sourceReconcileRouteWorkflow", () => { afterEach(() => { vi.clearAllMocks() @@ -64,7 +89,7 @@ describe("sourceReconcileRouteWorkflow", () => { }) }) - it("writes a parsed snapshot before marking the source ready", async () => { + it("syncs the parsed document then marks the source ready", async () => { const context = createWorkflowContext() const continuations: ContinuationTriggerInput[] = [] const restore = @@ -73,36 +98,19 @@ describe("sourceReconcileRouteWorkflow", () => { continuations.push(input) }, ) - const storageAdapter = { - adapter: { - writeObject: vi.fn(), - }, - keyPrefix: "workspaces/workspace_1/sources/source_1/parsed-result", - } - const loadJobResult = vi.fn().mockResolvedValue({ - assetUrlsByFilePath: { - "page_citation_assets/page-1.png": - "https://blob.example/page_citation_assets/page-1.png", - }, - parsedSnapshot: { - manifestKey: - "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - manifestUrl: - "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - }, + const wired = createClient({}) + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: wired.client, + knowledge: wired.knowledge, + }) + mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ + status: "ready", }) - const client = { jobs: {}, knowledge: { loadJobResult } } - mocks.createParsedResultStorageAdapter.mockReturnValue(storageAdapter) - mocks.makeKnowhereClient.mockReturnValue(client) - mocks.saveParseResult.mockResolvedValue({ id: "parse_result_1" }) mocks.pollSourceReconciliation.mockResolvedValue({ kind: "ready-to-prepare", jobId: "job_1", documentId: "doc_1", }) - mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ - status: "ready", - }) try { await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ @@ -117,54 +125,70 @@ describe("sourceReconcileRouteWorkflow", () => { restore() } - expect(mocks.createParsedResultStorageAdapter).toHaveBeenCalledWith({ - workspaceId: "workspace_1", - sourceId: "source_1", - }) - expect(loadJobResult).toHaveBeenCalledWith({ - jobId: "job_1", - storageAdapter, + expect(wired.knowledge.syncParsedDocument).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "rev_1", }) - expect(mocks.saveParseResult).toHaveBeenCalledWith( + expect(mocks.updateSyncStatus).toHaveBeenCalledWith( "workspace_1", "source_1", - { - resultBlobUrl: - "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - snapshotManifestUrl: - "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - snapshotManifestKey: - "workspaces/workspace_1/sources/source_1/parsed-result/manifest/current.json", - assetUrlsByFilePath: { - "page_citation_assets/page-1.png": - "https://blob.example/page_citation_assets/page-1.png", - }, - }, + { revisionKey: "rev_1", syncStatus: "completed" }, ) expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ workspaceId: "workspace_1", sourceId: "source_1", documentId: "doc_1", }) - expect(mocks.loggerWarn).not.toHaveBeenCalled() + expect(mocks.enqueueParsedDocumentSync).not.toHaveBeenCalled() expect(continuations).toEqual([]) }) - it("does not mark ready when snapshot manifest storage is missing", async () => { + it("hands off to parsed-sync and stays parsing when sync is incomplete", async () => { const context = createWorkflowContext() - const storageAdapter = { - adapter: { - writeObject: vi.fn(), - }, - keyPrefix: "workspaces/workspace_1/sources/source_1/parsed-result", - } - const loadJobResult = vi.fn().mockResolvedValue({ - assetUrlsByFilePath: {}, + const syncParsedDocument = vi.fn(async () => ({ + documentId: "doc_1", + revisionKey: "rev_1", + completed: false, + })) + const wired = createClient({ syncParsedDocument }) + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: wired.client, + knowledge: wired.knowledge, }) - mocks.createParsedResultStorageAdapter.mockReturnValue(storageAdapter) - mocks.makeKnowhereClient.mockReturnValue({ - jobs: {}, - knowledge: { loadJobResult }, + mocks.pollSourceReconciliation.mockResolvedValue({ + kind: "ready-to-prepare", + jobId: "job_1", + documentId: "doc_1", + }) + + await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ + context, + payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + }), + }) + + expect(mocks.markSourceReadyAfterReconciliation).not.toHaveBeenCalled() + expect(mocks.enqueueParsedDocumentSync).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "source_1", + documentId: "doc_1", + apiKey: "jwt_1", + revisionKey: "rev_1", + }) + }) + + it("fails the source with storage_sync stage when sync throws", async () => { + const context = createWorkflowContext() + const syncParsedDocument = vi.fn(async () => { + throw new Error("blob write failed") + }) + const wired = createClient({ syncParsedDocument }) + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: wired.client, + knowledge: wired.knowledge, }) mocks.pollSourceReconciliation.mockResolvedValue({ kind: "ready-to-prepare", @@ -181,9 +205,15 @@ describe("sourceReconcileRouteWorkflow", () => { apiKey: "jwt_1", }), }), - ).rejects.toThrow("Parsed result snapshot was not written") + ).rejects.toThrow("blob write failed") - expect(mocks.saveParseResult).not.toHaveBeenCalled() + expect(mocks.markFailed).toHaveBeenCalledWith( + "workspace_1", + "source_1", + expect.stringContaining("storage sync failed"), + "parsing", + "storage_sync", + ) expect(mocks.markSourceReadyAfterReconciliation).not.toHaveBeenCalled() }) @@ -196,7 +226,10 @@ describe("sourceReconcileRouteWorkflow", () => { continuations.push(input) }, ) - mocks.makeKnowhereClient.mockReturnValue({ jobs: {} }) + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: { jobs: {}, documents: { listChunks: vi.fn() } }, + knowledge: { syncParsedDocument: vi.fn() }, + }) mocks.pollSourceReconciliation.mockResolvedValue({ kind: "waiting", jobId: "job_1", diff --git a/src/domains/sources/source-reconcile-route-workflow.ts b/src/domains/sources/source-reconcile-route-workflow.ts index 16f9054..cfa90b3 100644 --- a/src/domains/sources/source-reconcile-route-workflow.ts +++ b/src/domains/sources/source-reconcile-route-workflow.ts @@ -6,11 +6,9 @@ import { markSourceReadyAfterReconciliation, pollSourceReconciliation, } from "@/domains/sources/source-reconcile-workflow" -import { makeKnowhereClient } from "@/integrations/knowhere" +import { makeKnowhereClientWithParsedStorage } from "@/integrations/knowhere" import { logger } from "@/lib/logger" -import { - createParsedResultStorageAdapter, -} from "./parse-result-storage-adapter" +import { enqueueParsedDocumentSync } from "./parsed-document-sync-scheduler" import { sourceWorkflowRuntime } from "./workflow-runtime" type ReconcilePayload = { @@ -23,14 +21,6 @@ type ReconcilePayload = { type ReconcilePhase = "poll-and-ready" | "poll-and-mirror" | "asset-batches" -type ParsedSnapshotResponse = { - readonly assetUrlsByFilePath?: Readonly> - readonly parsedSnapshot?: { - readonly manifestKey: string - readonly manifestUrl?: string - } -} - type NormalizedReconcilePayload = { readonly workspaceId: string readonly sourceId: string @@ -53,6 +43,7 @@ type ContinuationTriggerInput = { const maxPollAttempts = 25 const initialDelaySeconds = 3 const maxDelaySeconds = 30 +const maxSyncStepsPerReconcile = 4 let triggerContinuation: typeof triggerReconcileContinuation = triggerReconcileContinuation @@ -63,7 +54,9 @@ async function runPollAndMirrorWorkflow(input: { }): Promise { const { context, payload } = input const { workspaceId, sourceId, apiKey } = payload - const client = makeKnowhereClient(apiKey) + const { client, knowledge } = makeKnowhereClientWithParsedStorage(apiKey, { + workspaceId, + }) let delay = initialDelaySeconds let completedJob: { readonly jobId: string @@ -129,35 +122,88 @@ async function runPollAndMirrorWorkflow(input: { return } - const snapshot = await context.run("load-job-result-snapshot", async () => { - const loadedResult = (await client.knowledge.loadJobResult({ - jobId: jobToPrepare.jobId, - storageAdapter: createParsedResultStorageAdapter({ + // Resolve the parsed revision and record that storage sync is in progress. + const revisionKey = await context.run("resolve-revision-key", async () => { + const firstPage = await client.documents.listChunks( + jobToPrepare.documentId, + { page: 1, pageSize: 1, includeAssetUrls: false }, + ) + return firstPage.jobResultId ?? firstPage.jobId ?? jobToPrepare.jobId + }) + await context.run("record-sync-running", async () => + sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { + revisionKey, + syncStatus: "running", + }), + ) + + // Sync the parsed snapshot into Blob before the source is considered ready. + // Each step is bounded by SDK limits; if it does not finish in this segment, + // hand off to the resumable parsed-sync workflow and stay `parsing`. A sync + // error stages the failure as `storage_sync` so a retry resumes sync without + // reparsing, then rethrows to fail the workflow run. + let syncCompleted = false + for (let step = 0; step < maxSyncStepsPerReconcile; step++) { + const result = await context.run(`parsed-sync-${step}`, async () => { + try { + return await knowledge.syncParsedDocument({ + documentId: jobToPrepare.documentId, + revisionKey, + }) + } catch (error) { + await sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { + revisionKey, + syncStatus: "failed", + syncError: error instanceof Error ? error.message : String(error), + }) + await sourceWorkflowRuntime.markFailed( + workspaceId, + sourceId, + `Parsed document storage sync failed: ${ + error instanceof Error ? error.message : String(error) + }`, + "parsing", + "storage_sync", + ) + throw error + } + }) + if (result.completed) { + syncCompleted = true + break + } + } + + if (!syncCompleted) { + await context.run("record-sync-progress", async () => + sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { + revisionKey, + syncStatus: "running", + }), + ) + await context.run("enqueue-parsed-sync-continuation", async () => + enqueueParsedDocumentSync({ workspaceId, sourceId, + documentId: jobToPrepare.documentId, + apiKey, + revisionKey, }), - })) as ParsedSnapshotResponse - const manifest = loadedResult.parsedSnapshot - if (!manifest?.manifestUrl || !manifest.manifestKey) { - throw new Error( - "Parsed result snapshot was not written; refusing to mark source ready.", - ) - } - return { - assetUrlsByFilePath: loadedResult.assetUrlsByFilePath ?? {}, - snapshotManifestUrl: manifest.manifestUrl, - snapshotManifestKey: manifest.manifestKey, - } - }) - await context.run("save-parse-result-assets", async () => - sourceWorkflowRuntime.saveParseResult(workspaceId, sourceId, { - resultBlobUrl: snapshot.snapshotManifestUrl, - snapshotManifestUrl: snapshot.snapshotManifestUrl, - snapshotManifestKey: snapshot.snapshotManifestKey, - assetUrlsByFilePath: snapshot.assetUrlsByFilePath, + ) + logger.info("workflow: parsed storage sync handed off to parsed-sync", { + sourceId, + documentId: jobToPrepare.documentId, + revisionKey, + }) + return + } + + await context.run("record-sync-completed", async () => + sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { + revisionKey, + syncStatus: "completed", }), ) - const ready = await context.run("source-ready", async () => markSourceReadyAfterReconciliation({ workspaceId, @@ -168,6 +214,7 @@ async function runPollAndMirrorWorkflow(input: { logger.info("workflow: source parse reconciliation finished", { sourceId, jobId: jobToPrepare.jobId, + revisionKey, status: ready.status, }) } diff --git a/src/domains/sources/source-reconcile-workflow.test.ts b/src/domains/sources/source-reconcile-workflow.test.ts index f1907ef..5757da8 100644 --- a/src/domains/sources/source-reconcile-workflow.test.ts +++ b/src/domains/sources/source-reconcile-workflow.test.ts @@ -24,6 +24,7 @@ function makeSource(overrides: Partial = {}): Source { sizeBytes: 1, status: "parsing", failureReason: null, + failureStage: null, knowhereJobId: "job_1", knowhereDocumentId: null, stagedBlobPathname: null, diff --git a/src/domains/sources/source-row-repository.ts b/src/domains/sources/source-row-repository.ts index 77c56dd..2c06765 100644 --- a/src/domains/sources/source-row-repository.ts +++ b/src/domains/sources/source-row-repository.ts @@ -26,6 +26,7 @@ type SourceUpdate = Partial< | "sizeBytes" | "status" | "failureReason" + | "failureStage" | "knowhereJobId" | "knowhereDocumentId" | "stagedBlobPathname" @@ -83,6 +84,7 @@ type SourceRowRepository = { sourceId: string, reason: string, requiredStatus?: string, + failureStage?: string, ) => Effect.Effect readonly clearStagedBlobEffect: ( workspaceId: string, @@ -193,6 +195,7 @@ const markParsingEffect: SourceRowRepository["markParsingEffect"] = ( knowhereJobId: jobId, knowhereDocumentId: documentId, failureReason: null, + failureStage: null, }, requiredStatus) const markReadyEffect: SourceRowRepository["markReadyEffect"] = ( @@ -204,6 +207,7 @@ const markReadyEffect: SourceRowRepository["markReadyEffect"] = ( status: "ready", knowhereDocumentId: documentId, failureReason: null, + failureStage: null, }, "parsing") const updateRevisionKeyEffect: SourceRowRepository["updateRevisionKeyEffect"] = ( @@ -220,10 +224,12 @@ const markFailedEffect: SourceRowRepository["markFailedEffect"] = ( sourceId: string, reason: string, requiredStatus?: string, + failureStage?: string, ) => updateInWorkspaceEffect(workspaceId, sourceId, { status: "failed", failureReason: reason, + failureStage: failureStage ?? null, }, requiredStatus) const clearStagedBlobEffect: SourceRowRepository["clearStagedBlobEffect"] = ( diff --git a/src/domains/sources/upload.test.ts b/src/domains/sources/upload.test.ts index 912ab9d..77c9022 100644 --- a/src/domains/sources/upload.test.ts +++ b/src/domains/sources/upload.test.ts @@ -23,6 +23,7 @@ function makeSource(overrides: Partial = {}): Source { sizeBytes: 12, status: "uploading", failureReason: null, + failureStage: null, knowhereJobId: null, knowhereDocumentId: null, stagedBlobPathname: null, diff --git a/src/domains/sources/view.test.ts b/src/domains/sources/view.test.ts index e854c29..cc5b223 100644 --- a/src/domains/sources/view.test.ts +++ b/src/domains/sources/view.test.ts @@ -12,6 +12,7 @@ function makeSource(overrides: Partial = {}): Source { sizeBytes: 1, status: "ready", failureReason: null, + failureStage: null, knowhereJobId: "job_1", knowhereDocumentId: "doc_1", stagedBlobPathname: null, diff --git a/src/domains/sources/workflow-runtime.test.ts b/src/domains/sources/workflow-runtime.test.ts index 3cc714a..08e9853 100644 --- a/src/domains/sources/workflow-runtime.test.ts +++ b/src/domains/sources/workflow-runtime.test.ts @@ -31,6 +31,7 @@ function makeSource(status: Source["status"]): Source { sizeBytes: 1024, status, failureReason: status === "failed" ? "failed" : null, + failureStage: null, knowhereJobId: status === "parsing" ? "job_1" : null, knowhereDocumentId: status === "ready" ? "document_1" : null, stagedBlobPathname: null, diff --git a/src/domains/sources/workflow-runtime.ts b/src/domains/sources/workflow-runtime.ts index 8527fab..24464e0 100644 --- a/src/domains/sources/workflow-runtime.ts +++ b/src/domains/sources/workflow-runtime.ts @@ -2,6 +2,7 @@ import "server-only" import { databaseRuntime } from "@/domains/workspace/database-runtime" import { sourceRepository } from "./repository" +import type { SourceParseSnapshotMetadata } from "./source-parse-result-repository" import type { Source, SourceParseResult } from "@/infrastructure/db/schema" import type { UploadSourceRepository } from "./upload" @@ -17,6 +18,10 @@ type SaveSourceParseResultInput = Parameters< typeof sourceRepository.saveParseResultEffect >[2] +type UpdateSyncStatusInput = Parameters< + typeof sourceRepository.updateSyncStatusEffect +>[2] + type UpsertMaterializedDemoSourceInput = Parameters< typeof sourceRepository.upsertMaterializedDemoSourceEffect >[1] @@ -38,6 +43,7 @@ type UploadRepositoryRuntime = { sourceId: string, reason: string, requiredStatus?: string, + failureStage?: string, ) => Promise } @@ -60,21 +66,11 @@ type SourceWorkflowRuntime = UploadRepositoryRuntime & { readonly getParseResultProgress: ( workspaceId: string, sourceId: string, - ) => Promise<{ - readonly resultBlobUrl: string - readonly snapshotManifestUrl?: string | null - readonly snapshotManifestKey?: string | null - readonly assetUrlsByFilePath: Readonly> - } | null> + ) => Promise readonly getParseSnapshotMetadata: ( workspaceId: string, sourceId: string, - ) => Promise<{ - readonly resultBlobUrl: string - readonly snapshotManifestUrl?: string | null - readonly snapshotManifestKey?: string | null - readonly assetUrlsByFilePath: Readonly> - } | null> + ) => Promise readonly listForWorkspace: (workspaceId: string) => Promise readonly localizeRemoteDocument: ( workspaceId: string, @@ -105,6 +101,11 @@ type SourceWorkflowRuntime = UploadRepositoryRuntime & { sourceId: string, input: SaveSourceParseResultInput, ) => Promise + readonly updateSyncStatus: ( + workspaceId: string, + sourceId: string, + input: UpdateSyncStatusInput, + ) => Promise readonly softDelete: ( workspaceId: string, sourceId: string, @@ -200,9 +201,16 @@ const markFailed: SourceWorkflowRuntime["markFailed"] = ( sourceId: string, reason: string, requiredStatus?: string, + failureStage?: string, ) => databaseRuntime.runPromise( - sourceRepository.markFailedEffect(workspaceId, sourceId, reason, requiredStatus), + sourceRepository.markFailedEffect( + workspaceId, + sourceId, + reason, + requiredStatus, + failureStage, + ), ) const clearStagedBlob: SourceWorkflowRuntime["clearStagedBlob"] = ( @@ -245,6 +253,15 @@ const mergeParseAssetUrls: SourceWorkflowRuntime["mergeParseAssetUrls"] = ( sourceRepository.mergeParseAssetUrlsEffect(workspaceId, sourceId, input), ) +const updateSyncStatus: SourceWorkflowRuntime["updateSyncStatus"] = ( + workspaceId: string, + sourceId: string, + input: UpdateSyncStatusInput, +) => + databaseRuntime.runPromise( + sourceRepository.updateSyncStatusEffect(workspaceId, sourceId, input), + ) + const getParseResultProgress: SourceWorkflowRuntime["getParseResultProgress"] = (workspaceId: string, sourceId: string) => databaseRuntime.runPromise( @@ -316,5 +333,6 @@ export const sourceWorkflowRuntime: SourceWorkflowRuntime = { mergeParseAssetUrls, saveParseResult, softDelete, + updateSyncStatus, upsertMaterializedDemoSource, } diff --git a/src/domains/workspace/initial-state.test.ts b/src/domains/workspace/initial-state.test.ts index 9ea24e3..23bd50f 100644 --- a/src/domains/workspace/initial-state.test.ts +++ b/src/domains/workspace/initial-state.test.ts @@ -647,6 +647,7 @@ function makeSource( sizeBytes: 1024, status: "ready", failureReason: null, + failureStage: null, knowhereJobId: "job_1", knowhereDocumentId: "document_1", stagedBlobPathname: null, diff --git a/src/domains/workspace/integration.test.ts b/src/domains/workspace/integration.test.ts index 9d27083..8b6d80b 100644 --- a/src/domains/workspace/integration.test.ts +++ b/src/domains/workspace/integration.test.ts @@ -145,7 +145,7 @@ describeIfDb("workspace helpers — integration", () => { markSourceReady: sourceWorkflowRuntime.markReady, markSourceFailed: sourceWorkflowRuntime.markFailed, saveSourceParseResult: sourceWorkflowRuntime.saveParseResult, - getParseAssetUrls: sourceService.getParseAssetUrls, + getParseAssetUrls: sourceWorkflowRuntime.getParseAssetUrls, hideDemoSource: sourceService.hideDemoSource, listHiddenDemoSourceIds: sourceService.listHiddenDemoSourceIds, upsertMaterializedDemoSource: diff --git a/src/infrastructure/db/schema.ts b/src/infrastructure/db/schema.ts index 8c909a9..c24b27c 100644 --- a/src/infrastructure/db/schema.ts +++ b/src/infrastructure/db/schema.ts @@ -66,6 +66,8 @@ export type NewWorkspace = typeof workspaces.$inferInsert; * - `size_bytes` — original upload size (for display + quota) * - `status` — lifecycle: uploading | parsing | ready | failed * - `failure_reason` — human-readable error text when status=failed + * - `failure_stage` — which stage failed: parse | storage_sync; drives + * whether a retry reparses or only resumes storage sync * - `knowhere_job_id` — set once the parse job is created * - `knowhere_document_id` — set when parsing completes; used to import * parsed snapshots and to exclude a source from a @@ -95,6 +97,7 @@ export const sources = pgTable( sizeBytes: bigint("size_bytes", { mode: "number" }).notNull(), status: text("status").notNull(), failureReason: text("failure_reason"), + failureStage: text("failure_stage"), knowhereJobId: text("knowhere_job_id"), knowhereDocumentId: text("knowhere_document_id"), stagedBlobPathname: text("staged_blob_pathname"), @@ -162,8 +165,18 @@ export type NewDemoSourceVisibility = typeof demoSourceVisibilities.$inferInsert * Notebook-owned parse-result artifact index for one source. * * Blob is the Notebook-owned read model for parsed chunks after source - * reconciliation completes. This row stores the current parsed snapshot - * manifest and the file-path-to-public-URL map for parsed media artifacts. + * reconciliation completes. The parsed snapshot itself (manifest, chunk pages, + * assets, and resumable sync progress) lives in Vercel Blob under + * `workspaces/{ws}/parsed-documents/{documentId}/{revisionKey}/...`, managed by + * the SDK `ParsedDocumentStorage`. This row records: + * - `revision_key` — current parsed revision (jobResultId ?? jobId); the + * storage fast-path key passed into SDK reads + * - `sync_status` — pending | running | completed | failed for the + * background/parse-time storage sync + * - `sync_error` — last storage-sync error detail when sync_status=failed + * - `result_blob_url` / `snapshot_manifest_*` — legacy columns retained for + * rows written by the pre-migration manifest format + * - `asset_urls` — legacy file-path-to-public-URL map for older rows */ export const sourceParseResults = pgTable( "source_parse_results", @@ -173,9 +186,12 @@ export const sourceParseResults = pgTable( .notNull() .references(() => sources.id, { onDelete: "cascade" }) .unique(), - resultBlobUrl: text("result_blob_url").notNull(), + resultBlobUrl: text("result_blob_url"), snapshotManifestUrl: text("snapshot_manifest_url"), snapshotManifestKey: text("snapshot_manifest_key"), + revisionKey: text("revision_key"), + syncStatus: text("sync_status"), + syncError: text("sync_error"), assetUrls: jsonb("asset_urls") .$type>>() .notNull(), diff --git a/src/integrations/knowhere.ts b/src/integrations/knowhere.ts index 81f6726..c0d074d 100644 --- a/src/integrations/knowhere.ts +++ b/src/integrations/knowhere.ts @@ -1,5 +1,35 @@ import Knowhere from "@ontos-ai/knowhere-sdk" +import type { + Knowledge, + ParsedDocumentStorageLimits, + ParsedDocumentSyncScheduler, +} from "@ontos-ai/knowhere-sdk" import { logger } from "@/lib/logger" +import { BlobParsedDocumentStorage } from "@/domains/sources/parsed-document-blob-storage" + +/** + * Vercel-safe defaults for parsed-document reads and background sync. Each read + * miss and each background sync step is bounded by page count and a deadline so + * a single serverless invocation stays well under the platform ceiling; the + * SDK returns `completed:false` and the caller re-enqueues to continue. + */ +const defaultParsedStorageLimits: ParsedDocumentStorageLimits = { + chunkPageSize: 200, + remotePageSize: 100, + maxPagesPerSync: 10, + maxAssetsPerSync: 20, + syncDeadlineMs: 8000, + grepMaxPages: 50, + grepDeadlineMs: 8000, + outlineMaxPages: 50, + outlineDeadlineMs: 8000, +} + +type ParsedStorageOptions = { + readonly workspaceId: string + readonly scheduler?: ParsedDocumentSyncScheduler + readonly limits?: ParsedDocumentStorageLimits +} /** * Create a Knowhere client with the given API key. @@ -14,6 +44,29 @@ export function makeKnowhereClient(apiKey: string): Knowhere { return wrapKnowhereClient(client) } +/** + * Create a Knowhere client plus a `Knowledge` configured with a Vercel-Blob + * `ParsedDocumentStorage`. Reads through `knowledge` serve from Blob first and + * fall back to Knowhere remote transparently; the `scheduler` (when provided) + * backfills Blob in the background. Use `client` for retrieval/documents/jobs. + * + * `withParsedStorage` is invoked through the logging Proxy so `this` binds to + * the real Knowledge; the returned Knowledge wraps the unwrapped inner client + * (its internal `documents.listChunks` calls are not logged, which is fine). + */ +export function makeKnowhereClientWithParsedStorage( + apiKey: string, + options: ParsedStorageOptions, +): { readonly client: Knowhere; readonly knowledge: Knowledge } { + const client = makeKnowhereClient(apiKey) + const knowledge = client.knowledge.withParsedStorage({ + storage: new BlobParsedDocumentStorage({ workspaceId: options.workspaceId }), + scheduler: options.scheduler, + limits: options.limits ?? defaultParsedStorageLimits, + }) + return { client, knowledge } +} + function wrapKnowhereClient(client: Knowhere): Knowhere { return new Proxy(client, { get(target, prop, receiver) { From 012d729526abeca8e92bf421f7fc639d3c0cab90 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 5 Jul 2026 14:09:21 +0800 Subject: [PATCH 11/51] Add parsed sync capacity guard --- drizzle/0012_lazy_wendell_rand.sql | 21 + drizzle/meta/0012_snapshot.json | 932 ++++++++++++++++++ drizzle/meta/_journal.json | 7 + .../parsed-document-sync-capacity.test.ts | 53 + .../sources/parsed-document-sync-capacity.ts | 304 ++++++ .../sources/parsed-document-sync-scheduler.ts | 4 + .../parsed-sync-route-workflow.test.ts | 113 ++- .../sources/parsed-sync-route-workflow.ts | 166 +++- .../source-reconcile-route-workflow.test.ts | 130 ++- .../source-reconcile-route-workflow.ts | 180 ++-- src/infrastructure/db/schema.ts | 55 ++ 11 files changed, 1853 insertions(+), 112 deletions(-) create mode 100644 drizzle/0012_lazy_wendell_rand.sql create mode 100644 drizzle/meta/0012_snapshot.json create mode 100644 src/domains/sources/parsed-document-sync-capacity.test.ts create mode 100644 src/domains/sources/parsed-document-sync-capacity.ts diff --git a/drizzle/0012_lazy_wendell_rand.sql b/drizzle/0012_lazy_wendell_rand.sql new file mode 100644 index 0000000..f50ab54 --- /dev/null +++ b/drizzle/0012_lazy_wendell_rand.sql @@ -0,0 +1,21 @@ +CREATE TABLE "parsed_document_sync_leases" ( + "id" uuid PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL, + "workspace_id" uuid NOT NULL, + "source_id" uuid NOT NULL, + "document_id" text NOT NULL, + "revision_key" text, + "lease_token" text NOT NULL, + "acquired_at" timestamp with time zone DEFAULT now() NOT NULL, + "expires_at" timestamp with time zone NOT NULL, + "released_at" timestamp with time zone, + "release_reason" text, + "created_at" timestamp with time zone DEFAULT now() NOT NULL, + "updated_at" timestamp with time zone DEFAULT now() NOT NULL +); +--> statement-breakpoint +ALTER TABLE "parsed_document_sync_leases" ADD CONSTRAINT "parsed_document_sync_leases_workspace_id_workspaces_id_fk" FOREIGN KEY ("workspace_id") REFERENCES "public"."workspaces"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint +ALTER TABLE "parsed_document_sync_leases" ADD CONSTRAINT "parsed_document_sync_leases_source_id_sources_id_fk" FOREIGN KEY ("source_id") REFERENCES "public"."sources"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint +CREATE UNIQUE INDEX "parsed_document_sync_leases_token_idx" ON "parsed_document_sync_leases" USING btree ("lease_token");--> statement-breakpoint +CREATE INDEX "parsed_document_sync_leases_active_idx" ON "parsed_document_sync_leases" USING btree ("expires_at") WHERE released_at IS NULL;--> statement-breakpoint +CREATE INDEX "parsed_document_sync_leases_workspace_active_idx" ON "parsed_document_sync_leases" USING btree ("workspace_id") WHERE released_at IS NULL;--> statement-breakpoint +CREATE INDEX "parsed_document_sync_leases_document_active_idx" ON "parsed_document_sync_leases" USING btree ("document_id") WHERE released_at IS NULL; \ No newline at end of file diff --git a/drizzle/meta/0012_snapshot.json b/drizzle/meta/0012_snapshot.json new file mode 100644 index 0000000..4011068 --- /dev/null +++ b/drizzle/meta/0012_snapshot.json @@ -0,0 +1,932 @@ +{ + "id": "a2a2f9a9-d567-4413-89d8-fa714b994351", + "prevId": "bab84445-8eb2-4d15-856e-b2f59edcaf47", + "version": "7", + "dialect": "postgresql", + "tables": { + "public.chat_messages": { + "name": "chat_messages", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "thread_id": { + "name": "thread_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "role": { + "name": "role", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "content": { + "name": "content", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "citations": { + "name": "citations", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "artifacts": { + "name": "artifacts", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "chat_messages_thread_created_idx": { + "name": "chat_messages_thread_created_idx", + "columns": [ + { + "expression": "thread_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "created_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "chat_messages_thread_id_chat_threads_id_fk": { + "name": "chat_messages_thread_id_chat_threads_id_fk", + "tableFrom": "chat_messages", + "tableTo": "chat_threads", + "columnsFrom": [ + "thread_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.chat_threads": { + "name": "chat_threads", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "workspace_id": { + "name": "workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "title": { + "name": "title", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "demo_key": { + "name": "demo_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + } + }, + "indexes": { + "chat_threads_workspace_updated_idx": { + "name": "chat_threads_workspace_updated_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "updated_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "where": "deleted_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + }, + "chat_threads_workspace_demo_key_idx": { + "name": "chat_threads_workspace_demo_key_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "demo_key", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "chat_threads_workspace_id_workspaces_id_fk": { + "name": "chat_threads_workspace_id_workspaces_id_fk", + "tableFrom": "chat_threads", + "tableTo": "workspaces", + "columnsFrom": [ + "workspace_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.demo_source_visibilities": { + "name": "demo_source_visibilities", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "workspace_id": { + "name": "workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "demo_source_id": { + "name": "demo_source_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "hidden_at": { + "name": "hidden_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "demo_source_visibilities_workspace_source_idx": { + "name": "demo_source_visibilities_workspace_source_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "demo_source_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "demo_source_visibilities_workspace_idx": { + "name": "demo_source_visibilities_workspace_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "demo_source_visibilities_workspace_id_workspaces_id_fk": { + "name": "demo_source_visibilities_workspace_id_workspaces_id_fk", + "tableFrom": "demo_source_visibilities", + "tableTo": "workspaces", + "columnsFrom": [ + "workspace_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.parsed_document_sync_leases": { + "name": "parsed_document_sync_leases", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "workspace_id": { + "name": "workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "source_id": { + "name": "source_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "document_id": { + "name": "document_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "revision_key": { + "name": "revision_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "lease_token": { + "name": "lease_token", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "acquired_at": { + "name": "acquired_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "expires_at": { + "name": "expires_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true + }, + "released_at": { + "name": "released_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "release_reason": { + "name": "release_reason", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "parsed_document_sync_leases_token_idx": { + "name": "parsed_document_sync_leases_token_idx", + "columns": [ + { + "expression": "lease_token", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "parsed_document_sync_leases_active_idx": { + "name": "parsed_document_sync_leases_active_idx", + "columns": [ + { + "expression": "expires_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "released_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + }, + "parsed_document_sync_leases_workspace_active_idx": { + "name": "parsed_document_sync_leases_workspace_active_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "released_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + }, + "parsed_document_sync_leases_document_active_idx": { + "name": "parsed_document_sync_leases_document_active_idx", + "columns": [ + { + "expression": "document_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "released_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "parsed_document_sync_leases_workspace_id_workspaces_id_fk": { + "name": "parsed_document_sync_leases_workspace_id_workspaces_id_fk", + "tableFrom": "parsed_document_sync_leases", + "tableTo": "workspaces", + "columnsFrom": [ + "workspace_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + }, + "parsed_document_sync_leases_source_id_sources_id_fk": { + "name": "parsed_document_sync_leases_source_id_sources_id_fk", + "tableFrom": "parsed_document_sync_leases", + "tableTo": "sources", + "columnsFrom": [ + "source_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.source_parse_results": { + "name": "source_parse_results", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "source_id": { + "name": "source_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "result_blob_url": { + "name": "result_blob_url", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "snapshot_manifest_url": { + "name": "snapshot_manifest_url", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "snapshot_manifest_key": { + "name": "snapshot_manifest_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "revision_key": { + "name": "revision_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "sync_status": { + "name": "sync_status", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "sync_error": { + "name": "sync_error", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "asset_urls": { + "name": "asset_urls", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "source_parse_results_source_id_idx": { + "name": "source_parse_results_source_id_idx", + "columns": [ + { + "expression": "source_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "source_parse_results_source_id_sources_id_fk": { + "name": "source_parse_results_source_id_sources_id_fk", + "tableFrom": "source_parse_results", + "tableTo": "sources", + "columnsFrom": [ + "source_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": { + "source_parse_results_source_id_unique": { + "name": "source_parse_results_source_id_unique", + "nullsNotDistinct": false, + "columns": [ + "source_id" + ] + } + }, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.sources": { + "name": "sources", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "workspace_id": { + "name": "workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "title": { + "name": "title", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "mime_type": { + "name": "mime_type", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "size_bytes": { + "name": "size_bytes", + "type": "bigint", + "primaryKey": false, + "notNull": true + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "failure_reason": { + "name": "failure_reason", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "failure_stage": { + "name": "failure_stage", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "knowhere_job_id": { + "name": "knowhere_job_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "knowhere_document_id": { + "name": "knowhere_document_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "staged_blob_pathname": { + "name": "staged_blob_pathname", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "staged_blob_url": { + "name": "staged_blob_url", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "original_blob_pathname": { + "name": "original_blob_pathname", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "original_blob_url": { + "name": "original_blob_url", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "demo_key": { + "name": "demo_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + } + }, + "indexes": { + "sources_workspace_created_idx": { + "name": "sources_workspace_created_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "created_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "where": "deleted_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + }, + "sources_workspace_status_idx": { + "name": "sources_workspace_status_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "status", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "sources_workspace_demo_key_idx": { + "name": "sources_workspace_demo_key_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "demo_key", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "sources_workspace_document_idx": { + "name": "sources_workspace_document_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "knowhere_document_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "where": "knowhere_document_id IS NOT NULL AND deleted_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "sources_workspace_id_workspaces_id_fk": { + "name": "sources_workspace_id_workspaces_id_fk", + "tableFrom": "sources", + "tableTo": "workspaces", + "columnsFrom": [ + "workspace_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.workspaces": { + "name": "workspaces", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "namespace": { + "name": "namespace", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "workspaces_user_id_idx": { + "name": "workspaces_user_id_idx", + "columns": [ + { + "expression": "user_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": { + "workspaces_user_id_unique": { + "name": "workspaces_user_id_unique", + "nullsNotDistinct": false, + "columns": [ + "user_id" + ] + }, + "workspaces_namespace_unique": { + "name": "workspaces_namespace_unique", + "nullsNotDistinct": false, + "columns": [ + "namespace" + ] + } + }, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + } + }, + "enums": {}, + "schemas": {}, + "sequences": {}, + "roles": {}, + "policies": {}, + "views": {}, + "_meta": { + "columns": {}, + "schemas": {}, + "tables": {} + } +} \ No newline at end of file diff --git a/drizzle/meta/_journal.json b/drizzle/meta/_journal.json index fc814c6..93f0ced 100644 --- a/drizzle/meta/_journal.json +++ b/drizzle/meta/_journal.json @@ -85,6 +85,13 @@ "when": 1783175614191, "tag": "0011_rare_deadpool", "breakpoints": true + }, + { + "idx": 12, + "version": "7", + "when": 1783231408481, + "tag": "0012_lazy_wendell_rand", + "breakpoints": true } ] } \ No newline at end of file diff --git a/src/domains/sources/parsed-document-sync-capacity.test.ts b/src/domains/sources/parsed-document-sync-capacity.test.ts new file mode 100644 index 0000000..5800e09 --- /dev/null +++ b/src/domains/sources/parsed-document-sync-capacity.test.ts @@ -0,0 +1,53 @@ +import { describe, expect, it } from "vitest" + +import { parsedDocumentSyncCapacityGuard } from "./parsed-document-sync-capacity" + +describe("parsedDocumentSyncCapacityGuard", () => { + it("uses the configured defaults when environment variables are absent", () => { + expect(parsedDocumentSyncCapacityGuard.readPolicy({})).toEqual({ + globalActiveLimit: 100, + workspaceActiveLimit: 5, + documentActiveLimit: 1, + waitSeconds: 60, + }) + }) + + it("reads positive integer capacity settings from the environment", () => { + expect( + parsedDocumentSyncCapacityGuard.readPolicy({ + SYNC_GLOBAL_ACTIVE_LIMIT: "12", + SYNC_WORKSPACE_ACTIVE_LIMIT: "3", + SYNC_DOCUMENT_ACTIVE_LIMIT: "2", + SYNC_WAIT_SECONDS: "90", + }), + ).toEqual({ + globalActiveLimit: 12, + workspaceActiveLimit: 3, + documentActiveLimit: 2, + waitSeconds: 90, + }) + }) + + it("denies capacity by document, workspace, then global cap", () => { + const policy = parsedDocumentSyncCapacityGuard.readPolicy({}) + + expect( + parsedDocumentSyncCapacityGuard.selectDenialReason( + { globalActive: 0, workspaceActive: 0, documentActive: 1 }, + policy, + ), + ).toBe("document") + expect( + parsedDocumentSyncCapacityGuard.selectDenialReason( + { globalActive: 0, workspaceActive: 5, documentActive: 0 }, + policy, + ), + ).toBe("workspace") + expect( + parsedDocumentSyncCapacityGuard.selectDenialReason( + { globalActive: 100, workspaceActive: 0, documentActive: 0 }, + policy, + ), + ).toBe("global") + }) +}) diff --git a/src/domains/sources/parsed-document-sync-capacity.ts b/src/domains/sources/parsed-document-sync-capacity.ts new file mode 100644 index 0000000..cff5584 --- /dev/null +++ b/src/domains/sources/parsed-document-sync-capacity.ts @@ -0,0 +1,304 @@ +import "server-only" + +import { randomUUID } from "node:crypto" + +import { and, eq, isNull, sql } from "drizzle-orm" +import { Effect } from "effect" + +import { databaseRuntime } from "@/domains/workspace/database-runtime" +import { DbClient } from "@/infrastructure/db" +import { parsedDocumentSyncLeases } from "@/infrastructure/db/schema" + +type SyncCapacityPolicy = { + readonly globalActiveLimit: number + readonly workspaceActiveLimit: number + readonly documentActiveLimit: number + readonly waitSeconds: number +} + +type ActiveSyncCounts = { + readonly globalActive: number + readonly workspaceActive: number + readonly documentActive: number +} + +type SyncCapacityDenialReason = "global" | "workspace" | "document" + +type AcquireSyncLeaseInput = { + readonly workspaceId: string + readonly sourceId: string + readonly documentId: string + readonly revisionKey?: string + readonly policy?: SyncCapacityPolicy +} + +type AcquireSyncLeaseResult = + | { + readonly kind: "acquired" + readonly leaseToken: string + readonly activeCounts: ActiveSyncCounts + } + | { + readonly kind: "capacity-full" + readonly reason: SyncCapacityDenialReason + readonly waitSeconds: number + readonly activeCounts: ActiveSyncCounts + } + | { + readonly kind: "source-missing" + } + +type ReleaseSyncLeaseInput = { + readonly leaseToken: string + readonly releaseReason: SyncLeaseReleaseReason +} + +type SyncLeaseReleaseReason = "completed" | "incomplete" | "failed" + +type AcquireSyncLeaseRow = { + readonly leaseId: string | null + readonly leaseToken: string | null + readonly hasSource: boolean + readonly globalActive: number | string + readonly workspaceActive: number | string + readonly documentActive: number | string +} + +type RawRowsResult = + | readonly Row[] + | { + readonly rows: readonly Row[] + } + +const defaultGlobalActiveLimit: number = 100 +const defaultWorkspaceActiveLimit: number = 5 +const defaultDocumentActiveLimit: number = 1 +const defaultWaitSeconds: number = 60 +const leaseDurationSeconds: number = 10 * 60 + +const capacityLockClassId: number = 24_071_105 +const capacityLockObjectId: number = 2_607_198 + +function readPolicy( + env: Readonly> = process.env, +): SyncCapacityPolicy { + return { + globalActiveLimit: readPositiveIntegerEnv( + "SYNC_GLOBAL_ACTIVE_LIMIT", + defaultGlobalActiveLimit, + env, + ), + workspaceActiveLimit: readPositiveIntegerEnv( + "SYNC_WORKSPACE_ACTIVE_LIMIT", + defaultWorkspaceActiveLimit, + env, + ), + documentActiveLimit: readPositiveIntegerEnv( + "SYNC_DOCUMENT_ACTIVE_LIMIT", + defaultDocumentActiveLimit, + env, + ), + waitSeconds: readPositiveIntegerEnv( + "SYNC_WAIT_SECONDS", + defaultWaitSeconds, + env, + ), + } +} + +function readPositiveIntegerEnv( + name: string, + defaultValue: number, + env: Readonly>, +): number { + const rawValue = env[name] + if (!rawValue) return defaultValue + + const value = Number(rawValue) + if (!Number.isInteger(value) || value < 1) return defaultValue + + return value +} + +function selectDenialReason( + counts: ActiveSyncCounts, + policy: SyncCapacityPolicy, +): SyncCapacityDenialReason | null { + if (counts.documentActive >= policy.documentActiveLimit) return "document" + if (counts.workspaceActive >= policy.workspaceActiveLimit) return "workspace" + if (counts.globalActive >= policy.globalActiveLimit) return "global" + + return null +} + +const acquireEffect = ( + input: AcquireSyncLeaseInput, +): Effect.Effect => + Effect.gen(function* () { + const db = yield* DbClient + const policy = input.policy ?? readPolicy() + const leaseToken = randomUUID() + const result = yield* Effect.promise(() => + db.execute(sql` + WITH lock AS ( + SELECT pg_advisory_xact_lock(${capacityLockClassId}, ${capacityLockObjectId}) + ), + expired AS ( + UPDATE parsed_document_sync_leases + SET released_at = now(), + release_reason = 'expired', + updated_at = now() + WHERE released_at IS NULL + AND expires_at <= now() + RETURNING id + ), + counts AS ( + SELECT + EXISTS ( + SELECT 1 + FROM sources + WHERE id = ${input.sourceId}::uuid + AND workspace_id = ${input.workspaceId}::uuid + AND deleted_at IS NULL + ) AS "hasSource", + ( + SELECT count(*)::int + FROM parsed_document_sync_leases + WHERE released_at IS NULL + AND expires_at > now() + ) AS "globalActive", + ( + SELECT count(*)::int + FROM parsed_document_sync_leases + WHERE workspace_id = ${input.workspaceId}::uuid + AND released_at IS NULL + AND expires_at > now() + ) AS "workspaceActive", + ( + SELECT count(*)::int + FROM parsed_document_sync_leases + WHERE document_id = ${input.documentId} + AND released_at IS NULL + AND expires_at > now() + ) AS "documentActive" + FROM lock + ), + inserted AS ( + INSERT INTO parsed_document_sync_leases ( + workspace_id, + source_id, + document_id, + revision_key, + lease_token, + expires_at + ) + SELECT + ${input.workspaceId}::uuid, + ${input.sourceId}::uuid, + ${input.documentId}, + ${input.revisionKey ?? null}, + ${leaseToken}, + now() + (${leaseDurationSeconds} * interval '1 second') + FROM counts + WHERE "hasSource" + AND "globalActive" < ${policy.globalActiveLimit} + AND "workspaceActive" < ${policy.workspaceActiveLimit} + AND "documentActive" < ${policy.documentActiveLimit} + RETURNING id AS "leaseId", lease_token AS "leaseToken" + ) + SELECT + (SELECT "leaseId" FROM inserted) AS "leaseId", + (SELECT "leaseToken" FROM inserted) AS "leaseToken", + counts."hasSource", + counts."globalActive", + counts."workspaceActive", + counts."documentActive" + FROM counts + `), + ) + const row = getFirstRow(getRawRows(result)) + if (!row) return { kind: "source-missing" } + if (!row.hasSource) return { kind: "source-missing" } + + const activeCounts = normalizeActiveCounts(row) + if (row.leaseToken) { + return { + kind: "acquired", + leaseToken: row.leaseToken, + activeCounts, + } + } + + return { + kind: "capacity-full", + reason: selectDenialReason(activeCounts, policy) ?? "global", + waitSeconds: policy.waitSeconds, + activeCounts, + } + }) + +const releaseEffect = ( + input: ReleaseSyncLeaseInput, +): Effect.Effect => + Effect.gen(function* () { + const db = yield* DbClient + yield* Effect.promise(() => + db + .update(parsedDocumentSyncLeases) + .set({ + releasedAt: sql`now()`, + releaseReason: input.releaseReason, + updatedAt: sql`now()`, + }) + .where( + and( + eq(parsedDocumentSyncLeases.leaseToken, input.leaseToken), + isNull(parsedDocumentSyncLeases.releasedAt), + ), + ), + ) + }) + +function normalizeActiveCounts(row: AcquireSyncLeaseRow): ActiveSyncCounts { + return { + globalActive: normalizeCount(row.globalActive), + workspaceActive: normalizeCount(row.workspaceActive), + documentActive: normalizeCount(row.documentActive), + } +} + +function normalizeCount(value: number | string): number { + return typeof value === "number" ? value : Number(value) +} + +function getRawRows(value: RawRowsResult): readonly Row[] { + if (isReadonlyArray(value)) return value + return value.rows +} + +function getFirstRow(rows: readonly Row[]): Row | undefined { + return rows[0] +} + +function isReadonlyArray( + value: RawRowsResult, +): value is readonly Row[] { + return Array.isArray(value) +} + +async function acquire( + input: AcquireSyncLeaseInput, +): Promise { + return databaseRuntime.runPromise(acquireEffect(input)) +} + +async function release(input: ReleaseSyncLeaseInput): Promise { + await databaseRuntime.runPromise(releaseEffect(input)) +} + +export const parsedDocumentSyncCapacityGuard = { + acquire, + readPolicy, + release, + selectDenialReason, +} diff --git a/src/domains/sources/parsed-document-sync-scheduler.ts b/src/domains/sources/parsed-document-sync-scheduler.ts index 5d2957d..c0ee4ce 100644 --- a/src/domains/sources/parsed-document-sync-scheduler.ts +++ b/src/domains/sources/parsed-document-sync-scheduler.ts @@ -24,6 +24,7 @@ export type ParsedSyncTrigger = (input: { readonly url: string readonly body: ParsedSyncPayload readonly workflowRunId: string + readonly delaySeconds?: number }) => Promise export type ParsedSyncPayload = { @@ -72,6 +73,7 @@ const defaultTrigger: ParsedSyncTrigger = async (input) => { body: input.body, workflowRunId: input.workflowRunId, retries: 3, + delay: input.delaySeconds, }) } @@ -91,6 +93,7 @@ export async function enqueueParsedDocumentSync( readonly documentId: string readonly apiKey: string readonly revisionKey?: string + readonly delaySeconds?: number }, trigger: ParsedSyncTrigger = defaultTrigger, ): Promise { @@ -102,6 +105,7 @@ export async function enqueueParsedDocumentSync( revisionKey: input.revisionKey ?? "initial", segmentIndex: 0, }), + delaySeconds: input.delaySeconds, }) } diff --git a/src/domains/sources/parsed-sync-route-workflow.test.ts b/src/domains/sources/parsed-sync-route-workflow.test.ts index 555c186..91095f5 100644 --- a/src/domains/sources/parsed-sync-route-workflow.test.ts +++ b/src/domains/sources/parsed-sync-route-workflow.test.ts @@ -1,7 +1,9 @@ -import { afterEach, describe, expect, it, vi } from "vitest" +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest" const mocks = vi.hoisted(() => ({ + acquireSyncCapacity: vi.fn(), makeKnowhereClientWithParsedStorage: vi.fn(), + releaseSyncCapacity: vi.fn(), updateSyncStatus: vi.fn(), findInWorkspace: vi.fn(), markFailed: vi.fn(), @@ -28,6 +30,13 @@ vi.mock("./source-reconcile-workflow", () => ({ markSourceReadyAfterReconciliation: mocks.markSourceReadyAfterReconciliation, })) +vi.mock("./parsed-document-sync-capacity", () => ({ + parsedDocumentSyncCapacityGuard: { + acquire: mocks.acquireSyncCapacity, + release: mocks.releaseSyncCapacity, + }, +})) + vi.mock("@/lib/logger", () => ({ logger: { info: mocks.loggerInfo, error: mocks.loggerError }, })) @@ -53,6 +62,22 @@ const basePayload = { } describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { + beforeEach(() => { + mocks.acquireSyncCapacity.mockResolvedValue({ + kind: "acquired", + leaseToken: "lease_1", + activeCounts: { + globalActive: 0, + workspaceActive: 0, + documentActive: 0, + }, + }) + mocks.releaseSyncCapacity.mockResolvedValue(undefined) + mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ + status: "ready", + }) + }) + afterEach(() => { vi.clearAllMocks() }) @@ -67,9 +92,6 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { client: {}, knowledge: { syncParsedDocument }, }) - mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ - status: "ready", - }) await parsedSyncRouteWorkflow.runParsedSyncWorkflow({ context: createContext(), @@ -88,9 +110,13 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { sourceId: "source_1", documentId: "doc_1", }) + expect(mocks.releaseSyncCapacity).toHaveBeenCalledWith({ + leaseToken: "lease_1", + releaseReason: "completed", + }) }) - it("triggers a continuation and does not mark ready when sync is incomplete", async () => { + it("marks ready and triggers a continuation when sync is incomplete", async () => { const syncParsedDocument = vi.fn(async () => ({ documentId: "doc_1", revisionKey: "rev_1", @@ -119,7 +145,11 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { restore() } - expect(mocks.markSourceReadyAfterReconciliation).not.toHaveBeenCalled() + expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "source_1", + documentId: "doc_1", + }) expect(triggered).toHaveLength(1) expect(triggered[0]?.segmentIndex).toBe(1) expect(triggered[0]?.workflowRunId).toBe("doc_1-sync-rev_1-1") @@ -128,6 +158,10 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { "source_1", { revisionKey: "rev_1", syncStatus: "running" }, ) + expect(mocks.releaseSyncCapacity).toHaveBeenCalledWith({ + leaseToken: "lease_1", + releaseReason: "incomplete", + }) }) it("passes an explicit revisionKey into syncParsedDocument when provided", async () => { @@ -140,9 +174,6 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { client: {}, knowledge: { syncParsedDocument }, }) - mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ - status: "ready", - }) await parsedSyncRouteWorkflow.runParsedSyncWorkflow({ context: createContext(), @@ -154,6 +185,70 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { revisionKey: "rev_9", }) }) + + it("schedules a delayed retry when capacity is full", async () => { + const syncParsedDocument = vi.fn(async () => ({ + documentId: "doc_1", + revisionKey: "rev_1", + completed: true, + })) + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: {}, + knowledge: { syncParsedDocument }, + }) + mocks.acquireSyncCapacity.mockResolvedValue({ + kind: "capacity-full", + reason: "workspace", + waitSeconds: 60, + activeCounts: { + globalActive: 10, + workspaceActive: 5, + documentActive: 0, + }, + }) + const triggered: Array<{ + readonly workflowRunId: string + readonly segmentIndex?: number + readonly delaySeconds?: number + }> = [] + const restore = parsedSyncRouteWorkflow.setContinuationTriggerForTesting( + async (input) => { + triggered.push({ + workflowRunId: input.workflowRunId, + segmentIndex: input.payload.segmentIndex, + delaySeconds: input.delaySeconds, + }) + }, + ) + + try { + await parsedSyncRouteWorkflow.runParsedSyncWorkflow({ + context: createContext(), + payload: { ...basePayload, revisionKey: "rev_1" }, + }) + } finally { + restore() + } + + expect(syncParsedDocument).not.toHaveBeenCalled() + expect(mocks.updateSyncStatus).toHaveBeenCalledWith( + "workspace_1", + "source_1", + { + revisionKey: "rev_1", + syncStatus: "pending", + syncError: null, + }, + ) + expect(triggered).toEqual([ + { + workflowRunId: "doc_1-sync-rev_1-1", + segmentIndex: 1, + delaySeconds: 60, + }, + ]) + expect(mocks.releaseSyncCapacity).not.toHaveBeenCalled() + }) }) describe("parsedSyncRouteWorkflow.markSyncFailedAfterWorkflowFailure", () => { diff --git a/src/domains/sources/parsed-sync-route-workflow.ts b/src/domains/sources/parsed-sync-route-workflow.ts index c33ace8..6be3b23 100644 --- a/src/domains/sources/parsed-sync-route-workflow.ts +++ b/src/domains/sources/parsed-sync-route-workflow.ts @@ -9,6 +9,7 @@ import { getParsedSyncWorkflowRunId, type ParsedSyncPayload, } from "./parsed-document-sync-scheduler" +import { parsedDocumentSyncCapacityGuard } from "./parsed-document-sync-capacity" import { markSourceReadyAfterReconciliation } from "./source-reconcile-workflow" import { sourceWorkflowRuntime } from "./workflow-runtime" @@ -30,8 +31,11 @@ type ContinuationTriggerInput = { readonly url: string readonly payload: ParsedSyncPayload readonly workflowRunId: string + readonly delaySeconds?: number } +type SyncLeaseReleaseReason = "completed" | "incomplete" | "failed" + // Sync steps per workflow segment. Each `syncParsedDocument` call is bounded by // the SDK limits (pages + deadline); this caps how many bounded steps we run in // one serverless invocation before handing off to a fresh continuation. @@ -68,37 +72,45 @@ async function runParsedSyncWorkflow(input: { workspaceId, }) + const preSyncReady = await context.run( + `source-ready-before-sync-${payload.segmentIndex}`, + async () => + markSourceReadyAfterReconciliation({ + workspaceId, + sourceId, + documentId, + }), + ) + if (preSyncReady.status === "gone") return + let revisionKey = payload.revisionKey let completed = false + let releaseReason: SyncLeaseReleaseReason = "incomplete" - for (let step = 0; step < maxSyncStepsPerSegment; step++) { - const result: KnowledgeSyncParsedDocumentResponse = await context.run( - `sync-${payload.segmentIndex}-${step}`, - async () => - knowledge.syncParsedDocument({ - documentId, - ...(revisionKey ? { revisionKey } : {}), - }), - ) - revisionKey = result.revisionKey - - await context.run(`record-progress-${payload.segmentIndex}-${step}`, () => - sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { + const capacity = await context.run( + `acquire-sync-capacity-${payload.segmentIndex}`, + async () => + parsedDocumentSyncCapacityGuard.acquire({ + workspaceId, + sourceId, + documentId, revisionKey, - syncStatus: result.completed ? "completed" : "running", }), + ) + if (capacity.kind === "source-missing") return + if (capacity.kind === "capacity-full") { + await context.run( + `record-capacity-wait-${payload.segmentIndex}`, + async () => + sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { + revisionKey, + syncStatus: "pending", + syncError: null, + }), ) - - if (result.completed) { - completed = true - break - } - } - - if (!completed) { const nextSegmentIndex = payload.segmentIndex + 1 await context.run( - `trigger-sync-continuation-${nextSegmentIndex}`, + `trigger-sync-capacity-retry-${nextSegmentIndex}`, async () => triggerContinuation({ url: context.url, @@ -115,33 +127,116 @@ async function runParsedSyncWorkflow(input: { revisionKey: revisionKey ?? "initial", segmentIndex: nextSegmentIndex, }), + delaySeconds: capacity.waitSeconds, }), ) - logger.info("parsed-sync: continuation triggered", { + logger.info("parsed-sync: capacity full; retry scheduled", { sourceId, documentId, - segmentIndex: nextSegmentIndex, + reason: capacity.reason, + waitSeconds: capacity.waitSeconds, + activeCounts: capacity.activeCounts, }) return } - // A parsing source is only readied once its parsed snapshot is fully synced. - // For an already-ready source (read-miss backfill) markReady is a no-op guard. - const ready = await context.run("source-ready", async () => - markSourceReadyAfterReconciliation({ - workspaceId, - sourceId, - documentId, - }), - ) + try { + for (let step = 0; step < maxSyncStepsPerSegment; step++) { + const result: KnowledgeSyncParsedDocumentResponse = await context.run( + `sync-${payload.segmentIndex}-${step}`, + async () => + knowledge.syncParsedDocument({ + documentId, + ...(revisionKey ? { revisionKey } : {}), + }), + ) + revisionKey = result.revisionKey + + await context.run(`record-progress-${payload.segmentIndex}-${step}`, () => + sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { + revisionKey, + syncStatus: result.completed ? "completed" : "running", + }), + ) + + if (result.completed) { + completed = true + releaseReason = "completed" + break + } + } + + if (!completed) { + const nextSegmentIndex = payload.segmentIndex + 1 + await context.run( + `trigger-sync-continuation-${nextSegmentIndex}`, + async () => + triggerContinuation({ + url: context.url, + payload: { + workspaceId, + sourceId, + documentId, + apiKey, + revisionKey, + segmentIndex: nextSegmentIndex, + }, + workflowRunId: getParsedSyncWorkflowRunId({ + documentId, + revisionKey: revisionKey ?? "initial", + segmentIndex: nextSegmentIndex, + }), + }), + ) + logger.info("parsed-sync: continuation triggered", { + sourceId, + documentId, + segmentIndex: nextSegmentIndex, + }) + return + } + } catch (error) { + releaseReason = "failed" + throw error + } finally { + await context.run(`release-sync-capacity-${payload.segmentIndex}`, async () => + releaseCapacityLease({ + leaseToken: capacity.leaseToken, + releaseReason, + sourceId, + documentId, + }), + ) + } + logger.info("parsed-sync: parsed document sync finished", { sourceId, documentId, revisionKey, - status: ready.status, + status: preSyncReady.status, }) } +async function releaseCapacityLease(input: { + readonly leaseToken: string + readonly releaseReason: SyncLeaseReleaseReason + readonly sourceId: string + readonly documentId: string +}): Promise { + try { + await parsedDocumentSyncCapacityGuard.release({ + leaseToken: input.leaseToken, + releaseReason: input.releaseReason, + }) + } catch (error) { + logger.error("parsed-sync: failed to release capacity lease", { + sourceId: input.sourceId, + documentId: input.documentId, + error: error instanceof Error ? error.message : String(error), + }) + } +} + async function triggerParsedSyncContinuation( input: ContinuationTriggerInput, ): Promise { @@ -155,6 +250,7 @@ async function triggerParsedSyncContinuation( body: input.payload, workflowRunId: input.workflowRunId, retries: 3, + delay: input.delaySeconds, }) } diff --git a/src/domains/sources/source-reconcile-route-workflow.test.ts b/src/domains/sources/source-reconcile-route-workflow.test.ts index 52f9c6b..ee415ff 100644 --- a/src/domains/sources/source-reconcile-route-workflow.test.ts +++ b/src/domains/sources/source-reconcile-route-workflow.test.ts @@ -1,7 +1,9 @@ -import { afterEach, describe, expect, it, vi } from "vitest" +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest" const mocks = vi.hoisted(() => ({ + acquireSyncCapacity: vi.fn(), enqueueParsedDocumentSync: vi.fn(), + releaseSyncCapacity: vi.fn(), updateSyncStatus: vi.fn(), markFailed: vi.fn(), loggerError: vi.fn(), @@ -28,6 +30,13 @@ vi.mock("./parsed-document-sync-scheduler", () => ({ enqueueParsedDocumentSync: mocks.enqueueParsedDocumentSync, })) +vi.mock("./parsed-document-sync-capacity", () => ({ + parsedDocumentSyncCapacityGuard: { + acquire: mocks.acquireSyncCapacity, + release: mocks.releaseSyncCapacity, + }, +})) + vi.mock("@/integrations/knowhere", () => ({ makeKnowhereClientWithParsedStorage: mocks.makeKnowhereClientWithParsedStorage, @@ -43,6 +52,12 @@ vi.mock("@/lib/logger", () => ({ import { sourceReconcileRouteWorkflow } from "./source-reconcile-route-workflow" +const activeCounts = { + globalActive: 0, + workspaceActive: 0, + documentActive: 0, +} + function createClient(overrides: { syncParsedDocument?: ReturnType jobResultId?: string @@ -68,6 +83,18 @@ function createClient(overrides: { } describe("sourceReconcileRouteWorkflow", () => { + beforeEach(() => { + mocks.acquireSyncCapacity.mockResolvedValue({ + kind: "acquired", + leaseToken: "lease_1", + activeCounts, + }) + mocks.releaseSyncCapacity.mockResolvedValue(undefined) + mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ + status: "ready", + }) + }) + afterEach(() => { vi.clearAllMocks() }) @@ -89,7 +116,7 @@ describe("sourceReconcileRouteWorkflow", () => { }) }) - it("syncs the parsed document then marks the source ready", async () => { + it("marks the source ready then syncs the parsed document", async () => { const context = createWorkflowContext() const continuations: ContinuationTriggerInput[] = [] const restore = @@ -103,9 +130,6 @@ describe("sourceReconcileRouteWorkflow", () => { client: wired.client, knowledge: wired.knowledge, }) - mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ - status: "ready", - }) mocks.pollSourceReconciliation.mockResolvedValue({ kind: "ready-to-prepare", jobId: "job_1", @@ -139,11 +163,21 @@ describe("sourceReconcileRouteWorkflow", () => { sourceId: "source_1", documentId: "doc_1", }) + expect(mocks.acquireSyncCapacity).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "source_1", + documentId: "doc_1", + revisionKey: "rev_1", + }) + expect(mocks.releaseSyncCapacity).toHaveBeenCalledWith({ + leaseToken: "lease_1", + releaseReason: "completed", + }) expect(mocks.enqueueParsedDocumentSync).not.toHaveBeenCalled() expect(continuations).toEqual([]) }) - it("hands off to parsed-sync and stays parsing when sync is incomplete", async () => { + it("hands off to parsed-sync when sync is incomplete", async () => { const context = createWorkflowContext() const syncParsedDocument = vi.fn(async () => ({ documentId: "doc_1", @@ -170,7 +204,11 @@ describe("sourceReconcileRouteWorkflow", () => { }), }) - expect(mocks.markSourceReadyAfterReconciliation).not.toHaveBeenCalled() + expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "source_1", + documentId: "doc_1", + }) expect(mocks.enqueueParsedDocumentSync).toHaveBeenCalledWith({ workspaceId: "workspace_1", sourceId: "source_1", @@ -178,9 +216,13 @@ describe("sourceReconcileRouteWorkflow", () => { apiKey: "jwt_1", revisionKey: "rev_1", }) + expect(mocks.releaseSyncCapacity).toHaveBeenCalledWith({ + leaseToken: "lease_1", + releaseReason: "incomplete", + }) }) - it("fails the source with storage_sync stage when sync throws", async () => { + it("records sync failure when Blob sync throws", async () => { const context = createWorkflowContext() const syncParsedDocument = vi.fn(async () => { throw new Error("blob write failed") @@ -207,14 +249,76 @@ describe("sourceReconcileRouteWorkflow", () => { }), ).rejects.toThrow("blob write failed") - expect(mocks.markFailed).toHaveBeenCalledWith( + expect(mocks.updateSyncStatus).toHaveBeenCalledWith( "workspace_1", "source_1", - expect.stringContaining("storage sync failed"), - "parsing", - "storage_sync", + { + revisionKey: "rev_1", + syncStatus: "failed", + syncError: "blob write failed", + }, + ) + expect(mocks.markFailed).not.toHaveBeenCalled() + expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "source_1", + documentId: "doc_1", + }) + expect(mocks.releaseSyncCapacity).toHaveBeenCalledWith( + { + leaseToken: "lease_1", + releaseReason: "failed", + }, ) - expect(mocks.markSourceReadyAfterReconciliation).not.toHaveBeenCalled() + }) + + it("marks the source ready and schedules delayed sync when capacity is full", async () => { + const context = createWorkflowContext() + const wired = createClient({}) + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: wired.client, + knowledge: wired.knowledge, + }) + mocks.acquireSyncCapacity.mockResolvedValue({ + kind: "capacity-full", + reason: "document", + waitSeconds: 60, + activeCounts: { + globalActive: 10, + workspaceActive: 1, + documentActive: 1, + }, + }) + mocks.pollSourceReconciliation.mockResolvedValue({ + kind: "ready-to-prepare", + jobId: "job_1", + documentId: "doc_1", + }) + + await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ + context, + payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + }), + }) + + expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "source_1", + documentId: "doc_1", + }) + expect(wired.knowledge.syncParsedDocument).not.toHaveBeenCalled() + expect(mocks.enqueueParsedDocumentSync).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "source_1", + documentId: "doc_1", + apiKey: "jwt_1", + revisionKey: "rev_1", + delaySeconds: 60, + }) + expect(mocks.releaseSyncCapacity).not.toHaveBeenCalled() }) it("triggers a fresh poll run when Knowhere is still running after the segment budget", async () => { diff --git a/src/domains/sources/source-reconcile-route-workflow.ts b/src/domains/sources/source-reconcile-route-workflow.ts index cfa90b3..1e45127 100644 --- a/src/domains/sources/source-reconcile-route-workflow.ts +++ b/src/domains/sources/source-reconcile-route-workflow.ts @@ -9,6 +9,7 @@ import { import { makeKnowhereClientWithParsedStorage } from "@/integrations/knowhere" import { logger } from "@/lib/logger" import { enqueueParsedDocumentSync } from "./parsed-document-sync-scheduler" +import { parsedDocumentSyncCapacityGuard } from "./parsed-document-sync-capacity" import { sourceWorkflowRuntime } from "./workflow-runtime" type ReconcilePayload = { @@ -40,6 +41,8 @@ type ContinuationTriggerInput = { readonly workflowRunId: string } +type SyncLeaseReleaseReason = "completed" | "incomplete" | "failed" + const maxPollAttempts = 25 const initialDelaySeconds = 3 const maxDelaySeconds = 30 @@ -122,7 +125,6 @@ async function runPollAndMirrorWorkflow(input: { return } - // Resolve the parsed revision and record that storage sync is in progress. const revisionKey = await context.run("resolve-revision-key", async () => { const firstPage = await client.documents.listChunks( jobToPrepare.documentId, @@ -130,87 +132,131 @@ async function runPollAndMirrorWorkflow(input: { ) return firstPage.jobResultId ?? firstPage.jobId ?? jobToPrepare.jobId }) - await context.run("record-sync-running", async () => + await context.run("record-sync-pending", async () => sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { revisionKey, - syncStatus: "running", + syncStatus: "pending", + syncError: null, }), ) - // Sync the parsed snapshot into Blob before the source is considered ready. - // Each step is bounded by SDK limits; if it does not finish in this segment, - // hand off to the resumable parsed-sync workflow and stay `parsing`. A sync - // error stages the failure as `storage_sync` so a retry resumes sync without - // reparsing, then rethrows to fail the workflow run. - let syncCompleted = false - for (let step = 0; step < maxSyncStepsPerReconcile; step++) { - const result = await context.run(`parsed-sync-${step}`, async () => { - try { - return await knowledge.syncParsedDocument({ - documentId: jobToPrepare.documentId, - revisionKey, - }) - } catch (error) { - await sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { - revisionKey, - syncStatus: "failed", - syncError: error instanceof Error ? error.message : String(error), - }) - await sourceWorkflowRuntime.markFailed( - workspaceId, - sourceId, - `Parsed document storage sync failed: ${ - error instanceof Error ? error.message : String(error) - }`, - "parsing", - "storage_sync", - ) - throw error - } - }) - if (result.completed) { - syncCompleted = true - break - } - } + const ready = await context.run("source-ready", async () => + markSourceReadyAfterReconciliation({ + workspaceId, + sourceId, + documentId: jobToPrepare.documentId, + }), + ) + if (ready.status === "gone") return - if (!syncCompleted) { - await context.run("record-sync-progress", async () => - sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { - revisionKey, - syncStatus: "running", - }), - ) - await context.run("enqueue-parsed-sync-continuation", async () => + const capacity = await context.run("acquire-sync-capacity", async () => + parsedDocumentSyncCapacityGuard.acquire({ + workspaceId, + sourceId, + documentId: jobToPrepare.documentId, + revisionKey, + }), + ) + if (capacity.kind === "source-missing") return + if (capacity.kind === "capacity-full") { + await context.run("enqueue-capacity-retry", async () => enqueueParsedDocumentSync({ workspaceId, sourceId, documentId: jobToPrepare.documentId, apiKey, revisionKey, + delaySeconds: capacity.waitSeconds, }), ) - logger.info("workflow: parsed storage sync handed off to parsed-sync", { + logger.info("workflow: parsed storage sync delayed by capacity guard", { sourceId, documentId: jobToPrepare.documentId, revisionKey, + reason: capacity.reason, + waitSeconds: capacity.waitSeconds, + activeCounts: capacity.activeCounts, }) return } + let syncCompleted = false + let releaseReason: SyncLeaseReleaseReason = "incomplete" + try { + await context.run("record-sync-running", async () => + sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { + revisionKey, + syncStatus: "running", + syncError: null, + }), + ) + + for (let step = 0; step < maxSyncStepsPerReconcile; step++) { + const result = await context.run(`parsed-sync-${step}`, async () => { + try { + return await knowledge.syncParsedDocument({ + documentId: jobToPrepare.documentId, + revisionKey, + }) + } catch (error) { + await sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { + revisionKey, + syncStatus: "failed", + syncError: getErrorMessage(error), + }) + throw error + } + }) + if (result.completed) { + syncCompleted = true + releaseReason = "completed" + break + } + } + + if (!syncCompleted) { + await context.run("record-sync-progress", async () => + sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { + revisionKey, + syncStatus: "running", + }), + ) + await context.run("enqueue-parsed-sync-continuation", async () => + enqueueParsedDocumentSync({ + workspaceId, + sourceId, + documentId: jobToPrepare.documentId, + apiKey, + revisionKey, + }), + ) + logger.info("workflow: parsed storage sync handed off to parsed-sync", { + sourceId, + documentId: jobToPrepare.documentId, + revisionKey, + }) + return + } + } catch (error) { + releaseReason = "failed" + throw error + } finally { + await context.run("release-sync-capacity", async () => + releaseCapacityLease({ + leaseToken: capacity.leaseToken, + releaseReason, + sourceId, + documentId: jobToPrepare.documentId, + }), + ) + } + await context.run("record-sync-completed", async () => sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { revisionKey, syncStatus: "completed", }), ) - const ready = await context.run("source-ready", async () => - markSourceReadyAfterReconciliation({ - workspaceId, - sourceId, - documentId: jobToPrepare.documentId, - }), - ) logger.info("workflow: source parse reconciliation finished", { sourceId, jobId: jobToPrepare.jobId, @@ -219,6 +265,26 @@ async function runPollAndMirrorWorkflow(input: { }) } +async function releaseCapacityLease(input: { + readonly leaseToken: string + readonly releaseReason: SyncLeaseReleaseReason + readonly sourceId: string + readonly documentId: string +}): Promise { + try { + await parsedDocumentSyncCapacityGuard.release({ + leaseToken: input.leaseToken, + releaseReason: input.releaseReason, + }) + } catch (error) { + logger.error("workflow: failed to release sync capacity lease", { + sourceId: input.sourceId, + documentId: input.documentId, + error: getErrorMessage(error), + }) + } +} + function normalizeReconcilePayload( payload: ReconcilePayload, ): NormalizedReconcilePayload { @@ -295,6 +361,10 @@ function getSafeFailureReason(value: string): string { return normalized.slice(0, 500) } +function getErrorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error) +} + export const sourceReconcileRouteWorkflow = { getPollWorkflowRunId, markSourceFailedAfterWorkflowFailure, diff --git a/src/infrastructure/db/schema.ts b/src/infrastructure/db/schema.ts index c24b27c..4e9470f 100644 --- a/src/infrastructure/db/schema.ts +++ b/src/infrastructure/db/schema.ts @@ -208,6 +208,61 @@ export const sourceParseResults = pgTable( export type SourceParseResult = typeof sourceParseResults.$inferSelect; export type NewSourceParseResult = typeof sourceParseResults.$inferInsert; +/** + * Durable active-work leases for Notebook-owned parsed-document Blob sync. + * + * The sync workers acquire one row before calling the Knowhere SDK's + * Vercel-Blob mirror. Active rows (`released_at IS NULL`) are counted globally, + * per workspace, and per document so Vercel can scale route invocations without + * allowing one user or one document to consume all sync capacity. Expired rows + * are released during the next acquire attempt; normal workers release in a + * `finally` block after their bounded sync segment exits. + */ +export const parsedDocumentSyncLeases = pgTable( + "parsed_document_sync_leases", + { + id: uuid("id").primaryKey().defaultRandom(), + workspaceId: uuid("workspace_id") + .notNull() + .references(() => workspaces.id, { onDelete: "cascade" }), + sourceId: uuid("source_id") + .notNull() + .references(() => sources.id, { onDelete: "cascade" }), + documentId: text("document_id").notNull(), + revisionKey: text("revision_key"), + leaseToken: text("lease_token").notNull(), + acquiredAt: timestamp("acquired_at", { withTimezone: true }) + .notNull() + .defaultNow(), + expiresAt: timestamp("expires_at", { withTimezone: true }).notNull(), + releasedAt: timestamp("released_at", { withTimezone: true }), + releaseReason: text("release_reason"), + createdAt: timestamp("created_at", { withTimezone: true }) + .notNull() + .defaultNow(), + updatedAt: timestamp("updated_at", { withTimezone: true }) + .notNull() + .defaultNow(), + }, + (t) => [ + uniqueIndex("parsed_document_sync_leases_token_idx").on(t.leaseToken), + index("parsed_document_sync_leases_active_idx") + .on(t.expiresAt) + .where(sql`released_at IS NULL`), + index("parsed_document_sync_leases_workspace_active_idx") + .on(t.workspaceId) + .where(sql`released_at IS NULL`), + index("parsed_document_sync_leases_document_active_idx") + .on(t.documentId) + .where(sql`released_at IS NULL`), + ], +); + +export type ParsedDocumentSyncLease = + typeof parsedDocumentSyncLeases.$inferSelect; +export type NewParsedDocumentSyncLease = + typeof parsedDocumentSyncLeases.$inferInsert; + /** * A chat thread is a conversation within a workspace. `demo_key` is retained * for legacy seeded demo conversations. From 1d0d6a9ec684ac41a4f967a4df45c470a4217e96 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 5 Jul 2026 15:46:58 +0800 Subject: [PATCH 12/51] Update Knowhere SDK to 2.1.1 --- package-lock.json | 10 +++++----- package.json | 2 +- pnpm-lock.yaml | 13 +++++-------- pnpm-workspace.yaml | 5 +---- 4 files changed, 12 insertions(+), 18 deletions(-) diff --git a/package-lock.json b/package-lock.json index b8d1d91..e465210 100644 --- a/package-lock.json +++ b/package-lock.json @@ -12,7 +12,7 @@ "@antv/chart-visualization-skills": "0.1.3", "@effect/platform": "^0.96.1", "@neondatabase/serverless": "^1.1.0", - "@ontos-ai/knowhere-sdk": "^2.0.0", + "@ontos-ai/knowhere-sdk": "^2.1.1", "@radix-ui/react-alert-dialog": "^1.1.15", "@radix-ui/react-checkbox": "^1.3.3", "@radix-ui/react-dialog": "^1.1.15", @@ -2275,16 +2275,16 @@ } }, "node_modules/@ontos-ai/knowhere-sdk": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/@ontos-ai/knowhere-sdk/-/knowhere-sdk-2.0.0.tgz", - "integrity": "sha512-Zhx8mn/8pYyhBaO9BOdQYsVZybGXyq15YctjHSU8TwgNjJEzX5lEiBEkp0kIGFPh3S+/7X3rZGakefEPf5G2tQ==", + "version": "2.1.1", + "resolved": "https://registry.npmjs.org/@ontos-ai/knowhere-sdk/-/knowhere-sdk-2.1.1.tgz", + "integrity": "sha512-K33ylB/QjVYlLgow+/Hw+uNhQbWWrMTIqZ3IAlx5buHSI2EHPVP9+kDsfslwpfffYSn85j9Ab0/E9qTG0QotTw==", "license": "MIT", "dependencies": { "axios": "^1.15.2", "jszip": "^3.10.0" }, "engines": { - "node": ">=20.19.0", + "node": ">=22.13.0", "npm": ">=10.0.0", "pnpm": ">=9.0.0" } diff --git a/package.json b/package.json index c998de4..94d1c3a 100644 --- a/package.json +++ b/package.json @@ -28,7 +28,7 @@ "@effect/platform": "^0.96.1", "@napi-rs/canvas": "^1.0.2", "@neondatabase/serverless": "^1.1.0", - "@ontos-ai/knowhere-sdk": "^2.0.0", + "@ontos-ai/knowhere-sdk": "^2.1.1", "@radix-ui/react-alert-dialog": "^1.1.15", "@radix-ui/react-checkbox": "^1.3.3", "@radix-ui/react-dialog": "^1.1.15", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 94e3f05..48aabe5 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -4,9 +4,6 @@ settings: autoInstallPeers: true excludeLinksFromLockfile: false -overrides: - '@ontos-ai/knowhere-sdk': file:../knowhere-node-sdk - importers: .: @@ -27,8 +24,8 @@ importers: specifier: ^1.1.0 version: 1.1.0 '@ontos-ai/knowhere-sdk': - specifier: file:../knowhere-node-sdk - version: file:../knowhere-node-sdk + specifier: ^2.1.1 + version: 2.1.1 '@radix-ui/react-alert-dialog': specifier: ^1.1.15 version: 1.1.15(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4) @@ -1606,8 +1603,8 @@ packages: resolution: {integrity: sha512-nn5ozdjYQpUCZlWGuxcJY/KpxkWQs4DcbMCmKojjyrYDEAGy4Ce19NN4v5MduafTwJlbKc99UA8YhSVqq9yPZA==} engines: {node: '>=12.4.0'} - '@ontos-ai/knowhere-sdk@file:../knowhere-node-sdk': - resolution: {directory: ../knowhere-node-sdk, type: directory} + '@ontos-ai/knowhere-sdk@2.1.1': + resolution: {integrity: sha512-K33ylB/QjVYlLgow+/Hw+uNhQbWWrMTIqZ3IAlx5buHSI2EHPVP9+kDsfslwpfffYSn85j9Ab0/E9qTG0QotTw==} engines: {node: '>=22.13.0', npm: '>=10.0.0', pnpm: '>=9.0.0'} '@open-draft/deferred-promise@2.2.0': @@ -6734,7 +6731,7 @@ snapshots: '@nolyfill/is-core-module@1.0.39': {} - '@ontos-ai/knowhere-sdk@file:../knowhere-node-sdk': + '@ontos-ai/knowhere-sdk@2.1.1': dependencies: axios: 1.18.1 jszip: 3.10.1 diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index a49c9d9..80ee5bb 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -3,7 +3,4 @@ ignoredBuiltDependencies: - unrs-resolver minimumReleaseAgeExclude: - - '@ontos-ai/knowhere-sdk@2.0.0' - -overrides: - '@ontos-ai/knowhere-sdk': file:../knowhere-node-sdk + - '@ontos-ai/knowhere-sdk@2.1.1' From f23ef69bc64ea40d1cf918ca8bc839ea67746910 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 5 Jul 2026 18:05:59 +0800 Subject: [PATCH 13/51] Align remote source processed row state --- src/components/source-row.test.ts | 38 +++++++++++++++++++++++ src/components/source-row.tsx | 5 ++- src/domains/chat/index.test.ts | 6 ++-- src/domains/chat/retrieval.ts | 9 +++++- src/domains/sources/remote-library.ts | 2 +- src/domains/sources/route-service.test.ts | 4 +-- 6 files changed, 54 insertions(+), 10 deletions(-) diff --git a/src/components/source-row.test.ts b/src/components/source-row.test.ts index 26e413f..427e09a 100644 --- a/src/components/source-row.test.ts +++ b/src/components/source-row.test.ts @@ -49,6 +49,44 @@ describe("SourceRow", () => { expect(screen.getByText("Processed · 3 chunks")).toBeTruthy(); }); + it("labels unlocalized remote documents as processed without chunk counts", () => { + const onToggleIncluded = vi.fn(); + + render( + React.createElement(SourceRow, { + isArchiving: false, + isSelected: false, + onSelect: vi.fn(), + onToggleIncluded, + source: { + id: "knowhere-doc:default:doc_remote", + kind: "remote", + mimeType: "application/pdf", + title: "remote.pdf", + status: "ready", + chunkCount: 4, + documentId: "doc_remote", + excludedFromQuery: false, + }, + }), + ); + + const checkbox = screen.getByRole("checkbox", { + name: "Use remote.pdf in answers", + }); + + expect(screen.getByText("Processed")).toBeTruthy(); + expect(screen.queryByText("Processed · 4 chunks")).toBeNull(); + expect(checkbox.getAttribute("aria-checked")).toBe("true"); + + fireEvent.click(checkbox); + + expect(onToggleIncluded).toHaveBeenCalledWith( + "knowhere-doc:default:doc_remote", + false, + ); + }); + it("shows source archive loading locally", () => { render( React.createElement(SourceRow, { diff --git a/src/components/source-row.tsx b/src/components/source-row.tsx index e964c93..9fb87d3 100644 --- a/src/components/source-row.tsx +++ b/src/components/source-row.tsx @@ -42,7 +42,6 @@ export function SourceRow({ const isFailed = source.status === "failed"; const canRetry = isFailed && source.originalFile !== undefined; const isLibrarySource = source.officialLibrary !== undefined; - const isRemoteSource = source.kind === "remote"; const iconBg = fileIconTint(source.title); @@ -63,7 +62,7 @@ export function SourceRow({ > onToggleIncluded?.(source.id, checked === true) } @@ -193,11 +192,11 @@ export function SourceRow({ function getReadySourceLabel(source: SourceView): string { if (source.officialLibrary !== undefined) return "Official Library"; - if (source.kind === "remote") return "Remote"; return "Processed"; } function getReadySourceStatusText(source: SourceView): string { + if (source.kind === "remote") return getReadySourceLabel(source); if (typeof source.chunkCount !== "number") return getReadySourceLabel(source); return `${getReadySourceLabel(source)} · ${source.chunkCount} chunks`; } diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 1b51180..59899b0 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -66,7 +66,7 @@ describe("answerQuestionWithRetrieval", () => { question: "What does the document say?", namespace: "notebook-workspace", sources, - excludedSourceIds: ["source_2"], + excludedSourceIds: ["source_2", "knowhere-doc:default:doc_remote"], retrieval, generateAnswer, messages: [], @@ -79,13 +79,13 @@ describe("answerQuestionWithRetrieval", () => { topK: 8, useAgentic: true, dataType: 1, - excludeDocumentIds: ["doc_excluded"], + excludeDocumentIds: ["doc_excluded", "doc_remote"], }); expect(generateAnswer).toHaveBeenCalledWith({ question: "What does the document say?", messages: [], sources, - excludedSourceIds: ["source_2"], + excludedSourceIds: ["source_2", "knowhere-doc:default:doc_remote"], searchSources: expect.any(Function), }); expect(answer).toEqual({ diff --git a/src/domains/chat/retrieval.ts b/src/domains/chat/retrieval.ts index b17add9..4d1103f 100644 --- a/src/domains/chat/retrieval.ts +++ b/src/domains/chat/retrieval.ts @@ -1,6 +1,7 @@ import type { RetrievalQueryParams } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" +import { decodeRemoteSourceId } from "@/domains/sources/remote-library" const RETRIEVAL_QUERY_CHAR_LIMIT = 600 @@ -25,10 +26,16 @@ export function excludeDocuments( excludedSourceIds: readonly string[], ): Pick { const excluded = new Set(excludedSourceIds) - const documentIds = sources + const localDocumentIds = sources .filter((source) => excluded.has(source.id)) .map((source) => source.knowhereDocumentId) .filter((documentId): documentId is string => Boolean(documentId)) + const remoteDocumentIds = excludedSourceIds + .map((sourceId) => decodeRemoteSourceId(sourceId)?.documentId) + .filter((documentId): documentId is string => Boolean(documentId)) + const documentIds = Array.from( + new Set([...localDocumentIds, ...remoteDocumentIds]), + ) return documentIds.length > 0 ? { excludeDocumentIds: documentIds } : {} } diff --git a/src/domains/sources/remote-library.ts b/src/domains/sources/remote-library.ts index 87d3fd8..ceb9c80 100644 --- a/src/domains/sources/remote-library.ts +++ b/src/domains/sources/remote-library.ts @@ -280,7 +280,7 @@ function toRemoteSourceView(document: RemoteDocument): SourceView { mimeType: document.mimeType ?? "application/octet-stream", status: document.status, documentId: document.documentId, - excludedFromQuery: true, + excludedFromQuery: false, } } diff --git a/src/domains/sources/route-service.test.ts b/src/domains/sources/route-service.test.ts index 26e60cb..b07439c 100644 --- a/src/domains/sources/route-service.test.ts +++ b/src/domains/sources/route-service.test.ts @@ -301,7 +301,7 @@ describe("source route service", () => { mimeType: "application/pdf", status: "ready", documentId: "doc_default", - excludedFromQuery: true, + excludedFromQuery: false, }, { id: "knowhere-doc:notebook-workspace_1:doc_legacy", @@ -311,7 +311,7 @@ describe("source route service", () => { mimeType: "application/octet-stream", status: "ready", documentId: "doc_legacy", - excludedFromQuery: true, + excludedFromQuery: false, }, ]); }); From b71073d6cb9b49dc34e791ca49beca9ea8c51356 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 5 Jul 2026 22:01:25 +0800 Subject: [PATCH 14/51] fix: harden chat assets on demand --- .../sources/[sourceId]/chunks/route.test.ts | 2 - src/components/workspace-citation-focus.ts | 3 + .../workspace-selected-chunks.test.ts | 52 +++ src/components/workspace-selected-chunks.ts | 15 +- src/components/workspace-shell.tsx | 1 + src/components/workspace-source-state.test.ts | 21 ++ src/components/workspace-source-state.ts | 23 ++ src/components/workspace-source-workflow.ts | 6 + src/domains/chat/contracts.ts | 4 +- src/domains/chat/index.test.ts | 60 ++-- src/domains/chat/index.ts | 4 +- .../chat/media-asset-hardening.test.ts | 32 +- src/domains/chat/media-asset-hardening.ts | 123 +++++--- src/domains/chat/media-assets.test.ts | 50 ++- src/domains/chat/media-assets.ts | 295 +++++++----------- src/domains/chat/page-citation-assets.test.ts | 33 +- src/domains/chat/page-citation-assets.ts | 56 +--- src/domains/chat/route-answer.ts | 148 +++++++-- src/domains/chat/route-service.test.ts | 103 +++--- src/domains/chat/service.ts | 4 +- src/domains/chunks/index.ts | 4 +- src/domains/chunks/read.test.ts | 4 +- src/domains/chunks/read.ts | 60 +--- .../sources/source-reconcile-workflow.test.ts | 64 +++- .../sources/source-reconcile-workflow.ts | 53 +++- 25 files changed, 707 insertions(+), 513 deletions(-) diff --git a/src/app/api/sources/[sourceId]/chunks/route.test.ts b/src/app/api/sources/[sourceId]/chunks/route.test.ts index 85fb1bd..2f0851c 100644 --- a/src/app/api/sources/[sourceId]/chunks/route.test.ts +++ b/src/app/api/sources/[sourceId]/chunks/route.test.ts @@ -498,7 +498,6 @@ describe("GET /api/sources/[sourceId]/chunks", () => { revisionKey: "job_1", page: 1, pageSize: 1, - assetUrlPolicy: "durable", }) }) @@ -660,7 +659,6 @@ describe("GET /api/sources/[sourceId]/chunks", () => { revisionKey: "job_result_1", page: 1, pageSize: 1, - assetUrlPolicy: "durable", }) }) }) diff --git a/src/components/workspace-citation-focus.ts b/src/components/workspace-citation-focus.ts index 3485dd2..357cd4e 100644 --- a/src/components/workspace-citation-focus.ts +++ b/src/components/workspace-citation-focus.ts @@ -21,6 +21,7 @@ type PrefetchedChunksUpdater = ( type WorkspaceCitationFocusInput = { readonly fetchChunks: (sourceId: string) => Promise readonly initialPrefetchedChunksBySourceId?: PrefetchedChunksBySourceId + readonly onRemoteSourceChunksLoaded?: (sourceId: string) => void readonly onSelectSource: (sourceId: string | null) => void readonly selectedSourceId: string | null readonly sources: readonly SourceView[] @@ -51,6 +52,7 @@ type WorkspaceCitationFocus = { export function useWorkspaceCitationFocus({ fetchChunks, initialPrefetchedChunksBySourceId = {}, + onRemoteSourceChunksLoaded, onSelectSource, selectedSourceId, sources, @@ -88,6 +90,7 @@ export function useWorkspaceCitationFocus({ selectedSourceId, sources, prefetchedChunksBySourceId, + onRemoteSourceChunksLoaded, }) const requestChunkFocus = useCallback( diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index fa08df5..45db25a 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -161,6 +161,58 @@ describe("useWorkspaceSelectedChunks", () => { expect(result.current.selectedChunks).toEqual([]); }); + it("requests a source refresh after loading an unlocalized remote source", async () => { + const onRemoteSourceChunksLoaded = vi.fn(); + const remoteSource: SourceView = { + ...readySource, + id: "knowhere-doc:default:doc_remote", + kind: "remote", + documentId: "doc_remote", + excludedFromQuery: false, + }; + fetchChunkPageMock.mockResolvedValue({ + chunks: [ + { + chunkId: "chunk_1", + type: "text", + content: "Remote content", + sourceTitle: "Remote.pdf", + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 1, + totalPages: 1, + }, + }); + + const { rerender } = renderHook( + (input: { + readonly onRemoteSourceChunksLoaded: (sourceId: string) => void; + }) => + useWorkspaceSelectedChunks({ + selectedSourceId: "knowhere-doc:default:doc_remote", + sources: [remoteSource], + prefetchedChunksBySourceId: {}, + onRemoteSourceChunksLoaded: input.onRemoteSourceChunksLoaded, + }), + { + initialProps: { onRemoteSourceChunksLoaded }, + wrapper: createSWRWrapper, + }, + ); + + await waitFor(() => + expect(onRemoteSourceChunksLoaded).toHaveBeenCalledWith( + "knowhere-doc:default:doc_remote", + ), + ); + + rerender({ onRemoteSourceChunksLoaded }); + expect(onRemoteSourceChunksLoaded).toHaveBeenCalledTimes(1); + }); + it("returns an empty chunk list when no source is selected", () => { const { result } = renderHook( () => diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index 6d79175..f00ac68 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -1,6 +1,6 @@ "use client" -import { useMemo } from "react" +import { useEffect, useMemo, useRef } from "react" import useSWRInfinite from "swr/infinite" import { workspaceClient } from "@/domains/workspace/client" @@ -17,6 +17,7 @@ type WorkspaceSelectedChunksInput = { readonly selectedSourceId: string | null readonly sources: readonly SourceView[] readonly prefetchedChunksBySourceId: Readonly> + readonly onRemoteSourceChunksLoaded?: (sourceId: string) => void } type WorkspaceSelectedChunks = { @@ -33,8 +34,10 @@ export function useWorkspaceSelectedChunks({ selectedSourceId, sources, prefetchedChunksBySourceId, + onRemoteSourceChunksLoaded, }: WorkspaceSelectedChunksInput): WorkspaceSelectedChunks { const selectedSource = sources.find((source) => source.id === selectedSourceId) + const remoteSourceRefreshRequestedIdsRef = useRef>(new Set()) const prefetchedSelectedChunks = selectedSourceId ? prefetchedChunksBySourceId[selectedSourceId] : undefined @@ -100,6 +103,16 @@ export function useWorkspaceSelectedChunks({ !selectedChunkPages && isChunksLoading) + useEffect(() => { + const sourceId = selectedSource?.id + if (!sourceId || selectedSource.kind !== "remote") return + if (!selectedChunkPages || selectedChunkPages.length === 0) return + if (remoteSourceRefreshRequestedIdsRef.current.has(sourceId)) return + + remoteSourceRefreshRequestedIdsRef.current.add(sourceId) + onRemoteSourceChunksLoaded?.(sourceId) + }, [onRemoteSourceChunksLoaded, selectedChunkPages, selectedSource]) + function handleLoadMoreChunks(): void { if (!hasMoreSelectedChunks || isSelectedChunksLoadingMore) return void setSelectedChunkPageCount(selectedChunkPageCount + 1) diff --git a/src/components/workspace-shell.tsx b/src/components/workspace-shell.tsx index 4dd8b9c..64663a0 100644 --- a/src/components/workspace-shell.tsx +++ b/src/components/workspace-shell.tsx @@ -114,6 +114,7 @@ function WorkspaceShellContent({ fetchChunks: workspaceClient.fetchChunks, initialPrefetchedChunksBySourceId: initialPrefetchedChunksBySourceId ?? undefined, + onRemoteSourceChunksLoaded: sourceWorkflow.handleSourcesRefresh, onSelectSource: handleCitationSourceSelected, selectedSourceId: sourceWorkflow.selectedSourceId, sources: sourceWorkflow.sources, diff --git a/src/components/workspace-source-state.test.ts b/src/components/workspace-source-state.test.ts index f1069d5..19e6cb1 100644 --- a/src/components/workspace-source-state.test.ts +++ b/src/components/workspace-source-state.test.ts @@ -83,6 +83,27 @@ describe("workspaceSourceState", () => { ).toBe("source_target"); }); + it("keeps a localized remote document selected after source refresh", () => { + const sources: readonly SourceView[] = [ + { + id: "source_localized", + kind: "workspace", + title: "remote.pdf", + status: "ready", + mimeType: "application/pdf", + documentId: "doc_remote", + excludedFromQuery: false, + }, + ]; + + expect( + workspaceSourceState.getResolvedSelectedSourceId( + sources, + "knowhere-doc:default:doc_remote", + ), + ).toBe("source_localized"); + }); + it("applies source query exclusions without mutating the source list", () => { const sources: readonly SourceView[] = [ { diff --git a/src/components/workspace-source-state.ts b/src/components/workspace-source-state.ts index 676e020..0373983 100644 --- a/src/components/workspace-source-state.ts +++ b/src/components/workspace-source-state.ts @@ -74,6 +74,15 @@ function getResolvedSelectedSourceId( return selectedSource.id } + const selectedDocumentId = getRemoteSourceDocumentId(selectedSourceId) + if (selectedDocumentId) { + const localizedSource = sources.find( + (source) => + source.documentId === selectedDocumentId && isReadySource(source), + ) + if (localizedSource) return localizedSource.id + } + return getFirstReadySourceId(sources) } @@ -134,6 +143,20 @@ function removeRecordKey( return remaining } +function getRemoteSourceDocumentId(sourceId: string | null): string | null { + if (!sourceId) return null + + const parts = sourceId.split(":") + if (parts.length !== 3 || parts[0] !== "knowhere-doc") return null + + try { + const documentId = decodeURIComponent(parts[2] ?? "") + return documentId.length > 0 ? documentId : null + } catch { + return null + } +} + export const workspaceSourceState: WorkspaceSourceStateModule = { getFirstReadySourceId, getInitialSelectedSourceId, diff --git a/src/components/workspace-source-workflow.ts b/src/components/workspace-source-workflow.ts index 9b9f1f8..7ad2871 100644 --- a/src/components/workspace-source-workflow.ts +++ b/src/components/workspace-source-workflow.ts @@ -22,6 +22,7 @@ type WorkspaceSourceWorkflow = { readonly handleRetrySource: (sourceId: string) => Promise readonly handleOfficialLibrarySourceAdd: (demoSourceId: string) => Promise readonly handleSelectedSourceChange: (sourceId: string | null) => void + readonly handleSourcesRefresh: () => void readonly handleSourcesMaterialized: ( demoSourceIds: readonly string[], materializedSources: readonly SourceView[], @@ -153,6 +154,10 @@ export function useWorkspaceSourceWorkflow({ setSelectedSourceId(sourceId) } + function handleSourcesRefresh(): void { + void mutateSources() + } + async function handleArchiveSource(sourceId: string): Promise { setArchivingSourceIds((current) => workspaceSourceState.addPendingId(current, sourceId), @@ -237,6 +242,7 @@ export function useWorkspaceSourceWorkflow({ handleRetrySource, handleOfficialLibrarySourceAdd, handleSelectedSourceChange, + handleSourcesRefresh, handleSourcesMaterialized, handleSourceUploaded, handleToggleIncluded, diff --git a/src/domains/chat/contracts.ts b/src/domains/chat/contracts.ts index 4c8d224..4752b86 100644 --- a/src/domains/chat/contracts.ts +++ b/src/domains/chat/contracts.ts @@ -10,7 +10,7 @@ import type { ChatCitationView, } from "@/domains/chat/types" import type { HardenMediaAssetUrls } from "./media-asset-hardening" -import type { LoadSourceAssetUrls } from "./media-assets" +import type { HardenChatAssetUrl } from "./media-assets" export type RetrievalClient = { query(params: RetrievalQueryParams): Promise @@ -67,7 +67,7 @@ export type AnswerQuestionInput = { excludedSourceIds: readonly string[] retrieval: RetrievalClient generateAnswer: GenerateAnswer - loadSourceAssetUrls?: LoadSourceAssetUrls + hardenChatAssetUrl?: HardenChatAssetUrl hardenMediaAssetUrls?: HardenMediaAssetUrls messages: readonly ChatHistoryMessage[] } diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 59899b0..0be6b78 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -452,10 +452,11 @@ describe("answerQuestionWithRetrieval", () => { }); return makeHarnessRunResult(`Use this launch photo. ${upstreamAssetUrl}`); }); - const loadSourceAssetUrls = vi.fn().mockResolvedValue({ - "images/image-9-Night Rocket Launch.jpg": + const hardenChatAssetUrl = vi + .fn() + .mockResolvedValue( "https://blob.example/images/image-9-Night%20Rocket%20Launch.jpg", - }); + ); const answer = await Effect.runPromise( answerQuestionWithRetrieval({ @@ -471,14 +472,16 @@ describe("answerQuestionWithRetrieval", () => { excludedSourceIds: [], retrieval, generateAnswer, - loadSourceAssetUrls, + hardenChatAssetUrl, messages: [], }), ); - expect(loadSourceAssetUrls).toHaveBeenCalledWith( - expect.objectContaining({ id: "source_spacex" }), - ); + expect(hardenChatAssetUrl).toHaveBeenCalledWith({ + source: expect.objectContaining({ id: "source_spacex" }), + sourcePath: "images/image-9-Night Rocket Launch.jpg", + assetUrl: upstreamAssetUrl, + }); expect(retrieval.query).toHaveBeenCalledWith({ namespace: "notebook-workspace", query: "SpaceX rocket photos", @@ -731,6 +734,7 @@ describe("answerQuestionWithRetrieval", () => { ...(artifacts ? { artifacts: [...artifacts] } : {}), }), ); + const hardenChatAssetUrl = vi.fn().mockResolvedValue(storedPageAssetUrl); const answer = await Effect.runPromise( answerQuestionWithRetrieval({ @@ -747,13 +751,17 @@ describe("answerQuestionWithRetrieval", () => { retrieval, generateAnswer, hardenMediaAssetUrls, - loadSourceAssetUrls: vi.fn(async () => ({ - "page_citation_assets/page-4.png": storedPageAssetUrl, - })), + hardenChatAssetUrl, messages: [], }), ); + expect(hardenChatAssetUrl).toHaveBeenCalledWith({ + source: expect.objectContaining({ id: "source_pages" }), + sourcePath: "page_citation_assets/page-4.png", + assetUrl: rawPageAssetUrl, + contentType: "image/png", + }); expect(hardenMediaAssetUrls).toHaveBeenCalledWith({ results: [ expect.objectContaining({ @@ -844,6 +852,7 @@ describe("answerQuestionWithRetrieval", () => { ...(artifacts ? { artifacts: [...artifacts] } : {}), }), ); + const hardenChatAssetUrl = vi.fn().mockResolvedValue(storedPageAssetUrl); const answer = await Effect.runPromise( answerQuestionWithRetrieval({ @@ -860,13 +869,17 @@ describe("answerQuestionWithRetrieval", () => { retrieval, generateAnswer, hardenMediaAssetUrls, - loadSourceAssetUrls: vi.fn(async () => ({ - "page_citation_assets/page-6.png": storedPageAssetUrl, - })), + hardenChatAssetUrl, messages: [], }), ); + expect(hardenChatAssetUrl).toHaveBeenCalledWith({ + source: expect.objectContaining({ id: "source_pages" }), + sourcePath: "page_citation_assets/page-6.png", + assetUrl: rawPageAssetUrl, + contentType: "image/png", + }); expect(hardenMediaAssetUrls).toHaveBeenCalledWith({ results: [ expect.objectContaining({ @@ -1406,7 +1419,7 @@ describe("answerQuestionWithRetrieval", () => { ); }); - it("turns retrieved evidence image filenames into image citations", async () => { + it("does not turn evidence-only image filenames into image citations", async () => { const result = makeRetrievalResult({ content: "This section contains identity proof attachments.", source: { @@ -1434,12 +1447,7 @@ describe("answerQuestionWithRetrieval", () => { }); return makeHarnessRunResult("这里是相关身份证明图片。"); }); - const loadSourceAssetUrls = vi.fn().mockResolvedValue({ - "images/image-6-中华人民共和国居民身份证.jpg": - "https://blob.example/images/image-6-id-front.jpg", - "images/image-7-中国居民身份证.jpg": - "https://blob.example/images/image-7-id-back.jpg", - }); + const hardenChatAssetUrl = vi.fn().mockResolvedValue(null); const sources = [ makeSource({ id: "source_identity", @@ -1456,7 +1464,7 @@ describe("answerQuestionWithRetrieval", () => { excludedSourceIds: [], retrieval, generateAnswer, - loadSourceAssetUrls, + hardenChatAssetUrl, messages: [], }), ); @@ -1478,14 +1486,8 @@ describe("answerQuestionWithRetrieval", () => { const imageCitations = answer.citations.filter( (citation) => citation.assetUrl, ) - expect(imageCitations.map((citation) => citation.assetUrl)).toEqual([ - "https://blob.example/images/image-6-id-front.jpg", - "https://blob.example/images/image-7-id-back.jpg", - ]); - expect(imageCitations.map((citation) => citation.chunkType)).toEqual([ - "image", - "image", - ]); + expect(imageCitations).toEqual([]); + expect(hardenChatAssetUrl).not.toHaveBeenCalled(); }); it("returns the agent answer without citations when retrieval has no results", async () => { diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index bbf95b9..dc10453 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -242,7 +242,7 @@ export const answerQuestionWithRetrieval = ( enrichRetrievalResultsWithAssetUrls({ results: useNotebookSourceTitles(rawResults, input.sources), sources: input.sources, - loadSourceAssetUrls: input.loadSourceAssetUrls, + hardenChatAssetUrl: input.hardenChatAssetUrl, evidenceText: formatRetrievalEvidenceText(retrievalResponses), }), ) @@ -250,7 +250,7 @@ export const answerQuestionWithRetrieval = ( enrichRetrievalResultsWithPageCitationAssetUrls({ results: enrichedResults, sources: input.sources, - loadSourceAssetUrls: input.loadSourceAssetUrls, + hardenChatAssetUrl: input.hardenChatAssetUrl, }), ) const artifacts = toChatArtifactViewsFromHarness(generatedAnswer, input.sources) diff --git a/src/domains/chat/media-asset-hardening.test.ts b/src/domains/chat/media-asset-hardening.test.ts index d3c87b2..7c106c9 100644 --- a/src/domains/chat/media-asset-hardening.test.ts +++ b/src/domains/chat/media-asset-hardening.test.ts @@ -22,10 +22,10 @@ afterEach(() => { }) describe("hardenChatMediaAssetUrls", () => { - it("keeps an already Notebook-owned asset URL without loading the asset map", async () => { + it("keeps an already Notebook-owned asset URL without calling the hardener", async () => { const ownedUrl = "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_1/rev_1/assets/images/a.png" - const loadSourceAssetUrls = vi.fn(async () => ({})) + const hardenChatAssetUrl = vi.fn(async () => null) const result = await hardenChatMediaAssetUrls({ workspaceId: "workspace_1", @@ -41,11 +41,11 @@ describe("hardenChatMediaAssetUrls", () => { }, }), ], - loadSourceAssetUrls, + hardenChatAssetUrl, }) expect(result.results[0]?.assetUrl).toBe(ownedUrl) - expect(loadSourceAssetUrls).not.toHaveBeenCalled() + expect(hardenChatAssetUrl).not.toHaveBeenCalled() }) it("resolves a raw asset URL to the durable parsed asset URL", async () => { @@ -53,9 +53,7 @@ describe("hardenChatMediaAssetUrls", () => { "https://knowhere-storage.example/results/job_1/images/id-front.jpg?AWSAccessKeyId=test" const durableUrl = "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/assets/images/id-front.jpg" - const loadSourceAssetUrls = vi.fn().mockResolvedValue({ - "images/id-front.jpg": durableUrl, - }) + const hardenChatAssetUrl = vi.fn().mockResolvedValue(durableUrl) const result = await hardenChatMediaAssetUrls({ workspaceId: "workspace_1", @@ -76,19 +74,21 @@ describe("hardenChatMediaAssetUrls", () => { }, }), ], - loadSourceAssetUrls, + hardenChatAssetUrl, }) - expect(loadSourceAssetUrls).toHaveBeenCalledWith( - expect.objectContaining({ id: "source_identity" }), - ) + expect(hardenChatAssetUrl).toHaveBeenCalledWith({ + source: expect.objectContaining({ id: "source_identity" }), + sourcePath: "images/id-front.jpg", + assetUrl: rawAssetUrl, + }) expect(result.results[0]?.assetUrl).toBe(durableUrl) }) it("omits an asset URL that cannot be resolved to a durable URL", async () => { const rawAssetUrl = "https://knowhere-storage.example/results/job_1/tables/table-1.html?AWSAccessKeyId=test" - const loadSourceAssetUrls = vi.fn().mockResolvedValue({}) + const hardenChatAssetUrl = vi.fn().mockResolvedValue(null) const result = await hardenChatMediaAssetUrls({ workspaceId: "workspace_1", @@ -106,7 +106,7 @@ describe("hardenChatMediaAssetUrls", () => { }, }), ], - loadSourceAssetUrls, + hardenChatAssetUrl, }) expect(result.results[0]?.assetUrl).toBeUndefined() @@ -117,9 +117,7 @@ describe("hardenChatMediaAssetUrls", () => { "https://knowhere-storage.example/results/job_1/images/front.jpg?AWSAccessKeyId=test" const durableUrl = "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/assets/images/front.jpg" - const loadSourceAssetUrls = vi.fn().mockResolvedValue({ - "images/front.jpg": durableUrl, - }) + const hardenChatAssetUrl = vi.fn().mockResolvedValue(durableUrl) const result = await hardenChatMediaAssetUrls({ workspaceId: "workspace_1", @@ -148,7 +146,7 @@ describe("hardenChatMediaAssetUrls", () => { }, }, ], - loadSourceAssetUrls, + hardenChatAssetUrl, }) const [artifact] = result.artifacts ?? [] diff --git a/src/domains/chat/media-asset-hardening.ts b/src/domains/chat/media-asset-hardening.ts index f768515..ba88bb1 100644 --- a/src/domains/chat/media-asset-hardening.ts +++ b/src/domains/chat/media-asset-hardening.ts @@ -6,8 +6,7 @@ import type { } from "@/domains/chat/types" import type { Source } from "@/infrastructure/db/schema" import { logger } from "@/lib/logger" -import type { LoadSourceAssetUrls } from "./media-assets" -import { resolveAssetUrlFromReferenceText } from "./media-assets" +import type { HardenChatAssetUrl } from "./media-assets" export type HardenableRetrievalResult = RetrievalResult & { readonly pageCitationAssetUrl?: string @@ -31,7 +30,7 @@ export type HardenChatMediaAssetUrlsForWorkspaceInput = HardenMediaAssetUrlsInput & { readonly workspaceId: string readonly sources: readonly Source[] - readonly loadSourceAssetUrls: LoadSourceAssetUrls + readonly hardenChatAssetUrl: HardenChatAssetUrl } type AssetReferenceSource = ChatCitationView["source"] @@ -44,11 +43,7 @@ type AssetUrlReference = { type HardeningContext = { readonly sourcesByDocumentId: ReadonlyMap - readonly loadSourceAssetUrls: LoadSourceAssetUrls - readonly assetUrlsBySourceId: Map< - string, - Promise>> - > + readonly hardenChatAssetUrl: HardenChatAssetUrl } const parsedResultDirectoryName = "parsed-result" @@ -56,9 +51,9 @@ const chatAssetsDirectoryName = "chat-assets" /** * Resolve chat citation/media asset URLs to durable Notebook Blob URLs. The - * single hardening path is the SDK's `assetUrlPolicy: "durable"` read that - * `loadSourceAssetUrls` performs; here we only map a retrieval result's - * reference text to the durable URL that read produced. + * hardening path is one asset at a time: derive the source path already present + * on the returned chat/citation asset, then let the route-level hardener check + * Blob and write only that missing asset. * * An asset that is already Notebook-owned is kept as-is. An asset that cannot * be resolved to a durable URL is omitted rather than exposing a presigned @@ -68,12 +63,11 @@ export async function hardenChatMediaAssetUrls({ results, artifacts, sources, - loadSourceAssetUrls, + hardenChatAssetUrl, }: HardenChatMediaAssetUrlsForWorkspaceInput): Promise { const context: HardeningContext = { sourcesByDocumentId: createSourcesByDocumentId(sources), - loadSourceAssetUrls, - assetUrlsBySourceId: new Map(), + hardenChatAssetUrl, } const hardenedResults = await Promise.all( @@ -195,7 +189,7 @@ async function hardenCitation( /** * Return a durable Notebook-owned URL for a reference: keep already-owned URLs, - * otherwise resolve against the source's durable parsed asset map. Returns + * otherwise harden only the referenced asset path. Returns * `undefined` when no durable URL is available so callers omit the URL rather * than leak a presigned Knowhere URL. */ @@ -210,15 +204,24 @@ async function resolveDurableAssetUrl( const source = resolveSourceForReference(reference, context) if (!source) return undefined - const assetUrlsByFilePath = await getCachedSourceAssetUrls(source, context) - const durableUrl = resolveAssetUrlFromReferenceText({ - values: [ - reference.source?.sectionPath, - reference.content, - getAssetUrlPathname(reference.assetUrl), - ], - assetUrlsByFilePath, - }) + const sourcePath = resolveSourcePathForReference(reference) + if (!sourcePath) return undefined + + const durableUrl = await context + .hardenChatAssetUrl({ + source, + sourcePath, + assetUrl: reference.assetUrl, + }) + .catch((error: unknown) => { + logger.warn("chat: failed to harden referenced asset", { + sourceId: source.id, + sourcePath, + error: formatUnknownError(error), + }) + return null + }) + return durableUrl ?? undefined } @@ -261,26 +264,6 @@ function applyAssetUrls< return next as T } -async function getCachedSourceAssetUrls( - source: Source, - context: HardeningContext, -): Promise>> { - const cached = context.assetUrlsBySourceId.get(source.id) - if (cached) return cached - - const loaded = context - .loadSourceAssetUrls(source) - .catch((error: unknown) => { - logger.warn("chat: failed to load durable parsed asset map", { - sourceId: source.id, - error: formatUnknownError(error), - }) - return {} - }) - context.assetUrlsBySourceId.set(source.id, loaded) - return loaded -} - export function isNotebookOwnedAssetUrl(assetUrl: string): boolean { const pathname = getAssetUrlPathname(assetUrl).toLowerCase() if ( @@ -313,6 +296,58 @@ function parseAbsoluteHttpUrl(assetUrl: string): URL | null { } } +function resolveSourcePathForReference( + reference: AssetUrlReference, +): string | null { + const candidates = [ + reference.source?.sectionPath, + reference.content, + getAssetUrlPathname(reference.assetUrl), + ] + + for (const candidate of candidates) { + const sourcePath = getSupportedAssetPath(candidate) + if (sourcePath) return sourcePath + } + + return null +} + +function getSupportedAssetPath(value: string | null | undefined): string | null { + const normalizedText = normalizeSourcePathCandidate(value) + if (!normalizedText) return null + + const match = + /(?:^|\/)((?:images|tables|pages|page_citation_assets)\/[^?#]+)/i.exec( + normalizedText, + ) + const matchedPath = match?.[1] + return matchedPath ? matchedPath.trim() : null +} + +function normalizeSourcePathCandidate( + value: string | null | undefined, +): string | null { + const trimmedValue = getTrimmedString(value) + if (!trimmedValue) return null + + const normalized = decodeUrlText(trimmedValue) + .replaceAll("\\", "/") + .replace(/\s*\/\s*/g, "/") + .replace(/\s+/g, " ") + .trim() + + return normalized.length > 0 ? normalized : null +} + +function decodeUrlText(value: string): string { + try { + return decodeURIComponent(value) + } catch { + return value + } +} + function resolveSourceForReference( reference: AssetUrlReference, context: HardeningContext, diff --git a/src/domains/chat/media-assets.test.ts b/src/domains/chat/media-assets.test.ts index 2fc84cb..470255d 100644 --- a/src/domains/chat/media-assets.test.ts +++ b/src/domains/chat/media-assets.test.ts @@ -11,10 +11,11 @@ import type { Source } from "@/infrastructure/db/schema" describe("chat media assets", () => { it("enriches retrieved image chunks from Notebook parsed asset URLs", async () => { - const loadSourceAssetUrls = vi.fn().mockResolvedValue({ - "images/image-9-Night Rocket Launch.jpg": + const hardenChatAssetUrl = vi + .fn() + .mockResolvedValue( "https://blob.example/images/image-9-Night%20Rocket%20Launch.jpg", - }) + ) const [result] = await enrichRetrievalResultsWithAssetUrls({ results: [ @@ -33,20 +34,25 @@ describe("chat media assets", () => { knowhereDocumentId: "doc_spacex", }), ], - loadSourceAssetUrls, + hardenChatAssetUrl, }) - expect(loadSourceAssetUrls).toHaveBeenCalledTimes(1) + expect(hardenChatAssetUrl).toHaveBeenCalledWith({ + source: expect.objectContaining({ id: "source_spacex" }), + sourcePath: "images/image-9-Night Rocket Launch.jpg", + assetUrl: null, + }) expect(result?.assetUrl).toBe( "https://blob.example/images/image-9-Night%20Rocket%20Launch.jpg", ) }) it("prefers Notebook parsed asset URLs over existing upstream asset URLs", async () => { - const loadSourceAssetUrls = vi.fn().mockResolvedValue({ - "images/image-6-情感分类模型.jpg": + const hardenChatAssetUrl = vi + .fn() + .mockResolvedValue( "https://blob.example/workspaces/workspace_1/sources/source_doc/parsed-result/images/image-6-model.jpg", - }) + ) const [result] = await enrichRetrievalResultsWithAssetUrls({ results: [ @@ -67,7 +73,7 @@ describe("chat media assets", () => { knowhereDocumentId: "doc_model", }), ], - loadSourceAssetUrls, + hardenChatAssetUrl, }) expect(result?.assetUrl).toBe( @@ -75,13 +81,8 @@ describe("chat media assets", () => { ) }) - it("adds image citation results for asset filenames that only appear in evidence text", async () => { - const loadSourceAssetUrls = vi.fn().mockResolvedValue({ - "images/image-6-中华人民共和国居民身份证.jpg": - "https://blob.example/images/image-6-id-front.jpg", - "images/image-7-中国居民身份证.jpg": - "https://blob.example/images/image-7-id-back.jpg", - }) + it("does not scan a source for asset filenames that only appear in evidence text", async () => { + const hardenChatAssetUrl = vi.fn().mockResolvedValue(null) const results = await enrichRetrievalResultsWithAssetUrls({ results: [ @@ -101,25 +102,14 @@ describe("chat media assets", () => { knowhereDocumentId: "doc_identity", }), ], - loadSourceAssetUrls, + hardenChatAssetUrl, evidenceText: "[image-6-中华人民共和国居民身份证.jpg]\n[image-7-中国居民身份证.jpg]", }) - expect(results).toHaveLength(3) + expect(results).toHaveLength(1) expect(results[0]?.assetUrl).toBeUndefined() - expect(results.slice(1).map((result) => result.assetUrl)).toEqual([ - "https://blob.example/images/image-6-id-front.jpg", - "https://blob.example/images/image-7-id-back.jpg", - ]) - expect(results.slice(1).map((result) => result.chunkType)).toEqual([ - "image", - "image", - ]) - expect(results.slice(1).map((result) => result.source.sectionPath)).toEqual([ - "images/image-6-中华人民共和国居民身份证.jpg", - "images/image-7-中国居民身份证.jpg", - ]) + expect(hardenChatAssetUrl).not.toHaveBeenCalled() }) it("deduplicates media citation assets globally by asset URL", async () => { diff --git a/src/domains/chat/media-assets.ts b/src/domains/chat/media-assets.ts index 85984cc..bafaae8 100644 --- a/src/domains/chat/media-assets.ts +++ b/src/domains/chat/media-assets.ts @@ -12,24 +12,30 @@ const internalMetadataKeys = new Set([ "chunk_id", ]) -export type LoadSourceAssetUrls = ( - source: Source, -) => Promise>> +export type HardenChatAssetUrlInput = { + readonly source: Source + readonly sourcePath: string + readonly assetUrl?: string | null + readonly contentType?: string | null +} + +export type HardenChatAssetUrl = ( + input: HardenChatAssetUrlInput, +) => Promise export type RetrievalResultAssetInput = { readonly results: readonly RetrievalResult[] readonly sources: readonly Source[] - readonly loadSourceAssetUrls?: LoadSourceAssetUrls + readonly hardenChatAssetUrl?: HardenChatAssetUrl readonly evidenceText?: string } export async function enrichRetrievalResultsWithAssetUrls({ results, sources, - loadSourceAssetUrls, - evidenceText, + hardenChatAssetUrl, }: RetrievalResultAssetInput): Promise { - if (!loadSourceAssetUrls || results.length === 0) { + if (!hardenChatAssetUrl || results.length === 0) { return dedupeMediaCitationResults(results) } @@ -38,23 +44,13 @@ export async function enrichRetrievalResultsWithAssetUrls({ source.knowhereDocumentId ? [[source.knowhereDocumentId, source]] : [], ), ) - const assetUrlsBySourceId = new Map< - string, - Promise>> - >() - const enrichedResults = await Promise.all( results.map(async (result): Promise => { const documentId = getTrimmedString(result.source.documentId) const source = documentId ? sourcesByDocumentId.get(documentId) : undefined if (!source) return [result] - const assetUrls = await getCachedSourceAssetUrls( - source, - loadSourceAssetUrls, - assetUrlsBySourceId, - ) - return addAssetCitationResults(result, assetUrls, evidenceText) + return addAssetCitationResults(result, source, hardenChatAssetUrl) }), ) @@ -234,56 +230,33 @@ function isNotebookParsedAssetUrl(assetUrl: string): boolean { ) === true } -async function getCachedSourceAssetUrls( +async function addAssetCitationResults( + result: RetrievalResult, source: Source, - loadSourceAssetUrls: LoadSourceAssetUrls, - cache: Map>>>, -): Promise>> { - let cached = cache.get(source.id) - if (!cached) { - cached = loadSourceAssetUrls(source).catch(() => ({})) - cache.set(source.id, cached) - } - return cached -} + hardenChatAssetUrl: HardenChatAssetUrl, +): Promise { + if (result.chunkType.toLowerCase() === "page") return [result] -function addAssetCitationResults( - result: RetrievalResult, - assetUrlsByFilePath: Readonly>, - evidenceText: string | undefined, -): readonly RetrievalResult[] { const existingAssetUrl = getTrimmedString(result.assetUrl) - const resultMatches = resolveAssetReferenceMatches(result, assetUrlsByFilePath) - const evidenceMatches = resolveAssetReferenceMatchesFromText( - evidenceText, - assetUrlsByFilePath, - ) - const seenAssetUrls = new Set() - const output: RetrievalResult[] = [] - - if (resultMatches.length > 0) { - const [firstMatch, ...remainingMatches] = resultMatches - seenAssetUrls.add(firstMatch.assetUrl) - output.push(toAssetResult(result, firstMatch)) - for (const match of remainingMatches) { - if (seenAssetUrls.has(match.assetUrl)) continue - seenAssetUrls.add(match.assetUrl) - output.push(toAssetResult(result, match)) - } - } else if (existingAssetUrl) { - seenAssetUrls.add(existingAssetUrl) - output.push(result) - } else { - output.push(result) - } - - for (const match of evidenceMatches) { - if (seenAssetUrls.has(match.assetUrl)) continue - seenAssetUrls.add(match.assetUrl) - output.push(toAssetResult(result, match)) - } - - return output + if (existingAssetUrl && isNotebookOwnedAssetUrl(existingAssetUrl)) return [result] + + const sourcePath = getAssetSourcePathFromResult(result, existingAssetUrl) + if (!sourcePath) return [result] + + const assetUrl = await hardenChatAssetUrl({ + source, + sourcePath, + assetUrl: existingAssetUrl, + }).catch(() => null) + if (!assetUrl) return [result] + + return [ + toAssetResult(result, { + assetPath: sourcePath, + assetUrl, + index: 0, + }), + ] } function toAssetResult( @@ -346,128 +319,12 @@ function isAssetFilePath(value: string): boolean { return normalizedPath ? /^(images|tables)\//.test(normalizedPath) : false } -function resolveAssetReferenceMatches( - result: RetrievalResult, - assetUrlsByFilePath: Readonly>, -): readonly AssetReferenceMatch[] { - const normalizedHaystacks = [ - result.source.sectionPath, - result.content, - ].flatMap((value): string[] => { - const normalized = normalizeAssetLookupText(value) - return normalized ? [normalized] : [] - }) - if (normalizedHaystacks.length === 0) return [] - - return resolveAssetReferenceMatchesFromHaystacks( - normalizedHaystacks, - assetUrlsByFilePath, - ) -} - -function resolveAssetReferenceMatchesFromText( - value: string | null | undefined, - assetUrlsByFilePath: Readonly>, -): readonly AssetReferenceMatch[] { - const normalized = normalizeAssetLookupText(value) - if (!normalized) return [] - - return resolveAssetReferenceMatchesFromHaystacks( - [normalized], - assetUrlsByFilePath, - ) -} - -export function resolveAssetUrlFromReferenceText(input: { - readonly values: readonly (string | null | undefined)[] - readonly assetUrlsByFilePath: Readonly> -}): string | null { - const normalizedHaystacks = input.values.flatMap((value): string[] => { - const normalized = normalizeAssetLookupText(value) - return normalized ? [normalized] : [] - }) - if (normalizedHaystacks.length === 0) return null - - const [match] = resolveAssetReferenceMatchesFromHaystacks( - normalizedHaystacks, - input.assetUrlsByFilePath, - ) - return match?.assetUrl ?? null -} - -function resolveAssetReferenceMatchesFromHaystacks( - normalizedHaystacks: readonly string[], - assetUrlsByFilePath: Readonly>, -): readonly AssetReferenceMatch[] { - const basenameCounts = getNormalizedBasenameCounts(assetUrlsByFilePath) - return Object.entries(assetUrlsByFilePath) - .flatMap(([assetPath, assetUrl]): readonly AssetReferenceMatch[] => { - const trimmedUrl = getTrimmedString(assetUrl) - if (!trimmedUrl || !isSupportedAssetPath(assetPath)) return [] - - const index = getAssetReferenceIndex( - normalizedHaystacks, - assetPath, - basenameCounts, - ) - return index === null ? [] : [{ assetPath, assetUrl: trimmedUrl, index }] - }) - .sort(compareAssetReferenceMatches) -} - type AssetReferenceMatch = { readonly assetPath: string readonly assetUrl: string readonly index: number } -function compareAssetReferenceMatches( - left: AssetReferenceMatch, - right: AssetReferenceMatch, -): number { - return left.index - right.index || left.assetPath.localeCompare(right.assetPath) -} - -function getAssetReferenceIndex( - normalizedHaystacks: readonly string[], - assetPath: string, - basenameCounts: ReadonlyMap, -): number | null { - const normalizedPath = normalizeAssetLookupText(assetPath) - if (!normalizedPath) return null - - const directIndex = getFirstIndex(normalizedHaystacks, normalizedPath) - if (directIndex !== null) return directIndex - - const basename = getNormalizedBasename(assetPath) - if (!basename || basenameCounts.get(basename) !== 1) return null - - return getFirstIndex(normalizedHaystacks, basename) -} - -function getFirstIndex( - normalizedHaystacks: readonly string[], - needle: string, -): number | null { - const indexes = normalizedHaystacks - .map((haystack): number => haystack.indexOf(needle)) - .filter((index): index is number => index >= 0) - - return indexes.length > 0 ? Math.min(...indexes) : null -} - -function getNormalizedBasenameCounts( - assetUrlsByFilePath: Readonly>, -): ReadonlyMap { - const counts = new Map() - for (const assetPath of Object.keys(assetUrlsByFilePath)) { - const basename = getNormalizedBasename(assetPath) - if (!basename) continue - counts.set(basename, (counts.get(basename) ?? 0) + 1) - } - return counts -} - function getNormalizedBasename(assetPath: string): string | null { const basename = assetPath.replaceAll("\\", "/").split("/").pop() return normalizeAssetLookupText(basename) @@ -495,14 +352,84 @@ function decodeUrlText(value: string): string { } } +function getAssetSourcePathFromResult( + result: RetrievalResult, + assetUrl: string | null, +): string | null { + const candidates = [ + result.filePath, + result.sourceChunkPath, + result.source.sectionPath, + assetUrl ? getUrlPathname(assetUrl) : null, + result.content, + ] + + for (const candidate of candidates) { + const sourcePath = getSupportedAssetPath(candidate) + if (sourcePath) return sourcePath + } + + return null +} + +function getSupportedAssetPath(value: string | null | undefined): string | null { + const normalizedText = normalizeSourcePathCandidate(value) + if (!normalizedText) return null + + const match = + /(?:^|\/)((?:images|tables|pages|page_citation_assets)\/[^?#]+)/i.exec( + normalizedText, + ) + const matchedPath = match?.[1] + return matchedPath ? matchedPath.trim() : null +} + +function normalizeSourcePathCandidate( + value: string | null | undefined, +): string | null { + const trimmedValue = getTrimmedString(value) + if (!trimmedValue) return null + + const normalized = decodeUrlText(trimmedValue) + .replaceAll("\\", "/") + .replace(/\s*\/\s*/g, "/") + .replace(/\s+/g, " ") + .trim() + + return normalized.length > 0 ? normalized : null +} + function isSupportedAssetPath(assetPath: string): boolean { - const normalizedPath = normalizeAssetLookupText(assetPath) + return getSupportedAssetPath(assetPath) !== null +} + +function isNotebookOwnedAssetUrl(assetUrl: string): boolean { + const pathname = getUrlPathname(assetUrl).toLowerCase() + if ( + pathname.includes("/parsed-result/") || + pathname.includes("/chat-assets/") || + pathname.includes("/parsed-documents/") + ) { + return true + } + + const absoluteUrl = parseAbsoluteHttpUrl(assetUrl) return ( - normalizedPath?.startsWith("images/") === true || - normalizedPath?.startsWith("tables/") === true + absoluteUrl?.hostname + .toLowerCase() + .endsWith(".blob.vercel-storage.com") === true ) } +function parseAbsoluteHttpUrl(assetUrl: string): URL | null { + try { + const url = new URL(assetUrl) + return url.protocol === "http:" || url.protocol === "https:" ? url : null + } catch { + return null + } +} + function isRenderableMediaAsset( result: RetrievalResult, assetUrl: string, diff --git a/src/domains/chat/page-citation-assets.test.ts b/src/domains/chat/page-citation-assets.test.ts index d99dbcb..4553d5d 100644 --- a/src/domains/chat/page-citation-assets.test.ts +++ b/src/domains/chat/page-citation-assets.test.ts @@ -29,10 +29,9 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { }) it("uses stored Blob URLs for page citation assets", async () => { - const loadSourceAssetUrls = vi.fn().mockResolvedValue({ - "page_citation_assets/page-2.png": - "https://blob.example/page_citation_assets/page-2.png", - }) + const hardenChatAssetUrl = vi + .fn() + .mockResolvedValue("https://blob.example/page_citation_assets/page-2.png") const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ results: [ @@ -51,22 +50,24 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { }), ], sources: [makeSource()], - loadSourceAssetUrls, + hardenChatAssetUrl, }) - expect(loadSourceAssetUrls).toHaveBeenCalledWith( - expect.objectContaining({ id: "source_1" }), - ) + expect(hardenChatAssetUrl).toHaveBeenCalledWith({ + source: expect.objectContaining({ id: "source_1" }), + sourcePath: "page_citation_assets/page-2.png", + assetUrl: "https://assets.example/pages/page-2.png", + contentType: undefined, + }) expect(result?.pageCitationAssetUrl).toBe( "https://blob.example/page_citation_assets/page-2.png", ) }) it("chooses the stored asset matching the citation page metadata", async () => { - const loadSourceAssetUrls = vi.fn().mockResolvedValue({ - "page_citation_assets/page-4.png": - "https://blob.example/page_citation_assets/page-4.png", - }) + const hardenChatAssetUrl = vi + .fn() + .mockResolvedValue("https://blob.example/page_citation_assets/page-4.png") const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ results: [ @@ -90,9 +91,15 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { }), ], sources: [makeSource()], - loadSourceAssetUrls, + hardenChatAssetUrl, }) + expect(hardenChatAssetUrl).toHaveBeenCalledWith({ + source: expect.objectContaining({ id: "source_1" }), + sourcePath: "page_citation_assets/page-4.png", + assetUrl: "https://assets.example/pages/page-4.png", + contentType: undefined, + }) expect(result?.pageCitationAssetUrl).toBe( "https://blob.example/page_citation_assets/page-4.png", ) diff --git a/src/domains/chat/page-citation-assets.ts b/src/domains/chat/page-citation-assets.ts index 77b8d53..566cead 100644 --- a/src/domains/chat/page-citation-assets.ts +++ b/src/domains/chat/page-citation-assets.ts @@ -3,7 +3,7 @@ import "server-only" import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" -import type { LoadSourceAssetUrls } from "./media-assets" +import type { HardenChatAssetUrl } from "./media-assets" export type PageCitationAssetRetrievalResult = RetrievalResult & { readonly pageCitationAssetUrl?: string @@ -12,19 +12,20 @@ export type PageCitationAssetRetrievalResult = RetrievalResult & { type EnrichRetrievalResultsWithPageCitationAssetUrlsInput = { readonly results: readonly RetrievalResult[] readonly sources: readonly Source[] - readonly loadSourceAssetUrls?: LoadSourceAssetUrls + readonly hardenChatAssetUrl?: HardenChatAssetUrl } type PageCitationAssetCandidate = { readonly pageNum: number readonly artifactRef?: string readonly assetUrl?: string + readonly contentType?: string } export async function enrichRetrievalResultsWithPageCitationAssetUrls({ results, sources, - loadSourceAssetUrls, + hardenChatAssetUrl, }: EnrichRetrievalResultsWithPageCitationAssetUrlsInput): Promise< PageCitationAssetRetrievalResult[] > { @@ -35,18 +36,12 @@ export async function enrichRetrievalResultsWithPageCitationAssetUrls({ source.knowhereDocumentId ? [[source.knowhereDocumentId, source]] : [], ), ) - const assetUrlsBySourceId = new Map< - string, - Promise>> - >() - return Promise.all( results.map((result) => enrichRetrievalResultWithPageCitationAssetUrl({ result, sourcesByDocumentId, - loadSourceAssetUrls, - assetUrlsBySourceId, + hardenChatAssetUrl, }), ), ) @@ -55,11 +50,7 @@ export async function enrichRetrievalResultsWithPageCitationAssetUrls({ async function enrichRetrievalResultWithPageCitationAssetUrl(input: { readonly result: RetrievalResult readonly sourcesByDocumentId: ReadonlyMap - readonly loadSourceAssetUrls?: LoadSourceAssetUrls - readonly assetUrlsBySourceId: Map< - string, - Promise>> - > + readonly hardenChatAssetUrl?: HardenChatAssetUrl }): Promise { if (!isPageResult(input.result)) return input.result @@ -69,8 +60,7 @@ async function enrichRetrievalResultWithPageCitationAssetUrl(input: { result: input.result, directAsset, sourcesByDocumentId: input.sourcesByDocumentId, - loadSourceAssetUrls: input.loadSourceAssetUrls, - assetUrlsBySourceId: input.assetUrlsBySourceId, + hardenChatAssetUrl: input.hardenChatAssetUrl, }) if (sourceAssetUrl) { return { @@ -86,25 +76,21 @@ async function getStoredPageCitationAssetUrl(input: { readonly result: RetrievalResult readonly directAsset: PageCitationAssetCandidate | null readonly sourcesByDocumentId: ReadonlyMap - readonly loadSourceAssetUrls?: LoadSourceAssetUrls - readonly assetUrlsBySourceId: Map< - string, - Promise>> - > + readonly hardenChatAssetUrl?: HardenChatAssetUrl }): Promise { const artifactRef = getTrimmedString(input.directAsset?.artifactRef) const documentId = getTrimmedString(input.result.source.documentId) - if (!artifactRef || !documentId || !input.loadSourceAssetUrls) return null + if (!artifactRef || !documentId || !input.hardenChatAssetUrl) return null const source = input.sourcesByDocumentId.get(documentId) if (!source) return null - const assetUrls = await getCachedSourceAssetUrls( + return input.hardenChatAssetUrl({ source, - input.loadSourceAssetUrls, - input.assetUrlsBySourceId, - ) - return getTrimmedString(assetUrls[artifactRef]) + sourcePath: artifactRef, + assetUrl: input.directAsset?.assetUrl, + contentType: input.directAsset?.contentType, + }).catch(() => null) } function isPageResult(result: RetrievalResult): boolean { @@ -142,24 +128,12 @@ function parsePageCitationAssetCandidates( pageNum, artifactRef: getTrimmedString(item.artifactRef) ?? undefined, assetUrl: getTrimmedString(item.assetUrl) ?? undefined, + contentType: getTrimmedString(item.contentType) ?? undefined, }, ] }) } -async function getCachedSourceAssetUrls( - source: Source, - loadSourceAssetUrls: LoadSourceAssetUrls, - cache: Map>>>, -): Promise>> { - let cached = cache.get(source.id) - if (!cached) { - cached = loadSourceAssetUrls(source).catch(() => ({})) - cache.set(source.id, cached) - } - return cached -} - function getPageNumbers( metadata: Readonly> | undefined, ): readonly number[] { diff --git a/src/domains/chat/route-answer.ts b/src/domains/chat/route-answer.ts index 9582b77..03c26ca 100644 --- a/src/domains/chat/route-answer.ts +++ b/src/domains/chat/route-answer.ts @@ -11,12 +11,12 @@ import { type ChatTurnValue, } from "@/domains/chat/service" import { chatTurnPersistence } from "@/domains/chat/chat-turn-persistence" -import { readSourceAssetUrls } from "@/domains/chunks/read" import { startBackgroundReconciliation } from "@/domains/sources/background-reconcile" +import { BlobParsedDocumentStorage } from "@/domains/sources/parsed-document-blob-storage" import { sourceWorkflowRuntime } from "@/domains/sources/workflow-runtime" -import { makeKnowhereClientWithParsedStorage } from "@/integrations/knowhere" import { notebookRequestContext } from "@/domains/workspace/request-context" import type { Source } from "@/infrastructure/db/schema" +import type { HardenChatAssetUrl } from "./media-assets" import { isAuthError } from "@/integrations/dashboard/api-key-service" import { summarizeUnknownError } from "@/lib/format-log-value" import { logger } from "@/lib/logger" @@ -69,30 +69,23 @@ const answerChatEffect = (input: AnswerChatInput) => apiKey, }), ) - const { knowledge } = makeKnowhereClientWithParsedStorage(apiKey, { + const parsedStorage = new BlobParsedDocumentStorage({ workspaceId: workspace.id, }) - const loadSourceAssetUrls = async ( - source: (typeof sources)[number], - ): Promise>> => { - if (source.status !== "ready" || !source.knowhereDocumentId) return {} - - try { - return await readSourceAssetUrls({ - knowledge, - documentId: source.knowhereDocumentId, - revisionKey: source.knowhereJobId, - }) - } catch (error) { - logger.warn("chat: durable parsed asset read failed", { - workspaceId: workspace.id, - sourceId: source.id, - documentId: source.knowhereDocumentId, - error: summarizeUnknownError(error), - }) - return {} - } - } + const hardenChatAssetUrl: HardenChatAssetUrl = async ({ + source, + sourcePath, + assetUrl, + contentType, + }): Promise => + hardenSingleChatAsset({ + workspaceId: workspace.id, + parsedStorage, + source, + sourcePath, + assetUrl, + contentType, + }) const result: Either.Either = yield* Effect.tryPromise(() => @@ -104,14 +97,14 @@ const answerChatEffect = (input: AnswerChatInput) => excludedSourceIds: body.value.excludedSourceIds, retrieval: client.retrieval, generateAnswer: generateAgenticOutputManifest, - loadSourceAssetUrls, + hardenChatAssetUrl, hardenMediaAssetUrls: ({ results, artifacts }) => hardenChatMediaAssetUrls({ workspaceId: workspace.id, sources, results, artifacts, - loadSourceAssetUrls, + hardenChatAssetUrl, }), repository: chatTurnPersistence.createRepository(), }), @@ -167,6 +160,104 @@ export const chatAnswerRouteService: ChatAnswerRouteService = { answerChat, } +async function hardenSingleChatAsset(input: { + readonly workspaceId: string + readonly parsedStorage: BlobParsedDocumentStorage + readonly source: Source + readonly sourcePath: string + readonly assetUrl?: string | null + readonly contentType?: string | null +}): Promise { + if ( + input.source.status !== "ready" || + !input.source.knowhereDocumentId || + !input.source.knowhereJobId + ) { + return null + } + + try { + const existingUrl = await input.parsedStorage.getAssetUrl({ + documentId: input.source.knowhereDocumentId, + revisionKey: input.source.knowhereJobId, + sourcePath: input.sourcePath, + }) + if (existingUrl) return existingUrl + + const sourceUrl = getTrimmedString(input.assetUrl) + if (!sourceUrl) return null + + const fetchedAsset = await fetchChatAsset({ + assetUrl: sourceUrl, + fallbackContentType: + getTrimmedString(input.contentType) ?? + inferContentTypeFromPath(input.sourcePath), + }) + const writtenAsset = await input.parsedStorage.writeAsset({ + documentId: input.source.knowhereDocumentId, + revisionKey: input.source.knowhereJobId, + sourcePath: input.sourcePath, + body: fetchedAsset.body, + contentType: fetchedAsset.contentType, + }) + + return ( + writtenAsset.url ?? + (await input.parsedStorage.getAssetUrl({ + documentId: input.source.knowhereDocumentId, + revisionKey: input.source.knowhereJobId, + sourcePath: input.sourcePath, + })) + ) + } catch (error) { + logger.warn("chat: single parsed asset hardening failed", { + workspaceId: input.workspaceId, + sourceId: input.source.id, + documentId: input.source.knowhereDocumentId, + sourcePath: input.sourcePath, + error: summarizeUnknownError(error), + }) + return null + } +} + +async function fetchChatAsset(input: { + readonly assetUrl: string + readonly fallbackContentType: string +}): Promise<{ + readonly body: Uint8Array + readonly contentType: string +}> { + const response = await fetch(input.assetUrl) + if (!response.ok) { + throw new Error( + `Failed to fetch chat asset (${response.status} ${response.statusText})`, + ) + } + + return { + body: new Uint8Array(await response.arrayBuffer()), + contentType: + getTrimmedString(response.headers.get("content-type")) ?? + input.fallbackContentType, + } +} + +function inferContentTypeFromPath(sourcePath: string): string { + const pathname = sourcePath.toLowerCase().split("?")[0] ?? sourcePath + if (pathname.endsWith(".png")) return "image/png" + if (pathname.endsWith(".jpg") || pathname.endsWith(".jpeg")) { + return "image/jpeg" + } + if (pathname.endsWith(".gif")) return "image/gif" + if (pathname.endsWith(".webp")) return "image/webp" + if (pathname.endsWith(".svg")) return "image/svg+xml" + if (pathname.endsWith(".html") || pathname.endsWith(".htm")) { + return "text/html; charset=utf-8" + } + return "application/octet-stream" +} + function triggerBackgroundReconciliationForParsingSources(input: { readonly workspaceId: string readonly sources: readonly Source[] @@ -246,3 +337,8 @@ function isMeaningfulSummary(value: string): boolean { normalized !== "An unknown error occurred in Effect.tryPromise" ) } + +function getTrimmedString(value: string | null | undefined): string | null { + const trimmedValue = value?.trim() ?? "" + return trimmedValue.length > 0 ? trimmedValue : null +} diff --git a/src/domains/chat/route-service.test.ts b/src/domains/chat/route-service.test.ts index f4df92a..6365edd 100644 --- a/src/domains/chat/route-service.test.ts +++ b/src/domains/chat/route-service.test.ts @@ -1,5 +1,5 @@ import { Either } from "effect" -import { beforeEach, describe, expect, it, vi } from "vitest" +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest" import type { ChatMessage, ChatThread, Source, Workspace } from "@/infrastructure/db/schema" @@ -18,8 +18,8 @@ const mocks = vi.hoisted(() => ({ loggerInfo: vi.fn(), loggerWarn: vi.fn(), listSourcesForWorkspace: vi.fn(), - makeKnowhereClientWithParsedStorage: vi.fn(), - readChunks: vi.fn(), + parsedStorageGetAssetUrl: vi.fn(), + parsedStorageWriteAsset: vi.fn(), softDeleteChatThread: vi.fn(), startBackgroundReconciliation: vi.fn(), })) @@ -53,9 +53,13 @@ vi.mock("@/domains/workspace/request-context", () => ({ }, })) -vi.mock("@/integrations/knowhere", () => ({ - makeKnowhereClientWithParsedStorage: - mocks.makeKnowhereClientWithParsedStorage, +vi.mock("@/domains/sources/parsed-document-blob-storage", () => ({ + BlobParsedDocumentStorage: vi.fn().mockImplementation(function () { + return { + getAssetUrl: mocks.parsedStorageGetAssetUrl, + writeAsset: mocks.parsedStorageWriteAsset, + } + }), })) vi.mock("@/domains/chat/thread-service", () => ({ @@ -84,18 +88,18 @@ import { chatThreadRouteService } from "./route-threads" describe("chat route services", () => { beforeEach(() => { vi.clearAllMocks() - mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ - client: { documents: { listChunks: vi.fn() } }, - knowledge: { readChunks: mocks.readChunks }, - }) - mocks.readChunks.mockResolvedValue({ - document: { localDocumentId: "doc" }, - chunks: [], - page: 1, - pageSize: 200, - totalChunks: 0, - totalPages: 1, - }) + mocks.parsedStorageGetAssetUrl.mockResolvedValue(null) + mocks.parsedStorageWriteAsset.mockResolvedValue({ url: null }) + vi.stubGlobal( + "fetch", + vi.fn(async () => new Response(new Uint8Array([1, 2, 3]), { + headers: { "content-type": "image/png" }, + })), + ) + }) + + afterEach(() => { + vi.unstubAllGlobals() }) it("orchestrates a chat turn from request body to response body", async () => { @@ -148,7 +152,7 @@ describe("chat route services", () => { excludedSourceIds: ["source_skipped"], retrieval: client.retrieval, generateAnswer: mocks.generateAgenticOutputManifest, - loadSourceAssetUrls: expect.any(Function), + hardenChatAssetUrl: expect.any(Function), repository: expect.objectContaining({ appendMessageToThread: expect.any(Function), ensureDefaultChatThread: expect.any(Function), @@ -159,7 +163,7 @@ describe("chat route services", () => { ) }) - it("builds durable citation asset URLs from SDK reads for a ready source", async () => { + it("hardens one chat asset through Notebook Blob for a ready source", async () => { const workspace = makeWorkspace() const client = { retrieval: { query: vi.fn() } } const readySource = makeSource({ @@ -167,29 +171,10 @@ describe("chat route services", () => { knowhereDocumentId: "doc_legacy", knowhereJobId: "job_1", }) + const rawUrl = "https://knowhere-storage.example/results/job_1/pages/page-1.png" const durableUrl = "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_legacy/job_1/assets/pages/page-1.png" - mocks.readChunks.mockResolvedValue({ - document: { localDocumentId: "doc_legacy" }, - chunks: [ - { - position: 1, - chunkId: "c1", - chunkType: "page", - content: "Page", - readableContent: "Page", - sectionPath: "Page 1", - sourceChunkPath: "Page 1", - filePath: "pages/page-1.png", - assetUrl: durableUrl, - metadata: {}, - }, - ], - page: 1, - pageSize: 200, - totalChunks: 1, - totalPages: 1, - }) + mocks.parsedStorageWriteAsset.mockResolvedValue({ url: durableUrl }) mocks.getAuthenticatedWithClient.mockResolvedValue({ user: { id: "user_1" }, workspace, @@ -199,12 +184,20 @@ describe("chat route services", () => { mocks.listSourcesForWorkspace.mockResolvedValue([readySource]) mocks.handleChatTurn.mockImplementation( async (input: { - readonly loadSourceAssetUrls?: ( - source: Source, - ) => Promise>> + readonly hardenChatAssetUrl?: (assetInput: { + readonly source: Source + readonly sourcePath: string + readonly assetUrl?: string | null + readonly contentType?: string | null + }) => Promise }) => { - const assetUrls = await input.loadSourceAssetUrls?.(readySource) - expect(assetUrls).toEqual({ "pages/page-1.png": durableUrl }) + const assetUrl = await input.hardenChatAssetUrl?.({ + source: readySource, + sourcePath: "pages/page-1.png", + assetUrl: rawUrl, + contentType: "image/png", + }) + expect(assetUrl).toBe(durableUrl) return Either.right({ threadId: "thread_1", messages: [ @@ -220,16 +213,18 @@ describe("chat route services", () => { }) expect(result.status).toBe(200) - expect(mocks.makeKnowhereClientWithParsedStorage).toHaveBeenCalledWith( - "jwt_123", - { workspaceId: workspace.id }, - ) - expect(mocks.readChunks).toHaveBeenCalledWith({ + expect(mocks.parsedStorageGetAssetUrl).toHaveBeenCalledWith({ + documentId: "doc_legacy", + revisionKey: "job_1", + sourcePath: "pages/page-1.png", + }) + expect(fetch).toHaveBeenCalledWith(rawUrl) + expect(mocks.parsedStorageWriteAsset).toHaveBeenCalledWith({ documentId: "doc_legacy", revisionKey: "job_1", - page: 1, - pageSize: 200, - assetUrlPolicy: "durable", + sourcePath: "pages/page-1.png", + body: new Uint8Array([1, 2, 3]), + contentType: "image/png", }) }) diff --git a/src/domains/chat/service.ts b/src/domains/chat/service.ts index 5361d5b..8b67e7c 100644 --- a/src/domains/chat/service.ts +++ b/src/domains/chat/service.ts @@ -67,7 +67,7 @@ type ChatTurnInput = { excludedSourceIds: readonly string[] retrieval: RetrievalClient generateAnswer: GenerateAnswer - loadSourceAssetUrls?: AnswerQuestionInput["loadSourceAssetUrls"] + hardenChatAssetUrl?: AnswerQuestionInput["hardenChatAssetUrl"] hardenMediaAssetUrls?: AnswerQuestionInput["hardenMediaAssetUrls"] repository: ChatRepository } @@ -126,7 +126,7 @@ export const handleChatTurnEffect = (input: ChatTurnInput) => excludedSourceIds: input.excludedSourceIds, retrieval: input.retrieval, generateAnswer: input.generateAnswer, - loadSourceAssetUrls: input.loadSourceAssetUrls, + hardenChatAssetUrl: input.hardenChatAssetUrl, hardenMediaAssetUrls: input.hardenMediaAssetUrls, messages: chatHistoryMessages, }).pipe(Effect.catchAllCause(Effect.die)) diff --git a/src/domains/chunks/index.ts b/src/domains/chunks/index.ts index 5e1b25a..56c2644 100644 --- a/src/domains/chunks/index.ts +++ b/src/domains/chunks/index.ts @@ -181,8 +181,8 @@ export function toParsedChunkView( /** * Map an SDK `KnowledgeReadChunk` (from `knowledge.readChunks`) to the view - * model. Asset URLs on durable reads are already hardened SDK-side, so no - * `assetUrlsByFilePath` remap is needed here. + * model. Display reads do not request durable asset URL hardening; chat + * hardens the specific assets it returns separately. */ export function toParsedChunkViewFromReadChunk( chunk: KnowledgeReadChunk, diff --git a/src/domains/chunks/read.test.ts b/src/domains/chunks/read.test.ts index 1d0786d..b6a3979 100644 --- a/src/domains/chunks/read.test.ts +++ b/src/domains/chunks/read.test.ts @@ -20,7 +20,7 @@ function makeReadChunk(overrides: Record = {}) { } describe("readSourceChunkPage", () => { - it("reads a durable page and maps chunks to the view model", async () => { + it("reads a page without durable asset hardening and maps chunks to the view model", async () => { const readChunks = vi.fn(async () => ({ document: { localDocumentId: "doc_1" }, chunks: [ @@ -48,7 +48,6 @@ describe("readSourceChunkPage", () => { revisionKey: "rev_1", page: 2, pageSize: 50, - assetUrlPolicy: "durable", }) expect(result.pagination).toEqual({ page: 2, @@ -85,7 +84,6 @@ describe("readSourceChunkPage", () => { documentId: "doc_1", page: 1, pageSize: 50, - assetUrlPolicy: "durable", }) }) }) diff --git a/src/domains/chunks/read.ts b/src/domains/chunks/read.ts index f375d48..320d559 100644 --- a/src/domains/chunks/read.ts +++ b/src/domains/chunks/read.ts @@ -16,8 +16,8 @@ type ReadableSource = { /** * Read a single display page of parsed chunks through the SDK. The SDK serves * from configured Blob storage when fresh and falls back to Knowhere remote - * otherwise, hardening visible asset URLs into durable Blob URLs - * (`assetUrlPolicy: "durable"`) and scheduling a background sync on a miss. + * otherwise. Display reads intentionally do not request durable asset URLs: + * chat hardens only the specific assets it sends back to the user. */ export async function readSourceChunkPage(input: { readonly knowledge: Knowledge @@ -29,7 +29,6 @@ export async function readSourceChunkPage(input: { ...(input.source.revisionKey ? { revisionKey: input.source.revisionKey } : {}), page: input.params.page, pageSize: input.params.pageSize, - assetUrlPolicy: "durable", }) const chunks = response.chunks.map((chunk) => @@ -69,7 +68,6 @@ export async function readAllSourceChunks(input: { : {}), page, pageSize: loadAllPageSize, - assetUrlPolicy: "durable", }) for (const chunk of response.chunks) { chunks.push( @@ -86,57 +84,3 @@ export async function readAllSourceChunks(input: { return chunks } - -/** - * Build a `filePath -> durable Blob URL` map for a source by paging durable - * reads to exhaustion. This is the single asset-hardening path for chat: the - * SDK writes any missing asset into Blob during the durable read and returns - * the durable URL, which we index by both the chunk file path and any - * `metadata.pageAssets[].artifactRef`. - */ -export async function readSourceAssetUrls(input: { - readonly knowledge: Knowledge - readonly documentId: string - readonly revisionKey?: string | null -}): Promise>> { - const assetUrlsByFilePath: Record = {} - let page = 1 - let totalPages = 1 - - do { - const response = await input.knowledge.readChunks({ - documentId: input.documentId, - ...(input.revisionKey ? { revisionKey: input.revisionKey } : {}), - page, - pageSize: loadAllPageSize, - assetUrlPolicy: "durable", - }) - for (const chunk of response.chunks) { - if (chunk.filePath && chunk.assetUrl) { - assetUrlsByFilePath[chunk.filePath] = chunk.assetUrl - } - collectPageAssetUrls(chunk.metadata, assetUrlsByFilePath) - } - totalPages = Math.max(1, response.totalPages ?? 1) - page += 1 - } while (page <= totalPages) - - return assetUrlsByFilePath -} - -function collectPageAssetUrls( - metadata: Record, - target: Record, -): void { - const pageAssets = metadata["pageAssets"] - if (!Array.isArray(pageAssets)) return - for (const pageAsset of pageAssets) { - if (typeof pageAsset !== "object" || pageAsset === null) continue - const record = pageAsset as Record - const artifactRef = record["artifactRef"] - const assetUrl = record["assetUrl"] - if (typeof artifactRef === "string" && typeof assetUrl === "string") { - target[artifactRef] = assetUrl - } - } -} diff --git a/src/domains/sources/source-reconcile-workflow.test.ts b/src/domains/sources/source-reconcile-workflow.test.ts index 5757da8..68769b7 100644 --- a/src/domains/sources/source-reconcile-workflow.test.ts +++ b/src/domains/sources/source-reconcile-workflow.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it, vi } from "vitest" -import type { JobResult } from "@ontos-ai/knowhere-sdk" +import { NotFoundError, type JobResult } from "@ontos-ai/knowhere-sdk" import type { Source, Workspace } from "@/infrastructure/db/schema" import { @@ -140,6 +140,68 @@ describe("pollSourceReconciliation", () => { ) }) + it("marks parsing sources failed when the Knowhere job no longer exists", async () => { + const source = makeSource({ + stagedBlobPathname: "source-uploads/upload_1/document.pdf", + }) + const repository = createRepository(source) + const blobStore = { + deleteStagedSourceBlob: vi.fn(async () => undefined), + } + const client = { + jobs: { + get: vi.fn(async () => { + throw new NotFoundError("Job not found") + }), + }, + } + + const result = await pollSourceReconciliation({ + workspaceId: workspace.id, + sourceId: "source_1", + client, + repository, + blobStore, + }) + + expect(result).toEqual({ kind: "resolved", status: "failed" }) + expect(repository.markFailed).toHaveBeenCalledWith( + workspace.id, + "source_1", + "Knowhere job job_1 was not found during reconciliation.", + "parsing", + ) + expect(blobStore.deleteStagedSourceBlob).toHaveBeenCalledWith( + "source-uploads/upload_1/document.pdf", + ) + expect(repository.clearStagedBlob).toHaveBeenCalledWith( + workspace.id, + "source_1", + ) + }) + + it("rethrows non-not-found job lookup errors for workflow retries", async () => { + const repository = createRepository() + const client = { + jobs: { + get: vi.fn(async () => { + throw new Error("Knowhere timeout") + }), + }, + } + + await expect( + pollSourceReconciliation({ + workspaceId: workspace.id, + sourceId: "source_1", + client, + repository, + }), + ).rejects.toThrow("Knowhere timeout") + + expect(repository.markFailed).not.toHaveBeenCalled() + }) + it("accepts done jobs without a result URL when a document id is published", async () => { const repository = createRepository() const jobWithoutResultUrl = makeDoneJob() diff --git a/src/domains/sources/source-reconcile-workflow.ts b/src/domains/sources/source-reconcile-workflow.ts index 4571038..7556c76 100644 --- a/src/domains/sources/source-reconcile-workflow.ts +++ b/src/domains/sources/source-reconcile-workflow.ts @@ -1,7 +1,7 @@ import "server-only" import { del } from "@vercel/blob" -import type { JobResult } from "@ontos-ai/knowhere-sdk" +import { NotFoundError, type JobResult } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" import { logger } from "@/lib/logger" @@ -13,6 +13,12 @@ type SourceReconcileWorkflowClient = { } } +type ApiErrorLike = { + readonly name?: string + readonly statusCode?: number + readonly message?: string +} + type SourceReconcileWorkflowRepository = { readonly findInWorkspace: ( workspaceId: string, @@ -112,7 +118,15 @@ export async function pollSourceReconciliation({ return { kind: "resolved", status: "failed" } } - const job = await client.jobs.get(jobId) + const job = await getJobOrFailMissingJob({ + workspaceId, + source, + jobId, + client, + repository, + blobStore, + }) + if (!job) return { kind: "resolved", status: "failed" } if (isFailedJob(job)) { await failSourceAndCleanup({ workspaceId, @@ -150,6 +164,30 @@ export async function pollSourceReconciliation({ } } +async function getJobOrFailMissingJob(input: { + readonly workspaceId: string + readonly source: Source + readonly jobId: string + readonly client: SourceReconcileWorkflowClient + readonly repository: SourceReconcileWorkflowRepository + readonly blobStore: SourceReconcileWorkflowBlobStore +}): Promise { + try { + return await input.client.jobs.get(input.jobId) + } catch (error) { + if (!isKnowhereJobNotFoundError(error)) throw error + + await failSourceAndCleanup({ + workspaceId: input.workspaceId, + source: input.source, + reason: `Knowhere job ${input.jobId} was not found during reconciliation.`, + repository: input.repository, + blobStore: input.blobStore, + }) + return null + } +} + export async function markSourceReadyAfterReconciliation({ workspaceId, sourceId, @@ -245,6 +283,17 @@ function isFailedJob(job: JobResult): boolean { return job.isFailed || job.status === "failed" } +function isKnowhereJobNotFoundError(error: unknown): boolean { + if (error instanceof NotFoundError) return true + if (!isApiErrorLike(error)) return false + + return error.name === "NotFoundError" || error.statusCode === 404 +} + +function isApiErrorLike(error: unknown): error is ApiErrorLike { + return error !== null && typeof error === "object" +} + const vercelBlobStore: SourceReconcileWorkflowBlobStore = { deleteStagedSourceBlob: del, } From 6fa2d046b394b6ed6555399bcb77359a83665383 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 5 Jul 2026 23:52:01 +0800 Subject: [PATCH 15/51] Fix source selection, archive idempotency, and tool metadata --- src/agent-harness/runtime.test.ts | 75 ++++++++++++ src/agent-harness/runtime.ts | 67 ++++++++++- src/components/workspace-source-state.test.ts | 47 ++++++++ src/components/workspace-source-state.ts | 6 +- src/domains/sources/route-archive.ts | 50 ++++++++ src/domains/sources/route-service.test.ts | 112 ++++++++++++++++++ 6 files changed, 354 insertions(+), 3 deletions(-) diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index 5657e5c..5470690 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -5,6 +5,7 @@ import { buildHarnessMessages, buildHarnessSystemPrompt, createHarnessTools, + sanitizeHarnessModelMessagesForStep, } from "./runtime" import { createEvidenceLedger } from "./ledger" import type { @@ -251,6 +252,80 @@ describe("agent harness runtime", () => { expect(JSON.stringify(messages)).toContain("id=turn_1 role=assistant") expect(JSON.stringify(messages)).not.toContain("searchSources.query") }) + + it("removes provider metadata from tool-result parts while preserving tool-call metadata", () => { + const messages = sanitizeHarnessModelMessagesForStep([ + { + role: "assistant", + content: [ + { + type: "tool-call", + toolCallId: "call_1", + toolName: "retrieve", + input: { query: "q4" }, + providerOptions: { + google: { + thoughtSignature: "signature-1", + }, + }, + }, + ], + }, + { + role: "tool", + content: [ + { + type: "tool-result", + toolCallId: "call_1", + toolName: "retrieve", + output: { + type: "json", + value: { + ok: true, + }, + }, + providerOptions: { + google: { + thoughtSignature: "signature-1", + }, + }, + }, + ], + }, + ]) + + expect(messages).toEqual([ + { + role: "assistant", + content: [ + expect.objectContaining({ + type: "tool-call", + providerOptions: { + google: { + thoughtSignature: "signature-1", + }, + }, + }), + ], + }, + { + role: "tool", + content: [ + { + type: "tool-result", + toolCallId: "call_1", + toolName: "retrieve", + output: { + type: "json", + value: { + ok: true, + }, + }, + }, + ], + }, + ]) + }) }) function executeTool(tool: unknown, input: unknown): Promise { diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index c03baac..895dfc5 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -1,4 +1,10 @@ -import { stepCountIs, ToolLoopAgent, tool, type ModelMessage } from "ai" +import { + stepCountIs, + ToolLoopAgent, + tool, + type ModelMessage, + type ToolResultPart, +} from "ai" import { z } from "zod" import { createEvidenceLedger } from "./ledger" @@ -147,6 +153,9 @@ export async function runAgentHarness( model: input.model, instructions: buildHarnessSystemPrompt(input.turn), tools, + prepareStep: ({ messages: stepMessages }) => ({ + messages: sanitizeHarnessModelMessagesForStep(stepMessages), + }), stopWhen: stepCountIs(input.maxSteps ?? defaultMaxSteps), }) @@ -205,6 +214,62 @@ export async function runAgentHarness( } } +export function sanitizeHarnessModelMessagesForStep( + messages: readonly ModelMessage[], +): ModelMessage[] { + return messages.map(sanitizeHarnessModelMessageForStep) +} + +function sanitizeHarnessModelMessageForStep( + message: ModelMessage, +): ModelMessage { + if (message.role === "tool") { + return { + ...message, + providerOptions: undefined, + content: message.content.map(sanitizeToolMessageContentPart), + } + } + + if (message.role === "assistant" && Array.isArray(message.content)) { + return { + ...message, + content: message.content.map(sanitizeAssistantMessageContentPart), + } + } + + return message +} + +function sanitizeToolMessageContentPart( + part: ModelMessageForRole<"tool">["content"][number], +): ModelMessageForRole<"tool">["content"][number] { + if (part.type !== "tool-result") return part + + return removeToolResultPartProviderOptions(part) +} + +function sanitizeAssistantMessageContentPart( + part: Exclude["content"], string>[number], +): Exclude["content"], string>[number] { + if (part.type !== "tool-result") return part + + return removeToolResultPartProviderOptions(part) +} + +function removeToolResultPartProviderOptions( + part: ToolResultPart, +): ToolResultPart { + const sanitizedPart = { ...part } + delete sanitizedPart.providerOptions + return sanitizedPart +} + +type ModelMessageForRole = Extract< + ModelMessage, + { readonly role: TRole } +> + function buildRevisionFeedback(errors: readonly string[]): string { return [ "Your finalize output did not satisfy the output contract:", diff --git a/src/components/workspace-source-state.test.ts b/src/components/workspace-source-state.test.ts index 19e6cb1..bfd0273 100644 --- a/src/components/workspace-source-state.test.ts +++ b/src/components/workspace-source-state.test.ts @@ -104,6 +104,53 @@ describe("workspaceSourceState", () => { ).toBe("source_localized"); }); + it("keeps an explicit non-ready Source selected instead of falling back", () => { + const sources: readonly SourceView[] = [ + { + id: "source_pending", + title: "pending.pdf", + status: "parsing", + mimeType: "application/pdf", + excludedFromQuery: false, + }, + { + id: "source_ready", + title: "ready.pdf", + status: "ready", + mimeType: "application/pdf", + excludedFromQuery: false, + }, + ]; + + expect( + workspaceSourceState.getResolvedSelectedSourceId( + sources, + "source_pending", + ), + ).toBe("source_pending"); + }); + + it("does not resolve a stale selected Source to an unrelated ready Source", () => { + const sources: readonly SourceView[] = [ + { + id: "demo_ready", + kind: "demo", + demoSourceId: "demo_ready", + title: "demo.pdf", + status: "ready", + mimeType: "application/pdf", + excludedFromQuery: false, + }, + ]; + + expect( + workspaceSourceState.getResolvedSelectedSourceId( + sources, + "source_stale", + ), + ).toBeNull(); + }); + it("applies source query exclusions without mutating the source list", () => { const sources: readonly SourceView[] = [ { diff --git a/src/components/workspace-source-state.ts b/src/components/workspace-source-state.ts index 0373983..97c0c55 100644 --- a/src/components/workspace-source-state.ts +++ b/src/components/workspace-source-state.ts @@ -69,8 +69,10 @@ function getResolvedSelectedSourceId( sources: readonly SourceView[], selectedSourceId: string | null, ): string | null { + if (!selectedSourceId) return getFirstReadySourceId(sources) + const selectedSource = sources.find((source) => source.id === selectedSourceId) - if (selectedSource && isReadySource(selectedSource)) { + if (selectedSource) { return selectedSource.id } @@ -83,7 +85,7 @@ function getResolvedSelectedSourceId( if (localizedSource) return localizedSource.id } - return getFirstReadySourceId(sources) + return null } function isReadySource(source: SourceView): boolean { diff --git a/src/domains/sources/route-archive.ts b/src/domains/sources/route-archive.ts index 897dee3..6047256 100644 --- a/src/domains/sources/route-archive.ts +++ b/src/domains/sources/route-archive.ts @@ -72,6 +72,8 @@ const archiveSourceEffect = ( ) yield* Effect.tryPromise(() => client.documents.archive(source.knowhereDocumentId!), + ).pipe( + Effect.catchIf(isKnowhereDocumentNotFoundError, () => Effect.void), ) } @@ -92,4 +94,52 @@ const archiveSourceEffect = ( return routeResult.ok({ id: input.sourceId, archived: true as const }) }) +function isKnowhereDocumentNotFoundError(error: unknown): boolean { + return readErrorMessages(error).some(isDocumentNotFoundMessage) +} + +function readErrorMessages(error: unknown): readonly string[] { + if (error instanceof Error) { + return [ + error.message, + ...readNestedErrorMessages(error), + ].filter(isNonEmptyString) + } + + if (typeof error === "string") return [error] + + if (!isRecord(error)) return [] + + const messages: string[] = [] + if (typeof error.message === "string") messages.push(error.message) + messages.push(...readNestedErrorMessages(error)) + return messages +} + +function readNestedErrorMessages(error: unknown): readonly string[] { + if (!isRecord(error)) return [] + + const nested = [ + error.error, + error.cause, + isRecord(error.body) ? error.body.error : undefined, + ] + + return nested.flatMap((value) => + value === undefined || value === error ? [] : readErrorMessages(value), + ) +} + +function isDocumentNotFoundMessage(message: string): boolean { + return /\bdocument\s+not\s+found\b/iu.test(message) +} + +function isNonEmptyString(value: string): boolean { + return value.trim().length > 0 +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null +} + export { createRouteArchive } diff --git a/src/domains/sources/route-service.test.ts b/src/domains/sources/route-service.test.ts index b07439c..a056427 100644 --- a/src/domains/sources/route-service.test.ts +++ b/src/domains/sources/route-service.test.ts @@ -834,6 +834,118 @@ describe("source route service", () => { expect(onUploadFinished).toHaveBeenCalledOnce(); }); + it("soft-deletes a source when Knowhere says its document is already missing", async () => { + const readySource: Source = { + ...source, + status: "ready", + knowhereJobId: null, + knowhereDocumentId: "doc_missing", + }; + const archiveDocument = vi.fn(async () => { + throw new Error("Document not found"); + }); + const softDelete = vi.fn(async () => true); + const service = createSourceRouteService({ + ensureApiKeyForWorkspace: vi.fn(async () => "jwt_123"), + ensureWorkspace: vi.fn(async () => workspace), + makeKnowhereClient: vi.fn(() => ({ + documents: { + archive: archiveDocument, + listChunks: vi.fn(async () => ({ + chunks: [], + pagination: { + page: 1, + pageSize: 1, + total: 0, + totalPages: 0, + }, + })), + }, + jobs: { + create: vi.fn(), + get: vi.fn(), + upload: vi.fn(), + }, + })), + requireUser: vi.fn(async () => ({ + id: "user_1", + email: null, + name: null, + })), + sourceService: { + findInWorkspace: vi.fn(async () => readySource), + softDelete, + }, + }); + + const result = await service.archiveSource({ + cookieHeader: "session=abc", + sourceId: "source_1", + }); + + expect(result).toEqual({ + status: 200, + body: { + id: "source_1", + archived: true, + }, + }); + expect(archiveDocument).toHaveBeenCalledWith("doc_missing"); + expect(softDelete).toHaveBeenCalledWith(workspace.id, "source_1"); + }); + + it("does not soft-delete a source when Knowhere archive fails unexpectedly", async () => { + const readySource: Source = { + ...source, + status: "ready", + knowhereJobId: null, + knowhereDocumentId: "doc_unavailable", + }; + const softDelete = vi.fn(async () => true); + const service = createSourceRouteService({ + ensureApiKeyForWorkspace: vi.fn(async () => "jwt_123"), + ensureWorkspace: vi.fn(async () => workspace), + makeKnowhereClient: vi.fn(() => ({ + documents: { + archive: vi.fn(async () => { + throw new Error("Knowhere unavailable"); + }), + listChunks: vi.fn(async () => ({ + chunks: [], + pagination: { + page: 1, + pageSize: 1, + total: 0, + totalPages: 0, + }, + })), + }, + jobs: { + create: vi.fn(), + get: vi.fn(), + upload: vi.fn(), + }, + })), + requireUser: vi.fn(async () => ({ + id: "user_1", + email: null, + name: null, + })), + sourceService: { + findInWorkspace: vi.fn(async () => readySource), + softDelete, + }, + }); + + await expect( + service.archiveSource({ + cookieHeader: "session=abc", + sourceId: "source_1", + }), + ).rejects.toThrow(); + expect(softDelete).not.toHaveBeenCalled(); + }); + it("retries a failed source from its saved original Blob", async () => { const failedSource: Source = { ...source, From 8c1a443362e85044736758a5557c1ba9384c31b2 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 6 Jul 2026 03:08:43 +0800 Subject: [PATCH 16/51] Fix chat retrieval payload guard --- src/agent-harness/runtime.test.ts | 67 +++++++++++++++++ src/agent-harness/runtime.ts | 11 ++- src/domains/chat/index.test.ts | 119 ++++++++++++++++++++++++++---- src/domains/chat/index.ts | 74 ++++++++++++++++--- 4 files changed, 247 insertions(+), 24 deletions(-) diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index 5470690..ef7e5e0 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -98,6 +98,73 @@ describe("agent harness runtime", () => { ]) }) + it("returns only newly retrieved evidence in each retrieve tool result", async () => { + const query = vi + .fn() + .mockResolvedValueOnce(makeRetrievalResponse()) + .mockResolvedValueOnce({ + ...makeRetrievalResponse(), + query: "second query", + evidenceText: "Second evidence", + results: [ + { + content: "Second retrieval evidence.", + chunkType: "text", + score: 0.8, + source: { + documentId: "doc_2", + sourceFileName: "second.pdf", + sectionPath: "Second", + }, + }, + ], + }) + const state: { + intent?: IntentFrame + contextPolicy?: ContextPolicy + toolCalls?: HarnessToolCallTrace[] + } = { + intent: { + task: "answer", + dependsOnPreviousTurn: false, + retrievalNeeded: "yes", + targetModalities: ["text"], + constraints: {}, + groundingPolicy: "must_use_sources", + }, + contextPolicy: { + carryHistory: "none", + reason: "Self-contained request.", + activePriorTurnIds: [], + }, + } + const ledger = createEvidenceLedger() + const tools = createHarnessTools({ + state, + ledger, + retrieval: { query }, + recentTurns: [], + }) + + const firstResult = await executeTool(tools.retrieve, { query: "first" }) + const secondResult = await executeTool(tools.retrieve, { query: "second" }) + + expect(firstResult).toMatchObject({ + retrievalCount: 1, + chunks: [{ ref: "r1:result:1" }], + }) + expect(secondResult).toMatchObject({ + retrievalCount: 2, + evidenceText: "Second evidence", + chunks: [{ ref: "r2:result:1" }], + }) + expect(JSON.stringify(secondResult)).not.toContain("r1:result:1") + expect(ledger.snapshot().chunks.map((chunk) => chunk.ref)).toEqual([ + "r1:result:1", + "r2:result:1", + ]) + }) + it("blocks finalize until intent and context policy are declared", async () => { const state: { intent?: IntentFrame diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index 895dfc5..fabc627 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -353,6 +353,7 @@ export function createHarnessTools(input: { } } + const beforeSnapshot = input.ledger.snapshot() const response = await input.retrieval.query({ query: request.query, modalities: request.modalities as TargetModality[], @@ -363,13 +364,19 @@ export function createHarnessTools(input: { threshold: request.threshold, }) const snapshot = input.ledger.addRetrievalResponse(response) + const currentChunks = snapshot.chunks.slice( + beforeSnapshot.chunks.length, + ) + const currentAssets = snapshot.assets.slice( + beforeSnapshot.assets.length, + ) return { ok: true, retrievalCount: snapshot.retrievalCount, evidenceText: response.evidenceText ?? "", stopReason: response.stopReason ?? null, failureReason: response.failureReason ?? null, - chunks: snapshot.chunks.map((chunk) => ({ + chunks: currentChunks.map((chunk) => ({ ref: chunk.ref, kind: chunk.kind, type: chunk.chunkType, @@ -377,7 +384,7 @@ export function createHarnessTools(input: { source: chunk.source, assetRef: chunk.assetRef, })), - assets: snapshot.assets.map((asset) => ({ + assets: currentAssets.map((asset) => ({ ref: asset.ref, type: asset.type, label: asset.label, diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 0be6b78..d16f26e 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -8,6 +8,7 @@ import { answerQuestionWithRetrieval, generateAgenticOutputManifest, parseChatRequestBody, + type SearchSources, } from "." import type { HardenableRetrievalResult, @@ -174,6 +175,91 @@ describe("answerQuestionWithRetrieval", () => { }); }); + it("bounds merged retrieval evidence before passing it to the answer agent", async () => { + const defaultResults = Array.from({ length: 40 }, (_, index) => + makeRetrievalResult({ + content: `Default namespace result ${index + 1}`, + source: { + documentId: `doc_default_${index + 1}`, + sourceFileName: "default.pdf", + sectionPath: `Default ${index + 1}`, + }, + }), + ); + const workspaceResults = Array.from({ length: 40 }, (_, index) => + makeRetrievalResult({ + content: `Workspace result ${index + 1}`, + source: { + documentId: `doc_workspace_${index + 1}`, + sourceFileName: "workspace.pdf", + sectionPath: `Workspace ${index + 1}`, + }, + }), + ); + const referencedChunks = Array.from({ length: 40 }, (_, index) => ({ + chunkId: `chunk_${index + 1}`, + documentId: `doc_reference_${index + 1}`, + chunkType: "text" as const, + sectionPath: `Reference ${index + 1}`, + })); + const retrieval = { + query: vi + .fn() + .mockResolvedValueOnce({ + results: defaultResults, + evidenceText: "Default evidence", + referencedChunks, + namespace: "default", + query: "large response", + routerUsed: "workflow_single_step", + answerText: null, + stopReason: "answer_done", + failureReason: null, + }) + .mockResolvedValueOnce({ + results: workspaceResults, + evidenceText: "Workspace evidence", + referencedChunks, + namespace: "notebook-workspace", + query: "large response", + routerUsed: "workflow_single_step", + answerText: null, + stopReason: "answer_done", + failureReason: null, + }), + }; + const generateAnswer = vi.fn( + async ({ searchSources }: { searchSources: SearchSources }) => { + const response = await searchSources({ + query: "large response", + topK: 3, + }); + expect(response.results).toHaveLength(6); + expect(response.referencedChunks).toHaveLength(6); + expect(response.results.map((result) => result.content)).toEqual( + [ + ...defaultResults.slice(0, 3), + ...workspaceResults.slice(0, 3), + ].map((result) => result.content), + ); + return makeHarnessRunResult("The answer is grounded."); + }, + ); + + await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "What does the document say?", + namespace: "notebook-workspace", + namespaces: ["default", "notebook-workspace"], + sources: [makeSource()], + excludedSourceIds: [], + retrieval, + generateAnswer, + messages: [], + }), + ); + }); + it("does not hide a failed namespace query behind an empty namespace result", async () => { const retrievalError = new Error("Legacy namespace query failed."); const retrieval = { @@ -235,7 +321,14 @@ describe("answerQuestionWithRetrieval", () => { }); const retrieval = { query: vi.fn().mockResolvedValue({ - results: [result], + results: [ + result, + ...Array.from({ length: 30 }, (_, index) => + makeRetrievalResult({ + content: `extra result ${index + 1}`, + }), + ), + ], evidenceText: `Evidence https://blob.example/evidence.jpg ${"evidence ".repeat( 80, )}`, @@ -251,6 +344,12 @@ describe("answerQuestionWithRetrieval", () => { jobId: "job_1", assetUrl: "https://blob.example/id.jpg", }, + ...Array.from({ length: 30 }, (_, index) => ({ + chunkId: `chunk_extra_${index + 1}`, + documentId: "doc_identity", + chunkType: "text" as const, + sectionPath: `Extra referenced chunk ${index + 1}`, + })), ], namespace: "notebook-workspace", query: "冯荣洲 身份证 ID card", @@ -284,19 +383,13 @@ describe("answerQuestionWithRetrieval", () => { const response = meta.response as KnowhereQueryResponseLogMeta; expect(response).toMatchObject({ query: "冯荣洲 身份证 ID card", - resultCount: 1, - referencedChunkCount: 1, - results: [ - { - chunkType: "image", - }, - ], - referencedChunks: [ - { - chunkType: "image", - }, - ], + resultCount: 31, + referencedChunkCount: 31, }); + expect(response.results).toHaveLength(20); + expect(response.referencedChunks).toHaveLength(20); + expect(response.results[0]).toMatchObject({ chunkType: "image" }); + expect(response.referencedChunks[0]).toMatchObject({ chunkType: "image" }); expect(response.answerText.length).toBeLessThanOrEqual(203); expect(response.evidenceText.length).toBeLessThanOrEqual(203); expect(response.results[0]?.content.length).toBeLessThanOrEqual(103); diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index dc10453..21e6ef8 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -43,9 +43,13 @@ import type { HardenableRetrievalResult } from "./media-asset-hardening" const DEFAULT_TOP_K = 8 const MAX_AGENTIC_TOP_K = 12 +const MAX_AGENTIC_MERGED_RESULT_COUNT = 24 +const MAX_AGENTIC_MERGED_REFERENCED_CHUNK_COUNT = 24 +const MAX_AGENTIC_MERGED_TEXT_CHARS = 12_000 const MAX_CITATION_RESULTS = 20 const KNOWHERE_RESPONSE_TEXT_LOG_LIMIT = 200 const KNOWHERE_CHUNK_LOG_LIMIT = 100 +const KNOWHERE_RESPONSE_LOG_ITEM_LIMIT = 20 const NO_RESULTS_ANSWER = "I couldn't find that in your sources." const HARNESS_VALIDATION_FAILURE_ANSWER = "I couldn't safely finish that response because the agent output did not pass Notebook's validation checks. Please try again." @@ -88,6 +92,13 @@ type KnowhereReferencedChunkLog = { readonly summary: string } +type AgenticMergedEvidenceLimits = { + readonly resultCountPerResponse: number + readonly referencedChunkCountPerResponse: number + readonly resultCount: number + readonly referencedChunkCount: number +} + export type { AnswerQuestionInput, AnswerQuestionResult, @@ -186,7 +197,14 @@ export const answerQuestionWithRetrieval = ( ) { throw queryFailures[0] } - return mergeRetrievalResponses(queryResponses, retrievalPlan) + return mergeRetrievalResponses( + queryResponses, + retrievalPlan, + getAgenticMergedEvidenceLimits({ + namespaceCount: queryResponses.length, + topK: queryInput.topK, + }), + ) } const generatedAnswer = yield* Effect.tryPromise(() => @@ -593,10 +611,12 @@ function formatKnowhereQueryResponseForLog( response.evidenceText ?? "", KNOWHERE_RESPONSE_TEXT_LOG_LIMIT, ), - results: response.results.map(formatKnowhereResultChunkForLog), - referencedChunks: response.referencedChunks.map( - formatKnowhereReferencedChunkForLog, - ), + results: response.results + .slice(0, KNOWHERE_RESPONSE_LOG_ITEM_LIMIT) + .map(formatKnowhereResultChunkForLog), + referencedChunks: response.referencedChunks + .slice(0, KNOWHERE_RESPONSE_LOG_ITEM_LIMIT) + .map(formatKnowhereReferencedChunkForLog), } } @@ -654,6 +674,7 @@ function getRetrievalNamespaces(input: AnswerQuestionInput): readonly string[] { function mergeRetrievalResponses( responses: readonly RetrievalQueryResponse[], retrievalPlan: AgenticRetrievalPlan, + evidenceLimits: AgenticMergedEvidenceLimits, ): AgenticRetrievalResponse { const [first] = responses if (!first) { @@ -661,16 +682,27 @@ function mergeRetrievalResponses( } const statusResponses = getRetrievalStatusResponses(responses) - const results = responses.flatMap((response) => response.results) - const referencedChunks = responses.flatMap( - (response) => response.referencedChunks, - ) + const results = responses + .flatMap((response) => + response.results.slice(0, evidenceLimits.resultCountPerResponse), + ) + .slice(0, evidenceLimits.resultCount) + const referencedChunks = responses + .flatMap((response) => + response.referencedChunks.slice( + 0, + evidenceLimits.referencedChunkCountPerResponse, + ), + ) + .slice(0, evidenceLimits.referencedChunkCount) const evidenceTexts = responses .map((response) => response.evidenceText) .filter((value): value is string => Boolean(value)) + .map(truncateAgenticModelText) const answerTexts = responses .map((response) => response.answerText) .filter((value): value is string => Boolean(value)) + .map(truncateAgenticModelText) return { ...first, @@ -693,6 +725,30 @@ function mergeRetrievalResponses( } } +function getAgenticMergedEvidenceLimits(input: { + readonly namespaceCount: number + readonly topK: number | undefined +}): AgenticMergedEvidenceLimits { + const namespaceCount = Math.max(input.namespaceCount, 1) + const perResponseCount = normalizeTopK(input.topK) + const requestedCount = perResponseCount * namespaceCount + return { + resultCountPerResponse: perResponseCount, + referencedChunkCountPerResponse: perResponseCount, + resultCount: Math.min(requestedCount, MAX_AGENTIC_MERGED_RESULT_COUNT), + referencedChunkCount: Math.min( + requestedCount, + MAX_AGENTIC_MERGED_REFERENCED_CHUNK_COUNT, + ), + } +} + +function truncateAgenticModelText(value: string): string { + const trimmed = value.trim() + if (trimmed.length <= MAX_AGENTIC_MERGED_TEXT_CHARS) return trimmed + return `${trimmed.slice(0, MAX_AGENTIC_MERGED_TEXT_CHARS)}\n...[truncated]` +} + function getRetrievalStatusResponses( responses: readonly RetrievalQueryResponse[], ): readonly RetrievalQueryResponse[] { From a8b66e6d11cc5b6994dc12c54b566b2b51a73055 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 6 Jul 2026 04:03:45 +0800 Subject: [PATCH 17/51] Fix chat harness finalization budget --- src/agent-harness/runtime.test.ts | 83 +++++++++++++++++++++++++++++++ src/agent-harness/runtime.ts | 65 ++++++++++++++++++++++-- 2 files changed, 143 insertions(+), 5 deletions(-) diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index ef7e5e0..f6abf7d 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -5,6 +5,7 @@ import { buildHarnessMessages, buildHarnessSystemPrompt, createHarnessTools, + prepareHarnessStep, sanitizeHarnessModelMessagesForStep, } from "./runtime" import { createEvidenceLedger } from "./ledger" @@ -393,6 +394,88 @@ describe("agent harness runtime", () => { }, ]) }) + + it("keeps normal steps unconstrained before the finalization step", () => { + const result = prepareHarnessStep({ + stepNumber: 11, + messages: [ + { + role: "user", + content: "Find the penalty amount.", + }, + ], + }) + + expect(result).toEqual({ + messages: [ + { + role: "user", + content: "Find the penalty amount.", + }, + ], + }) + }) + + it("forces finalize at step 12 using existing tool results", () => { + const result = prepareHarnessStep({ + stepNumber: 12, + messages: [ + { + role: "tool", + content: [ + { + type: "tool-result", + toolCallId: "call_1", + toolName: "retrieve", + output: { + type: "json", + value: { + ok: true, + chunks: [{ ref: "r1:result:1" }], + }, + }, + providerOptions: { + google: { + thoughtSignature: "signature-1", + }, + }, + }, + ], + }, + ], + }) + + expect(result.activeTools).toEqual(["finalize"]) + expect(result.toolChoice).toEqual({ + type: "tool", + toolName: "finalize", + }) + expect(result.messages).toEqual([ + { + role: "tool", + content: [ + { + type: "tool-result", + toolCallId: "call_1", + toolName: "retrieve", + output: { + type: "json", + value: { + ok: true, + chunks: [{ ref: "r1:result:1" }], + }, + }, + }, + ], + }, + { + role: "user", + content: expect.stringContaining( + "Use only the evidence and tool results already available", + ), + }, + ]) + }) }) function executeTool(tool: unknown, input: unknown): Promise { diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index fabc627..b565918 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -1,4 +1,5 @@ import { + hasToolCall, stepCountIs, ToolLoopAgent, tool, @@ -22,8 +23,9 @@ import type { } from "./types" import { validateOutputManifest } from "./validator" -const defaultMaxSteps = 10 +const defaultMaxSteps = 13 const defaultMaxRevisions = 1 +const forcedFinalizationStepNumber = 12 type ToolLoopAgentSettings = ConstructorParameters[0] @@ -50,6 +52,17 @@ type HarnessToolState = { toolCalls?: HarnessToolCallTrace[] } +type HarnessTools = ReturnType + +type HarnessStepPreparation = { + messages: ModelMessage[] + activeTools?: Array> + toolChoice?: { + type: "tool" + toolName: Extract + } +} + const targetModalitySchema = z.enum(["text", "image", "table"]) const intentFrameSchema = z.object({ @@ -153,10 +166,15 @@ export async function runAgentHarness( model: input.model, instructions: buildHarnessSystemPrompt(input.turn), tools, - prepareStep: ({ messages: stepMessages }) => ({ - messages: sanitizeHarnessModelMessagesForStep(stepMessages), - }), - stopWhen: stepCountIs(input.maxSteps ?? defaultMaxSteps), + prepareStep: ({ messages: stepMessages, stepNumber }) => + prepareHarnessStep({ + messages: stepMessages, + stepNumber, + }), + stopWhen: [ + hasToolCall("finalize"), + stepCountIs(input.maxSteps ?? defaultMaxSteps), + ], }) const maxRevisions = input.maxRevisions ?? defaultMaxRevisions @@ -214,6 +232,32 @@ export async function runAgentHarness( } } +export function prepareHarnessStep(input: { + readonly stepNumber: number + readonly messages: readonly ModelMessage[] +}): HarnessStepPreparation { + const messages = sanitizeHarnessModelMessagesForStep(input.messages) + + if (input.stepNumber < forcedFinalizationStepNumber) { + return { messages } + } + + return { + messages: [ + ...messages, + { + role: "user", + content: buildForcedFinalizationFeedback(), + }, + ], + activeTools: ["finalize"], + toolChoice: { + type: "tool", + toolName: "finalize", + }, + } +} + export function sanitizeHarnessModelMessagesForStep( messages: readonly ModelMessage[], ): ModelMessage[] { @@ -284,6 +328,17 @@ function buildRevisionFeedback(errors: readonly string[]): string { ].join("\n") } +function buildForcedFinalizationFeedback(): string { + return [ + "The retrieval step budget has been reached.", + "Do not search again or call any evidence-reading tools.", + "Use only the evidence and tool results already available in this turn.", + "Call finalize now with the best supported answer.", + "If the existing evidence is insufficient, explain the gap in unresolved", + "instead of making unsupported claims.", + ].join("\n") +} + export function createHarnessTools(input: { readonly state: HarnessToolState readonly ledger: ReturnType From f83660041f6fce601ae10b0440fe415262e7ac82 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 6 Jul 2026 16:09:57 +0800 Subject: [PATCH 18/51] Add agentic Gemini image inspection --- src/agent-harness/ledger.test.ts | 46 +++ src/agent-harness/ledger.ts | 239 +++++++++++- src/agent-harness/runtime.test.ts | 334 ++++++++++++++++- src/agent-harness/runtime.ts | 281 ++++++++++++++- src/agent-harness/types.ts | 39 +- src/agent-harness/validator.test.ts | 143 ++++++++ src/agent-harness/validator.ts | 81 +++++ src/domains/chat/contracts.ts | 4 +- src/domains/chat/index.test.ts | 291 +++++++++++++++ src/domains/chat/index.ts | 5 +- src/domains/chat/prompt.ts | 3 + src/domains/chat/route-answer.ts | 480 ++++++++++++++++++++++++- src/domains/chat/route-service.test.ts | 448 +++++++++++++++++++++++ src/domains/chat/service.ts | 2 + 14 files changed, 2369 insertions(+), 27 deletions(-) diff --git a/src/agent-harness/ledger.test.ts b/src/agent-harness/ledger.test.ts index 5685bc3..22e0f6a 100644 --- a/src/agent-harness/ledger.test.ts +++ b/src/agent-harness/ledger.test.ts @@ -43,6 +43,26 @@ describe("createEvidenceLedger", () => { contentSlice: "", }) }) + + it("creates page image assets from live page citation asset URLs without metadata", () => { + const ledger = createEvidenceLedger() + + const snapshot = ledger.addRetrievalResponse( + makePageAssetUrlRetrievalResponse(), + ) + + expect(snapshot.assets).toContainEqual( + expect.objectContaining({ + ref: "asset:r1:referenced:1", + chunkRef: "r1:referenced:1", + type: "image", + assetUrl: + "https://knowhere-storage.example/results/job_1/page_citation_assets/page-8.png?AWSAccessKeyId=test", + sourcePath: "page_citation_assets/page-8.png", + revisionKey: "job_1", + }), + ) + }) }) function makeRetrievalResponse(): RetrievalQueryResponse { @@ -89,3 +109,29 @@ function makeRetrievalResponse(): RetrievalQueryResponse { ], } } + +function makePageAssetUrlRetrievalResponse(): RetrievalQueryResponse { + return { + namespace: "notebook", + query: "承包人 进度计划 修改 违约金", + routerUsed: "workflow_single_step", + answerText: null, + evidenceText: + "Root / （6）现场工期进度管理方面的违约责任 [Page PDF (page 8)]", + stopReason: "answer_done", + failureReason: null, + results: [], + referencedChunks: [ + { + chunkId: "node_3a513cf7-77d7-5c62-a9bd-6a1109123e2c", + documentId: "doc_contract", + chunkType: "page", + sectionPath: "Root / （6）现场工期进度管理方面的违约责任", + filePath: null, + jobId: "job_1", + assetUrl: + "https://knowhere-storage.example/results/job_1/page_citation_assets/page-8.png?AWSAccessKeyId=test", + }, + ], + } +} diff --git a/src/agent-harness/ledger.ts b/src/agent-harness/ledger.ts index 103f9c0..f9d859d 100644 --- a/src/agent-harness/ledger.ts +++ b/src/agent-harness/ledger.ts @@ -22,6 +22,20 @@ type MutableLedger = { decisionTraces: unknown[] } +type EvidenceAssetCandidate = { + readonly type: EvidenceAsset["type"] + readonly assetUrl?: string + readonly sourcePath?: string + readonly label: string +} + +type PageCitationAssetCandidate = { + readonly pageNum: number + readonly artifactRef?: string + readonly assetUrl?: string + readonly contentType?: string +} + export type EvidenceLedger = ReturnType export function createEvidenceLedger() { @@ -81,6 +95,7 @@ export function createEvidenceLedger() { sourceFileName: null, sectionPath: chunk.sectionPath, }, + ...(chunk.jobId ? { revisionKey: chunk.jobId } : {}), ...(chunk.assetUrl ? { assetUrl: chunk.assetUrl } : {}), }, }) @@ -168,13 +183,12 @@ function addChunk(input: { readonly ledger: MutableLedger readonly chunk: Omit }): void { - const assetUrl = input.chunk.assetUrl?.trim() - if (!assetUrl || !isRenderableAsset(input.chunk.chunkType, assetUrl)) { + const asset = getEvidenceAssetCandidate(input.chunk) + if (!asset) { input.ledger.chunks.push(input.chunk) return } - const type = getAssetType(input.chunk.chunkType, assetUrl) const assetRef = `asset:${input.chunk.ref}` const chunk: EvidenceChunk = { ...input.chunk, @@ -184,10 +198,12 @@ function addChunk(input: { input.ledger.assets.push({ ref: assetRef, chunkRef: chunk.ref, - type, - assetUrl, + type: asset.type, + ...(asset.assetUrl ? { assetUrl: asset.assetUrl } : {}), + ...(asset.sourcePath ? { sourcePath: asset.sourcePath } : {}), + ...(chunk.revisionKey ? { revisionKey: chunk.revisionKey } : {}), source: chunk.source, - label: formatAssetLabel(chunk), + label: asset.label, }) } @@ -206,6 +222,56 @@ function isRenderableAsset(chunkType: string, assetUrl: string): boolean { ) } +function getEvidenceAssetCandidate( + chunk: Omit, +): EvidenceAssetCandidate | null { + const pageAsset = getPageCitationAssetCandidate(chunk) + if (pageAsset) { + return pageAsset + } + + const assetUrl = getTrimmedString(chunk.assetUrl) + if (!assetUrl || !isRenderableAsset(chunk.chunkType, assetUrl)) return null + + const sourcePath = getAssetSourcePath(chunk, assetUrl) + return { + type: getAssetType(chunk.chunkType, assetUrl), + assetUrl, + ...(sourcePath ? { sourcePath } : {}), + label: formatAssetLabel(chunk, sourcePath), + } +} + +function getPageCitationAssetCandidate( + chunk: Omit, +): EvidenceAssetCandidate | null { + if (chunk.chunkType.toLowerCase() !== "page") return null + + const candidates = parsePageCitationAssetCandidates(chunk.metadata?.pageAssets) + .filter(isSupportedPageCitationAsset) + if (candidates.length === 0) return null + + const pageNumbers = getPageNumbers(chunk.metadata) + const candidate = + pageNumbers.length > 0 + ? candidates.find((item) => pageNumbers.includes(item.pageNum)) ?? + candidates[0] + : candidates[0] + if (!candidate) return null + + const sourcePath = getTrimmedString(candidate.artifactRef) + const assetUrl = + getTrimmedString(candidate.assetUrl) ?? getTrimmedString(chunk.assetUrl) + if (!sourcePath && !assetUrl) return null + + return { + type: "image", + ...(assetUrl ? { assetUrl } : {}), + ...(sourcePath ? { sourcePath } : {}), + label: formatAssetLabel(chunk, sourcePath), + } +} + function getAssetType(chunkType: string, assetUrl: string): "image" | "table" { return chunkType.toLowerCase() === "table" && !isImageAssetUrl(assetUrl) ? "table" @@ -217,6 +283,37 @@ function isImageAssetUrl(assetUrl: string): boolean { return imageExtensions.some((extension) => pathname.endsWith(extension)) } +function getAssetSourcePath( + chunk: Omit, + assetUrl: string, +): string | null { + const candidates = [ + chunk.filePath, + chunk.sourceChunkPath, + chunk.source.sectionPath, + getUrlPathname(assetUrl), + ] + + for (const candidate of candidates) { + const sourcePath = getSupportedAssetPath(candidate) + if (sourcePath) return sourcePath + } + + return null +} + +function getSupportedAssetPath(value: string | null | undefined): string | null { + const normalizedText = normalizeSourcePathCandidate(value) + if (!normalizedText) return null + + const match = + /(?:^|\/)((?:images|tables|pages|page_citation_assets)\/[^?#]+)(?:[?#]|$)?/i.exec( + normalizedText, + ) + const matchedPath = match?.[1] + return matchedPath ? matchedPath.trim() : null +} + function getUrlPathname(assetUrl: string): string { try { return new URL(assetUrl).pathname @@ -225,14 +322,136 @@ function getUrlPathname(assetUrl: string): string { } } -function formatAssetLabel(chunk: EvidenceChunk): string { - return [ +function parsePageCitationAssetCandidates( + value: unknown, +): readonly PageCitationAssetCandidate[] { + if (!Array.isArray(value)) return [] + + return value.flatMap((item): PageCitationAssetCandidate[] => { + if (!isRecord(item)) return [] + const pageNum = getPositiveInteger(item.pageNum) + if (!pageNum) return [] + + return [ + { + pageNum, + ...(getTrimmedString(item.artifactRef) + ? { artifactRef: getTrimmedString(item.artifactRef) ?? undefined } + : {}), + ...(getTrimmedString(item.assetUrl) + ? { assetUrl: getTrimmedString(item.assetUrl) ?? undefined } + : {}), + ...(getTrimmedString(item.contentType) + ? { contentType: getTrimmedString(item.contentType) ?? undefined } + : {}), + }, + ] + }) +} + +function isSupportedPageCitationAsset( + candidate: PageCitationAssetCandidate, +): boolean { + const contentType = candidate.contentType?.toLowerCase() + return ( + contentType?.startsWith("image/") === true || + hasImageExtension(candidate.artifactRef) || + hasImageExtension(candidate.assetUrl) + ) +} + +function hasImageExtension(value: string | null | undefined): boolean { + const normalized = normalizeSourcePathCandidate(value)?.toLowerCase() + return ( + normalized !== undefined && + imageExtensions.some((extension) => normalized.endsWith(extension)) + ) +} + +function getPageNumbers( + metadata: Readonly> | undefined, +): readonly number[] { + if (!metadata) return [] + + const values = [metadata.pageNums, metadata.page_nums, metadata.pageNum] + const pageNumbers = new Set() + + for (const value of values) { + if (Array.isArray(value)) { + for (const item of value) { + const pageNum = getPositiveInteger(item) + if (pageNum) pageNumbers.add(pageNum) + } + continue + } + + const pageNum = getPositiveInteger(value) + if (pageNum) pageNumbers.add(pageNum) + } + + return [...pageNumbers].sort((left, right) => left - right) +} + +function getTrimmedString(value: unknown): string | null { + if (typeof value !== "string") return null + const trimmed = value.trim() + return trimmed.length > 0 ? trimmed : null +} + +function getPositiveInteger(value: unknown): number | null { + return typeof value === "number" && + Number.isSafeInteger(value) && + value > 0 + ? value + : null +} + +function normalizeSourcePathCandidate( + value: string | null | undefined, +): string | null { + const trimmedValue = getTrimmedString(value) + if (!trimmedValue) return null + + const normalized = decodeUrlText(trimmedValue) + .replaceAll("\\", "/") + .replace(/\s*\/\s*/g, "/") + .replace(/\s+/g, " ") + .trim() + + return normalized.length > 0 ? normalized : null +} + +function decodeUrlText(value: string): string { + try { + return decodeURIComponent(value) + } catch { + return value + } +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null +} + +function formatAssetLabel( + chunk: EvidenceChunk, + sourcePath?: string | null, +): string { + const labels = [ chunk.source.sourceFileName, chunk.source.sectionPath, + sourcePath, chunk.chunkType, ] - .filter((value): value is string => Boolean(value?.trim())) - .join(" / ") + const uniqueLabels: string[] = [] + + for (const label of labels) { + const normalized = label?.trim() + if (!normalized || uniqueLabels.includes(normalized)) continue + uniqueLabels.push(normalized) + } + + return uniqueLabels.join(" / ") } function snapshot(ledger: MutableLedger): EvidenceLedgerSnapshot { diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index f6abf7d..3e3c5fa 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -166,6 +166,198 @@ describe("agent harness runtime", () => { ]) }) + it("rejects image inspection before retrieval has returned image assets", async () => { + const inspectImages = vi.fn() + const tools = createHarnessTools({ + state: {}, + ledger: createEvidenceLedger(), + retrieval: { query: vi.fn() }, + inspectImages, + recentTurns: [], + }) + + const result = await executeTool(tools.inspectImage, { + refs: ["asset:r1:result:1"], + question: "What text is visible?", + }) + + expect(result).toEqual({ + ok: false, + message: "retrieve must be called before inspectImage.", + inspected: [], + skipped: [ + { + ref: "asset:r1:result:1", + reason: "No retrieval evidence is available yet.", + }, + ], + }) + expect(inspectImages).not.toHaveBeenCalled() + }) + + it("rejects image inspection refs that are unknown or not image assets", async () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse(makeTableRetrievalResponse()) + const inspectImages = vi.fn() + const tools = createHarnessTools({ + state: {}, + ledger, + retrieval: { query: vi.fn() }, + inspectImages, + recentTurns: [], + }) + + const result = await executeTool(tools.inspectImage, { + refs: ["asset:r1:result:1", "missing"], + question: "What is in these images?", + }) + + expect(result).toEqual({ + ok: false, + message: "No inspectable image asset refs were provided.", + inspected: [], + skipped: [ + { + ref: "asset:r1:result:1", + reason: "Ref is not an image asset.", + }, + { + ref: "missing", + reason: "Ref was not returned by retrieve as an asset.", + }, + ], + }) + expect(inspectImages).not.toHaveBeenCalled() + }) + + it("enforces the inspectImage six-ref limit", async () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse(makeImageRetrievalResponse(7)) + const tools = createHarnessTools({ + state: {}, + ledger, + retrieval: { query: vi.fn() }, + inspectImages: vi.fn(), + recentTurns: [], + }) + + const result = await executeTool(tools.inspectImage, { + refs: Array.from({ length: 7 }, (_, index) => `asset:r1:result:${index + 1}`), + question: "Compare these images.", + }) + + expect(result).toMatchObject({ + ok: false, + message: "inspectImage accepts at most 6 refs per call.", + }) + }) + + it("calls the visual inspection capability with retrieved image ledger assets", async () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse(makeRetrievalResponse()) + const inspectImages = vi.fn().mockResolvedValue({ + analysis: "The image shows a Q4 revenue chart with a rising line.", + inspected: [ + { + ref: "asset:r1:result:1", + label: "report.pdf / images/chart.png / image", + }, + ], + skipped: [], + }) + const tools = createHarnessTools({ + state: {}, + ledger, + retrieval: { query: vi.fn() }, + inspectImages, + recentTurns: [], + }) + + const result = await executeTool(tools.inspectImage, { + refs: ["asset:r1:result:1"], + question: "What does the chart show?", + }) + + expect(inspectImages).toHaveBeenCalledWith({ + question: "What does the chart show?", + assets: [ + { + ref: "asset:r1:result:1", + label: "report.pdf / images/chart.png / image", + assetUrl: "https://assets.example/chart.png", + sourcePath: "images/chart.png", + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "images/chart.png", + }, + }, + ], + }) + expect(result).toEqual({ + ok: true, + analysis: "The image shows a Q4 revenue chart with a rising line.", + inspected: [ + { + ref: "asset:r1:result:1", + label: "report.pdf / images/chart.png / image", + }, + ], + skipped: [], + }) + }) + + it("calls the visual inspection capability with retrieved page citation assets", async () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse(makePageCitationRetrievalResponse()) + const inspectImages = vi.fn().mockResolvedValue({ + analysis: "The clause says the contractor pays 5000 yuan per occurrence.", + inspected: [ + { + ref: "asset:r1:referenced:1", + label: + "Root / （6）现场工期进度管理方面的违约责任 / page_citation_assets/page-8.png / page", + }, + ], + skipped: [], + }) + const tools = createHarnessTools({ + state: {}, + ledger, + retrieval: { query: vi.fn() }, + inspectImages, + recentTurns: [], + }) + + const result = await executeTool(tools.inspectImage, { + refs: ["asset:r1:referenced:1"], + question: "What liquidated damages amount is visible in this clause?", + }) + + expect(inspectImages).toHaveBeenCalledWith({ + question: "What liquidated damages amount is visible in this clause?", + assets: [ + { + ref: "asset:r1:referenced:1", + label: + "Root / （6）现场工期进度管理方面的违约责任 / page_citation_assets/page-8.png / page", + assetUrl: "https://assets.example/page-8.png", + sourcePath: "page_citation_assets/page-8.png", + revisionKey: "job_contract", + source: { + documentId: "doc_contract", + sourceFileName: null, + sectionPath: "Root / （6）现场工期进度管理方面的违约责任", + }, + }, + ], + }) + expect(result).toMatchObject({ + ok: true, + analysis: "The clause says the contractor pays 5000 yuan per occurrence.", + }) + }) + it("blocks finalize until intent and context policy are declared", async () => { const state: { intent?: IntentFrame @@ -416,9 +608,64 @@ describe("agent harness runtime", () => { }) }) - it("forces finalize at step 12 using existing tool results", () => { + it("forces image inspection before forced finalization when image assets are available", () => { const result = prepareHarnessStep({ stepNumber: 12, + hasUninspectedImageAssets: true, + messages: [ + { + role: "tool", + content: [ + { + type: "tool-result", + toolCallId: "call_1", + toolName: "retrieve", + output: { + type: "json", + value: { + ok: true, + assets: [{ ref: "asset:r1:referenced:1", type: "image" }], + }, + }, + }, + ], + }, + ], + }) + + expect(result.activeTools).toEqual(["inspectImage"]) + expect(result.toolChoice).toEqual({ + type: "tool", + toolName: "inspectImage", + }) + expect(result.messages).toEqual([ + { + role: "tool", + content: [ + { + type: "tool-result", + toolCallId: "call_1", + toolName: "retrieve", + output: { + type: "json", + value: { + ok: true, + assets: [{ ref: "asset:r1:referenced:1", type: "image" }], + }, + }, + }, + ], + }, + { + role: "user", + content: expect.stringContaining("call inspectImage now"), + }, + ]) + }) + + it("forces finalize at step 13 using existing tool results", () => { + const result = prepareHarnessStep({ + stepNumber: 13, messages: [ { role: "tool", @@ -521,3 +768,88 @@ function makeRetrievalResponse(): RetrievalQueryResponse { referencedChunks: [], } } + +function makeTableRetrievalResponse(): RetrievalQueryResponse { + return { + namespace: "notebook", + query: "q4 table", + routerUsed: "workflow_single_step", + answerText: null, + evidenceText: "Table evidence", + stopReason: "answer_done", + failureReason: null, + results: [ + { + content: "", + chunkType: "table", + score: 0.8, + assetUrl: "https://assets.example/tables/revenue.html", + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "tables/revenue.html", + }, + }, + ], + referencedChunks: [], + } +} + +function makeImageRetrievalResponse(count: number): RetrievalQueryResponse { + return { + namespace: "notebook", + query: "q4 images", + routerUsed: "workflow_single_step", + answerText: null, + evidenceText: "Image evidence", + stopReason: "answer_done", + failureReason: null, + results: Array.from({ length: count }, (_, index) => ({ + content: "", + chunkType: "image", + score: 0.8, + assetUrl: `https://assets.example/images/chart-${index + 1}.png`, + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: `images/chart-${index + 1}.png`, + }, + })), + referencedChunks: [], + } +} + +function makePageCitationRetrievalResponse(): RetrievalQueryResponse { + return { + namespace: "notebook", + query: "进度计划", + routerUsed: "workflow_single_step", + answerText: null, + evidenceText: "Root / （6）现场工期进度管理方面的违约责任", + stopReason: "answer_done", + failureReason: null, + results: [], + referencedChunks: [ + { + chunkId: "chunk_page_8", + documentId: "doc_contract", + chunkType: "page", + sectionPath: "Root / （6）现场工期进度管理方面的违约责任", + filePath: null, + jobId: "job_contract", + assetUrl: "https://assets.example/page-8.png", + metadata: { + pageNums: [8], + pageAssets: [ + { + pageNum: 8, + artifactRef: "page_citation_assets/page-8.png", + assetUrl: "https://assets.example/page-8.png", + contentType: "image/png", + }, + ], + }, + }, + ], + } +} diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index b565918..d38028e 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -16,6 +16,9 @@ import type { HarnessRunResult, HarnessToolCallTrace, HarnessTrace, + ImageInspectionAsset, + ImageInspectionResponse, + InspectImages, IntentFrame, OutputManifest, RetrievalCapability, @@ -23,9 +26,11 @@ import type { } from "./types" import { validateOutputManifest } from "./validator" -const defaultMaxSteps = 13 +const defaultMaxSteps = 14 const defaultMaxRevisions = 1 -const forcedFinalizationStepNumber = 12 +const imageInspectionReminderStepNumber = 12 +const forcedFinalizationStepNumber = 13 +const imageInspectionRefLimit = 6 type ToolLoopAgentSettings = ConstructorParameters[0] @@ -35,6 +40,7 @@ export type RunAgentHarnessInput = { readonly model: AgentHarnessModel readonly turn: AgentTurnInput readonly retrieval: RetrievalCapability + readonly inspectImages?: InspectImages readonly maxSteps?: number /** * How many times the agent may revise after a failed validation pass before @@ -49,6 +55,7 @@ type HarnessToolState = { finalizedManifest?: OutputManifest finalized?: boolean priorTurnReads?: string[] + inspectedImageRefs?: string[] toolCalls?: HarnessToolCallTrace[] } @@ -160,6 +167,7 @@ export async function runAgentHarness( state, ledger, retrieval: input.retrieval, + inspectImages: input.inspectImages, recentTurns: input.turn.recentTurns, }) const agent = new ToolLoopAgent({ @@ -170,6 +178,9 @@ export async function runAgentHarness( prepareHarnessStep({ messages: stepMessages, stepNumber, + hasUninspectedImageAssets: + input.inspectImages !== undefined && + hasUninspectedImageAssets({ state, ledger }), }), stopWhen: [ hasToolCall("finalize"), @@ -194,6 +205,7 @@ export async function runAgentHarness( contextPolicy: state.contextPolicy, finalized: state.finalized === true, ledger: ledger.snapshot(), + toolCalls: state.toolCalls, surface: input.turn.surface, }) validationErrors = validation.errors @@ -235,9 +247,30 @@ export async function runAgentHarness( export function prepareHarnessStep(input: { readonly stepNumber: number readonly messages: readonly ModelMessage[] + readonly hasUninspectedImageAssets?: boolean }): HarnessStepPreparation { const messages = sanitizeHarnessModelMessagesForStep(input.messages) + if ( + input.stepNumber === imageInspectionReminderStepNumber && + input.hasUninspectedImageAssets === true + ) { + return { + messages: [ + ...messages, + { + role: "user", + content: buildImageInspectionReminderFeedback(), + }, + ], + activeTools: ["inspectImage"], + toolChoice: { + type: "tool", + toolName: "inspectImage", + }, + } + } + if (input.stepNumber < forcedFinalizationStepNumber) { return { messages } } @@ -339,10 +372,32 @@ function buildForcedFinalizationFeedback(): string { ].join("\n") } +function buildImageInspectionReminderFeedback(): string { + return [ + "The retrieval step budget is nearly reached and retrieved image assets are available.", + "If the exact answer depends on OCR, page-image text, visual details, or image verification, call inspectImage now with the most relevant retrieved image asset refs.", + "If image inspection is not needed for this answer, call finalize using the evidence already available.", + "Do not search again.", + ].join("\n") +} + +function hasUninspectedImageAssets(input: { + readonly state: HarnessToolState + readonly ledger: ReturnType +}): boolean { + const inspectedRefs = new Set(input.state.inspectedImageRefs ?? []) + return input.ledger + .snapshot() + .assets.some( + (asset) => asset.type === "image" && !inspectedRefs.has(asset.ref), + ) +} + export function createHarnessTools(input: { readonly state: HarnessToolState readonly ledger: ReturnType readonly retrieval: RetrievalCapability + readonly inspectImages?: InspectImages readonly recentTurns: readonly AgentTurn[] }) { return { @@ -451,6 +506,29 @@ export function createHarnessTools(input: { }), }), + inspectImage: tool({ + description: + "Inspect retrieved image asset refs visually for OCR, visual details, comparisons, or verification. Use only after retrieve has returned image assets.", + inputSchema: z.object({ + refs: z.array(z.string().min(1)).min(1).max(imageInspectionRefLimit), + question: z.string().min(1), + }), + execute: async (request) => + traceToolCall(input.state, { + toolName: "inspectImage", + inputSummary: summarizeInspectImageRequest(request), + execute: async () => + inspectRetrievedImages({ + state: input.state, + ledger: input.ledger, + inspectImages: input.inspectImages, + refs: request.refs, + question: request.question, + }), + summarizeOutput: summarizeInspectImageOutput, + }), + }), + readEvidence: tool({ description: "Read more text from an evidence chunk already returned by KNOWHERE.", @@ -569,6 +647,173 @@ export function createHarnessTools(input: { } as const } +async function inspectRetrievedImages(input: { + readonly state: HarnessToolState + readonly ledger: ReturnType + readonly inspectImages?: InspectImages + readonly refs: readonly string[] + readonly question: string +}): Promise< + | ({ readonly ok: true } & ImageInspectionResponse) + | { + readonly ok: false + readonly message: string + readonly inspected: readonly [] + readonly skipped: readonly { + readonly ref: string + readonly reason: string + }[] + } +> { + const refs = getUniqueTrimmedRefs(input.refs) + const question = input.question.trim() + + if (refs.length === 0 || question.length === 0) { + return { + ok: false, + message: "At least one image asset ref and a question are required.", + inspected: [], + skipped: [], + } + } + if (refs.length > imageInspectionRefLimit) { + return { + ok: false, + message: `inspectImage accepts at most ${imageInspectionRefLimit} refs per call.`, + inspected: [], + skipped: refs.map((ref) => ({ + ref, + reason: "Too many refs were requested in one inspectImage call.", + })), + } + } + + const snapshot = input.ledger.snapshot() + if (snapshot.retrievalCount === 0) { + return { + ok: false, + message: "retrieve must be called before inspectImage.", + inspected: [], + skipped: refs.map((ref) => ({ + ref, + reason: "No retrieval evidence is available yet.", + })), + } + } + if (!input.inspectImages) { + return { + ok: false, + message: "Image inspection is not available for this turn.", + inspected: [], + skipped: refs.map((ref) => ({ + ref, + reason: "No image inspection capability is configured.", + })), + } + } + + const assetsByRef = new Map( + snapshot.assets.map((asset) => [asset.ref, asset] as const), + ) + const skipped: { + readonly ref: string + readonly reason: string + }[] = [] + const selectedAssets: ImageInspectionAsset[] = [] + + for (const ref of refs) { + const asset = assetsByRef.get(ref) + if (!asset) { + skipped.push({ + ref, + reason: "Ref was not returned by retrieve as an asset.", + }) + continue + } + if (asset.type !== "image") { + skipped.push({ + ref, + reason: "Ref is not an image asset.", + }) + continue + } + + selectedAssets.push({ + ref: asset.ref, + label: asset.label, + ...(asset.assetUrl ? { assetUrl: asset.assetUrl } : {}), + ...(asset.sourcePath ? { sourcePath: asset.sourcePath } : {}), + ...(asset.revisionKey ? { revisionKey: asset.revisionKey } : {}), + source: asset.source, + }) + } + + if (selectedAssets.length === 0) { + return { + ok: false, + message: "No inspectable image asset refs were provided.", + inspected: [], + skipped, + } + } + + const inspectedImageRefs = input.state.inspectedImageRefs ?? [] + const inspectedCountAfterCall = + inspectedImageRefs.length + selectedAssets.length + if (inspectedCountAfterCall > imageInspectionRefLimit) { + return { + ok: false, + message: `inspectImage accepts at most ${imageInspectionRefLimit} image refs per turn.`, + inspected: [], + skipped: selectedAssets.map((asset) => ({ + ref: asset.ref, + reason: "The per-turn image inspection limit would be exceeded.", + })), + } + } + + input.state.inspectedImageRefs = [ + ...inspectedImageRefs, + ...selectedAssets.map((asset) => asset.ref), + ] + + try { + const response = await input.inspectImages({ + question, + assets: selectedAssets, + }) + return { + ok: true, + analysis: response.analysis, + inspected: response.inspected, + skipped: [...skipped, ...response.skipped], + } + } catch (error) { + return { + ok: false, + message: + error instanceof Error + ? `Image inspection failed: ${error.message}` + : "Image inspection failed.", + inspected: [], + skipped: selectedAssets.map((asset) => ({ + ref: asset.ref, + reason: "The image inspection request failed.", + })), + } + } +} + +function getUniqueTrimmedRefs(refs: readonly string[]): string[] { + const normalizedRefs: string[] = [] + for (const ref of refs) { + const normalizedRef = ref.trim() + if (!normalizedRef || normalizedRefs.includes(normalizedRef)) continue + normalizedRefs.push(normalizedRef) + } + return normalizedRefs +} + async function traceToolCall(input: { readonly toolCalls?: HarnessToolCallTrace[] }, call: { @@ -683,6 +928,30 @@ function summarizeReadEvidenceOutput(output: unknown): unknown { } } +function summarizeInspectImageRequest(request: { + readonly refs: readonly string[] + readonly question: string +}): unknown { + return { + refs: getUniqueTrimmedRefs(request.refs), + questionLength: request.question.trim().length, + } +} + +function summarizeInspectImageOutput(output: unknown): unknown { + if (!isRecord(output)) return output + return { + ok: output.ok, + analysisLength: + typeof output.analysis === "string" ? output.analysis.length : 0, + inspectedCount: Array.isArray(output.inspected) + ? output.inspected.length + : 0, + skippedCount: Array.isArray(output.skipped) ? output.skipped.length : 0, + message: output.message, + } +} + function summarizeReadPriorTurnOutput(output: unknown): unknown { if (!isRecord(output)) return output return { @@ -739,8 +1008,10 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "2. Call setContextPolicy next, deciding how prior turns should influence this turn.", "3. When the policy needs prior-turn detail (references or corrections), call readPriorTurn for the relevant ids.", "4. Call retrieve only when evidence is needed. The query must be concise and self-contained.", - "5. Use readEvidence only for chunk refs already in the evidence ledger.", - "6. Call finalize with text, citations, artifacts, and unresolved issues. finalize requires declareIntent and setContextPolicy first.", + "5. For pixel-level details, OCR, visual comparison, image verification, or when the likely answer is only visible on a retrieved page/image asset, call inspectImage only after retrieve returned image asset refs.", + `6. inspectImage accepts at most ${imageInspectionRefLimit} image asset refs per call and per turn.`, + "7. Use readEvidence only for chunk refs already in the evidence ledger.", + "8. Call finalize with text, citations, artifacts, and unresolved issues. finalize requires declareIntent and setContextPolicy first.", "", "Context rules:", "- If the current user request is unrelated to prior turns, set carryHistory to none and do not reuse prior topics.", @@ -752,6 +1023,8 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "- artifacts with display=true are the exact images/tables shown. Never display every candidate; honor constraints.desiredCount / maxCount.", "- Use type=derived_table only for tables you create from evidence; every derived_table.sourceRefs entry must reference evidence in the ledger.", "- citations and selected image/table artifact refs may only reference refs returned by retrieve (in the evidence ledger).", + "- inspectImage observations are inspection notes, not new source refs. Final citations and displayed image artifacts must use the original retrieved image asset refs.", + "- If text evidence identifies a relevant page/image but does not include the exact fact, inspect the returned image asset for OCR/detail before saying the answer is unavailable.", "- If evidence is insufficient, list it in unresolved instead of fabricating facts.", "- After a validation-feedback message, fix all listed issues and call finalize again.", `Surface: ${turn.surface}`, diff --git a/src/agent-harness/types.ts b/src/agent-harness/types.ts index c559623..4601f6f 100644 --- a/src/agent-harness/types.ts +++ b/src/agent-harness/types.ts @@ -107,6 +107,7 @@ export type EvidenceChunk = { readonly sourceFileName?: string | null readonly sectionPath?: string | null } + readonly revisionKey?: string | null readonly assetRef?: string readonly assetUrl?: string } @@ -115,11 +116,47 @@ export type EvidenceAsset = { readonly ref: string readonly chunkRef: string readonly type: "image" | "table" - readonly assetUrl: string + readonly assetUrl?: string + readonly sourcePath?: string + readonly revisionKey?: string | null readonly source: EvidenceChunk["source"] readonly label: string } +export type ImageInspectionAsset = { + readonly ref: string + readonly label: string + readonly assetUrl?: string | null + readonly sourcePath?: string | null + readonly revisionKey?: string | null + readonly source: EvidenceChunk["source"] +} + +export type ImageInspectionSkippedAsset = { + readonly ref: string + readonly reason: string +} + +export type ImageInspectionInspectedAsset = { + readonly ref: string + readonly label: string +} + +export type ImageInspectionRequest = { + readonly question: string + readonly assets: readonly ImageInspectionAsset[] +} + +export type ImageInspectionResponse = { + readonly analysis: string + readonly inspected: readonly ImageInspectionInspectedAsset[] + readonly skipped: readonly ImageInspectionSkippedAsset[] +} + +export type InspectImages = ( + input: ImageInspectionRequest, +) => Promise + export type EvidenceLedgerSnapshot = { readonly retrievalCount: number readonly chunks: readonly EvidenceChunk[] diff --git a/src/agent-harness/validator.test.ts b/src/agent-harness/validator.test.ts index 2815241..be8de3e 100644 --- a/src/agent-harness/validator.test.ts +++ b/src/agent-harness/validator.test.ts @@ -232,6 +232,149 @@ describe("validateOutputManifest", () => { "Typing compose output must be insertion-ready plain text.", ) }) + + it("rejects image inspection failure claims when inspectImage was not called", () => { + const validation = validateOutputManifest({ + manifest: makeManifest({ + text: "图像检查工具未能成功读取该页面的具体条款内容。", + citations: [ + { + ref: "r1:referenced:1", + label: "contract.pdf / page 8", + source: { + documentId: "doc_contract", + sourceFileName: "contract.pdf", + sectionPath: "Root / （6）现场工期进度管理方面的违约责任", + }, + }, + ], + }), + intent: makeIntent({}), + contextPolicy: unrelatedContextPolicy, + ledger: { + ...emptyLedger, + chunks: [ + { + ref: "r1:referenced:1", + kind: "referenced_chunk", + content: "", + contentPreview: "", + chunkType: "page", + score: null, + source: { + documentId: "doc_contract", + sourceFileName: "contract.pdf", + sectionPath: "Root / （6）现场工期进度管理方面的违约责任", + }, + assetRef: "asset:r1:referenced:1", + }, + ], + assets: [makeAsset("asset:r1:referenced:1")], + }, + toolCalls: [], + surface: "notebook_chat", + }) + + expect(validation.errors).toContain( + "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", + ) + }) + + it("allows image inspection failure claims after inspectImage was called", () => { + const validation = validateOutputManifest({ + manifest: makeManifest({ + text: "对该页面的图像识别未成功提取具体数值。", + citations: [ + { + ref: "r1:referenced:1", + label: "contract.pdf / page 8", + source: { + documentId: "doc_contract", + sourceFileName: "contract.pdf", + sectionPath: "Root / （6）现场工期进度管理方面的违约责任", + }, + }, + ], + }), + intent: makeIntent({}), + contextPolicy: unrelatedContextPolicy, + ledger: { + ...emptyLedger, + chunks: [ + { + ref: "r1:referenced:1", + kind: "referenced_chunk", + content: "", + contentPreview: "", + chunkType: "page", + score: null, + source: { + documentId: "doc_contract", + sourceFileName: "contract.pdf", + sectionPath: "Root / （6）现场工期进度管理方面的违约责任", + }, + assetRef: "asset:r1:referenced:1", + }, + ], + assets: [makeAsset("asset:r1:referenced:1")], + }, + toolCalls: [ + { + tool: "inspectImage", + ok: false, + inputSummary: { refs: ["asset:r1:referenced:1"] }, + outputSummary: { inspectedCount: 0 }, + startedAt: "2026-07-06T00:00:00.000Z", + durationMs: 1, + }, + ], + surface: "notebook_chat", + }) + + expect(validation.errors).not.toContain( + "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", + ) + }) + + it("rejects unreadable page claims when image assets are available but uninspected", () => { + const validation = validateOutputManifest({ + manifest: makeManifest({ + text: "相关的页面图像无法通过目前的检测工具获取。", + unresolved: ["由于无法直接读取该页面的详细条款内容，目前无法给出确切的赔偿金额。"], + }), + intent: makeIntent({}), + contextPolicy: unrelatedContextPolicy, + ledger: { + ...emptyLedger, + chunks: [ + { + ref: "r1:referenced:1", + kind: "referenced_chunk", + content: "", + contentPreview: "", + chunkType: "page", + score: null, + source: { + documentId: "doc_contract", + sourceFileName: "contract.pdf", + sectionPath: "Root / （6）现场工期进度管理方面的违约责任", + }, + assetRef: "asset:r1:referenced:1", + }, + ], + assets: [makeAsset("asset:r1:referenced:1")], + }, + toolCalls: [], + surface: "notebook_chat", + }) + + expect(validation.errors).toContain( + "Final output must inspect available image assets before claiming retrieved page/image content cannot be read.", + ) + expect(validation.errors).toContain( + "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", + ) + }) }) const emptyLedger: EvidenceLedgerSnapshot = { diff --git a/src/agent-harness/validator.ts b/src/agent-harness/validator.ts index 89b5e43..dd67343 100644 --- a/src/agent-harness/validator.ts +++ b/src/agent-harness/validator.ts @@ -1,6 +1,7 @@ import type { ContextPolicy, EvidenceLedgerSnapshot, + HarnessToolCallTrace, IntentFrame, OutputManifest, } from "./types" @@ -11,6 +12,7 @@ export type ManifestValidationInput = { readonly contextPolicy?: ContextPolicy readonly finalized?: boolean readonly ledger: EvidenceLedgerSnapshot + readonly toolCalls?: readonly HarnessToolCallTrace[] readonly surface: "notebook_chat" | "typing_compose" | "typing_quick_ask" } @@ -34,6 +36,8 @@ export function validateOutputManifest( validateArtifactCounts(input, errors) validateGrounding(input, errors) validateTaskEvidence(input, errors) + validateImageInspectionClaims(input, errors) + validateUnavailableImageContentClaims(input, errors) validateTypingText(input, errors) return { @@ -201,3 +205,80 @@ function validateTypingText( errors.push("Typing compose output must be insertion-ready plain text.") } } + +function validateImageInspectionClaims( + input: ManifestValidationInput, + errors: string[], +): void { + const claimText = [input.manifest.text, ...input.manifest.unresolved].join("\n") + if (!mentionsImageInspectionResult(claimText)) return + + if (hasImageInspectionToolCall(input)) return + + errors.push( + "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", + ) +} + +function validateUnavailableImageContentClaims( + input: ManifestValidationInput, + errors: string[], +): void { + const hasImageAssets = input.ledger.assets.some((asset) => asset.type === "image") + if (!hasImageAssets) return + if (hasImageInspectionToolCall(input)) return + + const claimText = [input.manifest.text, ...input.manifest.unresolved].join("\n") + if (!mentionsUnavailableImageContent(claimText)) return + + errors.push( + "Final output must inspect available image assets before claiming retrieved page/image content cannot be read.", + ) +} + +function hasImageInspectionToolCall(input: ManifestValidationInput): boolean { + return input.toolCalls?.some((call) => call.tool === "inspectImage") === true +} + +function mentionsImageInspectionResult(value: string): boolean { + const normalized = value.replace(/\s+/g, " ").trim() + if (!normalized) return false + + return ( + /\b(?:image|visual)\s+(?:inspection|recognition|analysis)\s+(?:failed|did not|could not|was unable|found|showed|confirmed)/iu.test( + normalized, + ) || + /\bOCR\s+(?:failed|did not|could not|was unable|found|showed|confirmed|extracted)/iu.test( + normalized, + ) || + /(?:图像|图片|视觉).{0,8}(?:识别|检查|检视|查看|分析).{0,12}(?:未|没|无法|不能|不成功|失败|成功|显示|发现|提取|读取)/u.test( + normalized, + ) || + /(?:图像|图片|视觉).{0,12}(?:未|没|无法|不能|不成功|失败).{0,16}(?:识别|检查|检视|查看|分析|检测|提取|读取|获取)/u.test( + normalized, + ) || + /(?:未|没|无法|不能|不成功|失败).{0,12}(?:图像|图片|视觉|OCR).{0,12}(?:识别|检查|检视|查看|分析|检测|提取|读取|获取)/u.test( + normalized, + ) + ) +} + +function mentionsUnavailableImageContent(value: string): boolean { + const normalized = value.replace(/\s+/g, " ").trim() + if (!normalized) return false + + return ( + /\b(?:cannot|can't|could not|unable to|was unable to|did not)\s+(?:directly\s+)?(?:read|inspect|access|extract|see|view)\b.{0,48}\b(?:page|image|visual|OCR|content|clause|details|text)\b/iu.test( + normalized, + ) || + /\b(?:page|image|visual|OCR|content|clause|details|text)\b.{0,48}\b(?:cannot|can't|could not|unable to|was unable to|did not)\s+(?:directly\s+)?(?:read|inspect|access|extract|see|view)\b/iu.test( + normalized, + ) || + /(?:无法|不能|未能|没法|没有办法).{0,8}(?:直接)?(?:读取|查看|识别|检测|提取|看清|访问|获取).{0,16}(?:页面|页|图片|图像|条款|内容|细节|文字)/u.test( + normalized, + ) || + /(?:页面|页|图片|图像|条款|内容|细节|文字).{0,16}(?:无法|不能|未能|没法|没有办法).{0,16}(?:直接)?(?:读取|查看|识别|检测|提取|看清|访问|获取)/u.test( + normalized, + ) + ) +} diff --git a/src/domains/chat/contracts.ts b/src/domains/chat/contracts.ts index 4752b86..3400af5 100644 --- a/src/domains/chat/contracts.ts +++ b/src/domains/chat/contracts.ts @@ -4,7 +4,7 @@ import type { } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" -import type { HarnessRunResult } from "@/agent-harness" +import type { HarnessRunResult, InspectImages } from "@/agent-harness" import type { ChatArtifactView, ChatCitationView, @@ -57,6 +57,7 @@ export type GenerateAnswer = (input: { sources: readonly Source[] excludedSourceIds: readonly string[] searchSources: SearchSources + inspectImages?: InspectImages }) => Promise export type AnswerQuestionInput = { @@ -69,6 +70,7 @@ export type AnswerQuestionInput = { generateAnswer: GenerateAnswer hardenChatAssetUrl?: HardenChatAssetUrl hardenMediaAssetUrls?: HardenMediaAssetUrls + inspectImages?: InspectImages messages: readonly ChatHistoryMessage[] } diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index d16f26e..8c239a9 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -1930,6 +1930,297 @@ describe("generateAgenticOutputManifest", () => { expect(JSON.stringify(capturedGenerateInput)).toContain("tax.pdf / deadline"); }); + it("lets the agent inspect retrieved image assets before finalizing cited image output", async () => { + process.env.AI_GATEWAY_API_KEY = "test_gateway_key"; + vi.spyOn(ToolLoopAgent.prototype, "generate").mockImplementation( + async function mockGenerate( + this: ToolLoopAgent, + ): ReturnType { + const tools = this.tools as unknown as Record< + string, + { execute: (input: unknown) => Promise } + >; + + await tools.declareIntent?.execute({ + task: "show_media", + dependsOnPreviousTurn: false, + retrievalNeeded: "yes", + targetModalities: ["image"], + constraints: { desiredCount: 1, maxCount: 1 }, + groundingPolicy: "must_use_sources", + }); + await tools.setContextPolicy?.execute({ + carryHistory: "none", + reason: "The current request is self-contained.", + activePriorTurnIds: [], + }); + await tools.retrieve?.execute({ + query: "identity card front image", + modalities: ["image"], + topK: 1, + purpose: "Find the ID card image to inspect.", + }); + await tools.inspectImage?.execute({ + refs: ["asset:r1:result:1"], + question: "What text is visible on the ID card?", + }); + await tools.finalize?.execute({ + text: "The inspected image appears to show the requested ID card.", + citations: [ + { + ref: "asset:r1:result:1", + label: "identity.pdf / images/id-front.png", + source: { + documentId: "doc_identity", + sourceFileName: "identity.pdf", + sectionPath: "images/id-front.png", + }, + }, + ], + artifacts: [ + { + type: "image", + ref: "asset:r1:result:1", + display: true, + reason: "Requested inspected ID card image.", + }, + ], + unresolved: [], + }); + + return { + text: "ignored", + } as Awaited>; + }, + ); + const searchSources = vi.fn().mockResolvedValue({ + results: [ + makeRetrievalResult({ + chunkType: "image", + assetUrl: "https://blob.example/images/id-front.png", + source: { + documentId: "doc_identity", + sourceFileName: "generated.pdf", + sectionPath: "images/id-front.png", + }, + }), + ], + evidenceText: "Identity image evidence.", + referencedChunks: [], + namespace: "notebook-workspace", + query: "identity card front image", + routerUsed: "workflow_single_step", + answerText: null, + stopReason: "answer_done", + failureReason: null, + }); + const inspectImages = vi.fn().mockResolvedValue({ + analysis: "The image contains a visible identity card number.", + inspected: [ + { + ref: "asset:r1:result:1", + label: "generated.pdf / images/id-front.png / image", + }, + ], + skipped: [], + }); + + const result = await generateAgenticOutputManifest({ + question: "Inspect and show the ID card image.", + messages: [], + sources: [ + makeSource({ title: "identity.pdf", knowhereDocumentId: "doc_identity" }), + ], + excludedSourceIds: [], + searchSources, + inspectImages, + }); + + expect(inspectImages).toHaveBeenCalledWith({ + question: "What text is visible on the ID card?", + assets: [ + { + ref: "asset:r1:result:1", + label: "generated.pdf / images/id-front.png / image", + assetUrl: "https://blob.example/images/id-front.png", + sourcePath: "images/id-front.png", + source: { + documentId: "doc_identity", + sourceFileName: "generated.pdf", + sectionPath: "images/id-front.png", + }, + }, + ], + }); + expect(result.manifest.citations.map((citation) => citation.ref)).toEqual([ + "asset:r1:result:1", + ]); + expect(result.manifest.artifacts).toEqual([ + { + type: "image", + ref: "asset:r1:result:1", + display: true, + reason: "Requested inspected ID card image.", + }, + ]); + expect(result.trace.toolCalls.map((call) => call.tool)).toContain( + "inspectImage", + ); + expect(result.trace.validationErrors).toEqual([]); + }); + + it("lets the agent inspect retrieved page assets before finalizing an OCR answer", async () => { + process.env.AI_GATEWAY_API_KEY = "test_gateway_key"; + vi.spyOn(ToolLoopAgent.prototype, "generate").mockImplementation( + async function mockGenerate( + this: ToolLoopAgent, + ): ReturnType { + const tools = this.tools as unknown as Record< + string, + { execute: (input: unknown) => Promise } + >; + + await tools.declareIntent?.execute({ + task: "answer", + dependsOnPreviousTurn: false, + retrievalNeeded: "yes", + targetModalities: ["text"], + constraints: { citationRequired: true, language: "zh-CN" }, + groundingPolicy: "must_use_sources", + }); + await tools.setContextPolicy?.execute({ + carryHistory: "none", + reason: "The current request is self-contained.", + activePriorTurnIds: [], + }); + await tools.retrieve?.execute({ + query: "进度计划 违约金 承包人", + modalities: ["text"], + topK: 6, + purpose: "Find the contract clause and page for the liquidated damages amount.", + }); + await tools.inspectImage?.execute({ + refs: ["asset:r1:referenced:1"], + question: + "OCR this clause and identify the liquidated damages amount for unauthorized schedule changes.", + }); + await tools.finalize?.execute({ + text: "承包人自行修改发包人审批的进度计划，应按每次 5000 元赔偿违约金。", + citations: [ + { + ref: "asset:r1:referenced:1", + label: "投标书 / （6）现场工期进度管理方面的违约责任", + source: { + documentId: "doc_contract", + sourceFileName: "投标书.pdf", + sectionPath: "Root / （6）现场工期进度管理方面的违约责任", + }, + }, + ], + artifacts: [], + unresolved: [], + }); + + return { + text: "ignored", + } as Awaited>; + }, + ); + const searchSources = vi.fn().mockResolvedValue({ + results: [], + evidenceText: "Root / （6）现场工期进度管理方面的违约责任", + referencedChunks: [ + { + chunkId: "chunk_page_8", + documentId: "doc_contract", + chunkType: "page", + sectionPath: "Root / （6）现场工期进度管理方面的违约责任", + filePath: null, + jobId: "job_contract", + metadata: { + pageNums: [8], + pageAssets: [ + { + pageNum: 8, + artifactRef: "page_citation_assets/page-8.png", + assetUrl: "https://blob.example/page-8.png", + contentType: "image/png", + }, + ], + }, + }, + ], + namespace: "notebook-workspace", + query: "进度计划 违约金 承包人", + routerUsed: "workflow_single_step", + answerText: null, + stopReason: "answer_done", + failureReason: null, + }); + const inspectImages = vi.fn().mockResolvedValue({ + analysis: "The clause states 5000 yuan per occurrence.", + inspected: [ + { + ref: "asset:r1:referenced:1", + label: + "Root / （6）现场工期进度管理方面的违约责任 / page_citation_assets/page-8.png / page", + }, + ], + skipped: [], + }); + + const result = await generateAgenticOutputManifest({ + question: "承包人自行修改发包人审批的进度时需要赔偿多少违约金？", + messages: [], + sources: [ + makeSource({ + title: "投标书.pdf", + knowhereDocumentId: "doc_contract", + }), + ], + excludedSourceIds: [], + searchSources, + inspectImages, + }); + + expect(searchSources).toHaveBeenCalledWith({ + query: "进度计划 违约金 承包人", + targetContent: "text", + purpose: "Find the contract clause and page for the liquidated damages amount.", + topK: 6, + signalPaths: undefined, + filterMode: undefined, + threshold: undefined, + }); + expect(inspectImages).toHaveBeenCalledWith({ + question: + "OCR this clause and identify the liquidated damages amount for unauthorized schedule changes.", + assets: [ + { + ref: "asset:r1:referenced:1", + label: + "Root / （6）现场工期进度管理方面的违约责任 / page_citation_assets/page-8.png / page", + assetUrl: "https://blob.example/page-8.png", + sourcePath: "page_citation_assets/page-8.png", + revisionKey: "job_contract", + source: { + documentId: "doc_contract", + sourceFileName: null, + sectionPath: "Root / （6）现场工期进度管理方面的违约责任", + }, + }, + ], + }); + expect(result.manifest.text).toContain("5000 元"); + expect(result.manifest.citations.map((citation) => citation.ref)).toEqual([ + "asset:r1:referenced:1", + ]); + expect(result.trace.toolCalls.map((call) => call.tool)).toContain( + "inspectImage", + ); + expect(result.trace.validationErrors).toEqual([]); + }); + it("self-corrects an over-budget manifest via a validation-feedback revision", async () => { process.env.AI_GATEWAY_API_KEY = "test_gateway_key"; let generateCallCount = 0; diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index 21e6ef8..fa07579 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -214,6 +214,7 @@ export const answerQuestionWithRetrieval = ( sources: input.sources, excludedSourceIds: input.excludedSourceIds, searchSources, + ...(input.inspectImages ? { inspectImages: input.inspectImages } : {}), }), ) @@ -414,12 +415,12 @@ function toChatArtifactView(input: { ref: input.artifact.ref, display: input.artifact.display, reason: input.artifact.reason, - assetUrl: input.asset.assetUrl, + ...(input.asset.assetUrl ? { assetUrl: input.asset.assetUrl } : {}), label: input.asset.label, citation: { chunkType: input.asset.type, score: null, - assetUrl: input.asset.assetUrl, + ...(input.asset.assetUrl ? { assetUrl: input.asset.assetUrl } : {}), source, }, } diff --git a/src/domains/chat/prompt.ts b/src/domains/chat/prompt.ts index b13831e..efee008 100644 --- a/src/domains/chat/prompt.ts +++ b/src/domains/chat/prompt.ts @@ -10,6 +10,7 @@ import { type AgentTurnInput, type HarnessRetrievalRequest, type HarnessRunResult, + type InspectImages, type TargetModality, } from "@/agent-harness" import type { @@ -29,6 +30,7 @@ type GenerateAgenticOutputManifestInput = { sources: readonly Source[] excludedSourceIds: readonly string[] searchSources: SearchSources + inspectImages?: InspectImages } export const generateAgenticOutputManifestEffect = ( @@ -61,6 +63,7 @@ export const generateAgenticOutputManifestEffect = ( query: (request) => input.searchSources(toAgenticRetrievalQuery(request)), }, + ...(input.inspectImages ? { inspectImages: input.inspectImages } : {}), }), ) diff --git a/src/domains/chat/route-answer.ts b/src/domains/chat/route-answer.ts index 03c26ca..2f8ae6b 100644 --- a/src/domains/chat/route-answer.ts +++ b/src/domains/chat/route-answer.ts @@ -1,9 +1,17 @@ import { Cause, Effect, Either, Option } from "effect" +import { generateText } from "ai" import { generateAgenticOutputManifest, parseChatRequestBody, } from "@/domains/chat" +import type { + ImageInspectionAsset, + ImageInspectionRequest, + ImageInspectionResponse, + ImageInspectionSkippedAsset, + InspectImages, +} from "@/agent-harness" import { hardenChatMediaAssetUrls } from "@/domains/chat/media-asset-hardening" import { handleChatTurn, @@ -16,6 +24,7 @@ import { BlobParsedDocumentStorage } from "@/domains/sources/parsed-document-blo import { sourceWorkflowRuntime } from "@/domains/sources/workflow-runtime" import { notebookRequestContext } from "@/domains/workspace/request-context" import type { Source } from "@/infrastructure/db/schema" +import { CHAT_MODEL } from "@/lib/ai" import type { HardenChatAssetUrl } from "./media-assets" import { isAuthError } from "@/integrations/dashboard/api-key-service" import { summarizeUnknownError } from "@/lib/format-log-value" @@ -24,6 +33,23 @@ import { routeResult, type RouteResult } from "@/lib/route-result" type RouteResponse = RouteResult +const VISION_MODEL = process.env.VISION_MODEL ?? CHAT_MODEL +const IMAGE_INSPECTION_URL_REPLACEMENT = "[image URL hidden]" +const SUPPORTED_IMAGE_CONTENT_TYPES: Readonly> = { + ".jpeg": "image/jpeg", + ".jpg": "image/jpeg", + ".png": "image/png", + ".webp": "image/webp", +} as const + +type HardenImageInspectionAsset = (input: { + readonly documentId: string + readonly revisionKey: string + readonly sourcePath: string + readonly assetUrl?: string | null + readonly contentType?: string | null +}) => Promise + type MessageBody = { readonly message: string } @@ -86,6 +112,24 @@ const answerChatEffect = (input: AnswerChatInput) => assetUrl, contentType, }) + const hardenImageInspectionAsset: HardenImageInspectionAsset = (asset) => + hardenChatAssetByDocument({ + workspaceId: workspace.id, + parsedStorage, + documentId: asset.documentId, + revisionKey: asset.revisionKey, + sourcePath: asset.sourcePath, + assetUrl: asset.assetUrl, + contentType: asset.contentType, + }) + const inspectImages: InspectImages = (request) => + inspectChatImages({ + workspaceId: workspace.id, + sources, + hardenChatAssetUrl, + hardenImageInspectionAsset, + request, + }) const result: Either.Either = yield* Effect.tryPromise(() => @@ -106,6 +150,7 @@ const answerChatEffect = (input: AnswerChatInput) => artifacts, hardenChatAssetUrl, }), + inspectImages, repository: chatTurnPersistence.createRepository(), }), ).pipe( @@ -160,6 +205,355 @@ export const chatAnswerRouteService: ChatAnswerRouteService = { answerChat, } +async function inspectChatImages(input: { + readonly workspaceId: string + readonly sources: readonly Source[] + readonly hardenChatAssetUrl: HardenChatAssetUrl + readonly hardenImageInspectionAsset: HardenImageInspectionAsset + readonly request: ImageInspectionRequest +}): Promise { + const preparedAssets: PreparedImageInspectionAsset[] = [] + const skippedAssets: ImageInspectionSkippedAsset[] = [] + + for (const asset of input.request.assets) { + const preparedAsset = await prepareImageInspectionAsset({ + asset, + sources: input.sources, + hardenChatAssetUrl: input.hardenChatAssetUrl, + hardenImageInspectionAsset: input.hardenImageInspectionAsset, + }) + if (preparedAsset.ok) { + preparedAssets.push(preparedAsset.asset) + } else { + skippedAssets.push({ + ref: asset.ref, + reason: preparedAsset.reason, + }) + } + } + + if (preparedAssets.length === 0) { + return { + analysis: "", + inspected: [], + skipped: skippedAssets, + } + } + + logger.info("chat: image inspection request", { + workspaceId: input.workspaceId, + model: VISION_MODEL, + requestedCount: input.request.assets.length, + inspectedCount: preparedAssets.length, + skippedCount: skippedAssets.length, + refs: preparedAssets.map((asset) => asset.ref), + }) + + const response = await generateImageInspectionText({ + workspaceId: input.workspaceId, + question: input.request.question, + assets: preparedAssets, + }) + const analysis = removeImageInspectionUrls({ + text: response.text.trim(), + assets: preparedAssets, + }) + + logger.info("chat: image inspection response", { + workspaceId: input.workspaceId, + model: VISION_MODEL, + inspectedCount: preparedAssets.length, + skippedCount: skippedAssets.length, + analysisLength: analysis.length, + }) + + return { + analysis, + inspected: preparedAssets.map((asset) => ({ + ref: asset.ref, + label: asset.label, + })), + skipped: skippedAssets, + } +} + +type PreparedImageInspectionAsset = { + readonly ref: string + readonly label: string + readonly url: URL + readonly body: Uint8Array + readonly contentType: string + readonly fileName?: string +} + +type PrepareImageInspectionAssetResult = + | { + readonly ok: true + readonly asset: PreparedImageInspectionAsset + } + | { + readonly ok: false + readonly reason: string + } + +async function prepareImageInspectionAsset(input: { + readonly asset: ImageInspectionAsset + readonly sources: readonly Source[] + readonly hardenChatAssetUrl: HardenChatAssetUrl + readonly hardenImageInspectionAsset: HardenImageInspectionAsset +}): Promise { + const sourcePath = resolveImageInspectionSourcePath(input.asset) + if (!sourcePath) { + return { + ok: false, + reason: "The image asset path could not be resolved.", + } + } + + const contentType = getSupportedImageContentType(sourcePath) + if (!contentType) { + return { + ok: false, + reason: "Only PNG, JPEG, and WebP image assets can be inspected.", + } + } + + const source = resolveImageInspectionSource(input.asset, input.sources) + const durableUrlFromSource = source + ? await input.hardenChatAssetUrl({ + source, + sourcePath, + assetUrl: input.asset.assetUrl, + contentType, + }) + : null + const durableUrl = + durableUrlFromSource ?? + (await hardenImageInspectionAssetByRef({ + asset: input.asset, + sourcePath, + contentType, + hardenImageInspectionAsset: input.hardenImageInspectionAsset, + })) + if (!durableUrl) { + return { + ok: false, + reason: source + ? "The image asset was unavailable in Notebook storage." + : "No ready Notebook source matched the image asset.", + } + } + + const url = parseAbsoluteHttpUrl(durableUrl) + if (!url) { + return { + ok: false, + reason: "Notebook storage did not return a valid image URL.", + } + } + const image = await fetchPreparedInspectionImage({ + url, + contentType, + }) + if (!image.ok) { + return { + ok: false, + reason: image.reason, + } + } + + return { + ok: true, + asset: { + ref: input.asset.ref, + label: input.asset.label, + url, + body: image.body, + contentType, + fileName: getFileNameFromPath(sourcePath), + }, + } +} + +async function generateImageInspectionText(input: { + readonly workspaceId: string + readonly question: string + readonly assets: readonly PreparedImageInspectionAsset[] +}): Promise>> { + try { + return await generateText({ + model: VISION_MODEL, + messages: [ + { + role: "user", + content: [ + { + type: "text", + text: buildImageInspectionPrompt({ + question: input.question, + assets: input.assets, + }), + }, + ...input.assets.map((asset) => ({ + type: "image" as const, + image: asset.body, + mediaType: asset.contentType, + })), + ], + }, + ], + experimental_include: { + requestBody: false, + responseBody: false, + }, + }) + } catch (error) { + logger.warn("chat: image inspection model call failed", { + workspaceId: input.workspaceId, + model: VISION_MODEL, + inspectedCount: input.assets.length, + error: summarizeUnknownError(error), + }) + throw error + } +} + +async function fetchPreparedInspectionImage(input: { + readonly url: URL + readonly contentType: string +}): Promise< + | { + readonly ok: true + readonly body: Uint8Array + } + | { + readonly ok: false + readonly reason: string + } +> { + try { + const image = await fetchChatAsset({ + assetUrl: input.url.toString(), + fallbackContentType: input.contentType, + }) + return { + ok: true, + body: image.body, + } + } catch { + return { + ok: false, + reason: "The Notebook image asset could not be read for inspection.", + } + } +} + +async function hardenImageInspectionAssetByRef(input: { + readonly asset: ImageInspectionAsset + readonly sourcePath: string + readonly contentType: string + readonly hardenImageInspectionAsset: HardenImageInspectionAsset +}): Promise { + const documentId = getTrimmedString(input.asset.source.documentId) + const revisionKey = getTrimmedString(input.asset.revisionKey) + if (!documentId || !revisionKey) return null + + return input.hardenImageInspectionAsset({ + documentId, + revisionKey, + sourcePath: input.sourcePath, + assetUrl: input.asset.assetUrl, + contentType: input.contentType, + }) +} + +function resolveImageInspectionSource( + asset: ImageInspectionAsset, + sources: readonly Source[], +): Source | undefined { + const documentId = getTrimmedString(asset.source.documentId) + if (!documentId) return undefined + + return sources.find( + (source) => + source.status === "ready" && source.knowhereDocumentId === documentId, + ) +} + +function resolveImageInspectionSourcePath( + asset: ImageInspectionAsset, +): string | null { + const candidates = [ + asset.sourcePath, + asset.source.sectionPath, + getAssetUrlPathname(asset.assetUrl), + ] + for (const candidate of candidates) { + const sourcePath = getSupportedImageAssetPath(candidate) + if (sourcePath) return sourcePath + } + + return null +} + +function getSupportedImageAssetPath( + value: string | null | undefined, +): string | null { + const normalizedText = normalizeSourcePathCandidate(value) + if (!normalizedText) return null + + const match = + /(?:^|\/)((?:images|pages|page_citation_assets)\/[^?#]+)(?:[?#]|$)?/i.exec( + normalizedText, + ) + const matchedPath = match?.[1] + return matchedPath ? matchedPath.trim() : null +} + +function getSupportedImageContentType(sourcePath: string): string | null { + const lowerPath = sourcePath.toLowerCase() + const extension = Object.keys(SUPPORTED_IMAGE_CONTENT_TYPES).find((candidate) => + lowerPath.endsWith(candidate), + ) + return extension ? SUPPORTED_IMAGE_CONTENT_TYPES[extension] ?? null : null +} + +function buildImageInspectionPrompt(input: { + readonly question: string + readonly assets: readonly PreparedImageInspectionAsset[] +}): string { + return [ + "Inspect the attached Notebook image assets selected from retrieved Knowhere evidence.", + "Answer the inspection question using concise visual observations only.", + "Use the provided refs and labels to identify images. Do not include image URLs.", + "If OCR text is unclear, say it is unclear instead of guessing.", + "Do not create citations. The calling agent will cite the original retrieved asset refs.", + "", + "Inspection question:", + input.question, + "", + "Images:", + ...input.assets.map((asset) => `- ref=${asset.ref} label=${asset.label}`), + ].join("\n") +} + +function removeImageInspectionUrls(input: { + readonly text: string + readonly assets: readonly PreparedImageInspectionAsset[] +}): string { + const knownUrls = input.assets.map((asset) => asset.url.toString()) + const withoutKnownUrls = knownUrls.reduce( + (text, assetUrl) => + text.replaceAll(assetUrl, IMAGE_INSPECTION_URL_REPLACEMENT), + input.text, + ) + + return withoutKnownUrls + .replace(/https?:\/\/[^\s)\]}>"']+/g, IMAGE_INSPECTION_URL_REPLACEMENT) + .replace(/[ \t]{2,}/g, " ") + .trim() +} + async function hardenSingleChatAsset(input: { readonly workspaceId: string readonly parsedStorage: BlobParsedDocumentStorage @@ -176,10 +570,32 @@ async function hardenSingleChatAsset(input: { return null } + return hardenChatAssetByDocument({ + workspaceId: input.workspaceId, + parsedStorage: input.parsedStorage, + documentId: input.source.knowhereDocumentId, + revisionKey: input.source.knowhereJobId, + sourcePath: input.sourcePath, + assetUrl: input.assetUrl, + contentType: input.contentType, + sourceId: input.source.id, + }) +} + +async function hardenChatAssetByDocument(input: { + readonly workspaceId: string + readonly parsedStorage: BlobParsedDocumentStorage + readonly documentId: string + readonly revisionKey: string + readonly sourcePath: string + readonly assetUrl?: string | null + readonly contentType?: string | null + readonly sourceId?: string | null +}): Promise { try { const existingUrl = await input.parsedStorage.getAssetUrl({ - documentId: input.source.knowhereDocumentId, - revisionKey: input.source.knowhereJobId, + documentId: input.documentId, + revisionKey: input.revisionKey, sourcePath: input.sourcePath, }) if (existingUrl) return existingUrl @@ -194,8 +610,8 @@ async function hardenSingleChatAsset(input: { inferContentTypeFromPath(input.sourcePath), }) const writtenAsset = await input.parsedStorage.writeAsset({ - documentId: input.source.knowhereDocumentId, - revisionKey: input.source.knowhereJobId, + documentId: input.documentId, + revisionKey: input.revisionKey, sourcePath: input.sourcePath, body: fetchedAsset.body, contentType: fetchedAsset.contentType, @@ -204,16 +620,16 @@ async function hardenSingleChatAsset(input: { return ( writtenAsset.url ?? (await input.parsedStorage.getAssetUrl({ - documentId: input.source.knowhereDocumentId, - revisionKey: input.source.knowhereJobId, + documentId: input.documentId, + revisionKey: input.revisionKey, sourcePath: input.sourcePath, })) ) } catch (error) { logger.warn("chat: single parsed asset hardening failed", { workspaceId: input.workspaceId, - sourceId: input.source.id, - documentId: input.source.knowhereDocumentId, + sourceId: input.sourceId ?? null, + documentId: input.documentId, sourcePath: input.sourcePath, error: summarizeUnknownError(error), }) @@ -258,6 +674,54 @@ function inferContentTypeFromPath(sourcePath: string): string { return "application/octet-stream" } +function getAssetUrlPathname(assetUrl: string | null | undefined): string | null { + const normalizedAssetUrl = getTrimmedString(assetUrl) + if (!normalizedAssetUrl) return null + + try { + return new URL(normalizedAssetUrl).pathname + } catch { + return normalizedAssetUrl.split("?")[0] ?? normalizedAssetUrl + } +} + +function normalizeSourcePathCandidate( + value: string | null | undefined, +): string | null { + const trimmedValue = getTrimmedString(value) + if (!trimmedValue) return null + + const normalized = decodeUrlText(trimmedValue) + .replaceAll("\\", "/") + .replace(/\s*\/\s*/g, "/") + .replace(/\s+/g, " ") + .trim() + + return normalized.length > 0 ? normalized : null +} + +function decodeUrlText(value: string): string { + try { + return decodeURIComponent(value) + } catch { + return value + } +} + +function parseAbsoluteHttpUrl(assetUrl: string): URL | null { + try { + const url = new URL(assetUrl) + return url.protocol === "http:" || url.protocol === "https:" ? url : null + } catch { + return null + } +} + +function getFileNameFromPath(sourcePath: string): string | undefined { + const fileName = sourcePath.replaceAll("\\", "/").split("/").pop()?.trim() + return fileName && fileName.length > 0 ? fileName : undefined +} + function triggerBackgroundReconciliationForParsingSources(input: { readonly workspaceId: string readonly sources: readonly Source[] diff --git a/src/domains/chat/route-service.test.ts b/src/domains/chat/route-service.test.ts index 6365edd..c771f12 100644 --- a/src/domains/chat/route-service.test.ts +++ b/src/domains/chat/route-service.test.ts @@ -9,6 +9,7 @@ const mocks = vi.hoisted(() => ({ ensureDefaultChatThread: vi.fn(), findChatThreadInWorkspace: vi.fn(), generateAgenticOutputManifest: vi.fn(), + generateText: vi.fn(), getAuthenticated: vi.fn(), getAuthenticatedWithClient: vi.fn(), handleChatTurn: vi.fn(), @@ -24,6 +25,14 @@ const mocks = vi.hoisted(() => ({ startBackgroundReconciliation: vi.fn(), })) +vi.mock("ai", async (importOriginal) => { + const original = await importOriginal() + return { + ...original, + generateText: mocks.generateText, + } +}) + vi.mock("@/domains/chat", async (importOriginal) => { const original = await importOriginal() return { @@ -90,6 +99,7 @@ describe("chat route services", () => { vi.clearAllMocks() mocks.parsedStorageGetAssetUrl.mockResolvedValue(null) mocks.parsedStorageWriteAsset.mockResolvedValue({ url: null }) + mocks.generateText.mockResolvedValue({ text: "The image shows a chart." }) vi.stubGlobal( "fetch", vi.fn(async () => new Response(new Uint8Array([1, 2, 3]), { @@ -228,6 +238,444 @@ describe("chat route services", () => { }) }) + it("hardens image inspection assets before sending Notebook URLs to Gemini", async () => { + const workspace = makeWorkspace() + const client = { retrieval: { query: vi.fn() } } + const readySource = makeSource({ + status: "ready", + knowhereDocumentId: "doc_identity", + knowhereJobId: "job_1", + }) + const rawUrl = + "https://knowhere-storage.example/results/job_1/images/id-front.png?AWSAccessKeyId=test" + const durableUrl = + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/job_1/assets/images/id-front.png" + mocks.parsedStorageWriteAsset.mockResolvedValue({ url: durableUrl }) + mocks.generateText.mockResolvedValue({ + text: `The card number is visible. ${durableUrl} ${rawUrl}`, + }) + mocks.getAuthenticatedWithClient.mockResolvedValue({ + user: { id: "user_1" }, + workspace, + apiKey: "jwt_123", + client, + }) + mocks.listSourcesForWorkspace.mockResolvedValue([readySource]) + mocks.handleChatTurn.mockImplementation( + async (input: { + readonly inspectImages?: (request: { + readonly question: string + readonly assets: readonly { + readonly ref: string + readonly label: string + readonly assetUrl: string + readonly source: { + readonly documentId?: string | null + readonly sourceFileName?: string | null + readonly sectionPath?: string | null + } + }[] + }) => Promise<{ + readonly analysis: string + readonly inspected: readonly { + readonly ref: string + readonly label: string + }[] + readonly skipped: readonly { + readonly ref: string + readonly reason: string + }[] + }> + }) => { + const inspection = await input.inspectImages?.({ + question: "Read the ID card text.", + assets: [ + { + ref: "asset:r1:result:1", + label: "identity.pdf / images/id-front.png / image", + assetUrl: rawUrl, + source: { + documentId: "doc_identity", + sourceFileName: "identity.pdf", + sectionPath: "images/id-front.png", + }, + }, + ], + }) + expect(inspection).toEqual({ + analysis: "The card number is visible. [image URL hidden] [image URL hidden]", + inspected: [ + { + ref: "asset:r1:result:1", + label: "identity.pdf / images/id-front.png / image", + }, + ], + skipped: [], + }) + return Either.right({ + threadId: "thread_1", + messages: [ + { id: "message_user", role: "user", content: "Inspect it" }, + { id: "message_assistant", role: "assistant", content: "Answer" }, + ], + }) + }, + ) + + const result = await chatAnswerRouteService.answerChat({ + body: { message: "Inspect the ID card image" }, + }) + + expect(result.status).toBe(200) + expect(fetch).toHaveBeenCalledWith(rawUrl) + expect(fetch).toHaveBeenCalledWith(durableUrl) + expect(mocks.parsedStorageWriteAsset).toHaveBeenCalledWith({ + documentId: "doc_identity", + revisionKey: "job_1", + sourcePath: "images/id-front.png", + body: new Uint8Array([1, 2, 3]), + contentType: "image/png", + }) + const generateInput = mocks.generateText.mock.calls[0]?.[0] + expect(generateInput).toMatchObject({ + model: "google/gemini-3-flash", + experimental_include: { + requestBody: false, + responseBody: false, + }, + }) + expect(JSON.stringify(generateInput)).not.toContain( + "knowhere-storage.example", + ) + const content = generateInput.messages[0].content + const imagePart = content.find( + (part: { readonly type: string }) => part.type === "image", + ) + expect(imagePart.image).toEqual(new Uint8Array([1, 2, 3])) + expect(imagePart.mediaType).toBe("image/png") + }) + + it("uses image inspection source paths for retrieved page assets", async () => { + const workspace = makeWorkspace() + const client = { retrieval: { query: vi.fn() } } + const readySource = makeSource({ + status: "ready", + knowhereDocumentId: "doc_contract", + knowhereJobId: "job_1", + }) + const durableUrl = + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_contract/job_1/assets/page_citation_assets/page-8.png" + mocks.parsedStorageGetAssetUrl.mockResolvedValue(durableUrl) + mocks.generateText.mockResolvedValue({ + text: "The page states 5000 yuan per occurrence.", + }) + mocks.getAuthenticatedWithClient.mockResolvedValue({ + user: { id: "user_1" }, + workspace, + apiKey: "jwt_123", + client, + }) + mocks.listSourcesForWorkspace.mockResolvedValue([readySource]) + mocks.handleChatTurn.mockImplementation( + async (input: { + readonly inspectImages?: (request: { + readonly question: string + readonly assets: readonly { + readonly ref: string + readonly label: string + readonly assetUrl?: string | null + readonly sourcePath?: string | null + readonly source: { + readonly documentId?: string | null + readonly sourceFileName?: string | null + readonly sectionPath?: string | null + } + }[] + }) => Promise<{ + readonly analysis: string + readonly inspected: readonly { + readonly ref: string + readonly label: string + }[] + readonly skipped: readonly { + readonly ref: string + readonly reason: string + }[] + }> + }) => { + const inspection = await input.inspectImages?.({ + question: "Read the damages amount.", + assets: [ + { + ref: "asset:r1:referenced:1", + label: + "Root / （6）现场工期进度管理方面的违约责任 / page_citation_assets/page-8.png / page", + sourcePath: "page_citation_assets/page-8.png", + source: { + documentId: "doc_contract", + sourceFileName: "投标书.pdf", + sectionPath: "Root / （6）现场工期进度管理方面的违约责任", + }, + }, + ], + }) + expect(inspection).toEqual({ + analysis: "The page states 5000 yuan per occurrence.", + inspected: [ + { + ref: "asset:r1:referenced:1", + label: + "Root / （6）现场工期进度管理方面的违约责任 / page_citation_assets/page-8.png / page", + }, + ], + skipped: [], + }) + return Either.right({ + threadId: "thread_1", + messages: [ + { id: "message_user", role: "user", content: "Inspect it" }, + { id: "message_assistant", role: "assistant", content: "Answer" }, + ], + }) + }, + ) + + const result = await chatAnswerRouteService.answerChat({ + body: { message: "Inspect the clause page" }, + }) + + expect(result.status).toBe(200) + expect(mocks.parsedStorageGetAssetUrl).toHaveBeenCalledWith({ + documentId: "doc_contract", + revisionKey: "job_1", + sourcePath: "page_citation_assets/page-8.png", + }) + expect(fetch).toHaveBeenCalledWith(durableUrl) + expect(mocks.parsedStorageWriteAsset).not.toHaveBeenCalled() + const generateInput = mocks.generateText.mock.calls[0]?.[0] + const content = generateInput.messages[0].content + const imagePart = content.find( + (part: { readonly type: string }) => part.type === "image", + ) + expect(imagePart.image).toEqual(new Uint8Array([1, 2, 3])) + expect(imagePart.mediaType).toBe("image/png") + }) + + it("hardens remote page inspection assets by document revision when no local source matches", async () => { + const workspace = makeWorkspace() + const client = { retrieval: { query: vi.fn() } } + const rawUrl = + "https://knowhere-storage.example/results/job_remote/page_citation_assets/page-8.png?AWSAccessKeyId=test" + const durableUrl = + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_remote/job_remote/assets/page_citation_assets/page-8.png" + mocks.parsedStorageWriteAsset.mockResolvedValue({ url: durableUrl }) + mocks.generateText.mockResolvedValue({ + text: "The page states 5000 yuan per occurrence.", + }) + mocks.getAuthenticatedWithClient.mockResolvedValue({ + user: { id: "user_1" }, + workspace, + apiKey: "jwt_123", + client, + }) + mocks.listSourcesForWorkspace.mockResolvedValue([]) + mocks.handleChatTurn.mockImplementation( + async (input: { + readonly inspectImages?: (request: { + readonly question: string + readonly assets: readonly { + readonly ref: string + readonly label: string + readonly assetUrl?: string | null + readonly sourcePath?: string | null + readonly revisionKey?: string | null + readonly source: { + readonly documentId?: string | null + readonly sourceFileName?: string | null + readonly sectionPath?: string | null + } + }[] + }) => Promise<{ + readonly analysis: string + readonly inspected: readonly { + readonly ref: string + readonly label: string + }[] + readonly skipped: readonly { + readonly ref: string + readonly reason: string + }[] + }> + }) => { + const inspection = await input.inspectImages?.({ + question: "Read the damages amount.", + assets: [ + { + ref: "asset:r1:referenced:5", + label: + "Root / （6）现场工期进度管理方面的违约责任 / page_citation_assets/page-8.png / page", + assetUrl: rawUrl, + sourcePath: "page_citation_assets/page-8.png", + revisionKey: "job_remote", + source: { + documentId: "doc_remote", + sourceFileName: null, + sectionPath: "Root / （6）现场工期进度管理方面的违约责任", + }, + }, + ], + }) + expect(inspection).toEqual({ + analysis: "The page states 5000 yuan per occurrence.", + inspected: [ + { + ref: "asset:r1:referenced:5", + label: + "Root / （6）现场工期进度管理方面的违约责任 / page_citation_assets/page-8.png / page", + }, + ], + skipped: [], + }) + return Either.right({ + threadId: "thread_1", + messages: [ + { id: "message_user", role: "user", content: "Inspect it" }, + { id: "message_assistant", role: "assistant", content: "Answer" }, + ], + }) + }, + ) + + const result = await chatAnswerRouteService.answerChat({ + body: { message: "Inspect the remote clause page" }, + }) + + expect(result.status).toBe(200) + expect(mocks.parsedStorageGetAssetUrl).toHaveBeenCalledWith({ + documentId: "doc_remote", + revisionKey: "job_remote", + sourcePath: "page_citation_assets/page-8.png", + }) + expect(fetch).toHaveBeenCalledWith(rawUrl) + expect(fetch).toHaveBeenCalledWith(durableUrl) + expect(mocks.parsedStorageWriteAsset).toHaveBeenCalledWith({ + documentId: "doc_remote", + revisionKey: "job_remote", + sourcePath: "page_citation_assets/page-8.png", + body: new Uint8Array([1, 2, 3]), + contentType: "image/png", + }) + const generateInput = mocks.generateText.mock.calls[0]?.[0] + expect(JSON.stringify(generateInput)).not.toContain( + "knowhere-storage.example", + ) + const content = generateInput.messages[0].content + const imagePart = content.find( + (part: { readonly type: string }) => part.type === "image", + ) + expect(imagePart.image).toEqual(new Uint8Array([1, 2, 3])) + expect(imagePart.mediaType).toBe("image/png") + }) + + it("skips unsupported or unavailable image inspection assets without failing chat", async () => { + const workspace = makeWorkspace() + const client = { retrieval: { query: vi.fn() } } + const readySource = makeSource({ + status: "ready", + knowhereDocumentId: "doc_identity", + knowhereJobId: "job_1", + }) + mocks.getAuthenticatedWithClient.mockResolvedValue({ + user: { id: "user_1" }, + workspace, + apiKey: "jwt_123", + client, + }) + mocks.listSourcesForWorkspace.mockResolvedValue([readySource]) + mocks.handleChatTurn.mockImplementation( + async (input: { + readonly inspectImages?: (request: { + readonly question: string + readonly assets: readonly { + readonly ref: string + readonly label: string + readonly assetUrl: string + readonly source: { + readonly documentId?: string | null + readonly sourceFileName?: string | null + readonly sectionPath?: string | null + } + }[] + }) => Promise<{ + readonly analysis: string + readonly inspected: readonly { + readonly ref: string + readonly label: string + }[] + readonly skipped: readonly { + readonly ref: string + readonly reason: string + }[] + }> + }) => { + const inspection = await input.inspectImages?.({ + question: "Inspect these.", + assets: [ + { + ref: "asset:r1:result:1", + label: "identity.pdf / images/animated.gif / image", + assetUrl: "https://knowhere-storage.example/results/job_1/images/animated.gif", + source: { + documentId: "doc_identity", + sourceFileName: "identity.pdf", + sectionPath: "images/animated.gif", + }, + }, + { + ref: "asset:r1:result:2", + label: "missing.pdf / images/missing.png / image", + assetUrl: "https://knowhere-storage.example/results/job_2/images/missing.png", + source: { + documentId: "doc_missing", + sourceFileName: "missing.pdf", + sectionPath: "images/missing.png", + }, + }, + ], + }) + expect(inspection).toEqual({ + analysis: "", + inspected: [], + skipped: [ + { + ref: "asset:r1:result:1", + reason: "Only PNG, JPEG, and WebP image assets can be inspected.", + }, + { + ref: "asset:r1:result:2", + reason: "No ready Notebook source matched the image asset.", + }, + ], + }) + return Either.right({ + threadId: "thread_1", + messages: [ + { id: "message_user", role: "user", content: "Inspect it" }, + { id: "message_assistant", role: "assistant", content: "Answer" }, + ], + }) + }, + ) + + const result = await chatAnswerRouteService.answerChat({ + body: { message: "Inspect images" }, + }) + + expect(result.status).toBe(200) + expect(mocks.generateText).not.toHaveBeenCalled() + expect(fetch).not.toHaveBeenCalled() + }) + it("triggers background reconciliation for parsing sources without blocking chat", async () => { const workspace = makeWorkspace() const client = { retrieval: { query: vi.fn() } } diff --git a/src/domains/chat/service.ts b/src/domains/chat/service.ts index 8b67e7c..bfa7e42 100644 --- a/src/domains/chat/service.ts +++ b/src/domains/chat/service.ts @@ -69,6 +69,7 @@ type ChatTurnInput = { generateAnswer: GenerateAnswer hardenChatAssetUrl?: AnswerQuestionInput["hardenChatAssetUrl"] hardenMediaAssetUrls?: AnswerQuestionInput["hardenMediaAssetUrls"] + inspectImages?: AnswerQuestionInput["inspectImages"] repository: ChatRepository } @@ -128,6 +129,7 @@ export const handleChatTurnEffect = (input: ChatTurnInput) => generateAnswer: input.generateAnswer, hardenChatAssetUrl: input.hardenChatAssetUrl, hardenMediaAssetUrls: input.hardenMediaAssetUrls, + inspectImages: input.inspectImages, messages: chatHistoryMessages, }).pipe(Effect.catchAllCause(Effect.die)) From 88d6b47be876c02d8c39efe2bf284a08e7c8a60a Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 6 Jul 2026 21:43:26 +0800 Subject: [PATCH 19/51] Keep parsed storage sync non-blocking --- .../sources/[sourceId]/chunks/route.test.ts | 2 + src/domains/chunks/index.ts | 4 +- src/domains/chunks/read.test.ts | 11 +- src/domains/chunks/read.ts | 6 +- .../parsed-sync-route-workflow.test.ts | 86 +++----- .../sources/parsed-sync-route-workflow.ts | 67 +++--- .../source-reconcile-route-workflow.test.ts | 152 +++----------- .../source-reconcile-route-workflow.ts | 195 +++++++----------- src/proxy.test.ts | 10 + src/proxy.ts | 1 + 10 files changed, 185 insertions(+), 349 deletions(-) diff --git a/src/app/api/sources/[sourceId]/chunks/route.test.ts b/src/app/api/sources/[sourceId]/chunks/route.test.ts index 2f0851c..85fb1bd 100644 --- a/src/app/api/sources/[sourceId]/chunks/route.test.ts +++ b/src/app/api/sources/[sourceId]/chunks/route.test.ts @@ -498,6 +498,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { revisionKey: "job_1", page: 1, pageSize: 1, + assetUrlPolicy: "durable", }) }) @@ -659,6 +660,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { revisionKey: "job_result_1", page: 1, pageSize: 1, + assetUrlPolicy: "durable", }) }) }) diff --git a/src/domains/chunks/index.ts b/src/domains/chunks/index.ts index 56c2644..eca1830 100644 --- a/src/domains/chunks/index.ts +++ b/src/domains/chunks/index.ts @@ -181,8 +181,8 @@ export function toParsedChunkView( /** * Map an SDK `KnowledgeReadChunk` (from `knowledge.readChunks`) to the view - * model. Display reads do not request durable asset URL hardening; chat - * hardens the specific assets it returns separately. + * model. Display reads request durable asset URLs from the SDK; chat still + * hardens only the specific assets it returns separately. */ export function toParsedChunkViewFromReadChunk( chunk: KnowledgeReadChunk, diff --git a/src/domains/chunks/read.test.ts b/src/domains/chunks/read.test.ts index b6a3979..e3c1000 100644 --- a/src/domains/chunks/read.test.ts +++ b/src/domains/chunks/read.test.ts @@ -20,7 +20,7 @@ function makeReadChunk(overrides: Record = {}) { } describe("readSourceChunkPage", () => { - it("reads a page without durable asset hardening and maps chunks to the view model", async () => { + it("reads a page with durable asset URLs and maps chunks to the view model", async () => { const readChunks = vi.fn(async () => ({ document: { localDocumentId: "doc_1" }, chunks: [ @@ -48,6 +48,7 @@ describe("readSourceChunkPage", () => { revisionKey: "rev_1", page: 2, pageSize: 50, + assetUrlPolicy: "durable", }) expect(result.pagination).toEqual({ page: 2, @@ -84,6 +85,7 @@ describe("readSourceChunkPage", () => { documentId: "doc_1", page: 1, pageSize: 50, + assetUrlPolicy: "durable", }) }) }) @@ -116,6 +118,13 @@ describe("readAllSourceChunks", () => { }) expect(readChunks).toHaveBeenCalledTimes(2) + expect(readChunks).toHaveBeenNthCalledWith(1, { + documentId: "doc_1", + revisionKey: "rev_1", + page: 1, + pageSize: 200, + assetUrlPolicy: "durable", + }) expect(chunks.map((chunk) => chunk.parserChunkId)).toEqual(["c1", "c2"]) }) }) diff --git a/src/domains/chunks/read.ts b/src/domains/chunks/read.ts index 320d559..1bf4790 100644 --- a/src/domains/chunks/read.ts +++ b/src/domains/chunks/read.ts @@ -16,8 +16,8 @@ type ReadableSource = { /** * Read a single display page of parsed chunks through the SDK. The SDK serves * from configured Blob storage when fresh and falls back to Knowhere remote - * otherwise. Display reads intentionally do not request durable asset URLs: - * chat hardens only the specific assets it sends back to the user. + * otherwise. Display reads request durable asset URLs so media chunks remain + * usable whether the page came from Blob or the remote Knowhere fallback. */ export async function readSourceChunkPage(input: { readonly knowledge: Knowledge @@ -29,6 +29,7 @@ export async function readSourceChunkPage(input: { ...(input.source.revisionKey ? { revisionKey: input.source.revisionKey } : {}), page: input.params.page, pageSize: input.params.pageSize, + assetUrlPolicy: "durable", }) const chunks = response.chunks.map((chunk) => @@ -68,6 +69,7 @@ export async function readAllSourceChunks(input: { : {}), page, pageSize: loadAllPageSize, + assetUrlPolicy: "durable", }) for (const chunk of response.chunks) { chunks.push( diff --git a/src/domains/sources/parsed-sync-route-workflow.test.ts b/src/domains/sources/parsed-sync-route-workflow.test.ts index 91095f5..670ee67 100644 --- a/src/domains/sources/parsed-sync-route-workflow.test.ts +++ b/src/domains/sources/parsed-sync-route-workflow.test.ts @@ -5,10 +5,6 @@ const mocks = vi.hoisted(() => ({ makeKnowhereClientWithParsedStorage: vi.fn(), releaseSyncCapacity: vi.fn(), updateSyncStatus: vi.fn(), - findInWorkspace: vi.fn(), - markFailed: vi.fn(), - markReady: vi.fn(), - markSourceReadyAfterReconciliation: vi.fn(), loggerInfo: vi.fn(), loggerError: vi.fn(), })) @@ -20,16 +16,9 @@ vi.mock("@/integrations/knowhere", () => ({ vi.mock("./workflow-runtime", () => ({ sourceWorkflowRuntime: { updateSyncStatus: mocks.updateSyncStatus, - findInWorkspace: mocks.findInWorkspace, - markFailed: mocks.markFailed, - markReady: mocks.markReady, }, })) -vi.mock("./source-reconcile-workflow", () => ({ - markSourceReadyAfterReconciliation: mocks.markSourceReadyAfterReconciliation, -})) - vi.mock("./parsed-document-sync-capacity", () => ({ parsedDocumentSyncCapacityGuard: { acquire: mocks.acquireSyncCapacity, @@ -73,16 +62,13 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { }, }) mocks.releaseSyncCapacity.mockResolvedValue(undefined) - mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ - status: "ready", - }) }) afterEach(() => { vi.clearAllMocks() }) - it("marks the source ready when sync completes in one segment", async () => { + it("records completion when sync completes in one segment", async () => { const syncParsedDocument = vi.fn(async () => ({ documentId: "doc_1", revisionKey: "rev_1", @@ -105,18 +91,13 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { "source_1", { revisionKey: "rev_1", syncStatus: "completed" }, ) - expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ - workspaceId: "workspace_1", - sourceId: "source_1", - documentId: "doc_1", - }) expect(mocks.releaseSyncCapacity).toHaveBeenCalledWith({ leaseToken: "lease_1", releaseReason: "completed", }) }) - it("marks ready and triggers a continuation when sync is incomplete", async () => { + it("triggers a continuation when sync is incomplete", async () => { const syncParsedDocument = vi.fn(async () => ({ documentId: "doc_1", revisionKey: "rev_1", @@ -145,11 +126,6 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { restore() } - expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ - workspaceId: "workspace_1", - sourceId: "source_1", - documentId: "doc_1", - }) expect(triggered).toHaveLength(1) expect(triggered[0]?.segmentIndex).toBe(1) expect(triggered[0]?.workflowRunId).toBe("doc_1-sync-rev_1-1") @@ -249,6 +225,37 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { ]) expect(mocks.releaseSyncCapacity).not.toHaveBeenCalled() }) + + it("does not release capacity when Upstash aborts during a planned step", async () => { + const workflowAbort = new Error("planned workflow step") + workflowAbort.name = "WorkflowAbort" + const syncParsedDocument = vi.fn(async () => ({ + documentId: "doc_1", + revisionKey: "rev_1", + completed: true, + })) + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: {}, + knowledge: { syncParsedDocument }, + }) + const run: RunStep = async (id, task) => { + if (id === "sync-0-0") throw workflowAbort + return task() + } + + await expect( + parsedSyncRouteWorkflow.runParsedSyncWorkflow({ + context: { + run, + url: "https://notebook.example/api/sources/parsed-sync", + }, + payload: basePayload, + }), + ).rejects.toThrow("planned workflow step") + + expect(mocks.releaseSyncCapacity).not.toHaveBeenCalled() + expect(mocks.updateSyncStatus).not.toHaveBeenCalled() + }) }) describe("parsedSyncRouteWorkflow.markSyncFailedAfterWorkflowFailure", () => { @@ -256,9 +263,7 @@ describe("parsedSyncRouteWorkflow.markSyncFailedAfterWorkflowFailure", () => { vi.clearAllMocks() }) - it("fails a parsing source with failure_stage storage_sync", async () => { - mocks.findInWorkspace.mockResolvedValue({ id: "source_1", status: "parsing" }) - + it("records failed sync metadata without failing the source", async () => { await parsedSyncRouteWorkflow.markSyncFailedAfterWorkflowFailure( { ...basePayload, revisionKey: "rev_1" }, "boom", @@ -269,28 +274,5 @@ describe("parsedSyncRouteWorkflow.markSyncFailedAfterWorkflowFailure", () => { "source_1", { revisionKey: "rev_1", syncStatus: "failed", syncError: "boom" }, ) - expect(mocks.markFailed).toHaveBeenCalledWith( - "workspace_1", - "source_1", - expect.stringContaining("storage sync failed"), - "parsing", - "storage_sync", - ) - }) - - it("does not fail an already-ready source, only records sync_status", async () => { - mocks.findInWorkspace.mockResolvedValue({ id: "source_1", status: "ready" }) - - await parsedSyncRouteWorkflow.markSyncFailedAfterWorkflowFailure( - basePayload, - "boom", - ) - - expect(mocks.updateSyncStatus).toHaveBeenCalledWith( - "workspace_1", - "source_1", - { revisionKey: undefined, syncStatus: "failed", syncError: "boom" }, - ) - expect(mocks.markFailed).not.toHaveBeenCalled() }) }) diff --git a/src/domains/sources/parsed-sync-route-workflow.ts b/src/domains/sources/parsed-sync-route-workflow.ts index 6be3b23..e832b01 100644 --- a/src/domains/sources/parsed-sync-route-workflow.ts +++ b/src/domains/sources/parsed-sync-route-workflow.ts @@ -1,6 +1,6 @@ import "server-only" -import { Client, type WorkflowContext } from "@upstash/workflow" +import { Client, WorkflowAbort, type WorkflowContext } from "@upstash/workflow" import type { KnowledgeSyncParsedDocumentResponse } from "@ontos-ai/knowhere-sdk" import { makeKnowhereClientWithParsedStorage } from "@/integrations/knowhere" @@ -10,7 +10,6 @@ import { type ParsedSyncPayload, } from "./parsed-document-sync-scheduler" import { parsedDocumentSyncCapacityGuard } from "./parsed-document-sync-capacity" -import { markSourceReadyAfterReconciliation } from "./source-reconcile-workflow" import { sourceWorkflowRuntime } from "./workflow-runtime" type ParsedSyncWorkflowContext = Pick< @@ -72,20 +71,10 @@ async function runParsedSyncWorkflow(input: { workspaceId, }) - const preSyncReady = await context.run( - `source-ready-before-sync-${payload.segmentIndex}`, - async () => - markSourceReadyAfterReconciliation({ - workspaceId, - sourceId, - documentId, - }), - ) - if (preSyncReady.status === "gone") return - let revisionKey = payload.revisionKey let completed = false let releaseReason: SyncLeaseReleaseReason = "incomplete" + let shouldReleaseLease = true const capacity = await context.run( `acquire-sync-capacity-${payload.segmentIndex}`, @@ -196,24 +185,29 @@ async function runParsedSyncWorkflow(input: { return } } catch (error) { + if (isWorkflowControlAbort(error)) { + shouldReleaseLease = false + throw error + } releaseReason = "failed" throw error } finally { - await context.run(`release-sync-capacity-${payload.segmentIndex}`, async () => - releaseCapacityLease({ - leaseToken: capacity.leaseToken, - releaseReason, - sourceId, - documentId, - }), - ) + if (shouldReleaseLease) { + await context.run(`release-sync-capacity-${payload.segmentIndex}`, async () => + releaseCapacityLease({ + leaseToken: capacity.leaseToken, + releaseReason, + sourceId, + documentId, + }), + ) + } } logger.info("parsed-sync: parsed document sync finished", { sourceId, documentId, revisionKey, - status: preSyncReady.status, }) } @@ -261,10 +255,9 @@ async function markSyncFailedAfterWorkflowFailure( const normalized = normalizeParsedSyncPayload(payload) const reason = getSafeFailureReason(failResponse) - // Record the storage-sync failure. A source still `parsing` is failed with - // failure_stage=storage_sync so a retry resumes sync without reparsing; an - // already-ready source is left ready (it still serves via remote fallback), - // only its sync_status is marked failed for observability. + // Parsed storage is a cache/read model. Exhausted sync failure is recorded for + // observability, but the source remains ready and SDK reads can fall back to + // Knowhere remote. await sourceWorkflowRuntime.updateSyncStatus( normalized.workspaceId, normalized.sourceId, @@ -275,26 +268,11 @@ async function markSyncFailedAfterWorkflowFailure( }, ) - const source = await sourceWorkflowRuntime.findInWorkspace( - normalized.workspaceId, - normalized.sourceId, - ) - if (source?.status === "parsing") { - await sourceWorkflowRuntime.markFailed( - normalized.workspaceId, - normalized.sourceId, - `Parsed document storage sync failed: ${reason}`, - "parsing", - "storage_sync", - ) - } - logger.error("parsed-sync: marked sync failed after workflow failure", { workspaceId: normalized.workspaceId, sourceId: normalized.sourceId, documentId: normalized.documentId, segmentIndex: normalized.segmentIndex, - sourceStatus: source?.status, }) } @@ -314,6 +292,13 @@ function setContinuationTriggerForTesting( } } +function isWorkflowControlAbort(error: unknown): boolean { + return ( + (error instanceof WorkflowAbort && error.constructor === WorkflowAbort) || + (error instanceof Error && error.name === "WorkflowAbort") + ) +} + export const parsedSyncRouteWorkflow = { markSyncFailedAfterWorkflowFailure, normalizeParsedSyncPayload, diff --git a/src/domains/sources/source-reconcile-route-workflow.test.ts b/src/domains/sources/source-reconcile-route-workflow.test.ts index ee415ff..730fcf0 100644 --- a/src/domains/sources/source-reconcile-route-workflow.test.ts +++ b/src/domains/sources/source-reconcile-route-workflow.test.ts @@ -5,6 +5,7 @@ const mocks = vi.hoisted(() => ({ enqueueParsedDocumentSync: vi.fn(), releaseSyncCapacity: vi.fn(), updateSyncStatus: vi.fn(), + updateRevisionKey: vi.fn(), markFailed: vi.fn(), loggerError: vi.fn(), loggerInfo: vi.fn(), @@ -23,6 +24,7 @@ vi.mock("@/domains/sources/workflow-runtime", () => ({ sourceWorkflowRuntime: { markFailed: mocks.markFailed, updateSyncStatus: mocks.updateSyncStatus, + updateRevisionKey: mocks.updateRevisionKey, }, })) @@ -90,9 +92,11 @@ describe("sourceReconcileRouteWorkflow", () => { activeCounts, }) mocks.releaseSyncCapacity.mockResolvedValue(undefined) + mocks.enqueueParsedDocumentSync.mockResolvedValue(undefined) mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ status: "ready", }) + mocks.updateRevisionKey.mockResolvedValue({ id: "source_1" }) }) afterEach(() => { @@ -116,7 +120,7 @@ describe("sourceReconcileRouteWorkflow", () => { }) }) - it("marks the source ready then syncs the parsed document", async () => { + it("marks the source ready, records pending sync, and enqueues parsed-sync", async () => { const context = createWorkflowContext() const continuations: ContinuationTriggerInput[] = [] const restore = @@ -149,46 +153,44 @@ describe("sourceReconcileRouteWorkflow", () => { restore() } - expect(wired.knowledge.syncParsedDocument).toHaveBeenCalledWith({ + expect(wired.knowledge.syncParsedDocument).not.toHaveBeenCalled() + expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "source_1", documentId: "doc_1", - revisionKey: "rev_1", }) + expect(mocks.updateRevisionKey).toHaveBeenCalledWith( + "workspace_1", + "source_1", + "rev_1", + ) expect(mocks.updateSyncStatus).toHaveBeenCalledWith( "workspace_1", "source_1", - { revisionKey: "rev_1", syncStatus: "completed" }, + { revisionKey: "rev_1", syncStatus: "pending", syncError: null }, ) - expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ - workspaceId: "workspace_1", - sourceId: "source_1", - documentId: "doc_1", - }) - expect(mocks.acquireSyncCapacity).toHaveBeenCalledWith({ + expect(mocks.enqueueParsedDocumentSync).toHaveBeenCalledWith({ workspaceId: "workspace_1", sourceId: "source_1", documentId: "doc_1", + apiKey: "jwt_1", revisionKey: "rev_1", }) - expect(mocks.releaseSyncCapacity).toHaveBeenCalledWith({ - leaseToken: "lease_1", - releaseReason: "completed", - }) - expect(mocks.enqueueParsedDocumentSync).not.toHaveBeenCalled() + expect(mocks.acquireSyncCapacity).not.toHaveBeenCalled() + expect(mocks.releaseSyncCapacity).not.toHaveBeenCalled() expect(continuations).toEqual([]) }) - it("hands off to parsed-sync when sync is incomplete", async () => { + it("keeps the source ready when parsed-sync enqueue fails", async () => { const context = createWorkflowContext() - const syncParsedDocument = vi.fn(async () => ({ - documentId: "doc_1", - revisionKey: "rev_1", - completed: false, - })) - const wired = createClient({ syncParsedDocument }) + const wired = createClient({}) mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ client: wired.client, knowledge: wired.knowledge, }) + mocks.enqueueParsedDocumentSync.mockRejectedValue( + new Error("qstash unavailable"), + ) mocks.pollSourceReconciliation.mockResolvedValue({ kind: "ready-to-prepare", jobId: "job_1", @@ -204,120 +206,18 @@ describe("sourceReconcileRouteWorkflow", () => { }), }) + expect(wired.knowledge.syncParsedDocument).not.toHaveBeenCalled() expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ workspaceId: "workspace_1", sourceId: "source_1", documentId: "doc_1", }) - expect(mocks.enqueueParsedDocumentSync).toHaveBeenCalledWith({ - workspaceId: "workspace_1", - sourceId: "source_1", - documentId: "doc_1", - apiKey: "jwt_1", - revisionKey: "rev_1", - }) - expect(mocks.releaseSyncCapacity).toHaveBeenCalledWith({ - leaseToken: "lease_1", - releaseReason: "incomplete", - }) - }) - - it("records sync failure when Blob sync throws", async () => { - const context = createWorkflowContext() - const syncParsedDocument = vi.fn(async () => { - throw new Error("blob write failed") - }) - const wired = createClient({ syncParsedDocument }) - mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ - client: wired.client, - knowledge: wired.knowledge, - }) - mocks.pollSourceReconciliation.mockResolvedValue({ - kind: "ready-to-prepare", - jobId: "job_1", - documentId: "doc_1", - }) - - await expect( - sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ - context, - payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ - workspaceId: "workspace_1", - sourceId: "source_1", - apiKey: "jwt_1", - }), - }), - ).rejects.toThrow("blob write failed") - expect(mocks.updateSyncStatus).toHaveBeenCalledWith( "workspace_1", "source_1", - { - revisionKey: "rev_1", - syncStatus: "failed", - syncError: "blob write failed", - }, + { revisionKey: "rev_1", syncStatus: "pending", syncError: null }, ) expect(mocks.markFailed).not.toHaveBeenCalled() - expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ - workspaceId: "workspace_1", - sourceId: "source_1", - documentId: "doc_1", - }) - expect(mocks.releaseSyncCapacity).toHaveBeenCalledWith( - { - leaseToken: "lease_1", - releaseReason: "failed", - }, - ) - }) - - it("marks the source ready and schedules delayed sync when capacity is full", async () => { - const context = createWorkflowContext() - const wired = createClient({}) - mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ - client: wired.client, - knowledge: wired.knowledge, - }) - mocks.acquireSyncCapacity.mockResolvedValue({ - kind: "capacity-full", - reason: "document", - waitSeconds: 60, - activeCounts: { - globalActive: 10, - workspaceActive: 1, - documentActive: 1, - }, - }) - mocks.pollSourceReconciliation.mockResolvedValue({ - kind: "ready-to-prepare", - jobId: "job_1", - documentId: "doc_1", - }) - - await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ - context, - payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ - workspaceId: "workspace_1", - sourceId: "source_1", - apiKey: "jwt_1", - }), - }) - - expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ - workspaceId: "workspace_1", - sourceId: "source_1", - documentId: "doc_1", - }) - expect(wired.knowledge.syncParsedDocument).not.toHaveBeenCalled() - expect(mocks.enqueueParsedDocumentSync).toHaveBeenCalledWith({ - workspaceId: "workspace_1", - sourceId: "source_1", - documentId: "doc_1", - apiKey: "jwt_1", - revisionKey: "rev_1", - delaySeconds: 60, - }) expect(mocks.releaseSyncCapacity).not.toHaveBeenCalled() }) diff --git a/src/domains/sources/source-reconcile-route-workflow.ts b/src/domains/sources/source-reconcile-route-workflow.ts index 1e45127..5d2cb2c 100644 --- a/src/domains/sources/source-reconcile-route-workflow.ts +++ b/src/domains/sources/source-reconcile-route-workflow.ts @@ -9,7 +9,6 @@ import { import { makeKnowhereClientWithParsedStorage } from "@/integrations/knowhere" import { logger } from "@/lib/logger" import { enqueueParsedDocumentSync } from "./parsed-document-sync-scheduler" -import { parsedDocumentSyncCapacityGuard } from "./parsed-document-sync-capacity" import { sourceWorkflowRuntime } from "./workflow-runtime" type ReconcilePayload = { @@ -41,12 +40,25 @@ type ContinuationTriggerInput = { readonly workflowRunId: string } -type SyncLeaseReleaseReason = "completed" | "incomplete" | "failed" +type RevisionKeyClient = { + readonly documents: { + readonly listChunks: ( + documentId: string, + params: { + readonly page: number + readonly pageSize: number + readonly includeAssetUrls: boolean + }, + ) => Promise<{ + readonly jobResultId?: string | null + readonly jobId?: string | null + }> + } +} const maxPollAttempts = 25 const initialDelaySeconds = 3 const maxDelaySeconds = 30 -const maxSyncStepsPerReconcile = 4 let triggerContinuation: typeof triggerReconcileContinuation = triggerReconcileContinuation @@ -57,7 +69,7 @@ async function runPollAndMirrorWorkflow(input: { }): Promise { const { context, payload } = input const { workspaceId, sourceId, apiKey } = payload - const { client, knowledge } = makeKnowhereClientWithParsedStorage(apiKey, { + const { client } = makeKnowhereClientWithParsedStorage(apiKey, { workspaceId, }) let delay = initialDelaySeconds @@ -125,21 +137,6 @@ async function runPollAndMirrorWorkflow(input: { return } - const revisionKey = await context.run("resolve-revision-key", async () => { - const firstPage = await client.documents.listChunks( - jobToPrepare.documentId, - { page: 1, pageSize: 1, includeAssetUrls: false }, - ) - return firstPage.jobResultId ?? firstPage.jobId ?? jobToPrepare.jobId - }) - await context.run("record-sync-pending", async () => - sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { - revisionKey, - syncStatus: "pending", - syncError: null, - }), - ) - const ready = await context.run("source-ready", async () => markSourceReadyAfterReconciliation({ workspaceId, @@ -149,139 +146,87 @@ async function runPollAndMirrorWorkflow(input: { ) if (ready.status === "gone") return - const capacity = await context.run("acquire-sync-capacity", async () => - parsedDocumentSyncCapacityGuard.acquire({ - workspaceId, + const revisionKey = await context.run("resolve-revision-key", async () => + resolveParsedRevisionKey({ + client, sourceId, documentId: jobToPrepare.documentId, + fallbackRevisionKey: jobToPrepare.jobId, + }), + ) + + await context.run("record-source-revision-key", async () => + sourceWorkflowRuntime.updateRevisionKey(workspaceId, sourceId, revisionKey), + ) + await context.run("record-sync-pending", async () => + sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { revisionKey, + syncStatus: "pending", + syncError: null, }), ) - if (capacity.kind === "source-missing") return - if (capacity.kind === "capacity-full") { - await context.run("enqueue-capacity-retry", async () => - enqueueParsedDocumentSync({ - workspaceId, - sourceId, - documentId: jobToPrepare.documentId, - apiKey, - revisionKey, - delaySeconds: capacity.waitSeconds, - }), - ) - logger.info("workflow: parsed storage sync delayed by capacity guard", { + const enqueueResult = await context.run("enqueue-parsed-sync", async () => + enqueueParsedSyncBestEffort({ + workspaceId, sourceId, documentId: jobToPrepare.documentId, + apiKey, revisionKey, - reason: capacity.reason, - waitSeconds: capacity.waitSeconds, - activeCounts: capacity.activeCounts, - }) - return - } - - let syncCompleted = false - let releaseReason: SyncLeaseReleaseReason = "incomplete" - try { - await context.run("record-sync-running", async () => - sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { - revisionKey, - syncStatus: "running", - syncError: null, - }), - ) - - for (let step = 0; step < maxSyncStepsPerReconcile; step++) { - const result = await context.run(`parsed-sync-${step}`, async () => { - try { - return await knowledge.syncParsedDocument({ - documentId: jobToPrepare.documentId, - revisionKey, - }) - } catch (error) { - await sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { - revisionKey, - syncStatus: "failed", - syncError: getErrorMessage(error), - }) - throw error - } - }) - if (result.completed) { - syncCompleted = true - releaseReason = "completed" - break - } - } - - if (!syncCompleted) { - await context.run("record-sync-progress", async () => - sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { - revisionKey, - syncStatus: "running", - }), - ) - await context.run("enqueue-parsed-sync-continuation", async () => - enqueueParsedDocumentSync({ - workspaceId, - sourceId, - documentId: jobToPrepare.documentId, - apiKey, - revisionKey, - }), - ) - logger.info("workflow: parsed storage sync handed off to parsed-sync", { - sourceId, - documentId: jobToPrepare.documentId, - revisionKey, - }) - return - } - } catch (error) { - releaseReason = "failed" - throw error - } finally { - await context.run("release-sync-capacity", async () => - releaseCapacityLease({ - leaseToken: capacity.leaseToken, - releaseReason, - sourceId, - documentId: jobToPrepare.documentId, - }), - ) - } - - await context.run("record-sync-completed", async () => - sourceWorkflowRuntime.updateSyncStatus(workspaceId, sourceId, { - revisionKey, - syncStatus: "completed", }), ) + logger.info("workflow: source parse reconciliation finished", { sourceId, jobId: jobToPrepare.jobId, revisionKey, status: ready.status, + parsedSyncEnqueued: enqueueResult.enqueued, }) } -async function releaseCapacityLease(input: { - readonly leaseToken: string - readonly releaseReason: SyncLeaseReleaseReason +async function resolveParsedRevisionKey(input: { + readonly client: RevisionKeyClient readonly sourceId: string readonly documentId: string -}): Promise { + readonly fallbackRevisionKey: string +}): Promise { try { - await parsedDocumentSyncCapacityGuard.release({ - leaseToken: input.leaseToken, - releaseReason: input.releaseReason, + const firstPage = await input.client.documents.listChunks(input.documentId, { + page: 1, + pageSize: 1, + includeAssetUrls: false, + }) + return firstPage.jobResultId ?? firstPage.jobId ?? input.fallbackRevisionKey + } catch (error) { + logger.warn("workflow: failed to resolve parsed revision key", { + sourceId: input.sourceId, + documentId: input.documentId, + fallbackRevisionKey: input.fallbackRevisionKey, + error: getErrorMessage(error), }) + return input.fallbackRevisionKey + } +} + +async function enqueueParsedSyncBestEffort(input: { + readonly workspaceId: string + readonly sourceId: string + readonly documentId: string + readonly apiKey: string + readonly revisionKey: string +}): Promise<{ readonly enqueued: boolean }> { + try { + await enqueueParsedDocumentSync(input) + return { enqueued: true } } catch (error) { - logger.error("workflow: failed to release sync capacity lease", { + logger.error("workflow: failed to enqueue parsed storage sync", { + workspaceId: input.workspaceId, sourceId: input.sourceId, documentId: input.documentId, + revisionKey: input.revisionKey, error: getErrorMessage(error), }) + return { enqueued: false } } } diff --git a/src/proxy.test.ts b/src/proxy.test.ts index b965699..0d80c19 100644 --- a/src/proxy.test.ts +++ b/src/proxy.test.ts @@ -63,6 +63,16 @@ describe("proxy", () => { ); }); + it("allows anonymous parsed-sync workflow callbacks", () => { + const response = proxy( + new NextRequest("http://localhost:3001/api/sources/parsed-sync", { + method: "POST", + }), + ); + + expect(response.headers.get("x-middleware-next")).toBe("1"); + }); + it("allows protected app routes without a session when KNOWHERE_API_KEY is configured", () => { process.env.KNOWHERE_API_KEY = "sk_dev_key"; diff --git a/src/proxy.ts b/src/proxy.ts index 6177d3b..7527048 100644 --- a/src/proxy.ts +++ b/src/proxy.ts @@ -27,6 +27,7 @@ const PUBLIC_PATHS: readonly string[] = [ "/favicon.ico", "/api/internal/health", "/api/sources/reconcile", + "/api/sources/parsed-sync", ] const STATIC_EXTENSIONS = /\.(?:svg|png|jpe?g|gif|webp|ico|woff2?|ttf|eot|css|js|map|txt|xml|webmanifest|json|pdf)$/i From 59737a00ae379393a76efd034efa59f81523b2cf Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 6 Jul 2026 22:38:34 +0800 Subject: [PATCH 20/51] Fix hidden library source auto-selection --- src/components/workspace-source-state.test.ts | 37 ++++++++++++++++++- src/components/workspace-source-state.ts | 18 +++++++-- 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/src/components/workspace-source-state.test.ts b/src/components/workspace-source-state.test.ts index bfd0273..5345646 100644 --- a/src/components/workspace-source-state.test.ts +++ b/src/components/workspace-source-state.test.ts @@ -28,7 +28,7 @@ describe("workspaceSourceState", () => { ); }); - it("can select an unmaterialized Official Library row for preview", () => { + it("skips unmaterialized Official Library rows when choosing the initial Source", () => { const sources: readonly SourceView[] = [ { id: "demo-spacex-s1", @@ -54,10 +54,43 @@ describe("workspaceSourceState", () => { ]; expect(workspaceSourceState.getInitialSelectedSourceId(sources)).toBe( - "demo-spacex-s1", + "source_ready", ); }); + it("does not resolve an unmaterialized Official Library row as the selected Source", () => { + const sources: readonly SourceView[] = [ + { + id: "demo-spacex-s1", + kind: "demo", + demoSourceId: "demo-spacex-s1", + title: "spacex-s1.pdf", + status: "ready", + mimeType: "application/pdf", + excludedFromQuery: false, + officialLibrary: { + librarySourceId: "financial-spacex-s1", + categoryId: "financial-reports", + sourceUrl: "https://example.com/spacex-s1.pdf", + }, + }, + { + id: "source_ready", + title: "ready.pdf", + status: "ready", + mimeType: "application/pdf", + excludedFromQuery: false, + }, + ]; + + expect( + workspaceSourceState.getResolvedSelectedSourceId( + sources, + "demo-spacex-s1", + ), + ).toBeNull(); + }); + it("selects a preferred document source when opening a chunk-tree link", () => { const sources: readonly SourceView[] = [ { diff --git a/src/components/workspace-source-state.ts b/src/components/workspace-source-state.ts index 97c0c55..5d33744 100644 --- a/src/components/workspace-source-state.ts +++ b/src/components/workspace-source-state.ts @@ -53,7 +53,8 @@ function getInitialSelectedSourceId( if (preferredDocumentId) { const preferredSource = sources.find( (source) => - source.documentId === preferredDocumentId && isReadySource(source), + source.documentId === preferredDocumentId && + isReadyVisibleSource(source), ) if (preferredSource) return preferredSource.id } @@ -62,7 +63,7 @@ function getInitialSelectedSourceId( } function getFirstReadySourceId(sources: readonly SourceView[]): string | null { - return sources.find(isReadySource)?.id ?? null + return sources.find(isReadyVisibleSource)?.id ?? null } function getResolvedSelectedSourceId( @@ -73,14 +74,15 @@ function getResolvedSelectedSourceId( const selectedSource = sources.find((source) => source.id === selectedSourceId) if (selectedSource) { - return selectedSource.id + return isVisibleSource(selectedSource) ? selectedSource.id : null } const selectedDocumentId = getRemoteSourceDocumentId(selectedSourceId) if (selectedDocumentId) { const localizedSource = sources.find( (source) => - source.documentId === selectedDocumentId && isReadySource(source), + source.documentId === selectedDocumentId && + isReadyVisibleSource(source), ) if (localizedSource) return localizedSource.id } @@ -92,6 +94,14 @@ function isReadySource(source: SourceView): boolean { return source.status === "ready" } +function isReadyVisibleSource(source: SourceView): boolean { + return isReadySource(source) && isVisibleSource(source) +} + +function isVisibleSource(source: SourceView): boolean { + return source.officialLibrary === undefined +} + function applyQueryExclusions( sources: readonly SourceView[], sourceExclusionById: SourceExclusionState, From 660204fed6544786848ac22752918b5ecc1d634b Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 13:35:59 +0800 Subject: [PATCH 21/51] Present page-mode sources as page assets --- .../[sourceId]/page-assets/route.test.ts | 271 +++++++++++++++++ .../sources/[sourceId]/page-assets/route.ts | 32 ++ src/components/chat-message-list.test.ts | 14 +- src/components/chat-message-list.tsx | 34 +-- src/components/chunks-panel.test.ts | 110 +++++++ src/components/chunks-panel.tsx | 286 +++++++++++++++++- src/components/source-row.test.ts | 21 ++ src/components/source-row.tsx | 3 + src/components/sources-panel.tsx | 1 + .../workspace-citation-focus.test.ts | 43 +++ src/components/workspace-citation-focus.ts | 32 ++ src/components/workspace-citation-state.ts | 37 +++ .../workspace-selected-chunks.test.ts | 22 ++ src/components/workspace-selected-chunks.ts | 4 +- src/components/workspace-shell-layout.tsx | 16 + src/components/workspace-shell.tsx | 1 + src/domains/chat/chat-citation-persistence.ts | 1 + src/domains/chat/citations.ts | 3 + src/domains/chat/page-citation-assets.ts | 8 +- src/domains/chat/types.ts | 1 + src/domains/chat/view.ts | 4 + src/domains/sources/counts.test.ts | 61 ++++ src/domains/sources/counts.ts | 138 ++++++++- src/domains/sources/page-assets.ts | 104 +++++++ src/domains/sources/route-page-assets.ts | 161 ++++++++++ src/domains/sources/route-service.ts | 5 + src/domains/sources/route-types.ts | 32 +- src/domains/sources/types.ts | 17 ++ src/domains/sources/view.test.ts | 12 + src/domains/sources/view.ts | 9 +- src/domains/workspace/client.ts | 36 ++- 31 files changed, 1448 insertions(+), 71 deletions(-) create mode 100644 src/app/api/sources/[sourceId]/page-assets/route.test.ts create mode 100644 src/app/api/sources/[sourceId]/page-assets/route.ts create mode 100644 src/domains/sources/page-assets.ts create mode 100644 src/domains/sources/route-page-assets.ts diff --git a/src/app/api/sources/[sourceId]/page-assets/route.test.ts b/src/app/api/sources/[sourceId]/page-assets/route.test.ts new file mode 100644 index 0000000..4bf5c5e --- /dev/null +++ b/src/app/api/sources/[sourceId]/page-assets/route.test.ts @@ -0,0 +1,271 @@ +import { NextRequest } from "next/server" +import { beforeEach, describe, expect, it, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + deleteBlob: vi.fn(), + ensureApiKeyForWorkspace: vi.fn(), + ensureWorkspace: vi.fn(), + findSourceInWorkspace: vi.fn(), + getCurrentUser: vi.fn(), + makeKnowhereClient: vi.fn(), + makeKnowhereClientWithParsedStorage: vi.fn(), + readChunks: vi.fn(), + requireUser: vi.fn(), +})) + +vi.mock("next/headers", () => ({ + headers: vi.fn(async () => new Headers({ cookie: "session=abc" })), +})) + +vi.mock("@/integrations/dashboard/api-key-service", () => ({ + ensureApiKeyForWorkspace: mocks.ensureApiKeyForWorkspace, +})) + +vi.mock("@/integrations/knowhere-demo", () => ({ + knowhereDemoApi: { + fetchCatalog: vi.fn(), + fetchChunkPage: vi.fn(), + }, +})) + +vi.mock("@/infrastructure/auth", () => ({ + getCurrentUser: mocks.getCurrentUser, + requireUser: mocks.requireUser, +})) + +vi.mock("@/integrations/knowhere", () => ({ + makeKnowhereClient: mocks.makeKnowhereClient, + makeKnowhereClientWithParsedStorage: + mocks.makeKnowhereClientWithParsedStorage, +})) + +vi.mock("@vercel/blob", () => ({ + del: mocks.deleteBlob, +})) + +vi.mock("@/domains/sources/service", () => ({ + sourceService: { + findInWorkspace: mocks.findSourceInWorkspace, + localizeRemoteDocument: vi.fn(), + }, +})) + +vi.mock("@/domains/workspace/service", () => ({ + workspaceService: { + ensureWorkspace: mocks.ensureWorkspace, + }, +})) + +import { GET } from "./route" + +describe("GET /api/sources/[sourceId]/page-assets", () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: { documents: { listChunks: vi.fn() } }, + knowledge: { readChunks: mocks.readChunks }, + }) + }) + + it("returns durable page assets for a ready workspace source", async () => { + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockResolvedValue({ + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + assetUrl: "https://assets.example/fallback.png", + metadata: { + pageAssets: [ + { + pageNum: 1, + assetUrl: "https://assets.example/page-000001.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 3, + totalPages: 3, + }) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + pages: [ + { + pageNumber: 1, + assetUrl: "https://assets.example/page-000001.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 3, + totalPages: 3, + }, + }) + expect(response.status).toBe(200) + expect(mocks.readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "job_1", + chunkType: "page", + page: 1, + pageSize: 1, + assetUrlPolicy: "durable", + }) + }) + + it("rejects non-ready workspace sources", async () => { + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + status: "parsing", + knowhereDocumentId: "doc_1", + }), + ) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + message: "Source is not ready.", + }) + expect(response.status).toBe(409) + expect(mocks.readChunks).not.toHaveBeenCalled() + }) + + it("returns an empty page list when page chunks have no usable assets", async () => { + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereDocumentId: "doc_1", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockResolvedValue({ + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + metadata: { pageAssets: [] }, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, + }) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + pages: [], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + }) + expect(response.status).toBe(200) + }) +}) + +function makeReadySource(overrides: Record) { + return { + id: "00000000-0000-0000-0000-000000000002", + workspaceId: "workspace_1", + title: "notes.pdf", + mimeType: "application/pdf", + sizeBytes: 1024, + status: "ready", + failureReason: null, + failureStage: null, + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: null, + originalBlobUrl: null, + demoKey: null, + createdAt: new Date("2026-05-10T00:00:00.000Z"), + updatedAt: new Date("2026-05-10T00:00:00.000Z"), + deletedAt: null, + ...overrides, + } +} diff --git a/src/app/api/sources/[sourceId]/page-assets/route.ts b/src/app/api/sources/[sourceId]/page-assets/route.ts new file mode 100644 index 0000000..b7606f7 --- /dev/null +++ b/src/app/api/sources/[sourceId]/page-assets/route.ts @@ -0,0 +1,32 @@ +import type { NextRequest, NextResponse } from "next/server" + +import { getChunkPageParams } from "@/domains/chunks" +import { createSourceRouteService } from "@/domains/sources/route-service" +import { withApiErrorResponse } from "@/lib/api-error-response" +import { nextRouteContext } from "@/lib/next-route-context" +import { nextRouteResponse } from "@/lib/next-route-response" + +type RouteContext = { + params: Promise<{ + sourceId: string + }> +} + +const sourceRouteService = createSourceRouteService() + +export async function GET( + request: NextRequest, + context: RouteContext, +): Promise { + return withApiErrorResponse("sources:page-assets", async () => { + const { sourceId } = await context.params + const routeContext = await nextRouteContext.read() + const result = await sourceRouteService.loadSourcePageAssets({ + cookieHeader: routeContext.cookieHeader, + pageParams: getChunkPageParams(request.nextUrl.searchParams), + sourceId, + }) + + return nextRouteResponse.toNextResponse(result) + }) +} diff --git a/src/components/chat-message-list.test.ts b/src/components/chat-message-list.test.ts index f94478e..9444128 100644 --- a/src/components/chat-message-list.test.ts +++ b/src/components/chat-message-list.test.ts @@ -146,6 +146,7 @@ describe("ChatMessageList", () => { score: 0.9, pageCitationAssetUrl: "https://blob.example/pages/page-000004.png", + pageCitationPageNumber: 4, source: { documentId: "doc_1", sourceFileName: "report.pdf", @@ -162,18 +163,17 @@ describe("ChatMessageList", () => { const citationButton = screen.getByRole("button", { name: "Open source report.pdf", }); - const pageImageLink = screen.getByRole("link", { - name: "Open page image for report.pdf", - }); - - expect(pageImageLink.getAttribute("href")).toBe( - "https://blob.example/pages/page-000004.png", - ); + expect( + screen.queryByRole("link", { + name: "Open page image for report.pdf", + }), + ).toBeNull(); await user.click(citationButton); expect(onCitationClick).toHaveBeenCalledWith( expect.objectContaining({ pageCitationAssetUrl: "https://blob.example/pages/page-000004.png", + pageCitationPageNumber: 4, }), "assistant_1:0", ); diff --git a/src/components/chat-message-list.tsx b/src/components/chat-message-list.tsx index 531c71f..38c5059 100644 --- a/src/components/chat-message-list.tsx +++ b/src/components/chat-message-list.tsx @@ -2,7 +2,7 @@ import { type CSSProperties, type ReactElement } from "react"; import { type VirtualItem } from "@tanstack/react-virtual"; -import { ExternalLink, ImageIcon, MessageCircle } from "lucide-react"; +import { ImageIcon, MessageCircle } from "lucide-react"; import ReactMarkdown, { defaultUrlTransform, type Components, @@ -633,10 +633,6 @@ function CitationChip({ citationId: string, ) => void; }): ReactElement { - const pageCitationAssetUrl = getTrimmedCitationField( - citation.pageCitationAssetUrl, - ); - return ( @@ -646,11 +642,7 @@ function CitationChip({ disabled={!onCitationClick || isPending} onClick={() => onCitationClick?.(citation, citationId)} aria-busy={isPending} - className={`inline-flex h-8 max-w-[250px] cursor-pointer items-center border border-primary/20 bg-primary/10 px-3 text-left font-mono text-xs font-semibold leading-none text-primary shadow-[0_1px_0_rgba(15,23,42,0.06)] transition-[background-color,border-color,color,box-shadow,transform] hover:border-primary/35 hover:bg-primary/15 hover:text-primary hover:shadow-[0_0_0_2px_rgba(37,99,235,0.12)] active:translate-y-px active:bg-primary/20 focus:outline-none focus:ring-4 focus:ring-ring/15 focus:ring-offset-2 focus:ring-offset-background disabled:cursor-wait disabled:opacity-75 disabled:hover:border-primary/20 disabled:hover:bg-primary/10 disabled:hover:text-primary disabled:hover:shadow-[0_1px_0_rgba(15,23,42,0.06)] dark:border-transparent dark:bg-[#5c606b] dark:text-[#cfd3dc] dark:shadow-none dark:hover:border-[#8f96a8] dark:hover:bg-[#4f535e] dark:hover:text-white dark:hover:shadow-[0_0_0_2px_rgba(143,150,168,0.22)] dark:active:bg-[#454955] dark:disabled:hover:border-transparent dark:disabled:hover:bg-[#5c606b] dark:disabled:hover:text-[#cfd3dc] dark:disabled:hover:shadow-none ${ - pageCitationAssetUrl - ? "rounded-l-md rounded-r-none" - : "rounded-md" - }`} + className="inline-flex h-8 max-w-[250px] cursor-pointer items-center rounded-md border border-primary/20 bg-primary/10 px-3 text-left font-mono text-xs font-semibold leading-none text-primary shadow-[0_1px_0_rgba(15,23,42,0.06)] transition-[background-color,border-color,color,box-shadow,transform] hover:border-primary/35 hover:bg-primary/15 hover:text-primary hover:shadow-[0_0_0_2px_rgba(37,99,235,0.12)] active:translate-y-px active:bg-primary/20 focus:outline-none focus:ring-4 focus:ring-ring/15 focus:ring-offset-2 focus:ring-offset-background disabled:cursor-wait disabled:opacity-75 disabled:hover:border-primary/20 disabled:hover:bg-primary/10 disabled:hover:text-primary disabled:hover:shadow-[0_1px_0_rgba(15,23,42,0.06)] dark:border-transparent dark:bg-[#5c606b] dark:text-[#cfd3dc] dark:shadow-none dark:hover:border-[#8f96a8] dark:hover:bg-[#4f535e] dark:hover:text-white dark:hover:shadow-[0_0_0_2px_rgba(143,150,168,0.22)] dark:active:bg-[#454955] dark:disabled:hover:border-transparent dark:disabled:hover:bg-[#5c606b] dark:disabled:hover:text-[#cfd3dc] dark:disabled:hover:shadow-none" aria-label={`Open source ${label}`} > {label} @@ -664,28 +656,6 @@ function CitationChip({ {tooltipLabel} - {pageCitationAssetUrl ? ( - - - - - - - - Open page image - - - ) : null} ); } diff --git a/src/components/chunks-panel.test.ts b/src/components/chunks-panel.test.ts index f99a68b..9e5bb4b 100644 --- a/src/components/chunks-panel.test.ts +++ b/src/components/chunks-panel.test.ts @@ -14,6 +14,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { ChunksPanel } from "./chunks-panel"; import { sourceOriginalPreviewRequest } from "./source-original-preview-request"; +const fetchPageAssetPageMock = vi.hoisted(() => vi.fn()); const C = ChunksPanel as React.FC>; const virtualizerScrollResetDelayMs = 150; @@ -29,14 +30,31 @@ vi.mock("react-pdf", () => ({ Page: () => React.createElement("div", { "data-testid": "pdf-page" }), })); +vi.mock("@/domains/workspace/client", () => ({ + workspaceClient: { + fetchPageAssetPage: fetchPageAssetPageMock, + }, +})); + describe("ChunksPanel", () => { beforeEach(() => { shouldFlushVirtualizerTimers = false; + fetchPageAssetPageMock.mockReset(); + fetchPageAssetPageMock.mockResolvedValue({ + pages: [], + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 1, + }, + }); globalThis.ResizeObserver = class ResizeObserver { observe() {} unobserve() {} disconnect() {} }; + Element.prototype.scrollIntoView = vi.fn(); }); afterEach(async () => { @@ -71,6 +89,98 @@ describe("ChunksPanel", () => { expect(screen.getByText(/Showing all parsed chunks from/)).toBeTruthy(); }); + it("renders page assets instead of parsed chunk controls for page-mode documents", async () => { + fetchPageAssetPageMock.mockResolvedValue({ + pages: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 4, + totalPages: 1, + }, + }); + + render( + React.createElement(C, { + chunks: [], + selectedSource: "report.pdf", + selectedSourceView: { + id: "source_1", + title: "report.pdf", + mimeType: "application/pdf", + status: "ready", + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }, + }), + ); + + await waitFor(() => + expect( + screen.getByTestId("page-asset-document-viewer"), + ).toBeTruthy(), + ); + expect(screen.getByRole("heading", { name: "Original File" })).toBeTruthy(); + expect(screen.getByText("Page 4")).toBeTruthy(); + expect(screen.queryByRole("button", { name: "Parsed" })).toBeNull(); + expect(screen.queryByRole("button", { name: "Tree" })).toBeNull(); + expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 1); + }); + + it("loads missing page buckets after focusing a later page asset", async () => { + const user = userEvent.setup(); + fetchPageAssetPageMock.mockImplementation( + async (_sourceId: string, page: number) => ({ + pages: [ + { + pageNumber: page === 3 ? 101 : page, + assetUrl: `https://assets.example/page-${page}.png`, + contentType: "image/png", + }, + ], + pagination: { + page, + pageSize: 50, + total: 120, + totalPages: 3, + }, + }), + ); + + render( + React.createElement(C, { + chunks: [], + selectedSource: "report.pdf", + selectedSourceView: { + id: "source_1", + title: "report.pdf", + mimeType: "application/pdf", + status: "ready", + documentPresentation: { kind: "page-assets", pageCount: 120 }, + }, + focusedPageNumber: 101, + focusedPageRequestId: 1, + }), + ); + + await waitFor(() => + expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 3), + ); + + await user.click(screen.getByRole("button", { name: "Load more pages" })); + + await waitFor(() => + expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 2), + ); + }); + it("defaults parsed chunks into a section tree view", () => { render( React.createElement(C, { diff --git a/src/components/chunks-panel.tsx b/src/components/chunks-panel.tsx index 78cd170..5561c4e 100644 --- a/src/components/chunks-panel.tsx +++ b/src/components/chunks-panel.tsx @@ -41,17 +41,25 @@ import { chunksPanelState } from "@/components/chunks-panel-state"; import { MAX_UPLOAD_MB } from "@/domains/sources/validation"; import { useSourceOriginalPreviewWarmup } from "@/components/source-original-preview-warmup"; import { sourceOriginalPreviewModel } from "@/components/source-original-preview-model"; +import { workspaceClient } from "@/domains/workspace/client"; import type { ParsedChunkView } from "@/domains/chunks/types"; -import type { SourceOriginalFileView, SourceView } from "@/domains/sources/types"; +import type { + SourceOriginalFileView, + SourcePageAssetView, + SourceView, +} from "@/domains/sources/types"; import type { AnalyticsContext } from "@/lib/posthog"; import { cn } from "@/lib/utils"; export type ChunksPanelProps = { chunks: ParsedChunkView[]; selectedSource?: string | null; + selectedSourceView?: SourceView | null; selectedSourceFile?: SourceOriginalFileView | null; focusedChunkId?: string | null; focusedChunkRequestId?: number; + focusedPageNumber?: number | null; + focusedPageRequestId?: number; citationListViewRequestId?: number; isLoading?: boolean; isLoadingMore?: boolean; @@ -77,9 +85,12 @@ type ChunkDisplayModeState = { export function ChunksPanel({ chunks = [], selectedSource = null, + selectedSourceView = null, selectedSourceFile = null, focusedChunkId = null, focusedChunkRequestId = 0, + focusedPageNumber = null, + focusedPageRequestId = 0, citationListViewRequestId = 0, isLoading = false, isLoadingMore = false, @@ -93,6 +104,8 @@ export function ChunksPanel({ analyticsContext, sourceCountSnapshot = 0, }: Partial = {}) { + const isPageAssetSource = + selectedSourceView?.documentPresentation?.kind === "page-assets"; const originalPreviewCacheKey = selectedSourceFile?.url ?? null; const isOriginalPreviewAvailable = sourceOriginalPreviewModel.canPreviewOriginalFile( @@ -240,14 +253,29 @@ export function ChunksPanel({ citationListViewRequestId ? "list" : chunkDisplayModeState.mode; - const headerTitle = focusedChunkId ? "Referenced Chunks" : "Parsed Chunks"; + const headerTitle = isPageAssetSource || visibleView === "original" + ? "Original File" + : focusedChunkId + ? "Referenced Chunks" + : "Parsed Chunks"; const shouldMountOriginalPreview = visibleView === "original" || (originalPreviewCacheKey !== null && mountedOriginalPreviewKey === originalPreviewCacheKey); const isTreeModeVisible = - visibleView === "parsed" && chunkDisplayMode === "tree"; - const headerSubtitle = visibleView === "original" ? ( + !isPageAssetSource && visibleView === "parsed" && chunkDisplayMode === "tree"; + const headerSubtitle = isPageAssetSource ? ( + selectedSource ? ( + <> + Showing page images for{" "} + + {selectedSource} + + + ) : ( + "Select a source to preview its page images." + ) + ) : visibleView === "original" ? ( selectedSource ? ( <> Showing the original file for{" "} @@ -299,14 +327,14 @@ export function ChunksPanel({

- {visibleView === "original" ? "Original File" : headerTitle} + {headerTitle}

{headerSubtitle}

- {visibleView === "parsed" && chunks.length > 0 ? ( + {!isPageAssetSource && visibleView === "parsed" && chunks.length > 0 ? (
) : null} - {hasOriginalView ? ( + {!isPageAssetSource && hasOriginalView ? (
+
+ ) : null} + {pageState.isLoading ? ( +
+ Loading page images... +
+ ) : null} +
+ ); +} + +function PageAssetImage({ + isFocused, + page, + refCallback, +}: { + readonly isFocused: boolean; + readonly page: SourcePageAssetView; + readonly refCallback: (element: HTMLDivElement | null) => void; +}): ReactNode { + const aspectRatio = + page.width && page.height ? `${page.width} / ${page.height}` : undefined; + + return ( +
+
+ Page {page.pageNumber} + {page.contentType} +
+
+ {`Page +
+
+ ); +} + +function getNextPageAssetPageIndex( + loadedPageIndexes: ReadonlySet, + totalPages: number, +): number | null { + for (let pageIndex = 1; pageIndex <= totalPages; pageIndex += 1) { + if (!loadedPageIndexes.has(pageIndex)) return pageIndex; + } + + return null; +} + +function EmptyPageAssets(): ReactNode { + return ( +
+ +

+ No page images available. +

+
+ ); +} + function truncateTreeLabel(value: string): string { const maxLength = 42; if (value.length <= maxLength) return value; diff --git a/src/components/source-row.test.ts b/src/components/source-row.test.ts index 427e09a..97c4bf9 100644 --- a/src/components/source-row.test.ts +++ b/src/components/source-row.test.ts @@ -87,6 +87,27 @@ describe("SourceRow", () => { ); }); + it("labels page-asset documents with page counts", () => { + render( + React.createElement(SourceRow, { + isArchiving: false, + isSelected: false, + onSelect: vi.fn(), + source: { + id: "source_1", + mimeType: "application/pdf", + title: "scan.pdf", + status: "ready", + chunkCount: 4, + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }, + }), + ); + + expect(screen.getByText("Processed · 4 pages")).toBeTruthy(); + expect(screen.queryByText("Processed · 4 chunks")).toBeNull(); + }); + it("shows source archive loading locally", () => { render( React.createElement(SourceRow, { diff --git a/src/components/source-row.tsx b/src/components/source-row.tsx index 9fb87d3..473453f 100644 --- a/src/components/source-row.tsx +++ b/src/components/source-row.tsx @@ -197,6 +197,9 @@ function getReadySourceLabel(source: SourceView): string { function getReadySourceStatusText(source: SourceView): string { if (source.kind === "remote") return getReadySourceLabel(source); + if (source.documentPresentation?.kind === "page-assets") { + return `${getReadySourceLabel(source)} · ${source.documentPresentation.pageCount} pages`; + } if (typeof source.chunkCount !== "number") return getReadySourceLabel(source); return `${getReadySourceLabel(source)} · ${source.chunkCount} chunks`; } diff --git a/src/components/sources-panel.tsx b/src/components/sources-panel.tsx index 70a0285..d1950ef 100644 --- a/src/components/sources-panel.tsx +++ b/src/components/sources-panel.tsx @@ -351,6 +351,7 @@ function getSelectedSourcePage( } function getChunkTreeHref(source: SourceView): string | undefined { + if (source.documentPresentation?.kind === "page-assets") return undefined; return source.documentId ? `/inspect/${encodeURIComponent(source.documentId)}/chunks` : undefined; diff --git a/src/components/workspace-citation-focus.test.ts b/src/components/workspace-citation-focus.test.ts index baa5a1a..bfc4654 100644 --- a/src/components/workspace-citation-focus.test.ts +++ b/src/components/workspace-citation-focus.test.ts @@ -169,6 +169,49 @@ describe("useWorkspaceCitationFocus", () => { expect(result.current.pendingCitationId).toBeNull(); }); + it("focuses page-asset citations without fetching full chunks", async () => { + const fetchChunks = vi.fn(async () => [prefetchedChunk]); + const selectSource = vi.fn(); + const pageAssetSource: SourceView = { + ...readySource, + documentPresentation: { kind: "page-assets", pageCount: 8 }, + }; + const pageCitation: ChatCitationView = { + chunkType: "page", + score: 0.9, + pageCitationAssetUrl: "https://assets.example/page-000004.png", + pageCitationPageNumber: 4, + source: { + documentId: "document_1", + sourceFileName: "Contract.pdf", + sectionPath: "Page 4", + }, + }; + + const { result } = renderHook(() => + useWorkspaceCitationFocus({ + fetchChunks, + onSelectSource: selectSource, + selectedSourceId: null, + sources: [pageAssetSource], + }), + { wrapper: createSWRWrapper }, + ); + + await act(async () => { + await result.current.handleCitationClick(pageCitation, "message_1:0"); + }); + + expect(fetchChunks).not.toHaveBeenCalled(); + expect(selectSource).toHaveBeenCalledWith("source_1"); + expect(result.current.focusedChunk.chunkId).toBeNull(); + expect(result.current.focusedPage).toEqual({ + pageNumber: 4, + requestId: 1, + }); + expect(result.current.citationListViewRequestId).toBe(0); + }); + it("reuses cached chunks for a different source without refetching", async () => { const fetchChunks = vi.fn(async () => [prefetchedChunk]); const selectSource = vi.fn(); diff --git a/src/components/workspace-citation-focus.ts b/src/components/workspace-citation-focus.ts index 357cd4e..c233671 100644 --- a/src/components/workspace-citation-focus.ts +++ b/src/components/workspace-citation-focus.ts @@ -13,6 +13,11 @@ type FocusedChunkState = { readonly requestId: number } +type FocusedPageState = { + readonly pageNumber: number | null + readonly requestId: number +} + type PrefetchedChunksBySourceId = Readonly> type PrefetchedChunksUpdater = ( current: PrefetchedChunksBySourceId, @@ -30,6 +35,7 @@ type WorkspaceCitationFocusInput = { type WorkspaceCitationFocus = { readonly citationListViewRequestId: number readonly focusedChunk: FocusedChunkState + readonly focusedPage: FocusedPageState readonly handleCitationClick: ( citation: ChatCitationView, citationId: string, @@ -61,6 +67,10 @@ export function useWorkspaceCitationFocus({ chunkId: null, requestId: 0, }) + const [focusedPage, setFocusedPage] = useState({ + pageNumber: null, + requestId: 0, + }) const [pendingCitationId, setPendingCitationId] = useState( null, ) @@ -102,6 +112,12 @@ export function useWorkspaceCitationFocus({ }, [], ) + const requestPageFocus = useCallback((pageNumber: number | null): void => { + setFocusedPage((current) => ({ + pageNumber, + requestId: current.requestId + 1, + })) + }, []) const updatePrefetchedChunksBySourceId = useCallback( (updater: PrefetchedChunksUpdater): void => { @@ -122,10 +138,12 @@ export function useWorkspaceCitationFocus({ ) } requestChunkFocus(null) + requestPageFocus(null) }, [ onSelectSource, requestChunkFocus, + requestPageFocus, selectedSourceId, updatePrefetchedChunksBySourceId, ], @@ -192,6 +210,18 @@ export function useWorkspaceCitationFocus({ citation, ) if (!source) return + + if ( + workspaceCitationState.isPageAssetCitationTarget(source, citation) + ) { + if (selectedSourceId !== source.id) onSelectSource(source.id) + requestChunkFocus(null) + requestPageFocus( + workspaceCitationState.getCitationPageNumber(citation), + ) + return + } + setCitationListViewRequestId((current) => current + 1) const loadedChunkId = workspaceCitationState.getLoadedCitationChunkId({ @@ -261,6 +291,7 @@ export function useWorkspaceCitationFocus({ loadAllChunksForSource, onSelectSource, requestChunkFocus, + requestPageFocus, selectedChunks, selectedSourceId, sources, @@ -271,6 +302,7 @@ export function useWorkspaceCitationFocus({ return { citationListViewRequestId, focusedChunk, + focusedPage, handleCitationClick, handleLoadAllChunks, handleLoadMoreChunks, diff --git a/src/components/workspace-citation-state.ts b/src/components/workspace-citation-state.ts index d8331af..cf33dad 100644 --- a/src/components/workspace-citation-state.ts +++ b/src/components/workspace-citation-state.ts @@ -27,6 +27,11 @@ type WorkspaceCitationStateModule = { readonly hasExactCitationTargetHint: ( citation: ChatCitationView, ) => boolean + readonly isPageAssetCitationTarget: ( + source: SourceView, + citation: ChatCitationView, + ) => boolean + readonly getCitationPageNumber: (citation: ChatCitationView) => number | null readonly upsertPrefetchedChunks: ( current: PrefetchedChunksBySourceId, sourceId: string, @@ -79,6 +84,36 @@ function hasExactCitationTargetHint(citation: ChatCitationView): boolean { return true } +function isPageAssetCitationTarget( + source: SourceView, + citation: ChatCitationView, +): boolean { + return ( + source.documentPresentation?.kind === "page-assets" && + (getCitationPageNumber(citation) !== null || + typeof citation.pageCitationAssetUrl === "string") + ) +} + +function getCitationPageNumber(citation: ChatCitationView): number | null { + if ( + typeof citation.pageCitationPageNumber === "number" && + Number.isSafeInteger(citation.pageCitationPageNumber) && + citation.pageCitationPageNumber > 0 + ) { + return citation.pageCitationPageNumber + } + + const sectionPath = citation.source.sectionPath + if (typeof sectionPath !== "string") return null + + const match = /\bpage\s+(\d+)\b/i.exec(sectionPath) + if (!match) return null + + const pageNumber = Number.parseInt(match[1]!, 10) + return Number.isSafeInteger(pageNumber) && pageNumber > 0 ? pageNumber : null +} + function upsertPrefetchedChunks( current: PrefetchedChunksBySourceId, sourceId: string, @@ -120,7 +155,9 @@ function removePrefetchedChunks( export const workspaceCitationState: WorkspaceCitationStateModule = { findCitationSource, getLoadedCitationChunkId, + getCitationPageNumber, hasExactCitationTargetHint, + isPageAssetCitationTarget, upsertPrefetchedChunks, removePrefetchedChunks, } diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index 45db25a..b8675d4 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -161,6 +161,28 @@ describe("useWorkspaceSelectedChunks", () => { expect(result.current.selectedChunks).toEqual([]); }); + it("does not fetch chunk pages for page-asset sources", () => { + const pageAssetSource: SourceView = { + ...readySource, + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }; + + const { result } = renderHook( + () => + useWorkspaceSelectedChunks({ + selectedSourceId: "source_1", + sources: [pageAssetSource], + prefetchedChunksBySourceId: {}, + }), + { wrapper: createSWRWrapper }, + ); + + expect(result.current.selectedSource?.id).toBe("source_1"); + expect(result.current.selectedChunks).toEqual([]); + expect(result.current.isSelectedChunksLoading).toBe(false); + expect(fetchChunkPageMock).not.toHaveBeenCalled(); + }); + it("requests a source refresh after loading an unlocalized remote source", async () => { const onRemoteSourceChunksLoaded = vi.fn(); const remoteSource: SourceView = { diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index f00ac68..0b920de 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -42,7 +42,9 @@ export function useWorkspaceSelectedChunks({ ? prefetchedChunksBySourceId[selectedSourceId] : undefined const selectedChunkSourceId = - selectedSource && selectedSource.status === "ready" + selectedSource && + selectedSource.status === "ready" && + selectedSource.documentPresentation?.kind !== "page-assets" ? selectedSource.id : null const { diff --git a/src/components/workspace-shell-layout.tsx b/src/components/workspace-shell-layout.tsx index 3b9e997..b9ccc6c 100644 --- a/src/components/workspace-shell-layout.tsx +++ b/src/components/workspace-shell-layout.tsx @@ -40,6 +40,11 @@ type FocusedChunkState = { readonly requestId: number } +type FocusedPageState = { + readonly pageNumber: number | null + readonly requestId: number +} + type WorkspaceShellUser = { readonly id: string readonly name: string | null @@ -66,6 +71,7 @@ export type WorkspaceShellLayoutProps = { readonly dashboardUrl?: string readonly desktopPanelWidths: Readonly readonly focusedChunk: FocusedChunkState + readonly focusedPage?: FocusedPageState readonly hasMessages: boolean readonly hasMoreSelectedChunks: boolean readonly contentView: ContentView @@ -146,6 +152,10 @@ export function WorkspaceShellLayout( props.desktopPanelWidths.chat <= workspaceShellState.desktopSidePanelCompactThreshold const isSourcesPanelNarrow = props.desktopPanelWidths.sources < 220 + const selectedSource = props.sources.find( + (source) => source.id === props.selectedSourceId, + ) + const focusedPage = props.focusedPage ?? { pageNumber: null, requestId: 0 } const handleDesktopLayoutRef = useCallback( (element: HTMLDivElement | null): void => { onDesktopLayoutElementChange(element) @@ -265,10 +275,13 @@ export function WorkspaceShellLayout( = {}): Source { describe("countChunksBySourceId", () => { it("counts ready source chunks from the document total", async () => { const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const readChunks = vi.fn(async () => ({ + chunks: [], + totalChunks: 0, + })) const mockClient = { documents: { listChunks }, + knowledge: { readChunks }, } as unknown as Knowhere const { countChunksBySourceId } = await import("./counts") @@ -54,6 +59,14 @@ describe("countChunksBySourceId", () => { ) expect(listChunks).toHaveBeenCalledTimes(1) + expect(readChunks).toHaveBeenCalledWith({ + documentId: "doc_ready", + revisionKey: "job_1", + chunkType: "page", + page: 1, + pageSize: 1, + assetUrlPolicy: "durable", + }) expect(listChunks).toHaveBeenCalledWith("doc_ready", { page: 1, pageSize: 1, @@ -65,8 +78,10 @@ describe("countChunksBySourceId", () => { const listChunks = vi.fn(async () => { throw new Error("temporary outage") }) + const readChunks = vi.fn(async () => ({ chunks: [], totalChunks: 0 })) const mockClient = { documents: { listChunks }, + knowledge: { readChunks }, } as unknown as Knowhere const { countChunksBySourceId } = await import("./counts") @@ -86,8 +101,10 @@ describe("countChunksBySourceId", () => { const listChunks = vi.fn().mockResolvedValue({ pagination: { total: 70 }, }) + const readChunks = vi.fn(async () => ({ chunks: [], totalChunks: 0 })) const mockClient = { documents: { listChunks }, + knowledge: { readChunks }, } as unknown as Knowhere const { countChunksBySourceId } = await import("./counts") @@ -106,6 +123,50 @@ describe("countChunksBySourceId", () => { ) expect(listChunks).not.toHaveBeenCalled() + expect(readChunks).not.toHaveBeenCalled() expect(counts.size).toBe(0) }) }) + +describe("sourceViewOptionsBySourceId", () => { + it("detects page-asset documents from SDK page chunks", async () => { + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const readChunks = vi.fn(async () => ({ + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + metadata: { + pageAssets: [ + { + pageNum: 1, + artifactRef: "pages/page-000001.png", + assetUrl: "https://assets.example/page-000001.png", + }, + ], + }, + }, + ], + totalChunks: 4, + })) + const mockClient = { + documents: { listChunks }, + knowledge: { readChunks }, + } as unknown as Knowhere + + const { sourceViewOptionsBySourceId } = await import("./counts") + + const options = await Effect.runPromise( + sourceViewOptionsBySourceId( + [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], + mockClient, + ), + ) + + expect(options.get("ready")).toEqual({ + chunkCount: 4, + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }) + expect(listChunks).not.toHaveBeenCalled() + }) +}) diff --git a/src/domains/sources/counts.ts b/src/domains/sources/counts.ts index e6e0c9b..110b3ae 100644 --- a/src/domains/sources/counts.ts +++ b/src/domains/sources/counts.ts @@ -2,8 +2,15 @@ import "server-only" import { Effect } from "effect" import type Knowhere from "@ontos-ai/knowhere-sdk" +import type { KnowledgeReadChunk } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" +import type { SourceDocumentPresentation } from "./types" + +type PageAssetDocumentPresentation = Extract< + SourceDocumentPresentation, + { readonly kind: "page-assets" } +> type CountChunksClient = { readonly documents: { @@ -14,9 +21,27 @@ type CountChunksClient = { readonly pagination?: { readonly total?: number } }> } + readonly knowledge: { + readChunks(params: { + readonly documentId: string + readonly revisionKey?: string + readonly chunkType: "page" + readonly page: number + readonly pageSize: number + readonly assetUrlPolicy: "durable" + }): Promise<{ + readonly chunks: readonly KnowledgeReadChunk[] + readonly totalChunks?: number + }> + } } -export const countChunksBySourceId = ( +export type SourceViewOptions = { + readonly chunkCount?: number + readonly documentPresentation?: SourceDocumentPresentation +} + +export const sourceViewOptionsBySourceId = ( sources: readonly Source[], client: Knowhere, ) => @@ -28,15 +53,19 @@ export const countChunksBySourceId = ( source.status === "ready" && source.knowhereDocumentId, ) - if (readySources.length === 0) return new Map() + if (readySources.length === 0) return new Map() const entries = yield* Effect.all( readySources.map((source) => Effect.gen(function* () { - const total = yield* Effect.tryPromise(() => - loadSourceChunkCount(countClient, source.knowhereDocumentId!), - ).pipe(Effect.catchAll(() => Effect.succeed(undefined))) - return [source.id, total] as const + const options = yield* Effect.tryPromise(() => + loadSourceViewOptions(countClient, source), + ).pipe( + Effect.catchAll(() => + Effect.sync((): SourceViewOptions | undefined => undefined), + ), + ) + return [source.id, options] as const }), ), { concurrency: "unbounded" }, @@ -44,26 +73,81 @@ export const countChunksBySourceId = ( return new Map( entries.filter( - (entry): entry is readonly [string, number] => - typeof entry[1] === "number", + (entry): entry is readonly [string, SourceViewOptions] => + entry[1] !== undefined, ), ) }) -export const sourceViewOptionsBySourceId = ( +export const countChunksBySourceId = ( sources: readonly Source[], client: Knowhere, ) => Effect.gen(function* () { - const counts = yield* countChunksBySourceId(sources, client) - return new Map( - sources.map((source) => [ - source.id, - { chunkCount: counts.get(source.id) }, - ]), - ) + const sourceOptions = yield* sourceViewOptionsBySourceId(sources, client) + const countEntries: [string, number][] = [] + for (const [sourceId, options] of sourceOptions.entries()) { + if (typeof options.chunkCount === "number") { + countEntries.push([sourceId, options.chunkCount]) + } + } + return new Map(countEntries) }) +async function loadSourceViewOptions( + client: CountChunksClient, + source: Source, +): Promise { + const documentId = source.knowhereDocumentId + if (!documentId) return undefined + + const pagePresentation = await loadPageAssetPresentation( + client, + documentId, + source.knowhereJobId, + ) + if (pagePresentation) { + return { + chunkCount: pagePresentation.pageCount, + documentPresentation: pagePresentation, + } + } + + const chunkCount = await loadSourceChunkCount(client, documentId) + return typeof chunkCount === "number" ? { chunkCount } : undefined +} + +async function loadPageAssetPresentation( + client: CountChunksClient, + documentId: string, + revisionKey: string | null, +): Promise { + try { + const response = await client.knowledge.readChunks({ + documentId, + ...(revisionKey ? { revisionKey } : {}), + chunkType: "page", + page: 1, + pageSize: 1, + assetUrlPolicy: "durable", + }) + const firstChunk = response.chunks[0] + if (!firstChunk || firstChunk.chunkType !== "page") return undefined + if (!hasUsablePageAssets(firstChunk.metadata.pageAssets)) return undefined + + const pageCount = + typeof response.totalChunks === "number" && + Number.isFinite(response.totalChunks) && + response.totalChunks > 0 + ? response.totalChunks + : 1 + + return { kind: "page-assets", pageCount } + } catch { + return undefined + } +} + async function loadSourceChunkCount( client: CountChunksClient, documentId: string, @@ -75,3 +159,25 @@ async function loadSourceChunkCount( const total = response.pagination?.total return typeof total === "number" && Number.isFinite(total) ? total : undefined } + +function hasUsablePageAssets(value: unknown): boolean { + if (!Array.isArray(value)) return false + + return value.some((item) => { + if (!isRecord(item)) return false + const pageNum = item.pageNum + const artifactRef = item.artifactRef + const assetUrl = item.assetUrl + return ( + typeof pageNum === "number" && + Number.isSafeInteger(pageNum) && + pageNum > 0 && + ((typeof artifactRef === "string" && artifactRef.trim().length > 0) || + (typeof assetUrl === "string" && assetUrl.trim().length > 0)) + ) + }) +} + +function isRecord(value: unknown): value is Readonly> { + return typeof value === "object" && value !== null +} diff --git a/src/domains/sources/page-assets.ts b/src/domains/sources/page-assets.ts new file mode 100644 index 0000000..55c108b --- /dev/null +++ b/src/domains/sources/page-assets.ts @@ -0,0 +1,104 @@ +import "server-only" + +import type { Knowledge } from "@ontos-ai/knowhere-sdk" + +import type { ChunkPageParams } from "@/domains/chunks" +import type { SourcePageAssetView } from "./route-types" + +type ReadableSource = { + readonly documentId: string + readonly revisionKey?: string | null +} + +export type SourcePageAssetsPage = { + readonly pages: readonly SourcePageAssetView[] + readonly pagination: { + readonly page: number + readonly pageSize: number + readonly total: number + readonly totalPages: number + } +} + +export async function readSourcePageAssets(input: { + readonly knowledge: Knowledge + readonly source: ReadableSource + readonly params: ChunkPageParams +}): Promise { + const response = await input.knowledge.readChunks({ + documentId: input.source.documentId, + ...(input.source.revisionKey ? { revisionKey: input.source.revisionKey } : {}), + chunkType: "page", + page: input.params.page, + pageSize: input.params.pageSize, + assetUrlPolicy: "durable", + }) + const pages = response.chunks.flatMap((chunk): SourcePageAssetView[] => + readPageAssetViews(chunk.metadata.pageAssets, chunk.assetUrl), + ) + + return { + pages, + pagination: { + page: response.page ?? input.params.page, + pageSize: response.pageSize ?? input.params.pageSize, + total: response.totalChunks ?? pages.length, + totalPages: + response.totalPages ?? + Math.max(1, Math.ceil(pages.length / input.params.pageSize)), + }, + } +} + +function readPageAssetViews( + value: unknown, + fallbackAssetUrl?: string, +): SourcePageAssetView[] { + if (!Array.isArray(value)) return [] + + return value.flatMap((item): SourcePageAssetView[] => { + if (!isRecord(item)) return [] + const pageNumber = getPositiveInteger(item.pageNum) + const assetUrl = getTrimmedString(item.assetUrl) ?? fallbackAssetUrl + const contentType = getTrimmedString(item.contentType) + if (!pageNumber || !assetUrl || !contentType) return [] + + return [ + { + pageNumber, + assetUrl, + contentType, + ...(getPositiveNumber(item.width) !== undefined + ? { width: getPositiveNumber(item.width) } + : {}), + ...(getPositiveNumber(item.height) !== undefined + ? { height: getPositiveNumber(item.height) } + : {}), + }, + ] + }) +} + +function getTrimmedString(value: unknown): string | undefined { + if (typeof value !== "string") return undefined + const trimmed = value.trim() + return trimmed.length > 0 ? trimmed : undefined +} + +function getPositiveInteger(value: unknown): number | undefined { + return typeof value === "number" && + Number.isSafeInteger(value) && + value > 0 + ? value + : undefined +} + +function getPositiveNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value > 0 + ? value + : undefined +} + +function isRecord(value: unknown): value is Readonly> { + return typeof value === "object" && value !== null +} diff --git a/src/domains/sources/route-page-assets.ts b/src/domains/sources/route-page-assets.ts new file mode 100644 index 0000000..e356702 --- /dev/null +++ b/src/domains/sources/route-page-assets.ts @@ -0,0 +1,161 @@ +import { Effect } from "effect" + +import { readSourcePageAssets } from "./page-assets" +import { routeResult } from "@/lib/route-result" +import { + decodeRemoteSourceId, + findRemoteLibraryDocumentBySourceId, +} from "./remote-library" +import { + getClientForWorkspace, + getKnowledgeForSource, +} from "./route-dependencies" +import { sourceRowRepository } from "./source-row-repository" +import type { + JsonRouteResult, + LoadSourcePageAssetsInput, + SourcePageAssetsBody, + SourceRouteServiceDependencies, +} from "./route-types" + +type RoutePageAssetsDependencies = Pick< + SourceRouteServiceDependencies, + | "ensureApiKeyForWorkspace" + | "ensureWorkspace" + | "getCurrentUser" + | "makeKnowhereClient" + | "sourceService" +> + +type RoutePageAssets = { + readonly loadSourcePageAssets: ( + input: LoadSourcePageAssetsInput, + ) => Promise> +} + +function createRoutePageAssets( + deps: RoutePageAssetsDependencies, +): RoutePageAssets { + return { + loadSourcePageAssets: (input: LoadSourcePageAssetsInput) => + Effect.runPromise(loadSourcePageAssetsEffect(input, deps)), + } +} + +const loadSourcePageAssetsEffect = ( + input: LoadSourcePageAssetsInput, + deps: RoutePageAssetsDependencies, +) => + Effect.gen(function* () { + if (!sourceRowRepository.isWorkspaceSourceId(input.sourceId)) { + const remoteResult = yield* loadRemotePageAssetsEffect(input, deps) + return remoteResult ?? sourceNotFound() + } + + const user = yield* Effect.tryPromise(() => deps.getCurrentUser()) + if (!user) return sourceNotFound() + + const workspace = yield* Effect.tryPromise(() => + deps.ensureWorkspace(user.id), + ) + const source = yield* Effect.tryPromise(() => + deps.sourceService.findInWorkspace(workspace.id, input.sourceId), + ) + if (!source) return sourceNotFound() + if (source.status !== "ready" || !source.knowhereDocumentId) { + return sourceNotReady() + } + + const documentId = source.knowhereDocumentId + const apiKey = yield* Effect.tryPromise(() => + deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), + ) + const knowledge = getKnowledgeForSource({ + apiKey, + workspaceId: workspace.id, + sourceId: source.id, + documentId, + revisionKey: source.knowhereJobId, + }) + const pageAssets = yield* Effect.tryPromise(() => + readSourcePageAssets({ + knowledge, + source: { + documentId, + revisionKey: source.knowhereJobId, + }, + params: input.pageParams, + }), + ) + + return routeResult.ok(pageAssets) + }) + +const loadRemotePageAssetsEffect = ( + input: LoadSourcePageAssetsInput, + deps: RoutePageAssetsDependencies, +) => + Effect.gen(function* () { + if (!decodeRemoteSourceId(input.sourceId)) return null + + const user = yield* Effect.tryPromise(() => deps.getCurrentUser()) + if (!user) return null + + const workspace = yield* Effect.tryPromise(() => + deps.ensureWorkspace(user.id), + ) + const apiKey = yield* Effect.tryPromise(() => + deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), + ) + const client = yield* Effect.tryPromise(() => + getClientForWorkspace(workspace.id, input.cookieHeader, deps), + ) + const remoteDocument = yield* findRemoteLibraryDocumentBySourceId({ + sourceId: input.sourceId, + workspace, + client, + localSources: [], + }) + if (!remoteDocument) return null + + const source = yield* Effect.tryPromise(() => + deps.sourceService.localizeRemoteDocument(workspace.id, { + documentId: remoteDocument.documentId, + namespace: remoteDocument.namespace, + status: remoteDocument.status, + title: remoteDocument.title, + mimeType: remoteDocument.mimeType, + sizeBytes: remoteDocument.sizeBytes, + revisionKey: remoteDocument.revisionKey ?? null, + }), + ) + const documentId = source.knowhereDocumentId ?? remoteDocument.documentId + const revisionKey = + source.knowhereJobId ?? remoteDocument.revisionKey ?? null + const knowledge = getKnowledgeForSource({ + apiKey, + workspaceId: workspace.id, + sourceId: source.id, + documentId, + revisionKey, + }) + const pageAssets = yield* Effect.tryPromise(() => + readSourcePageAssets({ + knowledge, + source: { documentId, revisionKey }, + params: input.pageParams, + }), + ) + + return routeResult.ok(pageAssets) + }) + +function sourceNotFound(): JsonRouteResult<{ readonly message: string }> { + return routeResult.error(404, "Source not found.") +} + +function sourceNotReady(): JsonRouteResult<{ readonly message: string }> { + return routeResult.error(409, "Source is not ready.") +} + +export { createRoutePageAssets } diff --git a/src/domains/sources/route-service.ts b/src/domains/sources/route-service.ts index 9a19bef..7e3b85f 100644 --- a/src/domains/sources/route-service.ts +++ b/src/domains/sources/route-service.ts @@ -4,12 +4,14 @@ import { createRouteArchive } from "./route-archive" import { createRouteChunks } from "./route-chunks" import { createSourceRouteDependencies } from "./route-dependencies" import { createRouteListing } from "./route-listing" +import { createRoutePageAssets } from "./route-page-assets" import { createRouteRetry } from "./route-retry" import { createRouteUpload } from "./route-upload" import type { ArchiveSourceInput, ListSourcesInput, LoadSourceChunksInput, + LoadSourcePageAssetsInput, RetrySourceInput, SourceRouteService, SourceRouteServiceOverrides, @@ -25,6 +27,7 @@ export function createSourceRouteService( const archive = createRouteArchive(deps) const retry = createRouteRetry(deps) const chunks = createRouteChunks(deps) + const pageAssets = createRoutePageAssets(deps) return { listSources: (input: ListSourcesInput) => listing.listSources(input), @@ -33,5 +36,7 @@ export function createSourceRouteService( retrySource: (input: RetrySourceInput) => retry.retrySource(input), loadSourceChunks: (input: LoadSourceChunksInput) => chunks.loadSourceChunks(input), + loadSourcePageAssets: (input: LoadSourcePageAssetsInput) => + pageAssets.loadSourcePageAssets(input), } } diff --git a/src/domains/sources/route-types.ts b/src/domains/sources/route-types.ts index 96bd80c..48401a0 100644 --- a/src/domains/sources/route-types.ts +++ b/src/domains/sources/route-types.ts @@ -4,7 +4,11 @@ import type { ChunkPageParams, } from "@/domains/chunks" import type { ParsedChunkView } from "@/domains/chunks/types" -import type { SourceStatus, SourceView } from "@/domains/sources/types" +import type { + SourcePageAssetView, + SourceStatus, + SourceView, +} from "@/domains/sources/types" import type { AuthUser } from "@/infrastructure/auth" import type { Source, Workspace } from "@/infrastructure/db/schema" import type { @@ -106,6 +110,20 @@ type SourceChunksBody = readonly message: string } +type SourcePageAssetsBody = + | { + readonly pages: readonly SourcePageAssetView[] + readonly pagination: { + readonly page: number + readonly pageSize: number + readonly total: number + readonly totalPages: number + } + } + | { + readonly message: string + } + type ListSourcesInput = { readonly cookieHeader: string } @@ -133,6 +151,12 @@ type LoadSourceChunksInput = { readonly pageParams: ChunkPageParams } +type LoadSourcePageAssetsInput = { + readonly cookieHeader: string + readonly sourceId: string + readonly pageParams: ChunkPageParams +} + type SourceRouteService = { readonly listSources: ( input: ListSourcesInput, @@ -149,6 +173,9 @@ type SourceRouteService = { readonly loadSourceChunks: ( input: LoadSourceChunksInput, ) => Promise> + readonly loadSourcePageAssets: ( + input: LoadSourcePageAssetsInput, + ) => Promise> } type SourceWorkflowService = { @@ -251,9 +278,12 @@ export type { ListSourcesBody, ListSourcesInput, LoadSourceChunksInput, + LoadSourcePageAssetsInput, RetrySourceBody, RetrySourceInput, SourceChunksBody, + SourcePageAssetsBody, + SourcePageAssetView, SourceRouteDemoApi, SourceRouteKnowhereClient, SourceRouteService, diff --git a/src/domains/sources/types.ts b/src/domains/sources/types.ts index 8479f87..e0ff9be 100644 --- a/src/domains/sources/types.ts +++ b/src/domains/sources/types.ts @@ -16,6 +16,18 @@ export type SourceOfficialLibraryView = { readonly sourceUrl: string } +export type SourceDocumentPresentation = + | { readonly kind: "parsed-chunks" } + | { readonly kind: "page-assets"; readonly pageCount: number } + +export type SourcePageAssetView = { + readonly pageNumber: number + readonly assetUrl: string + readonly contentType: string + readonly width?: number + readonly height?: number +} + export type OfficialLibrarySourceView = { readonly librarySourceId: string readonly categoryId: string @@ -50,6 +62,11 @@ export type SourceView = { readonly officialLibrary?: SourceOfficialLibraryView /** Count from the Notebook parsed snapshot manifest when available. */ readonly chunkCount?: number + /** + * Preferred source content presentation. Missing values are treated as + * parsed chunks so older views and cached responses remain compatible. + */ + readonly documentPresentation?: SourceDocumentPresentation /** User opt-out for this query session. Drives excludeDocumentIds. */ readonly excludedFromQuery?: boolean } diff --git a/src/domains/sources/view.test.ts b/src/domains/sources/view.test.ts index cc5b223..0069f6a 100644 --- a/src/domains/sources/view.test.ts +++ b/src/domains/sources/view.test.ts @@ -121,6 +121,18 @@ describe("toSourceView", () => { }); }); + it("includes document presentation options when provided", () => { + expect( + toSourceView(makeSource(), { + chunkCount: 4, + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }), + ).toMatchObject({ + chunkCount: 4, + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }); + }); + it("does not expose legacy demo original proxy routes", () => { const view = toSourceView( makeSource({ diff --git a/src/domains/sources/view.ts b/src/domains/sources/view.ts index 4096a81..5c6fe50 100644 --- a/src/domains/sources/view.ts +++ b/src/domains/sources/view.ts @@ -3,6 +3,7 @@ import { Schema } from "effect"; import type { Source } from "@/infrastructure/db/schema"; import type { SourceView } from "@/domains/sources/types"; import { sourceFailureMessage } from "./failure-message"; +import type { SourceDocumentPresentation } from "./types"; const SourceStatus = Schema.Literal( "uploading", @@ -13,7 +14,10 @@ const SourceStatus = Schema.Literal( export function toSourceView( source: Source, - options: { chunkCount?: number } = {}, + options: { + readonly chunkCount?: number + readonly documentPresentation?: SourceDocumentPresentation + } = {}, ): SourceView { const originalFile = getSourceOriginalFile(source) const status = toSourceStatus(source.status) @@ -35,6 +39,9 @@ export function toSourceView( ...(options.chunkCount !== undefined ? { chunkCount: options.chunkCount } : {}), + ...(options.documentPresentation !== undefined + ? { documentPresentation: options.documentPresentation } + : {}), }; } diff --git a/src/domains/workspace/client.ts b/src/domains/workspace/client.ts index a650b6f..efd8e32 100644 --- a/src/domains/workspace/client.ts +++ b/src/domains/workspace/client.ts @@ -4,7 +4,10 @@ import type { ChatThreadView, } from "@/domains/chat/types" import type { ParsedChunkView } from "@/domains/chunks/types" -import type { SourceView } from "@/domains/sources/types" +import type { + SourcePageAssetView, + SourceView, +} from "@/domains/sources/types" import { workspaceRouteClient } from "./route-client" const workspaceClientKeys = { @@ -34,6 +37,17 @@ type SourceChunksResponse = { } } +type SourcePageAssetsResponse = { + pages?: SourcePageAssetView[] + message?: string + pagination?: { + page: number + pageSize: number + total: number + totalPages: number + } +} + type ChatThreadResponse = { thread?: ChatThreadView messages?: ChatMessageView[] @@ -91,6 +105,7 @@ export const workspaceClient = { keys: workspaceClientKeys, fetchChunks, fetchChunkPage, + fetchPageAssetPage, fetchSources, fetchChatThreads, fetchChatThread, @@ -136,6 +151,25 @@ async function fetchChunkPage( } } +async function fetchPageAssetPage( + sourceId: string, + page: number, +): Promise { + const searchParams = new URLSearchParams({ + page: String(page), + pageSize: String(workspaceClientConfig.sourceChunkPageSize), + }) + const body = await workspaceRouteClient.getJson( + `/api/sources/${encodeURIComponent(sourceId)}/page-assets?${searchParams.toString()}`, + ) + + return { + pages: Array.isArray(body.pages) ? body.pages : [], + ...(typeof body.message === "string" ? { message: body.message } : {}), + pagination: body.pagination, + } +} + async function fetchSources(): Promise { const body = await workspaceRouteClient.getJson( workspaceClientKeys.sources, From d955a8104398fc7af07dbf9f26f685cbf9d54efb Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 15:00:17 +0800 Subject: [PATCH 22/51] Tolerate missing remote display reads --- .../sources/[sourceId]/chunks/route.test.ts | 50 +++++++++++++ .../[sourceId]/page-assets/route.test.ts | 49 +++++++++++++ src/components/chunks-panel.test.ts | 38 ++++++++++ src/components/chunks-panel.tsx | 25 +++++++ .../workspace-selected-chunks.test.ts | 33 +++++++++ src/components/workspace-selected-chunks.ts | 7 +- .../sources/display-read-unavailable.ts | 47 ++++++++++++ src/domains/sources/route-chunks.ts | 72 ++++++++++++++++--- src/domains/sources/route-page-assets.ts | 57 +++++++++++++-- src/domains/sources/route-types.ts | 22 ++++++ src/domains/workspace/client.test.ts | 60 ++++++++++++++++ src/domains/workspace/client.ts | 8 ++- 12 files changed, 451 insertions(+), 17 deletions(-) create mode 100644 src/domains/sources/display-read-unavailable.ts diff --git a/src/app/api/sources/[sourceId]/chunks/route.test.ts b/src/app/api/sources/[sourceId]/chunks/route.test.ts index 85fb1bd..d13677d 100644 --- a/src/app/api/sources/[sourceId]/chunks/route.test.ts +++ b/src/app/api/sources/[sourceId]/chunks/route.test.ts @@ -548,6 +548,56 @@ describe("GET /api/sources/[sourceId]/chunks", () => { expect(mocks.readChunks).not.toHaveBeenCalled() }) + it("marks workspace chunks unavailable when the parsed document is missing remotely", async () => { + const notFoundError = new Error("Document not found") + notFoundError.name = "NotFoundError" + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereJobId: "job_1", + knowhereDocumentId: "doc_missing", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockRejectedValue(notFoundError) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/chunks?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + chunks: [], + pagination: { + page: 1, + pageSize: 1, + total: 0, + totalPages: 0, + }, + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + isUnavailable: true, + }) + expect(response.status).toBe(200) + }) + it("materializes a remote source id and reads chunks through the SDK", async () => { const knowhereClient = { documents: { diff --git a/src/app/api/sources/[sourceId]/page-assets/route.test.ts b/src/app/api/sources/[sourceId]/page-assets/route.test.ts index 4bf5c5e..073a20e 100644 --- a/src/app/api/sources/[sourceId]/page-assets/route.test.ts +++ b/src/app/api/sources/[sourceId]/page-assets/route.test.ts @@ -244,6 +244,55 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { }) expect(response.status).toBe(200) }) + + it("marks page assets unavailable when the parsed document is missing remotely", async () => { + const notFoundError = new Error("Document not found") + notFoundError.name = "NotFoundError" + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereDocumentId: "doc_missing", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockRejectedValue(notFoundError) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + pages: [], + pagination: { + page: 1, + pageSize: 1, + total: 0, + totalPages: 0, + }, + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + isUnavailable: true, + }) + expect(response.status).toBe(200) + }) }) function makeReadySource(overrides: Record) { diff --git a/src/components/chunks-panel.test.ts b/src/components/chunks-panel.test.ts index 9e5bb4b..238ca0c 100644 --- a/src/components/chunks-panel.test.ts +++ b/src/components/chunks-panel.test.ts @@ -181,6 +181,44 @@ describe("ChunksPanel", () => { ); }); + it("shows page asset unavailable messages", async () => { + fetchPageAssetPageMock.mockResolvedValue({ + pages: [], + isUnavailable: true, + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }); + + render( + React.createElement(C, { + chunks: [], + selectedSource: "report.pdf", + selectedSourceView: { + id: "source_1", + title: "report.pdf", + mimeType: "application/pdf", + status: "ready", + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }, + }), + ); + + await waitFor(() => + expect( + screen.getByText( + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + ), + ).toBeTruthy(), + ); + expect(screen.queryByText("No page images available.")).toBeNull(); + }); + it("defaults parsed chunks into a section tree view", () => { render( React.createElement(C, { diff --git a/src/components/chunks-panel.tsx b/src/components/chunks-panel.tsx index 5561c4e..70f6448 100644 --- a/src/components/chunks-panel.tsx +++ b/src/components/chunks-panel.tsx @@ -397,6 +397,8 @@ export function ChunksPanel({ /> ) : isLoading ? ( + ) : chunks.length === 0 && processingMessage ? ( + ) : chunks.length === 0 ? ( selectedSource ? ( @@ -1036,6 +1038,7 @@ type PageAssetPageState = { readonly pages: readonly SourcePageAssetView[]; readonly loadedPageIndexes: ReadonlySet; readonly isLoading: boolean; + readonly message: string | null; readonly totalPages: number; }; @@ -1054,6 +1057,7 @@ function PageAssetDocumentViewer({ pages: [], loadedPageIndexes: new Set(), isLoading: false, + message: null, totalPages: Math.max( 1, Math.ceil( @@ -1092,6 +1096,7 @@ function PageAssetDocumentViewer({ ), loadedPageIndexes, isLoading: false, + message: response.message ?? null, totalPages: response.pagination?.totalPages ?? current.totalPages, }; @@ -1115,6 +1120,7 @@ function PageAssetDocumentViewer({ pages: [], loadedPageIndexes: new Set(), isLoading: false, + message: null, totalPages: Math.max( 1, Math.ceil( @@ -1158,6 +1164,10 @@ function PageAssetDocumentViewer({ return ; } + if (pageState.pages.length === 0 && pageState.message) { + return ; + } + if (pageState.pages.length === 0) { return ; } @@ -1264,6 +1274,21 @@ function EmptyPageAssets(): ReactNode { ); } +function UnavailableSourceMessage({ + message, +}: { + readonly message: string; +}): ReactNode { + return ( +
+ +

+ {message} +

+
+ ); +} + function truncateTreeLabel(value: string): string { const maxLength = 42; if (value.length <= maxLength) return value; diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index b8675d4..b57a788 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -161,6 +161,39 @@ describe("useWorkspaceSelectedChunks", () => { expect(result.current.selectedChunks).toEqual([]); }); + it("surfaces unavailable chunk messages without a loading state", async () => { + fetchChunkPageMock.mockResolvedValue({ + chunks: [], + isUnavailable: true, + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }); + + const { result } = renderHook( + () => + useWorkspaceSelectedChunks({ + selectedSourceId: "source_1", + sources: [readySource], + prefetchedChunksBySourceId: {}, + }), + { wrapper: createSWRWrapper }, + ); + + await waitFor(() => + expect(result.current.selectedChunksMessage).toBe( + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + ), + ); + expect(result.current.isSelectedChunksLoading).toBe(false); + expect(result.current.selectedChunks).toEqual([]); + }); + it("does not fetch chunk pages for page-asset sources", () => { const pageAssetSource: SourceView = { ...readySource, diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index 0b920de..f1a1283 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -68,6 +68,8 @@ export function useWorkspaceSelectedChunks({ }, ) const selectedChunksMessage = getSelectedChunksMessage(selectedChunkPages) + const hasProcessingSelectedChunkPage = + hasProcessingChunkPage(selectedChunkPages) const pagedSelectedChunks = useMemo( () => resolveChunkConnectionTargets( @@ -99,7 +101,7 @@ export function useWorkspaceSelectedChunks({ typeof selectedChunkPages[selectedChunkPageCount - 1] === "undefined", ) const isSelectedChunksLoading = - Boolean(selectedChunksMessage) || + hasProcessingSelectedChunkPage || (selectedChunkSourceId !== null && !prefetchedSelectedChunks && !selectedChunkPages && @@ -150,7 +152,8 @@ function getSelectedChunksMessage( ): string | null { const page = pages?.find( (candidate) => - candidate.isProcessing && typeof candidate.message === "string", + (candidate.isProcessing || candidate.isUnavailable) && + typeof candidate.message === "string", ) return page?.message ?? null } diff --git a/src/domains/sources/display-read-unavailable.ts b/src/domains/sources/display-read-unavailable.ts new file mode 100644 index 0000000..302cd85 --- /dev/null +++ b/src/domains/sources/display-read-unavailable.ts @@ -0,0 +1,47 @@ +const sourceUnavailableMessage: string = + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere." + +function isDisplayReadUnavailableError(error: unknown): boolean { + const details = readErrorDetails(error) + if ( + details.name === "NotFoundError" && + details.message.toLowerCase().includes("document not found") + ) { + return true + } + + const causeDetails = readErrorDetails(details.cause) + return ( + causeDetails.name === "NotFoundError" && + causeDetails.message.toLowerCase().includes("document not found") + ) +} + +function readErrorDetails(error: unknown): { + readonly name: string + readonly message: string + readonly cause: unknown +} { + if (!isRecord(error)) { + return { + name: "", + message: String(error), + cause: undefined, + } + } + + return { + name: typeof error.name === "string" ? error.name : "", + message: typeof error.message === "string" ? error.message : "", + cause: error.cause, + } +} + +function isRecord(value: unknown): value is Readonly> { + return typeof value === "object" && value !== null +} + +export const displayReadUnavailable = { + isError: isDisplayReadUnavailableError, + message: sourceUnavailableMessage, +} as const diff --git a/src/domains/sources/route-chunks.ts b/src/domains/sources/route-chunks.ts index 7d97920..8d9d315 100644 --- a/src/domains/sources/route-chunks.ts +++ b/src/domains/sources/route-chunks.ts @@ -6,6 +6,7 @@ import { resolveChunkConnectionTargets } from "@/domains/chunks" import type { DemoChunkPage } from "@/integrations/knowhere-demo" import { logger } from "@/lib/logger" import { routeResult } from "@/lib/route-result" +import { displayReadUnavailable } from "./display-read-unavailable" import { decodeRemoteSourceId, findRemoteLibraryDocumentBySourceId, @@ -112,20 +113,26 @@ const loadSourceChunksEffect = ( } if (input.shouldLoadAll) { - const chunks = yield* Effect.tryPromise(() => + return yield* Effect.tryPromise(() => readAllSourceChunks({ knowledge, source: readableSource }), + ).pipe( + Effect.map((chunks) => + routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }), + ), + Effect.catchAll((error) => recoverUnavailableChunks(input, error)), ) - return routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }) } - const chunkPage = yield* Effect.tryPromise(() => + return yield* Effect.tryPromise(() => readSourceChunkPage({ knowledge, source: readableSource, params: input.pageParams, }), + ).pipe( + Effect.map((chunkPage) => routeResult.ok(chunkPage)), + Effect.catchAll((error) => recoverUnavailableChunks(input, error)), ) - return routeResult.ok(chunkPage) }) const loadRemoteChunkPageEffect = ( @@ -182,20 +189,26 @@ const loadRemoteChunkPageEffect = ( } if (input.shouldLoadAll) { - const chunks = yield* Effect.tryPromise(() => + return yield* Effect.tryPromise(() => readAllSourceChunks({ knowledge, source: readableSource }), + ).pipe( + Effect.map((chunks) => + routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }), + ), + Effect.catchAll((error) => recoverUnavailableChunks(input, error)), ) - return routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }) } - const chunkPage = yield* Effect.tryPromise(() => + return yield* Effect.tryPromise(() => readSourceChunkPage({ knowledge, source: readableSource, params: input.pageParams, }), + ).pipe( + Effect.map((chunkPage) => routeResult.ok(chunkPage)), + Effect.catchAll((error) => recoverUnavailableChunks(input, error)), ) - return routeResult.ok(chunkPage) }) const loadDemoChunkPageEffect = ( @@ -337,4 +350,47 @@ function sourceSnapshotProcessing( ) } +function sourceChunksUnavailable( + input: LoadSourceChunksInput, +): JsonRouteResult<{ + readonly chunks: [] + readonly pagination?: { + readonly page: number + readonly pageSize: number + readonly total: 0 + readonly totalPages: 0 + } + readonly message: string + readonly isUnavailable: true +}> { + if (input.shouldLoadAll) { + return routeResult.ok({ + chunks: [], + message: displayReadUnavailable.message, + isUnavailable: true, + }) + } + + return routeResult.ok({ + chunks: [], + pagination: { + page: input.pageParams.page, + pageSize: input.pageParams.pageSize, + total: 0, + totalPages: 0, + }, + message: displayReadUnavailable.message, + isUnavailable: true, + }) +} + +function recoverUnavailableChunks( + input: LoadSourceChunksInput, + error: unknown, +): Effect.Effect, unknown> { + return displayReadUnavailable.isError(error) + ? Effect.succeed(sourceChunksUnavailable(input)) + : Effect.fail(error) +} + export { createRouteChunks } diff --git a/src/domains/sources/route-page-assets.ts b/src/domains/sources/route-page-assets.ts index e356702..3077a83 100644 --- a/src/domains/sources/route-page-assets.ts +++ b/src/domains/sources/route-page-assets.ts @@ -2,6 +2,7 @@ import { Effect } from "effect" import { readSourcePageAssets } from "./page-assets" import { routeResult } from "@/lib/route-result" +import { displayReadUnavailable } from "./display-read-unavailable" import { decodeRemoteSourceId, findRemoteLibraryDocumentBySourceId, @@ -77,7 +78,7 @@ const loadSourcePageAssetsEffect = ( documentId, revisionKey: source.knowhereJobId, }) - const pageAssets = yield* Effect.tryPromise(() => + return yield* Effect.tryPromise(() => readSourcePageAssets({ knowledge, source: { @@ -86,9 +87,12 @@ const loadSourcePageAssetsEffect = ( }, params: input.pageParams, }), + ).pipe( + Effect.map((pageAssets) => routeResult.ok(pageAssets)), + Effect.catchAll((error) => + recoverUnavailablePageAssets(input, error), + ), ) - - return routeResult.ok(pageAssets) }) const loadRemotePageAssetsEffect = ( @@ -139,15 +143,18 @@ const loadRemotePageAssetsEffect = ( documentId, revisionKey, }) - const pageAssets = yield* Effect.tryPromise(() => + return yield* Effect.tryPromise(() => readSourcePageAssets({ knowledge, source: { documentId, revisionKey }, params: input.pageParams, }), + ).pipe( + Effect.map((pageAssets) => routeResult.ok(pageAssets)), + Effect.catchAll((error) => + recoverUnavailablePageAssets(input, error), + ), ) - - return routeResult.ok(pageAssets) }) function sourceNotFound(): JsonRouteResult<{ readonly message: string }> { @@ -158,4 +165,42 @@ function sourceNotReady(): JsonRouteResult<{ readonly message: string }> { return routeResult.error(409, "Source is not ready.") } +function sourcePageAssetsUnavailable( + input: LoadSourcePageAssetsInput, +): { + readonly pages: [] + readonly pagination: { + readonly page: number + readonly pageSize: number + readonly total: 0 + readonly totalPages: 0 + } + readonly message: string + readonly isUnavailable: true +} { + return { + pages: [], + pagination: { + page: input.pageParams.page, + pageSize: input.pageParams.pageSize, + total: 0, + totalPages: 0, + }, + message: displayReadUnavailable.message, + isUnavailable: true, + } +} + +function recoverUnavailablePageAssets( + input: LoadSourcePageAssetsInput, + error: unknown, +): Effect.Effect< + JsonRouteResult>, + unknown +> { + return displayReadUnavailable.isError(error) + ? Effect.succeed(routeResult.ok(sourcePageAssetsUnavailable(input))) + : Effect.fail(error) +} + export { createRoutePageAssets } diff --git a/src/domains/sources/route-types.ts b/src/domains/sources/route-types.ts index 48401a0..bd8ea61 100644 --- a/src/domains/sources/route-types.ts +++ b/src/domains/sources/route-types.ts @@ -106,6 +106,17 @@ type SourceChunksBody = readonly chunks: readonly ParsedChunkView[] } | ChunkPage + | { + readonly chunks: readonly [] + readonly pagination?: { + readonly page: number + readonly pageSize: number + readonly total: 0 + readonly totalPages: 0 + } + readonly message: string + readonly isUnavailable: true + } | { readonly message: string } @@ -120,6 +131,17 @@ type SourcePageAssetsBody = readonly totalPages: number } } + | { + readonly pages: readonly [] + readonly pagination: { + readonly page: number + readonly pageSize: number + readonly total: 0 + readonly totalPages: 0 + } + readonly message: string + readonly isUnavailable: true + } | { readonly message: string } diff --git a/src/domains/workspace/client.test.ts b/src/domains/workspace/client.test.ts index bc19bba..0c9df2b 100644 --- a/src/domains/workspace/client.test.ts +++ b/src/domains/workspace/client.test.ts @@ -90,6 +90,66 @@ describe("workspaceClient", () => { }) }) + it("preserves source chunk unavailable messages without processing state", async () => { + mockRouteClient.getJson.mockResolvedValue({ + chunks: [], + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + isUnavailable: true, + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }) + + const page = await workspaceClient.fetchChunkPage("source_1", 1) + + expect(page).toEqual({ + chunks: [], + isUnavailable: true, + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }) + }) + + it("preserves page asset unavailable messages", async () => { + mockRouteClient.getJson.mockResolvedValue({ + pages: [], + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + isUnavailable: true, + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }) + + const page = await workspaceClient.fetchPageAssetPage("source_1", 1) + + expect(page).toEqual({ + pages: [], + isUnavailable: true, + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }) + }) + it("throws materialization route errors instead of treating them as empty sources", async () => { mockRouteClient.postJsonWithStatus.mockResolvedValue({ status: 502, diff --git a/src/domains/workspace/client.ts b/src/domains/workspace/client.ts index efd8e32..5490d33 100644 --- a/src/domains/workspace/client.ts +++ b/src/domains/workspace/client.ts @@ -28,6 +28,7 @@ const workspaceClientConfig = { type SourceChunksResponse = { chunks?: ParsedChunkView[] isProcessing?: boolean + isUnavailable?: boolean message?: string pagination?: { page: number @@ -39,6 +40,7 @@ type SourceChunksResponse = { type SourcePageAssetsResponse = { pages?: SourcePageAssetView[] + isUnavailable?: boolean message?: string pagination?: { page: number @@ -146,7 +148,10 @@ async function fetchChunkPage( return { chunks: Array.isArray(body.chunks) ? body.chunks : [], ...(typeof body.message === "string" ? { message: body.message } : {}), - ...(typeof body.message === "string" ? { isProcessing: true } : {}), + ...(body.isUnavailable === true ? { isUnavailable: true } : {}), + ...(typeof body.message === "string" && body.isUnavailable !== true + ? { isProcessing: true } + : {}), pagination: body.pagination, } } @@ -165,6 +170,7 @@ async function fetchPageAssetPage( return { pages: Array.isArray(body.pages) ? body.pages : [], + ...(body.isUnavailable === true ? { isUnavailable: true } : {}), ...(typeof body.message === "string" ? { message: body.message } : {}), pagination: body.pagination, } From b5e0c2d30413b1bff8073bb2926c70f456836242 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 17:40:49 +0800 Subject: [PATCH 23/51] Lazy detect page asset documents --- .../workspace-selected-chunks.test.ts | 59 +++++++ src/components/workspace-selected-chunks.ts | 78 ++++++++- src/components/workspace-shell-layout.tsx | 7 +- src/components/workspace-shell.tsx | 1 + src/domains/sources/counts.test.ts | 163 ++++++++++++++++++ src/domains/sources/counts.ts | 80 +++++---- src/domains/sources/page-assets.ts | 21 ++- src/domains/sources/route-dependencies.ts | 3 +- src/domains/sources/route-listing.ts | 3 + src/domains/sources/route-service.test.ts | 12 ++ src/domains/sources/route-types.ts | 6 +- src/domains/workspace/client-cache.ts | 4 + src/domains/workspace/initial-state.test.ts | 27 ++- src/domains/workspace/initial-state.ts | 4 + 14 files changed, 423 insertions(+), 45 deletions(-) diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index b57a788..be575fb 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -10,10 +10,12 @@ import type { ParsedChunkView } from "@/domains/chunks/types"; import type { SourceView } from "@/domains/sources/types"; const fetchChunkPageMock = vi.hoisted(() => vi.fn()); +const fetchPageAssetPageMock = vi.hoisted(() => vi.fn()); vi.mock("@/domains/workspace/client", () => ({ workspaceClient: { fetchChunkPage: fetchChunkPageMock, + fetchPageAssetPage: fetchPageAssetPageMock, }, })); @@ -28,6 +30,7 @@ const readySource: SourceView = { describe("useWorkspaceSelectedChunks", () => { beforeEach(() => { fetchChunkPageMock.mockReset(); + fetchPageAssetPageMock.mockReset(); fetchChunkPageMock.mockResolvedValue({ chunks: [], pagination: { @@ -37,6 +40,15 @@ describe("useWorkspaceSelectedChunks", () => { totalPages: 1, }, }); + fetchPageAssetPageMock.mockResolvedValue({ + pages: [], + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }); }); it("returns prefetched chunks while checking the visible page for media", async () => { @@ -59,6 +71,9 @@ describe("useWorkspaceSelectedChunks", () => { { wrapper: createSWRWrapper }, ); + await waitFor(() => + expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 1), + ); await waitFor(() => expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1), ); @@ -213,6 +228,50 @@ describe("useWorkspaceSelectedChunks", () => { expect(result.current.selectedSource?.id).toBe("source_1"); expect(result.current.selectedChunks).toEqual([]); expect(result.current.isSelectedChunksLoading).toBe(false); + expect(fetchPageAssetPageMock).not.toHaveBeenCalled(); + expect(fetchChunkPageMock).not.toHaveBeenCalled(); + }); + + it("detects page assets for selected sources without presentation metadata", async () => { + fetchPageAssetPageMock.mockResolvedValue({ + pages: [ + { + pageNumber: 1, + assetUrl: "https://blob.example/page-1.png", + contentType: "image/png", + }, + { + pageNumber: 20, + assetUrl: "https://blob.example/page-20.png", + contentType: "image/png", + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 1, + totalPages: 1, + }, + }); + + const { result } = renderHook( + () => + useWorkspaceSelectedChunks({ + selectedSourceId: "source_1", + sources: [readySource], + prefetchedChunksBySourceId: {}, + }), + { wrapper: createSWRWrapper }, + ); + + await waitFor(() => + expect(result.current.selectedSource?.documentPresentation).toEqual({ + kind: "page-assets", + pageCount: 20, + }), + ); + expect(result.current.selectedChunks).toEqual([]); + expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 1); expect(fetchChunkPageMock).not.toHaveBeenCalled(); }); diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index f1a1283..acbcd50 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -1,6 +1,6 @@ "use client" -import { useEffect, useMemo, useRef } from "react" +import { useEffect, useMemo, useRef, useState } from "react" import useSWRInfinite from "swr/infinite" import { workspaceClient } from "@/domains/workspace/client" @@ -8,6 +8,7 @@ import { workspaceClientCache, type SourceChunksKey, type SourceChunksResponse, + type SourcePageAssetsResponse, } from "@/domains/workspace/client-cache" import { resolveChunkConnectionTargets } from "@/domains/chunks" import type { ParsedChunkView } from "@/domains/chunks/types" @@ -30,21 +31,51 @@ type WorkspaceSelectedChunks = { readonly selectedSource: SourceView | undefined } +type PageAssetProbeState = + | { readonly status: "page-assets"; readonly pageCount: number } + | { readonly status: "parsed-chunks" } + export function useWorkspaceSelectedChunks({ selectedSourceId, sources, prefetchedChunksBySourceId, onRemoteSourceChunksLoaded, }: WorkspaceSelectedChunksInput): WorkspaceSelectedChunks { - const selectedSource = sources.find((source) => source.id === selectedSourceId) + const rawSelectedSource = sources.find( + (source) => source.id === selectedSourceId, + ) const remoteSourceRefreshRequestedIdsRef = useRef>(new Set()) + const requestedPageAssetProbeIdsRef = useRef>(new Set()) + const [pageAssetProbeBySourceId, setPageAssetProbeBySourceId] = useState< + Readonly> + >({}) + const pageAssetProbeState = rawSelectedSource + ? pageAssetProbeBySourceId[rawSelectedSource.id] + : undefined + const selectedSource = + rawSelectedSource && pageAssetProbeState?.status === "page-assets" + ? { + ...rawSelectedSource, + chunkCount: pageAssetProbeState.pageCount, + documentPresentation: { + kind: "page-assets" as const, + pageCount: pageAssetProbeState.pageCount, + }, + } + : rawSelectedSource const prefetchedSelectedChunks = selectedSourceId ? prefetchedChunksBySourceId[selectedSourceId] : undefined + const shouldProbePageAssets = + rawSelectedSource !== undefined && + rawSelectedSource.status === "ready" && + rawSelectedSource.documentPresentation === undefined && + pageAssetProbeState === undefined const selectedChunkSourceId = selectedSource && selectedSource.status === "ready" && - selectedSource.documentPresentation?.kind !== "page-assets" + selectedSource.documentPresentation?.kind !== "page-assets" && + !shouldProbePageAssets ? selectedSource.id : null const { @@ -101,12 +132,40 @@ export function useWorkspaceSelectedChunks({ typeof selectedChunkPages[selectedChunkPageCount - 1] === "undefined", ) const isSelectedChunksLoading = + shouldProbePageAssets || hasProcessingSelectedChunkPage || (selectedChunkSourceId !== null && !prefetchedSelectedChunks && !selectedChunkPages && isChunksLoading) + useEffect(() => { + const source = rawSelectedSource + if (!source || !shouldProbePageAssets) return + if (requestedPageAssetProbeIdsRef.current.has(source.id)) return + + requestedPageAssetProbeIdsRef.current.add(source.id) + void workspaceClient + .fetchPageAssetPage(source.id, 1) + .then((response) => { + setPageAssetProbeBySourceId((current) => ({ + ...current, + [source.id]: getPageAssetProbeState(response), + })) + + if (source.kind === "remote" && (response.pages?.length ?? 0) > 0) { + onRemoteSourceChunksLoaded?.(source.id) + } + }) + .catch(() => { + requestedPageAssetProbeIdsRef.current.delete(source.id) + setPageAssetProbeBySourceId((current) => ({ + ...current, + [source.id]: { status: "parsed-chunks" }, + })) + }) + }, [onRemoteSourceChunksLoaded, rawSelectedSource, shouldProbePageAssets]) + useEffect(() => { const sourceId = selectedSource?.id if (!sourceId || selectedSource.kind !== "remote") return @@ -133,6 +192,19 @@ export function useWorkspaceSelectedChunks({ } } +function getPageAssetProbeState( + response: SourcePageAssetsResponse, +): PageAssetProbeState { + const pages = response.pages ?? [] + if (pages.length === 0) return { status: "parsed-chunks" } + + const maxPageNumber = Math.max( + ...pages.map((page) => page.pageNumber), + ) + const pageCount = Math.max(response.pagination?.total ?? 0, maxPageNumber) + return { status: "page-assets", pageCount } +} + function fetchChunksByKey([ , sourceId, diff --git a/src/components/workspace-shell-layout.tsx b/src/components/workspace-shell-layout.tsx index b9ccc6c..0a71d90 100644 --- a/src/components/workspace-shell-layout.tsx +++ b/src/components/workspace-shell-layout.tsx @@ -90,6 +90,7 @@ export type WorkspaceShellLayoutProps = { readonly selectedSourceFile: SourceOriginalFileView | null readonly selectedSourceId: string | null readonly selectedSourceTitle: string | null + readonly selectedSourceView?: SourceView readonly sourceTitlesByDocumentId: Readonly> readonly sources: readonly SourceView[] readonly officialLibrarySources: readonly OfficialLibrarySourceView[] @@ -152,9 +153,9 @@ export function WorkspaceShellLayout( props.desktopPanelWidths.chat <= workspaceShellState.desktopSidePanelCompactThreshold const isSourcesPanelNarrow = props.desktopPanelWidths.sources < 220 - const selectedSource = props.sources.find( - (source) => source.id === props.selectedSourceId, - ) + const selectedSource = + props.selectedSourceView ?? + props.sources.find((source) => source.id === props.selectedSourceId) const focusedPage = props.focusedPage ?? { pageNumber: null, requestId: 0 } const handleDesktopLayoutRef = useCallback( (element: HTMLDivElement | null): void => { diff --git a/src/components/workspace-shell.tsx b/src/components/workspace-shell.tsx index 6fc842a..bd45e2f 100644 --- a/src/components/workspace-shell.tsx +++ b/src/components/workspace-shell.tsx @@ -242,6 +242,7 @@ function WorkspaceShellContent({ selectedSourceFile={citationFocus.selectedSource?.originalFile ?? null} selectedSourceId={sourceWorkflow.selectedSourceId} selectedSourceTitle={selectedSourceTitle} + selectedSourceView={citationFocus.selectedSource} sourceTitlesByDocumentId={sourceWorkflow.sourceTitlesByDocumentId} sources={sourceWorkflow.sources} officialLibrarySources={officialLibrarySources ?? []} diff --git a/src/domains/sources/counts.test.ts b/src/domains/sources/counts.test.ts index 652f404..72d132f 100644 --- a/src/domains/sources/counts.test.ts +++ b/src/domains/sources/counts.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it, vi } from "vitest" import { Effect } from "effect" import type Knowhere from "@ontos-ai/knowhere-sdk" +import type { Knowledge } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" @@ -129,6 +130,45 @@ describe("countChunksBySourceId", () => { }) describe("sourceViewOptionsBySourceId", () => { + it("detects page count from many page assets in a single SDK page chunk", async () => { + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const readChunks = vi.fn(async () => ({ + chunks: [ + { + chunkId: "page_bundle", + chunkType: "page", + metadata: { + pageAssets: Array.from({ length: 20 }, (_, index) => ({ + pageNum: index + 1, + artifactRef: `pages/page-${String(index + 1).padStart(6, "0")}.png`, + assetUrl: `https://assets.example/page-${index + 1}.png`, + })), + }, + }, + ], + totalChunks: 1, + })) + const mockClient = { + documents: { listChunks }, + knowledge: { readChunks }, + } as unknown as Knowhere + + const { sourceViewOptionsBySourceId } = await import("./counts") + + const options = await Effect.runPromise( + sourceViewOptionsBySourceId( + [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], + mockClient, + ), + ) + + expect(options.get("ready")).toEqual({ + chunkCount: 20, + documentPresentation: { kind: "page-assets", pageCount: 20 }, + }) + expect(listChunks).not.toHaveBeenCalled() + }) + it("detects page-asset documents from SDK page chunks", async () => { const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) const readChunks = vi.fn(async () => ({ @@ -169,4 +209,127 @@ describe("sourceViewOptionsBySourceId", () => { }) expect(listChunks).not.toHaveBeenCalled() }) + + it("uses a source-specific knowledge reader for presentation detection", async () => { + const defaultReadChunks = vi.fn(async () => ({ chunks: [], totalChunks: 0 })) + const parsedStorageReadChunks = vi.fn(async () => ({ + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + metadata: { + pageAssets: [ + { + pageNum: 1, + artifactRef: "pages/page-000001.png", + assetUrl: "https://assets.example/page-000001.png", + }, + ], + }, + }, + ], + totalChunks: 1, + })) + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const mockClient = { + documents: { listChunks }, + knowledge: { readChunks: defaultReadChunks }, + } as unknown as Knowhere + + const { sourceViewOptionsBySourceId } = await import("./counts") + + const options = await Effect.runPromise( + sourceViewOptionsBySourceId( + [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], + mockClient, + { + getKnowledgeForSource: () => + ({ readChunks: parsedStorageReadChunks }) as unknown as Knowledge, + }, + ), + ) + + expect(options.get("ready")).toEqual({ + chunkCount: 1, + documentPresentation: { kind: "page-assets", pageCount: 1 }, + }) + expect(parsedStorageReadChunks).toHaveBeenCalledWith({ + documentId: "doc_ready", + revisionKey: "job_1", + chunkType: "page", + page: 1, + pageSize: 1, + assetUrlPolicy: "durable", + }) + expect(defaultReadChunks).not.toHaveBeenCalled() + }) + + it("falls back to chunk counts when page presentation detection fails", async () => { + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const readChunks = vi.fn(async () => { + throw new Error("parsed storage and remote unavailable") + }) + const mockClient = { + documents: { listChunks }, + knowledge: { readChunks }, + } as unknown as Knowhere + + const { sourceViewOptionsBySourceId } = await import("./counts") + + const options = await Effect.runPromise( + sourceViewOptionsBySourceId( + [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], + mockClient, + ), + ) + + expect(options.get("ready")).toEqual({ chunkCount: 12 }) + expect(listChunks).toHaveBeenCalledWith("doc_ready", { + page: 1, + pageSize: 1, + }) + }) + + it("skips page presentation reads when detection is disabled", async () => { + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const readChunks = vi.fn(async () => ({ + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + metadata: { + pageAssets: [ + { + pageNum: 1, + artifactRef: "pages/page-000001.png", + assetUrl: "https://assets.example/page-000001.png", + }, + ], + }, + }, + ], + totalChunks: 1, + })) + const mockClient = { + documents: { listChunks }, + knowledge: { readChunks }, + } as unknown as Knowhere + + const { sourceViewOptionsBySourceId } = await import("./counts") + + const options = await Effect.runPromise( + sourceViewOptionsBySourceId( + [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], + mockClient, + { documentPresentationDetection: "disabled" }, + ), + ) + + expect(readChunks).not.toHaveBeenCalled() + expect(listChunks).toHaveBeenCalledWith("doc_ready", { + page: 1, + pageSize: 1, + }) + expect(options.get("ready")).toEqual({ chunkCount: 12 }) + }) }) diff --git a/src/domains/sources/counts.ts b/src/domains/sources/counts.ts index 110b3ae..d161adf 100644 --- a/src/domains/sources/counts.ts +++ b/src/domains/sources/counts.ts @@ -2,7 +2,7 @@ import "server-only" import { Effect } from "effect" import type Knowhere from "@ontos-ai/knowhere-sdk" -import type { KnowledgeReadChunk } from "@ontos-ai/knowhere-sdk" +import type { Knowledge, KnowledgeReadChunk } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" import type { SourceDocumentPresentation } from "./types" @@ -36,6 +36,11 @@ type CountChunksClient = { } } +export type SourceViewOptionsLoadOptions = { + readonly documentPresentationDetection?: "enabled" | "disabled" + readonly getKnowledgeForSource?: (source: Source) => Knowledge +} + export type SourceViewOptions = { readonly chunkCount?: number readonly documentPresentation?: SourceDocumentPresentation @@ -44,6 +49,7 @@ export type SourceViewOptions = { export const sourceViewOptionsBySourceId = ( sources: readonly Source[], client: Knowhere, + options: SourceViewOptionsLoadOptions = {}, ) => Effect.gen(function* () { const countClient = client as unknown as CountChunksClient @@ -58,14 +64,14 @@ export const sourceViewOptionsBySourceId = ( const entries = yield* Effect.all( readySources.map((source) => Effect.gen(function* () { - const options = yield* Effect.tryPromise(() => - loadSourceViewOptions(countClient, source), + const loadedOptions = yield* Effect.tryPromise(() => + loadSourceViewOptions(countClient, source, options), ).pipe( Effect.catchAll(() => Effect.sync((): SourceViewOptions | undefined => undefined), ), ) - return [source.id, options] as const + return [source.id, loadedOptions] as const }), ), { concurrency: "unbounded" }, @@ -97,19 +103,22 @@ export const countChunksBySourceId = ( async function loadSourceViewOptions( client: CountChunksClient, source: Source, + options: SourceViewOptionsLoadOptions, ): Promise { const documentId = source.knowhereDocumentId if (!documentId) return undefined - const pagePresentation = await loadPageAssetPresentation( - client, - documentId, - source.knowhereJobId, - ) - if (pagePresentation) { - return { - chunkCount: pagePresentation.pageCount, - documentPresentation: pagePresentation, + if (options.documentPresentationDetection !== "disabled") { + const pagePresentation = await loadPageAssetPresentation( + client, + source, + options, + ) + if (pagePresentation) { + return { + chunkCount: pagePresentation.pageCount, + documentPresentation: pagePresentation, + } } } @@ -119,13 +128,17 @@ async function loadSourceViewOptions( async function loadPageAssetPresentation( client: CountChunksClient, - documentId: string, - revisionKey: string | null, + source: Source, + options: SourceViewOptionsLoadOptions, ): Promise { + const documentId = source.knowhereDocumentId + if (!documentId) return undefined + try { - const response = await client.knowledge.readChunks({ + const knowledge = options.getKnowledgeForSource?.(source) ?? client.knowledge + const response = await knowledge.readChunks({ documentId, - ...(revisionKey ? { revisionKey } : {}), + ...(source.knowhereJobId ? { revisionKey: source.knowhereJobId } : {}), chunkType: "page", page: 1, pageSize: 1, @@ -133,14 +146,13 @@ async function loadPageAssetPresentation( }) const firstChunk = response.chunks[0] if (!firstChunk || firstChunk.chunkType !== "page") return undefined - if (!hasUsablePageAssets(firstChunk.metadata.pageAssets)) return undefined + const maxPageAssetNumber = getMaxUsablePageAssetNumber( + firstChunk.metadata.pageAssets, + ) + if (!maxPageAssetNumber) return undefined - const pageCount = - typeof response.totalChunks === "number" && - Number.isFinite(response.totalChunks) && - response.totalChunks > 0 - ? response.totalChunks - : 1 + const totalChunks = getPositiveFiniteNumber(response.totalChunks) ?? 0 + const pageCount = Math.max(totalChunks, maxPageAssetNumber) return { kind: "page-assets", pageCount } } catch { @@ -160,22 +172,30 @@ async function loadSourceChunkCount( return typeof total === "number" && Number.isFinite(total) ? total : undefined } -function hasUsablePageAssets(value: unknown): boolean { - if (!Array.isArray(value)) return false +function getMaxUsablePageAssetNumber(value: unknown): number | undefined { + if (!Array.isArray(value)) return undefined - return value.some((item) => { - if (!isRecord(item)) return false + const pageNumbers = value.flatMap((item): number[] => { + if (!isRecord(item)) return [] const pageNum = item.pageNum const artifactRef = item.artifactRef const assetUrl = item.assetUrl - return ( + const isUsable = typeof pageNum === "number" && Number.isSafeInteger(pageNum) && pageNum > 0 && ((typeof artifactRef === "string" && artifactRef.trim().length > 0) || (typeof assetUrl === "string" && assetUrl.trim().length > 0)) - ) + return isUsable ? [pageNum] : [] }) + if (pageNumbers.length === 0) return undefined + return Math.max(...pageNumbers) +} + +function getPositiveFiniteNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value > 0 + ? value + : undefined } function isRecord(value: unknown): value is Readonly> { diff --git a/src/domains/sources/page-assets.ts b/src/domains/sources/page-assets.ts index 55c108b..cfb2765 100644 --- a/src/domains/sources/page-assets.ts +++ b/src/domains/sources/page-assets.ts @@ -36,16 +36,22 @@ export async function readSourcePageAssets(input: { const pages = response.chunks.flatMap((chunk): SourcePageAssetView[] => readPageAssetViews(chunk.metadata.pageAssets, chunk.assetUrl), ) + const maxPageNumber = getMaxPageNumber(pages) + const total = Math.max(response.totalChunks ?? 0, maxPageNumber ?? 0) + const resolvedTotal = total > 0 ? total : pages.length + const computedTotalPages = Math.max( + 1, + Math.ceil(resolvedTotal / input.params.pageSize), + ) + const totalPages = Math.max(response.totalPages ?? 0, computedTotalPages) return { pages, pagination: { page: response.page ?? input.params.page, pageSize: response.pageSize ?? input.params.pageSize, - total: response.totalChunks ?? pages.length, - totalPages: - response.totalPages ?? - Math.max(1, Math.ceil(pages.length / input.params.pageSize)), + total: resolvedTotal, + totalPages, }, } } @@ -99,6 +105,13 @@ function getPositiveNumber(value: unknown): number | undefined { : undefined } +function getMaxPageNumber( + pages: readonly SourcePageAssetView[], +): number | undefined { + if (pages.length === 0) return undefined + return Math.max(...pages.map((page) => page.pageNumber)) +} + function isRecord(value: unknown): value is Readonly> { return typeof value === "object" && value !== null } diff --git a/src/domains/sources/route-dependencies.ts b/src/domains/sources/route-dependencies.ts index a39635f..a51a9b7 100644 --- a/src/domains/sources/route-dependencies.ts +++ b/src/domains/sources/route-dependencies.ts @@ -28,10 +28,11 @@ const defaultDependencies: SourceRouteServiceDependencies = { ensureApiKeyForWorkspace, ensureWorkspace: workspaceService.ensureWorkspace, getCurrentUser, - getSourceViewOptionsBySourceId: (sources, client) => + getSourceViewOptionsBySourceId: (sources, client, options) => getDefaultSourceViewOptionsBySourceId( sources, client as ReturnType, + options, ), makeKnowhereClient: (apiKey: string) => makeDefaultKnowhereClient(apiKey) as SourceRouteKnowhereClient, diff --git a/src/domains/sources/route-listing.ts b/src/domains/sources/route-listing.ts index 35c0c72..08c17b2 100644 --- a/src/domains/sources/route-listing.ts +++ b/src/domains/sources/route-listing.ts @@ -113,6 +113,9 @@ const listSourcesEffect = ( const sourceOptions = yield* deps.getSourceViewOptionsBySourceId( sourcesNeedingChunkCount, client, + { + documentPresentationDetection: "disabled", + }, ) const hiddenDemoSourceIds = new Set( yield* Effect.tryPromise(() => diff --git a/src/domains/sources/route-service.test.ts b/src/domains/sources/route-service.test.ts index a056427..cb9c7da 100644 --- a/src/domains/sources/route-service.test.ts +++ b/src/domains/sources/route-service.test.ts @@ -472,6 +472,9 @@ describe("source route service", () => { expect(getSourceViewOptionsBySourceId).toHaveBeenCalledWith( [source], knowhereClient, + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), ); expect(result).toEqual({ status: 200, @@ -546,6 +549,9 @@ describe("source route service", () => { expect(getSourceViewOptionsBySourceId).toHaveBeenCalledWith( [], knowhereClient, + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), ); expect(result).toEqual({ status: 200, @@ -628,6 +634,9 @@ describe("source route service", () => { expect(getSourceViewOptionsBySourceId).toHaveBeenCalledWith( [], knowhereClient, + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), ); expect(result).toEqual({ status: 200, @@ -700,6 +709,9 @@ describe("source route service", () => { expect(getSourceViewOptionsBySourceId).toHaveBeenCalledWith( [], knowhereClient, + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), ); expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled(); expect(result).toEqual({ diff --git a/src/domains/sources/route-types.ts b/src/domains/sources/route-types.ts index bd8ea61..8ce806d 100644 --- a/src/domains/sources/route-types.ts +++ b/src/domains/sources/route-types.ts @@ -17,7 +17,10 @@ import type { } from "@/integrations/knowhere-demo" import type { RouteResult } from "@/lib/route-result" import type { SourceBlobUploadInput } from "./blob-upload" -import type { sourceViewOptionsBySourceId } from "./counts" +import type { + sourceViewOptionsBySourceId, + SourceViewOptionsLoadOptions, +} from "./counts" import type { UploadKnowhereClient } from "./upload" type SourceRouteKnowhereClient = UploadKnowhereClient & @@ -275,6 +278,7 @@ type SourceRouteServiceDependencies = { readonly getSourceViewOptionsBySourceId: ( sources: readonly Source[], client: SourceRouteKnowhereClient, + options?: SourceViewOptionsLoadOptions, ) => ReturnType readonly makeKnowhereClient: (apiKey: string) => SourceRouteKnowhereClient readonly listSourcesForWorkspace: (workspaceId: string) => Promise diff --git a/src/domains/workspace/client-cache.ts b/src/domains/workspace/client-cache.ts index af5baae..beaa41d 100644 --- a/src/domains/workspace/client-cache.ts +++ b/src/domains/workspace/client-cache.ts @@ -11,6 +11,9 @@ import type { SourceView } from "@/domains/sources/types" type SourceChunksResponse = Awaited< ReturnType > +type SourcePageAssetsResponse = Awaited< + ReturnType +> type ChatThreadDetailResponse = Awaited< ReturnType > @@ -114,4 +117,5 @@ export type { ChatThreadKey, SourceChunksKey, SourceChunksResponse, + SourcePageAssetsResponse, } diff --git a/src/domains/workspace/initial-state.test.ts b/src/domains/workspace/initial-state.test.ts index 23bd50f..9fc1450 100644 --- a/src/domains/workspace/initial-state.test.ts +++ b/src/domains/workspace/initial-state.test.ts @@ -172,6 +172,9 @@ describe("loadWorkspaceShellInitialState", () => { expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith( [source], expect.any(Object), + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), ) expect(state.sources).toEqual([ { @@ -203,7 +206,13 @@ describe("loadWorkspaceShellInitialState", () => { const state = await loadWorkspaceShellInitialState(deps) - expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith([], expect.any(Object)) + expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith( + [], + expect.any(Object), + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), + ) expect(state.sources).toEqual([ expect.objectContaining({ id: "source_demo", @@ -231,7 +240,13 @@ describe("loadWorkspaceShellInitialState", () => { const state = await loadWorkspaceShellInitialState(deps) - expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith([], expect.any(Object)) + expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith( + [], + expect.any(Object), + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), + ) expect(state.sources).toEqual([ expect.objectContaining({ id: "demo-tsla-q4-2025", @@ -259,7 +274,13 @@ describe("loadWorkspaceShellInitialState", () => { const state = await loadWorkspaceShellInitialState(deps) - expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith([], expect.any(Object)) + expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith( + [], + expect.any(Object), + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), + ) expect(state.sources).toEqual([ expect.objectContaining({ id: "demo-tsla-q4-2025", diff --git a/src/domains/workspace/initial-state.ts b/src/domains/workspace/initial-state.ts index 851a34b..f959ad9 100644 --- a/src/domains/workspace/initial-state.ts +++ b/src/domains/workspace/initial-state.ts @@ -149,6 +149,7 @@ type WorkspaceShellInitialStateDependencies = { readonly sourceViewOptionsBySourceId: ( sources: readonly Source[], client: WorkspaceShellInitialStateClient, + options?: Parameters[2], ) => ReturnType } @@ -356,6 +357,9 @@ export const loadWorkspaceShellInitialStateEffect = ( deps.sourceViewOptionsBySourceId( sourcesNeedingChunkCount, client, + { + documentPresentationDetection: "disabled", + }, ), ) From f6198d93543d1962ab2a0b698325da589176c11f Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 18:17:14 +0800 Subject: [PATCH 24/51] Make page asset viewing non-blocking --- .../sources/[sourceId]/chunks/route.test.ts | 125 ++++++++++++- .../[sourceId]/page-assets/route.test.ts | 153 +++++++++++++++- src/components/chunks-panel.test.ts | 42 +++++ src/components/chunks-panel.tsx | 41 ++++- src/domains/chunks/index.ts | 1 + src/domains/chunks/read.test.ts | 109 +++++++++++- src/domains/chunks/read.ts | 167 ++++++++++++++++-- src/domains/sources/page-assets.ts | 111 +++++++++++- src/domains/sources/route-chunks.ts | 24 ++- src/domains/sources/route-dependencies.ts | 15 +- src/domains/sources/route-page-assets.ts | 12 +- 11 files changed, 743 insertions(+), 57 deletions(-) diff --git a/src/app/api/sources/[sourceId]/chunks/route.test.ts b/src/app/api/sources/[sourceId]/chunks/route.test.ts index d13677d..f410285 100644 --- a/src/app/api/sources/[sourceId]/chunks/route.test.ts +++ b/src/app/api/sources/[sourceId]/chunks/route.test.ts @@ -8,6 +8,7 @@ const mocks = vi.hoisted(() => ({ fetchDemoChunkPage: vi.fn(), findSourceInWorkspace: vi.fn(), getCurrentUser: vi.fn(), + listChunks: vi.fn(), localizeRemoteDocument: vi.fn(), makeKnowhereClient: vi.fn(), makeKnowhereClientWithParsedStorage: vi.fn(), @@ -64,7 +65,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { beforeEach(() => { vi.clearAllMocks() mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ - client: { documents: { listChunks: vi.fn() } }, + client: { documents: { listChunks: mocks.listChunks } }, knowledge: { readChunks: mocks.readChunks }, }) }) @@ -421,7 +422,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { } }) - it("loads authenticated workspace chunks through the SDK durable read", async () => { + it("loads authenticated workspace chunks from parsed storage without durable hardening", async () => { mocks.getCurrentUser.mockResolvedValue({ id: "user_1", email: null, @@ -442,7 +443,10 @@ describe("GET /api/sources/[sourceId]/chunks", () => { ) mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") mocks.readChunks.mockResolvedValue({ - document: { localDocumentId: "doc_1" }, + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [ { position: 1, @@ -498,7 +502,114 @@ describe("GET /api/sources/[sourceId]/chunks", () => { revisionKey: "job_1", page: 1, pageSize: 1, - assetUrlPolicy: "durable", + }) + expect(mocks.listChunks).not.toHaveBeenCalled() + }) + + it("falls back to raw Knowhere asset URLs for workspace chunk display when storage misses", async () => { + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockResolvedValue({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "remote:doc_1", + }, + chunks: [ + { + position: 1, + chunkId: "parser_1", + chunkType: "image", + content: "Workspace chunk", + readableContent: "Workspace chunk", + sectionPath: "Summary", + sourceChunkPath: "Summary", + filePath: "images/chart.png", + metadata: {}, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, + }) + mocks.listChunks.mockResolvedValue({ + chunks: [ + { + id: "document_chunk_1", + chunkId: "parser_1", + chunkType: "image", + content: "Workspace chunk", + sectionId: null, + sectionPath: "Summary", + sourceChunkPath: "Summary", + filePath: "images/chart.png", + sortOrder: 0, + metadata: {}, + assetUrl: "https://knowhere.example/assets/chart.png", + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + }) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/chunks?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toMatchObject({ + chunks: [ + { + parserChunkId: "parser_1", + documentId: "doc_1", + assetUrl: "https://knowhere.example/assets/chart.png", + sourceTitle: "notes.pdf", + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + }, + }) + expect(response.status).toBe(200) + expect(mocks.readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "job_1", + page: 1, + pageSize: 1, + }) + expect(mocks.listChunks).toHaveBeenCalledWith("doc_1", { + page: 1, + pageSize: 1, + includeAssetUrls: true, }) }) @@ -641,7 +752,10 @@ describe("GET /api/sources/[sourceId]/chunks", () => { }), ) mocks.readChunks.mockResolvedValue({ - document: { localDocumentId: "doc_remote" }, + document: { + localDocumentId: "doc_remote", + resultDirectoryPath: "parsed-storage:doc_remote", + }, chunks: [ { position: 1, @@ -710,7 +824,6 @@ describe("GET /api/sources/[sourceId]/chunks", () => { revisionKey: "job_result_1", page: 1, pageSize: 1, - assetUrlPolicy: "durable", }) }) }) diff --git a/src/app/api/sources/[sourceId]/page-assets/route.test.ts b/src/app/api/sources/[sourceId]/page-assets/route.test.ts index 073a20e..76ef16a 100644 --- a/src/app/api/sources/[sourceId]/page-assets/route.test.ts +++ b/src/app/api/sources/[sourceId]/page-assets/route.test.ts @@ -7,6 +7,7 @@ const mocks = vi.hoisted(() => ({ ensureWorkspace: vi.fn(), findSourceInWorkspace: vi.fn(), getCurrentUser: vi.fn(), + listChunks: vi.fn(), makeKnowhereClient: vi.fn(), makeKnowhereClientWithParsedStorage: vi.fn(), readChunks: vi.fn(), @@ -62,12 +63,12 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { beforeEach(() => { vi.clearAllMocks() mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ - client: { documents: { listChunks: vi.fn() } }, + client: { documents: { listChunks: mocks.listChunks } }, knowledge: { readChunks: mocks.readChunks }, }) }) - it("returns durable page assets for a ready workspace source", async () => { + it("returns stored page assets for a ready workspace source without durable hardening", async () => { mocks.getCurrentUser.mockResolvedValue({ id: "user_1", email: null, @@ -88,6 +89,10 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { ) mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") mocks.readChunks.mockResolvedValue({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [ { chunkId: "page_1", @@ -147,8 +152,8 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { chunkType: "page", page: 1, pageSize: 1, - assetUrlPolicy: "durable", }) + expect(mocks.listChunks).not.toHaveBeenCalled() }) it("rejects non-ready workspace sources", async () => { @@ -209,6 +214,10 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { ) mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") mocks.readChunks.mockResolvedValue({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "remote:doc_1", + }, chunks: [ { chunkId: "page_1", @@ -221,6 +230,15 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { totalChunks: 1, totalPages: 1, }) + mocks.listChunks.mockResolvedValue({ + chunks: [], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + }) const response = await GET( new NextRequest( @@ -243,6 +261,135 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { }, }) expect(response.status).toBe(200) + expect(mocks.listChunks).toHaveBeenCalledWith("doc_1", { + page: 1, + pageSize: 1, + chunkType: "page", + includeAssetUrls: true, + }) + }) + + it("returns Knowhere page asset URLs when the storage probe falls through to remote chunks", async () => { + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereDocumentId: "doc_1", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockResolvedValue({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "remote:doc_1", + }, + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + metadata: { + pageAssets: [ + { + pageNum: 1, + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, + }) + mocks.listChunks.mockResolvedValue({ + chunks: [ + { + id: "document_page_1", + chunkId: "page_1", + chunkType: "page", + content: "Page 1", + sectionId: null, + sectionPath: "pages/1", + sourceChunkPath: "pages/1", + filePath: "pages/page-000001.png", + sortOrder: 0, + metadata: { + pageAssets: [ + { + pageNum: 1, + assetUrl: "https://knowhere.example/page-000001.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + assetUrl: "https://knowhere.example/page-000001.png", + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + }) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + pages: [ + { + pageNumber: 1, + assetUrl: "https://knowhere.example/page-000001.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + }) + expect(response.status).toBe(200) + expect(mocks.readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "job_1", + chunkType: "page", + page: 1, + pageSize: 1, + }) + expect(mocks.listChunks).toHaveBeenCalledWith("doc_1", { + page: 1, + pageSize: 1, + chunkType: "page", + includeAssetUrls: true, + }) }) it("marks page assets unavailable when the parsed document is missing remotely", async () => { diff --git a/src/components/chunks-panel.test.ts b/src/components/chunks-panel.test.ts index 238ca0c..8e06155 100644 --- a/src/components/chunks-panel.test.ts +++ b/src/components/chunks-panel.test.ts @@ -134,6 +134,48 @@ describe("ChunksPanel", () => { expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 1); }); + it("shows a page-level placeholder when a page asset image fails to load", async () => { + fetchPageAssetPageMock.mockResolvedValue({ + pages: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 4, + totalPages: 1, + }, + }); + + render( + React.createElement(C, { + chunks: [], + selectedSource: "report.pdf", + selectedSourceView: { + id: "source_1", + title: "report.pdf", + mimeType: "application/pdf", + status: "ready", + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }, + }), + ); + + const pageImage = await screen.findByRole("img", { name: "Page 4" }); + fireEvent.error(pageImage); + + expect( + screen.getByTestId("page-asset-image-unavailable-4"), + ).toBeTruthy(); + expect(screen.getByText("Page image unavailable.")).toBeTruthy(); + }); + it("loads missing page buckets after focusing a later page asset", async () => { const user = userEvent.setup(); fetchPageAssetPageMock.mockImplementation( diff --git a/src/components/chunks-panel.tsx b/src/components/chunks-panel.tsx index 70f6448..fc785a2 100644 --- a/src/components/chunks-panel.tsx +++ b/src/components/chunks-panel.tsx @@ -19,6 +19,7 @@ import { } from "d3-hierarchy"; import { FilePlus2, + ImageOff, Layers, RotateCcw, UploadCloud, @@ -1221,6 +1222,8 @@ function PageAssetImage({ readonly page: SourcePageAssetView; readonly refCallback: (element: HTMLDivElement | null) => void; }): ReactNode { + const [failedAssetUrl, setFailedAssetUrl] = useState(null); + const hasImageError = failedAssetUrl === page.assetUrl; const aspectRatio = page.width && page.height ? `${page.width} / ${page.height}` : undefined; @@ -1241,17 +1244,43 @@ function PageAssetImage({ className="overflow-hidden rounded-md bg-muted/30" style={aspectRatio ? { aspectRatio } : undefined} > - {`Page + {hasImageError ? ( + + ) : ( + // eslint-disable-next-line @next/next/no-img-element -- Page assets can be short-lived Knowhere URLs outside Next image optimization. + {`Page setFailedAssetUrl(page.assetUrl)} + /> + )} ); } +function PageAssetImageUnavailable({ + pageNumber, +}: { + readonly pageNumber: number; +}): ReactNode { + return ( +
+
+ +
+

+ Page image unavailable. +

+
+ ); +} + function getNextPageAssetPageIndex( loadedPageIndexes: ReadonlySet, totalPages: number, diff --git a/src/domains/chunks/index.ts b/src/domains/chunks/index.ts index eca1830..e78293e 100644 --- a/src/domains/chunks/index.ts +++ b/src/domains/chunks/index.ts @@ -17,6 +17,7 @@ export type ChunkKnowhereClient = { params: { page: number pageSize: number + chunkType?: "text" | "image" | "table" | "page" includeAssetUrls: boolean }, ): Promise<{ diff --git a/src/domains/chunks/read.test.ts b/src/domains/chunks/read.test.ts index e3c1000..9ba2203 100644 --- a/src/domains/chunks/read.test.ts +++ b/src/domains/chunks/read.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it, vi } from "vitest" -import type { Knowledge } from "@ontos-ai/knowhere-sdk" +import type { DocumentChunk, Knowledge } from "@ontos-ai/knowhere-sdk" import { readAllSourceChunks, readSourceChunkPage } from "./read" @@ -20,9 +20,13 @@ function makeReadChunk(overrides: Record = {}) { } describe("readSourceChunkPage", () => { - it("reads a page with durable asset URLs and maps chunks to the view model", async () => { + it("uses parsed-storage chunks without durable asset hardening", async () => { + const listChunks = vi.fn() const readChunks = vi.fn(async () => ({ - document: { localDocumentId: "doc_1" }, + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [ makeReadChunk({ chunkType: "image", @@ -38,6 +42,7 @@ describe("readSourceChunkPage", () => { const knowledge = { readChunks } as unknown as Knowledge const result = await readSourceChunkPage({ + client: { documents: { listChunks } }, knowledge, source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, params: { page: 2, pageSize: 50 }, @@ -48,8 +53,8 @@ describe("readSourceChunkPage", () => { revisionKey: "rev_1", page: 2, pageSize: 50, - assetUrlPolicy: "durable", }) + expect(listChunks).not.toHaveBeenCalled() expect(result.pagination).toEqual({ page: 2, pageSize: 50, @@ -64,9 +69,68 @@ describe("readSourceChunkPage", () => { }) }) + it("falls back to Knowhere asset URLs when the parsed-storage probe reads remote chunks", async () => { + const readChunks = vi.fn(async () => ({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "remote:doc_1", + }, + chunks: [makeReadChunk({ chunkType: "image", filePath: "images/a.png" })], + page: 1, + pageSize: 50, + totalChunks: 1, + totalPages: 1, + })) + const listChunks = vi.fn(async () => ({ + chunks: [ + makeRemoteDocumentChunk({ + id: "document_chunk_1", + chunkId: "parser_1", + chunkType: "image", + content: "Body", + sectionId: null, + sectionPath: "Summary", + sourceChunkPath: "Summary", + filePath: "images/a.png", + sortOrder: 0, + metadata: {}, + assetUrl: "https://knowhere.example/assets/a.png", + }), + ], + pagination: { page: 1, pageSize: 50, total: 1, totalPages: 1 }, + })) + const knowledge = { readChunks } as unknown as Knowledge + + const result = await readSourceChunkPage({ + client: { documents: { listChunks } }, + knowledge, + source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, + params: { page: 1, pageSize: 50 }, + }) + + expect(readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "rev_1", + page: 1, + pageSize: 50, + }) + expect(listChunks).toHaveBeenCalledWith("doc_1", { + page: 1, + pageSize: 50, + includeAssetUrls: true, + }) + expect(result.chunks[0]?.assetUrl).toBe( + "https://knowhere.example/assets/a.png", + ) + }) + it("omits revisionKey when the source has none", async () => { + const listChunks = vi.fn() const readChunks = vi.fn(async () => ({ - document: { localDocumentId: "doc_1" }, + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [], page: 1, pageSize: 50, @@ -76,6 +140,7 @@ describe("readSourceChunkPage", () => { const knowledge = { readChunks } as unknown as Knowledge await readSourceChunkPage({ + client: { documents: { listChunks } }, knowledge, source: { documentId: "doc_1", title: "notes.pdf", revisionKey: null }, params: { page: 1, pageSize: 50 }, @@ -85,17 +150,39 @@ describe("readSourceChunkPage", () => { documentId: "doc_1", page: 1, pageSize: 50, - assetUrlPolicy: "durable", }) }) }) +function makeRemoteDocumentChunk( + overrides: Partial = {}, +): DocumentChunk { + return { + id: "document_chunk_1", + chunkId: "parser_1", + chunkType: "text", + content: "Body", + sectionId: null, + sectionPath: "Summary", + sourceChunkPath: "Summary", + filePath: null, + sortOrder: 0, + metadata: {}, + assetUrl: null, + ...overrides, + } +} + describe("readAllSourceChunks", () => { it("pages the SDK to exhaustion", async () => { + const listChunks = vi.fn() const readChunks = vi .fn() .mockResolvedValueOnce({ - document: { localDocumentId: "doc_1" }, + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [makeReadChunk({ chunkId: "c1" })], page: 1, pageSize: 200, @@ -103,7 +190,10 @@ describe("readAllSourceChunks", () => { totalPages: 2, }) .mockResolvedValueOnce({ - document: { localDocumentId: "doc_1" }, + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [makeReadChunk({ chunkId: "c2" })], page: 2, pageSize: 200, @@ -113,6 +203,7 @@ describe("readAllSourceChunks", () => { const knowledge = { readChunks } as unknown as Knowledge const chunks = await readAllSourceChunks({ + client: { documents: { listChunks } }, knowledge, source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, }) @@ -123,8 +214,8 @@ describe("readAllSourceChunks", () => { revisionKey: "rev_1", page: 1, pageSize: 200, - assetUrlPolicy: "durable", }) + expect(listChunks).not.toHaveBeenCalled() expect(chunks.map((chunk) => chunk.parserChunkId)).toEqual(["c1", "c2"]) }) }) diff --git a/src/domains/chunks/read.ts b/src/domains/chunks/read.ts index 1bf4790..ae474cf 100644 --- a/src/domains/chunks/read.ts +++ b/src/domains/chunks/read.ts @@ -1,25 +1,61 @@ import "server-only" -import type { Knowledge } from "@ontos-ai/knowhere-sdk" +import type { + DocumentChunk, + Knowledge, + KnowledgeReadChunk, + KnowledgeReadResponse, +} from "@ontos-ai/knowhere-sdk" -import { toParsedChunkViewFromReadChunk, type ChunkPage, type ChunkPageParams } from "@/domains/chunks" +import { + toParsedChunkView, + toParsedChunkViewFromReadChunk, + type ChunkPage, + type ChunkPageParams, +} from "@/domains/chunks" import type { ParsedChunkView } from "@/domains/chunks/types" const loadAllPageSize = 200 +type DisplayReadChunkType = "text" | "image" | "table" | "page" + type ReadableSource = { readonly documentId: string readonly title: string readonly revisionKey?: string | null } +type DisplayReadClient = { + readonly documents: { + listChunks( + documentId: string, + params: { + readonly page: number + readonly pageSize: number + readonly chunkType?: DisplayReadChunkType + readonly includeAssetUrls: true + }, + ): Promise<{ + readonly chunks: readonly DocumentChunk[] + readonly pagination?: { + readonly page?: number + readonly pageSize?: number + readonly total?: number + readonly totalPages?: number + } + }> + } +} + /** * Read a single display page of parsed chunks through the SDK. The SDK serves - * from configured Blob storage when fresh and falls back to Knowhere remote - * otherwise. Display reads request durable asset URLs so media chunks remain - * usable whether the page came from Blob or the remote Knowhere fallback. + * from configured Blob storage when fresh and falls back to Knowhere remote. + * If the probe falls through to remote without asset URLs, fetch the same page + * from Knowhere with short-lived asset URLs instead of hardening assets into + * Notebook storage. */ export async function readSourceChunkPage(input: { + readonly client: DisplayReadClient readonly knowledge: Knowledge readonly source: ReadableSource readonly params: ChunkPageParams @@ -29,22 +65,59 @@ export async function readSourceChunkPage(input: { ...(input.source.revisionKey ? { revisionKey: input.source.revisionKey } : {}), page: input.params.page, pageSize: input.params.pageSize, - assetUrlPolicy: "durable", }) + if (shouldUseKnowledgeChunkResponse(response)) { + return toChunkPageFromKnowledgeResponse(response, input.source, input.params) + } + + const remoteResponse = await input.client.documents.listChunks( + input.source.documentId, + { + page: input.params.page, + pageSize: input.params.pageSize, + includeAssetUrls: true, + }, + ) + const chunks = remoteResponse.chunks.map((chunk) => + toParsedChunkView( + chunk, + input.source.title, + input.source.documentId, + ), + ) + + return { + chunks, + pagination: { + page: remoteResponse.pagination?.page ?? input.params.page, + pageSize: remoteResponse.pagination?.pageSize ?? input.params.pageSize, + total: remoteResponse.pagination?.total ?? chunks.length, + totalPages: + remoteResponse.pagination?.totalPages ?? + Math.max(1, Math.ceil(chunks.length / input.params.pageSize)), + }, + } +} + +function toChunkPageFromKnowledgeResponse( + response: KnowledgeReadResponse, + source: ReadableSource, + params: ChunkPageParams, +): ChunkPage { const chunks = response.chunks.map((chunk) => - toParsedChunkViewFromReadChunk(chunk, input.source.title, input.source.documentId), + toParsedChunkViewFromReadChunk(chunk, source.title, source.documentId), ) return { chunks, pagination: { - page: response.page ?? input.params.page, - pageSize: response.pageSize ?? input.params.pageSize, + page: response.page ?? params.page, + pageSize: response.pageSize ?? params.pageSize, total: response.totalChunks ?? chunks.length, totalPages: response.totalPages ?? - Math.max(1, Math.ceil(chunks.length / input.params.pageSize)), + Math.max(1, Math.ceil(chunks.length / params.pageSize)), }, } } @@ -54,6 +127,7 @@ export async function readSourceChunkPage(input: { * the tree view and load-all display mode. */ export async function readAllSourceChunks(input: { + readonly client: DisplayReadClient readonly knowledge: Knowledge readonly source: ReadableSource }): Promise { @@ -69,20 +143,85 @@ export async function readAllSourceChunks(input: { : {}), page, pageSize: loadAllPageSize, - assetUrlPolicy: "durable", }) - for (const chunk of response.chunks) { + if (shouldUseKnowledgeChunkResponse(response)) { + for (const chunk of response.chunks) { + chunks.push( + toParsedChunkViewFromReadChunk( + chunk, + input.source.title, + input.source.documentId, + ), + ) + } + totalPages = Math.max(1, response.totalPages ?? 1) + page += 1 + continue + } + + const remoteResponse = await input.client.documents.listChunks( + input.source.documentId, + { + page, + pageSize: loadAllPageSize, + includeAssetUrls: true, + }, + ) + for (const chunk of remoteResponse.chunks) { chunks.push( - toParsedChunkViewFromReadChunk( + toParsedChunkView( chunk, input.source.title, input.source.documentId, ), ) } - totalPages = Math.max(1, response.totalPages ?? 1) + totalPages = Math.max(1, remoteResponse.pagination?.totalPages ?? 1) page += 1 } while (page <= totalPages) return chunks } + +function shouldUseKnowledgeChunkResponse( + response: KnowledgeReadResponse, +): boolean { + return isParsedStorageReadResponse(response) || hasUsableReadChunkAssetUrl(response.chunks) +} + +function isParsedStorageReadResponse(response: KnowledgeReadResponse): boolean { + const resultDirectoryPath = response.document.resultDirectoryPath + return ( + typeof resultDirectoryPath === "string" && + resultDirectoryPath.startsWith("parsed-storage:") + ) +} + +function hasUsableReadChunkAssetUrl( + chunks: readonly KnowledgeReadChunk[], +): boolean { + return chunks.some( + (chunk) => + hasNonEmptyString(chunk.assetUrl) || + hasMetadataPageAssetUrl(chunk.metadata), + ) +} + +function hasMetadataPageAssetUrl( + metadata: Readonly>, +): boolean { + const value = metadata.pageAssets + if (!Array.isArray(value)) return false + + return value.some( + (item) => isRecord(item) && hasNonEmptyString(item.assetUrl), + ) +} + +function hasNonEmptyString(value: unknown): value is string { + return typeof value === "string" && value.trim().length > 0 +} + +function isRecord(value: unknown): value is Readonly> { + return typeof value === "object" && value !== null +} diff --git a/src/domains/sources/page-assets.ts b/src/domains/sources/page-assets.ts index cfb2765..c23a702 100644 --- a/src/domains/sources/page-assets.ts +++ b/src/domains/sources/page-assets.ts @@ -1,6 +1,10 @@ import "server-only" -import type { Knowledge } from "@ontos-ai/knowhere-sdk" +import type { + DocumentChunk, + Knowledge, + KnowledgeReadResponse, +} from "@ontos-ai/knowhere-sdk" import type { ChunkPageParams } from "@/domains/chunks" import type { SourcePageAssetView } from "./route-types" @@ -20,7 +24,30 @@ export type SourcePageAssetsPage = { } } +type PageAssetReadClient = { + readonly documents: { + listChunks( + documentId: string, + params: { + readonly page: number + readonly pageSize: number + readonly chunkType: "page" + readonly includeAssetUrls: true + }, + ): Promise<{ + readonly chunks: readonly DocumentChunk[] + readonly pagination?: { + readonly page?: number + readonly pageSize?: number + readonly total?: number + readonly totalPages?: number + } + }> + } +} + export async function readSourcePageAssets(input: { + readonly client: PageAssetReadClient readonly knowledge: Knowledge readonly source: ReadableSource readonly params: ChunkPageParams @@ -31,8 +58,36 @@ export async function readSourcePageAssets(input: { chunkType: "page", page: input.params.page, pageSize: input.params.pageSize, - assetUrlPolicy: "durable", }) + + const knowledgePage = toSourcePageAssetsPageFromKnowledgeResponse( + response, + input.params, + ) + if ( + isParsedStorageReadResponse(response) && + knowledgePage.pages.length > 0 + ) { + return knowledgePage + } + if (knowledgePage.pages.length > 0) return knowledgePage + + const remoteResponse = await input.client.documents.listChunks( + input.source.documentId, + { + page: input.params.page, + pageSize: input.params.pageSize, + chunkType: "page", + includeAssetUrls: true, + }, + ) + return toSourcePageAssetsPageFromRemoteResponse(remoteResponse, input.params) +} + +function toSourcePageAssetsPageFromKnowledgeResponse( + response: KnowledgeReadResponse, + params: ChunkPageParams, +): SourcePageAssetsPage { const pages = response.chunks.flatMap((chunk): SourcePageAssetView[] => readPageAssetViews(chunk.metadata.pageAssets, chunk.assetUrl), ) @@ -41,21 +96,67 @@ export async function readSourcePageAssets(input: { const resolvedTotal = total > 0 ? total : pages.length const computedTotalPages = Math.max( 1, - Math.ceil(resolvedTotal / input.params.pageSize), + Math.ceil(resolvedTotal / params.pageSize), ) const totalPages = Math.max(response.totalPages ?? 0, computedTotalPages) return { pages, pagination: { - page: response.page ?? input.params.page, - pageSize: response.pageSize ?? input.params.pageSize, + page: response.page ?? params.page, + pageSize: response.pageSize ?? params.pageSize, total: resolvedTotal, totalPages, }, } } +function toSourcePageAssetsPageFromRemoteResponse( + response: { + readonly chunks: readonly DocumentChunk[] + readonly pagination?: { + readonly page?: number + readonly pageSize?: number + readonly total?: number + readonly totalPages?: number + } + }, + params: ChunkPageParams, +): SourcePageAssetsPage { + const pages = response.chunks.flatMap((chunk): SourcePageAssetView[] => + readPageAssetViews(chunk.metadata.pageAssets, chunk.assetUrl ?? undefined), + ) + const maxPageNumber = getMaxPageNumber(pages) + const total = Math.max(response.pagination?.total ?? 0, maxPageNumber ?? 0) + const resolvedTotal = total > 0 ? total : pages.length + const computedTotalPages = Math.max( + 1, + Math.ceil(resolvedTotal / params.pageSize), + ) + const totalPages = Math.max( + response.pagination?.totalPages ?? 0, + computedTotalPages, + ) + + return { + pages, + pagination: { + page: response.pagination?.page ?? params.page, + pageSize: response.pagination?.pageSize ?? params.pageSize, + total: resolvedTotal, + totalPages, + }, + } +} + +function isParsedStorageReadResponse(response: KnowledgeReadResponse): boolean { + const resultDirectoryPath = response.document.resultDirectoryPath + return ( + typeof resultDirectoryPath === "string" && + resultDirectoryPath.startsWith("parsed-storage:") + ) +} + function readPageAssetViews( value: unknown, fallbackAssetUrl?: string, diff --git a/src/domains/sources/route-chunks.ts b/src/domains/sources/route-chunks.ts index 8d9d315..c636cb6 100644 --- a/src/domains/sources/route-chunks.ts +++ b/src/domains/sources/route-chunks.ts @@ -13,7 +13,7 @@ import { } from "./remote-library" import { getClientForWorkspace, - getKnowledgeForSource, + getKnowledgeResourcesForSource, } from "./route-dependencies" import { sourceRowRepository } from "./source-row-repository" import type { @@ -99,7 +99,7 @@ const loadSourceChunksEffect = ( const apiKey = yield* Effect.tryPromise(() => deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), ) - const knowledge = getKnowledgeForSource({ + const readResources = getKnowledgeResourcesForSource({ apiKey, workspaceId: workspace.id, sourceId: source.id, @@ -114,7 +114,11 @@ const loadSourceChunksEffect = ( if (input.shouldLoadAll) { return yield* Effect.tryPromise(() => - readAllSourceChunks({ knowledge, source: readableSource }), + readAllSourceChunks({ + client: readResources.client, + knowledge: readResources.knowledge, + source: readableSource, + }), ).pipe( Effect.map((chunks) => routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }), @@ -125,7 +129,8 @@ const loadSourceChunksEffect = ( return yield* Effect.tryPromise(() => readSourceChunkPage({ - knowledge, + client: readResources.client, + knowledge: readResources.knowledge, source: readableSource, params: input.pageParams, }), @@ -175,7 +180,7 @@ const loadRemoteChunkPageEffect = ( ) const documentId = source.knowhereDocumentId ?? remoteDocument.documentId - const knowledge = getKnowledgeForSource({ + const readResources = getKnowledgeResourcesForSource({ apiKey, workspaceId: workspace.id, sourceId: source.id, @@ -190,7 +195,11 @@ const loadRemoteChunkPageEffect = ( if (input.shouldLoadAll) { return yield* Effect.tryPromise(() => - readAllSourceChunks({ knowledge, source: readableSource }), + readAllSourceChunks({ + client: readResources.client, + knowledge: readResources.knowledge, + source: readableSource, + }), ).pipe( Effect.map((chunks) => routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }), @@ -201,7 +210,8 @@ const loadRemoteChunkPageEffect = ( return yield* Effect.tryPromise(() => readSourceChunkPage({ - knowledge, + client: readResources.client, + knowledge: readResources.knowledge, source: readableSource, params: input.pageParams, }), diff --git a/src/domains/sources/route-dependencies.ts b/src/domains/sources/route-dependencies.ts index a51a9b7..b0cc887 100644 --- a/src/domains/sources/route-dependencies.ts +++ b/src/domains/sources/route-dependencies.ts @@ -99,6 +99,16 @@ function getKnowledgeForSource(input: { readonly documentId: string readonly revisionKey?: string | null }): Knowledge { + return getKnowledgeResourcesForSource(input).knowledge +} + +function getKnowledgeResourcesForSource(input: { + readonly apiKey: string + readonly workspaceId: string + readonly sourceId: string + readonly documentId: string + readonly revisionKey?: string | null +}): { readonly client: SourceRouteKnowhereClient; readonly knowledge: Knowledge } { const scheduler = createParsedDocumentSyncScheduler({ workspaceId: input.workspaceId, sourceId: input.sourceId, @@ -106,15 +116,16 @@ function getKnowledgeForSource(input: { apiKey: input.apiKey, revisionKey: input.revisionKey ?? undefined, }) - const { knowledge } = makeKnowhereClientWithParsedStorage(input.apiKey, { + const resources = makeKnowhereClientWithParsedStorage(input.apiKey, { workspaceId: input.workspaceId, scheduler, }) - return knowledge + return { client: resources.client, knowledge: resources.knowledge } } export { createSourceRouteDependencies, getClientForWorkspace, getKnowledgeForSource, + getKnowledgeResourcesForSource, } diff --git a/src/domains/sources/route-page-assets.ts b/src/domains/sources/route-page-assets.ts index 3077a83..a2316fd 100644 --- a/src/domains/sources/route-page-assets.ts +++ b/src/domains/sources/route-page-assets.ts @@ -9,7 +9,7 @@ import { } from "./remote-library" import { getClientForWorkspace, - getKnowledgeForSource, + getKnowledgeResourcesForSource, } from "./route-dependencies" import { sourceRowRepository } from "./source-row-repository" import type { @@ -71,7 +71,7 @@ const loadSourcePageAssetsEffect = ( const apiKey = yield* Effect.tryPromise(() => deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), ) - const knowledge = getKnowledgeForSource({ + const readResources = getKnowledgeResourcesForSource({ apiKey, workspaceId: workspace.id, sourceId: source.id, @@ -80,7 +80,8 @@ const loadSourcePageAssetsEffect = ( }) return yield* Effect.tryPromise(() => readSourcePageAssets({ - knowledge, + client: readResources.client, + knowledge: readResources.knowledge, source: { documentId, revisionKey: source.knowhereJobId, @@ -136,7 +137,7 @@ const loadRemotePageAssetsEffect = ( const documentId = source.knowhereDocumentId ?? remoteDocument.documentId const revisionKey = source.knowhereJobId ?? remoteDocument.revisionKey ?? null - const knowledge = getKnowledgeForSource({ + const readResources = getKnowledgeResourcesForSource({ apiKey, workspaceId: workspace.id, sourceId: source.id, @@ -145,7 +146,8 @@ const loadRemotePageAssetsEffect = ( }) return yield* Effect.tryPromise(() => readSourcePageAssets({ - knowledge, + client: readResources.client, + knowledge: readResources.knowledge, source: { documentId, revisionKey }, params: input.pageParams, }), From 2d825942e48de7cdded607837e43d6cdb1f783d6 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 20:08:06 +0800 Subject: [PATCH 25/51] Adapt parsed document storage layout --- .../[sourceId]/page-assets/route.test.ts | 87 +++++++++ .../chat/media-asset-hardening.test.ts | 8 +- src/domains/chat/route-service.test.ts | 8 +- src/domains/chunks/read.test.ts | 32 ++++ .../parsed-document-blob-storage.test.ts | 73 ++++++- .../sources/parsed-document-blob-storage.ts | 179 ++++++++++++++---- 6 files changed, 335 insertions(+), 52 deletions(-) diff --git a/src/app/api/sources/[sourceId]/page-assets/route.test.ts b/src/app/api/sources/[sourceId]/page-assets/route.test.ts index 76ef16a..4f00bd3 100644 --- a/src/app/api/sources/[sourceId]/page-assets/route.test.ts +++ b/src/app/api/sources/[sourceId]/page-assets/route.test.ts @@ -156,6 +156,93 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { expect(mocks.listChunks).not.toHaveBeenCalled() }) + it("uses SDK page asset URLs directly when the SDK remote fallback is usable", async () => { + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockResolvedValue({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "remote:doc_1", + }, + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + metadata: { + pageAssets: [ + { + pageNum: 1, + assetUrl: "https://sdk.example/page-000001.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, + }) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + pages: [ + { + pageNumber: 1, + assetUrl: "https://sdk.example/page-000001.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + }) + expect(response.status).toBe(200) + expect(mocks.readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "job_1", + chunkType: "page", + page: 1, + pageSize: 1, + }) + expect(mocks.listChunks).not.toHaveBeenCalled() + }) + it("rejects non-ready workspace sources", async () => { mocks.getCurrentUser.mockResolvedValue({ id: "user_1", diff --git a/src/domains/chat/media-asset-hardening.test.ts b/src/domains/chat/media-asset-hardening.test.ts index 7c106c9..2afa6d9 100644 --- a/src/domains/chat/media-asset-hardening.test.ts +++ b/src/domains/chat/media-asset-hardening.test.ts @@ -24,7 +24,7 @@ afterEach(() => { describe("hardenChatMediaAssetUrls", () => { it("keeps an already Notebook-owned asset URL without calling the hardener", async () => { const ownedUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_1/rev_1/assets/images/a.png" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_1/rev_1/images/a.png" const hardenChatAssetUrl = vi.fn(async () => null) const result = await hardenChatMediaAssetUrls({ @@ -52,7 +52,7 @@ describe("hardenChatMediaAssetUrls", () => { const rawAssetUrl = "https://knowhere-storage.example/results/job_1/images/id-front.jpg?AWSAccessKeyId=test" const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/assets/images/id-front.jpg" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/images/id-front.jpg" const hardenChatAssetUrl = vi.fn().mockResolvedValue(durableUrl) const result = await hardenChatMediaAssetUrls({ @@ -116,7 +116,7 @@ describe("hardenChatMediaAssetUrls", () => { const rawAssetUrl = "https://knowhere-storage.example/results/job_1/images/front.jpg?AWSAccessKeyId=test" const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/assets/images/front.jpg" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/images/front.jpg" const hardenChatAssetUrl = vi.fn().mockResolvedValue(durableUrl) const result = await hardenChatMediaAssetUrls({ @@ -164,7 +164,7 @@ describe("isNotebookOwnedAssetUrl", () => { ).toBe(true) expect( isNotebookOwnedAssetUrl( - "https://cdn.example/workspaces/w/parsed-documents/d/r/assets/a.png", + "https://cdn.example/workspaces/w/parsed-documents/d/r/images/a.png", ), ).toBe(true) expect( diff --git a/src/domains/chat/route-service.test.ts b/src/domains/chat/route-service.test.ts index c771f12..d66ae68 100644 --- a/src/domains/chat/route-service.test.ts +++ b/src/domains/chat/route-service.test.ts @@ -183,7 +183,7 @@ describe("chat route services", () => { }) const rawUrl = "https://knowhere-storage.example/results/job_1/pages/page-1.png" const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_legacy/job_1/assets/pages/page-1.png" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_legacy/job_1/pages/page-1.png" mocks.parsedStorageWriteAsset.mockResolvedValue({ url: durableUrl }) mocks.getAuthenticatedWithClient.mockResolvedValue({ user: { id: "user_1" }, @@ -249,7 +249,7 @@ describe("chat route services", () => { const rawUrl = "https://knowhere-storage.example/results/job_1/images/id-front.png?AWSAccessKeyId=test" const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/job_1/assets/images/id-front.png" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/job_1/images/id-front.png" mocks.parsedStorageWriteAsset.mockResolvedValue({ url: durableUrl }) mocks.generateText.mockResolvedValue({ text: `The card number is visible. ${durableUrl} ${rawUrl}`, @@ -364,7 +364,7 @@ describe("chat route services", () => { knowhereJobId: "job_1", }) const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_contract/job_1/assets/page_citation_assets/page-8.png" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_contract/job_1/page_citation_assets/page-8.png" mocks.parsedStorageGetAssetUrl.mockResolvedValue(durableUrl) mocks.generateText.mockResolvedValue({ text: "The page states 5000 yuan per occurrence.", @@ -467,7 +467,7 @@ describe("chat route services", () => { const rawUrl = "https://knowhere-storage.example/results/job_remote/page_citation_assets/page-8.png?AWSAccessKeyId=test" const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_remote/job_remote/assets/page_citation_assets/page-8.png" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_remote/job_remote/page_citation_assets/page-8.png" mocks.parsedStorageWriteAsset.mockResolvedValue({ url: durableUrl }) mocks.generateText.mockResolvedValue({ text: "The page states 5000 yuan per occurrence.", diff --git a/src/domains/chunks/read.test.ts b/src/domains/chunks/read.test.ts index 9ba2203..f63c897 100644 --- a/src/domains/chunks/read.test.ts +++ b/src/domains/chunks/read.test.ts @@ -124,6 +124,38 @@ describe("readSourceChunkPage", () => { ) }) + it("uses SDK remote chunks directly when the SDK returns usable asset URLs", async () => { + const listChunks = vi.fn() + const readChunks = vi.fn(async () => ({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "remote:doc_1", + }, + chunks: [ + makeReadChunk({ + chunkType: "image", + filePath: "images/a.png", + assetUrl: "https://sdk.example/assets/a.png", + }), + ], + page: 1, + pageSize: 50, + totalChunks: 1, + totalPages: 1, + })) + const knowledge = { readChunks } as unknown as Knowledge + + const result = await readSourceChunkPage({ + client: { documents: { listChunks } }, + knowledge, + source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, + params: { page: 1, pageSize: 50 }, + }) + + expect(listChunks).not.toHaveBeenCalled() + expect(result.chunks[0]?.assetUrl).toBe("https://sdk.example/assets/a.png") + }) + it("omits revisionKey when the source has none", async () => { const listChunks = vi.fn() const readChunks = vi.fn(async () => ({ diff --git a/src/domains/sources/parsed-document-blob-storage.test.ts b/src/domains/sources/parsed-document-blob-storage.test.ts index 8a3c427..4c3ee16 100644 --- a/src/domains/sources/parsed-document-blob-storage.test.ts +++ b/src/domains/sources/parsed-document-blob-storage.test.ts @@ -34,6 +34,7 @@ function createFakeBlobStore(): { return { statusCode: 200, stream: new Response(body).body as ReadableStream, + contentType: object.contentType, } }, put: async (pathname, body, options) => { @@ -93,7 +94,7 @@ const chunkPage: KnowhereParsedSnapshotChunkPage = { } describe("BlobParsedDocumentStorage", () => { - it("round-trips a manifest keyed by workspace/document/revision", async () => { + it("round-trips a manifest at the result-relative manifest path", async () => { const { store, objects } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ workspaceId: "ws_1", @@ -103,20 +104,23 @@ describe("BlobParsedDocumentStorage", () => { await storage.writeManifest({ documentId, revisionKey, manifest }) expect([...objects.keys()]).toContain( - "workspaces/ws_1/parsed-documents/doc_123/job_result_1/manifest/current.json", + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/manifest.json", ) const read = await storage.readManifest({ documentId, revisionKey }) expect(read).toEqual(manifest) }) - it("round-trips a chunk page", async () => { - const { store } = createFakeBlobStore() + it("round-trips a chunk page at its result-relative chunk path", async () => { + const { store, objects } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ workspaceId: "ws_1", blobStore: store, }) await storage.writeChunkPage({ documentId, revisionKey, page: chunkPage }) + expect([...objects.keys()]).toContain( + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/chunks/page-1.json", + ) const read = await storage.readChunkPage({ documentId, revisionKey, page: 1 }) expect(read).toEqual(chunkPage) }) @@ -141,7 +145,7 @@ describe("BlobParsedDocumentStorage", () => { expect(read).toEqual(progress) }) - it("writes an asset and resolves a durable URL", async () => { + it("writes an asset without adding a Notebook logical assets prefix", async () => { const { store } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ workspaceId: "ws_1", @@ -156,7 +160,7 @@ describe("BlobParsedDocumentStorage", () => { contentType: "image/png", }) expect(written.url).toContain( - "workspaces/ws_1/parsed-documents/doc_123/job_result_1/assets/images/fig-1.png", + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/images/fig-1.png", ) const url = await storage.getAssetUrl({ @@ -167,6 +171,56 @@ describe("BlobParsedDocumentStorage", () => { expect(url).toBe(written.url) }) + it("round-trips arbitrary result-relative objects", async () => { + const { store, objects } = createFakeBlobStore() + const storage = new BlobParsedDocumentStorage({ + workspaceId: "ws_1", + blobStore: store, + }) + + await storage.writeObject({ + documentId, + revisionKey, + path: "page_citation_assets/page-1.png", + body: new Uint8Array([4, 5, 6]), + contentType: "image/png", + }) + + expect([...objects.keys()]).toContain( + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/page_citation_assets/page-1.png", + ) + const read = await storage.readObject({ + documentId, + revisionKey, + path: "page_citation_assets/page-1.png", + }) + expect(read?.body).toEqual(Buffer.from([4, 5, 6])) + expect(read?.contentType).toBe("image/png") + }) + + it("resolves legacy asset-prefixed URLs without writing new assets there", async () => { + const { store, objects } = createFakeBlobStore() + const storage = new BlobParsedDocumentStorage({ + workspaceId: "ws_1", + blobStore: store, + }) + const legacyKey = + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/assets/images/fig-1.png" + objects.set(legacyKey, { + body: Buffer.from([1]), + contentType: "image/png", + }) + + const url = await storage.getAssetUrl({ + documentId, + revisionKey, + sourcePath: "images/fig-1.png", + }) + expect(url).toBe( + "https://fake.public.blob.vercel-storage.com/workspaces/ws_1/parsed-documents/doc_123/job_result_1/assets/images/fig-1.png", + ) + }) + it("returns null for a missing manifest, chunk page, progress, and asset", async () => { const { store } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ @@ -219,5 +273,12 @@ describe("BlobParsedDocumentStorage", () => { sourcePath: "../../etc/passwd", }), ).rejects.toThrow(/Invalid parsed storage path/) + await expect( + storage.readObject({ + documentId, + revisionKey, + path: "images/../escape.png", + }), + ).rejects.toThrow(/Invalid parsed storage path/) }) }) diff --git a/src/domains/sources/parsed-document-blob-storage.ts b/src/domains/sources/parsed-document-blob-storage.ts index 23a93a0..8f19837 100644 --- a/src/domains/sources/parsed-document-blob-storage.ts +++ b/src/domains/sources/parsed-document-blob-storage.ts @@ -26,14 +26,17 @@ import type { */ const parsedDocumentsDirectoryName = "parsed-documents" -const manifestStoragePath = "manifest/current.json" +const manifestStoragePath = "manifest.json" +const legacyManifestStoragePath = "manifest/current.json" const syncProgressStoragePath = "sync-progress.json" const jsonContentType = "application/json; charset=utf-8" +const binaryContentType = "application/octet-stream" type BlobGetResult = | { readonly statusCode: 200 readonly stream: ReadableStream + readonly contentType?: string } | { readonly statusCode: 304 @@ -77,8 +80,34 @@ export type BlobParsedDocumentStorageInput = { readonly blobStore?: ParsedDocumentBlobStore } +export type ParsedDocumentStorageObject = ParsedDocumentStorageDocument & { + readonly path: string +} + +export type ParsedDocumentStorageWritableObject = + ParsedDocumentStorageObject & { + readonly body: string | Buffer | Uint8Array + readonly contentType?: string + } + +export type ParsedDocumentStorageReadableObject = { + readonly body: Buffer + readonly contentType?: string +} + const vercelBlobStore: ParsedDocumentBlobStore = { - get: (pathname, options) => get(pathname, options), + get: async (pathname, options) => { + const blob = await get(pathname, options) + if (!blob) return null + if (blob.statusCode === 304) { + return { statusCode: 304, stream: null } + } + return { + statusCode: 200, + stream: blob.stream, + contentType: blob.blob.contentType, + } + }, put: async (pathname, body, options) => { const blob = await put(pathname, body, { access: options.access, @@ -113,8 +142,21 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { async readManifest( params: ParsedDocumentStorageManifestParams, ): Promise { + const current = await this.readJson( + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: manifestStoragePath, + }), + ) + if (current) return current + return this.readJson( - this.getManifestKey(params.documentId, params.revisionKey), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: legacyManifestStoragePath, + }), ) } @@ -124,7 +166,11 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { readonly manifest: KnowhereParsedSnapshotManifest }): Promise { await this.writeJson( - this.getManifestKey(params.documentId, params.revisionKey), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: manifestStoragePath, + }), params.manifest, ) } @@ -134,7 +180,11 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { ): Promise { // chunkType filtering happens SDK-side after read; storage returns the full page. return this.readJson( - this.getChunkPageKey(params.documentId, params.revisionKey, params.page), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: getChunkPageStoragePath(params.page), + }), ) } @@ -144,11 +194,11 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { readonly page: KnowhereParsedSnapshotChunkPage }): Promise { await this.writeJson( - this.getChunkPageKey( - params.documentId, - params.revisionKey, - params.page.page, - ), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: getChunkPageStoragePath(params.page.page), + }), params.page, ) } @@ -157,7 +207,11 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { params: ParsedDocumentStorageDocument & ParsedDocumentStorageAsset, ): Promise<{ readonly sourcePath: string; readonly url?: string }> { const blob = await this.blobStore.put( - this.getAssetKey(params.documentId, params.revisionKey, params.sourcePath), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: params.sourcePath, + }), Buffer.from(params.body), { access: "public", @@ -173,26 +227,80 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { params: ParsedDocumentStorageAssetParams, ): Promise { const result = await this.blobStore.head( - this.getAssetKey(params.documentId, params.revisionKey, params.sourcePath), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: params.sourcePath, + }), ) - return result?.url ?? null + if (result) return result.url + + const legacyResult = await this.blobStore.head( + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: `assets/${params.sourcePath}`, + }), + ) + return legacyResult?.url ?? null } async readSyncProgress( params: ParsedDocumentStorageDocument, ): Promise { return this.readJson( - this.getSyncProgressKey(params.documentId, params.revisionKey), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: syncProgressStoragePath, + }), ) } async writeSyncProgress(params: ParsedDocumentSyncProgress): Promise { await this.writeJson( - this.getSyncProgressKey(params.documentId, params.revisionKey), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: syncProgressStoragePath, + }), params, ) } + async readObject( + params: ParsedDocumentStorageObject, + ): Promise { + return this.readBlobObject(this.getObjectKey(params)) + } + + async writeObject( + params: ParsedDocumentStorageWritableObject, + ): Promise<{ readonly path: string; readonly url?: string }> { + const blob = await this.blobStore.put( + this.getObjectKey(params), + Buffer.from(params.body), + { + access: "public", + allowOverwrite: true, + contentType: params.contentType ?? binaryContentType, + multipart: true, + }, + ) + return { path: params.path, url: blob.url } + } + + async getObjectUrl( + params: ParsedDocumentStorageObject, + ): Promise { + const result = await this.blobStore.head(this.getObjectKey(params)) + return result?.url ?? null + } + + async deleteObject(params: ParsedDocumentStorageObject): Promise { + await this.blobStore.del(this.getObjectKey(params)) + } + private getRevisionPrefix(documentId: string, revisionKey: string): string { return [ "workspaces", @@ -203,28 +311,8 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { ].join("/") } - private getManifestKey(documentId: string, revisionKey: string): string { - return `${this.getRevisionPrefix(documentId, revisionKey)}/${manifestStoragePath}` - } - - private getChunkPageKey( - documentId: string, - revisionKey: string, - page: number, - ): string { - return `${this.getRevisionPrefix(documentId, revisionKey)}/chunks/page-${page}.json` - } - - private getSyncProgressKey(documentId: string, revisionKey: string): string { - return `${this.getRevisionPrefix(documentId, revisionKey)}/${syncProgressStoragePath}` - } - - private getAssetKey( - documentId: string, - revisionKey: string, - sourcePath: string, - ): string { - return `${this.getRevisionPrefix(documentId, revisionKey)}/assets/${normalizeRelativeStoragePath(sourcePath)}` + private getObjectKey(params: ParsedDocumentStorageObject): string { + return `${this.getRevisionPrefix(params.documentId, params.revisionKey)}/${normalizeRelativeStoragePath(params.path)}` } private async readJson(key: string): Promise { @@ -246,10 +334,21 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { } private async readBlobText(key: string): Promise { + const object = await this.readBlobObject(key) + return object ? object.body.toString("utf8") : null + } + + private async readBlobObject( + key: string, + ): Promise { try { const result = await this.blobStore.get(key, { access: "public" }) if (!result || result.statusCode !== 200) return null - return await new Response(result.stream).text() + const body = Buffer.from(await new Response(result.stream).arrayBuffer()) + return { + body, + ...(result.contentType ? { contentType: result.contentType } : {}), + } } catch (error) { if (error instanceof BlobNotFoundError) return null throw error @@ -257,6 +356,10 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { } } +function getChunkPageStoragePath(page: number): string { + return `chunks/page-${page}.json` +} + /** * Reject path segments that could traverse outside the intended prefix. Mirrors * the SDK `DiskParsedDocumentStorage` guard so blob keys stay well-formed. From fde80272615ef99f5ff5800cfd31c5080a7ca106 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 23:24:40 +0800 Subject: [PATCH 26/51] Align Notebook with SDK result layout storage --- src/domains/chat/index.test.ts | 10 +- src/domains/chat/index.ts | 5 +- src/domains/chat/service.test.ts | 4 +- .../parsed-document-blob-storage.test.ts | 128 +++++++++----- .../sources/parsed-document-blob-storage.ts | 165 +++++++----------- src/integrations/knowhere.ts | 2 - 6 files changed, 154 insertions(+), 160 deletions(-) diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 8c239a9..9a34a4b 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -78,7 +78,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "What does the document say?", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 1, excludeDocumentIds: ["doc_excluded", "doc_remote"], }); @@ -579,7 +579,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "SpaceX rocket photos", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 3, }); expect(answer.answer).toBe("Use this launch photo."); @@ -1573,7 +1573,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "公民身份证明 图片", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 3, }); const imageCitations = answer.citations.filter( @@ -1664,7 +1664,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "Tesla Q4 2025 Update energy generation and storage deployments", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 1, }); expect(generateAnswer).toHaveBeenCalledWith({ @@ -1720,7 +1720,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "Tesla energy storage deployments", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 1, }); expect(JSON.stringify(queryInput)).not.toContain( diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index fa07579..4f5dee6 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -42,6 +42,9 @@ import { enrichRetrievalResultsWithPageCitationAssetUrls } from "./page-citation import type { HardenableRetrievalResult } from "./media-asset-hardening" const DEFAULT_TOP_K = 8 +const NOTEBOOK_USE_AGENTIC_RETRIEVAL: NonNullable< + RetrievalQueryParams["useAgentic"] +> = false const MAX_AGENTIC_TOP_K = 12 const MAX_AGENTIC_MERGED_RESULT_COUNT = 24 const MAX_AGENTIC_MERGED_REFERENCED_CHUNK_COUNT = 24 @@ -795,7 +798,7 @@ function buildRetrievalQueryParams(input: { namespace: input.namespace, query, topK: normalizeTopK(input.input.topK), - useAgentic: true, + useAgentic: NOTEBOOK_USE_AGENTIC_RETRIEVAL, dataType, ...(input.input.signalPaths && input.input.signalPaths.length > 0 ? { signalPaths: input.input.signalPaths } diff --git a/src/domains/chat/service.test.ts b/src/domains/chat/service.test.ts index 7c0ce92..12df745 100644 --- a/src/domains/chat/service.test.ts +++ b/src/domains/chat/service.test.ts @@ -57,7 +57,7 @@ describe("handleChatTurn", () => { namespace: "notebook-namespace", query: "What does the document say?", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 1, excludeDocumentIds: ["doc_excluded"], }); @@ -221,7 +221,7 @@ describe("handleChatTurn", () => { namespace: "notebook-namespace", query: "Tesla Q4 2025 Update energy generation and storage deployments", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 1, }); }); diff --git a/src/domains/sources/parsed-document-blob-storage.test.ts b/src/domains/sources/parsed-document-blob-storage.test.ts index 4c3ee16..b5fb005 100644 --- a/src/domains/sources/parsed-document-blob-storage.test.ts +++ b/src/domains/sources/parsed-document-blob-storage.test.ts @@ -1,9 +1,5 @@ import { describe, expect, it } from "vitest" -import type { - KnowhereParsedSnapshotChunkPage, - KnowhereParsedSnapshotManifest, - ParsedDocumentSyncProgress, -} from "@ontos-ai/knowhere-sdk" +import type { ParsedDocumentSyncProgress } from "@ontos-ai/knowhere-sdk" import { BlobParsedDocumentStorage, @@ -34,6 +30,7 @@ function createFakeBlobStore(): { return { statusCode: 200, stream: new Response(body).body as ReadableStream, + url: toUrl(pathname), contentType: object.contentType, } }, @@ -56,42 +53,30 @@ function createFakeBlobStore(): { const documentId = "doc_123" const revisionKey = "job_result_1" -const manifest: KnowhereParsedSnapshotManifest = { - version: 1, - kind: "knowhere-parsed-result-snapshot", +const manifest = { + version: "2.0", jobId: "job_1", - revisionKey, - documentId, sourceFileName: "example.pdf", - totalChunks: 2, - chunkPageSize: 200, - chunkPages: [{ page: 1, pageSize: 200, chunkCount: 2, key: "chunks/page-1.json" }], - assetUrlsByFilePath: {}, - createdAt: "2026-07-04T00:00:00.000Z", -} + statistics: { + totalChunks: 2, + textChunks: 1, + imageChunks: 1, + tableChunks: 0, + pageChunks: 0, + }, +} satisfies Record -const chunkPage: KnowhereParsedSnapshotChunkPage = { - version: 1, - jobId: "job_1", - revisionKey, - documentId, - sourceFileName: "example.pdf", - page: 1, - pageSize: 200, - total: 2, - totalPages: 1, +const chunks = { chunks: [ { - id: "c1", - chunkId: "c1", - chunkType: "text", + chunk_id: "c1", + type: "text", content: "hello", - sourceChunkPath: "example.pdf", - sortOrder: 0, + path: "example.pdf", metadata: {}, }, ], -} +} satisfies Record describe("BlobParsedDocumentStorage", () => { it("round-trips a manifest at the result-relative manifest path", async () => { @@ -101,28 +86,52 @@ describe("BlobParsedDocumentStorage", () => { blobStore: store, }) - await storage.writeManifest({ documentId, revisionKey, manifest }) + await storage.writeObject({ + documentId, + revisionKey, + path: "manifest.json", + body: Buffer.from(JSON.stringify(manifest), "utf8"), + contentType: "application/json; charset=utf-8", + }) expect([...objects.keys()]).toContain( "workspaces/ws_1/parsed-documents/doc_123/job_result_1/manifest.json", ) - const read = await storage.readManifest({ documentId, revisionKey }) - expect(read).toEqual(manifest) + const read = await storage.readObject({ + documentId, + revisionKey, + path: "manifest.json", + }) + expect(read ? Buffer.from(read.body).toString("utf8") : null).toBe( + JSON.stringify(manifest), + ) }) - it("round-trips a chunk page at its result-relative chunk path", async () => { + it("round-trips chunks at the result-relative chunks path", async () => { const { store, objects } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ workspaceId: "ws_1", blobStore: store, }) - await storage.writeChunkPage({ documentId, revisionKey, page: chunkPage }) + await storage.writeObject({ + documentId, + revisionKey, + path: "chunks.json", + body: Buffer.from(JSON.stringify(chunks), "utf8"), + contentType: "application/json; charset=utf-8", + }) expect([...objects.keys()]).toContain( - "workspaces/ws_1/parsed-documents/doc_123/job_result_1/chunks/page-1.json", + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/chunks.json", + ) + const read = await storage.readObject({ + documentId, + revisionKey, + path: "chunks.json", + }) + expect(read ? Buffer.from(read.body).toString("utf8") : null).toBe( + JSON.stringify(chunks), ) - const read = await storage.readChunkPage({ documentId, revisionKey, page: 1 }) - expect(read).toEqual(chunkPage) }) it("round-trips sync progress", async () => { @@ -135,7 +144,6 @@ describe("BlobParsedDocumentStorage", () => { documentId, revisionKey, nextChunkPage: 3, - nextAssetIndex: 0, status: "running", updatedAt: "2026-07-04T00:00:00.000Z", } @@ -185,10 +193,20 @@ describe("BlobParsedDocumentStorage", () => { body: new Uint8Array([4, 5, 6]), contentType: "image/png", }) + await storage.writeObject({ + documentId, + revisionKey, + path: "images/fig-1.png", + body: new Uint8Array([7, 8, 9]), + contentType: "image/png", + }) expect([...objects.keys()]).toContain( "workspaces/ws_1/parsed-documents/doc_123/job_result_1/page_citation_assets/page-1.png", ) + expect([...objects.keys()]).toContain( + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/images/fig-1.png", + ) const read = await storage.readObject({ documentId, revisionKey, @@ -221,16 +239,15 @@ describe("BlobParsedDocumentStorage", () => { ) }) - it("returns null for a missing manifest, chunk page, progress, and asset", async () => { + it("returns null for a missing object, progress, and asset", async () => { const { store } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ workspaceId: "ws_1", blobStore: store, }) - expect(await storage.readManifest({ documentId, revisionKey })).toBeNull() expect( - await storage.readChunkPage({ documentId, revisionKey, page: 9 }), + await storage.readObject({ documentId, revisionKey, path: "manifest.json" }), ).toBeNull() expect(await storage.readSyncProgress({ documentId, revisionKey })).toBeNull() expect( @@ -245,10 +262,17 @@ describe("BlobParsedDocumentStorage", () => { blobStore: store, }) - await storage.writeManifest({ documentId, revisionKey, manifest }) - const staleRead = await storage.readManifest({ + await storage.writeObject({ + documentId, + revisionKey, + path: "manifest.json", + body: Buffer.from(JSON.stringify(manifest), "utf8"), + contentType: "application/json; charset=utf-8", + }) + const staleRead = await storage.readObject({ documentId, revisionKey: "job_result_2", + path: "manifest.json", }) expect(staleRead).toBeNull() }) @@ -261,10 +285,18 @@ describe("BlobParsedDocumentStorage", () => { }) await expect( - storage.readManifest({ documentId: "../escape", revisionKey }), + storage.readObject({ + documentId: "../escape", + revisionKey, + path: "manifest.json", + }), ).rejects.toThrow(/Invalid parsed storage segment/) await expect( - storage.readManifest({ documentId, revisionKey: "a/b" }), + storage.readObject({ + documentId, + revisionKey: "a/b", + path: "manifest.json", + }), ).rejects.toThrow(/Invalid parsed storage segment/) await expect( storage.getAssetUrl({ diff --git a/src/domains/sources/parsed-document-blob-storage.ts b/src/domains/sources/parsed-document-blob-storage.ts index 8f19837..5a05e5c 100644 --- a/src/domains/sources/parsed-document-blob-storage.ts +++ b/src/domains/sources/parsed-document-blob-storage.ts @@ -2,15 +2,15 @@ import "server-only" import { del, get, head, put, BlobNotFoundError } from "@vercel/blob" import type { - KnowhereParsedSnapshotChunkPage, - KnowhereParsedSnapshotManifest, + ParsedDocumentObject, + ParsedDocumentObjectHead, + ParsedDocumentObjectParams, + ParsedDocumentRevisionParams, ParsedDocumentStorage, - ParsedDocumentStorageAsset, - ParsedDocumentStorageAssetParams, - ParsedDocumentStorageChunkPageParams, ParsedDocumentStorageDocument, - ParsedDocumentStorageManifestParams, ParsedDocumentSyncProgress, + ParsedDocumentWriteObjectParams, + ParsedDocumentWriteObjectResult, } from "@ontos-ai/knowhere-sdk" /** @@ -26,8 +26,6 @@ import type { */ const parsedDocumentsDirectoryName = "parsed-documents" -const manifestStoragePath = "manifest.json" -const legacyManifestStoragePath = "manifest/current.json" const syncProgressStoragePath = "sync-progress.json" const jsonContentType = "application/json; charset=utf-8" const binaryContentType = "application/octet-stream" @@ -36,6 +34,7 @@ type BlobGetResult = | { readonly statusCode: 200 readonly stream: ReadableStream + readonly url: string readonly contentType?: string } | { @@ -80,19 +79,19 @@ export type BlobParsedDocumentStorageInput = { readonly blobStore?: ParsedDocumentBlobStore } -export type ParsedDocumentStorageObject = ParsedDocumentStorageDocument & { +type ParsedDocumentStorageObject = ParsedDocumentStorageDocument & { readonly path: string } -export type ParsedDocumentStorageWritableObject = - ParsedDocumentStorageObject & { - readonly body: string | Buffer | Uint8Array - readonly contentType?: string - } +type LegacyParsedDocumentStorageAsset = { + readonly sourcePath: string + readonly body: string | Buffer | Uint8Array + readonly contentType: string + readonly metadata?: Readonly> +} -export type ParsedDocumentStorageReadableObject = { - readonly body: Buffer - readonly contentType?: string +type LegacyParsedDocumentStorageAssetParams = ParsedDocumentStorageDocument & { + readonly sourcePath: string } const vercelBlobStore: ParsedDocumentBlobStore = { @@ -105,6 +104,7 @@ const vercelBlobStore: ParsedDocumentBlobStore = { return { statusCode: 200, stream: blob.stream, + url: blob.blob.url, contentType: blob.blob.contentType, } }, @@ -139,72 +139,8 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { this.blobStore = input.blobStore ?? vercelBlobStore } - async readManifest( - params: ParsedDocumentStorageManifestParams, - ): Promise { - const current = await this.readJson( - this.getObjectKey({ - documentId: params.documentId, - revisionKey: params.revisionKey, - path: manifestStoragePath, - }), - ) - if (current) return current - - return this.readJson( - this.getObjectKey({ - documentId: params.documentId, - revisionKey: params.revisionKey, - path: legacyManifestStoragePath, - }), - ) - } - - async writeManifest(params: { - readonly documentId: string - readonly revisionKey: string - readonly manifest: KnowhereParsedSnapshotManifest - }): Promise { - await this.writeJson( - this.getObjectKey({ - documentId: params.documentId, - revisionKey: params.revisionKey, - path: manifestStoragePath, - }), - params.manifest, - ) - } - - async readChunkPage( - params: ParsedDocumentStorageChunkPageParams, - ): Promise { - // chunkType filtering happens SDK-side after read; storage returns the full page. - return this.readJson( - this.getObjectKey({ - documentId: params.documentId, - revisionKey: params.revisionKey, - path: getChunkPageStoragePath(params.page), - }), - ) - } - - async writeChunkPage(params: { - readonly documentId: string - readonly revisionKey: string - readonly page: KnowhereParsedSnapshotChunkPage - }): Promise { - await this.writeJson( - this.getObjectKey({ - documentId: params.documentId, - revisionKey: params.revisionKey, - path: getChunkPageStoragePath(params.page.page), - }), - params.page, - ) - } - async writeAsset( - params: ParsedDocumentStorageDocument & ParsedDocumentStorageAsset, + params: ParsedDocumentStorageDocument & LegacyParsedDocumentStorageAsset, ): Promise<{ readonly sourcePath: string; readonly url?: string }> { const blob = await this.blobStore.put( this.getObjectKey({ @@ -224,7 +160,7 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { } async getAssetUrl( - params: ParsedDocumentStorageAssetParams, + params: LegacyParsedDocumentStorageAssetParams, ): Promise { const result = await this.blobStore.head( this.getObjectKey({ @@ -246,7 +182,7 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { } async readSyncProgress( - params: ParsedDocumentStorageDocument, + params: ParsedDocumentRevisionParams, ): Promise { return this.readJson( this.getObjectKey({ @@ -269,14 +205,25 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { } async readObject( - params: ParsedDocumentStorageObject, - ): Promise { - return this.readBlobObject(this.getObjectKey(params)) + params: ParsedDocumentObjectParams, + ): Promise { + const object = await this.readBlobObject(this.getObjectKey(params)) + if (!object) return null + + return { + documentId: params.documentId, + revisionKey: params.revisionKey, + path: params.path, + body: object.body, + ...(object.contentType ? { contentType: object.contentType } : {}), + contentLength: object.body.byteLength, + url: object.url, + } } async writeObject( - params: ParsedDocumentStorageWritableObject, - ): Promise<{ readonly path: string; readonly url?: string }> { + params: ParsedDocumentWriteObjectParams, + ): Promise { const blob = await this.blobStore.put( this.getObjectKey(params), Buffer.from(params.body), @@ -287,20 +234,35 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { multipart: true, }, ) - return { path: params.path, url: blob.url } + return { + documentId: params.documentId, + revisionKey: params.revisionKey, + path: params.path, + url: blob.url, + } + } + + async headObject( + params: ParsedDocumentObjectParams, + ): Promise { + const result = await this.blobStore.head(this.getObjectKey(params)) + if (!result) return null + + return { + documentId: params.documentId, + revisionKey: params.revisionKey, + path: params.path, + url: result.url, + } } async getObjectUrl( - params: ParsedDocumentStorageObject, + params: ParsedDocumentObjectParams, ): Promise { const result = await this.blobStore.head(this.getObjectKey(params)) return result?.url ?? null } - async deleteObject(params: ParsedDocumentStorageObject): Promise { - await this.blobStore.del(this.getObjectKey(params)) - } - private getRevisionPrefix(documentId: string, revisionKey: string): string { return [ "workspaces", @@ -338,9 +300,11 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { return object ? object.body.toString("utf8") : null } - private async readBlobObject( - key: string, - ): Promise { + private async readBlobObject(key: string): Promise<{ + readonly body: Buffer + readonly contentType?: string + readonly url?: string + } | null> { try { const result = await this.blobStore.get(key, { access: "public" }) if (!result || result.statusCode !== 200) return null @@ -348,6 +312,7 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { return { body, ...(result.contentType ? { contentType: result.contentType } : {}), + url: result.url, } } catch (error) { if (error instanceof BlobNotFoundError) return null @@ -356,10 +321,6 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { } } -function getChunkPageStoragePath(page: number): string { - return `chunks/page-${page}.json` -} - /** * Reject path segments that could traverse outside the intended prefix. Mirrors * the SDK `DiskParsedDocumentStorage` guard so blob keys stay well-formed. diff --git a/src/integrations/knowhere.ts b/src/integrations/knowhere.ts index c0d074d..e55fc0b 100644 --- a/src/integrations/knowhere.ts +++ b/src/integrations/knowhere.ts @@ -14,10 +14,8 @@ import { BlobParsedDocumentStorage } from "@/domains/sources/parsed-document-blo * SDK returns `completed:false` and the caller re-enqueues to continue. */ const defaultParsedStorageLimits: ParsedDocumentStorageLimits = { - chunkPageSize: 200, remotePageSize: 100, maxPagesPerSync: 10, - maxAssetsPerSync: 20, syncDeadlineMs: 8000, grepMaxPages: 50, grepDeadlineMs: 8000, From 7c2bba4df77544f5b1cbd2593200fbe610ab4915 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 23:27:53 +0800 Subject: [PATCH 27/51] Keep home page route exports valid --- src/app/home-content.tsx | 28 ++++++++++++++++++++++++++++ src/app/page.test.ts | 2 +- src/app/page.tsx | 29 +---------------------------- 3 files changed, 30 insertions(+), 29 deletions(-) create mode 100644 src/app/home-content.tsx diff --git a/src/app/home-content.tsx b/src/app/home-content.tsx new file mode 100644 index 0000000..c2f1a64 --- /dev/null +++ b/src/app/home-content.tsx @@ -0,0 +1,28 @@ +import { connection } from "next/server" +import { WorkspaceShell } from "@/components/workspace-shell" +import { loadWorkspaceShellInitialState } from "@/domains/workspace/initial-state" +import { effectOperation } from "@/lib/effect-operation" +import { summarizeUnknownError } from "@/lib/format-log-value" +import { logger } from "@/lib/logger" + +export async function HomeContent() { + await connection() + const initialState = await loadWorkspaceInitialState() + return +} + +async function loadWorkspaceInitialState(): ReturnType< + typeof loadWorkspaceShellInitialState +> { + try { + return await loadWorkspaceShellInitialState() + } catch (error) { + logger.error("workspace: initial state failed", { + error: summarizeUnknownError(error), + }) + throw effectOperation.createBoundaryError( + "Workspace initial state failed", + error, + ) + } +} diff --git a/src/app/page.test.ts b/src/app/page.test.ts index f7e3525..3e9f811 100644 --- a/src/app/page.test.ts +++ b/src/app/page.test.ts @@ -20,7 +20,7 @@ vi.mock("@/domains/workspace/initial-state", () => ({ loadWorkspaceShellInitialState: mocks.loadWorkspaceShellInitialState, })) -import { HomeContent } from "./page" +import { HomeContent } from "./home-content" import { makeWorkspaceInitialStateFailureFixture } from "@/test/workspace-initial-state-failure-fixture" describe("Home", () => { diff --git a/src/app/page.tsx b/src/app/page.tsx index 2ebc597..0a3c967 100644 --- a/src/app/page.tsx +++ b/src/app/page.tsx @@ -1,10 +1,5 @@ import { Suspense } from "react" -import { WorkspaceShell } from "@/components/workspace-shell" -import { loadWorkspaceShellInitialState } from "@/domains/workspace/initial-state" -import { effectOperation } from "@/lib/effect-operation" -import { summarizeUnknownError } from "@/lib/format-log-value" -import { logger } from "@/lib/logger" -import { connection } from "next/server" +import { HomeContent } from "./home-content" export default function Home() { return ( @@ -13,25 +8,3 @@ export default function Home() { ) } - -export async function HomeContent() { - await connection() - const initialState = await loadWorkspaceInitialState() - return -} - -async function loadWorkspaceInitialState(): ReturnType< - typeof loadWorkspaceShellInitialState -> { - try { - return await loadWorkspaceShellInitialState() - } catch (error) { - logger.error("workspace: initial state failed", { - error: summarizeUnknownError(error), - }) - throw effectOperation.createBoundaryError( - "Workspace initial state failed", - error, - ) - } -} From 203f72cbad08237e4e41c814e0b95fcd612f6d6a Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 23:48:45 +0800 Subject: [PATCH 28/51] Install published Knowhere SDK 2.1.2 --- package.json | 2 +- pnpm-lock.yaml | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/package.json b/package.json index 94d1c3a..cf48c4a 100644 --- a/package.json +++ b/package.json @@ -28,7 +28,7 @@ "@effect/platform": "^0.96.1", "@napi-rs/canvas": "^1.0.2", "@neondatabase/serverless": "^1.1.0", - "@ontos-ai/knowhere-sdk": "^2.1.1", + "@ontos-ai/knowhere-sdk": "^2.1.2", "@radix-ui/react-alert-dialog": "^1.1.15", "@radix-ui/react-checkbox": "^1.3.3", "@radix-ui/react-dialog": "^1.1.15", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 48aabe5..12dc39f 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -24,8 +24,8 @@ importers: specifier: ^1.1.0 version: 1.1.0 '@ontos-ai/knowhere-sdk': - specifier: ^2.1.1 - version: 2.1.1 + specifier: ^2.1.2 + version: 2.1.2 '@radix-ui/react-alert-dialog': specifier: ^1.1.15 version: 1.1.15(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4) @@ -1603,8 +1603,8 @@ packages: resolution: {integrity: sha512-nn5ozdjYQpUCZlWGuxcJY/KpxkWQs4DcbMCmKojjyrYDEAGy4Ce19NN4v5MduafTwJlbKc99UA8YhSVqq9yPZA==} engines: {node: '>=12.4.0'} - '@ontos-ai/knowhere-sdk@2.1.1': - resolution: {integrity: sha512-K33ylB/QjVYlLgow+/Hw+uNhQbWWrMTIqZ3IAlx5buHSI2EHPVP9+kDsfslwpfffYSn85j9Ab0/E9qTG0QotTw==} + '@ontos-ai/knowhere-sdk@2.1.2': + resolution: {integrity: sha512-m1wbQZNesExcg2yoEnmqFRlML/tprv11y/v7tlTLXEKTqKOY5zpG67idxXtbmw2qjOjQkR02HhgUwo/SdrFWPA==} engines: {node: '>=22.13.0', npm: '>=10.0.0', pnpm: '>=9.0.0'} '@open-draft/deferred-promise@2.2.0': @@ -6731,7 +6731,7 @@ snapshots: '@nolyfill/is-core-module@1.0.39': {} - '@ontos-ai/knowhere-sdk@2.1.1': + '@ontos-ai/knowhere-sdk@2.1.2': dependencies: axios: 1.18.1 jszip: 3.10.1 From 3fe820627c47f61807f31fef192311186c841082 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 8 Jul 2026 01:03:55 +0800 Subject: [PATCH 29/51] fix(sources): render page assets from chunks --- .../[sourceId]/page-assets/route.test.ts | 554 ------------------ .../sources/[sourceId]/page-assets/route.ts | 32 - src/components/chunks-panel.test.ts | 215 +++---- src/components/chunks-panel.tsx | 332 +---------- src/components/parsed-chunk-card.test.ts | 44 +- src/components/parsed-chunk-card.tsx | 112 +++- .../workspace-citation-focus.test.ts | 28 +- src/components/workspace-citation-focus.ts | 12 - .../workspace-citation-state.test.ts | 54 ++ src/components/workspace-citation-state.ts | 65 +- .../workspace-selected-chunks.test.ts | 89 +-- src/components/workspace-selected-chunks.ts | 123 ++-- src/domains/chunks/index.test.ts | 36 ++ src/domains/chunks/normalization.ts | 51 ++ src/domains/chunks/types.ts | 3 + src/domains/sources/page-assets.ts | 218 ------- src/domains/sources/route-page-assets.ts | 208 ------- src/domains/sources/route-service.ts | 5 - src/domains/sources/route-types.ts | 38 -- src/domains/workspace/client-cache.ts | 4 - src/domains/workspace/client.test.ts | 30 - src/domains/workspace/client.ts | 34 -- 22 files changed, 564 insertions(+), 1723 deletions(-) delete mode 100644 src/app/api/sources/[sourceId]/page-assets/route.test.ts delete mode 100644 src/app/api/sources/[sourceId]/page-assets/route.ts delete mode 100644 src/domains/sources/page-assets.ts delete mode 100644 src/domains/sources/route-page-assets.ts diff --git a/src/app/api/sources/[sourceId]/page-assets/route.test.ts b/src/app/api/sources/[sourceId]/page-assets/route.test.ts deleted file mode 100644 index 4f00bd3..0000000 --- a/src/app/api/sources/[sourceId]/page-assets/route.test.ts +++ /dev/null @@ -1,554 +0,0 @@ -import { NextRequest } from "next/server" -import { beforeEach, describe, expect, it, vi } from "vitest" - -const mocks = vi.hoisted(() => ({ - deleteBlob: vi.fn(), - ensureApiKeyForWorkspace: vi.fn(), - ensureWorkspace: vi.fn(), - findSourceInWorkspace: vi.fn(), - getCurrentUser: vi.fn(), - listChunks: vi.fn(), - makeKnowhereClient: vi.fn(), - makeKnowhereClientWithParsedStorage: vi.fn(), - readChunks: vi.fn(), - requireUser: vi.fn(), -})) - -vi.mock("next/headers", () => ({ - headers: vi.fn(async () => new Headers({ cookie: "session=abc" })), -})) - -vi.mock("@/integrations/dashboard/api-key-service", () => ({ - ensureApiKeyForWorkspace: mocks.ensureApiKeyForWorkspace, -})) - -vi.mock("@/integrations/knowhere-demo", () => ({ - knowhereDemoApi: { - fetchCatalog: vi.fn(), - fetchChunkPage: vi.fn(), - }, -})) - -vi.mock("@/infrastructure/auth", () => ({ - getCurrentUser: mocks.getCurrentUser, - requireUser: mocks.requireUser, -})) - -vi.mock("@/integrations/knowhere", () => ({ - makeKnowhereClient: mocks.makeKnowhereClient, - makeKnowhereClientWithParsedStorage: - mocks.makeKnowhereClientWithParsedStorage, -})) - -vi.mock("@vercel/blob", () => ({ - del: mocks.deleteBlob, -})) - -vi.mock("@/domains/sources/service", () => ({ - sourceService: { - findInWorkspace: mocks.findSourceInWorkspace, - localizeRemoteDocument: vi.fn(), - }, -})) - -vi.mock("@/domains/workspace/service", () => ({ - workspaceService: { - ensureWorkspace: mocks.ensureWorkspace, - }, -})) - -import { GET } from "./route" - -describe("GET /api/sources/[sourceId]/page-assets", () => { - beforeEach(() => { - vi.clearAllMocks() - mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ - client: { documents: { listChunks: mocks.listChunks } }, - knowledge: { readChunks: mocks.readChunks }, - }) - }) - - it("returns stored page assets for a ready workspace source without durable hardening", async () => { - mocks.getCurrentUser.mockResolvedValue({ - id: "user_1", - email: null, - name: null, - }) - mocks.ensureWorkspace.mockResolvedValue({ - id: "workspace_1", - userId: "user_1", - namespace: "notebook-workspace_1", - createdAt: new Date("2026-05-10T00:00:00.000Z"), - }) - mocks.findSourceInWorkspace.mockResolvedValue( - makeReadySource({ - id: "00000000-0000-0000-0000-000000000002", - knowhereJobId: "job_1", - knowhereDocumentId: "doc_1", - }), - ) - mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") - mocks.readChunks.mockResolvedValue({ - document: { - localDocumentId: "doc_1", - resultDirectoryPath: "parsed-storage:doc_1", - }, - chunks: [ - { - chunkId: "page_1", - chunkType: "page", - assetUrl: "https://assets.example/fallback.png", - metadata: { - pageAssets: [ - { - pageNum: 1, - assetUrl: "https://assets.example/page-000001.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - }, - }, - ], - page: 1, - pageSize: 1, - totalChunks: 3, - totalPages: 3, - }) - - const response = await GET( - new NextRequest( - "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", - ), - { - params: Promise.resolve({ - sourceId: "00000000-0000-0000-0000-000000000002", - }), - }, - ) - - await expect(response.json()).resolves.toEqual({ - pages: [ - { - pageNumber: 1, - assetUrl: "https://assets.example/page-000001.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - pagination: { - page: 1, - pageSize: 1, - total: 3, - totalPages: 3, - }, - }) - expect(response.status).toBe(200) - expect(mocks.readChunks).toHaveBeenCalledWith({ - documentId: "doc_1", - revisionKey: "job_1", - chunkType: "page", - page: 1, - pageSize: 1, - }) - expect(mocks.listChunks).not.toHaveBeenCalled() - }) - - it("uses SDK page asset URLs directly when the SDK remote fallback is usable", async () => { - mocks.getCurrentUser.mockResolvedValue({ - id: "user_1", - email: null, - name: null, - }) - mocks.ensureWorkspace.mockResolvedValue({ - id: "workspace_1", - userId: "user_1", - namespace: "notebook-workspace_1", - createdAt: new Date("2026-05-10T00:00:00.000Z"), - }) - mocks.findSourceInWorkspace.mockResolvedValue( - makeReadySource({ - id: "00000000-0000-0000-0000-000000000002", - knowhereJobId: "job_1", - knowhereDocumentId: "doc_1", - }), - ) - mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") - mocks.readChunks.mockResolvedValue({ - document: { - localDocumentId: "doc_1", - resultDirectoryPath: "remote:doc_1", - }, - chunks: [ - { - chunkId: "page_1", - chunkType: "page", - metadata: { - pageAssets: [ - { - pageNum: 1, - assetUrl: "https://sdk.example/page-000001.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - }, - }, - ], - page: 1, - pageSize: 1, - totalChunks: 1, - totalPages: 1, - }) - - const response = await GET( - new NextRequest( - "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", - ), - { - params: Promise.resolve({ - sourceId: "00000000-0000-0000-0000-000000000002", - }), - }, - ) - - await expect(response.json()).resolves.toEqual({ - pages: [ - { - pageNumber: 1, - assetUrl: "https://sdk.example/page-000001.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - }) - expect(response.status).toBe(200) - expect(mocks.readChunks).toHaveBeenCalledWith({ - documentId: "doc_1", - revisionKey: "job_1", - chunkType: "page", - page: 1, - pageSize: 1, - }) - expect(mocks.listChunks).not.toHaveBeenCalled() - }) - - it("rejects non-ready workspace sources", async () => { - mocks.getCurrentUser.mockResolvedValue({ - id: "user_1", - email: null, - name: null, - }) - mocks.ensureWorkspace.mockResolvedValue({ - id: "workspace_1", - userId: "user_1", - namespace: "notebook-workspace_1", - createdAt: new Date("2026-05-10T00:00:00.000Z"), - }) - mocks.findSourceInWorkspace.mockResolvedValue( - makeReadySource({ - id: "00000000-0000-0000-0000-000000000002", - status: "parsing", - knowhereDocumentId: "doc_1", - }), - ) - - const response = await GET( - new NextRequest( - "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", - ), - { - params: Promise.resolve({ - sourceId: "00000000-0000-0000-0000-000000000002", - }), - }, - ) - - await expect(response.json()).resolves.toEqual({ - message: "Source is not ready.", - }) - expect(response.status).toBe(409) - expect(mocks.readChunks).not.toHaveBeenCalled() - }) - - it("returns an empty page list when page chunks have no usable assets", async () => { - mocks.getCurrentUser.mockResolvedValue({ - id: "user_1", - email: null, - name: null, - }) - mocks.ensureWorkspace.mockResolvedValue({ - id: "workspace_1", - userId: "user_1", - namespace: "notebook-workspace_1", - createdAt: new Date("2026-05-10T00:00:00.000Z"), - }) - mocks.findSourceInWorkspace.mockResolvedValue( - makeReadySource({ - id: "00000000-0000-0000-0000-000000000002", - knowhereDocumentId: "doc_1", - }), - ) - mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") - mocks.readChunks.mockResolvedValue({ - document: { - localDocumentId: "doc_1", - resultDirectoryPath: "remote:doc_1", - }, - chunks: [ - { - chunkId: "page_1", - chunkType: "page", - metadata: { pageAssets: [] }, - }, - ], - page: 1, - pageSize: 1, - totalChunks: 1, - totalPages: 1, - }) - mocks.listChunks.mockResolvedValue({ - chunks: [], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - }) - - const response = await GET( - new NextRequest( - "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", - ), - { - params: Promise.resolve({ - sourceId: "00000000-0000-0000-0000-000000000002", - }), - }, - ) - - await expect(response.json()).resolves.toEqual({ - pages: [], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - }) - expect(response.status).toBe(200) - expect(mocks.listChunks).toHaveBeenCalledWith("doc_1", { - page: 1, - pageSize: 1, - chunkType: "page", - includeAssetUrls: true, - }) - }) - - it("returns Knowhere page asset URLs when the storage probe falls through to remote chunks", async () => { - mocks.getCurrentUser.mockResolvedValue({ - id: "user_1", - email: null, - name: null, - }) - mocks.ensureWorkspace.mockResolvedValue({ - id: "workspace_1", - userId: "user_1", - namespace: "notebook-workspace_1", - createdAt: new Date("2026-05-10T00:00:00.000Z"), - }) - mocks.findSourceInWorkspace.mockResolvedValue( - makeReadySource({ - id: "00000000-0000-0000-0000-000000000002", - knowhereDocumentId: "doc_1", - }), - ) - mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") - mocks.readChunks.mockResolvedValue({ - document: { - localDocumentId: "doc_1", - resultDirectoryPath: "remote:doc_1", - }, - chunks: [ - { - chunkId: "page_1", - chunkType: "page", - metadata: { - pageAssets: [ - { - pageNum: 1, - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - }, - }, - ], - page: 1, - pageSize: 1, - totalChunks: 1, - totalPages: 1, - }) - mocks.listChunks.mockResolvedValue({ - chunks: [ - { - id: "document_page_1", - chunkId: "page_1", - chunkType: "page", - content: "Page 1", - sectionId: null, - sectionPath: "pages/1", - sourceChunkPath: "pages/1", - filePath: "pages/page-000001.png", - sortOrder: 0, - metadata: { - pageAssets: [ - { - pageNum: 1, - assetUrl: "https://knowhere.example/page-000001.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - }, - assetUrl: "https://knowhere.example/page-000001.png", - }, - ], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - }) - - const response = await GET( - new NextRequest( - "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", - ), - { - params: Promise.resolve({ - sourceId: "00000000-0000-0000-0000-000000000002", - }), - }, - ) - - await expect(response.json()).resolves.toEqual({ - pages: [ - { - pageNumber: 1, - assetUrl: "https://knowhere.example/page-000001.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - }) - expect(response.status).toBe(200) - expect(mocks.readChunks).toHaveBeenCalledWith({ - documentId: "doc_1", - revisionKey: "job_1", - chunkType: "page", - page: 1, - pageSize: 1, - }) - expect(mocks.listChunks).toHaveBeenCalledWith("doc_1", { - page: 1, - pageSize: 1, - chunkType: "page", - includeAssetUrls: true, - }) - }) - - it("marks page assets unavailable when the parsed document is missing remotely", async () => { - const notFoundError = new Error("Document not found") - notFoundError.name = "NotFoundError" - mocks.getCurrentUser.mockResolvedValue({ - id: "user_1", - email: null, - name: null, - }) - mocks.ensureWorkspace.mockResolvedValue({ - id: "workspace_1", - userId: "user_1", - namespace: "notebook-workspace_1", - createdAt: new Date("2026-05-10T00:00:00.000Z"), - }) - mocks.findSourceInWorkspace.mockResolvedValue( - makeReadySource({ - id: "00000000-0000-0000-0000-000000000002", - knowhereDocumentId: "doc_missing", - }), - ) - mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") - mocks.readChunks.mockRejectedValue(notFoundError) - - const response = await GET( - new NextRequest( - "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", - ), - { - params: Promise.resolve({ - sourceId: "00000000-0000-0000-0000-000000000002", - }), - }, - ) - - await expect(response.json()).resolves.toEqual({ - pages: [], - pagination: { - page: 1, - pageSize: 1, - total: 0, - totalPages: 0, - }, - message: - "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", - isUnavailable: true, - }) - expect(response.status).toBe(200) - }) -}) - -function makeReadySource(overrides: Record) { - return { - id: "00000000-0000-0000-0000-000000000002", - workspaceId: "workspace_1", - title: "notes.pdf", - mimeType: "application/pdf", - sizeBytes: 1024, - status: "ready", - failureReason: null, - failureStage: null, - knowhereJobId: "job_1", - knowhereDocumentId: "doc_1", - stagedBlobPathname: null, - stagedBlobUrl: null, - originalBlobPathname: null, - originalBlobUrl: null, - demoKey: null, - createdAt: new Date("2026-05-10T00:00:00.000Z"), - updatedAt: new Date("2026-05-10T00:00:00.000Z"), - deletedAt: null, - ...overrides, - } -} diff --git a/src/app/api/sources/[sourceId]/page-assets/route.ts b/src/app/api/sources/[sourceId]/page-assets/route.ts deleted file mode 100644 index b7606f7..0000000 --- a/src/app/api/sources/[sourceId]/page-assets/route.ts +++ /dev/null @@ -1,32 +0,0 @@ -import type { NextRequest, NextResponse } from "next/server" - -import { getChunkPageParams } from "@/domains/chunks" -import { createSourceRouteService } from "@/domains/sources/route-service" -import { withApiErrorResponse } from "@/lib/api-error-response" -import { nextRouteContext } from "@/lib/next-route-context" -import { nextRouteResponse } from "@/lib/next-route-response" - -type RouteContext = { - params: Promise<{ - sourceId: string - }> -} - -const sourceRouteService = createSourceRouteService() - -export async function GET( - request: NextRequest, - context: RouteContext, -): Promise { - return withApiErrorResponse("sources:page-assets", async () => { - const { sourceId } = await context.params - const routeContext = await nextRouteContext.read() - const result = await sourceRouteService.loadSourcePageAssets({ - cookieHeader: routeContext.cookieHeader, - pageParams: getChunkPageParams(request.nextUrl.searchParams), - sourceId, - }) - - return nextRouteResponse.toNextResponse(result) - }) -} diff --git a/src/components/chunks-panel.test.ts b/src/components/chunks-panel.test.ts index 8e06155..78a35f7 100644 --- a/src/components/chunks-panel.test.ts +++ b/src/components/chunks-panel.test.ts @@ -14,7 +14,6 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { ChunksPanel } from "./chunks-panel"; import { sourceOriginalPreviewRequest } from "./source-original-preview-request"; -const fetchPageAssetPageMock = vi.hoisted(() => vi.fn()); const C = ChunksPanel as React.FC>; const virtualizerScrollResetDelayMs = 150; @@ -30,25 +29,9 @@ vi.mock("react-pdf", () => ({ Page: () => React.createElement("div", { "data-testid": "pdf-page" }), })); -vi.mock("@/domains/workspace/client", () => ({ - workspaceClient: { - fetchPageAssetPage: fetchPageAssetPageMock, - }, -})); - describe("ChunksPanel", () => { beforeEach(() => { shouldFlushVirtualizerTimers = false; - fetchPageAssetPageMock.mockReset(); - fetchPageAssetPageMock.mockResolvedValue({ - pages: [], - pagination: { - page: 1, - pageSize: 50, - total: 0, - totalPages: 1, - }, - }); globalThis.ResizeObserver = class ResizeObserver { observe() {} unobserve() {} @@ -89,28 +72,35 @@ describe("ChunksPanel", () => { expect(screen.getByText(/Showing all parsed chunks from/)).toBeTruthy(); }); - it("renders page assets instead of parsed chunk controls for page-mode documents", async () => { - fetchPageAssetPageMock.mockResolvedValue({ - pages: [ - { - pageNumber: 4, - assetUrl: "https://assets.example/page-000004.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - pagination: { - page: 1, - pageSize: 50, - total: 4, - totalPages: 1, - }, - }); + it("renders page chunk assets inside the normal chunk list", async () => { + mockVisibleVirtualViewport(); render( React.createElement(C, { - chunks: [], + chunks: [ + { + chunkId: "page_4", + type: "page", + content: "Page 4 summary", + sourceTitle: "report.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + { + chunkId: "table_1", + type: "table", + content: "
Budget
", + sourceTitle: "report.pdf", + }, + ], selectedSource: "report.pdf", selectedSourceView: { id: "source_1", @@ -122,124 +112,74 @@ describe("ChunksPanel", () => { }), ); - await waitFor(() => - expect( - screen.getByTestId("page-asset-document-viewer"), - ).toBeTruthy(), - ); - expect(screen.getByRole("heading", { name: "Original File" })).toBeTruthy(); - expect(screen.getByText("Page 4")).toBeTruthy(); - expect(screen.queryByRole("button", { name: "Parsed" })).toBeNull(); + expect(screen.getByRole("heading", { name: "Parsed Chunks" })).toBeTruthy(); + expect( + await screen.findByRole("img", { name: "Page 4" }), + ).toBeTruthy(); + expect(screen.getByText("image/png")).toBeTruthy(); + expect(screen.getByText("Budget")).toBeTruthy(); + expect( + screen.queryByRole("button", { name: "Original" }), + ).toBeNull(); expect(screen.queryByRole("button", { name: "Tree" })).toBeNull(); - expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 1); + expect( + screen.queryByRole("button", { name: /original file/i }), + ).toBeNull(); }); - it("shows a page-level placeholder when a page asset image fails to load", async () => { - fetchPageAssetPageMock.mockResolvedValue({ - pages: [ - { - pageNumber: 4, - assetUrl: "https://assets.example/page-000004.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - pagination: { - page: 1, - pageSize: 50, - total: 4, - totalPages: 1, - }, - }); + it("renders page chunks normally when no page assets exist", async () => { + mockVisibleVirtualViewport(); render( React.createElement(C, { - chunks: [], - selectedSource: "report.pdf", - selectedSourceView: { - id: "source_1", - title: "report.pdf", - mimeType: "application/pdf", - status: "ready", - documentPresentation: { kind: "page-assets", pageCount: 4 }, - }, - }), - ); - - const pageImage = await screen.findByRole("img", { name: "Page 4" }); - fireEvent.error(pageImage); - - expect( - screen.getByTestId("page-asset-image-unavailable-4"), - ).toBeTruthy(); - expect(screen.getByText("Page image unavailable.")).toBeTruthy(); - }); - - it("loads missing page buckets after focusing a later page asset", async () => { - const user = userEvent.setup(); - fetchPageAssetPageMock.mockImplementation( - async (_sourceId: string, page: number) => ({ - pages: [ + chunks: [ { - pageNumber: page === 3 ? 101 : page, - assetUrl: `https://assets.example/page-${page}.png`, - contentType: "image/png", + chunkId: "page_4", + type: "page", + content: "Page 4 summary", + readableContent: "Page 4 summary", + sourceTitle: "report.pdf", + pageNums: [4], }, ], - pagination: { - page, - pageSize: 50, - total: 120, - totalPages: 3, - }, - }), - ); - - render( - React.createElement(C, { - chunks: [], selectedSource: "report.pdf", selectedSourceView: { id: "source_1", title: "report.pdf", mimeType: "application/pdf", status: "ready", - documentPresentation: { kind: "page-assets", pageCount: 120 }, + documentPresentation: { kind: "page-assets", pageCount: 4 }, }, - focusedPageNumber: 101, - focusedPageRequestId: 1, }), ); - await waitFor(() => - expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 3), - ); - - await user.click(screen.getByRole("button", { name: "Load more pages" })); - - await waitFor(() => - expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 2), - ); + expect(await screen.findByText("Page 4 summary")).toBeTruthy(); + expect(screen.queryByRole("img", { name: "Page 4" })).toBeNull(); }); - it("shows page asset unavailable messages", async () => { - fetchPageAssetPageMock.mockResolvedValue({ - pages: [], - isUnavailable: true, - message: - "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", - pagination: { - page: 1, - pageSize: 50, - total: 0, - totalPages: 0, - }, - }); + it("shows a page-level placeholder when a page asset image fails to load", async () => { + mockVisibleVirtualViewport(); render( React.createElement(C, { - chunks: [], + chunks: [ + { + chunkId: "page_4", + type: "page", + content: "Page 4 summary", + sourceTitle: "report.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + ], selectedSource: "report.pdf", selectedSourceView: { id: "source_1", @@ -251,14 +191,13 @@ describe("ChunksPanel", () => { }), ); - await waitFor(() => - expect( - screen.getByText( - "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", - ), - ).toBeTruthy(), - ); - expect(screen.queryByText("No page images available.")).toBeNull(); + const pageImage = await screen.findByRole("img", { name: "Page 4" }); + fireEvent.error(pageImage); + + expect( + screen.getByTestId("page-asset-image-unavailable-4"), + ).toBeTruthy(); + expect(screen.getByText("Page image unavailable.")).toBeTruthy(); }); it("defaults parsed chunks into a section tree view", () => { diff --git a/src/components/chunks-panel.tsx b/src/components/chunks-panel.tsx index fc785a2..71037c9 100644 --- a/src/components/chunks-panel.tsx +++ b/src/components/chunks-panel.tsx @@ -19,7 +19,6 @@ import { } from "d3-hierarchy"; import { FilePlus2, - ImageOff, Layers, RotateCcw, UploadCloud, @@ -42,11 +41,9 @@ import { chunksPanelState } from "@/components/chunks-panel-state"; import { MAX_UPLOAD_MB } from "@/domains/sources/validation"; import { useSourceOriginalPreviewWarmup } from "@/components/source-original-preview-warmup"; import { sourceOriginalPreviewModel } from "@/components/source-original-preview-model"; -import { workspaceClient } from "@/domains/workspace/client"; import type { ParsedChunkView } from "@/domains/chunks/types"; import type { SourceOriginalFileView, - SourcePageAssetView, SourceView, } from "@/domains/sources/types"; import type { AnalyticsContext } from "@/lib/posthog"; @@ -90,8 +87,6 @@ export function ChunksPanel({ selectedSourceFile = null, focusedChunkId = null, focusedChunkRequestId = 0, - focusedPageNumber = null, - focusedPageRequestId = 0, citationListViewRequestId = 0, isLoading = false, isLoadingMore = false, @@ -105,10 +100,16 @@ export function ChunksPanel({ analyticsContext, sourceCountSnapshot = 0, }: Partial = {}) { + const hasPageAssetChunks = chunks.some( + (chunk) => (chunk.pageAssets?.length ?? 0) > 0, + ); const isPageAssetSource = - selectedSourceView?.documentPresentation?.kind === "page-assets"; + selectedSourceView?.documentPresentation?.kind === "page-assets" || + hasPageAssetChunks; + const effectiveVisibleView = isPageAssetSource ? "parsed" : undefined; const originalPreviewCacheKey = selectedSourceFile?.url ?? null; const isOriginalPreviewAvailable = + !isPageAssetSource && sourceOriginalPreviewModel.canPreviewOriginalFile( selectedSource, selectedSourceFile, @@ -153,6 +154,7 @@ export function ChunksPanel({ isLoadingMore, onLoadMore, }); + const activeVisibleView = effectiveVisibleView ?? visibleView; useSourceOriginalPreviewWarmup({ sourceTitle: selectedSource, @@ -254,29 +256,21 @@ export function ChunksPanel({ citationListViewRequestId ? "list" : chunkDisplayModeState.mode; - const headerTitle = isPageAssetSource || visibleView === "original" + const headerTitle = activeVisibleView === "original" ? "Original File" : focusedChunkId ? "Referenced Chunks" : "Parsed Chunks"; const shouldMountOriginalPreview = - visibleView === "original" || - (originalPreviewCacheKey !== null && - mountedOriginalPreviewKey === originalPreviewCacheKey); + !isPageAssetSource && + (activeVisibleView === "original" || + (originalPreviewCacheKey !== null && + mountedOriginalPreviewKey === originalPreviewCacheKey)); const isTreeModeVisible = - !isPageAssetSource && visibleView === "parsed" && chunkDisplayMode === "tree"; - const headerSubtitle = isPageAssetSource ? ( - selectedSource ? ( - <> - Showing page images for{" "} - - {selectedSource} - - - ) : ( - "Select a source to preview its page images." - ) - ) : visibleView === "original" ? ( + !isPageAssetSource && + activeVisibleView === "parsed" && + chunkDisplayMode === "tree"; + const headerSubtitle = activeVisibleView === "original" ? ( selectedSource ? ( <> Showing the original file for{" "} @@ -335,7 +329,9 @@ export function ChunksPanel({

- {!isPageAssetSource && visibleView === "parsed" && chunks.length > 0 ? ( + {!isPageAssetSource && + activeVisibleView === "parsed" && + chunks.length > 0 ? (
Parsed @@ -379,7 +375,7 @@ export function ChunksPanel({
- + - {isPageAssetSource && selectedSourceView ? ( - - ) : isLoading ? ( + {isLoading ? ( ) : chunks.length === 0 && processingMessage ? ( @@ -437,7 +427,9 @@ export function ChunksPanel({ isOriginalPreviewAvailable={isOriginalPreviewAvailable} measureElement={measureVirtualChunkElement} onChunkClick={ - hasOriginalFile ? handleChunkSelected : undefined + hasOriginalFile && !isPageAssetSource + ? handleChunkSelected + : undefined } onReferenceClick={requestChunkFocus} selectedSourceFile={selectedSourceFile} @@ -472,7 +464,7 @@ export function ChunksPanel({ ) : null} {shouldMountOriginalPreview ? ( - + ; - readonly isLoading: boolean; - readonly message: string | null; - readonly totalPages: number; -}; - -const pageAssetPageSize = 50; - -function PageAssetDocumentViewer({ - focusedPageNumber, - focusedPageRequestId, - source, -}: { - readonly focusedPageNumber: number | null; - readonly focusedPageRequestId: number; - readonly source: SourceView; -}): ReactNode { - const [pageState, setPageState] = useState(() => ({ - pages: [], - loadedPageIndexes: new Set(), - isLoading: false, - message: null, - totalPages: Math.max( - 1, - Math.ceil( - (source.documentPresentation?.kind === "page-assets" - ? source.documentPresentation.pageCount - : 0) / pageAssetPageSize, - ), - ), - })); - const requestedPageIndexesRef = useRef>(new Set()); - const pageElementsRef = useRef>(new Map()); - const sourceIdRef = useRef(source.id); - - const loadPageIndex = useCallback( - (pageIndex: number): void => { - if (requestedPageIndexesRef.current.has(pageIndex)) return; - requestedPageIndexesRef.current.add(pageIndex); - setPageState((current) => ({ ...current, isLoading: true })); - - void workspaceClient - .fetchPageAssetPage(source.id, pageIndex) - .then((response) => { - setPageState((current) => { - const pagesByPageNumber = new Map( - current.pages.map((page) => [page.pageNumber, page]), - ); - for (const page of response.pages ?? []) { - pagesByPageNumber.set(page.pageNumber, page); - } - const loadedPageIndexes = new Set(current.loadedPageIndexes); - loadedPageIndexes.add(pageIndex); - - return { - pages: [...pagesByPageNumber.values()].sort( - (left, right) => left.pageNumber - right.pageNumber, - ), - loadedPageIndexes, - isLoading: false, - message: response.message ?? null, - totalPages: - response.pagination?.totalPages ?? current.totalPages, - }; - }); - }) - .catch(() => { - requestedPageIndexesRef.current.delete(pageIndex); - setPageState((current) => ({ ...current, isLoading: false })); - }); - }, - [source.id], - ); - - useEffect(() => { - if (sourceIdRef.current === source.id) return; - - sourceIdRef.current = source.id; - requestedPageIndexesRef.current = new Set(); - pageElementsRef.current = new Map(); - setPageState({ - pages: [], - loadedPageIndexes: new Set(), - isLoading: false, - message: null, - totalPages: Math.max( - 1, - Math.ceil( - (source.documentPresentation?.kind === "page-assets" - ? source.documentPresentation.pageCount - : 0) / pageAssetPageSize, - ), - ), - }); - }, [source.documentPresentation, source.id]); - - useEffect(() => { - loadPageIndex(1); - }, [loadPageIndex]); - - useEffect(() => { - if (!focusedPageNumber) return; - - loadPageIndex(Math.max(1, Math.ceil(focusedPageNumber / pageAssetPageSize))); - }, [focusedPageNumber, focusedPageRequestId, loadPageIndex]); - - useEffect(() => { - if (!focusedPageNumber) return; - - const element = pageElementsRef.current.get(focusedPageNumber); - if (!element) return; - - element.scrollIntoView({ behavior: "smooth", block: "start" }); - }, [focusedPageNumber, focusedPageRequestId, pageState.pages]); - - const nextPageIndex = getNextPageAssetPageIndex( - pageState.loadedPageIndexes, - pageState.totalPages, - ); - const canLoadMore = - nextPageIndex !== null && - nextPageIndex <= pageState.totalPages && - !pageState.isLoading; - - if (pageState.pages.length === 0 && pageState.isLoading) { - return ; - } - - if (pageState.pages.length === 0 && pageState.message) { - return ; - } - - if (pageState.pages.length === 0) { - return ; - } - - return ( -
- {pageState.pages.map((page) => ( - { - if (element) { - pageElementsRef.current.set(page.pageNumber, element); - return; - } - pageElementsRef.current.delete(page.pageNumber); - }} - /> - ))} - {canLoadMore ? ( -
- -
- ) : null} - {pageState.isLoading ? ( -
- Loading page images... -
- ) : null} -
- ); -} - -function PageAssetImage({ - isFocused, - page, - refCallback, -}: { - readonly isFocused: boolean; - readonly page: SourcePageAssetView; - readonly refCallback: (element: HTMLDivElement | null) => void; -}): ReactNode { - const [failedAssetUrl, setFailedAssetUrl] = useState(null); - const hasImageError = failedAssetUrl === page.assetUrl; - const aspectRatio = - page.width && page.height ? `${page.width} / ${page.height}` : undefined; - - return ( -
-
- Page {page.pageNumber} - {page.contentType} -
-
- {hasImageError ? ( - - ) : ( - // eslint-disable-next-line @next/next/no-img-element -- Page assets can be short-lived Knowhere URLs outside Next image optimization. - {`Page setFailedAssetUrl(page.assetUrl)} - /> - )} -
-
- ); -} - -function PageAssetImageUnavailable({ - pageNumber, -}: { - readonly pageNumber: number; -}): ReactNode { - return ( -
-
- -
-

- Page image unavailable. -

-
- ); -} - -function getNextPageAssetPageIndex( - loadedPageIndexes: ReadonlySet, - totalPages: number, -): number | null { - for (let pageIndex = 1; pageIndex <= totalPages; pageIndex += 1) { - if (!loadedPageIndexes.has(pageIndex)) return pageIndex; - } - - return null; -} - -function EmptyPageAssets(): ReactNode { - return ( -
- -

- No page images available. -

-
- ); -} - function UnavailableSourceMessage({ message, }: { diff --git a/src/components/parsed-chunk-card.test.ts b/src/components/parsed-chunk-card.test.ts index 946cffb..ff4fbef 100644 --- a/src/components/parsed-chunk-card.test.ts +++ b/src/components/parsed-chunk-card.test.ts @@ -1,6 +1,6 @@ // @vitest-environment jsdom import React from "react"; -import { cleanup, render, screen } from "@testing-library/react"; +import { cleanup, fireEvent, render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, describe, expect, it, vi } from "vitest"; @@ -81,6 +81,48 @@ describe("ParsedChunkCard", () => { ).toBeTruthy(); }); + it("renders page citation assets instead of page summary content", () => { + render( + React.createElement(ParsedChunkCard, { + chunk: { + chunkId: "page_4", + type: "page", + content: "The summary should not be primary when an image exists.", + readableContent: "The summary should not be primary when an image exists.", + sourceTitle: "manual.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + isFocused: false, + isOriginalPreviewAvailable: true, + onChunkClick: vi.fn(), + onReferenceClick: vi.fn(), + }), + ); + + const pageImage = screen.getByRole("img", { name: "Page 4" }); + + expect(pageImage.getAttribute("src")).toBe( + "https://assets.example/page-4.png", + ); + expect(screen.getByText("image/png")).toBeTruthy(); + expect(screen.queryByText(/summary should not be primary/i)).toBeNull(); + expect( + screen.queryByRole("button", { name: /original file/i }), + ).toBeNull(); + + fireEvent.error(pageImage); + expect(screen.getByTestId("page-asset-image-unavailable-4")).toBeTruthy(); + }); + it("routes resolved artifact reference clicks to the target chunk", async () => { const user = userEvent.setup(); const onReferenceClick = vi.fn(); diff --git a/src/components/parsed-chunk-card.tsx b/src/components/parsed-chunk-card.tsx index 1ba233d..d794687 100644 --- a/src/components/parsed-chunk-card.tsx +++ b/src/components/parsed-chunk-card.tsx @@ -1,7 +1,15 @@ "use client"; -import { useMemo, type MouseEvent, type ReactNode } from "react"; -import { FileSearch, FileText, ImageIcon, Table2, Tags, TextQuote } from "lucide-react"; +import { useMemo, useState, type MouseEvent, type ReactNode } from "react"; +import { + FileSearch, + FileText, + ImageIcon, + ImageOff, + Table2, + Tags, + TextQuote, +} from "lucide-react"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; @@ -185,7 +193,9 @@ function ChunkSourcePanel({ ) : null}
- {onChunkClick && firstPageNumber !== null ? ( + {onChunkClick && + firstPageNumber !== null && + !hasPageCitationAssets(chunk) ? ( void; }): ReactNode { + const pageAssets = chunk.pageAssets ?? []; + return ( - -

- {chunk.readableContent ?? chunk.content} -

- + {pageAssets.length > 0 ? ( + + + + ) : ( + +

+ {chunk.readableContent ?? chunk.content} +

+ + )} ); } +function PageCitationAssets({ + assets, +}: { + readonly assets: NonNullable; +}): ReactNode { + return ( +
+ {assets.map((asset) => ( + + ))} +
+ ); +} + +function PageCitationAssetImage({ + asset, +}: { + readonly asset: NonNullable[number]; +}): ReactNode { + const [failedAssetUrl, setFailedAssetUrl] = useState(null); + const hasImageError = failedAssetUrl === asset.assetUrl; + const aspectRatio = + asset.width && asset.height ? `${asset.width} / ${asset.height}` : undefined; + + return ( +
+
+ Page {asset.pageNumber} + {asset.contentType} +
+
+ {hasImageError ? ( + + ) : ( + // eslint-disable-next-line @next/next/no-img-element -- Page assets can be short-lived Knowhere URLs outside Next image optimization. + {`Page setFailedAssetUrl(asset.assetUrl)} + /> + )} +
+
+ ); +} + +function PageCitationAssetUnavailable({ + pageNumber, +}: { + readonly pageNumber: number; +}): ReactNode { + return ( +
+
+ +
+

+ Page image unavailable. +

+
+ ); +} + function ImageChunkCard({ chunk, isFocused, @@ -467,6 +559,10 @@ function getImageChunkAssetUrl( return sourceOriginalFile.url; } +function hasPageCitationAssets(chunk: ParsedChunkView): boolean { + return (chunk.pageAssets?.length ?? 0) > 0; +} + function renderTextChunkContent( chunk: ParsedChunkView, onReferenceClick: (chunkId: string) => void, diff --git a/src/components/workspace-citation-focus.test.ts b/src/components/workspace-citation-focus.test.ts index bfc4654..21fce5a 100644 --- a/src/components/workspace-citation-focus.test.ts +++ b/src/components/workspace-citation-focus.test.ts @@ -169,8 +169,22 @@ describe("useWorkspaceCitationFocus", () => { expect(result.current.pendingCitationId).toBeNull(); }); - it("focuses page-asset citations without fetching full chunks", async () => { - const fetchChunks = vi.fn(async () => [prefetchedChunk]); + it("focuses page-asset citations through the loaded page chunk", async () => { + const pageChunk: ParsedChunkView = { + chunkId: "page_4", + documentId: "document_1", + type: "page", + content: "Page 4 summary", + sourceTitle: "Contract.pdf", + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004.png", + contentType: "image/png", + }, + ], + }; + const fetchChunks = vi.fn(async () => [pageChunk]); const selectSource = vi.fn(); const pageAssetSource: SourceView = { ...readySource, @@ -202,14 +216,14 @@ describe("useWorkspaceCitationFocus", () => { await result.current.handleCitationClick(pageCitation, "message_1:0"); }); - expect(fetchChunks).not.toHaveBeenCalled(); + expect(fetchChunks).toHaveBeenCalledWith("source_1"); expect(selectSource).toHaveBeenCalledWith("source_1"); - expect(result.current.focusedChunk.chunkId).toBeNull(); + expect(result.current.focusedChunk.chunkId).toBe("page_4"); expect(result.current.focusedPage).toEqual({ - pageNumber: 4, - requestId: 1, + pageNumber: null, + requestId: 0, }); - expect(result.current.citationListViewRequestId).toBe(0); + expect(result.current.citationListViewRequestId).toBe(1); }); it("reuses cached chunks for a different source without refetching", async () => { diff --git a/src/components/workspace-citation-focus.ts b/src/components/workspace-citation-focus.ts index c233671..5469cfa 100644 --- a/src/components/workspace-citation-focus.ts +++ b/src/components/workspace-citation-focus.ts @@ -211,17 +211,6 @@ export function useWorkspaceCitationFocus({ ) if (!source) return - if ( - workspaceCitationState.isPageAssetCitationTarget(source, citation) - ) { - if (selectedSourceId !== source.id) onSelectSource(source.id) - requestChunkFocus(null) - requestPageFocus( - workspaceCitationState.getCitationPageNumber(citation), - ) - return - } - setCitationListViewRequestId((current) => current + 1) const loadedChunkId = workspaceCitationState.getLoadedCitationChunkId({ @@ -291,7 +280,6 @@ export function useWorkspaceCitationFocus({ loadAllChunksForSource, onSelectSource, requestChunkFocus, - requestPageFocus, selectedChunks, selectedSourceId, sources, diff --git a/src/components/workspace-citation-state.test.ts b/src/components/workspace-citation-state.test.ts index ce5851e..cce0e6b 100644 --- a/src/components/workspace-citation-state.test.ts +++ b/src/components/workspace-citation-state.test.ts @@ -91,6 +91,44 @@ describe("workspaceCitationState", () => { ).toBeNull() }) + it("focuses loaded page chunks from page citation metadata", () => { + const citation: ChatCitationView = { + chunkType: "page", + score: 0.93, + pageCitationPageNumber: 4, + pageCitationAssetUrl: "https://assets.example/page-4.png", + source: { + documentId: "document_1", + sectionPath: "Page 4", + }, + } + + expect( + workspaceCitationState.getLoadedCitationChunkId({ + citation, + selectedSourceId: "source_1", + sourceId: "source_1", + selectedChunks: [ + { + chunkId: "page_4", + documentId: "document_1", + type: "page", + content: "Page 4 summary", + sourceTitle: "Contract.pdf", + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4.png", + contentType: "image/png", + }, + ], + }, + ], + hasMoreSelectedChunks: true, + }), + ).toBe("page_4") + }) + describe("hasExactCitationTargetHint", () => { it("returns true when the citation has content text", () => { const citation: ChatCitationView = { @@ -120,6 +158,22 @@ describe("workspaceCitationState", () => { ).toBe(true) }) + it("returns true when the citation has page asset metadata", () => { + const citation: ChatCitationView = { + chunkType: "page", + score: 0.5, + pageCitationPageNumber: 4, + source: { + documentId: "document_1", + sectionPath: "Root", + }, + } + + expect( + workspaceCitationState.hasExactCitationTargetHint(citation), + ).toBe(true) + }) + it("returns false for source-only citations with no useful target hint", () => { const citation: ChatCitationView = { chunkType: "text", diff --git a/src/components/workspace-citation-state.ts b/src/components/workspace-citation-state.ts index cf33dad..0353409 100644 --- a/src/components/workspace-citation-state.ts +++ b/src/components/workspace-citation-state.ts @@ -27,10 +27,6 @@ type WorkspaceCitationStateModule = { readonly hasExactCitationTargetHint: ( citation: ChatCitationView, ) => boolean - readonly isPageAssetCitationTarget: ( - source: SourceView, - citation: ChatCitationView, - ) => boolean readonly getCitationPageNumber: (citation: ChatCitationView) => number | null readonly upsertPrefetchedChunks: ( current: PrefetchedChunksBySourceId, @@ -66,7 +62,10 @@ function getLoadedCitationChunkId( ? resolveCitationChunkByContent(input.citation, input.selectedChunks) : resolveCitationChunk(input.citation, input.selectedChunks) - return focusedChunk?.chunkId ?? null + return ( + focusedChunk?.chunkId ?? + resolveCitationPageChunkId(input.citation, input.selectedChunks) + ) } function hasExactCitationTargetHint(citation: ChatCitationView): boolean { @@ -74,6 +73,13 @@ function hasExactCitationTargetHint(citation: ChatCitationView): boolean { return true } + if ( + getCitationPageNumber(citation) !== null || + typeof citation.pageCitationAssetUrl === "string" + ) { + return true + } + const sectionPath = citation.source.sectionPath if (typeof sectionPath !== "string") return false @@ -84,17 +90,6 @@ function hasExactCitationTargetHint(citation: ChatCitationView): boolean { return true } -function isPageAssetCitationTarget( - source: SourceView, - citation: ChatCitationView, -): boolean { - return ( - source.documentPresentation?.kind === "page-assets" && - (getCitationPageNumber(citation) !== null || - typeof citation.pageCitationAssetUrl === "string") - ) -} - function getCitationPageNumber(citation: ChatCitationView): number | null { if ( typeof citation.pageCitationPageNumber === "number" && @@ -114,6 +109,43 @@ function getCitationPageNumber(citation: ChatCitationView): number | null { return Number.isSafeInteger(pageNumber) && pageNumber > 0 ? pageNumber : null } +function resolveCitationPageChunkId( + citation: ChatCitationView, + chunks: readonly ParsedChunkView[], +): string | null { + const pageNumber = getCitationPageNumber(citation) + if (pageNumber === null) return null + + const documentChunks = citation.source.documentId + ? chunks.filter((chunk) => chunk.documentId === citation.source.documentId) + : chunks + const matches = documentChunks.filter((chunk) => + isChunkForPageNumber(chunk, pageNumber), + ) + if (matches.length === 0) return null + + const pageAssetChunk = matches.find( + (chunk) => + chunk.type === "page" && + (chunk.pageAssets ?? []).some( + (pageAsset) => pageAsset.pageNumber === pageNumber, + ), + ) + + return pageAssetChunk?.chunkId ?? matches[0]?.chunkId ?? null +} + +function isChunkForPageNumber( + chunk: ParsedChunkView, + pageNumber: number, +): boolean { + return ( + (chunk.pageAssets ?? []).some( + (pageAsset) => pageAsset.pageNumber === pageNumber, + ) || (chunk.pageNums ?? []).includes(pageNumber) + ) +} + function upsertPrefetchedChunks( current: PrefetchedChunksBySourceId, sourceId: string, @@ -157,7 +189,6 @@ export const workspaceCitationState: WorkspaceCitationStateModule = { getLoadedCitationChunkId, getCitationPageNumber, hasExactCitationTargetHint, - isPageAssetCitationTarget, upsertPrefetchedChunks, removePrefetchedChunks, } diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index be575fb..f1edec1 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -10,12 +10,10 @@ import type { ParsedChunkView } from "@/domains/chunks/types"; import type { SourceView } from "@/domains/sources/types"; const fetchChunkPageMock = vi.hoisted(() => vi.fn()); -const fetchPageAssetPageMock = vi.hoisted(() => vi.fn()); vi.mock("@/domains/workspace/client", () => ({ workspaceClient: { fetchChunkPage: fetchChunkPageMock, - fetchPageAssetPage: fetchPageAssetPageMock, }, })); @@ -30,7 +28,6 @@ const readySource: SourceView = { describe("useWorkspaceSelectedChunks", () => { beforeEach(() => { fetchChunkPageMock.mockReset(); - fetchPageAssetPageMock.mockReset(); fetchChunkPageMock.mockResolvedValue({ chunks: [], pagination: { @@ -40,18 +37,9 @@ describe("useWorkspaceSelectedChunks", () => { totalPages: 1, }, }); - fetchPageAssetPageMock.mockResolvedValue({ - pages: [], - pagination: { - page: 1, - pageSize: 50, - total: 0, - totalPages: 0, - }, - }); }); - it("returns prefetched chunks while checking the visible page for media", async () => { + it("returns prefetched chunks while loading the visible chunk page", async () => { const { result } = renderHook( () => useWorkspaceSelectedChunks({ @@ -71,9 +59,6 @@ describe("useWorkspaceSelectedChunks", () => { { wrapper: createSWRWrapper }, ); - await waitFor(() => - expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 1), - ); await waitFor(() => expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1), ); @@ -209,11 +194,34 @@ describe("useWorkspaceSelectedChunks", () => { expect(result.current.selectedChunks).toEqual([]); }); - it("does not fetch chunk pages for page-asset sources", () => { + it("loads chunk pages for page-asset sources", async () => { const pageAssetSource: SourceView = { ...readySource, documentPresentation: { kind: "page-assets", pageCount: 4 }, }; + fetchChunkPageMock.mockResolvedValue({ + chunks: [ + { + chunkId: "page_1", + type: "page", + content: "Page summary", + sourceTitle: "lecture.pdf", + pageAssets: [ + { + pageNumber: 1, + assetUrl: "https://blob.example/page-1.png", + contentType: "image/png", + }, + ], + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 1, + totalPages: 1, + }, + }); const { result } = renderHook( () => @@ -225,25 +233,35 @@ describe("useWorkspaceSelectedChunks", () => { { wrapper: createSWRWrapper }, ); - expect(result.current.selectedSource?.id).toBe("source_1"); - expect(result.current.selectedChunks).toEqual([]); + await waitFor(() => + expect(result.current.selectedChunks.map((chunk) => chunk.chunkId)).toEqual([ + "page_1", + ]), + ); expect(result.current.isSelectedChunksLoading).toBe(false); - expect(fetchPageAssetPageMock).not.toHaveBeenCalled(); - expect(fetchChunkPageMock).not.toHaveBeenCalled(); + expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1); }); - it("detects page assets for selected sources without presentation metadata", async () => { - fetchPageAssetPageMock.mockResolvedValue({ - pages: [ - { - pageNumber: 1, - assetUrl: "https://blob.example/page-1.png", - contentType: "image/png", - }, + it("detects page assets from selected source chunks without a second request", async () => { + fetchChunkPageMock.mockResolvedValue({ + chunks: [ { - pageNumber: 20, - assetUrl: "https://blob.example/page-20.png", - contentType: "image/png", + chunkId: "page_1", + type: "page", + content: "Page summary", + sourceTitle: "lecture.pdf", + pageAssets: [ + { + pageNumber: 1, + assetUrl: "https://blob.example/page-1.png", + contentType: "image/png", + }, + { + pageNumber: 20, + assetUrl: "https://blob.example/page-20.png", + contentType: "image/png", + }, + ], }, ], pagination: { @@ -270,9 +288,10 @@ describe("useWorkspaceSelectedChunks", () => { pageCount: 20, }), ); - expect(result.current.selectedChunks).toEqual([]); - expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 1); - expect(fetchChunkPageMock).not.toHaveBeenCalled(); + expect(result.current.selectedChunks.map((chunk) => chunk.chunkId)).toEqual([ + "page_1", + ]); + expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1); }); it("requests a source refresh after loading an unlocalized remote source", async () => { diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index acbcd50..45b19a3 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -1,6 +1,6 @@ "use client" -import { useEffect, useMemo, useRef, useState } from "react" +import { useEffect, useMemo, useRef } from "react" import useSWRInfinite from "swr/infinite" import { workspaceClient } from "@/domains/workspace/client" @@ -8,7 +8,6 @@ import { workspaceClientCache, type SourceChunksKey, type SourceChunksResponse, - type SourcePageAssetsResponse, } from "@/domains/workspace/client-cache" import { resolveChunkConnectionTargets } from "@/domains/chunks" import type { ParsedChunkView } from "@/domains/chunks/types" @@ -31,10 +30,6 @@ type WorkspaceSelectedChunks = { readonly selectedSource: SourceView | undefined } -type PageAssetProbeState = - | { readonly status: "page-assets"; readonly pageCount: number } - | { readonly status: "parsed-chunks" } - export function useWorkspaceSelectedChunks({ selectedSourceId, sources, @@ -45,37 +40,12 @@ export function useWorkspaceSelectedChunks({ (source) => source.id === selectedSourceId, ) const remoteSourceRefreshRequestedIdsRef = useRef>(new Set()) - const requestedPageAssetProbeIdsRef = useRef>(new Set()) - const [pageAssetProbeBySourceId, setPageAssetProbeBySourceId] = useState< - Readonly> - >({}) - const pageAssetProbeState = rawSelectedSource - ? pageAssetProbeBySourceId[rawSelectedSource.id] - : undefined - const selectedSource = - rawSelectedSource && pageAssetProbeState?.status === "page-assets" - ? { - ...rawSelectedSource, - chunkCount: pageAssetProbeState.pageCount, - documentPresentation: { - kind: "page-assets" as const, - pageCount: pageAssetProbeState.pageCount, - }, - } - : rawSelectedSource + const selectedSource = rawSelectedSource const prefetchedSelectedChunks = selectedSourceId ? prefetchedChunksBySourceId[selectedSourceId] : undefined - const shouldProbePageAssets = - rawSelectedSource !== undefined && - rawSelectedSource.status === "ready" && - rawSelectedSource.documentPresentation === undefined && - pageAssetProbeState === undefined const selectedChunkSourceId = - selectedSource && - selectedSource.status === "ready" && - selectedSource.documentPresentation?.kind !== "page-assets" && - !shouldProbePageAssets + selectedSource && selectedSource.status === "ready" ? selectedSource.id : null const { @@ -118,9 +88,17 @@ export function useWorkspaceSelectedChunks({ : undefined, [pagedSelectedChunks, prefetchedSelectedChunks], ) - const selectedChunks = selectedSourceId - ? (resolvedPrefetchedChunks ?? pagedSelectedChunks) - : [] + const selectedChunks = useMemo( + () => + selectedSourceId + ? (resolvedPrefetchedChunks ?? pagedSelectedChunks) + : [], + [pagedSelectedChunks, resolvedPrefetchedChunks, selectedSourceId], + ) + const resolvedSelectedSource = useMemo( + () => getResolvedSelectedSource(selectedSource, selectedChunks), + [selectedChunks, selectedSource], + ) const hasMoreSelectedChunks = !prefetchedSelectedChunks && workspaceClientCache.hasMoreChunkPages(selectedChunkPages) @@ -132,40 +110,12 @@ export function useWorkspaceSelectedChunks({ typeof selectedChunkPages[selectedChunkPageCount - 1] === "undefined", ) const isSelectedChunksLoading = - shouldProbePageAssets || hasProcessingSelectedChunkPage || (selectedChunkSourceId !== null && !prefetchedSelectedChunks && !selectedChunkPages && isChunksLoading) - useEffect(() => { - const source = rawSelectedSource - if (!source || !shouldProbePageAssets) return - if (requestedPageAssetProbeIdsRef.current.has(source.id)) return - - requestedPageAssetProbeIdsRef.current.add(source.id) - void workspaceClient - .fetchPageAssetPage(source.id, 1) - .then((response) => { - setPageAssetProbeBySourceId((current) => ({ - ...current, - [source.id]: getPageAssetProbeState(response), - })) - - if (source.kind === "remote" && (response.pages?.length ?? 0) > 0) { - onRemoteSourceChunksLoaded?.(source.id) - } - }) - .catch(() => { - requestedPageAssetProbeIdsRef.current.delete(source.id) - setPageAssetProbeBySourceId((current) => ({ - ...current, - [source.id]: { status: "parsed-chunks" }, - })) - }) - }, [onRemoteSourceChunksLoaded, rawSelectedSource, shouldProbePageAssets]) - useEffect(() => { const sourceId = selectedSource?.id if (!sourceId || selectedSource.kind !== "remote") return @@ -188,23 +138,10 @@ export function useWorkspaceSelectedChunks({ isSelectedChunksLoadingMore, selectedChunksMessage, selectedChunks, - selectedSource, + selectedSource: resolvedSelectedSource, } } -function getPageAssetProbeState( - response: SourcePageAssetsResponse, -): PageAssetProbeState { - const pages = response.pages ?? [] - if (pages.length === 0) return { status: "parsed-chunks" } - - const maxPageNumber = Math.max( - ...pages.map((page) => page.pageNumber), - ) - const pageCount = Math.max(response.pagination?.total ?? 0, maxPageNumber) - return { status: "page-assets", pageCount } -} - function fetchChunksByKey([ , sourceId, @@ -213,6 +150,36 @@ function fetchChunksByKey([ return workspaceClient.fetchChunkPage(sourceId, page) } +function getResolvedSelectedSource( + source: SourceView | undefined, + chunks: readonly ParsedChunkView[], +): SourceView | undefined { + if (!source) return undefined + if (source.documentPresentation?.kind === "page-assets") return source + + const pageCount = getLoadedPageAssetCount(chunks) + if (!pageCount) return source + + return { + ...source, + documentPresentation: { + kind: "page-assets", + pageCount, + }, + } +} + +function getLoadedPageAssetCount( + chunks: readonly ParsedChunkView[], +): number | null { + const pageNumbers = chunks.flatMap((chunk) => + (chunk.pageAssets ?? []).map((pageAsset) => pageAsset.pageNumber), + ) + if (pageNumbers.length === 0) return null + + return Math.max(...pageNumbers) +} + function hasProcessingChunkPage( pages: readonly SourceChunksResponse[] | undefined, ): boolean { diff --git a/src/domains/chunks/index.test.ts b/src/domains/chunks/index.test.ts index 316bfd6..586d32b 100644 --- a/src/domains/chunks/index.test.ts +++ b/src/domains/chunks/index.test.ts @@ -169,6 +169,42 @@ describe("toParsedChunkView", () => { }); }); + it("maps usable page citation assets on page chunks", () => { + const chunk = makeDocumentChunk({ + id: "document_page_1", + chunkId: "parser_page_1", + chunkType: "page" as DocumentChunk["chunkType"], + metadata: { + pageAssets: [ + { + pageNum: 4, + assetUrl: "https://assets.example/page-4.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + { + pageNum: 5, + assetUrl: " ", + contentType: "image/png", + }, + ], + }, + }); + + expect(toParsedChunkView(chunk, "manual.pdf", "doc_123")).toMatchObject({ + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }); + }); + it("maps SDK-normalized page number metadata", () => { const chunk = makeDocumentChunk({ metadata: { diff --git a/src/domains/chunks/normalization.ts b/src/domains/chunks/normalization.ts index c6590b4..3ab6bae 100644 --- a/src/domains/chunks/normalization.ts +++ b/src/domains/chunks/normalization.ts @@ -35,6 +35,10 @@ function createParsedChunkView( metadata: input.metadata, }) const summary = getStringMetadata(input.metadata, "summary") + const pageAssets = + type === "page" + ? getPageAssetViews(input.metadata["pageAssets"], assetUrl) + : undefined return { chunkId: input.chunkId, @@ -47,6 +51,7 @@ function createParsedChunkView( readableContent: getReadableContent({ type, content, summary }), filePath, assetUrl, + pageAssets, summary, keywords: getStringArrayMetadata(input.metadata, "keywords"), pageNums: getPageNumbers( @@ -58,6 +63,38 @@ function createParsedChunkView( } } +function getPageAssetViews( + value: unknown, + fallbackAssetUrl?: string, +): ParsedChunkView["pageAssets"] | undefined { + if (!Array.isArray(value)) return undefined + + const pageAssets = value.flatMap((item): NonNullable => { + if (!isRecord(item)) return [] + const pageNumber = + getPositiveInteger(item["pageNum"]) ?? getPositiveInteger(item["pageNumber"]) + const assetUrl = getString(item["assetUrl"]) ?? fallbackAssetUrl + const contentType = getString(item["contentType"]) + if (!pageNumber || !assetUrl || !contentType) return [] + + return [ + { + pageNumber, + assetUrl, + contentType, + ...(getPositiveNumber(item["width"]) !== undefined + ? { width: getPositiveNumber(item["width"]) } + : {}), + ...(getPositiveNumber(item["height"]) !== undefined + ? { height: getPositiveNumber(item["height"]) } + : {}), + }, + ] + }) + + return pageAssets.length > 0 ? pageAssets : undefined +} + function resolveConnectionTargets( chunks: readonly ParsedChunkView[], ): ParsedChunkView[] { @@ -281,6 +318,20 @@ function getPageNumbers(value: unknown): number[] | undefined { return uniquePageNumbers.length > 0 ? uniquePageNumbers : undefined } +function getPositiveInteger(value: unknown): number | undefined { + return typeof value === "number" && + Number.isSafeInteger(value) && + value > 0 + ? value + : undefined +} + +function getPositiveNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value > 0 + ? value + : undefined +} + function isRecord(value: unknown): value is Readonly> { return typeof value === "object" && value !== null } diff --git a/src/domains/chunks/types.ts b/src/domains/chunks/types.ts index 2e84189..fedb17a 100644 --- a/src/domains/chunks/types.ts +++ b/src/domains/chunks/types.ts @@ -1,3 +1,5 @@ +import type { SourcePageAssetView } from "@/domains/sources/types" + export type ChunkType = "text" | "image" | "table" | "page" export type ParsedChunkConnection = { @@ -30,6 +32,7 @@ export type ParsedChunkView = { readonly filePath?: string /** Public Blob URL for parsed media/table artifacts when Notebook stored it. */ readonly assetUrl?: string + readonly pageAssets?: readonly SourcePageAssetView[] readonly summary?: string readonly keywords?: readonly string[] readonly pageNums?: readonly number[] diff --git a/src/domains/sources/page-assets.ts b/src/domains/sources/page-assets.ts deleted file mode 100644 index c23a702..0000000 --- a/src/domains/sources/page-assets.ts +++ /dev/null @@ -1,218 +0,0 @@ -import "server-only" - -import type { - DocumentChunk, - Knowledge, - KnowledgeReadResponse, -} from "@ontos-ai/knowhere-sdk" - -import type { ChunkPageParams } from "@/domains/chunks" -import type { SourcePageAssetView } from "./route-types" - -type ReadableSource = { - readonly documentId: string - readonly revisionKey?: string | null -} - -export type SourcePageAssetsPage = { - readonly pages: readonly SourcePageAssetView[] - readonly pagination: { - readonly page: number - readonly pageSize: number - readonly total: number - readonly totalPages: number - } -} - -type PageAssetReadClient = { - readonly documents: { - listChunks( - documentId: string, - params: { - readonly page: number - readonly pageSize: number - readonly chunkType: "page" - readonly includeAssetUrls: true - }, - ): Promise<{ - readonly chunks: readonly DocumentChunk[] - readonly pagination?: { - readonly page?: number - readonly pageSize?: number - readonly total?: number - readonly totalPages?: number - } - }> - } -} - -export async function readSourcePageAssets(input: { - readonly client: PageAssetReadClient - readonly knowledge: Knowledge - readonly source: ReadableSource - readonly params: ChunkPageParams -}): Promise { - const response = await input.knowledge.readChunks({ - documentId: input.source.documentId, - ...(input.source.revisionKey ? { revisionKey: input.source.revisionKey } : {}), - chunkType: "page", - page: input.params.page, - pageSize: input.params.pageSize, - }) - - const knowledgePage = toSourcePageAssetsPageFromKnowledgeResponse( - response, - input.params, - ) - if ( - isParsedStorageReadResponse(response) && - knowledgePage.pages.length > 0 - ) { - return knowledgePage - } - if (knowledgePage.pages.length > 0) return knowledgePage - - const remoteResponse = await input.client.documents.listChunks( - input.source.documentId, - { - page: input.params.page, - pageSize: input.params.pageSize, - chunkType: "page", - includeAssetUrls: true, - }, - ) - return toSourcePageAssetsPageFromRemoteResponse(remoteResponse, input.params) -} - -function toSourcePageAssetsPageFromKnowledgeResponse( - response: KnowledgeReadResponse, - params: ChunkPageParams, -): SourcePageAssetsPage { - const pages = response.chunks.flatMap((chunk): SourcePageAssetView[] => - readPageAssetViews(chunk.metadata.pageAssets, chunk.assetUrl), - ) - const maxPageNumber = getMaxPageNumber(pages) - const total = Math.max(response.totalChunks ?? 0, maxPageNumber ?? 0) - const resolvedTotal = total > 0 ? total : pages.length - const computedTotalPages = Math.max( - 1, - Math.ceil(resolvedTotal / params.pageSize), - ) - const totalPages = Math.max(response.totalPages ?? 0, computedTotalPages) - - return { - pages, - pagination: { - page: response.page ?? params.page, - pageSize: response.pageSize ?? params.pageSize, - total: resolvedTotal, - totalPages, - }, - } -} - -function toSourcePageAssetsPageFromRemoteResponse( - response: { - readonly chunks: readonly DocumentChunk[] - readonly pagination?: { - readonly page?: number - readonly pageSize?: number - readonly total?: number - readonly totalPages?: number - } - }, - params: ChunkPageParams, -): SourcePageAssetsPage { - const pages = response.chunks.flatMap((chunk): SourcePageAssetView[] => - readPageAssetViews(chunk.metadata.pageAssets, chunk.assetUrl ?? undefined), - ) - const maxPageNumber = getMaxPageNumber(pages) - const total = Math.max(response.pagination?.total ?? 0, maxPageNumber ?? 0) - const resolvedTotal = total > 0 ? total : pages.length - const computedTotalPages = Math.max( - 1, - Math.ceil(resolvedTotal / params.pageSize), - ) - const totalPages = Math.max( - response.pagination?.totalPages ?? 0, - computedTotalPages, - ) - - return { - pages, - pagination: { - page: response.pagination?.page ?? params.page, - pageSize: response.pagination?.pageSize ?? params.pageSize, - total: resolvedTotal, - totalPages, - }, - } -} - -function isParsedStorageReadResponse(response: KnowledgeReadResponse): boolean { - const resultDirectoryPath = response.document.resultDirectoryPath - return ( - typeof resultDirectoryPath === "string" && - resultDirectoryPath.startsWith("parsed-storage:") - ) -} - -function readPageAssetViews( - value: unknown, - fallbackAssetUrl?: string, -): SourcePageAssetView[] { - if (!Array.isArray(value)) return [] - - return value.flatMap((item): SourcePageAssetView[] => { - if (!isRecord(item)) return [] - const pageNumber = getPositiveInteger(item.pageNum) - const assetUrl = getTrimmedString(item.assetUrl) ?? fallbackAssetUrl - const contentType = getTrimmedString(item.contentType) - if (!pageNumber || !assetUrl || !contentType) return [] - - return [ - { - pageNumber, - assetUrl, - contentType, - ...(getPositiveNumber(item.width) !== undefined - ? { width: getPositiveNumber(item.width) } - : {}), - ...(getPositiveNumber(item.height) !== undefined - ? { height: getPositiveNumber(item.height) } - : {}), - }, - ] - }) -} - -function getTrimmedString(value: unknown): string | undefined { - if (typeof value !== "string") return undefined - const trimmed = value.trim() - return trimmed.length > 0 ? trimmed : undefined -} - -function getPositiveInteger(value: unknown): number | undefined { - return typeof value === "number" && - Number.isSafeInteger(value) && - value > 0 - ? value - : undefined -} - -function getPositiveNumber(value: unknown): number | undefined { - return typeof value === "number" && Number.isFinite(value) && value > 0 - ? value - : undefined -} - -function getMaxPageNumber( - pages: readonly SourcePageAssetView[], -): number | undefined { - if (pages.length === 0) return undefined - return Math.max(...pages.map((page) => page.pageNumber)) -} - -function isRecord(value: unknown): value is Readonly> { - return typeof value === "object" && value !== null -} diff --git a/src/domains/sources/route-page-assets.ts b/src/domains/sources/route-page-assets.ts deleted file mode 100644 index a2316fd..0000000 --- a/src/domains/sources/route-page-assets.ts +++ /dev/null @@ -1,208 +0,0 @@ -import { Effect } from "effect" - -import { readSourcePageAssets } from "./page-assets" -import { routeResult } from "@/lib/route-result" -import { displayReadUnavailable } from "./display-read-unavailable" -import { - decodeRemoteSourceId, - findRemoteLibraryDocumentBySourceId, -} from "./remote-library" -import { - getClientForWorkspace, - getKnowledgeResourcesForSource, -} from "./route-dependencies" -import { sourceRowRepository } from "./source-row-repository" -import type { - JsonRouteResult, - LoadSourcePageAssetsInput, - SourcePageAssetsBody, - SourceRouteServiceDependencies, -} from "./route-types" - -type RoutePageAssetsDependencies = Pick< - SourceRouteServiceDependencies, - | "ensureApiKeyForWorkspace" - | "ensureWorkspace" - | "getCurrentUser" - | "makeKnowhereClient" - | "sourceService" -> - -type RoutePageAssets = { - readonly loadSourcePageAssets: ( - input: LoadSourcePageAssetsInput, - ) => Promise> -} - -function createRoutePageAssets( - deps: RoutePageAssetsDependencies, -): RoutePageAssets { - return { - loadSourcePageAssets: (input: LoadSourcePageAssetsInput) => - Effect.runPromise(loadSourcePageAssetsEffect(input, deps)), - } -} - -const loadSourcePageAssetsEffect = ( - input: LoadSourcePageAssetsInput, - deps: RoutePageAssetsDependencies, -) => - Effect.gen(function* () { - if (!sourceRowRepository.isWorkspaceSourceId(input.sourceId)) { - const remoteResult = yield* loadRemotePageAssetsEffect(input, deps) - return remoteResult ?? sourceNotFound() - } - - const user = yield* Effect.tryPromise(() => deps.getCurrentUser()) - if (!user) return sourceNotFound() - - const workspace = yield* Effect.tryPromise(() => - deps.ensureWorkspace(user.id), - ) - const source = yield* Effect.tryPromise(() => - deps.sourceService.findInWorkspace(workspace.id, input.sourceId), - ) - if (!source) return sourceNotFound() - if (source.status !== "ready" || !source.knowhereDocumentId) { - return sourceNotReady() - } - - const documentId = source.knowhereDocumentId - const apiKey = yield* Effect.tryPromise(() => - deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), - ) - const readResources = getKnowledgeResourcesForSource({ - apiKey, - workspaceId: workspace.id, - sourceId: source.id, - documentId, - revisionKey: source.knowhereJobId, - }) - return yield* Effect.tryPromise(() => - readSourcePageAssets({ - client: readResources.client, - knowledge: readResources.knowledge, - source: { - documentId, - revisionKey: source.knowhereJobId, - }, - params: input.pageParams, - }), - ).pipe( - Effect.map((pageAssets) => routeResult.ok(pageAssets)), - Effect.catchAll((error) => - recoverUnavailablePageAssets(input, error), - ), - ) - }) - -const loadRemotePageAssetsEffect = ( - input: LoadSourcePageAssetsInput, - deps: RoutePageAssetsDependencies, -) => - Effect.gen(function* () { - if (!decodeRemoteSourceId(input.sourceId)) return null - - const user = yield* Effect.tryPromise(() => deps.getCurrentUser()) - if (!user) return null - - const workspace = yield* Effect.tryPromise(() => - deps.ensureWorkspace(user.id), - ) - const apiKey = yield* Effect.tryPromise(() => - deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), - ) - const client = yield* Effect.tryPromise(() => - getClientForWorkspace(workspace.id, input.cookieHeader, deps), - ) - const remoteDocument = yield* findRemoteLibraryDocumentBySourceId({ - sourceId: input.sourceId, - workspace, - client, - localSources: [], - }) - if (!remoteDocument) return null - - const source = yield* Effect.tryPromise(() => - deps.sourceService.localizeRemoteDocument(workspace.id, { - documentId: remoteDocument.documentId, - namespace: remoteDocument.namespace, - status: remoteDocument.status, - title: remoteDocument.title, - mimeType: remoteDocument.mimeType, - sizeBytes: remoteDocument.sizeBytes, - revisionKey: remoteDocument.revisionKey ?? null, - }), - ) - const documentId = source.knowhereDocumentId ?? remoteDocument.documentId - const revisionKey = - source.knowhereJobId ?? remoteDocument.revisionKey ?? null - const readResources = getKnowledgeResourcesForSource({ - apiKey, - workspaceId: workspace.id, - sourceId: source.id, - documentId, - revisionKey, - }) - return yield* Effect.tryPromise(() => - readSourcePageAssets({ - client: readResources.client, - knowledge: readResources.knowledge, - source: { documentId, revisionKey }, - params: input.pageParams, - }), - ).pipe( - Effect.map((pageAssets) => routeResult.ok(pageAssets)), - Effect.catchAll((error) => - recoverUnavailablePageAssets(input, error), - ), - ) - }) - -function sourceNotFound(): JsonRouteResult<{ readonly message: string }> { - return routeResult.error(404, "Source not found.") -} - -function sourceNotReady(): JsonRouteResult<{ readonly message: string }> { - return routeResult.error(409, "Source is not ready.") -} - -function sourcePageAssetsUnavailable( - input: LoadSourcePageAssetsInput, -): { - readonly pages: [] - readonly pagination: { - readonly page: number - readonly pageSize: number - readonly total: 0 - readonly totalPages: 0 - } - readonly message: string - readonly isUnavailable: true -} { - return { - pages: [], - pagination: { - page: input.pageParams.page, - pageSize: input.pageParams.pageSize, - total: 0, - totalPages: 0, - }, - message: displayReadUnavailable.message, - isUnavailable: true, - } -} - -function recoverUnavailablePageAssets( - input: LoadSourcePageAssetsInput, - error: unknown, -): Effect.Effect< - JsonRouteResult>, - unknown -> { - return displayReadUnavailable.isError(error) - ? Effect.succeed(routeResult.ok(sourcePageAssetsUnavailable(input))) - : Effect.fail(error) -} - -export { createRoutePageAssets } diff --git a/src/domains/sources/route-service.ts b/src/domains/sources/route-service.ts index 7e3b85f..9a19bef 100644 --- a/src/domains/sources/route-service.ts +++ b/src/domains/sources/route-service.ts @@ -4,14 +4,12 @@ import { createRouteArchive } from "./route-archive" import { createRouteChunks } from "./route-chunks" import { createSourceRouteDependencies } from "./route-dependencies" import { createRouteListing } from "./route-listing" -import { createRoutePageAssets } from "./route-page-assets" import { createRouteRetry } from "./route-retry" import { createRouteUpload } from "./route-upload" import type { ArchiveSourceInput, ListSourcesInput, LoadSourceChunksInput, - LoadSourcePageAssetsInput, RetrySourceInput, SourceRouteService, SourceRouteServiceOverrides, @@ -27,7 +25,6 @@ export function createSourceRouteService( const archive = createRouteArchive(deps) const retry = createRouteRetry(deps) const chunks = createRouteChunks(deps) - const pageAssets = createRoutePageAssets(deps) return { listSources: (input: ListSourcesInput) => listing.listSources(input), @@ -36,7 +33,5 @@ export function createSourceRouteService( retrySource: (input: RetrySourceInput) => retry.retrySource(input), loadSourceChunks: (input: LoadSourceChunksInput) => chunks.loadSourceChunks(input), - loadSourcePageAssets: (input: LoadSourcePageAssetsInput) => - pageAssets.loadSourcePageAssets(input), } } diff --git a/src/domains/sources/route-types.ts b/src/domains/sources/route-types.ts index 8ce806d..01d1883 100644 --- a/src/domains/sources/route-types.ts +++ b/src/domains/sources/route-types.ts @@ -5,7 +5,6 @@ import type { } from "@/domains/chunks" import type { ParsedChunkView } from "@/domains/chunks/types" import type { - SourcePageAssetView, SourceStatus, SourceView, } from "@/domains/sources/types" @@ -124,31 +123,6 @@ type SourceChunksBody = readonly message: string } -type SourcePageAssetsBody = - | { - readonly pages: readonly SourcePageAssetView[] - readonly pagination: { - readonly page: number - readonly pageSize: number - readonly total: number - readonly totalPages: number - } - } - | { - readonly pages: readonly [] - readonly pagination: { - readonly page: number - readonly pageSize: number - readonly total: 0 - readonly totalPages: 0 - } - readonly message: string - readonly isUnavailable: true - } - | { - readonly message: string - } - type ListSourcesInput = { readonly cookieHeader: string } @@ -176,12 +150,6 @@ type LoadSourceChunksInput = { readonly pageParams: ChunkPageParams } -type LoadSourcePageAssetsInput = { - readonly cookieHeader: string - readonly sourceId: string - readonly pageParams: ChunkPageParams -} - type SourceRouteService = { readonly listSources: ( input: ListSourcesInput, @@ -198,9 +166,6 @@ type SourceRouteService = { readonly loadSourceChunks: ( input: LoadSourceChunksInput, ) => Promise> - readonly loadSourcePageAssets: ( - input: LoadSourcePageAssetsInput, - ) => Promise> } type SourceWorkflowService = { @@ -304,12 +269,9 @@ export type { ListSourcesBody, ListSourcesInput, LoadSourceChunksInput, - LoadSourcePageAssetsInput, RetrySourceBody, RetrySourceInput, SourceChunksBody, - SourcePageAssetsBody, - SourcePageAssetView, SourceRouteDemoApi, SourceRouteKnowhereClient, SourceRouteService, diff --git a/src/domains/workspace/client-cache.ts b/src/domains/workspace/client-cache.ts index beaa41d..af5baae 100644 --- a/src/domains/workspace/client-cache.ts +++ b/src/domains/workspace/client-cache.ts @@ -11,9 +11,6 @@ import type { SourceView } from "@/domains/sources/types" type SourceChunksResponse = Awaited< ReturnType > -type SourcePageAssetsResponse = Awaited< - ReturnType -> type ChatThreadDetailResponse = Awaited< ReturnType > @@ -117,5 +114,4 @@ export type { ChatThreadKey, SourceChunksKey, SourceChunksResponse, - SourcePageAssetsResponse, } diff --git a/src/domains/workspace/client.test.ts b/src/domains/workspace/client.test.ts index 0c9df2b..47878bf 100644 --- a/src/domains/workspace/client.test.ts +++ b/src/domains/workspace/client.test.ts @@ -120,36 +120,6 @@ describe("workspaceClient", () => { }) }) - it("preserves page asset unavailable messages", async () => { - mockRouteClient.getJson.mockResolvedValue({ - pages: [], - message: - "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", - isUnavailable: true, - pagination: { - page: 1, - pageSize: 50, - total: 0, - totalPages: 0, - }, - }) - - const page = await workspaceClient.fetchPageAssetPage("source_1", 1) - - expect(page).toEqual({ - pages: [], - isUnavailable: true, - message: - "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", - pagination: { - page: 1, - pageSize: 50, - total: 0, - totalPages: 0, - }, - }) - }) - it("throws materialization route errors instead of treating them as empty sources", async () => { mockRouteClient.postJsonWithStatus.mockResolvedValue({ status: 502, diff --git a/src/domains/workspace/client.ts b/src/domains/workspace/client.ts index 5490d33..a415d6a 100644 --- a/src/domains/workspace/client.ts +++ b/src/domains/workspace/client.ts @@ -5,7 +5,6 @@ import type { } from "@/domains/chat/types" import type { ParsedChunkView } from "@/domains/chunks/types" import type { - SourcePageAssetView, SourceView, } from "@/domains/sources/types" import { workspaceRouteClient } from "./route-client" @@ -38,18 +37,6 @@ type SourceChunksResponse = { } } -type SourcePageAssetsResponse = { - pages?: SourcePageAssetView[] - isUnavailable?: boolean - message?: string - pagination?: { - page: number - pageSize: number - total: number - totalPages: number - } -} - type ChatThreadResponse = { thread?: ChatThreadView messages?: ChatMessageView[] @@ -107,7 +94,6 @@ export const workspaceClient = { keys: workspaceClientKeys, fetchChunks, fetchChunkPage, - fetchPageAssetPage, fetchSources, fetchChatThreads, fetchChatThread, @@ -156,26 +142,6 @@ async function fetchChunkPage( } } -async function fetchPageAssetPage( - sourceId: string, - page: number, -): Promise { - const searchParams = new URLSearchParams({ - page: String(page), - pageSize: String(workspaceClientConfig.sourceChunkPageSize), - }) - const body = await workspaceRouteClient.getJson( - `/api/sources/${encodeURIComponent(sourceId)}/page-assets?${searchParams.toString()}`, - ) - - return { - pages: Array.isArray(body.pages) ? body.pages : [], - ...(body.isUnavailable === true ? { isUnavailable: true } : {}), - ...(typeof body.message === "string" ? { message: body.message } : {}), - pagination: body.pagination, - } -} - async function fetchSources(): Promise { const body = await workspaceRouteClient.getJson( workspaceClientKeys.sources, From c83a0f38cbe1c0efc4a9a344d0b105a091c4fc42 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 8 Jul 2026 02:49:35 +0800 Subject: [PATCH 30/51] Fix remote document refresh and page dedupe --- src/components/chunks-panel-state.test.ts | 59 +++++++++++++ src/components/chunks-panel-state.ts | 34 ++++++++ src/components/chunks-panel.test.ts | 67 +++++++++++++++ src/components/chunks-panel.tsx | 17 ++-- src/components/workspace-citation-focus.ts | 3 - .../workspace-selected-chunks.test.ts | 27 +++--- src/components/workspace-selected-chunks.ts | 15 +--- src/components/workspace-shell.test.ts | 85 +++++++++++++++++++ src/components/workspace-shell.tsx | 1 - 9 files changed, 268 insertions(+), 40 deletions(-) diff --git a/src/components/chunks-panel-state.test.ts b/src/components/chunks-panel-state.test.ts index bcee0f0..53cb15a 100644 --- a/src/components/chunks-panel-state.test.ts +++ b/src/components/chunks-panel-state.test.ts @@ -139,6 +139,65 @@ describe("chunksPanelState", () => { ]) }) + it("deduplicates page-asset chunks by page number", () => { + const chunks: ParsedChunkView[] = [ + { + chunkId: "page_4_first", + type: "page", + content: "First page 4 summary.", + sourceTitle: "manual.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4-a.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "page_4_duplicate", + type: "page", + content: "Duplicate page 4 summary.", + sourceTitle: "manual.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4-b.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "page_5", + type: "page", + content: "Page 5 summary.", + sourceTitle: "manual.pdf", + pageAssets: [ + { + pageNumber: 5, + assetUrl: "https://assets.example/page-5.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "table_1", + type: "table", + content: "", + sourceTitle: "manual.pdf", + pageNums: [4], + }, + ] + + expect( + chunksPanelState + .getPageAssetChunksWithoutDuplicatePages(chunks) + .map((chunk) => chunk.chunkId), + ).toEqual(["page_4_first", "page_5", "table_1"]) + }) + it("formats Knowhere section paths and reference labels for display", () => { expect( chunksPanelState.formatChunkSectionPath( diff --git a/src/components/chunks-panel-state.ts b/src/components/chunks-panel-state.ts index dfd2ac7..0c64515 100644 --- a/src/components/chunks-panel-state.ts +++ b/src/components/chunks-panel-state.ts @@ -42,6 +42,9 @@ type ChunksPanelStateModule = { chunks: readonly ParsedChunkView[], focusedChunkId: string | null, ) => readonly ParsedChunkView[] + readonly getPageAssetChunksWithoutDuplicatePages: ( + chunks: readonly ParsedChunkView[], + ) => readonly ParsedChunkView[] readonly getReferenceLabel: (connection: ParsedChunkConnection) => string readonly getRenderableReferences: ( chunk: ParsedChunkView, @@ -157,6 +160,32 @@ function dedupeChunksById( return uniqueChunks } +function getPageAssetChunksWithoutDuplicatePages( + chunks: readonly ParsedChunkView[], +): readonly ParsedChunkView[] { + const seenPageNumbers = new Set() + + return chunks.filter((chunk) => { + if (chunk.type !== "page") return true + + const pageNumber = getPageAssetChunkPageNumber(chunk) + if (pageNumber === null) return true + if (seenPageNumbers.has(pageNumber)) return false + + seenPageNumbers.add(pageNumber) + return true + }) +} + +function getPageAssetChunkPageNumber(chunk: ParsedChunkView): number | null { + const pageAssetNumbers = (chunk.pageAssets ?? []) + .map((pageAsset) => pageAsset.pageNumber) + .filter(isPositivePageNumber) + if (pageAssetNumbers.length > 0) return Math.min(...pageAssetNumbers) + + return getFirstPageNumber(chunk) +} + function createMutableSectionTreeNode(input: { readonly id: string readonly kind: ChunkSectionTreeNodeKind @@ -336,6 +365,10 @@ function getFirstPageNumber(chunk: ParsedChunkView): number | null { return Math.min(...finitePageNumbers) } +function isPositivePageNumber(pageNumber: number): boolean { + return Number.isFinite(pageNumber) && pageNumber > 0 +} + function formatChunkSectionPath( sectionPath: ParsedChunkView["sectionPath"], ): string | null { @@ -454,6 +487,7 @@ export const chunksPanelState: ChunksPanelStateModule = { formatChunkSectionPath, formatReferenceLabel, getChunksWithFocusedFirst, + getPageAssetChunksWithoutDuplicatePages, getReferenceLabel, getRenderableReferences, } diff --git a/src/components/chunks-panel.test.ts b/src/components/chunks-panel.test.ts index 78a35f7..3031cb8 100644 --- a/src/components/chunks-panel.test.ts +++ b/src/components/chunks-panel.test.ts @@ -127,6 +127,73 @@ describe("ChunksPanel", () => { ).toBeNull(); }); + it("renders only the first page-asset chunk for each page number", async () => { + mockVisibleVirtualViewport(); + + render( + React.createElement(C, { + chunks: [ + { + chunkId: "page_4_first", + type: "page", + content: "First page 4 summary", + sourceTitle: "report.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004-a.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "page_4_duplicate", + type: "page", + content: "Duplicate page 4 summary", + sourceTitle: "report.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004-b.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "page_5", + type: "page", + content: "Page 5 summary", + sourceTitle: "report.pdf", + pageNums: [5], + pageAssets: [ + { + pageNumber: 5, + assetUrl: "https://assets.example/page-000005.png", + contentType: "image/png", + }, + ], + }, + ], + selectedSource: "report.pdf", + selectedSourceView: { + id: "source_1", + title: "report.pdf", + mimeType: "application/pdf", + status: "ready", + documentPresentation: { kind: "page-assets", pageCount: 5 }, + }, + }), + ); + + expect(await screen.findAllByRole("img", { name: "Page 4" })) + .toHaveLength(1); + expect(screen.getByRole("img", { name: "Page 5" })).toBeTruthy(); + expect(screen.queryByTestId("chunk-card-shell-page_4_duplicate")) + .toBeNull(); + }); + it("renders page chunks normally when no page assets exist", async () => { mockVisibleVirtualViewport(); diff --git a/src/components/chunks-panel.tsx b/src/components/chunks-panel.tsx index 71037c9..f49df7e 100644 --- a/src/components/chunks-panel.tsx +++ b/src/components/chunks-panel.tsx @@ -106,6 +106,13 @@ export function ChunksPanel({ const isPageAssetSource = selectedSourceView?.documentPresentation?.kind === "page-assets" || hasPageAssetChunks; + const displayChunks = useMemo( + () => + isPageAssetSource + ? chunksPanelState.getPageAssetChunksWithoutDuplicatePages(chunks) + : chunks, + [chunks, isPageAssetSource], + ); const effectiveVisibleView = isPageAssetSource ? "parsed" : undefined; const originalPreviewCacheKey = selectedSourceFile?.url ?? null; const isOriginalPreviewAvailable = @@ -144,7 +151,7 @@ export function ChunksPanel({ visibleChunks, visibleView, } = useChunksPanelWorkflow({ - chunks, + chunks: displayChunks, selectedSource, selectedSourceFile, focusedChunkId, @@ -331,7 +338,7 @@ export function ChunksPanel({
{!isPageAssetSource && activeVisibleView === "parsed" && - chunks.length > 0 ? ( + displayChunks.length > 0 ? (
{isLoading ? ( - ) : chunks.length === 0 && processingMessage ? ( + ) : displayChunks.length === 0 && processingMessage ? ( - ) : chunks.length === 0 ? ( + ) : displayChunks.length === 0 ? ( selectedSource ? ( ) : ( @@ -404,7 +411,7 @@ export function ChunksPanel({ ) : isTreeModeVisible ? ( Promise readonly initialPrefetchedChunksBySourceId?: PrefetchedChunksBySourceId - readonly onRemoteSourceChunksLoaded?: (sourceId: string) => void readonly onSelectSource: (sourceId: string | null) => void readonly selectedSourceId: string | null readonly sources: readonly SourceView[] @@ -58,7 +57,6 @@ type WorkspaceCitationFocus = { export function useWorkspaceCitationFocus({ fetchChunks, initialPrefetchedChunksBySourceId = {}, - onRemoteSourceChunksLoaded, onSelectSource, selectedSourceId, sources, @@ -100,7 +98,6 @@ export function useWorkspaceCitationFocus({ selectedSourceId, sources, prefetchedChunksBySourceId, - onRemoteSourceChunksLoaded, }) const requestChunkFocus = useCallback( diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index f1edec1..59cf03a 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -294,8 +294,7 @@ describe("useWorkspaceSelectedChunks", () => { expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1); }); - it("requests a source refresh after loading an unlocalized remote source", async () => { - const onRemoteSourceChunksLoaded = vi.fn(); + it("loads an unlocalized remote source without requesting a source refresh", async () => { const remoteSource: SourceView = { ...readySource, id: "knowhere-doc:default:doc_remote", @@ -320,30 +319,24 @@ describe("useWorkspaceSelectedChunks", () => { }, }); - const { rerender } = renderHook( - (input: { - readonly onRemoteSourceChunksLoaded: (sourceId: string) => void; - }) => + const { result } = renderHook( + () => useWorkspaceSelectedChunks({ selectedSourceId: "knowhere-doc:default:doc_remote", sources: [remoteSource], prefetchedChunksBySourceId: {}, - onRemoteSourceChunksLoaded: input.onRemoteSourceChunksLoaded, }), - { - initialProps: { onRemoteSourceChunksLoaded }, - wrapper: createSWRWrapper, - }, + { wrapper: createSWRWrapper }, ); await waitFor(() => - expect(onRemoteSourceChunksLoaded).toHaveBeenCalledWith( - "knowhere-doc:default:doc_remote", - ), + expect(result.current.selectedChunks.map((chunk) => chunk.chunkId)).toEqual([ + "chunk_1", + ]), + ); + expect(result.current.selectedSource?.id).toBe( + "knowhere-doc:default:doc_remote", ); - - rerender({ onRemoteSourceChunksLoaded }); - expect(onRemoteSourceChunksLoaded).toHaveBeenCalledTimes(1); }); it("returns an empty chunk list when no source is selected", () => { diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index 45b19a3..260a6f5 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -1,6 +1,6 @@ "use client" -import { useEffect, useMemo, useRef } from "react" +import { useMemo } from "react" import useSWRInfinite from "swr/infinite" import { workspaceClient } from "@/domains/workspace/client" @@ -17,7 +17,6 @@ type WorkspaceSelectedChunksInput = { readonly selectedSourceId: string | null readonly sources: readonly SourceView[] readonly prefetchedChunksBySourceId: Readonly> - readonly onRemoteSourceChunksLoaded?: (sourceId: string) => void } type WorkspaceSelectedChunks = { @@ -34,12 +33,10 @@ export function useWorkspaceSelectedChunks({ selectedSourceId, sources, prefetchedChunksBySourceId, - onRemoteSourceChunksLoaded, }: WorkspaceSelectedChunksInput): WorkspaceSelectedChunks { const rawSelectedSource = sources.find( (source) => source.id === selectedSourceId, ) - const remoteSourceRefreshRequestedIdsRef = useRef>(new Set()) const selectedSource = rawSelectedSource const prefetchedSelectedChunks = selectedSourceId ? prefetchedChunksBySourceId[selectedSourceId] @@ -116,16 +113,6 @@ export function useWorkspaceSelectedChunks({ !selectedChunkPages && isChunksLoading) - useEffect(() => { - const sourceId = selectedSource?.id - if (!sourceId || selectedSource.kind !== "remote") return - if (!selectedChunkPages || selectedChunkPages.length === 0) return - if (remoteSourceRefreshRequestedIdsRef.current.has(sourceId)) return - - remoteSourceRefreshRequestedIdsRef.current.add(sourceId) - onRemoteSourceChunksLoaded?.(sourceId) - }, [onRemoteSourceChunksLoaded, selectedChunkPages, selectedSource]) - function handleLoadMoreChunks(): void { if (!hasMoreSelectedChunks || isSelectedChunksLoadingMore) return void setSelectedChunkPageCount(selectedChunkPageCount + 1) diff --git a/src/components/workspace-shell.test.ts b/src/components/workspace-shell.test.ts index 651f4ba..308b60e 100644 --- a/src/components/workspace-shell.test.ts +++ b/src/components/workspace-shell.test.ts @@ -274,6 +274,91 @@ describe("WorkspaceShell", () => { expect(countFetches(fetch, "/api/sources/source_2/chunks")).toBe(0); }); + it("keeps a remote document open without refreshing sources after chunks load", async () => { + const remoteSourceId = "knowhere-doc:default:doc_remote"; + const encodedRemoteSourceId = encodeURIComponent(remoteSourceId); + const fetch = vi.fn(async (input) => { + const url = getRequestURL(input); + + if ( + url.pathname === `/api/sources/${encodedRemoteSourceId}/chunks` + ) { + return Response.json({ + chunks: [ + { + chunkId: "remote_page_1", + documentId: "doc_remote", + sectionPath: "Page 1", + type: "page", + content: "Remote page summary.", + sourceTitle: "remote.pdf", + pageNums: [1], + pageAssets: [ + { + pageNumber: 1, + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + ], + }, + ], + pagination: { + page: Number(url.searchParams.get("page") ?? "1"), + pageSize: 50, + total: 1, + totalPages: 1, + }, + }); + } + + if (url.pathname === "/api/sources") { + return Response.json({ + sources: [ + { + id: "source_localized", + kind: "workspace", + title: "remote.pdf", + status: "ready", + documentId: "doc_remote", + }, + ], + }); + } + + return Response.json({ message: "Unexpected request" }, { status: 404 }); + }); + vi.stubGlobal("fetch", fetch); + + render( + React.createElement(C, { + sources: [ + { + id: remoteSourceId, + kind: "remote", + title: "remote.pdf", + status: "ready", + documentId: "doc_remote", + excludedFromQuery: false, + }, + ], + }), + ); + + const desktopChunksPanel = within(screen.getByTestId("desktop-chunks-panel")); + await waitFor(() => { + expect(desktopChunksPanel.getByRole("img", { name: "Page 1" })) + .toBeTruthy(); + }); + + expect(countFetches(fetch, "/api/sources")).toBe(0); + expect( + countFetches( + fetch, + `/api/sources/${encodedRemoteSourceId}/chunks`, + ), + ).toBe(1); + }); + it("focuses guest citations on desktop using loaded demo chunks", async () => { const fetch = vi.fn(async (input) => { const url = getRequestURL(input); diff --git a/src/components/workspace-shell.tsx b/src/components/workspace-shell.tsx index bd45e2f..7dd43dc 100644 --- a/src/components/workspace-shell.tsx +++ b/src/components/workspace-shell.tsx @@ -114,7 +114,6 @@ function WorkspaceShellContent({ fetchChunks: workspaceClient.fetchChunks, initialPrefetchedChunksBySourceId: initialPrefetchedChunksBySourceId ?? undefined, - onRemoteSourceChunksLoaded: sourceWorkflow.handleSourcesRefresh, onSelectSource: handleCitationSourceSelected, selectedSourceId: sourceWorkflow.selectedSourceId, sources: sourceWorkflow.sources, From 15db165812fd1c8c38df77ee1871291e44d7bb90 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 8 Jul 2026 03:08:49 +0800 Subject: [PATCH 31/51] Fix inline parsed image viewing --- .../assets/[...assetPath]/route.test.ts | 51 ++++++ .../assets/[...assetPath]/route.ts | 1 + .../api/parsed-assets/inline/route.test.ts | 96 +++++++++++ src/app/api/parsed-assets/inline/route.ts | 159 ++++++++++++++++++ src/components/parsed-chunk-card.test.ts | 57 +++++++ src/components/parsed-chunk-card.tsx | 37 +++- 6 files changed, 396 insertions(+), 5 deletions(-) create mode 100644 src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.test.ts create mode 100644 src/app/api/parsed-assets/inline/route.test.ts create mode 100644 src/app/api/parsed-assets/inline/route.ts diff --git a/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.test.ts b/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.test.ts new file mode 100644 index 0000000..aa1d09b --- /dev/null +++ b/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.test.ts @@ -0,0 +1,51 @@ +import { beforeEach, describe, expect, it, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + fetch: vi.fn(), +})) + +vi.mock("@/integrations/knowhere-demo", () => ({ + knowhereDemoApi: { + resolveApiURL: (pathname: string) => `https://demo.example${pathname}`, + }, +})) + +import { GET } from "./route" + +describe("GET /api/demo-sources/[demoSourceId]/assets/[...assetPath]", () => { + beforeEach(() => { + vi.clearAllMocks() + vi.stubGlobal("fetch", mocks.fetch) + }) + + it("serves demo image assets inline", async () => { + mocks.fetch.mockResolvedValue( + new Response("image bytes", { + status: 200, + headers: { + "content-type": "image/png", + }, + }), + ) + + const response = await GET( + new Request( + "http://localhost:3001/api/demo-sources/demo_1/assets/images/page-1.png", + ), + { + params: Promise.resolve({ + demoSourceId: "demo_1", + assetPath: ["images", "page-1.png"], + }), + }, + ) + + expect(response.status).toBe(200) + expect(response.headers.get("content-type")).toBe("image/png") + expect(response.headers.get("content-disposition")).toBe("inline") + expect(mocks.fetch).toHaveBeenCalledWith( + "https://demo.example/api/v1/demo/sources/demo_1/assets/images/page-1.png", + { cache: "no-store" }, + ) + }) +}) diff --git a/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.ts b/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.ts index 247f645..7dc89c3 100644 --- a/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.ts +++ b/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.ts @@ -33,6 +33,7 @@ export async function GET( status: 200, headers: { "content-type": response.headers.get("content-type") ?? "application/octet-stream", + "content-disposition": "inline", "cache-control": "public, max-age=3600", }, }) diff --git a/src/app/api/parsed-assets/inline/route.test.ts b/src/app/api/parsed-assets/inline/route.test.ts new file mode 100644 index 0000000..1bff9d3 --- /dev/null +++ b/src/app/api/parsed-assets/inline/route.test.ts @@ -0,0 +1,96 @@ +import { NextRequest } from "next/server" +import { beforeEach, describe, expect, it, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + getAuthenticated: vi.fn(), + getBlob: vi.fn(), +})) + +vi.mock("@/domains/workspace/request-context", () => ({ + notebookRequestContext: { + getAuthenticated: mocks.getAuthenticated, + }, +})) + +vi.mock("@vercel/blob", () => ({ + BlobNotFoundError: class BlobNotFoundError extends Error {}, + get: mocks.getBlob, +})) + +import { GET } from "./route" + +describe("GET /api/parsed-assets/inline", () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.getAuthenticated.mockResolvedValue({ + workspace: { id: "workspace_1" }, + }) + }) + + it("serves Notebook parsed document image blobs inline", async () => { + const assetUrl = + "https://store.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_1/rev_1/page_citation_assets/page-1.png" + mocks.getBlob.mockResolvedValue({ + statusCode: 200, + stream: new Response("image bytes").body, + blob: { + contentType: "image/png", + }, + }) + + const response = await GET( + new NextRequest( + `http://localhost:3001/api/parsed-assets/inline?url=${encodeURIComponent( + assetUrl, + )}`, + ), + ) + + expect(response.status).toBe(200) + expect(response.headers.get("content-type")).toBe("image/png") + expect(response.headers.get("content-disposition")).toBe("inline") + await expect(response.text()).resolves.toBe("image bytes") + expect(mocks.getBlob).toHaveBeenCalledWith( + "workspaces/workspace_1/parsed-documents/doc_1/rev_1/page_citation_assets/page-1.png", + { access: "public" }, + ) + }) + + it("rejects blob URLs outside the authenticated workspace", async () => { + const assetUrl = + "https://store.public.blob.vercel-storage.com/workspaces/workspace_2/parsed-documents/doc_1/rev_1/page_citation_assets/page-1.png" + + const response = await GET( + new NextRequest( + `http://localhost:3001/api/parsed-assets/inline?url=${encodeURIComponent( + assetUrl, + )}`, + ), + ) + + expect(response.status).toBe(404) + expect(mocks.getBlob).not.toHaveBeenCalled() + }) + + it("does not serve non-image parsed assets through the inline image route", async () => { + const assetUrl = + "https://store.public.blob.vercel-storage.com/workspaces/workspace_1/sources/source_1/parsed-result/tables/table-1.html" + mocks.getBlob.mockResolvedValue({ + statusCode: 200, + stream: new Response("
").body, + blob: { + contentType: "text/html; charset=utf-8", + }, + }) + + const response = await GET( + new NextRequest( + `http://localhost:3001/api/parsed-assets/inline?url=${encodeURIComponent( + assetUrl, + )}`, + ), + ) + + expect(response.status).toBe(415) + }) +}) diff --git a/src/app/api/parsed-assets/inline/route.ts b/src/app/api/parsed-assets/inline/route.ts new file mode 100644 index 0000000..9e9beb3 --- /dev/null +++ b/src/app/api/parsed-assets/inline/route.ts @@ -0,0 +1,159 @@ +import { BlobNotFoundError, get } from "@vercel/blob" +import type { NextRequest } from "next/server" + +import { notebookRequestContext } from "@/domains/workspace/request-context" + +type ParsedAssetBlobPath = { + readonly pathname: string + readonly assetPath: string +} + +const notebookBlobHostSuffix = ".blob.vercel-storage.com" +const inlineImageContentTypes = new Set([ + "image/gif", + "image/jpeg", + "image/png", + "image/webp", +]) +const parsedDocumentsPathPattern = + /^\/?workspaces\/([^/]+)\/parsed-documents\/[^/]+\/[^/]+\/(.+)$/u +const parsedResultPathPattern = + /^\/?workspaces\/([^/]+)\/sources\/[^/]+\/parsed-result\/(.+)$/u + +export async function GET(request: NextRequest): Promise { + const { workspace } = await notebookRequestContext.getAuthenticated() + const parsedAssetBlobPath = getParsedAssetBlobPath( + request.nextUrl.searchParams.get("url"), + workspace.id, + ) + if (!parsedAssetBlobPath) { + return Response.json({ message: "Parsed asset not found." }, { status: 404 }) + } + + try { + const blob = await get(parsedAssetBlobPath.pathname, { access: "public" }) + if (!blob) { + return Response.json( + { message: "Parsed asset not found." }, + { status: 404 }, + ) + } + if (blob.statusCode !== 200) { + return new Response(null, { status: 304 }) + } + + const contentType = getInlineImageContentType( + blob.blob.contentType, + parsedAssetBlobPath.assetPath, + ) + if (!contentType) { + return Response.json( + { message: "Parsed asset is not an image." }, + { status: 415 }, + ) + } + + return new Response(blob.stream, { + status: 200, + headers: { + "content-type": contentType, + "content-disposition": "inline", + "cache-control": "public, max-age=3600", + }, + }) + } catch (error) { + if (error instanceof BlobNotFoundError) { + return Response.json( + { message: "Parsed asset not found." }, + { status: 404 }, + ) + } + throw error + } +} + +function getParsedAssetBlobPath( + assetUrl: string | null, + workspaceId: string, +): ParsedAssetBlobPath | null { + if (!assetUrl) return null + + const url = parseAbsoluteUrl(assetUrl) + if (!url) return null + if (url.protocol !== "https:") return null + if (!url.hostname.toLowerCase().endsWith(notebookBlobHostSuffix)) return null + + const pathname = decodePathname(url.pathname) + if (!pathname) return null + + const documentsMatch = parsedDocumentsPathPattern.exec(pathname) + if (documentsMatch) { + return getWorkspaceScopedAssetPath(documentsMatch, workspaceId, pathname) + } + + const parsedResultMatch = parsedResultPathPattern.exec(pathname) + if (parsedResultMatch) { + return getWorkspaceScopedAssetPath(parsedResultMatch, workspaceId, pathname) + } + + return null +} + +function getWorkspaceScopedAssetPath( + match: RegExpExecArray, + workspaceId: string, + pathname: string, +): ParsedAssetBlobPath | null { + const matchedWorkspaceId = match[1] + const assetPath = match[2] + if (matchedWorkspaceId !== workspaceId || !assetPath) return null + + return { + pathname: pathname.replace(/^\/+/u, ""), + assetPath, + } +} + +function getInlineImageContentType( + contentType: string | null, + assetPath: string, +): string | null { + const normalizedContentType = getBaseContentType(contentType) + if (normalizedContentType && inlineImageContentTypes.has(normalizedContentType)) { + return contentType + } + + return inferImageContentType(assetPath) +} + +function getBaseContentType(contentType: string | null): string | null { + const baseContentType = contentType?.split(";")[0]?.trim().toLowerCase() + return baseContentType && baseContentType.length > 0 ? baseContentType : null +} + +function inferImageContentType(assetPath: string): string | null { + const pathname = assetPath.toLowerCase().split("?")[0] ?? assetPath + if (pathname.endsWith(".png")) return "image/png" + if (pathname.endsWith(".jpg") || pathname.endsWith(".jpeg")) { + return "image/jpeg" + } + if (pathname.endsWith(".gif")) return "image/gif" + if (pathname.endsWith(".webp")) return "image/webp" + return null +} + +function parseAbsoluteUrl(value: string): URL | null { + try { + return new URL(value) + } catch { + return null + } +} + +function decodePathname(pathname: string): string | null { + try { + return decodeURIComponent(pathname) + } catch { + return null + } +} diff --git a/src/components/parsed-chunk-card.test.ts b/src/components/parsed-chunk-card.test.ts index ff4fbef..8bf5a42 100644 --- a/src/components/parsed-chunk-card.test.ts +++ b/src/components/parsed-chunk-card.test.ts @@ -123,6 +123,63 @@ describe("ParsedChunkCard", () => { expect(screen.getByTestId("page-asset-image-unavailable-4")).toBeTruthy(); }); + it("routes Notebook Blob page assets through the inline image endpoint", () => { + const assetUrl = + "https://store.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_1/rev_1/page_citation_assets/page-4.png"; + + render( + React.createElement(ParsedChunkCard, { + chunk: { + chunkId: "page_4", + type: "page", + content: "The summary should not be primary when an image exists.", + readableContent: "The summary should not be primary when an image exists.", + sourceTitle: "manual.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl, + contentType: "image/png", + }, + ], + }, + isFocused: false, + onReferenceClick: vi.fn(), + }), + ); + + const pageImage = screen.getByRole("img", { name: "Page 4" }); + expect(pageImage.getAttribute("src")).toBe( + `/api/parsed-assets/inline?url=${encodeURIComponent(assetUrl)}`, + ); + }); + + it("routes Notebook Blob image chunks through the inline image endpoint", () => { + const assetUrl = + "https://store.public.blob.vercel-storage.com/workspaces/workspace_1/sources/source_1/parsed-result/images/image-1.jpg"; + + render( + React.createElement(ParsedChunkCard, { + chunk: { + chunkId: "image_1", + type: "image", + content: "", + sourceTitle: "manual.pdf", + assetUrl, + summary: "A scanned diagram", + }, + isFocused: false, + onReferenceClick: vi.fn(), + }), + ); + + const image = screen.getByRole("img", { name: "A scanned diagram" }); + expect(image.getAttribute("src")).toBe( + `/api/parsed-assets/inline?url=${encodeURIComponent(assetUrl)}`, + ); + }); + it("routes resolved artifact reference clicks to the target chunk", async () => { const user = userEvent.setup(); const onReferenceClick = vi.fn(); diff --git a/src/components/parsed-chunk-card.tsx b/src/components/parsed-chunk-card.tsx index d794687..a8ddd6d 100644 --- a/src/components/parsed-chunk-card.tsx +++ b/src/components/parsed-chunk-card.tsx @@ -444,7 +444,8 @@ function PageCitationAssetImage({ readonly asset: NonNullable[number]; }): ReactNode { const [failedAssetUrl, setFailedAssetUrl] = useState(null); - const hasImageError = failedAssetUrl === asset.assetUrl; + const imageAssetUrl = getInlineImageAssetUrl(asset.assetUrl); + const hasImageError = failedAssetUrl === imageAssetUrl; const aspectRatio = asset.width && asset.height ? `${asset.width} / ${asset.height}` : undefined; @@ -463,11 +464,11 @@ function PageCitationAssetImage({ ) : ( // eslint-disable-next-line @next/next/no-img-element -- Page assets can be short-lived Knowhere URLs outside Next image optimization. {`Page setFailedAssetUrl(asset.assetUrl)} + onError={() => setFailedAssetUrl(imageAssetUrl)} /> )} @@ -509,6 +510,9 @@ function ImageChunkCard({ readonly sourceOriginalFile: SourceOriginalFileView | null; }): ReactNode { const imageAssetUrl = getImageChunkAssetUrl(chunk, sourceOriginalFile); + const inlineImageAssetUrl = imageAssetUrl + ? getInlineImageAssetUrl(imageAssetUrl) + : null; return ( - {imageAssetUrl ? ( + {inlineImageAssetUrl ? (
{/* eslint-disable-next-line @next/next/no-img-element -- Parsed artifact dimensions are not known before render. */} {chunk.summary @@ -559,6 +563,29 @@ function getImageChunkAssetUrl( return sourceOriginalFile.url; } +function getInlineImageAssetUrl(assetUrl: string): string { + if (!isNotebookBlobAssetUrl(assetUrl)) return assetUrl; + + return `/api/parsed-assets/inline?url=${encodeURIComponent(assetUrl)}`; +} + +function isNotebookBlobAssetUrl(assetUrl: string): boolean { + try { + const url = new URL(assetUrl); + if (!url.hostname.toLowerCase().endsWith(".blob.vercel-storage.com")) { + return false; + } + + const pathname = decodeURIComponent(url.pathname).toLowerCase(); + return ( + pathname.includes("/parsed-documents/") || + pathname.includes("/parsed-result/") + ); + } catch { + return false; + } +} + function hasPageCitationAssets(chunk: ParsedChunkView): boolean { return (chunk.pageAssets?.length ?? 0) > 0; } From 6d01583527e34ffe189f7c6629ec61a29d02ddcf Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 8 Jul 2026 03:31:45 +0800 Subject: [PATCH 32/51] Refactor chat to Knowhere tool runtime --- src/agent-harness/index.ts | 1 + src/agent-harness/knowhere-text.test.ts | 258 ++++++++++++ src/agent-harness/knowhere-text.ts | 300 ++++++++++++++ src/agent-harness/ledger.test.ts | 124 +++++- src/agent-harness/ledger.ts | 103 +++++ src/agent-harness/runtime.test.ts | 153 ++++--- src/agent-harness/runtime.ts | 504 ++++++++++++++++++------ src/agent-harness/types.ts | 52 ++- src/domains/chat/contracts.ts | 11 +- src/domains/chat/index.test.ts | 236 ++++++++++- src/domains/chat/index.ts | 9 + src/domains/chat/knowhere-tools.ts | 155 ++++++++ src/domains/chat/prompt.ts | 46 +-- src/domains/chat/route-answer.ts | 6 + src/domains/chat/service.test.ts | 2 + src/domains/chat/service.ts | 4 + 16 files changed, 1702 insertions(+), 262 deletions(-) create mode 100644 src/agent-harness/knowhere-text.test.ts create mode 100644 src/agent-harness/knowhere-text.ts create mode 100644 src/domains/chat/knowhere-tools.ts diff --git a/src/agent-harness/index.ts b/src/agent-harness/index.ts index bca503f..a1feff1 100644 --- a/src/agent-harness/index.ts +++ b/src/agent-harness/index.ts @@ -1,4 +1,5 @@ export * from "./ledger" +export * from "./knowhere-text" export * from "./runtime" export * from "./types" export * from "./validator" diff --git a/src/agent-harness/knowhere-text.test.ts b/src/agent-harness/knowhere-text.test.ts new file mode 100644 index 0000000..88159f5 --- /dev/null +++ b/src/agent-harness/knowhere-text.test.ts @@ -0,0 +1,258 @@ +import { describe, expect, it } from "vitest" +import type { + KnowledgeGrepResponse, + KnowledgeOutline, + KnowledgeReadResponse, + RetrievalQueryResponse, +} from "@ontos-ai/knowhere-sdk" + +import { createEvidenceLedger } from "./ledger" +import { knowhereToolText } from "./knowhere-text" + +describe("knowhereToolText", () => { + it("formats search results as tagged text with refs and no raw asset URLs", () => { + const ledger = createEvidenceLedger() + const response = makeSearchResponse() + const snapshot = ledger.addRetrievalResponse(response) + + const text = knowhereToolText.formatSearch({ + response, + retrievalCount: snapshot.retrievalCount, + chunks: snapshot.chunks, + assets: snapshot.assets, + }) + + expect(text).toContain('') + expect(text).toContain('ref="r1:result:1"') + expect(text).toContain('ref="asset:r1:result:1"') + expect(text).toContain("Call inspectImage") + expect(text).not.toContain("https://assets.example/page-1.png") + }) + + it("formats list and outline responses", () => { + const listText = knowhereToolText.formatListDocuments({ + documents: [ + { + documentId: "doc_1", + revisionKey: "job_1", + namespace: "notebook", + sourceFileName: "report.pdf", + title: "Report", + status: "ready", + }, + ], + }) + const outlineText = knowhereToolText.formatOutline(makeOutlineResponse()) + + expect(listText).toContain('') + expect(listText).toContain('documentId="doc_1"') + expect(outlineText).toContain( + '', + ) + expect(outlineText).toContain('sectionPath="Root / Revenue"') + }) + + it("formats full read chunk bodies without truncating large content", () => { + const ledger = createEvidenceLedger() + const largeContent = `BEGIN ${"full chunk body ".repeat(700)} END` + const response = makeReadResponse(largeContent) + const snapshot = ledger.addReadChunksResponse(response) + + const text = knowhereToolText.formatReadChunks({ + response, + chunks: snapshot.chunks, + assets: snapshot.assets, + }) + + expect(text).toContain('') + expect(text).toContain('ref="read1:chunk:1"') + expect(text).toContain(largeContent) + expect(text).not.toContain("...[truncated]") + expect(text).not.toContain("https://assets.example/page-1.png") + }) + + it("formats grep matches with continuation metadata", () => { + const ledger = createEvidenceLedger() + const response = makeGrepResponse() + const snapshot = ledger.addGrepChunksResponse(response) + + const text = knowhereToolText.formatGrepChunks({ + response, + chunks: snapshot.chunks, + assets: snapshot.assets, + }) + + expect(text).toContain('') + expect(text).toContain('truncated="true"') + expect(text).toContain('continuationCursor="cursor_2"') + expect(text).toContain('ref="grep1:match:1"') + expect(text).toContain("matched penalty snippet") + }) + + it("formats errors as tagged text", () => { + expect( + knowhereToolText.formatError({ + operation: "read_chunks", + message: "A documentId is required.", + }), + ).toBe( + [ + '', + "", + "A documentId is required.", + "", + "", + ].join("\n"), + ) + }) +}) + +function makeSearchResponse(): RetrievalQueryResponse { + return { + namespace: "notebook", + query: "page one", + routerUsed: "workflow_single_step", + answerText: null, + evidenceText: "Page one evidence.", + stopReason: "answer_done", + failureReason: null, + results: [ + { + content: "Page one summary.", + chunkType: "page", + score: 0.91, + assetUrl: "https://assets.example/page-1.png", + metadata: { + pageNums: [1], + pageAssets: [ + { + pageNum: 1, + artifactRef: "page_citation_assets/page-1.png", + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + ], + }, + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "Root / Page 1", + }, + }, + ], + referencedChunks: [], + } +} + +function makeReadResponse(content: string): KnowledgeReadResponse { + return { + document: { + localDocumentId: "doc_1", + documentId: "doc_1", + jobId: "job_1", + namespace: "notebook", + sourceFileName: "report.pdf", + chunkCount: 1, + typeCounts: { text: 0, image: 0, table: 0, page: 1 }, + resultDirectoryPath: "parsed-storage:doc_1/job_1", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }, + chunks: [ + { + position: 1, + chunkId: "chunk_page_1", + chunkType: "page", + contentSource: "content", + content, + readableContent: content, + sectionPath: "Root / Page 1", + sourceChunkPath: "pages/page-1.md", + filePath: "pages/page-1.png", + assetUrl: "https://assets.example/page-1.png", + pageNumbers: [1], + metadata: { + pageNums: [1], + pageAssets: [ + { + pageNum: 1, + artifactRef: "page_citation_assets/page-1.png", + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + ], + }, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, + } +} + +function makeOutlineResponse(): KnowledgeOutline { + return { + document: { + localDocumentId: "doc_1", + documentId: "doc_1", + jobId: "job_1", + namespace: "notebook", + sourceFileName: "report.pdf", + chunkCount: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + resultDirectoryPath: "parsed-storage:doc_1/job_1", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }, + totalChunks: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + sections: [ + { + sectionPath: "Root / Revenue", + sectionTitle: "Revenue", + sectionLevel: 2, + summary: "Revenue summary.", + startChunk: 1, + endChunk: 1, + chunkCount: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + children: [], + }, + ], + sectionTree: [], + } +} + +function makeGrepResponse(): KnowledgeGrepResponse { + return { + document: { + localDocumentId: "doc_1", + documentId: "doc_1", + jobId: "job_1", + namespace: "notebook", + sourceFileName: "contract.pdf", + chunkCount: 4, + typeCounts: { text: 4, image: 0, table: 0, page: 0 }, + resultDirectoryPath: "parsed-storage:doc_1/job_1", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }, + matches: [ + { + position: 3, + chunkId: "chunk_3", + chunkType: "text", + sectionPath: "Root / Penalties", + sourceChunkPath: "chunks/chunk-3.md", + filePath: "contract.pdf", + startOffset: 10, + endOffset: 17, + snippet: "matched penalty snippet", + }, + ], + scannedChunks: 4, + truncated: true, + continuationCursor: "cursor_2", + } +} diff --git a/src/agent-harness/knowhere-text.ts b/src/agent-harness/knowhere-text.ts new file mode 100644 index 0000000..69de846 --- /dev/null +++ b/src/agent-harness/knowhere-text.ts @@ -0,0 +1,300 @@ +import type { + KnowledgeGrepResponse, + KnowledgeOutline, + KnowledgeReadResponse, + RetrievalQueryResponse, +} from "@ontos-ai/knowhere-sdk" + +import type { + EvidenceAsset, + EvidenceChunk, + KnowhereListDocumentsResponse, +} from "./types" + +type EvidenceDelta = { + readonly chunks: readonly EvidenceChunk[] + readonly assets: readonly EvidenceAsset[] +} + +type SearchTextInput = EvidenceDelta & { + readonly response: RetrievalQueryResponse + readonly retrievalCount: number +} + +type ReadChunksTextInput = EvidenceDelta & { + readonly response: KnowledgeReadResponse +} + +type GrepChunksTextInput = EvidenceDelta & { + readonly response: KnowledgeGrepResponse +} + +type ErrorTextInput = { + readonly operation: KnowhereOperation + readonly message: string +} + +type KnowhereOperation = + | "search" + | "list_documents" + | "get_document_outline" + | "read_chunks" + | "grep_chunks" + +const assetInstruction = + "Notebook returned image/page asset refs. Call inspectImage with the asset refs when OCR, visual details, or verification are needed. Do not expose raw asset URLs." + +export const knowhereToolText = { + formatSearch(input: SearchTextInput): string { + return wrapKnowhereBlock("search", [ + formatTag("summary", { + retrievalCount: String(input.retrievalCount), + namespace: input.response.namespace, + query: input.response.query, + resultCount: String(input.response.results.length), + referencedChunkCount: String(input.response.referencedChunks.length), + stopReason: input.response.stopReason ?? undefined, + failureReason: input.response.failureReason ?? undefined, + }), + formatOptionalTextTag("evidence", input.response.evidenceText), + formatEvidenceChunks(input.chunks), + formatEvidenceAssets(input.assets), + formatAssetInstruction(input.assets), + ]) + }, + + formatListDocuments(response: KnowhereListDocumentsResponse): string { + return wrapKnowhereBlock("list_documents", [ + formatTag("summary", { documentCount: String(response.documents.length) }), + ...response.documents.map((document, index) => + formatSelfClosingTag("document", { + index: String(index + 1), + documentId: document.documentId, + localDocumentId: document.localDocumentId, + revisionKey: document.revisionKey, + namespace: document.namespace, + sourceFileName: document.sourceFileName, + title: document.title, + status: document.status, + chunkCount: + typeof document.chunkCount === "number" + ? String(document.chunkCount) + : undefined, + }), + ), + ]) + }, + + formatOutline(response: KnowledgeOutline): string { + return wrapKnowhereBlock("get_document_outline", [ + formatTag("document", { + documentId: response.document.documentId, + localDocumentId: response.document.localDocumentId, + revisionKey: response.document.jobId, + sourceFileName: response.document.sourceFileName, + totalChunks: String(response.totalChunks), + truncated: response.truncated === true ? "true" : undefined, + continuationCursor: response.continuationCursor, + }), + ...response.sections.map((section) => formatSection(section, 0)), + ]) + }, + + formatReadChunks(input: ReadChunksTextInput): string { + return wrapKnowhereBlock("read_chunks", [ + formatTag("document", { + documentId: input.response.document.documentId, + localDocumentId: input.response.document.localDocumentId, + revisionKey: input.response.document.jobId, + sourceFileName: input.response.document.sourceFileName, + page: + typeof input.response.page === "number" + ? String(input.response.page) + : undefined, + pageSize: + typeof input.response.pageSize === "number" + ? String(input.response.pageSize) + : undefined, + totalChunks: + typeof input.response.totalChunks === "number" + ? String(input.response.totalChunks) + : undefined, + totalPages: + typeof input.response.totalPages === "number" + ? String(input.response.totalPages) + : undefined, + nextChunk: + typeof input.response.nextChunk === "number" + ? String(input.response.nextChunk) + : undefined, + }), + formatEvidenceChunks(input.chunks), + formatEvidenceAssets(input.assets), + formatAssetInstruction(input.assets), + ]) + }, + + formatGrepChunks(input: GrepChunksTextInput): string { + return wrapKnowhereBlock("grep_chunks", [ + formatTag("document", { + documentId: input.response.document.documentId, + localDocumentId: input.response.document.localDocumentId, + revisionKey: input.response.document.jobId, + sourceFileName: input.response.document.sourceFileName, + matchCount: String(input.response.matches.length), + scannedChunks: String(input.response.scannedChunks), + truncated: input.response.truncated ? "true" : "false", + continuationCursor: input.response.continuationCursor, + }), + formatEvidenceChunks(input.chunks), + ]) + }, + + formatError(input: ErrorTextInput): string { + return [ + formatOpenTag("knowhere", { + operation: input.operation, + status: "error", + }), + formatTextTag("message", input.message), + "", + ].join("\n") + }, +} as const + +function wrapKnowhereBlock( + operation: KnowhereOperation, + parts: readonly string[], +): string { + return [ + formatOpenTag("knowhere", { operation, status: "ok" }), + ...parts.filter((part) => part.trim().length > 0), + "", + ].join("\n") +} + +function formatEvidenceChunks(chunks: readonly EvidenceChunk[]): string { + if (chunks.length === 0) return "" + + return [ + "", + ...chunks.map((chunk) => + [ + formatOpenTag("chunk", { + ref: chunk.ref, + kind: chunk.kind, + chunkId: chunk.chunkId, + chunkType: chunk.chunkType, + score: chunk.score === null ? undefined : String(chunk.score), + documentId: chunk.source.documentId ?? undefined, + sourceFileName: chunk.source.sourceFileName ?? undefined, + sectionPath: chunk.source.sectionPath ?? undefined, + sourceChunkPath: chunk.sourceChunkPath ?? undefined, + filePath: chunk.filePath ?? undefined, + assetRef: chunk.assetRef, + }), + formatTextTag("content", chunk.content), + "", + ].join("\n"), + ), + "", + ].join("\n") +} + +function formatEvidenceAssets(assets: readonly EvidenceAsset[]): string { + if (assets.length === 0) return "" + + return [ + "", + ...assets.map((asset) => + formatSelfClosingTag("asset", { + ref: asset.ref, + chunkRef: asset.chunkRef, + type: asset.type, + label: asset.label, + sourcePath: asset.sourcePath, + documentId: asset.source.documentId ?? undefined, + sectionPath: asset.source.sectionPath ?? undefined, + }), + ), + "", + ].join("\n") +} + +function formatAssetInstruction(assets: readonly EvidenceAsset[]): string { + if (!assets.some((asset) => asset.type === "image")) return "" + return formatTextTag("asset_instruction", assetInstruction) +} + +function formatSection( + section: KnowledgeOutline["sections"][number], + depth: number, +): string { + return [ + formatOpenTag("section", { + depth: String(depth), + sectionPath: section.sectionPath, + sectionTitle: section.sectionTitle, + sectionLevel: String(section.sectionLevel), + startChunk: + typeof section.startChunk === "number" + ? String(section.startChunk) + : undefined, + endChunk: + typeof section.endChunk === "number" + ? String(section.endChunk) + : undefined, + chunkCount: String(section.chunkCount), + }), + formatOptionalTextTag("summary", section.summary), + ...section.children.map((child) => formatSection(child, depth + 1)), + "", + ] + .filter((part) => part.trim().length > 0) + .join("\n") +} + +function formatOptionalTextTag( + tagName: string, + value: string | null | undefined, +): string { + const trimmedValue = value?.trim() + return trimmedValue ? formatTextTag(tagName, trimmedValue) : "" +} + +function formatTextTag(tagName: string, value: string): string { + return [`<${tagName}>`, value, ``].join("\n") +} + +function formatTag( + tagName: string, + attrs: Readonly>, +): string { + return `${formatOpenTag(tagName, attrs)}` +} + +function formatSelfClosingTag( + tagName: string, + attrs: Readonly>, +): string { + return `${formatOpenTag(tagName, attrs).slice(0, -1)} />` +} + +function formatOpenTag( + tagName: string, + attrs: Readonly>, +): string { + const serializedAttrs = Object.entries(attrs) + .filter((entry): entry is [string, string] => typeof entry[1] === "string") + .map(([key, value]) => `${key}="${escapeAttribute(value)}"`) + .join(" ") + return serializedAttrs ? `<${tagName} ${serializedAttrs}>` : `<${tagName}>` +} + +function escapeAttribute(value: string): string { + return value + .replaceAll("&", "&") + .replaceAll('"', """) + .replaceAll("<", "<") + .replaceAll(">", ">") +} diff --git a/src/agent-harness/ledger.test.ts b/src/agent-harness/ledger.test.ts index 22e0f6a..061a892 100644 --- a/src/agent-harness/ledger.test.ts +++ b/src/agent-harness/ledger.test.ts @@ -1,5 +1,9 @@ import { describe, expect, it } from "vitest" -import type { RetrievalQueryResponse } from "@ontos-ai/knowhere-sdk" +import type { + KnowledgeGrepResponse, + KnowledgeReadResponse, + RetrievalQueryResponse, +} from "@ontos-ai/knowhere-sdk" import { createEvidenceLedger } from "./ledger" @@ -63,6 +67,50 @@ describe("createEvidenceLedger", () => { }), ) }) + + it("adds read chunk refs and page image assets", () => { + const ledger = createEvidenceLedger() + + const snapshot = ledger.addReadChunksResponse(makeReadResponse()) + + expect(snapshot.chunks).toEqual([ + expect.objectContaining({ + ref: "read1:chunk:1", + kind: "read_chunk", + content: "Full page content.", + contentPreview: "Full page content.", + assetRef: "asset:read1:chunk:1", + }), + ]) + expect(snapshot.assets).toEqual([ + expect.objectContaining({ + ref: "asset:read1:chunk:1", + chunkRef: "read1:chunk:1", + type: "image", + sourcePath: "page_citation_assets/page-1.png", + }), + ]) + }) + + it("adds grep match refs", () => { + const ledger = createEvidenceLedger() + + const snapshot = ledger.addGrepChunksResponse(makeGrepResponse()) + + expect(snapshot.chunks).toEqual([ + expect.objectContaining({ + ref: "grep1:match:1", + kind: "grep_match", + chunkId: "chunk_3", + content: "matched penalty snippet", + source: expect.objectContaining({ + documentId: "doc_contract", + sourceFileName: "contract.pdf", + sectionPath: "Root / Penalties", + }), + }), + ]) + }) }) function makeRetrievalResponse(): RetrievalQueryResponse { @@ -135,3 +183,77 @@ function makePageAssetUrlRetrievalResponse(): RetrievalQueryResponse { ], } } + +function makeReadResponse(): KnowledgeReadResponse { + return { + document: { + localDocumentId: "doc_contract", + documentId: "doc_contract", + jobId: "job_contract", + namespace: "notebook", + sourceFileName: "contract.pdf", + chunkCount: 1, + typeCounts: { text: 0, image: 0, table: 0, page: 1 }, + resultDirectoryPath: "parsed-storage:doc_contract/job_contract", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }, + chunks: [ + { + position: 1, + chunkId: "chunk_page_1", + chunkType: "page", + content: "Full page content.", + readableContent: "Full page content.", + sectionPath: "Root / Page 1", + sourceChunkPath: "pages/page-1.md", + filePath: "pages/page-1.png", + assetUrl: "https://assets.example/page-1.png", + pageNumbers: [1], + metadata: { + pageNums: [1], + pageAssets: [ + { + pageNum: 1, + artifactRef: "page_citation_assets/page-1.png", + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + ], + }, + }, + ], + } +} + +function makeGrepResponse(): KnowledgeGrepResponse { + return { + document: { + localDocumentId: "doc_contract", + documentId: "doc_contract", + jobId: "job_contract", + namespace: "notebook", + sourceFileName: "contract.pdf", + chunkCount: 4, + typeCounts: { text: 4, image: 0, table: 0, page: 0 }, + resultDirectoryPath: "parsed-storage:doc_contract/job_contract", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }, + matches: [ + { + position: 3, + chunkId: "chunk_3", + chunkType: "text", + sectionPath: "Root / Penalties", + sourceChunkPath: "chunks/chunk-3.md", + filePath: "contract.pdf", + startOffset: 10, + endOffset: 17, + snippet: "matched penalty snippet", + }, + ], + scannedChunks: 4, + truncated: false, + } +} diff --git a/src/agent-harness/ledger.ts b/src/agent-harness/ledger.ts index f9d859d..d782ae3 100644 --- a/src/agent-harness/ledger.ts +++ b/src/agent-harness/ledger.ts @@ -1,4 +1,8 @@ import type { + KnowledgeGrepMatch, + KnowledgeGrepResponse, + KnowledgeReadChunk, + KnowledgeReadResponse, RetrievalQueryResponse, RetrievalResult, } from "@ontos-ai/knowhere-sdk" @@ -14,6 +18,8 @@ const imageExtensions = [".jpg", ".jpeg", ".png", ".gif", ".webp", ".svg"] as co type MutableLedger = { retrievalCount: number + readCount: number + grepCount: number chunks: EvidenceChunk[] assets: EvidenceAsset[] evidenceText: string[] @@ -41,6 +47,8 @@ export type EvidenceLedger = ReturnType export function createEvidenceLedger() { const ledger: MutableLedger = { retrievalCount: 0, + readCount: 0, + grepCount: 0, chunks: [], assets: [], evidenceText: [], @@ -104,6 +112,38 @@ export function createEvidenceLedger() { return snapshot(ledger) }, + addReadChunksResponse(response: KnowledgeReadResponse): EvidenceLedgerSnapshot { + ledger.readCount += 1 + const readIndex = ledger.readCount + + response.chunks.forEach((chunk, index) => { + addChunkFromReadChunk({ + ledger, + response, + chunk, + ref: `read${readIndex}:chunk:${index + 1}`, + }) + }) + + return snapshot(ledger) + }, + + addGrepChunksResponse(response: KnowledgeGrepResponse): EvidenceLedgerSnapshot { + ledger.grepCount += 1 + const grepIndex = ledger.grepCount + + response.matches.forEach((match, index) => { + addChunkFromGrepMatch({ + ledger, + response, + match, + ref: `grep${grepIndex}:match:${index + 1}`, + }) + }) + + return snapshot(ledger) + }, + read(ref: string, offset = 0, limit = 4_000) { const chunk = ledger.chunks.find((candidate) => candidate.ref === ref) if (!chunk) { @@ -179,6 +219,69 @@ function addChunkFromResult(input: { }) } +function addChunkFromReadChunk(input: { + readonly ledger: MutableLedger + readonly response: KnowledgeReadResponse + readonly chunk: KnowledgeReadChunk + readonly ref: string +}): void { + addChunk({ + ledger: input.ledger, + chunk: { + ref: input.ref, + kind: "read_chunk", + chunkId: input.chunk.chunkId, + content: input.chunk.content, + contentPreview: buildContentPreview(input.chunk.content), + chunkType: input.chunk.chunkType, + score: null, + sourceChunkPath: input.chunk.sourceChunkPath, + filePath: input.chunk.filePath, + metadata: input.chunk.metadata, + source: { + documentId: input.response.document.documentId, + sourceFileName: input.response.document.sourceFileName, + sectionPath: input.chunk.sectionPath, + }, + revisionKey: input.response.document.jobId, + ...(input.chunk.assetUrl ? { assetUrl: input.chunk.assetUrl } : {}), + }, + }) +} + +function addChunkFromGrepMatch(input: { + readonly ledger: MutableLedger + readonly response: KnowledgeGrepResponse + readonly match: KnowledgeGrepMatch + readonly ref: string +}): void { + addChunk({ + ledger: input.ledger, + chunk: { + ref: input.ref, + kind: "grep_match", + chunkId: input.match.chunkId, + content: input.match.snippet, + contentPreview: buildContentPreview(input.match.snippet), + chunkType: input.match.chunkType, + score: null, + sourceChunkPath: input.match.sourceChunkPath, + filePath: input.match.filePath, + metadata: { + position: input.match.position, + startOffset: input.match.startOffset, + endOffset: input.match.endOffset, + }, + source: { + documentId: input.response.document.documentId, + sourceFileName: input.response.document.sourceFileName, + sectionPath: input.match.sectionPath, + }, + revisionKey: input.response.document.jobId, + }, + }) +} + function addChunk(input: { readonly ledger: MutableLedger readonly chunk: Omit diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index 3e3c5fa..e78de33 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -14,8 +14,8 @@ import type { ContextPolicy, HarnessToolCallTrace, IntentFrame, + KnowhereToolRuntime, OutputManifest, - RetrievalCapability, } from "./types" describe("agent harness runtime", () => { @@ -29,7 +29,7 @@ describe("agent harness runtime", () => { }) it("passes only outer retrieval parameters to KNOWHERE after intent and context policy are declared", async () => { - const query = vi.fn().mockResolvedValue( + const query = vi.fn().mockResolvedValue( makeRetrievalResponse(), ) const state: { @@ -40,14 +40,12 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state, ledger: createEvidenceLedger(), - retrieval: { query }, + knowhereTools: makeKnowhereTools(query), recentTurns: [], }) - expect(await executeTool(tools.retrieve, { query: "q4 chart" })).toEqual({ - ok: false, - message: "declareIntent must be called before retrieve.", - }) + expect(await executeTool(tools.knowhere_search, { query: "q4 chart" })) + .toContain('status="error"') await executeTool(tools.declareIntent, { task: "show_media", @@ -57,30 +55,27 @@ describe("agent harness runtime", () => { constraints: { desiredCount: 2, maxCount: 2 }, groundingPolicy: "must_use_sources", }) - expect(await executeTool(tools.retrieve, { query: "q4 chart" })).toEqual({ - ok: false, - message: "setContextPolicy must be called before retrieve.", - }) + expect(await executeTool(tools.knowhere_search, { query: "q4 chart" })) + .toContain("setContextPolicy must be called before knowhere_search.") await executeTool(tools.setContextPolicy, { carryHistory: "none", reason: "The current request is unrelated to previous turns.", activePriorTurnIds: [], }) - const result = await executeTool(tools.retrieve, { + const result = await executeTool(tools.knowhere_search, { query: "q4 chart", - modalities: ["image"], + targetContent: "image", topK: 2, purpose: "Find the two requested charts.", }) - expect(result).toMatchObject({ - ok: true, - retrievalCount: 1, - }) + expect(result).toContain('') + expect(result).toContain('ref="r1:result:1"') + expect(result).toContain('ref="asset:r1:result:1"') expect(query).toHaveBeenCalledWith({ query: "q4 chart", - modalities: ["image"], + targetContent: "image", topK: 2, purpose: "Find the two requested charts.", signalPaths: undefined, @@ -91,17 +86,17 @@ describe("agent harness runtime", () => { "LegalAction", ) expect(state.toolCalls?.map((call) => [call.tool, call.ok])).toEqual([ - ["retrieve", false], + ["knowhere_search", false], ["declareIntent", true], - ["retrieve", false], + ["knowhere_search", false], ["setContextPolicy", true], - ["retrieve", true], + ["knowhere_search", true], ]) }) - it("returns only newly retrieved evidence in each retrieve tool result", async () => { + it("returns only newly searched evidence in each knowhere_search tool result", async () => { const query = vi - .fn() + .fn() .mockResolvedValueOnce(makeRetrievalResponse()) .mockResolvedValueOnce({ ...makeRetrievalResponse(), @@ -143,22 +138,17 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state, ledger, - retrieval: { query }, + knowhereTools: makeKnowhereTools(query), recentTurns: [], }) - const firstResult = await executeTool(tools.retrieve, { query: "first" }) - const secondResult = await executeTool(tools.retrieve, { query: "second" }) + const firstResult = await executeTool(tools.knowhere_search, { query: "first" }) + const secondResult = await executeTool(tools.knowhere_search, { query: "second" }) - expect(firstResult).toMatchObject({ - retrievalCount: 1, - chunks: [{ ref: "r1:result:1" }], - }) - expect(secondResult).toMatchObject({ - retrievalCount: 2, - evidenceText: "Second evidence", - chunks: [{ ref: "r2:result:1" }], - }) + expect(firstResult).toContain('ref="r1:result:1"') + expect(secondResult).toContain('retrievalCount="2"') + expect(secondResult).toContain("Second evidence") + expect(secondResult).toContain('ref="r2:result:1"') expect(JSON.stringify(secondResult)).not.toContain("r1:result:1") expect(ledger.snapshot().chunks.map((chunk) => chunk.ref)).toEqual([ "r1:result:1", @@ -171,7 +161,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state: {}, ledger: createEvidenceLedger(), - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), inspectImages, recentTurns: [], }) @@ -183,12 +173,13 @@ describe("agent harness runtime", () => { expect(result).toEqual({ ok: false, - message: "retrieve must be called before inspectImage.", + message: + "Knowhere evidence tools must return image assets before inspectImage.", inspected: [], skipped: [ { ref: "asset:r1:result:1", - reason: "No retrieval evidence is available yet.", + reason: "No Knowhere evidence is available yet.", }, ], }) @@ -202,7 +193,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state: {}, ledger, - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), inspectImages, recentTurns: [], }) @@ -223,7 +214,7 @@ describe("agent harness runtime", () => { }, { ref: "missing", - reason: "Ref was not returned by retrieve as an asset.", + reason: "Ref was not returned by Knowhere as an asset.", }, ], }) @@ -236,7 +227,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state: {}, ledger, - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), inspectImages: vi.fn(), recentTurns: [], }) @@ -268,7 +259,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state: {}, ledger, - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), inspectImages, recentTurns: [], }) @@ -324,7 +315,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state: {}, ledger, - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), inspectImages, recentTurns: [], }) @@ -368,7 +359,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state, ledger: createEvidenceLedger(), - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), recentTurns: [], }) @@ -427,7 +418,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state, ledger: createEvidenceLedger(), - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), recentTurns: [ { id: "turn_1", @@ -459,7 +450,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state, ledger: createEvidenceLedger(), - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), recentTurns: [ { id: "turn_1", @@ -521,7 +512,7 @@ describe("agent harness runtime", () => { { type: "tool-call", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", input: { query: "q4" }, providerOptions: { google: { @@ -537,12 +528,10 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - }, + type: "text", + value: '', }, providerOptions: { google: { @@ -574,12 +563,10 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - }, + type: "text", + value: '', }, }, ], @@ -619,13 +606,11 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - assets: [{ ref: "asset:r1:referenced:1", type: "image" }], - }, + type: "text", + value: + '', }, }, ], @@ -645,13 +630,11 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - assets: [{ ref: "asset:r1:referenced:1", type: "image" }], - }, + type: "text", + value: + '', }, }, ], @@ -673,13 +656,11 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - chunks: [{ ref: "r1:result:1" }], - }, + type: "text", + value: + '', }, providerOptions: { google: { @@ -704,13 +685,11 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - chunks: [{ ref: "r1:result:1" }], - }, + type: "text", + value: + '', }, }, ], @@ -729,6 +708,18 @@ function executeTool(tool: unknown, input: unknown): Promise { return (tool as { execute: (input: unknown) => Promise }).execute(input) } +function makeKnowhereTools( + search: KnowhereToolRuntime["search"] = vi.fn(), +): KnowhereToolRuntime { + return { + search, + listDocuments: vi.fn().mockResolvedValue({ documents: [] }), + getDocumentOutline: vi.fn().mockRejectedValue(new Error("Not configured.")), + readChunks: vi.fn().mockRejectedValue(new Error("Not configured.")), + grepChunks: vi.fn().mockRejectedValue(new Error("Not configured.")), + } +} + function makeTurnInput(overrides: Partial = {}): AgentTurnInput { return { surface: "notebook_chat", diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index d38028e..e2b2783 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -9,6 +9,7 @@ import { import { z } from "zod" import { createEvidenceLedger } from "./ledger" +import { knowhereToolText } from "./knowhere-text" import type { AgentTurn, AgentTurnInput, @@ -20,9 +21,10 @@ import type { ImageInspectionResponse, InspectImages, IntentFrame, + KnowhereSearchRequest, + KnowhereSearchTargetContent, + KnowhereToolRuntime, OutputManifest, - RetrievalCapability, - TargetModality, } from "./types" import { validateOutputManifest } from "./validator" @@ -39,7 +41,7 @@ export type AgentHarnessModel = ToolLoopAgentSettings["model"] export type RunAgentHarnessInput = { readonly model: AgentHarnessModel readonly turn: AgentTurnInput - readonly retrieval: RetrievalCapability + readonly knowhereTools: KnowhereToolRuntime readonly inspectImages?: InspectImages readonly maxSteps?: number /** @@ -71,6 +73,53 @@ type HarnessStepPreparation = { } const targetModalitySchema = z.enum(["text", "image", "table"]) +const knowhereSearchTargetContentSchema = z.enum([ + "all", + "text", + "image", + "table", + "text_image", + "text_table", +]) +const knowledgeChunkTypeSchema = z.enum(["text", "image", "table", "page"]) + +const knowhereDocumentReferenceSchema = z.object({ + localDocumentId: z.string().min(1).optional(), + documentId: z.string().min(1).optional(), + jobId: z.string().min(1).optional(), + revisionKey: z.string().min(1).optional(), +}) + +const knowhereSearchSchema = z.object({ + query: z.string().min(1), + targetContent: knowhereSearchTargetContentSchema.default("all"), + purpose: z.string().optional(), + topK: z.number().int().min(1).max(12).optional(), + signalPaths: z.array(z.string().min(1)).max(8).optional(), + filterMode: z.enum(["keep", "delete"]).optional(), + threshold: z.number().min(0).max(1).optional(), +}) + +const knowhereReadChunksSchema = knowhereDocumentReferenceSchema.extend({ + page: z.number().int().min(1).optional(), + pageSize: z.number().int().min(1).max(50).optional(), + sectionPath: z.string().min(1).optional(), + startChunk: z.number().int().min(0).optional(), + endChunk: z.number().int().min(0).optional(), + chunkId: z.string().min(1).optional(), + chunkType: knowledgeChunkTypeSchema.optional(), +}) + +const knowhereGrepChunksSchema = knowhereDocumentReferenceSchema.extend({ + pattern: z.string().min(1), + continuationCursor: z.string().min(1).optional(), + isRegex: z.boolean().optional(), + isCaseSensitive: z.boolean().optional(), + maxResults: z.number().int().min(1).max(50).optional(), + chunkType: knowledgeChunkTypeSchema.optional(), + sectionPathPrefix: z.string().min(1).optional(), + contextChars: z.number().int().min(0).max(2_000).optional(), +}) const intentFrameSchema = z.object({ task: z.enum([ @@ -166,7 +215,7 @@ export async function runAgentHarness( const tools = createHarnessTools({ state, ledger, - retrieval: input.retrieval, + knowhereTools: input.knowhereTools, inspectImages: input.inspectImages, recentTurns: input.turn.recentTurns, }) @@ -364,7 +413,7 @@ function buildRevisionFeedback(errors: readonly string[]): string { function buildForcedFinalizationFeedback(): string { return [ "The retrieval step budget has been reached.", - "Do not search again or call any evidence-reading tools.", + "Do not search again or call any Knowhere evidence-reading tools.", "Use only the evidence and tool results already available in this turn.", "Call finalize now with the best supported answer.", "If the existing evidence is insufficient, explain the gap in unresolved", @@ -396,7 +445,7 @@ function hasUninspectedImageAssets(input: { export function createHarnessTools(input: { readonly state: HarnessToolState readonly ledger: ReturnType - readonly retrieval: RetrievalCapability + readonly knowhereTools: KnowhereToolRuntime readonly inspectImages?: InspectImages readonly recentTurns: readonly AgentTurn[] }) { @@ -433,82 +482,109 @@ export function createHarnessTools(input: { }), }), - retrieve: tool({ + knowhere_search: tool({ description: - "Ask KNOWHERE for evidence context. KNOWHERE handles internal navigation; this tool only submits a concise query and records returned evidence.", - inputSchema: z.object({ - query: z.string().min(1), - modalities: z.array(targetModalitySchema).default(["text"]), - purpose: z.string().optional(), - topK: z.number().int().min(1).max(12).optional(), - signalPaths: z.array(z.string().min(1)).max(8).optional(), - filterMode: z.enum(["keep", "delete"]).optional(), - threshold: z.number().min(0).max(1).optional(), - }), + "Search Knowhere for relevant Notebook evidence. Returns tagged text with evidence refs such as r1:result:1 and asset refs such as asset:r1:result:1.", + inputSchema: knowhereSearchSchema, execute: async (request) => traceToolCall(input.state, { - toolName: "retrieve", - inputSummary: summarizeRetrievalRequest(request), - execute: async () => { - if (!input.state.intent) { - return { - ok: false, - message: "declareIntent must be called before retrieve.", - } - } - if (!input.state.contextPolicy) { - return { - ok: false, - message: "setContextPolicy must be called before retrieve.", - } - } + toolName: "knowhere_search", + inputSummary: summarizeKnowhereSearchRequest(request), + execute: async () => + executeKnowhereSearch({ + state: input.state, + ledger: input.ledger, + knowhereTools: input.knowhereTools, + request, + }), + summarizeOutput: summarizeKnowhereTextOutput, + }), + }), - const beforeSnapshot = input.ledger.snapshot() - const response = await input.retrieval.query({ - query: request.query, - modalities: request.modalities as TargetModality[], - purpose: request.purpose, - topK: request.topK, - signalPaths: request.signalPaths, - filterMode: request.filterMode, - threshold: request.threshold, - }) - const snapshot = input.ledger.addRetrievalResponse(response) - const currentChunks = snapshot.chunks.slice( - beforeSnapshot.chunks.length, - ) - const currentAssets = snapshot.assets.slice( - beforeSnapshot.assets.length, - ) - return { - ok: true, - retrievalCount: snapshot.retrievalCount, - evidenceText: response.evidenceText ?? "", - stopReason: response.stopReason ?? null, - failureReason: response.failureReason ?? null, - chunks: currentChunks.map((chunk) => ({ - ref: chunk.ref, - kind: chunk.kind, - type: chunk.chunkType, - preview: chunk.contentPreview, - source: chunk.source, - assetRef: chunk.assetRef, - })), - assets: currentAssets.map((asset) => ({ - ref: asset.ref, - type: asset.type, - label: asset.label, - source: asset.source, - })), - } - }, - summarizeOutput: summarizeRetrieveOutput, + knowhere_list_documents: tool({ + description: + "List ready visible Notebook/Knowhere documents available for this chat turn. Use this to discover documentId and revisionKey before outline/read/grep.", + inputSchema: z.object({}), + execute: async () => + traceToolCall(input.state, { + toolName: "knowhere_list_documents", + inputSummary: {}, + execute: async () => + executeKnowhereTextTool({ + operation: "list_documents", + state: input.state, + execute: async () => + knowhereToolText.formatListDocuments( + await input.knowhereTools.listDocuments(), + ), + }), + summarizeOutput: summarizeKnowhereTextOutput, + }), + }), + + knowhere_get_document_outline: tool({ + description: + "Read a document outline from Knowhere parsed storage. Use documentId/revisionKey from knowhere_list_documents or search refs.", + inputSchema: knowhereDocumentReferenceSchema, + execute: async (request) => + traceToolCall(input.state, { + toolName: "knowhere_get_document_outline", + inputSummary: summarizeDocumentReference(request), + execute: async () => + executeKnowhereTextTool({ + operation: "get_document_outline", + state: input.state, + validate: () => validateDocumentReference(request), + execute: async () => + knowhereToolText.formatOutline( + await input.knowhereTools.getDocumentOutline(request), + ), + }), + summarizeOutput: summarizeKnowhereTextOutput, + }), + }), + + knowhere_read_chunks: tool({ + description: + "Read complete chunk bodies from Knowhere parsed storage. This tool never slices individual chunk content; control read size with page/pageSize, sectionPath, startChunk/endChunk, chunkId, and chunkType.", + inputSchema: knowhereReadChunksSchema, + execute: async (request) => + traceToolCall(input.state, { + toolName: "knowhere_read_chunks", + inputSummary: summarizeReadChunksRequest(request), + execute: async () => + executeKnowhereReadChunks({ + state: input.state, + ledger: input.ledger, + knowhereTools: input.knowhereTools, + request, + }), + summarizeOutput: summarizeKnowhereTextOutput, + }), + }), + + knowhere_grep_chunks: tool({ + description: + "Search chunk text with a literal or regex pattern. Returns bounded match snippets as grep refs such as grep1:match:1 and may include truncated=true with a continuationCursor.", + inputSchema: knowhereGrepChunksSchema, + execute: async (request) => + traceToolCall(input.state, { + toolName: "knowhere_grep_chunks", + inputSummary: summarizeGrepChunksRequest(request), + execute: async () => + executeKnowhereGrepChunks({ + state: input.state, + ledger: input.ledger, + knowhereTools: input.knowhereTools, + request, + }), + summarizeOutput: summarizeKnowhereTextOutput, }), }), inspectImage: tool({ description: - "Inspect retrieved image asset refs visually for OCR, visual details, comparisons, or verification. Use only after retrieve has returned image assets.", + "Inspect Knowhere image asset refs visually for OCR, visual details, comparisons, or verification. Use only after a Knowhere tool has returned image assets.", inputSchema: z.object({ refs: z.array(z.string().min(1)).min(1).max(imageInspectionRefLimit), question: z.string().min(1), @@ -529,32 +605,6 @@ export function createHarnessTools(input: { }), }), - readEvidence: tool({ - description: - "Read more text from an evidence chunk already returned by KNOWHERE.", - inputSchema: z.object({ - ref: z.string().min(1), - offset: z.number().int().min(0).optional(), - limit: z.number().int().min(1).max(8_000).optional(), - }), - execute: async (request) => - traceToolCall(input.state, { - toolName: "readEvidence", - inputSummary: { - ref: request.ref, - offset: request.offset ?? 0, - limit: request.limit ?? 4_000, - }, - execute: async () => - input.ledger.read( - request.ref, - request.offset ?? 0, - request.limit ?? 4_000, - ), - summarizeOutput: summarizeReadEvidenceOutput, - }), - }), - readPriorTurn: tool({ description: "Read the full text and citation labels of a specific prior turn by id " + @@ -689,14 +739,15 @@ async function inspectRetrievedImages(input: { } const snapshot = input.ledger.snapshot() - if (snapshot.retrievalCount === 0) { + if (snapshot.chunks.length === 0 && snapshot.assets.length === 0) { return { ok: false, - message: "retrieve must be called before inspectImage.", + message: + "Knowhere evidence tools must return image assets before inspectImage.", inspected: [], skipped: refs.map((ref) => ({ ref, - reason: "No retrieval evidence is available yet.", + reason: "No Knowhere evidence is available yet.", })), } } @@ -726,7 +777,7 @@ async function inspectRetrievedImages(input: { if (!asset) { skipped.push({ ref, - reason: "Ref was not returned by retrieve as an asset.", + reason: "Ref was not returned by Knowhere as an asset.", }) continue } @@ -804,6 +855,161 @@ async function inspectRetrievedImages(input: { } } +type KnowhereToolOperation = + | "search" + | "list_documents" + | "get_document_outline" + | "read_chunks" + | "grep_chunks" + +type KnowhereSearchToolRequest = z.infer +type KnowhereDocumentReferenceRequest = z.infer< + typeof knowhereDocumentReferenceSchema +> +type KnowhereReadChunksToolRequest = z.infer +type KnowhereGrepChunksToolRequest = z.infer +type DocumentReferenceSummary = { + readonly documentId?: string + readonly localDocumentId?: string + readonly hasJobId: boolean + readonly hasRevisionKey: boolean +} + +async function executeKnowhereSearch(input: { + readonly state: HarnessToolState + readonly ledger: ReturnType + readonly knowhereTools: KnowhereToolRuntime + readonly request: KnowhereSearchToolRequest +}): Promise { + return executeKnowhereTextTool({ + operation: "search", + state: input.state, + execute: async () => { + const beforeSnapshot = input.ledger.snapshot() + const response = await input.knowhereTools.search({ + query: input.request.query, + targetContent: input.request.targetContent, + purpose: input.request.purpose, + topK: input.request.topK, + signalPaths: input.request.signalPaths, + filterMode: input.request.filterMode, + threshold: input.request.threshold, + }) + const snapshot = input.ledger.addRetrievalResponse(response) + return knowhereToolText.formatSearch({ + response, + retrievalCount: snapshot.retrievalCount, + chunks: snapshot.chunks.slice(beforeSnapshot.chunks.length), + assets: snapshot.assets.slice(beforeSnapshot.assets.length), + }) + }, + }) +} + +async function executeKnowhereReadChunks(input: { + readonly state: HarnessToolState + readonly ledger: ReturnType + readonly knowhereTools: KnowhereToolRuntime + readonly request: KnowhereReadChunksToolRequest +}): Promise { + return executeKnowhereTextTool({ + operation: "read_chunks", + state: input.state, + validate: () => validateDocumentReference(input.request), + execute: async () => { + const beforeSnapshot = input.ledger.snapshot() + const response = await input.knowhereTools.readChunks(input.request) + const snapshot = input.ledger.addReadChunksResponse(response) + return knowhereToolText.formatReadChunks({ + response, + chunks: snapshot.chunks.slice(beforeSnapshot.chunks.length), + assets: snapshot.assets.slice(beforeSnapshot.assets.length), + }) + }, + }) +} + +async function executeKnowhereGrepChunks(input: { + readonly state: HarnessToolState + readonly ledger: ReturnType + readonly knowhereTools: KnowhereToolRuntime + readonly request: KnowhereGrepChunksToolRequest +}): Promise { + return executeKnowhereTextTool({ + operation: "grep_chunks", + state: input.state, + validate: () => validateDocumentReference(input.request), + execute: async () => { + const beforeSnapshot = input.ledger.snapshot() + const response = await input.knowhereTools.grepChunks(input.request) + const snapshot = input.ledger.addGrepChunksResponse(response) + return knowhereToolText.formatGrepChunks({ + response, + chunks: snapshot.chunks.slice(beforeSnapshot.chunks.length), + assets: snapshot.assets.slice(beforeSnapshot.assets.length), + }) + }, + }) +} + +async function executeKnowhereTextTool(input: { + readonly operation: KnowhereToolOperation + readonly state: HarnessToolState + readonly validate?: () => string | null + readonly execute: () => Promise +}): Promise { + const workflowError = validateKnowhereWorkflow(input.state, input.operation) + if (workflowError) { + return knowhereToolText.formatError({ + operation: input.operation, + message: workflowError, + }) + } + + const validationError = input.validate?.() + if (validationError) { + return knowhereToolText.formatError({ + operation: input.operation, + message: validationError, + }) + } + + try { + return await input.execute() + } catch (error) { + return knowhereToolText.formatError({ + operation: input.operation, + message: formatUnknownError(error), + }) + } +} + +function validateKnowhereWorkflow( + state: HarnessToolState, + operation: KnowhereToolOperation, +): string | null { + const toolName = `knowhere_${operation}` + if (!state.intent) return `declareIntent must be called before ${toolName}.` + if (!state.contextPolicy) { + return `setContextPolicy must be called before ${toolName}.` + } + return null +} + +function validateDocumentReference( + request: KnowhereDocumentReferenceRequest, +): string | null { + if ( + request.documentId || + request.localDocumentId || + request.jobId + ) { + return null + } + + return "A documentId, localDocumentId, or jobId is required." +} + function getUniqueTrimmedRefs(refs: readonly string[]): string[] { const normalizedRefs: string[] = [] for (const ref of refs) { @@ -860,6 +1066,7 @@ function recordToolCall( } function getToolTraceOk(output: unknown): boolean { + if (typeof output === "string") return !output.includes('status="error"') if (!isRecord(output)) return true if (typeof output.ok === "boolean") return output.ok if (typeof output.found === "boolean") return output.found @@ -884,9 +1091,9 @@ function summarizeContextPolicy(policy: ContextPolicy): unknown { } } -function summarizeRetrievalRequest(request: { +function summarizeKnowhereSearchRequest(request: { readonly query: string - readonly modalities?: readonly TargetModality[] + readonly targetContent?: KnowhereSearchTargetContent readonly purpose?: string readonly topK?: number readonly signalPaths?: readonly string[] @@ -895,7 +1102,7 @@ function summarizeRetrievalRequest(request: { }): unknown { return { query: request.query, - modalities: request.modalities ?? ["text"], + targetContent: request.targetContent ?? "all", purpose: request.purpose, topK: request.topK, signalPathCount: request.signalPaths?.length ?? 0, @@ -904,27 +1111,67 @@ function summarizeRetrievalRequest(request: { } } -function summarizeRetrieveOutput(output: unknown): unknown { - if (!isRecord(output)) return output +function summarizeDocumentReference( + request: KnowhereDocumentReferenceRequest, +): DocumentReferenceSummary { return { - ok: output.ok, - retrievalCount: output.retrievalCount, - stopReason: output.stopReason, - failureReason: output.failureReason, - chunkCount: Array.isArray(output.chunks) ? output.chunks.length : 0, - assetCount: Array.isArray(output.assets) ? output.assets.length : 0, + documentId: request.documentId, + localDocumentId: request.localDocumentId, + hasJobId: typeof request.jobId === "string", + hasRevisionKey: typeof request.revisionKey === "string", } } -function summarizeReadEvidenceOutput(output: unknown): unknown { - if (!isRecord(output)) return output +function summarizeReadChunksRequest( + request: KnowhereReadChunksToolRequest, +): unknown { return { - found: output.found, - ref: output.ref, - contentLength: output.contentLength, - offset: output.offset, - limit: output.limit, - hasMoreContent: output.hasMoreContent, + ...summarizeDocumentReference(request), + page: request.page, + pageSize: request.pageSize, + sectionPath: request.sectionPath, + startChunk: request.startChunk, + endChunk: request.endChunk, + chunkId: request.chunkId, + chunkType: request.chunkType, + } +} + +function summarizeGrepChunksRequest( + request: KnowhereGrepChunksToolRequest, +): unknown { + return { + ...summarizeDocumentReference(request), + patternLength: request.pattern.trim().length, + continuationCursor: request.continuationCursor, + isRegex: request.isRegex, + isCaseSensitive: request.isCaseSensitive, + maxResults: request.maxResults, + chunkType: request.chunkType, + sectionPathPrefix: request.sectionPathPrefix, + contextChars: request.contextChars, + } +} + +function summarizeKnowhereTextOutput(output: unknown): unknown { + if (typeof output !== "string") return output + return { + ok: !output.includes('status="error"'), + textLength: output.length, + chunkCount: countOccurrences(output, " { return typeof value === "object" && value !== null } +function formatUnknownError(error: unknown): string { + if (error instanceof Error) return error.message + return String(error) +} + export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { return [ "You are the outer Knowhere Agent Harness.", @@ -1007,10 +1259,10 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "1. Call declareIntent first. Capture constraints like a requested image/table count in constraints.desiredCount.", "2. Call setContextPolicy next, deciding how prior turns should influence this turn.", "3. When the policy needs prior-turn detail (references or corrections), call readPriorTurn for the relevant ids.", - "4. Call retrieve only when evidence is needed. The query must be concise and self-contained.", - "5. For pixel-level details, OCR, visual comparison, image verification, or when the likely answer is only visible on a retrieved page/image asset, call inspectImage only after retrieve returned image asset refs.", + "4. Call knowhere_search when relevance search is needed. Use knowhere_list_documents, knowhere_get_document_outline, knowhere_read_chunks, and knowhere_grep_chunks for focused document reads.", + "5. For pixel-level details, OCR, visual comparison, image verification, or when the likely answer is only visible on a returned page/image asset, call inspectImage only after a Knowhere tool returned image asset refs.", `6. inspectImage accepts at most ${imageInspectionRefLimit} image asset refs per call and per turn.`, - "7. Use readEvidence only for chunk refs already in the evidence ledger.", + "7. knowhere_read_chunks returns complete chunk bodies; control size with page/pageSize, sectionPath, startChunk/endChunk, chunkId, and chunkType.", "8. Call finalize with text, citations, artifacts, and unresolved issues. finalize requires declareIntent and setContextPolicy first.", "", "Context rules:", @@ -1022,7 +1274,7 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "- Final output is the OutputManifest passed to finalize, not freeform tool JSON or trailing text.", "- artifacts with display=true are the exact images/tables shown. Never display every candidate; honor constraints.desiredCount / maxCount.", "- Use type=derived_table only for tables you create from evidence; every derived_table.sourceRefs entry must reference evidence in the ledger.", - "- citations and selected image/table artifact refs may only reference refs returned by retrieve (in the evidence ledger).", + "- citations and selected image/table artifact refs may only reference refs returned by Knowhere tools in the evidence ledger.", "- inspectImage observations are inspection notes, not new source refs. Final citations and displayed image artifacts must use the original retrieved image asset refs.", "- If text evidence identifies a relevant page/image but does not include the exact fact, inspect the returned image asset for OCR/detail before saying the answer is unavailable.", "- If evidence is insufficient, list it in unresolved instead of fabricating facts.", diff --git a/src/agent-harness/types.ts b/src/agent-harness/types.ts index 4601f6f..83bb08d 100644 --- a/src/agent-harness/types.ts +++ b/src/agent-harness/types.ts @@ -1,4 +1,10 @@ import type { + KnowledgeDocumentReference, + KnowledgeGrepParams, + KnowledgeGrepResponse, + KnowledgeOutline, + KnowledgeReadParams, + KnowledgeReadResponse, RetrievalQueryParams, RetrievalQueryResponse, } from "@ontos-ai/knowhere-sdk" @@ -77,23 +83,57 @@ export type AgentTurnInput = { } } -export type HarnessRetrievalRequest = Pick< +export type KnowhereSearchTargetContent = + | "all" + | "text" + | "image" + | "table" + | "text_image" + | "text_table" + +export type KnowhereSearchRequest = Pick< RetrievalQueryParams, "query" | "topK" | "signalPaths" | "filterMode" | "threshold" > & { - readonly modalities: readonly TargetModality[] + readonly targetContent?: KnowhereSearchTargetContent readonly purpose?: string } -export type RetrievalCapability = { - readonly query: ( - input: HarnessRetrievalRequest, +export type KnowhereDocumentSummary = { + readonly documentId?: string + readonly localDocumentId?: string + readonly revisionKey?: string + readonly namespace?: string + readonly sourceFileName: string + readonly title?: string + readonly status?: string + readonly chunkCount?: number + readonly typeCounts?: Readonly> +} + +export type KnowhereListDocumentsResponse = { + readonly documents: readonly KnowhereDocumentSummary[] +} + +export type KnowhereToolRuntime = { + readonly search: ( + input: KnowhereSearchRequest, ) => Promise + readonly listDocuments: () => Promise + readonly getDocumentOutline: ( + input: KnowledgeDocumentReference, + ) => Promise + readonly readChunks: ( + input: KnowledgeReadParams, + ) => Promise + readonly grepChunks: ( + input: KnowledgeGrepParams, + ) => Promise } export type EvidenceChunk = { readonly ref: string - readonly kind: "result" | "referenced_chunk" + readonly kind: "result" | "referenced_chunk" | "read_chunk" | "grep_match" readonly chunkId?: string readonly content: string readonly contentPreview: string diff --git a/src/domains/chat/contracts.ts b/src/domains/chat/contracts.ts index 3400af5..e43b232 100644 --- a/src/domains/chat/contracts.ts +++ b/src/domains/chat/contracts.ts @@ -1,16 +1,22 @@ import type { + Knowledge, RetrievalQueryParams, RetrievalQueryResponse, } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" -import type { HarnessRunResult, InspectImages } from "@/agent-harness" +import type { + HarnessRunResult, + InspectImages, + KnowhereToolRuntime, +} from "@/agent-harness" import type { ChatArtifactView, ChatCitationView, } from "@/domains/chat/types" import type { HardenMediaAssetUrls } from "./media-asset-hardening" import type { HardenChatAssetUrl } from "./media-assets" +import type { NotebookKnowhereRemoteDocumentClient } from "./knowhere-tools" export type RetrievalClient = { query(params: RetrievalQueryParams): Promise @@ -57,6 +63,7 @@ export type GenerateAnswer = (input: { sources: readonly Source[] excludedSourceIds: readonly string[] searchSources: SearchSources + knowhereTools?: KnowhereToolRuntime inspectImages?: InspectImages }) => Promise @@ -67,6 +74,8 @@ export type AnswerQuestionInput = { sources: readonly Source[] excludedSourceIds: readonly string[] retrieval: RetrievalClient + knowledge?: Knowledge + remoteDocumentClient?: NotebookKnowhereRemoteDocumentClient generateAnswer: GenerateAnswer hardenChatAssetUrl?: HardenChatAssetUrl hardenMediaAssetUrls?: HardenMediaAssetUrls diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 9a34a4b..313d3ff 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -1,5 +1,11 @@ import { afterEach, describe, expect, it, vi } from "vitest" -import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" +import type { + Knowledge, + KnowledgeGrepResponse, + KnowledgeOutline, + KnowledgeReadResponse, + RetrievalResult, +} from "@ontos-ai/knowhere-sdk" import { Effect } from "effect" import { ToolLoopAgent } from "ai" import type { HarnessRunResult } from "@/agent-harness" @@ -8,6 +14,7 @@ import { answerQuestionWithRetrieval, generateAgenticOutputManifest, parseChatRequestBody, + type GenerateAnswer, type SearchSources, } from "." import type { @@ -88,6 +95,7 @@ describe("answerQuestionWithRetrieval", () => { sources, excludedSourceIds: ["source_2", "knowhere-doc:default:doc_remote"], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); expect(answer).toEqual({ answer: "The answer is grounded.", @@ -96,6 +104,132 @@ describe("answerQuestionWithRetrieval", () => { }); }); + it("exposes search, list, outline, read, and grep through the Knowhere tool runtime", async () => { + const result = makeRetrievalResult({ + chunkType: "image", + source: { + documentId: "doc_included", + sourceFileName: "notes.txt", + sectionPath: "images/diagram.png", + }, + }); + const retrieval = { + query: vi.fn().mockResolvedValue({ + results: [result], + evidenceText: "Diagram evidence.", + referencedChunks: [], + namespace: "notebook-workspace", + query: "diagram", + routerUsed: "workflow_single_step", + answerText: null, + }), + }; + const getDocumentOutline = vi.fn().mockResolvedValue(makeKnowledgeOutline()); + const readChunks = vi.fn().mockResolvedValue( + makeKnowledgeReadResponse("Full diagram chunk body."), + ); + const grepChunks = vi.fn().mockResolvedValue(makeKnowledgeGrepResponse()); + const knowledge = { + getDocumentOutline, + readChunks, + grepChunks, + } as unknown as Knowledge; + const listDocuments = vi.fn().mockResolvedValue({ + documents: [ + { + documentId: "doc_remote", + namespace: "default", + status: "ready", + currentJobResultId: "job_remote", + sourceFileName: "remote.pdf", + }, + ], + }); + const generateAnswer = vi.fn( + async ({ knowhereTools }: Parameters[0]) => { + if (!knowhereTools) throw new Error("Knowhere tools were not provided."); + + const searchResponse = await knowhereTools.search({ + query: "diagram", + targetContent: "image", + topK: 2, + }); + const documents = await knowhereTools.listDocuments(); + await knowhereTools.getDocumentOutline({ + documentId: "doc_included", + revisionKey: "job_123", + }); + await knowhereTools.readChunks({ + documentId: "doc_included", + revisionKey: "job_123", + page: 1, + pageSize: 2, + }); + await knowhereTools.grepChunks({ + documentId: "doc_included", + revisionKey: "job_123", + pattern: "diagram", + maxResults: 3, + }); + + expect(searchResponse.results).toEqual([result]); + expect( + documents.documents.map((document) => document.documentId), + ).toEqual(["doc_included", "doc_remote"]); + return makeHarnessRunResult("Runtime answer."); + }, + ); + const sources = [ + makeSource({ id: "source_included", knowhereDocumentId: "doc_included" }), + makeSource({ id: "source_excluded", knowhereDocumentId: "doc_excluded" }), + ]; + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "Show the diagram.", + namespace: "notebook-workspace", + sources, + excludedSourceIds: ["source_excluded"], + retrieval, + knowledge, + remoteDocumentClient: { documents: { list: listDocuments } }, + generateAnswer, + messages: [], + }), + ); + + expect(retrieval.query).toHaveBeenCalledWith({ + namespace: "notebook-workspace", + query: "diagram", + topK: 2, + useAgentic: false, + dataType: 3, + excludeDocumentIds: ["doc_excluded"], + }); + expect(listDocuments).toHaveBeenCalledWith({ + namespace: "default", + page: 1, + pageSize: 200, + }); + expect(getDocumentOutline).toHaveBeenCalledWith({ + documentId: "doc_included", + revisionKey: "job_123", + }); + expect(readChunks).toHaveBeenCalledWith({ + documentId: "doc_included", + revisionKey: "job_123", + page: 1, + pageSize: 2, + }); + expect(grepChunks).toHaveBeenCalledWith({ + documentId: "doc_included", + revisionKey: "job_123", + pattern: "diagram", + maxResults: 3, + }); + expect(answer.answer).toBe("Runtime answer."); + }); + it("does not carry no-evidence metadata from default into a successful legacy namespace result", async () => { const legacyResult = makeRetrievalResult({ source: { @@ -501,6 +635,7 @@ describe("answerQuestionWithRetrieval", () => { sources, excludedSourceIds: [], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); const expectedResult = { ...result, @@ -1568,6 +1703,7 @@ describe("answerQuestionWithRetrieval", () => { sources, excludedSourceIds: [], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); expect(retrieval.query).toHaveBeenCalledWith({ namespace: "notebook-workspace", @@ -1673,6 +1809,7 @@ describe("answerQuestionWithRetrieval", () => { sources: [makeSource({ title: "TSLA-Q4-2025-Update.pdf" })], excludedSourceIds: [], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); }); @@ -1821,9 +1958,9 @@ describe("generateAgenticOutputManifest", () => { reason: "The current request is self-contained.", activePriorTurnIds: [], }); - await tools.retrieve?.execute({ + await tools.knowhere_search?.execute({ query: "冯荣洲 身份证 图片", - modalities: ["text", "image"], + targetContent: "text_image", topK: 2, purpose: "Find exactly the requested identity-card images.", }); @@ -1954,9 +2091,9 @@ describe("generateAgenticOutputManifest", () => { reason: "The current request is self-contained.", activePriorTurnIds: [], }); - await tools.retrieve?.execute({ + await tools.knowhere_search?.execute({ query: "identity card front image", - modalities: ["image"], + targetContent: "image", topK: 1, purpose: "Find the ID card image to inspect.", }); @@ -2093,9 +2230,9 @@ describe("generateAgenticOutputManifest", () => { reason: "The current request is self-contained.", activePriorTurnIds: [], }); - await tools.retrieve?.execute({ + await tools.knowhere_search?.execute({ query: "进度计划 违约金 承包人", - modalities: ["text"], + targetContent: "text", topK: 6, purpose: "Find the contract clause and page for the liquidated damages amount.", }); @@ -2248,9 +2385,9 @@ describe("generateAgenticOutputManifest", () => { reason: "Self-contained request.", activePriorTurnIds: [], }); - await tools.retrieve?.execute({ + await tools.knowhere_search?.execute({ query: "身份证 图片", - modalities: ["image"], + targetContent: "image", topK: 3, purpose: "Find requested identity images.", }); @@ -2421,6 +2558,87 @@ function makeHarnessRunResult(text: string): HarnessRunResult { }; } +function makeKnowledgeOutline(): KnowledgeOutline { + return { + document: makeLocalKnowledgeDocument(), + totalChunks: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + sections: [ + { + sectionPath: "Root / Diagram", + sectionTitle: "Diagram", + sectionLevel: 2, + summary: "Diagram section.", + startChunk: 1, + endChunk: 1, + chunkCount: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + children: [], + }, + ], + sectionTree: [], + }; +} + +function makeKnowledgeReadResponse(content: string): KnowledgeReadResponse { + return { + document: makeLocalKnowledgeDocument(), + chunks: [ + { + position: 1, + chunkId: "chunk_1", + chunkType: "text", + content, + readableContent: content, + sectionPath: "Root / Diagram", + sourceChunkPath: "chunks/chunk-1.md", + filePath: "notes.txt", + metadata: {}, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, + }; +} + +function makeKnowledgeGrepResponse(): KnowledgeGrepResponse { + return { + document: makeLocalKnowledgeDocument(), + matches: [ + { + position: 1, + chunkId: "chunk_1", + chunkType: "text", + sectionPath: "Root / Diagram", + sourceChunkPath: "chunks/chunk-1.md", + filePath: "notes.txt", + startOffset: 0, + endOffset: 7, + snippet: "diagram", + }, + ], + scannedChunks: 1, + truncated: false, + }; +} + +function makeLocalKnowledgeDocument() { + return { + localDocumentId: "doc_included", + documentId: "doc_included", + jobId: "job_123", + namespace: "notebook-workspace", + sourceFileName: "notes.txt", + chunkCount: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + resultDirectoryPath: "parsed-storage:doc_included/job_123", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }; +} + type KnowhereQueryResponseLogMeta = { readonly query: string readonly resultCount: number diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index 4f5dee6..c0dc51b 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -40,6 +40,7 @@ import { } from "./media-assets" import { enrichRetrievalResultsWithPageCitationAssetUrls } from "./page-citation-assets" import type { HardenableRetrievalResult } from "./media-asset-hardening" +import { notebookKnowhereTools } from "./knowhere-tools" const DEFAULT_TOP_K = 8 const NOTEBOOK_USE_AGENTIC_RETRIEVAL: NonNullable< @@ -217,6 +218,14 @@ export const answerQuestionWithRetrieval = ( sources: input.sources, excludedSourceIds: input.excludedSourceIds, searchSources, + knowhereTools: notebookKnowhereTools.createRuntime({ + namespace: input.namespace, + sources: input.sources, + excludedSourceIds: input.excludedSourceIds, + searchSources, + knowledge: input.knowledge, + remoteDocumentClient: input.remoteDocumentClient, + }), ...(input.inspectImages ? { inspectImages: input.inspectImages } : {}), }), ) diff --git a/src/domains/chat/knowhere-tools.ts b/src/domains/chat/knowhere-tools.ts new file mode 100644 index 0000000..210b487 --- /dev/null +++ b/src/domains/chat/knowhere-tools.ts @@ -0,0 +1,155 @@ +import { Effect } from "effect" +import type { Knowledge } from "@ontos-ai/knowhere-sdk" + +import type { + KnowhereDocumentSummary, + KnowhereToolRuntime, +} from "@/agent-harness" +import type { Source } from "@/infrastructure/db/schema" +import { + listRemoteLibraryDocuments, + type RemoteLibraryDocument, +} from "@/domains/sources/remote-library" +import type { SearchSources } from "./contracts" +import { excludeDocuments } from "./retrieval" + +type RemoteDocumentClient = Parameters< + typeof listRemoteLibraryDocuments +>[0]["client"] + +export type NotebookKnowhereRemoteDocumentClient = RemoteDocumentClient + +type NotebookKnowhereToolsInput = { + readonly namespace: string + readonly sources: readonly Source[] + readonly excludedSourceIds: readonly string[] + readonly searchSources: SearchSources + readonly knowledge?: Knowledge + readonly remoteDocumentClient?: RemoteDocumentClient +} + +type SearchOnlyRuntimeInput = { + readonly searchSources: SearchSources +} + +export const notebookKnowhereTools = { + createRuntime(input: NotebookKnowhereToolsInput): KnowhereToolRuntime { + return { + search: (request) => input.searchSources(request), + listDocuments: async () => ({ + documents: await listVisibleDocuments(input), + }), + getDocumentOutline: async (request) => { + const knowledge = requireKnowledge(input.knowledge) + return knowledge.getDocumentOutline(request) + }, + readChunks: async (request) => { + const knowledge = requireKnowledge(input.knowledge) + return knowledge.readChunks(request) + }, + grepChunks: async (request) => { + const knowledge = requireKnowledge(input.knowledge) + return knowledge.grepChunks(request) + }, + } + }, + + createSearchOnlyRuntime(input: SearchOnlyRuntimeInput): KnowhereToolRuntime { + return { + search: (request) => input.searchSources(request), + listDocuments: async () => ({ documents: [] }), + getDocumentOutline: async () => { + throw new Error("Knowhere document outline is not configured.") + }, + readChunks: async () => { + throw new Error("Knowhere chunk reads are not configured.") + }, + grepChunks: async () => { + throw new Error("Knowhere grep is not configured.") + }, + } + }, +} as const + +async function listVisibleDocuments( + input: NotebookKnowhereToolsInput, +): Promise { + const excludedSourceIds = new Set(input.excludedSourceIds) + const excludedDocumentIds = new Set( + excludeDocuments(input.sources, input.excludedSourceIds) + .excludeDocumentIds ?? [], + ) + const localDocuments = input.sources + .filter( + (source): source is Source & { readonly knowhereDocumentId: string } => + source.status === "ready" && + Boolean(source.knowhereDocumentId) && + !excludedSourceIds.has(source.id) && + !excludedDocumentIds.has(source.knowhereDocumentId ?? ""), + ) + .map((source): KnowhereDocumentSummary => ({ + documentId: source.knowhereDocumentId, + revisionKey: source.knowhereJobId ?? undefined, + namespace: input.namespace, + sourceFileName: source.title, + title: source.title, + status: source.status, + })) + + const remoteDocuments = await listVisibleRemoteDocuments({ + input, + localDocuments, + excludedDocumentIds, + }) + + return [...localDocuments, ...remoteDocuments] +} + +async function listVisibleRemoteDocuments(input: { + readonly input: NotebookKnowhereToolsInput + readonly localDocuments: readonly KnowhereDocumentSummary[] + readonly excludedDocumentIds: ReadonlySet +}): Promise { + if (!input.input.remoteDocumentClient) return [] + + const localDocumentIds = new Set( + input.localDocuments.flatMap((document): string[] => + document.documentId ? [document.documentId] : [], + ), + ) + const documents = await Effect.runPromise( + listRemoteLibraryDocuments({ + workspace: { namespace: input.input.namespace }, + client: input.input.remoteDocumentClient, + localSources: input.input.sources, + }), + ) + + return documents + .filter( + (document) => + document.status === "ready" && + !localDocumentIds.has(document.documentId) && + !input.excludedDocumentIds.has(document.documentId), + ) + .map(toRemoteDocumentSummary) +} + +function toRemoteDocumentSummary( + document: RemoteLibraryDocument, +): KnowhereDocumentSummary { + return { + documentId: document.documentId, + revisionKey: document.revisionKey, + namespace: document.namespace, + sourceFileName: + document.sourceFileName ?? document.title ?? document.documentId, + title: document.title, + status: document.status, + } +} + +function requireKnowledge(knowledge: Knowledge | undefined): Knowledge { + if (knowledge) return knowledge + throw new Error("Knowhere parsed-document reads are not configured.") +} diff --git a/src/domains/chat/prompt.ts b/src/domains/chat/prompt.ts index efee008..f424252 100644 --- a/src/domains/chat/prompt.ts +++ b/src/domains/chat/prompt.ts @@ -8,17 +8,15 @@ import { runAgentHarness, type AgentTurn, type AgentTurnInput, - type HarnessRetrievalRequest, type HarnessRunResult, type InspectImages, - type TargetModality, + type KnowhereToolRuntime, } from "@/agent-harness" import type { - AgenticRetrievalQuery, - AgenticRetrievalTargetContent, ChatHistoryMessage, SearchSources, } from "./contracts" +import { notebookKnowhereTools } from "./knowhere-tools" const RECENT_CONTEXT_MESSAGE_LIMIT = 8 const CONTEXT_CONTENT_CHAR_LIMIT = 900 @@ -30,6 +28,7 @@ type GenerateAgenticOutputManifestInput = { sources: readonly Source[] excludedSourceIds: readonly string[] searchSources: SearchSources + knowhereTools?: KnowhereToolRuntime inspectImages?: InspectImages } @@ -59,10 +58,11 @@ export const generateAgenticOutputManifestEffect = ( runAgentHarness({ model: CHAT_MODEL, turn, - retrieval: { - query: (request) => - input.searchSources(toAgenticRetrievalQuery(request)), - }, + knowhereTools: + input.knowhereTools ?? + notebookKnowhereTools.createSearchOnlyRuntime({ + searchSources: input.searchSources, + }), ...(input.inspectImages ? { inspectImages: input.inspectImages } : {}), }), ) @@ -124,36 +124,6 @@ function getCitationLabels( .filter((label) => label.length > 0) } -function toAgenticRetrievalQuery( - request: HarnessRetrievalRequest, -): AgenticRetrievalQuery { - return { - query: request.query, - targetContent: toAgenticRetrievalTargetContent(request.modalities), - purpose: request.purpose, - topK: request.topK, - signalPaths: request.signalPaths, - filterMode: request.filterMode, - threshold: request.threshold, - } -} - -function toAgenticRetrievalTargetContent( - modalities: readonly TargetModality[], -): AgenticRetrievalTargetContent { - const requestedModalities = new Set(modalities) - if (requestedModalities.has("image") && requestedModalities.has("text")) { - return "text_image" - } - if (requestedModalities.has("table") && requestedModalities.has("text")) { - return "text_table" - } - if (requestedModalities.has("image")) return "image" - if (requestedModalities.has("table")) return "table" - if (requestedModalities.has("text")) return "text" - return "all" -} - function formatSourceContext( sources: readonly Source[], excludedSourceIds: readonly string[], diff --git a/src/domains/chat/route-answer.ts b/src/domains/chat/route-answer.ts index 2f8ae6b..3f9fc9d 100644 --- a/src/domains/chat/route-answer.ts +++ b/src/domains/chat/route-answer.ts @@ -27,6 +27,7 @@ import type { Source } from "@/infrastructure/db/schema" import { CHAT_MODEL } from "@/lib/ai" import type { HardenChatAssetUrl } from "./media-assets" import { isAuthError } from "@/integrations/dashboard/api-key-service" +import { makeKnowhereClientWithParsedStorage } from "@/integrations/knowhere" import { summarizeUnknownError } from "@/lib/format-log-value" import { logger } from "@/lib/logger" import { routeResult, type RouteResult } from "@/lib/route-result" @@ -98,6 +99,9 @@ const answerChatEffect = (input: AnswerChatInput) => const parsedStorage = new BlobParsedDocumentStorage({ workspaceId: workspace.id, }) + const knowhereResources = makeKnowhereClientWithParsedStorage(apiKey, { + workspaceId: workspace.id, + }) const hardenChatAssetUrl: HardenChatAssetUrl = async ({ source, sourcePath, @@ -140,6 +144,8 @@ const answerChatEffect = (input: AnswerChatInput) => threadId: body.value.threadId, excludedSourceIds: body.value.excludedSourceIds, retrieval: client.retrieval, + knowledge: knowhereResources.knowledge, + remoteDocumentClient: client, generateAnswer: generateAgenticOutputManifest, hardenChatAssetUrl, hardenMediaAssetUrls: ({ results, artifacts }) => diff --git a/src/domains/chat/service.test.ts b/src/domains/chat/service.test.ts index 12df745..441a12c 100644 --- a/src/domains/chat/service.test.ts +++ b/src/domains/chat/service.test.ts @@ -67,6 +67,7 @@ describe("handleChatTurn", () => { sources, excludedSourceIds: ["source_excluded"], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); expect(repository.appendMessageToThread).toHaveBeenNthCalledWith(1, "workspace_1", { threadId: "thread_1", @@ -216,6 +217,7 @@ describe("handleChatTurn", () => { sources, excludedSourceIds: [], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); expect(retrieval.query).toHaveBeenCalledWith({ namespace: "notebook-namespace", diff --git a/src/domains/chat/service.ts b/src/domains/chat/service.ts index bfa7e42..7a7b2b5 100644 --- a/src/domains/chat/service.ts +++ b/src/domains/chat/service.ts @@ -66,6 +66,8 @@ type ChatTurnInput = { threadId?: string excludedSourceIds: readonly string[] retrieval: RetrievalClient + knowledge?: AnswerQuestionInput["knowledge"] + remoteDocumentClient?: AnswerQuestionInput["remoteDocumentClient"] generateAnswer: GenerateAnswer hardenChatAssetUrl?: AnswerQuestionInput["hardenChatAssetUrl"] hardenMediaAssetUrls?: AnswerQuestionInput["hardenMediaAssetUrls"] @@ -126,6 +128,8 @@ export const handleChatTurnEffect = (input: ChatTurnInput) => sources: readySources, excludedSourceIds: input.excludedSourceIds, retrieval: input.retrieval, + knowledge: input.knowledge, + remoteDocumentClient: input.remoteDocumentClient, generateAnswer: input.generateAnswer, hardenChatAssetUrl: input.hardenChatAssetUrl, hardenMediaAssetUrls: input.hardenMediaAssetUrls, From b59db5aaacf259b6a2e4a35008fe61be10f1cca2 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 8 Jul 2026 04:10:00 +0800 Subject: [PATCH 33/51] fix(chat): validate citation source refs --- src/agent-harness/knowhere-text.ts | 1 + src/agent-harness/runtime.test.ts | 11 + src/agent-harness/runtime.ts | 10 +- src/agent-harness/validator.test.ts | 136 +++++++++++ src/agent-harness/validator.ts | 79 ++++++ src/domains/chat/index.test.ts | 364 ++++++++++++++++++++++++++-- src/domains/chat/index.ts | 55 +---- src/domains/chat/service.test.ts | 4 +- 8 files changed, 580 insertions(+), 80 deletions(-) diff --git a/src/agent-harness/knowhere-text.ts b/src/agent-harness/knowhere-text.ts index 69de846..1c19368 100644 --- a/src/agent-harness/knowhere-text.ts +++ b/src/agent-harness/knowhere-text.ts @@ -214,6 +214,7 @@ function formatEvidenceAssets(assets: readonly EvidenceAsset[]): string { label: asset.label, sourcePath: asset.sourcePath, documentId: asset.source.documentId ?? undefined, + sourceFileName: asset.source.sourceFileName ?? undefined, sectionPath: asset.source.sectionPath ?? undefined, }), ), diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index e78de33..eec6064 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -28,6 +28,17 @@ describe("agent harness runtime", () => { expect(prompt).not.toContain("navigation action") }) + it("tells the agent to cite only refs with matching source metadata", () => { + const prompt = buildHarnessSystemPrompt(makeTurnInput()) + + expect(prompt).toContain( + "source.documentId, sourceFileName, and sectionPath must match the selected evidence ref exactly", + ) + expect(prompt).toContain( + "Omit citations when you cannot identify a supporting evidence ref", + ) + }) + it("passes only outer retrieval parameters to KNOWHERE after intent and context policy are declared", async () => { const query = vi.fn().mockResolvedValue( makeRetrievalResponse(), diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index e2b2783..e272340 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -21,7 +21,6 @@ import type { ImageInspectionResponse, InspectImages, IntentFrame, - KnowhereSearchRequest, KnowhereSearchTargetContent, KnowhereToolRuntime, OutputManifest, @@ -406,7 +405,9 @@ function buildRevisionFeedback(errors: readonly string[]): string { "answer contract.", "Do not exceed the user's requested artifact count, only cite or display", "evidence refs that exist in the evidence ledger, and do not fabricate", - "facts when evidence is missing.", + "facts when evidence is missing. Citation source.documentId,", + "sourceFileName, and sectionPath must match the selected evidence ref", + "exactly.", ].join("\n") } @@ -668,7 +669,8 @@ export function createHarnessTools(input: { description: "Finalize the user-facing output manifest. This is the only final answer " + "contract. Artifacts listed here with display=true are the exact set of " + - "images/tables shown to the user; cite only refs from the evidence ledger.", + "images/tables shown to the user; cite only refs from the evidence ledger " + + "and copy citation source metadata exactly from the selected ref.", inputSchema: outputManifestSchema, execute: async (manifest) => traceToolCall(input.state, { @@ -1275,6 +1277,8 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "- artifacts with display=true are the exact images/tables shown. Never display every candidate; honor constraints.desiredCount / maxCount.", "- Use type=derived_table only for tables you create from evidence; every derived_table.sourceRefs entry must reference evidence in the ledger.", "- citations and selected image/table artifact refs may only reference refs returned by Knowhere tools in the evidence ledger.", + "- For every citation, source.documentId, sourceFileName, and sectionPath must match the selected evidence ref exactly. Use null or omit the field only when the ref omits it.", + "- Omit citations when you cannot identify a supporting evidence ref with matching source metadata.", "- inspectImage observations are inspection notes, not new source refs. Final citations and displayed image artifacts must use the original retrieved image asset refs.", "- If text evidence identifies a relevant page/image but does not include the exact fact, inspect the returned image asset for OCR/detail before saying the answer is unavailable.", "- If evidence is insufficient, list it in unresolved instead of fabricating facts.", diff --git a/src/agent-harness/validator.test.ts b/src/agent-harness/validator.test.ts index be8de3e..89bf0ba 100644 --- a/src/agent-harness/validator.test.ts +++ b/src/agent-harness/validator.test.ts @@ -115,6 +115,142 @@ describe("validateOutputManifest", () => { ) }) + it("rejects citation source metadata that conflicts with the resolved evidence ref", () => { + const validation = validateOutputManifest({ + manifest: makeManifest({ + text: "The source discusses information hiding.", + citations: [ + { + ref: "r1:result:1", + label: "wrong source", + source: { + documentId: "doc_claimed", + sourceFileName: "claimed-source.pdf", + sectionPath: "Claimed Section", + }, + }, + ], + }), + intent: makeIntent({}), + contextPolicy: unrelatedContextPolicy, + ledger: { + ...emptyLedger, + chunks: [ + { + ref: "r1:result:1", + kind: "result", + content: "Information hiding is a module design principle.", + contentPreview: "Information hiding is a module design principle.", + chunkType: "text", + score: 0.9, + source: { + documentId: "doc_information_hiding", + sourceFileName: "information_hiding.pdf", + sectionPath: "Root / Module Design", + }, + }, + ], + }, + surface: "notebook_chat", + }) + + expect(validation.errors).toContain( + "Citation ref 'r1:result:1' source.documentId must match resolved evidence source. Expected 'doc_information_hiding', received 'doc_claimed'.", + ) + expect(validation.errors).toContain( + "Citation ref 'r1:result:1' source.sourceFileName must match resolved evidence source. Expected 'information_hiding.pdf', received 'claimed-source.pdf'.", + ) + expect(validation.errors).toContain( + "Citation ref 'r1:result:1' source.sectionPath must match resolved evidence source. Expected 'Root / Module Design', received 'Claimed Section'.", + ) + }) + + it("accepts citation source metadata that exactly matches the resolved evidence ref", () => { + const validation = validateOutputManifest({ + manifest: makeManifest({ + text: "Revenue increased.", + citations: [ + { + ref: "r1:result:1", + label: "report.pdf / Q4", + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "Q4", + }, + }, + ], + }), + intent: makeIntent({}), + contextPolicy: unrelatedContextPolicy, + ledger: { + ...emptyLedger, + chunks: [ + { + ref: "r1:result:1", + kind: "result", + content: "Revenue increased.", + contentPreview: "Revenue increased.", + chunkType: "text", + score: 0.9, + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "Q4", + }, + }, + ], + }, + surface: "notebook_chat", + }) + + expect(validation.ok).toBe(true) + expect(validation.errors).toEqual([]) + }) + + it("rejects missing citation documentId when the resolved evidence has one", () => { + const validation = validateOutputManifest({ + manifest: makeManifest({ + text: "Revenue increased.", + citations: [ + { + ref: "r1:result:1", + label: "report.pdf / Q4", + source: { + sourceFileName: "report.pdf", + sectionPath: "Q4", + }, + }, + ], + }), + intent: makeIntent({}), + contextPolicy: unrelatedContextPolicy, + ledger: { + ...emptyLedger, + chunks: [ + { + ref: "r1:result:1", + kind: "result", + content: "Revenue increased.", + contentPreview: "Revenue increased.", + chunkType: "text", + score: 0.9, + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "Q4", + }, + }, + ], + }, + surface: "notebook_chat", + }) + + expect(validation.errors).toContain( + "Citation ref 'r1:result:1' source.documentId must match resolved evidence source. Expected 'doc_1', received missing.", + ) + }) + it("accepts source-backed derived tables and rejects missing source refs", () => { const validation = validateOutputManifest({ manifest: makeManifest({ diff --git a/src/agent-harness/validator.ts b/src/agent-harness/validator.ts index dd67343..fcd8c0e 100644 --- a/src/agent-harness/validator.ts +++ b/src/agent-harness/validator.ts @@ -3,6 +3,7 @@ import type { EvidenceLedgerSnapshot, HarnessToolCallTrace, IntentFrame, + OutputCitation, OutputManifest, } from "./types" @@ -21,6 +22,15 @@ export type ManifestValidationResult = { readonly errors: readonly string[] } +const citationSourceFields = [ + "documentId", + "sourceFileName", + "sectionPath", +] as const + +type CitationSourceField = (typeof citationSourceFields)[number] +type EvidenceSource = EvidenceLedgerSnapshot["chunks"][number]["source"] + export function validateOutputManifest( input: ManifestValidationInput, ): ManifestValidationResult { @@ -71,6 +81,14 @@ function validateArtifactRefs( ...input.ledger.chunks.map((chunk) => chunk.ref), ...input.ledger.assets.map((asset) => asset.ref), ]) + const sourcesByRef = new Map([ + ...input.ledger.chunks.map( + (chunk): readonly [string, EvidenceSource] => [chunk.ref, chunk.source], + ), + ...input.ledger.assets.map( + (asset): readonly [string, EvidenceSource] => [asset.ref, asset.source], + ), + ]) for (const artifact of input.manifest.artifacts) { if (artifact.type === "derived_table") { @@ -102,10 +120,71 @@ function validateArtifactRefs( for (const citation of input.manifest.citations) { if (!knownRefs.has(citation.ref)) { errors.push(`Citation ref '${citation.ref}' was not found in the evidence ledger.`) + continue } + + const source = sourcesByRef.get(citation.ref) + if (!source) continue + validateCitationSource( + { + ref: citation.ref, + declared: citation.source, + resolved: source, + }, + errors, + ) + } +} + +function validateCitationSource( + input: { + readonly ref: string + readonly declared: OutputCitation["source"] + readonly resolved: EvidenceSource + }, + errors: string[], +): void { + for (const field of citationSourceFields) { + validateCitationSourceField( + { + ref: input.ref, + field, + declaredValue: input.declared[field], + resolvedValue: input.resolved[field], + }, + errors, + ) } } +function validateCitationSourceField( + input: { + readonly ref: string + readonly field: CitationSourceField + readonly declaredValue: string | null | undefined + readonly resolvedValue: string | null | undefined + }, + errors: string[], +): void { + const declaredValue = normalizeSourceValue(input.declaredValue) + const resolvedValue = normalizeSourceValue(input.resolvedValue) + if (declaredValue === resolvedValue) return + + errors.push( + `Citation ref '${input.ref}' source.${input.field} must match resolved evidence source. Expected ${formatSourceValue( + resolvedValue, + )}, received ${formatSourceValue(declaredValue)}.`, + ) +} + +function normalizeSourceValue(value: string | null | undefined): string | null { + return value ?? null +} + +function formatSourceValue(value: string | null): string { + return value === null ? "missing" : `'${value}'` +} + function validateArtifactCounts( input: ManifestValidationInput, errors: string[], diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 313d3ff..4281520 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -99,7 +99,51 @@ describe("answerQuestionWithRetrieval", () => { }); expect(answer).toEqual({ answer: "The answer is grounded.", - citations: [result], + citations: [], + artifacts: [], + }); + }); + + it("does not create source chips from retrieval results when the manifest has no citations", async () => { + const unrelatedResult = makeRetrievalResult({ + content: "Information hiding is unrelated to the requested source.", + source: { + documentId: "doc_information_hiding", + sourceFileName: "information_hiding.pdf", + sectionPath: "Root", + }, + }); + const retrieval = { + query: vi.fn().mockResolvedValue({ + results: [unrelatedResult], + evidenceText: "Information hiding evidence.", + referencedChunks: [], + namespace: "notebook-workspace", + query: "requested fact", + routerUsed: "workflow_single_step", + answerText: null, + }), + }; + const generateAnswer = vi.fn(async ({ searchSources }) => { + await searchSources({ query: "requested fact" }); + return makeHarnessRunResult("The answer omits citations."); + }); + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "What does the selected source say?", + namespace: "notebook-workspace", + sources: [makeSource()], + excludedSourceIds: [], + retrieval, + generateAnswer, + messages: [], + }), + ); + + expect(answer).toEqual({ + answer: "The answer omits citations.", + citations: [], artifacts: [], }); }); @@ -304,7 +348,7 @@ describe("answerQuestionWithRetrieval", () => { ); expect(answer).toEqual({ answer: "The legacy answer is grounded.", - citations: [legacyResult], + citations: [], artifacts: [], }); }); @@ -562,8 +606,18 @@ describe("answerQuestionWithRetrieval", () => { }; const generateAnswer = vi.fn(async ({ searchSources }) => { await searchSources({ query: "What improved?" }); - return makeHarnessRunResult( + return makeHarnessRunResultWithLedger( "Revenue improved [Source 1: revenue growth]. Margins expanded [Source 2: margin expansion].", + { + citations: [ + makeOutputCitation("r1:result:1", firstResult), + makeOutputCitation("r1:result:2", secondResult), + ], + chunks: [ + makeEvidenceChunkFromRetrievalResult("r1:result:1", firstResult), + makeEvidenceChunkFromRetrievalResult("r1:result:2", secondResult), + ], + }, ); }); @@ -606,8 +660,9 @@ describe("answerQuestionWithRetrieval", () => { }; const generateAnswer = vi.fn(async ({ searchSources }) => { await searchSources({ query: "Tesla xAI investment" }); - return makeHarnessRunResult( + return makeCitedHarnessRunResult( "Tesla invested in xAI [Source 1: xAI investment].", + result, ); }); const sources = [ @@ -678,7 +733,10 @@ describe("answerQuestionWithRetrieval", () => { targetContent: "image", purpose: "Find visual rocket launch chunks.", }); - return makeHarnessRunResult(`Use this launch photo. ${upstreamAssetUrl}`); + return makeCitedHarnessRunResult( + `Use this launch photo. ${upstreamAssetUrl}`, + result, + ); }); const hardenChatAssetUrl = vi .fn() @@ -938,7 +996,10 @@ describe("answerQuestionWithRetrieval", () => { }; const generateAnswer = vi.fn(async ({ searchSources }) => { await searchSources({ query: "page four evidence" }); - return makeHarnessRunResult(`This page has the answer. ${storedPageAssetUrl}`); + return makeCitedHarnessRunResult( + `This page has the answer. ${storedPageAssetUrl}`, + result, + ); }); const hardenMediaAssetUrls = vi.fn( async ({ @@ -1022,6 +1083,35 @@ describe("answerQuestionWithRetrieval", () => { "https://blob.example/workspaces/workspace_1/sources/source_pages/parsed-result/page_citation_assets/page-6.png"; const hardenedPageAssetUrl = "https://blob.example/workspaces/workspace_1/chat-assets/source-source_pages/page-6.png"; + const pageMetadata = { + pageNums: [6], + pageAssets: [ + { + pageNum: 6, + artifactRef: "page_citation_assets/page-6.png", + assetUrl: rawPageAssetUrl, + contentType: "image/png", + }, + ], + }; + const referencedPageChunk: HarnessRunResult["trace"]["ledger"]["chunks"][number] = { + ref: "r1:referenced:1", + kind: "referenced_chunk", + chunkId: "chunk_page_6", + content: "", + contentPreview: "", + chunkType: "page", + score: null, + filePath: null, + metadata: pageMetadata, + source: { + documentId: "doc_pages", + sourceFileName: null, + sectionPath: "Page 6", + }, + revisionKey: "job_1", + assetUrl: rawPageAssetUrl, + }; const retrieval = { query: vi.fn().mockResolvedValue({ results: [], @@ -1035,17 +1125,7 @@ describe("answerQuestionWithRetrieval", () => { filePath: null, jobId: "job_1", assetUrl: rawPageAssetUrl, - metadata: { - pageNums: [6], - pageAssets: [ - { - pageNum: 6, - artifactRef: "page_citation_assets/page-6.png", - assetUrl: rawPageAssetUrl, - contentType: "image/png", - }, - ], - }, + metadata: pageMetadata, }, ], namespace: "notebook-workspace", @@ -1056,7 +1136,16 @@ describe("answerQuestionWithRetrieval", () => { }; const generateAnswer = vi.fn(async ({ searchSources }) => { await searchSources({ query: "page six evidence" }); - return makeHarnessRunResult("This page has referenced evidence."); + return makeHarnessRunResultWithLedger("This page has referenced evidence.", { + citations: [ + { + ref: "r1:referenced:1", + label: "Page 6", + source: referencedPageChunk.source, + }, + ], + chunks: [referencedPageChunk], + }); }); const hardenMediaAssetUrls = vi.fn( async ({ @@ -1403,6 +1492,103 @@ describe("answerQuestionWithRetrieval", () => { }); }); + it("returns a safe fallback when a manifest citation ref declares the wrong source", async () => { + process.env.AI_GATEWAY_API_KEY = "test_gateway_key"; + const result = makeRetrievalResult({ + content: "Information hiding is a module design principle.", + source: { + documentId: "doc_information_hiding", + sourceFileName: "information_hiding.pdf", + sectionPath: "Root / Module Design", + }, + }); + const retrieval = { + query: vi.fn().mockResolvedValue({ + results: [result], + evidenceText: "Information hiding evidence.", + referencedChunks: [], + namespace: "notebook-workspace", + query: "information hiding", + routerUsed: "workflow_single_step", + answerText: null, + }), + }; + let generateCallCount = 0; + vi.spyOn(ToolLoopAgent.prototype, "generate").mockImplementation( + async function mockGenerate( + this: ToolLoopAgent, + ): ReturnType { + generateCallCount += 1; + const tools = this.tools as unknown as Record< + string, + { execute: (input: unknown) => Promise } + >; + + if (generateCallCount === 1) { + await tools.declareIntent?.execute({ + task: "answer", + dependsOnPreviousTurn: false, + retrievalNeeded: "yes", + targetModalities: ["text"], + constraints: { citationRequired: true }, + groundingPolicy: "must_use_sources", + }); + await tools.setContextPolicy?.execute({ + carryHistory: "none", + reason: "Self-contained request.", + activePriorTurnIds: [], + }); + await tools.knowhere_search?.execute({ + query: "information hiding", + targetContent: "text", + }); + } + + await tools.finalize?.execute({ + text: "Information hiding is a module design principle.", + citations: [ + { + ref: "r1:result:1", + label: "claimed-source.pdf / Claimed", + source: { + documentId: "doc_claimed", + sourceFileName: "claimed-source.pdf", + sectionPath: "Claimed", + }, + }, + ], + artifacts: [], + unresolved: [], + }); + + return { + text: "ignored", + response: { messages: [] }, + } as unknown as Awaited>; + }, + ); + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "What is information hiding?", + namespace: "notebook-workspace", + sources: [makeSource()], + excludedSourceIds: [], + retrieval, + generateAnswer: generateAgenticOutputManifest, + messages: [], + }), + ); + + expect(generateCallCount).toBe(2); + expect(answer).toEqual({ + answer: + "I couldn't safely finish that response because the agent output did not pass Notebook's validation checks. Please try again.", + citations: [], + artifacts: [], + }); + }); + it("keeps image-only harness output instead of treating it as no results", async () => { const assetUrl = "https://blob.example/images/diagram.png"; const retrieval = { @@ -1866,6 +2052,23 @@ describe("answerQuestionWithRetrieval", () => { }); it("uses structured referenced chunks from RetrievalQueryResponse as citations", async () => { + const referencedImageChunk: HarnessRunResult["trace"]["ledger"]["chunks"][number] = { + ref: "r1:referenced:1", + kind: "referenced_chunk", + chunkId: "chunk_1", + content: "", + contentPreview: "", + chunkType: "image", + score: null, + filePath: "images/launch.jpg", + source: { + documentId: "doc_spacex", + sourceFileName: null, + sectionPath: "Assets / images / launch.jpg", + }, + revisionKey: "job_1", + assetUrl: "https://blob.example/images/launch.jpg", + }; const retrieval = { query: vi.fn().mockResolvedValue({ results: [], @@ -1892,7 +2095,16 @@ describe("answerQuestionWithRetrieval", () => { query: "SpaceX launch image", targetContent: "image", }); - return makeHarnessRunResult("Here is the launch image."); + return makeHarnessRunResultWithLedger("Here is the launch image.", { + citations: [ + { + ref: "r1:referenced:1", + label: "launch image", + source: referencedImageChunk.source, + }, + ], + chunks: [referencedImageChunk], + }); }); const answer = await Effect.runPromise( @@ -1972,7 +2184,7 @@ describe("generateAgenticOutputManifest", () => { label: "商务标文件.pdf / 身份证正面", source: { documentId: "doc_identity", - sourceFileName: "商务标文件.pdf", + sourceFileName: "document-generated.pdf", sectionPath: "身份证正面", }, }, @@ -2109,7 +2321,7 @@ describe("generateAgenticOutputManifest", () => { label: "identity.pdf / images/id-front.png", source: { documentId: "doc_identity", - sourceFileName: "identity.pdf", + sourceFileName: "generated.pdf", sectionPath: "images/id-front.png", }, }, @@ -2249,7 +2461,7 @@ describe("generateAgenticOutputManifest", () => { label: "投标书 / （6）现场工期进度管理方面的违约责任", source: { documentId: "doc_contract", - sourceFileName: "投标书.pdf", + sourceFileName: null, sectionPath: "Root / （6）现场工期进度管理方面的违约责任", }, }, @@ -2393,7 +2605,17 @@ describe("generateAgenticOutputManifest", () => { }); await tools.finalize?.execute({ text: "见下方图片。", - citations: [{ ref: "r1:result:1", label: "id" }], + citations: [ + { + ref: "r1:result:1", + label: "ids.pdf / 身份证 1", + source: { + documentId: "doc_identity", + sourceFileName: "ids.pdf", + sectionPath: "身份证 1", + }, + }, + ], artifacts: [1, 2, 3].map((index) => ({ type: "image", ref: `asset:r1:result:${index}`, @@ -2405,7 +2627,17 @@ describe("generateAgenticOutputManifest", () => { } else { await tools.finalize?.execute({ text: "见下方图片。", - citations: [{ ref: "r1:result:1", label: "id" }], + citations: [ + { + ref: "r1:result:1", + label: "ids.pdf / 身份证 1", + source: { + documentId: "doc_identity", + sourceFileName: "ids.pdf", + sectionPath: "身份证 1", + }, + }, + ], artifacts: [1, 2].map((index) => ({ type: "image", ref: `asset:r1:result:${index}`, @@ -2558,6 +2790,90 @@ function makeHarnessRunResult(text: string): HarnessRunResult { }; } +function makeCitedHarnessRunResult( + text: string, + result: RetrievalResult, + ref = "r1:result:1", +): HarnessRunResult { + return makeHarnessRunResultWithLedger(text, { + citations: [makeOutputCitation(ref, result)], + chunks: [makeEvidenceChunkFromRetrievalResult(ref, result)], + }); +} + +function makeHarnessRunResultWithLedger( + text: string, + input: { + readonly citations?: HarnessRunResult["manifest"]["citations"] + readonly chunks?: HarnessRunResult["trace"]["ledger"]["chunks"] + readonly assets?: HarnessRunResult["trace"]["ledger"]["assets"] + readonly artifacts?: HarnessRunResult["manifest"]["artifacts"] + }, +): HarnessRunResult { + const chunks = input.chunks ?? []; + return { + manifest: { + text, + citations: input.citations ?? [], + artifacts: input.artifacts ?? [], + unresolved: [], + }, + trace: { + ...makeHarnessRunResult("").trace, + ledger: { + retrievalCount: chunks.length > 0 ? 1 : 0, + chunks, + assets: input.assets ?? [], + evidenceText: [], + stopReasons: [], + failureReasons: [], + decisionTraces: [], + }, + }, + }; +} + +function makeOutputCitation( + ref: string, + result: RetrievalResult, +): HarnessRunResult["manifest"]["citations"][number] { + return { + ref, + label: [result.source.sourceFileName, result.source.sectionPath] + .filter(Boolean) + .join(" / "), + source: { + documentId: result.source.documentId, + sourceFileName: result.source.sourceFileName, + sectionPath: result.source.sectionPath, + }, + }; +} + +function makeEvidenceChunkFromRetrievalResult( + ref: string, + result: RetrievalResult, +): HarnessRunResult["trace"]["ledger"]["chunks"][number] { + return { + ref, + kind: "result", + ...(result.chunkId ? { chunkId: result.chunkId } : {}), + content: result.content, + contentPreview: result.content, + chunkType: result.chunkType, + score: result.score, + ...(result.sourceChunkPath ? { sourceChunkPath: result.sourceChunkPath } : {}), + ...(result.filePath ? { filePath: result.filePath } : {}), + ...(result.metadata ? { metadata: result.metadata } : {}), + source: { + documentId: result.source.documentId, + sourceFileName: result.source.sourceFileName, + sectionPath: result.source.sectionPath, + }, + ...(result.assetUrl ? { assetUrl: result.assetUrl } : {}), + }; +} + function makeKnowledgeOutline(): KnowledgeOutline { return { document: makeLocalKnowledgeDocument(), diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index c0dc51b..d5cd82c 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -254,8 +254,6 @@ export const answerQuestionWithRetrieval = ( const rawResults = selectCitationRawResults({ generatedAnswer, - retrievalResponses, - sources: input.sources, }) if ( rawResults.length === 0 && @@ -858,14 +856,12 @@ function normalizeTopK(value: number | undefined): number { /** * Display citations come from the agent-curated manifest (the refs it chose to - * cite), resolved against the evidence ledger. Only when the agent cited - * nothing do we fall back to the full set of retrieved results, so a grounded - * answer still shows its sources instead of appearing unsupported. + * cite), resolved against the evidence ledger. If the manifest has no citations, + * only displayed artifacts produce citation chips; arbitrary retrieval results + * stay hidden so Notebook does not imply support from an unselected source. */ function selectCitationRawResults(input: { readonly generatedAnswer: HarnessRunResult - readonly retrievalResponses: readonly RetrievalQueryResponse[] - readonly sources: readonly AnswerQuestionInput["sources"][number][] }): RetrievalResult[] { const curated = mapManifestCitationsToResults(input.generatedAnswer) if (curated.length > 0) return curated @@ -873,7 +869,7 @@ function selectCitationRawResults(input: { input.generatedAnswer, ) if (displayedArtifacts.length > 0) return displayedArtifacts - return collectRetrievalResults(input.retrievalResponses, input.sources) + return [] } function mapManifestCitationsToResults( @@ -1008,49 +1004,6 @@ function hasDisplayedManifestArtifacts(result: HarnessRunResult): boolean { return result.manifest.artifacts.some((artifact) => artifact.display) } -function collectRetrievalResults( - responses: readonly RetrievalQueryResponse[], - sources: readonly AnswerQuestionInput["sources"][number][], -): RetrievalResult[] { - const results: RetrievalResult[] = [] - const seenKeys = new Set() - const sourceTitlesByDocumentId = new Map( - sources.flatMap((source): readonly [string, string][] => - source.knowhereDocumentId ? [[source.knowhereDocumentId, source.title]] : [], - ), - ) - - for (const response of responses) { - for (const result of [ - ...response.results, - ...response.referencedChunks.map((chunk): RetrievalResult => ({ - chunkId: chunk.chunkId, - content: "", - chunkType: chunk.chunkType, - score: null, - ...(chunk.assetUrl ? { assetUrl: chunk.assetUrl } : {}), - ...(chunk.sourceChunkPath ? { sourceChunkPath: chunk.sourceChunkPath } : {}), - ...(chunk.filePath ? { filePath: chunk.filePath } : {}), - ...(chunk.metadata ? { metadata: chunk.metadata } : {}), - source: { - documentId: chunk.documentId, - sourceFileName: sourceTitlesByDocumentId.get(chunk.documentId), - sectionPath: chunk.sectionPath, - }, - })), - ]) { - const key = getRetrievalResultKey(result) - if (seenKeys.has(key)) continue - - seenKeys.add(key) - results.push(result) - if (results.length >= MAX_CITATION_RESULTS) return results - } - } - - return results -} - function formatRetrievalEvidenceText( responses: readonly RetrievalQueryResponse[], ): string | undefined { diff --git a/src/domains/chat/service.test.ts b/src/domains/chat/service.test.ts index 441a12c..3a1ab82 100644 --- a/src/domains/chat/service.test.ts +++ b/src/domains/chat/service.test.ts @@ -48,7 +48,7 @@ describe("handleChatTurn", () => { { role: "assistant", content: "Grounded answer.", - citations: [makeRetrievalResult()], + citations: undefined, }, ], }); @@ -78,7 +78,7 @@ describe("handleChatTurn", () => { threadId: "thread_1", role: "assistant", content: "Grounded answer.", - citations: [makeRetrievalResult()], + citations: [], artifacts: [], }); }); From e4bc2b6da96a47ed7a6ac47d4149417b8e7a0c29 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 8 Jul 2026 04:53:31 +0800 Subject: [PATCH 34/51] Trust agent chat output --- src/agent-harness/index.ts | 1 - src/agent-harness/runtime.test.ts | 63 +-- src/agent-harness/runtime.ts | 150 ++------ src/agent-harness/types.ts | 4 +- src/agent-harness/validator.test.ts | 570 ---------------------------- src/agent-harness/validator.ts | 363 ------------------ src/domains/chat/index.test.ts | 32 +- src/domains/chat/index.ts | 22 +- src/domains/chat/prompt.ts | 2 +- 9 files changed, 52 insertions(+), 1155 deletions(-) delete mode 100644 src/agent-harness/validator.test.ts delete mode 100644 src/agent-harness/validator.ts diff --git a/src/agent-harness/index.ts b/src/agent-harness/index.ts index a1feff1..5fbdf7a 100644 --- a/src/agent-harness/index.ts +++ b/src/agent-harness/index.ts @@ -2,4 +2,3 @@ export * from "./ledger" export * from "./knowhere-text" export * from "./runtime" export * from "./types" -export * from "./validator" diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index eec6064..04614a7 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -28,24 +28,23 @@ describe("agent harness runtime", () => { expect(prompt).not.toContain("navigation action") }) - it("tells the agent to cite only refs with matching source metadata", () => { + it("tells the agent citation metadata is optional and ledger-resolved", () => { const prompt = buildHarnessSystemPrompt(makeTurnInput()) expect(prompt).toContain( - "source.documentId, sourceFileName, and sectionPath must match the selected evidence ref exactly", + "Citation label and source metadata are optional", ) expect(prompt).toContain( - "Omit citations when you cannot identify a supporting evidence ref", + "Notebook resolves citation metadata from evidence refs when possible", ) + expect(prompt).not.toContain("must match the selected evidence ref exactly") }) - it("passes only outer retrieval parameters to KNOWHERE after intent and context policy are declared", async () => { + it("passes only outer retrieval parameters to KNOWHERE without planning-tool gating", async () => { const query = vi.fn().mockResolvedValue( makeRetrievalResponse(), ) const state: { - intent?: IntentFrame - contextPolicy?: ContextPolicy toolCalls?: HarnessToolCallTrace[] } = {} const tools = createHarnessTools({ @@ -55,25 +54,6 @@ describe("agent harness runtime", () => { recentTurns: [], }) - expect(await executeTool(tools.knowhere_search, { query: "q4 chart" })) - .toContain('status="error"') - - await executeTool(tools.declareIntent, { - task: "show_media", - dependsOnPreviousTurn: false, - retrievalNeeded: "yes", - targetModalities: ["image"], - constraints: { desiredCount: 2, maxCount: 2 }, - groundingPolicy: "must_use_sources", - }) - expect(await executeTool(tools.knowhere_search, { query: "q4 chart" })) - .toContain("setContextPolicy must be called before knowhere_search.") - - await executeTool(tools.setContextPolicy, { - carryHistory: "none", - reason: "The current request is unrelated to previous turns.", - activePriorTurnIds: [], - }) const result = await executeTool(tools.knowhere_search, { query: "q4 chart", targetContent: "image", @@ -97,10 +77,6 @@ describe("agent harness runtime", () => { "LegalAction", ) expect(state.toolCalls?.map((call) => [call.tool, call.ok])).toEqual([ - ["knowhere_search", false], - ["declareIntent", true], - ["knowhere_search", false], - ["setContextPolicy", true], ["knowhere_search", true], ]) }) @@ -360,10 +336,8 @@ describe("agent harness runtime", () => { }) }) - it("blocks finalize until intent and context policy are declared", async () => { + it("accepts finalize output without planning-tool gating", async () => { const state: { - intent?: IntentFrame - contextPolicy?: ContextPolicy finalizedManifest?: OutputManifest finalized?: boolean } = {} @@ -381,31 +355,6 @@ describe("agent harness runtime", () => { unresolved: [], } - expect(await executeTool(tools.finalize, manifest)).toEqual({ - ok: false, - message: "declareIntent must be called before finalize.", - }) - expect(state.finalizedManifest).toBeUndefined() - - await executeTool(tools.declareIntent, { - task: "answer_question", - dependsOnPreviousTurn: false, - retrievalNeeded: "no", - targetModalities: ["text"], - constraints: {}, - groundingPolicy: "may_use_sources", - }) - expect(await executeTool(tools.finalize, manifest)).toEqual({ - ok: false, - message: "setContextPolicy must be called before finalize.", - }) - expect(state.finalizedManifest).toBeUndefined() - - await executeTool(tools.setContextPolicy, { - carryHistory: "none", - reason: "Self-contained request.", - activePriorTurnIds: [], - }) expect(await executeTool(tools.finalize, manifest)).toMatchObject({ ok: true, text: "Answer.", diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index e272340..8963fdf 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -25,10 +25,8 @@ import type { KnowhereToolRuntime, OutputManifest, } from "./types" -import { validateOutputManifest } from "./validator" const defaultMaxSteps = 14 -const defaultMaxRevisions = 1 const imageInspectionReminderStepNumber = 12 const forcedFinalizationStepNumber = 13 const imageInspectionRefLimit = 6 @@ -43,11 +41,6 @@ export type RunAgentHarnessInput = { readonly knowhereTools: KnowhereToolRuntime readonly inspectImages?: InspectImages readonly maxSteps?: number - /** - * How many times the agent may revise after a failed validation pass before - * the harness gives up and returns the last manifest with recorded errors. - */ - readonly maxRevisions?: number } type HarnessToolState = { @@ -164,12 +157,14 @@ const contextPolicySchema = z.object({ const outputCitationSchema = z.object({ ref: z.string().min(1), - label: z.string().min(1), - source: z.object({ - documentId: z.string().nullable().optional(), - sourceFileName: z.string().nullable().optional(), - sectionPath: z.string().nullable().optional(), - }), + label: z.string().min(1).optional(), + source: z + .object({ + documentId: z.string().nullable().optional(), + sourceFileName: z.string().nullable().optional(), + sectionPath: z.string().nullable().optional(), + }) + .optional(), }) const selectedOutputArtifactSchema = z.object({ @@ -236,46 +231,11 @@ export async function runAgentHarness( ], }) - const maxRevisions = input.maxRevisions ?? defaultMaxRevisions - let messages = buildHarnessMessages(input.turn) - let manifest = buildFallbackManifest("") - let validationErrors: readonly string[] = [] - let revisionsUsed = 0 - - for (let attempt = 0; ; attempt += 1) { - const response = await agent.generate({ messages }) - manifest = - state.finalizedManifest ?? buildFallbackManifest(response.text.trim()) - - const validation = validateOutputManifest({ - manifest, - intent: state.intent, - contextPolicy: state.contextPolicy, - finalized: state.finalized === true, - ledger: ledger.snapshot(), - toolCalls: state.toolCalls, - surface: input.turn.surface, - }) - validationErrors = validation.errors - - if (validation.ok || attempt >= maxRevisions) break - - // Self-correction (reflexion): continue the same conversation with - // structured validator feedback and require a fresh finalize so the agent - // can repair its own contract violations instead of shipping them. - revisionsUsed += 1 - state.finalizedManifest = undefined - state.finalized = false - messages = [ - ...messages, - ...(response.response.messages as ModelMessage[]), - { - role: "user", - content: buildRevisionFeedback(validation.errors), - }, - ] - } - + const response = await agent.generate({ + messages: buildHarnessMessages(input.turn), + }) + const manifest = + state.finalizedManifest ?? buildFallbackManifest(response.text.trim()) const ledgerSnapshot = ledger.snapshot() return { manifest, @@ -286,8 +246,8 @@ export async function runAgentHarness( finalized: state.finalized === true, priorTurnReads: [...(state.priorTurnReads ?? [])], toolCalls: [...(state.toolCalls ?? [])], - validationErrors, - revisionsUsed, + validationErrors: [], + revisionsUsed: 0, }, } } @@ -395,22 +355,6 @@ type ModelMessageForRole = Extract< { readonly role: TRole } > -function buildRevisionFeedback(errors: readonly string[]): string { - return [ - "Your finalize output did not satisfy the output contract:", - ...errors.map((error) => `- ${error}`), - "", - "Fix every issue and call finalize again with a corrected manifest.", - "You must call finalize; freeform assistant text is not a valid final", - "answer contract.", - "Do not exceed the user's requested artifact count, only cite or display", - "evidence refs that exist in the evidence ledger, and do not fabricate", - "facts when evidence is missing. Citation source.documentId,", - "sourceFileName, and sectionPath must match the selected evidence ref", - "exactly.", - ].join("\n") -} - function buildForcedFinalizationFeedback(): string { return [ "The retrieval step budget has been reached.", @@ -453,7 +397,7 @@ export function createHarnessTools(input: { return { declareIntent: tool({ description: - "Declare the user's intent before any other action. This is working memory, not a final answer.", + "Declare the user's intent when it helps plan the response. This is working memory, not a final answer.", inputSchema: intentFrameSchema, execute: async (intent): Promise => traceToolCall(input.state, { @@ -493,7 +437,6 @@ export function createHarnessTools(input: { inputSummary: summarizeKnowhereSearchRequest(request), execute: async () => executeKnowhereSearch({ - state: input.state, ledger: input.ledger, knowhereTools: input.knowhereTools, request, @@ -513,7 +456,6 @@ export function createHarnessTools(input: { execute: async () => executeKnowhereTextTool({ operation: "list_documents", - state: input.state, execute: async () => knowhereToolText.formatListDocuments( await input.knowhereTools.listDocuments(), @@ -534,7 +476,6 @@ export function createHarnessTools(input: { execute: async () => executeKnowhereTextTool({ operation: "get_document_outline", - state: input.state, validate: () => validateDocumentReference(request), execute: async () => knowhereToolText.formatOutline( @@ -555,7 +496,6 @@ export function createHarnessTools(input: { inputSummary: summarizeReadChunksRequest(request), execute: async () => executeKnowhereReadChunks({ - state: input.state, ledger: input.ledger, knowhereTools: input.knowhereTools, request, @@ -574,7 +514,6 @@ export function createHarnessTools(input: { inputSummary: summarizeGrepChunksRequest(request), execute: async () => executeKnowhereGrepChunks({ - state: input.state, ledger: input.ledger, knowhereTools: input.knowhereTools, request, @@ -669,26 +608,13 @@ export function createHarnessTools(input: { description: "Finalize the user-facing output manifest. This is the only final answer " + "contract. Artifacts listed here with display=true are the exact set of " + - "images/tables shown to the user; cite only refs from the evidence ledger " + - "and copy citation source metadata exactly from the selected ref.", + "images/tables shown to the user; cite evidence refs when available.", inputSchema: outputManifestSchema, execute: async (manifest) => traceToolCall(input.state, { toolName: "finalize", inputSummary: summarizeManifest(manifest), execute: async () => { - if (!input.state.intent) { - return { - ok: false as const, - message: "declareIntent must be called before finalize.", - } - } - if (!input.state.contextPolicy) { - return { - ok: false as const, - message: "setContextPolicy must be called before finalize.", - } - } input.state.finalizedManifest = manifest input.state.finalized = true return { ok: true as const, ...manifest } @@ -878,14 +804,12 @@ type DocumentReferenceSummary = { } async function executeKnowhereSearch(input: { - readonly state: HarnessToolState readonly ledger: ReturnType readonly knowhereTools: KnowhereToolRuntime readonly request: KnowhereSearchToolRequest }): Promise { return executeKnowhereTextTool({ operation: "search", - state: input.state, execute: async () => { const beforeSnapshot = input.ledger.snapshot() const response = await input.knowhereTools.search({ @@ -909,14 +833,12 @@ async function executeKnowhereSearch(input: { } async function executeKnowhereReadChunks(input: { - readonly state: HarnessToolState readonly ledger: ReturnType readonly knowhereTools: KnowhereToolRuntime readonly request: KnowhereReadChunksToolRequest }): Promise { return executeKnowhereTextTool({ operation: "read_chunks", - state: input.state, validate: () => validateDocumentReference(input.request), execute: async () => { const beforeSnapshot = input.ledger.snapshot() @@ -932,14 +854,12 @@ async function executeKnowhereReadChunks(input: { } async function executeKnowhereGrepChunks(input: { - readonly state: HarnessToolState readonly ledger: ReturnType readonly knowhereTools: KnowhereToolRuntime readonly request: KnowhereGrepChunksToolRequest }): Promise { return executeKnowhereTextTool({ operation: "grep_chunks", - state: input.state, validate: () => validateDocumentReference(input.request), execute: async () => { const beforeSnapshot = input.ledger.snapshot() @@ -956,18 +876,9 @@ async function executeKnowhereGrepChunks(input: { async function executeKnowhereTextTool(input: { readonly operation: KnowhereToolOperation - readonly state: HarnessToolState readonly validate?: () => string | null readonly execute: () => Promise }): Promise { - const workflowError = validateKnowhereWorkflow(input.state, input.operation) - if (workflowError) { - return knowhereToolText.formatError({ - operation: input.operation, - message: workflowError, - }) - } - const validationError = input.validate?.() if (validationError) { return knowhereToolText.formatError({ @@ -986,18 +897,6 @@ async function executeKnowhereTextTool(input: { } } -function validateKnowhereWorkflow( - state: HarnessToolState, - operation: KnowhereToolOperation, -): string | null { - const toolName = `knowhere_${operation}` - if (!state.intent) return `declareIntent must be called before ${toolName}.` - if (!state.contextPolicy) { - return `setContextPolicy must be called before ${toolName}.` - } - return null -} - function validateDocumentReference( request: KnowhereDocumentReferenceRequest, ): string | null { @@ -1257,15 +1156,15 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "KNOWHERE is only an evidence provider. Do not infer or control its internal navigation algorithm.", "Your job is to understand intent, decide context use, optionally retrieve evidence, select evidence/artifacts, create source-backed derived tables when useful, and finalize an output manifest.", "", - "Required workflow:", - "1. Call declareIntent first. Capture constraints like a requested image/table count in constraints.desiredCount.", - "2. Call setContextPolicy next, deciding how prior turns should influence this turn.", + "Recommended workflow:", + "1. Call declareIntent when it helps you plan the response. Capture constraints like a requested image/table count in constraints.desiredCount.", + "2. Call setContextPolicy when prior turns may influence this turn.", "3. When the policy needs prior-turn detail (references or corrections), call readPriorTurn for the relevant ids.", "4. Call knowhere_search when relevance search is needed. Use knowhere_list_documents, knowhere_get_document_outline, knowhere_read_chunks, and knowhere_grep_chunks for focused document reads.", "5. For pixel-level details, OCR, visual comparison, image verification, or when the likely answer is only visible on a returned page/image asset, call inspectImage only after a Knowhere tool returned image asset refs.", `6. inspectImage accepts at most ${imageInspectionRefLimit} image asset refs per call and per turn.`, "7. knowhere_read_chunks returns complete chunk bodies; control size with page/pageSize, sectionPath, startChunk/endChunk, chunkId, and chunkType.", - "8. Call finalize with text, citations, artifacts, and unresolved issues. finalize requires declareIntent and setContextPolicy first.", + "8. Call finalize with text, citations, artifacts, and unresolved issues when you are ready to answer.", "", "Context rules:", "- If the current user request is unrelated to prior turns, set carryHistory to none and do not reuse prior topics.", @@ -1276,13 +1175,12 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "- Final output is the OutputManifest passed to finalize, not freeform tool JSON or trailing text.", "- artifacts with display=true are the exact images/tables shown. Never display every candidate; honor constraints.desiredCount / maxCount.", "- Use type=derived_table only for tables you create from evidence; every derived_table.sourceRefs entry must reference evidence in the ledger.", - "- citations and selected image/table artifact refs may only reference refs returned by Knowhere tools in the evidence ledger.", - "- For every citation, source.documentId, sourceFileName, and sectionPath must match the selected evidence ref exactly. Use null or omit the field only when the ref omits it.", - "- Omit citations when you cannot identify a supporting evidence ref with matching source metadata.", + "- Prefer citation and selected image/table artifact refs returned by Knowhere tools in the evidence ledger.", + "- Citation label and source metadata are optional. Notebook resolves citation metadata from evidence refs when possible.", + "- If evidence is relevant but you cannot identify a supporting evidence ref, answer with unresolved issues instead of fabricating a ref.", "- inspectImage observations are inspection notes, not new source refs. Final citations and displayed image artifacts must use the original retrieved image asset refs.", "- If text evidence identifies a relevant page/image but does not include the exact fact, inspect the returned image asset for OCR/detail before saying the answer is unavailable.", "- If evidence is insufficient, list it in unresolved instead of fabricating facts.", - "- After a validation-feedback message, fix all listed issues and call finalize again.", `Surface: ${turn.surface}`, `Output capabilities: ${JSON.stringify(turn.outputCapabilities)}`, turn.sourceContext ? `Searchable source context:\n${turn.sourceContext}` : "", diff --git a/src/agent-harness/types.ts b/src/agent-harness/types.ts index 83bb08d..36b6a27 100644 --- a/src/agent-harness/types.ts +++ b/src/agent-harness/types.ts @@ -209,8 +209,8 @@ export type EvidenceLedgerSnapshot = { export type OutputCitation = { readonly ref: string - readonly label: string - readonly source: EvidenceChunk["source"] + readonly label?: string + readonly source?: EvidenceChunk["source"] } export type OutputArtifact = { diff --git a/src/agent-harness/validator.test.ts b/src/agent-harness/validator.test.ts deleted file mode 100644 index 89bf0ba..0000000 --- a/src/agent-harness/validator.test.ts +++ /dev/null @@ -1,570 +0,0 @@ -import { describe, expect, it } from "vitest" - -import { validateOutputManifest } from "./validator" -import type { - ContextPolicy, - EvidenceLedgerSnapshot, - IntentFrame, - OutputManifest, -} from "./types" - -describe("validateOutputManifest", () => { - it("requires finalize to be the successful output path", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ text: "Freeform answer." }), - intent: makeIntent({ groundingPolicy: "no_retrieval" }), - contextPolicy: unrelatedContextPolicy, - finalized: false, - ledger: emptyLedger, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Agent must call finalize to produce the output manifest.", - ) - }) - - it("requires the agent to declare intent and context policy before finalizing", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ text: "Answer." }), - ledger: emptyLedger, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Agent must declare intent before finalizing.", - ) - expect(validation.errors).toContain( - "Agent must set context policy before finalizing.", - ) - }) - - it("limits displayed artifacts using the declared intent instead of hard-coded media rules", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - artifacts: [ - { - type: "image", - ref: "asset:r1:result:1", - display: true, - reason: "front", - }, - { - type: "image", - ref: "asset:r1:result:2", - display: true, - reason: "back", - }, - { - type: "image", - ref: "asset:r1:result:3", - display: true, - reason: "extra candidate", - }, - ], - }), - intent: makeIntent({ desiredCount: 2, maxCount: 2 }), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - assets: [ - makeAsset("asset:r1:result:1"), - makeAsset("asset:r1:result:2"), - makeAsset("asset:r1:result:3"), - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Displayed artifact count 3 exceeds desired count 2.", - ) - expect(validation.errors).toContain( - "Displayed artifact count 3 exceeds maximum count 2.", - ) - }) - - it("rejects grounded answers that use evidence without citations or selected artifacts", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ text: "Revenue increased." }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:result:1", - kind: "result", - content: "Revenue increased.", - contentPreview: "Revenue increased.", - chunkType: "text", - score: 0.9, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "Q4", - }, - }, - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Grounded output used evidence but did not cite or display any selected evidence.", - ) - }) - - it("rejects citation source metadata that conflicts with the resolved evidence ref", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "The source discusses information hiding.", - citations: [ - { - ref: "r1:result:1", - label: "wrong source", - source: { - documentId: "doc_claimed", - sourceFileName: "claimed-source.pdf", - sectionPath: "Claimed Section", - }, - }, - ], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:result:1", - kind: "result", - content: "Information hiding is a module design principle.", - contentPreview: "Information hiding is a module design principle.", - chunkType: "text", - score: 0.9, - source: { - documentId: "doc_information_hiding", - sourceFileName: "information_hiding.pdf", - sectionPath: "Root / Module Design", - }, - }, - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Citation ref 'r1:result:1' source.documentId must match resolved evidence source. Expected 'doc_information_hiding', received 'doc_claimed'.", - ) - expect(validation.errors).toContain( - "Citation ref 'r1:result:1' source.sourceFileName must match resolved evidence source. Expected 'information_hiding.pdf', received 'claimed-source.pdf'.", - ) - expect(validation.errors).toContain( - "Citation ref 'r1:result:1' source.sectionPath must match resolved evidence source. Expected 'Root / Module Design', received 'Claimed Section'.", - ) - }) - - it("accepts citation source metadata that exactly matches the resolved evidence ref", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "Revenue increased.", - citations: [ - { - ref: "r1:result:1", - label: "report.pdf / Q4", - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "Q4", - }, - }, - ], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:result:1", - kind: "result", - content: "Revenue increased.", - contentPreview: "Revenue increased.", - chunkType: "text", - score: 0.9, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "Q4", - }, - }, - ], - }, - surface: "notebook_chat", - }) - - expect(validation.ok).toBe(true) - expect(validation.errors).toEqual([]) - }) - - it("rejects missing citation documentId when the resolved evidence has one", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "Revenue increased.", - citations: [ - { - ref: "r1:result:1", - label: "report.pdf / Q4", - source: { - sourceFileName: "report.pdf", - sectionPath: "Q4", - }, - }, - ], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:result:1", - kind: "result", - content: "Revenue increased.", - contentPreview: "Revenue increased.", - chunkType: "text", - score: 0.9, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "Q4", - }, - }, - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Citation ref 'r1:result:1' source.documentId must match resolved evidence source. Expected 'doc_1', received missing.", - ) - }) - - it("accepts source-backed derived tables and rejects missing source refs", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - artifacts: [ - { - type: "derived_table", - ref: "derived:table:1", - title: "Revenue comparison", - columns: ["Metric", "Value"], - rows: [["Revenue", "$10M"]], - sourceRefs: ["r1:result:1"], - display: true, - reason: "Structured comparison requested by the user.", - }, - { - type: "derived_table", - ref: "derived:table:2", - title: "Invalid table", - columns: ["Metric", "Value"], - rows: [["Revenue"]], - sourceRefs: ["missing"], - display: true, - reason: "Demonstrates validation.", - }, - ], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:result:1", - kind: "result", - content: "Revenue was $10M.", - contentPreview: "Revenue was $10M.", - chunkType: "text", - score: 0.9, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "Revenue", - }, - }, - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Derived table source ref 'missing' was not found in the evidence ledger.", - ) - expect(validation.errors).toContain( - "Derived table row 1 has 1 cells but expected 2.", - ) - }) - - it("requires compare outputs to cite at least two evidence refs", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "A is stronger than B.", - citations: [ - { - ref: "r1:result:1", - label: "report.pdf / A", - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "A", - }, - }, - ], - }), - intent: { - ...makeIntent({}), - task: "compare", - }, - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:result:1", - kind: "result", - content: "A is strong.", - contentPreview: "A is strong.", - chunkType: "text", - score: 0.9, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "A", - }, - }, - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Compare outputs that must use sources require at least two evidence refs or an explicit unresolved reason.", - ) - }) - - it("keeps typing compose output insertion-ready", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ text: "- bullet\n- list" }), - intent: makeIntent({ groundingPolicy: "no_retrieval" }), - contextPolicy: unrelatedContextPolicy, - ledger: emptyLedger, - surface: "typing_compose", - }) - - expect(validation.errors).toContain( - "Typing compose output must be insertion-ready plain text.", - ) - }) - - it("rejects image inspection failure claims when inspectImage was not called", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "图像检查工具未能成功读取该页面的具体条款内容。", - citations: [ - { - ref: "r1:referenced:1", - label: "contract.pdf / page 8", - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - }, - ], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:referenced:1", - kind: "referenced_chunk", - content: "", - contentPreview: "", - chunkType: "page", - score: null, - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - assetRef: "asset:r1:referenced:1", - }, - ], - assets: [makeAsset("asset:r1:referenced:1")], - }, - toolCalls: [], - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", - ) - }) - - it("allows image inspection failure claims after inspectImage was called", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "对该页面的图像识别未成功提取具体数值。", - citations: [ - { - ref: "r1:referenced:1", - label: "contract.pdf / page 8", - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - }, - ], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:referenced:1", - kind: "referenced_chunk", - content: "", - contentPreview: "", - chunkType: "page", - score: null, - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - assetRef: "asset:r1:referenced:1", - }, - ], - assets: [makeAsset("asset:r1:referenced:1")], - }, - toolCalls: [ - { - tool: "inspectImage", - ok: false, - inputSummary: { refs: ["asset:r1:referenced:1"] }, - outputSummary: { inspectedCount: 0 }, - startedAt: "2026-07-06T00:00:00.000Z", - durationMs: 1, - }, - ], - surface: "notebook_chat", - }) - - expect(validation.errors).not.toContain( - "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", - ) - }) - - it("rejects unreadable page claims when image assets are available but uninspected", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "相关的页面图像无法通过目前的检测工具获取。", - unresolved: ["由于无法直接读取该页面的详细条款内容，目前无法给出确切的赔偿金额。"], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:referenced:1", - kind: "referenced_chunk", - content: "", - contentPreview: "", - chunkType: "page", - score: null, - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - assetRef: "asset:r1:referenced:1", - }, - ], - assets: [makeAsset("asset:r1:referenced:1")], - }, - toolCalls: [], - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Final output must inspect available image assets before claiming retrieved page/image content cannot be read.", - ) - expect(validation.errors).toContain( - "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", - ) - }) -}) - -const emptyLedger: EvidenceLedgerSnapshot = { - retrievalCount: 0, - chunks: [], - assets: [], - evidenceText: [], - stopReasons: [], - failureReasons: [], - decisionTraces: [], -} - -const unrelatedContextPolicy: ContextPolicy = { - carryHistory: "none", - reason: "The current turn is self-contained.", - activePriorTurnIds: [], -} - -function makeIntent( - constraints: IntentFrame["constraints"] & { - readonly groundingPolicy?: IntentFrame["groundingPolicy"] - }, -): IntentFrame { - return { - task: "answer", - dependsOnPreviousTurn: false, - retrievalNeeded: constraints.groundingPolicy === "no_retrieval" ? "no" : "yes", - targetModalities: ["text"], - constraints, - groundingPolicy: constraints.groundingPolicy ?? "must_use_sources", - } -} - -function makeManifest(overrides: Partial): OutputManifest { - return { - text: "", - citations: [], - artifacts: [], - unresolved: [], - ...overrides, - } -} - -function makeAsset(ref: string): EvidenceLedgerSnapshot["assets"][number] { - return { - ref, - chunkRef: ref.replace("asset:", ""), - type: "image", - assetUrl: `https://assets.example/${ref}.png`, - label: ref, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: ref, - }, - } -} diff --git a/src/agent-harness/validator.ts b/src/agent-harness/validator.ts deleted file mode 100644 index fcd8c0e..0000000 --- a/src/agent-harness/validator.ts +++ /dev/null @@ -1,363 +0,0 @@ -import type { - ContextPolicy, - EvidenceLedgerSnapshot, - HarnessToolCallTrace, - IntentFrame, - OutputCitation, - OutputManifest, -} from "./types" - -export type ManifestValidationInput = { - readonly manifest: OutputManifest - readonly intent?: IntentFrame - readonly contextPolicy?: ContextPolicy - readonly finalized?: boolean - readonly ledger: EvidenceLedgerSnapshot - readonly toolCalls?: readonly HarnessToolCallTrace[] - readonly surface: "notebook_chat" | "typing_compose" | "typing_quick_ask" -} - -export type ManifestValidationResult = { - readonly ok: boolean - readonly errors: readonly string[] -} - -const citationSourceFields = [ - "documentId", - "sourceFileName", - "sectionPath", -] as const - -type CitationSourceField = (typeof citationSourceFields)[number] -type EvidenceSource = EvidenceLedgerSnapshot["chunks"][number]["source"] - -export function validateOutputManifest( - input: ManifestValidationInput, -): ManifestValidationResult { - const errors: string[] = [] - const text = input.manifest.text.trim() - - if (!text && input.manifest.artifacts.every((artifact) => !artifact.display)) { - errors.push("Final output must contain text or at least one displayed artifact.") - } - - validateWorkflow(input, errors) - validateArtifactRefs(input, errors) - validateArtifactCounts(input, errors) - validateGrounding(input, errors) - validateTaskEvidence(input, errors) - validateImageInspectionClaims(input, errors) - validateUnavailableImageContentClaims(input, errors) - validateTypingText(input, errors) - - return { - ok: errors.length === 0, - errors, - } -} - -function validateWorkflow( - input: ManifestValidationInput, - errors: string[], -): void { - if (input.finalized === false) { - errors.push("Agent must call finalize to produce the output manifest.") - } - - if (!input.intent) { - errors.push("Agent must declare intent before finalizing.") - } - - if (!input.contextPolicy) { - errors.push("Agent must set context policy before finalizing.") - } -} - -function validateArtifactRefs( - input: ManifestValidationInput, - errors: string[], -): void { - const knownRefs = new Set([ - ...input.ledger.chunks.map((chunk) => chunk.ref), - ...input.ledger.assets.map((asset) => asset.ref), - ]) - const sourcesByRef = new Map([ - ...input.ledger.chunks.map( - (chunk): readonly [string, EvidenceSource] => [chunk.ref, chunk.source], - ), - ...input.ledger.assets.map( - (asset): readonly [string, EvidenceSource] => [asset.ref, asset.source], - ), - ]) - - for (const artifact of input.manifest.artifacts) { - if (artifact.type === "derived_table") { - artifact.rows.forEach((row, index) => { - if (row.length !== artifact.columns.length) { - errors.push( - `Derived table row ${index + 1} has ${row.length} cells but expected ${artifact.columns.length}.`, - ) - } - }) - - for (const ref of artifact.sourceRefs) { - if (!knownRefs.has(ref)) { - errors.push( - `Derived table source ref '${ref}' was not found in the evidence ledger.`, - ) - } - } - continue - } - - if (!knownRefs.has(artifact.ref)) { - errors.push( - `Artifact ref '${artifact.ref}' was not found in the evidence ledger.`, - ) - } - } - - for (const citation of input.manifest.citations) { - if (!knownRefs.has(citation.ref)) { - errors.push(`Citation ref '${citation.ref}' was not found in the evidence ledger.`) - continue - } - - const source = sourcesByRef.get(citation.ref) - if (!source) continue - validateCitationSource( - { - ref: citation.ref, - declared: citation.source, - resolved: source, - }, - errors, - ) - } -} - -function validateCitationSource( - input: { - readonly ref: string - readonly declared: OutputCitation["source"] - readonly resolved: EvidenceSource - }, - errors: string[], -): void { - for (const field of citationSourceFields) { - validateCitationSourceField( - { - ref: input.ref, - field, - declaredValue: input.declared[field], - resolvedValue: input.resolved[field], - }, - errors, - ) - } -} - -function validateCitationSourceField( - input: { - readonly ref: string - readonly field: CitationSourceField - readonly declaredValue: string | null | undefined - readonly resolvedValue: string | null | undefined - }, - errors: string[], -): void { - const declaredValue = normalizeSourceValue(input.declaredValue) - const resolvedValue = normalizeSourceValue(input.resolvedValue) - if (declaredValue === resolvedValue) return - - errors.push( - `Citation ref '${input.ref}' source.${input.field} must match resolved evidence source. Expected ${formatSourceValue( - resolvedValue, - )}, received ${formatSourceValue(declaredValue)}.`, - ) -} - -function normalizeSourceValue(value: string | null | undefined): string | null { - return value ?? null -} - -function formatSourceValue(value: string | null): string { - return value === null ? "missing" : `'${value}'` -} - -function validateArtifactCounts( - input: ManifestValidationInput, - errors: string[], -): void { - const displayedCount = input.manifest.artifacts.filter( - (artifact) => artifact.display, - ).length - const desiredCount = input.intent?.constraints.desiredCount - const maxCount = input.intent?.constraints.maxCount - - if (typeof desiredCount === "number" && displayedCount > desiredCount) { - errors.push( - `Displayed artifact count ${displayedCount} exceeds desired count ${desiredCount}.`, - ) - } - - if (typeof maxCount === "number" && displayedCount > maxCount) { - errors.push( - `Displayed artifact count ${displayedCount} exceeds maximum count ${maxCount}.`, - ) - } -} - -function validateGrounding( - input: ManifestValidationInput, - errors: string[], -): void { - if (input.intent?.groundingPolicy !== "must_use_sources") return - - const hasLedgerEvidence = - input.ledger.chunks.length > 0 || input.ledger.evidenceText.length > 0 - const hasOutputEvidence = - input.manifest.citations.length > 0 || - input.manifest.artifacts.some((artifact) => artifact.display) - const hasUnresolved = input.manifest.unresolved.length > 0 - - if (!hasLedgerEvidence && !hasUnresolved) { - errors.push( - "Grounded output requires evidence or an explicit unresolved reason.", - ) - } - - if (hasLedgerEvidence && !hasOutputEvidence && !hasUnresolved) { - errors.push( - "Grounded output used evidence but did not cite or display any selected evidence.", - ) - } -} - -function validateTaskEvidence( - input: ManifestValidationInput, - errors: string[], -): void { - if (input.intent?.groundingPolicy !== "must_use_sources") return - if (input.manifest.unresolved.length > 0) return - - const refs = getOutputEvidenceRefs(input.manifest) - - if (input.intent.task === "compare" && refs.size < 2) { - errors.push( - "Compare outputs that must use sources require at least two evidence refs or an explicit unresolved reason.", - ) - } - - if (input.intent.task === "summarize" && refs.size < 1) { - errors.push( - "Summaries that must use sources require at least one evidence ref or an explicit unresolved reason.", - ) - } -} - -function getOutputEvidenceRefs(manifest: OutputManifest): Set { - const refs = new Set() - - for (const citation of manifest.citations) refs.add(citation.ref) - - for (const artifact of manifest.artifacts) { - if (!artifact.display) continue - if (artifact.type === "derived_table") { - artifact.sourceRefs.forEach((ref) => refs.add(ref)) - } else { - refs.add(artifact.ref) - } - } - - return refs -} - -function validateTypingText( - input: ManifestValidationInput, - errors: string[], -): void { - if (input.surface !== "typing_compose") return - - const text = input.manifest.text - if (/```|^\s*#{1,6}\s|^\s*[-*]\s/m.test(text)) { - errors.push("Typing compose output must be insertion-ready plain text.") - } -} - -function validateImageInspectionClaims( - input: ManifestValidationInput, - errors: string[], -): void { - const claimText = [input.manifest.text, ...input.manifest.unresolved].join("\n") - if (!mentionsImageInspectionResult(claimText)) return - - if (hasImageInspectionToolCall(input)) return - - errors.push( - "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", - ) -} - -function validateUnavailableImageContentClaims( - input: ManifestValidationInput, - errors: string[], -): void { - const hasImageAssets = input.ledger.assets.some((asset) => asset.type === "image") - if (!hasImageAssets) return - if (hasImageInspectionToolCall(input)) return - - const claimText = [input.manifest.text, ...input.manifest.unresolved].join("\n") - if (!mentionsUnavailableImageContent(claimText)) return - - errors.push( - "Final output must inspect available image assets before claiming retrieved page/image content cannot be read.", - ) -} - -function hasImageInspectionToolCall(input: ManifestValidationInput): boolean { - return input.toolCalls?.some((call) => call.tool === "inspectImage") === true -} - -function mentionsImageInspectionResult(value: string): boolean { - const normalized = value.replace(/\s+/g, " ").trim() - if (!normalized) return false - - return ( - /\b(?:image|visual)\s+(?:inspection|recognition|analysis)\s+(?:failed|did not|could not|was unable|found|showed|confirmed)/iu.test( - normalized, - ) || - /\bOCR\s+(?:failed|did not|could not|was unable|found|showed|confirmed|extracted)/iu.test( - normalized, - ) || - /(?:图像|图片|视觉).{0,8}(?:识别|检查|检视|查看|分析).{0,12}(?:未|没|无法|不能|不成功|失败|成功|显示|发现|提取|读取)/u.test( - normalized, - ) || - /(?:图像|图片|视觉).{0,12}(?:未|没|无法|不能|不成功|失败).{0,16}(?:识别|检查|检视|查看|分析|检测|提取|读取|获取)/u.test( - normalized, - ) || - /(?:未|没|无法|不能|不成功|失败).{0,12}(?:图像|图片|视觉|OCR).{0,12}(?:识别|检查|检视|查看|分析|检测|提取|读取|获取)/u.test( - normalized, - ) - ) -} - -function mentionsUnavailableImageContent(value: string): boolean { - const normalized = value.replace(/\s+/g, " ").trim() - if (!normalized) return false - - return ( - /\b(?:cannot|can't|could not|unable to|was unable to|did not)\s+(?:directly\s+)?(?:read|inspect|access|extract|see|view)\b.{0,48}\b(?:page|image|visual|OCR|content|clause|details|text)\b/iu.test( - normalized, - ) || - /\b(?:page|image|visual|OCR|content|clause|details|text)\b.{0,48}\b(?:cannot|can't|could not|unable to|was unable to|did not)\s+(?:directly\s+)?(?:read|inspect|access|extract|see|view)\b/iu.test( - normalized, - ) || - /(?:无法|不能|未能|没法|没有办法).{0,8}(?:直接)?(?:读取|查看|识别|检测|提取|看清|访问|获取).{0,16}(?:页面|页|图片|图像|条款|内容|细节|文字)/u.test( - normalized, - ) || - /(?:页面|页|图片|图像|条款|内容|细节|文字).{0,16}(?:无法|不能|未能|没法|没有办法).{0,16}(?:直接)?(?:读取|查看|识别|检测|提取|看清|访问|获取)/u.test( - normalized, - ) - ) -} diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 4281520..697fe77 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -1446,7 +1446,7 @@ describe("answerQuestionWithRetrieval", () => { ]); }); - it("returns a safe fallback when the harness still has validation errors", async () => { + it("returns agent output when a legacy harness trace has validation errors", async () => { const retrieval = { query: vi.fn().mockResolvedValue({ results: [makeRetrievalResult()], @@ -1485,14 +1485,13 @@ describe("answerQuestionWithRetrieval", () => { ); expect(answer).toEqual({ - answer: - "I couldn't safely finish that response because the agent output did not pass Notebook's validation checks. Please try again.", + answer: "This invalid answer should not ship.", citations: [], artifacts: [], }); }); - it("returns a safe fallback when a manifest citation ref declares the wrong source", async () => { + it("renders answer and resolves citations when manifest source metadata is wrong", async () => { process.env.AI_GATEWAY_API_KEY = "test_gateway_key"; const result = makeRetrievalResult({ content: "Information hiding is a module design principle.", @@ -1580,13 +1579,16 @@ describe("answerQuestionWithRetrieval", () => { }), ); - expect(generateCallCount).toBe(2); - expect(answer).toEqual({ - answer: - "I couldn't safely finish that response because the agent output did not pass Notebook's validation checks. Please try again.", - citations: [], - artifacts: [], - }); + expect(generateCallCount).toBe(1); + expect(answer.answer).toBe("Information hiding is a module design principle."); + expect(answer.artifacts).toEqual([]); + expect(answer.citations.map((citation) => citation.source)).toEqual([ + { + documentId: "doc_information_hiding", + sourceFileName: "information_hiding.pdf", + sectionPath: "Root / Module Design", + }, + ]); }); it("keeps image-only harness output instead of treating it as no results", async () => { @@ -2570,7 +2572,7 @@ describe("generateAgenticOutputManifest", () => { expect(result.trace.validationErrors).toEqual([]); }); - it("self-corrects an over-budget manifest via a validation-feedback revision", async () => { + it("keeps an over-budget manifest after the first generation", async () => { process.env.AI_GATEWAY_API_KEY = "test_gateway_key"; let generateCallCount = 0; vi.spyOn(ToolLoopAgent.prototype, "generate").mockImplementation( @@ -2687,12 +2689,12 @@ describe("generateAgenticOutputManifest", () => { searchSources, }); - expect(generateCallCount).toBe(2); - expect(result.trace.revisionsUsed).toBe(1); + expect(generateCallCount).toBe(1); + expect(result.trace.revisionsUsed).toBe(0); expect(result.trace.validationErrors).toEqual([]); expect( result.manifest.artifacts.filter((artifact) => artifact.display).length, - ).toBe(2); + ).toBe(3); }); }); diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index d5cd82c..d4dbcc5 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -16,7 +16,6 @@ import type { EvidenceChunk, HarnessRunResult, OutputArtifact, - OutputCitation, } from "@/agent-harness" import { toChatCitationViews, @@ -55,8 +54,6 @@ const KNOWHERE_RESPONSE_TEXT_LOG_LIMIT = 200 const KNOWHERE_CHUNK_LOG_LIMIT = 100 const KNOWHERE_RESPONSE_LOG_ITEM_LIMIT = 20 const NO_RESULTS_ANSWER = "I couldn't find that in your sources." -const HARNESS_VALIDATION_FAILURE_ANSWER = - "I couldn't safely finish that response because the agent output did not pass Notebook's validation checks. Please try again." const RAW_URL_PATTERN = /https?:\/\/[^\s)\]}>"']+/g const REDACTED_MEDIA_URL = "[media asset URL hidden]" const RETRIEVAL_TARGET_CONTENT_DATA_TYPES: Readonly< @@ -234,23 +231,8 @@ export const answerQuestionWithRetrieval = ( answerLength: generatedAnswer.manifest.text.length, retrievalCallCount: retrievalResponses.length, citationCount: generatedAnswer.manifest.citations.length, - harnessValidationErrorCount: generatedAnswer.trace.validationErrors.length, - revisionsUsed: generatedAnswer.trace.revisionsUsed, + finalized: generatedAnswer.trace.finalized, }) - if (generatedAnswer.trace.validationErrors.length > 0) { - logger.warn("chat-agent: validation failed; returning safe fallback", { - validationErrors: generatedAnswer.trace.validationErrors, - revisionsUsed: generatedAnswer.trace.revisionsUsed, - finalized: generatedAnswer.trace.finalized, - intentTask: generatedAnswer.trace.intent?.task ?? null, - retrievalCallCount: retrievalResponses.length, - }) - return { - answer: HARNESS_VALIDATION_FAILURE_ANSWER, - citations: [] as ChatCitationView[], - artifacts: [] as ChatArtifactView[], - } - } const rawResults = selectCitationRawResults({ generatedAnswer, @@ -437,7 +419,7 @@ function toChatArtifactView(input: { } function normalizeHarnessSource( - source: OutputCitation["source"], + source: EvidenceChunk["source"], sources: readonly AnswerQuestionInput["sources"][number][], ): ChatCitationView["source"] { const sourceTitle = source.documentId diff --git a/src/domains/chat/prompt.ts b/src/domains/chat/prompt.ts index f424252..a4ae326 100644 --- a/src/domains/chat/prompt.ts +++ b/src/domains/chat/prompt.ts @@ -74,7 +74,7 @@ export const generateAgenticOutputManifestEffect = ( citationCount: result.manifest.citations.length, artifactCount: result.manifest.artifacts.length, unresolvedCount: result.manifest.unresolved.length, - validationErrorCount: result.trace.validationErrors.length, + finalized: result.trace.finalized, intentTask: result.trace.intent?.task ?? null, carryHistory: result.trace.contextPolicy?.carryHistory ?? null, }) From df83e1b39a8b3109ffff0795a82432883b1950a4 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 13 Jul 2026 15:53:01 +0800 Subject: [PATCH 35/51] Attach notebook client metadata on Knowhere job creates. Share official-client defaults for created_by_client so OSS telemetry can attribute notebook uploads without leaking document content. Co-authored-by: Cursor --- src/domains/sources/document-metadata.test.ts | 49 +++++++++++++++++++ src/domains/sources/document-metadata.ts | 26 ++++++++++ src/domains/sources/knowhere-upload.ts | 15 +----- src/domains/sources/retry.test.ts | 1 + src/domains/sources/retry.ts | 15 +----- src/domains/sources/upload.test.ts | 2 + 6 files changed, 80 insertions(+), 28 deletions(-) create mode 100644 src/domains/sources/document-metadata.test.ts create mode 100644 src/domains/sources/document-metadata.ts diff --git a/src/domains/sources/document-metadata.test.ts b/src/domains/sources/document-metadata.test.ts new file mode 100644 index 0000000..ba64346 --- /dev/null +++ b/src/domains/sources/document-metadata.test.ts @@ -0,0 +1,49 @@ +import { describe, expect, it } from "vitest" + +import { + createNotebookDocumentMetadata, + NOTEBOOK_DOCUMENT_METADATA_DEFAULTS, +} from "./document-metadata" + +describe("createNotebookDocumentMetadata", () => { + it("attaches notebook client identity and display fields", () => { + expect( + createNotebookDocumentMetadata({ + title: "notes.pdf", + mimeType: "application/pdf", + sizeBytes: 12, + }), + ).toEqual({ + createdByClient: "notebook", + clientVersion: "0.1.0", + sourceFileName: "notes.pdf", + title: "notes.pdf", + mimeType: "application/pdf", + sizeBytes: 12, + }) + }) + + it("lets caller overrides win for identity fields", () => { + expect( + createNotebookDocumentMetadata({ + title: "notes.pdf", + mimeType: "application/pdf", + sizeBytes: 12, + overrides: { + createdByClient: "api", + clientVersion: "9.9.9", + }, + }), + ).toMatchObject({ + createdByClient: "api", + clientVersion: "9.9.9", + }) + }) + + it("exports notebook defaults", () => { + expect(NOTEBOOK_DOCUMENT_METADATA_DEFAULTS).toEqual({ + createdByClient: "notebook", + clientVersion: "0.1.0", + }) + }) +}) diff --git a/src/domains/sources/document-metadata.ts b/src/domains/sources/document-metadata.ts new file mode 100644 index 0000000..d6b608f --- /dev/null +++ b/src/domains/sources/document-metadata.ts @@ -0,0 +1,26 @@ +import packageJson from "../../../package.json"; + +/** + * Official notebook client identity for job `document_metadata`. + * Caller overrides win; defaults fill missing keys only when merged. + */ +export const NOTEBOOK_DOCUMENT_METADATA_DEFAULTS = { + createdByClient: "notebook", + clientVersion: packageJson.version, +} as const; + +export function createNotebookDocumentMetadata(input: { + readonly title: string; + readonly mimeType: string; + readonly sizeBytes: number; + readonly overrides?: Readonly>; +}): Readonly> { + return { + ...NOTEBOOK_DOCUMENT_METADATA_DEFAULTS, + sourceFileName: input.title, + title: input.title, + mimeType: input.mimeType, + sizeBytes: input.sizeBytes, + ...input.overrides, + }; +} diff --git a/src/domains/sources/knowhere-upload.ts b/src/domains/sources/knowhere-upload.ts index cdfdca5..1ea8396 100644 --- a/src/domains/sources/knowhere-upload.ts +++ b/src/domains/sources/knowhere-upload.ts @@ -15,6 +15,7 @@ import { validateUploadFile } from "./validation" import { TempFile, tempFileLayer } from "@/lib/temp-files" import { getUploadNamespace } from "./namespace" import { sourceFailureMessage } from "./failure-message" +import { createNotebookDocumentMetadata } from "./document-metadata" /** * Upload a browser file to Knowhere for parsing. @@ -213,20 +214,6 @@ const markSourceParsingEffect = (input: { ), ) -function createNotebookDocumentMetadata(input: { - readonly title: string - readonly mimeType: string - readonly sizeBytes: number -}): Readonly> { - return { - createdByClient: "notebook", - sourceFileName: input.title, - title: input.title, - mimeType: input.mimeType, - sizeBytes: input.sizeBytes, - } -} - function getDocumentId(job: UploadJobResult): string | null { return typeof job.documentId === "string" && job.documentId.length > 0 ? job.documentId diff --git a/src/domains/sources/retry.test.ts b/src/domains/sources/retry.test.ts index 19eb565..36c0eec 100644 --- a/src/domains/sources/retry.test.ts +++ b/src/domains/sources/retry.test.ts @@ -55,6 +55,7 @@ describe("retrySourceToKnowhereEffect", () => { namespace: "default", documentMetadata: { createdByClient: "notebook", + clientVersion: "0.1.0", sourceFileName: "notes.pdf", title: "notes.pdf", mimeType: "application/pdf", diff --git a/src/domains/sources/retry.ts b/src/domains/sources/retry.ts index 15ada9a..4ba37b4 100644 --- a/src/domains/sources/retry.ts +++ b/src/domains/sources/retry.ts @@ -9,6 +9,7 @@ import type { UploadKnowhereClient, } from "./source-upload-contracts" import { sourceFailureMessage } from "./failure-message" +import { createNotebookDocumentMetadata } from "./document-metadata" type RetrySourceRepository = { readonly markSourceParsing: ( @@ -123,20 +124,6 @@ const tryGetPlannedDocumentIdEffect = ( }).pipe(Effect.catchAll(() => Effect.succeed(null))) } -function createNotebookDocumentMetadata(input: { - readonly title: string - readonly mimeType: string - readonly sizeBytes: number -}): Readonly> { - return { - createdByClient: "notebook", - sourceFileName: input.title, - title: input.title, - mimeType: input.mimeType, - sizeBytes: input.sizeBytes, - } -} - function getDocumentId(job: UploadJobResult): string | null { return typeof job.documentId === "string" && job.documentId.length > 0 ? job.documentId diff --git a/src/domains/sources/upload.test.ts b/src/domains/sources/upload.test.ts index 77c9022..a9ad74c 100644 --- a/src/domains/sources/upload.test.ts +++ b/src/domains/sources/upload.test.ts @@ -113,6 +113,7 @@ describe("uploadSourceToKnowhere", () => { namespace: "default", documentMetadata: { createdByClient: "notebook", + clientVersion: "0.1.0", sourceFileName: "notes.pdf", title: "notes.pdf", mimeType: "application/pdf", @@ -324,6 +325,7 @@ describe("uploadSourceToKnowhere", () => { namespace: "default", documentMetadata: { createdByClient: "notebook", + clientVersion: "0.1.0", sourceFileName: "large.pdf", title: "large.pdf", mimeType: "application/pdf", From e90d0afbe6c16a037206374ac2a38d730dd6fa05 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 10 Aug 2026 00:35:43 +0800 Subject: [PATCH 36/51] feat(chat): add agentic retrieval toggle to chat composer --- src/components/chat-composer.test.ts | 43 +++++++++- src/components/chat-composer.tsx | 79 ++++++++++++++----- src/components/chat-panel.test.ts | 4 +- src/components/chat-panel.tsx | 14 +++- .../workspace-chat-workflow.test.ts | 13 ++- src/components/workspace-chat-workflow.ts | 12 ++- src/components/workspace-shell-layout.tsx | 6 +- src/domains/chat/contracts.ts | 1 + src/domains/chat/index.test.ts | 28 +++++-- src/domains/chat/index.ts | 8 +- src/domains/chat/request.ts | 3 + src/domains/chat/route-answer.ts | 1 + src/domains/chat/route-service.test.ts | 2 + src/domains/chat/service.test.ts | 4 +- src/domains/chat/service.ts | 2 + src/domains/workspace/client.ts | 1 + 16 files changed, 179 insertions(+), 42 deletions(-) diff --git a/src/components/chat-composer.test.ts b/src/components/chat-composer.test.ts index 1d7051b..0853e00 100644 --- a/src/components/chat-composer.test.ts +++ b/src/components/chat-composer.test.ts @@ -28,10 +28,51 @@ describe("ChatComposer", () => { await user.type(input, " Summarize this document "); await user.click(screen.getByRole("button", { name: "Send message" })); - expect(onSend).toHaveBeenCalledWith("Summarize this document"); + expect(onSend).toHaveBeenCalledWith("Summarize this document", { + useAgentic: true, + }); expect(input.value).toBe(""); }); + it("defaults to agentic retrieval enabled and explains the toggle", async () => { + const user = userEvent.setup(); + + render(React.createElement(ChatComposer)); + + const toggle = screen.getByRole("button", { + name: "Toggle agentic retrieval", + }); + expect(toggle.getAttribute("aria-pressed")).toBe("true"); + + await user.hover(toggle); + + const tooltip = await screen.findByRole("tooltip"); + expect(tooltip.textContent).toContain( + "Agentic retrieval plans document selection and navigation", + ); + }); + + it("sends useAgentic false after toggling agentic retrieval off", async () => { + const user = userEvent.setup(); + const onSend = vi.fn(); + + render(React.createElement(ChatComposer, { onSend })); + + const toggle = screen.getByRole("button", { + name: "Toggle agentic retrieval", + }); + await user.click(toggle); + expect(toggle.getAttribute("aria-pressed")).toBe("false"); + + const input = getComposerTextArea(); + await user.type(input, "Quick summary"); + await user.click(screen.getByRole("button", { name: "Send message" })); + + expect(onSend).toHaveBeenCalledWith("Quick summary", { + useAgentic: false, + }); + }); + it("caps long prompts and resets the composer after sending", async () => { const user = userEvent.setup(); const onSend = vi.fn(); diff --git a/src/components/chat-composer.tsx b/src/components/chat-composer.tsx index 4cfa62c..e9ee085 100644 --- a/src/components/chat-composer.tsx +++ b/src/components/chat-composer.tsx @@ -10,7 +10,7 @@ import { type MouseEvent, type ReactElement, } from "react"; -import { BarChart3, FileText, Plus, Send } from "lucide-react"; +import { BarChart3, FileText, Plus, Send, Sparkles } from "lucide-react"; import { Button } from "@/components/ui/button"; import { @@ -22,6 +22,12 @@ import { } from "@/components/ui/dropdown-menu"; import { Spinner } from "@/components/ui/spinner"; import { Textarea } from "@/components/ui/textarea"; +import { + Tooltip, + TooltipContent, + TooltipProvider, + TooltipTrigger, +} from "@/components/ui/tooltip"; import { chatPromptTemplates } from "@/domains/chat/prompt-templates"; const chatComposerName = "chat-composer"; const chatComposerTextAreaMinHeight = 128; @@ -33,6 +39,10 @@ type TextRange = { readonly end: number; }; +export type ChatSendOptions = { + readonly useAgentic: boolean; +}; + export type ChatComposerProps = { readonly canCreateDiagram?: boolean; readonly isDisabled?: boolean; @@ -40,7 +50,7 @@ export type ChatComposerProps = { readonly isSending?: boolean; readonly onCreateDiagram?: () => void; readonly onLoginClick?: () => void; - readonly onSend?: (text: string) => void; + readonly onSend?: (text: string, options: ChatSendOptions) => void; }; export function ChatComposer({ @@ -53,6 +63,7 @@ export function ChatComposer({ onSend, }: ChatComposerProps): ReactElement { const [input, setInput] = useState(""); + const [useAgentic, setUseAgentic] = useState(true); const composerInputId = useId(); const pendingTemplatePromptRef = useRef(null); const textareaRef = useRef(null); @@ -79,7 +90,7 @@ export function ChatComposer({ function handleSend(): void { if (!canSend) return; - onSend?.(trimmedInput); + onSend?.(trimmedInput, { useAgentic }); setInput(""); } @@ -177,22 +188,52 @@ export function ChatComposer({ onCreateDiagram={onCreateDiagram} onTemplateSelect={handleTemplateSelect} /> - +
+ + + + + + + Agentic retrieval plans document selection and navigation + for more thorough answers. Turn off for faster classic + search. + + + + +
)} diff --git a/src/components/chat-panel.test.ts b/src/components/chat-panel.test.ts index 54ccce5..f1bd6d1 100644 --- a/src/components/chat-panel.test.ts +++ b/src/components/chat-panel.test.ts @@ -254,7 +254,9 @@ describe("ChatPanel", () => { ); await user.click(screen.getByRole("button", { name: "Send message" })); - expect(onSend).toHaveBeenCalledWith("Summarize revenue"); + expect(onSend).toHaveBeenCalledWith("Summarize revenue", { + useAgentic: true, + }); expect( analyticsMocks.trackNotebookAssistantQuestionSubmitted, ).toHaveBeenCalledWith({ diff --git a/src/components/chat-panel.tsx b/src/components/chat-panel.tsx index 92d3cbf..1f0d1e8 100644 --- a/src/components/chat-panel.tsx +++ b/src/components/chat-panel.tsx @@ -6,7 +6,10 @@ import { type ReactElement, } from "react"; import { History, Plus } from "lucide-react"; -import { ChatComposer } from "@/components/chat-composer"; +import { + ChatComposer, + type ChatSendOptions, +} from "@/components/chat-composer"; import { ChatHistorySheet } from "@/components/chat-history-sheet"; import { ChatMessageList, @@ -47,7 +50,7 @@ export type ChatPanelProps = { messages: ChatMessageView[]; threads: ChatThreadView[]; activeThreadId?: string | null; - onSend?: (text: string) => void; + onSend?: (text: string, options: ChatSendOptions) => void; onNewChat?: () => void; onThreadSelect?: (threadId: string) => void; onThreadArchive?: (threadId: string) => void; @@ -152,7 +155,10 @@ export function ChatPanel({ } } - function handleComposerSend(text: string): void { + function handleComposerSend( + text: string, + options: ChatSendOptions, + ): void { if (isCreateDiagramCommand(text)) { void handleCreateDiagramCommand(); return; @@ -165,7 +171,7 @@ export function ChatPanel({ sourceCountSnapshot: sourceCount, messageLength: text.length, }); - onSend?.(text); + onSend?.(text, options); } return ( diff --git a/src/components/workspace-chat-workflow.test.ts b/src/components/workspace-chat-workflow.test.ts index f6b36a3..babf01a 100644 --- a/src/components/workspace-chat-workflow.test.ts +++ b/src/components/workspace-chat-workflow.test.ts @@ -84,12 +84,15 @@ describe("useWorkspaceChatWorkflow", () => { }) await act(async () => { - await result.current.handleChatSend("Summarize it") + await result.current.handleChatSend("Summarize it", { + useAgentic: true, + }) }) expect(mocks.sendChatMessage).toHaveBeenCalledWith({ message: "Summarize it", threadId: undefined, + useAgentic: true, excludedSourceIds: ["source_excluded"], }) await waitFor(() => { @@ -127,7 +130,9 @@ describe("useWorkspaceChatWorkflow", () => { }) await act(async () => { - await result.current.handleChatSend("What changed in Q4?") + await result.current.handleChatSend("What changed in Q4?", { + useAgentic: true, + }) }) expect(mocks.materializeDemoSources).toHaveBeenCalledWith({ @@ -172,7 +177,9 @@ describe("useWorkspaceChatWorkflow", () => { }) await act(async () => { - await result.current.handleChatSend("Summarize it") + await result.current.handleChatSend("Summarize it", { + useAgentic: true, + }) }) expect(mocks.materializeDemoSources).not.toHaveBeenCalled() diff --git a/src/components/workspace-chat-workflow.ts b/src/components/workspace-chat-workflow.ts index 2943a23..dd85576 100644 --- a/src/components/workspace-chat-workflow.ts +++ b/src/components/workspace-chat-workflow.ts @@ -12,6 +12,7 @@ import { type AnalyticsContext, } from "@/lib/posthog" import { workspaceClient } from "@/domains/workspace/client" +import type { ChatSendOptions } from "@/components/chat-composer" import { workspaceClientCache, type ChatThreadDetailResponse, @@ -43,7 +44,10 @@ type WorkspaceChatWorkflow = { readonly chat: ReturnType readonly chatThreads: ChatThreadView[] readonly handleArchiveChatThread: (threadId: string) => Promise - readonly handleChatSend: (text: string) => Promise + readonly handleChatSend: ( + text: string, + options: ChatSendOptions, + ) => Promise readonly handleCreateChatThread: () => Promise readonly handleRefreshActiveChatThread: () => Promise readonly handleSelectChatThread: (threadId: string) => void @@ -247,7 +251,10 @@ export function useWorkspaceChatWorkflow({ } } - async function handleChatSend(text: string): Promise { + async function handleChatSend( + text: string, + options: ChatSendOptions, + ): Promise { const sendStart = Date.now() const selectedSourcesCount = sources.filter( (source) => @@ -297,6 +304,7 @@ export function useWorkspaceChatWorkflow({ const body = await sendChatMessage({ message: text, threadId: chat.threadId ?? undefined, + useAgentic: options.useAgentic, excludedSourceIds: sources .filter((source) => source.excludedFromQuery) .map((source) => source.id), diff --git a/src/components/workspace-shell-layout.tsx b/src/components/workspace-shell-layout.tsx index 0a71d90..72c6724 100644 --- a/src/components/workspace-shell-layout.tsx +++ b/src/components/workspace-shell-layout.tsx @@ -8,6 +8,7 @@ import { } from "lucide-react" import { ChatPanel } from "@/components/chat-panel" +import type { ChatSendOptions } from "@/components/chat-composer" import { ChunksPanel } from "@/components/chunks-panel" import { MobileTabBar } from "@/components/mobile-tab-bar" import { OfficialLibraryPanel } from "@/components/official-library-panel" @@ -99,7 +100,10 @@ export type WorkspaceShellLayoutProps = { readonly onArchiveChatThread: (threadId: string) => void | Promise readonly onArchiveSource: (sourceId: string) => void | Promise readonly onRetrySource?: (sourceId: string) => void | Promise - readonly onChatSend: (text: string) => void | Promise + readonly onChatSend: ( + text: string, + options: ChatSendOptions, + ) => void | Promise readonly onCitationClick: ( citation: ChatCitationView, citationId: string, diff --git a/src/domains/chat/contracts.ts b/src/domains/chat/contracts.ts index e43b232..527b52c 100644 --- a/src/domains/chat/contracts.ts +++ b/src/domains/chat/contracts.ts @@ -73,6 +73,7 @@ export type AnswerQuestionInput = { namespaces?: readonly string[] sources: readonly Source[] excludedSourceIds: readonly string[] + useAgentic?: boolean retrieval: RetrievalClient knowledge?: Knowledge remoteDocumentClient?: NotebookKnowhereRemoteDocumentClient diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 697fe77..f026dc6 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -85,7 +85,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "What does the document say?", topK: 8, - useAgentic: false, + useAgentic: true, dataType: 1, excludeDocumentIds: ["doc_excluded", "doc_remote"], }); @@ -246,7 +246,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "diagram", topK: 2, - useAgentic: false, + useAgentic: true, dataType: 3, excludeDocumentIds: ["doc_excluded"], }); @@ -772,7 +772,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "SpaceX rocket photos", topK: 8, - useAgentic: false, + useAgentic: true, dataType: 3, }); expect(answer.answer).toBe("Use this launch photo."); @@ -1897,7 +1897,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "公民身份证明 图片", topK: 8, - useAgentic: false, + useAgentic: true, dataType: 3, }); const imageCitations = answer.citations.filter( @@ -1988,7 +1988,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "Tesla Q4 2025 Update energy generation and storage deployments", topK: 8, - useAgentic: false, + useAgentic: true, dataType: 1, }); expect(generateAnswer).toHaveBeenCalledWith({ @@ -2034,6 +2034,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", sources: [makeSource()], excludedSourceIds: [], + useAgentic: false, retrieval, generateAnswer, messages, @@ -2711,11 +2712,28 @@ describe("parseChatRequestBody", () => { value: { question: "What changed?", threadId: "thread_1", + useAgentic: true, excludedSourceIds: ["source_1", "source_2"], }, }); }); + it("keeps an explicit useAgentic choice from the request body", () => { + expect( + parseChatRequestBody({ + message: "Quick summary", + useAgentic: false, + }), + ).toEqual({ + ok: true, + value: { + question: "Quick summary", + useAgentic: false, + excludedSourceIds: [], + }, + }); + }); + it("rejects empty questions before retrieval or model calls", () => { expect(parseChatRequestBody({ message: " " })).toEqual({ ok: false, diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index d4dbcc5..f45c88d 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -42,9 +42,6 @@ import type { HardenableRetrievalResult } from "./media-asset-hardening" import { notebookKnowhereTools } from "./knowhere-tools" const DEFAULT_TOP_K = 8 -const NOTEBOOK_USE_AGENTIC_RETRIEVAL: NonNullable< - RetrievalQueryParams["useAgentic"] -> = false const MAX_AGENTIC_TOP_K = 12 const MAX_AGENTIC_MERGED_RESULT_COUNT = 24 const MAX_AGENTIC_MERGED_REFERENCED_CHUNK_COUNT = 24 @@ -146,6 +143,7 @@ export const answerQuestionWithRetrieval = ( input: queryInput, fallbackQuestion: question, namespace, + useAgentic: input.useAgentic ?? true, sources: input.sources, excludedSourceIds: input.excludedSourceIds, }) @@ -153,6 +151,7 @@ export const answerQuestionWithRetrieval = ( namespace, query: retrievalQueryParams.query, topK: retrievalQueryParams.topK, + useAgentic: retrievalQueryParams.useAgentic, dataType: retrievalQueryParams.dataType ?? null, signalPathCount: retrievalQueryParams.signalPaths?.length ?? 0, filterMode: retrievalQueryParams.filterMode ?? null, @@ -775,6 +774,7 @@ function buildRetrievalQueryParams(input: { readonly input: AgenticRetrievalQuery readonly fallbackQuestion: string readonly namespace: string + readonly useAgentic: boolean readonly sources: AnswerQuestionInput["sources"] readonly excludedSourceIds: readonly string[] }): RetrievalQueryParams { @@ -787,7 +787,7 @@ function buildRetrievalQueryParams(input: { namespace: input.namespace, query, topK: normalizeTopK(input.input.topK), - useAgentic: NOTEBOOK_USE_AGENTIC_RETRIEVAL, + useAgentic: input.useAgentic, dataType, ...(input.input.signalPaths && input.input.signalPaths.length > 0 ? { signalPaths: input.input.signalPaths } diff --git a/src/domains/chat/request.ts b/src/domains/chat/request.ts index c15a76e..3e56f06 100644 --- a/src/domains/chat/request.ts +++ b/src/domains/chat/request.ts @@ -3,6 +3,7 @@ import { Either, Schema } from "effect" export type ParsedChatRequest = { question: string threadId?: string + useAgentic: boolean excludedSourceIds: string[] } @@ -13,6 +14,7 @@ export type ParseChatRequestResult = const ChatRequestBody = Schema.Struct({ message: Schema.String, threadId: Schema.optional(Schema.String), + useAgentic: Schema.optional(Schema.Boolean), excludedSourceIds: Schema.optional(Schema.Array(Schema.Unknown)), }) @@ -43,6 +45,7 @@ export function parseChatRequestBody(body: unknown): ParseChatRequestResult { parsed.threadId !== undefined && parsed.threadId.length > 0 ? parsed.threadId : undefined, + useAgentic: parsed.useAgentic ?? true, excludedSourceIds, }, } diff --git a/src/domains/chat/route-answer.ts b/src/domains/chat/route-answer.ts index 3f9fc9d..ae8b13d 100644 --- a/src/domains/chat/route-answer.ts +++ b/src/domains/chat/route-answer.ts @@ -142,6 +142,7 @@ const answerChatEffect = (input: AnswerChatInput) => sources, question: body.value.question, threadId: body.value.threadId, + useAgentic: body.value.useAgentic, excludedSourceIds: body.value.excludedSourceIds, retrieval: client.retrieval, knowledge: knowhereResources.knowledge, diff --git a/src/domains/chat/route-service.test.ts b/src/domains/chat/route-service.test.ts index d66ae68..97c9f6d 100644 --- a/src/domains/chat/route-service.test.ts +++ b/src/domains/chat/route-service.test.ts @@ -137,6 +137,7 @@ describe("chat route services", () => { body: { message: " Summarize it ", threadId: "thread_1", + useAgentic: true, excludedSourceIds: ["source_skipped", null], }, }) @@ -159,6 +160,7 @@ describe("chat route services", () => { sources: [readySource], question: "Summarize it", threadId: "thread_1", + useAgentic: true, excludedSourceIds: ["source_skipped"], retrieval: client.retrieval, generateAnswer: mocks.generateAgenticOutputManifest, diff --git a/src/domains/chat/service.test.ts b/src/domains/chat/service.test.ts index 3a1ab82..6001171 100644 --- a/src/domains/chat/service.test.ts +++ b/src/domains/chat/service.test.ts @@ -57,7 +57,7 @@ describe("handleChatTurn", () => { namespace: "notebook-namespace", query: "What does the document say?", topK: 8, - useAgentic: false, + useAgentic: true, dataType: 1, excludeDocumentIds: ["doc_excluded"], }); @@ -223,7 +223,7 @@ describe("handleChatTurn", () => { namespace: "notebook-namespace", query: "Tesla Q4 2025 Update energy generation and storage deployments", topK: 8, - useAgentic: false, + useAgentic: true, dataType: 1, }); }); diff --git a/src/domains/chat/service.ts b/src/domains/chat/service.ts index 7a7b2b5..23bfa68 100644 --- a/src/domains/chat/service.ts +++ b/src/domains/chat/service.ts @@ -64,6 +64,7 @@ type ChatTurnInput = { sources: readonly Source[] question: string threadId?: string + useAgentic?: boolean excludedSourceIds: readonly string[] retrieval: RetrievalClient knowledge?: AnswerQuestionInput["knowledge"] @@ -126,6 +127,7 @@ export const handleChatTurnEffect = (input: ChatTurnInput) => namespace: input.workspace.namespace, namespaces: getCompatibleNamespaces(input.workspace), sources: readySources, + useAgentic: input.useAgentic ?? true, excludedSourceIds: input.excludedSourceIds, retrieval: input.retrieval, knowledge: input.knowledge, diff --git a/src/domains/workspace/client.ts b/src/domains/workspace/client.ts index a415d6a..7720383 100644 --- a/src/domains/workspace/client.ts +++ b/src/domains/workspace/client.ts @@ -50,6 +50,7 @@ type ChatThreadDetailResponse = ChatThreadResponse & { type ChatMessageRequest = { message: string threadId?: string + useAgentic: boolean excludedSourceIds: string[] } From edd8a79682fe58b7b4642603b1ee3e02d6548800 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 10 Aug 2026 09:37:35 +0800 Subject: [PATCH 37/51] feat(chat): refine deep search toggle UI --- src/components/chat-composer.test.ts | 18 ++++++------- src/components/chat-composer.tsx | 40 +++++++++++++++------------- 2 files changed, 30 insertions(+), 28 deletions(-) diff --git a/src/components/chat-composer.test.ts b/src/components/chat-composer.test.ts index 0853e00..a12264a 100644 --- a/src/components/chat-composer.test.ts +++ b/src/components/chat-composer.test.ts @@ -34,35 +34,35 @@ describe("ChatComposer", () => { expect(input.value).toBe(""); }); - it("defaults to agentic retrieval enabled and explains the toggle", async () => { + it("defaults to deep search enabled and explains the toggle", async () => { const user = userEvent.setup(); render(React.createElement(ChatComposer)); - const toggle = screen.getByRole("button", { - name: "Toggle agentic retrieval", + const toggle = screen.getByRole("checkbox", { + name: "Deep search", }); - expect(toggle.getAttribute("aria-pressed")).toBe("true"); + expect(toggle.getAttribute("aria-checked")).toBe("true"); await user.hover(toggle); const tooltip = await screen.findByRole("tooltip"); expect(tooltip.textContent).toContain( - "Agentic retrieval plans document selection and navigation", + "Deep search plans document selection and navigation", ); }); - it("sends useAgentic false after toggling agentic retrieval off", async () => { + it("sends useAgentic false after toggling deep search off", async () => { const user = userEvent.setup(); const onSend = vi.fn(); render(React.createElement(ChatComposer, { onSend })); - const toggle = screen.getByRole("button", { - name: "Toggle agentic retrieval", + const toggle = screen.getByRole("checkbox", { + name: "Deep search", }); await user.click(toggle); - expect(toggle.getAttribute("aria-pressed")).toBe("false"); + expect(toggle.getAttribute("aria-checked")).toBe("false"); const input = getComposerTextArea(); await user.type(input, "Quick summary"); diff --git a/src/components/chat-composer.tsx b/src/components/chat-composer.tsx index e9ee085..4aab263 100644 --- a/src/components/chat-composer.tsx +++ b/src/components/chat-composer.tsx @@ -10,9 +10,10 @@ import { type MouseEvent, type ReactElement, } from "react"; -import { BarChart3, FileText, Plus, Send, Sparkles } from "lucide-react"; +import { BarChart3, FileText, Plus, Send } from "lucide-react"; import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; import { DropdownMenu, DropdownMenuContent, @@ -192,28 +193,29 @@ export function ChatComposer({ - + + setUseAgentic(checked === true) + } + /> + Deep search + - Agentic retrieval plans document selection and navigation - for more thorough answers. Turn off for faster classic - search. + Deep search plans document selection and navigation for + more thorough answers. Turn off for faster classic search. From 112dc2903e0babb13de0230d3d8d4173347f07f9 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 19 Aug 2026 22:34:51 +0800 Subject: [PATCH 38/51] feat: add parsed-results citation UX with page chips Show {title}/pN chips, numbered SOURCES, page-card jump, Brain chrome, and turn export. Hydrate grep citation pages until the SDK match payload includes them. Co-authored-by: Cursor --- docs/parsed-results-citation-ux.md | 72 ++++ e2e/chunk-original-preview.e2e.ts | 4 +- e2e/citation-dedupe.e2e.ts | 21 + src/agent-harness/ledger.test.ts | 71 ++++ src/agent-harness/ledger.ts | 36 ++ src/agent-harness/runtime.test.ts | 9 + src/agent-harness/runtime.ts | 3 + src/app/e2e/citation-same-page/page.tsx | 78 ++++ src/app/global-error.tsx | 4 +- src/app/login/page.tsx | 2 +- src/components/chat-citation-model.test.ts | 153 +++++++ src/components/chat-citation-model.ts | 284 +++++++++++++ .../chat-message-export-pdf.test.ts | 16 + src/components/chat-message-export-pdf.ts | 136 ++++++ src/components/chat-message-list.test.ts | 124 +++++- src/components/chat-message-list.tsx | 392 +++++++++--------- src/components/chat-panel.test.ts | 37 +- src/components/chunks-panel-state.test.ts | 45 +- src/components/chunks-panel-state.ts | 37 +- src/components/chunks-panel-workflow.test.ts | 43 ++ src/components/chunks-panel-workflow.ts | 27 +- src/components/chunks-panel.test.ts | 20 +- src/components/chunks-panel.tsx | 32 +- .../parsed-chunk-card-model.test.ts | 24 ++ src/components/parsed-chunk-card-model.ts | 20 + src/components/parsed-chunk-card.test.ts | 61 +++ src/components/parsed-chunk-card.tsx | 38 +- src/components/sources-panel.test.ts | 35 ++ src/components/sources-panel.tsx | 31 +- src/components/top-nav.tsx | 2 +- .../workspace-citation-focus.test.ts | 188 ++++++++- src/components/workspace-citation-focus.ts | 218 +++++++--- .../workspace-selected-chunks.test.ts | 4 +- src/components/workspace-selected-chunks.ts | 19 +- src/components/workspace-shell-layout.test.ts | 15 + src/components/workspace-shell-layout.tsx | 20 + src/components/workspace-shell.test.ts | 34 +- src/domains/chat/citations.test.ts | 19 + src/domains/chat/citations.ts | 10 +- src/domains/chat/index.test.ts | 136 ++++++ src/domains/chat/index.ts | 131 +++++- src/domains/chat/page-citation-assets.test.ts | 46 +- src/domains/chat/page-citation-assets.ts | 130 ++++-- src/domains/chunks/index.test.ts | 49 +++ src/domains/chunks/index.ts | 13 + src/domains/chunks/normalization.ts | 17 +- src/domains/chunks/read.test.ts | 37 ++ src/domains/chunks/read.ts | 10 +- src/domains/sources/route-chunks.ts | 26 +- src/domains/workspace/client-cache.test.ts | 39 ++ src/domains/workspace/client-cache.ts | 63 +++ src/domains/workspace/client.test.ts | 42 ++ src/domains/workspace/client.ts | 65 ++- src/integrations/knowhere-demo.test.ts | 74 ++++ src/integrations/knowhere-demo.ts | 74 +++- src/lib/app-metadata.ts | 2 +- 56 files changed, 2897 insertions(+), 441 deletions(-) create mode 100644 docs/parsed-results-citation-ux.md create mode 100644 src/app/e2e/citation-same-page/page.tsx create mode 100644 src/components/chat-citation-model.test.ts create mode 100644 src/components/chat-citation-model.ts create mode 100644 src/components/chat-message-export-pdf.test.ts create mode 100644 src/components/chat-message-export-pdf.ts diff --git a/docs/parsed-results-citation-ux.md b/docs/parsed-results-citation-ux.md new file mode 100644 index 0000000..4730107 --- /dev/null +++ b/docs/parsed-results-citation-ux.md @@ -0,0 +1,72 @@ +# Parsed Results + citation UX + +Recorded 2026-08-18 so we do not relitigate the Figma / issue discussion. + +Sources: + +- Figma [Knowhere (Copy) / brain](https://www.figma.com/design/0jjj7SQ1N0zo4D9WzJoNJh/Knowhere--Copy-?node-id=4000-37742) (`4000:37742`) +- GitHub [knowhere#222](https://github.com/Ontos-AI/knowhere/issues/222), [knowhere#223](https://github.com/Ontos-AI/knowhere/issues/223) +- Product lock from that conversation (issues and Figma notes are partly stale) + +This work lands in **Knowhere Notebook**. Region geometry later lands in **Knowhere retrieval** (`Ontos-AI/knowhere`) and is only consumed here when the payload exists. + +## Locked product + +Vision / page-image mode only. There is **no original PDF preview** as a citation destination. + +| Topic | Decision | +| --- | --- | +| Layout | Sources \| Parsed Results \| Assistant (today’s Notebook columns). Ignore the Figma frame that put Assistant in the middle. | +| Rename | Knowhere Notebook → **Knowhere Brain** | +| Canvas | Title **Parsed Results**, subtitle **From {source title}** | +| Page card | **Page N + parse path**, page image, keywords. No keywords → hide the keyword row, still show the image. Screen 2 “text” is a crop overlay on the page image, not a text chunk. | +| List / Tree | Keep **both**. Tree is the **existing** section tree, also in page-image mode. **Default Tree**. Citation click may switch to List so the page card is visible. | +| Inline chips | One chip **per citation**, not per page. Label `{source title}/pN` (same name as the Sources row). Max width 250px, truncate. Same page can be two chips; later they get different crop boxes. | +| Chip / `pN` click | Switch the **viewed** source if needed (not query checkboxes), open that file’s parsed page, scroll to it. Chat may cite any workspace file while the canvas shows another. | +| Footer SOURCES | Numbered list, one row per file, `{title}` + cited page links `p25 p26 p27`. Those links do the **same jump** as chips. | +| Copy / download | Copy answer / code / tables. Download **this assistant turn** as **Markdown and PDF**, **without citations** (no chips, no SOURCES list). | +| Region box | **Required product**, **not this slice**. When retrieval sends a crop, draw `#8E51FF` / 25% / 8px radius on that page image (flash twice, 3s, fade). Until then, page jump only. Do not fake a full-page box. | + +Out of this slice: original-file viewer for citations, region overlay, faking geometry. + +## Current code (gaps) + +Almost all of this is UI in this repo. Page images, keywords, page numbers, and source-title chips already exist in pieces. + +| Area | Today | Need | +| --- | --- | --- | +| Branding | “Knowhere Notebook” in `top-nav.tsx`, `lib/app-metadata.ts`, login, `global-error.tsx` | Brain | +| Canvas copy | “Parsed Chunks” / “Showing all parsed chunks from …” | Parsed Results / From {title} | +| Original preview | Hidden for `page-assets` sources (`chunks-panel.tsx`) | Keep hidden. Citations must not reopen it. | +| List / Tree | Tree **hidden** for page-asset sources (`!isPageAssetSource`). Default mode is already `"tree"` but unused in page mode. | Show the toggle for page mode; Tree uses existing `ChunkSectionTree`. | +| Page cards | `PageChunkCard` already shows page image + `ChunkKeywords` | Match Figma title (`Page N` + path). Image-first. Hide keywords when empty (already returns null). | +| Inline chips | `[Source N: …]` tokens are **stripped** (`buildCitationContentMarkdown`). Chips live in a footer `AssistantSources`. | Replace tokens with inline `{title}/pN` chips in the answer. | +| Chip label | Filename / source title only, **no `/pN`**. | `{title}/pN` | +| Dedupe | `getDisplayCitations` and `selectCitationRawResults` drop same-key citations. | Keep one chip per citation so two regions on one page stay two chips. | +| Footer SOURCES | Flat chips | Numbered per-file list + page links | +| Click target | `handleCitationClick` focuses a **chunk**. `focusedPage` is plumbed to `ChunksPanel` but **not used**. | Focus/scroll the **page** card; `onSelectSource` for file B while viewing A already exists. | +| Export | None | Per-turn copy + MD/PDF without citations | + +Later API (do not block Notebook UI): + +- Retrieval / agent citation payload has page number and page image today (`pageCitationPageNumber`, `pageCitationAssetUrl`). +- It does **not** have a citation crop box. Worker bboxes are for asset crops, not “this quote lives here.” +- When Knowhere adds a region on the citation, extend `ChatCitationView` + persistence, then overlay on `PageCitationAssetImage`. + +## Implementation slices (Notebook) + +Do these in order. Tests follow existing ownership: model tests for labels/grouping, component tests for chips/cards/export, panel tests for click → source + page. + +1. **Chrome** — rename Brain; Parsed Results header/subtitle; enable List/Tree on page-asset sources; default Tree. +2. **Page cards** — Figma card header (page + path); image + optional keywords; citation click forces List (already does via `citationListViewRequestId`) and scrolls to the page chunk. +3. **Citations** — stop stripping without replacement; inline chips `{title}/pN`; stop collapsing same-page citations; numbered SOURCES with `pN` links; click switches viewed source and page-focuses. +4. **Copy / download** — message actions: clipboard of visible answer (code/tables included as rendered); MD + PDF of this turn with the same citation-stripped body we already produce for display (`buildCitationContentMarkdown` / equivalent). PDF is client-side from that Markdown, not the original document. +5. **Later** — region overlay when retrieval sends a box. Keep citation identity stable (do not key chips by `title+page` only). + +## Do not + +- Route citation clicks to original PDF preview. +- Deduplicate inline chips by page. +- Draw a highlight without retrieval geometry. +- Change which sources are **checked for the next query** when jumping the canvas. +- Put this feature in archived `knowhere-api`; new retrieval fields go in `Ontos-AI/knowhere`. diff --git a/e2e/chunk-original-preview.e2e.ts b/e2e/chunk-original-preview.e2e.ts index 1a6502c..c7e0426 100644 --- a/e2e/chunk-original-preview.e2e.ts +++ b/e2e/chunk-original-preview.e2e.ts @@ -7,7 +7,7 @@ test("reopens the demo PDF preview after returning from original to parsed chunk }) => { await page.goto("/") await expect( - page.getByRole("heading", { name: "Parsed Chunks" }), + page.getByRole("heading", { name: "Parsed Results" }), ).toBeVisible() await openFirstChunk(page) @@ -18,7 +18,7 @@ test("reopens the demo PDF preview after returning from original to parsed chunk .getByRole("button", { name: "Parsed", exact: true }) .click() await expect( - page.getByRole("heading", { name: "Parsed Chunks" }), + page.getByRole("heading", { name: "Parsed Results" }), ).toBeVisible() await openFirstChunk(page) diff --git a/e2e/citation-dedupe.e2e.ts b/e2e/citation-dedupe.e2e.ts index 78de15e..071a986 100644 --- a/e2e/citation-dedupe.e2e.ts +++ b/e2e/citation-dedupe.e2e.ts @@ -87,3 +87,24 @@ test("keeps duplicate source labels clickable for separate documents", async ({ await expect(page.getByText("First report source content.")).toBeVisible() expect(firstSourceChunkRequests).toBeGreaterThan(0) }) + +test("keeps two chips to the same title/pN as separate buttons", async ({ + context, + page, +}) => { + await context.addCookies([ + { + name: "better-auth.session_token", + value: "playwright", + url: "http://localhost:3000", + }, + ]) + + await page.goto("/e2e/citation-same-page") + + const chatPanel = page.getByTestId("desktop-chat-panel") + await expect(chatPanel.getByTestId("citation-chip")).toHaveCount(2) + await expect( + chatPanel.getByRole("button", { name: "Open source spacex-s1.pdf/p26" }), + ).toHaveCount(2) +}) diff --git a/src/agent-harness/ledger.test.ts b/src/agent-harness/ledger.test.ts index 061a892..15f1b20 100644 --- a/src/agent-harness/ledger.test.ts +++ b/src/agent-harness/ledger.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it } from "vitest" import type { + KnowledgeGrepMatch, KnowledgeGrepResponse, KnowledgeReadResponse, RetrievalQueryResponse, @@ -111,6 +112,76 @@ describe("createEvidenceLedger", () => { }), ]) }) + + it("copies page metadata onto grep matches from the same chunk id", () => { + const ledger = createEvidenceLedger() + ledger.addReadChunksResponse(makeReadResponse()) + + const snapshot = ledger.addGrepChunksResponse({ + ...makeGrepResponse(), + matches: [ + { + position: 1, + chunkId: "chunk_page_1", + chunkType: "page", + sectionPath: "Root / Page 1", + sourceChunkPath: "pages/page-1.md", + startOffset: 0, + endOffset: 12, + snippet: "Full page snip", + }, + ], + }) + + expect(snapshot.chunks[1]).toEqual( + expect.objectContaining({ + ref: "grep1:match:1", + kind: "grep_match", + chunkId: "chunk_page_1", + metadata: expect.objectContaining({ + pageNums: [1], + position: 1, + startOffset: 0, + endOffset: 12, + }), + }), + ) + }) + + it("copies pageNumbers from the grep match when the SDK provides them", () => { + const ledger = createEvidenceLedger() + + const snapshot = ledger.addGrepChunksResponse({ + ...makeGrepResponse(), + matches: [ + { + position: 1, + chunkId: "chunk_page_4", + chunkType: "page", + sectionPath: "FINANCIAL SUMMARY", + sourceChunkPath: "pages/page-4.md", + startOffset: 0, + endOffset: 12, + snippet: "automotive revenues", + pageNumbers: [4], + } as KnowledgeGrepMatch, + ], + }) + + expect(snapshot.chunks[0]).toEqual( + expect.objectContaining({ + ref: "grep1:match:1", + kind: "grep_match", + chunkId: "chunk_page_4", + metadata: expect.objectContaining({ + pageNums: [4], + position: 1, + startOffset: 0, + endOffset: 12, + }), + }), + ) + }) }) function makeRetrievalResponse(): RetrievalQueryResponse { diff --git a/src/agent-harness/ledger.ts b/src/agent-harness/ledger.ts index d782ae3..aa9d58e 100644 --- a/src/agent-harness/ledger.ts +++ b/src/agent-harness/ledger.ts @@ -255,6 +255,19 @@ function addChunkFromGrepMatch(input: { readonly match: KnowledgeGrepMatch readonly ref: string }): void { + const matchWithPages = input.match as KnowledgeGrepMatch & { + readonly pageNumbers?: readonly number[] + readonly metadata?: Readonly> + } + const donor = input.ledger.chunks.find( + (chunk) => + chunk.chunkId === input.match.chunkId && hasPageMetadata(chunk.metadata), + ) + const pageNums = + matchWithPages.pageNumbers && matchWithPages.pageNumbers.length > 0 + ? [...matchWithPages.pageNumbers] + : undefined + addChunk({ ledger: input.ledger, chunk: { @@ -268,6 +281,9 @@ function addChunkFromGrepMatch(input: { sourceChunkPath: input.match.sourceChunkPath, filePath: input.match.filePath, metadata: { + ...(donor?.metadata ?? {}), + ...(matchWithPages.metadata ?? {}), + ...(pageNums ? { pageNums } : {}), position: input.match.position, startOffset: input.match.startOffset, endOffset: input.match.endOffset, @@ -282,6 +298,26 @@ function addChunkFromGrepMatch(input: { }) } +function hasPageMetadata( + metadata: Readonly> | undefined, +): boolean { + if (!metadata) return false + const values = [ + metadata.pageNums, + metadata.page_nums, + metadata.pageNum, + metadata.page_num, + metadata.pageAssets, + metadata.page_assets, + ] + return values.some((value) => { + if (Array.isArray(value)) return value.length > 0 + if (typeof value === "number") return Number.isSafeInteger(value) && value > 0 + if (typeof value === "string") return value.trim().length > 0 + return false + }) +} + function addChunk(input: { readonly ledger: MutableLedger readonly chunk: Omit diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index 04614a7..12873d3 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -40,6 +40,15 @@ describe("agent harness runtime", () => { expect(prompt).not.toContain("must match the selected evidence ref exactly") }) + it("tells the agent to emit [[cite:n]] markers instead of title/pN or [1]", () => { + const prompt = buildHarnessSystemPrompt(makeTurnInput()) + + expect(prompt).toContain("[[cite:n]]") + expect(prompt).toContain("1-based index into the citations array") + expect(prompt).toContain("Do not write title/pN, [1], Markdown footnotes") + expect(prompt).toContain("Do not collapse same-page citations") + }) + it("passes only outer retrieval parameters to KNOWHERE without planning-tool gating", async () => { const query = vi.fn().mockResolvedValue( makeRetrievalResponse(), diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index 8963fdf..0ad95d8 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -1176,6 +1176,9 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "- artifacts with display=true are the exact images/tables shown. Never display every candidate; honor constraints.desiredCount / maxCount.", "- Use type=derived_table only for tables you create from evidence; every derived_table.sourceRefs entry must reference evidence in the ledger.", "- Prefer citation and selected image/table artifact refs returned by Knowhere tools in the evidence ledger.", + "- Place [[cite:n]] immediately after the supported claim. n is the 1-based index into the citations array passed to finalize.", + "- Do not write title/pN, [1], Markdown footnotes, or [Source N: ...] in the answer text. Notebook renders chips from [[cite:n]] and citation metadata.", + "- Repeat [[cite:n]] when another claim uses the same page. Do not collapse same-page citations to one row.", "- Citation label and source metadata are optional. Notebook resolves citation metadata from evidence refs when possible.", "- If evidence is relevant but you cannot identify a supporting evidence ref, answer with unresolved issues instead of fabricating a ref.", "- inspectImage observations are inspection notes, not new source refs. Final citations and displayed image artifacts must use the original retrieved image asset refs.", diff --git a/src/app/e2e/citation-same-page/page.tsx b/src/app/e2e/citation-same-page/page.tsx new file mode 100644 index 0000000..c825f91 --- /dev/null +++ b/src/app/e2e/citation-same-page/page.tsx @@ -0,0 +1,78 @@ +import { notFound } from "next/navigation" + +import { WorkspaceShell } from "@/components/workspace-shell" +import type { ChatMessageView, ChatThreadView } from "@/domains/chat/types" +import type { SourceView } from "@/domains/sources/types" + +const sources: SourceView[] = [ + { + id: "source_spacex", + title: "spacex-s1.pdf", + mimeType: "application/pdf", + status: "ready", + documentId: "doc_spacex", + chunkCount: 1, + documentPresentation: { kind: "page-assets", pageCount: 26 }, + }, +] + +const chatThreads: ChatThreadView[] = [ + { + id: "thread_1", + title: "Same page citations", + createdAt: "2026-05-17T00:00:00.000Z", + updatedAt: "2026-05-17T00:00:00.000Z", + }, +] + +const chatMessages: ChatMessageView[] = [ + { + id: "assistant_1", + role: "assistant", + content: "Revenue grew [[cite:1]] and later expanded [[cite:2]].", + citations: [ + { + chunkType: "page", + score: 0.91, + pageCitationPageNumber: 26, + source: { + documentId: "doc_spacex", + sourceFileName: "spacex-s1.pdf", + sectionPath: "Page 26", + }, + }, + { + chunkType: "page", + score: 0.89, + pageCitationPageNumber: 26, + source: { + documentId: "doc_spacex", + sourceFileName: "spacex-s1.pdf", + sectionPath: "Page 26", + }, + }, + ], + }, +] + +export default function CitationSamePageTestPage() { + if (process.env.NODE_ENV === "production") notFound() + + return ( + + ) +} diff --git a/src/app/global-error.tsx b/src/app/global-error.tsx index 5fd5cc7..409a3af 100644 --- a/src/app/global-error.tsx +++ b/src/app/global-error.tsx @@ -15,11 +15,11 @@ export default function GlobalError({ return ( - Something went wrong - Knowhere Notebook + Something went wrong - Knowhere Brain

- Notebook error + Brain error

Something went wrong diff --git a/src/app/login/page.tsx b/src/app/login/page.tsx index c802af3..8ae3bc7 100644 --- a/src/app/login/page.tsx +++ b/src/app/login/page.tsx @@ -32,7 +32,7 @@ export async function LoginContent() {

- Knowhere Notebook + Knowhere Brain

{ + it("labels chips as title/pN and keeps two same-page citations as two identities", () => { + expect( + chatCitationModel.getCitationChipLabel(pageCitation, { + doc_1: "spacex-s1.pdf", + }), + ).toBe("spacex-s1.pdf/p26") + + const groups = chatCitationModel.groupCitationsByFile( + "assistant_1", + [pageCitation, samePageCitation, otherFileCitation], + { + doc_1: "spacex-s1.pdf", + doc_2: "report.pdf", + }, + ) + + expect(groups).toHaveLength(2) + expect(groups[0]?.entries.map((entry) => entry.citationId)).toEqual([ + "assistant_1:0", + "assistant_1:1", + ]) + expect(groups[0]?.entries.map((entry) => entry.pageNumber)).toEqual([26, 26]) + expect( + chatCitationModel + .uniquePageLinkEntries(groups[0]!.entries) + .map((entry) => entry.pageNumber), + ).toEqual([26]) + expect(groups[1]?.title).toBe("report.pdf") + }) + + it("embeds cite markers as links and leaves fenced code unchanged", () => { + const markdown = chatCitationModel.embedCitationMarkersAsLinks( + [ + "Revenue grew [[cite:1]] and later [[cite:2]].", + "", + "```ts", + "const value = '[[cite:1]]'", + "```", + "", + "Legacy [Source 1: spacex-s1.pdf / Page 26] remains.", + ].join("\n"), + [pageCitation, samePageCitation], + { doc_1: "spacex-s1.pdf" }, + ) + + expect(markdown).toContain( + "[spacex-s1.pdf/p26](knowhere-cite://1)", + ) + expect(markdown).toContain( + "[spacex-s1.pdf/p26](knowhere-cite://2)", + ) + expect(markdown).toContain("const value = '[[cite:1]]'") + expect(markdown).not.toContain("[Source 1:") + expect(markdown).not.toContain("[[cite:1]] and later") + }) + + it("drops unknown cite markers from display markdown", () => { + const markdown = chatCitationModel.embedCitationMarkersAsLinks( + "Unsupported claim [[cite:9]].", + [pageCitation], + { doc_1: "spacex-s1.pdf" }, + ) + + expect(markdown).toBe("Unsupported claim.") + expect(markdown).not.toContain("knowhere-cite://9") + }) + + it("strips markers from export markdown and keeps copy labels", () => { + const message: ChatMessageView = { + id: "assistant_1", + role: "assistant", + content: [ + "Revenue grew [[cite:1]].", + "", + "```ts", + "const value = 1;", + "```", + ].join("\n"), + citations: [pageCitation], + artifacts: [ + { + type: "derived_table", + display: true, + title: "Revenue", + columns: ["Year", "Amount"], + rows: [["2025", "4.4B"]], + }, + ], + } + + expect(chatCitationModel.getExportMarkdown(message)).toBe( + [ + "Revenue grew.", + "", + "```ts", + "const value = 1;", + "```", + "", + "### Revenue", + "| Year | Amount |", + "| --- | --- |", + "| 2025 | 4.4B |", + ].join("\n"), + ) + expect(chatCitationModel.getCopyMarkdown(message, { doc_1: "spacex-s1.pdf" })) + .toContain("Revenue grew spacex-s1.pdf/p26.") + expect(chatCitationModel.getCopyMarkdown(message, { doc_1: "spacex-s1.pdf" })) + .toContain("const value = 1;") + }) + + it("keeps knowhere cite hrefs out of the default markdown URL transform", () => { + expect( + chatCitationModel.transformMarkdownUrl("knowhere-cite://2", () => ""), + ).toBe("knowhere-cite://2") + expect(chatCitationModel.parseKnowhereCiteIndex("knowhere-cite://2")).toBe(2) + expect(chatCitationModel.isKnowhereCiteHref("https://example.com")).toBe( + false, + ) + }) +}) diff --git a/src/components/chat-citation-model.ts b/src/components/chat-citation-model.ts new file mode 100644 index 0000000..c3bac75 --- /dev/null +++ b/src/components/chat-citation-model.ts @@ -0,0 +1,284 @@ +import { workspaceCitationState } from "@/components/workspace-citation-state" +import type { + ChatArtifactView, + ChatCitationView, + ChatMessageView, +} from "@/domains/chat/types" + +export const knowhereCiteHrefPrefix = "knowhere-cite://" + +export type CitationSourceGroupEntry = { + readonly citation: ChatCitationView + readonly citationId: string + readonly citationIndex: number + readonly chipLabel: string + readonly pageNumber: number | null +} + +export type CitationSourceGroup = { + readonly key: string + readonly title: string + readonly documentId: string | null + readonly entries: readonly CitationSourceGroupEntry[] +} + +export const chatCitationModel = { + knowhereCiteHrefPrefix, + embedCitationMarkersAsLinks, + getCitationChipLabel, + getCopyMarkdown, + getExportMarkdown, + getSourceTitle, + groupCitationsByFile, + isKnowhereCiteHref, + parseKnowhereCiteIndex, + stripCitationMarkers, + transformMarkdownUrl, + uniquePageLinkEntries, +} as const + +function embedCitationMarkersAsLinks( + content: string, + citations: readonly ChatCitationView[], + sourceTitlesByDocumentId: Readonly>, +): string { + return mapOutsideCodeFences(content, (text) => + collapseRemovedCitationSpace( + replaceCitationTokens(text, (index) => { + const citation = citations[index - 1] + if (!citation) return "" + return toCiteMarkdownLink( + getCitationChipLabel(citation, sourceTitlesByDocumentId), + index, + ) + }), + ), + ) +} + +function getCopyMarkdown( + message: ChatMessageView, + sourceTitlesByDocumentId: Readonly>, +): string { + const copiedBody = mapOutsideCodeFences(message.content, (text) => + collapseRemovedCitationSpace( + replaceCitationTokens(text, (index) => { + const citation = message.citations?.[index - 1] + if (!citation) return "" + return ` ${getCitationChipLabel(citation, sourceTitlesByDocumentId)}` + }), + ), + ) + + return appendDerivedTables(copiedBody, message.artifacts) +} + +function getExportMarkdown(message: ChatMessageView): string { + return appendDerivedTables(stripCitationMarkers(message.content), message.artifacts) +} + +function stripCitationMarkers(content: string): string { + return mapOutsideCodeFences(content, (text) => + collapseRemovedCitationSpace(replaceCitationTokens(text, () => "")), + ) +} + +function groupCitationsByFile( + messageId: string, + citations: readonly ChatCitationView[], + sourceTitlesByDocumentId: Readonly>, +): readonly CitationSourceGroup[] { + const groups: CitationSourceGroup[] = [] + const groupIndexByKey = new Map() + + for (const [citationIndex, citation] of citations.entries()) { + const title = getSourceTitle(citation, sourceTitlesByDocumentId) + const documentId = getTrimmedField(citation.source.documentId) + const key = documentId + ? `document:${documentId}` + : `file:${citationIndex}:${title}` + const entry: CitationSourceGroupEntry = { + citation, + citationId: `${messageId}:${citationIndex}`, + citationIndex, + chipLabel: getCitationChipLabel(citation, sourceTitlesByDocumentId), + pageNumber: workspaceCitationState.getCitationPageNumber(citation), + } + const existingIndex = groupIndexByKey.get(key) + if (existingIndex === undefined) { + groupIndexByKey.set(key, groups.length) + groups.push({ + key, + title, + documentId, + entries: [entry], + }) + continue + } + + const existing = groups[existingIndex] + if (!existing) continue + groups[existingIndex] = { + ...existing, + entries: [...existing.entries, entry], + } + } + + return groups +} + +function uniquePageLinkEntries( + entries: readonly CitationSourceGroupEntry[], +): readonly CitationSourceGroupEntry[] { + const seenPageNumbers = new Set() + const uniqueEntries: CitationSourceGroupEntry[] = [] + + for (const entry of entries) { + if (entry.pageNumber === null || seenPageNumbers.has(entry.pageNumber)) { + continue + } + seenPageNumbers.add(entry.pageNumber) + uniqueEntries.push(entry) + } + + return uniqueEntries +} + +function getCitationChipLabel( + citation: ChatCitationView, + sourceTitlesByDocumentId: Readonly>, +): string { + const title = getSourceTitle(citation, sourceTitlesByDocumentId) + const pageNumber = workspaceCitationState.getCitationPageNumber(citation) + return pageNumber === null ? title : `${title}/p${pageNumber}` +} + +function getSourceTitle( + citation: ChatCitationView, + sourceTitlesByDocumentId: Readonly>, +): string { + const documentId = getTrimmedField(citation.source.documentId) + const sourceTitle = documentId + ? getTrimmedField(sourceTitlesByDocumentId[documentId]) + : null + if (sourceTitle) return sourceTitle + + const sourceFileName = getTrimmedField(citation.source.sourceFileName) + if (sourceFileName && !isGeneratedKnowhereFileName(sourceFileName)) { + return sourceFileName + } + + return "Source" +} + +function isKnowhereCiteHref(href: string | undefined): boolean { + return typeof href === "string" && href.startsWith(knowhereCiteHrefPrefix) +} + +function parseKnowhereCiteIndex(href: string | undefined): number | null { + if (!href || !href.startsWith(knowhereCiteHrefPrefix)) return null + const value = href.slice(knowhereCiteHrefPrefix.length) + const index = Number.parseInt(value, 10) + if (!Number.isSafeInteger(index) || index < 1) return null + return index +} + +function transformMarkdownUrl(value: string, transformDefault: (next: string) => string): string { + if (isKnowhereCiteHref(value)) return value + return transformDefault(value) +} + +function replaceCitationTokens( + text: string, + replaceIndex: (index: number) => string, +): string { + return text + .replace(/\[\[cite:(\d+)\]\]/g, (_match, rawIndex: string) => { + const index = Number.parseInt(rawIndex, 10) + if (!Number.isSafeInteger(index) || index < 1) return "" + return replaceIndex(index) + }) + .replace(/\[Source\s+(\d+)\s*:\s*[^\]]*\]/g, (_match, rawIndex: string) => { + const index = Number.parseInt(rawIndex, 10) + if (!Number.isSafeInteger(index) || index < 1) return "" + return replaceIndex(index) + }) +} + +function toCiteMarkdownLink(label: string, index: number): string { + return `[${escapeMarkdownLinkLabel(label)}](${knowhereCiteHrefPrefix}${index})` +} + +function escapeMarkdownLinkLabel(label: string): string { + return label.replace(/\\/g, "\\\\").replace(/\[/g, "\\[").replace(/\]/g, "\\]") +} + +function collapseRemovedCitationSpace(text: string): string { + return text.replaceAll(" ", " ").replace(/ +([.,;:])/g, "$1").trimEnd() +} + +function mapOutsideCodeFences( + content: string, + mapText: (text: string) => string, +): string { + const lines = content.split("\n") + const output: string[] = [] + let fenceMarker: string | null = null + + for (const line of lines) { + if (fenceMarker) { + output.push(line) + if (line.startsWith(fenceMarker)) fenceMarker = null + continue + } + + const fenceMatch = /^(```+|~~~+)/.exec(line) + if (fenceMatch) { + fenceMarker = fenceMatch[1] ?? null + output.push(line) + continue + } + + output.push(mapText(line)) + } + + return output.join("\n") +} + +function appendDerivedTables( + markdown: string, + artifacts: readonly ChatArtifactView[] | undefined, +): string { + const tables = (artifacts ?? []).flatMap((artifact, index) => { + if (artifact.display === false || artifact.type !== "derived_table") return [] + if (!artifact.title || !artifact.columns || !artifact.rows) return [] + return [formatDerivedTable(artifact, index)] + }) + if (tables.length === 0) return markdown + const body = markdown.trimEnd() + return body.length > 0 ? `${body}\n\n${tables.join("\n\n")}` : tables.join("\n\n") +} + +function formatDerivedTable( + artifact: ChatArtifactView, + index: number, +): string { + const columns = artifact.columns ?? [] + const header = `| ${columns.join(" | ")} |` + const separator = `| ${columns.map(() => "---").join(" | ")} |` + const rows = (artifact.rows ?? []).map( + (row) => `| ${columns.map((_, columnIndex) => row[columnIndex] ?? "").join(" | ")} |`, + ) + return [`### ${artifact.title ?? `Table ${index + 1}`}`, header, separator, ...rows].join( + "\n", + ) +} + +function getTrimmedField(value: string | null | undefined): string | null { + const trimmed = value?.trim() ?? "" + return trimmed.length > 0 ? trimmed : null +} + +function isGeneratedKnowhereFileName(value: string): boolean { + return /^document-[A-Za-z0-9_-]{16,}\.[A-Za-z0-9]+$/u.test(value) +} diff --git a/src/components/chat-message-export-pdf.test.ts b/src/components/chat-message-export-pdf.test.ts new file mode 100644 index 0000000..8bbe53c --- /dev/null +++ b/src/components/chat-message-export-pdf.test.ts @@ -0,0 +1,16 @@ +import { describe, expect, it } from "vitest" + +import { createPdfBlobFromMarkdown } from "./chat-message-export-pdf" + +describe("createPdfBlobFromMarkdown", () => { + it("builds a PDF without citation markers", async () => { + const blob = createPdfBlobFromMarkdown("Revenue grew in the quarter.") + const text = await blob.text() + + expect(blob.type).toBe("application/pdf") + expect(text.startsWith("%PDF-1.4")).toBe(true) + expect(text).toContain("Revenue grew in the quarter.") + expect(text).not.toContain("[[cite:") + expect(text).not.toContain("SOURCES") + }) +}) diff --git a/src/components/chat-message-export-pdf.ts b/src/components/chat-message-export-pdf.ts new file mode 100644 index 0000000..4c4aa3b --- /dev/null +++ b/src/components/chat-message-export-pdf.ts @@ -0,0 +1,136 @@ +const pdfPageWidth = 612 +const pdfPageHeight = 792 +const pdfMargin = 54 +const pdfFontSize = 11 +const pdfLineHeight = 14 +const pdfCharsPerLine = 90 + +export async function downloadAnswerPdf( + filename: string, + markdown: string, +): Promise { + const blob = createPdfBlobFromMarkdown(markdown) + downloadBlob(filename, blob) +} + +export function createPdfBlobFromMarkdown(markdown: string): Blob { + const lines = wrapPdfLines(toWinAnsi(markdown), pdfCharsPerLine) + const linesPerPage = Math.max( + 1, + Math.floor((pdfPageHeight - pdfMargin * 2) / pdfLineHeight), + ) + const pageCount = Math.max(1, Math.ceil(lines.length / linesPerPage)) + const objects: string[] = [] + const pageObjectNumbers: number[] = [] + const fontObjectNumber = 3 + pageCount * 2 + + for (let pageIndex = 0; pageIndex < pageCount; pageIndex += 1) { + const pageLines = lines.slice( + pageIndex * linesPerPage, + (pageIndex + 1) * linesPerPage, + ) + const content = buildPageContent(pageLines) + const contentObjectNumber = 3 + pageIndex * 2 + const pageObjectNumber = contentObjectNumber + 1 + objects.push(pdfObject(contentObjectNumber, `<< /Length ${content.length} >>\nstream\n${content}\nendstream`)) + objects.push( + pdfObject( + pageObjectNumber, + `<< /Type /Page /Parent 2 0 R /MediaBox [0 0 ${pdfPageWidth} ${pdfPageHeight}] /Contents ${contentObjectNumber} 0 R /Resources << /Font << /F1 ${fontObjectNumber} 0 R >> >> >>`, + ), + ) + pageObjectNumbers.push(pageObjectNumber) + } + + objects.unshift( + pdfObject(1, "<< /Type /Catalog /Pages 2 0 R >>"), + pdfObject( + 2, + `<< /Type /Pages /Count ${pageCount} /Kids [${pageObjectNumbers + .map((objectNumber) => `${objectNumber} 0 R`) + .join(" ")}] >>`, + ), + ) + objects.push( + pdfObject(fontObjectNumber, "<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>"), + ) + + return buildPdf(objects) +} + +function buildPageContent(lines: readonly string[]): string { + const commands = [ + "BT", + `/F1 ${pdfFontSize} Tf`, + `${pdfMargin} ${pdfPageHeight - pdfMargin} Td`, + `${pdfLineHeight} TL`, + ] + for (const [index, line] of lines.entries()) { + commands.push(`(${escapePdfLiteral(line)}) Tj`) + if (index < lines.length - 1) commands.push("T*") + } + commands.push("ET") + return commands.join("\n") +} + +function wrapPdfLines(text: string, maxChars: number): string[] { + const lines: string[] = [] + for (const paragraph of text.split(/\r?\n/)) { + if (paragraph.length === 0) { + lines.push("") + continue + } + + let remaining = paragraph + while (remaining.length > maxChars) { + let breakAt = remaining.lastIndexOf(" ", maxChars) + if (breakAt < 1) breakAt = maxChars + lines.push(remaining.slice(0, breakAt)) + remaining = remaining.slice(breakAt).trimStart() + } + lines.push(remaining) + } + return lines.length > 0 ? lines : [""] +} + +function toWinAnsi(text: string): string { + return Array.from(text, (character) => + character.charCodeAt(0) < 256 ? character : "?", + ).join("") +} + +function escapePdfLiteral(text: string): string { + return text.replace(/\\/g, "\\\\").replace(/\(/g, "\\(").replace(/\)/g, "\\)") +} + +function pdfObject(objectNumber: number, body: string): string { + return `${objectNumber} 0 obj\n${body}\nendobj` +} + +function buildPdf(objects: readonly string[]): Blob { + const xrefOffsets: number[] = [0] + let body = "%PDF-1.4\n" + for (const object of objects) { + xrefOffsets.push(body.length) + body += `${object}\n` + } + const xrefStart = body.length + const xrefEntries = [ + "xref", + `0 ${objects.length + 1}`, + "0000000000 65535 f ", + ...xrefOffsets.slice(1).map((offset) => `${String(offset).padStart(10, "0")} 00000 n `), + ] + body += `${xrefEntries.join("\n")}\n` + body += `trailer\n<< /Size ${objects.length + 1} /Root 1 0 R >>\nstartxref\n${xrefStart}\n%%EOF\n` + return new Blob([body], { type: "application/pdf" }) +} + +function downloadBlob(filename: string, blob: Blob): void { + const objectUrl = URL.createObjectURL(blob) + const anchor = document.createElement("a") + anchor.href = objectUrl + anchor.download = filename + anchor.click() + URL.revokeObjectURL(objectUrl) +} diff --git a/src/components/chat-message-list.test.ts b/src/components/chat-message-list.test.ts index 9444128..e59fd54 100644 --- a/src/components/chat-message-list.test.ts +++ b/src/components/chat-message-list.test.ts @@ -70,13 +70,14 @@ describe("ChatMessageList", () => { id: "assistant_1", role: "assistant", content: [ - "Capital expenditure appears in the appendix. [Source 1: spacex-s1.pdf / Assets / tables / table-25 Capital Expenditures.html]", - "Drivers are discussed elsewhere. [Source 3: spacex-s1.pdf / MD&A / Drivers of Our Performance]", + "Capital expenditure appears in the appendix. [[cite:1]]", + "Drivers are discussed elsewhere. [[cite:3]]", ].join("\n\n"), citations: [ { chunkType: "table", score: 0.9, + pageCitationPageNumber: 25, source: { documentId: "doc_1", sourceFileName: "spacex-s1.pdf", @@ -87,6 +88,7 @@ describe("ChatMessageList", () => { { chunkType: "table", score: 0.91, + pageCitationPageNumber: 25, source: { documentId: "doc_1", sourceFileName: "spacex-s1.pdf", @@ -97,6 +99,7 @@ describe("ChatMessageList", () => { { chunkType: "text", score: 0.8, + pageCitationPageNumber: 40, source: { documentId: "doc_1", sourceFileName: "spacex-s1.pdf", @@ -110,16 +113,27 @@ describe("ChatMessageList", () => { }), ); - expect(screen.getByText("Capital expenditure appears in the appendix.")) + expect(screen.getByText(/Capital expenditure appears in the appendix./u)) .toBeTruthy(); expect(screen.queryByText(/Source 1/u)).toBeNull(); - expect(screen.queryByText(/Source 3/u)).toBeNull(); + expect(screen.queryByText(/\[\[cite:/u)).toBeNull(); expect(screen.getByText("Sources")).toBeTruthy(); - const sourceChips = screen.getAllByRole("button", { - name: "Open source spacex-s1.pdf", - }); + const sourceChips = screen.getAllByTestId("citation-chip"); expect(sourceChips).toHaveLength(2); + expect(sourceChips[0]?.textContent).toBe("spacex-s1.pdf/p25"); + expect(sourceChips[1]?.textContent).toBe("spacex-s1.pdf/p40"); + expect(sourceChips[0]?.className).toContain("bg-muted"); + expect(sourceChips[0]?.className).toContain("h-5"); + expect( + screen.getByRole("button", { name: "Open source spacex-s1.pdf" }), + ).toBeTruthy(); + expect( + screen.getAllByRole("button", { name: "Open page 25 of spacex-s1.pdf" }), + ).toHaveLength(1); + expect( + screen.getByRole("button", { name: "Open page 40 of spacex-s1.pdf" }), + ).toBeTruthy(); await user.hover(sourceChips[0]!); @@ -161,7 +175,7 @@ describe("ChatMessageList", () => { ); const citationButton = screen.getByRole("button", { - name: "Open source report.pdf", + name: "Open page 4 of report.pdf", }); expect( screen.queryByRole("link", { @@ -179,7 +193,7 @@ describe("ChatMessageList", () => { ); }); - it("removes description-only source labels without changing other markdown whitespace", () => { + it("renders inline chips for leftover Source tokens without rewriting fenced code", () => { render( React.createElement(ChatMessageList, { messages: [ @@ -198,6 +212,7 @@ describe("ChatMessageList", () => { chunkType: "text", score: 0.9, description: "revenue growth", + pageCitationPageNumber: 2, source: { documentId: "doc_1", sourceFileName: "notes.pdf", @@ -210,13 +225,56 @@ describe("ChatMessageList", () => { }), ); - expect(screen.getByText("Revenue improved.")).toBeTruthy(); + expect(screen.getByText(/Revenue improved/u)).toBeTruthy(); expect(screen.queryByText(/Source 1/u)).toBeNull(); + expect(screen.getByRole("button", { name: "Open source notes.pdf/p2" })) + .toBeTruthy(); expect(document.querySelector("code.language-ts")?.textContent).toContain( "const value = 1;", ); }); + it("keeps two same-page citation chips as separate buttons", () => { + render( + React.createElement(ChatMessageList, { + messages: [ + { + id: "assistant_1", + role: "assistant", + content: "First claim [[cite:1]] and second claim [[cite:2]].", + citations: [ + { + chunkType: "page", + score: 0.9, + pageCitationPageNumber: 2, + source: { + documentId: "doc_1", + sourceFileName: "spacex-s1.pdf", + sectionPath: "Page 2", + }, + }, + { + chunkType: "page", + score: 0.88, + pageCitationPageNumber: 2, + source: { + documentId: "doc_1", + sourceFileName: "spacex-s1.pdf", + sectionPath: "Page 2", + }, + }, + ], + }, + ], + }), + ); + + expect( + screen.getAllByRole("button", { name: "Open source spacex-s1.pdf/p2" }), + ).toHaveLength(2); + expect(screen.getAllByTestId("citation-chip")).toHaveLength(2); + }); + it("preserves repeated spaces when there are no citation tokens to remove", () => { render( React.createElement(ChatMessageList, { @@ -533,4 +591,50 @@ describe("ChatMessageList", () => { within(screen.getByTestId("chat-scroll")).getByText("What changed?"), ).toBeTruthy(); }); + + it("copies visible answer text and keeps export markdown citation-stripped", async () => { + const user = userEvent.setup(); + const writeText = vi.fn().mockResolvedValue(undefined); + Object.defineProperty(navigator, "clipboard", { + configurable: true, + value: { writeText }, + }); + + render( + React.createElement(ChatMessageList, { + messages: [ + { + id: "assistant_1", + role: "assistant", + content: "Revenue grew [[cite:1]].", + citations: [ + { + chunkType: "page", + score: 0.9, + pageCitationPageNumber: 26, + source: { + documentId: "doc_1", + sourceFileName: "spacex-s1.pdf", + sectionPath: "Page 26", + }, + }, + ], + }, + ], + sourceTitlesByDocumentId: { doc_1: "spacex-s1.pdf" }, + }), + ); + + await user.click(screen.getByRole("button", { name: "Copy answer" })); + expect(writeText).toHaveBeenCalledWith( + expect.stringContaining("Revenue grew spacex-s1.pdf/p26."), + ); + expect(writeText.mock.calls[0]?.[0]).not.toContain("[[cite:"); + expect( + screen.getByRole("button", { name: "Download answer as Markdown" }), + ).toBeTruthy(); + expect( + screen.getByRole("button", { name: "Download answer as PDF" }), + ).toBeTruthy(); + }); }); diff --git a/src/components/chat-message-list.tsx b/src/components/chat-message-list.tsx index 38c5059..64ab794 100644 --- a/src/components/chat-message-list.tsx +++ b/src/components/chat-message-list.tsx @@ -2,16 +2,18 @@ import { type CSSProperties, type ReactElement } from "react"; import { type VirtualItem } from "@tanstack/react-virtual"; -import { ImageIcon, MessageCircle } from "lucide-react"; +import { Copy, Download, FileText, ImageIcon, MessageCircle } from "lucide-react"; import ReactMarkdown, { defaultUrlTransform, type Components, } from "react-markdown"; import remarkGfm from "remark-gfm"; +import { chatCitationModel } from "@/components/chat-citation-model"; import { ChatDiagramCard } from "@/components/chat-diagram-card"; import { useChatMessageListWorkflow } from "@/components/chat-message-list-workflow"; import { chatPanelModel } from "@/components/chat-panel-model"; +import { Button } from "@/components/ui/button"; import { ScrollArea } from "@/components/ui/scroll-area"; import { Spinner } from "@/components/ui/spinner"; import { @@ -27,14 +29,11 @@ import type { ChatMessageView, } from "@/domains/chat/types"; -type DisplayCitation = { +type DisplayImageCitation = { readonly citation: ChatCitationView; readonly citationId: string; readonly label: string; readonly tooltipLabel: string; -}; - -type DisplayImageCitation = DisplayCitation & { readonly assetUrl: string; }; @@ -306,10 +305,6 @@ function MessageBubble({ ); } - const displayCitations = getDisplayCitations( - message, - sourceTitlesByDocumentId, - ); const displayImageArtifacts = getDisplayImageArtifacts( message, sourceTitlesByDocumentId, @@ -319,7 +314,7 @@ function MessageBubble({ ? displayImageArtifacts : getDisplayImageCitations(message, sourceTitlesByDocumentId); const displayDerivedTables = getDisplayDerivedTableArtifacts(message); - const citationContentMarkdown = buildCitationContentMarkdown( + const citationContentMarkdown = chatCitationModel.embedCitationMarkersAsLinks( message.content, message.citations ?? [], sourceTitlesByDocumentId, @@ -327,9 +322,14 @@ function MessageBubble({ return (
+
{displayDerivedTables.length > 0 && (
@@ -371,11 +371,17 @@ function MessageBubble({
)}
+ +
); } @@ -418,78 +424,6 @@ function AssistantDiagram({ ); } -function buildCitationContentMarkdown( - content: string, - citations: readonly ChatCitationView[], - sourceTitlesByDocumentId: Readonly>, -): string { - let rewrittenContent = content; - - for (const [index, citation] of citations.entries()) { - const displayCitation = { - citation, - citationId: "", - label: getCitationSourceChipLabel(citation, sourceTitlesByDocumentId), - tooltipLabel: chatPanelModel.getCitationLabel( - citation, - sourceTitlesByDocumentId, - ), - }; - for (const token of getInlineCitationTokens(displayCitation, index)) { - if (!rewrittenContent.includes(token)) continue; - - rewrittenContent = removeInlineCitationToken(rewrittenContent, token); - } - } - - return rewrittenContent; -} - -function removeInlineCitationToken(content: string, token: string): string { - return content - .replaceAll(` ${token}`, "") - .replaceAll(`${token} `, "") - .replaceAll(token, ""); -} - -function getInlineCitationTokens( - displayCitation: DisplayCitation, - index: number, -): readonly string[] { - const label = displayCitation.label; - const slashLabel = label.replace(/\s+·\s+/gu, " / "); - const sourceName = getCitationSourceName(label); - const sectionPath = getTrimmedCitationField( - displayCitation.citation.source.sectionPath, - ); - const description = getTrimmedCitationField(displayCitation.citation.description); - const sectionLabel = sectionPath ? `${sourceName} / ${sectionPath}` : null; - const descriptionLabel = description ? `${sourceName} / ${description}` : null; - const citationNumber = index + 1; - const tokens = [ - `[${label}]`, - `[${slashLabel}]`, - `[Source ${citationNumber}: ${label}]`, - `[Source ${citationNumber}: ${slashLabel}]`, - sectionLabel ? `[${sectionLabel}]` : null, - sectionLabel ? `[Source ${citationNumber}: ${sectionLabel}]` : null, - descriptionLabel ? `[${descriptionLabel}]` : null, - descriptionLabel ? `[Source ${citationNumber}: ${descriptionLabel}]` : null, - description ? `[${description}]` : null, - description ? `[Source ${citationNumber}: ${description}]` : null, - ]; - - return Array.from( - new Set(tokens.filter((token): token is string => Boolean(token))), - ).sort((left, right): number => right.length - left.length); -} - -function getCitationSourceName(label: string): string { - const [sourceName] = label.split(/\s+·\s+/u); - const normalized = sourceName?.trim(); - return normalized && normalized.length > 0 ? normalized : label; -} - function DerivedTableArtifactView({ artifact, }: { @@ -540,12 +474,49 @@ function DerivedTableArtifactView({ function AssistantMessageContent({ content, + message, + onCitationClick, + pendingCitationId, + sourceTitlesByDocumentId, }: { readonly content: string; + readonly message: ChatMessageView; + readonly onCitationClick?: ( + citation: ChatCitationView, + citationId: string, + ) => void; + readonly pendingCitationId?: string | null; + readonly sourceTitlesByDocumentId: Readonly>; }): ReactElement { const markdownComponents: Components = { ...assistantMarkdownComponents, a: ({ href, children }) => { + const citationIndex = chatCitationModel.parseKnowhereCiteIndex(href); + if (citationIndex !== null) { + const citation = (message.citations ?? [])[citationIndex - 1]; + if (!citation) return null; + const citationId = chatPanelModel.getCitationId( + message.id, + citationIndex - 1, + ); + return ( + + ); + } + return ( void; readonly pendingCitationId?: string | null; + readonly sourceTitlesByDocumentId: Readonly>; }): ReactElement | null { - if (displayCitations.length === 0) return null; + const groups = chatCitationModel.groupCitationsByFile( + message.id, + message.citations ?? [], + sourceTitlesByDocumentId, + ); + if (groups.length === 0) return null; return (
@@ -593,26 +571,148 @@ function AssistantSources({ Sources

-
- {displayCitations.map((displayCitation) => ( - +
    + {groups.map((group, groupIndex) => ( +
  1. + + {groupIndex + 1}. + + entry.citationId === pendingCitationId, + )} + onCitationClick={onCitationClick} + /> + {chatCitationModel.uniquePageLinkEntries(group.entries).map( + (entry) => ( + + ), + )} +
    2. ))} -
+
); } +function CitationSourceTitle({ + group, + isPending, + onCitationClick, +}: { + readonly group: ReturnType[number]; + readonly isPending: boolean; + readonly onCitationClick?: ( + citation: ChatCitationView, + citationId: string, + ) => void; +}): ReactElement { + const firstEntry = group.entries[0]; + if (!firstEntry) { + return {group.title}; + } + + return ( + + ); +} + +function AssistantMessageActions({ + message, + sourceTitlesByDocumentId, +}: { + readonly message: ChatMessageView; + readonly sourceTitlesByDocumentId: Readonly>; +}): ReactElement { + const copyMarkdown = chatCitationModel.getCopyMarkdown( + message, + sourceTitlesByDocumentId, + ); + const exportMarkdown = chatCitationModel.getExportMarkdown(message); + + return ( +
+ + + +
+ ); +} + +function downloadTextFile(filename: string, contents: string): void { + const blob = new Blob([contents], { type: "text/markdown;charset=utf-8" }); + const objectUrl = URL.createObjectURL(blob); + const anchor = document.createElement("a"); + anchor.href = objectUrl; + anchor.download = filename; + anchor.click(); + URL.revokeObjectURL(objectUrl); +} + function transformAssistantMarkdownUrl(value: string): string { - return defaultUrlTransform(value); + return chatCitationModel.transformMarkdownUrl(value, defaultUrlTransform); } function CitationChip({ @@ -634,15 +734,17 @@ function CitationChip({ ) => void; }): ReactElement { return ( - + - ) : null} + {workspaceSources.length === 0 ? ( diff --git a/src/components/top-nav.tsx b/src/components/top-nav.tsx index 5c32df8..641fcf8 100644 --- a/src/components/top-nav.tsx +++ b/src/components/top-nav.tsx @@ -34,7 +34,7 @@ export function TopNav({

- Knowhere Notebook + Knowhere Brain

{ await result.current.handleCitationClick(pageCitation, "message_1:0"); }); - expect(fetchChunks).toHaveBeenCalledWith("source_1"); + expect(fetchChunks).toHaveBeenCalledWith("source_1", { + chunkType: "page", + untilPageNumber: 4, + }); expect(selectSource).toHaveBeenCalledWith("source_1"); expect(result.current.focusedChunk.chunkId).toBe("page_4"); expect(result.current.focusedPage).toEqual({ - pageNumber: null, - requestId: 0, + pageNumber: 4, + requestId: 1, + citationId: "message_1:0", }); expect(result.current.citationListViewRequestId).toBe(1); }); + it("uses cached SWR pages for file B while viewing A without loading all chunk types", async () => { + const pageChunk: ParsedChunkView = { + chunkId: "page_4", + documentId: "document_1", + type: "page", + content: "Page 4 summary", + sourceTitle: "Contract.pdf", + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004.png", + contentType: "image/png", + }, + ], + }; + const fetchChunks = vi.fn(async () => [pageChunk]); + const selectSource = vi.fn(); + const viewingSource: SourceView = { + id: "source_2", + title: "Other.pdf", + mimeType: "application/pdf", + status: "ready", + documentId: "document_2", + }; + const pageAssetSource: SourceView = { + ...readySource, + documentPresentation: { kind: "page-assets", pageCount: 8 }, + }; + const pageCitation: ChatCitationView = { + chunkType: "page", + score: 0.9, + pageCitationAssetUrl: "https://assets.example/page-000004.png", + pageCitationPageNumber: 4, + source: { + documentId: "document_1", + sourceFileName: "Contract.pdf", + sectionPath: "Page 4", + }, + }; + const cache = new Map([ + [ + unstable_serialize(["source-chunks", "source_1", 1]), + { + data: { + chunks: [pageChunk], + pagination: { + page: 1, + pageSize: 50, + total: 1, + totalPages: 1, + }, + }, + }, + ], + ]); + + const { result } = renderHook( + () => + useWorkspaceCitationFocus({ + fetchChunks, + onSelectSource: selectSource, + selectedSourceId: "source_2", + sources: [pageAssetSource, viewingSource], + }), + { + wrapper: ({ children }: { readonly children: ReactNode }) => + React.createElement( + SWRConfig, + { value: { provider: () => cache } }, + children, + ), + }, + ); + + await act(async () => { + await result.current.handleCitationClick(pageCitation, "message_1:0"); + }); + + expect(fetchChunks).not.toHaveBeenCalled(); + expect(selectSource).toHaveBeenCalledWith("source_1"); + expect(result.current.focusedChunk.chunkId).toBe("page_4"); + expect(result.current.focusedPage.pageNumber).toBe(4); + expect(result.current.focusedPage.citationId).toBe("message_1:0"); + }); + + it("loads only page-type chunks when the tree asks for a full page-asset source", async () => { + const pageChunk: ParsedChunkView = { + chunkId: "page_4", + documentId: "document_1", + type: "page", + content: "Page 4 summary", + sourceTitle: "Contract.pdf", + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004.png", + contentType: "image/png", + }, + ], + }; + const fetchChunks = vi.fn(async () => [pageChunk]); + const pageAssetSource: SourceView = { + ...readySource, + documentPresentation: { kind: "page-assets", pageCount: 8 }, + }; + + const { result } = renderHook( + () => + useWorkspaceCitationFocus({ + fetchChunks, + onSelectSource: vi.fn(), + selectedSourceId: "source_1", + sources: [pageAssetSource], + }), + { wrapper: createSWRWrapper }, + ); + + await act(async () => { + result.current.handleLoadAllChunks(); + }); + + await waitFor(() => { + expect(fetchChunks).toHaveBeenCalledWith("source_1", { chunkType: "page" }); + }); + expect(result.current.prefetchedChunksBySourceId).toEqual({ + source_1: [pageChunk], + }); + }); + it("reuses cached chunks for a different source without refetching", async () => { const fetchChunks = vi.fn(async () => [prefetchedChunk]); const selectSource = vi.fn(); @@ -261,6 +394,51 @@ describe("useWorkspaceCitationFocus", () => { "source_1", ); }); + + it("focuses a second citation on the selected source without refetching", async () => { + const secondChunk: ParsedChunkView = { + chunkId: "chunk_2", + documentId: "document_1", + sectionPath: "Costs", + type: "text", + content: "Costs rose in the quarter.", + sourceTitle: "Contract.pdf", + }; + const secondCitation: ChatCitationView = { + chunkType: "text", + score: 0.9, + content: "Costs rose in the quarter.", + source: { + documentId: "document_1", + sourceFileName: "Contract.pdf", + sectionPath: "Costs", + }, + }; + const fetchChunks = vi.fn(async () => [prefetchedChunk, secondChunk]); + const selectSource = vi.fn(); + + const { result } = renderHook(() => + useWorkspaceCitationFocus({ + fetchChunks, + onSelectSource: selectSource, + selectedSourceId: "source_1", + sources: [readySource], + }), + { wrapper: createSWRWrapper }, + ); + + await act(async () => { + await result.current.handleCitationClick(citation, "message_1:0"); + }); + await act(async () => { + await result.current.handleCitationClick(secondCitation, "message_1:1"); + }); + + expect(fetchChunks).toHaveBeenCalledTimes(1); + expect(fetchChunks).toHaveBeenCalledWith("source_1"); + expect(result.current.focusedChunk.chunkId).toBe("chunk_2"); + expect(result.current.focusedPage.citationId).toBe("message_1:1"); + }); }); function createSWRWrapper({ diff --git a/src/components/workspace-citation-focus.ts b/src/components/workspace-citation-focus.ts index 70b0acf..379dc29 100644 --- a/src/components/workspace-citation-focus.ts +++ b/src/components/workspace-citation-focus.ts @@ -1,12 +1,17 @@ "use client" import { useCallback, useRef, useState } from "react" +import { useSWRConfig } from "swr" import { workspaceCitationState } from "@/components/workspace-citation-state" import { useWorkspaceSelectedChunks } from "@/components/workspace-selected-chunks" import type { ChatCitationView } from "@/domains/chat/types" import type { ParsedChunkView } from "@/domains/chunks/types" import type { SourceView } from "@/domains/sources/types" +import { + type FetchChunksOptions, +} from "@/domains/workspace/client" +import { workspaceClientCache } from "@/domains/workspace/client-cache" type FocusedChunkState = { readonly chunkId: string | null @@ -16,6 +21,7 @@ type FocusedChunkState = { type FocusedPageState = { readonly pageNumber: number | null readonly requestId: number + readonly citationId: string | null } type PrefetchedChunksBySourceId = Readonly> @@ -24,7 +30,10 @@ type PrefetchedChunksUpdater = ( ) => PrefetchedChunksBySourceId type WorkspaceCitationFocusInput = { - readonly fetchChunks: (sourceId: string) => Promise + readonly fetchChunks: ( + sourceId: string, + options?: FetchChunksOptions, + ) => Promise readonly initialPrefetchedChunksBySourceId?: PrefetchedChunksBySourceId readonly onSelectSource: (sourceId: string | null) => void readonly selectedSourceId: string | null @@ -68,6 +77,7 @@ export function useWorkspaceCitationFocus({ const [focusedPage, setFocusedPage] = useState({ pageNumber: null, requestId: 0, + citationId: null, }) const [pendingCitationId, setPendingCitationId] = useState( null, @@ -77,6 +87,8 @@ export function useWorkspaceCitationFocus({ const [fullChunkLoadingSourceId, setFullChunkLoadingSourceId] = useState< string | null >(null) + const { cache, mutate } = useSWRConfig() + const partialPrefetchSourceIdsRef = useRef>(new Set()) const fullChunkRequestsBySourceIdRef = useRef< Map> >(new Map()) @@ -109,12 +121,16 @@ export function useWorkspaceCitationFocus({ }, [], ) - const requestPageFocus = useCallback((pageNumber: number | null): void => { - setFocusedPage((current) => ({ - pageNumber, - requestId: current.requestId + 1, - })) - }, []) + const requestPageFocus = useCallback( + (pageNumber: number | null, citationId: string | null = null): void => { + setFocusedPage((current) => ({ + pageNumber, + citationId, + requestId: current.requestId + 1, + })) + }, + [], + ) const updatePrefetchedChunksBySourceId = useCallback( (updater: PrefetchedChunksUpdater): void => { @@ -147,14 +163,31 @@ export function useWorkspaceCitationFocus({ ) const loadAllChunksForSource = useCallback( - (sourceId: string): Promise => { + ( + sourceId: string, + options?: FetchChunksOptions, + ): Promise => { const existingRequest = fullChunkRequestsBySourceIdRef.current.get(sourceId) if (existingRequest) return existingRequest setFullChunkLoadingSourceId(sourceId) - const request = fetchChunks(sourceId) + const source = sources.find((candidate) => candidate.id === sourceId) + const chunkType = + options?.chunkType ?? + (source?.documentPresentation?.kind === "page-assets" + ? "page" + : undefined) + const fetchOptions: FetchChunksOptions | undefined = chunkType + ? { chunkType, ...options } + : options + const request = ( + fetchOptions + ? fetchChunks(sourceId, fetchOptions) + : fetchChunks(sourceId) + ) .then((chunks) => { + partialPrefetchSourceIdsRef.current.delete(sourceId) updatePrefetchedChunksBySourceId((current) => workspaceCitationState.upsertPrefetchedChunks( current, @@ -162,6 +195,13 @@ export function useWorkspaceCitationFocus({ chunks, ), ) + workspaceClientCache.hydrateSourceChunks( + (key, data) => { + void mutate(key, data, { revalidate: false }) + }, + sourceId, + chunks, + ) return chunks }) .finally(() => { @@ -174,25 +214,32 @@ export function useWorkspaceCitationFocus({ fullChunkRequestsBySourceIdRef.current.set(sourceId, request) return request }, - [fetchChunks, updatePrefetchedChunksBySourceId], + [fetchChunks, mutate, sources, updatePrefetchedChunksBySourceId], ) const handleLoadAllChunks = useCallback((): void => { + if (!selectedSourceId) return + + const hasPartialPrefetch = + partialPrefetchSourceIdsRef.current.has(selectedSourceId) if ( - !selectedSourceId || - prefetchedChunksBySourceIdRef.current[selectedSourceId] || - fullChunkRequestedSourceIdsRef.current.has(selectedSourceId) || - fullChunkRequestsBySourceIdRef.current.has(selectedSourceId) + !hasPartialPrefetch && + (prefetchedChunksBySourceIdRef.current[selectedSourceId] || + fullChunkRequestedSourceIdsRef.current.has(selectedSourceId) || + fullChunkRequestsBySourceIdRef.current.has(selectedSourceId)) ) { return } fullChunkRequestedSourceIdsRef.current.add(selectedSourceId) - void loadAllChunksForSource(selectedSourceId) - }, [ - loadAllChunksForSource, - selectedSourceId, - ]) + const source = sources.find((candidate) => candidate.id === selectedSourceId) + void loadAllChunksForSource( + selectedSourceId, + source?.documentPresentation?.kind === "page-assets" + ? { chunkType: "page" } + : undefined, + ) + }, [loadAllChunksForSource, selectedSourceId, sources]) const handleCitationClick = useCallback( async ( @@ -209,17 +256,62 @@ export function useWorkspaceCitationFocus({ if (!source) return setCitationListViewRequestId((current) => current + 1) + const pageNumber = workspaceCitationState.getCitationPageNumber(citation) + const isPageAssetSource = + source.documentPresentation?.kind === "page-assets" - const loadedChunkId = workspaceCitationState.getLoadedCitationChunkId({ - citation, - selectedSourceId, - sourceId: source.id, - selectedChunks, - hasMoreSelectedChunks, - }) - if (loadedChunkId) { - requestChunkFocus(loadedChunkId) - return + const focusFromChunks = ( + chunks: readonly ParsedChunkView[], + hasMore: boolean, + ): string | null => + workspaceCitationState.getLoadedCitationChunkId({ + citation, + selectedSourceId: source.id, + sourceId: source.id, + selectedChunks: chunks, + hasMoreSelectedChunks: hasMore, + }) + + const applyFocus = (chunkId: string | null): void => { + if (selectedSourceId !== source.id) onSelectSource(source.id) + requestChunkFocus(chunkId) + requestPageFocus(pageNumber, citationId) + } + + if (selectedSourceId === source.id) { + const loadedChunkId = workspaceCitationState.getLoadedCitationChunkId({ + citation, + selectedSourceId, + sourceId: source.id, + selectedChunks, + hasMoreSelectedChunks, + }) + if (loadedChunkId) { + requestChunkFocus(loadedChunkId) + requestPageFocus(pageNumber, citationId) + return + } + } + + const prefetchedChunks = prefetchedChunksBySourceIdRef.current[source.id] + if (prefetchedChunks) { + const prefetchedChunkId = focusFromChunks(prefetchedChunks, false) + if (prefetchedChunkId) { + applyFocus(prefetchedChunkId) + return + } + } + + const swrChunks = workspaceClientCache.getCachedSourceChunks( + cache, + source.id, + ) + if (swrChunks && swrChunks.length > 0) { + const cachedChunkId = focusFromChunks(swrChunks, true) + if (cachedChunkId) { + applyFocus(cachedChunkId) + return + } } if (!workspaceCitationState.hasExactCitationTargetHint(citation)) { @@ -229,43 +321,45 @@ export function useWorkspaceCitationFocus({ ) if (selectedSourceId !== source.id) onSelectSource(source.id) requestChunkFocus(null) + requestPageFocus(null) return } - const cachedChunks = prefetchedChunksBySourceIdRef.current[source.id] - if (cachedChunks) { - const cachedFocusId = - workspaceCitationState.getLoadedCitationChunkId({ - citation, - selectedSourceId: source.id, - sourceId: source.id, - selectedChunks: cachedChunks, - hasMoreSelectedChunks: false, - }) - updatePrefetchedChunksBySourceId((current) => - workspaceCitationState.upsertPrefetchedChunks( - current, + requestChunkFocus(null) + const fetchOptions: FetchChunksOptions | undefined = isPageAssetSource + ? { + chunkType: "page", + ...(pageNumber !== null ? { untilPageNumber: pageNumber } : {}), + } + : undefined + const chunks = + isPageAssetSource && pageNumber !== null + ? await fetchChunks(source.id, fetchOptions) + : await loadAllChunksForSource(source.id, fetchOptions) + + if (isPageAssetSource && pageNumber !== null) { + if (!fullChunkRequestsBySourceIdRef.current.has(source.id)) { + partialPrefetchSourceIdsRef.current.add(source.id) + updatePrefetchedChunksBySourceId((current) => + workspaceCitationState.upsertPrefetchedChunks( + current, + source.id, + chunks, + ), + ) + workspaceClientCache.hydrateSourceChunks( + (key, data) => { + void mutate(key, data, { revalidate: false }) + }, source.id, - cachedChunks, - ), - ) - if (selectedSourceId !== source.id) onSelectSource(source.id) - requestChunkFocus(cachedFocusId) - return + chunks, + ) + } } - requestChunkFocus(null) - const chunks = await loadAllChunksForSource(source.id) - const prefetchedChunkId = - workspaceCitationState.getLoadedCitationChunkId({ - citation, - selectedSourceId: source.id, - sourceId: source.id, - selectedChunks: chunks, - hasMoreSelectedChunks: false, - }) - onSelectSource(source.id) - requestChunkFocus(prefetchedChunkId) + applyFocus( + focusFromChunks(chunks, false), + ) } finally { setPendingCitationId((current) => current === citationId ? null : current, @@ -273,10 +367,14 @@ export function useWorkspaceCitationFocus({ } }, [ + cache, + fetchChunks, hasMoreSelectedChunks, loadAllChunksForSource, + mutate, onSelectSource, requestChunkFocus, + requestPageFocus, selectedChunks, selectedSourceId, sources, diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index 59cf03a..7adb3ac 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -239,7 +239,9 @@ describe("useWorkspaceSelectedChunks", () => { ]), ); expect(result.current.isSelectedChunksLoading).toBe(false); - expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1); + expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1, { + chunkType: "page", + }); }); it("detects page assets from selected source chunks without a second request", async () => { diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index 260a6f5..da357e2 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -45,6 +45,10 @@ export function useWorkspaceSelectedChunks({ selectedSource && selectedSource.status === "ready" ? selectedSource.id : null + const pageChunkType = + selectedSource?.documentPresentation?.kind === "page-assets" + ? "page" + : undefined const { data: selectedChunkPages, isLoading: isChunksLoading, @@ -57,7 +61,12 @@ export function useWorkspaceSelectedChunks({ pageIndex, previousPageData, ), - fetchChunksByKey, + ([, sourceId, page]: SourceChunksKey) => + pageChunkType + ? workspaceClient.fetchChunkPage(sourceId, page, { + chunkType: pageChunkType, + }) + : workspaceClient.fetchChunkPage(sourceId, page), { revalidateIfStale: false, keepPreviousData: false, @@ -129,14 +138,6 @@ export function useWorkspaceSelectedChunks({ } } -function fetchChunksByKey([ - , - sourceId, - page, -]: SourceChunksKey): Promise { - return workspaceClient.fetchChunkPage(sourceId, page) -} - function getResolvedSelectedSource( source: SourceView | undefined, chunks: readonly ParsedChunkView[], diff --git a/src/components/workspace-shell-layout.test.ts b/src/components/workspace-shell-layout.test.ts index d414f1a..f32d24d 100644 --- a/src/components/workspace-shell-layout.test.ts +++ b/src/components/workspace-shell-layout.test.ts @@ -77,6 +77,11 @@ describe("WorkspaceShellLayout", () => { "350px", ) expect(screen.getByTestId("desktop-chat-panel").style.width).toBe("420px") + expect( + within(screen.getByTestId("desktop-sources-panel")).getByRole("button", { + name: "Open library", + }), + ).toBeTruthy() }) it("renders compact sidebars when the side panels are collapsed", () => { @@ -175,6 +180,11 @@ describe("WorkspaceShellLayout", () => { screen.getByRole("button", { name: "Show sources panel" }), ).toBeTruthy() expect(screen.getByRole("button", { name: "Show chat panel" })).toBeTruthy() + expect( + within(screen.getByTestId("desktop-sources-panel")).getByRole("button", { + name: "Open library", + }), + ).toBeTruthy() screen.getByRole("button", { name: "Open source Quarterly Report.pdf" }).click() expect(handleSourceSelected).toHaveBeenCalledWith("source_report") @@ -273,5 +283,10 @@ describe("WorkspaceShellLayout", () => { name: "Delete Very Long Quarterly Report Filename.pdf", }), ).toBeTruthy() + expect( + within(desktopSourcesPanel).getByRole("button", { + name: "Open library", + }), + ).toBeTruthy() }) }) diff --git a/src/components/workspace-shell-layout.tsx b/src/components/workspace-shell-layout.tsx index 72c6724..33c9394 100644 --- a/src/components/workspace-shell-layout.tsx +++ b/src/components/workspace-shell-layout.tsx @@ -1,5 +1,6 @@ import { useCallback, type ReactElement } from "react" import { + BookOpen, Database, FileText, MessageSquare, @@ -44,6 +45,7 @@ type FocusedChunkState = { type FocusedPageState = { readonly pageNumber: number | null readonly requestId: number + readonly citationId?: string | null } type WorkspaceShellUser = { @@ -209,7 +211,9 @@ export function WorkspaceShellLayout( props.onDesktopPanelExpand("sources")} + onLibraryOpen={props.onLibraryOpen} onSourceSelected={props.onSourceSelected} /> ) : ( @@ -287,6 +291,7 @@ export function WorkspaceShellLayout( focusedChunkRequestId={props.focusedChunk.requestId} focusedPageNumber={focusedPage.pageNumber} focusedPageRequestId={focusedPage.requestId} + focusedCitationId={focusedPage.citationId ?? null} isLoading={props.isSelectedChunksLoading} isLoadingAllChunks={props.isSelectedAllChunksLoading} isLoadingMore={props.isSelectedChunksLoadingMore} @@ -432,6 +437,7 @@ export function WorkspaceShellLayout( focusedChunkRequestId={props.focusedChunk.requestId} focusedPageNumber={focusedPage.pageNumber} focusedPageRequestId={focusedPage.requestId} + focusedCitationId={focusedPage.citationId ?? null} isLoading={props.isSelectedChunksLoading} isLoadingAllChunks={props.isSelectedAllChunksLoading} isLoadingMore={props.isSelectedChunksLoadingMore} @@ -564,12 +570,16 @@ function DesktopPanelRestoreButton({ } function CompactSourcesSidebar({ + isLibraryOpen = false, onExpand, + onLibraryOpen, onSourceSelected, selectedSourceId, sources, }: { + readonly isLibraryOpen?: boolean readonly onExpand: () => void + readonly onLibraryOpen?: () => void readonly onSourceSelected: (sourceId: string | null) => void readonly selectedSourceId: string | null readonly sources: readonly SourceView[] @@ -582,6 +592,16 @@ function CompactSourcesSidebar({ onClick={onExpand} />
+ onLibraryOpen?.()} + > + + +
{sources.length === 0 ? ( { ); const desktopChunksPanel = within(screen.getByTestId("desktop-chunks-panel")); + await waitFor(() => { + expect(desktopChunksPanel.getByRole("button", { name: "List" })).toBeTruthy(); + }); + fireEvent.click(desktopChunksPanel.getByRole("button", { name: "List" })); await waitFor(() => { expect(desktopChunksPanel.getByRole("img", { name: "Page 1" })) .toBeTruthy(); @@ -537,7 +541,7 @@ describe("WorkspaceShell", () => { { id: "assistant_1", role: "assistant", - content: "The answer uses two sections.", + content: "The answer uses two sections. [[cite:1]] [[cite:2]]", citations: [ { content: "First cited section", @@ -620,19 +624,17 @@ describe("WorkspaceShell", () => { }); await user.click(sendButton); - await desktopChatPanel.findAllByRole("button", { - name: "Open source doc.pdf", - }); - const citationButtons = desktopChatPanel.getAllByRole( - "button", - { - name: "Open source doc.pdf", - }, - ); - expect(citationButtons).toHaveLength(2); - const firstCitation = citationButtons[0] as HTMLButtonElement; - const secondCitation = citationButtons[1] as HTMLButtonElement; - await user.click(firstCitation); + await desktopChatPanel.findAllByTestId("citation-chip"); + const getCitationChip = (citationId: string): HTMLButtonElement => { + const chip = desktopChatPanel + .getAllByTestId("citation-chip") + .find((element) => element.getAttribute("data-citation-id") === citationId); + expect(chip).toBeTruthy(); + return chip as HTMLButtonElement; + }; + expect(getCitationChip("assistant_1:0")).toBeTruthy(); + expect(getCitationChip("assistant_1:1")).toBeTruthy(); + await user.click(getCitationChip("assistant_1:0")); await waitFor(() => { expect( @@ -650,9 +652,9 @@ describe("WorkspaceShell", () => { expect(countFetches(fetch, "/api/sources/source_1/chunks")).toBe(1); await waitFor(() => { - expect(secondCitation.disabled).toBe(false); + expect(getCitationChip("assistant_1:1").disabled).toBe(false); }); - await user.click(secondCitation); + await user.click(getCitationChip("assistant_1:1")); await waitFor(() => { const topRow = screen diff --git a/src/domains/chat/citations.test.ts b/src/domains/chat/citations.test.ts index b6daa0a..c07a69f 100644 --- a/src/domains/chat/citations.test.ts +++ b/src/domains/chat/citations.test.ts @@ -44,6 +44,25 @@ describe("toChatCitationViews", () => { "https://blob.example/pages/page-000004.png", ) }) + + it("copies the page number from retrieval metadata onto the citation chip", () => { + const citations = toChatCitationViews( + [ + makeRetrievalResult({ + chunkType: "page", + metadata: { page_nums: [26] }, + source: { + documentId: "doc_1", + sourceFileName: "spacex-s1.pdf", + sectionPath: "spacex-s1.pdf / Overview", + }, + }), + ], + "Revenue grew.", + ) + + expect(citations[0]?.pageCitationPageNumber).toBe(26) + }) }) function makeRetrievalResult( diff --git a/src/domains/chat/citations.ts b/src/domains/chat/citations.ts index cdc9dda..ea203f7 100644 --- a/src/domains/chat/citations.ts +++ b/src/domains/chat/citations.ts @@ -2,7 +2,10 @@ import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" import type { ChatCitationView } from "@/domains/chat/types" -import type { PageCitationAssetRetrievalResult } from "./page-citation-assets" +import { + resolvePageCitationPageNumber, + type PageCitationAssetRetrievalResult, +} from "./page-citation-assets" export function toChatCitationViews( results: readonly PageCitationAssetRetrievalResult[], @@ -12,6 +15,7 @@ export function toChatCitationViews( return results.map((result, index) => { const description = descriptionsBySourceNumber.get(index + 1) + const pageCitationPageNumber = resolvePageCitationPageNumber(result) return { content: result.content, chunkType: result.chunkType, @@ -20,8 +24,8 @@ export function toChatCitationViews( ...(result.pageCitationAssetUrl ? { pageCitationAssetUrl: result.pageCitationAssetUrl } : {}), - ...(result.pageCitationPageNumber - ? { pageCitationPageNumber: result.pageCitationPageNumber } + ...(pageCitationPageNumber + ? { pageCitationPageNumber } : {}), ...(description ? { description } : {}), source: { diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index f026dc6..e9556ac 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -639,6 +639,61 @@ describe("answerQuestionWithRetrieval", () => { ]); }); + it("keeps two manifest citations to the same evidence chunk as two answer citations", async () => { + const result = makeRetrievalResult({ + content: "Revenue grew on the same page twice.", + chunkType: "page", + source: { + documentId: "doc_included", + sourceFileName: "spacex-s1.pdf", + sectionPath: "Page 26", + }, + }); + const retrieval = { + query: vi.fn().mockResolvedValue({ + results: [result], + evidenceText: "Revenue page evidence.", + referencedChunks: [], + namespace: "notebook-workspace", + query: "What grew?", + routerUsed: "workflow_single_step", + answerText: null, + }), + }; + const generateAnswer = vi.fn(async ({ searchSources }) => { + await searchSources({ query: "What grew?" }); + return makeHarnessRunResultWithLedger( + "Revenue grew [[cite:1]] and later expanded [[cite:2]].", + { + citations: [ + makeOutputCitation("r1:result:1", result), + makeOutputCitation("r1:result:1", result), + ], + chunks: [makeEvidenceChunkFromRetrievalResult("r1:result:1", result)], + }, + ); + }); + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "What grew?", + namespace: "notebook-workspace", + sources: [makeSource()], + excludedSourceIds: [], + retrieval, + generateAnswer, + messages: [], + }), + ); + + expect(answer.answer).toBe( + "Revenue grew [[cite:1]] and later expanded [[cite:2]].", + ); + expect(answer.citations).toHaveLength(2); + expect(answer.citations[0]?.source.sectionPath).toBe("Page 26"); + expect(answer.citations[1]?.source.sectionPath).toBe("Page 26"); + }); + it("uses Notebook source titles instead of generated Knowhere filenames", async () => { const result = makeRetrievalResult({ source: { @@ -1215,6 +1270,87 @@ describe("answerQuestionWithRetrieval", () => { expect(answer.citations[0]?.pageCitationAssetUrl).toBe(hardenedPageAssetUrl); }); + it("hydrates page numbers for grep citations from the matching parsed chunk", async () => { + const grepChunk = { + ref: "grep1:match:1", + kind: "grep_match" as const, + chunkId: "chunk_financial_summary", + content: "ept percentages and per share data)\nTotal automotive revenues\n17,693", + contentPreview: "ept percentages and per share data)", + chunkType: "page", + score: null, + metadata: { + position: 1, + startOffset: 12, + endOffset: 80, + }, + source: { + documentId: "doc_tsla", + sourceFileName: "TSLA-Q4-2025-Update.pdf", + sectionPath: "FINANCIAL SUMMARY", + }, + }; + const retrieval = { + query: vi.fn(), + }; + const readChunks = vi.fn().mockResolvedValue({ + document: { + documentId: "doc_tsla", + sourceFileName: "TSLA-Q4-2025-Update.pdf", + }, + chunks: [ + { + position: 1, + chunkId: "chunk_financial_summary", + chunkType: "page", + content: "Full financial summary page.", + readableContent: "Full financial summary page.", + sectionPath: "FINANCIAL SUMMARY", + sourceChunkPath: "pages/page-4.md", + pageNumbers: [4, 5], + metadata: { + pageNums: [4, 5], + }, + }, + ], + }); + const generateAnswer = vi.fn(async () => + makeHarnessRunResultWithLedger( + "Automotive revenue was $17,693 million [[cite:1]].", + { + citations: [{ ref: "grep1:match:1" }], + chunks: [grepChunk], + }, + ), + ); + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "Tesla automotive revenue", + namespace: "notebook-workspace", + sources: [ + makeSource({ + id: "source_tsla", + title: "TSLA-Q4-2025-Update.pdf", + knowhereDocumentId: "doc_tsla", + }), + ], + excludedSourceIds: [], + retrieval, + knowledge: { readChunks } as unknown as Knowledge, + generateAnswer, + messages: [], + }), + ); + + expect(readChunks).toHaveBeenCalledWith({ + documentId: "doc_tsla", + chunkId: "chunk_financial_summary", + }); + expect(answer.citations[0]?.pageCitationPageNumber).toBe(4); + expect(answer.citations[0]?.source.sectionPath).toBe("FINANCIAL SUMMARY"); + }); + it("returns only harness-selected artifacts when retrieval has extra media candidates", async () => { const frontAssetUrl = "https://blob.example/images/id-front.jpg"; const backAssetUrl = "https://blob.example/images/id-back.jpg"; diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index f45c88d..9f735e6 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -37,7 +37,10 @@ import { enrichRetrievalResultsWithAssetUrls, removeRetrievedMediaAssetUrls, } from "./media-assets" -import { enrichRetrievalResultsWithPageCitationAssetUrls } from "./page-citation-assets" +import { + enrichRetrievalResultsWithPageCitationAssetUrls, + resolvePageCitationPageNumber, +} from "./page-citation-assets" import type { HardenableRetrievalResult } from "./media-asset-hardening" import { notebookKnowhereTools } from "./knowhere-tools" @@ -233,9 +236,15 @@ export const answerQuestionWithRetrieval = ( finalized: generatedAnswer.trace.finalized, }) - const rawResults = selectCitationRawResults({ - generatedAnswer, - }) + const rawResults = yield* Effect.tryPromise(() => + hydrateMissingCitationPageMetadata({ + results: selectCitationRawResults({ + generatedAnswer, + }), + ledgerChunks: generatedAnswer.trace.ledger.chunks, + knowledge: input.knowledge, + }), + ) if ( rawResults.length === 0 && generatedAnswer.manifest.text.trim().length === 0 && @@ -871,7 +880,6 @@ function mapManifestCitationsToResults( ) const results: RetrievalResult[] = [] - const seenKeys = new Set() for (const citation of result.manifest.citations) { const chunk = @@ -879,12 +887,11 @@ function mapManifestCitationsToResults( resolveChunkForAssetRef(citation.ref, assetsByRef, chunksByRef) if (!chunk) continue - const retrievalResult = toRetrievalResultFromEvidenceChunk(chunk) - const key = getRetrievalResultKey(retrievalResult) - if (seenKeys.has(key)) continue - - seenKeys.add(key) - results.push(retrievalResult) + results.push( + toRetrievalResultFromEvidenceChunk( + mergeChunkPageMetadata(chunk, result.trace.ledger.chunks), + ), + ) if (results.length >= MAX_CITATION_RESULTS) break } @@ -934,7 +941,9 @@ function mapDisplayedManifestArtifactsToResults( resolveChunkForAssetRef(sourceRef, assetsByRef, chunksByRef) if (!chunk) continue - const retrievalResult = toRetrievalResultFromEvidenceChunk(chunk) + const retrievalResult = toRetrievalResultFromEvidenceChunk( + mergeChunkPageMetadata(chunk, result.trace.ledger.chunks), + ) const key = getRetrievalResultKey(retrievalResult) if (seenKeys.has(key)) continue @@ -950,7 +959,9 @@ function mapDisplayedManifestArtifactsToResults( resolveChunkForAssetRef(artifact.ref, assetsByRef, chunksByRef) if (!chunk) continue - const retrievalResult = toRetrievalResultFromEvidenceChunk(chunk) + const retrievalResult = toRetrievalResultFromEvidenceChunk( + mergeChunkPageMetadata(chunk, result.trace.ledger.chunks), + ) const key = getRetrievalResultKey(retrievalResult) if (seenKeys.has(key)) continue @@ -962,6 +973,100 @@ function mapDisplayedManifestArtifactsToResults( return results } +function mergeChunkPageMetadata( + chunk: EvidenceChunk, + ledgerChunks: readonly EvidenceChunk[], +): EvidenceChunk { + if (hasResolvablePageNumber(chunk) || !chunk.chunkId) return chunk + + const donor = ledgerChunks.find( + (candidate) => + candidate.ref !== chunk.ref && + candidate.chunkId === chunk.chunkId && + hasResolvablePageNumber(candidate), + ) + if (!donor?.metadata) return chunk + + return { + ...chunk, + metadata: { ...donor.metadata, ...chunk.metadata }, + } +} + +async function hydrateMissingCitationPageMetadata(input: { + readonly results: readonly RetrievalResult[] + readonly ledgerChunks: readonly EvidenceChunk[] + readonly knowledge: AnswerQuestionInput["knowledge"] +}): Promise { + const results = input.results.map((result) => { + const donor = input.ledgerChunks.find( + (chunk) => + Boolean(result.chunkId) && + chunk.chunkId === result.chunkId && + hasResolvablePageNumber(chunk), + ) + if (!donor?.metadata || resolvePageCitationPageNumber(result)) { + return result + } + return { + ...result, + metadata: { ...donor.metadata, ...result.metadata }, + } + }) + + if (!input.knowledge) return results + + return Promise.all( + results.map((result) => + hydrateResultPageMetadataFromKnowledge(result, input.knowledge), + ), + ) +} + +async function hydrateResultPageMetadataFromKnowledge( + result: RetrievalResult, + knowledge: NonNullable, +): Promise { + if (resolvePageCitationPageNumber(result)) return result + const documentId = result.source.documentId + const chunkId = result.chunkId + if (!documentId || !chunkId) return result + + try { + const response = await knowledge.readChunks({ documentId, chunkId }) + const chunk = response.chunks[0] + if (!chunk) return result + return { + ...result, + metadata: { + ...(chunk.metadata ?? {}), + ...(chunk.pageNumbers && chunk.pageNumbers.length > 0 + ? { pageNums: chunk.pageNumbers } + : {}), + ...result.metadata, + }, + } + } catch { + return result + } +} + +function hasResolvablePageNumber(chunk: EvidenceChunk): boolean { + return ( + resolvePageCitationPageNumber({ + content: chunk.content, + chunkType: chunk.chunkType, + score: chunk.score, + metadata: chunk.metadata, + source: { + documentId: chunk.source.documentId ?? undefined, + sourceFileName: chunk.source.sourceFileName ?? undefined, + sectionPath: chunk.source.sectionPath ?? undefined, + }, + }) !== undefined + ) +} + function toRetrievalResultFromEvidenceChunk( chunk: EvidenceChunk, ): RetrievalResult { diff --git a/src/domains/chat/page-citation-assets.test.ts b/src/domains/chat/page-citation-assets.test.ts index 4553d5d..5b6c2b2 100644 --- a/src/domains/chat/page-citation-assets.test.ts +++ b/src/domains/chat/page-citation-assets.test.ts @@ -125,9 +125,52 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { }) expect(result?.pageCitationAssetUrl).toBeUndefined() + expect(result?.pageCitationPageNumber).toBe(4) }) - it("leaves non-page results unchanged even when they have page metadata", async () => { + it("sets the page number from page_nums even when page assets are missing", async () => { + const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ + results: [ + makeRetrievalResult({ + chunkType: "page", + metadata: { page_nums: "25" }, + source: { + documentId: "doc_1", + sourceFileName: "spacex-s1.pdf", + sectionPath: "spacex-s1.pdf / MD&A", + }, + }), + ], + sources: [makeSource()], + }) + + expect(result?.pageCitationAssetUrl).toBeUndefined() + expect(result?.pageCitationPageNumber).toBe(25) + }) + + it("reads snake_case page_assets for the cited page number", async () => { + const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ + results: [ + makeRetrievalResult({ + chunkType: "page", + metadata: { + page_nums: [11], + page_assets: [ + { + page_num: 11, + artifact_ref: "page_citation_assets/page-11.png", + }, + ], + }, + }), + ], + sources: [makeSource()], + }) + + expect(result?.pageCitationPageNumber).toBe(11) + }) + + it("attaches page numbers to text citations without turning them into page images", async () => { const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ results: [ makeRetrievalResult({ @@ -148,6 +191,7 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { }) expect(result?.pageCitationAssetUrl).toBeUndefined() + expect(result?.pageCitationPageNumber).toBe(4) }) }) diff --git a/src/domains/chat/page-citation-assets.ts b/src/domains/chat/page-citation-assets.ts index 223498b..7a388a7 100644 --- a/src/domains/chat/page-citation-assets.ts +++ b/src/domains/chat/page-citation-assets.ts @@ -48,32 +48,50 @@ export async function enrichRetrievalResultsWithPageCitationAssetUrls({ ) } +export function resolvePageCitationPageNumber( + result: RetrievalResult & { readonly pageCitationPageNumber?: number }, +): number | undefined { + const existing = getPositiveInteger(result.pageCitationPageNumber) + if (existing) return existing + + const pageNumbers = getPageNumbers(result.metadata) + const directAsset = getDirectPageCitationAsset(result, pageNumbers) + return ( + directAsset?.pageNum ?? + pageNumbers[0] ?? + parseSectionPathPageNumber(result.source.sectionPath) ?? + undefined + ) +} + async function enrichRetrievalResultWithPageCitationAssetUrl(input: { readonly result: RetrievalResult readonly sourcesByDocumentId: ReadonlyMap readonly hardenChatAssetUrl?: HardenChatAssetUrl }): Promise { - if (!isPageResult(input.result)) return input.result - const pageNumbers = getPageNumbers(input.result.metadata) const directAsset = getDirectPageCitationAsset(input.result, pageNumbers) - const sourceAssetUrl = await getStoredPageCitationAssetUrl({ - result: input.result, - directAsset, - sourcesByDocumentId: input.sourcesByDocumentId, - hardenChatAssetUrl: input.hardenChatAssetUrl, - }) - if (sourceAssetUrl || directAsset?.pageNum) { - return { - ...input.result, - ...(sourceAssetUrl ? { pageCitationAssetUrl: sourceAssetUrl } : {}), - ...(directAsset?.pageNum - ? { pageCitationPageNumber: directAsset.pageNum } - : {}), - } + const pageCitationPageNumber = resolvePageCitationPageNumber(input.result) + const sourceAssetUrl = isPageResult(input.result) + ? await getStoredPageCitationAssetUrl({ + result: input.result, + directAsset, + sourcesByDocumentId: input.sourcesByDocumentId, + hardenChatAssetUrl: input.hardenChatAssetUrl, + }) + : null + + if (!sourceAssetUrl && pageCitationPageNumber === undefined) { + return input.result } - return input.result + return { + ...input.result, + ...(sourceAssetUrl ? { pageCitationAssetUrl: sourceAssetUrl } : {}), + ...(pageCitationPageNumber + ? { pageCitationPageNumber } + : {}), + } } async function getStoredPageCitationAssetUrl(input: { @@ -105,7 +123,9 @@ function getDirectPageCitationAsset( result: RetrievalResult, pageNumbers: readonly number[], ): PageCitationAssetCandidate | null { - const candidates = parsePageCitationAssetCandidates(result.metadata?.pageAssets) + const candidates = parsePageCitationAssetCandidates( + result.metadata?.pageAssets ?? result.metadata?.page_assets, + ) if (pageNumbers.length > 0) { const matchingCandidates = candidates.filter((candidate) => @@ -124,15 +144,27 @@ function parsePageCitationAssetCandidates( return value.flatMap((item): PageCitationAssetCandidate[] => { if (!isRecord(item)) return [] - const pageNum = getPositiveInteger(item.pageNum) + const pageNum = + getPositiveInteger(item.pageNum) ?? + getPositiveInteger(item.page_num) ?? + getPositiveInteger(item.pageNumber) if (!pageNum) return [] return [ { pageNum, - artifactRef: getTrimmedString(item.artifactRef) ?? undefined, - assetUrl: getTrimmedString(item.assetUrl) ?? undefined, - contentType: getTrimmedString(item.contentType) ?? undefined, + artifactRef: + getTrimmedString(item.artifactRef) ?? + getTrimmedString(item.artifact_ref) ?? + undefined, + assetUrl: + getTrimmedString(item.assetUrl) ?? + getTrimmedString(item.asset_url) ?? + undefined, + contentType: + getTrimmedString(item.contentType) ?? + getTrimmedString(item.content_type) ?? + undefined, }, ] }) @@ -143,23 +175,38 @@ function getPageNumbers( ): readonly number[] { if (!metadata) return [] - const values = [metadata.pageNums, metadata.page_nums, metadata.pageNum] const pageNumbers = new Set() + collectPageNumbers(metadata.pageNums, pageNumbers) + collectPageNumbers(metadata.page_nums, pageNumbers) + collectPageNumbers(metadata.pageNum, pageNumbers) + collectPageNumbers(metadata.page_num, pageNumbers) + return [...pageNumbers].sort((left, right) => left - right) +} - for (const value of values) { - if (Array.isArray(value)) { - for (const item of value) { - const pageNum = getPositiveInteger(item) - if (pageNum) pageNumbers.add(pageNum) - } - continue - } - - const pageNum = getPositiveInteger(value) - if (pageNum) pageNumbers.add(pageNum) +function collectPageNumbers(value: unknown, pageNumbers: Set): void { + if (Array.isArray(value)) { + for (const item of value) collectPageNumbers(item, pageNumbers) + return } - return [...pageNumbers].sort((left, right) => left - right) + if (typeof value === "string" && value.includes(",")) { + for (const part of value.split(",")) collectPageNumbers(part.trim(), pageNumbers) + return + } + + const pageNum = getPositiveInteger(value) + if (pageNum) pageNumbers.add(pageNum) +} + +function parseSectionPathPageNumber( + sectionPath: string | null | undefined, +): number | null { + if (typeof sectionPath !== "string") return null + const match = + /\bpage\s+(\d+)\b/i.exec(sectionPath) ?? + /(?:^|[^\w])p(\d+)(?:[^\w]|$)/i.exec(sectionPath) + if (!match) return null + return getPositiveInteger(match[1]) } function getTrimmedString(value: unknown): string | null { @@ -169,11 +216,14 @@ function getTrimmedString(value: unknown): string | null { } function getPositiveInteger(value: unknown): number | null { - return typeof value === "number" && - Number.isSafeInteger(value) && - value > 0 - ? value - : null + if (typeof value === "number" && Number.isSafeInteger(value) && value > 0) { + return value + } + if (typeof value === "string") { + const parsed = Number.parseInt(value.trim(), 10) + if (Number.isSafeInteger(parsed) && parsed > 0) return parsed + } + return null } function isRecord(value: unknown): value is Record { diff --git a/src/domains/chunks/index.test.ts b/src/domains/chunks/index.test.ts index 586d32b..e4ab72c 100644 --- a/src/domains/chunks/index.test.ts +++ b/src/domains/chunks/index.test.ts @@ -205,6 +205,39 @@ describe("toParsedChunkView", () => { }); }); + it("maps demo API snake_case page citation assets on page chunks", () => { + const chunk = makeDocumentChunk({ + id: "document_page_1", + chunkId: "parser_page_1", + chunkType: "page" as DocumentChunk["chunkType"], + metadata: { + page_assets: [ + { + page_num: 8, + asset_url: + "/api/demo-sources/demo-tsla-q4-2025/assets/page_citation_assets/page-8.png", + content_type: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + }); + + expect(toParsedChunkView(chunk, "manual.pdf", "doc_123")).toMatchObject({ + pageAssets: [ + { + pageNumber: 8, + assetUrl: + "/api/demo-sources/demo-tsla-q4-2025/assets/page_citation_assets/page-8.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }); + }); + it("maps SDK-normalized page number metadata", () => { const chunk = makeDocumentChunk({ metadata: { @@ -445,6 +478,22 @@ describe("loadChunkPageForSource", () => { }); }); + it("reads optional chunkType from the chunks query", () => { + expect( + getChunkPageParams( + new URLSearchParams({ + page: "2", + pageSize: "50", + chunkType: "page", + }), + ), + ).toEqual({ + page: 2, + pageSize: 50, + chunkType: "page", + }); + }); + it("resolves connection targets across already-loaded infinite pages", () => { const chunks = resolveChunkConnectionTargets([ makeParsedChunkView({ diff --git a/src/domains/chunks/index.ts b/src/domains/chunks/index.ts index e78293e..7a38b88 100644 --- a/src/domains/chunks/index.ts +++ b/src/domains/chunks/index.ts @@ -40,9 +40,12 @@ export type LoadChunksOptions = { assetUrlsByFilePath?: Readonly> } +export type ChunkReadType = "text" | "image" | "table" | "page" + export type ChunkPageParams = { page: number pageSize: number + chunkType?: ChunkReadType } export type ChunkPagePagination = { @@ -60,6 +63,7 @@ export type ChunkPage = { export function getChunkPageParams( searchParams: URLSearchParams, ): ChunkPageParams { + const chunkType = parseChunkReadType(searchParams.get("chunkType")) return { page: normalizePositiveInteger(searchParams.get("page"), 1), pageSize: normalizePageSize( @@ -68,7 +72,15 @@ export function getChunkPageParams( defaultChunkPageSize, ), ), + ...(chunkType ? { chunkType } : {}), + } +} + +function parseChunkReadType(value: string | null): ChunkReadType | undefined { + if (value === "text" || value === "image" || value === "table" || value === "page") { + return value } + return undefined } export const loadChunksForSource = ( @@ -130,6 +142,7 @@ export const loadChunkPageForSource = ( client.documents.listChunks(source.knowhereDocumentId!, { page: params.page, pageSize: params.pageSize, + ...(params.chunkType ? { chunkType: params.chunkType } : {}), includeAssetUrls: true, }), ) diff --git a/src/domains/chunks/normalization.ts b/src/domains/chunks/normalization.ts index 3ab6bae..53fe88b 100644 --- a/src/domains/chunks/normalization.ts +++ b/src/domains/chunks/normalization.ts @@ -37,7 +37,10 @@ function createParsedChunkView( const summary = getStringMetadata(input.metadata, "summary") const pageAssets = type === "page" - ? getPageAssetViews(input.metadata["pageAssets"], assetUrl) + ? getPageAssetViews( + input.metadata["pageAssets"] ?? input.metadata["page_assets"], + assetUrl, + ) : undefined return { @@ -72,9 +75,15 @@ function getPageAssetViews( const pageAssets = value.flatMap((item): NonNullable => { if (!isRecord(item)) return [] const pageNumber = - getPositiveInteger(item["pageNum"]) ?? getPositiveInteger(item["pageNumber"]) - const assetUrl = getString(item["assetUrl"]) ?? fallbackAssetUrl - const contentType = getString(item["contentType"]) + getPositiveInteger(item["pageNum"]) ?? + getPositiveInteger(item["pageNumber"]) ?? + getPositiveInteger(item["page_num"]) + const assetUrl = + getString(item["assetUrl"]) ?? + getString(item["asset_url"]) ?? + fallbackAssetUrl + const contentType = + getString(item["contentType"]) ?? getString(item["content_type"]) if (!pageNumber || !assetUrl || !contentType) return [] return [ diff --git a/src/domains/chunks/read.test.ts b/src/domains/chunks/read.test.ts index f63c897..32bf8cd 100644 --- a/src/domains/chunks/read.test.ts +++ b/src/domains/chunks/read.test.ts @@ -69,6 +69,43 @@ describe("readSourceChunkPage", () => { }) }) + it("passes chunkType through to knowledge.readChunks", async () => { + const listChunks = vi.fn() + const readChunks = vi.fn(async () => ({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, + chunks: [ + makeReadChunk({ + chunkType: "page", + filePath: "pages/page-000004.png", + }), + ], + page: 1, + pageSize: 50, + totalChunks: 8, + totalPages: 1, + })) + const knowledge = { readChunks } as unknown as Knowledge + + await readSourceChunkPage({ + client: { documents: { listChunks } }, + knowledge, + source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, + params: { page: 1, pageSize: 50, chunkType: "page" }, + }) + + expect(readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "rev_1", + page: 1, + pageSize: 50, + chunkType: "page", + }) + expect(listChunks).not.toHaveBeenCalled() + }) + it("falls back to Knowhere asset URLs when the parsed-storage probe reads remote chunks", async () => { const readChunks = vi.fn(async () => ({ document: { diff --git a/src/domains/chunks/read.ts b/src/domains/chunks/read.ts index ae474cf..552ba38 100644 --- a/src/domains/chunks/read.ts +++ b/src/domains/chunks/read.ts @@ -12,13 +12,12 @@ import { toParsedChunkViewFromReadChunk, type ChunkPage, type ChunkPageParams, + type ChunkReadType, } from "@/domains/chunks" import type { ParsedChunkView } from "@/domains/chunks/types" const loadAllPageSize = 200 -type DisplayReadChunkType = "text" | "image" | "table" | "page" - type ReadableSource = { readonly documentId: string readonly title: string @@ -32,7 +31,7 @@ type DisplayReadClient = { params: { readonly page: number readonly pageSize: number - readonly chunkType?: DisplayReadChunkType + readonly chunkType?: ChunkReadType readonly includeAssetUrls: true }, ): Promise<{ @@ -65,6 +64,7 @@ export async function readSourceChunkPage(input: { ...(input.source.revisionKey ? { revisionKey: input.source.revisionKey } : {}), page: input.params.page, pageSize: input.params.pageSize, + ...(input.params.chunkType ? { chunkType: input.params.chunkType } : {}), }) if (shouldUseKnowledgeChunkResponse(response)) { @@ -76,6 +76,7 @@ export async function readSourceChunkPage(input: { { page: input.params.page, pageSize: input.params.pageSize, + ...(input.params.chunkType ? { chunkType: input.params.chunkType } : {}), includeAssetUrls: true, }, ) @@ -130,6 +131,7 @@ export async function readAllSourceChunks(input: { readonly client: DisplayReadClient readonly knowledge: Knowledge readonly source: ReadableSource + readonly chunkType?: ChunkReadType }): Promise { const chunks: ParsedChunkView[] = [] let page = 1 @@ -143,6 +145,7 @@ export async function readAllSourceChunks(input: { : {}), page, pageSize: loadAllPageSize, + ...(input.chunkType ? { chunkType: input.chunkType } : {}), }) if (shouldUseKnowledgeChunkResponse(response)) { for (const chunk of response.chunks) { @@ -164,6 +167,7 @@ export async function readAllSourceChunks(input: { { page, pageSize: loadAllPageSize, + ...(input.chunkType ? { chunkType: input.chunkType } : {}), includeAssetUrls: true, }, ) diff --git a/src/domains/sources/route-chunks.ts b/src/domains/sources/route-chunks.ts index c636cb6..9cfc99e 100644 --- a/src/domains/sources/route-chunks.ts +++ b/src/domains/sources/route-chunks.ts @@ -2,7 +2,8 @@ import { Effect } from "effect" import { demoView } from "@/domains/demo/view" import { readAllSourceChunks, readSourceChunkPage } from "@/domains/chunks/read" -import { resolveChunkConnectionTargets } from "@/domains/chunks" +import { resolveChunkConnectionTargets, type ChunkReadType } from "@/domains/chunks" +import type { ParsedChunkView } from "@/domains/chunks/types" import type { DemoChunkPage } from "@/integrations/knowhere-demo" import { logger } from "@/lib/logger" import { routeResult } from "@/lib/route-result" @@ -118,6 +119,9 @@ const loadSourceChunksEffect = ( client: readResources.client, knowledge: readResources.knowledge, source: readableSource, + ...(input.pageParams.chunkType + ? { chunkType: input.pageParams.chunkType } + : {}), }), ).pipe( Effect.map((chunks) => @@ -199,6 +203,9 @@ const loadRemoteChunkPageEffect = ( client: readResources.client, knowledge: readResources.knowledge, source: readableSource, + ...(input.pageParams.chunkType + ? { chunkType: input.pageParams.chunkType } + : {}), }), ).pipe( Effect.map((chunks) => @@ -252,10 +259,13 @@ const loadDemoChunkPageEffect = ( status: "ready" as const, documentId: documentIdOverride ?? page.canonicalDocumentId, } - const chunks = pages.flatMap((demoChunkPage) => - demoChunkPage.chunks.map((chunk) => - demoView.toParsedChunkView(source, chunk), + const chunks = filterChunksByType( + pages.flatMap((demoChunkPage) => + demoChunkPage.chunks.map((chunk) => + demoView.toParsedChunkView(source, chunk), + ), ), + input.pageParams.chunkType, ) return routeResult.ok( @@ -319,6 +329,14 @@ function getErrorMessage(error: unknown): string { return String(error) } +function filterChunksByType( + chunks: readonly ParsedChunkView[], + chunkType: ChunkReadType | undefined, +): ParsedChunkView[] { + if (!chunkType) return [...chunks] + return chunks.filter((chunk) => chunk.type === chunkType) +} + function sourceNotFound(): JsonRouteResult<{ readonly message: string }> { return routeResult.error(404, "Source not found.") } diff --git a/src/domains/workspace/client-cache.test.ts b/src/domains/workspace/client-cache.test.ts index a04080a..e2ee2cc 100644 --- a/src/domains/workspace/client-cache.test.ts +++ b/src/domains/workspace/client-cache.test.ts @@ -52,4 +52,43 @@ describe("workspaceClientCache", () => { messages: [], }); }); + + it("reads cached source chunk pages for citation jumps", () => { + const cache = new Map([ + [ + unstable_serialize(["source-chunks", "source_1", 1]), + { + data: { + chunks: [ + { + chunkId: "page_4", + type: "page", + content: "Page 4", + sourceTitle: "report.pdf", + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 1, + totalPages: 1, + }, + }, + }, + ], + ]) as unknown as Cache; + + expect(workspaceClientCache.getCachedSourceChunks(cache, "source_1")).toEqual( + [ + { + chunkId: "page_4", + type: "page", + content: "Page 4", + sourceTitle: "report.pdf", + }, + ], + ); + expect(workspaceClientCache.getCachedSourceChunks(cache, "source_missing")) + .toBeNull(); + }); }); diff --git a/src/domains/workspace/client-cache.ts b/src/domains/workspace/client-cache.ts index af5baae..1902db9 100644 --- a/src/domains/workspace/client-cache.ts +++ b/src/domains/workspace/client-cache.ts @@ -6,6 +6,7 @@ import type { ChatMessageView, ChatThreadView, } from "@/domains/chat/types" +import type { ParsedChunkView } from "@/domains/chunks/types" import type { SourceView } from "@/domains/sources/types" type SourceChunksResponse = Awaited< @@ -27,6 +28,10 @@ type WorkspaceClientCache = { cache: Cache, threadId: string, ) => ChatThreadDetailResponse | null + readonly getCachedSourceChunks: ( + cache: Cache, + sourceId: string, + ) => ParsedChunkView[] | null readonly getChatThreadKey: (threadId: string) => ChatThreadKey readonly getSourceChunksKey: ( sourceId: string | null, @@ -40,6 +45,11 @@ type WorkspaceClientCache = { pages: readonly SourceChunksResponse[] | undefined, ) => boolean readonly hasPendingSources: (sources: readonly SourceView[]) => boolean + readonly hydrateSourceChunks: ( + mutate: (key: SourceChunksKey, data: SourceChunksResponse) => unknown, + sourceId: string, + chunks: readonly ParsedChunkView[], + ) => void } function getSourceChunksKey( @@ -90,6 +100,57 @@ function hasMoreChunkPage(page: SourceChunksResponse): boolean { return page.pagination.page < page.pagination.totalPages } +function isSourceChunksResponse( + value: unknown, +): value is SourceChunksResponse { + if (!value || typeof value !== "object") return false + return Array.isArray((value as Partial).chunks) +} + +function getCachedSourceChunks( + cache: Cache, + sourceId: string, +): ParsedChunkView[] | null { + const chunks: ParsedChunkView[] = [] + let previous: SourceChunksResponse | null = null + + for (let pageIndex = 0; pageIndex < 200; pageIndex += 1) { + const key = getSourceChunksKey(sourceId, pageIndex, previous) + if (!key) break + + const cachedState = cache.get(unstable_serialize(key)) + const cachedData = cachedState?.data + if (!isSourceChunksResponse(cachedData)) { + return pageIndex === 0 ? null : chunks + } + + chunks.push(...(cachedData.chunks ?? [])) + previous = cachedData + } + + return chunks.length > 0 ? chunks : null +} + +function hydrateSourceChunks( + mutate: (key: SourceChunksKey, data: SourceChunksResponse) => unknown, + sourceId: string, + chunks: readonly ParsedChunkView[], +): void { + const key = getSourceChunksKey(sourceId, 0, null) + if (!key) return + + const total = chunks.length + mutate(key, { + chunks: [...chunks], + pagination: { + page: 1, + pageSize: Math.max(total, 1), + total, + totalPages: 1, + }, + }) +} + function isChatThreadDetailResponse( value: unknown, threadId: string, @@ -102,11 +163,13 @@ function isChatThreadDetailResponse( export const workspaceClientCache: WorkspaceClientCache = { getCachedChatThreadData, + getCachedSourceChunks, getChatThreadKey, getSourceChunksKey, hasLoadedChatThreadData, hasMoreChunkPages, hasPendingSources, + hydrateSourceChunks, } export type { diff --git a/src/domains/workspace/client.test.ts b/src/domains/workspace/client.test.ts index 47878bf..f92cd32 100644 --- a/src/domains/workspace/client.test.ts +++ b/src/domains/workspace/client.test.ts @@ -63,6 +63,48 @@ describe("workspaceClient", () => { }) }) + it("fetches all chunks filtered to page type until the cited page is found", async () => { + mockRouteClient.getJson.mockResolvedValueOnce({ + chunks: [ + { + chunkId: "page_1", + type: "page", + content: "Page 1", + sourceTitle: "report.pdf", + pageAssets: [{ pageNumber: 1, assetUrl: "https://a/1.png", contentType: "image/png" }], + }, + ], + pagination: { page: 1, pageSize: 50, total: 2, totalPages: 2 }, + }) + mockRouteClient.getJson.mockResolvedValueOnce({ + chunks: [ + { + chunkId: "page_4", + type: "page", + content: "Page 4", + sourceTitle: "report.pdf", + pageAssets: [{ pageNumber: 4, assetUrl: "https://a/4.png", contentType: "image/png" }], + }, + ], + pagination: { page: 2, pageSize: 50, total: 2, totalPages: 2 }, + }) + + const chunks = await workspaceClient.fetchChunks("source_1", { + chunkType: "page", + untilPageNumber: 4, + }) + + expect(mockRouteClient.getJson).toHaveBeenNthCalledWith( + 1, + "/api/sources/source_1/chunks?page=1&pageSize=50&chunkType=page", + ) + expect(mockRouteClient.getJson).toHaveBeenNthCalledWith( + 2, + "/api/sources/source_1/chunks?page=2&pageSize=50&chunkType=page", + ) + expect(chunks.map((chunk) => chunk.chunkId)).toEqual(["page_1", "page_4"]) + }) + it("preserves source chunk processing messages", async () => { mockRouteClient.getJson.mockResolvedValue({ chunks: [], diff --git a/src/domains/workspace/client.ts b/src/domains/workspace/client.ts index 7720383..1f6fbf4 100644 --- a/src/domains/workspace/client.ts +++ b/src/domains/workspace/client.ts @@ -24,6 +24,11 @@ const workspaceClientConfig = { sourceChunkPageSize: 50, } as const +export type FetchChunksOptions = { + readonly chunkType?: "text" | "image" | "table" | "page" + readonly untilPageNumber?: number +} + type SourceChunksResponse = { chunks?: ParsedChunkView[] isProcessing?: boolean @@ -107,12 +112,23 @@ export const workspaceClient = { archiveChatThread, } as const -async function fetchChunks(sourceId: string): Promise { +async function fetchChunks( + sourceId: string, + options?: FetchChunksOptions, +): Promise { + if (options?.untilPageNumber != null) { + const { chunks } = await fetchChunkPagesUntil(sourceId, options) + return chunks + } + try { + const searchParams = new URLSearchParams() + if (options?.chunkType) searchParams.set("chunkType", options.chunkType) + const query = searchParams.toString() const body = await workspaceRouteClient.getJson<{ chunks?: ParsedChunkView[] }>( - `/api/sources/${encodeURIComponent(sourceId)}/chunks`, + `/api/sources/${encodeURIComponent(sourceId)}/chunks${query ? `?${query}` : ""}`, ) return Array.isArray(body.chunks) ? body.chunks : [] } catch { @@ -123,11 +139,13 @@ async function fetchChunks(sourceId: string): Promise { async function fetchChunkPage( sourceId: string, page: number, + options?: Pick, ): Promise { const searchParams = new URLSearchParams({ page: String(page), pageSize: String(workspaceClientConfig.sourceChunkPageSize), }) + if (options?.chunkType) searchParams.set("chunkType", options.chunkType) const body = await workspaceRouteClient.getJson( `/api/sources/${encodeURIComponent(sourceId)}/chunks?${searchParams.toString()}`, ) @@ -143,6 +161,49 @@ async function fetchChunkPage( } } +async function fetchChunkPagesUntil( + sourceId: string, + options: FetchChunksOptions, +): Promise<{ + readonly chunks: ParsedChunkView[] + readonly pages: SourceChunksResponse[] +}> { + const pages: SourceChunksResponse[] = [] + const chunks: ParsedChunkView[] = [] + let page = 1 + let totalPages = 1 + + do { + const result = await fetchChunkPage(sourceId, page, { + chunkType: options.chunkType, + }) + pages.push(result) + chunks.push(...(result.chunks ?? [])) + if ( + options.untilPageNumber != null && + chunksContainPage(chunks, options.untilPageNumber) + ) { + return { chunks, pages } + } + totalPages = result.pagination?.totalPages ?? 1 + page += 1 + } while (page <= totalPages) + + return { chunks, pages } +} + +function chunksContainPage( + chunks: readonly ParsedChunkView[], + pageNumber: number, +): boolean { + return chunks.some( + (chunk) => + (chunk.pageAssets ?? []).some( + (pageAsset) => pageAsset.pageNumber === pageNumber, + ) || (chunk.pageNums ?? []).includes(pageNumber), + ) +} + async function fetchSources(): Promise { const body = await workspaceRouteClient.getJson( workspaceClientKeys.sources, diff --git a/src/integrations/knowhere-demo.test.ts b/src/integrations/knowhere-demo.test.ts index 7926818..fe4f7e9 100644 --- a/src/integrations/knowhere-demo.test.ts +++ b/src/integrations/knowhere-demo.test.ts @@ -106,6 +106,80 @@ describe("knowhereDemoApi", () => { ) }) + it("rewrites demo page citation assets onto the Notebook asset proxy", async () => { + globalThis.fetch = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + demo_source_id: "demo-tsla-q4-2025", + canonical_document_id: "demo-doc-tsla-q4-2025", + title: "TSLA-Q4-2025-Update.pdf", + mime_type: "application/pdf", + chunks: [ + { + id: "demo-tsla-q4-2025:page-1", + chunk_id: "page-1", + chunk_type: "page", + content: "Tesla energy storage deployments.", + section_path: "TSLA-Q4-2025-Update.pdf/SUMMARY", + source_chunk_path: "TSLA-Q4-2025-Update.pdf/SUMMARY", + file_path: "page_citation_assets/page-8.png", + sort_order: 8, + metadata: { + page_nums: [8], + page_assets: [ + { + page_num: 8, + artifact_ref: "page_citation_assets/page-8.png", + content_type: "image/png", + source: "knowhere-rendered-page-citation-source", + asset_url: + "/api/v1/demo/sources/demo-tsla-q4-2025/assets/page_citation_assets/page-8.png", + width: 1200, + height: 1600, + }, + ], + }, + asset_url: + "/api/v1/demo/sources/demo-tsla-q4-2025/assets/page_citation_assets/page-8.png", + }, + ], + pagination: { + page: 1, + page_size: 100, + total: 1, + total_pages: 1, + }, + }), + { status: 200, headers: { "content-type": "application/json" } }, + ), + ) + + const page = await knowhereDemoApi.fetchChunkPage({ + demoSourceId: "demo-tsla-q4-2025", + page: 1, + pageSize: 100, + }) + + expect(page.chunks[0]).toMatchObject({ + chunkType: "page", + filePath: "page_citation_assets/page-8.png", + assetUrl: + "/api/demo-sources/demo-tsla-q4-2025/assets/page_citation_assets/page-8.png", + metadata: { + pageAssets: [ + { + pageNum: 8, + contentType: "image/png", + assetUrl: + "/api/demo-sources/demo-tsla-q4-2025/assets/page_citation_assets/page-8.png", + width: 1200, + height: 1600, + }, + ], + }, + }) + }) + it("maps Official Library metadata from the demo catalog", async () => { globalThis.fetch = vi.fn().mockResolvedValue( new Response( diff --git a/src/integrations/knowhere-demo.ts b/src/integrations/knowhere-demo.ts index 8507dab..f4342a5 100644 --- a/src/integrations/knowhere-demo.ts +++ b/src/integrations/knowhere-demo.ts @@ -481,6 +481,10 @@ function toDemoChunk( demoSourceId: string, chunk: DemoChunkResponse, ): DemoChunk { + const metadata = rewriteDemoPageAssetMetadata( + demoSourceId, + toRecord(chunk.metadata), + ) return { id: requireString(chunk.id), chunkId: requireString(chunk.chunk_id), @@ -490,8 +494,10 @@ function toDemoChunk( sourceChunkPath: optionalString(chunk.source_chunk_path) ?? null, filePath: optionalString(chunk.file_path) ?? null, sortOrder: requireNumber(chunk.sort_order), - metadata: toRecord(chunk.metadata), - assetUrl: toDemoAssetUrl(demoSourceId, optionalString(chunk.asset_url)), + metadata, + assetUrl: + toDemoAssetUrl(demoSourceId, optionalString(chunk.asset_url)) ?? + firstDemoPageAssetUrl(metadata), } } @@ -603,6 +609,70 @@ function toRecord(value: unknown): Readonly> { return value as Readonly> } +function rewriteDemoPageAssetMetadata( + demoSourceId: string, + metadata: Readonly>, +): Readonly> { + const rawAssets = metadata.pageAssets ?? metadata.page_assets + if (!Array.isArray(rawAssets)) return metadata + + const pageAssets = rawAssets.flatMap((item) => { + if (typeof item !== "object" || item === null || Array.isArray(item)) { + return [] + } + const asset = item as Readonly> + const pageNum = + optionalPositiveNumber(asset.pageNum) ?? + optionalPositiveNumber(asset.pageNumber) ?? + optionalPositiveNumber(asset.page_num) + const contentType = + optionalString(asset.contentType) ?? optionalString(asset.content_type) + const artifactRef = optionalString(asset.artifact_ref) + const remoteAssetUrl = + optionalString(asset.assetUrl) ?? + optionalString(asset.asset_url) ?? + (artifactRef + ? `/api/v1/demo/sources/${encodeURIComponent(demoSourceId)}/assets/${artifactRef}` + : undefined) + const assetUrl = toDemoAssetUrl(demoSourceId, remoteAssetUrl) + if (!pageNum || !contentType || !assetUrl) return [] + + return [ + { + pageNum, + assetUrl, + contentType, + ...(optionalPositiveNumber(asset.width) !== undefined + ? { width: optionalPositiveNumber(asset.width) } + : {}), + ...(optionalPositiveNumber(asset.height) !== undefined + ? { height: optionalPositiveNumber(asset.height) } + : {}), + }, + ] + }) + + const next: Record = { ...metadata, pageAssets } + delete next.page_assets + return next +} + +function firstDemoPageAssetUrl( + metadata: Readonly>, +): string | null { + const pageAssets = metadata.pageAssets + if (!Array.isArray(pageAssets) || pageAssets.length === 0) return null + const first = pageAssets[0] + if (typeof first !== "object" || first === null) return null + return optionalString((first as Readonly>).assetUrl) ?? null +} + +function optionalPositiveNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value > 0 + ? value + : undefined +} + function toDemoAssetUrl( demoSourceId: string, assetUrl: string | undefined, diff --git a/src/lib/app-metadata.ts b/src/lib/app-metadata.ts index 705b93d..db9b41d 100644 --- a/src/lib/app-metadata.ts +++ b/src/lib/app-metadata.ts @@ -1,7 +1,7 @@ import type { Metadata } from "next"; export const appMetadata: Metadata = { - title: "Knowhere Notebook", + title: "Knowhere Brain", description: "Upload documents, explore parsed content, and ask questions about your knowledge.", icons: { From f06f8d99f490a9ff526eba4d19486b2242c47e9a Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 20 Aug 2026 10:18:58 +0800 Subject: [PATCH 39/51] fix: narrow knowledge before hydrating citation pages Typecheck failed because `input.knowledge` stayed `Knowledge | undefined` inside the map callback after the guard. Co-authored-by: Cursor --- src/domains/chat/index.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index 9f735e6..a6990c4 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -1014,11 +1014,12 @@ async function hydrateMissingCitationPageMetadata(input: { } }) - if (!input.knowledge) return results + const knowledge = input.knowledge + if (!knowledge) return results return Promise.all( results.map((result) => - hydrateResultPageMetadataFromKnowledge(result, input.knowledge), + hydrateResultPageMetadataFromKnowledge(result, knowledge), ), ) } From af497f639d9c281039add238c40646f32d05e3ca Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 20 Aug 2026 15:20:55 +0800 Subject: [PATCH 40/51] chore: bump knowhere-sdk to 2.2.0 for grep page numbers Use the published KnowledgeGrepMatch.pageNumbers field instead of a type cast so citation chips can take page numbers from grep without a follow-up read. Co-authored-by: Cursor --- package-lock.json | 1024 ++++++++++++++++++++++-------- package.json | 2 +- pnpm-lock.yaml | 10 +- src/agent-harness/ledger.test.ts | 3 +- src/agent-harness/ledger.ts | 9 +- 5 files changed, 781 insertions(+), 267 deletions(-) diff --git a/package-lock.json b/package-lock.json index e465210..81a0dbe 100644 --- a/package-lock.json +++ b/package-lock.json @@ -11,8 +11,9 @@ "@ai-sdk/react": "^3.0.177", "@antv/chart-visualization-skills": "0.1.3", "@effect/platform": "^0.96.1", + "@napi-rs/canvas": "^1.0.2", "@neondatabase/serverless": "^1.1.0", - "@ontos-ai/knowhere-sdk": "^2.1.1", + "@ontos-ai/knowhere-sdk": "^2.2.0", "@radix-ui/react-alert-dialog": "^1.1.15", "@radix-ui/react-checkbox": "^1.3.3", "@radix-ui/react-dialog": "^1.1.15", @@ -39,6 +40,7 @@ "next": "16.2.4", "next-themes": "^0.4.6", "pdfjs-dist": "5.4.296", + "piscina": "^5.2.0", "postgres": "^3.4.9", "posthog-js": "^1.386.8", "react": "19.2.4", @@ -72,6 +74,9 @@ "tailwindcss": "^4", "typescript": "^6.0.3", "vitest": "^4.1.5" + }, + "engines": { + "node": ">=22.13.0" } }, "node_modules/@ai-sdk/gateway": { @@ -1903,9 +1908,10 @@ "license": "MIT" }, "node_modules/@napi-rs/canvas": { - "version": "0.1.100", + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas/-/canvas-1.0.7.tgz", + "integrity": "sha512-26wVFEgs6gbe7wmzCrud1pK8q6oOgcUu7OOF24BuazB8ZUCskU9ZSLrCjoWIFVxx09rjAxsXxPleaWowHvdPCA==", "license": "MIT", - "optional": true, "workspaces": [ "e2e/*" ], @@ -1917,28 +1923,30 @@ "url": "https://github.com/sponsors/Brooooooklyn" }, "optionalDependencies": { - "@napi-rs/canvas-android-arm64": "0.1.100", - "@napi-rs/canvas-darwin-arm64": "0.1.100", - "@napi-rs/canvas-darwin-x64": "0.1.100", - "@napi-rs/canvas-linux-arm-gnueabihf": "0.1.100", - "@napi-rs/canvas-linux-arm64-gnu": "0.1.100", - "@napi-rs/canvas-linux-arm64-musl": "0.1.100", - "@napi-rs/canvas-linux-riscv64-gnu": "0.1.100", - "@napi-rs/canvas-linux-x64-gnu": "0.1.100", - "@napi-rs/canvas-linux-x64-musl": "0.1.100", - "@napi-rs/canvas-win32-arm64-msvc": "0.1.100", - "@napi-rs/canvas-win32-x64-msvc": "0.1.100" - } - }, - "node_modules/@napi-rs/canvas-darwin-arm64": { - "version": "0.1.100", + "@napi-rs/canvas-android-arm64": "1.0.7", + "@napi-rs/canvas-darwin-arm64": "1.0.7", + "@napi-rs/canvas-darwin-x64": "1.0.7", + "@napi-rs/canvas-linux-arm-gnueabihf": "1.0.7", + "@napi-rs/canvas-linux-arm64-gnu": "1.0.7", + "@napi-rs/canvas-linux-arm64-musl": "1.0.7", + "@napi-rs/canvas-linux-riscv64-gnu": "1.0.7", + "@napi-rs/canvas-linux-x64-gnu": "1.0.7", + "@napi-rs/canvas-linux-x64-musl": "1.0.7", + "@napi-rs/canvas-win32-arm64-msvc": "1.0.7", + "@napi-rs/canvas-win32-x64-msvc": "1.0.7" + } + }, + "node_modules/@napi-rs/canvas-android-arm64": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-android-arm64/-/canvas-android-arm64-1.0.7.tgz", + "integrity": "sha512-d5p4hHTykc/9PjiBXkvCH/IC5hT5jH7BjQ1u8ITq+G8x4VmvveOHdy/4BWYcC3ebWRmrG9zEc/oCTy+Yf3iZ4A==", "cpu": [ "arm64" ], "license": "MIT", "optional": true, "os": [ - "darwin" + "android" ], "engines": { "node": ">= 10" @@ -1948,17 +1956,17 @@ "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/@napi-rs/canvas/node_modules/@napi-rs/canvas-android-arm64": { - "version": "0.1.100", - "resolved": "https://registry.npmjs.org/@napi-rs/canvas-android-arm64/-/canvas-android-arm64-0.1.100.tgz", - "integrity": "sha512-hjhCKhntPv9+t4ckHymdx0phYNcVW+GKQR6Lzw2zE+pOVjOplSmtx9nNNknTjbEDLcuLZqA1y8ufKg1XfgftzQ==", + "node_modules/@napi-rs/canvas-darwin-arm64": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-darwin-arm64/-/canvas-darwin-arm64-1.0.7.tgz", + "integrity": "sha512-jKfZl2QDqBr6/Ap/8NKkX0Po9SFfVjUPBJbQQmYGVtQQIazWWIAO60riH3Mz2sOyysa2oJO39sLEfXerypu0vg==", "cpu": [ "arm64" ], "license": "MIT", "optional": true, "os": [ - "android" + "darwin" ], "engines": { "node": ">= 10" @@ -1968,10 +1976,10 @@ "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/@napi-rs/canvas/node_modules/@napi-rs/canvas-darwin-x64": { - "version": "0.1.100", - "resolved": "https://registry.npmjs.org/@napi-rs/canvas-darwin-x64/-/canvas-darwin-x64-0.1.100.tgz", - "integrity": "sha512-ePNZtj7pNIva/siZMg+HmbeozkIjqUIYdoymH8HaA3qK7LfzFN4WMBM8G6HQ9ZC+H3+Dnn5pqtiXpgLykaPOhw==", + "node_modules/@napi-rs/canvas-darwin-x64": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-darwin-x64/-/canvas-darwin-x64-1.0.7.tgz", + "integrity": "sha512-v1asrnKBu0tD+sdSD2qVIUfhoXSrr8LFTn7z76pN+8xsiOrmFcAVty57R/5DB8ZNqNLKsUBDXFSrzf4Wt9NaSg==", "cpu": [ "x64" ], @@ -1988,10 +1996,10 @@ "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/@napi-rs/canvas/node_modules/@napi-rs/canvas-linux-arm-gnueabihf": { - "version": "0.1.100", - "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm-gnueabihf/-/canvas-linux-arm-gnueabihf-0.1.100.tgz", - "integrity": "sha512-d5cDB48oWFGU8/XPhUOFAlySgb/VAu7D+s8fi55K1Pcfg8aPplHWqMgibhVLU8ky7Pyg/fuiVLz4Nf3JrSTuUA==", + "node_modules/@napi-rs/canvas-linux-arm-gnueabihf": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm-gnueabihf/-/canvas-linux-arm-gnueabihf-1.0.7.tgz", + "integrity": "sha512-c4Hcu4L5bXCgLY8jrZ2hy/Avl65MsbfH9DqNWrrd9InbpBZZF/rmrh7XkgmTUmOFcUrt/PaFSBinW2C0moDgcA==", "cpu": [ "arm" ], @@ -2008,13 +2016,16 @@ "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/@napi-rs/canvas/node_modules/@napi-rs/canvas-linux-arm64-gnu": { - "version": "0.1.100", - "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm64-gnu/-/canvas-linux-arm64-gnu-0.1.100.tgz", - "integrity": "sha512-rDxgxRu69RvDlX/bh9o22DxLsGr8EqsNgotL9+RwQE1S0b0cqeatqsw6aW45mukm0B42DIAaAacKaYQ8cqS1nw==", + "node_modules/@napi-rs/canvas-linux-arm64-gnu": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm64-gnu/-/canvas-linux-arm64-gnu-1.0.7.tgz", + "integrity": "sha512-OVW6T1x65BTVfWb//xylCoWxbdII+nzHu3L5T035aerBAnZ5e0nmeTMkMEO9qQrVJq4WcWW8hp+T0dF+JvJPUQ==", "cpu": [ "arm64" ], + "libc": [ + "glibc" + ], "license": "MIT", "optional": true, "os": [ @@ -2028,13 +2039,16 @@ "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/@napi-rs/canvas/node_modules/@napi-rs/canvas-linux-arm64-musl": { - "version": "0.1.100", - "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm64-musl/-/canvas-linux-arm64-musl-0.1.100.tgz", - "integrity": "sha512-K3mDW66N+xT2/V439u1alFANiBUjdEx2gLiNYnCmUsva5jZMxWTjafBYwTzYK+EMFMHrUoabuU+T1BIP5CgbYQ==", + "node_modules/@napi-rs/canvas-linux-arm64-musl": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm64-musl/-/canvas-linux-arm64-musl-1.0.7.tgz", + "integrity": "sha512-KX/k1UO61XdKlXxhtIkq1ZUH8uP+NNK4vSrBcM2/4wWUUCZaG93BU5s0KuE9n+TFn0jEoqLSlXOuH1pYuV5dcQ==", "cpu": [ "arm64" ], + "libc": [ + "musl" + ], "license": "MIT", "optional": true, "os": [ @@ -2048,12 +2062,367 @@ "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/@napi-rs/canvas/node_modules/@napi-rs/canvas-linux-riscv64-gnu": { - "version": "0.1.100", - "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-riscv64-gnu/-/canvas-linux-riscv64-gnu-0.1.100.tgz", - "integrity": "sha512-mooqUBTIsccZpnoQC4NgrC1v6C1vof39etLNMnBwCY+p0gajWJvAHLGQ6g/gGyS5YrpDW+GefSN4+Cvcr08UWw==", + "node_modules/@napi-rs/canvas-linux-riscv64-gnu": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-riscv64-gnu/-/canvas-linux-riscv64-gnu-1.0.7.tgz", + "integrity": "sha512-r+0Bg+fK8r2AsrbeT7JwBUAERZSjlyhfAMQfZE/zxYGmbswS92hN0FPjynVWbeuz5fIrLfZ6JA5pH8V6drN6qw==", + "cpu": [ + "riscv64" + ], + "libc": [ + "glibc" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@napi-rs/canvas-linux-x64-gnu": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-x64-gnu/-/canvas-linux-x64-gnu-1.0.7.tgz", + "integrity": "sha512-0tT9KzEcTfb6MWdUWIDfy6uq6UNF691r/9NfXxxl8s9+G5QrD7VP1UC+PAwzsCzJTClIKyidNau/grYTL+QmHw==", + "cpu": [ + "x64" + ], + "libc": [ + "glibc" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@napi-rs/canvas-linux-x64-musl": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-x64-musl/-/canvas-linux-x64-musl-1.0.7.tgz", + "integrity": "sha512-rvT0xo1xA+C7e+U/2ngsxWsl9gC5knHU+z8KTT5VK0Zos4AQbWZvtjvegrmzxm4o2yHE32Lexj6CDEJNvuTczA==", + "cpu": [ + "x64" + ], + "libc": [ + "musl" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@napi-rs/canvas-win32-arm64-msvc": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-win32-arm64-msvc/-/canvas-win32-arm64-msvc-1.0.7.tgz", + "integrity": "sha512-N6s3yoFFPUDkfRpjUiKERVYkli7ex5Sf0Bu/My6D6gOGYxYolC1KvL6x8U0aN+ScCWAzbkIWcMYR9g0OGsFuVw==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@napi-rs/canvas-win32-x64-msvc": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-win32-x64-msvc/-/canvas-win32-x64-msvc-1.0.7.tgz", + "integrity": "sha512-SKTNdBV/ljfhkPsLhXjm4x2PQ0ILpc0PhkBzRHuroJXidZUaF5yh0j3s3dIzB8PrRmW+nEASzyMTaWT3nH1ywQ==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@napi-rs/nice": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice/-/nice-1.1.1.tgz", + "integrity": "sha512-xJIPs+bYuc9ASBl+cvGsKbGrJmS6fAKaSZCnT0lhahT5rhA2VVy9/EcIgd2JhtEuFOJNx7UHNn/qiTPTY4nrQw==", + "license": "MIT", + "optional": true, + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + }, + "optionalDependencies": { + "@napi-rs/nice-android-arm-eabi": "1.1.1", + "@napi-rs/nice-android-arm64": "1.1.1", + "@napi-rs/nice-darwin-arm64": "1.1.1", + "@napi-rs/nice-darwin-x64": "1.1.1", + "@napi-rs/nice-freebsd-x64": "1.1.1", + "@napi-rs/nice-linux-arm-gnueabihf": "1.1.1", + "@napi-rs/nice-linux-arm64-gnu": "1.1.1", + "@napi-rs/nice-linux-arm64-musl": "1.1.1", + "@napi-rs/nice-linux-ppc64-gnu": "1.1.1", + "@napi-rs/nice-linux-riscv64-gnu": "1.1.1", + "@napi-rs/nice-linux-s390x-gnu": "1.1.1", + "@napi-rs/nice-linux-x64-gnu": "1.1.1", + "@napi-rs/nice-linux-x64-musl": "1.1.1", + "@napi-rs/nice-openharmony-arm64": "1.1.1", + "@napi-rs/nice-win32-arm64-msvc": "1.1.1", + "@napi-rs/nice-win32-ia32-msvc": "1.1.1", + "@napi-rs/nice-win32-x64-msvc": "1.1.1" + } + }, + "node_modules/@napi-rs/nice-android-arm-eabi": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-android-arm-eabi/-/nice-android-arm-eabi-1.1.1.tgz", + "integrity": "sha512-kjirL3N6TnRPv5iuHw36wnucNqXAO46dzK9oPb0wj076R5Xm8PfUVA9nAFB5ZNMmfJQJVKACAPd/Z2KYMppthw==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@napi-rs/nice-android-arm64": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-android-arm64/-/nice-android-arm64-1.1.1.tgz", + "integrity": "sha512-blG0i7dXgbInN5urONoUCNf+DUEAavRffrO7fZSeoRMJc5qD+BJeNcpr54msPF6qfDD6kzs9AQJogZvT2KD5nw==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@napi-rs/nice-darwin-arm64": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-darwin-arm64/-/nice-darwin-arm64-1.1.1.tgz", + "integrity": "sha512-s/E7w45NaLqTGuOjC2p96pct4jRfo61xb9bU1unM/MJ/RFkKlJyJDx7OJI/O0ll/hrfpqKopuAFDV8yo0hfT7A==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@napi-rs/nice-darwin-x64": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-darwin-x64/-/nice-darwin-x64-1.1.1.tgz", + "integrity": "sha512-dGoEBnVpsdcC+oHHmW1LRK5eiyzLwdgNQq3BmZIav+9/5WTZwBYX7r5ZkQC07Nxd3KHOCkgbHSh4wPkH1N1LiQ==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@napi-rs/nice-freebsd-x64": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-freebsd-x64/-/nice-freebsd-x64-1.1.1.tgz", + "integrity": "sha512-kHv4kEHAylMYmlNwcQcDtXjklYp4FCf0b05E+0h6nDHsZ+F0bDe04U/tXNOqrx5CmIAth4vwfkjjUmp4c4JktQ==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "freebsd" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@napi-rs/nice-linux-arm-gnueabihf": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-linux-arm-gnueabihf/-/nice-linux-arm-gnueabihf-1.1.1.tgz", + "integrity": "sha512-E1t7K0efyKXZDoZg1LzCOLxgolxV58HCkaEkEvIYQx12ht2pa8hoBo+4OB3qh7e+QiBlp1SRf+voWUZFxyhyqg==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@napi-rs/nice-linux-arm64-gnu": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-linux-arm64-gnu/-/nice-linux-arm64-gnu-1.1.1.tgz", + "integrity": "sha512-CIKLA12DTIZlmTaaKhQP88R3Xao+gyJxNWEn04wZwC2wmRapNnxCUZkVwggInMJvtVElA+D4ZzOU5sX4jV+SmQ==", + "cpu": [ + "arm64" + ], + "libc": [ + "glibc" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@napi-rs/nice-linux-arm64-musl": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-linux-arm64-musl/-/nice-linux-arm64-musl-1.1.1.tgz", + "integrity": "sha512-+2Rzdb3nTIYZ0YJF43qf2twhqOCkiSrHx2Pg6DJaCPYhhaxbLcdlV8hCRMHghQ+EtZQWGNcS2xF4KxBhSGeutg==", + "cpu": [ + "arm64" + ], + "libc": [ + "musl" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@napi-rs/nice-linux-ppc64-gnu": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-linux-ppc64-gnu/-/nice-linux-ppc64-gnu-1.1.1.tgz", + "integrity": "sha512-4FS8oc0GeHpwvv4tKciKkw3Y4jKsL7FRhaOeiPei0X9T4Jd619wHNe4xCLmN2EMgZoeGg+Q7GY7BsvwKpL22Tg==", + "cpu": [ + "ppc64" + ], + "libc": [ + "glibc" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@napi-rs/nice-linux-riscv64-gnu": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-linux-riscv64-gnu/-/nice-linux-riscv64-gnu-1.1.1.tgz", + "integrity": "sha512-HU0nw9uD4FO/oGCCk409tCi5IzIZpH2agE6nN4fqpwVlCn5BOq0MS1dXGjXaG17JaAvrlpV5ZeyZwSon10XOXw==", + "cpu": [ + "riscv64" + ], + "libc": [ + "glibc" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@napi-rs/nice-linux-s390x-gnu": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-linux-s390x-gnu/-/nice-linux-s390x-gnu-1.1.1.tgz", + "integrity": "sha512-2YqKJWWl24EwrX0DzCQgPLKQBxYDdBxOHot1KWEq7aY2uYeX+Uvtv4I8xFVVygJDgf6/92h9N3Y43WPx8+PAgQ==", + "cpu": [ + "s390x" + ], + "libc": [ + "glibc" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@napi-rs/nice-linux-x64-gnu": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-linux-x64-gnu/-/nice-linux-x64-gnu-1.1.1.tgz", + "integrity": "sha512-/gaNz3R92t+dcrfCw/96pDopcmec7oCcAQ3l/M+Zxr82KT4DljD37CpgrnXV+pJC263JkW572pdbP3hP+KjcIg==", + "cpu": [ + "x64" + ], + "libc": [ + "glibc" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + } + }, + "node_modules/@napi-rs/nice-linux-x64-musl": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-linux-x64-musl/-/nice-linux-x64-musl-1.1.1.tgz", + "integrity": "sha512-xScCGnyj/oppsNPMnevsBe3pvNaoK7FGvMjT35riz9YdhB2WtTG47ZlbxtOLpjeO9SqqQ2J2igCmz6IJOD5JYw==", "cpu": [ - "riscv64" + "x64" + ], + "libc": [ + "musl" ], "license": "MIT", "optional": true, @@ -2062,58 +2431,46 @@ ], "engines": { "node": ">= 10" - }, - "funding": { - "type": "github", - "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/@napi-rs/canvas/node_modules/@napi-rs/canvas-linux-x64-gnu": { - "version": "0.1.100", - "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-x64-gnu/-/canvas-linux-x64-gnu-0.1.100.tgz", - "integrity": "sha512-1eCvkDCazm7FFhsT7DfGOdSaHgZVK3bt/dSBl5EWHOWmnz+I7j8tPseJqqD81NF+MH21jKUK4wQSDjN0mdhnTg==", + "node_modules/@napi-rs/nice-openharmony-arm64": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-openharmony-arm64/-/nice-openharmony-arm64-1.1.1.tgz", + "integrity": "sha512-6uJPRVwVCLDeoOaNyeiW0gp2kFIM4r7PL2MczdZQHkFi9gVlgm+Vn+V6nTWRcu856mJ2WjYJiumEajfSm7arPQ==", "cpu": [ - "x64" + "arm64" ], "license": "MIT", "optional": true, "os": [ - "linux" + "openharmony" ], "engines": { "node": ">= 10" - }, - "funding": { - "type": "github", - "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/@napi-rs/canvas/node_modules/@napi-rs/canvas-linux-x64-musl": { - "version": "0.1.100", - "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-x64-musl/-/canvas-linux-x64-musl-0.1.100.tgz", - "integrity": "sha512-20arT6lnI19S68qNlii73TSEDbECNgzMz2EpldC1V3mZFuRkeujXkcebRk0LRJe9SEUAooYiLokfMViY8IX7yA==", + "node_modules/@napi-rs/nice-win32-arm64-msvc": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-win32-arm64-msvc/-/nice-win32-arm64-msvc-1.1.1.tgz", + "integrity": "sha512-uoTb4eAvM5B2aj/z8j+Nv8OttPf2m+HVx3UjA5jcFxASvNhQriyCQF1OB1lHL43ZhW+VwZlgvjmP5qF3+59atA==", "cpu": [ - "x64" + "arm64" ], "license": "MIT", "optional": true, "os": [ - "linux" + "win32" ], "engines": { "node": ">= 10" - }, - "funding": { - "type": "github", - "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/@napi-rs/canvas/node_modules/@napi-rs/canvas-win32-arm64-msvc": { - "version": "0.1.100", - "resolved": "https://registry.npmjs.org/@napi-rs/canvas-win32-arm64-msvc/-/canvas-win32-arm64-msvc-0.1.100.tgz", - "integrity": "sha512-DZFFT1wIAg37LJw37yhMRFfjATd3vTQzjZ1Yki8u2vhO6Hi5VE6BVaGQ1aaDu7xb4iMErz+9EOwjpS7xcxFeBw==", + "node_modules/@napi-rs/nice-win32-ia32-msvc": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-win32-ia32-msvc/-/nice-win32-ia32-msvc-1.1.1.tgz", + "integrity": "sha512-CNQqlQT9MwuCsg1Vd/oKXiuH+TcsSPJmlAFc5frFyX/KkOh0UpBLEj7aoY656d5UKZQMQFP7vJNa1DNUNORvug==", "cpu": [ - "arm64" + "ia32" ], "license": "MIT", "optional": true, @@ -2122,16 +2479,12 @@ ], "engines": { "node": ">= 10" - }, - "funding": { - "type": "github", - "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/@napi-rs/canvas/node_modules/@napi-rs/canvas-win32-x64-msvc": { - "version": "0.1.100", - "resolved": "https://registry.npmjs.org/@napi-rs/canvas-win32-x64-msvc/-/canvas-win32-x64-msvc-0.1.100.tgz", - "integrity": "sha512-MyT1j3mHC2+Lu4pBi9mKyMJhtP6U7k7EldY7sj/uS5gJA65gTXt8MefJQXLJo5d/vZbuWmfxzkEUNc/urV3pHA==", + "node_modules/@napi-rs/nice-win32-x64-msvc": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@napi-rs/nice-win32-x64-msvc/-/nice-win32-x64-msvc-1.1.1.tgz", + "integrity": "sha512-vB+4G/jBQCAh0jelMTY3+kgFy00Hlx2f2/1zjMoH821IbplbWZOkLiTYXQkygNTzQJTq5cvwBDgn2ppHD+bglQ==", "cpu": [ "x64" ], @@ -2142,10 +2495,6 @@ ], "engines": { "node": ">= 10" - }, - "funding": { - "type": "github", - "url": "https://github.com/sponsors/Brooooooklyn" } }, "node_modules/@napi-rs/wasm-runtime": { @@ -2275,9 +2624,9 @@ } }, "node_modules/@ontos-ai/knowhere-sdk": { - "version": "2.1.1", - "resolved": "https://registry.npmjs.org/@ontos-ai/knowhere-sdk/-/knowhere-sdk-2.1.1.tgz", - "integrity": "sha512-K33ylB/QjVYlLgow+/Hw+uNhQbWWrMTIqZ3IAlx5buHSI2EHPVP9+kDsfslwpfffYSn85j9Ab0/E9qTG0QotTw==", + "version": "2.2.0", + "resolved": "https://registry.npmjs.org/@ontos-ai/knowhere-sdk/-/knowhere-sdk-2.2.0.tgz", + "integrity": "sha512-+kFykBxJ6BzcTAiyl7s7vgWO6Gtay7VH3n6/QLXRmemOjm8qk00gfE105wFFrkrkYzBXu4rFQaX5hxHM9uRWCA==", "license": "MIT", "dependencies": { "axios": "^1.15.2", @@ -7817,7 +8166,6 @@ }, "node_modules/fsevents": { "version": "2.3.2", - "dev": true, "license": "MIT", "optional": true, "os": [ @@ -11389,93 +11737,358 @@ "lines-and-columns": "^1.1.6" }, "engines": { - "node": ">=8" + "node": ">=8" + }, + "funding": { + "url": "https://github.com/sponsors/sindresorhus" + } + }, + "node_modules/parse-ms": { + "version": "4.0.0", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=18" + }, + "funding": { + "url": "https://github.com/sponsors/sindresorhus" + } + }, + "node_modules/parse5": { + "version": "8.0.1", + "dev": true, + "license": "MIT", + "dependencies": { + "entities": "^8.0.0" + }, + "funding": { + "url": "https://github.com/inikulin/parse5?sponsor=1" + } + }, + "node_modules/parseurl": { + "version": "1.3.3", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, + "node_modules/path-browserify": { + "version": "1.0.1", + "dev": true, + "license": "MIT" + }, + "node_modules/path-exists": { + "version": "4.0.0", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=8" + } + }, + "node_modules/path-is-absolute": { + "version": "1.0.1", + "license": "MIT", + "engines": { + "node": ">=0.10.0" + } + }, + "node_modules/path-key": { + "version": "3.1.1", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=8" + } + }, + "node_modules/path-parse": { + "version": "1.0.7", + "dev": true, + "license": "MIT" + }, + "node_modules/path-to-regexp": { + "version": "6.3.0", + "dev": true, + "license": "MIT" + }, + "node_modules/pathe": { + "version": "2.0.3", + "dev": true, + "license": "MIT" + }, + "node_modules/pdfjs-dist": { + "version": "5.4.296", + "license": "Apache-2.0", + "engines": { + "node": ">=20.16.0 || >=22.3.0" + }, + "optionalDependencies": { + "@napi-rs/canvas": "^0.1.80" + } + }, + "node_modules/pdfjs-dist/node_modules/@napi-rs/canvas": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas/-/canvas-0.1.100.tgz", + "integrity": "sha512-xglYA6q3XO5P3BNJYxVZ1IV7DLVjp1Py6nwag88YntrS+3vKHyYcMqXVS4ZztJmwz2uGvz1FWhI/4LgbR5uQDA==", + "license": "MIT", + "optional": true, + "workspaces": [ + "e2e/*" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + }, + "optionalDependencies": { + "@napi-rs/canvas-android-arm64": "0.1.100", + "@napi-rs/canvas-darwin-arm64": "0.1.100", + "@napi-rs/canvas-darwin-x64": "0.1.100", + "@napi-rs/canvas-linux-arm-gnueabihf": "0.1.100", + "@napi-rs/canvas-linux-arm64-gnu": "0.1.100", + "@napi-rs/canvas-linux-arm64-musl": "0.1.100", + "@napi-rs/canvas-linux-riscv64-gnu": "0.1.100", + "@napi-rs/canvas-linux-x64-gnu": "0.1.100", + "@napi-rs/canvas-linux-x64-musl": "0.1.100", + "@napi-rs/canvas-win32-arm64-msvc": "0.1.100", + "@napi-rs/canvas-win32-x64-msvc": "0.1.100" + } + }, + "node_modules/pdfjs-dist/node_modules/@napi-rs/canvas-android-arm64": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-android-arm64/-/canvas-android-arm64-0.1.100.tgz", + "integrity": "sha512-hjhCKhntPv9+t4ckHymdx0phYNcVW+GKQR6Lzw2zE+pOVjOplSmtx9nNNknTjbEDLcuLZqA1y8ufKg1XfgftzQ==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/pdfjs-dist/node_modules/@napi-rs/canvas-darwin-arm64": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-darwin-arm64/-/canvas-darwin-arm64-0.1.100.tgz", + "integrity": "sha512-2PcswRaC7Ly645DGt88///zuFDhJxJYdKAs1uU3mfk1atYkXufgcgLfBpk6Tm12nCQBaNt1wpybuPZ4qOhTo8A==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/pdfjs-dist/node_modules/@napi-rs/canvas-darwin-x64": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-darwin-x64/-/canvas-darwin-x64-0.1.100.tgz", + "integrity": "sha512-ePNZtj7pNIva/siZMg+HmbeozkIjqUIYdoymH8HaA3qK7LfzFN4WMBM8G6HQ9ZC+H3+Dnn5pqtiXpgLykaPOhw==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/pdfjs-dist/node_modules/@napi-rs/canvas-linux-arm-gnueabihf": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm-gnueabihf/-/canvas-linux-arm-gnueabihf-0.1.100.tgz", + "integrity": "sha512-d5cDB48oWFGU8/XPhUOFAlySgb/VAu7D+s8fi55K1Pcfg8aPplHWqMgibhVLU8ky7Pyg/fuiVLz4Nf3JrSTuUA==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/pdfjs-dist/node_modules/@napi-rs/canvas-linux-arm64-gnu": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm64-gnu/-/canvas-linux-arm64-gnu-0.1.100.tgz", + "integrity": "sha512-rDxgxRu69RvDlX/bh9o22DxLsGr8EqsNgotL9+RwQE1S0b0cqeatqsw6aW45mukm0B42DIAaAacKaYQ8cqS1nw==", + "cpu": [ + "arm64" + ], + "libc": [ + "glibc" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" }, "funding": { - "url": "https://github.com/sponsors/sindresorhus" + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/parse-ms": { - "version": "4.0.0", - "dev": true, + "node_modules/pdfjs-dist/node_modules/@napi-rs/canvas-linux-arm64-musl": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm64-musl/-/canvas-linux-arm64-musl-0.1.100.tgz", + "integrity": "sha512-K3mDW66N+xT2/V439u1alFANiBUjdEx2gLiNYnCmUsva5jZMxWTjafBYwTzYK+EMFMHrUoabuU+T1BIP5CgbYQ==", + "cpu": [ + "arm64" + ], + "libc": [ + "musl" + ], "license": "MIT", + "optional": true, + "os": [ + "linux" + ], "engines": { - "node": ">=18" + "node": ">= 10" }, "funding": { - "url": "https://github.com/sponsors/sindresorhus" + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/parse5": { - "version": "8.0.1", - "dev": true, + "node_modules/pdfjs-dist/node_modules/@napi-rs/canvas-linux-riscv64-gnu": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-riscv64-gnu/-/canvas-linux-riscv64-gnu-0.1.100.tgz", + "integrity": "sha512-mooqUBTIsccZpnoQC4NgrC1v6C1vof39etLNMnBwCY+p0gajWJvAHLGQ6g/gGyS5YrpDW+GefSN4+Cvcr08UWw==", + "cpu": [ + "riscv64" + ], + "libc": [ + "glibc" + ], "license": "MIT", - "dependencies": { - "entities": "^8.0.0" + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">= 10" }, "funding": { - "url": "https://github.com/inikulin/parse5?sponsor=1" + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/parseurl": { - "version": "1.3.3", - "dev": true, + "node_modules/pdfjs-dist/node_modules/@napi-rs/canvas-linux-x64-gnu": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-x64-gnu/-/canvas-linux-x64-gnu-0.1.100.tgz", + "integrity": "sha512-1eCvkDCazm7FFhsT7DfGOdSaHgZVK3bt/dSBl5EWHOWmnz+I7j8tPseJqqD81NF+MH21jKUK4wQSDjN0mdhnTg==", + "cpu": [ + "x64" + ], + "libc": [ + "glibc" + ], "license": "MIT", + "optional": true, + "os": [ + "linux" + ], "engines": { - "node": ">= 0.8" + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/path-browserify": { - "version": "1.0.1", - "dev": true, - "license": "MIT" - }, - "node_modules/path-exists": { - "version": "4.0.0", - "dev": true, + "node_modules/pdfjs-dist/node_modules/@napi-rs/canvas-linux-x64-musl": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-x64-musl/-/canvas-linux-x64-musl-0.1.100.tgz", + "integrity": "sha512-20arT6lnI19S68qNlii73TSEDbECNgzMz2EpldC1V3mZFuRkeujXkcebRk0LRJe9SEUAooYiLokfMViY8IX7yA==", + "cpu": [ + "x64" + ], + "libc": [ + "musl" + ], "license": "MIT", + "optional": true, + "os": [ + "linux" + ], "engines": { - "node": ">=8" + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/path-is-absolute": { - "version": "1.0.1", + "node_modules/pdfjs-dist/node_modules/@napi-rs/canvas-win32-arm64-msvc": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-win32-arm64-msvc/-/canvas-win32-arm64-msvc-0.1.100.tgz", + "integrity": "sha512-DZFFT1wIAg37LJw37yhMRFfjATd3vTQzjZ1Yki8u2vhO6Hi5VE6BVaGQ1aaDu7xb4iMErz+9EOwjpS7xcxFeBw==", + "cpu": [ + "arm64" + ], "license": "MIT", + "optional": true, + "os": [ + "win32" + ], "engines": { - "node": ">=0.10.0" + "node": ">= 10" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" } }, - "node_modules/path-key": { - "version": "3.1.1", - "dev": true, + "node_modules/pdfjs-dist/node_modules/@napi-rs/canvas-win32-x64-msvc": { + "version": "0.1.100", + "resolved": "https://registry.npmjs.org/@napi-rs/canvas-win32-x64-msvc/-/canvas-win32-x64-msvc-0.1.100.tgz", + "integrity": "sha512-MyT1j3mHC2+Lu4pBi9mKyMJhtP6U7k7EldY7sj/uS5gJA65gTXt8MefJQXLJo5d/vZbuWmfxzkEUNc/urV3pHA==", + "cpu": [ + "x64" + ], "license": "MIT", + "optional": true, + "os": [ + "win32" + ], "engines": { - "node": ">=8" - } - }, - "node_modules/path-parse": { - "version": "1.0.7", - "dev": true, - "license": "MIT" - }, - "node_modules/path-to-regexp": { - "version": "6.3.0", - "dev": true, - "license": "MIT" - }, - "node_modules/pathe": { - "version": "2.0.3", - "dev": true, - "license": "MIT" - }, - "node_modules/pdfjs-dist": { - "version": "5.4.296", - "license": "Apache-2.0", - "engines": { - "node": ">=20.16.0 || >=22.3.0" + "node": ">= 10" }, - "optionalDependencies": { - "@napi-rs/canvas": "^0.1.80" + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" } }, "node_modules/picocolors": { @@ -11493,6 +12106,18 @@ "url": "https://github.com/sponsors/jonschlinkert" } }, + "node_modules/piscina": { + "version": "5.3.0", + "resolved": "https://registry.npmjs.org/piscina/-/piscina-5.3.0.tgz", + "integrity": "sha512-9D3tPBayfoal6l9l4Y8NrMn1jkV718qJoYrla6P51+hh9h0Q+9/1c8KgEnoBQqhguyfgjay4biADI8HJG3pkoA==", + "license": "MIT", + "engines": { + "node": ">=20.x" + }, + "optionalDependencies": { + "@napi-rs/nice": "^1.0.4" + } + }, "node_modules/pkce-challenge": { "version": "5.0.1", "dev": true, @@ -15453,111 +16078,6 @@ "type": "github", "url": "https://github.com/sponsors/wooorm" } - }, - "node_modules/@next/swc-darwin-x64": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/swc-darwin-x64/-/swc-darwin-x64-16.2.4.tgz", - "integrity": "sha512-XhpVnUfmYWvD3YrXu55XdcAkQtOnvaI6wtQa8fuF5fGoKoxIUZ0kWPtcOfqJEWngFF/lOS9l3+O9CcownhiQxQ==", - "cpu": [ - "x64" - ], - "optional": true, - "os": [ - "darwin" - ], - "engines": { - "node": ">= 10" - } - }, - "node_modules/@next/swc-linux-arm64-gnu": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-gnu/-/swc-linux-arm64-gnu-16.2.4.tgz", - "integrity": "sha512-Mx/tjlNA3G8kg14QvuGAJ4xBwPk1tUHq56JxZ8CXnZwz1Etz714soCEzGQQzVMz4bEnGPowzkV6Xrp6wAkEWOQ==", - "cpu": [ - "arm64" - ], - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": ">= 10" - } - }, - "node_modules/@next/swc-linux-arm64-musl": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-musl/-/swc-linux-arm64-musl-16.2.4.tgz", - "integrity": "sha512-iVMMp14514u7Nup2umQS03nT/bN9HurK8ufylC3FZNykrwjtx7V1A7+4kvhbDSCeonTVqV3Txnv0Lu+m2oDXNg==", - "cpu": [ - "arm64" - ], - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": ">= 10" - } - }, - "node_modules/@next/swc-linux-x64-gnu": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-gnu/-/swc-linux-x64-gnu-16.2.4.tgz", - "integrity": "sha512-EZOvm1aQWgnI/N/xcWOlnS3RQBk0VtVav5Zo7n4p0A7UKyTDx047k8opDbXgBpHl4CulRqRfbw3QrX2w5UOXMQ==", - "cpu": [ - "x64" - ], - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": ">= 10" - } - }, - "node_modules/@next/swc-linux-x64-musl": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-musl/-/swc-linux-x64-musl-16.2.4.tgz", - "integrity": "sha512-h9FxsngCm9cTBf71AR4fGznDEDx1hS7+kSEiIRjq5kO1oXWm07DxVGZjCvk0SGx7TSjlUqhI8oOyz7NfwAdPoA==", - "cpu": [ - "x64" - ], - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": ">= 10" - } - }, - "node_modules/@next/swc-win32-arm64-msvc": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/swc-win32-arm64-msvc/-/swc-win32-arm64-msvc-16.2.4.tgz", - "integrity": "sha512-3NdJV5OXMSOeJYijX+bjaLge3mJBlh4ybydbT4GFoB/2hAojWHtMhl3CYlYoMrjPuodp0nzFVi4Tj2+WaMg+Ow==", - "cpu": [ - "arm64" - ], - "optional": true, - "os": [ - "win32" - ], - "engines": { - "node": ">= 10" - } - }, - "node_modules/@next/swc-win32-x64-msvc": { - "version": "16.2.4", - "resolved": "https://registry.npmjs.org/@next/swc-win32-x64-msvc/-/swc-win32-x64-msvc-16.2.4.tgz", - "integrity": "sha512-kMVGgsqhO5YTYODD9IPGGhA6iprWidQckK3LmPeW08PIFENRmgfb4MjXHO+p//d+ts2rpjvK5gXWzXSMrPl9cw==", - "cpu": [ - "x64" - ], - "optional": true, - "os": [ - "win32" - ], - "engines": { - "node": ">= 10" - } } } } diff --git a/package.json b/package.json index cf48c4a..f377889 100644 --- a/package.json +++ b/package.json @@ -28,7 +28,7 @@ "@effect/platform": "^0.96.1", "@napi-rs/canvas": "^1.0.2", "@neondatabase/serverless": "^1.1.0", - "@ontos-ai/knowhere-sdk": "^2.1.2", + "@ontos-ai/knowhere-sdk": "^2.2.0", "@radix-ui/react-alert-dialog": "^1.1.15", "@radix-ui/react-checkbox": "^1.3.3", "@radix-ui/react-dialog": "^1.1.15", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 12dc39f..8847510 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -24,8 +24,8 @@ importers: specifier: ^1.1.0 version: 1.1.0 '@ontos-ai/knowhere-sdk': - specifier: ^2.1.2 - version: 2.1.2 + specifier: ^2.2.0 + version: 2.2.0 '@radix-ui/react-alert-dialog': specifier: ^1.1.15 version: 1.1.15(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4) @@ -1603,8 +1603,8 @@ packages: resolution: {integrity: sha512-nn5ozdjYQpUCZlWGuxcJY/KpxkWQs4DcbMCmKojjyrYDEAGy4Ce19NN4v5MduafTwJlbKc99UA8YhSVqq9yPZA==} engines: {node: '>=12.4.0'} - '@ontos-ai/knowhere-sdk@2.1.2': - resolution: {integrity: sha512-m1wbQZNesExcg2yoEnmqFRlML/tprv11y/v7tlTLXEKTqKOY5zpG67idxXtbmw2qjOjQkR02HhgUwo/SdrFWPA==} + '@ontos-ai/knowhere-sdk@2.2.0': + resolution: {integrity: sha512-+kFykBxJ6BzcTAiyl7s7vgWO6Gtay7VH3n6/QLXRmemOjm8qk00gfE105wFFrkrkYzBXu4rFQaX5hxHM9uRWCA==} engines: {node: '>=22.13.0', npm: '>=10.0.0', pnpm: '>=9.0.0'} '@open-draft/deferred-promise@2.2.0': @@ -6731,7 +6731,7 @@ snapshots: '@nolyfill/is-core-module@1.0.39': {} - '@ontos-ai/knowhere-sdk@2.1.2': + '@ontos-ai/knowhere-sdk@2.2.0': dependencies: axios: 1.18.1 jszip: 3.10.1 diff --git a/src/agent-harness/ledger.test.ts b/src/agent-harness/ledger.test.ts index 15f1b20..8a77142 100644 --- a/src/agent-harness/ledger.test.ts +++ b/src/agent-harness/ledger.test.ts @@ -1,6 +1,5 @@ import { describe, expect, it } from "vitest" import type { - KnowledgeGrepMatch, KnowledgeGrepResponse, KnowledgeReadResponse, RetrievalQueryResponse, @@ -164,7 +163,7 @@ describe("createEvidenceLedger", () => { endOffset: 12, snippet: "automotive revenues", pageNumbers: [4], - } as KnowledgeGrepMatch, + }, ], }) diff --git a/src/agent-harness/ledger.ts b/src/agent-harness/ledger.ts index aa9d58e..4e13f81 100644 --- a/src/agent-harness/ledger.ts +++ b/src/agent-harness/ledger.ts @@ -255,17 +255,13 @@ function addChunkFromGrepMatch(input: { readonly match: KnowledgeGrepMatch readonly ref: string }): void { - const matchWithPages = input.match as KnowledgeGrepMatch & { - readonly pageNumbers?: readonly number[] - readonly metadata?: Readonly> - } const donor = input.ledger.chunks.find( (chunk) => chunk.chunkId === input.match.chunkId && hasPageMetadata(chunk.metadata), ) const pageNums = - matchWithPages.pageNumbers && matchWithPages.pageNumbers.length > 0 - ? [...matchWithPages.pageNumbers] + input.match.pageNumbers && input.match.pageNumbers.length > 0 + ? [...input.match.pageNumbers] : undefined addChunk({ @@ -282,7 +278,6 @@ function addChunkFromGrepMatch(input: { filePath: input.match.filePath, metadata: { ...(donor?.metadata ?? {}), - ...(matchWithPages.metadata ?? {}), ...(pageNums ? { pageNums } : {}), position: input.match.position, startOffset: input.match.startOffset, From 6671a0a3e38296bdf9f7e71eea05baef56e0d34c Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 20 Aug 2026 19:32:48 +0800 Subject: [PATCH 41/51] feat: archive remote Knowhere sources and filter by client origin Remote-only sidebar ids now call Knowhere archive (and clean a matching local row if one exists). Listing and chat only show remotes created by notebook, CLI, or MCP. Co-authored-by: Cursor --- src/app/api/sources/[sourceId]/route.test.ts | 69 ++++++++++- src/domains/chat/index.test.ts | 9 ++ src/domains/chat/knowhere-tools.ts | 2 + src/domains/sources/document-metadata.test.ts | 32 +++++ src/domains/sources/document-metadata.ts | 35 ++++++ src/domains/sources/remote-library.test.ts | 117 ++++++++++++++++++ src/domains/sources/remote-library.ts | 13 +- src/domains/sources/repository.ts | 3 + src/domains/sources/route-archive.ts | 87 +++++++++---- src/domains/sources/route-dependencies.ts | 1 + src/domains/sources/route-service.test.ts | 28 +++++ src/domains/sources/route-types.ts | 4 + src/domains/sources/service.ts | 5 + src/domains/sources/source-row-repository.ts | 36 ++++++ src/domains/sources/workflow-runtime.ts | 11 ++ 15 files changed, 426 insertions(+), 26 deletions(-) create mode 100644 src/domains/sources/remote-library.test.ts diff --git a/src/app/api/sources/[sourceId]/route.test.ts b/src/app/api/sources/[sourceId]/route.test.ts index 6f82574..f531c23 100644 --- a/src/app/api/sources/[sourceId]/route.test.ts +++ b/src/app/api/sources/[sourceId]/route.test.ts @@ -10,6 +10,7 @@ const mocks = vi.hoisted(() => { ensureWorkspace: vi.fn(), fetchDemoCatalog: vi.fn(), findSourceInWorkspace: vi.fn(), + findByKnowhereDocumentId: vi.fn(), getCurrentUser: vi.fn(), hideDemoSource: vi.fn(), makeKnowhereClient: vi.fn(), @@ -55,6 +56,7 @@ vi.mock("@/domains/sources/background-reconcile", () => ({ vi.mock("@/domains/sources/service", () => ({ sourceService: { findInWorkspace: mocks.findSourceInWorkspace, + findByKnowhereDocumentId: mocks.findByKnowhereDocumentId, hideDemoSource: mocks.hideDemoSource, retrySourceToKnowhere: mocks.retrySourceToKnowhere, softDelete: mocks.softDeleteSource, @@ -116,11 +118,17 @@ describe("PATCH /api/sources/[sourceId]", () => { ); }); - it("rejects archive requests for unlocalized remote source ids", async () => { + it("archives unlocalized remote Knowhere documents", async () => { mocks.requireUser.mockResolvedValue({ id: "user_1" }); mocks.ensureWorkspace.mockResolvedValue({ id: "workspace_1" }); mocks.findSourceInWorkspace.mockResolvedValue(null); + mocks.findByKnowhereDocumentId.mockResolvedValue(null); mocks.fetchDemoCatalog.mockResolvedValue({ sources: [] }); + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123"); + mocks.makeKnowhereClient.mockReturnValue({ + documents: { archive: mocks.archive }, + }); + mocks.archive.mockResolvedValue(undefined); const response = await PATCH( new NextRequest( @@ -138,18 +146,71 @@ describe("PATCH /api/sources/[sourceId]", () => { ); await expect(response.json()).resolves.toEqual({ - message: "Source not found.", + id: "knowhere-doc:default:doc_remote", + archived: true, }); - expect(response.status).toBe(404); + expect(response.status).toBe(200); expect(mocks.findSourceInWorkspace).toHaveBeenCalledWith( "workspace_1", "knowhere-doc:default:doc_remote", ); - expect(mocks.archive).not.toHaveBeenCalled(); + expect(mocks.archive).toHaveBeenCalledWith("doc_remote"); + expect(mocks.findByKnowhereDocumentId).toHaveBeenCalledWith( + "workspace_1", + "doc_remote", + ); expect(mocks.softDeleteSource).not.toHaveBeenCalled(); expect(mocks.deleteBlob).not.toHaveBeenCalled(); }); + it("soft-deletes a matching local row when archiving a remote source id", async () => { + mocks.requireUser.mockResolvedValue({ id: "user_1" }); + mocks.ensureWorkspace.mockResolvedValue({ id: "workspace_1" }); + mocks.findSourceInWorkspace.mockResolvedValue(null); + mocks.findByKnowhereDocumentId.mockResolvedValue({ + id: "source_1", + knowhereDocumentId: "doc_remote", + originalBlobPathname: "source-uploads/upload_1/document.pdf", + demoKey: null, + }); + mocks.fetchDemoCatalog.mockResolvedValue({ sources: [] }); + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123"); + mocks.makeKnowhereClient.mockReturnValue({ + documents: { archive: mocks.archive }, + }); + mocks.archive.mockResolvedValue(undefined); + mocks.softDeleteSource.mockResolvedValue(true); + + const response = await PATCH( + new NextRequest( + "http://localhost:3001/api/sources/knowhere-doc:default:doc_remote", + { + method: "PATCH", + body: JSON.stringify({ archived: true }), + }, + ), + { + params: Promise.resolve({ + sourceId: "knowhere-doc:default:doc_remote", + }), + }, + ); + + await expect(response.json()).resolves.toEqual({ + id: "knowhere-doc:default:doc_remote", + archived: true, + }); + expect(response.status).toBe(200); + expect(mocks.archive).toHaveBeenCalledWith("doc_remote"); + expect(mocks.softDeleteSource).toHaveBeenCalledWith( + "workspace_1", + "source_1", + ); + expect(mocks.deleteBlob).toHaveBeenCalledWith( + "source-uploads/upload_1/document.pdf", + ); + }); + it("does not fail an already-soft-deleted source when original Blob cleanup fails", async () => { mocks.requireUser.mockResolvedValue({ id: "user_1" }); mocks.ensureWorkspace.mockResolvedValue({ id: "workspace_1" }); diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index e9556ac..e26a670 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -186,6 +186,15 @@ describe("answerQuestionWithRetrieval", () => { status: "ready", currentJobResultId: "job_remote", sourceFileName: "remote.pdf", + documentMetadata: { + createdByClient: "cli", + }, + }, + { + documentId: "doc_untagged", + namespace: "default", + status: "ready", + sourceFileName: "dummy.pdf", }, ], }); diff --git a/src/domains/chat/knowhere-tools.ts b/src/domains/chat/knowhere-tools.ts index 210b487..5c70fd7 100644 --- a/src/domains/chat/knowhere-tools.ts +++ b/src/domains/chat/knowhere-tools.ts @@ -8,6 +8,7 @@ import type { import type { Source } from "@/infrastructure/db/schema" import { listRemoteLibraryDocuments, + isNotebookVisibleRemoteDocument, type RemoteLibraryDocument, } from "@/domains/sources/remote-library" import type { SearchSources } from "./contracts" @@ -128,6 +129,7 @@ async function listVisibleRemoteDocuments(input: { return documents .filter( (document) => + isNotebookVisibleRemoteDocument(document) && document.status === "ready" && !localDocumentIds.has(document.documentId) && !input.excludedDocumentIds.has(document.documentId), diff --git a/src/domains/sources/document-metadata.test.ts b/src/domains/sources/document-metadata.test.ts index ba64346..6c69ac6 100644 --- a/src/domains/sources/document-metadata.test.ts +++ b/src/domains/sources/document-metadata.test.ts @@ -2,7 +2,11 @@ import { describe, expect, it } from "vitest" import { createNotebookDocumentMetadata, + getCreatedByClient, + isNotebookVisibleRemoteClient, + isNotebookVisibleRemoteMetadata, NOTEBOOK_DOCUMENT_METADATA_DEFAULTS, + NOTEBOOK_VISIBLE_CREATED_BY_CLIENTS, } from "./document-metadata" describe("createNotebookDocumentMetadata", () => { @@ -46,4 +50,32 @@ describe("createNotebookDocumentMetadata", () => { clientVersion: "0.1.0", }) }) + + it("reads camelCase and snake_case created-by-client metadata", () => { + expect(getCreatedByClient({ createdByClient: "cli" })).toBe("cli") + expect(getCreatedByClient({ created_by_client: "mcp" })).toBe("mcp") + expect(getCreatedByClient({})).toBeUndefined() + expect(getCreatedByClient(undefined)).toBeUndefined() + }) + + it("allowlists notebook, cli, and mcp remote documents", () => { + expect(NOTEBOOK_VISIBLE_CREATED_BY_CLIENTS).toEqual([ + "notebook", + "cli", + "mcp", + ]) + expect(isNotebookVisibleRemoteClient("notebook")).toBe(true) + expect(isNotebookVisibleRemoteClient("cli")).toBe(true) + expect(isNotebookVisibleRemoteClient("mcp")).toBe(true) + expect(isNotebookVisibleRemoteClient("node-sdk")).toBe(false) + expect(isNotebookVisibleRemoteClient("api")).toBe(false) + expect(isNotebookVisibleRemoteClient(undefined)).toBe(false) + expect(isNotebookVisibleRemoteMetadata({ createdByClient: "cli" })).toBe( + true, + ) + expect(isNotebookVisibleRemoteMetadata({ created_by_client: "api" })).toBe( + false, + ) + expect(isNotebookVisibleRemoteMetadata({})).toBe(false) + }) }) diff --git a/src/domains/sources/document-metadata.ts b/src/domains/sources/document-metadata.ts index d6b608f..b2be4d3 100644 --- a/src/domains/sources/document-metadata.ts +++ b/src/domains/sources/document-metadata.ts @@ -9,6 +9,19 @@ export const NOTEBOOK_DOCUMENT_METADATA_DEFAULTS = { clientVersion: packageJson.version, } as const; +export const NOTEBOOK_VISIBLE_CREATED_BY_CLIENTS = [ + "notebook", + "cli", + "mcp", +] as const; + +export type NotebookVisibleCreatedByClient = + (typeof NOTEBOOK_VISIBLE_CREATED_BY_CLIENTS)[number]; + +const notebookVisibleCreatedByClientSet = new Set( + NOTEBOOK_VISIBLE_CREATED_BY_CLIENTS, +); + export function createNotebookDocumentMetadata(input: { readonly title: string; readonly mimeType: string; @@ -24,3 +37,25 @@ export function createNotebookDocumentMetadata(input: { ...input.overrides, }; } + +export function getCreatedByClient( + metadata: Readonly> | undefined, +): string | undefined { + if (!metadata) return undefined; + const value = metadata.createdByClient ?? metadata.created_by_client; + return typeof value === "string" && value.length > 0 ? value : undefined; +} + +export function isNotebookVisibleRemoteClient( + client: string | undefined, +): client is NotebookVisibleCreatedByClient { + return ( + client !== undefined && notebookVisibleCreatedByClientSet.has(client) + ); +} + +export function isNotebookVisibleRemoteMetadata( + metadata: Readonly> | undefined, +): boolean { + return isNotebookVisibleRemoteClient(getCreatedByClient(metadata)); +} diff --git a/src/domains/sources/remote-library.test.ts b/src/domains/sources/remote-library.test.ts new file mode 100644 index 0000000..1a0df68 --- /dev/null +++ b/src/domains/sources/remote-library.test.ts @@ -0,0 +1,117 @@ +import { describe, expect, it, vi } from "vitest" +import { Effect } from "effect" + +import type { Source } from "@/infrastructure/db/schema" +import { + isNotebookVisibleRemoteDocument, + listRemoteLibrarySourceViews, +} from "./remote-library" + +const localSource: Source = { + id: "source_local", + workspaceId: "workspace_1", + title: "notes.pdf", + mimeType: "application/pdf", + sizeBytes: 5, + status: "ready", + failureReason: null, + failureStage: null, + knowhereJobId: null, + knowhereDocumentId: "doc_local", + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: null, + originalBlobUrl: null, + demoKey: null, + createdAt: new Date("2026-05-10T00:00:00Z"), + updatedAt: new Date("2026-05-10T00:00:00Z"), + deletedAt: null, +} + +describe("listRemoteLibrarySourceViews", () => { + it("keeps notebook, cli, and mcp remotes and hides other origins", async () => { + const list = vi.fn( + async (params?: { readonly namespace?: string }) => { + if (params?.namespace !== "default") { + return { documents: [] } + } + + return { + documents: [ + { + documentId: "doc_cli", + namespace: "default", + status: "active", + sourceFileName: "cli.pdf", + documentMetadata: { createdByClient: "cli" }, + }, + { + documentId: "doc_mcp", + namespace: "default", + status: "active", + sourceFileName: "mcp.pdf", + documentMetadata: { created_by_client: "mcp" }, + }, + { + documentId: "doc_notebook", + namespace: "default", + status: "active", + sourceFileName: "notebook.pdf", + documentMetadata: { createdByClient: "notebook" }, + }, + { + documentId: "doc_untagged", + namespace: "default", + status: "active", + sourceFileName: "dummy.pdf", + }, + { + documentId: "doc_sdk", + namespace: "default", + status: "active", + sourceFileName: "sdk.pdf", + documentMetadata: { createdByClient: "node-sdk" }, + }, + { + documentId: "doc_api", + namespace: "default", + status: "active", + sourceFileName: "api.pdf", + documentMetadata: { created_by_client: "api" }, + }, + { + documentId: "doc_local", + namespace: "default", + status: "active", + sourceFileName: "notes.pdf", + documentMetadata: { createdByClient: "cli" }, + }, + ], + } + }) + + const views = await Effect.runPromise( + listRemoteLibrarySourceViews({ + workspace: { namespace: "notebook-workspace_1" }, + client: { documents: { list } }, + localSources: [localSource], + }), + ) + + expect(views.map((view) => view.documentId)).toEqual([ + "doc_cli", + "doc_mcp", + "doc_notebook", + ]) + expect( + isNotebookVisibleRemoteDocument({ + documentMetadata: { createdByClient: "cli" }, + }), + ).toBe(true) + expect( + isNotebookVisibleRemoteDocument({ + documentMetadata: {}, + }), + ).toBe(false) + }) +}) diff --git a/src/domains/sources/remote-library.ts b/src/domains/sources/remote-library.ts index ceb9c80..77746d1 100644 --- a/src/domains/sources/remote-library.ts +++ b/src/domains/sources/remote-library.ts @@ -3,6 +3,10 @@ import { Effect } from "effect" import type { Source } from "@/infrastructure/db/schema" import type { SourceView } from "./types" import type { SourceStatus } from "./types" +import { + getCreatedByClient, + isNotebookVisibleRemoteMetadata, +} from "./document-metadata" import { getCompatibleNamespaces, sharedLibraryNamespace } from "./namespace" type RemoteDocument = { @@ -152,6 +156,7 @@ export function listRemoteLibrarySourceViews( ) const remoteDocuments = (yield* listRemoteLibraryDocuments(input)).filter( (document) => + isNotebookVisibleRemoteDocument(document) && !localDocumentIds.has(document.documentId) && !matchesActiveNotebookParsingSource(document, input.localSources), ) @@ -271,6 +276,12 @@ function normalizeRemoteDocument( } } +export function isNotebookVisibleRemoteDocument( + document: Pick, +): boolean { + return isNotebookVisibleRemoteMetadata(document.documentMetadata) +} + function toRemoteSourceView(document: RemoteDocument): SourceView { return { id: encodeRemoteSourceId(document), @@ -317,7 +328,7 @@ function matchesActiveNotebookParsingSource( document: RemoteDocument, localSources: readonly Source[], ): boolean { - if (document.documentMetadata?.createdByClient !== "notebook") return false + if (getCreatedByClient(document.documentMetadata) !== "notebook") return false if (!document.title || !document.mimeType || document.sizeBytes === undefined) { return false } diff --git a/src/domains/sources/repository.ts b/src/domains/sources/repository.ts index 3c6d01c..f26f9f7 100644 --- a/src/domains/sources/repository.ts +++ b/src/domains/sources/repository.ts @@ -6,6 +6,7 @@ import { sourceRowRepository } from "./source-row-repository" type SourceRepository = { readonly findInWorkspaceEffect: typeof sourceRowRepository.findInWorkspaceEffect + readonly findByKnowhereDocumentIdEffect: typeof sourceRowRepository.findByKnowhereDocumentIdEffect readonly listForWorkspaceEffect: typeof sourceRowRepository.listForWorkspaceEffect readonly createUploadingEffect: typeof sourceRowRepository.createUploadingEffect readonly localizeRemoteDocumentEffect: typeof sourceRowRepository.localizeRemoteDocumentEffect @@ -28,6 +29,8 @@ type SourceRepository = { export const sourceRepository: SourceRepository = { findInWorkspaceEffect: sourceRowRepository.findInWorkspaceEffect, + findByKnowhereDocumentIdEffect: + sourceRowRepository.findByKnowhereDocumentIdEffect, listForWorkspaceEffect: sourceRowRepository.listForWorkspaceEffect, createUploadingEffect: sourceRowRepository.createUploadingEffect, localizeRemoteDocumentEffect: diff --git a/src/domains/sources/route-archive.ts b/src/domains/sources/route-archive.ts index 6047256..dc968ec 100644 --- a/src/domains/sources/route-archive.ts +++ b/src/domains/sources/route-archive.ts @@ -1,7 +1,9 @@ import { Effect } from "effect" +import type { Source } from "@/infrastructure/db/schema" import { routeResult } from "@/lib/route-result" import { getClientForWorkspace } from "./route-dependencies" +import { decodeRemoteSourceId } from "./remote-library" import type { ArchiveSourceBody, ArchiveSourceInput, @@ -51,38 +53,83 @@ const archiveSourceEffect = ( deps.sourceService.findInWorkspace(workspace.id, input.sourceId), ) - if (!source) { - const catalog = yield* Effect.tryPromise(() => deps.demoApi.fetchCatalog()) - const isDemoSource = catalog.sources.some( - (candidate) => candidate.demoSourceId === input.sourceId, - ) - if (isDemoSource) { - yield* Effect.tryPromise(() => - deps.sourceService.hideDemoSource(workspace.id, input.sourceId), + if (source) { + if (source.knowhereDocumentId) { + yield* archiveKnowhereDocument( + workspace.id, + input.cookieHeader, + source.knowhereDocumentId, + deps, ) - return routeResult.ok({ id: input.sourceId, archived: true as const }) } - return routeResult.error(404, "Source not found.") + yield* cleanupLocalSource(workspace.id, source, deps) + return routeResult.ok({ id: input.sourceId, archived: true as const }) } - if (source.knowhereDocumentId) { - const client = yield* Effect.tryPromise(() => - getClientForWorkspace(workspace.id, input.cookieHeader, deps), - ) + const catalog = yield* Effect.tryPromise(() => deps.demoApi.fetchCatalog()) + const isDemoSource = catalog.sources.some( + (candidate) => candidate.demoSourceId === input.sourceId, + ) + if (isDemoSource) { yield* Effect.tryPromise(() => - client.documents.archive(source.knowhereDocumentId!), - ).pipe( - Effect.catchIf(isKnowhereDocumentNotFoundError, () => Effect.void), + deps.sourceService.hideDemoSource(workspace.id, input.sourceId), ) + return routeResult.ok({ id: input.sourceId, archived: true as const }) + } + + const remoteSource = decodeRemoteSourceId(input.sourceId) + if (!remoteSource) { + return routeResult.error(404, "Source not found.") } + yield* archiveKnowhereDocument( + workspace.id, + input.cookieHeader, + remoteSource.documentId, + deps, + ) + + const localSource = yield* Effect.tryPromise(() => + deps.sourceService.findByKnowhereDocumentId( + workspace.id, + remoteSource.documentId, + ), + ) + if (localSource) { + yield* cleanupLocalSource(workspace.id, localSource, deps) + } + + return routeResult.ok({ id: input.sourceId, archived: true as const }) + }) + +const archiveKnowhereDocument = ( + workspaceId: string, + cookieHeader: string, + documentId: string, + deps: RouteArchiveDependencies, +) => + Effect.gen(function* () { + const client = yield* Effect.tryPromise(() => + getClientForWorkspace(workspaceId, cookieHeader, deps), + ) + yield* Effect.tryPromise(() => client.documents.archive(documentId)).pipe( + Effect.catchIf(isKnowhereDocumentNotFoundError, () => Effect.void), + ) + }) + +const cleanupLocalSource = ( + workspaceId: string, + source: Source, + deps: RouteArchiveDependencies, +) => + Effect.gen(function* () { yield* Effect.tryPromise(() => - deps.sourceService.softDelete(workspace.id, input.sourceId), + deps.sourceService.softDelete(workspaceId, source.id), ) if (source.demoKey) { yield* Effect.tryPromise(() => - deps.sourceService.hideDemoSource(workspace.id, source.demoKey!), + deps.sourceService.hideDemoSource(workspaceId, source.demoKey!), ) } if (source.originalBlobPathname) { @@ -90,8 +137,6 @@ const archiveSourceEffect = ( deps.deleteBlob(source.originalBlobPathname!), ).pipe(Effect.catchAllCause(() => Effect.void)) } - - return routeResult.ok({ id: input.sourceId, archived: true as const }) }) function isKnowhereDocumentNotFoundError(error: unknown): boolean { diff --git a/src/domains/sources/route-dependencies.ts b/src/domains/sources/route-dependencies.ts index b0cc887..f55c131 100644 --- a/src/domains/sources/route-dependencies.ts +++ b/src/domains/sources/route-dependencies.ts @@ -45,6 +45,7 @@ const defaultDependencies: SourceRouteServiceDependencies = { requireUser, sourceService: { findInWorkspace: defaultSourceService.findInWorkspace, + findByKnowhereDocumentId: defaultSourceService.findByKnowhereDocumentId, hideDemoSource: defaultSourceService.hideDemoSource, listHiddenDemoSourceIds: defaultSourceService.listHiddenDemoSourceIds, localizeRemoteDocument: defaultSourceService.localizeRemoteDocument, diff --git a/src/domains/sources/route-service.test.ts b/src/domains/sources/route-service.test.ts index cb9c7da..64c902a 100644 --- a/src/domains/sources/route-service.test.ts +++ b/src/domains/sources/route-service.test.ts @@ -183,9 +183,34 @@ describe("source route service", () => { status: "active", sourceFileName: "cli.pdf", documentMetadata: { + createdByClient: "cli", mimeType: "application/pdf", }, }, + { + documentId: "doc_untagged", + namespace: "default", + status: "active", + sourceFileName: "dummy.pdf", + }, + { + documentId: "doc_sdk", + namespace: "default", + status: "active", + sourceFileName: "sdk.pdf", + documentMetadata: { + createdByClient: "node-sdk", + }, + }, + { + documentId: "doc_api", + namespace: "default", + status: "active", + sourceFileName: "api.pdf", + documentMetadata: { + created_by_client: "api", + }, + }, ], pagination: { page: 1, @@ -217,6 +242,9 @@ describe("source route service", () => { namespace: workspace.namespace, status: "active", sourceFileName: "legacy.pdf", + documentMetadata: { + createdByClient: "mcp", + }, }, ], pagination: { diff --git a/src/domains/sources/route-types.ts b/src/domains/sources/route-types.ts index 01d1883..6787ae4 100644 --- a/src/domains/sources/route-types.ts +++ b/src/domains/sources/route-types.ts @@ -188,6 +188,10 @@ type SourceWorkflowService = { workspaceId: string, sourceId: string, ) => Promise + readonly findByKnowhereDocumentId: ( + workspaceId: string, + documentId: string, + ) => Promise readonly softDelete: ( workspaceId: string, sourceId: string, diff --git a/src/domains/sources/service.ts b/src/domains/sources/service.ts index 4fcbfd5..1e4801c 100644 --- a/src/domains/sources/service.ts +++ b/src/domains/sources/service.ts @@ -17,6 +17,10 @@ type SourceService = { workspaceId: string, sourceId: string, ) => Promise + readonly findByKnowhereDocumentId: ( + workspaceId: string, + documentId: string, + ) => Promise readonly listForWorkspace: (workspaceId: string) => Promise readonly localizeRemoteDocument: ( workspaceId: string, @@ -96,6 +100,7 @@ const retrySourceToKnowhere: SourceService["retrySourceToKnowhere"] = ( export const sourceService: SourceService = { findInWorkspace: sourceWorkflowRuntime.findInWorkspace, + findByKnowhereDocumentId: sourceWorkflowRuntime.findByKnowhereDocumentId, hideDemoSource: sourceWorkflowRuntime.hideDemoSource, listHiddenDemoSourceIds: sourceWorkflowRuntime.listHiddenDemoSourceIds, listForWorkspace: sourceWorkflowRuntime.listForWorkspace, diff --git a/src/domains/sources/source-row-repository.ts b/src/domains/sources/source-row-repository.ts index 2c06765..256158f 100644 --- a/src/domains/sources/source-row-repository.ts +++ b/src/domains/sources/source-row-repository.ts @@ -51,6 +51,10 @@ type SourceRowRepository = { workspaceId: string, sourceId: string, ) => Effect.Effect + readonly findByKnowhereDocumentIdEffect: ( + workspaceId: string, + documentId: string, + ) => Effect.Effect readonly listForWorkspaceEffect: ( workspaceId: string, ) => Effect.Effect @@ -126,6 +130,15 @@ const findInWorkspaceEffect: SourceRowRepository["findInWorkspaceEffect"] = ( ) }) +const findByKnowhereDocumentIdEffect: SourceRowRepository["findByKnowhereDocumentIdEffect"] = + (workspaceId: string, documentId: string) => + Effect.gen(function* () { + const db = yield* DbClient + return yield* Effect.promise(() => + findByKnowhereDocumentIdWithDb(db, workspaceId, documentId), + ) + }) + const listForWorkspaceEffect: SourceRowRepository["listForWorkspaceEffect"] = ( workspaceId: string, ) => @@ -301,6 +314,28 @@ async function findInWorkspaceWithDb( return row[0] ?? null } +async function findByKnowhereDocumentIdWithDb( + db: Db, + workspaceId: string, + documentId: string, +): Promise { + if (documentId.length === 0) return null + + const row = await db + .select() + .from(sources) + .where( + and( + eq(sources.workspaceId, workspaceId), + eq(sources.knowhereDocumentId, documentId), + isNull(sources.deletedAt), + ), + ) + .limit(1) + + return row[0] ?? null +} + async function updateInWorkspaceWithDb( db: Db, workspaceId: string, @@ -414,6 +449,7 @@ function requireSource(source: Source | null, message: string): Source { export const sourceRowRepository: SourceRowRepository = { findInWorkspaceEffect, + findByKnowhereDocumentIdEffect, listForWorkspaceEffect, createUploadingEffect, localizeRemoteDocumentEffect, diff --git a/src/domains/sources/workflow-runtime.ts b/src/domains/sources/workflow-runtime.ts index 24464e0..9b0487e 100644 --- a/src/domains/sources/workflow-runtime.ts +++ b/src/domains/sources/workflow-runtime.ts @@ -59,6 +59,10 @@ type SourceWorkflowRuntime = UploadRepositoryRuntime & { workspaceId: string, sourceId: string, ) => Promise + readonly findByKnowhereDocumentId: ( + workspaceId: string, + documentId: string, + ) => Promise readonly getParseAssetUrls: ( workspaceId: string, sourceId: string, @@ -124,6 +128,12 @@ const findInWorkspace: SourceWorkflowRuntime["findInWorkspace"] = ( sourceRepository.findInWorkspaceEffect(workspaceId, sourceId), ) +const findByKnowhereDocumentId: SourceWorkflowRuntime["findByKnowhereDocumentId"] = + (workspaceId: string, documentId: string) => + databaseRuntime.runPromise( + sourceRepository.findByKnowhereDocumentIdEffect(workspaceId, documentId), + ) + const listForWorkspace: SourceWorkflowRuntime["listForWorkspace"] = ( workspaceId: string, ) => @@ -319,6 +329,7 @@ export const sourceWorkflowRuntime: SourceWorkflowRuntime = { createUploadRepository, createUploading, findInWorkspace, + findByKnowhereDocumentId, getParseAssetUrls, getParseResultProgress, getParseSnapshotMetadata, From f4569eaaaf9ad190ae8d6937157110f8b9b03396 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 20 Aug 2026 19:55:42 +0800 Subject: [PATCH 42/51] fix: keep overlapping page-memory sections in parsed results Only drop singleton page cards that repeat the same page number. Multi-page section chunks that share a boundary page stay visible. Co-authored-by: Cursor --- src/components/chunks-panel-state.test.ts | 88 ++++++++++++++++++++++- src/components/chunks-panel-state.ts | 35 ++++++--- src/components/chunks-panel.test.ts | 59 +++++++++++++++ 3 files changed, 170 insertions(+), 12 deletions(-) diff --git a/src/components/chunks-panel-state.test.ts b/src/components/chunks-panel-state.test.ts index ae75c71..febfd71 100644 --- a/src/components/chunks-panel-state.test.ts +++ b/src/components/chunks-panel-state.test.ts @@ -149,7 +149,7 @@ describe("chunksPanelState", () => { ]) }) - it("deduplicates page-asset chunks by page number", () => { + it("deduplicates singleton page-asset chunks with the same page number", () => { const chunks: ParsedChunkView[] = [ { chunkId: "page_4_first", @@ -208,6 +208,92 @@ describe("chunksPanelState", () => { ).toEqual(["page_4_first", "page_5"]) }) + it("keeps overlapping page-memory section chunks that share a boundary page", () => { + const chunks: ParsedChunkView[] = [ + { + chunkId: "kenneth", + type: "page", + content: "IR introduction.", + sectionPath: "call.pdf/Root/Kenneth Dorell", + sourceTitle: "call.pdf", + pageNums: [1], + pageAssets: [ + { + pageNumber: 1, + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "zuckerberg", + type: "page", + content: "[SAME-AS call.pdf/Root/Kenneth Dorell p1] CEO remarks.", + sectionPath: "call.pdf/Root/Mark Zuckerberg, CEO", + sourceTitle: "call.pdf", + pageNums: [1, 2], + pageAssets: [ + { + pageNumber: 1, + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + { + pageNumber: 2, + assetUrl: "https://assets.example/page-2.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "outlook", + type: "page", + content: "Q2 outlook.", + sectionPath: "call.pdf/Root/Moving to our financial outlook.", + sourceTitle: "call.pdf", + pageNums: [8], + pageAssets: [ + { + pageNumber: 8, + assetUrl: "https://assets.example/page-8.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "capex", + type: "page", + content: "[SAME-AS call.pdf/Root/Moving to our financial outlook. p8] Q&A.", + sectionPath: "call.pdf/Root/Turning to the expense and capex outlooks.", + sourceTitle: "call.pdf", + pageNums: [8, 9, 10], + pageAssets: [ + { + pageNumber: 8, + assetUrl: "https://assets.example/page-8.png", + contentType: "image/png", + }, + { + pageNumber: 9, + assetUrl: "https://assets.example/page-9.png", + contentType: "image/png", + }, + { + pageNumber: 10, + assetUrl: "https://assets.example/page-10.png", + contentType: "image/png", + }, + ], + }, + ] + + expect( + chunksPanelState + .getPageAssetChunksWithoutDuplicatePages(chunks) + .map((chunk) => chunk.chunkId), + ).toEqual(["kenneth", "zuckerberg", "outlook", "capex"]) + }) + it("hides table asset chunks from page-asset lists", () => { const chunks: ParsedChunkView[] = [ { diff --git a/src/components/chunks-panel-state.ts b/src/components/chunks-panel-state.ts index f11e855..3afccbf 100644 --- a/src/components/chunks-panel-state.ts +++ b/src/components/chunks-panel-state.ts @@ -192,29 +192,42 @@ function dedupeChunksById( function getPageAssetChunksWithoutDuplicatePages( chunks: readonly ParsedChunkView[], ): readonly ParsedChunkView[] { - const seenPageNumbers = new Set() + const seenSingletonPageNumbers = new Set() return chunks.filter((chunk) => { // Page-memory table assets currently store a file path, not HTML. if (chunk.type === "table") return false if (chunk.type !== "page") return true - const pageNumber = getPageAssetChunkPageNumber(chunk) - if (pageNumber === null) return true - if (seenPageNumbers.has(pageNumber)) return false + const pageNumbers = getPageAssetChunkPageNumbers(chunk) + if (pageNumbers.length === 0) return true + if (pageNumbers.length > 1) return true - seenPageNumbers.add(pageNumber) + const pageNumber = pageNumbers[0]! + if (seenSingletonPageNumbers.has(pageNumber)) return false + + seenSingletonPageNumbers.add(pageNumber) return true }) } -function getPageAssetChunkPageNumber(chunk: ParsedChunkView): number | null { - const pageAssetNumbers = (chunk.pageAssets ?? []) - .map((pageAsset) => pageAsset.pageNumber) - .filter(isPositivePageNumber) - if (pageAssetNumbers.length > 0) return Math.min(...pageAssetNumbers) +function getPageAssetChunkPageNumbers( + chunk: ParsedChunkView, +): readonly number[] { + const pageAssetNumbers = uniquePositivePageNumbers( + (chunk.pageAssets ?? []).map((pageAsset) => pageAsset.pageNumber), + ) + if (pageAssetNumbers.length > 0) return pageAssetNumbers + + return uniquePositivePageNumbers(chunk.pageNums ?? []) +} - return getFirstPageNumber(chunk) +function uniquePositivePageNumbers( + pageNumbers: readonly number[], +): readonly number[] { + return [...new Set(pageNumbers.filter(isPositivePageNumber))].sort( + (left, right) => left - right, + ) } function createMutableSectionTreeNode(input: { diff --git a/src/components/chunks-panel.test.ts b/src/components/chunks-panel.test.ts index 33367f1..a56da3f 100644 --- a/src/components/chunks-panel.test.ts +++ b/src/components/chunks-panel.test.ts @@ -198,6 +198,65 @@ describe("ChunksPanel", () => { .toBeNull(); }); + it("keeps overlapping page-memory sections that share a first page image", async () => { + mockVisibleVirtualViewport(); + + render( + React.createElement(C, { + chunks: [ + { + chunkId: "kenneth", + type: "page", + content: "IR introduction", + sectionPath: "call.pdf/Root/Kenneth Dorell", + sourceTitle: "call.pdf", + pageNums: [1], + pageAssets: [ + { + pageNumber: 1, + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "zuckerberg", + type: "page", + content: "CEO remarks", + sectionPath: "call.pdf/Root/Mark Zuckerberg, CEO", + sourceTitle: "call.pdf", + pageNums: [1, 2], + pageAssets: [ + { + pageNumber: 1, + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + { + pageNumber: 2, + assetUrl: "https://assets.example/page-2.png", + contentType: "image/png", + }, + ], + }, + ], + selectedSource: "call.pdf", + selectedSourceView: { + id: "source_1", + title: "call.pdf", + mimeType: "application/pdf", + status: "ready", + documentPresentation: { kind: "page-assets", pageCount: 2 }, + }, + }), + ); + + selectListView(); + expect(await screen.findByTestId("chunk-card-shell-kenneth")).toBeTruthy(); + expect(screen.getByTestId("chunk-card-shell-zuckerberg")).toBeTruthy(); + expect(screen.getByRole("img", { name: "Page 2" })).toBeTruthy(); + }); + it("renders page chunks normally when no page assets exist", async () => { mockVisibleVirtualViewport(); From 11d6f71eb32b87d0d4449e6cf4753601cd9d05e2 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 20 Aug 2026 16:09:45 +0800 Subject: [PATCH 43/51] feat: render demo seed answers with citation chips Map catalog page metadata onto demo chat citations and keep [[cite:n]] in the answer so seed Tesla/SpaceX turns match live {title}/pN chips. Co-authored-by: Cursor --- src/domains/demo/view.test.ts | 104 +++++++++++++++++++++++++ src/domains/demo/view.ts | 19 ++++- src/integrations/knowhere-demo.test.ts | 62 +++++++++++++++ src/integrations/knowhere-demo.ts | 15 ++++ 4 files changed, 199 insertions(+), 1 deletion(-) create mode 100644 src/domains/demo/view.test.ts diff --git a/src/domains/demo/view.test.ts b/src/domains/demo/view.test.ts new file mode 100644 index 0000000..04d2c7c --- /dev/null +++ b/src/domains/demo/view.test.ts @@ -0,0 +1,104 @@ +import { describe, expect, it } from "vitest" + +import { demoView } from "@/domains/demo/view" +import type { DemoCatalog, DemoCitation, DemoSource } from "@/integrations/knowhere-demo" + +describe("demoView.toChatMessages", () => { + it("embeds cite markers and page metadata so demo answers render title/pN chips", () => { + const messages = demoView.toChatMessages( + makeCatalog({ + answer: + "Tesla entered an agreement to invest about $2 billion in xAI. [[cite:1]]", + citations: [makeCitation({ pageCitationPageNumber: 12 })], + }), + ) + const assistant = messages.find((message) => message.role === "assistant") + + expect(assistant?.content).toContain("[[cite:1]]") + expect(assistant?.citations).toEqual([ + expect.objectContaining({ + pageCitationPageNumber: 12, + pageCitationAssetUrl: + "/api/demo-sources/demo-tsla-q4-2025/assets/page_citation_assets/page-12.png", + source: expect.objectContaining({ + sourceFileName: "TSLA-Q4-2025-Update.pdf", + sectionPath: "TSLA-Q4-2025-Update.pdf/OTHER UPDATES", + }), + }), + ]) + }) + + it("appends cite markers when the catalog answer still omits them", () => { + const messages = demoView.toChatMessages( + makeCatalog({ + answer: "Tesla entered an agreement to invest about $2 billion in xAI.", + citations: [makeCitation({ pageCitationPageNumber: 12 })], + }), + ) + const assistant = messages.find((message) => message.role === "assistant") + + expect(assistant?.content).toBe( + "Tesla entered an agreement to invest about $2 billion in xAI. [[cite:1]]", + ) + }) +}) + +function makeCatalog(input: { + readonly answer: string + readonly citations: readonly DemoCitation[] +}): DemoCatalog { + return { + officialLibrary: { categories: [], sources: [] }, + sources: [makeDemoSource(input)], + } +} + +function makeDemoSource(input: { + readonly answer: string + readonly citations: readonly DemoCitation[] +}): DemoSource { + return { + demoSourceId: "demo-tsla-q4-2025", + canonicalDocumentId: "demo-doc-tsla-q4-2025", + title: "TSLA-Q4-2025-Update.pdf", + mimeType: "application/pdf", + sizeBytes: 1024, + status: "ready", + chunkCount: 71, + originalFile: { + url: "/api/v1/demo/sources/demo-tsla-q4-2025/original", + mimeType: "application/pdf", + sizeBytes: 1024, + canDownload: false, + }, + examples: [ + { + id: "demo-tsla-q4-2025-xai", + question: "What does the document say about Tesla's xAI investment?", + answer: input.answer, + citations: input.citations, + }, + ], + } +} + +function makeCitation( + overrides: Partial = {}, +): DemoCitation { + return { + demoSourceId: "demo-tsla-q4-2025", + canonicalDocumentId: "demo-doc-tsla-q4-2025", + canonicalChunkId: "demo-tsla-q4-2025:chunk", + chunkId: "chunk", + chunkType: "page", + content: "Tesla entered into an agreement to invest approximately", + pageCitationAssetUrl: + "/api/demo-sources/demo-tsla-q4-2025/assets/page_citation_assets/page-12.png", + source: { + documentId: "demo-doc-tsla-q4-2025", + sourceFileName: "TSLA-Q4-2025-Update.pdf", + sectionPath: "TSLA-Q4-2025-Update.pdf/OTHER UPDATES", + }, + ...overrides, + } +} diff --git a/src/domains/demo/view.ts b/src/domains/demo/view.ts index 0b9a39b..fa35f9e 100644 --- a/src/domains/demo/view.ts +++ b/src/domains/demo/view.ts @@ -51,7 +51,7 @@ function toChatMessages(catalog: DemoCatalog): ChatMessageView[] { { id: `${example.id}-assistant`, role: "assistant", - content: example.answer, + content: withCitationMarkers(example.answer, example.citations.length), citations: example.citations.map((citation) => ({ chunkType: citation.chunkType, score: 0.95, @@ -59,6 +59,12 @@ function toChatMessages(catalog: DemoCatalog): ChatMessageView[] { ...(citation.description ? { description: citation.description } : {}), + ...(citation.pageCitationPageNumber + ? { pageCitationPageNumber: citation.pageCitationPageNumber } + : {}), + ...(citation.pageCitationAssetUrl + ? { pageCitationAssetUrl: citation.pageCitationAssetUrl } + : {}), source: { documentId: citation.canonicalDocumentId, sourceFileName: citation.source.sourceFileName, @@ -87,3 +93,14 @@ function toParsedChunkView( sourceTitle: source.title, }) } + +const citeMarkerPattern = /\[\[cite:\d+\]\]/ + +function withCitationMarkers(answer: string, citationCount: number): string { + if (citationCount < 1 || citeMarkerPattern.test(answer)) return answer + const markers = Array.from( + { length: citationCount }, + (_, index) => `[[cite:${index + 1}]]`, + ).join(" ") + return `${answer.trimEnd()} ${markers}` +} diff --git a/src/integrations/knowhere-demo.test.ts b/src/integrations/knowhere-demo.test.ts index fe4f7e9..be47459 100644 --- a/src/integrations/knowhere-demo.test.ts +++ b/src/integrations/knowhere-demo.test.ts @@ -252,6 +252,68 @@ describe("knowhereDemoApi", () => { chunkCount: 922, }) }) + + it("maps demo example citation page metadata onto Notebook citation fields", async () => { + globalThis.fetch = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + sources: [ + { + demo_source_id: "demo-tsla-q4-2025", + canonical_document_id: "demo-doc-tsla-q4-2025", + title: "TSLA-Q4-2025-Update.pdf", + mime_type: "application/pdf", + size_bytes: 1024, + status: "ready", + chunk_count: 71, + original_file: { + url: "/api/v1/demo/sources/demo-tsla-q4-2025/original", + mime_type: "application/pdf", + size_bytes: 1024, + can_download: false, + }, + examples: [ + { + id: "demo-tsla-q4-2025-xai", + question: "What does the document say about Tesla's xAI investment?", + answer: + "Tesla entered an agreement to invest about $2 billion. [[cite:1]]", + citations: [ + { + demo_source_id: "demo-tsla-q4-2025", + canonical_document_id: "demo-doc-tsla-q4-2025", + canonical_chunk_id: "demo-tsla-q4-2025:chunk", + chunk_id: "chunk", + chunk_type: "page", + content: "Tesla entered into an agreement", + page_citation_page_number: 12, + page_citation_asset_url: + "/api/v1/demo/sources/demo-tsla-q4-2025/assets/page_citation_assets/page-12.png", + source: { + document_id: "demo-doc-tsla-q4-2025", + source_file_name: "TSLA-Q4-2025-Update.pdf", + section_path: "TSLA-Q4-2025-Update.pdf/OTHER UPDATES", + }, + }, + ], + }, + ], + }, + ], + official_library: { categories: [], sources: [] }, + }), + { status: 200, headers: { "content-type": "application/json" } }, + ), + ) + + const catalog = await knowhereDemoApi.fetchCatalog() + + expect(catalog.sources[0]?.examples[0]?.citations[0]).toMatchObject({ + pageCitationPageNumber: 12, + pageCitationAssetUrl: + "/api/demo-sources/demo-tsla-q4-2025/assets/page_citation_assets/page-12.png", + }) + }) }) function restoreEnv(key: string, value: string | undefined): void { diff --git a/src/integrations/knowhere-demo.ts b/src/integrations/knowhere-demo.ts index f4342a5..a63dcac 100644 --- a/src/integrations/knowhere-demo.ts +++ b/src/integrations/knowhere-demo.ts @@ -11,6 +11,8 @@ export type DemoCitation = { readonly chunkType: string readonly content: string readonly description?: string + readonly pageCitationPageNumber?: number + readonly pageCitationAssetUrl?: string readonly source: { readonly documentId: string readonly sourceFileName: string @@ -155,6 +157,8 @@ type DemoCitationResponse = { readonly chunk_type?: unknown readonly content?: unknown readonly description?: unknown + readonly page_citation_page_number?: unknown + readonly page_citation_asset_url?: unknown readonly source?: { readonly document_id?: unknown readonly source_file_name?: unknown @@ -442,6 +446,13 @@ function toDemoExample(example: DemoExampleResponse): DemoExample { function toDemoCitation(citation: DemoCitationResponse): DemoCitation { const source = citation.source ?? {} const description = optionalString(citation.description) + const pageCitationPageNumber = optionalPositiveNumber( + citation.page_citation_page_number, + ) + const pageCitationAssetUrl = toDemoAssetUrl( + requireString(citation.demo_source_id), + optionalString(citation.page_citation_asset_url), + ) return { demoSourceId: requireString(citation.demo_source_id), canonicalDocumentId: requireString(citation.canonical_document_id), @@ -450,6 +461,10 @@ function toDemoCitation(citation: DemoCitationResponse): DemoCitation { chunkType: requireString(citation.chunk_type), content: requireString(citation.content), ...(description ? { description } : {}), + ...(pageCitationPageNumber !== undefined + ? { pageCitationPageNumber } + : {}), + ...(pageCitationAssetUrl ? { pageCitationAssetUrl } : {}), source: { documentId: requireString(source.document_id), sourceFileName: requireString(source.source_file_name), From 2692fbcd22c17e3f99773f329eb476ff7c045c9f Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 21 Aug 2026 00:32:42 +0800 Subject: [PATCH 44/51] feat: map demo seed citations through the live chat citation path Reuse toChatCitationViews so demo Q&A chips get the same title/pN page numbers as normal retrieval answers. Co-authored-by: Cursor --- src/domains/demo/view.test.ts | 40 ++++++++++- src/domains/demo/view.ts | 92 +++++++++++++++++--------- src/integrations/knowhere-demo.test.ts | 2 + src/integrations/knowhere-demo.ts | 23 ++++++- 4 files changed, 120 insertions(+), 37 deletions(-) diff --git a/src/domains/demo/view.test.ts b/src/domains/demo/view.test.ts index 04d2c7c..593ff12 100644 --- a/src/domains/demo/view.test.ts +++ b/src/domains/demo/view.test.ts @@ -1,10 +1,11 @@ import { describe, expect, it } from "vitest" +import { chatCitationModel } from "@/components/chat-citation-model" import { demoView } from "@/domains/demo/view" import type { DemoCatalog, DemoCitation, DemoSource } from "@/integrations/knowhere-demo" describe("demoView.toChatMessages", () => { - it("embeds cite markers and page metadata so demo answers render title/pN chips", () => { + it("uses the live citation mapper so demo answers render title/pN chips", () => { const messages = demoView.toChatMessages( makeCatalog({ answer: @@ -26,6 +27,39 @@ describe("demoView.toChatMessages", () => { }), }), ]) + expect( + chatCitationModel.getCitationChipLabel(assistant!.citations![0]!, {}), + ).toBe("TSLA-Q4-2025-Update.pdf/p12") + }) + + it("resolves page chips from catalog page_nums when the explicit page field is omitted", () => { + const messages = demoView.toChatMessages( + makeCatalog({ + title: "spacex-s1.pdf", + answer: + "The filing says SpaceX operates about 9,600 Starlink satellites. [[cite:1]]", + citations: [ + makeCitation({ + demoSourceId: "demo-spacex-s1", + canonicalDocumentId: "demo-doc-spacex-s1", + pageCitationPageNumber: undefined, + pageCitationAssetUrl: undefined, + pageNums: [28], + source: { + documentId: "demo-doc-spacex-s1", + sourceFileName: "spacex-s1.pdf", + sectionPath: "spacex-s1.pdf/Root", + }, + }), + ], + }), + ) + const assistant = messages.find((message) => message.role === "assistant") + + expect(assistant?.citations?.[0]?.pageCitationPageNumber).toBe(28) + expect( + chatCitationModel.getCitationChipLabel(assistant!.citations![0]!, {}), + ).toBe("spacex-s1.pdf/p28") }) it("appends cite markers when the catalog answer still omits them", () => { @@ -44,6 +78,7 @@ describe("demoView.toChatMessages", () => { }) function makeCatalog(input: { + readonly title?: string readonly answer: string readonly citations: readonly DemoCitation[] }): DemoCatalog { @@ -54,13 +89,14 @@ function makeCatalog(input: { } function makeDemoSource(input: { + readonly title?: string readonly answer: string readonly citations: readonly DemoCitation[] }): DemoSource { return { demoSourceId: "demo-tsla-q4-2025", canonicalDocumentId: "demo-doc-tsla-q4-2025", - title: "TSLA-Q4-2025-Update.pdf", + title: input.title ?? "TSLA-Q4-2025-Update.pdf", mimeType: "application/pdf", sizeBytes: 1024, status: "ready", diff --git a/src/domains/demo/view.ts b/src/domains/demo/view.ts index fa35f9e..129c7fe 100644 --- a/src/domains/demo/view.ts +++ b/src/domains/demo/view.ts @@ -1,10 +1,13 @@ -import { parsedChunkNormalization } from "@/domains/chunks/normalization" +import { toChatCitationViews } from "@/domains/chat/citations" +import type { PageCitationAssetRetrievalResult } from "@/domains/chat/page-citation-assets" import type { ChatMessageView } from "@/domains/chat/types" +import { parsedChunkNormalization } from "@/domains/chunks/normalization" import type { ParsedChunkView } from "@/domains/chunks/types" import { demoOriginalFile } from "@/domains/demo/original-file" import type { SourceView } from "@/domains/sources/types" import type { DemoCatalog, + DemoCitation, DemoChunk, DemoSource, } from "@/integrations/knowhere-demo" @@ -42,40 +45,65 @@ function toSourceView(source: DemoSource): SourceView { function toChatMessages(catalog: DemoCatalog): ChatMessageView[] { return catalog.sources.flatMap((source) => - source.examples.flatMap((example): ChatMessageView[] => [ - { - id: `${example.id}-user`, - role: "user", - content: example.question, - }, - { - id: `${example.id}-assistant`, - role: "assistant", - content: withCitationMarkers(example.answer, example.citations.length), - citations: example.citations.map((citation) => ({ - chunkType: citation.chunkType, - score: 0.95, - content: citation.content, - ...(citation.description - ? { description: citation.description } - : {}), - ...(citation.pageCitationPageNumber - ? { pageCitationPageNumber: citation.pageCitationPageNumber } - : {}), - ...(citation.pageCitationAssetUrl - ? { pageCitationAssetUrl: citation.pageCitationAssetUrl } - : {}), - source: { - documentId: citation.canonicalDocumentId, - sourceFileName: citation.source.sourceFileName, - sectionPath: citation.source.sectionPath, - }, - })), - }, - ]), + source.examples.flatMap((example): ChatMessageView[] => { + const answer = withCitationMarkers( + example.answer, + example.citations.length, + ) + return [ + { + id: `${example.id}-user`, + role: "user", + content: example.question, + }, + { + id: `${example.id}-assistant`, + role: "assistant", + content: answer, + citations: toDemoChatCitationViews(example.citations, answer), + }, + ] + }), ) } +function toDemoChatCitationViews( + citations: readonly DemoCitation[], + answer: string, +) { + return toChatCitationViews( + citations.map(toDemoRetrievalResult), + answer, + ).map((citation, index) => { + const description = citations[index]?.description + return description ? { ...citation, description } : citation + }) +} + +function toDemoRetrievalResult( + citation: DemoCitation, +): PageCitationAssetRetrievalResult { + return { + content: citation.content, + chunkType: citation.chunkType, + score: 0.95, + ...(citation.pageCitationAssetUrl + ? { pageCitationAssetUrl: citation.pageCitationAssetUrl } + : {}), + ...(citation.pageCitationPageNumber + ? { pageCitationPageNumber: citation.pageCitationPageNumber } + : {}), + ...(citation.pageNums && citation.pageNums.length > 0 + ? { metadata: { page_nums: [...citation.pageNums] } } + : {}), + source: { + documentId: citation.canonicalDocumentId, + sourceFileName: citation.source.sourceFileName, + sectionPath: citation.source.sectionPath, + }, + } +} + function toParsedChunkView( source: SourceView, chunk: DemoChunk, diff --git a/src/integrations/knowhere-demo.test.ts b/src/integrations/knowhere-demo.test.ts index be47459..aa56041 100644 --- a/src/integrations/knowhere-demo.test.ts +++ b/src/integrations/knowhere-demo.test.ts @@ -287,6 +287,7 @@ describe("knowhereDemoApi", () => { chunk_type: "page", content: "Tesla entered into an agreement", page_citation_page_number: 12, + page_nums: [12], page_citation_asset_url: "/api/v1/demo/sources/demo-tsla-q4-2025/assets/page_citation_assets/page-12.png", source: { @@ -310,6 +311,7 @@ describe("knowhereDemoApi", () => { expect(catalog.sources[0]?.examples[0]?.citations[0]).toMatchObject({ pageCitationPageNumber: 12, + pageNums: [12], pageCitationAssetUrl: "/api/demo-sources/demo-tsla-q4-2025/assets/page_citation_assets/page-12.png", }) diff --git a/src/integrations/knowhere-demo.ts b/src/integrations/knowhere-demo.ts index a63dcac..83c56ee 100644 --- a/src/integrations/knowhere-demo.ts +++ b/src/integrations/knowhere-demo.ts @@ -13,6 +13,7 @@ export type DemoCitation = { readonly description?: string readonly pageCitationPageNumber?: number readonly pageCitationAssetUrl?: string + readonly pageNums?: readonly number[] readonly source: { readonly documentId: string readonly sourceFileName: string @@ -159,6 +160,7 @@ type DemoCitationResponse = { readonly description?: unknown readonly page_citation_page_number?: unknown readonly page_citation_asset_url?: unknown + readonly page_nums?: unknown readonly source?: { readonly document_id?: unknown readonly source_file_name?: unknown @@ -446,9 +448,9 @@ function toDemoExample(example: DemoExampleResponse): DemoExample { function toDemoCitation(citation: DemoCitationResponse): DemoCitation { const source = citation.source ?? {} const description = optionalString(citation.description) - const pageCitationPageNumber = optionalPositiveNumber( - citation.page_citation_page_number, - ) + const pageNums = toPositiveIntegers(citation.page_nums) + const pageCitationPageNumber = + optionalPositiveInteger(citation.page_citation_page_number) ?? pageNums[0] const pageCitationAssetUrl = toDemoAssetUrl( requireString(citation.demo_source_id), optionalString(citation.page_citation_asset_url), @@ -465,6 +467,7 @@ function toDemoCitation(citation: DemoCitationResponse): DemoCitation { ? { pageCitationPageNumber } : {}), ...(pageCitationAssetUrl ? { pageCitationAssetUrl } : {}), + ...(pageNums.length > 0 ? { pageNums } : {}), source: { documentId: requireString(source.document_id), sourceFileName: requireString(source.source_file_name), @@ -688,6 +691,20 @@ function optionalPositiveNumber(value: unknown): number | undefined { : undefined } +function optionalPositiveInteger(value: unknown): number | undefined { + return typeof value === "number" && Number.isSafeInteger(value) && value > 0 + ? value + : undefined +} + +function toPositiveIntegers(value: unknown): readonly number[] { + if (!Array.isArray(value)) return [] + return value.flatMap((item) => { + const page = optionalPositiveInteger(item) + return page === undefined ? [] : [page] + }) +} + function toDemoAssetUrl( demoSourceId: string, assetUrl: string | undefined, From 0ef6314e4b4e8ca5cf884260c7d28cbbd696a6b6 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 21 Aug 2026 00:45:40 +0800 Subject: [PATCH 45/51] test: expect guest demo seed answers to include cite markers withCitationMarkers appends [[cite:n]] when the catalog answer omits them so chips can render on guest demo chat. Co-authored-by: Cursor --- src/domains/workspace/initial-state.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/domains/workspace/initial-state.test.ts b/src/domains/workspace/initial-state.test.ts index 9fc1450..0ed05e2 100644 --- a/src/domains/workspace/initial-state.test.ts +++ b/src/domains/workspace/initial-state.test.ts @@ -85,7 +85,7 @@ describe("loadWorkspaceShellInitialState", () => { { id: "demo-example-1-assistant", role: "assistant", - content: "Tesla delivered higher revenue.", + content: "Tesla delivered higher revenue. [[cite:1]]", citations: [ { chunkType: "text", From 16ee7133addf3bc428c73a2bed4db40baf8117ec Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 21 Aug 2026 01:04:35 +0800 Subject: [PATCH 46/51] fix: keep Official Library demo sources selected on citation click Tesla and SpaceX seed citations point at catalog rows tagged as Official Library. Resolving those ids back to null left the parsed panel on the guest empty login state instead of the cited chunk. Co-authored-by: Cursor --- src/components/workspace-shell.test.ts | 91 +++++++++++++++++++ src/components/workspace-source-state.test.ts | 4 +- src/components/workspace-source-state.ts | 2 +- 3 files changed, 94 insertions(+), 3 deletions(-) diff --git a/src/components/workspace-shell.test.ts b/src/components/workspace-shell.test.ts index e8a4fcb..660471f 100644 --- a/src/components/workspace-shell.test.ts +++ b/src/components/workspace-shell.test.ts @@ -449,6 +449,97 @@ describe("WorkspaceShell", () => { ).toBe(false); }); + it("focuses guest Official Library demo citations instead of the empty login state", async () => { + const fetch = vi.fn(async (input) => { + const url = getRequestURL(input); + + if (url.pathname === "/api/sources/demo-tsla-q4-2025/chunks") { + return Response.json({ + chunks: [ + { + chunkId: "demo-tsla-q4-2025:chunk_1", + documentId: "demo-doc-tsla-q4-2025", + sectionPath: "TSLA-Q4-2025-Update.pdf/OTHER UPDATES", + type: "text", + content: "Tesla entered into an agreement to invest approximately", + sourceTitle: "TSLA-Q4-2025-Update.pdf", + }, + ], + pagination: { + page: Number(url.searchParams.get("page") ?? "1"), + pageSize: 100, + total: 1, + totalPages: 1, + }, + }); + } + + return Response.json({ message: "Unexpected request" }, { status: 404 }); + }); + vi.stubGlobal("fetch", fetch); + + render( + React.createElement(C, { + isGuest: true, + sources: [ + { + id: "demo-tsla-q4-2025", + kind: "demo", + demoSourceId: "demo-tsla-q4-2025", + title: "TSLA-Q4-2025-Update.pdf", + status: "ready", + mimeType: "application/pdf", + documentId: "demo-doc-tsla-q4-2025", + officialLibrary: { + librarySourceId: "financial-tsla-q4-2025", + categoryId: "financial-reports", + sourceUrl: "https://example.com/tsla-q4-2025.pdf", + }, + }, + ], + chatMessages: [ + { + id: "assistant_1", + role: "assistant", + content: "Tesla invested in xAI. [[cite:1]]", + citations: [ + { + content: "Tesla entered into an agreement to invest approximately", + description: "xAI investment", + chunkType: "text", + score: 0.95, + pageCitationPageNumber: 12, + source: { + documentId: "demo-doc-tsla-q4-2025", + sourceFileName: "TSLA-Q4-2025-Update.pdf", + sectionPath: "TSLA-Q4-2025-Update.pdf/OTHER UPDATES", + }, + }, + ], + }, + ], + }), + ); + + const citationButton = await findStableConnectedElement(() => { + const desktopChatPanel = within(screen.getByTestId("desktop-chat-panel")); + return desktopChatPanel.getByRole("button", { + name: "Open source TSLA-Q4-2025-Update.pdf/p12", + }); + }); + fireEvent.click(citationButton); + + await waitFor(() => { + const chunksPanel = screen.getByTestId("desktop-chunks-panel"); + expect(within(chunksPanel).queryByText("Log in to add documents")).toBeNull(); + const topRow = chunksPanel.querySelector('[data-index="0"]'); + expect(topRow?.getAttribute("data-chunk-id")).toBe( + "demo-tsla-q4-2025:chunk_1", + ); + expect(topRow?.getAttribute("data-focused-chunk")).toBe("true"); + }); + }); + it("focuses guest citations from the mobile chat panel", async () => { const fetch = vi.fn(async (input) => { const url = getRequestURL(input); diff --git a/src/components/workspace-source-state.test.ts b/src/components/workspace-source-state.test.ts index 5345646..03d9319 100644 --- a/src/components/workspace-source-state.test.ts +++ b/src/components/workspace-source-state.test.ts @@ -58,7 +58,7 @@ describe("workspaceSourceState", () => { ); }); - it("does not resolve an unmaterialized Official Library row as the selected Source", () => { + it("keeps an explicitly selected Official Library demo source for citation jumps", () => { const sources: readonly SourceView[] = [ { id: "demo-spacex-s1", @@ -88,7 +88,7 @@ describe("workspaceSourceState", () => { sources, "demo-spacex-s1", ), - ).toBeNull(); + ).toBe("demo-spacex-s1"); }); it("selects a preferred document source when opening a chunk-tree link", () => { diff --git a/src/components/workspace-source-state.ts b/src/components/workspace-source-state.ts index 5d33744..41607f3 100644 --- a/src/components/workspace-source-state.ts +++ b/src/components/workspace-source-state.ts @@ -74,7 +74,7 @@ function getResolvedSelectedSourceId( const selectedSource = sources.find((source) => source.id === selectedSourceId) if (selectedSource) { - return isVisibleSource(selectedSource) ? selectedSource.id : null + return selectedSource.id } const selectedDocumentId = getRemoteSourceDocumentId(selectedSourceId) From af700557cf35e8e0e8a6e7c7d8675a71861a7616 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 21 Aug 2026 12:22:27 +0800 Subject: [PATCH 47/51] feat: overlay answer provenance highlights on displayed page images Extend inspectImage with structured bbox regions, persist them on finalize artifacts, and render red translucent overlays in chat. Fall back to analysis-only when structured vision output is unavailable. Co-authored-by: Cursor --- src/agent-harness/image-highlights.test.ts | 93 +++++ src/agent-harness/image-highlights.ts | 104 +++++ src/agent-harness/index.ts | 1 + src/agent-harness/runtime.ts | 9 +- src/agent-harness/types.ts | 19 + src/components/chat-message-list.test.ts | 56 +++ src/components/chat-message-list.tsx | 73 +++- src/domains/chat/chat-citation-persistence.ts | 1 + .../chat/image-inspection-model.test.ts | 298 +++++++++++++ src/domains/chat/image-inspection-model.ts | 392 ++++++++++++++++++ src/domains/chat/index.test.ts | 1 + src/domains/chat/index.ts | 16 + src/domains/chat/route-answer.ts | 80 +--- src/domains/chat/route-service.test.ts | 45 +- src/domains/chat/service.test.ts | 1 + src/domains/chat/types.ts | 12 + src/domains/chat/view.ts | 33 ++ 17 files changed, 1152 insertions(+), 82 deletions(-) create mode 100644 src/agent-harness/image-highlights.test.ts create mode 100644 src/agent-harness/image-highlights.ts create mode 100644 src/domains/chat/image-inspection-model.test.ts create mode 100644 src/domains/chat/image-inspection-model.ts diff --git a/src/agent-harness/image-highlights.test.ts b/src/agent-harness/image-highlights.test.ts new file mode 100644 index 0000000..631ccc1 --- /dev/null +++ b/src/agent-harness/image-highlights.test.ts @@ -0,0 +1,93 @@ +import { describe, expect, it } from "vitest" + +import { + mergeImageInspectionHighlights, + normalizeHighlightBoxes, + normalizeImageInspectionHighlights, +} from "./image-highlights" + +describe("image highlights", () => { + it("normalizes multi-page multi-region boxes and drops invalid refs", () => { + const highlights = normalizeImageInspectionHighlights({ + allowedRefs: new Set(["asset:page-1", "asset:page-2"]), + pages: [ + { + ref: "asset:page-1", + regions: [ + { x: 0.1, y: 0.2, w: 0.3, h: 0.1 }, + { x: 0.5, y: 0.6, w: 0.2, h: 0.15 }, + ], + }, + { + ref: "asset:page-2", + regions: [{ x: -0.1, y: 0.9, w: 0.3, h: 0.2 }], + }, + { + ref: "asset:unknown", + regions: [{ x: 0.1, y: 0.1, w: 0.2, h: 0.2 }], + }, + ], + }) + + expect(highlights).toEqual([ + { + ref: "asset:page-1", + regions: [ + { x: 0.1, y: 0.2, w: 0.3, h: 0.1 }, + { x: 0.5, y: 0.6, w: 0.2, h: 0.15 }, + ], + }, + { + ref: "asset:page-2", + regions: [{ x: 0, y: 0.9, w: 0.2, h: 0.1 }], + }, + ]) + }) + + it("clamps standalone highlight boxes for persistence reload", () => { + expect( + normalizeHighlightBoxes([ + { x: -0.2, y: 0.5, w: 0.4, h: 0.2 }, + { x: 0.1, y: 0.1, w: 0.001, h: 0.2 }, + ]), + ).toEqual([{ x: 0, y: 0.5, w: 0.2, h: 0.2 }]) + }) + + it("replaces prior regions for the same ref on later inspects", () => { + expect( + mergeImageInspectionHighlights( + [ + { + ref: "asset:page-1", + regions: [{ x: 0.1, y: 0.1, w: 0.2, h: 0.2 }], + }, + ], + [ + { + ref: "asset:page-1", + regions: [ + { x: 0.2, y: 0.3, w: 0.1, h: 0.1 }, + { x: 0.5, y: 0.5, w: 0.1, h: 0.1 }, + ], + }, + { + ref: "asset:page-3", + regions: [{ x: 0.1, y: 0.1, w: 0.2, h: 0.2 }], + }, + ], + ), + ).toEqual([ + { + ref: "asset:page-1", + regions: [ + { x: 0.2, y: 0.3, w: 0.1, h: 0.1 }, + { x: 0.5, y: 0.5, w: 0.1, h: 0.1 }, + ], + }, + { + ref: "asset:page-3", + regions: [{ x: 0.1, y: 0.1, w: 0.2, h: 0.2 }], + }, + ]) + }) +}) diff --git a/src/agent-harness/image-highlights.ts b/src/agent-harness/image-highlights.ts new file mode 100644 index 0000000..8adc037 --- /dev/null +++ b/src/agent-harness/image-highlights.ts @@ -0,0 +1,104 @@ +import type { ImageHighlightBox, ImageInspectionHighlights } from "./types" + +const MIN_BOX_SIZE = 0.01 + +export function normalizeImageInspectionHighlights(input: { + readonly pages: readonly { + readonly ref: string + readonly regions: readonly { + readonly x: number + readonly y: number + readonly w: number + readonly h: number + }[] + }[] + readonly allowedRefs: ReadonlySet +}): ImageInspectionHighlights[] { + const byRef = new Map() + + for (const page of input.pages) { + const ref = page.ref.trim() + if (!ref || !input.allowedRefs.has(ref)) continue + + const regions = normalizeHighlightBoxes(page.regions) + if (regions.length === 0) continue + + const existing = byRef.get(ref) ?? [] + byRef.set(ref, [...existing, ...regions]) + } + + return Array.from(byRef.entries()).map(([ref, regions]) => ({ + ref, + regions, + })) +} + +/** Clamp/normalize boxes for persistence and reload paths. */ +export function normalizeHighlightBoxes( + regions: readonly { + readonly x: number + readonly y: number + readonly w: number + readonly h: number + }[], +): ImageHighlightBox[] { + return regions + .map(clampHighlightBox) + .filter((box): box is ImageHighlightBox => box !== null) +} + +export function mergeImageInspectionHighlights( + existing: readonly ImageInspectionHighlights[] | undefined, + incoming: readonly ImageInspectionHighlights[] | undefined, +): ImageInspectionHighlights[] { + const byRef = new Map() + + for (const page of existing ?? []) { + byRef.set(page.ref, [...page.regions]) + } + for (const page of incoming ?? []) { + // Latest inspect for a ref replaces prior regions for that ref. + byRef.set(page.ref, [...page.regions]) + } + + return Array.from(byRef.entries()).map(([ref, regions]) => ({ + ref, + regions, + })) +} + +function clampHighlightBox(input: { + readonly x: number + readonly y: number + readonly w: number + readonly h: number +}): ImageHighlightBox | null { + if ( + !Number.isFinite(input.x) || + !Number.isFinite(input.y) || + !Number.isFinite(input.w) || + !Number.isFinite(input.h) + ) { + return null + } + + const x1 = clamp01(Math.min(input.x, input.x + input.w)) + const y1 = clamp01(Math.min(input.y, input.y + input.h)) + const x2 = clamp01(Math.max(input.x, input.x + input.w)) + const y2 = clamp01(Math.max(input.y, input.y + input.h)) + const w = x2 - x1 + const h = y2 - y1 + if (w < MIN_BOX_SIZE || h < MIN_BOX_SIZE) return null + + return { x: round01(x1), y: round01(y1), w: round01(w), h: round01(h) } +} + +function clamp01(value: number): number { + if (value <= 0) return 0 + if (value >= 1) return 1 + return value +} + +function round01(value: number): number { + return Math.round(value * 1_000_000) / 1_000_000 +} diff --git a/src/agent-harness/index.ts b/src/agent-harness/index.ts index 5fbdf7a..a2666bd 100644 --- a/src/agent-harness/index.ts +++ b/src/agent-harness/index.ts @@ -1,3 +1,4 @@ +export * from "./image-highlights" export * from "./ledger" export * from "./knowhere-text" export * from "./runtime" diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index 0ad95d8..8dbe18d 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -10,6 +10,7 @@ import { z } from "zod" import { createEvidenceLedger } from "./ledger" import { knowhereToolText } from "./knowhere-text" +import { mergeImageInspectionHighlights } from "./image-highlights" import type { AgentTurn, AgentTurnInput, @@ -50,6 +51,7 @@ type HarnessToolState = { finalized?: boolean priorTurnReads?: string[] inspectedImageRefs?: string[] + imageHighlights?: ImageInspectionHighlights[] toolCalls?: HarnessToolCallTrace[] } @@ -246,6 +248,7 @@ export async function runAgentHarness( finalized: state.finalized === true, priorTurnReads: [...(state.priorTurnReads ?? [])], toolCalls: [...(state.toolCalls ?? [])], + imageHighlights: [...(state.imageHighlights ?? [])], validationErrors: [], revisionsUsed: 0, }, @@ -761,8 +764,12 @@ async function inspectRetrievedImages(input: { question, assets: selectedAssets, }) + input.state.imageHighlights = mergeImageInspectionHighlights( + input.state.imageHighlights, + response.highlights, + ) return { - ok: true, + ok: true as const, analysis: response.analysis, inspected: response.inspected, skipped: [...skipped, ...response.skipped], diff --git a/src/agent-harness/types.ts b/src/agent-harness/types.ts index 36b6a27..edd9c2f 100644 --- a/src/agent-harness/types.ts +++ b/src/agent-harness/types.ts @@ -182,6 +182,23 @@ export type ImageInspectionInspectedAsset = { readonly label: string } +/** + * Normalized box relative to image width/height. Origin is top-left. + * Values are clamped to [0, 1] before render. + */ +export type ImageHighlightBox = { + readonly x: number + readonly y: number + readonly w: number + readonly h: number +} + +/** One page/image may contain multiple answer regions; no per-region labels. */ +export type ImageInspectionHighlights = { + readonly ref: string + readonly regions: readonly ImageHighlightBox[] +} + export type ImageInspectionRequest = { readonly question: string readonly assets: readonly ImageInspectionAsset[] @@ -191,6 +208,7 @@ export type ImageInspectionResponse = { readonly analysis: string readonly inspected: readonly ImageInspectionInspectedAsset[] readonly skipped: readonly ImageInspectionSkippedAsset[] + readonly highlights?: readonly ImageInspectionHighlights[] } export type InspectImages = ( @@ -256,6 +274,7 @@ export type HarnessTrace = { readonly finalized: boolean readonly priorTurnReads: readonly string[] readonly toolCalls: readonly HarnessToolCallTrace[] + readonly imageHighlights: readonly ImageInspectionHighlights[] readonly validationErrors: readonly string[] readonly revisionsUsed: number } diff --git a/src/components/chat-message-list.test.ts b/src/components/chat-message-list.test.ts index e59fd54..1f665fe 100644 --- a/src/components/chat-message-list.test.ts +++ b/src/components/chat-message-list.test.ts @@ -401,6 +401,62 @@ describe("ChatMessageList", () => { expect(screen.queryByRole("img", { name: "其他候选图片" })).toBeNull(); }); + it("renders multi-region answer highlights on displayed page artifacts", () => { + render( + React.createElement(ChatMessageList, { + messages: [ + { + id: "assistant_1", + role: "assistant", + content: "风险辨识要求建立分级管控制度。", + artifacts: [ + { + type: "image", + display: true, + assetUrl: "https://blob.example/pages/page-225.png", + label: "page 225", + highlightRegions: [ + { x: 0.1, y: 0.2, w: 0.4, h: 0.1 }, + { x: 0.2, y: 0.5, w: 0.5, h: 0.12 }, + ], + }, + ], + }, + ], + }), + ); + + expect(screen.getByTestId("chat-image-highlights")).toBeTruthy(); + expect(screen.getAllByTestId("chat-image-highlight-region")).toHaveLength(2); + }); + + it("keeps the original image layout when artifacts have no highlight regions", () => { + render( + React.createElement(ChatMessageList, { + messages: [ + { + id: "assistant_1", + role: "assistant", + content: "Here is the page.", + artifacts: [ + { + type: "image", + display: true, + assetUrl: "https://blob.example/pages/page-1.png", + label: "page 1", + }, + ], + }, + ], + }), + ); + + const image = screen.getByRole("img", { name: "page 1" }); + expect(image.className).toContain("object-contain"); + expect(image.className).toContain("w-full"); + expect(screen.queryByTestId("chat-image-highlights")).toBeNull(); + }); + it("does not fall back to image citations when a harness message has empty artifacts", () => { render( React.createElement(ChatMessageList, { diff --git a/src/components/chat-message-list.tsx b/src/components/chat-message-list.tsx index 64ab794..6007f94 100644 --- a/src/components/chat-message-list.tsx +++ b/src/components/chat-message-list.tsx @@ -35,12 +35,14 @@ type DisplayImageCitation = { readonly label: string; readonly tooltipLabel: string; readonly assetUrl: string; + readonly highlightRegions?: ChatArtifactView["highlightRegions"]; }; type DisplayImageArtifact = { readonly assetUrl: string; readonly citationId: string; readonly label: string; + readonly highlightRegions?: ChatArtifactView["highlightRegions"]; }; type DisplayDerivedTableArtifact = { @@ -349,16 +351,16 @@ function MessageBubble({ Images

- {displayImageCitations.map(({ assetUrl, citationId, label }) => ( + {displayImageCitations.map( + ({ assetUrl, citationId, label, highlightRegions }) => (
- {/* eslint-disable-next-line @next/next/no-img-element -- Chat image citation dimensions are not known before render. */} - {label}
@@ -366,7 +368,8 @@ function MessageBubble({
- ))} + ), + )}
)} @@ -809,12 +812,68 @@ function getDisplayImageArtifacts( assetUrl, citationId: `${message.id}:artifact:${index}`, label: getArtifactLabel(artifact, sourceTitlesByDocumentId), + ...(artifact.highlightRegions && artifact.highlightRegions.length > 0 + ? { highlightRegions: artifact.highlightRegions } + : {}), }); } return imageArtifacts; } +function HighlightedChatImage({ + assetUrl, + label, + highlightRegions, +}: { + readonly assetUrl: string; + readonly label: string; + readonly highlightRegions?: ChatArtifactView["highlightRegions"]; +}): ReactElement { + const regions = highlightRegions ?? []; + + if (regions.length === 0) { + return ( + // eslint-disable-next-line @next/next/no-img-element -- Chat image citation dimensions are not known before render. + {label} + ); + } + + return ( +
+ {/* eslint-disable-next-line @next/next/no-img-element -- Chat image citation dimensions are not known before render. */} + {label} + +
+ ); +} + function getDisplayDerivedTableArtifacts( message: ChatMessageView, ): readonly DisplayDerivedTableArtifact[] { diff --git a/src/domains/chat/chat-citation-persistence.ts b/src/domains/chat/chat-citation-persistence.ts index 5b0b26e..3753a9a 100644 --- a/src/domains/chat/chat-citation-persistence.ts +++ b/src/domains/chat/chat-citation-persistence.ts @@ -51,6 +51,7 @@ function toArtifactView(artifact: ChatArtifactView): ChatArtifactView { label: artifact.label, display: artifact.display, reason: artifact.reason, + highlightRegions: artifact.highlightRegions, citation: artifact.citation ? toCitationView(artifact.citation) : undefined, diff --git a/src/domains/chat/image-inspection-model.test.ts b/src/domains/chat/image-inspection-model.test.ts new file mode 100644 index 0000000..5e123a7 --- /dev/null +++ b/src/domains/chat/image-inspection-model.test.ts @@ -0,0 +1,298 @@ +import { + generateObject, + generateText, + NoObjectGeneratedError, + UnsupportedFunctionalityError, +} from "ai" +import { beforeEach, describe, expect, it, vi } from "vitest" + +import { + generateImageInspectionModelResult, + isRecoverableStructuredOutputError, + parseStructuredInspectionText, + salvageAnalysisFromFailedStructuredText, +} from "./image-inspection-model" + +vi.mock("ai", async (importOriginal) => { + const original = await importOriginal() + return { + ...original, + generateObject: vi.fn(), + generateText: vi.fn(), + } +}) + +describe("image inspection model", () => { + const assets = [ + { + ref: "asset:page-1", + label: "page 1", + body: new Uint8Array([1, 2, 3]), + contentType: "image/png", + }, + ] as const + + beforeEach(() => { + vi.mocked(generateObject).mockReset() + vi.mocked(generateText).mockReset() + }) + + it("returns structured analysis and pages on success", async () => { + vi.mocked(generateObject).mockResolvedValue({ + object: { + analysis: "The page states approval is required.", + pages: [ + { + ref: "asset:page-1", + regions: [{ x: 0.1, y: 0.2, w: 0.3, h: 0.4 }], + }, + ], + }, + } as Awaited>) + + const result = await generateImageInspectionModelResult({ + workspaceId: "ws_1", + question: "When is approval required?", + assets, + }) + + expect(result).toEqual({ + analysis: "The page states approval is required.", + pages: [ + { + ref: "asset:page-1", + regions: [{ x: 0.1, y: 0.2, w: 0.3, h: 0.4 }], + }, + ], + source: "structured", + }) + expect(generateText).not.toHaveBeenCalled() + }) + + it("salvages valid structured JSON from failed generateObject text", async () => { + vi.mocked(generateObject).mockRejectedValue( + makeNoObjectGeneratedError( + JSON.stringify({ + analysis: "Visible caption says 5000 yuan.", + pages: [ + { + ref: "asset:page-1", + regions: [{ x: 0.2, y: 0.3, w: 0.4, h: 0.1 }], + }, + ], + }), + ), + ) + + const result = await generateImageInspectionModelResult({ + workspaceId: "ws_1", + question: "What amount is shown?", + assets, + }) + + expect(result).toEqual({ + analysis: "Visible caption says 5000 yuan.", + pages: [ + { + ref: "asset:page-1", + regions: [{ x: 0.2, y: 0.3, w: 0.4, h: 0.1 }], + }, + ], + source: "structured_text", + }) + expect(generateText).not.toHaveBeenCalled() + }) + + it("salvages analysis-only when failed structured text has analysis but invalid pages", async () => { + vi.mocked(generateObject).mockRejectedValue( + makeNoObjectGeneratedError( + '{"analysis":"Visible caption says 5000 yuan.","pages":[{"ref":"bad"}]}', + ), + ) + + const result = await generateImageInspectionModelResult({ + workspaceId: "ws_1", + question: "What amount is shown?", + assets, + }) + + expect(result).toEqual({ + analysis: "Visible caption says 5000 yuan.", + pages: [], + source: "analysis_fallback", + }) + expect(generateText).not.toHaveBeenCalled() + }) + + it("falls back to generateText when structured output fails without salvageable analysis", async () => { + vi.mocked(generateObject).mockRejectedValue( + makeNoObjectGeneratedError('{"pages":[]}'), + ) + vi.mocked(generateText) + .mockResolvedValueOnce({ + text: "not-json", + } as Awaited>) + .mockResolvedValueOnce({ + text: "The image shows a fee table.", + } as Awaited>) + + const result = await generateImageInspectionModelResult({ + workspaceId: "ws_1", + question: "What does the table show?", + assets, + }) + + expect(result).toEqual({ + analysis: "The image shows a fee table.", + pages: [], + source: "analysis_fallback", + }) + expect(generateText).toHaveBeenCalledTimes(2) + }) + + it("uses generateText JSON when the model does not support structured object generation", async () => { + vi.mocked(generateObject).mockRejectedValue( + new UnsupportedFunctionalityError({ + functionality: "object generation", + }), + ) + vi.mocked(generateText).mockResolvedValue({ + text: JSON.stringify({ + analysis: "Approval is required on City land.", + pages: [ + { + ref: "asset:page-1", + regions: [{ x: 0.05, y: 0.1, w: 0.9, h: 0.2 }], + }, + ], + }), + } as Awaited>) + + const result = await generateImageInspectionModelResult({ + workspaceId: "ws_1", + question: "When is approval required?", + assets, + }) + + expect(result).toEqual({ + analysis: "Approval is required on City land.", + pages: [ + { + ref: "asset:page-1", + regions: [{ x: 0.05, y: 0.1, w: 0.9, h: 0.2 }], + }, + ], + source: "structured_text", + }) + expect(generateText).toHaveBeenCalledTimes(1) + }) + + it("falls back to analysis-only text when unsupported structured models cannot emit JSON", async () => { + vi.mocked(generateObject).mockRejectedValue( + new UnsupportedFunctionalityError({ + functionality: "object generation", + }), + ) + vi.mocked(generateText) + .mockResolvedValueOnce({ + text: "not-json", + } as Awaited>) + .mockResolvedValueOnce({ + text: "The diagram labels a 40km/h zone.", + } as Awaited>) + + const result = await generateImageInspectionModelResult({ + workspaceId: "ws_1", + question: "What speed is shown?", + assets, + }) + + expect(result).toEqual({ + analysis: "The diagram labels a 40km/h zone.", + pages: [], + source: "analysis_fallback", + }) + expect(generateText).toHaveBeenCalledTimes(2) + }) + + it("rethrows non-schema model failures without falling back", async () => { + vi.mocked(generateObject).mockRejectedValue(new Error("network down")) + + await expect( + generateImageInspectionModelResult({ + workspaceId: "ws_1", + question: "What does the table show?", + assets, + }), + ).rejects.toThrow("network down") + expect(generateText).not.toHaveBeenCalled() + }) + + it("classifies recoverable structured-output errors", () => { + expect( + isRecoverableStructuredOutputError( + new UnsupportedFunctionalityError({ + functionality: "object generation", + }), + ), + ).toBe(true) + expect( + isRecoverableStructuredOutputError(makeNoObjectGeneratedError("{}")), + ).toBe(true) + expect(isRecoverableStructuredOutputError(new Error("network down"))).toBe( + false, + ) + }) + + it("parses and salvages structured inspection text", () => { + expect( + parseStructuredInspectionText( + JSON.stringify({ + analysis: "Box A is labeled red.", + pages: [{ ref: "asset:page-1", regions: [{ x: 0, y: 0, w: 1, h: 1 }] }], + }), + ), + ).toEqual({ + analysis: "Box A is labeled red.", + pages: [{ ref: "asset:page-1", regions: [{ x: 0, y: 0, w: 1, h: 1 }] }], + }) + expect( + salvageAnalysisFromFailedStructuredText( + 'prefix {"analysis":"Box A is labeled red."} suffix', + ), + ).toBe("Box A is labeled red.") + expect(salvageAnalysisFromFailedStructuredText("Plain OCR notes.")).toBe( + "Plain OCR notes.", + ) + expect(salvageAnalysisFromFailedStructuredText('{"pages":[]}')).toBeNull() + expect(salvageAnalysisFromFailedStructuredText("")).toBeNull() + }) +}) + +function makeNoObjectGeneratedError(text: string): NoObjectGeneratedError { + return new NoObjectGeneratedError({ + message: "No object generated: response did not match schema.", + cause: new Error("schema mismatch"), + text, + response: { + id: "response_1", + modelId: "test-model", + timestamp: new Date("2026-01-01T00:00:00Z"), + }, + usage: { + inputTokens: 1, + inputTokenDetails: { + noCacheTokens: 1, + cacheReadTokens: 0, + cacheWriteTokens: 0, + }, + outputTokens: 1, + outputTokenDetails: { + textTokens: 1, + reasoningTokens: 0, + }, + totalTokens: 2, + }, + finishReason: "stop", + }) +} diff --git a/src/domains/chat/image-inspection-model.ts b/src/domains/chat/image-inspection-model.ts new file mode 100644 index 0000000..e89804d --- /dev/null +++ b/src/domains/chat/image-inspection-model.ts @@ -0,0 +1,392 @@ +import { + APICallError, + generateObject, + generateText, + NoObjectGeneratedError, + UnsupportedFunctionalityError, +} from "ai" +import { z } from "zod" + +import { CHAT_MODEL } from "@/lib/ai" +import { summarizeUnknownError } from "@/lib/format-log-value" +import { logger } from "@/lib/logger" + +const VISION_MODEL = process.env.VISION_MODEL ?? CHAT_MODEL + +export const imageInspectionResultSchema = z.object({ + analysis: z.string(), + pages: z + .array( + z.object({ + ref: z.string().min(1), + regions: z + .array( + z.object({ + x: z.number(), + y: z.number(), + w: z.number(), + h: z.number(), + }), + ) + .max(12), + }), + ) + .max(6) + .default([]), +}) + +export type ImageInspectionModelPage = z.infer< + typeof imageInspectionResultSchema +>["pages"][number] + +export type ImageInspectionModelAsset = { + readonly ref: string + readonly label: string + readonly body: Uint8Array + readonly contentType: string +} + +export type ImageInspectionModelResult = { + readonly analysis: string + readonly pages: readonly ImageInspectionModelPage[] + readonly source: "structured" | "structured_text" | "analysis_fallback" +} + +/** + * Prefer native structured output (analysis + highlight boxes). + * If the model/provider cannot do generateObject, degrade through: + * 1) salvage/parse JSON from failed text + * 2) generateText + JSON prompt (may still yield boxes) + * 3) analysis-only generateText + * so inspectImage still succeeds. + */ +export async function generateImageInspectionModelResult(input: { + readonly workspaceId: string + readonly question: string + readonly assets: readonly ImageInspectionModelAsset[] +}): Promise { + let structuredFailureText: string | undefined + + try { + const response = await generateObject({ + model: VISION_MODEL, + schema: imageInspectionResultSchema, + messages: [ + { + role: "user", + content: [ + { + type: "text", + text: buildImageInspectionStructuredPrompt({ + question: input.question, + assets: input.assets, + }), + }, + ...toImageParts(input.assets), + ], + }, + ], + experimental_include: { + requestBody: false, + responseBody: false, + }, + }) + + return { + analysis: response.object.analysis.trim(), + pages: response.object.pages, + source: "structured", + } + } catch (error) { + if (!isRecoverableStructuredOutputError(error)) { + logger.warn("chat: image inspection model call failed", { + workspaceId: input.workspaceId, + model: VISION_MODEL, + inspectedCount: input.assets.length, + error: summarizeUnknownError(error), + }) + throw error + } + + structuredFailureText = getStructuredFailureText(error) + logger.warn("chat: image inspection structured output unavailable; falling back", { + workspaceId: input.workspaceId, + model: VISION_MODEL, + inspectedCount: input.assets.length, + generatedTextLength: structuredFailureText?.length ?? 0, + error: summarizeUnknownError(error), + }) + } + + const salvagedStructured = parseStructuredInspectionText(structuredFailureText) + if (salvagedStructured) { + return { + ...salvagedStructured, + source: "structured_text", + } + } + + const salvagedAnalysis = salvageAnalysisFromFailedStructuredText( + structuredFailureText, + ) + if (salvagedAnalysis !== null) { + return { + analysis: salvagedAnalysis, + pages: [], + source: "analysis_fallback", + } + } + + try { + const jsonTextResponse = await generateText({ + model: VISION_MODEL, + messages: [ + { + role: "user", + content: [ + { + type: "text", + text: buildImageInspectionStructuredPrompt({ + question: input.question, + assets: input.assets, + }), + }, + ...toImageParts(input.assets), + ], + }, + ], + experimental_include: { + requestBody: false, + responseBody: false, + }, + }) + + const parsedFromText = parseStructuredInspectionText(jsonTextResponse.text) + if (parsedFromText) { + return { + ...parsedFromText, + source: "structured_text", + } + } + + // Structured prompt asked for JSON; only keep an analysis field from JSON, + // never treat arbitrary freeform text as a successful inspection here. + const analysisFromJsonAttempt = salvageAnalysisFieldFromJsonText( + jsonTextResponse.text, + ) + if (analysisFromJsonAttempt !== null) { + return { + analysis: analysisFromJsonAttempt, + pages: [], + source: "analysis_fallback", + } + } + } catch (error) { + logger.warn("chat: image inspection structured-text fallback failed", { + workspaceId: input.workspaceId, + model: VISION_MODEL, + inspectedCount: input.assets.length, + error: summarizeUnknownError(error), + }) + } + + try { + const response = await generateText({ + model: VISION_MODEL, + messages: [ + { + role: "user", + content: [ + { + type: "text", + text: buildImageInspectionAnalysisPrompt({ + question: input.question, + assets: input.assets, + }), + }, + ...toImageParts(input.assets), + ], + }, + ], + experimental_include: { + requestBody: false, + responseBody: false, + }, + }) + + return { + analysis: response.text.trim(), + pages: [], + source: "analysis_fallback", + } + } catch (error) { + logger.warn("chat: image inspection analysis fallback failed", { + workspaceId: input.workspaceId, + model: VISION_MODEL, + inspectedCount: input.assets.length, + error: summarizeUnknownError(error), + }) + throw error + } +} + +export function isRecoverableStructuredOutputError(error: unknown): boolean { + if (NoObjectGeneratedError.isInstance(error)) return true + if (UnsupportedFunctionalityError.isInstance(error)) return true + + if (APICallError.isInstance(error)) { + const haystack = [ + error.message, + error.data === undefined ? "" : JSON.stringify(error.data), + typeof error.responseBody === "string" ? error.responseBody : "", + ] + .join(" ") + .toLowerCase() + + return ( + /response[_\s-]?format/.test(haystack) || + /json[_\s-]?schema/.test(haystack) || + /structured[_\s-]?output/.test(haystack) || + /tool[_\s-]?choice/.test(haystack) || + /unsupported/.test(haystack) + ) + } + + return false +} + +export function buildImageInspectionStructuredPrompt(input: { + readonly question: string + readonly assets: readonly Pick[] +}): string { + return [ + "Inspect the attached Notebook image assets selected from retrieved Knowhere evidence.", + "Answer the inspection question using concise visual observations only.", + "Use the provided refs to identify images. Do not include image URLs.", + "If OCR text is unclear, say it is unclear instead of guessing.", + "Do not create citations. The calling agent will cite the original retrieved asset refs.", + "", + "Also return answer provenance boxes for every page that supports the answer.", + "One answer may span multiple pages; each page may have one or more regions.", + "Coordinate system: origin top-left; x,y,w,h are fractions of image width/height in [0,1].", + "Do not add labels or captions for regions. Omit pages with no useful region.", + "", + "Return ONLY valid JSON with this shape:", + '{"analysis":"string","pages":[{"ref":"asset:page-12","regions":[{"x":0,"y":0,"w":0.2,"h":0.1}]}]}', + "", + "Inspection question:", + input.question, + "", + "Images:", + ...input.assets.map((asset) => `- ref=${asset.ref} label=${asset.label}`), + ].join("\n") +} + +export function buildImageInspectionAnalysisPrompt(input: { + readonly question: string + readonly assets: readonly Pick[] +}): string { + return [ + "Inspect the attached Notebook image assets selected from retrieved Knowhere evidence.", + "Answer the inspection question using concise visual observations only.", + "Use the provided refs and labels to identify images. Do not include image URLs.", + "If OCR text is unclear, say it is unclear instead of guessing.", + "Do not create citations. The calling agent will cite the original retrieved asset refs.", + "", + "Inspection question:", + input.question, + "", + "Images:", + ...input.assets.map((asset) => `- ref=${asset.ref} label=${asset.label}`), + ].join("\n") +} + +export function parseStructuredInspectionText( + text: string | undefined, +): { readonly analysis: string; readonly pages: ImageInspectionModelPage[] } | null { + const trimmed = text?.trim() + if (!trimmed) return null + + const start = trimmed.indexOf("{") + const end = trimmed.lastIndexOf("}") + if (start < 0 || end <= start) return null + + try { + const parsed = JSON.parse(trimmed.slice(start, end + 1)) as unknown + const result = imageInspectionResultSchema.safeParse(parsed) + if (!result.success) return null + + const analysis = result.data.analysis.trim() + if (!analysis) return null + + return { + analysis, + pages: result.data.pages, + } + } catch { + return null + } +} + +export function salvageAnalysisFieldFromJsonText( + text: string | undefined, +): string | null { + const trimmed = text?.trim() + if (!trimmed) return null + + try { + const start = trimmed.indexOf("{") + const end = trimmed.lastIndexOf("}") + if (start < 0 || end <= start) return null + + const parsed = JSON.parse(trimmed.slice(start, end + 1)) as unknown + if ( + !parsed || + typeof parsed !== "object" || + !("analysis" in parsed) || + typeof (parsed as { analysis: unknown }).analysis !== "string" + ) { + return null + } + + const analysis = (parsed as { analysis: string }).analysis.trim() + return analysis.length > 0 ? analysis : null + } catch { + return null + } +} + +export function salvageAnalysisFromFailedStructuredText( + text: string | undefined, +): string | null { + const structured = parseStructuredInspectionText(text) + if (structured) return structured.analysis + + const fromJson = salvageAnalysisFieldFromJsonText(text) + if (fromJson !== null) return fromJson + + const trimmed = text?.trim() + if (!trimmed) return null + + // Reject obvious raw JSON blobs without a usable analysis field. + if (trimmed.startsWith("{") || trimmed.startsWith("[")) { + return null + } + + return trimmed +} + +function getStructuredFailureText(error: unknown): string | undefined { + if (NoObjectGeneratedError.isInstance(error)) { + return error.text + } + return undefined +} + +function toImageParts(assets: readonly ImageInspectionModelAsset[]) { + return assets.map((asset) => ({ + type: "image" as const, + image: asset.body, + mediaType: asset.contentType, + })) +} diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index e26a670..1138fe2 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -2949,6 +2949,7 @@ function makeHarnessRunResult(text: string): HarnessRunResult { finalized: true, priorTurnReads: [], toolCalls: [], + imageHighlights: [], validationErrors: [], revisionsUsed: 0, }, diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index a6990c4..c4f7f85 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -9,6 +9,7 @@ import { logger } from "@/lib/logger" import type { ChatArtifactView, ChatCitationView, + ChatImageHighlightBox, } from "@/domains/chat/types" import type { DerivedTableArtifact, @@ -320,6 +321,11 @@ function toChatArtifactViewsFromHarness( chunk, ]), ) + const highlightsByRef = new Map( + (result.trace.imageHighlights ?? []).map( + (page) => [page.ref, page.regions] as const, + ), + ) const displayLimit = getHarnessArtifactDisplayLimit(result) const artifacts: ChatArtifactView[] = [] @@ -333,6 +339,7 @@ function toChatArtifactViewsFromHarness( artifact, assetsByRef, chunksByRef, + highlightsByRef, sources, }) if (!artifactView) continue @@ -381,6 +388,7 @@ function resolveHarnessArtifactView(input: { readonly artifact: OutputArtifact readonly assetsByRef: ReadonlyMap readonly chunksByRef: ReadonlyMap + readonly highlightsByRef: ReadonlyMap readonly sources: readonly AnswerQuestionInput["sources"][number][] }): ChatArtifactView | null { const asset = input.assetsByRef.get(input.artifact.ref) @@ -388,6 +396,7 @@ function resolveHarnessArtifactView(input: { return toChatArtifactView({ artifact: input.artifact, asset, + highlightRegions: input.highlightsByRef.get(asset.ref), sources: input.sources, }) } @@ -399,6 +408,9 @@ function resolveHarnessArtifactView(input: { ? toChatArtifactView({ artifact: input.artifact, asset: chunkAsset, + highlightRegions: + input.highlightsByRef.get(chunkAsset.ref) ?? + input.highlightsByRef.get(input.artifact.ref), sources: input.sources, }) : null @@ -407,6 +419,7 @@ function resolveHarnessArtifactView(input: { function toChatArtifactView(input: { readonly artifact: OutputArtifact readonly asset: EvidenceAsset + readonly highlightRegions?: readonly ChatImageHighlightBox[] readonly sources: readonly AnswerQuestionInput["sources"][number][] }): ChatArtifactView { const source = normalizeHarnessSource(input.asset.source, input.sources) @@ -417,6 +430,9 @@ function toChatArtifactView(input: { reason: input.artifact.reason, ...(input.asset.assetUrl ? { assetUrl: input.asset.assetUrl } : {}), label: input.asset.label, + ...(input.highlightRegions && input.highlightRegions.length > 0 + ? { highlightRegions: input.highlightRegions } + : {}), citation: { chunkType: input.asset.type, score: null, diff --git a/src/domains/chat/route-answer.ts b/src/domains/chat/route-answer.ts index ae8b13d..eaa8064 100644 --- a/src/domains/chat/route-answer.ts +++ b/src/domains/chat/route-answer.ts @@ -1,5 +1,4 @@ import { Cause, Effect, Either, Option } from "effect" -import { generateText } from "ai" import { generateAgenticOutputManifest, @@ -12,6 +11,8 @@ import type { ImageInspectionSkippedAsset, InspectImages, } from "@/agent-harness" +import { normalizeImageInspectionHighlights } from "@/agent-harness/image-highlights" +import { generateImageInspectionModelResult } from "@/domains/chat/image-inspection-model" import { hardenChatMediaAssetUrls } from "@/domains/chat/media-asset-hardening" import { handleChatTurn, @@ -256,15 +257,19 @@ async function inspectChatImages(input: { refs: preparedAssets.map((asset) => asset.ref), }) - const response = await generateImageInspectionText({ + const response = await generateImageInspectionModelResult({ workspaceId: input.workspaceId, question: input.request.question, assets: preparedAssets, }) const analysis = removeImageInspectionUrls({ - text: response.text.trim(), + text: response.analysis, assets: preparedAssets, }) + const highlights = normalizeImageInspectionHighlights({ + pages: response.pages, + allowedRefs: new Set(preparedAssets.map((asset) => asset.ref)), + }) logger.info("chat: image inspection response", { workspaceId: input.workspaceId, @@ -272,6 +277,12 @@ async function inspectChatImages(input: { inspectedCount: preparedAssets.length, skippedCount: skippedAssets.length, analysisLength: analysis.length, + inspectionSource: response.source, + highlightPageCount: highlights.length, + highlightRegionCount: highlights.reduce( + (count, page) => count + page.regions.length, + 0, + ), }) return { @@ -281,6 +292,7 @@ async function inspectChatImages(input: { label: asset.label, })), skipped: skippedAssets, + ...(highlights.length > 0 ? { highlights } : {}), } } @@ -382,49 +394,6 @@ async function prepareImageInspectionAsset(input: { } } -async function generateImageInspectionText(input: { - readonly workspaceId: string - readonly question: string - readonly assets: readonly PreparedImageInspectionAsset[] -}): Promise>> { - try { - return await generateText({ - model: VISION_MODEL, - messages: [ - { - role: "user", - content: [ - { - type: "text", - text: buildImageInspectionPrompt({ - question: input.question, - assets: input.assets, - }), - }, - ...input.assets.map((asset) => ({ - type: "image" as const, - image: asset.body, - mediaType: asset.contentType, - })), - ], - }, - ], - experimental_include: { - requestBody: false, - responseBody: false, - }, - }) - } catch (error) { - logger.warn("chat: image inspection model call failed", { - workspaceId: input.workspaceId, - model: VISION_MODEL, - inspectedCount: input.assets.length, - error: summarizeUnknownError(error), - }) - throw error - } -} - async function fetchPreparedInspectionImage(input: { readonly url: URL readonly contentType: string @@ -525,25 +494,6 @@ function getSupportedImageContentType(sourcePath: string): string | null { return extension ? SUPPORTED_IMAGE_CONTENT_TYPES[extension] ?? null : null } -function buildImageInspectionPrompt(input: { - readonly question: string - readonly assets: readonly PreparedImageInspectionAsset[] -}): string { - return [ - "Inspect the attached Notebook image assets selected from retrieved Knowhere evidence.", - "Answer the inspection question using concise visual observations only.", - "Use the provided refs and labels to identify images. Do not include image URLs.", - "If OCR text is unclear, say it is unclear instead of guessing.", - "Do not create citations. The calling agent will cite the original retrieved asset refs.", - "", - "Inspection question:", - input.question, - "", - "Images:", - ...input.assets.map((asset) => `- ref=${asset.ref} label=${asset.label}`), - ].join("\n") -} - function removeImageInspectionUrls(input: { readonly text: string readonly assets: readonly PreparedImageInspectionAsset[] diff --git a/src/domains/chat/route-service.test.ts b/src/domains/chat/route-service.test.ts index 97c9f6d..0772ca0 100644 --- a/src/domains/chat/route-service.test.ts +++ b/src/domains/chat/route-service.test.ts @@ -9,6 +9,7 @@ const mocks = vi.hoisted(() => ({ ensureDefaultChatThread: vi.fn(), findChatThreadInWorkspace: vi.fn(), generateAgenticOutputManifest: vi.fn(), + generateObject: vi.fn(), generateText: vi.fn(), getAuthenticated: vi.fn(), getAuthenticatedWithClient: vi.fn(), @@ -19,6 +20,7 @@ const mocks = vi.hoisted(() => ({ loggerInfo: vi.fn(), loggerWarn: vi.fn(), listSourcesForWorkspace: vi.fn(), + makeKnowhereClientWithParsedStorage: vi.fn(), parsedStorageGetAssetUrl: vi.fn(), parsedStorageWriteAsset: vi.fn(), softDeleteChatThread: vi.fn(), @@ -29,10 +31,15 @@ vi.mock("ai", async (importOriginal) => { const original = await importOriginal() return { ...original, + generateObject: mocks.generateObject, generateText: mocks.generateText, } }) +vi.mock("@/integrations/knowhere", () => ({ + makeKnowhereClientWithParsedStorage: mocks.makeKnowhereClientWithParsedStorage, +})) + vi.mock("@/domains/chat", async (importOriginal) => { const original = await importOriginal() return { @@ -99,6 +106,16 @@ describe("chat route services", () => { vi.clearAllMocks() mocks.parsedStorageGetAssetUrl.mockResolvedValue(null) mocks.parsedStorageWriteAsset.mockResolvedValue({ url: null }) + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: {}, + knowledge: {}, + }) + mocks.generateObject.mockResolvedValue({ + object: { + analysis: "The image shows a chart.", + pages: [], + }, + }) mocks.generateText.mockResolvedValue({ text: "The image shows a chart." }) vi.stubGlobal( "fetch", @@ -253,8 +270,11 @@ describe("chat route services", () => { const durableUrl = "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/job_1/images/id-front.png" mocks.parsedStorageWriteAsset.mockResolvedValue({ url: durableUrl }) - mocks.generateText.mockResolvedValue({ - text: `The card number is visible. ${durableUrl} ${rawUrl}`, + mocks.generateObject.mockResolvedValue({ + object: { + analysis: `The card number is visible. ${durableUrl} ${rawUrl}`, + pages: [], + }, }) mocks.getAuthenticatedWithClient.mockResolvedValue({ user: { id: "user_1" }, @@ -338,7 +358,7 @@ describe("chat route services", () => { body: new Uint8Array([1, 2, 3]), contentType: "image/png", }) - const generateInput = mocks.generateText.mock.calls[0]?.[0] + const generateInput = mocks.generateObject.mock.calls[0]?.[0] expect(generateInput).toMatchObject({ model: "google/gemini-3-flash", experimental_include: { @@ -368,8 +388,11 @@ describe("chat route services", () => { const durableUrl = "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_contract/job_1/page_citation_assets/page-8.png" mocks.parsedStorageGetAssetUrl.mockResolvedValue(durableUrl) - mocks.generateText.mockResolvedValue({ - text: "The page states 5000 yuan per occurrence.", + mocks.generateObject.mockResolvedValue({ + object: { + analysis: "The page states 5000 yuan per occurrence.", + pages: [], + }, }) mocks.getAuthenticatedWithClient.mockResolvedValue({ user: { id: "user_1" }, @@ -454,7 +477,7 @@ describe("chat route services", () => { }) expect(fetch).toHaveBeenCalledWith(durableUrl) expect(mocks.parsedStorageWriteAsset).not.toHaveBeenCalled() - const generateInput = mocks.generateText.mock.calls[0]?.[0] + const generateInput = mocks.generateObject.mock.calls[0]?.[0] const content = generateInput.messages[0].content const imagePart = content.find( (part: { readonly type: string }) => part.type === "image", @@ -471,8 +494,11 @@ describe("chat route services", () => { const durableUrl = "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_remote/job_remote/page_citation_assets/page-8.png" mocks.parsedStorageWriteAsset.mockResolvedValue({ url: durableUrl }) - mocks.generateText.mockResolvedValue({ - text: "The page states 5000 yuan per occurrence.", + mocks.generateObject.mockResolvedValue({ + object: { + analysis: "The page states 5000 yuan per occurrence.", + pages: [], + }, }) mocks.getAuthenticatedWithClient.mockResolvedValue({ user: { id: "user_1" }, @@ -567,7 +593,7 @@ describe("chat route services", () => { body: new Uint8Array([1, 2, 3]), contentType: "image/png", }) - const generateInput = mocks.generateText.mock.calls[0]?.[0] + const generateInput = mocks.generateObject.mock.calls[0]?.[0] expect(JSON.stringify(generateInput)).not.toContain( "knowhere-storage.example", ) @@ -674,6 +700,7 @@ describe("chat route services", () => { }) expect(result.status).toBe(200) + expect(mocks.generateObject).not.toHaveBeenCalled() expect(mocks.generateText).not.toHaveBeenCalled() expect(fetch).not.toHaveBeenCalled() }) diff --git a/src/domains/chat/service.test.ts b/src/domains/chat/service.test.ts index 6001171..b065251 100644 --- a/src/domains/chat/service.test.ts +++ b/src/domains/chat/service.test.ts @@ -350,6 +350,7 @@ function makeHarnessRunResult(text: string): HarnessRunResult { finalized: true, priorTurnReads: [], toolCalls: [], + imageHighlights: [], validationErrors: [], revisionsUsed: 0, }, diff --git a/src/domains/chat/types.ts b/src/domains/chat/types.ts index e430e4c..e5179b0 100644 --- a/src/domains/chat/types.ts +++ b/src/domains/chat/types.ts @@ -31,6 +31,17 @@ export type ChatCitationView = CitationView & { readonly content?: string } +/** + * Normalized highlight box for page/image answer provenance. + * Origin top-left; values in [0, 1]. No per-region labels. + */ +export type ChatImageHighlightBox = { + readonly x: number + readonly y: number + readonly w: number + readonly h: number +} + export type ChatArtifactView = { readonly type: "image" | "table" | "derived_table" readonly ref?: string @@ -42,6 +53,7 @@ export type ChatArtifactView = { readonly label?: string readonly display?: boolean readonly reason?: string + readonly highlightRegions?: readonly ChatImageHighlightBox[] readonly citation?: ChatCitationView } diff --git a/src/domains/chat/view.ts b/src/domains/chat/view.ts index f5c7f3a..f74c287 100644 --- a/src/domains/chat/view.ts +++ b/src/domains/chat/view.ts @@ -1,3 +1,4 @@ +import { normalizeHighlightBoxes } from "@/agent-harness/image-highlights" import { deriveChatThreadTitle } from "./title" import type { ChatMessage, ChatThread } from "@/infrastructure/db/schema" import type { @@ -137,6 +138,7 @@ function toPersistedArtifactViews(value: unknown): ChatArtifactView[] | undefine label: getString(item.label), display: typeof item.display === "boolean" ? item.display : undefined, reason: getString(item.reason), + highlightRegions: getHighlightRegions(item.highlightRegions), ...(citation ? { citation } : {}), }, ] @@ -145,6 +147,37 @@ function toPersistedArtifactViews(value: unknown): ChatArtifactView[] | undefine return artifacts.length > 0 ? artifacts : undefined } +function getHighlightRegions( + value: unknown, +): ChatArtifactView["highlightRegions"] { + if (!Array.isArray(value) || value.length === 0) return undefined + + const candidates = value.flatMap((item): Array<{ + x: number + y: number + w: number + h: number + }> => { + if (!isRecord(item)) return [] + const x = getNumber(item.x) + const y = getNumber(item.y) + const w = getNumber(item.w) + const h = getNumber(item.h) + if ( + x === undefined || + y === undefined || + w === undefined || + h === undefined + ) { + return [] + } + return [{ x, y, w, h }] + }) + + const regions = normalizeHighlightBoxes(candidates) + return regions.length > 0 ? regions : undefined +} + function getString(value: unknown): string | undefined { if (typeof value !== "string") return undefined return value.length > 0 ? value : undefined From c33739d6d1eef696d0348ee316b5ca48e1abd869 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 21 Aug 2026 12:28:50 +0800 Subject: [PATCH 48/51] fix: resolve page-highlight typecheck failures for CI Import ImageInspectionHighlights, drop unsupported generateObject experimental_include, and complete HarnessTrace fixtures with imageHighlights. Co-authored-by: Cursor --- src/agent-harness/runtime.ts | 1 + src/domains/chat/image-inspection-model.ts | 4 ---- src/domains/chat/index.test.ts | 1 + src/domains/chat/route-service.test.ts | 4 ---- 4 files changed, 2 insertions(+), 8 deletions(-) diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index 8dbe18d..3930146 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -19,6 +19,7 @@ import type { HarnessToolCallTrace, HarnessTrace, ImageInspectionAsset, + ImageInspectionHighlights, ImageInspectionResponse, InspectImages, IntentFrame, diff --git a/src/domains/chat/image-inspection-model.ts b/src/domains/chat/image-inspection-model.ts index e89804d..8589dda 100644 --- a/src/domains/chat/image-inspection-model.ts +++ b/src/domains/chat/image-inspection-model.ts @@ -86,10 +86,6 @@ export async function generateImageInspectionModelResult(input: { ], }, ], - experimental_include: { - requestBody: false, - responseBody: false, - }, }) return { diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 1138fe2..114ade8 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -1545,6 +1545,7 @@ describe("answerQuestionWithRetrieval", () => { finalized: true, priorTurnReads: [], toolCalls: [], + imageHighlights: [], }, }; return harnessResult; diff --git a/src/domains/chat/route-service.test.ts b/src/domains/chat/route-service.test.ts index 0772ca0..4eab67d 100644 --- a/src/domains/chat/route-service.test.ts +++ b/src/domains/chat/route-service.test.ts @@ -361,10 +361,6 @@ describe("chat route services", () => { const generateInput = mocks.generateObject.mock.calls[0]?.[0] expect(generateInput).toMatchObject({ model: "google/gemini-3-flash", - experimental_include: { - requestBody: false, - responseBody: false, - }, }) expect(JSON.stringify(generateInput)).not.toContain( "knowhere-storage.example", From f29dcbdac36b08d5ca4f37623655cd74556bb0e0 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 21 Aug 2026 21:16:06 +0800 Subject: [PATCH 49/51] fix: require inspected citation regions Force cited page images through visual inspection, preserve the resulting boxes through persistence and navigation, and avoid misleading full-page highlights. Co-authored-by: Cursor --- e2e/citation-dedupe.e2e.ts | 91 ++++++ src/agent-harness/knowhere-text.test.ts | 1 + src/agent-harness/knowhere-text.ts | 2 +- src/agent-harness/ledger.test.ts | 46 +++ src/agent-harness/ledger.ts | 40 ++- src/agent-harness/runtime.test.ts | 289 +++++++++++++++++- src/agent-harness/runtime.ts | 229 ++++++++++---- src/app/e2e/citation-same-page/page.tsx | 5 + src/app/globals.css | 31 ++ src/components/chat-citation-model.test.ts | 84 +++++ src/components/chat-citation-model.ts | 51 +++- src/components/chat-message-list.test.ts | 44 +++ src/components/chat-message-list.tsx | 18 +- src/components/chat-panel.test.ts | 7 +- src/components/chat-panel.tsx | 7 +- src/components/chunks-panel.tsx | 11 + src/components/citation-region-highlight.tsx | 38 +++ src/components/parsed-chunk-card.test.ts | 34 +++ src/components/parsed-chunk-card.tsx | 56 +++- .../workspace-citation-focus.test.ts | 12 + src/components/workspace-citation-focus.ts | 24 +- src/components/workspace-shell-layout.tsx | 16 +- src/domains/chat/chat-citation-persistence.ts | 5 + src/domains/chat/citations.test.ts | 21 ++ src/domains/chat/citations.ts | 3 + src/domains/chat/index.test.ts | 210 +++++++++++++ src/domains/chat/index.ts | 86 +++++- src/domains/chat/page-citation-assets.ts | 6 + src/domains/chat/route-answer.ts | 8 + src/domains/chat/types.ts | 23 +- src/domains/chat/view.ts | 4 + 31 files changed, 1397 insertions(+), 105 deletions(-) create mode 100644 src/components/citation-region-highlight.tsx diff --git a/e2e/citation-dedupe.e2e.ts b/e2e/citation-dedupe.e2e.ts index 071a986..a21d7d0 100644 --- a/e2e/citation-dedupe.e2e.ts +++ b/e2e/citation-dedupe.e2e.ts @@ -108,3 +108,94 @@ test("keeps two chips to the same title/pN as separate buttons", async ({ chatPanel.getByRole("button", { name: "Open source spacex-s1.pdf/p26" }), ).toHaveCount(2) }) + +test("draws citation regions instead of a full-page highlight", async ({ + context, + page, +}) => { + await context.addCookies([ + { + name: "better-auth.session_token", + value: "playwright", + url: "http://localhost:3000", + }, + ]) + await page.setViewportSize({ width: 1280, height: 832 }) + await page.route("**/api/sources/source_spacex/chunks**", async (route) => { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + chunks: [ + { + chunkId: "page_26", + documentId: "doc_spacex", + sectionPath: "Page 26", + type: "page", + content: "Revenue evidence on page 26.", + readableContent: "Revenue evidence on page 26.", + pageNums: [26], + pageAssets: [ + { + pageNumber: 26, + assetUrl: "/images/knowhere/logo-icon.png", + contentType: "image/png", + width: 1000, + height: 1400, + }, + ], + sourceTitle: "spacex-s1.pdf", + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 1, + totalPages: 1, + }, + }), + }) + }) + + await page.goto("/e2e/citation-same-page") + const chips = page + .getByTestId("desktop-chat-panel") + .getByTestId("citation-chip") + await expect(chips).toHaveCount(2) + + await chips.first().click() + const firstRegions = page.getByTestId("citation-region-highlight") + await expect(firstRegions).toHaveCount(1) + await expect + .poll(() => + firstRegions.first().evaluate((element) => ({ + left: element.style.left, + top: element.style.top, + width: element.style.width, + height: element.style.height, + })), + ) + .toEqual({ left: "12%", top: "18%", width: "46%", height: "8%" }) + + await chips.nth(1).click() + const secondRegions = page.getByTestId("citation-region-highlight") + await expect(secondRegions).toHaveCount(2) + const regionStyles = await secondRegions.evaluateAll((elements) => + elements.map((element) => ({ + left: (element as HTMLElement).style.left, + top: (element as HTMLElement).style.top, + width: (element as HTMLElement).style.width, + height: (element as HTMLElement).style.height, + })), + ) + expect(regionStyles).toEqual([ + { left: "62%", top: "52%", width: "24%", height: "6%" }, + { left: "15%", top: "68%", width: "32%", height: "5%" }, + ]) + expect(regionStyles).not.toContainEqual({ + left: "0%", + top: "0%", + width: "100%", + height: "100%", + }) +}) diff --git a/src/agent-harness/knowhere-text.test.ts b/src/agent-harness/knowhere-text.test.ts index 88159f5..e916cad 100644 --- a/src/agent-harness/knowhere-text.test.ts +++ b/src/agent-harness/knowhere-text.test.ts @@ -26,6 +26,7 @@ describe("knowhereToolText", () => { expect(text).toContain('ref="r1:result:1"') expect(text).toContain('ref="asset:r1:result:1"') expect(text).toContain("Call inspectImage") + expect(text).toContain("before finalize") expect(text).not.toContain("https://assets.example/page-1.png") }) diff --git a/src/agent-harness/knowhere-text.ts b/src/agent-harness/knowhere-text.ts index 1c19368..f5b19da 100644 --- a/src/agent-harness/knowhere-text.ts +++ b/src/agent-harness/knowhere-text.ts @@ -42,7 +42,7 @@ type KnowhereOperation = | "grep_chunks" const assetInstruction = - "Notebook returned image/page asset refs. Call inspectImage with the asset refs when OCR, visual details, or verification are needed. Do not expose raw asset URLs." + "Notebook returned image/page asset refs. Call inspectImage with the asset refs you will cite before finalize so OCR/visual context and provenance boxes exist. Do not expose raw asset URLs." export const knowhereToolText = { formatSearch(input: SearchTextInput): string { diff --git a/src/agent-harness/ledger.test.ts b/src/agent-harness/ledger.test.ts index 8a77142..90cada3 100644 --- a/src/agent-harness/ledger.test.ts +++ b/src/agent-harness/ledger.test.ts @@ -68,6 +68,52 @@ describe("createEvidenceLedger", () => { ) }) + it("creates page image assets from snake-case retrieval metadata", () => { + const ledger = createEvidenceLedger() + + const snapshot = ledger.addRetrievalResponse({ + namespace: "notebook", + query: "revenue", + routerUsed: "mapnav", + answerText: null, + evidenceText: "Revenue evidence", + stopReason: "completed", + failureReason: null, + results: [ + { + chunkId: "chunk_page_4", + content: "Revenue was $24.9B.", + chunkType: "page", + score: 0.9, + metadata: { + page_nums: [4], + page_assets: [ + { + page_num: 4, + artifact_ref: "page_citation_assets/page-4.png", + content_type: "image/png", + }, + ], + }, + source: { + documentId: "doc_tsla", + sourceFileName: "TSLA-Q4-2025-Update.pdf", + sectionPath: "FINANCIAL SUMMARY", + }, + }, + ], + referencedChunks: [], + }) + + expect(snapshot.assets).toContainEqual( + expect.objectContaining({ + ref: "asset:r1:result:1", + sourcePath: "page_citation_assets/page-4.png", + type: "image", + }), + ) + }) + it("adds read chunk refs and page image assets", () => { const ledger = createEvidenceLedger() diff --git a/src/agent-harness/ledger.ts b/src/agent-harness/ledger.ts index 4e13f81..64700e1 100644 --- a/src/agent-harness/ledger.ts +++ b/src/agent-harness/ledger.ts @@ -381,8 +381,9 @@ function getPageCitationAssetCandidate( ): EvidenceAssetCandidate | null { if (chunk.chunkType.toLowerCase() !== "page") return null - const candidates = parsePageCitationAssetCandidates(chunk.metadata?.pageAssets) - .filter(isSupportedPageCitationAsset) + const candidates = parsePageCitationAssetCandidates( + chunk.metadata?.pageAssets ?? chunk.metadata?.page_assets, + ).filter(isSupportedPageCitationAsset) if (candidates.length === 0) return null const pageNumbers = getPageNumbers(chunk.metadata) @@ -463,20 +464,34 @@ function parsePageCitationAssetCandidates( return value.flatMap((item): PageCitationAssetCandidate[] => { if (!isRecord(item)) return [] - const pageNum = getPositiveInteger(item.pageNum) + const pageNum = + getPositiveInteger(item.pageNum) ?? + getPositiveInteger(item.page_num) ?? + getPositiveInteger(item.pageNumber) if (!pageNum) return [] return [ { pageNum, - ...(getTrimmedString(item.artifactRef) - ? { artifactRef: getTrimmedString(item.artifactRef) ?? undefined } + ...(getTrimmedString(item.artifactRef ?? item.artifact_ref) + ? { + artifactRef: + getTrimmedString(item.artifactRef ?? item.artifact_ref) ?? + undefined, + } : {}), - ...(getTrimmedString(item.assetUrl) - ? { assetUrl: getTrimmedString(item.assetUrl) ?? undefined } + ...(getTrimmedString(item.assetUrl ?? item.asset_url) + ? { + assetUrl: + getTrimmedString(item.assetUrl ?? item.asset_url) ?? undefined, + } : {}), - ...(getTrimmedString(item.contentType) - ? { contentType: getTrimmedString(item.contentType) ?? undefined } + ...(getTrimmedString(item.contentType ?? item.content_type) + ? { + contentType: + getTrimmedString(item.contentType ?? item.content_type) ?? + undefined, + } : {}), }, ] @@ -507,7 +522,12 @@ function getPageNumbers( ): readonly number[] { if (!metadata) return [] - const values = [metadata.pageNums, metadata.page_nums, metadata.pageNum] + const values = [ + metadata.pageNums, + metadata.page_nums, + metadata.pageNum, + metadata.page_num, + ] const pageNumbers = new Set() for (const value of values) { diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index 12873d3..5be40e4 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -46,9 +46,21 @@ describe("agent harness runtime", () => { expect(prompt).toContain("[[cite:n]]") expect(prompt).toContain("1-based index into the citations array") expect(prompt).toContain("Do not write title/pN, [1], Markdown footnotes") + expect(prompt).toContain("Never group indices as [[cite:1, 3, 5]]") expect(prompt).toContain("Do not collapse same-page citations") }) + it("requires inspectImage on cited page/image assets before finalize", () => { + const prompt = buildHarnessSystemPrompt(makeTurnInput()) + + expect(prompt).toContain( + "call inspectImage on the page/image assets you will cite before finalize", + ) + expect(prompt).toContain( + "Do not finalize cited page/image assets from chunk text alone", + ) + }) + it("passes only outer retrieval parameters to KNOWHERE without planning-tool gating", async () => { const query = vi.fn().mockResolvedValue( makeRetrievalResponse(), @@ -217,14 +229,21 @@ describe("agent harness runtime", () => { expect(inspectImages).not.toHaveBeenCalled() }) - it("enforces the inspectImage six-ref limit", async () => { + it("lets retrieval bound the number of images inspected in one call", async () => { const ledger = createEvidenceLedger() ledger.addRetrievalResponse(makeImageRetrievalResponse(7)) + const inspectImages = vi.fn(async (request: { + readonly assets: readonly { readonly ref: string; readonly label: string }[] + }) => ({ + analysis: "Compared all retrieved images.", + inspected: request.assets.map(({ ref, label }) => ({ ref, label })), + skipped: [], + })) const tools = createHarnessTools({ state: {}, ledger, knowhereTools: makeKnowhereTools(), - inspectImages: vi.fn(), + inspectImages, recentTurns: [], }) @@ -233,10 +252,8 @@ describe("agent harness runtime", () => { question: "Compare these images.", }) - expect(result).toMatchObject({ - ok: false, - message: "inspectImage accepts at most 6 refs per call.", - }) + expect(result).toMatchObject({ ok: true }) + expect(inspectImages.mock.calls[0]?.[0].assets).toHaveLength(7) }) it("calls the visual inspection capability with retrieved image ledger assets", async () => { @@ -345,6 +362,147 @@ describe("agent harness runtime", () => { }) }) + it("inspects duplicate retrieved pages only once across namespaces", async () => { + const ledger = createEvidenceLedger() + const pageMetadata = { + pageNums: [4], + pageAssets: [ + { + pageNum: 4, + artifactRef: "page_citation_assets/page-4.png", + assetUrl: "https://assets.example/page-4.png", + contentType: "image/png", + }, + ], + } + ledger.addRetrievalResponse({ + namespace: "default,notebook", + query: "revenue", + routerUsed: "mapnav", + answerText: null, + evidenceText: "Revenue evidence", + stopReason: "completed", + failureReason: null, + results: [ + { + chunkId: "chunk_page_4", + content: "Revenue was $24.9B.", + chunkType: "page", + score: 0.9, + metadata: pageMetadata, + source: { + documentId: "doc_catalog", + sourceFileName: "original.pdf", + sectionPath: "FINANCIAL SUMMARY", + }, + }, + { + chunkId: "chunk_page_4", + content: "Revenue was $24.9B.", + chunkType: "page", + score: 0.9, + metadata: pageMetadata, + source: { + documentId: "doc_workspace", + sourceFileName: "TSLA-Q4-2025-Update.pdf", + sectionPath: "FINANCIAL SUMMARY", + }, + }, + ], + referencedChunks: [], + }) + const inspectImages = vi.fn().mockResolvedValue({ + analysis: "Revenue is visible in the financial summary.", + inspected: [ + { + ref: "asset:r1:result:1", + label: "original.pdf / page_citation_assets/page-4.png / page", + }, + ], + skipped: [], + highlights: [ + { + ref: "asset:r1:result:1", + regions: [{ x: 0.1, y: 0.2, w: 0.8, h: 0.1 }], + }, + ], + }) + const state: { + inspectedImageRefs?: string[] + } = {} + const tools = createHarnessTools({ + state, + ledger, + knowhereTools: makeKnowhereTools(), + inspectImages, + recentTurns: [], + }) + + const inspection = await executeTool(tools.inspectImage, { + refs: ["asset:r1:result:1", "asset:r1:result:2"], + question: "Locate the cited revenue.", + }) + + expect(inspection).toMatchObject({ ok: true }) + expect(inspectImages.mock.calls[0]?.[0].assets).toHaveLength(1) + expect( + await executeTool(tools.finalize, { + text: "Revenue was $24.9B [[cite:1]] [[cite:2]].", + citations: [{ ref: "r1:result:1" }, { ref: "r1:result:2" }], + artifacts: [], + unresolved: [], + }), + ).toMatchObject({ ok: true }) + }) + + it("does not treat skipped image assets as successfully inspected", async () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse(makePageCitationRetrievalResponse()) + const state: { + inspectedImageRefs?: string[] + finalized?: boolean + } = {} + const tools = createHarnessTools({ + state, + ledger, + knowhereTools: makeKnowhereTools(), + inspectImages: vi.fn().mockResolvedValue({ + analysis: "", + inspected: [], + skipped: [ + { + ref: "asset:r1:referenced:1", + reason: "The image asset was unavailable in Notebook storage.", + }, + ], + }), + recentTurns: [], + }) + + const inspection = await executeTool(tools.inspectImage, { + refs: ["asset:r1:referenced:1"], + question: "Locate the cited amount.", + }) + + expect(inspection).toMatchObject({ + ok: false, + inspected: [], + }) + expect(state.inspectedImageRefs).toEqual([]) + + const finalize = await executeTool(tools.finalize, { + text: "The amount is 5000 yuan [[cite:1]].", + citations: [{ ref: "r1:referenced:1" }], + artifacts: [], + unresolved: [], + }) + expect(finalize).toMatchObject({ + ok: false, + inspectRefs: ["asset:r1:referenced:1"], + }) + expect(state.finalized).not.toBe(true) + }) + it("accepts finalize output without planning-tool gating", async () => { const state: { finalizedManifest?: OutputManifest @@ -372,6 +530,104 @@ describe("agent harness runtime", () => { expect(state.finalized).toBe(true) }) + it("rejects finalize of cited page images until inspectImage has run", async () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse(makePageCitationRetrievalResponse()) + const state: { + finalizedManifest?: OutputManifest + finalized?: boolean + inspectedImageRefs?: string[] + } = {} + const tools = createHarnessTools({ + state, + ledger, + knowhereTools: makeKnowhereTools(), + inspectImages: vi.fn().mockResolvedValue({ + analysis: "The clause shows 5000 yuan per occurrence.", + inspected: [{ ref: "asset:r1:referenced:1", label: "page 8" }], + skipped: [], + highlights: [ + { + ref: "asset:r1:referenced:1", + regions: [{ x: 0.1, y: 0.2, w: 0.4, h: 0.15 }], + }, + ], + }), + recentTurns: [], + }) + const manifest = { + text: "The contractor pays 5000 yuan per occurrence [[cite:1]].", + citations: [{ ref: "r1:referenced:1" }], + artifacts: [], + unresolved: [], + } + + const blocked = await executeTool(tools.finalize, manifest) + + expect(blocked).toMatchObject({ + ok: false, + inspectRefs: ["asset:r1:referenced:1"], + }) + expect(blocked).toEqual( + expect.objectContaining({ + message: expect.stringContaining("inspectImage"), + }), + ) + expect(state.finalized).not.toBe(true) + + await executeTool(tools.inspectImage, { + refs: ["asset:r1:referenced:1"], + question: "Locate the cited liquidated-damages amount on this page.", + }) + + expect(await executeTool(tools.finalize, manifest)).toMatchObject({ + ok: true, + text: manifest.text, + }) + expect(state.finalized).toBe(true) + }) + + it("requires inspection when the cited result shares a page asset with a referenced chunk", async () => { + const response = makePageCitationRetrievalResponse() + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse({ + ...response, + results: [ + { + chunkId: "chunk_page_8", + content: "The contractor pays 5000 yuan per occurrence.", + chunkType: "page", + score: 0.9, + metadata: { pageNums: [8] }, + source: { + documentId: "doc_contract", + sourceFileName: "contract.pdf", + sectionPath: "Root / Liquidated damages", + }, + }, + ], + }) + const tools = createHarnessTools({ + state: {}, + ledger, + knowhereTools: makeKnowhereTools(), + inspectImages: vi.fn(), + recentTurns: [], + }) + + const result = await executeTool(tools.finalize, { + text: "The contractor pays 5000 yuan [[cite:1]].", + citations: [{ ref: "r1:result:1" }], + artifacts: [], + unresolved: [], + }) + + expect(result).toMatchObject({ + ok: false, + inspectRefs: ["asset:r1:referenced:1"], + }) + }) + it("exposes full prior-turn content through policy-approved readPriorTurn", async () => { const state: { contextPolicy?: ContextPolicy @@ -610,11 +866,30 @@ describe("agent harness runtime", () => { }, { role: "user", - content: expect.stringContaining("call inspectImage now"), + content: expect.stringContaining("Call inspectImage now"), }, ]) }) + it("keeps forcing inspectImage at the finalization step while page assets are uninspected", () => { + const result = prepareHarnessStep({ + stepNumber: 13, + hasUninspectedImageAssets: true, + messages: [ + { + role: "user", + content: "What is the penalty amount?", + }, + ], + }) + + expect(result.activeTools).toEqual(["inspectImage"]) + expect(result.toolChoice).toEqual({ + type: "tool", + toolName: "inspectImage", + }) + }) + it("forces finalize at step 13 using existing tool results", () => { const result = prepareHarnessStep({ stepNumber: 13, diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index 3930146..38cc9fc 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -1,5 +1,4 @@ import { - hasToolCall, stepCountIs, ToolLoopAgent, tool, @@ -15,6 +14,8 @@ import type { AgentTurn, AgentTurnInput, ContextPolicy, + EvidenceAsset, + EvidenceChunk, HarnessRunResult, HarnessToolCallTrace, HarnessTrace, @@ -31,7 +32,6 @@ import type { const defaultMaxSteps = 14 const imageInspectionReminderStepNumber = 12 const forcedFinalizationStepNumber = 13 -const imageInspectionRefLimit = 6 type ToolLoopAgentSettings = ConstructorParameters[0] @@ -229,7 +229,7 @@ export async function runAgentHarness( hasUninspectedImageAssets({ state, ledger }), }), stopWhen: [ - hasToolCall("finalize"), + () => state.finalized === true, stepCountIs(input.maxSteps ?? defaultMaxSteps), ], }) @@ -263,10 +263,12 @@ export function prepareHarnessStep(input: { }): HarnessStepPreparation { const messages = sanitizeHarnessModelMessagesForStep(input.messages) - if ( - input.stepNumber === imageInspectionReminderStepNumber && - input.hasUninspectedImageAssets === true - ) { + const shouldForceImageInspection = + input.hasUninspectedImageAssets === true && + input.stepNumber >= imageInspectionReminderStepNumber && + input.stepNumber <= forcedFinalizationStepNumber + + if (shouldForceImageInspection) { return { messages: [ ...messages, @@ -372,9 +374,10 @@ function buildForcedFinalizationFeedback(): string { function buildImageInspectionReminderFeedback(): string { return [ - "The retrieval step budget is nearly reached and retrieved image assets are available.", - "If the exact answer depends on OCR, page-image text, visual details, or image verification, call inspectImage now with the most relevant retrieved image asset refs.", - "If image inspection is not needed for this answer, call finalize using the evidence already available.", + "Retrieved image/page assets are available and have not been inspected.", + "Call inspectImage now with the page/image asset refs you will cite.", + "Use a question that locates the cited evidence on those pages for OCR/visual context and provenance boxes.", + "Do not finalize until those cited image assets have been inspected.", "Do not search again.", ].join("\n") } @@ -383,12 +386,24 @@ function hasUninspectedImageAssets(input: { readonly state: HarnessToolState readonly ledger: ReturnType }): boolean { - const inspectedRefs = new Set(input.state.inspectedImageRefs ?? []) - return input.ledger - .snapshot() - .assets.some( - (asset) => asset.type === "image" && !inspectedRefs.has(asset.ref), - ) + const snapshot = input.ledger.snapshot() + const chunksByRef = new Map( + snapshot.chunks.map((chunk) => [chunk.ref, chunk] as const), + ) + const assetsByRef = new Map( + snapshot.assets.map((asset) => [asset.ref, asset] as const), + ) + const inspectedKeys = new Set( + (input.state.inspectedImageRefs ?? []).map((ref) => { + const asset = assetsByRef.get(ref) + return asset ? getCanonicalImageAssetKey(asset, chunksByRef) : ref + }), + ) + return snapshot.assets.some( + (asset) => + asset.type === "image" && + !inspectedKeys.has(getCanonicalImageAssetKey(asset, chunksByRef)), + ) } export function createHarnessTools(input: { @@ -528,9 +543,9 @@ export function createHarnessTools(input: { inspectImage: tool({ description: - "Inspect Knowhere image asset refs visually for OCR, visual details, comparisons, or verification. Use only after a Knowhere tool has returned image assets.", + "Inspect cited Knowhere page/image asset refs for OCR, visual details, and provenance boxes. Call this after retrieval and before finalize whenever the answer cites page or image assets.", inputSchema: z.object({ - refs: z.array(z.string().min(1)).min(1).max(imageInspectionRefLimit), + refs: z.array(z.string().min(1)).min(1), question: z.string().min(1), }), execute: async (request) => @@ -612,13 +627,28 @@ export function createHarnessTools(input: { description: "Finalize the user-facing output manifest. This is the only final answer " + "contract. Artifacts listed here with display=true are the exact set of " + - "images/tables shown to the user; cite evidence refs when available.", + "images/tables shown to the user; cite evidence refs when available. " + + "Cited page/image assets must be inspected with inspectImage first.", inputSchema: outputManifestSchema, execute: async (manifest) => traceToolCall(input.state, { toolName: "finalize", inputSummary: summarizeManifest(manifest), execute: async () => { + const inspectRefs = getUninspectedCitedImageRefs({ + manifest, + ledger: input.ledger, + inspectedImageRefs: input.state.inspectedImageRefs ?? [], + inspectImagesAvailable: input.inspectImages !== undefined, + }) + if (inspectRefs.length > 0) { + return { + ok: false as const, + message: buildFinalizeRequiresInspectionMessage(inspectRefs), + inspectRefs, + } + } + input.state.finalizedManifest = manifest input.state.finalized = true return { ok: true as const, ...manifest } @@ -658,18 +688,6 @@ async function inspectRetrievedImages(input: { skipped: [], } } - if (refs.length > imageInspectionRefLimit) { - return { - ok: false, - message: `inspectImage accepts at most ${imageInspectionRefLimit} refs per call.`, - inspected: [], - skipped: refs.map((ref) => ({ - ref, - reason: "Too many refs were requested in one inspectImage call.", - })), - } - } - const snapshot = input.ledger.snapshot() if (snapshot.chunks.length === 0 && snapshot.assets.length === 0) { return { @@ -698,11 +716,15 @@ async function inspectRetrievedImages(input: { const assetsByRef = new Map( snapshot.assets.map((asset) => [asset.ref, asset] as const), ) + const chunksByRef = new Map( + snapshot.chunks.map((chunk) => [chunk.ref, chunk] as const), + ) const skipped: { readonly ref: string readonly reason: string }[] = [] const selectedAssets: ImageInspectionAsset[] = [] + const selectedKeys = new Set() for (const ref of refs) { const asset = assetsByRef.get(ref) @@ -721,6 +743,15 @@ async function inspectRetrievedImages(input: { continue } + const assetKey = getCanonicalImageAssetKey(asset, chunksByRef) + if (selectedKeys.has(assetKey)) { + skipped.push({ + ref, + reason: "Duplicate of another retrieved page selected for inspection.", + }) + continue + } + selectedKeys.add(assetKey) selectedAssets.push({ ref: asset.ref, label: asset.label, @@ -741,34 +772,32 @@ async function inspectRetrievedImages(input: { } const inspectedImageRefs = input.state.inspectedImageRefs ?? [] - const inspectedCountAfterCall = - inspectedImageRefs.length + selectedAssets.length - if (inspectedCountAfterCall > imageInspectionRefLimit) { - return { - ok: false, - message: `inspectImage accepts at most ${imageInspectionRefLimit} image refs per turn.`, - inspected: [], - skipped: selectedAssets.map((asset) => ({ - ref: asset.ref, - reason: "The per-turn image inspection limit would be exceeded.", - })), - } - } - - input.state.inspectedImageRefs = [ - ...inspectedImageRefs, - ...selectedAssets.map((asset) => asset.ref), - ] try { const response = await input.inspectImages({ question, assets: selectedAssets, }) + const selectedRefs = new Set(selectedAssets.map((asset) => asset.ref)) + const successfulRefs = response.inspected + .map((asset) => asset.ref) + .filter((ref) => selectedRefs.has(ref)) + input.state.inspectedImageRefs = [ + ...inspectedImageRefs, + ...successfulRefs.filter((ref) => !inspectedImageRefs.includes(ref)), + ] input.state.imageHighlights = mergeImageInspectionHighlights( input.state.imageHighlights, response.highlights, ) + if (successfulRefs.length === 0) { + return { + ok: false as const, + message: "Image inspection skipped every requested asset.", + inspected: [], + skipped: [...skipped, ...response.skipped], + } + } return { ok: true as const, analysis: response.analysis, @@ -791,6 +820,102 @@ async function inspectRetrievedImages(input: { } } +function getUninspectedCitedImageRefs(input: { + readonly manifest: OutputManifest + readonly ledger: ReturnType + readonly inspectedImageRefs: readonly string[] + readonly inspectImagesAvailable: boolean +}): string[] { + if (!input.inspectImagesAvailable) return [] + + const snapshot = input.ledger.snapshot() + const chunksByRef = new Map( + snapshot.chunks.map((chunk) => [chunk.ref, chunk] as const), + ) + const assetsByRef = new Map( + snapshot.assets.map((asset) => [asset.ref, asset] as const), + ) + const inspected = new Set(input.inspectedImageRefs) + const inspectedKeys = new Set( + input.inspectedImageRefs.map((ref) => { + const asset = assetsByRef.get(ref) + return asset ? getCanonicalImageAssetKey(asset, chunksByRef) : ref + }), + ) + + const refs: string[] = [] + const seenKeys = new Set() + const addRef = (ref: string | null): void => { + if (!ref || inspected.has(ref)) return + const asset = assetsByRef.get(ref) + const key = asset ? getCanonicalImageAssetKey(asset, chunksByRef) : ref + if (seenKeys.has(key) || inspectedKeys.has(key)) return + seenKeys.add(key) + refs.push(ref) + } + + for (const citation of input.manifest.citations) { + addRef(resolveImageAssetRef(citation.ref, chunksByRef, assetsByRef)) + } + + for (const artifact of input.manifest.artifacts) { + if (artifact.type !== "image" || artifact.display === false) continue + addRef(resolveImageAssetRef(artifact.ref, chunksByRef, assetsByRef)) + } + + return refs +} + +function resolveImageAssetRef( + ref: string, + chunksByRef: ReadonlyMap, + assetsByRef: ReadonlyMap, +): string | null { + const directAsset = assetsByRef.get(ref) + if (directAsset?.type === "image") return directAsset.ref + + const chunk = chunksByRef.get(ref) + if (!chunk) return null + + if (chunk.assetRef) { + const chunkAsset = assetsByRef.get(chunk.assetRef) + if (chunkAsset?.type === "image") return chunkAsset.ref + } + + if (!chunk.chunkId) return null + const sibling = Array.from(chunksByRef.values()).find( + (candidate) => + candidate.ref !== chunk.ref && + candidate.chunkId === chunk.chunkId && + candidate.source.documentId === chunk.source.documentId && + candidate.assetRef !== undefined && + assetsByRef.get(candidate.assetRef)?.type === "image", + ) + return sibling?.assetRef ?? null +} + +function getCanonicalImageAssetKey( + asset: EvidenceAsset, + chunksByRef: ReadonlyMap, +): string { + const chunk = chunksByRef.get(asset.chunkRef) + const chunkId = chunk?.chunkId?.trim() + const sourcePath = asset.sourcePath?.trim().toLowerCase() + if (chunkId && sourcePath) return `${chunkId}\u0000${sourcePath}` + return asset.ref +} + +function buildFinalizeRequiresInspectionMessage( + inspectRefs: readonly string[], +): string { + return [ + "Cited page/image assets must be inspected before finalize.", + `Call inspectImage with refs: ${inspectRefs.join(" ")}.`, + "Use a question that locates the cited evidence on those pages.", + "Then call finalize again, using the inspection notes in the answer.", + ].join(" ") +} + type KnowhereToolOperation = | "search" | "list_documents" @@ -1169,8 +1294,8 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "2. Call setContextPolicy when prior turns may influence this turn.", "3. When the policy needs prior-turn detail (references or corrections), call readPriorTurn for the relevant ids.", "4. Call knowhere_search when relevance search is needed. Use knowhere_list_documents, knowhere_get_document_outline, knowhere_read_chunks, and knowhere_grep_chunks for focused document reads.", - "5. For pixel-level details, OCR, visual comparison, image verification, or when the likely answer is only visible on a returned page/image asset, call inspectImage only after a Knowhere tool returned image asset refs.", - `6. inspectImage accepts at most ${imageInspectionRefLimit} image asset refs per call and per turn.`, + "5. After Knowhere returns image/page asset refs, call inspectImage on the page/image assets you will cite before finalize. This supplies OCR/visual context and provenance boxes.", + "6. Inspect each unique cited page once; retrieval already bounds the available evidence set.", "7. knowhere_read_chunks returns complete chunk bodies; control size with page/pageSize, sectionPath, startChunk/endChunk, chunkId, and chunkType.", "8. Call finalize with text, citations, artifacts, and unresolved issues when you are ready to answer.", "", @@ -1185,11 +1310,13 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "- Use type=derived_table only for tables you create from evidence; every derived_table.sourceRefs entry must reference evidence in the ledger.", "- Prefer citation and selected image/table artifact refs returned by Knowhere tools in the evidence ledger.", "- Place [[cite:n]] immediately after the supported claim. n is the 1-based index into the citations array passed to finalize.", + "- Write one marker per index: [[cite:1]] [[cite:3]] [[cite:5]]. Never group indices as [[cite:1, 3, 5]].", "- Do not write title/pN, [1], Markdown footnotes, or [Source N: ...] in the answer text. Notebook renders chips from [[cite:n]] and citation metadata.", "- Repeat [[cite:n]] when another claim uses the same page. Do not collapse same-page citations to one row.", "- Citation label and source metadata are optional. Notebook resolves citation metadata from evidence refs when possible.", "- If evidence is relevant but you cannot identify a supporting evidence ref, answer with unresolved issues instead of fabricating a ref.", "- inspectImage observations are inspection notes, not new source refs. Final citations and displayed image artifacts must use the original retrieved image asset refs.", + "- Do not finalize cited page/image assets from chunk text alone when inspectImage is available. Inspect those asset refs first, then write the answer using the inspection notes.", "- If text evidence identifies a relevant page/image but does not include the exact fact, inspect the returned image asset for OCR/detail before saying the answer is unavailable.", "- If evidence is insufficient, list it in unresolved instead of fabricating facts.", `Surface: ${turn.surface}`, diff --git a/src/app/e2e/citation-same-page/page.tsx b/src/app/e2e/citation-same-page/page.tsx index c825f91..d2c16f9 100644 --- a/src/app/e2e/citation-same-page/page.tsx +++ b/src/app/e2e/citation-same-page/page.tsx @@ -35,6 +35,7 @@ const chatMessages: ChatMessageView[] = [ chunkType: "page", score: 0.91, pageCitationPageNumber: 26, + highlightRegions: [{ x: 0.12, y: 0.18, w: 0.46, h: 0.08 }], source: { documentId: "doc_spacex", sourceFileName: "spacex-s1.pdf", @@ -45,6 +46,10 @@ const chatMessages: ChatMessageView[] = [ chunkType: "page", score: 0.89, pageCitationPageNumber: 26, + highlightRegions: [ + { x: 0.62, y: 0.52, w: 0.24, h: 0.06 }, + { x: 0.15, y: 0.68, w: 0.32, h: 0.05 }, + ], source: { documentId: "doc_spacex", sourceFileName: "spacex-s1.pdf", diff --git a/src/app/globals.css b/src/app/globals.css index 42112a3..38bea15 100644 --- a/src/app/globals.css +++ b/src/app/globals.css @@ -43,6 +43,7 @@ --color-secondary: var(--secondary); --color-primary-foreground: var(--primary-foreground); --color-primary: var(--primary); + --animate-citation-region-flash: citation-region-flash 4.6s ease-in-out forwards; --color-primary-light: var(--primary-light); --color-primary-dark: var(--primary-dark); --color-popover-foreground: var(--popover-foreground); @@ -96,6 +97,36 @@ animation: citation-pulse 1.2s ease-out; } +@keyframes citation-region-flash { + 0% { opacity: 0; } + 4.3% { opacity: 1; } + 8.7% { opacity: 0; } + 13% { opacity: 1; } + 17.4% { opacity: 0; } + 21.7% { opacity: 1; } + 82.6% { opacity: 1; } + 100% { opacity: 0; } +} + +.citation-region-flash { + z-index: 10; + display: block; + border-radius: 8px; + background-color: rgb(142 81 255 / 0.25); + animation: citation-region-flash 4.6s ease-in-out forwards; +} + +@media (prefers-reduced-motion: reduce) { + .citation-region-flash { + animation: citation-region-flash-reduced 3.6s ease-out forwards; + } + + @keyframes citation-region-flash-reduced { + 0%, 83% { opacity: 1; } + 100% { opacity: 0; } + } +} + @layer base { * { diff --git a/src/components/chat-citation-model.test.ts b/src/components/chat-citation-model.test.ts index a7de7ce..f8e8709 100644 --- a/src/components/chat-citation-model.test.ts +++ b/src/components/chat-citation-model.test.ts @@ -61,6 +61,70 @@ describe("chatCitationModel", () => { expect(groups[1]?.title).toBe("report.pdf") }) + it("uses citation or same-page artifact boxes without inventing a region", () => { + const citation: ChatCitationView = { + chunkType: "page", + score: 0.9, + pageCitationPageNumber: 4, + source: { + documentId: "doc_1", + sourceFileName: "TSLA-Q4-2025-Update.pdf", + sectionPath: "Page 4", + }, + } + + expect( + chatCitationModel.getListHighlightRegions( + { + ...citation, + highlightRegions: [{ x: 0.12, y: 0.18, w: 0.4, h: 0.1 }], + }, + [ + { + type: "image", + display: true, + highlightRegions: [{ x: 0.5, y: 0.5, w: 0.2, h: 0.2 }], + citation: { + chunkType: "page", + score: 0.8, + pageCitationPageNumber: 4, + source: { documentId: "doc_1" }, + }, + }, + ], + ), + ).toEqual([{ x: 0.12, y: 0.18, w: 0.4, h: 0.1 }]) + + expect( + chatCitationModel.getListHighlightRegions(citation, [ + { + type: "image", + display: true, + highlightRegions: [{ x: 0.1, y: 0.2, w: 0.3, h: 0.15 }], + citation: { + chunkType: "page", + score: 0.8, + pageCitationPageNumber: 4, + source: { documentId: "doc_1" }, + }, + }, + { + type: "image", + display: true, + highlightRegions: [{ x: 0.5, y: 0.5, w: 0.2, h: 0.2 }], + citation: { + chunkType: "page", + score: 0.7, + pageCitationPageNumber: 12, + source: { documentId: "doc_1" }, + }, + }, + ]), + ).toEqual([{ x: 0.1, y: 0.2, w: 0.3, h: 0.15 }]) + + expect(chatCitationModel.getListHighlightRegions(citation, [])).toEqual([]) + }) + it("embeds cite markers as links and leaves fenced code unchanged", () => { const markdown = chatCitationModel.embedCitationMarkersAsLinks( [ @@ -87,6 +151,26 @@ describe("chatCitationModel", () => { expect(markdown).not.toContain("[[cite:1]] and later") }) + it("expands grouped cite markers into one chip per index", () => { + const markdown = chatCitationModel.embedCitationMarkersAsLinks( + "Total revenue was $24,901 million [[cite:1, 3, 5]].", + [pageCitation, samePageCitation, otherFileCitation, pageCitation, samePageCitation], + { + doc_1: "TSLA-Q4-2025-Update.pdf", + doc_2: "report.pdf", + }, + ) + + expect(markdown).toContain( + "[TSLA-Q4-2025-Update.pdf/p26](knowhere-cite://1)", + ) + expect(markdown).toContain("[report.pdf/p3](knowhere-cite://3)") + expect(markdown).toContain( + "[TSLA-Q4-2025-Update.pdf/p26](knowhere-cite://5)", + ) + expect(markdown).not.toContain("[[cite:") + }) + it("drops unknown cite markers from display markdown", () => { const markdown = chatCitationModel.embedCitationMarkersAsLinks( "Unsupported claim [[cite:9]].", diff --git a/src/components/chat-citation-model.ts b/src/components/chat-citation-model.ts index c3bac75..daa9c6f 100644 --- a/src/components/chat-citation-model.ts +++ b/src/components/chat-citation-model.ts @@ -2,6 +2,7 @@ import { workspaceCitationState } from "@/components/workspace-citation-state" import type { ChatArtifactView, ChatCitationView, + ChatImageHighlightBox, ChatMessageView, } from "@/domains/chat/types" @@ -28,6 +29,7 @@ export const chatCitationModel = { getCitationChipLabel, getCopyMarkdown, getExportMarkdown, + getListHighlightRegions, getSourceTitle, groupCitationsByFile, isKnowhereCiteHref, @@ -144,6 +146,34 @@ function uniquePageLinkEntries( return uniqueEntries } +function getListHighlightRegions( + citation: ChatCitationView, + artifacts: readonly ChatArtifactView[] | undefined, +): readonly ChatImageHighlightBox[] { + if (citation.highlightRegions && citation.highlightRegions.length > 0) { + return citation.highlightRegions + } + + const pageNumber = workspaceCitationState.getCitationPageNumber(citation) + const matchingRegions = (artifacts ?? []).flatMap((artifact) => { + if (!artifact.highlightRegions || artifact.highlightRegions.length === 0) { + return [] + } + + const artifactPage = artifact.citation + ? workspaceCitationState.getCitationPageNumber(artifact.citation) + : null + if (pageNumber !== null && artifactPage !== pageNumber) { + return [] + } + + return [...artifact.highlightRegions] + }) + + if (matchingRegions.length > 0) return matchingRegions + return [] +} + function getCitationChipLabel( citation: ChatCitationView, sourceTitlesByDocumentId: Readonly>, @@ -193,11 +223,14 @@ function replaceCitationTokens( replaceIndex: (index: number) => string, ): string { return text - .replace(/\[\[cite:(\d+)\]\]/g, (_match, rawIndex: string) => { - const index = Number.parseInt(rawIndex, 10) - if (!Number.isSafeInteger(index) || index < 1) return "" - return replaceIndex(index) - }) + .replace( + /\[\[cite:\s*(\d+(?:\s*,\s*\d+)*)\s*\]\]/g, + (_match, rawIndices: string) => + parseCiteIndices(rawIndices) + .map((index) => replaceIndex(index)) + .filter((part) => part.length > 0) + .join(" "), + ) .replace(/\[Source\s+(\d+)\s*:\s*[^\]]*\]/g, (_match, rawIndex: string) => { const index = Number.parseInt(rawIndex, 10) if (!Number.isSafeInteger(index) || index < 1) return "" @@ -205,6 +238,14 @@ function replaceCitationTokens( }) } +function parseCiteIndices(rawIndices: string): readonly number[] { + return rawIndices.split(",").flatMap((rawIndex) => { + const index = Number.parseInt(rawIndex.trim(), 10) + if (!Number.isSafeInteger(index) || index < 1) return [] + return [index] + }) +} + function toCiteMarkdownLink(label: string, index: number): string { return `[${escapeMarkdownLinkLabel(label)}](${knowhereCiteHrefPrefix}${index})` } diff --git a/src/components/chat-message-list.test.ts b/src/components/chat-message-list.test.ts index 1f665fe..f62632e 100644 --- a/src/components/chat-message-list.test.ts +++ b/src/components/chat-message-list.test.ts @@ -190,6 +190,50 @@ describe("ChatMessageList", () => { pageCitationPageNumber: 4, }), "assistant_1:0", + [], + ); + }); + + it("passes persisted citation regions to page navigation", async () => { + const user = userEvent.setup(); + const onCitationClick = vi.fn(); + const regions = [ + { x: 0.19, y: 0.16, w: 0.51, h: 0.03 }, + { x: 0.04, y: 0.29, w: 0.15, h: 0.03 }, + ]; + + render( + React.createElement(ChatMessageList, { + messages: [ + { + id: "assistant_regions", + role: "assistant", + content: "Revenue was $24.9B [[cite:1]].", + citations: [ + { + chunkType: "page", + score: 0.9, + pageCitationPageNumber: 4, + highlightRegions: regions, + source: { + documentId: "doc_tsla", + sourceFileName: "TSLA-Q4-2025-Update.pdf", + sectionPath: "FINANCIAL SUMMARY", + }, + }, + ], + }, + ], + onCitationClick, + }), + ); + + await user.click(screen.getByTestId("citation-chip")); + + expect(onCitationClick).toHaveBeenCalledWith( + expect.objectContaining({ pageCitationPageNumber: 4 }), + "assistant_regions:0", + regions, ); }); diff --git a/src/components/chat-message-list.tsx b/src/components/chat-message-list.tsx index 6007f94..796f2ea 100644 --- a/src/components/chat-message-list.tsx +++ b/src/components/chat-message-list.tsx @@ -26,6 +26,7 @@ import type { ChatDiagramChartSpec } from "@/domains/chat/diagram"; import type { ChatArtifactView, ChatCitationView, + ChatImageHighlightBox, ChatMessageView, } from "@/domains/chat/types"; @@ -89,6 +90,7 @@ export type ChatMessageListProps = { readonly onCitationClick?: ( citation: ChatCitationView, citationId: string, + highlightRegions?: readonly ChatImageHighlightBox[], ) => void; readonly pendingCitationId?: string | null; readonly pendingStatusText?: string | null; @@ -223,6 +225,7 @@ function VirtualMessageRow({ readonly onCitationClick?: ( citation: ChatCitationView, citationId: string, + highlightRegions?: readonly ChatImageHighlightBox[], ) => void; readonly pendingCitationId?: string | null; readonly sourceTitlesByDocumentId: Readonly>; @@ -231,6 +234,19 @@ function VirtualMessageRow({ return null; } + const boundCitationClick = onCitationClick + ? (citation: ChatCitationView, citationId: string): void => { + onCitationClick( + citation, + citationId, + chatCitationModel.getListHighlightRegions( + citation, + message.artifacts, + ), + ) + } + : undefined + const rowStyle: CSSProperties = { position: "absolute", transform: `translateY(${virtualItem.start}px)`, @@ -247,7 +263,7 @@ function VirtualMessageRow({ diff --git a/src/components/chat-panel.test.ts b/src/components/chat-panel.test.ts index 62cce02..6ca70f9 100644 --- a/src/components/chat-panel.test.ts +++ b/src/components/chat-panel.test.ts @@ -376,7 +376,11 @@ describe("ChatPanel", () => { await user.click(duplicatedSourceLinks[0]); - expect(onCitationClick).toHaveBeenCalledWith(firstCitation, "assistant_1:0"); + expect(onCitationClick).toHaveBeenCalledWith( + firstCitation, + "assistant_1:0", + [], + ); }); it("keeps separate source links when different documents share one displayed label", async () => { @@ -430,6 +434,7 @@ describe("ChatPanel", () => { expect(onCitationClick).toHaveBeenCalledWith( secondCitation, "assistant_1:1", + [], ); }); diff --git a/src/components/chat-panel.tsx b/src/components/chat-panel.tsx index 1f0d1e8..85efa91 100644 --- a/src/components/chat-panel.tsx +++ b/src/components/chat-panel.tsx @@ -37,6 +37,7 @@ import { } from "@/components/ui/tooltip"; import type { ChatCitationView, + ChatImageHighlightBox, ChatMessageView, ChatThreadView, } from "@/domains/chat/types"; @@ -54,7 +55,11 @@ export type ChatPanelProps = { onNewChat?: () => void; onThreadSelect?: (threadId: string) => void; onThreadArchive?: (threadId: string) => void; - onCitationClick?: (citation: ChatCitationView, citationId: string) => void; + onCitationClick?: ( + citation: ChatCitationView, + citationId: string, + highlightRegions?: readonly ChatImageHighlightBox[], + ) => void; onLoginClick?: () => void; sourceTitlesByDocumentId?: Readonly>; sourceCount?: number; diff --git a/src/components/chunks-panel.tsx b/src/components/chunks-panel.tsx index 825f819..84298bd 100644 --- a/src/components/chunks-panel.tsx +++ b/src/components/chunks-panel.tsx @@ -37,6 +37,7 @@ import { } from "@/components/ui/tooltip"; import { useChunksPanelWorkflow } from "@/components/chunks-panel-workflow"; import { ParsedChunkCard } from "@/components/parsed-chunk-card"; +import type { ChatImageHighlightBox } from "@/domains/chat/types"; import { chunksPanelState } from "@/components/chunks-panel-state"; import { MAX_UPLOAD_MB } from "@/domains/sources/validation"; import { useSourceOriginalPreviewWarmup } from "@/components/source-original-preview-warmup"; @@ -59,6 +60,7 @@ export type ChunksPanelProps = { focusedPageNumber?: number | null; focusedPageRequestId?: number; focusedCitationId?: string | null; + focusedHighlightRegions?: readonly ChatImageHighlightBox[]; citationListViewRequestId?: number; isLoading?: boolean; isLoadingMore?: boolean; @@ -91,6 +93,7 @@ export function ChunksPanel({ focusedPageNumber = null, focusedPageRequestId = 0, focusedCitationId = null, + focusedHighlightRegions = [], citationListViewRequestId = 0, isLoading = false, isLoadingMore = false, @@ -433,6 +436,8 @@ export function ChunksPanel({ focusedChunkId={activeFocusedChunkId} focusedCitationId={focusedCitationId} focusedPageNumber={focusedPageNumber} + focusedPageRequestId={focusedPageRequestId} + highlightRegions={focusedHighlightRegions} isOriginalPreviewAvailable={isOriginalPreviewAvailable} measureElement={measureVirtualChunkElement} onChunkClick={ @@ -1216,6 +1221,8 @@ function VirtualChunkRow({ focusedChunkId, focusedCitationId, focusedPageNumber, + focusedPageRequestId, + highlightRegions, isOriginalPreviewAvailable, measureElement, onChunkClick, @@ -1227,6 +1234,8 @@ function VirtualChunkRow({ focusedChunkId: string | null; focusedCitationId: string | null; focusedPageNumber: number | null; + focusedPageRequestId: number; + highlightRegions: readonly ChatImageHighlightBox[]; isOriginalPreviewAvailable: boolean; measureElement: (node: HTMLDivElement | null) => void; onChunkClick?: (chunk: ParsedChunkView) => void; @@ -1257,6 +1266,8 @@ function VirtualChunkRow({ isFocused={chunk.chunkId === focusedChunkId} focusedCitationId={focusedCitationId} focusedPageNumber={focusedPageNumber} + focusedPageRequestId={focusedPageRequestId} + highlightRegions={highlightRegions} isOriginalPreviewAvailable={isOriginalPreviewAvailable} onChunkClick={onChunkClick} onReferenceClick={onReferenceClick} diff --git a/src/components/citation-region-highlight.tsx b/src/components/citation-region-highlight.tsx new file mode 100644 index 0000000..6eecf0e --- /dev/null +++ b/src/components/citation-region-highlight.tsx @@ -0,0 +1,38 @@ +import type { ReactNode } from "react" + +import type { ChatImageHighlightBox } from "@/domains/chat/types" + +export function CitationRegionHighlight({ + regions, + requestId, +}: { + readonly regions: readonly ChatImageHighlightBox[] + readonly requestId: number +}): ReactNode { + if (regions.length === 0) return null + + return ( + + ) +} diff --git a/src/components/parsed-chunk-card.test.ts b/src/components/parsed-chunk-card.test.ts index 2a600f5..bcdb31e 100644 --- a/src/components/parsed-chunk-card.test.ts +++ b/src/components/parsed-chunk-card.test.ts @@ -157,6 +157,40 @@ describe("ParsedChunkCard", () => { ); }); + it("draws a temporary region highlight on the focused page image", () => { + render( + React.createElement(ParsedChunkCard, { + chunk: { + chunkId: "page_4", + type: "page", + content: "Page 4", + sourceTitle: "manual.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4.png", + contentType: "image/png", + }, + ], + }, + isFocused: true, + focusedCitationId: "assistant_1:0", + focusedPageNumber: 4, + focusedPageRequestId: 2, + highlightRegions: [{ x: 0.1, y: 0.2, w: 0.3, h: 0.15 }], + onReferenceClick: vi.fn(), + }), + ); + + const highlight = screen.getByTestId("citation-region-highlight"); + expect(highlight.className).toContain("citation-region-flash"); + expect(highlight.getAttribute("style")).toContain("10%"); + expect(highlight.getAttribute("style")).toContain("20%"); + expect(highlight.getAttribute("style")).toContain("30%"); + expect(highlight.getAttribute("style")).toContain("15%"); + }); + it("hides the keywords row when a page card has none", () => { render( React.createElement(ParsedChunkCard, { diff --git a/src/components/parsed-chunk-card.tsx b/src/components/parsed-chunk-card.tsx index 4dba392..9b5c1a3 100644 --- a/src/components/parsed-chunk-card.tsx +++ b/src/components/parsed-chunk-card.tsx @@ -15,6 +15,8 @@ import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Card, CardContent } from "@/components/ui/card"; import { parsedChunkCardModel } from "@/components/parsed-chunk-card-model"; +import { CitationRegionHighlight } from "@/components/citation-region-highlight"; +import type { ChatImageHighlightBox } from "@/domains/chat/types"; import type { ParsedChunkView } from "@/domains/chunks/types"; import type { SourceOriginalFileView } from "@/domains/sources/types"; import { cn } from "@/lib/utils"; @@ -33,6 +35,8 @@ export function ParsedChunkCard({ isFocused, focusedCitationId = null, focusedPageNumber = null, + focusedPageRequestId = 0, + highlightRegions = [], isOriginalPreviewAvailable = false, onChunkClick, onReferenceClick, @@ -42,6 +46,8 @@ export function ParsedChunkCard({ readonly isFocused: boolean; readonly focusedCitationId?: string | null; readonly focusedPageNumber?: number | null; + readonly focusedPageRequestId?: number; + readonly highlightRegions?: readonly ChatImageHighlightBox[]; readonly isOriginalPreviewAvailable?: boolean; readonly onChunkClick?: (chunk: ParsedChunkView) => void; readonly onReferenceClick: (chunkId: string) => void; @@ -55,6 +61,8 @@ export function ParsedChunkCard({ isFocused={isFocused} focusedCitationId={focusedCitationId} focusedPageNumber={focusedPageNumber} + focusedPageRequestId={focusedPageRequestId} + highlightRegions={highlightRegions} isOriginalPreviewAvailable={isOriginalPreviewAvailable} onChunkClick={onChunkClick} /> @@ -67,6 +75,8 @@ export function ParsedChunkCard({ void; }): ReactNode { @@ -424,6 +438,8 @@ function PageChunkCard({ assets={pageAssets} focusedCitationId={focusedCitationId} focusedPageNumber={focusedPageNumber} + focusedPageRequestId={focusedPageRequestId} + highlightRegions={highlightRegions} /> ) : ( @@ -442,10 +458,14 @@ function PageCitationAssets({ assets, focusedCitationId, focusedPageNumber, + focusedPageRequestId, + highlightRegions, }: { readonly assets: NonNullable; readonly focusedCitationId: string | null; readonly focusedPageNumber: number | null; + readonly focusedPageRequestId: number; + readonly highlightRegions: readonly ChatImageHighlightBox[]; }): ReactNode { return (
@@ -454,6 +474,8 @@ function PageCitationAssets({ key={asset.pageNumber} asset={asset} focusedCitationId={focusedCitationId} + focusedPageRequestId={focusedPageRequestId} + highlightRegions={highlightRegions} isCitationFocus={ focusedCitationId !== null && focusedPageNumber === asset.pageNumber @@ -467,10 +489,14 @@ function PageCitationAssets({ function PageCitationAssetImage({ asset, focusedCitationId, + focusedPageRequestId, + highlightRegions, isCitationFocus, }: { readonly asset: NonNullable[number]; readonly focusedCitationId: string | null; + readonly focusedPageRequestId: number; + readonly highlightRegions: readonly ChatImageHighlightBox[]; readonly isCitationFocus: boolean; }): ReactNode { const [failedAssetUrl, setFailedAssetUrl] = useState(null); @@ -503,6 +529,12 @@ function PageCitationAssetImage({ onError={() => setFailedAssetUrl(imageAssetUrl)} /> )} + {isCitationFocus ? ( + + ) : null}
); @@ -531,12 +563,16 @@ function PageCitationAssetUnavailable({ function ImageChunkCard({ chunk, isFocused, + focusedPageRequestId, + highlightRegions, isOriginalPreviewAvailable, onChunkClick, sourceOriginalFile, }: { readonly chunk: ParsedChunkView; readonly isFocused: boolean; + readonly focusedPageRequestId: number; + readonly highlightRegions: readonly ChatImageHighlightBox[]; readonly isOriginalPreviewAvailable: boolean; readonly onChunkClick?: (chunk: ParsedChunkView) => void; readonly sourceOriginalFile: SourceOriginalFileView | null; @@ -557,12 +593,20 @@ function ImageChunkCard({ {inlineImageAssetUrl ? (
- {/* eslint-disable-next-line @next/next/no-img-element -- Parsed artifact dimensions are not known before render. */} - {chunk.summary +
+ {/* eslint-disable-next-line @next/next/no-img-element -- Parsed artifact dimensions are not known before render. */} + {chunk.summary + {isFocused ? ( + + ) : null} +
) : (
diff --git a/src/components/workspace-citation-focus.test.ts b/src/components/workspace-citation-focus.test.ts index 1386c88..c3508b1 100644 --- a/src/components/workspace-citation-focus.test.ts +++ b/src/components/workspace-citation-focus.test.ts @@ -226,8 +226,20 @@ describe("useWorkspaceCitationFocus", () => { pageNumber: 4, requestId: 1, citationId: "message_1:0", + highlightRegions: [], }); expect(result.current.citationListViewRequestId).toBe(1); + + await act(async () => { + await result.current.handleCitationClick(pageCitation, "message_1:0", [ + { x: 0.1, y: 0.2, w: 0.3, h: 0.15 }, + ]); + }); + + expect(result.current.focusedPage.highlightRegions).toEqual([ + { x: 0.1, y: 0.2, w: 0.3, h: 0.15 }, + ]); + expect(result.current.focusedPage.requestId).toBe(2); }); it("uses cached SWR pages for file B while viewing A without loading all chunk types", async () => { diff --git a/src/components/workspace-citation-focus.ts b/src/components/workspace-citation-focus.ts index 379dc29..f12d277 100644 --- a/src/components/workspace-citation-focus.ts +++ b/src/components/workspace-citation-focus.ts @@ -5,7 +5,7 @@ import { useSWRConfig } from "swr" import { workspaceCitationState } from "@/components/workspace-citation-state" import { useWorkspaceSelectedChunks } from "@/components/workspace-selected-chunks" -import type { ChatCitationView } from "@/domains/chat/types" +import type { ChatCitationView, ChatImageHighlightBox } from "@/domains/chat/types" import type { ParsedChunkView } from "@/domains/chunks/types" import type { SourceView } from "@/domains/sources/types" import { @@ -22,6 +22,7 @@ type FocusedPageState = { readonly pageNumber: number | null readonly requestId: number readonly citationId: string | null + readonly highlightRegions: readonly ChatImageHighlightBox[] } type PrefetchedChunksBySourceId = Readonly> @@ -47,6 +48,7 @@ type WorkspaceCitationFocus = { readonly handleCitationClick: ( citation: ChatCitationView, citationId: string, + highlightRegions?: readonly ChatImageHighlightBox[], ) => Promise readonly handleLoadMoreChunks: () => void readonly handleLoadAllChunks: () => void @@ -78,6 +80,7 @@ export function useWorkspaceCitationFocus({ pageNumber: null, requestId: 0, citationId: null, + highlightRegions: [], }) const [pendingCitationId, setPendingCitationId] = useState( null, @@ -122,10 +125,15 @@ export function useWorkspaceCitationFocus({ [], ) const requestPageFocus = useCallback( - (pageNumber: number | null, citationId: string | null = null): void => { + ( + pageNumber: number | null, + citationId: string | null = null, + highlightRegions: readonly ChatImageHighlightBox[] = [], + ): void => { setFocusedPage((current) => ({ pageNumber, citationId, + highlightRegions: resolveFocusHighlightRegions(highlightRegions), requestId: current.requestId + 1, })) }, @@ -245,6 +253,7 @@ export function useWorkspaceCitationFocus({ async ( citation: ChatCitationView, citationId: string, + highlightRegions?: readonly ChatImageHighlightBox[], ): Promise => { setPendingCitationId(citationId) @@ -275,7 +284,7 @@ export function useWorkspaceCitationFocus({ const applyFocus = (chunkId: string | null): void => { if (selectedSourceId !== source.id) onSelectSource(source.id) requestChunkFocus(chunkId) - requestPageFocus(pageNumber, citationId) + requestPageFocus(pageNumber, citationId, highlightRegions) } if (selectedSourceId === source.id) { @@ -288,7 +297,7 @@ export function useWorkspaceCitationFocus({ }) if (loadedChunkId) { requestChunkFocus(loadedChunkId) - requestPageFocus(pageNumber, citationId) + requestPageFocus(pageNumber, citationId, highlightRegions) return } } @@ -402,3 +411,10 @@ export function useWorkspaceCitationFocus({ selectedSource, } } + +function resolveFocusHighlightRegions( + highlightRegions: readonly ChatImageHighlightBox[] | undefined, +): readonly ChatImageHighlightBox[] { + if (highlightRegions && highlightRegions.length > 0) return highlightRegions + return [] +} diff --git a/src/components/workspace-shell-layout.tsx b/src/components/workspace-shell-layout.tsx index 33c9394..5a8082c 100644 --- a/src/components/workspace-shell-layout.tsx +++ b/src/components/workspace-shell-layout.tsx @@ -20,6 +20,7 @@ import { useWorkspaceResizeHandleWorkflow } from "@/components/workspace-resize- import { workspaceShellState } from "@/components/workspace-shell-state" import type { ChatCitationView, + ChatImageHighlightBox, ChatMessageView, ChatThreadView, } from "@/domains/chat/types" @@ -46,6 +47,7 @@ type FocusedPageState = { readonly pageNumber: number | null readonly requestId: number readonly citationId?: string | null + readonly highlightRegions?: readonly ChatImageHighlightBox[] } type WorkspaceShellUser = { @@ -109,6 +111,7 @@ export type WorkspaceShellLayoutProps = { readonly onCitationClick: ( citation: ChatCitationView, citationId: string, + highlightRegions?: readonly ChatImageHighlightBox[], ) => void | Promise readonly onCreateChatThread: () => void | Promise readonly onDesktopLayoutElementChange: (element: HTMLDivElement | null) => void @@ -162,7 +165,12 @@ export function WorkspaceShellLayout( const selectedSource = props.selectedSourceView ?? props.sources.find((source) => source.id === props.selectedSourceId) - const focusedPage = props.focusedPage ?? { pageNumber: null, requestId: 0 } + const focusedPage = props.focusedPage ?? { + pageNumber: null, + requestId: 0, + citationId: null, + highlightRegions: [], + } const handleDesktopLayoutRef = useCallback( (element: HTMLDivElement | null): void => { onDesktopLayoutElementChange(element) @@ -292,6 +300,7 @@ export function WorkspaceShellLayout( focusedPageNumber={focusedPage.pageNumber} focusedPageRequestId={focusedPage.requestId} focusedCitationId={focusedPage.citationId ?? null} + focusedHighlightRegions={focusedPage.highlightRegions ?? []} isLoading={props.isSelectedChunksLoading} isLoadingAllChunks={props.isSelectedAllChunksLoading} isLoadingMore={props.isSelectedChunksLoadingMore} @@ -438,6 +447,7 @@ export function WorkspaceShellLayout( focusedPageNumber={focusedPage.pageNumber} focusedPageRequestId={focusedPage.requestId} focusedCitationId={focusedPage.citationId ?? null} + focusedHighlightRegions={focusedPage.highlightRegions ?? []} isLoading={props.isSelectedChunksLoading} isLoadingAllChunks={props.isSelectedAllChunksLoading} isLoadingMore={props.isSelectedChunksLoadingMore} @@ -481,9 +491,9 @@ export function WorkspaceShellLayout( onThreadArchive={props.isGuest ? undefined : props.onArchiveChatThread} onLoginClick={props.isGuest ? props.onLoginClick : undefined} sourceTitlesByDocumentId={props.sourceTitlesByDocumentId} - onCitationClick={(citation, citationId) => { + onCitationClick={(citation, citationId, highlightRegions) => { props.onMobilePanelChange("content") - props.onCitationClick(citation, citationId) + props.onCitationClick(citation, citationId, highlightRegions) }} />
diff --git a/src/domains/chat/chat-citation-persistence.ts b/src/domains/chat/chat-citation-persistence.ts index 3753a9a..d48281d 100644 --- a/src/domains/chat/chat-citation-persistence.ts +++ b/src/domains/chat/chat-citation-persistence.ts @@ -83,6 +83,8 @@ function replaceDemoCitationDocumentId( function toCitationView( citation: ChatCitationView | CitationView | RetrievalResultView, ): CitationView { + const highlightRegions = + "highlightRegions" in citation ? citation.highlightRegions : undefined return { chunkType: citation.chunkType, score: citation.score, @@ -90,6 +92,9 @@ function toCitationView( pageCitationAssetUrl: citation.pageCitationAssetUrl, pageCitationPageNumber: citation.pageCitationPageNumber, description: "description" in citation ? citation.description : undefined, + ...(highlightRegions && highlightRegions.length > 0 + ? { highlightRegions } + : {}), source: { documentId: citation.source.documentId, sourceFileName: citation.source.sourceFileName, diff --git a/src/domains/chat/citations.test.ts b/src/domains/chat/citations.test.ts index c07a69f..b381a83 100644 --- a/src/domains/chat/citations.test.ts +++ b/src/domains/chat/citations.test.ts @@ -63,6 +63,27 @@ describe("toChatCitationViews", () => { expect(citations[0]?.pageCitationPageNumber).toBe(26) }) + + it("copies inspect-image provenance boxes onto the citation", () => { + const citations = toChatCitationViews( + [ + makeRetrievalResult({ + chunkType: "page", + highlightRegions: [{ x: 0.1, y: 0.2, w: 0.3, h: 0.15 }], + source: { + documentId: "doc_1", + sourceFileName: "TSLA-Q4-2025-Update.pdf", + sectionPath: "Page 4", + }, + }), + ], + "Automotive revenue was $17.7B.", + ) + + expect(citations[0]?.highlightRegions).toEqual([ + { x: 0.1, y: 0.2, w: 0.3, h: 0.15 }, + ]) + }) }) function makeRetrievalResult( diff --git a/src/domains/chat/citations.ts b/src/domains/chat/citations.ts index ea203f7..9fb6c8d 100644 --- a/src/domains/chat/citations.ts +++ b/src/domains/chat/citations.ts @@ -28,6 +28,9 @@ export function toChatCitationViews( ? { pageCitationPageNumber } : {}), ...(description ? { description } : {}), + ...(result.highlightRegions && result.highlightRegions.length > 0 + ? { highlightRegions: result.highlightRegions } + : {}), source: { documentId: result.source.documentId, sourceFileName: result.source.sourceFileName, diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 114ade8..9efcbb5 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -1360,6 +1360,214 @@ describe("answerQuestionWithRetrieval", () => { expect(answer.citations[0]?.source.sectionPath).toBe("FINANCIAL SUMMARY"); }); + it("copies inspect-image provenance boxes onto cited page results", async () => { + const pageResult = makeRetrievalResult({ + chunkType: "page", + metadata: { page_nums: [4] }, + source: { + documentId: "doc_tsla", + sourceFileName: "TSLA-Q4-2025-Update.pdf", + sectionPath: "Page 4", + }, + }); + const generateAnswer = vi.fn(async () => + makeHarnessRunResultWithLedger("Automotive revenue was $17.7B [[cite:1]].", { + citations: [makeOutputCitation("r1:result:1", pageResult)], + chunks: [ + { + ...makeEvidenceChunkFromRetrievalResult("r1:result:1", pageResult), + assetRef: "asset:r1:result:1", + }, + ], + assets: [ + { + ref: "asset:r1:result:1", + chunkRef: "r1:result:1", + type: "image", + source: pageResult.source, + label: "TSLA-Q4-2025-Update.pdf / Page 4", + }, + ], + imageHighlights: [ + { + ref: "asset:r1:result:1", + regions: [{ x: 0.1, y: 0.2, w: 0.3, h: 0.15 }], + }, + ], + }), + ); + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "Tesla automotive revenue", + namespace: "notebook-workspace", + sources: [ + makeSource({ + id: "source_tsla", + title: "TSLA-Q4-2025-Update.pdf", + knowhereDocumentId: "doc_tsla", + }), + ], + excludedSourceIds: [], + retrieval: { query: vi.fn() }, + generateAnswer, + messages: [], + }), + ); + + expect(answer.citations[0]?.highlightRegions).toEqual([ + { x: 0.1, y: 0.2, w: 0.3, h: 0.15 }, + ]); + }); + + it("copies boxes from a referenced page asset onto its cited result chunk", async () => { + const pageResult = makeRetrievalResult({ + chunkId: "chunk_page_4", + chunkType: "page", + metadata: { page_nums: [4] }, + source: { + documentId: "doc_tsla", + sourceFileName: "TSLA-Q4-2025-Update.pdf", + sectionPath: "FINANCIAL SUMMARY", + }, + }); + const resultChunk = makeEvidenceChunkFromRetrievalResult( + "r1:result:1", + pageResult, + ); + const referencedChunk = { + ...resultChunk, + ref: "r1:referenced:1", + kind: "referenced_chunk" as const, + assetRef: "asset:r1:referenced:1", + }; + const generateAnswer = vi.fn(async () => + makeHarnessRunResultWithLedger("Revenue was $24.9B [[cite:1]].", { + citations: [makeOutputCitation("r1:result:1", pageResult)], + chunks: [resultChunk, referencedChunk], + assets: [ + { + ref: "asset:r1:referenced:1", + chunkRef: "r1:referenced:1", + type: "image", + source: pageResult.source, + label: "TSLA-Q4-2025-Update.pdf / Page 4", + }, + ], + imageHighlights: [ + { + ref: "asset:r1:referenced:1", + regions: [{ x: 0.08, y: 0.42, w: 0.84, h: 0.12 }], + }, + ], + }), + ); + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "What was Tesla's Q4 revenue?", + namespace: "notebook-workspace", + sources: [ + makeSource({ + id: "source_tsla", + title: "TSLA-Q4-2025-Update.pdf", + knowhereDocumentId: "doc_tsla", + }), + ], + excludedSourceIds: [], + retrieval: { query: vi.fn() }, + generateAnswer, + messages: [], + }), + ); + + expect(answer.citations[0]?.highlightRegions).toEqual([ + { x: 0.08, y: 0.42, w: 0.84, h: 0.12 }, + ]); + }); + + it("reuses boxes from a deduplicated page alias across namespaces", async () => { + const catalogResult = makeRetrievalResult({ + chunkId: "chunk_page_4", + chunkType: "page", + metadata: { page_nums: [4] }, + source: { + documentId: "doc_catalog", + sourceFileName: "original.pdf", + sectionPath: "FINANCIAL SUMMARY", + }, + }); + const workspaceResult = makeRetrievalResult({ + ...catalogResult, + source: { + documentId: "doc_workspace", + sourceFileName: "TSLA-Q4-2025-Update.pdf", + sectionPath: "FINANCIAL SUMMARY", + }, + }); + const catalogChunk = { + ...makeEvidenceChunkFromRetrievalResult("r1:result:1", catalogResult), + assetRef: "asset:r1:result:1", + }; + const workspaceChunk = { + ...makeEvidenceChunkFromRetrievalResult("r1:result:2", workspaceResult), + assetRef: "asset:r1:result:2", + }; + const assets = [ + { + ref: "asset:r1:result:1", + chunkRef: "r1:result:1", + type: "image" as const, + sourcePath: "page_citation_assets/page-4.png", + source: catalogResult.source, + label: "original.pdf / Page 4", + }, + { + ref: "asset:r1:result:2", + chunkRef: "r1:result:2", + type: "image" as const, + sourcePath: "page_citation_assets/page-4.png", + source: workspaceResult.source, + label: "TSLA-Q4-2025-Update.pdf / Page 4", + }, + ]; + const generateAnswer = vi.fn(async () => + makeHarnessRunResultWithLedger("Revenue was $24.9B [[cite:1]].", { + citations: [makeOutputCitation("r1:result:2", workspaceResult)], + chunks: [catalogChunk, workspaceChunk], + assets, + imageHighlights: [ + { + ref: "asset:r1:result:1", + regions: [{ x: 0.2, y: 0.3, w: 0.5, h: 0.08 }], + }, + ], + }), + ); + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "What was Tesla's Q4 revenue?", + namespace: "notebook-workspace", + sources: [ + makeSource({ + id: "source_tsla", + title: "TSLA-Q4-2025-Update.pdf", + knowhereDocumentId: "doc_workspace", + }), + ], + excludedSourceIds: [], + retrieval: { query: vi.fn() }, + generateAnswer, + messages: [], + }), + ); + + expect(answer.citations[0]?.highlightRegions).toEqual([ + { x: 0.2, y: 0.3, w: 0.5, h: 0.08 }, + ]); + }); + it("returns only harness-selected artifacts when retrieval has extra media candidates", async () => { const frontAssetUrl = "https://blob.example/images/id-front.jpg"; const backAssetUrl = "https://blob.example/images/id-back.jpg"; @@ -2975,6 +3183,7 @@ function makeHarnessRunResultWithLedger( readonly chunks?: HarnessRunResult["trace"]["ledger"]["chunks"] readonly assets?: HarnessRunResult["trace"]["ledger"]["assets"] readonly artifacts?: HarnessRunResult["manifest"]["artifacts"] + readonly imageHighlights?: HarnessRunResult["trace"]["imageHighlights"] }, ): HarnessRunResult { const chunks = input.chunks ?? []; @@ -2987,6 +3196,7 @@ function makeHarnessRunResultWithLedger( }, trace: { ...makeHarnessRunResult("").trace, + imageHighlights: input.imageHighlights ?? [], ledger: { retrievalCount: chunks.length > 0 ? 1 : 0, chunks, diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index c4f7f85..9fc711c 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -903,11 +903,20 @@ function mapManifestCitationsToResults( resolveChunkForAssetRef(citation.ref, assetsByRef, chunksByRef) if (!chunk) continue - results.push( - toRetrievalResultFromEvidenceChunk( - mergeChunkPageMetadata(chunk, result.trace.ledger.chunks), - ), + const retrievalResult = toRetrievalResultFromEvidenceChunk( + mergeChunkPageMetadata(chunk, result.trace.ledger.chunks), + ) + const highlightRegions = getHighlightRegionsForChunk( + chunk, + result.trace.imageHighlights, + chunksByRef, + assetsByRef, ) + const resultWithHighlights = + highlightRegions && highlightRegions.length > 0 + ? { ...retrievalResult, highlightRegions } + : retrievalResult + results.push(resultWithHighlights as RetrievalResult) if (results.length >= MAX_CITATION_RESULTS) break } @@ -924,6 +933,75 @@ function resolveChunkForAssetRef( return chunksByRef.get(asset.chunkRef) } +function getHighlightRegionsForChunk( + chunk: EvidenceChunk, + imageHighlights: HarnessRunResult["trace"]["imageHighlights"], + chunksByRef: ReadonlyMap, + assetsByRef: ReadonlyMap, +): ChatImageHighlightBox[] | undefined { + if (!imageHighlights || imageHighlights.length === 0) return undefined + + const candidateRefs = new Set([chunk.ref]) + const assetRef = resolveCitationImageAssetRef( + chunk, + chunksByRef, + assetsByRef, + ) + if (assetRef) candidateRefs.add(assetRef) + const canonicalAssetKey = assetRef + ? getCanonicalCitationAssetKey(assetRef, chunksByRef, assetsByRef) + : null + + for (const page of imageHighlights) { + const isDirectMatch = candidateRefs.has(page.ref) + const isCanonicalMatch = + canonicalAssetKey !== null && + getCanonicalCitationAssetKey(page.ref, chunksByRef, assetsByRef) === + canonicalAssetKey + if ((!isDirectMatch && !isCanonicalMatch) || page.regions.length === 0) { + continue + } + return [...page.regions] + } + + return undefined +} + +function resolveCitationImageAssetRef( + chunk: EvidenceChunk, + chunksByRef: ReadonlyMap, + assetsByRef: ReadonlyMap, +): string | null { + if (chunk.assetRef && assetsByRef.get(chunk.assetRef)?.type === "image") { + return chunk.assetRef + } + if (!chunk.chunkId) return null + + const sibling = Array.from(chunksByRef.values()).find( + (candidate) => + candidate.ref !== chunk.ref && + candidate.chunkId === chunk.chunkId && + candidate.source.documentId === chunk.source.documentId && + candidate.assetRef !== undefined && + assetsByRef.get(candidate.assetRef)?.type === "image", + ) + return sibling?.assetRef ?? null +} + +function getCanonicalCitationAssetKey( + assetRef: string, + chunksByRef: ReadonlyMap, + assetsByRef: ReadonlyMap, +): string { + const asset = assetsByRef.get(assetRef) + if (!asset) return assetRef + const chunk = chunksByRef.get(asset.chunkRef) + const chunkId = chunk?.chunkId?.trim() + const sourcePath = asset.sourcePath?.trim().toLowerCase() + if (chunkId && sourcePath) return `${chunkId}\u0000${sourcePath}` + return assetRef +} + function mapDisplayedManifestArtifactsToResults( result: HarnessRunResult, ): RetrievalResult[] { diff --git a/src/domains/chat/page-citation-assets.ts b/src/domains/chat/page-citation-assets.ts index 7a388a7..d411b68 100644 --- a/src/domains/chat/page-citation-assets.ts +++ b/src/domains/chat/page-citation-assets.ts @@ -8,6 +8,12 @@ import type { HardenChatAssetUrl } from "./media-assets" export type PageCitationAssetRetrievalResult = RetrievalResult & { readonly pageCitationAssetUrl?: string readonly pageCitationPageNumber?: number + readonly highlightRegions?: readonly { + readonly x: number + readonly y: number + readonly w: number + readonly h: number + }[] } type EnrichRetrievalResultsWithPageCitationAssetUrlsInput = { diff --git a/src/domains/chat/route-answer.ts b/src/domains/chat/route-answer.ts index eaa8064..612c256 100644 --- a/src/domains/chat/route-answer.ts +++ b/src/domains/chat/route-answer.ts @@ -241,6 +241,14 @@ async function inspectChatImages(input: { } if (preparedAssets.length === 0) { + logger.warn("chat: image inspection skipped all assets", { + workspaceId: input.workspaceId, + requestedCount: input.request.assets.length, + skipped: skippedAssets.map((asset) => ({ + ref: asset.ref, + reason: asset.reason, + })), + }) return { analysis: "", inspected: [], diff --git a/src/domains/chat/types.ts b/src/domains/chat/types.ts index e5179b0..1d6f88f 100644 --- a/src/domains/chat/types.ts +++ b/src/domains/chat/types.ts @@ -15,12 +15,24 @@ export type RetrievalResultView = { } } +/** + * Normalized highlight box for page/image answer provenance. + * Origin top-left; values in [0, 1]. No per-region labels. + */ +export type ChatImageHighlightBox = { + readonly x: number + readonly y: number + readonly w: number + readonly h: number +} + /** * Persisted chat citation metadata. This deliberately excludes source chunk * text so Notebook never stores upstream chunk content in Postgres. */ export type CitationView = Omit & { readonly description?: string + readonly highlightRegions?: readonly ChatImageHighlightBox[] } /** @@ -31,17 +43,6 @@ export type ChatCitationView = CitationView & { readonly content?: string } -/** - * Normalized highlight box for page/image answer provenance. - * Origin top-left; values in [0, 1]. No per-region labels. - */ -export type ChatImageHighlightBox = { - readonly x: number - readonly y: number - readonly w: number - readonly h: number -} - export type ChatArtifactView = { readonly type: "image" | "table" | "derived_table" readonly ref?: string diff --git a/src/domains/chat/view.ts b/src/domains/chat/view.ts index f74c287..c11ea79 100644 --- a/src/domains/chat/view.ts +++ b/src/domains/chat/view.ts @@ -55,6 +55,7 @@ function toPersistedCitationViews(value: unknown): ChatCitationView[] | undefine pageCitationAssetUrl: getString(item.pageCitationAssetUrl), pageCitationPageNumber: getNumber(item.pageCitationPageNumber), description: getString(item.description), + highlightRegions: getHighlightRegions(item.highlightRegions), source: { documentId: getString(item.source.documentId), sourceFileName: getString(item.source.sourceFileName), @@ -91,6 +92,9 @@ function toPersistedArtifactViews(value: unknown): ChatArtifactView[] | undefine item.citation.pageCitationPageNumber, ), description: getString(item.citation.description), + highlightRegions: getHighlightRegions( + item.citation.highlightRegions, + ), source: { documentId: getString(item.citation.source.documentId), sourceFileName: getString(item.citation.source.sourceFileName), From cf9aac6d3d2e515a36b9cdbe5ce90d66d9bc86ce Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 22 Aug 2026 02:08:13 +0800 Subject: [PATCH 50/51] fix: harden inspected citation provenance Prevent cross-document region reuse and keep inspection enforcement and rendered overlays accurate. Co-authored-by: Cursor --- e2e/citation-dedupe.e2e.ts | 40 +++++++++ src/agent-harness/image-asset-identity.ts | 46 ++++++++++ src/agent-harness/ledger.test.ts | 3 + src/agent-harness/ledger.ts | 38 ++++---- src/agent-harness/runtime.test.ts | 87 ++++++++++++++++++- src/agent-harness/runtime.ts | 15 +--- src/app/globals.css | 4 +- src/components/chat-citation-model.test.ts | 32 +++++++ src/components/chat-citation-model.ts | 22 +++-- src/components/citation-region-highlight.tsx | 5 +- src/components/parsed-chunk-card.test.ts | 10 ++- src/components/parsed-chunk-card.tsx | 69 ++++++++------- .../chat/image-inspection-model.test.ts | 45 ++++++++++ src/domains/chat/image-inspection-model.ts | 46 ++++++++++ src/domains/chat/index.test.ts | 2 + src/domains/chat/index.ts | 7 +- src/domains/chat/page-citation-assets.test.ts | 12 ++- src/domains/chat/page-citation-assets.ts | 7 +- 18 files changed, 400 insertions(+), 90 deletions(-) create mode 100644 src/agent-harness/image-asset-identity.ts diff --git a/e2e/citation-dedupe.e2e.ts b/e2e/citation-dedupe.e2e.ts index a21d7d0..03c3b1a 100644 --- a/e2e/citation-dedupe.e2e.ts +++ b/e2e/citation-dedupe.e2e.ts @@ -176,6 +176,22 @@ test("draws citation regions instead of a full-page highlight", async ({ })), ) .toEqual({ left: "12%", top: "18%", width: "46%", height: "8%" }) + const pageImage = page.getByRole("img", { name: "Page 26" }) + await expect(pageImage).toBeVisible() + const imageBox = await pageImage.boundingBox() + const firstRegionBox = await firstRegions.first().boundingBox() + expect(imageBox).not.toBeNull() + expect(firstRegionBox).not.toBeNull() + expect(firstRegionBox!.x).toBeCloseTo( + imageBox!.x + imageBox!.width * 0.12, + 1, + ) + expect(firstRegionBox!.y).toBeCloseTo( + imageBox!.y + imageBox!.height * 0.18, + 1, + ) + expect(firstRegionBox!.width).toBeCloseTo(imageBox!.width * 0.46, 1) + expect(firstRegionBox!.height).toBeCloseTo(imageBox!.height * 0.08, 1) await chips.nth(1).click() const secondRegions = page.getByTestId("citation-region-highlight") @@ -198,4 +214,28 @@ test("draws citation regions instead of a full-page highlight", async ({ width: "100%", height: "100%", }) + const expectedSecondRegions = [ + { x: 0.62, y: 0.52, w: 0.24, h: 0.06 }, + { x: 0.15, y: 0.68, w: 0.32, h: 0.05 }, + ] + for (const [index, expectedRegion] of expectedSecondRegions.entries()) { + const regionBox = await secondRegions.nth(index).boundingBox() + expect(regionBox).not.toBeNull() + expect(regionBox!.x).toBeCloseTo( + imageBox!.x + imageBox!.width * expectedRegion.x, + 1, + ) + expect(regionBox!.y).toBeCloseTo( + imageBox!.y + imageBox!.height * expectedRegion.y, + 1, + ) + expect(regionBox!.width).toBeCloseTo( + imageBox!.width * expectedRegion.w, + 1, + ) + expect(regionBox!.height).toBeCloseTo( + imageBox!.height * expectedRegion.h, + 1, + ) + } }) diff --git a/src/agent-harness/image-asset-identity.ts b/src/agent-harness/image-asset-identity.ts new file mode 100644 index 0000000..66037b1 --- /dev/null +++ b/src/agent-harness/image-asset-identity.ts @@ -0,0 +1,46 @@ +import type { EvidenceAsset, EvidenceChunk } from "./types" + +export function getCanonicalImageAssetKey( + asset: EvidenceAsset, + chunksByRef: ReadonlyMap, +): string { + const chunk = chunksByRef.get(asset.chunkRef) + const documentId = getTrimmedValue( + asset.source.documentId ?? chunk?.source.documentId, + ) + const revisionKey = getTrimmedValue(asset.revisionKey ?? chunk?.revisionKey) + const sourcePath = getTrimmedValue(asset.sourcePath)?.toLowerCase() + + if (documentId && revisionKey && sourcePath) { + return `document:${documentId}\u0000${revisionKey}\u0000${sourcePath}` + } + + const assetUrl = getNormalizedAssetUrl(asset.assetUrl) + if (assetUrl) return `url:${assetUrl}` + + if (documentId && sourcePath) { + return `document:${documentId}\u0000${sourcePath}` + } + + return asset.ref +} + +function getNormalizedAssetUrl(value: string | undefined): string | null { + const trimmed = getTrimmedValue(value) + if (!trimmed) return null + + try { + const url = new URL(trimmed) + url.hash = "" + return url.toString() + } catch { + return trimmed + } +} + +function getTrimmedValue( + value: string | null | undefined, +): string | null { + const trimmed = value?.trim() + return trimmed ? trimmed : null +} diff --git a/src/agent-harness/ledger.test.ts b/src/agent-harness/ledger.test.ts index 90cada3..3af1275 100644 --- a/src/agent-harness/ledger.test.ts +++ b/src/agent-harness/ledger.test.ts @@ -87,10 +87,13 @@ describe("createEvidenceLedger", () => { score: 0.9, metadata: { page_nums: [4], + pageAssets: [], page_assets: [ { page_num: 4, + artifactRef: "", artifact_ref: "page_citation_assets/page-4.png", + contentType: "", content_type: "image/png", }, ], diff --git a/src/agent-harness/ledger.ts b/src/agent-harness/ledger.ts index 64700e1..42f3120 100644 --- a/src/agent-harness/ledger.ts +++ b/src/agent-harness/ledger.ts @@ -381,9 +381,10 @@ function getPageCitationAssetCandidate( ): EvidenceAssetCandidate | null { if (chunk.chunkType.toLowerCase() !== "page") return null - const candidates = parsePageCitationAssetCandidates( - chunk.metadata?.pageAssets ?? chunk.metadata?.page_assets, - ).filter(isSupportedPageCitationAsset) + const candidates = [ + ...parsePageCitationAssetCandidates(chunk.metadata?.pageAssets), + ...parsePageCitationAssetCandidates(chunk.metadata?.page_assets), + ].filter(isSupportedPageCitationAsset) if (candidates.length === 0) return null const pageNumbers = getPageNumbers(chunk.metadata) @@ -469,30 +470,21 @@ function parsePageCitationAssetCandidates( getPositiveInteger(item.page_num) ?? getPositiveInteger(item.pageNumber) if (!pageNum) return [] + const artifactRef = + getTrimmedString(item.artifactRef) ?? + getTrimmedString(item.artifact_ref) + const assetUrl = + getTrimmedString(item.assetUrl) ?? getTrimmedString(item.asset_url) + const contentType = + getTrimmedString(item.contentType) ?? + getTrimmedString(item.content_type) return [ { pageNum, - ...(getTrimmedString(item.artifactRef ?? item.artifact_ref) - ? { - artifactRef: - getTrimmedString(item.artifactRef ?? item.artifact_ref) ?? - undefined, - } - : {}), - ...(getTrimmedString(item.assetUrl ?? item.asset_url) - ? { - assetUrl: - getTrimmedString(item.assetUrl ?? item.asset_url) ?? undefined, - } - : {}), - ...(getTrimmedString(item.contentType ?? item.content_type) - ? { - contentType: - getTrimmedString(item.contentType ?? item.content_type) ?? - undefined, - } - : {}), + ...(artifactRef ? { artifactRef } : {}), + ...(assetUrl ? { assetUrl } : {}), + ...(contentType ? { contentType } : {}), }, ] }) diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index 5be40e4..0a5ca19 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -13,6 +13,7 @@ import type { AgentTurnInput, ContextPolicy, HarnessToolCallTrace, + ImageInspectionRequest, IntentFrame, KnowhereToolRuntime, OutputManifest, @@ -455,6 +456,86 @@ describe("agent harness runtime", () => { ).toMatchObject({ ok: true }) }) + it("does not deduplicate unrelated documents with generic page paths", async () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse({ + namespace: "notebook", + query: "revenue", + routerUsed: "mapnav", + answerText: null, + evidenceText: "Revenue evidence", + stopReason: "completed", + failureReason: null, + results: [ + { + chunkId: "chunk_page_4", + content: "Revenue evidence from document A.", + chunkType: "page", + score: 0.9, + metadata: { + pageNums: [4], + pageAssets: [ + { + pageNum: 4, + artifactRef: "page_citation_assets/page-4.png", + assetUrl: "https://assets.example/doc-a/page-4.png", + contentType: "image/png", + }, + ], + }, + source: { + documentId: "doc_a", + sourceFileName: "a.pdf", + sectionPath: "Page 4", + }, + }, + { + chunkId: "chunk_page_4", + content: "Revenue evidence from document B.", + chunkType: "page", + score: 0.8, + metadata: { + pageNums: [4], + pageAssets: [ + { + pageNum: 4, + artifactRef: "page_citation_assets/page-4.png", + assetUrl: "https://assets.example/doc-b/page-4.png", + contentType: "image/png", + }, + ], + }, + source: { + documentId: "doc_b", + sourceFileName: "b.pdf", + sectionPath: "Page 4", + }, + }, + ], + referencedChunks: [], + }) + const inspectImages = vi.fn(async (request: ImageInspectionRequest) => ({ + analysis: "Inspected both pages.", + inspected: request.assets.map(({ ref, label }) => ({ ref, label })), + skipped: [], + })) + const tools = createHarnessTools({ + state: {}, + ledger, + knowhereTools: makeKnowhereTools(), + inspectImages, + recentTurns: [], + }) + + expect( + await executeTool(tools.inspectImage, { + refs: ["asset:r1:result:1", "asset:r1:result:2"], + question: "Compare the cited revenue.", + }), + ).toMatchObject({ ok: true }) + expect(inspectImages.mock.calls[0]?.[0].assets).toHaveLength(2) + }) + it("does not treat skipped image assets as successfully inspected", async () => { const ledger = createEvidenceLedger() ledger.addRetrievalResponse(makePageCitationRetrievalResponse()) @@ -871,7 +952,7 @@ describe("agent harness runtime", () => { ]) }) - it("keeps forcing inspectImage at the finalization step while page assets are uninspected", () => { + it("reserves the finalization step even when unused page assets remain", () => { const result = prepareHarnessStep({ stepNumber: 13, hasUninspectedImageAssets: true, @@ -883,10 +964,10 @@ describe("agent harness runtime", () => { ], }) - expect(result.activeTools).toEqual(["inspectImage"]) + expect(result.activeTools).toEqual(["finalize"]) expect(result.toolChoice).toEqual({ type: "tool", - toolName: "inspectImage", + toolName: "finalize", }) }) diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index 38cc9fc..073051d 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -8,6 +8,7 @@ import { import { z } from "zod" import { createEvidenceLedger } from "./ledger" +import { getCanonicalImageAssetKey } from "./image-asset-identity" import { knowhereToolText } from "./knowhere-text" import { mergeImageInspectionHighlights } from "./image-highlights" import type { @@ -265,8 +266,7 @@ export function prepareHarnessStep(input: { const shouldForceImageInspection = input.hasUninspectedImageAssets === true && - input.stepNumber >= imageInspectionReminderStepNumber && - input.stepNumber <= forcedFinalizationStepNumber + input.stepNumber === imageInspectionReminderStepNumber if (shouldForceImageInspection) { return { @@ -894,17 +894,6 @@ function resolveImageAssetRef( return sibling?.assetRef ?? null } -function getCanonicalImageAssetKey( - asset: EvidenceAsset, - chunksByRef: ReadonlyMap, -): string { - const chunk = chunksByRef.get(asset.chunkRef) - const chunkId = chunk?.chunkId?.trim() - const sourcePath = asset.sourcePath?.trim().toLowerCase() - if (chunkId && sourcePath) return `${chunkId}\u0000${sourcePath}` - return asset.ref -} - function buildFinalizeRequiresInspectionMessage( inspectRefs: readonly string[], ): string { diff --git a/src/app/globals.css b/src/app/globals.css index 38bea15..d2b9a67 100644 --- a/src/app/globals.css +++ b/src/app/globals.css @@ -111,8 +111,8 @@ .citation-region-flash { z-index: 10; display: block; - border-radius: 8px; - background-color: rgb(142 81 255 / 0.25); + border-radius: var(--radius); + background-color: hsl(var(--ring) / 0.25); animation: citation-region-flash 4.6s ease-in-out forwards; } diff --git a/src/components/chat-citation-model.test.ts b/src/components/chat-citation-model.test.ts index f8e8709..b037bdf 100644 --- a/src/components/chat-citation-model.test.ts +++ b/src/components/chat-citation-model.test.ts @@ -122,6 +122,38 @@ describe("chatCitationModel", () => { ]), ).toEqual([{ x: 0.1, y: 0.2, w: 0.3, h: 0.15 }]) + expect( + chatCitationModel.getListHighlightRegions(citation, [ + { + type: "image", + display: true, + highlightRegions: [{ x: 0.7, y: 0.1, w: 0.2, h: 0.1 }], + citation: { + chunkType: "page", + score: 0.7, + pageCitationPageNumber: 4, + source: { documentId: "doc_2" }, + }, + }, + ]), + ).toEqual([]) + + expect( + chatCitationModel.getListHighlightRegions(citation, [ + { + type: "image", + display: false, + highlightRegions: [{ x: 0.7, y: 0.1, w: 0.2, h: 0.1 }], + citation: { + chunkType: "page", + score: 0.7, + pageCitationPageNumber: 4, + source: { documentId: "doc_1" }, + }, + }, + ]), + ).toEqual([]) + expect(chatCitationModel.getListHighlightRegions(citation, [])).toEqual([]) }) diff --git a/src/components/chat-citation-model.ts b/src/components/chat-citation-model.ts index daa9c6f..31181d5 100644 --- a/src/components/chat-citation-model.ts +++ b/src/components/chat-citation-model.ts @@ -155,15 +155,27 @@ function getListHighlightRegions( } const pageNumber = workspaceCitationState.getCitationPageNumber(citation) + const documentId = getTrimmedField(citation.source.documentId) + if (pageNumber === null || !documentId) return [] + const matchingRegions = (artifacts ?? []).flatMap((artifact) => { - if (!artifact.highlightRegions || artifact.highlightRegions.length === 0) { + if ( + artifact.type !== "image" || + artifact.display === false || + !artifact.citation || + !artifact.highlightRegions || + artifact.highlightRegions.length === 0 + ) { return [] } - const artifactPage = artifact.citation - ? workspaceCitationState.getCitationPageNumber(artifact.citation) - : null - if (pageNumber !== null && artifactPage !== pageNumber) { + const artifactDocumentId = getTrimmedField( + artifact.citation.source.documentId, + ) + const artifactPage = workspaceCitationState.getCitationPageNumber( + artifact.citation, + ) + if (artifactDocumentId !== documentId || artifactPage !== pageNumber) { return [] } diff --git a/src/components/citation-region-highlight.tsx b/src/components/citation-region-highlight.tsx index 6eecf0e..a5b56a2 100644 --- a/src/components/citation-region-highlight.tsx +++ b/src/components/citation-region-highlight.tsx @@ -21,15 +21,12 @@ export function CitationRegionHighlight({ ))} diff --git a/src/components/parsed-chunk-card.test.ts b/src/components/parsed-chunk-card.test.ts index bcdb31e..a4fa48d 100644 --- a/src/components/parsed-chunk-card.test.ts +++ b/src/components/parsed-chunk-card.test.ts @@ -151,7 +151,9 @@ describe("ParsedChunkCard", () => { '[data-citation-page="4"]', ); expect(citationTarget).not.toBeNull(); - expect(citationTarget?.className).toContain("relative"); + expect(screen.getByTestId("citation-image-stage").className).toContain( + "relative", + ); expect(citationTarget?.getAttribute("data-focused-citation-id")).toBe( "assistant_1:0", ); @@ -189,6 +191,12 @@ describe("ParsedChunkCard", () => { expect(highlight.getAttribute("style")).toContain("20%"); expect(highlight.getAttribute("style")).toContain("30%"); expect(highlight.getAttribute("style")).toContain("15%"); + expect( + screen.getByTestId("citation-region-highlights").parentElement, + ).toBe(screen.getByTestId("citation-image-stage")); + + fireEvent.error(screen.getByRole("img", { name: "Page 4" })); + expect(screen.queryByTestId("citation-region-highlight")).toBeNull(); }); it("hides the keywords row when a page card has none", () => { diff --git a/src/components/parsed-chunk-card.tsx b/src/components/parsed-chunk-card.tsx index 9b5c1a3..611ebad 100644 --- a/src/components/parsed-chunk-card.tsx +++ b/src/components/parsed-chunk-card.tsx @@ -502,8 +502,6 @@ function PageCitationAssetImage({ const [failedAssetUrl, setFailedAssetUrl] = useState(null); const imageAssetUrl = getInlineImageAssetUrl(asset.assetUrl); const hasImageError = failedAssetUrl === imageAssetUrl; - const aspectRatio = - asset.width && asset.height ? `${asset.width} / ${asset.height}` : undefined; return (
@@ -512,29 +510,35 @@ function PageCitationAssetImage({ {asset.contentType}
{hasImageError ? ( ) : ( - // eslint-disable-next-line @next/next/no-img-element -- Page assets can be short-lived Knowhere URLs outside Next image optimization. - {`Page setFailedAssetUrl(imageAssetUrl)} - /> +
+ {/* eslint-disable-next-line @next/next/no-img-element -- Page assets can be short-lived Knowhere URLs outside Next image optimization. */} + {`Page setFailedAssetUrl(imageAssetUrl)} + /> + {isCitationFocus ? ( + + ) : null} +
)} - {isCitationFocus ? ( - - ) : null}
); @@ -593,19 +597,24 @@ function ImageChunkCard({ {inlineImageAssetUrl ? (
-
- {/* eslint-disable-next-line @next/next/no-img-element -- Parsed artifact dimensions are not known before render. */} - {chunk.summary - {isFocused ? ( - +
+ {/* eslint-disable-next-line @next/next/no-img-element -- Parsed artifact dimensions are not known before render. */} + {chunk.summary - ) : null} + {isFocused ? ( + + ) : null} +
) : ( diff --git a/src/domains/chat/image-inspection-model.test.ts b/src/domains/chat/image-inspection-model.test.ts index 5e123a7..6181c6d 100644 --- a/src/domains/chat/image-inspection-model.test.ts +++ b/src/domains/chat/image-inspection-model.test.ts @@ -69,6 +69,51 @@ describe("image inspection model", () => { expect(generateText).not.toHaveBeenCalled() }) + it("batches more than six images while preserving every page region", async () => { + const sevenAssets = Array.from({ length: 7 }, (_, index) => ({ + ref: `asset:page-${index + 1}`, + label: `page ${index + 1}`, + body: new Uint8Array([index + 1]), + contentType: "image/png", + })) + vi.mocked(generateObject) + .mockResolvedValueOnce({ + object: { + analysis: "Pages one through six.", + pages: sevenAssets.slice(0, 6).map((asset) => ({ + ref: asset.ref, + regions: [{ x: 0.1, y: 0.2, w: 0.3, h: 0.1 }], + })), + }, + } as Awaited>) + .mockResolvedValueOnce({ + object: { + analysis: "Page seven.", + pages: [ + { + ref: sevenAssets[6]!.ref, + regions: [{ x: 0.4, y: 0.5, w: 0.2, h: 0.1 }], + }, + ], + }, + } as Awaited>) + + const result = await generateImageInspectionModelResult({ + workspaceId: "ws_1", + question: "Locate evidence on every cited page.", + assets: sevenAssets, + }) + + expect(generateObject).toHaveBeenCalledTimes(2) + expect(result.analysis).toBe("Pages one through six.\n\nPage seven.") + expect(result.pages).toHaveLength(7) + expect(result.pages[6]).toEqual({ + ref: "asset:page-7", + regions: [{ x: 0.4, y: 0.5, w: 0.2, h: 0.1 }], + }) + expect(result.source).toBe("structured") + }) + it("salvages valid structured JSON from failed generateObject text", async () => { vi.mocked(generateObject).mockRejectedValue( makeNoObjectGeneratedError( diff --git a/src/domains/chat/image-inspection-model.ts b/src/domains/chat/image-inspection-model.ts index 8589dda..8a498af 100644 --- a/src/domains/chat/image-inspection-model.ts +++ b/src/domains/chat/image-inspection-model.ts @@ -12,6 +12,7 @@ import { summarizeUnknownError } from "@/lib/format-log-value" import { logger } from "@/lib/logger" const VISION_MODEL = process.env.VISION_MODEL ?? CHAT_MODEL +const IMAGE_INSPECTION_BATCH_SIZE = 6 export const imageInspectionResultSchema = z.object({ analysis: z.string(), @@ -64,6 +65,39 @@ export async function generateImageInspectionModelResult(input: { readonly workspaceId: string readonly question: string readonly assets: readonly ImageInspectionModelAsset[] +}): Promise { + if (input.assets.length <= IMAGE_INSPECTION_BATCH_SIZE) { + return generateImageInspectionBatchResult(input) + } + + const results: ImageInspectionModelResult[] = [] + for ( + let index = 0; + index < input.assets.length; + index += IMAGE_INSPECTION_BATCH_SIZE + ) { + results.push( + await generateImageInspectionBatchResult({ + ...input, + assets: input.assets.slice(index, index + IMAGE_INSPECTION_BATCH_SIZE), + }), + ) + } + + return { + analysis: results + .map((result) => result.analysis.trim()) + .filter(Boolean) + .join("\n\n"), + pages: results.flatMap((result) => result.pages), + source: getCombinedInspectionSource(results), + } +} + +async function generateImageInspectionBatchResult(input: { + readonly workspaceId: string + readonly question: string + readonly assets: readonly ImageInspectionModelAsset[] }): Promise { let structuredFailureText: string | undefined @@ -226,6 +260,18 @@ export async function generateImageInspectionModelResult(input: { } } +function getCombinedInspectionSource( + results: readonly ImageInspectionModelResult[], +): ImageInspectionModelResult["source"] { + if (results.every((result) => result.source === "structured")) { + return "structured" + } + if (results.some((result) => result.source === "analysis_fallback")) { + return "analysis_fallback" + } + return "structured_text" +} + export function isRecoverableStructuredOutputError(error: unknown): boolean { if (NoObjectGeneratedError.isInstance(error)) return true if (UnsupportedFunctionalityError.isInstance(error)) return true diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 9efcbb5..003e436 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -1518,6 +1518,7 @@ describe("answerQuestionWithRetrieval", () => { ref: "asset:r1:result:1", chunkRef: "r1:result:1", type: "image" as const, + assetUrl: "https://assets.example/tsla/page-4.png", sourcePath: "page_citation_assets/page-4.png", source: catalogResult.source, label: "original.pdf / Page 4", @@ -1526,6 +1527,7 @@ describe("answerQuestionWithRetrieval", () => { ref: "asset:r1:result:2", chunkRef: "r1:result:2", type: "image" as const, + assetUrl: "https://assets.example/tsla/page-4.png", sourcePath: "page_citation_assets/page-4.png", source: workspaceResult.source, label: "TSLA-Q4-2025-Update.pdf / Page 4", diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index 9fc711c..446b591 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -6,6 +6,7 @@ import type { } from "@ontos-ai/knowhere-sdk" import { logger } from "@/lib/logger" +import { getCanonicalImageAssetKey } from "@/agent-harness/image-asset-identity" import type { ChatArtifactView, ChatCitationView, @@ -995,11 +996,7 @@ function getCanonicalCitationAssetKey( ): string { const asset = assetsByRef.get(assetRef) if (!asset) return assetRef - const chunk = chunksByRef.get(asset.chunkRef) - const chunkId = chunk?.chunkId?.trim() - const sourcePath = asset.sourcePath?.trim().toLowerCase() - if (chunkId && sourcePath) return `${chunkId}\u0000${sourcePath}` - return assetRef + return getCanonicalImageAssetKey(asset, chunksByRef) } function mapDisplayedManifestArtifactsToResults( diff --git a/src/domains/chat/page-citation-assets.test.ts b/src/domains/chat/page-citation-assets.test.ts index 5b6c2b2..5b16264 100644 --- a/src/domains/chat/page-citation-assets.test.ts +++ b/src/domains/chat/page-citation-assets.test.ts @@ -148,13 +148,17 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { expect(result?.pageCitationPageNumber).toBe(25) }) - it("reads snake_case page_assets for the cited page number", async () => { + it("reads snake_case page_assets when camel-case metadata is empty", async () => { + const hardenChatAssetUrl = vi + .fn() + .mockResolvedValue("https://blob.example/page_citation_assets/page-11.png") const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ results: [ makeRetrievalResult({ chunkType: "page", metadata: { page_nums: [11], + pageAssets: [], page_assets: [ { page_num: 11, @@ -165,9 +169,15 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { }), ], sources: [makeSource()], + hardenChatAssetUrl, }) expect(result?.pageCitationPageNumber).toBe(11) + expect(hardenChatAssetUrl).toHaveBeenCalledWith( + expect.objectContaining({ + sourcePath: "page_citation_assets/page-11.png", + }), + ) }) it("attaches page numbers to text citations without turning them into page images", async () => { diff --git a/src/domains/chat/page-citation-assets.ts b/src/domains/chat/page-citation-assets.ts index d411b68..3007f60 100644 --- a/src/domains/chat/page-citation-assets.ts +++ b/src/domains/chat/page-citation-assets.ts @@ -129,9 +129,10 @@ function getDirectPageCitationAsset( result: RetrievalResult, pageNumbers: readonly number[], ): PageCitationAssetCandidate | null { - const candidates = parsePageCitationAssetCandidates( - result.metadata?.pageAssets ?? result.metadata?.page_assets, - ) + const candidates = [ + ...parsePageCitationAssetCandidates(result.metadata?.pageAssets), + ...parsePageCitationAssetCandidates(result.metadata?.page_assets), + ] if (pageNumbers.length > 0) { const matchingCandidates = candidates.filter((candidate) => From 20db3ecc047670ebf1ed27ad8425b6f71e49b2f3 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 24 Aug 2026 09:26:08 +0800 Subject: [PATCH 51/51] fix: recognize prefixed auth session cookies Allow Notebook's proxy to recognize staging and production Better Auth token names so authenticated inspect requests are not redirected to login. Co-authored-by: Cursor --- .env.local.example | 5 +-- src/infrastructure/auth/index.test.ts | 23 ++++++++++++- src/infrastructure/auth/index.ts | 12 +++---- .../auth/session-cookie-names.ts | 27 +++++++++++---- src/proxy.test.ts | 33 +++++++++++++++++++ src/proxy.ts | 6 ++-- 6 files changed, 87 insertions(+), 19 deletions(-) diff --git a/.env.local.example b/.env.local.example index 7340c1a..1307872 100644 --- a/.env.local.example +++ b/.env.local.example @@ -39,8 +39,9 @@ DASHBOARD_ORIGIN=http://dashboard.local.knowhereto.ai:3000 # local: http://notebook.local.knowhereto.ai:3001 NOTEBOOK_PUBLIC_URL=http://notebook.local.knowhereto.ai:3001 -# Optional override of the session cookie names Dashboard sets. Defaults -# to Better Auth's standard cookie names. +# Optional additional session cookie names Dashboard sets. +# Without an override, Notebook recognizes Better Auth session-token cookies +# with standard or environment-prefixed names. # SESSION_COOKIE_NAMES=better-auth.session_token,__Secure-better-auth.session_token # --- Product analytics (client-side) --- diff --git a/src/infrastructure/auth/index.test.ts b/src/infrastructure/auth/index.test.ts index 781c80c..ad2e22b 100644 --- a/src/infrastructure/auth/index.test.ts +++ b/src/infrastructure/auth/index.test.ts @@ -24,7 +24,7 @@ vi.mock("next/cache", () => nextCacheMocks) * Playwright flow added in a later PR. */ -import { extractUser, sessionCookieNames } from "." +import { extractUser, isSessionCookieName, sessionCookieNames } from "." const SESSION_PATH = "/api/orpc/users/getCurrentUser" @@ -147,6 +147,27 @@ describe("sessionCookieNames", () => { "__Secure-better-auth.session_token", ]) }) + + it("recognizes environment-prefixed Better Auth session tokens", () => { + delete process.env.SESSION_COOKIE_NAMES + + expect( + isSessionCookieName("__Secure-better-auth-staging-session_token"), + ).toBe(true) + expect(isSessionCookieName("__Secure-better-auth-session_token")).toBe(true) + expect( + isSessionCookieName("__Secure-better-auth-staging-session_data"), + ).toBe(false) + }) + + it("adds an explicit cookie-name override without hiding Better Auth tokens", () => { + process.env.SESSION_COOKIE_NAMES = "custom-session" + + expect(isSessionCookieName("custom-session")).toBe(true) + expect( + isSessionCookieName("__Secure-better-auth-staging-session_token"), + ).toBe(true) + }) }) describe("getCurrentUser", () => { diff --git a/src/infrastructure/auth/index.ts b/src/infrastructure/auth/index.ts index b28ac70..afc2b2f 100644 --- a/src/infrastructure/auth/index.ts +++ b/src/infrastructure/auth/index.ts @@ -9,13 +9,16 @@ import { HttpClientRequest, } from "@effect/platform" import { authURLs } from "./urls" -import { sessionCookieNames } from "./session-cookie-names" +import { + isSessionCookieName, + sessionCookieNames, +} from "./session-cookie-names" import { logger } from "@/lib/logger" import { knowhereApiKeyOverride } from "@/integrations/knowhere-api-key" import { setEmptyJsonBody } from "@/integrations/dashboard/orpc-request" import { formatUnknownForLog } from "@/lib/format-log-value" -export { sessionCookieNames } +export { isSessionCookieName, sessionCookieNames } /** * Auth helpers for Knowhere Notebook. @@ -222,10 +225,7 @@ export async function hasSessionCookie(): Promise { if (knowhereApiKeyOverride.hasApiKey()) return true const jar = await cookies() - for (const name of sessionCookieNames()) { - if (jar.get(name) !== undefined) return true - } - return false + return jar.getAll().some((cookie) => isSessionCookieName(cookie.name)) } /** diff --git a/src/infrastructure/auth/session-cookie-names.ts b/src/infrastructure/auth/session-cookie-names.ts index 8b81aee..c0cbf83 100644 --- a/src/infrastructure/auth/session-cookie-names.ts +++ b/src/infrastructure/auth/session-cookie-names.ts @@ -3,13 +3,26 @@ const DEFAULT_SESSION_COOKIE_NAMES = [ "__Secure-better-auth.session_token", ] as const +const BETTER_AUTH_SESSION_COOKIE_NAME = + /^(?:__Secure-)?better-auth(?:[-.][A-Za-z0-9_]+)*[-.]session_token$/u + export function sessionCookieNames(): readonly string[] { + return configuredSessionCookieNames() ?? DEFAULT_SESSION_COOKIE_NAMES +} + +export function isSessionCookieName(name: string): boolean { + return ( + sessionCookieNames().includes(name) || + BETTER_AUTH_SESSION_COOKIE_NAME.test(name) + ) +} + +function configuredSessionCookieNames(): readonly string[] | null { const override = process.env.SESSION_COOKIE_NAMES - if (override !== undefined && override.trim().length > 0) { - return override - .split(",") - .map((s) => s.trim()) - .filter(Boolean) - } - return DEFAULT_SESSION_COOKIE_NAMES + if (override === undefined || override.trim().length === 0) return null + + return override + .split(",") + .map((name) => name.trim()) + .filter(Boolean) } diff --git a/src/proxy.test.ts b/src/proxy.test.ts index 0d80c19..d6b1f55 100644 --- a/src/proxy.test.ts +++ b/src/proxy.test.ts @@ -6,10 +6,12 @@ import { proxy } from "./proxy"; describe("proxy", () => { const originalDashboardOrigin = process.env.DASHBOARD_ORIGIN; const originalKnowhereApiKey = process.env.KNOWHERE_API_KEY; + const originalSessionCookieNames = process.env.SESSION_COOKIE_NAMES; beforeEach(() => { delete process.env.DASHBOARD_ORIGIN; delete process.env.KNOWHERE_API_KEY; + delete process.env.SESSION_COOKIE_NAMES; }); afterEach(() => { @@ -23,6 +25,11 @@ describe("proxy", () => { } else { process.env.KNOWHERE_API_KEY = originalKnowhereApiKey; } + if (originalSessionCookieNames === undefined) { + delete process.env.SESSION_COOKIE_NAMES; + } else { + process.env.SESSION_COOKIE_NAMES = originalSessionCookieNames; + } }); it("allows anonymous guest source reads", () => { @@ -84,4 +91,30 @@ describe("proxy", () => { expect(response.headers.get("x-middleware-next")).toBe("1"); }); + + it("recognizes Dashboard session tokens with an environment prefix", () => { + const response = proxy( + new NextRequest("http://localhost:3001/inspect/doc-1/chunks", { + headers: { + cookie: "__Secure-better-auth-staging-session_token=token", + }, + }), + ); + + expect(response.headers.get("x-middleware-next")).toBe("1"); + }); + + it("does not mistake Better Auth session data for a session token", () => { + const response = proxy( + new NextRequest("http://localhost:3001/inspect/doc-1/chunks", { + headers: { + cookie: "__Secure-better-auth-staging-session_data=data", + }, + }), + ); + + expect(response.headers.get("location")).toBe( + "http://localhost:3001/login", + ); + }); }); diff --git a/src/proxy.ts b/src/proxy.ts index 7527048..6c10573 100644 --- a/src/proxy.ts +++ b/src/proxy.ts @@ -1,6 +1,6 @@ import { NextResponse, type NextRequest } from "next/server" import { authURLs } from "@/infrastructure/auth/urls" -import { sessionCookieNames } from "@/infrastructure/auth/session-cookie-names" +import { isSessionCookieName } from "@/infrastructure/auth/session-cookie-names" import { knowhereApiKeyOverride } from "@/integrations/knowhere-api-key" /** @@ -59,8 +59,8 @@ export function proxy(req: NextRequest): NextResponse { if (isPublicPath(req)) return NextResponse.next() - for (const name of sessionCookieNames()) { - if (req.cookies.get(name)) return NextResponse.next() + if (req.cookies.getAll().some((cookie) => isSessionCookieName(cookie.name))) { + return NextResponse.next() } const origin = process.env.DASHBOARD_ORIGIN