Plug in RPC client

docker · ColinMcNeil · Oct 29, 2024 · Oct 28, 2024 · Oct 28, 2024 · Oct 28, 2024
commit 7d1959edd4b3a7033dbc69bda8a11ff41e84d9c2
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
@@ -171,24 +171,20 @@
         await spawnPromptImage(promptOption.id, runningLocal ? inputWorkspace! : workspaceFolder!.uri.fsPath, Username || 'vscode-user', Password, process.platform, async (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');
@@ -219,8 +215,9 @@
                     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));

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';
@@ -82,7 +82,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,8 @@
+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'),
+    functionsDone: new rpc.NotificationType<{ id: string, function: { name: string, arguments: string } }>('functions-done'),
+}
diff --git a/src/utils/promptRunner.ts b/src/utils/promptRunner.ts
@@ -1,6 +1,9 @@
 import { spawn } from "child_process";
 import { CancellationToken, commands, window, workspace } from "vscode";
 import { setThreadId } from "../commands/setThreadId";
+import * as rpc from 'vscode-jsonrpc/node';
+import { notifications } from "./notifications";
+
 const output = window.createOutputChannel("Docker Labs: AI Tools");
 
 export const getRunArgs = async (promptRef: string, projectDir: string, username: string, pat: string, platform: string, render = false) => {
@@ -44,87 +47,49 @@ 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);
-        });
+    }
+
+    connection.onNotification(notifications.message, (params) => pushNotification('message', params));
+    connection.onNotification(notifications.error, (params) => pushNotification('error', params));
+    connection.onNotification(notifications.functions, (params) => pushNotification('functions', params));
+    connection.onNotification(notifications.functionsDone, (params) => pushNotification('functions-done', 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) => {
@@ -133,14 +98,26 @@ export const writeKeyToVolume = async (key: string) => {
     const args2 = [
         "run",
         "-v",
+        "--rm",
         "openai_key:/secret",
         "--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) => {
+        output.appendLine(data.toString());
+    });
+    child1.stderr.on('data', (data) => {
+        output.appendLine(data.toString());
+    });
+
+    const child2 = spawn("docker", args2);
+    child2.stdout.on('data', (data) => {
+        output.appendLine(data.toString());
+    });
+    child2.stderr.on('data', (data) => {
+        output.appendLine(data.toString());
+    });
 };