@@ -4,6 +4,7 @@ import type {
44Workspace ,
55WorkspaceAgent ,
66WorkspaceAgentDevcontainer ,
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} ) => {
5254const { browser_only} = useFeatureVisibility ( ) ;
5355const { 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> = ({
8077showVSCode ||
8178appSections . 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 ( ) => {
9582const response = await fetch (
9683`/api/v2/workspaceagents/${ parentAgent . id } ${ devcontainer . container ?. id } ,
97- {
98- method :"POST" ,
99- } ,
84+ { method :"POST" } ,
10085) ;
10186if ( ! response . ok ) {
10287const errorData = await response . json ( ) . catch ( ( ) => ( { } ) ) ;
10388throw new Error (
104- errorData . message || `Failed torecreate :${ response . statusText } ,
89+ errorData . message || `Failed torebuild :${ 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 ) {
114147const errorMessage =
115148error 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+ ) ;
125159useEffect ( ( ) => {
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
142180const appsClasses = "flex flex-wrap gap-4 empty:hidden md:justify-start" ;
143181
@@ -172,15 +210,15 @@ export const AgentDevcontainerCard: FC<AgentDevcontainerCardProps> = ({
172210md: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
186224devcontainer = { 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" />
196234< Skeleton width = { 36 } variant = "text" />
@@ -203,9 +241,9 @@ export const AgentDevcontainerCard: FC<AgentDevcontainerCardProps> = ({
203241variant = "outline"
204242size = "sm"
205243onClick = { handleRebuildDevcontainer }
206- disabled = { isRebuilding }
244+ disabled = { devcontainer . status === "starting" }
207245>
208- < Spinner loading = { isRebuilding } />
246+ < Spinner loading = { devcontainer . status === "starting" } />
209247Rebuild
210248</ Button >
211249