Feature/vscode 0326 (UKGovernmentBEIS#193)

* decorate the html tag for the logview so it can detect vscode immediately

* Properly preserve focus when showing log viewer

src/inspect_ai/_view/www/App.mjs

@@ -344,7 +344,6 @@ export function App() {
     if (showFind) {
       setShowFind(false);
     }
-    mainAppRef.current.focus();
   }, [showFind, setShowFind]);
 
   return html`

tools/vscode/CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## 0.3.26
+
+- Properly preserve focus When showing the log viewer upon task completion.
+
 ## 0.3.25
 
 - Poll more frequently to show new log files

tools/vscode/src/components/focus.ts

@@ -1,4 +1,4 @@
-import { commands, window } from "vscode";
+import { commands, ExtensionContext, window } from "vscode";
 
 export function scheduleReturnFocus(command: string) {
   setTimeout(() => {
@@ -14,3 +14,85 @@ export function scheduleFocusActiveEditor() {
     }
   }, 200);
 }
+
+export class FocusManager {
+  private lastFocused: 'editor' | 'terminal' | 'notebook' | 'none' = 'none';
+
+  constructor(context: ExtensionContext) {
+    this.initialize(context);
+  }
+
+  private initialize(context: ExtensionContext) {
+    // Track editor focus changes
+    context.subscriptions.push(
+      window.onDidChangeActiveTextEditor((editor) => {
+        if (editor) {
+          this.lastFocused = 'editor';
+        }
+      })
+    );
+
+    // Track terminal focus changes
+    context.subscriptions.push(
+      window.onDidOpenTerminal(async (terminal) => {
+        const pid = await terminal.processId;
+        if (window.activeTerminal?.processId === pid) {
+          this.lastFocused = 'terminal';
+        }
+      })
+    );
+
+    // Handle terminal focus changes (when terminal is in focus and user types)
+    context.subscriptions.push(
+      window.onDidChangeTerminalState((terminal) => {
+        if (terminal.state.isInteractedWith) {
+          this.lastFocused = 'terminal';
+        }
+      })
+    );
+
+    // Handle when terminal becomes active
+    context.subscriptions.push(
+      window.onDidChangeActiveTerminal((terminal) => {
+        if (terminal) {
+          this.lastFocused = 'terminal';
+        }
+      })
+    );
+
+    // Track when window focus changes to ensure robustness
+    context.subscriptions.push(
+      window.onDidChangeWindowState((windowState) => {
+        if (windowState.focused) {
+          if (window.activeTextEditor) {
+            this.lastFocused = 'editor';
+          } else if (window.activeTerminal) {
+            this.lastFocused = 'terminal';
+          }
+        }
+      })
+    );
+
+    // Track when editors gain or lose focus
+    context.subscriptions.push(
+      window.onDidChangeTextEditorSelection((e) => {
+        if (e.textEditor === window.activeTextEditor) {
+          this.lastFocused = 'editor';
+        }
+      })
+    );
+
+    // Track notebook editor focus changes
+    context.subscriptions.push(
+      window.onDidChangeActiveNotebookEditor((notebookEditor) => {
+        if (notebookEditor) {
+          this.lastFocused = 'notebook';
+        }
+      })
+    );
+  }
+
+  public getLastFocused(): 'editor' | 'terminal' | 'notebook' | 'none' {
+    return this.lastFocused;
+  }
+}

tools/vscode/src/components/webview.ts

@@ -1,8 +1,18 @@
-import { Uri, ViewColumn, EventEmitter, ExtensionContext, window } from "vscode";
+import {
+  Uri,
+  ViewColumn,
+  EventEmitter,
+  ExtensionContext,
+  window,
+  commands,
+} from "vscode";
 
 import { Disposable } from "../core/dispose";
 import { getNonce } from "../core/nonce";
 import { ExtensionHost, HostWebviewPanel } from "../hooks";
+import { isNotebook } from "./notebook";
+import { FocusManager } from "./focus";
+import { log } from "../core/log";
 
 export interface ShowOptions {
   readonly preserveFocus?: boolean;
@@ -35,7 +45,10 @@ export class InspectWebviewManager<T extends InspectWebview<S>, S> {
         },
       })
     );
+
+    this.focusManager_ = new FocusManager(context);
   }
+  private focusManager_: FocusManager;
 
   public setOnShow(f: () => void) {
     this.onShow_ = f;
@@ -90,10 +103,47 @@ export class InspectWebviewManager<T extends InspectWebview<S>, S> {
   }
 
   private preserveEditorFocus() {
-    // No need to take action here - we are already setting preserveFocus
-    // and ensuring that focus ends up in the correct places
+    // Replace focus to the correct spot
+    const lastFocused = this.focusManager_.getLastFocused();
+    if (lastFocused === "terminal") {
+      // The terminal
+      setTimeout(() => {
+        commands.executeCommand('workbench.action.terminal.focus').then(
+          () => {
+            // Command executed successfully
+          },
+          (error) => {
+            log.append("Couldn't focus terminal.\n" + error);
+          }
+        );
+      }, 50);
+    } else if (lastFocused === "editor") {
+      // The editor
+      const editor = window.activeTextEditor;
+      if (editor) {
+        if (!isNotebook(editor.document.uri)) {
+          setTimeout(() => {
+            // Refocus the active document by calling showTextDocument with the active editor
+            window.showTextDocument(editor.document, editor.viewColumn).then(() => {
+
+            }, (error) => {
+              log.append("Couldn't focus editor.\n" + error);
+            });
+          }, 50);
+        }
+
+      }
+    } else if (lastFocused === "notebook") {
+      // A notebook
+      setTimeout(() => {
+        if (window.activeNotebookEditor) {
+          window.activeNotebookEditor.revealRange(window.activeNotebookEditor.selection);
+        }
+      }, 50);
+    }
   }
 
+
   private restoreWebview(panel: HostWebviewPanel, state: S): void {
     const view = new this.webviewType_(this.context, state, panel);
     this.registerWebviewListeners(view);

tools/vscode/src/providers/logview/logview-webview.ts

@@ -269,6 +269,9 @@ class InspectLogviewWebview extends InspectWebview<LogviewState> {
           )}</script>`
           : "";
 
+      // decorate the html tag
+      indexHtml = indexHtml.replace("<html ", '<html class="vscode" ');
+
       // add content security policy
       indexHtml = indexHtml.replace(
         "<head>\n",