Skip to content

Commit 2b7e259

Browse files
ctehannesrudolph
authored andcommitted
Refactor the extension bridge (#7515)
1 parent af00c62 commit 2b7e259

20 files changed

+505
-2260
lines changed

packages/cloud/src/WebAuthService.ts

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,7 @@ export class WebAuthService extends EventEmitter<AuthServiceEvents> implements A
129129
private changeState(newState: AuthState): void {
130130
const previousState = this.state
131131
this.state = newState
132+
this.log(`[auth] changeState: ${previousState} -> ${newState}`)
132133
this.emit("auth-state-changed", { state: newState, previousState })
133134
}
134135

@@ -162,8 +163,6 @@ export class WebAuthService extends EventEmitter<AuthServiceEvents> implements A
162163
this.userInfo = null
163164

164165
this.changeState("logged-out")
165-
166-
this.log("[auth] Transitioned to logged-out state")
167166
}
168167

169168
private transitionToAttemptingSession(credentials: AuthCredentials): void {
@@ -176,17 +175,13 @@ export class WebAuthService extends EventEmitter<AuthServiceEvents> implements A
176175
this.changeState("attempting-session")
177176

178177
this.timer.start()
179-
180-
this.log("[auth] Transitioned to attempting-session state")
181178
}
182179

183180
private transitionToInactiveSession(): void {
184181
this.sessionToken = null
185182
this.userInfo = null
186183

187184
this.changeState("inactive-session")
188-
189-
this.log("[auth] Transitioned to inactive-session state")
190185
}
191186

192187
/**
@@ -422,7 +417,6 @@ export class WebAuthService extends EventEmitter<AuthServiceEvents> implements A
422417

423418
if (previousState !== "active-session") {
424419
this.changeState("active-session")
425-
this.log("[auth] Transitioned to active-session state")
426420
this.fetchUserInfo()
427421
} else {
428422
this.state = "active-session"

packages/cloud/src/bridge/BaseChannel.ts

Lines changed: 4 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,4 @@
11
import type { Socket } from "socket.io-client"
2-
import * as vscode from "vscode"
3-
4-
import type { StaticAppProperties, GitProperties } from "@roo-code/types"
5-
6-
export interface BaseChannelOptions {
7-
instanceId: string
8-
appProperties: StaticAppProperties
9-
gitProperties?: GitProperties
10-
}
112

123
/**
134
* Abstract base class for communication channels in the bridge system.
@@ -20,13 +11,9 @@ export interface BaseChannelOptions {
2011
export abstract class BaseChannel<TCommand = unknown, TEventName extends string = string, TEventData = unknown> {
2112
protected socket: Socket | null = null
2213
protected readonly instanceId: string
23-
protected readonly appProperties: StaticAppProperties
24-
protected readonly gitProperties?: GitProperties
2514

26-
constructor(options: BaseChannelOptions) {
27-
this.instanceId = options.instanceId
28-
this.appProperties = options.appProperties
29-
this.gitProperties = options.gitProperties
15+
constructor(instanceId: string) {
16+
this.instanceId = instanceId
3017
}
3118

3219
/**
@@ -94,26 +81,9 @@ export abstract class BaseChannel<TCommand = unknown, TEventName extends string
9481
}
9582

9683
/**
97-
* Handle incoming commands - template method that ensures common functionality
98-
* is executed before subclass-specific logic.
99-
*
100-
* This method should be called by subclasses to handle commands.
101-
* It will execute common functionality and then delegate to the abstract
102-
* handleCommandImplementation method.
103-
*/
104-
public async handleCommand(command: TCommand): Promise<void> {
105-
// Common functionality: focus the sidebar.
106-
await vscode.commands.executeCommand(`${this.appProperties.appName}.SidebarProvider.focus`)
107-
108-
// Delegate to subclass-specific implementation.
109-
await this.handleCommandImplementation(command)
110-
}
111-
112-
/**
113-
* Handle command-specific logic - must be implemented by subclasses.
114-
* This method is called after common functionality has been executed.
84+
* Handle incoming commands - must be implemented by subclasses.
11585
*/
116-
protected abstract handleCommandImplementation(command: TCommand): Promise<void>
86+
public abstract handleCommand(command: TCommand): void
11787

