Skip to content

Commit b309a6a

Browse files
ctecte
authored
Disable Roomote Control on logout (RooCodeInc#7976)
1 parent 9ea7173 commit b309a6a

File tree

4 files changed

+292
-9
lines changed

4 files changed

+292
-9
lines changed

‎packages/cloud/src/bridge/BridgeOrchestrator.ts‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -59,13 +59,28 @@ export class BridgeOrchestrator {
5959
return BridgeOrchestrator.instance
6060
}
6161

62-
public static isEnabled(user?: CloudUserInfo | null, remoteControlEnabled?: boolean): boolean {
63-
return !!(user?.id && user.extensionBridgeEnabled && remoteControlEnabled)
62+
public static isEnabled(user: CloudUserInfo | null, remoteControlEnabled: boolean): boolean {
63+
// Always disabled if signed out.
64+
if (!user) {
65+
return false
66+
}
67+
68+
// Disabled by the user's organization?
69+
if (!user.extensionBridgeEnabled) {
70+
return false
71+
}
72+
73+
// Disabled by the user?
74+
if (!remoteControlEnabled) {
75+
return false
76+
}
77+
78+
return true
6479
}
6580

6681
public static async connectOrDisconnect(
67-
userInfo: CloudUserInfo | null,
68-
remoteControlEnabled: boolean | undefined,
82+
userInfo: CloudUserInfo,
83+
remoteControlEnabled: boolean,
6984
options: BridgeOrchestratorOptions,
7085
): Promise<void> {
7186
if (BridgeOrchestrator.isEnabled(userInfo, remoteControlEnabled)) {

‎src/__tests__/extension.spec.ts‎

Lines changed: 257 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,257 @@
1+
// npx vitest run __tests__/extension.spec.ts
2+
3+
import type * as vscode from "vscode"
4+
import type { AuthState } from "@roo-code/types"
5+
6+
vi.mock("vscode", () => ({
7+
window: {
8+
createOutputChannel: vi.fn().mockReturnValue({
9+
appendLine: vi.fn(),
10+
}),
11+
registerWebviewViewProvider: vi.fn(),
12+
registerUriHandler: vi.fn(),
13+
tabGroups: {
14+
onDidChangeTabs: vi.fn(),
15+
},
16+
onDidChangeActiveTextEditor: vi.fn(),
17+
},
18+
workspace: {
19+
registerTextDocumentContentProvider: vi.fn(),
20+
getConfiguration: vi.fn().mockReturnValue({
21+
get: vi.fn().mockReturnValue([]),
22+
}),
23+
createFileSystemWatcher: vi.fn().mockReturnValue({
24+
onDidCreate: vi.fn(),
25+
onDidChange: vi.fn(),
26+
onDidDelete: vi.fn(),
27+
dispose: vi.fn(),
28+
}),
29+
onDidChangeWorkspaceFolders: vi.fn(),
30+
},
31+
languages: {
32+
registerCodeActionsProvider: vi.fn(),
33+
},
34+
commands: {
35+
executeCommand: vi.fn(),
36+
},
37+
env: {
38+
language: "en",
39+
},
40+
ExtensionMode: {
41+
Production: 1,
42+
},
43+
}))
44+
45+
vi.mock("@dotenvx/dotenvx", () => ({
46+
config: vi.fn(),
47+
}))
48+
49+
const mockBridgeOrchestratorDisconnect = vi.fn().mockResolvedValue(undefined)
50+
51+
vi.mock("@roo-code/cloud", () => ({
52+
CloudService: {
53+
createInstance: vi.fn(),
54+
hasInstance: vi.fn().mockReturnValue(true),
55+
get instance() {
56+
return {
57+
off: vi.fn(),
58+
on: vi.fn(),
59+
getUserInfo: vi.fn().mockReturnValue(null),
60+
isTaskSyncEnabled: vi.fn().mockReturnValue(false),
61+
}
62+
},
63+
},
64+
BridgeOrchestrator: {
65+
disconnect: mockBridgeOrchestratorDisconnect,
66+
},
67+
getRooCodeApiUrl: vi.fn().mockReturnValue("https://app.roocode.com"),
68+
}))
69+
70+
vi.mock("@roo-code/telemetry", () => ({
71+
TelemetryService: {
72+
createInstance: vi.fn().mockReturnValue({
73+
register: vi.fn(),
74+
setProvider: vi.fn(),
75+
shutdown: vi.fn(),
76+
}),
77+
get instance() {
78+
return {
79+
register: vi.fn(),
80+
setProvider: vi.fn(),
81+
shutdown: vi.fn(),
82+
}
83+
},
84+
},
85+
PostHogTelemetryClient: vi.fn(),
86+
}))
87+
88+
vi.mock("../utils/outputChannelLogger", () => ({
89+
createOutputChannelLogger: vi.fn().mockReturnValue(vi.fn()),
90+
createDualLogger: vi.fn().mockReturnValue(vi.fn()),
91+
}))
92+
93+
vi.mock("../shared/package", () => ({
94+
Package: {
95+
name: "test-extension",
96+
outputChannel: "Test Output",
97+
version: "1.0.0",
98+
},
99+
}))
100+
101+
vi.mock("../shared/language", () => ({
102+
formatLanguage: vi.fn().mockReturnValue("en"),
103+
}))
104+
105+
vi.mock("../core/config/ContextProxy", () => ({
106+
ContextProxy: {
107+
getInstance: vi.fn().mockResolvedValue({
108+
getValue: vi.fn(),
109+
setValue: vi.fn(),
110+
getValues: vi.fn().mockReturnValue({}),
111+
getProviderSettings: vi.fn().mockReturnValue({}),
112+
}),
113+
},
114+
}))
115+
116+
vi.mock("../integrations/editor/DiffViewProvider", () => ({
117+
DIFF_VIEW_URI_SCHEME: "test-diff-scheme",
118+
}))
119+
120+
vi.mock("../integrations/terminal/TerminalRegistry", () => ({
121+
TerminalRegistry: {
122+
initialize: vi.fn(),
123+
cleanup: vi.fn(),
124+
},
125+
}))
126+
127+
vi.mock("../services/mcp/McpServerManager", () => ({
128+
McpServerManager: {
129+
cleanup: vi.fn().mockResolvedValue(undefined),
130+
getInstance: vi.fn().mockResolvedValue(null),
131+
unregisterProvider: vi.fn(),
132+
},
133+
}))
134+
135+
vi.mock("../services/code-index/manager", () => ({
136+
CodeIndexManager: {
137+
getInstance: vi.fn().mockReturnValue(null),
138+
},
139+
}))
140+
141+
vi.mock("../services/mdm/MdmService", () => ({
142+
MdmService: {
143+
createInstance: vi.fn().mockResolvedValue(null),
144+
},
145+
}))
146+
147+
vi.mock("../utils/migrateSettings", () => ({
148+
migrateSettings: vi.fn().mockResolvedValue(undefined),
149+
}))
150+
151+
vi.mock("../utils/autoImportSettings", () => ({
152+
autoImportSettings: vi.fn().mockResolvedValue(undefined),
153+
}))
154+
155+
vi.mock("../extension/api", () => ({
156+
API: vi.fn().mockImplementation(() => ({})),
157+
}))
158+
159+
vi.mock("../activate", () => ({
160+
handleUri: vi.fn(),
161+
registerCommands: vi.fn(),
162+
registerCodeActions: vi.fn(),
163+
registerTerminalActions: vi.fn(),
164+
CodeActionProvider: vi.fn().mockImplementation(() => ({
165+
providedCodeActionKinds: [],
166+
})),
167+
}))
168+
169+
vi.mock("../i18n", () => ({
170+
initializeI18n: vi.fn(),
171+
}))
172+
173+
describe("extension.ts", () => {
174+
let mockContext: vscode.ExtensionContext
175+
let authStateChangedHandler:
176+
| ((data: { state: AuthState; previousState: AuthState }) => void | Promise<void>)
177+
| undefined
178+
179+
beforeEach(() => {
180+
vi.clearAllMocks()
181+
mockBridgeOrchestratorDisconnect.mockClear()
182+
183+
mockContext = {
184+
extensionPath: "/test/path",
185+
globalState: {
186+
get: vi.fn().mockReturnValue(undefined),
187+
update: vi.fn(),
188+
},
189+
subscriptions: [],
190+
} as unknown as vscode.ExtensionContext
191+
192+
authStateChangedHandler = undefined
193+
})
194+
195+
test("authStateChangedHandler calls BridgeOrchestrator.disconnect when logged-out event fires", async () => {
196+
const { CloudService, BridgeOrchestrator } = await import("@roo-code/cloud")
197+
198+
// Capture the auth state changed handler.
199+
vi.mocked(CloudService.createInstance).mockImplementation(async (_context, _logger, handlers) => {
200+
if (handlers?.["auth-state-changed"]) {
201+
authStateChangedHandler = handlers["auth-state-changed"]
202+
}
203+
204+
return {
205+
off: vi.fn(),
206+
on: vi.fn(),
207+
telemetryClient: null,
208+
} as any
209+
})
210+
211+
// Activate the extension.
212+
const { activate } = await import("../extension")
213+
await activate(mockContext)
214+
215+
// Verify handler was registered.
216+
expect(authStateChangedHandler).toBeDefined()
217+
218+
// Trigger logout.
219+
await authStateChangedHandler!({
220+
state: "logged-out" as AuthState,
221+
previousState: "logged-in" as AuthState,
222+
})
223+
224+
// Verify BridgeOrchestrator.disconnect was called
225+
expect(mockBridgeOrchestratorDisconnect).toHaveBeenCalled()
226+
})
227+
228+
test("authStateChangedHandler does not call BridgeOrchestrator.disconnect for other states", async () => {
229+
const { CloudService } = await import("@roo-code/cloud")
230+
231+
// Capture the auth state changed handler.
232+
vi.mocked(CloudService.createInstance).mockImplementation(async (_context, _logger, handlers) => {
233+
if (handlers?.["auth-state-changed"]) {
234+
authStateChangedHandler = handlers["auth-state-changed"]
235+
}
236+
237+
return {
238+
off: vi.fn(),
239+
on: vi.fn(),
240+
telemetryClient: null,
241+
} as any
242+
})
243+
244+
// Activate the extension.
245+
const { activate } = await import("../extension")
246+
await activate(mockContext)
247+
248+
// Trigger login.
249+
await authStateChangedHandler!({
250+
state: "logged-in" as AuthState,
251+
previousState: "logged-out" as AuthState,
252+
})
253+
254+
// Verify BridgeOrchestrator.disconnect was NOT called.
255+
expect(mockBridgeOrchestratorDisconnect).not.toHaveBeenCalled()
256+
})
257+
})

‎src/core/webview/ClineProvider.ts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2262,7 +2262,19 @@ export class ClineProvider
22622262
}
22632263

22642264
public async remoteControlEnabled(enabled: boolean) {
2265+
if (!enabled) {
2266+
await BridgeOrchestrator.disconnect()
2267+
return
2268+
}
2269+
22652270
const userInfo = CloudService.instance.getUserInfo()
2271+
2272+
if (!userInfo) {
2273+
this.log("[ClineProvider#remoteControlEnabled] Failed to get user info, disconnecting")
2274+
await BridgeOrchestrator.disconnect()
2275+
return
2276+
}
2277+
22662278
const config = await CloudService.instance.cloudAPI?.bridgeConfig().catch(() => undefined)
22672279

22682280
if (!config) {

‎src/extension.ts‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -134,10 +134,9 @@ export async function activate(context: vscode.ExtensionContext) {
134134
if (data.state === "logged-out") {
135135
try {
136136
await provider.remoteControlEnabled(false)
137-
cloudLogger("[CloudService] BridgeOrchestrator disconnected on logout")
138137
} catch (error) {
139138
cloudLogger(
140-
`[CloudService] Failed to disconnect BridgeOrchestrator on logout: ${error instanceof Error ? error.message : String(error)}`,
139+
`[authStateChangedHandler] remoteControlEnabled(false) failed: ${error instanceof Error ? error.message : String(error)}`,
141140
)
142141
}
143142
}
@@ -151,7 +150,7 @@ export async function activate(context: vscode.ExtensionContext) {
151150
provider.remoteControlEnabled(CloudService.instance.isTaskSyncEnabled())
152151
} catch (error) {
153152
cloudLogger(
154-
`[CloudService] BridgeOrchestrator#connectOrDisconnect failed on settings change: ${error instanceof Error ? error.message : String(error)}`,
153+
`[settingsUpdatedHandler] remoteControlEnabled failed: ${error instanceof Error ? error.message : String(error)}`,
155154
)
156155
}
157156
}
@@ -163,15 +162,15 @@ export async function activate(context: vscode.ExtensionContext) {
163162
postStateListener()
164163

165164
if (!CloudService.instance.cloudAPI) {
166-
cloudLogger("[CloudService] CloudAPI is not initialized")
165+
cloudLogger("[userInfoHandler] CloudAPI is not initialized")
167166
return
168167
}
169168

170169
try {
171170
provider.remoteControlEnabled(CloudService.instance.isTaskSyncEnabled())
172171
} catch (error) {
173172
cloudLogger(
174-
`[CloudService] BridgeOrchestrator#connectOrDisconnect failed on user change: ${error instanceof Error ? error.message : String(error)}`,
173+
`[userInfoHandler] remoteControlEnabled failed: ${error instanceof Error ? error.message : String(error)}`,
175174
)
176175
}
177176
}

0 commit comments

Comments
 (0)