From 372979556e63222b25ab00a886f4c511013918c6 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Thu, 4 Jun 2026 09:29:18 +0000 Subject: [PATCH 1/5] feat: add shared workspaces view Add a Shared Workspaces tree view alongside My and All Workspaces, with search and manual refresh. The view queries `shared:true`, which on the server matches any workspace with non-empty user or group ACLs. That set also includes workspaces the signed-in user owns and shared out, so the extension filters those out client-side by `owner_id` to leave only workspaces shared with the user. `shared_with_user:me` is not used: the server does not implement `me` there, and direct-user filtering would miss group shares. --- package.json | 27 +++ src/core/commandManager.ts | 1 + src/deployment/deploymentManager.ts | 85 +++++++- src/extension.ts | 61 +++++- src/workspace/workspacesProvider.ts | 20 +- test/mocks/vscode.runtime.ts | 9 +- .../unit/deployment/deploymentManager.test.ts | 43 +++- .../unit/workspace/workspacesProvider.test.ts | 185 ++++++++++++++++++ 8 files changed, 394 insertions(+), 37 deletions(-) create mode 100644 test/unit/workspace/workspacesProvider.test.ts diff --git a/package.json b/package.json index cb0e8d01fd..7c7c694d7a 100644 --- a/package.json +++ b/package.json @@ -317,6 +317,13 @@ "visibility": "visible", "icon": "media/logo-white.svg" }, + { + "id": "sharedWorkspaces", + "name": "Shared Workspaces", + "visibility": "visible", + "icon": "media/logo-white.svg", + "when": "coder.authenticated" + }, { "id": "allWorkspaces", "name": "All Workspaces", @@ -459,6 +466,12 @@ "category": "Coder", "icon": "$(search)" }, + { + "command": "coder.searchSharedWorkspaces", + "title": "Search", + "category": "Coder", + "icon": "$(search)" + }, { "command": "coder.searchAllWorkspaces", "title": "Search", @@ -585,6 +598,10 @@ "command": "coder.searchMyWorkspaces", "when": "false" }, + { + "command": "coder.searchSharedWorkspaces", + "when": "false" + }, { "command": "coder.searchAllWorkspaces", "when": "false" @@ -628,6 +645,16 @@ "when": "coder.authenticated && view == myWorkspaces", "group": "navigation@3" }, + { + "command": "coder.refreshWorkspaces", + "when": "coder.authenticated && view == sharedWorkspaces", + "group": "navigation@2" + }, + { + "command": "coder.searchSharedWorkspaces", + "when": "coder.authenticated && view == sharedWorkspaces", + "group": "navigation@3" + }, { "command": "coder.searchAllWorkspaces", "when": "coder.authenticated && view == allWorkspaces", diff --git a/src/core/commandManager.ts b/src/core/commandManager.ts index 33c0035193..1310e8a246 100644 --- a/src/core/commandManager.ts +++ b/src/core/commandManager.ts @@ -22,6 +22,7 @@ export const CODER_COMMAND_IDS = [ "coder.viewLogs", "coder.exportTelemetry", "coder.searchMyWorkspaces", + "coder.searchSharedWorkspaces", "coder.searchAllWorkspaces", "coder.manageCredentials", "coder.applyRecommendedSettings", diff --git a/src/deployment/deploymentManager.ts b/src/deployment/deploymentManager.ts index 8ef71d56e7..20da669dc5 100644 --- a/src/deployment/deploymentManager.ts +++ b/src/deployment/deploymentManager.ts @@ -48,6 +48,8 @@ export class DeploymentManager implements vscode.Disposable { private readonly deploymentTelemetry: DeploymentTelemetry; #deployment: Deployment | null = null; + #authedUser: User | null = null; + #authStateVersion = 0; #disposed = false; #authListenerDisposable: vscode.Disposable | undefined; #authConfigDisposable: vscode.Disposable | undefined; @@ -93,6 +95,14 @@ export class DeploymentManager implements vscode.Disposable { return this.#deployment; } + public getCurrentUserId(): string | undefined { + return this.#authedUser?.id; + } + + public getAuthStateVersion(): number { + return this.#authStateVersion; + } + /** * Check if we have an authenticated deployment (does not guarantee that the current auth data is valid). */ @@ -110,6 +120,7 @@ export class DeploymentManager implements vscode.Disposable { deployment: Deployment & { token?: string }, ): Promise { const deploymentBefore = this.#deployment; + const authStateVersionBefore = this.#authStateVersion; const token = deployment.token ?? (await this.secretsManager.getSessionAuth(deployment.safeHostname)) @@ -118,7 +129,9 @@ export class DeploymentManager implements vscode.Disposable { try { const user = await tempClient.getAuthenticatedUser(); - if (this.#hasStateChangedSince(deploymentBefore)) { + if ( + this.#hasStateChangedSince(deploymentBefore, authStateVersionBefore) + ) { return false; } await this.setDeployment({ ...deployment, token, user }); @@ -132,11 +145,15 @@ export class DeploymentManager implements vscode.Disposable { } /** True if disposal, login, or a deployment switch raced our await. */ - #hasStateChangedSince(deploymentBefore: Deployment | null): boolean { + #hasStateChangedSince( + deploymentBefore: Deployment | null, + authStateVersionBefore: number, + ): boolean { return ( this.#disposed || this.isAuthenticated() || - this.#deployment !== deploymentBefore + this.#deployment !== deploymentBefore || + this.#authStateVersion !== authStateVersionBefore ); } @@ -151,8 +168,12 @@ export class DeploymentManager implements vscode.Disposable { hostname: deployment.safeHostname, user: deployment.user.username, }); - this.#deployment = { ...deployment }; - const ourRef = this.#deployment; + const deploymentWithoutAuth = DeploymentSchema.parse(deployment); + const ourRef = this.commitDeploymentState( + deploymentWithoutAuth, + deployment.user, + ); + const authStateVersion = this.#authStateVersion; this.telemetryService.setDeploymentUrl(deployment.url); // Updates client credentials @@ -169,11 +190,12 @@ export class DeploymentManager implements vscode.Disposable { this.updateAuthContexts(deployment.user); this.refreshWorkspaces(); - const deploymentWithoutAuth: Deployment = - DeploymentSchema.parse(deployment); await this.oauthSessionManager.setDeployment(deploymentWithoutAuth); // Bail if a concurrent write took over during the await. - if (this.#deployment !== ourRef) { + if ( + this.#deployment !== ourRef || + this.#authStateVersion !== authStateVersion + ) { return; } await this.persistDeployment(deploymentWithoutAuth); @@ -187,7 +209,7 @@ export class DeploymentManager implements vscode.Disposable { this.suspendSession(reason); this.#authListenerDisposable?.dispose(); this.#authListenerDisposable = undefined; - this.#deployment = null; + this.commitDeploymentState(null, null); this.telemetryService.setDeploymentUrl(""); await this.secretsManager.setCurrentDeployment(undefined); @@ -199,6 +221,7 @@ export class DeploymentManager implements vscode.Disposable { */ public suspendSession(reason: DeploymentSuspendReason): void { const wasAuthenticated = this.isAuthenticated(); + this.commitDeploymentState(this.#deployment, null); this.oauthSessionManager.clearDeployment(); this.client.setCredentials(undefined, undefined); this.updateAuthContexts(undefined); @@ -248,7 +271,11 @@ export class DeploymentManager implements vscode.Disposable { if (auth) { if (this.isAuthenticated()) { - this.client.setCredentials(auth.url, auth.token); + await this.verifyAndUpdateAuthenticatedSession({ + url: auth.url, + safeHostname, + token: auth.token, + }); } else { this.logger.debug( "Token updated after session suspended, recovering", @@ -269,6 +296,44 @@ export class DeploymentManager implements vscode.Disposable { ); } + private async verifyAndUpdateAuthenticatedSession( + deployment: Deployment & { token: string }, + ): Promise { + const deploymentBefore = this.#deployment; + const authStateVersionBefore = this.#authStateVersion; + const tempClient = CoderApi.create( + deployment.url, + deployment.token, + this.logger, + ); + + try { + const user = await tempClient.getAuthenticatedUser(); + if ( + this.#disposed || + this.#deployment !== deploymentBefore || + this.#authStateVersion !== authStateVersionBefore + ) { + return; + } + await this.setDeployment({ ...deployment, user }); + } catch (e) { + this.logger.warn("Failed to authenticate updated session:", e); + } finally { + tempClient.dispose(); + } + } + + private commitDeploymentState( + deployment: Deployment | null, + authedUser: User | null, + ): Deployment | null { + this.#deployment = deployment; + this.#authedUser = authedUser; + this.#authStateVersion += 1; + return this.#deployment; + } + private subscribeToAuthConfigChanges(): void { this.#authConfigDisposable = watchConfigurationChanges( getAuthConfigWatchSettings(), diff --git a/src/extension.ts b/src/extension.ts index 5903d79cfd..a1832131bd 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -30,7 +30,10 @@ import { WorkspaceQuery, } from "./workspace/workspacesProvider"; +import type { Workspace } from "coder/site/src/api/typesGenerated"; + const MY_WORKSPACES_TREE_ID = "myWorkspaces"; +const SHARED_WORKSPACES_TREE_ID = "sharedWorkspaces"; const ALL_WORKSPACES_TREE_ID = "allWorkspaces"; export async function activate(ctx: vscode.ExtensionContext): Promise { @@ -100,7 +103,6 @@ async function doActivate( ? await secretsManager.getSessionAuth(deployment.safeHostname) : undefined; tracer.setAuthState(deploymentSessionAuth ? "stored" : "none"); - // Shared handler for auth failures (used by interceptor + session manager) const handleAuthFailure = (): Promise => { deploymentManager.suspendSession("auth_failure"); @@ -155,6 +157,13 @@ async function doActivate( ctx.subscriptions.push(authInterceptor); const isAuthenticated = () => contextManager.get("coder.authenticated"); + const getAuthStateVersion = () => deploymentManager.getAuthStateVersion(); + const filterSharedWorkspaces = (workspaces: readonly Workspace[]) => { + const currentUserId = deploymentManager.getCurrentUserId(); + return currentUserId + ? workspaces.filter((workspace) => workspace.owner_id !== currentUserId) + : []; + }; const myWorkspacesProvider = new WorkspaceProvider( WorkspaceQuery.Mine, @@ -162,6 +171,8 @@ async function doActivate( output, isAuthenticated, 5, + undefined, + getAuthStateVersion, ); ctx.subscriptions.push(myWorkspacesProvider); @@ -170,9 +181,31 @@ async function doActivate( client, output, isAuthenticated, + undefined, + undefined, + getAuthStateVersion, ); ctx.subscriptions.push(allWorkspacesProvider); + const sharedWorkspacesProvider = new WorkspaceProvider( + WorkspaceQuery.Shared, + client, + output, + isAuthenticated, + undefined, + filterSharedWorkspaces, + getAuthStateVersion, + ); + ctx.subscriptions.push(sharedWorkspacesProvider); + + const deploymentManager = DeploymentManager.create( + serviceContainer, + client, + oauthSessionManager, + [myWorkspacesProvider, sharedWorkspacesProvider, allWorkspacesProvider], + ); + ctx.subscriptions.push(deploymentManager); + // createTreeView, unlike registerTreeDataProvider, gives us the tree view API // (so we can see when it is visible) but otherwise they have the same effect. const myWsTree = vscode.window.createTreeView(MY_WORKSPACES_TREE_ID, { @@ -188,6 +221,19 @@ async function doActivate( ctx.subscriptions, ); + const sharedWsTree = vscode.window.createTreeView(SHARED_WORKSPACES_TREE_ID, { + treeDataProvider: sharedWorkspacesProvider, + }); + ctx.subscriptions.push(sharedWsTree); + sharedWorkspacesProvider.setVisibility(sharedWsTree.visible); + sharedWsTree.onDidChangeVisibility( + (event) => { + sharedWorkspacesProvider.setVisibility(event.visible); + }, + undefined, + ctx.subscriptions, + ); + const allWsTree = vscode.window.createTreeView(ALL_WORKSPACES_TREE_ID, { treeDataProvider: allWorkspacesProvider, }); @@ -201,15 +247,6 @@ async function doActivate( ctx.subscriptions, ); - // Create deployment manager to centralize deployment state management - const deploymentManager = DeploymentManager.create( - serviceContainer, - client, - oauthSessionManager, - [myWorkspacesProvider, allWorkspacesProvider], - ); - ctx.subscriptions.push(deploymentManager); - // Register globally available commands. Many of these have visibility // controlled by contexts, see `when` in the package.json. const commands = new Commands(serviceContainer, client, deploymentManager); @@ -318,6 +355,7 @@ async function doActivate( ); commandManager.register("coder.refreshWorkspaces", () => { void myWorkspacesProvider.fetchAndRefresh(); + void sharedWorkspacesProvider.fetchAndRefresh(); void allWorkspacesProvider.fetchAndRefresh(); }); commandManager.register("coder.viewLogs", commands.viewLogs.bind(commands)); @@ -328,6 +366,9 @@ async function doActivate( commandManager.register("coder.searchMyWorkspaces", async () => showTreeViewSearch(MY_WORKSPACES_TREE_ID), ); + commandManager.register("coder.searchSharedWorkspaces", async () => + showTreeViewSearch(SHARED_WORKSPACES_TREE_ID), + ); commandManager.register("coder.searchAllWorkspaces", async () => showTreeViewSearch(ALL_WORKSPACES_TREE_ID), ); diff --git a/src/workspace/workspacesProvider.ts b/src/workspace/workspacesProvider.ts index 7ead0dfafd..c039700ba8 100644 --- a/src/workspace/workspacesProvider.ts +++ b/src/workspace/workspacesProvider.ts @@ -23,8 +23,13 @@ import { type Logger } from "../logging/logger"; export enum WorkspaceQuery { Mine = "owner:me", All = "", + Shared = "shared:true", } +type WorkspaceFilter = ( + workspaces: readonly Workspace[], +) => readonly Workspace[]; + /** * Polls workspaces using the provided REST client and renders them in a tree. * @@ -53,6 +58,9 @@ export class WorkspaceProvider private readonly logger: Logger, private readonly isAuthenticated: () => boolean, private readonly timerSeconds?: number, + private readonly filterWorkspaces: WorkspaceFilter = (workspaces) => + workspaces, + private readonly getStateVersion: () => number = () => 0, ) { // No initialization. } @@ -108,6 +116,7 @@ export class WorkspaceProvider throw new Error("not logged in"); } + const stateVersion = this.getStateVersion(); const resp = await this.client.getWorkspaces({ q: this.getWorkspacesQuery, }); @@ -117,7 +126,7 @@ export class WorkspaceProvider const url2 = this.client.getAxiosInstance().defaults.baseURL; if (!url2) { throw new Error("not logged in"); - } else if (url !== url2) { + } else if (url !== url2 || stateVersion !== this.getStateVersion()) { // In this case we need to fetch from the new deployment instead. // TODO: It would be better to cancel this fetch when that happens, // because this means we have to wait for the old fetch to finish before @@ -125,6 +134,7 @@ export class WorkspaceProvider return this.fetch(); } + const workspaces = this.filterWorkspaces(resp.workspaces); const oldWatcherIds = [...this.agentWatchers.keys()]; const reusedWatcherIds: string[] = []; @@ -133,7 +143,7 @@ export class WorkspaceProvider // have this separate map held outside the tree. const showMetadata = this.getWorkspacesQuery === WorkspaceQuery.Mine; if (showMetadata) { - const agents = extractAllAgents(resp.workspaces); + const agents = extractAllAgents(workspaces); for (const agent of agents) { // If we have an existing watcher, re-use it. const oldWatcher = this.agentWatchers.get(agent.id); @@ -159,11 +169,13 @@ export class WorkspaceProvider } } + const showOwner = this.getWorkspacesQuery !== WorkspaceQuery.Mine; + // Create tree items for each workspace - const workspaceTreeItems = resp.workspaces.map((workspace: Workspace) => { + const workspaceTreeItems = workspaces.map((workspace: Workspace) => { const workspaceTreeItem = new WorkspaceTreeItem( workspace, - this.getWorkspacesQuery === WorkspaceQuery.All, + showOwner, showMetadata, ); diff --git a/test/mocks/vscode.runtime.ts b/test/mocks/vscode.runtime.ts index 8eee245214..03dc548d25 100644 --- a/test/mocks/vscode.runtime.ts +++ b/test/mocks/vscode.runtime.ts @@ -67,16 +67,19 @@ export class ThemeColor { export class TreeItem { id?: string; - contextValue?: string; + label?: string; description?: string; tooltip?: string; + contextValue?: string; command?: unknown; iconPath?: unknown; constructor( - public label: string, + label: string, public collapsibleState?: number, - ) {} + ) { + this.label = label; + } } export class Uri { diff --git a/test/unit/deployment/deploymentManager.test.ts b/test/unit/deployment/deploymentManager.test.ts index e47c75cacb..1078700243 100644 --- a/test/unit/deployment/deploymentManager.test.ts +++ b/test/unit/deployment/deploymentManager.test.ts @@ -115,12 +115,13 @@ describe("DeploymentManager", () => { const { manager } = createTestContext(); expect(manager.getCurrentDeployment()).toBeNull(); + expect(manager.getCurrentUserId()).toBeUndefined(); expect(manager.isAuthenticated()).toBe(false); }); it("returns deployment and isAuthenticated=true after setDeployment", async () => { const { manager } = createTestContext(); - const user = createMockUser(); + const user = createMockUser({ id: "current-user" }); await manager.setDeployment({ url: TEST_URL, @@ -133,6 +134,7 @@ describe("DeploymentManager", () => { url: TEST_URL, safeHostname: TEST_HOSTNAME, }); + expect(manager.getCurrentUserId()).toBe("current-user"); expect(manager.isAuthenticated()).toBe(true); }); @@ -150,6 +152,7 @@ describe("DeploymentManager", () => { await manager.clearDeployment("credentials_removed"); expect(manager.getCurrentDeployment()).toBeNull(); + expect(manager.getCurrentUserId()).toBeUndefined(); expect(manager.isAuthenticated()).toBe(false); }); }); @@ -221,6 +224,7 @@ describe("DeploymentManager", () => { expect(result).toBe(true); expect(mockClient.host).toBe(TEST_URL); expect(mockClient.token).toBe("test-token"); + expect(manager.getCurrentUserId()).toBe(user.id); expect(manager.isAuthenticated()).toBe(true); }); @@ -238,6 +242,7 @@ describe("DeploymentManager", () => { expect(result).toBe(false); expect(manager.getCurrentDeployment()).toBeNull(); + expect(manager.getCurrentUserId()).toBeUndefined(); expect(manager.isAuthenticated()).toBe(false); }); @@ -377,11 +382,16 @@ describe("DeploymentManager", () => { }); describe("auth listener", () => { - it("updates credentials on token change", async () => { - const { mockClient, secretsManager, manager } = createTestContext(); - const user = createMockUser(); + it("updates credentials and user on token change", async () => { + const { mockClient, validationMockClient, secretsManager, manager } = + createTestContext(); + const user = createMockUser({ id: "initial-user" }); + const refreshedUser = createMockUser({ + id: "refreshed-user", + username: "refresheduser", + }); + validationMockClient.setAuthenticatedUserResponse(refreshedUser); - // Set up authenticated deployment await manager.setDeployment({ url: TEST_URL, safeHostname: TEST_HOSTNAME, @@ -390,18 +400,17 @@ describe("DeploymentManager", () => { }); expect(mockClient.token).toBe("initial-token"); + expect(manager.getCurrentUserId()).toBe("initial-user"); expect(manager.isAuthenticated()).toBe(true); - // Simulate token refresh via secrets change await secretsManager.setSessionAuth(TEST_HOSTNAME, { url: TEST_URL, token: "refreshed-token", }); - - // Wait for async handler await new Promise((resolve) => setImmediate(resolve)); expect(mockClient.token).toBe("refreshed-token"); + expect(manager.getCurrentUserId()).toBe("refreshed-user"); expect(manager.isAuthenticated()).toBe(true); }); }); @@ -424,6 +433,7 @@ describe("DeploymentManager", () => { expect(mockClient.token).toBeUndefined(); expect(contextManager.get("coder.authenticated")).toBe(false); expect(contextManager.get("coder.isOwner")).toBe(false); + expect(manager.getCurrentUserId()).toBeUndefined(); }); it("resets the telemetry deployment URL on clearDeployment", async () => { @@ -484,6 +494,7 @@ describe("DeploymentManager", () => { expect(mockClient.host).toBeUndefined(); expect(mockClient.token).toBeUndefined(); expect(contextManager.get("coder.authenticated")).toBe(false); + expect(manager.getCurrentUserId()).toBeUndefined(); expect(manager.isAuthenticated()).toBe(false); expect(mockWorkspaceProvider.clear).toHaveBeenCalled(); @@ -499,12 +510,21 @@ describe("DeploymentManager", () => { it("recovers from suspended state when auth settings change", async () => { vi.useFakeTimers(); try { - const { mockClient, validationMockClient, telemetrySink, manager } = - createTestContext(); + const { + mockClient, + validationMockClient, + secretsManager, + telemetrySink, + manager, + } = createTestContext(); const config = new MockConfigurationProvider(); const user = createMockUser(); validationMockClient.setAuthenticatedUserResponse(user); + await secretsManager.setSessionAuth(TEST_HOSTNAME, { + url: TEST_URL, + token: "", + }); await manager.setDeployment({ url: TEST_URL, safeHostname: TEST_HOSTNAME, @@ -521,6 +541,7 @@ describe("DeploymentManager", () => { expect(mockClient.host).toBe(TEST_URL); expect(mockClient.token).toBe(""); expect(manager.isAuthenticated()).toBe(true); + expect(manager.getCurrentUserId()).toBe(user.id); expect(validationMockClient.getAuthenticatedUser).toHaveBeenCalledTimes( 1, ); @@ -562,6 +583,7 @@ describe("DeploymentManager", () => { await vi.runAllTimersAsync(); expect(manager.getCurrentDeployment()).toBeNull(); + expect(manager.getCurrentUserId()).toBeUndefined(); expect(manager.isAuthenticated()).toBe(false); } finally { vi.useRealTimers(); @@ -601,6 +623,7 @@ describe("DeploymentManager", () => { // Should recover and be authenticated again expect(mockClient.token).toBe("recovered-token"); + expect(manager.getCurrentUserId()).toBe(user.id); expect(manager.isAuthenticated()).toBe(true); expect(telemetrySink.expectOne("deployment.recovered")).toMatchObject({ properties: { trigger: "token_update" }, diff --git a/test/unit/workspace/workspacesProvider.test.ts b/test/unit/workspace/workspacesProvider.test.ts new file mode 100644 index 0000000000..eb1d4236ad --- /dev/null +++ b/test/unit/workspace/workspacesProvider.test.ts @@ -0,0 +1,185 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { + WorkspaceProvider, + WorkspaceQuery, +} from "@/workspace/workspacesProvider"; + +import { workspace as createWorkspace } from "@repo/mocks"; + +import { createMockLogger } from "../../mocks/testHelpers"; + +import type { Workspace } from "coder/site/src/api/typesGenerated"; + +import type { CoderApi } from "@/api/coderApi"; + +const baseUrl = "https://coder.example.com"; + +class MockClient { + baseURL: string | undefined = baseUrl; + readonly getWorkspaces = vi.fn( + (_req: { + q: string; + }): Promise<{ + workspaces: readonly Workspace[]; + count: number; + }> => Promise.resolve({ workspaces: [], count: 0 }), + ); + + getAxiosInstance() { + return { + defaults: { + baseURL: this.baseURL, + }, + }; + } +} + +describe("WorkspaceProvider", () => { + let client: MockClient; + let isAuthenticated: () => boolean; + + beforeEach(() => { + client = new MockClient(); + isAuthenticated = vi.fn<() => boolean>(() => true); + }); + + function createProvider( + query: WorkspaceQuery, + options: { + filterWorkspaces?: ( + workspaces: readonly Workspace[], + ) => readonly Workspace[]; + getStateVersion?: () => number; + timerSeconds?: number; + } = {}, + ) { + return new WorkspaceProvider( + query, + client as unknown as CoderApi, + createMockLogger(), + isAuthenticated, + options.timerSeconds, + options.filterWorkspaces, + options.getStateVersion, + ); + } + + async function fetchVisible(provider: WorkspaceProvider) { + provider.setVisibility(true); + await new Promise((resolve) => setImmediate(resolve)); + } + + it("queries shared workspaces", async () => { + const provider = createProvider(WorkspaceQuery.Shared); + + await fetchVisible(provider); + + expect(client.getWorkspaces).toHaveBeenCalledWith({ + q: WorkspaceQuery.Shared, + }); + }); + + it("applies the workspace filter before rendering tree items", async () => { + client.getWorkspaces.mockResolvedValueOnce({ + workspaces: [ + createWorkspace({ + id: "owned-workspace", + owner_id: "current-user", + owner_name: "current", + name: "owned", + }), + createWorkspace({ + id: "shared-workspace", + owner_id: "other-user", + owner_name: "alice", + name: "shared", + }), + ], + count: 2, + }); + const provider = createProvider(WorkspaceQuery.Shared, { + filterWorkspaces: (workspaces) => + workspaces.filter((workspace) => workspace.owner_id !== "current-user"), + }); + + await fetchVisible(provider); + const children = await provider.getChildren(); + + expect(children).toHaveLength(1); + expect(children[0]?.label).toBe("alice / shared"); + }); + + it("fails closed when the shared workspace filter cannot identify the current user", async () => { + client.getWorkspaces.mockResolvedValueOnce({ + workspaces: [ + createWorkspace({ + id: "shared-workspace", + owner_id: "other-user", + owner_name: "alice", + name: "shared", + }), + ], + count: 1, + }); + const provider = createProvider(WorkspaceQuery.Shared, { + filterWorkspaces: () => [], + }); + + await fetchVisible(provider); + const children = await provider.getChildren(); + + expect(children).toEqual([]); + }); + + it("refetches when auth state changes while a request is pending", async () => { + let stateVersion = 1; + let resolveFirst!: (response: { + workspaces: readonly Workspace[]; + count: number; + }) => void; + client.getWorkspaces + .mockReturnValueOnce( + new Promise((resolve) => { + resolveFirst = resolve; + }), + ) + .mockResolvedValueOnce({ + workspaces: [ + createWorkspace({ + id: "fresh-workspace", + owner_id: "other-user", + owner_name: "alice", + name: "fresh", + }), + ], + count: 1, + }); + const provider = createProvider(WorkspaceQuery.Shared, { + getStateVersion: () => stateVersion, + }); + + provider.setVisibility(true); + await new Promise((resolve) => setImmediate(resolve)); + stateVersion = 2; + resolveFirst({ + workspaces: [ + createWorkspace({ + id: "stale-workspace", + owner_id: "other-user", + owner_name: "bob", + name: "stale", + }), + ], + count: 1, + }); + await new Promise((resolve) => setImmediate(resolve)); + await new Promise((resolve) => setImmediate(resolve)); + + const children = await provider.getChildren(); + + expect(client.getWorkspaces).toHaveBeenCalledTimes(2); + expect(children).toHaveLength(1); + expect(children[0]?.label).toBe("alice / fresh"); + }); +}); From 8ac3fc2fa6a01896de655ba02037ec9db77fff3d Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Thu, 4 Jun 2026 14:34:05 +0000 Subject: [PATCH 2/5] refactor: simplify workspace session state Replace the WorkspaceProvider filter callback and ad-hoc state-version getter with a small WorkspaceSessionState abstraction. Providers now read a point-in-time snapshot (kind, revision, userId) and subscribe to changes, instead of being handed a filter function and a version getter. This puts current-user knowledge and stale-fetch detection behind one interface: DeploymentManager exposes session state in one place, and the provider compares snapshot revisions to discard responses from a session that changed mid-request. --- src/deployment/deploymentManager.ts | 181 ++--- src/extension.ts | 40 +- src/workspace/session.ts | 14 + src/workspace/workspacesProvider.ts | 172 +++-- .../unit/deployment/deploymentManager.test.ts | 128 +++- .../unit/workspace/workspacesProvider.test.ts | 671 +++++++++++++++--- 6 files changed, 895 insertions(+), 311 deletions(-) create mode 100644 src/workspace/session.ts diff --git a/src/deployment/deploymentManager.ts b/src/deployment/deploymentManager.ts index 20da669dc5..58941eeff5 100644 --- a/src/deployment/deploymentManager.ts +++ b/src/deployment/deploymentManager.ts @@ -1,3 +1,5 @@ +import * as vscode from "vscode"; + import { CoderApi } from "../api/coderApi"; import { CONFIG_CHANGE_DEBOUNCE_MS, @@ -16,7 +18,6 @@ import { type Logger } from "../logging/logger"; import { type OAuthSessionManager } from "../oauth/sessionManager"; import { getAuthConfigWatchSettings } from "../settings/authConfig"; import { type TelemetryService } from "../telemetry/service"; -import { type WorkspaceProvider } from "../workspace/workspacesProvider"; import { DeploymentSchema, @@ -25,7 +26,24 @@ import { } from "./types"; import type { User } from "coder/site/src/api/typesGenerated"; -import type * as vscode from "vscode"; + +import type { + WorkspaceSessionSnapshot, + WorkspaceSessionState, +} from "../workspace/session"; + +type DeploymentSessionSnapshot = + | { + readonly kind: "signedOut"; + readonly revision: number; + readonly deployment: Deployment | null; + } + | { + readonly kind: "signedIn"; + readonly revision: number; + readonly deployment: Deployment; + readonly user: User; + }; /** * Manages deployment state for the extension. @@ -36,10 +54,11 @@ import type * as vscode from "vscode"; * - OAuth session management * - Auth listener registration * - Context updates (coder.authenticated, coder.isOwner) - * - Workspace provider refresh * - Cross-window sync handling */ -export class DeploymentManager implements vscode.Disposable { +export class DeploymentManager + implements vscode.Disposable, WorkspaceSessionState +{ private readonly secretsManager: SecretsManager; private readonly mementoManager: MementoManager; private readonly contextManager: ContextManager; @@ -47,9 +66,14 @@ export class DeploymentManager implements vscode.Disposable { private readonly telemetryService: TelemetryService; private readonly deploymentTelemetry: DeploymentTelemetry; - #deployment: Deployment | null = null; - #authedUser: User | null = null; - #authStateVersion = 0; + #session: DeploymentSessionSnapshot = { + kind: "signedOut", + revision: 0, + deployment: null, + }; + readonly #onDidChangeWorkspaceSession = + new vscode.EventEmitter(); + public readonly onDidChange = this.#onDidChangeWorkspaceSession.event; #disposed = false; #authListenerDisposable: vscode.Disposable | undefined; #authConfigDisposable: vscode.Disposable | undefined; @@ -61,7 +85,6 @@ export class DeploymentManager implements vscode.Disposable { serviceContainer: ServiceContainer, private readonly client: CoderApi, private readonly oauthSessionManager: OAuthSessionManager, - private readonly workspaceProviders: WorkspaceProvider[], ) { this.secretsManager = serviceContainer.getSecretsManager(); this.mementoManager = serviceContainer.getMementoManager(); @@ -75,13 +98,11 @@ export class DeploymentManager implements vscode.Disposable { serviceContainer: ServiceContainer, client: CoderApi, oauthSessionManager: OAuthSessionManager, - workspaceProviders: WorkspaceProvider[], ): DeploymentManager { const manager = new DeploymentManager( serviceContainer, client, oauthSessionManager, - workspaceProviders, ); manager.subscribeToAuthConfigChanges(); manager.subscribeToCrossWindowChanges(); @@ -92,22 +113,31 @@ export class DeploymentManager implements vscode.Disposable { * Get the current deployment state. */ public getCurrentDeployment(): Deployment | null { - return this.#deployment; + return this.#session.deployment; } public getCurrentUserId(): string | undefined { - return this.#authedUser?.id; + return this.#session.kind === "signedIn" + ? this.#session.user.id + : undefined; } - public getAuthStateVersion(): number { - return this.#authStateVersion; + public getSnapshot(): WorkspaceSessionSnapshot { + if (this.#session.kind === "signedIn") { + return { + kind: "signedIn", + revision: this.#session.revision, + userId: this.#session.user.id, + }; + } + return { kind: "signedOut", revision: this.#session.revision }; } /** * Check if we have an authenticated deployment (does not guarantee that the current auth data is valid). */ public isAuthenticated(): boolean { - return this.contextManager.get("coder.authenticated"); + return this.#session.kind === "signedIn"; } /** @@ -119,8 +149,7 @@ export class DeploymentManager implements vscode.Disposable { public async verifyAndApplyDeployment( deployment: Deployment & { token?: string }, ): Promise { - const deploymentBefore = this.#deployment; - const authStateVersionBefore = this.#authStateVersion; + const sessionBefore = this.#session; const token = deployment.token ?? (await this.secretsManager.getSessionAuth(deployment.safeHostname)) @@ -129,9 +158,7 @@ export class DeploymentManager implements vscode.Disposable { try { const user = await tempClient.getAuthenticatedUser(); - if ( - this.#hasStateChangedSince(deploymentBefore, authStateVersionBefore) - ) { + if (this.#hasStateChangedSince(sessionBefore)) { return false; } await this.setDeployment({ ...deployment, token, user }); @@ -145,15 +172,11 @@ export class DeploymentManager implements vscode.Disposable { } /** True if disposal, login, or a deployment switch raced our await. */ - #hasStateChangedSince( - deploymentBefore: Deployment | null, - authStateVersionBefore: number, - ): boolean { + #hasStateChangedSince(sessionBefore: DeploymentSessionSnapshot): boolean { return ( this.#disposed || this.isAuthenticated() || - this.#deployment !== deploymentBefore || - this.#authStateVersion !== authStateVersionBefore + this.#session !== sessionBefore ); } @@ -169,33 +192,21 @@ export class DeploymentManager implements vscode.Disposable { user: deployment.user.username, }); const deploymentWithoutAuth = DeploymentSchema.parse(deployment); - const ourRef = this.commitDeploymentState( - deploymentWithoutAuth, - deployment.user, - ); - const authStateVersion = this.#authStateVersion; this.telemetryService.setDeploymentUrl(deployment.url); - - // Updates client credentials if (deployment.token === undefined) { this.client.setHost(deployment.url); } else { this.client.setCredentials(deployment.url, deployment.token); } - // Register auth listener before setDeployment so background token refresh - // can update client credentials via the listener + const ourRef = this.setSignedIn(deploymentWithoutAuth, deployment.user); + // Register before OAuth setup so background token refresh can update client credentials. this.registerAuthListener(); - // Contexts must be set before refresh (providers check isAuthenticated) this.updateAuthContexts(deployment.user); - this.refreshWorkspaces(); await this.oauthSessionManager.setDeployment(deploymentWithoutAuth); // Bail if a concurrent write took over during the await. - if ( - this.#deployment !== ourRef || - this.#authStateVersion !== authStateVersion - ) { + if (this.#session !== ourRef) { return; } await this.persistDeployment(deploymentWithoutAuth); @@ -205,12 +216,19 @@ export class DeploymentManager implements vscode.Disposable { * Clears the current deployment. */ public async clearDeployment(reason: DeploymentSuspendReason): Promise { - this.logger.debug("Clearing deployment", this.#deployment?.safeHostname); - this.suspendSession(reason); + this.logger.debug( + "Clearing deployment", + this.#session.deployment?.safeHostname, + ); + const wasAuthenticated = this.isAuthenticated(); this.#authListenerDisposable?.dispose(); this.#authListenerDisposable = undefined; - this.commitDeploymentState(null, null); + this.setSignedOut(null); + this.clearSessionSideEffects(); this.telemetryService.setDeploymentUrl(""); + if (wasAuthenticated) { + this.deploymentTelemetry.suspended(reason); + } await this.secretsManager.setCurrentDeployment(undefined); } @@ -221,23 +239,17 @@ export class DeploymentManager implements vscode.Disposable { */ public suspendSession(reason: DeploymentSuspendReason): void { const wasAuthenticated = this.isAuthenticated(); - this.commitDeploymentState(this.#deployment, null); - this.oauthSessionManager.clearDeployment(); - this.client.setCredentials(undefined, undefined); - this.updateAuthContexts(undefined); - this.clearWorkspaces(); + this.setSignedOut(this.#session.deployment); + this.clearSessionSideEffects(); if (wasAuthenticated) { this.deploymentTelemetry.suspended(reason); } } - /** - * Clear all workspace providers without fetching. - */ - private clearWorkspaces(): void { - for (const provider of this.workspaceProviders) { - provider.clear(); - } + private clearSessionSideEffects(): void { + this.oauthSessionManager.clearDeployment(); + this.client.setCredentials(undefined, undefined); + this.updateAuthContexts(undefined); } public dispose(): void { @@ -245,6 +257,7 @@ export class DeploymentManager implements vscode.Disposable { this.#authListenerDisposable?.dispose(); this.#authConfigDisposable?.dispose(); this.#crossWindowSyncDisposable?.dispose(); + this.#onDidChangeWorkspaceSession.dispose(); } /** @@ -253,19 +266,19 @@ export class DeploymentManager implements vscode.Disposable { * Also handles recovery from suspended session state. */ private registerAuthListener(): void { - if (!this.#deployment) { + if (!this.#session.deployment) { return; } // Capture hostname at registration time for the guard clause - const safeHostname = this.#deployment.safeHostname; + const safeHostname = this.#session.deployment.safeHostname; this.#authListenerDisposable?.dispose(); this.logger.debug("Registering auth listener for hostname", safeHostname); this.#authListenerDisposable = this.secretsManager.onDidChangeSessionAuth( safeHostname, async (auth) => { - if (this.#deployment?.safeHostname !== safeHostname) { + if (this.#session.deployment?.safeHostname !== safeHostname) { return; } @@ -299,8 +312,7 @@ export class DeploymentManager implements vscode.Disposable { private async verifyAndUpdateAuthenticatedSession( deployment: Deployment & { token: string }, ): Promise { - const deploymentBefore = this.#deployment; - const authStateVersionBefore = this.#authStateVersion; + const sessionBefore = this.#session; const tempClient = CoderApi.create( deployment.url, deployment.token, @@ -309,11 +321,7 @@ export class DeploymentManager implements vscode.Disposable { try { const user = await tempClient.getAuthenticatedUser(); - if ( - this.#disposed || - this.#deployment !== deploymentBefore || - this.#authStateVersion !== authStateVersionBefore - ) { + if (this.#disposed || this.#session !== sessionBefore) { return; } await this.setDeployment({ ...deployment, user }); @@ -324,14 +332,30 @@ export class DeploymentManager implements vscode.Disposable { } } - private commitDeploymentState( + private setSignedIn( + deployment: Deployment, + user: User, + ): DeploymentSessionSnapshot { + this.#session = { + kind: "signedIn", + revision: this.#session.revision + 1, + deployment, + user, + }; + this.#onDidChangeWorkspaceSession.fire(this.getSnapshot()); + return this.#session; + } + + private setSignedOut( deployment: Deployment | null, - authedUser: User | null, - ): Deployment | null { - this.#deployment = deployment; - this.#authedUser = authedUser; - this.#authStateVersion += 1; - return this.#deployment; + ): DeploymentSessionSnapshot { + this.#session = { + kind: "signedOut", + revision: this.#session.revision + 1, + deployment, + }; + this.#onDidChangeWorkspaceSession.fire(this.getSnapshot()); + return this.#session; } private subscribeToAuthConfigChanges(): void { @@ -357,7 +381,7 @@ export class DeploymentManager implements vscode.Disposable { try { do { this.#recoveryPending = false; - const snapshot = this.#deployment; + const snapshot = this.#session.deployment; if (this.#disposed || !snapshot || this.isAuthenticated()) { return; } @@ -416,15 +440,6 @@ export class DeploymentManager implements vscode.Disposable { this.contextManager.set("coder.isOwner", isOwner); } - /** - * Refresh all workspace providers asynchronously. - */ - private refreshWorkspaces(): void { - for (const provider of this.workspaceProviders) { - void provider.fetchAndRefresh(); - } - } - /** * Persist deployment to storage for cross-window sync. */ diff --git a/src/extension.ts b/src/extension.ts index a1832131bd..a6b44448eb 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -30,8 +30,6 @@ import { WorkspaceQuery, } from "./workspace/workspacesProvider"; -import type { Workspace } from "coder/site/src/api/typesGenerated"; - const MY_WORKSPACES_TREE_ID = "myWorkspaces"; const SHARED_WORKSPACES_TREE_ID = "sharedWorkspaces"; const ALL_WORKSPACES_TREE_ID = "allWorkspaces"; @@ -156,23 +154,19 @@ async function doActivate( ); ctx.subscriptions.push(authInterceptor); - const isAuthenticated = () => contextManager.get("coder.authenticated"); - const getAuthStateVersion = () => deploymentManager.getAuthStateVersion(); - const filterSharedWorkspaces = (workspaces: readonly Workspace[]) => { - const currentUserId = deploymentManager.getCurrentUserId(); - return currentUserId - ? workspaces.filter((workspace) => workspace.owner_id !== currentUserId) - : []; - }; + const deploymentManager = DeploymentManager.create( + serviceContainer, + client, + oauthSessionManager, + ); + ctx.subscriptions.push(deploymentManager); const myWorkspacesProvider = new WorkspaceProvider( WorkspaceQuery.Mine, client, output, - isAuthenticated, - 5, - undefined, - getAuthStateVersion, + deploymentManager, + { refreshIntervalMs: 5_000 }, ); ctx.subscriptions.push(myWorkspacesProvider); @@ -180,10 +174,7 @@ async function doActivate( WorkspaceQuery.All, client, output, - isAuthenticated, - undefined, - undefined, - getAuthStateVersion, + deploymentManager, ); ctx.subscriptions.push(allWorkspacesProvider); @@ -191,21 +182,10 @@ async function doActivate( WorkspaceQuery.Shared, client, output, - isAuthenticated, - undefined, - filterSharedWorkspaces, - getAuthStateVersion, + deploymentManager, ); ctx.subscriptions.push(sharedWorkspacesProvider); - const deploymentManager = DeploymentManager.create( - serviceContainer, - client, - oauthSessionManager, - [myWorkspacesProvider, sharedWorkspacesProvider, allWorkspacesProvider], - ); - ctx.subscriptions.push(deploymentManager); - // createTreeView, unlike registerTreeDataProvider, gives us the tree view API // (so we can see when it is visible) but otherwise they have the same effect. const myWsTree = vscode.window.createTreeView(MY_WORKSPACES_TREE_ID, { diff --git a/src/workspace/session.ts b/src/workspace/session.ts new file mode 100644 index 0000000000..67b47d6fd9 --- /dev/null +++ b/src/workspace/session.ts @@ -0,0 +1,14 @@ +import type * as vscode from "vscode"; + +export type WorkspaceSessionSnapshot = + | { readonly kind: "signedOut"; readonly revision: number } + | { + readonly kind: "signedIn"; + readonly revision: number; + readonly userId: string; + }; + +export interface WorkspaceSessionState { + getSnapshot(): WorkspaceSessionSnapshot; + onDidChange: vscode.Event; +} diff --git a/src/workspace/workspacesProvider.ts b/src/workspace/workspacesProvider.ts index c039700ba8..e0f7368ac8 100644 --- a/src/workspace/workspacesProvider.ts +++ b/src/workspace/workspacesProvider.ts @@ -20,15 +20,20 @@ import { import { type CoderApi } from "../api/coderApi"; import { type Logger } from "../logging/logger"; +import type { + WorkspaceSessionSnapshot, + WorkspaceSessionState, +} from "./session"; + export enum WorkspaceQuery { Mine = "owner:me", All = "", Shared = "shared:true", } -type WorkspaceFilter = ( - workspaces: readonly Workspace[], -) => readonly Workspace[]; +export interface WorkspaceProviderOptions { + readonly refreshIntervalMs?: number; +} /** * Polls workspaces using the provided REST client and renders them in a tree. @@ -47,8 +52,10 @@ export class WorkspaceProvider WorkspaceAgent["id"], AgentMetadataWatcher >(); + private readonly sessionChangeDisposable: vscode.Disposable; private timeout: NodeJS.Timeout | undefined; private fetching = false; + private refreshQueued = false; private visible = false; private disposed = false; @@ -56,85 +63,113 @@ export class WorkspaceProvider private readonly getWorkspacesQuery: WorkspaceQuery, private readonly client: CoderApi, private readonly logger: Logger, - private readonly isAuthenticated: () => boolean, - private readonly timerSeconds?: number, - private readonly filterWorkspaces: WorkspaceFilter = (workspaces) => - workspaces, - private readonly getStateVersion: () => number = () => 0, + private readonly sessionState: WorkspaceSessionState, + private readonly options: WorkspaceProviderOptions = {}, ) { - // No initialization. + this.sessionChangeDisposable = this.sessionState.onDidChange(() => { + this.clear(); + this.requestRefresh(); + void this.runRefreshLoop(); + }); } // fetchAndRefresh fetches new workspaces, re-renders the entire tree, then // keeps refreshing (if a timer length was provided) as long as the user is // still logged in and no errors were encountered fetching workspaces. - // Calling this while already refreshing or not visible is a no-op and will - // return immediately. + // Calling this while not visible is a no-op and will return immediately. public async fetchAndRefresh() { - if ( - this.disposed || - this.fetching || - !this.visible || - !this.isAuthenticated() - ) { + this.requestRefresh(); + await this.runRefreshLoop(); + } + + private requestRefresh(): void { + if (this.disposed || !this.visible) { return; } - this.fetching = true; + this.refreshQueued = true; + } - // It is possible we called fetchAndRefresh() manually (through the button - // for example), in which case we might still have a pending refresh that - // needs to be cleared. + private async runRefreshLoop(): Promise { + if (this.disposed || this.fetching || !this.visible) { + return; + } + + this.fetching = true; this.cancelPendingRefresh(); + let shouldScheduleRefresh = false; - let hadError = false; try { - this.workspaces = await this.fetch(); - } catch (error) { - this.logger.warn("Failed to fetch workspaces:", error); - hadError = true; - this.workspaces = []; - } - - this.fetching = false; + while (this.refreshQueued && !this.disposed && this.visible) { + this.refreshQueued = false; + shouldScheduleRefresh = false; + const session = this.sessionState.getSnapshot(); + + if (session.kind !== "signedIn") { + this.setWorkspaces([]); + continue; + } - this.refresh(); + let hadError = false; + try { + const workspaces = await this.fetch(session); + if (workspaces && !this.disposed) { + this.setWorkspaces(workspaces); + } + } catch (error) { + this.logger.warn("Failed to fetch workspaces:", error); + hadError = true; + this.setWorkspaces([]); + } - // As long as there was no error we can schedule the next refresh. - if (!hadError) { - this.maybeScheduleRefresh(); + shouldScheduleRefresh = !hadError && !this.refreshQueued; + } + } finally { + this.fetching = false; + if (this.refreshQueued && !this.disposed && this.visible) { + void this.runRefreshLoop(); + } else if (shouldScheduleRefresh && !this.disposed && this.visible) { + this.maybeScheduleRefresh(); + } } } + private setWorkspaces(workspaces: WorkspaceTreeItem[]): void { + this.workspaces = workspaces; + this.refresh(); + } + /** * Fetch workspaces and turn them into tree items. Throw an error if not * logged in or the query fails. */ - private async fetch(): Promise { + private async fetch( + session: Extract, + ): Promise { // If there is no URL configured, assume we are logged out. const url = this.client.getAxiosInstance().defaults.baseURL; if (!url) { throw new Error("not logged in"); } - const stateVersion = this.getStateVersion(); const resp = await this.client.getWorkspaces({ q: this.getWorkspacesQuery, }); - // We could have logged out while waiting for the query, or logged into a - // different deployment. + const latestSession = this.sessionState.getSnapshot(); const url2 = this.client.getAxiosInstance().defaults.baseURL; if (!url2) { throw new Error("not logged in"); - } else if (url !== url2 || stateVersion !== this.getStateVersion()) { - // In this case we need to fetch from the new deployment instead. - // TODO: It would be better to cancel this fetch when that happens, - // because this means we have to wait for the old fetch to finish before - // finally getting workspaces for the new one. - return this.fetch(); + } + if ( + url !== url2 || + latestSession.kind !== "signedIn" || + latestSession.revision !== session.revision + ) { + this.refreshQueued = true; + return undefined; } - const workspaces = this.filterWorkspaces(resp.workspaces); + const workspaces = this.filterWorkspaces(resp.workspaces, session); const oldWatcherIds = [...this.agentWatchers.keys()]; const reusedWatcherIds: string[] = []; @@ -172,17 +207,22 @@ export class WorkspaceProvider const showOwner = this.getWorkspacesQuery !== WorkspaceQuery.Mine; // Create tree items for each workspace - const workspaceTreeItems = workspaces.map((workspace: Workspace) => { - const workspaceTreeItem = new WorkspaceTreeItem( - workspace, - showOwner, - showMetadata, - ); - - return workspaceTreeItem; - }); + return workspaces.map( + (workspace: Workspace) => + new WorkspaceTreeItem(workspace, showOwner, showMetadata), + ); + } - return workspaceTreeItems; + private filterWorkspaces( + workspaces: readonly Workspace[], + session: Extract, + ): readonly Workspace[] { + if (this.getWorkspacesQuery !== WorkspaceQuery.Shared) { + return workspaces; + } + return workspaces.filter( + (workspace) => workspace.owner_id !== session.userId, + ); } /** @@ -197,10 +237,13 @@ export class WorkspaceProvider this.visible = visible; if (!visible) { this.cancelPendingRefresh(); + } else if (this.refreshQueued) { + void this.runRefreshLoop(); } else if (this.workspaces) { this.maybeScheduleRefresh(); } else { - void this.fetchAndRefresh(); + this.requestRefresh(); + void this.runRefreshLoop(); } } @@ -216,10 +259,11 @@ export class WorkspaceProvider * timeout length was provided. */ private maybeScheduleRefresh() { - if (this.timerSeconds && !this.timeout && !this.fetching) { + if (this.options.refreshIntervalMs && !this.timeout) { this.timeout = setTimeout(() => { - void this.fetchAndRefresh(); - }, this.timerSeconds * 1000); + this.requestRefresh(); + void this.runRefreshLoop(); + }, this.options.refreshIntervalMs); } } @@ -320,18 +364,24 @@ export class WorkspaceProvider * Clear all workspaces from the tree without fetching. */ public clear(): void { + this.clearState(); + this.refresh(); + } + + private clearState(): void { this.cancelPendingRefresh(); for (const watcher of this.agentWatchers.values()) { watcher.dispose(); } this.agentWatchers.clear(); this.workspaces = undefined; - this.refresh(); + this.refreshQueued = false; } public dispose() { this.disposed = true; - this.clear(); + this.clearState(); + this.sessionChangeDisposable.dispose(); } } diff --git a/test/unit/deployment/deploymentManager.test.ts b/test/unit/deployment/deploymentManager.test.ts index 1078700243..97c35cd15b 100644 --- a/test/unit/deployment/deploymentManager.test.ts +++ b/test/unit/deployment/deploymentManager.test.ts @@ -20,7 +20,6 @@ import { } from "../../mocks/testHelpers"; import type { OAuthSessionManager } from "@/oauth/sessionManager"; -import type { WorkspaceProvider } from "@/workspace/workspacesProvider"; // Mock CoderApi.create to return our mock client for validation vi.mock("@/api/coderApi", async (importOriginal) => { @@ -34,18 +33,25 @@ vi.mock("@/api/coderApi", async (importOriginal) => { }; }); -/** - * Mock WorkspaceProvider for deployment tests. - */ -class MockWorkspaceProvider { - readonly fetchAndRefresh = vi.fn(); - readonly clear = vi.fn(); -} - const TEST_URL = "https://coder.example.com"; const TEST_HOSTNAME = "coder.example.com"; const managers: DeploymentManager[] = []; +function deferred(): { + promise: Promise; + resolve: (value: T | PromiseLike) => void; +} { + let resolve!: (value: T | PromiseLike) => void; + const promise = new Promise((res) => { + resolve = res; + }); + return { promise, resolve }; +} + +async function flush(): Promise { + await new Promise((resolve) => setImmediate(resolve)); +} + afterEach(() => { for (const manager of managers) { manager.dispose(); @@ -63,7 +69,6 @@ function createTestContext() { const mockClient = new MockCoderApi(); // For verifyAndApplyDeployment, we use a separate mock for validation const validationMockClient = new MockCoderApi(); - const mockWorkspaceProvider = new MockWorkspaceProvider(); const mockOAuthSessionManager = new MockOAuthSessionManager(); const secretStorage = new InMemorySecretStorage(); const memento = new InMemoryMemento(); @@ -91,7 +96,6 @@ function createTestContext() { }), mockClient as unknown as CoderApi, mockOAuthSessionManager as unknown as OAuthSessionManager, - [mockWorkspaceProvider as unknown as WorkspaceProvider], ); managers.push(manager); @@ -101,7 +105,6 @@ function createTestContext() { secretsManager, contextManager, mockOAuthSessionManager, - mockWorkspaceProvider, telemetrySink, telemetryService, setDeploymentUrlSpy, @@ -179,6 +182,29 @@ describe("DeploymentManager", () => { expect(persisted?.url).toBe(TEST_URL); }); + it("does not persist a deployment after a newer state wins", async () => { + const { mockOAuthSessionManager, secretsManager, manager } = + createTestContext(); + const oauthSetDeployment = deferred(); + mockOAuthSessionManager.setDeployment.mockReturnValueOnce( + oauthSetDeployment.promise, + ); + const firstSetDeployment = manager.setDeployment({ + url: TEST_URL, + safeHostname: TEST_HOSTNAME, + token: "test-token", + user: createMockUser(), + }); + + await flush(); + await manager.clearDeployment("logout"); + oauthSetDeployment.resolve(); + await firstSetDeployment; + + expect(await secretsManager.getCurrentDeployment()).toBeNull(); + expect(manager.isAuthenticated()).toBe(false); + }); + it("notifies telemetry of the deployment URL", async () => { const { setDeploymentUrlSpy, manager } = createTestContext(); @@ -297,6 +323,35 @@ describe("DeploymentManager", () => { expect(validationMockClient.disposed).toBe(true); }); + + it("does not apply stale validation after a concurrent login", async () => { + const { mockClient, validationMockClient, manager } = createTestContext(); + const remoteUser = createMockUser({ id: "remote-user" }); + const localUser = createMockUser({ id: "local-user" }); + const validation = deferred(); + validationMockClient.getAuthenticatedUser.mockReturnValueOnce( + validation.promise, + ); + const remoteLogin = manager.verifyAndApplyDeployment({ + url: TEST_URL, + safeHostname: TEST_HOSTNAME, + token: "remote-token", + }); + + await flush(); + await manager.setDeployment({ + url: "https://local.example.com", + safeHostname: "local.example.com", + token: "local-token", + user: localUser, + }); + validation.resolve(remoteUser); + + expect(await remoteLogin).toBe(false); + expect(mockClient.host).toBe("https://local.example.com"); + expect(mockClient.token).toBe("local-token"); + expect(manager.getCurrentUserId()).toBe("local-user"); + }); }); describe("cross-window sync", () => { @@ -345,7 +400,7 @@ describe("DeploymentManager", () => { }); // Wait for async handler - await new Promise((resolve) => setImmediate(resolve)); + await flush(); expect(mockClient.host).toBe(TEST_URL); expect(mockClient.token).toBe("synced-token"); @@ -374,7 +429,7 @@ describe("DeploymentManager", () => { }); // Wait for async handler - await new Promise((resolve) => setImmediate(resolve)); + await flush(); expect(mockClient.host).toBe(TEST_URL); expect(mockClient.token).toBe(""); @@ -407,12 +462,43 @@ describe("DeploymentManager", () => { url: TEST_URL, token: "refreshed-token", }); - await new Promise((resolve) => setImmediate(resolve)); + await flush(); expect(mockClient.token).toBe("refreshed-token"); expect(manager.getCurrentUserId()).toBe("refreshed-user"); expect(manager.isAuthenticated()).toBe(true); }); + + it("does not apply stale token validation after logout", async () => { + const { mockClient, validationMockClient, secretsManager, manager } = + createTestContext(); + const user = createMockUser({ id: "initial-user" }); + const refreshedUser = createMockUser({ id: "refreshed-user" }); + const validation = deferred(); + validationMockClient.getAuthenticatedUser.mockReturnValueOnce( + validation.promise, + ); + await manager.setDeployment({ + url: TEST_URL, + safeHostname: TEST_HOSTNAME, + token: "initial-token", + user, + }); + + await secretsManager.setSessionAuth(TEST_HOSTNAME, { + url: TEST_URL, + token: "refreshed-token", + }); + await flush(); + await manager.clearDeployment("logout"); + validation.resolve(refreshedUser); + await flush(); + + expect(mockClient.host).toBeUndefined(); + expect(mockClient.token).toBeUndefined(); + expect(manager.getCurrentDeployment()).toBeNull(); + expect(manager.getCurrentUserId()).toBeUndefined(); + }); }); describe("logout", () => { @@ -471,13 +557,8 @@ describe("DeploymentManager", () => { }); it("clears auth state but keeps deployment for re-login", async () => { - const { - mockClient, - contextManager, - mockOAuthSessionManager, - mockWorkspaceProvider, - manager, - } = createTestContext(); + const { mockClient, contextManager, mockOAuthSessionManager, manager } = + createTestContext(); await manager.setDeployment({ url: TEST_URL, @@ -496,7 +577,6 @@ describe("DeploymentManager", () => { expect(contextManager.get("coder.authenticated")).toBe(false); expect(manager.getCurrentUserId()).toBeUndefined(); expect(manager.isAuthenticated()).toBe(false); - expect(mockWorkspaceProvider.clear).toHaveBeenCalled(); // Deployment is retained for easy re-login expect(manager.getCurrentDeployment()).toMatchObject({ @@ -619,7 +699,7 @@ describe("DeploymentManager", () => { token: "recovered-token", }); - await new Promise((resolve) => setImmediate(resolve)); + await flush(); // Should recover and be authenticated again expect(mockClient.token).toBe("recovered-token"); diff --git a/test/unit/workspace/workspacesProvider.test.ts b/test/unit/workspace/workspacesProvider.test.ts index eb1d4236ad..15e9e8aeef 100644 --- a/test/unit/workspace/workspacesProvider.test.ts +++ b/test/unit/workspace/workspacesProvider.test.ts @@ -1,29 +1,58 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; +import * as vscode from "vscode"; import { WorkspaceProvider, WorkspaceQuery, + type AgentTreeItem, + type WorkspaceTreeItem, } from "@/workspace/workspacesProvider"; -import { workspace as createWorkspace } from "@repo/mocks"; +import { + agent as createAgent, + resource as createResource, + workspace as createWorkspace, +} from "@repo/mocks"; import { createMockLogger } from "../../mocks/testHelpers"; -import type { Workspace } from "coder/site/src/api/typesGenerated"; +import type { + Workspace, + WorkspaceAgent, + WorkspaceApp, + WorkspaceAppStatus, +} from "coder/site/src/api/typesGenerated"; +import type { AgentMetadataWatcher } from "@/api/agentMetadataHelper"; +import type { AgentMetadataEvent } from "@/api/api-helper"; import type { CoderApi } from "@/api/coderApi"; +import type { WorkspaceSessionSnapshot } from "@/workspace/session"; + +vi.mock("@/api/agentMetadataHelper", async (importOriginal) => { + const original = + await importOriginal(); + return { + ...original, + createAgentMetadataWatcher: vi.fn(), + }; +}); + +const { createAgentMetadataWatcher } = + await import("@/api/agentMetadataHelper"); const baseUrl = "https://coder.example.com"; +const currentUserId = "current-user"; + +interface WorkspacesResponse { + workspaces: readonly Workspace[]; + count: number; +} -class MockClient { +class TestClient { baseURL: string | undefined = baseUrl; readonly getWorkspaces = vi.fn( - (_req: { - q: string; - }): Promise<{ - workspaces: readonly Workspace[]; - count: number; - }> => Promise.resolve({ workspaces: [], count: 0 }), + (_req: { q: string }): Promise => + Promise.resolve({ workspaces: [], count: 0 }), ); getAxiosInstance() { @@ -35,151 +64,567 @@ class MockClient { } } +class TestSessionState { + private revision = 0; + private readonly onDidChangeEmitter = + new vscode.EventEmitter(); + readonly onDidChange = this.onDidChangeEmitter.event; + private snapshot: WorkspaceSessionSnapshot = { + kind: "signedIn", + revision: this.revision, + userId: currentUserId, + }; + + getSnapshot(): WorkspaceSessionSnapshot { + return this.snapshot; + } + + signIn(userId = currentUserId): void { + this.revision += 1; + this.snapshot = { kind: "signedIn", revision: this.revision, userId }; + this.onDidChangeEmitter.fire(this.snapshot); + } + + signOut(): void { + this.revision += 1; + this.snapshot = { kind: "signedOut", revision: this.revision }; + this.onDidChangeEmitter.fire(this.snapshot); + } +} + +type TestWatcher = AgentMetadataWatcher & { + onChangeEmitter: vscode.EventEmitter; + dispose: ReturnType void>>; +}; + describe("WorkspaceProvider", () => { - let client: MockClient; - let isAuthenticated: () => boolean; + let client: TestClient; + let logger: ReturnType; + let session: TestSessionState; + let watchers: TestWatcher[]; beforeEach(() => { - client = new MockClient(); - isAuthenticated = vi.fn<() => boolean>(() => true); + client = new TestClient(); + logger = createMockLogger(); + session = new TestSessionState(); + watchers = []; + vi.mocked(createAgentMetadataWatcher).mockImplementation(() => { + const onChangeEmitter = new vscode.EventEmitter(); + const watcher: TestWatcher = { + onChangeEmitter, + onChange: onChangeEmitter.event, + dispose: vi.fn<() => void>(), + }; + watchers.push(watcher); + return Promise.resolve(watcher); + }); }); - function createProvider( + function makeProvider( query: WorkspaceQuery, - options: { - filterWorkspaces?: ( - workspaces: readonly Workspace[], - ) => readonly Workspace[]; - getStateVersion?: () => number; - timerSeconds?: number; - } = {}, - ) { + options?: { refreshIntervalMs?: number }, + ): WorkspaceProvider { return new WorkspaceProvider( query, client as unknown as CoderApi, - createMockLogger(), - isAuthenticated, - options.timerSeconds, - options.filterWorkspaces, - options.getStateVersion, + logger, + session, + options, ); } - async function fetchVisible(provider: WorkspaceProvider) { - provider.setVisibility(true); - await new Promise((resolve) => setImmediate(resolve)); + function workspace( + overrides: Parameters[0] = {}, + ): Workspace { + return createWorkspace(overrides); + } + + function agent(overrides: Partial = {}): WorkspaceAgent { + return createAgent(overrides); } - it("queries shared workspaces", async () => { - const provider = createProvider(WorkspaceQuery.Shared); + function workspaceWithAgents( + workspaceOverrides: Parameters[0] = {}, + agents: WorkspaceAgent[], + ): Workspace { + return workspace({ + ...workspaceOverrides, + latest_build: { + ...workspaceOverrides.latest_build, + resources: [createResource({ agents })], + }, + }); + } - await fetchVisible(provider); + function appStatus( + overrides: Partial = {}, + ): WorkspaceAppStatus { + return { + id: "status-1", + created_at: "2024-01-01T00:00:00Z", + workspace_id: "workspace-1", + agent_id: "agent-1", + app_id: "app-1", + state: "working", + message: "Opening pull request", + uri: "https://example.com/pr/1", + icon: "", + needs_user_attention: false, + ...overrides, + }; + } - expect(client.getWorkspaces).toHaveBeenCalledWith({ - q: WorkspaceQuery.Shared, + function app(overrides: Partial = {}): WorkspaceApp { + return { + id: "app-1", + external: false, + slug: "app", + subdomain: false, + sharing_level: "owner", + health: "healthy", + hidden: false, + open_in: "tab", + statuses: [], + ...overrides, + }; + } + + function metadata( + overrides: Partial = {}, + ): AgentMetadataEvent { + return { + result: { + collected_at: "2024-01-01T00:00:00Z", + age: 0, + value: "42", + error: "", + }, + description: { + display_name: "CPU", + key: "cpu", + script: "cpu.sh", + interval: 5, + timeout: 1, + }, + ...overrides, + }; + } + + function respondOnce(workspaces: readonly Workspace[]): void { + client.getWorkspaces.mockResolvedValueOnce({ + workspaces, + count: workspaces.length, }); + } + + function pendingWorkspaces(): { + resolve: (workspaces: readonly Workspace[]) => void; + } { + let resolve!: (workspaces: readonly Workspace[]) => void; + client.getWorkspaces.mockReturnValueOnce( + new Promise((res) => { + resolve = (workspaces) => res({ workspaces, count: workspaces.length }); + }), + ); + return { resolve }; + } + + async function flush(): Promise { + await new Promise((resolve) => setImmediate(resolve)); + } + + async function flushPromises(): Promise { + await Promise.resolve(); + await Promise.resolve(); + await Promise.resolve(); + } + + async function show(provider: WorkspaceProvider): Promise { + provider.setVisibility(true); + await flush(); + } + + async function labels(provider: WorkspaceProvider): Promise { + return (await provider.getChildren()).map((item) => item.label); + } + + it("does not fetch while signed out", async () => { + session.signOut(); + const provider = makeProvider(WorkspaceQuery.Mine); + + await show(provider); + + expect(client.getWorkspaces).not.toHaveBeenCalled(); + expect(await provider.getChildren()).toEqual([]); }); - it("applies the workspace filter before rendering tree items", async () => { - client.getWorkspaces.mockResolvedValueOnce({ - workspaces: [ - createWorkspace({ - id: "owned-workspace", - owner_id: "current-user", - owner_name: "current", - name: "owned", - }), - createWorkspace({ - id: "shared-workspace", - owner_id: "other-user", + it.each([ + [WorkspaceQuery.Mine, "owner:me"], + [WorkspaceQuery.Shared, "shared:true"], + [WorkspaceQuery.All, ""], + ])("fetches %s with the expected query", async (query, expectedQuery) => { + const provider = makeProvider(query); + + await show(provider); + + expect(client.getWorkspaces).toHaveBeenCalledWith({ q: expectedQuery }); + }); + + it.each([ + { + query: WorkspaceQuery.Mine, + label: "dev", + collapsibleState: vscode.TreeItemCollapsibleState.Expanded, + }, + { + query: WorkspaceQuery.Shared, + label: "alice / dev", + collapsibleState: vscode.TreeItemCollapsibleState.Collapsed, + }, + { + query: WorkspaceQuery.All, + label: "alice / dev", + collapsibleState: vscode.TreeItemCollapsibleState.Collapsed, + }, + ])( + "renders top-level workspace items for $query", + async ({ query, label, collapsibleState }) => { + respondOnce([ + workspace({ + id: "workspace-1", + name: "dev", + owner_id: "alice-id", owner_name: "alice", - name: "shared", }), - ], - count: 2, - }); - const provider = createProvider(WorkspaceQuery.Shared, { - filterWorkspaces: (workspaces) => - workspaces.filter((workspace) => workspace.owner_id !== "current-user"), - }); + ]); + const provider = makeProvider(query); + + await show(provider); + const [item] = (await provider.getChildren()) as WorkspaceTreeItem[]; + + expect(item?.label).toBe(label); + expect(item?.description).toBe("running"); + expect(item?.collapsibleState).toBe(collapsibleState); + expect(item?.contextValue).toContain("running"); + }, + ); + + it("filters current-user-owned workspaces from shared results", async () => { + respondOnce([ + workspace({ + id: "owned-shared-out", + name: "owned", + owner_id: currentUserId, + owner_name: "current", + }), + workspace({ + id: "shared-with-me", + name: "shared", + owner_id: "alice-id", + owner_name: "alice", + }), + ]); + const provider = makeProvider(WorkspaceQuery.Shared); + + await show(provider); + + expect(await labels(provider)).toEqual(["alice / shared"]); + }); + + it.each([WorkspaceQuery.Mine, WorkspaceQuery.All])( + "does not apply shared ownership filtering to %s", + async (query) => { + respondOnce([ + workspace({ + id: "owned", + name: "owned", + owner_id: currentUserId, + owner_name: "current", + }), + ]); + const provider = makeProvider(query); + + await show(provider); + + expect(await labels(provider)).toHaveLength(1); + }, + ); - await fetchVisible(provider); - const children = await provider.getChildren(); + it("clears rendered workspaces when the session signs out", async () => { + respondOnce([workspace({ name: "dev" })]); + const provider = makeProvider(WorkspaceQuery.Mine); - expect(children).toHaveLength(1); - expect(children[0]?.label).toBe("alice / shared"); + await show(provider); + expect(await labels(provider)).toEqual(["dev"]); + + session.signOut(); + await flush(); + + expect(await provider.getChildren()).toEqual([]); }); - it("fails closed when the shared workspace filter cannot identify the current user", async () => { + it("does not render a pending response after sign-out", async () => { + const pending = pendingWorkspaces(); + const provider = makeProvider(WorkspaceQuery.Shared); + + provider.setVisibility(true); + await flush(); + session.signOut(); + pending.resolve([workspace({ owner_id: "alice-id", owner_name: "alice" })]); + await flush(); + + expect(await provider.getChildren()).toEqual([]); + }); + + it("refetches when the session changes while a request is pending", async () => { + const pending = pendingWorkspaces(); client.getWorkspaces.mockResolvedValueOnce({ workspaces: [ - createWorkspace({ - id: "shared-workspace", - owner_id: "other-user", + workspace({ + id: "fresh-workspace", + owner_id: "alice-id", owner_name: "alice", - name: "shared", + name: "fresh", }), ], count: 1, }); - const provider = createProvider(WorkspaceQuery.Shared, { - filterWorkspaces: () => [], - }); + const provider = makeProvider(WorkspaceQuery.Shared); - await fetchVisible(provider); - const children = await provider.getChildren(); + provider.setVisibility(true); + await flush(); + session.signIn("second-user"); + pending.resolve([ + workspace({ + id: "stale-workspace", + owner_id: "bob-id", + owner_name: "bob", + name: "stale", + }), + ]); + await flush(); + await flush(); - expect(children).toEqual([]); + expect(client.getWorkspaces).toHaveBeenCalledTimes(2); + expect(await labels(provider)).toEqual(["alice / fresh"]); }); - it("refetches when auth state changes while a request is pending", async () => { - let stateVersion = 1; - let resolveFirst!: (response: { - workspaces: readonly Workspace[]; - count: number; - }) => void; - client.getWorkspaces - .mockReturnValueOnce( - new Promise((resolve) => { - resolveFirst = resolve; - }), - ) - .mockResolvedValueOnce({ - workspaces: [ - createWorkspace({ - id: "fresh-workspace", - owner_id: "other-user", - owner_name: "alice", - name: "fresh", - }), - ], - count: 1, - }); - const provider = createProvider(WorkspaceQuery.Shared, { - getStateVersion: () => stateVersion, - }); + it("does not fetch while hidden", async () => { + const provider = makeProvider(WorkspaceQuery.Mine); + + await provider.fetchAndRefresh(); + + expect(client.getWorkspaces).not.toHaveBeenCalled(); + }); + + it("renders a response that completes after the tree is hidden", async () => { + const pending = pendingWorkspaces(); + const provider = makeProvider(WorkspaceQuery.Mine); provider.setVisibility(true); - await new Promise((resolve) => setImmediate(resolve)); - stateVersion = 2; - resolveFirst({ - workspaces: [ - createWorkspace({ - id: "stale-workspace", - owner_id: "other-user", - owner_name: "bob", - name: "stale", + await flush(); + provider.setVisibility(false); + pending.resolve([workspace({ name: "dev" })]); + await flush(); + + expect(await labels(provider)).toEqual(["dev"]); + }); + + it("fetches queued session changes when the tree is shown again", async () => { + const pending = pendingWorkspaces(); + respondOnce([workspace({ name: "fresh", owner_id: "alice-id" })]); + const provider = makeProvider(WorkspaceQuery.Mine); + + provider.setVisibility(true); + await flush(); + provider.setVisibility(false); + session.signIn("second-user"); + pending.resolve([workspace({ name: "stale" })]); + await flush(); + + provider.setVisibility(true); + await flush(); + + expect(client.getWorkspaces).toHaveBeenCalledTimes(2); + expect(await labels(provider)).toEqual(["fresh"]); + }); + + it("clears and logs when fetch fails", async () => { + respondOnce([workspace({ name: "dev" })]); + client.getWorkspaces.mockRejectedValueOnce(new Error("network down")); + const provider = makeProvider(WorkspaceQuery.Mine); + + await show(provider); + expect(await labels(provider)).toEqual(["dev"]); + await provider.fetchAndRefresh(); + + expect(await provider.getChildren()).toEqual([]); + expect(logger.warn).toHaveBeenCalledWith( + "Failed to fetch workspaces:", + expect.any(Error), + ); + }); + + it("schedules polling after a successful refresh", async () => { + vi.useFakeTimers(); + try { + respondOnce([workspace({ name: "first" })]); + respondOnce([workspace({ name: "second" })]); + const provider = makeProvider(WorkspaceQuery.Mine, { + refreshIntervalMs: 5_000, + }); + + provider.setVisibility(true); + await flushPromises(); + expect(await labels(provider)).toEqual(["first"]); + + await vi.advanceTimersByTimeAsync(5_000); + await flushPromises(); + + expect(client.getWorkspaces).toHaveBeenCalledTimes(2); + expect(await labels(provider)).toEqual(["second"]); + } finally { + vi.useRealTimers(); + } + }); + + it("renders workspace child agents", async () => { + respondOnce([ + workspaceWithAgents({ name: "dev" }, [ + agent({ id: "agent-1", name: "main", status: "connected" }), + agent({ id: "agent-2", name: "sidecar", status: "disconnected" }), + ]), + ]); + const provider = makeProvider(WorkspaceQuery.Mine); + + await show(provider); + const [workspaceItem] = + (await provider.getChildren()) as WorkspaceTreeItem[]; + const agentItems = (await provider.getChildren( + workspaceItem, + )) as AgentTreeItem[]; + + expect(agentItems.map((item) => item.label)).toEqual(["main", "sidecar"]); + expect(agentItems.map((item) => item.description)).toEqual([ + "connected", + "disconnected", + ]); + }); + + it("renders app status children for an agent", async () => { + respondOnce([ + workspaceWithAgents({ name: "dev" }, [ + agent({ + id: "agent-1", + apps: [ + app({ + command: "open-pr", + statuses: [ + appStatus({ id: "status-1", message: "First" }), + appStatus({ id: "status-2", message: "Second" }), + ], + }), + ], }), - ], - count: 1, + ]), + ]); + const provider = makeProvider(WorkspaceQuery.Mine); + + await show(provider); + const [workspaceItem] = + (await provider.getChildren()) as WorkspaceTreeItem[]; + const [agentItem] = (await provider.getChildren( + workspaceItem, + )) as AgentTreeItem[]; + const [section] = await provider.getChildren(agentItem); + const statuses = await provider.getChildren(section); + + expect(section?.label).toBe("App Statuses"); + expect(statuses.map((item) => item.description)).toEqual([ + "Second", + "First", + ]); + expect(statuses[0]?.command).toMatchObject({ + command: "coder.openAppStatus", }); - await new Promise((resolve) => setImmediate(resolve)); - await new Promise((resolve) => setImmediate(resolve)); + }); - const children = await provider.getChildren(); + it("renders metadata and metadata errors for watched agents", async () => { + respondOnce([ + workspaceWithAgents({ name: "dev" }, [ + agent({ id: "agent-1", name: "main" }), + ]), + ]); + const provider = makeProvider(WorkspaceQuery.Mine); - expect(client.getWorkspaces).toHaveBeenCalledTimes(2); - expect(children).toHaveLength(1); - expect(children[0]?.label).toBe("alice / fresh"); + await show(provider); + const [workspaceItem] = + (await provider.getChildren()) as WorkspaceTreeItem[]; + const [agentItem] = (await provider.getChildren( + workspaceItem, + )) as AgentTreeItem[]; + + watchers[0].metadata = [metadata()]; + let [section] = await provider.getChildren(agentItem); + const metadataItems = await provider.getChildren(section); + expect(section?.label).toBe("Agent Metadata"); + expect(metadataItems[0]?.label).toBe("CPU: 42"); + + watchers[0].error = new Error("boom"); + [section] = await provider.getChildren(agentItem); + expect(section?.label).toBe("Failed to query metadata: boom"); + }); + + it("reuses and disposes metadata watchers as agents change", async () => { + respondOnce([ + workspaceWithAgents({ name: "dev" }, [agent({ id: "agent-1" })]), + ]); + respondOnce([ + workspaceWithAgents({ name: "dev" }, [agent({ id: "agent-2" })]), + ]); + const provider = makeProvider(WorkspaceQuery.Mine); + + await show(provider); + expect(vi.mocked(createAgentMetadataWatcher)).toHaveBeenCalledWith( + "agent-1", + client, + ); + + await provider.fetchAndRefresh(); + await flush(); + + expect(watchers[0].dispose).toHaveBeenCalled(); + expect(vi.mocked(createAgentMetadataWatcher)).toHaveBeenCalledWith( + "agent-2", + client, + ); + }); + + it("clear removes workspaces and disposes metadata watchers", async () => { + respondOnce([ + workspaceWithAgents({ name: "dev" }, [agent({ id: "agent-1" })]), + ]); + const provider = makeProvider(WorkspaceQuery.Mine); + + await show(provider); + provider.clear(); + + expect(await provider.getChildren()).toEqual([]); + expect(watchers[0].dispose).toHaveBeenCalled(); + }); + + it("dispose removes metadata watchers without firing a tree refresh", async () => { + respondOnce([ + workspaceWithAgents({ name: "dev" }, [agent({ id: "agent-1" })]), + ]); + const provider = makeProvider(WorkspaceQuery.Mine); + const onDidChangeTreeData = vi.fn(); + + await show(provider); + provider.onDidChangeTreeData(onDidChangeTreeData); + provider.dispose(); + + expect(watchers[0].dispose).toHaveBeenCalled(); + expect(onDidChangeTreeData).not.toHaveBeenCalled(); }); }); From ea22b418be3764df42f75cb5bb9a113ab95d2e07 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Fri, 5 Jun 2026 12:22:07 +0300 Subject: [PATCH 3/5] refactor: simplify shared workspaces provider and tests - collapse per-query branches into a readonly WORKSPACE_QUERY_CONFIG table - drop the unreachable re-entry branch in the workspace refresh loop - remove test-only getCurrentUserId; tests read getSnapshot via a helper - extract the repeated tree-view wiring in extension.ts into a helper - drive metadata tests through the real watcher over MockEventStream - move shared test doubles (session, workspaces client, flush) into testHelpers - replace describe-scoped state + beforeEach/afterEach with a setup() helper --- src/deployment/deploymentManager.ts | 167 +++---- src/deployment/sessionStore.ts | 76 +++ src/extension.ts | 60 +-- src/workspace/workspacesProvider.ts | 172 +++---- test/mocks/testHelpers.ts | 111 ++++- .../unit/deployment/deploymentManager.test.ts | 67 ++- .../unit/workspace/workspacesProvider.test.ts | 458 ++++++------------ 7 files changed, 527 insertions(+), 584 deletions(-) create mode 100644 src/deployment/sessionStore.ts diff --git a/src/deployment/deploymentManager.ts b/src/deployment/deploymentManager.ts index 58941eeff5..5c02437cf1 100644 --- a/src/deployment/deploymentManager.ts +++ b/src/deployment/deploymentManager.ts @@ -1,5 +1,3 @@ -import * as vscode from "vscode"; - import { CoderApi } from "../api/coderApi"; import { CONFIG_CHANGE_DEBOUNCE_MS, @@ -19,6 +17,7 @@ import { type OAuthSessionManager } from "../oauth/sessionManager"; import { getAuthConfigWatchSettings } from "../settings/authConfig"; import { type TelemetryService } from "../telemetry/service"; +import { SessionStore, type SessionData } from "./sessionStore"; import { DeploymentSchema, type Deployment, @@ -26,25 +25,13 @@ import { } from "./types"; import type { User } from "coder/site/src/api/typesGenerated"; +import type * as vscode from "vscode"; import type { WorkspaceSessionSnapshot, WorkspaceSessionState, } from "../workspace/session"; -type DeploymentSessionSnapshot = - | { - readonly kind: "signedOut"; - readonly revision: number; - readonly deployment: Deployment | null; - } - | { - readonly kind: "signedIn"; - readonly revision: number; - readonly deployment: Deployment; - readonly user: User; - }; - /** * Manages deployment state for the extension. * @@ -66,14 +53,8 @@ export class DeploymentManager private readonly telemetryService: TelemetryService; private readonly deploymentTelemetry: DeploymentTelemetry; - #session: DeploymentSessionSnapshot = { - kind: "signedOut", - revision: 0, - deployment: null, - }; - readonly #onDidChangeWorkspaceSession = - new vscode.EventEmitter(); - public readonly onDidChange = this.#onDidChangeWorkspaceSession.event; + readonly #sessionStore = new SessionStore(); + public readonly onDidChange = this.#sessionStore.onDidChange; #disposed = false; #authListenerDisposable: vscode.Disposable | undefined; #authConfigDisposable: vscode.Disposable | undefined; @@ -113,51 +94,37 @@ export class DeploymentManager * Get the current deployment state. */ public getCurrentDeployment(): Deployment | null { - return this.#session.deployment; - } - - public getCurrentUserId(): string | undefined { - return this.#session.kind === "signedIn" - ? this.#session.user.id - : undefined; + return this.#sessionStore.current.deployment; } public getSnapshot(): WorkspaceSessionSnapshot { - if (this.#session.kind === "signedIn") { - return { - kind: "signedIn", - revision: this.#session.revision, - userId: this.#session.user.id, - }; - } - return { kind: "signedOut", revision: this.#session.revision }; + return this.#sessionStore.getSnapshot(); } /** * Check if we have an authenticated deployment (does not guarantee that the current auth data is valid). */ public isAuthenticated(): boolean { - return this.#session.kind === "signedIn"; + return this.#sessionStore.current.kind === "signedIn"; } /** - * Verify credentials and apply the deployment on success. Used for - * fresh logins and for un-suspending a session after auth settings or - * a token become valid again. Bails if state moved during the verify - * (logout, another login, dispose), so callers don't need a race guard. + * Verify credentials and apply the deployment on success, signing in. Used + * for fresh logins and for un-suspending a session once auth settings or a + * token become valid again. Bails if state moved during the verify (logout, + * another login, dispose), so callers don't need a race guard. */ - public async verifyAndApplyDeployment( + public async verifyAndApplySession( deployment: Deployment & { token?: string }, ): Promise { - const sessionBefore = this.#session; + const sessionBefore = this.#sessionStore.current; const token = deployment.token ?? (await this.secretsManager.getSessionAuth(deployment.safeHostname)) ?.token; - const tempClient = CoderApi.create(deployment.url, token, this.logger); try { - const user = await tempClient.getAuthenticatedUser(); + const user = await this.#verifyCredentials(deployment.url, token); if (this.#hasStateChangedSince(sessionBefore)) { return false; } @@ -166,17 +133,31 @@ export class DeploymentManager } catch (e) { this.logger.warn("Failed to authenticate with deployment:", e); return false; + } + } + + /** + * Verify credentials with a throwaway client and return the authenticated + * user. Throws if the credentials are rejected. + */ + async #verifyCredentials( + url: string, + token: string | undefined, + ): Promise { + const tempClient = CoderApi.create(url, token, this.logger); + try { + return await tempClient.getAuthenticatedUser(); } finally { tempClient.dispose(); } } /** True if disposal, login, or a deployment switch raced our await. */ - #hasStateChangedSince(sessionBefore: DeploymentSessionSnapshot): boolean { + #hasStateChangedSince(sessionBefore: SessionData): boolean { return ( this.#disposed || this.isAuthenticated() || - this.#session !== sessionBefore + this.#sessionStore.current !== sessionBefore ); } @@ -199,14 +180,17 @@ export class DeploymentManager this.client.setCredentials(deployment.url, deployment.token); } - const ourRef = this.setSignedIn(deploymentWithoutAuth, deployment.user); + const ourRef = this.#sessionStore.signIn( + deploymentWithoutAuth, + deployment.user, + ); // Register before OAuth setup so background token refresh can update client credentials. this.registerAuthListener(); this.updateAuthContexts(deployment.user); await this.oauthSessionManager.setDeployment(deploymentWithoutAuth); // Bail if a concurrent write took over during the await. - if (this.#session !== ourRef) { + if (this.#sessionStore.current !== ourRef) { return; } await this.persistDeployment(deploymentWithoutAuth); @@ -218,13 +202,13 @@ export class DeploymentManager public async clearDeployment(reason: DeploymentSuspendReason): Promise { this.logger.debug( "Clearing deployment", - this.#session.deployment?.safeHostname, + this.#sessionStore.current.deployment?.safeHostname, ); const wasAuthenticated = this.isAuthenticated(); this.#authListenerDisposable?.dispose(); this.#authListenerDisposable = undefined; - this.setSignedOut(null); - this.clearSessionSideEffects(); + this.#sessionStore.signOut(null); + this.clearSideEffects(); this.telemetryService.setDeploymentUrl(""); if (wasAuthenticated) { this.deploymentTelemetry.suspended(reason); @@ -239,14 +223,14 @@ export class DeploymentManager */ public suspendSession(reason: DeploymentSuspendReason): void { const wasAuthenticated = this.isAuthenticated(); - this.setSignedOut(this.#session.deployment); - this.clearSessionSideEffects(); + this.#sessionStore.signOut(this.#sessionStore.current.deployment); + this.clearSideEffects(); if (wasAuthenticated) { this.deploymentTelemetry.suspended(reason); } } - private clearSessionSideEffects(): void { + private clearSideEffects(): void { this.oauthSessionManager.clearDeployment(); this.client.setCredentials(undefined, undefined); this.updateAuthContexts(undefined); @@ -257,7 +241,7 @@ export class DeploymentManager this.#authListenerDisposable?.dispose(); this.#authConfigDisposable?.dispose(); this.#crossWindowSyncDisposable?.dispose(); - this.#onDidChangeWorkspaceSession.dispose(); + this.#sessionStore.dispose(); } /** @@ -266,25 +250,28 @@ export class DeploymentManager * Also handles recovery from suspended session state. */ private registerAuthListener(): void { - if (!this.#session.deployment) { + const deployment = this.#sessionStore.current.deployment; + if (!deployment) { return; } // Capture hostname at registration time for the guard clause - const safeHostname = this.#session.deployment.safeHostname; + const safeHostname = deployment.safeHostname; this.#authListenerDisposable?.dispose(); this.logger.debug("Registering auth listener for hostname", safeHostname); this.#authListenerDisposable = this.secretsManager.onDidChangeSessionAuth( safeHostname, async (auth) => { - if (this.#session.deployment?.safeHostname !== safeHostname) { + if ( + this.#sessionStore.current.deployment?.safeHostname !== safeHostname + ) { return; } if (auth) { if (this.isAuthenticated()) { - await this.verifyAndUpdateAuthenticatedSession({ + await this.verifyAndUpdateSession({ url: auth.url, safeHostname, token: auth.token, @@ -309,55 +296,24 @@ export class DeploymentManager ); } - private async verifyAndUpdateAuthenticatedSession( + private async verifyAndUpdateSession( deployment: Deployment & { token: string }, ): Promise { - const sessionBefore = this.#session; - const tempClient = CoderApi.create( - deployment.url, - deployment.token, - this.logger, - ); - + const sessionBefore = this.#sessionStore.current; try { - const user = await tempClient.getAuthenticatedUser(); - if (this.#disposed || this.#session !== sessionBefore) { + const user = await this.#verifyCredentials( + deployment.url, + deployment.token, + ); + if (this.#disposed || this.#sessionStore.current !== sessionBefore) { return; } await this.setDeployment({ ...deployment, user }); } catch (e) { this.logger.warn("Failed to authenticate updated session:", e); - } finally { - tempClient.dispose(); } } - private setSignedIn( - deployment: Deployment, - user: User, - ): DeploymentSessionSnapshot { - this.#session = { - kind: "signedIn", - revision: this.#session.revision + 1, - deployment, - user, - }; - this.#onDidChangeWorkspaceSession.fire(this.getSnapshot()); - return this.#session; - } - - private setSignedOut( - deployment: Deployment | null, - ): DeploymentSessionSnapshot { - this.#session = { - kind: "signedOut", - revision: this.#session.revision + 1, - deployment, - }; - this.#onDidChangeWorkspaceSession.fire(this.getSnapshot()); - return this.#session; - } - private subscribeToAuthConfigChanges(): void { this.#authConfigDisposable = watchConfigurationChanges( getAuthConfigWatchSettings(), @@ -381,14 +337,17 @@ export class DeploymentManager try { do { this.#recoveryPending = false; - const snapshot = this.#session.deployment; - if (this.#disposed || !snapshot || this.isAuthenticated()) { + const deployment = this.#sessionStore.current.deployment; + if (this.#disposed || !deployment || this.isAuthenticated()) { return; } this.logger.debug( "Authentication settings changed after session suspended, recovering", ); - const recovered = await this.recoverDeployment(snapshot, "auth_config"); + const recovered = await this.recoverDeployment( + deployment, + "auth_config", + ); if (!recovered) { this.deploymentTelemetry.authConfigRecoveryFailed(); } @@ -424,7 +383,7 @@ export class DeploymentManager deployment: Deployment & { token?: string }, trigger: DeploymentRecoveryTrigger, ): Promise { - const recovered = await this.verifyAndApplyDeployment(deployment); + const recovered = await this.verifyAndApplySession(deployment); if (recovered) { this.deploymentTelemetry.recovered(trigger); } diff --git a/src/deployment/sessionStore.ts b/src/deployment/sessionStore.ts new file mode 100644 index 0000000000..f9141cdf83 --- /dev/null +++ b/src/deployment/sessionStore.ts @@ -0,0 +1,76 @@ +import * as vscode from "vscode"; + +import type { User } from "coder/site/src/api/typesGenerated"; + +import type { + WorkspaceSessionSnapshot, + WorkspaceSessionState, +} from "../workspace/session"; + +import type { Deployment } from "./types"; + +/** + * The deployment session: signed out (optionally keeping the last deployment + * for re-login) or signed in with an authenticated user. + * + * Every transition makes a new object, so callers can spot a change by + * comparing identity against an earlier value. + */ +export type SessionData = + | { readonly kind: "signedOut"; readonly deployment: Deployment | null } + | { + readonly kind: "signedIn"; + readonly deployment: Deployment; + readonly user: User; + }; + +/** + * Owns the deployment session. State changes only through signIn()/signOut(), + * each of which bumps the revision and notifies listeners. + * + * Consumers that only need auth status (like the workspace tree) take the lean + * WorkspaceSessionState projection instead of the full session. + */ +export class SessionStore implements WorkspaceSessionState { + #data: SessionData = { kind: "signedOut", deployment: null }; + #revision = 0; + readonly #onDidChange = new vscode.EventEmitter(); + + public readonly onDidChange = this.#onDidChange.event; + + /** The full current session, including deployment and user. */ + public get current(): SessionData { + return this.#data; + } + + /** The lean projection exposed to session-state consumers. */ + public getSnapshot(): WorkspaceSessionSnapshot { + if (this.#data.kind === "signedIn") { + return { + kind: "signedIn", + revision: this.#revision, + userId: this.#data.user.id, + }; + } + return { kind: "signedOut", revision: this.#revision }; + } + + public signIn(deployment: Deployment, user: User): SessionData { + return this.transition({ kind: "signedIn", deployment, user }); + } + + public signOut(deployment: Deployment | null): SessionData { + return this.transition({ kind: "signedOut", deployment }); + } + + private transition(data: SessionData): SessionData { + this.#data = data; + this.#revision++; + this.#onDidChange.fire(this.getSnapshot()); + return data; + } + + public dispose(): void { + this.#onDidChange.dispose(); + } +} diff --git a/src/extension.ts b/src/extension.ts index a6b44448eb..3de6031b10 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -188,44 +188,28 @@ async function doActivate( // createTreeView, unlike registerTreeDataProvider, gives us the tree view API // (so we can see when it is visible) but otherwise they have the same effect. - const myWsTree = vscode.window.createTreeView(MY_WORKSPACES_TREE_ID, { - treeDataProvider: myWorkspacesProvider, - }); - ctx.subscriptions.push(myWsTree); - myWorkspacesProvider.setVisibility(myWsTree.visible); - myWsTree.onDidChangeVisibility( - (event) => { - myWorkspacesProvider.setVisibility(event.visible); - }, - undefined, - ctx.subscriptions, - ); - - const sharedWsTree = vscode.window.createTreeView(SHARED_WORKSPACES_TREE_ID, { - treeDataProvider: sharedWorkspacesProvider, - }); - ctx.subscriptions.push(sharedWsTree); - sharedWorkspacesProvider.setVisibility(sharedWsTree.visible); - sharedWsTree.onDidChangeVisibility( - (event) => { - sharedWorkspacesProvider.setVisibility(event.visible); - }, - undefined, - ctx.subscriptions, - ); + const registerWorkspaceTreeView = ( + viewId: string, + provider: WorkspaceProvider, + ) => { + const tree = vscode.window.createTreeView(viewId, { + treeDataProvider: provider, + }); + ctx.subscriptions.push(tree); + provider.setVisibility(tree.visible); + tree.onDidChangeVisibility( + (event) => provider.setVisibility(event.visible), + undefined, + ctx.subscriptions, + ); + }; - const allWsTree = vscode.window.createTreeView(ALL_WORKSPACES_TREE_ID, { - treeDataProvider: allWorkspacesProvider, - }); - ctx.subscriptions.push(allWsTree); - allWorkspacesProvider.setVisibility(allWsTree.visible); - allWsTree.onDidChangeVisibility( - (event) => { - allWorkspacesProvider.setVisibility(event.visible); - }, - undefined, - ctx.subscriptions, + registerWorkspaceTreeView(MY_WORKSPACES_TREE_ID, myWorkspacesProvider); + registerWorkspaceTreeView( + SHARED_WORKSPACES_TREE_ID, + sharedWorkspacesProvider, ); + registerWorkspaceTreeView(ALL_WORKSPACES_TREE_ID, allWorkspacesProvider); // Register globally available commands. Many of these have visibility // controlled by contexts, see `when` in the package.json. @@ -427,7 +411,7 @@ async function doActivate( if (details) { ctx.subscriptions.push(details); - const deploymentSet = await deploymentManager.verifyAndApplyDeployment({ + const deploymentSet = await deploymentManager.verifyAndApplySession({ safeHostname: details.safeHostname, url: details.url, token: details.token, @@ -482,7 +466,7 @@ async function doActivate( output.info(`Initializing deployment: ${deployment.url}`); tracer .traceDeploymentInit(() => - deploymentManager.verifyAndApplyDeployment(deployment), + deploymentManager.verifyAndApplySession(deployment), ) .then((success) => { if (success) { diff --git a/src/workspace/workspacesProvider.ts b/src/workspace/workspacesProvider.ts index e0f7368ac8..03a5cb01dd 100644 --- a/src/workspace/workspacesProvider.ts +++ b/src/workspace/workspacesProvider.ts @@ -31,6 +31,31 @@ export enum WorkspaceQuery { Shared = "shared:true", } +/** Per-view rendering behavior, keyed by workspace query. */ +interface WorkspaceQueryConfig { + readonly showOwner: boolean; + readonly showMetadata: boolean; + readonly excludeOwn: boolean; +} + +const WORKSPACE_QUERY_CONFIG = { + [WorkspaceQuery.Mine]: { + showOwner: false, + showMetadata: true, + excludeOwn: false, + }, + [WorkspaceQuery.All]: { + showOwner: true, + showMetadata: false, + excludeOwn: false, + }, + [WorkspaceQuery.Shared]: { + showOwner: true, + showMetadata: false, + excludeOwn: true, + }, +} as const satisfies Record; + export interface WorkspaceProviderOptions { readonly refreshIntervalMs?: number; } @@ -40,8 +65,8 @@ export interface WorkspaceProviderOptions { * * Polling does not start until fetchAndRefresh() is called at least once. * - * If the poll fails or the client has no URL configured, clear the tree and - * abort polling until fetchAndRefresh() is called again. + * If a poll fails or the session is signed out, the tree is cleared and polling + * stops until the next fetchAndRefresh() or session change. */ export class WorkspaceProvider implements vscode.TreeDataProvider, vscode.Disposable @@ -53,9 +78,9 @@ export class WorkspaceProvider AgentMetadataWatcher >(); private readonly sessionChangeDisposable: vscode.Disposable; + private readonly config: WorkspaceQueryConfig; private timeout: NodeJS.Timeout | undefined; private fetching = false; - private refreshQueued = false; private visible = false; private disposed = false; @@ -66,107 +91,72 @@ export class WorkspaceProvider private readonly sessionState: WorkspaceSessionState, private readonly options: WorkspaceProviderOptions = {}, ) { + this.config = WORKSPACE_QUERY_CONFIG[getWorkspacesQuery]; this.sessionChangeDisposable = this.sessionState.onDidChange(() => { this.clear(); - this.requestRefresh(); - void this.runRefreshLoop(); + void this.fetchAndRefresh(); }); } - // fetchAndRefresh fetches new workspaces, re-renders the entire tree, then - // keeps refreshing (if a timer length was provided) as long as the user is - // still logged in and no errors were encountered fetching workspaces. - // Calling this while not visible is a no-op and will return immediately. - public async fetchAndRefresh() { - this.requestRefresh(); - await this.runRefreshLoop(); - } - - private requestRefresh(): void { - if (this.disposed || !this.visible) { - return; - } - this.refreshQueued = true; - } - - private async runRefreshLoop(): Promise { + // Fetch workspaces, render them, and queue the next poll. Does nothing when + // hidden, disposed, or already fetching. Never rejects, so it is safe as void. + public async fetchAndRefresh(): Promise { if (this.disposed || this.fetching || !this.visible) { return; } - this.fetching = true; + // A manual refresh may race a scheduled one, so drop any pending timer. this.cancelPendingRefresh(); - let shouldScheduleRefresh = false; + let hadError = false; try { - while (this.refreshQueued && !this.disposed && this.visible) { - this.refreshQueued = false; - shouldScheduleRefresh = false; - const session = this.sessionState.getSnapshot(); - - if (session.kind !== "signedIn") { - this.setWorkspaces([]); - continue; - } - - let hadError = false; - try { - const workspaces = await this.fetch(session); - if (workspaces && !this.disposed) { - this.setWorkspaces(workspaces); - } - } catch (error) { - this.logger.warn("Failed to fetch workspaces:", error); - hadError = true; - this.setWorkspaces([]); - } - - shouldScheduleRefresh = !hadError && !this.refreshQueued; - } + this.setWorkspaces(await this.fetch()); + } catch (error) { + this.logger.warn("Failed to fetch workspaces:", error); + hadError = true; + this.setWorkspaces([]); } finally { this.fetching = false; - if (this.refreshQueued && !this.disposed && this.visible) { - void this.runRefreshLoop(); - } else if (shouldScheduleRefresh && !this.disposed && this.visible) { - this.maybeScheduleRefresh(); - } + } + + // Keep polling only after a clean fetch while still signed in and visible. + if ( + !hadError && + !this.disposed && + this.visible && + this.sessionState.getSnapshot().kind === "signedIn" + ) { + this.maybeScheduleRefresh(); } } private setWorkspaces(workspaces: WorkspaceTreeItem[]): void { + if (this.disposed) { + return; + } this.workspaces = workspaces; - this.refresh(); + this.refreshTree(); } /** - * Fetch workspaces and turn them into tree items. Throw an error if not - * logged in or the query fails. + * Fetch workspaces and turn them into tree items. Returns an empty list when + * signed out, and throws if the query fails. */ - private async fetch( - session: Extract, - ): Promise { - // If there is no URL configured, assume we are logged out. - const url = this.client.getAxiosInstance().defaults.baseURL; - if (!url) { - throw new Error("not logged in"); + private async fetch(): Promise { + const session = this.sessionState.getSnapshot(); + if (session.kind !== "signedIn") { + return []; } const resp = await this.client.getWorkspaces({ q: this.getWorkspacesQuery, }); - const latestSession = this.sessionState.getSnapshot(); - const url2 = this.client.getAxiosInstance().defaults.baseURL; - if (!url2) { - throw new Error("not logged in"); - } - if ( - url !== url2 || - latestSession.kind !== "signedIn" || - latestSession.revision !== session.revision - ) { - this.refreshQueued = true; - return undefined; + // If the session changed while the request was in flight, this result is + // stale. Drop it and fetch again for the current session. + const latest = this.sessionState.getSnapshot(); + if (latest.kind !== "signedIn" || latest.revision !== session.revision) { + return this.fetch(); } const workspaces = this.filterWorkspaces(resp.workspaces, session); @@ -176,8 +166,7 @@ export class WorkspaceProvider // TODO: I think it might make more sense for the tree items to contain // their own watchers, rather than recreate the tree items every time and // have this separate map held outside the tree. - const showMetadata = this.getWorkspacesQuery === WorkspaceQuery.Mine; - if (showMetadata) { + if (this.config.showMetadata) { const agents = extractAllAgents(workspaces); for (const agent of agents) { // If we have an existing watcher, re-use it. @@ -190,7 +179,7 @@ export class WorkspaceProvider agent.id, this.client, ); - watcher.onChange(() => this.refresh()); + watcher.onChange(() => this.refreshTree()); this.agentWatchers.set(agent.id, watcher); } } @@ -204,12 +193,14 @@ export class WorkspaceProvider } } - const showOwner = this.getWorkspacesQuery !== WorkspaceQuery.Mine; - // Create tree items for each workspace return workspaces.map( (workspace: Workspace) => - new WorkspaceTreeItem(workspace, showOwner, showMetadata), + new WorkspaceTreeItem( + workspace, + this.config.showOwner, + this.config.showMetadata, + ), ); } @@ -217,9 +208,11 @@ export class WorkspaceProvider workspaces: readonly Workspace[], session: Extract, ): readonly Workspace[] { - if (this.getWorkspacesQuery !== WorkspaceQuery.Shared) { + if (!this.config.excludeOwn) { return workspaces; } + // `shared:true` also returns workspaces we own and shared out; drop them + // to leave only those shared with us. return workspaces.filter( (workspace) => workspace.owner_id !== session.userId, ); @@ -237,13 +230,10 @@ export class WorkspaceProvider this.visible = visible; if (!visible) { this.cancelPendingRefresh(); - } else if (this.refreshQueued) { - void this.runRefreshLoop(); } else if (this.workspaces) { this.maybeScheduleRefresh(); } else { - this.requestRefresh(); - void this.runRefreshLoop(); + void this.fetchAndRefresh(); } } @@ -261,8 +251,7 @@ export class WorkspaceProvider private maybeScheduleRefresh() { if (this.options.refreshIntervalMs && !this.timeout) { this.timeout = setTimeout(() => { - this.requestRefresh(); - void this.runRefreshLoop(); + void this.fetchAndRefresh(); }, this.options.refreshIntervalMs); } } @@ -274,8 +263,8 @@ export class WorkspaceProvider vscode.TreeItem | undefined | null | void > = this._onDidChangeTreeData.event; - // refresh causes the tree to re-render. It does not fetch fresh workspaces. - public refresh(item?: vscode.TreeItem): void { + // Re-render the tree from the current workspaces. Does not fetch. + public refreshTree(item?: vscode.TreeItem): void { if (this.disposed) { return; } @@ -365,7 +354,7 @@ export class WorkspaceProvider */ public clear(): void { this.clearState(); - this.refresh(); + this.refreshTree(); } private clearState(): void { @@ -375,7 +364,6 @@ export class WorkspaceProvider } this.agentWatchers.clear(); this.workspaces = undefined; - this.refreshQueued = false; } public dispose() { diff --git a/test/mocks/testHelpers.ts b/test/mocks/testHelpers.ts index f715d9b4a6..a7725cdfaa 100644 --- a/test/mocks/testHelpers.ts +++ b/test/mocks/testHelpers.ts @@ -12,10 +12,15 @@ import * as vscode from "vscode"; import { createTestTelemetryService } from "./telemetry"; import { window as vscodeWindow } from "./vscode.runtime"; -import type { Experiment, User } from "coder/site/src/api/typesGenerated"; +import type { + Experiment, + User, + Workspace, +} from "coder/site/src/api/typesGenerated"; import type { WebSocketEventType } from "coder/site/src/utils/OneWayWebSocket"; import type { IncomingMessage } from "node:http"; +import type { AgentMetadataEvent } from "@/api/api-helper"; import type { CoderApi } from "@/api/coderApi"; import type { CliCredentialManager } from "@/core/cliCredentialManager"; import type { ServiceContainer } from "@/core/container"; @@ -31,6 +36,10 @@ import type { ParsedMessageEvent, UnidirectionalStream, } from "@/websocket/eventStreamConnection"; +import type { + WorkspaceSessionSnapshot, + WorkspaceSessionState, +} from "@/workspace/session"; /** * Subset of `ContextManager`'s public API that mocks (e.g. `MockContextManager`) @@ -461,6 +470,21 @@ export function createMockLogger(): Logger { }; } +/** Resolve once pending microtasks and the macrotask queue have drained. */ +export async function flush(): Promise { + await new Promise((resolve) => setImmediate(resolve)); +} + +/** + * Drain only the microtask queue. Use instead of flush() under fake timers, + * which leave the macrotask queue (setImmediate) untouched. + */ +export async function flushPromises(): Promise { + await Promise.resolve(); + await Promise.resolve(); + await Promise.resolve(); +} + /** * Backdate a file's atime/mtime by `daysAgo` days so tests can exercise the * support-bundle age cutoff without waiting. @@ -1184,3 +1208,88 @@ export class MockTerminalOutputChannel { this.write.mockClear(); } } + +/** Default user id reported by MockWorkspaceSessionState's signed-in snapshot. */ +export const TEST_CURRENT_USER_ID = "current-user"; + +/** In-memory WorkspaceSessionState double with sign-in/out controls. */ +export class MockWorkspaceSessionState implements WorkspaceSessionState { + private revision = 0; + private readonly emitter = + new vscode.EventEmitter(); + private snapshot: WorkspaceSessionSnapshot = { + kind: "signedIn", + revision: 0, + userId: TEST_CURRENT_USER_ID, + }; + + readonly onDidChange = this.emitter.event; + + getSnapshot(): WorkspaceSessionSnapshot { + return this.snapshot; + } + + signIn(userId = TEST_CURRENT_USER_ID): void { + this.revision += 1; + this.snapshot = { kind: "signedIn", revision: this.revision, userId }; + this.emitter.fire(this.snapshot); + } + + signOut(): void { + this.revision += 1; + this.snapshot = { kind: "signedOut", revision: this.revision }; + this.emitter.fire(this.snapshot); + } +} + +interface WorkspacesResponse { + workspaces: readonly Workspace[]; + count: number; +} + +/** + * Stands in for CoderApi at the boundaries WorkspaceProvider touches: the + * workspaces query and the per-agent metadata socket. Program responses with + * respondOnce()/pending(); metadata flows through a real MockEventStream so + * tests drive the production watcher rather than a stub. + */ +export class MockWorkspacesClient { + baseURL: string | undefined = "https://coder.example.com"; + readonly metadataStreams = new Map< + string, + MockEventStream<{ data: AgentMetadataEvent[] }> + >(); + + readonly getWorkspaces = vi.fn( + (_req: { q: string }): Promise => + Promise.resolve({ workspaces: [], count: 0 }), + ); + + watchAgentMetadata(agentId: string) { + const stream = new MockEventStream<{ data: AgentMetadataEvent[] }>(); + this.metadataStreams.set(agentId, stream); + return Promise.resolve(stream); + } + + getAxiosInstance() { + return { defaults: { baseURL: this.baseURL } }; + } + + /** Resolve the next getWorkspaces call with these workspaces. */ + respondOnce(workspaces: readonly Workspace[]): void { + this.getWorkspaces.mockResolvedValueOnce({ + workspaces, + count: workspaces.length, + }); + } + + /** Make the next getWorkspaces call hang until the returned resolve() runs. */ + pending(): { resolve: (workspaces: readonly Workspace[]) => void } { + const { promise, resolve } = Promise.withResolvers(); + this.getWorkspaces.mockReturnValueOnce(promise); + return { + resolve: (workspaces) => + resolve({ workspaces, count: workspaces.length }), + }; + } +} diff --git a/test/unit/deployment/deploymentManager.test.ts b/test/unit/deployment/deploymentManager.test.ts index 97c35cd15b..96595af335 100644 --- a/test/unit/deployment/deploymentManager.test.ts +++ b/test/unit/deployment/deploymentManager.test.ts @@ -11,6 +11,7 @@ import { createMockLogger, createMockServiceContainer, createMockUser, + flush, InMemoryMemento, InMemorySecretStorage, MockCoderApi, @@ -37,19 +38,9 @@ const TEST_URL = "https://coder.example.com"; const TEST_HOSTNAME = "coder.example.com"; const managers: DeploymentManager[] = []; -function deferred(): { - promise: Promise; - resolve: (value: T | PromiseLike) => void; -} { - let resolve!: (value: T | PromiseLike) => void; - const promise = new Promise((res) => { - resolve = res; - }); - return { promise, resolve }; -} - -async function flush(): Promise { - await new Promise((resolve) => setImmediate(resolve)); +function currentUserId(manager: DeploymentManager): string | undefined { + const snapshot = manager.getSnapshot(); + return snapshot.kind === "signedIn" ? snapshot.userId : undefined; } afterEach(() => { @@ -67,7 +58,7 @@ function createTestContext() { new MockConfigurationProvider(); const mockClient = new MockCoderApi(); - // For verifyAndApplyDeployment, we use a separate mock for validation + // For verifyAndApplySession, we use a separate mock for validation const validationMockClient = new MockCoderApi(); const mockOAuthSessionManager = new MockOAuthSessionManager(); const secretStorage = new InMemorySecretStorage(); @@ -118,7 +109,7 @@ describe("DeploymentManager", () => { const { manager } = createTestContext(); expect(manager.getCurrentDeployment()).toBeNull(); - expect(manager.getCurrentUserId()).toBeUndefined(); + expect(currentUserId(manager)).toBeUndefined(); expect(manager.isAuthenticated()).toBe(false); }); @@ -137,7 +128,7 @@ describe("DeploymentManager", () => { url: TEST_URL, safeHostname: TEST_HOSTNAME, }); - expect(manager.getCurrentUserId()).toBe("current-user"); + expect(currentUserId(manager)).toBe("current-user"); expect(manager.isAuthenticated()).toBe(true); }); @@ -155,7 +146,7 @@ describe("DeploymentManager", () => { await manager.clearDeployment("credentials_removed"); expect(manager.getCurrentDeployment()).toBeNull(); - expect(manager.getCurrentUserId()).toBeUndefined(); + expect(currentUserId(manager)).toBeUndefined(); expect(manager.isAuthenticated()).toBe(false); }); }); @@ -185,7 +176,7 @@ describe("DeploymentManager", () => { it("does not persist a deployment after a newer state wins", async () => { const { mockOAuthSessionManager, secretsManager, manager } = createTestContext(); - const oauthSetDeployment = deferred(); + const oauthSetDeployment = Promise.withResolvers(); mockOAuthSessionManager.setDeployment.mockReturnValueOnce( oauthSetDeployment.promise, ); @@ -235,13 +226,13 @@ describe("DeploymentManager", () => { }); }); - describe("verifyAndApplyDeployment", () => { + describe("verifyAndApplySession", () => { it("returns true and sets deployment on auth success", async () => { const { mockClient, validationMockClient, manager } = createTestContext(); const user = createMockUser(); validationMockClient.setAuthenticatedUserResponse(user); - const result = await manager.verifyAndApplyDeployment({ + const result = await manager.verifyAndApplySession({ url: TEST_URL, safeHostname: TEST_HOSTNAME, token: "test-token", @@ -250,7 +241,7 @@ describe("DeploymentManager", () => { expect(result).toBe(true); expect(mockClient.host).toBe(TEST_URL); expect(mockClient.token).toBe("test-token"); - expect(manager.getCurrentUserId()).toBe(user.id); + expect(currentUserId(manager)).toBe(user.id); expect(manager.isAuthenticated()).toBe(true); }); @@ -260,7 +251,7 @@ describe("DeploymentManager", () => { new Error("Auth failed"), ); - const result = await manager.verifyAndApplyDeployment({ + const result = await manager.verifyAndApplySession({ url: TEST_URL, safeHostname: TEST_HOSTNAME, token: "test-token", @@ -268,7 +259,7 @@ describe("DeploymentManager", () => { expect(result).toBe(false); expect(manager.getCurrentDeployment()).toBeNull(); - expect(manager.getCurrentUserId()).toBeUndefined(); + expect(currentUserId(manager)).toBeUndefined(); expect(manager.isAuthenticated()).toBe(false); }); @@ -277,7 +268,7 @@ describe("DeploymentManager", () => { const user = createMockUser(); validationMockClient.setAuthenticatedUserResponse(user); - const result = await manager.verifyAndApplyDeployment({ + const result = await manager.verifyAndApplySession({ url: TEST_URL, safeHostname: TEST_HOSTNAME, token: "", @@ -301,7 +292,7 @@ describe("DeploymentManager", () => { token: "stored-token", }); - const result = await manager.verifyAndApplyDeployment({ + const result = await manager.verifyAndApplySession({ url: TEST_URL, safeHostname: TEST_HOSTNAME, }); @@ -315,7 +306,7 @@ describe("DeploymentManager", () => { const user = createMockUser(); validationMockClient.setAuthenticatedUserResponse(user); - await manager.verifyAndApplyDeployment({ + await manager.verifyAndApplySession({ url: TEST_URL, safeHostname: TEST_HOSTNAME, token: "test-token", @@ -328,11 +319,11 @@ describe("DeploymentManager", () => { const { mockClient, validationMockClient, manager } = createTestContext(); const remoteUser = createMockUser({ id: "remote-user" }); const localUser = createMockUser({ id: "local-user" }); - const validation = deferred(); + const validation = Promise.withResolvers(); validationMockClient.getAuthenticatedUser.mockReturnValueOnce( validation.promise, ); - const remoteLogin = manager.verifyAndApplyDeployment({ + const remoteLogin = manager.verifyAndApplySession({ url: TEST_URL, safeHostname: TEST_HOSTNAME, token: "remote-token", @@ -350,7 +341,7 @@ describe("DeploymentManager", () => { expect(await remoteLogin).toBe(false); expect(mockClient.host).toBe("https://local.example.com"); expect(mockClient.token).toBe("local-token"); - expect(manager.getCurrentUserId()).toBe("local-user"); + expect(currentUserId(manager)).toBe("local-user"); }); }); @@ -455,7 +446,7 @@ describe("DeploymentManager", () => { }); expect(mockClient.token).toBe("initial-token"); - expect(manager.getCurrentUserId()).toBe("initial-user"); + expect(currentUserId(manager)).toBe("initial-user"); expect(manager.isAuthenticated()).toBe(true); await secretsManager.setSessionAuth(TEST_HOSTNAME, { @@ -465,7 +456,7 @@ describe("DeploymentManager", () => { await flush(); expect(mockClient.token).toBe("refreshed-token"); - expect(manager.getCurrentUserId()).toBe("refreshed-user"); + expect(currentUserId(manager)).toBe("refreshed-user"); expect(manager.isAuthenticated()).toBe(true); }); @@ -474,7 +465,7 @@ describe("DeploymentManager", () => { createTestContext(); const user = createMockUser({ id: "initial-user" }); const refreshedUser = createMockUser({ id: "refreshed-user" }); - const validation = deferred(); + const validation = Promise.withResolvers(); validationMockClient.getAuthenticatedUser.mockReturnValueOnce( validation.promise, ); @@ -497,7 +488,7 @@ describe("DeploymentManager", () => { expect(mockClient.host).toBeUndefined(); expect(mockClient.token).toBeUndefined(); expect(manager.getCurrentDeployment()).toBeNull(); - expect(manager.getCurrentUserId()).toBeUndefined(); + expect(currentUserId(manager)).toBeUndefined(); }); }); @@ -519,7 +510,7 @@ describe("DeploymentManager", () => { expect(mockClient.token).toBeUndefined(); expect(contextManager.get("coder.authenticated")).toBe(false); expect(contextManager.get("coder.isOwner")).toBe(false); - expect(manager.getCurrentUserId()).toBeUndefined(); + expect(currentUserId(manager)).toBeUndefined(); }); it("resets the telemetry deployment URL on clearDeployment", async () => { @@ -575,7 +566,7 @@ describe("DeploymentManager", () => { expect(mockClient.host).toBeUndefined(); expect(mockClient.token).toBeUndefined(); expect(contextManager.get("coder.authenticated")).toBe(false); - expect(manager.getCurrentUserId()).toBeUndefined(); + expect(currentUserId(manager)).toBeUndefined(); expect(manager.isAuthenticated()).toBe(false); // Deployment is retained for easy re-login @@ -621,7 +612,7 @@ describe("DeploymentManager", () => { expect(mockClient.host).toBe(TEST_URL); expect(mockClient.token).toBe(""); expect(manager.isAuthenticated()).toBe(true); - expect(manager.getCurrentUserId()).toBe(user.id); + expect(currentUserId(manager)).toBe(user.id); expect(validationMockClient.getAuthenticatedUser).toHaveBeenCalledTimes( 1, ); @@ -663,7 +654,7 @@ describe("DeploymentManager", () => { await vi.runAllTimersAsync(); expect(manager.getCurrentDeployment()).toBeNull(); - expect(manager.getCurrentUserId()).toBeUndefined(); + expect(currentUserId(manager)).toBeUndefined(); expect(manager.isAuthenticated()).toBe(false); } finally { vi.useRealTimers(); @@ -703,7 +694,7 @@ describe("DeploymentManager", () => { // Should recover and be authenticated again expect(mockClient.token).toBe("recovered-token"); - expect(manager.getCurrentUserId()).toBe(user.id); + expect(currentUserId(manager)).toBe(user.id); expect(manager.isAuthenticated()).toBe(true); expect(telemetrySink.expectOne("deployment.recovered")).toMatchObject({ properties: { trigger: "token_update" }, diff --git a/test/unit/workspace/workspacesProvider.test.ts b/test/unit/workspace/workspacesProvider.test.ts index 15e9e8aeef..87b053b8f0 100644 --- a/test/unit/workspace/workspacesProvider.test.ts +++ b/test/unit/workspace/workspacesProvider.test.ts @@ -1,4 +1,4 @@ -import { beforeEach, describe, expect, it, vi } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import * as vscode from "vscode"; import { @@ -8,13 +8,16 @@ import { type WorkspaceTreeItem, } from "@/workspace/workspacesProvider"; -import { - agent as createAgent, - resource as createResource, - workspace as createWorkspace, -} from "@repo/mocks"; +import { agent, resource, workspace } from "@repo/mocks"; -import { createMockLogger } from "../../mocks/testHelpers"; +import { + createMockLogger, + flush, + flushPromises, + MockWorkspaceSessionState, + MockWorkspacesClient, + TEST_CURRENT_USER_ID, +} from "../../mocks/testHelpers"; import type { Workspace, @@ -23,232 +26,106 @@ import type { WorkspaceAppStatus, } from "coder/site/src/api/typesGenerated"; -import type { AgentMetadataWatcher } from "@/api/agentMetadataHelper"; import type { AgentMetadataEvent } from "@/api/api-helper"; import type { CoderApi } from "@/api/coderApi"; -import type { WorkspaceSessionSnapshot } from "@/workspace/session"; - -vi.mock("@/api/agentMetadataHelper", async (importOriginal) => { - const original = - await importOriginal(); - return { - ...original, - createAgentMetadataWatcher: vi.fn(), - }; -}); -const { createAgentMetadataWatcher } = - await import("@/api/agentMetadataHelper"); - -const baseUrl = "https://coder.example.com"; -const currentUserId = "current-user"; - -interface WorkspacesResponse { - workspaces: readonly Workspace[]; - count: number; -} - -class TestClient { - baseURL: string | undefined = baseUrl; - readonly getWorkspaces = vi.fn( - (_req: { q: string }): Promise => - Promise.resolve({ workspaces: [], count: 0 }), - ); - - getAxiosInstance() { - return { - defaults: { - baseURL: this.baseURL, - }, - }; - } -} - -class TestSessionState { - private revision = 0; - private readonly onDidChangeEmitter = - new vscode.EventEmitter(); - readonly onDidChange = this.onDidChangeEmitter.event; - private snapshot: WorkspaceSessionSnapshot = { - kind: "signedIn", - revision: this.revision, - userId: currentUserId, - }; - - getSnapshot(): WorkspaceSessionSnapshot { - return this.snapshot; - } - - signIn(userId = currentUserId): void { - this.revision += 1; - this.snapshot = { kind: "signedIn", revision: this.revision, userId }; - this.onDidChangeEmitter.fire(this.snapshot); - } - - signOut(): void { - this.revision += 1; - this.snapshot = { kind: "signedOut", revision: this.revision }; - this.onDidChangeEmitter.fire(this.snapshot); - } -} - -type TestWatcher = AgentMetadataWatcher & { - onChangeEmitter: vscode.EventEmitter; - dispose: ReturnType void>>; -}; - -describe("WorkspaceProvider", () => { - let client: TestClient; - let logger: ReturnType; - let session: TestSessionState; - let watchers: TestWatcher[]; - - beforeEach(() => { - client = new TestClient(); - logger = createMockLogger(); - session = new TestSessionState(); - watchers = []; - vi.mocked(createAgentMetadataWatcher).mockImplementation(() => { - const onChangeEmitter = new vscode.EventEmitter(); - const watcher: TestWatcher = { - onChangeEmitter, - onChange: onChangeEmitter.event, - dispose: vi.fn<() => void>(), - }; - watchers.push(watcher); - return Promise.resolve(watcher); - }); - }); - - function makeProvider( +function setup() { + const logger = createMockLogger(); + const client = new MockWorkspacesClient(); + const session = new MockWorkspaceSessionState(); + const makeProvider = ( query: WorkspaceQuery, options?: { refreshIntervalMs?: number }, - ): WorkspaceProvider { - return new WorkspaceProvider( + ): WorkspaceProvider => + new WorkspaceProvider( query, client as unknown as CoderApi, logger, session, options, ); - } - - function workspace( - overrides: Parameters[0] = {}, - ): Workspace { - return createWorkspace(overrides); - } - - function agent(overrides: Partial = {}): WorkspaceAgent { - return createAgent(overrides); - } - - function workspaceWithAgents( - workspaceOverrides: Parameters[0] = {}, - agents: WorkspaceAgent[], - ): Workspace { - return workspace({ - ...workspaceOverrides, - latest_build: { - ...workspaceOverrides.latest_build, - resources: [createResource({ agents })], - }, - }); - } - - function appStatus( - overrides: Partial = {}, - ): WorkspaceAppStatus { - return { - id: "status-1", - created_at: "2024-01-01T00:00:00Z", - workspace_id: "workspace-1", - agent_id: "agent-1", - app_id: "app-1", - state: "working", - message: "Opening pull request", - uri: "https://example.com/pr/1", - icon: "", - needs_user_attention: false, - ...overrides, - }; - } - - function app(overrides: Partial = {}): WorkspaceApp { - return { - id: "app-1", - external: false, - slug: "app", - subdomain: false, - sharing_level: "owner", - health: "healthy", - hidden: false, - open_in: "tab", - statuses: [], - ...overrides, - }; - } - - function metadata( - overrides: Partial = {}, - ): AgentMetadataEvent { - return { - result: { - collected_at: "2024-01-01T00:00:00Z", - age: 0, - value: "42", - error: "", - }, - description: { - display_name: "CPU", - key: "cpu", - script: "cpu.sh", - interval: 5, - timeout: 1, - }, - ...overrides, - }; - } - - function respondOnce(workspaces: readonly Workspace[]): void { - client.getWorkspaces.mockResolvedValueOnce({ - workspaces, - count: workspaces.length, - }); - } - - function pendingWorkspaces(): { - resolve: (workspaces: readonly Workspace[]) => void; - } { - let resolve!: (workspaces: readonly Workspace[]) => void; - client.getWorkspaces.mockReturnValueOnce( - new Promise((res) => { - resolve = (workspaces) => res({ workspaces, count: workspaces.length }); - }), - ); - return { resolve }; - } + return { logger, client, session, makeProvider }; +} - async function flush(): Promise { - await new Promise((resolve) => setImmediate(resolve)); - } +function workspaceWithAgents( + workspaceOverrides: Parameters[0] = {}, + agents: WorkspaceAgent[], +): Workspace { + return workspace({ + ...workspaceOverrides, + latest_build: { + ...workspaceOverrides.latest_build, + resources: [resource({ agents })], + }, + }); +} - async function flushPromises(): Promise { - await Promise.resolve(); - await Promise.resolve(); - await Promise.resolve(); - } +function appStatus( + overrides: Partial = {}, +): WorkspaceAppStatus { + return { + id: "status-1", + created_at: "2024-01-01T00:00:00Z", + workspace_id: "workspace-1", + agent_id: "agent-1", + app_id: "app-1", + state: "working", + message: "Opening pull request", + uri: "https://example.com/pr/1", + icon: "", + needs_user_attention: false, + ...overrides, + }; +} - async function show(provider: WorkspaceProvider): Promise { - provider.setVisibility(true); - await flush(); - } +function app(overrides: Partial = {}): WorkspaceApp { + return { + id: "app-1", + external: false, + slug: "app", + subdomain: false, + sharing_level: "owner", + health: "healthy", + hidden: false, + open_in: "tab", + statuses: [], + ...overrides, + }; +} - async function labels(provider: WorkspaceProvider): Promise { - return (await provider.getChildren()).map((item) => item.label); - } +function metadata( + overrides: Partial = {}, +): AgentMetadataEvent { + return { + result: { + collected_at: "2024-01-01T00:00:00Z", + age: 0, + value: "42", + error: "", + }, + description: { + display_name: "CPU", + key: "cpu", + script: "cpu.sh", + interval: 5, + timeout: 1, + }, + ...overrides, + }; +} + +async function show(provider: WorkspaceProvider): Promise { + provider.setVisibility(true); + await flush(); +} +async function labels(provider: WorkspaceProvider): Promise { + return (await provider.getChildren()).map((item) => item.label); +} + +describe("WorkspaceProvider", () => { it("does not fetch while signed out", async () => { + const { client, session, makeProvider } = setup(); session.signOut(); const provider = makeProvider(WorkspaceQuery.Mine); @@ -263,6 +140,7 @@ describe("WorkspaceProvider", () => { [WorkspaceQuery.Shared, "shared:true"], [WorkspaceQuery.All, ""], ])("fetches %s with the expected query", async (query, expectedQuery) => { + const { client, makeProvider } = setup(); const provider = makeProvider(query); await show(provider); @@ -289,7 +167,8 @@ describe("WorkspaceProvider", () => { ])( "renders top-level workspace items for $query", async ({ query, label, collapsibleState }) => { - respondOnce([ + const { client, makeProvider } = setup(); + client.respondOnce([ workspace({ id: "workspace-1", name: "dev", @@ -310,11 +189,12 @@ describe("WorkspaceProvider", () => { ); it("filters current-user-owned workspaces from shared results", async () => { - respondOnce([ + const { client, makeProvider } = setup(); + client.respondOnce([ workspace({ id: "owned-shared-out", name: "owned", - owner_id: currentUserId, + owner_id: TEST_CURRENT_USER_ID, owner_name: "current", }), workspace({ @@ -334,11 +214,12 @@ describe("WorkspaceProvider", () => { it.each([WorkspaceQuery.Mine, WorkspaceQuery.All])( "does not apply shared ownership filtering to %s", async (query) => { - respondOnce([ + const { client, makeProvider } = setup(); + client.respondOnce([ workspace({ id: "owned", name: "owned", - owner_id: currentUserId, + owner_id: TEST_CURRENT_USER_ID, owner_name: "current", }), ]); @@ -351,7 +232,8 @@ describe("WorkspaceProvider", () => { ); it("clears rendered workspaces when the session signs out", async () => { - respondOnce([workspace({ name: "dev" })]); + const { client, session, makeProvider } = setup(); + client.respondOnce([workspace({ name: "dev" })]); const provider = makeProvider(WorkspaceQuery.Mine); await show(provider); @@ -364,7 +246,8 @@ describe("WorkspaceProvider", () => { }); it("does not render a pending response after sign-out", async () => { - const pending = pendingWorkspaces(); + const { client, session, makeProvider } = setup(); + const pending = client.pending(); const provider = makeProvider(WorkspaceQuery.Shared); provider.setVisibility(true); @@ -376,40 +259,28 @@ describe("WorkspaceProvider", () => { expect(await provider.getChildren()).toEqual([]); }); - it("refetches when the session changes while a request is pending", async () => { - const pending = pendingWorkspaces(); - client.getWorkspaces.mockResolvedValueOnce({ - workspaces: [ - workspace({ - id: "fresh-workspace", - owner_id: "alice-id", - owner_name: "alice", - name: "fresh", - }), - ], - count: 1, - }); + it("renders fresh results when the session changes mid-request", async () => { + const { client, session, makeProvider } = setup(); + const pending = client.pending(); + client.respondOnce([ + workspace({ owner_id: "alice-id", owner_name: "alice", name: "fresh" }), + ]); const provider = makeProvider(WorkspaceQuery.Shared); provider.setVisibility(true); await flush(); session.signIn("second-user"); pending.resolve([ - workspace({ - id: "stale-workspace", - owner_id: "bob-id", - owner_name: "bob", - name: "stale", - }), + workspace({ owner_id: "bob-id", owner_name: "bob", name: "stale" }), ]); await flush(); await flush(); - expect(client.getWorkspaces).toHaveBeenCalledTimes(2); expect(await labels(provider)).toEqual(["alice / fresh"]); }); it("does not fetch while hidden", async () => { + const { client, makeProvider } = setup(); const provider = makeProvider(WorkspaceQuery.Mine); await provider.fetchAndRefresh(); @@ -418,7 +289,8 @@ describe("WorkspaceProvider", () => { }); it("renders a response that completes after the tree is hidden", async () => { - const pending = pendingWorkspaces(); + const { client, makeProvider } = setup(); + const pending = client.pending(); const provider = makeProvider(WorkspaceQuery.Mine); provider.setVisibility(true); @@ -430,9 +302,10 @@ describe("WorkspaceProvider", () => { expect(await labels(provider)).toEqual(["dev"]); }); - it("fetches queued session changes when the tree is shown again", async () => { - const pending = pendingWorkspaces(); - respondOnce([workspace({ name: "fresh", owner_id: "alice-id" })]); + it("renders fresh results for a session change queued while hidden", async () => { + const { client, session, makeProvider } = setup(); + const pending = client.pending(); + client.respondOnce([workspace({ name: "fresh" })]); const provider = makeProvider(WorkspaceQuery.Mine); provider.setVisibility(true); @@ -445,31 +318,29 @@ describe("WorkspaceProvider", () => { provider.setVisibility(true); await flush(); - expect(client.getWorkspaces).toHaveBeenCalledTimes(2); expect(await labels(provider)).toEqual(["fresh"]); }); - it("clears and logs when fetch fails", async () => { - respondOnce([workspace({ name: "dev" })]); + it("clears the tree when a fetch fails", async () => { + const { client, makeProvider } = setup(); + client.respondOnce([workspace({ name: "dev" })]); client.getWorkspaces.mockRejectedValueOnce(new Error("network down")); const provider = makeProvider(WorkspaceQuery.Mine); await show(provider); expect(await labels(provider)).toEqual(["dev"]); + await provider.fetchAndRefresh(); expect(await provider.getChildren()).toEqual([]); - expect(logger.warn).toHaveBeenCalledWith( - "Failed to fetch workspaces:", - expect.any(Error), - ); }); - it("schedules polling after a successful refresh", async () => { + it("refreshes content on the polling interval", async () => { vi.useFakeTimers(); try { - respondOnce([workspace({ name: "first" })]); - respondOnce([workspace({ name: "second" })]); + const { client, makeProvider } = setup(); + client.respondOnce([workspace({ name: "first" })]); + client.respondOnce([workspace({ name: "second" })]); const provider = makeProvider(WorkspaceQuery.Mine, { refreshIntervalMs: 5_000, }); @@ -481,7 +352,6 @@ describe("WorkspaceProvider", () => { await vi.advanceTimersByTimeAsync(5_000); await flushPromises(); - expect(client.getWorkspaces).toHaveBeenCalledTimes(2); expect(await labels(provider)).toEqual(["second"]); } finally { vi.useRealTimers(); @@ -489,7 +359,8 @@ describe("WorkspaceProvider", () => { }); it("renders workspace child agents", async () => { - respondOnce([ + const { client, makeProvider } = setup(); + client.respondOnce([ workspaceWithAgents({ name: "dev" }, [ agent({ id: "agent-1", name: "main", status: "connected" }), agent({ id: "agent-2", name: "sidecar", status: "disconnected" }), @@ -512,7 +383,8 @@ describe("WorkspaceProvider", () => { }); it("renders app status children for an agent", async () => { - respondOnce([ + const { client, makeProvider } = setup(); + client.respondOnce([ workspaceWithAgents({ name: "dev" }, [ agent({ id: "agent-1", @@ -549,8 +421,9 @@ describe("WorkspaceProvider", () => { }); }); - it("renders metadata and metadata errors for watched agents", async () => { - respondOnce([ + it("renders agent metadata and surfaces metadata errors", async () => { + const { client, makeProvider } = setup(); + client.respondOnce([ workspaceWithAgents({ name: "dev" }, [ agent({ id: "agent-1", name: "main" }), ]), @@ -563,68 +436,31 @@ describe("WorkspaceProvider", () => { const [agentItem] = (await provider.getChildren( workspaceItem, )) as AgentTreeItem[]; + const stream = client.metadataStreams.get("agent-1")!; - watchers[0].metadata = [metadata()]; - let [section] = await provider.getChildren(agentItem); - const metadataItems = await provider.getChildren(section); - expect(section?.label).toBe("Agent Metadata"); + stream.pushMessage({ data: [metadata()] }); + const [metadataSection] = await provider.getChildren(agentItem); + const metadataItems = await provider.getChildren(metadataSection); + expect(metadataSection?.label).toBe("Agent Metadata"); expect(metadataItems[0]?.label).toBe("CPU: 42"); - watchers[0].error = new Error("boom"); - [section] = await provider.getChildren(agentItem); - expect(section?.label).toBe("Failed to query metadata: boom"); + stream.pushError(new Error("boom")); + const [errorSection] = await provider.getChildren(agentItem); + expect(errorSection?.label).toBe("Failed to query metadata: boom"); }); - it("reuses and disposes metadata watchers as agents change", async () => { - respondOnce([ + it("empties the tree when cleared", async () => { + const { client, makeProvider } = setup(); + client.respondOnce([ workspaceWithAgents({ name: "dev" }, [agent({ id: "agent-1" })]), ]); - respondOnce([ - workspaceWithAgents({ name: "dev" }, [agent({ id: "agent-2" })]), - ]); const provider = makeProvider(WorkspaceQuery.Mine); await show(provider); - expect(vi.mocked(createAgentMetadataWatcher)).toHaveBeenCalledWith( - "agent-1", - client, - ); - - await provider.fetchAndRefresh(); - await flush(); - - expect(watchers[0].dispose).toHaveBeenCalled(); - expect(vi.mocked(createAgentMetadataWatcher)).toHaveBeenCalledWith( - "agent-2", - client, - ); - }); - - it("clear removes workspaces and disposes metadata watchers", async () => { - respondOnce([ - workspaceWithAgents({ name: "dev" }, [agent({ id: "agent-1" })]), - ]); - const provider = makeProvider(WorkspaceQuery.Mine); + expect(await labels(provider)).toEqual(["dev"]); - await show(provider); provider.clear(); expect(await provider.getChildren()).toEqual([]); - expect(watchers[0].dispose).toHaveBeenCalled(); - }); - - it("dispose removes metadata watchers without firing a tree refresh", async () => { - respondOnce([ - workspaceWithAgents({ name: "dev" }, [agent({ id: "agent-1" })]), - ]); - const provider = makeProvider(WorkspaceQuery.Mine); - const onDidChangeTreeData = vi.fn(); - - await show(provider); - provider.onDidChangeTreeData(onDidChangeTreeData); - provider.dispose(); - - expect(watchers[0].dispose).toHaveBeenCalled(); - expect(onDidChangeTreeData).not.toHaveBeenCalled(); }); }); From e54f977160fc21f669d413f0f3e4ac252d237b35 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Mon, 8 Jun 2026 18:42:01 +0300 Subject: [PATCH 4/5] fix: address shared workspaces review findings verifyAndUpdateSession keeps verify-before-apply: it verifies the rotated token before touching the live client, rotates credentials in place when the user is unchanged (no revision bump or tree rebuild), and recovers a session that was suspended mid-verify instead of getting stuck signed out (CRF-1, CRF-2). Also: - bound fetch() retries with a per-attempt disposed check, and dispose a metadata watcher created while dispose() races the fetch (CRF-3, CRF-8) - dispose the tree-data EventEmitter on dispose() (CRF-15) - rename clearSideEffects to clearAuthState (CRF-12) - document the session snapshot revision contract (CRF-13) - drop the dead getAxiosInstance mock helper and restore an unrelated blank line (CRF-14, CRF-16) - trim comments that restated the code (CRF-9) - add tests for dispose lifecycle, verify failure, signed-out startup, and stale-fetch recovery (CRF-4 through CRF-7) --- src/deployment/deploymentManager.ts | 34 +++++- src/deployment/sessionStore.ts | 4 +- src/extension.ts | 1 + src/workspace/session.ts | 7 ++ src/workspace/workspacesProvider.ts | 110 ++++++++++-------- test/mocks/testHelpers.ts | 5 - .../unit/deployment/deploymentManager.test.ts | 88 ++++++++++++++ .../unit/workspace/workspacesProvider.test.ts | 53 +++++++++ 8 files changed, 240 insertions(+), 62 deletions(-) diff --git a/src/deployment/deploymentManager.ts b/src/deployment/deploymentManager.ts index 5c02437cf1..bee26f6ae8 100644 --- a/src/deployment/deploymentManager.ts +++ b/src/deployment/deploymentManager.ts @@ -208,7 +208,7 @@ export class DeploymentManager this.#authListenerDisposable?.dispose(); this.#authListenerDisposable = undefined; this.#sessionStore.signOut(null); - this.clearSideEffects(); + this.clearAuthState(); this.telemetryService.setDeploymentUrl(""); if (wasAuthenticated) { this.deploymentTelemetry.suspended(reason); @@ -224,13 +224,13 @@ export class DeploymentManager public suspendSession(reason: DeploymentSuspendReason): void { const wasAuthenticated = this.isAuthenticated(); this.#sessionStore.signOut(this.#sessionStore.current.deployment); - this.clearSideEffects(); + this.clearAuthState(); if (wasAuthenticated) { this.deploymentTelemetry.suspended(reason); } } - private clearSideEffects(): void { + private clearAuthState(): void { this.oauthSessionManager.clearDeployment(); this.client.setCredentials(undefined, undefined); this.updateAuthContexts(undefined); @@ -296,6 +296,10 @@ export class DeploymentManager ); } + /** + * Handle a token change while already signed in: verify the new token before + * applying it, so an unverified token never reaches the live client. + */ private async verifyAndUpdateSession( deployment: Deployment & { token: string }, ): Promise { @@ -305,10 +309,30 @@ export class DeploymentManager deployment.url, deployment.token, ); - if (this.#disposed || this.#sessionStore.current !== sessionBefore) { + if (this.#disposed) { return; } - await this.setDeployment({ ...deployment, user }); + const current = this.#sessionStore.current; + + if (current === sessionBefore) { + // Same-user rotation only needs fresh credentials; skipping + // setDeployment avoids a revision bump and tree rebuild. + if (current.kind === "signedIn" && current.user.id === user.id) { + this.client.setCredentials(deployment.url, deployment.token); + } else { + await this.setDeployment({ ...deployment, user }); + } + return; + } + + // Session changed under us: recover only a same-deployment suspension + // (e.g. a concurrent 401); a logout or switch wins. + if ( + current.kind === "signedOut" && + current.deployment?.safeHostname === deployment.safeHostname + ) { + await this.setDeployment({ ...deployment, user }); + } } catch (e) { this.logger.warn("Failed to authenticate updated session:", e); } diff --git a/src/deployment/sessionStore.ts b/src/deployment/sessionStore.ts index f9141cdf83..2c373e5cc0 100644 --- a/src/deployment/sessionStore.ts +++ b/src/deployment/sessionStore.ts @@ -38,12 +38,12 @@ export class SessionStore implements WorkspaceSessionState { public readonly onDidChange = this.#onDidChange.event; - /** The full current session, including deployment and user. */ + /** Full session state, including deployment and user. */ public get current(): SessionData { return this.#data; } - /** The lean projection exposed to session-state consumers. */ + /** Lean projection for consumers that only track auth status and revision. */ public getSnapshot(): WorkspaceSessionSnapshot { if (this.#data.kind === "signedIn") { return { diff --git a/src/extension.ts b/src/extension.ts index 3de6031b10..897cc29c68 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -101,6 +101,7 @@ async function doActivate( ? await secretsManager.getSessionAuth(deployment.safeHostname) : undefined; tracer.setAuthState(deploymentSessionAuth ? "stored" : "none"); + // Shared handler for auth failures (used by interceptor + session manager) const handleAuthFailure = (): Promise => { deploymentManager.suspendSession("auth_failure"); diff --git a/src/workspace/session.ts b/src/workspace/session.ts index 67b47d6fd9..ab180e8197 100644 --- a/src/workspace/session.ts +++ b/src/workspace/session.ts @@ -1,5 +1,12 @@ import type * as vscode from "vscode"; +/** + * A point-in-time view of the deployment session for workspace providers. + * + * `revision` increments on every sign-in or sign-out. Consumers snapshot the + * revision before an async call and compare afterward to detect that the + * session changed while the call was in flight. + */ export type WorkspaceSessionSnapshot = | { readonly kind: "signedOut"; readonly revision: number } | { diff --git a/src/workspace/workspacesProvider.ts b/src/workspace/workspacesProvider.ts index 03a5cb01dd..6d9225dcff 100644 --- a/src/workspace/workspacesProvider.ts +++ b/src/workspace/workspacesProvider.ts @@ -60,6 +60,9 @@ export interface WorkspaceProviderOptions { readonly refreshIntervalMs?: number; } +// Bounds fetch() retries when the session keeps changing mid-request. +const MAX_FETCH_ATTEMPTS = 3; + /** * Polls workspaces using the provided REST client and renders them in a tree. * @@ -119,7 +122,6 @@ export class WorkspaceProvider this.fetching = false; } - // Keep polling only after a clean fetch while still signed in and visible. if ( !hadError && !this.disposed && @@ -143,65 +145,75 @@ export class WorkspaceProvider * signed out, and throws if the query fails. */ private async fetch(): Promise { - const session = this.sessionState.getSnapshot(); - if (session.kind !== "signedIn") { - return []; - } + for (let attempt = 0; attempt < MAX_FETCH_ATTEMPTS; attempt++) { + if (this.disposed) { + return []; + } + const session = this.sessionState.getSnapshot(); + if (session.kind !== "signedIn") { + return []; + } - const resp = await this.client.getWorkspaces({ - q: this.getWorkspacesQuery, - }); + const resp = await this.client.getWorkspaces({ + q: this.getWorkspacesQuery, + }); - // If the session changed while the request was in flight, this result is - // stale. Drop it and fetch again for the current session. - const latest = this.sessionState.getSnapshot(); - if (latest.kind !== "signedIn" || latest.revision !== session.revision) { - return this.fetch(); - } + // Session changed mid-request; this result is stale, so retry. + const latest = this.sessionState.getSnapshot(); + if (latest.kind !== "signedIn" || latest.revision !== session.revision) { + continue; + } - const workspaces = this.filterWorkspaces(resp.workspaces, session); - const oldWatcherIds = [...this.agentWatchers.keys()]; - const reusedWatcherIds: string[] = []; - - // TODO: I think it might make more sense for the tree items to contain - // their own watchers, rather than recreate the tree items every time and - // have this separate map held outside the tree. - if (this.config.showMetadata) { - const agents = extractAllAgents(workspaces); - for (const agent of agents) { - // If we have an existing watcher, re-use it. - const oldWatcher = this.agentWatchers.get(agent.id); - if (oldWatcher) { - reusedWatcherIds.push(agent.id); - } else { - // Otherwise create a new watcher. + const workspaces = this.filterWorkspaces(resp.workspaces, session); + const oldWatcherIds = [...this.agentWatchers.keys()]; + const reusedWatcherIds: string[] = []; + + // TODO: I think it might make more sense for the tree items to contain + // their own watchers, rather than recreate the tree items every time + // and have this separate map held outside the tree. + if (this.config.showMetadata) { + const agents = extractAllAgents(workspaces); + for (const agent of agents) { + // If we have an existing watcher, re-use it. + const oldWatcher = this.agentWatchers.get(agent.id); + if (oldWatcher) { + reusedWatcherIds.push(agent.id); + continue; + } const watcher = await createAgentMetadataWatcher( agent.id, this.client, ); + // dispose() may have cleared the map mid-create; don't leak + // this watcher. + if (this.disposed) { + watcher.dispose(); + return []; + } watcher.onChange(() => this.refreshTree()); this.agentWatchers.set(agent.id, watcher); } } - } - // Dispose of watchers we ended up not reusing. - for (const id of oldWatcherIds) { - if (!reusedWatcherIds.includes(id)) { - this.agentWatchers.get(id)?.dispose(); - this.agentWatchers.delete(id); + // Dispose of watchers we ended up not reusing. + for (const id of oldWatcherIds) { + if (!reusedWatcherIds.includes(id)) { + this.agentWatchers.get(id)?.dispose(); + this.agentWatchers.delete(id); + } } - } - // Create tree items for each workspace - return workspaces.map( - (workspace: Workspace) => - new WorkspaceTreeItem( - workspace, - this.config.showOwner, - this.config.showMetadata, - ), - ); + return workspaces.map( + (workspace: Workspace) => + new WorkspaceTreeItem( + workspace, + this.config.showOwner, + this.config.showMetadata, + ), + ); + } + // Session changed on every attempt; the next refresh will catch up. + return []; } private filterWorkspaces( @@ -244,10 +256,7 @@ export class WorkspaceProvider } } - /** - * Schedule a refresh if one is not already scheduled or underway and a - * timeout length was provided. - */ + /** Schedule the next poll, unless one is pending or no interval is set. */ private maybeScheduleRefresh() { if (this.options.refreshIntervalMs && !this.timeout) { this.timeout = setTimeout(() => { @@ -370,6 +379,7 @@ export class WorkspaceProvider this.disposed = true; this.clearState(); this.sessionChangeDisposable.dispose(); + this._onDidChangeTreeData.dispose(); } } diff --git a/test/mocks/testHelpers.ts b/test/mocks/testHelpers.ts index a7725cdfaa..912fed72a0 100644 --- a/test/mocks/testHelpers.ts +++ b/test/mocks/testHelpers.ts @@ -1254,7 +1254,6 @@ interface WorkspacesResponse { * tests drive the production watcher rather than a stub. */ export class MockWorkspacesClient { - baseURL: string | undefined = "https://coder.example.com"; readonly metadataStreams = new Map< string, MockEventStream<{ data: AgentMetadataEvent[] }> @@ -1271,10 +1270,6 @@ export class MockWorkspacesClient { return Promise.resolve(stream); } - getAxiosInstance() { - return { defaults: { baseURL: this.baseURL } }; - } - /** Resolve the next getWorkspaces call with these workspaces. */ respondOnce(workspaces: readonly Workspace[]): void { this.getWorkspaces.mockResolvedValueOnce({ diff --git a/test/unit/deployment/deploymentManager.test.ts b/test/unit/deployment/deploymentManager.test.ts index 96595af335..b5600d4995 100644 --- a/test/unit/deployment/deploymentManager.test.ts +++ b/test/unit/deployment/deploymentManager.test.ts @@ -490,6 +490,94 @@ describe("DeploymentManager", () => { expect(manager.getCurrentDeployment()).toBeNull(); expect(currentUserId(manager)).toBeUndefined(); }); + + it("rotates the token in place without a revision bump when the user is unchanged", async () => { + const { mockClient, validationMockClient, secretsManager, manager } = + createTestContext(); + const user = createMockUser({ id: "stable-user" }); + validationMockClient.setAuthenticatedUserResponse(user); + + await manager.setDeployment({ + url: TEST_URL, + safeHostname: TEST_HOSTNAME, + token: "initial-token", + user, + }); + const revisionBefore = manager.getSnapshot().revision; + + await secretsManager.setSessionAuth(TEST_HOSTNAME, { + url: TEST_URL, + token: "rotated-token", + }); + await flush(); + + expect(mockClient.token).toBe("rotated-token"); + expect(currentUserId(manager)).toBe("stable-user"); + // No sign-in fired, so the workspace trees are not rebuilt. + expect(manager.getSnapshot().revision).toBe(revisionBefore); + }); + + it("keeps the existing session when verifying a rotated token fails", async () => { + const { mockClient, validationMockClient, secretsManager, manager } = + createTestContext(); + const user = createMockUser({ id: "stable-user" }); + + await manager.setDeployment({ + url: TEST_URL, + safeHostname: TEST_HOSTNAME, + token: "initial-token", + user, + }); + + // Verifying the rotated token fails (e.g. a transient network blip). + validationMockClient.getAuthenticatedUser.mockRejectedValueOnce( + new Error("network down"), + ); + await secretsManager.setSessionAuth(TEST_HOSTNAME, { + url: TEST_URL, + token: "rotated-token", + }); + await flush(); + + // Verify-before-apply: the unverified token never reaches the client. + expect(mockClient.token).toBe("initial-token"); + expect(manager.isAuthenticated()).toBe(true); + expect(currentUserId(manager)).toBe("stable-user"); + }); + + it("recovers a session suspended while a rotated token is verified", async () => { + const { mockClient, validationMockClient, secretsManager, manager } = + createTestContext(); + const user = createMockUser({ id: "stable-user" }); + const validation = Promise.withResolvers(); + validationMockClient.getAuthenticatedUser.mockReturnValueOnce( + validation.promise, + ); + + await manager.setDeployment({ + url: TEST_URL, + safeHostname: TEST_HOSTNAME, + token: "initial-token", + user, + }); + await secretsManager.setSessionAuth(TEST_HOSTNAME, { + url: TEST_URL, + token: "rotated-token", + }); + await flush(); + + // A concurrent 401 on the old token suspends the session mid-verify. + manager.suspendSession("auth_failure"); + expect(manager.isAuthenticated()).toBe(false); + + validation.resolve(user); + await flush(); + + // The verified token recovers the session instead of staying suspended. + expect(manager.isAuthenticated()).toBe(true); + expect(mockClient.token).toBe("rotated-token"); + expect(currentUserId(manager)).toBe("stable-user"); + }); }); describe("logout", () => { diff --git a/test/unit/workspace/workspacesProvider.test.ts b/test/unit/workspace/workspacesProvider.test.ts index 87b053b8f0..67af0ed7fa 100644 --- a/test/unit/workspace/workspacesProvider.test.ts +++ b/test/unit/workspace/workspacesProvider.test.ts @@ -463,4 +463,57 @@ describe("WorkspaceProvider", () => { expect(await provider.getChildren()).toEqual([]); }); + + it("fetches when the session signs in after starting signed out", async () => { + const { client, session, makeProvider } = setup(); + session.signOut(); + client.respondOnce([workspace({ name: "dev" })]); + const provider = makeProvider(WorkspaceQuery.Mine); + + await show(provider); + expect(client.getWorkspaces).not.toHaveBeenCalled(); + + session.signIn(); + await flush(); + + expect(await labels(provider)).toEqual(["dev"]); + }); + + it("stops reacting to session changes after dispose", async () => { + const { client, session, makeProvider } = setup(); + client.respondOnce([workspace({ name: "dev" })]); + const provider = makeProvider(WorkspaceQuery.Mine); + + await show(provider); + expect(await labels(provider)).toEqual(["dev"]); + + provider.dispose(); + session.signIn("another-user"); + await flush(); + + expect(client.getWorkspaces).toHaveBeenCalledTimes(1); + expect(await provider.getChildren()).toEqual([]); + }); + + it("leaves the tree empty when a fetch fails after a mid-request session change", async () => { + const { client, session, makeProvider } = setup(); + const pending = client.pending(); + client.getWorkspaces.mockRejectedValueOnce(new Error("network down")); + const provider = makeProvider(WorkspaceQuery.Mine); + + provider.setVisibility(true); + await flush(); + session.signIn("second-user"); + pending.resolve([workspace({ name: "stale" })]); + await flush(); + await flush(); + + // Failed retry leaves the tree empty until a manual refresh. + expect(await provider.getChildren()).toEqual([]); + + // A manual refresh recovers. + client.respondOnce([workspace({ name: "recovered" })]); + await provider.fetchAndRefresh(); + expect(await labels(provider)).toEqual(["recovered"]); + }); }); From 400b4d27227e9495ff102d0c35f5ad7edef8786e Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Mon, 8 Jun 2026 19:49:16 +0300 Subject: [PATCH 5/5] fix: address round 2 review findings - re-check the session revision after the watcher-creation await so a session change mid-await cannot overwrite the cleared tree with a stale session's workspaces (CRF-18) - add tests for the MAX_FETCH_ATTEMPTS exhaustion path and the dispose-during-watcher-creation guard (CRF-17, CRF-19) --- src/workspace/workspacesProvider.ts | 19 +++++--- .../unit/workspace/workspacesProvider.test.ts | 43 +++++++++++++++++++ 2 files changed, 56 insertions(+), 6 deletions(-) diff --git a/src/workspace/workspacesProvider.ts b/src/workspace/workspacesProvider.ts index 6d9225dcff..89fd132fbd 100644 --- a/src/workspace/workspacesProvider.ts +++ b/src/workspace/workspacesProvider.ts @@ -61,7 +61,7 @@ export interface WorkspaceProviderOptions { } // Bounds fetch() retries when the session keeps changing mid-request. -const MAX_FETCH_ATTEMPTS = 3; +export const MAX_FETCH_ATTEMPTS = 3; /** * Polls workspaces using the provided REST client and renders them in a tree. @@ -159,8 +159,7 @@ export class WorkspaceProvider }); // Session changed mid-request; this result is stale, so retry. - const latest = this.sessionState.getSnapshot(); - if (latest.kind !== "signedIn" || latest.revision !== session.revision) { + if (this.sessionChangedSince(session)) { continue; } @@ -184,9 +183,9 @@ export class WorkspaceProvider agent.id, this.client, ); - // dispose() may have cleared the map mid-create; don't leak - // this watcher. - if (this.disposed) { + // dispose() or a session change may have raced this await; + // drop the watcher rather than leak it or render stale data. + if (this.disposed || this.sessionChangedSince(session)) { watcher.dispose(); return []; } @@ -216,6 +215,14 @@ export class WorkspaceProvider return []; } + /** True if the session signed out or changed revision since `session`. */ + private sessionChangedSince( + session: Extract, + ): boolean { + const latest = this.sessionState.getSnapshot(); + return latest.kind !== "signedIn" || latest.revision !== session.revision; + } + private filterWorkspaces( workspaces: readonly Workspace[], session: Extract, diff --git a/test/unit/workspace/workspacesProvider.test.ts b/test/unit/workspace/workspacesProvider.test.ts index 67af0ed7fa..a9ba758a68 100644 --- a/test/unit/workspace/workspacesProvider.test.ts +++ b/test/unit/workspace/workspacesProvider.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it, vi } from "vitest"; import * as vscode from "vscode"; import { + MAX_FETCH_ATTEMPTS, WorkspaceProvider, WorkspaceQuery, type AgentTreeItem, @@ -14,6 +15,7 @@ import { createMockLogger, flush, flushPromises, + MockEventStream, MockWorkspaceSessionState, MockWorkspacesClient, TEST_CURRENT_USER_ID, @@ -516,4 +518,45 @@ describe("WorkspaceProvider", () => { await provider.fetchAndRefresh(); expect(await labels(provider)).toEqual(["recovered"]); }); + + it("gives up after the retry cap when the session keeps changing", async () => { + const { client, session, makeProvider } = setup(); + // Every response arrives stale: each call bumps the revision first. + client.getWorkspaces.mockImplementation(() => { + session.signIn(); + return Promise.resolve({ + workspaces: [workspace({ name: "ws" })], + count: 1, + }); + }); + const provider = makeProvider(WorkspaceQuery.Mine); + + await show(provider); + + expect(client.getWorkspaces).toHaveBeenCalledTimes(MAX_FETCH_ATTEMPTS); + expect(await provider.getChildren()).toEqual([]); + }); + + it("disposes a watcher created after dispose() races fetch()", async () => { + const { client, makeProvider } = setup(); + client.respondOnce([ + workspaceWithAgents({ name: "dev" }, [agent({ id: "agent-1" })]), + ]); + const stream = new MockEventStream<{ data: AgentMetadataEvent[] }>(); + const watch = + Promise.withResolvers>(); + client.watchAgentMetadata = vi.fn((_agentId: string) => watch.promise); + const provider = makeProvider(WorkspaceQuery.Mine); + + provider.setVisibility(true); + await flush(); + + // dispose() while the watcher socket is still opening. + provider.dispose(); + watch.resolve(stream); + await flush(); + + expect(stream.close).toHaveBeenCalled(); + expect(await provider.getChildren()).toEqual([]); + }); });