11888
/**
11989
* Handle connection-specific logic.

packages/cloud/src/bridge/BridgeOrchestrator.ts

Lines changed: 49 additions & 110 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,11 @@
11
import crypto from "crypto"
2-
import os from "os"
32

43
import {
54
type TaskProviderLike,
65
type TaskLike,
76
type CloudUserInfo,
87
type ExtensionBridgeCommand,
98
type TaskBridgeCommand,
10-
type StaticAppProperties,
11-
type GitProperties,
129
ConnectionState,
1310
ExtensionSocketEvents,
1411
TaskSocketEvents,
@@ -34,16 +31,12 @@ export interface BridgeOrchestratorOptions {
3431
export class BridgeOrchestrator {
3532
private static instance: BridgeOrchestrator | null = null
3633

37-
private static pendingTask: TaskLike | null = null
38-
3934
// Core
4035
private readonly userId: string
4136
private readonly socketBridgeUrl: string
4237
private readonly token: string
4338
private readonly provider: TaskProviderLike
4439
private readonly instanceId: string
45-
private readonly appProperties: StaticAppProperties
46-
private readonly gitProperties?: GitProperties
4740

4841
// Components
4942
private socketTransport: SocketTransport
@@ -68,86 +61,58 @@ export class BridgeOrchestrator {
6861
remoteControlEnabled: boolean | undefined,
6962
options: BridgeOrchestratorOptions,
7063
): Promise<void> {
71-
if (BridgeOrchestrator.isEnabled(userInfo, remoteControlEnabled)) {
72-
await BridgeOrchestrator.connect(options)
73-
} else {
74-
await BridgeOrchestrator.disconnect()
75-
}
76-
}
77-
78-
public static async connect(options: BridgeOrchestratorOptions) {
64+
const isEnabled = BridgeOrchestrator.isEnabled(userInfo, remoteControlEnabled)
7965
const instance = BridgeOrchestrator.instance
8066

81-
if (!instance) {
82-
try {
83-
console.log(`[BridgeOrchestrator#connectOrDisconnect] Connecting...`)
84-
// Populate telemetry properties before registering the instance.
85-
await options.provider.getTelemetryProperties()
86-
87-
BridgeOrchestrator.instance = new BridgeOrchestrator(options)
88-
await BridgeOrchestrator.instance.connect()
89-
} catch (error) {
90-
console.error(
91-
`[BridgeOrchestrator#connectOrDisconnect] connect() failed: ${error instanceof Error ? error.message : String(error)}`,
92-
)
93-
}
94-
} else {
95-
if (
96-
instance.connectionState === ConnectionState.FAILED ||
97-
instance.connectionState === ConnectionState.DISCONNECTED
98-
) {
99-
console.log(
100-
`[BridgeOrchestrator#connectOrDisconnect] Re-connecting... (state: ${instance.connectionState})`,
101-
)
102-
103-
instance.reconnect().catch((error) => {
67+
if (isEnabled) {
68+
if (!instance) {
69+
try {
70+
console.log(`[BridgeOrchestrator#connectOrDisconnect] Connecting...`)
71+
BridgeOrchestrator.instance = new BridgeOrchestrator(options)
72+
await BridgeOrchestrator.instance.connect()
73+
} catch (error) {
10474
console.error(
105-
`[BridgeOrchestrator#connectOrDisconnect] reconnect() failed: ${error instanceof Error ? error.message : String(error)}`,
75+
`[BridgeOrchestrator#connectOrDisconnect] connect() failed: ${error instanceof Error ? error.message : String(error)}`,
10676
)
107-
})
77+
}
10878
} else {
109-
console.log(
110-
`[BridgeOrchestrator#connectOrDisconnect] Already connected or connecting (state: ${instance.connectionState})`,
111-
)
112-
}
113-
}
114-
}
115-
116-
public static async disconnect() {
117-
const instance = BridgeOrchestrator.instance
79+
if (
80+
instance.connectionState === ConnectionState.FAILED ||
81+
instance.connectionState === ConnectionState.DISCONNECTED
82+
) {
83+
console.log(
84+
`[BridgeOrchestrator#connectOrDisconnect] Re-connecting... (state: ${instance.connectionState})`,
85+
)
11886

119-
if (instance) {
120-
try {
121-
console.log(
122-
`[BridgeOrchestrator#connectOrDisconnect] Disconnecting... (state: ${instance.connectionState})`,
123-
)
124-
125-
await instance.disconnect()
126-
} catch (error) {
127-
console.error(
128-
`[BridgeOrchestrator#connectOrDisconnect] disconnect() failed: ${error instanceof Error ? error.message : String(error)}`,
129-
)
130-
} finally {
131-
BridgeOrchestrator.instance = null
87+
instance.reconnect().catch((error) => {
88+
console.error(
89+
`[BridgeOrchestrator#connectOrDisconnect] reconnect() failed: ${error instanceof Error ? error.message : String(error)}`,
90+
)
91+
})
92+
} else {
93+
console.log(
94+
`[BridgeOrchestrator#connectOrDisconnect] Already connected or connecting (state: ${instance.connectionState})`,
95+
)
96+
}
13297
}
13398
} else {
134-
console.log(`[BridgeOrchestrator#connectOrDisconnect] Already disconnected`)
135-
}
136-
}
137-
138-
/**
139-
* @TODO: What if subtasks also get spawned? We'd probably want deferred
140-
* subscriptions for those too.
141-
*/
142-
public static async subscribeToTask(task: TaskLike): Promise<void> {
143-
const instance = BridgeOrchestrator.instance
99+
if (instance) {
100+
try {
101+
console.log(
102+
`[BridgeOrchestrator#connectOrDisconnect] Disconnecting... (state: ${instance.connectionState})`,
103+
)
144104

145-
if (instance && instance.socketTransport.isConnected()) {
146-
console.log(`[BridgeOrchestrator#subscribeToTask] Subscribing to task ${task.taskId}`)
147-
await instance.subscribeToTask(task)
148-
} else {
149-
console.log(`[BridgeOrchestrator#subscribeToTask] Deferring subscription for task ${task.taskId}`)
150-
BridgeOrchestrator.pendingTask = task
105+
await instance.disconnect()
106+
} catch (error) {
107+
console.error(
108+
`[BridgeOrchestrator#connectOrDisconnect] disconnect() failed: ${error instanceof Error ? error.message : String(error)}`,
109+
)
110+
} finally {
111+
BridgeOrchestrator.instance = null
112+
}
113+
} else {
114+
console.log(`[BridgeOrchestrator#connectOrDisconnect] Already disconnected`)
115+
}
151116
}
152117
}
153118

