From f1e37b371c54a4d367a321b44ce599327e736b2b Mon Sep 17 00:00:00 2001 From: Kurt Overmier Date: Sat, 18 Apr 2026 03:00:37 -0500 Subject: [PATCH] fix(oauth): bake tenantId and tier into grant props at mint (closes #34) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both completeAuthorization call sites (auto-approve + consent-approve) now call provisionTenant once at grant mint and thread tenantId + tier through props, so resolveAuth reads them directly instead of re-calling AUTH_SERVICE on every request. Eliminates the grant-time vs. session-time tier drift that silently downgraded pro/enterprise callers to blessed templates on scaffold_create. Legacy grants that predate this change keep working — resolveAuth falls through to the provisionTenant call when props carry neither tenantId nor tier. A provisionTenant failure during mint is logged and the grant is minted without the fields, leaving the legacy fallback to cover it so a transient AUTH_SERVICE blip can't block OAuth consent. Regression guards: oauth-handler.test.ts asserts both mint sites bake the fields in and survive provisionTenant failure; gateway.test.ts asserts resolveAuth skips the provisionTenant call when props carry the baked values and still invokes it for the legacy cohort. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/gateway.ts | 21 ++++++- src/index.ts | 2 +- src/oauth-handler.ts | 35 +++++++++++ test/gateway.test.ts | 85 +++++++++++++++++++++++++ test/oauth-handler.test.ts | 125 +++++++++++++++++++++++++++++++++++++ 5 files changed, 265 insertions(+), 3 deletions(-) diff --git a/src/gateway.ts b/src/gateway.ts index 53ed4f9..2a6bc71 100644 --- a/src/gateway.ts +++ b/src/gateway.ts @@ -962,12 +962,16 @@ async function proxyToolCall( // The OAuth provider library then passes them to the apiHandler as ctx.props // on every authenticated call. `scopes` carries the actual scope grant from // the original OAuth request so that resolveAuth can enforce them at call -// time (see C-1a remediation in resolveAuth). +// time (see C-1a remediation in resolveAuth). `tenantId` and `tier` are +// baked in at grant mint (#34 remediation) so session init reads them from +// props instead of re-calling provisionTenant on every request. export interface OAuthProps { userId?: string; email?: string; name?: string; scopes?: string[]; + tenantId?: string; + tier?: string; } // ─── Main request handler ───────────────────────────────────── @@ -1070,7 +1074,20 @@ async function resolveAuth( effectiveScopes = await resolveLegacyGrantScopes(request, env); } - // Resolve tenant info from AUTH_SERVICE for proper tier + // #34 remediation: if the grant was minted with tenantId and tier baked + // into props, skip the per-request provisionTenant round trip. Falls + // through to the legacy path only for pre-#34 grants that carry neither. + if (oauthProps.tenantId && oauthProps.tier) { + return { + authenticated: true, + userId: oauthProps.userId, + tenantId: oauthProps.tenantId, + tier: oauthProps.tier as Tier, + scopes: effectiveScopes, + }; + } + + // Legacy cohort (pre-#34) — resolve tenant info from AUTH_SERVICE. try { const tenant = await env.AUTH_SERVICE.provisionTenant({ userId: oauthProps.userId, diff --git a/src/index.ts b/src/index.ts index 9d76393..7dee23f 100644 --- a/src/index.ts +++ b/src/index.ts @@ -37,7 +37,7 @@ const oauthProvider = new OAuthProvider({ // See src/oauth-handler.ts — we pass userId, email, name, and (per // C-1a remediation) the actual scopes the token was granted so // resolveAuth can enforce them at call time. - const oauthProps = (ctx as unknown as { props?: { userId?: string; email?: string; name?: string; scopes?: string[] } }).props; + const oauthProps = (ctx as unknown as { props?: { userId?: string; email?: string; name?: string; scopes?: string[]; tenantId?: string; tier?: string } }).props; const response = await handleMcpRequest(request, env, oauthProps); return addCorsHeaders(response); }, diff --git a/src/oauth-handler.ts b/src/oauth-handler.ts index 5eacefc..6cfb2db 100644 --- a/src/oauth-handler.ts +++ b/src/oauth-handler.ts @@ -485,6 +485,32 @@ function buildAuthorizeRedirect( return authorizeUrl.toString(); } +// --- Helper: resolve tenantId/tier to bake into grant props (#34) --- +// +// Called from both completeAuthorization sites (auto-approve + consent). +// On AUTH_SERVICE failure we return an empty object instead of throwing — +// that preserves the pre-#34 behavior where a grant was minted without +// tenant info and resolveAuth did the lookup per-request. The legacy +// fallback in gateway.ts:resolveAuth still handles that cohort, so we +// don't want a transient AUTH_SERVICE blip to block OAuth consent. +async function resolveTenantPropsForGrant( + env: GatewayEnv, + userId: string, +): Promise<{ tenantId?: string; tier?: string }> { + try { + const tenant = await env.AUTH_SERVICE.provisionTenant({ + userId, + source: 'oauth', + }); + return { tenantId: tenant.tenantId, tier: tenant.tier }; + } catch (err) { + console.error( + `[oauth-handler] provisionTenant failed for userId=${userId} during grant mint: ${err instanceof Error ? err.message : String(err)}`, + ); + return {}; + } +} + // --- "Coming soon" page when signups are gated --- function renderComingSoonPage(): string { @@ -636,6 +662,7 @@ async function handleGetAuthorize(request: Request, env: GatewayEnv): Promise { }); }); + describe('OAuth grant props (#34) — tenantId/tier baked in at mint', () => { + // The gateway's OAuth grant path used to mint grants with only + // {userId, email, name, scopes} and call provisionTenant on every + // request to resolve tenantId/tier. That re-resolution was both + // wasteful and a drift vector: if AUTH_SERVICE's view of the tenant + // tier changed between grant mint and session init, callers could + // silently downgrade (pro/enterprise → blessed on scaffold_create). + // + // Post-#34 grants bake tenantId and tier into props at mint time; + // resolveAuth reads them directly and must NOT call provisionTenant + // on the hot path. Legacy grants (no tenantId/tier in props) still + // fall through to the provisionTenant call for compatibility. + + it('uses tenantId/tier from oauthProps and does not call provisionTenant', async () => { + let provisionTenantCalls = 0; + const env = makeEnv({ + AUTH_SERVICE: { + ...mockAuthService(), + provisionTenant: async () => { + provisionTenantCalls += 1; + return { tenantId: 'drifted-tenant', userId: 'user-1', tier: 'free', delinquent: false, createdAt: '' }; + }, + }, + }); + const oauthProps = { + userId: 'user-1', + email: 'user@stackbilt.dev', + name: 'User', + scopes: ['generate', 'read'], + tenantId: 'tenant-mint-time', + tier: 'pro', + }; + + const req = new Request('https://mcp.stackbilt.dev/mcp', { + method: 'POST', + headers: { 'Content-Type': 'application/json', Accept: 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'initialize', + params: { protocolVersion: '2025-03-26', capabilities: {}, clientInfo: { name: 'test' } }, + }), + }); + const res = await handleMcpRequest(req, env, oauthProps); + + expect(res.status).toBe(200); + expect(provisionTenantCalls).toBe(0); + }); + + it('falls back to provisionTenant for legacy grants missing tenantId/tier', async () => { + let provisionTenantCalls = 0; + const env = makeEnv({ + AUTH_SERVICE: { + ...mockAuthService(), + provisionTenant: async () => { + provisionTenantCalls += 1; + return { tenantId: 'legacy-tenant', userId: 'user-1', tier: 'team', delinquent: false, createdAt: '' }; + }, + }, + }); + const oauthProps = { + userId: 'user-1', + email: 'user@stackbilt.dev', + name: 'User', + scopes: ['generate', 'read'], + // tenantId and tier deliberately omitted — legacy cohort + }; + + const req = new Request('https://mcp.stackbilt.dev/mcp', { + method: 'POST', + headers: { 'Content-Type': 'application/json', Accept: 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'initialize', + params: { protocolVersion: '2025-03-26', capabilities: {}, clientInfo: { name: 'test' } }, + }), + }); + const res = await handleMcpRequest(req, env, oauthProps); + + expect(res.status).toBe(200); + expect(provisionTenantCalls).toBe(1); + }); + }); + describe('DELETE session', () => { it('terminates an active session', async () => { const env = makeEnv(); diff --git a/test/oauth-handler.test.ts b/test/oauth-handler.test.ts index 3b6317d..7e57e98 100644 --- a/test/oauth-handler.test.ts +++ b/test/oauth-handler.test.ts @@ -593,6 +593,131 @@ describe('OAuth handler — #28 empty-scope consent flow', () => { }); }); +describe('OAuth handler — #34 tenantId/tier baked into grant props', () => { + // Before #34, completeAuthorization was called with props limited to + // {userId, email, name, scopes}. Every request then re-resolved tier + // via provisionTenant — a drift vector where grant-time vs. session-time + // tier divergence silently downgraded pro/enterprise callers to blessed + // engine templates on scaffold_create. + // + // Both grant-mint sites (auto-approve and consent-approve) must call + // provisionTenant once at mint time and bake tenantId + tier into props + // so the gateway can read them directly on every request. + + async function mkIdentityToken(clientId: string, redirectUri: string): Promise { + return signIdentityToken(TEST_SECRET, { + userId: 'user-42', + email: 'kurt@stackbilt.dev', + name: 'Kurt', + clientId, + redirectUri, + }); + } + + it('auto-approve bakes tenantId and tier into grant props', async () => { + const provider = mockOAuthProvider(); + const authService = mockAuthService({ + provisionTenant: vi.fn(async () => ({ + tenantId: 'tenant-pro', + userId: 'user-42', + tier: 'pro', + delinquent: false, + createdAt: '2026-01-01T00:00:00Z', + })), + }); + const env = makeEnv({ + OAUTH_PROVIDER: provider as unknown as GatewayEnv['OAUTH_PROVIDER'], + AUTH_SERVICE: authService as unknown as GatewayEnv['AUTH_SERVICE'], + }); + + const identityToken = await mkIdentityToken( + MOCK_AUTH_REQUEST.clientId, + MOCK_AUTH_REQUEST.redirectUri, + ); + const req = getRequest(`/authorize?identity_token=${encodeURIComponent(identityToken)}`); + const res = await callHandler(req, env); + + expect(res.status).toBe(302); + expect(authService.provisionTenant).toHaveBeenCalledWith({ userId: 'user-42', source: 'oauth' }); + expect(provider.completeAuthorization).toHaveBeenCalledWith( + expect.objectContaining({ + props: expect.objectContaining({ + userId: 'user-42', + tenantId: 'tenant-pro', + tier: 'pro', + }), + }), + ); + }); + + it('consent-approve bakes tenantId and tier into grant props', async () => { + const provider = mockOAuthProvider(); + const authService = mockAuthService({ + provisionTenant: vi.fn(async () => ({ + tenantId: 'tenant-enterprise', + userId: 'user-42', + tier: 'enterprise', + delinquent: false, + createdAt: '2026-01-01T00:00:00Z', + })), + }); + const env = makeEnv({ + OAUTH_PROVIDER: provider as unknown as GatewayEnv['OAUTH_PROVIDER'], + AUTH_SERVICE: authService as unknown as GatewayEnv['AUTH_SERVICE'], + }); + + const identityToken = await mkIdentityToken( + MOCK_AUTH_REQUEST.clientId, + MOCK_AUTH_REQUEST.redirectUri, + ); + const req = formRequest('/authorize', { + action: 'approve', + oauth_params: makeOAuthParamsB64(), + identity_token: identityToken, + }); + const res = await callHandler(req, env); + + expect(res.status).toBe(302); + expect(authService.provisionTenant).toHaveBeenCalledWith({ userId: 'user-42', source: 'oauth' }); + expect(provider.completeAuthorization).toHaveBeenCalledWith( + expect.objectContaining({ + props: expect.objectContaining({ + tenantId: 'tenant-enterprise', + tier: 'enterprise', + }), + }), + ); + }); + + it('mints a grant without tenantId/tier when provisionTenant fails (graceful degradation)', async () => { + const provider = mockOAuthProvider(); + const authService = mockAuthService({ + provisionTenant: vi.fn(async () => { + throw new Error('auth service unreachable'); + }), + }); + const env = makeEnv({ + OAUTH_PROVIDER: provider as unknown as GatewayEnv['OAUTH_PROVIDER'], + AUTH_SERVICE: authService as unknown as GatewayEnv['AUTH_SERVICE'], + }); + + const identityToken = await mkIdentityToken( + MOCK_AUTH_REQUEST.clientId, + MOCK_AUTH_REQUEST.redirectUri, + ); + const req = getRequest(`/authorize?identity_token=${encodeURIComponent(identityToken)}`); + const res = await callHandler(req, env); + + expect(res.status).toBe(302); + expect(provider.completeAuthorization).toHaveBeenCalledOnce(); + const call = (provider.completeAuthorization as ReturnType).mock.calls[0]?.[0] as { + props: { tenantId?: string; tier?: string }; + }; + expect(call.props.tenantId).toBeUndefined(); + expect(call.props.tier).toBeUndefined(); + }); +}); + describe('Cookie helpers (M-1 / HD-2)', () => { it('getCookie returns named value from multi-cookie header', () => { const req = new Request('https://x/', {