🐛 fix(model-runtime): include tool_calls in speed metrics & add getActiveTraceId (#11927)

packages/model-runtime/src/core/streams/protocol.ts

@@ -472,12 +472,14 @@ export const createTokenSpeedCalculator = (
     // - text/reasoning: standard text output events
     // - content_part/reasoning_part: multimodal output events used by Gemini 3+ models
     //   which emit structured parts instead of plain text events
+    // - tool_calls: function calling output events
     if (
       !outputStartAt &&
       (chunk.type === 'text' ||
         chunk.type === 'reasoning' ||
         chunk.type === 'content_part' ||
-        chunk.type === 'reasoning_part')
+        chunk.type === 'reasoning_part' ||
+        chunk.type === 'tool_calls')
     ) {
       outputStartAt = Date.now();
     }

src/layout/GlobalProvider/useUserStateRedirect.ts

@@ -73,9 +73,13 @@ export const useWebUserStateRedirect = () =>
     }
 
     // Redirect away from invite-code page if no longer required
+    // Skip redirect if force=true is present (for re-entering invite code)
     if (pathname.startsWith('/invite-code')) {
-      window.location.href = '/';
-      return;
+      const params = new URLSearchParams(window.location.search);
+      if (params.get('force') !== 'true') {
+        window.location.href = '/';
+        return;
+      }
     }
 
     if (!onboardingSelectors.needsOnboarding(state)) return;

src/libs/observability/traceparent.test.ts

@@ -1,22 +1,23 @@
 import type { Mock } from 'vitest';
 import { afterEach, describe, expect, it, vi } from 'vitest';
 
+// eslint-disable-next-line import/first
+import { getActiveTraceId, injectSpanTraceHeaders } from './traceparent';
+
 vi.mock('@lobechat/observability-otel/api', () => {
   const inject = vi.fn();
   const setSpan = vi.fn((_ctx, span) => span);
+  const getActiveSpan = vi.fn();
 
   return {
     context: {
       active: vi.fn(() => ({})),
     },
     propagation: { inject },
-    trace: { setSpan },
+    trace: { getActiveSpan, setSpan },
   };
 });
 
