diff --git a/genkit-tools/telemetry-server/src/localFileTraceStore.ts b/genkit-tools/telemetry-server/src/localFileTraceStore.ts index df51494e2f..ee08f0cbe8 100644 --- a/genkit-tools/telemetry-server/src/localFileTraceStore.ts +++ b/genkit-tools/telemetry-server/src/localFileTraceStore.ts @@ -115,18 +115,31 @@ export class LocalFileTraceStore implements TraceStore { async save(id: string, rawTrace: TraceData): Promise { let trace = this.filter(rawTrace); - if (Object.keys(trace.spans).length === 0) { - return; - } + // if everything is filtered, it's probably the root. + const possibleRoot = Object.keys(trace.spans).length === 0; const mutex = this.getMutex(id); await mutex.waitForUnlock(); const release = await mutex.acquire(); try { - const existing = await this.load(id); + const existing = (await this.load(id)) || trace; if (existing) { Object.keys(trace.spans).forEach( (spanId) => (existing.spans[spanId] = trace.spans[spanId]) ); + // If it's one of those weird roots (internal span that we filter) we try to fix + // whoever was referencing it by making them root. + if (possibleRoot) { + Object.keys(existing.spans).forEach((spanId) => { + const span = existing.spans[spanId]; + if ( + possibleRoot && + span.parentSpanId && + !existing.spans[span.parentSpanId] + ) { + delete span.parentSpanId; + } + }); + } existing.displayName = trace.displayName; existing.startTime = trace.startTime; existing.endTime = trace.endTime; @@ -136,7 +149,7 @@ export class LocalFileTraceStore implements TraceStore { path.resolve(this.storeRoot, `${id}`), JSON.stringify(trace) ); - const hasRootSpan = !!Object.values(rawTrace.spans).find( + const hasRootSpan = !!Object.values(trace.spans).find( (s) => !s.parentSpanId ); if (this.index && hasRootSpan) { diff --git a/genkit-tools/telemetry-server/tests/file_store_test.ts b/genkit-tools/telemetry-server/tests/file_store_test.ts index 6d222b3690..9c05069d20 100644 --- a/genkit-tools/telemetry-server/tests/file_store_test.ts +++ b/genkit-tools/telemetry-server/tests/file_store_test.ts @@ -319,6 +319,7 @@ describe('index', () => { name: 'spanB', start: 2345, end: 3456, + status: 0, }, { id: TRACE_ID_1, @@ -326,6 +327,7 @@ describe('index', () => { name: 'spanA', start: 1234, end: 2345, + status: 0, }, ]); }); @@ -374,6 +376,7 @@ describe('index', () => { name: 'flowA', start: 1, end: 2, + status: 0, }, { id: TRACE_ID_1, @@ -381,6 +384,7 @@ describe('index', () => { name: 'flowA', start: 1, end: 2, + status: 0, }, ] ); @@ -399,6 +403,7 @@ describe('index', () => { name: 'flowA', start: 1, end: 2, + status: 0, }, ] ); @@ -418,6 +423,7 @@ describe('index', () => { name: 'flowA', start: 1, end: 2, + status: 0, }, ] ); @@ -469,6 +475,7 @@ describe('index', () => { name: 'flowA', start: 1, end: 2, + status: 0, }, ] ); diff --git a/genkit-tools/telemetry-server/tests/utils.ts b/genkit-tools/telemetry-server/tests/utils.ts index 69e2648fdb..607dede66c 100644 --- a/genkit-tools/telemetry-server/tests/utils.ts +++ b/genkit-tools/telemetry-server/tests/utils.ts @@ -40,6 +40,7 @@ export function span( instrumentationLibrary: { name: 'genkit' }, spanKind: 'INTERNAL', attributes, + status: { code: 0 }, } as SpanData; }