Skip to content

Commit cac8a15

Browse files
encodedzencodedz
authored
feat: atp-client-opt-in-refresh-token-behaviour (#62)
* feat: atp-client-opt-in-refresh-token * Remove active updatetoken from both ends. * Shared. * r. * B. * Cleaned. * New client id issueing
1 parent 48e9d7d commit cac8a15

20 files changed

Lines changed: 618 additions & 426 deletions

File tree

‎examples/test-server/server.ts‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,10 @@
11
/**
22
* Minimal ATP Test Server for LangChain Integration Testing
3+
*
4+
* This server uses SHORT token TTL (5 seconds) to demonstrate and test
5+
* the automatic token refresh feature. In production, use longer TTLs.
6+
*
7+
* Run with: npx tsx server.ts
38
*/
49
import { config } from 'dotenv';
510
config({ path: '../../.env' });
@@ -12,12 +17,17 @@ if (!process.env.ATP_JWT_SECRET) {
1217
);
1318
}
1419

15-
import { AgentToolProtocolServer, loadOpenAPI } from '@mondaydotcomorg/atp-server';
20+
import { AgentToolProtocolServer } from '@mondaydotcomorg/atp-server';
1621

1722
async function main() {
18-
// Create ATP server
23+
// Create ATP server with SHORT token TTL for testing auto-refresh
24+
// In production, use longer TTLs (e.g., 1 hour default)
1925
const server = new AgentToolProtocolServer({
2026
execution: { timeout: 30000 },
27+
/* clientInit: {
28+
tokenTTL: 10000, // 10 seconds for testing
29+
tokenRotation: 5000,
30+
}*/
2131
});
2232

2333
// Register tools

‎examples/token-refresh/server.ts‎

Lines changed: 77 additions & 125 deletions
Original file line numberDiff line numberDiff line change
@@ -1,159 +1,111 @@
11
/**
22
* Token Refresh Example
33
*
4-
* Demonstrates how to use preRequestHook to automatically refresh
5-
* short-lived authentication tokens (e.g., 3-minute TTL bearer tokens)
4+
* Demonstrates ATP's built-in automatic token refresh feature.
5+
* The ATP client automatically refreshes tokens before they expire,
6+
* eliminating the need for manual token management in most cases.
7+
*
8+
* Run this with the test-server example (which has short token TTL):
9+
* 1. In one terminal: cd examples/test-server && npx tsx server.ts with these following config:
10+
* `clientInit: { tokenTTL: 5000, tokenRotation: 2500 }`
11+
* 2. In another terminal: cd examples/token-refresh && npx tsx server.ts
612
*/
713

814
import { AgentToolProtocolClient } from '@mondaydotcomorg/atp-client';
9-
import type { ClientHooks } from '@mondaydotcomorg/atp-client';
10-
11-
/**
12-
* Token Manager - Handles token lifecycle with caching
13-
*/
14-
class TokenManager {
15-
private currentToken: string | null = null;
16-
private tokenExpiry: number = 0;
17-
private refreshPromise: Promise<void> | null = null;
18-
19-
constructor(
20-
private authEndpoint: string,
21-
private credentials: { clientId: string; clientSecret: string }
22-
) {}
23-
24-
/**
25-
* Gets a valid token, refreshing if necessary
26-
* Thread-safe: multiple concurrent calls will share the same refresh
27-
*/
28-
async getValidToken(): Promise<string> {
29-
const now = Date.now();
30-
31-
// Refresh if expired or about to expire (30 second buffer)
32-
if (!this.currentToken || now >= this.tokenExpiry - 30000) {
33-
// Prevent multiple concurrent refreshes
34-
if (!this.refreshPromise) {
35-
this.refreshPromise = this.refreshToken().finally(() => {
36-
this.refreshPromise = null;
37-
});
38-
}
39-
await this.refreshPromise;
40-
}
41-
42-
return this.currentToken!;
43-
}
44-
45-
/**
46-
* Refreshes the token by calling the auth service
47-
*/
48-
private async refreshToken(): Promise<void> {
49-
console.log('[TokenManager] Refreshing token...');
50-
51-
try {
52-
const response = await fetch(this.authEndpoint, {
53-
method: 'POST',
54-
headers: { 'Content-Type': 'application/json' },
55-
body: JSON.stringify({
56-
grant_type: 'client_credentials',
57-
client_id: this.credentials.clientId,
58-
client_secret: this.credentials.clientSecret,
59-
}),
60-
});
6115

62-
if (!response.ok) {
63-
throw new Error(`Token refresh failed: ${response.status} ${response.statusText}`);
64-
}
65-
66-
const data: any = await response.json();
67-
this.currentToken = data.access_token;
68-
69-
// Calculate expiry with buffer
70-
const expiresIn = data.expires_in || 180; // Default to 3 minutes
71-
this.tokenExpiry = Date.now() + expiresIn * 1000;
72-
73-
console.log(`[TokenManager] Token refreshed. Expires in ${expiresIn} seconds`);
74-
} catch (error) {
75-
console.error('[TokenManager] Failed to refresh token:', error);
76-
throw error;
77-
}
78-
}
79-
80-
/**
81-
* Simulates getting an initial token (for demo purposes)
82-
*/
83-
async initialize(): Promise<void> {
84-
// For demo: simulate getting initial token
85-
this.currentToken = 'initial-token-' + Date.now();
86-
this.tokenExpiry = Date.now() + 180000; // 3 minutes
87-
console.log('[TokenManager] Initialized with demo token');
88-
}
89-
}
16+
const wait = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
9017

9118
/**
92-
* Example: Using ATP Client with automatic token refresh
19+
* Example: ATP Client with automatic token refresh (default behavior)
9320
*/
9421
async function main() {
95-
// Setup token manager
96-
const tokenManager = new TokenManager('https://auth.example.com/oauth/token', {
97-
clientId: process.env.CLIENT_ID || 'demo-client',
98-
clientSecret: process.env.CLIENT_SECRET || 'demo-secret',
99-
});
100-
101-
// Initialize token
102-
await tokenManager.initialize();
103-
104-
// Create hooks object with token refresh
105-
const hooks: ClientHooks = {
106-
preRequest: async (context) => {
107-
console.log(`[Hook] ${context.method} ${context.url}`);
108-
109-
// Get fresh token (will refresh if needed)
110-
const token = await tokenManager.getValidToken();
22+
console.log('='.repeat(60));
23+
console.log('ATP Automatic Token Refresh Demo');
24+
console.log('='.repeat(60));
11125

112-
// Return updated headers with fresh token
113-
return {
114-
headers: {
115-
...context.currentHeaders,
116-
Authorization: `Bearer ${token}`,
117-
'X-Request-Time': new Date().toISOString(),
118-
},
119-
};
120-
},
121-
};
122-
123-
// Create ATP client with hooks
26+
// Create ATP client - automatic token refresh is enabled by default
12427
const client = new AgentToolProtocolClient({
12528
baseUrl: process.env.ATP_SERVER_URL || 'http://localhost:3333',
126-
hooks,
29+
tokenRefresh: { enabled: true },
30+
hooks: {
31+
preRequest: async (context) => {
32+
console.log('[Hook] Request to:', context.url);
33+
return { headers: context.currentHeaders };
34+
},
35+
},
12736
});
12837

12938
console.log('\n=== Initializing ATP Client ===');
130-
await client.init({ name: 'token-refresh-example', version: '1.0.0' });
39+
const initResult = await client.init({ name: 'token-refresh-example', version: '1.0.0' });
40+
41+
console.log('Current time:', new Date());
42+
console.log('Client ID:', initResult.clientId);
43+
console.log('Token expires at:', new Date(initResult.expiresAt));
44+
console.log('Token rotates at:', new Date(initResult.tokenRotateAt));
45+
46+
const tokenTTL = initResult.expiresAt - Date.now();
47+
const rotateIn = initResult.tokenRotateAt - Date.now();
48+
console.log(`Token TTL: ${Math.round(tokenTTL / 1000)}s, Rotate in: ${Math.round(rotateIn / 1000)}s`);
13149

13250
console.log('\n=== Connecting to Server ===');
13351
await client.connect();
13452

135-
console.log('\n=== Executing Code ===');
136-
const result = await client.execute(`
137-
// Example code that uses ATP tools
53+
// First execution - should use original token
54+
console.log('\n=== First Execution (using original token) ===');
55+
const result1 = await client.execute(`
56+
const t = api.custom.add({ a: 2, b: 3 });
57+
const result = {
58+
timestamp: Date.now(),
59+
message: "First call with original token"
60+
};
61+
return result;
62+
`);
63+
console.log('Result:', JSON.stringify(result1.result, null, 2));
64+
65+
// Wait past the rotation time (test-server uses 2.5s rotation for 5s TTL)
66+
const waitTime = Math.max(rotateIn + 500, 30000);
67+
console.log(`\n=== Waiting ${waitTime / 1000}s to trigger token rotation ===`);
68+
await wait(waitTime);
69+
70+
// Second execution - should automatically refresh token before calling
71+
console.log('\n=== Second Execution (token should auto-refresh) ===');
72+
const result2 = await client.execute(`
13873
const result = {
13974
timestamp: Date.now(),
140-
message: "Hello from ATP with auto-refreshed token!"
75+
message: "Second call - token was auto-refreshed!"
14176
};
14277
return result;
14378
`);
79+
console.log('Result:', JSON.stringify(result2.result, null, 2));
14480

145-
console.log('\n=== Execution Result ===');
146-
console.log(JSON.stringify(result, null, 2));
81+
// Third execution - should still work
82+
console.log('\n=== Third Execution (continued use) ===');
83+
const result3 = await client.execute(`
84+
const result = {
85+
timestamp: Date.now(),
86+
message: "Third call - everything still works!"
87+
};
88+
return result;
89+
`);
90+
console.log('Result:', JSON.stringify(result3.result, null, 2));
14791

14892
console.log('\n=== Getting Server Info ===');
14993
const info = await client.getServerInfo();
15094
console.log('Server version:', info.version);
15195

152-
console.log('\n✅ All requests completed with automatic token refresh!');
96+
console.log('\n' + '='.repeat(60));
97+
console.log('✅ All requests completed with automatic token refresh!');
98+
console.log('='.repeat(60));
99+
console.log('\nKey takeaways:');
100+
console.log('1. Token refresh happens automatically before each request');
101+
console.log('2. No manual token management code needed');
102+
console.log('3. Requests never fail due to expired tokens');
103+
console.log('4. Works with short-lived tokens (even 5-second TTL)');
153104
}
154105

155-
// Run example
156-
main().catch((error) => {
157-
console.error('Error:', error);
158-
process.exit(1);
159-
});
106+
// Run examples
107+
main()
108+
.catch((error) => {
109+
console.error('Error:', error);
110+
process.exit(1);
111+
});

‎package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
"test": "NODE_ENV=test npm run jest -- --runInBand --forceExit --logHeapUsage",
1414
"test:unit": "npm run jest -- __tests__/unit --runInBand --logHeapUsage && cd packages/atp-compiler && npm test",
1515
"test:e2e": "npm run jest -- __tests__/e2e --runInBand --forceExit --testTimeout=120000 --logHeapUsage",
16+
"test:e2e:core": "npm run jest -- __tests__/e2e/core-flows --runInBand --forceExit --testTimeout=120000 --logHeapUsage",
1617
"test:e2e:checkpointer": "npm run jest -- __tests__/e2e/checkpoint --runInBand --forceExit --testTimeout=120000 --logHeapUsage",
1718
"test:e2e:runtime": "npm run jest -- __tests__/e2e/runtime --runInBand --forceExit --testTimeout=120000 --logHeapUsage",
1819
"test:e2e:server": "npm run jest -- __tests__/e2e/server --runInBand --forceExit --testTimeout=120000 --logHeapUsage",

‎packages/client/README.md‎

Lines changed: 7 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -210,24 +210,19 @@ const result = await client.execute({
210210

211211
### Pre-Request Hooks
212212

213-
Intercept and modify requests (e.g., token refresh):
213+
Intercept and modify requests (e.g., custom tokens set):
214214

215215
```typescript
216216
const client = new AgentToolProtocolClient({
217217
baseUrl: 'http://localhost:3333',
218218
hooks: {
219219
preRequest: async (context) => {
220-
// Refresh token if needed
221-
if (tokenExpired()) {
222-
const newToken = await refreshToken();
223-
return {
224-
headers: {
225-
...context.currentHeaders,
226-
Authorization: `Bearer ${newToken}`,
227-
},
228-
};
229-
}
230-
return {};
220+
return {
221+
headers: {
222+
...context.currentHeaders,
223+
'X-Custom-Token': `Bearer ${newToken}`,
224+
},
225+
};
231226
},
232227
},
233228
});

0 commit comments

Comments
 (0)