Andreas/LibreChat
1
0
Fork 0
mirror of https://github.com/danny-avila/LibreChat.git synced 2025-12-18 17:30:16 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

refactor: Enhance GenerationJobManager and event transports for improved resource management

Browse source 
- Updated GenerationJobManager to prevent immediate cleanup of eventTransport upon job completion, allowing final events to transmit fully before cleanup.
- Added orphaned stream cleanup logic in GenerationJobManager to handle streams without corresponding jobs.
- Introduced getTrackedStreamIds method in both InMemoryEventTransport and RedisEventTransport for better management of orphaned streams.
- Improved comments for clarity on resource management and cleanup processes.
This commit is contained in:
Danny Avila 2025-12-15 09:44:27 -05:00
parent fe37f428bc
commit f82483c6b6
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
4 changed files with 36 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

18
packages/api/src/stream/GenerationJobManager.ts
View file

@ -357,7 +357,10 @@ class GenerationJobManagerClass {
/**
* Mark job as complete.
* If cleanupOnComplete is true (default), immediately cleans up all job resources.
* If cleanupOnComplete is true (default), immediately cleans up job resources.
* Note: eventTransport is NOT cleaned up here to allow the final event to be
* fully transmitted. It will be cleaned up when subscribers disconnect or
* by the periodic cleanup job.
*/
async completeJob(streamId: string, error?: string): Promise<void> {
// Clear content state and run step buffer (Redis only)
@ -367,7 +370,8 @@ class GenerationJobManagerClass {
// Immediate cleanup if configured (default: true)
if (this._cleanupOnComplete) {
this.runtimeState.delete(streamId);
this.eventTransport.cleanup(streamId);
// Don't cleanup eventTransport here - let the done event fully transmit first.
// EventTransport will be cleaned up when subscribers disconnect or by periodic cleanup.
await this.jobStore.deleteJob(streamId);
} else {
// Only update status if keeping the job around
@ -443,7 +447,7 @@ class GenerationJobManagerClass {
// Immediate cleanup if configured (default: true)
if (this._cleanupOnComplete) {
this.runtimeState.delete(streamId);
this.eventTransport.cleanup(streamId);
// Don't cleanup eventTransport here - let the abort event fully transmit first.
await this.jobStore.deleteJob(streamId);
} else {
// Only update status if keeping the job around
@ -806,6 +810,14 @@ class GenerationJobManagerClass {
}
}
// Check eventTransport for orphaned streams (e.g., connections dropped without clean close)
// These are streams that exist in eventTransport but have no corresponding job
for (const streamId of this.eventTransport.getTrackedStreamIds()) {
if (!(await this.jobStore.hasJob(streamId)) && !this.runtimeState.has(streamId)) {
this.eventTransport.cleanup(streamId);
}
}
if (count > 0) {
logger.debug(`[GenerationJobManager] Cleaned up ${count} expired jobs`);
}