Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit b24da44

Browse files
tjx666tjx666
authored
🐛 fix(model-runtime): include tool_calls in speed metrics & add getActiveTraceId (#11927)
1 parent 74b8fb6 commit b24da44

File tree

5 files changed

+72
-20
lines changed

5 files changed

+72
-20
lines changed

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -472,12 +472,14 @@ export const createTokenSpeedCalculator = (
472472
// - text/reasoning: standard text output events
473473
// - content_part/reasoning_part: multimodal output events used by Gemini 3+ models
474474
// which emit structured parts instead of plain text events
475+
// - tool_calls: function calling output events
475476
if (
476477
!outputStartAt &&
477478
(chunk.type === 'text' ||
478479
chunk.type === 'reasoning' ||
479480
chunk.type === 'content_part' ||
480-
chunk.type === 'reasoning_part')
481+
chunk.type === 'reasoning_part' ||
482+
chunk.type === 'tool_calls')
481483
) {
482484
outputStartAt = Date.now();
483485
}

‎src/layout/GlobalProvider/useUserStateRedirect.ts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,9 +73,13 @@ export const useWebUserStateRedirect = () =>
7373
}
7474

7575
// Redirect away from invite-code page if no longer required
76+
// Skip redirect if force=true is present (for re-entering invite code)
7677
if (pathname.startsWith('/invite-code')) {
77-
window.location.href = '/';
78-
return;
78+
const params = new URLSearchParams(window.location.search);
79+
if (params.get('force') !== 'true') {
80+
window.location.href = '/';
81+
return;
82+
}
7983
}
8084

8185
if (!onboardingSelectors.needsOnboarding(state)) return;

‎src/libs/observability/traceparent.test.ts‎

Lines changed: 46 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,23 @@
11
import type { Mock } from 'vitest';
22
import { afterEach, describe, expect, it, vi } from 'vitest';
33

4+
// eslint-disable-next-line import/first
5+
import { getActiveTraceId, injectSpanTraceHeaders } from './traceparent';
6+
47
vi.mock('@lobechat/observability-otel/api', () => {
58
const inject = vi.fn();
69
const setSpan = vi.fn((_ctx, span) => span);
10+
const getActiveSpan = vi.fn();
711

812
return {
913
context: {
1014
active: vi.fn(() => ({})),
1115
},
1216
propagation: { inject },
13-
trace: { setSpan },
17+
trace: { getActiveSpan, setSpan },
1418
};
1519
});
1620

