Ask Claude to extend tokenExchangeCallback to allow customizing the access token TTL.

Transcript (same as previous commit): https://claude-workerd-transcript.pages.dev/oauth-provider-token-exchange-callback2

Author: kentonv

README.md

@@ -242,7 +242,9 @@ new OAuthProvider({
         newProps: {
           ...options.props,
           upstreamRefreshToken: upstreamTokens.refresh_token || options.props.upstreamRefreshToken
-        }
+        },
+        // Optionally override the default access token TTL to match the upstream token
+        accessTokenTTL: upstreamTokens.expires_in
       };
     }
   }
@@ -253,8 +255,11 @@ The callback can:
 - Return both `accessTokenProps` and `newProps` to update both
 - Return only `accessTokenProps` to update just the current access token
 - Return only `newProps` to update both the grant and access token (the access token inherits these props)
+- Return `accessTokenTTL` to override the default TTL for this specific access token
 - Return nothing to keep the original props unchanged
 
+The `accessTokenTTL` override is particularly useful when the application is also an OAuth client to another service and wants to match its access token TTL to the upstream access token TTL. This helps prevent situations where the downstream token is still valid but the upstream token has expired.
+
 The `props` values are end-to-end encrypted, so they can safely contain sensitive information.
 
 ## Written by Claude

__tests__/oauth-provider.test.ts

@@ -1539,6 +1539,111 @@ describe('OAuthProvider', () => {
       });
     });
 
