Skip to content

Commit 4032530

Browse files
mafredrimafredri
committed
refactor to use mutation
1 parent e1ea4bf commit 4032530

File tree

1 file changed

+98
-56
lines changed

1 file changed

+98
-56
lines changed

‎site/src/modules/resources/AgentDevcontainerCard.tsx

Lines changed: 98 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import type {
44
Workspace,
55
WorkspaceAgent,
66
WorkspaceAgentDevcontainer,
7+
WorkspaceAgentListContainersResponse,
78
} from "api/typesGenerated";
89
import { Button } from "components/Button/Button";
910
import { displayError } from "components/GlobalSnackbar/utils";
@@ -20,7 +21,8 @@ import { Container, ExternalLinkIcon } from "lucide-react";
2021
import { useFeatureVisibility } from "modules/dashboard/useFeatureVisibility";
2122
import { AppStatuses } from "pages/WorkspacePage/AppStatuses";
2223
import type { FC } from "react";
23-
import { useEffect, useState } from "react";
24+
import { useEffect, useMemo } from "react";
25+
import { useMutation, useQueryClient } from "react-query";
2426
import { portForwardURL } from "utils/portForward";
2527
import { AgentApps, organizeAgentApps } from "./AgentApps/AgentApps";
2628
import { AgentButton } from "./AgentButton";
@@ -51,18 +53,16 @@ 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
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

6767
const appSections = (subAgent && organizeAgentApps(subAgent.apps)) || [];
6868
const displayApps =
@@ -80,64 +80,106 @@ export const AgentDevcontainerCard: FC<AgentDevcontainerCardProps> = ({
8080
showVSCode ||
8181
appSections.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 () => {
9585
const 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
);
10189
if (!response.ok) {
10290
const errorData = await response.json().catch(() => ({}));
10391
throw new Error(
104-
errorData.message || `Failed to recreate: ${response.statusText}`,
92+
errorData.message || `Failed to rebuild: ${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) {
114150
const errorMessage =
115151
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-
};
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+
);
125163
useEffect(() => {
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

142184
const appsClasses = "flex flex-wrap gap-4 empty:hidden md:justify-start";
143185

@@ -172,15 +214,15 @@ export const AgentDevcontainerCard: FC<AgentDevcontainerCardProps> = ({
172214
md: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
186228
devcontainer={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> = ({
203245
variant="outline"
204246
size="sm"
205247
onClick={handleRebuildDevcontainer}
206-
disabled={isRebuilding}
248+
disabled={devcontainer.status === "starting"}
207249
>
208-
<Spinner loading={isRebuilding} />
250+
<Spinner loading={devcontainer.status === "starting"} />
209251
Rebuild
210252
</Button>
211253

0 commit comments

Comments
 (0)