mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-27 04:07:05 +00:00
🛰️ feat: Route Live Subagent Controls Across Replicas (#14971)
* feat: route live subagent controls across replicas * fix: initialize task routing in cluster workers * fix: harden cross-replica task routing * fix: expire routed task owners independently * fix: close cross-replica routing edge cases * fix: bound owner refresh and close routed cancellation gaps Refresh owned task registrations in bounded parallel batches so a full heartbeat pass stays well inside the 30-second directory lease instead of serializing one Redis EVAL per registration. Route conversation-deletion cancellation through a dedicated owner-side scope operation. The owner applies the deletion predicate to its complete local task set, so a scope holding more children than the model-facing list cap no longer leaves live executors running after their parent is removed. Key a consumed claim's retained response by its operation rather than by one caller's correlation id, so a later poll recovers a terminal result whose responses were all lost. Live claim statuses stay uncached so a poll always observes the task's current state. Type the model-facing `maxLength` bounds with a narrow local string schema; the SDK's JsonSchemaType does not declare the keyword, and the runtime checks continue to enforce the same limits. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: retain claimed results apart from control replays A consumed claim is the only routed response whose loss destroys data, so it no longer shares one bounded cache with control replays that unrelated command traffic can evict. Claims are retained under their own budget, and the requester acknowledges a result it received so the owner releases the copy immediately instead of holding it for the full replay window. Resolve the post-delete cancellation pass from durable leases. The deleted conversations cannot be read back, so re-reading each one only scaled the cascade while probing the owner directory once per removed id; one lease read now resolves every live child address instead. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: never consume a result the owner cannot replay Retention for consumed claims is bounded, so a burst of undelivered results could evict an earlier one and lose it for good. The owner now admits a claim only while it can retain a worst-case result, and refuses the routed claim otherwise instead of consuming it, leaving the result on the task for a later poll. Retained claims are never displaced; control replays keep evicting. Key a control replay by the command itself rather than by one caller's correlation id. The transport's own retry reuses a single envelope, but a caller that saw the owner as unavailable reissues the command under a new id, which steered, queued, or interrupted the child a second time. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: own a claimed result until it is acknowledged A consumed terminal result is task-owned state, not a cache entry. It now carries no expiry at all: the owner holds it until a caller acknowledges receipt, and only then is it released. Retention stays bounded by the existing admission gate, which refuses a claim the owner could not keep rather than consuming a result it might drop. Identify a control by the caller's invocation instead of by its content. The tool mints one id per invocation and routing carries it, so a routed retransmission of that invocation replays the owner's result while two deliberate identical commands arrive under distinct ids and both apply. Content-derived identity could not tell those apart and would have answered the second from a stale snapshot. Wait for the dpkg frontend lock in the best-effort Playwright font step. Its timeout kills npx while the apt-get it spawned keeps the lock, which then failed the fatal Redis install and ended the MCP replica jobs before any test ran (#14983). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: treat acknowledgement as part of delivering a result Publishing an acknowledgement once and ignoring the outcome meant a result could be reported as delivered while the owner never learned it could let go, and since that retention neither expires nor evicts, enough lost acknowledgements would fill it and refuse every later remote claim. An acknowledgement is now confirmed: publishing to zero subscribers is not success, it retries inside the ordinary request window, and a claim whose acknowledgement cannot be confirmed reports the retryable unavailable path instead of handing back a result the owner still holds. A later poll recovers that result and acknowledges it, and releasing is idempotent. Owner registration also outlives the task while a result is unacknowledged, so the retained result cannot become unreachable. Take the control invocation identity from the provider's tool-call id rather than minting one per execution, so replaying the same tool call stays idempotent while two distinct calls with identical payloads both apply. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * style: sort the widened node:crypto import Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: own control invocations and cancellation plans at the task seam Applies one logical control exactly once for its owning task rather than in the transport, so a local caller and a routed caller of the same invocation agree, and reusing an invocation id for different content is refused instead of silently applied. Invocation identity now comes from the run, agent, and provider tool-call id hashed to a bounded 32 characters, so a repeated `call_0` never bleeds across tasks and no id can overrun the routed bound. Cancellation for conversation deletion is now resolved into a plan while those rows are still readable, then replayed against the owner directory after the cascade is deleted. Owner registration is awaited before any provider work, so a child that cannot be addressed never starts. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * style: separate the control invocation map from the next member Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: close subagent deletion, claim, and control invocation gaps Bulk conversation deletion now runs behind a durable owner admission fence. Draining alone could not close the race: a child admitted on another replica after the drain read its leases would start provider work against a parent about to disappear. The fence is written before any lease is read and each child revalidates it after its own lease is written, so one of the two always observes the other. It expires on its own, so a process lost mid-deletion cannot leave an account unable to run subagents. A terminal child result is no longer kept alive in the owning replica's memory until someone acknowledges it. Collection is recorded durably on the child's own message against the polling invocation, so the poll whose response was lost recovers its own result while a different invocation is told the result was already collected. Owner-side retention returns to an ordinary bounded cache that expires, which is what abandoned polls needed: they can no longer occupy claim capacity until the process restarts. The deletion drain now cancels each task under one invocation held for the whole drain, stops re-sending once the owner answers, and retries only deliveries it could not confirm. A routed control replay also validates the command fingerprint, so one invocation id carrying different content reaches the owner to be refused instead of collecting the earlier command's success. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * test: assert the drain's calls before restoring its spies Restoring a spy also clears its recorded calls, so the drain assertions ran against an emptied mock. Formats the durable claim method tests. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: close the follow-on gaps in the deletion fence and result claim The admission fence now carries an ownership token, so an overlapping deletion's fence is never lifted by the one that finishes first, and both fence writes invalidate the cached auth user document. It also covers the other bulk-delete path: `DELETE /` with no conversation filter removes every conversation, so it runs behind the same fence rather than a bare drain. The durable record now decides who holds a one-shot result. An owner replaying a retained response could hand the same terminal claim to a second invocation; that invocation is told the result was already collected, while the one that consumed it still recovers its own. A task with no durable record to arbitrate keeps whatever the owner answered. Drain cancellation treats `not_found` as unconfirmed: a missing registration while the durable lease is still live means the child may be running, so the command is retried under its invocation once the owner republishes itself. Control fingerprints are hashed, so retaining one per invocation costs a fixed few bytes instead of a bounded message. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: hold every deletion fence and keep live idempotency records An owner now holds one admission fence per concurrent bulk deletion instead of one at a time, so admission reopens only when the last deletion finishes regardless of completion order. Expired fences are pruned as new ones arrive and the set is bounded, so an abandoned fence cannot accumulate or lock an account out. A failed durable claim write is no longer read as an absent record. Handing a terminal result over without recording its claimant would let another invocation collect the same one-shot output once the database recovered, so the collection reports the retryable unavailable path and leaves the result for a later poll. Control invocation records now evict tasks the store no longer holds before live ones, over a bounded scan. Dropping a live task's record would let a caller retry apply its queue, steer, or interrupt a second time once the transport replay had also expired. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: keep the deletion fence portable and never drop a live record The admission fence is written with plain update operators again. DocumentDB rejects pipeline-form updates, and this runs before any deletion, so the pipeline form would have failed both bulk-delete endpoints outright on a supported database target. An excess deletion is now refused rather than silently displacing the oldest active fence, which would have reopened admission for a deletion still running. Expired fences are pruned before the cap is tested, so only genuinely concurrent deletions count against it. Control invocation records now sweep every settled task's entry when the window fills, and a window of entirely live records refuses the new control before touching the child instead of evicting one. Applying a command with no room to record it would let the caller's own retry apply it twice. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: hold the fence, bound recovered results, and expire stale commands The admission fence is renewed for as long as its deletion runs, so a very large account or a stalled database cannot let it lapse while conversations are still being removed. Only the deletion's own fence is renewed, and the renewal stops with the operation. Cancellation now covers every conversation the cascade removed, not only the ones a plan named: a grandchild lives in its own parent's scope, which a plan naming the deleted root never reaches. A routed request carries the deadline its caller waits for, and an owner drops one that arrives past it. A publisher disconnected mid-request queues the envelope offline and delivers it after the caller was told the owner was unavailable, which would otherwise steer a child the caller believes untouched. A result recovered from its durable child message is bounded like a routed one. The message keeps the child's untruncated output, so recovery could otherwise return far more than the routed result limit allows. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * test: size the fence window so a renewal can be observed The renewal test set a 30ms drain timeout but the five-minute grace window dominates it, so the interval was 100 seconds and no renewal could fire inside the test's deletion. The grace window is an option now, matching the store's other timings, and the test sizes the window to 90ms. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: wire the durable claim method and close the fence follow-ons The production store never received `claimSubagentTaskResult`, so every terminal result would have surfaced as unavailable once a task settled. The host wires that object from JavaScript, where the factory's parameter type checks nothing, so the factory now refuses a store missing any method it calls rather than failing at the first claim. The routing transport takes a dedicated publisher with the offline queue disabled. The shared client held commands issued during a disconnect and delivered them after the caller had given up, which the request deadline narrowed but could not close inside the clock-skew allowance. Fence renewal invalidates the cached auth document like the fence and release paths, and a renewal reporting its entry gone re-takes the fence instead of letting the deletion run on unfenced. The post-delete cancellation retries a transiently unreachable owner: the conversations are already gone, so it is the only pass that can still stop a late-admitted child. A replaced replay entry no longer leaves its bytes counted, which would have inflated the cache's total until unrelated responses were evicted. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * test: wait on observed lease renewal instead of a fixed delay The shared-lease renewal test held a 60ms lease and slept 100ms before asserting an overlapping worker was refused, so a loaded runner that starved the 10ms heartbeat past the TTL let the lease lapse and the second worker run. Spy on acquisition and renewal, then wait until a renewal succeeds past the acquired lease's own deadline — direct evidence the heartbeat carried it past expiry, with no timing assumption — and give the lease enough headroom that a stalled timer no longer decides the outcome. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix(agents): close the routing, fence, and cache gaps found in review Five separate seams, each with its own failure: `Cluster.duplicate` reads its first argument as a startup-node list and its second as the overrides, unlike `Redis.duplicate`, so the publisher's `enableOfflineQueue: false` was silently dropped under `USE_REDIS_CLUSTER` and a command issued mid-disconnect could still reach a child after its caller was told `unavailable`. Route both through `duplicateIoRedisClient`. The control window's capacity refusal ran before the store knew whether it owned the task, so unrelated local load could veto a cancellation bound for another replica. Establish that the task is local first and leave a remote one to its owner's window. `clearInterval` stops only future fence renewals. One already waiting on the database could resolve after the release, read its own lifted fence as expiry, and write a replacement that nothing remained to lift — closing subagent admission for the account until it aged out. Track the in-flight renewal, refuse overlapping passes, and await it before releasing. Every owner bounds its own task list, but the aggregation appended each batch whole, so the model-facing list grew with the number of replicas holding the scope. Cap the merged list while still reading every reply for the stale-registration sweep. The admission-fence prune commits independently of the fence that follows it, so a refused or failed push left the cached auth document describing entries the collection no longer held. Invalidate whichever way the second write goes. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix(agents): cap the merged task list the poll tool actually reads Each owner bounds its own reply and the remote aggregation bounds their sum, but `listTasks` merged that bounded remote list with however many children this replica owns and returned it whole. `check_background_task` could therefore still receive roughly twice the advertised cap. Bound the deduplicated, sorted result and export the cap so both seams share one number. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * test: admit every task the merged-list cap test starts The base store admits ten concurrent runs per scope by default, so starting 150 at once left most refused for capacity and the assertion never reached the merge it was written to check. Raise the cap for this store only; admission is a different invariant with its own tests. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix(agents): let a deletion notice its admission fence lapsing Renewal failures were logged and swallowed, so a run of rejected writes let the last confirmed `fencedUntil` pass while the deletion carried on believing admission was still closed — long enough for another replica to admit a child against conversations about to be removed. Track the deadline only a confirmed write advances, and check it after the drain, before anything is deleted: nothing has been removed at that point, so the operation fails closed and the caller retries once the fence can be held. A lapse detected after the rows are gone is logged instead, since reporting failure there would invite a retry against conversations that no longer exist. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * test: raise both concurrency caps the merged-list test trips Raising the per-scope limit left the store-wide `maxRunningTotal` at its default hundred, so fifty of the hundred and fifty starts were still refused. Verified against the base store directly this time: with only the per-scope cap raised it admits a hundred, and with both raised it admits all hundred and fifty. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix(agents): close the fence renewal gap and keep running tasks listed A renewal that started before its deadline but landed after it was still credited with extending the fence from its own start time, so a window in which admission stood open was papered over: a child could take a lease the drain had already read past and the deletion would proceed without cancelling it. The deadline now only advances when the write lands while the previous one still holds; anything later records a lapse the fence cannot be restored backwards over. The model-facing cap sorted oldest-first and sliced, which dropped the newest tasks — including children that had only just started running, and which the poll tool offers no other way to discover. Bound by status instead: running children first, then the most recent settled results. Both caps share one helper, and the routed aggregation now bounds after its loop so the choice is made across every owner's reply rather than by whichever answered first. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix(agents): finish the cap and the fence at the seams they still missed The status-aware cap only reached the requester: an owner's own reply still sliced positionally, so a replica holding more than the cap dropped its running children before the requester could bound anything. Both sides now share `boundedTaskList`. A fence that lapsed during the deletion itself was only logged. The rows are gone by then, so failing is still wrong, but the child another replica admitted while the fence was down is not: the fence is retaken and the drain repeated to cancel it. A child's lease renewal had the same retroactive hole the admission fence had — Mongo filters on the `now` captured before the call, so a write landing after the lease expired still moves the row forward, while an owner drain reading active leases in that gap saw the thread as free. The lease now carries its own deadline and a late renewal stops the executor. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * test: cover the lease lapse and the post-deletion re-drain The owner-side cap shipped with a regression test; these two did not. One drives a lease renewal that succeeds only after the lease it was extending had expired and asserts the executor stops; the other lets the fence lapse during the deletion itself and asserts a second drain runs while the request still reports success. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix(agents): close live-task lifecycle gaps * test(redis): exercise cluster node discovery * fix(test): type cluster discovery seam * fix(ci): wait for orphaned apt processes * fix(ci): reserve time for apt drain * fix(ci): skip optional fonts in MCP jobs * fix(agents): recover tasks after owner loss * fix(agents): preserve local task discovery * fix(agents): initialize fail-fast cluster publisher * style(agents): sort routing test imports --------- Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Danny Avila
GitHub
parent
5f95631283
commit
259f1e0c32
35 changed files with 5839 additions and 102 deletions
6
.github/workflows/playwright-mock.yml
vendored
6
.github/workflows/playwright-mock.yml
vendored
|
|
@ -251,12 +251,6 @@ jobs:
|
|||
continue-on-error: true
|
||||
run: timeout -k 10 90 npx playwright install ffmpeg
|
||||
|
||||
# Optional fonts only — see the note in the e2e_shards job.
|
||||
- name: Install optional Playwright font dependencies (best effort)
|
||||
timeout-minutes: 4
|
||||
continue-on-error: true
|
||||
run: .github/scripts/install-playwright-fonts.sh
|
||||
|
||||
# Redis is a hard requirement for this job, so this step stays fatal.
|
||||
- name: Install Redis runtime dependencies
|
||||
timeout-minutes: 5
|
||||
|
|
|
|||
|
|
@ -2,4 +2,5 @@
|
|||
|
||||
- **Agent run envelope**: the versioned, JSON-safe request contract created after ingress authentication and protocol validation but before agent, provider, tool, or MCP initialization. It carries only the validated protocol payload and the minimum trusted principal identifiers. The execution host rehydrates all runtime state from those identifiers.
|
||||
- **Subagent thread**: a durable, view-only child conversation owned by one parent conversation and subagent identity. A parent agent may continue it by stable `threadId`; each continuation uses a fresh execution lease restored from the canonical child transcript. It is not an ordinary human-writable chat.
|
||||
- **Live subagent task owner**: the one API process holding a detached child execution, its abort controller, and its bounded control queue. Redis may route trusted poll/control envelopes to that owner, but it does not migrate or persist the executor; Mongo persists only the logical child thread and its continuation fence.
|
||||
- **Theme definition**: a versioned, data-only description of LibreChat semantic colors and shared appearance roles, optionally specialized by light or dark mode. The theme module validates and resolves partial definitions against bundled defaults before adapters apply them. A theme definition does not contain arbitrary CSS, application behavior, or alternate feature layouts.
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ const {
|
|||
setupGracefulShutdown,
|
||||
configureMessageFilterRegexValidator,
|
||||
configureFileConfigRegexEngine,
|
||||
waitForKeyvRedisClient,
|
||||
} = require('@librechat/api');
|
||||
const { connectDb, indexSync } = require('~/db');
|
||||
const initializeOAuthReconnectManager = require('./services/initializeOAuthReconnectManager');
|
||||
|
|
@ -36,6 +37,7 @@ const createValidateImageRequest = require('./middleware/validateImageRequest');
|
|||
const { startExpiredFileSweep } = require('./services/Files/process');
|
||||
const { initializeGitHubSkillSync } = require('./services/Skills/sync');
|
||||
const { initializeAgentTriggerService } = require('./services/Agents/triggers');
|
||||
const { configureSubagentTaskRouting } = require('./services/Endpoints/agents/subagentThreadStore');
|
||||
const { jwtLogin, ldapLogin, passportLogin } = require('~/strategies');
|
||||
const { updateInterfacePermissions: updateInterfacePerms } = require('@librechat/api');
|
||||
const {
|
||||
|
|
@ -304,6 +306,9 @@ if (cluster.isMaster) {
|
|||
const startServer = async () => {
|
||||
logger.info(`Worker ${process.pid} initializing...`);
|
||||
|
||||
await waitForKeyvRedisClient();
|
||||
await configureSubagentTaskRouting();
|
||||
|
||||
if (typeof Bun !== 'undefined') {
|
||||
axios.defaults.headers.common['Accept-Encoding'] = 'gzip';
|
||||
}
|
||||
|
|
|
|||
|
|
@ -30,6 +30,16 @@ describe('Experimental server configuration', () => {
|
|||
);
|
||||
});
|
||||
|
||||
it('configures routed subagent controls before a worker accepts requests', () => {
|
||||
const redisReadyIndex = source.indexOf('await waitForKeyvRedisClient();');
|
||||
const routingIndex = source.indexOf('await configureSubagentTaskRouting();');
|
||||
const listenIndex = source.indexOf('const server = app.listen');
|
||||
|
||||
expect(redisReadyIndex).toBeGreaterThan(-1);
|
||||
expect(routingIndex).toBeGreaterThan(redisReadyIndex);
|
||||
expect(listenIndex).toBeGreaterThan(routingIndex);
|
||||
});
|
||||
|
||||
it('matches the standard server pre-authentication tenant routes', () => {
|
||||
expect(source).toContain("app.use('/oauth', preAuthTenantMiddleware, routes.oauth);");
|
||||
expect(source).toContain("app.use('/api/auth', preAuthTenantMiddleware, routes.auth);");
|
||||
|
|
|
|||
|
|
@ -62,6 +62,7 @@ const { startExpiredFileSweep } = require('./services/Files/process');
|
|||
const { checkMigrations } = require('./services/start/migration');
|
||||
const optionalJwtAuth = require('./middleware/optionalJwtAuth');
|
||||
const initializeMCPs = require('./services/initializeMCPs');
|
||||
const { configureSubagentTaskRouting } = require('./services/Endpoints/agents/subagentThreadStore');
|
||||
const configureSocialLogins = require('./socialLogins');
|
||||
const createSpaFallback = require('./utils/fallback');
|
||||
const { getAppConfig } = require('./services/Config');
|
||||
|
|
@ -124,6 +125,7 @@ const configureGenerationStreams = () => {
|
|||
|
||||
const startServer = async () => {
|
||||
await waitForKeyvRedisClient();
|
||||
await configureSubagentTaskRouting();
|
||||
const { metricsMiddleware, metricsRouter } = createMetrics();
|
||||
if (!process.env.METRICS_SECRET) {
|
||||
logger.warn('[metrics] METRICS_SECRET is not set - /metrics will return 401 for all requests');
|
||||
|
|
|
|||
|
|
@ -136,6 +136,14 @@ describe('Startup readiness wiring', () => {
|
|||
expect(streamConfigIndex).toBeLessThan(postListenMcpIndex);
|
||||
});
|
||||
|
||||
it('configures subagent task routing before the server accepts requests', () => {
|
||||
const routingIndex = source.indexOf('await configureSubagentTaskRouting();');
|
||||
const listenIndex = source.indexOf('const server = app.listen');
|
||||
|
||||
expect(routingIndex).toBeGreaterThan(-1);
|
||||
expect(listenIndex).toBeGreaterThan(routingIndex);
|
||||
});
|
||||
|
||||
it('registers generation stream cleanup with the graceful shutdown coordinator', () => {
|
||||
const shutdownRegistrationIndex = source.indexOf(
|
||||
"registerShutdownTask('generation job manager'",
|
||||
|
|
|
|||
|
|
@ -123,7 +123,20 @@ module.exports = {
|
|||
assistantEndpoint: () => ({ initializeClient: jest.fn() }),
|
||||
|
||||
subagentThreadStore: () => ({
|
||||
cancelForConversations: jest.fn(),
|
||||
cancelAndDrainForOwner: jest.fn().mockResolvedValue(undefined),
|
||||
withOwnerDeletionFence: jest.fn().mockImplementation(async (_userId, _tenantId, deletion) => {
|
||||
return deletion();
|
||||
}),
|
||||
planCancellationForConversations: jest
|
||||
.fn()
|
||||
.mockImplementation(async (userId, conversationIds, tenantId) => ({
|
||||
userId,
|
||||
tenantId,
|
||||
conversationIds: [...conversationIds],
|
||||
scopes: [],
|
||||
leases: [],
|
||||
})),
|
||||
cancelPlan: jest.fn().mockResolvedValue(0),
|
||||
cancelForOwner: jest.fn(),
|
||||
}),
|
||||
};
|
||||
|
|
|
|||
|
|
@ -65,7 +65,13 @@ describe('Convos Routes', () => {
|
|||
expect(response.status).toBe(201);
|
||||
expect(deleteAgentCheckpoints).toHaveBeenCalledTimes(1);
|
||||
expect(deleteAgentCheckpoints.mock.calls[0][0]).toEqual(conversationIds);
|
||||
expect(subagentThreadStore.cancelForOwner).toHaveBeenCalledWith('test-user-123', undefined);
|
||||
/** The deletion runs inside the owner admission fence, not around it. */
|
||||
expect(subagentThreadStore.withOwnerDeletionFence).toHaveBeenCalledTimes(1);
|
||||
const [fencedUserId, fencedTenantId] =
|
||||
subagentThreadStore.withOwnerDeletionFence.mock.calls[0];
|
||||
expect(fencedUserId).toBe('test-user-123');
|
||||
expect(fencedTenantId).toBeUndefined();
|
||||
expect(subagentThreadStore.cancelAndDrainForOwner).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should delete all conversations, tool calls, and shared links for a user', async () => {
|
||||
|
|
@ -132,6 +138,21 @@ describe('Convos Routes', () => {
|
|||
expect(logger.error).toHaveBeenCalledWith('Error clearing conversations', expect.any(Error));
|
||||
});
|
||||
|
||||
it('does not delete conversations when cross-replica task draining fails', async () => {
|
||||
/** Draining happens inside the admission fence, so its failure fails the fence. */
|
||||
subagentThreadStore.withOwnerDeletionFence.mockRejectedValueOnce(
|
||||
new Error('task owner unavailable'),
|
||||
);
|
||||
|
||||
const response = await request(app).delete('/api/convos/all');
|
||||
|
||||
expect(response.status).toBe(500);
|
||||
expect(deleteConvos).not.toHaveBeenCalled();
|
||||
expect(deleteAgentCheckpoints).not.toHaveBeenCalled();
|
||||
expect(deleteToolCalls).not.toHaveBeenCalled();
|
||||
expect(deleteAllSharedLinksWithCleanup).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should return 500 if deleteToolCalls fails', async () => {
|
||||
deleteConvos.mockResolvedValue({ deletedCount: 5 });
|
||||
deleteToolCalls.mockRejectedValue(new Error('Tool calls deletion failed'));
|
||||
|
|
@ -239,6 +260,21 @@ describe('Convos Routes', () => {
|
|||
});
|
||||
|
||||
describe('DELETE /', () => {
|
||||
it('fences the owner when DELETE / is called without a conversation filter', async () => {
|
||||
deleteConvos.mockResolvedValue({ deletedCount: 3, conversationIds: ['a', 'b', 'c'] });
|
||||
|
||||
const response = await request(app)
|
||||
.delete('/api/convos')
|
||||
.send({ arg: { thread_id: 'thread-abc' } });
|
||||
|
||||
expect(response.status).toBe(201);
|
||||
/** An empty filter deletes everything, so it takes the same admission fence. */
|
||||
expect(subagentThreadStore.withOwnerDeletionFence).toHaveBeenCalledTimes(1);
|
||||
expect(subagentThreadStore.withOwnerDeletionFence.mock.calls[0][0]).toBe('test-user-123');
|
||||
expect(subagentThreadStore.cancelAndDrainForOwner).not.toHaveBeenCalled();
|
||||
expect(deleteConvos).toHaveBeenCalledWith('test-user-123', {});
|
||||
});
|
||||
|
||||
it('cancels root and descendant leases and cleans every cascaded conversation', async () => {
|
||||
deleteConvos.mockResolvedValue({
|
||||
deletedCount: 2,
|
||||
|
|
@ -254,18 +290,22 @@ describe('Convos Routes', () => {
|
|||
});
|
||||
|
||||
expect(response.status).toBe(201);
|
||||
expect(subagentThreadStore.cancelForConversations).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
/** The plan is resolved before deletion, while those rows can still be read. */
|
||||
expect(subagentThreadStore.planCancellationForConversations).toHaveBeenCalledWith(
|
||||
'test-user-123',
|
||||
['parent-conversation'],
|
||||
undefined,
|
||||
);
|
||||
expect(subagentThreadStore.cancelForConversations).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
'test-user-123',
|
||||
['parent-conversation', 'child-conversation'],
|
||||
undefined,
|
||||
);
|
||||
expect(
|
||||
subagentThreadStore.planCancellationForConversations.mock.invocationCallOrder[0],
|
||||
).toBeLessThan(deleteConvos.mock.invocationCallOrder[0]);
|
||||
/** It is applied once before deletion and replayed after with the cascade. */
|
||||
expect(subagentThreadStore.cancelPlan).toHaveBeenCalledTimes(2);
|
||||
expect(subagentThreadStore.cancelPlan.mock.calls[0][1]).toBeUndefined();
|
||||
expect(subagentThreadStore.cancelPlan.mock.calls[1][1]).toEqual([
|
||||
'parent-conversation',
|
||||
'child-conversation',
|
||||
]);
|
||||
expect(deleteToolCalls.mock.calls.map((call) => call[1])).toEqual([
|
||||
'parent-conversation',
|
||||
'child-conversation',
|
||||
|
|
|
|||
|
|
@ -118,6 +118,26 @@ router.get('/gen_title/:conversationId', async (req, res) => {
|
|||
}
|
||||
});
|
||||
|
||||
const POST_DELETE_CANCEL_ATTEMPTS = 3;
|
||||
const POST_DELETE_CANCEL_BACKOFF_MS = 250;
|
||||
|
||||
/** Replays a cancellation plan after deletion, retrying a transiently unreachable
|
||||
* owner rather than losing the only pass that can stop a late-admitted child. */
|
||||
async function retryPostDeleteCancellation(cancellationPlan, deletedConversationIds) {
|
||||
for (let attempt = 1; attempt <= POST_DELETE_CANCEL_ATTEMPTS; attempt += 1) {
|
||||
try {
|
||||
await subagentThreadTaskStore.cancelPlan(cancellationPlan, deletedConversationIds);
|
||||
return;
|
||||
} catch (error) {
|
||||
if (attempt === POST_DELETE_CANCEL_ATTEMPTS) {
|
||||
logger.warn('Post-delete subagent cancellation failed', error);
|
||||
return;
|
||||
}
|
||||
await new Promise((resolve) => setTimeout(resolve, POST_DELETE_CANCEL_BACKOFF_MS * attempt));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
router.delete('/', configMiddleware, async (req, res) => {
|
||||
let filter = {};
|
||||
const { conversationId, source, thread_id, endpoint } = req.body?.arg ?? {};
|
||||
|
|
@ -154,17 +174,36 @@ router.delete('/', configMiddleware, async (req, res) => {
|
|||
typeof req.user.tenantId === 'string' && req.user.tenantId !== ''
|
||||
? req.user.tenantId
|
||||
: undefined;
|
||||
let cancellationPlan;
|
||||
let dbResponse;
|
||||
if (filter.conversationId) {
|
||||
subagentThreadTaskStore.cancelForConversations(
|
||||
/** Resolve the targets while the conversations still exist: the second pass
|
||||
* runs after their rows are gone and can only reach registered owners. */
|
||||
cancellationPlan = await subagentThreadTaskStore.planCancellationForConversations(
|
||||
req.user.id,
|
||||
[filter.conversationId],
|
||||
tenantId,
|
||||
);
|
||||
await subagentThreadTaskStore.cancelPlan(cancellationPlan);
|
||||
dbResponse = await db.deleteConvos(req.user.id, filter);
|
||||
} else {
|
||||
/** An empty filter deletes every conversation this owner has, so it runs behind
|
||||
* the same admission fence as `DELETE /all` rather than a bare drain. */
|
||||
dbResponse = await subagentThreadTaskStore.withOwnerDeletionFence(req.user.id, tenantId, () =>
|
||||
db.deleteConvos(req.user.id, filter),
|
||||
);
|
||||
}
|
||||
const dbResponse = await db.deleteConvos(req.user.id, filter);
|
||||
const deletedConversationIds =
|
||||
dbResponse.conversationIds ?? (filter.conversationId ? [filter.conversationId] : []);
|
||||
subagentThreadTaskStore.cancelForConversations(req.user.id, deletedConversationIds, tenantId);
|
||||
/** Root deletion closes new child admission. Replay the plan to catch a task
|
||||
* admitted after the first pass but before that fence, extended with the cascade
|
||||
* this deletion reported. */
|
||||
if (cancellationPlan != null && deletedConversationIds.length > 0) {
|
||||
/** The conversations are gone, so this pass is the only thing that can still
|
||||
* stop a child admitted after the first one. It cannot fail the request — the
|
||||
* deletion already committed — so it retries briefly before giving up. */
|
||||
await retryPostDeleteCancellation(cancellationPlan, deletedConversationIds);
|
||||
}
|
||||
// HITL: prune the deleted conversations' durable checkpoints — a paused run's
|
||||
// checkpoint would otherwise persist until the Mongo TTL. Never throws.
|
||||
await deleteAgentCheckpoints(
|
||||
|
|
@ -186,13 +225,18 @@ router.delete('/', configMiddleware, async (req, res) => {
|
|||
|
||||
router.delete('/all', configMiddleware, async (req, res) => {
|
||||
try {
|
||||
subagentThreadTaskStore.cancelForOwner(
|
||||
req.user.id,
|
||||
const tenantId =
|
||||
typeof req.user.tenantId === 'string' && req.user.tenantId !== ''
|
||||
? req.user.tenantId
|
||||
: undefined,
|
||||
: undefined;
|
||||
/** Fences new child admission for this owner, drains the live ones, and deletes
|
||||
* inside that fence: a child admitted on another replica mid-deletion would
|
||||
* otherwise keep running against conversations that no longer exist. */
|
||||
const dbResponse = await subagentThreadTaskStore.withOwnerDeletionFence(
|
||||
req.user.id,
|
||||
tenantId,
|
||||
() => db.deleteConvos(req.user.id, {}),
|
||||
);
|
||||
const dbResponse = await db.deleteConvos(req.user.id, {});
|
||||
// HITL: prune ALL the deleted conversations' durable checkpoints in one bulk pass.
|
||||
await deleteAgentCheckpoints(
|
||||
dbResponse.conversationIds,
|
||||
|
|
|
|||
|
|
@ -1210,7 +1210,7 @@ const initializeClient = async ({
|
|||
}
|
||||
|
||||
/** Build detached execution only for an attributable owner/thread. New
|
||||
* tasks still require a spawnable child, while an existing process-local
|
||||
* tasks still require a spawnable child, while an existing registered live
|
||||
* task keeps its poll/control seam after agent configuration changes. The
|
||||
* SDK receives only this trusted host scope; models can select a child
|
||||
* `threadId`, never the owner or parent-thread namespace. */
|
||||
|
|
@ -1236,9 +1236,20 @@ const initializeClient = async ({
|
|||
: {}),
|
||||
})
|
||||
: undefined;
|
||||
const hasExistingSubagentTask =
|
||||
trustedSubagentTasks != null &&
|
||||
trustedSubagentTasks.store.list(trustedSubagentTasks.scopeId).length > 0;
|
||||
let hasExistingSubagentTask = false;
|
||||
if (trustedSubagentTasks != null && !(subagentsAvailableForRun && hasSpawnableSubagent)) {
|
||||
try {
|
||||
hasExistingSubagentTask = await subagentThreadTaskStore.hasTasks(
|
||||
trustedSubagentTasks.scopeId,
|
||||
);
|
||||
} catch (error) {
|
||||
/** Keep the poll/control tool visible when the owner directory is briefly
|
||||
* unavailable. The tool then returns an honest `unavailable` status
|
||||
* instead of making a live task look nonexistent. */
|
||||
logger.warn('[initializeClient] Failed to inspect routed subagent tasks', error);
|
||||
hasExistingSubagentTask = true;
|
||||
}
|
||||
}
|
||||
const subagentTasks =
|
||||
trustedSubagentTasks != null &&
|
||||
((subagentsAvailableForRun && hasSpawnableSubagent) || hasExistingSubagentTask)
|
||||
|
|
|
|||
|
|
@ -756,19 +756,7 @@ describe('initializeClient — subagent loading', () => {
|
|||
endpointOption: makeEndpointOption(),
|
||||
});
|
||||
const existingConfig = agentClientArgs.subagentTasks;
|
||||
const listSpy = jest.spyOn(existingConfig.store, 'list').mockReturnValueOnce([
|
||||
{
|
||||
taskId: 'existing-task',
|
||||
threadId: 'existing-thread',
|
||||
subagentType: 'researcher',
|
||||
status: 'running',
|
||||
createdAt: Date.now(),
|
||||
updatedAt: Date.now(),
|
||||
resultAvailable: false,
|
||||
resultClaimed: false,
|
||||
pendingControls: 0,
|
||||
},
|
||||
]);
|
||||
const hasTasksSpy = jest.spyOn(existingConfig.store, 'hasTasks').mockResolvedValueOnce(true);
|
||||
mockInitializeAgent.mockResolvedValue(makePrimaryConfig({}));
|
||||
const changedReq = makeSubagentReq();
|
||||
changedReq.config.endpoints.agents.capabilities.push('run_in_background');
|
||||
|
|
@ -783,7 +771,7 @@ describe('initializeClient — subagent loading', () => {
|
|||
expect(agentClientArgs.subagentTasks).toEqual(existingConfig);
|
||||
expect(capturedToolExecuteOptions.subagentTasks).toEqual(existingConfig);
|
||||
expect(agentClientArgs.agent.subagents).toBeUndefined();
|
||||
listSpy.mockRestore();
|
||||
hasTasksSpy.mockRestore();
|
||||
});
|
||||
|
||||
it('disables every nested subagent path at the durable child-thread depth limit', async () => {
|
||||
|
|
|
|||
|
|
@ -1,16 +1,25 @@
|
|||
const { createSubagentThreadTaskStore } = require('@librechat/api');
|
||||
const {
|
||||
cacheConfig,
|
||||
ioredisClient,
|
||||
registerShutdownTask,
|
||||
duplicateIoRedisClient,
|
||||
createSubagentThreadTaskStore,
|
||||
RedisSubagentTaskControlTransport,
|
||||
} = require('@librechat/api');
|
||||
const db = require('~/models');
|
||||
|
||||
/** Durable logical threads use normal LibreChat conversations/messages. Live
|
||||
* controls stay process-local; Mongo fences continuation across API replicas. */
|
||||
/** Durable logical threads use normal LibreChat conversations/messages. Mongo
|
||||
* fences continuation; optional Redis routing reaches the live owning process. */
|
||||
const subagentThreadTaskStore = createSubagentThreadTaskStore(
|
||||
{
|
||||
acquireSubagentThreadLease: db.acquireSubagentThreadLease,
|
||||
claimSubagentTaskResult: db.claimSubagentTaskResult,
|
||||
countActiveSubagentThreadLeases: db.countActiveSubagentThreadLeases,
|
||||
deleteConvos: db.deleteConvos,
|
||||
deleteMessages: db.deleteMessages,
|
||||
getConvo: db.getConvo,
|
||||
getMessages: db.getMessages,
|
||||
listActiveSubagentThreadLeases: db.listActiveSubagentThreadLeases,
|
||||
releaseSubagentThreadLease: db.releaseSubagentThreadLease,
|
||||
reserveSubagentThread: db.reserveSubagentThread,
|
||||
renewSubagentThreadLease: db.renewSubagentThreadLease,
|
||||
|
|
@ -18,8 +27,49 @@ const subagentThreadTaskStore = createSubagentThreadTaskStore(
|
|||
saveMessage: db.saveMessage,
|
||||
},
|
||||
{
|
||||
isOwnerActive: db.isAgentTriggerPrincipalActive,
|
||||
isOwnerActive: db.isSubagentOwnerAdmissible,
|
||||
fenceOwnerAdmission: db.fenceSubagentAdmission,
|
||||
renewOwnerAdmission: db.renewSubagentAdmission,
|
||||
releaseOwnerAdmission: db.releaseSubagentAdmission,
|
||||
},
|
||||
);
|
||||
|
||||
let taskRoutingConfigured = false;
|
||||
|
||||
/** Starts the optional Redis owner directory before HTTP admission opens. */
|
||||
async function configureSubagentTaskRouting() {
|
||||
if (taskRoutingConfigured || !cacheConfig.USE_REDIS) {
|
||||
return;
|
||||
}
|
||||
if (ioredisClient == null || typeof ioredisClient.duplicate !== 'function') {
|
||||
throw new Error('Redis subagent task routing requires a dedicated subscriber connection.');
|
||||
}
|
||||
const subscriber = ioredisClient.duplicate();
|
||||
/** A dedicated publisher without the offline queue: the shared client would hold a
|
||||
* command issued during a disconnect and deliver it after the caller gave up, so a
|
||||
* steer the caller was told had failed could still reach the child. Failing fast
|
||||
* turns that into the honest `unavailable` the caller already handles. */
|
||||
const publisher = duplicateIoRedisClient(ioredisClient, { enableOfflineQueue: false });
|
||||
const transport = new RedisSubagentTaskControlTransport(publisher, subscriber, {
|
||||
namespace: cacheConfig.REDIS_KEY_PREFIX,
|
||||
});
|
||||
try {
|
||||
await subagentThreadTaskStore.configureTaskControlTransport(transport);
|
||||
} catch (error) {
|
||||
subscriber.disconnect();
|
||||
publisher.disconnect();
|
||||
throw error;
|
||||
}
|
||||
taskRoutingConfigured = true;
|
||||
registerShutdownTask(
|
||||
'subagent task control transport',
|
||||
async () => {
|
||||
await subagentThreadTaskStore.destroyTaskControlTransport();
|
||||
publisher.disconnect();
|
||||
},
|
||||
{ priority: 90 },
|
||||
);
|
||||
}
|
||||
|
||||
module.exports = subagentThreadTaskStore;
|
||||
module.exports.configureSubagentTaskRouting = configureSubagentTaskRouting;
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ import {
|
|||
CHECK_BACKGROUND_TASK_NAME,
|
||||
RUN_IN_BACKGROUND_ARG,
|
||||
} from './background';
|
||||
import { SubagentTaskOwnerUnavailableError } from './subagentTaskRouting';
|
||||
import { TOOL_SELECTION_WILDCARD } from './selection';
|
||||
import { toolOptionsSchema } from './validation';
|
||||
|
||||
|
|
@ -1056,8 +1057,8 @@ describe('getBackgroundCodeDelivery (singleton)', () => {
|
|||
});
|
||||
|
||||
describe('runCheckBackgroundTask (singleton)', () => {
|
||||
it('returns not_found for an unknown id', () => {
|
||||
const content = runCheckBackgroundTask({
|
||||
it('returns not_found for an unknown id', async () => {
|
||||
const content = await runCheckBackgroundTask({
|
||||
userId: 'poll_user',
|
||||
conversationId: 'poll_convo',
|
||||
args: { background_task_id: 'nope' },
|
||||
|
|
@ -1067,7 +1068,29 @@ describe('runCheckBackgroundTask (singleton)', () => {
|
|||
);
|
||||
});
|
||||
|
||||
it('returns a single task by id and lists all when omitted', () => {
|
||||
it('rejects an oversized task id before local or cross-replica lookup', async () => {
|
||||
const store = Object.assign(new InMemorySubagentTaskStore(), {
|
||||
claimTask: jest.fn(),
|
||||
controlTask: jest.fn(),
|
||||
listTasks: jest.fn(),
|
||||
});
|
||||
const content = await runCheckBackgroundTask({
|
||||
userId: 'owner',
|
||||
conversationId: 'parent-thread',
|
||||
args: { background_task_id: 'x'.repeat(257) },
|
||||
subagentTasks: { store, scopeId: 'owner:parent-thread' },
|
||||
});
|
||||
|
||||
expect(JSON.parse(content)).toEqual({
|
||||
status: 'invalid',
|
||||
message: 'A background_task_id cannot exceed 256 characters.',
|
||||
});
|
||||
expect(store.claimTask).not.toHaveBeenCalled();
|
||||
expect(store.controlTask).not.toHaveBeenCalled();
|
||||
expect(store.listTasks).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('returns a single task by id and lists all when omitted', async () => {
|
||||
const created = backgroundTaskRegistry.create({
|
||||
userId: 'poll_user',
|
||||
conversationId: 'poll_convo2',
|
||||
|
|
@ -1082,7 +1105,7 @@ describe('runCheckBackgroundTask (singleton)', () => {
|
|||
});
|
||||
|
||||
const single = JSON.parse(
|
||||
runCheckBackgroundTask({
|
||||
await runCheckBackgroundTask({
|
||||
userId: 'poll_user',
|
||||
conversationId: 'poll_convo2',
|
||||
args: { background_task_id: created.task.id },
|
||||
|
|
@ -1097,7 +1120,11 @@ describe('runCheckBackgroundTask (singleton)', () => {
|
|||
);
|
||||
|
||||
const listed = JSON.parse(
|
||||
runCheckBackgroundTask({ userId: 'poll_user', conversationId: 'poll_convo2', args: {} }),
|
||||
await runCheckBackgroundTask({
|
||||
userId: 'poll_user',
|
||||
conversationId: 'poll_convo2',
|
||||
args: {},
|
||||
}),
|
||||
);
|
||||
expect(listed.tasks).toHaveLength(1);
|
||||
expect(listed.tasks[0].background_task_id).toBe(created.task.id);
|
||||
|
|
@ -1109,7 +1136,7 @@ describe('runCheckBackgroundTask (singleton)', () => {
|
|||
|
||||
// stringified args must still resolve the specific task (with its full result)
|
||||
const singleFromString = JSON.parse(
|
||||
runCheckBackgroundTask({
|
||||
await runCheckBackgroundTask({
|
||||
userId: 'poll_user',
|
||||
conversationId: 'poll_convo2',
|
||||
args: `{"background_task_id":"${created.task.id}"}`,
|
||||
|
|
@ -1120,7 +1147,68 @@ describe('runCheckBackgroundTask (singleton)', () => {
|
|||
);
|
||||
});
|
||||
|
||||
it('retrieves a task across turns: the poll is keyed only by id, not the dispatch run/turn', () => {
|
||||
it('preserves local task lists when cross-replica subagent discovery is unavailable', async () => {
|
||||
const ordinary = backgroundTaskRegistry.create({
|
||||
userId: 'partial-list-owner',
|
||||
conversationId: 'partial-list-parent',
|
||||
toolCallId: 'ordinary-call',
|
||||
toolName: 'search_mcp_docs',
|
||||
});
|
||||
if ('atCapacity' in ordinary) {
|
||||
throw new Error('unexpected capacity');
|
||||
}
|
||||
|
||||
const store = new InMemorySubagentTaskStore();
|
||||
const started = store.start({
|
||||
scopeId: 'partial-list-owner:partial-list-parent',
|
||||
idempotencyKey: 'partial-list-run:parent-agent:subagent-call',
|
||||
parentRunId: 'partial-list-run',
|
||||
parentAgentId: 'parent-agent',
|
||||
parentToolCallId: 'subagent-call',
|
||||
input: 'Keep working locally.',
|
||||
subagentKind: 'agent',
|
||||
subagentType: 'researcher',
|
||||
run: async () => ({ content: 'local result' }),
|
||||
});
|
||||
if (!started.accepted) {
|
||||
throw new Error('Expected subagent task to start.');
|
||||
}
|
||||
await waitForSubagentTaskToSettle(
|
||||
store,
|
||||
'partial-list-owner:partial-list-parent',
|
||||
started.task.taskId,
|
||||
);
|
||||
|
||||
const routedStore = Object.assign(store, {
|
||||
claimTask: jest.fn(),
|
||||
controlTask: jest.fn(),
|
||||
listTasks: jest.fn().mockRejectedValue(new SubagentTaskOwnerUnavailableError()),
|
||||
});
|
||||
const listed = JSON.parse(
|
||||
await runCheckBackgroundTask({
|
||||
userId: 'partial-list-owner',
|
||||
conversationId: 'partial-list-parent',
|
||||
args: {},
|
||||
subagentTasks: {
|
||||
store: routedStore,
|
||||
scopeId: 'partial-list-owner:partial-list-parent',
|
||||
},
|
||||
}),
|
||||
);
|
||||
|
||||
expect(listed).toEqual(
|
||||
expect.objectContaining({
|
||||
partial: true,
|
||||
warning:
|
||||
'Cross-replica subagent tasks could not be listed: The process running this subagent task is temporarily unavailable.',
|
||||
}),
|
||||
);
|
||||
expect(
|
||||
listed.tasks.map((task: { background_task_id: string }) => task.background_task_id),
|
||||
).toEqual(expect.arrayContaining([ordinary.task.id, started.task.taskId]));
|
||||
});
|
||||
|
||||
it('retrieves a task across turns: the poll is keyed only by id, not the dispatch run/turn', async () => {
|
||||
// Turn 1 dispatches under run-turn-1 and the result lands after the turn.
|
||||
const dispatched = backgroundTaskRegistry.create({
|
||||
userId: 'poll_user',
|
||||
|
|
@ -1139,7 +1227,7 @@ describe('runCheckBackgroundTask (singleton)', () => {
|
|||
|
||||
// Turn 2 (a later run) polls with just the id; get/list carry no run/turn scope.
|
||||
const polled = JSON.parse(
|
||||
runCheckBackgroundTask({
|
||||
await runCheckBackgroundTask({
|
||||
userId: 'poll_user',
|
||||
conversationId: 'poll_xturn',
|
||||
args: { background_task_id: dispatched.task.id },
|
||||
|
|
@ -1174,7 +1262,7 @@ describe('runCheckBackgroundTask (singleton)', () => {
|
|||
await waitForSubagentTaskToSettle(store, subagentTasks.scopeId, started.task.taskId);
|
||||
|
||||
const first = JSON.parse(
|
||||
runCheckBackgroundTask({
|
||||
await runCheckBackgroundTask({
|
||||
userId: 'owner',
|
||||
conversationId: 'parent-thread',
|
||||
args: { background_task_id: started.task.taskId },
|
||||
|
|
@ -1192,7 +1280,7 @@ describe('runCheckBackgroundTask (singleton)', () => {
|
|||
);
|
||||
|
||||
const second = JSON.parse(
|
||||
runCheckBackgroundTask({
|
||||
await runCheckBackgroundTask({
|
||||
userId: 'owner',
|
||||
conversationId: 'parent-thread',
|
||||
args: { background_task_id: started.task.taskId },
|
||||
|
|
@ -1227,7 +1315,7 @@ describe('runCheckBackgroundTask (singleton)', () => {
|
|||
await Promise.resolve();
|
||||
|
||||
const queued = JSON.parse(
|
||||
runCheckBackgroundTask({
|
||||
await runCheckBackgroundTask({
|
||||
userId: 'owner',
|
||||
conversationId: 'parent-thread',
|
||||
args: {
|
||||
|
|
@ -1243,7 +1331,7 @@ describe('runCheckBackgroundTask (singleton)', () => {
|
|||
);
|
||||
|
||||
const cancelledMessage = JSON.parse(
|
||||
runCheckBackgroundTask({
|
||||
await runCheckBackgroundTask({
|
||||
userId: 'owner',
|
||||
conversationId: 'parent-thread',
|
||||
args: {
|
||||
|
|
@ -1257,7 +1345,7 @@ describe('runCheckBackgroundTask (singleton)', () => {
|
|||
expect(cancelledMessage.status).toBe('accepted');
|
||||
|
||||
const cancelledTask = JSON.parse(
|
||||
runCheckBackgroundTask({
|
||||
await runCheckBackgroundTask({
|
||||
userId: 'owner',
|
||||
conversationId: 'parent-thread',
|
||||
args: { background_task_id: started.task.taskId, action: 'cancel' },
|
||||
|
|
@ -1267,6 +1355,102 @@ describe('runCheckBackgroundTask (singleton)', () => {
|
|||
expect(cancelledTask.status).toBe('cancelled');
|
||||
finish({ content: 'late result' });
|
||||
});
|
||||
|
||||
it('derives a bounded control invocation identity from the tool call', async () => {
|
||||
const controlTask = jest.fn().mockResolvedValue({
|
||||
status: 'not_running',
|
||||
task: {
|
||||
taskId: 'remote-task',
|
||||
subagentType: 'researcher',
|
||||
status: 'completed',
|
||||
createdAt: 1,
|
||||
updatedAt: 2,
|
||||
resultAvailable: false,
|
||||
resultClaimed: true,
|
||||
pendingControls: 0,
|
||||
},
|
||||
});
|
||||
const store = Object.assign(new InMemorySubagentTaskStore(), {
|
||||
claimTask: jest.fn(),
|
||||
controlTask,
|
||||
listTasks: jest.fn(),
|
||||
});
|
||||
const control = (toolCallId: string | undefined) =>
|
||||
runCheckBackgroundTask({
|
||||
userId: 'owner',
|
||||
conversationId: 'parent-thread',
|
||||
args: {
|
||||
background_task_id: 'remote-task',
|
||||
action: 'queue',
|
||||
message: 'Check one more source.',
|
||||
},
|
||||
toolCallId,
|
||||
subagentTasks: { store, scopeId: 'owner:parent-thread' },
|
||||
});
|
||||
|
||||
/** Replaying one tool call keeps its identity, so routing can replay the result. */
|
||||
await control('call_abc');
|
||||
await control('call_abc');
|
||||
const [firstInvocation, replayedInvocation] = controlTask.mock.calls.map((call) => call[3]);
|
||||
expect(firstInvocation).toBe(replayedInvocation);
|
||||
expect(firstInvocation).toHaveLength(32);
|
||||
|
||||
/** A separate tool call is a separate command even with an identical payload. */
|
||||
await control('call_def');
|
||||
expect(controlTask.mock.calls[2][3]).not.toBe(firstInvocation);
|
||||
|
||||
/** The same provider id in another run or agent is a different command. */
|
||||
await runCheckBackgroundTask({
|
||||
userId: 'owner',
|
||||
conversationId: 'parent-thread',
|
||||
args: {
|
||||
background_task_id: 'remote-task',
|
||||
action: 'queue',
|
||||
message: 'Check one more source.',
|
||||
},
|
||||
toolCallId: 'call_abc',
|
||||
runId: 'run-2:0',
|
||||
subagentTasks: { store, scopeId: 'owner:parent-thread' },
|
||||
});
|
||||
expect(controlTask.mock.calls[3][3]).not.toBe(firstInvocation);
|
||||
|
||||
/** A provider id far past the protocol bound still routes as a bounded identity. */
|
||||
const longToolCallId = `call_${'x'.repeat(200)}`;
|
||||
await control(longToolCallId);
|
||||
await control(longToolCallId);
|
||||
const [longInvocation, replayedLongInvocation] = controlTask.mock.calls
|
||||
.slice(4)
|
||||
.map((call) => call[3]);
|
||||
expect(longInvocation).toHaveLength(32);
|
||||
expect(replayedLongInvocation).toBe(longInvocation);
|
||||
|
||||
/** Without a tool-call id each invocation stays distinct rather than colliding. */
|
||||
await control(undefined);
|
||||
await control(undefined);
|
||||
const [fallback, otherFallback] = controlTask.mock.calls.slice(6).map((call) => call[3]);
|
||||
expect(fallback).not.toBe(otherFallback);
|
||||
expect(fallback.length).toBeLessThanOrEqual(128);
|
||||
});
|
||||
|
||||
it('reports an unreachable remote subagent owner without pretending the task is missing', async () => {
|
||||
const store = Object.assign(new InMemorySubagentTaskStore(), {
|
||||
claimTask: jest.fn().mockRejectedValue(new SubagentTaskOwnerUnavailableError()),
|
||||
controlTask: jest.fn().mockRejectedValue(new SubagentTaskOwnerUnavailableError()),
|
||||
listTasks: jest.fn().mockRejectedValue(new SubagentTaskOwnerUnavailableError()),
|
||||
});
|
||||
const content = await runCheckBackgroundTask({
|
||||
userId: 'owner',
|
||||
conversationId: 'parent-thread',
|
||||
args: { background_task_id: 'remote-task' },
|
||||
subagentTasks: { store, scopeId: 'owner:parent-thread' },
|
||||
});
|
||||
|
||||
expect(JSON.parse(content)).toEqual({
|
||||
status: 'unavailable',
|
||||
background_task_id: 'remote-task',
|
||||
message: 'The process running this subagent task is temporarily unavailable.',
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('stripBackgroundFromToolRegistry', () => {
|
||||
|
|
|
|||
|
|
@ -17,7 +17,10 @@
|
|||
* are lost on restart and are not shared across replicas (durable follow-up),
|
||||
* and ephemeral request-scoped MCP tools (runtime `{{LIBRECHAT_BODY_*}}`
|
||||
* placeholders) are never backgrounded — their connection is torn down at
|
||||
* request end, so the executor runs them in the foreground instead.
|
||||
* request end, so the executor runs them in the foreground instead. Detached
|
||||
* subagents use the separate host task store; Redis-backed hosts may route
|
||||
* their poll/control operations to the owning process without moving the live
|
||||
* executor or making ordinary background tool results durable.
|
||||
*
|
||||
* Opt-in mirrors `deferred_tools`: an admin capability
|
||||
* (`AgentCapabilities.run_in_background`) gates the feature, and a per-tool
|
||||
|
|
@ -30,8 +33,8 @@
|
|||
* @module packages/api/src/agents/background
|
||||
*/
|
||||
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import { logger } from '@librechat/data-schemas';
|
||||
import { createHash, randomUUID } from 'node:crypto';
|
||||
import { Constants as AgentConstants } from '@librechat/agents';
|
||||
import { Tools, Constants, imageGenTools } from 'librechat-data-provider';
|
||||
import type {
|
||||
|
|
@ -43,6 +46,7 @@ import type {
|
|||
SubagentTaskSnapshot,
|
||||
SubagentTaskControlCommand,
|
||||
SubagentTaskControlResult,
|
||||
SubagentTaskStore,
|
||||
} from '@librechat/agents';
|
||||
import type { AgentToolOptions } from 'librechat-data-provider';
|
||||
import type { CapabilityToolNames } from './selection';
|
||||
|
|
@ -52,6 +56,7 @@ import {
|
|||
warnUnmatchedSelectionNames,
|
||||
synthesizeSelectionToolOptions,
|
||||
} from './selection';
|
||||
import { SubagentTaskOwnerUnavailableError } from './subagentTaskRouting';
|
||||
import { SET_MEMORY_TOOL_NAME, DELETE_MEMORY_TOOL_NAME } from './memory';
|
||||
import { ASK_USER_QUESTION_TOOL_NAME } from './hitl/askUserQuestionTool';
|
||||
import { CREATE_FILE_TOOL_NAME, EDIT_FILE_TOOL_NAME } from './tools';
|
||||
|
|
@ -62,6 +67,9 @@ export const RUN_IN_BACKGROUND_ARG = 'run_in_background';
|
|||
|
||||
/** Log prefix for selection diagnostics, phrased in the spec's own field name. */
|
||||
const BACKGROUND_SELECTION_LABEL = '[background] runInBackground';
|
||||
const MAX_BACKGROUND_TASK_ID_CHARS = 256;
|
||||
const MAX_BACKGROUND_CONTROL_ID_CHARS = 256;
|
||||
const MAX_BACKGROUND_CONTROL_MESSAGE_CHARS = 64 * 1024;
|
||||
|
||||
/**
|
||||
* `type` of the synthetic attachment emitted on a poll turn when a harvested
|
||||
|
|
@ -297,13 +305,36 @@ export function stripBackgroundFromToolRegistry(
|
|||
|
||||
const CHECK_BACKGROUND_TASK_DESCRIPTION = `Check, control, and retrieve tool or subagent tasks previously dispatched in the background (with run_in_background: true).
|
||||
|
||||
Provide a background_task_id to poll one task; omit it to list every background task in this thread. A task is only finished when its status is "completed", "error", or "cancelled" — never assume completion without polling. Results are not pushed to you; you must call this tool to collect them. Subagent tasks additionally accept steer, queue, interrupt, cancel, and cancel_message actions while running. Execution leases remain available only while requests reach the owning server process; they do not survive a restart or cross-worker routing. A completed subagent thread may be continued later through the subagent tool's durable thread id.`;
|
||||
Provide a background_task_id to poll one task; omit it to list every background task in this thread. A task is only finished when its status is "completed", "error", or "cancelled" — never assume completion without polling. Results are not pushed to you; you must call this tool to collect them. Subagent tasks additionally accept steer, queue, interrupt, cancel, and cancel_message actions while running. Live subagent controls route across API replicas but do not survive a restart of the process that owns the executor. A completed subagent thread may be continued later through the subagent tool's durable thread id.`;
|
||||
|
||||
const CHECK_BACKGROUND_TASK_PARAMETERS: JsonSchemaType = Object.freeze<JsonSchemaType>({
|
||||
/**
|
||||
* `maxLength` is valid JSON Schema and is honored by providers, but the SDK's
|
||||
* `JsonSchemaType` does not declare it, so the model-facing bounds are typed here.
|
||||
* Runtime argument validation enforces the same limits as defense in depth.
|
||||
*/
|
||||
interface BoundedStringSchema {
|
||||
type: 'string';
|
||||
maxLength: number;
|
||||
description: string;
|
||||
}
|
||||
|
||||
interface CheckBackgroundTaskParameters {
|
||||
type: 'object';
|
||||
properties: {
|
||||
background_task_id: BoundedStringSchema;
|
||||
action: { type: 'string'; enum: string[]; description: string };
|
||||
message: BoundedStringSchema;
|
||||
control_id: BoundedStringSchema;
|
||||
};
|
||||
required: string[];
|
||||
}
|
||||
|
||||
const CHECK_BACKGROUND_TASK_PARAMETERS = Object.freeze<CheckBackgroundTaskParameters>({
|
||||
type: 'object',
|
||||
properties: {
|
||||
background_task_id: {
|
||||
type: 'string',
|
||||
maxLength: MAX_BACKGROUND_TASK_ID_CHARS,
|
||||
description:
|
||||
'The id returned when the tool or subagent was dispatched. Omit to list all background tasks in this thread.',
|
||||
},
|
||||
|
|
@ -314,10 +345,12 @@ const CHECK_BACKGROUND_TASK_PARAMETERS: JsonSchemaType = Object.freeze<JsonSchem
|
|||
},
|
||||
message: {
|
||||
type: 'string',
|
||||
maxLength: MAX_BACKGROUND_CONTROL_MESSAGE_CHARS,
|
||||
description: 'Required for steer, queue, or interrupt.',
|
||||
},
|
||||
control_id: {
|
||||
type: 'string',
|
||||
maxLength: MAX_BACKGROUND_CONTROL_ID_CHARS,
|
||||
description: 'Required for cancel_message; use the id returned by a prior control action.',
|
||||
},
|
||||
},
|
||||
|
|
@ -1080,28 +1113,86 @@ function buildSubagentControlCommand(
|
|||
return { action: 'cancel' };
|
||||
}
|
||||
if (action === 'cancel_message') {
|
||||
return typeof args.control_id === 'string'
|
||||
return typeof args.control_id === 'string' &&
|
||||
args.control_id.length <= MAX_BACKGROUND_CONTROL_ID_CHARS
|
||||
? { action: 'cancel_message', controlId: args.control_id }
|
||||
: undefined;
|
||||
}
|
||||
if (action === 'steer' || action === 'queue' || action === 'interrupt') {
|
||||
return typeof args.message === 'string' ? { action, message: args.message } : undefined;
|
||||
return typeof args.message === 'string' &&
|
||||
args.message.length <= MAX_BACKGROUND_CONTROL_MESSAGE_CHARS
|
||||
? { action, message: args.message }
|
||||
: undefined;
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* One tool call is one invocation, of a control or of the poll that collects a result.
|
||||
* A provider tool-call id such as `call_0` repeats across runs and agents, so the
|
||||
* identity also carries the run and executing agent; replaying that same call stays
|
||||
* idempotent while a later run's identical id is a new invocation. Hashing keeps every
|
||||
* derived identity inside the routed bound.
|
||||
*/
|
||||
function controlInvocationId(params: {
|
||||
toolCallId?: string;
|
||||
agentId?: string;
|
||||
runId?: string;
|
||||
}): string {
|
||||
const toolCallId = params.toolCallId?.trim();
|
||||
if (toolCallId == null || toolCallId === '') {
|
||||
return randomUUID();
|
||||
}
|
||||
return createHash('sha256')
|
||||
.update(`${params.runId ?? ''}\u0000${params.agentId ?? ''}\u0000${toolCallId}`)
|
||||
.digest('base64url')
|
||||
.slice(0, 32);
|
||||
}
|
||||
|
||||
/** Executes a `check_background_task` call and returns the ToolMessage content. */
|
||||
export function runCheckBackgroundTask(params: {
|
||||
interface RoutedSubagentTaskStore {
|
||||
claimTask(scopeId: string, taskId: string, invocationId: string): Promise<SubagentTaskClaim>;
|
||||
controlTask(
|
||||
scopeId: string,
|
||||
taskId: string,
|
||||
command: SubagentTaskControlCommand,
|
||||
invocationId: string,
|
||||
): Promise<SubagentTaskControlResult>;
|
||||
listTasks(scopeId: string): Promise<SubagentTaskSnapshot[]>;
|
||||
}
|
||||
|
||||
function routedSubagentStore(store: SubagentTaskStore): RoutedSubagentTaskStore | undefined {
|
||||
const candidate = store as SubagentTaskStore & Partial<RoutedSubagentTaskStore>;
|
||||
return typeof candidate.claimTask === 'function' &&
|
||||
typeof candidate.controlTask === 'function' &&
|
||||
typeof candidate.listTasks === 'function'
|
||||
? (candidate as RoutedSubagentTaskStore)
|
||||
: undefined;
|
||||
}
|
||||
|
||||
export async function runCheckBackgroundTask(params: {
|
||||
userId: string;
|
||||
conversationId: string;
|
||||
args: unknown;
|
||||
/** The provider's tool-call id: one control invocation, stable across replays. */
|
||||
toolCallId?: string;
|
||||
/** Scopes that tool-call id, whose provider ids repeat across runs and agents. */
|
||||
agentId?: string;
|
||||
runId?: string;
|
||||
subagentTasks?: SubagentTaskConfig;
|
||||
}): string {
|
||||
}): Promise<string> {
|
||||
const { userId, conversationId } = params;
|
||||
const args = coerceArgsObject(params.args) ?? {};
|
||||
const rawId = args.background_task_id;
|
||||
if (typeof rawId === 'string' && rawId.trim().length > MAX_BACKGROUND_TASK_ID_CHARS) {
|
||||
return JSON.stringify({
|
||||
status: 'invalid',
|
||||
message: `A background_task_id cannot exceed ${MAX_BACKGROUND_TASK_ID_CHARS} characters.`,
|
||||
});
|
||||
}
|
||||
const taskId = typeof rawId === 'string' && rawId.trim() !== '' ? rawId.trim() : undefined;
|
||||
const action = typeof args.action === 'string' && args.action !== '' ? args.action : 'poll';
|
||||
const invocationId = controlInvocationId(params);
|
||||
|
||||
if (taskId) {
|
||||
const task = backgroundTaskRegistry.get(userId, conversationId, taskId);
|
||||
|
|
@ -1118,28 +1209,44 @@ export function runCheckBackgroundTask(params: {
|
|||
|
||||
const subagentTasks = params.subagentTasks;
|
||||
if (subagentTasks != null) {
|
||||
if (action === 'poll') {
|
||||
const claimed = serializeSubagentClaim(
|
||||
subagentTasks.store.claim(subagentTasks.scopeId, taskId),
|
||||
);
|
||||
if (claimed != null) {
|
||||
return JSON.stringify(claimed);
|
||||
try {
|
||||
const routedStore = routedSubagentStore(subagentTasks.store);
|
||||
if (action === 'poll') {
|
||||
const claim =
|
||||
routedStore == null
|
||||
? subagentTasks.store.claim(subagentTasks.scopeId, taskId)
|
||||
: await routedStore.claimTask(subagentTasks.scopeId, taskId, invocationId);
|
||||
const claimed = serializeSubagentClaim(claim);
|
||||
if (claimed != null) {
|
||||
return JSON.stringify(claimed);
|
||||
}
|
||||
} else {
|
||||
const command = buildSubagentControlCommand(args, action);
|
||||
if (command == null) {
|
||||
return JSON.stringify({
|
||||
status: 'invalid',
|
||||
background_task_id: taskId,
|
||||
message: 'This subagent control action is unknown or missing its required argument.',
|
||||
});
|
||||
}
|
||||
const result =
|
||||
routedStore == null
|
||||
? subagentTasks.store.control(subagentTasks.scopeId, taskId, command)
|
||||
: await routedStore.controlTask(subagentTasks.scopeId, taskId, command, invocationId);
|
||||
const controlled = serializeSubagentControl(result);
|
||||
if (controlled != null) {
|
||||
return JSON.stringify(controlled);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
const command = buildSubagentControlCommand(args, action);
|
||||
if (command == null) {
|
||||
} catch (error) {
|
||||
if (error instanceof SubagentTaskOwnerUnavailableError) {
|
||||
return JSON.stringify({
|
||||
status: 'invalid',
|
||||
status: 'unavailable',
|
||||
background_task_id: taskId,
|
||||
message: 'This subagent control action is unknown or missing its required argument.',
|
||||
message: error.message,
|
||||
});
|
||||
}
|
||||
const controlled = serializeSubagentControl(
|
||||
subagentTasks.store.control(subagentTasks.scopeId, taskId, command),
|
||||
);
|
||||
if (controlled != null) {
|
||||
return JSON.stringify(controlled);
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1158,10 +1265,30 @@ export function runCheckBackgroundTask(params: {
|
|||
}
|
||||
|
||||
const tasks = backgroundTaskRegistry.list(userId, conversationId);
|
||||
const subagentTasks =
|
||||
params.subagentTasks?.store
|
||||
.list(params.subagentTasks.scopeId)
|
||||
.map((task) => serializeSubagentSnapshot(task)) ?? [];
|
||||
let subagentTasks: SerializedSubagentTask[] = [];
|
||||
let listWarning: string | undefined;
|
||||
if (params.subagentTasks != null) {
|
||||
try {
|
||||
const routedStore = routedSubagentStore(params.subagentTasks.store);
|
||||
const snapshots =
|
||||
routedStore == null
|
||||
? params.subagentTasks.store.list(params.subagentTasks.scopeId)
|
||||
: await routedStore.listTasks(params.subagentTasks.scopeId);
|
||||
subagentTasks = snapshots.map((task) => serializeSubagentSnapshot(task));
|
||||
} catch (error) {
|
||||
if (error instanceof SubagentTaskOwnerUnavailableError) {
|
||||
/** Cross-replica discovery is an additive source. A Redis outage must not
|
||||
* hide ordinary tasks or subagents owned by this process; surface the
|
||||
* incomplete view explicitly so the caller can retry for remote tasks. */
|
||||
subagentTasks = params.subagentTasks.store
|
||||
.list(params.subagentTasks.scopeId)
|
||||
.map((task) => serializeSubagentSnapshot(task));
|
||||
listWarning = `Cross-replica subagent tasks could not be listed: ${error.message}`;
|
||||
} else {
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
logger.debug(
|
||||
`[background] check_background_task listed ${tasks.length + subagentTasks.length} task(s)`,
|
||||
);
|
||||
|
|
@ -1170,6 +1297,7 @@ export function runCheckBackgroundTask(params: {
|
|||
...tasks.map((task) => serializeTask(task, { includeResult: false })),
|
||||
...subagentTasks,
|
||||
],
|
||||
...(listWarning != null && { partial: true, warning: listWarning }),
|
||||
});
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -31,11 +31,13 @@ function makeStore(): SubagentThreadTaskStore {
|
|||
const unused = jest.fn();
|
||||
return new SubagentThreadTaskStore({
|
||||
acquireSubagentThreadLease: unused as AllMethods['acquireSubagentThreadLease'],
|
||||
claimSubagentTaskResult: unused as AllMethods['claimSubagentTaskResult'],
|
||||
countActiveSubagentThreadLeases: unused as AllMethods['countActiveSubagentThreadLeases'],
|
||||
deleteConvos: unused as AllMethods['deleteConvos'],
|
||||
deleteMessages: unused as AllMethods['deleteMessages'],
|
||||
getConvo: unused as AllMethods['getConvo'],
|
||||
getMessages: unused as AllMethods['getMessages'],
|
||||
listActiveSubagentThreadLeases: unused as AllMethods['listActiveSubagentThreadLeases'],
|
||||
releaseSubagentThreadLease: unused as AllMethods['releaseSubagentThreadLease'],
|
||||
reserveSubagentThread: unused as AllMethods['reserveSubagentThread'],
|
||||
renewSubagentThreadLease: unused as AllMethods['renewSubagentThreadLease'],
|
||||
|
|
|
|||
|
|
@ -4089,10 +4089,13 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand
|
|||
const results: ToolExecuteResult[] = await Promise.all(
|
||||
toolCalls.map(async (tc: ToolCallRequest) => {
|
||||
if (backgroundControlEnabled && tc.name === CHECK_BACKGROUND_TASK_NAME) {
|
||||
const pollContent = runCheckBackgroundTask({
|
||||
const pollContent = await runCheckBackgroundTask({
|
||||
userId: backgroundUserId,
|
||||
conversationId: backgroundConversationId,
|
||||
args: tc.args,
|
||||
toolCallId: tc.id,
|
||||
agentId,
|
||||
runId: `${backgroundRunId ?? ''}:${tc.turn ?? ''}`,
|
||||
subagentTasks,
|
||||
});
|
||||
/** Deliver a completed task's artifact through THIS live poll
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ export * from './skills';
|
|||
export * from './phases';
|
||||
export * from './startup';
|
||||
export * from './subagentThreads';
|
||||
export * from './subagentTaskRouting';
|
||||
export * from './skillConfigurable';
|
||||
export * from './skillFiles';
|
||||
export * from './codeFilesSession';
|
||||
|
|
|
|||
1301
packages/api/src/agents/subagentTaskRouting.spec.ts
Normal file
1301
packages/api/src/agents/subagentTaskRouting.spec.ts
Normal file
File diff suppressed because it is too large
Load diff
1379
packages/api/src/agents/subagentTaskRouting.ts
Normal file
1379
packages/api/src/agents/subagentTaskRouting.ts
Normal file
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
|
|
@ -8,24 +8,36 @@ import {
|
|||
} from '@librechat/agents/langchain/messages';
|
||||
import type {
|
||||
InMemorySubagentTaskStoreOptions,
|
||||
SubagentTaskClaim,
|
||||
SubagentTaskConfig,
|
||||
SubagentTaskControlCommand,
|
||||
SubagentTaskControlResult,
|
||||
SubagentTaskRuntime,
|
||||
SubagentTaskSnapshot,
|
||||
SubagentTaskStartRequest,
|
||||
SubagentTaskStartResult,
|
||||
} from '@librechat/agents';
|
||||
import type {
|
||||
AllMethods,
|
||||
IActiveSubagentThreadLease,
|
||||
IConversation,
|
||||
IMessage,
|
||||
MessageMethods,
|
||||
ConversationMethods,
|
||||
SubagentTaskResultClaim,
|
||||
} from '@librechat/data-schemas';
|
||||
import type { BaseMessage, StoredMessage } from '@librechat/agents/langchain/messages';
|
||||
import type { SubagentTaskControlTransport } from './subagentTaskRouting';
|
||||
import type { UsageMetadata } from '~/stream/interfaces/IJobStore';
|
||||
import {
|
||||
boundedClaim,
|
||||
boundedTaskList,
|
||||
controlFingerprint,
|
||||
SubagentTaskOwnerUnavailableError,
|
||||
} from './subagentTaskRouting';
|
||||
import { createSubagentAttemptKey, createSubagentThreadId } from './subagentThreadIds';
|
||||
import { runWithDetachedSubagentUsage } from './subagentTaskContext';
|
||||
import { createConcurrencyLimiter } from '~/utils/promise';
|
||||
import { aggregateEmittedUsage } from './usage';
|
||||
|
||||
const SCOPE_VERSION = 1;
|
||||
|
|
@ -33,10 +45,29 @@ const DEFAULT_MAX_THREAD_DEPTH = 1;
|
|||
const DEFAULT_LEASE_TTL_MS = 30_000;
|
||||
const DEFAULT_LEASE_HEARTBEAT_MS = 10_000;
|
||||
const DEFAULT_OWNER_DRAIN_TIMEOUT_MS = 45_000;
|
||||
/** Keeps the admission fence alive across the deletion that follows the drain. */
|
||||
const OWNER_FENCE_GRACE_MS = 5 * 60_000;
|
||||
const DEFAULT_OWNER_DRAIN_POLL_MS = 100;
|
||||
/** Matches the deletion drain batch so cancellation cannot burst Redis. */
|
||||
const DELETION_CANCEL_CONCURRENCY = 32;
|
||||
/** Bounds retained control invocations; one entry per applied command. */
|
||||
const MAX_CONTROL_INVOCATIONS = 4_096;
|
||||
|
||||
/** A cancellation target set resolved before the conversations are removed. */
|
||||
export interface SubagentCancellationPlan {
|
||||
userId: string;
|
||||
tenantId?: string;
|
||||
conversationIds: string[];
|
||||
scopes: Array<{ scopeId: string; threadIds: string[] | null }>;
|
||||
leases: IActiveSubagentThreadLease[];
|
||||
}
|
||||
/** Three missed 10-second transport heartbeats retire a crashed owner. */
|
||||
const DEFAULT_TASK_ROUTING_TTL_MS = 30_000;
|
||||
const MAX_TRANSCRIPT_BYTES = 12 * 1024 * 1024;
|
||||
const TRANSCRIPT_SELECT =
|
||||
'messageId parentMessageId text createdAt +subagentTranscript +subagentTask';
|
||||
const DURABLE_RESULT_SELECT =
|
||||
'messageId conversationId sender text createdAt updatedAt +subagentTask';
|
||||
|
||||
class SubagentThreadPublicError extends Error {}
|
||||
class SubagentThreadDeletedError extends SubagentThreadPublicError {}
|
||||
|
|
@ -44,11 +75,13 @@ class SubagentThreadDeletedError extends SubagentThreadPublicError {}
|
|||
type SubagentThreadMethods = Pick<
|
||||
AllMethods,
|
||||
| 'acquireSubagentThreadLease'
|
||||
| 'claimSubagentTaskResult'
|
||||
| 'countActiveSubagentThreadLeases'
|
||||
| 'deleteConvos'
|
||||
| 'deleteMessages'
|
||||
| 'getConvo'
|
||||
| 'getMessages'
|
||||
| 'listActiveSubagentThreadLeases'
|
||||
| 'reserveSubagentThread'
|
||||
| 'releaseSubagentThreadLease'
|
||||
| 'renewSubagentThreadLease'
|
||||
|
|
@ -88,6 +121,8 @@ interface TaskThreadLease {
|
|||
shared?: {
|
||||
token: string;
|
||||
lost: boolean;
|
||||
/** Epoch ms this lease is durable until, advanced only by a confirmed renewal. */
|
||||
expiresAt: number;
|
||||
heartbeat?: ReturnType<typeof setInterval>;
|
||||
heartbeatInFlight?: Promise<void>;
|
||||
};
|
||||
|
|
@ -99,7 +134,13 @@ export interface SubagentThreadTaskStoreOptions extends InMemorySubagentTaskStor
|
|||
leaseHeartbeatMs?: number;
|
||||
ownerDrainTimeoutMs?: number;
|
||||
ownerDrainPollMs?: number;
|
||||
taskRoutingTtlMs?: number;
|
||||
isOwnerActive?: (userId: string) => Promise<boolean>;
|
||||
maxControlInvocations?: number;
|
||||
ownerFenceGraceMs?: number;
|
||||
fenceOwnerAdmission?: (userId: string, token: string, fencedUntil: Date) => Promise<void>;
|
||||
renewOwnerAdmission?: (userId: string, token: string, fencedUntil: Date) => Promise<boolean>;
|
||||
releaseOwnerAdmission?: (userId: string, token: string) => Promise<void>;
|
||||
}
|
||||
|
||||
function positiveInteger(value: number | undefined, fallback: number): number {
|
||||
|
|
@ -142,6 +183,10 @@ function parseScope(scopeId: string): SubagentThreadScope {
|
|||
};
|
||||
}
|
||||
|
||||
function serializeScope(scope: Omit<SubagentThreadScope, 'version'>): string {
|
||||
return JSON.stringify({ version: SCOPE_VERSION, ...scope });
|
||||
}
|
||||
|
||||
function matchesTenant(actual: string | undefined, expected: string | undefined): boolean {
|
||||
return actual === expected;
|
||||
}
|
||||
|
|
@ -284,21 +329,65 @@ function publicFailureDetail(error: unknown): string {
|
|||
: 'The child run could not be completed.';
|
||||
}
|
||||
|
||||
/** Rebuilds the terminal claim a recovered durable result stands for. */
|
||||
function recoveredClaim(
|
||||
message: IMessage,
|
||||
claim: Extract<SubagentTaskClaim, { status: 'claimed' }>,
|
||||
): SubagentTaskClaim | undefined {
|
||||
const status = message.subagentTask?.status;
|
||||
const content = message.text ?? '';
|
||||
/** A durable child message keeps the untruncated output, so recovering one applies
|
||||
* the same bounds a routed response would have. */
|
||||
if (status === 'completed') {
|
||||
return boundedClaim({ status: 'completed', task: claim.task, result: content });
|
||||
}
|
||||
if (status === 'error' || status === 'cancelled') {
|
||||
return boundedClaim({ status, task: claim.task, error: content });
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function drainKey(parentConversationId: string, taskId: string): string {
|
||||
return `${parentConversationId}\u0000${taskId}`;
|
||||
}
|
||||
|
||||
function safeErrorMessage(error: unknown): string {
|
||||
return `Subagent task failed: ${publicFailureDetail(error).slice(0, 2_000)}`;
|
||||
}
|
||||
|
||||
/** Persists view-only logical child threads with process-local controls and a shared execution fence. */
|
||||
/** Persists view-only logical child threads with owner-routed controls and a shared execution fence. */
|
||||
export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
||||
readonly supportsThreadContinuation = true;
|
||||
private readonly activeThreads = new Map<string, TaskThreadLease>();
|
||||
private readonly controlInvocations = new Map<
|
||||
string,
|
||||
{ scopeId: string; taskId: string; fingerprint: string; result: SubagentTaskControlResult }
|
||||
>();
|
||||
|
||||
private readonly parentPersistence = new Map<string, Promise<unknown>>();
|
||||
private readonly maxThreadDepth: number;
|
||||
private readonly leaseTtlMs: number;
|
||||
private readonly leaseHeartbeatMs: number;
|
||||
private readonly ownerDrainTimeoutMs: number;
|
||||
private readonly ownerDrainPollMs: number;
|
||||
private readonly taskRoutingTtlMs: number;
|
||||
private readonly maxControlInvocations: number;
|
||||
private readonly ownerFenceGraceMs: number;
|
||||
private readonly isOwnerActive: (userId: string) => Promise<boolean>;
|
||||
private readonly fenceOwnerAdmission?: (
|
||||
userId: string,
|
||||
token: string,
|
||||
fencedUntil: Date,
|
||||
) => Promise<void>;
|
||||
|
||||
private readonly renewOwnerAdmission?: (
|
||||
userId: string,
|
||||
token: string,
|
||||
fencedUntil: Date,
|
||||
) => Promise<boolean>;
|
||||
|
||||
private readonly releaseOwnerAdmission?: (userId: string, token: string) => Promise<void>;
|
||||
private taskControlTransport?: SubagentTaskControlTransport;
|
||||
|
||||
constructor(
|
||||
private readonly methods: SubagentThreadMethods,
|
||||
|
|
@ -319,7 +408,37 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
DEFAULT_OWNER_DRAIN_TIMEOUT_MS,
|
||||
);
|
||||
this.ownerDrainPollMs = positiveInteger(options.ownerDrainPollMs, DEFAULT_OWNER_DRAIN_POLL_MS);
|
||||
this.taskRoutingTtlMs = positiveInteger(options.taskRoutingTtlMs, DEFAULT_TASK_ROUTING_TTL_MS);
|
||||
this.maxControlInvocations = positiveInteger(
|
||||
options.maxControlInvocations,
|
||||
MAX_CONTROL_INVOCATIONS,
|
||||
);
|
||||
this.ownerFenceGraceMs = positiveInteger(options.ownerFenceGraceMs, OWNER_FENCE_GRACE_MS);
|
||||
this.isOwnerActive = options.isOwnerActive ?? (async () => true);
|
||||
this.fenceOwnerAdmission = options.fenceOwnerAdmission;
|
||||
this.renewOwnerAdmission = options.renewOwnerAdmission;
|
||||
this.releaseOwnerAdmission = options.releaseOwnerAdmission;
|
||||
}
|
||||
|
||||
/** Enables optional cross-replica lookup after the host's Redis service is ready. */
|
||||
async configureTaskControlTransport(transport: SubagentTaskControlTransport): Promise<void> {
|
||||
if (this.taskControlTransport != null) {
|
||||
throw new Error('Subagent task control transport is already configured.');
|
||||
}
|
||||
await transport.bind({
|
||||
claim: (scopeId, taskId) => super.claim(scopeId, taskId),
|
||||
control: (scopeId, taskId, command, invocationId) =>
|
||||
this.controlInvocation(scopeId, taskId, command, invocationId),
|
||||
list: (scopeId) => super.list(scopeId),
|
||||
cancelScope: (scopeId, threadIds) => this.cancelForScope(scopeId, threadIds),
|
||||
});
|
||||
this.taskControlTransport = transport;
|
||||
}
|
||||
|
||||
async destroyTaskControlTransport(): Promise<void> {
|
||||
const transport = this.taskControlTransport;
|
||||
this.taskControlTransport = undefined;
|
||||
await transport?.destroy();
|
||||
}
|
||||
|
||||
/** Gates child creation on the ordinary parent write without retaining request state. */
|
||||
|
|
@ -381,6 +500,15 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
if (runtime.signal.aborted) {
|
||||
throw runtime.signal.reason ?? new Error('Subagent task was cancelled.');
|
||||
}
|
||||
/** Publish the owner address before any provider work: a child running
|
||||
* while unaddressable cannot be polled, controlled, or cancelled, and its
|
||||
* side effects would already have happened by the time a heartbeat
|
||||
* republished it. A failed registration fails the task closed instead. */
|
||||
await this.taskControlTransport?.registerTask(
|
||||
request.scopeId,
|
||||
runtime.taskId,
|
||||
this.taskRoutingTtlMs,
|
||||
);
|
||||
await parentReady;
|
||||
const prepared = await this.prepareThread(
|
||||
request.scopeId,
|
||||
|
|
@ -490,6 +618,285 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
return started;
|
||||
}
|
||||
|
||||
/**
|
||||
* Claims locally when possible, otherwise asks the registered owning replica.
|
||||
*
|
||||
* A child's terminal result is durable in its own thread, so collection is recorded
|
||||
* there against the polling invocation rather than kept alive in the owner's memory.
|
||||
* The invocation that lost a response re-acquires its own result on the next poll;
|
||||
* a different invocation is told the result was already collected. Owner-side
|
||||
* retention stays a fast path, free to expire, instead of the only copy.
|
||||
*/
|
||||
async claimTask(
|
||||
scopeId: string,
|
||||
taskId: string,
|
||||
invocationId?: string,
|
||||
): Promise<SubagentTaskClaim> {
|
||||
const local = super.claim(scopeId, taskId);
|
||||
const claim =
|
||||
local.status !== 'not_found'
|
||||
? local
|
||||
: ((await this.taskControlTransport?.claim(scopeId, taskId)) ?? local);
|
||||
if (invocationId == null || claim.status === 'running') {
|
||||
return claim;
|
||||
}
|
||||
if (claim.status === 'not_found') {
|
||||
return this.claimDurableTaskResult(scopeId, taskId, invocationId);
|
||||
}
|
||||
const threadId = claim.task.threadId;
|
||||
if (threadId == null || threadId === '') {
|
||||
return claim;
|
||||
}
|
||||
/** The durable record decides who holds this one-shot result. The invocation that
|
||||
* already consumed it re-acquires and is handed it again, a second invocation is
|
||||
* told it was collected instead of being given a duplicate, and a task with no
|
||||
* durable record to arbitrate keeps whatever the owner just answered. */
|
||||
const collected = await this.assignResultClaim(
|
||||
parseScope(scopeId).userId,
|
||||
threadId,
|
||||
claim.task.taskId,
|
||||
invocationId,
|
||||
);
|
||||
if (collected.status === 'claimed') {
|
||||
return { status: 'claimed', task: claim.task };
|
||||
}
|
||||
if (collected.status === 'not_found') {
|
||||
return claim;
|
||||
}
|
||||
return claim.status === 'claimed' ? (recoveredClaim(collected.message, claim) ?? claim) : claim;
|
||||
}
|
||||
|
||||
/**
|
||||
* Recovers a terminal task after its owning process and Redis registration are gone.
|
||||
* The task id locates only a candidate; durable child lineage re-establishes the
|
||||
* trusted parent scope before the one-shot result is claimed.
|
||||
*/
|
||||
private async claimDurableTaskResult(
|
||||
scopeId: string,
|
||||
taskId: string,
|
||||
invocationId: string,
|
||||
): Promise<SubagentTaskClaim> {
|
||||
const scope = parseScope(scopeId);
|
||||
let message: IMessage | undefined;
|
||||
try {
|
||||
[message] = await this.methods.getMessages(
|
||||
{
|
||||
user: scope.userId,
|
||||
messageId: `${taskId}:assistant`,
|
||||
'subagentTask.status': { $in: ['completed', 'error', 'cancelled'] },
|
||||
},
|
||||
DURABLE_RESULT_SELECT,
|
||||
{ limit: 1, sort: false },
|
||||
);
|
||||
} catch (error) {
|
||||
logger.warn('[subagentThreads] Failed to locate a durable child result', error);
|
||||
throw new SubagentTaskOwnerUnavailableError();
|
||||
}
|
||||
const threadId = message?.conversationId;
|
||||
const status = message?.subagentTask?.status;
|
||||
if (
|
||||
message == null ||
|
||||
!isNonEmptyString(threadId) ||
|
||||
!isNonEmptyString(message.sender) ||
|
||||
(status !== 'completed' && status !== 'error' && status !== 'cancelled')
|
||||
) {
|
||||
return { status: 'not_found' };
|
||||
}
|
||||
|
||||
let parent: IConversation | null;
|
||||
let conversation: IConversation | null;
|
||||
try {
|
||||
[parent, conversation] = await Promise.all([
|
||||
this.methods.getConvo(scope.userId, scope.parentConversationId),
|
||||
this.methods.getConvo(scope.userId, threadId),
|
||||
]);
|
||||
} catch (error) {
|
||||
logger.warn('[subagentThreads] Failed to verify durable child lineage', error);
|
||||
throw new SubagentTaskOwnerUnavailableError();
|
||||
}
|
||||
const lineage = conversation?.subagentThread;
|
||||
if (
|
||||
parent == null ||
|
||||
conversation == null ||
|
||||
lineage == null ||
|
||||
conversation.endpoint !== EModelEndpoint.agents ||
|
||||
lineage.parentConversationId !== scope.parentConversationId ||
|
||||
lineage.subagentType !== message.sender ||
|
||||
lineage.depth > this.maxThreadDepth ||
|
||||
!matchesTenant(parent.tenantId, scope.tenantId) ||
|
||||
!matchesTenant(conversation.tenantId, scope.tenantId)
|
||||
) {
|
||||
return { status: 'not_found' };
|
||||
}
|
||||
|
||||
const createdAt = message.createdAt?.getTime();
|
||||
const updatedAt = message.updatedAt?.getTime() ?? createdAt;
|
||||
if (createdAt == null || updatedAt == null) {
|
||||
return { status: 'not_found' };
|
||||
}
|
||||
const task: SubagentTaskSnapshot = {
|
||||
taskId,
|
||||
threadId,
|
||||
subagentType: lineage.subagentType,
|
||||
status,
|
||||
createdAt,
|
||||
updatedAt,
|
||||
resultAvailable: true,
|
||||
resultClaimed: true,
|
||||
pendingControls: 0,
|
||||
...(status === 'completed' ? {} : { error: message.text ?? '' }),
|
||||
};
|
||||
const collected = await this.assignResultClaim(scope.userId, threadId, taskId, invocationId);
|
||||
if (collected.status === 'not_found') {
|
||||
return { status: 'not_found' };
|
||||
}
|
||||
if (collected.status === 'claimed') {
|
||||
return { status: 'claimed', task };
|
||||
}
|
||||
return (
|
||||
recoveredClaim(collected.message, { status: 'claimed', task }) ?? {
|
||||
status: 'not_found',
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Assigns one durable terminal result to the invocation collecting it. A failed
|
||||
* write is not an absent record: handing the result over without recording its
|
||||
* claimant would let another invocation acquire the same one-shot output once the
|
||||
* database recovers, so this reports the retryable path and leaves the result
|
||||
* unclaimed for a later poll.
|
||||
*/
|
||||
private async assignResultClaim(
|
||||
userId: string,
|
||||
threadId: string,
|
||||
taskId: string,
|
||||
invocationId: string,
|
||||
): Promise<SubagentTaskResultClaim> {
|
||||
try {
|
||||
return await this.methods.claimSubagentTaskResult({
|
||||
userId,
|
||||
conversationId: threadId,
|
||||
taskId,
|
||||
claimId: invocationId,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.warn('[subagentThreads] Failed to record a collected child result', error);
|
||||
throw new SubagentTaskOwnerUnavailableError();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Controls locally when possible, otherwise asks the registered owning replica.
|
||||
* `invocationId` identifies one caller invocation: a routed retransmission of that
|
||||
* invocation replays the owner's result, while a fresh invocation applies again even
|
||||
* when its action and message are identical.
|
||||
*/
|
||||
async controlTask(
|
||||
scopeId: string,
|
||||
taskId: string,
|
||||
command: SubagentTaskControlCommand,
|
||||
invocationId: string = randomUUID(),
|
||||
): Promise<SubagentTaskControlResult> {
|
||||
const local = this.controlInvocation(scopeId, taskId, command, invocationId);
|
||||
if (local.status !== 'not_found') {
|
||||
return local;
|
||||
}
|
||||
return (
|
||||
(await this.taskControlTransport?.control(scopeId, taskId, command, invocationId)) ?? local
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Applies one logical control exactly once for its owning task. Idempotency lives
|
||||
* here rather than in the transport so a local and a routed caller of the same
|
||||
* invocation agree, and it is keyed by task as well as invocation because provider
|
||||
* tool-call ids repeat across runs and agents.
|
||||
*/
|
||||
controlInvocation(
|
||||
scopeId: string,
|
||||
taskId: string,
|
||||
command: SubagentTaskControlCommand,
|
||||
invocationId: string,
|
||||
): SubagentTaskControlResult {
|
||||
const key = `${scopeId}\u0000${taskId}\u0000${invocationId}`;
|
||||
const fingerprint = controlFingerprint(command);
|
||||
const applied = this.controlInvocations.get(key);
|
||||
if (applied != null) {
|
||||
/** One invocation is one command; reusing its id for different content is a
|
||||
* caller error rather than a retry, so it is refused instead of applied. */
|
||||
return applied.fingerprint === fingerprint
|
||||
? applied.result
|
||||
: {
|
||||
status: 'invalid',
|
||||
message: 'This control invocation id was already used for a different command.',
|
||||
};
|
||||
}
|
||||
if (this.get(scopeId, taskId) == null) {
|
||||
/** Not this replica's task. Refusing here would keep the command from ever
|
||||
* reaching its owner, so local load cannot veto a remote cancellation: the
|
||||
* owner applies its own window to the routed request. */
|
||||
return this.control(scopeId, taskId, command);
|
||||
}
|
||||
if (!this.makeRoomForInvocation()) {
|
||||
/** Every tracked invocation belongs to a task this store still holds. Applying
|
||||
* this command without room to record it would let a caller retry apply it a
|
||||
* second time, so it is refused before the child is touched at all. */
|
||||
logger.warn('[subagentThreads] Refused a control; live invocation records are full');
|
||||
return {
|
||||
status: 'invalid',
|
||||
message: 'Too many control invocations are in flight for this process; retry shortly.',
|
||||
};
|
||||
}
|
||||
const result = this.control(scopeId, taskId, command);
|
||||
if (result.status === 'not_found') {
|
||||
return result;
|
||||
}
|
||||
this.controlInvocations.set(key, { scopeId, taskId, fingerprint, result });
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
* Frees invocation slots by dropping records whose task the store no longer holds:
|
||||
* a settled task cannot be controlled again, so its record is worthless, while a
|
||||
* live one is exactly what a caller retry needs to replay instead of applying its
|
||||
* command twice. The sweep runs only when the window is full and clears every dead
|
||||
* record at once, so it is amortized rather than repeated per control.
|
||||
*/
|
||||
private makeRoomForInvocation(): boolean {
|
||||
if (this.controlInvocations.size < this.maxControlInvocations) {
|
||||
return true;
|
||||
}
|
||||
for (const [key, invocation] of this.controlInvocations) {
|
||||
if (this.get(invocation.scopeId, invocation.taskId) == null) {
|
||||
this.controlInvocations.delete(key);
|
||||
}
|
||||
}
|
||||
return this.controlInvocations.size < this.maxControlInvocations;
|
||||
}
|
||||
|
||||
/** Returns this process's tasks plus tasks reported by registered remote owners. */
|
||||
async listTasks(scopeId: string): Promise<SubagentTaskSnapshot[]> {
|
||||
const local = super.list(scopeId);
|
||||
const remote = (await this.taskControlTransport?.list(scopeId)) ?? [];
|
||||
const byId = new Map(local.map((task) => [task.taskId, task]));
|
||||
for (const task of remote) {
|
||||
byId.set(task.taskId, task);
|
||||
}
|
||||
/** The remote aggregation and each owner's reply carry their own bound, but this
|
||||
* merge is what the poll tool reads: without a cap here the list the model sees is
|
||||
* that bound plus however many children this replica happens to own. */
|
||||
return boundedTaskList([...byId.values()]);
|
||||
}
|
||||
|
||||
/** Fast capability probe used while deciding whether a later turn needs the poll tool. */
|
||||
async hasTasks(scopeId: string): Promise<boolean> {
|
||||
if (super.list(scopeId).length > 0) {
|
||||
return true;
|
||||
}
|
||||
return (await this.taskControlTransport?.hasTasks(scopeId)) ?? false;
|
||||
}
|
||||
|
||||
override control(
|
||||
scopeId: string,
|
||||
taskId: string,
|
||||
|
|
@ -544,6 +951,148 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolves every cancellation target while the conversations still exist. The plan is
|
||||
* replayed after deletion, when those rows can no longer be read back, so the second
|
||||
* pass only has to reach registered owners through Redis.
|
||||
*/
|
||||
async planCancellationForConversations(
|
||||
userId: string,
|
||||
conversationIds: Iterable<string>,
|
||||
tenantId?: string,
|
||||
): Promise<SubagentCancellationPlan> {
|
||||
const targetIds = [...new Set(conversationIds)];
|
||||
const plan: SubagentCancellationPlan = {
|
||||
userId,
|
||||
...(tenantId == null ? {} : { tenantId }),
|
||||
conversationIds: targetIds,
|
||||
scopes: [],
|
||||
leases: [],
|
||||
};
|
||||
if (targetIds.length === 0 || this.taskControlTransport == null) {
|
||||
return plan;
|
||||
}
|
||||
const targets = new Set(targetIds);
|
||||
const scopeIdFor = (parentConversationId: string): string =>
|
||||
serializeScope({
|
||||
userId,
|
||||
parentConversationId,
|
||||
...(tenantId ? { tenantId } : {}),
|
||||
});
|
||||
/** Deleting a conversation takes its whole scope; a deleted child only cancels its
|
||||
* own thread inside a parent scope that survives. */
|
||||
const conversations = await Promise.all(
|
||||
targetIds.map((conversationId) => this.methods.getConvo(userId, conversationId)),
|
||||
);
|
||||
const threadTargetsByParent = new Map<string, Set<string>>();
|
||||
for (const [index, conversation] of conversations.entries()) {
|
||||
const parentConversationId = conversation?.subagentThread?.parentConversationId;
|
||||
if (
|
||||
parentConversationId == null ||
|
||||
targets.has(parentConversationId) ||
|
||||
!matchesTenant(conversation?.tenantId, tenantId)
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
const threadIds = threadTargetsByParent.get(parentConversationId) ?? new Set<string>();
|
||||
threadIds.add(targetIds[index]);
|
||||
threadTargetsByParent.set(parentConversationId, threadIds);
|
||||
}
|
||||
plan.scopes = [
|
||||
...targetIds.map((parentConversationId) => ({
|
||||
scopeId: scopeIdFor(parentConversationId),
|
||||
threadIds: null,
|
||||
})),
|
||||
...[...threadTargetsByParent].map(([parentConversationId, threadIds]) => ({
|
||||
scopeId: scopeIdFor(parentConversationId),
|
||||
threadIds: [...threadIds],
|
||||
})),
|
||||
];
|
||||
/** Captured now so descendants removed by the cascade stay reachable afterwards. */
|
||||
plan.leases = await this.methods.listActiveSubagentThreadLeases({
|
||||
user: userId,
|
||||
now: new Date(),
|
||||
...(tenantId == null ? {} : { tenantId }),
|
||||
});
|
||||
return plan;
|
||||
}
|
||||
|
||||
/**
|
||||
* Cancels local children and replays a plan against registered remote owners.
|
||||
* `removedConversationIds` extends it with the cascade a deletion reported, matched
|
||||
* against leases captured before those rows were removed.
|
||||
*/
|
||||
async cancelPlan(
|
||||
plan: SubagentCancellationPlan,
|
||||
removedConversationIds: Iterable<string> = [],
|
||||
): Promise<number> {
|
||||
const { userId, tenantId } = plan;
|
||||
const planned = new Set(plan.conversationIds);
|
||||
const removed = new Set(removedConversationIds);
|
||||
/** A cascade can remove descendants the plan never named — a grandchild lives in
|
||||
* its own parent's scope, not the deleted root's — so every removed conversation
|
||||
* is cancelled as a scope of its own. */
|
||||
const targets = [...new Set([...planned, ...removed])];
|
||||
let cancelled = this.cancelForConversations(userId, targets, tenantId);
|
||||
const transport = this.taskControlTransport;
|
||||
if (transport == null) {
|
||||
return cancelled;
|
||||
}
|
||||
const cancelSlot = createConcurrencyLimiter(DELETION_CANCEL_CONCURRENCY);
|
||||
const cascadeScopes = [...removed]
|
||||
.filter((conversationId) => !planned.has(conversationId))
|
||||
.map((parentConversationId) => ({
|
||||
scopeId: serializeScope({
|
||||
userId,
|
||||
parentConversationId,
|
||||
...(tenantId ? { tenantId } : {}),
|
||||
}),
|
||||
threadIds: null,
|
||||
}));
|
||||
const scopeCancellations = [...plan.scopes, ...cascadeScopes].map((scope) =>
|
||||
cancelSlot(() => transport.cancelScope(scope.scopeId, scope.threadIds)),
|
||||
);
|
||||
const leaseCancellations = plan.leases
|
||||
.filter(
|
||||
(lease) => removed.has(lease.parentConversationId) || removed.has(lease.conversationId),
|
||||
)
|
||||
.map((lease) =>
|
||||
cancelSlot(() =>
|
||||
this.controlTask(
|
||||
serializeScope({
|
||||
userId,
|
||||
parentConversationId: lease.parentConversationId,
|
||||
...(tenantId ? { tenantId } : {}),
|
||||
}),
|
||||
lease.taskId,
|
||||
{ action: 'cancel' },
|
||||
),
|
||||
),
|
||||
);
|
||||
for (const count of await Promise.all(scopeCancellations)) {
|
||||
cancelled += count;
|
||||
}
|
||||
for (const result of await Promise.all(leaseCancellations)) {
|
||||
if (result.status === 'cancelled') {
|
||||
cancelled += 1;
|
||||
}
|
||||
}
|
||||
return cancelled;
|
||||
}
|
||||
|
||||
/** Cancels this process's live children for one scope, optionally narrowed to threads. */
|
||||
private cancelForScope(scopeId: string, threadIds: string[] | null): number {
|
||||
const scope = parseScope(scopeId);
|
||||
const targets = threadIds == null ? null : new Set(threadIds);
|
||||
return this.cancelMatchingThreads(
|
||||
(candidate, threadId) =>
|
||||
candidate.userId === scope.userId &&
|
||||
candidate.parentConversationId === scope.parentConversationId &&
|
||||
matchesTenant(candidate.tenantId, scope.tenantId) &&
|
||||
(targets == null || targets.has(threadId)),
|
||||
);
|
||||
}
|
||||
|
||||
/** Cancels every active child owned by a user before a delete-all operation. */
|
||||
cancelForOwner(userId: string, tenantId?: string): number {
|
||||
return this.cancelMatchingThreads(
|
||||
|
|
@ -551,26 +1100,202 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
);
|
||||
}
|
||||
|
||||
/** Cancels local work and waits for every replica's durable lease to drain. */
|
||||
/**
|
||||
* Deletes an owner's conversations behind a durable admission fence. Draining alone
|
||||
* cannot close the race: a child admitted on another replica after the drain read
|
||||
* its leases would begin provider work against a parent that is about to disappear.
|
||||
* Fencing first inverts that — the fence is written before any lease is read, and a
|
||||
* child validates the fence after its own lease is written, so one of the two always
|
||||
* observes the other. The fence expires by itself, so a process lost mid-deletion
|
||||
* cannot leave the account unable to run subagents.
|
||||
*/
|
||||
async withOwnerDeletionFence<T>(
|
||||
userId: string,
|
||||
tenantId: string | undefined,
|
||||
deletion: () => Promise<T>,
|
||||
): Promise<T> {
|
||||
const fenceWindowMs = this.ownerDrainTimeoutMs + this.ownerFenceGraceMs;
|
||||
const token = randomUUID();
|
||||
/** Only a confirmed write moves this, so a run of failed renewals leaves it in the
|
||||
* past and the deletion can tell that its fence is no longer guaranteed. */
|
||||
let fencedUntil = Date.now() + fenceWindowMs;
|
||||
let fenceLapsed = false;
|
||||
await this.fenceOwnerAdmission?.(userId, token, new Date(fencedUntil));
|
||||
/** A very large account, or a stalled database, can outlast one fence window, and
|
||||
* a fence that expires mid-deletion lets another replica admit a child against
|
||||
* conversations being deleted. It is renewed for as long as the work runs. */
|
||||
let releasing = false;
|
||||
let inFlight: Promise<void> | undefined;
|
||||
const renewal = setInterval(
|
||||
() => {
|
||||
if (inFlight != null) {
|
||||
return;
|
||||
}
|
||||
inFlight = (async () => {
|
||||
const deadline = fencedUntil;
|
||||
const renewedUntil = Date.now() + fenceWindowMs;
|
||||
const held = await this.renewOwnerAdmission?.(userId, token, new Date(renewedUntil));
|
||||
if (held === false) {
|
||||
/** The durable entry was absent, so admission may already have opened even
|
||||
* when the local deadline has not passed. Reacquire for containment, but
|
||||
* retain the lapse so the enclosing deletion re-drains before success. */
|
||||
fenceLapsed = true;
|
||||
if (releasing) {
|
||||
return;
|
||||
}
|
||||
/** The entry is gone — expired, or pruned by another deletion — so this
|
||||
* deletion takes its fence again rather than running on unfenced. */
|
||||
await this.fenceOwnerAdmission?.(userId, token, new Date(renewedUntil));
|
||||
}
|
||||
if (Date.now() >= deadline) {
|
||||
/** The write only landed after the deadline it was meant to extend, so
|
||||
* admission stood open in between and a child could have taken a lease the
|
||||
* drain had already read past. A fence cannot be restored backwards over
|
||||
* that gap, so the lapse is recorded rather than papered over. */
|
||||
fenceLapsed = true;
|
||||
return;
|
||||
}
|
||||
fencedUntil = renewedUntil;
|
||||
})()
|
||||
.catch((error) => {
|
||||
logger.warn('[subagentThreads] Failed to hold the owner admission fence', error);
|
||||
})
|
||||
.finally(() => {
|
||||
inFlight = undefined;
|
||||
});
|
||||
},
|
||||
Math.max(1, Math.floor(fenceWindowMs / 3)),
|
||||
);
|
||||
renewal.unref?.();
|
||||
const stopRenewal = async (): Promise<void> => {
|
||||
clearInterval(renewal);
|
||||
await inFlight;
|
||||
};
|
||||
const fenceHeld = (): boolean =>
|
||||
this.fenceOwnerAdmission == null || (!fenceLapsed && Date.now() < fencedUntil);
|
||||
try {
|
||||
await this.cancelAndDrainForOwner(userId, tenantId);
|
||||
/** The drain can outlast the fence window when the database is unreachable, and
|
||||
* renewals that keep failing leave the account open to admitting a child against
|
||||
* conversations about to disappear. Nothing has been removed yet, so this fails
|
||||
* closed and the caller retries once the fence can be held again. */
|
||||
if (!fenceHeld()) {
|
||||
throw new Error('The subagent admission fence expired before this deletion began.');
|
||||
}
|
||||
const deleted = await deletion();
|
||||
/** Settle a renewal already in flight before deciding whether deletion crossed a
|
||||
* gap. Otherwise a late write can report the lapse only after this check and the
|
||||
* finally block would release the fence without re-draining. */
|
||||
await stopRenewal();
|
||||
if (!fenceHeld()) {
|
||||
/** The rows are gone, but the gap can leave a child another replica admitted
|
||||
* while the fence was down. Re-take the fence and drain that work before this
|
||||
* operation may report success. */
|
||||
logger.error(
|
||||
'[subagentThreads] Owner deletion outlived its admission fence; draining children admitted in the gap',
|
||||
);
|
||||
const recoveryUntil = Date.now() + fenceWindowMs;
|
||||
const reheld = await this.renewOwnerAdmission?.(userId, token, new Date(recoveryUntil));
|
||||
if (reheld !== true) {
|
||||
await this.fenceOwnerAdmission?.(userId, token, new Date(recoveryUntil));
|
||||
}
|
||||
if (Date.now() >= recoveryUntil) {
|
||||
throw new Error('The subagent admission fence expired while it was being restored.');
|
||||
}
|
||||
fencedUntil = recoveryUntil;
|
||||
fenceLapsed = false;
|
||||
await this.cancelAndDrainForOwner(userId, tenantId);
|
||||
if (!fenceHeld()) {
|
||||
throw new Error('The subagent admission fence expired while recovering this deletion.');
|
||||
}
|
||||
}
|
||||
return deleted;
|
||||
} finally {
|
||||
releasing = true;
|
||||
await stopRenewal();
|
||||
/** `clearInterval` stops only future passes. A renewal still waiting on the
|
||||
* database would otherwise find its fence released, read that as expiry, and
|
||||
* write a fresh one that nothing is left to lift. */
|
||||
/** Only this deletion's own fence is lifted: an overlapping deletion that took a
|
||||
* later one keeps admission closed until it finishes. */
|
||||
await this.releaseOwnerAdmission?.(userId, token).catch((error) => {
|
||||
logger.warn('[subagentThreads] Failed to release the owner admission fence', error);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Cancels local work and waits for every replica's durable lease to drain. Each task
|
||||
* is cancelled under one invocation held for the whole drain and only while its
|
||||
* owner has not answered: a fresh invocation per poll would retain a replay entry on
|
||||
* the owner for every pass, and a task already reported cancelled needs no second
|
||||
* command, only its lease to disappear.
|
||||
*/
|
||||
async cancelAndDrainForOwner(userId: string, tenantId?: string): Promise<void> {
|
||||
this.cancelForOwner(userId, tenantId);
|
||||
const deadline = Date.now() + this.ownerDrainTimeoutMs;
|
||||
const invocations = new Map<string, string>();
|
||||
const answered = new Set<string>();
|
||||
while (true) {
|
||||
const active = await this.methods.countActiveSubagentThreadLeases({
|
||||
const activeLeases = await this.methods.listActiveSubagentThreadLeases({
|
||||
user: userId,
|
||||
now: new Date(),
|
||||
...(tenantId == null ? {} : { tenantId }),
|
||||
});
|
||||
if (active === 0) {
|
||||
if (activeLeases.length === 0) {
|
||||
return;
|
||||
}
|
||||
if (Date.now() >= deadline) {
|
||||
throw new Error('Timed out draining detached subagent tasks for account deletion.');
|
||||
}
|
||||
const unanswered = activeLeases.filter(
|
||||
({ parentConversationId, taskId }) => !answered.has(drainKey(parentConversationId, taskId)),
|
||||
);
|
||||
for (let index = 0; index < unanswered.length; index += DELETION_CANCEL_CONCURRENCY) {
|
||||
await Promise.all(
|
||||
unanswered
|
||||
.slice(index, index + DELETION_CANCEL_CONCURRENCY)
|
||||
.map(({ parentConversationId, taskId }) =>
|
||||
this.cancelDrainedTask(
|
||||
{ userId, parentConversationId, taskId, tenantId },
|
||||
invocations,
|
||||
answered,
|
||||
),
|
||||
),
|
||||
);
|
||||
}
|
||||
await new Promise<void>((resolve) => setTimeout(resolve, this.ownerDrainPollMs));
|
||||
}
|
||||
}
|
||||
|
||||
/** Sends one drained task's cancellation, retrying only unconfirmed deliveries. */
|
||||
private async cancelDrainedTask(
|
||||
target: { userId: string; parentConversationId: string; taskId: string; tenantId?: string },
|
||||
invocations: Map<string, string>,
|
||||
answered: Set<string>,
|
||||
): Promise<void> {
|
||||
const { userId, parentConversationId, taskId, tenantId } = target;
|
||||
const key = drainKey(parentConversationId, taskId);
|
||||
const invocationId = invocations.get(key) ?? randomUUID();
|
||||
invocations.set(key, invocationId);
|
||||
const scopeId = serializeScope({
|
||||
userId,
|
||||
parentConversationId,
|
||||
...(tenantId == null ? {} : { tenantId }),
|
||||
});
|
||||
try {
|
||||
const result = await this.controlTask(scopeId, taskId, { action: 'cancel' }, invocationId);
|
||||
/** Only the owner confirming the task is stopped ends the commands for it. A
|
||||
* `not_found` means its registration is missing while its lease is still live —
|
||||
* an unconfirmed delivery, retried once the owner republishes itself. */
|
||||
if (result.status === 'cancelled' || result.status === 'not_running') {
|
||||
answered.add(key);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.warn('[subagentThreads] Retrying an unconfirmed child cancellation', error);
|
||||
}
|
||||
}
|
||||
|
||||
private startSharedLeaseHeartbeat(
|
||||
scopeId: string,
|
||||
scope: SubagentThreadScope,
|
||||
|
|
@ -645,19 +1370,32 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
return false;
|
||||
}
|
||||
try {
|
||||
const deadline = shared.expiresAt;
|
||||
const now = new Date();
|
||||
const renewedUntil = now.getTime() + this.leaseTtlMs;
|
||||
const renewed = await this.methods.renewSubagentThreadLease({
|
||||
user: scope.userId,
|
||||
conversationId: threadId,
|
||||
token: shared.token,
|
||||
now,
|
||||
expiresAt: new Date(now.getTime() + this.leaseTtlMs),
|
||||
expiresAt: new Date(renewedUntil),
|
||||
...(scope.tenantId == null ? {} : { tenantId: scope.tenantId }),
|
||||
});
|
||||
if (!renewed) {
|
||||
shared.lost = true;
|
||||
return false;
|
||||
}
|
||||
return renewed;
|
||||
if (Date.now() >= deadline) {
|
||||
/** The renewal filter compares against the `now` captured before the call, so a
|
||||
* write that only lands after this lease had expired still succeeds and moves
|
||||
* the row forward. An owner drain reading active leases in that gap saw this
|
||||
* thread as free, so the executor stops rather than run past a deletion that
|
||||
* may already have stepped over it. */
|
||||
shared.lost = true;
|
||||
return false;
|
||||
}
|
||||
shared.expiresAt = renewedUntil;
|
||||
return true;
|
||||
} catch (error) {
|
||||
shared.lost = true;
|
||||
logger.warn('[subagentThreads] Lost the shared child-thread lease', error);
|
||||
|
|
@ -779,7 +1517,11 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
'This child thread is already being continued by another run.',
|
||||
);
|
||||
}
|
||||
lease.shared = { token: sharedToken, lost: false };
|
||||
lease.shared = {
|
||||
token: sharedToken,
|
||||
lost: false,
|
||||
expiresAt: now.getTime() + this.leaseTtlMs,
|
||||
};
|
||||
this.startSharedLeaseHeartbeat(scopeId, scope, threadId, lease);
|
||||
/** Account deletion can fence the owner after the optimistic probe but before
|
||||
* this lease exists. Once the lease is visible, revalidate so deletion either
|
||||
|
|
@ -1226,6 +1968,22 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
}
|
||||
}
|
||||
|
||||
const REQUIRED_THREAD_METHODS = [
|
||||
'acquireSubagentThreadLease',
|
||||
'claimSubagentTaskResult',
|
||||
'countActiveSubagentThreadLeases',
|
||||
'deleteConvos',
|
||||
'deleteMessages',
|
||||
'getConvo',
|
||||
'getMessages',
|
||||
'listActiveSubagentThreadLeases',
|
||||
'releaseSubagentThreadLease',
|
||||
'renewSubagentThreadLease',
|
||||
'reserveSubagentThread',
|
||||
'saveConvo',
|
||||
'saveMessage',
|
||||
] as const;
|
||||
|
||||
export function createSubagentThreadTaskStore(
|
||||
methods: Pick<
|
||||
ConversationMethods,
|
||||
|
|
@ -1233,14 +1991,27 @@ export function createSubagentThreadTaskStore(
|
|||
| 'countActiveSubagentThreadLeases'
|
||||
| 'deleteConvos'
|
||||
| 'getConvo'
|
||||
| 'listActiveSubagentThreadLeases'
|
||||
| 'releaseSubagentThreadLease'
|
||||
| 'reserveSubagentThread'
|
||||
| 'renewSubagentThreadLease'
|
||||
| 'saveConvo'
|
||||
> &
|
||||
Pick<MessageMethods, 'deleteMessages' | 'getMessages' | 'saveMessage'>,
|
||||
Pick<
|
||||
MessageMethods,
|
||||
'claimSubagentTaskResult' | 'deleteMessages' | 'getMessages' | 'saveMessage'
|
||||
>,
|
||||
options?: SubagentThreadTaskStoreOptions,
|
||||
): SubagentThreadTaskStore {
|
||||
/** The host wires this from JavaScript, where the parameter type checks nothing. A
|
||||
* method missing there would otherwise surface as a routed failure at claim time,
|
||||
* long after startup, so the omission is caught here instead. */
|
||||
const missing = REQUIRED_THREAD_METHODS.filter(
|
||||
(name) => typeof (methods as Record<string, unknown>)[name] !== 'function',
|
||||
);
|
||||
if (missing.length > 0) {
|
||||
throw new Error(`Subagent thread task store is missing methods: ${missing.join(', ')}`);
|
||||
}
|
||||
return new SubagentThreadTaskStore(methods, options);
|
||||
}
|
||||
|
||||
|
|
@ -1250,6 +2021,6 @@ export function buildSubagentThreadTaskConfig(
|
|||
): SubagentTaskConfig {
|
||||
return {
|
||||
store,
|
||||
scopeId: JSON.stringify({ version: SCOPE_VERSION, ...scope }),
|
||||
scopeId: serializeScope(scope),
|
||||
};
|
||||
}
|
||||
|
|
|
|||
81
packages/api/src/cache/redisUtils.spec.ts
vendored
Normal file
81
packages/api/src/cache/redisUtils.spec.ts
vendored
Normal file
|
|
@ -0,0 +1,81 @@
|
|||
import IoRedis from 'ioredis';
|
||||
import { duplicateIoRedisClient } from './redisUtils';
|
||||
|
||||
describe('duplicateIoRedisClient', () => {
|
||||
it('applies overrides to a single-node duplicate', () => {
|
||||
const client = new IoRedis({ host: '127.0.0.1', port: 6379, lazyConnect: true });
|
||||
const duplicate = duplicateIoRedisClient(client, { enableOfflineQueue: false });
|
||||
try {
|
||||
expect(duplicate.options.enableOfflineQueue).toBe(false);
|
||||
expect(client.options.enableOfflineQueue).not.toBe(false);
|
||||
} finally {
|
||||
duplicate.disconnect();
|
||||
client.disconnect();
|
||||
}
|
||||
});
|
||||
|
||||
it('applies overrides to a cluster duplicate, whose options come second', () => {
|
||||
const client = new IoRedis.Cluster([{ host: '127.0.0.1', port: 6379 }], {
|
||||
lazyConnect: true,
|
||||
});
|
||||
const duplicate = duplicateIoRedisClient(client, { enableOfflineQueue: false });
|
||||
try {
|
||||
/** `Cluster.duplicate` reads its first argument as startup nodes, so passing the
|
||||
* overrides positionally silently keeps the original's queueing behaviour. */
|
||||
expect(duplicate.options.enableOfflineQueue).toBe(false);
|
||||
expect(client.options.enableOfflineQueue).not.toBe(false);
|
||||
} finally {
|
||||
duplicate.disconnect();
|
||||
client.disconnect();
|
||||
}
|
||||
});
|
||||
|
||||
it('disables the offline queue only after a cluster node is ready', () => {
|
||||
const client = new IoRedis.Cluster([{ host: '127.0.0.1', port: 6379 }], {
|
||||
lazyConnect: true,
|
||||
});
|
||||
const duplicate = duplicateIoRedisClient(client, { enableOfflineQueue: false });
|
||||
try {
|
||||
/** ioredis emits from its private pool and synchronously forwards `+node` from
|
||||
* `Cluster`; drive that real discovery path so the test cannot pass merely
|
||||
* because a synthetic event happened to share the public event name. */
|
||||
const pool = (
|
||||
duplicate as unknown as {
|
||||
connectionPool: {
|
||||
findOrCreate(options: { host: string; port: number }): InstanceType<typeof IoRedis>;
|
||||
};
|
||||
}
|
||||
).connectionPool;
|
||||
const node = pool.findOrCreate({ host: '127.0.0.1', port: 6380 });
|
||||
/** Topology discovery needs the node queue until this connection is ready. */
|
||||
expect(node.options.enableOfflineQueue).toBe(true);
|
||||
node.emit('ready');
|
||||
expect(node.options.enableOfflineQueue).toBe(false);
|
||||
} finally {
|
||||
duplicate.disconnect();
|
||||
client.disconnect();
|
||||
}
|
||||
});
|
||||
|
||||
it('disables the offline queue immediately on nodes discovered after cluster readiness', () => {
|
||||
const client = new IoRedis.Cluster([{ host: '127.0.0.1', port: 6379 }], {
|
||||
lazyConnect: true,
|
||||
});
|
||||
const duplicate = duplicateIoRedisClient(client, { enableOfflineQueue: false });
|
||||
try {
|
||||
duplicate.emit('ready');
|
||||
const pool = (
|
||||
duplicate as unknown as {
|
||||
connectionPool: {
|
||||
findOrCreate(options: { host: string; port: number }): InstanceType<typeof IoRedis>;
|
||||
};
|
||||
}
|
||||
).connectionPool;
|
||||
const replacement = pool.findOrCreate({ host: '127.0.0.1', port: 6381 });
|
||||
expect(replacement.options.enableOfflineQueue).toBe(false);
|
||||
} finally {
|
||||
duplicate.disconnect();
|
||||
client.disconnect();
|
||||
}
|
||||
});
|
||||
});
|
||||
44
packages/api/src/cache/redisUtils.ts
vendored
44
packages/api/src/cache/redisUtils.ts
vendored
|
|
@ -1,7 +1,49 @@
|
|||
import type { RedisClientType, RedisClusterType } from '@redis/client';
|
||||
import { logger } from '@librechat/data-schemas';
|
||||
import type { ClusterOptions, RedisOptions, Cluster, Redis } from 'ioredis';
|
||||
import type { RedisClientType, RedisClusterType } from '@redis/client';
|
||||
import { cacheConfig } from './cacheConfig';
|
||||
|
||||
/**
|
||||
* Duplicates an ioredis connection with option overrides. `Cluster.duplicate` reads its
|
||||
* first argument as an optional startup-node list and its second as the overrides,
|
||||
* unlike `Redis.duplicate`, so options passed positionally to a cluster are silently
|
||||
* dropped and the duplicate quietly inherits the original's behaviour.
|
||||
*/
|
||||
export function duplicateIoRedisClient(
|
||||
client: Redis | Cluster,
|
||||
options: RedisOptions & ClusterOptions = {},
|
||||
): Redis | Cluster {
|
||||
if (client.isCluster) {
|
||||
const duplicate = (client as Cluster).duplicate([], options);
|
||||
if (options.enableOfflineQueue !== false) {
|
||||
return duplicate;
|
||||
}
|
||||
let clusterHasBeenReady = duplicate.status === 'ready';
|
||||
duplicate.once('ready', () => {
|
||||
clusterHasBeenReady = true;
|
||||
});
|
||||
/** ioredis deliberately forces `enableOfflineQueue: true` on every Cluster node
|
||||
* after applying `redisOptions`. It needs that queue while a new node discovers
|
||||
* topology, so changing it at `+node` prevents the cluster from ever becoming
|
||||
* ready. Initial nodes switch once connected; nodes discovered after the cluster
|
||||
* was usable fail fast immediately, including during a slot-owner replacement. */
|
||||
const disableNodeOfflineQueue = (node: Redis): void => {
|
||||
const disable = (): void => {
|
||||
node.options.enableOfflineQueue = false;
|
||||
};
|
||||
if (node.status === 'ready' || clusterHasBeenReady) {
|
||||
disable();
|
||||
} else {
|
||||
node.once('ready', disable);
|
||||
}
|
||||
};
|
||||
duplicate.on('+node', disableNodeOfflineQueue);
|
||||
duplicate.nodes('all').forEach(disableNodeOfflineQueue);
|
||||
return duplicate;
|
||||
}
|
||||
return (client as Redis).duplicate(options);
|
||||
}
|
||||
|
||||
/**
|
||||
* Efficiently deletes multiple Redis keys with support for both cluster and single-node modes.
|
||||
*
|
||||
|
|
|
|||
|
|
@ -3338,6 +3338,15 @@ describe('Conversation Operations', () => {
|
|||
const winner = claims[0] ? 'token-a' : 'token-b';
|
||||
const loser = winner === 'token-a' ? 'token-b' : 'token-a';
|
||||
expect(await methods.countActiveSubagentThreadLeases({ user: 'lease-user', now })).toBe(1);
|
||||
await expect(
|
||||
methods.listActiveSubagentThreadLeases({ user: 'lease-user', now }),
|
||||
).resolves.toEqual([
|
||||
{
|
||||
conversationId,
|
||||
parentConversationId: 'parent',
|
||||
taskId: `task-${winner}`,
|
||||
},
|
||||
]);
|
||||
expect(await methods.getConvo('lease-user', conversationId)).not.toHaveProperty(
|
||||
'subagentThreadLease',
|
||||
);
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import type { DeleteResult } from 'mongoose';
|
|||
import type {
|
||||
AppConfig,
|
||||
IChatProjectDocument,
|
||||
IActiveSubagentThreadLease,
|
||||
IConversation,
|
||||
ISharedLink,
|
||||
ISubagentThreadReservation,
|
||||
|
|
@ -183,6 +184,11 @@ export interface ConversationMethods {
|
|||
now: Date;
|
||||
tenantId?: string;
|
||||
}): Promise<number>;
|
||||
listActiveSubagentThreadLeases(input: {
|
||||
user: string;
|
||||
now: Date;
|
||||
tenantId?: string;
|
||||
}): Promise<IActiveSubagentThreadLease[]>;
|
||||
getConvoOwnership(
|
||||
user: string,
|
||||
conversationId: string,
|
||||
|
|
@ -389,6 +395,37 @@ export function createConversationMethods(
|
|||
});
|
||||
}
|
||||
|
||||
/** Resolves only live task addresses so account-wide cancellation stays O(active tasks). */
|
||||
async function listActiveSubagentThreadLeases(input: {
|
||||
user: string;
|
||||
now: Date;
|
||||
tenantId?: string;
|
||||
}): Promise<IActiveSubagentThreadLease[]> {
|
||||
const Conversation = mongoose.models.Conversation as Model<IConversation>;
|
||||
const conversations = await Conversation.find({
|
||||
user: input.user,
|
||||
...subagentLeaseTenantFilter(input.tenantId),
|
||||
'subagentThreadLease.expiresAt': { $gt: input.now },
|
||||
})
|
||||
.select('conversationId subagentThread.parentConversationId +subagentThreadLease')
|
||||
.lean<
|
||||
Array<Pick<IConversation, 'conversationId' | 'subagentThread' | 'subagentThreadLease'>>
|
||||
>();
|
||||
return conversations.flatMap((conversation) => {
|
||||
const { conversationId } = conversation;
|
||||
const parentConversationId = conversation.subagentThread?.parentConversationId;
|
||||
const taskId = conversation.subagentThreadLease?.taskId;
|
||||
return typeof conversationId === 'string' &&
|
||||
conversationId !== '' &&
|
||||
typeof parentConversationId === 'string' &&
|
||||
parentConversationId !== '' &&
|
||||
typeof taskId === 'string' &&
|
||||
taskId !== ''
|
||||
? [{ conversationId, parentConversationId, taskId }]
|
||||
: [];
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Ownership probe for request validation: resolves only the owning user id
|
||||
* instead of materializing the full conversation document (preset spread +
|
||||
|
|
@ -1525,6 +1562,7 @@ export function createConversationMethods(
|
|||
renewSubagentThreadLease,
|
||||
releaseSubagentThreadLease,
|
||||
countActiveSubagentThreadLeases,
|
||||
listActiveSubagentThreadLeases,
|
||||
getConvoOwnership,
|
||||
getConvoRetention,
|
||||
getConvoTitle,
|
||||
|
|
|
|||
|
|
@ -44,7 +44,12 @@ import { createCategoriesMethods, type CategoriesMethods } from './categories';
|
|||
import { createPresetMethods, type PresetMethods } from './preset';
|
||||
/* Tier 2 — Moderate (service deps injected) */
|
||||
import { createConversationTagMethods, type ConversationTagMethods } from './conversationTag';
|
||||
import { createMessageMethods, CLIENT_MESSAGE_SELECT, type MessageMethods } from './message';
|
||||
import {
|
||||
createMessageMethods,
|
||||
CLIENT_MESSAGE_SELECT,
|
||||
type MessageMethods,
|
||||
type SubagentTaskResultClaim,
|
||||
} from './message';
|
||||
import { createConversationMethods, type ConversationMethods } from './conversation';
|
||||
import { createChatProjectMethods, type ChatProjectMethods } from './chatProject';
|
||||
export type {
|
||||
|
|
@ -368,6 +373,7 @@ export type {
|
|||
PresetMethods,
|
||||
ConversationTagMethods,
|
||||
MessageMethods,
|
||||
SubagentTaskResultClaim,
|
||||
ConversationMethods,
|
||||
ChatProjectMethods,
|
||||
TxMethods,
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ let bulkSaveMessages: ReturnType<typeof createMessageMethods>['bulkSaveMessages'
|
|||
let updateMessageText: ReturnType<typeof createMessageMethods>['updateMessageText'];
|
||||
let deleteMessagesSince: ReturnType<typeof createMessageMethods>['deleteMessagesSince'];
|
||||
let recordMessage: ReturnType<typeof createMessageMethods>['recordMessage'];
|
||||
let claimSubagentTaskResult: ReturnType<typeof createMessageMethods>['claimSubagentTaskResult'];
|
||||
|
||||
beforeAll(async () => {
|
||||
mongoServer = await MongoMemoryServer.create();
|
||||
|
|
@ -47,6 +48,7 @@ beforeAll(async () => {
|
|||
updateMessageText = methods.updateMessageText;
|
||||
deleteMessagesSince = methods.deleteMessagesSince;
|
||||
recordMessage = methods.recordMessage;
|
||||
claimSubagentTaskResult = methods.claimSubagentTaskResult;
|
||||
|
||||
await mongoose.connect(mongoUri);
|
||||
});
|
||||
|
|
@ -1549,4 +1551,81 @@ describe('Message Operations', () => {
|
|||
expect(doc?.tenantId).toBeUndefined();
|
||||
});
|
||||
});
|
||||
describe('claimSubagentTaskResult', () => {
|
||||
const terminalResult = async (taskId: string, conversationId: string, status: string) =>
|
||||
saveMessage({ userId: 'user123' }, {
|
||||
messageId: `${taskId}:assistant`,
|
||||
conversationId,
|
||||
text: 'child result',
|
||||
subagentTask: { attemptKey: `${taskId}:attempt`, status },
|
||||
} as Partial<IMessage>);
|
||||
|
||||
it('hands one terminal result to a single polling invocation', async () => {
|
||||
const taskId = uuidv4();
|
||||
const conversationId = uuidv4();
|
||||
await terminalResult(taskId, conversationId, 'completed');
|
||||
|
||||
const first = await claimSubagentTaskResult({
|
||||
userId: 'user123',
|
||||
conversationId,
|
||||
taskId,
|
||||
claimId: 'poll-1',
|
||||
});
|
||||
expect(first.status).toBe('acquired');
|
||||
expect(first.status === 'acquired' && first.message.text).toBe('child result');
|
||||
|
||||
/** The same invocation retrying recovers the result it never received. */
|
||||
const retried = await claimSubagentTaskResult({
|
||||
userId: 'user123',
|
||||
conversationId,
|
||||
taskId,
|
||||
claimId: 'poll-1',
|
||||
});
|
||||
expect(retried.status).toBe('acquired');
|
||||
|
||||
/** Another invocation is told it was collected instead of handed a copy. */
|
||||
await expect(
|
||||
claimSubagentTaskResult({ userId: 'user123', conversationId, taskId, claimId: 'poll-2' }),
|
||||
).resolves.toEqual({ status: 'claimed' });
|
||||
});
|
||||
|
||||
it('reports a result that is missing or still running as not found', async () => {
|
||||
const runningTaskId = uuidv4();
|
||||
const conversationId = uuidv4();
|
||||
await terminalResult(runningTaskId, conversationId, 'running');
|
||||
|
||||
await expect(
|
||||
claimSubagentTaskResult({
|
||||
userId: 'user123',
|
||||
conversationId,
|
||||
taskId: runningTaskId,
|
||||
claimId: 'poll-1',
|
||||
}),
|
||||
).resolves.toEqual({ status: 'not_found' });
|
||||
|
||||
await expect(
|
||||
claimSubagentTaskResult({
|
||||
userId: 'user123',
|
||||
conversationId,
|
||||
taskId: uuidv4(),
|
||||
claimId: 'poll-1',
|
||||
}),
|
||||
).resolves.toEqual({ status: 'not_found' });
|
||||
});
|
||||
|
||||
it('never hands one owner’s result to another user', async () => {
|
||||
const taskId = uuidv4();
|
||||
const conversationId = uuidv4();
|
||||
await terminalResult(taskId, conversationId, 'completed');
|
||||
|
||||
await expect(
|
||||
claimSubagentTaskResult({
|
||||
userId: 'other-user',
|
||||
conversationId,
|
||||
taskId,
|
||||
claimId: 'poll-1',
|
||||
}),
|
||||
).resolves.toEqual({ status: 'not_found' });
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -48,6 +48,11 @@ interface MessageQueryOptions {
|
|||
sort?: Record<string, 1 | -1> | false;
|
||||
}
|
||||
|
||||
export type SubagentTaskResultClaim =
|
||||
| { status: 'not_found' }
|
||||
| { status: 'claimed' }
|
||||
| { status: 'acquired'; message: IMessage };
|
||||
|
||||
export interface MessageMethods {
|
||||
saveMessage(
|
||||
ctx: { userId: string; isTemporary?: boolean; interfaceConfig?: AppConfig['interfaceConfig'] },
|
||||
|
|
@ -82,6 +87,12 @@ export interface MessageMethods {
|
|||
message: Partial<IMessage> & { newMessageId?: string },
|
||||
metadata?: { context?: string },
|
||||
): Promise<Partial<IMessage>>;
|
||||
claimSubagentTaskResult(params: {
|
||||
userId: string;
|
||||
conversationId: string;
|
||||
taskId: string;
|
||||
claimId: string;
|
||||
}): Promise<SubagentTaskResultClaim>;
|
||||
deleteMessagesSince(
|
||||
userId: string,
|
||||
params: { messageId: string; conversationId: string },
|
||||
|
|
@ -518,6 +529,62 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa
|
|||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Assigns one durable terminal child result to the polling invocation that collects
|
||||
* it. The same invocation may re-acquire, so a poll whose response was lost recovers
|
||||
* the result it never received; a different invocation is told it was already
|
||||
* collected rather than handed a second copy.
|
||||
*/
|
||||
async function claimSubagentTaskResult({
|
||||
userId,
|
||||
conversationId,
|
||||
taskId,
|
||||
claimId,
|
||||
}: {
|
||||
userId: string;
|
||||
conversationId: string;
|
||||
taskId: string;
|
||||
claimId: string;
|
||||
}): Promise<SubagentTaskResultClaim> {
|
||||
if (
|
||||
taskId.length === 0 ||
|
||||
taskId.length > 256 ||
|
||||
conversationId.length === 0 ||
|
||||
conversationId.length > 256 ||
|
||||
claimId.length === 0 ||
|
||||
claimId.length > 128
|
||||
) {
|
||||
throw new TypeError('Invalid subagent task result claim');
|
||||
}
|
||||
const Message = mongoose.models.Message as Model<IMessage>;
|
||||
const filter = {
|
||||
user: userId,
|
||||
conversationId,
|
||||
messageId: `${taskId}:assistant`,
|
||||
'subagentTask.status': { $in: ['completed', 'error', 'cancelled'] },
|
||||
};
|
||||
const acquired = await Message.findOneAndUpdate(
|
||||
{
|
||||
...filter,
|
||||
$or: [
|
||||
{ 'subagentTask.resultClaim': { $exists: false } },
|
||||
{ 'subagentTask.resultClaim.claimId': claimId },
|
||||
],
|
||||
},
|
||||
{ $set: { 'subagentTask.resultClaim': { claimId, claimedAt: new Date() } } },
|
||||
{
|
||||
new: true,
|
||||
timestamps: false,
|
||||
projection: { messageId: 1, conversationId: 1, text: 1, subagentTask: 1 },
|
||||
},
|
||||
).lean<IMessage | null>();
|
||||
if (acquired != null) {
|
||||
return { status: 'acquired', message: acquired };
|
||||
}
|
||||
const existing = await Message.exists(filter);
|
||||
return existing == null ? { status: 'not_found' } : { status: 'claimed' };
|
||||
}
|
||||
|
||||
/**
|
||||
* Deletes messages in a conversation since a specific message.
|
||||
*/
|
||||
|
|
@ -655,6 +722,7 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa
|
|||
updateMessageText,
|
||||
updateToolCallResult,
|
||||
updateMessage,
|
||||
claimSubagentTaskResult,
|
||||
deleteMessagesSince,
|
||||
getMessages,
|
||||
getMessage,
|
||||
|
|
|
|||
|
|
@ -707,6 +707,159 @@ describe('User Methods - Database Tests', () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe('subagent admission fence', () => {
|
||||
test('closes admission until the deletion that took the fence releases it', async () => {
|
||||
const user = await User.create({
|
||||
name: 'Subagent Fence',
|
||||
email: 'subagent-fence@example.com',
|
||||
provider: 'local',
|
||||
});
|
||||
const userId = user._id.toString();
|
||||
const fencedUntil = new Date(Date.now() + 60_000);
|
||||
|
||||
await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(true);
|
||||
await methods.fenceSubagentAdmission(userId, 'deletion-a', fencedUntil);
|
||||
await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(false);
|
||||
|
||||
/** Each overlapping deletion holds its own fence, so admission reopens only
|
||||
* once the last one finishes — in either completion order. */
|
||||
await methods.fenceSubagentAdmission(userId, 'deletion-b', fencedUntil);
|
||||
await methods.releaseSubagentAdmission(userId, 'deletion-a');
|
||||
await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(false);
|
||||
|
||||
await methods.releaseSubagentAdmission(userId, 'deletion-b');
|
||||
await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(true);
|
||||
});
|
||||
|
||||
test('keeps admission closed when the later deletion finishes first', async () => {
|
||||
const user = await User.create({
|
||||
name: 'Reverse Fence',
|
||||
email: 'reverse-fence@example.com',
|
||||
provider: 'local',
|
||||
});
|
||||
const userId = user._id.toString();
|
||||
const fencedUntil = new Date(Date.now() + 60_000);
|
||||
|
||||
await methods.fenceSubagentAdmission(userId, 'deletion-a', fencedUntil);
|
||||
await methods.fenceSubagentAdmission(userId, 'deletion-b', fencedUntil);
|
||||
|
||||
/** The deletion that started second finishes first; the first is still running. */
|
||||
await methods.releaseSubagentAdmission(userId, 'deletion-b');
|
||||
await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(false);
|
||||
|
||||
await methods.releaseSubagentAdmission(userId, 'deletion-a');
|
||||
await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(true);
|
||||
});
|
||||
|
||||
test('prunes an expired fence when the next deletion takes one', async () => {
|
||||
const user = await User.create({
|
||||
name: 'Pruned Fence',
|
||||
email: 'pruned-fence@example.com',
|
||||
provider: 'local',
|
||||
});
|
||||
const userId = user._id.toString();
|
||||
|
||||
await methods.fenceSubagentAdmission(userId, 'abandoned', new Date(Date.now() - 1));
|
||||
await methods.fenceSubagentAdmission(userId, 'deletion-a', new Date(Date.now() + 60_000));
|
||||
|
||||
const stored = await User.findById(userId).select('+subagentAdmissionFences').lean();
|
||||
expect(stored?.subagentAdmissionFences).toHaveLength(1);
|
||||
expect(stored?.subagentAdmissionFences?.[0]?.token).toBe('deletion-a');
|
||||
});
|
||||
|
||||
test('reopens admission once an abandoned fence expires', async () => {
|
||||
const user = await User.create({
|
||||
name: 'Expired Fence',
|
||||
email: 'expired-fence@example.com',
|
||||
provider: 'local',
|
||||
});
|
||||
const userId = user._id.toString();
|
||||
|
||||
await methods.fenceSubagentAdmission(userId, 'deletion-a', new Date(Date.now() - 1));
|
||||
await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(true);
|
||||
});
|
||||
|
||||
test('refuses an excess deletion instead of discarding an active fence', async () => {
|
||||
const user = await User.create({
|
||||
name: 'Saturated Fence',
|
||||
email: 'saturated-fence@example.com',
|
||||
provider: 'local',
|
||||
});
|
||||
const userId = user._id.toString();
|
||||
const fencedUntil = new Date(Date.now() + 60_000);
|
||||
|
||||
for (let index = 0; index < 32; index += 1) {
|
||||
await methods.fenceSubagentAdmission(userId, `deletion-${index}`, fencedUntil);
|
||||
}
|
||||
await expect(
|
||||
methods.fenceSubagentAdmission(userId, 'deletion-overflow', fencedUntil),
|
||||
).rejects.toThrow('Too many concurrent bulk deletions');
|
||||
|
||||
/** The first deletion still owns its fence, so admission stays closed for it. */
|
||||
const stored = await User.findById(userId).select('+subagentAdmissionFences').lean();
|
||||
expect(stored?.subagentAdmissionFences).toHaveLength(32);
|
||||
expect(stored?.subagentAdmissionFences?.[0]?.token).toBe('deletion-0');
|
||||
await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(false);
|
||||
});
|
||||
|
||||
test('invalidates the cached auth document when a refused fence still pruned', async () => {
|
||||
enableAuthUserDocCache();
|
||||
const user = await User.create({
|
||||
name: 'Refused Fence',
|
||||
email: 'refused-fence@example.com',
|
||||
provider: 'local',
|
||||
});
|
||||
const userId = user._id?.toString() ?? '';
|
||||
const indexKey = `${AUTH_USER_DOC_BY_ID_PREFIX}:${userId}`;
|
||||
const fencedUntil = new Date(Date.now() + 60_000);
|
||||
/** A saturated owner that has since abandoned one fence: the next attempt prunes
|
||||
* the expired entry and is then refused by the cap, so the two writes disagree. */
|
||||
await User.updateOne(
|
||||
{ _id: userId },
|
||||
{
|
||||
$set: {
|
||||
subagentAdmissionFences: [
|
||||
...Array.from({ length: 32 }, (_unused, index) => ({
|
||||
token: `deletion-${index}`,
|
||||
expiresAt: fencedUntil,
|
||||
})),
|
||||
{ token: 'abandoned', expiresAt: new Date(Date.now() - 1) },
|
||||
],
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
const cache = {
|
||||
get: jest.fn().mockResolvedValue(['auth-cache-key-a']),
|
||||
delete: jest.fn().mockResolvedValue(true),
|
||||
};
|
||||
const methodsWithCache = createUserMethods(mongoose, {
|
||||
getCache: jest.fn().mockReturnValue(cache),
|
||||
});
|
||||
|
||||
await expect(
|
||||
methodsWithCache.fenceSubagentAdmission(userId, 'deletion-overflow', fencedUntil),
|
||||
).rejects.toThrow('Too many concurrent bulk deletions');
|
||||
const stored = await User.findById(userId).select('+subagentAdmissionFences').lean();
|
||||
expect(stored?.subagentAdmissionFences).toHaveLength(32);
|
||||
/** The prune committed, so leaving the cached document in place would serve the
|
||||
* pruned fence until its own TTL expired. */
|
||||
expect(cache.delete).toHaveBeenCalledWith('auth-cache-key-a');
|
||||
expect(cache.delete).toHaveBeenCalledWith(indexKey);
|
||||
});
|
||||
|
||||
test('refuses an unbounded or invalid fence', async () => {
|
||||
const userId = new mongoose.Types.ObjectId().toString();
|
||||
|
||||
await expect(
|
||||
methods.fenceSubagentAdmission(userId, 'deletion-a', new Date(Number.NaN)),
|
||||
).rejects.toThrow('fencedUntil must be a valid Date');
|
||||
await expect(
|
||||
methods.fenceSubagentAdmission(userId, '', new Date(Date.now() + 60_000)),
|
||||
).rejects.toThrow('bounded owner token');
|
||||
});
|
||||
});
|
||||
|
||||
describe('countUsers', () => {
|
||||
test('should count all users', async () => {
|
||||
await User.create([
|
||||
|
|
|
|||
|
|
@ -14,6 +14,8 @@ import { signPayload } from '~/crypto';
|
|||
export const DEFAULT_SESSION_EXPIRY: number = 1000 * 60 * 15;
|
||||
/** Minimum age before an explicitly offline operator may recover an abandoned deletion fence. */
|
||||
export const USER_DELETION_FENCE_STALE_MS: number = 15 * 60_000;
|
||||
/** Bounds concurrent bulk deletions held for one owner at any moment. */
|
||||
const MAX_SUBAGENT_ADMISSION_FENCES = 32;
|
||||
|
||||
interface UserMethodDeps {
|
||||
getCache?: (key: string) => CacheStore | undefined;
|
||||
|
|
@ -129,6 +131,10 @@ export function createUserMethods(
|
|||
) => Promise<'acquired' | 'in_progress' | 'missing'>;
|
||||
cancelAgentTriggerUserDeletion: (userId: string, startedAt: Date) => Promise<boolean>;
|
||||
isAgentTriggerPrincipalActive: (userId: string) => Promise<boolean>;
|
||||
fenceSubagentAdmission: (userId: string, token: string, fencedUntil: Date) => Promise<void>;
|
||||
renewSubagentAdmission: (userId: string, token: string, fencedUntil: Date) => Promise<boolean>;
|
||||
releaseSubagentAdmission: (userId: string, token: string) => Promise<void>;
|
||||
isSubagentOwnerAdmissible: (userId: string) => Promise<boolean>;
|
||||
deleteUserById: (userId: string) => Promise<UserDeleteResult>;
|
||||
updateUserPlugins: (
|
||||
userId: string,
|
||||
|
|
@ -452,6 +458,101 @@ export function createUserMethods(
|
|||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Closes subagent admission for one owner while a bulk conversation deletion drains
|
||||
* its live children. Every concurrent deletion holds its own fence, so admission
|
||||
* reopens only once the last one finishes, in whatever order they complete. Each
|
||||
* fence expires on its own, so a process that dies mid-delete cannot lock the
|
||||
* account out of running subagents, and expired fences are pruned as new ones
|
||||
* arrive rather than accumulating.
|
||||
*/
|
||||
async function fenceSubagentAdmission(
|
||||
userId: string,
|
||||
token: string,
|
||||
fencedUntil: Date,
|
||||
): Promise<void> {
|
||||
if (!(fencedUntil instanceof Date) || !Number.isFinite(fencedUntil.getTime())) {
|
||||
throw new TypeError('fencedUntil must be a valid Date');
|
||||
}
|
||||
if (token.length === 0 || token.length > 128) {
|
||||
throw new TypeError('A subagent admission fence needs a bounded owner token');
|
||||
}
|
||||
const User = mongoose.models.User;
|
||||
/** Plain update operators only: DocumentDB rejects pipeline-form updates, and
|
||||
* this runs before any deletion, so using one would fail the whole endpoint. */
|
||||
await User.updateOne(
|
||||
{ _id: userId },
|
||||
{ $pull: { subagentAdmissionFences: { expiresAt: { $lte: new Date() } } } },
|
||||
{ timestamps: false },
|
||||
);
|
||||
try {
|
||||
/** Admitted only while the owner is under the concurrent-deletion cap. Dropping
|
||||
* an active fence to make room would reopen admission for a deletion that is
|
||||
* still running, so an excess deletion is refused instead. */
|
||||
const fenced = await User.updateOne(
|
||||
{
|
||||
_id: userId,
|
||||
[`subagentAdmissionFences.${MAX_SUBAGENT_ADMISSION_FENCES - 1}`]: { $exists: false },
|
||||
},
|
||||
{ $push: { subagentAdmissionFences: { token, expiresAt: fencedUntil } } },
|
||||
{ timestamps: false },
|
||||
);
|
||||
if (fenced.matchedCount !== 1) {
|
||||
throw new Error('Too many concurrent bulk deletions are already fencing this owner.');
|
||||
}
|
||||
} finally {
|
||||
/** The prune above commits on its own, so a refused or failed fence still leaves
|
||||
* the cached document describing entries the collection no longer holds. */
|
||||
await invalidateAuthUserDocCache(userId);
|
||||
}
|
||||
}
|
||||
|
||||
/** Extends only this deletion's own fence while its work is still running. */
|
||||
async function renewSubagentAdmission(
|
||||
userId: string,
|
||||
token: string,
|
||||
fencedUntil: Date,
|
||||
): Promise<boolean> {
|
||||
if (!(fencedUntil instanceof Date) || !Number.isFinite(fencedUntil.getTime())) {
|
||||
throw new TypeError('fencedUntil must be a valid Date');
|
||||
}
|
||||
const User = mongoose.models.User;
|
||||
const result = await User.updateOne(
|
||||
{ _id: userId, 'subagentAdmissionFences.token': token },
|
||||
{ $set: { 'subagentAdmissionFences.$.expiresAt': fencedUntil } },
|
||||
{ timestamps: false },
|
||||
);
|
||||
if (result.modifiedCount === 1) {
|
||||
await invalidateAuthUserDocCache(userId);
|
||||
}
|
||||
return result.matchedCount === 1;
|
||||
}
|
||||
|
||||
/** Lifts only this deletion's fence, so an overlapping one keeps admission closed. */
|
||||
async function releaseSubagentAdmission(userId: string, token: string): Promise<void> {
|
||||
const User = mongoose.models.User;
|
||||
const result = await User.updateOne(
|
||||
{ _id: userId },
|
||||
{ $pull: { subagentAdmissionFences: { token } } },
|
||||
{ timestamps: false },
|
||||
);
|
||||
if (result.modifiedCount === 1) {
|
||||
await invalidateAuthUserDocCache(userId);
|
||||
}
|
||||
}
|
||||
|
||||
/** True while this owner may admit a new child: no account deletion, no live fence. */
|
||||
async function isSubagentOwnerAdmissible(userId: string): Promise<boolean> {
|
||||
const User = mongoose.models.User;
|
||||
return (
|
||||
(await User.exists({
|
||||
_id: userId,
|
||||
agentTriggerDeletionStartedAt: { $exists: false },
|
||||
subagentAdmissionFences: { $not: { $elemMatch: { expiresAt: { $gt: new Date() } } } },
|
||||
})) != null
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates a JWT token for a given user.
|
||||
* @param user - The user object
|
||||
|
|
@ -707,6 +808,10 @@ export function createUserMethods(
|
|||
recoverStaleAgentTriggerUserDeletion,
|
||||
cancelAgentTriggerUserDeletion,
|
||||
isAgentTriggerPrincipalActive,
|
||||
fenceSubagentAdmission,
|
||||
renewSubagentAdmission,
|
||||
releaseSubagentAdmission,
|
||||
isSubagentOwnerAdmissible,
|
||||
deleteUserById,
|
||||
updateUserPlugins,
|
||||
toggleUserMemories,
|
||||
|
|
|
|||
|
|
@ -144,6 +144,14 @@ const messageSchema: Schema<IMessage> = new Schema(
|
|||
enum: ['running', 'completed', 'error', 'cancelled'],
|
||||
required: true,
|
||||
},
|
||||
resultClaim: {
|
||||
type: {
|
||||
claimId: { type: String, required: true },
|
||||
claimedAt: { type: Date, required: true },
|
||||
},
|
||||
_id: false,
|
||||
default: undefined,
|
||||
},
|
||||
},
|
||||
_id: false,
|
||||
select: false,
|
||||
|
|
|
|||
|
|
@ -135,6 +135,17 @@ const userSchema: Schema<IUser> = new Schema<IUser>(
|
|||
type: Date,
|
||||
select: false,
|
||||
},
|
||||
subagentAdmissionFences: {
|
||||
type: [
|
||||
{
|
||||
token: { type: String, required: true },
|
||||
expiresAt: { type: Date, required: true },
|
||||
},
|
||||
],
|
||||
_id: false,
|
||||
select: false,
|
||||
default: undefined,
|
||||
},
|
||||
personalization: {
|
||||
type: {
|
||||
memories: {
|
||||
|
|
|
|||
|
|
@ -7,6 +7,12 @@ export interface ISubagentThreadLease {
|
|||
expiresAt: Date;
|
||||
}
|
||||
|
||||
export interface IActiveSubagentThreadLease {
|
||||
conversationId: string;
|
||||
parentConversationId: string;
|
||||
taskId: string;
|
||||
}
|
||||
|
||||
export interface ISubagentThreadReservation {
|
||||
conversation: IConversation;
|
||||
created: boolean;
|
||||
|
|
|
|||
|
|
@ -53,6 +53,11 @@ export interface IMessage extends Document {
|
|||
attemptKey: string;
|
||||
requestFingerprint?: string;
|
||||
status: 'running' | 'completed' | 'error' | 'cancelled';
|
||||
/** Records which polling invocation collected this terminal result. */
|
||||
resultClaim?: {
|
||||
claimId: string;
|
||||
claimedAt: Date;
|
||||
};
|
||||
};
|
||||
contextMeta?: {
|
||||
calibrationRatio?: number;
|
||||
|
|
|
|||
|
|
@ -55,6 +55,11 @@ export interface IUser extends Document {
|
|||
termsAcceptedAt?: Date | null;
|
||||
/** Internal fence that prevents agent-trigger admission during account deletion. */
|
||||
agentTriggerDeletionStartedAt?: Date;
|
||||
/** Expiring fences closing subagent admission while bulk deletions drain. */
|
||||
subagentAdmissionFences?: Array<{
|
||||
token: string;
|
||||
expiresAt: Date;
|
||||
}>;
|
||||
personalization?: {
|
||||
memories?: boolean;
|
||||
statefulCodeEnvironment?: StatefulCodeEnvironment;
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue