refactor(slack): store the workspace bot token on the workspace

apps/web/src/lib/bot/platforms/slack.tsx

@@ -13,7 +13,7 @@ import {
 import type { BotPlatform, RequesterInfo } from '@/lib/bot/platforms/types';
 import { BOT_CONTEXT_MESSAGE_LIMIT } from '@/lib/bot/constants';
 import { APP_URL } from '@/lib/constants';
-import { getAccessTokenFromInstallation } from '@/lib/integrations/slack-service';
+import { getSlackBotToken } from '@/lib/integrations/slack-workspace-installation';
 import { PLATFORM } from '@/lib/integrations/core/constants';
 import { getSlackMessagePermalink } from '@/lib/slack-bot/slack-utils';
 import { captureException } from '@sentry/nextjs';
@@ -219,7 +219,7 @@ async function getSlackRequesterInfo(
   platformIntegration: PlatformIntegration,
   displayName: string
 ): Promise<RequesterInfo> {
-  const accessToken = getAccessTokenFromInstallation(platformIntegration);
+  const accessToken = await getSlackBotToken(platformIntegration);
   if (!accessToken) {
     return { displayName, platform: PLATFORM.SLACK };
   }

apps/web/src/lib/integrations/oauth/platforms/slack-callback.ts

@@ -129,7 +129,12 @@ export async function handleSlackOAuthCallback(request: NextRequest) {
 
     // 8. Store installation in database
     try {
-      await upsertSlackInstallation({ owner, teamId, installation });
+      await upsertSlackInstallation({
+        owner,
+        teamId,
+        installation,
+        installedByUserId: user.id,
+      });
     } catch (error) {
       if (error instanceof SlackWorkspaceAlreadyConnectedError) {
         return NextResponse.redirect(

apps/web/src/lib/integrations/slack-service.test.ts

@@ -10,9 +10,14 @@ const mockInsertValues = jest.fn();
 const mockInsertReturning = jest.fn();
 const mockAuthRevoke = jest.fn();
 const mockAuthTest = jest.fn();
-
-jest.mock('@/lib/drizzle', () => ({
-  db: {
+const mockUpsertWorkspaceInstallation = jest.fn();
+const mockDeleteWorkspaceInstallation = jest.fn();
+const mockDeleteWorkspaceInstallationIfUnreferenced = jest.fn();
+const mockGetSlackBotToken = jest.fn();
+const mockTransaction = jest.fn();
+
+jest.mock('@/lib/drizzle', () => {
+  const client = {
     select: jest.fn(() => ({
       from: jest.fn(() => ({
         where: jest.fn(() => ({
@@ -29,8 +34,13 @@ jest.mock('@/lib/drizzle', () => ({
     insert: jest.fn(() => ({
       values: mockInsertValues,
     })),
-  },
-}));
+    // Hands the callback the same client, so statements issued inside a
+    // transaction land on the same assertions as statements issued outside one.
+    transaction: (...args: unknown[]) => mockTransaction(...args),
+  };
+
+  return { db: client };
+});
 
 jest.mock('@slack/web-api', () => ({
   WebClient: jest.fn(() => ({
@@ -41,6 +51,19 @@ jest.mock('@slack/web-api', () => ({
   })),
 }));
 
+// The factory runs while `slack-service` is being imported, which happens before
+// the `const` declarations above are initialized, so each export forwards lazily
+// instead of capturing the mock function directly.
+jest.mock('@/lib/integrations/slack-workspace-installation', () => ({
+  upsertSlackWorkspaceInstallation: (...args: unknown[]) =>
+    mockUpsertWorkspaceInstallation(...args),
+  deleteSlackWorkspaceInstallation: (...args: unknown[]) =>
+    mockDeleteWorkspaceInstallation(...args),
+  deleteSlackWorkspaceInstallationIfUnreferenced: (...args: unknown[]) =>
+    mockDeleteWorkspaceInstallationIfUnreferenced(...args),
+  getSlackBotToken: (...args: unknown[]) => mockGetSlackBotToken(...args),
+}));
+
 import type { Owner } from '@/lib/integrations/core/types';
 import type { SlackInstallation } from '@chat-adapter/slack';
 import { DEFAULT_BOT_MODEL } from '@/lib/bot/constants';
@@ -69,6 +92,23 @@ function buildSlackIntegration(overrides: Record<string, unknown> = {}) {
   };
 }
 
+function resetTransactionMock() {
+  const { db } = jest.requireMock('@/lib/drizzle') as { db: unknown };
+  mockTransaction.mockReset();
+  mockTransaction.mockImplementation(async (fn: (tx: unknown) => Promise<unknown>) => fn(db));
+}
+
+function resetWorkspaceInstallationMocks() {
+  mockUpsertWorkspaceInstallation.mockReset();
+  mockDeleteWorkspaceInstallation.mockReset();
+  mockDeleteWorkspaceInstallationIfUnreferenced.mockReset();
+  mockGetSlackBotToken.mockReset();
+  mockUpsertWorkspaceInstallation.mockResolvedValue({ team_id: 'T123' });
+  mockDeleteWorkspaceInstallation.mockResolvedValue(undefined);
+  mockDeleteWorkspaceInstallationIfUnreferenced.mockResolvedValue(undefined);
+  mockGetSlackBotToken.mockResolvedValue('xoxb-token');
+}
+
 describe('slack-service uninstallApp', () => {
   beforeEach(() => {
     mockLimit.mockReset();
@@ -82,6 +122,28 @@ describe('slack-service uninstallApp', () => {
     mockUpdateWhere.mockReturnValue({ returning: mockUpdateReturning });
     mockUpdateReturning.mockResolvedValue([buildSlackIntegration()]);
     mockDeleteWhere.mockResolvedValue(undefined);
+    resetWorkspaceInstallationMocks();
+    resetTransactionMock();
+  });
+
+  // Whether the record is actually removed depends on the reference check inside
+  // deleteSlackWorkspaceInstallationIfUnreferenced, which is covered against a real
+  // database in slack-workspace-installation.test.ts.
+  it('asks for the workspace installation to be removed', async () => {
+    mockLimit.mockResolvedValue([buildSlackIntegration()]);
+
+    await uninstallApp(owner);
+
+    expect(mockDeleteWorkspaceInstallationIfUnreferenced).toHaveBeenCalledWith('T123');
+  });
+
+  it('revokes the token resolved from the workspace installation', async () => {
+    mockLimit.mockResolvedValue([buildSlackIntegration({ metadata: {} })]);
+    mockGetSlackBotToken.mockResolvedValue('xoxb-workspace-token');
+
+    await uninstallApp(owner);
+
+    expect(mockAuthRevoke).toHaveBeenCalledTimes(1);
   });
 
   it('deletes Chat SDK Slack state before removing the platform integration row', async () => {
@@ -165,6 +227,7 @@ describe('slack-service uninstallApp', () => {
     expect(deleteChatSdkInstallation).not.toHaveBeenCalled();
     expect(deleteChatSdkIdentityCache).not.toHaveBeenCalled();
     expect(mockDeleteWhere).toHaveBeenCalledTimes(1);
+    expect(mockDeleteWorkspaceInstallationIfUnreferenced).not.toHaveBeenCalled();
   });
 });
 
@@ -173,6 +236,8 @@ describe('slack-service deleteInstallationByTeamId', () => {
     mockLimit.mockReset();
     mockDeleteWhere.mockReset();
     mockDeleteWhere.mockResolvedValue(undefined);
+    resetWorkspaceInstallationMocks();
+    resetTransactionMock();
   });
 
   it('deletes the platform integration and Chat SDK state for a Slack team', async () => {
@@ -184,13 +249,28 @@ describe('slack-service deleteInstallationByTeamId', () => {
     });
 
     expect(mockDeleteWhere).toHaveBeenCalledTimes(1);
+    expect(mockDeleteWorkspaceInstallationIfUnreferenced).toHaveBeenCalledWith('T123');
+  });
+
+  it('still clears workspace state when the platform integration is already gone', async () => {
+    mockLimit.mockResolvedValue([]);
+
+    await expect(deleteInstallationByTeamId('T123')).resolves.toEqual({
+      success: true,
+      deleted: false,
+    });
+
+    expect(mockDeleteWhere).not.toHaveBeenCalled();
+    expect(mockDeleteWorkspaceInstallation).toHaveBeenCalledWith('T123');
   });
 });
 
 describe('slack-service testConnection', () => {
   beforeEach(() => {
     mockLimit.mockReset();
     mockAuthTest.mockReset();
+    resetWorkspaceInstallationMocks();
+    resetTransactionMock();
   });
 
   it('returns success when auth.test succeeds', async () => {
@@ -211,8 +291,9 @@ describe('slack-service testConnection', () => {
     expect(mockAuthTest).not.toHaveBeenCalled();
   });
 
-  it('returns failure when the access token is missing from metadata', async () => {
+  it('returns failure when no bot token can be resolved', async () => {
     mockLimit.mockResolvedValue([buildSlackIntegration({ metadata: {} })]);
+    mockGetSlackBotToken.mockResolvedValue(undefined);
 
     await expect(testConnection(owner)).resolves.toEqual({
       success: false,
@@ -267,6 +348,166 @@ describe('upsertSlackInstallation', () => {
     mockInsertReturning.mockReset();
     mockInsertValues.mockReturnValue({ returning: mockInsertReturning });
     mockInsertReturning.mockResolvedValue([buildSlackIntegration()]);
+    resetWorkspaceInstallationMocks();
+    resetTransactionMock();
+  });
+
+  it('writes the bot token to the workspace installation store', async () => {
+    mockLimit.mockResolvedValueOnce([]).mockResolvedValueOnce([]);
+
+    const installation = {
+      botToken: 'xoxb-new-token',
+      botUserId: 'U_NEW_BOT',
+      teamName: 'Kilo Team',
+    } satisfies SlackInstallation;
+
+    await upsertSlackInstallation({
+      owner,
+      teamId: 'T123',
+      installation,
+      installedByUserId: 'user-1',
+    });
+
+    expect(mockUpsertWorkspaceInstallation).toHaveBeenCalledWith(
+      expect.objectContaining({
+        teamId: 'T123',
+        teamName: 'Kilo Team',
+        botToken: 'xoxb-new-token',
+        botUserId: 'U_NEW_BOT',
+        scopes: SLACK_SCOPES,
+        installedByUserId: 'user-1',
+      })
+    );
+  });
+
+  it('writes the workspace installation before the platform integration row', async () => {
+    mockLimit.mockResolvedValueOnce([]).mockResolvedValueOnce([]);
+
+    const installation = {
+      botToken: 'xoxb-new-token',
+      botUserId: 'U_NEW_BOT',
+      teamName: 'Kilo Team',
+    } satisfies SlackInstallation;
+
+    await upsertSlackInstallation({ owner, teamId: 'T123', installation });
+
+    expect(mockUpsertWorkspaceInstallation.mock.invocationCallOrder[0]).toBeLessThan(
+      mockInsertValues.mock.invocationCallOrder[0]
+    );
+  });
+
+  it('removes the previous workspace installation when an owner switches workspaces', async () => {
+    mockLimit.mockResolvedValue([
+      buildSlackIntegration({ platform_installation_id: 'T_OLD', platform_account_id: 'T_OLD' }),
+    ]);
+
+    const installation = {
+      botToken: 'xoxb-new-token',
+      botUserId: 'U_NEW_BOT',
+      teamName: 'Kilo Team',
+    } satisfies SlackInstallation;
+
+    await upsertSlackInstallation({ owner, teamId: 'T_NEW', installation });
+
+    expect(mockUpsertWorkspaceInstallation).toHaveBeenCalledWith(
+      expect.objectContaining({ teamId: 'T_NEW' })
+    );
+    expect(mockDeleteWorkspaceInstallationIfUnreferenced).toHaveBeenCalledWith('T_OLD');
+  });
+
+  it('keeps the previous workspace installation when the workspace is unchanged', async () => {
+    mockLimit.mockResolvedValue([
+      buildSlackIntegration({ platform_installation_id: 'T123', platform_account_id: 'T123' }),
+    ]);
+
+    const installation = {
+      botToken: 'xoxb-new-token',
+      botUserId: 'U_NEW_BOT',
+      teamName: 'Kilo Team',
+    } satisfies SlackInstallation;
+
+    await upsertSlackInstallation({ owner, teamId: 'T123', installation });
+
+    expect(mockDeleteWorkspaceInstallationIfUnreferenced).not.toHaveBeenCalled();
+  });
+
+  it('writes the workspace installation inside the integration transaction', async () => {
+    mockLimit.mockResolvedValueOnce([]).mockResolvedValueOnce([]);
+
+    const installation = {
+      botToken: 'xoxb-new-token',
+      botUserId: 'U_NEW_BOT',
+      teamName: 'Kilo Team',
+    } satisfies SlackInstallation;
+
+    await upsertSlackInstallation({ owner, teamId: 'T123', installation });
+
+    expect(mockTransaction).toHaveBeenCalledTimes(1);
+    expect(mockUpsertWorkspaceInstallation).toHaveBeenCalledWith(
+      expect.objectContaining({ executor: expect.anything() })
+    );
+  });
+
+  // The transaction rolls the workspace record back, so there is no compensating
+  // delete to issue and no risk of one masking the original error.
+  it('propagates the original error without a compensating delete when the write fails', async () => {
+    mockLimit.mockResolvedValueOnce([]).mockResolvedValueOnce([]);
+    mockInsertReturning.mockRejectedValue(new Error('insert exploded'));
+
+    const installation = {
+      botToken: 'xoxb-new-token',
+      botUserId: 'U_NEW_BOT',
+      teamName: 'Kilo Team',
+    } satisfies SlackInstallation;
+
+    await expect(upsertSlackInstallation({ owner, teamId: 'T123', installation })).rejects.toThrow(
+      'insert exploded'
+    );
+
+    expect(mockDeleteWorkspaceInstallationIfUnreferenced).not.toHaveBeenCalled();
+    expect(mockDeleteWorkspaceInstallation).not.toHaveBeenCalled();
+  });
+
+  it('still translates a workspace unique violation when the write fails', async () => {
+    mockLimit.mockResolvedValueOnce([]).mockResolvedValueOnce([]);
+    mockInsertReturning.mockRejectedValue({
+      constraint: 'UQ_platform_integrations_slack_platform_inst',
+    });
+
+    const installation = {
+      botToken: 'xoxb-new-token',
+      botUserId: 'U_NEW_BOT',
+      teamName: 'Kilo Team',
+    } satisfies SlackInstallation;
+
+    await expect(upsertSlackInstallation({ owner, teamId: 'T123', installation })).rejects.toThrow(
+      SlackWorkspaceAlreadyConnectedError
+    );
+  });
+
+  it('clears an earlier suspension when refreshing an existing installation', async () => {
+    mockLimit.mockResolvedValue([
+      buildSlackIntegration({
+        integration_status: 'suspended',
+        suspended_by: 'duplicate_slack_workspace_migration',
+      }),
+    ]);
+
+    const installation = {
+      botToken: 'xoxb-new-token',
+      botUserId: 'U_NEW_BOT',
+      teamName: 'Kilo Team',
+    } satisfies SlackInstallation;
+
+    await upsertSlackInstallation({ owner, teamId: 'T123', installation });
+
+    expect(mockUpdateSet).toHaveBeenCalledWith(
+      expect.objectContaining({
+        integration_status: 'active',
+        suspended_at: null,
+        suspended_by: null,
+      })
+    );
   });
 
   it('preserves the selected model when refreshing an existing installation', async () => {