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 type { Response } from 'express'; /** * Event type for span broadcasts. */ export interface SpanEvent { type: 'span_start' | 'span_end'; traceId: string; span: SpanData; } /** * Broadcast manager for SSE connections. * Tracks active connections per traceId and broadcasts updates. */ export class BroadcastManager { private connections: Map<string, Set<Response>> = new Map(); /** * Register a new SSE connection for a traceId. */ subscribe(traceId: string, response: Response): void { if (!this.connections.has(traceId)) { this.connections.set(traceId, new Set()); } this.connections.get(traceId)!.add(response); // Clean up when connection closes response.on('close', () => { this.unsubscribe(traceId, response); }); } /** * Remove a connection from subscriptions. */ unsubscribe(traceId: string, response: Response): void { const connections = this.connections.get(traceId); if (connections) { connections.delete(response); if (connections.size === 0) { this.connections.delete(traceId); } } } /** * Broadcast span updates to all subscribers of a traceId. */ broadcast(traceId: string, event: SpanEvent): void { const connections = this.connections.get(traceId); if (!connections || connections.size === 0) { return; } const data = JSON.stringify(event); const messageToSend = `data: ${data}\n\n`; // Note: response.write() doesn't throw on dead sockets in Express - it // returns false and the 'close' event handles cleanup. No try/catch needed. for (const connection of connections) { connection.write(messageToSend); } } /** * Close all connections for a traceId. */ close(traceId: string): void { const connections = this.connections.get(traceId); if (connections) { for (const connection of connections) { try { connection.end(); } catch (error) { // Ignore errors when closing } } this.connections.delete(traceId); } } /** * Get the number of active connections for a traceId. * Useful for testing. */ getConnectionCount(traceId: string): number { return this.connections.get(traceId)?.size ?? 0; } /** * Check if a traceId has any active connections. * Useful for testing. */ hasConnections(traceId: string): boolean { return this.connections.has(traceId); } }