Skip to content

Commit 0b09b43

Browse files
alisa-alisamattKorwel
authored and
liamhelmer
committed
feat(a2a): implement standardized normalization and streaming reassembly (google-gemini#21402)
Co-authored-by: matt korwel <matt.korwel@gmail.com>
1 parent b99f718 commit 0b09b43

2 files changed

Lines changed: 571 additions & 32 deletions

File tree

‎packages/core/src/agents/a2aUtils.test.ts‎

Lines changed: 282 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,17 @@
44
* SPDX-License-Identifier: Apache-2.0
55
*/
66

7-
import { describe, it, expect } from 'vitest';
7+
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
88
import {
99
extractMessageText,
1010
extractIdsFromResponse,
1111
isTerminalState,
1212
A2AResultReassembler,
1313
AUTH_REQUIRED_MSG,
14+
normalizeAgentCard,
15+
getGrpcCredentials,
16+
pinUrlToIp,
17+
splitAgentCardUrl,
1418
} from './a2aUtils.js';
1519
import type { SendMessageResult } from './a2a-client-manager.js';
1620
import type {
@@ -22,8 +26,105 @@ import type {
2226
TaskStatusUpdateEvent,
2327
TaskArtifactUpdateEvent,
2428
} from '@a2a-js/sdk';
29+
import * as dnsPromises from 'node:dns/promises';
30+
import type { LookupAddress } from 'node:dns';
31+
32+
vi.mock('node:dns/promises', () => ({
33+
lookup: vi.fn(),
34+
}));
2535

2636
describe('a2aUtils', () => {
37+
beforeEach(() => {
38+
vi.clearAllMocks();
39+
});
40+
41+
afterEach(() => {
42+
vi.restoreAllMocks();
43+
});
44+
45+
describe('getGrpcCredentials', () => {
46+
it('should return secure credentials for https', () => {
47+
const credentials = getGrpcCredentials('https://test.agent');
48+
expect(credentials).toBeDefined();
49+
});
50+
51+
it('should return insecure credentials for http', () => {
52+
const credentials = getGrpcCredentials('http://test.agent');
53+
expect(credentials).toBeDefined();
54+
});
55+
});
56+
57+
describe('pinUrlToIp', () => {
58+
it('should resolve and pin hostname to IP', async () => {
59+
vi.mocked(
60+
dnsPromises.lookup as unknown as (
61+
hostname: string,
62+
options: { all: true },
63+
) => Promise<LookupAddress[]>,
64+
).mockResolvedValue([{ address: '93.184.216.34', family: 4 }]);
65+
66+
const { pinnedUrl, hostname } = await pinUrlToIp(
67+
'http://example.com:9000',
68+
'test-agent',
69+
);
70+
expect(hostname).toBe('example.com');
71+
expect(pinnedUrl).toBe('http://93.184.216.34:9000/');
72+
});
73+
74+
it('should handle raw host:port strings (standard for gRPC)', async () => {
75+
vi.mocked(
76+
dnsPromises.lookup as unknown as (
77+
hostname: string,
78+
options: { all: true },
79+
) => Promise<LookupAddress[]>,
80+
).mockResolvedValue([{ address: '93.184.216.34', family: 4 }]);
81+
82+
const { pinnedUrl, hostname } = await pinUrlToIp(
83+
'example.com:9000',
84+
'test-agent',
85+
);
86+
expect(hostname).toBe('example.com');
87+
expect(pinnedUrl).toBe('93.184.216.34:9000');
88+
});
89+
90+
it('should throw error if resolution fails (fail closed)', async () => {
91+
vi.mocked(dnsPromises.lookup).mockRejectedValue(new Error('DNS Error'));
92+
93+
await expect(
94+
pinUrlToIp('http://unreachable.com', 'test-agent'),
95+
).rejects.toThrow("Failed to resolve host for agent 'test-agent'");
96+
});
97+
98+
it('should throw error if resolved to private IP', async () => {
99+
vi.mocked(
100+
dnsPromises.lookup as unknown as (
101+
hostname: string,
102+
options: { all: true },
103+
) => Promise<LookupAddress[]>,
104+
).mockResolvedValue([{ address: '10.0.0.1', family: 4 }]);
105+
106+
await expect(
107+
pinUrlToIp('http://malicious.com', 'test-agent'),
108+
).rejects.toThrow('resolves to private IP range');
109+
});
110+
111+
it('should allow localhost/127.0.0.1/::1 exceptions', async () => {
112+
vi.mocked(
113+
dnsPromises.lookup as unknown as (
114+
hostname: string,
115+
options: { all: true },
116+
) => Promise<LookupAddress[]>,
117+
).mockResolvedValue([{ address: '127.0.0.1', family: 4 }]);
118+
119+
const { pinnedUrl, hostname } = await pinUrlToIp(
120+
'http://localhost:9000',
121+
'test-agent',
122+
);
123+
expect(hostname).toBe('localhost');
124+
expect(pinnedUrl).toBe('http://127.0.0.1:9000/');
125+
});
126+
});
127+
27128
describe('isTerminalState', () => {
28129
it('should return true for completed, failed, canceled, and rejected', () => {
29130
expect(isTerminalState('completed')).toBe(true);
@@ -223,6 +324,173 @@ describe('a2aUtils', () => {
223324
} as Message),
224325
).toBe('');
225326
});
327+
328+
it('should handle file parts with neither name nor uri', () => {
329+
const message: Message = {
330+
kind: 'message',
331+
role: 'user',
332+
messageId: '1',
333+
parts: [
334+
{
335+
kind: 'file',
336+
file: {
337+
mimeType: 'text/plain',
338+
},
339+
} as FilePart,
340+
],
341+
};
342+
expect(extractMessageText(message)).toBe('File: [binary/unnamed]');
343+
});
344+
});
345+
346+
describe('normalizeAgentCard', () => {
347+
it('should throw if input is not an object', () => {
348+
expect(() => normalizeAgentCard(null)).toThrow('Agent card is missing.');
349+
expect(() => normalizeAgentCard(undefined)).toThrow(
350+
'Agent card is missing.',
351+
);
352+
expect(() => normalizeAgentCard('not an object')).toThrow(
353+
'Agent card is missing.',
354+
);
355+
});
356+
357+
it('should preserve unknown fields while providing defaults for mandatory ones', () => {
358+
const raw = {
359+
name: 'my-agent',
360+
customField: 'keep-me',
361+
};
362+
363+
const normalized = normalizeAgentCard(raw);
364+
365+
expect(normalized.name).toBe('my-agent');
366+
// @ts-expect-error - testing dynamic preservation
367+
expect(normalized.customField).toBe('keep-me');
368+
expect(normalized.description).toBe('');
369+
expect(normalized.skills).toEqual([]);
370+
expect(normalized.defaultInputModes).toEqual([]);
371+
});
372+
373+
it('should normalize and synchronize interfaces while preserving other fields', () => {
374+
const raw = {
375+
name: 'test',
376+
supportedInterfaces: [
377+
{
378+
url: 'grpc://test',
379+
protocolBinding: 'GRPC',
380+
protocolVersion: '1.0',
381+
},
382+
],
383+
};
384+
385+
const normalized = normalizeAgentCard(raw);
386+
387+
// Should exist in both fields
388+
expect(normalized.additionalInterfaces).toHaveLength(1);
389+
expect(
390+
(normalized as unknown as Record<string, unknown>)[
391+
'supportedInterfaces'
392+
],
393+
).toHaveLength(1);
394+
395+
const intf = normalized.additionalInterfaces?.[0] as unknown as Record<
396+
string,
397+
unknown
398+
>;
399+
400+
expect(intf['transport']).toBe('GRPC');
401+
expect(intf['url']).toBe('grpc://test');
402+
403+
// Should fallback top-level url
404+
expect(normalized.url).toBe('grpc://test');
405+
});
406+
407+
it('should preserve existing top-level url if present', () => {
408+
const raw = {
409+
name: 'test',
410+
url: 'http://existing',
411+
supportedInterfaces: [{ url: 'http://other', transport: 'REST' }],
412+
};
413+
414+
const normalized = normalizeAgentCard(raw);
415+
expect(normalized.url).toBe('http://existing');
416+
});
417+
418+
it('should NOT prepend http:// scheme to raw IP:port strings for gRPC interfaces', () => {
419+
const raw = {
420+
name: 'raw-ip-grpc',
421+
supportedInterfaces: [{ url: '127.0.0.1:9000', transport: 'GRPC' }],
422+
};
423+
424+
const normalized = normalizeAgentCard(raw);
425+
expect(normalized.additionalInterfaces?.[0].url).toBe('127.0.0.1:9000');
426+
expect(normalized.url).toBe('127.0.0.1:9000');
427+
});
428+
429+
it('should prepend http:// scheme to raw IP:port strings for REST interfaces', () => {
430+
const raw = {
431+
name: 'raw-ip-rest',
432+
supportedInterfaces: [{ url: '127.0.0.1:8080', transport: 'REST' }],
433+
};
434+
435+
const normalized = normalizeAgentCard(raw);
436+
expect(normalized.additionalInterfaces?.[0].url).toBe(
437+
'http://127.0.0.1:8080',
438+
);
439+
});
440+
441+
it('should NOT override existing transport if protocolBinding is also present', () => {
442+
const raw = {
443+
name: 'priority-test',
444+
supportedInterfaces: [
445+
{ url: 'foo', transport: 'GRPC', protocolBinding: 'REST' },
446+
],
447+
};
448+
const normalized = normalizeAgentCard(raw);
449+
expect(normalized.additionalInterfaces?.[0].transport).toBe('GRPC');
450+
});
451+
});
452+
453+
describe('splitAgentCardUrl', () => {
454+
const standard = '.well-known/agent-card.json';
455+
456+
it('should return baseUrl as-is if it does not end with standard path', () => {
457+
const url = 'http://localhost:9001/custom/path';
458+
expect(splitAgentCardUrl(url)).toEqual({ baseUrl: url });
459+
});
460+
461+
it('should split correctly if URL ends with standard path', () => {
462+
const url = `http://localhost:9001/${standard}`;
463+
expect(splitAgentCardUrl(url)).toEqual({
464+
baseUrl: 'http://localhost:9001/',
465+
path: undefined,
466+
});
467+
});
468+
469+
it('should handle trailing slash in baseUrl when splitting', () => {
470+
const url = `http://example.com/api/${standard}`;
471+
expect(splitAgentCardUrl(url)).toEqual({
472+
baseUrl: 'http://example.com/api/',
473+
path: undefined,
474+
});
475+
});
476+
477+
it('should ignore hashes and query params when splitting', () => {
478+
const url = `http://localhost:9001/${standard}?foo=bar#baz`;
479+
expect(splitAgentCardUrl(url)).toEqual({
480+
baseUrl: 'http://localhost:9001/',
481+
path: undefined,
482+
});
483+
});
484+
485+
it('should return original URL if parsing fails', () => {
486+
const url = 'not-a-url';
487+
expect(splitAgentCardUrl(url)).toEqual({ baseUrl: url });
488+
});
489+
490+
it('should handle standard path appearing earlier in the path', () => {
491+
const url = `http://localhost:9001/${standard}/something-else`;
492+
expect(splitAgentCardUrl(url)).toEqual({ baseUrl: url });
493+
});
226494
});
227495

228496
describe('A2AResultReassembler', () => {
@@ -233,6 +501,7 @@ describe('a2aUtils', () => {
233501
reassembler.update({
234502
kind: 'status-update',
235503
taskId: 't1',
504+
contextId: 'ctx1',
236505
status: {
237506
state: 'working',
238507
message: {
@@ -247,6 +516,7 @@ describe('a2aUtils', () => {
247516
reassembler.update({
248517
kind: 'artifact-update',
249518
taskId: 't1',
519+
contextId: 'ctx1',
250520
append: false,
251521
artifact: {
252522
artifactId: 'a1',
@@ -259,6 +529,7 @@ describe('a2aUtils', () => {
259529
reassembler.update({
260530
kind: 'status-update',
261531
taskId: 't1',
532+
contextId: 'ctx1',
262533
status: {
263534
state: 'working',
264535
message: {
@@ -273,6 +544,7 @@ describe('a2aUtils', () => {
273544
reassembler.update({
274545
kind: 'artifact-update',
275546
taskId: 't1',
547+
contextId: 'ctx1',
276548
append: true,
277549
artifact: {
278550
artifactId: 'a1',
@@ -291,6 +563,7 @@ describe('a2aUtils', () => {
291563

292564
reassembler.update({
293565
kind: 'status-update',
566+
contextId: 'ctx1',
294567
status: {
295568
state: 'auth-required',
296569
message: {
@@ -310,6 +583,7 @@ describe('a2aUtils', () => {
310583

311584
reassembler.update({
312585
kind: 'status-update',
586+
contextId: 'ctx1',
313587
status: {
314588
state: 'auth-required',
315589
},
@@ -323,6 +597,7 @@ describe('a2aUtils', () => {
323597

324598
const chunk = {
325599
kind: 'status-update',
600+
contextId: 'ctx1',
326601
status: {
327602
state: 'auth-required',
328603
message: {
@@ -351,6 +626,8 @@ describe('a2aUtils', () => {
351626

352627
reassembler.update({
353628
kind: 'task',
629+
id: 'task-1',
630+
contextId: 'ctx1',
354631
status: { state: 'completed' },
355632
history: [
356633
{
@@ -369,6 +646,8 @@ describe('a2aUtils', () => {
369646

370647
reassembler.update({
371648
kind: 'task',
649+
id: 'task-1',
650+
contextId: 'ctx1',
372651
status: { state: 'working' },
373652
history: [
374653
{
@@ -387,6 +666,8 @@ describe('a2aUtils', () => {
387666

388667
reassembler.update({
389668
kind: 'task',
669+
id: 'task-1',
670+
contextId: 'ctx1',
390671
status: { state: 'completed' },
391672
artifacts: [
392673
{

0 commit comments

Comments
 (0)