mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-01-23 02:36:12 +01:00
🏁 fix: Message Race Condition if Cancelled Early (#11462)
* 🔧 fix: Prevent race conditions in message saving during abort scenarios * Added logic to save partial responses before returning from the abort endpoint to ensure parentMessageId exists in the database. * Updated the ResumableAgentController to save response messages before emitting final events, preventing orphaned parentMessageIds. * Enhanced handling of unfinished responses to improve stability and data integrity in agent interactions. * 🔧 fix: logging and job replacement handling in ResumableAgentController * Added detailed logging for job creation and final event emissions to improve traceability. * Implemented logic to check for job replacement before emitting events, preventing stale requests from affecting newer jobs. * Updated abort handling to log additional context about the abort result, enhancing debugging capabilities. * refactor: abort handling and token spending logic in AgentStream * Added authorization check for abort attempts to prevent unauthorized access. * Improved response message saving logic to ensure valid message IDs are stored. * Implemented token spending for aborted requests to prevent double-spending across parallel agents. * Enhanced logging for better traceability of token spending operations during abort scenarios. * refactor: remove TODO comments for token spending in abort handling * Removed outdated TODO comments regarding token spending for aborted requests in the abort endpoint. * This change streamlines the code and clarifies the current implementation status. * ✅ test: Add comprehensive tests for job replacement and abort handling * Introduced unit tests for job replacement detection in ResumableAgentController, covering job creation timestamp tracking, stale job detection, and response message saving order. * Added tests for the agent abort endpoint, ensuring proper authorization checks, early abort handling, and partial response saving. * Enhanced logging and error handling in tests to improve traceability and robustness of the abort functionality.
This commit is contained in:
Danny Avila
GitHub
parent
dea246934e
commit
11210d8b98
4 changed files with 682 additions and 12 deletions
281
api/server/controllers/agents/__tests__/jobReplacement.spec.js
Normal file
281
api/server/controllers/agents/__tests__/jobReplacement.spec.js
Normal file
|
|
@ -0,0 +1,281 @@
|
|||
/**
|
||||
* Tests for job replacement detection in ResumableAgentController
|
||||
*
|
||||
* Tests the following fixes from PR #11462:
|
||||
* 1. Job creation timestamp tracking
|
||||
* 2. Stale job detection and event skipping
|
||||
* 3. Response message saving before final event emission
|
||||
*/
|
||||
|
||||
const mockLogger = {
|
||||
debug: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
error: jest.fn(),
|
||||
info: jest.fn(),
|
||||
};
|
||||
|
||||
const mockGenerationJobManager = {
|
||||
createJob: jest.fn(),
|
||||
getJob: jest.fn(),
|
||||
emitDone: jest.fn(),
|
||||
emitChunk: jest.fn(),
|
||||
completeJob: jest.fn(),
|
||||
updateMetadata: jest.fn(),
|
||||
setContentParts: jest.fn(),
|
||||
subscribe: jest.fn(),
|
||||
};
|
||||
|
||||
const mockSaveMessage = jest.fn();
|
||||
const mockDecrementPendingRequest = jest.fn();
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: mockLogger,
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/api', () => ({
|
||||
isEnabled: jest.fn().mockReturnValue(false),
|
||||
GenerationJobManager: mockGenerationJobManager,
|
||||
checkAndIncrementPendingRequest: jest.fn().mockResolvedValue({ allowed: true }),
|
||||
decrementPendingRequest: (...args) => mockDecrementPendingRequest(...args),
|
||||
getViolationInfo: jest.fn(),
|
||||
sanitizeMessageForTransmit: jest.fn((msg) => msg),
|
||||
sanitizeFileForTransmit: jest.fn((file) => file),
|
||||
Constants: { NO_PARENT: '00000000-0000-0000-0000-000000000000' },
|
||||
}));
|
||||
|
||||
jest.mock('~/models', () => ({
|
||||
saveMessage: (...args) => mockSaveMessage(...args),
|
||||
}));
|
||||
|
||||
describe('Job Replacement Detection', () => {
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('Job Creation Timestamp Tracking', () => {
|
||||
it('should capture createdAt when job is created', async () => {
|
||||
const streamId = 'test-stream-123';
|
||||
const createdAt = Date.now();
|
||||
|
||||
mockGenerationJobManager.createJob.mockResolvedValue({
|
||||
createdAt,
|
||||
readyPromise: Promise.resolve(),
|
||||
abortController: new AbortController(),
|
||||
emitter: { on: jest.fn() },
|
||||
});
|
||||
|
||||
const job = await mockGenerationJobManager.createJob(streamId, 'user-123', streamId);
|
||||
|
||||
expect(job.createdAt).toBe(createdAt);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Job Replacement Detection Logic', () => {
|
||||
/**
|
||||
* Simulates the job replacement detection logic from request.js
|
||||
* This is extracted for unit testing since the full controller is complex
|
||||
*/
|
||||
const detectJobReplacement = async (streamId, originalCreatedAt) => {
|
||||
const currentJob = await mockGenerationJobManager.getJob(streamId);
|
||||
return !currentJob || currentJob.createdAt !== originalCreatedAt;
|
||||
};
|
||||
|
||||
it('should detect when job was replaced (different createdAt)', async () => {
|
||||
const streamId = 'test-stream-123';
|
||||
const originalCreatedAt = 1000;
|
||||
const newCreatedAt = 2000;
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
createdAt: newCreatedAt,
|
||||
});
|
||||
|
||||
const wasReplaced = await detectJobReplacement(streamId, originalCreatedAt);
|
||||
|
||||
expect(wasReplaced).toBe(true);
|
||||
});
|
||||
|
||||
it('should detect when job was deleted', async () => {
|
||||
const streamId = 'test-stream-123';
|
||||
const originalCreatedAt = 1000;
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue(null);
|
||||
|
||||
const wasReplaced = await detectJobReplacement(streamId, originalCreatedAt);
|
||||
|
||||
expect(wasReplaced).toBe(true);
|
||||
});
|
||||
|
||||
it('should not detect replacement when same job (same createdAt)', async () => {
|
||||
const streamId = 'test-stream-123';
|
||||
const originalCreatedAt = 1000;
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
createdAt: originalCreatedAt,
|
||||
});
|
||||
|
||||
const wasReplaced = await detectJobReplacement(streamId, originalCreatedAt);
|
||||
|
||||
expect(wasReplaced).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Event Emission Behavior', () => {
|
||||
/**
|
||||
* Simulates the final event emission logic from request.js
|
||||
*/
|
||||
const emitFinalEventIfNotReplaced = async ({
|
||||
streamId,
|
||||
originalCreatedAt,
|
||||
finalEvent,
|
||||
userId,
|
||||
}) => {
|
||||
const currentJob = await mockGenerationJobManager.getJob(streamId);
|
||||
const jobWasReplaced = !currentJob || currentJob.createdAt !== originalCreatedAt;
|
||||
|
||||
if (jobWasReplaced) {
|
||||
mockLogger.debug('Skipping FINAL emit - job was replaced', {
|
||||
streamId,
|
||||
originalCreatedAt,
|
||||
currentCreatedAt: currentJob?.createdAt,
|
||||
});
|
||||
await mockDecrementPendingRequest(userId);
|
||||
return false;
|
||||
}
|
||||
|
||||
mockGenerationJobManager.emitDone(streamId, finalEvent);
|
||||
mockGenerationJobManager.completeJob(streamId);
|
||||
await mockDecrementPendingRequest(userId);
|
||||
return true;
|
||||
};
|
||||
|
||||
it('should skip emitting when job was replaced', async () => {
|
||||
const streamId = 'test-stream-123';
|
||||
const originalCreatedAt = 1000;
|
||||
const newCreatedAt = 2000;
|
||||
const userId = 'user-123';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
createdAt: newCreatedAt,
|
||||
});
|
||||
|
||||
const emitted = await emitFinalEventIfNotReplaced({
|
||||
streamId,
|
||||
originalCreatedAt,
|
||||
finalEvent: { final: true },
|
||||
userId,
|
||||
});
|
||||
|
||||
expect(emitted).toBe(false);
|
||||
expect(mockGenerationJobManager.emitDone).not.toHaveBeenCalled();
|
||||
expect(mockGenerationJobManager.completeJob).not.toHaveBeenCalled();
|
||||
expect(mockDecrementPendingRequest).toHaveBeenCalledWith(userId);
|
||||
expect(mockLogger.debug).toHaveBeenCalledWith(
|
||||
'Skipping FINAL emit - job was replaced',
|
||||
expect.objectContaining({
|
||||
streamId,
|
||||
originalCreatedAt,
|
||||
currentCreatedAt: newCreatedAt,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should emit when job was not replaced', async () => {
|
||||
const streamId = 'test-stream-123';
|
||||
const originalCreatedAt = 1000;
|
||||
const userId = 'user-123';
|
||||
const finalEvent = { final: true, conversation: { conversationId: streamId } };
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
createdAt: originalCreatedAt,
|
||||
});
|
||||
|
||||
const emitted = await emitFinalEventIfNotReplaced({
|
||||
streamId,
|
||||
originalCreatedAt,
|
||||
finalEvent,
|
||||
userId,
|
||||
});
|
||||
|
||||
expect(emitted).toBe(true);
|
||||
expect(mockGenerationJobManager.emitDone).toHaveBeenCalledWith(streamId, finalEvent);
|
||||
expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith(streamId);
|
||||
expect(mockDecrementPendingRequest).toHaveBeenCalledWith(userId);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Response Message Saving Order', () => {
|
||||
/**
|
||||
* Tests that response messages are saved BEFORE final events are emitted
|
||||
* This prevents race conditions where clients send follow-up messages
|
||||
* before the response is in the database
|
||||
*/
|
||||
it('should save message before emitting final event', async () => {
|
||||
const callOrder = [];
|
||||
|
||||
mockSaveMessage.mockImplementation(async () => {
|
||||
callOrder.push('saveMessage');
|
||||
});
|
||||
|
||||
mockGenerationJobManager.emitDone.mockImplementation(() => {
|
||||
callOrder.push('emitDone');
|
||||
});
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
createdAt: 1000,
|
||||
});
|
||||
|
||||
// Simulate the order of operations from request.js
|
||||
const streamId = 'test-stream-123';
|
||||
const originalCreatedAt = 1000;
|
||||
const response = { messageId: 'response-123' };
|
||||
const userId = 'user-123';
|
||||
|
||||
// Step 1: Save message
|
||||
await mockSaveMessage({}, { ...response, user: userId }, { context: 'test' });
|
||||
|
||||
// Step 2: Check for replacement
|
||||
const currentJob = await mockGenerationJobManager.getJob(streamId);
|
||||
const jobWasReplaced = !currentJob || currentJob.createdAt !== originalCreatedAt;
|
||||
|
||||
// Step 3: Emit if not replaced
|
||||
if (!jobWasReplaced) {
|
||||
mockGenerationJobManager.emitDone(streamId, { final: true });
|
||||
}
|
||||
|
||||
expect(callOrder).toEqual(['saveMessage', 'emitDone']);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Aborted Request Handling', () => {
|
||||
it('should use unfinished: true instead of error: true for aborted requests', () => {
|
||||
const response = { messageId: 'response-123', content: [] };
|
||||
|
||||
// The new format for aborted responses
|
||||
const abortedResponse = { ...response, unfinished: true };
|
||||
|
||||
expect(abortedResponse.unfinished).toBe(true);
|
||||
expect(abortedResponse.error).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should include unfinished flag in final event for aborted requests', () => {
|
||||
const response = { messageId: 'response-123', content: [] };
|
||||
|
||||
// Old format (deprecated)
|
||||
const _oldFinalEvent = {
|
||||
final: true,
|
||||
responseMessage: { ...response, error: true },
|
||||
error: { message: 'Request was aborted' },
|
||||
};
|
||||
|
||||
// New format (PR #11462)
|
||||
const newFinalEvent = {
|
||||
final: true,
|
||||
responseMessage: { ...response, unfinished: true },
|
||||
};
|
||||
|
||||
expect(newFinalEvent.responseMessage.unfinished).toBe(true);
|
||||
expect(newFinalEvent.error).toBeUndefined();
|
||||
expect(newFinalEvent.responseMessage.error).toBeUndefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -67,7 +67,15 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
|
|||
let client = null;
|
||||
|
||||
try {
|
||||
logger.debug(`[ResumableAgentController] Creating job`, {
|
||||
streamId,
|
||||
conversationId,
|
||||
reqConversationId,
|
||||
userId,
|
||||
});
|
||||
|
||||
const job = await GenerationJobManager.createJob(streamId, userId, conversationId);
|
||||
const jobCreatedAt = job.createdAt; // Capture creation time to detect job replacement
|
||||
req._resumableStreamId = streamId;
|
||||
|
||||
// Send JSON response IMMEDIATELY so client can connect to SSE stream
|
||||
|
|
@ -272,6 +280,33 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
|
|||
});
|
||||
}
|
||||
|
||||
// CRITICAL: Save response message BEFORE emitting final event.
|
||||
// This prevents race conditions where the client sends a follow-up message
|
||||
// before the response is saved to the database, causing orphaned parentMessageIds.
|
||||
if (client.savedMessageIds && !client.savedMessageIds.has(messageId)) {
|
||||
await saveMessage(
|
||||
req,
|
||||
{ ...response, user: userId, unfinished: wasAbortedBeforeComplete },
|
||||
{ context: 'api/server/controllers/agents/request.js - resumable response end' },
|
||||
);
|
||||
}
|
||||
|
||||
// Check if our job was replaced by a new request before emitting
|
||||
// This prevents stale requests from emitting events to newer jobs
|
||||
const currentJob = await GenerationJobManager.getJob(streamId);
|
||||
const jobWasReplaced = !currentJob || currentJob.createdAt !== jobCreatedAt;
|
||||
|
||||
if (jobWasReplaced) {
|
||||
logger.debug(`[ResumableAgentController] Skipping FINAL emit - job was replaced`, {
|
||||
streamId,
|
||||
originalCreatedAt: jobCreatedAt,
|
||||
currentCreatedAt: currentJob?.createdAt,
|
||||
});
|
||||
// Still decrement pending request since we incremented at start
|
||||
await decrementPendingRequest(userId);
|
||||
return;
|
||||
}
|
||||
|
||||
if (!wasAbortedBeforeComplete) {
|
||||
const finalEvent = {
|
||||
final: true,
|
||||
|
|
@ -281,26 +316,34 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
|
|||
responseMessage: { ...response },
|
||||
};
|
||||
|
||||
logger.debug(`[ResumableAgentController] Emitting FINAL event`, {
|
||||
streamId,
|
||||
wasAbortedBeforeComplete,
|
||||
userMessageId: userMessage?.messageId,
|
||||
responseMessageId: response?.messageId,
|
||||
conversationId: conversation?.conversationId,
|
||||
});
|
||||
|
||||
GenerationJobManager.emitDone(streamId, finalEvent);
|
||||
GenerationJobManager.completeJob(streamId);
|
||||
await decrementPendingRequest(userId);
|
||||
|
||||
if (client.savedMessageIds && !client.savedMessageIds.has(messageId)) {
|
||||
await saveMessage(
|
||||
req,
|
||||
{ ...response, user: userId },
|
||||
{ context: 'api/server/controllers/agents/request.js - resumable response end' },
|
||||
);
|
||||
}
|
||||
} else {
|
||||
const finalEvent = {
|
||||
final: true,
|
||||
conversation,
|
||||
title: conversation.title,
|
||||
requestMessage: sanitizeMessageForTransmit(userMessage),
|
||||
responseMessage: { ...response, error: true },
|
||||
error: { message: 'Request was aborted' },
|
||||
responseMessage: { ...response, unfinished: true },
|
||||
};
|
||||
|
||||
logger.debug(`[ResumableAgentController] Emitting ABORTED FINAL event`, {
|
||||
streamId,
|
||||
wasAbortedBeforeComplete,
|
||||
userMessageId: userMessage?.messageId,
|
||||
responseMessageId: response?.messageId,
|
||||
conversationId: conversation?.conversationId,
|
||||
});
|
||||
|
||||
GenerationJobManager.emitDone(streamId, finalEvent);
|
||||
GenerationJobManager.completeJob(streamId, 'Request aborted');
|
||||
await decrementPendingRequest(userId);
|
||||
|
|
|
|||
301
api/server/routes/agents/__tests__/abort.spec.js
Normal file
301
api/server/routes/agents/__tests__/abort.spec.js
Normal file
|
|
@ -0,0 +1,301 @@
|
|||
/**
|
||||
* Tests for the agent abort endpoint
|
||||
*
|
||||
* Tests the following fixes from PR #11462:
|
||||
* 1. Authorization check - only job owner can abort
|
||||
* 2. Early abort handling - skip save when no responseMessageId
|
||||
* 3. Partial response saving - save message before returning
|
||||
*/
|
||||
|
||||
const express = require('express');
|
||||
const request = require('supertest');
|
||||
|
||||
const mockLogger = {
|
||||
debug: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
error: jest.fn(),
|
||||
info: jest.fn(),
|
||||
};
|
||||
|
||||
const mockGenerationJobManager = {
|
||||
getJob: jest.fn(),
|
||||
abortJob: jest.fn(),
|
||||
getActiveJobIdsForUser: jest.fn(),
|
||||
};
|
||||
|
||||
const mockSaveMessage = jest.fn();
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: mockLogger,
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/api', () => ({
|
||||
isEnabled: jest.fn().mockReturnValue(false),
|
||||
GenerationJobManager: mockGenerationJobManager,
|
||||
}));
|
||||
|
||||
jest.mock('~/models', () => ({
|
||||
saveMessage: (...args) => mockSaveMessage(...args),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/middleware', () => ({
|
||||
uaParser: (req, res, next) => next(),
|
||||
checkBan: (req, res, next) => next(),
|
||||
requireJwtAuth: (req, res, next) => {
|
||||
req.user = { id: 'test-user-123' };
|
||||
next();
|
||||
},
|
||||
messageIpLimiter: (req, res, next) => next(),
|
||||
configMiddleware: (req, res, next) => next(),
|
||||
messageUserLimiter: (req, res, next) => next(),
|
||||
}));
|
||||
|
||||
// Mock the chat module - needs to be a router
|
||||
jest.mock('~/server/routes/agents/chat', () => require('express').Router());
|
||||
|
||||
// Mock the v1 module - v1 is directly used as middleware
|
||||
jest.mock('~/server/routes/agents/v1', () => ({
|
||||
v1: require('express').Router(),
|
||||
}));
|
||||
|
||||
// Import after mocks
|
||||
const agentRoutes = require('~/server/routes/agents/index');
|
||||
|
||||
describe('Agent Abort Endpoint', () => {
|
||||
let app;
|
||||
|
||||
beforeAll(() => {
|
||||
app = express();
|
||||
app.use(express.json());
|
||||
app.use('/api/agents', agentRoutes);
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('POST /chat/abort', () => {
|
||||
describe('Authorization', () => {
|
||||
it("should return 403 when user tries to abort another user's job", async () => {
|
||||
const jobStreamId = 'test-stream-123';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
metadata: { userId: 'other-user-456' },
|
||||
});
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: jobStreamId });
|
||||
|
||||
expect(response.status).toBe(403);
|
||||
expect(response.body).toEqual({ error: 'Unauthorized' });
|
||||
expect(mockLogger.warn).toHaveBeenCalledWith(
|
||||
expect.stringContaining('Unauthorized abort attempt'),
|
||||
);
|
||||
expect(mockGenerationJobManager.abortJob).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should allow abort when user owns the job', async () => {
|
||||
const jobStreamId = 'test-stream-123';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
metadata: { userId: 'test-user-123' },
|
||||
});
|
||||
|
||||
mockGenerationJobManager.abortJob.mockResolvedValue({
|
||||
success: true,
|
||||
jobData: null,
|
||||
content: [],
|
||||
text: '',
|
||||
});
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: jobStreamId });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body).toEqual({ success: true, aborted: jobStreamId });
|
||||
expect(mockGenerationJobManager.abortJob).toHaveBeenCalledWith(jobStreamId);
|
||||
});
|
||||
|
||||
it('should allow abort when job has no userId metadata (backwards compatibility)', async () => {
|
||||
const jobStreamId = 'test-stream-123';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
metadata: {},
|
||||
});
|
||||
|
||||
mockGenerationJobManager.abortJob.mockResolvedValue({
|
||||
success: true,
|
||||
jobData: null,
|
||||
content: [],
|
||||
text: '',
|
||||
});
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: jobStreamId });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body).toEqual({ success: true, aborted: jobStreamId });
|
||||
});
|
||||
});
|
||||
|
||||
describe('Early Abort Handling', () => {
|
||||
it('should skip message saving when responseMessageId is missing (early abort)', async () => {
|
||||
const jobStreamId = 'test-stream-123';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
metadata: { userId: 'test-user-123' },
|
||||
});
|
||||
|
||||
mockGenerationJobManager.abortJob.mockResolvedValue({
|
||||
success: true,
|
||||
jobData: {
|
||||
userMessage: { messageId: 'user-msg-123' },
|
||||
// No responseMessageId - early abort before generation started
|
||||
conversationId: jobStreamId,
|
||||
},
|
||||
content: [],
|
||||
text: '',
|
||||
});
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: jobStreamId });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(mockSaveMessage).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should skip message saving when userMessage is missing', async () => {
|
||||
const jobStreamId = 'test-stream-123';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
metadata: { userId: 'test-user-123' },
|
||||
});
|
||||
|
||||
mockGenerationJobManager.abortJob.mockResolvedValue({
|
||||
success: true,
|
||||
jobData: {
|
||||
// No userMessage
|
||||
responseMessageId: 'response-msg-123',
|
||||
conversationId: jobStreamId,
|
||||
},
|
||||
content: [],
|
||||
text: '',
|
||||
});
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: jobStreamId });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(mockSaveMessage).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Partial Response Saving', () => {
|
||||
it('should save partial response when both userMessage and responseMessageId exist', async () => {
|
||||
const jobStreamId = 'test-stream-123';
|
||||
const userMessageId = 'user-msg-123';
|
||||
const responseMessageId = 'response-msg-456';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
metadata: { userId: 'test-user-123' },
|
||||
});
|
||||
|
||||
mockGenerationJobManager.abortJob.mockResolvedValue({
|
||||
success: true,
|
||||
jobData: {
|
||||
userMessage: { messageId: userMessageId },
|
||||
responseMessageId,
|
||||
conversationId: jobStreamId,
|
||||
sender: 'TestAgent',
|
||||
endpoint: 'anthropic',
|
||||
model: 'claude-3',
|
||||
},
|
||||
content: [{ type: 'text', text: 'Partial response...' }],
|
||||
text: 'Partial response...',
|
||||
});
|
||||
|
||||
mockSaveMessage.mockResolvedValue();
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: jobStreamId });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(mockSaveMessage).toHaveBeenCalledWith(
|
||||
expect.anything(),
|
||||
expect.objectContaining({
|
||||
messageId: responseMessageId,
|
||||
parentMessageId: userMessageId,
|
||||
conversationId: jobStreamId,
|
||||
content: [{ type: 'text', text: 'Partial response...' }],
|
||||
text: 'Partial response...',
|
||||
sender: 'TestAgent',
|
||||
endpoint: 'anthropic',
|
||||
model: 'claude-3',
|
||||
unfinished: true,
|
||||
error: false,
|
||||
isCreatedByUser: false,
|
||||
user: 'test-user-123',
|
||||
}),
|
||||
expect.objectContaining({
|
||||
context: 'api/server/routes/agents/index.js - abort endpoint',
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle saveMessage errors gracefully', async () => {
|
||||
const jobStreamId = 'test-stream-123';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
metadata: { userId: 'test-user-123' },
|
||||
});
|
||||
|
||||
mockGenerationJobManager.abortJob.mockResolvedValue({
|
||||
success: true,
|
||||
jobData: {
|
||||
userMessage: { messageId: 'user-msg-123' },
|
||||
responseMessageId: 'response-msg-456',
|
||||
conversationId: jobStreamId,
|
||||
},
|
||||
content: [],
|
||||
text: '',
|
||||
});
|
||||
|
||||
mockSaveMessage.mockRejectedValue(new Error('Database error'));
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: jobStreamId });
|
||||
|
||||
// Should still return success even if save fails
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body).toEqual({ success: true, aborted: jobStreamId });
|
||||
expect(mockLogger.error).toHaveBeenCalledWith(
|
||||
expect.stringContaining('Failed to save partial response'),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Job Not Found', () => {
|
||||
it('should return 404 when job is not found', async () => {
|
||||
mockGenerationJobManager.getJob.mockResolvedValue(null);
|
||||
mockGenerationJobManager.getActiveJobIdsForUser.mockResolvedValue([]);
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: 'non-existent-job' });
|
||||
|
||||
expect(response.status).toBe(404);
|
||||
expect(response.body).toEqual({
|
||||
error: 'Job not found',
|
||||
streamId: 'non-existent-job',
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -9,6 +9,7 @@ const {
|
|||
configMiddleware,
|
||||
messageUserLimiter,
|
||||
} = require('~/server/middleware');
|
||||
const { saveMessage } = require('~/models');
|
||||
const { v1 } = require('./v1');
|
||||
const chat = require('./chat');
|
||||
|
||||
|
|
@ -194,9 +195,53 @@ router.post('/chat/abort', async (req, res) => {
|
|||
logger.debug(`[AgentStream] Computed jobStreamId: ${jobStreamId}`);
|
||||
|
||||
if (job && jobStreamId) {
|
||||
if (job.metadata?.userId && job.metadata.userId !== userId) {
|
||||
logger.warn(`[AgentStream] Unauthorized abort attempt for ${jobStreamId} by user ${userId}`);
|
||||
return res.status(403).json({ error: 'Unauthorized' });
|
||||
}
|
||||
|
||||
logger.debug(`[AgentStream] Job found, aborting: ${jobStreamId}`);
|
||||
await GenerationJobManager.abortJob(jobStreamId);
|
||||
logger.debug(`[AgentStream] Job aborted successfully: ${jobStreamId}`);
|
||||
const abortResult = await GenerationJobManager.abortJob(jobStreamId);
|
||||
logger.debug(`[AgentStream] Job aborted successfully: ${jobStreamId}`, {
|
||||
abortResultSuccess: abortResult.success,
|
||||
abortResultUserMessageId: abortResult.jobData?.userMessage?.messageId,
|
||||
abortResultResponseMessageId: abortResult.jobData?.responseMessageId,
|
||||
});
|
||||
|
||||
// CRITICAL: Save partial response BEFORE returning to prevent race condition.
|
||||
// If user sends a follow-up immediately after abort, the parentMessageId must exist in DB.
|
||||
// Only save if we have a valid responseMessageId (skip early aborts before generation started)
|
||||
if (
|
||||
abortResult.success &&
|
||||
abortResult.jobData?.userMessage?.messageId &&
|
||||
abortResult.jobData?.responseMessageId
|
||||
) {
|
||||
const { jobData, content, text } = abortResult;
|
||||
const responseMessage = {
|
||||
messageId: jobData.responseMessageId,
|
||||
parentMessageId: jobData.userMessage.messageId,
|
||||
conversationId: jobData.conversationId,
|
||||
content: content || [],
|
||||
text: text || '',
|
||||
sender: jobData.sender || 'AI',
|
||||
endpoint: jobData.endpoint,
|
||||
model: jobData.model,
|
||||
unfinished: true,
|
||||
error: false,
|
||||
isCreatedByUser: false,
|
||||
user: userId,
|
||||
};
|
||||
|
||||
try {
|
||||
await saveMessage(req, responseMessage, {
|
||||
context: 'api/server/routes/agents/index.js - abort endpoint',
|
||||
});
|
||||
logger.debug(`[AgentStream] Saved partial response for: ${jobStreamId}`);
|
||||
} catch (saveError) {
|
||||
logger.error(`[AgentStream] Failed to save partial response: ${saveError.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
return res.json({ success: true, aborted: jobStreamId });
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue