| 1 | import * as vscode from "vscode"; |
| 2 | |
| 3 | /** A piece of editor context attached to the next prompt. */ |
| 4 | export interface ContextChip { |
| 5 | id: string; |
| 6 | kind: "selection" | "file" | "diagnostics"; |
| 7 | label: string; |
| 8 | detail?: string; |
| 9 | /** Full text included in the assembled prompt; may be truncated. */ |
| 10 | body: string; |
| 11 | } |
| 12 | |
| 13 | const MAX_BODY_CHARS = 8000; |
| 14 | |
| 15 | let chipCounter = 0; |
| 16 | |
| 17 | export function collectSelectionContext(): ContextChip | undefined { |
| 18 | const editor = vscode.window.activeTextEditor; |
| 19 | if (!editor || editor.selection.isEmpty) { |
| 20 | return undefined; |
| 21 | } |
| 22 | const selection = editor.selection; |
| 23 | const text = editor.document.getText(selection); |
| 24 | if (text.trim().length === 0) { |
| 25 | return undefined; |
| 26 | } |
| 27 | const relative = workspaceRelativePath(editor.document.uri); |
| 28 | const startLine = selection.start.line + 1; |
| 29 | const endLine = selection.end.line + 1; |
| 30 | const lines = endLine - startLine + 1; |
| 31 | return { |
| 32 | id: `chip-${++chipCounter}`, |
| 33 | kind: "selection", |
| 34 | label: `${relative}:${startLine}-${endLine}`, |
| 35 | detail: `${lines} line${lines === 1 ? "" : "s"}`, |
| 36 | body: clip(`${text}\n`), |
| 37 | }; |
| 38 | } |
| 39 | |
| 40 | export function collectActiveFileContext(): ContextChip | undefined { |
| 41 | const editor = vscode.window.activeTextEditor; |
| 42 | if (!editor) { |
| 43 | return undefined; |
| 44 | } |
| 45 | const relative = workspaceRelativePath(editor.document.uri); |
| 46 | const text = editor.document.getText(); |
| 47 | return { |
| 48 | id: `chip-${++chipCounter}`, |
| 49 | kind: "file", |
| 50 | label: relative, |
| 51 | detail: `${editor.document.lineCount} lines`, |
| 52 | body: clip(text), |
| 53 | }; |
| 54 | } |
| 55 | |
| 56 | export function collectDiagnosticsContext(): ContextChip | undefined { |
| 57 | const editor = vscode.window.activeTextEditor; |
| 58 | const entries: string[] = []; |
| 59 | const uris = editor |
| 60 | ? [editor.document.uri] |
| 61 | : vscode.workspace.textDocuments.map((document) => document.uri); |
| 62 | for (const uri of uris) { |
| 63 | for (const diagnostic of vscode.languages.getDiagnostics(uri)) { |
| 64 | if (diagnostic.severity > vscode.DiagnosticSeverity.Warning) { |
| 65 | continue; |
| 66 | } |
| 67 | const position = `${uri.path.split("/").pop()}:${diagnostic.range.start.line + 1}`; |
| 68 | const severity = diagnostic.severity === vscode.DiagnosticSeverity.Error ? "error" : "warn"; |
| 69 | entries.push(`- [${severity}] ${position} ${diagnostic.message.split("\n")[0]}`); |
| 70 | if (entries.length >= 20) { |
| 71 | break; |
| 72 | } |
| 73 | } |
| 74 | } |
| 75 | if (entries.length === 0) { |
| 76 | return undefined; |
| 77 | } |
| 78 | return { |
| 79 | id: `chip-${++chipCounter}`, |
| 80 | kind: "diagnostics", |
| 81 | label: "Problems", |
| 82 | detail: `${entries.length} entrie${entries.length === 1 ? "" : "s"}`, |
| 83 | body: clip(entries.join("\n")), |
| 84 | }; |
| 85 | } |
| 86 | |
| 87 | /** Assemble the final prompt with context blocks the runtime model can read. */ |
| 88 | export function assemblePrompt(prompt: string, chips: ContextChip[]): string { |
| 89 | if (chips.length === 0) { |
| 90 | return prompt; |
| 91 | } |
| 92 | const sections = chips.map((chip) => { |
| 93 | const header = |
| 94 | chip.kind === "diagnostics" |
| 95 | ? "Diagnostics (problems panel)" |
| 96 | : `File \`${chip.label}\`${chip.kind === "selection" ? " (selected lines)" : ""}`; |
| 97 | return `--- ${header} ---\n${chip.body}`; |
| 98 | }); |
| 99 | return `${prompt}\n\n[Attached context]\n${sections.join("\n\n")}`; |
| 100 | } |
| 101 | |
| 102 | export function workspaceRelativePath(uri: vscode.Uri): string { |
| 103 | const folder = vscode.workspace.getWorkspaceFolder(uri); |
| 104 | if (!folder) { |
| 105 | return uri.fsPath; |
| 106 | } |
| 107 | const prefix = folder.uri.fsPath.endsWith("/") |
| 108 | ? folder.uri.fsPath |
| 109 | : `${folder.uri.fsPath}/`; |
| 110 | return uri.fsPath.startsWith(prefix) ? uri.fsPath.slice(prefix.length) : uri.fsPath; |
| 111 | } |
| 112 | |
| 113 | function clip(text: string): string { |
| 114 | if (text.length <= MAX_BODY_CHARS) { |
| 115 | return text; |
| 116 | } |
| 117 | return `${text.slice(0, MAX_BODY_CHARS)}\n… [truncated ${text.length - MAX_BODY_CHARS} chars]`; |
| 118 | } |
| 119 |