Skip to content

feat(telemetry): unify span creation paths for hierarchical trace tree #4126

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
doudouOUC wants to merge 13 commits into
base: main
Choose a base branch
from
Open

feat(telemetry): unify span creation paths for hierarchical trace tree #4126

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
b7eb7cf
feat(telemetry): unify span creation paths for hierarchical trace tre…
doudouOUC May 13, 2026
dbe81f9
fix(telemetry): address CI build errors in session-tracing tests
doudouOUC May 13, 2026
e561f9e
fix(telemetry): update coreToolScheduler and loggingContentGenerator …
doudouOUC May 14, 2026
b58de59
fix(telemetry): address review feedback from wenshao
doudouOUC May 14, 2026
a586655
fix(telemetry): use partial mock for telemetry/index.js in coreToolSc…
doudouOUC May 14, 2026
c14c553
fix(telemetry): getLastToolSpan must skip tool.execution sub-spans
doudouOUC May 14, 2026
7273fef
fix(telemetry): address second round review feedback
doudouOUC May 14, 2026
2e1a832
fix(telemetry): endToolSpan mock writes to record directly
doudouOUC May 14, 2026
26a43dd
fix(telemetry): mock session-tracing.js directly instead of telemetry…
doudouOUC May 14, 2026
eef8ce7
fix(telemetry): fix endToolSpan status on success — toolCalls is empt…
doudouOUC May 14, 2026
f62d91f
fix(telemetry): address Codex review — activate OTel context, end spa…
doudouOUC May 15, 2026
154be37
fix(telemetry): address Codex review v2 — session-root fallback + exe…
doudouOUC May 15, 2026
daa06b3
fix(telemetry): address Codex review v3 — sync throw, idle timeout ra…
doudouOUC May 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
376 changes: 376 additions & 0 deletions docs/design/workflow-tracing-gaps.md
View file
Open in desktop

Large diffs are not rendered by default.

81 changes: 2 additions & 79 deletions packages/core/src/core/client.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@ import {
} from 'vitest';

import type { Content, GenerateContentResponse, Part } from '@google/genai';
import { SpanStatusCode } from '@opentelemetry/api';
import { GeminiClient, SendMessageType } from './client.js';
import { findCompressSplitPoint } from '../services/chatCompressionService.js';
import {
Expand Down Expand Up @@ -167,29 +166,9 @@ const mockUiTelemetryService = vi.hoisted(() => ({
reset: vi.fn(),
addEvent: vi.fn(),
}));
const clientSpanCalls = vi.hoisted(
(): Array<{
name: string;
attributes: Record<string, string | number | boolean>;
statuses: Array<{ code: number; message?: string }>;
}> => [],
);
const mockWithSpan = vi.hoisted(() => vi.fn());

vi.mock('../telemetry/tracer.js', () => ({
API_CALL_ABORTED_SPAN_STATUS_MESSAGE: 'API call aborted',
API_CALL_FAILED_SPAN_STATUS_MESSAGE: 'API call failed',
safeSetStatus: (
span: { setStatus: (status: { code: number; message?: string }) => void },
status: { code: number; message?: string },
) => {
try {
span.setStatus(status);
} catch {
// Match production best-effort telemetry behavior.
}
},
withSpan: mockWithSpan,
}));

vi.mock('../telemetry/index.js', async (importOriginal) => {
Expand Down Expand Up @@ -356,32 +335,6 @@ describe('Gemini Client (client.ts)', () => {
};
beforeEach(async () => {
vi.resetAllMocks();
clientSpanCalls.length = 0;
mockWithSpan.mockImplementation(
async (
name: string,
attributes: Record<string, string | number | boolean>,
fn: (span: {
setStatus: ReturnType<typeof vi.fn>;
setAttribute: ReturnType<typeof vi.fn>;
end: ReturnType<typeof vi.fn>;
}) => Promise<unknown>,
) => {
const spanCall = {
name,
attributes,
statuses: [] as Array<{ code: number; message?: string }>,
};
clientSpanCalls.push(spanCall);
return fn({
setStatus: vi.fn((status: { code: number; message?: string }) => {
spanCall.statuses.push(status);
}),
setAttribute: vi.fn(),
end: vi.fn(),
});
},
);
vi.mocked(uiTelemetryService.setLastPromptTokenCount).mockClear();

// Default: createContentGenerator rejects (simulates test env without auth).
Expand Down Expand Up @@ -4494,15 +4447,6 @@ Other open files:
}),
'btw-prompt-id',
);
expect(clientSpanCalls.at(-1)).toEqual(
expect.objectContaining({
name: 'client.generateContent',
attributes: {
model: DEFAULT_QWEN_FLASH_MODEL,
prompt_id: 'btw-prompt-id',
},
}),
);
});

it('should prefer an explicit prompt id override over the current context', async () => {
Expand Down Expand Up @@ -4530,15 +4474,6 @@ Other open files:
}),
'override-prompt-id',
);
expect(clientSpanCalls.at(-1)).toEqual(
expect.objectContaining({
name: 'client.generateContent',
attributes: {
model: DEFAULT_QWEN_FLASH_MODEL,
prompt_id: 'override-prompt-id',
},
}),
);
});

