feat(js): Add telemetry server stream endpoint and Update tools server

genkit-tools/cli/src/commands/start.ts

@@ -26,6 +26,7 @@ interface RunOptions {
  noui?: boolean;
  port?: string;
  open?: boolean;
+ experimentalRealtimeTelemetry?: boolean;
 }
 
 /** Command to run code in dev mode and/or the Dev UI. */
@@ -34,6 +35,10 @@ export const start = new Command('start')
  .option('-n, --noui', 'do not start the Dev UI', false)
  .option('-p, --port <port>', 'port for the Dev UI')
  .option('-o, --open', 'Open the browser on UI start up')
+ .option(
+ '--experimental-realtime-telemetry',
+ 'Enable real-time telemetry streaming (experimental)'
+ )
  .action(async (options: RunOptions) => {
  const projectRoot = await findProjectRoot();
  if (projectRoot.includes('/.Trash/')) {
@@ -49,7 +54,8 @@ export const start = new Command('start')
  const result = await startDevProcessManager(
  projectRoot,
  start.args[0],
- start.args.slice(1)
+ start.args.slice(1),
+ { enableRealtimeTelemetry: options.experimentalRealtimeTelemetry }
  );
  manager = result.manager;
  processPromise = result.processPromise;

genkit-tools/cli/src/utils/manager-utils.ts

@@ -66,21 +66,31 @@ export async function startManager(
  return manager;
 }
 
+export interface DevProcessManagerOptions {
+ enableRealtimeTelemetry?: boolean;
+}
+
 export async function startDevProcessManager(
  projectRoot: string,
  command: string,
- args: string[]
+ args: string[],
+ options?: DevProcessManagerOptions
 ): Promise<{ manager: RuntimeManager; processPromise: Promise<void> }> {
  const telemetryServerUrl = await resolveTelemetryServer(projectRoot);
- const processManager = new ProcessManager(command, args, {
+ const envVars: Record<string, string> = {
  GENKIT_TELEMETRY_SERVER: telemetryServerUrl,
  GENKIT_ENV: 'dev',
- });
+ };
+ if (options?.enableRealtimeTelemetry) {
+ envVars.GENKIT_ENABLE_REALTIME_TELEMETRY = 'true';
+ }
+ const processManager = new ProcessManager(command, args, envVars);
  const manager = await RuntimeManager.create({
  telemetryServerUrl,
  manageHealth: true,
  projectRoot,
  processManager,
+ enableRealtimeTelemetry: options?.enableRealtimeTelemetry,
  });
  const processPromise = processManager.start();
  return { manager, processPromise };

genkit-tools/common/package.json

@@ -10,6 +10,7 @@
  },
  "dependencies": {
  "@asteasolutions/zod-to-openapi": "^7.0.0",
+ "@inquirer/prompts": "^7.8.0",
  "@trpc/server": "^10.45.2",
  "adm-zip": "^0.5.12",
  "ajv": "^8.12.0",
@@ -20,10 +21,10 @@
  "colorette": "^2.0.20",
  "commander": "^11.1.0",
  "configstore": "^5.0.1",
+ "cors": "^2.8.5",
  "express": "^4.21.0",
  "get-port": "5.1.1",
  "glob": "^10.3.12",
- "@inquirer/prompts": "^7.8.0",
  "js-yaml": "^4.1.0",
  "json-2-csv": "^5.5.1",
  "json-schema": "^0.4.0",
@@ -41,6 +42,7 @@
  "@types/body-parser": "^1.19.5",
  "@types/cli-color": "^2.0.6",
  "@types/configstore": "^6.0.2",
+ "@types/cors": "^2.8.19",
  "@types/express": "^4.17.21",
  "@types/inquirer": "^8.1.3",
  "@types/jest": "^29.5.12",

genkit-tools/common/src/manager/manager.ts

@@ -60,10 +60,13 @@ interface RuntimeManagerOptions {
  projectRoot: string;
  /** An optional process manager for the main application process. */
  processManager?: ProcessManager;
+ /** Whether realtime telemetry streaming is enabled. */
+ enableRealtimeTelemetry?: boolean;
 }
 
 export class RuntimeManager {
  readonly processManager?: ProcessManager;
+ readonly enableRealtimeTelemetry: boolean;
  private filenameToRuntimeMap: Record<string, RuntimeInfo> = {};
  private filenameToDevUiMap: Record<string, DevToolsInfo> = {};
  private idToFileMap: Record<string, string> = {};
@@ -75,9 +78,11 @@ export class RuntimeManager {
  readonly telemetryServerUrl: string | undefined,
  private manageHealth: boolean,
  readonly projectRoot: string,
- processManager?: ProcessManager
+ processManager?: ProcessManager,
+ enableRealtimeTelemetry?: boolean
  ) {
  this.processManager = processManager;
+ this.enableRealtimeTelemetry = enableRealtimeTelemetry ?? false;
  }
 
  /**
@@ -88,7 +93,8 @@ export class RuntimeManager {
  options.telemetryServerUrl,
  options.manageHealth ?? true,
  options.projectRoot,
- options.processManager
+ options.processManager,
+ options.enableRealtimeTelemetry
  );
  await manager.setupRuntimesWatcher();
  await manager.setupDevUiWatcher();
@@ -197,7 +203,8 @@ export class RuntimeManager {
  */
  async runAction(
  input: apis.RunActionRequest,
- streamingCallback?: StreamingCallback<any>
+ streamingCallback?: StreamingCallback<any>,
+ onTraceId?: (traceId: string) => void
  ): Promise<RunActionResponse> {
  const runtime = input.runtimeId
  ? this.getRuntimeById(input.runtimeId)
@@ -226,6 +233,12 @@ export class RuntimeManager {
  if (response.headers['x-genkit-version']) {
  genkitVersion = response.headers['x-genkit-version'];
  }
+
+ const traceId = response.headers['x-genkit-trace-id'];
+ if (traceId && onTraceId) {
+ onTraceId(traceId);
+ }
+
  const stream = response.data;
 
  let buffer = '';
@@ -280,20 +293,115 @@ export class RuntimeManager {
  });
  return promise;
  } else {
+ // runAction should use chunked JSON streaming to send early headers
  const response = await axios
  .post(`${runtime.reflectionServerUrl}/api/runAction`, input, {
  headers: {
  'Content-Type': 'application/json',
  },
+ responseType: 'stream', // Use stream to get early headers
  })
  .catch((err) =>
  this.httpErrorHandler(err, `Error running action key='${input.key}'.`)
  );
- const resp = RunActionResponseSchema.parse(response.data);
- if (response.headers['x-genkit-version']) {
- resp.genkitVersion = response.headers['x-genkit-version'];
+
+ const traceId = response.headers['x-genkit-trace-id'];
+ if (traceId && onTraceId) {
+ onTraceId(traceId);
  }
- return resp;
+
+ return new Promise<RunActionResponse>((resolve, reject) => {
+ let buffer = '';
+
+ response.data.on('data', (chunk: Buffer) => {
+ buffer += chunk.toString();
+ });
+
+ response.data.on('end', () => {
+ try {
+ const responseData = JSON.parse(buffer);
+
+ if (responseData.error) {
+ const err = new GenkitToolsError(
+ `Error running action key='${input.key}'.`
+ );
+ // massage the error into a shape dev ui expects
+ err.data = {
+ ...responseData.error,
+ stack: (responseData.error?.details as any).stack,
+ data: {
+ genkitErrorMessage: responseData.error?.message,
+ genkitErrorDetails: responseData.error?.details,
+ },
+ };
+ reject(err);
+ return;
+ }
+
+ // Handle backward compatibility - add trace ID from header if not in body
+ if (!responseData.telemetry && traceId) {
+ responseData.telemetry = { traceId: traceId };
+ }
+
+ const parsed = RunActionResponseSchema.parse(responseData);
+ if (response.headers['x-genkit-version']) {
+ parsed.genkitVersion = response.headers['x-genkit-version'];
+ }
+ resolve(parsed);
+ } catch (err) {
+ reject(new GenkitToolsError(`Failed to parse response: ${err}`));
+ }
+ });
+
+ response.data.on('error', (err: Error) => {
+ reject(err);
+ });
+ });
+ }
+ }
+
+ /**
+ * Cancels an in-flight action by trace ID
+ */
+ async cancelAction(input: {
+ traceId: string;
+ runtimeId?: string;
+ }): Promise<{ message: string }> {
+ const runtime = input.runtimeId
+ ? this.getRuntimeById(input.runtimeId)
+ : this.getMostRecentRuntime();
+ if (!runtime) {
+ throw new Error(
+ input.runtimeId
+ ? `No runtime found with ID ${input.runtimeId}.`
+ : 'No runtimes found. Make sure your app is running.'
+ );
+ }
+
+ try {
+ const response = await axios.post(
+ `${runtime.reflectionServerUrl}/api/cancelAction`,
+ { traceId: input.traceId },
+ {
+ headers: {
+ 'Content-Type': 'application/json',
+ },
+ }
+ );
+ return response.data;
+ } catch (err) {
+ const axiosError = err as AxiosError;
+ if (axiosError.response?.status === 404) {
+ const error = new GenkitToolsError(
+ 'Action not found or already completed'
+ );
+ error.data = {
+ message: 'Action not found or already completed',
+ } as any;
+ (error.data as any).statusCode = 404;
+ throw error;
+ }
+ throw this.httpErrorHandler(axiosError);
  }
  }
 
@@ -347,6 +455,78 @@ export class RuntimeManager {
  return response.data as TraceData;
  }
 
+ /**
+ * Streams trace updates in real-time from the telemetry server.
+ * Connects to the telemetry server's SSE endpoint and forwards updates via callback.
+ */
+ async streamTrace(
+ input: apis.StreamTraceRequest,
+ streamingCallback: StreamingCallback<any>
+ ): Promise<void> {
+ const { traceId } = input;
+
+ if (!this.telemetryServerUrl) {
+ throw new Error(
+ 'Telemetry server URL not configured. Cannot stream trace updates.'
+ );
+ }
+
+ const response = await axios
+ .get(`${this.telemetryServerUrl}/api/traces/${traceId}/stream`, {
+ headers: {
+ Accept: 'text/event-stream',
+ },
+ responseType: 'stream',
+ })
+ .catch((err) =>
+ this.httpErrorHandler(
+ err,
+ `Error streaming trace for traceId='${traceId}'`
+ )
+ );
+
+ const stream = response.data;
+ let buffer = '';
+
+ // Return a promise that resolves when the stream ends
+ return new Promise<void>((resolve, reject) => {
+ stream.on('data', (chunk: Buffer) => {
+ buffer += chunk.toString();
+
+ // Process complete messages (ending with \n\n)
+ while (buffer.includes('\n\n')) {
+ const messageEnd = buffer.indexOf('\n\n');
+ const message = buffer.substring(0, messageEnd).trim();
+ buffer = buffer.substring(messageEnd + 2);
+
+ // Skip empty messages
+ if (!message) {
+ continue;
+ }
+ // Parse SSE data line - strip "data: " prefix
+ try {
+ const jsonData = message.startsWith('data: ')
+ ? message.slice(6)
+ : message;
+ const parsed = JSON.parse(jsonData);
+ streamingCallback(parsed);
+ } catch (err) {
+ logger.error(`Error parsing stream data: ${err}`);
+ }
+ }
+ });
+
+ stream.on('end', () => {
+ resolve();
+ });
+
+ stream.on('error', (err: Error) => {
+ logger.error(`Stream error for traceId='${traceId}': ${err}`);
+ reject(err);
+ });
+ });
+ }
+
  /**
  * Adds a trace to the trace store
  */

genkit-tools/common/src/server/router.ts

@@ -27,6 +27,7 @@ import { GenkitToolsError, type RuntimeInfo } from '../manager/types';
 import { TraceDataSchema } from '../types';
 import type { Action } from '../types/action';
 import * as apis from '../types/apis';
+import { CancelActionRequestSchema } from '../types/apis';
 import type { EnvironmentVariable } from '../types/env';
 import * as evals from '../types/eval';
 import type { PromptFrontmatter } from '../types/prompt';
@@ -133,13 +134,6 @@ export const TOOLS_SERVER_ROUTER = (manager: RuntimeManager) =>
  return manager.listActions(input);
  }),
 
- /** Runs an action. */
- runAction: loggedProcedure
- .input(apis.RunActionRequestSchema)
- .mutation(async ({ input }) => {
- return manager.runAction(input);
- }),
-
  /** Generate a .prompt file from messages and model config. */
  createPrompt: loggedProcedure
  .input(apis.CreatePromptRequestSchema)
@@ -328,6 +322,13 @@ export const TOOLS_SERVER_ROUTER = (manager: RuntimeManager) =>
  await manager.processManager?.kill();
  return true;
  }),
+
+ /** Cancels a long-running action. */
+ cancelAction: loggedProcedure
+ .input(CancelActionRequestSchema)
+ .mutation(async ({ input }) => {
+ return manager.cancelAction(input);
+ }),
  });
 
 export type ToolsServerRouter = ReturnType<typeof TOOLS_SERVER_ROUTER>;