17-
// eslint-disable-next-line import/first
18-
import { injectSpanTraceHeaders } from './traceparent';
19-
2021
const mockSpan = (traceId: string, spanId: string) =>
2122
({
2223
spanContext: () => ({
@@ -39,7 +40,9 @@ describe('injectSpanTraceHeaders', () => {
3940

4041
it('uses propagator output when available', async () => {
4142
const { propagation } = await api;
42-
(propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>).mockImplementation((_ctx, carrier) => {
43+
(
44+
propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>
45+
).mockImplementation((_ctx, carrier) => {
4346
carrier.traceparent = 'from-propagator';
4447
carrier.tracestate = 'state';
4548
});
@@ -56,14 +59,50 @@ describe('injectSpanTraceHeaders', () => {
5659

5760
it('falls back to manual traceparent formatting when propagator gives none', async () => {
5861
const { propagation } = await api;
59-
(propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>).mockImplementation(() => undefined);
62+
(
63+
propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>
64+
).mockImplementation(() => undefined);
6065

6166
const headers = headersWith();
6267
const span = mockSpan('1'.repeat(32), '2'.repeat(16));
6368

6469
const tp = injectSpanTraceHeaders(headers, span);
6570

6671
expect(tp).toBe('00-11111111111111111111111111111111-2222222222222222-01');
67-
expect(headers.get('traceparent')).toBe('00-11111111111111111111111111111111-2222222222222222-01');
72+
expect(headers.get('traceparent')).toBe(
73+
'00-11111111111111111111111111111111-2222222222222222-01',
74+
);
75+
});
76+
});
77+
78+
describe('getActiveTraceId', () => {
79+
const api = vi.importMock<typeof import('@lobechat/observability-otel/api')>(
80+
'@lobechat/observability-otel/api',
81+
);
82+
83+
afterEach(() => {
84+
vi.resetAllMocks();
85+
});
86+
87+
it('returns traceId from active span', async () => {
88+
const { trace } = await api;
89+
const expectedTraceId = 'a'.repeat(32);
90+
(trace.getActiveSpan as Mock).mockReturnValue(mockSpan(expectedTraceId, 'b'.repeat(16)));
91+
92+
expect(getActiveTraceId()).toBe(expectedTraceId);
93+
});
94+
95+
it('returns undefined when no active span', async () => {
96+
const { trace } = await api;
97+
(trace.getActiveSpan as Mock).mockReturnValue(undefined);
98+
99+
expect(getActiveTraceId()).toBeUndefined();
100+
});
101+
102+
it('returns undefined when traceId is all zeros', async () => {
103+
const { trace } = await api;
104+
(trace.getActiveSpan as Mock).mockReturnValue(mockSpan('0'.repeat(32), 'b'.repeat(16)));
105+
106+
expect(getActiveTraceId()).toBeUndefined();
68107
});
69108
});

‎src/libs/observability/traceparent.ts‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,5 @@
1-
import type {
2-
Span,
3-
Context as OtContext,
4-
TextMapGetter
5-
} from '@lobechat/observability-otel/api';
6-
import {
7-
context as otContext,
8-
propagation,
9-
trace,
10-
} from '@lobechat/observability-otel/api';
1+
import type { Context as OtContext, Span, TextMapGetter } from '@lobechat/observability-otel/api';
2+
import { context as otContext, propagation, trace } from '@lobechat/observability-otel/api';
113

124
// NOTICE: do not try to optimize this into .repeat(...) or similar,
135
// here served for better search / semantic search purpose for further diagnostic
@@ -47,6 +39,16 @@ export const getActiveTraceparent = () => {
4739
return toTraceparent(span as Span);
4840
};
4941

42+
/**
43+
* Get the traceId from the active span.
44+
*/
45+
export const getActiveTraceId = () => {
46+
const span = trace.getActiveSpan();
47+
if (!isValidContext(span)) return undefined;
48+
49+
return span!.spanContext().traceId;
50+
};
51+
5052
/**
5153
* Injects the active context into headers using the configured propagator (W3C by default).
5254
* Also returns the traceparent for convenience.

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,19 +14,22 @@ var aiInfraMocks:
1414
| {
1515
getAiProviderRuntimeState: ReturnType<typeof vi.fn>;
1616
tryMatchingModelFrom: ReturnType<typeof vi.fn>;
17+
tryMatchingProviderFrom: ReturnType<typeof vi.fn>;
1718
};
1819

1920
vi.mock('@/database/repositories/aiInfra', () => {
2021
aiInfraMocks = {
2122
getAiProviderRuntimeState: vi.fn(),
2223
tryMatchingModelFrom: vi.fn(),
24+
tryMatchingProviderFrom: vi.fn(),
2325
};
2426

2527
const AiInfraRepos = vi.fn().mockImplementation(() => ({
2628
getAiProviderRuntimeState: aiInfraMocks!.getAiProviderRuntimeState,
2729
})) as unknown as typeof import('@/database/repositories/aiInfra').AiInfraRepos;
2830

2931
(AiInfraRepos as any).tryMatchingModelFrom = aiInfraMocks!.tryMatchingModelFrom;
32+
(AiInfraRepos as any).tryMatchingProviderFrom = aiInfraMocks!.tryMatchingProviderFrom;
3033

3134
return { AiInfraRepos };
3235
});
@@ -85,7 +88,9 @@ beforeEach(async () => {
8588
toolCall.mockClear();
8689
aiInfraMocks!.getAiProviderRuntimeState.mockReset();
8790
aiInfraMocks!.tryMatchingModelFrom.mockReset();
91+
aiInfraMocks!.tryMatchingProviderFrom.mockReset();
8892
aiInfraMocks!.tryMatchingModelFrom.mockResolvedValue('openai');
93+
aiInfraMocks!.tryMatchingProviderFrom.mockResolvedValue('openai');
8994
aiInfraMocks!.getAiProviderRuntimeState.mockResolvedValue({
9095
enabledAiModels: [
9196
{ abilities: {}, enabled: true, id: 'gpt-mock', providerId: 'openai', type: 'chat' },

0 commit comments

Comments
 (0)