refactor types

Author: betegon

packages/core/src/integrations/mcp-server/correlation.ts

@@ -4,27 +4,20 @@
  */
 
 import { getClient } from '../../currentScopes';
-import { withActiveSpan } from '../../tracing';
-import type { Span } from '../../types-hoist/span';
+import { SPAN_STATUS_ERROR, withActiveSpan } from '../../tracing';
+import type { Span, SpanAttributeValue } from '../../types-hoist/span';
 import {
   MCP_TOOL_RESULT_CONTENT_ATTRIBUTE,
   MCP_TOOL_RESULT_CONTENT_COUNT_ATTRIBUTE,
   MCP_TOOL_RESULT_IS_ERROR_ATTRIBUTE,
 } from './attributes';
 import { captureError } from './errorCapture';
 import { filterMcpPiiFromSpanData } from './piiFiltering';
-import type { RequestId, SessionId } from './types';
+import type { RequestId, RequestSpanMapValue, SessionId } from './types';
 
 // Simplified correlation system that works with or without sessionId
 // Maps requestId directly to span data for stateless operation
-const requestIdToSpanMap = new Map<
-  RequestId,
-  {
-    span: Span;
-    method: string;
-    startTime: number;
-  }
->();
+const requestIdToSpanMap = new Map<RequestId, RequestSpanMapValue>();
 
 /**
  * Stores span context for later correlation with handler execution
@@ -69,37 +62,28 @@ export function completeSpanWithResults(requestId: RequestId, result: unknown):
   if (spanData) {
     const { span, method } = spanData;
 
-    const spanWithMethods = span as Span & {
-      setAttributes: (attrs: Record<string, unknown>) => void;
-      setStatus: (status: { code: number; message: string }) => void;
-      end: () => void;
-    };
-
-    if (spanWithMethods.setAttributes && method === 'tools/call') {
+    if (method === 'tools/call') {
       // Add tool-specific attributes with PII filtering
       const rawToolAttributes = extractToolResultAttributes(result);
       const client = getClient();
       const sendDefaultPii = Boolean(client?.getOptions().sendDefaultPii);
       const toolAttributes = filterMcpPiiFromSpanData(rawToolAttributes, sendDefaultPii);
 
-      spanWithMethods.setAttributes(toolAttributes);
+      span.setAttributes(toolAttributes);
 
       const isToolError = rawToolAttributes[MCP_TOOL_RESULT_IS_ERROR_ATTRIBUTE] === true;
 
       if (isToolError) {
-        spanWithMethods.setStatus({
-          code: 2, // ERROR
-          message: 'Tool execution failed',
+        span.setStatus({
+          code: SPAN_STATUS_ERROR,
+          message: 'Tool returned error result',
         });
 
         captureError(new Error('Tool returned error result'), 'tool_execution');
       }
     }
 
-    if (spanWithMethods.end) {
-      spanWithMethods.end();
-    }
-
+    span.end();
     requestIdToSpanMap.delete(requestId);
   }
 }
@@ -111,17 +95,11 @@ export function cleanupAllPendingSpans(): number {
   const pendingCount = requestIdToSpanMap.size;
 
   for (const [, spanData] of requestIdToSpanMap) {
-    const spanWithEnd = spanData.span as Span & {
-      end: () => void;
-      setStatus: (status: { code: number; message: string }) => void;
-    };
-    if (spanWithEnd.setStatus && spanWithEnd.end) {
-      spanWithEnd.setStatus({
-        code: 2, // ERROR
-        message: 'Transport closed before request completion',
-      });
-      spanWithEnd.end();
-    }
+    spanData.span.setStatus({
+      code: SPAN_STATUS_ERROR,
+      message: 'Transport closed before request completion',
+    });
+    spanData.span.end();
   }
 
   requestIdToSpanMap.clear();
@@ -131,8 +109,8 @@ export function cleanupAllPendingSpans(): number {
 /**
  * Simplified tool result attribute extraction
  */
-function extractToolResultAttributes(result: unknown): Record<string, string | number | boolean> {
-  const attributes: Record<string, string | number | boolean> = {};
+function extractToolResultAttributes(result: unknown): Record<string, SpanAttributeValue | undefined> {
+  const attributes: Record<string, SpanAttributeValue | undefined> = {};
 
   if (typeof result === 'object' && result !== null) {
     const resultObj = result as Record<string, unknown>;

packages/core/src/integrations/mcp-server/piiFiltering.ts

@@ -2,7 +2,7 @@
  * PII filtering for MCP server spans
  * Removes sensitive data when sendDefaultPii is false
  */
-
+import type { SpanAttributeValue } from '../../types-hoist/span';
 import { MCP_TOOL_RESULT_CONTENT_ATTRIBUTE } from './attributes';
 
 /** PII attributes that should be removed when sendDefaultPii is false */
@@ -20,9 +20,9 @@ const PII_ATTRIBUTES = new Set([
 export function filterMcpPiiFromSpanData(
   spanData: Record<string, unknown>,
   sendDefaultPii: boolean,
-): Record<string, unknown> {
+): Record<string, SpanAttributeValue | undefined> {
   if (sendDefaultPii) {
-    return spanData; // Send everything when PII is allowed
+    return spanData as Record<string, SpanAttributeValue | undefined>; // Cast for type safety
   }
 
   // Use Object.fromEntries with filter for a more functional approach
@@ -31,5 +31,5 @@ export function filterMcpPiiFromSpanData(
       const isPiiAttribute = PII_ATTRIBUTES.has(key) || key.startsWith('mcp.request.argument.');
       return !isPiiAttribute;
     }),
-  );
+  ) as Record<string, SpanAttributeValue | undefined>;
 }

packages/core/src/integrations/mcp-server/types.ts

@@ -1,3 +1,5 @@
+import type { Span } from '../../types-hoist/span';
+
 /**
  * types for MCP server instrumentation
  */
@@ -143,6 +145,14 @@ export interface McpSpanConfig {
 
 export type SessionId = string;
 export type RequestId = string | number;
+
+// Type for the value stored in requestIdToSpanMap
+export type RequestSpanMapValue = {
+  span: Span;
+  method: string;
+  startTime: number;
+};
+
 export type MCPHandler = (...args: unknown[]) => unknown;
 export interface HandlerExtraData {
   sessionId?: SessionId;