From e6ff2530c344e5507179450f32d6cdf20d417a49 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 30 Jul 2026 15:47:36 -0400 Subject: [PATCH] fix: decide capability inside the atomic arm, neutralize the lost-race toast MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex round 5, both findings, both edges of the new arm design rather than its mechanism. P2, capability TOCTOU. A HITL resume on a rolling deploy rewrites preemptCapable for the SAME generation, so the handler's read could go stale between validation and the flag flip, arming a steer the live owner cannot seal. armSteer now returns armed | missing | incapable, with the owner's live capability part of the same atomic predicate as the generation fence (HGET preemptCapable inside the Lua; the flat job field, not a metadata blob — the in-memory store reads the same field). The handler's pre-check is deleted rather than kept alongside; the store predicate is the single source. New handler test rewrites the capability after queueing and expects PREEMPT_UNSUPPORTED with the item left unflagged; the Redis guards test now asserts the incapable refusal against real Redis. P2, ambiguous toast. armed:false covers injected, cancelled, re-homed, and run-over alike, so telling the user the message "already reached the agent" claimed one specific outcome. The lost-race branch now uses a neutral message (com_ui_steer_arm_lost_race) and defers to the events for what actually happened. --- .../components/Chat/Input/InFlightSteers.tsx | 6 +++++- .../Input/__tests__/InFlightSteers.test.tsx | 6 ++++-- client/src/locales/en/translation.json | 1 + .../agents/steering/__tests__/request.spec.ts | 19 +++++++++++++++++- packages/api/src/agents/steering/request.ts | 20 ++++++++++++------- packages/api/src/stream/SteeringLifecycle.ts | 9 +++++---- .../RedisJobStore.stream_integration.spec.ts | 17 +++++++++++----- .../implementations/InMemoryJobStore.ts | 18 ++++++++++++----- .../stream/implementations/RedisJobStore.ts | 20 +++++++++++++++---- .../api/src/stream/interfaces/IJobStore.ts | 17 ++++++++++++---- 10 files changed, 100 insertions(+), 33 deletions(-) diff --git a/client/src/components/Chat/Input/InFlightSteers.tsx b/client/src/components/Chat/Input/InFlightSteers.tsx index b71db39971..e087f39aad 100644 --- a/client/src/components/Chat/Input/InFlightSteers.tsx +++ b/client/src/components/Chat/Input/InFlightSteers.tsx @@ -244,11 +244,15 @@ const InFlightSteer = memo(function InFlightSteer({ markSteerPreempt(steer.steerId); return; } + /* `armed: false` is deliberately ambiguous — injected, + * cancelled, re-homed, or run over — so the message only + * says the escalation lost, and the chip defers to the + * events for whatever actually happened. */ showToast({ message: localize( response.code === 'PREEMPT_UNSUPPORTED' ? 'com_ui_steer_preempt_unsupported' - : 'com_ui_steer_already_applied', + : 'com_ui_steer_arm_lost_race', ), status: 'info', }); diff --git a/client/src/components/Chat/Input/__tests__/InFlightSteers.test.tsx b/client/src/components/Chat/Input/__tests__/InFlightSteers.test.tsx index 908ee2d0df..8b4e9c91a8 100644 --- a/client/src/components/Chat/Input/__tests__/InFlightSteers.test.tsx +++ b/client/src/components/Chat/Input/__tests__/InFlightSteers.test.tsx @@ -646,14 +646,16 @@ describe('InFlightSteers — interrupt-now escalation', () => { expect(await screen.findByText('com_ui_interrupt_steer_now')).toBeInTheDocument(); }); - it('defers to the events when the steer already left the queue', async () => { + it('stays neutral when the arm loses its race, whatever the reason', async () => { + /** `armed: false` covers injected, cancelled, re-homed, and run-over + * alike, so the toast must not claim one specific outcome. */ mockArmMutateAsync.mockResolvedValue({ armed: false }); renderSteers([{ steerId: 's1', text: 'too late', status: 'pending', createdAt: 1 }]); await clickMenuItem('com_ui_interrupt_steer_now'); expect(mockRetrySteer).not.toHaveBeenCalled(); expect(mockShowToast).toHaveBeenCalledWith( - expect.objectContaining({ message: 'com_ui_steer_already_applied' }), + expect.objectContaining({ message: 'com_ui_steer_arm_lost_race' }), ); }); diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index f4603ecb12..bd8f29355d 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -1918,6 +1918,7 @@ "com_ui_steer": "Steer", "com_ui_steer_already_applied": "That steering message already reached the agent, so it was left in the response", "com_ui_steer_arm_failed": "Couldn't arm the interrupt, so that steering message still lands at the next tool step", + "com_ui_steer_arm_lost_race": "That steering message is no longer waiting, so there was nothing to interrupt", "com_ui_steer_cancel": "Cancel steering message", "com_ui_steer_cancel_failed": "Could not cancel the steering message — it may still reach the agent", "com_ui_steer_edit_queued": "Your composer already has a draft, so that steering message was queued for after the response instead", diff --git a/packages/api/src/agents/steering/__tests__/request.spec.ts b/packages/api/src/agents/steering/__tests__/request.spec.ts index 380291f499..42e30e3341 100644 --- a/packages/api/src/agents/steering/__tests__/request.spec.ts +++ b/packages/api/src/agents/steering/__tests__/request.spec.ts @@ -823,6 +823,23 @@ describe('handleSteerArm (real in-memory job manager)', () => { expect(result.status).toBe(403); }); + it('refuses when capability was rewritten for the same generation mid-flight', async () => { + /** A HITL resume on a rolling deploy rewrites `preemptCapable` for the + * SAME createdAt, so the capability must live inside the atomic store + * predicate — a value read before the arm is not trustworthy. */ + const streamId = 'arm-capability-rewritten'; + await createCapableJob(streamId); + const posted = await handleSteerRequest(user, { conversationId: streamId, text: 'waiting' }); + await GenerationJobManager.updateMetadata(streamId, { preemptCapable: false }); + + const result = await handleSteerArm(user, { + conversationId: streamId, + steerId: posted.body.steerId as string, + }); + expect(result).toEqual({ status: 200, body: { armed: false, code: 'PREEMPT_UNSUPPORTED' } }); + expect((await GenerationJobManager.steering.peek(streamId))[0].preempt).toBeUndefined(); + }); + it("never arms another generation's steer", async () => { const streamId = 'arm-stale-generation'; await createCapableJob(streamId); @@ -834,7 +851,7 @@ describe('handleSteerArm (real in-memory job manager)', () => { posted.body.steerId as string, (live?.createdAt as number) + 999, ); - expect(armed).toBe(false); + expect(armed).toBe('missing'); expect((await GenerationJobManager.steering.peek(streamId))[0].preempt).toBeUndefined(); }); }); diff --git a/packages/api/src/agents/steering/request.ts b/packages/api/src/agents/steering/request.ts index f9c0374a88..3af3dfca2e 100644 --- a/packages/api/src/agents/steering/request.ts +++ b/packages/api/src/agents/steering/request.ts @@ -492,15 +492,21 @@ export async function handleSteerArm( if (hasTenantMismatch(job.metadata, user)) { return { status: 403, body: { code: 'UNAUTHORIZED' } }; } - if (job.metadata?.preemptCapable !== true) { - /** Same honesty rule as the POST's echo: an owner that cannot seal must - * not have its steer relabelled "interrupting". The steer stays queued - * for the next tool boundary, which is the documented degradation. */ + + /** + * Capability is decided INSIDE the atomic store predicate, not from the job + * read above: a HITL resume on a rolling deploy rewrites `preemptCapable` + * for the SAME generation, so a value read here can be stale by the time + * the flag flips. An incapable owner answers `PREEMPT_UNSUPPORTED` (same + * honesty rule as the POST's echo — the chip must not read "interrupting" + * for a run that can only inject at a tool boundary), with the item left + * unflagged and still queued. + */ + const outcome = await GenerationJobManager.steering.arm(streamId, body.steerId, job.createdAt); + if (outcome === 'incapable') { return { status: 200, body: { armed: false, code: 'PREEMPT_UNSUPPORTED' } }; } - - const armed = await GenerationJobManager.steering.arm(streamId, body.steerId, job.createdAt); - if (!armed) { + if (outcome !== 'armed') { return { status: 200, body: { armed: false } }; } /** NOT awaited, exactly like the POST: the durable flag is already the diff --git a/packages/api/src/stream/SteeringLifecycle.ts b/packages/api/src/stream/SteeringLifecycle.ts index 0ce530ab08..74d3b7a8d2 100644 --- a/packages/api/src/stream/SteeringLifecycle.ts +++ b/packages/api/src/stream/SteeringLifecycle.ts @@ -1,7 +1,7 @@ import { logger } from '@librechat/data-schemas'; import { ContentTypes, SteerEvents } from 'librechat-data-provider'; import type { TPendingSteer } from 'librechat-data-provider'; -import type { IJobStore, SteerQueueItem } from '~/stream/interfaces/IJobStore'; +import type { IJobStore, SteerArmOutcome, SteerQueueItem } from '~/stream/interfaces/IJobStore'; import type { ServerSentEvent } from '~/types'; /** Client-safe projection of a queued steer (drops the server-only userId). */ @@ -154,10 +154,11 @@ export class SteeringLifecycle { * Escalate a still-queued steer to an interrupt IN PLACE — the durable * `preempt` flag flips on the existing item, so its FIFO position survives * (the whole queue drains at the seal, in order). Races with a drain, - * cancel, or replacement run settle inside the store's atomic update: - * `false` simply means the steer is no longer this generation's to arm. + * cancel, or replacement run settle inside the store's atomic update as + * `missing`, and the owner's LIVE capability is part of the same predicate + * (`incapable`) — see {@link IJobStore.armSteer}. */ - arm(streamId: string, steerId: string, expectedCreatedAt?: number): Promise { + arm(streamId: string, steerId: string, expectedCreatedAt?: number): Promise { return this.store.armSteer(streamId, steerId, expectedCreatedAt); } diff --git a/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts b/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts index f47b8969a9..fc0e322a71 100644 --- a/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts +++ b/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts @@ -399,7 +399,9 @@ describe('RedisJobStore Integration Tests', () => { const streamId = `arm-steer-${Date.now()}`; try { - const job = await store.createJob(streamId, 'user-1', streamId); + const job = await store.createJob(streamId, 'user-1', streamId, undefined, { + preemptCapable: true, + }); await store.enqueueSteer(streamId, { steerId: 'first', text: 'earlier instruction', @@ -414,7 +416,7 @@ describe('RedisJobStore Integration Tests', () => { createdAt: 2, }); - await expect(store.armSteer(streamId, 'first', job.createdAt)).resolves.toBe(true); + await expect(store.armSteer(streamId, 'first', job.createdAt)).resolves.toBe('armed'); const queue = await store.peekSteers(streamId); expect(queue.map((item) => item.steerId)).toEqual(['first', 'second']); @@ -451,8 +453,13 @@ describe('RedisJobStore Integration Tests', () => { createdAt: 1, }); - await expect(store.armSteer(streamId, 'absent', job.createdAt)).resolves.toBe(false); - await expect(store.armSteer(streamId, 'kept', job.createdAt + 999)).resolves.toBe(false); + await expect(store.armSteer(streamId, 'absent', job.createdAt)).resolves.toBe('missing'); + await expect(store.armSteer(streamId, 'kept', job.createdAt + 999)).resolves.toBe( + 'missing', + ); + /** Live-capability predicate: the job above carries no preemptCapable, + * so an otherwise-valid arm answers `incapable` and leaves the item. */ + await expect(store.armSteer(streamId, 'kept', job.createdAt)).resolves.toBe('incapable'); expect((await store.peekSteers(streamId))[0].preempt).toBeUndefined(); /** `enqueueSteer` refuses once closed, so plant a raw item directly to @@ -467,7 +474,7 @@ describe('RedisJobStore Integration Tests', () => { createdAt: 1, }), ); - await expect(store.armSteer(streamId, 'kept', job.createdAt)).resolves.toBe(false); + await expect(store.armSteer(streamId, 'kept', job.createdAt)).resolves.toBe('missing'); } finally { await store.destroy(); } diff --git a/packages/api/src/stream/implementations/InMemoryJobStore.ts b/packages/api/src/stream/implementations/InMemoryJobStore.ts index 459a2f57e2..9eda0217f2 100644 --- a/packages/api/src/stream/implementations/InMemoryJobStore.ts +++ b/packages/api/src/stream/implementations/InMemoryJobStore.ts @@ -3,6 +3,7 @@ import type { StandardGraph } from '@librechat/agents'; import type { Agents } from 'librechat-data-provider'; import type { SerializableJobData, + SteerArmOutcome, SteerQueueItem, UsageMetadata, IJobStore, @@ -661,20 +662,27 @@ export class InMemoryJobStore implements IJobStore { return true; } - async armSteer(streamId: string, steerId: string, expectedCreatedAt?: number): Promise { + async armSteer( + streamId: string, + steerId: string, + expectedCreatedAt?: number, + ): Promise { const job = this.jobs.get(streamId); if (!job || this.closedSteerQueues.has(streamId)) { - return false; + return 'missing'; } if (expectedCreatedAt != null && job.createdAt !== expectedCreatedAt) { - return false; + return 'missing'; } const item = this.steerQueues.get(streamId)?.find((entry) => entry.steerId === steerId); if (item == null) { - return false; + return 'missing'; + } + if (job.preemptCapable !== true) { + return 'incapable'; } item.preempt = true; - return true; + return 'armed'; } async parkSteers(streamId: string, payload: string, expectedCreatedAt?: number): Promise { diff --git a/packages/api/src/stream/implementations/RedisJobStore.ts b/packages/api/src/stream/implementations/RedisJobStore.ts index 0b76894128..a436422b54 100644 --- a/packages/api/src/stream/implementations/RedisJobStore.ts +++ b/packages/api/src/stream/implementations/RedisJobStore.ts @@ -13,6 +13,7 @@ import type { JobStatusTransition, IdempotencyClaimValue, IdempotencyClaimResult, + SteerArmOutcome, } from '~/stream/interfaces/IJobStore'; import { STEER_ENQUEUE_NOT_RUNNING, @@ -395,11 +396,14 @@ const STEER_REMOVE_LUA = * set `preempt`, and LSET it back at its index, so its FIFO position is * untouched (the entire queue drains at the seal, in order). Guarded like * {@link STEER_ENQUEUE_LUA}: a closed queue or a generation mismatch refuses, - * so a stale request can never arm a replacement run's steer. + * so a stale request can never arm a replacement run's steer. The owner's + * LIVE `preemptCapable` is part of the same atomic predicate — a HITL resume + * on a rolling deploy rewrites it for the SAME generation, so a value the + * caller read earlier is not trustworthy. * * KEYS: [job, steers] * ARGV: [steerIdFragment, expectedCreatedAt or ""] - * Returns: 1 armed, 0 not found / closed / fenced + * Returns: 1 armed, 0 not found / closed / fenced, -1 owner cannot seal */ const STEER_ARM_LUA = 'if ARGV[2] ~= "" and redis.call("HGET", KEYS[1], "createdAt") ~= ARGV[2] then return 0 end ' + @@ -407,6 +411,7 @@ const STEER_ARM_LUA = 'local items = redis.call("LRANGE", KEYS[2], 0, -1) ' + 'for i = 1, #items do ' + 'if string.find(items[i], ARGV[1], 1, true) then ' + + 'if redis.call("HGET", KEYS[1], "preemptCapable") ~= "1" then return -1 end ' + 'local decoded, item = pcall(cjson.decode, items[i]) ' + 'if not decoded then return 0 end ' + 'item.preempt = true ' + @@ -1905,7 +1910,11 @@ export class RedisJobStore implements IJobStore { return removed === 1; } - async armSteer(streamId: string, steerId: string, expectedCreatedAt?: number): Promise { + async armSteer( + streamId: string, + steerId: string, + expectedCreatedAt?: number, + ): Promise { const armed = (await this.redis.eval( STEER_ARM_LUA, 2, @@ -1914,7 +1923,10 @@ export class RedisJobStore implements IJobStore { `"steerId":"${steerId}"`, expectedCreatedAt != null ? String(expectedCreatedAt) : '', )) as number; - return armed === 1; + if (armed === 1) { + return 'armed'; + } + return armed === -1 ? 'incapable' : 'missing'; } async parkSteers(streamId: string, payload: string, expectedCreatedAt?: number): Promise { diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index 14a82c5543..c7e8f8ec12 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -215,6 +215,11 @@ export interface PreemptMessage { steerIds: string[]; } +/** {@link IJobStore.armSteer}: `armed` flipped the flag in place; `missing` + * covers every left-the-queue interleaving (drained, cancelled, closed, + * replaced generation); `incapable` means the live owner cannot seal. */ +export type SteerArmOutcome = 'armed' | 'missing' | 'incapable'; + /** Maximum steers a single run can have queued at once. */ export const STEER_QUEUE_MAX_DEPTH = 10; @@ -687,11 +692,15 @@ export interface IJobStore { * Atomically set `preempt: true` on ONE queued steer IN PLACE, preserving * its FIFO position (the user escalated a waiting steer to an interrupt; * the whole queue drains at the seal, so its order must not change). - * Guarded like {@link enqueueSteer}: refuses when the queue is closed or, - * with `expectedCreatedAt`, when the stream belongs to another generation. - * False when the steer is no longer queued — drained, cancelled, or fenced. + * Guarded like {@link enqueueSteer}: `missing` when the steer is no longer + * queued, the queue is closed, or (with `expectedCreatedAt`) the stream + * belongs to another generation. The owner's LIVE `preemptCapable` is part + * of the same atomic predicate — a HITL resume on a rolling deploy can + * rewrite it for the SAME generation, so a value read before the call is + * not trustworthy — and an incapable owner answers `incapable` with the + * item left unflagged. */ - armSteer(streamId: string, steerId: string, expectedCreatedAt?: number): Promise; + armSteer(streamId: string, steerId: string, expectedCreatedAt?: number): Promise; /** * Persist terminally-drained steers under their OWN bounded-TTL key so a