From 3bc36ad23674d608e10f803f801158a06aa45c8a Mon Sep 17 00:00:00 2001 From: Oded Goldglas Date: Sun, 1 Feb 2026 07:41:15 +0200 Subject: [PATCH 1/7] feat: atp-client-opt-in-refresh-token --- examples/test-server/server.ts | 13 +- examples/token-refresh/server.ts | 205 +++++++----------- packages/client/src/client.ts | 40 +++- .../client/src/core/in-process-session.ts | 141 ++++++++++-- packages/client/src/core/session.ts | 161 +++++++++++++- packages/server/src/client-sessions.ts | 161 +++++++------- packages/server/src/create-server.ts | 10 +- packages/server/src/handlers/token.handler.ts | 59 +++++ packages/server/src/http/request-handler.ts | 6 +- packages/server/src/http/router.ts | 2 + 10 files changed, 554 insertions(+), 244 deletions(-) create mode 100644 packages/server/src/handlers/token.handler.ts diff --git a/examples/test-server/server.ts b/examples/test-server/server.ts index b57420a..bdcd70e 100644 --- a/examples/test-server/server.ts +++ b/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' }); @@ -15,9 +20,15 @@ if (!process.env.ATP_JWT_SECRET) { import { AgentToolProtocolServer, loadOpenAPI } 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: { + // Short TTL for testing auto-refresh behavior + tokenTTL: 5000, // 5 seconds until token expires + tokenRotation: 2500, // Rotate at 2.5 seconds (halfway) + } }); // Register tools diff --git a/examples/token-refresh/server.ts b/examples/token-refresh/server.ts index 22f2de4..dfb5891 100644 --- a/examples/token-refresh/server.ts +++ b/examples/token-refresh/server.ts @@ -1,159 +1,114 @@ /** * 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 + * 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 | 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 { - 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 { - 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 { - // 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(); - - // Return updated headers with fresh token - return { - headers: { - ...context.currentHeaders, - Authorization: `Bearer ${token}`, - 'X-Request-Time': new Date().toISOString(), - }, - }; - }, - }; + console.log('='.repeat(60)); + console.log('ATP Automatic Token Refresh Demo'); + console.log('='.repeat(60)); - // 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, + 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 + console.log(await client.getTypeDefinitions()); + console.log(await client.exploreAPI('/custom')); + console.log(await client.searchAPI('add')); + console.log(await client.searchAPI('echo')); + + // 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, 10000); + 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); + }); diff --git a/packages/client/src/client.ts b/packages/client/src/client.ts index 7b104ac..2f4cc17 100644 --- a/packages/client/src/client.ts +++ b/packages/client/src/client.ts @@ -18,6 +18,7 @@ import { type ClientServiceProviders, type ClientHooks, type ISession, + type TokenRefreshConfig, ClientSession, InProcessSession, APIOperations, @@ -42,6 +43,7 @@ interface InProcessServer { handleExplore(ctx: unknown): Promise; handleExecute(ctx: unknown): Promise; handleResume(ctx: unknown, executionId: string): Promise; + handleTokenRefresh(ctx: unknown): Promise; } /** @@ -58,6 +60,14 @@ export interface AgentToolProtocolClientOptions { serviceProviders?: ClientServiceProviders; /** Optional hooks for intercepting and modifying client behavior */ hooks?: ClientHooks; + /** + * Configuration for automatic token refresh behavior. + * By default, ATP client automatically refreshes tokens before they expire. + * Set `enabled: false` to disable and manage tokens manually via preRequest hook. + * + * @default { enabled: true, bufferMs: 1000 } + */ + tokenRefresh?: Partial; } /** @@ -76,13 +86,25 @@ export class AgentToolProtocolClient { * * @example * ```typescript - * // HTTP mode + * // HTTP mode with automatic token refresh (default) * const client = new AgentToolProtocolClient({ * baseUrl: 'http://localhost:3333', * headers: { Authorization: 'Bearer token' }, + * }); + * + * // HTTP mode with custom token refresh behavior + * const client = new AgentToolProtocolClient({ + * baseUrl: 'http://localhost:3333', + * tokenRefresh: { enabled: true, bufferMs: 5000 }, // Refresh 5s before rotateAt + * }); + * + * // Disable automatic token refresh (use preRequest hook instead) + * const client = new AgentToolProtocolClient({ + * baseUrl: 'http://localhost:3333', + * tokenRefresh: { enabled: false }, * hooks: { * preRequest: async (context) => { - * const token = await refreshToken(); + * const token = await myCustomRefresh(); * return { headers: { ...context.currentHeaders, Authorization: `Bearer ${token}` } }; * } * } @@ -95,7 +117,7 @@ export class AgentToolProtocolClient { * ``` */ constructor(options: AgentToolProtocolClientOptions) { - const { baseUrl, server, headers, serviceProviders, hooks } = options; + const { baseUrl, server, headers, serviceProviders, hooks, tokenRefresh } = options; if (!baseUrl && !server) { throw new Error('Either baseUrl or server must be provided'); @@ -108,7 +130,7 @@ export class AgentToolProtocolClient { this.serviceProviders = new ServiceProviders(serviceProviders); if (server) { - this.inProcessSession = new InProcessSession(server); + this.inProcessSession = new InProcessSession(server, tokenRefresh); this.session = this.inProcessSession; this.apiOps = new APIOperations(this.session, this.inProcessSession); this.execOps = new ExecutionOperations( @@ -117,7 +139,7 @@ export class AgentToolProtocolClient { this.inProcessSession ); } else { - this.session = new ClientSession(baseUrl!, headers, hooks); + this.session = new ClientSession(baseUrl!, headers, hooks, tokenRefresh); this.apiOps = new APIOperations(this.session); this.execOps = new ExecutionOperations(this.session, this.serviceProviders); } @@ -150,6 +172,14 @@ export class AgentToolProtocolClient { return this.session.getClientId(); } + /** + * Configures automatic token refresh behavior. + * Call this to enable/disable auto-refresh or adjust timing. + */ + setTokenRefreshConfig(config: Partial): void { + this.session.setTokenRefreshConfig(config); + } + /** * Provides an LLM implementation for server to use during execution. */ diff --git a/packages/client/src/core/in-process-session.ts b/packages/client/src/core/in-process-session.ts index c765e92..481287e 100644 --- a/packages/client/src/core/in-process-session.ts +++ b/packages/client/src/core/in-process-session.ts @@ -1,5 +1,5 @@ import type { ClientToolDefinition } from '@mondaydotcomorg/atp-protocol'; -import type { ISession } from './session.js'; +import type { ISession, TokenRefreshConfig } from './session.js'; interface InProcessServer { start(): Promise; @@ -11,6 +11,7 @@ interface InProcessServer { handleExplore(ctx: InProcessRequestContext): Promise; handleExecute(ctx: InProcessRequestContext): Promise; handleResume(ctx: InProcessRequestContext, executionId: string): Promise; + handleTokenRefresh(ctx: InProcessRequestContext): Promise; } interface InProcessRequestContext { @@ -40,11 +41,21 @@ export class InProcessSession implements ISession { private server: InProcessServer; private clientId?: string; private clientToken?: string; + private tokenExpiresAt?: number; + private tokenRotateAt?: number; private initialized: boolean = false; private initPromise?: Promise; + private refreshPromise?: Promise; + private tokenRefreshConfig: TokenRefreshConfig = { + enabled: true, + bufferMs: 1000, + }; - constructor(server: InProcessServer) { + constructor(server: InProcessServer, tokenRefreshConfig?: Partial) { this.server = server; + if (tokenRefreshConfig) { + this.tokenRefreshConfig = { ...this.tokenRefreshConfig, ...tokenRefreshConfig }; + } } async init( @@ -62,15 +73,22 @@ export class InProcessSession implements ISession { return { clientId: this.clientId!, token: this.clientToken!, - expiresAt: 0, - tokenRotateAt: 0, + expiresAt: this.tokenExpiresAt!, + tokenRotateAt: this.tokenRotateAt!, }; } + let initResult: { + clientId: string; + token: string; + expiresAt: number; + tokenRotateAt: number; + }; + this.initPromise = (async () => { await this.server.start(); - const ctx = this.createContext({ + const ctx = await this.createContext({ method: 'POST', path: '/api/init', body: { @@ -89,17 +107,15 @@ export class InProcessSession implements ISession { this.clientId = result.clientId; this.clientToken = result.token; + this.tokenExpiresAt = result.expiresAt; + this.tokenRotateAt = result.tokenRotateAt; this.initialized = true; + initResult = result; })(); await this.initPromise; - return { - clientId: this.clientId!, - token: this.clientToken!, - expiresAt: 0, - tokenRotateAt: 0, - }; + return initResult!; } getClientId(): string { @@ -136,14 +152,93 @@ export class InProcessSession implements ISession { } updateToken(_response: Response): void { - // No-op for in-process - tokens are managed directly + // No-op for in-process - tokens are managed via refreshTokenIfNeeded + } + + /** + * Configure automatic token refresh behavior + */ + setTokenRefreshConfig(config: Partial): void { + this.tokenRefreshConfig = { ...this.tokenRefreshConfig, ...config }; + } + + /** + * Refresh token if needed (past rotateAt time or expired). + * For in-process mode, this directly calls the server's handleTokenRefresh. + * + * Note: Even expired tokens can be refreshed as long as the server session + * still exists. The server accepts expired JWTs for the refresh endpoint. + */ + async refreshTokenIfNeeded(): Promise { + // Skip if auto-refresh is disabled + if (!this.tokenRefreshConfig.enabled) { + return; + } + + // Skip if not initialized + if (!this.clientId || !this.clientToken) { + return; + } + + // Check if we need to refresh: + // - Past rotateAt time (proactive refresh), OR + // - Token has expired (reactive refresh - still allowed by server) + const now = Date.now(); + const needsRefresh = + (this.tokenRotateAt && now >= this.tokenRotateAt - this.tokenRefreshConfig.bufferMs) || + (this.tokenExpiresAt && now >= this.tokenExpiresAt); + + if (!needsRefresh) { + return; // Token is still fresh + } + + // Prevent concurrent refresh requests + if (this.refreshPromise) { + await this.refreshPromise; + return; + } + + this.refreshPromise = this.doRefreshToken(); + + try { + await this.refreshPromise; + } finally { + this.refreshPromise = undefined; + } + } + + /** + * Perform the actual token refresh via in-process server call + */ + private async doRefreshToken(): Promise { + const ctx = await this.createContext({ + method: 'POST', + path: '/api/token/refresh', + body: { clientId: this.clientId }, + }); + + const result = (await this.server.handleTokenRefresh(ctx)) as { + clientId: string; + token: string; + expiresAt: number; + tokenRotateAt: number; + }; + + this.clientToken = result.token; + this.tokenExpiresAt = result.expiresAt; + this.tokenRotateAt = result.tokenRotateAt; } async prepareHeaders( _method: string, - _url: string, + url: string, _body?: unknown ): Promise> { + // Refresh token if needed BEFORE preparing headers + // Skip for token refresh and init endpoints + if (!url.includes('/api/token/refresh') && !url.includes('/api/init')) { + await this.refreshTokenIfNeeded(); + } return this.getHeaders(); } @@ -154,7 +249,7 @@ export class InProcessSession implements ISession { }> { await this.ensureInitialized(); - const ctx = this.createContext({ + const ctx = await this.createContext({ method: 'GET', path: '/api/definitions', query: options?.apiGroups ? { apiGroups: options.apiGroups.join(',') } : {}, @@ -170,7 +265,7 @@ export class InProcessSession implements ISession { async getRuntimeDefinitions(options?: { apis?: string[] }): Promise { await this.ensureInitialized(); - const ctx = this.createContext({ + const ctx = await this.createContext({ method: 'GET', path: '/api/runtime', query: options?.apis?.length ? { apis: options.apis.join(',') } : {}, @@ -193,7 +288,7 @@ export class InProcessSession implements ISession { async search(query: string, options?: Record): Promise<{ results: unknown[] }> { await this.ensureInitialized(); - const ctx = this.createContext({ + const ctx = await this.createContext({ method: 'POST', path: '/api/search', body: { query, ...options }, @@ -205,7 +300,7 @@ export class InProcessSession implements ISession { async explore(path: string, options?: Record): Promise { await this.ensureInitialized(); - const ctx = this.createContext({ + const ctx = await this.createContext({ method: 'POST', path: '/api/explore', body: { path, ...options }, @@ -217,7 +312,7 @@ export class InProcessSession implements ISession { async execute(code: string, config?: Record): Promise { await this.ensureInitialized(); - const ctx = this.createContext({ + const ctx = await this.createContext({ method: 'POST', path: '/api/execute', body: { code, config }, @@ -229,7 +324,7 @@ export class InProcessSession implements ISession { async resume(executionId: string, callbackResult: unknown): Promise { await this.ensureInitialized(); - const ctx = this.createContext({ + const ctx = await this.createContext({ method: 'POST', path: `/api/resume/${executionId}`, body: { result: callbackResult }, @@ -244,7 +339,7 @@ export class InProcessSession implements ISession { ): Promise { await this.ensureInitialized(); - const ctx = this.createContext({ + const ctx = await this.createContext({ method: 'POST', path: `/api/resume/${executionId}`, body: { results: batchResults }, @@ -253,12 +348,12 @@ export class InProcessSession implements ISession { return await this.server.handleResume(ctx, executionId); } - private createContext(options: { + private async createContext(options: { method: string; path: string; query?: Record; body?: unknown; - }): InProcessRequestContext { + }): Promise { const noopLogger = { debug: () => {}, info: () => {}, @@ -270,7 +365,7 @@ export class InProcessSession implements ISession { method: options.method, path: options.path, query: options.query || {}, - headers: this.getHeaders(), + headers: await this.prepareHeaders(options.method, options.path, options.body), body: options.body, clientId: this.clientId, clientToken: this.clientToken, diff --git a/packages/client/src/core/session.ts b/packages/client/src/core/session.ts index 694d74b..9e6dda0 100644 --- a/packages/client/src/core/session.ts +++ b/packages/client/src/core/session.ts @@ -1,6 +1,16 @@ import type { ClientHooks } from './types.js'; import type { ClientToolDefinition } from '@mondaydotcomorg/atp-protocol'; +/** + * Configuration for automatic token refresh behavior + */ +export interface TokenRefreshConfig { + /** Enable automatic token refresh (default: true) */ + enabled: boolean; + /** Buffer time in ms before rotateAt to trigger refresh (default: 1000ms) */ + bufferMs: number; +} + export interface ISession { init( clientInfo?: { name?: string; version?: string; [key: string]: unknown }, @@ -18,6 +28,10 @@ export interface ISession { getBaseUrl(): string; updateToken(response: Response): void; prepareHeaders(method: string, url: string, body?: unknown): Promise>; + /** Refresh token if needed (past rotateAt time) */ + refreshTokenIfNeeded(): Promise; + /** Configure automatic token refresh behavior */ + setTokenRefreshConfig(config: Partial): void; } export class ClientSession implements ISession { @@ -25,13 +39,28 @@ export class ClientSession implements ISession { private customHeaders: Record; private clientId?: string; private clientToken?: string; + private tokenExpiresAt?: number; + private tokenRotateAt?: number; private initPromise?: Promise; + private refreshPromise?: Promise; private hooks?: ClientHooks; + private tokenRefreshConfig: TokenRefreshConfig = { + enabled: true, + bufferMs: 1000, + }; - constructor(baseUrl: string, headers?: Record, hooks?: ClientHooks) { + constructor( + baseUrl: string, + headers?: Record, + hooks?: ClientHooks, + tokenRefreshConfig?: Partial + ) { this.baseUrl = baseUrl; this.customHeaders = headers || {}; this.hooks = hooks; + if (tokenRefreshConfig) { + this.tokenRefreshConfig = { ...this.tokenRefreshConfig, ...tokenRefreshConfig }; + } } /** @@ -57,11 +86,18 @@ export class ClientSession implements ISession { return { clientId: this.clientId!, token: this.clientToken!, - expiresAt: 0, - tokenRotateAt: 0, + expiresAt: this.tokenExpiresAt!, + tokenRotateAt: this.tokenRotateAt!, }; } + let initResult: { + clientId: string; + token: string; + expiresAt: number; + tokenRotateAt: number; + }; + this.initPromise = (async () => { const url = `${this.baseUrl}/api/init`; const body = JSON.stringify({ @@ -69,6 +105,7 @@ export class ClientSession implements ISession { tools: tools || [], services, }); + const headers = await this.prepareHeaders('POST', url, body); const response = await fetch(url, { @@ -90,16 +127,14 @@ export class ClientSession implements ISession { this.clientId = data.clientId; this.clientToken = data.token; + this.tokenExpiresAt = data.expiresAt; + this.tokenRotateAt = data.tokenRotateAt; + initResult = data; })(); await this.initPromise; - return { - clientId: this.clientId!, - token: this.clientToken!, - expiresAt: 0, - tokenRotateAt: 0, - }; + return initResult!; } /** @@ -150,19 +185,125 @@ export class ClientSession implements ISession { */ updateToken(response: Response): void { const newToken = response.headers.get('X-ATP-Token'); + const newExpiresAt = response.headers.get('X-ATP-Token-Expires'); + if (newToken) { this.clientToken = newToken; } + if (newExpiresAt) { + this.tokenExpiresAt = parseInt(newExpiresAt, 10); + // Estimate tokenRotateAt as halfway between now and expiry + const now = Date.now(); + const ttl = this.tokenExpiresAt - now; + this.tokenRotateAt = now + ttl / 2; + } } /** - * Prepares headers for a request, calling preRequest hook if configured + * Configure automatic token refresh behavior + */ + setTokenRefreshConfig(config: Partial): void { + this.tokenRefreshConfig = { ...this.tokenRefreshConfig, ...config }; + } + + /** + * Refresh token if needed (past rotateAt time or expired). + * This is called automatically before requests when autoRefresh is enabled. + * Uses a shared promise to prevent concurrent refresh requests. + * + * Note: Even expired tokens can be refreshed as long as the server session + * still exists. The server accepts expired JWTs for the refresh endpoint. + */ + async refreshTokenIfNeeded(): Promise { + // Skip if auto-refresh is disabled + if (!this.tokenRefreshConfig.enabled) { + return; + } + + // Skip if not initialized + if (!this.clientId || !this.clientToken) { + return; + } + + // Check if we need to refresh: + // - Past rotateAt time (proactive refresh), OR + // - Token has expired (reactive refresh - still allowed by server) + const now = Date.now(); + const needsRefresh = + (this.tokenRotateAt && now >= this.tokenRotateAt - this.tokenRefreshConfig.bufferMs) || + (this.tokenExpiresAt && now >= this.tokenExpiresAt); + + if (!needsRefresh) { + return; // Token is still fresh + } + + // Prevent concurrent refresh requests + if (this.refreshPromise) { + await this.refreshPromise; + return; + } + + this.refreshPromise = this.doRefreshToken(); + + try { + await this.refreshPromise; + } finally { + this.refreshPromise = undefined; + } + } + + /** + * Perform the actual token refresh + */ + private async doRefreshToken(): Promise { + const url = `${this.baseUrl}/api/token/refresh`; + const body = JSON.stringify({ clientId: this.clientId }); + + // Use current token for auth, but don't recursively try to refresh + const headers: Record = { + 'Content-Type': 'application/json', + ...this.customHeaders, + 'X-Client-ID': this.clientId!, + Authorization: `Bearer ${this.clientToken}`, + }; + + const response = await fetch(url, { + method: 'POST', + headers, + body, + }); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error(`Token refresh failed: ${response.status} ${response.statusText} - ${errorText}`); + } + + const data = (await response.json()) as { + clientId: string; + token: string; + expiresAt: number; + tokenRotateAt: number; + }; + + this.clientToken = data.token; + this.tokenExpiresAt = data.expiresAt; + this.tokenRotateAt = data.tokenRotateAt; + } + + /** + * Prepares headers for a request, refreshing token if needed and calling preRequest hook if configured */ async prepareHeaders( method: string, url: string, body?: unknown ): Promise> { + // Refresh token if needed BEFORE preparing headers + // Skip for token refresh endpoint to avoid infinite recursion + if (!url.includes('/api/token/refresh') && !url.includes('/api/init')) { + await this.refreshTokenIfNeeded(); + } + let headers: Record = { 'Content-Type': 'application/json', ...this.customHeaders, diff --git a/packages/server/src/client-sessions.ts b/packages/server/src/client-sessions.ts index 91eebe6..b8eac34 100644 --- a/packages/server/src/client-sessions.ts +++ b/packages/server/src/client-sessions.ts @@ -45,22 +45,22 @@ export interface ClientInitResponse { clientId: string; token: string; expiresAt: number; - tokenRotateAt?: number; + tokenRotateAt: number; } /** * Client session manager with JWT-based authentication */ export class ClientSessionManager { - private cache?: CacheProvider; - private inMemorySessions: Map = new Map(); - private cleanupTimers: Map = new Map(); - private tokenTTL: number; - private jwtSecret: string; + private cache: CacheProvider; + private readonly tokenTTL: number; + private readonly tokenRotation: number; + private readonly jwtSecret: string; - constructor(options: { cache?: CacheProvider; tokenTTL: number; tokenRotation: number }) { + constructor(options: { cache: CacheProvider; tokenTTL: number; tokenRotation: number }) { this.cache = options.cache; this.tokenTTL = options.tokenTTL; + this.tokenRotation = options.tokenRotation; const secret = process.env.ATP_JWT_SECRET; if (!secret) { @@ -74,6 +74,19 @@ export class ClientSessionManager { } } + private ensureClientJWT(token: string, clientId: string, ignoreExpiration = false) { + const decoded = jwt.verify(token, this.jwtSecret, { + algorithms: ['HS256'], + ignoreExpiration, + }) as { clientId: string; type: string }; + + if (decoded.clientId !== clientId || decoded.type !== 'client') { + return false; + } + + return decoded; + } + /** * Initialize a new client session */ @@ -82,27 +95,22 @@ export class ClientSessionManager { const now = Date.now(); const expiresAt = now + this.tokenTTL; - const tokenRotateAt = now + this.tokenTTL / 2; + const tokenRotateAt = now + this.tokenRotation; const token = this.generateToken(clientId); const session: ClientSession = { clientId, createdAt: now, - expiresAt, + expiresAt: expiresAt, clientInfo: request.clientInfo, guidance: request.guidance, tools: request.tools || [], services: request.services, }; - if (this.cache) { - const ttlSeconds = Math.floor(this.tokenTTL / 1000); - await this.cache.set(`session:${clientId}`, session, ttlSeconds); - } else { - this.inMemorySessions.set(clientId, session); - this.scheduleCleanup(clientId, this.tokenTTL); - } + // Caching client session with default cache provider ttl, sessionExpiresAt is enforced programmatically on get. + await this.cache.set(`session:${clientId}`, session); return { clientId, @@ -117,44 +125,50 @@ export class ClientSessionManager { */ async verifyClient(clientId: string, token: string): Promise { try { - // Prevent algorithm confusion attacks - only allow HS256 - const decoded = jwt.verify(token, this.jwtSecret, { - algorithms: ['HS256'], - }) as { clientId: string; type: string }; - - if (decoded.clientId !== clientId || decoded.type !== 'client') { + if (!this.ensureClientJWT(token, clientId)) { return false; } const session = await this.getSession(clientId); return session !== null; - } catch (error) { + } catch { return false; } } /** - * Get client session + * Verify client token for refresh purposes - allows expired JWT tokens. + * This is used during token refresh when the JWT may have expired but + * the session still exists in cache. */ - async getSession(clientId: string): Promise { - let session: ClientSession | null = null; + async verifyClientForRefresh(clientId: string, token: string): Promise { + try { + // Verify token structure but ignore expiration - token refresh should work + // even if the JWT has expired, as long as the session still exists in cache + if (!this.ensureClientJWT(token, clientId, true)) { + return false; + } - if (this.cache) { - session = await this.cache.get(`session:${clientId}`); - } else { - session = this.inMemorySessions.get(clientId) || null; + // Check if session exists in cache - don't check session.expiresAt + const session = await this.cache.get(`session:${clientId}`); + return session !== null; + } catch { + return false; } + } + + /** + * Get client session + */ + async getSession(clientId: string): Promise { + const session = await this.cache.get(`session:${clientId}`); if (!session) { return null; } if (Date.now() > session.expiresAt) { - if (this.cache) { - await this.cache.delete(`session:${clientId}`); - } else { - this.inMemorySessions.delete(clientId); - } + await this.cache.delete(`session:${clientId}`); return null; } @@ -165,11 +179,7 @@ export class ClientSessionManager { * Revoke client session */ async revokeClient(clientId: string): Promise { - if (this.cache) { - await this.cache.delete(`session:${clientId}`); - } else { - this.inMemorySessions.delete(clientId); - } + await this.cache.delete(`session:${clientId}`); } /** @@ -192,56 +202,55 @@ export class ClientSessionManager { }, this.jwtSecret, { - expiresIn: 3600, + expiresIn: Math.ceil(this.tokenTTL / 1000), } ); } /** - * Schedule cleanup of expired in-memory session + * Refresh token for an existing client session. + * Returns new token credentials if session exists in cache. + * This works even if the session's expiresAt has passed - the refresh + * will update expiresAt to extend the session. */ - private scheduleCleanup(clientId: string, ttl: number): void { - const existingTimer = this.cleanupTimers.get(clientId); - if (existingTimer) { - clearTimeout(existingTimer); + async refreshToken(clientId: string): Promise { + // Get session directly from cache without expiry check + // Refresh should work even for "expired" sessions as long as they exist in cache + const session = await this.cache.get(`session:${clientId}`); + if (!session) { + return null; } - const timer = setTimeout(() => { - this.inMemorySessions.delete(clientId); - this.cleanupTimers.delete(clientId); - }, ttl); + const now = Date.now(); + const newExpiresAt = now + this.tokenTTL; + const newTokenRotateAt = now + this.tokenRotation; - this.cleanupTimers.set(clientId, timer); - } + // Update session with new expiry in cache + const updatedSession: ClientSession = { + ...session, + expiresAt: newExpiresAt, + }; - /** - * Manually cleanup a session (useful for tests and explicit cleanup) - */ - async cleanup(clientId: string): Promise { - const timer = this.cleanupTimers.get(clientId); - if (timer) { - clearTimeout(timer); - this.cleanupTimers.delete(clientId); - } + const ttlSeconds = Math.floor(this.tokenTTL / 1000); + await this.cache.set(`session:${clientId}`, updatedSession, ttlSeconds); - if (this.cache) { - await this.cache.delete(`session:${clientId}`); - } else { - this.inMemorySessions.delete(clientId); - } + const newToken = this.generateToken(clientId); + + return { + clientId, + token: newToken, + expiresAt: newExpiresAt, + tokenRotateAt: newTokenRotateAt, + }; } /** - * Cleanup all sessions (useful for shutdown) + * Get token TTL and rotation settings (useful for clients) */ - async cleanupAll(): Promise { - for (const timer of this.cleanupTimers.values()) { - clearTimeout(timer); - } - this.cleanupTimers.clear(); - - if (!this.cache) { - this.inMemorySessions.clear(); - } + getTokenSettings(): { tokenTTL: number; tokenRotation: number } { + return { + tokenTTL: this.tokenTTL, + tokenRotation: this.tokenRotation, + }; } } diff --git a/packages/server/src/create-server.ts b/packages/server/src/create-server.ts index 218ddb6..2d6cd15 100644 --- a/packages/server/src/create-server.ts +++ b/packages/server/src/create-server.ts @@ -39,6 +39,7 @@ import { handleSearch, handleSearchQuery } from './handlers/search.handler.js'; import { handleExplore } from './handlers/explorer.handler.js'; import { handleExecute } from './handlers/execute.handler.js'; import { handleResume } from './handlers/resume.handler.js'; +import { handleTokenRefresh } from './handlers/token.handler.js'; import { getDefinitions } from './handlers/definitions.handler.js'; import { shutdownAudit } from './middleware/audit.js'; import { @@ -397,8 +398,9 @@ export class AgentToolProtocolServer { async start(): Promise { if (this.isRunning) return; + // Cache provider is always set (defaults to MemoryCache in constructor) this.sessionManager = new ClientSessionManager({ - cache: this.cacheProvider, + cache: this.cacheProvider!, tokenTTL: this.config.clientInit.tokenTTL, tokenRotation: this.config.clientInit.tokenRotation, }); @@ -503,7 +505,6 @@ export class AgentToolProtocolServer { this.authProvider?.disconnect?.(), this.auditSink?.disconnect?.(), this.stateManager?.close(), - this.sessionManager?.cleanupAll(), ]); this.isRunning = false; @@ -611,6 +612,11 @@ export class AgentToolProtocolServer { ); } + async handleTokenRefresh(ctx: RequestContext): Promise { + if (!this.sessionManager) ctx.throw(503, 'Session manager not initialized'); + return await handleTokenRefresh(ctx, this.sessionManager); + } + /** * Update server components with new API groups (internal method) * @private diff --git a/packages/server/src/handlers/token.handler.ts b/packages/server/src/handlers/token.handler.ts new file mode 100644 index 0000000..f03fc19 --- /dev/null +++ b/packages/server/src/handlers/token.handler.ts @@ -0,0 +1,59 @@ +import type { RequestContext } from '../core/config.js'; +import type { ClientSessionManager } from '../client-sessions.js'; +import { log } from '@mondaydotcomorg/atp-runtime'; + +export interface TokenRefreshRequest { + clientId: string; +} + +export interface TokenRefreshResponse { + clientId: string; + token: string; + expiresAt: number; + tokenRotateAt: number; +} + +/** + * Handle token refresh requests. + * Allows clients to refresh their token, even if the JWT has expired. + * The session must still exist in the cache for refresh to succeed. + */ +export async function handleTokenRefresh( + ctx: RequestContext, + sessionManager: ClientSessionManager +): Promise { + // Get clientId from header or body + const clientId = ctx.clientId || (ctx.body as TokenRefreshRequest)?.clientId; + + if (!clientId) { + ctx.throw(400, 'Client ID is required for token refresh'); + } + + // Verify the current token (from Authorization header) + const authHeader = ctx.headers['authorization']; + if (!authHeader || !authHeader.startsWith('Bearer ')) { + ctx.throw(401, 'Bearer token required for refresh'); + } + + const currentToken = authHeader.substring(7); + + // Verify the token belongs to this client - allows expired JWT tokens + const isValid = await sessionManager.verifyClientForRefresh(clientId, currentToken); + if (!isValid) { + ctx.throw(401, 'Invalid token or session expired'); + } + + // Refresh the token + const refreshResult = await sessionManager.refreshToken(clientId); + if (!refreshResult) { + ctx.throw(401, 'Session not found or expired'); + } + + log.debug('Token refreshed', { + clientId, + newExpiresAt: refreshResult.expiresAt, + newRotateAt: refreshResult.tokenRotateAt, + }); + + return refreshResult; +} diff --git a/packages/server/src/http/request-handler.ts b/packages/server/src/http/request-handler.ts index 7a6ed8f..9a8de08 100644 --- a/packages/server/src/http/request-handler.ts +++ b/packages/server/src/http/request-handler.ts @@ -60,8 +60,9 @@ export async function handleHTTPRequest( try { if (ctx.clientId && deps.sessionManager && ctx.path !== '/api/init') { try { + const tokenSettings = deps.sessionManager.getTokenSettings(); const newToken = deps.sessionManager.generateToken(ctx.clientId); - const expiresAt = Date.now() + 60 * 60 * 1000; + const expiresAt = Date.now() + (tokenSettings.tokenTTL ?? 60 * 60 * 1000); headers.set('X-ATP-Token', newToken); headers.set('X-ATP-Token-Expires', expiresAt.toString()); @@ -79,8 +80,9 @@ export async function handleHTTPRequest( try { if (ctx.clientId && deps.sessionManager && ctx.path !== '/api/init') { try { + const tokenSettings = deps.sessionManager.getTokenSettings(); const newToken = deps.sessionManager.generateToken(ctx.clientId); - const expiresAt = Date.now() + 60 * 60 * 1000; + const expiresAt = Date.now() + tokenSettings.tokenTTL; headers.set('X-ATP-Token', newToken); headers.set('X-ATP-Token-Expires', expiresAt.toString()); diff --git a/packages/server/src/http/router.ts b/packages/server/src/http/router.ts index c722ff5..13fc029 100644 --- a/packages/server/src/http/router.ts +++ b/packages/server/src/http/router.ts @@ -24,6 +24,8 @@ export async function handleRoute( } else if (ctx.path.startsWith('/api/resume/') && ctx.method === 'POST') { const executionId = ctx.path.substring('/api/resume/'.length); ctx.responseBody = await server.handleResume(ctx, executionId); + } else if (ctx.path === '/api/token/refresh' && ctx.method === 'POST') { + ctx.responseBody = await server.handleTokenRefresh(ctx); } else { ctx.status = 404; ctx.responseBody = { error: 'Not found' }; From 89ea5ee2b2be5125ae54adbc99badc9854521171 Mon Sep 17 00:00:00 2001 From: Oded Goldglas Date: Sun, 1 Feb 2026 07:45:33 +0200 Subject: [PATCH 2/7] Remove active updatetoken from both ends. --- .../client/src/core/execution-operations.ts | 6 --- .../client/src/core/in-process-session.ts | 7 +--- packages/client/src/core/session.ts | 34 +---------------- packages/client/src/core/types.ts | 10 +++++ packages/server/src/http/request-handler.ts | 37 ------------------- 5 files changed, 13 insertions(+), 81 deletions(-) diff --git a/packages/client/src/core/execution-operations.ts b/packages/client/src/core/execution-operations.ts index 311a690..3ee52fa 100644 --- a/packages/client/src/core/execution-operations.ts +++ b/packages/client/src/core/execution-operations.ts @@ -169,8 +169,6 @@ export class ExecutionOperations { body, }); - this.session.updateToken(response); - if (!response.ok) { const error = (await response.json()) as { error: string }; throw new Error(`Execution failed: ${error.error || response.statusText}`); @@ -426,8 +424,6 @@ export class ExecutionOperations { body, }); - this.session.updateToken(response); - if (!response.ok) { const error = (await response.json()) as { error: string }; throw new Error(`Resume failed: ${error.error || response.statusText}`); @@ -480,8 +476,6 @@ export class ExecutionOperations { body, }); - this.session.updateToken(response); - if (!response.ok) { const error = (await response.json()) as { error: string }; throw new Error(`Batch resume failed: ${error.error || response.statusText}`); diff --git a/packages/client/src/core/in-process-session.ts b/packages/client/src/core/in-process-session.ts index 481287e..9da803c 100644 --- a/packages/client/src/core/in-process-session.ts +++ b/packages/client/src/core/in-process-session.ts @@ -1,5 +1,6 @@ import type { ClientToolDefinition } from '@mondaydotcomorg/atp-protocol'; -import type { ISession, TokenRefreshConfig } from './session.js'; +import type { TokenRefreshConfig } from './types.js'; +import type { ISession } from './session.js'; interface InProcessServer { start(): Promise; @@ -151,10 +152,6 @@ export class InProcessSession implements ISession { return ''; } - updateToken(_response: Response): void { - // No-op for in-process - tokens are managed via refreshTokenIfNeeded - } - /** * Configure automatic token refresh behavior */ diff --git a/packages/client/src/core/session.ts b/packages/client/src/core/session.ts index 9e6dda0..b57b7e7 100644 --- a/packages/client/src/core/session.ts +++ b/packages/client/src/core/session.ts @@ -1,16 +1,6 @@ -import type { ClientHooks } from './types.js'; +import type { ClientHooks, TokenRefreshConfig } from './types.js'; import type { ClientToolDefinition } from '@mondaydotcomorg/atp-protocol'; -/** - * Configuration for automatic token refresh behavior - */ -export interface TokenRefreshConfig { - /** Enable automatic token refresh (default: true) */ - enabled: boolean; - /** Buffer time in ms before rotateAt to trigger refresh (default: 1000ms) */ - bufferMs: number; -} - export interface ISession { init( clientInfo?: { name?: string; version?: string; [key: string]: unknown }, @@ -26,11 +16,8 @@ export interface ISession { ensureInitialized(): Promise; getHeaders(): Record; getBaseUrl(): string; - updateToken(response: Response): void; prepareHeaders(method: string, url: string, body?: unknown): Promise>; - /** Refresh token if needed (past rotateAt time) */ refreshTokenIfNeeded(): Promise; - /** Configure automatic token refresh behavior */ setTokenRefreshConfig(config: Partial): void; } @@ -180,25 +167,6 @@ export class ClientSession implements ISession { return this.baseUrl; } - /** - * Updates the client token from response headers (token refresh). - */ - updateToken(response: Response): void { - const newToken = response.headers.get('X-ATP-Token'); - const newExpiresAt = response.headers.get('X-ATP-Token-Expires'); - - if (newToken) { - this.clientToken = newToken; - } - if (newExpiresAt) { - this.tokenExpiresAt = parseInt(newExpiresAt, 10); - // Estimate tokenRotateAt as halfway between now and expiry - const now = Date.now(); - const ttl = this.tokenExpiresAt - now; - this.tokenRotateAt = now + ttl / 2; - } - } - /** * Configure automatic token refresh behavior */ diff --git a/packages/client/src/core/types.ts b/packages/client/src/core/types.ts index 4018776..e0b4f02 100644 --- a/packages/client/src/core/types.ts +++ b/packages/client/src/core/types.ts @@ -77,3 +77,13 @@ export interface ClientHooks { // onError?: ErrorHook; // onRetry?: RetryHook; } + +/** + * Configuration for automatic token refresh behavior + */ +export interface TokenRefreshConfig { + /** Enable automatic token refresh (default: true) */ + enabled: boolean; + /** Buffer time in ms before rotateAt to trigger refresh (default: 1000ms) */ + bufferMs: number; +} diff --git a/packages/server/src/http/request-handler.ts b/packages/server/src/http/request-handler.ts index 9a8de08..80c9de5 100644 --- a/packages/server/src/http/request-handler.ts +++ b/packages/server/src/http/request-handler.ts @@ -58,17 +58,6 @@ export async function handleHTTPRequest( await runMiddleware(ctx, deps.middleware, deps.routeHandler); try { - if (ctx.clientId && deps.sessionManager && ctx.path !== '/api/init') { - try { - const tokenSettings = deps.sessionManager.getTokenSettings(); - const newToken = deps.sessionManager.generateToken(ctx.clientId); - const expiresAt = Date.now() + (tokenSettings.tokenTTL ?? 60 * 60 * 1000); - - headers.set('X-ATP-Token', newToken); - headers.set('X-ATP-Token-Expires', expiresAt.toString()); - } catch (error) {} - } - const isStringResponse = typeof ctx.responseBody === 'string'; res.writeHead(ctx.status, { 'Content-Type': isStringResponse ? 'text/plain; charset=utf-8' : 'application/json', @@ -78,32 +67,6 @@ export async function handleHTTPRequest( } catch (writeError) {} } catch (error) { try { - if (ctx.clientId && deps.sessionManager && ctx.path !== '/api/init') { - try { - const tokenSettings = deps.sessionManager.getTokenSettings(); - const newToken = deps.sessionManager.generateToken(ctx.clientId); - const expiresAt = Date.now() + tokenSettings.tokenTTL; - - headers.set('X-ATP-Token', newToken); - headers.set('X-ATP-Token-Expires', expiresAt.toString()); - - log.debug('Token refresh headers set on error', { - clientId: ctx.clientId, - path: ctx.path, - hasSessionManager: !!deps.sessionManager, - headerCount: headers.size, - }); - } catch (tokenError) { - log.warn('Token refresh failed on error', { error: tokenError }); - } - } else { - log.debug('Token refresh skipped on error', { - hasClientId: !!ctx.clientId, - hasSessionManager: !!deps.sessionManager, - path: ctx.path, - }); - } - handleError(res, error as Error, randomUUID(), headers); } catch (handlerError) { try { From 52ba85e4b974e24bbfabe16974792c33610e6bd3 Mon Sep 17 00:00:00 2001 From: Oded Goldglas Date: Sun, 1 Feb 2026 07:52:35 +0200 Subject: [PATCH 3/7] Shared. --- package.json | 1 + packages/client/src/core/base-session.ts | 171 +++++++++++++++ .../client/src/core/in-process-session.ts | 154 ++++---------- packages/client/src/core/index.ts | 1 + packages/client/src/core/session.ts | 200 +++--------------- 5 files changed, 250 insertions(+), 277 deletions(-) create mode 100644 packages/client/src/core/base-session.ts diff --git a/package.json b/package.json index 36fef32..9ba461e 100644 --- a/package.json +++ b/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", diff --git a/packages/client/src/core/base-session.ts b/packages/client/src/core/base-session.ts new file mode 100644 index 0000000..750a46f --- /dev/null +++ b/packages/client/src/core/base-session.ts @@ -0,0 +1,171 @@ +import type { ClientToolDefinition } from '@mondaydotcomorg/atp-protocol'; +import type { TokenRefreshConfig } from './types.js'; + +/** + * Token credentials returned from init and refresh operations + */ +export interface TokenCredentials { + clientId: string; + token: string; + expiresAt: number; + tokenRotateAt: number; +} + +/** + * Session interface for ATP client operations + */ +export interface ISession { + init( + clientInfo?: { name?: string; version?: string; [key: string]: unknown }, + tools?: ClientToolDefinition[], + services?: { hasLLM: boolean; hasApproval: boolean; hasEmbedding: boolean; hasTools: boolean } + ): Promise; + getClientId(): string; + ensureInitialized(): Promise; + getHeaders(): Record; + getBaseUrl(): string; + prepareHeaders(method: string, url: string, body?: unknown): Promise>; + refreshTokenIfNeeded(): Promise; + setTokenRefreshConfig(config: Partial): void; +} + +/** + * Base session class with shared token management logic. + * Subclasses implement the transport-specific operations (HTTP vs in-process). + */ +export abstract class BaseSession implements ISession { + protected clientId?: string; + protected clientToken?: string; + protected tokenExpiresAt?: number; + protected tokenRotateAt?: number; + protected initPromise?: Promise; + protected refreshPromise?: Promise; + protected tokenRefreshConfig: TokenRefreshConfig = { + enabled: true, + bufferMs: 1000, + }; + + constructor(tokenRefreshConfig?: Partial) { + if (tokenRefreshConfig) { + this.tokenRefreshConfig = { ...this.tokenRefreshConfig, ...tokenRefreshConfig }; + } + } + + /** + * Initialize the session - must be implemented by subclass + */ + abstract init( + clientInfo?: { name?: string; version?: string; [key: string]: unknown }, + tools?: ClientToolDefinition[], + services?: { hasLLM: boolean; hasApproval: boolean; hasEmbedding: boolean; hasTools: boolean } + ): Promise; + + /** + * Get headers for requests - must be implemented by subclass + */ + abstract getHeaders(): Record; + + /** + * Get base URL - must be implemented by subclass + */ + abstract getBaseUrl(): string; + + /** + * Ensure client is initialized - must be implemented by subclass + */ + abstract ensureInitialized(): Promise; + + /** + * Prepare headers for a request - must be implemented by subclass + */ + abstract prepareHeaders(method: string, url: string, body?: unknown): Promise>; + + /** + * Perform the actual token refresh - must be implemented by subclass + */ + protected abstract doRefreshToken(): Promise; + + /** + * Gets the unique client ID. + */ + getClientId(): string { + if (!this.clientId) { + throw new Error('Client not initialized. Call init() first.'); + } + return this.clientId; + } + + /** + * Configure automatic token refresh behavior + */ + setTokenRefreshConfig(config: Partial): void { + this.tokenRefreshConfig = { ...this.tokenRefreshConfig, ...config }; + } + + /** + * Check if token needs refresh based on current time and config + */ + protected needsTokenRefresh(): boolean { + const now = Date.now(); + return ( + (this.tokenRotateAt !== undefined && now >= this.tokenRotateAt - this.tokenRefreshConfig.bufferMs) || + (this.tokenExpiresAt !== undefined && now >= this.tokenExpiresAt) + ); + } + + /** + * Refresh token if needed (past rotateAt time or expired). + * This is called automatically before requests when autoRefresh is enabled. + * Uses a shared promise to prevent concurrent refresh requests. + * + * Note: Even expired tokens can be refreshed as long as the server session + * still exists. The server accepts expired JWTs for the refresh endpoint. + */ + async refreshTokenIfNeeded(): Promise { + // Skip if auto-refresh is disabled + if (!this.tokenRefreshConfig.enabled) { + return; + } + + // Skip if not initialized + if (!this.clientId || !this.clientToken) { + return; + } + + // Check if we need to refresh + if (!this.needsTokenRefresh()) { + return; // Token is still fresh + } + + // Prevent concurrent refresh requests + if (this.refreshPromise) { + await this.refreshPromise; + return; + } + + this.refreshPromise = this.doRefreshToken(); + + try { + await this.refreshPromise; + } finally { + this.refreshPromise = undefined; + } + } + + /** + * Update token state after successful init or refresh + */ + protected updateTokenState(credentials: TokenCredentials): void { + this.clientId = credentials.clientId; + this.clientToken = credentials.token; + this.tokenExpiresAt = credentials.expiresAt; + this.tokenRotateAt = credentials.tokenRotateAt; + } + + /** + * Check if URL should skip token refresh (to avoid infinite recursion) + */ + protected shouldSkipRefreshForUrl(url: string): boolean { + return url.includes('/api/token/refresh') || url.includes('/api/init'); + } +} diff --git a/packages/client/src/core/in-process-session.ts b/packages/client/src/core/in-process-session.ts index 9da803c..9e4abf0 100644 --- a/packages/client/src/core/in-process-session.ts +++ b/packages/client/src/core/in-process-session.ts @@ -1,8 +1,11 @@ import type { ClientToolDefinition } from '@mondaydotcomorg/atp-protocol'; import type { TokenRefreshConfig } from './types.js'; -import type { ISession } from './session.js'; +import { BaseSession, type TokenCredentials } from './base-session.js'; -interface InProcessServer { +/** + * Server interface for in-process communication + */ +export interface InProcessServer { start(): Promise; handleInit(ctx: InProcessRequestContext): Promise; getDefinitions(ctx?: InProcessRequestContext): Promise; @@ -15,7 +18,10 @@ interface InProcessServer { handleTokenRefresh(ctx: InProcessRequestContext): Promise; } -interface InProcessRequestContext { +/** + * Request context for in-process server calls + */ +export interface InProcessRequestContext { method: string; path: string; query: Record; @@ -38,37 +44,27 @@ interface InProcessRequestContext { set(header: string, value: string): void; } -export class InProcessSession implements ISession { +/** + * In-process session for direct server communication without HTTP. + * Used when the client and server run in the same process. + */ +export class InProcessSession extends BaseSession { private server: InProcessServer; - private clientId?: string; - private clientToken?: string; - private tokenExpiresAt?: number; - private tokenRotateAt?: number; private initialized: boolean = false; - private initPromise?: Promise; - private refreshPromise?: Promise; - private tokenRefreshConfig: TokenRefreshConfig = { - enabled: true, - bufferMs: 1000, - }; constructor(server: InProcessServer, tokenRefreshConfig?: Partial) { + super(tokenRefreshConfig); this.server = server; - if (tokenRefreshConfig) { - this.tokenRefreshConfig = { ...this.tokenRefreshConfig, ...tokenRefreshConfig }; - } } + /** + * Initializes the client session with the in-process server. + */ async init( clientInfo?: { name?: string; version?: string; [key: string]: unknown }, tools?: ClientToolDefinition[], services?: { hasLLM: boolean; hasApproval: boolean; hasEmbedding: boolean; hasTools: boolean } - ): Promise<{ - clientId: string; - token: string; - expiresAt: number; - tokenRotateAt: number; - }> { + ): Promise { if (this.initPromise) { await this.initPromise; return { @@ -79,12 +75,7 @@ export class InProcessSession implements ISession { }; } - let initResult: { - clientId: string; - token: string; - expiresAt: number; - tokenRotateAt: number; - }; + let initResult: TokenCredentials; this.initPromise = (async () => { await this.server.start(); @@ -99,17 +90,8 @@ export class InProcessSession implements ISession { }, }); - const result = (await this.server.handleInit(ctx)) as { - clientId: string; - token: string; - expiresAt: number; - tokenRotateAt: number; - }; - - this.clientId = result.clientId; - this.clientToken = result.token; - this.tokenExpiresAt = result.expiresAt; - this.tokenRotateAt = result.tokenRotateAt; + const result = (await this.server.handleInit(ctx)) as TokenCredentials; + this.updateTokenState(result); this.initialized = true; initResult = result; })(); @@ -119,19 +101,18 @@ export class InProcessSession implements ISession { return initResult!; } - getClientId(): string { - if (!this.clientId) { - throw new Error('Client not initialized. Call init() first.'); - } - return this.clientId; - } - + /** + * Ensures the client is initialized before making requests. + */ async ensureInitialized(): Promise { if (!this.initialized) { throw new Error('Client not initialized. Call init() first.'); } } + /** + * Creates headers for in-process requests (lowercase for consistency with Node.js) + */ getHeaders(): Record { const headers: Record = { 'content-type': 'application/json', @@ -152,93 +133,39 @@ export class InProcessSession implements ISession { return ''; } - /** - * Configure automatic token refresh behavior - */ - setTokenRefreshConfig(config: Partial): void { - this.tokenRefreshConfig = { ...this.tokenRefreshConfig, ...config }; - } - - /** - * Refresh token if needed (past rotateAt time or expired). - * For in-process mode, this directly calls the server's handleTokenRefresh. - * - * Note: Even expired tokens can be refreshed as long as the server session - * still exists. The server accepts expired JWTs for the refresh endpoint. - */ - async refreshTokenIfNeeded(): Promise { - // Skip if auto-refresh is disabled - if (!this.tokenRefreshConfig.enabled) { - return; - } - - // Skip if not initialized - if (!this.clientId || !this.clientToken) { - return; - } - - // Check if we need to refresh: - // - Past rotateAt time (proactive refresh), OR - // - Token has expired (reactive refresh - still allowed by server) - const now = Date.now(); - const needsRefresh = - (this.tokenRotateAt && now >= this.tokenRotateAt - this.tokenRefreshConfig.bufferMs) || - (this.tokenExpiresAt && now >= this.tokenExpiresAt); - - if (!needsRefresh) { - return; // Token is still fresh - } - - // Prevent concurrent refresh requests - if (this.refreshPromise) { - await this.refreshPromise; - return; - } - - this.refreshPromise = this.doRefreshToken(); - - try { - await this.refreshPromise; - } finally { - this.refreshPromise = undefined; - } - } - /** * Perform the actual token refresh via in-process server call */ - private async doRefreshToken(): Promise { + protected async doRefreshToken(): Promise { const ctx = await this.createContext({ method: 'POST', path: '/api/token/refresh', body: { clientId: this.clientId }, }); - const result = (await this.server.handleTokenRefresh(ctx)) as { - clientId: string; - token: string; - expiresAt: number; - tokenRotateAt: number; - }; - - this.clientToken = result.token; - this.tokenExpiresAt = result.expiresAt; - this.tokenRotateAt = result.tokenRotateAt; + const result = (await this.server.handleTokenRefresh(ctx)) as TokenCredentials; + this.updateTokenState(result); } + /** + * Prepares headers for a request, refreshing token if needed + */ async prepareHeaders( _method: string, url: string, _body?: unknown ): Promise> { // Refresh token if needed BEFORE preparing headers - // Skip for token refresh and init endpoints - if (!url.includes('/api/token/refresh') && !url.includes('/api/init')) { + if (!this.shouldSkipRefreshForUrl(url)) { await this.refreshTokenIfNeeded(); } return this.getHeaders(); } + // ============================================ + // In-process specific methods for direct server calls + // ============================================ + async getDefinitions(options?: { apiGroups?: string[] }): Promise<{ typescript: string; version: string; @@ -345,6 +272,9 @@ export class InProcessSession implements ISession { return await this.server.handleResume(ctx, executionId); } + /** + * Creates a request context for in-process server calls + */ private async createContext(options: { method: string; path: string; diff --git a/packages/client/src/core/index.ts b/packages/client/src/core/index.ts index fb7baae..6ff1a38 100644 --- a/packages/client/src/core/index.ts +++ b/packages/client/src/core/index.ts @@ -1,3 +1,4 @@ +export * from './base-session.js'; export * from './session.js'; export * from './in-process-session.js'; export * from './api-operations.js'; diff --git a/packages/client/src/core/session.ts b/packages/client/src/core/session.ts index b57b7e7..a165121 100644 --- a/packages/client/src/core/session.ts +++ b/packages/client/src/core/session.ts @@ -1,40 +1,17 @@ -import type { ClientHooks, TokenRefreshConfig } from './types.js'; import type { ClientToolDefinition } from '@mondaydotcomorg/atp-protocol'; +import type { ClientHooks, TokenRefreshConfig } from './types.js'; +import { BaseSession, type TokenCredentials, type ISession } from './base-session.js'; -export interface ISession { - init( - clientInfo?: { name?: string; version?: string; [key: string]: unknown }, - tools?: ClientToolDefinition[], - services?: { hasLLM: boolean; hasApproval: boolean; hasEmbedding: boolean; hasTools: boolean } - ): Promise<{ - clientId: string; - token: string; - expiresAt: number; - tokenRotateAt: number; - }>; - getClientId(): string; - ensureInitialized(): Promise; - getHeaders(): Record; - getBaseUrl(): string; - prepareHeaders(method: string, url: string, body?: unknown): Promise>; - refreshTokenIfNeeded(): Promise; - setTokenRefreshConfig(config: Partial): void; -} +// Re-export for backward compatibility +export type { ISession, TokenCredentials, TokenRefreshConfig }; -export class ClientSession implements ISession { +/** + * HTTP-based session for connecting to remote ATP servers. + */ +export class ClientSession extends BaseSession { private baseUrl: string; private customHeaders: Record; - private clientId?: string; - private clientToken?: string; - private tokenExpiresAt?: number; - private tokenRotateAt?: number; - private initPromise?: Promise; - private refreshPromise?: Promise; private hooks?: ClientHooks; - private tokenRefreshConfig: TokenRefreshConfig = { - enabled: true, - bufferMs: 1000, - }; constructor( baseUrl: string, @@ -42,32 +19,22 @@ export class ClientSession implements ISession { hooks?: ClientHooks, tokenRefreshConfig?: Partial ) { + super(tokenRefreshConfig); this.baseUrl = baseUrl; this.customHeaders = headers || {}; this.hooks = hooks; - if (tokenRefreshConfig) { - this.tokenRefreshConfig = { ...this.tokenRefreshConfig, ...tokenRefreshConfig }; - } } /** * Initializes the client session with the server. * This MUST be called before any other operations. * The server generates and returns a unique client ID and token. - * @param clientInfo - Optional client information - * @param tools - Optional client tool definitions to register with the server - * @param services - Optional client service capabilities (LLM, approval, embedding) */ async init( clientInfo?: { name?: string; version?: string; [key: string]: unknown }, tools?: ClientToolDefinition[], services?: { hasLLM: boolean; hasApproval: boolean; hasEmbedding: boolean; hasTools: boolean } - ): Promise<{ - clientId: string; - token: string; - expiresAt: number; - tokenRotateAt: number; - }> { + ): Promise { if (this.initPromise) { await this.initPromise; return { @@ -78,12 +45,7 @@ export class ClientSession implements ISession { }; } - let initResult: { - clientId: string; - token: string; - expiresAt: number; - tokenRotateAt: number; - }; + let initResult: TokenCredentials; this.initPromise = (async () => { const url = `${this.baseUrl}/api/init`; @@ -93,7 +55,11 @@ export class ClientSession implements ISession { services, }); - const headers = await this.prepareHeaders('POST', url, body); + // Don't call prepareHeaders here to avoid token refresh before init + const headers: Record = { + 'Content-Type': 'application/json', + ...this.customHeaders, + }; const response = await fetch(url, { method: 'POST', @@ -105,17 +71,8 @@ export class ClientSession implements ISession { throw new Error(`Client initialization failed: ${response.status} ${response.statusText}`); } - const data = (await response.json()) as { - clientId: string; - token: string; - expiresAt: number; - tokenRotateAt: number; - }; - - this.clientId = data.clientId; - this.clientToken = data.token; - this.tokenExpiresAt = data.expiresAt; - this.tokenRotateAt = data.tokenRotateAt; + const data = (await response.json()) as TokenCredentials; + this.updateTokenState(data); initResult = data; })(); @@ -124,16 +81,6 @@ export class ClientSession implements ISession { return initResult!; } - /** - * Gets the unique client ID. - */ - getClientId(): string { - if (!this.clientId) { - throw new Error('Client not initialized. Call init() first.'); - } - return this.clientId; - } - /** * Ensures the client is initialized before making requests. */ @@ -168,62 +115,9 @@ export class ClientSession implements ISession { } /** - * Configure automatic token refresh behavior + * Perform the actual token refresh via HTTP */ - setTokenRefreshConfig(config: Partial): void { - this.tokenRefreshConfig = { ...this.tokenRefreshConfig, ...config }; - } - - /** - * Refresh token if needed (past rotateAt time or expired). - * This is called automatically before requests when autoRefresh is enabled. - * Uses a shared promise to prevent concurrent refresh requests. - * - * Note: Even expired tokens can be refreshed as long as the server session - * still exists. The server accepts expired JWTs for the refresh endpoint. - */ - async refreshTokenIfNeeded(): Promise { - // Skip if auto-refresh is disabled - if (!this.tokenRefreshConfig.enabled) { - return; - } - - // Skip if not initialized - if (!this.clientId || !this.clientToken) { - return; - } - - // Check if we need to refresh: - // - Past rotateAt time (proactive refresh), OR - // - Token has expired (reactive refresh - still allowed by server) - const now = Date.now(); - const needsRefresh = - (this.tokenRotateAt && now >= this.tokenRotateAt - this.tokenRefreshConfig.bufferMs) || - (this.tokenExpiresAt && now >= this.tokenExpiresAt); - - if (!needsRefresh) { - return; // Token is still fresh - } - - // Prevent concurrent refresh requests - if (this.refreshPromise) { - await this.refreshPromise; - return; - } - - this.refreshPromise = this.doRefreshToken(); - - try { - await this.refreshPromise; - } finally { - this.refreshPromise = undefined; - } - } - - /** - * Perform the actual token refresh - */ - private async doRefreshToken(): Promise { + protected async doRefreshToken(): Promise { const url = `${this.baseUrl}/api/token/refresh`; const body = JSON.stringify({ clientId: this.clientId }); @@ -246,16 +140,8 @@ export class ClientSession implements ISession { throw new Error(`Token refresh failed: ${response.status} ${response.statusText} - ${errorText}`); } - const data = (await response.json()) as { - clientId: string; - token: string; - expiresAt: number; - tokenRotateAt: number; - }; - - this.clientToken = data.token; - this.tokenExpiresAt = data.expiresAt; - this.tokenRotateAt = data.tokenRotateAt; + const data = (await response.json()) as TokenCredentials; + this.updateTokenState(data); } /** @@ -267,42 +153,26 @@ export class ClientSession implements ISession { body?: unknown ): Promise> { // Refresh token if needed BEFORE preparing headers - // Skip for token refresh endpoint to avoid infinite recursion - if (!url.includes('/api/token/refresh') && !url.includes('/api/init')) { + if (!this.shouldSkipRefreshForUrl(url)) { await this.refreshTokenIfNeeded(); } - let headers: Record = { - 'Content-Type': 'application/json', - ...this.customHeaders, - }; - - if (this.clientId) { - headers['X-Client-ID'] = this.clientId; - } - - if (this.clientToken) { - headers['Authorization'] = `Bearer ${this.clientToken}`; - } + let headers = this.getHeaders(); if (this.hooks?.preRequest) { - try { - const result = await this.hooks.preRequest({ - url, - method, - currentHeaders: headers, - body, - }); + const result = await this.hooks.preRequest({ + url, + method, + currentHeaders: headers, + body, + }); - if (result.abort) { - throw new Error(result.abortReason || 'Request aborted by preRequest hook'); - } + if (result.abort) { + throw new Error(result.abortReason || 'Request aborted by preRequest hook'); + } - if (result.headers) { - headers = result.headers; - } - } catch (error) { - throw error; + if (result.headers) { + headers = result.headers; } } From 2c757c93fe3e25c25239f908b805061151333d07 Mon Sep 17 00:00:00 2001 From: Oded Goldglas Date: Sun, 1 Feb 2026 08:04:27 +0200 Subject: [PATCH 4/7] r. --- packages/client/src/index.ts | 1 + packages/langchain/src/langgraph-client.ts | 2 +- packages/vercel-ai-sdk/src/types.ts | 14 +------------- 3 files changed, 3 insertions(+), 14 deletions(-) diff --git a/packages/client/src/index.ts b/packages/client/src/index.ts index 5fd0961..8c60d7a 100644 --- a/packages/client/src/index.ts +++ b/packages/client/src/index.ts @@ -20,4 +20,5 @@ export { } from './tools/index.js'; export type { AgentToolProtocolClientOptions } from './client.js'; export { InProcessSession } from './core/in-process-session.js'; +export type { InProcessServer } from './core/in-process-session.js'; export type { ISession } from './core/session.js'; diff --git a/packages/langchain/src/langgraph-client.ts b/packages/langchain/src/langgraph-client.ts index b172b15..93461ed 100644 --- a/packages/langchain/src/langgraph-client.ts +++ b/packages/langchain/src/langgraph-client.ts @@ -18,7 +18,7 @@ import type { ExecutionResult, ExecutionConfig, ClientTool } from '@mondaydotcom import { log } from '@mondaydotcomorg/atp-runtime'; import type { BaseChatModel } from '@langchain/core/language_models/chat_models'; import type { BaseMessage } from '@langchain/core/messages'; -import { HumanMessage, SystemMessage, AIMessage } from '@langchain/core/messages'; +import { HumanMessage, SystemMessage } from '@langchain/core/messages'; import type { Embeddings } from '@langchain/core/embeddings'; /** diff --git a/packages/vercel-ai-sdk/src/types.ts b/packages/vercel-ai-sdk/src/types.ts index 9d2dadb..97e1427 100644 --- a/packages/vercel-ai-sdk/src/types.ts +++ b/packages/vercel-ai-sdk/src/types.ts @@ -1,5 +1,5 @@ import type { ExecutionConfig, ClientTool } from '@mondaydotcomorg/atp-protocol'; -import type { ClientHooks } from '@mondaydotcomorg/atp-client'; +import type { ClientHooks, InProcessServer } from '@mondaydotcomorg/atp-client'; import type { UIMessageStreamWriter } from './event-adapter.js'; import type { generateText, generateObject } from 'ai'; @@ -66,18 +66,6 @@ export type GenerateObjectFunction = ( options: GenerateObjectOptions ) => Promise; -export interface InProcessServer { - start(): Promise; - handleInit(ctx: unknown): Promise; - getDefinitions(ctx?: unknown): Promise; - getRuntimeDefinitions(ctx?: unknown): Promise; - getInfo(): unknown; - handleSearch(ctx: unknown): Promise; - handleExplore(ctx: unknown): Promise; - handleExecute(ctx: unknown): Promise; - handleResume(ctx: unknown, executionId: string): Promise; -} - interface BaseClientOptions { model: any; embeddings?: EmbeddingProvider; From 477718ba9c10329a216b6db1373a9e92570a5aaf Mon Sep 17 00:00:00 2001 From: Oded Goldglas Date: Sun, 1 Feb 2026 08:08:02 +0200 Subject: [PATCH 5/7] B. --- packages/client/README.md | 19 +++++++------------ packages/client/src/core/api-operations.ts | 2 +- .../client/src/core/execution-operations.ts | 2 +- packages/client/src/core/session.ts | 5 +---- packages/client/src/index.ts | 2 +- 5 files changed, 11 insertions(+), 19 deletions(-) diff --git a/packages/client/README.md b/packages/client/README.md index 4bb73cf..f8c5243 100644 --- a/packages/client/README.md +++ b/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}`, + }, + }; }, }, }); diff --git a/packages/client/src/core/api-operations.ts b/packages/client/src/core/api-operations.ts index 136dda5..358df56 100644 --- a/packages/client/src/core/api-operations.ts +++ b/packages/client/src/core/api-operations.ts @@ -1,6 +1,6 @@ import type { SearchOptions, SearchResult, ExploreResult, ApiGroupRules } from '@mondaydotcomorg/atp-protocol'; import type { RuntimeAPIName } from '@mondaydotcomorg/atp-runtime'; -import type { ISession } from './session.js'; +import type { ISession } from './base-session.js'; import type { InProcessSession } from './in-process-session.js'; export class APIOperations { diff --git a/packages/client/src/core/execution-operations.ts b/packages/client/src/core/execution-operations.ts index 3ee52fa..04dbc96 100644 --- a/packages/client/src/core/execution-operations.ts +++ b/packages/client/src/core/execution-operations.ts @@ -1,7 +1,7 @@ import type { ExecutionResult, ExecutionConfig, ATPEvent } from '@mondaydotcomorg/atp-protocol'; import { ExecutionStatus, CallbackType } from '@mondaydotcomorg/atp-protocol'; import { log } from '@mondaydotcomorg/atp-runtime'; -import type { ISession } from './session.js'; +import type { ISession } from './base-session.js'; import type { InProcessSession } from './in-process-session.js'; import type { ServiceProviders } from './service-providers'; import { ClientCallbackError } from '../errors.js'; diff --git a/packages/client/src/core/session.ts b/packages/client/src/core/session.ts index a165121..1e5a95c 100644 --- a/packages/client/src/core/session.ts +++ b/packages/client/src/core/session.ts @@ -1,9 +1,6 @@ import type { ClientToolDefinition } from '@mondaydotcomorg/atp-protocol'; import type { ClientHooks, TokenRefreshConfig } from './types.js'; -import { BaseSession, type TokenCredentials, type ISession } from './base-session.js'; - -// Re-export for backward compatibility -export type { ISession, TokenCredentials, TokenRefreshConfig }; +import { BaseSession, type TokenCredentials } from './base-session.js'; /** * HTTP-based session for connecting to remote ATP servers. diff --git a/packages/client/src/index.ts b/packages/client/src/index.ts index 8c60d7a..b2a8eac 100644 --- a/packages/client/src/index.ts +++ b/packages/client/src/index.ts @@ -21,4 +21,4 @@ export { export type { AgentToolProtocolClientOptions } from './client.js'; export { InProcessSession } from './core/in-process-session.js'; export type { InProcessServer } from './core/in-process-session.js'; -export type { ISession } from './core/session.js'; +export type { ISession } from './core/base-session.js'; From af6433ec9442aca50398c6d71cd71a4e0302dcf6 Mon Sep 17 00:00:00 2001 From: Oded Goldglas Date: Sun, 1 Feb 2026 08:14:02 +0200 Subject: [PATCH 6/7] Cleaned. --- examples/test-server/server.ts | 7 +------ examples/token-refresh/server.ts | 10 +++------- packages/client/src/core/session.ts | 6 +----- 3 files changed, 5 insertions(+), 18 deletions(-) diff --git a/examples/test-server/server.ts b/examples/test-server/server.ts index bdcd70e..616f955 100644 --- a/examples/test-server/server.ts +++ b/examples/test-server/server.ts @@ -17,18 +17,13 @@ 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 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: { - // Short TTL for testing auto-refresh behavior - tokenTTL: 5000, // 5 seconds until token expires - tokenRotation: 2500, // Rotate at 2.5 seconds (halfway) - } }); // Register tools diff --git a/examples/token-refresh/server.ts b/examples/token-refresh/server.ts index dfb5891..00131de 100644 --- a/examples/token-refresh/server.ts +++ b/examples/token-refresh/server.ts @@ -6,7 +6,8 @@ * 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 + * 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 */ @@ -48,11 +49,6 @@ async function main() { console.log('\n=== Connecting to Server ==='); await client.connect(); - console.log(await client.getTypeDefinitions()); - console.log(await client.exploreAPI('/custom')); - console.log(await client.searchAPI('add')); - console.log(await client.searchAPI('echo')); - // First execution - should use original token console.log('\n=== First Execution (using original token) ==='); const result1 = await client.execute(` @@ -66,7 +62,7 @@ async function main() { 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, 10000); + const waitTime = Math.max(rotateIn + 500, 30000); console.log(`\n=== Waiting ${waitTime / 1000}s to trigger token rotation ===`); await wait(waitTime); diff --git a/packages/client/src/core/session.ts b/packages/client/src/core/session.ts index 1e5a95c..d67f01d 100644 --- a/packages/client/src/core/session.ts +++ b/packages/client/src/core/session.ts @@ -52,11 +52,7 @@ export class ClientSession extends BaseSession { services, }); - // Don't call prepareHeaders here to avoid token refresh before init - const headers: Record = { - 'Content-Type': 'application/json', - ...this.customHeaders, - }; + const headers = await this.prepareHeaders('POST', url, body); const response = await fetch(url, { method: 'POST', From 5900b0cad279665aa5f533e9600b3bfc7279ab6a Mon Sep 17 00:00:00 2001 From: Oded Goldglas Date: Sun, 1 Feb 2026 12:25:20 +0200 Subject: [PATCH 7/7] New client id issueing --- examples/test-server/server.ts | 4 ++++ examples/token-refresh/server.ts | 1 + packages/server/src/client-sessions.ts | 16 ++++++++++------ 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/examples/test-server/server.ts b/examples/test-server/server.ts index 616f955..c0f5ec7 100644 --- a/examples/test-server/server.ts +++ b/examples/test-server/server.ts @@ -24,6 +24,10 @@ async function main() { // 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 diff --git a/examples/token-refresh/server.ts b/examples/token-refresh/server.ts index 00131de..6e46bc0 100644 --- a/examples/token-refresh/server.ts +++ b/examples/token-refresh/server.ts @@ -26,6 +26,7 @@ async function main() { // Create ATP client - automatic token refresh is enabled by default const client = new AgentToolProtocolClient({ baseUrl: process.env.ATP_SERVER_URL || 'http://localhost:3333', + tokenRefresh: { enabled: true }, hooks: { preRequest: async (context) => { console.log('[Hook] Request to:', context.url); diff --git a/packages/server/src/client-sessions.ts b/packages/server/src/client-sessions.ts index b8eac34..08c3079 100644 --- a/packages/server/src/client-sessions.ts +++ b/packages/server/src/client-sessions.ts @@ -215,29 +215,33 @@ export class ClientSessionManager { */ async refreshToken(clientId: string): Promise { // Get session directly from cache without expiry check - // Refresh should work even for "expired" sessions as long as they exist in cache const session = await this.cache.get(`session:${clientId}`); if (!session) { + // Throw happens in handler return null; } + // Remove old client session entry. + await this.cache.delete(`session:${clientId}`);; + + const newClientId = this.generateClientId(); const now = Date.now(); const newExpiresAt = now + this.tokenTTL; const newTokenRotateAt = now + this.tokenRotation; - // Update session with new expiry in cache + // Update session with both new clientId and new expiresAt const updatedSession: ClientSession = { ...session, + clientId, expiresAt: newExpiresAt, }; - const ttlSeconds = Math.floor(this.tokenTTL / 1000); - await this.cache.set(`session:${clientId}`, updatedSession, ttlSeconds); + await this.cache.set(`session:${newClientId}`, updatedSession); - const newToken = this.generateToken(clientId); + const newToken = this.generateToken(newClientId); return { - clientId, + clientId: newClientId, token: newToken, expiresAt: newExpiresAt, tokenRotateAt: newTokenRotateAt,