diff --git a/src/main/search/code_index.test.ts b/src/main/search/code_index.test.ts index 2b9d383..6a1c6e3 100644 --- a/src/main/search/code_index.test.ts +++ b/src/main/search/code_index.test.ts @@ -3,6 +3,7 @@ import { tmpdir } from 'node:os'; import { join } from 'node:path'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { SettingsStore, type SecretCipher } from '../settings'; +import { REDACTION_MARK } from '../tools/redact'; import { ShellRunner } from '../tools/shell'; import { Workspace } from '../tools/workspace'; import { chunkFile } from './chunker'; @@ -485,6 +486,21 @@ describe('CodeIndex with a reranker', () => { expect(documents).not.toContain('hunter2hunter2'); }); + it('masks secrets in the query before it leaves to the embedding and the reranker (#286)', async () => { + const embedder = new FakeEmbedder(); + const reranker = new FakeReranker('addToCart'); + const index = new CodeIndex(new Workspace(root), embedder, indexDir, () => 1000, reranker); + const token = `ghp_${'a'.repeat(36)}`; + + // A model can copy a token it read straight into its search query, and both services are third parties. + await index.searchDetailed(`fix the login that stores ${token}`, 2, signal); + + expect(embedder.inputs.filter((input) => input === 'query')).toHaveLength(1); + const sent = embedder.calls[embedder.inputs.indexOf('query')]; + expect(sent).toEqual([`fix the login that stores ${REDACTION_MARK}`]); + expect(reranker.seen!.query).toBe(`fix the login that stores ${REDACTION_MARK}`); + }); + it('falls back to the embedding order when reranking fails, and says so', async () => { const reranker = new FakeReranker('addToCart'); reranker.fail = new Error('OpenRouter rerank request failed (503)'); diff --git a/src/main/search/code_index.ts b/src/main/search/code_index.ts index 3993686..2996ff0 100644 --- a/src/main/search/code_index.ts +++ b/src/main/search/code_index.ts @@ -196,7 +196,10 @@ export class CodeIndex implements CodeSearch { onProgress?: (progress: UpdateProgress) => void, ): Promise { await this.update(signal, onProgress); - const [queryVector] = await this.embedder.embed([query], 'query', signal); + // The query leaves the machine for the embedding and the rerank service, so a secret the model copied into it + // is masked first, the way a web search query is (#73) and the way the documents below already are. + const sent = redactSecrets(query); + const [queryVector] = await this.embedder.embed([sent], 'query', signal); if (!queryVector) throw new Error('The embedding service returned no vector for the query.'); const q = normalize(Float32Array.from(queryVector)); @@ -231,7 +234,7 @@ export class CodeIndex implements CodeSearch { try { const ranked = await this.reranker.rerank( - query, + sent, candidates.map((hit) => redactSecrets(`${hit.path}\n${hit.text}`)), Math.min(limit, candidates.length), signal,