Genkit MCP

/** * 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 type { SpanData } from '@genkit-ai/tools-common'; import * as assert from 'assert'; import { EventEmitter } from 'events'; import type { Response } from 'express'; import { beforeEach, describe, it } from 'node:test'; import { BroadcastManager, type SpanEvent } from '../src/broadcast-manager'; /** * Creates a mock Response object for testing. */ function createMockResponse(): Response & { writtenData: string[]; ended: boolean; emit: (event: string) => boolean; } { const emitter = new EventEmitter(); const mock = { writtenData: [] as string[], ended: false, write(data: string) { if (mock.ended) { throw new Error('Cannot write to ended response'); } mock.writtenData.push(data); return true; }, end() { mock.ended = true; }, on(event: string, handler: (...args: any[]) => void) { emitter.on(event, handler); return mock; }, emit(event: string) { return emitter.emit(event); }, }; return mock as any; } function createSpanEvent( traceId: string, type: 'span_start' | 'span_end' = 'span_start' ): SpanEvent { return { type, traceId, span: { traceId, spanId: 'test-span', displayName: 'Test Span', startTime: 1000, endTime: type === 'span_end' ? 2000 : 0, instrumentationLibrary: { name: 'genkit' }, spanKind: 'INTERNAL', attributes: {}, status: { code: 0 }, } as SpanData, }; } describe('BroadcastManager', () => { let manager: BroadcastManager; beforeEach(() => { manager = new BroadcastManager(); }); it('subscribes connections and broadcasts events to them', () => { const response1 = createMockResponse(); const response2 = createMockResponse(); manager.subscribe('trace-1', response1); manager.subscribe('trace-1', response2); const event = createSpanEvent('trace-1'); manager.broadcast('trace-1', event); const expectedData = `data: ${JSON.stringify(event)}\n\n`; assert.strictEqual(response1.writtenData[0], expectedData); assert.strictEqual(response2.writtenData[0], expectedData); }); it('auto-unsubscribes when connection closes', () => { const response = createMockResponse(); manager.subscribe('trace-1', response); assert.strictEqual(manager.getConnectionCount('trace-1'), 1); response.emit('close'); assert.strictEqual(manager.getConnectionCount('trace-1'), 0); }); it('close() ends all connections for a traceId', () => { const response1 = createMockResponse(); const response2 = createMockResponse(); manager.subscribe('trace-1', response1); manager.subscribe('trace-1', response2); manager.close('trace-1'); assert.strictEqual(response1.ended, true); assert.strictEqual(response2.ended, true); assert.strictEqual(manager.hasConnections('trace-1'), false); }); });