context: Align diagnostics retrieval context

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<AutocompleteRequest> {
 		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<FileChunk[]> {
+		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<FileChunk[]> {
+		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<FileChunk[]> {
+		try {
+			const results =
+				(await vscode.commands.executeCommand<
+					Array<vscode.Location | vscode.LocationLink> | 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<FileChunk[]> {
+		try {
+			const results =
+				(await vscode.commands.executeCommand<vscode.Location[] | undefined>(
+					"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<FileChunk[]> {
+		const seen = new Set<string>();
+		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<FileChunk | null> {
+		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 {

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),
+	};
+}

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<typeof FileChunkSchema>;
 export type UserAction = z.infer<typeof UserActionSchema>;
+export type EditorDiagnostic = z.infer<typeof EditorDiagnosticSchema>;
 export type AutocompleteRequest = z.infer<typeof AutocompleteRequestSchema>;
 export type AutocompleteResponse = z.infer<typeof AutocompleteResponseSchema>;
 export type AutocompleteMetricsRequest = z.infer<