Merge pull request #49 from docker/cm/0.1.8

Plug in RPC client
docker · Oct 29, 2024 · 3bda8b1 · 3bda8b1
2 parents b7acc6b + e30779f
commit 3bda8b1
Show file tree

Hide file tree

Showing 7 changed files with 527 additions and 512 deletions.
diff --git a/package.json b/package.json
@@ -2,7 +2,7 @@
   "name": "labs-ai-tools-vscode",
   "displayName": "Labs: AI Tools for VSCode",
   "description": "Run & Debug AI Prompts with Dockerized tools",
-  "version": "0.1.7",
+  "version": "0.1.8",
   "publisher": "docker",
   "repository": {
     "type": "git",
@@ -106,12 +106,13 @@
     "@typescript-eslint/parser": "^7.4.0",
     "@vscode/test-cli": "^0.0.8",
     "@vscode/test-electron": "^2.3.9",
+    "@vscode/vsce": "^3.1.1",
     "eslint": "^8.57.0",
-    "typescript": "^5.3.3",
-    "@vscode/vsce": "^3.1.1"
+    "typescript": "^5.3.3"
   },
   "dependencies": {
     "semver": "^7.6.3",
+    "vscode-jsonrpc": "^8.2.1",
     "vscode-languageclient": "^8.1.0"
   },
   "packageManager": "yarn@1.22.22+sha512.a6b2f7906b721bba3d67d4aff083df04dad64c399707841b7acf00f6b133b7ac24255f2652fa22ae3534329dc6180534e98d17432037ff6fd140556e2bb3137e"

diff --git a/src/commands/runPrompt.ts b/src/commands/runPrompt.ts
@@ -9,6 +9,8 @@ import { verifyHasOpenAIKey } from "./setOpenAIKey";
 import { getCredential } from "../utils/credential";
 import { setProjectDir } from "./setProjectDir";
 import { postToBackendSocket } from "../utils/ddSocket";
+import { extensionOutput } from "../extension";
+import { randomUUID } from "crypto";
 
 type PromptOption = 'local-dir' | 'local-file' | 'remote';
 
@@ -132,13 +134,15 @@ export const runPrompt: (secrets: vscode.SecretStorage, mode: PromptOption) => v
         return vscode.window.showErrorMessage("No project path set. Please set the project path in settings or run a local prompt from a workspace.");
     }
 
+    const hostDir = runningLocal ? inputWorkspace! : workspaceFolder!.uri.fsPath;
+
     progress.report({ increment: 5, message: "Checking for project path..." });
 
     progress.report({ increment: 5, message: "Writing prompt output file..." });
 
     const apiKey = await secrets.get("openAIKey");
 
