Skip to content

Commit 9702d1a

Browse files
celestial-vaultcelestial-vault
authored
migrate relinquishControl (RooCodeInc#4098)
1 parent 922c795 commit 9702d1a

File tree

11 files changed

+140
-89
lines changed

11 files changed

+140
-89
lines changed

‎proto/ui.proto‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -252,4 +252,7 @@ service UiService {
252252

253253
// Initialize webview when it launches
254254
rpc initializeWebview(EmptyRequest) returns (Empty);
255+
256+
// Subscribe to relinquish control events
257+
rpc subscribeToRelinquishControl(EmptyRequest) returns (stream Empty);
255258
}

‎src/core/controller/index.ts‎

Lines changed: 2 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ import { sendStateUpdate } from "./state/subscribeToState"
4242
import { sendAddToInputEvent } from "./ui/subscribeToAddToInput"
4343
import { sendAuthCallbackEvent } from "./account/subscribeToAuthCallback"
4444
import { sendMcpMarketplaceCatalogEvent } from "./mcp/subscribeToMcpMarketplaceCatalog"
45+
import { sendRelinquishControlEvent } from "./ui/subscribeToRelinquishControl"
4546

4647
/*
4748
https://github.com/microsoft/vscode-webview-ui-toolkit-samples/blob/main/default/weather-webview/src/providers/WeatherViewProvider.ts
@@ -220,41 +221,6 @@ export class Controller {
220221
await this.fetchMcpMarketplace(message.bool)
221222
break
222223
}
223-
// case "openMcpMarketplaceServerDetails": {
224-
// if (message.text) {
225-
// const response = await fetch(`https://api.cline.bot/v1/mcp/marketplace/item?mcpId=${message.mcpId}`)
226-
// const details: McpDownloadResponse = await response.json()
227-
228-
// if (details.readmeContent) {
229-
// // Disable markdown preview markers
230-
// const config = vscode.workspace.getConfiguration("markdown")
231-
// await config.update("preview.markEditorSelection", false, true)
232-
233-
// // Create URI with base64 encoded markdown content
234-
// const uri = vscode.Uri.parse(
235-
// `${DIFF_VIEW_URI_SCHEME}:${details.name} README?${Buffer.from(details.readmeContent).toString("base64")}`,
236-
// )
237-
238-
// // close existing
239-
// const tabs = vscode.window.tabGroups.all
240-
// .flatMap((tg) => tg.tabs)
241-
// .filter((tab) => tab.label && tab.label.includes("README") && tab.label.includes("Preview"))
242-
// for (const tab of tabs) {
243-
// await vscode.window.tabGroups.close(tab)
244-
// }
245-
246-
// // Show only the preview
247-
// await vscode.commands.executeCommand("markdown.showPreview", uri, {
248-
// sideBySide: true,
249-
// preserveFocus: true,
250-
// })
251-
// }
252-
// }
253-
254-
// this.postMessageToWebview({ type: "relinquishControl" })
255-
256-
// break
257-
// }
258224

259225
// telemetry
260226
case "telemetrySetting": {
@@ -280,7 +246,7 @@ export class Controller {
280246
await this.deleteAllTaskHistory()
281247
await this.postStateToWebview()
282248
}
283-
this.postMessageToWebview({ type: "relinquishControl" })
249+
sendRelinquishControlEvent()
284250
break
285251
}
286252
case "grpc_request": {

‎src/core/controller/ui/subscribeToRelinquishControl.ts‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
import { Controller } from "../index"
2+
import { EmptyRequest, Empty } from "@shared/proto/common"
3+
import { StreamingResponseHandler, getRequestRegistry } from "../grpc-handler"
4+
5+
// Keep track of active subscriptions
6+
const activeRelinquishControlSubscriptions = new Set<StreamingResponseHandler>()
7+
8+
/**
9+
* Subscribe to relinquish control events
10+
* @param controller The controller instance
11+
* @param request The empty request
12+
* @param responseStream The streaming response handler
13+
* @param requestId The ID of the request (passed by the gRPC handler)
14+
*/
15+
export async function subscribeToRelinquishControl(
16+
controller: Controller,
17+
request: EmptyRequest,
18+
responseStream: StreamingResponseHandler,
19+
requestId?: string,
20+
): Promise<void> {
21+
// Add this subscription to the active subscriptions
22+
activeRelinquishControlSubscriptions.add(responseStream)
23+
24+
// Register cleanup when the connection is closed
25+
const cleanup = () => {
26+
activeRelinquishControlSubscriptions.delete(responseStream)
27+
}
28+
29+
// Register the cleanup function with the request registry if we have a requestId
30+
if (requestId) {
31+
getRequestRegistry().registerRequest(requestId, cleanup, { type: "relinquish_control_subscription" }, responseStream)
32+
}
33+
}
34+
35+
/**
36+
* Send a relinquish control event to all active subscribers
37+
*/
38+
export async function sendRelinquishControlEvent(): Promise<void> {
39+
// Send the event to all active subscribers
40+
const promises = Array.from(activeRelinquishControlSubscriptions).map(async (responseStream) => {
41+
try {
42+
const event = Empty.create({})
43+
await responseStream(
44+
event,
45+
false, // Not the last message
46+
)
47+
} catch (error) {
48+
console.error("Error sending relinquish control event:", error)
49+
// Remove the subscription if there was an error
50+
activeRelinquishControlSubscriptions.delete(responseStream)
51+
}
52+
})
53+
54+
await Promise.all(promises)
55+
}

‎src/core/task/index.ts‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ import { parseMentions } from "@core/mentions"
7373
import { formatResponse } from "@core/prompts/responses"
7474
import { addUserInstructions, SYSTEM_PROMPT } from "@core/prompts/system"
7575
import { sendPartialMessageEvent } from "@core/controller/ui/subscribeToPartialMessage"
76+
import { sendRelinquishControlEvent } from "@core/controller/ui/subscribeToRelinquishControl"
7677
import { convertClineMessageToProto } from "@shared/proto-conversions/cline-message"
7778
import { getContextWindowInfo } from "@core/context/context-management/context-window-utils"
7879
import { FileContextTracker } from "@core/context/context-tracking/FileContextTracker"
@@ -517,11 +518,11 @@ export class Task {
517518

518519
await this.saveClineMessagesAndUpdateHistory()
519520

520-
await this.postMessageToWebview({ type: "relinquishControl" })
521+
sendRelinquishControlEvent()
521522

522523
this.cancelTask() // the task is already cancelled by the provider beforehand, but we need to re-init to get the updated messages
523524
} else {
524-
await this.postMessageToWebview({ type: "relinquishControl" })
525+
sendRelinquishControlEvent()
525526
}
526527
}
527528

