Skip to content

Commit f150bad

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

File tree

2 files changed

+95
-60
lines changed

2 files changed

+95
-60
lines changed

‎agent/agentcontainers/api.go

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -663,11 +663,7 @@ func (api *API) getContainers() (codersdk.WorkspaceAgentListContainersResponse,
663663
for _, dc := range api.knownDevcontainers {
664664
// Include the agent if it's been created (we're iterating over
665665
// copies, so mutating is fine).
666-
//
667-
// NOTE(mafredri): We could filter on "proc.containerID == dc.Container.ID"
668-
// here but not doing so allows us to do some tricks in the UI to
669-
// make the experience more responsive for now.
670-
if proc := api.injectedSubAgentProcs[dc.WorkspaceFolder]; proc.agent.ID != uuid.Nil {
666+
if proc := api.injectedSubAgentProcs[dc.WorkspaceFolder]; proc.agent.ID != uuid.Nil && dc.Container != nil && proc.containerID == dc.Container.ID {
671667
dc.Agent = &codersdk.WorkspaceAgentDevcontainerAgent{
672668
ID: proc.agent.ID,
673669
Name: proc.agent.Name,
@@ -762,6 +758,7 @@ func (api *API) handleDevcontainerRecreate(w http.ResponseWriter, r *http.Reques
762758
// Update the status so that we don't try to recreate the
763759
// devcontainer multiple times in parallel.
764760
dc.Status = codersdk.WorkspaceAgentDevcontainerStatusStarting
761+
dc.Container = nil
765762
api.knownDevcontainers[dc.WorkspaceFolder] = dc
766763
api.asyncWg.Add(1)
767764
go api.recreateDevcontainer(dc, configPath)

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

Lines changed: 93 additions & 55 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 } 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,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}/containers/devcontainers/container/${devcontainer.container?.id}/recreate`,
97-
{
98-
method: "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" />
196234
<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

Comments
 (0)