feat(rag): add encrypted embeddings foundation#211
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdds a user-embedding storage model, a cache-backed RAG search pipeline with ranking and filters, OpenAI embeddings request/billing refactoring, and cache eviction wiring on sensitive account changes. ChangesRAG Embedding Pipeline
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant search_user_embeddings
participant RagCache
participant Database
participant EmbeddingService
Client->>search_user_embeddings: query text + filters
search_user_embeddings->>EmbeddingService: embed query text
EmbeddingService-->>search_user_embeddings: query vector
search_user_embeddings->>RagCache: get(user_id)
alt cache miss
search_user_embeddings->>Database: load embeddings
Database-->>search_user_embeddings: vector rows
search_user_embeddings->>RagCache: put(user_id, embeddings)
end
search_user_embeddings->>search_user_embeddings: rank candidates and fetch content
search_user_embeddings-->>Client: RagSearchOutcome
sequenceDiagram
participant Client
participant proxy_embeddings
participant request_embeddings_with_billing
participant request_embeddings_from_primary
participant PrimaryProvider
Client->>proxy_embeddings: embeddings request
proxy_embeddings->>request_embeddings_with_billing: forward request
request_embeddings_with_billing->>request_embeddings_from_primary: canonicalized request
request_embeddings_from_primary->>PrimaryProvider: POST /v1/embeddings
PrimaryProvider-->>request_embeddings_from_primary: response_json + usage
request_embeddings_with_billing-->>proxy_embeddings: response_json + prompt_tokens
proxy_embeddings-->>Client: embeddings response
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
84e1c61 to
55ecf39
Compare
55ecf39 to
a111aa3
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/main.rs (1)
1706-1712: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftDelete stale encrypted embeddings on destructive password reset.
Evicting the cache is not enough here: the reset creates a new user seed/key, so persisted embeddings encrypted with the old key become undecryptable but still count toward storage and scans. Delete the user’s embedding rows as part of the destructive reset flow, or explicitly re-encrypt them if that is intended.
🤖 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 `@src/main.rs` around lines 1706 - 1712, The destructive password reset flow in complete_destructive_password_reset currently only evicts the RAG cache, which leaves persisted embeddings encrypted with the old user key behind. Update this path in the reset handling around self.db.complete_destructive_password_reset and self.rag_cache.lock().await.evict_user(...) to also remove the user’s stored embedding rows, or re-encrypt them if that is the intended behavior, so stale encrypted embeddings do not remain in storage.src/rag.rs (2)
624-638: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winMirror the id tie-break in bounded heap replacement.
Ordprefers lower ids after score/token ties, but the replacement check omits that final tie-break, so exact ties can keep the wrong candidate.Proposed fix
let ordering = score .total_cmp(&min.score) - .then_with(|| min.token_count.cmp(&e.token_count)); + .then_with(|| min.token_count.cmp(&e.token_count)) + .then_with(|| min.id.cmp(&e.id));🤖 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 `@src/rag.rs` around lines 624 - 638, The bounded heap replacement logic in the ranking loop is missing the final id tie-break that HeapItem::Ord uses, so exact score and token_count ties can preserve the wrong entry. Update the comparison in the heap replacement path that builds and compares HeapItem values to mirror the full ordering from Ord, including the id comparison after score and token_count, so the same candidate wins in both the heap and replacement check.
902-929: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winCheck cheap quota limits before generating embeddings.
Both insert paths call the external embedding/billing flow before enforcing insert-rate and row limits, so over-quota requests still incur embedding cost before being rejected. Split the cheap rate/row checks ahead of
embed_text_via_tinfoil, then keep the exact byte check before insert.Also applies to: 981-996
🤖 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 `@src/rag.rs` around lines 902 - 929, In the insert flow in rag.rs, the cheap quota checks are happening too late because embed_text_via_tinfoil is called before insert-rate and row limits are enforced, causing avoidable embedding/billing work for requests that will be rejected. Move the inexpensive quota validation ahead of embedding in the insert path, keeping the exact text-size and projected-byte checks where they are needed before the actual insert. Apply the same reordering to the other insert path mentioned in the review so both paths use the same early-check pattern around embed_text_via_tinfoil, ensure_storage_limits, and the insert validation logic.
♻️ Duplicate comments (1)
migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql (1)
1-3: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winMirror the
CONCURRENTLYfix for rollback drops.Same lock-blocking concern applies to
DROP INDEXhere — a normal drop takes anACCESS EXCLUSIVElock. This should be addressed consistently with the up.sql fix (same transaction-wrapping caveat applies).🔧 Proposed fix
-DROP INDEX IF EXISTS idx_user_embeddings_user_model_conversation_created_id; -DROP INDEX IF EXISTS idx_user_embeddings_user_model_source_created_id; -DROP INDEX IF EXISTS idx_user_embeddings_user_model_created_id; +DROP INDEX CONCURRENTLY IF EXISTS idx_user_embeddings_user_model_conversation_created_id; +DROP INDEX CONCURRENTLY IF EXISTS idx_user_embeddings_user_model_source_created_id; +DROP INDEX CONCURRENTLY IF EXISTS idx_user_embeddings_user_model_created_id;🤖 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 `@migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql` around lines 1 - 3, The rollback in down.sql drops the same indexes without the non-blocking form, so mirror the up.sql approach by changing the DROP INDEX statements for idx_user_embeddings_user_model_conversation_created_id, idx_user_embeddings_user_model_source_created_id, and idx_user_embeddings_user_model_created_id to use CONCURRENTLY. Keep in mind the migration must not be wrapped in a transaction when using the concurrent drop, so adjust the migration setup consistently with the existing up.sql fix.Source: Linters/SAST tools
🤖 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 `@migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/up.sql`:
- Around line 1-9: The three index builds in the user_embeddings hot path
migration are currently blocking because they use plain CREATE INDEX. Update the
migration’s CREATE INDEX statements for
idx_user_embeddings_user_model_created_id,
idx_user_embeddings_user_model_source_created_id, and
idx_user_embeddings_user_model_conversation_created_id to use CREATE INDEX
CONCURRENTLY, and add a metadata.toml for this migration with run_in_transaction
= false so Diesel can execute it outside a transaction.
In `@src/main.rs`:
- Line 24: The application router is missing the RAG routes, so the new web APIs
are not being exposed. Update the router setup in main.rs where the existing
route groups are mounted to import and merge the RAG route set alongside
instructions_routes, login_routes, oauth_routes, openai_routes, and
protected_routes, ensuring the new subsystem is reachable through the app’s
router configuration.
In `@src/rag.rs`:
- Around line 1360-1371: The duplicate-load path in
load_user_embeddings_for_search is leaving the in-flight marker behind when the
awaited load fails, which can cause later callers to keep timing out. Update the
branch around the await in src/rag.rs so the cache cleanup in
rag_cache.finish_load(user_id) (or equivalent in-flight removal) runs even when
load_user_embeddings_for_search returns an error, using a guard/defer-style
cleanup tied to the load_user_embeddings_for_search / cache.put flow.
---
Outside diff comments:
In `@src/main.rs`:
- Around line 1706-1712: The destructive password reset flow in
complete_destructive_password_reset currently only evicts the RAG cache, which
leaves persisted embeddings encrypted with the old user key behind. Update this
path in the reset handling around self.db.complete_destructive_password_reset
and self.rag_cache.lock().await.evict_user(...) to also remove the user’s stored
embedding rows, or re-encrypt them if that is the intended behavior, so stale
encrypted embeddings do not remain in storage.
In `@src/rag.rs`:
- Around line 624-638: The bounded heap replacement logic in the ranking loop is
missing the final id tie-break that HeapItem::Ord uses, so exact score and
token_count ties can preserve the wrong entry. Update the comparison in the heap
replacement path that builds and compares HeapItem values to mirror the full
ordering from Ord, including the id comparison after score and token_count, so
the same candidate wins in both the heap and replacement check.
- Around line 902-929: In the insert flow in rag.rs, the cheap quota checks are
happening too late because embed_text_via_tinfoil is called before insert-rate
and row limits are enforced, causing avoidable embedding/billing work for
requests that will be rejected. Move the inexpensive quota validation ahead of
embedding in the insert path, keeping the exact text-size and projected-byte
checks where they are needed before the actual insert. Apply the same reordering
to the other insert path mentioned in the review so both paths use the same
early-check pattern around embed_text_via_tinfoil, ensure_storage_limits, and
the insert validation logic.
---
Duplicate comments:
In `@migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql`:
- Around line 1-3: The rollback in down.sql drops the same indexes without the
non-blocking form, so mirror the up.sql approach by changing the DROP INDEX
statements for idx_user_embeddings_user_model_conversation_created_id,
idx_user_embeddings_user_model_source_created_id, and
idx_user_embeddings_user_model_created_id to use CONCURRENTLY. Keep in mind the
migration must not be wrapped in a transaction when using the concurrent drop,
so adjust the migration setup consistently with the existing up.sql fix.
🪄 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: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 2221ba81-ea9e-4b86-9599-94927732874f
📒 Files selected for processing (5)
migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sqlmigrations/2026-07-07-000001_user_embeddings_hot_path_indexes/up.sqlsrc/main.rssrc/rag.rssrc/web/mod.rs
💤 Files with no reviewable changes (1)
- src/web/mod.rs
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/main.rs (1)
1706-1712: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftDelete stale encrypted embeddings on destructive password reset.
Evicting the cache is not enough here: the reset creates a new user seed/key, so persisted embeddings encrypted with the old key become undecryptable but still count toward storage and scans. Delete the user’s embedding rows as part of the destructive reset flow, or explicitly re-encrypt them if that is intended.
🤖 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 `@src/main.rs` around lines 1706 - 1712, The destructive password reset flow in complete_destructive_password_reset currently only evicts the RAG cache, which leaves persisted embeddings encrypted with the old user key behind. Update this path in the reset handling around self.db.complete_destructive_password_reset and self.rag_cache.lock().await.evict_user(...) to also remove the user’s stored embedding rows, or re-encrypt them if that is the intended behavior, so stale encrypted embeddings do not remain in storage.src/rag.rs (2)
624-638: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winMirror the id tie-break in bounded heap replacement.
Ordprefers lower ids after score/token ties, but the replacement check omits that final tie-break, so exact ties can keep the wrong candidate.Proposed fix
let ordering = score .total_cmp(&min.score) - .then_with(|| min.token_count.cmp(&e.token_count)); + .then_with(|| min.token_count.cmp(&e.token_count)) + .then_with(|| min.id.cmp(&e.id));🤖 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 `@src/rag.rs` around lines 624 - 638, The bounded heap replacement logic in the ranking loop is missing the final id tie-break that HeapItem::Ord uses, so exact score and token_count ties can preserve the wrong entry. Update the comparison in the heap replacement path that builds and compares HeapItem values to mirror the full ordering from Ord, including the id comparison after score and token_count, so the same candidate wins in both the heap and replacement check.
902-929: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winCheck cheap quota limits before generating embeddings.
Both insert paths call the external embedding/billing flow before enforcing insert-rate and row limits, so over-quota requests still incur embedding cost before being rejected. Split the cheap rate/row checks ahead of
embed_text_via_tinfoil, then keep the exact byte check before insert.Also applies to: 981-996
🤖 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 `@src/rag.rs` around lines 902 - 929, In the insert flow in rag.rs, the cheap quota checks are happening too late because embed_text_via_tinfoil is called before insert-rate and row limits are enforced, causing avoidable embedding/billing work for requests that will be rejected. Move the inexpensive quota validation ahead of embedding in the insert path, keeping the exact text-size and projected-byte checks where they are needed before the actual insert. Apply the same reordering to the other insert path mentioned in the review so both paths use the same early-check pattern around embed_text_via_tinfoil, ensure_storage_limits, and the insert validation logic.
♻️ Duplicate comments (1)
migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql (1)
1-3: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winMirror the
CONCURRENTLYfix for rollback drops.Same lock-blocking concern applies to
DROP INDEXhere — a normal drop takes anACCESS EXCLUSIVElock. This should be addressed consistently with the up.sql fix (same transaction-wrapping caveat applies).🔧 Proposed fix
-DROP INDEX IF EXISTS idx_user_embeddings_user_model_conversation_created_id; -DROP INDEX IF EXISTS idx_user_embeddings_user_model_source_created_id; -DROP INDEX IF EXISTS idx_user_embeddings_user_model_created_id; +DROP INDEX CONCURRENTLY IF EXISTS idx_user_embeddings_user_model_conversation_created_id; +DROP INDEX CONCURRENTLY IF EXISTS idx_user_embeddings_user_model_source_created_id; +DROP INDEX CONCURRENTLY IF EXISTS idx_user_embeddings_user_model_created_id;🤖 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 `@migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql` around lines 1 - 3, The rollback in down.sql drops the same indexes without the non-blocking form, so mirror the up.sql approach by changing the DROP INDEX statements for idx_user_embeddings_user_model_conversation_created_id, idx_user_embeddings_user_model_source_created_id, and idx_user_embeddings_user_model_created_id to use CONCURRENTLY. Keep in mind the migration must not be wrapped in a transaction when using the concurrent drop, so adjust the migration setup consistently with the existing up.sql fix.Source: Linters/SAST tools
🤖 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 `@migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/up.sql`:
- Around line 1-9: The three index builds in the user_embeddings hot path
migration are currently blocking because they use plain CREATE INDEX. Update the
migration’s CREATE INDEX statements for
idx_user_embeddings_user_model_created_id,
idx_user_embeddings_user_model_source_created_id, and
idx_user_embeddings_user_model_conversation_created_id to use CREATE INDEX
CONCURRENTLY, and add a metadata.toml for this migration with run_in_transaction
= false so Diesel can execute it outside a transaction.
In `@src/main.rs`:
- Line 24: The application router is missing the RAG routes, so the new web APIs
are not being exposed. Update the router setup in main.rs where the existing
route groups are mounted to import and merge the RAG route set alongside
instructions_routes, login_routes, oauth_routes, openai_routes, and
protected_routes, ensuring the new subsystem is reachable through the app’s
router configuration.
In `@src/rag.rs`:
- Around line 1360-1371: The duplicate-load path in
load_user_embeddings_for_search is leaving the in-flight marker behind when the
awaited load fails, which can cause later callers to keep timing out. Update the
branch around the await in src/rag.rs so the cache cleanup in
rag_cache.finish_load(user_id) (or equivalent in-flight removal) runs even when
load_user_embeddings_for_search returns an error, using a guard/defer-style
cleanup tied to the load_user_embeddings_for_search / cache.put flow.
---
Outside diff comments:
In `@src/main.rs`:
- Around line 1706-1712: The destructive password reset flow in
complete_destructive_password_reset currently only evicts the RAG cache, which
leaves persisted embeddings encrypted with the old user key behind. Update this
path in the reset handling around self.db.complete_destructive_password_reset
and self.rag_cache.lock().await.evict_user(...) to also remove the user’s stored
embedding rows, or re-encrypt them if that is the intended behavior, so stale
encrypted embeddings do not remain in storage.
In `@src/rag.rs`:
- Around line 624-638: The bounded heap replacement logic in the ranking loop is
missing the final id tie-break that HeapItem::Ord uses, so exact score and
token_count ties can preserve the wrong entry. Update the comparison in the heap
replacement path that builds and compares HeapItem values to mirror the full
ordering from Ord, including the id comparison after score and token_count, so
the same candidate wins in both the heap and replacement check.
- Around line 902-929: In the insert flow in rag.rs, the cheap quota checks are
happening too late because embed_text_via_tinfoil is called before insert-rate
and row limits are enforced, causing avoidable embedding/billing work for
requests that will be rejected. Move the inexpensive quota validation ahead of
embedding in the insert path, keeping the exact text-size and projected-byte
checks where they are needed before the actual insert. Apply the same reordering
to the other insert path mentioned in the review so both paths use the same
early-check pattern around embed_text_via_tinfoil, ensure_storage_limits, and
the insert validation logic.
---
Duplicate comments:
In `@migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql`:
- Around line 1-3: The rollback in down.sql drops the same indexes without the
non-blocking form, so mirror the up.sql approach by changing the DROP INDEX
statements for idx_user_embeddings_user_model_conversation_created_id,
idx_user_embeddings_user_model_source_created_id, and
idx_user_embeddings_user_model_created_id to use CONCURRENTLY. Keep in mind the
migration must not be wrapped in a transaction when using the concurrent drop,
so adjust the migration setup consistently with the existing up.sql fix.
🪄 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: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 2221ba81-ea9e-4b86-9599-94927732874f
📒 Files selected for processing (5)
migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sqlmigrations/2026-07-07-000001_user_embeddings_hot_path_indexes/up.sqlsrc/main.rssrc/rag.rssrc/web/mod.rs
💤 Files with no reviewable changes (1)
- src/web/mod.rs
🛑 Comments failed to post (3)
migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/up.sql (1)
1-9: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Does Diesel ORM support disabling transaction wrapping for a single migration to allow CREATE INDEX CONCURRENTLY?💡 Result:
Yes, Diesel ORM supports disabling transaction wrapping for individual migrations, which is required for operations like CREATE INDEX CONCURRENTLY in PostgreSQL [1][2][3]. To disable the default transaction wrapping for a migration, you can use the following methods: For file-based migrations: Create a metadata.toml file within the specific migration directory and set the run_in_transaction flag to false [1][2]: # migrations/_<migration_name>/metadata.toml run_in_transaction = false For Rust-based migrations: If you are implementing migrations using Rust code, you can call the without_transaction method on the builder [4][2]: // Example for Rust-based migration RustMigration::new.up(|conn| { /* your SQL here */ }).without_transaction // Disables transaction wrapping.build Alternatively, if you are implementing the TypedMigration trait directly, you can return false from the run_in_transaction method [2][5]. Because CREATE INDEX CONCURRENTLY cannot run inside a transaction block in PostgreSQL, using these settings ensures the migration executes correctly without triggering an error [6][3][7].
Citations:
- 1: https://docs.rs/crate/diesel_migrations/latest
- 2: https://docs.diesel.rs/main/diesel_migrations/index.html
- 3: https://docs.rs/crate/diesel-guard/0.1.1
- 4: https://docs.diesel.rs/main/diesel_migrations/struct.RustMigration.html
- 5: https://docs.diesel.rs/master/diesel/migration/trait.MigrationMetadata.html
- 6: diesel-rs/diesel#2061
- 7: https://crates.io/crates/diesel-guard/0.6.0
🏁 Script executed:
#!/bin/bash set -euo pipefail echo "== migration files ==" git ls-files 'migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/**' || true echo echo "== nearby Diesel config references ==" rg -n --hidden --glob '!target' --glob '!node_modules' \ 'run_in_transaction|metadata\.toml|diesel_migrations|migrations' \ Cargo.toml . 2>/dev/null | sed -n '1,200p' echo echo "== migration directory contents ==" find migrations/2026-07-07-000001_user_embeddings_hot_path_indexes -maxdepth 2 -type f | sortRepository: OpenSecretCloud/opensecret
Length of output: 3157
🏁 Script executed:
#!/bin/bash set -euo pipefail echo "== diesel.toml ==" cat -n diesel.toml echo echo "== src/migrations.rs ==" cat -n src/migrations.rs echo echo "== migration directory metadata files ==" find migrations/2026-07-07-000001_user_embeddings_hot_path_indexes -maxdepth 1 -type f -name 'metadata.toml' -o -name '*.toml' -o -name '*.sql' | sort | xargs -r -I{} sh -c 'echo "--- {}"; cat -n "{}"; echo'Repository: OpenSecretCloud/opensecret
Length of output: 9026
🏁 Script executed:
#!/bin/bash set -euo pipefail rg -n --hidden --glob '!target' --glob '!node_modules' \ 'MigrationHarness|run_pending_migrations|embed_migrations|embedded_migrations|diesel_migrations|metadata\.toml|run_in_transaction' \ src migrations Cargo.toml . 2>/dev/null | sed -n '1,240p'Repository: OpenSecretCloud/opensecret
Length of output: 164
Use
CONCURRENTLYfor these indexes
TheseCREATE INDEXstatements will block writes while they build. Switch them toCREATE INDEX CONCURRENTLYand addmetadata.tomlwithrun_in_transaction = falseso Diesel can run this migration outside a transaction.🧰 Tools
🪛 Squawk (2.59.0)
[warning] 2-3: During normal index creation, table updates are blocked, but reads are still allowed. Use
concurrentlyto avoid blocking writes.(require-concurrent-index-creation)
[warning] 5-6: During normal index creation, table updates are blocked, but reads are still allowed. Use
concurrentlyto avoid blocking writes.(require-concurrent-index-creation)
[warning] 8-9: During normal index creation, table updates are blocked, but reads are still allowed. Use
concurrentlyto avoid blocking writes.(require-concurrent-index-creation)
🤖 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 `@migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/up.sql` around lines 1 - 9, The three index builds in the user_embeddings hot path migration are currently blocking because they use plain CREATE INDEX. Update the migration’s CREATE INDEX statements for idx_user_embeddings_user_model_created_id, idx_user_embeddings_user_model_source_created_id, and idx_user_embeddings_user_model_conversation_created_id to use CREATE INDEX CONCURRENTLY, and add a metadata.toml for this migration with run_in_transaction = false so Diesel can execute it outside a transaction.Source: Linters/SAST tools
src/main.rs (1)
24-24: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Mount the RAG routes in the application router.
The RAG route import/merge is absent, so the new RAG web APIs are not reachable even though the PR adds the subsystem.
🤖 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 `@src/main.rs` at line 24, The application router is missing the RAG routes, so the new web APIs are not being exposed. Update the router setup in main.rs where the existing route groups are mounted to import and merge the RAG route set alongside instructions_routes, login_routes, oauth_routes, openai_routes, and protected_routes, ensuring the new subsystem is reachable through the app’s router configuration.src/rag.rs (1)
1360-1371: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Always clear the in-flight marker on duplicate load errors.
The timeout-duplicate branch returns early on
load_user_embeddings_for_search(...).await?, so an error leavesin_flight_loadsstuck and future callers can loop through wait timeouts indefinitely.Proposed fix
- let loaded = load_user_embeddings_for_search( + let loaded = load_user_embeddings_for_search( state, user_id, user_key, filters, limits.scan_limit, ) - .await?; + .await; let mut cache = state.rag_cache.lock().await; - let cached = cache.put(user_id, loaded.embeddings.clone(), loaded.scan_limit_hit); + if let Ok(loaded) = &loaded { + let cached = cache.put(user_id, loaded.embeddings.clone(), loaded.scan_limit_hit); + info!( + target: "rag", + user_id = %user_id, + cached, + loaded_rows = loaded.embeddings.len(), + cache_bytes_per_user = cached_embeddings_bytes(&loaded.embeddings), + total_cache_bytes = cache.total_bytes, + scan_limit_hit = loaded.scan_limit_hit, + reason = "timeout_duplicate", + "rag_cache_store_after_load" + ); + } cache.finish_load(user_id); - info!( - target: "rag", - user_id = %user_id, - cached, - loaded_rows = loaded.embeddings.len(), - cache_bytes_per_user = cached_embeddings_bytes(&loaded.embeddings), - total_cache_bytes = cache.total_bytes, - scan_limit_hit = loaded.scan_limit_hit, - reason = "timeout_duplicate", - "rag_cache_store_after_load" - ); - return Ok(loaded); + return loaded;📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.let loaded = load_user_embeddings_for_search( state, user_id, user_key, filters, limits.scan_limit, ) .await; let mut cache = state.rag_cache.lock().await; if let Ok(loaded) = &loaded { let cached = cache.put(user_id, loaded.embeddings.clone(), loaded.scan_limit_hit); info!( target: "rag", user_id = %user_id, cached, loaded_rows = loaded.embeddings.len(), cache_bytes_per_user = cached_embeddings_bytes(&loaded.embeddings), total_cache_bytes = cache.total_bytes, scan_limit_hit = loaded.scan_limit_hit, reason = "timeout_duplicate", "rag_cache_store_after_load" ); } cache.finish_load(user_id); return loaded;🤖 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 `@src/rag.rs` around lines 1360 - 1371, The duplicate-load path in load_user_embeddings_for_search is leaving the in-flight marker behind when the awaited load fails, which can cause later callers to keep timing out. Update the branch around the await in src/rag.rs so the cache cleanup in rag_cache.finish_load(user_id) (or equivalent in-flight removal) runs even when load_user_embeddings_for_search returns an error, using a guard/defer-style cleanup tied to the load_user_embeddings_for_search / cache.put flow.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
|
@coderabbitai review |
✅ Action performedReview finished.
|
| pub async fn get_embedding_vector( | ||
| state: &Arc<AppState>, | ||
| user: &User, | ||
| auth_method: AuthMethod, | ||
| model: &str, | ||
| input: &str, | ||
| dimensions: Option<i32>, | ||
| ) -> Result<(Vec<f32>, i32), ApiError> { | ||
| let request = EmbeddingRequest { | ||
| input: Value::String(input.to_string()), | ||
| model: model.to_string(), | ||
| encoding_format: Some("float".to_string()), | ||
| dimensions, | ||
| user: None, | ||
| }; | ||
|
|
||
| let (response_json, _provider, prompt_tokens) = | ||
| request_embeddings_with_billing(state, user, auth_method, &request).await?; | ||
|
|
||
| let embedding = response_json | ||
| .get("data") | ||
| .and_then(|v| v.as_array()) | ||
| .and_then(|arr| arr.first()) | ||
| .and_then(|obj| obj.get("embedding")) | ||
| .and_then(|v| v.as_array()) | ||
| .ok_or(ApiError::InternalServerError)?; | ||
|
|
||
| let mut vector: Vec<f32> = Vec::with_capacity(embedding.len()); | ||
| for v in embedding { | ||
| let f = v.as_f64().ok_or(ApiError::InternalServerError)?; | ||
| vector.push(f as f32); | ||
| } | ||
|
|
||
| if let Some(dim) = dimensions { | ||
| if vector.len() != dim as usize { | ||
| error!( | ||
| "Embeddings response dim mismatch: expected={}, got={}", | ||
| dim, | ||
| vector.len() | ||
| ); | ||
| return Err(ApiError::InternalServerError); | ||
| } | ||
| } | ||
|
|
||
| Ok((vector, prompt_tokens)) | ||
| } |
There was a problem hiding this comment.
🔍 New get_embedding_vector public API bills every RAG internal embedding call
The newly exported get_embedding_vector function (src/web/openai.rs:2146) calls request_embeddings_with_billing, which means every RAG insert and search operation triggers a billing event via publish_usage_event_internal. This is called from embed_text_via_tinfoil (src/rag.rs:697-712) for both insert_archival_embedding and search_user_embeddings_with_options. This is likely intentional (embedding API calls have real cost), but it means RAG search queries incur billing for the query embedding in addition to any completion billing. Worth confirming this is the desired billing behavior.
Was this helpful? React with 👍 or 👎 to provide feedback.
There was a problem hiding this comment.
🧹 Nitpick comments (2)
src/rag.rs (2)
839-880: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winAvoid the duplicate project storage-stats scan.
ensure_pre_embedding_limits(Line 928) already runsload_project_embedding_storage_stats, thenensure_storage_limitsimmediately reloads it again at Line 860. On the insert path this is a full aggregate scan repeated needlessly (and the whole limit block already runs once as a pre-check and again at insert time). Return the already-loaded stats fromensure_pre_embedding_limitsand reuse them here.♻️ Sketch
// ensure_pre_embedding_limits already computes project_stats; // return (user_stats, project_stats) and reuse in ensure_storage_limits // instead of calling load_project_embedding_storage_stats a second time.🤖 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 `@src/rag.rs` around lines 839 - 880, The storage limit check in ensure_storage_limits is reloading project embedding stats even though ensure_pre_embedding_limits already computed them, causing a duplicate aggregate scan. Update ensure_pre_embedding_limits to return both user_stats and project_stats, then change ensure_storage_limits to reuse the returned project_stats instead of calling load_project_embedding_storage_stats again; keep the existing limit comparisons and warning logs in ensure_storage_limits.
1406-1443: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCache-hit path recomputes byte size O(n) just for a log field.
cached_embeddings_bytes(&embeddings)at Line 1421 re-sums every cached row on each cache hit, even thoughCacheEntry.bytesis already tracked. For large users (up toscan_limitrows) this adds avoidable per-search CPU on the hottest path. Consider havinggetalso return the stored byte count (or omitcache_bytes_per_userfrom this hit log).♻️ Sketch
// return the cached byte size from get(): pub fn get(&mut self, user_id: Uuid) -> Option<(Arc<Vec<CachedEmbedding>>, bool, usize)> { // ... include entry.bytes in the tuple, then log it directly. }🤖 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 `@src/rag.rs` around lines 1406 - 1443, The cache-hit logging in load_cacheable_embeddings_with_coordination is recomputing cached_embeddings_bytes(&embeddings) on every hit, which is avoidable work on a hot path. Update Cache::get (or the cache-hit branch in load_cacheable_embeddings_with_coordination) to return the stored byte count from CacheEntry.bytes and log that value directly, or remove cache_bytes_per_user from the hit log if it is not needed.
🤖 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.
Nitpick comments:
In `@src/rag.rs`:
- Around line 839-880: The storage limit check in ensure_storage_limits is
reloading project embedding stats even though ensure_pre_embedding_limits
already computed them, causing a duplicate aggregate scan. Update
ensure_pre_embedding_limits to return both user_stats and project_stats, then
change ensure_storage_limits to reuse the returned project_stats instead of
calling load_project_embedding_storage_stats again; keep the existing limit
comparisons and warning logs in ensure_storage_limits.
- Around line 1406-1443: The cache-hit logging in
load_cacheable_embeddings_with_coordination is recomputing
cached_embeddings_bytes(&embeddings) on every hit, which is avoidable work on a
hot path. Update Cache::get (or the cache-hit branch in
load_cacheable_embeddings_with_coordination) to return the stored byte count
from CacheEntry.bytes and log that value directly, or remove
cache_bytes_per_user from the hit log if it is not needed.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 87dc8337-e39b-4062-82d4-fe6882c97692
📒 Files selected for processing (6)
migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sqlmigrations/2026-07-07-000001_user_embeddings_hot_path_indexes/metadata.tomlmigrations/2026-07-07-000001_user_embeddings_hot_path_indexes/up.sqlsrc/main.rssrc/rag.rssrc/web/mod.rs
💤 Files with no reviewable changes (1)
- src/web/mod.rs
✅ Files skipped from review due to trivial changes (1)
- migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/metadata.toml
🚧 Files skipped from review as they are similar to previous changes (2)
- migrations/2026-07-07-000001_user_embeddings_hot_path_indexes/down.sql
- src/main.rs
1 participant
Summary
Verification
nix develop -c cargo checkNotes
master.Summary by CodeRabbit