@amaster.ai/dsh-a2a 0.1.37 → 0.1.39

package/lib/inbox.js

@@ -5,7 +5,7 @@ import { RequestMalformedError } from '@a2a-js/sdk/errors';
 import { AgentEvent, ServerCallContext, STATE_HEADERS_KEY, } from '@a2a-js/sdk/server';
 import { createUserMessage } from '@deepseek-ai/dsh-llm/message';
 import { InboxChanges } from './inbox-changes.js';
-import { commitInboxInput, markInboxInputPending } from './task-store.js';
+import { commitInboxInput, findRecoveryCandidates, isTerminalTask, markInboxInputPending, taskScopeKey, } from './task-store.js';
 import { agentTextMessage, SessionTranslator, terminalStatusUpdate } from './translator.js';
 export const INBOX_EXTENSION = 'https://github.com/TGYD-helige/dsh-plugins/blob/main/packages/dsh-a2a/README.md';
 const RETAINED_INPUTS = `${INBOX_EXTENSION}#retained-inputs`;
@@ -131,11 +131,14 @@ export class NativeInbox {
             fiber = parent;
         }
     }
-    content(message, contextId) {
-        return this.bridge.buildContent(message, contextId);
+    content(message, contextId, signal) {
+        return this.bridge.buildContent(message, contextId, signal);
     }
-    async resolve(contextId) {
-        return (await this.context(contextId)).entry.handle.agent;
+    async resolve(contextId, context) {
+        return (await this.context(contextId, context)).entry.handle.agent;
+    }
+    recoveryCandidates(context = new ServerCallContext()) {
+        return findRecoveryCandidates(this.store, context);
     }
     async get(contextId, page) {
         validateInboxPage(page);
@@ -255,7 +258,7 @@ export class NativeInbox {
                 !message.parts.length ||
                 message.parts.some((part) => !part.content))
                 throw new RequestMalformedError('A user Message with unchanged messageId, taskId and contextId is required.');
-            const content = await this.content(message, contextId);
+            const content = await this.content(message, contextId, task.controller.signal);
             if (this.stopping ||
                 this.bridge.isClearing(contextId) ||
                 this.queued(run, id).message !== queued.message)
@@ -523,7 +526,7 @@ export class NativeInbox {
             if (this.stopping || this.bridge.isClearing(contextId))
                 throw new Error('Context is unavailable.');
             const previous = await this.store.load(taskId, context);
-            if (!previous || previous.contextId !== contextId || !this.terminal(previous))
+            if (!previous || previous.contextId !== contextId || !isTerminalTask(previous))
                 throw new Error('Resume requires a terminal retained task.');
             const retained = retainedSnapshot(previous);
             if (typeof retained !== 'string')
@@ -658,7 +661,7 @@ export class NativeInbox {
             if (message.source.promotedFrom === message.source.taskId)
                 continue;
             const former = await this.store.load(message.source.promotedFrom, context);
-            if (!former || this.terminal(former))
+            if (!former || isTerminalTask(former))
                 continue;
             former.status = {
                 state: TaskState.TASK_STATE_CANCELED,
@@ -694,7 +697,7 @@ export class NativeInbox {
                     continue;
                 // Terminal snapshots belong to explicit resume(); do not consume their
                 // retained metadata during ordinary context recovery.
-                if (this.terminal(task))
+                if (isTerminalTask(task))
                     continue;
                 snapshots.push(task);
                 for (const input of JSON.parse(retained)) {
@@ -797,18 +800,25 @@ export class NativeInbox {
             await this.checkpoint(task, context);
         }
     }
-    context(contextId, context = new ServerCallContext()) {
+    context(contextId, context) {
         if (this.stopping)
             return Promise.reject(new Error('Native inbox is unloading.'));
+        const checkScope = (run) => {
+            if (context && run.scope !== taskScopeKey(context))
+                throw new Error('Context is already loaded in a different tenant or owner scope.');
+            return run;
+        };
         const existing = this.contexts.get(contextId);
         if (existing)
-            return Promise.resolve(existing);
+            return Promise.resolve(existing).then(checkScope);
         let creating = this.resolving.get(contextId);
         if (!creating) {
             creating = (async () => {
+                const call = context ?? new ServerCallContext();
                 const { entry } = await this.bridge.ensureTask(randomUUID(), contextId);
                 const run = {
                     entry,
+                    scope: taskScopeKey(call),
                     paused: false,
                     tasks: new Map(),
                     claimed: new Set(),
@@ -819,23 +829,26 @@ export class NativeInbox {
                     rolledBack: new Set(),
                     retained: { 'next-step': [], 'next-turn': [] },
                 };
-                await this.recoverRetained(run, context);
-                await this.recoverPromotions(run, context);
+                await this.recoverRetained(run, call);
+                await this.recoverPromotions(run, call);
                 const snapshots = new Map();
                 for (const message of this.pending(run)) {
                     if (message.source.kind !== 'a2a' || snapshots.has(message.source.taskId))
                         continue;
-                    const snapshot = await this.store.load(message.source.taskId, context);
+                    const snapshot = await this.store.load(message.source.taskId, call);
                     if (!snapshot || snapshot.contextId !== contextId)
                         throw new Error('Queued A2A task is not accessible.');
-                    if (this.terminal(snapshot))
+                    if (isTerminalTask(snapshot))
                         entry.handle.agent.inbox.remove(message.id);
                     else
                         snapshots.set(snapshot.id, snapshot);
                 }
                 await this.flush(entry.handle.agent.session);
+                const committed = snapshots.size
+                    ? this.committedInputs(run, new Set(snapshots.keys()))
+                    : new Map();
                 for (const snapshot of snapshots.values())
-                    this.restoreTask(run, snapshot, context);
+                    this.restoreTask(run, snapshot, call, committed.get(snapshot.id) ?? []);
                 run.retained = structuredClone({
                     'next-step': [...entry.handle.agent.inbox.nextStep],
                     'next-turn': [...entry.handle.agent.inbox.nextTurn],
@@ -848,7 +861,7 @@ export class NativeInbox {
             })().finally(() => this.resolving.delete(contextId));
             this.resolving.set(contextId, creating);
         }
-        return creating;
+        return creating.then(checkScope);
     }
     observe(phase, run, fields = {}) {
         try {
@@ -951,12 +964,10 @@ export class NativeInbox {
             context: new ServerCallContext({ tenant: context.tenant, user: context.user }),
         };
     }
-    controls(message, context) {
+    controls(message) {
         const value = message.metadata?.[INBOX_EXTENSION];
         if (value === undefined)
             return { target: 'next-turn', wakeup: true };
-        if (!context.requestedExtensions?.includes(INBOX_EXTENSION))
-            throw new Error('Inbox controls require A2A-Extensions negotiation.');
         if (!value ||
             typeof value !== 'object' ||
             Array.isArray(value) ||
@@ -966,7 +977,6 @@ export class NativeInbox {
             throw new Error('Invalid inbox target.');
         if (value.wakeup !== undefined && typeof value.wakeup !== 'boolean')
             throw new Error('Invalid inbox wakeup.');
-        context.addActivatedExtension(INBOX_EXTENSION);
         return {
             target: value.target === 'next-step' ? 'next-step' : 'next-turn',
             wakeup: value.wakeup !== false,
@@ -976,7 +986,7 @@ export class NativeInbox {
     async beforeMessage(message, context) {
         if (!message.messageId)
             throw new Error('message.messageId is required.');
-        const controls = this.controls(message, context);
+        const controls = this.controls(message);
         if (!message.taskId) {
             if (controls.target === 'next-step')
                 throw new Error('Steering must address an active task.');
@@ -985,7 +995,7 @@ export class NativeInbox {
         const snapshot = await this.store.load(message.taskId, context);
         if (!snapshot ||
             (message.contextId && snapshot.contextId !== message.contextId) ||
-            this.terminal(snapshot))
+            isTerminalTask(snapshot))
             return;
         const run = await this.context(snapshot.contextId, context);
         let task = run.tasks.get(snapshot.id);
@@ -1016,7 +1026,7 @@ export class NativeInbox {
     /** Reattach a cold task before the SDK decides whether cancellation has a live executor. */
     async beforeTask(taskId, context) {
         const snapshot = await this.store.load(taskId, context);
-        if (!snapshot || this.terminal(snapshot))
+        if (!snapshot || isTerminalTask(snapshot))
             return;
         const run = await this.context(snapshot.contextId, context);
         const task = run.tasks.get(taskId) ?? this.restoreTask(run, snapshot, context);
@@ -1034,15 +1044,6 @@ export class NativeInbox {
             task.status?.state === TaskState.TASK_STATE_REJECTED)
             return 'failed';
     }
-    terminal(task) {
-        return (!!task.status &&
-            [
-                TaskState.TASK_STATE_COMPLETED,
-                TaskState.TASK_STATE_CANCELED,
-                TaskState.TASK_STATE_FAILED,
-                TaskState.TASK_STATE_REJECTED,
-            ].includes(task.status.state));
-    }
     async admit(run, task, a2aMessage, context, target, wakeup, accepted) {
         if (run.paused)
             throw new Error('Inbox is paused.');
@@ -1060,7 +1061,7 @@ export class NativeInbox {
         const headers = context.state.get(STATE_HEADERS_KEY);
         const work = run.admissionTail.then(async () => {
             abort.signal.throwIfAborted();
-            const content = await this.content(a2aMessage, run.entry.sessionId);
+            const content = await this.content(a2aMessage, run.entry.sessionId, abort.signal);
             abort.signal.throwIfAborted();
             const message = createUserMessage({
                 content,
@@ -1149,7 +1150,7 @@ export class NativeInbox {
         let task;
         let ownsTask = false;
         try {
-            const controls = this.controls(userMessage, context);
+            const controls = this.controls(userMessage);
             const run = await this.context(contextId, context);
             task = run.tasks.get(taskId);
             if (!task) {
@@ -1212,13 +1213,23 @@ export class NativeInbox {
             bus.finished();
         }
     }
-    restoreTask(run, snapshot, context) {
+    committedInputs(run, taskIds) {
+        const inputs = new Map();
         for (const event of run.entry.handle.agent.session.snapshotEvents()) {
-            if (event.type === 'user/message' &&
-                event.data.source.kind === 'a2a' &&
-                event.data.source.taskId === snapshot.id)
-                commitInboxInput(snapshot.history, Message.fromJSON(JSON.parse(event.data.source.message)));
+            if (event.type !== 'user/message' || event.data.source.kind !== 'a2a')
+                continue;
+            const source = event.data.source;
+            if (!taskIds.has(source.taskId))
+                continue;
+            const messages = inputs.get(source.taskId) ?? [];
+            messages.push(Message.fromJSON(JSON.parse(source.message)));
+            inputs.set(source.taskId, messages);
         }
+        return inputs;
+    }
+    restoreTask(run, snapshot, context, committed = this.committedInputs(run, new Set([snapshot.id])).get(snapshot.id) ?? []) {
+        for (const message of committed)
+            commitInboxInput(snapshot.history, message);
         if (!this.server)
             throw new Error('A2A task recovery requires the running server.');
         const bus = this.server.restoreTask(snapshot, context);

package/lib/stores/filesystem.js

@@ -3,7 +3,8 @@ import { createHash, randomUUID } from 'node:crypto';
 import { mkdir, open, readdir, readFile, rename, rm } from 'node:fs/promises';
 import { join, resolve } from 'node:path';
 import { Task } from '@a2a-js/sdk';
-import { listShells, taskScopeKey } from '../task-store.js';
+import { recoveryContextIds, taskScopeKey } from '../task-store.js';
+import { LIST_PREFIX_BYTES, loadTaskPage, serializeTaskSnapshot, snapshotListShell, taskListShell, } from './listing.js';
 export class FilesystemTaskStore {
     root;
     constructor(config) {
@@ -24,7 +25,7 @@ export class FilesystemTaskStore {
         try {
             const file = await open(temporary, 'wx', 0o600);
             try {
-                await file.writeFile(JSON.stringify(Task.toJSON(task)));
+                await file.writeFile(serializeTaskSnapshot(task, Task.toJSON(task)));
                 await file.sync();
             }
             finally {
@@ -59,7 +60,7 @@ export class FilesystemTaskStore {
     async delete(id, context) {
         await rm(this.path(id, context), { force: true });
     }
-    async list(params, context) {
+    async shells(context) {
         let files;
         try {
             files = await readdir(this.directory(context));
@@ -69,18 +70,42 @@ export class FilesystemTaskStore {
                 throw error;
             files = [];
         }
-        const tasks = [];
-        for (const file of files) {
-            if (!/^[a-f0-9]{64}\.json$/.test(file))
-                continue;
-            try {
-                tasks.push(Task.fromJSON(JSON.parse(await readFile(join(this.directory(context), file), 'utf8'))));
-            }
-            catch (error) {
-                if (error.code !== 'ENOENT')
-                    throw error;
-            }
+        const snapshots = files.filter((file) => /^[a-f0-9]{64}\.json$/.test(file));
+        const shells = [];
+        for (let start = 0; start < snapshots.length; start += 16) {
+            const batch = await Promise.all(snapshots.slice(start, start + 16).map(async (file) => {
+                const path = join(this.directory(context), file);
+                try {
+                    const handle = await open(path, 'r');
+                    let shell;
+                    try {
+                        const prefix = Buffer.alloc(LIST_PREFIX_BYTES);
+                        const { bytesRead } = await handle.read(prefix, 0, prefix.length, 0);
+                        shell = snapshotListShell(prefix.toString('utf8', 0, bytesRead));
+                        // Read legacy or oversized headers through the same open inode.
+                        if (!shell)
+                            shell = taskListShell(Task.fromJSON(JSON.parse(await handle.readFile('utf8'))));
+                    }
+                    finally {
+                        await handle.close();
+                    }
+                    return shell;
+                }
+                catch (error) {
+                    if (error.code !== 'ENOENT')
+                        throw error;
+                }
+            }));
+            for (const shell of batch)
+                if (shell)
+                    shells.push(shell);
         }
-        return listShells(tasks, params);
+        return shells;
+    }
+    async recoveryCandidates(context) {
+        return recoveryContextIds(await this.shells(context));
+    }
+    async list(params, context) {
+        return loadTaskPage(await this.shells(context), params, (ids) => Promise.all(ids.map((id) => this.load(id, context))));
     }
 }

package/lib/stores/gcs.js

@@ -17,8 +17,11 @@ import { tmpdir } from 'node:os';
 import { join } from 'node:path';
 import { pipeline } from 'node:stream/promises';
 import { promisify } from 'node:util';
-import { gunzipSync, gzipSync } from 'node:zlib';
-import { listShells, taskScopeKey } from '../task-store.js';
+import { gunzip, gzip } from 'node:zlib';
+import { recoveryContextIds, taskScopeKey } from '../task-store.js';
+import { loadTaskPage, parseListShell, taskListShell } from './listing.js';
+const compress = promisify(gzip);
+const decompress = promisify(gunzip);
 export class GcsTaskStore {
     config;
     bucket = null;
@@ -35,19 +38,29 @@ export class GcsTaskStore {
     async save(task, context) {
         if (!this.bucket)
             return;
+        const shell = JSON.stringify(taskListShell(task));
         await this.bucket
             .file(`${this.prefix}/${taskScopeKey(context)}/${task.id}/metadata.json.gz`)
-            .save(gzipSync(JSON.stringify(task)), { contentType: 'application/gzip', resumable: false });
+            .save(await compress(JSON.stringify(task)), {
+            contentType: 'application/gzip',
+            resumable: false,
+            // GCS custom metadata is bounded; oversized identifiers use the legacy read path.
+            metadata: { metadata: Buffer.byteLength(shell) <= 4096 ? { dshTaskList: shell } : {} },
+        });
     }
     async load(taskId, context) {
         if (!this.bucket)
             return undefined;
         const file = this.bucket.file(`${this.prefix}/${taskScopeKey(context)}/${taskId}/metadata.json.gz`);
-        const [exists] = await file.exists();
-        if (!exists)
-            return undefined;
-        const [buf] = await file.download();
-        return JSON.parse(gunzipSync(buf).toString('utf8'));
+        try {
+            const [buf] = await file.download();
+            return JSON.parse((await decompress(buf)).toString('utf8'));
+        }
+        catch (error) {
+            if (error.code === 404)
+                return undefined;
+            throw error;
+        }
     }
     async delete(taskId, context) {
         if (!this.bucket)
@@ -61,21 +74,54 @@ export class GcsTaskStore {
                 .delete({ ignoreNotFound: true }),
         ]);
     }
-    async list(params, context) {
+    async shells(context) {
         if (!this.bucket)
-            return { tasks: [], nextPageToken: '', pageSize: 0, totalSize: 0 };
+            return [];
         const [files] = await this.bucket.getFiles({
             prefix: `${this.prefix}/${taskScopeKey(context)}/`,
         });
         const shells = [];
         const metadata = files.filter((file) => file.name.endsWith('/metadata.json.gz'));
-        // Bound in-flight downloads while avoiding one network round-trip per task.
         for (let start = 0; start < metadata.length; start += 16) {
-            const buffers = await Promise.all(metadata.slice(start, start + 16).map((file) => file.download()));
-            for (const [buf] of buffers)
-                shells.push(JSON.parse(gunzipSync(buf).toString('utf8')));
+            const batch = await Promise.all(metadata.slice(start, start + 16).map(async (file) => {
+                const indexed = parseListShell(file.metadata?.metadata?.dshTaskList);
+                if (indexed)
+                    return indexed;
+                // Older snapshots have no summary; keep them visible without migration.
+                try {
+                    const [buf] = await file.download();
+                    return taskListShell(JSON.parse((await decompress(buf)).toString('utf8')));
+                }
+                catch (error) {
+                    if (error.code === 404)
+                        return;
+                    throw error;
+                }
+            }));
+            for (const shell of batch)
+                if (shell)
+                    shells.push(shell);
         }
-        return listShells(shells, params);
+        return shells;
+    }
+    async recoveryCandidates(context) {
+        if (!this.bucket)
+            throw new Error('[dsh-a2a] GCS TaskStore is not initialized.');
+        return recoveryContextIds(await this.shells(context));
+    }
+    async list(params, context) {
+        if (!this.bucket)
+            return { tasks: [], nextPageToken: '', pageSize: 0, totalSize: 0 };
+        return loadTaskPage(await this.shells(context), params, (ids) => Promise.all(ids.map(async (id) => {
+            try {
+                return await this.load(id, context);
+            }
+            catch (error) {
+                if (error.code === 404)
+                    return;
+                throw error;
+            }
+        })));
     }
     /**
      * Tar + gzip the workspace directory to `tasks/{scope}/{taskId}/workspace.tar.gz`.

package/lib/stores/listing.js

@@ -0,0 +1,63 @@
+import { listShells } from '../task-store.js';
+export const LIST_PREFIX_BYTES = 8192;
+const HEADER = '{"_dshTaskList":';
+/** Only the fields needed by listShells; never retain history or status-message content. */
+export function taskListShell(task) {
+    return {
+        id: task.id,
+        contextId: task.contextId,
+        status: task.status && {
+            state: task.status.state,
+            message: undefined,
+            timestamp: task.status.timestamp,
+        },
+        history: [],
+        artifacts: [],
+        metadata: undefined,
+    };
+}
+export function parseListShell(value) {
+    if (typeof value !== 'string')
+        return;
+    try {
+        const shell = JSON.parse(value);
+        if (!shell ||
+            typeof shell.id !== 'string' ||
+            typeof shell.contextId !== 'string' ||
+            (shell.status !== undefined &&
+                (!shell.status ||
+                    !Number.isInteger(shell.status.state) ||
+                    (shell.status.timestamp !== undefined && typeof shell.status.timestamp !== 'string'))))
+            return;
+        return taskListShell(shell);
+    }
+    catch {
+        return;
+    }
+}
+/** A JSON-compatible header shares the snapshot's atomic write, including on older readers. */
+export function serializeTaskSnapshot(task, json = task) {
+    return `${HEADER}${JSON.stringify(taskListShell(task))},\n${JSON.stringify(json).slice(1)}`;
+}
+export function snapshotListShell(prefix) {
+    if (!prefix.startsWith(HEADER))
+        return;
+    const newline = prefix.indexOf('\n');
+    if (newline < 0 || prefix[newline - 1] !== ',')
+        return;
+    return parseListShell(prefix.slice(HEADER.length, newline - 1));
+}
+/** Hydrate only the selected page, with bounded I/O and the existing artifact semantics. */
+export async function loadTaskPage(shells, params, load) {
+    const page = listShells(shells, params);
+    const tasks = [];
+    for (let start = 0; start < page.tasks.length; start += 16) {
+        const batch = await load(page.tasks.slice(start, start + 16).map((task) => task.id));
+        for (const task of batch) {
+            // A task can expire or be deleted between enumeration and hydration.
+            if (task)
+                tasks.push(params.includeArtifacts ? task : { ...task, artifacts: [] });
+        }
+    }
+    return { ...page, tasks };
+}

package/lib/stores/redis.js

@@ -6,7 +6,12 @@
  * snapshots arrive from SanitizedTaskStore at state transitions;
  * dsh-storage separately owns the full conversation log (ai_messages).
  */
-import { listShells, taskScopeKey } from '../task-store.js';
+import { recoveryContextIds, taskScopeKey } from '../task-store.js';
+import { LIST_PREFIX_BYTES, loadTaskPage, serializeTaskSnapshot, snapshotListShell, taskListShell, } from './listing.js';
+function parseSnapshot(raw) {
+    const { _dshTaskList: _, ...task } = JSON.parse(raw);
+    return task;
+}
 export class RedisTaskStore {
     config;
     redis = null;
@@ -27,34 +32,72 @@ export class RedisTaskStore {
     async save(task, context) {
         if (!this.redis)
             return;
-        await this.redis.set(this.taskKey(task.id, context), JSON.stringify(task), 'EX', this.ttl);
+        await this.redis.set(this.taskKey(task.id, context), serializeTaskSnapshot(task), 'EX', this.ttl);
     }
     async load(taskId, context) {
         if (!this.redis)
             return undefined;
         const raw = await this.redis.get(this.taskKey(taskId, context));
-        return raw ? JSON.parse(raw) : undefined;
+        return raw ? parseSnapshot(raw) : undefined;
     }
     async delete(taskId, context) {
         await this.redis?.del(this.taskKey(taskId, context));
     }
-    async list(params, context) {
+    async shells(context) {
         if (!this.redis)
-            return { tasks: [], nextPageToken: '', pageSize: 0, totalSize: 0 };
-        // Task shells are few and TTL-bound — a paged SCAN over the prefix is ample.
+            return [];
+        // SCAN only names; bounded prefixes supply the lightweight list summaries.
         const shells = [];
+        const seen = new Set();
         let cursor = '0';
         do {
-            const [next, keys] = await this.redis.scan(cursor, 'MATCH', `${this.prefix}:tasks:${taskScopeKey(context)}:*`, 'COUNT', 100);
+            const [next, scanned] = await this.redis.scan(cursor, 'MATCH', `${this.prefix}:tasks:${taskScopeKey(context)}:*`, 'COUNT', 100);
             cursor = next;
+            // SCAN may repeat keys during rehashing, including within one batch.
+            const keys = scanned.filter((key) => {
+                if (seen.has(key))
+                    return false;
+                seen.add(key);
+                return true;
+            });
             if (keys.length) {
-                const raws = await this.redis.mget(...keys);
-                for (const raw of raws)
-                    if (raw)
-                        shells.push(JSON.parse(raw));
+                const prefixes = await this.redis
+                    .pipeline(keys.map((key) => ['getrange', key, 0, LIST_PREFIX_BYTES - 1]))
+                    .exec();
+                const legacy = [];
+                prefixes.forEach(([error, prefix], index) => {
+                    if (error)
+                        throw error;
+                    if (!prefix)
+                        return; // Expired since SCAN.
+                    const shell = snapshotListShell(prefix);
+                    if (shell)
+                        shells.push(shell);
+                    else
+                        legacy.push(keys[index]);
+                });
+                if (legacy.length) {
+                    const raws = await this.redis.mget(...legacy);
+                    for (const raw of raws)
+                        if (raw)
+                            shells.push(taskListShell(parseSnapshot(raw)));
+                }
             }
         } while (cursor !== '0');
-        return listShells(shells, params);
+        return shells;
+    }
+    async recoveryCandidates(context) {
+        if (!this.redis)
+            throw new Error('[dsh-a2a] Redis TaskStore is not initialized.');
+        return recoveryContextIds(await this.shells(context));
+    }
+    async list(params, context) {
+        if (!this.redis)
+            return { tasks: [], nextPageToken: '', pageSize: 0, totalSize: 0 };
+        return loadTaskPage(await this.shells(context), params, async (ids) => {
+            const raws = await this.redis.mget(...ids.map((id) => this.taskKey(id, context)));
+            return raws.map((raw) => (raw ? parseSnapshot(raw) : undefined));
+        });
     }
     async close() {
         await this.redis?.quit?.();

package/lib/task-store.js

@@ -52,6 +52,40 @@ export function sanitizeTask(task) {
         ? { ...task, history: task.history.filter((message) => message.messageId !== current) }
         : task;
 }
+export function isTerminalTask(task) {
+    return [
+        TaskState.TASK_STATE_COMPLETED,
+        TaskState.TASK_STATE_CANCELED,
+        TaskState.TASK_STATE_FAILED,
+        TaskState.TASK_STATE_REJECTED,
+    ].includes(task.status?.state ?? TaskState.TASK_STATE_UNSPECIFIED);
+}
+export function recoveryContextIds(tasks) {
+    return [
+        ...new Set(tasks.filter((task) => !isTerminalTask(task)).map((task) => task.contextId)),
+    ].sort();
+}
+/** Custom stores retain a complete, if more expensive, discovery path. */
+export async function findRecoveryCandidates(store, context) {
+    if (store.recoveryCandidates)
+        return store.recoveryCandidates(context);
+    const ids = new Set();
+    let pageToken = '';
+    do {
+        const page = await store.list({
+            tenant: '',
+            contextId: '',
+            pageSize: 100,
+            pageToken,
+            status: TaskState.TASK_STATE_UNSPECIFIED,
+            statusTimestampAfter: undefined,
+        }, context);
+        for (const id of recoveryContextIds(page.tasks))
+            ids.add(id);
+        pageToken = page.nextPageToken;
+    } while (pageToken);
+    return [...ids].sort();
+}
 function scopeOf(context) {
     return JSON.stringify([context?.tenant ?? '', context ? resolveUserScope(context) : 'unknown']);
 }
@@ -81,6 +115,9 @@ export class MemoryTaskStore {
     async list(params, context) {
         return listShells([...this.bucket(context).values()], params);
     }
+    async recoveryCandidates(context) {
+        return recoveryContextIds([...this.bucket(context).values()]);
+    }
     async delete(taskId, context) {
         this.bucket(context).delete(taskId);
     }
@@ -137,13 +174,7 @@ export class SanitizedTaskStore {
             this.terminalStates.delete(task.id);
         else if (terminal && state !== terminal.state)
             return;
-        if (state !== undefined &&
-            [
-                TaskState.TASK_STATE_COMPLETED,
-                TaskState.TASK_STATE_CANCELED,
-                TaskState.TASK_STATE_FAILED,
-                TaskState.TASK_STATE_REJECTED,
-            ].includes(state)) {
+        if (state !== undefined && isTerminalTask(task)) {
             // Map order matches expiry order: only the expired prefix needs scanning.
             for (const [id, guard] of this.terminalStates) {
                 if (guard.expiresAt > now)
@@ -237,6 +268,9 @@ export class SanitizedTaskStore {
     list(params, context) {
         return this.inner.list(params, context);
     }
+    recoveryCandidates(context) {
+        return findRecoveryCandidates(this.inner, context);
+    }
 }
 /**
  * Shared `listTasks` semantics for the backends: filter by contextId / status /