+    it('should allow customizing access token TTL via callback', async () => {
+      // Create a provider with a callback that customizes TTL
+      const customTtlCallback = async (options: any) => {
+        if (options.grantType === 'refresh_token') {
+          // Return custom TTL for the access token
+          return {
+            accessTokenProps: { ...options.props, customTtl: true },
+            accessTokenTTL: 7200 // 2 hours instead of default
+          };
+        }
+        return undefined;
+      };
+
+      const customTtlProvider = new OAuthProvider({
+        apiRoute: ['/api/'],
+        apiHandler: TestApiHandler,
+        defaultHandler: testDefaultHandler,
+        authorizeEndpoint: '/authorize',
+        tokenEndpoint: '/oauth/token',
+        clientRegistrationEndpoint: '/oauth/register',
+        scopesSupported: ['read', 'write'],
+        accessTokenTTL: 3600, // Default 1 hour
+        tokenExchangeCallback: customTtlCallback
+      });
+
+      // Create a client
+      const clientData = {
+        redirect_uris: ['https://client.example.com/callback'],
+        client_name: 'Custom TTL Test',
+        token_endpoint_auth_method: 'client_secret_basic'
+      };
+
+      const registerRequest = createMockRequest(
+        'https://example.com/oauth/register',
+        'POST',
+        { 'Content-Type': 'application/json' },
+        JSON.stringify(clientData)
+      );
+
+      const registerResponse = await customTtlProvider.fetch(registerRequest, mockEnv, mockCtx);
+      const client = await registerResponse.json();
+      const testClientId = client.client_id;
+      const testClientSecret = client.client_secret;
+      const testRedirectUri = 'https://client.example.com/callback';
+
+      // Get an auth code
+      const authRequest = createMockRequest(
+        `https://example.com/authorize?response_type=code&client_id=${testClientId}` +
+        `&redirect_uri=${encodeURIComponent(testRedirectUri)}` +
+        `&scope=read%20write&state=xyz123`
+      );
+
+      const authResponse = await customTtlProvider.fetch(authRequest, mockEnv, mockCtx);
+      const code = new URL(authResponse.headers.get('Location')!).searchParams.get('code')!;
+
+      // Exchange code for tokens
+      const params = new URLSearchParams();
+      params.append('grant_type', 'authorization_code');
+      params.append('code', code);
+      params.append('redirect_uri', testRedirectUri);
+      params.append('client_id', testClientId);
+      params.append('client_secret', testClientSecret);
+
+      const tokenRequest = createMockRequest(
+        'https://example.com/oauth/token',
+        'POST',
+        { 'Content-Type': 'application/x-www-form-urlencoded' },
+        params.toString()
+      );
+
+      const tokenResponse = await customTtlProvider.fetch(tokenRequest, mockEnv, mockCtx);
+      const tokens = await tokenResponse.json();
+
+      // Now do a refresh
+      const refreshParams = new URLSearchParams();
+      refreshParams.append('grant_type', 'refresh_token');
+      refreshParams.append('refresh_token', tokens.refresh_token);
+      refreshParams.append('client_id', testClientId);
+      refreshParams.append('client_secret', testClientSecret);
+
+      const refreshRequest = createMockRequest(
+        'https://example.com/oauth/token',
+        'POST',
+        { 'Content-Type': 'application/x-www-form-urlencoded' },
+        refreshParams.toString()
+      );
+
+      const refreshResponse = await customTtlProvider.fetch(refreshRequest, mockEnv, mockCtx);
+      const newTokens = await refreshResponse.json();
+
+      // Verify that the TTL is from the callback, not the default
+      expect(newTokens.expires_in).toBe(7200);
+      
+      // Verify the token contains our custom property
+      const apiRequest = createMockRequest(
+        'https://example.com/api/test',
+        'GET',
+        { 'Authorization': `Bearer ${newTokens.access_token}` }
+      );
+
+      const apiResponse = await customTtlProvider.fetch(apiRequest, mockEnv, mockCtx);
+      const apiData = await apiResponse.json();
+      expect(apiData.user.customTtl).toBe(true);
+    });
+
     it('should handle callback that returns undefined (keeping original props)', async () => {
       // Create a provider with a callback that returns undefined
       const noopCallback = async (options: any) => {

src/oauth-provider.ts

@@ -50,6 +50,14 @@ export interface TokenExchangeCallbackResult {
    * If not provided, the original props will be used.
    */
   newProps?: any;
+  
+  /**
+   * Override the default access token TTL (time-to-live) for this specific token.
+   * This is especially useful when the application is also an OAuth client to another service
+   * and wants to match its access token TTL to the upstream access token TTL.
+   * Value should be in seconds.
+   */
+  accessTokenTTL?: number;
 }
 
 /**
@@ -1227,7 +1235,9 @@ class OAuthProviderImpl {
     const refreshTokenId = await generateTokenId(refreshToken);
 
     const now = Math.floor(Date.now() / 1000);
-    const accessTokenExpiresAt = now + this.options.accessTokenTTL!;
+    
+    // Define the access token TTL, may be updated by callback if provided
+    let accessTokenTTL = this.options.accessTokenTTL!;
 
     // Get the encryption key for props by unwrapping it using the auth code
     const encryptionKey = await unwrapKeyWithToken(code, grantData.authCodeWrappedKey!);
@@ -1272,6 +1282,11 @@ class OAuthProviderImpl {
         if (callbackResult.accessTokenProps) {
           accessTokenProps = callbackResult.accessTokenProps;
         }
+        
+        // If accessTokenTTL was specified, use that for this token
+        if (callbackResult.accessTokenTTL !== undefined) {
+          accessTokenTTL = callbackResult.accessTokenTTL;
+        }
       }
 
       // Re-encrypt the potentially updated grant props
@@ -1291,6 +1306,9 @@ class OAuthProviderImpl {
       }
     }
 
+    // Calculate the access token expiration time (after callback might have updated TTL)
+    const accessTokenExpiresAt = now + accessTokenTTL;
+    
     // Wrap the keys for the new tokens
     const accessTokenWrappedKey = await wrapKeyWithToken(accessToken, accessTokenEncryptionKey);
     const refreshTokenWrappedKey = await wrapKeyWithToken(refreshToken, grantEncryptionKey);
@@ -1328,18 +1346,18 @@ class OAuthProviderImpl {
       }
     };
 
-    // Save access token with TTL
+    // Save access token with TTL (using the potentially callback-provided TTL)
     await env.OAUTH_KV.put(
       `token:${userId}:${grantId}:${accessTokenId}`,
       JSON.stringify(accessTokenData),
-      { expirationTtl: this.options.accessTokenTTL }
+      { expirationTtl: accessTokenTTL }
     );
 
     // Return the tokens
     return new Response(JSON.stringify({
       access_token: accessToken,
       token_type: 'bearer',
-      expires_in: this.options.accessTokenTTL,
+      expires_in: accessTokenTTL,
       refresh_token: refreshToken,
       scope: grantData.scope.join(' ')
     }), {
@@ -1424,7 +1442,9 @@ class OAuthProviderImpl {
     const newRefreshTokenId = await generateTokenId(newRefreshToken);
 
     const now = Math.floor(Date.now() / 1000);
-    const accessTokenExpiresAt = now + this.options.accessTokenTTL!;
+    
+    // Define the access token TTL, may be updated by callback if provided
+    let accessTokenTTL = this.options.accessTokenTTL!;
 
     // Determine which wrapped key to use for unwrapping
     let wrappedKeyToUse: string;
@@ -1479,6 +1499,11 @@ class OAuthProviderImpl {
         if (callbackResult.accessTokenProps) {
           accessTokenProps = callbackResult.accessTokenProps;
         }
+        
+        // If accessTokenTTL was specified, use that for this token
+        if (callbackResult.accessTokenTTL !== undefined) {
+          accessTokenTTL = callbackResult.accessTokenTTL;
+        }
       }
 
       // Only re-encrypt the grant props if they've changed
@@ -1508,6 +1533,9 @@ class OAuthProviderImpl {
       }
     }
 
+    // Calculate the access token expiration time (after callback might have updated TTL)
+    const accessTokenExpiresAt = now + accessTokenTTL;
+    
     // Wrap the key for both the new access token and refresh token
     const accessTokenWrappedKey = await wrapKeyWithToken(newAccessToken, accessTokenEncryptionKey);
     const newRefreshTokenWrappedKey = await wrapKeyWithToken(newRefreshToken, grantEncryptionKey);
@@ -1549,18 +1577,18 @@ class OAuthProviderImpl {
       }
     };
 
-    // Save access token with TTL
+    // Save access token with TTL (using the potentially callback-provided TTL)
     await env.OAUTH_KV.put(
       `token:${userId}:${grantId}:${accessTokenId}`,
       JSON.stringify(accessTokenData),
-      { expirationTtl: this.options.accessTokenTTL }
+      { expirationTtl: accessTokenTTL }
     );
 
     // Return the new access token and refresh token
     return new Response(JSON.stringify({
       access_token: newAccessToken,
       token_type: 'bearer',
-      expires_in: this.options.accessTokenTTL,
+      expires_in: accessTokenTTL,
       refresh_token: newRefreshToken,
       scope: grantData.scope.join(' ')
     }), {