CERBERUS/LibreChat
1
0
Fork 0
mirror of https://github.com/danny-avila/LibreChat.git synced 2026-09-29 04:12:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 9

🧠 feat: Add Live Reasoning Labels (#14893)

Browse source 
* feat: add live reasoning labels

* fix: Stabilize reasoning label checks

* fix: Address reasoning label review findings

* chore: Bump Agents SDK for reasoning labels

* fix: Reset reused reasoning step evidence

* fix: Reconcile cleared reasoning labels

* fix: Fence reasoning label resets

* fix: Reset reasoning ownership before gap labels

* fix: Preserve THINK type through label reset

* test: Expect run-global reasoning revision
This commit is contained in:
Danny Avila 2026-08-16 18:11:55 -04:00 • committed by GitHub
parent 3bd2358805
commit 7d850c308a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
39 changed files with 3822 additions and 75 deletions
Download patch file Download diff file Expand all files Collapse all files

239
api/server/controllers/agents/client.js
View file

@ -58,13 +58,18 @@ const {
stampSteerPartMedia,
createActivityLabelWiring,
createActivityPhaseWiring,
createReasoningLabelHostWiring,
generateReasoningLabelRevision,
getLabelUsageSequenceSeed,
createAssistantPhaseStampingHandlers,
resolveActivityConfig,
resolveActivityPhaseConfig,
resolveReasoningLabelConfig,
getCustomEndpointConfig,
mapCollectedMetadataToUsage,
resolveActivityLabelModel,
resolveActivityPhaseLabelModel,
resolveReasoningLabelModel,
traceIdForMessage,
settlePendingLabelFills,
stripActivityLabelParts,
@ -450,6 +455,35 @@ class AgentClient extends BaseClient {
return this.activityPhaseLabelLLMPromise;
}
/** Reasoning-label resolution is independently configurable and memoized per response. */
async resolveReasoningLabelLLM() {
this.reasoningLabelLLMPromise =
this.reasoningLabelLLMPromise ??
resolveReasoningLabelModel({
req: this.options.req,
agent: this.options.agent,
publicEndpoint: this.options.endpoint,
ids: {
messageId: this.responseMessageId,
conversationId: this.conversationId,
parentMessageId: this.parentMessageId,
},
db: { getUserKey: db.getUserKey, getUserKeyValues: db.getUserKeyValues },
}).catch((error) => {
this.reasoningLabelLLMPromise = null;
throw error;
});
return this.reasoningLabelLLMPromise;
}
/** Seeds the shared negative usage sequence from durable label-call state. */
seedActivityLabelUsageSequence() {
this.activityLabelUsageSeq = getLabelUsageSequenceSeed(
this.contentParts ?? [],
this.activityLabelUsageSeq,
);
}
/**
* Bills the label call and folds its usage into the response rollup with
* an `activity-label` tag (subagent precedent) so `metadata.usage` and the
@ -857,6 +891,61 @@ class AgentClient extends BaseClient {
};
}
/** SDK bridge for one revision of a live reasoning-step title. */
async generateReasoningLabelViaRun({
visibleReasoning,
reasoningStepId,
revision,
status,
previousLabel,
agentId,
charLimit,
prompt,
signal,
}) {
return generateReasoningLabelRevision({
payload: {
visibleReasoning,
reasoningStepId,
revision,
status,
...(previousLabel != null && { previousLabel }),
...(agentId != null && { agentId }),
charLimit,
...(prompt != null && { prompt }),
signal,
},
run: this.run,
resolveModel: () => this.resolveReasoningLabelLLM(),
sourceRunId: this.responseMessageId,
sourceTraceId: traceIdForMessage(this.responseMessageId),
responseId: this.responseMessageId,
sessionId: this.conversationId,
userId: this.user ?? this.options.req?.user?.id,
parentMessageId: this.parentMessageId,
recordUsage: ({
collectedMetadata,
model,
endpointTokenConfig,
sameEndpoint,
provider,
promptText,
completionText,
}) =>
this.recordActivityLabelUsage(
collectedMetadata,
model,
endpointTokenConfig,
sameEndpoint,
undefined,
provider,
() => ({ promptText, completionText }),
'reasoning-label',
),
onError: (error) => logger.warn('[AgentClient] Reasoning label generation failed', error),
});
}
/** Bounded settle for in-flight label fills before finalization. On
* timeout the label scope is closed and its abort controller fired, so a
* straggler cannot mutate the saved response or emit into a dead job. */
@ -871,20 +960,34 @@ class AgentClient extends BaseClient {
scope.detach?.();
}
};
const pending = this.pendingActivityLabelFills;
if (!pending || pending.length === 0) {
detachScopeListeners();
return;
}
this.pendingActivityLabelFills = [];
await settlePendingLabelFills(pending, timeoutMs, () => {
const closeScopes = () => {
/** Close EVERY generation's scope: a pre-pause wiring's straggler must
* stay closed even though a resume built a newer one. */
for (const scope of this.activityLabelScopes ?? []) {
scope.closed = true;
scope.abort.abort();
}
});
};
const deadline = Date.now() + timeoutMs;
while ((this.pendingActivityLabelFills?.length ?? 0) > 0) {
const pending = this.pendingActivityLabelFills;
this.pendingActivityLabelFills = [];
const remainingMs = Math.max(0, deadline - Date.now());
let timedOut = remainingMs === 0;
if (!timedOut) {
await settlePendingLabelFills(pending, remainingMs, () => {
timedOut = true;
closeScopes();
});
}
if (timedOut) {
closeScopes();
break;
}
/** A reasoning revision can synchronously enqueue its trailing final
* revision from the settled task's `finally`; drain it under the same
* deadline before any content reshaping can invalidate its index. */
}
detachScopeListeners();
}
@ -975,9 +1078,7 @@ class AgentClient extends BaseClient {
* the client's `runId:seq` deduper would discard the post-approval
* label's usage as already counted. Each label generation is a single
* non-streaming invoke, so one existing label part == one consumed seq. */
this.activityLabelUsageSeq =
this.activityLabelUsageSeq ??
(this.contentParts ?? []).filter((part) => part?.type === ContentTypes.ACTIVITY_LABEL).length;
this.seedActivityLabelUsageSequence();
this.activityLabelAbort = labelScope.abort;
/** An abort CLOSES the scope, not just cancels the call. The rejected
* generation still runs its catch and calls `fill(null)`; with the scope
@ -1124,9 +1225,7 @@ class AgentClient extends BaseClient {
scope.detach = () => abortSignal.removeEventListener('abort', closeOnAbort);
}
}
this.activityLabelUsageSeq =
this.activityLabelUsageSeq ??
(this.contentParts ?? []).filter((part) => part?.type === ContentTypes.ACTIVITY_LABEL).length;
this.seedActivityLabelUsageSequence();
const wiring = createActivityPhaseWiring({
maxPerRun: phaseConfig.maxPerRun,
@ -1164,6 +1263,72 @@ class AgentClient extends BaseClient {
return wiring;
}
/** Builds the independently opt-in live reasoning-label controller. */
buildReasoningLabelWiring(streamId, abortSignal, seedFromContent = false) {
if (!streamId || typeof Run?.prototype?.generateReasoningLabel !== 'function') {
return undefined;
}
const agentEndpoint = this.options.agent?.endpoint ?? '';
const appConfig = this.options.req?.config;
let customEndpointConfig;
try {
customEndpointConfig = getCustomEndpointConfig({ endpoint: agentEndpoint, appConfig });
} catch {
customEndpointConfig = undefined;
}
const config = resolveReasoningLabelConfig(
appConfig,
agentEndpoint,
customEndpointConfig,
this.options.endpoint,
);
if (!config.enabled) {
return undefined;
}
const shouldMarkResumable = this.activityLabelsMarkedPromise == null;
const { wiring, scope, markedPromise } = createReasoningLabelHostWiring({
config,
seedFromContent,
abortSignal,
...(shouldMarkResumable && {
markResumable: () => GenerationJobManager.markActivityLabels(streamId, this.jobCreatedAt),
onMarkFailure: () =>
logger.warn(
`[AgentClient] Could not flag reasoning labels for ${streamId}; an update resolving during a resume gap may not be reconciled.`,
),
}),
getContentParts: () => this.contentParts,
getStepIndex: (stepId) => this.stepMap?.get(stepId)?.index,
emitEvent: (event, data) =>
GenerationJobManager.emitChunk(
streamId,
{
event,
data: {
...data,
responseMessageId: this.responseMessageId,
conversationId: this.conversationId,
},
},
{ durable: true, expectedCreatedAt: this.jobCreatedAt },
),
trackPendingFill: (fillDone) => {
this.pendingActivityLabelFills = this.pendingActivityLabelFills ?? [];
this.pendingActivityLabelFills.push(fillDone);
},
generateLabel: (payload) => this.generateReasoningLabelViaRun(payload),
});
if (markedPromise != null) {
this.activityLabelsMarkedPromise = markedPromise;
}
this.activityLabelScopes = this.activityLabelScopes ?? [];
this.activityLabelScopes.push(scope);
this.seedActivityLabelUsageSequence();
this.reasoningLabelWiring = wiring;
return wiring;
}
/**
* `AgentClient` is not opinionated about vision requests, so we don't do anything here
* @param {MongoFile[]} attachments
@ -2875,10 +3040,14 @@ class AgentClient extends BaseClient {
const activityLabel = this.buildActivityLabelWiring(streamId, abortController.signal);
const activityPhase = this.buildActivityPhaseWiring(streamId, abortController.signal);
const reasoningLabel = this.buildReasoningLabelWiring(streamId, abortController.signal);
const offsetHandlers = createSteerIndexOffsetHandlers(
this.options.eventHandlers,
this.steerOffsetState,
);
const activityHandlers =
activityPhase?.handlers(offsetHandlers) ??
(activityLabel ? createAssistantPhaseStampingHandlers(offsetHandlers) : offsetHandlers);
const createRunPromise = createRun({
agents,
messages,
@ -2902,9 +3071,7 @@ class AgentClient extends BaseClient {
signal: abortController.signal,
/** The phase wrapper stays outermost: it claims and offsets the
* parent slot before the text step reaches the normal handlers. */
customHandlers:
activityPhase?.handlers(offsetHandlers) ??
(activityLabel ? createAssistantPhaseStampingHandlers(offsetHandlers) : offsetHandlers),
customHandlers: reasoningLabel?.handlers(activityHandlers) ?? activityHandlers,
requestBody: config.configurable.requestBody,
user: createSafeUser(this.options.req?.user),
tenantId: this.options.req?.user?.tenantId,
@ -2967,11 +3134,15 @@ class AgentClient extends BaseClient {
if (this.activityLabelsMarkedPromise != null) {
await this.activityLabelsMarkedPromise;
}
await run.processStream({ messages }, config, {
callbacks: {
[Callback.TOOL_ERROR]: logToolError,
},
});
try {
await run.processStream({ messages }, config, {
callbacks: {
[Callback.TOOL_ERROR]: logToolError,
},
});
} finally {
reasoningLabel?.complete();
}
this.completeActivityPhase(run, activityPhase);
// HITL: if the run paused for tool approval, mark the job
@ -3252,6 +3423,7 @@ class AgentClient extends BaseClient {
abortController.signal,
activityPhaseSnapshot,
);
const reasoningLabel = this.buildReasoningLabelWiring(streamId, abortController.signal, true);
const offsetHandlers = createSteerIndexOffsetHandlers(
createContentIndexOffsetHandlers(
this.options.eventHandlers,
@ -3259,6 +3431,9 @@ class AgentClient extends BaseClient {
),
this.steerOffsetState,
);
const activityHandlers =
activityPhase?.handlers(offsetHandlers) ??
(activityLabel ? createAssistantPhaseStampingHandlers(offsetHandlers) : offsetHandlers);
run = await createRun({
agents,
// State (messages, tool calls) is rehydrated from the checkpoint by
@ -3291,9 +3466,7 @@ class AgentClient extends BaseClient {
// type mismatch, is silently dropped against) the pre-pause content. The
// steer wrapper composes on top: resumed indices shift by seed + any
// steer parts spliced in while the resumed segment streams.
customHandlers:
activityPhase?.handlers(offsetHandlers) ??
(activityLabel ? createAssistantPhaseStampingHandlers(offsetHandlers) : offsetHandlers),
customHandlers: reasoningLabel?.handlers(activityHandlers) ?? activityHandlers,
requestBody: config.configurable.requestBody,
user: createSafeUser(this.options.req?.user),
tenantId: this.options.req?.user?.tenantId,
@ -3341,12 +3514,16 @@ class AgentClient extends BaseClient {
if (this.activityLabelsMarkedPromise != null) {
await this.activityLabelsMarkedPromise;
}
await run.resume(
resumeValue,
config,
{ callbacks: { [Callback.TOOL_ERROR]: logToolError } },
commandOptions,
);
try {
await run.resume(
resumeValue,
config,
{ callbacks: { [Callback.TOOL_ERROR]: logToolError } },
commandOptions,
);
} finally {
reasoningLabel?.complete();
}
this.completeActivityPhase(run, activityPhase);
config.signal = null;

105
api/server/controllers/agents/client.test.js
View file

@ -59,6 +59,111 @@ jest.mock('@librechat/api', () => ({
maybePrewarmCodeSandbox: jest.fn(),
}));
describe('AgentClient - label settlement', () => {
it('drains a trailing fill enqueued by an in-flight reasoning revision', async () => {
const client = Object.create(AgentClient.prototype);
const first = deferred();
const trailing = deferred();
const scope = { closed: false, abort: new AbortController(), detach: jest.fn() };
client.activityLabelScopes = [scope];
client.pendingActivityLabelFills = [
first.promise.finally(() => {
client.pendingActivityLabelFills.push(trailing.promise);
}),
];
let settled = false;
const settlement = client.settleActivityLabels(1_000).then(() => {
settled = true;
});
first.resolve();
await new Promise((resolve) => setImmediate(resolve));
expect(settled).toBe(false);
trailing.resolve();
await settlement;
expect(scope.closed).toBe(false);
expect(scope.detach).toHaveBeenCalledTimes(1);
expect(client.pendingActivityLabelFills).toEqual([]);
});
});
describe('AgentClient - reasoning label accounting', () => {
function createReasoningLabelClient(generateReasoningLabel) {
const client = Object.create(AgentClient.prototype);
client.options = { req: { user: { id: 'user-123' } } };
client.conversationId = 'conversation-123';
client.parentMessageId = 'parent-123';
client.responseMessageId = 'response-123';
client.run = { generateReasoningLabel };
client.resolveReasoningLabelLLM = jest.fn(async () => ({
provider: Providers.OPENAI,
clientOptions: { model: 'reasoning-label-model' },
endpointTokenConfig: { input: 1, output: 2 },
sameEndpoint: false,
}));
client.recordActivityLabelUsage = jest.fn(async () => undefined);
return client;
}
it('estimates output tokens from the raw model completion before title normalization', async () => {
const rawCompletion = 'Inspecting the cache race\nThis extra explanation also consumed tokens.';
const client = createReasoningLabelClient(
jest.fn(async ({ chainOptions }) => {
const callback = chainOptions.callbacks[0];
callback.handleChatModelStart(undefined, [[{ content: 'captured SDK prompt' }]]);
callback.handleLLMEnd({
generations: [
[
{
text: rawCompletion,
message: { content: [{ type: 'text', text: rawCompletion }] },
},
],
],
});
return { label: 'Inspecting the cache race' };
}),
);
const generated = await client.generateReasoningLabelViaRun({
visibleReasoning: 'x'.repeat(500),
reasoningStepId: 'reasoning-1',
revision: 1,
status: 'streaming',
signal: new AbortController().signal,
});
await generated.collectUsage(generated.label);
const usageCall = client.recordActivityLabelUsage.mock.calls[0];
expect(usageCall[6]()).toEqual({
promptText: 'captured SDK prompt',
completionText: rawCompletion,
});
expect(usageCall[7]).toBe('reasoning-label');
});
it('falls back to the returned label when no raw completion callback is available', async () => {
const client = createReasoningLabelClient(
jest.fn(async () => ({ label: 'Inspecting the cache race' })),
);
const generated = await client.generateReasoningLabelViaRun({
visibleReasoning: 'x'.repeat(500),
reasoningStepId: 'reasoning-1',
revision: 1,
status: 'streaming',
signal: new AbortController().signal,
});
await generated.collectUsage(generated.label);
expect(client.recordActivityLabelUsage.mock.calls[0][6]()).toMatchObject({
completionText: 'Inspecting the cache race',
});
});
});
describe('AgentClient - interrupt discovery persistence', () => {
beforeEach(async () => {
await GenerationJobManager.destroy();