@@ -157,8 +122,6 @@ export class BridgeOrchestrator {
157122
this.token = options.token
158123
this.provider = options.provider
159124
this.instanceId = options.sessionId || crypto.randomUUID()
160-
this.appProperties = { ...options.provider.appProperties, hostname: os.hostname() }
161-
this.gitProperties = options.provider.gitProperties
162125

163126
this.socketTransport = new SocketTransport({
164127
url: this.socketBridgeUrl,
@@ -179,19 +142,8 @@ export class BridgeOrchestrator {
179142
onReconnect: () => this.handleReconnect(),
180143
})
181144

182-
this.extensionChannel = new ExtensionChannel({
183-
instanceId: this.instanceId,
184-
appProperties: this.appProperties,
185-
gitProperties: this.gitProperties,
186-
userId: this.userId,
187-
provider: this.provider,
188-
})
189-
190-
this.taskChannel = new TaskChannel({
191-
instanceId: this.instanceId,
192-
appProperties: this.appProperties,
193-
gitProperties: this.gitProperties,
194-
})
145+
this.extensionChannel = new ExtensionChannel(this.instanceId, this.userId, this.provider)
146+
this.taskChannel = new TaskChannel(this.instanceId)
195147
}
196148

197149
private setupSocketListeners() {
@@ -228,27 +180,12 @@ export class BridgeOrchestrator {
228180
const socket = this.socketTransport.getSocket()
229181

230182
if (!socket) {
231-
console.error("[BridgeOrchestrator#handleConnect] Socket not available")
183+
console.error("[BridgeOrchestrator] Socket not available after connect")
232184
return
233185
}
234186

235187
await this.extensionChannel.onConnect(socket)
236188
await this.taskChannel.onConnect(socket)
237-
238-
if (BridgeOrchestrator.pendingTask) {
239-
console.log(
240-
`[BridgeOrchestrator#handleConnect] Subscribing to task ${BridgeOrchestrator.pendingTask.taskId}`,
241-
)
242-
243-
try {
244-
await this.subscribeToTask(BridgeOrchestrator.pendingTask)
245-
BridgeOrchestrator.pendingTask = null
246-
} catch (error) {
247-
console.error(
248-
`[BridgeOrchestrator#handleConnect] subscribeToTask() failed: ${error instanceof Error ? error.message : String(error)}`,
249-
)
250-
}
251-
}
252189
}
253190

254191
private handleDisconnect() {
@@ -312,6 +249,9 @@ export class BridgeOrchestrator {
312249
}
313250

314251
private async connect(): Promise<void> {
252+
// Populate the app and git properties before registering the instance.
253+
await this.provider.getTelemetryProperties()
254+
315255
await this.socketTransport.connect()
316256
this.setupSocketListeners()
317257
}
@@ -321,7 +261,6 @@ export class BridgeOrchestrator {
321261
await this.taskChannel.cleanup(this.socketTransport.getSocket())
322262
await this.socketTransport.disconnect()
323263
BridgeOrchestrator.instance = null
324-
BridgeOrchestrator.pendingTask = null
325264
}
326265

327266
public async reconnect(): Promise<void> {

0 commit comments

Comments
 (0)