refactor: Enhance GenerationJobManager with In-Memory Implementations

- Introduced InMemoryJobStore, InMemoryEventTransport, and InMemoryContentState for improved job management and event handling.
- Updated GenerationJobManager to utilize these new implementations, allowing for better separation of concerns and easier maintenance.
- Enhanced job metadata handling to support user messages and response IDs for resumable functionality.
- Improved cleanup and state management processes to prevent memory leaks and ensure efficient resource usage.
This commit is contained in:
Danny Avila 2025-12-12 02:16:24 -05:00
parent 5ff66f2d77
commit ff86f96416
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
14 changed files with 892 additions and 321 deletions

View file

@ -0,0 +1,139 @@
import type { Agents } from 'librechat-data-provider';
import type { StandardGraph } from '@librechat/agents';
/**
* Job status enum
*/
export type JobStatus = 'running' | 'complete' | 'error' | 'aborted';
/**
* Serializable job data - no object references, suitable for Redis/external storage
*/
export interface SerializableJobData {
streamId: string;
userId: string;
status: JobStatus;
createdAt: number;
completedAt?: number;
conversationId?: string;
error?: string;
/** User message metadata */
userMessage?: {
messageId: string;
parentMessageId?: string;
conversationId?: string;
text?: string;
};
/** Response message ID for reconnection */
responseMessageId?: string;
/** Sender name for UI display */
sender?: string;
/** Whether sync has been sent to a client */
syncSent: boolean;
/** Serialized final event for replay */
finalEvent?: string;
}
/**
* Resume state for reconnecting clients
*/
export interface ResumeState {
runSteps: Agents.RunStep[];
aggregatedContent: Agents.MessageContentComplex[];
userMessage?: SerializableJobData['userMessage'];
responseMessageId?: string;
conversationId?: string;
sender?: string;
}
/**
* Interface for job storage backend.
* Implementations can use in-memory Map, Redis, KV store, etc.
*/
export interface IJobStore {
/** Create a new job */
createJob(
streamId: string,
userId: string,
conversationId?: string,
): Promise<SerializableJobData>;
/** Get a job by streamId */
getJob(streamId: string): Promise<SerializableJobData | null>;
/** Find active job by conversationId */
getJobByConversation(conversationId: string): Promise<SerializableJobData | null>;
/** Update job data */
updateJob(streamId: string, updates: Partial<SerializableJobData>): Promise<void>;
/** Delete a job */
deleteJob(streamId: string): Promise<void>;
/** Check if job exists */
hasJob(streamId: string): Promise<boolean>;
/** Get all running jobs (for cleanup) */
getRunningJobs(): Promise<SerializableJobData[]>;
/** Cleanup expired jobs */
cleanup(): Promise<number>;
}
/**
* Interface for pub/sub event transport.
* Implementations can use EventEmitter, Redis Pub/Sub, etc.
*/
export interface IEventTransport {
/** Subscribe to events for a stream */
subscribe(
streamId: string,
handlers: {
onChunk: (event: unknown) => void;
onDone?: (event: unknown) => void;
onError?: (error: string) => void;
},
): { unsubscribe: () => void };
/** Publish a chunk event */
emitChunk(streamId: string, event: unknown): void;
/** Publish a done event */
emitDone(streamId: string, event: unknown): void;
/** Publish an error event */
emitError(streamId: string, error: string): void;
/** Get subscriber count for a stream */
getSubscriberCount(streamId: string): number;
/** Listen for all subscribers leaving */
onAllSubscribersLeft(streamId: string, callback: () => void): void;
}
/**
* Interface for content state management.
* Separates volatile content state from persistent job data.
* In-memory only - not persisted to external storage.
*/
export interface IContentStateManager {
/** Set content parts reference (in-memory only) */
setContentParts(streamId: string, contentParts: Agents.MessageContentComplex[]): void;
/** Get content parts */
getContentParts(streamId: string): Agents.MessageContentComplex[] | null;
/** Set graph reference for run steps */
setGraph(streamId: string, graph: StandardGraph): void;
/** Get run steps from graph */
getRunSteps(streamId: string): Agents.RunStep[];
/** Clear content state for a job */
clearContentState(streamId: string): void;
}