Repository navigation
Use Workers AI for vector embeddings - #613
Conversation
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
📝 WalkthroughWalkthroughThis PR introduces an optional Workers AI binding for capability embeddings with Cloudflare AI Gateway routing and offline fallback, expands capability maintenance reindexing to cover memories, jobs, and saved packages via a new ChangesWorkers AI embeddings and multi-vector reindex
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant handleCapabilityReindexRequest
participant reindexAllCapabilitySearchVectors
participant embedTextsForVectorize
participant WorkersAI as "runtime.AI"
Client->>handleCapabilityReindexRequest: POST /__maintenance/reindex-capabilities
handleCapabilityReindexRequest->>reindexAllCapabilitySearchVectors: run()
reindexAllCapabilitySearchVectors->>embedTextsForVectorize: embed capability/memory/job/package text
alt AI binding available
embedTextsForVectorize->>WorkersAI: run(CAPABILITY_EMBEDDING_MODEL)
WorkersAI-->>embedTextsForVectorize: embedding rows
else offline/non-production
embedTextsForVectorize->>embedTextsForVectorize: deterministicEmbedding fallback
end
embedTextsForVectorize-->>reindexAllCapabilitySearchVectors: embeddings
reindexAllCapabilitySearchVectors-->>handleCapabilityReindexRequest: aggregated {upserted, error?} per kind
handleCapabilityReindexRequest-->>Client: 200 {ok:true, ...results} or 500 with error summary
Related PRs: None identified. Suggested labels: enhancement, documentation, worker Suggested reviewers: kentcdodds 🐰 A carrot-sized model now hums, 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
|
🔎 Preview deployed: https://kody-pr-613.kentcdodds.workers.dev Worker: Mocks:
|
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (5)
packages/worker/src/package-registry/package-reindex.ts (1)
30-41: 🧹 Nitpick | 🔵 TrivialPossible subrequest-limit exposure on large batches.
Each package in a batch triggers a manifest fetch via
loadPackageManifestBySourceId, likely afetch()againstbaseUrl. With enough saved packages, cumulative subrequests across all batches (plus the other reindex steps run in the same request) could approach the free-plan external subrequest ceiling. Worth confirming the deployment's Workers plan tier and expected saved-package volume.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/worker/src/package-registry/package-reindex.ts` around lines 30 - 41, The reindex flow in package-reindex.ts may create too many external subrequests because each batch calls loadPackageManifestBySourceId in parallel and the overall request also performs other reindex work. Review the batch size and the PackageReindex path around the loop using upsertBatchSize and Promise.all, then add a safeguard by reducing concurrency, lowering batch size, or splitting work so the total fetch volume stays within the Workers plan’s subrequest limits.packages/worker/src/capability-maintenance.ts (2)
27-45: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winConsider running independent reindex steps concurrently.
runReindexStepalready isolates failures per step (it never rejects), so the fourawaitcalls don't need to be sequential. Since this endpoint is invoked from deploy/preview automation per the PR description, reducing total wall-clock latency is worthwhile.- const capabilities = await runReindexStep(() => - reindexCapabilityVectors(env, capabilitySpecs), - ) - const memories = await runReindexStep(() => reindexMemoryVectors(env)) - const jobs = await runReindexStep(() => reindexJobVectors(env)) - const packages = await runReindexStep(() => - reindexSavedPackageVectors(env, { baseUrl: input.baseUrl }), - ) - return { - capabilities, - memories, - jobs, - packages, - } + const [capabilities, memories, jobs, packages] = await Promise.all([ + runReindexStep(() => reindexCapabilityVectors(env, capabilitySpecs)), + runReindexStep(() => reindexMemoryVectors(env)), + runReindexStep(() => reindexJobVectors(env)), + runReindexStep(() => + reindexSavedPackageVectors(env, { baseUrl: input.baseUrl }), + ), + ]) + return { capabilities, memories, jobs, packages }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/worker/src/capability-maintenance.ts` around lines 27 - 45, `reindexAllCapabilitySearchVectors` is running four independent `runReindexStep` calls sequentially even though each step is failure-isolated and never rejects. Update the function to start the `reindexCapabilityVectors`, `reindexMemoryVectors`, `reindexJobVectors`, and `reindexSavedPackageVectors` steps concurrently (while still passing `env` and `input.baseUrl` as needed), then wait for all of them before returning the `{ capabilities, memories, jobs, packages }` result.
55-73: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winPartial success detail is discarded on failure.
When any step fails, the response body only contains
{ ok: false, error }— the counts from steps that did succeed (e.g.,capabilities: {upserted: 3}) are lost. For an endpoint invoked from deploy automation, surfacing partial progress alongside the error would ease triage.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/worker/src/capability-maintenance.ts` around lines 55 - 73, The `run` handler in `capability-maintenance.ts` throws away successful step results when any reindex step fails. Update the `run` logic so it preserves and returns the partial `result` alongside the failure information instead of only surfacing an error string; keep using `reindexAllCapabilitySearchVectors` and the `failed` detection, but have the response include both the error and the successful step payloads (for example, the `capabilities` counts) so deploy automation can see partial progress.packages/worker/src/package-registry/package-reindex.node.test.ts (1)
39-101: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing coverage for batching and manifest-load failures.
Only a single-row, single-batch success case is tested. Consider adding cases for: multiple batches (>16 rows), the offline/missing-index early returns, and a rejected
loadPackageManifestBySourceIdcall — the last would have caught the failure-propagation issue flagged inpackage-reindex.ts.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/worker/src/package-registry/package-reindex.node.test.ts` around lines 39 - 101, Add tests in reindexSavedPackageVectors to cover the missing control-flow paths: verify the early return cases when isCapabilitySearchOffline() is true and when getCapabilityVectorIndex() yields no index, then add a multi-row fixture that forces batching beyond 16 items and asserts multiple upsert calls. Also add a failure case where loadPackageManifestBySourceId rejects so the test confirms the error is propagated rather than swallowed, using the existing reindexSavedPackageVectors and mockModule helpers to locate the behavior.packages/worker/src/maintenance-handler.ts (1)
3-3: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winGuard against a future
run()result overriding theokfield.
MaintenanceResultis nowRecord<string, unknown>, and{ ok: true, ...result }lets anyokkey inresultsilently override the success flag. Not an issue with today's only caller, but nothing prevents a future maintenance handler from returningokand corrupting the response contract shared across allhandleSecretMaintenanceRequestconsumers.-type MaintenanceResult = Record<string, unknown> +type MaintenanceResult = Record<string, unknown> & { ok?: never }Also applies to: 36-37
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/worker/src/maintenance-handler.ts` at line 3, The maintenance response shape in MaintenanceResult can be corrupted because handleSecretMaintenanceRequest spreads the run() result after setting ok: true, allowing a future result.ok to override the success flag. Update the response assembly in handleSecretMaintenanceRequest so the contract-owned ok field cannot be overwritten, and keep the safeguard local to the code path that merges the maintenance handler result.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/worker/src/package-registry/package-reindex.ts`:
- Around line 32-56: The saved-package reindex flow in packageReindex currently
uses Promise.all over loadPackageManifestBySourceId, so one failed manifest load
aborts the whole batch and stops later packages from being indexed. Update the
manifest-loading step to handle failures per row inside the batch loop, using
loadPackageManifestBySourceId and the subsequent
buildSavedPackageEmbedText/embedTextsForVectorize flow only for successfully
loaded manifests, while skipping and continuing past failed rows so the rest of
the batch and later batches still run.
---
Nitpick comments:
In `@packages/worker/src/capability-maintenance.ts`:
- Around line 27-45: `reindexAllCapabilitySearchVectors` is running four
independent `runReindexStep` calls sequentially even though each step is
failure-isolated and never rejects. Update the function to start the
`reindexCapabilityVectors`, `reindexMemoryVectors`, `reindexJobVectors`, and
`reindexSavedPackageVectors` steps concurrently (while still passing `env` and
`input.baseUrl` as needed), then wait for all of them before returning the `{
capabilities, memories, jobs, packages }` result.
- Around line 55-73: The `run` handler in `capability-maintenance.ts` throws
away successful step results when any reindex step fails. Update the `run` logic
so it preserves and returns the partial `result` alongside the failure
information instead of only surfacing an error string; keep using
`reindexAllCapabilitySearchVectors` and the `failed` detection, but have the
response include both the error and the successful step payloads (for example,
the `capabilities` counts) so deploy automation can see partial progress.
In `@packages/worker/src/maintenance-handler.ts`:
- Line 3: The maintenance response shape in MaintenanceResult can be corrupted
because handleSecretMaintenanceRequest spreads the run() result after setting
ok: true, allowing a future result.ok to override the success flag. Update the
response assembly in handleSecretMaintenanceRequest so the contract-owned ok
field cannot be overwritten, and keep the safeguard local to the code path that
merges the maintenance handler result.
In `@packages/worker/src/package-registry/package-reindex.node.test.ts`:
- Around line 39-101: Add tests in reindexSavedPackageVectors to cover the
missing control-flow paths: verify the early return cases when
isCapabilitySearchOffline() is true and when getCapabilityVectorIndex() yields
no index, then add a multi-row fixture that forces batching beyond 16 items and
asserts multiple upsert calls. Also add a failure case where
loadPackageManifestBySourceId rejects so the test confirms the error is
propagated rather than swallowed, using the existing reindexSavedPackageVectors
and mockModule helpers to locate the behavior.
In `@packages/worker/src/package-registry/package-reindex.ts`:
- Around line 30-41: The reindex flow in package-reindex.ts may create too many
external subrequests because each batch calls loadPackageManifestBySourceId in
parallel and the overall request also performs other reindex work. Review the
batch size and the PackageReindex path around the loop using upsertBatchSize and
Promise.all, then add a safeguard by reducing concurrency, lowering batch size,
or splitting work so the total fetch volume stays within the Workers plan’s
subrequest limits.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: e2886509-a461-4aab-a5db-ab0e218c8620
📒 Files selected for processing (17)
.github/workflows/deploy.yml.github/workflows/preview.ymldocs/contributing/adding-capabilities.mddocs/contributing/environment-variables.mddocs/contributing/setup-manifest.mddocs/contributing/setup.mdpackages/worker/.env.examplepackages/worker/src/capability-maintenance.node.test.tspackages/worker/src/capability-maintenance.tspackages/worker/src/env-schema.tspackages/worker/src/maintenance-handler.tspackages/worker/src/mcp/capabilities/capability-search.tspackages/worker/src/mcp/capabilities/capability-search.workers.test.tspackages/worker/src/package-registry/package-reindex.node.test.tspackages/worker/src/package-registry/package-reindex.tspackages/worker/worker-configuration.d.tspackages/worker/wrangler.jsonc
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit cffb6da. Configure here.
| upserted: 0, | ||
| error: getErrorMessage(error), | ||
| } | ||
| } |
There was a problem hiding this comment.
Reindex errors hide partial upserts
Medium Severity
When a reindex step throws after earlier batches succeeded, runReindexStep always reports upserted: 0 for that kind. Vectorize may already contain those vectors, but the maintenance JSON and failure payload imply nothing was written for that step.
Reviewed by Cursor Bugbot for commit cffb6da. Configure here.
3 participants
Summary
@cf/baai/bge-small-en-v1.5(384 dimensions,clspooling), with strict response shape/dimension validation.AIbindings and optional AI Gateway routing viaAI_GATEWAY_ID; local/test and non-production no-AI paths keep deterministic offline fallback behavior.CAPABILITY_REINDEX_SECRET/__maintenance/reindex-capabilitiesoperation to rebuild built-in capabilities, memories, jobs, and saved packages, preservinguserIdmetadata filters for user-owned vectors.Operational reindex note
The existing Vectorize indexes remain
--dimensions=384 --metric=cosine, matching@cf/baai/bge-small-en-v1.5. After deploy, run (or let the production workflow run)POST /__maintenance/reindex-capabilitieswithAuthorization: Bearer $CAPABILITY_REINDEX_SECRETto rebuild all shared-index vectors: built-in capabilities, memories, jobs, and saved packages. This is required after the fake-to-real embedding switch so old pseudo-vectors are replaced with Workers AI vectors.Validation
npm run test -- --run packages/worker/src/mcp/capabilities/capability-search.workers.test.ts packages/worker/src/capability-maintenance.node.test.ts packages/worker/src/maintenance-handler.node.test.ts packages/worker/src/package-registry/package-reindex.node.test.ts✅npm run typecheck✅npm run format:check✅npm run lint✅ (existing warnings only)npm run validate✅ (run before initial push and after each reviewer-fix commit)System recap — extends existing primitives (medium risk)
Mode: recap · Base:
main@cf8c22e· Head:cffb6daClassification: extends — changes the embedding implementation and maintenance rebuild behavior for the existing Vectorize search primitive.
Primitives touched
vectorize-searchmcp-serverSystem map
Change flow
Invariants
metadata.userId, and existing search paths continue to filter byuserIdfor user-owned data.kind: builtin) and are not user content.Summary by CodeRabbit
New Features
Documentation