vscode-pi-companion 0.1.2 → 0.1.3

package/extensions/index.ts

@@ -1,8 +1,7 @@
 /**
  * PI Companion Extension
- * 
+ *
  * Connects to the VS Code extension via SSE for real-time context updates.
- * No polling, instant updates, no tool calls visible.
  */
 
 import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
@@ -13,78 +12,11 @@ import * as http from "http";
 import { execSync } from "child_process";
 import { Type } from "@sinclair/typebox";
 
-function isLocalPackageInstall(): boolean {
-  const filename = __filename;
-  return !filename.includes(`${path.sep}.pi${path.sep}agent${path.sep}git${path.sep}`)
-    && !filename.includes(`${path.sep}.pi${path.sep}git${path.sep}`);
-}
-
-// Check if running in VS Code integrated terminal
-function isRunningInVSCode(): boolean {
-  return process.env.TERM_PROGRAM === 'vscode';
-}
-
-// Install VS Code extension from marketplace
-function installFromMarketplace(ctx: ExtensionContext): void {
-  try {
-    execSync('code --install-extension ravshansbox.vscode-pi-companion --force', { stdio: 'pipe' });
-    ctx.ui.notify("PI Companion: VS Code extension installed. Reload window to activate.", "info");
-  } catch {
-    // Silently ignore
-  }
-}
-
-// Check if extension is installed
-function isExtensionInstalled(): boolean {
-  try {
-    const result = execSync('code --list-extensions', { encoding: 'utf8' });
-    return result.includes('ravshansbox.vscode-pi-companion');
-  } catch {
-    return false;
-  }
-}
-
-// Install/update VS Code extension if running in VS Code terminal
-async function installVSCodeExtension(ctx: ExtensionContext): Promise<void> {
-  if (!isRunningInVSCode()) {
-    return;
-  }
-  
-  const isInstalled = isExtensionInstalled();
-  
-  if (isInstalled) {
-    // Already installed - check for updates, but skip for local package installs
-    if (!isLocalPackageInstall()) {
-      const latestVersion = await getLatestVersion();
-      if (latestVersion) {
-        console.log(`[PI Companion] Update available: v${latestVersion}`);
-        console.log(`[PI Companion] Run: pi install git:github.com/ravshansbox/vscode-pi-companion`);
-      }
-    }
-    return;
-  }
-  
-  // Not installed - install from marketplace
-  installFromMarketplace(ctx);
-}
-
-// Get latest version from GitHub
-async function getLatestVersion(): Promise<string | null> {
-  try {
-    const response = await fetch('https://api.github.com/repos/ravshansbox/vscode-pi-companion/releases/latest', {
-      headers: { 'Accept': 'application/vnd.github.v3+json' }
-    });
-    
-    if (!response.ok) {
-      return null;
-    }
-    
-    const release = await response.json() as { tag_name?: string };
-    return release.tag_name?.replace('v', '') || null;
-  } catch {
-    return null;
-  }
-}
+const CONNECTION_DIR = path.join(os.tmpdir(), "pi-companion");
+const WIDGET_DEBOUNCE_MS = 50;
+const RECONNECT_DELAY_MS = 2000;
+const PORT_VALIDATE_TIMEOUT_MS = 750;
+const MAX_INJECTED_LINES = 100;
 
 interface IdeContext {
   workspaceState: {
@@ -103,6 +35,12 @@ interface IdeContext {
   };
 }
 
+interface ConnectionInfo {
+  port: number;
+  workspacePath: string;
+  authToken: string;
+}
+
 let currentContext: IdeContext | undefined;
 let currentCtx: ExtensionContext | undefined;
 let lastWidgetKey = "";
@@ -110,34 +48,95 @@ let widgetDebounceTimer: NodeJS.Timeout | undefined;
 let reconnectTimer: NodeJS.Timeout | undefined;
 let sseRequest: http.ClientRequest | undefined;
 
-function getPort(): number | undefined {
-  const configDir = path.join(os.tmpdir(), "pi-companion");
+function isRunningInVSCode(): boolean {
+  return process.env.TERM_PROGRAM === "vscode";
+}
 
+function isExtensionInstalled(): boolean {
   try {
-    const files = fs.readdirSync(configDir);
+    const result = execSync("code --list-extensions", { encoding: "utf8" });
+    return result.includes("ravshansbox.vscode-pi-companion");
+  } catch {
+    return false;
+  }
+}
+
+function installFromMarketplace(ctx: ExtensionContext): void {
+  try {
+    execSync("code --install-extension ravshansbox.vscode-pi-companion --force", { stdio: "pipe" });
+    ctx.ui.notify("PI Companion: VS Code extension installed. Reload window to activate.", "info");
+  } catch {
+    ctx.ui.notify("PI Companion: Failed to auto-install VS Code extension.", "warning");
+  }
+}
+
+function ensureVSCodeExtensionInstalled(ctx: ExtensionContext): void {
+  if (!isRunningInVSCode()) return;
+  if (!isExtensionInstalled()) {
+    installFromMarketplace(ctx);
+  }
+}
+
+function readLatestConnectionInfo(): ConnectionInfo | undefined {
+  try {
+    const files = fs.readdirSync(CONNECTION_DIR);
     const connectionFiles = files
       .filter((f) => f.startsWith("connection-") && f.endsWith(".json"))
       .map((f) => ({
-        name: f,
-        path: path.join(configDir, f),
-        mtime: fs.statSync(path.join(configDir, f)).mtime.getTime(),
+        path: path.join(CONNECTION_DIR, f),
+        mtime: fs.statSync(path.join(CONNECTION_DIR, f)).mtime.getTime(),
       }))
       .sort((a, b) => b.mtime - a.mtime);
 
-    if (connectionFiles.length > 0) {
-      const content = fs.readFileSync(connectionFiles[0].path, "utf-8");
-      const info = JSON.parse(content);
-      return info.port;
+    for (const file of connectionFiles) {
+      try {
+        const content = fs.readFileSync(file.path, "utf-8");
+        const info = JSON.parse(content) as ConnectionInfo;
+        if (typeof info.port === "number" && info.port > 0) {
+          return info;
+        }
+      } catch {
+        // Ignore malformed files and try the next one.
+      }
     }
   } catch {
-    // No connection files found
+    // No connection files found.
   }
   return undefined;
 }
 
+function isPortAlive(port: number): Promise<boolean> {
+  return new Promise((resolve) => {
+    const req = http.get(
+      {
+        hostname: "127.0.0.1",
+        port,
+        path: "/context",
+        timeout: PORT_VALIDATE_TIMEOUT_MS,
+      },
+      (res) => {
+        res.resume();
+        resolve(res.statusCode === 200);
+      },
+    );
+
+    req.on("timeout", () => {
+      req.destroy();
+      resolve(false);
+    });
+    req.on("error", () => resolve(false));
+  });
+}
+
+async function getLivePort(): Promise<number | undefined> {
+  const info = readLatestConnectionInfo();
+  if (!info) return undefined;
+  return (await isPortAlive(info.port)) ? info.port : undefined;
+}
+
 function updateWidget(): void {
   if (!currentCtx) return;
-  
+
   if (!currentContext || !currentContext.workspaceState.openFiles.length) {
     currentCtx.ui.setWidget("pi-companion", ["PI Companion: No files open"]);
     return;
@@ -145,25 +144,16 @@ function updateWidget(): void {
 
   const { workspaceState } = currentContext;
   const lines: string[] = [];
-
   const activeFile = workspaceState.openFiles.find((f) => f.isActive);
 
   if (activeFile) {
     const filename = activeFile.path.split("/").pop() || activeFile.path;
-    
-    // Build location info - ONLY show when text is actually selected
     let location = "";
     if (activeFile.selectionStart && activeFile.selectionEnd) {
       const start = activeFile.selectionStart.line;
       const end = activeFile.selectionEnd.line;
-      if (start === end) {
-        location = ` L${start}`;
-      } else {
-        location = ` L${start}:${end}`;
-      }
+      location = start === end ? ` L${start}` : ` L${start}:${end}`;
     }
-    // Don't show cursor position when nothing is selected
-    
     lines.push(`📄 ${filename}${location}`);
   }
 
@@ -173,129 +163,105 @@ function updateWidget(): void {
   }
 
   const widgetKey = lines.join("|");
-  
-  // Debounce widget updates to avoid flickering during rapid selection changes
-  if (widgetDebounceTimer) {
-    clearTimeout(widgetDebounceTimer);
-  }
-  
+  if (widgetDebounceTimer) clearTimeout(widgetDebounceTimer);
+
   widgetDebounceTimer = setTimeout(() => {
     if (widgetKey !== lastWidgetKey) {
       lastWidgetKey = widgetKey;
       currentCtx!.ui.setWidget("pi-companion", lines);
     }
     widgetDebounceTimer = undefined;
-  }, 50); // Wait 50ms after last update
+  }, WIDGET_DEBOUNCE_MS);
 }
 
-function connectSSE(): void {
-  const port = getPort();
+async function connectSSE(): Promise<void> {
+  const port = await getLivePort();
   if (!port) return;
 
-  // Close existing connection
   if (sseRequest) {
     sseRequest.destroy();
     sseRequest = undefined;
   }
 
-  const url = new URL(`http://127.0.0.1:${port}/context/stream`);
-  
-  sseRequest = http.get({
-    hostname: url.hostname,
-    port: url.port,
-    path: url.pathname,
-  }, (res) => {
-    let buffer = "";
-    
-    res.on("data", (chunk: string) => {
-      buffer += chunk;
-      
-      // Process complete SSE messages
-      const lines = buffer.split("\n");
-      buffer = lines.pop() || "";
-      
-      for (const line of lines) {
-        if (line.startsWith("data: ")) {
+  sseRequest = http.get(
+    {
+      hostname: "127.0.0.1",
+      port,
+      path: "/context/stream",
+    },
+    (res) => {
+      let buffer = "";
+
+      res.on("data", (chunk: string) => {
+        buffer += chunk;
+        const lines = buffer.split("\n");
+        buffer = lines.pop() || "";
+
+        for (const line of lines) {
+          if (!line.startsWith("data: ")) continue;
           const data = line.slice(6);
-          if (data && data !== "[object Object]") {
-            try {
-              const parsed = JSON.parse(data);
-              currentContext = parsed;
-              updateWidget();
-            } catch {
-              // Ignore parse errors
-            }
+          if (!data || data === "[object Object]") continue;
+          try {
+            currentContext = JSON.parse(data);
+            updateWidget();
+          } catch {
+            // Ignore malformed SSE payloads.
           }
         }
-      }
-    });
+      });
 
-    res.on("end", () => {
-      scheduleReconnect();
-    });
-
-    res.on("error", () => {
-      scheduleReconnect();
-    });
-  });
+      res.on("end", scheduleReconnect);
+      res.on("error", scheduleReconnect);
+    },
+  );
 
-  sseRequest.on("error", () => {
-    scheduleReconnect();
-  });
+  sseRequest.on("error", scheduleReconnect);
 }
 
 function scheduleReconnect(): void {
   if (reconnectTimer || !currentCtx) return;
   reconnectTimer = setTimeout(() => {
     reconnectTimer = undefined;
-    connectSSE();
-  }, 2000);
+    void connectSSE();
+  }, RECONNECT_DELAY_MS);
 }
 
 export default function (pi: ExtensionAPI) {
-  // Inject IDE context before each prompt
-  pi.on("before_agent_start", async (event, ctx) => {
+  pi.on("before_agent_start", async () => {
     if (!currentContext) return;
-    
+
     const activeFile = currentContext.workspaceState.openFiles.find((f) => f.isActive);
     if (!activeFile) return;
 
     let contextText = "";
-    
+
     if (activeFile.selectedText) {
-      // Inject selected text
       contextText = `Current selection in ${activeFile.path}:\n\`\`\`\n${activeFile.selectedText}\n\`\`\``;
     } else if (activeFile.content) {
-      // Inject full file content for small files, or first 100 lines
       const lines = activeFile.content.split("\n");
-      if (lines.length <= 100) {
-        contextText = `Current file ${activeFile.path}:\n\`\`\`\n${activeFile.content}\n\`\`\``;
-      } else {
-        contextText = `Current file ${activeFile.path} (first 100 lines):\n\`\`\`\n${lines.slice(0, 100).join("\n")}\n\`\`\``;
-      }
+      contextText = lines.length <= MAX_INJECTED_LINES
+        ? `Current file ${activeFile.path}:\n\`\`\`\n${activeFile.content}\n\`\`\``
+        : `Current file ${activeFile.path} (first ${MAX_INJECTED_LINES} lines):\n\`\`\`\n${lines.slice(0, MAX_INJECTED_LINES).join("\n")}\n\`\`\``;
     }
 
-    if (contextText) {
-      return {
-        message: {
-          customType: "pi-companion",
-          content: contextText,
-          display: false, // Don't show in chat history
-        },
-      };
-    }
+    if (!contextText) return;
+    return {
+      message: {
+        customType: "pi-companion",
+        content: contextText,
+        display: false,
+      },
+    };
   });
 
   pi.on("session_start", async (_event, ctx) => {
     currentCtx = ctx;
+    ensureVSCodeExtensionInstalled(ctx);
 
-    // Auto-install/update VS Code extension if needed
-    installVSCodeExtension(ctx);
-
-    const port = getPort();
+    const port = await getLivePort();
     if (port) {
       ctx.ui.notify("PI Companion: Connected to VS Code", "info");
-      connectSSE();
+      await connectSSE();
     } else {
       ctx.ui.notify("PI Companion: VS Code companion server not running. Run 'PI Companion: Start Server' in VS Code.", "warning");
     }
@@ -312,7 +278,6 @@ export default function (pi: ExtensionAPI) {
     }
   });
 
-  // Register tools
   pi.registerTool({
     name: "get_ide_context",
     label: "Get IDE Context",
@@ -346,12 +311,15 @@ export default function (pi: ExtensionAPI) {
         return { content: [{ type: "text", text: "No active file" }], details: {} };
       }
       return {
-        content: [{ type: "text", text: JSON.stringify({
-          path: activeFile.path,
-          cursor: activeFile.cursor,
-          selection: activeFile.selectedText,
-          content: activeFile.content,
-        }, null, 2) }],
+        content: [{
+          type: "text",
+          text: JSON.stringify({
+            path: activeFile.path,
+            cursor: activeFile.cursor,
+            selection: activeFile.selectedText,
+            content: activeFile.content,
+          }, null, 2),
+        }],
         details: {},
       };
     },

package/package.json

@@ -3,7 +3,7 @@
   "displayName": "PI Companion",
   "description": "Share VS Code context with pi coding agent",
   "repository": "https://github.com/ravshansbox/vscode-pi-companion",
-  "version": "0.1.2",
+  "version": "0.1.3",
   "publisher": "ravshansbox",
   "keywords": [
     "pi-package"