diff --git a/src/cli/aws/__tests__/agentcore-a2a-bearer.test.ts b/src/cli/aws/__tests__/agentcore-a2a-bearer.test.ts index 71b5af6d9..9b2f78745 100644 --- a/src/cli/aws/__tests__/agentcore-a2a-bearer.test.ts +++ b/src/cli/aws/__tests__/agentcore-a2a-bearer.test.ts @@ -1,4 +1,4 @@ -import { invokeA2ARuntime } from '../agentcore.js'; +import { invokeA2ARuntime, invokeA2ARuntimeStreaming } from '../agentcore.js'; import type { A2AInvokeOptions } from '../agentcore.js'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; @@ -95,3 +95,158 @@ describe('invokeA2ARuntime bearer-token auth path', () => { expect(fetchSpy).not.toHaveBeenCalled(); }); }); + +// --------------------------------------------------------------------------- +// Shared helpers for streaming tests +// --------------------------------------------------------------------------- + +function makeReaderMock(frames: string[]): { + getReader: () => { + read: () => Promise<{ done: boolean; value: Uint8Array | undefined }>; + releaseLock: ReturnType; + }; +} { + const encoder = new TextEncoder(); + let i = 0; + return { + getReader: () => ({ + read: () => { + if (i < frames.length) { + return Promise.resolve({ done: false as const, value: encoder.encode(frames[i++]) }); + } + return Promise.resolve({ done: true as const, value: undefined }); + }, + releaseLock: vi.fn(), + }), + }; +} + +describe('invokeA2ARuntimeStreaming SigV4', () => { + afterEach(() => { + vi.clearAllMocks(); + }); + + it('yields incremental chunks from status-update SSE events', async () => { + const frames = [ + 'data: {"kind":"status-update","status":{"state":"working","message":{"parts":[{"kind":"text","text":"Hello "}]}}}\n\n', + 'data: {"kind":"status-update","status":{"state":"working","message":{"parts":[{"kind":"text","text":"world"}]}}}\n\n', + ]; + const stream = makeReaderMock(frames); + mockSdkSend.mockResolvedValue({ + runtimeSessionId: 'sigv4-stream-session', + response: { transformToWebStream: () => stream }, + }); + + const result = await invokeA2ARuntimeStreaming(baseOpts, 'hello'); + const text = await drain(result.stream); + + expect(mockSdkSend).toHaveBeenCalledTimes(1); + expect(text).toBe('Hello world'); + expect(result.sessionId).toBe('sigv4-stream-session'); + }); + + it('yields artifact-update text with type:"text" parts (backward compat) when no status-update text', async () => { + const frames = [ + 'data: {"kind":"artifact-update","artifact":{"parts":[{"type":"text","text":"legacy output"}]}}\n\n', + ]; + const stream = makeReaderMock(frames); + mockSdkSend.mockResolvedValue({ + runtimeSessionId: 'sigv4-stream-session-2', + response: { transformToWebStream: () => stream }, + }); + + const result = await invokeA2ARuntimeStreaming(baseOpts, 'hello'); + const text = await drain(result.stream); + + expect(text).toBe('legacy output'); + }); + + it('suppresses artifact-update when status-update already emitted text', async () => { + const frames = [ + 'data: {"kind":"status-update","status":{"state":"working","message":{"parts":[{"kind":"text","text":"status text"}]}}}\n\n', + 'data: {"kind":"artifact-update","artifact":{"parts":[{"kind":"text","text":"should be skipped"}]}}\n\n', + ]; + const stream = makeReaderMock(frames); + mockSdkSend.mockResolvedValue({ + runtimeSessionId: 'sigv4-stream-session-3', + response: { transformToWebStream: () => stream }, + }); + + const result = await invokeA2ARuntimeStreaming(baseOpts, 'hello'); + const text = await drain(result.stream); + + expect(text).toBe('status text'); + }); + + it('yields trailing SSE data line without final newline', async () => { + const frames = ['data: {"kind":"status-update","status":{"state":"working","message":{"parts":[{"kind":"text","text":"tail chunk"}]}}}']; + const stream = makeReaderMock(frames); + mockSdkSend.mockResolvedValue({ + runtimeSessionId: 'sigv4-stream-session-4', + response: { transformToWebStream: () => stream }, + }); + + const result = await invokeA2ARuntimeStreaming(baseOpts, 'hello'); + const text = await drain(result.stream); + + expect(text).toBe('tail chunk'); + }); +}); + +describe('invokeA2ARuntimeStreaming bearer-token', () => { + let fetchSpy: { mockRestore: () => void } | undefined; + let capturedRequests: { url: string; init: RequestInit }[]; + + afterEach(() => { + fetchSpy?.mockRestore(); + vi.clearAllMocks(); + }); + + function mockFetchStream(frames: string[], sessionId: string | null = 'bearer-stream-session'): void { + capturedRequests = []; + const stream = makeReaderMock(frames); + fetchSpy = vi.spyOn(globalThis, 'fetch').mockImplementation((input, init) => { + capturedRequests.push({ url: input as string, init: init! }); + return Promise.resolve({ + ok: true, + status: 200, + body: stream, + headers: { + get: (h: string) => (h === 'X-Amzn-Bedrock-AgentCore-Runtime-Session-Id' ? sessionId : null), + }, + } as unknown as Response); + }); + } + + it('uses message/stream method with Bearer token and yields SSE chunks', async () => { + mockFetchStream([ + 'data: {"kind":"status-update","status":{"state":"working","message":{"parts":[{"kind":"text","text":"streamed result"}]}}}\n\n', + ]); + + const result = await invokeA2ARuntimeStreaming({ ...baseOpts, bearerToken: 'my-jwt' }, 'hello'); + const text = await drain(result.stream); + + expect(fetchSpy).toHaveBeenCalledTimes(1); + expect(mockSdkSend).not.toHaveBeenCalled(); + + const headers = capturedRequests[0]!.init.headers as Record; + expect(headers.Authorization).toBe('Bearer my-jwt'); + + const reqBody = JSON.parse(capturedRequests[0]!.init.body as string); + expect(reqBody.method).toBe('message/stream'); + + expect(text).toBe('streamed result'); + expect(result.sessionId).toBe('bearer-stream-session'); + }); + + it('yields artifact-update text with type:"text" parts (backward compat)', async () => { + mockFetchStream([ + 'data: {"kind":"artifact-update","artifact":{"parts":[{"type":"text","text":"legacy bearer output"}]}}\n\n', + ]); + + const result = await invokeA2ARuntimeStreaming({ ...baseOpts, bearerToken: 'my-jwt' }, 'hello'); + const text = await drain(result.stream); + + expect(text).toBe('legacy bearer output'); + }); +}); diff --git a/src/cli/aws/agentcore.ts b/src/cli/aws/agentcore.ts index 49b0cfb60..5d495f581 100644 --- a/src/cli/aws/agentcore.ts +++ b/src/cli/aws/agentcore.ts @@ -960,7 +960,6 @@ export async function invokeA2ARuntime(options: A2AInvokeOptions, message: strin if (options.bearerToken) { const url = buildInvokeUrl(options.region, options.runtimeArn); const headers = buildBearerInvokeHeaders(options, 'application/json, text/event-stream'); - const res = await fetch(url, { method: 'POST', headers, body: JSON.stringify(body) }); if (!res.ok) { const errBody = await res.text().catch(() => ''); @@ -1006,6 +1005,147 @@ export async function invokeA2ARuntime(options: A2AInvokeOptions, message: strin }; } +/** + * Invoke a deployed A2A agent via InvokeAgentRuntime with JSON-RPC message/stream. + * Yields text parts incrementally from SSE events. + */ +export async function invokeA2ARuntimeStreaming( + options: A2AInvokeOptions, + message: string +): Promise { + const body = { + jsonrpc: '2.0', + id: a2aRequestId++, + method: 'message/stream', + params: { + message: { + role: 'user', + parts: [{ kind: 'text', text: message }], + messageId: `msg-${Date.now()}`, + }, + }, + }; + + options.logger?.logSSEEvent(`A2A streaming request: ${JSON.stringify(body)}`); + + async function* a2aStreamGenerator( + reader: ReadableStreamDefaultReader + ): AsyncGenerator { + const decoder = new TextDecoder(); + let buffer = ''; + let streamedFromStatus = false; + + function extractData(line: string): string { + if (!line.startsWith('data: ')) return ''; + return line.slice(6).trim(); + } + + function collectTextParts( + parts?: { kind?: string; type?: string; text?: string }[] + ): string { + if (!parts) return ''; + return parts + .filter(p => (p.kind === 'text' || p.type === 'text') && p.text) + .map(p => p.text!) + .join(''); + } + + function* parseDataLine(line: string): Generator { + const data = extractData(line); + if (!data) return; + + options.logger?.logSSEEvent(line); + + try { + const event = JSON.parse(data) as Record; + // Unwrap JSON-RPC result envelope if present + const target = (event.result as Record) ?? event; + const kind = target.kind as string | undefined; + + if (kind === 'status-update') { + const status = target.status as + | { state?: string; message?: { parts?: { kind?: string; type?: string; text?: string }[] } } + | undefined; + const text = collectTextParts(status?.message?.parts); + if (text) { + streamedFromStatus = true; + yield text; + } + } else if (kind === 'artifact-update' && !streamedFromStatus) { + const artifact = target.artifact as + | { parts?: { kind?: string; type?: string; text?: string }[] } + | undefined; + const text = collectTextParts(artifact?.parts); + if (text) yield text; + } + } catch { + // Non-JSON SSE line, skip + } + } + + try { + while (true) { + const result = await reader.read(); + if (result.done) break; + + buffer += decoder.decode(result.value, { stream: true }); + const lines = buffer.split('\n'); + buffer = lines.pop() ?? ''; + + for (const line of lines) { + yield* parseDataLine(line); + } + } + + // Flush any buffered decoder state so split multi-byte characters are preserved. + const trailingDecoded = decoder.decode(); + if (trailingDecoded) { + buffer += trailingDecoded; + } + + if (buffer) { + yield* parseDataLine(buffer); + } + } finally { + reader.releaseLock(); + } + } + + if (options.bearerToken) { + const url = buildInvokeUrl(options.region, options.runtimeArn); + const streamHeaders = buildBearerInvokeHeaders(options, 'text/event-stream'); + const res = await fetch(url, { method: 'POST', headers: streamHeaders, body: JSON.stringify(body) }); + if (!res.ok) { + const errBody = await res.text().catch(() => ''); + throw new Error(`Invoke failed (${res.status}): ${errBody || res.statusText}`); + } + if (!res.body) throw new Error('No response body for A2A streaming'); + const sessionId = res.headers.get('X-Amzn-Bedrock-AgentCore-Runtime-Session-Id') ?? undefined; + return { stream: a2aStreamGenerator(res.body.getReader()), sessionId }; + } + + const client = createAgentCoreClient(options.region, options.headers); + + const streamCommand = new InvokeAgentRuntimeCommand({ + agentRuntimeArn: options.runtimeArn, + payload: new TextEncoder().encode(JSON.stringify(body)), + contentType: 'application/json', + accept: 'text/event-stream', + runtimeUserId: options.userId ?? DEFAULT_RUNTIME_USER_ID, + ...(options.sessionId && { runtimeSessionId: options.sessionId }), + }); + + const streamResponse = await client.send(streamCommand); + const sessionId = streamResponse.runtimeSessionId; + + if (!streamResponse.response) { + throw new Error('No response from AgentCore Runtime'); + } + + const reader = streamResponse.response.transformToWebStream().getReader(); + return { stream: a2aStreamGenerator(reader), sessionId }; +} + /** Wrap a single string value as an AsyncGenerator for StreamingInvokeResult compatibility. */ // eslint-disable-next-line @typescript-eslint/require-await async function* singleValueStream(value: string): AsyncGenerator { diff --git a/src/cli/aws/index.ts b/src/cli/aws/index.ts index 09851e678..6ab52a132 100644 --- a/src/cli/aws/index.ts +++ b/src/cli/aws/index.ts @@ -55,6 +55,7 @@ export { DEFAULT_RUNTIME_USER_ID, executeBashCommand, invokeA2ARuntime, + invokeA2ARuntimeStreaming, invokeAgentRuntime, invokeAgentRuntimeStreaming, mcpInitSession, diff --git a/src/cli/commands/invoke/__tests__/action-a2a-routing.test.ts b/src/cli/commands/invoke/__tests__/action-a2a-routing.test.ts new file mode 100644 index 000000000..873c95f44 --- /dev/null +++ b/src/cli/commands/invoke/__tests__/action-a2a-routing.test.ts @@ -0,0 +1,102 @@ +import type { InvokeContext } from '../action'; +import { handleInvoke } from '../action'; +import type { InvokeOptions } from '../types'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +const mockResolveInvokeTarget = vi.fn(); +const mockInvokeA2ARuntime = vi.fn(); +const mockInvokeA2ARuntimeStreaming = vi.fn(); + +vi.mock('../resolve', () => ({ + resolveInvokeTarget: (...args: unknown[]) => mockResolveInvokeTarget(...args), +})); + +vi.mock('../../../aws', () => ({ + DEFAULT_RUNTIME_USER_ID: 'default-user', + buildAguiRunInput: vi.fn(), + executeBashCommand: vi.fn(), + extractResult: vi.fn(), + getOrCreatePaymentSession: vi.fn(), + invokeA2ARuntime: (...args: unknown[]) => mockInvokeA2ARuntime(...args), + invokeA2ARuntimeStreaming: (...args: unknown[]) => mockInvokeA2ARuntimeStreaming(...args), + invokeAgentRuntime: vi.fn(), + invokeAgentRuntimeStreaming: vi.fn(), + invokeAguiRuntime: vi.fn(), + mcpCallTool: vi.fn(), + mcpInitSession: vi.fn(), + mcpListTools: vi.fn(), + parseSSE: vi.fn(), +})); + +function resolvedA2A(): Record { + return { + success: true, + agentSpec: { name: 'A2AAgent', protocol: 'A2A' }, + targetName: 'default', + targetConfig: { name: 'default', region: 'us-east-1' }, + region: 'us-east-1', + runtimeArn: 'arn:aws:bedrock-agentcore:us-east-1:123:runtime/r', + baggage: undefined, + }; +} + +function makeContext(): InvokeContext { + return { + project: { name: 'p', runtimes: [{ name: 'A2AAgent', protocol: 'A2A' }] } as never, + deployedState: { targets: { default: { resources: {} } } } as never, + awsTargets: [{ name: 'default', region: 'us-east-1' }] as never, + }; +} + +// eslint-disable-next-line @typescript-eslint/require-await +async function* streamOf(text: string): AsyncGenerator { + yield text; +} + +describe('handleInvoke — A2A stream routing', () => { + let stdoutSpy: ReturnType; + + beforeEach(() => { + mockResolveInvokeTarget.mockResolvedValue(resolvedA2A()); + mockInvokeA2ARuntime.mockResolvedValue({ stream: streamOf('non-stream path'), sessionId: undefined }); + mockInvokeA2ARuntimeStreaming.mockResolvedValue({ stream: streamOf('stream path'), sessionId: undefined }); + stdoutSpy = vi.spyOn(process.stdout, 'write').mockImplementation(() => true); + }); + + afterEach(() => { + vi.clearAllMocks(); + stdoutSpy.mockRestore(); + }); + + it('routes --stream to invokeA2ARuntimeStreaming for A2A agents', async () => { + const options: InvokeOptions = { prompt: 'hello', stream: true }; + const result = await handleInvoke(makeContext(), options); + + expect(result.success).toBe(true); + expect(mockInvokeA2ARuntimeStreaming).toHaveBeenCalledTimes(1); + expect(mockInvokeA2ARuntime).not.toHaveBeenCalled(); + expect(mockInvokeA2ARuntimeStreaming).toHaveBeenCalledWith( + expect.objectContaining({ + region: 'us-east-1', + runtimeArn: 'arn:aws:bedrock-agentcore:us-east-1:123:runtime/r', + }), + 'hello' + ); + }); + + it('routes non-stream invokes to invokeA2ARuntime for A2A agents', async () => { + const options: InvokeOptions = { prompt: 'hello', stream: false }; + const result = await handleInvoke(makeContext(), options); + + expect(result.success).toBe(true); + expect(mockInvokeA2ARuntime).toHaveBeenCalledTimes(1); + expect(mockInvokeA2ARuntimeStreaming).not.toHaveBeenCalled(); + expect(mockInvokeA2ARuntime).toHaveBeenCalledWith( + expect.objectContaining({ + region: 'us-east-1', + runtimeArn: 'arn:aws:bedrock-agentcore:us-east-1:123:runtime/r', + }), + 'hello' + ); + }); +}); diff --git a/src/cli/commands/invoke/__tests__/action-payments.test.ts b/src/cli/commands/invoke/__tests__/action-payments.test.ts index 724955766..6577429f9 100644 --- a/src/cli/commands/invoke/__tests__/action-payments.test.ts +++ b/src/cli/commands/invoke/__tests__/action-payments.test.ts @@ -38,6 +38,7 @@ vi.mock('../../../aws', () => ({ buildAguiRunInput: vi.fn(), executeBashCommand: vi.fn(), invokeA2ARuntime: vi.fn(), + invokeA2ARuntimeStreaming: vi.fn(), invokeAguiRuntime: vi.fn(), mcpCallTool: vi.fn(), mcpInitSession: vi.fn(), diff --git a/src/cli/commands/invoke/action.ts b/src/cli/commands/invoke/action.ts index 2a6cb139d..72abf7416 100644 --- a/src/cli/commands/invoke/action.ts +++ b/src/cli/commands/invoke/action.ts @@ -7,6 +7,7 @@ import { extractResult, getOrCreatePaymentSession, invokeA2ARuntime, + invokeA2ARuntimeStreaming, invokeAgentRuntime, invokeAgentRuntimeStreaming, invokeAguiRuntime, @@ -583,20 +584,20 @@ export async function handleInvoke(context: InvokeContext, options: InvokeOption return { success: false, error: new ValidationError('No prompt provided. Usage: agentcore invoke "your prompt"') }; } - // A2A protocol handling — send JSON-RPC message/send via InvokeAgentRuntime + // A2A protocol handling — send JSON-RPC via InvokeAgentRuntime if (agentSpec.protocol === 'A2A') { try { - const a2aResult = await invokeA2ARuntime( - { - region: targetConfig.region, - runtimeArn: runtimeArn, - userId: options.userId, - sessionId: options.sessionId, - headers: options.headers, - bearerToken: options.bearerToken, - }, - options.prompt - ); + const a2aOpts = { + region: targetConfig.region, + runtimeArn: runtimeArn, + userId: options.userId, + sessionId: options.sessionId, + headers: options.headers, + bearerToken: options.bearerToken, + }; + const a2aResult = options.stream + ? await invokeA2ARuntimeStreaming(a2aOpts, options.prompt) + : await invokeA2ARuntime(a2aOpts, options.prompt); let response = ''; for await (const chunk of a2aResult.stream) { response += chunk;