From b354e443ee1583a4dcb62e9d794d97d75a431a4b Mon Sep 17 00:00:00 2001 From: Matt Rubens Date: Sun, 31 Aug 2025 00:10:03 -0400 Subject: [PATCH 1/3] Disconnect extension bridge on logout --- src/__tests__/extension.logout-bridge.test.ts | 226 ++++++++++++++++++ src/extension.ts | 32 ++- 2 files changed, 255 insertions(+), 3 deletions(-) create mode 100644 src/__tests__/extension.logout-bridge.test.ts diff --git a/src/__tests__/extension.logout-bridge.test.ts b/src/__tests__/extension.logout-bridge.test.ts new file mode 100644 index 000000000000..2a1e0c92401d --- /dev/null +++ b/src/__tests__/extension.logout-bridge.test.ts @@ -0,0 +1,226 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest" +import type { ExtensionContext } from "vscode" +import type { AuthState, CloudUserInfo } from "@roo-code/types" +import { CloudService, BridgeOrchestrator } from "@roo-code/cloud" + +// Mock modules +vi.mock("vscode", () => ({ + window: { + createOutputChannel: vi.fn(() => ({ + appendLine: vi.fn(), + })), + }, + env: { + sessionId: "test-session-id", + }, + workspace: { + getConfiguration: vi.fn(() => ({ + get: vi.fn(() => []), + })), + }, + commands: { + executeCommand: vi.fn(), + }, +})) + +vi.mock("@roo-code/cloud", () => { + const mockBridgeOrchestrator = { + connectOrDisconnect: vi.fn(), + getInstance: vi.fn(), + disconnect: vi.fn(), + } + + const mockCloudService = { + instance: { + off: vi.fn(), + getUserInfo: vi.fn(), + cloudAPI: { + bridgeConfig: vi.fn(() => ({ + userId: "test-user", + socketBridgeUrl: "wss://test.bridge.url", + token: "test-token", + })), + }, + }, + hasInstance: vi.fn(() => true), + createInstance: vi.fn(), + } + + return { + CloudService: mockCloudService, + BridgeOrchestrator: mockBridgeOrchestrator, + } +}) + +describe("Extension Bridge Logout", () => { + let authStateChangedHandler: ((data: { state: AuthState; previousState: AuthState }) => Promise) | undefined + let mockCloudLogger: ReturnType + let mockProvider: any + + beforeEach(() => { + // Reset mocks + vi.clearAllMocks() + + // Setup mock logger + mockCloudLogger = vi.fn() + + // Setup mock provider + mockProvider = { + postStateToWebview: vi.fn(), + } + + // Reset handlers + authStateChangedHandler = undefined + }) + + afterEach(() => { + vi.restoreAllMocks() + }) + + it("should disconnect BridgeOrchestrator when user logs out", async () => { + // Create the auth state changed handler (simulating what extension.ts does) + const postStateListener = () => mockProvider.postStateToWebview() + + authStateChangedHandler = async (data: { state: AuthState; previousState: AuthState }) => { + postStateListener() + + // Check if user has logged out + if (data.state === "logged-out") { + try { + // Disconnect the bridge when user logs out + await BridgeOrchestrator.connectOrDisconnect(null, false, { + userId: "", + socketBridgeUrl: "", + token: "", + provider: mockProvider, + sessionId: "test-session-id", + }) + + mockCloudLogger("[CloudService] BridgeOrchestrator disconnected on logout") + } catch (error) { + mockCloudLogger( + `[CloudService] Failed to disconnect BridgeOrchestrator on logout: ${ + error instanceof Error ? error.message : String(error) + }`, + ) + } + } + } + + // Simulate logout event + await authStateChangedHandler({ + state: "logged-out", + previousState: "active-session", + }) + + // Verify that postStateToWebview was called + expect(mockProvider.postStateToWebview).toHaveBeenCalledOnce() + + // Verify that BridgeOrchestrator.connectOrDisconnect was called with correct params to disconnect + expect(BridgeOrchestrator.connectOrDisconnect).toHaveBeenCalledWith( + null, // userInfo is null when logged out + false, // remoteControlEnabled is false to trigger disconnection + { + userId: "", + socketBridgeUrl: "", + token: "", + provider: mockProvider, + sessionId: "test-session-id", + }, + ) + + // Verify success log message + expect(mockCloudLogger).toHaveBeenCalledWith("[CloudService] BridgeOrchestrator disconnected on logout") + }) + + it("should handle errors when disconnecting BridgeOrchestrator fails", async () => { + // Make connectOrDisconnect throw an error + const errorMessage = "Failed to disconnect" + vi.mocked(BridgeOrchestrator.connectOrDisconnect).mockRejectedValueOnce(new Error(errorMessage)) + + // Create the auth state changed handler + const postStateListener = () => mockProvider.postStateToWebview() + + authStateChangedHandler = async (data: { state: AuthState; previousState: AuthState }) => { + postStateListener() + + if (data.state === "logged-out") { + try { + await BridgeOrchestrator.connectOrDisconnect(null, false, { + userId: "", + socketBridgeUrl: "", + token: "", + provider: mockProvider, + sessionId: "test-session-id", + }) + + mockCloudLogger("[CloudService] BridgeOrchestrator disconnected on logout") + } catch (error) { + mockCloudLogger( + `[CloudService] Failed to disconnect BridgeOrchestrator on logout: ${ + error instanceof Error ? error.message : String(error) + }`, + ) + } + } + } + + // Simulate logout event + await authStateChangedHandler({ + state: "logged-out", + previousState: "active-session", + }) + + // Verify error was caught and logged + expect(mockCloudLogger).toHaveBeenCalledWith( + `[CloudService] Failed to disconnect BridgeOrchestrator on logout: ${errorMessage}`, + ) + + // Verify that postStateToWebview was still called despite the error + expect(mockProvider.postStateToWebview).toHaveBeenCalledOnce() + }) + + it("should not disconnect BridgeOrchestrator for non-logout state changes", async () => { + // Create the auth state changed handler + const postStateListener = () => mockProvider.postStateToWebview() + + authStateChangedHandler = async (data: { state: AuthState; previousState: AuthState }) => { + postStateListener() + + if (data.state === "logged-out") { + try { + await BridgeOrchestrator.connectOrDisconnect(null, false, { + userId: "", + socketBridgeUrl: "", + token: "", + provider: mockProvider, + sessionId: "test-session-id", + }) + + mockCloudLogger("[CloudService] BridgeOrchestrator disconnected on logout") + } catch (error) { + mockCloudLogger( + `[CloudService] Failed to disconnect BridgeOrchestrator on logout: ${ + error instanceof Error ? error.message : String(error) + }`, + ) + } + } + } + + // Simulate state change that is NOT a logout + await authStateChangedHandler({ + state: "active-session", + previousState: "attempting-session", + }) + + // Verify that postStateToWebview was called + expect(mockProvider.postStateToWebview).toHaveBeenCalledOnce() + + // Verify that BridgeOrchestrator.connectOrDisconnect was NOT called + expect(BridgeOrchestrator.connectOrDisconnect).not.toHaveBeenCalled() + + // Verify no log messages about disconnection + expect(mockCloudLogger).not.toHaveBeenCalled() + }) +}) diff --git a/src/extension.ts b/src/extension.ts index cb765ad718e8..07f9d7dacf63 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -12,7 +12,7 @@ try { console.warn("Failed to load environment variables:", e) } -import type { CloudUserInfo } from "@roo-code/types" +import type { CloudUserInfo, AuthState } from "@roo-code/types" import { CloudService, BridgeOrchestrator } from "@roo-code/cloud" import { TelemetryService, PostHogTelemetryClient } from "@roo-code/telemetry" @@ -53,7 +53,7 @@ let outputChannel: vscode.OutputChannel let extensionContext: vscode.ExtensionContext let cloudService: CloudService | undefined -let authStateChangedHandler: (() => void) | undefined +let authStateChangedHandler: ((data: { state: AuthState; previousState: AuthState }) => Promise) | undefined let settingsUpdatedHandler: (() => void) | undefined let userInfoHandler: ((data: { userInfo: CloudUserInfo }) => Promise) | undefined @@ -127,7 +127,33 @@ export async function activate(context: vscode.ExtensionContext) { // Initialize Roo Code Cloud service. const postStateListener = () => ClineProvider.getVisibleInstance()?.postStateToWebview() - authStateChangedHandler = postStateListener + + authStateChangedHandler = async (data: { state: AuthState; previousState: AuthState }) => { + postStateListener() + + // Check if user has logged out + if (data.state === "logged-out") { + try { + // Disconnect the bridge when user logs out + // Pass null userInfo and false for remoteControlEnabled to trigger disconnection + await BridgeOrchestrator.connectOrDisconnect(null, false, { + userId: "", + socketBridgeUrl: "", + token: "", + provider, + sessionId: vscode.env.sessionId, + }) + + cloudLogger("[CloudService] BridgeOrchestrator disconnected on logout") + } catch (error) { + cloudLogger( + `[CloudService] Failed to disconnect BridgeOrchestrator on logout: ${ + error instanceof Error ? error.message : String(error) + }`, + ) + } + } + } settingsUpdatedHandler = async () => { const userInfo = CloudService.instance.getUserInfo() From a16185bb7f91bce92f26380f80570ae1cf892769 Mon Sep 17 00:00:00 2001 From: Matt Rubens Date: Sun, 31 Aug 2025 00:17:24 -0400 Subject: [PATCH 2/3] Remove bad test --- src/__tests__/extension.logout-bridge.test.ts | 226 ------------------ 1 file changed, 226 deletions(-) delete mode 100644 src/__tests__/extension.logout-bridge.test.ts diff --git a/src/__tests__/extension.logout-bridge.test.ts b/src/__tests__/extension.logout-bridge.test.ts deleted file mode 100644 index 2a1e0c92401d..000000000000 --- a/src/__tests__/extension.logout-bridge.test.ts +++ /dev/null @@ -1,226 +0,0 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from "vitest" -import type { ExtensionContext } from "vscode" -import type { AuthState, CloudUserInfo } from "@roo-code/types" -import { CloudService, BridgeOrchestrator } from "@roo-code/cloud" - -// Mock modules -vi.mock("vscode", () => ({ - window: { - createOutputChannel: vi.fn(() => ({ - appendLine: vi.fn(), - })), - }, - env: { - sessionId: "test-session-id", - }, - workspace: { - getConfiguration: vi.fn(() => ({ - get: vi.fn(() => []), - })), - }, - commands: { - executeCommand: vi.fn(), - }, -})) - -vi.mock("@roo-code/cloud", () => { - const mockBridgeOrchestrator = { - connectOrDisconnect: vi.fn(), - getInstance: vi.fn(), - disconnect: vi.fn(), - } - - const mockCloudService = { - instance: { - off: vi.fn(), - getUserInfo: vi.fn(), - cloudAPI: { - bridgeConfig: vi.fn(() => ({ - userId: "test-user", - socketBridgeUrl: "wss://test.bridge.url", - token: "test-token", - })), - }, - }, - hasInstance: vi.fn(() => true), - createInstance: vi.fn(), - } - - return { - CloudService: mockCloudService, - BridgeOrchestrator: mockBridgeOrchestrator, - } -}) - -describe("Extension Bridge Logout", () => { - let authStateChangedHandler: ((data: { state: AuthState; previousState: AuthState }) => Promise) | undefined - let mockCloudLogger: ReturnType - let mockProvider: any - - beforeEach(() => { - // Reset mocks - vi.clearAllMocks() - - // Setup mock logger - mockCloudLogger = vi.fn() - - // Setup mock provider - mockProvider = { - postStateToWebview: vi.fn(), - } - - // Reset handlers - authStateChangedHandler = undefined - }) - - afterEach(() => { - vi.restoreAllMocks() - }) - - it("should disconnect BridgeOrchestrator when user logs out", async () => { - // Create the auth state changed handler (simulating what extension.ts does) - const postStateListener = () => mockProvider.postStateToWebview() - - authStateChangedHandler = async (data: { state: AuthState; previousState: AuthState }) => { - postStateListener() - - // Check if user has logged out - if (data.state === "logged-out") { - try { - // Disconnect the bridge when user logs out - await BridgeOrchestrator.connectOrDisconnect(null, false, { - userId: "", - socketBridgeUrl: "", - token: "", - provider: mockProvider, - sessionId: "test-session-id", - }) - - mockCloudLogger("[CloudService] BridgeOrchestrator disconnected on logout") - } catch (error) { - mockCloudLogger( - `[CloudService] Failed to disconnect BridgeOrchestrator on logout: ${ - error instanceof Error ? error.message : String(error) - }`, - ) - } - } - } - - // Simulate logout event - await authStateChangedHandler({ - state: "logged-out", - previousState: "active-session", - }) - - // Verify that postStateToWebview was called - expect(mockProvider.postStateToWebview).toHaveBeenCalledOnce() - - // Verify that BridgeOrchestrator.connectOrDisconnect was called with correct params to disconnect - expect(BridgeOrchestrator.connectOrDisconnect).toHaveBeenCalledWith( - null, // userInfo is null when logged out - false, // remoteControlEnabled is false to trigger disconnection - { - userId: "", - socketBridgeUrl: "", - token: "", - provider: mockProvider, - sessionId: "test-session-id", - }, - ) - - // Verify success log message - expect(mockCloudLogger).toHaveBeenCalledWith("[CloudService] BridgeOrchestrator disconnected on logout") - }) - - it("should handle errors when disconnecting BridgeOrchestrator fails", async () => { - // Make connectOrDisconnect throw an error - const errorMessage = "Failed to disconnect" - vi.mocked(BridgeOrchestrator.connectOrDisconnect).mockRejectedValueOnce(new Error(errorMessage)) - - // Create the auth state changed handler - const postStateListener = () => mockProvider.postStateToWebview() - - authStateChangedHandler = async (data: { state: AuthState; previousState: AuthState }) => { - postStateListener() - - if (data.state === "logged-out") { - try { - await BridgeOrchestrator.connectOrDisconnect(null, false, { - userId: "", - socketBridgeUrl: "", - token: "", - provider: mockProvider, - sessionId: "test-session-id", - }) - - mockCloudLogger("[CloudService] BridgeOrchestrator disconnected on logout") - } catch (error) { - mockCloudLogger( - `[CloudService] Failed to disconnect BridgeOrchestrator on logout: ${ - error instanceof Error ? error.message : String(error) - }`, - ) - } - } - } - - // Simulate logout event - await authStateChangedHandler({ - state: "logged-out", - previousState: "active-session", - }) - - // Verify error was caught and logged - expect(mockCloudLogger).toHaveBeenCalledWith( - `[CloudService] Failed to disconnect BridgeOrchestrator on logout: ${errorMessage}`, - ) - - // Verify that postStateToWebview was still called despite the error - expect(mockProvider.postStateToWebview).toHaveBeenCalledOnce() - }) - - it("should not disconnect BridgeOrchestrator for non-logout state changes", async () => { - // Create the auth state changed handler - const postStateListener = () => mockProvider.postStateToWebview() - - authStateChangedHandler = async (data: { state: AuthState; previousState: AuthState }) => { - postStateListener() - - if (data.state === "logged-out") { - try { - await BridgeOrchestrator.connectOrDisconnect(null, false, { - userId: "", - socketBridgeUrl: "", - token: "", - provider: mockProvider, - sessionId: "test-session-id", - }) - - mockCloudLogger("[CloudService] BridgeOrchestrator disconnected on logout") - } catch (error) { - mockCloudLogger( - `[CloudService] Failed to disconnect BridgeOrchestrator on logout: ${ - error instanceof Error ? error.message : String(error) - }`, - ) - } - } - } - - // Simulate state change that is NOT a logout - await authStateChangedHandler({ - state: "active-session", - previousState: "attempting-session", - }) - - // Verify that postStateToWebview was called - expect(mockProvider.postStateToWebview).toHaveBeenCalledOnce() - - // Verify that BridgeOrchestrator.connectOrDisconnect was NOT called - expect(BridgeOrchestrator.connectOrDisconnect).not.toHaveBeenCalled() - - // Verify no log messages about disconnection - expect(mockCloudLogger).not.toHaveBeenCalled() - }) -}) From 1f8be2b1c0f98642e003900bc2d46026d6d0a123 Mon Sep 17 00:00:00 2001 From: Matt Rubens Date: Sun, 31 Aug 2025 00:25:54 -0400 Subject: [PATCH 3/3] Cleanup --- packages/cloud/src/bridge/BridgeOrchestrator.ts | 8 +++++++- src/extension.ts | 11 +++-------- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/packages/cloud/src/bridge/BridgeOrchestrator.ts b/packages/cloud/src/bridge/BridgeOrchestrator.ts index 952c4b3e217d..69a6f5a57d02 100644 --- a/packages/cloud/src/bridge/BridgeOrchestrator.ts +++ b/packages/cloud/src/bridge/BridgeOrchestrator.ts @@ -61,13 +61,19 @@ export class BridgeOrchestrator { public static async connectOrDisconnect( userInfo: CloudUserInfo | null, remoteControlEnabled: boolean | undefined, - options: BridgeOrchestratorOptions, + options?: BridgeOrchestratorOptions, ): Promise { const isEnabled = BridgeOrchestrator.isEnabled(userInfo, remoteControlEnabled) const instance = BridgeOrchestrator.instance if (isEnabled) { if (!instance) { + if (!options) { + console.error( + `[BridgeOrchestrator#connectOrDisconnect] Cannot connect: options are required for connection`, + ) + return + } try { console.log(`[BridgeOrchestrator#connectOrDisconnect] Connecting...`) BridgeOrchestrator.instance = new BridgeOrchestrator(options) diff --git a/src/extension.ts b/src/extension.ts index 07f9d7dacf63..8cb739922edf 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -135,14 +135,9 @@ export async function activate(context: vscode.ExtensionContext) { if (data.state === "logged-out") { try { // Disconnect the bridge when user logs out - // Pass null userInfo and false for remoteControlEnabled to trigger disconnection - await BridgeOrchestrator.connectOrDisconnect(null, false, { - userId: "", - socketBridgeUrl: "", - token: "", - provider, - sessionId: vscode.env.sessionId, - }) + // When userInfo is null and remoteControlEnabled is false, BridgeOrchestrator + // will disconnect. The options parameter is not needed for disconnection. + await BridgeOrchestrator.connectOrDisconnect(null, false) cloudLogger("[CloudService] BridgeOrchestrator disconnected on logout") } catch (error) {