[UI] Fixes for sessions on UI - ensure errors have a session and use 1 session for test key (#10342)

* ui fix use same trace id on test key page

* _get_session_id_for_spend_log

* fix linting

litellm/proxy/spend_tracking/spend_tracking_utils.py

@@ -290,10 +290,9 @@ def get_logging_payload(  # noqa: PLR0915
             proxy_server_request=_get_proxy_server_request_for_spend_logs_payload(
                 metadata=metadata, litellm_params=litellm_params
             ),
-            session_id=(
-                standard_logging_payload.get("trace_id")
-                if standard_logging_payload is not None
-                else None
+            session_id=_get_session_id_for_spend_log(
+                kwargs=kwargs,
+                standard_logging_payload=standard_logging_payload,
             ),
         )
 
@@ -310,6 +309,32 @@ def get_logging_payload(  # noqa: PLR0915
         raise e
 
 
+def _get_session_id_for_spend_log(
+    kwargs: dict,
+    standard_logging_payload: Optional[StandardLoggingPayload],
+) -> str:
+    """
+    Get the session id for the spend log.
+
+    This ensures each spend log is associated with a unique session id.
+
+    """
+    import uuid
+
+    if (
+        standard_logging_payload is not None
+        and standard_logging_payload.get("trace_id") is not None
+    ):
+        return str(standard_logging_payload.get("trace_id"))
+
+    # Users can dynamically set the trace_id for each request by passing `litellm_trace_id` in kwargs
+    if kwargs.get("litellm_trace_id") is not None:
+        return str(kwargs.get("litellm_trace_id"))
+
+    # Ensure we always have a session id, if none is provided
+    return str(uuid.uuid4())
+
+
 def _ensure_datetime_utc(timestamp: datetime) -> datetime:
     """Helper to ensure datetime is in UTC"""
     timestamp = timestamp.astimezone(timezone.utc)

ui/litellm-dashboard/package-lock.json

@@ -17,6 +17,7 @@
         "@tanstack/react-table": "^8.20.6",
         "@tremor/react": "^3.13.3",
         "@types/papaparse": "^5.3.15",
+        "@types/uuid": "^10.0.0",
         "antd": "^5.13.2",
         "cva": "^1.0.0-beta.3",
         "fs": "^0.0.1-security",
@@ -31,7 +32,8 @@
         "react-dom": "^18",
         "react-markdown": "^9.0.1",
         "react-syntax-highlighter": "^15.6.1",
-        "tailwind-merge": "^3.2.0"
+        "tailwind-merge": "^3.2.0",
+        "uuid": "^11.1.0"
       },
       "devDependencies": {
         "@tailwindcss/forms": "^0.5.7",
@@ -1100,6 +1102,11 @@
       "resolved": "https://registry.npmjs.org/@types/unist/-/unist-3.0.2.tgz",
       "integrity": "sha512-dqId9J8K/vGi5Zr7oo212BGii5m3q5Hxlkwy3WpYuKPklmBEvsbMYYyLxAQpSffdLl/gdW0XUpKWFvYmyoWCoQ=="
     },
+    "node_modules/@types/uuid": {
+      "version": "10.0.0",
+      "resolved": "https://registry.npmjs.org/@types/uuid/-/uuid-10.0.0.tgz",
+      "integrity": "sha512-7gqG38EyHgyP1S+7+xomFtL+ZNHcKv6DwNaCZmJmo1vgMugyF3TCnXVg4t1uk89mLNwnLtnY3TpOpCOyp1/xHQ=="
+    },
     "node_modules/@typescript-eslint/parser": {
       "version": "6.19.1",
       "resolved": "https://registry.npmjs.org/@typescript-eslint/parser/-/parser-6.19.1.tgz",
@@ -7687,6 +7694,18 @@
       "resolved": "https://registry.npmjs.org/util-deprecate/-/util-deprecate-1.0.2.tgz",
       "integrity": "sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw=="
     },
+    "node_modules/uuid": {
+      "version": "11.1.0",
+      "resolved": "https://registry.npmjs.org/uuid/-/uuid-11.1.0.tgz",
+      "integrity": "sha512-0/A9rDy9P7cJ+8w1c9WD9V//9Wj15Ce2MPz8Ri6032usz+NfePxx5AcN3bN+r6ZL6jEo066/yNYB3tn4pQEx+A==",
+      "funding": [
+        "https://github.com/sponsors/broofa",
+        "https://github.com/sponsors/ctavan"
+      ],
+      "bin": {
+        "uuid": "dist/esm/bin/uuid"
+      }
+    },
     "node_modules/vfile": {
       "version": "6.0.1",
       "resolved": "https://registry.npmjs.org/vfile/-/vfile-6.0.1.tgz",

ui/litellm-dashboard/package.json

@@ -18,6 +18,7 @@
     "@tanstack/react-table": "^8.20.6",
     "@tremor/react": "^3.13.3",
     "@types/papaparse": "^5.3.15",
+    "@types/uuid": "^10.0.0",
     "antd": "^5.13.2",
     "cva": "^1.0.0-beta.3",
     "fs": "^0.0.1-security",
@@ -32,7 +33,8 @@
     "react-dom": "^18",
     "react-markdown": "^9.0.1",
     "react-syntax-highlighter": "^15.6.1",
-    "tailwind-merge": "^3.2.0"
+    "tailwind-merge": "^3.2.0",
+    "uuid": "^11.1.0"
   },
   "devDependencies": {
     "@tailwindcss/forms": "^0.5.7",

ui/litellm-dashboard/src/components/chat_ui.tsx

@@ -22,6 +22,7 @@ import {
   Button,
   Divider,
 } from "@tremor/react";
+import { v4 as uuidv4 } from 'uuid';
 
 import { message, Select, Spin, Typography, Tooltip, Input } from "antd";
 import { makeOpenAIChatCompletionRequest } from "./chat_ui/llm_calls/chat_completion";
@@ -82,6 +83,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
   const [isLoading, setIsLoading] = useState<boolean>(false);
   const abortControllerRef = useRef<AbortController | null>(null);
   const [selectedTags, setSelectedTags] = useState<string[]>([]);
+  const [messageTraceId, setMessageTraceId] = useState<string | null>(null);
 
   const chatEndRef = useRef<HTMLDivElement>(null);
 
@@ -300,6 +302,12 @@ const ChatUI: React.FC<ChatUIProps> = ({
     // Create message object without model field for API call
     const newUserMessage = { role: "user", content: inputMessage };
     
+    // Generate new trace ID for a new conversation or use existing one
+    const traceId = messageTraceId || uuidv4();
+    if (!messageTraceId) {
+      setMessageTraceId(traceId);
+    }
+    
     // Update UI with full message object
     setChatHistory([...chatHistory, newUserMessage]);
     setIsLoading(true);
@@ -319,7 +327,8 @@ const ChatUI: React.FC<ChatUIProps> = ({
             signal,
             updateReasoningContent,
             updateTimingData,
-            updateUsageData
+            updateUsageData,
+            traceId
           );
         } else if (endpointType === EndpointType.IMAGE) {
           // For image generation
@@ -344,7 +353,8 @@ const ChatUI: React.FC<ChatUIProps> = ({
             signal,
             updateReasoningContent,
             updateTimingData,
-            updateUsageData
+            updateUsageData,
+            traceId
           );
         }
       }
@@ -365,6 +375,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
 
   const clearChatHistory = () => {
     setChatHistory([]);
+    setMessageTraceId(null);
     message.success("Chat history cleared.");
   };
 

ui/litellm-dashboard/src/components/chat_ui/llm_calls/chat_completion.tsx

@@ -12,7 +12,8 @@ export async function makeOpenAIChatCompletionRequest(
     signal?: AbortSignal,
     onReasoningContent?: (content: string) => void,
     onTimingData?: (timeToFirstToken: number) => void,
-    onUsageData?: (usage: TokenUsage) => void
+    onUsageData?: (usage: TokenUsage) => void,
+    traceId?: string
   ) {
     // base url should be the current base_url
     const isLocal = process.env.NODE_ENV === "development";
@@ -23,11 +24,18 @@ export async function makeOpenAIChatCompletionRequest(
     const proxyBaseUrl = isLocal
       ? "http://localhost:4000"
       : window.location.origin;
+      
+    // Prepare headers with tags and trace ID
+    const headers: Record<string, string> = {};
+    if (tags && tags.length > 0) {
+      headers['x-litellm-tags'] = tags.join(',');
+    }
+    
     const client = new openai.OpenAI({
       apiKey: accessToken,
       baseURL: proxyBaseUrl,
       dangerouslyAllowBrowser: true,
-      defaultHeaders: tags && tags.length > 0 ? { 'x-litellm-tags': tags.join(',') } : undefined,
+      defaultHeaders: headers,
     });
   
     try {
@@ -45,6 +53,7 @@ export async function makeOpenAIChatCompletionRequest(
         stream_options: {
           include_usage: true,
         },
+        litellm_trace_id: traceId,
         messages: chatHistory as ChatCompletionMessageParam[],
       }, { signal });
   

ui/litellm-dashboard/src/components/chat_ui/llm_calls/responses_api.tsx

@@ -12,7 +12,8 @@ export async function makeOpenAIResponsesRequest(
   signal?: AbortSignal,
   onReasoningContent?: (content: string) => void,
   onTimingData?: (timeToFirstToken: number) => void,
-  onUsageData?: (usage: TokenUsage) => void
+  onUsageData?: (usage: TokenUsage) => void,
+  traceId?: string
 ) {
   if (!accessToken) {
     throw new Error("API key is required");
@@ -28,11 +29,17 @@ export async function makeOpenAIResponsesRequest(
     ? "http://localhost:4000"
     : window.location.origin;
   
+  // Prepare headers with tags and trace ID
+  const headers: Record<string, string> = {};
+  if (tags && tags.length > 0) {
+    headers['x-litellm-tags'] = tags.join(',');
+  }
+  
   const client = new openai.OpenAI({
     apiKey: accessToken,
     baseURL: proxyBaseUrl,
     dangerouslyAllowBrowser: true,
-    defaultHeaders: tags && tags.length > 0 ? { 'x-litellm-tags': tags.join(',') } : undefined,
+    defaultHeaders: headers,
   });
 
   try {
@@ -52,6 +59,7 @@ export async function makeOpenAIResponsesRequest(
       model: selectedModel,
       input: formattedInput,
       stream: true,
+      litellm_trace_id: traceId,
     }, { signal });
 
     for await (const event of response) {
@@ -60,7 +68,7 @@ export async function makeOpenAIResponsesRequest(
       // Use a type-safe approach to handle events
       if (typeof event === 'object' && event !== null) {
         // Handle output text delta
-        // 1) drop any “role” streams
+        // 1) drop any "role" streams
         if (event.type === "response.role.delta") {
             continue;
         }