From 2b62319ca0af03725ad8d600988d1609ca914868 Mon Sep 17 00:00:00 2001 From: Roomote Date: Sat, 5 Sep 2026 03:54:43 +0000 Subject: [PATCH] feat(worker): attribute inference usage to workflows --- .../runtime-envelope-subscription.test.ts | 2 + .../run-task/subscribe-harness-callbacks.ts | 3 + apps/worker/src/sandbox-server/lib/harness.ts | 2 + .../__tests__/opencode-server.test.ts | 70 +++++++++++++++++++ .../lib/harnesses/opencode-server/harness.ts | 11 ++- 5 files changed, 87 insertions(+), 1 deletion(-) diff --git a/apps/worker/src/run-task/__tests__/runtime-envelope-subscription.test.ts b/apps/worker/src/run-task/__tests__/runtime-envelope-subscription.test.ts index 01dc3c6798..7e1bc1f084 100644 --- a/apps/worker/src/run-task/__tests__/runtime-envelope-subscription.test.ts +++ b/apps/worker/src/run-task/__tests__/runtime-envelope-subscription.test.ts @@ -313,6 +313,7 @@ describe('subscribeHarnessCallbacks', () => { providerId: 'openrouter', modelId: 'openai/gpt-5.4', agent: 'explore', + workflowSkill: 'implement-changes', inputTokens: 10, outputTokens: 5, reasoningTokens: 2, @@ -344,6 +345,7 @@ describe('subscribeHarnessCallbacks', () => { costSource: 'opencode_message', messageCreatedAt: new Date('2026-07-01T12:00:00.000Z'), messageCompletedAt: new Date('2026-07-01T12:00:01.000Z'), + details: { workflowSkill: 'implement-changes' }, }); const unsubscribePromise = unsubscribe(); diff --git a/apps/worker/src/run-task/subscribe-harness-callbacks.ts b/apps/worker/src/run-task/subscribe-harness-callbacks.ts index 3f29c0f555..ef25e5ff32 100644 --- a/apps/worker/src/run-task/subscribe-harness-callbacks.ts +++ b/apps/worker/src/run-task/subscribe-harness-callbacks.ts @@ -206,6 +206,9 @@ export function subscribeHarnessCallbacks({ costSource: event.costSource, messageCreatedAt: event.messageCreatedAt ?? null, messageCompletedAt: event.messageCompletedAt ?? null, + details: event.workflowSkill + ? { workflowSkill: event.workflowSkill } + : null, }) .then(() => { consecutivePersistenceFailures = 0; diff --git a/apps/worker/src/sandbox-server/lib/harness.ts b/apps/worker/src/sandbox-server/lib/harness.ts index 51065ff763..4c6428eb93 100644 --- a/apps/worker/src/sandbox-server/lib/harness.ts +++ b/apps/worker/src/sandbox-server/lib/harness.ts @@ -267,6 +267,8 @@ export interface HarnessInferenceUsageEvent { * `build`) or a subagent name (for example `explore` or `visual`). */ agent?: string; + /** Packaged workflow active for a primary-session inference request. */ + workflowSkill?: string; inputTokens: number; outputTokens: number; reasoningTokens: number; diff --git a/apps/worker/src/sandbox-server/lib/harnesses/__tests__/opencode-server.test.ts b/apps/worker/src/sandbox-server/lib/harnesses/__tests__/opencode-server.test.ts index a6ee34b5e4..a3e2ccbc43 100644 --- a/apps/worker/src/sandbox-server/lib/harnesses/__tests__/opencode-server.test.ts +++ b/apps/worker/src/sandbox-server/lib/harnesses/__tests__/opencode-server.test.ts @@ -1266,6 +1266,76 @@ describe('OpenCodeServerHarness', () => { } }); + it('attributes primary-session inference usage to the active workflow skill', async () => { + const { client, harness } = createHarness(new FakeOpenCodeServerClient()); + const inferenceUsageEvents: HarnessInferenceUsageEvent[] = []; + + harness.subscribeRuntimeInferenceUsage((event) => + inferenceUsageEvents.push(event), + ); + + try { + await connectHarness(harness, client); + + expect( + harness.sendCommand({ + commandName: TaskCommandName.StartNewTask, + data: { + text: 'Implement the change.', + visibleInTranscript: true, + source: 'web', + }, + }), + ).toBe(true); + await vi.waitFor(() => { + expect(client.promptAsync).toHaveBeenCalledTimes(1); + }); + + await client.emit({ + type: 'message.part.updated', + properties: { + part: { + id: 'skill_part_usage', + sessionID: 'ses_1', + messageID: 'msg_skill_usage', + type: 'tool', + callID: 'skill_call_usage', + tool: 'skill', + state: { + status: 'completed', + input: { name: 'implement-changes' }, + title: 'Load skill', + }, + }, + }, + }); + + client.message.mockResolvedValueOnce( + createFinalAssistantMessage({ messageId: 'msg_skill_usage' }), + ); + await client.emit({ + type: 'message.updated', + properties: { + info: { + id: 'msg_skill_usage', + sessionID: 'ses_1', + role: 'assistant', + time: { completed: 1 }, + }, + }, + }); + + expect(inferenceUsageEvents).toHaveLength(1); + expect(inferenceUsageEvents[0]).toMatchObject({ + sessionId: 'ses_1', + messageId: 'msg_skill_usage', + workflowSkill: 'implement-changes', + }); + } finally { + harness.dispose(); + } + }); + it('auto-submits one hidden continuation on the build agent after an in-flight plan turn loads implement-changes', async () => { const { client, harness } = createHarness(new FakeOpenCodeServerClient()); const runtimeOutputEvents: AcpMessage[] = []; diff --git a/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/harness.ts b/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/harness.ts index e854a79ddc..06faaabcc4 100644 --- a/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/harness.ts +++ b/apps/worker/src/sandbox-server/lib/harnesses/opencode-server/harness.ts @@ -753,6 +753,7 @@ function createInferenceUsageEvent( info: OpenCodeMessageInfo, tokenUsage: Record, fallbackAgent?: string, + workflowSkill?: string, ): HarnessInferenceUsageEvent { const messageCreatedAt = openCodeTimestampToDate(info.time?.created); const messageCompletedAt = openCodeTimestampToDate(info.time?.completed); @@ -766,6 +767,7 @@ function createInferenceUsageEvent( : {}), ...(typeof info.modelID === 'string' ? { modelId: info.modelID } : {}), ...(agent ? { agent } : {}), + ...(workflowSkill ? { workflowSkill } : {}), inputTokens: Number(tokenUsage.inputTokens ?? 0), outputTokens: Number(tokenUsage.outputTokens ?? 0), reasoningTokens: Number(tokenUsage.reasoningTokens ?? 0), @@ -5445,7 +5447,14 @@ export class OpenCodeServerHarness }); this.emit( 'runtimeInferenceUsage', - createInferenceUsageEvent(message.info, tokenUsage, options?.agentType), + createInferenceUsageEvent( + message.info, + tokenUsage, + options?.agentType, + message.info.sessionID === this.sessionId + ? (this.activeWorkflowSkill ?? undefined) + : undefined, + ), ); if (options?.finalizeParentTurn !== false) { this.finalizedAssistantTurn = finalized;