[ide-mode] Add support for in-IDE diff handling in the CLI (#5603)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

2 files changed, 49 insertions, 40 deletions

diff --git a/packages/vscode-ide-companion/src/diff-manager.ts b/packages/vscode-ide-companion/src/diff-manager.ts
index 159a6101..0dad03a6 100644
--- a/packages/vscode-ide-companion/src/diff-manager.ts
+++ b/packages/vscode-ide-companion/src/diff-manager.ts
@@ -4,10 +4,14 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import * as vscode from 'vscode';
+import {
+  IdeDiffAcceptedNotificationSchema,
+  IdeDiffClosedNotificationSchema,
+} from '@google/gemini-cli-core';
+import { type JSONRPCNotification } from '@modelcontextprotocol/sdk/types.js';
 import * as path from 'node:path';
+import * as vscode from 'vscode';
 import { DIFF_SCHEME } from './extension.js';
-import { type JSONRPCNotification } from '@modelcontextprotocol/sdk/types.js';
 
 export class DiffContentProvider implements vscode.TextDocumentContentProvider {
   private content = new Map<string, string>();
@@ -126,18 +130,19 @@ export class DiffManager {
       const rightDoc = await vscode.workspace.openTextDocument(uriToClose);
       const modifiedContent = rightDoc.getText();
       await this.closeDiffEditor(uriToClose);
-      this.onDidChangeEmitter.fire({
-        jsonrpc: '2.0',
-        method: 'ide/diffClosed',
-        params: {
-          filePath,
-          content: modifiedContent,
-        },
-      });
-      vscode.window.showInformationMessage(`Diff for ${filePath} closed.`);
-    } else {
-      vscode.window.showWarningMessage(`No open diff found for ${filePath}.`);
+      this.onDidChangeEmitter.fire(
+        IdeDiffClosedNotificationSchema.parse({
+          jsonrpc: '2.0',
+          method: 'ide/diffClosed',
+          params: {
+            filePath,
+            content: modifiedContent,
+          },
+        }),
+      );
+      return modifiedContent;
     }
+    return;
   }
 
   /**
@@ -156,14 +161,16 @@ export class DiffManager {
     const modifiedContent = rightDoc.getText();
     await this.closeDiffEditor(rightDocUri);
 
-    this.onDidChangeEmitter.fire({
-      jsonrpc: '2.0',
-      method: 'ide/diffAccepted',
-      params: {
-        filePath: diffInfo.originalFilePath,
-        content: modifiedContent,
-      },
-    });
+    this.onDidChangeEmitter.fire(
+      IdeDiffAcceptedNotificationSchema.parse({
+        jsonrpc: '2.0',
+        method: 'ide/diffAccepted',
+        params: {
+          filePath: diffInfo.originalFilePath,
+          content: modifiedContent,
+        },
+      }),
+    );
   }
 
   /**
@@ -184,14 +191,16 @@ export class DiffManager {
     const modifiedContent = rightDoc.getText();
     await this.closeDiffEditor(rightDocUri);
 
-    this.onDidChangeEmitter.fire({
-      jsonrpc: '2.0',
-      method: 'ide/diffClosed',
-      params: {
-        filePath: diffInfo.originalFilePath,
-        content: modifiedContent,
-      },
-    });
+    this.onDidChangeEmitter.fire(
+      IdeDiffClosedNotificationSchema.parse({
+        jsonrpc: '2.0',
+        method: 'ide/diffClosed',
+        params: {
+          filePath: diffInfo.originalFilePath,
+          content: modifiedContent,
+        },
+      }),
+    );
   }
 
   private addDiffDocument(uri: vscode.Uri, diffInfo: DiffInfo) {
diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts
index 30215ccc..eec99cb3 100644
--- a/packages/vscode-ide-companion/src/ide-server.ts
+++ b/packages/vscode-ide-companion/src/ide-server.ts
@@ -5,15 +5,13 @@
  */
 
 import * as vscode from 'vscode';
+import { IdeContextNotificationSchema } from '@google/gemini-cli-core';
+import { isInitializeRequest } from '@modelcontextprotocol/sdk/types.js';
 import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
 import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
-import express, { Request, Response } from 'express';
+import express, { type Request, type Response } from 'express';
 import { randomUUID } from 'node:crypto';
-import {
-  isInitializeRequest,
-  type JSONRPCNotification,
-} from '@modelcontextprotocol/sdk/types.js';
-import { Server as HTTPServer } from 'node:http';
+import { type Server as HTTPServer } from 'node:http';
 import { z } from 'zod';
 import { DiffManager } from './diff-manager.js';
 import { OpenFilesManager } from './open-files-manager.js';
@@ -28,11 +26,12 @@ function sendIdeContextUpdateNotification(
 ) {
   const ideContext = openFilesManager.state;
 
-  const notification: JSONRPCNotification = {
+  const notification = IdeContextNotificationSchema.parse({
     jsonrpc: '2.0',
     method: 'ide/contextUpdate',
     params: ideContext,
-  };
+  });
+
   log(
     `Sending IDE context update notification: ${JSON.stringify(
       notification,
@@ -76,7 +75,7 @@ export class IDEServer {
     });
     context.subscriptions.push(onDidChangeSubscription);
     const onDidChangeDiffSubscription = this.diffManager.onDidChange(
-      (notification: JSONRPCNotification) => {
+      (notification) => {
         for (const transport of Object.values(transports)) {
           transport.send(notification);
         }
@@ -269,12 +268,13 @@ const createMcpServer = (diffManager: DiffManager) => {
       }).shape,
     },
     async ({ filePath }: { filePath: string }) => {
-      await diffManager.closeDiff(filePath);
+      const content = await diffManager.closeDiff(filePath);
+      const response = { content: content ?? undefined };
       return {
         content: [
           {
             type: 'text',
-            text: `Closed diff for ${filePath}`,
+            text: JSON.stringify(response),
           },
         ],
       };