@@ -4,6 +4,7 @@ import type {
44 Workspace ,
55 WorkspaceAgent ,
66 WorkspaceAgentDevcontainer ,
7+ WorkspaceAgentListContainersResponse ,
78} from "api/typesGenerated" ;
89import { Button } from "components/Button/Button" ;
910import { displayError } from "components/GlobalSnackbar/utils" ;
@@ -20,7 +21,8 @@ import { Container, ExternalLinkIcon } from "lucide-react";
2021import { useFeatureVisibility } from "modules/dashboard/useFeatureVisibility" ;
2122import { AppStatuses } from "pages/WorkspacePage/AppStatuses" ;
2223import type { FC } from "react" ;
23- import { useEffect , useState } from "react" ;
24+ import { useEffect } from "react" ;
25+ import { useMutation , useQueryClient } from "react-query" ;
2426import { portForwardURL } from "utils/portForward" ;
2527import { AgentApps , organizeAgentApps } from "./AgentApps/AgentApps" ;
2628import { AgentButton } from "./AgentButton" ;
@@ -51,12 +53,7 @@ export const AgentDevcontainerCard: FC<AgentDevcontainerCardProps> = ({
5153} ) => {
5254 const { browser_only } = useFeatureVisibility ( ) ;
5355 const { proxy } = useProxy ( ) ;
54-
55- const [ isRebuilding , setIsRebuilding ] = useState ( false ) ;
56-
57- // Track sub agent removal state to improve UX. This will not be needed once
58- // the devcontainer and agent responses are aligned.
59- const [ subAgentRemoved , setSubAgentRemoved ] = useState ( false ) ;
56+ const queryClient = useQueryClient ( ) ;
6057
6158 // The sub agent comes from the workspace response whereas the devcontainer
6259 // comes from the agent containers endpoint. We need alignment between the
@@ -80,64 +77,105 @@ export const AgentDevcontainerCard: FC<AgentDevcontainerCardProps> = ({
8077 showVSCode ||
8178 appSections . some ( ( it ) => it . apps . length > 0 ) ;
8279
83- const showDevcontainerControls =
84- ! subAgentRemoved && subAgent && devcontainer . container ;
85- const showSubAgentApps =
86- ! subAgentRemoved && subAgent ?. status === "connected" && hasAppsToDisplay ;
87- const showSubAgentAppsPlaceholders =
88- subAgentRemoved || subAgent ?. status === "connecting" ;
89-
90- const handleRebuildDevcontainer = async ( ) => {
91- setIsRebuilding ( true ) ;
92- setSubAgentRemoved ( true ) ;
93- let rebuildSucceeded = false ;
94- try {
80+ const rebuildDevcontainerMutation = useMutation ( {
81+ mutationFn : async ( ) => {
9582 const response = await fetch (
9683 `/api/v2/workspaceagents/${ parentAgent . id } ${ devcontainer . container ?. id } ,
97- {
98- m
A3E2
ethod : "POST" ,
99- } ,
84+ { method : "POST" } ,
10085 ) ;
10186 if ( ! response . ok ) {
10287 const errorData = await response . json ( ) . catch ( ( ) => ( { } ) ) ;
10388 throw new Error (
104- errorData . message || `Failed to recreate : ${ response . statusText } ,
89+ errorData . message || `Failed to rebuild : ${ response . statusText } ,
10590 ) ;
10691 }
107- // If the request was accepted (e.g. 202), we mark it as succeeded.
108- // Once complete, the component will unmount, so the spinner will
109- // disappear with it.
110- if ( response . status === 202 ) {
111- rebuildSucceeded = true ;
92+ return response ;
93+ } ,
94+ onMutate : async ( ) => {
95+ await queryClient . cancelQueries ( {
96+ queryKey : [ "agents" , parentAgent . id , "containers" ] ,
97+ } ) ;
98+
99+ // Snapshot the previous data for rollback in case of error.
100+ const previousData = queryClient . getQueryData ( [
101+ "agents" ,
102+ parentAgent . id ,
103+ "containers" ,
104+ ] ) ;
105+
106+ // Optimistically update the devcontainer status to
107+ // "starting" and zero the agent and container to mimic what
108+ // the API does.
109+ queryClient . setQueryData (
110+ [ "agents" , parentAgent . id , "containers" ] ,
111+ ( oldData ?: WorkspaceAgentListContainersResponse ) => {
112+ if ( ! oldData ?. devcontainers ) return oldData ;
113+ return {
114+ ...oldData ,
115+ devcontainers : oldData . devcontainers . map ( ( dc ) => {
116+ if ( dc . id === devcontainer . id ) {
117+ return {
118+ ...dc ,
119+ agent : null ,
120+ container : null ,
121+ status : "starting" ,
122+ } ;
123+ }
124+ return dc ;
125+ } ) ,
126+ } ;
127+ } ,
128+ ) ;
129+
130+ return { previousData } ;
131+ } ,
132+ onSuccess : async ( ) => {
133+ // Invalidate the containers query to refetch updated data.
134+ await queryClient . invalidateQueries ( {
135+ queryKey : [ "agents" , parentAgent . id , "containers" ] ,
136+ } ) ;
137+ } ,
138+ onError : ( error , _ , context ) => {
139+ // If the mutation fails, use the context returned from
140+ // onMutate to roll back.
141+ if ( context ?. previousData ) {
142+ queryClient . setQueryData (
143+ [ "agents" , parentAgent . id , "containers" ] ,
144+ context . previousData ,
145+ ) ;
112146 }
113- } catch ( error ) {
114147 const errorMessage =
115148 error instanceof Error ? error . message : "An unknown error occurred." ;
116- displayError ( `Failed to recreate devcontainer: ${ errorMessage } ) ;
117- console . error ( "Failed to recreate devcontainer:" , error ) ;
118- } finally {
119- if ( ! rebuildSucceeded ) {
120- setIsRebuilding ( false ) ;
121- }
122- }
123- } ;
149+ displayError ( `Failed to rebuild devcontainer: ${ errorMessage } ) ;
150+ console . error ( "Failed to rebuild devcontainer:" , error ) ;
151+ } ,
152+ } ) ;
124153
154+ // Re-fetch containers when the subAgent changes to ensure data is
155+ // in sync.
156+ const latestSubAgentByName = subAgents . find (
157+ ( agent ) => agent . name === devcontainer . name ,
158+ ) ;
125159 useEffect ( ( ) => {
126- if ( subAgent ?. id ) {
127- setSubAgentRemoved ( false ) ;
128- } else {
129- setSubAgentRemoved ( true ) ;
160+ if ( ! latestSubAgentByName ) {
161+ return ;
130162 }
131- } , [ subAgent ?. id ] ) ;
163+ queryClient . invalidateQueries ( {
164+ queryKey : [ "agents" , parentAgent . id , "containers" ] ,
165+ } ) ;
166+ } , [ latestSubAgentByName , queryClient , parentAgent . id ] ) ;
132167
133- // If the devcontainer is starting, reflect this in the recreate button.
134- useEffect ( ( ) => {
135- if ( devcontainer . status === "starting" ) {
136- setIsRebuilding ( true ) ;
137- } else {
138- setIsRebuilding ( false ) ;
139- }
140- } , [ devcontainer ] ) ;
168+ const showDevcontainerControls = subAgent && devcontainer . container ;
169+ const showSubAgentApps =
170+ devcontainer . status !== "starting" &&
171+ subAgent ?. status === "connected" &&
172+ hasAppsToDisplay ;
173+ const showSubAgentAppsPlaceholders =
174+ devcontainer . status === "starting" || subAgent ?. status === "connecting" ;
175+
176+ const handleRebuildDevcontainer = ( ) => {
177+ rebuildDevcontainerMutation . mutate ( ) ;
178+ } ;
141179
142180 const appsClasses = "flex flex-wrap gap-4 empty:hidden md:justify-start" ;
143181
@@ -172,15 +210,15 @@ export const AgentDevcontainerCard: FC<AgentDevcontainerCardProps> = ({
172210 md:overflow-visible"
173211 >
174212 { subAgent ?. name ?? devcontainer . name }
175- { ! isRebuilding && devcontainer . container && (
213+ { devcontainer . container && (
176214 < span className = "text-content-tertiary" >
177215 { " " }
178216 ({ devcontainer . container . name } )
179217 </ span >
180218 ) }
181219 </ span >
182220 </ div >
183- { ! subAgentRemoved && subAgent ?. status === "connected" && (
221+ { subAgent ?. status === "connected" && (
184222 < >
185223 < SubAgentOutdatedTooltip
186224 devcontainer = { devcontainer }
@@ -190,7 +228,7 @@ export const AgentDevcontainerCard: FC<AgentDevcontainerCardProps> = ({
190228 < AgentLatency agent = { subAgent } />
191229 </ >
192230 ) }
193- { ! subAgentRemoved && subAgent ?. status === "connecting" && (
231+ { subAgent ?. status === "connecting" && (
194232 < >
195233 < Skeleton width = { 160 } variant = "text" />
196
179B
234 < Skeleton width = { 36 } variant = "text" />
@@ -203,9 +241,9 @@ export const AgentDevcontainerCard: FC<AgentDevcontainerCardProps> = ({
203241 variant = "outline"
204242 size = "sm"
205243 onClick = { handleRebuildDevcontainer }
206- disabled = { isRebuilding }
244+ disabled = { devcontainer . status === "starting" }
207245 >
208- < Spinner loading = { isRebuilding } />
246+ < Spinner loading = { devcontainer . status === "starting" } />
209247 Rebuild
210248 </ Button >
211249
0 commit comments