feat/odedgo/reprise-token-refresh-init

examples/test-server/server.ts

@@ -24,10 +24,10 @@ async function main() {
 	// In production, use longer TTLs (e.g., 1 hour default)
 	const server = new AgentToolProtocolServer({
 		execution: { timeout: 30000 },
-/*		clientInit: {
+		clientInit: {
 			tokenTTL: 10000, // 10 seconds for testing
 			tokenRotation: 5000,
-		}*/
+		}
 	});
 
 	// Register tools

packages/client/src/client.ts

@@ -43,7 +43,6 @@ interface InProcessServer {
 	handleExplore(ctx: unknown): Promise<unknown>;
 	handleExecute(ctx: unknown): Promise<unknown>;
 	handleResume(ctx: unknown, executionId: string): Promise<unknown>;
-	handleTokenRefresh(ctx: unknown): Promise<unknown>;
 }
 
 /**

packages/client/src/core/base-session.ts

@@ -29,6 +29,15 @@ export interface ISession {
 	setTokenRefreshConfig(config: Partial<TokenRefreshConfig>): void;
 }
 
+/**
+ * Stored init parameters for re-initialization during token refresh.
+ */
+export interface StoredInitParams {
+	clientInfo?: { name?: string; version?: string; [key: string]: unknown };
+	tools?: ClientToolDefinition[];
+	services?: { hasLLM: boolean; hasApproval: boolean; hasEmbedding: boolean; hasTools: boolean };
+}
+
 /**
  * Base session class with shared token management logic.
  * Subclasses implement the transport-specific operations (HTTP vs in-process).
@@ -40,12 +49,13 @@ export abstract class BaseSession implements ISession {
 	protected tokenRotateAt?: number;
 	protected initPromise?: Promise<void>;
 	protected refreshPromise?: Promise<void>;
+	protected storedInitParams?: StoredInitParams;
 	protected tokenRefreshConfig: TokenRefreshConfig = {
 		enabled: true,
 		bufferMs: 1000,
 	};
 
-	constructor(tokenRefreshConfig?: Partial<TokenRefreshConfig>) {
+	protected constructor(tokenRefreshConfig?: Partial<TokenRefreshConfig>) {
 		if (tokenRefreshConfig) {
 			this.tokenRefreshConfig = { ...this.tokenRefreshConfig, ...tokenRefreshConfig };
 		}
@@ -81,9 +91,25 @@ export abstract class BaseSession implements ISession {
 	abstract prepareHeaders(method: string, url: string, body?: unknown): Promise<Record<string, string>>;
 
 	/**
-	 * Perform the actual token refresh - must be implemented by subclass
+	 * Perform token refresh by re-initializing the session.
+	 * Resets the init guard and calls init() again with the stored params,
+	 * effectively creating a fresh session without depending on the old
+	 * session still existing in the server's cache.
 	 */
-	protected abstract doRefreshToken(): Promise<void>;
+	protected async doRefreshToken(): Promise<void> {
+		if (!this.storedInitParams) {
+			throw new Error('Cannot refresh token: init params not stored. Was init() called?');
+		}
+
+		// Reset the init guard so init() runs a fresh handshake
+		this.initPromise = undefined;
+
+		await this.init(
+			this.storedInitParams.clientInfo,
+			this.storedInitParams.tools,
+			this.storedInitParams.services,
+		);
+	}
 
 	/**
 	 * Gets the unique client ID.
@@ -118,8 +144,8 @@ export abstract class BaseSession implements ISession {
 	 * 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.
+	 * Refresh works by re-initializing the session (calling init() again),
+	 * so it does not depend on the old session still existing in the server's cache.
 	 */
 	async refreshTokenIfNeeded(): Promise<void> {
 		// Skip if auto-refresh is disabled
@@ -163,9 +189,10 @@ export abstract class BaseSession implements ISession {
 	}
 
 	/**
-	 * Check if URL should skip token refresh (to avoid infinite recursion)
+	 * Check if URL should skip token refresh (to avoid infinite recursion).
+	 * Since refresh now calls init(), we only need to guard the init path.
 	 */
 	protected shouldSkipRefreshForUrl(url: string): boolean {
-		return url.includes('/api/token/refresh') || url.includes('/api/init');
+		return url.includes('/api/init');
 	}
 }

