@genkit-ai/core
Version:
Genkit AI framework core libraries.
1,293 lines (1,242 loc) • 39.9 kB
text/typescript
/**
* Copyright 2024 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import * as assert from 'assert';
import { afterEach, beforeEach, describe, it } from 'node:test';
import { WebSocket, WebSocketServer } from 'ws';
import { z } from 'zod';
import { action, bidiAction } from '../src/action.js';
import { initNodeFeatures } from '../src/node.js';
import { ReflectionServerV2 } from '../src/reflection-v2.js';
import { Registry } from '../src/registry.js';
initNodeFeatures();
describe('ReflectionServerV2', () => {
let wss: WebSocketServer;
let server: ReflectionServerV2;
let registry: Registry;
let port: number;
let connections: WebSocket[] = [];
beforeEach(() => {
return new Promise<void>((resolve) => {
wss = new WebSocketServer({ port: 0 });
wss.on('listening', () => {
port = (wss.address() as any).port;
resolve();
});
wss.on('connection', (ws) => {
connections.push(ws); // Track all connections
});
registry = new Registry();
});
});
afterEach(async () => {
if (server) {
await server.stop();
}
// Terminate all connections to let wss.close() proceed
for (const ws of connections) {
ws.terminate();
}
connections = [];
await new Promise<void>((resolve) => {
wss.close(() => resolve());
});
});
it('should connect to the server and register', async () => {
const connected = new Promise<void>((resolve, reject) => {
const timer = setTimeout(
() => reject(new Error('connect timeout')),
2000
);
wss.on('connection', (ws) => {
ws.on('message', (data) => {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
assert.strictEqual(msg.params.name, 'test-app');
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
clearTimeout(timer);
resolve();
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
name: 'test-app',
});
await server.start();
await connected;
});
it('should handle listActions', async () => {
// Register a dummy action
const testAction = action(
{
name: 'testAction',
description: 'A test action',
inputSchema: z.object({ foo: z.string() }),
outputSchema: z.object({ bar: z.string() }),
actionType: 'custom',
},
async (input) => ({ bar: input.foo })
);
registry.registerAction('custom', testAction);
const gotListActions = new Promise<void>((resolve, reject) => {
const timer = setTimeout(
() => reject(new Error('listActions timeout')),
2000
);
wss.on('connection', (ws) => {
ws.on('message', (data) => {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
// After registration, request listActions
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'listActions',
id: '123',
})
);
} else if (msg.id === '123') {
const actions = msg.result.actions;
assert.ok(actions['/custom/testAction']);
assert.strictEqual(
actions['/custom/testAction'].name,
'testAction'
);
clearTimeout(timer);
resolve();
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
});
await server.start();
await gotListActions;
});
it('should handle listValues', async () => {
registry.registerValue('middleware', 'my-mw', { template: 'foo' });
const gotListValues = new Promise<void>((resolve, reject) => {
const timer = setTimeout(
() => reject(new Error('listValues timeout')),
2000
);
wss.on('connection', (ws) => {
ws.on('message', (data) => {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'listValues',
params: { type: 'middleware' },
id: '124',
})
);
} else if (msg.id === '124') {
assert.ok(msg.result.values['my-mw']);
assert.strictEqual(msg.result.values['my-mw'].template, 'foo');
clearTimeout(timer);
resolve();
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
});
await server.start();
await gotListValues;
});
it('should handle listValues with toJson mapping', async () => {
registry.registerValue('middleware', 'mw1', {
toJson: () => ({ name: 'mw1' }),
});
registry.registerValue('middleware', 'mw2', {
name: 'mw2',
});
const gotListValues = new Promise<void>((resolve, reject) => {
const timer = setTimeout(
() => reject(new Error('listValues toJson timeout')),
2000
);
wss.on('connection', (ws) => {
ws.on('message', (data) => {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'listValues',
params: { type: 'middleware' },
id: '125',
})
);
} else if (msg.id === '125') {
assert.ok(msg.result.values['mw1']);
assert.strictEqual(msg.result.values['mw1'].name, 'mw1');
assert.ok(msg.result.values['mw2']);
assert.strictEqual(msg.result.values['mw2'].name, 'mw2');
clearTimeout(timer);
resolve();
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
});
await server.start();
await gotListValues;
});
it('should reject unsupported type parameter for listValues in V2', async () => {
const gotError = new Promise<void>((resolve, reject) => {
const timer = setTimeout(
() => reject(new Error('listValues error timeout')),
2000
);
wss.on('connection', (ws) => {
ws.on('message', (data) => {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'listValues',
params: { type: 'unsupported_type' },
id: '126',
})
);
} else if (msg.id === '126') {
assert.ok(msg.error);
assert.strictEqual(msg.error.code, -32602);
assert.match(msg.error.message, /is not supported/);
clearTimeout(timer);
resolve();
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
});
await server.start();
await gotError;
});
it('should handle runAction', async () => {
const testAction = action(
{
name: 'testAction',
inputSchema: z.object({ foo: z.string() }),
outputSchema: z.object({ bar: z.string() }),
actionType: 'custom',
},
async (input) => ({ bar: input.foo })
);
registry.registerAction('custom', testAction);
const actionRun = new Promise<void>((resolve, reject) => {
const timeout = setTimeout(
() => reject(new Error('runAction timeout')),
2000
);
wss.on('connection', (ws) => {
ws.on('message', (data) => {
try {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'runAction',
params: {
key: '/custom/testAction',
input: { foo: 'baz' },
},
id: '456',
})
);
} else if (msg.id === '456') {
if (msg.error) {
reject(
new Error(`runAction error: ${JSON.stringify(msg.error)}`)
);
return;
}
assert.strictEqual(msg.result.result.bar, 'baz');
clearTimeout(timeout);
resolve();
}
} catch (e) {
clearTimeout(timeout);
reject(e);
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
});
await server.start();
await actionRun;
});
it('should handle streaming runAction', async () => {
const streamAction = action(
{
name: 'streamAction',
inputSchema: z.object({ foo: z.string() }),
outputSchema: z.string(),
actionType: 'custom',
},
async (input, { sendChunk }) => {
sendChunk('chunk1');
sendChunk('chunk2');
return 'done';
}
);
registry.registerAction('custom', streamAction);
const chunks: any[] = [];
const actionRun = new Promise<void>((resolve, reject) => {
const timeout = setTimeout(
() => reject(new Error('streamAction timeout')),
2000
);
wss.on('connection', (ws) => {
ws.on('message', (data) => {
try {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'runAction',
params: {
key: '/custom/streamAction',
input: { foo: 'baz' },
stream: true,
},
id: '789',
})
);
} else if (msg.method === 'streamChunk') {
chunks.push(msg.params.chunk);
} else if (msg.id === '789') {
if (msg.error) {
reject(
new Error(`streamAction error: ${JSON.stringify(msg.error)}`)
);
return;
}
assert.strictEqual(msg.result.result, 'done');
assert.deepStrictEqual(chunks, ['chunk1', 'chunk2']);
clearTimeout(timeout);
resolve();
}
} catch (e) {
clearTimeout(timeout);
reject(e);
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
});
await server.start();
await actionRun;
});
it('should handle cancelAction', async () => {
let cancelSignal: AbortSignal | undefined;
const longAction = action(
{
name: 'longAction',
inputSchema: z.any(),
outputSchema: z.any(),
actionType: 'custom',
},
async (_, { abortSignal }) => {
cancelSignal = abortSignal;
await new Promise((resolve, reject) => {
const timer = setTimeout(resolve, 5000);
if (abortSignal.aborted) {
clearTimeout(timer);
reject(new Error('Action cancelled'));
return;
}
abortSignal.addEventListener('abort', () => {
clearTimeout(timer);
reject(new Error('Action cancelled'));
});
});
}
);
registry.registerAction('custom', longAction);
const actionCancelled = new Promise<void>((resolve, reject) => {
const timeout = setTimeout(
() => reject(new Error('cancelAction timeout')),
2000
);
wss.on('connection', (ws) => {
ws.on('message', (data) => {
try {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
// Start action
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'runAction',
params: {
key: '/custom/longAction',
input: {},
},
id: '999',
})
);
} else if (msg.method === 'runActionState') {
// Got traceId, send cancel
const traceId = msg.params.state.traceId;
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'cancelAction',
params: { traceId },
id: '1000',
})
);
} else if (msg.id === '1000') {
// Cancel response
assert.strictEqual(msg.result.message, 'Action cancelled');
} else if (msg.id === '999') {
// Run action response (should be error)
if (msg.error) {
// Ensure code indicates cancellation if possible, or just error
// In implementation we send code -32000 and message 'Action was cancelled'
assert.match(msg.error.message, /cancelled/);
assert.ok(cancelSignal?.aborted);
clearTimeout(timeout);
resolve();
} else {
reject(new Error('Action should have failed'));
}
}
} catch (e) {
clearTimeout(timeout);
reject(e);
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
});
await server.start();
await actionCancelled;
});
it('should reconnect when lost connection and register again', async () => {
let connectionCount = 0;
const reconnected = new Promise<void>((resolve, reject) => {
const timeout = setTimeout(
() => reject(new Error('reconnect timeout')),
5000
);
wss.on('connection', (ws) => {
connectionCount++;
ws.on('message', (data) => {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
if (connectionCount === 1) {
ws.terminate(); // Simulate server drop
} else if (connectionCount === 2) {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
clearTimeout(timeout);
resolve();
}
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
});
(server as any).baseDelayMs = 10; // Fast for testing
await server.start();
await reconnected;
});
it('should handle bidi streaming action with input and output streams', async () => {
// Create a bidi action that echoes input chunks back as output chunks,
// and returns the count of received chunks as the final result.
const echoBidiAction = bidiAction(
{
name: 'echoBidi',
inputSchema: z.string(),
outputSchema: z.object({ count: z.number() }),
streamSchema: z.string(),
actionType: 'custom',
},
async function* ({ inputStream }) {
let count = 0;
for await (const chunk of inputStream) {
count++;
yield `echo:${chunk}`;
}
return { count };
}
);
registry.registerAction('custom', echoBidiAction);
const outputChunks: any[] = [];
const actionRun = new Promise<void>((resolve, reject) => {
const timeout = setTimeout(
() => reject(new Error('bidi action timeout')),
3000
);
wss.on('connection', (ws) => {
ws.on('message', (data) => {
try {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
// Start the bidi action with stream + streamInput
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'runAction',
params: {
key: '/custom/echoBidi',
input: null,
stream: true,
streamInput: true,
},
id: 'bidi-1',
})
);
// Send input stream chunks
setTimeout(() => {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'sendInputStreamChunk',
params: { requestId: 'bidi-1', chunk: 'hello' },
})
);
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'sendInputStreamChunk',
params: { requestId: 'bidi-1', chunk: 'world' },
})
);
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'endInputStream',
params: { requestId: 'bidi-1' },
})
);
}, 50);
} else if (msg.method === 'streamChunk') {
outputChunks.push(msg.params.chunk);
} else if (msg.id === 'bidi-1') {
if (msg.error) {
reject(
new Error(`bidi action error: ${JSON.stringify(msg.error)}`)
);
return;
}
assert.deepStrictEqual(msg.result.result, { count: 2 });
assert.deepStrictEqual(outputChunks, [
'echo:hello',
'echo:world',
]);
clearTimeout(timeout);
resolve();
}
} catch (e) {
clearTimeout(timeout);
reject(e);
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
});
await server.start();
await actionRun;
});
it('should handle bidi streaming with multiple interleaved chunks', async () => {
// Bidi action that transforms each input chunk and yields it
const transformBidiAction = bidiAction(
{
name: 'transformBidi',
inputSchema: z.number(),
outputSchema: z.object({ sum: z.number() }),
streamSchema: z.number(),
actionType: 'custom',
},
async function* ({ inputStream }) {
let sum = 0;
for await (const n of inputStream) {
sum += n;
yield n * 2; // output is double the input
}
return { sum };
}
);
registry.registerAction('custom', transformBidiAction);
const outputChunks: any[] = [];
const actionRun = new Promise<void>((resolve, reject) => {
const timeout = setTimeout(
() => reject(new Error('transform bidi timeout')),
3000
);
wss.on('connection', (ws) => {
ws.on('message', (data) => {
try {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'runAction',
params: {
key: '/custom/transformBidi',
input: null,
stream: true,
streamInput: true,
},
id: 'bidi-2',
})
);
// Send chunks with slight delays to simulate real-time input
setTimeout(() => {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'sendInputStreamChunk',
params: { requestId: 'bidi-2', chunk: 1 },
})
);
}, 20);
setTimeout(() => {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'sendInputStreamChunk',
params: { requestId: 'bidi-2', chunk: 2 },
})
);
}, 40);
setTimeout(() => {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'sendInputStreamChunk',
params: { requestId: 'bidi-2', chunk: 3 },
})
);
}, 60);
setTimeout(() => {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'endInputStream',
params: { requestId: 'bidi-2' },
})
);
}, 80);
} else if (msg.method === 'streamChunk') {
if (msg.params.requestId === 'bidi-2') {
outputChunks.push(msg.params.chunk);
}
} else if (msg.id === 'bidi-2') {
if (msg.error) {
reject(
new Error(
`transform bidi error: ${JSON.stringify(msg.error)}`
)
);
return;
}
assert.deepStrictEqual(msg.result.result, { sum: 6 });
assert.deepStrictEqual(outputChunks, [2, 4, 6]);
clearTimeout(timeout);
resolve();
}
} catch (e) {
clearTimeout(timeout);
reject(e);
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
});
await server.start();
await actionRun;
});
it('should handle bidi action with no input chunks (immediate close)', async () => {
const emptyBidiAction = bidiAction(
{
name: 'emptyBidi',
inputSchema: z.string(),
outputSchema: z.object({ received: z.number() }),
streamSchema: z.string(),
actionType: 'custom',
},
async function* ({ inputStream }) {
let received = 0;
for await (const _ of inputStream) {
received++;
}
return { received };
}
);
registry.registerAction('custom', emptyBidiAction);
const actionRun = new Promise<void>((resolve, reject) => {
const timeout = setTimeout(
() => reject(new Error('empty bidi timeout')),
3000
);
wss.on('connection', (ws) => {
ws.on('message', (data) => {
try {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'runAction',
params: {
key: '/custom/emptyBidi',
input: null,
stream: true,
streamInput: true,
},
id: 'bidi-3',
})
);
// Immediately close the input stream without sending any chunks
setTimeout(() => {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'endInputStream',
params: { requestId: 'bidi-3' },
})
);
}, 50);
} else if (msg.id === 'bidi-3') {
if (msg.error) {
reject(
new Error(`empty bidi error: ${JSON.stringify(msg.error)}`)
);
return;
}
assert.deepStrictEqual(msg.result.result, { received: 0 });
clearTimeout(timeout);
resolve();
}
} catch (e) {
clearTimeout(timeout);
reject(e);
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
});
await server.start();
await actionRun;
});
it('should handle bidi action with init data', async () => {
const initBidiAction = bidiAction(
{
name: 'initBidi',
inputSchema: z.string(),
outputSchema: z.object({
prefix: z.string(),
messages: z.array(z.string()),
}),
streamSchema: z.string(),
initSchema: z.object({ prefix: z.string() }),
actionType: 'custom',
},
async function* ({ inputStream, init }) {
const messages: string[] = [];
for await (const msg of inputStream) {
const prefixed = `${init.prefix}:${msg}`;
messages.push(prefixed);
yield prefixed;
}
return { prefix: init.prefix, messages };
}
);
registry.registerAction('custom', initBidiAction);
const outputChunks: any[] = [];
const actionRun = new Promise<void>((resolve, reject) => {
const timeout = setTimeout(
() => reject(new Error('init bidi timeout')),
3000
);
wss.on('connection', (ws) => {
ws.on('message', (data) => {
try {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
// Init data is passed via the init param
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'runAction',
params: {
key: '/custom/initBidi',
input: null,
init: { prefix: 'test' },
stream: true,
streamInput: true,
},
id: 'bidi-4',
})
);
setTimeout(() => {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'sendInputStreamChunk',
params: { requestId: 'bidi-4', chunk: 'a' },
})
);
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'sendInputStreamChunk',
params: { requestId: 'bidi-4', chunk: 'b' },
})
);
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'endInputStream',
params: { requestId: 'bidi-4' },
})
);
}, 50);
} else if (msg.method === 'streamChunk') {
if (msg.params.requestId === 'bidi-4') {
outputChunks.push(msg.params.chunk);
}
} else if (msg.id === 'bidi-4') {
if (msg.error) {
reject(
new Error(`init bidi error: ${JSON.stringify(msg.error)}`)
);
return;
}
assert.deepStrictEqual(msg.result.result, {
prefix: 'test',
messages: ['test:a', 'test:b'],
});
assert.deepStrictEqual(outputChunks, ['test:a', 'test:b']);
clearTimeout(timeout);
resolve();
}
} catch (e) {
clearTimeout(timeout);
reject(e);
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
});
await server.start();
await actionRun;
});
it('should handle bidi streaming without output streaming (streamInput only)', async () => {
// Test bidi with streamInput but no output streaming (stream: false)
const collectBidiAction = bidiAction(
{
name: 'collectBidi',
inputSchema: z.string(),
outputSchema: z.object({ items: z.array(z.string()) }),
streamSchema: z.string(),
actionType: 'custom',
},
async function* ({ inputStream }) {
const items: string[] = [];
for await (const item of inputStream) {
items.push(item);
}
return { items };
}
);
registry.registerAction('custom', collectBidiAction);
const actionRun = new Promise<void>((resolve, reject) => {
const timeout = setTimeout(
() => reject(new Error('collect bidi timeout')),
3000
);
wss.on('connection', (ws) => {
ws.on('message', (data) => {
try {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
// stream: false but streamInput: true
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'runAction',
params: {
key: '/custom/collectBidi',
input: null,
stream: false,
streamInput: true,
},
id: 'bidi-5',
})
);
setTimeout(() => {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'sendInputStreamChunk',
params: { requestId: 'bidi-5', chunk: 'x' },
})
);
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'sendInputStreamChunk',
params: { requestId: 'bidi-5', chunk: 'y' },
})
);
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'endInputStream',
params: { requestId: 'bidi-5' },
})
);
}, 50);
} else if (msg.id === 'bidi-5') {
if (msg.error) {
reject(
new Error(`collect bidi error: ${JSON.stringify(msg.error)}`)
);
return;
}
assert.deepStrictEqual(msg.result.result, {
items: ['x', 'y'],
});
clearTimeout(timeout);
resolve();
}
} catch (e) {
clearTimeout(timeout);
reject(e);
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
});
await server.start();
await actionRun;
});
it('should use explicit init param over input for bidi action', async () => {
// When both init and input are provided for a bidi action,
// the explicit init should take precedence over input.
const bidiWithInit = bidiAction(
{
name: 'bidiWithInit',
inputSchema: z.string(),
outputSchema: z.object({
initValue: z.string(),
messages: z.array(z.string()),
}),
streamSchema: z.string(),
initSchema: z.object({ sessionId: z.string() }),
actionType: 'custom',
},
async function* ({ inputStream, init }) {
const messages: string[] = [];
for await (const msg of inputStream) {
messages.push(`${init.sessionId}:${msg}`);
yield `${init.sessionId}:${msg}`;
}
return { initValue: init.sessionId, messages };
}
);
registry.registerAction('custom', bidiWithInit);
const outputChunks: any[] = [];
const actionRun = new Promise<void>((resolve, reject) => {
const timeout = setTimeout(
() => reject(new Error('bidi with init timeout')),
3000
);
wss.on('connection', (ws) => {
ws.on('message', (data) => {
try {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
// Send both input and init; init should win
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'runAction',
params: {
key: '/custom/bidiWithInit',
input: { sessionId: 'wrong-from-input' },
init: { sessionId: 'correct-from-init' },
stream: true,
streamInput: true,
},
id: 'bidi-init-1',
})
);
setTimeout(() => {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'sendInputStreamChunk',
params: { requestId: 'bidi-init-1', chunk: 'hello' },
})
);
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'endInputStream',
params: { requestId: 'bidi-init-1' },
})
);
}, 50);
} else if (msg.method === 'streamChunk') {
if (msg.params.requestId === 'bidi-init-1') {
outputChunks.push(msg.params.chunk);
}
} else if (msg.id === 'bidi-init-1') {
if (msg.error) {
reject(
new Error(
`bidi with init error: ${JSON.stringify(msg.error)}`
)
);
return;
}
// init should be { sessionId: 'correct-from-init' }, not input
assert.deepStrictEqual(msg.result.result, {
initValue: 'correct-from-init',
messages: ['correct-from-init:hello'],
});
assert.deepStrictEqual(outputChunks, ['correct-from-init:hello']);
clearTimeout(timeout);
resolve();
}
} catch (e) {
clearTimeout(timeout);
reject(e);
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
});
await server.start();
await actionRun;
});
it('should pass init param from runAction to the action handler', async () => {
const initAction = action(
{
name: 'initAction',
inputSchema: z.string(),
outputSchema: z.object({
receivedInput: z.string(),
receivedInit: z.object({ sessionId: z.string() }),
}),
initSchema: z.object({ sessionId: z.string() }),
actionType: 'custom',
},
async (input, { init }) => {
return { receivedInput: input, receivedInit: init };
}
);
registry.registerAction('custom', initAction);
const actionRun = new Promise<void>((resolve, reject) => {
const timeout = setTimeout(
() => reject(new Error('init action timeout')),
3000
);
wss.on('connection', (ws) => {
ws.on('message', (data) => {
try {
const msg = JSON.parse(data.toString());
if (msg.method === 'register') {
ws.send(
JSON.stringify({
jsonrpc: '2.0',
result: {},
id: msg.id,
})
);
ws.send(
JSON.stringify({
jsonrpc: '2.0',
method: 'runAction',
params: {
key: '/custom/initAction',
input: 'hello',
init: { sessionId: 'sess-42' },
stream: false,
},
id: 'init-1',
})
);
} else if (msg.id === 'init-1') {
if (msg.error) {
reject(
new Error(`init action error: ${JSON.stringify(msg.error)}`)
);
return;
}
assert.deepStrictEqual(msg.result.result, {
receivedInput: 'hello',
receivedInit: { sessionId: 'sess-42' },
});
clearTimeout(timeout);
resolve();
}
} catch (e) {
clearTimeout(timeout);
reject(e);
}
});
});
});
server = new ReflectionServerV2(registry, {
url: `ws://localhost:${port}`,
});
await server.start();
await actionRun;
});
});