MCP Terminal Server
by dillip285
/**
* 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 {
Span as ApiSpan,
Link,
ROOT_CONTEXT,
SpanStatusCode,
trace,
} from '@opentelemetry/api';
import { performance } from 'node:perf_hooks';
import { HasRegistry, Registry } from '../registry.js';
import { ensureBasicTelemetryInstrumentation } from '../tracing.js';
import { PathMetadata, SpanMetadata, TraceMetadata } from './types.js';
export const spanMetadataAlsKey = 'core.tracing.instrumentation.span';
export const traceMetadataAlsKey = 'core.tracing.instrumentation.trace';
export const ATTR_PREFIX = 'genkit';
/** @hidden */
export const SPAN_TYPE_ATTR = ATTR_PREFIX + ':type';
const TRACER_NAME = 'genkit-tracer';
const TRACER_VERSION = 'v1';
/**
* @hidden
*/
export async function newTrace<T>(
registry: Registry | HasRegistry,
opts: {
name: string;
labels?: Record<string, string>;
links?: Link[];
},
fn: (metadata: SpanMetadata, rootSpan: ApiSpan) => Promise<T>
) {
registry = (registry as HasRegistry).registry
? (registry as HasRegistry).registry
: (registry as Registry);
await ensureBasicTelemetryInstrumentation();
const traceMetadata: TraceMetadata = registry.asyncStore.getStore(
traceMetadataAlsKey
) || {
paths: new Set<PathMetadata>(),
timestamp: performance.now(),
featureName: opts.name,
};
return await registry.asyncStore.run(traceMetadataAlsKey, traceMetadata, () =>
runInNewSpan(
registry,
{
metadata: {
name: opts.name,
},
labels: opts.labels,
links: opts.links,
},
async (metadata, otSpan) => {
return await fn(metadata, otSpan);
}
)
);
}
/**
* Runs the provided function in a new span.
*
* @hidden
*/
export async function runInNewSpan<T>(
registry: Registry | HasRegistry,
opts: {
metadata: SpanMetadata;
labels?: Record<string, string>;
links?: Link[];
},
fn: (metadata: SpanMetadata, otSpan: ApiSpan, isRoot: boolean) => Promise<T>
): Promise<T> {
await ensureBasicTelemetryInstrumentation();
const resolvedRegistry = (registry as HasRegistry).registry
? (registry as HasRegistry).registry
: (registry as Registry);
const tracer = trace.getTracer(TRACER_NAME, TRACER_VERSION);
const parentStep =
resolvedRegistry.asyncStore.getStore<SpanMetadata>(spanMetadataAlsKey);
const isInRoot = parentStep?.isRoot === true;
if (!parentStep) opts.metadata.isRoot ||= true;
return await tracer.startActiveSpan(
opts.metadata.name,
{ links: opts.links, root: opts.metadata.isRoot },
async (otSpan) => {
if (opts.labels) otSpan.setAttributes(opts.labels);
try {
opts.metadata.path = buildPath(
opts.metadata.name,
parentStep?.path || '',
opts.labels
);
const output = await resolvedRegistry.asyncStore.run(
spanMetadataAlsKey,
opts.metadata,
() => fn(opts.metadata, otSpan, isInRoot)
);
if (opts.metadata.state !== 'error') {
opts.metadata.state = 'success';
}
recordPath(resolvedRegistry, opts.metadata);
return output;
} catch (e) {
recordPath(resolvedRegistry, opts.metadata, e);
opts.metadata.state = 'error';
otSpan.setStatus({
code: SpanStatusCode.ERROR,
message: getErrorMessage(e),
});
if (e instanceof Error) {
otSpan.recordException(e);
}
throw e;
} finally {
otSpan.setAttributes(metadataToAttributes(opts.metadata));
otSpan.end();
}
}
);
}
/**
* Creates a new child span and attaches it to a previously created trace. This
* is useful, for example, for adding deferred user engagement metadata.
*
* @hidden
*/
export async function appendSpan(
traceId: string,
parentSpanId: string,
metadata: SpanMetadata,
labels?: Record<string, string>
) {
await ensureBasicTelemetryInstrumentation();
const tracer = trace.getTracer(TRACER_NAME, TRACER_VERSION);
const spanContext = trace.setSpanContext(ROOT_CONTEXT, {
traceId: traceId,
traceFlags: 1, // sampled
spanId: parentSpanId,
});
// TODO(abrook): add explicit start time to align with parent
const span = tracer.startSpan(metadata.name, {}, spanContext);
span.setAttributes(metadataToAttributes(metadata));
if (labels) {
span.setAttributes(labels);
}
span.end();
}
function getErrorMessage(e: any): string {
if (e instanceof Error) {
return e.message;
}
return `${e}`;
}
function metadataToAttributes(metadata: SpanMetadata): Record<string, string> {
const out = {} as Record<string, string>;
Object.keys(metadata).forEach((key) => {
if (
key === 'metadata' &&
typeof metadata[key] === 'object' &&
metadata.metadata
) {
Object.entries(metadata.metadata).forEach(([metaKey, value]) => {
out[ATTR_PREFIX + ':metadata:' + metaKey] = value;
});
} else if (key === 'input' || typeof metadata[key] === 'object') {
out[ATTR_PREFIX + ':' + key] = JSON.stringify(metadata[key]);
} else {
out[ATTR_PREFIX + ':' + key] = metadata[key];
}
});
return out;
}
/**
* Sets provided attribute value in the current span.
*
* @hidden
*/
export function setCustomMetadataAttribute(
registry: Registry,
key: string,
value: string
) {
const currentStep = getCurrentSpan(registry);
if (!currentStep) {
return;
}
if (!currentStep.metadata) {
currentStep.metadata = {};
}
currentStep.metadata[key] = value;
}
/**
* Sets provided attribute values in the current span.
*
* @hidden
*/
export function setCustomMetadataAttributes(
registry: Registry,
values: Record<string, string>
) {
const currentStep = getCurrentSpan(registry);
if (!currentStep) {
return;
}
if (!currentStep.metadata) {
currentStep.metadata = {};
}
for (const [key, value] of Object.entries(values)) {
currentStep.metadata[key] = value;
}
}
/**
* Converts a fully annotated path to a friendly display version for logs
*
* @hidden
*/
export function toDisplayPath(path: string): string {
const pathPartRegex = /\{([^\,}]+),[^\}]+\}/g;
return Array.from(path.matchAll(pathPartRegex), (m) => m[1]).join(' > ');
}
function getCurrentSpan(registry: Registry): SpanMetadata {
const step = registry.asyncStore.getStore<SpanMetadata>(spanMetadataAlsKey);
if (!step) {
throw new Error('running outside step context');
}
return step;
}
function buildPath(
name: string,
parentPath: string,
labels?: Record<string, string>
) {
const stepType =
labels && labels['genkit:type']
? `,t:${labels['genkit:metadata:subtype'] === 'flow' ? 'flow' : labels['genkit:type']}`
: '';
return parentPath + `/{${name}${stepType}}`;
}
function recordPath(registry: Registry, spanMeta: SpanMetadata, err?: any) {
const path = spanMeta.path || '';
const decoratedPath = decoratePathWithSubtype(spanMeta);
// Only add the path if a child has not already been added. In the event that
// an error is rethrown, we don't want to add each step in the unwind.
const paths = Array.from(
registry.asyncStore.getStore<TraceMetadata>(traceMetadataAlsKey)?.paths ||
new Set<PathMetadata>()
);
const status = err ? 'failure' : 'success';
if (!paths.some((p) => p.path.startsWith(path) && p.status === status)) {
const now = performance.now();
const start =
registry.asyncStore.getStore<TraceMetadata>(traceMetadataAlsKey)
?.timestamp || now;
registry.asyncStore
.getStore<TraceMetadata>(traceMetadataAlsKey)
?.paths?.add({
path: decoratedPath,
error: err?.name,
latency: now - start,
status,
});
}
spanMeta.path = decoratedPath;
}
function decoratePathWithSubtype(metadata: SpanMetadata): string {
if (!metadata.path) {
return '';
}
const pathComponents = metadata.path.split('}/{');
if (pathComponents.length == 1) {
return metadata.path;
}
const stepSubtype =
metadata.metadata && metadata.metadata['subtype']
? `,s:${metadata.metadata['subtype']}`
: '';
const root = `${pathComponents.slice(0, -1).join('}/{')}}/`;
const decoratedStep = `{${pathComponents.at(-1)?.slice(0, -1)}${stepSubtype}}`;
return root + decoratedStep;
}