diff --git a/lib/mcp/gateway.ts b/lib/mcp/gateway.ts index 40af89c..aa41cb6 100644 --- a/lib/mcp/gateway.ts +++ b/lib/mcp/gateway.ts @@ -1,8 +1,7 @@ import { createGateway, - SignerRefreshRequired, type InferenceRequest, - type InferenceResult + type InferenceResult, } from "@pymthouse/gateway-web"; import { pymthouseSignerUrl } from "./env"; import type { McpPrincipal } from "./jwt"; @@ -12,28 +11,24 @@ export async function runInference( principal: McpPrincipal, request: InferenceRequest ): Promise { - let session = await resolveSignerSession(principal); + const session = await resolveSignerSession(principal); const timeoutMs = request.timeoutMs ?? request.timeout ?? 120_000; + let seeded: typeof session | null = session; - const attempt = async (signerJwt: string) => { - const gw = createGateway({ - signerUrl: session.signer_url || pymthouseSignerUrl(), - signerHeaders: { Authorization: `Bearer ${signerJwt}` }, - discoveryUrl: session.discovery_url, - insecureTls: true, - timeoutMs, - attributionSource: "pymthouse_gateway" - }); - return gw.runInference(request); - }; - - try { - return await attempt(session.access_token); - } catch (err) { - if (err instanceof SignerRefreshRequired) { - session = await resolveSignerSession(principal); - return attempt(session.access_token); - } - throw err; - } + const gw = createGateway({ + signerUrl: session.signer_url || pymthouseSignerUrl(), + discoveryUrl: session.discovery_url, + signerHeaders: async () => { + const next = seeded ?? (await resolveSignerSession(principal)); + seeded = null; + return { + headers: { Authorization: `Bearer ${next.access_token}` }, + expiresInSeconds: next.expires_in, + }; + }, + insecureTls: true, + timeoutMs, + attributionSource: "pymthouse_gateway", + }); + return gw.runInference(request); } diff --git a/package.json b/package.json index ddc61f2..d615cc1 100644 --- a/package.json +++ b/package.json @@ -16,7 +16,7 @@ "@auth0/nextjs-auth0": "^4.27.0", "@modelcontextprotocol/sdk": "^1.30.0", "@pymthouse/builder-sdk": "^0.6.5", - "@pymthouse/gateway-web": "^0.3.0", + "@pymthouse/gateway-web": "^0.4.0", "framer-motion": "^11.15.0", "geist": "^1.7.0", "jose": "^6.2.10",