@@ -531,7 +532,7 @@ export class Task {
531532
}
532533

533534
const relinquishButton = () => {
534-
this.postMessageToWebview({ type: "relinquishControl" })
535+
sendRelinquishControlEvent()
535536
}
536537
if (!this.enableCheckpoints) {
537538
vscode.window.showInformationMessage("Checkpoints are disabled in settings. Cannot show diff.")

‎src/shared/ExtensionMessage.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ export interface ExtensionMessage {
2020
| "openAiModels"
2121
| "requestyModels"
2222
| "mcpServers"
23-
| "relinquishControl"
2423
| "mcpDownloadDetails"
2524
| "userCreditsBalance"
2625
| "userCreditsUsage"

‎webview-ui/src/components/chat/ChatRow.tsx‎

Lines changed: 7 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -237,7 +237,7 @@ export const ChatRowContent = ({
237237
sendMessageFromChatRow,
238238
onSetQuote,
239239
}: ChatRowContentProps) => {
240-
const { mcpServers, mcpMarketplaceCatalog } = useExtensionState()
240+
const { mcpServers, mcpMarketplaceCatalog, onRelinquishControl } = useExtensionState()
241241
const [seeNewChangesDisabled, setSeeNewChangesDisabled] = useState(false)
242242
const [quoteButtonState, setQuoteButtonState] = useState<QuoteButtonState>({
243243
visible: false,
@@ -269,17 +269,12 @@ export const ChatRowContent = ({
269269

270270
const type = message.type === "ask" ? message.ask : message.say
271271

272-
const handleMessage = useCallback((event: MessageEvent) => {
273-
const message: ExtensionMessage = event.data
274-
switch (message.type) {
275-
case "relinquishControl": {
276-
setSeeNewChangesDisabled(false)
277-
break
278-
}
279-
}
280-
}, [])
281-
282-
useEvent("message", handleMessage)
272+
// Use the onRelinquishControl hook instead of message event
273+
useEffect(() => {
274+
return onRelinquishControl(() => {
275+
setSeeNewChangesDisabled(false)
276+
})
277+
}, [onRelinquishControl])
283278

284279
// --- Quote Button Logic ---
285280
// MOVE handleQuoteClick INSIDE ChatRowContent

‎webview-ui/src/components/common/CheckmarkControl.tsx‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,14 @@
11
import { CODE_BLOCK_BG_COLOR } from "@/components/common/CodeBlock"
22
import { CheckpointsServiceClient } from "@/services/grpc-client"
33
import { flip, offset, shift, useFloating } from "@floating-ui/react"
4-
import { ExtensionMessage } from "@shared/ExtensionMessage"
54
import { CheckpointRestoreRequest } from "@shared/proto/checkpoints"
65
import { Int64Request } from "@shared/proto/common"
76
import { ClineCheckpointRestore } from "@shared/WebviewMessage"
87
import { VSCodeButton } from "@vscode/webview-ui-toolkit/react"
9-
import { useCallback, useEffect, useRef, useState } from "react"
8+
import { useEffect, useRef, useState } from "react"
109
import { createPortal } from "react-dom"
11-
import { useEvent } from "react-use"
1210
import styled from "styled-components"
11+
import { useExtensionState } from "@/context/ExtensionStateContext"
1312

1413
interface CheckmarkControlProps {
1514
messageTs?: number
@@ -25,6 +24,7 @@ export const CheckmarkControl = ({ messageTs, isCheckpointCheckedOut }: Checkmar
2524
const [hasMouseEntered, setHasMouseEntered] = useState(false)
2625
const containerRef = useRef<HTMLDivElement>(null)
2726
const tooltipRef = useRef<HTMLDivElement>(null)
27+
const { onRelinquishControl } = useExtensionState()
2828

2929
const { refs, floatingStyles, update, placement } = useFloating({
3030
placement: "bottom-end",
@@ -52,15 +52,16 @@ export const CheckmarkControl = ({ messageTs, isCheckpointCheckedOut }: Checkmar
5252
}
5353
}, [showRestoreConfirm, update])
5454

55-
const handleMessage = useCallback((event: MessageEvent<ExtensionMessage>) => {
56-
if (event.data.type === "relinquishControl") {
55+
// Use the onRelinquishControl hook instead of message event
56+
useEffect(() => {
57+
return onRelinquishControl(() => {
5758
setCompareDisabled(false)
5859
setRestoreTaskDisabled(false)
5960
setRestoreWorkspaceDisabled(false)
6061
setRestoreBothDisabled(false)
6162
setShowRestoreConfirm(false)
62-
}
63-
}, [])
63+
})
64+
}, [onRelinquishControl])
6465

6566
const handleRestoreTask = async () => {
6667
setRestoreTaskDisabled(true)
@@ -141,8 +142,6 @@ export const CheckmarkControl = ({ messageTs, isCheckpointCheckedOut }: Checkmar
141142
setHasMouseEntered(false)
142143
}
143144

144-
useEvent("message", handleMessage)
145-
146145
return (
147146
<Container isMenuOpen={showRestoreConfirm} $isCheckedOut={isCheckpointCheckedOut} onMouseLeave={handleControlsMouseLeave}>
148147
<i

‎webview-ui/src/components/common/CheckpointControls.tsx‎

Lines changed: 14 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
import { CODE_BLOCK_BG_COLOR } from "@/components/common/CodeBlock"
22
import { CheckpointsServiceClient } from "@/services/grpc-client"
3-
import { ExtensionMessage } from "@shared/ExtensionMessage"
43
import { CheckpointRestoreRequest } from "@shared/proto/checkpoints"
54
import { Int64Request } from "@shared/proto/common"
65
import { VSCodeButton } from "@vscode/webview-ui-toolkit/react"
7-
import { useCallback, useRef, useState } from "react"
8-
import { useClickAway, useEvent } from "react-use"
6+
import { useEffect, useRef, useState } from "react"
7+
import { useClickAway } from "react-use"
98
import styled from "styled-components"
9+
import { useExtensionState } from "@/context/ExtensionStateContext"
1010

1111
interface CheckpointOverlayProps {
1212
messageTs?: number
@@ -21,6 +21,7 @@ export const CheckpointOverlay = ({ messageTs }: CheckpointOverlayProps) => {
2121
const [hasMouseEntered, setHasMouseEntered] = useState(false)
2222
const containerRef = useRef<HTMLDivElement>(null)
2323
const tooltipRef = useRef<HTMLDivElement>(null)
24+
const { onRelinquishControl } = useExtensionState()
2425

2526
useClickAway(containerRef, () => {
2627
if (showRestoreConfirm) {
@@ -29,21 +30,16 @@ export const CheckpointOverlay = ({ messageTs }: CheckpointOverlayProps) => {
2930
}
3031
})
3132

32-
const handleMessage = useCallback((event: MessageEvent) => {
33-
const message: ExtensionMessage = event.data
34-
switch (message.type) {
35-
case "relinquishControl": {
36-
setCompareDisabled(false)
37-
setRestoreTaskDisabled(false)
38-
setRestoreWorkspaceDisabled(false)
39-
setRestoreBothDisabled(false)
40-
setShowRestoreConfirm(false)
41-
break
42-
}
43-
}
44-
}, [])
45-
46-
useEvent("message", handleMessage)
33+
// Use the onRelinquishControl hook instead of message event
34+
useEffect(() => {
35+
return onRelinquishControl(() => {
36+
setCompareDisabled(false)
37+
setRestoreTaskDisabled(false)
38+
setRestoreWorkspaceDisabled(false)
39+
setRestoreBothDisabled(false)
40+
setShowRestoreConfirm(false)
41+
})
42+
}, [onRelinquishControl])
4743

4844
const handleRestoreTask = async () => {
4945
setRestoreTaskDisabled(true)

‎webview-ui/src/components/history/HistoryView.tsx‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ const CustomFilterRadio = ({ checked, onChange, icon, label }: CustomFilterRadio
4949

5050
const HistoryView = ({ onDone }: HistoryViewProps) => {
5151
const extensionStateContext = useExtensionState()
52-
const { taskHistory, filePaths } = extensionStateContext
52+
const { taskHistory, filePaths, onRelinquishControl } = extensionStateContext
5353
const [searchQuery, setSearchQuery] = useState("")
5454
const [sortOption, setSortOption] = useState<SortOption>("newest")
5555
const [lastNonRelevantSort, setLastNonRelevantSort] = useState<SortOption | null>("newest")
@@ -130,12 +130,12 @@ const HistoryView = ({ onDone }: HistoryViewProps) => {
130130
[showFavoritesOnly, loadTaskHistory],
131131
)
132132

133-
const handleMessage = useCallback((event: MessageEvent<ExtensionMessage>) => {
134-
if (event.data.type === "relinquishControl") {
133+
// Use the onRelinquishControl hook instead of message event
134+
useEffect(() => {
135+
return onRelinquishControl(() => {
135136
setDeleteAllDisabled(false)
136-
}
137-
}, [])
138-
useEvent("message", handleMessage)
137+
})
138+
}, [onRelinquishControl])
139139

140140
const { totalTasksSize, setTotalTasksSize } = extensionStateContext
141141

‎webview-ui/src/components/mcp/configuration/tabs/marketplace/McpMarketplaceCard.tsx‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
import { McpServiceClient } from "@/services/grpc-client"
22
import { McpMarketplaceItem, McpServer } from "@shared/mcp"
33
import { StringRequest } from "@shared/proto/common"
4-
import { useCallback, useMemo, useRef, useState } from "react"
4+
import { useCallback, useEffect, useMemo, useRef, useState } from "react"
55
import { useEvent } from "react-use"
66
import styled from "styled-components"
7+
import { useExtensionState } from "@/context/ExtensionStateContext"
78

89
interface McpMarketplaceCardProps {
910
item: McpMarketplaceItem
@@ -15,21 +16,25 @@ const McpMarketplaceCard = ({ item, installedServers }: McpMarketplaceCardProps)
1516
const [isDownloading, setIsDownloading] = useState(false)
1617
const [isLoading, setIsLoading] = useState(false)
1718
const githubLinkRef = useRef<HTMLDivElement>(null)
19+
const { onRelinquishControl } = useExtensionState()
1820

1921
const handleMessage = useCallback((event: MessageEvent) => {
2022
const message = event.data
2123
switch (message.type) {
2224
case "mcpDownloadDetails":
2325
setIsDownloading(false)
2426
break
25-
case "relinquishControl":
26-
setIsLoading(false)
27-
break
2827
}
2928
}, [])
3029

3130
useEvent("message", handleMessage)
3231

32+
useEffect(() => {
33+
return onRelinquishControl(() => {
34+
setIsLoading(false)
35+
})
36+
}, [onRelinquishControl])
37+
3338
const githubAuthorUrl = useMemo(() => {
3439
const url = new URL(item.githubUrl)
3540
const pathParts = url.pathname.split("/")

0 commit comments

Comments
 (0)