Something went wrong. Try again.
My version of https://github.com/vercel-labs/fx
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310231123122313231423152316231723182319232023212322232323242325232623272328232923302331233223332334233523362337233823392340234123422343234423452346234723482349235023512352235323542355235623572358235923602361236223632364236523662367236823692370237123722373237423752376237723782379238023812382238323842385238623872388238923902391239223932394239523962397239823992400240124022403240424052406240724082409241024112412241324142415241624172418241924202421242224232424242524262427242824292430243124322433243424352436243724382439244024412442244324442445244624472448244924502451245224532454245524562457245824592460246124622463246424652466246724682469247024712472247324742475247624772478247924802481248224832484248524862487248824892490249124922493249424952496249724982499250025012502250325042505250625072508250925102511251225132514251525162517251825192520252125222523252425252526252725282529253025312532253325342535253625372538253925402541254225432544254525462547254825492550255125522553255425552556255725582559256025612562256325642565256625672568256925702571257225732574257525762577257825792580258125822583258425852586258725882589259025912592259325942595259625972598259926002601260226032604260526062607260826092610261126122613261426152616261726182619262026212622262326242625262626272628262926302631263226332634263526362637263826392640264126422643264426452646264726482649265026512652265326542655265626572658265926602661266226632664266526662667266826692670267126722673267426752676267726782679268026812682268326842685268626872688268926902691269226932694269526962697269826992700270127022703270427052706270727082709271027112712271327142715271627172718271927202721272227232724272527262727272827292730273127322733273427352736273727382739274027412742274327442745274627472748274927502751275227532754275527562757275827592760276127622763276427652766276727682769277027712772277327742775277627772778277927802781278227832784278527862787278827892790279127922793279427952796279727982799280028012802280328042805280628072808280928102811281228132814281528162817281828192820282128222823282428252826282728282829283028312832283328342835283628372838283928402841284228432844284528462847284828492850285128522853285428552856285728582859286028612862286328642865286628672868286928702871287228732874287528762877287828792880288128822883288428852886288728882889289028912892289328942895289628972898289929002901290229032904290529062907290829092910291129122913291429152916291729182919292029212922292329242925292629272928292929302931293229332934293529362937293829392940294129422943294429452946294729482949295029512952295329542955295629572958295929602961296229632964296529662967296829692970297129722973297429752976297729782979298029812982298329842985298629872988298929902991299229932994299529962997299829993000300130023003300430053006300730083009301030113012301330143015301630173018301930203021302230233024302530263027302830293030303130323033303430353036303730383039304030413042304330443045304630473048304930503051305230533054305530563057305830593060306130623063306430653066306730683069307030713072307330743075307630773078307930803081308230833084308530863087308830893090309130923093309430953096309730983099310031013102310331043105310631073108310931103111311231133114311531163117311831193120312131223123312431253126312731283129313031313132313331343135313631373138313931403141314231433144314531463147314831493150315131523153315431553156315731583159316031613162316331643165316631673168316931703171317231733174317531763177317831793180318131823183318431853186318731883189319031913192319331943195319631973198319932003201320232033204320532063207320832093210321132123213321432153216321732183219322032213222322332243225322632273228322932303231323232333234323532363237323832393240324132423243324432453246324732483249325032513252325332543255325632573258325932603261326232633264326532663267326832693270327132723273327432753276327732783279328032813282328332843285328632873288328932903291329232933294329532963297329832993300330133023303330433053306330733083309331033113312331333143315331633173318331933203321332233233324332533263327332833293330333133323333333433353336333733383339334033413342334333443345334633473348334933503351335233533354335533563357335833593360336133623363336433653366336733683369337033713372337333743375337633773378337933803381338233833384338533863387338833893390339133923393339433953396339733983399340034013402340334043405340634073408340934103411341234133414341534163417341834193420342134223423342434253426342734283429343034313432343334343435343634373438343934403441344234433444344534463447344834493450345134523453345434553456345734583459346034613462346334643465346634673468346934703471347234733474347534763477347834793480348134823483348434853486348734883489349034913492349334943495349634973498349935003501350235033504350535063507350835093510351135123513351435153516351735183519352035213522352335243525352635273528352935303531353235333534353535363537353835393540354135423543354435453546354735483549355035513552355335543555355635573558355935603561356235633564356535663567356835693570357135723573357435753576357735783579358035813582358335843585358635873588358935903591359235933594359535963597359835993600360136023603360436053606360736083609361036113612361336143615361636173618361936203621362236233624362536263627362836293630363136323633363436353636363736383639364036413642364336443645364636473648364936503651365236533654365536563657365836593660366136623663366436653666366736683669367036713672367336743675367636773678367936803681368236833684368536863687368836893690369136923693369436953696369736983699370037013702370337043705370637073708370937103711371237133714371537163717371837193720372137223723372437253726372737283729373037313732373337343735373637373738373937403741374237433744374537463747374837493750375137523753375437553756375737583759376037613762376337643765376637673768376937703771377237733774377537763777377837793780378137823783378437853786378737883789379037913792379337943795379637973798379938003801380238033804380538063807380838093810381138123813381438153816381738183819382038213822382338243825382638273828382938303831383238333834383538363837383838393840384138423843384438453846384738483849385038513852385338543855385638573858385938603861386238633864386538663867386838693870387138723873387438753876387738783879388038813882388338843885388638873888388938903891389238933894389538963897389838993900390139023903390439053906390739083909391039113912391339143915391639173918391939203921392239233924392539263927392839293930393139323933393439353936393739383939394039413942394339443945394639473948394939503951395239533954395539563957395839593960396139623963396439653966396739683969397039713972397339743975397639773978397939803981398239833984398539863987398839893990399139923993399439953996399739983999400040014002400340044005400640074008400940104011401240134014401540164017401840194020402140224023402440254026402740284029403040314032403340344035403640374038403940404041404240434044404540464047404840494050405140524053405440554056405740584059406040614062406340644065406640674068406940704071407240734074407540764077407840794080408140824083408440854086408740884089409040914092409340944095409640974098409941004101410241034104410541064107410841094110411141124113411441154116411741184119412041214122412341244125412641274128412941304131413241334134413541364137413841394140414141424143414441454146414741484149415041514152415341544155415641574158415941604161416241634164416541664167416841694170417141724173417441754176417741784179418041814182418341844185418641874188418941904191419241934194419541964197419841994200420142024203420442054206420742084209421042114212421342144215421642174218421942204221422242234224422542264227422842294230423142324233423442354236423742384239424042414242424342444245424642474248424942504251425242534254425542564257425842594260426142624263426442654266426742684269427042714272427342744275427642774278427942804281428242834284428542864287428842894290429142924293429442954296429742984299430043014302430343044305430643074308430943104311431243134314431543164317431843194320432143224323432443254326432743284329433043314332433343344335433643374338433943404341434243434344434543464347434843494350435143524353435443554356435743584359436043614362436343644365436643674368436943704371437243734374437543764377437843794380438143824383438443854386438743884389439043914392439343944395439643974398439944004401440244034404440544064407440844094410441144124413441444154416441744184419442044214422442344244425442644274428442944304431443244334434443544364437443844394440444144424443444444454446444744484449445044514452445344544455445644574458445944604461446244634464446544664467446844694470447144724473447444754476447744784479448044814482448344844485448644874488448944904491449244934494449544964497449844994500450145024503450445054506450745084509451045114512451345144515451645174518451945204521452245234524452545264527452845294530453145324533453445354536453745384539454045414542454345444545454645474548454945504551455245534554455545564557455845594560456145624563456445654566456745684569457045714572457345744575457645774578457945804581458245834584458545864587458845894590459145924593459445954596459745984599460046014602460346044605460646074608460946104611461246134614461546164617461846194620462146224623462446254626462746284629463046314632463346344635463646374638463946404641464246434644464546464647464846494650465146524653465446554656465746584659466046614662466346644665466646674668466946704671467246734674467546764677467846794680468146824683468446854686468746884689469046914692469346944695469646974698469947004701470247034704470547064707470847094710471147124713471447154716471747184719472047214722472347244725472647274728472947304731473247334734473547364737473847394740474147424743474447454746474747484749475047514752475347544755475647574758475947604761476247634764476547664767476847694770477147724773477447754776477747784779478047814782478347844785478647874788478947904791479247934794479547964797479847994800480148024803480448054806480748084809481048114812481348144815481648174818481948204821482248234824482548264827482848294830483148324833483448354836483748384839484048414842484348444845484648474848484948504851485248534854485548564857485848594860486148624863486448654866486748684869487048714872487348744875487648774878487948804881488248834884488548864887488848894890489148924893489448954896489748984899490049014902490349044905490649074908490949104911491249134914491549164917491849194920492149224923492449254926492749284929493049314932493349344935493649374938493949404941494249434944494549464947494849494950495149524953495449554956495749584959496049614962496349644965496649674968496949704971497249734974497549764977497849794980498149824983498449854986498749884989499049914992499349944995499649974998499950005001500250035004500550065007500850095010501150125013501450155016501750185019502050215022502350245025502650275028502950305031503250335034503550365037503850395040504150425043504450455046504750485049505050515052505350545055505650575058505950605061506250635064506550665067506850695070507150725073507450755076507750785079508050815082508350845085508650875088508950905091509250935094509550965097509850995100510151025103510451055106510751085109511051115112511351145115511651175118511951205121512251235124512551265127512851295130513151325133513451355136513751385139514051415142514351445145514651475148514951505151515251535154515551565157515851595160516151625163516451655166516751685169517051715172517351745175517651775178517951805181518251835184518551865187518851895190519151925193519451955196519751985199520052015202520352045205520652075208520952105211521252135214521552165217521852195220522152225223522452255226522752285229523052315232523352345235523652375238523952405241524252435244524552465247524852495250525152525253525452555256525752585259526052615262526352645265526652675268526952705271527252735274527552765277527852795280528152825283528452855286528752885289529052915292529352945295529652975298529953005301530253035304530553065307530853095310531153125313531453155316531753185319532053215322532353245325532653275328532953305331533253335334533553365337533853395340534153425343534453455346534753485349535053515352535353545355535653575358535953605361536253635364536553665367536853695370537153725373537453755376537753785379538053815382538353845385538653875388538953905391539253935394539553965397539853995400540154025403540454055406540754085409541054115412541354145415541654175418541954205421542254235424542554265427542854295430543154325433543454355436543754385439544054415442544354445445544654475448544954505451545254535454545554565457545854595460546154625463546454655466546754685469547054715472547354745475547654775478547954805481548254835484548554865487548854895490549154925493549454955496549754985499550055015502550355045505550655075508550955105511551255135514551555165517551855195520552155225523552455255526552755285529553055315532553355345535553655375538553955405541554255435544554555465547554855495550555155525553555455555556555755585559556055615562556355645565556655675568556955705571557255735574557555765577557855795580558155825583558455855586558755885589559055915592559355945595559655975598559956005601560256035604560556065607560856095610561156125613561456155616561756185619562056215622562356245625562656275628562956305631563256335634563556365637563856395640564156425643564456455646564756485649565056515652565356545655565656575658565956605661566256635664566556665667566856695670567156725673567456755676567756785679568056815682568356845685568656875688568956905691569256935694569556965697569856995700570157025703570457055706570757085709571057115712571357145715571657175718571957205721572257235724572557265727572857295730573157325733573457355736573757385739574057415742574357445745574657475748574957505751575257535754575557565757575857595760576157625763576457655766576757685769577057715772577357745775577657775778577957805781578257835784578557865787578857895790579157925793579457955796579757985799580058015802580358045805580658075808580958105811581258135814581558165817581858195820582158225823582458255826582758285829583058315832583358345835583658375838583958405841584258435844584558465847584858495850585158525853585458555856585758585859586058615862586358645865586658675868586958705871587258735874587558765877587858795880588158825883588458855886588758885889589058915892589358945895589658975898589959005901590259035904590559065907590859095910591159125913591459155916591759185919592059215922592359245925592659275928592959305931593259335934593559365937593859395940594159425943594459455946594759485949595059515952595359545955595659575958595959605961596259635964596559665967596859695970597159725973597459755976597759785979598059815982598359845985598659875988598959905991599259935994599559965997599859996000600160026003600460056006600760086009601060116012601360146015601660176018601960206021602260236024602560266027602860296030603160326033603460356036603760386039604060416042604360446045604660476048604960506051605260536054605560566057605860596060606160626063606460656066606760686069607060716072607360746075607660776078607960806081608260836084608560866087608860896090609160926093609460956096609760986099610061016102610361046105610661076108610961106111611261136114611561166117611861196120612161226123612461256126612761286129613061316132613361346135613661376138613961406141614261436144614561466147614861496150615161526153615461556156615761586159616061616162616361646165616661676168616961706171617261736174617561766177617861796180618161826183618461856186618761886189619061916192619361946195619661976198const std = @import("std");const builtin = @import("builtin");const secret = @import("../auth/secret.zig");const debug_trace = @import("../shared/debug_trace.zig");const mem_utils = @import("../shared/mem_utils.zig");const generation_fact_codec = @import("generation_fact_codec.zig");const generation_usage = @import("generation_usage_provider.zig");const io_mod = @import("../shared/io.zig");const types = @import("../shared/types.zig");const stream_provider = @import("../agent/stream_provider.zig");const model_provider = @import("../config/model_provider.zig");const credential_authority = @import("../auth/credential_authority.zig");pub const usage_report = @import("usage_report.zig");
const Allocator = std.mem.Allocator;const Sha256 = std.crypto.hash.sha2.Sha256;
const max_models: usize = 32;const max_pending_generations: usize = 16;const max_publication_backlog: usize = 16;const max_usage_incidents: usize = 16;const max_model_bytes: usize = 1024;const max_origin_bytes: usize = 2048;const max_team_bytes: usize = 255;const max_identifier_bytes: usize = 8 * 1024;const max_active_invocations: usize = 64;const shutdown_reconciliation_budget_ms: usize = 250;pub const max_snapshot_bytes: usize = 256 * 1024;
/// Completeness of the saved billing window rendered by `/cost`.pub const Availability = enum { /// Every billable invocation has an authoritative generation record. complete, /// Known generation identities still require reconciliation. pending, /// Billing may have occurred without a recoverable authoritative record. incomplete, /// The persisted session predates usage accounting. legacy,};
/// Billing consequence of one logical Gateway invocation across retries.pub const DeliveryOutcome = enum { /// Delivery was proven not to create a billable generation. unbilled, /// A validated generation identity can be reconciled. observed_generation, /// Delivery succeeded but no usable generation identity was returned. possibly_billed_without_identity, /// Transport evidence cannot prove whether a billable generation was created. ambiguous_delivery,};
/// Host-owned durable sink. Its context and allocator remain borrowed until/// the sink is unset; `persist` must not retain the snapshot, checkpoint, or/// reconfigure the sink.pub const UsageCheckpointSink = struct { context: *anyopaque, allocator: Allocator, persist: *const fn (context: *anyopaque, snapshot: Snapshot) anyerror!void,};
/// Host-owned local-profile publication sink. The callback receives a borrowed/// event and must return only after it is durably accepted or deduplicated.pub const ProfilePublicationSink = struct { context: *anyopaque, allocator: Allocator, publish: *const fn ( context: *anyopaque, event: usage_report.ProfileEvent, ) anyerror!void,};
const ProfilePublicationBatch = struct { pending: []usage_report.PendingMarker, incidents: []usage_report.Incident, facts: []usage_report.GenerationFact, checkpoint_changed: bool,
fn deinit(self: *ProfilePublicationBatch, alloc: Allocator) void { for (self.pending) |*marker| marker.deinit(alloc); alloc.free(self.pending); alloc.free(self.incidents); for (self.facts) |*fact| fact.deinit(alloc); alloc.free(self.facts); self.* = undefined; }};
/// One admitted provider invocation. `begin` durably reserves it before network I/O;/// every successful reservation must terminate through `fail` or `complete`.pub const InvocationObservation = struct { usage: ?*Usage, sequence: u64 = 0, started_at_ms: i64,
pub fn begin(usage: ?*Usage) !InvocationObservation { return .{ .usage = usage, .sequence = if (usage) |ledger| try ledger.reserveInvocationDurably() else 0, .started_at_ms = io_mod.milliTimestamp(), }; }
pub fn fail(self: InvocationObservation, outcome: DeliveryOutcome) !void { const ledger = self.usage orelse return; try ledger.finishInvocationDurably(self.sequence, self.elapsedMs(), outcome); }
pub fn complete( self: InvocationObservation, alloc: Allocator, completion: types.ModelCompletion, usage_outcome: stream_provider.UsageOutcome, ) !void { const ledger = self.usage orelse return; ledger.observeContextUsage(self.sequence, completion.usage); switch (usage_outcome) { .unavailable => |availability| { const delivery: DeliveryOutcome = switch (availability) { .unbilled => .unbilled, .possibly_billed => if (completion.delivery_ambiguous) .ambiguous_delivery else .possibly_billed_without_identity, }; try ledger.finishInvocationDurably(self.sequence, self.elapsedMs(), delivery); if (delivery == .possibly_billed_without_identity) { debug_trace.logf( "session", "usage billing incomplete sequence={d} reason=generation_identity_missing", .{self.sequence}, ); } return; }, .exact, .deferred => {}, } const delivery: DeliveryOutcome = if (completion.delivery_ambiguous) .ambiguous_delivery else if (!completion.generation_metadata_invalid) .observed_generation else .possibly_billed_without_identity; if (completion.generation_metadata_invalid) { try ledger.finishInvocationDurably(self.sequence, self.elapsedMs(), delivery); debug_trace.logf( "session", "usage generation ignored sequence={d} metadata_invalid=true", .{self.sequence}, ); return; } switch (usage_outcome) { .exact => |provider| { const generation_id = completion.generation_id orelse { try ledger.finishInvocationDurably( self.sequence, self.elapsedMs(), if (completion.delivery_ambiguous) .ambiguous_delivery else .possibly_billed_without_identity, ); debug_trace.logf( "session", "usage billing incomplete sequence={d} reason=generation_identity_missing", .{self.sequence}, ); return; }; const billing = completion.billing orelse { try ledger.finishInvocationDurably( self.sequence, self.elapsedMs(), if (completion.delivery_ambiguous) .ambiguous_delivery else .possibly_billed_without_identity, ); debug_trace.logf( "session", "usage billing incomplete sequence={d} reason=exact_usage_missing", .{self.sequence}, ); return; }; const accepted = try ledger.finishExactInvocationDurably( alloc, self.sequence, self.elapsedMs(), delivery, provider, generation_id, billing, ); if (!accepted) return; debug_trace.logf( "session", "usage exact generation settled sequence={d} id={s}", .{ self.sequence, generation_id }, ); }, .deferred => |reference| { const accepted = try ledger.finishDeferredInvocationDurably( alloc, self.sequence, self.elapsedMs(), delivery, reference, ); if (!accepted) return; debug_trace.logf( "session", "usage generation queued sequence={d} id={s}", .{ self.sequence, reference.generation_id }, ); ledger.scheduleProfilePublicationDrain(); }, .unavailable => unreachable, } }
fn elapsedMs(self: InvocationObservation) u64 { const ended_at_ms = io_mod.milliTimestamp(); if (ended_at_ms <= self.started_at_ms) return 0; return std.math.cast(u64, ended_at_ms - self.started_at_ms) orelse std.math.maxInt(u64); }};
const GenerationRecord = generation_usage.Record;
pub const ModelAggregate = struct { model: []u8, first_sequence: u64, total_cost: f64 = 0, input_tokens: u64 = 0, output_tokens: u64 = 0, cache_read_tokens: u64 = 0, cache_write_tokens: u64 = 0, reasoning_tokens: ?u64 = null, request_count: ?u64 = null, billable_web_search_calls: u64 = 0,
fn deinit(self: *ModelAggregate, alloc: Allocator) void { alloc.free(self.model); self.* = undefined; }
fn dupe(self: ModelAggregate, alloc: Allocator) Allocator.Error!ModelAggregate { return .{ .model = try alloc.dupe(u8, self.model), .first_sequence = self.first_sequence, .total_cost = self.total_cost, .input_tokens = self.input_tokens, .output_tokens = self.output_tokens, .cache_read_tokens = self.cache_read_tokens, .cache_write_tokens = self.cache_write_tokens, .reasoning_tokens = self.reasoning_tokens, .request_count = self.request_count, .billable_web_search_calls = self.billable_web_search_calls, }; }};
pub const PendingGeneration = struct { id: []u8, sequence: u64, provider: model_provider.ProviderId = .gateway, origin: []u8, team: ?[]u8, credential_source: ?types.CredentialSource = null, credential_identity: ?credential_authority.Identity = null, account_id: ?[]u8 = null, observed_at_ms: ?i64 = null,
fn deinit(self: *PendingGeneration, alloc: Allocator) void { alloc.free(self.id); alloc.free(self.origin); if (self.team) |team| alloc.free(team); if (self.account_id) |account_id| alloc.free(account_id); self.* = undefined; }
fn dupe(self: PendingGeneration, alloc: Allocator) Allocator.Error!PendingGeneration { const id = try alloc.dupe(u8, self.id); errdefer mem_utils.free(alloc, id); const origin = try alloc.dupe(u8, self.origin); errdefer mem_utils.free(alloc, origin); const team = if (self.team) |value| try alloc.dupe(u8, value) else null; errdefer if (team) |value| mem_utils.free(alloc, value); const account_id = if (self.account_id) |value| try alloc.dupe(u8, value) else null; return .{ .id = id, .sequence = self.sequence, .provider = self.provider, .origin = origin, .team = team, .credential_source = self.credential_source, .credential_identity = self.credential_identity, .account_id = account_id, .observed_at_ms = self.observed_at_ms, }; }};
pub const Snapshot = struct { billing: Availability, api_duration_complete: bool, wall_duration_complete: bool, code_complete: bool, next_sequence: u64, settled_through_sequence: u64, api_duration_ms: u64, wall_duration_ms: u64, total_cost: f64, input_tokens: u64, output_tokens: u64, cache_read_tokens: u64, cache_write_tokens: u64, reasoning_tokens: ?u64 = null, request_count: ?u64 = null, billable_web_search_calls: u64, lines_added: u64, lines_removed: u64, models: []ModelAggregate, pending: []PendingGeneration, publication_backlog: []usage_report.GenerationFact = &.{}, incidents: []usage_report.Incident = &.{},
pub fn deinit(self: *Snapshot, alloc: Allocator) void { for (self.models) |*model| model.deinit(alloc); alloc.free(self.models); for (self.pending) |*generation| generation.deinit(alloc); alloc.free(self.pending); for (self.publication_backlog) |*fact| fact.deinit(alloc); if (self.publication_backlog.len > 0) alloc.free(self.publication_backlog); if (self.incidents.len > 0) alloc.free(self.incidents); self.* = undefined; }};
/// True when a durable session can still contribute profile usage that is not/// proven present in the profile ledger.pub fn needsProfileRecovery(snapshot: Snapshot) bool { return snapshot.settled_through_sequence != snapshot.next_sequence - 1 or snapshot.pending.len > 0 or snapshot.publication_backlog.len > 0 or snapshot.incidents.len > 0;}
pub const Usage = struct { mutex: std.Io.Mutex = .init, checkpoint_mutex: std.Io.Mutex = .init, checkpoint_sink: ?UsageCheckpointSink = null, publication_mutex: std.Io.Mutex = .init, publication_sink: ?ProfilePublicationSink = null, billing: Availability, api_duration_complete: bool, wall_duration_complete: bool, code_complete: bool, next_sequence: u64 = 1, settled_through_sequence: u64 = 0, api_duration_ms: u64 = 0, wall_duration_ms: u64 = 0, active_started_at_ms: i64 = 0, active_sequences: [max_active_invocations]u64 = undefined, active_sequence_count: usize = 0, total_cost: f64 = 0, input_tokens: u64 = 0, output_tokens: u64 = 0, cache_read_tokens: u64 = 0, cache_write_tokens: u64 = 0, reasoning_tokens: ?u64, request_count: ?u64, billable_web_search_calls: u64 = 0, latest_context_sequence: u64 = 0, latest_context_used: ?u64 = null, lines_added: u64 = 0, lines_removed: u64 = 0, models: std.ArrayList(ModelAggregate) = .empty, pending: std.ArrayList(PendingGeneration) = .empty, publication_backlog: std.ArrayList(usage_report.GenerationFact) = .empty, incidents: [max_usage_incidents]usage_report.Incident = undefined, incident_count: usize = 0, dirty: bool = false, reconciliation_mutex: std.Io.Mutex = .init, reconciliation_thread: ?std.Thread = null, reconciliation_done: std.atomic.Value(bool) = .init(true), reconciliation_cancel: std.atomic.Value(bool) = .init(false), reconciliation_work_epoch: std.atomic.Value(u64) = .init(0), reconciliation_key_digest: ?[Sha256.digest_length]u8 = null, reconciliation_authority: ?ReconciliationAuthority = null, reconciliation_credential_blocked: bool = false, generation_usage_providers: generation_usage.Set = .{}, publication_drain_mutex: std.Io.Mutex = .init, publication_drain_thread: ?std.Thread = null, publication_drain_done: std.atomic.Value(bool) = .init(true), publication_drain_cancel: std.atomic.Value(bool) = .init(false), publication_drain_epoch: std.atomic.Value(u64) = .init(0),
pub fn initFresh() Usage { return .{ .billing = .complete, .api_duration_complete = true, .wall_duration_complete = true, .code_complete = true, .reasoning_tokens = 0, .request_count = 0, }; }
pub fn initFreshWithProvider(provider: generation_usage.Provider) Usage { var usage = initFresh(); usage.generation_usage_providers = generation_usage.Set.gatewayOnly(provider); return usage; }
pub fn initFreshWithProviders(providers: generation_usage.Set) Usage { var usage = initFresh(); usage.generation_usage_providers = providers; return usage; }
pub fn initLegacy() Usage { return .{ .billing = .legacy, .api_duration_complete = false, .wall_duration_complete = false, .code_complete = false, .reasoning_tokens = null, .request_count = null, }; }
pub fn deinit(self: *Usage, alloc: Allocator) void { self.stopReconciliation(); self.stopPublicationDrain(); self.clearOwned(alloc); self.* = undefined; }
pub fn resetFresh(self: *Usage, alloc: Allocator) void { self.stopReconciliation(); self.stopPublicationDrain(); self.reset(alloc, true); }
pub fn resetLegacy(self: *Usage, alloc: Allocator) void { self.stopReconciliation(); self.stopPublicationDrain(); self.reset(alloc, false); }
pub fn restoreLegacyWallDuration( self: *Usage, session_started_at_ms: i64, ) void { const now_ms = io_mod.milliTimestamp(); self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); if (session_started_at_ms <= 0 or session_started_at_ms > now_ms) return; self.wall_duration_complete = true; self.active_started_at_ms = session_started_at_ms; }
pub fn reserveInvocation(self: *Usage) !u64 { self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); if (self.active_started_at_ms == 0) { self.active_started_at_ms = io_mod.milliTimestamp(); } const sequence = self.next_sequence; if (self.active_sequence_count == max_active_invocations) { debug_trace.logf( "session", "usage invocation reservation failed reason=active_capacity limit={d}", .{max_active_invocations}, ); return error.UsageCapacityExceeded; } const next_sequence = std.math.add(u64, sequence, 1) catch { debug_trace.logf( "session", "usage invocation reservation failed reason=sequence_overflow sequence={d}", .{sequence}, ); return error.UsageSequenceOverflow; }; self.next_sequence = next_sequence; self.active_sequences[self.active_sequence_count] = sequence; self.active_sequence_count += 1; self.dirty = true; return sequence; }
pub fn configureCheckpointSink(self: *Usage, sink: ?UsageCheckpointSink) void { self.checkpoint_mutex.lockUncancelable(io_mod.getIo()); defer self.checkpoint_mutex.unlock(io_mod.getIo()); self.checkpoint_sink = sink; }
pub fn configurePublicationSink( self: *Usage, sink: ?ProfilePublicationSink, ) void { self.publication_mutex.lockUncancelable(io_mod.getIo()); self.publication_sink = sink; self.publication_mutex.unlock(io_mod.getIo()); if (sink != null) self.flushProfilePublications(); }
pub fn persistCheckpoint(self: *Usage) bool { self.checkpoint_mutex.lockUncancelable(io_mod.getIo()); defer self.checkpoint_mutex.unlock(io_mod.getIo()); return self.persistCheckpointBestEffortLocked(); }
fn reserveInvocationDurably(self: *Usage) !u64 { self.checkpoint_mutex.lockUncancelable(io_mod.getIo()); defer self.checkpoint_mutex.unlock(io_mod.getIo()); const sequence = try self.reserveInvocation(); self.persistCheckpointRequiredLocked() catch |err| { self.finishInvocation(sequence, 0, .unbilled); return err; }; return sequence; }
fn finishInvocationDurably( self: *Usage, sequence: u64, duration_ms: u64, outcome: DeliveryOutcome, ) !void { self.checkpoint_mutex.lockUncancelable(io_mod.getIo()); errdefer self.checkpoint_mutex.unlock(io_mod.getIo()); self.finishInvocation(sequence, duration_ms, outcome); _ = try self.persistCheckpointForContinuationLocked(); self.checkpoint_mutex.unlock(io_mod.getIo()); self.scheduleProfilePublicationDrain(); }
fn finishObservedInvocationDurably( self: *Usage, alloc: Allocator, sequence: u64, duration_ms: u64, outcome: DeliveryOutcome, id: []const u8, origin: []const u8, team: ?[]const u8, ) !bool { self.checkpoint_mutex.lockUncancelable(io_mod.getIo()); errdefer self.checkpoint_mutex.unlock(io_mod.getIo()); const accepted = self.finishObservedInvocationAccepted( alloc, sequence, duration_ms, outcome, id, origin, team, ) catch |err| { self.markBillingIncomplete(); _ = try self.persistCheckpointForContinuationLocked(); debug_trace.logf( "session", "usage generation checkpointed incomplete reason={s}", .{@errorName(err)}, ); self.checkpoint_mutex.unlock(io_mod.getIo()); self.scheduleProfilePublicationDrain(); return false; }; _ = try self.persistCheckpointForContinuationLocked(); self.checkpoint_mutex.unlock(io_mod.getIo()); return accepted; }
fn finishDeferredInvocationDurably( self: *Usage, alloc: Allocator, sequence: u64, duration_ms: u64, outcome: DeliveryOutcome, reference: stream_provider.DeferredUsageReference, ) !bool { self.checkpoint_mutex.lockUncancelable(io_mod.getIo()); errdefer self.checkpoint_mutex.unlock(io_mod.getIo()); const accepted = self.finishDeferredInvocationAccepted( alloc, sequence, duration_ms, outcome, reference, ) catch |err| { self.markBillingIncomplete(); _ = try self.persistCheckpointForContinuationLocked(); debug_trace.logf( "session", "usage generation checkpointed incomplete reason={s}", .{@errorName(err)}, ); self.checkpoint_mutex.unlock(io_mod.getIo()); self.scheduleProfilePublicationDrain(); return false; }; _ = try self.persistCheckpointForContinuationLocked(); self.checkpoint_mutex.unlock(io_mod.getIo()); return accepted; }
fn finishExactInvocationDurably( self: *Usage, alloc: Allocator, sequence: u64, duration_ms: u64, outcome: DeliveryOutcome, provider: model_provider.ProviderId, external_id: []const u8, billing: types.ProviderBilling, ) !bool { var canonical_buffer: [30]u8 = undefined; const canonical_id = canonicalExactGenerationId( provider, external_id, &canonical_buffer, ) catch |err| { try self.finishInvocationDurably( sequence, duration_ms, if (outcome == .ambiguous_delivery) .ambiguous_delivery else .possibly_billed_without_identity, ); debug_trace.logf( "session", "usage exact generation rejected reason={s}", .{@errorName(err)}, ); return false; }; const record = GenerationRecord{ .id = canonical_id, .created_at_ms = billing.created_at_ms, .model = billing.model, .total_cost = billing.total_cost, .input_tokens = billing.input_tokens, .output_tokens = billing.output_tokens, .cache_read_tokens = billing.cache_read_tokens, .cache_write_tokens = billing.cache_write_tokens, .reasoning_tokens = billing.reasoning_tokens, .billable_web_search_calls = billing.billable_web_search_calls, }; validateGenerationRecord(record) catch |err| { try self.finishInvocationDurably( sequence, duration_ms, if (outcome == .ambiguous_delivery) .ambiguous_delivery else .possibly_billed_without_identity, ); debug_trace.logf( "session", "usage exact generation rejected reason={s}", .{@errorName(err)}, ); return false; };
self.checkpoint_mutex.lockUncancelable(io_mod.getIo()); errdefer self.checkpoint_mutex.unlock(io_mod.getIo()); const durable_bridge = self.checkpoint_sink != null; const accepted = self.finishExactInvocationAccepted( alloc, sequence, duration_ms, outcome, provider, canonical_id, record, durable_bridge, ) catch |err| { self.markBillingIncomplete(); _ = try self.persistCheckpointForContinuationLocked(); self.checkpoint_mutex.unlock(io_mod.getIo()); debug_trace.logf( "session", "usage exact generation checkpointed incomplete reason={s}", .{@errorName(err)}, ); return false; }; if (!accepted) { _ = try self.persistCheckpointForContinuationLocked(); self.checkpoint_mutex.unlock(io_mod.getIo()); return false; } if (durable_bridge and !(try self.persistCheckpointForContinuationLocked())) { self.checkpoint_mutex.unlock(io_mod.getIo()); return false; } self.checkpoint_mutex.unlock(io_mod.getIo()); if (durable_bridge) self.scheduleProfilePublicationDrain(); return true; }
fn finishExactInvocationAccepted( self: *Usage, alloc: Allocator, sequence: u64, duration_ms: u64, outcome: DeliveryOutcome, provider: model_provider.ProviderId, generation_id: []const u8, record: GenerationRecord, durable_bridge: bool, ) !bool { self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); if (!self.finishInvocationUnlocked(sequence, duration_ms, outcome)) return false; try self.observeGenerationFieldsUnlocked( alloc, sequence, generation_id, provider, exactUsageOrigin(provider), null, null, null, null, ); if (durable_bridge) { try self.stagePublicationBacklogUnlocked( alloc, generationFactBorrowed(record), ); } else { try self.applyGenerationUnlocked(alloc, record, false); } return true; }
fn persistCheckpointRequiredLocked(self: *Usage) !void { const sink = self.checkpoint_sink orelse return; var persisted = try self.snapshotCurrent(sink.allocator); defer persisted.deinit(sink.allocator); try sink.persist(sink.context, persisted); self.markClean(persisted); }
fn persistCheckpointBestEffortLocked(self: *Usage) bool { return self.persistCheckpointForContinuationLocked() catch false; }
fn persistCheckpointForContinuationLocked(self: *Usage) error{ SessionPersistenceUncertain, SessionWriterChanged, SessionWriterParked, SessionCommitFailed }!bool { const sink = self.checkpoint_sink orelse return true; var persisted = self.snapshotCurrent(sink.allocator) catch |err| { self.markBillingIncomplete(); debug_trace.logf( "session", "usage checkpoint unavailable; billing marked incomplete reason={s}", .{@errorName(err)}, ); return false; }; defer persisted.deinit(sink.allocator); sink.persist(sink.context, persisted) catch |err| { self.markBillingIncomplete(); debug_trace.logf( "session", "usage checkpoint unavailable; billing marked incomplete reason={s}", .{@errorName(err)}, ); switch (err) { error.SessionPersistenceUncertain, error.SessionWriterChanged, error.SessionWriterParked, error.SessionCommitFailed, => return @errorCast(err), else => {}, } return false; }; self.markClean(persisted); return true; }
pub fn finishInvocation( self: *Usage, sequence: u64, duration_ms: u64, outcome: DeliveryOutcome, ) void { self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); _ = self.finishInvocationUnlocked(sequence, duration_ms, outcome); }
pub fn finishObservedInvocation( self: *Usage, alloc: Allocator, sequence: u64, duration_ms: u64, outcome: DeliveryOutcome, id: []const u8, origin: []const u8, team: ?[]const u8, ) !void { _ = try self.finishObservedInvocationAccepted( alloc, sequence, duration_ms, outcome, id, origin, team, ); }
fn finishObservedInvocationAccepted( self: *Usage, alloc: Allocator, sequence: u64, duration_ms: u64, outcome: DeliveryOutcome, id: []const u8, origin: []const u8, team: ?[]const u8, ) !bool { self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); if (!self.finishInvocationUnlocked(sequence, duration_ms, outcome)) return false; try validateGenerationId(id); try validateOrigin(origin); if (team) |value| try validateTeam(value); self.observeGenerationUnlocked(alloc, sequence, id, origin, team) catch |err| { self.billing = .incomplete; self.dirty = true; return err; }; return true; }
pub fn finishDeferredInvocation( self: *Usage, alloc: Allocator, sequence: u64, duration_ms: u64, outcome: DeliveryOutcome, reference: stream_provider.DeferredUsageReference, ) !void { _ = try self.finishDeferredInvocationAccepted( alloc, sequence, duration_ms, outcome, reference, ); }
fn finishDeferredInvocationAccepted( self: *Usage, alloc: Allocator, sequence: u64, duration_ms: u64, outcome: DeliveryOutcome, reference: stream_provider.DeferredUsageReference, ) !bool { self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); if (!self.finishInvocationUnlocked(sequence, duration_ms, outcome)) return false; try validateGenerationId(reference.generation_id); try validateOrigin(reference.scope); if (reference.tenant) |value| try validateTeam(value); if (reference.account_id) |value| try validateIdentifier(value); self.observeGenerationAuthorityUnlocked(alloc, sequence, reference) catch |err| { self.billing = .incomplete; self.dirty = true; return err; }; return true; }
fn finishInvocationUnlocked( self: *Usage, sequence: u64, duration_ms: u64, outcome: DeliveryOutcome, ) bool { if (sequence == 0 or sequence >= self.next_sequence) { self.billing = .incomplete; self.api_duration_complete = false; self.dirty = true; return false; } const active_index = for (self.active_sequences[0..self.active_sequence_count], 0..) |active, index| { if (active == sequence) break index; } else return false; self.active_sequence_count -= 1; self.active_sequences[active_index] = self.active_sequences[self.active_sequence_count]; self.api_duration_ms = std.math.add(u64, self.api_duration_ms, duration_ms) catch { self.billing = .incomplete; self.api_duration_complete = false; self.dirty = true; return false; }; switch (outcome) { .unbilled, .observed_generation => {}, .possibly_billed_without_identity, .ambiguous_delivery => { self.billing = .incomplete; self.recordIncidentUnlocked(.incomplete, io_mod.milliTimestamp()); }, } if (self.active_sequence_count == 0) { self.settled_through_sequence = self.next_sequence - 1; } self.dirty = true; return true; }
fn observeGenerationUnlocked( self: *Usage, alloc: Allocator, sequence: u64, id: []const u8, origin: []const u8, team: ?[]const u8, ) !void { return self.observeGenerationFieldsUnlocked( alloc, sequence, id, .gateway, origin, team, null, null, null, ); }
fn observeGenerationAuthorityUnlocked( self: *Usage, alloc: Allocator, sequence: u64, reference: stream_provider.DeferredUsageReference, ) !void { return self.observeGenerationFieldsUnlocked( alloc, sequence, reference.generation_id, reference.provider, reference.scope, reference.tenant, reference.credential_source, reference.credential_identity, reference.account_id, ); }
fn observeGenerationFieldsUnlocked( self: *Usage, alloc: Allocator, sequence: u64, id: []const u8, provider: model_provider.ProviderId, origin: []const u8, team: ?[]const u8, credential_source: ?types.CredentialSource, credential_identity: ?credential_authority.Identity, account_id: ?[]const u8, ) !void { for (self.pending.items) |pending| { if (pending.sequence == sequence and !std.mem.eql(u8, pending.id, id)) { self.billing = .incomplete; self.dirty = true; return error.ConflictingGenerationIdentity; } if (std.mem.eql(u8, pending.id, id)) { if (pending.sequence != sequence or !pending.provider.same_authority(provider) or !std.mem.eql(u8, pending.origin, origin) or !optionalStringsEqual(pending.team, team) or pending.credential_source != credential_source or !optionalCredentialIdentitiesEqual(pending.credential_identity, credential_identity) or !optionalStringsEqual(pending.account_id, account_id)) { self.billing = .incomplete; self.dirty = true; return error.ConflictingGenerationIdentity; } return; } } if (self.pending.items.len >= max_pending_generations) { self.billing = .incomplete; self.dirty = true; return error.UsageCapacityExceeded; } var added_identifier_bytes = std.math.add(usize, id.len, origin.len) catch return self.failOverflow(); if (team) |value| { added_identifier_bytes = std.math.add( usize, added_identifier_bytes, value.len, ) catch return self.failOverflow(); } if (account_id) |value| { added_identifier_bytes = std.math.add( usize, added_identifier_bytes, value.len, ) catch return self.failOverflow(); } const next_identifier_bytes = std.math.add( usize, self.identifierBytesUnlocked(), added_identifier_bytes, ) catch return self.failOverflow(); if (next_identifier_bytes > max_identifier_bytes) { self.billing = .incomplete; self.dirty = true; return error.UsageCapacityExceeded; } const owned_id = try alloc.dupe(u8, id); errdefer mem_utils.free(alloc, owned_id); const owned_origin = try alloc.dupe(u8, origin); errdefer mem_utils.free(alloc, owned_origin); const owned_team = if (team) |value| try alloc.dupe(u8, value) else null; errdefer if (owned_team) |value| mem_utils.free(alloc, value); const owned_account_id = if (account_id) |value| try alloc.dupe(u8, value) else null; errdefer if (owned_account_id) |value| mem_utils.free(alloc, value); try self.pending.append(alloc, .{ .id = owned_id, .sequence = sequence, .provider = provider, .origin = owned_origin, .team = owned_team, .credential_source = credential_source, .credential_identity = credential_identity, .account_id = owned_account_id, .observed_at_ms = io_mod.milliTimestamp(), }); if (self.billing == .complete) self.billing = .pending; self.dirty = true; }
pub fn applyGeneration(self: *Usage, alloc: Allocator, record: GenerationRecord) !void { try validateGenerationRecord(record); const fact = generationFactBorrowed(record);
self.checkpoint_mutex.lockUncancelable(io_mod.getIo()); const publication = self.publishFact(fact); if (publication != .durable and self.checkpoint_sink != null) { self.stagePublicationBacklog(alloc, fact) catch |err| { self.checkpoint_mutex.unlock(io_mod.getIo()); self.flushProfilePublications(); return err; }; if (!self.persistCheckpointBestEffortLocked()) { self.checkpoint_mutex.unlock(io_mod.getIo()); self.flushProfilePublications(); return error.ProfilePublicationRecoveryUnavailable; } self.checkpoint_mutex.unlock(io_mod.getIo()); self.flushProfilePublications(); // The rollback-readable event omits the richer publication fact, // so the pending generation remains the canonical recovery bridge // even when the sidecar preserves the full backlog. return; } if (publication == .failed) { self.stagePublicationBacklog(alloc, fact) catch |err| { self.checkpoint_mutex.unlock(io_mod.getIo()); self.flushProfilePublications(); return err; }; }
self.mutex.lockUncancelable(io_mod.getIo()); self.applyGenerationUnlocked( alloc, record, publication == .durable, ) catch |err| { self.mutex.unlock(io_mod.getIo()); self.checkpoint_mutex.unlock(io_mod.getIo()); self.flushProfilePublications(); return err; }; self.mutex.unlock(io_mod.getIo()); _ = self.persistCheckpointBestEffortLocked(); self.checkpoint_mutex.unlock(io_mod.getIo()); if (publication == .failed) self.flushProfilePublications(); }
fn applyGenerationUnlocked( self: *Usage, alloc: Allocator, record: GenerationRecord, published: bool, ) !void { const pending_index = for (self.pending.items, 0..) |pending, index| { if (std.mem.eql(u8, pending.id, record.id)) break index; } else { if (published) self.removePublicationBacklogUnlocked(alloc, record.id); return; }; const pending_sequence = self.pending.items[pending_index].sequence;
const model_index = for (self.models.items, 0..) |model, index| { if (std.mem.eql(u8, model.model, record.model)) break index; } else null; if (model_index == null and self.models.items.len >= max_models) { self.billing = .incomplete; self.dirty = true; return error.UsageCapacityExceeded; } if (model_index == null) { const resolved_pending = self.pending.items[pending_index]; var removed_identifier_bytes = resolved_pending.id.len +| resolved_pending.origin.len; if (resolved_pending.team) |team| removed_identifier_bytes +|= team.len; if (resolved_pending.account_id) |account_id| removed_identifier_bytes +|= account_id.len; const next_identifier_bytes = std.math.add( usize, self.identifierBytesUnlocked() -| removed_identifier_bytes, record.model.len, ) catch return self.failOverflow(); if (next_identifier_bytes > max_identifier_bytes) { self.billing = .incomplete; self.dirty = true; return error.UsageCapacityExceeded; } }
const next_total_cost = self.total_cost + record.total_cost; if (!std.math.isFinite(next_total_cost)) { self.billing = .incomplete; self.dirty = true; return error.UsageOverflow; } const next_input_tokens = std.math.add(u64, self.input_tokens, record.input_tokens) catch return self.failOverflow(); const next_output_tokens = std.math.add(u64, self.output_tokens, record.output_tokens) catch return self.failOverflow(); const next_cache_read_tokens = std.math.add(u64, self.cache_read_tokens, record.cache_read_tokens) catch return self.failOverflow(); const next_cache_write_tokens = std.math.add(u64, self.cache_write_tokens, record.cache_write_tokens) catch return self.failOverflow(); const next_reasoning_tokens = try addOptionalCounter( self.reasoning_tokens, record.reasoning_tokens, ); const next_request_count = if (self.request_count) |requests| std.math.add(u64, requests, 1) catch return self.failOverflow() else null; const next_web_search_calls = std.math.add( u64, self.billable_web_search_calls, record.billable_web_search_calls, ) catch return self.failOverflow();
var owned_model: ?[]u8 = null; errdefer if (owned_model) |model| mem_utils.free(alloc, model); if (model_index == null) { owned_model = try alloc.dupe(u8, record.model); try self.models.ensureUnusedCapacity(alloc, 1); }
if (model_index) |index| { try addRecordToModel(&self.models.items[index], record, pending_sequence); } else { var model = ModelAggregate{ .model = owned_model.?, .first_sequence = pending_sequence, .reasoning_tokens = 0, .request_count = 0, }; owned_model = null; try addRecordToModel(&model, record, pending_sequence); self.models.appendAssumeCapacity(model); } self.sortModelsBySequenceUnlocked();
self.total_cost = next_total_cost; self.input_tokens = next_input_tokens; self.output_tokens = next_output_tokens; self.cache_read_tokens = next_cache_read_tokens; self.cache_write_tokens = next_cache_write_tokens; self.reasoning_tokens = next_reasoning_tokens; self.request_count = next_request_count; self.billable_web_search_calls = next_web_search_calls;
var removed = self.pending.orderedRemove(pending_index); removed.deinit(alloc); if (published) self.removePublicationBacklogUnlocked(alloc, record.id); if (self.billing == .pending and self.pending.items.len == 0) { self.billing = .complete; } self.dirty = true; }
const PublicationResult = enum { not_configured, durable, failed, };
fn publishFact( self: *Usage, fact: usage_report.GenerationFact, ) PublicationResult { self.publication_mutex.lockUncancelable(io_mod.getIo()); defer self.publication_mutex.unlock(io_mod.getIo()); const sink = self.publication_sink orelse return .not_configured; sink.publish(sink.context, .{ .generation = fact }) catch |err| { debug_trace.logf( "session", "usage profile publication failed id={s} reason={s}", .{ fact.id, @errorName(err) }, ); return .failed; }; return .durable; }
fn stagePublicationBacklog( self: *Usage, alloc: Allocator, fact: usage_report.GenerationFact, ) !void { self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); try self.stagePublicationBacklogUnlocked(alloc, fact); }
fn stagePublicationBacklogUnlocked( self: *Usage, alloc: Allocator, fact: usage_report.GenerationFact, ) !void { const pending_exists = for (self.pending.items) |pending| { if (std.mem.eql(u8, pending.id, fact.id)) break true; } else false; if (!pending_exists) return; for (self.publication_backlog.items) |existing| { if (!std.mem.eql(u8, existing.id, fact.id)) continue; if (!usage_report.GenerationFact.eql(existing, fact)) { self.billing = .incomplete; self.recordIncidentUnlocked(.incomplete, fact.created_at_ms); self.dirty = true; return error.ConflictingGenerationRecord; } return; } if (self.publication_backlog.items.len == max_publication_backlog) { self.billing = .incomplete; self.recordIncidentUnlocked(.incomplete, fact.created_at_ms); self.dirty = true; return error.UsageCapacityExceeded; } var owned = try fact.dupe(alloc); errdefer owned.deinit(alloc); try self.publication_backlog.append(alloc, owned); self.dirty = true; }
fn snapshotProfilePublications( self: *Usage, alloc: Allocator, ) !ProfilePublicationBatch { self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo());
const observed_at_ms = @max(io_mod.milliTimestamp(), 0); var checkpoint_changed = false; for (self.pending.items) |*pending| { if (pending.observed_at_ms != null) continue; pending.observed_at_ms = observed_at_ms; self.dirty = true; checkpoint_changed = true; }
const pending = try alloc.alloc( usage_report.PendingMarker, self.pending.items.len, ); errdefer mem_utils.free(alloc, pending); var copied_pending: usize = 0; errdefer for (pending[0..copied_pending]) |*marker| marker.deinit(alloc); for (self.pending.items, 0..) |generation, index| { pending[index] = .{ .id = try alloc.dupe(u8, generation.id), .observed_at_ms = generation.observed_at_ms.?, }; copied_pending += 1; }
const incidents = try alloc.dupe( usage_report.Incident, self.incidents[0..self.incident_count], ); errdefer mem_utils.free(alloc, incidents);
const facts = try alloc.alloc( usage_report.GenerationFact, self.publication_backlog.items.len, ); errdefer mem_utils.free(alloc, facts); var copied_facts: usize = 0; errdefer for (facts[0..copied_facts]) |*fact| fact.deinit(alloc); for (self.publication_backlog.items, 0..) |fact, index| { facts[index] = try fact.dupe(alloc); copied_facts += 1; }
return .{ .pending = pending, .incidents = incidents, .facts = facts, .checkpoint_changed = checkpoint_changed, }; }
/// Queues a profile-publication drain so invocation finish paths never run /// the append-only store's full-file parse on the agent worker thread. The /// session checkpoint is already durable before this is called; the drain /// only updates the profile-level usage ledger. Tests keep the legacy /// synchronous flush so assertions stay deterministic. fn scheduleProfilePublicationDrain(self: *Usage) void { if (builtin.is_test or comptime builtin.os.tag == .wasi) { self.flushProfilePublications(); return; } _ = self.publication_drain_epoch.fetchAdd(1, .seq_cst); self.publication_drain_mutex.lockUncancelable(io_mod.getIo()); defer self.publication_drain_mutex.unlock(io_mod.getIo()); if (self.publication_drain_thread) |thread| { if (!self.publication_drain_done.load(.seq_cst)) return; thread.join(); self.publication_drain_thread = null; self.publication_drain_done.store(true, .seq_cst); } self.publication_mutex.lockUncancelable(io_mod.getIo()); const has_sink = self.publication_sink != null; self.publication_mutex.unlock(io_mod.getIo()); if (!has_sink) return; self.publication_drain_cancel.store(false, .seq_cst); self.publication_drain_done.store(false, .seq_cst); self.publication_drain_thread = std.Thread.spawn(.{}, publicationDrainThreadMain, .{self}) catch |err| { self.publication_drain_done.store(true, .seq_cst); debug_trace.logf( "session", "usage profile publication drain start failed reason={s}", .{@errorName(err)}, ); self.flushProfilePublications(); return; }; }
/// Stops the background profile-publication drain, joining any live worker. /// A later schedule starts a fresh worker. fn stopPublicationDrain(self: *Usage) void { if (builtin.is_test or comptime builtin.os.tag == .wasi) return; self.publication_drain_cancel.store(true, .seq_cst); self.publication_drain_mutex.lockUncancelable(io_mod.getIo()); defer self.publication_drain_mutex.unlock(io_mod.getIo()); const thread = self.publication_drain_thread; self.publication_drain_thread = null; if (thread) |handle| handle.join(); self.publication_drain_done.store(true, .seq_cst); self.publication_drain_cancel.store(false, .seq_cst); }
fn flushProfilePublications(self: *Usage) void { self.publication_mutex.lockUncancelable(io_mod.getIo()); const sink = self.publication_sink orelse { self.publication_mutex.unlock(io_mod.getIo()); return; };
var batch = self.snapshotProfilePublications(sink.allocator) catch |err| { self.publication_mutex.unlock(io_mod.getIo()); debug_trace.logf( "session", "usage profile publication copy failed reason={s}", .{@errorName(err)}, ); return; }; defer batch.deinit(sink.allocator);
for (batch.pending) |marker| { sink.publish(sink.context, .{ .pending = marker }) catch |err| { debug_trace.logf( "session", "usage profile pending publication failed id={s} reason={s}", .{ marker.id, @errorName(err) }, ); }; } var published_any = false; for (batch.incidents) |incident| { sink.publish(sink.context, .{ .incident = incident }) catch |err| { debug_trace.logf( "session", "usage profile incident publication failed reason={s}", .{@errorName(err)}, ); continue; }; self.mutex.lockUncancelable(io_mod.getIo()); self.removeIncidentUnlocked(incident); self.mutex.unlock(io_mod.getIo()); published_any = true; }
for (batch.facts) |fact| { sink.publish(sink.context, .{ .generation = fact }) catch |err| { debug_trace.logf( "session", "usage profile backlog retry failed id={s} reason={s}", .{ fact.id, @errorName(err) }, ); continue; }; self.mutex.lockUncancelable(io_mod.getIo()); self.applyGenerationUnlocked( sink.allocator, generationRecordBorrowed(fact), true, ) catch |err| { self.billing = .incomplete; self.recordIncidentUnlocked(.incomplete, fact.created_at_ms); self.removePublicationBacklogUnlocked(sink.allocator, fact.id); self.dirty = true; debug_trace.logf( "session", "usage profile backlog settlement failed id={s} reason={s}", .{ fact.id, @errorName(err) }, ); }; self.mutex.unlock(io_mod.getIo()); published_any = true; } self.publication_mutex.unlock(io_mod.getIo());
if (batch.checkpoint_changed or published_any) { self.checkpoint_mutex.lockUncancelable(io_mod.getIo()); _ = self.persistCheckpointBestEffortLocked(); self.checkpoint_mutex.unlock(io_mod.getIo()); } }
fn removePublicationBacklogUnlocked( self: *Usage, alloc: Allocator, id: []const u8, ) void { const index = for (self.publication_backlog.items, 0..) |fact, fact_index| { if (std.mem.eql(u8, fact.id, id)) break fact_index; } else return; var removed = self.publication_backlog.orderedRemove(index); removed.deinit(alloc); }
fn removeIncidentUnlocked( self: *Usage, incident: usage_report.Incident, ) void { const index = for (self.incidents[0..self.incident_count], 0..) |existing, incident_index| { if (existing.occurred_at_ms == incident.occurred_at_ms and existing.completeness == incident.completeness) { break incident_index; } } else return; if (index + 1 < self.incident_count) { std.mem.copyForwards( usage_report.Incident, self.incidents[index .. self.incident_count - 1], self.incidents[index + 1 .. self.incident_count], ); } self.incident_count -= 1; self.dirty = true; }
fn recordIncidentUnlocked( self: *Usage, completeness: usage_report.Completeness, occurred_at_ms: i64, ) void { if (occurred_at_ms < 0) return; for (self.incidents[0..self.incident_count]) |incident| { if (incident.occurred_at_ms == occurred_at_ms and incident.completeness == completeness) { return; } } if (self.incident_count == max_usage_incidents) { var newest_at_ms = occurred_at_ms; for (self.incidents[0..self.incident_count]) |incident| { newest_at_ms = @max(newest_at_ms, incident.occurred_at_ms); } self.incidents[0] = .{ .occurred_at_ms = newest_at_ms, .completeness = .incomplete, }; self.incident_count = 1; return; } self.incidents[self.incident_count] = .{ .occurred_at_ms = occurred_at_ms, .completeness = completeness, }; self.incident_count += 1; }
fn sortModelsBySequenceUnlocked(self: *Usage) void { var index: usize = 1; while (index < self.models.items.len) : (index += 1) { var current = index; while (current > 0 and self.models.items[current - 1].first_sequence > self.models.items[current].first_sequence) { std.mem.swap( ModelAggregate, &self.models.items[current - 1], &self.models.items[current], ); current -= 1; } } }
pub fn recordCommittedLines(self: *Usage, additions: usize, deletions: usize) !void { const additions_u64 = std.math.cast(u64, additions) orelse return error.UsageOverflow; const deletions_u64 = std.math.cast(u64, deletions) orelse return error.UsageOverflow; self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); self.lines_added = std.math.add(u64, self.lines_added, additions_u64) catch { self.code_complete = false; self.dirty = true; return error.UsageOverflow; }; self.lines_removed = std.math.add(u64, self.lines_removed, deletions_u64) catch { self.code_complete = false; self.dirty = true; return error.UsageOverflow; }; self.dirty = true; }
pub const LiveContextSnapshot = struct { used: u64, complete_cost: ?f64, };
/// Returns the latest provider-reported context usage. This state is /// intentionally runtime-only: a restored billing aggregate cannot prove /// how much context the next model request currently occupies. pub fn liveContextSnapshot(self: *Usage) ?LiveContextSnapshot { self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); return .{ .used = self.latest_context_used orelse return null, .complete_cost = if (self.billing == .complete and std.math.isFinite(self.total_cost)) self.total_cost else null, }; }
fn observeContextUsage(self: *Usage, sequence: u64, provider_usage: types.Usage) void { const used: ?u64 = if (provider_usage.input_tokens) |input| if (provider_usage.output_tokens) |output| std.math.add(u64, input, output) catch null else null else null; self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); if (sequence < self.latest_context_sequence) return; self.latest_context_sequence = sequence; self.latest_context_used = used; }
/// Returns an owned point-in-time snapshot. The caller must call `deinit`. pub fn snapshot(self: *Usage, alloc: Allocator) !Snapshot { self.finishReconciliationIfDone(); return self.snapshotCurrent(alloc); }
fn snapshotCurrent(self: *Usage, alloc: Allocator) !Snapshot { self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo());
const models = try alloc.alloc(ModelAggregate, self.models.items.len); errdefer mem_utils.free(alloc, models); var copied_models: usize = 0; errdefer for (models[0..copied_models]) |*model| model.deinit(alloc); for (self.models.items, 0..) |model, index| { models[index] = try model.dupe(alloc); copied_models += 1; }
const pending = try alloc.alloc(PendingGeneration, self.pending.items.len); errdefer mem_utils.free(alloc, pending); var copied_pending: usize = 0; errdefer for (pending[0..copied_pending]) |*generation| generation.deinit(alloc); for (self.pending.items, 0..) |generation, index| { pending[index] = try generation.dupe(alloc); copied_pending += 1; }
const publication_backlog = try alloc.alloc( usage_report.GenerationFact, self.publication_backlog.items.len, ); errdefer mem_utils.free(alloc, publication_backlog); var copied_publications: usize = 0; errdefer for (publication_backlog[0..copied_publications]) |*fact| { fact.deinit(alloc); }; for (self.publication_backlog.items, 0..) |fact, index| { publication_backlog[index] = try fact.dupe(alloc); copied_publications += 1; }
const incidents = try alloc.dupe( usage_report.Incident, self.incidents[0..self.incident_count], ); errdefer mem_utils.free(alloc, incidents);
const now_ms = io_mod.milliTimestamp(); if (self.active_started_at_ms == 0) self.active_started_at_ms = now_ms; const active_elapsed_ms = if (now_ms > self.active_started_at_ms) std.math.cast(u64, now_ms - self.active_started_at_ms) orelse std.math.maxInt(u64) else 0; const wall_duration_ms = std.math.add( u64, self.wall_duration_ms, active_elapsed_ms, ) catch std.math.maxInt(u64); return .{ .billing = if (self.active_sequence_count > 0) .incomplete else self.billing, .api_duration_complete = self.api_duration_complete and self.active_sequence_count == 0, .wall_duration_complete = self.wall_duration_complete and wall_duration_ms != std.math.maxInt(u64), .code_complete = self.code_complete, .next_sequence = self.next_sequence, .settled_through_sequence = self.settled_through_sequence, .api_duration_ms = self.api_duration_ms, .wall_duration_ms = wall_duration_ms, .total_cost = self.total_cost, .input_tokens = self.input_tokens, .output_tokens = self.output_tokens, .cache_read_tokens = self.cache_read_tokens, .cache_write_tokens = self.cache_write_tokens, .reasoning_tokens = self.reasoning_tokens, .request_count = self.request_count, .billable_web_search_calls = self.billable_web_search_calls, .lines_added = self.lines_added, .lines_removed = self.lines_removed, .models = models, .pending = pending, .publication_backlog = publication_backlog, .incidents = incidents, }; }
/// Returns the current session through the shared usage-report contract. pub fn reportSnapshot( self: *Usage, alloc: Allocator, ) !usage_report.Snapshot { var source_snapshot = try self.snapshot(alloc); defer source_snapshot.deinit(alloc);
const models = try alloc.alloc( usage_report.SessionModelSource, source_snapshot.models.len, ); defer alloc.free(models); for (source_snapshot.models, 0..) |model, index| { models[index] = .{ .model = model.model, .total_cost = model.total_cost, .input_tokens = model.input_tokens, .output_tokens = model.output_tokens, .cache_read_tokens = model.cache_read_tokens, .cache_write_tokens = model.cache_write_tokens, .reasoning_tokens = model.reasoning_tokens, .request_count = model.request_count, }; }
const snapshot_time_ms = @max(io_mod.milliTimestamp(), 0); const wall_duration_ms = std.math.cast(i64, source_snapshot.wall_duration_ms) orelse snapshot_time_ms; return usage_report.buildSessionSnapshot(alloc, .{ .snapshot_time_ms = snapshot_time_ms, .session_started_at_ms = @max(snapshot_time_ms -| wall_duration_ms, 0), .completeness = switch (source_snapshot.billing) { .complete => .complete, .pending => .pending, .incomplete => .incomplete, .legacy => .legacy, }, .total_cost = source_snapshot.total_cost, .input_tokens = source_snapshot.input_tokens, .output_tokens = source_snapshot.output_tokens, .cache_read_tokens = source_snapshot.cache_read_tokens, .cache_write_tokens = source_snapshot.cache_write_tokens, .reasoning_tokens = source_snapshot.reasoning_tokens, .request_count = source_snapshot.request_count, .models = models, .activity = .{ .api_duration_complete = source_snapshot.api_duration_complete, .wall_duration_complete = source_snapshot.wall_duration_complete, .code_complete = source_snapshot.code_complete, .api_duration_ms = source_snapshot.api_duration_ms, .wall_duration_ms = source_snapshot.wall_duration_ms, .lines_added = source_snapshot.lines_added, .lines_removed = source_snapshot.lines_removed, }, }); }
pub fn restore( self: *Usage, alloc: Allocator, source: Snapshot, session_started_at_ms: i64, ) !void { self.stopReconciliation(); self.stopPublicationDrain(); try validateSnapshot(source);
var copied = try dupeSnapshotOwned(alloc, source); errdefer copied.deinit(alloc);
self.mutex.lockUncancelable(io_mod.getIo()); self.clearOwnedUnlocked(alloc); self.billing = copied.billing; self.api_duration_complete = copied.api_duration_complete; self.wall_duration_complete = copied.wall_duration_complete; self.code_complete = copied.code_complete; self.next_sequence = copied.next_sequence; self.settled_through_sequence = copied.settled_through_sequence; self.api_duration_ms = copied.api_duration_ms; // Wall time is derived from durable session creation; restoring a // checkpoint's elapsed value as well would count the same span twice. self.wall_duration_ms = 0; const now_ms = io_mod.milliTimestamp(); if (session_started_at_ms <= 0 or session_started_at_ms > now_ms) { self.wall_duration_complete = false; self.active_started_at_ms = now_ms; } else { self.active_started_at_ms = session_started_at_ms; } self.active_sequence_count = 0; self.latest_context_sequence = 0; self.latest_context_used = null; self.total_cost = copied.total_cost; self.input_tokens = copied.input_tokens; self.output_tokens = copied.output_tokens; self.cache_read_tokens = copied.cache_read_tokens; self.cache_write_tokens = copied.cache_write_tokens; self.reasoning_tokens = copied.reasoning_tokens; self.request_count = copied.request_count; self.billable_web_search_calls = copied.billable_web_search_calls; self.lines_added = copied.lines_added; self.lines_removed = copied.lines_removed; self.models = .fromOwnedSlice(copied.models); self.pending = .fromOwnedSlice(copied.pending); self.publication_backlog = .fromOwnedSlice(copied.publication_backlog); self.incident_count = copied.incidents.len; std.mem.copyForwards( usage_report.Incident, self.incidents[0..self.incident_count], copied.incidents, ); if (copied.incidents.len > 0) alloc.free(copied.incidents); copied = undefined; self.dirty = false; self.mutex.unlock(io_mod.getIo()); self.flushProfilePublications(); }
pub fn isDirty(self: *Usage) bool { self.finishReconciliationIfDone(); self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); return self.dirty; }
pub fn startReconciliation( self: *Usage, alloc: Allocator, api_key: []const u8, ) void { self.startReconciliationWithCredential( alloc, api_key, .{ .provider = .gateway, .credential_identity = null }, false, null, ); }
pub fn startDeferredReconciliation( self: *Usage, alloc: Allocator, reference: stream_provider.DeferredUsageReference, credential: []const u8, ) void { self.startReconciliationWithCredential( alloc, credential, .{ .provider = reference.provider, .credential_identity = reference.credential_identity, }, false, null, ); }
pub fn startHostManagedDeferredReconciliation( self: *Usage, alloc: Allocator, reference: stream_provider.DeferredUsageReference, ) void { self.startReconciliationWithCredential( alloc, null, .{ .provider = reference.provider, .credential_identity = reference.credential_identity, }, false, null, ); }
/// Installs the host's authoritative credential regardless of the prior key. pub fn replaceReconciliationCredential( self: *Usage, alloc: Allocator, api_key: []const u8, ) void { self.startReconciliationWithCredential( alloc, api_key, self.reconciliation_authority orelse .{ .provider = .gateway, .credential_identity = null }, true, null, ); }
pub fn replaceProviderReconciliationCredential( self: *Usage, alloc: Allocator, provider: model_provider.ProviderId, source: types.CredentialSource, account_id: ?[]const u8, credential: []const u8, ) void { const identity = credential_authority.derive(source, account_id) orelse { self.mutex.lockUncancelable(io_mod.getIo()); const has_pending = self.pending.items.len > 0; self.mutex.unlock(io_mod.getIo()); if (!has_pending) { self.startReconciliationWithCredential( alloc, credential, .{ .provider = provider, .credential_identity = null }, true, null, ); return; } self.clearReconciliationCredential(); debug_trace.logf( "session", "usage reconciliation withheld provider={s} source={s} reason=stable_credential_identity_unavailable", .{ @tagName(provider), @tagName(source) }, ); return; }; self.startReconciliationWithCredential( alloc, credential, .{ .provider = provider, .credential_identity = identity, }, true, null, ); }
pub fn replaceHostManagedReconciliationAuthority( self: *Usage, alloc: Allocator, provider: model_provider.ProviderId, ) void { self.startReconciliationWithCredential( alloc, null, .{ .provider = provider, .credential_identity = credential_authority.derive(.host_managed, null), }, true, null, ); }
/// Replaces a producer's key only while that key is still authoritative. pub fn refreshReconciliationCredential( self: *Usage, alloc: Allocator, expected_api_key: []const u8, refreshed_api_key: []const u8, ) void { self.startReconciliationWithCredential( alloc, refreshed_api_key, self.reconciliation_authority orelse .{ .provider = .gateway, .credential_identity = null }, true, expected_api_key, ); }
/// Cancels credential work and rejects stale producer starts until replaced. pub fn clearReconciliationCredential(self: *Usage) void { self.reconciliation_cancel.store(true, .seq_cst); self.reconciliation_mutex.lockUncancelable(io_mod.getIo()); defer self.reconciliation_mutex.unlock(io_mod.getIo()); self.reconciliation_cancel.store(true, .seq_cst); _ = self.reconciliation_work_epoch.fetchAdd(1, .seq_cst); const thread = self.reconciliation_thread; self.reconciliation_thread = null; if (thread) |handle| handle.join(); self.reconciliation_done.store(true, .seq_cst); self.reconciliation_cancel.store(false, .seq_cst); self.reconciliation_key_digest = null; self.reconciliation_authority = null; self.reconciliation_credential_blocked = true; }
fn startReconciliationWithCredential( self: *Usage, alloc: Allocator, credential: ?[]const u8, authority: ReconciliationAuthority, replace_existing: bool, expected_api_key: ?[]const u8, ) void { if (credential) |api_key| if (api_key.len == 0) return; const key_digest = if (credential) |api_key| reconciliationKeyDigest(api_key) else hostManagedReconciliationDigest(); const expected_digest = if (expected_api_key) |expected| reconciliationKeyDigest(expected) else null; self.reconciliation_mutex.lockUncancelable(io_mod.getIo()); defer self.reconciliation_mutex.unlock(io_mod.getIo()); const authoritative_replace = replace_existing and expected_api_key == null; if (self.reconciliation_credential_blocked and !authoritative_replace) { return; } if (authoritative_replace) { self.reconciliation_credential_blocked = false; } const same_key = if (self.reconciliation_key_digest) |active_digest| std.mem.eql(u8, &active_digest, &key_digest) else false; const same_authority = if (self.reconciliation_authority) |active| active.eql(authority) else false; if (self.reconciliation_key_digest) |active_digest| { if (!same_key) { if (!replace_existing) return; if (expected_digest) |expected| { if (!std.mem.eql(u8, &active_digest, &expected)) return; } } } _ = self.reconciliation_work_epoch.fetchAdd(1, .seq_cst); if (self.reconciliation_thread) |thread| { if (same_key and same_authority and !self.reconciliation_done.load(.seq_cst)) return; self.reconciliation_cancel.store(true, .seq_cst); self.reconciliation_thread = null; thread.join(); self.reconciliation_done.store(true, .seq_cst); self.reconciliation_cancel.store(false, .seq_cst); } self.reconciliation_key_digest = key_digest; self.reconciliation_authority = authority; if (builtin.is_test) return;
self.mutex.lockUncancelable(io_mod.getIo()); const still_has_pending = self.pending.items.len > 0; self.mutex.unlock(io_mod.getIo()); if (!still_has_pending) return;
const credential_copy = if (credential) |api_key| alloc.dupe(u8, api_key) catch |err| { debug_trace.logf( "session", "usage reconciliation start failed reason={s}", .{@errorName(err)}, ); return; } else null; self.reconciliation_cancel.store(false, .seq_cst); self.reconciliation_done.store(false, .seq_cst); self.reconciliation_key_digest = key_digest; self.reconciliation_thread = std.Thread.spawn( .{}, reconciliationThreadMain, .{ self, alloc, credential_copy, authority, self.generation_usage_providers }, ) catch |err| { self.reconciliation_done.store(true, .seq_cst); debug_trace.logf( "session", "usage reconciliation start failed reason={s}", .{@errorName(err)}, ); if (credential_copy) |api_key| secret.zeroAndFree(alloc, api_key); return; }; }
pub fn cancelReconciliation(self: *Usage) void { self.stopReconciliation(); }
/// Gives the current worker a bounded drain, then cancels it so one-shot /// hosts can persist prompt results without turning exit into an API wait. pub fn finishReconciliationBeforeShutdown(self: *Usage) void { for (0..shutdown_reconciliation_budget_ms / 5) |_| { if (self.reconciliation_done.load(.seq_cst)) { self.finishReconciliationIfDone(); return; } io_mod.sleep(5 * std.time.ns_per_ms); } self.stopReconciliation(); }
pub fn finishProfilePublicationsBeforeShutdown(self: *Usage) void { self.stopPublicationDrain(); self.flushProfilePublications(); self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); if (self.publication_backlog.items.len > 0) { debug_trace.logf( "session", "usage profile publication incomplete facts={d}", .{self.publication_backlog.items.len}, ); } }
fn finishReconciliationIfDone(self: *Usage) void { if (!self.reconciliation_done.load(.seq_cst)) return; self.reconciliation_mutex.lockUncancelable(io_mod.getIo()); if (!self.reconciliation_done.load(.seq_cst)) { self.reconciliation_mutex.unlock(io_mod.getIo()); return; } const thread = self.reconciliation_thread; self.reconciliation_thread = null; if (thread) |handle| { self.reconciliation_cancel.store(true, .seq_cst); handle.join(); self.reconciliation_cancel.store(false, .seq_cst); } self.reconciliation_done.store(true, .seq_cst); self.reconciliation_mutex.unlock(io_mod.getIo()); }
fn stopReconciliation(self: *Usage) void { self.reconciliation_cancel.store(true, .seq_cst); self.reconciliation_mutex.lockUncancelable(io_mod.getIo()); self.reconciliation_cancel.store(true, .seq_cst); const thread = self.reconciliation_thread; self.reconciliation_thread = null; if (thread) |handle| handle.join(); self.reconciliation_done.store(true, .seq_cst); self.reconciliation_cancel.store(false, .seq_cst); self.reconciliation_mutex.unlock(io_mod.getIo()); }
fn markBillingIncomplete(self: *Usage) void { self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); self.billing = .incomplete; self.recordIncidentUnlocked(.incomplete, @max(io_mod.milliTimestamp(), 0)); self.dirty = true; }
fn rejectPendingGeneration(self: *Usage, alloc: Allocator, id: []const u8) void { self.mutex.lockUncancelable(io_mod.getIo()); const pending_index = for (self.pending.items, 0..) |pending, index| { if (std.mem.eql(u8, pending.id, id)) break index; } else null; if (pending_index) |index| { var removed = self.pending.orderedRemove(index); self.recordIncidentUnlocked( .incomplete, removed.observed_at_ms orelse io_mod.milliTimestamp(), ); removed.deinit(alloc); } self.billing = .incomplete; self.dirty = true; self.mutex.unlock(io_mod.getIo()); self.flushProfilePublications(); }
pub fn markCodeIncomplete(self: *Usage) void { self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); self.code_complete = false; self.dirty = true; }
pub fn markClean(self: *Usage, persisted: Snapshot) void { self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); if (!self.matchesSnapshotUnlocked(persisted)) return; self.dirty = false; }
fn failOverflow(self: *Usage) error{UsageOverflow} { self.billing = .incomplete; self.dirty = true; return error.UsageOverflow; }
fn identifierBytesUnlocked(self: *const Usage) usize { var total: usize = 0; for (self.models.items) |model| total +|= model.model.len; for (self.pending.items) |generation| { total +|= generation.id.len; total +|= generation.origin.len; if (generation.team) |team| total +|= team.len; if (generation.account_id) |account_id| total +|= account_id.len; } for (self.publication_backlog.items) |fact| { total +|= fact.id.len; total +|= fact.model.len; } return total; }
fn matchesSnapshotUnlocked(self: *const Usage, snapshot_value: Snapshot) bool { if (self.active_sequence_count > 0) return false; if (self.billing != snapshot_value.billing or self.api_duration_complete != snapshot_value.api_duration_complete or self.wall_duration_complete != snapshot_value.wall_duration_complete or self.code_complete != snapshot_value.code_complete or self.next_sequence != snapshot_value.next_sequence or self.settled_through_sequence != snapshot_value.settled_through_sequence or self.api_duration_ms != snapshot_value.api_duration_ms or self.total_cost != snapshot_value.total_cost or self.input_tokens != snapshot_value.input_tokens or self.output_tokens != snapshot_value.output_tokens or self.cache_read_tokens != snapshot_value.cache_read_tokens or self.cache_write_tokens != snapshot_value.cache_write_tokens or self.reasoning_tokens != snapshot_value.reasoning_tokens or self.request_count != snapshot_value.request_count or self.billable_web_search_calls != snapshot_value.billable_web_search_calls or self.lines_added != snapshot_value.lines_added or self.lines_removed != snapshot_value.lines_removed or self.models.items.len != snapshot_value.models.len or self.pending.items.len != snapshot_value.pending.len or self.publication_backlog.items.len != snapshot_value.publication_backlog.len or self.incident_count != snapshot_value.incidents.len) { return false; } for (self.models.items, snapshot_value.models) |model, saved| { if (!std.mem.eql(u8, model.model, saved.model) or model.first_sequence != saved.first_sequence or model.total_cost != saved.total_cost or model.input_tokens != saved.input_tokens or model.output_tokens != saved.output_tokens or model.cache_read_tokens != saved.cache_read_tokens or model.cache_write_tokens != saved.cache_write_tokens or model.reasoning_tokens != saved.reasoning_tokens or model.request_count != saved.request_count or model.billable_web_search_calls != saved.billable_web_search_calls) { return false; } } for (self.pending.items, snapshot_value.pending) |generation, saved| { if (!std.mem.eql(u8, generation.id, saved.id) or generation.sequence != saved.sequence or !std.mem.eql(u8, generation.origin, saved.origin) or !optionalStringsEqual(generation.team, saved.team) or !generation.provider.same_authority(saved.provider) or generation.credential_source != saved.credential_source or !optionalCredentialIdentitiesEqual(generation.credential_identity, saved.credential_identity) or !optionalStringsEqual(generation.account_id, saved.account_id) or generation.observed_at_ms != saved.observed_at_ms) { return false; } } for (self.publication_backlog.items, snapshot_value.publication_backlog) |fact, saved| { if (!usage_report.GenerationFact.eql(fact, saved)) return false; } for (self.incidents[0..self.incident_count], snapshot_value.incidents) |incident, saved| { if (incident.occurred_at_ms != saved.occurred_at_ms or incident.completeness != saved.completeness) { return false; } } return true; }
fn clearOwned(self: *Usage, alloc: Allocator) void { self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); self.clearOwnedUnlocked(alloc); }
fn reset(self: *Usage, alloc: Allocator, fresh: bool) void { self.mutex.lockUncancelable(io_mod.getIo()); defer self.mutex.unlock(io_mod.getIo()); if (self.pending.items.len > 0 or self.active_sequence_count > 0) { debug_trace.logf( "session", "usage reset drops live state pending={d} active={d}", .{ self.pending.items.len, self.active_sequence_count }, ); } self.clearOwnedUnlocked(alloc); self.billing = if (fresh) .complete else .legacy; self.api_duration_complete = fresh; self.wall_duration_complete = fresh; self.code_complete = fresh; self.next_sequence = 1; self.settled_through_sequence = 0; self.api_duration_ms = 0; self.wall_duration_ms = 0; self.active_started_at_ms = io_mod.milliTimestamp(); self.active_sequence_count = 0; self.latest_context_sequence = 0; self.latest_context_used = null; self.total_cost = 0; self.input_tokens = 0; self.output_tokens = 0; self.cache_read_tokens = 0; self.cache_write_tokens = 0; self.reasoning_tokens = if (fresh) 0 else null; self.request_count = if (fresh) 0 else null; self.billable_web_search_calls = 0; self.lines_added = 0; self.lines_removed = 0; self.incident_count = 0; self.dirty = false; }
fn clearOwnedUnlocked(self: *Usage, alloc: Allocator) void { for (self.models.items) |*model| model.deinit(alloc); self.models.deinit(alloc); self.models = .empty; for (self.pending.items) |*generation| generation.deinit(alloc); self.pending.deinit(alloc); self.pending = .empty; for (self.publication_backlog.items) |*fact| fact.deinit(alloc); self.publication_backlog.deinit(alloc); self.publication_backlog = .empty; }};
const ReconciliationAuthority = struct { provider: model_provider.ProviderId, credential_identity: ?credential_authority.Identity,
fn eql(self: ReconciliationAuthority, other: ReconciliationAuthority) bool { return self.provider.same_authority(other.provider) and optionalCredentialIdentitiesEqual(self.credential_identity, other.credential_identity); }};
pub fn validateSnapshot(snapshot: Snapshot) !void { _ = try validateSnapshotContract(snapshot, false);}
fn validateSnapshotContract(snapshot: Snapshot, allow_legacy_cache: bool) !bool { var cache_totals_valid = true; if (snapshot.next_sequence == 0) return error.InvalidUsageSnapshot; if (snapshot.settled_through_sequence >= snapshot.next_sequence) { return error.InvalidUsageSnapshot; } if (!std.math.isFinite(snapshot.total_cost) or snapshot.total_cost < 0) { return error.InvalidUsageSnapshot; } if (snapshot.models.len > max_models or snapshot.pending.len > max_pending_generations or snapshot.publication_backlog.len > max_publication_backlog or snapshot.incidents.len > max_usage_incidents) { return error.UsageCapacityExceeded; } if (snapshot.billing == .complete and snapshot.pending.len != 0) { return error.InvalidUsageSnapshot; } if (snapshot.billing == .pending and snapshot.pending.len == 0) { return error.InvalidUsageSnapshot; } var identifier_bytes: usize = 0; var input_tokens: u64 = 0; var output_tokens: u64 = 0; var cache_read_tokens: u64 = 0; var cache_write_tokens: u64 = 0; var reasoning_tokens: ?u64 = if (snapshot.reasoning_tokens == null) null else 0; var request_count: ?u64 = if (snapshot.request_count == null) null else 0; var billable_web_search_calls: u64 = 0; var total_cost: f64 = 0; for (snapshot.models, 0..) |model, index| { try validateModel(model.model); if (model.first_sequence == 0 or model.first_sequence >= snapshot.next_sequence) { return error.InvalidUsageSnapshot; } if (!std.math.isFinite(model.total_cost) or model.total_cost < 0) { return error.InvalidUsageSnapshot; } total_cost += model.total_cost; if (!std.math.isFinite(total_cost)) return error.InvalidUsageSnapshot; identifier_bytes = std.math.add(usize, identifier_bytes, model.model.len) catch return error.UsageCapacityExceeded; input_tokens = std.math.add(u64, input_tokens, model.input_tokens) catch return error.InvalidUsageSnapshot; output_tokens = std.math.add(u64, output_tokens, model.output_tokens) catch return error.InvalidUsageSnapshot; cache_read_tokens = std.math.add(u64, cache_read_tokens, model.cache_read_tokens) catch return error.InvalidUsageSnapshot; cache_write_tokens = std.math.add(u64, cache_write_tokens, model.cache_write_tokens) catch return error.InvalidUsageSnapshot; reasoning_tokens = addOptionalCounter( reasoning_tokens, model.reasoning_tokens, ) catch return error.InvalidUsageSnapshot; request_count = addOptionalCounter( request_count, model.request_count, ) catch return error.InvalidUsageSnapshot; if (model.cache_read_tokens > model.input_tokens or model.cache_write_tokens > model.input_tokens) { if (!allow_legacy_cache) return error.InvalidUsageSnapshot; cache_totals_valid = false; } if (model.reasoning_tokens) |reasoning| { if (reasoning > model.output_tokens) return error.InvalidUsageSnapshot; } billable_web_search_calls = std.math.add( u64, billable_web_search_calls, model.billable_web_search_calls, ) catch return error.InvalidUsageSnapshot; for (snapshot.models[0..index]) |prior| { if (std.mem.eql(u8, prior.model, model.model)) return error.InvalidUsageSnapshot; } if (index > 0 and snapshot.models[index - 1].first_sequence >= model.first_sequence) { return error.InvalidUsageSnapshot; } } if (input_tokens != snapshot.input_tokens or output_tokens != snapshot.output_tokens or cache_read_tokens != snapshot.cache_read_tokens or cache_write_tokens != snapshot.cache_write_tokens or reasoning_tokens != snapshot.reasoning_tokens or request_count != snapshot.request_count or billable_web_search_calls != snapshot.billable_web_search_calls) { return error.InvalidUsageSnapshot; } const cost_tolerance = @max(1e-12, snapshot.total_cost * 1e-12); if (@abs(total_cost - snapshot.total_cost) > cost_tolerance) { return error.InvalidUsageSnapshot; } for (snapshot.pending, 0..) |generation, index| { try validateGenerationId(generation.id); try validateOrigin(generation.origin); if (generation.team) |team| try validateTeam(team); if (generation.account_id) |account_id| try validateIdentifier(account_id); if (generation.credential_identity != null and generation.credential_source == null) { return error.InvalidUsageSnapshot; } if (generation.credential_source) |source| { if (!model_provider.authorizesCredential(generation.provider, source)) { return error.InvalidUsageSnapshot; } } if (generation.sequence == 0 or generation.sequence >= snapshot.next_sequence) { return error.InvalidUsageSnapshot; } if (generation.observed_at_ms) |observed_at_ms| { if (observed_at_ms < 0) return error.InvalidUsageSnapshot; } identifier_bytes = std.math.add(usize, identifier_bytes, generation.id.len) catch return error.UsageCapacityExceeded; identifier_bytes = std.math.add(usize, identifier_bytes, generation.origin.len) catch return error.UsageCapacityExceeded; if (generation.team) |team| { identifier_bytes = std.math.add(usize, identifier_bytes, team.len) catch return error.UsageCapacityExceeded; } if (generation.account_id) |account_id| { identifier_bytes = std.math.add(usize, identifier_bytes, account_id.len) catch return error.UsageOverflow; } for (snapshot.pending[0..index]) |prior| { if (std.mem.eql(u8, prior.id, generation.id) or prior.sequence == generation.sequence) { return error.InvalidUsageSnapshot; } } } for (snapshot.publication_backlog, 0..) |fact, index| { usage_report.validateFact(fact) catch return error.InvalidUsageSnapshot; identifier_bytes = std.math.add(usize, identifier_bytes, fact.id.len) catch return error.UsageCapacityExceeded; identifier_bytes = std.math.add(usize, identifier_bytes, fact.model.len) catch return error.UsageCapacityExceeded; for (snapshot.publication_backlog[0..index]) |prior| { if (!std.mem.eql(u8, prior.id, fact.id)) continue; if (!usage_report.GenerationFact.eql(prior, fact)) { return error.InvalidUsageSnapshot; } return error.InvalidUsageSnapshot; } } for (snapshot.incidents) |incident| { if (incident.occurred_at_ms < 0 or (incident.completeness != .pending and incident.completeness != .incomplete)) { return error.InvalidUsageSnapshot; } } if (identifier_bytes > max_identifier_bytes) return error.UsageCapacityExceeded; return cache_totals_valid;}
pub fn appendIncidentOwned( alloc: Allocator, snapshot: *Snapshot, incident: usage_report.Incident,) !void { if (incident.occurred_at_ms < 0 or (incident.completeness != .pending and incident.completeness != .incomplete)) { return error.InvalidUsageSnapshot; } for (snapshot.incidents) |existing| { if (existing.occurred_at_ms == incident.occurred_at_ms and existing.completeness == incident.completeness) { return; } }
if (snapshot.incidents.len == max_usage_incidents) { var newest_at_ms = incident.occurred_at_ms; for (snapshot.incidents) |existing| { newest_at_ms = @max(newest_at_ms, existing.occurred_at_ms); } const next = try alloc.alloc(usage_report.Incident, 1); next[0] = .{ .occurred_at_ms = newest_at_ms, .completeness = .incomplete, }; alloc.free(snapshot.incidents); snapshot.incidents = next; return; }
const next = try alloc.alloc( usage_report.Incident, snapshot.incidents.len + 1, ); if (snapshot.incidents.len > 0) { std.mem.copyForwards( usage_report.Incident, next[0..snapshot.incidents.len], snapshot.incidents, ); } next[snapshot.incidents.len] = incident; if (snapshot.incidents.len > 0) alloc.free(snapshot.incidents); snapshot.incidents = next;}
/// Returns whether two snapshots have the same rollback-readable usage shape.pub fn rollbackProjectionEql(first: Snapshot, second: Snapshot) bool { return billingProjectionEql(first, second) and first.api_duration_complete == second.api_duration_complete and first.wall_duration_complete == second.wall_duration_complete and first.code_complete == second.code_complete and first.api_duration_ms == second.api_duration_ms and first.wall_duration_ms == second.wall_duration_ms and first.lines_added == second.lines_added and first.lines_removed == second.lines_removed;}
/// Returns whether two current-schema snapshots contain the same usage state.pub fn snapshotEql(first: Snapshot, second: Snapshot) bool { if (!rollbackProjectionEql(first, second) or first.reasoning_tokens != second.reasoning_tokens or first.request_count != second.request_count or first.publication_backlog.len != second.publication_backlog.len or first.incidents.len != second.incidents.len) { return false; } for (first.models, second.models) |left, right| { if (left.reasoning_tokens != right.reasoning_tokens or left.request_count != right.request_count) { return false; } } for (first.pending, second.pending) |left, right| { if (left.observed_at_ms != right.observed_at_ms) return false; } for (first.publication_backlog, second.publication_backlog) |left, right| { if (!usage_report.GenerationFact.eql(left, right)) return false; } for (first.incidents, second.incidents) |left, right| { if (left.occurred_at_ms != right.occurred_at_ms or left.completeness != right.completeness) { return false; } } return true;}
/// Returns whether two snapshots have the same billing-relevant projection.pub fn billingProjectionEql(first: Snapshot, second: Snapshot) bool { if (first.billing != second.billing or first.next_sequence != second.next_sequence or first.settled_through_sequence != second.settled_through_sequence or first.total_cost != second.total_cost or first.input_tokens != second.input_tokens or first.output_tokens != second.output_tokens or first.cache_read_tokens != second.cache_read_tokens or first.cache_write_tokens != second.cache_write_tokens or first.billable_web_search_calls != second.billable_web_search_calls or first.models.len != second.models.len or first.pending.len != second.pending.len) { return false; } for (first.models, second.models) |left, right| { if (!std.mem.eql(u8, left.model, right.model) or left.first_sequence != right.first_sequence or left.total_cost != right.total_cost or left.input_tokens != right.input_tokens or left.output_tokens != right.output_tokens or left.cache_read_tokens != right.cache_read_tokens or left.cache_write_tokens != right.cache_write_tokens or left.billable_web_search_calls != right.billable_web_search_calls) { return false; } } for (first.pending, second.pending) |left, right| { if (!std.mem.eql(u8, left.id, right.id) or left.sequence != right.sequence or !std.mem.eql(u8, left.origin, right.origin) or !optionalStringsEqual(left.team, right.team) or !left.provider.same_authority(right.provider) or left.credential_source != right.credential_source or !optionalCredentialIdentitiesEqual(left.credential_identity, right.credential_identity) or !optionalStringsEqual(left.account_id, right.account_id)) { return false; } } return true;}
pub fn writeSnapshot(writer: *std.Io.Writer, snapshot: Snapshot) !void { try validateSnapshot(snapshot); // Keep the durable session payload in the exact pre-usage-dashboard // shape. Older fx binaries reject unknown snapshot fields instead of // ignoring them; richer metrics and recovery hints live in the validated // session sidecar. try writer.writeAll("{\"billing\":"); try std.json.Stringify.value(@tagName(snapshot.billing), .{}, writer); try writer.print( ",\"api_duration_complete\":{s},\"wall_duration_complete\":{s},\"code_complete\":{s},\"next_sequence\":{d},\"settled_through_sequence\":{d},\"api_duration_ms\":{d},\"wall_duration_ms\":{d},\"total_cost\":{d},\"input_tokens\":{d},\"output_tokens\":{d},\"cache_read_tokens\":{d},\"cache_write_tokens\":{d},\"billable_web_search_calls\":{d},\"lines_added\":{d},\"lines_removed\":{d},\"models\":[", .{ if (snapshot.api_duration_complete) "true" else "false", if (snapshot.wall_duration_complete) "true" else "false", if (snapshot.code_complete) "true" else "false", snapshot.next_sequence, snapshot.settled_through_sequence, snapshot.api_duration_ms, snapshot.wall_duration_ms, snapshot.total_cost, snapshot.input_tokens, snapshot.output_tokens, snapshot.cache_read_tokens, snapshot.cache_write_tokens, snapshot.billable_web_search_calls, snapshot.lines_added, snapshot.lines_removed, }, ); for (snapshot.models, 0..) |model, index| { if (index > 0) try writer.writeByte(','); try writer.writeAll("{\"model\":"); try std.json.Stringify.value(model.model, .{}, writer); try writer.print( ",\"first_sequence\":{d},\"total_cost\":{d},\"input_tokens\":{d},\"output_tokens\":{d},\"cache_read_tokens\":{d},\"cache_write_tokens\":{d},\"billable_web_search_calls\":{d}}}", .{ model.first_sequence, model.total_cost, model.input_tokens, model.output_tokens, model.cache_read_tokens, model.cache_write_tokens, model.billable_web_search_calls, }, ); } try writer.writeAll("],\"pending\":["); for (snapshot.pending, 0..) |pending, index| { if (index > 0) try writer.writeByte(','); try writer.writeAll("{\"id\":"); try std.json.Stringify.value(pending.id, .{}, writer); try writer.print(",\"sequence\":{d},\"provider\":", .{pending.sequence}); try std.json.Stringify.value(@tagName(pending.provider), .{}, writer); try writer.writeAll(",\"origin\":"); try std.json.Stringify.value(pending.origin, .{}, writer); try writer.writeAll(",\"team\":"); if (pending.team) |team| { try std.json.Stringify.value(team, .{}, writer); } else { try writer.writeAll("null"); } try writePendingAuthority(writer, pending); try writer.writeByte('}'); } try writer.writeAll("]}");}
/// Writes the current usage schema for storage outside the rollback-readable/// session event stream.pub fn writeRichSnapshot(writer: *std.Io.Writer, snapshot: Snapshot) !void { try validateSnapshot(snapshot); try writer.writeAll("{\"schema_version\":3,\"billing\":"); try std.json.Stringify.value(@tagName(snapshot.billing), .{}, writer); try writer.print( ",\"api_duration_complete\":{s},\"wall_duration_complete\":{s},\"code_complete\":{s},\"next_sequence\":{d},\"settled_through_sequence\":{d},\"api_duration_ms\":{d},\"wall_duration_ms\":{d},\"total_cost\":{d},\"input_tokens\":{d},\"output_tokens\":{d},\"cache_read_tokens\":{d},\"cache_write_tokens\":{d},\"reasoning_tokens\":", .{ if (snapshot.api_duration_complete) "true" else "false", if (snapshot.wall_duration_complete) "true" else "false", if (snapshot.code_complete) "true" else "false", snapshot.next_sequence, snapshot.settled_through_sequence, snapshot.api_duration_ms, snapshot.wall_duration_ms, snapshot.total_cost, snapshot.input_tokens, snapshot.output_tokens, snapshot.cache_read_tokens, snapshot.cache_write_tokens, }, ); try writeOptionalU64(writer, snapshot.reasoning_tokens); try writer.writeAll(",\"request_count\":"); try writeOptionalU64(writer, snapshot.request_count); try writer.print( ",\"billable_web_search_calls\":{d},\"lines_added\":{d},\"lines_removed\":{d},\"models\":[", .{ snapshot.billable_web_search_calls, snapshot.lines_added, snapshot.lines_removed, }, ); for (snapshot.models, 0..) |model, index| { if (index > 0) try writer.writeByte(','); try writer.writeAll("{\"model\":"); try std.json.Stringify.value(model.model, .{}, writer); try writer.print( ",\"first_sequence\":{d},\"total_cost\":{d},\"input_tokens\":{d},\"output_tokens\":{d},\"cache_read_tokens\":{d},\"cache_write_tokens\":{d},\"reasoning_tokens\":", .{ model.first_sequence, model.total_cost, model.input_tokens, model.output_tokens, model.cache_read_tokens, model.cache_write_tokens, }, ); try writeOptionalU64(writer, model.reasoning_tokens); try writer.writeAll(",\"request_count\":"); try writeOptionalU64(writer, model.request_count); try writer.print( ",\"billable_web_search_calls\":{d}}}", .{model.billable_web_search_calls}, ); } try writer.writeAll("],\"pending\":["); for (snapshot.pending, 0..) |pending, index| { if (index > 0) try writer.writeByte(','); try writer.writeAll("{\"id\":"); try std.json.Stringify.value(pending.id, .{}, writer); try writer.print(",\"sequence\":{d},\"provider\":", .{pending.sequence}); try std.json.Stringify.value(@tagName(pending.provider), .{}, writer); try writer.writeAll(",\"origin\":"); try std.json.Stringify.value(pending.origin, .{}, writer); try writer.writeAll(",\"team\":"); if (pending.team) |team| { try std.json.Stringify.value(team, .{}, writer); } else { try writer.writeAll("null"); } try writePendingAuthority(writer, pending); try writer.writeAll(",\"observed_at_ms\":"); if (pending.observed_at_ms) |observed_at_ms| { try writer.print("{d}", .{observed_at_ms}); } else { try writer.writeAll("null"); } try writer.writeByte('}'); } try writer.writeAll("],\"publication_backlog\":["); for (snapshot.publication_backlog, 0..) |fact, index| { if (index > 0) try writer.writeByte(','); try generation_fact_codec.write(writer, fact); } try writer.writeAll("],\"incidents\":["); for (snapshot.incidents, 0..) |incident, index| { if (index > 0) try writer.writeByte(','); try writer.print( "{{\"occurred_at_ms\":{d},\"completeness\":", .{incident.occurred_at_ms}, ); try std.json.Stringify.value(@tagName(incident.completeness), .{}, writer); try writer.writeByte('}'); } try writer.writeAll("]}");}
fn writePendingAuthority(writer: *std.Io.Writer, pending: PendingGeneration) !void { try writer.writeAll(",\"credential_source\":"); if (pending.credential_source) |source| { try std.json.Stringify.value(@tagName(source), .{}, writer); } else { try writer.writeAll("null"); } try writer.writeAll(",\"credential_identity\":"); if (pending.credential_identity) |identity| { const hex = std.fmt.bytesToHex(identity.bytes, .lower); try std.json.Stringify.value(&hex, .{}, writer); } else { try writer.writeAll("null"); } try writer.writeAll(",\"account_id\":"); if (pending.account_id) |account_id| { try std.json.Stringify.value(account_id, .{}, writer); } else { try writer.writeAll("null"); }}
/// Parses either the rollback-readable or current usage snapshot schema.pub fn parseSnapshotValue(alloc: Allocator, value: std.json.Value) !Snapshot { var snapshot = try parseSnapshotFields(alloc, value); errdefer snapshot.deinit(alloc); try validateSnapshot(snapshot); return snapshot;}
/// Caller owns the result. Only historical separate-cache accounting may/// become unavailable; malformed or versioned snapshots remain errors.pub fn parseLegacySnapshotValue(alloc: Allocator, value: std.json.Value) !Snapshot { var snapshot = try parseSnapshotFields(alloc, value); errdefer snapshot.deinit(alloc); if (try validateSnapshotContract(snapshot, isUnversionedSnapshot(value))) return snapshot;
snapshot.deinit(alloc); return .{ .billing = .legacy, .api_duration_complete = false, .wall_duration_complete = false, .code_complete = false, .next_sequence = 1, .settled_through_sequence = 0, .api_duration_ms = 0, .wall_duration_ms = 0, .total_cost = 0, .input_tokens = 0, .output_tokens = 0, .cache_read_tokens = 0, .cache_write_tokens = 0, .billable_web_search_calls = 0, .lines_added = 0, .lines_removed = 0, .models = &.{}, .pending = &.{}, };}
fn isUnversionedSnapshot(value: std.json.Value) bool { return value == .object and value.object.count() == 18;}
pub fn supports_snapshot_schema(schema_version: u64) bool { return schema_version == 2 or schema_version == 3;}
fn parseSnapshotFields(alloc: Allocator, value: std.json.Value) !Snapshot { if (value != .object) { return error.InvalidUsageSnapshot; } const legacy = isUnversionedSnapshot(value); const schema_version = if (legacy) @as(u64, 1) else try parseNonNegativeInteger(value.object.get("schema_version")); if (!legacy) { if (value.object.count() != 23 or !supports_snapshot_schema(schema_version)) { return error.InvalidUsageSnapshot; } } const billing_value = value.object.get("billing") orelse return error.InvalidUsageSnapshot; if (billing_value != .string) return error.InvalidUsageSnapshot; const billing = std.meta.stringToEnum(Availability, billing_value.string) orelse return error.InvalidUsageSnapshot; const api_complete = try parseBool(value.object.get("api_duration_complete")); const wall_complete = try parseBool(value.object.get("wall_duration_complete")); const code_complete = try parseBool(value.object.get("code_complete")); const next_sequence = try parseNonNegativeInteger(value.object.get("next_sequence")); const settled_through_sequence = try parseNonNegativeInteger( value.object.get("settled_through_sequence"), ); const api_duration_ms = try parseNonNegativeInteger(value.object.get("api_duration_ms")); const wall_duration_ms = try parseNonNegativeInteger(value.object.get("wall_duration_ms")); const total_cost = try parseNonNegativeNumber(value.object.get("total_cost")); const input_tokens = try parseNonNegativeInteger(value.object.get("input_tokens")); const output_tokens = try parseNonNegativeInteger(value.object.get("output_tokens")); const cache_read_tokens = try parseNonNegativeInteger(value.object.get("cache_read_tokens")); const cache_write_tokens = try parseNonNegativeInteger(value.object.get("cache_write_tokens")); const reasoning_tokens = if (legacy) null else blk: { const field = value.object.get("reasoning_tokens") orelse return error.InvalidUsageSnapshot; break :blk try parseOptionalNonNegativeInteger(field); }; const request_count = if (legacy) null else blk: { const field = value.object.get("request_count") orelse return error.InvalidUsageSnapshot; break :blk try parseOptionalNonNegativeInteger(field); }; const billable_web_search_calls = try parseNonNegativeInteger( value.object.get("billable_web_search_calls"), ); const lines_added = try parseNonNegativeInteger(value.object.get("lines_added")); const lines_removed = try parseNonNegativeInteger(value.object.get("lines_removed")); const models_value = value.object.get("models") orelse return error.InvalidUsageSnapshot; const pending_value = value.object.get("pending") orelse return error.InvalidUsageSnapshot; if (models_value != .array or pending_value != .array) return error.InvalidUsageSnapshot; if (models_value.array.items.len > max_models or pending_value.array.items.len > max_pending_generations) { return error.UsageCapacityExceeded; }
const models = try alloc.alloc(ModelAggregate, models_value.array.items.len); errdefer mem_utils.free(alloc, models); var model_count: usize = 0; errdefer for (models[0..model_count]) |*model| model.deinit(alloc); for (models_value.array.items, 0..) |model_value, index| { const expected_model_fields: usize = if (legacy) 8 else 10; if (model_value != .object or model_value.object.count() != expected_model_fields) { return error.InvalidUsageSnapshot; } const model_name = model_value.object.get("model") orelse return error.InvalidUsageSnapshot; if (model_name != .string) return error.InvalidUsageSnapshot; const first_sequence = try parseNonNegativeInteger(model_value.object.get("first_sequence")); const model_total_cost = try parseNonNegativeNumber(model_value.object.get("total_cost")); const model_input_tokens = try parseNonNegativeInteger(model_value.object.get("input_tokens")); const model_output_tokens = try parseNonNegativeInteger(model_value.object.get("output_tokens")); const model_cache_read_tokens = try parseNonNegativeInteger(model_value.object.get("cache_read_tokens")); const model_cache_write_tokens = try parseNonNegativeInteger(model_value.object.get("cache_write_tokens")); const model_reasoning_tokens = if (legacy) null else blk: { const field = model_value.object.get("reasoning_tokens") orelse return error.InvalidUsageSnapshot; break :blk try parseOptionalNonNegativeInteger(field); }; const model_request_count = if (legacy) null else blk: { const field = model_value.object.get("request_count") orelse return error.InvalidUsageSnapshot; break :blk try parseOptionalNonNegativeInteger(field); }; const model_web_search_calls = try parseNonNegativeInteger( model_value.object.get("billable_web_search_calls"), ); models[index] = .{ .model = try alloc.dupe(u8, model_name.string), .first_sequence = first_sequence, .total_cost = model_total_cost, .input_tokens = model_input_tokens, .output_tokens = model_output_tokens, .cache_read_tokens = model_cache_read_tokens, .cache_write_tokens = model_cache_write_tokens, .reasoning_tokens = model_reasoning_tokens, .request_count = model_request_count, .billable_web_search_calls = model_web_search_calls, }; model_count += 1; }
const pending = try alloc.alloc(PendingGeneration, pending_value.array.items.len); errdefer mem_utils.free(alloc, pending); var pending_count: usize = 0; errdefer for (pending[0..pending_count]) |*generation| generation.deinit(alloc); for (pending_value.array.items, 0..) |pending_entry, index| { if (pending_entry != .object) return error.InvalidUsageSnapshot; const provider_scoped = pending_entry.object.contains("provider"); const has_observed_at = pending_entry.object.contains("observed_at_ms"); const expected_pending_fields: usize = if (provider_scoped) if (has_observed_at) 9 else 8 else if (has_observed_at) 5 else 4; if (pending_entry != .object or pending_entry.object.count() != expected_pending_fields) { return error.InvalidUsageSnapshot; } const id_value = pending_entry.object.get("id") orelse return error.InvalidUsageSnapshot; const origin_value = pending_entry.object.get("origin") orelse return error.InvalidUsageSnapshot; const team_value = pending_entry.object.get("team") orelse return error.InvalidUsageSnapshot; if (id_value != .string or origin_value != .string) return error.InvalidUsageSnapshot; const sequence = try parseNonNegativeInteger(pending_entry.object.get("sequence")); const observed_at_ms = if (!has_observed_at) null else blk: { const field = pending_entry.object.get("observed_at_ms") orelse return error.InvalidUsageSnapshot; break :blk try parseOptionalNonNegativeI64(field); }; if (team_value != .null and team_value != .string) return error.InvalidUsageSnapshot; const id = try alloc.dupe(u8, id_value.string); errdefer mem_utils.free(alloc, id); const origin = try alloc.dupe(u8, origin_value.string); errdefer mem_utils.free(alloc, origin); const team = switch (team_value) { .null => null, .string => |text| try alloc.dupe(u8, text), else => unreachable, }; errdefer if (team) |text| mem_utils.free(alloc, text); const provider = if (provider_scoped) provider: { const field = pending_entry.object.get("provider") orelse return error.InvalidUsageSnapshot; if (field != .string) return error.InvalidUsageSnapshot; break :provider model_provider.parse(field.string) orelse return error.InvalidUsageSnapshot; } else .gateway; const credential_source = if (provider_scoped) try parseCredentialSourceOptional(pending_entry.object.get("credential_source")) else null; const credential_identity = if (provider_scoped) try parseCredentialIdentityOptional(pending_entry.object.get("credential_identity")) else null; const account_id = if (provider_scoped) try parseOptionalOwnedString(alloc, pending_entry.object.get("account_id")) else null; errdefer if (account_id) |text| mem_utils.free(alloc, text); pending[index] = .{ .id = id, .sequence = sequence, .provider = provider, .origin = origin, .team = team, .credential_source = credential_source, .credential_identity = credential_identity, .account_id = account_id, .observed_at_ms = observed_at_ms, }; pending_count += 1; }
const publication_backlog = if (legacy) blk: { break :blk try alloc.alloc(usage_report.GenerationFact, 0); } else blk: { const backlog_value = value.object.get("publication_backlog") orelse return error.InvalidUsageSnapshot; if (backlog_value != .array or backlog_value.array.items.len > max_publication_backlog) { return error.InvalidUsageSnapshot; } const backlog = try alloc.alloc( usage_report.GenerationFact, backlog_value.array.items.len, ); errdefer mem_utils.free(alloc, backlog); var backlog_count: usize = 0; errdefer for (backlog[0..backlog_count]) |*fact| fact.deinit(alloc); for (backlog_value.array.items, 0..) |fact_value, index| { backlog[index] = generation_fact_codec.parse( alloc, fact_value, ) catch |err| switch (err) { error.OutOfMemory => return error.OutOfMemory, error.InvalidGenerationFact => return error.InvalidUsageSnapshot, }; backlog_count += 1; } break :blk backlog; }; errdefer { for (publication_backlog) |*fact| fact.deinit(alloc); if (publication_backlog.len > 0) alloc.free(publication_backlog); }
const incidents = if (legacy) blk: { break :blk try alloc.alloc(usage_report.Incident, 0); } else blk: { const incidents_value = value.object.get("incidents") orelse return error.InvalidUsageSnapshot; if (incidents_value != .array or incidents_value.array.items.len > max_usage_incidents) { return error.InvalidUsageSnapshot; } const owned_incidents = try alloc.alloc( usage_report.Incident, incidents_value.array.items.len, ); errdefer mem_utils.free(alloc, owned_incidents); for (incidents_value.array.items, 0..) |incident_value, index| { if (incident_value != .object or incident_value.object.count() != 2) { return error.InvalidUsageSnapshot; } const completeness_value = incident_value.object.get("completeness") orelse return error.InvalidUsageSnapshot; if (completeness_value != .string) return error.InvalidUsageSnapshot; const completeness = std.meta.stringToEnum( usage_report.Completeness, completeness_value.string, ) orelse return error.InvalidUsageSnapshot; owned_incidents[index] = .{ .occurred_at_ms = try parseNonNegativeI64( incident_value.object.get("occurred_at_ms"), ), .completeness = completeness, }; } break :blk owned_incidents; }; errdefer if (incidents.len > 0) mem_utils.free(alloc, incidents);
const snapshot = Snapshot{ .billing = billing, .api_duration_complete = api_complete, .wall_duration_complete = wall_complete, .code_complete = code_complete, .next_sequence = next_sequence, .settled_through_sequence = settled_through_sequence, .api_duration_ms = api_duration_ms, .wall_duration_ms = wall_duration_ms, .total_cost = total_cost, .input_tokens = input_tokens, .output_tokens = output_tokens, .cache_read_tokens = cache_read_tokens, .cache_write_tokens = cache_write_tokens, .reasoning_tokens = reasoning_tokens, .request_count = request_count, .billable_web_search_calls = billable_web_search_calls, .lines_added = lines_added, .lines_removed = lines_removed, .models = models, .pending = pending, .publication_backlog = publication_backlog, .incidents = incidents, }; return snapshot;}
fn parseBool(value: ?std.json.Value) !bool { const actual = value orelse return error.InvalidUsageSnapshot; if (actual != .bool) return error.InvalidUsageSnapshot; return actual.bool;}
fn writeOptionalU64(writer: *std.Io.Writer, value: ?u64) !void { if (value) |number| { try writer.print("{d}", .{number}); } else { try writer.writeAll("null"); }}
fn reconciliationThreadMain( usage: *Usage, alloc: Allocator, credential: ?[]u8, authority: ReconciliationAuthority, providers: generation_usage.Set,) void { defer if (credential) |api_key| secret.zeroAndFree(alloc, api_key); defer usage.reconciliation_done.store(true, .seq_cst); var observed_epoch = usage.reconciliation_work_epoch.load(.seq_cst); while (!usage.reconciliation_cancel.load(.seq_cst)) { reconcilePendingBlocking( usage, alloc, credential, &usage.reconciliation_cancel, authority, providers, 30, ); if (usage.reconciliation_cancel.load(.seq_cst)) return; const current_epoch = usage.reconciliation_work_epoch.load(.seq_cst); if (current_epoch != observed_epoch) { observed_epoch = current_epoch; continue; } usage.reconciliation_done.store(true, .seq_cst); if (usage.reconciliation_cancel.load(.seq_cst)) return; const confirmed_epoch = usage.reconciliation_work_epoch.load(.seq_cst); if (confirmed_epoch == observed_epoch) return; usage.reconciliation_done.store(false, .seq_cst); observed_epoch = confirmed_epoch; }}
fn publicationDrainThreadMain(usage: *Usage) void { defer usage.publication_drain_done.store(true, .seq_cst); var observed_epoch = usage.publication_drain_epoch.load(.seq_cst); while (!usage.publication_drain_cancel.load(.seq_cst)) { usage.flushProfilePublications(); if (usage.publication_drain_cancel.load(.seq_cst)) return; const current_epoch = usage.publication_drain_epoch.load(.seq_cst); if (current_epoch != observed_epoch) { observed_epoch = current_epoch; continue; } usage.publication_drain_done.store(true, .seq_cst); if (usage.publication_drain_cancel.load(.seq_cst)) return; const confirmed_epoch = usage.publication_drain_epoch.load(.seq_cst); if (confirmed_epoch == observed_epoch) return; usage.publication_drain_done.store(false, .seq_cst); observed_epoch = confirmed_epoch; }}
fn reconciliationKeyDigest(api_key: []const u8) [Sha256.digest_length]u8 { var digest: [Sha256.digest_length]u8 = undefined; Sha256.hash(api_key, &digest, .{}); return digest;}
fn hostManagedReconciliationDigest() [Sha256.digest_length]u8 { return reconciliationKeyDigest("fx-host-managed-auth-v1");}
fn reconcilePendingBlocking( usage: *Usage, alloc: Allocator, credential: ?[]const u8, cancel_flag: *std.atomic.Value(bool), authority: ReconciliationAuthority, providers: generation_usage.Set, max_attempts: usize,) void { if (credential) |api_key| if (api_key.len == 0) return; if (cancel_flag.load(.seq_cst)) return; var attempt: usize = 0; while (attempt < max_attempts and !cancel_flag.load(.seq_cst)) : (attempt += 1) { var current = usage.snapshot(alloc) catch |err| { debug_trace.logf( "session", "usage reconciliation snapshot failed reason={s}", .{@errorName(err)}, ); usage.markBillingIncomplete(); return; }; defer current.deinit(alloc); if (current.pending.len == 0) return;
var retry_needed = false; for (current.pending) |pending| { if (!pending.provider.same_authority(authority.provider) or !optionalCredentialIdentitiesEqual( pending.credential_identity, authority.credential_identity, )) { continue; } const provider = providers.select(pending.provider) orelse { retry_needed = true; continue; }; var outcome = provider.lookup(alloc, .{ .credential = credential, .tenant = pending.team, .origin = pending.origin, .generation_id = pending.id, .cancel_flag = cancel_flag, }) catch |err| { if (err == error.Cancelled) return; debug_trace.logf( "session", "usage generation lookup failed attempt={d} reason={s}", .{ attempt + 1, @errorName(err) }, ); retry_needed = true; continue; }; defer outcome.deinit(alloc); switch (outcome) { .retry => { debug_trace.logf( "session", "usage generation lookup retry attempt={d}", .{attempt + 1}, ); retry_needed = true; continue; }, .preserve_pending => { debug_trace.logf( "session", "usage generation lookup preserved attempt={d}", .{attempt + 1}, ); continue; }, .reject => { debug_trace.logf( "session", "usage generation lookup rejected attempt={d}", .{attempt + 1}, ); usage.rejectPendingGeneration(alloc, pending.id); }, .found => |record| { usage.applyGeneration(alloc, record) catch |err| { debug_trace.logf( "session", "usage generation apply failed reason={s}", .{@errorName(err)}, ); usage.markBillingIncomplete(); }; }, } } if (!retry_needed) return; if (attempt + 1 < max_attempts and !waitForReconciliationRetry(cancel_flag)) { return; } }}
fn waitForReconciliationRetry(cancel_flag: *std.atomic.Value(bool)) bool { const slices: usize = 200; for (0..slices) |_| { if (cancel_flag.load(.seq_cst)) return false; io_mod.getIo().sleep(.fromMilliseconds(5), .awake) catch |err| { debug_trace.logf( "session", "usage reconciliation retry wait failed reason={s}", .{@errorName(err)}, ); return false; }; } return !cancel_flag.load(.seq_cst);}
fn addRecordToModel(model: *ModelAggregate, record: GenerationRecord, sequence: u64) !void { const total_cost = model.total_cost + record.total_cost; if (!std.math.isFinite(total_cost)) return error.UsageOverflow; const input_tokens = std.math.add(u64, model.input_tokens, record.input_tokens) catch return error.UsageOverflow; const output_tokens = std.math.add(u64, model.output_tokens, record.output_tokens) catch return error.UsageOverflow; const cache_read_tokens = std.math.add(u64, model.cache_read_tokens, record.cache_read_tokens) catch return error.UsageOverflow; const cache_write_tokens = std.math.add(u64, model.cache_write_tokens, record.cache_write_tokens) catch return error.UsageOverflow; const reasoning_tokens = try addOptionalCounter( model.reasoning_tokens, record.reasoning_tokens, ); const request_count = if (model.request_count) |requests| std.math.add(u64, requests, 1) catch return error.UsageOverflow else null; const web_search_calls = std.math.add( u64, model.billable_web_search_calls, record.billable_web_search_calls, ) catch return error.UsageOverflow; model.total_cost = total_cost; model.input_tokens = input_tokens; model.output_tokens = output_tokens; model.cache_read_tokens = cache_read_tokens; model.cache_write_tokens = cache_write_tokens; model.reasoning_tokens = reasoning_tokens; model.request_count = request_count; model.billable_web_search_calls = web_search_calls; model.first_sequence = @min(model.first_sequence, sequence);}
fn validateGenerationRecord(record: GenerationRecord) !void { usage_report.validateFact(generationFactBorrowed(record)) catch { return error.InvalidGenerationRecord; };}
fn generationFactBorrowed(record: GenerationRecord) usage_report.GenerationFact { return .{ .id = @constCast(record.id), .created_at_ms = record.created_at_ms, .model = @constCast(record.model), .input_tokens = record.input_tokens, .output_tokens = record.output_tokens, .cache_read_tokens = record.cache_read_tokens, .cache_write_tokens = record.cache_write_tokens, .reasoning_tokens = record.reasoning_tokens, .billable_web_search_calls = record.billable_web_search_calls, .total_cost = record.total_cost, };}
fn generationRecordBorrowed( fact: usage_report.GenerationFact,) GenerationRecord { return .{ .id = fact.id, .created_at_ms = fact.created_at_ms, .model = fact.model, .total_cost = fact.total_cost, .input_tokens = fact.input_tokens, .output_tokens = fact.output_tokens, .cache_read_tokens = fact.cache_read_tokens, .cache_write_tokens = fact.cache_write_tokens, .reasoning_tokens = fact.reasoning_tokens, .billable_web_search_calls = fact.billable_web_search_calls, };}
fn addOptionalCounter(first: ?u64, second: ?u64) error{UsageOverflow}!?u64 { if (first == null or second == null) return null; return std.math.add(u64, first.?, second.?) catch error.UsageOverflow;}
pub fn dupeSnapshotOwned(alloc: Allocator, source: Snapshot) !Snapshot { const models = try alloc.alloc(ModelAggregate, source.models.len); errdefer mem_utils.free(alloc, models); var copied_models: usize = 0; errdefer for (models[0..copied_models]) |*model| model.deinit(alloc); for (source.models, 0..) |model, index| { models[index] = try model.dupe(alloc); copied_models += 1; }
const pending = try alloc.alloc(PendingGeneration, source.pending.len); errdefer mem_utils.free(alloc, pending); var copied_pending: usize = 0; errdefer for (pending[0..copied_pending]) |*generation| generation.deinit(alloc); for (source.pending, 0..) |generation, index| { pending[index] = try generation.dupe(alloc); copied_pending += 1; }
const publication_backlog = try alloc.alloc( usage_report.GenerationFact, source.publication_backlog.len, ); errdefer mem_utils.free(alloc, publication_backlog); var copied_publications: usize = 0; errdefer for (publication_backlog[0..copied_publications]) |*fact| { fact.deinit(alloc); }; for (source.publication_backlog, 0..) |fact, index| { publication_backlog[index] = try fact.dupe(alloc); copied_publications += 1; }
const incidents = try alloc.dupe(usage_report.Incident, source.incidents); errdefer mem_utils.free(alloc, incidents);
return .{ .billing = source.billing, .api_duration_complete = source.api_duration_complete, .wall_duration_complete = source.wall_duration_complete, .code_complete = source.code_complete, .next_sequence = source.next_sequence, .settled_through_sequence = source.settled_through_sequence, .api_duration_ms = source.api_duration_ms, .wall_duration_ms = source.wall_duration_ms, .total_cost = source.total_cost, .input_tokens = source.input_tokens, .output_tokens = source.output_tokens, .cache_read_tokens = source.cache_read_tokens, .cache_write_tokens = source.cache_write_tokens, .reasoning_tokens = source.reasoning_tokens, .request_count = source.request_count, .billable_web_search_calls = source.billable_web_search_calls, .lines_added = source.lines_added, .lines_removed = source.lines_removed, .models = models, .pending = pending, .publication_backlog = publication_backlog, .incidents = incidents, };}
fn canonicalExactGenerationId( provider: model_provider.ProviderId, external_id: []const u8, buffer: *[30]u8,) ![]const u8 { if (provider == .gateway) { try validateGenerationId(external_id); return external_id; } try validateExternalGenerationId(external_id); var digest: [Sha256.digest_length]u8 = undefined; var hash = Sha256.init(.{}); hash.update(provider.label()); if (provider == .configured) if (provider.configured.binding) |binding| hash.update(&binding); hash.update(&.{0}); hash.update(external_id); hash.final(&digest); const encoded = std.fmt.bytesToHex(digest[0..13].*, .upper); @memcpy(buffer[0..4], "gen_"); @memcpy(buffer[4..], &encoded); return buffer;}
fn exactUsageOrigin(provider: model_provider.ProviderId) []const u8 { return switch (provider) { .gateway => "exact/gateway", .codex => "exact/codex", .grok => "exact/grok", .configured => "exact/configured", };}
test "live context usage keeps the newest completed provider observation" { var usage = Usage.initFresh(); defer usage.deinit(std.testing.allocator);
usage.observeContextUsage(2, .{ .input_tokens = 30, .output_tokens = 7 }); usage.observeContextUsage(1, .{ .input_tokens = 1, .output_tokens = 1 }); const snapshot = usage.liveContextSnapshot().?; try std.testing.expectEqual(@as(u64, 37), snapshot.used); try std.testing.expectEqual(@as(?f64, 0), snapshot.complete_cost);}
test "newer missing context usage clears the prior observation" { var usage = Usage.initFresh(); defer usage.deinit(std.testing.allocator);
usage.observeContextUsage(1, .{ .input_tokens = 30, .output_tokens = 7 }); usage.observeContextUsage(2, .{ .input_tokens = 40 }); try std.testing.expect(usage.liveContextSnapshot() == null);
usage.observeContextUsage(1, .{ .input_tokens = 90, .output_tokens = 10 }); try std.testing.expect(usage.liveContextSnapshot() == null);}
test "direct exact generation IDs are deterministic and provider scoped" { var first_buffer: [30]u8 = undefined; var replay_buffer: [30]u8 = undefined; var other_provider_buffer: [30]u8 = undefined; const first = try canonicalExactGenerationId( .codex, "response-shared-id", &first_buffer, ); const replay = try canonicalExactGenerationId( .codex, "response-shared-id", &replay_buffer, ); const other_provider = try canonicalExactGenerationId( .grok, "response-shared-id", &other_provider_buffer, );
try std.testing.expectEqualStrings(first, replay); try std.testing.expect(!std.mem.eql(u8, first, other_provider)); try std.testing.expect(types.validGatewayGenerationId(first)); try std.testing.expect(types.validGatewayGenerationId(other_provider));
var gateway_buffer: [30]u8 = undefined; const gateway_id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV"; try std.testing.expectEqualStrings( gateway_id, try canonicalExactGenerationId(.gateway, gateway_id, &gateway_buffer), ); try std.testing.expectError( error.InvalidGenerationId, canonicalExactGenerationId(.codex, "", &gateway_buffer), ); try std.testing.expectError( error.InvalidGenerationId, canonicalExactGenerationId(.grok, "response\ninvalid", &gateway_buffer), );}
fn validateGenerationId(id: []const u8) !void { if (!types.validGatewayGenerationId(id)) return error.InvalidGenerationId;}
fn validateExternalGenerationId(id: []const u8) !void { if (id.len == 0 or id.len > max_identifier_bytes or !std.unicode.utf8ValidateSlice(id)) { return error.InvalidGenerationId; } for (id) |byte| if (std.ascii.isControl(byte)) return error.InvalidGenerationId;}
fn validateModel(model: []const u8) !void { if (model.len == 0 or model.len > max_model_bytes) { return error.InvalidModel; } for (model) |char| { if (char < 0x21 or char > 0x7e) return error.InvalidModel; }}
fn validateOrigin(origin: []const u8) !void { if (origin.len == 0 or origin.len > max_origin_bytes) { return error.InvalidGatewayOrigin; } for (origin) |char| { if (char < 0x21 or char > 0x7e) return error.InvalidGatewayOrigin; }}
fn validateTeam(team: []const u8) !void { if (team.len == 0 or team.len > max_team_bytes) { return error.InvalidGatewayTeam; } for (team) |char| { if (char < 0x21 or char > 0x7e) return error.InvalidGatewayTeam; }}
fn validateIdentifier(value: []const u8) !void { if (value.len == 0 or value.len > max_model_bytes or !std.unicode.utf8ValidateSlice(value)) { return error.InvalidCredentialIdentity; } for (value) |byte| if (std.ascii.isControl(byte)) return error.InvalidCredentialIdentity;}
fn parseNonNegativeNumber(value: ?std.json.Value) !f64 { const actual = value orelse return error.InvalidGenerationRecord; const number: f64 = switch (actual) { .integer => |integer| @floatFromInt(integer), .float => |float| float, .number_string => |text| std.fmt.parseFloat(f64, text) catch return error.InvalidGenerationRecord, else => return error.InvalidGenerationRecord, }; if (!std.math.isFinite(number) or number < 0) return error.InvalidGenerationRecord; return number;}
fn parseNonNegativeInteger(value: ?std.json.Value) !u64 { const actual = value orelse return error.InvalidGenerationRecord; return switch (actual) { .integer => |integer| if (integer >= 0) std.math.cast(u64, integer) orelse error.InvalidGenerationRecord else error.InvalidGenerationRecord, .number_string => |text| std.fmt.parseInt(u64, text, 10) catch error.InvalidGenerationRecord, else => error.InvalidGenerationRecord, };}
fn parseOptionalNonNegativeInteger(value: ?std.json.Value) !?u64 { const actual = value orelse return null; if (actual == .null) return null; return try parseNonNegativeInteger(actual);}
fn parseNonNegativeI64(value: ?std.json.Value) !i64 { const parsed = try parseNonNegativeInteger(value); return std.math.cast(i64, parsed) orelse error.InvalidGenerationRecord;}
fn parseOptionalNonNegativeI64(value: ?std.json.Value) !?i64 { const actual = value orelse return null; if (actual == .null) return null; return try parseNonNegativeI64(actual);}
fn parseCredentialSourceOptional(value: ?std.json.Value) !?types.CredentialSource { const actual = value orelse return error.InvalidUsageSnapshot; return switch (actual) { .null => null, .string => |text| types.parseRuntimeCredentialSource(text) orelse return error.InvalidUsageSnapshot, else => error.InvalidUsageSnapshot, };}
fn parseCredentialIdentityOptional(value: ?std.json.Value) !?credential_authority.Identity { const actual = value orelse return error.InvalidUsageSnapshot; return switch (actual) { .null => null, .string => |hex| identity: { var bytes: [Sha256.digest_length]u8 = undefined; _ = std.fmt.hexToBytes(&bytes, hex) catch return error.InvalidUsageSnapshot; const canonical = std.fmt.bytesToHex(bytes, .lower); if (!std.mem.eql(u8, &canonical, hex)) return error.InvalidUsageSnapshot; break :identity credential_authority.Identity{ .bytes = bytes }; }, else => error.InvalidUsageSnapshot, };}
fn parseOptionalOwnedString(alloc: Allocator, value: ?std.json.Value) !?[]u8 { const actual = value orelse return error.InvalidUsageSnapshot; return switch (actual) { .null => null, .string => |text| try alloc.dupe(u8, text), else => error.InvalidUsageSnapshot, };}
fn optionalStringsEqual(first: ?[]const u8, second: ?[]const u8) bool { if (first == null or second == null) return first == null and second == null; return std.mem.eql(u8, first.?, second.?);}
fn optionalCredentialIdentitiesEqual( first: ?credential_authority.Identity, second: ?credential_authority.Identity,) bool { if (first == null or second == null) return first == null and second == null; return first.?.eql(second.?);}
fn testGatewayUsageOutcome( generation_id: []const u8, immediate: bool,) stream_provider.UsageOutcome { return if (immediate) .{ .exact = .gateway } else .{ .deferred = testGatewayUsageReference( generation_id, "https://ai-gateway.vercel.sh", ) };}
fn testGatewayUsageReference( generation_id: []const u8, scope: []const u8,) stream_provider.DeferredUsageReference { return .{ .provider = .gateway, .generation_id = generation_id, .scope = scope, .credential_source = .ai_gateway_api_key, .credential_identity = credential_authority.derive( .ai_gateway_api_key, null, ), };}
test "usage snapshot JSON round trips" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc); var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc);
var encoded: std.Io.Writer.Allocating = .init(alloc); defer encoded.deinit(); try writeSnapshot(&encoded.writer, snapshot); var parsed = try std.json.parseFromSlice(std.json.Value, alloc, encoded.written(), .{}); defer parsed.deinit(); var decoded = try parseSnapshotValue(alloc, parsed.value); defer decoded.deinit(alloc); try std.testing.expectEqual(snapshot.billing, decoded.billing); try std.testing.expectEqual(snapshot.wall_duration_ms, decoded.wall_duration_ms);}
test "host-managed deferred usage authority round trips" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc); const identity = credential_authority.derive(.host_managed, null).?; const observation = try InvocationObservation.begin(&usage); try observation.complete(alloc, .{}, .{ .deferred = .{ .provider = .gateway, .generation_id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .scope = "https://ai-gateway.vercel.sh", .credential_source = .host_managed, .credential_identity = identity, } });
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); var encoded: std.Io.Writer.Allocating = .init(alloc); defer encoded.deinit(); try writeSnapshot(&encoded.writer, snapshot); var parsed = try std.json.parseFromSlice(std.json.Value, alloc, encoded.written(), .{}); defer parsed.deinit(); var decoded = try parseSnapshotValue(alloc, parsed.value); defer decoded.deinit(alloc);
try std.testing.expectEqual(types.CredentialSource.host_managed, decoded.pending[0].credential_source.?); try std.testing.expect(decoded.pending[0].credential_identity.?.eql(identity));}
test "profile recovery hint follows unresolved durable usage" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
var settled = try usage.snapshot(alloc); defer settled.deinit(alloc); try std.testing.expect(!needsProfileRecovery(settled));
const sequence = try usage.reserveInvocation(); var active = try usage.snapshot(alloc); defer active.deinit(alloc); try std.testing.expect(needsProfileRecovery(active));
usage.finishInvocation(sequence, 1, .unbilled); var finished = try usage.snapshot(alloc); defer finished.deinit(alloc); try std.testing.expect(!needsProfileRecovery(finished));}
test "durable incident publication retires profile recovery" { const alloc = std.testing.allocator; const CheckpointProbe = struct { calls: usize = 0, recovery_pending: bool = true,
fn persist(raw: *anyopaque, snapshot: Snapshot) !void { const self: *@This() = @ptrCast(@alignCast(raw)); self.calls += 1; self.recovery_pending = needsProfileRecovery(snapshot); } }; const PublicationProbe = struct { fail: bool = true, incidents: usize = 0,
fn publish( raw: *anyopaque, event: usage_report.ProfileEvent, ) !void { const self: *@This() = @ptrCast(@alignCast(raw)); switch (event) { .incident => {}, .pending, .generation => return, } if (self.fail) return error.InjectedPublicationFailure; self.incidents += 1; } };
var checkpoint = CheckpointProbe{}; var publication = PublicationProbe{}; var usage = Usage.initFresh(); defer usage.deinit(alloc); usage.configureCheckpointSink(.{ .context = &checkpoint, .allocator = alloc, .persist = CheckpointProbe.persist, }); usage.markBillingIncomplete(); usage.configurePublicationSink(.{ .context = &publication, .allocator = alloc, .publish = PublicationProbe.publish, });
var unresolved = try usage.snapshot(alloc); defer unresolved.deinit(alloc); try std.testing.expect(needsProfileRecovery(unresolved)); try std.testing.expectEqual(@as(usize, 1), unresolved.incidents.len);
usage.configurePublicationSink(null); publication.fail = false; usage.configurePublicationSink(.{ .context = &publication, .allocator = alloc, .publish = PublicationProbe.publish, });
var durable = try usage.snapshot(alloc); defer durable.deinit(alloc); try std.testing.expectEqual(Availability.incomplete, durable.billing); try std.testing.expectEqual(@as(usize, 0), durable.incidents.len); try std.testing.expect(!needsProfileRecovery(durable)); try std.testing.expectEqual(@as(usize, 1), publication.incidents); try std.testing.expect(checkpoint.calls > 0); try std.testing.expect(!checkpoint.recovery_pending);}
test "profile publication failure preserves session totals and retries backlog" { const alloc = std.testing.allocator; const PublicationProbe = struct { fail: bool = true, calls: usize = 0,
fn publish( raw: *anyopaque, _: usage_report.ProfileEvent, ) !void { const self: *@This() = @ptrCast(@alignCast(raw)); self.calls += 1; if (self.fail) return error.InjectedPublicationFailure; } };
var probe = PublicationProbe{}; var usage = Usage.initFresh(); defer usage.deinit(alloc); usage.configurePublicationSink(.{ .context = &probe, .allocator = alloc, .publish = PublicationProbe.publish, });
const sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, sequence, 1, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", null, ); try usage.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .created_at_ms = 1000, .model = "provider/model", .total_cost = 0.25, .input_tokens = 10, .output_tokens = 2, .cache_read_tokens = 1, .cache_write_tokens = 0, .reasoning_tokens = 1, .billable_web_search_calls = 0, });
var failed = try usage.snapshot(alloc); defer failed.deinit(alloc); try std.testing.expectEqual(Availability.complete, failed.billing); try std.testing.expectEqual(@as(u64, 12), failed.input_tokens + failed.output_tokens); try std.testing.expectEqual(@as(usize, 1), failed.publication_backlog.len); try std.testing.expectEqual(@as(usize, 2), probe.calls);
probe.fail = false; usage.configurePublicationSink(.{ .context = &probe, .allocator = alloc, .publish = PublicationProbe.publish, }); var retried = try usage.snapshot(alloc); defer retried.deinit(alloc); try std.testing.expectEqual(@as(usize, 0), retried.publication_backlog.len); try std.testing.expectEqual(@as(usize, 3), probe.calls);}
test "restored publication backlog settles the pending generation exactly" { const alloc = std.testing.allocator; const Checkpoint = struct { fn persist(_: *anyopaque, _: Snapshot) !void {} }; const PublicationProbe = struct { fail_generation: bool = true, successful_generations: usize = 0,
fn publish( raw: *anyopaque, event: usage_report.ProfileEvent, ) !void { const self: *@This() = @ptrCast(@alignCast(raw)); switch (event) { .generation => { if (self.fail_generation) return error.InjectedPublicationFailure; self.successful_generations += 1; }, .pending, .incident => {}, } } };
var checkpoint_context: u8 = 0; var probe = PublicationProbe{}; var source = Usage.initFresh(); defer source.deinit(alloc); source.configureCheckpointSink(.{ .context = &checkpoint_context, .allocator = alloc, .persist = Checkpoint.persist, }); source.configurePublicationSink(.{ .context = &probe, .allocator = alloc, .publish = PublicationProbe.publish, });
const sequence = try source.reserveInvocationDurably(); _ = try source.finishObservedInvocationDurably( alloc, sequence, 1, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", null, ); try source.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .created_at_ms = 1000, .model = "provider/model", .total_cost = 0.25, .input_tokens = 10, .output_tokens = 2, .cache_read_tokens = 1, .cache_write_tokens = 0, .reasoning_tokens = 1, .billable_web_search_calls = 2, });
var saved = try source.snapshot(alloc); defer saved.deinit(alloc); try std.testing.expectEqual(Availability.pending, saved.billing); try std.testing.expectEqual(@as(usize, 1), saved.pending.len); try std.testing.expectEqual(@as(usize, 1), saved.publication_backlog.len); try std.testing.expectEqual(@as(u64, 0), saved.input_tokens); source.configurePublicationSink(null);
probe.fail_generation = false; var resumed = Usage.initFresh(); defer resumed.deinit(alloc); resumed.configurePublicationSink(.{ .context = &probe, .allocator = alloc, .publish = PublicationProbe.publish, }); try resumed.restore(alloc, saved, 1);
var settled = try resumed.snapshot(alloc); defer settled.deinit(alloc); try std.testing.expectEqual(Availability.complete, settled.billing); try std.testing.expectEqual(@as(usize, 0), settled.pending.len); try std.testing.expectEqual(@as(usize, 0), settled.publication_backlog.len); try std.testing.expectEqual(@as(u64, 10), settled.input_tokens); try std.testing.expectEqual(@as(u64, 2), settled.output_tokens); try std.testing.expectEqual(@as(u64, 2), settled.billable_web_search_calls); try std.testing.expectEqual(@as(usize, 1), probe.successful_generations);}
test "no-checkpoint usage drains a transient publication failure" { const alloc = std.testing.allocator; const PublicationProbe = struct { generation_failures: usize = 1, successful_generations: usize = 0, published_search_calls: u64 = 0,
fn publish( raw: *anyopaque, event: usage_report.ProfileEvent, ) !void { const self: *@This() = @ptrCast(@alignCast(raw)); switch (event) { .generation => |fact| { if (self.generation_failures > 0) { self.generation_failures -= 1; return error.InjectedPublicationFailure; } self.successful_generations += 1; self.published_search_calls = fact.billable_web_search_calls; }, .pending, .incident => {}, } } };
var probe = PublicationProbe{}; var usage = Usage.initFresh(); defer usage.deinit(alloc); usage.configurePublicationSink(.{ .context = &probe, .allocator = alloc, .publish = PublicationProbe.publish, });
const sequence = try usage.reserveInvocationDurably(); _ = try usage.finishObservedInvocationDurably( alloc, sequence, 1, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", null, ); try usage.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .created_at_ms = 1000, .model = "provider/model", .total_cost = 0.25, .input_tokens = 10, .output_tokens = 2, .cache_read_tokens = 1, .cache_write_tokens = 0, .reasoning_tokens = 1, .billable_web_search_calls = 3, });
var settled = try usage.snapshot(alloc); defer settled.deinit(alloc); try std.testing.expectEqual(Availability.complete, settled.billing); try std.testing.expectEqual(@as(usize, 0), settled.pending.len); try std.testing.expectEqual(@as(usize, 0), settled.publication_backlog.len); try std.testing.expectEqual(@as(usize, 1), probe.successful_generations); try std.testing.expectEqual(@as(u64, 3), probe.published_search_calls);}
test "missing profile publication sink keeps the durable pending bridge" { const alloc = std.testing.allocator; const Checkpoint = struct { calls: usize = 0,
fn persist(raw: *anyopaque, _: Snapshot) !void { const self: *@This() = @ptrCast(@alignCast(raw)); self.calls += 1; } };
var checkpoint = Checkpoint{}; var usage = Usage.initFresh(); defer usage.deinit(alloc); usage.configureCheckpointSink(.{ .context = &checkpoint, .allocator = alloc, .persist = Checkpoint.persist, });
const sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, sequence, 1, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", null, ); try usage.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .created_at_ms = 1000, .model = "provider/model", .total_cost = 0.25, .input_tokens = 10, .output_tokens = 2, .cache_read_tokens = 1, .cache_write_tokens = 0, .reasoning_tokens = 1, .billable_web_search_calls = 0, });
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(Availability.pending, snapshot.billing); try std.testing.expectEqual(@as(usize, 1), snapshot.pending.len); try std.testing.expectEqual(@as(usize, 1), snapshot.publication_backlog.len); try std.testing.expectEqual(@as(u64, 0), snapshot.input_tokens); try std.testing.expectEqualStrings( "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", snapshot.publication_backlog[0].id, ); try std.testing.expectEqual(@as(usize, 1), checkpoint.calls);}
test "session report preserves usage-first totals and session activity" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc); try usage.recordCommittedLines(3, 1); var report = try usage.reportSnapshot(alloc); defer report.deinit(alloc); try std.testing.expectEqual(usage_report.Scope.session, report.scope); try std.testing.expectEqual(@as(u64, 0), report.totals.?.total_tokens); try std.testing.expectEqual(@as(u64, 3), report.session_activity.?.lines_added);}
test "incident overflow retains the newest incomplete boundary" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const newest_at_ms: i64 = 10_000; for (0..max_usage_incidents) |index| { usage.recordIncidentUnlocked( .incomplete, newest_at_ms - @as(i64, @intCast(index)), ); } usage.recordIncidentUnlocked(.incomplete, 100);
try std.testing.expectEqual(@as(usize, 1), usage.incident_count); try std.testing.expectEqual( newest_at_ms, usage.incidents[0].occurred_at_ms, ); try std.testing.expectEqual( usage_report.Completeness.incomplete, usage.incidents[0].completeness, );}
test "usage snapshot parsing releases every partial allocation" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const resolved_sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, resolved_sequence, 1, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", null, ); try usage.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .model = "anthropic/claude-opus-4.8", .total_cost = 0.1, .input_tokens = 10, .output_tokens = 2, .cache_read_tokens = 1, .cache_write_tokens = 0, .billable_web_search_calls = 0, }); const pending_sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, pending_sequence, 1, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAW", "https://ai-gateway.vercel.sh", "team_1", );
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); var encoded: std.Io.Writer.Allocating = .init(alloc); defer encoded.deinit(); try writeSnapshot(&encoded.writer, snapshot); var parsed = try std.json.parseFromSlice( std.json.Value, alloc, encoded.written(), .{}, ); defer parsed.deinit();
var probe = std.testing.FailingAllocator.init(alloc, .{}); var decoded = try parseSnapshotValue(probe.allocator(), parsed.value); decoded.deinit(probe.allocator()); try std.testing.expectEqual(probe.allocated_bytes, probe.freed_bytes);
for (0..probe.alloc_index) |fail_index| { var failing = std.testing.FailingAllocator.init( alloc, .{ .fail_index = fail_index }, ); if (parseSnapshotValue(failing.allocator(), parsed.value)) |value| { var owned = value; owned.deinit(failing.allocator()); } else |err| switch (err) { error.OutOfMemory => {}, else => return err, } try std.testing.expect(failing.has_induced_failure); try std.testing.expectEqual(failing.allocated_bytes, failing.freed_bytes); }}
test "fresh usage aggregates authoritative generations in invocation order" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const first_sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, first_sequence, 125, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", "team_1", ); try usage.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .model = "anthropic/claude-opus-4.8", .total_cost = 0.00123, .input_tokens = 100, .output_tokens = 50, .cache_read_tokens = 20, .cache_write_tokens = 10, .billable_web_search_calls = 1, });
const second_sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, second_sequence, 75, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAW", "https://ai-gateway.vercel.sh", "team_1", ); try usage.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAW", .model = "openai/gpt-5.5", .total_cost = 0.004, .input_tokens = 200, .output_tokens = 25, .cache_read_tokens = 0, .cache_write_tokens = 0, .billable_web_search_calls = 0, });
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(Availability.complete, snapshot.billing); try std.testing.expect(snapshot.api_duration_complete); try std.testing.expectApproxEqAbs(@as(f64, 0.00523), snapshot.total_cost, 1e-12); try std.testing.expectEqual(@as(u64, 200), snapshot.api_duration_ms); try std.testing.expectEqual(@as(u64, 300), snapshot.input_tokens); try std.testing.expectEqual(@as(u64, 75), snapshot.output_tokens); try std.testing.expectEqual(@as(u64, 20), snapshot.cache_read_tokens); try std.testing.expectEqual(@as(u64, 10), snapshot.cache_write_tokens); try std.testing.expectEqual(@as(u64, 1), snapshot.billable_web_search_calls); try std.testing.expectEqual(@as(usize, 2), snapshot.models.len); try std.testing.expectEqualStrings("anthropic/claude-opus-4.8", snapshot.models[0].model); try std.testing.expectEqual(first_sequence, snapshot.models[0].first_sequence); try std.testing.expectEqualStrings("openai/gpt-5.5", snapshot.models[1].model); try std.testing.expectEqual(second_sequence, snapshot.models[1].first_sequence); try std.testing.expectEqual(@as(usize, 0), snapshot.pending.len);
snapshot.input_tokens += 1; try std.testing.expectError(error.InvalidUsageSnapshot, validateSnapshot(snapshot)); snapshot.input_tokens -= 1; snapshot.total_cost += 1; try std.testing.expectError(error.InvalidUsageSnapshot, validateSnapshot(snapshot));}
fn legacyUsageForTest(alloc: Allocator) !std.json.Parsed(std.json.Value) { return std.json.parseFromSlice(std.json.Value, alloc, \\{"billing":"complete","api_duration_complete":true,"wall_duration_complete":true,"code_complete":true,"next_sequence":2,"settled_through_sequence":1, \\"api_duration_ms":10,"wall_duration_ms":20,"total_cost":1,"input_tokens":10,"output_tokens":3,"cache_read_tokens":2,"cache_write_tokens":0,"billable_web_search_calls":0,"lines_added":0,"lines_removed":0, \\"models":[{"model":"test/model","first_sequence":1,"total_cost":1,"input_tokens":10,"output_tokens":3,"cache_read_tokens":2,"cache_write_tokens":0,"billable_web_search_calls":0}],"pending":[]} , .{});}
test "legacy usage compatibility preserves valid snapshots and strict parsing" { const alloc = std.testing.allocator; var parsed = try legacyUsageForTest(alloc); defer parsed.deinit(); var strict = try parseSnapshotValue(alloc, parsed.value); defer strict.deinit(alloc); var compatible = try parseLegacySnapshotValue(alloc, parsed.value); defer compatible.deinit(alloc); try std.testing.expect(snapshotEql(strict, compatible));
for ([_][]const u8{ "cache_read_tokens", "cache_write_tokens" }) |field| { const global = parsed.value.object.getPtr(field).?; const model = parsed.value.object.getPtr("models").?.array.items[0].object.getPtr(field).?; const saved = global.*; global.* = .{ .integer = 11 }; model.* = global.*; var unavailable = try parseLegacySnapshotValue(alloc, parsed.value); defer unavailable.deinit(alloc); try validateSnapshot(unavailable); try std.testing.expectEqual(Availability.legacy, unavailable.billing); try std.testing.expectEqual(@as(usize, 0), unavailable.models.len); try std.testing.expectEqual(@as(usize, 0), unavailable.pending.len); try std.testing.expectError(error.InvalidUsageSnapshot, parseSnapshotValue(alloc, parsed.value)); global.* = saved; model.* = saved; }
var encoded: std.Io.Writer.Allocating = .init(alloc); defer encoded.deinit(); try writeRichSnapshot(&encoded.writer, strict); var rich = try std.json.parseFromSlice(std.json.Value, alloc, encoded.written(), .{}); defer rich.deinit(); var rich_copy = try parseLegacySnapshotValue(alloc, rich.value); defer rich_copy.deinit(alloc); try std.testing.expect(snapshotEql(strict, rich_copy)); rich.value.object.getPtr("cache_read_tokens").?.* = .{ .integer = 11 }; rich.value.object.getPtr("models").?.array.items[0].object.getPtr("cache_read_tokens").?.* = .{ .integer = 11 }; try std.testing.expectError(error.InvalidUsageSnapshot, parseLegacySnapshotValue(alloc, rich.value));}
test "legacy usage compatibility does not hide malformed accounting" { const alloc = std.testing.allocator; var parsed = try legacyUsageForTest(alloc); defer parsed.deinit(); parsed.value.object.getPtr("cache_read_tokens").?.* = .{ .integer = 11 }; const model = &parsed.value.object.getPtr("models").?.array.items[0]; model.object.getPtr("cache_read_tokens").?.* = .{ .integer = 11 };
const cases = [_]struct { field: []const u8, value: std.json.Value, want: error{ InvalidGenerationRecord, InvalidUsageSnapshot }, }{ .{ .field = "total_cost", .value = .{ .float = -1 }, .want = error.InvalidGenerationRecord }, .{ .field = "total_cost", .value = .{ .float = std.math.inf(f64) }, .want = error.InvalidGenerationRecord }, .{ .field = "input_tokens", .value = .{ .integer = 9 }, .want = error.InvalidUsageSnapshot }, .{ .field = "next_sequence", .value = .{ .integer = 0 }, .want = error.InvalidUsageSnapshot }, }; for (cases) |case| { const field = parsed.value.object.getPtr(case.field).?; const saved = field.*; field.* = case.value; try std.testing.expectError(case.want, parseLegacySnapshotValue(alloc, parsed.value)); try std.testing.expectError(case.want, parseSnapshotValue(alloc, parsed.value)); field.* = saved; } model.object.getPtr("first_sequence").?.* = .{ .integer = 0 }; try std.testing.expectError(error.InvalidUsageSnapshot, parseLegacySnapshotValue(alloc, parsed.value)); model.object.getPtr("first_sequence").?.* = .{ .integer = 1 }; const models = parsed.value.object.getPtr("models").?; try models.array.append(models.array.items[0]); const totals = [_]struct { field: []const u8, original: i64, doubled: i64 }{ .{ .field = "total_cost", .original = 1, .doubled = 2 }, .{ .field = "input_tokens", .original = 10, .doubled = 20 }, .{ .field = "output_tokens", .original = 3, .doubled = 6 }, .{ .field = "cache_read_tokens", .original = 11, .doubled = 22 }, }; for (totals) |field| parsed.value.object.getPtr(field.field).?.* = .{ .integer = field.doubled }; try std.testing.expectError(error.InvalidUsageSnapshot, parseLegacySnapshotValue(alloc, parsed.value)); models.array.items.len = 1; for (totals) |field| parsed.value.object.getPtr(field.field).?.* = .{ .integer = field.original };
var pending = try std.json.parseFromSlice(std.json.Value, alloc, "[{\"id\":\"gen_01ARZ3NDEKTSV4RRFFQ69G5FAV\",\"sequence\":0,\"origin\":\"https://ai-gateway.vercel.sh\",\"team\":null}]", .{}); defer pending.deinit(); const pending_field = parsed.value.object.getPtr("pending").?; const old_pending = pending_field.*; pending_field.* = pending.value; parsed.value.object.getPtr("billing").?.* = .{ .string = "pending" }; try std.testing.expectError(error.InvalidUsageSnapshot, parseLegacySnapshotValue(alloc, parsed.value)); pending_field.* = old_pending; parsed.value.object.getPtr("billing").?.* = .{ .string = "complete" }; try parsed.value.object.put(parsed.arena.allocator(), "unknown", .null); try std.testing.expectError(error.InvalidGenerationRecord, parseLegacySnapshotValue(alloc, parsed.value)); try std.testing.expectError(error.InvalidUsageSnapshot, parseLegacySnapshotValue(alloc, .null));}
test "legacy usage compatibility releases rejected and unavailable allocations" { const alloc = std.testing.allocator; var parsed = try legacyUsageForTest(alloc); defer parsed.deinit(); parsed.value.object.getPtr("cache_read_tokens").?.* = .{ .integer = 11 }; parsed.value.object.getPtr("models").?.array.items[0].object.getPtr("cache_read_tokens").?.* = .{ .integer = 11 }; const Check = struct { fn run(a: Allocator, value: std.json.Value, reject: bool) !void { if (reject) { if (parseLegacySnapshotValue(a, value)) |result| { var owned = result; owned.deinit(a); return error.ExpectedInvalidUsage; } else |err| switch (err) { error.InvalidUsageSnapshot => return, else => return err, } } var snapshot = try parseLegacySnapshotValue(a, value); defer snapshot.deinit(a); try std.testing.expectEqual(Availability.legacy, snapshot.billing); } }; try std.testing.checkAllAllocationFailures(alloc, Check.run, .{ parsed.value, false }); parsed.value.object.getPtr("input_tokens").?.* = .{ .integer = 9 }; try std.testing.checkAllAllocationFailures(alloc, Check.run, .{ parsed.value, true });}
test "usage deduplicates terminal and generation callbacks" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, sequence, 10, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", null, ); try usage.finishObservedInvocation( alloc, sequence, 10, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", null, ); const record = GenerationRecord{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .model = "anthropic/claude-opus-4.8", .total_cost = 1.5, .input_tokens = 10, .output_tokens = 2, .cache_read_tokens = 3, .cache_write_tokens = 1, .billable_web_search_calls = 0, }; try usage.applyGeneration(alloc, record); try usage.applyGeneration(alloc, record);
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(@as(u64, 10), snapshot.api_duration_ms); try std.testing.expectApproxEqAbs(@as(f64, 1.5), snapshot.total_cost, 1e-12); try std.testing.expectEqual(@as(u64, 10), snapshot.input_tokens); try std.testing.expectEqual(@as(usize, 1), snapshot.models.len);}
test "usage deduplicates callbacks after the durable settlement boundary" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, sequence, 10, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", null, ); try usage.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .model = "anthropic/claude-opus-4.8", .total_cost = 1.5, .input_tokens = 10, .output_tokens = 2, .cache_read_tokens = 3, .cache_write_tokens = 1, .billable_web_search_calls = 0, });
var persisted = try usage.snapshot(alloc); defer persisted.deinit(alloc); usage.markClean(persisted);
usage.finishInvocation(sequence, 10, .observed_generation); try usage.finishObservedInvocation( alloc, sequence, 10, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", null, );
var current = try usage.snapshot(alloc); defer current.deinit(alloc); try std.testing.expectEqual(@as(u64, 10), current.api_duration_ms); try std.testing.expectApproxEqAbs(@as(f64, 1.5), current.total_cost, 1e-12); try std.testing.expectEqual(@as(usize, 0), current.pending.len);}
test "model aggregates remain ordered by invocation when reconciliation finishes out of order" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const first_sequence = try usage.reserveInvocation(); const second_sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, first_sequence, 1, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", null, ); try usage.finishObservedInvocation( alloc, second_sequence, 1, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAW", "https://ai-gateway.vercel.sh", null, );
try usage.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAW", .model = "openai/gpt-5.5", .total_cost = 1, .input_tokens = 1, .output_tokens = 1, .cache_read_tokens = 0, .cache_write_tokens = 0, .billable_web_search_calls = 0, }); try usage.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .model = "anthropic/claude-opus-4.8", .total_cost = 1, .input_tokens = 1, .output_tokens = 1, .cache_read_tokens = 0, .cache_write_tokens = 0, .billable_web_search_calls = 0, });
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqualStrings("anthropic/claude-opus-4.8", snapshot.models[0].model); try std.testing.expectEqualStrings("openai/gpt-5.5", snapshot.models[1].model);}
test "active invocation does not render partial usage as complete" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const sequence = try usage.reserveInvocation(); var active = try usage.snapshot(alloc); defer active.deinit(alloc); try std.testing.expectEqual(Availability.incomplete, active.billing); try std.testing.expect(!active.api_duration_complete);
usage.finishInvocation(sequence, 10, .unbilled); var finished = try usage.snapshot(alloc); defer finished.deinit(alloc); try std.testing.expectEqual(Availability.complete, finished.billing); try std.testing.expect(finished.api_duration_complete);}
test "active invocation dominates separate pending and publication state" { const alloc = std.testing.allocator; const Checkpoint = struct { fn persist(_: *anyopaque, _: Snapshot) !void {} }; var checkpoint_context: u8 = 0; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const active_sequence = try usage.reserveInvocation(); const observed_sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, observed_sequence, 1, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", null, );
var pending = try usage.snapshot(alloc); defer pending.deinit(alloc); try std.testing.expectEqual(Availability.incomplete, pending.billing); try std.testing.expectEqual(@as(u64, 0), pending.settled_through_sequence); try std.testing.expectEqual(@as(usize, 1), pending.pending.len);
usage.configureCheckpointSink(.{ .context = &checkpoint_context, .allocator = alloc, .persist = Checkpoint.persist, }); try usage.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .created_at_ms = 1000, .model = "provider/model", .total_cost = 0.25, .input_tokens = 10, .output_tokens = 2, .cache_read_tokens = 1, .cache_write_tokens = 0, .reasoning_tokens = 1, .billable_web_search_calls = 0, });
var backlogged = try usage.snapshot(alloc); defer backlogged.deinit(alloc); try std.testing.expectEqual(Availability.incomplete, backlogged.billing); try std.testing.expectEqual(@as(u64, 0), backlogged.settled_through_sequence); try std.testing.expectEqual(@as(usize, 1), backlogged.pending.len); try std.testing.expectEqual( @as(usize, 1), backlogged.publication_backlog.len, );
usage.finishInvocation(active_sequence, 1, .unbilled); var settled = try usage.snapshot(alloc); defer settled.deinit(alloc); try std.testing.expectEqual(Availability.pending, settled.billing); try std.testing.expectEqual(@as(u64, 2), settled.settled_through_sequence);}
test "active invocation capacity fails before provider admission" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc); var observations: [max_active_invocations]InvocationObservation = undefined;
for (&observations) |*observation| { observation.* = try InvocationObservation.begin(&usage); } try std.testing.expectError( error.UsageCapacityExceeded, InvocationObservation.begin(&usage), );
try observations[0].fail(.unbilled); const replacement = try InvocationObservation.begin(&usage); for (observations[1..]) |observation| try observation.fail(.unbilled); try replacement.fail(.unbilled); var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(Availability.complete, snapshot.billing); try std.testing.expect(snapshot.api_duration_complete);}
test "generation allocation failure marks billing incomplete before snapshot" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc); const sequence = try usage.reserveInvocation(); var failing = std.testing.FailingAllocator.init( alloc, .{ .fail_index = 0 }, );
try std.testing.expectError( error.OutOfMemory, usage.finishObservedInvocation( failing.allocator(), sequence, 1, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", null, ), ); var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(Availability.incomplete, snapshot.billing); try std.testing.expect(snapshot.api_duration_complete); try std.testing.expectEqual(@as(u64, 1), snapshot.settled_through_sequence); try std.testing.expectEqual(@as(usize, 0), snapshot.pending.len);}
test "invalid generation identity settles the provider observation" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const observation = try InvocationObservation.begin(&usage); try observation.complete( alloc, .{ .generation_id = "resp_provider_local", .billing = .{ .created_at_ms = 1, .model = "provider/model", .total_cost = 0, .input_tokens = 1, .output_tokens = 1, .cache_read_tokens = 0, .cache_write_tokens = 0, .reasoning_tokens = null, .billable_web_search_calls = 0, }, }, .{ .exact = .gateway }, );
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(Availability.incomplete, snapshot.billing); try std.testing.expect(snapshot.api_duration_complete); try std.testing.expectEqual(@as(u64, 1), snapshot.settled_through_sequence); try std.testing.expectEqual(@as(usize, 0), snapshot.pending.len);}
test "rejected observed generation settles without publishing its identity or billing" { const alloc = std.testing.allocator; const rejected_id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV"; const PublicationProbe = struct { rejected_id: []const u8, pending: usize = 0, generations: usize = 0, incidents: usize = 0,
fn publish(raw: *anyopaque, event: usage_report.ProfileEvent) !void { const self: *@This() = @ptrCast(@alignCast(raw)); switch (event) { .pending => |marker| { if (std.mem.eql(u8, marker.id, self.rejected_id)) self.pending += 1; }, .generation => |fact| { if (std.mem.eql(u8, fact.id, self.rejected_id)) self.generations += 1; }, .incident => self.incidents += 1, } } };
var probe = PublicationProbe{ .rejected_id = rejected_id }; var usage = Usage.initFresh(); defer usage.deinit(alloc); usage.configurePublicationSink(.{ .context = &probe, .allocator = alloc, .publish = PublicationProbe.publish, });
const observation = try InvocationObservation.begin(&usage); try observation.complete( alloc, .{ .generation_id = rejected_id, .generation_metadata_invalid = true, .billing = .{ .created_at_ms = 1, .model = "provider/model", .total_cost = 0.25, .input_tokens = 10, .output_tokens = 2, .cache_read_tokens = 1, .cache_write_tokens = 0, .reasoning_tokens = 1, .billable_web_search_calls = 0, }, }, .{ .exact = .gateway }, );
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(Availability.incomplete, snapshot.billing); try std.testing.expect(snapshot.api_duration_complete); try std.testing.expectEqual(@as(u64, 1), snapshot.settled_through_sequence); try std.testing.expectEqual(@as(usize, 0), snapshot.pending.len); try std.testing.expectEqual(@as(u64, 0), snapshot.input_tokens); try std.testing.expectEqual(@as(u64, 0), snapshot.output_tokens); try std.testing.expectEqual(@as(f64, 0), snapshot.total_cost); try std.testing.expectEqual(@as(usize, 0), probe.pending); try std.testing.expectEqual(@as(usize, 0), probe.generations); try std.testing.expectEqual(@as(usize, 1), probe.incidents);}
test "concurrent invocation completion and snapshots stay internally consistent" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc); var start = std.atomic.Value(bool).init(false); var reservation_failed = std.atomic.Value(bool).init(false); const Context = struct { ledger: *Usage, gate: *std.atomic.Value(bool), failed: *std.atomic.Value(bool),
fn run(ctx: *@This()) void { while (!ctx.gate.load(.seq_cst)) std.atomic.spinLoopHint(); const sequence = ctx.ledger.reserveInvocation() catch { ctx.failed.store(true, .seq_cst); return; }; ctx.ledger.finishInvocation(sequence, 1, .unbilled); } }; var contexts: [8]Context = undefined; var threads: [contexts.len]std.Thread = undefined; for (&contexts, &threads) |*context, *thread| { context.* = .{ .ledger = &usage, .gate = &start, .failed = &reservation_failed, }; thread.* = try std.Thread.spawn(.{}, Context.run, .{context}); }
start.store(true, .seq_cst); for (0..8) |_| { var current = try usage.snapshot(alloc); defer current.deinit(alloc); if (current.billing == .complete) { try std.testing.expect(current.api_duration_complete); } } for (threads) |thread| thread.join(); try std.testing.expect(!reservation_failed.load(.seq_cst));
var finished = try usage.snapshot(alloc); defer finished.deinit(alloc); try std.testing.expectEqual(Availability.complete, finished.billing); try std.testing.expect(finished.api_duration_complete); try std.testing.expectEqual(@as(u64, contexts.len), finished.api_duration_ms);}
test "pending and ambiguous generation states remain honest" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const pending_sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, pending_sequence, 8, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", null, );
var pending = try usage.snapshot(alloc); defer pending.deinit(alloc); try std.testing.expectEqual(Availability.pending, pending.billing); try std.testing.expectEqual(@as(usize, 1), pending.pending.len);
const ambiguous_sequence = try usage.reserveInvocation(); usage.finishInvocation(ambiguous_sequence, 12, .ambiguous_delivery);
var incomplete = try usage.snapshot(alloc); defer incomplete.deinit(alloc); try std.testing.expectEqual(Availability.incomplete, incomplete.billing); try std.testing.expectEqual(@as(u64, 20), incomplete.api_duration_ms);}
test "successful response without generation identity marks billing incomplete" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const observation = try InvocationObservation.begin(&usage); try observation.complete( alloc, .{}, .{ .unavailable = .possibly_billed }, );
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(Availability.incomplete, snapshot.billing); try std.testing.expectEqual(@as(usize, 0), snapshot.pending.len);}
test "deferred usage preserves provider and credential authority" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc); const identity = @import("../auth/credential_authority.zig").derive( .fx_login, "acct_1", ).?; const observation = try InvocationObservation.begin(&usage); try observation.complete(alloc, .{}, .{ .deferred = .{ .provider = .gateway, .generation_id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .scope = "https://ai-gateway.vercel.sh", .tenant = "team_1", .account_id = "acct_1", .credential_source = .fx_login, .credential_identity = identity, } });
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(model_provider.ProviderId.gateway, snapshot.pending[0].provider); try std.testing.expectEqual(types.CredentialSource.fx_login, snapshot.pending[0].credential_source.?); try std.testing.expect(snapshot.pending[0].credential_identity.?.eql(identity));}
const TestGenerationUsageProvider = struct { outcome: enum { found, retry, preserve_pending, reject, unavailable, cancelled, }, calls: usize = 0, saw_expected_input: bool = false,
fn lookup( raw_context: ?*anyopaque, alloc: Allocator, input: generation_usage.LookupInput, ) generation_usage.LookupError!generation_usage.LookupOutcome { const self: *@This() = @ptrCast(@alignCast(raw_context.?)); self.calls += 1; self.saw_expected_input = input.credential != null and std.mem.eql(u8, input.credential.?, "credential") and std.mem.eql( u8, input.generation_id, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", ) and std.mem.eql(u8, input.origin, "https://provider.example"); return switch (self.outcome) { .found => blk: { const id = try alloc.dupe(u8, input.generation_id); errdefer mem_utils.free(alloc, id); const model = try alloc.dupe(u8, "provider/model"); break :blk .{ .found = .{ .id = id, .model = model, .total_cost = 0.25, .input_tokens = 10, .output_tokens = 2, .cache_read_tokens = 3, .cache_write_tokens = 1, .billable_web_search_calls = 0, } }; }, .retry => .retry, .preserve_pending => .preserve_pending, .reject => .reject, .unavailable => error.Unavailable, .cancelled => error.Cancelled, }; }
fn provider(self: *@This()) generation_usage.Provider { return .{ .context = self, .lookup_fn = lookup, }; }};
test "usage restore and reset retain the injected generation provider" { const alloc = std.testing.allocator; var fake = TestGenerationUsageProvider{ .outcome = .found }; const provider = fake.provider(); var usage = Usage.initFreshWithProvider(provider); defer usage.deinit(alloc);
var saved = try usage.snapshot(alloc); defer saved.deinit(alloc);
usage.resetFresh(alloc); try std.testing.expect( usage.generation_usage_providers.select(.gateway).?.lookup_fn == provider.lookup_fn, ); try usage.restore(alloc, saved, 1); try std.testing.expect( usage.generation_usage_providers.select(.gateway).?.lookup_fn == provider.lookup_fn, ); usage.resetLegacy(alloc); try std.testing.expect( usage.generation_usage_providers.select(.gateway).?.lookup_fn == provider.lookup_fn, );}
test "reconciliation settles usage through the injected provider" { const alloc = std.testing.allocator; var fake = TestGenerationUsageProvider{ .outcome = .found }; var usage = Usage.initFreshWithProvider(fake.provider()); defer usage.deinit(alloc);
const sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, sequence, 5, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://provider.example", null, ); var cancel = std.atomic.Value(bool).init(false); reconcilePendingBlocking( &usage, alloc, "credential", &cancel, .{ .provider = .gateway, .credential_identity = null }, usage.generation_usage_providers, 1, );
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(@as(usize, 1), fake.calls); try std.testing.expect(fake.saw_expected_input); try std.testing.expectEqual(Availability.complete, snapshot.billing); try std.testing.expectEqual(@as(usize, 0), snapshot.pending.len); try std.testing.expectApproxEqAbs(@as(f64, 0.25), snapshot.total_cost, 1e-12); try std.testing.expectEqual(@as(u64, 10), snapshot.input_tokens); try std.testing.expectEqual(@as(u64, 2), snapshot.output_tokens);}
test "terminal Gateway billing settles the durable observation immediately" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc); const observation = try InvocationObservation.begin(&usage); try observation.complete( alloc, .{ .generation_id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .billing = .{ .created_at_ms = 100, .model = "provider/model", .total_cost = 0.0123, .input_tokens = 130, .output_tokens = 25, .cache_read_tokens = 20, .cache_write_tokens = 10, .reasoning_tokens = 5, .billable_web_search_calls = 2, }, }, testGatewayUsageOutcome("gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", true), );
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(Availability.complete, snapshot.billing); try std.testing.expectEqual(@as(usize, 0), snapshot.pending.len); try std.testing.expectEqual(@as(u64, 130), snapshot.input_tokens); try std.testing.expectEqual(@as(u64, 25), snapshot.output_tokens); try std.testing.expectEqual(@as(u64, 2), snapshot.billable_web_search_calls);}
test "duplicate Gateway terminal callback does not republish inline billing" { const alloc = std.testing.allocator; const CheckpointProbe = struct { fn persist(_: *anyopaque, _: Snapshot) !void {} }; const PublicationProbe = struct { generations: usize = 0,
fn publish(raw: *anyopaque, event: usage_report.ProfileEvent) !void { const self: *@This() = @ptrCast(@alignCast(raw)); switch (event) { .generation => self.generations += 1, .pending, .incident => {}, } } };
var probe = PublicationProbe{}; var checkpoint_context: u8 = 0; var usage = Usage.initFresh(); defer usage.deinit(alloc); usage.configureCheckpointSink(.{ .context = &checkpoint_context, .allocator = alloc, .persist = CheckpointProbe.persist, }); usage.configurePublicationSink(.{ .context = &probe, .allocator = alloc, .publish = PublicationProbe.publish, });
const observation = try InvocationObservation.begin(&usage); const completion: types.ModelCompletion = .{ .generation_id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .billing = .{ .created_at_ms = 100, .model = "provider/model", .total_cost = 0.0123, .input_tokens = 130, .output_tokens = 25, .cache_read_tokens = 20, .cache_write_tokens = 10, .reasoning_tokens = 5, .billable_web_search_calls = 2, }, }; try observation.complete( alloc, completion, testGatewayUsageOutcome("gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", true), ); try observation.complete( alloc, completion, testGatewayUsageOutcome("gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", true), );
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(@as(usize, 1), probe.generations); try std.testing.expectEqual(@as(f64, 0.0123), snapshot.total_cost); try std.testing.expectEqual(@as(u64, 130), snapshot.input_tokens); try std.testing.expectEqual(@as(u64, 25), snapshot.output_tokens);}
test "successful retry after ambiguous delivery retains known generation" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const observation = try InvocationObservation.begin(&usage); try observation.complete( alloc, .{ .generation_id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .delivery_ambiguous = true, }, testGatewayUsageOutcome("gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", false), ); try usage.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .model = "anthropic/claude-opus-4.8", .total_cost = 0.25, .input_tokens = 10, .output_tokens = 2, .cache_read_tokens = 0, .cache_write_tokens = 0, .billable_web_search_calls = 0, });
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(Availability.incomplete, snapshot.billing); try std.testing.expectEqual(@as(usize, 1), snapshot.models.len); try std.testing.expectApproxEqAbs(@as(f64, 0.25), snapshot.total_cost, 1e-12);}
test "wall duration resumes from the session creation time" { const alloc = std.testing.allocator; var first = Usage.initFresh(); defer first.deinit(alloc); var saved = try first.snapshot(alloc); defer saved.deinit(alloc); const session_started_at_ms = io_mod.milliTimestamp() - 500;
var resumed = Usage.initLegacy(); defer resumed.deinit(alloc); try resumed.restore(alloc, saved, session_started_at_ms); var current = try resumed.snapshot(alloc); defer current.deinit(alloc); try std.testing.expect(current.wall_duration_complete); try std.testing.expect(current.wall_duration_ms >= 500); try std.testing.expect(!resumed.isDirty());}
test "legacy sessions retain exact wall duration from durable creation" { const alloc = std.testing.allocator; var usage = Usage.initLegacy(); defer usage.deinit(alloc); usage.restoreLegacyWallDuration(io_mod.milliTimestamp() - 500);
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expect(snapshot.wall_duration_complete); try std.testing.expect(snapshot.wall_duration_ms >= 500); try std.testing.expectEqual(Availability.legacy, snapshot.billing);}
test "mark clean preserves mutations newer than the persisted snapshot" { const alloc = std.testing.allocator; var usage = Usage.initLegacy(); defer usage.deinit(alloc);
try usage.recordCommittedLines(3, 1); var persisted = try usage.snapshot(alloc); defer persisted.deinit(alloc); try usage.recordCommittedLines(2, 4);
usage.markClean(persisted); try std.testing.expect(usage.isDirty());
var latest = try usage.snapshot(alloc); defer latest.deinit(alloc); usage.markClean(latest); try std.testing.expect(!usage.isDirty());}
test "usage keeps more than sixteen exact resolved models" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const suffixes = "ABCDEFGHJKMNPQRST"; for (suffixes, 0..) |suffix, index| { var id_buf: [30]u8 = undefined; const id = try std.fmt.bufPrint( &id_buf, "gen_01ARZ3NDEKTSV4RRFFQ69G5FA{c}", .{suffix}, ); var model_buf: [64]u8 = undefined; const model = try std.fmt.bufPrint(&model_buf, "provider/model-{d}", .{index}); const sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, sequence, 1, .observed_generation, id, "https://ai-gateway.vercel.sh", null, ); try usage.applyGeneration(alloc, .{ .id = id, .model = model, .total_cost = 0.001, .input_tokens = 1, .output_tokens = 1, .cache_read_tokens = 0, .cache_write_tokens = 0, .billable_web_search_calls = 0, }); }
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(@as(usize, 17), snapshot.models.len); try std.testing.expectEqualStrings("provider/model-16", snapshot.models[16].model);}
test "usage records committed lines once and restores a bounded snapshot" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc); try usage.recordCommittedLines(4, 2);
var saved = try usage.snapshot(alloc); defer saved.deinit(alloc); try std.testing.expect(saved.code_complete); try std.testing.expectEqual(@as(u64, 4), saved.lines_added); try std.testing.expectEqual(@as(u64, 2), saved.lines_removed); try validateSnapshot(saved);
var restored = Usage.initLegacy(); defer restored.deinit(alloc); try restored.restore(alloc, saved, io_mod.milliTimestamp());
var round_trip = try restored.snapshot(alloc); defer round_trip.deinit(alloc); try std.testing.expectEqual(Availability.complete, round_trip.billing); try std.testing.expect(round_trip.code_complete); try std.testing.expectEqual(@as(u64, 4), round_trip.lines_added); try std.testing.expectEqual(@as(u64, 2), round_trip.lines_removed);}
test "sixteenth pending generation is the capacity boundary" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
// Every suffix is a valid Crockford identifier character, so each id is a // distinct, well-formed generation id. const suffixes = "ABCDEFGHJKMNPQRST"; try std.testing.expect(suffixes.len > max_pending_generations);
for (suffixes[0..max_pending_generations]) |suffix| { var id_buf: [30]u8 = undefined; const id = try std.fmt.bufPrint( &id_buf, "gen_01ARZ3NDEKTSV4RRFFQ69G5FA{c}", .{suffix}, ); const sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, sequence, 1, .observed_generation, id, "https://ai-gateway.vercel.sh", null, ); }
var at_capacity = try usage.snapshot(alloc); defer at_capacity.deinit(alloc); try std.testing.expectEqual(Availability.pending, at_capacity.billing); try std.testing.expectEqual( @as(usize, max_pending_generations), at_capacity.pending.len, );
var overflow_id_buf: [30]u8 = undefined; const overflow_id = try std.fmt.bufPrint( &overflow_id_buf, "gen_01ARZ3NDEKTSV4RRFFQ69G5FA{c}", .{suffixes[max_pending_generations]}, ); const overflow_sequence = try usage.reserveInvocation(); try std.testing.expectError( error.UsageCapacityExceeded, usage.finishObservedInvocation( alloc, overflow_sequence, 1, .observed_generation, overflow_id, "https://ai-gateway.vercel.sh", null, ), );
var overflowed = try usage.snapshot(alloc); defer overflowed.deinit(alloc); try std.testing.expectEqual(Availability.incomplete, overflowed.billing); try std.testing.expectEqual( @as(usize, max_pending_generations), overflowed.pending.len, );}
test "populated usage snapshot keeps the rollback-readable durable shape" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const first_sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, first_sequence, 120, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", "team_alpha", ); try usage.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .model = "anthropic/claude-opus-4.8", .total_cost = 0.0125, .input_tokens = 100, .output_tokens = 40, .cache_read_tokens = 12, .cache_write_tokens = 8, .billable_web_search_calls = 1, });
const second_sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, second_sequence, 90, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAW", "https://ai-gateway.vercel.sh", null, ); try usage.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAW", .model = "openai/gpt-5.5", .total_cost = 0.004, .input_tokens = 200, .output_tokens = 25, .cache_read_tokens = 0, .cache_write_tokens = 0, .billable_web_search_calls = 0, });
// Leave the third invocation unresolved so the snapshot carries pending // identity that validation cannot reconstruct. const pending_sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, pending_sequence, 30, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAX", "https://ai-gateway.vercel.sh", "team_beta", ); try usage.recordCommittedLines(5, 3);
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(Availability.pending, snapshot.billing); try std.testing.expectEqual(@as(usize, 2), snapshot.models.len); try std.testing.expectEqual(@as(usize, 1), snapshot.pending.len);
var encoded: std.Io.Writer.Allocating = .init(alloc); defer encoded.deinit(); try writeSnapshot(&encoded.writer, snapshot); var parsed = try std.json.parseFromSlice(std.json.Value, alloc, encoded.written(), .{}); defer parsed.deinit(); try std.testing.expectEqual(@as(usize, 18), parsed.value.object.count()); try std.testing.expect(parsed.value.object.get("schema_version") == null); const persisted_models = parsed.value.object.get("models").?.array.items; for (persisted_models) |model| { try std.testing.expectEqual(@as(usize, 8), model.object.count()); } const persisted_pending = parsed.value.object.get("pending").?.array.items; for (persisted_pending) |pending| { try std.testing.expectEqual(@as(usize, 8), pending.object.count()); } var decoded = try parseSnapshotValue(alloc, parsed.value); defer decoded.deinit(alloc);
try std.testing.expectEqual(snapshot.billing, decoded.billing); try std.testing.expectEqual(snapshot.api_duration_complete, decoded.api_duration_complete); try std.testing.expectEqual(snapshot.wall_duration_complete, decoded.wall_duration_complete); try std.testing.expectEqual(snapshot.code_complete, decoded.code_complete); try std.testing.expectEqual(snapshot.next_sequence, decoded.next_sequence); try std.testing.expectEqual(snapshot.settled_through_sequence, decoded.settled_through_sequence); try std.testing.expectEqual(snapshot.api_duration_ms, decoded.api_duration_ms); try std.testing.expectEqual(snapshot.wall_duration_ms, decoded.wall_duration_ms); try std.testing.expectEqual(snapshot.lines_added, decoded.lines_added); try std.testing.expectEqual(snapshot.lines_removed, decoded.lines_removed); try std.testing.expect(snapshot.request_count != null); try std.testing.expect(decoded.request_count == null); try std.testing.expectEqual(@as(usize, 0), decoded.publication_backlog.len); try std.testing.expectEqual(@as(usize, 0), decoded.incidents.len);
try std.testing.expectEqual(snapshot.models.len, decoded.models.len); for (snapshot.models, decoded.models) |expected, actual| { try std.testing.expectEqualStrings(expected.model, actual.model); try std.testing.expectEqual(expected.first_sequence, actual.first_sequence); try std.testing.expectApproxEqAbs(expected.total_cost, actual.total_cost, 1e-12); try std.testing.expectEqual(expected.input_tokens, actual.input_tokens); try std.testing.expectEqual(expected.output_tokens, actual.output_tokens); try std.testing.expectEqual(expected.cache_read_tokens, actual.cache_read_tokens); try std.testing.expectEqual(expected.cache_write_tokens, actual.cache_write_tokens); try std.testing.expectEqual( expected.billable_web_search_calls, actual.billable_web_search_calls, ); }
try std.testing.expectEqual(snapshot.pending.len, decoded.pending.len); for (snapshot.pending, decoded.pending) |expected, actual| { try std.testing.expectEqualStrings(expected.id, actual.id); try std.testing.expectEqual(expected.sequence, actual.sequence); try std.testing.expectEqualStrings(expected.origin, actual.origin); try std.testing.expect(optionalStringsEqual(expected.team, actual.team)); }
const Publisher = struct { fn publish(_: *anyopaque, _: usage_report.ProfileEvent) !void {} }; var publisher_context: u8 = 0; var resumed = Usage.initFresh(); defer resumed.deinit(alloc); resumed.configurePublicationSink(.{ .context = &publisher_context, .allocator = alloc, .publish = Publisher.publish, }); try resumed.restore(alloc, decoded, 1); try resumed.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAX", .model = "provider/new", .total_cost = 0.001, .input_tokens = 2, .output_tokens = 1, .cache_read_tokens = 0, .cache_write_tokens = 0, .reasoning_tokens = 1, .billable_web_search_calls = 0, }); var continued = try resumed.snapshot(alloc); defer continued.deinit(alloc); try validateSnapshot(continued); try std.testing.expect(continued.reasoning_tokens == null); try std.testing.expect(continued.request_count == null); try std.testing.expectEqual(Availability.complete, continued.billing);}
test "rich usage snapshot preserves optional metrics and recovery state" { const alloc = std.testing.allocator; const Publisher = struct { fn fail(_: *anyopaque, _: usage_report.ProfileEvent) !void { return error.InjectedPublicationFailure; } }; var publisher_context: u8 = 0; var usage = Usage.initFresh(); defer usage.deinit(alloc); usage.configurePublicationSink(.{ .context = &publisher_context, .allocator = alloc, .publish = Publisher.fail, });
const resolved_sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, resolved_sequence, 5, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", null, ); try usage.applyGeneration(alloc, .{ .id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", .created_at_ms = 100, .model = "provider/model", .total_cost = 0.01, .input_tokens = 10, .output_tokens = 4, .cache_read_tokens = 2, .cache_write_tokens = 1, .reasoning_tokens = 3, .billable_web_search_calls = 0, }); const pending_sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, pending_sequence, 6, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAW", "https://ai-gateway.vercel.sh", "team_alpha", );
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try appendIncidentOwned(alloc, &snapshot, .{ .occurred_at_ms = 90, .completeness = .incomplete, }); var encoded: std.Io.Writer.Allocating = .init(alloc); defer encoded.deinit(); try writeRichSnapshot(&encoded.writer, snapshot); var parsed = try std.json.parseFromSlice( std.json.Value, alloc, encoded.written(), .{}, ); defer parsed.deinit(); var decoded = try parseSnapshotValue(alloc, parsed.value); defer decoded.deinit(alloc);
try std.testing.expectEqual(snapshot.reasoning_tokens, decoded.reasoning_tokens); try std.testing.expectEqual(snapshot.request_count, decoded.request_count); try std.testing.expectEqual( snapshot.models[0].reasoning_tokens, decoded.models[0].reasoning_tokens, ); try std.testing.expectEqual( snapshot.models[0].request_count, decoded.models[0].request_count, ); try std.testing.expectEqual( snapshot.pending[0].observed_at_ms, decoded.pending[0].observed_at_ms, ); try std.testing.expectEqual(@as(usize, 1), decoded.publication_backlog.len); try std.testing.expect(usage_report.GenerationFact.eql( snapshot.publication_backlog[0], decoded.publication_backlog[0], )); try std.testing.expectEqualSlices( usage_report.Incident, snapshot.incidents, decoded.incidents, );}
test "owned snapshot incident overflow matches runtime retention" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc); var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc);
for (0..max_usage_incidents) |index| { try appendIncidentOwned(alloc, &snapshot, .{ .occurred_at_ms = @intCast(index + 1), .completeness = .pending, }); } try std.testing.expectEqual(max_usage_incidents, snapshot.incidents.len); try appendIncidentOwned(alloc, &snapshot, .{ .occurred_at_ms = max_usage_incidents + 1, .completeness = .pending, }); try std.testing.expectEqual(@as(usize, 1), snapshot.incidents.len); try std.testing.expectEqual( @as(i64, max_usage_incidents + 1), snapshot.incidents[0].occurred_at_ms, ); try std.testing.expectEqual( usage_report.Completeness.incomplete, snapshot.incidents[0].completeness, );}
test "provider rejection removes pending generation and marks billing incomplete" { const alloc = std.testing.allocator; var fake = TestGenerationUsageProvider{ .outcome = .reject }; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, sequence, 5, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://provider.example", null, ); var cancel = std.atomic.Value(bool).init(false); reconcilePendingBlocking( &usage, alloc, "credential", &cancel, .{ .provider = .gateway, .credential_identity = null }, generation_usage.Set.gatewayOnly(fake.provider()), 1, );
var rejected = try usage.snapshot(alloc); defer rejected.deinit(alloc); try std.testing.expectEqual(@as(usize, 1), fake.calls); try std.testing.expectEqual(Availability.incomplete, rejected.billing); try std.testing.expectEqual(@as(usize, 0), rejected.pending.len);}
test "provider preserve outcome keeps pending generation for a future credential" { const alloc = std.testing.allocator; var fake = TestGenerationUsageProvider{ .outcome = .preserve_pending }; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, sequence, 5, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://provider.example", null, ); var cancel = std.atomic.Value(bool).init(false); reconcilePendingBlocking( &usage, alloc, "credential", &cancel, .{ .provider = .gateway, .credential_identity = null }, generation_usage.Set.gatewayOnly(fake.provider()), 1, );
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(@as(usize, 1), fake.calls); try std.testing.expectEqual(Availability.pending, snapshot.billing); try std.testing.expectEqual(@as(usize, 1), snapshot.pending.len); try std.testing.expectEqualStrings( "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", snapshot.pending[0].id, );}
test "provider retry outcome leaves pending generation unchanged" { const alloc = std.testing.allocator; var fake = TestGenerationUsageProvider{ .outcome = .retry }; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, sequence, 5, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://provider.example", null, ); var cancel = std.atomic.Value(bool).init(false); reconcilePendingBlocking( &usage, alloc, "credential", &cancel, .{ .provider = .gateway, .credential_identity = null }, generation_usage.Set.gatewayOnly(fake.provider()), 1, );
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(@as(usize, 1), fake.calls); try std.testing.expectEqual(Availability.pending, snapshot.billing); try std.testing.expectEqual(@as(usize, 1), snapshot.pending.len);}
test "provider failure leaves pending generation unchanged" { const alloc = std.testing.allocator; var fake = TestGenerationUsageProvider{ .outcome = .unavailable }; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, sequence, 5, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://provider.example", null, ); var cancel = std.atomic.Value(bool).init(false); reconcilePendingBlocking( &usage, alloc, "credential", &cancel, .{ .provider = .gateway, .credential_identity = null }, generation_usage.Set.gatewayOnly(fake.provider()), 1, );
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(@as(usize, 1), fake.calls); try std.testing.expectEqual(Availability.pending, snapshot.billing); try std.testing.expectEqual(@as(usize, 1), snapshot.pending.len);}
test "provider cancellation stops reconciliation and preserves pending usage" { const alloc = std.testing.allocator; var fake = TestGenerationUsageProvider{ .outcome = .cancelled }; var usage = Usage.initFresh(); defer usage.deinit(alloc);
const sequence = try usage.reserveInvocation(); try usage.finishObservedInvocation( alloc, sequence, 5, .observed_generation, "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://provider.example", null, ); var cancel = std.atomic.Value(bool).init(false); reconcilePendingBlocking( &usage, alloc, "credential", &cancel, .{ .provider = .gateway, .credential_identity = null }, generation_usage.Set.gatewayOnly(fake.provider()), 2, );
var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(@as(usize, 1), fake.calls); try std.testing.expectEqual(Availability.pending, snapshot.billing); try std.testing.expectEqual(@as(usize, 1), snapshot.pending.len);}
test "checkpoint snapshot does not join a completed reconciliation worker" { const Capture = struct { calls: usize = 0,
fn persist(raw: *anyopaque, _: Snapshot) !void { const self: *@This() = @ptrCast(@alignCast(raw)); self.calls += 1; } }; const Worker = struct { gate: *std.atomic.Value(bool),
fn run(self: *@This()) void { while (!self.gate.load(.seq_cst)) std.atomic.spinLoopHint(); } };
const alloc = std.testing.allocator; var capture = Capture{}; var usage = Usage.initFresh(); defer usage.deinit(alloc); usage.configureCheckpointSink(.{ .context = &capture, .allocator = alloc, .persist = Capture.persist, });
var gate = std.atomic.Value(bool).init(false); var worker = Worker{ .gate = &gate }; usage.reconciliation_thread = try std.Thread.spawn( .{}, Worker.run, .{&worker}, ); usage.reconciliation_done.store(true, .seq_cst); defer gate.store(true, .seq_cst);
try std.testing.expect(usage.persistCheckpoint()); try std.testing.expectEqual(@as(usize, 1), capture.calls); try std.testing.expect(usage.reconciliation_thread != null);
gate.store(true, .seq_cst); usage.finishReconciliationIfDone(); try std.testing.expect(usage.reconciliation_thread == null);}
test "gateway observation checkpoints active and terminal usage states" { const Capture = struct { calls: usize = 0, billing: [2]Availability = undefined, api_duration_complete: [2]bool = undefined, pending_count: [2]usize = undefined,
fn persist(raw: *anyopaque, snapshot: Snapshot) !void { const self: *@This() = @ptrCast(@alignCast(raw)); const index = self.calls; self.billing[index] = snapshot.billing; self.api_duration_complete[index] = snapshot.api_duration_complete; self.pending_count[index] = snapshot.pending.len; self.calls += 1; } };
const alloc = std.testing.allocator; var capture = Capture{}; var usage = Usage.initFresh(); defer usage.deinit(alloc); usage.configureCheckpointSink(.{ .context = &capture, .allocator = alloc, .persist = Capture.persist, });
const observation = try InvocationObservation.begin(&usage); try std.testing.expectEqual(@as(usize, 1), capture.calls); try std.testing.expectEqual(Availability.incomplete, capture.billing[0]); try std.testing.expect(!capture.api_duration_complete[0]); try std.testing.expectEqual(@as(usize, 0), capture.pending_count[0]);
try observation.complete( alloc, .{ .generation_id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV" }, testGatewayUsageOutcome("gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", false), ); try std.testing.expectEqual(@as(usize, 2), capture.calls); try std.testing.expectEqual(Availability.pending, capture.billing[1]); try std.testing.expect(capture.api_duration_complete[1]); try std.testing.expectEqual(@as(usize, 1), capture.pending_count[1]);}
test "gateway observation does not proceed when active checkpoint fails" { const Reject = struct { calls: usize = 0,
fn persist(raw: *anyopaque, _: Snapshot) !void { const self: *@This() = @ptrCast(@alignCast(raw)); self.calls += 1; return error.CheckpointRejected; } };
const alloc = std.testing.allocator; var reject = Reject{}; var usage = Usage.initFresh(); defer usage.deinit(alloc); usage.configureCheckpointSink(.{ .context = &reject, .allocator = alloc, .persist = Reject.persist, });
try std.testing.expectError( error.CheckpointRejected, InvocationObservation.begin(&usage), ); try std.testing.expectEqual(@as(usize, 1), reject.calls); var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(Availability.complete, snapshot.billing); try std.testing.expect(snapshot.api_duration_complete); try std.testing.expectEqual(@as(u64, 1), snapshot.settled_through_sequence);}
test "terminal checkpoint failure preserves request progress as incomplete" { const Reject = struct { calls: usize = 0,
fn persist(raw: *anyopaque, _: Snapshot) !void { const self: *@This() = @ptrCast(@alignCast(raw)); self.calls += 1; if (self.calls == 2) return error.CheckpointRejected; } };
const alloc = std.testing.allocator; var reject = Reject{}; var usage = Usage.initFresh(); defer usage.deinit(alloc); usage.configureCheckpointSink(.{ .context = &reject, .allocator = alloc, .persist = Reject.persist, });
const observation = try InvocationObservation.begin(&usage); try observation.fail(.unbilled); try std.testing.expectEqual(@as(usize, 2), reject.calls); var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(Availability.incomplete, snapshot.billing); try std.testing.expect(snapshot.api_duration_complete); try std.testing.expectEqual(@as(u64, 1), snapshot.settled_through_sequence);}
test "pending generation origins remain bounded transport inputs" { try validateOrigin("https://ai-gateway.vercel.sh"); try validateOrigin("http://127.0.0.1:3000"); try validateOrigin("https://provider.example");
try std.testing.expectError( error.InvalidGatewayOrigin, validateOrigin(""), ); try std.testing.expectError( error.InvalidGatewayOrigin, validateOrigin("https://provider.example\ninjected"), ); var oversized: [max_origin_bytes + 1]u8 = @splat('a'); try std.testing.expectError( error.InvalidGatewayOrigin, validateOrigin(&oversized), );}
test "reconciliation retry delay observes cancellation promptly" { var cancel = std.atomic.Value(bool).init(false); var cancelled_at_ms = std.atomic.Value(i64).init(0); const Trigger = struct { fn run(flag: *std.atomic.Value(bool), signal_time: *std.atomic.Value(i64)) void { var io_backend: std.Io.Threaded = .init_single_threaded; io_backend.io().sleep(.fromMilliseconds(20), .real) catch {}; signal_time.store(io_mod.milliTimestamp(), .seq_cst); flag.store(true, .seq_cst); } }; const thread = try std.Thread.spawn(.{}, Trigger.run, .{ &cancel, &cancelled_at_ms }); const completed = waitForReconciliationRetry(&cancel); const completed_at_ms = io_mod.milliTimestamp(); thread.join();
try std.testing.expect(!completed); const signalled_at_ms = cancelled_at_ms.load(.seq_cst); try std.testing.expect(signalled_at_ms > 0); const cancellation_response_bound_ms = 250; try std.testing.expect(completed_at_ms - signalled_at_ms < cancellation_response_bound_ms);}
test "stale reconciliation credential cannot replace a refreshed credential" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
usage.startReconciliation(alloc, "initial-key"); usage.replaceReconciliationCredential(alloc, "refreshed-key"); usage.startReconciliation(alloc, "initial-key"); usage.refreshReconciliationCredential( alloc, "initial-key", "stale-producer-key", );
const expected = reconciliationKeyDigest("refreshed-key"); try std.testing.expectEqualSlices( u8, &expected, &(usage.reconciliation_key_digest orelse return error.TestUnexpectedResult), );
usage.refreshReconciliationCredential( alloc, "refreshed-key", "next-key", ); const next_expected = reconciliationKeyDigest("next-key"); try std.testing.expectEqualSlices( u8, &next_expected, &(usage.reconciliation_key_digest orelse return error.TestUnexpectedResult), );
usage.clearReconciliationCredential(); usage.refreshReconciliationCredential( alloc, "next-key", "stale-after-clear", ); try std.testing.expect(usage.reconciliation_key_digest == null); try std.testing.expect(usage.reconciliation_credential_blocked);
usage.replaceReconciliationCredential(alloc, "replacement-key"); const replacement_expected = reconciliationKeyDigest("replacement-key"); try std.testing.expectEqualSlices( u8, &replacement_expected, &(usage.reconciliation_key_digest orelse return error.TestUnexpectedResult), ); try std.testing.expect(!usage.reconciliation_credential_blocked);}
test "resumed provider reconciliation uses Gateway credential slot identity" { const alloc = std.testing.allocator; var usage = Usage.initFresh(); defer usage.deinit(alloc);
usage.replaceProviderReconciliationCredential( alloc, .gateway, .ai_gateway_api_key, null, "fresh-secret-key", ); try std.testing.expect(usage.reconciliation_key_digest != null); try std.testing.expect(!usage.reconciliation_credential_blocked);
const observation = try InvocationObservation.begin(&usage); try observation.complete(alloc, .{}, .{ .deferred = testGatewayUsageReference( "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", "https://ai-gateway.vercel.sh", ) });
usage.replaceProviderReconciliationCredential( alloc, .gateway, .ai_gateway_api_key, null, "secret-key", ); try std.testing.expect(usage.reconciliation_key_digest != null); try std.testing.expect(usage.reconciliation_authority.?.credential_identity != null); try std.testing.expect(!usage.reconciliation_credential_blocked);
usage.replaceProviderReconciliationCredential( alloc, .codex, .chatgpt_subscription, null, "subscription-token", ); try std.testing.expect(usage.reconciliation_key_digest == null); try std.testing.expect(usage.reconciliation_authority == null); try std.testing.expect(usage.reconciliation_credential_blocked);}
test "host-managed reconciliation records authority without credential bytes" { var usage = Usage.initFresh(); defer usage.deinit(std.testing.allocator);
usage.replaceHostManagedReconciliationAuthority(std.testing.allocator, .gateway);
try std.testing.expect(usage.reconciliation_key_digest != null); try std.testing.expectEqual(model_provider.ProviderId.gateway, usage.reconciliation_authority.?.provider); try std.testing.expect(usage.reconciliation_authority.?.credential_identity.?.eql( credential_authority.derive(.host_managed, null).?, )); try std.testing.expect(!usage.reconciliation_credential_blocked);}
test "terminal checkpoint writer failure stops invocation completion and releases its lock" { const Reject = struct { calls: usize = 0, fn persist(raw: *anyopaque, _: Snapshot) !void { const self: *@This() = @ptrCast(@alignCast(raw)); self.calls += 1; if (self.calls == 2) return error.SessionPersistenceUncertain; } }; const alloc = std.testing.allocator; for ([_]bool{ false, true }) |deferred| { var reject = Reject{}; var usage = Usage.initFresh(); defer usage.deinit(alloc); usage.configureCheckpointSink(.{ .context = &reject, .allocator = alloc, .persist = Reject.persist }); const observation = try InvocationObservation.begin(&usage); try std.testing.expectError(error.SessionPersistenceUncertain, observation.complete( alloc, .{ .generation_id = "gen_01ARZ3NDEKTSV4RRFFQ69G5FAV" }, if (deferred) testGatewayUsageOutcome("gen_01ARZ3NDEKTSV4RRFFQ69G5FAV", false) else .{ .unavailable = .unbilled }, )); try std.testing.expectEqual(@as(usize, 2), reject.calls); try std.testing.expect(usage.checkpoint_mutex.tryLock()); usage.checkpoint_mutex.unlock(io_mod.getIo()); var snapshot = try usage.snapshot(alloc); defer snapshot.deinit(alloc); try std.testing.expectEqual(Availability.incomplete, snapshot.billing); }}