|
|
@@ -4,6 +4,8 @@ import { Session } from "../../src/session"
|
|
|
import { Bus } from "../../src/bus"
|
|
|
import { Log } from "../../src/util/log"
|
|
|
import { Instance } from "../../src/project/instance"
|
|
|
+import { MessageV2 } from "../../src/session/message-v2"
|
|
|
+import { Identifier } from "../../src/id/id"
|
|
|
|
|
|
const projectRoot = path.join(__dirname, "../..")
|
|
|
Log.init({ print: false })
|
|
|
@@ -69,3 +71,68 @@ describe("session.started event", () => {
|
|
|
})
|
|
|
})
|
|
|
})
|
|
|
+
|
|
|
+describe("step-finish token propagation via Bus event", () => {
|
|
|
+ test("non-zero tokens propagate through PartUpdated event", async () => {
|
|
|
+ await Instance.provide({
|
|
|
+ directory: projectRoot,
|
|
|
+ fn: async () => {
|
|
|
+ const session = await Session.create({})
|
|
|
+
|
|
|
+ const messageID = Identifier.ascending("message")
|
|
|
+ await Session.updateMessage({
|
|
|
+ id: messageID,
|
|
|
+ sessionID: session.id,
|
|
|
+ role: "user",
|
|
|
+ time: { created: Date.now() },
|
|
|
+ agent: "user",
|
|
|
+ model: { providerID: "test", modelID: "test" },
|
|
|
+ tools: {},
|
|
|
+ mode: "",
|
|
|
+ } as unknown as MessageV2.Info)
|
|
|
+
|
|
|
+ let received: MessageV2.Part | undefined
|
|
|
+ const unsub = Bus.subscribe(MessageV2.Event.PartUpdated, (event) => {
|
|
|
+ received = event.properties.part
|
|
|
+ })
|
|
|
+
|
|
|
+ const tokens = {
|
|
|
+ total: 1500,
|
|
|
+ input: 500,
|
|
|
+ output: 800,
|
|
|
+ reasoning: 200,
|
|
|
+ cache: { read: 100, write: 50 },
|
|
|
+ }
|
|
|
+
|
|
|
+ const partInput = {
|
|
|
+ id: Identifier.ascending("part"),
|
|
|
+ messageID,
|
|
|
+ sessionID: session.id,
|
|
|
+ type: "step-finish" as const,
|
|
|
+ reason: "stop",
|
|
|
+ cost: 0.005,
|
|
|
+ tokens,
|
|
|
+ }
|
|
|
+
|
|
|
+ await Session.updatePart(partInput)
|
|
|
+
|
|
|
+ await new Promise((resolve) => setTimeout(resolve, 100))
|
|
|
+
|
|
|
+ expect(received).toBeDefined()
|
|
|
+ expect(received!.type).toBe("step-finish")
|
|
|
+ const finish = received as MessageV2.StepFinishPart
|
|
|
+ expect(finish.tokens.input).toBe(500)
|
|
|
+ expect(finish.tokens.output).toBe(800)
|
|
|
+ expect(finish.tokens.reasoning).toBe(200)
|
|
|
+ expect(finish.tokens.total).toBe(1500)
|
|
|
+ expect(finish.tokens.cache.read).toBe(100)
|
|
|
+ expect(finish.tokens.cache.write).toBe(50)
|
|
|
+ expect(finish.cost).toBe(0.005)
|
|
|
+ expect(received).not.toBe(partInput)
|
|
|
+
|
|
|
+ unsub()
|
|
|
+ await Session.remove(session.id)
|
|
|
+ },
|
|
|
+ })
|
|
|
+ }, { timeout: 30000 })
|
|
|
+})
|