-    const { editor, doc } = await createOutputBuffer();
+    const { editor, doc } = await createOutputBuffer('prompt-output' + randomUUID() + '.md', hostDir);
 
     if (!editor || !doc) {
         postToBackendSocket({ event: 'eventLabsPromptError', properties: { error: 'No editor or document found' } });
@@ -168,27 +172,24 @@ export const runPrompt: (secrets: vscode.SecretStorage, mode: PromptOption) => v
         progress.report({ increment: 5, message: "Running..." });
         const ranges: Record<string, vscode.Range> = {};
         const getBaseFunctionRange = () => new vscode.Range(doc.lineCount, 0, doc.lineCount, 0);
-        await spawnPromptImage(promptOption.id, runningLocal ? inputWorkspace! : workspaceFolder!.uri.fsPath, Username || 'vscode-user', Password, process.platform, async (json) => {
+        await spawnPromptImage(promptOption.id, hostDir, Username || 'vscode-user', Password, process.platform, async (json) => {
+            extensionOutput.appendLine(JSON.stringify(json))
             switch (json.method) {
                 case 'functions':
-                    const functions = json.params;
-                    for (const func of functions) {
-                        const { id, function: { arguments: args } } = func;
-
-                        const params_str = args;
-                        let functionRange = ranges[id] || getBaseFunctionRange();
-                        if (functionRange.isSingleLine) {
-                            // Add function to the end of the file and update the range
-                            await writeToEditor(`\`\`\`json\n${params_str}`);
-                            functionRange = new vscode.Range(functionRange.start.line, functionRange.start.character, doc.lineCount, 0);
-                        }
-                        else {
-                            // Replace existing function and update the range
-                            await writeToEditor(params_str, functionRange);
-                            functionRange = new vscode.Range(functionRange.start.line, functionRange.start.character, functionRange.end.line + params_str.split('\n').length, 0);
-                        }
-                        ranges[id] = functionRange;
+                    const { id, function: { arguments: args } } = json.params;
+                    const params_str = args;
+                    let functionRange = ranges[id] || getBaseFunctionRange();
+                    if (functionRange.isSingleLine) {
+                        // Add function to the end of the file and update the range
+                        await writeToEditor(`\`\`\`json\n${params_str}`);
+                        functionRange = new vscode.Range(functionRange.start.line, functionRange.start.character, doc.lineCount, 0);
+                    }
+                    else {
+                        // Replace existing function and update the range
+                        await writeToEditor(params_str, functionRange);
+                        functionRange = new vscode.Range(functionRange.start.line, functionRange.start.character, functionRange.end.line + params_str.split('\n').length, 0);
                     }
+                    ranges[id] = functionRange;
                     break;
                 case 'functions-done':
                     await writeToEditor('\n```\n\n');
@@ -199,7 +200,7 @@ export const runPrompt: (secrets: vscode.SecretStorage, mode: PromptOption) => v
                     await writeToEditor(`${header} ROLE ${role}${content ? ` (${content})` : ''}\n\n`);
                     break;
                 case 'functions-done':
-                    await writeToEditor(json.params.content);
+                    await writeToEditor(json.params.content+'\n\n');
                     break;
                 case 'message':
                     await writeToEditor(json.params.content);
@@ -219,13 +220,15 @@ export const runPrompt: (secrets: vscode.SecretStorage, mode: PromptOption) => v
                     await writeToEditor(json.params.messages.map((m: any) => `# ${m.role}\n${m.content}`).join('\n') + '\n');
                     break;
                 case 'error':
-                    await writeToEditor('```error\n' + json.params.content + '\n```\n');
-                    postToBackendSocket({ event: 'eventLabsPromptError', properties: { error: json.params.content } });
+                    const errorMSG = String(json.params.content) + String(json.params.message)
+                    await writeToEditor('```error\n' + errorMSG + '\n```\n');
+                    postToBackendSocket({ event: 'eventLabsPromptError', properties: { error: errorMSG } });
                     break;
                 default:
                     await writeToEditor(JSON.stringify(json, null, 2));
             }
         }, token);
+        await doc.save();
     } catch (e: unknown) {
         void vscode.window.showErrorMessage("Error running prompt");
         await writeToEditor('```json\n' + (e as Error).toString() + '\n```');

diff --git a/src/extension.ts b/src/extension.ts
@@ -1,7 +1,7 @@
 import * as vscode from 'vscode';
 import { setOpenAIKey } from './commands/setOpenAIKey';
 import { nativeClient } from './utils/lsp';
-import { spawnSync } from 'child_process';
+import { spawn, spawnSync } from 'child_process';
 import semver from 'semver';
 import commands from './commands';
 import { postToBackendSocket, setDefaultProperties } from './utils/ddSocket';
@@ -20,6 +20,8 @@ export const extensionId = 'docker.labs-ai-tools-vscode';
 
 export const packageJSON = vscode.extensions.getExtension(extensionId)?.packageJSON;
 
+export const extensionOutput = vscode.window.createOutputChannel('Docker Labs AI', 'json')
+
 
 const getLatestVersion = async () => {
 	const resp = (await fetch(
@@ -82,7 +84,17 @@ export async function activate(context: vscode.ExtensionContext) {
 	});
 	context.subscriptions.push(setOpenAIKeyCommand);
 
-	spawnSync('docker', ['pull', "vonwig/prompts:latest"]);
+	const pullPromptImage = () => {
+		const process = spawn('docker', ['pull', "vonwig/prompts:latest"]);
+		process.stdout.on('data', (data) => {
+			console.error(data.toString());
+		});
+		process.stderr.on('data', (data) => {
+			console.error(data.toString());
+		});
+	}
+
+	pullPromptImage();
 
 	const registeredCommands = commands(context)
 

diff --git a/src/utils/notifications.ts b/src/utils/notifications.ts
@@ -0,0 +1,9 @@
+import * as rpc from 'vscode-jsonrpc/node';
+
+export const notifications = {
+    message: new rpc.NotificationType<{ content: string }>('message'),
+    error: new rpc.NotificationType<{ content: string }>('error'),
+    functions: new rpc.NotificationType<{ function: { arguments: string, name: string }, id: string }>('functions'),
+    "functions-done": new rpc.NotificationType<{ id: string, function: { name: string, arguments: string } }>('functions-done'),
+    start: new rpc.NotificationType<{ id: string, function: { name: string, arguments: string } }>('start'),
+}
diff --git a/src/utils/promptFilename.ts b/src/utils/promptFilename.ts
@@ -1,7 +1,15 @@
 import * as vscode from "vscode";
 
-export const createOutputBuffer = async () => {
-    const doc = await vscode.workspace.openTextDocument({ language: 'markdown' });
+export const createOutputBuffer = async (fileName: string, hostDir: string) => {
+    const edit = new vscode.WorkspaceEdit();
+
+    const newURI = vscode.Uri.file(`${hostDir}/${fileName}`);
+
+    edit.createFile(newURI, { ignoreIfExists: true });
+
+    await vscode.workspace.applyEdit(edit);
+
+    const doc = await vscode.workspace.openTextDocument(newURI);
 
     const editor = await vscode.window.showTextDocument(doc, vscode.ViewColumn.Beside);
 

diff --git a/src/utils/promptRunner.ts b/src/utils/promptRunner.ts
@@ -1,7 +1,9 @@
 import { spawn } from "child_process";
 import { CancellationToken, commands, window, workspace } from "vscode";
 import { setThreadId } from "../commands/setThreadId";
-const output = window.createOutputChannel("Docker Labs: AI Tools");
+import * as rpc from 'vscode-jsonrpc/node';
+import { notifications } from "./notifications";
+import { extensionOutput } from "../extension";
 
 export const getRunArgs = async (promptRef: string, projectDir: string, username: string, pat: string, platform: string, render = false) => {
     const isLocal = promptRef.startsWith('local://');
@@ -44,103 +46,78 @@ export const getRunArgs = async (promptRef: string, projectDir: string, username
     return [...baseArgs, ...mountArgs, ...runArgs];
 };
 
-// const anonymizePAT = (args: string[]) => {
-//     if (!args.includes('--pat')) {
-//         return args
-//     }
-//     const patIndex = args.indexOf('--pat')
-//     const newArgs = [...args]
-//     newArgs[patIndex + 1] = args[patIndex + 1].slice(0, 10) + '****'
-//     return newArgs
-// }
-
-const runAndStream = async (command: string, args: string[], callback: (json: any) => Promise<any>, token?: CancellationToken) => {
-    // const argsWithPrivatePAT = anonymizePAT(args)
-    // output.appendLine(`Running ${command} with args ${argsWithPrivatePAT.join(' ')}`);
-    const child = spawn(command, args);
-    if (token) {
-        token.onCancellationRequested(() => {
-            child.kill()
-        })
-    }
-    let out: string[] = [];
-    let processing = false
-    const processSTDOUT = async (callback: (json: {}) => Promise<void>) => {
-        processing = true
-        while (out.length) {
-            const last = out.shift()!
-            let json;
-            try {
-                json = JSON.parse(last);
-            } catch (e) {
-                console.error(`Failed to parse JSON: ${last}, ${e}`)
-                callback({ method: 'error', params: { message: 'Error occured parsing JSON RPC. Please see error console.' } })
-                child.kill();
-            }
-            await callback(json);
+export const spawnPromptImage = async (promptArg: string, projectDir: string, username: string, platform: string, pat: string, callback: (json: any) => Promise<void>, token: CancellationToken) => {
+    const args = await getRunArgs(promptArg!, projectDir!, username, platform, pat);
+    callback({ method: 'message', params: { debug: `Running ${args.join(' ')}` } });
+    const childProcess = spawn("docker", args);
+
+    let connection = rpc.createMessageConnection(
+        new rpc.StreamMessageReader(childProcess.stdout),
+        new rpc.StreamMessageWriter(childProcess.stdin)
+    );
+
+    const notificationBuffer: { method: string, params: object }[] = []
+
+    let processingBuffer = false;
+
+    const processBuffer = async () => {
+        processingBuffer = true;
+        while (notificationBuffer.length > 0) {
+            await callback(notificationBuffer.shift());
         }
-        processing = false;
+        processingBuffer = false;
     }
 
-    const onChildSTDIO = async ({ stdout, stderr }: { stdout: string; stderr: string | null }) => {
-        if (stdout && stdout.startsWith('Content-Length:')) {
-            /**
-             * 
-                Content-Length: 61{}
-             * 
-             */
-            const messages = stdout.split('Content-Length: ').filter(Boolean)
-            const messagesJSON = messages.map(m => m.slice(m.indexOf('{')))
-            out.push(...messagesJSON)
-            if (!processing && out.length) {
-                await processSTDOUT(callback)
-            }
-        }
-        else if (stderr) {
-            callback({ method: 'error', params: { message: stderr } });
-        }
-        else {
-            callback({ method: 'message', params: { content: stdout } });
+
+    const pushNotification = (method: string, params: object) => {
+        notificationBuffer.push({ method, params });
+        if (!processingBuffer) {
+            processBuffer();
         }
-    };
-    return new Promise((resolve, reject) => {
-        child.stdout.on('data', (data) => {
-            onChildSTDIO({ stdout: data.toString(), stderr: '' });
-        });
-        child.stderr.on('data', (data) => {
-            onChildSTDIO({ stderr: data.toString(), stdout: '' });
-        });
-        child.on('close', (code) => {
-            callback({ method: 'message', params: { debug: `child process exited with code ${code}` } });
-            resolve(code);
-        });
-        child.on('error', (err) => {
-            callback({ method: 'error', params: { message: JSON.stringify(err) } });
-            reject(err);
-        });
+    }
+
+    for (const [type, properties] of Object.entries(notifications)){
+        // @ts-expect-error
+        connection.onNotification(properties, (params)=> pushNotification(type, params))
+    }
+
+    connection.listen();
+
+    token.onCancellationRequested(() => {
+        childProcess.kill();
+        connection.dispose();
     });
-};
 
-export const spawnPromptImage = async (promptArg: string, projectDir: string, username: string, platform: string, pat: string, callback: (json: any) => Promise<void>, token: CancellationToken) => {
-    const args = await getRunArgs(promptArg!, projectDir!, username, platform, pat);
-    callback({ method: 'message', params: { debug: `Running ${args.join(' ')}` } });
-    return runAndStream("docker", args, callback, token);
+
 };
 
 export const writeKeyToVolume = async (key: string) => {
 
     const args1 = ["pull", "vonwig/function_write_files"];
+
     const args2 = [
         "run",
-        "-v",
-        "openai_key:/secret",
+        "-v", "openai_key:/secret",
+        "--rm",
         "--workdir", "/secret",
         "vonwig/function_write_files",
         `'` + JSON.stringify({ files: [{ path: ".openai-api-key", content: key, executable: false }] }) + `'`
     ];
-    const callback = async (json: any) => {
-        output.appendLine(JSON.stringify(json, null, 2));
-    };
-    await runAndStream("docker", args1, callback);
-    await runAndStream("docker", args2, callback);
+
+    const child1 = spawn("docker", args1);
+
+    child1.stdout.on('data', (data) => {
+        extensionOutput.appendLine(data.toString());
+    });
+    child1.stderr.on('data', (data) => {
+        extensionOutput.appendLine(data.toString());
+    });
+
+    const child2 = spawn("docker", args2);
+    child2.stdout.on('data', (data) => {
+        extensionOutput.appendLine(data.toString());
+    });
+    child2.stderr.on('data', (data) => {
+        extensionOutput.appendLine(data.toString());
+    });
 };