@genkit-ai/telemetry-server
Version:
Genkit AI telemetry server
235 lines • 8.69 kB
JavaScript
import { TraceDataSchema, TraceQueryFilterSchema, } from '@genkit-ai/tools-common';
import { logger } from '@genkit-ai/tools-common/utils';
import cors from 'cors';
import express from 'express';
import { BroadcastManager } from './broadcast-manager.js';
import { logDataFromOtlp, traceDataFromOtlp } from './utils/otlp.js';
export { LocalFileLogStore } from './file-log-store.js';
export { LocalFileTraceStore } from './file-trace-store.js';
export { TraceQuerySchema } from './types';
let server;
const broadcastManager = new BroadcastManager();
export async function startTelemetryServer(params) {
await params.traceStore.init();
await params.logStore.init();
const api = express();
api.use(cors({
origin: params.corsOrigin || /^http:\/\/localhost:\d+$/,
allowedHeaders: ['Content-Type'],
exposedHeaders: ['X-Genkit-Trace-Id'],
}));
api.use(express.json({ limit: params.maxRequestBodySize ?? '100mb' }));
api.get('/api/__health', async (_, response) => {
response.status(200).send('OK');
});
api.get('/api/traces/:traceId', async (request, response, next) => {
try {
const { traceId } = request.params;
response.json(await params.traceStore.load(traceId));
}
catch (e) {
next(e);
}
});
api.get('/api/traces/:traceId/stream', async (request, response, next) => {
try {
const { traceId } = request.params;
response.setHeader('Content-Type', 'text/event-stream');
response.setHeader('Cache-Control', 'no-cache');
response.setHeader('Connection', 'keep-alive');
response.setHeader('Access-Control-Allow-Origin', '*');
response.setHeader('Access-Control-Allow-Headers', 'Content-Type');
const currentTrace = await params.traceStore.load(traceId);
if (currentTrace) {
const snapshot = JSON.stringify(currentTrace);
response.write(`data: ${snapshot}\n\n`);
}
broadcastManager.subscribe(traceId, response);
response.on('close', () => {
broadcastManager.unsubscribe(traceId, response);
});
}
catch (e) {
next(e);
}
});
api.post('/api/traces', async (request, response, next) => {
try {
const traceData = TraceDataSchema.parse(request.body);
await params.traceStore.save(traceData.traceId, traceData);
const allSpans = Object.values(traceData.spans);
const events = [];
for (const span of allSpans) {
events.push({
type: 'span_start',
traceId: traceData.traceId,
span,
});
if (span.endTime > 0) {
events.push({
type: 'span_end',
traceId: traceData.traceId,
span,
});
}
}
events.sort((a, b) => {
const aTime = a.type === 'span_start' ? a.span.startTime : a.span.endTime;
const bTime = b.type === 'span_start' ? b.span.startTime : b.span.endTime;
if (aTime !== bTime) {
return aTime - bTime;
}
return a.type === 'span_start' ? -1 : 1;
});
for (const event of events) {
broadcastManager.broadcast(traceData.traceId, event);
}
response.status(200).send('OK');
}
catch (e) {
next(e);
}
});
api.get('/api/traces', async (request, response, next) => {
try {
const { limit, continuationToken, filter } = request.query;
response.json(await params.traceStore.list({
limit: limit ? Number.parseInt(limit.toString()) : 10,
continuationToken: continuationToken
? continuationToken.toString()
: undefined,
filter: filter
? TraceQueryFilterSchema.parse(JSON.parse(filter))
: undefined,
}));
}
catch (e) {
next(e);
}
});
api.get('/api/logs', async (request, response, next) => {
try {
const { limit, continuationToken } = request.query;
response.json(await params.logStore.list({
limit: limit ? Number.parseInt(limit.toString()) : 100,
continuationToken: continuationToken
? continuationToken.toString()
: undefined,
}));
}
catch (e) {
next(e);
}
});
api.get('/api/traces/:traceId/logs', async (request, response, next) => {
try {
const { limit, continuationToken } = request.query;
const { traceId } = request.params;
response.json(await params.logStore.list({
limit: limit ? Number.parseInt(limit.toString()) : 100,
continuationToken: continuationToken
? continuationToken.toString()
: undefined,
traceId,
}));
}
catch (e) {
next(e);
}
});
api.get('/api/traces/:traceId/spans/:spanId/logs', async (request, response, next) => {
try {
const { limit, continuationToken } = request.query;
const { traceId, spanId } = request.params;
response.json(await params.logStore.list({
limit: limit ? Number.parseInt(limit.toString()) : 100,
continuationToken: continuationToken
? continuationToken.toString()
: undefined,
traceId,
spanId,
}));
}
catch (e) {
next(e);
}
});
api.post([
'/api/otlp',
'/api/otlp/v1/traces',
'/api/otlp/v1/logs',
'/api/otlp/v1/metrics',
], async (request, response) => {
try {
if (!request.body.resourceSpans?.length &&
!request.body.resourceLogs?.length) {
response.status(200).json({});
return;
}
const traces = traceDataFromOtlp(request.body);
for (const trace of traces) {
const traceData = TraceDataSchema.parse(trace);
await params.traceStore.save(traceData.traceId, traceData);
for (const [_, span] of Object.entries(traceData.spans)) {
const event = {
type: span.endTime > 0 ? 'span_end' : 'span_start',
traceId: traceData.traceId,
span,
};
broadcastManager.broadcast(traceData.traceId, event);
}
}
if (request.body.resourceLogs?.length) {
const logs = logDataFromOtlp(request.body);
if (logs.length > 0) {
await params.logStore.save(logs);
}
}
response.status(200).json({});
}
catch (err) {
logger.error(`Error processing OTLP payload: ${err}`);
response.status(500).json({
code: 13,
message: 'An internal error occurred while processing the OTLP payload.',
});
}
});
api.use((err, req, res, next) => {
logger.error(err.stack);
const error = err;
const { message, stack } = error;
const errorResponse = {
code: 13,
message,
details: {
stack,
traceId: err.traceId,
},
};
res.status(500).json(errorResponse);
});
server = api.listen(params.port, () => {
logger.info(`Telemetry API running on http://localhost:${params.port}`);
});
server.on('error', (error) => {
logger.error(error);
});
process.on('SIGTERM', async () => await stopTelemetryApi());
}
export async function stopTelemetryApi() {
await Promise.all([
new Promise((resolve) => {
if (server) {
server.close(() => {
logger.debug('Telemetry API has succesfully shut down.');
resolve();
});
}
else {
resolve();
}
}),
]);
}
//# sourceMappingURL=index.js.map