@@ -4,6 +4,8 @@ import { Session } from "../../src/session"
44import { Bus } from "../../src/bus"
55import { Log } from "../../src/util/log"
66import { Instance } from "../../src/project/instance"
7+ import { MessageV2 } from "../../src/session/message-v2"
8+ import { Identifier } from "../../src/id/id"
79
810const projectRoot = path . join ( __dirname , "../.." )
911Log . init ( { print : false } )
@@ -69,3 +71,68 @@ describe("session.started event", () => {
6971 } )
7072 } )
7173} )
74+
75+ describe ( "step-finish token propagation via Bus event" , ( ) => {
76+ test ( "non-zero tokens propagate through PartUpdated event" , async ( ) => {
77+ await Instance . provide ( {
78+ directory : projectRoot ,
79+ fn : async ( ) => {
80+ const session = await Session . create ( { } )
81+
82+ const messageID = Identifier . ascending ( "message" )
83+ await Session . updateMessage ( {
84+ id : messageID ,
85+ sessionID : session . id ,
86+ role : "user" ,
87+ time : { created : Date . now ( ) } ,
88+ agent : "user" ,
89+ model : { providerID : "test" , modelID : "test" } ,
90+ tools : { } ,
91+ mode : "" ,
92+ } as unknown as MessageV2 . Info )
93+
94+ let received : MessageV2 . Part | undefined
95+ const unsub = Bus . subscribe ( MessageV2 . Event . PartUpdated , ( event ) => {
96+ received = event . properties . part
97+ } )
98+
99+ const tokens = {
100+ total : 1500 ,
101+ input : 500 ,
102+ output : 800 ,
103+ reasoning : 200 ,
104+ cache : { read : 100 , write : 50 } ,
105+ }
106+
107+ const partInput = {
108+ id : Identifier . ascending ( "part" ) ,
109+ messageID,
110+ sessionID : session . id ,
111+ type : "step-finish" as const ,
112+ reason : "stop" ,
113+ cost : 0.005 ,
114+ tokens,
115+ }
116+
117+ await Session . updatePart ( partInput )
118+
119+ await new Promise ( ( resolve ) => setTimeout ( resolve , 100 ) )
120+
121+ expect ( received ) . toBeDefined ( )
122+ expect ( received ! . type ) . toBe ( "step-finish" )
123+ const finish = received as MessageV2 . StepFinishPart
124+ expect ( finish . tokens . input ) . toBe ( 500 )
125+ expect ( finish . tokens . output ) . toBe ( 800 )
126+ expect ( finish . tokens . reasoning ) . toBe ( 200 )
127+ expect ( finish . tokens . total ) . toBe ( 1500 )
128+ expect ( finish . tokens . cache . read ) . toBe ( 100 )
129+ expect ( finish . tokens . cache . write ) . toBe ( 50 )
130+ expect ( finish . cost ) . toBe ( 0.005 )
131+ expect ( received ) . not . toBe ( partInput )
132+
133+ unsub ( )
134+ await Session . remove ( session . id )
135+ } ,
136+ } )
137+ } , { timeout : 30000 } )
138+ } )
0 commit comments