@amaster.ai/dsh-a2a 0.1.33 → 0.1.35

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/lib/index.js CHANGED
@@ -32,12 +32,24 @@ export const inject = ['agents'];
32
32
  export const Config = Schema.object({
33
33
  inbox: Schema.object({
34
34
  enabled: Schema.boolean().default(false),
35
+ notifications: Schema.boolean()
36
+ .default(false)
37
+ .description('Publish checkpointed inbox invalidations to same-process host observers; SSE subscriptions also activate their own feed.'),
35
38
  shutdown: Schema.union([Schema.const('cancel'), Schema.const('retain-pending')]).default('cancel'),
36
39
  }).description('Use the official DSH inbox for queued turns and steering.'),
37
40
  enabled: Schema.boolean().default(false),
38
41
  host: Schema.string().default('127.0.0.1'),
39
42
  port: Schema.natural().default(41241),
40
43
  basePath: Schema.string().default('/a2a'),
44
+ rest: Schema.boolean().default(false),
45
+ grpc: Schema.object({
46
+ enabled: Schema.boolean().default(false),
47
+ port: Schema.natural().max(65535).default(41242),
48
+ publicUrl: Schema.string().default(''),
49
+ }),
50
+ pushNotifications: Schema.boolean()
51
+ .default(false)
52
+ .description('Enable standard A2A task webhook configuration and delivery through the SDK.'),
41
53
  heartbeat: Schema.object({
42
54
  enabled: Schema.boolean().default(false),
43
55
  intervalMs: Schema.natural().min(1).max(2_147_483_647).default(15_000),
@@ -49,6 +61,12 @@ export const Config = Schema.object({
49
61
  * below it). Empty = <OS temp>/dsh-a2a-uploads.
50
62
  */
51
63
  uploadsDir: Schema.string().default(''),
64
+ urlDownloads: Schema.object({
65
+ allowPrivate: Schema.boolean()
66
+ .default(false)
67
+ .description('allow every private/loopback URL host in a trusted intranet'),
68
+ allowedHosts: Schema.array(Schema.string()).default([]),
69
+ }).description('URL downloads; private hosts require allowPrivate or an allowedHosts entry.'),
52
70
  subagentTools: Schema.object({
53
71
  enabled: Schema.boolean().default(false),
54
72
  }),
@@ -99,9 +117,10 @@ export function apply(ctx, config) {
99
117
  // covers the bind.
100
118
  let ready;
101
119
  let serverPort = 0;
120
+ let grpcPort;
102
121
  ctx.effect(() => {
103
122
  const startup = start();
104
- ready = startup.then(() => ({ port: serverPort }));
123
+ ready = startup.then(() => ({ port: serverPort, grpcPort }));
105
124
  return startup;
106
125
  });
107
126
  return ready;
@@ -109,6 +128,8 @@ export function apply(ctx, config) {
109
128
  const bridge = new A2aBridge(ctx, {
110
129
  cwd: config.cwd,
111
130
  uploadsDir: config.uploadsDir || undefined,
131
+ downloadAllowPrivate: config.urlDownloads?.allowPrivate,
132
+ downloadAllowedHosts: config.urlDownloads?.allowedHosts,
112
133
  agentOptions: {
113
134
  provider: config.agent.provider || undefined,
114
135
  model: config.agent.model || undefined,
@@ -118,7 +139,7 @@ export function apply(ctx, config) {
118
139
  });
119
140
  const store = createTaskStore(config);
120
141
  const inbox = config.inbox?.enabled
121
- ? new NativeInbox(ctx, bridge, store, config.inbox.shutdown === 'retain-pending')
142
+ ? new NativeInbox(ctx, bridge, store, config.inbox.shutdown === 'retain-pending', config.inbox.notifications)
122
143
  : undefined;
123
144
  await store.init?.();
124
145
  if (inbox)
@@ -150,20 +171,27 @@ export function apply(ctx, config) {
150
171
  });
151
172
  const executor = new DshAgentExecutor(bridge, (taskId) => store.suppress(taskId));
152
173
  const server = await startA2aServer({
174
+ inbox,
153
175
  extensions: inbox
154
176
  ? [
155
177
  {
156
178
  uri: INBOX_EXTENSION,
157
- description: 'Native inbox steering and parked admission',
179
+ description: 'Native inbox management, steering and checkpointed invalidation',
158
180
  required: false,
159
- params: undefined,
181
+ params: {
182
+ documentationUrl: 'https://github.com/TGYD-helige/dsh-plugins/blob/main/packages/dsh-a2a/README.md#native-inbox-extension-v1',
183
+ },
160
184
  },
161
185
  ]
162
186
  : [],
163
187
  host: config.host,
164
188
  port: config.port,
165
189
  basePath: config.basePath,
190
+ rest: config.rest,
191
+ grpc: config.grpc,
192
+ extendedAgentCard: ctx.get('a2aExtendedCard'),
166
193
  heartbeat: config.heartbeat,
194
+ pushNotifications: config.pushNotifications,
167
195
  card: {
168
196
  name: config.card.name,
169
197
  description: config.card.description,
@@ -180,6 +208,7 @@ export function apply(ctx, config) {
180
208
  beforeWorkingTaskMessage: (message, headers) => bridge.beforeWorkingTaskMessage(message, headers),
181
209
  });
182
210
  serverPort = server.port;
211
+ grpcPort = server.grpcPort;
183
212
  if (inbox)
184
213
  inbox.server = server;
185
214
  console.log(`[dsh-a2a] A2A endpoint: http://${config.host}:${server.port}${config.basePath}/ ` +
@@ -32,6 +32,11 @@ export function legacyStreamHandler(requestHandler) {
32
32
  return;
33
33
  }
34
34
  const stream = response;
35
+ const resubscribing = req.body.method === 'tasks/resubscribe';
36
+ let closed;
37
+ const close = new Promise((resolve) => (closed = resolve));
38
+ if (resubscribing)
39
+ res.once('close', closed);
35
40
  try {
36
41
  const startStream = () => {
37
42
  if (context.activatedExtensions)
@@ -42,9 +47,12 @@ export function legacyStreamHandler(requestHandler) {
42
47
  res.flushHeaders();
43
48
  };
44
49
  // The SDK always streams resubscribe errors; sends validate before SSE headers.
45
- if (req.body.method === 'tasks/resubscribe')
50
+ if (resubscribing)
46
51
  startStream();
47
- const first = await stream.next();
52
+ const next = () => (resubscribing ? Promise.race([stream.next(), close]) : stream.next());
53
+ const first = await next();
54
+ if (!first)
55
+ return;
48
56
  if (!res.headersSent)
49
57
  startStream();
50
58
  let last;
@@ -60,10 +68,15 @@ export function legacyStreamHandler(requestHandler) {
60
68
  };
61
69
  if (!first.done && write(first.value))
62
70
  return;
63
- for await (const event of stream) {
71
+ while (true) {
72
+ const nextEvent = await next();
73
+ if (!nextEvent || nextEvent.done)
74
+ break;
75
+ if (res.destroyed && !resubscribing)
76
+ continue; // Sends still drain for persistence.
64
77
  if (res.destroyed)
65
- continue; // Keep the SDK consuming results after its client leaves.
66
- if (write(event))
78
+ break;
79
+ if (write(nextEvent.value))
67
80
  return;
68
81
  }
69
82
  // A cold subscription has no live bus: the SDK ends after the snapshot.
@@ -82,6 +95,8 @@ export function legacyStreamHandler(requestHandler) {
82
95
  });
83
96
  }
84
97
  finally {
98
+ if (resubscribing)
99
+ res.off('close', closed);
85
100
  await stream.return(undefined);
86
101
  }
87
102
  }
package/lib/server.js CHANGED
@@ -15,11 +15,15 @@ import { AsyncLocalStorage } from 'node:async_hooks';
15
15
  import { createServer } from 'node:http';
16
16
  import { A2A_PROTOCOL_VERSION, AGENT_CARD_PATH, TaskState, } from '@a2a-js/sdk';
17
17
  import { duplicateInterfacesForLegacy } from '@a2a-js/sdk/compat/v0_3';
18
- import { RequestMalformedError } from '@a2a-js/sdk/errors';
19
- import { DefaultExecutionEventBus, DefaultRequestHandler, ExecutionEventQueue, ResultManager, ServerCallContext, STATE_HEADERS_KEY, } from '@a2a-js/sdk/server';
20
- import { agentCardHandler, jsonRpcHandler, UserBuilder } from '@a2a-js/sdk/server/express';
18
+ import { createLegacyAwarePushNotificationSender } from '@a2a-js/sdk/compat/v0_3/server';
19
+ import { RequestMalformedError, UnsupportedOperationError } from '@a2a-js/sdk/errors';
20
+ import { DefaultExecutionEventBus, DefaultRequestHandler, ExecutionEventQueue, InMemoryPushNotificationStore, ResultManager, ServerCallContext, STATE_HEADERS_KEY, } from '@a2a-js/sdk/server';
21
+ import { agentCardHandler, jsonRpcHandler, restHandler, UserBuilder, } from '@a2a-js/sdk/server/express';
21
22
  import express from 'express';
23
+ import { AdmissionGate } from './admission.js';
24
+ import { inboxRpcHandler } from './inbox-api.js';
22
25
  import { legacyStreamHandler } from './legacy-stream.js';
26
+ import { connectedRequestHandler, subscriptionBus, withStreamConnection, } from './stream-connection.js';
23
27
  export async function startA2aServer(options) {
24
28
  const intervalMs = options.heartbeat?.intervalMs ?? 15_000;
25
29
  if (!Number.isInteger(intervalMs) || intervalMs < 1 || intervalMs > 2_147_483_647) {
@@ -48,19 +52,25 @@ export async function startA2aServer(options) {
48
52
  protocolVersion: A2A_PROTOCOL_VERSION,
49
53
  },
50
54
  ];
55
+ if (options.rest)
56
+ interfaces.push({
57
+ ...interfaces[0],
58
+ url: `${publicUrl}${base}/rest`,
59
+ protocolBinding: 'HTTP+JSON',
60
+ });
51
61
  const card = {
52
62
  name: options.card.name,
53
63
  description: options.card.description,
54
64
  version: options.card.version,
55
65
  // The v0.3 mirror entries let pre-1.0 clients keep working through the
56
66
  // legacyCompat layer (they discover the card without an A2A-Version header).
57
- supportedInterfaces: duplicateInterfacesForLegacy(interfaces, ['JSONRPC']),
67
+ supportedInterfaces: duplicateInterfacesForLegacy(interfaces, ['JSONRPC', 'HTTP+JSON']),
58
68
  provider: undefined,
59
69
  capabilities: {
60
70
  streaming: true,
61
- pushNotifications: false,
71
+ pushNotifications: options.pushNotifications ?? false,
62
72
  extensions: options.extensions ?? [],
63
- extendedAgentCard: false,
73
+ extendedAgentCard: !!options.extendedAgentCard,
64
74
  },
65
75
  securitySchemes: {},
66
76
  securityRequirements: [],
@@ -82,9 +92,25 @@ export async function startA2aServer(options) {
82
92
  // SDK's default manager gives concurrent sends one bus. Give each send its
83
93
  // own bus while forwarding events to a stable task bus for resubscribe/cancel.
84
94
  const approvalReply = new AsyncLocalStorage();
95
+ const admissionReply = new AsyncLocalStorage();
96
+ const admissions = new AdmissionGate();
97
+ async function acquireAdmission(message) {
98
+ const key = message?.taskId || message?.contextId || message?.messageId || '';
99
+ const admission = admissions.acquire(key);
100
+ await admission.ready;
101
+ const state = { release: admission.release, expectedSaves: message?.taskId ? 2 : 1, saves: 0 };
102
+ return {
103
+ release: admission.release,
104
+ run: (work) => admissionReply.run(state, work),
105
+ };
106
+ }
85
107
  const taskBuses = new Map();
108
+ const ownedTasks = new Set();
86
109
  const busManager = {
87
110
  createOrGetByTaskId(taskId) {
111
+ // restoreTask also creates buses, but only message admission proves ownership.
112
+ if (admissionReply.getStore())
113
+ ownedTasks.add(taskId);
88
114
  let task = taskBuses.get(taskId);
89
115
  if (!task) {
90
116
  const broadcast = new DefaultExecutionEventBus();
@@ -113,6 +139,7 @@ export async function startA2aServer(options) {
113
139
  priorMessage.messageId === message.messageId &&
114
140
  before?.$case === 'text' &&
115
141
  delta?.$case === 'text') {
142
+ // shortcut: replay compaction joins text per delta; use chunks if long streams make this measurable.
116
143
  created.replay[created.replay.length - 1] = {
117
144
  ...event,
118
145
  data: {
@@ -164,12 +191,13 @@ export async function startA2aServer(options) {
164
191
  requestBus.on('finished', () => task.requests.delete(requestBus));
165
192
  return requestBus;
166
193
  },
167
- getByTaskId: (taskId) => taskBuses.get(taskId)?.broadcast,
194
+ getByTaskId: (taskId) => subscriptionBus(taskBuses.get(taskId)?.broadcast),
168
195
  cleanupByTaskId(taskId) {
169
196
  const task = taskBuses.get(taskId);
170
197
  task?.broadcast.finished();
171
198
  task?.broadcast.removeAllListeners();
172
199
  taskBuses.delete(taskId);
200
+ ownedTasks.delete(taskId);
173
201
  },
174
202
  };
175
203
  const beforeWorkingTaskMessage = async (message, context) => {
@@ -188,7 +216,20 @@ export async function startA2aServer(options) {
188
216
  throw new RequestMalformedError(error instanceof Error ? error.message : String(error));
189
217
  }
190
218
  };
191
- const requestHandler = new (class extends DefaultRequestHandler {
219
+ const admissionStore = {
220
+ save: async (task, context) => {
221
+ await options.taskStore.save(task, context);
222
+ const admission = admissionReply.getStore();
223
+ if (admission && ++admission.saves >= admission.expectedSaves)
224
+ admission.release();
225
+ },
226
+ load: (taskId, context) => options.taskStore.load(taskId, context),
227
+ list: (params, context) => options.taskStore.list(params, context),
228
+ };
229
+ const pushStore = options.pushNotifications ? new InMemoryPushNotificationStore() : undefined;
230
+ const pushSender = pushStore ? createLegacyAwarePushNotificationSender(pushStore) : undefined;
231
+ const restoredPushFailed = (error) => console.error('[dsh-a2a] restored task push failed:', error);
232
+ const handler = new (class extends DefaultRequestHandler {
192
233
  async beforeMessage(message, context) {
193
234
  if (!message)
194
235
  return;
@@ -201,34 +242,35 @@ export async function startA2aServer(options) {
201
242
  }
202
243
  async sendMessage(...args) {
203
244
  const release = guardApproval(args[0].message);
245
+ const admission = await acquireAdmission(args[0].message);
204
246
  let releaseAdmission;
205
247
  try {
206
248
  if (!release)
207
249
  await beforeWorkingTaskMessage(args[0].message, args[1]);
208
250
  if (!release)
209
251
  releaseAdmission = await this.beforeMessage(args[0].message, args[1]);
210
- return await approvalReply.run(Boolean(release), () => super.sendMessage(...args));
252
+ return await approvalReply.run(Boolean(release), () => admission.run(() => super.sendMessage(...args)));
211
253
  }
212
254
  finally {
255
+ admission.release();
213
256
  releaseAdmission?.();
214
257
  release?.();
215
258
  }
216
259
  }
217
260
  async *sendMessageStream(...args) {
218
261
  const release = guardApproval(args[0].message);
262
+ const admission = await acquireAdmission(args[0].message);
219
263
  let releaseAdmission;
220
264
  try {
221
265
  if (!release)
222
266
  await beforeWorkingTaskMessage(args[0].message, args[1]);
223
267
  if (!release)
224
268
  releaseAdmission = await this.beforeMessage(args[0].message, args[1]);
225
- const stream = observeTask(super.sendMessageStream(...args));
269
+ const stream = admission.run(() => observeTask(super.sendMessageStream(...args)));
226
270
  try {
227
- if (release) {
228
- const first = await approvalReply.run(true, () => stream.next());
229
- if (!first.done)
230
- yield first.value;
231
- }
271
+ const first = await admission.run(() => approvalReply.run(Boolean(release), () => stream.next()));
272
+ if (!first.done)
273
+ yield first.value;
232
274
  yield* stream;
233
275
  }
234
276
  finally {
@@ -237,12 +279,21 @@ export async function startA2aServer(options) {
237
279
  }
238
280
  }
239
281
  finally {
282
+ admission.release();
240
283
  releaseAdmission?.();
241
284
  release?.();
242
285
  }
243
286
  }
244
287
  async cancelTask(...args) {
288
+ // A cold native task must reattach before ownership is decided: its restored
289
+ // bus proves this worker can cancel it, while an unattached remote task still fails closed.
245
290
  await options.beforeTask?.(args[0].id, args[1]);
291
+ const task = await options.taskStore.load(args[0].id, args[1]);
292
+ if (task?.status?.state === TaskState.TASK_STATE_WORKING &&
293
+ !ownedTasks.has(task.id) &&
294
+ !taskBuses.has(task.id)) {
295
+ throw new UnsupportedOperationError('Working task is owned by another worker; route cancellation to its owner.');
296
+ }
246
297
  return super.cancelTask(...args);
247
298
  }
248
299
  async *resubscribe(...args) {
@@ -339,14 +390,17 @@ export async function startA2aServer(options) {
339
390
  await stream.return(undefined);
340
391
  }
341
392
  }
342
- })(card, options.taskStore, options.executor, busManager);
393
+ })(card, admissionStore, options.executor, busManager, pushStore, pushSender, options.extendedAgentCard);
394
+ const requestHandler = connectedRequestHandler(handler);
343
395
  const app = express();
396
+ app.use((_req, res, next) => withStreamConnection(res, next));
344
397
  // The SDK's jsonRpcHandler parses bodies with express.json()'s 100kb default;
345
398
  // raise the ceiling here — body-parser skips re-parsing an already-read body.
346
- app.use(express.json({ limit: '16mb' }));
399
+ app.use(express.json({ limit: '16mb', type: ['application/json', 'application/a2a+json'] }));
347
400
  if (options.heartbeat?.enabled) {
348
401
  app.use((req, res, next) => {
349
402
  const task = {};
403
+ const rest = options.rest && req.path.startsWith(`${base}/rest/`);
350
404
  const legacy = req.body?.method === 'message/stream' || req.body?.method === 'tasks/resubscribe';
351
405
  const flushHeaders = res.flushHeaders;
352
406
  // Wait for the SDK to commit SSE headers: validation errors stay JSON.
@@ -370,13 +424,16 @@ export async function startA2aServer(options) {
370
424
  append: false,
371
425
  lastChunk: false,
372
426
  };
373
- res.write(`data: ${JSON.stringify({
374
- jsonrpc: '2.0',
375
- id: req.body.id ?? null,
376
- result: legacy
377
- ? { kind: 'artifact-update', ...update }
378
- : { artifactUpdate: update },
379
- })}\n\n`);
427
+ const result = legacy
428
+ ? { kind: 'artifact-update', ...update }
429
+ : { artifactUpdate: update };
430
+ res.write(`data: ${JSON.stringify(rest
431
+ ? result
432
+ : {
433
+ jsonrpc: '2.0',
434
+ id: req.body.id ?? null,
435
+ result,
436
+ })}\n\n`);
380
437
  }
381
438
  catch (error) {
382
439
  stop();
@@ -407,6 +464,12 @@ export async function startA2aServer(options) {
407
464
  app.use(`/${AGENT_CARD_PATH}`, cardHandler);
408
465
  // Pre-1.0 discovery path, kept as a convenience alias.
409
466
  app.use('/.well-known/agent.json', cardHandler);
467
+ if (options.rest)
468
+ app.use(`${base}/rest`, restHandler({
469
+ requestHandler,
470
+ userBuilder: UserBuilder.noAuthentication,
471
+ legacyCompat: { enabled: true },
472
+ }));
410
473
  const contextLookup = async (req, res, next) => {
411
474
  const body = req.body;
412
475
  const params = body?.params;
@@ -461,12 +524,63 @@ export async function startA2aServer(options) {
461
524
  res.sendStatus(404);
462
525
  }, contextLookup, rpcHandler);
463
526
  }
527
+ if (options.inbox)
528
+ app.post(base || '/', inboxRpcHandler(options.inbox, card));
464
529
  app.use(base, contextLookup, legacyStreamHandler(requestHandler), rpcHandler);
465
530
  const server = createServer(app);
466
531
  await new Promise((resolve, reject) => {
467
532
  server.once('error', reject);
468
533
  server.listen(options.port, options.host, () => resolve());
469
534
  });
535
+ const actualPort = server.address().port;
536
+ if (options.port === 0 && options.card.publicUrl === undefined) {
537
+ const actualBase = `http://${options.host}:${actualPort}`;
538
+ for (const endpoint of card.supportedInterfaces)
539
+ endpoint.url = endpoint.url.replace(publicUrl, actualBase);
540
+ }
541
+ let grpcServer;
542
+ let grpcPort;
543
+ if (options.grpc?.enabled) {
544
+ try {
545
+ const [grpc, sdk, legacy] = await Promise.all([
546
+ import('@grpc/grpc-js'),
547
+ import('@a2a-js/sdk/server/grpc'),
548
+ import('@a2a-js/sdk/compat/v0_3/server/grpc'),
549
+ ]);
550
+ grpcServer = new grpc.Server();
551
+ const handlers = { requestHandler, userBuilder: sdk.UserBuilder.noAuthentication };
552
+ for (const [descriptor, service] of [
553
+ [sdk.A2AService, sdk.grpcService(handlers)],
554
+ [legacy.LegacyA2AService, legacy.legacyGrpcService(handlers)],
555
+ ]) {
556
+ for (const [method, definition] of Object.entries(descriptor)) {
557
+ if (!definition.responseStream)
558
+ continue;
559
+ const stream = service[method];
560
+ service[method] = (call) => withStreamConnection(call, () => stream(call));
561
+ }
562
+ grpcServer.addService(descriptor, service);
563
+ }
564
+ const listener = grpcServer;
565
+ grpcPort = await new Promise((resolve, reject) => {
566
+ listener.bindAsync(`${options.host}:${options.grpc.port}`, grpc.ServerCredentials.createInsecure(), (error, port) => (error ? reject(error) : resolve(port)));
567
+ });
568
+ card.supportedInterfaces.push(...duplicateInterfacesForLegacy([
569
+ {
570
+ url: options.grpc.publicUrl || `${options.host}:${grpcPort}`,
571
+ protocolBinding: 'GRPC',
572
+ protocolVersion: A2A_PROTOCOL_VERSION,
573
+ tenant: '',
574
+ },
575
+ ], ['GRPC']));
576
+ }
577
+ catch (error) {
578
+ grpcServer?.forceShutdown();
579
+ server.closeAllConnections();
580
+ await new Promise((resolve) => server.close(() => resolve()));
581
+ throw error;
582
+ }
583
+ }
470
584
  return {
471
585
  restoreTask(task, context) {
472
586
  const bus = busManager.createOrGetByTaskId(task.id);
@@ -474,8 +588,25 @@ export async function startA2aServer(options) {
474
588
  const result = new ResultManager(options.taskStore, context);
475
589
  void (async () => {
476
590
  try {
477
- for await (const event of queue.events())
591
+ for await (const event of queue.events()) {
478
592
  await result.processEvent(event);
593
+ if (pushSender) {
594
+ try {
595
+ const snapshot = structuredClone(result.getCurrentTask());
596
+ // SDK execution-event kinds match the StreamResponse oneof cases.
597
+ const response = {
598
+ payload: {
599
+ $case: event.kind,
600
+ value: event.kind === 'task' ? (snapshot ?? event.data) : event.data,
601
+ },
602
+ };
603
+ void pushSender.send(response, context, snapshot).catch(restoredPushFailed);
604
+ }
605
+ catch (error) {
606
+ restoredPushFailed(error);
607
+ }
608
+ }
609
+ }
479
610
  }
480
611
  catch (error) {
481
612
  console.error('[dsh-a2a] restored task persistence failed:', error);
@@ -486,10 +617,12 @@ export async function startA2aServer(options) {
486
617
  })();
487
618
  return bus;
488
619
  },
489
- port: server.address().port,
620
+ port: actualPort,
621
+ grpcPort,
490
622
  close: () => new Promise((resolve) => {
491
623
  for (const stop of stopHeartbeats)
492
624
  stop();
625
+ grpcServer?.forceShutdown();
493
626
  server.closeAllConnections?.();
494
627
  server.close(() => resolve());
495
628
  }),
package/lib/stores/gcs.js CHANGED
@@ -68,11 +68,12 @@ export class GcsTaskStore {
68
68
  prefix: `${this.prefix}/${taskScopeKey(context)}/`,
69
69
  });
70
70
  const shells = [];
71
- for (const file of files) {
72
- if (!file.name.endsWith('/metadata.json.gz'))
73
- continue;
74
- const [buf] = await file.download();
75
- shells.push(JSON.parse(gunzipSync(buf).toString('utf8')));
71
+ const metadata = files.filter((file) => file.name.endsWith('/metadata.json.gz'));
72
+ // Bound in-flight downloads while avoiding one network round-trip per task.
73
+ for (let start = 0; start < metadata.length; start += 16) {
74
+ const buffers = await Promise.all(metadata.slice(start, start + 16).map((file) => file.download()));
75
+ for (const [buf] of buffers)
76
+ shells.push(JSON.parse(gunzipSync(buf).toString('utf8')));
76
77
  }
77
78
  return listShells(shells, params);
78
79
  }
@@ -47,10 +47,11 @@ export class RedisTaskStore {
47
47
  do {
48
48
  const [next, keys] = await this.redis.scan(cursor, 'MATCH', `${this.prefix}:tasks:${taskScopeKey(context)}:*`, 'COUNT', 100);
49
49
  cursor = next;
50
- for (const key of keys) {
51
- const raw = await this.redis.get(key);
52
- if (raw)
53
- shells.push(JSON.parse(raw));
50
+ if (keys.length) {
51
+ const raws = await this.redis.mget(...keys);
52
+ for (const raw of raws)
53
+ if (raw)
54
+ shells.push(JSON.parse(raw));
54
55
  }
55
56
  } while (cursor !== '0');
56
57
  return listShells(shells, params);
@@ -0,0 +1,75 @@
1
+ import { AsyncLocalStorage } from 'node:async_hooks';
2
+ import { DefaultExecutionEventBus } from '@a2a-js/sdk/server';
3
+ const connections = new AsyncLocalStorage();
4
+ /** SDK transports do not end their pending iterators on connection close. */
5
+ export function withStreamConnection(output, run) {
6
+ const controller = new AbortController();
7
+ const connection = { signal: controller.signal, subscription: false };
8
+ let backlogLimit;
9
+ const drained = () => {
10
+ backlogLimit = undefined;
11
+ };
12
+ const close = () => {
13
+ controller.abort();
14
+ output.off('close', close).off('finish', close).off('cancelled', close).off('drain', drained);
15
+ };
16
+ output.once('close', close).once('finish', close).once('cancelled', close);
17
+ const write = output.write;
18
+ output.write = ((...args) => {
19
+ const accepted = write.apply(output, args);
20
+ if (!accepted && !controller.signal.aborted) {
21
+ if (backlogLimit === undefined) {
22
+ // Allow one large frame, then at most 1 MiB or 16 additional gRPC messages.
23
+ backlogLimit = output.writableLength + (output.writableObjectMode ? 16 : 1024 * 1024);
24
+ output.once('drain', drained);
25
+ }
26
+ else if (output.writableLength >= backlogLimit) {
27
+ close();
28
+ output.destroy();
29
+ }
30
+ }
31
+ return accepted;
32
+ });
33
+ return connections.run(connection, run);
34
+ }
35
+ /** A disconnected subscriber stops only its SDK queue, never the task bus. */
36
+ export function subscriptionBus(source) {
37
+ const connection = connections.getStore();
38
+ if (!source || !connection?.subscription)
39
+ return source;
40
+ if (connection.signal.aborted)
41
+ return undefined;
42
+ const mirror = new DefaultExecutionEventBus();
43
+ const forward = (event) => mirror.publish(event);
44
+ const finish = () => {
45
+ source.off('event', forward).off('finished', finish);
46
+ connection.signal.removeEventListener('abort', finish);
47
+ mirror.finished();
48
+ };
49
+ source.on('event', forward).on('finished', finish);
50
+ connection.signal.addEventListener('abort', finish, { once: true });
51
+ return mirror;
52
+ }
53
+ async function* connectedStream(stream, subscription) {
54
+ const connection = connections.getStore();
55
+ if (connection)
56
+ connection.subscription = subscription;
57
+ for await (const event of stream) {
58
+ // Disconnected sends still drain the SDK ResultManager for persistence.
59
+ if (!connection?.signal.aborted)
60
+ yield event;
61
+ }
62
+ }
63
+ /** Keep SDK parsing, serialization and errors; adapt only stream consumption. */
64
+ export function connectedRequestHandler(handler) {
65
+ return new Proxy(handler, {
66
+ get(target, key) {
67
+ if (key === 'sendMessageStream')
68
+ return (...args) => connectedStream(target.sendMessageStream(...args), false);
69
+ if (key === 'resubscribe')
70
+ return (...args) => connectedStream(target.resubscribe(...args), true);
71
+ const value = Reflect.get(target, key, target);
72
+ return typeof value === 'function' ? value.bind(target) : value;
73
+ },
74
+ });
75
+ }