Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions src/main/search/code_index.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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)');
Expand Down
7 changes: 5 additions & 2 deletions src/main/search/code_index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -196,7 +196,10 @@ export class CodeIndex implements CodeSearch {
onProgress?: (progress: UpdateProgress) => void,
): Promise<SearchResult> {
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));

Expand Down Expand Up @@ -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,
Expand Down