packages/client/src/core/in-process-session.ts

@@ -15,7 +15,6 @@ export interface InProcessServer {
 	handleExplore(ctx: InProcessRequestContext): Promise<unknown>;
 	handleExecute(ctx: InProcessRequestContext): Promise<unknown>;
 	handleResume(ctx: InProcessRequestContext, executionId: string): Promise<unknown>;
-	handleTokenRefresh(ctx: InProcessRequestContext): Promise<unknown>;
 }
 
 /**
@@ -65,6 +64,9 @@ export class InProcessSession extends BaseSession {
 		tools?: ClientToolDefinition[],
 		services?: { hasLLM: boolean; hasApproval: boolean; hasEmbedding: boolean; hasTools: boolean }
 	): Promise<TokenCredentials> {
+		// Store init params so doRefreshToken() can re-init with the same data
+		this.storedInitParams = { clientInfo, tools, services };
+
 		if (this.initPromise) {
 			await this.initPromise;
 			return {
@@ -133,20 +135,6 @@ export class InProcessSession extends BaseSession {
 		return '';
 	}
 
-	/**
-	 * Perform the actual token refresh via in-process server call
-	 */
-	protected async doRefreshToken(): Promise<void> {
-		const ctx = await this.createContext({
-			method: 'POST',
-			path: '/api/token/refresh',
-			body: { clientId: this.clientId },
-		});
-
-		const result = (await this.server.handleTokenRefresh(ctx)) as TokenCredentials;
-		this.updateTokenState(result);
-	}
-
 	/**
 	 * Prepares headers for a request, refreshing token if needed
 	 */

packages/client/src/core/session.ts

