feat/odedgo/atp-client-opt-in-refresh-token-behaviour

examples/test-server/server.ts

@@ -1,5 +1,10 @@
 /**
  * Minimal ATP Test Server for LangChain Integration Testing
+ *
+ * This server uses SHORT token TTL (5 seconds) to demonstrate and test
+ * the automatic token refresh feature. In production, use longer TTLs.
+ *
+ * Run with: npx tsx server.ts
  */
 import { config } from 'dotenv';
 config({ path: '../../.env' });
@@ -12,12 +17,17 @@ if (!process.env.ATP_JWT_SECRET) {
 	);
 }
 
-import { AgentToolProtocolServer, loadOpenAPI } from '@mondaydotcomorg/atp-server';
+import { AgentToolProtocolServer } from '@mondaydotcomorg/atp-server';
 
 async function main() {
-	// Create ATP server
+	// Create ATP server with SHORT token TTL for testing auto-refresh
+	// In production, use longer TTLs (e.g., 1 hour default)
 	const server = new AgentToolProtocolServer({
 		execution: { timeout: 30000 },
+/*		clientInit: {
+			tokenTTL: 10000, // 10 seconds for testing
+			tokenRotation: 5000,
+		}*/
 	});
 
 	// Register tools

examples/token-refresh/server.ts

@@ -1,159 +1,111 @@
 /**
  * Token Refresh Example
  *
- * Demonstrates how to use preRequestHook to automatically refresh
- * short-lived authentication tokens (e.g., 3-minute TTL bearer tokens)
+ * Demonstrates ATP's built-in automatic token refresh feature.
+ * The ATP client automatically refreshes tokens before they expire,
+ * eliminating the need for manual token management in most cases.
+ *
+ * Run this with the test-server example (which has short token TTL):
+ *   1. In one terminal: cd examples/test-server && npx tsx server.ts with these following config:
+ *   `clientInit: { tokenTTL: 5000, tokenRotation: 2500 }`
+ *   2. In another terminal: cd examples/token-refresh && npx tsx server.ts
  */
 
 import { AgentToolProtocolClient } from '@mondaydotcomorg/atp-client';
-import type { ClientHooks } from '@mondaydotcomorg/atp-client';
-
-/**
- * Token Manager - Handles token lifecycle with caching
- */
-class TokenManager {
-	private currentToken: string | null = null;
-	private tokenExpiry: number = 0;
-	private refreshPromise: Promise<void> | null = null;
-
-	constructor(
-		private authEndpoint: string,
-		private credentials: { clientId: string; clientSecret: string }
-	) {}
-
-	/**
-	 * Gets a valid token, refreshing if necessary
-	 * Thread-safe: multiple concurrent calls will share the same refresh
-	 */
-	async getValidToken(): Promise<string> {
-		const now = Date.now();
-
-		// Refresh if expired or about to expire (30 second buffer)
-		if (!this.currentToken || now >= this.tokenExpiry - 30000) {
-			// Prevent multiple concurrent refreshes
-			if (!this.refreshPromise) {
-				this.refreshPromise = this.refreshToken().finally(() => {
-					this.refreshPromise = null;
-				});
-			}
-			await this.refreshPromise;
-		}
-
-		return this.currentToken!;
-	}
-
-	/**
-	 * Refreshes the token by calling the auth service
-	 */
-	private async refreshToken(): Promise<void> {
-		console.log('[TokenManager] Refreshing token...');
-
-		try {
-			const response = await fetch(this.authEndpoint, {
-				method: 'POST',
-				headers: { 'Content-Type': 'application/json' },
-				body: JSON.stringify({
-					grant_type: 'client_credentials',
-					client_id: this.credentials.clientId,
-					client_secret: this.credentials.clientSecret,
-				}),
-			});
 
-			if (!response.ok) {
-				throw new Error(`Token refresh failed: ${response.status} ${response.statusText}`);
-			}
-
-			const data: any = await response.json();
-			this.currentToken = data.access_token;
-
-			// Calculate expiry with buffer
-			const expiresIn = data.expires_in || 180; // Default to 3 minutes
-			this.tokenExpiry = Date.now() + expiresIn * 1000;
-
-			console.log(`[TokenManager] Token refreshed. Expires in ${expiresIn} seconds`);
-		} catch (error) {
-			console.error('[TokenManager] Failed to refresh token:', error);
-			throw error;
-		}
-	}
-
-	/**
-	 * Simulates getting an initial token (for demo purposes)
-	 */
-	async initialize(): Promise<void> {
-		// For demo: simulate getting initial token
-		this.currentToken = 'initial-token-' + Date.now();
-		this.tokenExpiry = Date.now() + 180000; // 3 minutes
-		console.log('[TokenManager] Initialized with demo token');
-	}
-}
+const wait = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
 
 /**
- * Example: Using ATP Client with automatic token refresh
+ * Example: ATP Client with automatic token refresh (default behavior)
  */
 async function main() {
-	// Setup token manager
-	const tokenManager = new TokenManager('https://auth.example.com/oauth/token', {
-		clientId: process.env.CLIENT_ID || 'demo-client',
-		clientSecret: process.env.CLIENT_SECRET || 'demo-secret',
-	});
-
-	// Initialize token
-	await tokenManager.initialize();
-
-	// Create hooks object with token refresh
-	const hooks: ClientHooks = {
-		preRequest: async (context) => {
-			console.log(`[Hook] ${context.method} ${context.url}`);
-
-			// Get fresh token (will refresh if needed)
-			const token = await tokenManager.getValidToken();
+	console.log('='.repeat(60));
+	console.log('ATP Automatic Token Refresh Demo');
+	console.log('='.repeat(60));
 
-			// Return updated headers with fresh token
-			return {
-				headers: {
-					...context.currentHeaders,
-					Authorization: `Bearer ${token}`,
-					'X-Request-Time': new Date().toISOString(),
-				},
-			};
-		},
-	};
-
-	// Create ATP client with hooks
+	// Create ATP client - automatic token refresh is enabled by default
 	const client = new AgentToolProtocolClient({
 		baseUrl: process.env.ATP_SERVER_URL || 'http://localhost:3333',
-		hooks,
+		tokenRefresh: { enabled: true },
+		hooks: {
+			preRequest: async (context) => {
+				console.log('[Hook] Request to:', context.url);
+				return { headers: context.currentHeaders };
+			},
+		},
 	});
 
 	console.log('\n=== Initializing ATP Client ===');
-	await client.init({ name: 'token-refresh-example', version: '1.0.0' });
+	const initResult = await client.init({ name: 'token-refresh-example', version: '1.0.0' });
+
+	console.log('Current time:', new Date());
+	console.log('Client ID:', initResult.clientId);
+	console.log('Token expires at:', new Date(initResult.expiresAt));
+	console.log('Token rotates at:', new Date(initResult.tokenRotateAt));
+
+	const tokenTTL = initResult.expiresAt - Date.now();
+	const rotateIn = initResult.tokenRotateAt - Date.now();
+	console.log(`Token TTL: ${Math.round(tokenTTL / 1000)}s, Rotate in: ${Math.round(rotateIn / 1000)}s`);
 
 	console.log('\n=== Connecting to Server ===');
 	await client.connect();
 
-	console.log('\n=== Executing Code ===');
-	const result = await client.execute(`
-		// Example code that uses ATP tools
+	// First execution - should use original token
+	console.log('\n=== First Execution (using original token) ===');
+	const result1 = await client.execute(`
+		const t = api.custom.add({ a: 2, b: 3 });  
+		const result = {
+			timestamp: Date.now(),
+			message: "First call with original token"
+		};
+		return result;
+	`);
+	console.log('Result:', JSON.stringify(result1.result, null, 2));
+
+	// Wait past the rotation time (test-server uses 2.5s rotation for 5s TTL)
+	const waitTime = Math.max(rotateIn + 500, 30000);
+	console.log(`\n=== Waiting ${waitTime / 1000}s to trigger token rotation ===`);
+	await wait(waitTime);
+
+	// Second execution - should automatically refresh token before calling
+	console.log('\n=== Second Execution (token should auto-refresh) ===');
+	const result2 = await client.execute(`
 		const result = {
 			timestamp: Date.now(),
-			message: "Hello from ATP with auto-refreshed token!"
+			message: "Second call - token was auto-refreshed!"
 		};
 		return result;
 	`);
+	console.log('Result:', JSON.stringify(result2.result, null, 2));
 
-	console.log('\n=== Execution Result ===');
-	console.log(JSON.stringify(result, null, 2));
+	// Third execution - should still work
+	console.log('\n=== Third Execution (continued use) ===');
+	const result3 = await client.execute(`
+		const result = {
+			timestamp: Date.now(),
+			message: "Third call - everything still works!"
+		};
+		return result;
+	`);
+	console.log('Result:', JSON.stringify(result3.result, null, 2));
 
 	console.log('\n=== Getting Server Info ===');
 	const info = await client.getServerInfo();
 	console.log('Server version:', info.version);
 
-	console.log('\n✅ All requests completed with automatic token refresh!');
+	console.log('\n' + '='.repeat(60));
+	console.log('✅ All requests completed with automatic token refresh!');
+	console.log('='.repeat(60));
+	console.log('\nKey takeaways:');
+	console.log('1. Token refresh happens automatically before each request');
+	console.log('2. No manual token management code needed');
+	console.log('3. Requests never fail due to expired tokens');
+	console.log('4. Works with short-lived tokens (even 5-second TTL)');
 }
 
-// Run example
-main().catch((error) => {
-	console.error('Error:', error);
-	process.exit(1);
-});
+// Run examples
+main()
+	.catch((error) => {
+		console.error('Error:', error);
+		process.exit(1);
+	});

package.json

@@ -13,6 +13,7 @@
 		"test": "NODE_ENV=test npm run jest -- --runInBand --forceExit --logHeapUsage",
 		"test:unit": "npm run jest -- __tests__/unit --runInBand --logHeapUsage && cd packages/atp-compiler && npm test",
 		"test:e2e": "npm run jest -- __tests__/e2e --runInBand --forceExit --testTimeout=120000 --logHeapUsage",
+		"test:e2e:core": "npm run jest -- __tests__/e2e/core-flows --runInBand --forceExit --testTimeout=120000 --logHeapUsage",
 		"test:e2e:checkpointer": "npm run jest -- __tests__/e2e/checkpoint --runInBand --forceExit --testTimeout=120000 --logHeapUsage",
 		"test:e2e:runtime": "npm run jest -- __tests__/e2e/runtime --runInBand --forceExit --testTimeout=120000 --logHeapUsage",
 		"test:e2e:server": "npm run jest -- __tests__/e2e/server --runInBand --forceExit --testTimeout=120000 --logHeapUsage",

packages/client/README.md

@@ -210,24 +210,19 @@ const result = await client.execute({
 
 ### Pre-Request Hooks
 
-Intercept and modify requests (e.g., token refresh):
+Intercept and modify requests (e.g., custom tokens set):
 
 ```typescript
 const client = new AgentToolProtocolClient({
 	baseUrl: 'http://localhost:3333',
 	hooks: {
 		preRequest: async (context) => {
-			// Refresh token if needed
-			if (tokenExpired()) {
-				const newToken = await refreshToken();
-				return {
-					headers: {
-						...context.currentHeaders,
-						Authorization: `Bearer ${newToken}`,
-					},
-				};
-			}
-			return {};
+            return {
+                headers: {
+                    ...context.currentHeaders,
+                    'X-Custom-Token': `Bearer ${newToken}`,
+                },
+            };
 		},
 	},
 });