From c62e9df39f3bc2e1eae124a79075e9bc5555fe1a Mon Sep 17 00:00:00 2001 From: andrewgu-oai Date: Wed, 28 Oct 2026 20:57:35 +0000 Subject: [PATCH] Preserve reasoning token accounting across resume and fork (#57839) ## Why Resuming or forking a thread loses the session flag indicating that reported token usage already includes retained reasoning. Counting that reasoning again can trigger automatic compaction too early. ## What changed Persist `last_usage_includes_reasoning` with `TokenUsageInfo` and use it when estimating context usage. Record the flag with completed sampling and compaction usage, and mark local estimates as including reasoning. Older records without the field retain the existing accounting behavior. ## Testing Add coverage for saved inclusive, exclusive, and legacy token counts. Extend the encrypted-reasoning compaction test to verify the compaction threshold for live, forked, and cold-resumed threads. GitOrigin-RevId: 9f22262a693111ec0406dce11804eb4844ad8983 --- .../app-server-exports-stable.json.zst | Bin 162792 -> 162836 bytes .../app-server/src/bespoke_event_handling.rs | 1 + codex-rs/app-server/tests/common/rollout.rs | 1 + .../tests/suite/v2/thread_resume.rs | 3 + codex-rs/core/src/compact.rs | 13 ++-- codex-rs/core/src/context_manager/history.rs | 22 +++--- .../core/src/context_manager/history_tests.rs | 30 +++++++- .../src/session/async_compaction_tests.rs | 1 + .../src/session/guardian_checkpoint_tests.rs | 1 + codex-rs/core/src/session/mod.rs | 22 +++--- codex-rs/core/src/session/tests.rs | 23 +++++- codex-rs/core/src/session/turn.rs | 7 +- codex-rs/core/src/state/session.rs | 19 ++--- codex-rs/core/src/tasks/mod_tests.rs | 11 ++- codex-rs/core/tests/suite/client.rs | 1 + codex-rs/core/tests/suite/compact.rs | 71 +++++++++++++----- .../ext/goal/tests/goal_extension_backend.rs | 1 + .../src/sessions/export.rs | 1 + codex-rs/protocol/src/protocol.rs | 53 +++++++------ codex-rs/state/src/runtime/threads.rs | 1 + 20 files changed, 191 insertions(+), 91 deletions(-) diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst index 4913a5dc765e925f35d18a1481dc3ccaa37068e7..3668daa519ed163e0a4448b8b549e1d37e602886 100644 GIT binary patch delta 8721 zcmV+sBJSPj_X(8v34nwFv;r$w1{zCpUOzNVDX~RmIo=JJ|k2QrFC~k-qED z6VmZ#&h3*%5QabgkdSK~D@HE5E*^j1P@#$DFs)VWBQuiOp#rf*YTNt{hlyejkCZ5k zmQ)E|TuesfR;sto=#ngd1s_FmRBDEaVRHcetU9SY$yJni$XdbW5<$&zY@4xPpfFKZ zBAm`mYGRrUaA|NSpax^46NFe79f#&ad zNT9i8r6~vr?pg$z&oYxQDe*bcA7dwLAAo}i61r#}A7FPn8R2t~4X}T!yA5-H_5kHD zLDTi576`M}`Ei`1|p6rEp7e=n_=g? zkACux{eeXusfapbkrjVf*d;Mh-@5RMPE|o%Y$)t{thCGrpU#OYf&d`Zx4SmD^6h2f z5eP6c4`OmbF#-q`gF<@xR+VfdL1JIK0hx(PhgY9GI}ML~1D#DQKK*H1TDdr#7~KM- zNuE3xs~yzrV?R)E02CM*U&?lMYU90l$s>`4$f6R{1-V<#(sh5kiXO*#@|wYZ>w5TR z!5B$k{jbH0RGh{}kxVqK8(&$b-kr-FKEi?aI@~wDI_uIhy zcQ`NP`+FkFc-Pp216>@d5?Q`~N=%qGFn+=nLwz5E3(3(NeZ#F4p$9Aa(w;L8tk;|r zHcqPDe0u_5(~Ex!EWxmQ4R7|&NFuWj>k%Btw=3cxRjJw!BFOCDXMk%@ug5W{T$H1A z?jq*0ru0zB-I`oz2RK8~`Xs&ASyVexQ}#qw`u5glVrKKFCd1Ml z(98QScTtcBKEm{SFQD?pOkp7TTg<+9qZx{iA#q1|?cO+&b-h}q6MG>42F5Gi_S(Wg zswwd0^kxLRh{G+FL;F=l48QC;&L_pQfz?Ec^PZOYA##67Z+l(IaM_z>i3Ys2t zhBQXo6BV@>FrR7w!n^<|w`5Wq>g<~ZXb!V{3yKGDLape}bfFIE(4J~@UQ#?a#^u`&21fb2ahA=j2+lrO0Dlo4>AUFU0fG4I;Q@u_m*zy(7n{K!kyo; zz%GAYpF8tzHOMNVOreUSwEi~#ofbU~MDt>ETcvq&$Ds12-KHC|n-|ydn6`*fwb=%xld;XITbj zp<}kz3DaC*hC`QkVA$XA0!4Cz=Z=tM zFHrh=#=!JP$XlQRoJVoGAWlhPTV(p7qv8F+AWG+hFn4c2nEs-*irJYd7{U?Jm3V*s zeB_xfs(%1D9XWN{o3l zCVuIo5oqaO*S|1B+^URu*OSW3t&Z4aCK_G7qqP_9*yK^;&97zZohYPgyJ_Kt9qSp(Yuqqne}WD6vNm2HjhXdg^%12of~JmW!opwF8V$HGf-0=z zw5tN_L==GSCwWoAx6gkf&k=u-(OS}8rk~M6BDza9lNFa23r2;3&(EXEwWA-k<58o0 zI=GZk+ol9~nwm=EWD0$uFEq0RD_p^5IJNK)<=Q>&S@l5CDLT+6@>_b$62p-s5soyA z<9zv2lN4E5*nq&bXPd=^hgdIkiHC~GSmB%N2?UV1OzS zU&Y2p$LiTzLQH0IzW~|*p!cNR%HD)54u3lR0qL)pXvW0M|H{cp7A!K3$Cf^CL z_S6A%ig`kdz#ZZlTi?>PVR;+?%CP z^DpH8S`GOba;AUw<;B$}m6uOqt@){%w5ihK6K@E~7pq>OD0w_z59MrdwD|kToq^RJEu}=zhR50cc9#3@*>$(;m zZMO8D7Dg7@$Ku0w9B1m_AJ_pS;0C#b7o)xZwHnrIT;tr<$3W2WFU}T>1FYW{@Q8Za z?(BjQrZ0c|tyV4~#O#*TGcm@8*1FuGs)Y!&8xpG0o&ndt9k;D7*gyp*WcVHrRmGz1 zi|a$wJxXiw)y76SWw1MhoE6|Og!$QhdJGQJ_NEPv%PB;t8xtt~s`sHtSd8qs`GZ5W zKzGv)gki#?Ef38cOw~)^ln3`_xZOtk~7-x3@_lvCzPKuT$1Kp}pP<$_aByhTs924tDqu9#uA z>8A-}I@y~1VQcf?!f{U-G?@62zQc@w@&X5J3i z;P(sTr-HG&xEdmCZrWN0=-x^^AP86t@h)3}>xH0?uss^*@DZ{oP~sPt?O@vweUtcvj zq-ZWe>(cK(_tfvD0Gi9r2}Rw@*@II{4&bNNTKnKALM-vEXQCdjxheNy7QpU4NP0Y` z;YJ2yi~u_?1CnOpJ80#jLNMh2^k{!jtPCwIMgcYZsSWR*IcX~<8T*|uj*nmm8I(|( z#7);(YKti(fR3sd&)P=5{!(q}aF5PrrzJV^BvskkG5q{e<9Nl=D4m2|4Y5rG{~X;hi^=^u^u^OL7kB-CjZ$uan2%!H9pZm%uS^ zCjOk3*g>0c!HFv=s31Y3vZq5y;4@x{55jCq8KR8qvS16-LxAJWloK-Y9GfUa(`vqz z0kD)Bbbtp`APc5T;YgcJ--ZChz0u&A)oD0u9#HWj)IWiT-y)*u0POsYz9->vQgkav z2H^NAdQTOyuk8k=lI-D!w;_M{VzUCd}H@w@)9_BF*ysN1FSjG6P4y?uT$AwR*!=qk8vVj zN8#gEw+{#T(e>Hl#y5XmDc-;2qeS_ra_f`sQ>nmQ+wo77HG=y zV5(=Cmn349e;xi_j*lUB3hca`R3t|X6kTrEa15@R2N|G%Z6$xU$+XSU?JYbBb`GAz zHQuT7C=m`0uGyv>hh$R6j9hs}2JX<2+nN2J4;m&O6pM} zZn<%B1JrE?TFKCk5%}CZ=M+0y_D;3(HKmB<_`#Kki00tYUrEUVT97|iS0=cUluOj? zv9-9M^8^^31=1z$j6#ssNJO+quCc(5&U1Z;i;ZnU#+ywKlk)GV{=~ z1ZvuyL{w~QZ`qh%0d>cWm2$xnnaAOJPguTz2m8y(ADw@76I2;GfC&~K_t8!$g*(~j~0yx3`lJ`hy2>*t zi;jyZlw2gd85%kTEsX5IEo7>v;xdVkHl)4nKvEU6Kp_tC$eE5V-BO9{;UORJY}qC6 z1Y&=mH3UdU?nP_|D7pVg&IwTRzy%?zvjLRcqt&ru}I0LsS8*C#t zb))EHrUy^At)QHJL-;_mfwz8EQLLwj&h_i)u254)OOF>YjAOdcV z2$EGbvc#;YN#g?)7qWp5P7tZMT+{y>MGJr6ny~ADsK}Fr-WMx|ScA=)#T0<~+-waY zMyOGasep4&I!dPDOrBs_1LdWQ+{DQ8g#EpOd6@AW$DJhRCaj>5;JN|yl#cV&RHtOF z8-`{%cGF2e6ED7w=1cBtzj8;r9xFp%l%~tW7z*5oTjj7YF2SKA!<AtTRPCgmykj5d|Ue?6TOxuZBSn*!*DPmG@YW0V2|s1iZCtJ zCimp0UsJo-k9Pyd=}ydZ{2s60waHOoUbOejA`IWwB9+l;8m?ok&|uquqcK~@=r)-Yi{)Vc>GphgzV9Q z`>Be|AA;r>VgprTdiX3swWYF>stoaz!N>CEO*`gWGf}bly?d4kMJ?9O2tVG*2CL(K z@E-5iizjiFwmguHVkl!@uM9s|)EaaKH04J6@27uqG{q)>d(_E@M5y8uK2aKD zUGeNzs}AgA5vn9&l+DTy7C43;ttgy5hRh}rB)DnxbU1+SU<@N+ar{`zW}1u*g2qog z5C*W5DY}bNY*bs+i>|8J8+zX&Yt}>S?bA%dIXgNnCReREVkHzYOj>MoUqvi+$djcF zJ9m8y`aXY#>B8BC_wh~Qk3MFGylJ4#drN!@%DXb8KY|#THl%@ZzYy#^%aa1_hWF(j zdu@EB*HLh145v#HSxTVpj$BoDd_~B+lU3DmwjDS0IPIQK4i!)4|8zNFG?{l>13eF? z0hO`w5ddN#r?@87hW0Z|=Fg0?Ajl6_s!mtY+Yf(C2lDwH&p9TF?td<H`EE@Lamz*u@3^s}b@X)a#9b+*^x-sOj_B2r2g9UREnju0a6Y)}KYvat2|3M_lx zdEdRM1y{C={H7(u-CU5hZ`tWfhd&tB2?1Uwuf`>|9Q6ArT71fEeN(JQKA%VRFIKDF z-v@tbss$T`a3+@olB*AoQ0ELY(QKxZc@5wWeS_8*1q3*3gr=dN2Xg z?H$F0r_{&(rCYEO{iXZMNgI*dVu=$sP+%kSZTJKV^O3r=E8F|Tyf8{oBXZQ_O(SxK z)Jr3BQ=v5)kyAKa1sDodRscecDHuT80S1>XaRQ=$FudZ4s zA7#x9jNH_Cw+1}WncgX;d@)#f)^N_`ioEg|0?tSJg45_7hb_o4=E8B??6*?@4 zGyX<=Ik&`luUWHPd;JP}YUEk)jNb@p#@kZ6(Z0PUY0yb%1|n&$;RXoUuC$TGHL&f= zxGUa&Rvz3mJs?i@YM8S1`b7<}(Z0gvvq2}rj2LHsIX5e=8zUjNe4l$eOVZJ%85zQ6 z-ZhPH1NK$VDi3Wprmr#Plyp-t=p7-=wHJmge&CkMh{5HUOA<%EZ<`R;g~#&&Fi29V zNOBV3{g_4mRufaro22QN96Q?ul(}<_=@ip{hA=uB%Y_-99}G-G?ilp{x3$}#9|h8N zWOF5|kOhF;F?sN9F7XC$HXi=)84T`ul8|Ya^vX@4#=SqEm7g@<={PCZQ}&)~Wno>U z1uWw)e38}iYC6C8e*SlOB_43iJ?~(0Oz4_S&2DLrfRjc2ieZQnk@7zKhSNb&$In)O zaN%-0u*e(=)#3~Q6TC&Id%iHcC`WnZ)KPtjECQwn4Uh&|p^`YP;Fb`qNf z)+1e;hS7m{09EMMDncX8$xi{YNwycDC8ol=NpOgN?Laaq zMXVWRIL7+}-^Qmrsf%o;ID^DNnlGQ(0M`&mQn#3WOOqsXiSZhNhT9D>mI+gKrf{vw z7vRIi0mBy39<~PPy_d^(#T`jP*Wp=-PGNyod_y-tIogZdzF9X)DN+s}3|5g%ggR6! zl6-a>V5YlSG7u9F-vz_yG(6mYpg2hD6p(IvM_`KWZ5x1RY*IkTg*cGM`>)=H$g|6B z(jyC;tcZQ7PZq2--_YGYddNUT38c8$Yl+|qHDZPqbCf6Cgk2IV%@rT)6cT{bSGiEa zEethvq+R1W)SXBM^?C<^o||1B^fr??7W}IS^CmQdCd7znW~VOG4z%fiNWx5cyGD&d zM3bxUC*eX4=S#v7G(f$mlKImS6mD@Bt1(s*4$n-kJ(xWy4+rld2Ju^u(qVu3WYQ&k z2y+N_E5vE4hoy%ArH$5;(79^1pF~|yTBVcI2qtTMglM6-*+fD5MJZ&^gAmb3E`iMv zA)+Jt2jfX-dU>YQ%}>>T_&g8To=R}R!YUkLKnb;$o}VND%T7Q5gtS2S((`0yChHYhbm1>BCNY#+0)S|*8i@GU-lat0%xj@05UzOrsq*O` zh&0NvB9v~yJ~{%U!3QF+H;0T?;CtX|Zt3N0!qU35*9rfBrQ1pYbQ#8?xd%K~61uMB z#y`li7!{Z_c{vM=!;bWK3rt3-9mB+Jpm;`910BclOIZTKvMj67q(Fd=A;m5d;AimN<~BTL&6v`4d+T zfO(VUqVcbPaQ9ct&pL|vjKGi65+WoxRegd8=W4)0)$=oTK0p-b`%88)9zkBAPySyv zp`DSWy2lg(KRbq3C@Rba1Hvm*kO46MuD?E*TFgqQ-WA>`JphS!g%ODx{D?~?^yC5~ z8VEeBN3!>OO?*@uIW^|jrwvB6`1((HRV1)*kvEusQn>E=6$=!nZ&!x+O4vTJ^gOKi zGs0~u`Ko_i@9!}A(MG5ruFN7=^}Ah=$XWqP2NjZSRzARy)#K^wYAdbY3;7$YMh2-e zsli})RYXHh%)XqS)q7%gO^&H~MpGgnW{c3@V!z?W%&%DH@A5Opc>L|7G&2wxe<$&m zzpbW!A0hz2%unf>`L^Zub`DuLUG^u8jx=X4v9fTxHRZpQ^(!P?SFQ|Xs0nua6}EK{ zl<$Q(l&>G=8JB(#SN5T00_tZ2^vRZ%rt+`C1If1#fcX@a>+u4pxK{-_bEEwG#)H} z6q#W$vu7iUnLX|a%pjS{^x=Yy6(dMsw>WCkh8iWz)k0`7JKV5k@_#SZqY!hA zVxret0wqLEhPCg=2G`A|oSNmUU9*RO^m@M9#ghEXU~|x3GXp`M-A>4uFp$%US2HH` z)r*WrB7K?lGM0qwFH^`EK7)HhU?<=4Gm-i+D(0PASgbphoY!dN7LZko$%KVCk5e24 zYZGKChC5*QXM;hvi0&ixD73v<9;!>=Xxw?1hhhn~sCxvD>B0gZE#->>USJh}$OHbG z5%eyT*-a;RvK+QPW+>o`s*H@U2i4UFPVm01zuK}Od4@Et1e~|C_L#Sd3HrooJ0c{& z7{T{uylJ8!re7H9mQwbYi!YygRvyr7qI`1gy!<2$?0ceWlxu$roTbB{eROy1^#O&A zV$l9nrVxcydwfl#qv7i%R_!l;G9!5F?8E{6V+O2gk>C_sAa4Pp{T9Pr?S`L#%51g? zxC)MTg6Hwz9K_BhIlzyj+Hq`H?nnJ2I09^9G2i$xw^PLuQv4u*0b|<%*O=#dDIdhn ztLru28d&1z+3&zFg)*o1~0T>dzl?JPdR08!YU{14jTh=G010Dx;uAg-n%!5W-! delta 8700 zcmV=0w=RDJg%Khu;T1<8s>mylgJGDs_`PaS{`lDbOb@u6dY}}+9Ti0;FHupU zM&J@DPOHd>C(Baz!0$O6KI1Rm+Hbwz3aa36S)L~0qeHJ*3X zNH2s>Mifp*Dmh|WisNxB^h-B+N0$3ykE-}#VusZ*Ie@x0olK_aU6d!bkx)4)lnPC! zn)L2InUOZF1sgBk^H|W0QHi94g;t zl>dCm-h9=78|=b9^JwF3_Bgy%3*@0>q#vTW*y1xddbsvedkFym~=ouDjz+?4u1e!B~#aqD8LTD z&mUX%Yrl=dB{O<*%sGJ&wjpVC z#gTeLf5()#n!ib7*o*hkf&XJ)VuMGjrp|n11t)e%?CM)GyrQ2fh!-0QyPlGkd8yMm znMaTbr26Ko4a9wWIe#_+!Jf>6fG=np02O0UJx<@g$wsb8?5hAEGrM$b!josez`$>$ z*~CanzY|LJXwfLRdTu@m+4u$ zZcx!KGOSPMyPk=SOD8Ukpcdy~e-WfGATVp)}0{M1X9Hb^y`xS!B z+4zh-+SBVO4Sy;Z1-8yz%U{-<9wNC*lMC$tTPRwer1v_9ibiVue#tC8kx~yyL+!0>8H;!ao zua@b=9mv0d@rt*-wlEN83V#!#p&3==lGG3X#OeLGeg<(iHqHB)QECTF2!|%Aq>vos zv%n(>{CZ4)Oe3K0Q8;JIewQ|FP{Br=VIn!hf=-vnn=Nm(VhIsYwKMP}+2m;gEf?|B ztLrv{4uAOrufni`razq_jnVu>MJ)!JPc?!tFBHlxm(+qf%Vq&O!z^Ed;sMl9>-1;3 zkVGQ1r`nvC6dbO?0KSSJakJtC6!GVb!3tvzLQL~E%*~~!C{&xQdal_W#c`EHbGTe5 z0L`_Q64r%+-~GY=iq`|R%~B$TgTnO<48GU2X>cI>$}o}j0w3e)y2$? zsekLerCU05FEy-i=eI1di`VDQyju;jN+?sP;wY`Z&3~sw&w#tsjnGu}PlDfOyJ}w8 zy?=dFF8N$>%jW5y>KfNrxCxD7MzWw516)D=9`f4nWe*ILm%=#HYEjb4i-#Fx&i29~ zIZY7F!5XANi4^+IuBLJX9SPEN;+;rzTW{VTbkt!R?yi9pZs|eFQ$5tI;{n?BbAUI5 zz!}db?S!`0EEFu;z$7|mWt}j^6=pbebALyN{Rb~FCO3E<5PC+wOnvkr$7f)zQDs~B z*T9_vlgL7s5900$ta&R8vwkii(}!D&bC1r&9mRl;>#=i0@5F? zsJm-kMUPh3krfv*XB)#iWAoYatbgwxut<0Z_YuJLOJyS?&{h0242;YDg5z2-Bjzt= z#gropXcL#wAp{D7w5PX!7i2Gt%{x$-=6DSy^cR}f#;pg+m{@o&baMV3)S2WU<#_#n z>2}XhOpI_X3lk%3Kj`XkKN@2LX;5I@c{ZX2z~E}vTk z3M1%iU@pCKAB9nPQS?6@7Gk6mg;Z@fExfQ}J!5%|8)ocJu)$u|#_OXovtFzoQ3@mI z=!jBS_=^;x0T)J4g>{^ERluK!3b36dFCu*VP)!~=B63xC}uP%$}+gv8iihz~Z@RzY3h zu~0qHP%%dwcYuJ4u>m1@M^honc0iaXM6r+$TZ}Fx+2ydNO&q|H;C~FcbglnM=Ou41 z;g*;qP$h?0A5Iq%wkQLdv7oSby-g~TYzp*zc;3R|Vl3vc#o71FZO2s7{HG&#p^@^D zr2HfmZ(@bmV$aP5pa|?h>rLQ_X7tUMAb8tJU)xT?HDdq(Xjvs-HZv+<2$dkFG*Dbp z#e|qIrN^fM0$ro;xqp(EiSWfPG2##8jt{D#CfkCRyoK0#hA{f747NFNR^`(m>o`3d zDBz$i$>!tecypG~&G6&CiVcX4-LtoZyv*c&3fcf_?@7a#y;)lvK6Uy7u*+h~ZYRl3 z|Kiy!{6P8QiUNKOFNFIbf3{B*2eiAr2gWa>O=q9u!>?M76n{$^0QCW@G-S`JpX30_ z7Q&_8fNQ9PkckbLd?&zaQ+VhUlZ%%1cZg>YeoNHy+Tp) zc)lclUXrHCMSnW2hOYaQ1E_g^zk(c1+)FF0N(Fq2FuRAKF@f8*Et#-Vkw)Tg)E0R2 zVY)H$#Swi>PlirL+=Oz%{y#VQ!Woj^?$y{m^!tHM69LDBFop zf#C2jZY&s=Sidjyh@J_ z-pmKdynXqTk0bMDnkbn!K7%InJ`w}Uybtrm&7LRo?%1<3?;o@*^A351%e*w1xjJtDLsdn;37S8V2bmsW&PBWch8r! z6*t+*P8ivbX$KjvQCd~a+1XBuNq-Wkj<9-0qo%IkoS2CBEJU_Ye?JAQ%}n)BeMUSy zXm*Nh+lxM`c|o$VVlm~SZrfK>N`$!mN{}pC_+FL768691@wmb8&RaMNwbznGox}QW z&mkAJlZW8Jgszt}F*8lroR-T$o035FE9SLs{7`Ua1Nuw`CP#i+}4<*%qcW zK+Dbhi6`=WlN>=)8Q;plUP|?Ma0WV-MWD-7wA7|=0|4UQIPuKNHJmk1ig*#!ufW4^ z5g+IPY_Lt(lkm7Gx@jW=a6CbJPea+)E&)?XtnkBoBluchfvuMu>EsRa06fhaxc1)n zW;iZNBvpaN7@)fmGsZwAvwv#}MyIla7-OK1M9iyCiSy%3E*n~|>R%UpDwEa3CcFk@ zC**>;YHa(6V^ZqhFmgd}M5Tlt1a`wVM(r|JUFoYlh`q~C>EFpsH|w|Dco1ErWO5h6 zN>X7)7&c%ZL>cnKQBWfBPdp1kM!-_W|>al>5)lg1B$C;-3$}nB0u_jtH^8n4Dw}aTpD~v3oIj2^_qboPy8+)|~Q0 zrPjqFZt+DK2W*!3Hwwkuom~8v*svC zAuL0ch$bT%R#+=HXn#wobhfz2%W#T<(iiKIagQ*acRhP#q*!NC`J#NK;ks!l(xXwz zH~k16Dn|~}@kGS}O<5jH^(^y}ELI7x!_mv(HpEWpK1)bN;-3Nb<%U&`!8QIM1CZEO zVslKJ0CYQ%H^N@!Nw~)wa~>t?!9irUDeocqykiDXo{`EOvVY@tW*_HD!=yYTuHtQs z9e3N~+|ef+wp19hD#z0;e=hE5b=z$v8FFI;CO6NEVn@v0RjqtRDPmoIa5W;LIe7F} zlCpp{<yN8JQb zh8$pm;RpU87V;BV^ih$_&s{&woCxvOvW4yJorYDsR70R|NA| zF*g4nhE~b}ATaKTHXGfm0Xn626v#&%p#Bgf*ke6nU<+B`M?6vf3(NA@W%O?|ns5;` z3iv^BG8v}62BSi}Q$s?CGlkTVi*pd$OP0#{z4F0$)E~scGcOjI(|he&mmiu}>rfNJs7>wgZ&>{3z!HDEXU$kl3>Ul-$tj*tspMED0lc ziGMEd>VtNi0os%evXT3`5p^;%hNs&`rJQ}c`0!iu)U@ol^?QL-Yupz;*~^C#M9N&Q(f>x#@^Ouqx;yH9$-?Nx ziczV-W^G{#z!h#bNeBfh9miC_2Pe6fX@9Vhm%gl_Zs{U7A+o%{{yt&W#d!X4CkeX= z3uq*SZa_Sx?fGg7P%_sILuWaz>9n7z7H>xLwfowyT%MhVm7(v0>GC}c1-8d+PFNV~ z)E#6%RM=G}Rh|*RCy6tf*SuJ5Z_VJHo)(@OLxUT(e%$7uMA2;-9?e0pQK=H;pML?4 zyrXAGTca^oXWQ1lADnntmL4HfqQwL)MZ!5afx!MtHw>^s1tjoD8FgSVxwgNUyCNUW zzAqO}K0WJ@#v$8a94Jg0^O_^2IPk=N!M#(16}O;X#i~MDsH0vr=xXo{|B+bDm99?b z8Bc|wlqgP9>)4D@qluoECVfy}C`rk1FcX?i5sP5Y>3WJVBh@DN=BHmnyV#F?1II~E z%yV2I%+~NvNRy=V%lS@uP|O=Xz-f9Jm#c{p%fAgjt23ek{SnR!fh6zP2w$2DU-pNL*6MpdC_v=NHxJs)%kcwiI zv9AXVKUdWnv<5e&qp#`PTNLlBe{wX-CXhYq*oZ`^{hGZ*@qCSAYznF z$Yz~)~FcAn))fp)|Da*y3M-qPzRxHAUdC5bFV zpzelT)jpm@$h&H)YDC+PGW7Tw&nKgbC-;B4oG_ZqyRCtqhtq({*!Tzlv5-?-lWIfz z87A{*##s>LhbvX5sp#znf2ITR{Env_6Gitw7uFXr_TaxW5J?jKzBre$6=Yy6Jq`L< z)BZFUuiiS_>u~S#!&VU~g+&Jk^0ZS4kqsNvo~`U}J)Q;2o_F4N&uXEREhBc)5^OLR z0^7H;`e5M?v+HgHuak#xiLJHvI|MC0uUp^b<&n?8qx#pm)tmQ0f0_b<4PkI5FM;I7 z36CHs!(=g=>0~Y=aEC8J>!SjK9CAMXUZ3QR6QecsWVD`|fa%tbV!{*kv482dY(#(Q z&gGc<#8Uo^NnC`KPu&kIBxwRo2eIItL> zDJP^Mi3$8M1TyAz78*ITwKp94%tKgkOZDcE^zm-I4?`++GU3gqjre6Q%Xx25jZ%A^ z74&o9S>_qs2r2K&Qn=AhxUFi?i)faDq+N$Z5wNv=8;fgT|Cf1Jynll~(6s#@9<9|# zW$S@NZLrY}!IiN=&&J#~u6{FaDz2Mh0XLJ+9X12j(N2zFge~rrmTP1GRfka?k~F5R zvBfLX*f5xSBATl%3=s{v4UiFI!x3SUIG(;FLhu@grw(8MP@#&RlYsjRi~L(mGM=wV z(@!aOwF_W#=bDo#rhg4#R5X4UW?((2ndW)&XS8oCb$>n%q!XXrjVPfOKyw%J;5VhD z2)>3a{P!!^+Vg~Evo7_@&B9H4Up{kxNxsE#B-hmS<_@bu^N4F$PG9)4tL2$=fbfkR z;DAcJ>l=ICL8X|`H5pso(jEa1i~2do4m~1fef1TigW_(VEq|fH<@aBaJrb%@C3|C_ zaSasKCV5YU_BazmrPRy5YTutDsn(|$$PN&Dyr~3E*yYS5HWR2vI`<5*31uK5^dX_Ut!b7^E(NzxBhkBf2z%>jA>M zUqykCuGbiU5Q*(zKpYw)K3fK`U^Pxl1|zt38C{Qf+hNHWzsE?uLiawMx(>k9)wTTq zhCsk~s7^JpRa^s9rCqBCjU+ESCD0<#Vt$sO25%-pp?|T%>68SpW))!!_X~VtSk;Lx z(uqPD3=Yz}d{qNpLfFZjV)iXfG>sFJ@Q71Rw|`kCOzDxrwJL9T50}RaTTFS_8ld@I zF5wh+ECpSM=p@?2h{qiBLO9l=47ueekat38Cnc7p77IkNvt$ifUZ+W08U@+J_%Pa)Rd5Rt?DfIKq9>B z)dYHOc6rd-OyXSbuOi%&(CnCzBBF<#x=gEIXMZCZwB+p?wh0kUuEL#!3pu1K2}jTX zWx+b;SwA4X#XaSwSTh`cG`sac;L3QvzJ~MxP(3mT{bkIiO8A894(wKl(^MEw4M|HI zE+^q~)o?wDT~t`*A*UTq)qaF%Lov3A0{TT+WKcv9p-Aq5&5|I3CjO)0Bs9G|Rq5uZ zYJYqm2W(FzxNun&j&Oj4T1(GQ5`cvlAOOQzpqp8F*PQQx3gEKU%4_sP0q^hi0H~DL znr|QDJm4N%qxT9%5GpKHkXYD(D&avX9OPw-KEo*w!{6j~lwUr!wYKSPbW+ydS$v6d z`mDV_ikAjiRjs|>-P>;O&UNPw&;B~3aDVLfPI=bRtSWNDMyb7bfo&{#`F)HZ6Fi~r zC6cjFgrR$dJpv3FX>(yU86Ks?mJQSF9aay`tNP5_53Qg>uU|4QBI+Bi$Dt*4@Ujyv zwN~@-V;qBvoyjPZY0a&RIrryjhqAtX90RTu(Zykj18zN_-MH5%?{rksAJl7K;G=oiaG#W1~P7jF7kRW-%TUzvSTDr=U7qR z*Q<|@`;&|1`^?oP6qf)X+N%a4K9+YmNjURbs0oBC-+rom`UfJ7a;ylYTd@C*zi99& z5!jT2CM*3xZgsa1aW+wE-SO%Kqkqy>r68&drP$n|JQpImDai#q$buOan9sbNOT*zu z`hVm}q|o7ST>*{mhTg2USF1($^~DQi(c=gOD}Fmti#B72VQ0YXn=fpOWAF=b@A}$^ z0RZ#+>DhMnyL~kL%DE1BBZ+)=)6Wlz!YgcXAgxsSo~t1+WN%kZk0hMt$bHXWe{ zz3g3zso6vmQaBtCqrRpD=>rls&=9Qtal8|5DhlkZ@nPGN73z*lt(Y)`2j-3nnqXYM2*X+C$vU zmyQW2U=7eGUs#@MaxlV1g$9%%2K!jutyw0j@hGJ&-tfJqc@ma<3i z7WXiYvS=R&ibL}$Wq)0LCrH#%7K`KwmOabFBhS0TRurTRkoXXod|ugOd~wSK^8KS? z*^9!^=NLPe`)MyMm+v}99U3pQwEd&gmU`NVWEY!LZ6+hQQx#qJ*2O5Q(=Y;2**hV+ z43xanHxqr7Wf-b#yZ}oAA+viY-oeNvi%KE0KLzj8G`wj%SbyL$LnCHSMHn%A=qbig zDizpA2_q+dNo2K@Xw#`114N+H{>LxVF%G2n}1CR-ZwpPKxfQ=JMDiwLp%sPIJe)@%cEWV8%$Z)F5$HUqgDAlQ78v7TIB{< z=caa&ExTP=`v^$@n^?{^zO3!kafBQ@2w*^&cEFmmJTK;ht^;E`%g+F`e)W13t<9P+ znB;=-vAOjJ7Hk5A Result< total_tokens: 150, ..Default::default() }, + last_usage_includes_reasoning: None, last_token_usage: TokenUsage { input_tokens: 70, output_tokens: 20, @@ -4361,6 +4362,7 @@ async fn cold_paginated_resume_omits_usage_when_its_turn_is_ambiguous() -> Resul total_tokens: 150, ..Default::default() }, + last_usage_includes_reasoning: None, last_token_usage: TokenUsage { total_tokens: 90, ..Default::default() @@ -4645,6 +4647,7 @@ async fn thread_resume_token_usage_replay_can_belong_to_interrupted_turn() -> Re total_tokens: 230, codex_rollout_budget_units: None, }, + last_usage_includes_reasoning: None, last_token_usage: TokenUsage { input_tokens: 90, cached_input_tokens: 30, diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 3c6642301f..2f6a91de70 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -830,8 +830,7 @@ async fn drain_to_completed( prompt: &Prompt, phase: CompactionPhase, ) -> CodexResult { - sess.set_server_reasoning_included(turn_context.model_info().use_responses_lite) - .await; + let mut includes_reasoning = turn_context.model_info().use_responses_lite; let mut stream = client_session .stream( prompt, @@ -862,7 +861,7 @@ async fn drain_to_completed( // Keep each attempt isolated until the response completes successfully. Ok(ResponseEvent::OutputItemDone(item)) => output.push(item), Ok(ResponseEvent::ServerReasoningIncluded(included)) => { - sess.set_server_reasoning_included(included).await; + includes_reasoning = included; } Ok(ResponseEvent::RateLimits(snapshot)) => { sess.update_rate_limits(turn_context, snapshot).await; @@ -896,8 +895,12 @@ async fn drain_to_completed( usage_metadata.as_ref(), ) .await; - sess.update_token_usage_info(turn_context, token_usage.as_ref()) - .await?; + sess.update_token_usage_info( + turn_context, + token_usage.as_ref(), + includes_reasoning, + ) + .await?; return Ok(CompactionResponse { response_id, output, diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 175f26e6ea..12fff896a8 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -246,9 +246,7 @@ impl ContextManager { reset_version: 0, user_message_revision: 0, guardian_review_context_revision: next_guardian_review_context_revision(), - token_info: TokenUsageInfo::new_or_append( - &None, &None, /*model_context_window*/ None, - ), + token_info: None, reference_context_item: None, world_state_baseline: None, } @@ -903,12 +901,14 @@ impl ContextManager { &mut self, usage: &TokenUsage, model_context_window: Option, + includes_reasoning: bool, ) { - self.token_info = TokenUsageInfo::new_or_append( + self.token_info = Some(TokenUsageInfo::new_or_append( &self.token_info, - &Some(usage.clone()), + usage, + includes_reasoning, model_context_window, - ); + )); } fn get_non_last_reasoning_items_tokens(&self) -> i64 { @@ -950,9 +950,7 @@ impl ContextManager { self.items[start..].iter().map(|envelope| &envelope.item) } - /// When true, the server already accounted for past reasoning tokens and - /// the client should not re-estimate them. - pub(crate) fn get_total_token_usage(&self, server_reasoning_included: bool) -> i64 { + pub(crate) fn get_total_token_usage(&self) -> i64 { let last_tokens = self .token_info .as_ref() @@ -962,7 +960,11 @@ impl ContextManager { .items_after_last_model_generated_item() .map(estimate_item_token_count) .fold(0i64, i64::saturating_add); - if server_reasoning_included { + if self + .token_info + .as_ref() + .is_some_and(|info| info.last_usage_includes_reasoning == Some(true)) + { last_tokens.saturating_add(items_after_last_model_generated_tokens) } else { last_tokens diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index b5eda994c8..7e80afe42b 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -937,6 +937,33 @@ fn non_last_reasoning_tokens_ignore_entries_after_last_user() { assert_eq!(history.get_non_last_reasoning_items_tokens(), 32); } +#[test_case(None, 125; "legacy count")] +#[test_case(Some(false), 125; "exclusive count")] +#[test_case(Some(true), 100; "inclusive count")] +fn total_token_usage_respects_saved_reasoning_accounting( + includes_reasoning: Option, + expected: i64, +) { + let mut history = create_history_with_items(vec![ + reasoning_with_encrypted_content(/*len*/ 1_000), + user_msg("second"), + assistant_msg("counted reply"), + ]); + let usage = TokenUsage { + total_tokens: 100, + ..Default::default() + }; + let mut saved = json!({ "last_token_usage": usage, "total_token_usage": usage }); + if let Some(included) = includes_reasoning { + saved["last_usage_includes_reasoning"] = json!(included); + } + history.set_token_info(Some( + serde_json::from_value(saved).expect("saved token count"), + )); + + assert_eq!(history.get_total_token_usage(), expected); +} + #[test] fn items_after_last_model_generated_tokens_include_user_and_tool_output() { let history = create_history_with_items(vec![ @@ -1160,6 +1187,7 @@ fn total_token_usage_includes_all_items_after_last_model_generated_item() { ..Default::default() }, /*model_context_window*/ None, + /*includes_reasoning*/ true, ); let added_user = user_msg("new user message"); let added_tool_output = custom_tool_call_output("tool-tail", "new tool output"); @@ -1169,7 +1197,7 @@ fn total_token_usage_includes_all_items_after_last_model_generated_item() { ); assert_eq!( - history.get_total_token_usage(/*server_reasoning_included*/ true), + history.get_total_token_usage(), 100 + estimate_item_token_count(&added_user) + estimate_item_token_count(&added_tool_output) ); diff --git a/codex-rs/core/src/session/async_compaction_tests.rs b/codex-rs/core/src/session/async_compaction_tests.rs index e4a048124b..d73b96c657 100644 --- a/codex-rs/core/src/session/async_compaction_tests.rs +++ b/codex-rs/core/src/session/async_compaction_tests.rs @@ -65,6 +65,7 @@ async fn cancelled_turn_skips_launch_but_preserves_existing_work(trigger: Compac total_tokens: 70_000, ..Default::default() }), + /*includes_reasoning*/ false, ) .await .expect("record usage"); diff --git a/codex-rs/core/src/session/guardian_checkpoint_tests.rs b/codex-rs/core/src/session/guardian_checkpoint_tests.rs index 6ff5ab3589..c64bf40429 100644 --- a/codex-rs/core/src/session/guardian_checkpoint_tests.rs +++ b/codex-rs/core/src/session/guardian_checkpoint_tests.rs @@ -78,6 +78,7 @@ async fn guardian_checkpoint_preserves_live_context_without_storage( ..Default::default() }, Some(32_000), + /*includes_reasoning*/ true, ); } let expected = session.clone_history().await; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index b8f3a0f078..cc418ca5ed 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -417,7 +417,7 @@ use codex_utils_stream_parser::ProposedPlanSegment; /// Whether reported usage updates the installed context baseline or only spending. #[derive(Clone, Copy)] pub(crate) enum TokenUsageContext { - Sampling, + Sampling { includes_reasoning: bool }, SpendingOnly, } @@ -1569,7 +1569,7 @@ impl Session { pub(crate) async fn get_total_token_usage(&self) -> i64 { let state = self.state.lock().await; - state.get_total_token_usage(state.server_reasoning_included()) + state.get_total_token_usage() } pub(crate) async fn auto_compact_window_snapshot(&self) -> AutoCompactWindowSnapshot { @@ -5244,13 +5244,14 @@ impl Session { &self, turn_context: &TurnContext, token_usage: Option<&TokenUsage>, + includes_reasoning: bool, ) -> CodexResult<()> { let result = self .record_token_usage_info( turn_context, &turn_context.initial_settings, token_usage, - TokenUsageContext::Sampling, + TokenUsageContext::Sampling { includes_reasoning }, ) .await; self.send_token_count_event(turn_context).await; @@ -5312,10 +5313,11 @@ impl Session { let token_info = { let mut state = self.state.lock().await; match context { - TokenUsageContext::Sampling => { + TokenUsageContext::Sampling { includes_reasoning } => { state.update_token_info_from_usage( token_usage, turn_context.model_context_window(), + includes_reasoning, ); if matches!( turn_context.config.model_auto_compact_token_limit_scope, @@ -5328,6 +5330,7 @@ impl Session { TokenUsageContext::SpendingOnly => { let mut info = state.token_info().unwrap_or(TokenUsageInfo { total_token_usage: TokenUsage::default(), + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window: turn_context.model_context_window(), }); @@ -5336,6 +5339,7 @@ impl Session { // Contributors receive the reported usage; stored context stays unchanged. info.last_token_usage = token_usage.clone(); + info.last_usage_includes_reasoning = None; info.model_context_window = turn_context .model_context_window() .or(info.model_context_window); @@ -5384,6 +5388,7 @@ impl Session { let mut state = self.state.lock().await; let mut info = state.token_info().unwrap_or(TokenUsageInfo { total_token_usage: TokenUsage::default(), + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window: None, }); @@ -5402,9 +5407,9 @@ impl Session { info.model_context_window = Some(model_context_window); } - state.set_token_info(Some(info)); // The local estimate already includes every retained reasoning item. - state.set_server_reasoning_included(/*included*/ true); + info.last_usage_includes_reasoning = Some(true); + state.set_token_info(Some(info)); } self.set_auto_compact_window_estimated_prefill_for_scope( turn_context, @@ -5443,11 +5448,6 @@ impl Session { state.record_mcp_dependency_prompted(names); } - pub(crate) async fn set_server_reasoning_included(&self, included: bool) { - let mut state = self.state.lock().await; - state.set_server_reasoning_included(included); - } - pub(crate) async fn send_token_count_event(&self, turn_context: &TurnContext) { let (info, rate_limits) = { let state = self.state.lock().await; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 579d109872..68022ae9d2 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3420,6 +3420,7 @@ async fn record_initial_history_seeds_token_info_from_rollout() { total_tokens: 30, codex_rollout_budget_units: None, }, + last_usage_includes_reasoning: None, last_token_usage: TokenUsage { input_tokens: 3, cached_input_tokens: 0, @@ -3441,6 +3442,7 @@ async fn record_initial_history_seeds_token_info_from_rollout() { total_tokens: 375, codex_rollout_budget_units: None, }, + last_usage_includes_reasoning: None, last_token_usage: TokenUsage { input_tokens: 10, cached_input_tokens: 0, @@ -3590,6 +3592,7 @@ async fn recompute_token_usage_updates_model_context_window() { let mut state = session.state.lock().await; state.set_token_info(Some(TokenUsageInfo { total_token_usage: TokenUsage::default(), + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window: Some(258_400), })); @@ -3606,7 +3609,7 @@ async fn recompute_token_usage_updates_model_context_window() { assert_eq!(actual.model_context_window, Some(128_000)); } -#[test_case(crate::session::TokenUsageContext::Sampling; "sampling")] +#[test_case(crate::session::TokenUsageContext::Sampling { includes_reasoning: false }; "sampling")] #[test_case(crate::session::TokenUsageContext::SpendingOnly; "spending_only")] #[tokio::test] async fn record_token_usage_info_notifies_extension_contributors( @@ -3693,7 +3696,9 @@ async fn record_token_usage_info_notifies_extension_contributors( &turn_context, &turn_context.initial_settings, Some(&first_usage), - crate::session::TokenUsageContext::Sampling, + crate::session::TokenUsageContext::Sampling { + includes_reasoning: true, + }, ) .await .expect("first usage should be recorded"); @@ -3711,8 +3716,12 @@ async fn record_token_usage_info_notifies_extension_contributors( expected_total_usage.add_assign(&second_usage); let expected_session_info = TokenUsageInfo { total_token_usage: expected_total_usage.clone(), + last_usage_includes_reasoning: Some(matches!( + context, + crate::session::TokenUsageContext::SpendingOnly + )), last_token_usage: match context { - crate::session::TokenUsageContext::Sampling => second_usage.clone(), + crate::session::TokenUsageContext::Sampling { .. } => second_usage.clone(), crate::session::TokenUsageContext::SpendingOnly => first_usage.clone(), }, model_context_window: turn_context.model_context_window(), @@ -3724,6 +3733,7 @@ async fn record_token_usage_info_notifies_extension_contributors( turn_level_id: turn_context.sub_id.clone(), token_usage: TokenUsageInfo { total_token_usage: first_usage.clone(), + last_usage_includes_reasoning: Some(true), last_token_usage: first_usage, model_context_window: turn_context.model_context_window(), }, @@ -3736,6 +3746,12 @@ async fn record_token_usage_info_notifies_extension_contributors( turn_level_id: turn_context.sub_id.clone(), token_usage: TokenUsageInfo { total_token_usage: expected_total_usage, + last_usage_includes_reasoning: match context { + crate::session::TokenUsageContext::Sampling { includes_reasoning } => { + Some(includes_reasoning) + } + crate::session::TokenUsageContext::SpendingOnly => None, + }, last_token_usage: second_usage, model_context_window: turn_context.model_context_window(), }, @@ -12972,6 +12988,7 @@ async fn set_total_token_usage(sess: &Session, total_token_usage: TokenUsage) { let mut state = sess.state.lock().await; state.set_token_info(Some(TokenUsageInfo { total_token_usage, + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window: None, })); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index b56e7ca2aa..e4a092ed65 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -2760,8 +2760,7 @@ async fn try_run_sampling_request( let sampling_timing_guard = turn_context.turn_timing_state.begin_sampling(); let mut response_completed = false; let mut request_token_usage = None; - sess.set_server_reasoning_included(step_context.settings.model_info.use_responses_lite) - .await; + let mut includes_reasoning = step_context.settings.model_info.use_responses_lite; let stream = client_session .stream( &prompt, @@ -3215,7 +3214,7 @@ async fn try_run_sampling_request( .await; } ResponseEvent::ServerReasoningIncluded(included) => { - sess.set_server_reasoning_included(included).await; + includes_reasoning = included; } ResponseEvent::RateLimits(snapshot) => { // Update internal state with latest rate limits, but defer sending until @@ -3267,7 +3266,7 @@ async fn try_run_sampling_request( &turn_context, &step_context.settings, token_usage.as_ref(), - super::TokenUsageContext::Sampling, + super::TokenUsageContext::Sampling { includes_reasoning }, ) .await; should_emit_token_count = true; diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index ebfc5817cb..968f3425fb 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -81,7 +81,6 @@ pub(crate) struct SessionState { pub(crate) history_reset: CancellationToken, pub(crate) latest_rate_limits: Option, pub(crate) latest_token_usage_record: Option, - pub(crate) server_reasoning_included: bool, pub(crate) mcp_dependency_prompted: HashSet, pub(crate) additional_context: AdditionalContextStore, /// Settings used by the latest regular user turn, used for turn-to-turn @@ -136,7 +135,6 @@ impl SessionState { history_reset: CancellationToken::new(), latest_rate_limits: None, latest_token_usage_record: None, - server_reasoning_included: false, mcp_dependency_prompted: HashSet::new(), additional_context: AdditionalContextStore::default(), previous_turn_settings: None, @@ -296,8 +294,10 @@ impl SessionState { &mut self, usage: &TokenUsage, model_context_window: Option, + includes_reasoning: bool, ) { - self.history.update_token_info(usage, model_context_window); + self.history + .update_token_info(usage, model_context_window, includes_reasoning); } pub(crate) fn ensure_auto_compact_window_server_prefill_from_usage( @@ -379,17 +379,8 @@ impl SessionState { self.history.set_token_usage_full(context_window); } - pub(crate) fn get_total_token_usage(&self, server_reasoning_included: bool) -> i64 { - self.history - .get_total_token_usage(server_reasoning_included) - } - - pub(crate) fn set_server_reasoning_included(&mut self, included: bool) { - self.server_reasoning_included = included; - } - - pub(crate) fn server_reasoning_included(&self) -> bool { - self.server_reasoning_included + pub(crate) fn get_total_token_usage(&self) -> i64 { + self.history.get_total_token_usage() } pub(crate) fn record_mcp_dependency_prompted(&mut self, names: I) diff --git a/codex-rs/core/src/tasks/mod_tests.rs b/codex-rs/core/src/tasks/mod_tests.rs index ba03d83a15..eb4bc4b5e2 100644 --- a/codex-rs/core/src/tasks/mod_tests.rs +++ b/codex-rs/core/src/tasks/mod_tests.rs @@ -171,6 +171,7 @@ async fn turn_completion_metrics_follow_model_switch(scenario: UsageScenario) { total_tokens: 1_000, ..TokenUsage::default() }), + /*includes_reasoning*/ false, ) .await .expect("earlier session usage should be recorded"); @@ -193,7 +194,11 @@ async fn turn_completion_metrics_follow_model_switch(scenario: UsageScenario) { }; // Local pre-turn compaction records usage with the previous model's context. session - .update_token_usage_info(&previous_context, Some(&compaction_usage)) + .update_token_usage_info( + &previous_context, + Some(&compaction_usage), + /*includes_reasoning*/ false, + ) .await .expect("compaction usage should be recorded"); expected_usage.push((previous_context.model_info().slug.clone(), compaction_usage)); @@ -225,7 +230,9 @@ async fn turn_completion_metrics_follow_model_switch(scenario: UsageScenario) { &turn_context, &response_context.initial_settings, Some(&response_usage), - crate::session::TokenUsageContext::Sampling, + crate::session::TokenUsageContext::Sampling { + includes_reasoning: false, + }, ) .await .expect("response usage should be recorded"); diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 2e192498cb..d20e1dfb26 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -3576,6 +3576,7 @@ async fn token_count_includes_rate_limits_snapshot() { "reasoning_output_tokens": 0, "total_tokens": 123 }, + "last_usage_includes_reasoning": false, // Default model is gpt-5.4 in tests → 95% usable context window "model_context_window": 258400 }, diff --git a/codex-rs/core/tests/suite/compact.rs b/codex-rs/core/tests/suite/compact.rs index b3522a09b2..c480acee0c 100644 --- a/codex-rs/core/tests/suite/compact.rs +++ b/codex-rs/core/tests/suite/compact.rs @@ -5480,8 +5480,18 @@ async fn auto_compact_body_after_prefix_still_caps_at_context_window() { ); } +#[derive(Clone, Copy)] +enum ContinueThread { + Live, + Fork, + Restart, +} + +#[test_case::test_case(ContinueThread::Live; "live")] +#[test_case::test_case(ContinueThread::Fork; "fork")] +#[test_case::test_case(ContinueThread::Restart; "cold resume")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn auto_compact_accounts_for_encrypted_reasoning() { +async fn auto_compact_accounts_for_encrypted_reasoning(continuation: ContinueThread) { skip_if_no_network!(); let server = start_mock_server().await; @@ -5519,26 +5529,53 @@ async fn auto_compact_accounts_for_encrypted_reasoning() { ) .await; - let test = test_codex() - .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) - .with_model_info_override("gpt-5.5", |model| { - model.use_responses_lite = true; - }) - .with_config(|config| { - config.model = Some("gpt-5.5".to_string()); - set_test_compact_prompt(config); - config.model_auto_compact_token_limit = Some(81); - }) + let builder = || { + test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_model_info_override("gpt-5.5", |model| { + model.use_responses_lite = true; + }) + .with_config(|config| { + config.model = Some("gpt-5.5".to_string()); + set_test_compact_prompt(config); + config.model_auto_compact_token_limit = Some(81); + }) + }; + let mut test = builder() .build_with_auto_env(&server) .await .expect("build codex"); - for user in [ - "FIRST_REASONING", - "SECOND_REASONING", - "BELOW_COMPACT_LIMIT", - "AT_COMPACT_LIMIT", - ] { + for user in ["FIRST_REASONING", "SECOND_REASONING"] { + test.submit_text_turn(user).await.expect("submit user turn"); + } + + match continuation { + ContinueThread::Live => {} + ContinueThread::Restart => { + test = builder() + .restart_with_auto_env(&server, &test) + .await + .expect("restart"); + } + ContinueThread::Fork => { + let path = test + .session_configured + .rollout_path + .clone() + .expect("rollout path"); + let options = test.start_thread_options().await; + test.codex.shutdown_and_wait().await.expect("flush rollout"); + test.codex = test + .thread_manager + .fork_legacy_thread(usize::MAX, options, path) + .await + .expect("fork") + .thread; + } + } + + for user in ["BELOW_COMPACT_LIMIT", "AT_COMPACT_LIMIT"] { test.submit_text_turn(user).await.expect("submit user turn"); } diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index cc88125491..8faa045d7e 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -1852,6 +1852,7 @@ impl GoalExtensionHarness { let turn_store = ExtensionData::new(turn_id); let token_usage = TokenUsageInfo { total_token_usage: usage.clone(), + last_usage_includes_reasoning: None, last_token_usage: last_usage.clone(), model_context_window: None, }; diff --git a/codex-rs/external-agent-migration/src/sessions/export.rs b/codex-rs/external-agent-migration/src/sessions/export.rs index d9067e2e4a..d5e7fc2c3a 100644 --- a/codex-rs/external-agent-migration/src/sessions/export.rs +++ b/codex-rs/external-agent-migration/src/sessions/export.rs @@ -200,6 +200,7 @@ fn token_count_item(last_model_visible_tokens: i64) -> RolloutItem { RolloutItem::EventMsg(EventMsg::TokenCount(TokenCountEvent { info: Some(TokenUsageInfo { total_token_usage: usage.clone(), + last_usage_includes_reasoning: None, last_token_usage: usage, model_context_window: None, }), diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 21d45f6b78..8c419c536c 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2506,6 +2506,10 @@ pub struct TokenUsageRecord { pub struct TokenUsageInfo { pub total_token_usage: TokenUsage, pub last_token_usage: TokenUsage, + /// Whether the last context count includes reasoning retained from earlier turns. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub last_usage_includes_reasoning: Option, // TODO(aibrahim): make this not optional #[ts(type = "number | null")] pub model_context_window: Option, @@ -2514,33 +2518,26 @@ pub struct TokenUsageInfo { impl TokenUsageInfo { pub fn new_or_append( info: &Option, - last: &Option, + last: &TokenUsage, + includes_reasoning: bool, model_context_window: Option, - ) -> Option { - if info.is_none() && last.is_none() { - return None; - } - + ) -> Self { let mut info = match info { Some(info) => info.clone(), None => Self { total_token_usage: TokenUsage::default(), + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window, }, }; - if let Some(last) = last { - info.append_last_usage(last); - } + info.total_token_usage.add_assign(last); + info.last_token_usage = last.clone(); + info.last_usage_includes_reasoning = Some(includes_reasoning); if let Some(model_context_window) = model_context_window { info.model_context_window = Some(model_context_window); } - Some(info) - } - - pub fn append_last_usage(&mut self, last: &TokenUsage) { - self.total_token_usage.add_assign(last); - self.last_token_usage = last.clone(); + info } pub fn fill_to_context_window(&mut self, context_window: i64) { @@ -2552,6 +2549,7 @@ impl TokenUsageInfo { total_tokens: context_window, ..TokenUsage::default() }; + self.last_usage_includes_reasoning = None; self.last_token_usage = TokenUsage { total_tokens: delta, ..TokenUsage::default() @@ -2561,6 +2559,7 @@ impl TokenUsageInfo { pub fn full_context_window(context_window: i64) -> Self { let mut info = Self { total_token_usage: TokenUsage::default(), + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window: Some(context_window), }; @@ -6627,10 +6626,11 @@ mod tests { fn token_usage_info_new_or_append_updates_context_window_when_provided() { let initial = Some(TokenUsageInfo { total_token_usage: TokenUsage::default(), + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window: Some(258_400), }); - let last = Some(TokenUsage { + let last = TokenUsage { input_tokens: 10, cached_input_tokens: 0, cache_write_input_tokens: 0, @@ -6638,10 +6638,14 @@ mod tests { reasoning_output_tokens: 0, total_tokens: 10, codex_rollout_budget_units: None, - }); + }; - let info = TokenUsageInfo::new_or_append(&initial, &last, Some(128_000)) - .expect("new_or_append should return info"); + let info = TokenUsageInfo::new_or_append( + &initial, + &last, + /*includes_reasoning*/ true, + Some(128_000), + ); assert_eq!(info.model_context_window, Some(128_000)); } @@ -6650,10 +6654,11 @@ mod tests { fn token_usage_info_new_or_append_preserves_context_window_when_not_provided() { let initial = Some(TokenUsageInfo { total_token_usage: TokenUsage::default(), + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window: Some(258_400), }); - let last = Some(TokenUsage { + let last = TokenUsage { input_tokens: 10, cached_input_tokens: 0, cache_write_input_tokens: 0, @@ -6661,11 +6666,11 @@ mod tests { reasoning_output_tokens: 0, total_tokens: 10, codex_rollout_budget_units: None, - }); + }; - let info = - TokenUsageInfo::new_or_append(&initial, &last, /*model_context_window*/ None) - .expect("new_or_append should return info"); + let info = TokenUsageInfo::new_or_append( + &initial, &last, /*includes_reasoning*/ true, /*model_context_window*/ None, + ); assert_eq!(info.model_context_window, Some(258_400)); } diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 1fc8b1e927..222fc83b3b 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -3756,6 +3756,7 @@ mod tests { total_tokens: 321, codex_rollout_budget_units: None, }, + last_usage_includes_reasoning: None, last_token_usage: codex_protocol::protocol::TokenUsage::default(), model_context_window: None, }),