fix(core): dynamic session ID injection to resolve resume bugs (#24972)

This commit is contained in:
Tommaso Sciortino
2026-04-08 23:27:24 +00:00
committed by GitHub
parent 80764c8bb5
commit d06dba3538
33 changed files with 165 additions and 189 deletions
+12 -9
View File
@@ -110,7 +110,7 @@ describe('runInDevTraceSpan', () => {
const fn = vi.fn(async () => 'result');
const result = await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
{ operation: GeminiCliOperation.LLMCall, sessionId: 'test-session-id' },
fn,
);
@@ -125,7 +125,7 @@ describe('runInDevTraceSpan', () => {
it('should set default attributes on the span metadata', async () => {
await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
{ operation: GeminiCliOperation.LLMCall, sessionId: 'test-session-id' },
async ({ metadata }) => {
expect(metadata.attributes[GEN_AI_OPERATION_NAME]).toBe(
GeminiCliOperation.LLMCall,
@@ -143,7 +143,7 @@ describe('runInDevTraceSpan', () => {
it('should set span attributes from metadata on completion', async () => {
await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
{ operation: GeminiCliOperation.LLMCall, sessionId: 'test-session-id' },
async ({ metadata }) => {
metadata.input = { query: 'hello' };
metadata.output = { response: 'world' };
@@ -169,9 +169,12 @@ describe('runInDevTraceSpan', () => {
it('should handle errors in the wrapped function', async () => {
const error = new Error('test error');
await expect(
runInDevTraceSpan({ operation: GeminiCliOperation.LLMCall }, async () => {
throw error;
}),
runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall, sessionId: 'test-session-id' },
async () => {
throw error;
},
),
).rejects.toThrow(error);
expect(mockSpan.setStatus).toHaveBeenCalledWith({
@@ -189,7 +192,7 @@ describe('runInDevTraceSpan', () => {
}
const resultStream = await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
{ operation: GeminiCliOperation.LLMCall, sessionId: 'test-session-id' },
async () => testStream(),
);
@@ -212,7 +215,7 @@ describe('runInDevTraceSpan', () => {
}
const resultStream = await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
{ operation: GeminiCliOperation.LLMCall, sessionId: 'test-session-id' },
async () => errorStream(),
);
@@ -231,7 +234,7 @@ describe('runInDevTraceSpan', () => {
});
await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
{ operation: GeminiCliOperation.LLMCall, sessionId: 'test-session-id' },
async ({ metadata }) => {
metadata.input = 'trigger error';
},
+6 -3
View File
@@ -23,7 +23,6 @@ import {
SERVICE_DESCRIPTION,
SERVICE_NAME,
} from './constants.js';
import { sessionId } from '../utils/session.js';
import { truncateString } from '../utils/textUtils.js';
@@ -96,10 +95,14 @@ export interface SpanMetadata {
* @returns The result of the function.
*/
export async function runInDevTraceSpan<R>(
opts: SpanOptions & { operation: GeminiCliOperation; logPrompts?: boolean },
opts: SpanOptions & {
operation: GeminiCliOperation;
logPrompts?: boolean;
sessionId: string;
},
fn: ({ metadata }: { metadata: SpanMetadata }) => Promise<R>,
): Promise<R> {
const { operation, logPrompts, ...restOfSpanOpts } = opts;
const { operation, logPrompts, sessionId, ...restOfSpanOpts } = opts;
const tracer = trace.getTracer(TRACER_NAME, TRACER_VERSION);
return tracer.startActiveSpan(operation, restOfSpanOpts, async (span) => {