Merge branch 'main' into sou-1870-scoped-access-tokens

Author: brendan-kellam

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - [EE] Added one-hour repository-scoped access tokens with public mint and revoke APIs. [#1549](https://github.com/sourcebot-dev/sourcebot/pull/1549)
+- [EE] Added guided reconnection for MCP connector authentication failures during Ask Sourcebot agent turns. [#1548](https://github.com/sourcebot-dev/sourcebot/pull/1548)
 
 ### Removed
 - Removed the Langfuse integration. [#1536](https://github.com/sourcebot-dev/sourcebot/pull/1536)

packages/web/src/app/api/(client)/client.ts

@@ -320,7 +320,7 @@ export const getOffers = async (): Promise<OffersResponse | ServiceError> => {
     return result as OffersResponse | ServiceError;
 }
 
-export const connectMcpToAsk = async (body: { serverId: string; returnTo?: string }): Promise<ConnectMcpResponse | ServiceError> => {
+export const connectMcpToAsk = async (body: { serverId: string; returnTo?: string; forceAuthorization?: boolean }): Promise<ConnectMcpResponse | ServiceError> => {
     const result = await fetch('/api/ee/askmcp/connect', {
         method: 'POST',
         headers: {

packages/web/src/app/api/(server)/ee/askmcp/connect/route.test.ts

@@ -48,7 +48,7 @@ vi.mock('@ai-sdk/mcp', () => ({
 const { POST } = await import('./route');
 const { getMcpOAuthReturnToFromState } = await import('@/ee/features/chat/mcp/mcpOAuthReturnTo');
 
-function createRequest(body: { serverId: string; returnTo?: string } = { serverId: 'server-1' }) {
+function createRequest(body: { serverId: string; returnTo?: string; forceAuthorization?: boolean } = { serverId: 'server-1' }) {
     return new NextRequest('https://sourcebot.example.com/api/ee/askmcp/connect', {
         method: 'POST',
         headers: { 'content-type': 'application/json' },
@@ -229,6 +229,60 @@ describe('POST /api/ee/askmcp/connect', () => {
         });
     });
 
+    test('forces an interactive OAuth redirect for reconnect recovery', async () => {
+        const prisma = createPrismaMock();
+        const tx = createTransactionMock();
+        tx.userMcpServer.findUnique.mockResolvedValue({
+            tokens: 'encrypted:{"access_token":"stale-token"}',
+            codeVerifier: null,
+            state: null,
+        });
+        mocks.authContext = {
+            org: { id: 1 },
+            user: { id: 'user-1' },
+            prisma,
+        };
+        mocks.unsafePrisma.$transaction.mockImplementation(async (callback, _options) => callback(tx));
+        mocks.mcpAuth.mockImplementation(async (provider) => {
+            await expect(provider.tokens()).resolves.toBeUndefined();
+            provider.authorizationUrl = 'https://oauth.example.com/authorize';
+            return 'REDIRECT';
+        });
+
+        const response = await POST(createRequest({
+            serverId: 'server-1',
+            returnTo: '/chat/abc123',
+            forceAuthorization: true,
+        }));
+
+        expect(await response.json()).toEqual({
+            authorizationUrl: 'https://oauth.example.com/authorize',
+        });
+    });
+
+    test('does not report a forced reconnect as successful without an OAuth redirect', async () => {
+        const prisma = createPrismaMock();
+        const tx = createTransactionMock();
+        mocks.authContext = {
+            org: { id: 1 },
+            user: { id: 'user-1' },
+            prisma,
+        };
+        mocks.unsafePrisma.$transaction.mockImplementation(async (callback, _options) => callback(tx));
+        mocks.mcpAuth.mockResolvedValue('AUTHORIZED');
+
+        const response = await POST(createRequest({
+            serverId: 'server-1',
+            returnTo: '/chat/abc123',
+            forceAuthorization: true,
+        }));
+
+        expect(response.status).toBe(502);
+        expect(await response.json()).toMatchObject({
+            message: 'Could not start connector reauthorization.',
+        });
+    });
+
     test('ignores unsafe return paths', async () => {
         const prisma = createPrismaMock();
         const tx = createTransactionMock();

packages/web/src/app/api/(server)/ee/askmcp/connect/route.ts

@@ -24,6 +24,7 @@ import { getEnabledMcpOAuthScopeNames } from '@/ee/features/chat/mcp/oauthScopeU
 const bodySchema = z.object({
     serverId: z.string(),
     returnTo: z.string().optional(),
+    forceAuthorization: z.boolean().optional().default(false),
 });
 const logger = createLogger('mcp-connect');
 const MCP_AUTH_FETCH_TIMEOUT_MS = Math.min(env.SOURCEBOT_MCP_TOOL_CALL_TIMEOUT_MS, 30000);
@@ -146,6 +147,7 @@ export const POST = apiHandler(async (request: NextRequest) => {
                     callbackReturnTo,
                     allowClientRegistration: true,
                     requestedOAuthScopes: getEnabledMcpOAuthScopeNames(mcpServer.oauthScopes),
+                    forceAuthorization: parsed.data.forceAuthorization,
                 });
 
                 let authResult: Awaited<ReturnType<typeof mcpAuth>>;
@@ -200,7 +202,7 @@ export const POST = apiHandler(async (request: NextRequest) => {
             throw error;
         }
 
-        if (connectResult.authResult === 'AUTHORIZED') {
+        if (connectResult.authResult === 'AUTHORIZED' && !parsed.data.forceAuthorization) {
             // Already has valid tokens (e.g., refreshed)
             void captureEvent('ask_mcp_connector_connection_completed', {
                 ...eventProperties,
@@ -209,6 +211,18 @@ export const POST = apiHandler(async (request: NextRequest) => {
             return { authorizationUrl: null } satisfies ConnectMcpResponse;
         }
 
+        if (connectResult.authResult === 'AUTHORIZED') {
+            void captureEvent('ask_mcp_connector_connection_failed', {
+                ...eventProperties,
+                failureReason: 'missing_authorization_url',
+            });
+            throw new ServiceErrorException({
+                statusCode: StatusCodes.BAD_GATEWAY,
+                errorCode: ErrorCode.UNEXPECTED_ERROR,
+                message: 'Could not start connector reauthorization.',
+            });
+        }
+
         if (!connectResult.authorizationUrl) {
             void captureEvent('ask_mcp_connector_connection_failed', {
                 ...eventProperties,

packages/web/src/ee/features/chat/agent.test.ts

@@ -321,6 +321,50 @@ describe('createMessageStream approval continuation', () => {
         });
     });
 
+    test('streams the connector ID when its tools fail to load', async () => {
+        const { getConnectedMcpClients } = await import('@/ee/features/chat/mcp/mcpClientFactory');
+        const { getMcpTools } = await import('@/ee/features/chat/mcp/mcpToolSets');
+        vi.mocked(getConnectedMcpClients).mockResolvedValueOnce([
+            { serverId: 'server-linear', serverName: 'Linear' },
+        ] as never);
+        vi.mocked(getMcpTools).mockResolvedValueOnce({
+            tools: {},
+            failedServers: [{ serverId: 'server-linear', serverName: 'Linear' }],
+            serverFaviconUrls: {},
+            toolDisplayNames: {},
+            cleanup: vi.fn(),
+        });
+        mockAi.streamText.mockReturnValue(createFakeStreamResult());
+
+        await createMessageStream({
+            chatId: 'chat-id',
+            messages: [createUserMessage()],
+            selectedRepos: [],
+            disabledMcpServerIds: [],
+            prisma: {},
+            model: {},
+            modelName: 'test-model',
+            promptCacheStrategy: noopStrategy,
+            onFinish: vi.fn(),
+            onError: () => 'error',
+            userId: 'user-id',
+            orgId: 1,
+        } as unknown as Parameters<typeof createMessageStream>[0]);
+
+        const execute = mockAi.latestCreateUIMessageStreamOptions?.execute;
+        if (!execute) {
+            throw new Error('Expected createUIMessageStream to capture execute callback.');
+        }
+
+        const write = vi.fn();
+        await execute({ writer: { merge: vi.fn(), write } });
+
+        expect(write).toHaveBeenCalledWith({
+            type: 'data-mcp-failed-server',
+            data: { serverId: 'server-linear', serverName: 'Linear' },
+        });
+    });
+
     test.each([
         ['dynamic', dynamicApprovalRespondedPart],
         ['static', staticApprovalRespondedPart],

packages/web/src/ee/features/chat/agent.ts

@@ -26,6 +26,12 @@ import { addLineNumbers, fileReferenceToString, formatAttachmentsForPrompt, getA
 import { createTools } from "./tools";
 import { getConnectedMcpClients } from "@/ee/features/chat/mcp/mcpClientFactory";
 import { getMcpTools, McpToolsResult } from "@/ee/features/chat/mcp/mcpToolSets";
+import {
+    createMcpAuthInterruptionDirective,
+    denyApprovedToolApprovalsForAuthInterruption,
+    getMcpAuthRequiredFailureFromAssistantMessage,
+    McpToolAuthFailure,
+} from "@/ee/features/chat/mcp/mcpAuthFailure";
 import { buildMcpToolRegistry, McpToolRegistryEntry } from "@/ee/features/chat/mcp/mcpToolRegistry";
 import { PromptCacheStrategy, mergeProviderOptions, detectPromptCacheBreak, detectUnexpectedCacheMiss } from "./promptCaching";
 import { hasEntitlement } from '@/lib/entitlements';
@@ -332,9 +338,21 @@ export const createMessageStream = async ({
         ? (lastMsg.metadata as SBChatMessageMetadata | undefined)
         : undefined;
 
+    // When the response was interrupted by a reconnect-required authentication
+    // failure (detected via the safe tool error's marker text), the
+    // continuation must run its final step with tool use disabled. Any
+    // approval that was still approved is rewritten to a denial: once the
+    // response is authentication-terminal, later approval actions are invalid.
+    const priorMcpAuthFailure = hasApprovalContinuationReady
+        ? getMcpAuthRequiredFailureFromAssistantMessage(lastMsg)
+        : undefined;
+
     if (hasApprovalContinuationReady) {
+        const continuationMessage = priorMcpAuthFailure
+            ? denyApprovedToolApprovalsForAuthInterruption(lastMsg, priorMcpAuthFailure.serverName)
+            : lastMsg;
         const fullLastTurn = await convertToModelMessages(
-            [lastMsg],
+            [continuationMessage],
             { ignoreIncompleteToolCalls: true }
         );
         messageHistory = [...messageHistory, ...fullLastTurn];
@@ -375,12 +393,22 @@ export const createMessageStream = async ({
                         data: { modelToolName, rawToolName },
                     });
                 },
-                onMcpServerFailed: (serverName) => {
+                onMcpServerFailed: (server) => {
                     writer.write({
                         type: 'data-mcp-failed-server',
-                        data: { serverName },
+                        data: server,
                     });
                 },
+                onMcpAuthRequired: (failure) => {
+                    // Transient: consumed live by the client to surface the
+                    // connector reconnect UI, never folded into persisted parts.
+                    writer.write({
+                        type: 'data-mcp-auth-required',
+                        data: failure,
+                        transient: true,
+                    });
+                },
+                priorMcpAuthFailure,
                 traceId,
                 chatId,
                 prisma,
@@ -508,7 +536,14 @@ interface AgentOptions {
     onWriteSource: (source: Source) => void;
     onMcpServerDiscovered: (sanitizedName: string, faviconUrl: string) => void;
     onMcpToolDiscovered: (modelToolName: string, rawToolName: string) => void;
-    onMcpServerFailed: (serverName: string) => void;
+    onMcpServerFailed: (server: { serverId: string; serverName: string }) => void;
+    // Fired at most once per connector per response when a tool call fails
+    // with a reconnect-required authentication failure.
+    onMcpAuthRequired: (failure: McpToolAuthFailure) => void;
+    // Set when the incoming messages show this response was already
+    // interrupted by an authentication failure (approval continuation): the
+    // stream must run its final step with tool use disabled from step one.
+    priorMcpAuthFailure?: { serverName: string };
     traceId: string;
     chatId: string;
     prisma: PrismaClient;
@@ -529,6 +564,8 @@ const createAgentStream = async ({
     onMcpServerDiscovered,
     onMcpToolDiscovered,
     onMcpServerFailed,
+    onMcpAuthRequired,
+    priorMcpAuthFailure,
     traceId,
     chatId,
     prisma,
@@ -564,6 +601,15 @@ const createAgentStream = async ({
         }))
     ).filter((source) => source !== undefined);
 
+    // Mutable, response-scoped authentication failure state. `serverName` is
+    // the first failed connector's display name (V1 supports recovery for a
+    // single failed connector). Failures are deduplicated by connector so the
+    // client sees at most one transient event per connector per response.
+    const mcpAuthFailureState: { failure?: { serverName: string } } = {
+        ...(priorMcpAuthFailure ? { failure: { serverName: priorMcpAuthFailure.serverName } } : {}),
+    };
+    const reportedMcpAuthFailureServerIds = new Set<string>();
+
     let mcpToolSetsObj: McpToolsResult = { tools: {}, failedServers: [], serverFaviconUrls: {}, toolDisplayNames: {}, cleanup: async () => {} };
     if (userId && orgId && await hasEntitlement('ask') && disabledMcpServerIds !== undefined) {
         try {
@@ -573,6 +619,16 @@ const createAgentStream = async ({
                 chatId,
                 traceId,
                 source: 'sourcebot-ask-agent',
+            }, {
+                onAuthFailure: (failure) => {
+                    if (!mcpAuthFailureState.failure) {
+                        mcpAuthFailureState.failure = { serverName: failure.serverName };
+                    }
+                    if (!reportedMcpAuthFailureServerIds.has(failure.serverId)) {
+                        reportedMcpAuthFailureServerIds.add(failure.serverId);
+                        onMcpAuthRequired(failure);
+                    }
+                },
             });
 
             for (const [sanitizedName, faviconUrl] of Object.entries(mcpToolSetsObj.serverFaviconUrls)) {
@@ -590,8 +646,8 @@ const createAgentStream = async ({
         }
     }
 
-    for (const serverName of mcpToolSetsObj.failedServers) {
-        onMcpServerFailed(serverName);
+    for (const server of mcpToolSetsObj.failedServers) {
+        onMcpServerFailed(server);
     }
 
     const mcpRegistry = buildMcpToolRegistry(mcpToolSetsObj.tools);
@@ -715,14 +771,34 @@ const createAgentStream = async ({
             // rebuilds the step's messages each time as the original input plus
             // its own accumulated response messages. Re-applying the moving tail marker
             // to the new last message each step is safe and does not accumulate.
-            prepareStep: (tailMarker || hasMcpTools) ? ({ steps, messages }) => {
+            prepareStep: (tailMarker || hasMcpTools || mcpAuthFailureState.failure) ? ({ steps, messages }) => {
                 const stepMessages = (tailMarker && messages.length > 0)
                     ? messages.map((message, index) =>
                         index === messages.length - 1
                             ? { ...message, providerOptions: mergeProviderOptions(message.providerOptions, tailMarker) }
                             : message)
                     : undefined;
 
+                // Once a reconnect-required authentication failure occurs, the
+                // response is terminal for tool use: every remaining step runs
+                // with tool calling disabled (`toolChoice: 'none'` keeps the
+                // tool definitions byte-stable for prompt caching) plus an
+                // ephemeral directive to summarize completed work and prompt
+                // the user to reconnect. In-flight tool calls of the failing
+                // step have already run to completion by the time this fires.
+                if (mcpAuthFailureState.failure) {
+                    return {
+                        messages: [
+                            ...(stepMessages ?? messages),
+                            {
+                                role: 'user' as const,
+                                content: createMcpAuthInterruptionDirective(mcpAuthFailureState.failure.serverName),
+                            },
+                        ],
+                        toolChoice: 'none' as const,
+                    };
+                }
+
                 if (!hasMcpTools) {
                     return stepMessages ? { messages: stepMessages } : {};
                 }