-// eslint-disable-next-line import/first
-import { injectSpanTraceHeaders } from './traceparent';
-
 const mockSpan = (traceId: string, spanId: string) =>
   ({
     spanContext: () => ({
@@ -39,7 +40,9 @@ describe('injectSpanTraceHeaders', () => {
 
   it('uses propagator output when available', async () => {
     const { propagation } = await api;
-    (propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>).mockImplementation((_ctx, carrier) => {
+    (
+      propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>
+    ).mockImplementation((_ctx, carrier) => {
       carrier.traceparent = 'from-propagator';
       carrier.tracestate = 'state';
     });
@@ -56,7 +59,9 @@ describe('injectSpanTraceHeaders', () => {
 
   it('falls back to manual traceparent formatting when propagator gives none', async () => {
     const { propagation } = await api;
-    (propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>).mockImplementation(() => undefined);
+    (
+      propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>
+    ).mockImplementation(() => undefined);
 
     const headers = headersWith();
     const span = mockSpan('1'.repeat(32), '2'.repeat(16));
@@ -64,6 +69,40 @@ describe('injectSpanTraceHeaders', () => {
     const tp = injectSpanTraceHeaders(headers, span);
 
     expect(tp).toBe('00-11111111111111111111111111111111-2222222222222222-01');
-    expect(headers.get('traceparent')).toBe('00-11111111111111111111111111111111-2222222222222222-01');
+    expect(headers.get('traceparent')).toBe(
+      '00-11111111111111111111111111111111-2222222222222222-01',
+    );
+  });
+});
+
+describe('getActiveTraceId', () => {
+  const api = vi.importMock<typeof import('@lobechat/observability-otel/api')>(
+    '@lobechat/observability-otel/api',
+  );
+
+  afterEach(() => {
+    vi.resetAllMocks();
+  });
+
+  it('returns traceId from active span', async () => {
+    const { trace } = await api;
+    const expectedTraceId = 'a'.repeat(32);
+    (trace.getActiveSpan as Mock).mockReturnValue(mockSpan(expectedTraceId, 'b'.repeat(16)));
+
+    expect(getActiveTraceId()).toBe(expectedTraceId);
+  });
+
+  it('returns undefined when no active span', async () => {
+    const { trace } = await api;
+    (trace.getActiveSpan as Mock).mockReturnValue(undefined);
+
+    expect(getActiveTraceId()).toBeUndefined();
+  });
+
+  it('returns undefined when traceId is all zeros', async () => {
+    const { trace } = await api;
+    (trace.getActiveSpan as Mock).mockReturnValue(mockSpan('0'.repeat(32), 'b'.repeat(16)));
+
+    expect(getActiveTraceId()).toBeUndefined();
   });
 });

src/libs/observability/traceparent.ts

@@ -1,13 +1,5 @@
-import type {
-  Span,
-  Context as OtContext,
-  TextMapGetter
-} from '@lobechat/observability-otel/api';
-import {
-  context as otContext,
-  propagation,
-  trace,
-} from '@lobechat/observability-otel/api';
+import type { Context as OtContext, Span, TextMapGetter } from '@lobechat/observability-otel/api';
+import { context as otContext, propagation, trace } from '@lobechat/observability-otel/api';
 
 // NOTICE: do not try to optimize this into .repeat(...) or similar,
 // here served for better search / semantic search purpose for further diagnostic
@@ -47,6 +39,16 @@ export const getActiveTraceparent = () => {
   return toTraceparent(span as Span);
 };
 
+/**
+ * Get the traceId from the active span.
+ */
+export const getActiveTraceId = () => {
+  const span = trace.getActiveSpan();
+  if (!isValidContext(span)) return undefined;
+
+  return span!.spanContext().traceId;
+};
+
 /**
  * Injects the active context into headers using the configured propagator (W3C by default).
  * Also returns the traceparent for convenience.

src/server/services/memory/userMemory/persona/__tests__/service.test.ts

@@ -14,12 +14,14 @@ var aiInfraMocks:
   | {
       getAiProviderRuntimeState: ReturnType<typeof vi.fn>;
       tryMatchingModelFrom: ReturnType<typeof vi.fn>;
+      tryMatchingProviderFrom: ReturnType<typeof vi.fn>;
     };
 
 vi.mock('@/database/repositories/aiInfra', () => {
   aiInfraMocks = {
     getAiProviderRuntimeState: vi.fn(),
     tryMatchingModelFrom: vi.fn(),
+    tryMatchingProviderFrom: vi.fn(),
   };
 
   const AiInfraRepos = vi.fn().mockImplementation(() => ({
@@ -27,6 +29,7 @@ vi.mock('@/database/repositories/aiInfra', () => {
   })) as unknown as typeof import('@/database/repositories/aiInfra').AiInfraRepos;
 
   (AiInfraRepos as any).tryMatchingModelFrom = aiInfraMocks!.tryMatchingModelFrom;
+  (AiInfraRepos as any).tryMatchingProviderFrom = aiInfraMocks!.tryMatchingProviderFrom;
 
   return { AiInfraRepos };
 });
@@ -85,7 +88,9 @@ beforeEach(async () => {
   toolCall.mockClear();
   aiInfraMocks!.getAiProviderRuntimeState.mockReset();
   aiInfraMocks!.tryMatchingModelFrom.mockReset();
+  aiInfraMocks!.tryMatchingProviderFrom.mockReset();
   aiInfraMocks!.tryMatchingModelFrom.mockResolvedValue('openai');
+  aiInfraMocks!.tryMatchingProviderFrom.mockResolvedValue('openai');
   aiInfraMocks!.getAiProviderRuntimeState.mockResolvedValue({
     enabledAiModels: [
       { abilities: {}, enabled: true, id: 'gpt-mock', providerId: 'openai', type: 'chat' },