Aryan D.

class JudgeSpanCaptureProcessor implements SpanProcessor {
  // keyed by "traceId:spanName" for concurrent-request safety
  private readonly capturedSpanIds = new Map<string, string>();
  private readonly pendingErrors = new Map<string, string>();

  onStart(_span: Span, _parentContext: Context): void {}

  onEnd(span: ReadableSpan): void {
    const key = `${span.spanContext().traceId}:${span.name}`;
    this.capturedSpanIds.set(key, span.spanContext().spanId);

    // Override status BEFORE BatchSpanProcessor exports it.
    // LangChain's handleChainEnd sets OK after our node runs — we fix it here.
    const errorMessage = this.pendingErrors.get(key);
    if (errorMessage !== undefined) {
      (span as unknown as Span).setStatus({ code: SpanStatusCode.ERROR, message: errorMessage });
      this.pendingErrors.delete(key);
    }
  }

  getSpanId(traceId: string, spanName: string): string | null {
    return this.capturedSpanIds.get(`${traceId}:${spanName}`) ?? null;
  }

  markSpanForError(traceId: string, spanName: string, message: string): void {
    this.pendingErrors.set(`${traceId}:${spanName}`, message);
  }

  clearSpan(traceId: string, spanName: string): void {
    this.capturedSpanIds.delete(`${traceId}:${spanName}`);
  }

  forceFlush(): Promise<void> { return Promise.resolve(); }
  shutdown(): Promise<void> { return Promise.resolve(); }
}