@agentuity/runtime 0.0.61 → 0.0.63

package/src/session.ts

@@ -4,6 +4,8 @@ import type { Context } from 'hono';
 import { getCookie, setCookie } from 'hono/cookie';
 import { type Env, fireEvent } from './app';
 import type { AppState } from './index';
+import { getServiceUrls } from '@agentuity/server';
+import { WebSocket } from 'ws';
 
 export type ThreadEventName = 'destroyed';
 export type SessionEventName = 'completed';
@@ -31,7 +33,7 @@ type SessionEventCallback<T extends Session> = (
  * const agent = createAgent({
  *   handler: async (ctx, input) => {
  *     // Get thread ID
- *     console.log('Thread:', ctx.thread.id);
+ *     ctx.logger.info('Thread: %s', ctx.thread.id);
  *
  *     // Store data in thread state (persists across sessions)
  *     ctx.thread.state.set('conversationCount',
@@ -40,7 +42,7 @@ type SessionEventCallback<T extends Session> = (
  *
  *     // Listen for thread destruction
  *     ctx.thread.addEventListener('destroyed', (eventName, thread) => {
- *       console.log('Thread destroyed:', thread.id);
+ *       ctx.logger.info('Thread destroyed: %s', thread.id);
  *     });
  *
  *     return 'Response';
@@ -79,7 +81,7 @@ export interface Thread {
 	 * @example
 	 * ```typescript
 	 * ctx.thread.addEventListener('destroyed', (eventName, thread) => {
-	 *   console.log('Cleaning up thread:', thread.id);
+	 *   ctx.logger.info('Cleaning up thread: %s', thread.id);
 	 * });
 	 * ```
 	 */
@@ -105,7 +107,7 @@ export interface Thread {
 	 *
 	 * @example
 	 * ```typescript
-	 * // End conversation
+	 * // Permanently delete the thread from storage
 	 * await ctx.thread.destroy();
 	 * ```
 	 */
@@ -123,18 +125,18 @@ export interface Thread {
  * const agent = createAgent({
  *   handler: async (ctx, input) => {
  *     // Get session ID (unique per request)
- *     console.log('Session:', ctx.session.id);
+ *     ctx.logger.info('Session: %s', ctx.session.id);
  *
  *     // Store data in session state (only for this request)
  *     ctx.session.state.set('startTime', Date.now());
  *
  *     // Access parent thread
- *     console.log('Thread:', ctx.session.thread.id);
+ *     ctx.logger.info('Thread: %s', ctx.session.thread.id);
  *
  *     // Listen for session completion
  *     ctx.session.addEventListener('completed', (eventName, session) => {
  *       const duration = Date.now() - (session.state.get('startTime') as number);
- *       console.log(`Session completed in ${duration}ms`);
+ *       ctx.logger.info('Session completed in %dms', duration);
  *     });
  *
  *     return 'Response';
@@ -176,7 +178,7 @@ export interface Session {
 	 * @example
 	 * ```typescript
 	 * ctx.session.addEventListener('completed', (eventName, session) => {
-	 *   console.log('Session finished:', session.id);
+	 *   ctx.logger.info('Session finished: %s', session.id);
 	 * });
 	 * ```
 	 */
@@ -195,6 +197,31 @@ export interface Session {
 		eventName: 'completed',
 		callback: (eventName: 'completed', session: Session) => Promise<void> | void
 	): void;
+
+	/**
+	 * Return the session data as a serializable string or return undefined if not
+	 * data should be serialized.
+	 */
+	serializeUserData(): string | undefined;
+}
+
+/**
+ * Represent an interface for handling how thread ids are generated or restored.
+ */
+export interface ThreadIDProvider {
+	/**
+	 * A function that should return a thread id to be used for the incoming request.
+	 * The returning thread id must be globally unique and must start with the prefix
+	 * thrd_ such as `thrd_212c16896b974ffeb21a748f0eeba620`. The max length of the
+	 * string is 64 characters and the min length is 32 characters long
+	 * (including the prefix). The characters after the prefix must match the
+	 * regular expression [a-zA-Z0-9-].
+	 *
+	 * @param appState - The app state from createApp setup function
+	 * @param ctx - Hono request context
+	 * @returns The thread id to use
+	 */
+	getThreadId(appState: AppState, ctx: Context<Env>): string;
 }
 
 /**
@@ -253,6 +280,14 @@ export interface ThreadProvider {
 	 */
 	initialize(appState: AppState): Promise<void>;
 
+	/**
+	 * Set the provider to use for generating / restoring the thread id
+	 * on new requests.  Overrides the built-in provider when set.
+	 *
+	 * @param provider - the provider implementation
+	 */
+	setThreadIDProvider(provider: ThreadIDProvider): void;
+
 	/**
 	 * Restore or create a thread from the HTTP request context.
 	 * Should check cookies for existing thread ID or create a new one.
@@ -394,17 +429,37 @@ export function generateId(prefix?: string): string {
 	return `${prefix}${prefix ? '_' : ''}${arr.toHex()}`;
 }
 
+/**
+ * DefaultThreadIDProvider will look for a cookie named `atid` and use that as
+ * the thread id or if not found, generate a new one.
+ */
+export class DefaultThreadIDProvider implements ThreadIDProvider {
+	getThreadId(_appState: AppState, ctx: Context<Env>): string {
+		const cookie = getCookie(ctx);
+		let threadId: string | undefined;
+
+		if (cookie.atid?.startsWith('thrd_')) {
+			threadId = cookie.atid;
+		}
+
+		threadId = threadId || generateId('thrd');
+
+		setCookie(ctx, 'atid', threadId);
+		return threadId;
+	}
+}
+
 export class DefaultThread implements Thread {
-	#lastUsed: number;
+	#initialStateJson: string | undefined;
 	readonly id: string;
 	readonly state: Map<string, unknown>;
 	private provider: ThreadProvider;
 
-	constructor(provider: ThreadProvider, id: string) {
+	constructor(provider: ThreadProvider, id: string, initialStateJson?: string) {
 		this.provider = provider;
 		this.id = id;
 		this.state = new Map();
-		this.#lastUsed = Date.now();
+		this.#initialStateJson = initialStateJson;
 	}
 
 	addEventListener(eventName: ThreadEventName, callback: ThreadEventCallback<any>): void {
@@ -434,15 +489,29 @@ export class DefaultThread implements Thread {
 	}
 
 	async destroy(): Promise<void> {
-		this.provider.destroy(this);
+		await this.provider.destroy(this);
 	}
 
-	touch() {
-		this.#lastUsed = Date.now();
+	/**
+	 * Check if thread state has been modified since restore
+	 * @internal
+	 */
+	isDirty(): boolean {
+		if (this.state.size === 0 && !this.#initialStateJson) {
+			return false;
+		}
+
+		const currentJson = JSON.stringify(Object.fromEntries(this.state));
+
+		return currentJson !== this.#initialStateJson;
 	}
 
-	expired() {
-		return Date.now() - this.#lastUsed >= 3.6e6; // 1 hour
+	/**
+	 * Get serialized state for saving
+	 * @internal
+	 */
+	getSerializedState(): string {
+		return JSON.stringify(Object.fromEntries(this.state));
 	}
 }
 
@@ -482,64 +551,420 @@ export class DefaultSession implements Session {
 	async fireEvent(eventName: SessionEventName): Promise<void> {
 		await fireSessionEvent(this, eventName);
 	}
+
+	/**
+	 * Serialize session state to JSON string for persistence.
+	 * Returns undefined if state is empty or exceeds 1MB limit.
+	 * @internal
+	 */
+	serializeUserData(): string | undefined {
+		if (this.state.size === 0) {
+			return undefined;
+		}
+
+		try {
+			const obj = Object.fromEntries(this.state);
+			const json = JSON.stringify(obj);
+
+			// Check 1MB limit (1,048,576 bytes)
+			const sizeInBytes = new TextEncoder().encode(json).length;
+			if (sizeInBytes > 1048576) {
+				console.error(
+					`Session ${this.id} user_data exceeds 1MB limit (${sizeInBytes} bytes), data will not be persisted`
+				);
+				return undefined;
+			}
+
+			return json;
+		} catch (err) {
+			console.error(`Failed to serialize session ${this.id} user_data:`, err);
+			return undefined;
+		}
+	}
 }
 
-export class DefaultThreadProvider implements ThreadProvider {
-	private threads = new Map<string, DefaultThread>();
+/**
+ * WebSocket client for thread state persistence
+ * @internal
+ */
+class ThreadWebSocketClient {
+	private ws: WebSocket | null = null;
+	private authenticated = false;
+	private pendingRequests = new Map<
+		string,
+		{ resolve: (data?: string) => void; reject: (err: Error) => void }
+	>();
+	private reconnectAttempts = 0;
+	private maxReconnectAttempts = 5;
+	private apiKey: string;
+	private wsUrl: string;
+	private wsConnecting: Promise<void> | null = null;
+
+	constructor(apiKey: string, wsUrl: string) {
+		this.apiKey = apiKey;
+		this.wsUrl = wsUrl;
+	}
 
-	async initialize(_appState: AppState): Promise<void> {
-		setInterval(() => {
-			for (const [, thread] of this.threads) {
-				if (thread.expired()) {
-					void (async () => {
-						try {
-							await this.destroy(thread);
-						} catch (err) {
-							console.error('Failed to destroy expired thread', err);
+	async connect(): Promise<void> {
+		return new Promise((resolve, reject) => {
+			// Set connection timeout
+			const connectionTimeout = setTimeout(() => {
+				this.cleanup();
+				reject(new Error('WebSocket connection timeout (10s)'));
+			}, 10_000);
+
+			try {
+				this.ws = new WebSocket(this.wsUrl);
+
+				this.ws.on('open', () => {
+					// Send authentication (do NOT clear timeout yet - wait for auth response)
+					this.ws?.send(JSON.stringify({ authorization: this.apiKey }));
+				});
+
+				this.ws.on('message', (data: any) => {
+					try {
+						const message = JSON.parse(data.toString());
+
+						// Handle auth response
+						if ('success' in message && !this.authenticated) {
+							clearTimeout(connectionTimeout);
+							if (message.success) {
+								this.authenticated = true;
+								this.reconnectAttempts = 0;
+								resolve();
+							} else {
+								const err = new Error(
+									`WebSocket authentication failed: ${message.error || 'Unknown error'}`
+								);
+								this.cleanup();
+								reject(err);
+							}
+							return;
 						}
-					})();
-				}
+
+						// Handle action response
+						if ('id' in message && this.pendingRequests.has(message.id)) {
+							const pending = this.pendingRequests.get(message.id)!;
+							this.pendingRequests.delete(message.id);
+
+							if (message.success) {
+								pending.resolve(message.data);
+							} else {
+								pending.reject(new Error(message.error || 'Request failed'));
+							}
+						}
+					} catch {
+						// Ignore parse errors
+					}
+				});
+
+				this.ws.on('error', (err: Error) => {
+					clearTimeout(connectionTimeout);
+					if (!this.authenticated) {
+						reject(new Error(`WebSocket error: ${err.message}`));
+					}
+				});
+
+				this.ws.on('close', () => {
+					clearTimeout(connectionTimeout);
+					const wasAuthenticated = this.authenticated;
+					this.authenticated = false;
+
+					// Reject all pending requests
+					for (const [id, pending] of this.pendingRequests) {
+						pending.reject(new Error('WebSocket connection closed'));
+						this.pendingRequests.delete(id);
+					}
+
+					// Reject connecting promise if still pending
+					if (!wasAuthenticated) {
+						reject(new Error('WebSocket closed before authentication'));
+					}
+
+					// Attempt reconnection only if we were previously authenticated
+					if (wasAuthenticated && this.reconnectAttempts < this.maxReconnectAttempts) {
+						this.reconnectAttempts++;
+						const delay = Math.min(1000 * Math.pow(2, this.reconnectAttempts), 30_000);
+
+						// Schedule reconnection with backoff delay
+						setTimeout(() => {
+							// Create new connection promise for reconnection
+							this.wsConnecting = this.connect().catch(() => {
+								// Reconnection failed, reset
+								this.wsConnecting = null;
+							});
+						}, delay);
+					}
+				});
+			} catch (err) {
+				clearTimeout(connectionTimeout);
+				reject(err);
 			}
-		}, 60_000).unref();
+		});
+	}
+
+	async restore(threadId: string): Promise<string | undefined> {
+		// Wait for connection/reconnection if in progress
+		if (this.wsConnecting) {
+			await this.wsConnecting;
+		}
+
+		if (!this.authenticated || !this.ws) {
+			throw new Error('WebSocket not connected or authenticated');
+		}
+
+		return new Promise((resolve, reject) => {
+			const requestId = crypto.randomUUID();
+			this.pendingRequests.set(requestId, { resolve, reject });
+
+			const message = {
+				id: requestId,
+				action: 'restore',
+				data: { thread_id: threadId },
+			};
+
+			this.ws!.send(JSON.stringify(message));
+
+			// Timeout after 10 seconds
+			setTimeout(() => {
+				if (this.pendingRequests.has(requestId)) {
+					this.pendingRequests.delete(requestId);
+					reject(new Error('Request timeout'));
+				}
+			}, 10000);
+		});
+	}
+
+	async save(threadId: string, userData: string): Promise<void> {
+		// Wait for connection/reconnection if in progress
+		if (this.wsConnecting) {
+			await this.wsConnecting;
+		}
+
+		if (!this.authenticated || !this.ws) {
+			throw new Error('WebSocket not connected or authenticated');
+		}
+
+		// Check 1MB limit
+		const sizeInBytes = new TextEncoder().encode(userData).length;
+		if (sizeInBytes > 1048576) {
+			console.error(
+				`Thread ${threadId} user_data exceeds 1MB limit (${sizeInBytes} bytes), data will not be persisted`
+			);
+			return;
+		}
+
+		return new Promise((resolve, reject) => {
+			const requestId = crypto.randomUUID();
+			this.pendingRequests.set(requestId, {
+				resolve: () => resolve(),
+				reject,
+			});
+
+			const message = {
+				id: requestId,
+				action: 'save',
+				data: { thread_id: threadId, user_data: userData },
+			};
+
+			this.ws!.send(JSON.stringify(message));
+
+			// Timeout after 10 seconds
+			setTimeout(() => {
+				if (this.pendingRequests.has(requestId)) {
+					this.pendingRequests.delete(requestId);
+					reject(new Error('Request timeout'));
+				}
+			}, 10_000);
+		});
+	}
+
+	async delete(threadId: string): Promise<void> {
+		// Wait for connection/reconnection if in progress
+		if (this.wsConnecting) {
+			await this.wsConnecting;
+		}
+
+		if (!this.authenticated || !this.ws) {
+			throw new Error('WebSocket not connected or authenticated');
+		}
+
+		return new Promise((resolve, reject) => {
+			const requestId = crypto.randomUUID();
+			this.pendingRequests.set(requestId, {
+				resolve: () => resolve(),
+				reject,
+			});
+
+			const message = {
+				id: requestId,
+				action: 'delete',
+				data: { thread_id: threadId },
+			};
+
+			this.ws!.send(JSON.stringify(message));
+
+			// Timeout after 10 seconds
+			setTimeout(() => {
+				if (this.pendingRequests.has(requestId)) {
+					this.pendingRequests.delete(requestId);
+					reject(new Error('Request timeout'));
+				}
+			}, 10_000);
+		});
+	}
+
+	cleanup(): void {
+		if (this.ws) {
+			this.ws.close();
+			this.ws = null;
+		}
+		this.authenticated = false;
+		this.pendingRequests.clear();
+	}
+}
+
+const validThreadIdCharacters = /^[a-zA-Z0-9-]+$/;
+
+export class DefaultThreadProvider implements ThreadProvider {
+	private appState: AppState | null = null;
+	private wsClient: ThreadWebSocketClient | null = null;
+	private wsConnecting: Promise<void> | null = null;
+	private threadIDProvider: ThreadIDProvider | null = null;
+
+	async initialize(appState: AppState): Promise<void> {
+		this.appState = appState;
+		this.threadIDProvider = new DefaultThreadIDProvider();
+
+		// Initialize WebSocket connection for thread persistence (async, non-blocking)
+		const apiKey = process.env.AGENTUITY_SDK_KEY;
+		if (apiKey) {
+			const serviceUrls = getServiceUrls();
+			const catalystUrl = serviceUrls.catalyst;
+			const wsUrl = new URL('/thread/ws', catalystUrl.replace(/^http/, 'ws'));
+
+			this.wsClient = new ThreadWebSocketClient(apiKey, wsUrl.toString());
+			// Connect in background, don't block initialization
+			this.wsConnecting = this.wsClient
+				.connect()
+				.then(() => {
+					this.wsConnecting = null;
+				})
+				.catch((err) => {
+					console.error('Failed to connect to thread WebSocket:', err);
+					this.wsClient = null;
+					this.wsConnecting = null;
+				});
+		}
+	}
+
+	setThreadIDProvider(provider: ThreadIDProvider): void {
+		this.threadIDProvider = provider;
 	}
 
 	async restore(ctx: Context<Env>): Promise<Thread> {
-		const cookie = getCookie(ctx);
-		let threadId: string | undefined;
+		const threadId = this.threadIDProvider!.getThreadId(this.appState!, ctx);
 
-		if (cookie.atid?.startsWith('thrd_')) {
-			threadId = cookie.atid;
+		if (!threadId) {
+			throw new Error(`the ThreadIDProvider returned an empty thread id for getThreadId`);
+		}
+		if (!threadId.startsWith('thrd_')) {
+			throw new Error(
+				`the ThreadIDProvider returned an invalid thread id (${threadId}) for getThreadId. The thread id must start with the prefix 'thrd_'.`
+			);
+		}
+		if (threadId.length > 64) {
+			throw new Error(
+				`the ThreadIDProvider returned an invalid thread id (${threadId}) for getThreadId. The thread id must be less than 64 characters long.`
+			);
+		}
+		if (threadId.length < 32) {
+			throw new Error(
+				`the ThreadIDProvider returned an invalid thread id (${threadId}) for getThreadId. The thread id must be at least 32 characters long.`
+			);
+		}
+		if (!validThreadIdCharacters.test(threadId.substring(5))) {
+			throw new Error(
+				`the ThreadIDProvider returned an invalid thread id (${threadId}) for getThreadId. The thread id must contain only characters that match the regular expression [a-zA-Z0-9-].`
+			);
 		}
 
-		threadId = threadId || generateId('thrd');
+		// Wait for WebSocket connection if still connecting
+		if (this.wsConnecting) {
+			await this.wsConnecting;
+		}
+
+		// Restore thread state from WebSocket if available
+		let initialStateJson: string | undefined;
+		if (this.wsClient) {
+			try {
+				const restoredData = await this.wsClient.restore(threadId);
+				if (restoredData) {
+					initialStateJson = restoredData;
+				}
+			} catch {
+				// Continue with empty state rather than failing
+			}
+		}
 
-		if (threadId) {
-			setCookie(ctx, 'atid', threadId);
-			const existing = this.threads.get(threadId);
-			if (existing) {
-				return existing;
+		const thread = new DefaultThread(this, threadId, initialStateJson);
+
+		// Populate thread state from restored data
+		if (initialStateJson) {
+			try {
+				const data = JSON.parse(initialStateJson);
+				for (const [key, value] of Object.entries(data)) {
+					thread.state.set(key, value);
+				}
+			} catch {
+				// Continue with empty state if parsing fails
 			}
 		}
 
-		const thread = new DefaultThread(this, threadId);
-		this.threads.set(thread.id, thread);
 		await fireEvent('thread.created', thread);
 		return thread;
 	}
 
 	async save(thread: Thread): Promise<void> {
 		if (thread instanceof DefaultThread) {
-			thread.touch();
+			// Wait for WebSocket connection if still connecting
+			if (this.wsConnecting) {
+				await this.wsConnecting;
+			}
+
+			// Only save to WebSocket if state has changed
+			if (this.wsClient && thread.isDirty()) {
+				try {
+					const serialized = thread.getSerializedState();
+					await this.wsClient.save(thread.id, serialized);
+				} catch {
+					// Don't throw - allow request to complete even if save fails
+				}
+			}
 		}
 	}
 
 	async destroy(thread: Thread): Promise<void> {
 		if (thread instanceof DefaultThread) {
 			try {
+				// Wait for WebSocket connection if still connecting
+				if (this.wsConnecting) {
+					await this.wsConnecting;
+				}
+
+				// Delete thread from remote storage
+				if (this.wsClient) {
+					try {
+						await this.wsClient.delete(thread.id);
+					} catch (err) {
+						console.error(`Failed to delete thread ${thread.id} from remote storage:`, err);
+						// Continue with local cleanup even if remote delete fails
+					}
+				}
+
 				await thread.fireEvent('destroyed');
 				await fireEvent('thread.destroyed', thread);
 			} finally {
-				this.threads.delete(thread.id);
 				threadEventListeners.delete(thread);
 			}
 		}