diff --git a/src/services/browser/BrowserSession.ts b/src/services/browser/BrowserSession.ts index 699b8c7315..e7a7b917ac 100644 --- a/src/services/browser/BrowserSession.ts +++ b/src/services/browser/BrowserSession.ts @@ -21,6 +21,7 @@ export class BrowserSession { private page?: Page private currentMousePosition?: string private lastConnectionAttempt?: number + private isUsingRemoteBrowser: boolean = false constructor(context: vscode.ExtensionContext) { this.context = context @@ -70,6 +71,7 @@ export class BrowserSession { defaultViewport: this.getViewport(), // headless: false, }) + this.isUsingRemoteBrowser = false } /** @@ -86,6 +88,7 @@ export class BrowserSession { console.log(`Connected to remote browser at ${chromeHostUrl}`) this.context.globalState.update("cachedChromeHostUrl", chromeHostUrl) this.lastConnectionAttempt = Date.now() + this.isUsingRemoteBrowser = true return true } catch (error) { @@ -192,15 +195,16 @@ export class BrowserSession { if (this.browser || this.page) { console.log("closing browser...") - const remoteBrowserEnabled = this.context.globalState.get("remoteBrowserEnabled") as boolean | undefined - if (remoteBrowserEnabled && this.browser) { + // Use the actual connection type, not just the setting + if (this.isUsingRemoteBrowser && this.browser) { + console.log("Disconnecting from remote browser") await this.browser.disconnect().catch(() => {}) } else { + console.log("Closing local browser") await this.browser?.close().catch(() => {}) - this.resetBrowserState() } - // this.resetBrowserState() + this.resetBrowserState() } return {} } @@ -212,6 +216,7 @@ export class BrowserSession { this.browser = undefined this.page = undefined this.currentMousePosition = undefined + this.isUsingRemoteBrowser = false } async doAction(action: (page: Page) => Promise): Promise { diff --git a/src/services/browser/__tests__/BrowserSession.test.ts b/src/services/browser/__tests__/BrowserSession.test.ts new file mode 100644 index 0000000000..d1a7a81b66 --- /dev/null +++ b/src/services/browser/__tests__/BrowserSession.test.ts @@ -0,0 +1,184 @@ +import * as vscode from "vscode" +import { BrowserSession } from "../BrowserSession" +import { Browser } from "puppeteer-core" + +// Mock vscode +jest.mock("vscode", () => ({ + ExtensionContext: jest.fn(), +})) + +// Mock fs/promises +jest.mock("fs/promises", () => ({ + mkdir: jest.fn().mockResolvedValue(undefined), +})) + +// Mock path +jest.mock("path", () => ({ + join: jest.fn((...args) => args.join("/")), +})) + +// Mock puppeteer-core +jest.mock("puppeteer-core", () => ({ + launch: jest.fn(), + connect: jest.fn(), +})) + +// Mock puppeteer-chromium-resolver +jest.mock("puppeteer-chromium-resolver", () => { + return jest.fn().mockResolvedValue({ + puppeteer: { + launch: jest.fn(), + }, + executablePath: "/mock/chrome/path", + }) +}) + +// Mock browserDiscovery +jest.mock("../browserDiscovery", () => ({ + discoverChromeHostUrl: jest.fn(), + tryChromeHostUrl: jest.fn(), +})) + +// Mock fs utilities +jest.mock("../../../utils/fs", () => ({ + fileExistsAtPath: jest.fn().mockResolvedValue(true), +})) + +// Mock delay and p-wait-for +jest.mock("delay", () => jest.fn().mockResolvedValue(undefined)) +jest.mock("p-wait-for", () => jest.fn().mockResolvedValue(undefined)) + +describe("BrowserSession", () => { + let mockContext: Partial + let mockBrowser: Partial + let browserSession: BrowserSession + let mockLaunch: jest.Mock + let mockConnect: jest.Mock + let mockDiscoverChromeHostUrl: jest.Mock + let mockTryChromeHostUrl: jest.Mock + let mockPCR: jest.Mock + + beforeEach(() => { + // Get the mocked functions + const puppeteerCore = require("puppeteer-core") + mockLaunch = puppeteerCore.launch as jest.Mock + mockConnect = puppeteerCore.connect as jest.Mock + + const browserDiscovery = require("../browserDiscovery") + mockDiscoverChromeHostUrl = browserDiscovery.discoverChromeHostUrl as jest.Mock + mockTryChromeHostUrl = browserDiscovery.tryChromeHostUrl as jest.Mock + + mockPCR = require("puppeteer-chromium-resolver") as jest.Mock + + // Reset all mocks + jest.clearAllMocks() + + // Mock browser instance + mockBrowser = { + close: jest.fn().mockResolvedValue(undefined), + disconnect: jest.fn().mockResolvedValue(undefined), + newPage: jest.fn(), + pages: jest.fn().mockResolvedValue([]), + } + + // Mock vscode context + mockContext = { + globalState: { + get: jest.fn(), + update: jest.fn(), + }, + globalStorageUri: { + fsPath: "/mock/storage/path", + }, + } as any + + // Setup PCR mock + mockPCR.mockResolvedValue({ + puppeteer: { + launch: mockLaunch, + }, + executablePath: "/mock/chrome/path", + }) + + // Setup default mock returns + mockLaunch.mockResolvedValue(mockBrowser) + mockConnect.mockResolvedValue(mockBrowser) + mockTryChromeHostUrl.mockResolvedValue(true) + + browserSession = new BrowserSession(mockContext as vscode.ExtensionContext) + }) + + describe("closeBrowser with remote browser fallback scenario", () => { + it("should close local browser when remote browser is enabled but fallback to local was used", async () => { + // Setup: Remote browser is enabled in settings + ;(mockContext.globalState!.get as jest.Mock).mockImplementation((key: string) => { + if (key === "remoteBrowserEnabled") return true + if (key === "browserViewportSize") return "900x600" + return undefined + }) + + // Mock discoverChromeHostUrl to fail (no remote browser available) + mockDiscoverChromeHostUrl.mockRejectedValue(new Error("No remote browser found")) + + // Launch browser - this should fallback to local browser + await browserSession.launchBrowser() + + // Verify that launch was called (local browser) + expect(mockLaunch).toHaveBeenCalled() + + // Now close the browser + await browserSession.closeBrowser() + + // Verify that close() was called instead of disconnect() + expect(mockBrowser.close).toHaveBeenCalled() + expect(mockBrowser.disconnect).not.toHaveBeenCalled() + }) + + it("should disconnect from remote browser when actually connected to remote", async () => { + // Setup: Remote browser is enabled in settings + ;(mockContext.globalState!.get as jest.Mock).mockImplementation((key: string) => { + if (key === "remoteBrowserEnabled") return true + if (key === "browserViewportSize") return "900x600" + return undefined + }) + + // Mock discoverChromeHostUrl to succeed + mockDiscoverChromeHostUrl.mockResolvedValue("ws://localhost:9222") + + // Launch browser - this should connect to remote browser + await browserSession.launchBrowser() + + // Verify that connect was called (remote browser) + expect(mockConnect).toHaveBeenCalled() + + // Now close the browser + await browserSession.closeBrowser() + + // Verify that disconnect() was called instead of close() + expect(mockBrowser.disconnect).toHaveBeenCalled() + expect(mockBrowser.close).not.toHaveBeenCalled() + }) + + it("should close local browser when remote browser is disabled", async () => { + // Setup: Remote browser is disabled in settings + ;(mockContext.globalState!.get as jest.Mock).mockImplementation((key: string) => { + if (key === "remoteBrowserEnabled") return false + if (key === "browserViewportSize") return "900x600" + return undefined + }) + + // Launch browser - this should use local browser + await browserSession.launchBrowser() + + // Verify that launch was called (local browser) + expect(mockLaunch).toHaveBeenCalled() + + // Now close the browser + await browserSession.closeBrowser() + + // Verify that close() was called + expect(mockBrowser.close).toHaveBeenCalled() + expect(mockBrowser.disconnect).not.toHaveBeenCalled() + }) + }) +})