mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-04 14:57:42 +00:00
⚡ perf: Reduce Agent Chat Startup Latency (#14423)
* 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
This commit is contained in:
Danny Avila
GitHub
parent
cd215150cc
commit
73699b5c25
64 changed files with 12134 additions and 1744 deletions
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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 });
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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 () => {
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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<typeof buildAgentScopedContext>} */
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue