diff --git a/src/api/client.ts b/src/api/client.ts index 8515d45..107778b 100644 --- a/src/api/client.ts +++ b/src/api/client.ts @@ -15,6 +15,10 @@ import { utf8ByteOffsetAt, utf8ByteOffsetToUtf16Offset, } from "~/utils/text.ts"; +import { + fuseAndDedupRetrievalSnippets, + truncateRetrievalChunk, +} from "./retrieval-chunks.ts"; import { type AutocompleteMetricsRequest, AutocompleteMetricsRequestSchema, @@ -23,6 +27,7 @@ import { type AutocompleteResponse, AutocompleteResponseSchema, type AutocompleteResult, + type EditorDiagnostic, type FileChunk, type RecentBuffer, type RecentChange, @@ -39,6 +44,15 @@ export interface AutocompleteInput { userActions: UserAction[]; } +const MAX_RETRIEVAL_CHUNKS = 16; +const MAX_DEFINITION_CHUNKS = 6; +const MAX_USAGE_CHUNKS = 6; +const MAX_RETRIEVAL_CHUNK_LINES = 200; +const RETRIEVAL_CONTEXT_LINES_ABOVE = 9; +const RETRIEVAL_CONTEXT_LINES_BELOW = 9; +const MAX_CLIPBOARD_LINES = 20; +const MAX_DIAGNOSTICS = 50; + export class ApiClient { private apiUrl: string; private metricsUrl: string; @@ -69,7 +83,7 @@ export class ApiClient { return null; } - const requestData = this.buildRequest(input); + const requestData = await this.buildRequest(input); const parsedRequest = AutocompleteRequestSchema.safeParse(requestData); if (!parsedRequest.success) { @@ -157,7 +171,9 @@ export class ApiClient { return config.apiKey; } - private buildRequest(input: AutocompleteInput): AutocompleteRequest { + private async buildRequest( + input: AutocompleteInput, + ): Promise { const { document, position, @@ -171,7 +187,16 @@ export class ApiClient { const filePath = toUnixPath(document.uri.fsPath) || "untitled"; const recentChangesText = this.formatRecentChanges(recentChanges); const fileChunks = this.buildFileChunks(recentBuffers); - const retrievalChunks = this.buildDiagnosticsChunk(filePath, diagnostics); + const retrievalChunks = await this.buildRetrievalChunks( + document, + position, + filePath, + diagnostics, + ); + const editorDiagnostics = this.buildEditorDiagnostics( + document, + diagnostics, + ); return { debug_info: this.getDebugInfo(), @@ -185,6 +210,7 @@ export class ApiClient { multiple_suggestions: true, file_chunks: fileChunks, retrieval_chunks: retrievalChunks, + editor_diagnostics: editorDiagnostics, recent_user_actions: userActions, use_bytes: true, privacy_mode_enabled: config.privacyMode, @@ -239,14 +265,40 @@ export class ApiClient { }); } - private buildDiagnosticsChunk( + private async buildRetrievalChunks( + document: vscode.TextDocument, + position: vscode.Position, + currentFilePath: string, + diagnostics: vscode.Diagnostic[], + ): Promise { + const [definitionChunks, usageChunks, clipboardChunks] = await Promise.all([ + this.buildDefinitionChunks(document, position), + this.buildUsageChunks(document, position), + this.buildClipboardChunks(), + ]); + + const chunks = [ + ...this.buildDiagnosticsTextChunk(currentFilePath, diagnostics), + ...clipboardChunks, + ...usageChunks, + ...definitionChunks, + ] + .filter((chunk) => chunk.file_path !== currentFilePath) + .map((chunk) => truncateRetrievalChunk(chunk, MAX_RETRIEVAL_CHUNK_LINES)) + .filter((chunk) => chunk.content.trim().length > 0); + + return fuseAndDedupRetrievalSnippets(chunks).slice(-MAX_RETRIEVAL_CHUNKS); + } + + private buildDiagnosticsTextChunk( filePath: string, diagnostics: vscode.Diagnostic[], ): FileChunk[] { if (diagnostics.length === 0) return []; let content = ""; - for (const d of diagnostics) { + const limitedDiagnostics = diagnostics.slice(0, MAX_DIAGNOSTICS); + for (const d of limitedDiagnostics) { const severity = this.formatSeverity(d.severity); const line = d.range.start.line + 1; const col = d.range.start.character + 1; @@ -257,12 +309,159 @@ export class ApiClient { { file_path: "diagnostics", start_line: 1, - end_line: diagnostics.length, + end_line: limitedDiagnostics.length, content, }, ]; } + private buildEditorDiagnostics( + document: vscode.TextDocument, + diagnostics: vscode.Diagnostic[], + ): EditorDiagnostic[] { + return diagnostics.slice(0, MAX_DIAGNOSTICS).map((diagnostic) => ({ + line: diagnostic.range.start.line + 1, + start_offset: document.offsetAt(diagnostic.range.start), + end_offset: document.offsetAt(diagnostic.range.end), + severity: this.formatSeverity(diagnostic.severity), + message: diagnostic.message, + timestamp: Date.now(), + })); + } + + private async buildClipboardChunks(): Promise { + try { + const clipboard = (await vscode.env.clipboard.readText()).trim(); + if (!clipboard) return []; + + const lines = clipboard.split(/\r?\n/).slice(0, MAX_CLIPBOARD_LINES); + const content = lines.join("\n").trim(); + if (!content) return []; + + return [ + { + file_path: "clipboard.txt", + start_line: 1, + end_line: lines.length, + content, + timestamp: Date.now(), + }, + ]; + } catch { + return []; + } + } + + private async buildDefinitionChunks( + document: vscode.TextDocument, + position: vscode.Position, + ): Promise { + try { + const results = + (await vscode.commands.executeCommand< + Array | undefined + >("vscode.executeDefinitionProvider", document.uri, position)) ?? []; + const locations = results + .map((result) => this.normalizeLocation(result)) + .filter((location): location is vscode.Location => location !== null); + return this.buildLocationChunks(locations, MAX_DEFINITION_CHUNKS); + } catch { + return []; + } + } + + private async buildUsageChunks( + document: vscode.TextDocument, + position: vscode.Position, + ): Promise { + try { + const results = + (await vscode.commands.executeCommand( + "vscode.executeReferenceProvider", + document.uri, + position, + )) ?? []; + return this.buildLocationChunks(results, MAX_USAGE_CHUNKS); + } catch { + return []; + } + } + + private normalizeLocation( + location: vscode.Location | vscode.LocationLink, + ): vscode.Location | null { + if ("uri" in location && "range" in location) { + return new vscode.Location(location.uri, location.range); + } + if ("targetUri" in location && "targetRange" in location) { + return new vscode.Location(location.targetUri, location.targetRange); + } + return null; + } + + private async buildLocationChunks( + locations: readonly vscode.Location[], + maxChunks: number, + ): Promise { + const seen = new Set(); + const chunks: FileChunk[] = []; + + for (const location of locations) { + if (chunks.length >= maxChunks) break; + const key = `${location.uri.toString()}:${location.range.start.line}:${location.range.end.line}`; + if (seen.has(key)) continue; + seen.add(key); + + const chunk = await this.buildChunkFromLocation(location); + if (!chunk) continue; + chunks.push(chunk); + } + + return chunks; + } + + private async buildChunkFromLocation( + location: vscode.Location, + ): Promise { + let targetDocument: vscode.TextDocument; + try { + targetDocument = await vscode.workspace.openTextDocument(location.uri); + } catch { + return null; + } + + const totalLines = targetDocument.lineCount; + if (totalLines === 0) return null; + + const startLine = Math.max( + 0, + location.range.start.line - RETRIEVAL_CONTEXT_LINES_ABOVE, + ); + const endLine = Math.min( + totalLines - 1, + location.range.end.line + RETRIEVAL_CONTEXT_LINES_BELOW, + ); + const endPosition = + endLine + 1 < totalLines + ? new vscode.Position(endLine + 1, 0) + : targetDocument.lineAt(endLine).range.end; + const range = new vscode.Range( + new vscode.Position(startLine, 0), + endPosition, + ); + const content = targetDocument.getText(range).trim(); + if (!content) return null; + + return { + file_path: + toUnixPath(targetDocument.uri.fsPath) || targetDocument.uri.toString(), + start_line: startLine + 1, + end_line: endLine + 1, + content, + timestamp: Date.now(), + }; + } + private formatSeverity( severity: vscode.DiagnosticSeverity | undefined, ): string { diff --git a/src/api/retrieval-chunks.ts b/src/api/retrieval-chunks.ts new file mode 100644 index 0000000..069e238 --- /dev/null +++ b/src/api/retrieval-chunks.ts @@ -0,0 +1,79 @@ +import type { FileChunk } from "./schemas.ts"; + +export function truncateRetrievalChunk( + chunk: FileChunk, + maxLines: number, +): FileChunk { + const lines = chunk.content.split("\n"); + if (lines.length <= maxLines) { + return chunk; + } + const truncatedLines = lines.slice(0, maxLines); + return { + ...chunk, + end_line: Math.min(chunk.end_line, chunk.start_line + maxLines - 1), + content: truncatedLines.join("\n"), + }; +} + +export function fuseAndDedupRetrievalSnippets( + snippets: FileChunk[], +): FileChunk[] { + const fused: FileChunk[] = []; + + snippetsLoop: for (const snippet of snippets) { + for (let i = 0; i < fused.length; i++) { + const existing = fused[i]; + if (!existing) continue; + + if ( + existing.file_path === snippet.file_path && + rangesTouch(existing, snippet) + ) { + fused[i] = mergeSnippet(existing, snippet); + continue snippetsLoop; + } + + if ( + existing.file_path === snippet.file_path && + existing.start_line === snippet.start_line && + existing.end_line === snippet.end_line && + existing.content === snippet.content + ) { + continue snippetsLoop; + } + } + + fused.push(snippet); + } + + return fused; +} + +function rangesTouch(a: FileChunk, b: FileChunk): boolean { + return b.start_line <= a.end_line + 1 && a.start_line <= b.end_line + 1; +} + +function mergeSnippet(a: FileChunk, b: FileChunk): FileChunk { + if (a.start_line <= b.start_line && a.end_line >= b.end_line) { + return a; + } + if (b.start_line <= a.start_line && b.end_line >= a.end_line) { + return b; + } + + const startLine = Math.min(a.start_line, b.start_line); + const endLine = Math.max(a.end_line, b.end_line); + const content = + a.start_line <= b.start_line + ? `${a.content}\n${b.content}` + : `${b.content}\n${a.content}`; + + return { + file_path: a.file_path, + start_line: startLine, + end_line: endLine, + content: content.trim(), + timestamp: Math.max(a.timestamp ?? 0, b.timestamp ?? 0), + }; +} diff --git a/src/api/schemas.ts b/src/api/schemas.ts index ba7e23b..95dfd01 100644 --- a/src/api/schemas.ts +++ b/src/api/schemas.ts @@ -24,6 +24,15 @@ export const UserActionSchema = z.object({ timestamp: z.number(), }); +export const EditorDiagnosticSchema = z.object({ + line: z.number(), + start_offset: z.number(), + end_offset: z.number(), + severity: z.string(), + message: z.string(), + timestamp: z.number(), +}); + export const AutocompleteRequestSchema = z.object({ debug_info: z.string(), repo_name: z.string(), @@ -37,6 +46,7 @@ export const AutocompleteRequestSchema = z.object({ multiple_suggestions: z.boolean(), file_chunks: z.array(FileChunkSchema), retrieval_chunks: z.array(FileChunkSchema), + editor_diagnostics: z.array(EditorDiagnosticSchema), recent_user_actions: z.array(UserActionSchema), use_bytes: z.boolean(), privacy_mode_enabled: z.boolean(), @@ -100,6 +110,7 @@ export const AutocompleteMetricsRequestSchema = z.object({ export type FileChunk = z.infer; export type UserAction = z.infer; +export type EditorDiagnostic = z.infer; export type AutocompleteRequest = z.infer; export type AutocompleteResponse = z.infer; export type AutocompleteMetricsRequest = z.infer< diff --git a/test/retrieval-chunks.test.ts b/test/retrieval-chunks.test.ts new file mode 100644 index 0000000..ed7b9b0 --- /dev/null +++ b/test/retrieval-chunks.test.ts @@ -0,0 +1,116 @@ +import { describe, expect, test } from "bun:test"; + +import { + fuseAndDedupRetrievalSnippets, + truncateRetrievalChunk, +} from "~/api/retrieval-chunks.ts"; +import type { FileChunk } from "~/api/schemas.ts"; + +function chunk(overrides: Partial): FileChunk { + return { + file_path: "src/example.ts", + start_line: 1, + end_line: 1, + content: "line", + timestamp: 1, + ...overrides, + }; +} + +describe("truncateRetrievalChunk", () => { + test("truncates long chunks and updates end_line", () => { + const input = chunk({ + start_line: 10, + end_line: 20, + content: "a\nb\nc\nd\ne", + }); + + const result = truncateRetrievalChunk(input, 3); + + expect(result.content).toBe("a\nb\nc"); + expect(result.end_line).toBe(12); + }); + + test("returns original chunk when already within max lines", () => { + const input = chunk({ + start_line: 2, + end_line: 3, + content: "a\nb", + }); + + const result = truncateRetrievalChunk(input, 3); + + expect(result).toEqual(input); + }); +}); + +describe("fuseAndDedupRetrievalSnippets", () => { + test("fuses touching ranges from the same file", () => { + const result = fuseAndDedupRetrievalSnippets([ + chunk({ start_line: 5, end_line: 6, content: "a\nb", timestamp: 10 }), + chunk({ start_line: 7, end_line: 8, content: "c\nd", timestamp: 15 }), + ]); + + expect(result).toHaveLength(1); + expect(result[0]).toEqual( + chunk({ + start_line: 5, + end_line: 8, + content: "a\nb\nc\nd", + timestamp: 15, + }), + ); + }); + + test("deduplicates identical chunks", () => { + const duplicate = chunk({ + start_line: 3, + end_line: 4, + content: "x\ny", + }); + + const result = fuseAndDedupRetrievalSnippets([duplicate, duplicate]); + + expect(result).toEqual([duplicate]); + }); + + test("does not fuse snippets across different files", () => { + const result = fuseAndDedupRetrievalSnippets([ + chunk({ + file_path: "src/a.ts", + start_line: 1, + end_line: 2, + content: "a", + }), + chunk({ + file_path: "src/b.ts", + start_line: 2, + end_line: 3, + content: "b", + }), + ]); + + expect(result).toHaveLength(2); + expect(result[0]?.file_path).toBe("src/a.ts"); + expect(result[1]?.file_path).toBe("src/b.ts"); + }); + + test("keeps larger containing range when one snippet fully covers the other", () => { + const large = chunk({ + start_line: 10, + end_line: 20, + content: "large", + timestamp: 2, + }); + const small = chunk({ + start_line: 12, + end_line: 13, + content: "small", + timestamp: 3, + }); + + const result = fuseAndDedupRetrievalSnippets([large, small]); + + expect(result).toEqual([large]); + }); +});