From 73699b5c259103168a83adc79cfe88f563f06cde Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sat, 25 Jul 2026 07:58:20 -0400 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=20perf:=20Reduce=20Agent=20Chat=20Sta?= =?UTF-8?q?rtup=20Latency=20(#14423)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf: reduce agent chat startup latency * test: align Redis stream readiness assertions * perf: overlap remaining agent startup work * perf: persist initial agent job metadata atomically * test: add agent startup latency benchmark * fix: harden resumable agent stream lifecycle * fix: isolate replacement stream lifecycles * fix: preserve terminal stream epochs --- .gitignore | 1 + api/app/clients/BaseClient.js | 8 +- api/app/clients/specs/BaseClient.test.js | 59 + api/server/controllers/ModelController.js | 6 +- .../controllers/ModelController.spec.js | 57 + .../__tests__/askUserQuestion.e2e.spec.js | 4 - .../__tests__/hitlCheckpoint.e2e.spec.js | 12 +- .../__tests__/request.resumeMetadata.spec.js | 330 +- .../agents/__tests__/resume.spec.js | 80 +- api/server/controllers/agents/client.js | 144 +- api/server/controllers/agents/client.test.js | 223 ++ api/server/controllers/agents/request.js | 424 ++- api/server/controllers/agents/resume.js | 102 +- api/server/experimental.js | 20 +- api/server/index.js | 32 +- api/server/index.spec.js | 28 + .../agents/__tests__/streamTenant.spec.js | 69 + api/server/routes/agents/index.js | 34 +- .../services/Endpoints/agents/initialize.js | 120 +- .../Endpoints/agents/initialize.spec.js | 44 + api/typedefs.js | 1 + e2e/benchmarks/README.md | 56 + e2e/benchmarks/agent-startup.latency.spec.ts | 369 +++ e2e/benchmarks/mongoose-latency-hook.cjs | 20 + e2e/playwright.config.benchmark.ts | 29 + package.json | 1 + .../agents/checkpointer.integration.spec.ts | 34 + packages/api/src/agents/checkpointer.spec.ts | 8 + packages/api/src/agents/checkpointer.ts | 78 +- packages/api/src/agents/conversation.spec.ts | 77 + packages/api/src/agents/conversation.ts | 53 + packages/api/src/agents/index.ts | 3 + packages/api/src/agents/phases.ts | 30 + packages/api/src/agents/startup.spec.ts | 322 ++ packages/api/src/agents/startup.ts | 285 ++ packages/api/src/app/metrics.spec.ts | 33 + packages/api/src/app/metrics.ts | 61 + packages/api/src/app/shutdown.spec.ts | 68 + packages/api/src/app/shutdown.ts | 76 +- packages/api/src/stream/ApprovalLifecycle.ts | 66 +- .../api/src/stream/GenerationJobManager.ts | 2667 ++++++++++++----- packages/api/src/stream/SteeringLifecycle.ts | 15 +- ...ationJobManager.stream_integration.spec.ts | 25 +- .../__tests__/RedisEventTransport.spec.ts | 1267 +++++++- ...sEventTransport.stream_integration.spec.ts | 402 ++- .../stream/__tests__/RedisJobStore.spec.ts | 732 +++++ .../RedisJobStore.stream_integration.spec.ts | 423 ++- .../src/stream/__tests__/helpers/publisher.ts | 4 + .../stream/__tests__/pendingAction.spec.ts | 297 +- ...-reorder-desync.stream_integration.spec.ts | 8 +- .../stream/__tests__/staleJobReaping.spec.ts | 79 + .../api/src/stream/__tests__/startup.spec.ts | 2353 +++++++++++++++ .../api/src/stream/__tests__/steering.spec.ts | 157 +- .../implementations/InMemoryEventTransport.ts | 72 +- .../implementations/InMemoryJobStore.ts | 144 +- .../implementations/RedisEventTransport.ts | 335 ++- .../stream/implementations/RedisJobStore.ts | 1173 +++++--- packages/api/src/stream/index.ts | 2 + .../api/src/stream/interfaces/IJobStore.ts | 143 +- .../src/stream/internal/chunkPublication.ts | 43 + packages/api/src/stream/metadata.ts | 40 + packages/api/src/telemetry/sdk.spec.ts | 4 +- packages/api/src/telemetry/sdk.ts | 12 +- packages/api/src/types/stream.ts | 14 +- 64 files changed, 12134 insertions(+), 1744 deletions(-) create mode 100644 api/server/controllers/ModelController.spec.js create mode 100644 e2e/benchmarks/README.md create mode 100644 e2e/benchmarks/agent-startup.latency.spec.ts create mode 100644 e2e/benchmarks/mongoose-latency-hook.cjs create mode 100644 e2e/playwright.config.benchmark.ts create mode 100644 packages/api/src/agents/conversation.spec.ts create mode 100644 packages/api/src/agents/conversation.ts create mode 100644 packages/api/src/agents/phases.ts create mode 100644 packages/api/src/agents/startup.spec.ts create mode 100644 packages/api/src/agents/startup.ts create mode 100644 packages/api/src/stream/__tests__/RedisJobStore.spec.ts create mode 100644 packages/api/src/stream/__tests__/startup.spec.ts create mode 100644 packages/api/src/stream/internal/chunkPublication.ts create mode 100644 packages/api/src/stream/metadata.ts diff --git a/.gitignore b/.gitignore index 264489c043..d28d18530e 100644 --- a/.gitignore +++ b/.gitignore @@ -88,6 +88,7 @@ archive .vscode/settings.json src/style - official.css /e2e/specs/.test-results/ +/e2e/benchmarks/.test-results/ /e2e/.generated/ /e2e/playwright-report/ /playwright/.cache/ diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js index 6d541e160f..5f426a56c9 100644 --- a/api/app/clients/BaseClient.js +++ b/api/app/clients/BaseClient.js @@ -389,6 +389,7 @@ class BaseClient { parentMessageId, responseMessageId, } = await this.setMessageOptions(opts); + this.options.startupTelemetry?.mark('history_loaded'); const userMessage = opts.isEdited ? this.currentMessages[this.currentMessages.length - 2] @@ -610,6 +611,7 @@ class BaseClient { this.getBuildMessagesOptions(opts), opts, ); + this.options.startupTelemetry?.mark('messages_built'); if (tokenCountMap && tokenCountMap[userMessage.messageId]) { userMessage.tokenCount = tokenCountMap[userMessage.messageId]; @@ -1521,8 +1523,10 @@ class BaseClient { return message; } - await this.addFileContextToMessage(message, contextFiles); - await this.processAttachments(message, contextFiles); + await Promise.all([ + this.addFileContextToMessage(message, contextFiles), + this.processAttachments(message, contextFiles), + ]); this.message_file_map[message.messageId] = contextFiles; return message; diff --git a/api/app/clients/specs/BaseClient.test.js b/api/app/clients/specs/BaseClient.test.js index d565f87012..da3f80bc0b 100644 --- a/api/app/clients/specs/BaseClient.test.js +++ b/api/app/clients/specs/BaseClient.test.js @@ -1,6 +1,14 @@ const { Constants } = require('librechat-data-provider'); const { FakeClient, initializeFakeClient } = require('./FakeClient'); +function deferred() { + let resolve; + const promise = new Promise((resolvePromise) => { + resolve = resolvePromise; + }); + return { promise, resolve }; +} + jest.mock('~/db/connect'); jest.mock('~/server/services/Config', () => ({ getAppConfig: jest.fn().mockResolvedValue({ @@ -905,6 +913,18 @@ describe('BaseClient', () => { expect(TestClient.sendCompletion).toHaveBeenCalledWith(payload, opts); }); + test('records history and message-build startup milestones', async () => { + const startupTelemetry = { mark: jest.fn() }; + TestClient.options.startupTelemetry = startupTelemetry; + + await TestClient.sendMessage('Hello, world!', {}); + + expect(startupTelemetry.mark.mock.calls.map(([milestone]) => milestone)).toEqual([ + 'history_loaded', + 'messages_built', + ]); + }); + test('getTokenCount for response is called with the correct arguments', async () => { const tokenCountMap = {}; // Mock tokenCountMap TestClient.buildMessages.mockReturnValue({ prompt: [], tokenCountMap }); @@ -1476,6 +1496,45 @@ describe('BaseClient', () => { expect(JSON.stringify(secondMessage)).not.toContain('second-forged'); }); + test('extracts historical file context while encoding provider attachments', async () => { + getFiles.mockResolvedValueOnce([ownerFile]); + const fileContext = deferred(); + const providerAttachments = deferred(); + let completed = false; + + TestClient.addFileContextToMessage.mockImplementation(async (message) => { + await fileContext.promise; + message.fileContext = 'authorized owner text'; + }); + TestClient.processAttachments.mockImplementation(() => providerAttachments.promise); + + const messagesPromise = TestClient.addPreviousAttachments([ + { + messageId: 'msg-concurrent-file-work', + files: [{ file_id: 'owner-file', filename: 'owner.txt' }], + }, + ]).then((messages) => { + completed = true; + return messages; + }); + + await Promise.resolve(); + await Promise.resolve(); + + expect(TestClient.addFileContextToMessage).toHaveBeenCalledTimes(1); + expect(TestClient.processAttachments).toHaveBeenCalledTimes(1); + + providerAttachments.resolve([ownerFile]); + await Promise.resolve(); + expect(completed).toBe(false); + + fileContext.resolve(); + const [message] = await messagesPromise; + + expect(message.fileContext).toBe('authorized owner text'); + expect(TestClient.message_file_map['msg-concurrent-file-work']).toEqual([ownerFile]); + }); + test('preserves download-only historical attachments without trusting file fields', async () => { const [message] = await TestClient.addPreviousAttachments([ { diff --git a/api/server/controllers/ModelController.js b/api/server/controllers/ModelController.js index 4738d45111..920306bfc5 100644 --- a/api/server/controllers/ModelController.js +++ b/api/server/controllers/ModelController.js @@ -4,8 +4,10 @@ const { loadDefaultModels, loadConfigModels } = require('~/server/services/Confi const getModelsConfig = (req) => loadModels(req); async function loadModels(req) { - const defaultModelsConfig = await loadDefaultModels(req); - const customModelsConfig = await loadConfigModels(req); + const [defaultModelsConfig, customModelsConfig] = await Promise.all([ + loadDefaultModels(req), + loadConfigModels(req), + ]); return { ...defaultModelsConfig, ...customModelsConfig }; } diff --git a/api/server/controllers/ModelController.spec.js b/api/server/controllers/ModelController.spec.js new file mode 100644 index 0000000000..b7920487bc --- /dev/null +++ b/api/server/controllers/ModelController.spec.js @@ -0,0 +1,57 @@ +const mockLoadDefaultModels = jest.fn(); +const mockLoadConfigModels = jest.fn(); + +jest.mock('@librechat/data-schemas', () => ({ + logger: { + error: jest.fn(), + }, +})); + +jest.mock('~/server/services/Config', () => ({ + loadDefaultModels: (...args) => mockLoadDefaultModels(...args), + loadConfigModels: (...args) => mockLoadConfigModels(...args), +})); + +const { loadModels } = require('./ModelController'); + +function deferred() { + let resolve; + const promise = new Promise((resolvePromise) => { + resolve = resolvePromise; + }); + return { promise, resolve }; +} + +describe('loadModels', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('loads default and configured models concurrently while preserving custom precedence', async () => { + const defaultModels = deferred(); + const configuredModels = deferred(); + const req = { user: { id: 'user-1' } }; + mockLoadDefaultModels.mockReturnValue(defaultModels.promise); + mockLoadConfigModels.mockReturnValue(configuredModels.promise); + + const resultPromise = loadModels(req); + + expect(mockLoadDefaultModels).toHaveBeenCalledWith(req); + expect(mockLoadConfigModels).toHaveBeenCalledWith(req); + + configuredModels.resolve({ + openAI: ['configured-model'], + custom: ['custom-model'], + }); + defaultModels.resolve({ + openAI: ['default-model'], + anthropic: ['default-anthropic'], + }); + + await expect(resultPromise).resolves.toEqual({ + openAI: ['configured-model'], + anthropic: ['default-anthropic'], + custom: ['custom-model'], + }); + }); +}); diff --git a/api/server/controllers/agents/__tests__/askUserQuestion.e2e.spec.js b/api/server/controllers/agents/__tests__/askUserQuestion.e2e.spec.js index 1976ae69a6..4613c50cbb 100644 --- a/api/server/controllers/agents/__tests__/askUserQuestion.e2e.spec.js +++ b/api/server/controllers/agents/__tests__/askUserQuestion.e2e.spec.js @@ -201,13 +201,9 @@ beforeAll(async () => { GenerationJobManager.configure({ ...createStreamServices(), cleanupOnComplete: false }); GenerationJobManager.initialize(); - GenerationJobManager.setApprovalExpiredHandler(async (conversationId) => { - await deleteAgentCheckpoint(conversationId, MONGO_CFG); - }); }, 60000); afterAll(async () => { - GenerationJobManager.setApprovalExpiredHandler(null); await GenerationJobManager.destroy(); await mongoose.disconnect(); await mongoServer.stop(); diff --git a/api/server/controllers/agents/__tests__/hitlCheckpoint.e2e.spec.js b/api/server/controllers/agents/__tests__/hitlCheckpoint.e2e.spec.js index 29e594bd7a..ada484c43e 100644 --- a/api/server/controllers/agents/__tests__/hitlCheckpoint.e2e.spec.js +++ b/api/server/controllers/agents/__tests__/hitlCheckpoint.e2e.spec.js @@ -147,14 +147,9 @@ beforeAll(async () => { GenerationJobManager.configure({ ...createStreamServices(), cleanupOnComplete: false }); GenerationJobManager.initialize(); - // Mirrors api/server/index.js: expiry prunes the paused run's durable checkpoint. - GenerationJobManager.setApprovalExpiredHandler(async (conversationId) => { - await deleteAgentCheckpoint(conversationId, MONGO_CFG); - }); }, 60000); afterAll(async () => { - GenerationJobManager.setApprovalExpiredHandler(null); await GenerationJobManager.destroy(); await mongoose.disconnect(); await mongoServer.stop(); @@ -313,7 +308,7 @@ describe('HITL checkpoint lifecycle (full wiring)', () => { expect(job).toBeDefined(); }); - test('an abandoned pause is pruned eagerly on approval EXPIRY (not left to the TTL)', async () => { + test('an abandoned pause expires without deleting a replacement-scoped checkpoint', async () => { const conversationId = `e2e-expiry-${Date.now()}`; const run = await buildHitlRun({ saver, @@ -336,12 +331,15 @@ describe('HITL checkpoint lifecycle (full wiring)', () => { }); await GenerationJobManager.approvals.pause(conversationId, pendingAction); - // The sweeper/stale-submit path: expiry fires the registered checkpoint prune. + // Expiry finalizes the stream, while checkpoint cleanup remains TTL-scoped. A + // thread-wide eager delete can race a replacement run on the same conversation. expect(await GenerationJobManager.expireApproval(conversationId, pendingAction.actionId)).toBe( true, ); expect(await GenerationJobManager.getJobStatus(conversationId)).toBe('aborted'); + expect((await checkpointCounts(conversationId)).checkpoints).toBeGreaterThan(0); + await deleteAgentCheckpoint(conversationId, MONGO_CFG); expect(await checkpointCounts(conversationId)).toEqual({ checkpoints: 0, writes: 0 }); }); }); diff --git a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js index addbdadc60..1d00544db4 100644 --- a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js +++ b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js @@ -16,6 +16,10 @@ const mockGenerationJobManager = { claimGeneration: jest.fn(), releaseGeneration: jest.fn(), hasJob: jest.fn(), + steering: { + closeAndDrain: jest.fn(), + park: jest.fn(), + }, }; const mockCheckAndIncrementPendingRequest = jest.fn(); @@ -32,6 +36,14 @@ const mockFilterPersistableAbortContent = jest.fn((content) => const mockGetConvo = jest.fn(); const mockGetMessages = jest.fn(); const mockSaveMessage = jest.fn(); +const mockStartupTelemetry = { + mark: jest.fn(), + setStreamId: jest.fn(), + recordGenerationEvent: jest.fn(), + end: jest.fn(), +}; +const mockGetAgentStartupTelemetry = jest.fn(() => mockStartupTelemetry); +const mockAcceptAgentStartupTelemetry = jest.fn(); let mockMCPContexts = new WeakMap(); const mockCreateMCPRequestContext = jest.fn(() => ({ @@ -94,6 +106,7 @@ jest.mock('@librechat/api', () => ({ getViolationInfo: (...args) => mockGetViolationInfo(...args), buildMessageFiles: jest.fn(() => []), resolveTitleTiming: jest.fn(() => 'immediate'), + resolveConversationAnchor: jest.requireActual('@librechat/api').resolveConversationAnchor, GenerationJobManager: mockGenerationJobManager, getReferencedQuotes: jest.fn((quotes) => { if (!Array.isArray(quotes)) { @@ -112,6 +125,8 @@ jest.mock('@librechat/api', () => ({ decrementPendingRequest: (...args) => mockDecrementPendingRequest(...args), sanitizeMessageForTransmit: jest.fn((message) => message), checkAndIncrementPendingRequest: (...args) => mockCheckAndIncrementPendingRequest(...args), + getAgentStartupTelemetry: (...args) => mockGetAgentStartupTelemetry(...args), + acceptAgentStartupTelemetry: (...args) => mockAcceptAgentStartupTelemetry(...args), isUnpersistedPreliminaryParent: async ({ userId, conversationId, @@ -155,6 +170,7 @@ jest.mock('~/models', () => ({ })); const AgentController = require('../request'); +const { disposeClient: mockDisposeClient } = require('~/server/cleanup'); const { getMCPRequestContext } = require('~/server/services/MCPRequestContext'); function createResumableResponse() { @@ -199,6 +215,8 @@ describe('ResumableAgentController resume metadata', () => { mockGenerationJobManager.claimGeneration.mockResolvedValue({ claimed: true }); mockGenerationJobManager.releaseGeneration.mockResolvedValue(undefined); mockGenerationJobManager.hasJob.mockResolvedValue(true); + mockGenerationJobManager.steering.closeAndDrain.mockResolvedValue([]); + mockGenerationJobManager.steering.park.mockResolvedValue(undefined); mockSaveMessage.mockResolvedValue({}); }); @@ -279,10 +297,17 @@ describe('ResumableAgentController resume metadata', () => { conversationId, 'user-123', conversationId, + expect.objectContaining({ + startupTelemetry: mockStartupTelemetry, + initialMetadata: expect.objectContaining({ + conversationId, + endpoint: 'agents', + }), + }), ); }); - it('stores the in-flight turn before MCP initialization can emit OAuth', async () => { + it('creates the job with the in-flight turn before MCP initialization can emit OAuth', async () => { const conversationId = 'conversation-123'; const initializeClient = jest.fn().mockRejectedValue(new Error('stop before tool loading')); const req = { @@ -292,6 +317,7 @@ describe('ResumableAgentController resume metadata', () => { messageId: 'follow-up-user', parentMessageId: 'original-response', conversationId, + isTemporary: true, endpointOption: { endpoint: 'agents', iconURL: 'https://example.com/spec-icon.png', @@ -310,25 +336,98 @@ describe('ResumableAgentController resume metadata', () => { await AgentController(req, res, jest.fn(), initializeClient, null); - expect(mockGenerationJobManager.updateMetadata).toHaveBeenCalledWith( + expect(mockGenerationJobManager.createJob).toHaveBeenCalledWith( conversationId, - expect.objectContaining({ - conversationId, - endpoint: 'agents', - iconURL: 'https://example.com/spec-icon.png', - model: 'gpt-3.5-turbo', - responseMessageId: 'follow-up-user_', - userMessage: { - messageId: 'follow-up-user', - parentMessageId: 'original-response', + 'user-123', + conversationId, + { + startupTelemetry: mockStartupTelemetry, + initialMetadata: { conversationId, - text: 'Check Google Workspace availability.', + endpoint: 'agents', + iconURL: 'https://example.com/spec-icon.png', + model: 'gpt-3.5-turbo', + agent_id: undefined, + isTemporary: true, + responseMessageId: 'follow-up-user_', + userMessage: { + messageId: 'follow-up-user', + parentMessageId: 'original-response', + conversationId, + text: 'Check Google Workspace availability.', + }, }, - }), + }, ); - expect(mockGenerationJobManager.updateMetadata.mock.invocationCallOrder[0]).toBeLessThan( + expect(mockGenerationJobManager.createJob.mock.invocationCallOrder[0]).toBeLessThan( initializeClient.mock.invocationCallOrder[0], ); + expect(mockGenerationJobManager.updateMetadata).not.toHaveBeenCalled(); + const startupMilestones = mockStartupTelemetry.mark.mock.calls.map(([milestone]) => milestone); + expect(startupMilestones.slice(0, 2)).toEqual(['request_admitted', 'job_created']); + expect(new Set(startupMilestones.slice(2))).toEqual( + new Set(['conversation_resolved', 'metadata_persisted']), + ); + expect(mockAcceptAgentStartupTelemetry).toHaveBeenCalledWith(req, conversationId); + expect(mockStartupTelemetry.end).toHaveBeenCalledWith('error', expect.any(Error)); + }); + + it('prefetches conversation state before admission and joins it with job metadata', async () => { + let resolveConversation; + let signalMetadataStarted; + const conversationPromise = new Promise((resolve) => { + resolveConversation = resolve; + }); + const metadataStarted = new Promise((resolve) => { + signalMetadataStarted = resolve; + }); + mockGetConvo.mockReturnValue(conversationPromise); + mockGenerationJobManager.createJob.mockImplementation(() => { + signalMetadataStarted(); + return Promise.resolve({ + createdAt: 1000, + readyPromise: Promise.resolve(), + abortController: new AbortController(), + emitter: { on: jest.fn() }, + }); + }); + const initializeClient = jest.fn().mockRejectedValue(new Error('stop after startup reads')); + const conversationId = 'conversation-123'; + const req = { + user: { id: 'user-123' }, + body: { + text: 'Run independent startup work together.', + messageId: 'user-message', + parentMessageId: 'parent-message', + conversationId, + endpointOption: { + endpoint: 'agents', + modelOptions: { model: 'gpt-4.1' }, + }, + }, + config: {}, + }; + const res = createResumableResponse(); + + const controllerPromise = AgentController(req, res, jest.fn(), initializeClient, null); + expect(mockGetConvo).toHaveBeenCalledWith('user-123', conversationId); + await metadataStarted; + await nextTick(); + + expect(mockGetConvo.mock.invocationCallOrder[0]).toBeLessThan( + mockCheckAndIncrementPendingRequest.mock.invocationCallOrder[0], + ); + expect(res.json).toHaveBeenCalledWith({ + streamId: conversationId, + conversationId, + status: 'started', + }); + expect(initializeClient).not.toHaveBeenCalled(); + + resolveConversation({ createdAt: '2026-06-07T00:00:00.000Z' }); + await controllerPromise; + + expect(initializeClient).toHaveBeenCalledTimes(1); }); it('keeps request-scoped MCP connections until resumable initialization finishes', async () => { @@ -382,6 +481,7 @@ describe('ResumableAgentController resume metadata', () => { messageId: 'follow-up-user', parentMessageId: 'original-response', conversationId, + isTemporary: true, endpointOption: { endpoint: 'agents', spec: 'agent-spec', @@ -413,11 +513,17 @@ describe('ResumableAgentController resume metadata', () => { await AgentController(req, res, jest.fn(), initializeClient, null); - expect(mockGenerationJobManager.updateMetadata).toHaveBeenCalledWith( + expect(mockGenerationJobManager.createJob).toHaveBeenCalledWith( + conversationId, + 'user-123', conversationId, expect.objectContaining({ - iconURL: 'https://example.com/preset-icon.png', - model: 'agent_resume_spec', + initialMetadata: expect.objectContaining({ + iconURL: 'https://example.com/preset-icon.png', + model: 'agent_resume_spec', + agent_id: 'agent_resume_spec', + isTemporary: true, + }), }), ); }); @@ -461,11 +567,15 @@ describe('ResumableAgentController resume metadata', () => { await AgentController(req, res, jest.fn(), initializeClient, null); - expect(mockGenerationJobManager.updateMetadata).toHaveBeenCalledWith( + expect(mockGenerationJobManager.createJob).toHaveBeenCalledWith( + conversationId, + 'user-123', conversationId, expect.objectContaining({ - iconURL: 'anthropic', - model: 'gpt-4.1', + initialMetadata: expect.objectContaining({ + iconURL: 'anthropic', + model: 'gpt-4.1', + }), }), ); }); @@ -662,6 +772,7 @@ describe('ResumableAgentController resume metadata', () => { expect(mockGenerationJobManager.createJob).not.toHaveBeenCalled(); expect(mockCheckAndIncrementPendingRequest).not.toHaveBeenCalled(); expect(initializeClient).not.toHaveBeenCalled(); + expect(mockStartupTelemetry.end).toHaveBeenCalledWith('deduplicated'); }); it('resumes when the job is missing but the claim is old (original completed and was cleaned up)', async () => { @@ -758,6 +869,32 @@ describe('ResumableAgentController resume metadata', () => { expect(mockGenerationJobManager.createJob).not.toHaveBeenCalled(); }); + it('does not finalize an unscoped generation when job creation rejects before returning', async () => { + mockGenerationJobManager.claimGeneration.mockResolvedValue({ claimed: true }); + mockGenerationJobManager.createJob.mockRejectedValue(new Error('create failed before return')); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Fail before receiving a job epoch.', + messageId: 'user-msg', + clientRequestId: 'req-abc', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = createResumableResponse(); + + await AgentController(req, res, jest.fn(), jest.fn(), null); + + expect(res.status).toHaveBeenCalledWith(500); + expect(res.json).toHaveBeenCalledWith({ error: 'create failed before return' }); + expect(mockGenerationJobManager.emitError).not.toHaveBeenCalled(); + expect(mockGenerationJobManager.completeJob).not.toHaveBeenCalled(); + expect(mockGenerationJobManager.releaseGeneration).toHaveBeenCalledWith('user-123', 'req-abc'); + expect(mockDecrementPendingRequest).toHaveBeenCalledWith('user-123'); + }); + it('finalizes the failed job before releasing the idempotency claim', async () => { mockGenerationJobManager.claimGeneration.mockResolvedValue({ claimed: true }); const initializeClient = jest.fn().mockRejectedValue(new Error('init boom after res.json')); @@ -779,6 +916,7 @@ describe('ResumableAgentController resume metadata', () => { expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith( 'conversation-123', expect.any(String), + 1000, ); expect(mockGenerationJobManager.releaseGeneration).toHaveBeenCalledWith('user-123', 'req-abc'); // completeJob must finalize the failed job BEFORE the claim is released, or a racing @@ -812,6 +950,150 @@ describe('ResumableAgentController resume metadata', () => { expect(mockDecrementPendingRequest).toHaveBeenCalledWith('user-123'); }); + it('still finalizes and releases when streaming the initialization error fails', async () => { + mockGenerationJobManager.claimGeneration.mockResolvedValue({ claimed: true }); + mockGenerationJobManager.emitError.mockRejectedValue(new Error('publish failed')); + const initializeClient = jest.fn().mockRejectedValue(new Error('init boom after res.json')); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Start fails while Redis publish is degraded.', + messageId: 'user-msg', + clientRequestId: 'req-abc', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = createResumableResponse(); + + await AgentController(req, res, jest.fn(), initializeClient, null); + + expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith( + 'conversation-123', + 'init boom after res.json', + 1000, + ); + expect(mockGenerationJobManager.releaseGeneration).toHaveBeenCalledWith('user-123', 'req-abc'); + expect(mockDecrementPendingRequest).toHaveBeenCalledWith('user-123'); + expect(mockStartupTelemetry.end).toHaveBeenCalledWith('error', expect.any(Error)); + }); + + it('finalizes and disposes a client aborted during initialization before releasing the slot', async () => { + const abortController = new AbortController(); + let resolveCompletion; + let signalCompletionStarted; + const completionStarted = new Promise((resolve) => { + signalCompletionStarted = resolve; + }); + mockGenerationJobManager.createJob.mockResolvedValue({ + createdAt: 1000, + readyPromise: Promise.resolve(), + abortController, + emitter: { on: jest.fn() }, + }); + mockGenerationJobManager.completeJob.mockImplementation(() => { + signalCompletionStarted(); + return new Promise((resolve) => { + resolveCompletion = resolve; + }); + }); + const client = { options: {} }; + const initializeClient = jest.fn(async ({ signal }) => { + expect(signal).toBe(abortController.signal); + abortController.abort(); + return { client }; + }); + const conversationId = 'conversation-123'; + const req = { + user: { id: 'user-123' }, + body: { + text: 'Stop during initialization.', + messageId: 'user-msg', + conversationId, + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = createResumableResponse(); + + const controllerPromise = AgentController(req, res, jest.fn(), initializeClient, null); + await completionStarted; + + expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith( + conversationId, + 'Request aborted during initialization', + 1000, + ); + expect(mockDecrementPendingRequest).not.toHaveBeenCalled(); + expect(mockDisposeClient).not.toHaveBeenCalled(); + + resolveCompletion(); + await controllerPromise; + + expect(mockDecrementPendingRequest).toHaveBeenCalledWith('user-123'); + expect(mockDisposeClient).toHaveBeenCalledTimes(1); + expect(mockDisposeClient).toHaveBeenCalledWith(client); + expect(mockStartupTelemetry.end).toHaveBeenCalledWith('aborted'); + }); + + it('awaits background error finalization before releasing the slot and always disposes', async () => { + const generationError = new Error('generation failed'); + let rejectCompletion; + let signalCompletionStarted; + const completionStarted = new Promise((resolve) => { + signalCompletionStarted = resolve; + }); + mockGenerationJobManager.emitError.mockRejectedValue(new Error('publish failed')); + mockGenerationJobManager.completeJob.mockImplementation(() => { + signalCompletionStarted(); + return new Promise((_, reject) => { + rejectCompletion = reject; + }); + }); + const client = { + options: {}, + sendMessage: jest.fn().mockRejectedValue(generationError), + }; + const initializeClient = jest.fn().mockResolvedValue({ client }); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Fail after initialization.', + messageId: 'user-msg', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = createResumableResponse(); + + await AgentController(req, res, jest.fn(), initializeClient, null); + await completionStarted; + + expect(mockGenerationJobManager.emitError).toHaveBeenCalledWith( + 'conversation-123', + generationError.message, + 1000, + ); + expect(mockDecrementPendingRequest).not.toHaveBeenCalled(); + expect(mockDisposeClient).not.toHaveBeenCalled(); + + rejectCompletion(new Error('store failed')); + await nextTick(); + + expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith( + 'conversation-123', + generationError.message, + 1000, + ); + expect(mockGenerationJobManager.completeJob.mock.invocationCallOrder[0]).toBeLessThan( + mockDecrementPendingRequest.mock.invocationCallOrder[0], + ); + expect(mockDecrementPendingRequest).toHaveBeenCalledWith('user-123'); + expect(mockDisposeClient).toHaveBeenCalledWith(client); + }); + it('proceeds to create the job when it wins the idempotency claim', async () => { mockGenerationJobManager.claimGeneration.mockResolvedValue({ claimed: true }); const initializeClient = jest.fn().mockRejectedValue(new Error('stop before tool loading')); @@ -842,6 +1124,13 @@ describe('ResumableAgentController resume metadata', () => { 'conversation-123', 'user-123', 'conversation-123', + expect.objectContaining({ + startupTelemetry: mockStartupTelemetry, + initialMetadata: expect.objectContaining({ + conversationId: 'conversation-123', + endpoint: 'agents', + }), + }), ); }); @@ -869,6 +1158,7 @@ describe('ResumableAgentController resume metadata', () => { expect(res.status).toHaveBeenCalledWith(429); expect(mockGenerationJobManager.releaseGeneration).toHaveBeenCalledWith('user-123', 'req-abc'); + expect(mockStartupTelemetry.end).toHaveBeenCalledWith('rejected'); }); it('does not release a claim it never won when a fail-open duplicate hits the limiter', async () => { diff --git a/api/server/controllers/agents/__tests__/resume.spec.js b/api/server/controllers/agents/__tests__/resume.spec.js index c7fb4dbef3..5ad47f8912 100644 --- a/api/server/controllers/agents/__tests__/resume.spec.js +++ b/api/server/controllers/agents/__tests__/resume.spec.js @@ -59,6 +59,7 @@ const mockGenerationJobManager = { }; const mockDeleteAgentCheckpoint = jest.fn(); +const mockCaptureAgentCheckpointGeneration = jest.fn(); const mockDecrementPendingRequest = jest.fn(); const mockCheckAndIncrementPendingRequest = jest.fn(); @@ -77,6 +78,7 @@ jest.mock('@librechat/data-schemas', () => ({ jest.mock('@librechat/api', () => ({ ...jest.requireActual('@librechat/api'), GenerationJobManager: mockGenerationJobManager, + captureAgentCheckpointGeneration: (...args) => mockCaptureAgentCheckpointGeneration(...args), deleteAgentCheckpoint: (...args) => mockDeleteAgentCheckpoint(...args), decrementPendingRequest: (...args) => mockDecrementPendingRequest(...args), checkAndIncrementPendingRequest: (...args) => mockCheckAndIncrementPendingRequest(...args), @@ -109,6 +111,7 @@ function makeToolApprovalJob(overrides = {}) { const pendingOverrides = metaOverrides.pendingAction ?? {}; return { status: 'requires_action', + createdAt: 1000, abortController: new AbortController(), ...overrides, metadata: { @@ -179,11 +182,19 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { mockCheckAndIncrementPendingRequest.mockResolvedValue({ allowed: true }); mockDecrementPendingRequest.mockResolvedValue(undefined); mockDeleteAgentCheckpoint.mockResolvedValue(undefined); + mockCaptureAgentCheckpointGeneration.mockResolvedValue({ + threadId: CONVO_ID, + checkpointIds: ['checkpoint-old'], + }); mockCleanupMCPRequestContextForReq.mockResolvedValue(undefined); mockSaveMessage.mockResolvedValue(undefined); mockGetConvo.mockResolvedValue(null); mockGetMessages.mockResolvedValue([]); - mockJobStore.getJob.mockResolvedValue({ tokenUsage: null, contextUsage: null }); + mockJobStore.getJob.mockResolvedValue({ + createdAt: 1000, + tokenUsage: null, + contextUsage: null, + }); mockJobStore.updateJob.mockResolvedValue(undefined); mockGenerationJobManager.getResumeState.mockResolvedValue({ aggregatedContent: [] }); mockGenerationJobManager.emitDone.mockResolvedValue(undefined); @@ -555,6 +566,21 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { expect(mockGenerationJobManager.approvals.resolve).not.toHaveBeenCalled(); }); + it('consumes a checkpoint-snapshot rejection on the 429 early-return path', async () => { + mockGenerationJobManager.getJob.mockResolvedValue(makeToolApprovalJob()); + mockCheckAndIncrementPendingRequest.mockResolvedValue({ allowed: false }); + mockCaptureAgentCheckpointGeneration.mockRejectedValue(new Error('mongo down')); + + const res = await post(approveBody()); + await flush(); + + expect(res.status).toBe(429); + expect(mockLogger.warn).toHaveBeenCalledWith( + '[ResumeAgentController] Failed to capture checkpoint generation', + expect.any(Error), + ); + }); + it('409 and releases the slot when the action was already claimed (single-winner)', async () => { mockGenerationJobManager.getJob.mockResolvedValue(makeToolApprovalJob()); mockGenerationJobManager.approvals.resolve.mockResolvedValue(false); @@ -589,7 +615,13 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { conversationId: CONVO_ID, status: 'resuming', }); + expect(mockCaptureAgentCheckpointGeneration).toHaveBeenCalledWith(CONVO_ID, { + type: 'mongo', + }); expect(mockGenerationJobManager.approvals.resolve).toHaveBeenCalledWith(CONVO_ID, ACTION_ID); + expect(mockCaptureAgentCheckpointGeneration.mock.invocationCallOrder[0]).toBeLessThan( + mockGenerationJobManager.approvals.resolve.mock.invocationCallOrder[0], + ); await settled; await flush(); }); @@ -732,12 +764,31 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { }); expect(typeof finalEvent.title).toBe('string'); - expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith(CONVO_ID); - expect(mockDeleteAgentCheckpoint).toHaveBeenCalledWith(CONVO_ID, { type: 'mongo' }); + expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith(CONVO_ID, undefined, 1000); + expect(mockDeleteAgentCheckpoint).toHaveBeenCalledWith( + CONVO_ID, + { type: 'mongo' }, + { threadId: CONVO_ID, checkpointIds: ['checkpoint-old'] }, + ); expect(mockDecrementPendingRequest).toHaveBeenCalledWith(USER_ID); expect(mockDisposeClient).toHaveBeenCalledTimes(1); }); + it('degrades a failed checkpoint snapshot to scoped no-op cleanup', async () => { + mockGenerationJobManager.getJob.mockResolvedValue(makeToolApprovalJob()); + mockCaptureAgentCheckpointGeneration.mockRejectedValue(new Error('mongo down')); + + await post(approveBody()); + await settled; + await flush(); + + expect(mockDeleteAgentCheckpoint).toHaveBeenCalledWith( + CONVO_ID, + { type: 'mongo' }, + { threadId: CONVO_ID, checkpointIds: [] }, + ); + }); + it('skips finalization (no save/emitDone/complete) when the job was replaced mid-resume', async () => { // The paused job has createdAt 1000; a concurrent request reused this conversationId, // so the live job now has a different createdAt — finalizing would clobber the newer @@ -941,7 +992,7 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { expect(client.resumeCompletion).toHaveBeenCalledWith( expect.objectContaining({ resumeValue: { answer: 'call it report.pdf' } }), ); - expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith(CONVO_ID); + expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith(CONVO_ID, undefined, 1000); }); it('generates a title for a first-turn pause before completing the stream', async () => { @@ -956,7 +1007,7 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { expect(mockAddTitle).toHaveBeenCalledTimes(1); // Title is emitted (and the job completed) — order matters but both must happen. - expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith(CONVO_ID); + expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith(CONVO_ID, undefined, 1000); }); it('still finalizes the turn when first-turn title generation throws', async () => { @@ -973,8 +1024,12 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { expect(mockLogger.error).toHaveBeenCalled(); expect(mockSaveMessage).toHaveBeenCalledTimes(1); - expect(mockGenerationJobManager.emitDone).toHaveBeenCalledWith(CONVO_ID, expect.any(Object)); - expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith(CONVO_ID); + expect(mockGenerationJobManager.emitDone).toHaveBeenCalledWith( + CONVO_ID, + expect.any(Object), + 1000, + ); + expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith(CONVO_ID, undefined, 1000); }); }); @@ -1095,9 +1150,13 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { await settled; await flush(); - expect(mockGenerationJobManager.emitError).toHaveBeenCalledWith(CONVO_ID, 'boom'); - expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith(CONVO_ID, 'boom'); - expect(mockDeleteAgentCheckpoint).toHaveBeenCalledWith(CONVO_ID, { type: 'mongo' }); + expect(mockGenerationJobManager.emitError).toHaveBeenCalledWith(CONVO_ID, 'boom', 1000); + expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith(CONVO_ID, 'boom', 1000); + expect(mockDeleteAgentCheckpoint).toHaveBeenCalledWith( + CONVO_ID, + { type: 'mongo' }, + { threadId: CONVO_ID, checkpointIds: ['checkpoint-old'] }, + ); expect(mockDecrementPendingRequest).toHaveBeenCalledWith(USER_ID); expect(mockSaveMessage).not.toHaveBeenCalled(); }); @@ -1121,6 +1180,7 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { expect(mockJobStore.updateJob).toHaveBeenCalledWith( CONVO_ID, expect.objectContaining({ status: 'error', error: 'Resume failed' }), + 1000, ); expect(mockDecrementPendingRequest).toHaveBeenCalledWith(USER_ID); }); diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index a4958713e7..7db770bf37 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -123,6 +123,8 @@ class AgentClient extends BaseClient { /** @deprecated @type {true} - Is a Chat Completion Request */ this.isChatCompletion = true; + /** @type {number | undefined} */ + this.jobCreatedAt = options.jobCreatedAt; /** @type {AgentRun} */ this.run; @@ -429,7 +431,37 @@ class AgentClient extends BaseClient { })) : []), ]; + + /** + * Memory authorization/loading and MCP config resolution do not depend on + * attachment hydration or prompt formatting. Start them before that work, + * but keep the existing context-application barrier below. + * + * Attach a rejection observer immediately because these operations may + * settle while request attachments are still being prepared. Awaiting the + * original promise later still propagates either error. + */ + const earlySharedContextPromise = Promise.all([ + this.useMemory(), + resolveConfigServers(this.options.req), + ]); + void earlySharedContextPromise.catch(() => {}); + const sharedRunAttachmentIds = new Set(); + /** @type {ReturnType} */ + let agentScopedContextPromise; + const startAgentScopedContext = () => { + const contextPromise = buildAgentScopedContext({ + agentIds: allAgents.map(({ agentId }) => agentId), + attachmentsByAgentId: this.options.agentContextAttachmentsByAgentId, + sharedRunAttachmentIds, + req: this.options.req, + tokenCountFn: (text) => countTokens(text), + }); + void contextPromise.catch(() => {}); + return contextPromise; + }; + if (this.options.attachments) { const attachments = await this.options.attachments; const latestMessage = orderedMessages[orderedMessages.length - 1]; @@ -438,6 +470,9 @@ class AgentClient extends BaseClient { sharedRunAttachmentIds.add(fileId); } + /** Agent-scoped extraction only depends on the shared attachment IDs. */ + agentScopedContextPromise = startAgentScopedContext(); + if (this.message_file_map) { this.message_file_map[latestMessage.messageId] = attachments; } else { @@ -446,10 +481,14 @@ class AgentClient extends BaseClient { }; } - await this.addFileContextToMessage(latestMessage, attachments); - const files = await this.processAttachments(latestMessage, attachments); + const [, files] = await Promise.all([ + this.addFileContextToMessage(latestMessage, attachments), + this.processAttachments(latestMessage, attachments), + ]); this.options.attachments = files; + } else { + agentScopedContextPromise = startAgentScopedContext(); } /** Note: Bedrock uses legacy RAG API handling */ @@ -655,19 +694,21 @@ class AgentClient extends BaseClient { * Memory context is handled separately and applied per-agent based on config. */ const sharedRunContextParts = []; + const [augmentedPrompt, [memories, configServers], agentScopedContext] = await Promise.all([ + this.contextHandlers?.createContext(), + earlySharedContextPromise, + agentScopedContextPromise, + ]); /** Augmented prompt from RAG/context handlers */ - if (this.contextHandlers) { - this.augmentedPrompt = await this.contextHandlers.createContext(); - if (this.augmentedPrompt) { - sharedRunContextParts.push(this.augmentedPrompt); - } + this.augmentedPrompt = augmentedPrompt; + if (this.augmentedPrompt) { + sharedRunContextParts.push(this.augmentedPrompt); } /** Memory context (user preferences/memories). Keyed context (with memory * keys + token metadata) is reserved for agents that can call * `delete_memory`; everyone else gets the unkeyed values only. */ - const memories = await this.useMemory(); /** Partition the loaded memories belong to (the primary agent's). */ const loadedMemoryAgentId = getMemoryAgentId(this.options.agent); const buildMemoryContext = (text) => @@ -700,14 +741,6 @@ class AgentClient extends BaseClient { const sharedRunContext = sharedRunContextParts.join('\n\n'); const memoryAgentEnabled = isMemoryAgentEnabled(this.options.req.config?.memory); - const agentScopedContext = await buildAgentScopedContext({ - agentIds: allAgents.map(({ agentId }) => agentId), - attachmentsByAgentId: this.options.agentContextAttachmentsByAgentId, - sharedRunAttachmentIds, - req: this.options.req, - tokenCountFn: (text) => countTokens(text), - }); - /** Preserve prompt token counts for graph formatting and pruning. */ this.indexTokenCountMap = indexTokenCountMap; @@ -741,8 +774,6 @@ class AgentClient extends BaseClient { const ephemeralAgent = this.options.req.body.ephemeralAgent; const mcpManager = getMCPManager(); - const configServers = await resolveConfigServers(this.options.req); - await Promise.all( allAgents.map(async ({ agent, agentId }) => { const agentRunContextParts = [sharedRunContext]; @@ -1471,9 +1502,13 @@ class AgentClient extends BaseClient { if (Array.isArray(runMessages) && runMessages.length > 0) { const discovered = extractDiscoveredToolsFromHistory(runMessages); if (discovered.size > 0) { - await GenerationJobManager.updateMetadata(streamId, { - discoveredTools: Array.from(discovered), - }); + await GenerationJobManager.updateMetadata( + streamId, + { + discoveredTools: Array.from(discovered), + }, + this.jobCreatedAt, + ); } } } catch (err) { @@ -1762,7 +1797,29 @@ class AgentClient extends BaseClient { } const streamId = this.options.req?._resumableStreamId; - run = await createRun({ + // HITL: clear any checkpoint orphaned by a prior paused turn in this + // conversation (one that expired or was aborted while paused) so this fresh + // turn starts clean instead of rehydrating a stale interrupt — thread_id is + // the stable conversationId. No-op when HITL is off or nothing is orphaned. + // Deliberately UNCONDITIONAL per HITL turn: any cheaper gate (job metadata, + // a Redis flag) can go stale across replicas/restarts and skip the prune + // exactly when an orphan exists, while these are two indexed, usually-empty + // deleteMany ops — correctness over a micro-optimization. + // The gate mirrors createRun's checkpointer condition: the approval policy + // OR an ask_user_question-capable agent (which attaches a checkpointer + // WITHOUT the approval policy) — an ask pause abandoned via job replacement + // or Stop would otherwise rehydrate here and silently duplicate context. + // + // Start the prune alongside graph construction. The all-settled barrier + // below still guarantees it completes before the graph is exposed or run. + const shouldPruneCheckpoint = + streamId && + (isHITLEnabled(agentsEConfig?.toolApproval) || agents.some(agentRequestsAskUserQuestion)); + const checkpointPrunePromise = shouldPruneCheckpoint + ? deleteAgentCheckpoint(this.conversationId, agentsEConfig?.checkpointer) + : Promise.resolve(); + + const createRunPromise = createRun({ agents, messages, // This controller implements the full HITL pause/resume lifecycle (handleRunInterrupt @@ -1802,11 +1859,25 @@ class AgentClient extends BaseClient { this.collectedUsage, this.buildSubagentUsageEmitter(appConfig), ), + }).then((createdRun) => { + if (!createdRun) { + throw new Error('Failed to create run'); + } + this.options.startupTelemetry?.mark('run_created'); + return createdRun; }); - if (!run) { - throw new Error('Failed to create run'); + const [createRunResult, checkpointPruneResult] = await Promise.allSettled([ + createRunPromise, + checkpointPrunePromise, + ]); + if (createRunResult.status === 'rejected') { + throw createRunResult.reason; } + if (checkpointPruneResult.status === 'rejected') { + throw checkpointPruneResult.reason; + } + run = createRunResult.value; this.run = run; if (this._resolveRun) { @@ -1815,7 +1886,7 @@ class AgentClient extends BaseClient { } if (streamId && run.Graph) { - GenerationJobManager.setGraph(streamId, run.Graph); + GenerationJobManager.setGraph(streamId, run.Graph, this.jobCreatedAt); } if (userMCPAuthMap != null) { @@ -1825,25 +1896,7 @@ class AgentClient extends BaseClient { /** @deprecated Agent Chain */ config.configurable.last_agent_id = agents[agents.length - 1].id; - // HITL: clear any checkpoint orphaned by a prior paused turn in this - // conversation (one that expired or was aborted while paused) so this fresh - // turn starts clean instead of rehydrating a stale interrupt — thread_id is - // the stable conversationId. No-op when HITL is off or nothing is orphaned. - // Deliberately UNCONDITIONAL per HITL turn: any cheaper gate (job metadata, - // a Redis flag) can go stale across replicas/restarts and skip the prune - // exactly when an orphan exists, while these are two indexed, usually-empty - // deleteMany ops — correctness over a micro-optimization. - // The gate mirrors createRun's checkpointer condition: the approval policy - // OR an ask_user_question-capable agent (which attaches a checkpointer - // WITHOUT the approval policy) — an ask pause abandoned via job replacement - // or Stop would otherwise rehydrate here and silently duplicate context. - if ( - streamId && - (isHITLEnabled(agentsEConfig?.toolApproval) || agents.some(agentRequestsAskUserQuestion)) - ) { - await deleteAgentCheckpoint(this.conversationId, agentsEConfig?.checkpointer); - } - + this.options.startupTelemetry?.mark('stream_processing_started'); await run.processStream({ messages }, config, { callbacks: { [Callback.TOOL_ERROR]: logToolError, @@ -1858,6 +1911,7 @@ class AgentClient extends BaseClient { config.signal = null; }; + this.options.startupTelemetry?.mark('run_input_prepared'); await runAgents(initialMessages); /** @@ -2162,7 +2216,7 @@ class AgentClient extends BaseClient { // introspection fall back to the durable chunk reconstruction, which is complete. // `setContentParts` still points the in-memory store at the seeded client content. if (streamId && this.contentParts) { - GenerationJobManager.setContentParts(streamId, this.contentParts); + GenerationJobManager.setContentParts(streamId, this.contentParts, this.jobCreatedAt); } // Carry the user's MCP auth into the rebuilt run so an approved MCP tool executes diff --git a/api/server/controllers/agents/client.test.js b/api/server/controllers/agents/client.test.js index 1646a448a7..f07ec03e9a 100644 --- a/api/server/controllers/agents/client.test.js +++ b/api/server/controllers/agents/client.test.js @@ -1,6 +1,28 @@ +const mockCreateRun = jest.fn(); +const mockDeleteAgentCheckpoint = jest.fn(); +const mockIsHITLEnabled = jest.fn().mockReturnValue(false); +const mockBuildAgentScopedContext = jest.fn((...args) => + jest.requireActual('@librechat/api').buildAgentScopedContext(...args), +); +const mockFormatAgentMessages = jest.fn(() => ({ + messages: [], + indexTokenCountMap: {}, + summary: undefined, + boundaryTokenAdjustment: undefined, +})); + const { Providers } = require('@librechat/agents'); const { Constants, ContentTypes, EModelEndpoint } = require('librechat-data-provider'); const AgentClient = require('./client'); +const { resolveConfigServers } = require('~/server/services/MCP'); + +function deferred() { + let resolve; + const promise = new Promise((resolvePromise) => { + resolve = resolvePromise; + }); + return { promise, resolve }; +} jest.mock('@librechat/agents', () => ({ ...jest.requireActual('@librechat/agents'), @@ -8,14 +30,20 @@ jest.mock('@librechat/agents', () => ({ handleLLMEnd: jest.fn(), collected: [], }), + formatAgentMessages: (...args) => mockFormatAgentMessages(...args), })); jest.mock('@librechat/api', () => ({ ...jest.requireActual('@librechat/api'), + buildAgentScopedContext: (...args) => mockBuildAgentScopedContext(...args), checkAccess: jest.fn(), + createRun: (...args) => mockCreateRun(...args), countFormattedMessageTokens: jest.fn(() => 42), countTokens: jest.fn((text) => Math.ceil(String(text ?? '').length / 4)), + createTokenCounter: jest.fn(() => jest.fn(() => 0)), + deleteAgentCheckpoint: (...args) => mockDeleteAgentCheckpoint(...args), initializeAgent: jest.fn(), + isHITLEnabled: (...args) => mockIsHITLEnabled(...args), createMemoryProcessor: jest.fn(), isMemoryAgentEnabled: jest.fn((config) => { if (!config || config.disabled === true) return false; @@ -24,6 +52,7 @@ jest.mock('@librechat/api', () => ({ return Boolean(agent.id || (agent.provider && agent.model)); }), loadAgent: jest.fn(), + maybePrewarmCodeSandbox: jest.fn(), })); jest.mock('~/server/services/Config', () => ({ @@ -74,6 +103,94 @@ describe('AgentClient - applyHideSequentialOutputsFilter', () => { }); }); +describe('AgentClient - startup telemetry', () => { + it('overlaps run creation with checkpoint pruning and joins both before stream processing', async () => { + let releaseCheckpoint; + let checkpointStarted; + const runCreation = deferred(); + const checkpointStartedPromise = new Promise((resolve) => { + checkpointStarted = resolve; + }); + const checkpointPromise = new Promise((resolve) => { + releaseCheckpoint = resolve; + }); + const processStream = jest.fn().mockResolvedValue(); + const run = { + Graph: null, + processStream, + getCalibrationRatio: jest.fn(() => 0), + }; + const startupTelemetry = { + mark: jest.fn(), + setStreamId: jest.fn(), + recordGenerationEvent: jest.fn(), + end: jest.fn(), + }; + mockCreateRun.mockReturnValue(runCreation.promise); + mockIsHITLEnabled.mockReturnValue(true); + mockDeleteAgentCheckpoint.mockImplementation(() => { + checkpointStarted(); + return checkpointPromise; + }); + + const client = new AgentClient({ + req: { + user: { id: 'user-123' }, + body: {}, + config: { endpoints: { [EModelEndpoint.agents]: { toolApproval: {} } } }, + _resumableStreamId: 'conversation-123', + }, + res: {}, + agent: { + id: 'agent-123', + endpoint: EModelEndpoint.openAI, + provider: EModelEndpoint.openAI, + model_parameters: { model: 'gpt-4' }, + hide_sequential_outputs: false, + }, + endpointTokenConfig: {}, + eventHandlers: {}, + contentParts: [], + collectedUsage: [], + artifactPromises: [], + startupTelemetry, + }); + client.conversationId = 'conversation-123'; + client.responseMessageId = 'response-123'; + client.parentMessageId = 'parent-123'; + client.recordCollectedUsage = jest.fn().mockResolvedValue(); + + const completionPromise = client.chatCompletion({ payload: [] }); + await checkpointStartedPromise; + + expect(mockCreateRun).toHaveBeenCalledTimes(1); + expect(mockDeleteAgentCheckpoint).toHaveBeenCalledWith('conversation-123', undefined); + expect(startupTelemetry.mark.mock.calls.map(([milestone]) => milestone)).toEqual([ + 'run_input_prepared', + ]); + expect(processStream).not.toHaveBeenCalled(); + + runCreation.resolve(run); + await Promise.resolve(); + + expect(startupTelemetry.mark.mock.calls.map(([milestone]) => milestone)).toEqual([ + 'run_input_prepared', + 'run_created', + ]); + expect(processStream).not.toHaveBeenCalled(); + + releaseCheckpoint(); + await completionPromise; + + expect(startupTelemetry.mark.mock.calls.map(([milestone]) => milestone)).toEqual([ + 'run_input_prepared', + 'run_created', + 'stream_processing_started', + ]); + expect(processStream).toHaveBeenCalledTimes(1); + }); +}); + describe('AgentClient - titleConvo', () => { let client; let mockRun; @@ -1423,6 +1540,112 @@ describe('AgentClient - titleConvo', () => { client.maxContextTokens = 4096; }); + it('loads RAG, memory, attachment, and MCP context without serial waits', async () => { + const ragContext = deferred(); + const memoryContext = deferred(); + const mcpConfig = deferred(); + client.contextHandlers = { + createContext: jest.fn(() => ragContext.promise), + }; + client.useMemory = jest.fn(() => memoryContext.promise); + resolveConfigServers.mockReturnValueOnce(mcpConfig.promise); + + const buildPromise = client.buildMessages( + [ + { + messageId: 'msg-1', + parentMessageId: null, + sender: 'User', + text: 'Load all context.', + isCreatedByUser: true, + }, + ], + null, + {}, + ); + + expect(client.contextHandlers.createContext).toHaveBeenCalledTimes(1); + expect(client.useMemory).toHaveBeenCalledTimes(1); + expect(resolveConfigServers).toHaveBeenCalledWith(mockReq); + + ragContext.resolve('Retrieved context'); + memoryContext.resolve(undefined); + mcpConfig.resolve({}); + await buildPromise; + + expect(client.augmentedPrompt).toBe('Retrieved context'); + expect(client.options.agent.additional_instructions).toContain('Retrieved context'); + }); + + it('starts independent context and current-file work at their earliest dependency barriers', async () => { + const requestAttachments = deferred(); + const memoryContext = deferred(); + const mcpConfig = deferred(); + const agentScopedContext = deferred(); + const fileContext = deferred(); + const providerAttachments = deferred(); + const requestFile = { + file_id: 'request-file', + filename: 'request.txt', + source: 'text', + type: 'text/plain', + }; + + client.options.attachments = requestAttachments.promise; + client.useMemory = jest.fn(() => memoryContext.promise); + resolveConfigServers.mockReturnValueOnce(mcpConfig.promise); + mockBuildAgentScopedContext.mockReturnValueOnce(agentScopedContext.promise); + client.addFileContextToMessage = jest.fn(() => fileContext.promise); + client.processAttachments = jest.fn(() => providerAttachments.promise); + + const buildPromise = client.buildMessages( + [ + { + messageId: 'msg-early-context', + parentMessageId: null, + sender: 'User', + text: 'Load the request file.', + isCreatedByUser: true, + }, + ], + 'msg-early-context', + {}, + ); + + expect(client.useMemory).toHaveBeenCalledTimes(1); + expect(resolveConfigServers).toHaveBeenCalledWith(mockReq); + expect(mockBuildAgentScopedContext).not.toHaveBeenCalled(); + expect(client.addFileContextToMessage).not.toHaveBeenCalled(); + expect(client.processAttachments).not.toHaveBeenCalled(); + + requestAttachments.resolve([requestFile]); + await Promise.resolve(); + + expect(mockBuildAgentScopedContext).toHaveBeenCalledTimes(1); + const scopedContextArgs = mockBuildAgentScopedContext.mock.calls[0][0]; + expect([...scopedContextArgs.sharedRunAttachmentIds]).toEqual(['request-file']); + expect(client.addFileContextToMessage).toHaveBeenCalledWith( + expect.objectContaining({ messageId: 'msg-early-context' }), + [requestFile], + ); + expect(client.processAttachments).toHaveBeenCalledWith( + expect.objectContaining({ messageId: 'msg-early-context' }), + [requestFile], + ); + + providerAttachments.resolve([requestFile]); + await Promise.resolve(); + expect(client.options.attachments).toBe(requestAttachments.promise); + + fileContext.resolve(); + memoryContext.resolve(undefined); + mcpConfig.resolve({}); + agentScopedContext.resolve(new Map()); + await buildPromise; + + expect(client.options.attachments).toEqual([requestFile]); + }); + it('should await MCP instructions and not include [object Promise] in agent instructions', async () => { // Set specific return value for this test mockFormatInstructions.mockResolvedValue( diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index 49d9329f0c..a77cdedf66 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -13,6 +13,9 @@ const { sanitizeMessageForTransmit, checkAndIncrementPendingRequest, isUnpersistedPreliminaryParent, + resolveConversationAnchor, + getAgentStartupTelemetry, + acceptAgentStartupTelemetry, } = require('@librechat/api'); const { disposeClient, clientRegistry, requestDataMap } = require('~/server/cleanup'); const { @@ -41,44 +44,24 @@ function createCloseHandler(abortController) { }; } -function toValidISOString(value) { - if (value == null) { - return null; - } - - const date = value instanceof Date ? value : new Date(value); - return Number.isNaN(date.getTime()) ? null : date.toISOString(); -} - -async function resolveConversationCreatedAt({ userId, conversationId, isNewConvo }) { - if (isNewConvo) { - return { createdAt: new Date().toISOString(), conversation: undefined }; - } - - try { - const conversation = await getConvo(userId, conversationId); - return { - conversation, - createdAt: toValidISOString(conversation?.createdAt) ?? new Date().toISOString(), - }; - } catch (error) { - logger.warn('[AgentController] Failed to resolve conversation timestamp anchor', { - conversationId, - error: error?.message ?? error, - }); - return { createdAt: new Date().toISOString(), conversation: undefined }; - } -} - -async function attachConversationCreatedAt(req, { userId, conversationId, isNewConvo }) { - req.body.conversationId = conversationId; - const resolved = await resolveConversationCreatedAt({ - userId, - conversationId, - isNewConvo, +function resolveConversationCreatedAt({ userId, conversationId, isNewConvo }) { + return resolveConversationAnchor({ + isNewConversation: isNewConvo, + loadConversation: () => getConvo(userId, conversationId), + onLoadError: (error) => { + logger.warn('[AgentController] Failed to resolve conversation timestamp anchor', { + conversationId, + error: error.message, + }); + }, }); +} + +async function attachConversationCreatedAt(req, conversationId, conversationAnchorPromise) { + req.body.conversationId = conversationId; + const resolved = await conversationAnchorPromise; req.conversationCreatedAt = resolved.createdAt; - if (!isNewConvo && resolved.conversation !== undefined) { + if (resolved.conversation !== undefined) { req.resolvedConversation = resolved.conversation ?? null; } } @@ -212,6 +195,7 @@ function rejectPreliminaryParentMessageId(res) { * Returns streamId immediately, client subscribes separately via SSE. */ const ResumableAgentController = async (req, res, next, initializeClient, addTitle) => { + const startupTelemetry = getAgentStartupTelemetry(req); const { text, isRegenerate, @@ -225,6 +209,13 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit } = req.body; const userId = req.user.id; + const isNewConvo = !reqConversationId || reqConversationId === 'new'; + const conversationId = isNewConvo ? crypto.randomUUID() : reqConversationId; + const conversationAnchorPromise = resolveConversationCreatedAt({ + userId, + conversationId, + isNewConvo, + }); if ( await isUnpersistedPreliminaryParent({ @@ -234,6 +225,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit getMessages, }) ) { + startupTelemetry?.end('rejected'); return rejectPreliminaryParentMessageId(res); } @@ -245,8 +237,6 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit // Generate conversationId upfront if not provided - streamId === conversationId always // Treat "new" as a placeholder that needs a real UUID (frontend may send "new" for new convos) - const isNewConvo = !reqConversationId || reqConversationId === 'new'; - const conversationId = isNewConvo ? crypto.randomUUID() : reqConversationId; const streamId = conversationId; req.body.conversationId = conversationId; @@ -296,6 +286,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit err, ); res.set('Retry-After', '1'); + startupTelemetry?.end('deduplicated'); return res.status(503).json({ code: 'SERVER_NOT_READY', error: 'Generation is still starting. Please retry shortly.', @@ -308,6 +299,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit // the winner goes on to generate and bill with no UI attached — ask the client to // retry via the readiness path instead. res.set('Retry-After', '1'); + startupTelemetry?.end('deduplicated'); return res.status(503).json({ code: 'SERVER_NOT_READY', error: 'Generation is still starting. Please retry shortly.', @@ -322,6 +314,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit clientRequestId, streamId: existingStreamId, }); + startupTelemetry?.end('deduplicated'); return res.json({ streamId: existingStreamId, conversationId: claim.existing.conversationId, @@ -337,10 +330,13 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit } const violationInfo = getViolationInfo(pendingRequests, limit); await logViolation(req, res, ViolationTypes.CONCURRENT, violationInfo, violationInfo.score); + startupTelemetry?.end('rejected'); return res.status(429).json(violationInfo); } + startupTelemetry?.mark('request_admitted'); let client = null; + let jobCreatedAt; try { logger.debug(`[ResumableAgentController] Creating job`, { @@ -350,8 +346,31 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit userId, }); - const job = await GenerationJobManager.createJob(streamId, userId, conversationId); - const jobCreatedAt = job.createdAt; // Capture creation time to detect job replacement + const endpointIconURL = getEndpointIconURL(req, endpointOption); + const responseModel = getAgentResponseModel(req, endpointOption); + const preliminaryUserMessage = getPreliminaryUserMessage(req.body, conversationId); + const preliminaryResponseMessageId = getPreliminaryResponseMessageId(req.body); + const job = await GenerationJobManager.createJob(streamId, userId, conversationId, { + startupTelemetry, + initialMetadata: { + conversationId, + endpoint: endpointOption.endpoint, + iconURL: endpointIconURL, + model: responseModel, + // Persist the originating agent so a HITL resume can refuse to rebuild this + // paused run on a different agent (see resume.js). + agent_id: endpointOption.agent_id ?? req.body?.agent_id, + // Persist temporary-chat state so a HITL resume keeps the resumed response + // non-persisted instead of trusting the resume request to re-send the flag. + isTemporary: req.body?.isTemporary, + responseMessageId: preliminaryResponseMessageId, + userMessage: preliminaryUserMessage, + }, + }); + startupTelemetry?.mark('job_created'); + acceptAgentStartupTelemetry(req, streamId); + startupTelemetry?.mark('metadata_persisted'); + jobCreatedAt = job.createdAt; // Capture creation time to detect job replacement req._resumableStreamId = streamId; getMCPRequestContext(req, undefined, { cleanupOnResponse: false }); @@ -359,26 +378,9 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit // This is critical: tool loading (MCP OAuth) may emit events that the client needs to receive res.json({ streamId, conversationId, status: 'started' }); - await attachConversationCreatedAt(req, { userId, conversationId, isNewConvo }); - - const endpointIconURL = getEndpointIconURL(req, endpointOption); - const responseModel = getAgentResponseModel(req, endpointOption); - const preliminaryUserMessage = getPreliminaryUserMessage(req.body, conversationId); - const preliminaryResponseMessageId = getPreliminaryResponseMessageId(req.body); - await GenerationJobManager.updateMetadata(streamId, { - conversationId, - endpoint: endpointOption.endpoint, - iconURL: endpointIconURL, - model: responseModel, - // Persist the originating agent so a HITL resume can refuse to rebuild this - // paused run on a different agent (see resume.js). - agent_id: endpointOption.agent_id ?? req.body?.agent_id, - // Persist temporary-chat state so a HITL resume keeps the resumed response - // non-persisted instead of trusting the resume request to re-send the flag. - isTemporary: req.body?.isTemporary, - responseMessageId: preliminaryResponseMessageId, - userMessage: preliminaryUserMessage, - }); + await attachConversationCreatedAt(req, conversationId, conversationAnchorPromise).then(() => + startupTelemetry?.mark('conversation_resolved'), + ); // Note: We no longer use res.on('close') to abort since we send JSON immediately. // The response closes normally after res.json(), which is not an abort condition. @@ -463,15 +465,34 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit endpointOption, // Use the job's abort controller signal - allows abort via GenerationJobManager.abortJob() signal: job.abortController.signal, + jobCreatedAt, }); + startupTelemetry?.mark('client_initialized'); + client = result.client; if (job.abortController.signal.aborted) { - GenerationJobManager.completeJob(streamId, 'Request aborted during initialization'); - await finishResumableRequest(req, userId); + await GenerationJobManager.completeJob( + streamId, + 'Request aborted during initialization', + jobCreatedAt, + ).catch((completeErr) => { + logger.warn( + '[ResumableAgentController] completeJob failed after initialization abort', + completeErr, + ); + }); + startupTelemetry?.end('aborted'); + try { + await finishResumableRequest(req, userId); + } finally { + if (client) { + disposeClient(client); + } + client = null; + } return; } - client = result.client; // Tag the client with THIS generation's identity so HITL terminal side-effects // (pause CAS, checkpoint prune) can tell whether a newer request has since replaced // this job on the same conversationId before acting on it. @@ -485,12 +506,18 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit }); if (client?.sender) { - GenerationJobManager.updateMetadata(streamId, { sender: client.sender }); + void GenerationJobManager.updateMetadata( + streamId, + { sender: client.sender }, + jobCreatedAt, + ).catch((err) => { + logger.warn('[ResumableAgentController] Failed to persist response sender', err); + }); } // Store reference to client's contentParts - graph will be set when run is created if (client?.contentParts) { - GenerationJobManager.setContentParts(streamId, client.contentParts); + GenerationJobManager.setContentParts(streamId, client.contentParts, jobCreatedAt); } let userMessage; @@ -502,24 +529,33 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit // conversationId is pre-generated, no need to update from callback }; - // Start background generation - readyPromise resolves immediately now - // (sync mechanism handles late subscribers) - const startGeneration = async () => { - try { - // Short timeout as safety net - promise should already be resolved - await Promise.race([job.readyPromise, new Promise((resolve) => setTimeout(resolve, 100))]); - } catch (waitError) { - logger.warn( - `[ResumableAgentController] Error waiting for subscriber: ${waitError.message}`, - ); + let immediateTitlePromise = null; + let backgroundClientCleanupScheduled = false; + const disposeBackgroundClient = () => { + if (backgroundClientCleanupScheduled) { + return; } + backgroundClientCleanupScheduled = true; + if (immediateTitlePromise) { + immediateTitlePromise.finally(() => { + if (client) { + disposeClient(client); + } + }); + } else if (client) { + disposeClient(client); + } + }; + + // Start background generation immediately. The stream layer buffers and persists events + // until an SSE subscriber attaches, so generation no longer waits on subscriber readiness. + const startGeneration = async () => { /** Immediate-mode title generation runs in parallel with the response, so * the conversation row may not exist when the title resolves. `convoReady` * resolves once the response (and thus the conversation) has been saved, * gating the title's `saveConvo`. Declared here so both the success tail * and the catch block can settle it and gate `disposeClient` on the title. */ - let immediateTitlePromise = null; let titleEventPromise = null; let acceptsTitleEvents = true; let resolveConvoReady; @@ -574,31 +610,37 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit userMessage = userMsg; // Store userMessage and responseMessageId upfront for resume capability - GenerationJobManager.updateMetadata(streamId, { - responseMessageId: respMsgId, - userMessage: { - messageId: userMsg.messageId, - parentMessageId: userMsg.parentMessageId, - conversationId: userMsg.conversationId, - text: userMsg.text, - quotes: userMsg.quotes, - // Persist the turn's uploaded files here (authoritative job metadata) so a - // HITL resume sources them from the job, not the user DB row — which the - // approval prompt can race (the row save may still be in flight when a fast - // /resume reads it). Without this an approved tool run can rebuild without the - // paused turn's files. - ...(Array.isArray(req.body?.files) && - req.body.files.length > 0 && { files: req.body.files }), - // Skill selections aren't on `userMsg` yet at onStart (BaseClient adds them - // later), so source them from the request — otherwise this update overwrites - // the preliminary metadata and a HITL-resumed turn loses its skill pills. - ...(Array.isArray(req.body?.manualSkills) && - req.body.manualSkills.length > 0 && { manualSkills: req.body.manualSkills }), - ...(Array.isArray(req.body?.alwaysAppliedSkills) && - req.body.alwaysAppliedSkills.length > 0 && { - alwaysAppliedSkills: req.body.alwaysAppliedSkills, - }), + GenerationJobManager.updateMetadata( + streamId, + { + responseMessageId: respMsgId, + userMessage: { + messageId: userMsg.messageId, + parentMessageId: userMsg.parentMessageId, + conversationId: userMsg.conversationId, + text: userMsg.text, + quotes: userMsg.quotes, + // Persist the turn's uploaded files here (authoritative job metadata) so a + // HITL resume sources them from the job, not the user DB row — which the + // approval prompt can race (the row save may still be in flight when a fast + // /resume reads it). Without this an approved tool run can rebuild without the + // paused turn's files. + ...(Array.isArray(req.body?.files) && + req.body.files.length > 0 && { files: req.body.files }), + // Skill selections aren't on `userMsg` yet at onStart (BaseClient adds them + // later), so source them from the request — otherwise this update overwrites + // the preliminary metadata and a HITL-resumed turn loses its skill pills. + ...(Array.isArray(req.body?.manualSkills) && + req.body.manualSkills.length > 0 && { manualSkills: req.body.manualSkills }), + ...(Array.isArray(req.body?.alwaysAppliedSkills) && + req.body.alwaysAppliedSkills.length > 0 && { + alwaysAppliedSkills: req.body.alwaysAppliedSkills, + }), + }, }, + jobCreatedAt, + ).catch((err) => { + logger.error('[ResumableAgentController] Failed to persist start metadata', err); }); GenerationJobManager.emitChunk(streamId, { @@ -621,6 +663,8 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit }), }, streamId, + }).catch((err) => { + logger.error('[ResumableAgentController] Failed to queue created event', err); }); }; @@ -757,6 +801,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit logger.debug( `[ResumableAgentController] Turn paused for approval; awaiting resume: ${streamId}`, ); + startupTelemetry?.end('paused'); return; } @@ -834,6 +879,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit resolveConvoReady(); // Still decrement pending request since we incremented at start await finishResumableRequest(req, userId); + startupTelemetry?.end('replaced'); if (immediateTitlePromise) { immediateTitlePromise.finally(() => { if (client) { @@ -881,10 +927,15 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit // Parked BEFORE the final event: a client with no live subscriber // recovers these via /chat/status (claim-on-read) within the // recovery TTL — the SSE copy alone is transient. - await GenerationJobManager.steering.park(streamId, pendingSteers, { - userId, - tenantId: req.user?.tenantId, - }); + await GenerationJobManager.steering.park( + streamId, + pendingSteers, + { + userId, + tenantId: req.user?.tenantId, + }, + jobCreatedAt, + ); } } catch (err) { logger.warn(`[ResumableAgentController] Failed to drain leftover steers`, err); @@ -908,8 +959,11 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit conversationId: conversation?.conversationId, }); - await GenerationJobManager.emitDone(streamId, finalEvent); - GenerationJobManager.completeJob(streamId); + await GenerationJobManager.emitDone(streamId, finalEvent, jobCreatedAt); + startupTelemetry?.end('completed_without_delta'); + void GenerationJobManager.completeJob(streamId, undefined, jobCreatedAt).catch((err) => { + logger.warn('[ResumableAgentController] Failed to finalize completed job', err); + }); await finishResumableRequest(req, userId); } else { const finalEvent = { @@ -929,8 +983,13 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit conversationId: conversation?.conversationId, }); - await GenerationJobManager.emitDone(streamId, finalEvent); - GenerationJobManager.completeJob(streamId, 'Request aborted'); + await GenerationJobManager.emitDone(streamId, finalEvent, jobCreatedAt); + startupTelemetry?.end('aborted'); + void GenerationJobManager.completeJob(streamId, 'Request aborted', jobCreatedAt).catch( + (err) => { + logger.warn('[ResumableAgentController] Failed to finalize aborted job', err); + }, + ); await finishResumableRequest(req, userId); } @@ -982,6 +1041,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit if (wasAborted) { logger.debug(`[ResumableAgentController] Generation aborted for ${streamId}`); + startupTelemetry?.end('aborted'); // abortJob already handled emitDone and completeJob } else { logger.error(`[ResumableAgentController] Generation error for ${streamId}:`, error); @@ -1002,6 +1062,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit streamId, erroredLeftovers.map(toPendingSteer), { userId, tenantId: req.user?.tenantId }, + jobCreatedAt, ); } } catch (drainErr) { @@ -1010,21 +1071,34 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit drainErr, ); } - await GenerationJobManager.emitError(streamId, error.message || 'Generation failed'); - GenerationJobManager.completeJob(streamId, error.message); + try { + await GenerationJobManager.emitError( + streamId, + error.message || 'Generation failed', + jobCreatedAt, + ); + } catch (notificationError) { + logger.warn( + '[ResumableAgentController] Failed to notify client of generation error', + notificationError, + ); + } finally { + startupTelemetry?.end('error', error); + } + await GenerationJobManager.completeJob(streamId, error.message, jobCreatedAt).catch( + (completeErr) => { + logger.warn( + '[ResumableAgentController] completeJob failed during generation-error cleanup', + completeErr, + ); + }, + ); } - await finishResumableRequest(req, userId); - - // Defer disposal until any immediate title settles (it holds the run/req). - if (immediateTitlePromise) { - immediateTitlePromise.finally(() => { - if (client) { - disposeClient(client); - } - }); - } else if (client) { - disposeClient(client); + try { + await finishResumableRequest(req, userId); + } finally { + disposeBackgroundClient(); } // Don't continue to title generation after error/abort @@ -1037,30 +1111,59 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit logger.error( `[ResumableAgentController] Unhandled error in background generation: ${err.message}`, ); - GenerationJobManager.completeJob(streamId, err.message); - await finishResumableRequest(req, userId); + startupTelemetry?.end('error', err); + await GenerationJobManager.completeJob(streamId, err.message, jobCreatedAt).catch( + (completeErr) => { + logger.warn( + '[ResumableAgentController] completeJob failed during background-error cleanup', + completeErr, + ); + }, + ); + try { + await finishResumableRequest(req, userId); + } finally { + disposeBackgroundClient(); + } }); } catch (error) { logger.error('[ResumableAgentController] Initialization error:', error); - if (!res.headersSent) { - res.status(500).json({ error: error.message || 'Failed to start generation' }); - } else { - // JSON already sent, emit error to stream so client can receive it - await GenerationJobManager.emitError(streamId, error.message || 'Failed to start generation'); + try { + if (!res.headersSent) { + res.status(500).json({ error: error.message || 'Failed to start generation' }); + } else if (jobCreatedAt != null) { + // JSON already sent, emit error to stream so client can receive it + await GenerationJobManager.emitError( + streamId, + error.message || 'Failed to start generation', + jobCreatedAt, + ); + } + } catch (notificationError) { + logger.warn( + '[ResumableAgentController] Failed to notify client of initialization error', + notificationError, + ); + } finally { + startupTelemetry?.end('error', error); } // Finalize THIS failed job before releasing the idempotency claim. Releasing first would // let the client's retry win the same key and createJob() the same streamId while we are - // still here — and completeJob() is not guarded by the original createdAt, so it would - // abort/error that replacement. A completeJob() rejection (store hiccup) must NOT skip the + // still here. The generation guard is defense-in-depth around that ordering. A + // completeJob() rejection (store hiccup) must NOT skip the // release + pending-request decrement below, or the retry stays wedged behind the claim // and the concurrency slot leaks — so swallow its error. (A failed completeJob did not // finalize anything, so releasing afterward can't let it abort a later replacement.) - await GenerationJobManager.completeJob(streamId, error.message).catch((completeErr) => { - logger.warn( - '[ResumableAgentController] completeJob failed during init-error cleanup', - completeErr, + if (jobCreatedAt != null) { + await GenerationJobManager.completeJob(streamId, error.message, jobCreatedAt).catch( + (completeErr) => { + logger.warn( + '[ResumableAgentController] completeJob failed during init-error cleanup', + completeErr, + ); + }, ); - }); + } if (ownsIdempotencyClaim) { await GenerationJobManager.releaseGeneration(userId, clientRequestId).catch(() => {}); } @@ -1107,6 +1210,7 @@ const _LegacyAgentController = async (req, res, next, initializeClient, addTitle let userMessageId; let responseMessageId; let client = null; + let jobCreatedAt; let cleanupHandlers = []; // Match the same logic used for conversationId generation above @@ -1135,9 +1239,9 @@ const _LegacyAgentController = async (req, res, next, initializeClient, addTitle responseMessageId = data[key]; } else if (key === 'promptTokens') { // Update job metadata with prompt tokens for abort handling - GenerationJobManager.updateMetadata(streamId, { promptTokens: data[key] }); + GenerationJobManager.updateMetadata(streamId, { promptTokens: data[key] }, jobCreatedAt); } else if (key === 'sender') { - GenerationJobManager.updateMetadata(streamId, { sender: data[key] }); + GenerationJobManager.updateMetadata(streamId, { sender: data[key] }, jobCreatedAt); } // conversationId is pre-generated, no need to update from callback } @@ -1159,9 +1263,9 @@ const _LegacyAgentController = async (req, res, next, initializeClient, addTitle } // Complete the job in GenerationJobManager - if (streamId) { + if (jobCreatedAt != null) { logger.debug('[AgentController] Completing job in GenerationJobManager'); - await GenerationJobManager.completeJob(streamId); + await GenerationJobManager.completeJob(streamId, undefined, jobCreatedAt); } // Dispose client properly @@ -1225,18 +1329,24 @@ const _LegacyAgentController = async (req, res, next, initializeClient, addTitle // Create job in GenerationJobManager for abort handling // streamId === conversationId (pre-generated above) const job = await GenerationJobManager.createJob(streamId, userId, conversationId); + jobCreatedAt = job.createdAt; + client.jobCreatedAt = jobCreatedAt; // Store endpoint metadata for abort handling - GenerationJobManager.updateMetadata(streamId, { - endpoint: endpointOption.endpoint, - iconURL: getEndpointIconURL(req, endpointOption), - model: getAgentResponseModel(req, endpointOption), - sender: client?.sender, - }); + GenerationJobManager.updateMetadata( + streamId, + { + endpoint: endpointOption.endpoint, + iconURL: getEndpointIconURL(req, endpointOption), + model: getAgentResponseModel(req, endpointOption), + sender: client?.sender, + }, + jobCreatedAt, + ); // Store content parts reference for abort if (client?.contentParts) { - GenerationJobManager.setContentParts(streamId, client.contentParts); + GenerationJobManager.setContentParts(streamId, client.contentParts, jobCreatedAt); } const closeHandler = createCloseHandler(job.abortController); @@ -1259,16 +1369,20 @@ const _LegacyAgentController = async (req, res, next, initializeClient, addTitle responseMessageId = respMsgId; // Store metadata for abort handling (conversationId is pre-generated) - GenerationJobManager.updateMetadata(streamId, { - responseMessageId: respMsgId, - userMessage: { - messageId: userMsg.messageId, - parentMessageId: userMsg.parentMessageId, - conversationId, - text: userMsg.text, - quotes: userMsg.quotes, + GenerationJobManager.updateMetadata( + streamId, + { + responseMessageId: respMsgId, + userMessage: { + messageId: userMsg.messageId, + parentMessageId: userMsg.parentMessageId, + conversationId, + text: userMsg.text, + quotes: userMsg.quotes, + }, }, - }); + jobCreatedAt, + ); }; const messageOptions = { diff --git a/api/server/controllers/agents/resume.js b/api/server/controllers/agents/resume.js index 7f0fc2d106..124a43ad55 100644 --- a/api/server/controllers/agents/resume.js +++ b/api/server/controllers/agents/resume.js @@ -10,6 +10,7 @@ const { findDisallowedDecisions, findIncompleteDecisions, computeAgentRequestFingerprint, + captureAgentCheckpointGeneration, deleteAgentCheckpoint, buildAbortedResponseMetadata, sanitizeMessageForTransmit, @@ -194,7 +195,15 @@ function resolveResumeValue(pendingAction, body) { * job, and prune the checkpoint. Mirrors the abort route's save shape but for a * successful finish. Best-effort title generation for a first-turn pause. */ -async function finalizeResumedTurn({ req, client, job, streamId, conversationId, addTitle }) { +async function finalizeResumedTurn({ + req, + client, + job, + streamId, + conversationId, + addTitle, + checkpointGeneration, +}) { const userId = req.user.id; const checkpointerCfg = req.config?.endpoints?.[EModelEndpoint.agents]?.checkpointer; const meta = job.metadata ?? {}; @@ -361,10 +370,15 @@ async function finalizeResumedTurn({ req, client, job, streamId, conversationId, // via /chat/status within the recovery TTL). NOTE: `job` is the manager // facade — owner fields live under `metadata` (a bare `job.userId` is // undefined and would make the parked payload unclaimable). - await GenerationJobManager.steering.park(streamId, pendingSteers, { - userId: job.metadata?.userId, - tenantId: job.metadata?.tenantId, - }); + await GenerationJobManager.steering.park( + streamId, + pendingSteers, + { + userId: job.metadata?.userId, + tenantId: job.metadata?.tenantId, + }, + job.createdAt, + ); } } catch (drainErr) { logger.warn('[ResumeAgentController] Failed to drain leftover steers', drainErr); @@ -391,15 +405,15 @@ async function finalizeResumedTurn({ req, client, job, streamId, conversationId, ...(pendingSteers && { pendingSteers }), }; - await GenerationJobManager.emitDone(streamId, finalEvent); + await GenerationJobManager.emitDone(streamId, finalEvent, job.createdAt); // Awaited (not fire-and-forget) so the job's terminal write lands before the // checkpoint prune, and so a failure here doesn't race the controller's error path. try { - await GenerationJobManager.completeJob(streamId); + await GenerationJobManager.completeJob(streamId, undefined, job.createdAt); } catch (completeErr) { logger.error('[ResumeAgentController] Failed to complete resumed turn', completeErr); } - await deleteAgentCheckpoint(conversationId, checkpointerCfg); + await deleteAgentCheckpoint(conversationId, checkpointerCfg, checkpointGeneration); } /** @@ -510,6 +524,23 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) }); } + // Snapshot the exact durable checkpoint ids before the atomic resume claim. The + // claim is the linearization point: a replacement that already owns this stream + // makes it fail, while one that starts afterward writes fresh ids outside the + // snapshot. Terminal cleanup can therefore delete this generation without a + // check-then-delete race against a later pause on the same conversation. + // + // Start the indexed read alongside the independent concurrency check so the + // generation guard adds minimal time to the resume ACK path. + const checkpointerCfg = req.config?.endpoints?.[EModelEndpoint.agents]?.checkpointer; + const checkpointGenerationPromise = captureAgentCheckpointGeneration( + conversationId, + checkpointerCfg, + ).catch((err) => { + logger.warn('[ResumeAgentController] Failed to capture checkpoint generation', err); + return { threadId: conversationId, checkpointIds: [] }; + }); + // Count the resume against the concurrency limit. The original turn released its slot // when it paused, so resuming must re-acquire one — otherwise pausing several turns // and resuming them at once would bypass LIMIT_CONCURRENT_MESSAGES. @@ -527,7 +558,9 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) // would leak the concurrency slot until the counter TTL expires — spuriously 429'ing // the user when they retry the still-paused approval. Release the slot on that path too. let claimed; + let checkpointGeneration; try { + checkpointGeneration = await checkpointGenerationPromise; claimed = await GenerationJobManager.approvals.resolve(streamId, pendingAction.actionId); } catch (err) { await decrementPendingRequest(userId); @@ -626,6 +659,7 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) res, endpointOption: req.body.endpointOption, signal: job.abortController.signal, + jobCreatedAt: job.createdAt, }); client = result.client; @@ -655,7 +689,7 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) ); } if (client.contentParts) { - GenerationJobManager.setContentParts(streamId, client.contentParts); + GenerationJobManager.setContentParts(streamId, client.contentParts, job.createdAt); } await client.resumeCompletion({ @@ -690,7 +724,15 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) return; } - await finalizeResumedTurn({ req, client, job, streamId, conversationId, addTitle }); + await finalizeResumedTurn({ + req, + client, + job, + streamId, + conversationId, + addTitle, + checkpointGeneration, + }); } catch (err) { logger.error('[ResumeAgentController] Resume failed', err); // Job-replacement guard (mirrors finalizeResumedTurn's success-path guard): if a @@ -720,30 +762,47 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) ); if (leftoverSteers.length > 0) { // Facade shape: owner fields are under `metadata` (see finalize). - await GenerationJobManager.steering.park(streamId, leftoverSteers.map(toPendingSteer), { - userId: job.metadata?.userId, - tenantId: job.metadata?.tenantId, - }); + await GenerationJobManager.steering.park( + streamId, + leftoverSteers.map(toPendingSteer), + { + userId: job.metadata?.userId, + tenantId: job.metadata?.tenantId, + }, + job.createdAt, + ); } } catch (drainErr) { logger.warn('[ResumeAgentController] Failed to drain steers on resume failure', drainErr); } try { - await GenerationJobManager.emitError(streamId, err?.message ?? 'Resume failed'); + await GenerationJobManager.emitError( + streamId, + err?.message ?? 'Resume failed', + job.createdAt, + ); } catch (emitErr) { logger.error('[ResumeAgentController] Failed to emit resume error', emitErr); } try { - await GenerationJobManager.completeJob(streamId, err?.message ?? 'Resume failed'); + await GenerationJobManager.completeJob( + streamId, + err?.message ?? 'Resume failed', + job.createdAt, + ); } catch (completeErr) { logger.error('[ResumeAgentController] Failed to finalize failed resume', completeErr); // Last resort: force a terminal state so the job isn't orphaned in `running`. await GenerationJobManager.getJobStore() - .updateJob(streamId, { - status: 'error', - completedAt: Date.now(), - error: 'Resume failed', - }) + .updateJob( + streamId, + { + status: 'error', + completedAt: Date.now(), + error: 'Resume failed', + }, + job.createdAt, + ) .catch((updErr) => logger.error('[ResumeAgentController] Fallback job finalize failed', updErr), ); @@ -751,6 +810,7 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) await deleteAgentCheckpoint( conversationId, req.config?.endpoints?.[EModelEndpoint.agents]?.checkpointer, + checkpointGeneration, ); } } finally { diff --git a/api/server/experimental.js b/api/server/experimental.js index e670dd07e0..dadb87dd0b 100644 --- a/api/server/experimental.js +++ b/api/server/experimental.js @@ -10,17 +10,15 @@ const express = require('express'); const passport = require('passport'); const compression = require('compression'); const cookieParser = require('cookie-parser'); -const { logger, runAsSystem, tenantStorage } = require('@librechat/data-schemas'); +const { logger, runAsSystem } = require('@librechat/data-schemas'); const mongoSanitize = require('express-mongo-sanitize'); const { isEnabled, apiNotFound, ErrorController, - GenerationJobManager, QUERY_DEVTOOLS_HEADER, performStartupChecks, handleJsonParseError, - deleteAgentCheckpoint, initializeFileStorage, loadToolApprovalHooks, maybeInjectQueryDevtoolsBootstrap, @@ -313,22 +311,6 @@ if (cluster.isMaster) { await loadToolApprovalHooks(toolApproval?.enabled ? toolApproval.hooks : undefined, { basePath: path.resolve(__dirname, '../..'), }); - // Prune the paused run's durable checkpoint when its approval EXPIRES (a stale submit — - // this startup never runs the periodic sweeper) instead of leaving it until the Mongo - // TTL. Mirrors api/server/index.js's configureGenerationStreams wiring; safe here even - // though this startup runs the manager on constructor defaults (the setter never resets - // services). streamId === conversationId === the LangGraph thread_id. - GenerationJobManager.setApprovalExpiredHandler(async (conversationId, job) => { - // Resolve config in the PAUSED JOB's tenant/user scope (mirrors index.js): enter the - // tenant ALS context — getAppConfig args alone only key the cache. - await tenantStorage.run({ tenantId: job?.tenantId, userId: job?.userId }, async () => { - const currentConfig = await getAppConfig({ - userId: job?.userId, - tenantId: job?.tenantId, - }); - await deleteAgentCheckpoint(conversationId, currentConfig?.endpoints?.agents?.checkpointer); - }); - }); expiredFileSweepOptions = { appConfig, loadAppConfig: getAppConfig }; startExpiredFileSweepOnce(); await performStartupChecks(appConfig); diff --git a/api/server/index.js b/api/server/index.js index f88b2fd258..be3b8a1fe0 100644 --- a/api/server/index.js +++ b/api/server/index.js @@ -9,7 +9,7 @@ const passport = require('passport'); const compression = require('compression'); const cookieParser = require('cookie-parser'); const mongoSanitize = require('express-mongo-sanitize'); -const { logger, runAsSystem, tenantStorage } = require('@librechat/data-schemas'); +const { logger, runAsSystem } = require('@librechat/data-schemas'); const { isEnabled, apiNotFound, @@ -21,12 +21,14 @@ const { GenerationJobManager, QUERY_DEVTOOLS_HEADER, createStreamServices, - deleteAgentCheckpoint, + agentStartupIngressMiddleware, + agentStartupTelemetryMiddleware, initializeFileStorage, initializeDeploymentSkills, loadToolApprovalHooks, maybeInjectQueryDevtoolsBootstrap, preAuthTenantMiddleware, + registerShutdownTask, setupGracefulShutdown, updateInterfacePermissions, } = require('@librechat/api'); @@ -85,18 +87,18 @@ const configureGenerationStreams = () => { cleanupOnComplete: !isEnabled(process.env.STREAM_KEEP_COMPLETED_JOBS), }); GenerationJobManager.initialize(); - // Prune the paused run's durable checkpoint when its approval EXPIRES (periodic sweeper - // or a stale submit) instead of leaving it until the Mongo TTL. streamId === conversationId - // === the LangGraph thread_id. Config is resolved lazily per expiry so the prune always - // targets the currently configured checkpoint collections. - GenerationJobManager.setApprovalExpiredHandler(async (conversationId, job) => { - // Resolve config in the PAUSED JOB's tenant/user scope — the expiry runs outside any - // request context. Passing ids to getAppConfig only keys the cache; the Config query - // itself is ALS-scoped by the tenant-isolation plugin, so ENTER the tenant context. - await tenantStorage.run({ tenantId: job?.tenantId, userId: job?.userId }, async () => { - const appConfig = await getAppConfig({ userId: job?.userId, tenantId: job?.tenantId }); - await deleteAgentCheckpoint(conversationId, appConfig?.endpoints?.agents?.checkpointer); - }); + // Stop active generations and close their SSE streams while the HTTP server drains. + registerShutdownTask( + 'generation job manager prepare', + () => GenerationJobManager.prepareForShutdown(), + { + phase: 'pre-drain', + priority: 100, + }, + ); + // Tear down stream resources before shared caches and telemetry exporters shut down. + registerShutdownTask('generation job manager', () => GenerationJobManager.destroy(), { + priority: 100, }); }; @@ -197,6 +199,7 @@ const startServer = async () => { }); /* Middleware */ + app.use('/api/agents/chat', agentStartupIngressMiddleware); app.use(metricsMiddleware); app.use(noIndex); app.use(express.json({ limit: '3mb' })); @@ -234,6 +237,7 @@ const startServer = async () => { if (telemetry.enabled) { app.use(telemetry.telemetryMiddleware); } + app.use('/api/agents/chat', agentStartupTelemetryMiddleware); if (!ALLOW_SOCIAL_LOGIN) { console.warn('Social logins are disabled. Set ALLOW_SOCIAL_LOGIN=true to enable them.'); diff --git a/api/server/index.spec.js b/api/server/index.spec.js index 3e0fc07127..e73ee40f60 100644 --- a/api/server/index.spec.js +++ b/api/server/index.spec.js @@ -81,6 +81,24 @@ describe('Telemetry wiring', () => { expect(errorControllerIndex).toBeGreaterThan(-1); expect(telemetryErrorMiddlewareIndex).toBeLessThan(errorControllerIndex); }); + + it('captures agent ingress before parsing and creates its recorder before auth routes', () => { + const ingressIndex = source.indexOf( + "app.use('/api/agents/chat', agentStartupIngressMiddleware);", + ); + const jsonParserIndex = source.indexOf("app.use(express.json({ limit: '3mb' }));"); + const recorderIndex = source.indexOf( + "app.use('/api/agents/chat', agentStartupTelemetryMiddleware);", + ); + const tracingIndex = source.indexOf('app.use(telemetry.telemetryMiddleware);'); + const agentsRouteIndex = source.indexOf("app.use('/api/agents', routes.agents);"); + + expect(ingressIndex).toBeGreaterThan(-1); + expect(recorderIndex).toBeGreaterThan(-1); + expect(ingressIndex).toBeLessThan(jsonParserIndex); + expect(tracingIndex).toBeLessThan(recorderIndex); + expect(recorderIndex).toBeLessThan(agentsRouteIndex); + }); }); describe('Startup readiness wiring', () => { @@ -98,6 +116,16 @@ describe('Startup readiness wiring', () => { expect(streamConfigIndex).toBeLessThan(postListenMcpIndex); }); + it('registers generation stream cleanup with the graceful shutdown coordinator', () => { + const shutdownRegistrationIndex = source.indexOf( + "registerShutdownTask('generation job manager'", + ); + const listenIndex = source.indexOf('const server = app.listen'); + + expect(shutdownRegistrationIndex).toBeGreaterThan(-1); + expect(shutdownRegistrationIndex).toBeLessThan(listenIndex); + }); + it('mounts the chat-start readiness gate before agent routes', () => { const readinessGateIndex = source.indexOf( "app.use('/api/agents/chat', rejectChatStartsUntilReady);", diff --git a/api/server/routes/agents/__tests__/streamTenant.spec.js b/api/server/routes/agents/__tests__/streamTenant.spec.js index 55d96d8cc3..ee55f5e9c8 100644 --- a/api/server/routes/agents/__tests__/streamTenant.spec.js +++ b/api/server/routes/agents/__tests__/streamTenant.spec.js @@ -4,7 +4,9 @@ const request = require('supertest'); const mockGenerationJobManager = { getJob: jest.fn(), subscribe: jest.fn(), + subscribeWithResume: jest.fn(), getResumeState: jest.fn(), + markSyncSent: jest.fn(), abortJob: jest.fn(), getActiveJobIdsForUser: jest.fn().mockResolvedValue([]), }; @@ -125,6 +127,73 @@ describe('SSE stream tenant isolation', () => { expect(mockGenerationJobManager.subscribe).toHaveBeenCalledTimes(1); }); + it('writes the resume sync frame before activating live delivery', async () => { + mockGenerationJobManager.getJob.mockResolvedValue({ + metadata: { userId: 'user-123' }, + status: 'running', + }); + const activate = jest.fn(); + mockGenerationJobManager.subscribeWithResume.mockImplementation( + async (_streamId, writeEvent, onDone) => { + activate.mockImplementation(() => { + writeEvent({ event: 'on_message_delta', data: { text: 'live' } }); + onDone({ final: true }); + }); + return { + subscription: { unsubscribe: jest.fn(), activate }, + resumeState: { runSteps: [], aggregatedContent: [] }, + pendingEvents: [], + }; + }, + ); + + const res = await request(app).get('/agents/chat/stream/stream-123?resume=true'); + const payloads = res.text + .trim() + .split('\n\n') + .map((frame) => JSON.parse(frame.split('\ndata: ')[1])); + + expect(res.status).toBe(200); + expect(payloads).toEqual([ + { + sync: true, + resumeState: { runSteps: [], aggregatedContent: [] }, + pendingEvents: [], + }, + { event: 'on_message_delta', data: { text: 'live' } }, + { final: true }, + ]); + expect(activate).toHaveBeenCalledTimes(1); + expect(mockGenerationJobManager.markSyncSent).toHaveBeenCalledWith('stream-123'); + }); + + it('detaches a paused resume subscription when the response ends before activation', async () => { + mockGenerationJobManager.getJob.mockResolvedValue({ + metadata: { userId: 'user-123' }, + status: 'running', + }); + const subscription = { + unsubscribe: jest.fn(), + activate: jest.fn(), + }; + mockGenerationJobManager.subscribeWithResume.mockImplementation( + async (_streamId, _writeEvent, onDone) => { + onDone({ final: true }); + return { + subscription, + resumeState: { runSteps: [], aggregatedContent: [] }, + pendingEvents: [], + }; + }, + ); + + const res = await request(app).get('/agents/chat/stream/stream-123?resume=true'); + + expect(res.status).toBe(200); + expect(subscription.unsubscribe).toHaveBeenCalled(); + expect(subscription.activate).not.toHaveBeenCalled(); + }); + it('returns 403 when job has tenantId but user has no tenantId', async () => { mockUserId = 'user-123'; mockTenantId = undefined; diff --git a/api/server/routes/agents/index.js b/api/server/routes/agents/index.js index 759d9c0177..217c94203b 100644 --- a/api/server/routes/agents/index.js +++ b/api/server/routes/agents/index.js @@ -71,8 +71,18 @@ router.use(uaParser); router.get('/chat/stream/:streamId', async (req, res) => { const { streamId } = req.params; const isResume = req.query.resume === 'true'; + let result; + const attachmentAbortController = new AbortController(); + req.on('close', () => { + logger.debug(`[AgentStream] Client disconnected from ${streamId}`); + attachmentAbortController.abort(); + result?.unsubscribe(); + }); const job = await GenerationJobManager.getJob(streamId); + if (attachmentAbortController.signal.aborted) { + return; + } if (!job) { return res.status(404).json({ error: 'Stream not found', @@ -129,13 +139,13 @@ router.get('/chat/stream/:streamId', async (req, res) => { } }; - let result; - if (isResume) { const { subscription, resumeState, pendingEvents } = - await GenerationJobManager.subscribeWithResume(streamId, writeEvent, onDone, onError); + await GenerationJobManager.subscribeWithResume(streamId, writeEvent, onDone, onError, { + signal: attachmentAbortController.signal, + }); - if (!res.writableEnded) { + if (subscription && !attachmentAbortController.signal.aborted && !res.writableEnded) { if (resumeState) { writeEvent({ sync: true, resumeState, pendingEvents }); GenerationJobManager.markSyncSent(streamId); @@ -150,23 +160,27 @@ router.get('/chat/stream/:streamId', async (req, res) => { `[AgentStream] Resume state null for ${streamId}, replayed ${pendingEvents.length} gap events directly`, ); } + subscription.activate(); + } else { + subscription?.unsubscribe(); } result = subscription; } else { - result = await GenerationJobManager.subscribe(streamId, writeEvent, onDone, onError); + result = await GenerationJobManager.subscribe(streamId, writeEvent, onDone, onError, { + signal: attachmentAbortController.signal, + }); } + if (attachmentAbortController.signal.aborted) { + result?.unsubscribe(); + return; + } if (!result) { streamTelemetry.recordSubscribeFailed(); onError('Failed to subscribe to stream'); return; } - - req.on('close', () => { - logger.debug(`[AgentStream] Client disconnected from ${streamId}`); - result.unsubscribe(); - }); }); /** diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index f655403f9a..50ab851aee 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -14,6 +14,7 @@ const { resolveAgentTokenConfig, resolveAgentScopedSkillIds, resolveModelSpecSkillIds, + getAgentStartupTelemetry, buildAgentContextAttachmentsByAgentId, } = require('@librechat/api'); const { @@ -111,12 +112,14 @@ function createToolLoader(signal, streamId = null, definitionsOnly = false) { * @param {Express.Response} params.res * @param {AbortSignal} params.signal * @param {Object} params.endpointOption + * @param {number} [params.jobCreatedAt] */ -const initializeClient = async ({ req, res, signal, endpointOption }) => { +const initializeClient = async ({ req, res, signal, endpointOption, jobCreatedAt }) => { if (!endpointOption) { throw new Error('Endpoint option not provided'); } const appConfig = req.config; + const startupTelemetry = getAgentStartupTelemetry(req); /** @type {string | null} */ const streamId = req._resumableStreamId || null; @@ -160,6 +163,10 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => { const ephemeralSkillsToggle = req.body?.ephemeralAgent?.skills === true; const skillDbMethods = getSkillDbMethods(); + if (!endpointOption.agent) { + throw new Error('No agent promise provided'); + } + /** Run-level gate for inline memory tools: the `memory` capability must be * enabled, memory must be configured, and the user must not have opted out. * Requires the memory WRITE permissions (CREATE + UPDATE) — both inline tools @@ -167,45 +174,66 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => { * read-only-memory roles that the runtime loader would then refuse to build. * Agents (or the ephemeral memory badge) opt in per-agent via the `memory` * marker on `tools`. */ - const memoryAvailable = + const memoryAvailablePromise = enabledCapabilities.has(AgentCapabilities.memory) && isMemoryEnabled(appConfig?.memory) && req.user?.personalization?.memories !== false && - (await checkAccess({ + checkAccess({ user: req.user, permissionType: PermissionTypes.MEMORIES, permissions: [Permissions.USE, Permissions.CREATE, Permissions.UPDATE], getRoleByName: db.getRoleByName, - })); + }); - const accessibleSkillIds = skillsCapabilityEnabled - ? withDeploymentSkillIds( - await findAccessibleResources({ - userId: req.user.id, - role: req.user.role, - resourceType: ResourceType.SKILL, - requiredPermissions: PermissionBits.VIEW, - }), - ) - : []; - const editableSkillIds = skillsCapabilityEnabled - ? await findAccessibleResources({ + const accessibleSkillIdsPromise = skillsCapabilityEnabled + ? findAccessibleResources({ + userId: req.user.id, + role: req.user.role, + resourceType: ResourceType.SKILL, + requiredPermissions: PermissionBits.VIEW, + }).then(withDeploymentSkillIds) + : Promise.resolve([]); + const editableSkillIdsPromise = skillsCapabilityEnabled + ? findAccessibleResources({ userId: req.user.id, role: req.user.role, resourceType: ResourceType.SKILL, requiredPermissions: PermissionBits.EDIT, }) - : []; - const skillCreateAllowed = skillsCapabilityEnabled - ? await getSkillToolDeps().canCreateSkill({ req }) - : false; + : Promise.resolve([]); + const skillCreateAllowedPromise = skillsCapabilityEnabled + ? getSkillToolDeps().canCreateSkill({ req }) + : Promise.resolve(false); + const skillStatesPromise = accessibleSkillIdsPromise.then((accessibleSkillIds) => + loadSkillStates({ + userId: req.user.id, + appConfig, + getUserById: db.getUserById, + accessibleSkillIds, + }), + ); + const primaryAgentPromise = endpointOption.agent; + const modelsConfigPromise = getModelsConfig(req); + const validatedPrimaryAgentPromise = Promise.all([primaryAgentPromise, modelsConfigPromise]).then( + async ([primaryAgent, modelsConfig]) => { + if (!primaryAgent) { + throw new Error('Agent not found'); + } - const { skillStates, defaultActiveOnShare } = await loadSkillStates({ - userId: req.user.id, - appConfig, - getUserById: db.getUserById, - accessibleSkillIds, - }); + const validationResult = await validateAgentModel({ + req, + res, + modelsConfig, + logViolation, + agent: primaryAgent, + }); + if (!validationResult.isValid) { + throw new Error(validationResult.error?.message); + } + + return { primaryAgent, modelsConfig }; + }, + ); /** * Agent context store - populated after initialization, accessed by callback via closure. @@ -315,28 +343,22 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => { usageEmitSink, }); - if (!endpointOption.agent) { - throw new Error('No agent promise provided'); - } - - const primaryAgent = await endpointOption.agent; + const [ + memoryAvailable, + accessibleSkillIds, + editableSkillIds, + skillCreateAllowed, + { skillStates, defaultActiveOnShare }, + { primaryAgent, modelsConfig }, + ] = await Promise.all([ + memoryAvailablePromise, + accessibleSkillIdsPromise, + editableSkillIdsPromise, + skillCreateAllowedPromise, + skillStatesPromise, + validatedPrimaryAgentPromise, + ]); delete endpointOption.agent; - if (!primaryAgent) { - throw new Error('Agent not found'); - } - - const modelsConfig = await getModelsConfig(req); - const validationResult = await validateAgentModel({ - req, - res, - modelsConfig, - logViolation, - agent: primaryAgent, - }); - - if (!validationResult.isValid) { - throw new Error(validationResult.error?.message); - } const agentConfigs = new Map(); const allowedProviders = new Set(appConfig?.endpoints?.[EModelEndpoint.agents]?.allowedProviders); @@ -1016,11 +1038,13 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => { * them to persist the breakdown + usage rollup on the response message. */ contextUsageSink, usageEmitSink, + startupTelemetry, toolInputValidationErrors, + jobCreatedAt, }); if (streamId) { - GenerationJobManager.setCollectedUsage(streamId, collectedUsage); + GenerationJobManager.setCollectedUsage(streamId, collectedUsage, jobCreatedAt); } return { client, userMCPAuthMap }; diff --git a/api/server/services/Endpoints/agents/initialize.spec.js b/api/server/services/Endpoints/agents/initialize.spec.js index 4063649b9f..c9f4cd4da9 100644 --- a/api/server/services/Endpoints/agents/initialize.spec.js +++ b/api/server/services/Endpoints/agents/initialize.spec.js @@ -13,6 +13,14 @@ const { MongoMemoryServer } = require('mongodb-memory-server'); const mockInitializeAgent = jest.fn(); const mockValidateAgentModel = jest.fn(); +function deferred() { + let resolve; + const promise = new Promise((resolvePromise) => { + resolve = resolvePromise; + }); + return { promise, resolve }; +} + jest.mock('@librechat/agents', () => ({ ...jest.requireActual('@librechat/agents'), createContentAggregator: jest.fn(() => ({ @@ -72,6 +80,7 @@ jest.mock('~/cache', () => ({ const { initializeClient } = require('./initialize'); const { getSkillToolDeps } = require('./skillDeps'); +const { getModelsConfig } = require('~/server/controllers/ModelController'); const { logger } = require('@librechat/data-schemas'); const { User, AclEntry } = require('~/db/models'); const { createAgent, createSkill } = require('~/models'); @@ -316,6 +325,41 @@ describe('initializeClient — processAgent ACL gate', () => { expect(initializeParams.agent.skills_enabled).toBe(true); expect(initializeParams.skillAuthoringAvailable).toBe(true); }); + + it('loads model validation and skill permissions without serial waits', async () => { + const models = deferred(); + const createPermission = deferred(); + const req = makeReq(); + req.config.endpoints.agents = { capabilities: ['skills'] }; + mockInitializeAgent.mockResolvedValue(makePrimaryConfig([])); + getModelsConfig.mockReturnValueOnce(models.promise); + const canCreateSkillSpy = jest + .spyOn(getSkillToolDeps(), 'canCreateSkill') + .mockReturnValue(createPermission.promise); + + try { + const initialization = initializeClient({ + req, + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + }); + + expect(getModelsConfig).toHaveBeenCalledWith(req); + expect(canCreateSkillSpy).toHaveBeenCalledWith({ req }); + + models.resolve({}); + await new Promise((resolve) => setImmediate(resolve)); + expect(mockValidateAgentModel).toHaveBeenCalledTimes(1); + expect(mockInitializeAgent).not.toHaveBeenCalled(); + + createPermission.resolve(false); + await initialization; + expect(mockInitializeAgent).toHaveBeenCalledTimes(1); + } finally { + canCreateSkillSpy.mockRestore(); + } + }); }); describe('initializeClient — subagent loading', () => { diff --git a/api/typedefs.js b/api/typedefs.js index 45c6efd107..6ec36b3573 100644 --- a/api/typedefs.js +++ b/api/typedefs.js @@ -1287,6 +1287,7 @@ * @property {string} [proxy] - Proxy configuration * @property {Object} [tools] - Available tools for the agent * @property {Object} [eventHandlers] - Custom event handlers + * @property {import('@librechat/api').AgentStartupTelemetry} [startupTelemetry] - Startup latency recorder * @property {Object} [addParams] - Additional parameters to add to requests * @property {string[]} [dropParams] - Parameters to remove from requests * @memberof typedefs diff --git a/e2e/benchmarks/README.md b/e2e/benchmarks/README.md new file mode 100644 index 0000000000..2578835fe9 --- /dev/null +++ b/e2e/benchmarks/README.md @@ -0,0 +1,56 @@ +# Agent startup latency benchmark + +This non-gating Playwright benchmark measures a fresh agent chat from the user's Enter keypress to: + +- the agent-chat POST response ending (`submitToAckMs`); +- the mock assistant token appearing in the message DOM, before browser paint + (`submitToFirstContentMs`); +- the interval between those events (`ackToFirstContentMs`). + +The first request is reported separately as `cold`. Warmups and measured samples each use a new +conversation, and measured conversations are deleted so history growth does not bias later samples. +Each report also captures host load and CPU utilization to make contaminated runs visible. + +Run the default in-memory, minimal-agent profile with: + +```sh +npm run e2e:benchmark:agents +``` + +Useful environment variables: + +| Variable | Default | Purpose | +| ---------------------------- | ----------- | ----------------------------------------------------------- | +| `E2E_LATENCY_PROFILE` | `minimal` | Use `mcp-memory` to exercise MCP and memory startup. | +| `E2E_LATENCY_TURN` | `first` | Use `follow-up` to measure a constant one-turn history. | +| `E2E_LATENCY_WARMUPS` | `5` | Number of unreported warmup samples after the cold request. | +| `E2E_LATENCY_SAMPLES` | `30` | Number of samples included in the summary. | +| `E2E_LATENCY_LABEL` | `unlabeled` | Identifies the revision or block in the JSON report. | +| `E2E_LATENCY_GIT_SHA` | `unknown` | Records the tested revision in the JSON report. | +| `E2E_LATENCY_STREAM_MODE` | `in-memory` | Describes the stream backend in the report. | +| `E2E_LATENCY_MONGO_DELAY_MS` | `0` | Adds a controlled delay before each Mongoose query. | +| `E2E_LATENCY_OUTPUT` | unset | Writes the complete report to this path. | + +To exercise Redis streams, point the E2E server at a disposable Redis instance: + +```sh +USE_REDIS=true \ +USE_REDIS_STREAMS=true \ +REDIS_URI=redis://127.0.0.1:16379 \ +E2E_LATENCY_STREAM_MODE=redis \ +E2E_LATENCY_PROFILE=mcp-memory \ +npm run e2e:benchmark:agents +``` + +For a base-versus-HEAD comparison, use identical dependencies and benchmark files, alternate blocks +in base/HEAD/HEAD/base order, and exclude the cold samples. Report both block medians as well as the +pooled median; do not remove outliers from an otherwise valid block. Avoid running builds, test +workers, or other CPU-heavy work at the same time. + +`E2E_LATENCY_MONGO_DELAY_MS` is useful for a separate simulated-I/O profile that reveals changes to +the request's asynchronous critical path. Always label and report that profile separately from the +zero-delay local result; it is a controlled workload, not a claim about production database latency. + +The `follow-up` turn profile creates one unmeasured seed exchange before every sample, then measures +the next request and deletes the conversation. This exercises conversation/history reads without +allowing the history to grow across samples. diff --git a/e2e/benchmarks/agent-startup.latency.spec.ts b/e2e/benchmarks/agent-startup.latency.spec.ts new file mode 100644 index 0000000000..12e2477dfe --- /dev/null +++ b/e2e/benchmarks/agent-startup.latency.spec.ts @@ -0,0 +1,369 @@ +import { mkdir, writeFile } from 'node:fs/promises'; +import { availableParallelism, cpus, loadavg } from 'node:os'; +import { dirname } from 'node:path'; +import { expect, test } from '@playwright/test'; +import type { Locator, Page, TestInfo } from '@playwright/test'; +import { cleanupAgent } from '../specs/mock/agents.helpers'; +import { NEW_CHAT_PATH, getAccessToken, messagesView, requestJson } from '../specs/mock/helpers'; + +type AgentResponse = { + id: string; + name?: string | null; + tools?: string[]; + mcpServerNames?: string[]; +}; + +type BrowserLatencyState = { + startedAt: number | null; + acknowledgedAt: number | null; + firstContentAt: number | null; +}; + +type LatencySample = { + submitToAckMs: number; + submitToFirstContentMs: number; + ackToFirstContentMs: number; +}; + +type Summary = { + p50: number; + p95: number; + mean: number; + min: number; + max: number; +}; + +type CpuSnapshot = { + idle: number; + total: number; +}; + +const BENCHMARK_REPLY = process.env.MOCK_LLM_REPLY ?? 'BENCH_TOKEN'; +const WARMUP_COUNT = parseCount('E2E_LATENCY_WARMUPS', 5); +const SAMPLE_COUNT = parseCount('E2E_LATENCY_SAMPLES', 30, 1); +const SIMULATED_MONGO_DELAY_MS = parseCount('E2E_LATENCY_MONGO_DELAY_MS', 0); +const BENCHMARK_PROFILE = process.env.E2E_LATENCY_PROFILE ?? 'minimal'; +if (!['minimal', 'mcp-memory'].includes(BENCHMARK_PROFILE)) { + throw new Error(`Unsupported E2E_LATENCY_PROFILE: ${BENCHMARK_PROFILE}`); +} +const BENCHMARK_TURN = process.env.E2E_LATENCY_TURN ?? 'first'; +if (!['first', 'follow-up'].includes(BENCHMARK_TURN)) { + throw new Error(`Unsupported E2E_LATENCY_TURN: ${BENCHMARK_TURN}`); +} +const MCP_SERVER_NAME = 'e2e-memory'; +const MCP_TOOLS = [ + 'memory', + `sys__server__sys_mcp_${MCP_SERVER_NAME}`, + `remember_fact_mcp_${MCP_SERVER_NAME}`, +]; + +function parseCount(name: string, fallback: number, minimum = 0) { + const parsed = Number.parseInt(process.env[name] ?? '', 10); + return Number.isInteger(parsed) && parsed >= minimum ? parsed : fallback; +} + +function round(value: number) { + return Math.round(value * 100) / 100; +} + +function captureCpuSnapshot(): CpuSnapshot { + return cpus().reduce( + (snapshot, cpu) => { + const total = Object.values(cpu.times).reduce((sum, value) => sum + value, 0); + snapshot.idle += cpu.times.idle; + snapshot.total += total; + return snapshot; + }, + { idle: 0, total: 0 }, + ); +} + +function calculateCpuUtilization(before: CpuSnapshot, after: CpuSnapshot) { + const idleDelta = after.idle - before.idle; + const totalDelta = after.total - before.total; + return totalDelta > 0 ? round(100 * (1 - idleDelta / totalDelta)) : 0; +} + +function percentile(sortedValues: number[], percentileValue: number) { + if (sortedValues.length === 1) { + return sortedValues[0]; + } + const position = (sortedValues.length - 1) * percentileValue; + const lowerIndex = Math.floor(position); + const upperIndex = Math.ceil(position); + const weight = position - lowerIndex; + return sortedValues[lowerIndex] * (1 - weight) + sortedValues[upperIndex] * weight; +} + +function summarize(values: number[]): Summary { + const sortedValues = [...values].sort((left, right) => left - right); + return { + p50: round(percentile(sortedValues, 0.5)), + p95: round(percentile(sortedValues, 0.95)), + mean: round(values.reduce((total, value) => total + value, 0) / values.length), + min: round(sortedValues[0]), + max: round(sortedValues.at(-1)!), + }; +} + +function summarizeSamples(samples: LatencySample[]) { + return { + submitToAckMs: summarize(samples.map((sample) => sample.submitToAckMs)), + submitToFirstContentMs: summarize(samples.map((sample) => sample.submitToFirstContentMs)), + ackToFirstContentMs: summarize(samples.map((sample) => sample.ackToFirstContentMs)), + }; +} + +function modelTrigger(page: Page) { + return page.getByRole('button', { name: 'Select a model' }).first(); +} + +async function createAgent(page: Page, name: string) { + const token = await getAccessToken(page); + return requestJson(page, { + path: '/api/agents', + token, + method: 'POST', + body: { + name, + provider: 'Mock Provider A', + model: 'mock-model-a', + model_parameters: {}, + ...(BENCHMARK_PROFILE === 'mcp-memory' ? { tools: MCP_TOOLS } : {}), + }, + }); +} + +async function selectAgent(page: Page, agentName: string) { + const trigger = modelTrigger(page); + await expect(trigger).toBeVisible(); + if ((await trigger.textContent())?.includes(agentName)) { + return; + } + await trigger.click(); + await page.getByRole('option', { name: 'My Agents' }).click(); + await page.getByRole('option', { name: agentName }).click(); + await expect(trigger).toContainText(agentName); +} + +async function prepareFreshChat(page: Page, agentName: string) { + await page.goto(NEW_CHAT_PATH, { timeout: 15000 }); + await selectAgent(page, agentName); + await expect(page.getByRole('textbox', { name: 'Message input' })).toBeVisible(); + await expect(messagesView(page).getByText(BENCHMARK_REPLY, { exact: true })).toHaveCount(0); +} + +async function prepareConversation(page: Page, agentName: string, sequence: number) { + await prepareFreshChat(page, agentName); + if (BENCHMARK_TURN === 'first') { + return; + } + + const input = page.getByRole('textbox', { name: 'Message input' }); + await input.fill(`agent startup latency seed ${sequence}`); + await input.press('Enter'); + await expect(messagesView(page).getByText(BENCHMARK_REPLY, { exact: true })).toHaveCount(1, { + timeout: 30000, + }); + await expect(page.getByTestId('stop-generation-button')).toHaveCount(0, { timeout: 10000 }); +} + +async function installBrowserObservers(input: Locator) { + return input.evaluate((inputElement, replyText) => { + const latencyWindow = window as typeof window & { + __agentStartupLatency?: BrowserLatencyState; + }; + const state: BrowserLatencyState = { + startedAt: null, + acknowledgedAt: null, + firstContentAt: null, + }; + latencyWindow.__agentStartupLatency = state; + performance.clearResourceTimings(); + const countReplies = () => + Array.from( + document.querySelectorAll('.message-render .agent-turn .message-content'), + ).filter((element) => element.textContent?.includes(replyText)).length; + const replyCountBefore = countReplies(); + + inputElement.addEventListener( + 'keydown', + (event) => { + if ( + event instanceof KeyboardEvent && + event.key === 'Enter' && + !event.shiftKey && + state.startedAt === null + ) { + state.startedAt = performance.now(); + } + }, + { capture: true }, + ); + + const resourceObserver = new PerformanceObserver((entries) => { + if (state.startedAt === null || state.acknowledgedAt !== null) { + return; + } + for (const entry of entries.getEntries()) { + const url = new URL(entry.name); + if (url.origin === location.origin && url.pathname === '/api/agents/chat/agents') { + state.acknowledgedAt = entry.responseEnd; + resourceObserver.disconnect(); + break; + } + } + }); + resourceObserver.observe({ type: 'resource', buffered: true }); + + const mutationObserver = new MutationObserver(() => { + if ( + state.startedAt !== null && + state.firstContentAt === null && + countReplies() > replyCountBefore + ) { + state.firstContentAt = performance.now(); + mutationObserver.disconnect(); + } + }); + mutationObserver.observe(document.body, { + childList: true, + characterData: true, + subtree: true, + }); + return replyCountBefore; + }, BENCHMARK_REPLY); +} + +async function deleteMeasuredConversation(page: Page, token: string) { + const match = new URL(page.url()).pathname.match(/^\/c\/([^/]+)$/); + const conversationId = match?.[1]; + if (!conversationId || conversationId === 'new') { + throw new Error(`Expected a persisted conversation URL, got: ${page.url()}`); + } + await requestJson(page, { + path: '/api/convos', + token, + method: 'DELETE', + body: { arg: { conversationId } }, + }); +} + +async function measureSample(page: Page, agentName: string, token: string, sequence: number) { + await prepareConversation(page, agentName, sequence); + const input = page.getByRole('textbox', { name: 'Message input' }); + await input.fill(`agent startup latency sample ${sequence}`); + await expect(page.getByTestId('send-button')).toBeEnabled(); + const replyCountBefore = await installBrowserObservers(input); + + await input.press('Enter'); + await page.waitForFunction( + () => { + const latencyWindow = window as typeof window & { + __agentStartupLatency?: BrowserLatencyState; + }; + const state = latencyWindow.__agentStartupLatency; + return ( + state?.startedAt != null && state.acknowledgedAt != null && state.firstContentAt != null + ); + }, + null, + { timeout: 30000 }, + ); + + const state = await page.evaluate(() => { + const latencyWindow = window as typeof window & { + __agentStartupLatency?: BrowserLatencyState; + }; + return latencyWindow.__agentStartupLatency; + }); + if (state?.startedAt == null || state.acknowledgedAt == null || state.firstContentAt == null) { + throw new Error('Browser latency observers did not capture all timestamps'); + } + + await expect(messagesView(page).getByText(BENCHMARK_REPLY, { exact: true })).toHaveCount( + replyCountBefore + 1, + ); + await expect(page.getByTestId('stop-generation-button')).toHaveCount(0, { timeout: 10000 }); + + const sample = { + submitToAckMs: round(state.acknowledgedAt - state.startedAt), + submitToFirstContentMs: round(state.firstContentAt - state.startedAt), + ackToFirstContentMs: round(state.firstContentAt - state.acknowledgedAt), + }; + await deleteMeasuredConversation(page, token); + return sample; +} + +async function saveReport(report: object, testInfo: TestInfo) { + const serialized = `${JSON.stringify(report, null, 2)}\n`; + await testInfo.attach('agent-startup-latency.json', { + body: Buffer.from(serialized), + contentType: 'application/json', + }); + + const outputPath = process.env.E2E_LATENCY_OUTPUT; + if (outputPath) { + await mkdir(dirname(outputPath), { recursive: true }); + await writeFile(outputPath, serialized, 'utf8'); + } +} + +test('measures agent-chat startup latency', async ({ page }, testInfo) => { + test.setTimeout(Math.max(120000, (WARMUP_COUNT + SAMPLE_COUNT + 1) * 30000)); + + const hostLoadBefore = loadavg(); + const hostCpuBefore = captureCpuSnapshot(); + const agentName = `E2E Agent Startup Benchmark ${Date.now()}`; + let agent: AgentResponse | undefined; + try { + await page.goto(NEW_CHAT_PATH, { timeout: 15000 }); + agent = await createAgent(page, agentName); + if (BENCHMARK_PROFILE === 'mcp-memory') { + expect(agent.tools).toEqual(expect.arrayContaining(MCP_TOOLS)); + expect(agent.mcpServerNames).toContain(MCP_SERVER_NAME); + } + const token = await getAccessToken(page); + + const cold = await measureSample(page, agentName, token, 0); + for (let index = 0; index < WARMUP_COUNT; index++) { + await measureSample(page, agentName, token, index + 1); + } + + const samples: LatencySample[] = []; + for (let index = 0; index < SAMPLE_COUNT; index++) { + samples.push(await measureSample(page, agentName, token, WARMUP_COUNT + index + 1)); + } + + const report = { + label: process.env.E2E_LATENCY_LABEL ?? 'unlabeled', + gitSha: process.env.E2E_LATENCY_GIT_SHA ?? 'unknown', + streamMode: process.env.E2E_LATENCY_STREAM_MODE ?? 'in-memory', + profile: BENCHMARK_PROFILE, + turn: BENCHMARK_TURN, + simulatedLatency: { + mongoQueryMs: SIMULATED_MONGO_DELAY_MS, + }, + cold, + warmups: WARMUP_COUNT, + samples: SAMPLE_COUNT, + host: { + logicalCpus: availableParallelism(), + loadAverageBefore: hostLoadBefore, + loadAverageAfter: loadavg(), + cpuUtilizationPct: calculateCpuUtilization(hostCpuBefore, captureCpuSnapshot()), + }, + raw: { + submitToAckMs: samples.map((sample) => sample.submitToAckMs), + submitToFirstContentMs: samples.map((sample) => sample.submitToFirstContentMs), + ackToFirstContentMs: samples.map((sample) => sample.ackToFirstContentMs), + }, + summary: summarizeSamples(samples), + }; + + console.log(`AGENT_STARTUP_LATENCY ${JSON.stringify(report)}`); + await saveReport(report, testInfo); + } finally { + await cleanupAgent(page, agent?.id); + } +}); diff --git a/e2e/benchmarks/mongoose-latency-hook.cjs b/e2e/benchmarks/mongoose-latency-hook.cjs new file mode 100644 index 0000000000..cfe2498874 --- /dev/null +++ b/e2e/benchmarks/mongoose-latency-hook.cjs @@ -0,0 +1,20 @@ +const mongoose = require('mongoose'); + +const delayMs = Number.parseInt(process.env.E2E_LATENCY_MONGO_DELAY_MS ?? '', 10); +const patched = Symbol.for('librechat.e2e.mongooseLatencyPatched'); + +function patchExec(prototype) { + if (!Number.isInteger(delayMs) || delayMs <= 0 || prototype[patched]) { + return; + } + + const originalExec = prototype.exec; + Object.defineProperty(prototype, patched, { value: true }); + prototype.exec = async function delayedExec(...args) { + await new Promise((resolve) => setTimeout(resolve, delayMs)); + return originalExec.apply(this, args); + }; +} + +patchExec(mongoose.Query.prototype); +patchExec(mongoose.Aggregate.prototype); diff --git a/e2e/playwright.config.benchmark.ts b/e2e/playwright.config.benchmark.ts new file mode 100644 index 0000000000..1f4c104d71 --- /dev/null +++ b/e2e/playwright.config.benchmark.ts @@ -0,0 +1,29 @@ +import { defineConfig } from '@playwright/test'; +import path from 'node:path'; +import mockConfig from './playwright.config.mock'; + +process.env.MOCK_LLM_REPLY ??= 'BENCH_TOKEN'; +process.env.MOCK_LLM_CHUNK_DELAY_MS ??= '1'; + +const mongoDelay = Number.parseInt(process.env.E2E_LATENCY_MONGO_DELAY_MS ?? '', 10); +if (Number.isInteger(mongoDelay) && mongoDelay > 0) { + const latencyHook = path.resolve(__dirname, 'benchmarks/mongoose-latency-hook.cjs'); + if (!process.env.NODE_OPTIONS?.includes(latencyHook)) { + process.env.NODE_OPTIONS = [process.env.NODE_OPTIONS, `--require=${latencyHook}`] + .filter(Boolean) + .join(' '); + } +} + +export default defineConfig({ + ...mockConfig, + testDir: 'benchmarks', + outputDir: 'benchmarks/.test-results', + timeout: 20 * 60 * 1000, + retries: 0, + reporter: [['line']], + /** The benchmark uses the stdio MCP fixture loaded by LibreChat, not the HTTP fixture. */ + webServer: Array.isArray(mockConfig.webServer) + ? mockConfig.webServer.slice(0, 1) + : mockConfig.webServer, +}); diff --git a/package.json b/package.json index 74b0702cd6..364b2e852a 100644 --- a/package.json +++ b/package.json @@ -59,6 +59,7 @@ "e2e:a11y": "npm run e2e:prepare && playwright test --config=e2e/playwright.config.a11y.ts --headed", "e2e:ci": "npm run e2e:prepare && playwright test --config=e2e/playwright.config.ts", "e2e:mock": "npm run e2e:prepare && playwright test --config=e2e/playwright.config.mock.ts", + "e2e:benchmark:agents": "npm run e2e:prepare && playwright test --config=e2e/playwright.config.benchmark.ts agent-startup.latency.spec.ts", "e2e:mock:enforce": "npm run e2e:prepare && cross-env E2E_MODEL_SPECS_ENFORCE=true playwright test --config=e2e/playwright.config.mock.ts enforced-model-specs.spec.ts", "e2e:mock:ci": "npm run e2e:prepare && playwright test --config=e2e/playwright.config.mock.ts", "e2e:debug": "npm run e2e:prepare && cross-env PWDEBUG=1 playwright test --config=e2e/playwright.config.local.ts", diff --git a/packages/api/src/agents/checkpointer.integration.spec.ts b/packages/api/src/agents/checkpointer.integration.spec.ts index f617a745bc..7e40d70675 100644 --- a/packages/api/src/agents/checkpointer.integration.spec.ts +++ b/packages/api/src/agents/checkpointer.integration.spec.ts @@ -5,6 +5,7 @@ import { MongoDBSaver } from '@langchain/langgraph-checkpoint-mongodb'; import { emptyCheckpoint, ERROR, INTERRUPT } from '@langchain/langgraph-checkpoint'; import { getAgentCheckpointer, + captureAgentCheckpointGeneration, deleteAgentCheckpoint, deleteAgentCheckpoints, LazyMongoSaver, @@ -121,6 +122,39 @@ describe('checkpointer (mongodb-memory-server integration)', () => { expect(await saver!.getTuple(readConfig(threadB))).toBeDefined(); }); + it('generation-scoped cleanup preserves a replacement checkpoint on the same thread', async () => { + const saver = await getAgentCheckpointer(MONGO_CFG); + const threadId = `convo-${new mongoose.Types.ObjectId().toString()}`; + const resumed = await seedInterruptCheckpoint(saver!, threadId); + const generation = await captureAgentCheckpointGeneration(threadId, MONGO_CFG); + + const replacement = await seedInterruptCheckpoint(saver!, threadId); + await deleteAgentCheckpoint(threadId, MONGO_CFG, generation); + + expect(await saver!.getTuple(readConfig(threadId))).toMatchObject({ + checkpoint: { id: replacement.id }, + }); + expect( + await saver!.getTuple({ + configurable: { + thread_id: threadId, + checkpoint_ns: '', + checkpoint_id: resumed.id, + }, + }), + ).toBeUndefined(); + expect( + await mongoose.connection + .db!.collection('agent_checkpoint_writes') + .countDocuments({ thread_id: threadId, checkpoint_id: replacement.id }), + ).toBe(1); + expect( + await mongoose.connection + .db!.collection('agent_checkpoint_writes') + .countDocuments({ thread_id: threadId, checkpoint_id: resumed.id }), + ).toBe(0); + }); + it('deleteAgentCheckpoint is a no-op for an undefined threadId', async () => { await expect(deleteAgentCheckpoint(undefined, MONGO_CFG)).resolves.toBeUndefined(); }); diff --git a/packages/api/src/agents/checkpointer.spec.ts b/packages/api/src/agents/checkpointer.spec.ts index cfadc2df73..3411b7202e 100644 --- a/packages/api/src/agents/checkpointer.spec.ts +++ b/packages/api/src/agents/checkpointer.spec.ts @@ -2,6 +2,7 @@ import { resolveCheckpointerConfig, getApprovalTtlMs, getAgentCheckpointer, + captureAgentCheckpointGeneration, deleteAgentCheckpoint, DEFAULT_CHECKPOINT_TTL_SECONDS, __resetCheckpointerForTests, @@ -69,4 +70,11 @@ describe('deleteAgentCheckpoint', () => { test('is a no-op (no throw) when no durable saver is available', async () => { await expect(deleteAgentCheckpoint('conversation-1')).resolves.toBeUndefined(); }); + + test('captures an empty generation when no durable saver is available', async () => { + await expect(captureAgentCheckpointGeneration('conversation-1')).resolves.toEqual({ + threadId: 'conversation-1', + checkpointIds: [], + }); + }); }); diff --git a/packages/api/src/agents/checkpointer.ts b/packages/api/src/agents/checkpointer.ts index 52a644d2cb..7ac8d39316 100644 --- a/packages/api/src/agents/checkpointer.ts +++ b/packages/api/src/agents/checkpointer.ts @@ -22,7 +22,8 @@ import type { RunnableConfig } from '@langchain/core/runnables'; * * Storage is bounded two ways: a Mongo TTL index reclaims runs that are never * resolved ({@link DEFAULT_CHECKPOINT_TTL_SECONDS}), and {@link deleteAgentCheckpoint} - * prunes a thread's checkpoints eagerly on every terminal transition. + * prunes a thread's checkpoints after ordinary terminal transitions. Approval + * expiry relies on the TTL because a thread-wide eager delete can race a replacement run. */ /** @@ -412,6 +413,18 @@ export interface ResolvedCheckpointerConfig { checkpointWritesCollectionName: string; } +/** + * Exact checkpoint ids present before a resumed generation is claimed. + * + * Terminal resume cleanup deletes only this immutable set. A replacement turn + * that later pauses on the same `thread_id` receives fresh checkpoint ids and + * therefore cannot be removed by the predecessor's delayed cleanup. + */ +export interface AgentCheckpointGeneration { + threadId: string; + checkpointIds: string[]; +} + /** * Apply defaults to the YAML `endpoints.agents.checkpointer` block. Mirrors * {@link resolveRecursionLimit} — the schema stays descriptive, defaults live here. @@ -513,6 +526,46 @@ async function buildMongoSaver( } } +/** + * Snapshot the durable checkpoint ids that belong to the generation about to + * resume. Capture this before atomically claiming the paused job; a replacement + * that wins before the claim makes that claim fail, while one that starts after + * the claim writes ids outside this snapshot. + */ +export async function captureAgentCheckpointGeneration( + threadId: string, + cfg?: TCheckpointerConfig, +): Promise { + const generation: AgentCheckpointGeneration = { threadId, checkpointIds: [] }; + if (!threadId) { + return generation; + } + try { + const saver = await getAgentCheckpointer(cfg); + const db = mongoose.connection.db; + if (!saver || !db) { + return generation; + } + const resolved = resolveCheckpointerConfig(cfg); + const checkpoints = await db + .collection<{ checkpoint_id?: string }>(resolved.checkpointCollectionName) + .find({ thread_id: threadId }, { projection: { _id: 0, checkpoint_id: 1 } }) + .toArray(); + generation.checkpointIds = checkpoints.reduce((ids, checkpoint) => { + if (typeof checkpoint.checkpoint_id === 'string') { + ids.push(checkpoint.checkpoint_id); + } + return ids; + }, []); + } catch (err) { + logger.warn( + `[checkpointer] Failed to capture checkpoint generation for thread ${threadId}:`, + err, + ); + } + return generation; +} + /** * Prune a thread's checkpoints on a terminal transition — natural completion, * abort, or expiry — so the durable store stays bounded. The TTL index is the @@ -520,10 +573,14 @@ async function buildMongoSaver( * has built the saver (nothing to delete). * * @param threadId - the LangGraph `thread_id` (LibreChat's conversationId). + * @param generation - when present, delete only the checkpoint ids captured for + * this resumed generation; omitted by legacy callers that intentionally prune + * the entire thread. */ export async function deleteAgentCheckpoint( threadId: string | undefined, cfg?: TCheckpointerConfig, + generation?: AgentCheckpointGeneration, ): Promise { if (!threadId) { return; @@ -533,6 +590,25 @@ export async function deleteAgentCheckpoint( return; } try { + if (generation) { + if (generation.threadId !== threadId || generation.checkpointIds.length === 0) { + return; + } + const db = mongoose.connection.db; + if (!db) { + return; + } + const resolved = resolveCheckpointerConfig(cfg); + const filter = { + thread_id: threadId, + checkpoint_id: { $in: generation.checkpointIds }, + }; + await Promise.all([ + db.collection(resolved.checkpointCollectionName).deleteMany(filter), + db.collection(resolved.checkpointWritesCollectionName).deleteMany(filter), + ]); + return; + } await saver.deleteThread(threadId); } catch (err) { logger.warn(`[checkpointer] Failed to delete checkpoints for thread ${threadId}:`, err); diff --git a/packages/api/src/agents/conversation.spec.ts b/packages/api/src/agents/conversation.spec.ts new file mode 100644 index 0000000000..1b912f88f8 --- /dev/null +++ b/packages/api/src/agents/conversation.spec.ts @@ -0,0 +1,77 @@ +import { resolveConversationAnchor } from './conversation'; + +describe('resolveConversationAnchor', () => { + const fallback = new Date('2026-07-24T12:00:00.000Z'); + + it('anchors a new conversation without loading an existing one', async () => { + const loadConversation = jest.fn(); + + const result = await resolveConversationAnchor({ + isNewConversation: true, + loadConversation, + now: () => fallback, + }); + + expect(result).toEqual({ + createdAt: fallback.toISOString(), + conversation: undefined, + }); + expect(loadConversation).not.toHaveBeenCalled(); + }); + + it('returns an existing conversation with a normalized creation time', async () => { + const conversation = { + conversationId: 'conversation-1', + createdAt: new Date('2025-01-02T03:04:05.000Z'), + }; + + const result = await resolveConversationAnchor({ + isNewConversation: false, + loadConversation: async () => conversation, + now: () => fallback, + }); + + expect(result).toEqual({ + createdAt: '2025-01-02T03:04:05.000Z', + conversation, + }); + }); + + it.each([null, undefined, 'not-a-date'])( + 'uses the fallback time for a missing or invalid creation time: %p', + async (createdAt) => { + const conversation = createdAt === undefined ? null : { createdAt }; + + const result = await resolveConversationAnchor({ + isNewConversation: false, + loadConversation: async () => conversation, + now: () => fallback, + }); + + expect(result).toEqual({ + createdAt: fallback.toISOString(), + conversation, + }); + }, + ); + + it('reports load failures and degrades to a fresh anchor', async () => { + const onLoadError = jest.fn(); + const failure = new Error('conversation store unavailable'); + + const result = await resolveConversationAnchor({ + isNewConversation: false, + loadConversation: async () => { + throw failure; + }, + now: () => fallback, + onLoadError, + }); + + expect(result).toEqual({ + createdAt: fallback.toISOString(), + conversation: undefined, + }); + expect(onLoadError).toHaveBeenCalledWith(failure); + }); +}); diff --git a/packages/api/src/agents/conversation.ts b/packages/api/src/agents/conversation.ts new file mode 100644 index 0000000000..a58b0b424d --- /dev/null +++ b/packages/api/src/agents/conversation.ts @@ -0,0 +1,53 @@ +export interface ConversationAnchorSource { + createdAt?: Date | string | number | null; +} + +export interface ConversationAnchor { + createdAt: string; + conversation: TConversation | null | undefined; +} + +interface ResolveConversationAnchorOptions { + isNewConversation: boolean; + loadConversation: () => Promise; + now?: () => Date; + onLoadError?: (error: Error) => void; +} + +function toValidISOString(value: Date | string | number | null | undefined): string | undefined { + if (value == null) { + return; + } + + const date = value instanceof Date ? value : new Date(value); + return Number.isNaN(date.getTime()) ? undefined : date.toISOString(); +} + +export async function resolveConversationAnchor({ + isNewConversation, + loadConversation, + now = () => new Date(), + onLoadError, +}: ResolveConversationAnchorOptions): Promise> { + if (isNewConversation) { + return { + createdAt: now().toISOString(), + conversation: undefined, + }; + } + + try { + const conversation = await loadConversation(); + return { + createdAt: toValidISOString(conversation?.createdAt) ?? now().toISOString(), + conversation, + }; + } catch (error) { + const normalizedError = error instanceof Error ? error : new Error(String(error)); + onLoadError?.(normalizedError); + return { + createdAt: now().toISOString(), + conversation: undefined, + }; + } +} diff --git a/packages/api/src/agents/index.ts b/packages/api/src/agents/index.ts index ced3c40fa2..115525fd25 100644 --- a/packages/api/src/agents/index.ts +++ b/packages/api/src/agents/index.ts @@ -6,6 +6,7 @@ export * from './config'; export * from './checkpointer'; export * from './contact'; export * from './context'; +export * from './conversation'; export * from './discovery'; export * from './edges'; export * from './handlers'; @@ -23,6 +24,8 @@ export * from './usage'; export * from './resources'; export * from './responses'; export * from './skills'; +export * from './phases'; +export * from './startup'; export * from './skillConfigurable'; export * from './skillFiles'; export * from './codeFilesSession'; diff --git a/packages/api/src/agents/phases.ts b/packages/api/src/agents/phases.ts new file mode 100644 index 0000000000..b811b79134 --- /dev/null +++ b/packages/api/src/agents/phases.ts @@ -0,0 +1,30 @@ +export const agentStartupMilestones = [ + 'request_admitted', + 'job_created', + 'ack_sent', + 'conversation_resolved', + 'metadata_persisted', + 'client_initialized', + 'history_loaded', + 'messages_built', + 'run_input_prepared', + 'run_created', + 'stream_processing_started', + 'request_message_queued', + 'first_response_event_queued', + 'first_content_delta_queued', +] as const; + +export const agentStartupResults = [ + 'content_queued', + 'completed_without_delta', + 'deduplicated', + 'rejected', + 'paused', + 'replaced', + 'aborted', + 'error', +] as const; + +export type AgentStartupMilestone = (typeof agentStartupMilestones)[number]; +export type AgentStartupResult = (typeof agentStartupResults)[number]; diff --git a/packages/api/src/agents/startup.spec.ts b/packages/api/src/agents/startup.spec.ts new file mode 100644 index 0000000000..42416cec89 --- /dev/null +++ b/packages/api/src/agents/startup.spec.ts @@ -0,0 +1,322 @@ +import { EventEmitter } from 'node:events'; +import { context, SpanKind, SpanStatusCode, trace } from '@opentelemetry/api'; +import type { Span, Tracer } from '@opentelemetry/api'; +import type { Response } from 'express'; +import type { ServerRequest } from '~/types'; +import { + acceptAgentStartupTelemetry, + agentStartupIngressMiddleware, + agentStartupTelemetryMiddleware, + createAgentStartupTelemetry, + getAgentStartupTelemetry, +} from './startup'; +import { + isMetricsConfigured, + recordAgentStartupMilestone, + recordAgentStartupResult, +} from '~/app/metrics'; + +jest.mock('~/app/metrics', () => ({ + isMetricsConfigured: jest.fn(() => true), + recordAgentStartupMilestone: jest.fn(), + recordAgentStartupResult: jest.fn(), +})); + +interface MockResponse extends EventEmitter { + statusCode: number; + locals: Record; +} + +function createSpan(): jest.Mocked { + const span = {} as jest.Mocked; + span.addEvent = jest.fn, Parameters>(() => span); + span.addLink = jest.fn, Parameters>(() => span); + span.addLinks = jest.fn, Parameters>(() => span); + span.end = jest.fn>(); + span.isRecording = jest.fn>(() => true); + span.recordException = jest.fn>(); + span.setAttribute = jest.fn, Parameters>(() => span); + span.setAttributes = jest.fn, Parameters>(() => span); + span.setStatus = jest.fn, Parameters>(() => span); + span.spanContext = jest.fn, Parameters>( + () => ({ + spanId: '0000000000000000', + traceFlags: 0, + traceId: '00000000000000000000000000000000', + }), + ); + span.updateName = jest.fn, Parameters>(() => span); + return span; +} + +function mockTracer(span: jest.Mocked): jest.Mock { + const startSpan = jest.fn(() => span); + jest.spyOn(trace, 'getTracer').mockReturnValue({ startSpan } as unknown as Tracer); + return startSpan; +} + +function createRequest(path = '/'): ServerRequest { + return { + method: 'POST', + path, + } as ServerRequest; +} + +function createResponse(statusCode = 200): MockResponse { + const res = new EventEmitter() as MockResponse; + res.statusCode = statusCode; + res.locals = {}; + return res; +} + +afterEach(() => { + jest.clearAllMocks(); + jest.restoreAllMocks(); +}); + +describe('createAgentStartupTelemetry', () => { + it('records cumulative milestones once and ends on the first renderable delta', () => { + const span = createSpan(); + const startSpan = mockTracer(span); + let now = 100; + const telemetry = createAgentStartupTelemetry({ now: () => now })!; + + now = 125; + telemetry.mark('job_created'); + now = 140; + telemetry.mark('job_created'); + now = 150; + telemetry.recordGenerationEvent({ event: 'on_run_step', data: {} }); + now = 165; + telemetry.recordGenerationEvent({ + event: 'on_message_delta', + data: { delta: { content: { text: '' } } }, + }); + now = 180; + telemetry.recordGenerationEvent({ + event: 'on_message_delta', + data: { delta: { content: [{ text: 'Hello' }] } }, + }); + now = 200; + telemetry.recordGenerationEvent({ + event: 'on_reasoning_delta', + data: { delta: { content: { think: 'Already ended' } } }, + }); + + expect(startSpan).toHaveBeenCalledWith( + 'librechat.agent.startup', + { kind: SpanKind.INTERNAL }, + context.active(), + ); + expect(span.addEvent).toHaveBeenNthCalledWith(1, 'job_created', { + 'librechat.agent.startup.elapsed_ms': 25, + }); + expect(span.addEvent).toHaveBeenNthCalledWith(2, 'first_response_event_queued', { + 'librechat.agent.startup.elapsed_ms': 50, + }); + expect(span.addEvent).toHaveBeenNthCalledWith(3, 'first_content_delta_queued', { + 'librechat.agent.startup.elapsed_ms': 80, + }); + expect(recordAgentStartupMilestone).toHaveBeenCalledTimes(3); + expect(recordAgentStartupMilestone).toHaveBeenNthCalledWith(1, 'job_created', 0.025); + expect(recordAgentStartupMilestone).toHaveBeenNthCalledWith( + 2, + 'first_response_event_queued', + 0.05, + ); + expect(recordAgentStartupMilestone).toHaveBeenNthCalledWith( + 3, + 'first_content_delta_queued', + 0.08, + ); + expect(recordAgentStartupResult).toHaveBeenCalledWith('content_queued'); + expect(span.setAttributes).toHaveBeenCalledWith({ + 'librechat.agent.startup.duration_ms': 80, + 'librechat.agent.startup.milestones.count': 3, + 'librechat.agent.startup.result': 'content_queued', + }); + expect(span.end).toHaveBeenCalledTimes(1); + }); + + it('records terminal errors once', () => { + const span = createSpan(); + mockTracer(span); + let now = 10; + const telemetry = createAgentStartupTelemetry({ now: () => now })!; + const error = new Error('startup failed'); + + now = 25; + telemetry.end('error', error); + now = 30; + telemetry.end('aborted'); + telemetry.mark('client_initialized'); + + expect(span.recordException).toHaveBeenCalledWith(error); + expect(span.setStatus).toHaveBeenCalledWith({ code: SpanStatusCode.ERROR }); + expect(recordAgentStartupResult).toHaveBeenCalledTimes(1); + expect(recordAgentStartupResult).toHaveBeenCalledWith('error'); + expect(span.end).toHaveBeenCalledTimes(1); + expect(span.addEvent).not.toHaveBeenCalled(); + }); + + it('drops untyped milestones and normalizes untyped terminal results', () => { + const span = createSpan(); + mockTracer(span); + const telemetry = createAgentStartupTelemetry({ now: () => 10 })!; + + Reflect.apply(telemetry.mark, undefined, ['unbounded-user-value']); + telemetry.mark('job_created'); + Reflect.apply(telemetry.end, undefined, ['unbounded-user-value']); + + expect(span.addEvent).toHaveBeenCalledTimes(1); + expect(span.addEvent).toHaveBeenCalledWith('job_created', expect.any(Object)); + expect(recordAgentStartupResult).toHaveBeenCalledWith('error'); + expect(span.setAttributes).toHaveBeenCalledWith( + expect.objectContaining({ + 'librechat.agent.startup.milestones.count': 1, + 'librechat.agent.startup.result': 'error', + }), + ); + expect(span.setStatus).toHaveBeenCalledWith({ code: SpanStatusCode.ERROR }); + expect(span.end).toHaveBeenCalledTimes(1); + }); + + it('returns no recorder when tracing and metrics are disabled', () => { + const span = createSpan(); + span.isRecording.mockReturnValue(false); + jest.mocked(isMetricsConfigured).mockReturnValueOnce(false); + mockTracer(span); + + const telemetry = createAgentStartupTelemetry(); + + expect(telemetry).toBeUndefined(); + expect(recordAgentStartupMilestone).not.toHaveBeenCalled(); + expect(recordAgentStartupResult).not.toHaveBeenCalled(); + expect(span.end).toHaveBeenCalledTimes(1); + }); + + it('skips span work when only metrics are enabled', () => { + const span = createSpan(); + span.isRecording.mockReturnValue(false); + mockTracer(span); + const telemetry = createAgentStartupTelemetry({ now: () => 10 })!; + + telemetry.mark('job_created'); + telemetry.end('content_queued'); + + expect(recordAgentStartupMilestone).toHaveBeenCalledWith('job_created', 0); + expect(recordAgentStartupResult).toHaveBeenCalledWith('content_queued'); + expect(span.addEvent).not.toHaveBeenCalled(); + expect(span.setAttributes).not.toHaveBeenCalled(); + expect(span.end).toHaveBeenCalledTimes(1); + }); + + it('skips metric work when only tracing is enabled', () => { + const span = createSpan(); + jest.mocked(isMetricsConfigured).mockReturnValueOnce(false); + mockTracer(span); + const telemetry = createAgentStartupTelemetry({ now: () => 10 })!; + + telemetry.mark('job_created'); + telemetry.end('content_queued'); + + expect(span.addEvent).toHaveBeenCalledWith('job_created', expect.any(Object)); + expect(span.setAttributes).toHaveBeenCalled(); + expect(recordAgentStartupMilestone).not.toHaveBeenCalled(); + expect(recordAgentStartupResult).not.toHaveBeenCalled(); + expect(span.end).toHaveBeenCalledTimes(1); + }); +}); + +describe('agentStartupTelemetryMiddleware', () => { + it('carries the outer ingress timestamp into the recorder', () => { + const span = createSpan(); + const startSpan = mockTracer(span); + const req = createRequest(); + const res = createResponse(); + const next = jest.fn(); + jest.spyOn(Date, 'now').mockReturnValue(1_750_000_000_000); + + agentStartupIngressMiddleware(req, res as Response, next); + agentStartupTelemetryMiddleware(req, res as Response, next); + + expect(next).toHaveBeenCalledTimes(2); + expect(getAgentStartupTelemetry(req)).toBeDefined(); + expect(startSpan).toHaveBeenCalledWith( + 'librechat.agent.startup', + { + kind: SpanKind.INTERNAL, + startTime: 1_750_000_000_000, + }, + context.active(), + ); + }); + + it('records the ACK without ending an accepted startup', () => { + const span = createSpan(); + mockTracer(span); + const req = createRequest(); + const res = createResponse(); + const next = jest.fn(); + + agentStartupTelemetryMiddleware(req, res as Response, next); + const telemetry = getAgentStartupTelemetry(req); + acceptAgentStartupTelemetry(req, 'stream-123'); + res.emit('finish'); + res.emit('close'); + + expect(next).toHaveBeenCalledTimes(1); + expect(telemetry).toBeDefined(); + expect(span.setAttribute).toHaveBeenCalledWith('librechat.stream.id', 'stream-123'); + expect(recordAgentStartupMilestone).toHaveBeenCalledWith('ack_sent', expect.any(Number)); + expect(recordAgentStartupResult).not.toHaveBeenCalled(); + expect(span.end).not.toHaveBeenCalled(); + }); + + it('finalizes requests rejected before job creation', () => { + const span = createSpan(); + mockTracer(span); + const req = createRequest(); + const res = createResponse(403); + const next = jest.fn(); + + agentStartupTelemetryMiddleware(req, res as Response, next); + res.emit('finish'); + res.emit('close'); + + expect(recordAgentStartupResult).toHaveBeenCalledTimes(1); + expect(recordAgentStartupResult).toHaveBeenCalledWith('rejected'); + expect(span.end).toHaveBeenCalledTimes(1); + }); + + it('skips resume requests', () => { + const span = createSpan(); + mockTracer(span); + const req = createRequest('/resume'); + const res = createResponse(); + const next = jest.fn(); + + agentStartupTelemetryMiddleware(req, res as Response, next); + + expect(next).toHaveBeenCalledTimes(1); + expect(getAgentStartupTelemetry(req)).toBeUndefined(); + expect(trace.getTracer).not.toHaveBeenCalled(); + }); + + it('does not retain listeners or request state when telemetry is disabled', () => { + const span = createSpan(); + span.isRecording.mockReturnValue(false); + jest.mocked(isMetricsConfigured).mockReturnValueOnce(false); + mockTracer(span); + const req = createRequest(); + const res = createResponse(); + const next = jest.fn(); + + agentStartupTelemetryMiddleware(req, res as Response, next); + + expect(next).toHaveBeenCalledTimes(1); + expect(getAgentStartupTelemetry(req)).toBeUndefined(); + expect(res.listenerCount('finish')).toBe(0); + expect(res.listenerCount('close')).toBe(0); + }); +}); diff --git a/packages/api/src/agents/startup.ts b/packages/api/src/agents/startup.ts new file mode 100644 index 0000000000..af70e61ea4 --- /dev/null +++ b/packages/api/src/agents/startup.ts @@ -0,0 +1,285 @@ +import { performance } from 'node:perf_hooks'; +import { ApprovalEvents, StepEvents } from 'librechat-data-provider'; +import { context, SpanKind, SpanStatusCode, trace } from '@opentelemetry/api'; +import type { NextFunction, Response } from 'express'; +import type { AgentStartupMilestone, AgentStartupResult } from './phases'; +import type { ServerRequest, ServerSentEvent } from '~/types'; +import { + isMetricsConfigured, + recordAgentStartupMilestone, + recordAgentStartupResult, +} from '~/app/metrics'; +import { agentStartupMilestones, agentStartupResults } from './phases'; + +const SPAN_NAME = 'librechat.agent.startup'; +const MILESTONES = new Set(agentStartupMilestones); +const RESULTS = new Set(agentStartupResults); +const RESPONSE_EVENTS = new Set([ + ...Object.values(StepEvents), + ApprovalEvents.ON_PENDING_ACTION, + 'attachment', +]); + +interface AgentGenerationEventData { + delta?: { + content?: + | { + text?: string; + think?: string; + } + | Array<{ + text?: string; + think?: string; + }>; + }; +} + +export interface AgentStartupTelemetry { + mark: (milestone: AgentStartupMilestone) => void; + setStreamId: (streamId: string) => void; + recordGenerationEvent: (event: ServerSentEvent) => boolean; + end: (result: AgentStartupResult, error?: Error) => void; +} + +interface AgentStartupTelemetryOptions { + now?: () => number; + startedAt?: number; + spanStartedAt?: number; +} + +interface AgentStartupState { + accepted: boolean; + telemetry: AgentStartupTelemetry; +} + +const requestTelemetry = new WeakMap(); +const AGENT_STARTUP_STARTED_AT = Symbol('agentStartupStartedAt'); +const EXCLUDED_AGENT_CHAT_PATHS = new Set(['/abort', '/resume', '/steer', '/steer/cancel']); + +interface AgentStartupIngressTime { + monotonic: number; + epoch: number; +} + +function isInitialAgentChatRequest(req: ServerRequest): boolean { + return req.method === 'POST' && !EXCLUDED_AGENT_CHAT_PATHS.has(req.path); +} + +function isRenderableDelta(event: ServerSentEvent): boolean { + if ( + !('event' in event) || + (event.event !== StepEvents.ON_MESSAGE_DELTA && + event.event !== StepEvents.ON_REASONING_DELTA) || + typeof event.data === 'string' + ) { + return false; + } + + const content = (event.data as AgentGenerationEventData).delta?.content; + const parts = Array.isArray(content) ? content : [content]; + return parts.some( + (part) => + (typeof part?.text === 'string' && part.text.length > 0) || + (typeof part?.think === 'string' && part.think.length > 0), + ); +} + +function isResponseEvent(event: ServerSentEvent, renderableDelta: boolean): boolean { + if ('final' in event) { + return true; + } + if (!('event' in event)) { + return false; + } + if ( + event.event === StepEvents.ON_MESSAGE_DELTA || + event.event === StepEvents.ON_REASONING_DELTA + ) { + return renderableDelta; + } + return RESPONSE_EVENTS.has(event.event); +} + +export function createAgentStartupTelemetry( + options: AgentStartupTelemetryOptions = {}, +): AgentStartupTelemetry | undefined { + const now = options.now ?? (() => performance.now()); + const startedAt = options.startedAt ?? now(); + const spanOptions = { + kind: SpanKind.INTERNAL, + ...(options.spanStartedAt != null && { startTime: options.spanStartedAt }), + }; + const span = trace + .getTracer('librechat.telemetry') + .startSpan(SPAN_NAME, spanOptions, context.active()); + const tracingEnabled = span.isRecording(); + const metricsEnabled = isMetricsConfigured(); + if (!tracingEnabled) { + span.end(); + } + if (!tracingEnabled && !metricsEnabled) { + return undefined; + } + const milestones = new Set(); + let ended = false; + + const elapsedMilliseconds = (): number => Math.max(0, now() - startedAt); + + const mark = (milestone: AgentStartupMilestone): void => { + if (ended || !MILESTONES.has(milestone) || milestones.has(milestone)) { + return; + } + + milestones.add(milestone); + const elapsedMs = elapsedMilliseconds(); + if (tracingEnabled) { + span.addEvent(milestone, { + 'librechat.agent.startup.elapsed_ms': elapsedMs, + }); + } + if (metricsEnabled) { + recordAgentStartupMilestone(milestone, elapsedMs / 1_000); + } + }; + + const setStreamId = (streamId: string): void => { + if (ended || !streamId) { + return; + } + if (tracingEnabled) { + span.setAttribute('librechat.stream.id', streamId); + } + }; + + const end = (result: AgentStartupResult, error?: Error): void => { + if (ended) { + return; + } + + ended = true; + const normalizedResult: AgentStartupResult = RESULTS.has(result) ? result : 'error'; + if (tracingEnabled) { + span.setAttributes({ + 'librechat.agent.startup.duration_ms': elapsedMilliseconds(), + 'librechat.agent.startup.milestones.count': milestones.size, + 'librechat.agent.startup.result': normalizedResult, + }); + } + if (metricsEnabled) { + recordAgentStartupResult(normalizedResult); + } + + if (tracingEnabled && error) { + span.recordException(error); + } + if (tracingEnabled && (normalizedResult === 'aborted' || normalizedResult === 'error')) { + span.setStatus({ code: SpanStatusCode.ERROR }); + } + + if (tracingEnabled) { + span.end(); + } + }; + + const recordGenerationEvent = (event: ServerSentEvent): boolean => { + if (ended) { + return true; + } + + const renderableDelta = isRenderableDelta(event); + if (!isResponseEvent(event, renderableDelta)) { + return false; + } + + mark('first_response_event_queued'); + if (!renderableDelta) { + return false; + } + + mark('first_content_delta_queued'); + end('content_queued'); + return true; + }; + + return { mark, setStreamId, recordGenerationEvent, end }; +} + +export function getAgentStartupTelemetry(req: ServerRequest): AgentStartupTelemetry | undefined { + return requestTelemetry.get(req)?.telemetry; +} + +export function acceptAgentStartupTelemetry(req: ServerRequest, streamId: string): void { + const state = requestTelemetry.get(req); + if (state) { + state.accepted = true; + state.telemetry.setStreamId(streamId); + } +} + +/** + * Capture the outer request timestamp before body parsing and auth. The recorder is + * created later, after the HTTP tracing middleware has installed its active context. + */ +export function agentStartupIngressMiddleware( + req: ServerRequest, + res: Response, + next: NextFunction, +): void { + if (isInitialAgentChatRequest(req)) { + const ingressTime: AgentStartupIngressTime = { + monotonic: performance.now(), + epoch: Date.now(), + }; + (res.locals as Record)[AGENT_STARTUP_STARTED_AT] = ingressTime; + } + next(); +} + +export function agentStartupTelemetryMiddleware( + req: ServerRequest, + res: Response, + next: NextFunction, +): void { + if (!isInitialAgentChatRequest(req)) { + next(); + return; + } + + const locals = res.locals as Record; + const ingressTime = locals[AGENT_STARTUP_STARTED_AT] as AgentStartupIngressTime | undefined; + delete locals[AGENT_STARTUP_STARTED_AT]; + const telemetry = createAgentStartupTelemetry({ + startedAt: ingressTime?.monotonic, + spanStartedAt: ingressTime?.epoch, + }); + if (!telemetry) { + next(); + return; + } + const state: AgentStartupState = { accepted: false, telemetry }; + requestTelemetry.set(req, state); + + let responseEnded = false; + const endBeforeAcceptance = (result: AgentStartupResult): void => { + if (responseEnded) { + return; + } + responseEnded = true; + if (!state.accepted) { + telemetry.end(result); + } + }; + + res.once('finish', () => { + if (state.accepted) { + telemetry.mark('ack_sent'); + return; + } + endBeforeAcceptance(res.statusCode >= 500 ? 'error' : 'rejected'); + }); + res.once('close', () => { + endBeforeAcceptance('aborted'); + }); + + next(); +} diff --git a/packages/api/src/app/metrics.spec.ts b/packages/api/src/app/metrics.spec.ts index 7714330f2f..7422ccaa37 100644 --- a/packages/api/src/app/metrics.spec.ts +++ b/packages/api/src/app/metrics.spec.ts @@ -7,6 +7,8 @@ import { createMetrics, instrumentMongooseQueryMetrics, normalizePath, + recordAgentStartupMilestone, + recordAgentStartupResult, recordGenerationJob, recordGenerationStreamResumePendingEvents, recordGenerationStreamSubscription, @@ -436,4 +438,35 @@ describe('createMetrics', () => { /generation_stream_resume_pending_events_total\{store="memory"\} 3/, ); }); + + it('tracks cumulative agent startup milestones and terminal results', async () => { + const app = express(); + process.env.METRICS_SECRET = 'test-secret'; + const { metricsRouter } = createMetrics(); + app.use('/metrics', metricsRouter); + + recordAgentStartupMilestone('job_created', 0.125); + recordAgentStartupMilestone('first_response_event_queued', 0.75); + recordAgentStartupResult('content_queued'); + Reflect.apply(recordAgentStartupMilestone, undefined, ['unbounded-user-value', 1]); + Reflect.apply(recordAgentStartupMilestone, undefined, ['job_created', Number.NaN]); + Reflect.apply(recordAgentStartupResult, undefined, ['unbounded-user-value']); + + const response = await request(app) + .get('/metrics') + .set('Authorization', 'Bearer test-secret') + .expect(200); + + expect(response.text).toMatch( + /agent_startup_milestone_duration_seconds_count\{milestone="job_created"\} 1/, + ); + expect(response.text).toMatch( + /agent_startup_milestone_duration_seconds_sum\{milestone="job_created"\} 0.125/, + ); + expect(response.text).toMatch( + /agent_startup_milestone_duration_seconds_count\{milestone="first_response_event_queued"\} 1/, + ); + expect(response.text).toMatch(/agent_startups_total\{result="content_queued"\} 1/); + expect(response.text).not.toContain('unbounded-user-value'); + }); }); diff --git a/packages/api/src/app/metrics.ts b/packages/api/src/app/metrics.ts index b10556cba8..5b02970dce 100644 --- a/packages/api/src/app/metrics.ts +++ b/packages/api/src/app/metrics.ts @@ -4,6 +4,8 @@ import { logger } from '@librechat/data-schemas'; import { Registry, collectDefaultMetrics, Counter, Gauge, Histogram } from 'prom-client'; import type { Request, Response, NextFunction, RequestHandler } from 'express'; import type { Mongoose } from 'mongoose'; +import type { AgentStartupMilestone, AgentStartupResult } from '~/agents/phases'; +import { agentStartupMilestones, agentStartupResults } from '~/agents/phases'; const PATH_NORMALIZATIONS: [RegExp, string][] = [ [/^\/api\/agents\/chat\/stream\/[^/]+(?=\/|$)/, '/api/agents/chat/stream/#id'], @@ -192,6 +194,19 @@ let generationJobMetrics: GenerationJobMetrics = { recordResumePendingEvents: () => undefined, }; +type AgentStartupMetrics = { + recordMilestone: (milestone: AgentStartupMilestone, durationSeconds: number) => void; + recordResult: (result: AgentStartupResult) => void; +}; + +const agentStartupMilestoneSet = new Set(agentStartupMilestones); +const agentStartupResultSet = new Set(agentStartupResults); + +let agentStartupMetrics: AgentStartupMetrics = { + recordMilestone: () => undefined, + recordResult: () => undefined, +}; + type RumProxyMetrics = { recordRequest: (endpoint: RumProxyEndpoint, result: RumProxyResult) => void; }; @@ -227,6 +242,10 @@ const resetMetricRecorders = (): void => { recordSubscription: () => undefined, recordResumePendingEvents: () => undefined, }; + agentStartupMetrics = { + recordMilestone: () => undefined, + recordResult: () => undefined, + }; rumProxyMetrics = { recordRequest: () => undefined, }; @@ -258,6 +277,27 @@ export function recordGenerationStreamResumePendingEvents( generationJobMetrics.recordResumePendingEvents(store, count); } +export function recordAgentStartupMilestone( + milestone: AgentStartupMilestone, + durationSeconds: number, +): void { + if ( + !agentStartupMilestoneSet.has(milestone) || + !Number.isFinite(durationSeconds) || + durationSeconds < 0 + ) { + return; + } + agentStartupMetrics.recordMilestone(milestone, durationSeconds); +} + +export function recordAgentStartupResult(result: AgentStartupResult): void { + if (!agentStartupResultSet.has(result)) { + return; + } + agentStartupMetrics.recordResult(result); +} + export function recordRumProxyRequest(endpoint: RumProxyEndpoint, result: RumProxyResult): void { rumProxyMetrics.recordRequest(endpoint, result); } @@ -548,6 +588,21 @@ export function createMetrics(): PrometheusMetrics { registers: [registry], }); + const agentStartupMilestoneDuration = new Histogram({ + name: 'agent_startup_milestone_duration_seconds', + help: 'Cumulative agent chat startup latency from request ingress to each milestone', + labelNames: ['milestone'] as const, + buckets: [0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 30, 60, 120, 300], + registers: [registry], + }); + + const agentStartups = new Counter({ + name: 'agent_startups_total', + help: 'Agent chat startup attempts by terminal result', + labelNames: ['result'] as const, + registers: [registry], + }); + const rumProxyRequests = new Counter({ name: 'rum_proxy_requests_total', help: 'RUM proxy requests by endpoint and result', @@ -579,6 +634,12 @@ export function createMetrics(): PrometheusMetrics { generationStreamResumePendingEvents.inc({ store }, count), }; + agentStartupMetrics = { + recordMilestone: (milestone, durationSeconds) => + agentStartupMilestoneDuration.observe({ milestone }, durationSeconds), + recordResult: (result) => agentStartups.inc({ result }), + }; + rumProxyMetrics = { recordRequest: (endpoint, result) => rumProxyRequests.inc({ endpoint, result }), }; diff --git a/packages/api/src/app/shutdown.spec.ts b/packages/api/src/app/shutdown.spec.ts index 2dba3d3fbe..7733bb5fed 100644 --- a/packages/api/src/app/shutdown.spec.ts +++ b/packages/api/src/app/shutdown.spec.ts @@ -192,6 +192,35 @@ describe('setupGracefulShutdown', () => { expect(exitSpy).toHaveBeenCalledWith(0); }); + it('runs pre-drain tasks while server.close is pending, then post-drain tasks', async () => { + const calls: string[] = []; + let finishClose: (() => void) | undefined; + jest.spyOn(server, 'close').mockImplementation((cb?: (err?: Error) => void) => { + calls.push('server.close'); + finishClose = () => cb?.(); + return server; + }); + registerShutdownTask( + 'release-held-close', + () => { + calls.push('pre-drain'); + finishClose?.(); + }, + { phase: 'pre-drain' }, + ); + registerShutdownTask('post-drain', () => { + calls.push('post-drain'); + }); + + setupGracefulShutdown(server); + triggerSignal('SIGTERM'); + await flush(); + await flush(); + + expect(calls).toEqual(['server.close', 'pre-drain', 'post-drain']); + expect(exitSpy).toHaveBeenCalledWith(0); + }); + it('runs tasks in registration order', async () => { const order: string[] = []; jest.spyOn(server, 'close').mockImplementation((cb?: (err?: Error) => void) => { @@ -212,6 +241,45 @@ describe('setupGracefulShutdown', () => { expect(order).toEqual(['first', 'second', 'third']); }); + it('runs higher-priority cleanup before telemetry while preserving ties', async () => { + const order: string[] = []; + jest.spyOn(server, 'close').mockImplementation((cb?: (err?: Error) => void) => { + if (cb) { + setImmediate(() => cb()); + } + return server; + }); + registerShutdownTask('default-first', () => { + order.push('default-first'); + }); + registerShutdownTask( + 'telemetry', + () => { + order.push('telemetry'); + }, + { priority: -100 }, + ); + registerShutdownTask( + 'generation streams', + () => { + order.push('generation streams'); + }, + { + priority: 100, + }, + ); + registerShutdownTask('default-second', () => { + order.push('default-second'); + }); + + setupGracefulShutdown(server); + triggerSignal('SIGTERM'); + await flush(); + await flush(); + + expect(order).toEqual(['generation streams', 'default-first', 'default-second', 'telemetry']); + }); + it('continues subsequent tasks and still exits if one task throws', async () => { const calls: string[] = []; jest.spyOn(server, 'close').mockImplementation((cb?: (err?: Error) => void) => { diff --git a/packages/api/src/app/shutdown.ts b/packages/api/src/app/shutdown.ts index 8ae79c979c..dd131fb7d9 100644 --- a/packages/api/src/app/shutdown.ts +++ b/packages/api/src/app/shutdown.ts @@ -4,33 +4,49 @@ import type { Server } from 'http'; const SHUTDOWN_TIMEOUT_MS = 60_000; const SIGNALS: NodeJS.Signals[] = ['SIGTERM', 'SIGINT', 'SIGQUIT', 'SIGHUP']; +export type ShutdownPhase = 'pre-drain' | 'post-drain'; + +export type ShutdownTaskOptions = { + priority?: number; + phase?: ShutdownPhase; +}; + type ShutdownTask = { name: string; fn: () => void | Promise; + phase: ShutdownPhase; + priority: number; + registrationOrder: number; }; const tasks: ShutdownTask[] = []; +let nextRegistrationOrder = 0; let isShuttingDown = false; let httpServer: Server | null = null; /** - * Register a cleanup task to run after the HTTP server has closed. - * Tasks run in registration order; if one throws, subsequent tasks - * and the final exit are not blocked. Use this instead of attaching - * `process.on('SIGTERM', ...)` handlers directly — multiple competing - * signal handlers race with the HTTP drain because Node dispatches - * listeners in registration order and any one of them can call - * `process.exit` before the HTTP server has finished closing. + * Register a cleanup task for graceful shutdown. Post-drain is the default phase. + * Higher-priority tasks run first; tasks at the same priority retain registration order. + * If one throws, subsequent tasks and the final exit are not blocked. Use this instead of + * attaching `process.on('SIGTERM', ...)` handlers directly — multiple competing signal + * handlers race with the HTTP drain because Node dispatches listeners in registration order + * and any one of them can call `process.exit` before the HTTP server has finished closing. */ -export function registerShutdownTask(name: string, fn: () => void | Promise): void { - tasks.push({ name, fn }); +export function registerShutdownTask( + name: string, + fn: () => void | Promise, + options: ShutdownTaskOptions = {}, +): void { + const phase: ShutdownPhase = options.phase === 'pre-drain' ? 'pre-drain' : 'post-drain'; + const priority = Number.isFinite(options.priority) ? (options.priority ?? 0) : 0; + tasks.push({ name, fn, phase, priority, registrationOrder: nextRegistrationOrder++ }); } /** * Wires SIGTERM, SIGINT, SIGQUIT, and SIGHUP to a graceful shutdown - * sequence: close the HTTP server (stop accepting new connections, let - * in-flight requests finish), run any tasks registered via - * `registerShutdownTask`, then `process.exit(0)`. After + * sequence: initiate HTTP server close to stop accepting new connections, + * run pre-drain tasks while in-flight requests settle, await the HTTP drain, + * run post-drain tasks, then `process.exit(0)`. After * SHUTDOWN_TIMEOUT_MS the process is force-exited with code 1 — a * safety net for long-lived connections such as SSE streams that may * not finish in time. @@ -49,10 +65,29 @@ export function setupGracefulShutdown(server: Server): void { */ export function __resetShutdownStateForTests(): void { tasks.length = 0; + nextRegistrationOrder = 0; isShuttingDown = false; httpServer = null; } +async function runShutdownTasks(phase: ShutdownPhase): Promise { + const orderedTasks = tasks + .filter((task) => task.phase === phase) + .sort( + (left, right) => + right.priority - left.priority || left.registrationOrder - right.registrationOrder, + ); + + for (const task of orderedTasks) { + try { + logger.info(`Running ${phase} shutdown task: ${task.name}`); + await task.fn(); + } catch (err) { + logger.error(`Shutdown task "${task.name}" failed:`, err); + } + } +} + async function shutdown(signal: NodeJS.Signals): Promise { if (isShuttingDown) { return; @@ -68,21 +103,14 @@ async function shutdown(signal: NodeJS.Signals): Promise { let exitCode = 0; - try { - await closeHttpServer(); - } catch (err) { + const serverClosePromise = closeHttpServer().catch((err) => { logger.error('Error closing HTTP server during graceful shutdown:', err); exitCode = 1; - } + }); - for (const task of tasks) { - try { - logger.info(`Running shutdown task: ${task.name}`); - await task.fn(); - } catch (err) { - logger.error(`Shutdown task "${task.name}" failed:`, err); - } - } + await runShutdownTasks('pre-drain'); + await serverClosePromise; + await runShutdownTasks('post-drain'); clearTimeout(forceExit); logger.info('Graceful shutdown complete, exiting'); diff --git a/packages/api/src/stream/ApprovalLifecycle.ts b/packages/api/src/stream/ApprovalLifecycle.ts index e117c7bee2..71ce0953e6 100644 --- a/packages/api/src/stream/ApprovalLifecycle.ts +++ b/packages/api/src/stream/ApprovalLifecycle.ts @@ -3,6 +3,12 @@ import type { Agents } from 'librechat-data-provider'; import type { IJobStore } from '~/stream/interfaces/IJobStore'; import { isPendingActionExpired, isPendingActionStale } from '~/stream/interfaces/IJobStore'; +export interface ApprovalLifecycleCallbacks { + onPaused?: (streamId: string, createdAt: number) => void; + onResumed?: (streamId: string, createdAt: number) => void; + onExpired?: (streamId: string, createdAt: number) => void; +} + /** * The guarded lifecycle of a run paused for human review (`requires_action`). * @@ -26,7 +32,10 @@ import { isPendingActionExpired, isPendingActionStale } from '~/stream/interface * ``` */ export class ApprovalLifecycle { - constructor(private readonly store: IJobStore) {} + constructor( + private readonly store: IJobStore, + private readonly callbacks: ApprovalLifecycleCallbacks = {}, + ) {} /** * `running → requires_action`, attaching the pending review record. @@ -34,13 +43,19 @@ export class ApprovalLifecycle { * so a late interrupt is dropped rather than pausing a dead job. */ async pause(streamId: string, pendingAction: Agents.PendingAction): Promise { + const job = await this.store.getJob(streamId); + if (!job || job.status !== 'running') { + return false; + } const ok = await this.store.transitionStatus(streamId, { from: 'running', to: 'requires_action', // pendingActionId is the flat mirror the atomic resolve/expire guard on. patch: { pendingAction, pendingActionId: pendingAction.actionId }, + expectCreatedAt: job.createdAt, }); if (ok) { + this.callbacks.onPaused?.(streamId, job.createdAt); logger.debug( `[ApprovalLifecycle] paused for review: ${streamId} action=${pendingAction.actionId}`, ); @@ -79,22 +94,25 @@ export class ApprovalLifecycle { */ async resolve(streamId: string, expectedActionId?: string): Promise { const job = await this.store.getJob(streamId); - if (job?.status === 'requires_action' && !job.pendingAction) { + if (!job || job.status !== 'requires_action') { + return false; + } + if (!job.pendingAction) { // The prompt was lost (e.g. a malformed record dropped on deserialize). // It can't be reviewed, so finalize the job instead of driving a resumed // run with no reviewed interrupt payload — consistent with how the active // listing and cleanup treat a stale pending action. - await this.expire(streamId); + await this.expire(streamId, undefined, job.createdAt); return false; } - if (job?.status === 'requires_action' && job.pendingAction && isPendingActionExpired(job)) { + if (isPendingActionExpired(job)) { // Target the exact record observed as expired. If the caller didn't pin an // actionId, fall back to the one just read — otherwise a concurrent // resume + re-pause for a new action could let this expire abort it. - await this.expire(streamId, expectedActionId ?? job.pendingAction.actionId); + await this.expire(streamId, expectedActionId ?? job.pendingAction.actionId, job.createdAt); return false; } - return this.store.transitionStatus(streamId, { + const resumed = await this.store.transitionStatus(streamId, { from: 'requires_action', to: 'running', clear: ['pendingAction', 'pendingActionId'], @@ -102,7 +120,12 @@ export class ApprovalLifecycle { // immediately after resuming (cleanup keys off lastActiveAt). patch: { lastActiveAt: Date.now() }, expectActionId: expectedActionId, + expectCreatedAt: job.createdAt, }); + if (resumed) { + this.callbacks.onResumed?.(streamId, job.createdAt); + } + return resumed; } /** @@ -111,7 +134,32 @@ export class ApprovalLifecycle { * transition. Returns `true` to the single caller that expired it. Honors * `expectedActionId` for the same stale-decision protection as `resolve`. */ - async expire(streamId: string, expectedActionId?: string): Promise { + async expire( + streamId: string, + expectedActionId?: string, + expectedCreatedAt?: number, + ): Promise { + return (await this.expireWithIdentity(streamId, expectedActionId, expectedCreatedAt)) != null; + } + + /** + * Expires the observed approval and returns the winning job identity. Callers + * that need to notify runtime-local subscribers can use the identity to avoid + * delivering the predecessor's terminal event to a replacement generation. + */ + async expireWithIdentity( + streamId: string, + expectedActionId?: string, + expectedCreatedAt?: number, + ): Promise { + let createdAt = expectedCreatedAt; + if (createdAt == null) { + const job = await this.store.getJob(streamId); + if (!job || job.status !== 'requires_action') { + return null; + } + createdAt = job.createdAt; + } const ok = await this.store.transitionStatus(streamId, { from: 'requires_action', to: 'aborted', @@ -120,10 +168,12 @@ export class ApprovalLifecycle { // it an expired approval lingers in the in-memory map indefinitely. patch: { error: 'Approval expired before a decision was made', completedAt: Date.now() }, expectActionId: expectedActionId, + expectCreatedAt: createdAt, }); if (ok) { + this.callbacks.onExpired?.(streamId, createdAt); logger.debug(`[ApprovalLifecycle] expired pending review: ${streamId}`); } - return ok; + return ok ? createdAt : null; } } diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index c31472b320..272b9d77a0 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -11,7 +11,6 @@ import type { TMessageContentParts, TContextUsageEvent, TTokenUsageEvent, - TPendingSteer, Agents, } from 'librechat-data-provider'; import type { StandardGraph } from '@librechat/agents'; @@ -23,7 +22,8 @@ import type { IJobStore, IdempotencyClaimResult, } from './interfaces/IJobStore'; -import type { SteerOwner, SteerContentView } from './SteeringLifecycle'; +import type { AgentStartupTelemetry } from '~/agents/startup'; +import type { SteerContentView } from './SteeringLifecycle'; import type { GenerationJobStore } from '~/app/metrics'; import type * as t from '~/types'; import { @@ -40,9 +40,11 @@ import { import { isPendingActionStale, isPendingActionExpired } from './interfaces/IJobStore'; import { InMemoryEventTransport } from './implementations/InMemoryEventTransport'; import { InMemoryJobStore } from './implementations/InMemoryJobStore'; +import { emitChunkWithReceipt } from './internal/chunkPublication'; import { filterPersistableAbortContent } from './abortContent'; import { toClientPendingAction } from '~/agents/hitl/policy'; import { ApprovalLifecycle } from './ApprovalLifecycle'; +import { sanitizeJobMetadata } from './metadata'; /** Terminal error surfaced to a client still attached when its approval window lapses. */ const APPROVAL_EXPIRED_ERROR = 'Approval expired before a decision was made'; @@ -54,6 +56,9 @@ const REAPED_JOB_ERROR = 'Generation timed out'; * so a late retry still dedups for the whole generation window. */ const IDEMPOTENCY_TTL_SECONDS = 1200; const OAUTH_TOOL_CALL_PREFIX = `oauth${Constants.mcp_delimiter}`; +const SHUTDOWN_SUBSCRIBER_ERROR = 'Server is shutting down'; +const SHUTDOWN_JOB_ERROR = 'Generation interrupted because its server shut down'; +const SHUTTING_DOWN_ERROR = 'Generation job manager is shutting down'; function getToolCallName(toolCall: unknown): unknown { return toolCall != null && typeof toolCall === 'object' && 'name' in toolCall @@ -141,6 +146,11 @@ export interface GenerationJobManagerOptions { cleanupOnComplete?: boolean; } +export interface CreateGenerationJobOptions { + startupTelemetry?: AgentStartupTelemetry; + initialMetadata?: Partial; +} + /** * Runtime state for active jobs - not serializable, kept in-memory per instance. * Contains AbortController, ready promise, and other non-serializable state. @@ -162,20 +172,56 @@ export interface GenerationJobManagerOptions { * when the real client connects, which would prevent readyPromise from resolving. */ interface RuntimeJobState { + createdAt: number; abortController: AbortController; + /** Removes this generation's cross-replica abort listener without touching a replacement. */ + abortUnsubscribe?: () => void; readyPromise: Promise; resolveReady: () => void; + startupTelemetry?: AgentStartupTelemetry; finalEvent?: t.ServerSentEvent; errorEvent?: string; - /** Approval-expired host cleanup already ran for this runtime (relay path is swept repeatedly). */ - approvalCleanupRan?: boolean; + /** Local, runtime-scoped terminal handlers. Avoids broadcasting predecessor errors to a + * replacement generation that reuses the same durable stream ID. */ + localErrorHandlers: Set; + /** Prevents a repeated approval-expiry sweep from republishing the same terminal event. */ + approvalExpiryPublished?: boolean; syncSent: boolean; earlyEventBuffer: t.ServerSentEvent[]; earlyEventSequencePromises: Array>; + /** Initial subscribers eligible to receive the local pre-attachment replay. */ + earlyReplayHandlers: Set; + /** Per-resume capture handlers that bridge an in-memory snapshot to transport attachment. */ + resumeCaptureHandlers: Set<(event: t.ServerSentEvent, sequence: number) => void>; + /** Monotonic local emission sequence used to establish an exact resume snapshot frontier. */ + emissionSequence: number; + /** Emissions that started before an in-memory resume snapshot and must become snapshot-visible + * before the graph/job state is read. The event identity also suppresses their later publish. */ + inFlightSnapshotEmissions: Map< + number, + { event: t.ServerSentEvent; snapshotReady: Promise } + >; + /** Prevents later events from overtaking the initial `created` metadata write and publish. */ + createdEventPublication?: Promise; hasSubscriber: boolean; - allSubscribersLeftHandlers?: Array<(...args: unknown[]) => void>; + /** Advances whenever every local SSE subscriber for one attachment generation leaves. */ + attachmentGeneration: number; + /** Attachment generation whose partial-response disconnect cleanup was most recently started. */ + lastSubscriberCleanupGeneration?: number; + allSubscribersLeftHandlers?: Array<(...args: unknown[]) => void | Promise>; } +interface PreparedSubscription { + runtime: RuntimeJobState; + jobData: SerializableJobData | null; + deferDeliveryUntilActivated: boolean; +} + +type DeferredDelivery = + | { type: 'chunk'; event: t.ServerSentEvent } + | { type: 'done'; event: t.ServerSentEvent } + | { type: 'error'; error: string }; + /** * Manages generation jobs for resumable LLM streams. * @@ -210,8 +256,8 @@ class GenerationJobManagerClass { /** Runtime state - always in-memory, not serializable */ private runtimeState = new Map(); - /** Jobs actively generating in this process. */ - private runningJobs = new Set(); + /** Jobs actively owned by this process, pinned to their durable creation epoch. */ + private ownedJobs = new Map(); /** Serializes replay-event read/modify/write updates per stream. */ private replayEventWriteQueues = new Map>(); @@ -219,27 +265,24 @@ class GenerationJobManagerClass { /** Serializes token-usage read/modify/write updates per stream. */ private tokenUsageWriteQueues = new Map>(); + /** Partial-response and disconnect-state writes still draining during shutdown. */ + private subscriberCleanupPromises = new Set>(); + private cleanupInterval: NodeJS.Timeout | null = null; + /** Rejects new jobs once graceful shutdown has started. */ + private shuttingDown = false; + /** Whether we're using Redis stores */ private _isRedis = false; /** Whether to cleanup event transport immediately on job completion */ private _cleanupOnComplete = true; - /** - * Host cleanup fired after an approval EXPIRES (periodic sweeper or a stale submit) — - * e.g. prune the paused run's durable checkpoint eagerly instead of letting it sit - * until its store TTL. Best-effort: failures are logged, never break the expiry. - */ - private _onApprovalExpired: - | ((streamId: string, job?: SerializableJobData | null) => void | Promise) - | null = null; - constructor(options?: GenerationJobManagerOptions) { this.jobStore = options?.jobStore ?? new InMemoryJobStore({ ttlAfterComplete: 0, maxJobs: 1000 }); - this._approvals = new ApprovalLifecycle(this.jobStore); + this._approvals = this.createApprovalLifecycle(this.jobStore); this._steering = new SteeringLifecycle(this.jobStore); this.eventTransport = options?.eventTransport ?? new InMemoryEventTransport(); this._cleanupOnComplete = options?.cleanupOnComplete ?? true; @@ -292,18 +335,49 @@ class GenerationJobManagerClass { logger.warn( '[GenerationJobManager] Reconfiguring after initialization - destroying existing services', ); - this.destroy(); + clearInterval(this.cleanupInterval); + this.cleanupInterval = null; + + const replacedJobStore = this.jobStore; + const replacedEventTransport = this.eventTransport; + const pendingSubscriberCleanups = [...this.subscriberCleanupPromises]; + for (const runtime of this.runtimeState.values()) { + runtime.startupTelemetry?.end('aborted'); + runtime.startupTelemetry = undefined; + this.releaseAbortSubscription(runtime); + runtime.abortController.abort(); + } + + // Detach the old transport synchronously so it cannot deliver into a replacement runtime. + // Its store can finish draining already-started disconnect persistence in the background. + try { + replacedEventTransport.destroy(); + } catch (err) { + logger.error('[GenerationJobManager] Failed to destroy replaced event transport:', err); + } + void Promise.allSettled(pendingSubscriberCleanups) + .then(() => replacedJobStore.destroy()) + .catch((err) => { + logger.error('[GenerationJobManager] Failed to destroy replaced job store:', err); + }); + + this.runtimeState.clear(); + this.subscriberCleanupPromises.clear(); + this.runStepBuffers?.clear(); + this.replayEventWriteQueues.clear(); + this.tokenUsageWriteQueues.clear(); } - this.runningJobs.clear(); + this.ownedJobs.clear(); setGenerationJobsInFlight(previousStore, 0); this.jobStore = services.jobStore; - this._approvals = new ApprovalLifecycle(this.jobStore); + this._approvals = this.createApprovalLifecycle(this.jobStore); this._steering = new SteeringLifecycle(this.jobStore); this.eventTransport = services.eventTransport; this._isRedis = services.isRedis ?? false; this._cleanupOnComplete = services.cleanupOnComplete ?? true; + this.shuttingDown = false; this.syncRunningJobMetrics(); logger.info( @@ -311,19 +385,6 @@ class GenerationJobManagerClass { ); } - /** - * Register a host callback fired after an approval EXPIRES — from the periodic sweeper or - * a stale submit — e.g. to prune the paused run's durable checkpoint eagerly instead of - * waiting out its TTL. Unlike {@link configure} this never resets services, so it is safe - * to call from any startup path (including ones that run on constructor defaults). The - * `streamId` argument equals the LangGraph `thread_id` (LibreChat's conversationId). - */ - setApprovalExpiredHandler( - handler: ((streamId: string, job?: SerializableJobData | null) => void | Promise) | null, - ): void { - this._onApprovalExpired = handler; - } - /** * Check if using Redis stores. */ @@ -336,7 +397,228 @@ class GenerationJobManagerClass { } private syncRunningJobMetrics(store: GenerationJobStore = this.storeLabel): void { - setGenerationJobsInFlight(store, this.runningJobs.size); + setGenerationJobsInFlight(store, this.ownedJobs.size); + } + + private createApprovalLifecycle(store: IJobStore): ApprovalLifecycle { + return new ApprovalLifecycle(store, { + onPaused: (streamId, createdAt) => this.releaseJobOwnership(streamId, createdAt), + onResumed: (streamId, createdAt) => this.acquireResumedJobOwnership(streamId, createdAt), + onExpired: (streamId, createdAt) => this.releaseJobOwnership(streamId, createdAt), + }); + } + + private acquireJobOwnership(streamId: string, createdAt: number): void { + this.ownedJobs.set(streamId, createdAt); + this.syncRunningJobMetrics(); + } + + private acquireResumedJobOwnership(streamId: string, createdAt: number): void { + const ownedCreatedAt = this.ownedJobs.get(streamId); + if (ownedCreatedAt != null && ownedCreatedAt !== createdAt) { + return; + } + this.acquireJobOwnership(streamId, createdAt); + } + + private releaseJobOwnership(streamId: string, expectedCreatedAt?: number): boolean { + if (expectedCreatedAt != null && this.ownedJobs.get(streamId) !== expectedCreatedAt) { + return false; + } + const released = this.ownedJobs.delete(streamId); + if (released) { + this.syncRunningJobMetrics(); + } + return released; + } + + private releaseAbortSubscription(runtime: RuntimeJobState): void { + const unsubscribe = runtime.abortUnsubscribe; + runtime.abortUnsubscribe = undefined; + if (!unsubscribe) { + return; + } + + try { + unsubscribe(); + } catch (err) { + logger.error('[GenerationJobManager] Failed to release abort subscription:', err); + } + } + + private reconcileInactiveGeneration( + streamId: string, + createdAt: number, + currentJob: SerializableJobData | null, + observedRuntime?: RuntimeJobState, + ): void { + if (currentJob?.createdAt === createdAt) { + if (currentJob.status === 'running') { + return; + } + if (currentJob.status === 'requires_action') { + this.releaseJobOwnership(streamId, createdAt); + return; + } + } + + if ( + observedRuntime?.createdAt === createdAt && + this.runtimeState.get(streamId) === observedRuntime + ) { + this.releaseAbortSubscription(observedRuntime); + observedRuntime.abortController.abort(); + } + this.releaseJobOwnership(streamId, createdAt); + } + + private async reconcileLostTerminalTransition( + streamId: string, + createdAt: number, + observedRuntime?: RuntimeJobState, + ): Promise { + const currentJob = await this.jobStore.getJob(streamId); + this.reconcileInactiveGeneration(streamId, createdAt, currentJob, observedRuntime); + } + + private async registerAbortSubscription( + streamId: string, + runtime: RuntimeJobState, + ): Promise { + if (!this.eventTransport.onAbort) { + return; + } + + const unsubscribe = await this.eventTransport.onAbort(streamId, (generationId) => { + const currentRuntime = this.runtimeState.get(streamId); + if ( + currentRuntime !== runtime || + (generationId != null && currentRuntime.createdAt !== generationId) || + currentRuntime.abortController.signal.aborted + ) { + return; + } + + logger.debug(`[GenerationJobManager] Received cross-replica abort for ${streamId}`); + currentRuntime.abortController.abort(); + this.releaseAbortSubscription(currentRuntime); + }); + + if (typeof unsubscribe === 'function') { + runtime.abortUnsubscribe = unsubscribe; + } + if (this.runtimeState.get(streamId) !== runtime || runtime.abortController.signal.aborted) { + this.releaseAbortSubscription(runtime); + } + } + + private rejectSubscriptionDuringShutdown( + subscriptionType: 'initial' | 'resume', + onError?: t.ErrorHandler, + ): boolean { + if (!this.shuttingDown) { + return false; + } + + recordGenerationStreamSubscription(this.storeLabel, subscriptionType, 'error'); + onError?.(SHUTDOWN_SUBSCRIBER_ERROR); + return true; + } + + private detachSubscriptionDuringShutdown( + subscription: { unsubscribe: t.UnsubscribeFn } | null, + ): boolean { + if (!this.shuttingDown) { + return false; + } + + subscription?.unsubscribe(); + return true; + } + + private registerAllSubscribersLeft(streamId: string): void { + this.eventTransport.onAllSubscribersLeft(streamId, () => { + const runtime = this.runtimeState.get(streamId); + if (!runtime) { + return; + } + + runtime.syncSent = false; + runtime.hasSubscriber = false; + runtime.attachmentGeneration++; + runtime.lastSubscriberCleanupGeneration = runtime.attachmentGeneration; + + // Terminal delivery closes the SSE subscription too, but it is not a user + // disconnect. Running partial-response handlers here can overwrite the + // already-saved final response as unfinished. + if (runtime.finalEvent || runtime.errorEvent) { + return; + } + + const cleanup = this.persistSubscriberCleanup(streamId, runtime); + this.subscriberCleanupPromises.add(cleanup); + void cleanup.then( + () => this.subscriberCleanupPromises.delete(cleanup), + (err) => { + this.subscriberCleanupPromises.delete(cleanup); + logger.error(`[GenerationJobManager] Failed to clean up disconnected subscriber:`, err); + }, + ); + }); + } + + private async persistSubscriberCleanup( + streamId: string, + runtime: RuntimeJobState, + ): Promise { + const persistSyncState = this.jobStore + .updateJob(streamId, { syncSent: false }, runtime.createdAt) + .catch((err) => { + logger.error(`[GenerationJobManager] Failed to persist syncSent=false:`, err); + }); + const handlers = runtime.allSubscribersLeftHandlers + ? [...runtime.allSubscribersLeftHandlers] + : []; + + if (handlers.length === 0) { + await persistSyncState; + return; + } + + try { + const result = await this.jobStore.getContentParts(streamId, runtime.createdAt); + const parts = result?.content ?? []; + const handlerResults = await Promise.allSettled( + handlers.map((handler) => Promise.resolve().then(() => handler(parts))), + ); + for (const handlerResult of handlerResults) { + if (handlerResult.status === 'rejected') { + logger.error( + `[GenerationJobManager] Error in allSubscribersLeft handler:`, + handlerResult.reason, + ); + } + } + } catch (err) { + logger.error( + `[GenerationJobManager] Failed to get content parts for allSubscribersLeft handlers:`, + err, + ); + } + + await persistSyncState; + } + + private async drainSubscriberCleanups(): Promise { + const pending = [...this.subscriberCleanupPromises]; + if (pending.length === 0) { + return; + } + + await Promise.allSettled(pending); + if (this.subscriberCleanupPromises.size > 0) { + await this.drainSubscriberCleanups(); + } } /** @@ -351,12 +633,11 @@ class GenerationJobManagerClass { * * This sets up: * 1. Serializable job data in the job store - * 2. Runtime state including readyPromise (resolves when first SSE client connects) + * 2. Runtime state including the legacy, immediately-resolved readyPromise facade * 3. allSubscribersLeft callback for handling client disconnections * - * The readyPromise mechanism ensures generation doesn't start before the client - * is ready to receive events. The controller awaits this promise (with a short timeout) - * before starting LLM generation. + * Generation starts independently of SSE attachment. Early events are buffered locally and, + * in Redis mode, persisted/published for replay when the client subscribes. * * @param streamId - Unique identifier for this stream * @param userId - User who initiated the request @@ -367,10 +648,50 @@ class GenerationJobManagerClass { streamId: string, userId: string, conversationId?: string, + options: CreateGenerationJobOptions = {}, ): Promise { + if (this.shuttingDown) { + throw new Error(SHUTTING_DOWN_ERROR); + } + const tenantId = getTenantId(); const safeTenantId = tenantId && tenantId !== SYSTEM_TENANT_ID ? tenantId : undefined; - const jobData = await this.jobStore.createJob(streamId, userId, conversationId, safeTenantId); + const initialMetadata = sanitizeJobMetadata(options.initialMetadata ?? {}); + const jobData = await this.jobStore.createJob( + streamId, + userId, + conversationId, + safeTenantId, + initialMetadata, + ); + const currentRuntimeBeforeInstall = this.runtimeState.get(streamId); + const ownedCreatedAtBeforeInstall = this.ownedJobs.get(streamId); + if ( + (currentRuntimeBeforeInstall != null && + currentRuntimeBeforeInstall.createdAt > jobData.createdAt) || + (ownedCreatedAtBeforeInstall != null && ownedCreatedAtBeforeInstall > jobData.createdAt) + ) { + throw new Error('Generation job was replaced during initialization'); + } + if (this.shuttingDown) { + await this.jobStore.transitionStatus(streamId, { + from: 'running', + to: 'error', + patch: { completedAt: Date.now(), error: SHUTDOWN_JOB_ERROR }, + expectCreatedAt: jobData.createdAt, + }); + throw new Error(SHUTTING_DOWN_ERROR); + } + this.acquireJobOwnership(streamId, jobData.createdAt); + recordGenerationJob(this.storeLabel, 'created'); + + const replacedRuntime = this.runtimeState.get(streamId); + if (replacedRuntime) { + replacedRuntime.startupTelemetry?.end('replaced'); + replacedRuntime.startupTelemetry = undefined; + this.releaseAbortSubscription(replacedRuntime); + replacedRuntime.abortController.abort(); + } /** * Create runtime state with readyPromise. @@ -383,80 +704,68 @@ class GenerationJobManagerClass { * We resolve readyPromise immediately to eliminate startup latency. * The sync mechanism handles late-connecting clients. */ - let resolveReady: () => void; - const readyPromise = new Promise((resolve) => { - resolveReady = resolve; - }); + const readyPromise = Promise.resolve(); + const resolveReady = (): void => undefined; const runtime: RuntimeJobState = { + createdAt: jobData.createdAt, abortController: new AbortController(), readyPromise, - resolveReady: resolveReady!, + resolveReady, + startupTelemetry: options.startupTelemetry, syncSent: false, earlyEventBuffer: [], earlyEventSequencePromises: [], + earlyReplayHandlers: new Set(), + resumeCaptureHandlers: new Set(), + localErrorHandlers: new Set(), + emissionSequence: 0, + inFlightSnapshotEmissions: new Map(), hasSubscriber: false, + attachmentGeneration: 0, }; this.runtimeState.set(streamId, runtime); - this.runningJobs.add(streamId); - this.syncRunningJobMetrics(); - recordGenerationJob(this.storeLabel, 'created'); - // Resolve immediately - early event buffer handles late subscribers - resolveReady!(); + try { + this.registerAllSubscribersLeft(streamId); - /** - * Set up all-subscribers-left callback. - * When all SSE clients disconnect, this: - * 1. Resets syncSent so reconnecting clients get sync event (persisted to Redis) - * 2. Calls any registered allSubscribersLeft handlers (e.g., to save partial responses) - */ - this.eventTransport.onAllSubscribersLeft(streamId, () => { - const currentRuntime = this.runtimeState.get(streamId); - if (currentRuntime) { - currentRuntime.syncSent = false; - currentRuntime.hasSubscriber = false; - // Persist syncSent=false to Redis for cross-replica consistency - this.jobStore.updateJob(streamId, { syncSent: false }).catch((err) => { - logger.error(`[GenerationJobManager] Failed to persist syncSent=false:`, err); - }); - // Call registered handlers (from job.emitter.on('allSubscribersLeft', ...)) - if (currentRuntime.allSubscribersLeftHandlers) { - this.jobStore - .getContentParts(streamId) - .then((result) => { - const parts = result?.content ?? []; - for (const handler of currentRuntime.allSubscribersLeftHandlers ?? []) { - try { - handler(parts); - } catch (err) { - logger.error(`[GenerationJobManager] Error in allSubscribersLeft handler:`, err); - } - } - }) - .catch((err) => { - logger.error( - `[GenerationJobManager] Failed to get content parts for allSubscribersLeft handlers:`, - err, - ); - }); - } + await this.registerAbortSubscription(streamId, runtime); + if (this.runtimeState.get(streamId) !== runtime) { + throw new Error('Generation job was replaced during initialization'); } - }); - - /** - * Set up cross-replica abort listener (Redis mode only). - * When abort is triggered on ANY replica, this replica receives the signal - * and aborts its local AbortController (if it's the one running generation). - */ - if (this.eventTransport.onAbort) { - this.eventTransport.onAbort(streamId, () => { - const currentRuntime = this.runtimeState.get(streamId); - if (currentRuntime && !currentRuntime.abortController.signal.aborted) { - logger.debug(`[GenerationJobManager] Received cross-replica abort for ${streamId}`); - currentRuntime.abortController.abort(); - } + const confirmedJobData = await this.jobStore.getJob(streamId); + if ( + this.runtimeState.get(streamId) !== runtime || + !confirmedJobData || + confirmedJobData.createdAt !== runtime.createdAt || + confirmedJobData.status !== 'running' + ) { + throw new Error('Generation job was replaced during initialization'); + } + if (this.shuttingDown) { + throw new Error(SHUTTING_DOWN_ERROR); + } + } catch (error) { + // The durable job already exists, but the caller has not received its + // generation identity yet. Finalize that exact epoch here so a controller + // catch never needs to issue an unsafe unscoped terminal mutation. + let message = SHUTDOWN_JOB_ERROR; + if (!this.shuttingDown) { + message = error instanceof Error ? error.message : String(error); + } + await this.completeJob(streamId, message, jobData.createdAt).catch((finalizeError) => { + logger.error( + `[GenerationJobManager] Failed to finalize partially initialized job ${streamId}:`, + finalizeError, + ); }); + if (this.runtimeState.get(streamId) === runtime) { + this.releaseAbortSubscription(runtime); + runtime.abortController.abort(); + this.runtimeState.delete(streamId); + this.releaseJobOwnership(streamId, runtime.createdAt); + } + throw error; } logger.debug(`[GenerationJobManager] Created job: ${streamId}`); @@ -500,7 +809,7 @@ class GenerationJobManagerClass { * incrementing subscriber count (see class JSDoc above). */ const emitterProxy = { - on: (event: string, handler: (...args: unknown[]) => void) => { + on: (event: string, handler: (...args: unknown[]) => void | Promise) => { if (event === 'allSubscribersLeft') { // Store handler for internal callback - don't use subscribe() to avoid counting as a subscriber if (!runtime.allSubscribersLeftHandlers) { @@ -576,30 +885,37 @@ class GenerationJobManagerClass { * @param streamId - The stream identifier * @returns Runtime state or null if job doesn't exist anywhere */ - private async getOrCreateRuntimeState(streamId: string): Promise { - const existingRuntime = this.runtimeState.get(streamId); - if (existingRuntime) { - return existingRuntime; - } - - // Job doesn't exist locally - check Redis - const jobData = await this.jobStore.getJob(streamId); + private async getOrCreateRuntimeState( + streamId: string, + knownJobData?: SerializableJobData | null, + ): Promise { + const jobData = + knownJobData === undefined ? await this.jobStore.getJob(streamId) : knownJobData; if (!jobData) { return null; } - // Cross-replica scenario: job exists in Redis but not locally - // Create minimal runtime state for handling reconnection/subscription + const concurrentRuntime = this.runtimeState.get(streamId); + if (concurrentRuntime?.createdAt === jobData.createdAt) { + this.reconcileInactiveGeneration(streamId, jobData.createdAt, jobData, concurrentRuntime); + return concurrentRuntime; + } + if (concurrentRuntime && concurrentRuntime.createdAt > jobData.createdAt) { + return concurrentRuntime; + } + if (concurrentRuntime) { + concurrentRuntime.startupTelemetry?.end('replaced'); + concurrentRuntime.startupTelemetry = undefined; + this.releaseAbortSubscription(concurrentRuntime); + concurrentRuntime.abortController.abort(); + } + + // Cross-replica scenario: create (or replace) the minimal runtime state + // from the durable generation currently owning this stream ID. logger.debug(`[GenerationJobManager] Creating cross-replica runtime for ${streamId}`); - let resolveReady: () => void; - const readyPromise = new Promise((resolve) => { - resolveReady = resolve; - }); - - // For jobs created on other replicas, readyPromise should be pre-resolved - // since generation has already started - resolveReady!(); + const readyPromise = Promise.resolve(); + const resolveReady = (): void => undefined; // Parse finalEvent from Redis if available let finalEvent: t.ServerSentEvent | undefined; @@ -612,67 +928,63 @@ class GenerationJobManagerClass { } const runtime: RuntimeJobState = { + createdAt: jobData.createdAt, abortController: new AbortController(), readyPromise, - resolveReady: resolveReady!, + resolveReady, syncSent: jobData.syncSent ?? false, earlyEventBuffer: [], earlyEventSequencePromises: [], + earlyReplayHandlers: new Set(), + resumeCaptureHandlers: new Set(), + localErrorHandlers: new Set(), + emissionSequence: 0, + inFlightSnapshotEmissions: new Map(), hasSubscriber: false, + attachmentGeneration: 0, finalEvent, errorEvent: jobData.error, }; this.runtimeState.set(streamId, runtime); - // Set up all-subscribers-left callback for this replica - this.eventTransport.onAllSubscribersLeft(streamId, () => { - const currentRuntime = this.runtimeState.get(streamId); - if (currentRuntime) { - currentRuntime.syncSent = false; - currentRuntime.hasSubscriber = false; - // Persist syncSent=false to Redis - this.jobStore.updateJob(streamId, { syncSent: false }).catch((err) => { - logger.error(`[GenerationJobManager] Failed to persist syncSent=false:`, err); - }); - // Call registered handlers - if (currentRuntime.allSubscribersLeftHandlers) { - this.jobStore - .getContentParts(streamId) - .then((result) => { - const parts = result?.content ?? []; - for (const handler of currentRuntime.allSubscribersLeftHandlers ?? []) { - try { - handler(parts); - } catch (err) { - logger.error(`[GenerationJobManager] Error in allSubscribersLeft handler:`, err); - } - } - }) - .catch((err) => { - logger.error( - `[GenerationJobManager] Failed to get content parts for allSubscribersLeft handlers:`, - err, - ); - }); - } - } - }); + this.registerAllSubscribersLeft(streamId); - // Set up cross-replica abort listener (Redis mode only) - // This ensures lazily-initialized jobs can receive abort signals - if (this.eventTransport.onAbort) { - this.eventTransport.onAbort(streamId, () => { - const currentRuntime = this.runtimeState.get(streamId); - if (currentRuntime && !currentRuntime.abortController.signal.aborted) { - logger.debug( - `[GenerationJobManager] Received cross-replica abort for lazily-init job ${streamId}`, - ); - currentRuntime.abortController.abort(); - } - }); + if (jobData.status === 'running' || jobData.status === 'requires_action') { + await this.registerAbortSubscription(streamId, runtime); } + const runtimeAfterAbortRegistration = this.runtimeState.get(streamId); + if (runtimeAfterAbortRegistration !== runtime) { + return runtimeAfterAbortRegistration ?? null; + } + + /** + * `onAbort` may require an asynchronous Redis subscription. A replacement can + * become durable on another replica while that subscription is activating, + * without changing this process's runtime map. Re-read the owner before + * exposing the facade, and reconcile again if the generation moved. + */ + const confirmedJobData = await this.jobStore.getJob(streamId); + if (this.runtimeState.get(streamId) !== runtime) { + return this.runtimeState.get(streamId) ?? null; + } + if (!confirmedJobData) { + this.releaseAbortSubscription(runtime); + runtime.abortController.abort(); + this.runtimeState.delete(streamId); + return null; + } + if (confirmedJobData.createdAt !== runtime.createdAt) { + return this.getOrCreateRuntimeState(streamId, confirmedJobData); + } + this.reconcileInactiveGeneration( + streamId, + confirmedJobData.createdAt, + confirmedJobData, + runtime, + ); + return runtime; } @@ -685,8 +997,12 @@ class GenerationJobManagerClass { return undefined; } - const runtime = await this.getOrCreateRuntimeState(streamId); - if (!runtime) { + const runtime = await this.getOrCreateRuntimeState(streamId, jobData); + if ( + !runtime || + this.runtimeState.get(streamId) !== runtime || + runtime.createdAt !== jobData.createdAt + ) { return undefined; } @@ -744,84 +1060,138 @@ class GenerationJobManagerClass { * fully transmitted. It will be cleaned up when subscribers disconnect or * by the periodic cleanup job. */ - async completeJob(streamId: string, error?: string): Promise { - const runtime = this.runtimeState.get(streamId); - - // Abort the controller to signal all pending operations (e.g., OAuth flow monitors) - // that the job is done and they should clean up - if (runtime) { - runtime.abortController.abort(); + async completeJob(streamId: string, error?: string, expectedCreatedAt?: number): Promise { + const observedRuntime = this.runtimeState.get(streamId); + const targetCreatedAt = expectedCreatedAt ?? observedRuntime?.createdAt; + const jobData = await this.jobStore.getJob(streamId); + if (!jobData || (targetCreatedAt != null && jobData.createdAt !== targetCreatedAt)) { + if (targetCreatedAt != null) { + this.reconcileInactiveGeneration(streamId, targetCreatedAt, jobData, observedRuntime); + } + logger.debug( + `[GenerationJobManager] Skipping stale completion for replaced job: ${streamId}`, + ); + return; + } + if (jobData.status !== 'running') { + this.reconcileInactiveGeneration(streamId, jobData.createdAt, jobData, observedRuntime); + logger.debug( + `[GenerationJobManager] Skipping completion for non-running job ${streamId}: ${jobData.status}`, + ); + return; } - // Clear content state and run step buffer (Redis only) - this.jobStore.clearContentState(streamId); - this.runStepBuffers?.delete(streamId); - this.replayEventWriteQueues.delete(streamId); - this.tokenUsageWriteQueues.delete(streamId); + const createdAt = jobData.createdAt; + const runtime = observedRuntime?.createdAt === createdAt ? observedRuntime : undefined; + // Backstop for direct terminal callers (init failures, unhandled errors) - // that never ran the controllers' close-and-park: close the queue, then - // park any 202-accepted leftovers for /chat/status claim-on-read instead - // of silently clearing them. Paths that already drained find an empty - // queue and no-op; the createdAt guard (re-checked inside the store's - // atomic drain) keeps a stale completion off a replacement job's queue. - // Runs BEFORE the terminal status write — the Redis terminal cleanup DELs - // the queue key. + // that never ran the controllers' close-and-park. Every mutation is pinned + // to this completion's immutable generation identity, so a predecessor can + // never drain or terminalize a replacement that reuses the stream ID. try { - const jobData = await this.jobStore.getJob(streamId); - if (jobData) { - const leftovers = ( - await this.jobStore.closeAndDrainSteers(streamId, jobData.createdAt) - ).map(toPendingSteer); - await this._steering.park(streamId, leftovers, { + const leftovers = (await this.jobStore.closeAndDrainSteers(streamId, createdAt)).map( + toPendingSteer, + ); + await this._steering.park( + streamId, + leftovers, + { userId: jobData.userId, tenantId: jobData.tenantId, - }); - } + }, + createdAt, + ); } catch (err) { logger.warn(`[GenerationJobManager] Failed to park leftover steers for ${streamId}:`, err); } - // For error jobs, DON'T delete immediately - keep around so late-connecting - // clients can receive the error. This handles the race condition where error - // occurs before client connects to SSE stream. - // - // Cleanup strategy: Error jobs are cleaned up by periodic cleanup (every 60s) - // via jobStore.cleanup() which checks for jobs with status 'error' and - // completedAt set. The TTL is configurable via jobStore options (default: 0, - // meaning cleanup on next interval). This gives clients ~60s to connect and - // receive the error before the job is removed. + // Error jobs stay durable long enough for late subscribers to receive the + // stored error. The status CAS also parks/cleans same-slot Redis state. if (error) { - await this.jobStore.updateJob(streamId, { - status: 'error', - completedAt: Date.now(), - error, + const finalized = await this.jobStore.transitionStatus(streamId, { + from: 'running', + to: 'error', + expectCreatedAt: createdAt, + patch: { completedAt: Date.now(), error }, }); - this.runningJobs.delete(streamId); - this.syncRunningJobMetrics(); + if (!finalized) { + await this.reconcileLostTerminalTransition(streamId, createdAt, runtime); + return; + } + if (runtime && this.runtimeState.get(streamId) === runtime) { + runtime.errorEvent = error; + } + try { + await this.eventTransport.emitError(streamId, error, createdAt); + } catch (publishError) { + logger.error( + `[GenerationJobManager] Failed to publish terminal error for ${streamId}:`, + publishError, + ); + if (runtime && this.runtimeState.get(streamId) === runtime) { + for (const notify of [...runtime.localErrorHandlers]) { + try { + notify(error); + } catch (notifyError) { + logger.error( + `[GenerationJobManager] Failed to notify terminal error for ${streamId}:`, + notifyError, + ); + } + } + } + } + if (runtime && this.runtimeState.get(streamId) === runtime) { + this.releaseAbortSubscription(runtime); + runtime.abortController.abort(); + runtime.startupTelemetry?.end('error', new Error(error)); + runtime.startupTelemetry = undefined; + this.jobStore.clearContentState(streamId, createdAt); + this.runStepBuffers?.delete(streamId); + this.replayEventWriteQueues.delete(streamId); + this.tokenUsageWriteQueues.delete(streamId); + } + this.releaseJobOwnership(streamId, createdAt); recordGenerationJob(this.storeLabel, 'error'); - // Keep runtime state so subscribe() can access errorEvent logger.debug( `[GenerationJobManager] Job completed with error (keeping for late subscribers): ${streamId}`, ); return; } - // Immediate cleanup if configured (default: true) - only for successful completions + // Successful completion deletes immediately by default. Both branches are + // guarded at the store boundary, closing the read→delete replacement race. + const completed = await this.jobStore.transitionStatus(streamId, { + from: 'running', + to: 'complete', + expectCreatedAt: createdAt, + patch: { completedAt: Date.now() }, + }); + if (!completed) { + await this.reconcileLostTerminalTransition(streamId, createdAt, runtime); + return; + } if (this._cleanupOnComplete) { - this.runtimeState.delete(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 - await this.jobStore.updateJob(streamId, { - status: 'complete', - completedAt: Date.now(), - }); + // A same-stream replacement created after the completion CAS makes this + // a safe no-op rather than deleting the replacement generation. + await this.jobStore.deleteJob(streamId, createdAt); } - this.runningJobs.delete(streamId); - this.syncRunningJobMetrics(); + if (runtime && this.runtimeState.get(streamId) === runtime) { + this.releaseAbortSubscription(runtime); + runtime.abortController.abort(); + runtime.startupTelemetry?.end('completed_without_delta'); + runtime.startupTelemetry = undefined; + this.jobStore.clearContentState(streamId, createdAt); + this.runStepBuffers?.delete(streamId); + this.replayEventWriteQueues.delete(streamId); + this.tokenUsageWriteQueues.delete(streamId); + if (this._cleanupOnComplete) { + this.runtimeState.delete(streamId); + } + } + + this.releaseJobOwnership(streamId, createdAt); recordGenerationJob(this.storeLabel, 'completed'); logger.debug(`[GenerationJobManager] Job completed: ${streamId}`); } @@ -846,10 +1216,18 @@ class GenerationJobManagerClass { transformAbortContent?: (content: TMessageContentParts[]) => TMessageContentParts[]; }, ): Promise { + const observedRuntime = this.runtimeState.get(streamId); const jobData = await this.jobStore.getJob(streamId); - const runtime = this.runtimeState.get(streamId); if (!jobData) { + if (observedRuntime) { + this.reconcileInactiveGeneration( + streamId, + observedRuntime.createdAt, + jobData, + observedRuntime, + ); + } logger.warn(`[GenerationJobManager] Cannot abort - job not found: ${streamId}`); recordGenerationJob(this.storeLabel, 'abort_failed'); return { @@ -862,16 +1240,24 @@ class GenerationJobManagerClass { }; } - // Emit abort signal for cross-replica support (Redis mode) - // This ensures the generating replica receives the abort signal - if (this.eventTransport.emitAbort) { - this.eventTransport.emitAbort(streamId); + const abortableStatus = jobData.status; + if (abortableStatus !== 'running' && abortableStatus !== 'requires_action') { + this.reconcileInactiveGeneration(streamId, jobData.createdAt, jobData, observedRuntime); + logger.debug( + `[GenerationJobManager] Cannot abort terminal job ${streamId}: ${jobData.status}`, + ); + recordGenerationJob(this.storeLabel, 'abort_failed'); + return { + text: '', + content: [], + jobData, + success: false, + finalEvent: null, + collectedUsage: [], + }; } - // Also abort local controller if we have it (same-replica abort) - if (runtime) { - runtime.abortController.abort(); - } + const runtime = observedRuntime?.createdAt === jobData.createdAt ? observedRuntime : undefined; /** Steers that never reached an injection boundary — reported on the abort * final event (and the abort route's JSON) so the client can restore them @@ -888,13 +1274,18 @@ class GenerationJobManagerClass { ).map(toPendingSteer); // No-subscriber recovery: the abort response/final are transient, so park // the leftovers for /chat/status claim-on-read within the recovery TTL. - await this.steering.park(streamId, pendingSteers, { - userId: jobData.userId, - tenantId: jobData.tenantId, - }); + await this.steering.park( + streamId, + pendingSteers, + { + userId: jobData.userId, + tenantId: jobData.tenantId, + }, + jobData.createdAt, + ); /** Content before clearing state */ - const result = await this.jobStore.getContentParts(streamId); + const result = await this.jobStore.getContentParts(streamId, jobData.createdAt); const content = result?.content ?? []; let abortContent = filterPersistableAbortContent(content); if (options?.transformAbortContent) { @@ -905,7 +1296,7 @@ class GenerationJobManagerClass { const shouldPersistAbortContent = abortContent.length > 0; /** Collected usage for all models */ - const collectedUsage = this.jobStore.getCollectedUsage(streamId); + const collectedUsage = this.jobStore.getCollectedUsage(streamId, jobData.createdAt); /** Text from content parts for fallback token counting; the persisted * abort record keeps steered words (they reached the model context). */ @@ -956,31 +1347,69 @@ class GenerationJobManagerClass { ...(pendingSteers.length > 0 && { pendingSteers }), } satisfies t.FinalEvent as t.ServerSentEvent; + // Claim the terminal state before publishing anything client-visible. A + // natural completion or approval resolution racing this abort can win the + // CAS, in which case this stale abort must not signal or close that winner. + const finalized = await this.jobStore.transitionStatus(streamId, { + from: abortableStatus, + to: 'aborted', + expectCreatedAt: jobData.createdAt, + patch: { completedAt: Date.now() }, + }); + if (!finalized) { + await this.reconcileLostTerminalTransition(streamId, jobData.createdAt, runtime); + return { + success: false, + jobData, + content: abortContent, + finalEvent: null, + text, + collectedUsage, + ...(pendingSteers.length > 0 && { pendingSteers }), + }; + } + + // Signal only the generation whose terminal transition won above. The + // transport tag prevents a delayed predecessor abort from reaching a + // same-stream replacement on another replica. + if (this.eventTransport.emitAbort) { + this.eventTransport.emitAbort(streamId, jobData.createdAt); + } + if (runtime) { + this.releaseAbortSubscription(runtime); + } + runtime?.abortController.abort(); + if (runtime) { runtime.finalEvent = abortFinalEvent; } - await this.eventTransport.emitDone(streamId, abortFinalEvent); - this.jobStore.clearContentState(streamId); - this.runStepBuffers?.delete(streamId); - this.replayEventWriteQueues.delete(streamId); - this.tokenUsageWriteQueues.delete(streamId); - - // Immediate cleanup if configured (default: true) + if (runtime?.createdEventPublication) { + await runtime.createdEventPublication; + } + await this.eventTransport.emitDone(streamId, abortFinalEvent, jobData.createdAt); + if (runtime?.startupTelemetry) { + this.recordStartupEvent(runtime, abortFinalEvent); + } + runtime?.startupTelemetry?.end('aborted'); + if (runtime) { + runtime.startupTelemetry = undefined; + } if (this._cleanupOnComplete) { - this.runtimeState.delete(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 - await this.jobStore.updateJob(streamId, { - status: 'aborted', - completedAt: Date.now(), - }); + // A replacement created after the abort CAS makes this a safe no-op. + await this.jobStore.deleteJob(streamId, jobData.createdAt); + } + if (runtime && this.runtimeState.get(streamId) === runtime) { + this.jobStore.clearContentState(streamId, jobData.createdAt); + this.runStepBuffers?.delete(streamId); + this.replayEventWriteQueues.delete(streamId); + this.tokenUsageWriteQueues.delete(streamId); + if (this._cleanupOnComplete) { + this.runtimeState.delete(streamId); + } } - this.runningJobs.delete(streamId); - this.syncRunningJobMetrics(); + this.releaseJobOwnership(streamId, jobData.createdAt); recordGenerationJob(this.storeLabel, 'aborted'); logger.debug(`[GenerationJobManager] Job aborted: ${streamId}`); @@ -1023,63 +1452,277 @@ class GenerationJobManagerClass { onDone?: t.DoneHandler, onError?: t.ErrorHandler, options?: t.SubscribeOptions, - ): Promise<{ unsubscribe: t.UnsubscribeFn } | null> { + ): Promise { + return this.attachSubscription(streamId, onChunk, onDone, onError, options); + } + + private async attachSubscription( + streamId: string, + onChunk: t.ChunkHandler, + onDone?: t.DoneHandler, + onError?: t.ErrorHandler, + options?: t.SubscribeOptions, + prepared?: PreparedSubscription, + ): Promise<(t.StreamSubscription & { activate?: () => void }) | null> { const subscriptionType = options?.skipBufferReplay ? 'resume' : 'initial'; - // Use lazy initialization to support cross-replica subscriptions - const runtime = await this.getOrCreateRuntimeState(streamId); + if (options?.signal?.aborted) { + recordGenerationStreamSubscription(this.storeLabel, subscriptionType, 'error'); + return null; + } + if (this.rejectSubscriptionDuringShutdown(subscriptionType, onError)) { + return null; + } + + // Read the durable generation first, then reconcile any lazily-created + // runtime against it. This also avoids the historical second Redis lookup + // when a cross-replica subscriber has no local runtime yet. + const jobData = prepared ? prepared.jobData : await this.jobStore.getJob(streamId); + if (options?.signal?.aborted) { + recordGenerationStreamSubscription(this.storeLabel, subscriptionType, 'error'); + return null; + } + if (this.rejectSubscriptionDuringShutdown(subscriptionType, onError)) { + return null; + } + + const runtime = prepared?.runtime ?? (await this.getOrCreateRuntimeState(streamId, jobData)); + if (options?.signal?.aborted) { + recordGenerationStreamSubscription(this.storeLabel, subscriptionType, 'error'); + return null; + } + if (this.rejectSubscriptionDuringShutdown(subscriptionType, onError)) { + return null; + } if (!runtime) { recordGenerationStreamSubscription(this.storeLabel, subscriptionType, 'not_found'); return null; } - const jobData = await this.jobStore.getJob(streamId); + if ( + this.runtimeState.get(streamId) !== runtime || + (jobData && runtime.createdAt !== jobData.createdAt) + ) { + recordGenerationStreamSubscription(this.storeLabel, subscriptionType, 'error'); + return null; + } - // If job already complete/error, send final event or error - // Error status takes precedence to ensure errors aren't misreported as successes - setImmediate(() => { - if (jobData && ['complete', 'error', 'aborted'].includes(jobData.status)) { - // Check for error status FIRST and prioritize error handling - if (jobData.status === 'error' && (runtime.errorEvent || jobData.error)) { - const errorToSend = runtime.errorEvent ?? jobData.error; - if (errorToSend) { - logger.debug( - `[GenerationJobManager] Sending stored error to late subscriber: ${streamId}`, - ); - onError?.(errorToSend); - } - } else if (runtime.finalEvent) { - onDone?.(runtime.finalEvent); + let subscriptionActive = true; + let createdEventDelivered = options?.skipBufferReplay === true; + let terminalEventDelivered = false; + let terminalEventQueued = false; + let deliveryActivated = prepared?.deferDeliveryUntilActivated !== true; + let deferredDeliveries: DeferredDelivery[] = []; + let subscription: { + ready?: Promise; + unsubscribe: t.UnsubscribeFn; + activate?: () => void; + } | null = null; + const releaseSubscriberOnlyAbortSubscription = (generationId?: number): void => { + if ( + generationId == null || + generationId !== runtime.createdAt || + this.ownedJobs.get(streamId) === runtime.createdAt + ) { + return; + } + this.releaseAbortSubscription(runtime); + runtime.abortController.abort(); + }; + const deliverChunk = (event: t.ServerSentEvent): void => { + if (!subscriptionActive || terminalEventDelivered) { + return; + } + if ('created' in event) { + if (createdEventDelivered) { + return; + } + createdEventDelivered = true; + } + onChunk(event); + }; + const deliverDone = (event: t.ServerSentEvent): void => { + if (!subscriptionActive || terminalEventDelivered) { + return; + } + terminalEventDelivered = true; + runtime.finalEvent = event; + try { + onDone?.(event); + } finally { + subscription?.unsubscribe(); + } + }; + const deliverError = (error: string): void => { + if (!subscriptionActive || terminalEventDelivered) { + return; + } + terminalEventDelivered = true; + // The pre-drain shutdown error only closes this process's SSE response; it is not a + // durable terminal job error. Leave errorEvent unset so all-subscribers-left cleanup + // still persists partial progress before the job store is destroyed. + if (error !== SHUTDOWN_SUBSCRIBER_ERROR) { + runtime.errorEvent = error; + } + try { + onError?.(error); + } finally { + subscription?.unsubscribe(); + } + }; + const queueChunk = (event: t.ServerSentEvent): void => { + if (!subscriptionActive || terminalEventDelivered || terminalEventQueued) { + return; + } + if (!deliveryActivated) { + deferredDeliveries.push({ type: 'chunk', event }); + return; + } + deliverChunk(event); + }; + const queueDone = (event: t.ServerSentEvent, generationId?: number): void => { + if (generationId != null && generationId !== runtime.createdAt) { + return; + } + if (!subscriptionActive || terminalEventDelivered || terminalEventQueued) { + return; + } + if (!deliveryActivated) { + terminalEventQueued = true; + runtime.finalEvent = event; + deferredDeliveries.push({ type: 'done', event }); + return; + } + deliverDone(event); + }; + const queueError = (error: string, generationId?: number): void => { + if (generationId != null && generationId !== runtime.createdAt) { + return; + } + if (!subscriptionActive || terminalEventDelivered || terminalEventQueued) { + return; + } + if (!deliveryActivated) { + terminalEventQueued = true; + if (error !== SHUTDOWN_SUBSCRIBER_ERROR) { + runtime.errorEvent = error; + } + deferredDeliveries.push({ type: 'error', error }); + return; + } + deliverError(error); + }; + const activateDelivery = (): void => { + if (!subscriptionActive || deliveryActivated) { + return; + } + + deliveryActivated = true; + const deliveries = deferredDeliveries; + deferredDeliveries = []; + + for (const delivery of deliveries) { + if (!subscriptionActive) { + return; + } + if (delivery.type === 'chunk') { + deliverChunk(delivery.event); + } else if (delivery.type === 'done') { + terminalEventQueued = false; + deliverDone(delivery.event); + } else { + terminalEventQueued = false; + deliverError(delivery.error); } } - }); + }; - const subscription = this.eventTransport.subscribe( + const deferSequenceDelivery = + this._isRedis && !runtime.hasSubscriber && !options?.skipBufferReplay; + const transportSubscription = this.eventTransport.subscribe( streamId, { - onChunk: (event) => { + onChunk: (event, generationId) => { + if ( + this.runtimeState.get(streamId) !== runtime || + (generationId != null && generationId !== runtime.createdAt) + ) { + return; + } const e = event as t.ServerSentEvent; if (!(e as Record)._internal) { - onChunk(e); + queueChunk(e); } }, - onDone: (event) => onDone?.(event as t.ServerSentEvent), - onError, + onDone: (event, generationId) => { + releaseSubscriberOnlyAbortSubscription(generationId); + queueDone(event as t.ServerSentEvent, generationId); + }, + onError: (error, generationId) => { + releaseSubscriberOnlyAbortSubscription(generationId); + queueError(error, generationId); + }, }, { // Redis can publish an early buffered event before the EVAL response carrying its // sequence reaches this process. Hold sequenced pub/sub delivery until replay and // sync establish the exact frontier, otherwise the new subscriber sees it twice. - deferSequenceDelivery: - this._isRedis && !runtime.hasSubscriber && !options?.skipBufferReplay, + deferSequenceDelivery, }, ); + runtime.localErrorHandlers.add(queueError); + if (!options?.skipBufferReplay) { + runtime.earlyReplayHandlers.add(queueChunk); + } + let resolveDetached!: () => void; + const detached = new Promise((resolve) => { + resolveDetached = resolve; + }); + const detachSignal = options?.signal; + const detachOnAbort = (): void => { + subscription?.unsubscribe(); + }; + subscription = { + ready: transportSubscription.ready, + ...(prepared?.deferDeliveryUntilActivated === true && { activate: activateDelivery }), + unsubscribe: (): void => { + if (!subscriptionActive) { + return; + } + subscriptionActive = false; + deferredDeliveries = []; + runtime.earlyReplayHandlers.delete(queueChunk); + runtime.localErrorHandlers.delete(queueError); + detachSignal?.removeEventListener('abort', detachOnAbort); + transportSubscription.unsubscribe(); + resolveDetached(); + }, + }; + if (detachSignal?.aborted) { + subscription.unsubscribe(); + } else { + detachSignal?.addEventListener('abort', detachOnAbort, { once: true }); + } + if (terminalEventDelivered) { + subscription.unsubscribe(); + } + + const waitWhileAttached = async (pending: Promise): Promise => { + await Promise.race([pending, detached]); + return subscriptionActive; + }; try { - if (subscription.ready) { - await subscription.ready; + if (subscription.ready && !(await waitWhileAttached(subscription.ready))) { + recordGenerationStreamSubscription(this.storeLabel, subscriptionType, 'error'); + return null; + } + if (this.detachSubscriptionDuringShutdown(subscription)) { + recordGenerationStreamSubscription(this.storeLabel, subscriptionType, 'error'); + return null; } recordGenerationStreamSubscription(this.storeLabel, subscriptionType, 'success'); } catch (err) { + subscription.unsubscribe(); recordGenerationStreamSubscription(this.storeLabel, subscriptionType, 'error'); throw err; } @@ -1088,73 +1731,109 @@ class GenerationJobManagerClass { if (!runtime.hasSubscriber) { runtime.hasSubscriber = true; + const attachmentGeneration = runtime.attachmentGeneration; + const earlyPublicationFence = this.waitForEarlyEventPublications(runtime); + if (!(await waitWhileAttached(earlyPublicationFence))) { + this.continueEarlyEventBootstrap( + streamId, + runtime, + earlyPublicationFence, + jobData, + attachmentGeneration, + deferSequenceDelivery, + ); + return null; + } + if (this.detachSubscriptionDuringShutdown(subscription)) { + return null; + } /** - * The Redis sequence is conversation-scoped and therefore may start this - * generation above zero. Synchronize with the absolute sequence frontier - * assigned to the events replayed below, never with their relative count. + * Redis sequences are conversation-scoped and may start above zero. Use the + * absolute sequences assigned to the exact events replayed below; a relative + * buffer count can skip live events from a later turn. * - * When skipBufferReplay is true, the resume sync payload delivers aggregated - * content up to the Redis counter, so syncReorderBuffer receives no local - * replay frontier and trusts the current counter. + * When no local replay occurs (including resume), undefined tells the transport + * to trust the current Redis counter. */ let replayedNextSeq: number | undefined; const bufferedEvents = runtime.earlyEventBuffer; const sequencePromises = runtime.earlyEventSequencePromises; - runtime.earlyEventBuffer = []; - runtime.earlyEventSequencePromises = []; - - if (bufferedEvents.length > 0) { - const sequences = await Promise.all(sequencePromises); - if (options?.skipBufferReplay) { - logger.debug( - `[GenerationJobManager] Skipping ${bufferedEvents.length} buffered events for ${streamId} (skipBufferReplay)`, - ); - } else { - const assignedSequences = sequences.filter( - (sequence): sequence is number => typeof sequence === 'number', - ); - if (assignedSequences.length > 0) { - replayedNextSeq = Math.max(...assignedSequences) + 1; - } - logger.debug( - `[GenerationJobManager] Replaying ${bufferedEvents.length} buffered events for ${streamId}`, - ); - for (const bufferedEvent of bufferedEvents) { - onChunk(bufferedEvent); - } - } - } else if (this._isRedis && !options?.skipBufferReplay && jobData?.userMessage) { - /** - * Cross-replica fallback: the created event was buffered on the generating - * instance and published via Redis pub/sub before this subscriber was active. - * Reconstruct from persisted metadata. Only fields stored by trackUserMessage() - * are available (messageId, parentMessageId, conversationId, text); - * sender/isCreatedByUser are invariant for user messages and added back here. - */ - logger.debug( - `[GenerationJobManager] Cross-replica subscribe: emitting created event from metadata for ${streamId}`, - ); - const createdEvent: t.CreatedEvent = { - created: true, - message: { - ...jobData.userMessage, - sender: 'User', - isCreatedByUser: true, - }, - streamId, - }; - onChunk(createdEvent); - } + const hasEarlyReplaySubscribers = runtime.earlyReplayHandlers.size > 0; try { - await this.eventTransport.syncReorderBuffer?.(streamId, replayedNextSeq); - } catch (err) { - logger.warn( - `[GenerationJobManager] Failed to sync reorder buffer for ${streamId}; proceeding with current nextSeq:`, - err, - ); + if (bufferedEvents.length > 0) { + if (!hasEarlyReplaySubscribers) { + logger.debug( + `[GenerationJobManager] Skipping ${bufferedEvents.length} buffered events for ${streamId} (skipBufferReplay)`, + ); + } else { + const sequences = await Promise.all(sequencePromises); + const assignedSequences = sequences.filter( + (sequence): sequence is number => typeof sequence === 'number', + ); + if (assignedSequences.length > 0) { + replayedNextSeq = Math.max(...assignedSequences) + 1; + } + logger.debug( + `[GenerationJobManager] Replaying ${bufferedEvents.length} buffered events for ${streamId}`, + ); + for (const bufferedEvent of bufferedEvents) { + for (const replayHandler of runtime.earlyReplayHandlers) { + replayHandler(bufferedEvent); + } + } + } + } else if (this._isRedis && hasEarlyReplaySubscribers && jobData?.userMessage) { + /** + * Cross-replica fallback: metadata can be visible before the generating + * replica publishes `created`. Emit the fallback before releasing buffered + * live events so `created` remains the first user-facing event. deliverChunk + * suppresses the original publication whether it is already pending or + * arrives after synchronization. + */ + logger.debug( + `[GenerationJobManager] Cross-replica subscribe: emitting created event from metadata for ${streamId}`, + ); + const fallbackCreatedEvent: t.ServerSentEvent = { + created: true, + message: { + ...jobData.userMessage, + sender: 'User', + isCreatedByUser: true, + }, + streamId, + }; + for (const replayHandler of runtime.earlyReplayHandlers) { + replayHandler(fallbackCreatedEvent); + } + } + } finally { + runtime.earlyEventBuffer = []; + runtime.earlyEventSequencePromises = []; + try { + const reorderSync = this.eventTransport.syncReorderBuffer?.(streamId, replayedNextSeq); + if (reorderSync) { + await waitWhileAttached(reorderSync); + } + } catch (err) { + logger.warn( + `[GenerationJobManager] Failed to sync reorder buffer for ${streamId}; proceeding with current nextSeq:`, + err, + ); + } } + + if (!subscriptionActive) { + return null; + } + if (this.detachSubscriptionDuringShutdown(subscription)) { + return null; + } + } + + if (this.detachSubscriptionDuringShutdown(subscription)) { + return null; } if (isFirst) { @@ -1164,122 +1843,464 @@ class GenerationJobManagerClass { ); } + // Only schedule stored terminal delivery after the attachment is fully prepared. + // The async function resolves before setImmediate runs, giving the route its + // unsubscribe handle before a terminal callback can end the response. + setImmediate(() => { + void (async () => { + if (this.shuttingDown || !subscriptionActive || terminalEventDelivered) { + return; + } + + let terminalJob = jobData; + if (!terminalJob || !['complete', 'error', 'aborted'].includes(terminalJob.status)) { + try { + terminalJob = await this.jobStore.getJob(streamId); + } catch (err) { + logger.warn( + `[GenerationJobManager] Failed to refresh terminal state for ${streamId}:`, + err, + ); + return; + } + } + if ( + this.shuttingDown || + !subscriptionActive || + terminalEventDelivered || + !terminalJob || + !['complete', 'error', 'aborted'].includes(terminalJob.status) + ) { + return; + } + + // A durable error takes precedence for every terminal status. Approval expiry uses + // `aborted` so a late subscriber still needs the stored terminal error. + if (runtime.errorEvent || terminalJob.error) { + const errorToSend = runtime.errorEvent ?? terminalJob.error; + if (errorToSend) { + logger.debug( + `[GenerationJobManager] Sending stored error to late subscriber: ${streamId}`, + ); + runtime.errorEvent = errorToSend; + queueError(errorToSend); + } + return; + } + + let finalEvent = runtime.finalEvent; + if (!finalEvent && terminalJob.finalEvent) { + try { + finalEvent = JSON.parse(terminalJob.finalEvent) as t.ServerSentEvent; + } catch (err) { + logger.warn( + `[GenerationJobManager] Failed to parse stored final event for ${streamId}:`, + err, + ); + } + } + if (finalEvent) { + runtime.finalEvent = finalEvent; + queueDone(finalEvent); + } + })(); + }); + return subscription; } /** - * Atomic resume + subscribe: snapshots resume state and drains the early event buffer - * in one synchronous step, then subscribes with skipBufferReplay. + * Wait until every buffered publication has an authoritative Redis sequence before replay. + * Replaying while a publication is unresolved can deliver the local copy and then deliver the + * same event again when its late pub/sub message arrives. + */ + private async waitForEarlyEventPublications(runtime: RuntimeJobState): Promise { + const pending = [...runtime.earlyEventSequencePromises]; + if (pending.length === 0) { + return; + } + + await Promise.all(pending); + } + + /** + * If the subscriber that owns Redis attachment bootstrap disconnects, finish the + * replay/sync for any concurrent subscriber. Otherwise the transport-wide reorder + * fence remains closed forever because later subscribers observe hasSubscriber=true. + */ + private continueEarlyEventBootstrap( + streamId: string, + runtime: RuntimeJobState, + publicationFence: Promise, + jobData: SerializableJobData | null, + attachmentGeneration: number, + sequenceDeliveryDeferred: boolean, + ): void { + if (this.eventTransport.getSubscriberCount(streamId) === 0) { + return; + } + + void publicationFence + .then(async () => { + if ( + this.shuttingDown || + this.eventTransport.getSubscriberCount(streamId) === 0 || + this.runtimeState.get(streamId) !== runtime || + runtime.attachmentGeneration !== attachmentGeneration + ) { + return; + } + + let replayedNextSeq: number | undefined; + try { + const hasEarlyReplaySubscribers = runtime.earlyReplayHandlers.size > 0; + if (hasEarlyReplaySubscribers && runtime.earlyEventBuffer.length > 0) { + const sequences = await Promise.all(runtime.earlyEventSequencePromises); + const assignedSequences = sequences.filter( + (sequence): sequence is number => typeof sequence === 'number', + ); + if (sequenceDeliveryDeferred && assignedSequences.length > 0) { + replayedNextSeq = Math.max(...assignedSequences) + 1; + } + for (const [index, bufferedEvent] of runtime.earlyEventBuffer.entries()) { + /** + * A canceled resume bootstrap does not defer Redis delivery. Any event with an + * assigned sequence was therefore already published to the surviving subscriber; + * replaying it locally would duplicate the event. Failed publications have no + * sequence and still need the local replay. When delivery was deferred, replay + * every buffered event and prune its pending pub/sub copy during synchronization. + */ + if (!sequenceDeliveryDeferred && typeof sequences[index] === 'number') { + continue; + } + for (const replayHandler of runtime.earlyReplayHandlers) { + replayHandler(bufferedEvent); + } + } + } else if (hasEarlyReplaySubscribers && jobData?.userMessage) { + const fallbackCreatedEvent: t.ServerSentEvent = { + created: true, + message: { + ...jobData.userMessage, + sender: 'User', + isCreatedByUser: true, + }, + streamId, + }; + for (const replayHandler of runtime.earlyReplayHandlers) { + replayHandler(fallbackCreatedEvent); + } + } + } finally { + runtime.earlyEventBuffer = []; + runtime.earlyEventSequencePromises = []; + await this.eventTransport.syncReorderBuffer?.(streamId, replayedNextSeq); + } + }) + .catch((err) => { + logger.warn( + `[GenerationJobManager] Failed to finish detached attachment bootstrap for ${streamId}:`, + err, + ); + }); + } + + /** + * Snapshots resume state and attaches a paused live subscription. * - * Closes the timing gap between separate `getResumeState()` and `subscribe()` calls - * where events could arrive in earlyEventBuffer after the snapshot but before subscribe - * clears the buffer. - * - * In-memory mode: drained buffer events are returned as `pendingEvents` since - * they exist nowhere else. The caller must deliver them after the sync payload. - * Redis mode: `pendingEvents` is empty — chunks are persisted via appendChunk - * and will appear in aggregatedContent on the next resume. + * In-memory emissions during the snapshot-to-attachment interval are captured per resume, + * so overlapping reconnects do not compete for the shared early-event buffer. Live delivery + * remains paused until the caller writes its sync frame and activates the subscription. */ async subscribeWithResume( streamId: string, onChunk: t.ChunkHandler, onDone?: t.DoneHandler, onError?: t.ErrorHandler, + options?: Pick, ): Promise { - const bufferLengthAtSnapshot = !this._isRedis - ? (this.runtimeState.get(streamId)?.earlyEventBuffer.length ?? 0) - : 0; - - const resumeState = await this.getResumeState(streamId); - recordGenerationStreamSubscription( - this.storeLabel, - 'resume_state', - resumeState ? 'found' : 'missing', - ); - - let pendingEvents: t.ServerSentEvent[] = []; - if (!this._isRedis) { - const runtime = this.runtimeState.get(streamId); - if (runtime) { - pendingEvents = runtime.earlyEventBuffer.slice(bufferLengthAtSnapshot); - runtime.earlyEventBuffer = []; - if (pendingEvents.length > 0) { - recordGenerationStreamResumePendingEvents(this.storeLabel, pendingEvents.length); - logger.debug( - `[GenerationJobManager] Captured ${pendingEvents.length} gap events for ${streamId}`, - ); - } - } + if (options?.signal?.aborted) { + recordGenerationStreamSubscription(this.storeLabel, 'resume', 'error'); + return { subscription: null, resumeState: null, pendingEvents: [] }; + } + if (this.rejectSubscriptionDuringShutdown('resume', onError)) { + return { subscription: null, resumeState: null, pendingEvents: [] }; } - const subscription = await this.subscribe(streamId, onChunk, onDone, onError, { - skipBufferReplay: true, - }); + const runtime = await this.getOrCreateRuntimeState(streamId); + if (options?.signal?.aborted) { + recordGenerationStreamSubscription(this.storeLabel, 'resume', 'error'); + return { subscription: null, resumeState: null, pendingEvents: [] }; + } + if (this.rejectSubscriptionDuringShutdown('resume', onError)) { + return { subscription: null, resumeState: null, pendingEvents: [] }; + } + if (!runtime) { + recordGenerationStreamSubscription(this.storeLabel, 'resume_state', 'missing'); + recordGenerationStreamSubscription(this.storeLabel, 'resume', 'not_found'); + return { subscription: null, resumeState: null, pendingEvents: [] }; + } - // Close the snapshot→subscribe race: getResumeState() snapshots BEFORE we attach the - // subscription, so a pause that becomes durable in that window is in neither - // resumeState.pendingAction nor (Redis mode) pendingEvents — and trackReplayEvent does - // not persist approval events — leaving the client attached to a paused job with no - // approval UI. Re-read the live job AFTER subscribing; if it is now requires_action and - // the snapshot didn't already carry the action, surface it as a pending event so the - // approval prompt renders. Idempotent: a pause landing AFTER attach is delivered live - // too, and the client's handler just sets the current action, so a duplicate is benign. - const liveJob = await this.jobStore.getJob(streamId); - if (!resumeState?.pendingAction) { - if ( - liveJob?.status === 'requires_action' && - liveJob.pendingAction != null && - !isPendingActionStale(liveJob) - ) { - pendingEvents = [ - ...pendingEvents, - { + const capturedPendingEvents: t.ServerSentEvent[] = []; + const pendingEvents: t.ServerSentEvent[] = []; + const capturedEventSet = new Set(); + const snapshotCoveredEventSet = new Set(); + const seenEmissionEvents = new Set(); + const unclassifiedEmissions: Array<{ event: t.ServerSentEvent; sequence: number }> = []; + let snapshotFrontier = 0; + let snapshotClassified = this._isRedis; + const classifyEmission = (event: t.ServerSentEvent, sequence: number): void => { + if (sequence <= snapshotFrontier) { + snapshotCoveredEventSet.add(event); + return; + } + capturedEventSet.add(event); + capturedPendingEvents.push(event); + pendingEvents.push(event); + }; + const capturePendingEvent = (event: t.ServerSentEvent, sequence: number): void => { + if (seenEmissionEvents.has(event)) { + return; + } + seenEmissionEvents.add(event); + if (!snapshotClassified) { + unclassifiedEmissions.push({ event, sequence }); + return; + } + classifyEmission(event, sequence); + }; + let resumeState: t.ResumeState | null = null; + let jobData: SerializableJobData | null = null; + const removeCaptureHandler = (): void => { + runtime.resumeCaptureHandlers.delete(capturePendingEvent); + }; + const restoreCapturedEvents = (): void => { + if (capturedPendingEvents.length === 0) { + return; + } + const currentRuntime = this.runtimeState.get(streamId); + if (currentRuntime && !currentRuntime.hasSubscriber) { + const bufferedEvents = new Set(currentRuntime.earlyEventBuffer); + const missingEvents = capturedPendingEvents.filter((event) => !bufferedEvents.has(event)); + if (missingEvents.length > 0) { + currentRuntime.earlyEventBuffer = [...missingEvents, ...currentRuntime.earlyEventBuffer]; + } + } + capturedPendingEvents.length = 0; + capturedEventSet.clear(); + }; + let subscription: (t.StreamSubscription & { activate?: () => void }) | null = null; + try { + if (!this._isRedis) { + runtime.resumeCaptureHandlers.add(capturePendingEvent); + while (true) { + const candidateFrontier = runtime.emissionSequence; + const preSnapshotEmissions = [...runtime.inFlightSnapshotEmissions.entries()].filter( + ([sequence]) => sequence <= candidateFrontier, + ); + await Promise.all(preSnapshotEmissions.map(([, emission]) => emission.snapshotReady)); + const [candidateState, candidateJob] = await Promise.all([ + this.getResumeState(streamId), + this.jobStore.getJob(streamId), + ]); + if ( + runtime.emissionSequence !== candidateFrontier && + !options?.signal?.aborted && + !this.shuttingDown + ) { + continue; + } + + snapshotFrontier = candidateFrontier; + resumeState = candidateState ? structuredClone(candidateState) : null; + jobData = candidateJob; + for (const [, emission] of preSnapshotEmissions) { + seenEmissionEvents.add(emission.event); + snapshotCoveredEventSet.add(emission.event); + } + for (const emission of unclassifiedEmissions) { + classifyEmission(emission.event, emission.sequence); + } + unclassifiedEmissions.length = 0; + snapshotClassified = true; + break; + } + } else { + [resumeState, jobData] = await Promise.all([ + this.getResumeState(streamId), + this.jobStore.getJob(streamId), + ]); + } + + if (options?.signal?.aborted) { + removeCaptureHandler(); + recordGenerationStreamSubscription(this.storeLabel, 'resume', 'error'); + return { subscription: null, resumeState, pendingEvents: [] }; + } + if (this.rejectSubscriptionDuringShutdown('resume', onError)) { + removeCaptureHandler(); + return { subscription: null, resumeState, pendingEvents: [] }; + } + recordGenerationStreamSubscription( + this.storeLabel, + 'resume_state', + resumeState ? 'found' : 'missing', + ); + + const forwardLiveChunk = (event: t.ServerSentEvent): void => { + if (capturedEventSet.has(event) || snapshotCoveredEventSet.has(event)) { + return; + } + onChunk(event); + }; + subscription = await this.attachSubscription( + streamId, + forwardLiveChunk, + onDone, + onError, + { + skipBufferReplay: true, + signal: options?.signal, + }, + { + runtime, + jobData, + deferDeliveryUntilActivated: true, + }, + ); + if (pendingEvents.length > 0) { + recordGenerationStreamResumePendingEvents(this.storeLabel, pendingEvents.length); + logger.debug( + `[GenerationJobManager] Captured ${pendingEvents.length} gap events for ${streamId}`, + ); + } + const cancelResumeSubscription = (): t.SubscribeWithResumeResult => { + removeCaptureHandler(); + subscription?.unsubscribe(); + restoreCapturedEvents(); + snapshotCoveredEventSet.clear(); + return { subscription: null, resumeState, pendingEvents: [] }; + }; + if (!subscription?.activate || options?.signal?.aborted) { + return cancelResumeSubscription(); + } + if (this.detachSubscriptionDuringShutdown(subscription)) { + return cancelResumeSubscription(); + } + + // Close the snapshot→subscribe race: getResumeState() snapshots BEFORE we attach the + // subscription, so a pause that becomes durable in that window is in neither + // resumeState.pendingAction nor (Redis mode) pendingEvents — and trackReplayEvent does + // not persist approval events — leaving the client attached to a paused job with no + // approval UI. Re-read the live job AFTER subscribing; if it is now requires_action and + // the snapshot didn't already carry the action, surface it as a pending event so the + // approval prompt renders. Idempotent: a pause landing AFTER attach is delivered live + // too, and the client's handler just sets the current action, so a duplicate is benign. + const liveJob = await this.jobStore.getJob(streamId); + if (options?.signal?.aborted || this.detachSubscriptionDuringShutdown(subscription)) { + return cancelResumeSubscription(); + } + if (!liveJob || liveJob.createdAt !== runtime.createdAt) { + return cancelResumeSubscription(); + } + if (!resumeState?.pendingAction) { + if ( + liveJob?.status === 'requires_action' && + liveJob.pendingAction != null && + !isPendingActionStale(liveJob) + ) { + pendingEvents.push({ event: ApprovalEvents.ON_PENDING_ACTION, data: toClientPendingAction(liveJob.pendingAction) as unknown as Record< string, unknown >, - }, - ]; - } - } - - // Same snapshot→subscribe race for steers: a steer accepted (and possibly - // applied) in the window is invisible to the snapshot, since the Redis - // `on_steer_applied` publish is fire-and-forget and the sync payload has no - // pendingSteers (in-memory covers it via the early buffer, where this - // re-check is a cheap no-op). Always re-peek for still-active jobs, - // treating a missing snapshot queue as empty; terminal jobs skip because - // the final event owns steer delivery. The content re-read runs only when - // the queue shows gap activity, and synthesis sources from the FRESH - // content view so an applied steer with no snapshot id still surfaces. - const jobActive = liveJob?.status === 'running' || liveJob?.status === 'requires_action'; - if (resumeState != null && jobActive) { - const snapshotSteers = resumeState.pendingSteers ?? []; - const liveQueue = await this.jobStore.peekSteers(streamId); - const liveIds = new Set(liveQueue.map((item) => item.steerId)); - const queueChanged = - liveQueue.length !== snapshotSteers.length || - snapshotSteers.some((steer) => !liveIds.has(steer.steerId)); - if (queueChanged) { - const livePending = liveQueue.map(toPendingSteer); - resumeState.pendingSteers = livePending.length > 0 ? livePending : undefined; - } - if (queueChanged || liveQueue.length > 0) { - const contentResult = await this.jobStore.getContentParts(streamId); - const gapEvents = synthesizeAppliedSteerEvents( - (resumeState.aggregatedContent ?? []) as SteerContentView, - liveQueue, - (contentResult?.content ?? []) as SteerContentView, - { conversationId: streamId, responseMessageId: resumeState.responseMessageId }, - ); - if (gapEvents.length > 0) { - pendingEvents = [...pendingEvents, ...gapEvents]; + }); } } - } - return { subscription, resumeState, pendingEvents }; + // Same snapshot→subscribe race for steers: a steer accepted (and possibly + // applied) in the window is invisible to the snapshot, since the Redis + // `on_steer_applied` publish is fire-and-forget and the sync payload has no + // pendingSteers (in-memory covers it via the early buffer, where this + // re-check is a cheap no-op). Always re-peek for still-active jobs, + // treating a missing snapshot queue as empty; terminal jobs skip because + // the final event owns steer delivery. The content re-read runs only when + // the queue shows gap activity, and synthesis sources from the FRESH + // content view so an applied steer with no snapshot id still surfaces. + const jobActive = liveJob?.status === 'running' || liveJob?.status === 'requires_action'; + if (resumeState != null && jobActive) { + const snapshotSteers = resumeState.pendingSteers ?? []; + const liveQueue = await this.jobStore.peekSteers(streamId, liveJob.createdAt); + if (options?.signal?.aborted || this.detachSubscriptionDuringShutdown(subscription)) { + return cancelResumeSubscription(); + } + const liveIds = new Set(liveQueue.map((item) => item.steerId)); + const queueChanged = + liveQueue.length !== snapshotSteers.length || + snapshotSteers.some((steer) => !liveIds.has(steer.steerId)); + if (queueChanged) { + const livePending = liveQueue.map(toPendingSteer); + resumeState.pendingSteers = livePending.length > 0 ? livePending : undefined; + } + if (queueChanged || liveQueue.length > 0) { + const contentResult = await this.jobStore.getContentParts(streamId, liveJob.createdAt); + if (options?.signal?.aborted || this.detachSubscriptionDuringShutdown(subscription)) { + return cancelResumeSubscription(); + } + const gapEvents = synthesizeAppliedSteerEvents( + (resumeState.aggregatedContent ?? []) as SteerContentView, + liveQueue, + (contentResult?.content ?? []) as SteerContentView, + { conversationId: streamId, responseMessageId: resumeState.responseMessageId }, + ); + if (gapEvents.length > 0) { + pendingEvents.push(...gapEvents); + } + } + } + + // Reconciliation is complete. Events that arrive after this point already belong to the + // paused transport subscription and must remain there for activation rather than being + // appended to a pending-events array the caller may already be serializing. + removeCaptureHandler(); + const activate = subscription.activate; + let activated = false; + let closed = false; + const resumeSubscription: t.ResumeSubscription = { + unsubscribe: () => { + if (closed) { + return; + } + closed = true; + removeCaptureHandler(); + subscription?.unsubscribe(); + if (!activated) { + restoreCapturedEvents(); + } + capturedEventSet.clear(); + snapshotCoveredEventSet.clear(); + }, + activate: () => { + if (closed || activated) { + return; + } + activated = true; + removeCaptureHandler(); + activate(); + capturedPendingEvents.length = 0; + capturedEventSet.clear(); + snapshotCoveredEventSet.clear(); + }, + }; + return { subscription: resumeSubscription, resumeState, pendingEvents }; + } catch (err) { + removeCaptureHandler(); + subscription?.unsubscribe(); + restoreCapturedEvents(); + snapshotCoveredEventSet.clear(); + throw err; + } } /** @@ -1305,20 +2326,93 @@ class GenerationJobManagerClass { options?: { durable?: boolean }, ): Promise { const runtime = this.runtimeState.get(streamId); - if (!runtime || runtime.abortController.signal.aborted) { + if (!runtime || !this.isCurrentRuntime(streamId, runtime)) { + return; + } + + const sequence = ++runtime.emissionSequence; + let signalSnapshotReady!: () => void; + const snapshotReady = new Promise((resolve) => { + signalSnapshotReady = resolve; + }); + let snapshotReadySignaled = false; + const markSnapshotReady = (): void => { + if (snapshotReadySignaled) { + return; + } + snapshotReadySignaled = true; + signalSnapshotReady(); + }; + runtime.inFlightSnapshotEmissions.set(sequence, { event, snapshotReady }); + + try { + const isCreatedEvent = 'created' in event; + const pendingCreatedEvent = runtime.createdEventPublication; + if (!isCreatedEvent) { + if (pendingCreatedEvent) { + await pendingCreatedEvent; + if (!this.isCurrentRuntime(streamId, runtime)) { + return; + } + } + await this.emitChunkNow(streamId, event, runtime, sequence, markSnapshotReady, options); + return; + } + + if (pendingCreatedEvent) { + await pendingCreatedEvent; + if (!this.isCurrentRuntime(streamId, runtime)) { + return; + } + } + + let releaseCreatedEvent!: () => void; + const createdEventPublication = new Promise((resolve) => { + releaseCreatedEvent = resolve; + }); + runtime.createdEventPublication = createdEventPublication; + + try { + await this.emitChunkNow(streamId, event, runtime, sequence, markSnapshotReady, options); + } finally { + releaseCreatedEvent(); + if (runtime.createdEventPublication === createdEventPublication) { + runtime.createdEventPublication = undefined; + } + } + } finally { + markSnapshotReady(); + if (runtime.inFlightSnapshotEmissions.get(sequence)?.event === event) { + runtime.inFlightSnapshotEmissions.delete(sequence); + } + } + } + + private async emitChunkNow( + streamId: string, + event: t.ServerSentEvent, + runtime: RuntimeJobState, + sequence: number, + markSnapshotReady: () => void, + options?: { durable?: boolean }, + ): Promise { + if (!this.isCurrentRuntime(streamId, runtime)) { return; } // Refresh job activity so the store's stale-job failsafe reaps on inactivity // (a hung generation), not on age (a long but live stream). Parity with // RedisJobStore refreshing the running TTL on each appendChunk. - this.jobStore.recordActivity?.(streamId); + this.jobStore.recordActivity?.(streamId, runtime.createdAt); - await this.trackUserMessage(streamId, event); - await this.trackTitleEvent(streamId, event); - await this.trackReplayEvent(streamId, event); - await this.trackContextUsage(streamId, event); - await this.trackTokenUsage(streamId, event); + const eventTracking = this.trackEvent(streamId, event, runtime.createdAt); + if (eventTracking) { + await eventTracking; + if (!this.isCurrentRuntime(streamId, runtime)) { + return; + } + } + markSnapshotReady(); // For Redis mode, persist chunk for later reconstruction (fire-and-forget for resumability) if (this._isRedis) { @@ -1331,44 +2425,137 @@ class GenerationJobManagerClass { if (eventType && eventData !== undefined) { // Store in format expected by aggregateContent: { event, data } const appendPromise = this.jobStore - .appendChunk(streamId, { event: eventType, data: eventData }) + .appendChunk(streamId, { event: eventType, data: eventData }, runtime.createdAt) .catch((err) => { logger.error(`[GenerationJobManager] Failed to append chunk:`, err); }); // For run step events, also save to run steps key for quick retrieval if (eventType === 'on_run_step' || eventType === 'on_run_step_completed') { - this.saveRunStepFromEvent(streamId, eventData as Record); + this.saveRunStepFromEvent( + streamId, + eventData as Record, + runtime.createdAt, + ); } if (options?.durable === true) { await appendPromise; + if (!this.isCurrentRuntime(streamId, runtime)) { + return; + } } } } - const shouldBuffer = !runtime.hasSubscriber; - if (shouldBuffer) { + if (!this.isCurrentRuntime(streamId, runtime)) { + return; + } + + if (!this._isRedis && runtime.resumeCaptureHandlers.size > 0) { + for (const captureHandler of runtime.resumeCaptureHandlers) { + captureHandler(event, sequence); + } + } + + const buffered = !runtime.hasSubscriber; + if (buffered) { runtime.earlyEventBuffer.push(event); if (!this._isRedis) { + if (runtime.startupTelemetry) { + this.recordStartupEvent(runtime, event); + } return; } } - const publishPromise = Promise.resolve(this.eventTransport.emitChunk(streamId, event)); - if (shouldBuffer) { - // Store the promise before yielding so subscribe() can wait for the exact - // sequence assignment that belongs to every event it replays. - runtime.earlyEventSequencePromises.push(publishPromise); + if (!buffered && !runtime.startupTelemetry) { + await this.eventTransport.emitChunk(streamId, event, runtime.createdAt); + return; + } + + const publication = emitChunkWithReceipt( + this.eventTransport, + streamId, + event, + runtime.createdAt, + ); + if (buffered) { + // Store a non-rejecting sequence receipt before yielding. The absolute value + // establishes the exact replay frontier; a failed/unsequenced publication + // contributes no frontier but can still be replayed from the local buffer. + runtime.earlyEventSequencePromises.push( + publication.then( + (published) => (typeof published === 'number' ? published : undefined), + () => undefined, + ), + ); + } + + const published = await publication; + if ((published !== false || buffered) && runtime.startupTelemetry) { + this.recordStartupEvent(runtime, event); + } + } + + private isCurrentRuntime(streamId: string, runtime: RuntimeJobState): boolean { + return this.runtimeState.get(streamId) === runtime && !runtime.abortController.signal.aborted; + } + + private recordStartupEvent(runtime: RuntimeJobState, event: t.ServerSentEvent): void { + const telemetry = runtime.startupTelemetry; + if (!telemetry) { + return; + } + if ('created' in event) { + telemetry.mark('request_message_queued'); + return; + } + if (!telemetry.recordGenerationEvent(event)) { + return; + } + runtime.startupTelemetry = undefined; + } + + private trackEvent( + streamId: string, + event: t.ServerSentEvent, + expectedCreatedAt: number, + ): Promise | undefined { + if ('created' in event) { + return this.trackUserMessage(streamId, event, expectedCreatedAt); + } + if (!('event' in event)) { + return; + } + if (event.event === 'title') { + return this.trackTitleEvent(streamId, event, expectedCreatedAt); + } + if (event.event === UsageEvents.ON_CONTEXT_USAGE) { + return this.trackContextUsage(streamId, event, expectedCreatedAt); + } + if (event.event === UsageEvents.ON_TOKEN_USAGE) { + return this.trackTokenUsage(streamId, event, expectedCreatedAt); + } + if ( + (event.event === 'on_run_step' || + event.event === 'on_run_step_delta' || + event.event === 'on_run_step_completed') && + isOAuthReplayEvent(event) + ) { + return this.trackReplayEvent(streamId, event, expectedCreatedAt); } - await publishPromise; } /** * Extract and save run step from event data. * The data is already the run step object from the event payload. */ - private saveRunStepFromEvent(streamId: string, data: Record): void { + private saveRunStepFromEvent( + streamId: string, + data: Record, + expectedCreatedAt: number, + ): void { // The data IS the run step object const runStep = data as Agents.RunStep; if (!runStep.id) { @@ -1376,7 +2563,7 @@ class GenerationJobManagerClass { } // Fire and forget - accumulate run steps - this.accumulateRunStep(streamId, runStep); + this.accumulateRunStep(streamId, runStep, expectedCreatedAt); } /** @@ -1384,19 +2571,24 @@ class GenerationJobManagerClass { * Uses a simple in-memory buffer that gets flushed to Redis. * Not used in in-memory mode - run steps come from live graph via WeakRef. */ - private runStepBuffers: Map | null = null; + private runStepBuffers: Map | null = null; - private accumulateRunStep(streamId: string, runStep: Agents.RunStep): void { + private accumulateRunStep( + streamId: string, + runStep: Agents.RunStep, + expectedCreatedAt: number, + ): void { // Lazy initialization - only create map when first used (Redis mode) if (!this.runStepBuffers) { this.runStepBuffers = new Map(); } - let buffer = this.runStepBuffers.get(streamId); - if (!buffer) { - buffer = []; - this.runStepBuffers.set(streamId, buffer); + let bufferState = this.runStepBuffers.get(streamId); + if (!bufferState || bufferState.createdAt !== expectedCreatedAt) { + bufferState = { createdAt: expectedCreatedAt, steps: [] }; + this.runStepBuffers.set(streamId, bufferState); } + const buffer = bufferState.steps; // Update or add run step const existingIdx = buffer.findIndex((rs) => rs.id === runStep.id); @@ -1408,7 +2600,7 @@ class GenerationJobManagerClass { // Save to Redis if (this.jobStore.saveRunSteps) { - this.jobStore.saveRunSteps(streamId, buffer).catch((err) => { + this.jobStore.saveRunSteps(streamId, buffer, expectedCreatedAt).catch((err) => { logger.error(`[GenerationJobManager] Failed to save run steps:`, err); }); } @@ -1419,14 +2611,22 @@ class GenerationJobManagerClass { * aggregation only reconstructs message parts, so UI-only events need their * own metadata slot. */ - private async trackTitleEvent(streamId: string, event: t.ServerSentEvent): Promise { + private async trackTitleEvent( + streamId: string, + event: t.ServerSentEvent, + expectedCreatedAt: number, + ): Promise { if (!('event' in event) || event.event !== 'title') { return; } - await this.jobStore.updateJob(streamId, { - titleEvent: JSON.stringify(event), - }); + await this.jobStore.updateJob( + streamId, + { + titleEvent: JSON.stringify(event), + }, + expectedCreatedAt, + ); } /** @@ -1434,7 +2634,11 @@ class GenerationJobManagerClass { * resuming client can restore the context gauge without waiting for the * next model call. */ - private async trackContextUsage(streamId: string, event: t.ServerSentEvent): Promise { + private async trackContextUsage( + streamId: string, + event: t.ServerSentEvent, + expectedCreatedAt: number, + ): Promise { if (!('event' in event) || event.event !== UsageEvents.ON_CONTEXT_USAGE) { return; } @@ -1446,9 +2650,13 @@ class GenerationJobManagerClass { * run's gauge when visible calls interleave. FIFO ordering keeps each call's * pre-invoke snapshot ahead of its own usage and behind the next snapshot. */ await this.queueJobWrite(this.tokenUsageWriteQueues, streamId, () => - this.jobStore.updateJob(streamId, { - contextUsage: JSON.stringify((event as { data?: unknown }).data ?? null), - }), + this.jobStore.updateJob( + streamId, + { + contextUsage: JSON.stringify((event as { data?: unknown }).data ?? null), + }, + expectedCreatedAt, + ), ); } @@ -1483,13 +2691,17 @@ class GenerationJobManagerClass { * Persist replay-only stream events that are needed to reconstruct active * UI state on resume but are not represented by aggregated message content. */ - private async trackReplayEvent(streamId: string, event: t.ServerSentEvent): Promise { + private async trackReplayEvent( + streamId: string, + event: t.ServerSentEvent, + expectedCreatedAt: number, + ): Promise { if (!isOAuthReplayEvent(event)) { return; } await this.queueJobWrite(this.replayEventWriteQueues, streamId, () => - this.persistReplayEvent(streamId, event), + this.persistReplayEvent(streamId, event, expectedCreatedAt), ); } @@ -1498,19 +2710,27 @@ class GenerationJobManagerClass { * usage totals on any replica (the live collectedUsage array only exists * on the generating instance). */ - private async trackTokenUsage(streamId: string, event: t.ServerSentEvent): Promise { + private async trackTokenUsage( + streamId: string, + event: t.ServerSentEvent, + expectedCreatedAt: number, + ): Promise { if (!('event' in event) || event.event !== UsageEvents.ON_TOKEN_USAGE) { return; } await this.queueJobWrite(this.tokenUsageWriteQueues, streamId, () => - this.persistTokenUsage(streamId, event as { data?: unknown }), + this.persistTokenUsage(streamId, event as { data?: unknown }, expectedCreatedAt), ); } - private async persistTokenUsage(streamId: string, event: { data?: unknown }): Promise { + private async persistTokenUsage( + streamId: string, + event: { data?: unknown }, + expectedCreatedAt: number, + ): Promise { const jobData = await this.jobStore.getJob(streamId); - if (!jobData || event.data == null) { + if (!jobData || jobData.createdAt !== expectedCreatedAt || event.data == null) { return; } @@ -1549,12 +2769,16 @@ class GenerationJobManagerClass { } } - await this.jobStore.updateJob(streamId, update); + await this.jobStore.updateJob(streamId, update, expectedCreatedAt); } - private async persistReplayEvent(streamId: string, event: t.ServerSentEvent): Promise { + private async persistReplayEvent( + streamId: string, + event: t.ServerSentEvent, + expectedCreatedAt: number, + ): Promise { const jobData = await this.jobStore.getJob(streamId); - if (!jobData) { + if (!jobData || jobData.createdAt !== expectedCreatedAt) { return; } @@ -1585,9 +2809,13 @@ class GenerationJobManagerClass { replayEvents.push(event); } - await this.jobStore.updateJob(streamId, { - replayEvents: JSON.stringify(replayEvents), - }); + await this.jobStore.updateJob( + streamId, + { + replayEvents: JSON.stringify(replayEvents), + }, + expectedCreatedAt, + ); } /** @@ -1596,7 +2824,11 @@ class GenerationJobManagerClass { * guaranteeing any cross-replica getJob() after the pub/sub window * finds userMessage in Redis. */ - private async trackUserMessage(streamId: string, event: t.ServerSentEvent): Promise { + private async trackUserMessage( + streamId: string, + event: t.ServerSentEvent, + expectedCreatedAt: number, + ): Promise { if (!('created' in event)) { return; } @@ -1635,7 +2867,7 @@ class GenerationJobManagerClass { updates.conversationId = message.conversationId; } - await this.jobStore.updateJob(streamId, updates); + await this.jobStore.updateJob(streamId, updates, expectedCreatedAt); } /** @@ -1644,76 +2876,52 @@ class GenerationJobManagerClass { async updateMetadata( streamId: string, metadata: Partial, + expectedCreatedAt?: number, ): Promise { - const updates: Partial = {}; - if (metadata.responseMessageId) { - updates.responseMessageId = metadata.responseMessageId; - } - if (metadata.sender) { - updates.sender = metadata.sender; - } - if (metadata.conversationId) { - updates.conversationId = metadata.conversationId; - } - if (metadata.userMessage) { - updates.userMessage = metadata.userMessage; - } - if (metadata.endpoint) { - updates.endpoint = metadata.endpoint; - } - if (metadata.iconURL) { - updates.iconURL = metadata.iconURL; - } - if (metadata.model) { - updates.model = metadata.model; - } - if (metadata.agent_id) { - updates.agent_id = metadata.agent_id; - } - if (metadata.isTemporary !== undefined) { - updates.isTemporary = metadata.isTemporary; - } - if (metadata.promptTokens !== undefined) { - updates.promptTokens = metadata.promptTokens; - } - if (metadata.discoveredTools) { - updates.discoveredTools = metadata.discoveredTools; - } - await this.jobStore.updateJob(streamId, updates); + const generationId = expectedCreatedAt ?? this.runtimeState.get(streamId)?.createdAt; + await this.jobStore.updateJob(streamId, sanitizeJobMetadata(metadata), generationId); } /** * Set reference to the graph's contentParts array. */ - setContentParts(streamId: string, contentParts: Agents.MessageContentComplex[]): void { - // Use runtime state check for performance (sync check) - if (!this.runtimeState.has(streamId)) { + setContentParts( + streamId: string, + contentParts: Agents.MessageContentComplex[], + expectedCreatedAt?: number, + ): void { + const runtime = this.runtimeState.get(streamId); + if (!runtime || (expectedCreatedAt != null && runtime.createdAt !== expectedCreatedAt)) { return; } - this.jobStore.setContentParts(streamId, contentParts); + this.jobStore.setContentParts(streamId, contentParts, runtime.createdAt); } /** * Set reference to the collectedUsage array. * This array accumulates token usage from all models during generation. */ - setCollectedUsage(streamId: string, collectedUsage: UsageMetadata[]): void { - // Use runtime state check for performance (sync check) - if (!this.runtimeState.has(streamId)) { + setCollectedUsage( + streamId: string, + collectedUsage: UsageMetadata[], + expectedCreatedAt?: number, + ): void { + const runtime = this.runtimeState.get(streamId); + if (!runtime || (expectedCreatedAt != null && runtime.createdAt !== expectedCreatedAt)) { return; } - this.jobStore.setCollectedUsage(streamId, collectedUsage); + this.jobStore.setCollectedUsage(streamId, collectedUsage, runtime.createdAt); } /** * Set reference to the graph instance. */ - setGraph(streamId: string, graph: StandardGraph): void { - // Use runtime state check for performance (sync check) - if (!this.runtimeState.has(streamId)) { + setGraph(streamId: string, graph: StandardGraph, expectedCreatedAt?: number): void { + const runtime = this.runtimeState.get(streamId); + if (!runtime || (expectedCreatedAt != null && runtime.createdAt !== expectedCreatedAt)) { return; } - this.jobStore.setGraph(streamId, graph); + this.jobStore.setGraph(streamId, graph, runtime.createdAt); } /** @@ -1757,9 +2965,9 @@ class GenerationJobManagerClass { * Safe despite readCachedGraph's cache-drop side effect — each call catches its own * unusable-graph throw and falls back to reconstruction, so ordering cannot change the result. */ const [result, runSteps, queuedSteers] = await Promise.all([ - this.jobStore.getContentParts(streamId), - this.jobStore.getRunSteps(streamId), - this.jobStore.peekSteers(streamId), + this.jobStore.getContentParts(streamId, jobData.createdAt), + this.jobStore.getRunSteps(streamId, jobData.createdAt), + this.jobStore.peekSteers(streamId, jobData.createdAt), ]); const aggregatedContent = result?.content ?? []; let titleEvent: t.ResumeState['titleEvent']; @@ -1831,7 +3039,7 @@ class GenerationJobManagerClass { ? toClientPendingAction(jobData.pendingAction) : undefined, pendingSteers: pendingSteers.length > 0 ? pendingSteers : undefined, - }; + } satisfies t.ResumeState; } /** @@ -1844,7 +3052,7 @@ class GenerationJobManagerClass { runtime.syncSent = true; } // Persist to Redis for cross-replica consistency - this.jobStore.updateJob(streamId, { syncSent: true }).catch((err) => { + this.jobStore.updateJob(streamId, { syncSent: true }, runtime?.createdAt).catch((err) => { logger.error(`[GenerationJobManager] Failed to persist syncSent flag:`, err); }); } @@ -1867,16 +3075,31 @@ class GenerationJobManagerClass { * Emit a done event. * Persists finalEvent to Redis for cross-replica access. */ - async emitDone(streamId: string, event: t.ServerSentEvent): Promise { + async emitDone( + streamId: string, + event: t.ServerSentEvent, + expectedCreatedAt?: number, + ): Promise { const runtime = this.runtimeState.get(streamId); - if (runtime) { - runtime.finalEvent = event; + const generationId = expectedCreatedAt ?? runtime?.createdAt; + const matchingRuntime = + runtime && (generationId == null || runtime.createdAt === generationId) ? runtime : undefined; + if (matchingRuntime) { + matchingRuntime.finalEvent = event; + } + if (matchingRuntime?.createdEventPublication) { + await matchingRuntime.createdEventPublication; } // Persist finalEvent to Redis for cross-replica consistency - this.jobStore.updateJob(streamId, { finalEvent: JSON.stringify(event) }).catch((err) => { - logger.error(`[GenerationJobManager] Failed to persist finalEvent:`, err); - }); - await this.eventTransport.emitDone(streamId, event); + this.jobStore + .updateJob(streamId, { finalEvent: JSON.stringify(event) }, generationId) + .catch((err) => { + logger.error(`[GenerationJobManager] Failed to persist finalEvent:`, err); + }); + await this.eventTransport.emitDone(streamId, event, generationId); + if (matchingRuntime?.startupTelemetry) { + this.recordStartupEvent(matchingRuntime, event); + } } /** @@ -1884,109 +3107,92 @@ class GenerationJobManagerClass { * Stores the error for late-connecting subscribers (race condition where error * occurs before client connects to SSE stream). */ - async emitError(streamId: string, error: string): Promise { + async emitError(streamId: string, error: string, expectedCreatedAt?: number): Promise { const runtime = this.runtimeState.get(streamId); - if (runtime) { - runtime.errorEvent = error; + const generationId = expectedCreatedAt ?? runtime?.createdAt; + const matchingRuntime = + runtime && (generationId == null || runtime.createdAt === generationId) ? runtime : undefined; + if (matchingRuntime) { + matchingRuntime.errorEvent = error; + } + if (matchingRuntime?.createdEventPublication) { + await matchingRuntime.createdEventPublication; } // Persist error to job store for cross-replica consistency - this.jobStore.updateJob(streamId, { error }).catch((err) => { + this.jobStore.updateJob(streamId, { error }, generationId).catch((err) => { logger.error(`[GenerationJobManager] Failed to persist error:`, err); }); - await this.eventTransport.emitError(streamId, error); + await this.eventTransport.emitError(streamId, error, generationId); + matchingRuntime?.startupTelemetry?.mark('first_response_event_queued'); } - /** - * Cleanup expired jobs. - * Also cleans up any orphaned runtime state, buffers, and event transport entries. - */ - /** - * Expire any locally-tracked approval whose window has lapsed: drive the atomic - * `requires_action → aborted` transition and, if this caller won it, emit a - * terminal error so a connected SSE client closes. Only streams this replica has - * runtime for are scanned — those are exactly the ones with a client subscribed - * here; a paused job on another replica is finalized by that replica's sweep (and - * the store's own cleanup). The durable checkpoint is reclaimed by its Mongo TTL - * index, which shares the approval window, so no cross-layer delete is needed here. - */ /** * Expire a single observed-stale pending approval NOW (immediate, not via the periodic * sweep): run the `requires_action → aborted` CAS — pinned to `actionId` so a concurrent * resolve + re-pause on a fresh action isn't aborted — and, on success, emit the terminal * `APPROVAL_EXPIRED_ERROR` so any attached SSE client gets a terminal event instead of a - * hung stream. Used by the periodic sweeper and by the resume route, which observes a - * just-expired action when the user submits a decision after the TTL lapsed. Returns true - * if this call expired the action. + * hung stream. The durable checkpoint remains bounded by its Mongo TTL; deleting the + * thread eagerly here could race a replacement generation. Used by the periodic sweeper + * and by the resume route, which observes a just-expired action when the user submits a + * decision after the TTL lapsed. Returns true if this call expired the action. */ async expireApproval(streamId: string, actionId?: string): Promise { - /** Steers accepted before the pause are frozen for its whole window - * (enqueue rejects while `requires_action`), so this pre-CAS snapshot is - * exactly what the expiry's terminal cleanup is about to delete. Read it - * BEFORE the transition — the store drops the queue key inside it — and - * park only if the CAS wins (a lost CAS means the run resumed and the - * live queue must stay untouched). */ - let parkableSteers: TPendingSteer[] = []; - let steerOwner: SteerOwner | undefined; + const observedRuntime = this.runtimeState.get(streamId); + let observedJob: SerializableJobData | null = null; try { - const job = await this.jobStore.getJob(streamId); - if (job) { - steerOwner = { userId: job.userId, tenantId: job.tenantId }; - parkableSteers = (await this.jobStore.peekSteers(streamId)).map(toPendingSteer); - } + observedJob = await this.jobStore.getJob(streamId); } catch (err) { - logger.warn(`[GenerationJobManager] Failed to snapshot steers pre-expiry ${streamId}`, err); + logger.warn(`[GenerationJobManager] Failed to read approval before expiry ${streamId}`, err); } - const expired = await this._approvals.expire(streamId, actionId); - if (!expired) { + const expiredCreatedAt = await this._approvals.expireWithIdentity( + streamId, + actionId, + observedJob?.createdAt, + ); + if (expiredCreatedAt == null) { return false; } - if (steerOwner && parkableSteers.length > 0) { - await this.steering.park(streamId, parkableSteers, steerOwner); - } - try { - await this.emitError(streamId, APPROVAL_EXPIRED_ERROR); - } catch (err) { - logger.error(`[GenerationJobManager] Failed to notify expired approval ${streamId}`, err); - } - await this.runApprovalExpiredHandler(streamId); - this.runningJobs.delete(streamId); + + await this.notifyApprovalExpiredRuntime(streamId, expiredCreatedAt, observedRuntime); return true; } - /** - * Invoke the host approval-expired cleanup, passing the job so the host can resolve - * tenant/user-scoped config (the expiry runs outside any request context). Best-effort: - * the job read and the handler itself may fail without breaking the expiry. - */ - private async runApprovalExpiredHandler( + private async notifyApprovalExpiredRuntime( streamId: string, - job?: SerializableJobData | null, + createdAt: number, + runtime: RuntimeJobState | undefined, ): Promise { - if (!this._onApprovalExpired) { + if (runtime?.createdAt === createdAt && runtime.approvalExpiryPublished) { return; } - // Dedup across the expiry paths: a locally expired approval (expireApproval) stays in - // the store/runtime for the completed-job TTL, so later sweeps re-enter the relay - // branch for the same aborted approval — run the cleanup once per runtime lifetime. - const runtime = this.runtimeState.get(streamId); - if (runtime?.approvalCleanupRan) { - return; - } - if (runtime) { - runtime.approvalCleanupRan = true; - } - let resolvedJob = job; - if (resolvedJob === undefined) { - try { - resolvedJob = await this.jobStore.getJob(streamId); - } catch { - resolvedJob = null; - } + + if (runtime?.createdAt === createdAt) { + runtime.errorEvent = APPROVAL_EXPIRED_ERROR; + runtime.startupTelemetry?.mark('first_response_event_queued'); } try { - await this._onApprovalExpired(streamId, resolvedJob); + await this.eventTransport.emitError(streamId, APPROVAL_EXPIRED_ERROR, createdAt); + if (runtime?.createdAt === createdAt) { + runtime.approvalExpiryPublished = true; + } } catch (err) { - logger.warn(`[GenerationJobManager] Approval-expired cleanup failed for ${streamId}`, err); + logger.error(`[GenerationJobManager] Failed to publish expired approval ${streamId}`, err); + if (runtime?.createdAt === createdAt) { + for (const notify of [...runtime.localErrorHandlers]) { + try { + notify(APPROVAL_EXPIRED_ERROR); + } catch (notifyError) { + logger.error( + `[GenerationJobManager] Failed to notify expired approval ${streamId}`, + notifyError, + ); + } + } + } + } + if (runtime?.createdAt === createdAt) { + this.releaseAbortSubscription(runtime); + runtime.abortController.abort(); } } @@ -2012,23 +3218,8 @@ class GenerationJobManagerClass { // The `errorEvent` flag (set by emitError) keeps this idempotent vs the win path. const runtime = this.runtimeState.get(streamId); if (job?.status === 'aborted' && job.error === APPROVAL_EXPIRED_ERROR) { - if (!runtime?.errorEvent) { - try { - await this.emitError(streamId, APPROVAL_EXPIRED_ERROR); - } catch (err) { - logger.error( - `[GenerationJobManager] Failed to relay expired approval ${streamId}`, - err, - ); - } - } - // The winning store cleanup (`cleanupRequiresActionIndex`) transitions status - // directly and can't run host cleanup — do it on relay. Deliberately NOT gated on - // `errorEvent`: a reconnect seeds that flag from the aborted job, which must not - // suppress the (idempotent) prune. The handler dedups per runtime lifetime, which - // also covers approvals expired LOCALLY via expireApproval. - await this.runApprovalExpiredHandler(streamId, job); - changed = this.runningJobs.delete(streamId) || changed; + await this.notifyApprovalExpiredRuntime(streamId, job.createdAt, runtime); + changed = this.releaseJobOwnership(streamId, job.createdAt) || changed; continue; } if (!job || job.status !== 'requires_action' || !isPendingActionExpired(job)) { @@ -2061,46 +3252,86 @@ class GenerationJobManagerClass { // Cleanup runtime state for deleted jobs for (const [streamId, observedRuntime] of this.runtimeState) { - if (!(await this.jobStore.hasJob(streamId))) { - // A replacement generation can reuse the same streamId while hasJob() - // is in flight. Never reap the replacement runtime based on the stale - // absence observed for its predecessor. - if (this.runtimeState.get(streamId) !== observedRuntime) { - if (!observedRuntime.abortController.signal.aborted) { - observedRuntime.abortController.abort(); - } + const jobExists = await this.jobStore.hasJob(streamId); + if (jobExists) { + const shouldInspectRemoteTerminal = + this.ownedJobs.get(streamId) !== observedRuntime.createdAt && + this.eventTransport.getSubscriberCount(streamId) === 0; + if (!shouldInspectRemoteTerminal) { continue; } - /** - * Abort any still-pending generation whose job has been reaped (e.g. a - * stale "running" job removed by the store's failsafe timeout). This - * unwinds the hung in-flight work so its client/graph references can be - * garbage collected, rather than leaking via the pending promise. - */ + + const currentJob = await this.jobStore.getJob(streamId); + const isRetainedTerminal = + currentJob?.createdAt === observedRuntime.createdAt && + currentJob.status !== 'running' && + currentJob.status !== 'requires_action'; + if (!isRetainedTerminal || this.runtimeState.get(streamId) !== observedRuntime) { + continue; + } + + this.reconcileInactiveGeneration( + streamId, + observedRuntime.createdAt, + currentJob, + observedRuntime, + ); + this.runtimeState.delete(streamId); + this.runStepBuffers?.delete(streamId); + this.replayEventWriteQueues.delete(streamId); + this.tokenUsageWriteQueues.delete(streamId); + this.jobStore.clearContentState(streamId, observedRuntime.createdAt); + this.eventTransport.cleanup(streamId); + continue; + } + + // A replacement generation can reuse the same streamId while hasJob() + // is in flight. Never reap the replacement runtime based on the stale + // absence observed for its predecessor. + if (this.runtimeState.get(streamId) !== observedRuntime) { + this.releaseAbortSubscription(observedRuntime); if (!observedRuntime.abortController.signal.aborted) { observedRuntime.abortController.abort(); } - // If a client is still attached when the job is reaped, send a terminal - // error first so the SSE connection closes instead of hanging open with no - // final/done event (the route only ends the response from onDone/onError). - if (this.eventTransport.getSubscriberCount(streamId) > 0) { - try { - await this.eventTransport.emitError(streamId, REAPED_JOB_ERROR); - } catch (err) { - logger.error(`[GenerationJobManager] Failed to notify reaped stream ${streamId}:`, err); - } - } - // emitError() is asynchronous; a replacement may have appeared while - // the terminal event was being published. - if (this.runtimeState.get(streamId) !== observedRuntime) { - continue; - } - this.runtimeState.delete(streamId); - runningJobsChanged = this.runningJobs.delete(streamId) || runningJobsChanged; - this.runStepBuffers?.delete(streamId); - this.jobStore.clearContentState(streamId); - this.eventTransport.cleanup(streamId); + continue; } + /** + * Abort any still-pending generation whose job has been reaped (e.g. a + * stale "running" job removed by the store's failsafe timeout). This + * unwinds the hung in-flight work so its client/graph references can be + * garbage collected, rather than leaking via the pending promise. + */ + if (!observedRuntime.abortController.signal.aborted) { + observedRuntime.abortController.abort(); + } + // If a client is still attached when the job is reaped, send a terminal + // error first so the SSE connection closes instead of hanging open with no + // final/done event (the route only ends the response from onDone/onError). + if (this.eventTransport.getSubscriberCount(streamId) > 0) { + try { + await this.eventTransport.emitError( + streamId, + REAPED_JOB_ERROR, + observedRuntime.createdAt, + ); + observedRuntime.startupTelemetry?.mark('first_response_event_queued'); + } catch (err) { + logger.error(`[GenerationJobManager] Failed to notify reaped stream ${streamId}:`, err); + } + } + // emitError() is asynchronous; a replacement may have appeared while + // the terminal event was being published. + if (this.runtimeState.get(streamId) !== observedRuntime) { + continue; + } + observedRuntime.startupTelemetry?.end('error', new Error(REAPED_JOB_ERROR)); + observedRuntime.startupTelemetry = undefined; + this.releaseAbortSubscription(observedRuntime); + this.runtimeState.delete(streamId); + runningJobsChanged = this.ownedJobs.delete(streamId) || runningJobsChanged; + this.runStepBuffers?.delete(streamId); + this.jobStore.clearContentState(streamId, observedRuntime.createdAt); + this.eventTransport.cleanup(streamId); } // Also check runStepBuffers for any orphaned entries (Redis mode only) @@ -2143,7 +3374,7 @@ class GenerationJobManagerClass { return null; } - const result = await this.jobStore.getContentParts(streamId); + const result = await this.jobStore.getContentParts(streamId, jobData.createdAt); const aggregatedContent = result?.content ?? []; return { @@ -2200,20 +3431,120 @@ class GenerationJobManagerClass { return this.jobStore.getActiveJobIdsByUser(userId, tenantId); } + private async finalizeOwnedJobsForShutdown(): Promise { + const ownedJobs = [...this.ownedJobs]; + if (ownedJobs.length === 0) { + return; + } + + const completedAt = Date.now(); + const results = await Promise.allSettled( + ownedJobs.map(async ([streamId, createdAt]) => { + const job = await this.jobStore.getJob(streamId); + if (!job || job.createdAt !== createdAt || job.status !== 'running') { + return; + } + const runtime = this.runtimeState.get(streamId); + if ( + runtime?.createdAt === createdAt && + runtime.allSubscribersLeftHandlers?.length && + runtime.lastSubscriberCleanupGeneration !== runtime.attachmentGeneration + ) { + runtime.lastSubscriberCleanupGeneration = runtime.attachmentGeneration; + await this.persistSubscriberCleanup(streamId, runtime); + } + const finalized = await this.jobStore.transitionStatus(streamId, { + from: 'running', + to: 'error', + expectCreatedAt: createdAt, + patch: { completedAt, error: SHUTDOWN_JOB_ERROR }, + }); + if (!finalized) { + return; + } + + if (runtime?.createdAt === createdAt) { + runtime.errorEvent = SHUTDOWN_JOB_ERROR; + this.releaseAbortSubscription(runtime); + } + try { + await this.eventTransport.emitError(streamId, SHUTDOWN_JOB_ERROR, createdAt); + } catch (err) { + logger.error( + `[GenerationJobManager] Failed to publish shutdown error for ${streamId}:`, + err, + ); + } + if (this.ownedJobs.get(streamId) === createdAt) { + this.ownedJobs.delete(streamId); + } + recordGenerationJob(this.storeLabel, 'error'); + }), + ); + + for (let index = 0; index < results.length; index++) { + const result = results[index]; + if (result.status === 'rejected') { + logger.error( + `[GenerationJobManager] Failed to finalize owned job ${ownedJobs[index][0]} during shutdown:`, + result.reason, + ); + } + } + this.syncRunningJobMetrics(); + } + + /** + * Stop accepting jobs and close only this process's attached SSE responses. + * + * This runs before HTTP drain. Durable finalization waits until post-drain, when only jobs + * still owned by this process are atomically moved to a terminal state. + */ + prepareForShutdown(): void { + if (this.shuttingDown) { + return; + } + this.shuttingDown = true; + + for (const runtime of this.runtimeState.values()) { + runtime.startupTelemetry?.end('aborted'); + runtime.startupTelemetry = undefined; + } + + const streamIds = new Set([ + ...this.runtimeState.keys(), + ...this.eventTransport.getTrackedStreamIds(), + ]); + for (const streamId of streamIds) { + this.eventTransport.closeLocalSubscribers?.(streamId, SHUTDOWN_SUBSCRIBER_ERROR); + } + } + /** * Destroy the manager. * Cleans up all resources including runtime state, buffers, and stores. */ async destroy(): Promise { + this.shuttingDown = true; + if (this.cleanupInterval) { clearInterval(this.cleanupInterval); this.cleanupInterval = null; } + for (const runtime of this.runtimeState.values()) { + runtime.startupTelemetry?.end('aborted'); + runtime.startupTelemetry = undefined; + this.releaseAbortSubscription(runtime); + runtime.abortController.abort(); + } + + await this.drainSubscriberCleanups(); + await this.finalizeOwnedJobsForShutdown(); await this.jobStore.destroy(); this.eventTransport.destroy(); this.runtimeState.clear(); - this.runningJobs.clear(); + this.ownedJobs.clear(); this.syncRunningJobMetrics(); this.runStepBuffers?.clear(); this.replayEventWriteQueues.clear(); diff --git a/packages/api/src/stream/SteeringLifecycle.ts b/packages/api/src/stream/SteeringLifecycle.ts index 7603f37bc7..bc3a500239 100644 --- a/packages/api/src/stream/SteeringLifecycle.ts +++ b/packages/api/src/stream/SteeringLifecycle.ts @@ -130,9 +130,9 @@ export class SteeringLifecycle { return this.store.closeAndDrainSteers(streamId, expectedCreatedAt); } - /** Non-destructive FIFO read (status/resume surfaces). */ - peek(streamId: string): Promise { - return this.store.peekSteers(streamId); + /** Non-destructive FIFO read, optionally scoped to one generation. */ + peek(streamId: string, expectedCreatedAt?: number): Promise { + return this.store.peekSteers(streamId, expectedCreatedAt); } /** @@ -160,7 +160,12 @@ export class SteeringLifecycle { * using the final/abort event copy; recovery is idempotent (queued chips * dedupe by steer id). */ - async park(streamId: string, steers: TPendingSteer[], owner: SteerOwner): Promise { + async park( + streamId: string, + steers: TPendingSteer[], + owner: SteerOwner, + expectedCreatedAt?: number, + ): Promise { if (steers.length === 0) { return; } @@ -170,7 +175,7 @@ export class SteeringLifecycle { steers, }; try { - await this.store.parkSteers(streamId, JSON.stringify(payload)); + await this.store.parkSteers(streamId, JSON.stringify(payload), expectedCreatedAt); } catch (error) { logger.warn(`[SteeringLifecycle] Failed to park leftover steers: ${streamId}`, error); } diff --git a/packages/api/src/stream/__tests__/GenerationJobManager.stream_integration.spec.ts b/packages/api/src/stream/__tests__/GenerationJobManager.stream_integration.spec.ts index 3e34d155d5..91f3d656d1 100644 --- a/packages/api/src/stream/__tests__/GenerationJobManager.stream_integration.spec.ts +++ b/packages/api/src/stream/__tests__/GenerationJobManager.stream_integration.spec.ts @@ -1678,7 +1678,7 @@ describe('GenerationJobManager Integration Tests', () => { await manager.destroy(); }); - test('should deliver live events after subscribeWithResume', async () => { + test('should defer live events until a resumed subscription is activated', async () => { const manager = createInMemoryManager(); const streamId = `atomic-live-${Date.now()}`; await manager.createJob(streamId, 'user-1'); @@ -1707,6 +1707,9 @@ describe('GenerationJobManager Integration Tests', () => { }); await new Promise((resolve) => setTimeout(resolve, 20)); + expect(liveEvents.length).toBe(0); + + subscription?.activate(); expect(liveEvents.length).toBe(1); const liveEvent = liveEvents[0] as { event: string; @@ -1751,6 +1754,9 @@ describe('GenerationJobManager Integration Tests', () => { }); await new Promise((resolve) => setTimeout(resolve, 200)); + expect(liveEvents.length).toBe(0); + + subscription?.activate(); expect(liveEvents.length).toBe(1); subscription?.unsubscribe(); @@ -2027,10 +2033,6 @@ describe('GenerationJobManager Integration Tests', () => { const streamId = `cross-live-${Date.now()}`; await replicaA.createJob(streamId, 'user-1'); - const replicaBJobStore = new RedisJobStore(ioredisClient!); - await replicaBJobStore.initialize(); - await replicaBJobStore.createJob(streamId, 'user-1'); - const receivedOnB: unknown[] = []; const subB = await replicaB.subscribe(streamId, (event: unknown) => receivedOnB.push(event)); @@ -2051,7 +2053,6 @@ describe('GenerationJobManager Integration Tests', () => { } subB?.unsubscribe(); - replicaBJobStore.destroy(); await replicaA.destroy(); await replicaB.destroy(); }); @@ -2076,9 +2077,6 @@ describe('GenerationJobManager Integration Tests', () => { const streamId = `cross-seq-safe-${Date.now()}`; await replicaA.createJob(streamId, 'user-1'); - const replicaBJobStore = new RedisJobStore(ioredisClient!); - await replicaBJobStore.initialize(); - await replicaBJobStore.createJob(streamId, 'user-1'); const receivedOnB: unknown[] = []; const subB = await replicaB.subscribe(streamId, (event: unknown) => receivedOnB.push(event)); @@ -2130,7 +2128,6 @@ describe('GenerationJobManager Integration Tests', () => { subA?.unsubscribe(); subB?.unsubscribe(); - replicaBJobStore.destroy(); await replicaA.destroy(); await replicaB.destroy(); }); @@ -2169,10 +2166,6 @@ describe('GenerationJobManager Integration Tests', () => { replicaB.configure(servicesB); replicaB.initialize(); - const replicaBJobStore = new RedisJobStore(ioredisClient!); - await replicaBJobStore.initialize(); - await replicaBJobStore.createJob(streamId, 'user-1'); - const receivedOnB: unknown[] = []; const subB = await replicaB.subscribe(streamId, (event: unknown) => receivedOnB.push(event)); @@ -2194,7 +2187,6 @@ describe('GenerationJobManager Integration Tests', () => { subA?.unsubscribe(); subB?.unsubscribe(); - replicaBJobStore.destroy(); await replicaA.destroy(); await replicaB.destroy(); }); @@ -2376,7 +2368,7 @@ describe('GenerationJobManager Integration Tests', () => { onDone: () => {}, }); - await sub1.ready; + await expect(sub1.ready).rejects.toThrow('Simulated Redis SUBSCRIBE failure'); const receivedEvents: unknown[] = []; sub1.unsubscribe(); @@ -2388,6 +2380,7 @@ describe('GenerationJobManager Integration Tests', () => { expect(sub2.ready).toBeDefined(); await sub2.ready; + expect(callCount).toBe(2); await transport.emitChunk(streamId, { event: 'test', data: { value: 'hello' } }); await new Promise((resolve) => setTimeout(resolve, 100)); diff --git a/packages/api/src/stream/__tests__/RedisEventTransport.spec.ts b/packages/api/src/stream/__tests__/RedisEventTransport.spec.ts index fa238b6628..e02b3e0249 100644 --- a/packages/api/src/stream/__tests__/RedisEventTransport.spec.ts +++ b/packages/api/src/stream/__tests__/RedisEventTransport.spec.ts @@ -1,6 +1,10 @@ -import type { Redis } from 'ioredis'; import { logger } from '@librechat/data-schemas'; +import type { Redis } from 'ioredis'; +import { InMemoryEventTransport } from '~/stream/implementations/InMemoryEventTransport'; import { RedisEventTransport } from '~/stream/implementations/RedisEventTransport'; +import { InMemoryJobStore } from '~/stream/implementations/InMemoryJobStore'; +import { emitChunkWithReceipt } from '~/stream/internal/chunkPublication'; +import { GenerationJobManagerClass } from '~/stream/GenerationJobManager'; import { createMockPublisher } from './helpers/publisher'; logger.silent = true; @@ -20,7 +24,1268 @@ function getMessageHandler(mockSubscriber: ReturnType void; } +interface SequencedTestMessage { + type: 'chunk' | 'done' | 'error'; + seq: number; + data?: object; + error?: string; + generationId?: number; +} + +function deliverSequencedMessage( + handler: ReturnType, + streamId: string, + message: SequencedTestMessage, +): void { + handler(`stream:{${streamId}}:events`, JSON.stringify(message)); +} + describe('RedisEventTransport', () => { + it('delivers tagged events and preserves legacy in-memory callback arity', () => { + const transport = new InMemoryEventTransport(); + const onChunk = jest.fn(); + const onDone = jest.fn(); + const onError = jest.fn(); + const streamId = 'in-memory-terminal-identity'; + const subscription = transport.subscribe(streamId, { + onChunk, + onDone, + onError, + }); + const taggedChunk = { delta: 'tagged' }; + const legacyChunk = { delta: 'legacy' }; + const taggedDone = { final: 'tagged' }; + const legacyDone = { final: 'legacy' }; + + transport.emitChunk(streamId, taggedChunk, 123456); + transport.emitChunk(streamId, legacyChunk); + transport.emitDone(streamId, taggedDone, 123456); + transport.emitDone(streamId, legacyDone); + transport.emitError(streamId, 'tagged error', 123456); + transport.emitError(streamId, 'legacy error'); + + expect(onChunk).toHaveBeenNthCalledWith(1, taggedChunk, 123456); + expect(onChunk).toHaveBeenNthCalledWith(2, legacyChunk); + expect(onDone).toHaveBeenNthCalledWith(1, taggedDone, 123456); + expect(onDone).toHaveBeenNthCalledWith(2, legacyDone); + expect(onError).toHaveBeenNthCalledWith(1, 'tagged error', 123456); + expect(onError).toHaveBeenNthCalledWith(2, 'legacy error'); + + subscription.unsubscribe(); + transport.destroy(); + }); + + it('round-trips tagged events and preserves legacy Redis callback arity', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const streamId = 'redis-terminal-identity'; + const onChunk = jest.fn(); + const onDone = jest.fn(); + const onError = jest.fn(); + const onAbort = jest.fn(); + const subscription = transport.subscribe(streamId, { + onChunk, + onDone, + onError, + }); + const messageHandler = getMessageHandler(mockSubscriber); + mockPublisher.publish.mockImplementation(async (channel: string, payload: string) => { + messageHandler(channel, payload); + return 1; + }); + + await subscription.ready; + await transport.onAbort(streamId, onAbort); + const taggedChunk = { delta: 'tagged' }; + const legacyChunk = { delta: 'legacy' }; + const taggedDone = { final: 'tagged' }; + const legacyDone = { final: 'legacy' }; + await transport.emitChunk(streamId, taggedChunk, 654321); + await transport.emitChunk(streamId, legacyChunk); + await transport.emitDone(streamId, taggedDone, 654321); + await transport.emitDone(streamId, legacyDone); + await transport.emitError(streamId, 'tagged error', 654321); + await transport.emitError(streamId, 'legacy error'); + transport.emitAbort(streamId, 654321); + transport.emitAbort(streamId); + + expect(onChunk).toHaveBeenNthCalledWith(1, taggedChunk, 654321); + expect(onChunk).toHaveBeenNthCalledWith(2, legacyChunk); + expect(onDone).toHaveBeenNthCalledWith(1, taggedDone, 654321); + expect(onDone).toHaveBeenNthCalledWith(2, legacyDone); + expect(onError).toHaveBeenNthCalledWith(1, 'tagged error', 654321); + expect(onError).toHaveBeenNthCalledWith(2, 'legacy error'); + expect(onAbort).toHaveBeenNthCalledWith(1, 654321); + expect(onAbort).toHaveBeenNthCalledWith(2); + + subscription.unsubscribe(); + transport.destroy(); + }); + + it('disposes only the owning abort callback and releases an unused channel', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const streamId = 'abort-registration-disposal'; + const firstAbort = jest.fn(); + const secondAbort = jest.fn(); + const disposeFirst = await transport.onAbort(streamId, firstAbort); + const disposeSecond = await transport.onAbort(streamId, secondAbort); + const messageHandler = getMessageHandler(mockSubscriber); + const channel = `stream:{${streamId}}:events`; + + disposeFirst(); + messageHandler(channel, JSON.stringify({ type: 'abort', generationId: 2 })); + + expect(firstAbort).not.toHaveBeenCalled(); + expect(secondAbort).toHaveBeenCalledWith(2); + expect(mockSubscriber.unsubscribe).not.toHaveBeenCalled(); + + disposeFirst(); + disposeSecond(); + + expect(mockSubscriber.unsubscribe).toHaveBeenCalledTimes(1); + expect(mockSubscriber.unsubscribe).toHaveBeenCalledWith(channel); + + transport.destroy(); + }); + + it('releases each generation abort subscription after successful completion', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const manager = new GenerationJobManagerClass({ + jobStore: new InMemoryJobStore(), + eventTransport: transport, + }); + const streamId = 'completed-generation-abort-cleanup'; + const channel = `stream:{${streamId}}:events`; + + const first = await manager.createJob(streamId, 'user-1'); + await manager.completeJob(streamId, undefined, first.createdAt); + + expect(mockSubscriber.unsubscribe).toHaveBeenNthCalledWith(1, channel); + + const second = await manager.createJob(streamId, 'user-1'); + await manager.completeJob(streamId, undefined, second.createdAt); + + expect(mockSubscriber.subscribe).toHaveBeenCalledTimes(2); + expect(mockSubscriber.unsubscribe).toHaveBeenCalledTimes(2); + expect(mockSubscriber.unsubscribe).toHaveBeenNthCalledWith(2, channel); + + await manager.destroy(); + }); + + it('does not register an abort listener for a lazily loaded terminal job', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const terminalJob = await jobStore.createJob('lazy-terminal-job', 'user-1'); + await jobStore.transitionStatus('lazy-terminal-job', { + from: 'running', + to: 'complete', + expectCreatedAt: terminalJob.createdAt, + patch: { completedAt: Date.now() }, + }); + const manager = new GenerationJobManagerClass({ + jobStore, + eventTransport: transport, + cleanupOnComplete: false, + }); + + await expect(manager.getJob('lazy-terminal-job')).resolves.toBeDefined(); + expect(mockSubscriber.subscribe).not.toHaveBeenCalled(); + + await manager.destroy(); + }); + + it('releases an equal-epoch lazy runtime when a later lookup observes it terminal', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const streamId = 'lazy-equal-epoch-terminal'; + const durableJob = await jobStore.createJob(streamId, 'user-1'); + const manager = new GenerationJobManagerClass({ + jobStore, + eventTransport: transport, + cleanupOnComplete: false, + }); + const lazyJob = await manager.getJob(streamId); + + await jobStore.transitionStatus(streamId, { + from: 'running', + to: 'error', + expectCreatedAt: durableJob.createdAt, + patch: { completedAt: Date.now(), error: 'remote terminal' }, + }); + await expect(manager.getJob(streamId)).resolves.toMatchObject({ + createdAt: durableJob.createdAt, + status: 'error', + }); + + expect(lazyJob?.abortController.signal.aborted).toBe(true); + expect(mockSubscriber.unsubscribe).toHaveBeenCalledWith(`stream:{${streamId}}:events`); + + await manager.destroy(); + }); + + it('releases a lazy abort runtime when cleanup observes a retained terminal job', async () => { + jest.useFakeTimers(); + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 3_600_000 }); + const streamId = 'lazy-remote-terminal-cleanup'; + const durableJob = await jobStore.createJob(streamId, 'user-1'); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport: transport, + isRedis: true, + cleanupOnComplete: false, + }); + manager.initialize(); + + try { + const lazyJob = await manager.getJob(streamId); + expect(lazyJob?.abortController.signal.aborted).toBe(false); + expect(mockSubscriber.subscribe).toHaveBeenCalledTimes(1); + + await jobStore.transitionStatus(streamId, { + from: 'running', + to: 'complete', + expectCreatedAt: durableJob.createdAt, + patch: { completedAt: Date.now() }, + }); + await jest.advanceTimersByTimeAsync(60_000); + + expect(lazyJob?.abortController.signal.aborted).toBe(true); + expect(mockSubscriber.unsubscribe).toHaveBeenCalledWith(`stream:{${streamId}}:events`); + expect(manager.getRuntimeStats().runtimeStateSize).toBe(0); + await expect(jobStore.getJob(streamId)).resolves.toMatchObject({ + createdAt: durableJob.createdAt, + status: 'complete', + }); + } finally { + await manager.destroy(); + jest.useRealTimers(); + } + }); + + it('releases a subscriber-only abort listener when tagged terminal delivery arrives', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const streamId = 'subscriber-only-terminal-cleanup'; + const durableJob = await jobStore.createJob(streamId, 'user-1'); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport: transport, + isRedis: true, + cleanupOnComplete: false, + }); + const onDone = jest.fn(); + const subscription = await manager.subscribe(streamId, () => undefined, onDone); + const channel = `stream:{${streamId}}:events`; + + deliverSequencedMessage(getMessageHandler(mockSubscriber), streamId, { + type: 'done', + seq: 0, + data: { final: true }, + generationId: durableJob.createdAt, + }); + + expect(onDone).toHaveBeenCalledWith({ final: true }); + expect(mockSubscriber.unsubscribe).toHaveBeenCalledTimes(1); + expect(mockSubscriber.unsubscribe).toHaveBeenCalledWith(channel); + + subscription?.unsubscribe(); + await manager.destroy(); + }); + + it('filters tagged predecessor chunks from a current generation subscription', async () => { + const now = jest.spyOn(Date, 'now').mockReturnValue(200); + const transport = new InMemoryEventTransport(); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport: transport, + isRedis: false, + }); + manager.initialize(); + + try { + const streamId = 'manager-chunk-generation-filter'; + const job = await manager.createJob(streamId, 'user-1'); + const received: unknown[] = []; + const subscription = await manager.subscribe(streamId, (event) => received.push(event)); + const stale = { event: 'on_message_delta', data: { text: 'stale' } }; + const current = { event: 'on_message_delta', data: { text: 'current' } }; + const legacy = { event: 'on_message_delta', data: { text: 'legacy' } }; + + transport.emitChunk(streamId, stale, job.createdAt - 1); + transport.emitChunk(streamId, current, job.createdAt); + transport.emitChunk(streamId, legacy); + + expect(received).toEqual([current, legacy]); + subscription?.unsubscribe(); + } finally { + now.mockRestore(); + await manager.destroy(); + } + }); + + it('defers ordered delivery until the replay frontier is synchronized', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const streamId = 'deferred-until-sync'; + const received: object[] = []; + const subscription = transport.subscribe( + streamId, + { onChunk: (event) => received.push(event as object) }, + { deferDeliveryUntilSynchronized: true }, + ); + + deliverSequencedMessage(getMessageHandler(mockSubscriber), streamId, { + type: 'chunk', + seq: 0, + data: { index: 0 }, + }); + + expect(received).toEqual([]); + + await transport.syncReorderBuffer(streamId); + + expect(received).toEqual([{ index: 0 }]); + + subscription.unsubscribe(); + transport.destroy(); + }); + + it('keeps the attachment fence closed across reorder timeout and buffer pressure', async () => { + jest.useFakeTimers(); + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const streamId = 'deferred-timeout-and-overflow'; + const received: object[] = []; + const subscription = transport.subscribe( + streamId, + { onChunk: (event) => received.push(event as object) }, + { deferDeliveryUntilSynchronized: true }, + ); + const messageHandler = getMessageHandler(mockSubscriber); + + try { + deliverSequencedMessage(messageHandler, streamId, { + type: 'chunk', + seq: 0, + data: { index: 0 }, + }); + await jest.advanceTimersByTimeAsync(501); + expect(received).toEqual([]); + + for (let i = 1; i < 100; i++) { + deliverSequencedMessage(messageHandler, streamId, { + type: 'chunk', + seq: i, + data: { index: i }, + }); + } + expect(received).toEqual([]); + + await transport.syncReorderBuffer(streamId); + expect(received).toHaveLength(100); + expect(received[0]).toEqual({ index: 0 }); + expect(received[99]).toEqual({ index: 99 }); + } finally { + subscription.unsubscribe(); + transport.destroy(); + jest.useRealTimers(); + } + }); + + it('waits at a same-replica frontier when the following sequence arrives first', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const streamId = 'same-replica-frontier-gap'; + const received: object[] = []; + const subscription = transport.subscribe( + streamId, + { onChunk: (event) => received.push(event as object) }, + { deferDeliveryUntilSynchronized: true }, + ); + const messageHandler = getMessageHandler(mockSubscriber); + mockPublisher.get.mockResolvedValueOnce('7'); + + deliverSequencedMessage(messageHandler, streamId, { + type: 'chunk', + seq: 6, + data: { index: 6 }, + }); + await transport.syncReorderBuffer(streamId, 5); + + expect(mockPublisher.get).toHaveBeenCalledWith(`stream:{${streamId}}:seq`); + expect(received).toEqual([]); + + deliverSequencedMessage(messageHandler, streamId, { + type: 'chunk', + seq: 5, + data: { index: 5 }, + }); + + expect(received).toEqual([{ index: 5 }, { index: 6 }]); + + subscription.unsubscribe(); + transport.destroy(); + }); + + it('uses the Redis sequence as the cross-replica attachment fence', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const streamId = 'cross-replica-attachment-fence'; + const received: object[] = []; + const subscription = transport.subscribe( + streamId, + { onChunk: (event) => received.push(event as object) }, + { deferDeliveryUntilSynchronized: true }, + ); + const messageHandler = getMessageHandler(mockSubscriber); + mockPublisher.get.mockResolvedValueOnce('7'); + + try { + deliverSequencedMessage(messageHandler, streamId, { + type: 'chunk', + seq: 6, + data: { index: 6 }, + }); + await transport.syncReorderBuffer(streamId); + + expect(mockPublisher.get).toHaveBeenCalledWith(`stream:{${streamId}}:seq`); + expect(received).toEqual([{ index: 6 }]); + + deliverSequencedMessage(messageHandler, streamId, { + type: 'chunk', + seq: 7, + data: { index: 7 }, + }); + deliverSequencedMessage(messageHandler, streamId, { + type: 'chunk', + seq: 5, + data: { index: 5 }, + }); + + expect(received).toEqual([{ index: 6 }, { index: 7 }]); + } finally { + subscription.unsubscribe(); + transport.destroy(); + } + }); + + it('adopts the Redis frontier after an empty cross-replica sync', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const streamId = 'cross-replica-empty-sync'; + const received: object[] = []; + const subscription = transport.subscribe( + streamId, + { onChunk: (event) => received.push(event as object) }, + { deferDeliveryUntilSynchronized: true }, + ); + mockPublisher.get.mockResolvedValueOnce('6'); + await transport.syncReorderBuffer(streamId); + expect(mockPublisher.get).toHaveBeenCalledWith(`stream:{${streamId}}:seq`); + + deliverSequencedMessage(getMessageHandler(mockSubscriber), streamId, { + type: 'chunk', + seq: 6, + data: { index: 6 }, + }); + expect(received).toEqual([{ index: 6 }]); + + deliverSequencedMessage(getMessageHandler(mockSubscriber), streamId, { + type: 'chunk', + seq: 7, + data: { index: 7 }, + }); + expect(received).toEqual([{ index: 6 }, { index: 7 }]); + + subscription.unsubscribe(); + transport.destroy(); + }); + + it('holds terminal events behind earlier chunks until synchronization', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const streamId = 'terminal-deferred-until-sync'; + const received: string[] = []; + const subscription = transport.subscribe( + streamId, + { + onChunk: (event) => received.push(`chunk:${(event as { index: number }).index}`), + onDone: () => received.push('done'), + }, + { deferDeliveryUntilSynchronized: true }, + ); + const messageHandler = getMessageHandler(mockSubscriber); + + deliverSequencedMessage(messageHandler, streamId, { + type: 'done', + seq: 1, + data: { final: true }, + }); + deliverSequencedMessage(messageHandler, streamId, { + type: 'chunk', + seq: 0, + data: { index: 0 }, + }); + + expect(received).toEqual([]); + + await transport.syncReorderBuffer(streamId); + + expect(mockPublisher.get).toHaveBeenCalledWith(`stream:{${streamId}}:seq`); + expect(received).toEqual(['chunk:0', 'done']); + + subscription.unsubscribe(); + transport.destroy(); + }); + + it('closes a snapshot of local subscribers without publishing', () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const streamId = 'close-local-subscribers'; + const secondOnError = jest.fn(); + const onAllSubscribersLeft = jest.fn(); + transport.onAllSubscribersLeft(streamId, onAllSubscribersLeft); + let unsubscribeSecond = (): void => undefined; + const firstSubscription = transport.subscribe(streamId, { + onChunk: () => undefined, + onError: () => { + unsubscribeSecond(); + throw new Error('first handler failed'); + }, + }); + const secondSubscription = transport.subscribe(streamId, { + onChunk: () => undefined, + onError: secondOnError, + }); + unsubscribeSecond = secondSubscription.unsubscribe; + + expect(() => transport.closeLocalSubscribers(streamId, 'stream closed')).not.toThrow(); + expect(secondOnError).toHaveBeenCalledWith('stream closed'); + expect(transport.getSubscriberCount(streamId)).toBe(0); + expect(onAllSubscribersLeft).toHaveBeenCalledTimes(1); + expect(mockPublisher.publish).not.toHaveBeenCalled(); + expect(mockPublisher.eval).not.toHaveBeenCalled(); + + firstSubscription.unsubscribe(); + transport.destroy(); + }); + + it('replaces the disconnect lifecycle callback for a replacement runtime', () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const streamId = 'replacement-disconnect-callback'; + const replacedRuntimeCallback = jest.fn(); + const currentRuntimeCallback = jest.fn(); + transport.onAllSubscribersLeft(streamId, replacedRuntimeCallback); + transport.onAllSubscribersLeft(streamId, currentRuntimeCallback); + const subscription = transport.subscribe(streamId, { + onChunk: () => undefined, + }); + + subscription.unsubscribe(); + + expect(replacedRuntimeCallback).not.toHaveBeenCalled(); + expect(currentRuntimeCallback).toHaveBeenCalledTimes(1); + transport.destroy(); + }); + + it('waits for the cross-replica abort channel before resolving job creation', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + let signalSubscriptionStarted: (() => void) | undefined; + const subscriptionStarted = new Promise((resolve) => { + signalSubscriptionStarted = resolve; + }); + let releaseSubscription: (() => void) | undefined; + const subscriptionGate = new Promise((resolve) => { + releaseSubscription = resolve; + }); + mockSubscriber.subscribe.mockImplementationOnce(() => { + signalSubscriptionStarted?.(); + return subscriptionGate; + }); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport: transport, + isRedis: true, + }); + + let createResolved = false; + const creating = manager.createJob('abort-readiness', 'user-1').then((job) => { + createResolved = true; + return job; + }); + + await subscriptionStarted; + await Promise.resolve(); + + expect(createResolved).toBe(false); + + releaseSubscription?.(); + await creating; + + expect(createResolved).toBe(true); + await manager.destroy(); + }); + + it('releases a registration that loses initialization without detaching its replacement', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + let signalSubscriptionStarted: (() => void) | undefined; + const subscriptionStarted = new Promise((resolve) => { + signalSubscriptionStarted = resolve; + }); + let releaseSubscription: (() => void) | undefined; + const subscriptionGate = new Promise((resolve) => { + releaseSubscription = resolve; + }); + mockSubscriber.subscribe.mockImplementationOnce(() => { + signalSubscriptionStarted?.(); + return subscriptionGate; + }); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const manager = new GenerationJobManagerClass({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport: transport, + }); + const streamId = 'abort-registration-replacement-race'; + const predecessorCreation = manager.createJob(streamId, 'user-1'); + + await subscriptionStarted; + + const replacementCreation = manager.createJob(streamId, 'user-1'); + await Promise.resolve(); + releaseSubscription?.(); + + await expect(predecessorCreation).rejects.toThrow( + 'Generation job was replaced during initialization', + ); + const replacement = await replacementCreation; + + expect(mockSubscriber.unsubscribe).not.toHaveBeenCalled(); + + getMessageHandler(mockSubscriber)( + `stream:{${streamId}}:events`, + JSON.stringify({ type: 'abort', generationId: replacement.createdAt }), + ); + + expect(replacement.abortController.signal.aborted).toBe(true); + expect(mockSubscriber.unsubscribe).toHaveBeenCalledTimes(1); + + await manager.destroy(); + }); + + it('keeps remote abort active after SSE disconnect and releases it after delivery', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport: transport, + isRedis: true, + }); + const streamId = 'remote-abort-after-disconnect'; + const job = await manager.createJob(streamId, 'user-1'); + const subscription = await manager.subscribe(streamId, () => undefined); + + subscription?.unsubscribe(); + + expect(mockSubscriber.unsubscribe).not.toHaveBeenCalled(); + + getMessageHandler(mockSubscriber)( + `stream:{${streamId}}:events`, + JSON.stringify({ type: 'abort' }), + ); + + expect(job.abortController.signal.aborted).toBe(true); + expect(mockSubscriber.unsubscribe).toHaveBeenCalledWith(`stream:{${streamId}}:events`); + + await manager.destroy(); + }); + + it('detaches a manager subscription when shutdown starts during transport readiness', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + let signalReadyStarted: (() => void) | undefined; + const readyStarted = new Promise((resolve) => { + signalReadyStarted = resolve; + }); + let releaseReady: (() => void) | undefined; + const readyGate = new Promise((resolve) => { + releaseReady = resolve; + }); + mockSubscriber.subscribe.mockImplementationOnce(() => { + signalReadyStarted?.(); + return readyGate; + }); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + Object.defineProperty(transport, 'onAbort', { value: undefined }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport: transport, + isRedis: true, + }); + manager.initialize(); + await manager.createJob('shutdown-during-ready', 'user-1'); + + const onError = jest.fn(); + const subscribing = manager.subscribe( + 'shutdown-during-ready', + () => undefined, + undefined, + onError, + ); + await readyStarted; + + manager.prepareForShutdown(); + + expect(onError).toHaveBeenCalledWith('Server is shutting down'); + expect(transport.getSubscriberCount('shutdown-during-ready')).toBe(0); + + releaseReady?.(); + await expect(subscribing).resolves.toBeNull(); + await manager.destroy(); + }); + + it('keeps publication receipts behind the internal transport capability', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + + await expect(transport.emitChunk('stream-1', { text: 'Hello' })).resolves.toBeUndefined(); + await expect(emitChunkWithReceipt(transport, 'stream-1', { text: 'World' }, 777)).resolves.toBe( + 1, + ); + const guardedPublish = mockPublisher.eval.mock.calls[1]; + expect(guardedPublish[0]).toContain('local currentCreatedAt = redis.call("HGET", KEYS[2]'); + expect(guardedPublish[9]).toBe('777'); + expect(guardedPublish[10]).toBe('0'); + expect(JSON.parse(`${guardedPublish[6]}1${guardedPublish[7]}`)).toMatchObject({ + type: 'chunk', + data: { text: 'World' }, + generationId: 777, + }); + + mockPublisher.eval.mockResolvedValueOnce(-1); + await expect( + emitChunkWithReceipt(transport, 'replaced-stream', { text: 'stale' }, 111), + ).resolves.toBe(false); + mockPublisher.eval.mockRejectedValue(new Error('publish failed')); + await expect(transport.emitChunk('failed-stream', { text: 'Hello' })).resolves.toBeUndefined(); + await expect(emitChunkWithReceipt(transport, 'failed-stream', { text: 'Hello' })).resolves.toBe( + false, + ); + + transport.destroy(); + }); + + it('guards terminal publications with the winning finalized epoch', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + + await transport.emitDone('terminal-guard', { final: true }, 777); + await transport.emitError('terminal-guard', 'failed', 777); + await emitChunkWithReceipt(transport, 'terminal-guard', { text: 'delta' }, 777); + + const [donePublish, errorPublish, chunkPublish] = mockPublisher.eval.mock.calls; + expect(donePublish[0]).toContain( + 'if redis.call("EXISTS", KEYS[2]) == 1 or ARGV[6] ~= "1" then return -1 end', + ); + expect(donePublish[0]).toContain( + 'redis.call("SET", KEYS[3], ARGV[5], "EX", tonumber(ARGV[7]), "NX")', + ); + expect(donePublish[0]).not.toContain('redis.call("DEL", KEYS[3])'); + expect(donePublish[4]).toBe('stream:{terminal-guard}:generation-epoch'); + expect(donePublish[9]).toBe('777'); + expect(donePublish[10]).toBe('1'); + expect(donePublish[11]).toBe('300'); + expect(errorPublish[9]).toBe('777'); + expect(errorPublish[10]).toBe('1'); + expect(errorPublish[11]).toBe('300'); + expect(chunkPublish[9]).toBe('777'); + expect(chunkPublish[10]).toBe('0'); + expect(chunkPublish[11]).toBe('300'); + + transport.destroy(); + }); + + it('replays an event exactly once when its Redis publish resolves during attachment', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const streamId = 'publish-resolves-during-attachment'; + const messageHandler = getMessageHandler(mockSubscriber); + mockPublisher.publish.mockImplementation(async (channel: string, payload: string) => { + if (transport.getSubscriberCount(streamId) > 0) { + messageHandler(channel, payload); + } + return 1; + }); + + const originalEval = mockPublisher.eval.getMockImplementation(); + let releasePublication: (() => void) | undefined; + const publicationGate = new Promise((resolve) => { + releasePublication = resolve; + }); + mockPublisher.eval.mockImplementationOnce(async (...args: unknown[]) => { + await publicationGate; + return originalEval?.(...args); + }); + + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport: transport, + isRedis: true, + }); + manager.initialize(); + await manager.createJob(streamId, 'user-1'); + + const earlyEvent = { + event: 'on_message_delta' as const, + data: { delta: { content: { type: 'text', text: 'early' } } }, + }; + const publication = manager.emitChunk(streamId, earlyEvent); + await Promise.resolve(); + + const received: unknown[] = []; + const attachment = manager.subscribe(streamId, (event) => received.push(event)); + let attachmentSettled = false; + void attachment.then(() => { + attachmentSettled = true; + }); + await new Promise((resolve) => setTimeout(resolve, 550)); + + expect(attachmentSettled).toBe(false); + expect(received).toEqual([]); + + releasePublication?.(); + + await publication; + const subscription = await attachment; + + expect(received).toEqual([earlyEvent]); + + const liveEvent = { + event: 'on_message_delta' as const, + data: { delta: { content: { type: 'text', text: 'live' } } }, + }; + await manager.emitChunk(streamId, liveEvent); + expect(received).toEqual([earlyEvent, liveEvent]); + + subscription?.unsubscribe(); + await manager.destroy(); + }); + + it('hands a canceled resume bootstrap to a surviving initial subscriber', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const streamId = 'abort-during-publication-fence'; + const messageHandler = getMessageHandler(mockSubscriber); + mockPublisher.publish.mockImplementation(async (channel: string, payload: string) => { + messageHandler(channel, payload); + return 1; + }); + const originalEval = mockPublisher.eval.getMockImplementation(); + let signalPublicationStarted: (() => void) | undefined; + const publicationStarted = new Promise((resolve) => { + signalPublicationStarted = resolve; + }); + let releasePublication: (() => void) | undefined; + const publicationGate = new Promise((resolve) => { + releasePublication = resolve; + }); + mockPublisher.eval.mockImplementationOnce(async (...args: unknown[]) => { + signalPublicationStarted?.(); + await publicationGate; + return originalEval?.(...args); + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport: transport, + isRedis: true, + }); + manager.initialize(); + await manager.createJob(streamId, 'user-1'); + + const earlyEvent = { + event: 'on_message_delta', + data: { delta: { content: [{ type: 'text', text: 'early' }] } }, + } as const; + const publication = manager.emitChunk(streamId, earlyEvent); + await publicationStarted; + + const attachmentAbortController = new AbortController(); + const subscribing = manager.subscribe(streamId, () => undefined, undefined, undefined, { + skipBufferReplay: true, + signal: attachmentAbortController.signal, + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(transport.getSubscriberCount(streamId)).toBe(1); + + const receivedBySurvivor: unknown[] = []; + const survivingSubscription = await manager.subscribe(streamId, (event) => + receivedBySurvivor.push(event), + ); + expect(transport.getSubscriberCount(streamId)).toBe(2); + + attachmentAbortController.abort(); + + await expect(subscribing).resolves.toBeNull(); + expect(transport.getSubscriberCount(streamId)).toBe(1); + + releasePublication?.(); + await publication; + await new Promise((resolve) => setImmediate(resolve)); + + expect(receivedBySurvivor).toEqual([earlyEvent]); + + const liveEvent = { + event: 'on_message_delta' as const, + data: { delta: { content: [{ type: 'text', text: 'live' }] } }, + }; + await manager.emitChunk(streamId, liveEvent); + expect(receivedBySurvivor).toEqual([earlyEvent, liveEvent]); + + survivingSubscription?.unsubscribe(); + await manager.destroy(); + }); + + it('deduplicates a cross-replica created fallback when the original publishes later', async () => { + const mockPublisher = createMockPublisher(); + const generatingSubscriber = createMockSubscriber(); + const attachingSubscriber = createMockSubscriber(); + const generatingTransport = new RedisEventTransport( + mockPublisher as unknown as Redis, + generatingSubscriber as unknown as Redis, + ); + const attachingTransport = new RedisEventTransport( + mockPublisher as unknown as Redis, + attachingSubscriber as unknown as Redis, + ); + const attachingMessageHandler = getMessageHandler(attachingSubscriber); + const streamId = 'cross-replica-created-fallback'; + mockPublisher.publish.mockImplementation(async (channel: string, payload: string) => { + if (attachingTransport.getSubscriberCount(streamId) > 0) { + attachingMessageHandler(channel, payload); + } + return 1; + }); + + const originalEval = mockPublisher.eval.getMockImplementation(); + let signalPublicationStarted: (() => void) | undefined; + const publicationStarted = new Promise((resolve) => { + signalPublicationStarted = resolve; + }); + let releasePublication: (() => void) | undefined; + const publicationGate = new Promise((resolve) => { + releasePublication = resolve; + }); + mockPublisher.eval.mockImplementationOnce(async (...args: unknown[]) => { + signalPublicationStarted?.(); + await publicationGate; + return originalEval?.(...args); + }); + + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const generatingManager = new GenerationJobManagerClass(); + generatingManager.configure({ + jobStore, + eventTransport: generatingTransport, + isRedis: true, + }); + const attachingManager = new GenerationJobManagerClass(); + attachingManager.configure({ + jobStore, + eventTransport: attachingTransport, + isRedis: true, + }); + await generatingManager.createJob(streamId, 'user-1', streamId); + + const createdEvent = { + created: true as const, + message: { + messageId: 'message-1', + conversationId: streamId, + text: 'Hello', + sender: 'User', + isCreatedByUser: true, + }, + streamId, + }; + const publication = generatingManager.emitChunk(streamId, createdEvent); + await publicationStarted; + + const received: unknown[] = []; + const subscription = await attachingManager.subscribe(streamId, (event) => + received.push(event), + ); + expect(received).toEqual([createdEvent]); + + releasePublication?.(); + await publication; + await Promise.resolve(); + + expect(received).toEqual([createdEvent]); + + subscription?.unsubscribe(); + await generatingManager.destroy(); + await attachingManager.destroy(); + }); + + it('reconstructs a missed cross-replica created event before a pending delta', async () => { + const mockPublisher = createMockPublisher(); + const generatingSubscriber = createMockSubscriber(); + const attachingSubscriber = createMockSubscriber(); + const generatingTransport = new RedisEventTransport( + mockPublisher as unknown as Redis, + generatingSubscriber as unknown as Redis, + ); + const attachingTransport = new RedisEventTransport( + mockPublisher as unknown as Redis, + attachingSubscriber as unknown as Redis, + ); + const streamId = 'cross-replica-created-before-delta'; + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const generatingManager = new GenerationJobManagerClass(); + generatingManager.configure({ + jobStore, + eventTransport: generatingTransport, + isRedis: true, + }); + const attachingManager = new GenerationJobManagerClass(); + attachingManager.configure({ + jobStore, + eventTransport: attachingTransport, + isRedis: true, + }); + await generatingManager.createJob(streamId, 'user-1', streamId); + + const createdEvent = { + created: true as const, + message: { + messageId: 'message-1', + conversationId: streamId, + text: 'Hello', + sender: 'User', + isCreatedByUser: true, + }, + streamId, + }; + await generatingManager.emitChunk(streamId, createdEvent); + + const deltaEvent = { + event: 'on_message_delta' as const, + data: { delta: { content: [{ type: 'text', text: 'World' }] } }, + }; + const messageHandler = getMessageHandler(attachingSubscriber); + const originalSync = attachingTransport.syncReorderBuffer.bind(attachingTransport); + jest + .spyOn(attachingTransport, 'syncReorderBuffer') + .mockImplementation(async (syncStreamId, replayedSequenceFrontier) => { + deliverSequencedMessage(messageHandler, streamId, { + type: 'chunk', + seq: 1, + data: deltaEvent, + }); + return originalSync(syncStreamId, replayedSequenceFrontier); + }); + + const received: unknown[] = []; + const subscription = await attachingManager.subscribe(streamId, (event) => + received.push(event), + ); + + expect(received).toEqual([createdEvent, deltaEvent]); + + subscription?.unsubscribe(); + await generatingManager.destroy(); + await attachingManager.destroy(); + }); + + it('keeps the replay frontier aligned after publication failure and reconnect', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const streamId = 'buffered-publish-failure'; + const messageHandler = getMessageHandler(mockSubscriber); + mockPublisher.publish.mockImplementation(async (channel: string, payload: string) => { + if (transport.getSubscriberCount(streamId) > 0) { + messageHandler(channel, payload); + } + return 1; + }); + + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport: transport, + isRedis: true, + }); + manager.initialize(); + + await manager.createJob(streamId, 'user-1'); + mockPublisher.eval.mockRejectedValueOnce( + new Error('publish failed before sequence allocation'), + ); + + await manager.emitChunk(streamId, { + event: 'on_message_delta', + data: { delta: { content: { type: 'text', text: 'buffered locally' } } }, + }); + + const received: unknown[] = []; + const subscription = await manager.subscribe(streamId, (event) => received.push(event)); + + await manager.emitChunk(streamId, { + event: 'on_message_delta', + data: { delta: { content: { type: 'text', text: 'first live chunk' } } }, + }); + + expect(received).toEqual([ + { + event: 'on_message_delta', + data: { delta: { content: { type: 'text', text: 'buffered locally' } } }, + }, + { + event: 'on_message_delta', + data: { delta: { content: { type: 'text', text: 'first live chunk' } } }, + }, + ]); + + subscription?.unsubscribe(); + + await manager.emitChunk(streamId, { + event: 'on_message_delta', + data: { delta: { content: { type: 'text', text: 'buffered after disconnect' } } }, + }); + + const resumed: unknown[] = []; + const resumedSubscription = await manager.subscribe(streamId, (event) => resumed.push(event)); + await manager.emitChunk(streamId, { + event: 'on_message_delta', + data: { delta: { content: { type: 'text', text: 'live after reconnect' } } }, + }); + + expect(resumed).toEqual([ + { + event: 'on_message_delta', + data: { delta: { content: { type: 'text', text: 'buffered after disconnect' } } }, + }, + { + event: 'on_message_delta', + data: { delta: { content: { type: 'text', text: 'live after reconnect' } } }, + }, + ]); + + resumedSubscription?.unsubscribe(); + await manager.destroy(); + }); + it('resets stale abort-listener reorder state before the next real subscriber', async () => { const mockPublisher = createMockPublisher(); const mockSubscriber = createMockSubscriber(); diff --git a/packages/api/src/stream/__tests__/RedisEventTransport.stream_integration.spec.ts b/packages/api/src/stream/__tests__/RedisEventTransport.stream_integration.spec.ts index ed4a382bf0..2e42a3f668 100644 --- a/packages/api/src/stream/__tests__/RedisEventTransport.stream_integration.spec.ts +++ b/packages/api/src/stream/__tests__/RedisEventTransport.stream_integration.spec.ts @@ -96,6 +96,286 @@ describe('RedisEventTransport Integration Tests', () => { subscriber.disconnect(); }); + test('should publish a zero-TTL terminal winner without leaking it into a replacement', async () => { + if (!ioredisClient) { + console.warn('Redis not available, skipping test'); + return; + } + + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const { RedisEventTransport } = await import('../implementations/RedisEventTransport'); + + const subscriber = (ioredisClient as Redis).duplicate(); + const transport = new RedisEventTransport(ioredisClient, subscriber); + const store = new RedisJobStore(ioredisClient, { completedTtl: 0 }); + await store.initialize(); + + const streamId = `zero-ttl-terminal-${Date.now()}`; + const received: unknown[] = []; + const subscription = transport.subscribe(streamId, { + onChunk: () => undefined, + onDone: (event) => received.push(event), + }); + await subscription.ready; + + let replacementCreatedAt: number | undefined; + try { + const terminalJob = await store.createJob(streamId, 'user-1', streamId); + await expect( + store.transitionStatus(streamId, { + from: 'running', + to: 'aborted', + expectCreatedAt: terminalJob.createdAt, + patch: { completedAt: Date.now() }, + }), + ).resolves.toBe(true); + await expect(store.getJob(streamId)).resolves.toBeNull(); + + const finalEvent = { final: true, aborted: true }; + await transport.emitDone(streamId, finalEvent, terminalJob.createdAt); + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(received).toEqual([finalEvent]); + + const regressedClock = jest + .spyOn(Date, 'now') + .mockReturnValue(terminalJob.createdAt - 1000); + const replacement = await (async () => { + try { + return await store.createJob(streamId, 'user-1', streamId); + } finally { + regressedClock.mockRestore(); + } + })(); + replacementCreatedAt = replacement.createdAt; + expect(replacement.createdAt).toBe(terminalJob.createdAt + 1); + + // A live replacement must reject the predecessor's delayed terminal event. + await transport.emitDone(streamId, { final: true, stale: true }, terminalJob.createdAt); + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(received).toEqual([finalEvent]); + + // The latest epoch marker survives zero-TTL finalization. It still rejects + // the predecessor after the replacement hash itself has disappeared. + await expect( + store.transitionStatus(streamId, { + from: 'running', + to: 'aborted', + expectCreatedAt: replacement.createdAt, + patch: { completedAt: Date.now() }, + }), + ).resolves.toBe(true); + await expect(store.getJob(streamId)).resolves.toBeNull(); + await transport.emitDone( + streamId, + { final: true, staleAfterReplacement: true }, + terminalJob.createdAt, + ); + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(received).toEqual([finalEvent]); + + const replacementFinalEvent = { final: true, replacement: true }; + await transport.emitDone(streamId, replacementFinalEvent, replacement.createdAt); + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(received).toEqual([finalEvent, replacementFinalEvent]); + } finally { + if (replacementCreatedAt != null) { + await store.deleteJob(streamId, replacementCreatedAt); + } + await ioredisClient.del(`stream:{${streamId}}:generation-epoch`); + subscription.unsubscribe(); + await store.destroy(); + transport.destroy(); + subscriber.disconnect(); + } + }); + + test('should publish the latest generation error after stale-job reaping removes its hash', async () => { + if (!ioredisClient) { + console.warn('Redis not available, skipping test'); + return; + } + + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const { RedisEventTransport } = await import('../implementations/RedisEventTransport'); + + const subscriber = (ioredisClient as Redis).duplicate(); + const transport = new RedisEventTransport(ioredisClient, subscriber); + const store = new RedisJobStore(ioredisClient, { runningTtl: 60 }); + await store.initialize(); + + const streamId = `reaped-generation-${Date.now()}`; + let resolveError!: (value: { error: string; generationId?: number }) => void; + const receivedError = new Promise<{ error: string; generationId?: number }>((resolve) => { + resolveError = resolve; + }); + const subscription = transport.subscribe(streamId, { + onChunk: () => undefined, + onError: (error, generationId) => resolveError({ error, generationId }), + }); + await subscription.ready; + + try { + const job = await store.createJob(streamId, 'user-1', streamId); + const generationEpochKey = `stream:{${streamId}}:generation-epoch`; + await ioredisClient.del(generationEpochKey); + await expect(ioredisClient.get(generationEpochKey)).resolves.toBeNull(); + await store.updateJob(streamId, { lastActiveAt: Date.now() - 61_000 }, job.createdAt); + + await expect(store.cleanup()).resolves.toBeGreaterThanOrEqual(1); + await expect(store.getJob(streamId)).resolves.toBeNull(); + await expect(ioredisClient.get(generationEpochKey)).resolves.toBe(String(job.createdAt)); + + await transport.emitError(streamId, 'Generation timed out', job.createdAt); + await expect(receivedError).resolves.toEqual({ + error: 'Generation timed out', + generationId: job.createdAt, + }); + } finally { + await ioredisClient.del(`stream:{${streamId}}:generation-epoch`); + subscription.unsubscribe(); + await store.destroy(); + transport.destroy(); + subscriber.disconnect(); + } + }); + + test('should atomically contain legacy terminal claims when both job and epoch have expired', async () => { + if (!ioredisClient) { + console.warn('Redis not available, skipping test'); + return; + } + + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const { RedisEventTransport } = await import('../implementations/RedisEventTransport'); + + const subscriber = (ioredisClient as Redis).duplicate(); + const transport = new RedisEventTransport(ioredisClient, subscriber); + const store = new RedisJobStore(ioredisClient); + await store.initialize(); + + const claimBeforeCreateId = `legacy-claim-first-${Date.now()}`; + const createBeforeClaimId = `legacy-create-first-${Date.now()}`; + const competingClaimsId = `legacy-competing-${Date.now()}`; + const claimedEvents: string[] = []; + const createFirstEvents: string[] = []; + const competingEvents: string[] = []; + const subscriptions = [ + transport.subscribe(claimBeforeCreateId, { + onChunk: () => undefined, + onError: (error) => claimedEvents.push(error), + }), + transport.subscribe(createBeforeClaimId, { + onChunk: () => undefined, + onError: (error) => createFirstEvents.push(error), + }), + transport.subscribe(competingClaimsId, { + onChunk: () => undefined, + onError: (error) => competingEvents.push(error), + }), + ]; + await Promise.all(subscriptions.map((subscription) => subscription.ready)); + + try { + // Natural expiry of a pre-deploy job leaves neither hash nor epoch. Its + // terminal event claims the empty marker, then replacement creation must + // allocate above the claimed epoch even if the local clock regresses. + const claimFirstLegacy = await store.createJob( + claimBeforeCreateId, + 'user-1', + claimBeforeCreateId, + ); + await ioredisClient.del( + `stream:{${claimBeforeCreateId}}:job`, + `stream:{${claimBeforeCreateId}}:generation-epoch`, + ); + await transport.emitError( + claimBeforeCreateId, + 'legacy claim won', + claimFirstLegacy.createdAt, + ); + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(claimedEvents).toEqual(['legacy claim won']); + await expect( + ioredisClient.get(`stream:{${claimBeforeCreateId}}:generation-epoch`), + ).resolves.toBe(String(claimFirstLegacy.createdAt)); + + const regressedClock = jest + .spyOn(Date, 'now') + .mockReturnValue(claimFirstLegacy.createdAt - 1000); + const claimedReplacement = await (async () => { + try { + return await store.createJob(claimBeforeCreateId, 'user-1', claimBeforeCreateId); + } finally { + regressedClock.mockRestore(); + } + })(); + expect(claimedReplacement.createdAt).toBe(claimFirstLegacy.createdAt + 1); + await transport.emitError( + claimBeforeCreateId, + 'stale after replacement', + claimFirstLegacy.createdAt, + ); + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(claimedEvents).toEqual(['legacy claim won']); + + // If replacement creation wins the Redis serialization order, its live + // hash rejects the legacy terminal before the fallback can claim. + const createFirstLegacy = await store.createJob( + createBeforeClaimId, + 'user-1', + createBeforeClaimId, + ); + await ioredisClient.del( + `stream:{${createBeforeClaimId}}:job`, + `stream:{${createBeforeClaimId}}:generation-epoch`, + ); + const forwardClock = jest + .spyOn(Date, 'now') + .mockReturnValue(createFirstLegacy.createdAt + 1000); + const createFirstReplacement = await (async () => { + try { + return await store.createJob(createBeforeClaimId, 'user-1', createBeforeClaimId); + } finally { + forwardClock.mockRestore(); + } + })(); + expect(createFirstReplacement.createdAt).toBe(createFirstLegacy.createdAt + 1000); + await transport.emitError( + createBeforeClaimId, + 'legacy claim lost', + createFirstLegacy.createdAt, + ); + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(createFirstEvents).toEqual([]); + + // With two unknowable pre-marker epochs and no live hash, Redis ordering + // gives the marker to the first claimant and rejects the differing second. + await transport.emitError(competingClaimsId, 'first claimant', 100); + await transport.emitError(competingClaimsId, 'second claimant', 200); + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(competingEvents).toEqual(['first claimant']); + await expect( + ioredisClient.get(`stream:{${competingClaimsId}}:generation-epoch`), + ).resolves.toBe('100'); + } finally { + for (const streamId of [claimBeforeCreateId, createBeforeClaimId, competingClaimsId]) { + await store.deleteJob(streamId); + } + await store.cleanup(); + await Promise.all( + [claimBeforeCreateId, createBeforeClaimId, competingClaimsId].map((streamId) => + ioredisClient!.del(`stream:{${streamId}}:generation-epoch`), + ), + ); + for (const subscription of subscriptions) { + subscription.unsubscribe(); + } + await store.destroy(); + transport.destroy(); + subscriber.disconnect(); + } + }); + test('should deliver events across transport instances (simulating different servers)', async () => { if (!ioredisClient) { console.warn('Redis not available, skipping test'); @@ -897,6 +1177,110 @@ describe('RedisEventTransport Integration Tests', () => { subscriber2.disconnect(); }); + test('should deliver a remote abort after the last SSE subscriber disconnects', async () => { + if (!ioredisClient) { + console.warn('Redis not available, skipping test'); + return; + } + + const { RedisEventTransport } = await import('../implementations/RedisEventTransport'); + + const subscriber1 = (ioredisClient as Redis).duplicate(); + const subscriber2 = (ioredisClient as Redis).duplicate(); + const transport1 = new RedisEventTransport(ioredisClient, subscriber1); + const transport2 = new RedisEventTransport(ioredisClient, subscriber2); + const streamId = `abort-after-disconnect-${Date.now()}`; + let remoteAbortReceived = false; + let signalAbortReceived: (() => void) | undefined; + const abortReceived = new Promise((resolve) => { + signalAbortReceived = resolve; + }); + let abortTimeout: ReturnType | undefined; + + try { + await transport1.onAbort(streamId, () => { + remoteAbortReceived = true; + signalAbortReceived?.(); + }); + const sseSubscription = transport1.subscribe(streamId, { onChunk: () => undefined }); + await sseSubscription.ready; + sseSubscription.unsubscribe(); + + transport2.emitAbort(streamId); + + await Promise.race([ + abortReceived, + new Promise((_, reject) => { + abortTimeout = setTimeout( + () => reject(new Error('Timed out waiting for remote abort')), + 2000, + ); + }), + ]); + expect(remoteAbortReceived).toBe(true); + } finally { + clearTimeout(abortTimeout); + transport1.cleanup(streamId); + transport1.destroy(); + transport2.destroy(); + subscriber1.disconnect(); + subscriber2.disconnect(); + } + }); + + test('should replace a disposed abort listener without stale cleanup detaching it', async () => { + if (!ioredisClient) { + console.warn('Redis not available, skipping test'); + return; + } + + const { RedisEventTransport } = await import('../implementations/RedisEventTransport'); + + const subscriber1 = (ioredisClient as Redis).duplicate(); + const subscriber2 = (ioredisClient as Redis).duplicate(); + const transport1 = new RedisEventTransport(ioredisClient, subscriber1); + const transport2 = new RedisEventTransport(ioredisClient, subscriber2); + const streamId = `abort-listener-reuse-${Date.now()}`; + let predecessorCalled = false; + let resolveReplacement!: () => void; + const replacementCalled = new Promise((resolve) => { + resolveReplacement = resolve; + }); + let abortTimeout: ReturnType | undefined; + + try { + const disposePredecessor = await transport1.onAbort(streamId, () => { + predecessorCalled = true; + }); + disposePredecessor(); + + const disposeReplacement = await transport1.onAbort(streamId, () => { + resolveReplacement(); + }); + disposePredecessor(); + transport2.emitAbort(streamId); + + await Promise.race([ + replacementCalled, + new Promise((_, reject) => { + abortTimeout = setTimeout( + () => reject(new Error('Timed out waiting for replacement abort listener')), + 2000, + ); + }), + ]); + + expect(predecessorCalled).toBe(false); + disposeReplacement(); + } finally { + clearTimeout(abortTimeout); + transport1.destroy(); + transport2.destroy(); + subscriber1.disconnect(); + subscriber2.disconnect(); + } + }); + test('should call multiple abort callbacks', async () => { if (!ioredisClient) { console.warn('Redis not available, skipping test'); @@ -972,11 +1356,11 @@ describe('RedisEventTransport Integration Tests', () => { /** * Cross-Replica Sequence Synchronization (#12575) * - * The core cross-replica sync logic (Redis INCR counter, async GET in - * syncReorderBuffer, pruneStaleEntries flag) is verified by: + * The core cross-replica sync logic (atomic sequence allocation, first-observed + * attachment baseline, and same-replica replay frontier) is verified by: * - Unit tests with mock publishers (deterministic, no cluster timing) * - GenerationJobManager integration tests (end-to-end with earlyEventBuffer) - * - The race-condition unit test (paused GET with injected message) + * - The race-condition unit test (delayed subscriber dispatch after publish) * * Transport-level integration tests with two real Redis transports are * inherently flaky in Redis Cluster: cluster pub/sub fan-out is async @@ -1054,6 +1438,7 @@ describe('RedisEventTransport Integration Tests', () => { describe('Publish Error Propagation', () => { test('should swallow emitChunk publish errors (callers fire-and-forget)', async () => { const { RedisEventTransport } = await import('../implementations/RedisEventTransport'); + const { emitChunkWithReceipt } = await import('../internal/chunkPublication'); const mockPublisher = createMockPublisher(); mockPublisher.publish.mockRejectedValue(new Error('Redis connection lost')); @@ -1070,15 +1455,19 @@ describe('RedisEventTransport Integration Tests', () => { const streamId = `error-prop-chunk-${Date.now()}`; - // emitChunk swallows errors because callers often fire-and-forget (no await). - // Throwing would cause unhandled promise rejections. + // Public callers retain Promise; the internal manager capability receives failure + // without creating an unhandled rejection. await expect(transport.emitChunk(streamId, { data: 'test' })).resolves.toBeUndefined(); + await expect(emitChunkWithReceipt(transport, streamId, { data: 'test' })).resolves.toBe( + false, + ); transport.destroy(); }); test('should swallow emitChunk incr errors (sequence allocation failure)', async () => { const { RedisEventTransport } = await import('../implementations/RedisEventTransport'); + const { emitChunkWithReceipt } = await import('../internal/chunkPublication'); const mockPublisher = createMockPublisher(); mockPublisher.incr.mockRejectedValue(new Error('INCR failed')); @@ -1096,6 +1485,9 @@ describe('RedisEventTransport Integration Tests', () => { const streamId = `error-prop-incr-${Date.now()}`; await expect(transport.emitChunk(streamId, { data: 'test' })).resolves.toBeUndefined(); + await expect(emitChunkWithReceipt(transport, streamId, { data: 'test' })).resolves.toBe( + false, + ); expect(mockPublisher.publish).not.toHaveBeenCalled(); transport.destroy(); diff --git a/packages/api/src/stream/__tests__/RedisJobStore.spec.ts b/packages/api/src/stream/__tests__/RedisJobStore.spec.ts new file mode 100644 index 0000000000..838fa49c23 --- /dev/null +++ b/packages/api/src/stream/__tests__/RedisJobStore.spec.ts @@ -0,0 +1,732 @@ +import type { Cluster } from 'ioredis'; +import { InMemoryJobStore } from '../implementations/InMemoryJobStore'; +import { RedisJobStore } from '../implementations/RedisJobStore'; + +jest.mock('~/cache/redisTelemetry', () => ({ + RedisUseCases: { GENERATION_STREAM: 'generation_stream' }, + instrumentIORedisClient: (client: unknown) => client, +})); + +type Deferred = { + promise: Promise; + resolve: (value: T) => void; +}; + +function createDeferred(): Deferred { + let resolve!: (value: T) => void; + const promise = new Promise((resolvePromise) => { + resolve = resolvePromise; + }); + return { promise, resolve }; +} + +async function waitFor(predicate: () => boolean): Promise { + for (let attempt = 0; attempt < 20; attempt++) { + if (predicate()) { + return; + } + await Promise.resolve(); + } + expect(predicate()).toBe(true); +} + +function jobHashFromCreationCall(call: unknown[]): Record { + const keyCount = Number(call[1]); + const fields = call.slice(5 + keyCount); + return Object.fromEntries( + Array.from({ length: fields.length / 2 }, (_, index) => [ + String(fields[index * 2]), + String(fields[index * 2 + 1]), + ]), + ); +} + +describe('RedisJobStore', () => { + test('guards the atomic status transition with the expected creation epoch', async () => { + const evalTransition = jest.fn().mockResolvedValue(0); + const redis = { + isCluster: true, + eval: evalTransition, + hgetall: jest.fn().mockResolvedValue({}), + } as unknown as Cluster; + const store = new RedisJobStore(redis); + + await expect( + store.transitionStatus('stream-epoch', { + from: 'running', + to: 'error', + expectCreatedAt: 123456, + }), + ).resolves.toBe(false); + + const [ + script, + keyCount, + jobKey, + sequenceKey, + chunksKey, + runStepsKey, + steersKey, + parkedSteersKey, + generationEpochKey, + from, + actionId, + createdAt, + ] = evalTransition.mock.calls[0]; + expect(script).toContain('HGET", KEYS[1], "createdAt"'); + expect(script).toContain('redis.call("DEL", KEYS[5])'); + expect(script).toContain('redis.call("SET", KEYS[6]'); + expect(script).toContain( + 'redis.call("SET", KEYS[7], currentCreatedAt, "EX", ttl + generationEpochGraceTtl)', + ); + expect(script.indexOf('local ownerUserId')).toBeLessThan( + script.indexOf('redis.call("EXPIRE", KEYS[1], ttl)'), + ); + expect([ + keyCount, + jobKey, + sequenceKey, + chunksKey, + runStepsKey, + steersKey, + parkedSteersKey, + generationEpochKey, + from, + actionId, + createdAt, + ]).toEqual([ + 7, + 'stream:{stream-epoch}:job', + 'stream:{stream-epoch}:seq', + 'stream:{stream-epoch}:chunks', + 'stream:{stream-epoch}:runsteps', + 'stream:{stream-epoch}:steers', + 'stream:{stream-epoch}:parked', + 'stream:{stream-epoch}:generation-epoch', + 'running', + '', + '123456', + ]); + }); + + test('retains the generation epoch beyond the paused job TTL', async () => { + const evalTransition = jest.fn().mockResolvedValue(0); + const redis = { + isCluster: true, + eval: evalTransition, + } as unknown as Cluster; + const store = new RedisJobStore(redis, { requiresActionTtl: 4321 }); + + await store.transitionStatus('stream-paused-epoch', { + from: 'running', + to: 'requires_action', + expectCreatedAt: 123456, + }); + + const transitionCall = evalTransition.mock.calls[0]; + expect(transitionCall[0]).toContain( + 'redis.call("SET", KEYS[7], currentCreatedAt, "EX", ttl + generationEpochGraceTtl)', + ); + expect(transitionCall[12]).toBe('4321'); + expect(transitionCall[13]).toBe('0'); + expect(transitionCall[17]).toBe('300'); + }); + + test('seeds the guarded epoch when reaping a legacy job without a marker', async () => { + const now = jest.spyOn(Date, 'now').mockReturnValue(100_000); + const evalCommand = jest.fn().mockResolvedValue(1); + const redis = { + isCluster: true, + eval: evalCommand, + hgetall: jest + .fn() + .mockResolvedValueOnce({ + streamId: 'legacy-reap', + userId: 'user-1', + status: 'running', + createdAt: '1', + syncSent: '0', + }) + .mockResolvedValue({}), + smembers: jest.fn().mockResolvedValueOnce(['legacy-reap']).mockResolvedValueOnce([]), + srem: jest.fn().mockResolvedValue(1), + } as unknown as Cluster; + const store = new RedisJobStore(redis, { runningTtl: 60 }); + + try { + await expect(store.cleanup()).resolves.toBe(1); + } finally { + now.mockRestore(); + } + + const reapCall = evalCommand.mock.calls[0]; + expect(reapCall[0]).toContain('redis.call("SET", KEYS[6], ARGV[1], "EX", tonumber(ARGV[5]))'); + expect(reapCall.slice(1)).toEqual([ + 6, + 'stream:{legacy-reap}:job', + 'stream:{legacy-reap}:chunks', + 'stream:{legacy-reap}:runsteps', + 'stream:{legacy-reap}:steers', + 'stream:{legacy-reap}:parked', + 'stream:{legacy-reap}:generation-epoch', + '1', + '100000', + '60000', + '300', + '300', + ]); + }); + + test('guards steer peeks with the expected creation epoch in one Redis script', async () => { + const queuedSteer = { + steerId: 'steer-1', + text: 'keep me', + userId: 'user-1', + createdAt: 123, + }; + const evalPeek = jest.fn().mockResolvedValue([JSON.stringify(queuedSteer)]); + const lrange = jest.fn(); + const redis = { + isCluster: true, + eval: evalPeek, + lrange, + } as unknown as Cluster; + const store = new RedisJobStore(redis); + + await expect(store.peekSteers('stream-peek', 123)).resolves.toEqual([queuedSteer]); + + const [script, keyCount, jobKey, steersKey, expectedCreatedAt] = evalPeek.mock.calls[0]; + expect(script).toContain('HGET", KEYS[1], "createdAt"'); + expect(script).toContain('LRANGE", KEYS[2], 0, -1'); + expect([keyCount, jobKey, steersKey, expectedCreatedAt]).toEqual([ + 2, + 'stream:{stream-peek}:job', + 'stream:{stream-peek}:steers', + '123', + ]); + expect(lrange).not.toHaveBeenCalled(); + }); + + test('writes initial metadata in the atomic job creation', async () => { + const evalJobCreation = jest + .fn() + .mockImplementation((...args: unknown[]) => ['', '', args[Number(args[1]) + 3]]); + const redis = { + isCluster: true, + eval: evalJobCreation, + hgetall: jest.fn(() => jobHashFromCreationCall(evalJobCreation.mock.calls[0])), + sadd: jest.fn().mockResolvedValue(1), + srem: jest.fn().mockResolvedValue(1), + expire: jest.fn().mockResolvedValue(1), + } as unknown as Cluster; + const store = new RedisJobStore(redis); + + const job = await store.createJob('stream-metadata', 'user-1', 'conversation-1', undefined, { + conversationId: 'untrusted-conversation', + userMessage: { + messageId: 'message-1', + parentMessageId: 'parent-1', + }, + responseMessageId: 'response-1', + sender: 'Agent', + endpoint: 'agents', + iconURL: 'https://example.com/icon.png', + model: 'test-model', + agent_id: 'agent-1', + isTemporary: false, + promptTokens: 0, + discoveredTools: [], + }); + + expect(job).toMatchObject({ + streamId: 'stream-metadata', + userId: 'user-1', + conversationId: 'conversation-1', + userMessage: { + messageId: 'message-1', + parentMessageId: 'parent-1', + }, + responseMessageId: 'response-1', + sender: 'Agent', + endpoint: 'agents', + iconURL: 'https://example.com/icon.png', + model: 'test-model', + agent_id: 'agent-1', + isTemporary: false, + promptTokens: 0, + discoveredTools: [], + }); + + const creationArgs = evalJobCreation.mock.calls[0]; + const storedFields = jobHashFromCreationCall(creationArgs); + expect(storedFields).toMatchObject({ + conversationId: 'conversation-1', + responseMessageId: 'response-1', + agent_id: 'agent-1', + isTemporary: '0', + promptTokens: '0', + discoveredTools: '[]', + }); + }); + + test('atomically resets predecessor state when creating a replacement', async () => { + const evalJobCreation = jest + .fn() + .mockImplementation((...args: unknown[]) => ['user-1', '', args[Number(args[1]) + 3]]); + const redis = { + isCluster: true, + eval: evalJobCreation, + hgetall: jest.fn(() => jobHashFromCreationCall(evalJobCreation.mock.calls[0])), + sadd: jest.fn().mockResolvedValue(1), + srem: jest.fn().mockResolvedValue(1), + expire: jest.fn().mockResolvedValue(1), + } as unknown as Cluster; + const store = new RedisJobStore(redis); + store.setCollectedUsage('stream-replacement', [{ input_tokens: 10 }]); + + await store.createJob('stream-replacement', 'user-1'); + + const [script, keyCount, ...args] = evalJobCreation.mock.calls[0]; + expect(script).toContain('local retainedEpoch = tonumber(redis.call("GET", KEYS[6]))'); + expect(script).toContain( + 'if previousCreatedAt and previousCreatedAt >= createdAt then createdAt = previousCreatedAt + 1 end', + ); + expect(script).toContain('redis.call("DEL", KEYS[1], KEYS[2], KEYS[3], KEYS[4], KEYS[5])'); + expect(script).toContain( + 'redis.call("SET", KEYS[6], tostring(createdAt), "EX", ttl + generationEpochGraceTtl)', + ); + expect([keyCount, ...args.slice(0, 6)]).toEqual([ + 6, + 'stream:{stream-replacement}:job', + 'stream:{stream-replacement}:chunks', + 'stream:{stream-replacement}:runsteps', + 'stream:{stream-replacement}:steers', + 'stream:{stream-replacement}:parked', + 'stream:{stream-replacement}:generation-epoch', + ]); + expect(args[8]).toBe('300'); + expect(store.getCollectedUsage('stream-replacement')).toEqual([]); + }); + + test('an older overlapping create cannot clear a newer local cache', async () => { + const now = jest.spyOn(Date, 'now').mockReturnValue(100); + const redis = { + isCluster: true, + eval: jest.fn().mockResolvedValue(['user-1', '', '100']), + hgetall: jest.fn().mockResolvedValue({ + streamId: 'stream-overlap', + userId: 'user-1', + status: 'running', + createdAt: '101', + syncSent: '0', + }), + sadd: jest.fn().mockResolvedValue(1), + srem: jest.fn().mockResolvedValue(1), + expire: jest.fn().mockResolvedValue(1), + } as unknown as Cluster; + const store = new RedisJobStore(redis); + const newerUsage = [{ input_tokens: 20 }]; + store.setCollectedUsage('stream-overlap', newerUsage, 101); + + try { + await expect(store.createJob('stream-overlap', 'user-1')).rejects.toThrow( + 'Generation job was replaced during creation', + ); + expect(store.getCollectedUsage('stream-overlap', 101)).toBe(newerUsage); + expect(store.getCollectedUsage('stream-overlap', 100)).toEqual([]); + } finally { + now.mockRestore(); + } + }); + + test('rejects creation when its durable epoch is already terminal', async () => { + const redis = { + isCluster: true, + eval: jest.fn().mockResolvedValue(['', '', '100']), + hgetall: jest.fn().mockResolvedValue({ + streamId: 'stream-terminal-create', + userId: 'user-1', + status: 'complete', + createdAt: '100', + syncSent: '0', + }), + sadd: jest.fn().mockResolvedValue(1), + srem: jest.fn().mockResolvedValue(1), + expire: jest.fn().mockResolvedValue(1), + } as unknown as Cluster; + const store = new RedisJobStore(redis); + + await expect(store.createJob('stream-terminal-create', 'user-1')).rejects.toThrow( + 'Generation job was replaced during creation', + ); + }); + + test('guards local content caches by creation epoch', async () => { + const evalRedis = jest.fn().mockResolvedValue(false); + const redis = { + isCluster: true, + eval: evalRedis, + } as unknown as Cluster; + const store = new RedisJobStore(redis); + const replacementContent = [{ type: 'text', text: 'replacement' }]; + const replacementUsage = [{ input_tokens: 20 }]; + const replacementGraph = { + getContentParts: () => replacementContent, + getRunSteps: () => [{ id: 'replacement-step' }], + }; + const staleGraph = { + getContentParts: () => [{ type: 'text', text: 'predecessor' }], + getRunSteps: () => [{ id: 'predecessor-step' }], + }; + + store.setContentParts('stream-local-epoch', replacementContent, 101); + store.setCollectedUsage('stream-local-epoch', replacementUsage, 101); + store.setGraph('stream-local-epoch', replacementGraph as never, 101); + + store.setContentParts('stream-local-epoch', [{ type: 'text', text: 'predecessor' }], 100); + store.setCollectedUsage('stream-local-epoch', [{ input_tokens: 10 }], 100); + store.setGraph('stream-local-epoch', staleGraph as never, 100); + store.setCollectedUsage('stream-local-epoch', [{ input_tokens: 30 }]); + store.clearContentState('stream-local-epoch', 100); + + await expect(store.getContentParts('stream-local-epoch', 101)).resolves.toEqual({ + content: replacementContent, + }); + expect(store.getCollectedUsage('stream-local-epoch', 101)).toBe(replacementUsage); + await expect(store.getRunSteps('stream-local-epoch', 101)).resolves.toEqual([ + { id: 'replacement-step' }, + ]); + await expect(store.getContentParts('stream-local-epoch', 100)).resolves.toBeNull(); + expect(store.getCollectedUsage('stream-local-epoch', 100)).toEqual([]); + await expect(store.getRunSteps('stream-local-epoch', 100)).resolves.toEqual([]); + + store.clearContentState('stream-local-epoch', 101); + await expect(store.getContentParts('stream-local-epoch', 101)).resolves.toBeNull(); + expect(store.getCollectedUsage('stream-local-epoch', 101)).toEqual([]); + + const chunkRead = evalRedis.mock.calls.find(([script]) => String(script).includes('XRANGE')); + expect(chunkRead).toEqual([ + expect.stringContaining('HGET", KEYS[1], "createdAt"'), + 2, + 'stream:{stream-local-epoch}:job', + 'stream:{stream-local-epoch}:chunks', + '100', + ]); + const runStepsRead = evalRedis.mock.calls.find(([script]) => + String(script).includes('GET", KEYS[2]'), + ); + expect(runStepsRead).toEqual([ + expect.stringContaining('HGET", KEYS[1], "createdAt"'), + 2, + 'stream:{stream-local-epoch}:job', + 'stream:{stream-local-epoch}:runsteps', + '100', + ]); + }); + + test('parallelizes Redis Cluster membership bookkeeping with ordered user TTL', async () => { + const evalResult = createDeferred(); + const runningMembership = createDeferred(); + const requiresActionRemoval = createDeferred(); + const userMembership = createDeferred(); + const userExpiry = createDeferred(); + const started: string[] = []; + + const expire = jest.fn(() => { + started.push('user_expiry'); + return userExpiry.promise; + }); + const evalJobCreation = jest.fn(() => { + started.push('job'); + return evalResult.promise; + }); + const redis = { + isCluster: true, + eval: evalJobCreation, + sadd: jest.fn((key: string) => { + if (key === 'stream:running') { + started.push('running'); + return runningMembership.promise; + } + started.push('user'); + return userMembership.promise; + }), + srem: jest.fn(() => { + started.push('requires_action'); + return requiresActionRemoval.promise; + }), + hgetall: jest.fn(() => jobHashFromCreationCall(evalJobCreation.mock.calls[0])), + expire, + } as unknown as Cluster; + const store = new RedisJobStore(redis, { userJobsSetTtl: 60 }); + + let settled = false; + const creating = store.createJob('stream-1', 'user-1').then((job) => { + settled = true; + return job; + }); + + expect(started).toEqual(['job']); + evalResult.resolve(1); + await waitFor(() => started.length === 4); + + expect(started).toEqual(['job', 'running', 'requires_action', 'user']); + expect(settled).toBe(false); + expect(expire).not.toHaveBeenCalled(); + + userMembership.resolve(1); + await waitFor(() => expire.mock.calls.length === 1); + + expect(started).toEqual(['job', 'running', 'requires_action', 'user', 'user_expiry']); + expect(expire).toHaveBeenCalledWith('stream:user:{user-1}:jobs', 60); + expect(settled).toBe(false); + + userExpiry.resolve(1); + await Promise.resolve(); + expect(settled).toBe(false); + + runningMembership.resolve(1); + await Promise.resolve(); + expect(settled).toBe(false); + + requiresActionRemoval.resolve(1); + await expect(creating).resolves.toMatchObject({ + streamId: 'stream-1', + userId: 'user-1', + status: 'running', + }); + expect(settled).toBe(true); + }); + + test('reconciles membership again when the generation changes at the post-write check', async () => { + const postWriteCheckStarted = createDeferred(); + const releasePostWriteCheck = createDeferred(); + const memberships = new Map>([ + ['stream:running', new Set(['stream-race'])], + ['stream:requires_action', new Set()], + ['stream:user:{user-old}:jobs', new Set(['stream-race'])], + ['stream:user:{user-new}:jobs', new Set()], + ]); + let durableHash: Record = { + streamId: 'stream-race', + userId: 'user-old', + status: 'running', + createdAt: '100', + syncSent: '0', + }; + let reads = 0; + const hgetall = jest.fn(async () => { + reads++; + if (reads === 2) { + postWriteCheckStarted.resolve(); + await releasePostWriteCheck.promise; + } + return { ...durableHash }; + }); + const sadd = jest.fn(async (key: string, streamId: string) => { + let members = memberships.get(key); + if (!members) { + members = new Set(); + memberships.set(key, members); + } + members.add(streamId); + return 1; + }); + const srem = jest.fn(async (key: string, streamId: string) => { + memberships.get(key)?.delete(streamId); + return 1; + }); + const redis = { + isCluster: true, + eval: jest.fn(async () => { + durableHash = { ...durableHash, status: 'requires_action' }; + return 1; + }), + hgetall, + sadd, + srem, + expire: jest.fn().mockResolvedValue(1), + } as unknown as Cluster; + const store = new RedisJobStore(redis); + + const transitioning = store.transitionStatus('stream-race', { + from: 'running', + to: 'requires_action', + }); + await postWriteCheckStarted.promise; + + durableHash = { + streamId: 'stream-race', + userId: 'user-new', + status: 'running', + createdAt: '200', + syncSent: '0', + }; + releasePostWriteCheck.resolve(); + + await expect(transitioning).resolves.toBe(true); + expect(reads).toBe(3); + expect(memberships.get('stream:running')).toContain('stream-race'); + expect(memberships.get('stream:requires_action')).not.toContain('stream-race'); + expect(memberships.get('stream:user:{user-old}:jobs')).not.toContain('stream-race'); + expect(memberships.get('stream:user:{user-new}:jobs')).toContain('stream-race'); + }); + + test('passes expected creation epochs to atomic Redis update and delete scripts', async () => { + const evalCommand = jest.fn().mockResolvedValue(0); + const redis = { + isCluster: true, + eval: evalCommand, + hgetall: jest.fn().mockResolvedValue({ + streamId: 'stream-guarded', + userId: 'user-1', + status: 'running', + createdAt: '200', + syncSent: '0', + }), + } as unknown as Cluster; + const store = new RedisJobStore(redis); + + await store.updateJob('stream-guarded', { error: 'late predecessor error' }, 100); + await expect(store.deleteJob('stream-guarded', 100)).resolves.toBe(false); + + const updateCall = evalCommand.mock.calls[0]; + expect(updateCall[0]).toContain('HGET", KEYS[1], "createdAt"'); + expect(updateCall[6]).toBe('100'); + const deleteCall = evalCommand.mock.calls[1]; + expect(deleteCall[0]).toContain('HGET", KEYS[1], "createdAt"'); + expect(deleteCall.slice(1)).toEqual([ + 4, + 'stream:{stream-guarded}:job', + 'stream:{stream-guarded}:chunks', + 'stream:{stream-guarded}:runsteps', + 'stream:{stream-guarded}:steers', + '100', + '0', + ]); + }); + + test('guards chunk appends with the expected creation epoch inside Redis Lua', async () => { + const evalCommand = jest.fn().mockResolvedValue(0); + const redis = { + isCluster: true, + eval: evalCommand, + } as unknown as Cluster; + const store = new RedisJobStore(redis); + const event = { event: 'on_message_delta', data: { text: 'stale' } }; + + await store.appendChunk('stream-chunk-guarded', event, 100); + + const appendCall = evalCommand.mock.calls[0]; + expect(appendCall[0]).toContain('redis.call("HGET", KEYS[2], "createdAt") ~= ARGV[3]'); + expect(appendCall.slice(1)).toEqual([ + 2, + 'stream:{stream-chunk-guarded}:chunks', + 'stream:{stream-chunk-guarded}:job', + JSON.stringify(event), + '1200', + '100', + ]); + }); + + test('guards run-step saves with the expected creation epoch inside Redis Lua', async () => { + const evalCommand = jest.fn().mockResolvedValue(0); + const redis = { + isCluster: true, + eval: evalCommand, + } as unknown as Cluster; + const store = new RedisJobStore(redis); + const runSteps = [{ id: 'step-1', type: 'tool_call' }]; + + await store.saveRunSteps?.('stream-runstep-guarded', runSteps as never, 100); + + const saveCall = evalCommand.mock.calls[0]; + expect(saveCall[0]).toContain('redis.call("HGET", KEYS[2], "createdAt") ~= ARGV[3]'); + expect(saveCall.slice(1)).toEqual([ + 2, + 'stream:{stream-runstep-guarded}:runsteps', + 'stream:{stream-runstep-guarded}:job', + JSON.stringify(runSteps), + '1200', + '100', + ]); + }); + + test('guards asynchronous content cleanup against a replacement epoch', async () => { + const evalCommand = jest.fn().mockResolvedValue(0); + const redis = { + isCluster: true, + eval: evalCommand, + } as unknown as Cluster; + const store = new RedisJobStore(redis); + + store.clearContentState('stream-content-guarded', 100); + await waitFor(() => evalCommand.mock.calls.length === 1); + + const clearCall = evalCommand.mock.calls[0]; + expect(clearCall[0]).toContain('redis.call("EXISTS", KEYS[3]) == 1'); + expect(clearCall.slice(1)).toEqual([ + 3, + 'stream:{stream-content-guarded}:chunks', + 'stream:{stream-content-guarded}:runsteps', + 'stream:{stream-content-guarded}:job', + '100', + ]); + }); + + test('in-memory update and delete guards preserve a replacement epoch', async () => { + const now = jest.spyOn(Date, 'now'); + try { + const store = new InMemoryJobStore(); + now.mockReturnValue(100); + const original = await store.createJob('stream-memory-guard', 'user-1'); + now.mockReturnValue(200); + const replacement = await store.createJob('stream-memory-guard', 'user-1'); + + await store.updateJob( + 'stream-memory-guard', + { error: 'late predecessor error' }, + original.createdAt, + ); + await expect(store.deleteJob('stream-memory-guard', original.createdAt)).resolves.toBe(false); + await expect(store.getJob('stream-memory-guard')).resolves.toMatchObject({ + createdAt: replacement.createdAt, + status: 'running', + }); + await expect(store.deleteJob('stream-memory-guard', replacement.createdAt)).resolves.toBe( + true, + ); + } finally { + now.mockRestore(); + } + }); + + test('in-memory replacement advances the creation epoch when the clock does not', async () => { + const now = jest.spyOn(Date, 'now').mockReturnValue(100); + try { + const store = new InMemoryJobStore({ maxJobs: 1 }); + const original = await store.createJob('stream-memory-collision', 'user-1'); + const replacement = await store.createJob('stream-memory-collision', 'user-1'); + + expect(original.createdAt).toBe(100); + expect(replacement.createdAt).toBe(101); + await expect(store.getJob('stream-memory-collision')).resolves.toMatchObject({ + createdAt: 101, + }); + } finally { + now.mockRestore(); + } + }); + + test('in-memory replacement clears predecessor content state', async () => { + const store = new InMemoryJobStore(); + await store.createJob('stream-memory-content', 'user-1'); + store.setContentParts('stream-memory-content', [{ type: 'text', text: 'predecessor' }]); + store.setCollectedUsage('stream-memory-content', [{ input_tokens: 10 }]); + + await store.createJob('stream-memory-content', 'user-1'); + + await expect(store.getContentParts('stream-memory-content')).resolves.toBeNull(); + expect(store.getCollectedUsage('stream-memory-content')).toEqual([]); + }); +}); diff --git a/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts b/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts index 338ce636a0..e48fb3bf34 100644 --- a/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts +++ b/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts @@ -168,6 +168,224 @@ describe('RedisJobStore Integration Tests', () => { await store.destroy(); }); + + test('stale generation update and delete cannot affect a same-stream replacement', async () => { + if (!ioredisClient) { + return; + } + + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const store = new RedisJobStore(ioredisClient); + await store.initialize(); + + const streamId = `stale-write-epoch-${Date.now()}`; + const now = jest.spyOn(Date, 'now').mockReturnValue(1000); + const replacementChunk = { + event: 'on_message_delta', + data: { text: 'replacement generation' }, + }; + const predecessorChunk = { + event: 'on_message_delta', + data: { text: 'stale predecessor generation' }, + }; + const replacementRunSteps = [{ id: 'replacement-step', type: 'tool_call' }]; + const predecessorRunSteps = [{ id: 'predecessor-step', type: 'tool_call' }]; + + try { + const original = await store.createJob(streamId, 'user-1', streamId); + await store.appendChunk(streamId, predecessorChunk, original.createdAt); + await store.saveRunSteps?.( + streamId, + predecessorRunSteps as Agents.RunStep[], + original.createdAt, + ); + await store.updateJob( + streamId, + { + finalEvent: JSON.stringify({ final: true, generation: 'predecessor' }), + titleEvent: JSON.stringify({ event: 'title', data: { title: 'Predecessor' } }), + completedAt: original.createdAt, + error: 'predecessor error', + }, + original.createdAt, + ); + store.setCollectedUsage(streamId, [{ input_tokens: 10 }], original.createdAt); + + const replacement = await store.createJob(streamId, 'user-1', streamId); + expect(replacement.createdAt).toBe(original.createdAt + 1); + expect(await ioredisClient.xlen(`stream:{${streamId}}:chunks`)).toBe(0); + await expect(store.getRunSteps(streamId)).resolves.toEqual([]); + const replacementJob = await store.getJob(streamId); + expect(replacementJob).toMatchObject({ + createdAt: replacement.createdAt, + status: 'running', + }); + expect(replacementJob?.finalEvent).toBeUndefined(); + expect(replacementJob?.titleEvent).toBeUndefined(); + expect(replacementJob?.completedAt).toBeUndefined(); + expect(replacementJob?.error).toBeUndefined(); + expect(store.getCollectedUsage(streamId, replacement.createdAt)).toEqual([]); + await store.appendChunk(streamId, replacementChunk, replacement.createdAt); + await store.appendChunk(streamId, predecessorChunk, original.createdAt); + await store.saveRunSteps?.( + streamId, + replacementRunSteps as Agents.RunStep[], + replacement.createdAt, + ); + await store.saveRunSteps?.( + streamId, + predecessorRunSteps as Agents.RunStep[], + original.createdAt, + ); + const replacementContent = [{ type: 'text', text: 'replacement local content' }]; + const replacementUsage = [{ input_tokens: 20 }]; + store.setContentParts(streamId, replacementContent, replacement.createdAt); + store.setCollectedUsage(streamId, replacementUsage, replacement.createdAt); + store.setContentParts( + streamId, + [{ type: 'text', text: 'stale predecessor content' }], + original.createdAt, + ); + store.setCollectedUsage(streamId, [{ input_tokens: 30 }], original.createdAt); + store.clearContentState(streamId, original.createdAt); + + await expect(store.getContentParts(streamId, replacement.createdAt)).resolves.toEqual({ + content: replacementContent, + }); + expect(store.getCollectedUsage(streamId, replacement.createdAt)).toBe(replacementUsage); + await expect(store.getContentParts(streamId, original.createdAt)).resolves.toBeNull(); + await expect(store.getRunSteps(streamId, original.createdAt)).resolves.toEqual([]); + + await store.updateJob( + streamId, + { status: 'complete', completedAt: 1000, sender: 'stale generation' }, + original.createdAt, + ); + await expect(store.deleteJob(streamId, original.createdAt)).resolves.toBe(false); + + await expect(store.getJob(streamId)).resolves.toMatchObject({ + createdAt: replacement.createdAt, + status: 'running', + }); + const chunks = await ioredisClient.xrange(`stream:{${streamId}}:chunks`, '-', '+'); + expect(chunks).toHaveLength(1); + expect(chunks[0]?.[1]).toContain(JSON.stringify(replacementChunk)); + await expect(store.getRunSteps(streamId, replacement.createdAt)).resolves.toEqual( + replacementRunSteps, + ); + + await expect(store.deleteJob(streamId, replacement.createdAt)).resolves.toBe(true); + await expect(store.getJob(streamId)).resolves.toBeNull(); + expect(await ioredisClient.xlen(`stream:{${streamId}}:chunks`)).toBe(0); + } finally { + now.mockRestore(); + await store.destroy(); + } + }); + + test('terminal CAS cleanup cannot delete a replacement job epoch', async () => { + if (!ioredisClient) { + return; + } + + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const store = new RedisJobStore(ioredisClient); + await store.initialize(); + + const streamId = `terminal-epoch-${Date.now()}`; + const userId = 'terminal-epoch-user'; + const now = jest.spyOn(Date, 'now').mockReturnValue(1000); + const originalEval = ioredisClient.eval.bind(ioredisClient) as ( + script: string | Buffer, + numberOfKeys: number, + ...args: Array + ) => Promise; + let signalCasApplied: (() => void) | undefined; + const casApplied = new Promise((resolve) => { + signalCasApplied = resolve; + }); + let releaseTransition: (() => void) | undefined; + const transitionGate = new Promise((resolve) => { + releaseTransition = resolve; + }); + let restoreEval: (() => void) | undefined; + + try { + const originalJob = await store.createJob(streamId, userId, streamId); + await store.appendChunk(streamId, { + event: 'on_message_delta', + data: { text: 'old generation' }, + }); + await store.enqueueSteer(streamId, { + steerId: 'old-steer', + text: 'do not leak this', + userId, + createdAt: Date.now(), + }); + + let gateFirstEval = true; + const evalSpy = jest.spyOn(ioredisClient, 'eval').mockImplementation((async ( + script, + numberOfKeys, + ...args + ) => { + const result = await originalEval( + script as string | Buffer, + Number(numberOfKeys), + ...(args as Array), + ); + if (gateFirstEval) { + gateFirstEval = false; + signalCasApplied?.(); + await transitionGate; + } + return result; + }) as typeof ioredisClient.eval); + restoreEval = () => evalSpy.mockRestore(); + + const finalizing = store.transitionStatus(streamId, { + from: 'running', + to: 'error', + expectCreatedAt: originalJob.createdAt, + patch: { error: 'old generation stopped', completedAt: Date.now() }, + }); + await casApplied; + + now.mockReturnValue(2000); + const replacement = await store.createJob(streamId, userId, streamId); + await store.appendChunk(streamId, { + event: 'on_message_delta', + data: { text: 'replacement generation' }, + }); + await store.enqueueSteer(streamId, { + steerId: 'replacement-steer', + text: 'keep replacement state', + userId, + createdAt: Date.now(), + }); + + releaseTransition?.(); + await expect(finalizing).resolves.toBe(true); + await expect(store.getJob(streamId)).resolves.toMatchObject({ + createdAt: replacement.createdAt, + status: 'running', + }); + expect(await ioredisClient.xlen(`stream:{${streamId}}:chunks`)).toBe(1); + expect((await store.peekSteers(streamId)).map((steer) => steer.steerId)).toEqual([ + 'replacement-steer', + ]); + await expect( + store.claimParkedSteers(streamId, `"userId":"${userId}"`), + ).resolves.toBeUndefined(); + expect(await ioredisClient.smembers('stream:running')).toContain(streamId); + expect(await store.getActiveJobIdsByUser(userId)).toContain(streamId); + } finally { + releaseTransition?.(); + restoreEval?.(); + now.mockRestore(); + await store.destroy(); + } + }); }); describe('Requires Action Status Tracking', () => { @@ -440,8 +658,7 @@ describe('RedisJobStore Integration Tests', () => { const streamId = `stale-agent-${Date.now()}`; // Turn 1: a saved agent in a temporary chat that discovered a deferred tool. - await store.createJob(streamId, 'user-1', streamId); - await store.updateJob(streamId, { + await store.createJob(streamId, 'user-1', streamId, undefined, { agent_id: 'saved-agent-1', isTemporary: true, discoveredTools: ['deep_tool'], @@ -1019,6 +1236,27 @@ describe('RedisJobStore Integration Tests', () => { await store.destroy(); }); + test('same-stream replacement transfers active membership to the new owner', async () => { + if (!ioredisClient) { + return; + } + + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const store = new RedisJobStore(ioredisClient); + await store.initialize(); + + const streamId = `owner-replacement-${Date.now()}`; + const oldUserId = `old-owner-${Date.now()}`; + const newUserId = `new-owner-${Date.now()}`; + await store.createJob(streamId, oldUserId, streamId); + await store.createJob(streamId, newUserId, streamId); + + await expect(store.getActiveJobIdsByUser(oldUserId)).resolves.not.toContain(streamId); + await expect(store.getActiveJobIdsByUser(newUserId)).resolves.toContain(streamId); + + await store.destroy(); + }); + test('should return empty array for user with no jobs', async () => { if (!ioredisClient) { return; @@ -1955,11 +2193,16 @@ describe('RedisJobStore Integration Tests', () => { await store.initialize(); const streamId = `steer-replace-${Date.now()}`; - await store.createJob(streamId, 'steer-user', streamId); + const predecessor = await store.createJob(streamId, 'steer-user', streamId); await store.enqueueSteer(streamId, buildSteer('s1', 'old run steer')); - await store.createJob(streamId, 'steer-user', streamId); - expect(await store.peekSteers(streamId)).toEqual([]); + const replacement = await store.createJob(streamId, 'steer-user', streamId); + await store.enqueueSteer(streamId, buildSteer('s2', 'replacement steer')); + + expect(await store.peekSteers(streamId, predecessor.createdAt)).toEqual([]); + expect( + (await store.peekSteers(streamId, replacement.createdAt)).map((steer) => steer.text), + ).toEqual(['replacement steer']); await store.destroy(); }); @@ -2183,6 +2426,87 @@ describe('RedisJobStore Integration Tests', () => { await store.destroy(); }); + test('terminal CAS with zero completed TTL keeps parked steers owner-claimable', async () => { + if (!ioredisClient) { + return; + } + + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const store = new RedisJobStore(ioredisClient, { completedTtl: 0 }); + await store.initialize(); + + const streamId = `steer-zero-terminal-${Date.now()}`; + const userId = 'zero-terminal-user'; + const job = await store.createJob(streamId, userId, streamId, 'tenant-1'); + await store.enqueueSteer(streamId, buildSteer('s1', 'survive immediate job expiry')); + + await expect( + store.transitionStatus(streamId, { + from: 'running', + to: 'error', + expectCreatedAt: job.createdAt, + patch: { error: 'stopped', completedAt: Date.now() }, + }), + ).resolves.toBe(true); + + await expect(store.getJob(streamId)).resolves.toBeNull(); + expect(await ioredisClient.smembers('stream:running')).not.toContain(streamId); + expect(await store.getActiveJobIdsByUser(userId, 'tenant-1')).not.toContain(streamId); + + const claimed = await store.claimParkedSteers(streamId, `"userId":"${userId}"`); + expect(claimed).toBeDefined(); + expect(JSON.parse(claimed as string)).toMatchObject({ + userId, + tenantId: 'tenant-1', + steers: [{ steerId: 's1', text: 'survive immediate job expiry' }], + }); + + await store.destroy(); + }); + + test('terminal CAS skips malformed steers while parking valid leftovers', async () => { + if (!ioredisClient) { + return; + } + + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const store = new RedisJobStore(ioredisClient); + await store.initialize(); + + const streamId = `steer-malformed-terminal-${Date.now()}`; + const userId = 'malformed-terminal-user'; + const job = await store.createJob(streamId, userId, streamId, 'tenant-1'); + await store.enqueueSteer(streamId, buildSteer('valid-steer', 'preserve valid input')); + await ioredisClient.rpush(`stream:{${streamId}}:steers`, '{malformed-json'); + + await expect( + store.transitionStatus(streamId, { + from: 'running', + to: 'error', + expectCreatedAt: job.createdAt, + patch: { error: 'stopped', completedAt: Date.now() }, + }), + ).resolves.toBe(true); + + await expect(store.getJob(streamId)).resolves.toMatchObject({ status: 'error' }); + expect(await ioredisClient.exists(`stream:{${streamId}}:steers`)).toBe(0); + + const claimed = await store.claimParkedSteers(streamId, `"userId":"${userId}"`); + expect(claimed).toBeDefined(); + expect(JSON.parse(claimed as string)).toMatchObject({ + userId, + tenantId: 'tenant-1', + steers: [ + { + steerId: 'valid-steer', + text: 'preserve valid input', + }, + ], + }); + + await store.destroy(); + }); + test('stale-running reap parks queued steers before deleting the job', async () => { if (!ioredisClient) { return; @@ -2223,6 +2547,95 @@ describe('RedisJobStore Integration Tests', () => { await store.destroy(); }); + test('stale-running reap cannot delete a replacement created after observation', async () => { + if (!ioredisClient) { + return; + } + + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const store = new RedisJobStore(ioredisClient, { runningTtl: 1 }); + await store.initialize(); + + const streamId = `stale-replacement-guard-${Date.now()}`; + const userId = 'stale-replacement-user'; + await store.createJob(streamId, userId, streamId); + await store.enqueueSteer(streamId, buildSteer('old-steer', 'old generation')); + await ioredisClient.hset( + `stream:{${streamId}}:job`, + 'createdAt', + String(Date.now() - 10_000), + ); + + const originalEval = ioredisClient.eval.bind(ioredisClient) as ( + script: string | Buffer, + numberOfKeys: number, + ...args: Array + ) => Promise; + let signalCleanupReady: (() => void) | undefined; + const cleanupReady = new Promise((resolve) => { + signalCleanupReady = resolve; + }); + let releaseCleanup: (() => void) | undefined; + const cleanupGate = new Promise((resolve) => { + releaseCleanup = resolve; + }); + let restoreEval: (() => void) | undefined; + + try { + let gated = false; + const evalSpy = jest.spyOn(ioredisClient, 'eval').mockImplementation((async ( + script, + numberOfKeys, + ...args + ) => { + if (!gated && String(script).includes('tonumber(ARGV[2]) - liveSince')) { + gated = true; + signalCleanupReady?.(); + await cleanupGate; + } + return originalEval( + script as string | Buffer, + Number(numberOfKeys), + ...(args as Array), + ); + }) as typeof ioredisClient.eval); + restoreEval = () => evalSpy.mockRestore(); + + const cleaning = store.cleanup(); + await cleanupReady; + + const replacement = await store.createJob(streamId, userId, streamId); + await store.appendChunk(streamId, { + event: 'on_message_delta', + data: { text: 'replacement generation' }, + }); + await store.enqueueSteer( + streamId, + buildSteer('replacement-steer', 'keep replacement state'), + ); + + releaseCleanup?.(); + await cleaning; + + await expect(store.getJob(streamId)).resolves.toMatchObject({ + createdAt: replacement.createdAt, + status: 'running', + }); + expect(await ioredisClient.xlen(`stream:{${streamId}}:chunks`)).toBe(1); + expect((await store.peekSteers(streamId)).map((steer) => steer.steerId)).toEqual([ + 'replacement-steer', + ]); + await expect( + store.claimParkedSteers(streamId, `"userId":"${userId}"`), + ).resolves.toBeUndefined(); + expect(await ioredisClient.smembers('stream:running')).toContain(streamId); + } finally { + releaseCleanup?.(); + restoreEval?.(); + await store.destroy(); + } + }); + test('parkSteers falls back to a positive recovery TTL when completedTtl is 0', async () => { if (!ioredisClient) { return; diff --git a/packages/api/src/stream/__tests__/helpers/publisher.ts b/packages/api/src/stream/__tests__/helpers/publisher.ts index 4fea5f9eac..c749bfd83e 100644 --- a/packages/api/src/stream/__tests__/helpers/publisher.ts +++ b/packages/api/src/stream/__tests__/helpers/publisher.ts @@ -57,10 +57,14 @@ export function createMockPublisher(): MockPublisher { _numKeys: number, seqKey: string, jobKey: string, + _generationEpochKey: string, channel: string, prefix: string, suffix: string, ttlSeconds: string, + _expectedGenerationId: string, + _allowRetainedEpoch: string, + _generationEpochGraceTtl: string, ) => { const val = (await publisher.incr(seqKey)) as number; let ttl = Number(ttlSeconds); diff --git a/packages/api/src/stream/__tests__/pendingAction.spec.ts b/packages/api/src/stream/__tests__/pendingAction.spec.ts index 67d6258529..ee5dd3df92 100644 --- a/packages/api/src/stream/__tests__/pendingAction.spec.ts +++ b/packages/api/src/stream/__tests__/pendingAction.spec.ts @@ -3,17 +3,22 @@ import { InMemoryEventTransport } from '~/stream/implementations/InMemoryEventTr import { buildPendingAction, buildToolApprovalPayload } from '~/agents/hitl/policy'; import { InMemoryJobStore } from '~/stream/implementations/InMemoryJobStore'; import { GenerationJobManagerClass } from '~/stream/GenerationJobManager'; +import { ApprovalLifecycle } from '~/stream/ApprovalLifecycle'; jest.spyOn(console, 'log').mockImplementation(); describe('ApprovalLifecycle via GenerationJobManager.approvals (in-memory)', () => { let manager: GenerationJobManagerClass; + let jobStore: InMemoryJobStore; + let eventTransport: InMemoryEventTransport; beforeEach(() => { + jobStore = new InMemoryJobStore({ ttlAfterComplete: 60000 }); + eventTransport = new InMemoryEventTransport(); manager = new GenerationJobManagerClass(); manager.configure({ - jobStore: new InMemoryJobStore({ ttlAfterComplete: 60000 }), - eventTransport: new InMemoryEventTransport(), + jobStore, + eventTransport, isRedis: false, cleanupOnComplete: false, }); @@ -223,108 +228,126 @@ describe('ApprovalLifecycle via GenerationJobManager.approvals (in-memory)', () }); }); - describe('expireApproval → approval-expired handler', () => { - // The host registers this to prune the paused run's durable checkpoint eagerly on - // expiry (sweeper or stale submit) instead of waiting out the checkpoint TTL. - test('fires the registered handler with the streamId after a successful expiry', async () => { - const streamId = 'stream-expire-handler'; - await manager.createJob(streamId, 'user-1'); - await manager.approvals.pause(streamId, buildAction(streamId, { actionId: 'action-A' })); + describe('expireApproval notification', () => { + test('publishes a generation-tagged expiry to local and remote subscribers', async () => { + const streamId = 'stream-expire-local-notification'; + const job = await manager.createJob(streamId, 'user-1'); + const onError = jest.fn(); + const subscription = await manager.subscribe(streamId, () => undefined, undefined, onError); + await manager.approvals.pause(streamId, buildAction(streamId)); + const broadcast = jest.spyOn(eventTransport, 'emitError'); - const handler = jest.fn(); - manager.setApprovalExpiredHandler(handler); + expect(await manager.expireApproval(streamId)).toBe(true); + expect(onError).toHaveBeenCalledWith('Approval expired before a decision was made'); + expect(broadcast).toHaveBeenCalledWith( + streamId, + 'Approval expired before a decision was made', + job.createdAt, + ); - expect(await manager.expireApproval(streamId, 'action-A')).toBe(true); - expect(handler).toHaveBeenCalledTimes(1); - // The expired job rides along so the host can resolve tenant/user-scoped config. - expect(handler).toHaveBeenCalledWith(streamId, expect.objectContaining({ userId: 'user-1' })); - - // The aborted job outlives the expiry (completed-job TTL), so the next sweep enters - // the relay branch for the SAME approval — the cleanup must not run a second time. - await ( - manager as unknown as { expireStaleApprovals(): Promise } - ).expireStaleApprovals(); - expect(handler).toHaveBeenCalledTimes(1); + subscription?.unsubscribe(); }); - test('relays a store-won expiry through the handler (multi-replica path)', async () => { - const streamId = 'stream-expire-relay'; + test('delivers the stored expiry error to a late subscriber', async () => { + const streamId = 'stream-expire-late-subscriber'; await manager.createJob(streamId, 'user-1'); await manager.approvals.pause(streamId, buildAction(streamId)); - - // Another replica's store cleanup wins the expiry CAS: the status flips via the - // lifecycle primitive with NO emit, NO handler, and no errorEvent on this replica. expect(await manager.approvals.expire(streamId)).toBe(true); - const handler = jest.fn(); - manager.setApprovalExpiredHandler(handler); - // This replica's sweep observes the already-aborted expiry and relays it. - await ( - manager as unknown as { expireStaleApprovals(): Promise } - ).expireStaleApprovals(); + const onError = jest.fn(); + const subscription = await manager.subscribe(streamId, () => undefined, undefined, onError); + await new Promise((resolve) => setImmediate(resolve)); - expect(handler).toHaveBeenCalledTimes(1); - expect(handler).toHaveBeenCalledWith( - streamId, - expect.objectContaining({ userId: 'user-1', status: 'aborted' }), - ); - - // Repeated sweeps must not re-run the (idempotent but not free) cleanup. - await ( - manager as unknown as { expireStaleApprovals(): Promise } - ).expireStaleApprovals(); - expect(handler).toHaveBeenCalledTimes(1); + expect(onError).toHaveBeenCalledWith('Approval expired before a decision was made'); + subscription?.unsubscribe(); }); - test('relay cleanup still runs when the terminal error is already cached (reconnect)', async () => { - const streamId = 'stream-expire-relay-cached'; + test('notifies the observed runtime when zero terminal TTL removes the job hash', async () => { + const streamId = 'stream-expire-zero-terminal-ttl'; await manager.createJob(streamId, 'user-1'); + const onError = jest.fn(); + const subscription = await manager.subscribe(streamId, () => undefined, undefined, onError); await manager.approvals.pause(streamId, buildAction(streamId)); - expect(await manager.approvals.expire(streamId)).toBe(true); // store-won CAS - - // A reconnect seeds runtime.errorEvent from the aborted job BEFORE any sweep — - // that must gate the relay emit, not the checkpoint cleanup. - const internals = manager as unknown as { - runtimeState: Map; - expireStaleApprovals(): Promise; - }; - const runtime = internals.runtimeState.get(streamId); - expect(runtime).toBeDefined(); - runtime!.errorEvent = 'cached-terminal-error'; - - const handler = jest.fn(); - manager.setApprovalExpiredHandler(handler); - await internals.expireStaleApprovals(); - - expect(handler).toHaveBeenCalledTimes(1); - expect(handler).toHaveBeenCalledWith( - streamId, - expect.objectContaining({ status: 'aborted' }), - ); - }); - - test('does NOT fire when nothing was expired (failed CAS)', async () => { - const streamId = 'stream-expire-handler-noop'; - await manager.createJob(streamId, 'user-1'); // running — no pending action to expire - - const handler = jest.fn(); - manager.setApprovalExpiredHandler(handler); - - expect(await manager.expireApproval(streamId)).toBe(false); - expect(handler).not.toHaveBeenCalled(); - }); - - test('a throwing handler never breaks the expiry itself', async () => { - const streamId = 'stream-expire-handler-throws'; - await manager.createJob(streamId, 'user-1'); - await manager.approvals.pause(streamId, buildAction(streamId)); - - manager.setApprovalExpiredHandler(() => { - throw new Error('prune failed'); + const originalGetJob = jobStore.getJob.bind(jobStore); + jest.spyOn(jobStore, 'getJob').mockImplementation(async (...args) => { + const job = await originalGetJob(...args); + return job?.status === 'aborted' ? null : job; }); expect(await manager.expireApproval(streamId)).toBe(true); - expect(await manager.getJobStatus(streamId)).toBe('aborted'); + expect(onError).toHaveBeenCalledWith('Approval expired before a decision was made'); + subscription?.unsubscribe(); + }); + + test('notifies after the manager pre-read fails and the expiry CAS removes the job hash', async () => { + const streamId = 'stream-expire-read-failure-zero-terminal-ttl'; + await manager.createJob(streamId, 'user-1'); + const onError = jest.fn(); + const subscription = await manager.subscribe(streamId, () => undefined, undefined, onError); + await manager.approvals.pause(streamId, buildAction(streamId)); + const originalGetJob = jobStore.getJob.bind(jobStore); + jest + .spyOn(jobStore, 'getJob') + .mockRejectedValueOnce(new Error('transient read failure')) + .mockImplementation(async (...args) => { + const job = await originalGetJob(...args); + return job?.status === 'aborted' ? null : job; + }); + + expect(await manager.expireApproval(streamId)).toBe(true); + expect(onError).toHaveBeenCalledWith('Approval expired before a decision was made'); + subscription?.unsubscribe(); + }); + + test('does not notify or mutate a replacement created after the expiry CAS', async () => { + const streamId = 'stream-expire-replacement-notification'; + const now = jest.spyOn(Date, 'now').mockReturnValue(1000); + const originalTransition = jobStore.transitionStatus.bind(jobStore); + let signalExpired: (() => void) | undefined; + const expired = new Promise((resolve) => { + signalExpired = resolve; + }); + let releaseTransition: (() => void) | undefined; + const transitionGate = new Promise((resolve) => { + releaseTransition = resolve; + }); + jest.spyOn(jobStore, 'transitionStatus').mockImplementation(async (...args) => { + const transitioned = await originalTransition(...args); + if (args[1].to === 'aborted' && transitioned) { + signalExpired?.(); + await transitionGate; + } + return transitioned; + }); + + try { + await manager.createJob(streamId, 'user-1'); + await manager.approvals.pause(streamId, buildAction(streamId)); + const expiring = manager.expireApproval(streamId); + await expired; + + now.mockReturnValue(2000); + await manager.createJob(streamId, 'user-1'); + const replacementError = jest.fn(); + const replacementSubscription = await manager.subscribe( + streamId, + () => undefined, + undefined, + replacementError, + ); + + releaseTransition?.(); + await expect(expiring).resolves.toBe(true); + await expect(jobStore.getJob(streamId)).resolves.toMatchObject({ + createdAt: 2000, + status: 'running', + }); + expect(replacementError).not.toHaveBeenCalled(); + replacementSubscription?.unsubscribe(); + } finally { + releaseTransition?.(); + now.mockRestore(); + } }); }); @@ -374,6 +397,29 @@ describe('ApprovalLifecycle via GenerationJobManager.approvals (in-memory)', () }); describe('InMemoryJobStore — approval expiry cleanup', () => { + test('guards status transitions against a replaced job epoch', async () => { + const store = new InMemoryJobStore({ ttlAfterComplete: 60000 }); + const job = await store.createJob('epoch-guard', 'u1'); + + expect( + await store.transitionStatus('epoch-guard', { + from: 'running', + to: 'error', + expectCreatedAt: job.createdAt + 1, + }), + ).toBe(false); + expect((await store.getJob('epoch-guard'))?.status).toBe('running'); + + expect( + await store.transitionStatus('epoch-guard', { + from: 'running', + to: 'error', + expectCreatedAt: job.createdAt, + }), + ).toBe(true); + expect((await store.getJob('epoch-guard'))?.status).toBe('error'); + }); + test('cleanup() finalizes and reclaims a past-expiry pending-approval job', async () => { const store = new InMemoryJobStore({ ttlAfterComplete: 0 }); await store.createJob('s1', 'u1'); @@ -394,6 +440,81 @@ describe('InMemoryJobStore — approval expiry cleanup', () => { }); }); +describe('ApprovalLifecycle ownership callbacks', () => { + test('notifies ownership changes only after successful lifecycle transitions', async () => { + const store = new InMemoryJobStore({ ttlAfterComplete: 60000 }); + const callbacks = { + onPaused: jest.fn(), + onResumed: jest.fn(), + onExpired: jest.fn(), + }; + const lifecycle = new ApprovalLifecycle(store, callbacks); + const streamId = 'ownership-callbacks'; + const job = await store.createJob(streamId, 'u1'); + const action = buildPendingAction( + buildToolApprovalPayload([{ name: 'shell', arguments: {}, tool_call_id: 'c1' }]), + { streamId }, + ); + + expect(await lifecycle.pause(streamId, action)).toBe(true); + expect(await lifecycle.pause(streamId, action)).toBe(false); + expect(callbacks.onPaused).toHaveBeenCalledTimes(1); + expect(callbacks.onPaused).toHaveBeenCalledWith(streamId, job.createdAt); + + expect(await lifecycle.resolve(streamId, action.actionId)).toBe(true); + expect(await lifecycle.resolve(streamId, action.actionId)).toBe(false); + expect(callbacks.onResumed).toHaveBeenCalledTimes(1); + expect(callbacks.onResumed).toHaveBeenCalledWith(streamId, job.createdAt); + + const nextAction = { ...action, actionId: 'next-action' }; + expect(await lifecycle.pause(streamId, nextAction)).toBe(true); + expect(await lifecycle.expire(streamId, nextAction.actionId)).toBe(true); + expect(await lifecycle.expire(streamId, nextAction.actionId)).toBe(false); + expect(callbacks.onExpired).toHaveBeenCalledTimes(1); + expect(callbacks.onExpired).toHaveBeenCalledWith(streamId, job.createdAt); + }); + + test('does not resume a replacement that appeared after the pending action was read', async () => { + const store = new InMemoryJobStore({ ttlAfterComplete: 60000 }); + const onResumed = jest.fn(); + const lifecycle = new ApprovalLifecycle(store, { onResumed }); + const streamId = 'resolve-replacement-guard'; + const now = jest.spyOn(Date, 'now').mockReturnValue(1000); + + try { + await store.createJob(streamId, 'u1'); + const action = buildPendingAction( + buildToolApprovalPayload([{ name: 'shell', arguments: {}, tool_call_id: 'c1' }]), + { streamId }, + ); + await lifecycle.pause(streamId, action); + const observedJob = await store.getJob(streamId); + if (!observedJob) { + throw new Error('Expected paused job'); + } + + now.mockReturnValue(2000); + const replacement = await store.createJob(streamId, 'u1'); + await store.transitionStatus(streamId, { + from: 'running', + to: 'requires_action', + patch: { pendingAction: action, pendingActionId: action.actionId }, + }); + + jest.spyOn(store, 'getJob').mockResolvedValueOnce(observedJob); + + expect(await lifecycle.resolve(streamId, action.actionId)).toBe(false); + expect(await store.getJob(streamId)).toMatchObject({ + createdAt: replacement.createdAt, + status: 'requires_action', + }); + expect(onResumed).not.toHaveBeenCalled(); + } finally { + now.mockRestore(); + } + }); +}); + describe('GenerationJobManager HITL resume metadata (round 19)', () => { let manager: GenerationJobManagerClass; diff --git a/packages/api/src/stream/__tests__/reconnect-reorder-desync.stream_integration.spec.ts b/packages/api/src/stream/__tests__/reconnect-reorder-desync.stream_integration.spec.ts index cd800dd0e9..ef091e3f47 100644 --- a/packages/api/src/stream/__tests__/reconnect-reorder-desync.stream_integration.spec.ts +++ b/packages/api/src/stream/__tests__/reconnect-reorder-desync.stream_integration.spec.ts @@ -933,10 +933,10 @@ describe('Reconnect Reorder Buffer Desync (Regression)', () => { * A producer replica can tear down its local transport after generation 1 while a * subscriber on another replica is still attached with nextSeq=10. Since stream IDs * are conversation IDs, generation 2 reuses the same Redis ordering namespace. The - * shared counter must continue at 10; resetting it to 0 makes the lingering consumer - * reject every regenerated chunk as an old duplicate. + * shared counter must continue at 10 for the new attachment, while generation tags + * keep the abandoned generation-1 subscriber from receiving generation-2 content. */ - test('regenerated turn reaches a lingering cross-replica subscriber after producer cleanup', async () => { + test('regenerated turn preserves sequence without leaking to a lingering predecessor subscriber', async () => { if (!ioredisClient) { console.warn('Redis not available, skipping test'); return; @@ -1021,7 +1021,7 @@ describe('Reconnect Reorder Buffer Desync (Regression)', () => { firstGeneration .filter((event) => JSON.stringify(event).includes('"generation":2')) .map((event) => (event as { data: { index: number } }).data.index), - ).toEqual([0, 1, 2, 3, 4]); + ).toEqual([]); expect(await ioredisClient.get(sequenceKey)).toBe('15'); lingering?.unsubscribe(); diff --git a/packages/api/src/stream/__tests__/staleJobReaping.spec.ts b/packages/api/src/stream/__tests__/staleJobReaping.spec.ts index 6c1ec8fddf..972683c97a 100644 --- a/packages/api/src/stream/__tests__/staleJobReaping.spec.ts +++ b/packages/api/src/stream/__tests__/staleJobReaping.spec.ts @@ -76,6 +76,31 @@ describe('InMemoryJobStore - stale running-job failsafe', () => { } }); + it('does not refresh replacement activity from a predecessor generation', async () => { + jest.useFakeTimers(); + try { + const { InMemoryJobStore } = await import('../implementations/InMemoryJobStore'); + const store = new InMemoryJobStore({ ttlAfterComplete: 0, staleJobTimeout: 1000 }); + await store.initialize(); + + const predecessor = await store.createJob('s1', 'u1', 's1'); + await jest.advanceTimersByTimeAsync(1); + const replacement = await store.createJob('s1', 'u1', 's1'); + await jest.advanceTimersByTimeAsync(900); + + store.recordActivity('s1', predecessor.createdAt); + await jest.advanceTimersByTimeAsync(101); + + expect(replacement.createdAt).not.toBe(predecessor.createdAt); + expect(await store.cleanup()).toBe(1); + expect(await store.hasJob('s1')).toBe(false); + + await store.destroy(); + } finally { + jest.useRealTimers(); + } + }); + it('does not reap a running job within the staleJobTimeout', async () => { const { InMemoryJobStore } = await import('../implementations/InMemoryJobStore'); const store = new InMemoryJobStore({ staleJobTimeout: 60000 }); @@ -179,6 +204,60 @@ describe('InMemoryJobStore - stale running-job failsafe', () => { await store.destroy(); }); + + it('does not delete a replacement created while cleanup awaits an earlier victim', async () => { + const { InMemoryJobStore } = await import('../implementations/InMemoryJobStore'); + const store = new InMemoryJobStore({ ttlAfterComplete: 0, staleJobTimeout: 60000 }); + await store.initialize(); + + const now = jest.spyOn(Date, 'now'); + const deleteJob = store.deleteJob.bind(store); + let releaseFirstDelete!: () => void; + let markFirstDeleteStarted!: () => void; + const firstDeleteStarted = new Promise((resolve) => { + markFirstDeleteStarted = resolve; + }); + const firstDeleteReleased = new Promise((resolve) => { + releaseFirstDelete = resolve; + }); + + try { + now.mockReturnValue(100); + await store.createJob('first-victim', 'u1'); + await store.updateJob('first-victim', { status: 'complete', completedAt: 100 }); + + now.mockReturnValue(200); + const originalLaterVictim = await store.createJob('later-victim', 'u1'); + await store.updateJob('later-victim', { status: 'complete', completedAt: 200 }); + + jest.spyOn(store, 'deleteJob').mockImplementation(async (streamId, expectedCreatedAt) => { + if (streamId === 'first-victim') { + markFirstDeleteStarted(); + await firstDeleteReleased; + } + return deleteJob(streamId, expectedCreatedAt); + }); + + now.mockReturnValue(1000); + const cleanup = store.cleanup(); + await firstDeleteStarted; + + now.mockReturnValue(2000); + const replacement = await store.createJob('later-victim', 'u1'); + releaseFirstDelete(); + await cleanup; + + expect(replacement.createdAt).not.toBe(originalLaterVictim.createdAt); + await expect(store.getJob('first-victim')).resolves.toBeNull(); + await expect(store.getJob('later-victim')).resolves.toMatchObject({ + createdAt: replacement.createdAt, + status: 'running', + }); + } finally { + now.mockRestore(); + await store.destroy(); + } + }); }); describe('GenerationJobManager - generation abort on reaping', () => { diff --git a/packages/api/src/stream/__tests__/startup.spec.ts b/packages/api/src/stream/__tests__/startup.spec.ts new file mode 100644 index 0000000000..d1a578a2dc --- /dev/null +++ b/packages/api/src/stream/__tests__/startup.spec.ts @@ -0,0 +1,2353 @@ +import type { AgentStartupTelemetry } from '~/agents/startup'; +import type { ServerSentEvent } from '~/types'; +import { InMemoryEventTransport } from '~/stream/implementations/InMemoryEventTransport'; +import { registerChunkPublicationCapability } from '~/stream/internal/chunkPublication'; +import { buildPendingAction, buildToolApprovalPayload } from '~/agents/hitl/policy'; +import { InMemoryJobStore } from '~/stream/implementations/InMemoryJobStore'; +import { GenerationJobManagerClass } from '~/stream/GenerationJobManager'; + +function createTelemetry(): jest.Mocked { + return { + mark: jest.fn(), + setStreamId: jest.fn(), + recordGenerationEvent: jest.fn().mockReturnValue(false), + end: jest.fn(), + }; +} + +function createManager(): GenerationJobManagerClass { + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport: new InMemoryEventTransport(), + isRedis: false, + }); + manager.initialize(); + return manager; +} + +function createPendingAction(streamId: string) { + return buildPendingAction( + buildToolApprovalPayload([ + { name: 'shell', arguments: { command: 'ls' }, tool_call_id: 'call-shutdown' }, + ]), + { + streamId, + conversationId: streamId, + runId: 'run-shutdown', + responseMessageId: 'response-shutdown', + }, + ); +} + +describe('GenerationJobManager startup telemetry', () => { + it('returns sanitized initial metadata from the atomic job creation', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const updateJob = jest.spyOn(jobStore, 'updateJob'); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + }); + manager.initialize(); + + const job = await manager.createJob('stream-initial-metadata', 'user-1', 'conversation-1', { + initialMetadata: { + userId: 'untrusted-user', + tenantId: 'untrusted-tenant', + conversationId: 'untrusted-conversation', + userMessage: { + messageId: 'message-1', + parentMessageId: 'parent-1', + }, + responseMessageId: 'response-1', + sender: 'Agent', + endpoint: 'agents', + iconURL: 'https://example.com/icon.png', + model: 'test-model', + agent_id: 'agent-1', + isTemporary: false, + promptTokens: 0, + discoveredTools: [], + pendingAction: { + actionId: 'untrusted-action', + streamId: 'stream-initial-metadata', + conversationId: 'conversation-1', + payload: { + type: 'ask_user_question', + question: { question: 'Should not be persisted?' }, + }, + createdAt: Date.now(), + }, + }, + }); + + expect(job.metadata).toMatchObject({ + userId: 'user-1', + conversationId: 'conversation-1', + userMessage: { + messageId: 'message-1', + parentMessageId: 'parent-1', + }, + responseMessageId: 'response-1', + sender: 'Agent', + endpoint: 'agents', + iconURL: 'https://example.com/icon.png', + model: 'test-model', + agent_id: 'agent-1', + isTemporary: false, + promptTokens: 0, + discoveredTools: [], + }); + expect(job.metadata.tenantId).toBeUndefined(); + expect(job.metadata.pendingAction).toBeUndefined(); + expect(updateJob).not.toHaveBeenCalled(); + + await manager.destroy(); + }); + + it('clears omitted per-turn identity when an in-memory job is replaced', async () => { + const manager = createManager(); + await manager.createJob('stream-replaced-metadata', 'user-1', 'conversation-1', { + initialMetadata: { + agent_id: 'agent-1', + isTemporary: true, + discoveredTools: ['deferred-tool'], + }, + }); + + const replacement = await manager.createJob( + 'stream-replaced-metadata', + 'user-1', + 'conversation-1', + ); + + expect(replacement.metadata.agent_id).toBeUndefined(); + expect(replacement.metadata.isTemporary).toBeUndefined(); + expect(replacement.metadata.discoveredTools).toBeUndefined(); + + await manager.destroy(); + }); + + it('preserves updateMetadata truthiness semantics through the shared sanitizer', async () => { + const manager = createManager(); + await manager.createJob('stream-metadata-update', 'user-1', 'conversation-1', { + initialMetadata: { + sender: 'Agent', + agent_id: 'agent-1', + isTemporary: true, + promptTokens: 42, + discoveredTools: ['deferred-tool'], + }, + }); + + await manager.updateMetadata('stream-metadata-update', { + sender: '', + agent_id: '', + isTemporary: false, + promptTokens: 0, + discoveredTools: [], + }); + + const updated = await manager.getJob('stream-metadata-update'); + expect(updated?.metadata).toMatchObject({ + sender: 'Agent', + agent_id: 'agent-1', + isTemporary: false, + promptTokens: 0, + discoveredTools: [], + }); + + await manager.destroy(); + }); + + it('records accepted events centrally and detaches after the first content delta', async () => { + const manager = createManager(); + const telemetry = createTelemetry(); + telemetry.recordGenerationEvent.mockReturnValueOnce(false).mockReturnValueOnce(true); + await manager.createJob('stream-1', 'user-1', 'conversation-1', { + startupTelemetry: telemetry, + }); + + await manager.emitChunk('stream-1', { + created: true, + message: { + messageId: 'message-1', + sender: 'User', + isCreatedByUser: true, + }, + streamId: 'stream-1', + }); + await manager.emitChunk('stream-1', { + event: 'on_run_step', + data: { id: 'step-1' }, + }); + await manager.emitChunk('stream-1', { + event: 'on_message_delta', + data: { delta: { content: [{ type: 'text', text: 'Hello' }] } }, + }); + await manager.emitChunk('stream-1', { + event: 'on_message_delta', + data: { delta: { content: [{ type: 'text', text: ' later token' }] } }, + }); + + expect(telemetry.mark).toHaveBeenCalledWith('request_message_queued'); + expect(telemetry.recordGenerationEvent).toHaveBeenCalledTimes(2); + expect(telemetry.recordGenerationEvent).toHaveBeenNthCalledWith(1, { + event: 'on_run_step', + data: { id: 'step-1' }, + }); + expect(telemetry.recordGenerationEvent).toHaveBeenNthCalledWith(2, { + event: 'on_message_delta', + data: { delta: { content: [{ type: 'text', text: 'Hello' }] } }, + }); + + await manager.destroy(); + }); + + it('records a final-only response event before completion', async () => { + const manager = createManager(); + const telemetry = createTelemetry(); + await manager.createJob('stream-2', 'user-1', 'conversation-2', { + startupTelemetry: telemetry, + }); + + const finalEvent: ServerSentEvent = { + final: true, + responseMessage: { + messageId: 'response-1', + content: [{ type: 'text', text: 'Complete response' }], + }, + }; + await manager.emitDone('stream-2', finalEvent); + await manager.completeJob('stream-2'); + + expect(telemetry.recordGenerationEvent).toHaveBeenCalledWith(finalEvent); + expect(telemetry.end).toHaveBeenCalledWith('completed_without_delta'); + + await manager.destroy(); + }); + + it('does not record an event when delivery is rejected for an active subscriber', async () => { + const eventTransport = new InMemoryEventTransport(); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport, + isRedis: false, + }); + manager.initialize(); + const telemetry = createTelemetry(); + await manager.createJob('stream-3', 'user-1', 'conversation-3', { + startupTelemetry: telemetry, + }); + const subscription = await manager.subscribe('stream-3', () => undefined); + registerChunkPublicationCapability(eventTransport, async () => false); + + await manager.emitChunk('stream-3', { + event: 'on_message_delta', + data: { delta: { content: [{ type: 'text', text: 'Dropped' }] } }, + }); + + expect(telemetry.recordGenerationEvent).not.toHaveBeenCalled(); + + subscription?.unsubscribe(); + await manager.destroy(); + }); + + it('bypasses publication receipts after startup telemetry completes', async () => { + const eventTransport = new InMemoryEventTransport(); + const emitChunk = jest.spyOn(eventTransport, 'emitChunk'); + const publishWithReceipt = jest.fn().mockResolvedValue(0); + registerChunkPublicationCapability(eventTransport, publishWithReceipt); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport, + isRedis: false, + }); + manager.initialize(); + const telemetry = createTelemetry(); + telemetry.recordGenerationEvent.mockReturnValue(true); + const job = await manager.createJob('stream-hot-path', 'user-1', 'conversation-1', { + startupTelemetry: telemetry, + }); + const subscription = await manager.subscribe('stream-hot-path', () => undefined); + const firstDelta: ServerSentEvent = { + event: 'on_message_delta', + data: { delta: { content: [{ type: 'text', text: 'First' }] } }, + }; + const laterDelta: ServerSentEvent = { + event: 'on_message_delta', + data: { delta: { content: [{ type: 'text', text: ' later' }] } }, + }; + + await manager.emitChunk('stream-hot-path', firstDelta); + await manager.emitChunk('stream-hot-path', laterDelta); + + expect(publishWithReceipt).toHaveBeenCalledTimes(1); + expect(publishWithReceipt).toHaveBeenCalledWith('stream-hot-path', firstDelta, job.createdAt); + expect(emitChunk).toHaveBeenCalledTimes(1); + expect(emitChunk).toHaveBeenCalledWith('stream-hot-path', laterDelta, job.createdAt); + + subscription?.unsubscribe(); + await manager.destroy(); + }); + + it('does not let a later delta overtake a created event metadata write', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const originalUpdateJob = jobStore.updateJob.bind(jobStore); + let signalCreatedWriteStarted: (() => void) | undefined; + const createdWriteStarted = new Promise((resolve) => { + signalCreatedWriteStarted = resolve; + }); + let releaseCreatedWrite: (() => void) | undefined; + const createdWriteGate = new Promise((resolve) => { + releaseCreatedWrite = resolve; + }); + jest + .spyOn(jobStore, 'updateJob') + .mockImplementation(async (streamId, updates, expectedCreatedAt) => { + if (updates.createdEventEmitted === true) { + signalCreatedWriteStarted?.(); + await createdWriteGate; + } + return originalUpdateJob(streamId, updates, expectedCreatedAt); + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + }); + manager.initialize(); + await manager.createJob('stream-created-order', 'user-1', 'conversation-1'); + + const createdEvent: ServerSentEvent = { + created: true, + message: { + messageId: 'message-1', + conversationId: 'conversation-1', + sender: 'User', + isCreatedByUser: true, + }, + streamId: 'stream-created-order', + }; + const deltaEvent: ServerSentEvent = { + event: 'on_message_delta', + data: { delta: { content: [{ type: 'text', text: 'Hello' }] } }, + }; + const createdPublication = manager.emitChunk('stream-created-order', createdEvent); + await createdWriteStarted; + + let deltaSettled = false; + const deltaPublication = manager.emitChunk('stream-created-order', deltaEvent).then(() => { + deltaSettled = true; + }); + await Promise.resolve(); + expect(deltaSettled).toBe(false); + + releaseCreatedWrite?.(); + await Promise.all([createdPublication, deltaPublication]); + + const received: ServerSentEvent[] = []; + const subscription = await manager.subscribe('stream-created-order', (event) => + received.push(event), + ); + expect(received).toEqual([createdEvent, deltaEvent]); + + subscription?.unsubscribe(); + await manager.destroy(); + }); + + it('ends an active startup when the manager shuts down', async () => { + const manager = createManager(); + const telemetry = createTelemetry(); + const job = await manager.createJob('stream-4', 'user-1', 'conversation-4', { + startupTelemetry: telemetry, + }); + + await manager.destroy(); + + expect(telemetry.end).toHaveBeenCalledWith('aborted'); + expect(job.abortController.signal.aborted).toBe(true); + }); + + it('ends replaced startup telemetry before aborting the old runtime', async () => { + const manager = createManager(); + const telemetry = createTelemetry(); + const oldJob = await manager.createJob('stream-5', 'user-1', 'conversation-5', { + startupTelemetry: telemetry, + }); + oldJob.abortController.signal.addEventListener('abort', () => { + expect(telemetry.end).toHaveBeenCalledWith('replaced'); + }); + + await manager.createJob('stream-5', 'user-1', 'conversation-5'); + + expect(oldJob.abortController.signal.aborted).toBe(true); + expect(telemetry.end).toHaveBeenCalledWith('replaced'); + await manager.destroy(); + }); + + it('shares one lazy runtime across concurrent first subscriptions', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const eventTransport = new InMemoryEventTransport(); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false }); + manager.initialize(); + await jobStore.createJob('stream-lazy-concurrent', 'user-1', 'conversation-1'); + + const originalGetJob = jobStore.getJob.bind(jobStore); + let lookupCount = 0; + let signalLookupsStarted: (() => void) | undefined; + const lookupsStarted = new Promise((resolve) => { + signalLookupsStarted = resolve; + }); + let releaseLookups: (() => void) | undefined; + const lookupGate = new Promise((resolve) => { + releaseLookups = resolve; + }); + jest.spyOn(jobStore, 'getJob').mockImplementation(async (...args) => { + const job = await originalGetJob(...args); + if (lookupCount < 2) { + lookupCount++; + if (lookupCount === 2) { + signalLookupsStarted?.(); + } + await lookupGate; + } + return job; + }); + const allSubscribersLeftSpy = jest.spyOn(eventTransport, 'onAllSubscribersLeft'); + + const firstSubscription = manager.subscribe('stream-lazy-concurrent', () => undefined); + const secondSubscription = manager.subscribe('stream-lazy-concurrent', () => undefined); + await lookupsStarted; + releaseLookups?.(); + + const [first, second] = await Promise.all([firstSubscription, secondSubscription]); + + expect(first).not.toBeNull(); + expect(second).not.toBeNull(); + expect(allSubscribersLeftSpy).toHaveBeenCalledTimes(1); + expect(eventTransport.getSubscriberCount('stream-lazy-concurrent')).toBe(2); + + first?.unsubscribe(); + second?.unsubscribe(); + await manager.destroy(); + }); + + it('does not attach a subscription to a runtime replaced during the job lookup', async () => { + const now = jest.spyOn(Date, 'now').mockReturnValue(1000); + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const eventTransport = new InMemoryEventTransport(); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false }); + manager.initialize(); + const oldJob = await manager.createJob('stream-runtime-replaced', 'user-1', 'conversation-1'); + + const originalGetJob = jobStore.getJob.bind(jobStore); + let signalLookupStarted: (() => void) | undefined; + const lookupStarted = new Promise((resolve) => { + signalLookupStarted = resolve; + }); + let releaseLookup: (() => void) | undefined; + const lookupGate = new Promise((resolve) => { + releaseLookup = resolve; + }); + jest + .spyOn(jobStore, 'getJob') + .mockImplementation(originalGetJob) + .mockImplementationOnce(async (...args) => { + const job = await originalGetJob(...args); + signalLookupStarted?.(); + await lookupGate; + return job; + }); + + const staleSubscription = manager.subscribe('stream-runtime-replaced', () => undefined); + await lookupStarted; + now.mockReturnValue(2000); + const replacementJob = await manager.createJob( + 'stream-runtime-replaced', + 'user-1', + 'conversation-1', + ); + now.mockRestore(); + releaseLookup?.(); + + await expect(staleSubscription).resolves.toBeNull(); + expect(oldJob.abortController.signal.aborted).toBe(true); + expect(replacementJob.abortController.signal.aborted).toBe(false); + expect(eventTransport.getSubscriberCount('stream-runtime-replaced')).toBe(0); + + const currentSubscription = await manager.subscribe('stream-runtime-replaced', () => undefined); + expect(currentSubscription).not.toBeNull(); + + currentSubscription?.unsubscribe(); + await manager.destroy(); + }); + + it('does not return a lazy runtime replaced while its abort listener activates', async () => { + const now = jest.spyOn(Date, 'now').mockReturnValue(1000); + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + let signalAbortRegistrationStarted: (() => void) | undefined; + const abortRegistrationStarted = new Promise((resolve) => { + signalAbortRegistrationStarted = resolve; + }); + let releaseAbortRegistration: (() => void) | undefined; + const abortRegistrationGate = new Promise((resolve) => { + releaseAbortRegistration = resolve; + }); + const eventTransport = Object.assign(new InMemoryEventTransport(), { + onAbort: jest + .fn() + .mockImplementationOnce(async () => { + signalAbortRegistrationStarted?.(); + await abortRegistrationGate; + }) + .mockResolvedValue(undefined), + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false }); + manager.initialize(); + + try { + await jobStore.createJob('stream-lazy-abort-race', 'user-1'); + const staleLookup = manager.getJob('stream-lazy-abort-race'); + await abortRegistrationStarted; + + now.mockReturnValue(2000); + const replacement = await jobStore.createJob('stream-lazy-abort-race', 'user-1'); + releaseAbortRegistration?.(); + + await expect(staleLookup).resolves.toBeUndefined(); + await expect(manager.getJob('stream-lazy-abort-race')).resolves.toMatchObject({ + createdAt: replacement.createdAt, + }); + expect(eventTransport.onAbort).toHaveBeenCalledTimes(2); + } finally { + releaseAbortRegistration?.(); + now.mockRestore(); + await manager.destroy(); + } + }); + + it('keeps a replacement generation intact when its predecessor completes late', async () => { + const now = jest.spyOn(Date, 'now').mockReturnValue(1000); + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const eventTransport = new InMemoryEventTransport(); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false }); + manager.initialize(); + + try { + const predecessor = await manager.createJob('stream-late-complete', 'user-1'); + now.mockReturnValue(2000); + const replacement = await manager.createJob('stream-late-complete', 'user-1'); + const replacementDone = jest.fn(); + const subscription = await manager.subscribe( + 'stream-late-complete', + () => undefined, + replacementDone, + ); + const staleFinal: ServerSentEvent = { + final: true, + responseMessage: { text: 'stale' }, + }; + + await manager.emitDone('stream-late-complete', staleFinal, predecessor.createdAt); + await manager.completeJob('stream-late-complete', undefined, predecessor.createdAt); + + await expect(jobStore.getJob('stream-late-complete')).resolves.toMatchObject({ + createdAt: replacement.createdAt, + status: 'running', + }); + expect((await jobStore.getJob('stream-late-complete'))?.finalEvent).toBeUndefined(); + expect(replacement.abortController.signal.aborted).toBe(false); + expect(replacementDone).not.toHaveBeenCalled(); + + subscription?.unsubscribe(); + } finally { + now.mockRestore(); + await manager.destroy(); + } + }); + + it('does not abort a replacement installed while an abort lookup is pending', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const abortDisposers: jest.Mock[] = []; + const eventTransport = Object.assign(new InMemoryEventTransport(), { + onAbort: jest.fn(async () => { + const dispose = jest.fn(); + abortDisposers.push(dispose); + return dispose; + }), + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false }); + manager.initialize(); + + let signalLookupStarted: (() => void) | undefined; + const lookupStarted = new Promise((resolve) => { + signalLookupStarted = resolve; + }); + let releaseLookup: (() => void) | undefined; + const lookupGate = new Promise((resolve) => { + releaseLookup = resolve; + }); + + try { + const predecessor = await manager.createJob('stream-abort-lookup-race', 'user-1'); + await jobStore.deleteJob('stream-abort-lookup-race', predecessor.createdAt); + jest.spyOn(jobStore, 'getJob').mockImplementationOnce(async () => { + signalLookupStarted?.(); + await lookupGate; + return null; + }); + + const aborting = manager.abortJob('stream-abort-lookup-race'); + await lookupStarted; + const replacement = await manager.createJob('stream-abort-lookup-race', 'user-1'); + releaseLookup?.(); + + await expect(aborting).resolves.toMatchObject({ success: false, jobData: null }); + expect(replacement.abortController.signal.aborted).toBe(false); + expect(abortDisposers).toHaveLength(2); + expect(abortDisposers[0]).toHaveBeenCalledTimes(1); + expect(abortDisposers[1]).not.toHaveBeenCalled(); + } finally { + releaseLookup?.(); + await manager.destroy(); + } + }); + + it('releases its exact runtime when completion loses to a terminal transition', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const abortDisposer = jest.fn(); + const eventTransport = Object.assign(new InMemoryEventTransport(), { + onAbort: jest.fn(async () => abortDisposer), + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport, + isRedis: false, + cleanupOnComplete: false, + }); + manager.initialize(); + const originalTransition = jobStore.transitionStatus.bind(jobStore); + let signalTransitionStarted: (() => void) | undefined; + const transitionStarted = new Promise((resolve) => { + signalTransitionStarted = resolve; + }); + let releaseTransition: (() => void) | undefined; + const transitionGate = new Promise((resolve) => { + releaseTransition = resolve; + }); + + try { + const job = await manager.createJob('stream-complete-cas-loss', 'user-1'); + const subscription = await manager.subscribe('stream-complete-cas-loss', () => undefined); + jest.spyOn(jobStore, 'transitionStatus').mockImplementationOnce(async (...args) => { + signalTransitionStarted?.(); + await transitionGate; + return originalTransition(...args); + }); + + const completing = manager.completeJob('stream-complete-cas-loss', undefined, job.createdAt); + await transitionStarted; + await originalTransition('stream-complete-cas-loss', { + from: 'running', + to: 'error', + expectCreatedAt: job.createdAt, + patch: { completedAt: Date.now(), error: 'competing terminal transition' }, + }); + releaseTransition?.(); + await completing; + + expect(job.abortController.signal.aborted).toBe(true); + expect(abortDisposer).toHaveBeenCalledTimes(1); + expect(eventTransport.getSubscriberCount('stream-complete-cas-loss')).toBe(1); + expect(manager.getRuntimeStats().runtimeStateSize).toBe(1); + subscription?.unsubscribe(); + } finally { + releaseTransition?.(); + await manager.destroy(); + } + }); + + it('releases its exact runtime when abort loses to a terminal transition', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const abortDisposer = jest.fn(); + const eventTransport = Object.assign(new InMemoryEventTransport(), { + onAbort: jest.fn(async () => abortDisposer), + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport, + isRedis: false, + cleanupOnComplete: false, + }); + manager.initialize(); + const originalTransition = jobStore.transitionStatus.bind(jobStore); + let signalTransitionStarted: (() => void) | undefined; + const transitionStarted = new Promise((resolve) => { + signalTransitionStarted = resolve; + }); + let releaseTransition: (() => void) | undefined; + const transitionGate = new Promise((resolve) => { + releaseTransition = resolve; + }); + + try { + const job = await manager.createJob('stream-abort-cas-loss', 'user-1'); + const subscription = await manager.subscribe('stream-abort-cas-loss', () => undefined); + jest.spyOn(jobStore, 'transitionStatus').mockImplementationOnce(async (...args) => { + signalTransitionStarted?.(); + await transitionGate; + return originalTransition(...args); + }); + + const aborting = manager.abortJob('stream-abort-cas-loss'); + await transitionStarted; + await originalTransition('stream-abort-cas-loss', { + from: 'running', + to: 'complete', + expectCreatedAt: job.createdAt, + patch: { completedAt: Date.now() }, + }); + releaseTransition?.(); + + await expect(aborting).resolves.toMatchObject({ success: false, finalEvent: null }); + expect(job.abortController.signal.aborted).toBe(true); + expect(abortDisposer).toHaveBeenCalledTimes(1); + expect(eventTransport.getSubscriberCount('stream-abort-cas-loss')).toBe(1); + expect(manager.getRuntimeStats().runtimeStateSize).toBe(1); + subscription?.unsubscribe(); + } finally { + releaseTransition?.(); + await manager.destroy(); + } + }); + + it('retains a same-epoch run that resumes while abort is claiming terminal state', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + const abortDisposer = jest.fn(); + const eventTransport = Object.assign(new InMemoryEventTransport(), { + onAbort: jest.fn(async () => abortDisposer), + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false }); + manager.initialize(); + const originalTransition = jobStore.transitionStatus.bind(jobStore); + let signalAbortTransitionStarted: (() => void) | undefined; + const abortTransitionStarted = new Promise((resolve) => { + signalAbortTransitionStarted = resolve; + }); + let releaseAbortTransition: (() => void) | undefined; + const abortTransitionGate = new Promise((resolve) => { + releaseAbortTransition = resolve; + }); + let gateAbortTransition = true; + jest.spyOn(jobStore, 'transitionStatus').mockImplementation(async (...args) => { + if (gateAbortTransition && args[1].to === 'aborted') { + gateAbortTransition = false; + signalAbortTransitionStarted?.(); + await abortTransitionGate; + } + return originalTransition(...args); + }); + + try { + const streamId = 'stream-abort-resume-race'; + const job = await manager.createJob(streamId, 'user-1'); + const pendingAction = createPendingAction(streamId); + await expect(manager.approvals.pause(streamId, pendingAction)).resolves.toBe(true); + + const aborting = manager.abortJob(streamId); + await abortTransitionStarted; + await expect(manager.approvals.resolve(streamId, pendingAction.actionId)).resolves.toBe(true); + releaseAbortTransition?.(); + + await expect(aborting).resolves.toMatchObject({ success: false, finalEvent: null }); + expect(job.abortController.signal.aborted).toBe(false); + expect(abortDisposer).not.toHaveBeenCalled(); + await expect(jobStore.getJob(streamId)).resolves.toMatchObject({ + createdAt: job.createdAt, + status: 'running', + }); + + manager.prepareForShutdown(); + await manager.destroy(); + await expect(jobStore.getJob(streamId)).resolves.toMatchObject({ + createdAt: job.createdAt, + status: 'error', + error: 'Generation interrupted because its server shut down', + }); + } finally { + releaseAbortTransition?.(); + if (manager.getRuntimeStats().runtimeStateSize > 0) { + await manager.destroy(); + } + } + }); + + it('releases ownership but keeps the abort runtime when completion observes a pause', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + const abortDisposer = jest.fn(); + const eventTransport = Object.assign(new InMemoryEventTransport(), { + onAbort: jest.fn(async () => abortDisposer), + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false }); + manager.initialize(); + const streamId = 'stream-complete-paused'; + const job = await manager.createJob(streamId, 'user-1'); + await jobStore.transitionStatus(streamId, { + from: 'running', + to: 'requires_action', + expectCreatedAt: job.createdAt, + }); + + await manager.completeJob(streamId, undefined, job.createdAt); + + expect(job.abortController.signal.aborted).toBe(false); + expect(abortDisposer).not.toHaveBeenCalled(); + await expect(jobStore.getJob(streamId)).resolves.toMatchObject({ + createdAt: job.createdAt, + status: 'requires_action', + }); + + await jobStore.transitionStatus(streamId, { + from: 'requires_action', + to: 'running', + expectCreatedAt: job.createdAt, + }); + manager.prepareForShutdown(); + await manager.destroy(); + await expect(jobStore.getJob(streamId)).resolves.toMatchObject({ + createdAt: job.createdAt, + status: 'running', + }); + }); + + it('releases ownership when completion observes that its durable job is missing', async () => { + const now = jest.spyOn(Date, 'now').mockReturnValue(1000); + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + const abortDisposer = jest.fn(); + const eventTransport = Object.assign(new InMemoryEventTransport(), { + onAbort: jest.fn(async () => abortDisposer), + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false }); + manager.initialize(); + + try { + const streamId = 'stream-complete-missing'; + const job = await manager.createJob(streamId, 'user-1'); + await jobStore.deleteJob(streamId, job.createdAt); + + await manager.completeJob(streamId, undefined, job.createdAt); + + expect(job.abortController.signal.aborted).toBe(true); + expect(abortDisposer).toHaveBeenCalledTimes(1); + const restored = await jobStore.createJob(streamId, 'user-1'); + expect(restored.createdAt).toBe(job.createdAt); + + manager.prepareForShutdown(); + await manager.destroy(); + await expect(jobStore.getJob(streamId)).resolves.toMatchObject({ + createdAt: job.createdAt, + status: 'running', + }); + } finally { + now.mockRestore(); + if (manager.getRuntimeStats().runtimeStateSize > 0) { + await manager.destroy(); + } + } + }); + + it('releases ownership when abort observes that its durable job is missing', async () => { + const now = jest.spyOn(Date, 'now').mockReturnValue(1000); + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + const abortDisposer = jest.fn(); + const eventTransport = Object.assign(new InMemoryEventTransport(), { + onAbort: jest.fn(async () => abortDisposer), + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false }); + manager.initialize(); + + try { + const streamId = 'stream-abort-missing'; + const job = await manager.createJob(streamId, 'user-1'); + await jobStore.deleteJob(streamId, job.createdAt); + + await expect(manager.abortJob(streamId)).resolves.toMatchObject({ + success: false, + jobData: null, + }); + + expect(job.abortController.signal.aborted).toBe(true); + expect(abortDisposer).toHaveBeenCalledTimes(1); + const restored = await jobStore.createJob(streamId, 'user-1'); + expect(restored.createdAt).toBe(job.createdAt); + + manager.prepareForShutdown(); + await manager.destroy(); + await expect(jobStore.getJob(streamId)).resolves.toMatchObject({ + createdAt: job.createdAt, + status: 'running', + }); + } finally { + now.mockRestore(); + if (manager.getRuntimeStats().runtimeStateSize > 0) { + await manager.destroy(); + } + } + }); + + it('releases ownership when abort observes an already-terminal durable job', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + const abortDisposer = jest.fn(); + const eventTransport = Object.assign(new InMemoryEventTransport(), { + onAbort: jest.fn(async () => abortDisposer), + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false }); + manager.initialize(); + const streamId = 'stream-abort-terminal'; + const job = await manager.createJob(streamId, 'user-1'); + await jobStore.transitionStatus(streamId, { + from: 'running', + to: 'complete', + expectCreatedAt: job.createdAt, + patch: { completedAt: Date.now() }, + }); + + await expect(manager.abortJob(streamId)).resolves.toMatchObject({ + success: false, + jobData: expect.objectContaining({ + createdAt: job.createdAt, + status: 'complete', + }), + }); + + expect(job.abortController.signal.aborted).toBe(true); + expect(abortDisposer).toHaveBeenCalledTimes(1); + await jobStore.transitionStatus(streamId, { + from: 'complete', + to: 'running', + expectCreatedAt: job.createdAt, + }); + + manager.prepareForShutdown(); + await manager.destroy(); + await expect(jobStore.getJob(streamId)).resolves.toMatchObject({ + createdAt: job.createdAt, + status: 'running', + }); + }); + + it('does not let a delayed predecessor metadata write contaminate a replacement', async () => { + const now = jest.spyOn(Date, 'now').mockReturnValue(1000); + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const originalUpdateJob = jobStore.updateJob.bind(jobStore); + let signalTitleWriteStarted: (() => void) | undefined; + const titleWriteStarted = new Promise((resolve) => { + signalTitleWriteStarted = resolve; + }); + let releaseTitleWrite: (() => void) | undefined; + const titleWriteGate = new Promise((resolve) => { + releaseTitleWrite = resolve; + }); + jest + .spyOn(jobStore, 'updateJob') + .mockImplementation(async (streamId, updates, expectedCreatedAt) => { + if (updates.titleEvent) { + signalTitleWriteStarted?.(); + await titleWriteGate; + } + return originalUpdateJob(streamId, updates, expectedCreatedAt); + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + cleanupOnComplete: false, + }); + manager.initialize(); + + try { + await manager.createJob('stream-metadata-epoch', 'user-1'); + const emitting = manager.emitChunk('stream-metadata-epoch', { + event: 'title', + data: { title: 'stale title' }, + }); + await titleWriteStarted; + now.mockReturnValue(2000); + const replacement = await manager.createJob('stream-metadata-epoch', 'user-1'); + releaseTitleWrite?.(); + await emitting; + + await expect(jobStore.getJob('stream-metadata-epoch')).resolves.toMatchObject({ + createdAt: replacement.createdAt, + status: 'running', + }); + expect((await jobStore.getJob('stream-metadata-epoch'))?.titleEvent).toBeUndefined(); + } finally { + releaseTitleWrite?.(); + now.mockRestore(); + await manager.destroy(); + } + }); + + it('does not let predecessor producers repopulate replacement volatile state', async () => { + const now = jest.spyOn(Date, 'now').mockReturnValue(1000); + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + cleanupOnComplete: false, + }); + manager.initialize(); + + try { + const predecessor = await manager.createJob('stream-volatile-epoch', 'user-1'); + now.mockReturnValue(2000); + const replacement = await manager.createJob('stream-volatile-epoch', 'user-1'); + const staleContent: Parameters[1] = [ + { type: 'text', text: 'predecessor' }, + ]; + const staleGraph = { + contentData: [{ id: 'predecessor-step' }], + }; + + manager.setContentParts('stream-volatile-epoch', staleContent, predecessor.createdAt); + manager.setCollectedUsage( + 'stream-volatile-epoch', + [{ input_tokens: 10 }], + predecessor.createdAt, + ); + manager.setGraph('stream-volatile-epoch', staleGraph as never, predecessor.createdAt); + + await expect(jobStore.getContentParts('stream-volatile-epoch')).resolves.toBeNull(); + expect(jobStore.getCollectedUsage('stream-volatile-epoch')).toEqual([]); + await expect(jobStore.getRunSteps('stream-volatile-epoch')).resolves.toEqual([]); + + const currentContent: Parameters[1] = [ + { type: 'text', text: 'replacement' }, + ]; + const currentGraph = { + contentData: [{ id: 'replacement-step' }], + }; + manager.setContentParts('stream-volatile-epoch', currentContent, replacement.createdAt); + manager.setCollectedUsage( + 'stream-volatile-epoch', + [{ input_tokens: 20 }], + replacement.createdAt, + ); + manager.setGraph('stream-volatile-epoch', currentGraph as never, replacement.createdAt); + + await expect(jobStore.getContentParts('stream-volatile-epoch')).resolves.toEqual({ + content: currentContent, + }); + expect(jobStore.getCollectedUsage('stream-volatile-epoch')).toEqual([{ input_tokens: 20 }]); + await expect(jobStore.getRunSteps('stream-volatile-epoch')).resolves.toEqual( + currentGraph.contentData, + ); + } finally { + now.mockRestore(); + await manager.destroy(); + } + }); + + it('reconciles a stale local runtime before delivering a replacement terminal event', async () => { + const now = jest.spyOn(Date, 'now').mockReturnValue(1000); + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const eventTransport = new InMemoryEventTransport(); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false, cleanupOnComplete: false }); + manager.initialize(); + + try { + const predecessor = await manager.createJob('stream-runtime-epoch', 'user-1'); + now.mockReturnValue(2000); + const replacement = await jobStore.createJob('stream-runtime-epoch', 'user-1'); + const onDone = jest.fn(); + const subscription = await manager.subscribe('stream-runtime-epoch', () => undefined, onDone); + const finalEvent: ServerSentEvent = { + final: true, + responseMessage: { text: 'replacement' }, + }; + + eventTransport.emitDone('stream-runtime-epoch', finalEvent, replacement.createdAt); + + expect(subscription).not.toBeNull(); + expect(predecessor.abortController.signal.aborted).toBe(true); + expect(onDone).toHaveBeenCalledWith(finalEvent); + subscription?.unsubscribe(); + } finally { + now.mockRestore(); + await manager.destroy(); + } + }); + + it('publishes a generation-scoped error when direct completion fails an active job', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + cleanupOnComplete: false, + }); + manager.initialize(); + const job = await manager.createJob('stream-direct-error', 'user-1'); + const onError = jest.fn(); + const subscription = await manager.subscribe( + 'stream-direct-error', + () => undefined, + undefined, + onError, + ); + + await manager.completeJob('stream-direct-error', 'initialization failed', job.createdAt); + + expect(onError).toHaveBeenCalledWith('initialization failed'); + await expect(jobStore.getJob('stream-direct-error')).resolves.toMatchObject({ + createdAt: job.createdAt, + status: 'error', + error: 'initialization failed', + }); + subscription?.unsubscribe(); + await manager.destroy(); + }); + + it('finalizes the exact durable job when abort-listener setup fails before createJob returns', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const eventTransport = Object.assign(new InMemoryEventTransport(), { + onAbort: jest.fn().mockRejectedValue(new Error('abort listener unavailable')), + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport, + isRedis: false, + cleanupOnComplete: false, + }); + manager.initialize(); + + await expect(manager.createJob('stream-partial-create', 'user-1')).rejects.toThrow( + 'abort listener unavailable', + ); + await expect(jobStore.getJob('stream-partial-create')).resolves.toMatchObject({ + status: 'error', + error: 'abort listener unavailable', + }); + + await manager.destroy(); + }); + + it('does not return a job replaced while its abort listener activates', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + let signalAbortRegistrationStarted: (() => void) | undefined; + const abortRegistrationStarted = new Promise((resolve) => { + signalAbortRegistrationStarted = resolve; + }); + let releaseAbortRegistration: (() => void) | undefined; + const abortRegistrationGate = new Promise((resolve) => { + releaseAbortRegistration = resolve; + }); + const eventTransport = Object.assign(new InMemoryEventTransport(), { + onAbort: jest + .fn() + .mockImplementationOnce(async () => { + signalAbortRegistrationStarted?.(); + await abortRegistrationGate; + }) + .mockResolvedValue(undefined), + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport, + isRedis: false, + cleanupOnComplete: false, + }); + manager.initialize(); + + try { + const predecessorCreation = manager.createJob('stream-create-replaced', 'user-1'); + await abortRegistrationStarted; + const replacement = await manager.createJob('stream-create-replaced', 'user-1'); + releaseAbortRegistration?.(); + + await expect(predecessorCreation).rejects.toThrow( + 'Generation job was replaced during initialization', + ); + await expect(jobStore.getJob('stream-create-replaced')).resolves.toMatchObject({ + createdAt: replacement.createdAt, + status: 'running', + }); + expect(replacement.abortController.signal.aborted).toBe(false); + } finally { + releaseAbortRegistration?.(); + await manager.destroy(); + } + }); + + it('rechecks durable ownership after its abort listener activates', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + let signalAbortRegistrationStarted: (() => void) | undefined; + const abortRegistrationStarted = new Promise((resolve) => { + signalAbortRegistrationStarted = resolve; + }); + let releaseAbortRegistration: (() => void) | undefined; + const abortRegistrationGate = new Promise((resolve) => { + releaseAbortRegistration = resolve; + }); + const eventTransport = Object.assign(new InMemoryEventTransport(), { + onAbort: jest.fn(async () => { + signalAbortRegistrationStarted?.(); + await abortRegistrationGate; + }), + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport, + isRedis: false, + cleanupOnComplete: false, + }); + manager.initialize(); + + try { + const predecessorCreation = manager.createJob('stream-remote-replacement', 'user-1'); + await abortRegistrationStarted; + const replacement = await jobStore.createJob('stream-remote-replacement', 'user-1'); + releaseAbortRegistration?.(); + + await expect(predecessorCreation).rejects.toThrow( + 'Generation job was replaced during initialization', + ); + await expect(jobStore.getJob('stream-remote-replacement')).resolves.toMatchObject({ + createdAt: replacement.createdAt, + status: 'running', + }); + await expect(manager.getJob('stream-remote-replacement')).resolves.toMatchObject({ + createdAt: replacement.createdAt, + }); + } finally { + releaseAbortRegistration?.(); + await manager.destroy(); + } + }); + + it('terminalizes a created epoch when shutdown starts during abort-listener readiness', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + let signalAbortRegistrationStarted: (() => void) | undefined; + const abortRegistrationStarted = new Promise((resolve) => { + signalAbortRegistrationStarted = resolve; + }); + let releaseAbortRegistration: (() => void) | undefined; + const abortRegistrationGate = new Promise((resolve) => { + releaseAbortRegistration = resolve; + }); + const eventTransport = Object.assign(new InMemoryEventTransport(), { + onAbort: jest.fn(async () => { + signalAbortRegistrationStarted?.(); + await abortRegistrationGate; + return jest.fn(); + }), + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport, + isRedis: false, + cleanupOnComplete: false, + }); + manager.initialize(); + const streamId = 'stream-create-shutdown-readiness'; + + try { + const creating = manager.createJob(streamId, 'user-1'); + await abortRegistrationStarted; + manager.prepareForShutdown(); + releaseAbortRegistration?.(); + + await expect(creating).rejects.toThrow('Generation job manager is shutting down'); + await manager.destroy(); + await expect(jobStore.getJob(streamId)).resolves.toMatchObject({ + status: 'error', + error: 'Generation interrupted because its server shut down', + completedAt: expect.any(Number), + }); + } finally { + releaseAbortRegistration?.(); + if (manager.getRuntimeStats().runtimeStateSize > 0) { + await manager.destroy(); + } + } + }); + + it('does not install a delayed older create over a completed replacement create', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + const createJob = jobStore.createJob.bind(jobStore); + let signalFirstDurableWrite: (() => void) | undefined; + const firstDurableWrite = new Promise((resolve) => { + signalFirstDurableWrite = resolve; + }); + let releaseFirstCreate: (() => void) | undefined; + const firstCreateGate = new Promise((resolve) => { + releaseFirstCreate = resolve; + }); + let creationCount = 0; + jest.spyOn(jobStore, 'createJob').mockImplementation(async (...args) => { + const job = await createJob(...args); + creationCount++; + if (creationCount === 1) { + signalFirstDurableWrite?.(); + await firstCreateGate; + } + return job; + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + cleanupOnComplete: false, + }); + manager.initialize(); + + try { + const predecessorCreation = manager.createJob('stream-delayed-create', 'user-1'); + await firstDurableWrite; + const replacement = await manager.createJob('stream-delayed-create', 'user-1'); + releaseFirstCreate?.(); + + await expect(predecessorCreation).rejects.toThrow( + 'Generation job was replaced during initialization', + ); + expect(replacement.abortController.signal.aborted).toBe(false); + await expect(manager.getJob('stream-delayed-create')).resolves.toMatchObject({ + createdAt: replacement.createdAt, + }); + await expect(jobStore.getJob('stream-delayed-create')).resolves.toMatchObject({ + createdAt: replacement.createdAt, + status: 'running', + }); + + manager.prepareForShutdown(); + await manager.destroy(); + await expect(jobStore.getJob('stream-delayed-create')).resolves.toMatchObject({ + createdAt: replacement.createdAt, + status: 'error', + }); + } finally { + releaseFirstCreate?.(); + if (manager.getRuntimeStats().runtimeStateSize > 0) { + await manager.destroy(); + } + } + }); + + it('closes local subscribers before drain without broadcasting an abort', async () => { + const manager = createManager(); + const telemetry = createTelemetry(); + const job = await manager.createJob('stream-6', 'user-1', 'conversation-6', { + startupTelemetry: telemetry, + }); + const onError = jest.fn(); + const subscription = await manager.subscribe('stream-6', () => undefined, undefined, onError); + + manager.prepareForShutdown(); + + expect(onError).toHaveBeenCalledWith('Server is shutting down'); + expect(telemetry.end).toHaveBeenCalledWith('aborted'); + expect(job.abortController.signal.aborted).toBe(false); + await expect(manager.createJob('stream-after-shutdown', 'user-1')).rejects.toThrow( + 'Generation job manager is shutting down', + ); + + subscription?.unsubscribe(); + await manager.destroy(); + }); + + it('durably finalizes locally owned running jobs after the HTTP drain', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + const eventTransport = new InMemoryEventTransport(); + const emitError = jest.spyOn(eventTransport, 'emitError'); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport, + isRedis: false, + }); + manager.initialize(); + const job = await manager.createJob('stream-shutdown-owned', 'user-1'); + + manager.prepareForShutdown(); + await manager.destroy(); + + await expect(jobStore.getJob('stream-shutdown-owned')).resolves.toMatchObject({ + status: 'error', + error: 'Generation interrupted because its server shut down', + completedAt: expect.any(Number), + }); + expect(emitError).toHaveBeenCalledWith( + 'stream-shutdown-owned', + 'Generation interrupted because its server shut down', + job.createdAt, + ); + }); + + it('persists a partial response before shutdown terminal cleanup without a local subscriber', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + }); + manager.initialize(); + const job = await manager.createJob('stream-shutdown-partial', 'user-1'); + const partial = [{ type: 'text', text: 'partial response' }]; + jest.spyOn(jobStore, 'getContentParts').mockResolvedValue({ content: partial }); + const persisted = jest.fn(async (parts) => { + expect(parts).toEqual(partial); + await expect(jobStore.getJob('stream-shutdown-partial')).resolves.toMatchObject({ + status: 'running', + }); + }); + job.emitter.on('allSubscribersLeft', persisted); + + manager.prepareForShutdown(); + await manager.destroy(); + + expect(persisted).toHaveBeenCalledTimes(1); + await expect(jobStore.getJob('stream-shutdown-partial')).resolves.toMatchObject({ + status: 'error', + error: 'Generation interrupted because its server shut down', + }); + }); + + it('does not let late success overwrite or delete a shutdown error', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + const originalCloseAndDrain = jobStore.closeAndDrainSteers.bind(jobStore); + let signalCompletionStarted: (() => void) | undefined; + const completionStarted = new Promise((resolve) => { + signalCompletionStarted = resolve; + }); + let releaseCompletion: (() => void) | undefined; + const completionGate = new Promise((resolve) => { + releaseCompletion = resolve; + }); + jest.spyOn(jobStore, 'closeAndDrainSteers').mockImplementationOnce(async (...args) => { + signalCompletionStarted?.(); + await completionGate; + return originalCloseAndDrain(...args); + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + }); + manager.initialize(); + const job = await manager.createJob('stream-shutdown-wins', 'user-1'); + + const completing = manager.completeJob('stream-shutdown-wins', undefined, job.createdAt); + await completionStarted; + manager.prepareForShutdown(); + await manager.destroy(); + releaseCompletion?.(); + await completing; + + await expect(jobStore.getJob('stream-shutdown-wins')).resolves.toMatchObject({ + createdAt: job.createdAt, + status: 'error', + error: 'Generation interrupted because its server shut down', + }); + }); + + it('parks accepted steers before shutdown terminal cleanup removes the queue', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + }); + manager.initialize(); + await manager.createJob('stream-shutdown-steer', 'user-1'); + await manager.steering.enqueue('stream-shutdown-steer', { + steerId: 'steer-shutdown', + text: 'keep this', + userId: 'user-1', + createdAt: Date.now(), + }); + + manager.prepareForShutdown(); + await manager.destroy(); + + await expect( + manager.steering.claim('stream-shutdown-steer', { userId: 'user-1' }), + ).resolves.toEqual([expect.objectContaining({ steerId: 'steer-shutdown', text: 'keep this' })]); + }); + + it('does not finalize a running job owned by a different replica', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + await jobStore.createJob('stream-remote-owner', 'user-1'); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + }); + manager.initialize(); + const subscription = await manager.subscribe('stream-remote-owner', () => undefined); + + manager.prepareForShutdown(); + await manager.destroy(); + + await expect(jobStore.getJob('stream-remote-owner')).resolves.toMatchObject({ + status: 'running', + }); + subscription?.unsubscribe(); + }); + + it('transfers shutdown ownership to the replica that resumes a paused job', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + const owner = new GenerationJobManagerClass(); + const resumer = new GenerationJobManagerClass(); + owner.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + }); + resumer.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + }); + + await owner.createJob('stream-owner-transfer', 'user-1'); + const pendingAction = createPendingAction('stream-owner-transfer'); + await expect(owner.approvals.pause('stream-owner-transfer', pendingAction)).resolves.toBe(true); + await expect( + resumer.approvals.resolve('stream-owner-transfer', pendingAction.actionId), + ).resolves.toBe(true); + + await owner.destroy(); + await expect(jobStore.getJob('stream-owner-transfer')).resolves.toMatchObject({ + status: 'running', + }); + + await resumer.destroy(); + await expect(jobStore.getJob('stream-owner-transfer')).resolves.toMatchObject({ + status: 'error', + error: 'Generation interrupted because its server shut down', + }); + }); + + it('rejects a job when shutdown starts while its store write is pending', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + const originalCreateJob = jobStore.createJob.bind(jobStore); + let releaseCreate: (() => void) | undefined; + const createGate = new Promise((resolve) => { + releaseCreate = resolve; + }); + jest.spyOn(jobStore, 'createJob').mockImplementation(async (...args) => { + await createGate; + return originalCreateJob(...args); + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + }); + manager.initialize(); + + const creating = manager.createJob('stream-7', 'user-1', 'conversation-7'); + manager.prepareForShutdown(); + releaseCreate?.(); + + await expect(creating).rejects.toThrow('Generation job manager is shutting down'); + await manager.destroy(); + await expect(jobStore.getJob('stream-7')).resolves.toMatchObject({ + status: 'error', + error: 'Generation interrupted because its server shut down', + }); + }); + + it('does not attach a subscriber when shutdown starts during the job lookup', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const eventTransport = new InMemoryEventTransport(); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false }); + manager.initialize(); + await manager.createJob('stream-8', 'user-1', 'conversation-8'); + + const originalGetJob = jobStore.getJob.bind(jobStore); + let signalReadStarted: (() => void) | undefined; + const readStarted = new Promise((resolve) => { + signalReadStarted = resolve; + }); + let releaseRead: (() => void) | undefined; + const readGate = new Promise((resolve) => { + releaseRead = resolve; + }); + jest.spyOn(jobStore, 'getJob').mockImplementation(async (...args) => { + signalReadStarted?.(); + await readGate; + return originalGetJob(...args); + }); + + const onError = jest.fn(); + const subscribing = manager.subscribe('stream-8', () => undefined, undefined, onError); + await readStarted; + manager.prepareForShutdown(); + releaseRead?.(); + + await expect(subscribing).resolves.toBeNull(); + expect(onError).toHaveBeenCalledWith('Server is shutting down'); + expect(eventTransport.getSubscriberCount('stream-8')).toBe(0); + await manager.destroy(); + }); + + it('drains asynchronous disconnect handlers before destroying the job store', async () => { + const manager = createManager(); + const job = await manager.createJob('stream-9', 'user-1', 'conversation-9'); + let signalHandlerStarted: (() => void) | undefined; + const handlerStarted = new Promise((resolve) => { + signalHandlerStarted = resolve; + }); + let releaseHandler: (() => void) | undefined; + const handlerGate = new Promise((resolve) => { + releaseHandler = resolve; + }); + job.emitter.on('allSubscribersLeft', async () => { + signalHandlerStarted?.(); + await handlerGate; + }); + const subscription = await manager.subscribe('stream-9', () => undefined); + + manager.prepareForShutdown(); + await handlerStarted; + + let destroySettled = false; + const destroying = manager.destroy().then(() => { + destroySettled = true; + }); + await Promise.resolve(); + + expect(destroySettled).toBe(false); + releaseHandler?.(); + await destroying; + expect(destroySettled).toBe(true); + subscription?.unsubscribe(); + }); + + it('drains replaced services without destroying the newly configured services', async () => { + const replacedJobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const replacedEventTransport = new InMemoryEventTransport(); + const replacedStoreDestroy = jest.spyOn(replacedJobStore, 'destroy'); + const replacedTransportDestroy = jest.spyOn(replacedEventTransport, 'destroy'); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: replacedJobStore, + eventTransport: replacedEventTransport, + isRedis: false, + }); + manager.initialize(); + const job = await manager.createJob('stream-reconfigure-old', 'user-1'); + let signalCleanupStarted: (() => void) | undefined; + const cleanupStarted = new Promise((resolve) => { + signalCleanupStarted = resolve; + }); + let releaseCleanup: (() => void) | undefined; + const cleanupGate = new Promise((resolve) => { + releaseCleanup = resolve; + }); + job.emitter.on('allSubscribersLeft', async () => { + signalCleanupStarted?.(); + await cleanupGate; + }); + const oldSubscription = await manager.subscribe('stream-reconfigure-old', () => undefined); + oldSubscription?.unsubscribe(); + await cleanupStarted; + + const currentJobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const currentEventTransport = new InMemoryEventTransport(); + const currentStoreDestroy = jest.spyOn(currentJobStore, 'destroy'); + const currentTransportDestroy = jest.spyOn(currentEventTransport, 'destroy'); + manager.configure({ + jobStore: currentJobStore, + eventTransport: currentEventTransport, + isRedis: false, + }); + manager.initialize(); + await manager.createJob('stream-reconfigure-current', 'user-1'); + + expect(replacedTransportDestroy).toHaveBeenCalledTimes(1); + expect(currentStoreDestroy).not.toHaveBeenCalled(); + expect(currentTransportDestroy).not.toHaveBeenCalled(); + + releaseCleanup?.(); + await new Promise((resolve) => setImmediate(resolve)); + await new Promise((resolve) => setImmediate(resolve)); + + expect(replacedStoreDestroy).toHaveBeenCalledTimes(1); + expect(currentStoreDestroy).not.toHaveBeenCalled(); + expect(currentTransportDestroy).not.toHaveBeenCalled(); + await expect(manager.hasJob('stream-reconfigure-current')).resolves.toBe(true); + + await manager.destroy(); + expect(currentStoreDestroy).toHaveBeenCalledTimes(1); + expect(currentTransportDestroy).toHaveBeenCalledTimes(1); + }); + + it('does not deliver a stored terminal event after the subscriber detaches', async () => { + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport: new InMemoryEventTransport(), + isRedis: false, + cleanupOnComplete: false, + }); + manager.initialize(); + await manager.createJob('stream-10', 'user-1', 'conversation-10'); + const finalEvent: ServerSentEvent = { + final: true, + responseMessage: { + messageId: 'response-10', + content: [{ type: 'text', text: 'Complete response' }], + }, + }; + await manager.emitDone('stream-10', finalEvent); + await manager.completeJob('stream-10'); + + const onDone = jest.fn(); + const subscription = await manager.subscribe('stream-10', () => undefined, onDone); + subscription?.unsubscribe(); + await new Promise((resolve) => setImmediate(resolve)); + + expect(onDone).not.toHaveBeenCalled(); + await manager.destroy(); + }); + + it('does not run partial-disconnect persistence after terminal delivery', async () => { + const eventTransport = new InMemoryEventTransport(); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport, + isRedis: false, + }); + manager.initialize(); + const job = await manager.createJob('stream-terminal-disconnect', 'user-1'); + const onAllSubscribersLeft = jest.fn(); + job.emitter.on('allSubscribersLeft', onAllSubscribersLeft); + const onDone = jest.fn(); + const subscription = await manager.subscribe( + 'stream-terminal-disconnect', + () => undefined, + onDone, + ); + const finalEvent: ServerSentEvent = { + final: true, + responseMessage: { + messageId: 'response-terminal', + content: [{ type: 'text', text: 'Complete response' }], + }, + }; + + await manager.emitDone('stream-terminal-disconnect', finalEvent); + await Promise.resolve(); + + expect(onDone).toHaveBeenCalledWith(finalEvent); + expect(onAllSubscribersLeft).not.toHaveBeenCalled(); + expect(eventTransport.getSubscriberCount('stream-terminal-disconnect')).toBe(0); + + subscription?.unsubscribe(); + await manager.destroy(); + }); + + it('waits for transport readiness before scheduling a stored terminal event', async () => { + const eventTransport = new InMemoryEventTransport(); + const originalSubscribe = eventTransport.subscribe.bind(eventTransport); + let releaseReady: (() => void) | undefined; + const readyGate = new Promise((resolve) => { + releaseReady = resolve; + }); + jest.spyOn(eventTransport, 'subscribe').mockImplementation((streamId, handlers) => ({ + ...originalSubscribe(streamId, handlers), + ready: readyGate, + })); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport, + isRedis: false, + cleanupOnComplete: false, + }); + manager.initialize(); + await manager.createJob('stream-terminal-ready', 'user-1'); + const finalEvent: ServerSentEvent = { + final: true, + responseMessage: { + messageId: 'response-ready', + content: [{ type: 'text', text: 'Complete response' }], + }, + }; + await manager.emitDone('stream-terminal-ready', finalEvent); + await manager.completeJob('stream-terminal-ready'); + + const onDone = jest.fn(); + let subscribeSettled = false; + const subscribing = manager + .subscribe('stream-terminal-ready', () => undefined, onDone) + .then((result) => { + subscribeSettled = true; + return result; + }); + await new Promise((resolve) => setImmediate(resolve)); + + expect(subscribeSettled).toBe(false); + expect(onDone).not.toHaveBeenCalled(); + expect(eventTransport.getSubscriberCount('stream-terminal-ready')).toBe(1); + + releaseReady?.(); + const subscription = await subscribing; + expect(onDone).not.toHaveBeenCalled(); + await new Promise((resolve) => setImmediate(resolve)); + + expect(onDone).toHaveBeenCalledWith(finalEvent); + expect(eventTransport.getSubscriberCount('stream-terminal-ready')).toBe(0); + + subscription?.unsubscribe(); + await manager.destroy(); + }); + + it('refreshes terminal state that changes while the transport attaches', async () => { + const eventTransport = new InMemoryEventTransport(); + const originalSubscribe = eventTransport.subscribe.bind(eventTransport); + let releaseReady: (() => void) | undefined; + const readyGate = new Promise((resolve) => { + releaseReady = resolve; + }); + jest.spyOn(eventTransport, 'subscribe').mockImplementation((streamId, handlers) => ({ + ...originalSubscribe(streamId, handlers), + ready: readyGate, + })); + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport, + isRedis: false, + cleanupOnComplete: false, + }); + manager.initialize(); + await manager.createJob('stream-terminal-race', 'user-1'); + const finalEvent: ServerSentEvent = { + final: true, + responseMessage: { + messageId: 'response-race', + content: [{ type: 'text', text: 'Complete response' }], + }, + }; + const onDone = jest.fn(); + const subscribing = manager.subscribe('stream-terminal-race', () => undefined, onDone); + await new Promise((resolve) => setImmediate(resolve)); + + await jobStore.updateJob('stream-terminal-race', { + status: 'complete', + completedAt: Date.now(), + finalEvent: JSON.stringify(finalEvent), + }); + releaseReady?.(); + const subscription = await subscribing; + await new Promise((resolve) => setImmediate(resolve)); + await Promise.resolve(); + + expect(onDone).toHaveBeenCalledWith(finalEvent); + expect(eventTransport.getSubscriberCount('stream-terminal-race')).toBe(0); + + subscription?.unsubscribe(); + await manager.destroy(); + }); + + it('detaches the transport when subscription readiness rejects', async () => { + const eventTransport = new InMemoryEventTransport(); + const originalSubscribe = eventTransport.subscribe.bind(eventTransport); + jest.spyOn(eventTransport, 'subscribe').mockImplementation((streamId, handlers) => ({ + ...originalSubscribe(streamId, handlers), + ready: Promise.reject(new Error('readiness failed')), + })); + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport, + isRedis: false, + cleanupOnComplete: false, + }); + manager.initialize(); + await manager.createJob('stream-terminal-reject', 'user-1'); + const finalEvent: ServerSentEvent = { + final: true, + responseMessage: { + messageId: 'response-reject', + content: [{ type: 'text', text: 'Complete response' }], + }, + }; + await manager.emitDone('stream-terminal-reject', finalEvent); + await manager.completeJob('stream-terminal-reject'); + const onDone = jest.fn(); + + await expect( + manager.subscribe('stream-terminal-reject', () => undefined, onDone), + ).rejects.toThrow('readiness failed'); + await new Promise((resolve) => setImmediate(resolve)); + + expect(onDone).not.toHaveBeenCalled(); + expect(eventTransport.getSubscriberCount('stream-terminal-reject')).toBe(0); + await manager.destroy(); + }); + + it('does not duplicate an event already included by the in-memory resume snapshot', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const originalGetRunSteps = jobStore.getRunSteps.bind(jobStore); + let signalSnapshotInProgress: (() => void) | undefined; + const snapshotInProgress = new Promise((resolve) => { + signalSnapshotInProgress = resolve; + }); + let releaseSnapshot: (() => void) | undefined; + const snapshotGate = new Promise((resolve) => { + releaseSnapshot = resolve; + }); + jest.spyOn(jobStore, 'getRunSteps').mockImplementationOnce(async (...args) => { + const runSteps = await originalGetRunSteps(...args); + signalSnapshotInProgress?.(); + await snapshotGate; + return runSteps; + }); + + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + }); + manager.initialize(); + await manager.createJob('stream-resume-snapshot', 'user-1'); + const contentParts: Parameters[1] = [ + { type: 'text', text: 'before' }, + ]; + manager.setContentParts('stream-resume-snapshot', contentParts); + + const liveEvents: ServerSentEvent[] = []; + const resuming = manager.subscribeWithResume('stream-resume-snapshot', (event) => + liveEvents.push(event), + ); + await snapshotInProgress; + + const snapshotEvent: ServerSentEvent = { + event: 'on_message_delta', + data: { delta: { content: [{ type: 'text', text: 'during snapshot' }] } }, + }; + contentParts.push({ type: 'text', text: 'during snapshot' }); + await manager.emitChunk('stream-resume-snapshot', snapshotEvent); + releaseSnapshot?.(); + + const result = await resuming; + expect(result.resumeState?.aggregatedContent).toEqual([ + { type: 'text', text: 'before' }, + { type: 'text', text: 'during snapshot' }, + ]); + expect(result.pendingEvents).toEqual([]); + result.subscription?.activate(); + expect(liveEvents).toEqual([]); + + result.subscription?.unsubscribe(); + await manager.destroy(); + }); + + it('does not recapture an emission that started before the snapshot frontier', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const originalUpdateJob = jobStore.updateJob.bind(jobStore); + let signalTitlePersisted: (() => void) | undefined; + const titlePersisted = new Promise((resolve) => { + signalTitlePersisted = resolve; + }); + let releaseTitleWrite: (() => void) | undefined; + const titleWriteGate = new Promise((resolve) => { + releaseTitleWrite = resolve; + }); + jest + .spyOn(jobStore, 'updateJob') + .mockImplementation(async (streamId, updates, expectedCreatedAt) => { + await originalUpdateJob(streamId, updates, expectedCreatedAt); + if (updates.titleEvent) { + signalTitlePersisted?.(); + await titleWriteGate; + } + }); + + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + }); + manager.initialize(); + await manager.createJob('stream-resume-inflight', 'user-1'); + const titleEvent: ServerSentEvent = { + event: 'title', + data: { title: 'Snapshot title' }, + }; + const emitting = manager.emitChunk('stream-resume-inflight', titleEvent); + await titlePersisted; + + const liveEvents: ServerSentEvent[] = []; + const resuming = manager.subscribeWithResume('stream-resume-inflight', (event) => + liveEvents.push(event), + ); + releaseTitleWrite?.(); + const result = await resuming; + await emitting; + + expect(result.resumeState?.titleEvent).toEqual(titleEvent); + expect(result.pendingEvents).toEqual([]); + result.subscription?.activate(); + expect(liveEvents).toEqual([]); + + result.subscription?.unsubscribe(); + await manager.destroy(); + }); + + it('preserves an in-memory resume event emitted after snapshot before transport readiness', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const originalGetContentParts = jobStore.getContentParts.bind(jobStore); + let signalSnapshotTaken: (() => void) | undefined; + const snapshotTaken = new Promise((resolve) => { + signalSnapshotTaken = resolve; + }); + let releaseSnapshot: (() => void) | undefined; + const snapshotGate = new Promise((resolve) => { + releaseSnapshot = resolve; + }); + jest.spyOn(jobStore, 'getContentParts').mockImplementationOnce(async (...args) => { + const snapshot = await originalGetContentParts(...args); + signalSnapshotTaken?.(); + await snapshotGate; + return snapshot; + }); + + const eventTransport = new InMemoryEventTransport(); + const originalSubscribe = eventTransport.subscribe.bind(eventTransport); + let signalTransportSubscribed: (() => void) | undefined; + const transportSubscribed = new Promise((resolve) => { + signalTransportSubscribed = resolve; + }); + let releaseReady: (() => void) | undefined; + const readyGate = new Promise((resolve) => { + releaseReady = resolve; + }); + jest.spyOn(eventTransport, 'subscribe').mockImplementationOnce((streamId, handlers) => { + const subscription = originalSubscribe(streamId, handlers); + signalTransportSubscribed?.(); + return { ...subscription, ready: readyGate }; + }); + + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false }); + manager.initialize(); + await manager.createJob('stream-resume-gap', 'user-1'); + + const liveEvents: ServerSentEvent[] = []; + const resuming = manager.subscribeWithResume('stream-resume-gap', (event) => + liveEvents.push(event), + ); + await snapshotTaken; + releaseSnapshot?.(); + await transportSubscribed; + + const gapEvent: ServerSentEvent = { + event: 'on_message_delta', + data: { delta: { content: [{ type: 'text', text: 'gap' }] } }, + }; + await manager.emitChunk('stream-resume-gap', gapEvent); + expect(liveEvents).toEqual([]); + + releaseReady?.(); + const result = await resuming; + expect(result.pendingEvents).toEqual([gapEvent]); + expect(liveEvents).toEqual([]); + + const deliveredEvents = [...result.pendingEvents]; + result.subscription?.activate(); + deliveredEvents.push(...liveEvents); + expect(deliveredEvents).toEqual([gapEvent]); + + result.subscription?.unsubscribe(); + await manager.destroy(); + }); + + it('captures the same in-memory gap independently for overlapping resume subscribers', async () => { + const eventTransport = new InMemoryEventTransport(); + const originalSubscribe = eventTransport.subscribe.bind(eventTransport); + let transportSubscriptions = 0; + let signalBothTransportSubscribed: (() => void) | undefined; + const bothTransportSubscribed = new Promise((resolve) => { + signalBothTransportSubscribed = resolve; + }); + let releaseReady: (() => void) | undefined; + const readyGate = new Promise((resolve) => { + releaseReady = resolve; + }); + jest.spyOn(eventTransport, 'subscribe').mockImplementation((streamId, handlers) => { + const subscription = originalSubscribe(streamId, handlers); + transportSubscriptions++; + if (transportSubscriptions === 2) { + signalBothTransportSubscribed?.(); + } + return { ...subscription, ready: readyGate }; + }); + + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60_000 }), + eventTransport, + isRedis: false, + }); + manager.initialize(); + await manager.createJob('stream-overlapping-resumes', 'user-1'); + + let snapshotsStarted = 0; + let signalSnapshotsStarted: (() => void) | undefined; + const bothSnapshotsStarted = new Promise((resolve) => { + signalSnapshotsStarted = resolve; + }); + let releaseSnapshots: (() => void) | undefined; + const snapshotGate = new Promise((resolve) => { + releaseSnapshots = resolve; + }); + jest.spyOn(manager, 'getResumeState').mockImplementation(async () => { + snapshotsStarted++; + if (snapshotsStarted === 2) { + signalSnapshotsStarted?.(); + } + await snapshotGate; + return { runSteps: [], aggregatedContent: [] }; + }); + + const firstLiveEvents: ServerSentEvent[] = []; + const secondLiveEvents: ServerSentEvent[] = []; + const firstResume = manager.subscribeWithResume('stream-overlapping-resumes', (event) => + firstLiveEvents.push(event), + ); + const secondResume = manager.subscribeWithResume('stream-overlapping-resumes', (event) => + secondLiveEvents.push(event), + ); + await bothSnapshotsStarted; + releaseSnapshots?.(); + await bothTransportSubscribed; + + const gapEvent: ServerSentEvent = { + event: 'on_message_delta', + data: { delta: { content: [{ type: 'text', text: 'gap' }] } }, + }; + await manager.emitChunk('stream-overlapping-resumes', gapEvent); + releaseReady?.(); + + const [firstResult, secondResult] = await Promise.all([firstResume, secondResume]); + expect(firstResult.pendingEvents).toEqual([gapEvent]); + expect(secondResult.pendingEvents).toEqual([gapEvent]); + + const liveEvent: ServerSentEvent = { + event: 'on_message_delta', + data: { delta: { content: [{ type: 'text', text: 'live' }] } }, + }; + await manager.emitChunk('stream-overlapping-resumes', liveEvent); + expect(firstLiveEvents).toEqual([]); + expect(secondLiveEvents).toEqual([]); + + firstResult.subscription?.activate(); + expect(firstLiveEvents).toEqual([liveEvent]); + expect(secondLiveEvents).toEqual([]); + secondResult.subscription?.activate(); + expect(secondLiveEvents).toEqual([liveEvent]); + + firstResult.subscription?.unsubscribe(); + secondResult.subscription?.unsubscribe(); + await manager.destroy(); + }); + + it('restores in-memory gap events when a resume attachment is canceled', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const originalGetContentParts = jobStore.getContentParts.bind(jobStore); + let signalSnapshotTaken: (() => void) | undefined; + const snapshotTaken = new Promise((resolve) => { + signalSnapshotTaken = resolve; + }); + let releaseSnapshot: (() => void) | undefined; + const snapshotGate = new Promise((resolve) => { + releaseSnapshot = resolve; + }); + jest.spyOn(jobStore, 'getContentParts').mockImplementationOnce(async (...args) => { + const snapshot = await originalGetContentParts(...args); + signalSnapshotTaken?.(); + await snapshotGate; + return snapshot; + }); + + const eventTransport = new InMemoryEventTransport(); + const originalSubscribe = eventTransport.subscribe.bind(eventTransport); + let signalTransportSubscribed: (() => void) | undefined; + const transportSubscribed = new Promise((resolve) => { + signalTransportSubscribed = resolve; + }); + let releaseReady: (() => void) | undefined; + const readyGate = new Promise((resolve) => { + releaseReady = resolve; + }); + jest.spyOn(eventTransport, 'subscribe').mockImplementationOnce((streamId, handlers) => { + const subscription = originalSubscribe(streamId, handlers); + signalTransportSubscribed?.(); + return { ...subscription, ready: readyGate }; + }); + + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false }); + manager.initialize(); + await manager.createJob('stream-resume-cancel', 'user-1'); + + const attachmentAbortController = new AbortController(); + const resuming = manager.subscribeWithResume( + 'stream-resume-cancel', + () => undefined, + undefined, + undefined, + { signal: attachmentAbortController.signal }, + ); + await snapshotTaken; + releaseSnapshot?.(); + await transportSubscribed; + + const gapEvent: ServerSentEvent = { + event: 'on_message_delta', + data: { delta: { content: [{ type: 'text', text: 'gap' }] } }, + }; + await manager.emitChunk('stream-resume-cancel', gapEvent); + + attachmentAbortController.abort(); + await expect(resuming).resolves.toEqual( + expect.objectContaining({ subscription: null, pendingEvents: [] }), + ); + + const replayed: ServerSentEvent[] = []; + const subscription = await manager.subscribe('stream-resume-cancel', (event) => + replayed.push(event), + ); + expect(replayed).toEqual([gapEvent]); + + releaseReady?.(); + subscription?.unsubscribe(); + await manager.destroy(); + }); + + it('restores gap events and detaches when resume reconciliation rejects', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const originalGetContentParts = jobStore.getContentParts.bind(jobStore); + let signalSnapshotTaken: (() => void) | undefined; + const snapshotTaken = new Promise((resolve) => { + signalSnapshotTaken = resolve; + }); + let releaseSnapshot: (() => void) | undefined; + const snapshotGate = new Promise((resolve) => { + releaseSnapshot = resolve; + }); + jest.spyOn(jobStore, 'getContentParts').mockImplementationOnce(async (...args) => { + const snapshot = await originalGetContentParts(...args); + signalSnapshotTaken?.(); + await snapshotGate; + return snapshot; + }); + const originalPeekSteers = jobStore.peekSteers.bind(jobStore); + jest + .spyOn(jobStore, 'peekSteers') + .mockImplementationOnce(originalPeekSteers) + .mockRejectedValueOnce(new Error('reconciliation failed')); + + const eventTransport = new InMemoryEventTransport(); + const originalSubscribe = eventTransport.subscribe.bind(eventTransport); + let signalTransportSubscribed: (() => void) | undefined; + const transportSubscribed = new Promise((resolve) => { + signalTransportSubscribed = resolve; + }); + let releaseReady: (() => void) | undefined; + const readyGate = new Promise((resolve) => { + releaseReady = resolve; + }); + jest.spyOn(eventTransport, 'subscribe').mockImplementationOnce((streamId, handlers) => { + const subscription = originalSubscribe(streamId, handlers); + signalTransportSubscribed?.(); + return { ...subscription, ready: readyGate }; + }); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false }); + manager.initialize(); + await manager.createJob('stream-resume-reject', 'user-1'); + + const resuming = manager.subscribeWithResume('stream-resume-reject', () => undefined); + await snapshotTaken; + releaseSnapshot?.(); + await transportSubscribed; + + const gapEvent: ServerSentEvent = { + event: 'on_message_delta', + data: { delta: { content: [{ type: 'text', text: 'gap' }] } }, + }; + await manager.emitChunk('stream-resume-reject', gapEvent); + releaseReady?.(); + + await expect(resuming).rejects.toThrow('reconciliation failed'); + expect(eventTransport.getSubscriberCount('stream-resume-reject')).toBe(0); + + const replayed: ServerSentEvent[] = []; + const subscription = await manager.subscribe('stream-resume-reject', (event) => + replayed.push(event), + ); + expect(replayed).toEqual([gapEvent]); + + subscription?.unsubscribe(); + await manager.destroy(); + }); +}); diff --git a/packages/api/src/stream/__tests__/steering.spec.ts b/packages/api/src/stream/__tests__/steering.spec.ts index 4f23e416d3..3f39de295f 100644 --- a/packages/api/src/stream/__tests__/steering.spec.ts +++ b/packages/api/src/stream/__tests__/steering.spec.ts @@ -138,6 +138,17 @@ describe('SteeringLifecycle via GenerationJobManager.steering (in-memory)', () = 'kept for the live run', ]); }); + + test('peek with a stale expectedCreatedAt hides and preserves the live queue', async () => { + const streamId = 'steer-peek-stale'; + const job = await manager.createJob(streamId, 'user-1'); + await manager.steering.enqueue(streamId, buildSteer('kept for the live run')); + + expect(await manager.steering.peek(streamId, job.createdAt - 1)).toEqual([]); + expect((await manager.steering.peek(streamId, job.createdAt)).map((s) => s.text)).toEqual([ + 'kept for the live run', + ]); + }); }); describe('cancel', () => { @@ -232,6 +243,22 @@ describe('SteeringLifecycle via GenerationJobManager.steering (in-memory)', () = expect(await manager.steering.claim(streamId, owner)).toEqual([]); }); + test('a stale generation cannot park leftovers onto a replacement', async () => { + const streamId = 'steer-park-stale'; + const oldJob = await manager.createJob(streamId, 'user-1'); + await new Promise((resolve) => setTimeout(resolve, 2)); + const replacement = await manager.createJob(streamId, 'user-1'); + const leftovers: TPendingSteer[] = [ + { steerId: 'old', text: 'belongs to predecessor', createdAt: Date.now() }, + ]; + + await manager.steering.park(streamId, leftovers, owner, oldJob.createdAt); + expect(await manager.steering.claim(streamId, owner)).toEqual([]); + + await manager.steering.park(streamId, leftovers, owner, replacement.createdAt); + expect(await manager.steering.claim(streamId, owner)).toEqual(leftovers); + }); + test('parked leftovers survive completeJob within the terminal TTL', async () => { const streamId = 'steer-park-terminal'; await manager.createJob(streamId, 'user-1'); @@ -411,6 +438,55 @@ describe('SteeringLifecycle via GenerationJobManager.steering (in-memory)', () = ).toEqual(['unsent one', 'unsent two']); expect(await manager.steering.peek(streamId)).toEqual([]); }); + + test('abortJob publishes nothing when natural completion wins its terminal CAS', async () => { + const streamId = 'steer-abort-loses-terminal-race'; + const eventTransport = new InMemoryEventTransport(); + const emitDone = jest.spyOn(eventTransport, 'emitDone'); + const racingManager = new GenerationJobManagerClass(); + racingManager.configure({ + jobStore, + eventTransport, + isRedis: false, + cleanupOnComplete: false, + }); + racingManager.initialize(); + const job = await racingManager.createJob(streamId, 'user-1'); + const originalGetContentParts = jobStore.getContentParts.bind(jobStore); + let signalSnapshotStarted: (() => void) | undefined; + const snapshotStarted = new Promise((resolve) => { + signalSnapshotStarted = resolve; + }); + let releaseSnapshot: (() => void) | undefined; + const snapshotGate = new Promise((resolve) => { + releaseSnapshot = resolve; + }); + jest.spyOn(jobStore, 'getContentParts').mockImplementationOnce(async (...args) => { + signalSnapshotStarted?.(); + await snapshotGate; + return originalGetContentParts(...args); + }); + + try { + const aborting = racingManager.abortJob(streamId); + await snapshotStarted; + await racingManager.completeJob(streamId, undefined, job.createdAt); + releaseSnapshot?.(); + + await expect(aborting).resolves.toMatchObject({ + success: false, + finalEvent: null, + }); + await expect(jobStore.getJob(streamId)).resolves.toMatchObject({ + createdAt: job.createdAt, + status: 'complete', + }); + expect(emitDone).not.toHaveBeenCalled(); + } finally { + releaseSnapshot?.(); + await racingManager.destroy(); + } + }); }); describe('synthesizeAppliedSteerEvents (snapshot→subscribe gap)', () => { @@ -591,6 +667,41 @@ describe('SteeringLifecycle via GenerationJobManager.steering (in-memory)', () = expect(result.resumeState?.pendingSteers).toBeUndefined(); expect(result.pendingEvents).toEqual([]); }); + + test('cancels when a replacement becomes durable after attachment', async () => { + const streamId = 'steer-gap-replaced'; + const predecessor = await manager.createJob(streamId, 'user-1'); + const predecessorSteer = buildSteer('predecessor queue'); + await manager.steering.enqueue(streamId, predecessorSteer); + jest + .spyOn(manager, 'getResumeState') + .mockResolvedValue(staleSnapshot(streamId, [toPendingSteer(predecessorSteer)])); + + const getJob = jobStore.getJob.bind(jobStore); + const peekSpy = jest.spyOn(jobStore, 'peekSteers'); + const contentSpy = jest.spyOn(jobStore, 'getContentParts'); + let jobReadCount = 0; + let replacementCreatedAt: number | undefined; + jest.spyOn(jobStore, 'getJob').mockImplementation(async (requestedStreamId) => { + jobReadCount++; + if (jobReadCount === 3) { + const replacement = await jobStore.createJob(requestedStreamId, 'user-1'); + replacementCreatedAt = replacement.createdAt; + await jobStore.enqueueSteer(requestedStreamId, buildSteer('replacement queue')); + return replacement; + } + return getJob(requestedStreamId); + }); + + const result = await manager.subscribeWithResume(streamId, jest.fn()); + + expect(result.subscription).toBeNull(); + expect(result.pendingEvents).toEqual([]); + expect(result.resumeState?.pendingSteers).toEqual([toPendingSteer(predecessorSteer)]); + expect(peekSpy).not.toHaveBeenCalled(); + expect(contentSpy).not.toHaveBeenCalled(); + expect(replacementCreatedAt).toBeGreaterThan(predecessor.createdAt); + }); }); describe('resume state', () => { @@ -644,7 +755,7 @@ describe('emitChunk durability (Redis-mode chunk log)', () => { const redisModeManager = buildRedisModeManager(store, transport); try { const streamId = 'steer-durable'; - await redisModeManager.createJob(streamId, 'user-1'); + const job = await redisModeManager.createJob(streamId, 'user-1'); let resolveAppend!: () => void; const appendGate = new Promise((resolve) => { @@ -665,7 +776,7 @@ describe('emitChunk durability (Redis-mode chunk log)', () => { resolveAppend(); await emit; expect(settled).toBe(true); - expect(publishSpy).toHaveBeenCalledWith(streamId, steerEvent); + expect(publishSpy).toHaveBeenCalledWith(streamId, steerEvent, job.createdAt); } finally { await redisModeManager.destroy(); } @@ -677,16 +788,54 @@ describe('emitChunk durability (Redis-mode chunk log)', () => { const redisModeManager = buildRedisModeManager(store, transport); try { const streamId = 'steer-fire-and-forget'; - await redisModeManager.createJob(streamId, 'user-1'); + const job = await redisModeManager.createJob(streamId, 'user-1'); // Never resolves: the per-delta hot path must not gate on durability. jest.spyOn(store, 'appendChunk').mockReturnValue(new Promise(() => undefined)); const publishSpy = jest.spyOn(transport, 'emitChunk'); await redisModeManager.emitChunk(streamId, steerEvent); - expect(publishSpy).toHaveBeenCalledWith(streamId, steerEvent); + expect(publishSpy).toHaveBeenCalledWith(streamId, steerEvent, job.createdAt); } finally { await redisModeManager.destroy(); } }); + + test('a durable predecessor emission stops when the runtime is replaced during append', async () => { + const now = jest.spyOn(Date, 'now').mockReturnValue(100); + const store = new InMemoryJobStore({ ttlAfterComplete: 60000 }); + const transport = new InMemoryEventTransport(); + const redisModeManager = buildRedisModeManager(store, transport); + let resolveAppend: (() => void) | undefined; + + try { + const streamId = 'steer-durable-replaced'; + const predecessor = await redisModeManager.createJob(streamId, 'user-1'); + const appendStarted = new Promise((resolve) => { + jest.spyOn(store, 'appendChunk').mockImplementationOnce( + () => + new Promise((resolveAppendPromise) => { + resolveAppend = resolveAppendPromise; + resolve(); + }), + ); + }); + const publishSpy = jest.spyOn(transport, 'emitChunk'); + const staleEmission = redisModeManager.emitChunk(streamId, steerEvent, { durable: true }); + await appendStarted; + + now.mockReturnValue(200); + const replacement = await redisModeManager.createJob(streamId, 'user-1'); + resolveAppend?.(); + await staleEmission; + + expect(predecessor.abortController.signal.aborted).toBe(true); + expect(replacement.abortController.signal.aborted).toBe(false); + expect(publishSpy).not.toHaveBeenCalled(); + } finally { + resolveAppend?.(); + now.mockRestore(); + await redisModeManager.destroy(); + } + }); }); diff --git a/packages/api/src/stream/implementations/InMemoryEventTransport.ts b/packages/api/src/stream/implementations/InMemoryEventTransport.ts index c2e7ba01fd..f9e953b9da 100644 --- a/packages/api/src/stream/implementations/InMemoryEventTransport.ts +++ b/packages/api/src/stream/implementations/InMemoryEventTransport.ts @@ -28,16 +28,34 @@ export class InMemoryEventTransport implements IEventTransport { subscribe( streamId: string, handlers: { - onChunk: (event: unknown) => void; - onDone?: (event: unknown) => void; - onError?: (error: string) => void; + onChunk: (event: unknown, generationId?: number) => void; + onDone?: (event: unknown, generationId?: number) => void; + onError?: (error: string, generationId?: number) => void; }, ): { unsubscribe: () => void; ready?: Promise } { const state = this.getOrCreateStream(streamId); - const chunkHandler = (event: unknown) => handlers.onChunk(event); - const doneHandler = (event: unknown) => handlers.onDone?.(event); - const errorHandler = (error: string) => handlers.onError?.(error); + const chunkHandler = (event: unknown, generationId?: number) => { + if (generationId == null) { + handlers.onChunk(event); + return; + } + handlers.onChunk(event, generationId); + }; + const doneHandler = (event: unknown, generationId?: number) => { + if (generationId == null) { + handlers.onDone?.(event); + return; + } + handlers.onDone?.(event, generationId); + }; + const errorHandler = (error: string, generationId?: number) => { + if (generationId == null) { + handlers.onError?.(error); + return; + } + handlers.onError?.(error, generationId); + }; state.emitter.on('chunk', chunkHandler); state.emitter.on('done', doneHandler); @@ -51,6 +69,10 @@ export class InMemoryEventTransport implements IEventTransport { unsubscribe: () => { const currentState = this.streams.get(streamId); if (currentState) { + if (!currentState.emitter.listeners('chunk').includes(chunkHandler)) { + return; + } + currentState.emitter.off('chunk', chunkHandler); currentState.emitter.off('done', doneHandler); currentState.emitter.off('error', errorHandler); @@ -69,22 +91,22 @@ export class InMemoryEventTransport implements IEventTransport { }; } - emitChunk(streamId: string, event: unknown): void { + emitChunk(streamId: string, event: unknown, generationId?: number): void { const state = this.streams.get(streamId); - state?.emitter.emit('chunk', event); + state?.emitter.emit('chunk', event, generationId); } - emitDone(streamId: string, event: unknown): void { + emitDone(streamId: string, event: unknown, generationId?: number): void { const state = this.streams.get(streamId); - state?.emitter.emit('done', event); + state?.emitter.emit('done', event, generationId); } - emitError(streamId: string, error: string): void { + emitError(streamId: string, error: string, generationId?: number): void { const state = this.streams.get(streamId); // Only emit if there are listeners - Node.js throws on unhandled 'error' events // This is intentional for the race condition where error occurs before client connects if (state?.emitter.listenerCount('error') ?? 0 > 0) { - state?.emitter.emit('error', error); + state?.emitter.emit('error', error, generationId); } } @@ -108,6 +130,32 @@ export class InMemoryEventTransport implements IEventTransport { return count === 1; } + closeLocalSubscribers(streamId: string, error: string): void { + const state = this.streams.get(streamId); + if (!state) { + return; + } + + const errorListeners = state.emitter.listeners('error'); + for (const listener of errorListeners) { + try { + listener(error); + } catch (err) { + logger.error( + `[InMemoryEventTransport] Failed to close local subscriber for ${streamId}:`, + err, + ); + } + } + + if (state.emitter.listenerCount('chunk') === 0) { + return; + } + + state.emitter.removeAllListeners(); + state.allSubscribersLeftCallback?.(); + } + /** * Cleanup a stream's event emitter */ diff --git a/packages/api/src/stream/implementations/InMemoryJobStore.ts b/packages/api/src/stream/implementations/InMemoryJobStore.ts index 1fae66e510..9d3d55a7e6 100644 --- a/packages/api/src/stream/implementations/InMemoryJobStore.ts +++ b/packages/api/src/stream/implementations/InMemoryJobStore.ts @@ -7,6 +7,7 @@ import type { UsageMetadata, IJobStore, JobStatus, + JobMetadataPatch, JobStatusTransition, IdempotencyClaimValue, IdempotencyClaimResult, @@ -120,22 +121,27 @@ export class InMemoryJobStore implements IJobStore { userId: string, conversationId?: string, tenantId?: string, + initialMetadata: JobMetadataPatch = {}, ): Promise { - if (this.jobs.size >= this.maxJobs) { + const previousCreatedAt = this.jobs.get(streamId)?.createdAt; + if (previousCreatedAt == null && this.jobs.size >= this.maxJobs) { await this.evictOldest(); } + const now = Date.now(); const job: SerializableJobData = { + ...initialMetadata, streamId, userId, ...(tenantId && { tenantId }), status: 'running', - createdAt: Date.now(), - conversationId, + createdAt: previousCreatedAt == null ? now : Math.max(now, previousCreatedAt + 1), + ...(conversationId !== undefined && { conversationId }), syncSent: false, }; this.jobs.set(streamId, job); + this.contentState.delete(streamId); // Clear any prior activity timestamp so a replacement reusing this streamId // (the controller handles job replacement) falls back to the fresh createdAt // and isn't reaped on the previous generation's stale last-activity time. @@ -165,9 +171,13 @@ export class InMemoryJobStore implements IJobStore { return this.jobs.get(streamId) ?? null; } - async updateJob(streamId: string, updates: Partial): Promise { + async updateJob( + streamId: string, + updates: Partial, + expectedCreatedAt?: number, + ): Promise { const job = this.jobs.get(streamId); - if (!job) { + if (!job || (expectedCreatedAt != null && job.createdAt !== expectedCreatedAt)) { return; } // Plain field writer. Membership-aware status transitions @@ -189,6 +199,12 @@ export class InMemoryJobStore implements IJobStore { if (args.expectActionId != null && job.pendingActionId !== args.expectActionId) { return false; } + if (args.expectCreatedAt != null && job.createdAt !== args.expectCreatedAt) { + return false; + } + if (['complete', 'error', 'aborted'].includes(args.to)) { + this.parkQueuedSteers(streamId, job, Date.now()); + } job.status = args.to; if (args.patch) { Object.assign(job, args.patch); @@ -217,13 +233,25 @@ export class InMemoryJobStore implements IJobStore { this.idempotencyClaims.delete(key); } - async deleteJob(streamId: string): Promise { + async deleteJob(streamId: string, expectedCreatedAt?: number): Promise { + const job = this.jobs.get(streamId); + if (!job || (expectedCreatedAt != null && job.createdAt !== expectedCreatedAt)) { + return false; + } + this.jobs.delete(streamId); this.contentState.delete(streamId); this.lastActivity.delete(streamId); this.steerQueues.delete(streamId); this.closedSteerQueues.delete(streamId); + const userKey = job.tenantId ? `${job.tenantId}:${job.userId}` : job.userId; + const userJobs = this.userJobMap.get(userKey); + userJobs?.delete(streamId); + if (userJobs?.size === 0) { + this.userJobMap.delete(userKey); + } logger.debug(`[InMemoryJobStore] Deleted job: ${streamId}`); + return true; } /** @@ -231,8 +259,9 @@ export class InMemoryJobStore implements IJobStore { * stale-job failsafe in cleanup() reaps on inactivity rather than total age, * mirroring RedisJobStore refreshing the running TTL on each appendChunk. */ - recordActivity(streamId: string): void { - if (this.jobs.has(streamId)) { + recordActivity(streamId: string, expectedCreatedAt?: number): void { + const job = this.jobs.get(streamId); + if (job && (expectedCreatedAt == null || job.createdAt === expectedCreatedAt)) { this.lastActivity.set(streamId, Date.now()); } } @@ -253,7 +282,7 @@ export class InMemoryJobStore implements IJobStore { async cleanup(): Promise { const now = Date.now(); - const toDelete: string[] = []; + const toDelete: Array<{ streamId: string; createdAt: number }> = []; let staleRunning = 0; // Expired parked steers are otherwise only purged by a claim. @@ -276,7 +305,7 @@ export class InMemoryJobStore implements IJobStore { if (isFinished && job.completedAt) { // TTL of 0 means immediate cleanup, otherwise wait for TTL to expire if (this.ttlAfterComplete === 0 || now - job.completedAt > this.ttlAfterComplete) { - toDelete.push(streamId); + toDelete.push({ streamId, createdAt: job.createdAt }); } } else if (job.status === 'requires_action' && isPendingActionStale(job)) { // Stale approval (expired, or missing/malformed pendingAction): @@ -292,7 +321,7 @@ export class InMemoryJobStore implements IJobStore { delete job.pendingAction; delete job.pendingActionId; if (this.ttlAfterComplete === 0) { - toDelete.push(streamId); + toDelete.push({ streamId, createdAt: job.createdAt }); } } else if (this.staleJobTimeout > 0 && job.status === 'running') { // Failsafe: reap jobs stuck in "running" with no generation activity for @@ -311,25 +340,14 @@ export class InMemoryJobStore implements IJobStore { // A crashed/hung run never reached a finalization drain — park the // 202-accepted queue before the delete drops it. this.parkQueuedSteers(streamId, job, now); - toDelete.push(streamId); + toDelete.push({ streamId, createdAt: job.createdAt }); staleRunning++; } } } - for (const id of toDelete) { - const job = this.jobs.get(id); - if (job) { - const userKey = job.tenantId ? `${job.tenantId}:${job.userId}` : job.userId; - const userJobs = this.userJobMap.get(userKey); - if (userJobs) { - userJobs.delete(id); - if (userJobs.size === 0) { - this.userJobMap.delete(userKey); - } - } - } - await this.deleteJob(id); + for (const { streamId, createdAt } of toDelete) { + await this.deleteJob(streamId, createdAt); } if (staleRunning > 0) { @@ -359,17 +377,10 @@ export class InMemoryJobStore implements IJobStore { if (oldestId) { logger.warn(`[InMemoryJobStore] Evicting oldest job: ${oldestId}`); const job = this.jobs.get(oldestId); - if (job) { - const userKey = job.tenantId ? `${job.tenantId}:${job.userId}` : job.userId; - const userJobs = this.userJobMap.get(userKey); - if (userJobs) { - userJobs.delete(oldestId); - if (userJobs.size === 0) { - this.userJobMap.delete(userKey); - } - } + if (!job) { + return; } - await this.deleteJob(oldestId); + await this.deleteJob(oldestId, job.createdAt); } } @@ -449,7 +460,10 @@ export class InMemoryJobStore implements IJobStore { * Set the graph reference for a job. * Uses WeakRef to allow garbage collection when graph is no longer needed. */ - setGraph(streamId: string, graph: StandardGraph): void { + setGraph(streamId: string, graph: StandardGraph, expectedCreatedAt?: number): void { + if (expectedCreatedAt != null && this.jobs.get(streamId)?.createdAt !== expectedCreatedAt) { + return; + } const existing = this.contentState.get(streamId); if (existing) { existing.graphRef = new WeakRef(graph); @@ -465,7 +479,14 @@ export class InMemoryJobStore implements IJobStore { /** * Set content parts reference for a job. */ - setContentParts(streamId: string, contentParts: Agents.MessageContentComplex[]): void { + setContentParts( + streamId: string, + contentParts: Agents.MessageContentComplex[], + expectedCreatedAt?: number, + ): void { + if (expectedCreatedAt != null && this.jobs.get(streamId)?.createdAt !== expectedCreatedAt) { + return; + } const existing = this.contentState.get(streamId); if (existing) { existing.contentParts = contentParts; @@ -477,7 +498,14 @@ export class InMemoryJobStore implements IJobStore { /** * Set collected usage reference for a job. */ - setCollectedUsage(streamId: string, collectedUsage: UsageMetadata[]): void { + setCollectedUsage( + streamId: string, + collectedUsage: UsageMetadata[], + expectedCreatedAt?: number, + ): void { + if (expectedCreatedAt != null && this.jobs.get(streamId)?.createdAt !== expectedCreatedAt) { + return; + } const existing = this.contentState.get(streamId); if (existing) { existing.collectedUsage = collectedUsage; @@ -489,7 +517,10 @@ export class InMemoryJobStore implements IJobStore { /** * Get collected usage for a job. */ - getCollectedUsage(streamId: string): UsageMetadata[] { + getCollectedUsage(streamId: string, expectedCreatedAt?: number): UsageMetadata[] { + if (expectedCreatedAt != null && this.jobs.get(streamId)?.createdAt !== expectedCreatedAt) { + return []; + } const state = this.contentState.get(streamId); return state?.collectedUsage ?? []; } @@ -498,9 +529,15 @@ export class InMemoryJobStore implements IJobStore { * Get content parts for a job. * Returns live content from stored reference. */ - async getContentParts(streamId: string): Promise<{ + async getContentParts( + streamId: string, + expectedCreatedAt?: number, + ): Promise<{ content: Agents.MessageContentComplex[]; } | null> { + if (expectedCreatedAt != null && this.jobs.get(streamId)?.createdAt !== expectedCreatedAt) { + return null; + } const state = this.contentState.get(streamId); if (!state?.contentParts) { return null; @@ -514,7 +551,10 @@ export class InMemoryJobStore implements IJobStore { * Get run steps for a job from graph.contentData. * Uses WeakRef - may return empty if graph has been GC'd. */ - async getRunSteps(streamId: string): Promise { + async getRunSteps(streamId: string, expectedCreatedAt?: number): Promise { + if (expectedCreatedAt != null && this.jobs.get(streamId)?.createdAt !== expectedCreatedAt) { + return []; + } const state = this.contentState.get(streamId); if (!state?.graphRef) { return []; @@ -528,14 +568,22 @@ export class InMemoryJobStore implements IJobStore { /** * No-op for in-memory - content available via graph reference. */ - async appendChunk(): Promise { + async appendChunk( + _streamId: string, + _event: unknown, + _expectedCreatedAt?: number, + ): Promise { // No-op: content available via graph reference } /** * Clear content state for a job. */ - clearContentState(streamId: string): void { + clearContentState(streamId: string, expectedCreatedAt?: number): void { + const job = this.jobs.get(streamId); + if (job && expectedCreatedAt != null && job.createdAt !== expectedCreatedAt) { + return; + } this.contentState.delete(streamId); } @@ -584,7 +632,10 @@ export class InMemoryJobStore implements IJobStore { return this.drainSteers(streamId); } - async peekSteers(streamId: string): Promise { + async peekSteers(streamId: string, expectedCreatedAt?: number): Promise { + if (expectedCreatedAt != null && this.jobs.get(streamId)?.createdAt !== expectedCreatedAt) { + return []; + } const queue = this.steerQueues.get(streamId); return queue ? [...queue] : []; } @@ -603,7 +654,10 @@ export class InMemoryJobStore implements IJobStore { return true; } - async parkSteers(streamId: string, payload: string): Promise { + async parkSteers(streamId: string, payload: string, expectedCreatedAt?: number): Promise { + if (expectedCreatedAt != null && this.jobs.get(streamId)?.createdAt !== expectedCreatedAt) { + return; + } this.parkedSteers.set(streamId, { payload, expiresAt: Date.now() + PARKED_STEERS_TTL_MS, diff --git a/packages/api/src/stream/implementations/RedisEventTransport.ts b/packages/api/src/stream/implementations/RedisEventTransport.ts index df710d25cd..1cf32f28d1 100644 --- a/packages/api/src/stream/implementations/RedisEventTransport.ts +++ b/packages/api/src/stream/implementations/RedisEventTransport.ts @@ -1,6 +1,7 @@ import { logger } from '@librechat/data-schemas'; import type { Redis, Cluster } from 'ioredis'; import type { IEventTransport } from '~/stream/interfaces/IJobStore'; +import { registerChunkPublicationCapability } from '~/stream/internal/chunkPublication'; import { instrumentIORedisClient, RedisUseCases } from '~/cache/redisTelemetry'; /** @@ -19,6 +20,8 @@ const KEYS = { sequence: (streamId: string) => `stream:{${streamId}}:seq`, /** Job metadata, used to keep the sequence counter alive for the full job lifetime */ job: (streamId: string) => `stream:{${streamId}}:job`, + /** Latest generation epoch, retained briefly beyond the live job hash. */ + generationEpoch: (streamId: string) => `stream:{${streamId}}:generation-epoch`, }; /** @@ -37,6 +40,8 @@ interface PubSubMessage { seq?: number; data?: unknown; error?: string; + /** Immutable identity of the generation that emitted the event. */ + generationId?: number; } /** @@ -54,6 +59,10 @@ interface ReorderBuffer { deliveryDeferred: boolean; } +interface AbortRegistration { + callback: (generationId?: number) => void; +} + /** * Allocate a sequence number and publish the event in a single round trip. * @@ -71,11 +80,37 @@ interface ReorderBuffer { * subscriber listens on. PUBLISH is broadcast cluster-wide rather than slot-routed, so it does * not need to be a key for Cluster correctness. * - * KEYS: [sequence, job] - * ARGV: [channel, payloadPrefix, payloadSuffix, sequenceTtlSeconds] - * RETURNS: the 0-indexed seq assigned to this event + * KEYS: [sequence, job, generationEpoch] + * ARGV: [ + * channel, + * payloadPrefix, + * payloadSuffix, + * sequenceTtlSeconds, + * expectCreatedAt | "", + * allowRetainedEpoch ("0" | "1"), + * generationEpochGraceTtl + * ] + * RETURNS: the 0-indexed seq assigned to this event, or -1 when the generation guard fails + * + * During a rolling deployment, a job created by the previous version can expire without + * leaving a generation marker. A tagged terminal event may claim that absent marker only + * while the job hash is also absent. The same bounded ambiguity exists for an extremely + * late event after a recovery marker expires; the generationId in the payload contains it, + * because active runtimes discard terminal events for another epoch. */ const PUBLISH_SEQ_LUA = + 'if ARGV[5] ~= "" then ' + + 'local currentCreatedAt = redis.call("HGET", KEYS[2], "createdAt") ' + + 'if currentCreatedAt ~= ARGV[5] then ' + + 'if redis.call("EXISTS", KEYS[2]) == 1 or ARGV[6] ~= "1" then return -1 end ' + + 'local retainedEpoch = redis.call("GET", KEYS[3]) ' + + 'if not retainedEpoch then ' + + 'redis.call("SET", KEYS[3], ARGV[5], "EX", tonumber(ARGV[7]), "NX") ' + + 'retainedEpoch = redis.call("GET", KEYS[3]) ' + + 'end ' + + 'if retainedEpoch ~= ARGV[5] then return -1 end ' + + 'end ' + + 'end ' + 'local val = redis.call("INCR", KEYS[1]) ' + 'local ttl = tonumber(ARGV[4]) ' + 'local seqTtl = redis.call("TTL", KEYS[1]) ' + @@ -92,6 +127,8 @@ const PUBLISH_SEQ_LUA = const REORDER_TIMEOUT_MS = 500; /** Max messages to buffer before force-flushing (prevents memory issues) */ const MAX_BUFFER_SIZE = 100; +/** Rolling-upgrade recovery window after a legacy job hash expires without an epoch marker. */ +const GENERATION_EPOCH_GRACE_TTL_SECONDS = 300; /** * Subscriber state for a stream @@ -101,14 +138,15 @@ interface StreamSubscribers { handlers: Map< string, { - onChunk: (event: unknown) => void; - onDone?: (event: unknown) => void; - onError?: (error: string) => void; + onChunk: (event: unknown, generationId?: number) => void; + onDone?: (event: unknown, generationId?: number) => void; + onError?: (error: string, generationId?: number) => void; } >; - allSubscribersLeftCallbacks: Array<() => void>; + /** Replaced when a stream runtime is replaced; only the current lifecycle owns cleanup. */ + allSubscribersLeftCallback?: () => void; /** Abort callbacks - called when abort signal is received from any replica */ - abortCallbacks: Array<() => void>; + abortCallbacks: Set; /** Reorder buffer for handling out-of-order delivery in Redis Cluster */ reorderBuffer: ReorderBuffer; } @@ -121,6 +159,9 @@ interface StreamSubscribers { * - Publisher: Emits events to Redis channel when chunks arrive * - Subscriber: Listens to Redis channel and forwards to SSE clients * - Decoupled: Generator and consumer don't need direct connection + * - Ordering: Every sequenced event for one stream must use PUBLISH_SEQ_LUA. Its hash-tagged + * counter routes all publishers through one Redis slot owner, so sequence and publish order + * share one authoritative FIFO origin. * * Note: Requires TWO Redis connections - one for publishing, one for subscribing. * This is a Redis limitation: a client in subscribe mode can't publish. @@ -153,6 +194,9 @@ export class RedisEventTransport implements IEventTransport { constructor(publisher: Redis | Cluster, subscriber: Redis | Cluster) { this.publisher = instrumentIORedisClient(publisher, RedisUseCases.GENERATION_STREAM); this.subscriber = instrumentIORedisClient(subscriber, RedisUseCases.GENERATION_STREAM); + registerChunkPublicationCapability(this, (streamId, event, generationId) => + this.publishChunkWithReceipt(streamId, event, generationId), + ); // Set up message handler for all subscriptions this.subscriber.on('message', (channel: string, message: string) => { @@ -187,21 +231,67 @@ export class RedisEventTransport implements IEventTransport { private async publishWithSequence( streamId: string, message: Omit, + expectedGenerationId?: number, + allowRetainedEpoch = false, ): Promise { const [prefix, suffix] = RedisEventTransport.buildPayloadParts(message); const seq = await this.publisher.eval( PUBLISH_SEQ_LUA, - 2, + 3, KEYS.sequence(streamId), KEYS.job(streamId), + KEYS.generationEpoch(streamId), CHANNELS.events(streamId), prefix, suffix, String(RedisEventTransport.SEQUENCE_TTL_SECONDS), + expectedGenerationId != null ? String(expectedGenerationId) : '', + allowRetainedEpoch ? '1' : '0', + String(GENERATION_EPOCH_GRACE_TTL_SECONDS), ); return seq as number; } + private publishChunkWithReceipt( + streamId: string, + event: unknown, + generationId?: number, + ): Promise { + return this.publishWithSequence( + streamId, + { + type: EventTypes.CHUNK, + data: event, + ...(generationId != null && { generationId }), + }, + generationId, + ) + .then((sequence) => (sequence === -1 ? false : sequence)) + .catch((err) => { + logger.error(`[RedisEventTransport] Failed to publish chunk:`, err); + return false; + }); + } + + private ensureChannelSubscription(channel: string): Promise { + const existing = this.channelSubscriptions.get(channel); + if (existing) { + return existing; + } + + const ready = this.subscriber.subscribe(channel).then(() => { + logger.debug(`[RedisEventTransport] Subscription active for channel ${channel}`); + }); + this.channelSubscriptions.set(channel, ready); + void ready.catch((err) => { + if (this.channelSubscriptions.get(channel) === ready) { + this.channelSubscriptions.delete(channel); + } + logger.error(`[RedisEventTransport] Failed to subscribe to ${channel}:`, err); + }); + return ready; + } + /** Reset subscriber reorder buffer state to initial values */ private resetReorderBuffer(streamId: string): void { const state = this.streams.get(streamId); @@ -217,7 +307,8 @@ export class RedisEventTransport implements IEventTransport { } /** - * Advance subscriber reorder buffer to the authoritative Redis sequence counter (cross-replica safe). + * Advance subscriber reorder buffer to the authoritative Redis sequence counter + * (cross-replica safe). * * @param replayedNextSeq - Absolute Redis sequence immediately after the last event replayed * from earlyEventBuffer. Pending entries below it were already delivered; entries at or @@ -273,7 +364,8 @@ export class RedisEventTransport implements IEventTransport { minPending = seq; } } - buffer.nextSeq = Math.max(buffer.nextSeq, Math.min(currentSeq, minPending)); + const replayOrRedisFrontier = replayedNextSeq ?? currentSeq; + buffer.nextSeq = Math.max(buffer.nextSeq, Math.min(replayOrRedisFrontier, minPending)); } buffer.deliveryDeferred = false; @@ -486,13 +578,25 @@ export class RedisEventTransport implements IEventTransport { for (const [, handlers] of streamState.handlers) { switch (message.type) { case EventTypes.CHUNK: - handlers.onChunk(message.data); + if (message.generationId == null) { + handlers.onChunk(message.data); + } else { + handlers.onChunk(message.data, message.generationId); + } break; case EventTypes.DONE: - handlers.onDone?.(message.data); + if (message.generationId == null) { + handlers.onDone?.(message.data); + } else { + handlers.onDone?.(message.data, message.generationId); + } break; case EventTypes.ERROR: - handlers.onError?.(message.error ?? 'Unknown error'); + if (message.generationId == null) { + handlers.onError?.(message.error ?? 'Unknown error'); + } else { + handlers.onError?.(message.error ?? 'Unknown error', message.generationId); + } break; case EventTypes.ABORT: break; @@ -500,9 +604,13 @@ export class RedisEventTransport implements IEventTransport { } if (message.type === EventTypes.ABORT) { - for (const callback of streamState.abortCallbacks) { + for (const registration of streamState.abortCallbacks) { try { - callback(); + if (message.generationId == null) { + registration.callback(); + } else { + registration.callback(message.generationId); + } } catch (err) { logger.error(`[RedisEventTransport] Error in abort callback:`, err); } @@ -510,20 +618,51 @@ export class RedisEventTransport implements IEventTransport { } } + private detachStreamSubscribers(streamId: string, state: StreamSubscribers): void { + this.resetReorderBuffer(streamId); + + this.unsubscribeUnusedChannel(streamId, state); + + try { + state.allSubscribersLeftCallback?.(); + } catch (err) { + logger.error(`[RedisEventTransport] Error in allSubscribersLeft callback:`, err); + } + } + + private unsubscribeUnusedChannel(streamId: string, state: StreamSubscribers): void { + if (this.streams.get(streamId) !== state || state.count > 0 || state.abortCallbacks.size > 0) { + return; + } + + const channel = CHANNELS.events(streamId); + if (!this.channelSubscriptions.has(channel)) { + return; + } + + this.subscriber.unsubscribe(channel).catch((err) => { + logger.error(`[RedisEventTransport] Failed to unsubscribe from ${channel}:`, err); + }); + this.channelSubscriptions.delete(channel); + } + /** * Subscribe to events for a stream. * - * On first subscriber for a stream, subscribes to the Redis channel. - * Returns unsubscribe function that cleans up when last subscriber leaves. + * Ensures the Redis channel is active and returns an SSE-specific unsubscribe function. */ subscribe( streamId: string, handlers: { - onChunk: (event: unknown) => void; - onDone?: (event: unknown) => void; - onError?: (error: string) => void; + onChunk: (event: unknown, generationId?: number) => void; + onDone?: (event: unknown, generationId?: number) => void; + onError?: (error: string, generationId?: number) => void; + }, + options?: { + deferSequenceDelivery?: boolean; + /** @deprecated Use deferSequenceDelivery. */ + deferDeliveryUntilSynchronized?: boolean; }, - options?: { deferSequenceDelivery?: boolean }, ): { unsubscribe: () => void; ready?: Promise } { const channel = CHANNELS.events(streamId); const subscriberId = `sub_${++this.subscriberIdCounter}`; @@ -533,8 +672,7 @@ export class RedisEventTransport implements IEventTransport { this.streams.set(streamId, { count: 0, handlers: new Map(), - allSubscribersLeftCallbacks: [], - abortCallbacks: [], + abortCallbacks: new Set(), reorderBuffer: { nextSeq: 0, pending: new Map(), @@ -550,25 +688,13 @@ export class RedisEventTransport implements IEventTransport { // attachment and must not inherit that prior generation's expected seq. if (streamState.count === 0) { this.resetReorderBuffer(streamId); - streamState.reorderBuffer.deliveryDeferred = options?.deferSequenceDelivery === true; + streamState.reorderBuffer.deliveryDeferred = + options?.deferSequenceDelivery === true || options?.deferDeliveryUntilSynchronized === true; } streamState.count++; streamState.handlers.set(subscriberId, handlers); - let readyPromise = this.channelSubscriptions.get(channel); - - if (!readyPromise) { - readyPromise = this.subscriber - .subscribe(channel) - .then(() => { - logger.debug(`[RedisEventTransport] Subscription active for channel ${channel}`); - }) - .catch((err) => { - this.channelSubscriptions.delete(channel); - logger.error(`[RedisEventTransport] Failed to subscribe to ${channel}:`, err); - }); - this.channelSubscriptions.set(channel, readyPromise); - } + const readyPromise = this.ensureChannelSubscription(channel); return { ready: readyPromise, @@ -585,7 +711,8 @@ export class RedisEventTransport implements IEventTransport { streamState.count--; - // If last subscriber left, unsubscribe from Redis and notify + // If the last SSE subscriber left, reset attachment state and notify. + // Keep the Redis channel active while the generation's abort listener owns it. if (streamState.count === 0) { /** * Preserve callbacks for reconnect, but drop ordering state from the @@ -593,25 +720,11 @@ export class RedisEventTransport implements IEventTransport { * keeping a detached subscriber's pending gaps or frontier here can * only delay the next attachment before that authoritative sync. */ - this.resetReorderBuffer(streamId); - - this.subscriber.unsubscribe(channel).catch((err) => { - logger.error(`[RedisEventTransport] Failed to unsubscribe from ${channel}:`, err); - }); - this.channelSubscriptions.delete(channel); - - // Call all-subscribers-left callbacks - for (const callback of streamState.allSubscribersLeftCallbacks) { - try { - callback(); - } catch (err) { - logger.error(`[RedisEventTransport] Error in allSubscribersLeft callback:`, err); - } - } + this.detachStreamSubscribers(streamId, streamState); /** * Preserve stream state (callbacks, abort handlers) for reconnection. * Previously this deleted the entire state, which lost the - * allSubscribersLeftCallbacks and abortCallbacks registered by + * allSubscribersLeft callback and abortCallbacks registered by * GenerationJobManager.createJob(). On the next subscribe() call, * fresh state was created without those callbacks, causing * hasSubscriber to never reset and syncReorderBuffer to be skipped. @@ -629,22 +742,26 @@ export class RedisEventTransport implements IEventTransport { * Performance: sequence allocation and publish share one round trip. This runs per streamed * delta, so the saved round trip is multiplied by the token count of every response. */ - async emitChunk(streamId: string, event: unknown): Promise { - try { - return await this.publishWithSequence(streamId, { type: EventTypes.CHUNK, data: event }); - } catch (err) { - logger.error(`[RedisEventTransport] Failed to publish chunk:`, err); - return undefined; - } + emitChunk(streamId: string, event: unknown, generationId?: number): Promise { + return this.publishChunkWithReceipt(streamId, event, generationId).then(() => undefined); } /** * Publish a done event to all subscribers. * Includes sequence number to ensure delivery after all chunks. */ - async emitDone(streamId: string, event: unknown): Promise { + async emitDone(streamId: string, event: unknown, generationId?: number): Promise { try { - await this.publishWithSequence(streamId, { type: EventTypes.DONE, data: event }); + await this.publishWithSequence( + streamId, + { + type: EventTypes.DONE, + data: event, + ...(generationId != null && { generationId }), + }, + generationId, + true, + ); } catch (err) { logger.error(`[RedisEventTransport] Failed to publish done:`, err); throw err; @@ -655,15 +772,51 @@ export class RedisEventTransport implements IEventTransport { * Publish an error event to all subscribers. * Includes sequence number to ensure delivery after all chunks. */ - async emitError(streamId: string, error: string): Promise { + async emitError(streamId: string, error: string, generationId?: number): Promise { try { - await this.publishWithSequence(streamId, { type: EventTypes.ERROR, error }); + await this.publishWithSequence( + streamId, + { + type: EventTypes.ERROR, + error, + ...(generationId != null && { generationId }), + }, + generationId, + true, + ); } catch (err) { logger.error(`[RedisEventTransport] Failed to publish error:`, err); throw err; } } + closeLocalSubscribers(streamId: string, error: string): void { + const state = this.streams.get(streamId); + if (!state) { + return; + } + + const localHandlers = [...state.handlers.values()]; + for (const handlers of localHandlers) { + try { + handlers.onError?.(error); + } catch (err) { + logger.error( + `[RedisEventTransport] Failed to close local subscriber for ${streamId}:`, + err, + ); + } + } + + if (state.handlers.size === 0) { + return; + } + + state.handlers.clear(); + state.count = 0; + this.detachStreamSubscribers(streamId, state); + } + /** * Get subscriber count for a stream (local instance only). * @@ -687,14 +840,14 @@ export class RedisEventTransport implements IEventTransport { onAllSubscribersLeft(streamId: string, callback: () => void): void { const state = this.streams.get(streamId); if (state) { - state.allSubscribersLeftCallbacks.push(callback); + state.allSubscribersLeftCallback = callback; } else { // Create state just for the callback this.streams.set(streamId, { count: 0, handlers: new Map(), - allSubscribersLeftCallbacks: [callback], - abortCallbacks: [], + allSubscribersLeftCallback: callback, + abortCallbacks: new Set(), reorderBuffer: { nextSeq: 0, pending: new Map(), @@ -710,9 +863,12 @@ export class RedisEventTransport implements IEventTransport { * This enables cross-replica abort: when a user aborts on Replica B, * the generating Replica A receives the signal and stops. */ - emitAbort(streamId: string): void { + emitAbort(streamId: string, generationId?: number): void { const channel = CHANNELS.events(streamId); - const message: PubSubMessage = { type: EventTypes.ABORT }; + const message: PubSubMessage = { + type: EventTypes.ABORT, + ...(generationId != null && { generationId }), + }; this.publisher.publish(channel, JSON.stringify(message)).catch((err) => { logger.error(`[RedisEventTransport] Failed to publish abort:`, err); @@ -722,11 +878,12 @@ export class RedisEventTransport implements IEventTransport { /** * Register callback for abort signals from any replica. * Called when abort is triggered on any replica (including this one). + * Resolves once the Redis channel is active so callers can safely expose the stream. * * @param streamId - The stream identifier * @param callback - Called when abort signal is received */ - onAbort(streamId: string, callback: () => void): void { + async onAbort(streamId: string, callback: (generationId?: number) => void): Promise<() => void> { const channel = CHANNELS.events(streamId); let state = this.streams.get(streamId); @@ -734,8 +891,7 @@ export class RedisEventTransport implements IEventTransport { state = { count: 0, handlers: new Map(), - allSubscribersLeftCallbacks: [], - abortCallbacks: [], + abortCallbacks: new Set(), reorderBuffer: { nextSeq: 0, pending: new Map(), @@ -746,18 +902,23 @@ export class RedisEventTransport implements IEventTransport { this.streams.set(streamId, state); } - state.abortCallbacks.push(callback); + const registration = { callback }; + state.abortCallbacks.add(registration); - if (!this.channelSubscriptions.has(channel)) { - const ready = this.subscriber - .subscribe(channel) - .then(() => {}) - .catch((err) => { - this.channelSubscriptions.delete(channel); - logger.error(`[RedisEventTransport] Failed to subscribe to ${channel}:`, err); - }); - this.channelSubscriptions.set(channel, ready); + try { + await this.ensureChannelSubscription(channel); + } catch (error) { + state.abortCallbacks.delete(registration); + this.unsubscribeUnusedChannel(streamId, state); + throw error; } + + return () => { + if (this.streams.get(streamId) !== state || !state.abortCallbacks.delete(registration)) { + return; + } + this.unsubscribeUnusedChannel(streamId, state); + }; } /** @@ -783,8 +944,8 @@ export class RedisEventTransport implements IEventTransport { if (state) { state.handlers.clear(); - state.allSubscribersLeftCallbacks = []; - state.abortCallbacks = []; + state.allSubscribersLeftCallback = undefined; + state.abortCallbacks.clear(); } this.resetReorderBuffer(streamId); diff --git a/packages/api/src/stream/implementations/RedisJobStore.ts b/packages/api/src/stream/implementations/RedisJobStore.ts index 2bce9fc704..32ff26f030 100644 --- a/packages/api/src/stream/implementations/RedisJobStore.ts +++ b/packages/api/src/stream/implementations/RedisJobStore.ts @@ -1,7 +1,7 @@ import { logger } from '@librechat/data-schemas'; import { createContentAggregator } from '@librechat/agents'; -import type { Agents, TPendingSteer } from 'librechat-data-provider'; import type { StandardGraph } from '@librechat/agents'; +import type { Agents } from 'librechat-data-provider'; import type { Redis, Cluster } from 'ioredis'; import type { SerializableJobData, @@ -9,6 +9,7 @@ import type { UsageMetadata, IJobStore, JobStatus, + JobMetadataPatch, JobStatusTransition, IdempotencyClaimValue, IdempotencyClaimResult, @@ -19,35 +20,84 @@ import { isPendingActionStale, } from '~/stream/interfaces/IJobStore'; import { instrumentIORedisClient, RedisUseCases } from '~/cache/redisTelemetry'; -import { toPendingSteer } from '~/stream/SteeringLifecycle'; /** * Atomic compare-and-set on the job hash — the single-winner decision for a - * status transition. The job and event-sequence keys share the stream hash tag, - * so updating the job and extending the counter TTL is atomic on both single-node - * Redis and Redis Cluster (cross-slot membership sets are reconciled afterward). + * status transition. All supplied keys share the stream hash tag, so updating + * the job and terminal stream cleanup are atomic on both single-node Redis and + * Redis Cluster (cross-slot membership sets self-heal during cleanup). * - * Guards on the current `status` and, when ARGV[2] is non-empty, on the flat - * `pendingActionId` field — so a stale decision targeting a different action - * loses. On success: removes `clear` fields, writes `status`+patch pairs, - * refreshes the job-hash TTL. Returns 1 if it fired, 0 otherwise. + * Guards on the current `status` and, when supplied, on the flat + * `pendingActionId` and `createdAt` fields — so a stale decision targeting a + * different action or replacement epoch loses. On success: removes `clear` + * fields, writes `status`+patch pairs, refreshes the job-hash TTL, and performs + * terminal cleanup of same-slot stream state. Returns 1 if it fired, 0 otherwise. * - * KEYS: [job, eventSequence] - * ARGV: [from, expectActionId | "", ttl, hdelCount, ...hdelFields, ...hsetPairs] + * KEYS: [job, eventSequence, chunks, runSteps, steers, parkedSteers, generationEpoch] + * ARGV: [ + * from, + * expectActionId | "", + * expectCreatedAt | "", + * ttl, + * terminal ("0" | "1"), + * chunksAfterComplete, + * runStepsAfterComplete, + * parkedSteersTtl, + * generationEpochGraceTtl, + * hdelCount, + * ...hdelFields, + * ...hsetPairs + * ] */ const JOB_CAS_LUA = 'if redis.call("HGET", KEYS[1], "status") ~= ARGV[1] then return 0 end ' + 'if ARGV[2] ~= "" and redis.call("HGET", KEYS[1], "pendingActionId") ~= ARGV[2] then return 0 end ' + - 'local ttl = tonumber(ARGV[3]) ' + - 'local hdelCount = tonumber(ARGV[4]) ' + - 'local idx = 5 ' + + 'if ARGV[3] ~= "" and redis.call("HGET", KEYS[1], "createdAt") ~= ARGV[3] then return 0 end ' + + 'local currentCreatedAt = redis.call("HGET", KEYS[1], "createdAt") ' + + 'local ttl = tonumber(ARGV[4]) ' + + 'local terminal = ARGV[5] == "1" ' + + 'local chunksTtl = tonumber(ARGV[6]) ' + + 'local runStepsTtl = tonumber(ARGV[7]) ' + + 'local parkedTtl = tonumber(ARGV[8]) ' + + 'local generationEpochGraceTtl = tonumber(ARGV[9]) ' + + 'local hdelCount = tonumber(ARGV[10]) ' + + 'local idx = 11 ' + 'for i = 1, hdelCount do redis.call("HDEL", KEYS[1], ARGV[idx]) idx = idx + 1 end ' + 'local hset = {} ' + 'for i = idx, #ARGV do hset[#hset + 1] = ARGV[i] end ' + 'if #hset > 0 then redis.call("HSET", KEYS[1], unpack(hset)) end ' + + 'local ownerUserId = redis.call("HGET", KEYS[1], "userId") ' + + 'local ownerTenantId = redis.call("HGET", KEYS[1], "tenantId") ' + 'redis.call("EXPIRE", KEYS[1], ttl) ' + 'local seqTtl = redis.call("TTL", KEYS[2]) ' + 'if seqTtl >= 0 and seqTtl < ttl then redis.call("EXPIRE", KEYS[2], ttl) end ' + + 'if currentCreatedAt then redis.call("SET", KEYS[7], currentCreatedAt, "EX", ttl + generationEpochGraceTtl) end ' + + 'if terminal then ' + + 'local queued = redis.call("LRANGE", KEYS[5], 0, -1) ' + + 'if #queued > 0 then ' + + 'local steers = {} ' + + 'for i = 1, #queued do ' + + 'local decoded, item = pcall(cjson.decode, queued[i]) ' + + 'if decoded and type(item) == "table" then ' + + 'local projected = { steerId = item.steerId, text = item.text, createdAt = item.createdAt } ' + + 'if item.files then projected.files = item.files end ' + + 'steers[#steers + 1] = projected ' + + 'end ' + + 'end ' + + 'if #steers > 0 then ' + + 'local parked = { userId = ownerUserId, steers = steers } ' + + 'if ownerTenantId then parked.tenantId = ownerTenantId end ' + + 'redis.call("SET", KEYS[6], cjson.encode(parked), "EX", parkedTtl) ' + + 'end ' + + 'end ' + + 'redis.call("DEL", KEYS[5]) ' + + 'if chunksTtl == 0 then redis.call("DEL", KEYS[3]) else redis.call("EXPIRE", KEYS[3], chunksTtl) end ' + + 'if runStepsTtl == 0 then redis.call("DEL", KEYS[4]) else redis.call("EXPIRE", KEYS[4], runStepsTtl) end ' + + 'else ' + + 'redis.call("EXPIRE", KEYS[3], ttl) ' + + 'redis.call("EXPIRE", KEYS[4], ttl) ' + + 'redis.call("EXPIRE", KEYS[5], ttl) ' + + 'end ' + 'return 1'; /** @@ -63,27 +113,118 @@ const IDEMPOTENCY_CLAIM_LUA = 'return redis.call("GET", KEYS[1])'; /** - * Atomic job (re)creation for the two same-slot keys: reset the steer queue - * and write the job hash in ONE script. A `/chat/steer` request can then - * never interleave between the queue reset and the hash write — the enqueue - * script observes either the old hash (its status guards apply and its list - * dies with it) or the fully initialized replacement with an empty queue, so - * a steer accepted against one run can never be drained into another. + * Atomic job (re)creation for all generation-scoped same-slot keys. The + * predecessor hash, content, run steps, steer queue, and parked steers are + * removed in the same script that installs the replacement hash. A reconnect + * that observes the replacement can therefore never reconstruct predecessor + * state, even when the predecessor's delayed completion loses its epoch guard. * - * KEYS: [job, steers, parkedSteers] - * ARGV: [ttl, hdelCount, ...hdelFields, ...hsetPairs] + * KEYS: [job, chunks, runSteps, steers, parkedSteers, generationEpoch] + * ARGV: [ttl, requestedCreatedAt, generationEpochGraceTtl, ...hsetPairs] + * Returns: [previousUserId | "", previousTenantId | "", createdAt] */ const JOB_CREATE_LUA = - 'redis.call("DEL", KEYS[2]) ' + - 'redis.call("DEL", KEYS[3]) ' + + 'local previousUserId = redis.call("HGET", KEYS[1], "userId") ' + + 'local previousTenantId = redis.call("HGET", KEYS[1], "tenantId") ' + + 'local previousCreatedAt = tonumber(redis.call("HGET", KEYS[1], "createdAt")) ' + + 'local retainedEpoch = tonumber(redis.call("GET", KEYS[6])) ' + + 'if retainedEpoch and (not previousCreatedAt or retainedEpoch > previousCreatedAt) then previousCreatedAt = retainedEpoch end ' + + 'local createdAt = tonumber(ARGV[2]) ' + + 'if previousCreatedAt and previousCreatedAt >= createdAt then createdAt = previousCreatedAt + 1 end ' + + 'redis.call("DEL", KEYS[1], KEYS[2], KEYS[3], KEYS[4], KEYS[5]) ' + 'local ttl = tonumber(ARGV[1]) ' + - 'local hdelCount = tonumber(ARGV[2]) ' + - 'local idx = 3 ' + - 'for i = 1, hdelCount do redis.call("HDEL", KEYS[1], ARGV[idx]) idx = idx + 1 end ' + + 'local generationEpochGraceTtl = tonumber(ARGV[3]) ' + 'local hset = {} ' + - 'for i = idx, #ARGV do hset[#hset + 1] = ARGV[i] end ' + + 'for i = 4, #ARGV do hset[#hset + 1] = ARGV[i] end ' + 'redis.call("HSET", KEYS[1], unpack(hset)) ' + + 'redis.call("HSET", KEYS[1], "createdAt", tostring(createdAt)) ' + 'redis.call("EXPIRE", KEYS[1], ttl) ' + + 'redis.call("SET", KEYS[6], tostring(createdAt), "EX", ttl + generationEpochGraceTtl) ' + + 'return { previousUserId or "", previousTenantId or "", tostring(createdAt) }'; + +/** + * Epoch-guarded field update. Terminal writes reclaim same-slot content in the + * same atomic step, so a replacement cannot appear between the guarded write + * and content cleanup. + * + * KEYS: [job, chunks, runSteps, steers] + * ARGV: [ + * expectCreatedAt | "", + * terminal ("0" | "1"), + * completedTtl, + * chunksAfterComplete, + * runStepsAfterComplete, + * ...hsetPairs + * ] + */ +const JOB_UPDATE_LUA = + 'if redis.call("EXISTS", KEYS[1]) == 0 then return 0 end ' + + 'if ARGV[1] ~= "" and redis.call("HGET", KEYS[1], "createdAt") ~= ARGV[1] then return 0 end ' + + 'local hset = {} ' + + 'for i = 6, #ARGV do hset[#hset + 1] = ARGV[i] end ' + + 'if #hset > 0 then redis.call("HSET", KEYS[1], unpack(hset)) end ' + + 'if ARGV[2] == "1" then ' + + 'local completedTtl = tonumber(ARGV[3]) ' + + 'local chunksTtl = tonumber(ARGV[4]) ' + + 'local runStepsTtl = tonumber(ARGV[5]) ' + + 'redis.call("EXPIRE", KEYS[1], completedTtl) ' + + 'redis.call("DEL", KEYS[4]) ' + + 'if chunksTtl == 0 then redis.call("DEL", KEYS[2]) else redis.call("EXPIRE", KEYS[2], chunksTtl) end ' + + 'if runStepsTtl == 0 then redis.call("DEL", KEYS[3]) else redis.call("EXPIRE", KEYS[3], runStepsTtl) end ' + + 'end ' + + 'return 1'; + +/** + * Epoch-guarded hard deletion. `expectMissing` makes an unguarded cleanup of + * already-absent state safe against a replacement appearing after the read. + * Parked steers intentionally survive: completion parks before deleting. + * + * KEYS: [job, chunks, runSteps, steers] + * ARGV: [expectCreatedAt | "", expectMissing ("0" | "1")] + */ +const JOB_DELETE_LUA = + 'if ARGV[1] ~= "" and redis.call("HGET", KEYS[1], "createdAt") ~= ARGV[1] then return 0 end ' + + 'local existed = redis.call("EXISTS", KEYS[1]) ' + + 'if ARGV[2] == "1" and existed == 1 then return 0 end ' + + 'redis.call("DEL", KEYS[1], KEYS[2], KEYS[3], KEYS[4]) ' + + 'return existed'; + +/** + * Atomic stale-running reap. The liveness and epoch checks, steer projection + * and parking, and same-slot deletion are one operation. A replacement either + * lands before the script and fails the guard, or lands afterward and clears + * the predecessor's parked payload in {@link JOB_CREATE_LUA}. + * + * KEYS: [job, chunks, runSteps, steers, parkedSteers, generationEpoch] + * ARGV: [expectCreatedAt, nowMs, staleAfterMs, parkedSteersTtl, generationEpochGraceTtl] + */ +const STALE_JOB_DELETE_LUA = + 'if redis.call("HGET", KEYS[1], "status") ~= "running" then return 0 end ' + + 'if redis.call("HGET", KEYS[1], "createdAt") ~= ARGV[1] then return 0 end ' + + 'local liveSince = tonumber(redis.call("HGET", KEYS[1], "lastActiveAt")) ' + + 'if not liveSince then liveSince = tonumber(redis.call("HGET", KEYS[1], "createdAt")) end ' + + 'if not liveSince or tonumber(ARGV[2]) - liveSince <= tonumber(ARGV[3]) then return 0 end ' + + 'local ownerUserId = redis.call("HGET", KEYS[1], "userId") ' + + 'local ownerTenantId = redis.call("HGET", KEYS[1], "tenantId") ' + + 'local queued = redis.call("LRANGE", KEYS[4], 0, -1) ' + + 'if #queued > 0 and ownerUserId then ' + + 'local steers = {} ' + + 'for i = 1, #queued do ' + + 'local decoded, item = pcall(cjson.decode, queued[i]) ' + + 'if decoded and type(item) == "table" then ' + + 'local projected = { steerId = item.steerId, text = item.text, createdAt = item.createdAt } ' + + 'if item.files then projected.files = item.files end ' + + 'steers[#steers + 1] = projected ' + + 'end ' + + 'end ' + + 'if #steers > 0 then ' + + 'local parked = { userId = ownerUserId, steers = steers } ' + + 'if ownerTenantId then parked.tenantId = ownerTenantId end ' + + 'redis.call("SET", KEYS[5], cjson.encode(parked), "EX", tonumber(ARGV[4])) ' + + 'end ' + + 'end ' + + 'redis.call("SET", KEYS[6], ARGV[1], "EX", tonumber(ARGV[5])) ' + + 'redis.call("DEL", KEYS[1], KEYS[2], KEYS[3], KEYS[4]) ' + 'return 1'; /** @@ -113,9 +254,10 @@ const JOB_CREATE_LUA = * stays `run`. * * KEYS: [chunks, job] - * ARGV: [eventJson, runningTtl] + * ARGV: [eventJson, runningTtl, expectCreatedAt | ""] */ const CHUNK_APPEND_LUA = + 'if ARGV[3] ~= "" and redis.call("HGET", KEYS[2], "createdAt") ~= ARGV[3] then return 0 end ' + 'redis.call("XADD", KEYS[1], "*", "event", ARGV[1]) ' + 'local run = tonumber(ARGV[2]) ' + 'local target = run ' + @@ -138,9 +280,10 @@ const CHUNK_APPEND_LUA = * `transitionStatus` set); a normally-running job keeps the short running TTL. * * KEYS: [runSteps, job] - * ARGV: [runStepsJson, runningTtl] + * ARGV: [runStepsJson, runningTtl, expectCreatedAt | ""] */ const RUNSTEPS_SAVE_LUA = + 'if ARGV[3] ~= "" and redis.call("HGET", KEYS[2], "createdAt") ~= ARGV[3] then return 0 end ' + 'redis.call("SET", KEYS[1], ARGV[1]) ' + 'local run = tonumber(ARGV[2]) ' + 'local target = run ' + @@ -151,6 +294,27 @@ const RUNSTEPS_SAVE_LUA = 'redis.call("EXPIRE", KEYS[1], target) ' + 'return 1'; +/** + * Clear same-slot content unless the stream already belongs to a replacement. + * A missing job is safe: terminal deletion may remove the hash before this + * best-effort cache cleanup, and Redis executes the check + deletes atomically. + * + * KEYS: [chunks, runSteps, job] + * ARGV: [expectCreatedAt | ""] + */ +const CONTENT_CLEAR_LUA = + 'if ARGV[1] ~= "" and redis.call("EXISTS", KEYS[3]) == 1 and redis.call("HGET", KEYS[3], "createdAt") ~= ARGV[1] then return 0 end ' + + 'redis.call("DEL", KEYS[1], KEYS[2]) ' + + 'return 1'; + +const CHUNKS_READ_LUA = + 'if redis.call("HGET", KEYS[1], "createdAt") ~= ARGV[1] then return {} end ' + + 'return redis.call("XRANGE", KEYS[2], "-", "+")'; + +const RUNSTEPS_READ_LUA = + 'if redis.call("HGET", KEYS[1], "createdAt") ~= ARGV[1] then return false end ' + + 'return redis.call("GET", KEYS[2])'; + /** * Atomically append a steer, guarded on the job hash still being `running` * AND the queue not being closed by a terminal drain (`steersClosed` field, @@ -189,6 +353,10 @@ const STEER_DRAIN_LUA = 'redis.call("DEL", KEYS[2]) ' + 'return items'; +const STEER_PEEK_LUA = + 'if redis.call("HGET", KEYS[1], "createdAt") ~= ARGV[1] then return {} end ' + + 'return redis.call("LRANGE", KEYS[2], 0, -1)'; + /** * Remove ONE queued steer by id without disturbing the rest: the list is * rebuilt atomically, so a concurrent drain either delivers the steer or the @@ -239,6 +407,20 @@ const CLAIM_PARKED_LUA = 'redis.call("DEL", KEYS[1]) ' + 'return v'; +/** + * Park leftovers only while the generation that drained them still owns the + * stream ID. If a replacement already exists, writing its parked key would + * leak predecessor state into the new run. A replacement created afterward + * atomically clears this key in {@link JOB_CREATE_LUA}. + * + * KEYS: [job, parkedSteers] + * ARGV: [expectedCreatedAt | "", payload, ttl] + */ +const PARK_STEERS_LUA = + 'if ARGV[1] ~= "" and redis.call("HGET", KEYS[1], "createdAt") ~= ARGV[1] then return 0 end ' + + 'redis.call("SET", KEYS[2], ARGV[2], "EX", ARGV[3]) ' + + 'return 1'; + /** * Terminal close-then-drain in one atomic step: mark the queue closed on the * job hash (only when the hash still exists — a bare HSET would resurrect a @@ -266,6 +448,11 @@ const KNOWN_INTERRUPT_TYPES = new Set(['tool_approval', 'ask_user_question']); * configured to 0 — Redis rejects `EX 0`, which would silently kill * park-based recovery. */ const PARKED_RECOVERY_TTL_S: number = 300; +/** Grace window for publishing terminal/reaper events after the live job hash expires. */ +const GENERATION_EPOCH_GRACE_TTL_S: number = 300; + +/** Bound pathological replacement churn without leaving the request unbounded. */ +const MEMBERSHIP_RECONCILE_MAX_ATTEMPTS: number = 8; /** * Key prefixes for Redis storage. @@ -291,6 +478,8 @@ const KEYS = { /** Parked terminally-drained steers (own TTL — must outlive the job hash, * which the default completeJob path deletes immediately) */ parkedSteers: (streamId: string) => `stream:{${streamId}}:parked`, + /** Latest generation epoch, retained briefly beyond the live job hash. */ + generationEpoch: (streamId: string) => `stream:{${streamId}}:generation-epoch`, /** Running jobs set for cleanup (global set - single slot) */ runningJobs: 'stream:running', /** Jobs paused for human review (global set - single slot) */ @@ -363,6 +552,11 @@ export interface RedisJobStoreOptions { requiresActionTtl?: number; } +interface LocalCacheEntry { + createdAt?: number; + value: T; +} + export class RedisJobStore implements IJobStore { private redis: Redis | Cluster; private cleanupInterval: NodeJS.Timeout | null = null; @@ -376,18 +570,21 @@ export class RedisJobStore implements IJobStore { * Enables fast reconnects when client returns to the same server. * Uses WeakRef to allow garbage collection when graph is no longer needed. */ - private localGraphCache = new Map>(); + private localGraphCache = new Map>>(); /** * Local cache for collectedUsage arrays. * Generation happens on a single instance, so collectedUsage is only available locally. * For cross-replica abort, the abort handler falls back to text-based token counting. */ - private localCollectedUsageCache = new Map(); + private localCollectedUsageCache = new Map>(); /** Same-instance HOST content view (includes host-authored parts like * steers, which the SDK graph never sees). Preferred over the graph cache * on local reads; cross-instance reads reconstruct from chunks. */ - private localContentParts = new Map>(); + private localContentParts = new Map< + string, + LocalCacheEntry> + >(); /** Cleanup interval in ms (1 minute) */ private cleanupIntervalMs = 60000; @@ -425,91 +622,144 @@ export class RedisJobStore implements IJobStore { logger.info('[RedisJobStore] Initialized with cleanup interval'); } + private getLocalEntry( + cache: Map>, + streamId: string, + expectedCreatedAt?: number, + ): LocalCacheEntry | undefined { + const entry = cache.get(streamId); + if (expectedCreatedAt != null && entry?.createdAt !== expectedCreatedAt) { + return undefined; + } + return entry; + } + + private setLocalEntry( + cache: Map>, + streamId: string, + entry: LocalCacheEntry, + ): void { + const currentEntry = cache.get(streamId); + if ( + currentEntry?.createdAt != null && + (entry.createdAt == null || entry.createdAt < currentEntry.createdAt) + ) { + return; + } + cache.set(streamId, entry); + } + + private deleteLocalEntry( + cache: Map>, + streamId: string, + expectedCreatedAt?: number, + observedEntry?: LocalCacheEntry, + ): void { + const entry = cache.get(streamId); + if ( + !entry || + (observedEntry != null && entry !== observedEntry) || + (expectedCreatedAt != null && entry.createdAt !== expectedCreatedAt) + ) { + return; + } + cache.delete(streamId); + } + + private clearLocalState(streamId: string, expectedCreatedAt?: number): void { + this.deleteLocalEntry(this.localGraphCache, streamId, expectedCreatedAt); + this.deleteLocalEntry(this.localContentParts, streamId, expectedCreatedAt); + this.deleteLocalEntry(this.localCollectedUsageCache, streamId, expectedCreatedAt); + } + + private clearPredecessorLocalState(streamId: string, createdAt: number): void { + const graphEntry = this.localGraphCache.get(streamId); + if (graphEntry && (graphEntry.createdAt == null || graphEntry.createdAt < createdAt)) { + this.deleteLocalEntry(this.localGraphCache, streamId, undefined, graphEntry); + } + const contentEntry = this.localContentParts.get(streamId); + if (contentEntry && (contentEntry.createdAt == null || contentEntry.createdAt < createdAt)) { + this.deleteLocalEntry(this.localContentParts, streamId, undefined, contentEntry); + } + const usageEntry = this.localCollectedUsageCache.get(streamId); + if (usageEntry && (usageEntry.createdAt == null || usageEntry.createdAt < createdAt)) { + this.deleteLocalEntry(this.localCollectedUsageCache, streamId, undefined, usageEntry); + } + } + async createJob( streamId: string, userId: string, conversationId?: string, tenantId?: string, + initialMetadata: JobMetadataPatch = {}, ): Promise { const job: SerializableJobData = { + ...initialMetadata, streamId, userId, ...(tenantId && { tenantId }), status: 'running', createdAt: Date.now(), - conversationId, + ...(conversationId !== undefined && { conversationId }), syncSent: false, }; const key = KEYS.job(streamId); - const userJobsKey = KEYS.userJobs(userId, tenantId); - - // A reused streamId overlays onto any existing hash, so per-turn fields from a - // prior generation could survive. Drop the HITL fields so the fresh running job - // never exposes stale approval metadata and cleanup keys off the new createdAt - // rather than a leftover lastActiveAt. `agent_id` is included because - // updateMetadata only writes it when truthy — without clearing it here, a - // conversation that switches from a saved agent to an ephemeral/no-agent turn - // would keep the old agent_id and the resume guard would reject the valid pause. - const staleHitlFields: Array = [ - 'pendingAction', - 'pendingActionId', - 'lastActiveAt', - 'agent_id', - // Same reasoning as agent_id: updateMetadata only writes isTemporary when the new - // metadata carries it, so a prior temporary turn's isTemporary=1 would otherwise - // survive and a later non-temporary resume would save its response as temporary. - 'isTemporary', - // Same reasoning again: handleRunInterrupt only writes discoveredTools when THIS - // turn discovered ≥1 deferred tool, so a replacement turn that later pauses without - // its own discovery would otherwise inherit the prior run's tool names and force-load - // deferred tools it never discovered on resume. - 'discoveredTools', - // A replacement must start with an open steer channel — the closed flag - // belongs to the finalized run this hash is being reused from. - 'steersClosed', - ]; // For cluster mode, we can't pipeline keys on different slots // The job key uses hash tag {streamId}, runningJobs and userJobs are on different slots - // Steer-queue reset + job-hash write happen ATOMICALLY (same-slot Lua): - // a steer POST can never land between them, so a replacement can neither - // inherit the replaced run's undrained steers nor lose/steal a steer - // 202-accepted against either run (see JOB_CREATE_LUA). + // Generation-state reset + job-hash write happen ATOMICALLY (same-slot Lua). const hsetPairs = Object.entries(this.serializeJob(job)).flat(); - await this.redis.eval( + const previousOwner = await this.redis.eval( JOB_CREATE_LUA, - 3, + 6, key, + KEYS.chunks(streamId), + KEYS.runSteps(streamId), KEYS.steers(streamId), KEYS.parkedSteers(streamId), + KEYS.generationEpoch(streamId), String(this.ttl.running), - String(staleHitlFields.length), - ...staleHitlFields, + String(job.createdAt), + String(GENERATION_EPOCH_GRACE_TTL_S), ...hsetPairs, ); - // Set-membership bookkeeping lives on other slots; ordering after the - // atomic hash write is safe (scanners tolerate momentary lag). - if (this.isCluster) { - await this.redis.sadd(KEYS.runningJobs, streamId); - await this.redis.srem(KEYS.requiresActionJobs, streamId); - await this.redis.sadd(userJobsKey, streamId); - if (this.ttl.userJobsSet > 0) { - await this.redis.expire(userJobsKey, this.ttl.userJobsSet); - } - } else { - const pipeline = this.redis.pipeline(); - pipeline.sadd(KEYS.runningJobs, streamId); - pipeline.srem(KEYS.requiresActionJobs, streamId); - pipeline.sadd(userJobsKey, streamId); - if (this.ttl.userJobsSet > 0) { - pipeline.expire(userJobsKey, this.ttl.userJobsSet); - } - await pipeline.exec(); + const previousUserId = + Array.isArray(previousOwner) && typeof previousOwner[0] === 'string' ? previousOwner[0] : ''; + const previousTenantId = + Array.isArray(previousOwner) && typeof previousOwner[1] === 'string' ? previousOwner[1] : ''; + const createdAt = + Array.isArray(previousOwner) && + (typeof previousOwner[2] === 'string' || typeof previousOwner[2] === 'number') + ? Number(previousOwner[2]) + : job.createdAt; + job.createdAt = Number.isFinite(createdAt) ? createdAt : job.createdAt; + const previousUserKeys = + previousUserId !== '' + ? [KEYS.userJobs(previousUserId, previousTenantId || undefined)] + : undefined; + // Cross-slot membership cannot join the creation Lua transaction. Reconcile + // from the durable hash and verify after writing so an overlapping status + // transition or same-stream replacement always gets the final word. + const currentJob = await this.reconcileJobMembership(streamId, { + initialJob: job, + previousUserKeys, + }); + if ( + currentJob !== undefined && + (!currentJob || currentJob.createdAt !== job.createdAt || currentJob.status !== 'running') + ) { + throw new Error('Generation job was replaced during creation'); } + if (currentJob === undefined) { + logger.warn(`[RedisJobStore] Created job without verified membership: ${streamId}`); + return job; + } + this.clearPredecessorLocalState(streamId, currentJob.createdAt); logger.debug(`[RedisJobStore] Created job: ${streamId}`); - return job; + return currentJob; } async getJob(streamId: string): Promise { @@ -520,89 +770,209 @@ export class RedisJobStore implements IJobStore { return this.deserializeJob(data); } - async updateJob(streamId: string, updates: Partial): Promise { + async updateJob( + streamId: string, + updates: Partial, + expectedCreatedAt?: number, + ): Promise { const key = KEYS.job(streamId); // Plain field writer. The membership-aware status transitions // (running ⇄ requires_action — sets, TTLs, the actionId guard) go solely - // through transitionStatus, the single race-safe path. updateJob still - // handles terminal status writes (complete/error/aborted) + their cleanup. + // through transitionStatus. The optional epoch guard keeps late metadata + // and terminal-event persistence from mutating a same-stream replacement. const serialized = this.serializeJob(updates as SerializableJobData); if (Object.keys(serialized).length === 0) { return; } + const terminal = + updates.status != null && ['complete', 'error', 'aborted'].includes(updates.status); + const observedJob = terminal ? await this.getJob(streamId) : null; const fields = Object.entries(serialized).flat(); - const updated = await this.updateExistingJobHash(key, fields); - if (!updated) { + const updated = await this.redis.eval( + JOB_UPDATE_LUA, + 4, + key, + KEYS.chunks(streamId), + KEYS.runSteps(streamId), + KEYS.steers(streamId), + expectedCreatedAt != null ? String(expectedCreatedAt) : '', + terminal ? '1' : '0', + String(this.ttl.completed), + String(this.ttl.chunksAfterComplete), + String(this.ttl.runStepsAfterComplete), + ...fields, + ); + if (updated !== 1) { return; } - if (updates.status && ['complete', 'error', 'aborted'].includes(updates.status)) { - await this.applyTerminalContentCleanup(streamId); + if (terminal) { + const currentJob = await this.reconcileJobMembership(streamId, { + previousJob: observedJob, + }); + this.clearLocalStateUnlessActive( + streamId, + currentJob, + expectedCreatedAt ?? observedJob?.createdAt, + ); } } - /** - * Terminal cleanup shared by `updateJob` (complete/error/aborted) and the - * terminal path of `transitionStatus` (approval expiry → aborted): drop the - * job from both membership sets and the user-active set, shorten the job-hash - * TTL to the completed window, and del/shorten the chunk + run-step keys per - * the configured after-complete TTLs. Without sharing this, an expired - * approval left Redis stream contents around for the full running TTL. - */ - private async applyTerminalContentCleanup(streamId: string): Promise { - const key = KEYS.job(streamId); - // Proactively remove from user's job set (requires reading userId from the job hash) - const job = await this.getJob(streamId); - const userJobsKey = job?.userId ? KEYS.userJobs(job.userId, job.tenantId) : null; + private sameMembershipSource( + left: SerializableJobData | null, + right: SerializableJobData | null, + ): boolean { + if (left == null || right == null) { + return left === right; + } + return ( + left.createdAt === right.createdAt && + left.status === right.status && + left.userId === right.userId && + left.tenantId === right.tenantId + ); + } + + private addObservedUserKey(keys: Set, job: SerializableJobData | null): void { + if (job?.userId) { + keys.add(KEYS.userJobs(job.userId, job.tenantId)); + } + } + + private async applyMembershipSnapshot( + streamId: string, + job: SerializableJobData | null, + observedUserKeys: Set, + ): Promise { + const statusKey = job ? this.statusSetKey(job.status) : null; + const activeUserKey = job && statusKey != null ? KEYS.userJobs(job.userId, job.tenantId) : null; if (this.isCluster) { - await this.redis.expire(key, this.ttl.completed); - await this.redis.srem(KEYS.runningJobs, streamId); - await this.redis.srem(KEYS.requiresActionJobs, streamId); - await this.redis.del(KEYS.steers(streamId)); - - if (this.ttl.chunksAfterComplete === 0) { - await this.redis.del(KEYS.chunks(streamId)); - } else { - await this.redis.expire(KEYS.chunks(streamId), this.ttl.chunksAfterComplete); + const operations: Promise[] = [ + statusKey === KEYS.runningJobs + ? this.redis.sadd(KEYS.runningJobs, streamId) + : this.redis.srem(KEYS.runningJobs, streamId), + statusKey === KEYS.requiresActionJobs + ? this.redis.sadd(KEYS.requiresActionJobs, streamId) + : this.redis.srem(KEYS.requiresActionJobs, streamId), + ]; + for (const userJobsKey of observedUserKeys) { + if (userJobsKey !== activeUserKey) { + operations.push(this.redis.srem(userJobsKey, streamId)); + } } - - if (this.ttl.runStepsAfterComplete === 0) { - await this.redis.del(KEYS.runSteps(streamId)); - } else { - await this.redis.expire(KEYS.runSteps(streamId), this.ttl.runStepsAfterComplete); + if (activeUserKey) { + operations.push( + (async () => { + await this.redis.sadd(activeUserKey, streamId); + if (this.ttl.userJobsSet > 0) { + await this.redis.expire(activeUserKey, this.ttl.userJobsSet); + } + })(), + ); } + await Promise.all(operations); + return this.getJob(streamId); + } - if (userJobsKey) { - await this.redis.srem(userJobsKey, streamId); - } + const pipeline = this.redis.pipeline(); + if (statusKey === KEYS.runningJobs) { + pipeline.sadd(KEYS.runningJobs, streamId); } else { - const pipeline = this.redis.pipeline(); - pipeline.expire(key, this.ttl.completed); pipeline.srem(KEYS.runningJobs, streamId); + } + if (statusKey === KEYS.requiresActionJobs) { + pipeline.sadd(KEYS.requiresActionJobs, streamId); + } else { pipeline.srem(KEYS.requiresActionJobs, streamId); - pipeline.del(KEYS.steers(streamId)); - - if (this.ttl.chunksAfterComplete === 0) { - pipeline.del(KEYS.chunks(streamId)); - } else { - pipeline.expire(KEYS.chunks(streamId), this.ttl.chunksAfterComplete); - } - - if (this.ttl.runStepsAfterComplete === 0) { - pipeline.del(KEYS.runSteps(streamId)); - } else { - pipeline.expire(KEYS.runSteps(streamId), this.ttl.runStepsAfterComplete); - } - - if (userJobsKey) { + } + for (const userJobsKey of observedUserKeys) { + if (userJobsKey !== activeUserKey) { pipeline.srem(userJobsKey, streamId); } - - await pipeline.exec(); } + if (activeUserKey) { + pipeline.sadd(activeUserKey, streamId); + if (this.ttl.userJobsSet > 0) { + pipeline.expire(activeUserKey, this.ttl.userJobsSet); + } + } + // Keep the verification read in this network flush. Redis executes it + // after the membership commands, preserving the guarded loop without an + // extra round trip on the default single-node deployment. + pipeline.hgetall(KEYS.job(streamId)); + const results = await pipeline.exec(); + const verification = results?.[results.length - 1]; + if (verification?.[0]) { + throw verification[0]; + } + const data = verification?.[1] as Record | null | undefined; + if (!data || Object.keys(data).length === 0) { + return null; + } + return this.deserializeJob(data); + } + + /** + * Cross-slot sets are derived state, so every mutation writes the membership + * implied by the durable job hash and then reads the hash again. If a status + * change or replacement crossed that window, the loop repairs from the newer + * source. Because every writer uses this path, the final writer always + * converges even when an older reconciliation finishes later. + */ + private async reconcileJobMembership( + streamId: string, + options: { + initialJob?: SerializableJobData | null; + previousJob?: SerializableJobData | null; + previousUserKeys?: string[]; + } = {}, + ): Promise { + const observedUserKeys = new Set(options.previousUserKeys ?? []); + this.addObservedUserKey(observedUserKeys, options.previousJob ?? null); + let currentJob = options.initialJob ?? null; + try { + if (options.initialJob === undefined) { + currentJob = await this.getJob(streamId); + } + + for (let attempt = 0; attempt < MEMBERSHIP_RECONCILE_MAX_ATTEMPTS; attempt++) { + this.addObservedUserKey(observedUserKeys, currentJob); + const verifiedJob = await this.applyMembershipSnapshot( + streamId, + currentJob, + observedUserKeys, + ); + this.addObservedUserKey(observedUserKeys, verifiedJob); + if (this.sameMembershipSource(currentJob, verifiedJob)) { + return verifiedJob; + } + currentJob = verifiedJob; + } + + logger.warn( + `[RedisJobStore] Membership reconciliation did not stabilize after ${MEMBERSHIP_RECONCILE_MAX_ATTEMPTS} attempts: ${streamId}`, + ); + } catch (err) { + logger.warn(`[RedisJobStore] Failed to reconcile job membership ${streamId}:`, err); + } + return undefined; + } + + private clearLocalStateUnlessActive( + streamId: string, + currentJob: SerializableJobData | null | undefined, + expectedCreatedAt?: number, + ): void { + if (currentJob === undefined) { + return; + } + if (expectedCreatedAt == null && currentJob && this.statusSetKey(currentJob.status)) { + return; + } + this.clearLocalState(streamId, expectedCreatedAt); } /** @@ -635,7 +1005,7 @@ export class RedisJobStore implements IJobStore { } async transitionStatus(streamId: string, args: JobStatusTransition): Promise { - const { from, to, patch, clear, expectActionId } = args; + const { from, to, patch, clear, expectActionId, expectCreatedAt } = args; const key = KEYS.job(streamId); // status + patch become HSET pairs; serializeJob skips undefined, so @@ -645,9 +1015,7 @@ export class RedisJobStore implements IJobStore { ).flat(); const clearFields = (clear ?? []).map(String); - const remSet = this.statusSetKey(from); - const addSet = this.statusSetKey(to); - const terminal = addSet === null; + const terminal = this.statusSetKey(to) === null; let ttl = terminal ? this.ttl.completed : this.ttl.running; if (to === 'requires_action') { // A paused job must outlive its approval window, even when that window is @@ -655,18 +1023,30 @@ export class RedisJobStore implements IJobStore { // decision can resume it. ttl = this.pauseTtlSeconds(patch?.pendingAction); } + const terminalJob = terminal ? await this.getJob(streamId) : null; // 1) Single-winner decision: an atomic CAS on the single-slot job hash. // Works identically on cluster and single-node, so two concurrent // resolves can never both win (and drive the run twice). const won = await this.redis.eval( JOB_CAS_LUA, - 2, + 7, key, KEYS.sequence(streamId), + KEYS.chunks(streamId), + KEYS.runSteps(streamId), + KEYS.steers(streamId), + KEYS.parkedSteers(streamId), + KEYS.generationEpoch(streamId), from, expectActionId ?? '', + expectCreatedAt != null ? String(expectCreatedAt) : '', String(ttl), + terminal ? '1' : '0', + String(this.ttl.chunksAfterComplete), + String(this.ttl.runStepsAfterComplete), + String(this.ttl.completed > 0 ? this.ttl.completed : PARKED_RECOVERY_TTL_S), + String(GENERATION_EPOCH_GRACE_TTL_S), String(clearFields.length), ...clearFields, ...fields, @@ -675,39 +1055,17 @@ export class RedisJobStore implements IJobStore { return false; } - // 2) Reconcile derived state. Only the winner reaches here; membership is - // self-healed by periodic cleanup, so this non-atomic cross-slot step is - // safe. A terminal target (e.g. approval expiry → aborted) gets the same - // content cleanup as updateJob's terminal path. + // 2) Same-slot TTL/content changes happened atomically in the CAS. Cross-slot + // indexes are reconciled last and verified against the durable hash. + const currentJob = await this.reconcileJobMembership(streamId, { + previousJob: terminalJob, + }); if (terminal) { - await this.applyTerminalContentCleanup(streamId); - return true; - } - if (this.isCluster) { - if (remSet) { - await this.redis.srem(remSet, streamId); - } - if (addSet) { - await this.redis.sadd(addSet, streamId); - } - await this.redis.expire(KEYS.chunks(streamId), ttl); - await this.redis.expire(KEYS.runSteps(streamId), ttl); - // Steers queued before a pause must survive the whole approval window - // (they inject at the first tool boundary after resume). EXPIRE on a - // missing key is a no-op. - await this.redis.expire(KEYS.steers(streamId), ttl); - } else { - const pipeline = this.redis.pipeline(); - if (remSet) { - pipeline.srem(remSet, streamId); - } - if (addSet) { - pipeline.sadd(addSet, streamId); - } - pipeline.expire(KEYS.chunks(streamId), ttl); - pipeline.expire(KEYS.runSteps(streamId), ttl); - pipeline.expire(KEYS.steers(streamId), ttl); - await pipeline.exec(); + this.clearLocalStateUnlessActive( + streamId, + currentJob, + expectCreatedAt ?? terminalJob?.createdAt, + ); } return true; } @@ -739,56 +1097,63 @@ export class RedisJobStore implements IJobStore { await this.redis.del(KEYS.idempotency(key)); } - private async updateExistingJobHash(key: string, fields: string[]): Promise { - const updated = await this.redis.eval( - 'if redis.call("EXISTS", KEYS[1]) == 1 then redis.call("HSET", KEYS[1], unpack(ARGV)) return 1 else return 0 end', - 1, - key, - ...fields, + async deleteJob(streamId: string, expectedCreatedAt?: number): Promise { + const observedJob = await this.getJob(streamId); + const targetCreatedAt = expectedCreatedAt ?? observedJob?.createdAt; + const expectMissing = expectedCreatedAt == null && observedJob == null; + const deleted = await this.redis.eval( + JOB_DELETE_LUA, + 4, + KEYS.job(streamId), + KEYS.chunks(streamId), + KEYS.runSteps(streamId), + KEYS.steers(streamId), + targetCreatedAt != null ? String(targetCreatedAt) : '', + expectMissing ? '1' : '0', ); - return updated === 1; - } - - async deleteJob(streamId: string): Promise { - this.localGraphCache.delete(streamId); - this.localContentParts.delete(streamId); - this.localCollectedUsageCache.delete(streamId); - const job = await this.getJob(streamId); - const userJobsKey = job?.userId ? KEYS.userJobs(job.userId, job.tenantId) : null; - return this.deleteJobInternal(streamId, userJobsKey); - } - - private async deleteJobInternal(streamId: string, userJobsKey: string | null): Promise { - this.localGraphCache.delete(streamId); - this.localContentParts.delete(streamId); - this.localCollectedUsageCache.delete(streamId); - - if (this.isCluster) { - const pipeline = this.redis.pipeline(); - pipeline.del(KEYS.job(streamId)); - pipeline.del(KEYS.chunks(streamId)); - pipeline.del(KEYS.runSteps(streamId)); - pipeline.del(KEYS.steers(streamId)); - await pipeline.exec(); - await this.redis.srem(KEYS.runningJobs, streamId); - await this.redis.srem(KEYS.requiresActionJobs, streamId); - if (userJobsKey) { - await this.redis.srem(userJobsKey, streamId); - } - } else { - const pipeline = this.redis.pipeline(); - pipeline.del(KEYS.job(streamId)); - pipeline.del(KEYS.chunks(streamId)); - pipeline.del(KEYS.runSteps(streamId)); - pipeline.del(KEYS.steers(streamId)); - pipeline.srem(KEYS.runningJobs, streamId); - pipeline.srem(KEYS.requiresActionJobs, streamId); - if (userJobsKey) { - pipeline.srem(userJobsKey, streamId); - } - await pipeline.exec(); + if (deleted !== 1) { + return false; } + + const currentJob = await this.reconcileJobMembership(streamId, { + initialJob: null, + previousJob: observedJob, + }); + this.clearLocalStateUnlessActive(streamId, currentJob, targetCreatedAt); logger.debug(`[RedisJobStore] Deleted job: ${streamId}`); + return true; + } + + private async deleteStaleRunningJob( + streamId: string, + observedJob: SerializableJobData, + now: number, + ): Promise { + const deleted = await this.redis.eval( + STALE_JOB_DELETE_LUA, + 6, + KEYS.job(streamId), + KEYS.chunks(streamId), + KEYS.runSteps(streamId), + KEYS.steers(streamId), + KEYS.parkedSteers(streamId), + KEYS.generationEpoch(streamId), + String(observedJob.createdAt), + String(now), + String(this.ttl.running * 1000), + String(this.ttl.completed > 0 ? this.ttl.completed : PARKED_RECOVERY_TTL_S), + String(GENERATION_EPOCH_GRACE_TTL_S), + ); + if (deleted !== 1) { + return false; + } + + const currentJob = await this.reconcileJobMembership(streamId, { + initialJob: null, + previousJob: observedJob, + }); + this.clearLocalStateUnlessActive(streamId, currentJob, observedJob.createdAt); + return true; } async hasJob(streamId: string): Promise { @@ -818,9 +1183,9 @@ export class RedisJobStore implements IJobStore { let cleaned = 0; // Clean up stale local graph cache entries (WeakRefs that were collected) - for (const [streamId, graphRef] of this.localGraphCache) { - if (!graphRef.deref()) { - this.localGraphCache.delete(streamId); + for (const [streamId, graphEntry] of this.localGraphCache) { + if (!graphEntry.value.deref()) { + this.deleteLocalEntry(this.localGraphCache, streamId, undefined, graphEntry); } } @@ -834,20 +1199,19 @@ export class RedisJobStore implements IJobStore { // Job no longer exists (TTL expired) - remove from set if (!job) { - await this.redis.srem(KEYS.runningJobs, streamId); - await this.redis.srem(KEYS.requiresActionJobs, streamId); - this.localGraphCache.delete(streamId); - this.localContentParts.delete(streamId); - this.localCollectedUsageCache.delete(streamId); + const currentJob = await this.reconcileJobMembership(streamId, { initialJob: null }); + this.clearLocalStateUnlessActive(streamId, currentJob); return 1; } if (job.status === 'requires_action') { - await this.redis.srem(KEYS.runningJobs, streamId); - await this.redis.sadd(KEYS.requiresActionJobs, streamId); - this.localGraphCache.delete(streamId); - this.localContentParts.delete(streamId); - this.localCollectedUsageCache.delete(streamId); + const currentJob = await this.reconcileJobMembership(streamId, { initialJob: job }); + if ( + currentJob !== undefined && + (!currentJob || currentJob.createdAt === job.createdAt) + ) { + this.clearLocalState(streamId, job.createdAt); + } return 1; } @@ -855,14 +1219,16 @@ export class RedisJobStore implements IJobStore { // Only remove from tracking sets — do NOT delete the job hash, which has // its own completedTtl so clients can still poll for final status. if (job.status !== 'running') { - await this.redis.srem(KEYS.runningJobs, streamId); - await this.redis.srem(KEYS.requiresActionJobs, streamId); - if (job.userId) { - await this.redis.srem(KEYS.userJobs(job.userId, job.tenantId), streamId); + const currentJob = await this.reconcileJobMembership(streamId, { + initialJob: job, + previousJob: job, + }); + if ( + currentJob !== undefined && + (!currentJob || currentJob.createdAt === job.createdAt) + ) { + this.clearLocalStateUnlessActive(streamId, currentJob, job.createdAt); } - this.localGraphCache.delete(streamId); - this.localContentParts.delete(streamId); - this.localCollectedUsageCache.delete(streamId); return 1; } @@ -872,17 +1238,10 @@ export class RedisJobStore implements IJobStore { const liveSince = job.lastActiveAt ?? job.createdAt; if (now - liveSince > this.ttl.running * 1000) { logger.warn(`[RedisJobStore] Cleaning up stale job: ${streamId}`); - // A crashed/hung run never reached a finalization drain — park the - // 202-accepted queue before the delete drops it, so the owner can - // still recover the words via /chat/status. - await this.parkSteerSnapshot( - streamId, - job, - await this.snapshotParkableSteers(streamId), - ); - const userJobsKey = job.userId ? KEYS.userJobs(job.userId, job.tenantId) : null; - await this.deleteJobInternal(streamId, userJobsKey); - return 1; + // Re-check liveness + epoch, park queued steers, and delete same-slot + // state in one script. A replacement cannot land in the old + // park-then-unconditional-delete gap. + return (await this.deleteStaleRunningJob(streamId, job, now)) ? 1 : 0; } return 0; @@ -918,32 +1277,26 @@ export class RedisJobStore implements IJobStore { const job = await this.getJob(streamId); if (!job) { - await this.redis.srem(KEYS.requiresActionJobs, streamId); - this.localGraphCache.delete(streamId); - this.localContentParts.delete(streamId); - this.localCollectedUsageCache.delete(streamId); + const currentJob = await this.reconcileJobMembership(streamId, { initialJob: null }); + this.clearLocalStateUnlessActive(streamId, currentJob); return 1; } if (job.status !== 'requires_action') { - await this.redis.srem(KEYS.requiresActionJobs, streamId); - if (job.status === 'running') { - await this.redis.sadd(KEYS.runningJobs, streamId); - } + await this.reconcileJobMembership(streamId, { + initialJob: job, + previousJob: job, + }); return 1; } // Stale approval (expired, or missing/malformed pendingAction): // finalize it (aborted) so it stops occupying the slot and its stream // contents are reclaimed, mirroring ApprovalLifecycle.expire(). - // transitionStatus runs the terminal content cleanup (sets, chunks, - // run-steps, userJobs, completed TTL). + // transitionStatus atomically applies the terminal state and same-slot + // content cleanup. Cross-slot membership indexes self-heal on read or + // during the next cleanup pass. if (isPendingActionStale(job)) { - // Snapshot 202-accepted steers BEFORE the terminal cleanup DELs the - // queue key; enqueue is frozen while paused, so this is exact. - // Parked only if the CAS wins — a lost CAS means the run resumed - // and the live queue must stay untouched. - const parkableSteers = await this.snapshotParkableSteers(streamId); const expired = await this.transitionStatus(streamId, { from: 'requires_action', to: 'aborted', @@ -958,11 +1311,9 @@ export class RedisJobStore implements IJobStore { // valid new pause. (Undefined for a missing/malformed pendingAction — nothing // to protect — so it falls back to the status-only check.) expectActionId: job.pendingAction?.actionId, + expectCreatedAt: job.createdAt, }); - if (expired) { - await this.parkSteerSnapshot(streamId, job, parkableSteers); - } - return 1; + return expired ? 1 : 0; } return 0; @@ -983,7 +1334,7 @@ export class RedisJobStore implements IJobStore { async getJobCount(): Promise { const [runningCount, requiresActionCount] = await Promise.all([ - this.redis.scard(KEYS.runningJobs), + this.countJobsInStatusSet(KEYS.runningJobs, 'running'), this.countJobsInStatusSet(KEYS.requiresActionJobs, 'requires_action'), ]); return runningCount + requiresActionCount; @@ -991,7 +1342,7 @@ export class RedisJobStore implements IJobStore { async getJobCountByStatus(status: JobStatus): Promise { if (status === 'running') { - return this.redis.scard(KEYS.runningJobs); + return this.countJobsInStatusSet(KEYS.runningJobs, status); } if (status === 'requires_action') { @@ -1008,20 +1359,21 @@ export class RedisJobStore implements IJobStore { } let count = 0; - const staleIds: string[] = []; for (const streamId of streamIds) { const job = await this.getJob(streamId); if (job?.status === status) { count++; } else { - staleIds.push(streamId); + const currentJob = await this.reconcileJobMembership(streamId, { + initialJob: job, + previousJob: job, + }); + if (currentJob?.status === status) { + count++; + } } } - if (staleIds.length > 0) { - await this.redis.srem(setKey, ...staleIds); - } - return count; } @@ -1042,7 +1394,7 @@ export class RedisJobStore implements IJobStore { } const activeIds: string[] = []; - const staleIds: string[] = []; + let healed = 0; for (const streamId of trackedIds) { const job = await this.getJob(streamId); @@ -1051,22 +1403,38 @@ export class RedisJobStore implements IJobStore { // only while its prompt is live: a past-`expiresAt` approval no longer // counts as active (cleanup/expiry will finalize it), so the client stops // polling and can complete. - if (job && (job.status === 'running' || job.status === 'requires_action')) { + const belongsToUser = + job?.userId === userId && (job.tenantId ?? undefined) === (tenantId ?? undefined); + if (belongsToUser && job && (job.status === 'running' || job.status === 'requires_action')) { if (job.status === 'requires_action' && isPendingActionStale(job)) { continue; } activeIds.push(streamId); } else { - // Self-healing: job completed/deleted but mapping wasn't cleaned - mark for removal - staleIds.push(streamId); + // Self-heal from durable state instead of a raw SREM, which could remove + // a replacement's membership after the read. + const currentJob = await this.reconcileJobMembership(streamId, { + initialJob: job, + previousJob: job, + previousUserKeys: [userJobsKey], + }); + const currentBelongsToUser = + currentJob?.userId === userId && + (currentJob.tenantId ?? undefined) === (tenantId ?? undefined); + if ( + currentBelongsToUser && + currentJob && + (currentJob.status === 'running' || currentJob.status === 'requires_action') && + !(currentJob.status === 'requires_action' && isPendingActionStale(currentJob)) + ) { + activeIds.push(streamId); + } + healed++; } } - if (staleIds.length > 0) { - await this.redis.srem(userJobsKey, ...staleIds); - logger.debug( - `[RedisJobStore] Self-healed ${staleIds.length} stale job entries for user ${userId}`, - ); + if (healed > 0) { + logger.debug(`[RedisJobStore] Self-healed ${healed} stale job entries for user ${userId}`); } return activeIds; @@ -1079,6 +1447,7 @@ export class RedisJobStore implements IJobStore { } // Clear local caches this.localGraphCache.clear(); + this.localContentParts.clear(); this.localCollectedUsageCache.clear(); // Don't close the Redis connection - it's shared logger.info('[RedisJobStore] Destroyed'); @@ -1096,8 +1465,11 @@ export class RedisJobStore implements IJobStore { * @param streamId - The stream identifier * @param graph - The graph instance (stored as WeakRef) */ - setGraph(streamId: string, graph: StandardGraph): void { - this.localGraphCache.set(streamId, new WeakRef(graph)); + setGraph(streamId: string, graph: StandardGraph, expectedCreatedAt?: number): void { + this.setLocalEntry(this.localGraphCache, streamId, { + createdAt: expectedCreatedAt, + value: new WeakRef(graph), + }); } /** Splice-inserts host-authored steer parts (from `on_steer_applied` @@ -1106,8 +1478,9 @@ export class RedisJobStore implements IJobStore { private async overlayHostSteerParts( streamId: string, parts: Agents.MessageContentComplex[], + expectedCreatedAt?: number, ): Promise { - const chunks = await this.getChunks(streamId); + const chunks = await this.getChunks(streamId, expectedCreatedAt); if (chunks.length === 0) { return parts; } @@ -1139,8 +1512,15 @@ export class RedisJobStore implements IJobStore { * live here but never inside the SDK graph, so preferring it over the graph * cache keeps same-instance reconnect/abort/status reads steer-complete. */ - setContentParts(streamId: string, contentParts: Agents.MessageContentComplex[]): void { - this.localContentParts.set(streamId, new WeakRef(contentParts)); + setContentParts( + streamId: string, + contentParts: Agents.MessageContentComplex[], + expectedCreatedAt?: number, + ): void { + this.setLocalEntry(this.localContentParts, streamId, { + createdAt: expectedCreatedAt, + value: new WeakRef(contentParts), + }); } /** @@ -1148,16 +1528,25 @@ export class RedisJobStore implements IJobStore { * This is used for abort handling to spend tokens for all models. * Note: Only available on the generating instance; cross-replica abort uses fallback. */ - setCollectedUsage(streamId: string, collectedUsage: UsageMetadata[]): void { - this.localCollectedUsageCache.set(streamId, collectedUsage); + setCollectedUsage( + streamId: string, + collectedUsage: UsageMetadata[], + expectedCreatedAt?: number, + ): void { + this.setLocalEntry(this.localCollectedUsageCache, streamId, { + createdAt: expectedCreatedAt, + value: collectedUsage, + }); } /** * Get collected usage for a job. * Only available if this is the generating instance. */ - getCollectedUsage(streamId: string): UsageMetadata[] { - return this.localCollectedUsageCache.get(streamId) ?? []; + getCollectedUsage(streamId: string, expectedCreatedAt?: number): UsageMetadata[] { + return ( + this.getLocalEntry(this.localCollectedUsageCache, streamId, expectedCreatedAt)?.value ?? [] + ); } /** @@ -1184,6 +1573,7 @@ export class RedisJobStore implements IJobStore { */ private readCachedGraph( streamId: string, + entry: LocalCacheEntry>, graph: StandardGraph, read: (graph: StandardGraph) => T, ): T | null { @@ -1194,24 +1584,27 @@ export class RedisJobStore implements IJobStore { `[RedisJobStore] Cached graph for ${streamId} is unusable (likely disposed); falling back to reconstruction:`, err instanceof Error ? err.message : err, ); - this.localGraphCache.delete(streamId); + this.deleteLocalEntry(this.localGraphCache, streamId, undefined, entry); return null; } } - async getContentParts(streamId: string): Promise<{ + async getContentParts( + streamId: string, + expectedCreatedAt?: number, + ): Promise<{ content: Agents.MessageContentComplex[]; } | null> { // 1. Prefer the HOST content array (same-instance fast path): it already // contains host-authored steer parts the SDK graph never sees. - const hostRef = this.localContentParts.get(streamId); - if (hostRef) { - const hostParts = hostRef.deref(); + const hostEntry = this.getLocalEntry(this.localContentParts, streamId, expectedCreatedAt); + if (hostEntry) { + const hostParts = hostEntry.value.deref(); if (hostParts && hostParts.length > 0) { return { content: hostParts }; } if (!hostParts) { - this.localContentParts.delete(streamId); + this.deleteLocalEntry(this.localContentParts, streamId, undefined, hostEntry); } } @@ -1219,24 +1612,26 @@ export class RedisJobStore implements IJobStore { // lacks host-authored steer parts, so overlay them from the chunk log — // insert (not assign): the graph array is UNSHIFTED, while recorded steer // indices are host-view positions that already account for prior steers. - const graphRef = this.localGraphCache.get(streamId); - if (graphRef) { - const graph = graphRef.deref(); + const graphEntry = this.getLocalEntry(this.localGraphCache, streamId, expectedCreatedAt); + if (graphEntry) { + const graph = graphEntry.value.deref(); if (graph) { - const localParts = this.readCachedGraph(streamId, graph, (g) => g.getContentParts()); + const localParts = this.readCachedGraph(streamId, graphEntry, graph, (g) => + g.getContentParts(), + ); if (localParts && localParts.length > 0) { return { - content: await this.overlayHostSteerParts(streamId, localParts), + content: await this.overlayHostSteerParts(streamId, localParts, expectedCreatedAt), }; } } else { // WeakRef was collected, remove from cache - this.localGraphCache.delete(streamId); + this.deleteLocalEntry(this.localGraphCache, streamId, undefined, graphEntry); } } // 2. Fall back to Redis chunk reconstruction (cross-instance reconnect) - const chunks = await this.getChunks(streamId); + const chunks = await this.getChunks(streamId, expectedCreatedAt); if (chunks.length === 0) { return null; } @@ -1303,13 +1698,15 @@ export class RedisJobStore implements IJobStore { * @param streamId - The stream identifier * @returns Run steps array */ - async getRunSteps(streamId: string): Promise { + async getRunSteps(streamId: string, expectedCreatedAt?: number): Promise { // 1. Try local graph cache first (fast path for same-instance reconnect) - const graphRef = this.localGraphCache.get(streamId); - if (graphRef) { - const graph = graphRef.deref(); + const graphEntry = this.getLocalEntry(this.localGraphCache, streamId, expectedCreatedAt); + if (graphEntry) { + const graph = graphEntry.value.deref(); if (graph) { - const localSteps = this.readCachedGraph(streamId, graph, (g) => g.getRunSteps()); + const localSteps = this.readCachedGraph(streamId, graphEntry, graph, (g) => + g.getRunSteps(), + ); if (localSteps && localSteps.length > 0) { return localSteps; } @@ -1319,8 +1716,7 @@ export class RedisJobStore implements IJobStore { } // 2. Fall back to Redis (cross-instance reconnect) - const key = KEYS.runSteps(streamId); - const data = await this.redis.get(key); + const data = await this.getRunStepsData(streamId, expectedCreatedAt); if (!data) { return []; } @@ -1331,18 +1727,33 @@ export class RedisJobStore implements IJobStore { } } + private async getRunStepsData( + streamId: string, + expectedCreatedAt?: number, + ): Promise { + if (expectedCreatedAt == null) { + return this.redis.get(KEYS.runSteps(streamId)); + } + const data = await this.redis.eval( + RUNSTEPS_READ_LUA, + 2, + KEYS.job(streamId), + KEYS.runSteps(streamId), + String(expectedCreatedAt), + ); + return typeof data === 'string' ? data : null; + } + /** * Clear content state for a job. * Removes both local cache and Redis data. */ - clearContentState(streamId: string): void { + clearContentState(streamId: string, expectedCreatedAt?: number): void { // Clear local caches immediately - this.localGraphCache.delete(streamId); - this.localContentParts.delete(streamId); - this.localCollectedUsageCache.delete(streamId); + this.clearLocalState(streamId, expectedCreatedAt); // Fire and forget - async cleanup for Redis - this.clearContentStateAsync(streamId).catch((err) => { + this.clearContentStateAsync(streamId, expectedCreatedAt).catch((err) => { logger.error(`[RedisJobStore] Failed to clear content state for ${streamId}:`, err); }); } @@ -1350,11 +1761,18 @@ export class RedisJobStore implements IJobStore { /** * Clear content state async. */ - private async clearContentStateAsync(streamId: string): Promise { - const pipeline = this.redis.pipeline(); - pipeline.del(KEYS.chunks(streamId)); - pipeline.del(KEYS.runSteps(streamId)); - await pipeline.exec(); + private async clearContentStateAsync( + streamId: string, + expectedCreatedAt?: number, + ): Promise { + await this.redis.eval( + CONTENT_CLEAR_LUA, + 3, + KEYS.chunks(streamId), + KEYS.runSteps(streamId), + KEYS.job(streamId), + expectedCreatedAt != null ? String(expectedCreatedAt) : '', + ); } // ===== Steering Queue Methods ===== @@ -1400,8 +1818,17 @@ export class RedisJobStore implements IJobStore { return this.parseSteerItems(raw); } - async peekSteers(streamId: string): Promise { - const raw = await this.redis.lrange(KEYS.steers(streamId), 0, -1); + async peekSteers(streamId: string, expectedCreatedAt?: number): Promise { + const raw = + expectedCreatedAt == null + ? await this.redis.lrange(KEYS.steers(streamId), 0, -1) + : await this.redis.eval( + STEER_PEEK_LUA, + 2, + KEYS.job(streamId), + KEYS.steers(streamId), + String(expectedCreatedAt), + ); return this.parseSteerItems(raw); } @@ -1419,46 +1846,17 @@ export class RedisJobStore implements IJobStore { return removed === 1; } - async parkSteers(streamId: string, payload: string): Promise { + async parkSteers(streamId: string, payload: string, expectedCreatedAt?: number): Promise { const ttl = this.ttl.completed > 0 ? this.ttl.completed : PARKED_RECOVERY_TTL_S; - await this.redis.set(KEYS.parkedSteers(streamId), payload, 'EX', ttl); - } - - /** Best-effort FIFO snapshot of the queued steers as client projections. */ - private async snapshotParkableSteers(streamId: string): Promise { - try { - return (await this.peekSteers(streamId)).map(toPendingSteer); - } catch (err) { - logger.warn(`[RedisJobStore] Failed to snapshot steers for parking ${streamId}:`, err); - return []; - } - } - - /** - * Parks an owner-carrying snapshot so /chat/status can recover 202-accepted - * steers after the job record is gone (approval expiry, stale-running reap). - * Best-effort: a park failure must not block the cleanup path that called it. - */ - private async parkSteerSnapshot( - streamId: string, - job: SerializableJobData, - steers: TPendingSteer[], - ): Promise { - if (steers.length === 0) { - return; - } - try { - await this.parkSteers( - streamId, - JSON.stringify({ - userId: job.userId, - ...(job.tenantId != null && { tenantId: job.tenantId }), - steers, - }), - ); - } catch (err) { - logger.warn(`[RedisJobStore] Failed to park steers ${streamId}:`, err); - } + await this.redis.eval( + PARK_STEERS_LUA, + 2, + KEYS.job(streamId), + KEYS.parkedSteers(streamId), + expectedCreatedAt != null ? String(expectedCreatedAt) : '', + payload, + String(ttl), + ); } async claimParkedSteers(streamId: string, ownerFragment: string): Promise { @@ -1496,7 +1894,7 @@ export class RedisJobStore implements IJobStore { * Uses XADD for efficient append-only storage. * Sets TTL on first chunk to ensure cleanup if job crashes. */ - async appendChunk(streamId: string, event: unknown): Promise { + async appendChunk(streamId: string, event: unknown, expectedCreatedAt?: number): Promise { const key = KEYS.chunks(streamId); const jobKey = KEYS.job(streamId); // XADD + derive-and-extend-only EXPIRE in a single atomic eval. Refreshing the TTL on @@ -1515,15 +1913,25 @@ export class RedisJobStore implements IJobStore { jobKey, JSON.stringify(event), String(this.ttl.running), + expectedCreatedAt != null ? String(expectedCreatedAt) : '', ); } /** * Get all chunks from Redis Stream. */ - private async getChunks(streamId: string): Promise { - const key = KEYS.chunks(streamId); - const entries = await this.redis.xrange(key, '-', '+'); + private async getChunks(streamId: string, expectedCreatedAt?: number): Promise { + const rawEntries = + expectedCreatedAt == null + ? await this.redis.xrange(KEYS.chunks(streamId), '-', '+') + : await this.redis.eval( + CHUNKS_READ_LUA, + 2, + KEYS.job(streamId), + KEYS.chunks(streamId), + String(expectedCreatedAt), + ); + const entries = Array.isArray(rawEntries) ? (rawEntries as Array<[string, string[]]>) : []; return entries .map(([, fields]) => { @@ -1546,7 +1954,11 @@ export class RedisJobStore implements IJobStore { * key to the short running TTL (which would drop the tool timeline on a reload of a * still-live approval — mirrors the chunk-stream no-shrink behavior). */ - async saveRunSteps(streamId: string, runSteps: Agents.RunStep[]): Promise { + async saveRunSteps( + streamId: string, + runSteps: Agents.RunStep[], + expectedCreatedAt?: number, + ): Promise { await this.redis.eval( RUNSTEPS_SAVE_LUA, 2, @@ -1554,6 +1966,7 @@ export class RedisJobStore implements IJobStore { KEYS.job(streamId), JSON.stringify(runSteps), String(this.ttl.running), + expectedCreatedAt != null ? String(expectedCreatedAt) : '', ); } diff --git a/packages/api/src/stream/index.ts b/packages/api/src/stream/index.ts index 722fce66fe..378325d265 100644 --- a/packages/api/src/stream/index.ts +++ b/packages/api/src/stream/index.ts @@ -1,6 +1,7 @@ export { GenerationJobManager, GenerationJobManagerClass, + type CreateGenerationJobOptions, type GenerationJobManagerOptions, } from './GenerationJobManager'; @@ -11,6 +12,7 @@ export type { UsageMetadata, AbortResult, JobStatus, + JobMetadataPatch, IJobStore, } from './interfaces/IJobStore'; // Canonical "is this approval live?" predicate — one definition shared by the diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index 3ee4755d64..2c2952aa73 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -126,6 +126,23 @@ export interface SerializableJobData { steersClosed?: boolean; } +export type JobMetadataPatch = Partial< + Pick< + SerializableJobData, + | 'responseMessageId' + | 'sender' + | 'conversationId' + | 'userMessage' + | 'endpoint' + | 'iconURL' + | 'model' + | 'agent_id' + | 'isTemporary' + | 'promptTokens' + | 'discoveredTools' + > +>; + /** * Whether a job's pending review has passed its `expiresAt`. Shared by the * stores so an expired approval is kept out of active-job listings (the client @@ -189,6 +206,12 @@ export interface JobStatusTransition { * resolve a job that has since paused for a different action. */ expectActionId?: string; + /** + * Additional guard: only fire if the job's creation epoch equals this value. + * Prevents a stale owner from transitioning a replacement job that reuses + * the same stream ID. + */ + expectCreatedAt?: number; } /** Value stored under an idempotency claim: the stream a retried request should attach to. */ @@ -351,13 +374,21 @@ export interface IJobStore { userId: string, conversationId?: string, tenantId?: string, + initialMetadata?: JobMetadataPatch, ): Promise; /** Get a job by streamId (streamId === conversationId) */ getJob(streamId: string): Promise; - /** Update job data */ - updateJob(streamId: string, updates: Partial): Promise; + /** + * Update job data. When `expectedCreatedAt` is supplied, apply the write only + * if the stream still belongs to that generation. + */ + updateJob( + streamId: string, + updates: Partial, + expectedCreatedAt?: number, + ): Promise; /** * Atomically transition a job's status, **only if** it is currently `from`. @@ -407,8 +438,11 @@ export interface IJobStore { */ releaseIdempotencyKey(key: string): Promise; - /** Delete a job */ - deleteJob(streamId: string): Promise; + /** + * Delete a job, optionally only when the stream still belongs to the expected + * generation. Returns true only when a matching job was actually deleted. + */ + deleteJob(streamId: string, expectedCreatedAt?: number): Promise; /** Check if job exists */ hasJob(streamId: string): Promise; @@ -428,8 +462,10 @@ export interface IJobStore { * Redis: no-op — the running-job TTL is already refreshed on each appendChunk. * * @param streamId - The stream identifier + * @param expectedCreatedAt - Optional generation identity. When supplied, replacement activity + * is not refreshed by a stale emitter. */ - recordActivity?(streamId: string): void; + recordActivity?(streamId: string, expectedCreatedAt?: number): void; /** Get total job count */ getJobCount(): Promise; @@ -465,7 +501,7 @@ export interface IJobStore { * @param streamId - The stream identifier * @param graph - The StandardGraph instance */ - setGraph(streamId: string, graph: StandardGraph): void; + setGraph(streamId: string, graph: StandardGraph, expectedCreatedAt?: number): void; /** * Set content parts reference for a job. @@ -476,7 +512,11 @@ export interface IJobStore { * @param streamId - The stream identifier * @param contentParts - The content parts array */ - setContentParts(streamId: string, contentParts: Agents.MessageContentComplex[]): void; + setContentParts( + streamId: string, + contentParts: Agents.MessageContentComplex[], + expectedCreatedAt?: number, + ): void; /** * Get aggregated content for a job. @@ -487,7 +527,10 @@ export interface IJobStore { * @param streamId - The stream identifier * @returns Content parts or null if not available */ - getContentParts(streamId: string): Promise<{ + getContentParts( + streamId: string, + expectedCreatedAt?: number, + ): Promise<{ content: Agents.MessageContentComplex[]; } | null>; @@ -500,7 +543,7 @@ export interface IJobStore { * @param streamId - The stream identifier * @returns Run steps or empty array */ - getRunSteps(streamId: string): Promise; + getRunSteps(streamId: string, expectedCreatedAt?: number): Promise; /** * Append a streaming chunk for later reconstruction. @@ -510,16 +553,20 @@ export interface IJobStore { * * @param streamId - The stream identifier * @param event - The SSE event to append + * @param expectedCreatedAt - Optional generation identity. When supplied, the append is + * refused if the stream ID now belongs to a replacement generation. */ - appendChunk(streamId: string, event: unknown): Promise; + appendChunk(streamId: string, event: unknown, expectedCreatedAt?: number): Promise; /** * Clear all content state for a job. * Called on job completion/cleanup. * * @param streamId - The stream identifier + * @param expectedCreatedAt - Optional generation identity. When supplied, replacement content + * is not cleared by a stale terminal cleanup. */ - clearContentState(streamId: string): void; + clearContentState(streamId: string, expectedCreatedAt?: number): void; /** * Save run steps to persistent storage. @@ -528,8 +575,14 @@ export interface IJobStore { * * @param streamId - The stream identifier * @param runSteps - Run steps to save + * @param expectedCreatedAt - Optional generation identity. When supplied, the save is refused + * if the stream ID now belongs to a replacement generation. */ - saveRunSteps?(streamId: string, runSteps: Agents.RunStep[]): Promise; + saveRunSteps?( + streamId: string, + runSteps: Agents.RunStep[], + expectedCreatedAt?: number, + ): Promise; /** * Set collected usage reference for a job. @@ -538,7 +591,11 @@ export interface IJobStore { * @param streamId - The stream identifier * @param collectedUsage - Array of usage metadata from all models */ - setCollectedUsage(streamId: string, collectedUsage: UsageMetadata[]): void; + setCollectedUsage( + streamId: string, + collectedUsage: UsageMetadata[], + expectedCreatedAt?: number, + ): void; /** * Get collected usage for a job. @@ -546,7 +603,7 @@ export interface IJobStore { * @param streamId - The stream identifier * @returns Array of usage metadata or empty array */ - getCollectedUsage(streamId: string): UsageMetadata[]; + getCollectedUsage(streamId: string, expectedCreatedAt?: number): UsageMetadata[]; // ===== Steering Queue Methods ===== // FIFO queue of mid-run user messages, keyed by streamId. Writable from any @@ -579,8 +636,12 @@ export interface IJobStore { */ closeAndDrainSteers(streamId: string, expectedCreatedAt?: number): Promise; - /** Non-destructive FIFO read of the queued steers (status/resume surfaces). */ - peekSteers(streamId: string): Promise; + /** + * Non-destructive FIFO read of the queued steers (status/resume surfaces). + * With `expectedCreatedAt`, returns an empty snapshot if the stream belongs + * to another generation. + */ + peekSteers(streamId: string, expectedCreatedAt?: number): Promise; /** Remove ONE queued steer by id (user-cancelled before injection). * False when it was no longer queued — already drained or run ended. */ @@ -593,8 +654,10 @@ export interface IJobStore { * path deletes the job immediately, and recovery must survive that. * Overwrites any prior payload; cleared by `createJob` (a replacement run * invalidates recovery — a live client had to start it). + * When `expectedCreatedAt` is supplied, the write is conditional on that + * generation still owning the stream ID. */ - parkSteers(streamId: string, payload: string): Promise; + parkSteers(streamId: string, payload: string, expectedCreatedAt?: number): Promise; /** * Claim-on-read: atomically return AND remove the parked payload, so a @@ -626,11 +689,17 @@ export interface IEventTransport { subscribe( streamId: string, handlers: { - onChunk: (event: unknown) => void; - onDone?: (event: unknown) => void; - onError?: (error: string) => void; + /** `generationId` identifies the immutable generation that emitted the chunk. */ + onChunk: (event: unknown, generationId?: number) => void; + /** `generationId` identifies the immutable generation that emitted the done event. */ + onDone?: (event: unknown, generationId?: number) => void; + /** `generationId` identifies the immutable generation that emitted the error. */ + onError?: (error: string, generationId?: number) => void; + }, + options?: { + /** Hold sequenced events until syncReorderBuffer establishes the replay frontier. */ + deferSequenceDelivery?: boolean; }, - options?: { deferSequenceDelivery?: boolean }, ): { unsubscribe: () => void; ready?: Promise }; /** @@ -638,13 +707,19 @@ export interface IEventTransport { * Redis returns the assigned absolute sequence so locally replayed events can * advance a subscriber to the exact ordering frontier. */ - emitChunk(streamId: string, event: unknown): void | Promise; + emitChunk(streamId: string, event: unknown, generationId?: number): void | Promise; - /** Publish a done event - returns Promise in Redis mode for ordered delivery */ - emitDone(streamId: string, event: unknown): void | Promise; + /** + * Publish a done event - returns Promise in Redis mode for ordered delivery. + * `generationId` is optional for compatibility with legacy, untagged publishers. + */ + emitDone(streamId: string, event: unknown, generationId?: number): void | Promise; - /** Publish an error event - returns Promise in Redis mode for ordered delivery */ - emitError(streamId: string, error: string): void | Promise; + /** + * Publish an error event - returns Promise in Redis mode for ordered delivery. + * `generationId` is optional for compatibility with legacy, untagged publishers. + */ + emitError(streamId: string, error: string, generationId?: number): void | Promise; /** * Publish an abort signal to all replicas (Redis mode). @@ -652,14 +727,20 @@ export interface IEventTransport { * generating Replica A receives signal and stops. * Optional - only implemented in Redis transport. */ - emitAbort?(streamId: string): void; + emitAbort?(streamId: string, generationId?: number): void; /** * Register callback for abort signals from any replica (Redis mode). * Called when abort is triggered from any replica. + * An async implementation resolves only after it can receive abort messages. + * The returned function removes only this registration, allowing a terminal + * generation to release its channel without affecting a same-stream replacement. * Optional - only implemented in Redis transport. */ - onAbort?(streamId: string, callback: () => void): void; + onAbort?( + streamId: string, + callback: (generationId?: number) => void, + ): void | (() => void) | Promise void)>; /** Get subscriber count for a stream */ getSubscriberCount(streamId: string): number; @@ -678,6 +759,12 @@ export interface IEventTransport { */ syncReorderBuffer?(streamId: string, replayedNextSeq?: number): void | Promise; + /** + * Notify and detach subscribers attached to this process without broadcasting a terminal event. + * Must trigger all-subscribers-left cleanup so graceful shutdown can drain partial persistence. + */ + closeLocalSubscribers?(streamId: string, error: string): void; + /** Cleanup transport resources for a specific stream */ cleanup(streamId: string): void; diff --git a/packages/api/src/stream/internal/chunkPublication.ts b/packages/api/src/stream/internal/chunkPublication.ts new file mode 100644 index 0000000000..0e66f13279 --- /dev/null +++ b/packages/api/src/stream/internal/chunkPublication.ts @@ -0,0 +1,43 @@ +import type { IEventTransport } from '../interfaces/IJobStore'; + +/** + * Internal publication result used to fence same-replica replay against Redis Pub/Sub. + * `number` is the zero-based Redis sequence and `false` means publication failed. + */ +export type ChunkPublicationReceipt = number | false | void; + +type ChunkPublicationCapability = ( + streamId: string, + event: unknown, + generationId?: number, +) => Promise; + +/** + * Keep transport-specific sequencing out of the exported IEventTransport contract. + * The WeakMap also keeps the capability off exported transport class declarations. + */ +const chunkPublicationCapabilities = new WeakMap(); + +export function registerChunkPublicationCapability( + transport: IEventTransport, + capability: ChunkPublicationCapability, +): void { + chunkPublicationCapabilities.set(transport, capability); +} + +/** + * Publish through a transport's internal receipt capability when available. + * Unsequenced and third-party transports retain the public emitChunk contract. + */ +export function emitChunkWithReceipt( + transport: IEventTransport, + streamId: string, + event: unknown, + generationId?: number, +): Promise { + const capability = chunkPublicationCapabilities.get(transport); + if (capability) { + return capability(streamId, event, generationId); + } + return Promise.resolve(transport.emitChunk(streamId, event, generationId)); +} diff --git a/packages/api/src/stream/metadata.ts b/packages/api/src/stream/metadata.ts new file mode 100644 index 0000000000..856e9c6413 --- /dev/null +++ b/packages/api/src/stream/metadata.ts @@ -0,0 +1,40 @@ +import type { JobMetadataPatch } from './interfaces/IJobStore'; +import type { GenerationJobMetadata } from '~/types'; + +export function sanitizeJobMetadata(metadata: Partial): JobMetadataPatch { + const patch: JobMetadataPatch = {}; + if (metadata.responseMessageId) { + patch.responseMessageId = metadata.responseMessageId; + } + if (metadata.sender) { + patch.sender = metadata.sender; + } + if (metadata.conversationId) { + patch.conversationId = metadata.conversationId; + } + if (metadata.userMessage) { + patch.userMessage = metadata.userMessage; + } + if (metadata.endpoint) { + patch.endpoint = metadata.endpoint; + } + if (metadata.iconURL) { + patch.iconURL = metadata.iconURL; + } + if (metadata.model) { + patch.model = metadata.model; + } + if (metadata.agent_id) { + patch.agent_id = metadata.agent_id; + } + if (metadata.isTemporary !== undefined) { + patch.isTemporary = metadata.isTemporary; + } + if (metadata.promptTokens !== undefined) { + patch.promptTokens = metadata.promptTokens; + } + if (metadata.discoveredTools) { + patch.discoveredTools = metadata.discoveredTools; + } + return patch; +} diff --git a/packages/api/src/telemetry/sdk.spec.ts b/packages/api/src/telemetry/sdk.spec.ts index c361a6cd5c..73b499e6ba 100644 --- a/packages/api/src/telemetry/sdk.spec.ts +++ b/packages/api/src/telemetry/sdk.spec.ts @@ -512,7 +512,9 @@ describe('telemetry SDK lifecycle', () => { const { registerShutdownTask } = await import('../app/shutdown'); initializeTelemetry({ OTEL_TRACING_ENABLED: 'true' }); - expect(registerShutdownTask).toHaveBeenCalledWith('telemetry', expect.any(Function)); + expect(registerShutdownTask).toHaveBeenCalledWith('telemetry', expect.any(Function), { + priority: -100, + }); }); it('the registered shutdown task warns when telemetry shutdown rejects', async () => { diff --git a/packages/api/src/telemetry/sdk.ts b/packages/api/src/telemetry/sdk.ts index c3fb90b852..6aa8a735fa 100644 --- a/packages/api/src/telemetry/sdk.ts +++ b/packages/api/src/telemetry/sdk.ts @@ -382,10 +382,14 @@ function ensureShutdownTaskRegistered(): void { // by Node, so a separate signal handler can let the coordinator // exit before the async OpenTelemetry flush completes, dropping // final spans during pod shutdowns. - registerShutdownTask('telemetry', () => - withTimeout(shutdownTelemetry(), SIGNAL_SHUTDOWN_TIMEOUT_MS).catch((error) => { - emitWarning(`OpenTelemetry shutdown failed: ${getErrorMessage(error)}`); - }), + registerShutdownTask( + 'telemetry', + () => + withTimeout(shutdownTelemetry(), SIGNAL_SHUTDOWN_TIMEOUT_MS).catch((error) => { + emitWarning(`OpenTelemetry shutdown failed: ${getErrorMessage(error)}`); + }), + // Exporters close after instrumented modules have emitted their final shutdown spans. + { priority: -100 }, ); } diff --git a/packages/api/src/types/stream.ts b/packages/api/src/types/stream.ts index 47a94d6ead..75afee28cb 100644 --- a/packages/api/src/types/stream.ts +++ b/packages/api/src/types/stream.ts @@ -61,6 +61,16 @@ export type DoneHandler = (event: ServerSentEvent) => void; export type ErrorHandler = (error: string) => void; export type UnsubscribeFn = () => void; +/** Active event-stream subscription. */ +export interface StreamSubscription { + unsubscribe: UnsubscribeFn; +} + +/** Resume subscription whose live delivery starts after the caller writes its sync frame. */ +export interface ResumeSubscription extends StreamSubscription { + activate: () => void; +} + /** Options for subscribing to a job event stream */ export interface SubscribeOptions { /** @@ -68,11 +78,13 @@ export interface SubscribeOptions { * Use for resume connections after a sync event has been sent. */ skipBufferReplay?: boolean; + /** Cancels attachment work when the HTTP client disconnects. */ + signal?: AbortSignal; } /** Result of an atomic subscribe-with-resume operation */ export interface SubscribeWithResumeResult { - subscription: { unsubscribe: UnsubscribeFn } | null; + subscription: ResumeSubscription | null; resumeState: ResumeState | null; /** * Events that arrived between the resume snapshot and the subscribe call.