fix(FN-2583): persist token usage on step completion
- Persist task token usage asynchronously whenever a step completes in StepSession execution - Log non-fatal warnings if per-step token usage persistence fails - Add executor integration coverage verifying incremental tokenUsage updates before moving tasks to in-review
This commit is contained in:
@@ -10875,6 +10875,59 @@ describe("StepSessionExecutor integration", () => {
|
|||||||
expect(Number.isNaN(new Date(tokenUsage.lastUsedAt as string).getTime())).toBe(false);
|
expect(Number.isNaN(new Date(tokenUsage.lastUsedAt as string).getTime())).toBe(false);
|
||||||
});
|
});
|
||||||
|
|
||||||
|
it("persists tokenUsage incrementally during step execution before in-review transition", async () => {
|
||||||
|
const { store } = createTokenUsageStepSessionStore();
|
||||||
|
const totals = { input: 100, output: 40 };
|
||||||
|
|
||||||
|
mockedStepSessionExecutor.mockImplementationOnce(((options: any) => ({
|
||||||
|
executeAll: vi.fn(async () => {
|
||||||
|
totals.input = 120;
|
||||||
|
totals.output = 50;
|
||||||
|
options.onStepComplete(0, { stepIndex: 0, success: true, retries: 0 });
|
||||||
|
await new Promise((resolve) => setTimeout(resolve, 0));
|
||||||
|
|
||||||
|
totals.input = 150;
|
||||||
|
totals.output = 70;
|
||||||
|
options.onStepComplete(1, { stepIndex: 1, success: true, retries: 0 });
|
||||||
|
await new Promise((resolve) => setTimeout(resolve, 0));
|
||||||
|
|
||||||
|
return [
|
||||||
|
{ stepIndex: 0, success: true, retries: 0 },
|
||||||
|
{ stepIndex: 1, success: true, retries: 0 },
|
||||||
|
];
|
||||||
|
}),
|
||||||
|
terminateAllSessions: mockTerminateAllSessions,
|
||||||
|
cleanup: mockCleanup,
|
||||||
|
})) as any);
|
||||||
|
|
||||||
|
const agentStore = {
|
||||||
|
getAgent: vi.fn().mockImplementation(async () => ({
|
||||||
|
id: "agent-001",
|
||||||
|
totalInputTokens: totals.input,
|
||||||
|
totalOutputTokens: totals.output,
|
||||||
|
})),
|
||||||
|
};
|
||||||
|
|
||||||
|
const executor = new TaskExecutor(store, "/tmp/test", { agentStore: agentStore as any });
|
||||||
|
await executor.execute(createTaskWithSteps());
|
||||||
|
|
||||||
|
const tokenUsageUpdates = store.updateTask.mock.calls
|
||||||
|
.filter(([, updates]: [string, Record<string, unknown>]) => updates.tokenUsage)
|
||||||
|
.map(([, updates]: [string, Record<string, unknown>]) => updates.tokenUsage as Record<string, unknown>);
|
||||||
|
|
||||||
|
expect(tokenUsageUpdates.length).toBeGreaterThanOrEqual(2);
|
||||||
|
expect(tokenUsageUpdates).toEqual(
|
||||||
|
expect.arrayContaining([
|
||||||
|
expect.objectContaining({
|
||||||
|
inputTokens: 20,
|
||||||
|
outputTokens: 10,
|
||||||
|
totalTokens: 30,
|
||||||
|
}),
|
||||||
|
]),
|
||||||
|
);
|
||||||
|
expect(store.moveTask).toHaveBeenCalledWith("FN-200", "in-review");
|
||||||
|
});
|
||||||
|
|
||||||
it("persists task tokenUsage on failure paths so partial usage is visible", async () => {
|
it("persists task tokenUsage on failure paths so partial usage is visible", async () => {
|
||||||
const { store } = createTokenUsageStepSessionStore();
|
const { store } = createTokenUsageStepSessionStore();
|
||||||
const totals = { input: 30, output: 10 };
|
const totals = { input: 30, output: 10 };
|
||||||
|
|||||||
@@ -1592,6 +1592,9 @@ export class TaskExecutor {
|
|||||||
} catch (err) {
|
} catch (err) {
|
||||||
executorLog.warn(`${task.id}: failed to update step ${stepIndex} status: ${err}`);
|
executorLog.warn(`${task.id}: failed to update step ${stepIndex} status: ${err}`);
|
||||||
}
|
}
|
||||||
|
this.persistTokenUsage(task.id).catch((err) => {
|
||||||
|
executorLog.warn(`${task.id}: failed to persist token usage on step ${stepIndex} complete: ${err}`);
|
||||||
|
});
|
||||||
},
|
},
|
||||||
});
|
});
|
||||||
this.activeStepExecutors.set(task.id, stepExecutor);
|
this.activeStepExecutors.set(task.id, stepExecutor);
|
||||||
|
|||||||
Reference in New Issue
Block a user