diff --git a/packages/server-utils/src/ai/pi-durable/index.ts b/packages/server-utils/src/ai/pi-durable/index.ts index b7ee007e6a02..d2187aae394e 100644 --- a/packages/server-utils/src/ai/pi-durable/index.ts +++ b/packages/server-utils/src/ai/pi-durable/index.ts @@ -220,9 +220,9 @@ function runPhase( const runInScope = async (scope: Scope): Promise => { // Read at the first phase of a run, which is a `pi.tool` phase when a restarted Harness reruns a - // tool call. + // tool call, and at each phase outside a run, such as an idle compaction or a task of an extension. let sessionId = run?.sessionId; - if (!sessionId && (run || kind === PI_TASK.COMPACTION)) { + if (!sessionId) { sessionId = await readProviderSessionId(runtime, context); if (run && sessionId) { run.sessionId = sessionId; diff --git a/packages/server-utils/test/ai/lib/tracing/pi-durable.test.ts b/packages/server-utils/test/ai/lib/tracing/pi-durable.test.ts index 5c35581138cf..e7bd701485cf 100644 --- a/packages/server-utils/test/ai/lib/tracing/pi-durable.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/pi-durable.test.ts @@ -1049,6 +1049,35 @@ describe('instrumentPiDurableHarnessOptions', () => { ); }); + it('uses the provider session id for a task of an extension on an idle conversation', async () => { + const sessionId = '019a2b3c-4d5e-7f60-8172-93a4b5c6d7e8'; + const task: PiTask = { + definition: { + name: 'app.title', + phases: { + request: async () => { + await harness.models!.completeSimple!(MODEL, { messages: [] }); + }, + }, + }, + }; + const harness: PiHarnessOptions = instrumentPiDurableHarnessOptions({ + models: { completeSimple: async () => ANSWER }, + registry: { snapshot: () => ({ task: () => task, tasks: () => [task] }), subscribe: () => () => undefined }, + }); + const runtime: PiTaskRuntime = { + taskId: 7, + conversationId: 1, + agent: async () => ({}), + commit: async () => undefined, + snapshot: async () => ({ sessionId }), + }; + + await harness.registry!.snapshot().task('app.title')!.definition.phases.request!(undefined, runtime, undefined); + + expect(spanToJSON(endedSpans[0]!).attributes[GEN_AI_CONVERSATION_ID]).toBe(sessionId); + }); + it('ends the open runs when the Harness closes', async () => { const harness = harnessFor(generationPhase([{ run: { inputs: [5] } }])); let onClose: (() => void) | undefined;