From 4c362499f23c0499394c3aa2efe3aaac1928b6be Mon Sep 17 00:00:00 2001 From: Ethan Date: Wed, 20 May 2026 12:30:28 +1000 Subject: [PATCH] fix(site): keep workspace quota meter live-updated in Agents chat (#25468) The workspace quota meter in the Agents chat sidebar previously only refreshed on a full page reload, while the AI cost-usage meter rendered next to it already polled every 60 seconds. This change gives the Agents usage indicator the same 60s workspace-quota polling and refreshes derived workspace caches immediately when workspace-affecting chat tools complete. `create_workspace`, `start_workspace`, and `stop_workspace` now invalidate the workspace quota and `workspaces` query family, so both the credit numbers and the workspace-count detail stay in sync. `create_workspace` still invalidates the chat record to resolve workspace bindings, and archive-and-delete uses the shared workspace mutation invalidation helper so deleting an agent workspace refreshes the same derived workspace data. Closes CODAGT-444 Tested manually and it works perfectly. --- site/src/api/queries/workspaces.test.ts | 165 ++++++++++ site/src/api/queries/workspaces.ts | 55 +++- site/src/modules/dashboard/useDashboard.ts | 4 + site/src/pages/AgentsPage/AgentChatPage.tsx | 16 +- site/src/pages/AgentsPage/AgentsPage.tsx | 17 +- .../useChatToolInvalidations.test.tsx | 297 +++++++++++++++++ .../useChatToolInvalidations.ts | 102 ++++++ .../useWorkspaceCreationWatcher.test.tsx | 301 ------------------ .../useWorkspaceCreationWatcher.ts | 76 ----- .../AgentsPage/components/UsageIndicator.tsx | 9 +- 10 files changed, 653 insertions(+), 389 deletions(-) create mode 100644 site/src/api/queries/workspaces.test.ts create mode 100644 site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.test.tsx create mode 100644 site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.ts delete mode 100644 site/src/pages/AgentsPage/components/ChatConversation/useWorkspaceCreationWatcher.test.tsx delete mode 100644 site/src/pages/AgentsPage/components/ChatConversation/useWorkspaceCreationWatcher.ts diff --git a/site/src/api/queries/workspaces.test.ts b/site/src/api/queries/workspaces.test.ts new file mode 100644 index 0000000000..a4ef076b1a --- /dev/null +++ b/site/src/api/queries/workspaces.test.ts @@ -0,0 +1,165 @@ +import { QueryClient } from "react-query"; +import { describe, expect, it } from "vitest"; +import type { WorkspacesResponse } from "#/api/typesGenerated"; +import { getWorkspaceQuotaQueryKey } from "./workspaceQuota"; +import { + autoCreateWorkspace, + buildLogsKey, + createWorkspace, + invalidateWorkspaceListQueries, + invalidateWorkspaceMutationQueries, + workspacesKey, + workspacesQueryKeyPrefix, + workspaceUsage, +} from "./workspaces"; + +const createTestQueryClient = (): QueryClient => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: Number.POSITIVE_INFINITY, + refetchOnWindowFocus: false, + networkMode: "offlineFirst", + }, + }, + }); + +const workspacesResponse = { + workspaces: [], + count: 0, +} satisfies WorkspacesResponse; + +const seedWorkspaceFamilyQueries = (queryClient: QueryClient) => { + const rawListKey = workspacesQueryKeyPrefix; + const defaultListKey = workspacesKey({}); + const filteredListKey = workspacesKey({ + q: "owner:me organization:default", + limit: 25, + offset: 50, + }); + const usageKey = workspaceUsage({ + usageApp: "reconnecting-pty", + connectionStatus: "connected", + workspaceId: "workspace-1", + agentId: "agent-1", + }).queryKey; + const buildLogs = buildLogsKey("workspace-1"); + const workspacePermissionsKey = [ + "workspaces", + "workspace-1", + "permissions", + ] as const; + const workspaceAgentCredentialsKey = [ + "workspaces", + "workspace-1", + "agents", + "main", + "credentials", + ] as const; + const organizationWorkspacePermissionsKey = [ + "workspaces", + ["organization-1"], + "permissions", + ] as const; + + queryClient.setQueryData(rawListKey, workspacesResponse); + queryClient.setQueryData(defaultListKey, workspacesResponse); + queryClient.setQueryData(filteredListKey, workspacesResponse); + queryClient.setQueryData(usageKey, { tracked: true }); + queryClient.setQueryData(buildLogs, []); + queryClient.setQueryData(workspacePermissionsKey, { read: true }); + queryClient.setQueryData(workspaceAgentCredentialsKey, { token: "secret" }); + queryClient.setQueryData(organizationWorkspacePermissionsKey, { read: true }); + + return { + listKeys: [rawListKey, defaultListKey, filteredListKey], + nonListKeys: [ + usageKey, + buildLogs, + workspacePermissionsKey, + workspaceAgentCredentialsKey, + organizationWorkspacePermissionsKey, + ], + }; +}; + +describe("invalidateWorkspaceListQueries", () => { + it("invalidates workspace list queries without touching side-effecting workspace-family queries", async () => { + const queryClient = createTestQueryClient(); + const { listKeys, nonListKeys } = seedWorkspaceFamilyQueries(queryClient); + + await invalidateWorkspaceListQueries(queryClient); + + for (const key of listKeys) { + expect( + queryClient.getQueryState(key)?.isInvalidated, + `${JSON.stringify(key)} should be invalidated`, + ).toBe(true); + } + for (const key of nonListKeys) { + expect( + queryClient.getQueryState(key)?.isInvalidated, + `${JSON.stringify(key)} should NOT be invalidated`, + ).not.toBe(true); + } + }); +}); + +describe("invalidateWorkspaceMutationQueries", () => { + it("uses narrowed list invalidation and keeps workspace usage queries untouched", async () => { + const queryClient = createTestQueryClient(); + const { listKeys, nonListKeys } = seedWorkspaceFamilyQueries(queryClient); + const quotaKey = getWorkspaceQuotaQueryKey("default", "me"); + queryClient.setQueryData(quotaKey, { credits_consumed: 1, budget: 10 }); + + await invalidateWorkspaceMutationQueries(queryClient, { + organizationName: "default", + username: "me", + }); + + for (const key of listKeys) { + expect( + queryClient.getQueryState(key)?.isInvalidated, + `${JSON.stringify(key)} should be invalidated`, + ).toBe(true); + } + expect(queryClient.getQueryState(quotaKey)?.isInvalidated).toBe(true); + for (const key of nonListKeys) { + expect( + queryClient.getQueryState(key)?.isInvalidated, + `${JSON.stringify(key)} should NOT be invalidated`, + ).not.toBe(true); + } + }); +}); + +describe("workspace creation mutations", () => { + it("use narrowed list invalidation for manual workspace creation", async () => { + const queryClient = createTestQueryClient(); + const { listKeys, nonListKeys } = seedWorkspaceFamilyQueries(queryClient); + + await createWorkspace(queryClient).onSuccess(); + + for (const key of listKeys) { + expect(queryClient.getQueryState(key)?.isInvalidated).toBe(true); + } + for (const key of nonListKeys) { + expect(queryClient.getQueryState(key)?.isInvalidated).not.toBe(true); + } + }); + + it("use narrowed list invalidation for auto workspace creation", async () => { + const queryClient = createTestQueryClient(); + const { listKeys, nonListKeys } = seedWorkspaceFamilyQueries(queryClient); + + await autoCreateWorkspace(queryClient).onSuccess(); + + for (const key of listKeys) { + expect(queryClient.getQueryState(key)?.isInvalidated).toBe(true); + } + for (const key of nonListKeys) { + expect(queryClient.getQueryState(key)?.isInvalidated).not.toBe(true); + } + }); +}); diff --git a/site/src/api/queries/workspaces.ts b/site/src/api/queries/workspaces.ts index 42e1988cfd..5782c32d18 100644 --- a/site/src/api/queries/workspaces.ts +++ b/site/src/api/queries/workspaces.ts @@ -32,6 +32,9 @@ import { import { checkAuthorization } from "./authCheck"; import { disabledRefetchOptions } from "./util"; import { workspaceBuildsKey } from "./workspaceBuilds"; +import { getWorkspaceQuotaQueryKey } from "./workspaceQuota"; + +export const workspacesQueryKeyPrefix = ["workspaces"] as const; export const workspaceByOwnerAndNameKey = ( ownerUsername: string, @@ -126,7 +129,7 @@ export const createWorkspace = (queryClient: QueryClient) => { return API.createWorkspace(userId, req); }, onSuccess: async () => { - await queryClient.invalidateQueries({ queryKey: ["workspaces"] }); + await invalidateWorkspaceListQueries(queryClient); }, }; }; @@ -185,7 +188,7 @@ export const autoCreateWorkspace = (queryClient: QueryClient) => { }); }, onSuccess: async () => { - await queryClient.invalidateQueries({ queryKey: ["workspaces"] }); + await invalidateWorkspaceListQueries(queryClient); }, }; }; @@ -212,7 +215,7 @@ async function findMatchWorkspace(q: string): Promise { } export function workspacesKey(req: WorkspacesRequest = {}) { - return ["workspaces", req] as const; + return [...workspacesQueryKeyPrefix, req] as const; } export function workspaces(req: WorkspacesRequest = {}) { @@ -222,6 +225,52 @@ export function workspaces(req: WorkspacesRequest = {}) { } as const satisfies QueryOptions; } +const isWorkspacesListQuery = (query: { + queryKey: readonly unknown[]; +}): boolean => { + const key = query.queryKey; + if (key.length === 1) { + return true; + } + if (key.length !== 2) { + return false; + } + const segment = key[1]; + return ( + segment !== null && typeof segment === "object" && !Array.isArray(segment) + ); +}; + +export const invalidateWorkspaceListQueries = (queryClient: QueryClient) => { + return queryClient.invalidateQueries({ + queryKey: workspacesQueryKeyPrefix, + predicate: isWorkspacesListQuery, + }); +}; + +interface WorkspaceMutationInvalidationOptions { + organizationName: string; + username: string; +} + +export async function invalidateWorkspaceMutationQueries( + queryClient: QueryClient, + { organizationName, username }: WorkspaceMutationInvalidationOptions, +): Promise { + const invalidations = [invalidateWorkspaceListQueries(queryClient)]; + + if (organizationName !== "") { + invalidations.push( + queryClient.invalidateQueries({ + queryKey: getWorkspaceQuotaQueryKey(organizationName, username), + exact: true, + }), + ); + } + + await Promise.all(invalidations); +} + export const updateDeadline = ( workspace: Workspace, ): UseMutationOptions => { diff --git a/site/src/modules/dashboard/useDashboard.ts b/site/src/modules/dashboard/useDashboard.ts index a7dcad5f4a..ebee00fba1 100644 --- a/site/src/modules/dashboard/useDashboard.ts +++ b/site/src/modules/dashboard/useDashboard.ts @@ -12,3 +12,7 @@ export const useDashboard = (): DashboardValue => { return context; }; + +export const getDefaultOrganizationName = ( + organizations: DashboardValue["organizations"], +): string => organizations.find((org) => org.is_default)?.name ?? ""; diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index 3e01bdb2a8..acade0b3d0 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -47,6 +47,10 @@ import type * as TypesGen from "#/api/typesGenerated"; import type { ChatMessagePart } from "#/api/typesGenerated"; import { useProxy } from "#/contexts/ProxyContext"; import { useAuthenticated } from "#/hooks/useAuthenticated"; +import { + getDefaultOrganizationName, + useDashboard, +} from "#/modules/dashboard/useDashboard"; import { isMobileViewport } from "#/utils/mobile"; import { pageTitle } from "#/utils/page"; import { rewriteLocalhostURL } from "#/utils/portForward"; @@ -71,7 +75,7 @@ import { useChatSelector, useChatStore, } from "./components/ChatConversation/chatStore"; -import { useWorkspaceCreationWatcher } from "./components/ChatConversation/useWorkspaceCreationWatcher"; +import { useChatToolInvalidations } from "./components/ChatConversation/useChatToolInvalidations"; import type { PendingAttachment } from "./components/ChatPageContent"; import { getDefaultMCPSelection, @@ -667,6 +671,8 @@ const AgentChatPage: FC = () => { } = useOutletContext(); const queryClient = useQueryClient(); const { permissions, user: currentUser } = useAuthenticated(); + const { organizations } = useDashboard(); + const organizationName = getDefaultOrganizationName(organizations); const [selectedModel, setSelectedModel] = useState(""); const scrollToBottomRef = useRef<(() => void) | null>(null); const chatInputRef = useRef(null); @@ -1003,11 +1009,13 @@ const AgentChatPage: FC = () => { agentStatus: workspaceAgent?.status, }); - // Detect workspace creation so the sidebar can resolve the - // workspace and display agent/git info. - useWorkspaceCreationWatcher({ + // Detect completed chat tool results so sidebar data stays in sync + // with the server state those tools may have changed. + useChatToolInvalidations({ store, chatID: agentId, + organizationName, + username: currentUser.username, }); const handleCommit = (repoRoot: string) => { diff --git a/site/src/pages/AgentsPage/AgentsPage.tsx b/site/src/pages/AgentsPage/AgentsPage.tsx index c8562c664f..056ecd6224 100644 --- a/site/src/pages/AgentsPage/AgentsPage.tsx +++ b/site/src/pages/AgentsPage/AgentsPage.tsx @@ -34,11 +34,18 @@ import { updateInfiniteChatsCache, userChatPersonalModelOverrides, } from "#/api/queries/chats"; -import { workspaceById } from "#/api/queries/workspaces"; +import { + invalidateWorkspaceMutationQueries, + workspaceById, +} from "#/api/queries/workspaces"; import type * as TypesGen from "#/api/typesGenerated"; import { ConfirmDialog } from "#/components/Dialogs/ConfirmDialog/ConfirmDialog"; import { DeleteDialog } from "#/components/Dialogs/DeleteDialog/DeleteDialog"; import { useAuthenticated } from "#/hooks/useAuthenticated"; +import { + getDefaultOrganizationName, + useDashboard, +} from "#/modules/dashboard/useDashboard"; import { createReconnectingWebSocket } from "#/utils/reconnectingWebSocket"; import { AgentsPageView } from "./AgentsPageView"; import { emptyInputStorageKey } from "./components/AgentCreateForm"; @@ -66,7 +73,9 @@ const AgentsPage: FC = () => { const navigate = useNavigate(); const location = useLocation(); const { agentId } = useParams(); - const { permissions } = useAuthenticated(); + const { permissions, user } = useAuthenticated(); + const { organizations } = useDashboard(); + const organizationName = getDefaultOrganizationName(organizations); const isAgentsAdmin = permissions.editDeploymentConfig; const [archivedFilter, setArchivedFilter] = useArchivedFilterParam(); @@ -200,6 +209,10 @@ const AgentsPage: FC = () => { await queryClient.invalidateQueries({ queryKey: chatsByWorkspaceKeyPrefix, }); + await invalidateWorkspaceMutationQueries(queryClient, { + organizationName, + username: user.username, + }); }, onError: (error) => { toast.error( diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.test.tsx new file mode 100644 index 0000000000..f6d393cb25 --- /dev/null +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.test.tsx @@ -0,0 +1,297 @@ +import { renderHook, waitFor } from "@testing-library/react"; +import type { FC, PropsWithChildren } from "react"; +import { act } from "react"; +import { QueryClient, QueryClientProvider } from "react-query"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { getWorkspaceQuotaQueryKey } from "#/api/queries/workspaceQuota"; +import { workspacesQueryKeyPrefix } from "#/api/queries/workspaces"; +import { createChatStore } from "./chatStore"; +import type { StreamState } from "./types"; +import { useChatToolInvalidations } from "./useChatToolInvalidations"; + +const ORGANIZATION_NAME = "coder"; +const USERNAME = "alice"; + +type ToolResultOverrides = Partial; + +const createStreamState = ( + name: string, + id = "tool-1", + resultOverrides: ToolResultOverrides = {}, +): StreamState => ({ + blocks: [], + toolCalls: { + [id]: { + id, + name, + args: {}, + }, + }, + toolResults: { + [id]: { + id, + name, + isError: false, + ...resultOverrides, + }, + }, + sources: [], +}); + +const createTestStore = (initial: StreamState | null = null) => { + const store = createChatStore(); + store.setStreamState(initial); + + return { + store, + setStreamState: store.setStreamState, + }; +}; + +const createWrapper = (queryClient: QueryClient): FC => { + return ({ children }) => ( + {children} + ); +}; + +describe("useChatToolInvalidations", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + }); + + const renderInvalidations = ({ + chatID = "chat-1", + organizationName = ORGANIZATION_NAME, + username = USERNAME, + }: { + chatID?: string; + organizationName?: string; + username?: string; + } = {}) => { + const { store, setStreamState } = createTestStore(); + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + const result = renderHook( + (props: { chatID: string; organizationName: string; username: string }) => + useChatToolInvalidations({ + store, + ...props, + }), + { + initialProps: { chatID, organizationName, username }, + wrapper: createWrapper(queryClient), + }, + ); + + return { + ...result, + invalidateSpy, + setStreamState, + }; + }; + + it("dispatches chat binding and workspace mutation invalidations on create_workspace completion", async () => { + const { invalidateSpy, setStreamState } = renderInvalidations(); + + await act(async () => { + setStreamState(createStreamState("create_workspace")); + }); + + await waitFor(() => { + expect(invalidateSpy).toHaveBeenCalledWith({ + queryKey: ["chats", "chat-1"], + }); + expect(invalidateSpy).toHaveBeenCalledWith( + expect.objectContaining({ + queryKey: workspacesQueryKeyPrefix, + predicate: expect.any(Function), + }), + ); + expect(invalidateSpy).toHaveBeenCalledWith({ + queryKey: getWorkspaceQuotaQueryKey(ORGANIZATION_NAME, USERNAME), + exact: true, + }); + expect(invalidateSpy).toHaveBeenCalledTimes(3); + }); + }); + + it("dispatches workspace mutation invalidations on start_workspace completion", async () => { + const { invalidateSpy, setStreamState } = renderInvalidations(); + + await act(async () => { + setStreamState(createStreamState("start_workspace")); + }); + + await waitFor(() => { + expect(invalidateSpy).toHaveBeenCalledWith( + expect.objectContaining({ + queryKey: workspacesQueryKeyPrefix, + predicate: expect.any(Function), + }), + ); + expect(invalidateSpy).toHaveBeenCalledWith({ + queryKey: getWorkspaceQuotaQueryKey(ORGANIZATION_NAME, USERNAME), + exact: true, + }); + expect(invalidateSpy).toHaveBeenCalledTimes(2); + }); + + expect(invalidateSpy).not.toHaveBeenCalledWith({ + queryKey: ["chats", "chat-1"], + }); + }); + + it("dispatches workspace mutation invalidations for errored workspace tools", async () => { + const { invalidateSpy, setStreamState } = renderInvalidations(); + + await act(async () => { + setStreamState( + createStreamState("stop_workspace", "tool-1", { isError: true }), + ); + }); + + await waitFor(() => { + expect(invalidateSpy).toHaveBeenCalledWith( + expect.objectContaining({ + queryKey: workspacesQueryKeyPrefix, + predicate: expect.any(Function), + }), + ); + expect(invalidateSpy).toHaveBeenCalledWith({ + queryKey: getWorkspaceQuotaQueryKey(ORGANIZATION_NAME, USERNAME), + exact: true, + }); + expect(invalidateSpy).toHaveBeenCalledTimes(2); + }); + }); + + it("invalidates workspace queries without quota when the quota key is incomplete", async () => { + const { invalidateSpy, setStreamState } = renderInvalidations({ + organizationName: "", + }); + + await act(async () => { + setStreamState(createStreamState("create_workspace")); + }); + + await waitFor(() => { + expect(invalidateSpy).toHaveBeenCalledWith({ + queryKey: ["chats", "chat-1"], + }); + expect(invalidateSpy).toHaveBeenCalledWith( + expect.objectContaining({ + queryKey: workspacesQueryKeyPrefix, + predicate: expect.any(Function), + }), + ); + expect(invalidateSpy).toHaveBeenCalledTimes(2); + }); + }); + + it("does not invalidate queries for tools without invalidation signals", async () => { + const { invalidateSpy, setStreamState } = renderInvalidations(); + + await act(async () => { + setStreamState(createStreamState("read_file")); + }); + + expect(invalidateSpy).not.toHaveBeenCalled(); + }); + + it("does not process the same tool call ID twice", async () => { + const { invalidateSpy, setStreamState } = renderInvalidations(); + const streamState = createStreamState("create_workspace"); + + await act(async () => { + setStreamState(streamState); + }); + + await waitFor(() => { + expect(invalidateSpy).toHaveBeenCalledTimes(3); + }); + + await act(async () => { + setStreamState(streamState); + }); + + await waitFor(() => { + expect(invalidateSpy).toHaveBeenCalledTimes(3); + }); + }); + + it("resets processed tool call IDs when chatID changes", async () => { + const { invalidateSpy, rerender, setStreamState } = renderInvalidations(); + + await act(async () => { + setStreamState(createStreamState("create_workspace")); + }); + + await waitFor(() => { + expect(invalidateSpy).toHaveBeenCalledTimes(3); + }); + + rerender({ + chatID: "chat-2", + organizationName: ORGANIZATION_NAME, + username: USERNAME, + }); + + await act(async () => { + setStreamState(createStreamState("create_workspace")); + }); + + await waitFor(() => { + expect(invalidateSpy).toHaveBeenCalledTimes(6); + expect(invalidateSpy).toHaveBeenCalledWith({ + queryKey: ["chats", "chat-2"], + }); + }); + }); + + it("waits for completed tool results before invalidating", async () => { + const { invalidateSpy, setStreamState } = renderInvalidations(); + + await act(async () => { + setStreamState( + createStreamState("create_workspace", "tool-1", { + isStreaming: true, + }), + ); + }); + + expect(invalidateSpy).not.toHaveBeenCalled(); + + await act(async () => { + setStreamState(createStreamState("create_workspace")); + }); + + await waitFor(() => { + expect(invalidateSpy).toHaveBeenCalledTimes(3); + }); + }); + + it("does nothing when streamState is null", () => { + const { store } = createTestStore(null); + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + + renderHook( + () => + useChatToolInvalidations({ + store, + chatID: "chat-1", + organizationName: ORGANIZATION_NAME, + username: USERNAME, + }), + { wrapper: createWrapper(queryClient) }, + ); + + expect(invalidateSpy).not.toHaveBeenCalled(); + }); +}); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.ts new file mode 100644 index 0000000000..41f90139b1 --- /dev/null +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.ts @@ -0,0 +1,102 @@ +import { useEffect, useRef } from "react"; +import { useQueryClient } from "react-query"; +import { chatKey } from "#/api/queries/chats"; +import { invalidateWorkspaceMutationQueries } from "#/api/queries/workspaces"; +import { type ChatStore, useChatSelector } from "./chatStore"; +import type { StreamState } from "./types"; + +type ChatToolResult = Pick< + StreamState["toolResults"][string], + "id" | "name" | "isStreaming" +>; + +// Only extract the toolResults record from the stream state. +// This reference is stable during pure text/thinking streaming +// and only changes when a tool result actually appears, avoiding +// a re-render of AgentChatPage on every token. +const selectStreamToolResults = (state: { + streamState: StreamState | null; +}): Record | null => + state.streamState?.toolResults ?? null; + +interface UseChatToolInvalidationsOptions { + store: ChatStore; + chatID: string | undefined; + organizationName: string; + username: string; +} + +const CHAT_WORKSPACE_BINDING_TOOL_NAMES = new Set(["create_workspace"]); +const WORKSPACE_MUTATION_TOOL_NAMES = new Set([ + "create_workspace", + "start_workspace", + "stop_workspace", +]); + +/** + * Watches completed chat tool results and invalidates derived UI data for the + * server state those tools may have changed. + */ +export function useChatToolInvalidations({ + store, + chatID, + organizationName, + username, +}: UseChatToolInvalidationsOptions): void { + const queryClient = useQueryClient(); + const toolResults = useChatSelector(store, selectStreamToolResults); + const processedToolCallIdsRef = useRef>(new Set()); + const chatIDRef = useRef(chatID); + + useEffect(() => { + if (chatIDRef.current !== chatID) { + chatIDRef.current = chatID; + processedToolCallIdsRef.current.clear(); + } + + if (!toolResults || !chatID) { + processedToolCallIdsRef.current.clear(); + return; + } + + let shouldInvalidateChat = false; + let shouldInvalidateWorkspace = false; + + for (const toolResult of Object.values(toolResults)) { + if ( + toolResult.isStreaming || + processedToolCallIdsRef.current.has(toolResult.id) + ) { + continue; + } + + const changesChatWorkspaceBinding = CHAT_WORKSPACE_BINDING_TOOL_NAMES.has( + toolResult.name, + ); + const changesWorkspace = WORKSPACE_MUTATION_TOOL_NAMES.has( + toolResult.name, + ); + if (!changesChatWorkspaceBinding && !changesWorkspace) { + continue; + } + + processedToolCallIdsRef.current.add(toolResult.id); + shouldInvalidateChat = + shouldInvalidateChat || changesChatWorkspaceBinding; + shouldInvalidateWorkspace = shouldInvalidateWorkspace || changesWorkspace; + } + + if (shouldInvalidateChat) { + void queryClient.invalidateQueries({ + queryKey: chatKey(chatID), + }); + } + + if (shouldInvalidateWorkspace) { + void invalidateWorkspaceMutationQueries(queryClient, { + organizationName, + username, + }); + } + }, [chatID, organizationName, queryClient, toolResults, username]); +} diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useWorkspaceCreationWatcher.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/useWorkspaceCreationWatcher.test.tsx deleted file mode 100644 index f8458597a2..0000000000 --- a/site/src/pages/AgentsPage/components/ChatConversation/useWorkspaceCreationWatcher.test.tsx +++ /dev/null @@ -1,301 +0,0 @@ -import { renderHook, waitFor } from "@testing-library/react"; -import type { FC, PropsWithChildren } from "react"; -import { act } from "react"; -import { QueryClient, QueryClientProvider } from "react-query"; -import { beforeEach, describe, expect, it, vi } from "vitest"; -import type { StreamState } from "./types"; -import { useWorkspaceCreationWatcher } from "./useWorkspaceCreationWatcher"; - -type ChatStoreHandle = Parameters< - typeof useWorkspaceCreationWatcher ->[0]["store"]; - -const createStreamState = ( - toolCalls: StreamState["toolCalls"], - toolResults: StreamState["toolResults"] = {}, -): StreamState => ({ - blocks: [], - toolCalls, - toolResults, - sources: [], -}); - -type MinimalChatStoreState = Pick< - ReturnType, - "streamState" ->; - -const createTestStore = (initial: StreamState | null = null) => { - let state: MinimalChatStoreState = { streamState: initial }; - const listeners = new Set<() => void>(); - - const store: Pick = { - getSnapshot: () => state as ReturnType, - subscribe: (listener: () => void) => { - listeners.add(listener); - return () => { - listeners.delete(listener); - }; - }, - }; - - return { - store: store as ChatStoreHandle, - setStreamState: (streamState: StreamState | null) => { - state = { streamState }; - for (const listener of listeners) { - listener(); - } - }, - }; -}; - -const createWrapper = (queryClient: QueryClient): FC => { - return ({ children }) => ( - {children} - ); -}; - -describe("useWorkspaceCreationWatcher", () => { - let queryClient: QueryClient; - - beforeEach(() => { - queryClient = new QueryClient({ - defaultOptions: { - queries: { - retry: false, - }, - }, - }); - }); - - it("invalidates chatKey on create_workspace tool result", async () => { - const { store, setStreamState } = createTestStore(); - - renderHook( - () => - useWorkspaceCreationWatcher({ - store, - chatID: "chat-1", - }), - { wrapper: createWrapper(queryClient) }, - ); - - const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); - - await act(async () => { - setStreamState( - createStreamState( - { - "tool-1": { - id: "tool-1", - name: "create_workspace", - args: { template: "some-template" }, - }, - }, - { - "tool-1": { - id: "tool-1", - name: "create_workspace", - isError: false, - }, - }, - ), - ); - }); - - await waitFor(() => { - expect(invalidateSpy).toHaveBeenCalledWith({ - queryKey: ["chats", "chat-1"], - }); - }); - - invalidateSpy.mockRestore(); - }); - - it("does not invalidate chat for non-workspace tools", async () => { - const { store, setStreamState } = createTestStore(); - - renderHook( - () => - useWorkspaceCreationWatcher({ - store, - chatID: "chat-1", - }), - { wrapper: createWrapper(queryClient) }, - ); - - const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); - - await act(async () => { - setStreamState( - createStreamState( - { - "tool-1": { - id: "tool-1", - name: "read_file", - args: { path: "/workspace/src/main.ts" }, - }, - }, - { - "tool-1": { - id: "tool-1", - name: "read_file", - isError: false, - }, - }, - ), - ); - }); - - await waitFor(() => { - expect(invalidateSpy).not.toHaveBeenCalled(); - }); - - invalidateSpy.mockRestore(); - }); - - it("does not process the same tool call ID twice", async () => { - const { store, setStreamState } = createTestStore(); - - renderHook( - () => - useWorkspaceCreationWatcher({ - store, - chatID: "chat-1", - }), - { wrapper: createWrapper(queryClient) }, - ); - - const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); - - const toolCalls = { - "tool-1": { - id: "tool-1", - name: "create_workspace", - args: { template: "some-template" }, - }, - }; - const toolResults = { - "tool-1": { - id: "tool-1", - name: "create_workspace", - isError: false, - }, - }; - - await act(async () => { - setStreamState(createStreamState(toolCalls, toolResults)); - }); - - await waitFor(() => { - expect(invalidateSpy).toHaveBeenCalledTimes(1); - }); - - // Re-emit the same stream state (simulates a re-render). - await act(async () => { - setStreamState(createStreamState(toolCalls, toolResults)); - }); - - // invalidateQueries should still have been called only once. - await waitFor(() => { - expect(invalidateSpy).toHaveBeenCalledTimes(1); - }); - - invalidateSpy.mockRestore(); - }); - - it("resets processed tool call IDs when chatID changes", async () => { - const { store, setStreamState } = createTestStore(); - - const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); - - const { rerender } = renderHook( - ({ chatID }: { chatID: string }) => - useWorkspaceCreationWatcher({ - store, - chatID, - }), - { - initialProps: { chatID: "chat-1" }, - wrapper: createWrapper(queryClient), - }, - ); - - await act(async () => { - setStreamState( - createStreamState( - { - "tool-1": { - id: "tool-1", - name: "create_workspace", - args: { template: "some-template" }, - }, - }, - { - "tool-1": { - id: "tool-1", - name: "create_workspace", - isError: false, - }, - }, - ), - ); - }); - - await waitFor(() => { - expect(invalidateSpy).toHaveBeenCalledTimes(1); - }); - - // Switch to a new chat and emit the same tool call ID. - rerender({ chatID: "chat-2" }); - - await act(async () => { - setStreamState( - createStreamState( - { - "tool-1": { - id: "tool-1", - name: "create_workspace", - args: { template: "some-template" }, - }, - }, - { - "tool-1": { - id: "tool-1", - name: "create_workspace", - isError: false, - }, - }, - ), - ); - }); - - await waitFor(() => { - expect(invalidateSpy).toHaveBeenCalledTimes(2); - }); - - invalidateSpy.mockRestore(); - }); - - it("does nothing when streamState is null", async () => { - const { store } = createTestStore(null); - - renderHook( - () => - useWorkspaceCreationWatcher({ - store, - chatID: "chat-1", - }), - { wrapper: createWrapper(queryClient) }, - ); - - const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); - - await waitFor(() => { - expect(invalidateSpy).not.toHaveBeenCalled(); - }); - - invalidateSpy.mockRestore(); - }); -}); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useWorkspaceCreationWatcher.ts b/site/src/pages/AgentsPage/components/ChatConversation/useWorkspaceCreationWatcher.ts deleted file mode 100644 index 9f6077c086..0000000000 --- a/site/src/pages/AgentsPage/components/ChatConversation/useWorkspaceCreationWatcher.ts +++ /dev/null @@ -1,76 +0,0 @@ -import { useEffect, useRef } from "react"; -import { useQueryClient } from "react-query"; -import { chatKey } from "#/api/queries/chats"; -import { useChatSelector } from "./chatStore"; -import type { StreamState } from "./types"; - -type ChatStoreHandle = Parameters[0]; - -// Only extract the toolResults record from the stream state. -// This reference is stable during pure text/thinking streaming -// and only changes when a tool result actually appears, avoiding -// a re-render of AgentChatPage on every token. -const selectStreamToolResults = (state: { - streamState: StreamState | null; -}): Record | null => - state.streamState?.toolResults ?? null; - -interface UseWorkspaceCreationWatcherOptions { - store: ChatStoreHandle; - chatID: string | undefined; -} - -// Triggers chat query invalidation to resolve the workspace/agent. -const WORKSPACE_TOOL_NAMES = new Set(["create_workspace"]); - -/** - * Watches stream tool results for create_workspace completions and - * invalidates the chat query so the sidebar can display workspace info. - * The agent now handles all path discovery and scan triggering via - * the PathStore — no frontend refresh needed. - */ -export function useWorkspaceCreationWatcher({ - store, - chatID, -}: UseWorkspaceCreationWatcherOptions): void { - const queryClient = useQueryClient(); - const toolResults = useChatSelector(store, selectStreamToolResults); - const processedToolCallIdsRef = useRef>(new Set()); - const chatIDRef = useRef(chatID); - - // Watch stream tool results for create_workspace completions. - useEffect(() => { - // Reset processed IDs when chatID changes. - if (chatIDRef.current !== chatID) { - chatIDRef.current = chatID; - processedToolCallIdsRef.current = new Set(); - } - - if (!toolResults || !chatID) { - processedToolCallIdsRef.current.clear(); - return; - } - - let shouldInvalidateChat = false; - - for (const toolResult of Object.values(toolResults)) { - if (processedToolCallIdsRef.current.has(toolResult.id)) { - continue; - } - - if (WORKSPACE_TOOL_NAMES.has(toolResult.name)) { - processedToolCallIdsRef.current.add(toolResult.id); - shouldInvalidateChat = true; - } - } - - if (shouldInvalidateChat) { - // Invalidate chatKey to trigger the workspace resolution - // cascade: chat refetch → workspaceId → workspace query → - // agent resolved → git watcher connects. - void queryClient.invalidateQueries({ - queryKey: chatKey(chatID), - }); - } - }, [toolResults, queryClient, chatID]); -} diff --git a/site/src/pages/AgentsPage/components/UsageIndicator.tsx b/site/src/pages/AgentsPage/components/UsageIndicator.tsx index fb5e0d8b58..2a07878f12 100644 --- a/site/src/pages/AgentsPage/components/UsageIndicator.tsx +++ b/site/src/pages/AgentsPage/components/UsageIndicator.tsx @@ -20,7 +20,10 @@ import { TooltipTrigger, } from "#/components/Tooltip/Tooltip"; import { useAuthenticated } from "#/hooks/useAuthenticated"; -import { useDashboard } from "#/modules/dashboard/useDashboard"; +import { + getDefaultOrganizationName, + useDashboard, +} from "#/modules/dashboard/useDashboard"; import { cn } from "#/utils/cn"; import { formatCostMicros } from "#/utils/currency"; import { getUsageLimitPeriodLabel } from "./ChatCostSummaryView"; @@ -46,11 +49,11 @@ export const UsageIndicator: FC = () => { ); const { user } = useAuthenticated(); const { organizations } = useDashboard(); - const organizationName = - organizations.find((org) => org.is_default)?.name ?? ""; + const organizationName = getDefaultOrganizationName(organizations); const username = user.username; const { data: quota, isError: isQuotaError } = useQuery({ ...workspaceQuota(organizationName, username), + refetchInterval: 60_000, enabled: organizationName !== "" && username !== "", }); const hasWorkspaceQuotaUsage =