mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-27 12:13:30 +00:00
🧾 fix: Honor Disabled Transactions on the Abort Paths (#15099)
Resolve the transactions config from the request and forward it to both abort write paths, so `transactions.enabled: false` is honored when a generation is stopped.
This commit is contained in:
parent
cc0111b3cf
commit
4d246469dd
2 changed files with 118 additions and 2 deletions
|
|
@ -6,6 +6,7 @@ const {
|
|||
countTokens,
|
||||
GenerationJobManager,
|
||||
recordCollectedUsage,
|
||||
getTransactionsConfig,
|
||||
sanitizeMessageForTransmit,
|
||||
buildAbortedResponseMetadata,
|
||||
} = require('@librechat/api');
|
||||
|
|
@ -59,6 +60,7 @@ const isAbortError = (error) => {
|
|||
* @param {Array<Object>} params.collectedUsage - Usage metadata from all models
|
||||
* @param {string} [params.fallbackModel] - Fallback model name if not in usage
|
||||
* @param {string} [params.messageId] - The response message ID for transaction correlation
|
||||
* @param {AppConfig['transactions']} [params.transactions] - Resolved transactions config
|
||||
*/
|
||||
async function spendCollectedUsage({
|
||||
userId,
|
||||
|
|
@ -66,6 +68,7 @@ async function spendCollectedUsage({
|
|||
collectedUsage,
|
||||
fallbackModel,
|
||||
messageId,
|
||||
transactions,
|
||||
}) {
|
||||
if (!collectedUsage || collectedUsage.length === 0) {
|
||||
return;
|
||||
|
|
@ -85,6 +88,7 @@ async function spendCollectedUsage({
|
|||
context: 'abort',
|
||||
messageId,
|
||||
model: fallbackModel,
|
||||
transactions,
|
||||
},
|
||||
);
|
||||
|
||||
|
|
@ -149,6 +153,8 @@ async function abortMessage(req, res) {
|
|||
responseMessage.metadata = abortMetadata;
|
||||
}
|
||||
|
||||
const transactions = getTransactionsConfig(req.config);
|
||||
|
||||
// Spend tokens for ALL models from collectedUsage (handles parallel agents/addedConvo)
|
||||
if (collectedUsage && collectedUsage.length > 0) {
|
||||
await spendCollectedUsage({
|
||||
|
|
@ -157,11 +163,12 @@ async function abortMessage(req, res) {
|
|||
collectedUsage,
|
||||
fallbackModel: jobData?.model,
|
||||
messageId: jobData?.responseMessageId,
|
||||
transactions,
|
||||
});
|
||||
} else {
|
||||
// Fallback: no collected usage, use text-based token counting for primary model only
|
||||
await db.spendTokens(
|
||||
{ ...responseMessage, context: 'incomplete', user: userId },
|
||||
{ ...responseMessage, context: 'incomplete', user: userId, transactions },
|
||||
{ promptTokens, completionTokens },
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ const mockRecordCollectedUsage = jest
|
|||
|
||||
const mockGetMultiplier = jest.fn().mockReturnValue(1);
|
||||
const mockGetCacheMultiplier = jest.fn().mockReturnValue(null);
|
||||
const mockGetTransactionsConfig = jest.fn().mockReturnValue({ enabled: false });
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: {
|
||||
|
|
@ -37,7 +38,9 @@ jest.mock('@librechat/api', () => ({
|
|||
abortJob: jest.fn(),
|
||||
},
|
||||
recordCollectedUsage: mockRecordCollectedUsage,
|
||||
getTransactionsConfig: (...args) => mockGetTransactionsConfig(...args),
|
||||
sanitizeMessageForTransmit: jest.fn((msg) => msg),
|
||||
buildAbortedResponseMetadata: jest.fn().mockReturnValue(null),
|
||||
}));
|
||||
|
||||
jest.mock('librechat-data-provider', () => ({
|
||||
|
|
@ -75,7 +78,9 @@ jest.mock('./abortRun', () => ({
|
|||
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { sendError } = require('~/server/middleware/error');
|
||||
const { handleAbortError, spendCollectedUsage } = require('./abortMiddleware');
|
||||
const { GenerationJobManager } = require('@librechat/api');
|
||||
const db = require('~/models');
|
||||
const { handleAbort, handleAbortError, spendCollectedUsage } = require('./abortMiddleware');
|
||||
|
||||
const buildAbortRequest = () => ({
|
||||
body: {
|
||||
|
|
@ -310,3 +315,107 @@ describe('abortMiddleware - handleAbortError', () => {
|
|||
expect(sendError).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* The transactions config is resolved from the request's app config and must reach
|
||||
* every write path in this file. `createTransaction` reads `transactions` from the
|
||||
* caller-supplied data, so an omitted value is indistinguishable from enabled and
|
||||
* the write proceeds even when `transactions.enabled` is false.
|
||||
*/
|
||||
describe('abortMiddleware - transactions config', () => {
|
||||
const buildJobData = () => ({
|
||||
model: 'gpt-4',
|
||||
responseMessageId: 'msg-123',
|
||||
conversationId: 'convo-123',
|
||||
endpoint: 'agents',
|
||||
sender: 'AI',
|
||||
promptTokens: 25,
|
||||
userMessage: {
|
||||
messageId: 'user-msg-123',
|
||||
parentMessageId: 'parent-123',
|
||||
conversationId: 'convo-123',
|
||||
text: 'hello',
|
||||
},
|
||||
});
|
||||
|
||||
const buildReq = () => ({
|
||||
body: { abortKey: 'convo-123:1', endpoint: 'agents' },
|
||||
user: { id: 'user-123', email: 'user@example.com' },
|
||||
config: { transactions: { enabled: false } },
|
||||
});
|
||||
|
||||
const buildRes = () => ({
|
||||
headersSent: false,
|
||||
setHeader: jest.fn(),
|
||||
send: jest.fn(),
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
mockGetTransactionsConfig.mockReturnValue({ enabled: false });
|
||||
mockRecordCollectedUsage.mockResolvedValue({ input_tokens: 100, output_tokens: 50 });
|
||||
db.getConvo.mockResolvedValue({ title: 'Test Chat' });
|
||||
});
|
||||
|
||||
it('forwards transactions through spendCollectedUsage to recordCollectedUsage', async () => {
|
||||
const collectedUsage = [{ input_tokens: 100, output_tokens: 50, model: 'gpt-4' }];
|
||||
|
||||
await spendCollectedUsage({
|
||||
userId: 'user-123',
|
||||
conversationId: 'convo-123',
|
||||
collectedUsage,
|
||||
fallbackModel: 'gpt-4',
|
||||
transactions: { enabled: false },
|
||||
});
|
||||
|
||||
expect(mockRecordCollectedUsage).toHaveBeenCalledTimes(1);
|
||||
expect(mockRecordCollectedUsage).toHaveBeenCalledWith(
|
||||
expect.any(Object),
|
||||
expect.objectContaining({ context: 'abort', transactions: { enabled: false } }),
|
||||
);
|
||||
});
|
||||
|
||||
it('resolves the config from req and forwards it on the collected-usage path', async () => {
|
||||
const collectedUsage = [{ input_tokens: 100, output_tokens: 50, model: 'gpt-4' }];
|
||||
GenerationJobManager.abortJob.mockResolvedValue({
|
||||
success: true,
|
||||
jobData: buildJobData(),
|
||||
content: [],
|
||||
text: 'partial',
|
||||
collectedUsage,
|
||||
});
|
||||
|
||||
const req = buildReq();
|
||||
await handleAbort()(req, buildRes());
|
||||
|
||||
expect(logger.error).not.toHaveBeenCalled();
|
||||
expect(mockGetTransactionsConfig).toHaveBeenCalledWith(req.config);
|
||||
expect(mockRecordCollectedUsage).toHaveBeenCalledTimes(1);
|
||||
expect(mockRecordCollectedUsage).toHaveBeenCalledWith(
|
||||
expect.any(Object),
|
||||
expect.objectContaining({ context: 'abort', transactions: { enabled: false } }),
|
||||
);
|
||||
});
|
||||
|
||||
it('resolves the config from req and forwards it on the token-count fallback path', async () => {
|
||||
GenerationJobManager.abortJob.mockResolvedValue({
|
||||
success: true,
|
||||
jobData: buildJobData(),
|
||||
content: [],
|
||||
text: 'partial',
|
||||
collectedUsage: [],
|
||||
});
|
||||
|
||||
const req = buildReq();
|
||||
await handleAbort()(req, buildRes());
|
||||
|
||||
expect(logger.error).not.toHaveBeenCalled();
|
||||
expect(mockGetTransactionsConfig).toHaveBeenCalledWith(req.config);
|
||||
expect(mockRecordCollectedUsage).not.toHaveBeenCalled();
|
||||
expect(mockSpendTokens).toHaveBeenCalledTimes(1);
|
||||
expect(mockSpendTokens).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ context: 'incomplete', transactions: { enabled: false } }),
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue