diff --git a/packages/vscode-ide-companion/src/diff-manager.ts b/packages/vscode-ide-companion/src/diff-manager.ts
index 9c7afc1d..abbc6eec 100644
--- a/packages/vscode-ide-companion/src/diff-manager.ts
+++ b/packages/vscode-ide-companion/src/diff-manager.ts
@@ -214,10 +214,18 @@ export class DiffManager {
   }
 
   private async onActiveEditorChange(editor: vscode.TextEditor | undefined) {
-    const isVisible =
-      !!editor &&
-      editor.document.uri.scheme === DIFF_SCHEME &&
-      this.diffDocuments.has(editor.document.uri.toString());
+    let isVisible = false;
+    if (editor) {
+      isVisible = this.diffDocuments.has(editor.document.uri.toString());
+      if (!isVisible) {
+        for (const document of this.diffDocuments.values()) {
+          if (document.originalFilePath === editor.document.uri.fsPath) {
+            isVisible = true;
+            break;
+          }
+        }
+      }
+    }
     await vscode.commands.executeCommand(
       'setContext',
       'gemini.diff.isVisible',