@@ -32,6 +32,9 @@ export class ClientSession extends BaseSession {
 		tools?: ClientToolDefinition[],
 		services?: { hasLLM: boolean; hasApproval: boolean; hasEmbedding: boolean; hasTools: boolean }
 	): Promise<TokenCredentials> {
+		// Store init params so doRefreshToken() can re-init with the same data
+		this.storedInitParams = { clientInfo, tools, services };
+
 		if (this.initPromise) {
 			await this.initPromise;
 			return {
@@ -107,36 +110,6 @@ export class ClientSession extends BaseSession {
 		return this.baseUrl;
 	}
 
-	/**
-	 * Perform the actual token refresh via HTTP
-	 */
-	protected async doRefreshToken(): Promise<void> {
-		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<string, string> = {
-			'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 TokenCredentials;
-		this.updateTokenState(data);
-	}
-
 	/**
 	 * Prepares headers for a request, refreshing token if needed and calling preRequest hook if configured
 	 */

packages/server/src/client-sessions.ts

@@ -74,10 +74,9 @@ export class ClientSessionManager {
 		}
 	}
 
-	private ensureClientJWT(token: string, clientId: string, ignoreExpiration = false) {
+	private ensureClientJWT(token: string, clientId: string) {
 		const decoded = jwt.verify(token, this.jwtSecret, {
 			algorithms: ['HS256'],
-			ignoreExpiration,
 		}) as { clientId: string; type: string };
 
 		if (decoded.clientId !== clientId || decoded.type !== 'client') {
@@ -136,27 +135,6 @@ export class ClientSessionManager {
 		}
 	}
 
-	/**
-	 * 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 verifyClientForRefresh(clientId: string, token: string): Promise<boolean> {
-		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;
-			}
-
-			// Check if session exists in cache - don't check session.expiresAt
-			const session = await this.cache.get<ClientSession>(`session:${clientId}`);
-			return session !== null;
-		} catch {
-			return false;
-		}
-	}
-
 	/**
 	 * Get client session
 	 */
@@ -207,47 +185,6 @@ export class ClientSessionManager {
 		);
 	}
 
-	/**
-	 * 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.
-	 */
-	async refreshToken(clientId: string): Promise<ClientInitResponse | null> {
-		// Get session directly from cache without expiry check
-		const session = await this.cache.get<ClientSession>(`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 both new clientId and new expiresAt
-		const updatedSession: ClientSession = {
-			...session,
-			clientId,
-			expiresAt: newExpiresAt,
-		};
-
-		await this.cache.set(`session:${newClientId}`, updatedSession);
-
-		const newToken = this.generateToken(newClientId);
-
-		return {
-			clientId: newClientId,
-			token: newToken,
-			expiresAt: newExpiresAt,
-			tokenRotateAt: newTokenRotateAt,
-		};
-	}
-
 	/**
 	 * Get token TTL and rotation settings (useful for clients)
 	 */

packages/server/src/create-server.ts

@@ -39,7 +39,6 @@ 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 {
@@ -159,7 +158,7 @@ export class AgentToolProtocolServer {
 		}
 
 		if (!this.cacheProvider) {
-			this.cacheProvider = new MemoryCache({ maxKeys: 1000, defaultTTL: 3600 });
+			this.cacheProvider = new MemoryCache({ maxKeys: 1000, defaultTTL: 24 * 3600 });
 			log.info('Cache provider configured (default)', { provider: 'memory' });
 		}
 
@@ -612,11 +611,6 @@ export class AgentToolProtocolServer {
 		);
 	}
 
-	async handleTokenRefresh(ctx: RequestContext): Promise<unknown> {
-		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

packages/server/src/graphql-loader.ts

@@ -204,6 +204,28 @@ export interface LoadGraphQLOptions {
 	auth?: AuthConfig;
 	depthLimit?: number;
 	queryDepthLimit?: number;
+
+	/**
+	 * Custom fetch function for full control over the HTTP transport layer.
+	 * Use this when the target API requires custom TLS certificates (mTLS, custom CA),
+	 * a proxy, or any other transport-level configuration that headers alone cannot provide.
+	 *
+	 * Applied to both schema loading and runtime GraphQL requests.
+	 *
+	 * @example
+	 * ```typescript
+	 * import { Agent } from 'undici';
+	 *
+	 * const mtlsAgent = new Agent({
+	 *   connect: { ca: fs.readFileSync('/certs/ca.pem') },
+	 * });
+	 *
+	 * await server.loadGraphQL('https://api.example.com/graphql', {
+	 *   fetcher: (url, init) => fetch(url, { ...init, dispatcher: mtlsAgent }),
+	 * });
+	 * ```
+	 */
+	fetcher?: (url: string, init: RequestInit) => Promise<Response>;
 }
 
 /**
@@ -244,10 +266,11 @@ export async function loadGraphQL(
 async function loadSchema(source: string, options: LoadGraphQLOptions): Promise<GraphQLSchema> {
 	// Check if source is a URL
 	if (source.startsWith('http://') || source.startsWith('https://')) {
+		const fetchFn = options.fetcher || fetch;
 		// Resolve headers for schema loading
 		const headers = await resolveHeaders(options);
 
-		const response = await fetch(source, {
+		const response = await fetchFn(source, {
 			method: 'POST',
 			headers: {
 				'Content-Type': 'application/json',
@@ -258,7 +281,8 @@ async function loadSchema(source: string, options: LoadGraphQLOptions): Promise<
 
 		if (!response.ok) {
 			// If POST with introspection fails, try GET assuming it might return SDL or JSON
-			const getResponse = await fetch(source, {
+			const getResponse = await fetchFn(source, {
+				method: 'GET',
 				headers,
 			});
 			if (getResponse.ok) {
@@ -352,7 +376,8 @@ function convertFieldToFunction(
 			// Pass execution context (including requestContext) to resolveHeaders
 			const headers = await resolveHeaders(options, paramsObj, context);
 
-			const response = await fetch(url, {
+			const fetchFn = options.fetcher || fetch;
+			const response = await fetchFn(url, {
 				method: 'POST',
 				headers: {
 					'Content-Type': 'application/json',