@@ -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 , useMemo } 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,18 +53,16 @@ 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
6360// two, so if the sub agent is not present or the IDs do not match, we
6461// assume it has been removed.
65- const subAgent = subAgents . find ( ( sub ) => sub . id === devcontainer . agent ?. id ) ;
62+ const subAgent = useMemo (
63+ ( ) => subAgents . find ( ( sub ) => sub . id === devcontainer . agent ?. id ) ,
64+ [ subAgents , devcontainer . agent ?. id ] ,
65+ ) ;
6666
6767const appSections = ( subAgent && organizeAgentApps ( subAgent . apps ) ) || [ ] ;
6868const displayApps =
@@ -80,64 +80,106 @@ export const AgentDevcontainerCard: FC<AgentDevcontainerCardProps> = ({
8080showVSCode ||
8181appSections . some ( ( it ) => it . apps . length > 0 ) ;
8282
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 {
83+ const rebuildDevcontainerMutation = useMutation ( {
84+ mutationFn :async ( ) => {
9585const response = await fetch (
9686`/api/v2/workspaceagents/${ parentAgent . id } /containers/devcontainers/container/${ devcontainer . container ?. id } /recreate` ,
97- {
98- method :"POST" ,
99- } ,
87+ { method :"POST" } ,
10088) ;
10189if ( ! response . ok ) {
10290const errorData = await response . json ( ) . catch ( ( ) => ( { } ) ) ;
10391throw new Error (
104- errorData . message || `Failed torecreate :${ response . statusText } ` ,
92+ errorData . message || `Failed torebuild :${ response . statusText } ` ,
10593) ;
10694}
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 ;
95+ return response ;
96+ } ,
97+ onMutate :async ( ) => {
98+ await queryClient . cancelQueries ( {
99+ queryKey :[ "agents" , parentAgent . id , "containers" ] ,
100+ } ) ;
101+
102+ // Snapshot the previous data for rollback in case of error.
103+ const previousData = queryClient . getQueryData ( [
104+ "agents" ,
105+ parentAgent . id ,
106+ "containers" ,
107+ ] ) ;
108+
109+ // Optimistically update the devcontainer status to
110+ // "starting" and zero the agent and container to mimic what
111+ // the API does.
112+ queryClient . setQueryData (
113+ [ "agents" , parentAgent . id , "containers" ] ,
114+ ( oldData ?:WorkspaceAgentListContainersResponse ) => {
115+ if ( ! oldData ?. devcontainers ) return oldData ;
116+ return {
117+ ...oldData ,
118+ devcontainers :oldData . devcontainers . map ( ( dc ) => {
119+ if ( dc . id === devcontainer . id ) {
120+ return {
121+ ...dc ,
122+ agent :null ,
123+ container :null ,
124+ status :"starting" ,
125+ } ;
126+ }
127+ return dc ;
128+ } ) ,
129+ } ;
130+ } ,
131+ ) ;
132+
133+ return { previousData} ;
134+ } ,
135+ onSuccess :async ( ) => {
136+ // Invalidate the containers query to refetch updated data.
137+ await queryClient . invalidateQueries ( {
138+ queryKey :[ "agents" , parentAgent . id , "containers" ] ,
139+ } ) ;
140+ } ,
141+ onError :( error , _ , context ) => {
142+ // If the mutation fails, use the context returned from
143+ // onMutate to roll back.
144+ if ( context ?. previousData ) {
145+ queryClient . setQueryData (
146+ [ "agents" , parentAgent . id , "containers" ] ,
147+ context . previousData ,
148+ ) ;
112149}
113- } catch ( error ) {
114150const errorMessage =
115151error 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- } ;
152+ displayError ( `Failed to rebuild devcontainer:${ errorMessage } ` ) ;
153+ console . error ( "Failed to rebuild devcontainer:" , error ) ;
154+ } ,
155+ } ) ;
124156
157+ // Re-fetch containers when the subAgent changes to ensure data is
158+ // in sync.
159+ const latestSubAgentByName = useMemo (
160+ ( ) => subAgents . find ( ( agent ) => agent . name === devcontainer . name ) ,
161+ [ subAgents , devcontainer . name ] ,
162+ ) ;
125163useEffect ( ( ) => {
126- if ( subAgent ?. id ) {
127- setSubAgentRemoved ( false ) ;
128- } else {
129- setSubAgentRemoved ( true ) ;
164+ if ( ! latestSubAgentByName ) {
165+ return ;
130166}
131- } , [ subAgent ?. id ] ) ;
167+ queryClient . invalidateQueries ( {
168+ queryKey :[ "agents" , parentAgent . id , "containers" ] ,
169+ } ) ;
170+ } , [ latestSubAgentByName , queryClient , parentAgent . id ] ) ;
132171
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 ] ) ;
172+ const showDevcontainerControls = subAgent && devcontainer . container ;
173+ const showSubAgentApps =
174+ devcontainer . status !== "starting" &&
175+ subAgent ?. status === "connected" &&
176+ hasAppsToDisplay ;
177+ const showSubAgentAppsPlaceholders =
178+ devcontainer . status === "starting" || subAgent ?. status === "connecting" ;
179+
180+ const handleRebuildDevcontainer = ( ) => {
181+ rebuildDevcontainerMutation . mutate ( ) ;
182+ } ;
141183
142184const appsClasses = "flex flex-wrap gap-4 empty:hidden md:justify-start" ;
143185
@@ -172,15 +214,15 @@ export const AgentDevcontainerCard: FC<AgentDevcontainerCardProps> = ({
172214md:overflow-visible"
173215>
174216{ subAgent ?. name ?? devcontainer . name }
175- { ! isRebuilding && devcontainer . container && (
217+ { devcontainer . container && (
176218< span className = "text-content-tertiary" >
177219{ " " }
178220({ devcontainer . container . name } )
179221</ span >
180222) }
181223</ span >
182224</ div >
183- { ! subAgentRemoved && subAgent ?. status === "connected" && (
225+ { subAgent ?. status === "connected" && (
184226< >
185227< SubAgentOutdatedTooltip
186228devcontainer = { devcontainer }
@@ -190,7 +232,7 @@ export const AgentDevcontainerCard: FC<AgentDevcontainerCardProps> = ({
190232< AgentLatency agent = { subAgent } />
191233</ >
192234) }
193- { ! subAgentRemoved && subAgent ?. status === "connecting" && (
235+ { subAgent ?. status === "connecting" && (
194236< >
195237< Skeleton width = { 160 } variant = "text" />
196238< Skeleton width = { 36 } variant = "text" />
@@ -203,9 +245,9 @@ export const AgentDevcontainerCard: FC<AgentDevcontainerCardProps> = ({
203245variant = "outline"
204246size = "sm"
205247onClick = { handleRebuildDevcontainer }
206- disabled = { isRebuilding }
248+ disabled = { devcontainer . status === "starting" }
207249>
208- < Spinner loading = { isRebuilding } />
250+ < Spinner loading = { devcontainer . status === "starting" } />
209251Rebuild
210252</ Button >
211253