it('should use config system prompt override when provided', async () => {
Expand Down Expand Up @@ -4642,7 +4577,7 @@ Other open files:
);
});

it('sets a generic span status when content generation fails', async () => {
it('propagates error when content generation fails', async () => {
const contents = [{ role: 'user', parts: [{ text: 'hello' }] }];
const abortSignal = new AbortController().signal;
mockGenerateContentFn.mockRejectedValueOnce(
Expand All @@ -4657,15 +4592,9 @@ Other open files:
DEFAULT_QWEN_FLASH_MODEL,
),
).rejects.toThrow('raw upstream 500 with sensitive details');

const spanCall = clientSpanCalls.at(-1);
expect(spanCall?.statuses).toEqual([
{ code: SpanStatusCode.ERROR, message: 'API call failed' },
]);
expect(JSON.stringify(spanCall?.statuses)).not.toContain('raw upstream');
});

it('sets a generic aborted span status when content generation is aborted', async () => {
it('propagates error when content generation is aborted', async () => {
const contents = [{ role: 'user', parts: [{ text: 'hello' }] }];
const abortController = new AbortController();
abortController.abort();
Expand All @@ -4681,12 +4610,6 @@ Other open files:
DEFAULT_QWEN_FLASH_MODEL,
),
).rejects.toThrow('raw abort reason with sensitive details');

const spanCall = clientSpanCalls.at(-1);
expect(spanCall?.statuses).toEqual([
{ code: SpanStatusCode.ERROR, message: 'API call aborted' },
]);
expect(JSON.stringify(spanCall?.statuses)).not.toContain('raw abort');
});

// Note: there is currently no "fallback mode" model routing; the model used
Expand Down
151 changes: 63 additions & 88 deletions packages/core/src/core/client.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@ import type {
PartListUnion,
Tool,
} from '@google/genai';
import { SpanStatusCode } from '@opentelemetry/api';

// Config
import { ApprovalMode, type Config } from '../config/config.js';
Expand Down Expand Up @@ -105,12 +104,6 @@ import { createHookOutput, SessionStartSource } from '../hooks/types.js';
import { ideContextStore } from '../ide/ideContext.js';
import { type File, type IdeContext } from '../ide/types.js';
import { PermissionMode, type StopHookOutput } from '../hooks/types.js';
import {
API_CALL_ABORTED_SPAN_STATUS_MESSAGE,
API_CALL_FAILED_SPAN_STATUS_MESSAGE,
safeSetStatus,
withSpan,
} from '../telemetry/tracer.js';

const MAX_TURNS = 100;

Expand Down Expand Up @@ -1684,91 +1677,73 @@ export class GeminiClient {
const promptId =
promptIdOverride ?? promptIdContext.getStore() ?? this.lastPromptId!;

return withSpan(
'client.generateContent',
{ model, prompt_id: promptId },
async (span) => {
let currentAttemptModel: string = model;
let currentAttemptModel: string = model;
Comment thread
doudouOUC marked this conversation as resolved.

try {
const userMemory = this.config.getUserMemory();
const finalSystemInstruction = generationConfig.systemInstruction
? getCustomSystemPrompt(
generationConfig.systemInstruction,
userMemory,
)
: this.getMainSessionSystemInstruction();

const requestConfig: GenerateContentConfig = {
abortSignal,
...generationConfig,
systemInstruction: finalSystemInstruction,
};
try {
const userMemory = this.config.getUserMemory();
const finalSystemInstruction = generationConfig.systemInstruction
? getCustomSystemPrompt(generationConfig.systemInstruction, userMemory)
: this.getMainSessionSystemInstruction();

const requestConfig: GenerateContentConfig = {
abortSignal,
...generationConfig,
systemInstruction: finalSystemInstruction,
};

// When the requested model differs from the main model (e.g. fast model
// side queries for session recap / title / summary), resolve the target
// model's own ContentGeneratorConfig so that per-model settings like
// extra_body, samplingParams, and reasoning are not inherited from the
// main model's config. The retry authType is resolved alongside so that
// provider-specific checks (e.g. QWEN_OAUTH quota detection) reference
// the target model's provider.
const {
contentGenerator,
retryAuthType,
// When the requested model differs from the main model (e.g. fast model
// side queries for session recap / title / summary), resolve the target
// model's own ContentGeneratorConfig so that per-model settings like
// extra_body, samplingParams, and reasoning are not inherited from the
// main model's config. The retry authType is resolved alongside so that
// provider-specific checks (e.g. QWEN_OAUTH quota detection) reference
// the target model's provider.
const {
contentGenerator,
retryAuthType,
model: requestModel,
} = await this.config.getBaseLlmClient().resolveForModel(model);

const apiCall = () => {
currentAttemptModel = requestModel;

return contentGenerator.generateContent(
{
model: requestModel,
} = await this.config.getBaseLlmClient().resolveForModel(model);

const apiCall = () => {
currentAttemptModel = requestModel;

return contentGenerator.generateContent(
{
model: requestModel,
config: requestConfig,
contents,
},
promptId,
);
};
const result = await retryWithBackoff(apiCall, {
authType: retryAuthType,
persistentMode: isUnattendedMode(),
signal: abortSignal,
heartbeatFn: (info) => {
process.stderr.write(
`[qwen-code] Waiting for API capacity... attempt ${info.attempt}, retry in ${Math.ceil(info.remainingMs / 1000)}s\n`,
);
},
});
return result;
} catch (error: unknown) {
if (abortSignal.aborted) {
safeSetStatus(span, {
code: SpanStatusCode.ERROR,
message: API_CALL_ABORTED_SPAN_STATUS_MESSAGE,
});
throw error;
}

safeSetStatus(span, {
code: SpanStatusCode.ERROR,
message: API_CALL_FAILED_SPAN_STATUS_MESSAGE,
});
await reportError(
error,
`Error generating content via API with model ${currentAttemptModel}.`,
{
requestContents: contents,
requestConfig: generationConfig,
},
'generateContent-api',
);
throw new Error(
`Failed to generate content with model ${currentAttemptModel}: ${getErrorMessage(error)}`,
config: requestConfig,
contents,
},
promptId,
);
};
const result = await retryWithBackoff(apiCall, {
authType: retryAuthType,
persistentMode: isUnattendedMode(),
signal: abortSignal,
heartbeatFn: (info) => {
process.stderr.write(
`[qwen-code] Waiting for API capacity... attempt ${info.attempt}, retry in ${Math.ceil(info.remainingMs / 1000)}s\n`,
);
}
},
);
},
});
return result;
} catch (error: unknown) {
if (abortSignal.aborted) {
throw error;
}
await reportError(
error,
`Error generating content via API with model ${currentAttemptModel}.`,
{
requestContents: contents,
requestConfig: generationConfig,
},
'generateContent-api',
);
throw new Error(
`Failed to generate content with model ${currentAttemptModel}: ${getErrorMessage(error)}`,
);
}
}

/**
Expand Down
Loading
Loading