fix(mcp): resolve userId before JWT generation for agent block auth (#3932)

* fix(mcp): resolve userId before JWT generation for agent block auth

* test(mcp): add regression test for agent block JWT userId resolution
This commit is contained in:
Waleed
2026-04-03 19:05:10 -07:00
committed by GitHub
parent 8ce0299400
commit 57e5bac121
2 changed files with 37 additions and 3 deletions
+34
View File
@@ -25,6 +25,7 @@ const {
mockGetCustomToolById,
mockListCustomTools,
mockGetCustomToolByIdOrTitle,
mockGenerateInternalToken,
} = vi.hoisted(() => ({
mockIsHosted: { value: false },
mockEnv: { NEXT_PUBLIC_APP_URL: 'http://localhost:3000' } as Record<string, string | undefined>,
@@ -38,6 +39,7 @@ const {
mockGetCustomToolById: vi.fn(),
mockListCustomTools: vi.fn(),
mockGetCustomToolByIdOrTitle: vi.fn(),
mockGenerateInternalToken: vi.fn(),
}))
// Mock feature flags
@@ -65,6 +67,10 @@ vi.mock('@/lib/api-key/byok', () => ({
getBYOKKey: (...args: unknown[]) => mockGetBYOKKey(...args),
}))
vi.mock('@/lib/auth/internal', () => ({
generateInternalToken: (...args: unknown[]) => mockGenerateInternalToken(...args),
}))
vi.mock('@/lib/billing/core/usage-log', () => ({}))
vi.mock('@/lib/core/rate-limiter/hosted-key', () => ({
@@ -1154,6 +1160,34 @@ describe('MCP Tool Execution', () => {
expect(result.timing).toBeDefined()
})
it('should embed userId in JWT when executionContext is undefined (agent block path)', async () => {
mockGenerateInternalToken.mockResolvedValue('test-token')
global.fetch = Object.assign(
vi.fn().mockImplementation(async () => ({
ok: true,
status: 200,
json: () =>
Promise.resolve({
success: true,
data: { output: { content: [{ type: 'text', text: 'OK' }] } },
}),
})),
{ preconnect: vi.fn() }
) as typeof fetch
await executeTool('mcp-123-test_tool', {
query: 'test',
_context: {
workspaceId: 'workspace-456',
workflowId: 'workflow-789',
userId: 'user-abc',
},
})
expect(mockGenerateInternalToken).toHaveBeenCalledWith('user-abc')
})
describe('Tool request retries', () => {
function makeJsonResponse(
status: number,
+3 -3
View File
@@ -1552,11 +1552,13 @@ async function executeMcpTool(
const baseUrl = getInternalApiBaseUrl()
const mcpScope = resolveToolScope(params, executionContext)
const headers: Record<string, string> = { 'Content-Type': 'application/json' }
if (typeof window === 'undefined') {
try {
const internalToken = await generateInternalToken(executionContext?.userId)
const internalToken = await generateInternalToken(mcpScope.userId)
headers.Authorization = `Bearer ${internalToken}`
} catch (error) {
logger.error(`[${actualRequestId}] Failed to generate internal token:`, error)
@@ -1587,8 +1589,6 @@ async function executeMcpTool(
)
}
const mcpScope = resolveToolScope(params, executionContext)
if (mcpScope.callChain && mcpScope.callChain.length > 0) {
headers[SIM_VIA_HEADER] = serializeCallChain(mcpScope.callChain)
}