Some cleanup in ExtensionHost (#10600)

Co-authored-by: Roo Code <roomote@roocode.com>

Author: cte

apps/cli/docs/AGENT_LOOP.md

@@ -180,12 +180,13 @@ function isStreaming(messages) {
 
 ### ExtensionClient
 
-The **single source of truth** for agent state. It:
+The **single source of truth** for agent state, including the current mode. It:
 
 - Receives all messages from the extension
 - Stores them in the `StateStore`
+- Tracks the current mode from state messages
 - Computes the current state via `detectAgentState()`
-- Emits events when state changes
+- Emits events when state changes (including mode changes)
 
 ```typescript
 const client = new ExtensionClient({
@@ -199,31 +200,41 @@ if (state.isWaitingForInput) {
 	console.log(`Agent needs: ${state.currentAsk}`)
 }
 
+// Query current mode
+const mode = client.getCurrentMode()
+console.log(`Current mode: ${mode}`) // e.g., "code", "architect", "ask"
+
 // Subscribe to events
 client.on("waitingForInput", (event) => {
 	console.log(`Waiting for: ${event.ask}`)
 })
+
+// Subscribe to mode changes
+client.on("modeChanged", (event) => {
+	console.log(`Mode changed: ${event.previousMode} -> ${event.currentMode}`)
+})
 ```
 
 ### StateStore
 
-Holds the `clineMessages` array and computed state:
+Holds the `clineMessages` array, computed state, and current mode:
 
 ```typescript
 interface StoreState {
 	messages: ClineMessage[] // The raw message array
 	agentState: AgentStateInfo // Computed state
 	isInitialized: boolean // Have we received any state?
+	currentMode: string | undefined // Current mode (e.g., "code", "architect")
 }
 ```
 
 ### MessageProcessor
 
 Handles incoming messages from the extension:
 
-- `"state"` messages → Update `clineMessages` array
+- `"state"` messages → Update `clineMessages` array and track mode
 - `"messageUpdated"` messages → Update single message in array
-- Emits events for state transitions
+- Emits events for state transitions and mode changes
 
 ### AskDispatcher
 

apps/cli/package.json

@@ -30,6 +30,7 @@
 		"commander": "^12.1.0",
 		"fuzzysort": "^3.1.0",
 		"ink": "^6.6.0",
+		"p-wait-for": "^5.0.2",
 		"react": "^19.1.0",
 		"superjson": "^2.2.6",
 		"zustand": "^5.0.0"

apps/cli/src/agent/__tests__/extension-client.test.ts

@@ -14,8 +14,8 @@ function createMessage(overrides: Partial<ClineMessage>): ClineMessage {
 	return { ts: Date.now() + Math.random() * 1000, type: "say", ...overrides }
 }
 
-function createStateMessage(messages: ClineMessage[]): ExtensionMessage {
-	return { type: "state", state: { clineMessages: messages } } as ExtensionMessage
+function createStateMessage(messages: ClineMessage[], mode?: string): ExtensionMessage {
+	return { type: "state", state: { clineMessages: messages, mode } } as ExtensionMessage
 }
 
 describe("detectAgentState", () => {
@@ -300,6 +300,44 @@ describe("ExtensionClient", () => {
 			client.handleMessage(createStateMessage([createMessage({ say: "text", ts: Date.now() + 1 })]))
 			expect(callCount).toBe(1) // Should not increase.
 		})
+
+		it("should emit modeChanged events", () => {
+			const { client } = createMockClient()
+			const modeChanges: { previousMode: string | undefined; currentMode: string }[] = []
+
+			client.onModeChanged((event) => {
+				modeChanges.push(event)
+			})
+
+			// Set initial mode
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "code"))
+
+			expect(modeChanges).toHaveLength(1)
+			expect(modeChanges[0]).toEqual({ previousMode: undefined, currentMode: "code" })
+
+			// Change mode
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "architect"))
+
+			expect(modeChanges).toHaveLength(2)
+			expect(modeChanges[1]).toEqual({ previousMode: "code", currentMode: "architect" })
+		})
+
+		it("should not emit modeChanged when mode stays the same", () => {
+			const { client } = createMockClient()
+			let modeChangeCount = 0
+
+			client.onModeChanged(() => {
+				modeChangeCount++
+			})
+
+			// Set initial mode
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "code"))
+			expect(modeChangeCount).toBe(1)
+
+			// Same mode - should not emit
+			client.handleMessage(createStateMessage([createMessage({ say: "text", ts: Date.now() + 1 })], "code"))
+			expect(modeChangeCount).toBe(1)
+		})
 	})
 
 	describe("Response methods", () => {
@@ -458,6 +496,65 @@ describe("ExtensionClient", () => {
 			expect(client.isInitialized()).toBe(false)
 			expect(client.getCurrentState()).toBe(AgentLoopState.NO_TASK)
 		})
+
+		it("should reset mode on reset", () => {
+			const { client } = createMockClient()
+
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "code"))
+			expect(client.getCurrentMode()).toBe("code")
+
+			client.reset()
+
+			expect(client.getCurrentMode()).toBeUndefined()
+		})
+	})
+
+	describe("Mode tracking", () => {
+		it("should return undefined mode when not initialized", () => {
+			const { client } = createMockClient()
+			expect(client.getCurrentMode()).toBeUndefined()
+		})
+
+		it("should track mode from state messages", () => {
+			const { client } = createMockClient()
+
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "code"))
+
+			expect(client.getCurrentMode()).toBe("code")
+		})
+
+		it("should update mode when it changes", () => {
+			const { client } = createMockClient()
+
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "code"))
+			expect(client.getCurrentMode()).toBe("code")
+
+			client.handleMessage(createStateMessage([createMessage({ say: "text", ts: Date.now() + 1 })], "architect"))
+			expect(client.getCurrentMode()).toBe("architect")
+		})
+
+		it("should preserve mode when state message has no mode", () => {
+			const { client } = createMockClient()
+
+			// Set initial mode
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "code"))
+			expect(client.getCurrentMode()).toBe("code")
+
+			// State update without mode - should preserve existing mode
+			client.handleMessage(createStateMessage([createMessage({ say: "text", ts: Date.now() + 1 })]))
+			expect(client.getCurrentMode()).toBe("code")
+		})
+
+		it("should preserve mode when task is cleared", () => {
+			const { client } = createMockClient()
+
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "architect"))
+			expect(client.getCurrentMode()).toBe("architect")
+
+			client.clearTask()
+			// Mode should be preserved after clear
+			expect(client.getCurrentMode()).toBe("architect")
+		})
 	})
 })