fix(oauth): wire post-request 401/403 handler into skill-mcp withOperationRetry
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
This commit is contained in:
@@ -0,0 +1,162 @@
|
|||||||
|
import { afterAll, beforeEach, describe, expect, it, mock } from "bun:test"
|
||||||
|
import type { ClaudeCodeMcpServer } from "../claude-code-mcp-loader/types"
|
||||||
|
import type { OAuthTokenData } from "../mcp-oauth/storage"
|
||||||
|
import type { SkillMcpClientInfo, SkillMcpServerContext } from "./types"
|
||||||
|
|
||||||
|
const mockGetOrCreateClient = mock(async () => {
|
||||||
|
throw new Error("not used")
|
||||||
|
})
|
||||||
|
|
||||||
|
const mockGetOrCreateClientWithRetryImpl = mock(async () => ({
|
||||||
|
callTool: mock(async () => ({ content: [{ type: "text", text: "unused" }] })),
|
||||||
|
close: mock(async () => {}),
|
||||||
|
}))
|
||||||
|
|
||||||
|
mock.module("./connection", () => ({
|
||||||
|
getOrCreateClient: mockGetOrCreateClient,
|
||||||
|
getOrCreateClientWithRetryImpl: mockGetOrCreateClientWithRetryImpl,
|
||||||
|
}))
|
||||||
|
|
||||||
|
mock.module("../mcp-oauth/provider", () => ({
|
||||||
|
McpOAuthProvider: class MockMcpOAuthProvider {},
|
||||||
|
}))
|
||||||
|
|
||||||
|
type ManagerModule = typeof import("./manager")
|
||||||
|
|
||||||
|
async function importFreshManagerModule(): Promise<ManagerModule> {
|
||||||
|
return await import(new URL(`./manager.ts?oauth-retry-test=${Date.now()}-${Math.random()}`, import.meta.url).href)
|
||||||
|
}
|
||||||
|
|
||||||
|
function createInfo(): SkillMcpClientInfo {
|
||||||
|
return {
|
||||||
|
serverName: "oauth-server",
|
||||||
|
skillName: "oauth-skill",
|
||||||
|
sessionID: "session-1",
|
||||||
|
scope: "builtin",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function createContext(): SkillMcpServerContext {
|
||||||
|
return {
|
||||||
|
skillName: "oauth-skill",
|
||||||
|
config: {
|
||||||
|
url: "https://mcp.example.com/mcp",
|
||||||
|
oauth: { clientId: "test-client" },
|
||||||
|
} satisfies ClaudeCodeMcpServer,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
afterAll(() => {
|
||||||
|
mock.restore()
|
||||||
|
})
|
||||||
|
|
||||||
|
describe("SkillMcpManager post-request OAuth retry", () => {
|
||||||
|
beforeEach(() => {
|
||||||
|
mockGetOrCreateClient.mockClear()
|
||||||
|
mockGetOrCreateClientWithRetryImpl.mockClear()
|
||||||
|
})
|
||||||
|
|
||||||
|
it("retries the operation after a 401 refresh succeeds", async () => {
|
||||||
|
// given
|
||||||
|
const { SkillMcpManager } = await importFreshManagerModule()
|
||||||
|
const refresh = mock(async () => ({ accessToken: "refreshed-token" } satisfies OAuthTokenData))
|
||||||
|
const manager = new SkillMcpManager({
|
||||||
|
createOAuthProvider: () => ({
|
||||||
|
tokens: () => ({ accessToken: "stale-token", refreshToken: "refresh-token" }),
|
||||||
|
login: mock(async () => ({ accessToken: "login-token" } satisfies OAuthTokenData)),
|
||||||
|
refresh,
|
||||||
|
}),
|
||||||
|
})
|
||||||
|
const callTool = mock(async () => {
|
||||||
|
if (callTool.mock.calls.length === 1) {
|
||||||
|
throw new Error("401 Unauthorized")
|
||||||
|
}
|
||||||
|
|
||||||
|
return { content: [{ type: "text", text: "success" }] }
|
||||||
|
})
|
||||||
|
mockGetOrCreateClientWithRetryImpl.mockResolvedValue({ callTool, close: mock(async () => {}) })
|
||||||
|
|
||||||
|
// when
|
||||||
|
const result = await manager.callTool(createInfo(), createContext(), "test-tool", {})
|
||||||
|
|
||||||
|
// then
|
||||||
|
expect(result).toEqual([{ type: "text", text: "success" }])
|
||||||
|
expect(refresh).toHaveBeenCalledTimes(1)
|
||||||
|
expect(callTool).toHaveBeenCalledTimes(2)
|
||||||
|
})
|
||||||
|
|
||||||
|
it("retries the operation after a 403 refresh succeeds without step-up scope", async () => {
|
||||||
|
// given
|
||||||
|
const { SkillMcpManager } = await importFreshManagerModule()
|
||||||
|
const refresh = mock(async () => ({ accessToken: "refreshed-token" } satisfies OAuthTokenData))
|
||||||
|
const manager = new SkillMcpManager({
|
||||||
|
createOAuthProvider: () => ({
|
||||||
|
tokens: () => ({ accessToken: "stale-token", refreshToken: "refresh-token" }),
|
||||||
|
login: mock(async () => ({ accessToken: "login-token" } satisfies OAuthTokenData)),
|
||||||
|
refresh,
|
||||||
|
}),
|
||||||
|
})
|
||||||
|
const callTool = mock(async () => {
|
||||||
|
if (callTool.mock.calls.length === 1) {
|
||||||
|
throw new Error("403 Forbidden")
|
||||||
|
}
|
||||||
|
|
||||||
|
return { content: [{ type: "text", text: "success" }] }
|
||||||
|
})
|
||||||
|
mockGetOrCreateClientWithRetryImpl.mockResolvedValue({ callTool, close: mock(async () => {}) })
|
||||||
|
|
||||||
|
// when
|
||||||
|
const result = await manager.callTool(createInfo(), createContext(), "test-tool", {})
|
||||||
|
|
||||||
|
// then
|
||||||
|
expect(result).toEqual([{ type: "text", text: "success" }])
|
||||||
|
expect(refresh).toHaveBeenCalledTimes(1)
|
||||||
|
expect(callTool).toHaveBeenCalledTimes(2)
|
||||||
|
})
|
||||||
|
|
||||||
|
it("propagates the auth error without retry when refresh fails", async () => {
|
||||||
|
// given
|
||||||
|
const { SkillMcpManager } = await importFreshManagerModule()
|
||||||
|
const refresh = mock(async () => {
|
||||||
|
throw new Error("refresh failed")
|
||||||
|
})
|
||||||
|
const manager = new SkillMcpManager({
|
||||||
|
createOAuthProvider: () => ({
|
||||||
|
tokens: () => ({ accessToken: "stale-token", refreshToken: "refresh-token" }),
|
||||||
|
login: mock(async () => ({ accessToken: "login-token" } satisfies OAuthTokenData)),
|
||||||
|
refresh,
|
||||||
|
}),
|
||||||
|
})
|
||||||
|
const callTool = mock(async () => {
|
||||||
|
throw new Error("401 Unauthorized")
|
||||||
|
})
|
||||||
|
mockGetOrCreateClientWithRetryImpl.mockResolvedValue({ callTool, close: mock(async () => {}) })
|
||||||
|
|
||||||
|
// when / then
|
||||||
|
await expect(manager.callTool(createInfo(), createContext(), "test-tool", {})).rejects.toThrow("401 Unauthorized")
|
||||||
|
expect(refresh).toHaveBeenCalledTimes(1)
|
||||||
|
expect(callTool).toHaveBeenCalledTimes(1)
|
||||||
|
})
|
||||||
|
|
||||||
|
it("only attempts one refresh when the retried operation returns 401 again", async () => {
|
||||||
|
// given
|
||||||
|
const { SkillMcpManager } = await importFreshManagerModule()
|
||||||
|
const refresh = mock(async () => ({ accessToken: "refreshed-token" } satisfies OAuthTokenData))
|
||||||
|
const manager = new SkillMcpManager({
|
||||||
|
createOAuthProvider: () => ({
|
||||||
|
tokens: () => ({ accessToken: "stale-token", refreshToken: "refresh-token" }),
|
||||||
|
login: mock(async () => ({ accessToken: "login-token" } satisfies OAuthTokenData)),
|
||||||
|
refresh,
|
||||||
|
}),
|
||||||
|
})
|
||||||
|
const callTool = mock(async () => {
|
||||||
|
throw new Error("401 Unauthorized")
|
||||||
|
})
|
||||||
|
mockGetOrCreateClientWithRetryImpl.mockResolvedValue({ callTool, close: mock(async () => {}) })
|
||||||
|
|
||||||
|
// when / then
|
||||||
|
await expect(manager.callTool(createInfo(), createContext(), "test-tool", {})).rejects.toThrow("401 Unauthorized")
|
||||||
|
expect(refresh).toHaveBeenCalledTimes(1)
|
||||||
|
expect(callTool).toHaveBeenCalledTimes(2)
|
||||||
|
})
|
||||||
|
})
|
||||||
@@ -4,7 +4,7 @@ import type { ClaudeCodeMcpServer } from "../claude-code-mcp-loader/types"
|
|||||||
import { McpOAuthProvider } from "../mcp-oauth/provider"
|
import { McpOAuthProvider } from "../mcp-oauth/provider"
|
||||||
import { disconnectAll, disconnectSession, forceReconnect } from "./cleanup"
|
import { disconnectAll, disconnectSession, forceReconnect } from "./cleanup"
|
||||||
import { getOrCreateClient, getOrCreateClientWithRetryImpl } from "./connection"
|
import { getOrCreateClient, getOrCreateClientWithRetryImpl } from "./connection"
|
||||||
import { handleStepUpIfNeeded } from "./oauth-handler"
|
import { handlePostRequestAuthError, handleStepUpIfNeeded } from "./oauth-handler"
|
||||||
import type {
|
import type {
|
||||||
OAuthProviderFactory,
|
OAuthProviderFactory,
|
||||||
SkillMcpClientInfo,
|
SkillMcpClientInfo,
|
||||||
@@ -110,6 +110,7 @@ export class SkillMcpManager {
|
|||||||
): Promise<T> {
|
): Promise<T> {
|
||||||
const maxRetries = 3
|
const maxRetries = 3
|
||||||
let lastError: Error | null = null
|
let lastError: Error | null = null
|
||||||
|
const refreshAttempted = new Set<string>()
|
||||||
|
|
||||||
for (let attempt = 1; attempt <= maxRetries; attempt++) {
|
for (let attempt = 1; attempt <= maxRetries; attempt++) {
|
||||||
try {
|
try {
|
||||||
@@ -130,6 +131,17 @@ export class SkillMcpManager {
|
|||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const postRequestRefreshHandled = await handlePostRequestAuthError({
|
||||||
|
error: lastError,
|
||||||
|
config,
|
||||||
|
authProviders: this.state.authProviders,
|
||||||
|
createOAuthProvider: this.state.createOAuthProvider,
|
||||||
|
refreshAttempted,
|
||||||
|
})
|
||||||
|
if (postRequestRefreshHandled) {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
if (!errorMessage.includes("not connected")) {
|
if (!errorMessage.includes("not connected")) {
|
||||||
throw lastError
|
throw lastError
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user