From 80d7ca34bce7497fa72a13af74695c759a48a8d4 Mon Sep 17 00:00:00 2001 From: felixxia-oai Date: Fri, 4 Sep 2026 13:05:30 +0000 Subject: [PATCH] Retain user instructions in guardian thread context (#42762) ## Why Guardian review needs the original user instructions alongside verified answers, even after compaction, resume, or rollback. A summarized or legacy checkpoint cannot establish a complete instruction history. ## What changed - Capture text user messages in retained context when `guardian_thread_context` is enabled, preserving their message identity through compaction and replay. - Keep retained instructions and verified answers in arrival order, with independent size and completeness tracking for each family. - Roll back retained evidence at the exact user-message boundary, including steers that share a turn ID, and apply the same semantics during rollout migration. - Treat checkpoints without retained user messages as incomplete while preserving their verified answers. ## Testing - Cover capture, compaction, resume, rollback, child forks, legacy checkpoints, storage bounds, and rollout migration. GitOrigin-RevId: 72706731b54e5ac8645db61b446d2401afb42235 --- .../app-server-exports-stable.json.zst | Bin 147820 -> 148006 bytes codex-rs/core/src/agent/control_tests.rs | 7 +- codex-rs/core/src/compact.rs | 28 +- codex-rs/core/src/compact_tests.rs | 17 +- codex-rs/core/src/context_manager/history.rs | 77 +++- .../core/src/context_manager/history_tests.rs | 49 +++ codex-rs/core/src/session/mod.rs | 4 +- .../src/session/rollout_reconstruction.rs | 20 +- .../session/rollout_reconstruction_tests.rs | 1 + codex-rs/core/src/session/session.rs | 3 + .../tests/suite/guardian_retained_context.rs | 369 +++++++++++++++++- .../guardian-context/src/verified_answers.rs | 2 +- codex-rs/history/src/lib.rs | 2 + codex-rs/history/src/retained_context.rs | 219 +++++++++-- .../history/src/retained_context_tests.rs | 112 +++++- codex-rs/rollout/src/lib.rs | 1 + .../local/rollout_migration/rollback_plan.rs | 114 ++++-- .../src/local/rollout_migration_tests.rs | 136 +++++++ 18 files changed, 1058 insertions(+), 103 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 520cb64c9c628cf9fe7ed49a6f97138fc8f9e48b..3b9f9d95f1f3cea4622fa59792808c68e4fee012 100644 GIT binary patch delta 10897 zcmV;CDsI*6f(fRC34nwFv;z1Se_X5pMxIt4FmoHvG?m9dvgetzZgOp6Kpc<>^@Gl^$kDk_L3?TM`{nCe{?A(L zJfk%?TD$Gq&*iGQP;;NjHdzB>6or1ZyNjH2e+tK>&3A?_RUg1AwO28HeKe;U4|{z|Hw#QkNW{MEGmCJ` z&11$WPhFpDI?k-)m61(QC>REX&660fx;~vc5#N?t@?*4E)Shb6M}BOOqlZBpW!RhA znXFu9bTu(?)u#c5d#K~xR!x#L2QwfLfWSaO0!3vqaWoW-1?L%%f5{&sZ~&peGzbRF zU?4CE1_FXWAgBm}AP7Q$umlK#AT)x|5E7xLPGgqa^{>-ICgEi)`I;FM>DXM}BJS*4 zEBs0>WBh8351)Zg-%UujL10Kscnm^o{ zuXObC5ibsJZ2o1O-$H47!-o9T{^rw*@J1r^n?zX^v7O?svX10V=Q`+F9I2=pOBn^0 zfQTk%H7-?Ue|oZgPv+7<*i^xq{1*i(p34zZJcMwF@9^rQibYT(RQc}6AUap8EGZ@D zL@2H+6nKOFJB>mCS6SHRDEt*N>9G?j8XjRz>4)b2`Mqpmo#1~cdob?3WQ3{Ed}h@e z4_}|69ngGf^{j?U?(w*AWA{Oehww&+f?(DuMseLZf3em3^~)u_U=i=Mq`7J)#o^5+ z?EjL!-Lpzt*&!qBgmLoFNyzX{Y10}tXL&6qHG&($yv|}Gnhcdhb1;S$;uh-9sUVSM zfo)J?8D$=9I7Jfe;E$d>s&7=|a`lS@t15FTaoaYKpsxs`@9kw0mHJlzg?kBi-eg9h z_8cv6e}A~R{c|odpJ#9|!}W4~OjeVT(^3tz$ukIEc>r~QX;iLsD2aVoD}6yKw^V~% zZ@=WUFrfm_a8vT6Bu{ur8fhvV#0tnusdb0pK#^q;y0nVsQ8BJp4p&+bV@5Z8c`*C& z=t>$NB)W)LtV;Y`Hcn6C^QNCgm6Wm`3~Jatf4l4v$c_OsGag3ksM>s?J1LWeYB*qd z1fFLNG_x1P1(HOeONl|6!l!By9ieiEcP%eB2lgMZmjc~h;^0Xbg#a|X8(!-i^U9w+Zs~XvWgj|A`1SN_rwg}zRVby zpE@_>SIyKu>W6NG@6?6=*Vw?Mh3O(-#0(E7RJW-OF~hTV->Wr9$Nn(GcwQrX{EpUv zPKW35k(R9E^MpKm^z;&!mEqqOp~1GUf4B}&loRLcVzKX>NJSFE01zc*5?C}o;>XSj zF$%1sxm9D5+8=qK*k_-_8@7hQsY%{x=3!O<>OB82`46H; z@laz&Pxkm7W<>NmWHB1Ll5-m5=6(P;JE~B4oD^rV=TS@pK(Um`EWxG+2GSfPe`xpa zuJ)Pd?32NeI+ja-DS2k>GJa2A^Fuml(gv{xp5$VEV0P?+`t{O=qb&56vaI~Tioy-d zV8?PV+|M(t3SRp*%M59`+mWuTD`1DUFoDWHCL8WJWy7(+xQM*yIOAOC;*YeCCRZF? zJ>M!P8ioz~LOYI)J$=YNi+h&pe_Il0aZ11=Pr_MhDBbIjTJjLF{16)ymLc<$VcD-l@&f7X+_`Uj=y zDW11Kc&uYY8Mk_KK3vYokw^3tnt~Xl77Tp~)v@Y(hRTaycLo=$zqhn-LNm`fmjuTd9uF{Q&$qY|i9>ucRhgyil$53E@SS^rVrCu7VHHxsAtlWgB*#tgL#5UYd`dSCgrn5Q%J%(oJ;P0J?3oznMJ(_zAbkov0ypMXax*(^QxA zs5kmaxa6X{3YN?yV}SQsfY!`Unl^lXfHj?-Hlo8yvdLEVIT1@CWkYZEnKPmyEcZV1 z6lmAC^T~t>ZOboJQrujSA<`nq92&^u6~z(z{u-b(`OCsDe-bFz>BqOeD_Zy1*b z?^AUt*87Z)wxjFNL5x&a-e=uM!elVY>2(U}aUlQc?VGg?^sQF6xNsvd0~n<$ zZ*ohsKcaEmgpKHwWU%}U90>i?)r!!zG`= zYkQ9J&J_Y8H_mp8D=TX{u#CjPyY<3vdurkJFpf13=2K` z{tTn`+MnT{gHwNofhy6TVFJO-pW!81Eb?dAizfaIM+M|=BStzCvF;nuGr5_F<$P94 zJO<6Je^2~HGG-!9-+}@Gbd`xn-c6Z^mee&t>72!{Q=m0Fkwl|n)VP#r6rxf;DnY=i z>^1krW4R<%X++Cta??_C@{yrusY;{N8W+{t-mN;w7cWK~H%x5;r_ix%I>gcsPtAfw z&_VRQ92@J-6X%C6s3@muFu3kH1;EF8x~}E|f6QKJy|3I@g6#U2OKJs9%YQ+m`HDnF zBdt^5h;db5G9BC7RKR*s5xtThw|gs89ygkyg`?k@owi8>BOf|)WhGb=asrKZt%qPI zBqA`6Q`-o)+L{SUsU{GaMB?WVQSZ ze54Qx9ED{OUnF_BnQ(&J`ufZ@!)ytDW{_#O%4>#;?z_rjVfCE5tM?(BrqU4o-#lK- z=AmYLRqn@Gx4&lC{Nh8oNce>Aqp zqF^4@;&3%c#5zij}Ikv3AHJ(#<-<)e_S4MZ;?m_Ofa$pZ$#k=@LpeC0lxaF2pfSGkYC2V z5uv9i&RDrp_EG9}&fw$2gy_A0PCmZ>##gA#Uqxty($gxV&)1>$Np+O~r^hTjAo_a* zumb)B*>!Uw)E3CZq7BlS#9=38&DNj|)+3KL_KS_t!5?9cAwa)H32TMPf6yR?2MtO$ zBtlD4)S!R7)M6FP6PmdLPG$E+x}bPm9W)1;9L*zJ6sUS;UE*#SH&7}I_yc{k9@|?n z*EYKP;o0Vin;2}Q9lwC9ZYOr;P$ynYN+>9LsH-8NnM9}TQg?pz$4Ny0dI->e@97u%^Yj^Fo`Km zfz)p2UB&(_TPDVhN4xkVOxO@0=gFg?NZ-kY-Vr65Lr;)0y4M|)DPnSTRA>XV_b5r( zu0&*k(VBG*<+DCirpjDXycfGhk!ICW##BALvA9^_+6tO6!N`W@EP%#=Zg|j_0n@w+ z7eLFi)<+otP>26ze<2~p>T3b&F-q@I1|f}qf6oEFzW{!y;(YGmKfm|ZRlF&;8VVfC zHxlFMQ_g7xaYJkM2%^VpKnzCo)s-4_`itcVZy?Gm3LOfdyS1@5y>(R#+8^gojUb8u z^ri_=Z)iCwJqnG6XyKSDT@*nR6pQ7Pa1sQ@86N3qf;X9Te-VJi@4t#SdvcAOi|m}X zJP6G7K5(**=MgC_jBYcWuSK_+vbK`G!6NjhZAD3f&RUl-Gl9f3wUPCQ}e9_w3+3pvOvaV z!F}u>kc=byB&kQWE$yCyd}j}XaSx)aJg#H4&}bX;n6^eQe?je~-r{464laspq;tDv$5UmY z(eaps@f*^Y7 z69f{Jfm(2!lRn7g4$x!>W?8RFSA#Z$g#p&csB4Sj*Vu=A*_QnA8=b6IBo|o0jte8XgClp zN&IpCQ7*Ck(#;<*)wzFLytqlbS`(JW-zq~8i(@Ok_uIVwQ<%liOG?@|md7PdTVUxK zax@DVe|cp1M1?cXOI%K3;BZZmO;WF>VnHaHxs>LLf{W~^+f|<#nE62O9-Jv(i9QMPx=1cbUT0a2~iNMSa?;H#N#1 z#N!~L@?@HU7!mT29JbruM}1oy0VN)F{{Z#z09UCiyjZHKrqe&172g0X58eO2bq}>6 zHDhHPObWH7W79MusS_N!WVc1qiMP?1_rzoWfcscceBgvfB8a3c};e+ z9y@ELvOSUNiN@s>3e|e<%iR$+z_-e5mSQ6r4W&Z{a%Gq}R zJN3WJHpnMlR7--S1Oz;9)EQsEg7(6o5J2_NOvHI#>fwuvQJ9CmL5q0z1^@}OEA;>v zcB3o??_I+I*51jxC!&XvUtr%>F%YtCmoTiK=n^YEg33B1z(ivfe33Vdwuvh2F2 zO8eN(3DTjuBMKPze|m<=IX%)oayuh_C~`-nxY#d?B!3{0hXMWww(t{g1-0P?=D=#D z{Ko^T7)*Rp(P>{hs%94Xb|&ql05)J>Bmk4!z`c&%qL#vE^ijxKoLQIc>01ggBZ=SU z@myL|RaI3}0#pE008;=Ej0BAWqJn_|nOAKGo z;8(E>oSaL?#i*-DYn#59!}caVt0N%LL_DHU+?vT(Qgdq&fM46Sr<5!5>#baWNj(ug z0cw&S<+vwxe<13c1Z5%Ofub!}$bwR_bbG^8EpZvtJ@uEP(u4wgw6JUyCYFeZaW~X* zoRcGweB337&Ah8&Vp+D#tD2<6RIDET6uf>ngP?alubQQ>bJh@Z2C-xi z!zA=T=klT+ADD*^RG9B2Y!-44@oF&Hp{f&d{H2nYg!;35cuAP9n>2!bF8 zi$Ej;e^lhVg#*t<$;Y;>8zotF;AdBb$JD1jA>-9DGe2>xHmrdG4atz92ijzetE?9} zOI<-P_{;*ZSW9)z12V!JcvHt&(21`ZA>`w?bBdhzJ!`aVzNY;0W7^^4Ihcvm#+RZ- zqkZC*>p{(@MiEK#6`no7TIHP4YXiGSHeK;%f5e9e!w8w6O9<&&JZdSeHQIY{(;IX~ zV#FFJKg>7c*S)!$HzJ>VH&fryUXLh;y{|jMe*^JX9a3LD;7cp`r=!PAE9s9jO#Xg} zb!UW$(z+>c#4tUssY+4wrJWGN;b>NXp~$db$Xw z40pn^CWhPR*;C`A_B#45haUaNQw2rof5IJ9(zt8CM!^qgX|EukZ%2l;brxZ;JPrmA zF*kNU9qrT`+D7RFLk^G`xopAkU^HFPfnFi-HfW#qHbgDYeGlD3(V24~L~0KlAJ}p4 zlmgAXy`7O1zU$Wk!nWae$*G7)Xla5^q^|bDV{jY43>GAJWEAwjA@q2VCQa#Te?O|? zrb(y_=K&7!%O8w2Tj4fP!XocsDI5d?uVEngY`?nKOtp`7BbLjTIk1AZTRI~f1|R>$ zh=7zhx>^RkZWEMTw>6x4tcqE>LSO%wx>JyetGbd5vQ`#lN^I&V17NWVT%aS#YXknAEbj;PSjCe~H=}--ARAFcY%at3}d|E zTc>_xH?ooU+yGiof)mi3{jKnYpi~hwkAh@WKQQ;$s=|J&*mr$JluUzjmtz+H08WHF z+)ii|RIsy>^VIGZTq=2x-+;KXZSF@T)R+hDT56rqjIgiAAd4KL^5C#(f1_WDfr>FQ z+77VM$#2I7if4^C>$FUJKKYY|_aalizv_ycx!)hik-OnEu1n;}kVqyws4j>#k#;+W z%QSC5+iF>!Q1?tLB zK1zah{VG9&lz+O|GObrdf6mb*MhkOadO1z6N#_hfp0C82%#d^_Z=E8cT}>FC8l@Z? ze>rS;wubC<_ev;Ws`Ox(^*kG;Dsdk7w!1(E$xD^~vw?6p53~peDxzRshpYBs2^R6I z@f4zM9390?NyNX#Z@E3K@2HVYnss7DnNeX1GmS-=uhm}e?>8vZ9CVg$uA)l z8xqRAQ8qjv^#HP(y?Q{Km$bTgr%qN5%5!AsT@vY%p*1CDRcM{e2jiB2iC`W+=dC2t zBgxw%vlP(7Y7awWnNcx)m`MxpD-_zrw7Lr`e^uWBF?b6pUL3|6gCRyF!C!);bb$r` zVAz~0<=fyehmT$Vf3OTmFIvl3ez?l(8$~tZD$D8JsS`|^Y|;#k&djDRh~{ob8U@e% zaG_~1Ze-V9284m8Wo~@)gUQPwL zx%Rj$OUFLH#jVw$aS^J!`>0m8Gnd$@BT8gCSY_5eZ?IIr5@r2qAk|8NZ6yLy?Qjs9 zL9v6GUAUL%7_s;&0BH9O=J6&E*!$sXKlm7_sV)#XuVn)V|MSAONKOmD5iC+8`=pmF zNs(Gn)Y$SDf3*4uFhM{FG!2!#=Z>EMqw4vLsec^l?sR50p=pKdbVjL!hV;yIgMAtU z5T>fcOY>}m>m()wZ_$)i-mv@MWF4VgO;{vM+b*+J?gmsiqaEb}nHBZm;0Qj8^5CDN z5JoZOtHI_S^%Y1)uL@rXkCf@_Opr&)bky|yy^8;{e=soi+ZNuSkAh6i`L-5ON=a{v zhfLik=54d#NBR6^^n+!nKl}OV#mgCs8>tB`p^w%7 zdN_t7HEa$ip7RtI&&-J`G0z7;7K`l&nddj}&K-~1(!mBE99i-c;qd%}YzUYJ2Is#l z5`v}=e{(L*+7A9VrY(-7VI|VEBb+xQKGv2dg3SZEV1u>iVzdbhTu0eQ;^}m?AzNmA zyNI&zAyQfAD98K7DYsQ1Gkr(7zM}*D$NdV*${yuo(F#&R0Zi^mTy$aBKdB!hSD`{l zM+AHDeCaF3r5J5sF9dhpME&)9w0?I83ap4Be*6%H|`hmxbx_Sa-D3s%hkD(Rtb2YeXFh;!9o-|njKm-*P@c^fc1qP6Au+|tZ0G$ ze=|CvdB)D=1yEI%F0h>uZ*xM5hp%Y{<`SJXO z(pOMg4%$1K(I(ICXF7IZ6_y_~@~}E*i|kUqL!Mop z&>(4Qmw*j?)mLF1nP}Iwh+#(4-ihucPPsIb?#-GRwlFn3b)pHN_-8W=2W?Lhe=5(g z$xU#}PuV5p50-$ZO171~1qP{gQ}+LEp=yw|f!&n-%VM$mD*L;ph#wB_y}ruc#YM-r zc=oxPW*{v`G3j?D<*hBoYTK$p*GDn@{M|J%wThR57PfxXaGWK$`V#8XSKb1pa_k$K zPNE0k*C&4hy5@m-EymvY+18?Me?v(=swj)zoeJZD18j2yo)%(@>z+S`5nWVjZr0^~ zuegaF^!qdoZw@_tN6;lS-a<$+UpDtB<&ZJ^io!v#k(UCm<`s~tZDGIwBXFUP7J=dc zU$73^-vXLVDLSJt+cpZ=%$xb|M%0`o6kY?sjq<R`|m{gq7Tf8aK;`i#Qve+T(5?@D?W&wW$KGw0cJmkR&A5}fp0Uq!0+iL=HOH=-M` z`w9``v)T9HifKH?)nb5O$S@()f96mB8Tj$^>%;+t5VwzHSsbSxF0$^PQ*3M>RWx3< zqRo%ir782!%jEF-co9&nKTUpYS0AdQI(JV1gAsOa+BcIU+%zBlf9M<0gpro6{o{KR zG@*BNft7ja{QpcbjSZ*mcPWQdZe-~Xfs+itbVScz#|MO0RW;BY4${%`o-1}&Ts!?F zw>yE?*5!e;rWh)UnV2I9ZCC|RA-GH09Weuh5p=~!JaKhk&Egs18UF!8hcEeMK49A+ zfGsJmqiOPgmvz!%f0Q2fZuix|htu1t+N&b_!2lu#B1Wj5c(cD#wS90I@!}PR4d!Il+HZ`EVJR!fr)-jc50{$2Wg0Kl1X`8wZym7_Qu0QlpQF7mzn92TA1* znR&bO^5u#$XX`|>jU!=NZXE5e8&%>p5VZv$ahgXW3S9Y1f69z*((qQ|HaxqPxE0Y9 zG^4W4t;AIiyFW^t?iwAXMQj(iaNWLQJ&%&1oD2YPu;wL?qJ$c>B@orz`JyLz+xB&d zi&x)fi&v8iksF|ZMk!DKiJ;))PLNr-M`57gS6pnvS{bD068`mrwXN*RB|?rkNe|q1 zMw#6fZ#n^!f5BkA8vX#Ptjs%gv>ZJF7%UW$rc*l8KyBRWAUkIQ^2 zq>2)bZ#-k!HOPK%XyHivD^pvt@dwGdi)g5ov-E`2B6SBKVeI+=fjS&U-VPYZooM&Y zZy}tTm+1-1Z3gRTXAve=Gs5YpLdJxS4<`U_4_(t=w3f z(B$j3xxm`QL!7-|1)8PJ4X&DWE}w9r#e1e`Q?5LZTGvpHM0sr$>s_jcN`A%2R&TYX z5dhC`dD;5`9Jd1%)xV#1Rr}QlqX>G&jrYP@T8Tg~W^S9DVXJ31^TS0_kyqaX2hPs3 z+>=7bf7>Tf)hL-h)n%>HdE=8F94)>f;tL$I{DS~R8olYhMhy-$}^Brt(QjhiX3y=`84}50~hv+;JPO22+^Sf{{(`As11A!5{4v zFh^NiG2+4x3V;4vz!`>wh7#zzj6d9nxf!Ay3#|%cY}P9*$7$6@VizO&D5S(?!jS)f z)K|GqpR#fnLak4NVs>Et^wGS1fyTy=# z;c!;Hy(fCPc2^9(fjQ?3eNOpF-R7lSU)v@E?^zUDtcwM1IgF3lf`O7SVniqIf4$-Y z_bhMEu-98ykpFT7Lt6-AG8P~%rk_Z_%V7~TbQZp+VG4K>r%Q@JHpu2wf0I*Pb&kg& zEGxc%SzN5IS^gJNTy|*u@<{F$LN^_Y>OdA-EhbqO_Qz;74``kmH7!BSIrlYT7^|(54u-nMU5I~6v2DggN!PDO~12sllSe}8sG;*hh@Du+oilkwY7mV$H@yl5ZFI)=()Tpjl%W&d&V z=1%NrW*&^~iE)NDae)$f1ur?J4!Umn?f~# z`rd#~04^%ca~{{UfS`aD#k#9`)@9I~XEy!j0|`~`MScw4))vRpeov$r3m^c?muYSi zr72f#O#K?&&yvHjk}N${)z&3QkaoSv8v{0{Qucc|S~@Z_cFIMlj!|uC1vgiXoYU?&+`txAy*$9V#WY1iu$sQ@&~g z5foMo{N-Y$2Z8H#3Mk;XL;@(#K)xLJnv#{iSvTIrc`=FN zxGh?BYXf2F6h&|~I}UvP$zIf#zX!Ds6zL>;ru}S#>NO;29!Xh%Gdx@dK7+yMKgLL# zR>Np-<(gxpe{ZXNN47GBNQ}jPaH4ZPB^}R5b>5LYS>kR&+2*9_& n!tlr;My^b*ijFE!Tu{oz?$=jf8BDqN83J9Ha{cRF1Gzu}1yLJe`he$6rVa%2~|m>!R) z*Bjb+O(;@J=lxQ+q`WuE4#=kpe|AuUnj5+mU=s+FJir2$0Amx%A5a*9e?8JE0|9Y3 z5DfnO);au~~T(3b^0I`^q)EA4f^Tvf{-+PKaW#bg;>m9APkwbIr~rLInWW#@X$2d8N#H_bmgZU(2e>`iK| zF(4H--}$*zeE_S}UbXZ!f5Mt-W365A1&T+`dVPg<8kUq2iGA6*WydjF7N@SxISuF4 zu-Zr_Fc=O4L&r&rS6!V>or&+LQ9hLRvfR|ZB;|6WtBZ=OI{izs$(#ozY4pnihBT#V*Z~&pOGzJBXU?4CEe*^-8Kp@Zvf*=S& zfUpDzf*>@4&=3-#s7`Bi0yJKyht#G^U-C6GI8yqz1g_k|w|Mb&boBUY)E|QNMjvuV zYYU|iIq+dk@1$))E&4H_Djms$IlIcG(l2SbWwJG1X7i9?rg_xe$ZI(c5A4e`t8%M5GR69CBS#!p8aS;9d%)fCPI;j0B_*6INp&MJ$h>I?xuK z-5L}!wS`hAB0iZW6z|*?S*to=pta%hg)Hj;huv+xg}dO^ zd!4^_ep@fwRNQ(`7?xYF5hl9zlA+YXd(W+xWZ&I-GkU!Be;%^ZTdxWcz4dZP?yZ;k z(YIb(r+(|D0UHAUXw9yG+JYfYGgC!6V6f{}DThFCf$L>`NitvMDD@*|4sX^;9ku?n z-6pjS+3VisQ)=Li)cfa-LnqRjadz^u15W2sq0V7OQ3AF}y*phk$`n~@7T2&=RbN>?We@}1f!41(NgR0aTL zf#WdvD=73_Nu+3Uf*0xM&aL))xs^Vm`Jt##fW2ChHy!kZsPz**N6`*wd}Z~lhEDE< zvAB2l(S?V4kEK*F{7XM!vyQxT;w-APSHhXoJT0Dze@a)Gtrqn8l0CZTR9jgl<8{Kg ze&|TZIHojCny$0?7Sm3kjWA$Y{}2nKUo$p1V-}(d^_Nvp46=*c;H8XC4>C@X+;^Bq zZ=S(7^0jETh+Kg&Y)V{|T~Cl%MC12Fk=(2QN|ggx06K4YqE+@t3voY`-+mSs>`tZIwQ0V0Kj#drg}25TNKdZIUg#845o%2p4( z;qaSe*;e55^3^u3`gam@UCDxriggYQX@7C;r@705_z+6* ze~+FPQOD7VTGi2O{D!yZ{6Y!%tevzm*q-87XytYx(CVauUzn$V6e1j*Ih1e~`MwqV ze_?rH6kbOTD{5-I{QCoz(-#iDHLSK}rZ3zg3IZ^1(iaA!SZ*@Vad~xA-zl8o*qpUw zj^jqCVT;Q;r7;XeU-)i1Zfkbv3;*AVS1x}=cF-5D&>C{!H>yLy4V&XbznE6r)4q`^}x?*wxe`VK6cC+TH|>e&!bI7ly%VC@!|48jy!8mAy+p@Kn$f4gL+Dq+LPouy_MA)h6Ta5IKiBp&P;=8t8z*E1}) z^ux}&o?)9dzQEhti>yaLyN-AD0ZreX470%ua@rePCNr0~BpS=YKGed)M=pW=0kA-N z@4XDGwfWMfuiW5Hvw=Pi(Kf8tmSj$0O88IwYbRkUO*TmYDdOXqw`D3WfA#|w?79Eg ze;}JbRk|cA2#dNOmbx1t;2W_|Q z{zUIWB?^5$0vjT@nHP?Kf7t!s*cdrPIX-?Yf-44ndc_@wiTY-{@9J?Xr%)n4Qo7kJ z8xY1DVwu^aT|dEXas@TSE_3j8)H>BAebyWOgqJjjE7+BpWIH$q%d0i>6Visyi)T$= zUPs)kB%4jjJ|}P~q-^w7`#B?1f`!_Ln*!}xaX#6N&<6UFLyAZoe=IjYIH-IEG&?2N7pXD?59$Uc}jOF z#oqs<_@iNC*I&M1zvGHrj#ddn82#u;$Kj;738Zc?X#`kl?PA}>G&}bf&kj7IWN8Ec z{g*E38vSYb*FK5df0<;k%m5B#%~q*9`1%Y& zQJQ`R*C^ORs8p5tQ>V{hxK`W7`3;uO;8okkXRxRg-5Hz7 ze}*eWvwwy&#?wE;8{p=j;mO&Ge}?UL-=AR~&f1^hg2Jgk!(Wo<&#=-k^JkbWFc$eU z%%zDx!%l&0f4Am{&O{7n^J;0P*8i5bYR`+)-BDSEmq zNC5&DS}TJ#27>Hr$R)LaRLft<(0oNsMffn;vU_bs}DaL=$KV${~u{=k?Vqa9QznLaP?)WU-sD?_-PSa*3l^iy?%@G zV0r2C+I_S^nV_Wg$YQqF3XJTsXHiwoM{7Q_q49)eGZbQE!o&z}iG=R0$|Ws6uYE5V zZZ8BD-@}Agp!LTk=%W|s6LfTizgL9hf7}Uh8fh*6{olt&kh+AfDX{X=k}WO|GH;Pc zdS<^_f=`ri1!}Laq5xk_tJugu3&`o>-iYt%NhVjW^!fPo${DcqVM4z5&#ABPzu^_v z=8qOMLWw2H=<{{1`=q*lfYUn`9uQqQ0$Aa{J?uJ)2-P>T%OW0Wp2UGUWld*Le;4cV z0v`)wjog7BA;%CP#zmQJ>1C`D!$X0j8xp0G6yK$P48wv6%@Z2u!N0P5I$aPIT%Dc+ z4HwQMI~`N?jF8rD7;jN33poRQv>yAdSgmbTGsCk@+D%OUru{|dY~vZ-hRrYL=}A;o z*^FqT^L}gaL&uLc>@vvZwxGj1e?zA!D2ERiXS;0EJz$P(X^~NCDC$^0Q&f@f9w1?n zFn2iPhGTf>MECr9Sp`fF?H`a~;!|>X@b42{Poex`ZO5JSe*|Rw;}%FoV}j!9sJ0y{`ezml_&ou2C-u77hkxUm)w;Fg4VJf3k+8^XEZa zhk8UaW@yq8=`sOtdB;YH1kD^s_%;?(z`&>7&Mn3MEjd$+8;=(HqkY&APtKD^!$RN5 zh5ivGnnQsjWvtg7h$&(+=c>@AJMU5GvR!G*`qVWd4rSH)P<55L#=J8|k0Sk}C4p(> z@W#Spg=(pdn_e}`^(GcNk876n|2BuP++jp(WiXI3gRAR^@vZ8*RWF<(I=7`^!OLE z1#bwISBy?5fNg4H3-{Jl2(-V_p&AJk0VbO!nBUMcmU<^qw*YhZd45Qmrsn??0oUE;+J77uu zY1{7pWGKToJBZtNd{aoHzOVzecn}(f1;n+U}A-eYkr;>fJztq8~^=H1e^MRme1`%UXz%L3yd$mz!r$>Kt=DXl^!-x z!a)As4fi-R9QeMV!?gv4@f@9K9SZVFVF~m+2vg^F@*7}Q2g}= zDoV{H!~SV$o?V>im0Jy!zu;rd4gv&iQeLdr!UTc9h()^_e<0uag2A|1=_;%1SUWY^ zeja~YBPFQC*IR6C|APCIjr418*%YemU!F76ntVHW(ujsX8`ENV=e{x?HRa`j3K2)} zjA6vTaAV|SneFw$XUBBwH*X0cDIlSj4g?v3GO!M=oAkjvd5}^WLcoXf_>chqK)X`; zPe?yutzF@||vo;DYdN7o3ZN z_E$k2T8q(Zu~iQ>lD$N!5G}pQK8m$MhR=wZ2;Iypq zTwaK>-j~!kYKgS$fP+tpcn1Py-ijpz;BAuYtPldB9O`f8xSUr}DzCv_r0gOaHX8<8eTL&(8qP$va+RC~sn<*1%}pvhe>#G&r|xYrFKDf>Xt%p)zC&lpm|=oi zJ-p6}-i1(Ush=8gguH?1aeI8WwjI#-!SiVy=I{!Ho-vn6jr)1Cw2d1SZVg$WHZD3_ z3;if^WITO$T7!qB{$mPw@8c_%Yl6jxqg^=cNmxxRqIDtQjG{G6ZWqgnG!0{#lMvUU zf72DtI4#;&=isa<0;0rQ=KBkP9f8<79IpCiB+uv7$L8x+CLyeP6S!rvWGi;ES~CMs zSG9NF-hnro9QNCD&HrU%GtQ#5*K9V&qx`?$F9iJ`Q$hir`Sa+{ut&51lC94}>5s}d z2s}GSU?YsS+Vl2AfcPwc>@Rn-&tBrkf7tIBlHkv=gFr>ka%MzQ1?@MR*Bih>^uzkB zfG9-?S#5#Uo-w-ay4rA8qCRcZni^#n#G4_ZV#ze)fe6{395&P5M?G5{!6Y7axeV&# zc~`01TP%H5(`nV56`laB3UogS&K_!`7=e{-l@w~biA~c(bjh01!aJA(({sp#e}*sc zI@UBEF9yb5ib53(pFh0A;rT*Yc}=#sp0BmXvOThTqJhB*1=m`}+AzV00Ow}!26qbx0rg~gqoW`z}rYM#*!2)*|5)~GQ`FWSoCz=SHu>R3L}woWXfKy88b^wBVdbLHORnArKOrn;04+{no-@&bWcJZ@OAXiAS$%{fsZ!fsdA**)3o zptD3U3+$jkHN%)GvFirxpz-AM6CeRl7g7LN0961Lf7BZOa-PLt(%&}< z*sICMHKYAq`BzTI$_rI*Up9nc4Ka1z$cYL&;eFsusX+yFs9#jp=e2OE53AdZ;(LSj z)yzWp@5u=^uM2IInD&N`2!K&_!Sq`Wb#z+We2aP1-n?hM3=A}mPZh|mnS7+0TYCZg z+O9pRT+XjIy8fb?e`0zBu%aGx+?8G;_-%s1AOTU)(DkR~Dq&D6mTvDTu}RS@XcZ9^ z4UI_aq7=mf`^2zp7b>2RcyB#$InL@Lk^I{&c+I@aVPaXf%!^IZULq?AH|-z%Q~)e<2(nJyAkXM$IOd65YjX z08PJ;H#b#gwSxo%1h?xzN`gqXew4LLAkrG*ghk-r*Dvf7L}fBH=jzyKEn{uGb5Y81 zE2UE%k}CArsB=;zE|!D$_un?Cd+G;<9;kyInMsmDngdfH5J*6P01<&e5|D!ef^j9K z`2!Px!Gz)>e=!&`1B2io7zhXgfgmFYf*=Tjpa_B>2#Y`@0xI#{VoX5E$4JqDl3sQ4 zt)4=p>H~g&;T0U>A91xAPQW0Im(aw6KV;0{SqyU48KR^9%paInOEvc4tm7Tiy2&i) zZ0J9%`goR{RGjx|X>vCo8_K{X)9~>e%tUI_yQtV`e`j%%@u1e!=&+>ugy#>it$j(t zu7TZ_jaIx}e5}c>JG$Zpf|eG~E~Pe&_889U20a?aEaTyC#pdq1v8`Wo zh#}Y^|5wOES8g?oZ3G_cFZe%=AG=aX4~}7Pwn#jl5$hSe@>M{P6E1Lu*hG=6m?$8wrNWb#HoX3@BGY6add7NUkrp+x*=xrSYKoq9RwXY zVSqP+Hi|8Ze5DAalfk&}2AHj5S%Uq~DEs?@1voifPXyopK&c`!1RB6p`h+V542ouWmD_=%zY_hB; zJTJ2IEPOmp0sx6J5Au(}!lfHYX2@>|V*s<;GY6Yp++Y7=;hTxjz` zf3+s)-ggy7yvaFwH(|j9MoU5E&CF+pC(ndVtqc|CsE>rHXdue1_#H-;D%G_L!#n=T zw1V@T0~qpvocMb5Pl18vft7uW01T8mEWHwCo26B}0E}H^TGa+^O@>3xjP39xrhO@3 zg}Cvtb-6-LtRA#x6>>)zvu>1$w*gSAf5}ns2B34Id*nB)Ii80xJI{acMllkIlr0b6 ze2fZH9HX=HLD3zaRXD1S;HHFP`I&?_@JLk|&kV)G;{Q65@ir&(pr`}P%q#YxlG*er zuzT59vjW9zwcpH4We%xZr(WkI@+er?06XZq*pw4QcBuQPx5hC#m2^{R8YXAh;n9UvOl zuJ|?3D%v{V?|xN)Yn$iGe2f^|e-VH2lHtRk>1NfP(DcI~hf2sbdI3=PU8VzEA8G#P z`@B-wOkNn z8#HMGFjhEaL+8Aw5Et?9^wpRh7(K>KNz=2;8=pU||JBIGp?Gq!m{TrtCf4G(a9Fsfd>9ysw5<1Ap8|I{;ZYO<3qi4HPoO^K~D zxG82T zi^Idt7FgrS1N;pvU0}$;1wN&M{P!+Q__2F`+`!d~)-KA>O<8;+2$gQi3VU}mfKf>e zy-_70Ghf&`8QT$qak0Z+Xc`PzMbS^vX2378LlI7sn+s-cIK0P1e;q{p;_2Jveq%8| zQ1e?eauaDue&qkI$;Z8#COt{liIZbNDgN`s%v&EgxX=eGA?YF^P6b!cr)sFgAUB6L zI+Ro@gFB})RPvgalY0(Lds3FAUC+O9i}%nt0M!{6snzXF+3i#nWta|DEwoPeyY5maDi7x>7WSuBtg;M^R6I^Mtl_C>ndZz29Pf&~)fwU*=9 z|GbtKiD?1&iTi0;f2eYO*-x8_8mau%KtBW~60k|9N#0h@^b;dh{i%r5r;73^XL;EA zl<}sV)B0sY&&b*!wK4Eus-QVwN~g0$WF6oh1BvAgk1FU^f9IF0+v|U`+C^D=<=IMg-hn5Y2k{S1-wAHp=(-yuEwK(1)-|C~ygitk%q&L)y zwWy7Apt})8jz|HhtGb(m46mnb(@ayZ(2wbGs%=RB3pD_J%7ds0V&|-GN2(b><)}qV z25ic$GxsUO_T0R?PdPQo-)A=m?MCK^l4sWve?vGWr%bLggac1VTML-tje&RwCa?_i zpNE*Fi1QH}AM@o!0VN>w4&|)r#nNR4n8Gp^3=-8+#QJ$oU=?XYC&jmnU_jEmwrY#T zpboZRBn2n>up8WqaHqusF8J!g26`l`4fN56lz4)iH()vDB-;u~XLCwX6f0h5LLG4d zf9q5<&agUTi>&TFehQCD!=s?IGE)*))rH~;dG8u4dOJ3DrL=m;M!XMG##@6C5qQ+T zVq!z^ZWd1-n5s%bE6*3n4c4}IvFY=PeSoZ5B9*-Y1nDJFK7Y56AV9E@NtC^orJ|)% z_FPjm8szK6bjqGarotDUopqXuB27oZe{6R};BF-5Y+aaw`6D^}&E@MbXNo653+#K< zfSDhaL+fa{XdW9vL|N)x>g0wT9b3*)7GNghEo1%nIz4p6h;LH*m4Ap zOq>;$oj-^XU965qm6!`RaC0}u&om9?4-I=;JPWnPifjCT*@*OQJTrR|g(u7Wf894< zZ3=yYZDA%s_8spXb$`5px`6ev;{`NRS#+j3+dLOgJ-f%{KfBfLuSJ7DLm>s1c^aOiOb|N7Rg}4Gg%uV>@hRiD1L?%{=^~&Igi!pPT`jC>T9G}GI>v8v)Ezp2 z;p2bcmv1aKMsHo0$G3~4N#Jn@R+*u5fivR9CcKL+|2(8DP?lma)_>rCt;!Q^vjSnb zP^How1*%_Jj4N(ea7DA+f9wgsN5P{46itypbWRLG34yM{Zz0?@S&pb6T+od%h)qd! z6vu4A%C+kQZw}$`3r=9+UI3s{-l7?Ze-qZtZj?^{E?|v;1E=?_+8su017bg{Jb7rd zQR2{3_3PvfVlFFO4^WtB`j_whMD>e%>ApXwbk8zsHQQF_jOVzpf5$ieTYgIai_1}$ zAy`7Wu~!W)CgZWsPijb2_%dU6H~!0&cRCgX-J*@m=^1V0s2?z5JqpO0CK1U;;&ijC zpGp~utm$LK(>;8Qc!}TDDr3LaW5m1PY@d}_dT8FHMJ#A=qOpChdfyv^Tm~3EXkU_J zk%a066ObV8>!K$*f6_)9Qpc_dsH~elHIAL&5%v{HRg_-JoA&wY z;A9s75;=4xExPT|af4+{sltR~9?y6@^w`-OS~=4Gz|^KTf4>*G_1MX9<*X!-I=rW! zHL`0tgcM0vk-SZ~I8oBiPtS15JIqBYpObTYg%57lb4w=B1uO9_e3A3`tNh(O$PQY! zN;R?JR@o2xFiixtxv@4OD${K=c&ic-IRjpWcxQ7uU1jK8E=6+!7bQ@p9dJHL*C;;{ z4ccnKyHpO9f4mg4wv=j=5kQ__bOh}|I0nZn3VuJ}s`l5?JSl#LV)r|>SrG$HOo4Oz zH3v;`_VbF^CNG1BG^Wk_-&3_MHKCv^Kk{&@%(I9I%x6V#O>Q4b4fpZ%>ZL@!z3IY{ z(Bu4!Lx?V3o5EwE9@)_1G5k(<4I57yq?6+HwObkaNM7VpH@Rt#5L zFu2K-6FaWb`*Y8h&xE_P%2P+(3HMKZ`iI~ThQ#0(_`3L^TN>RA?i>r(3IleGyJZ87 zX$gl+e+b?E1wx|YwSPit9_5fr*U|G}L{5_Oqrlp7>@L(M0m~=%a+0f%SQ*0_Pc6SH z#e2Z!AXH?G*cr|1-qCD>15P$ZLrotgOW%P7sVW*%)*`fm{P}%fBmpKw_!XZxZrg_E zn(~-N27%~(;R(zO31@2qa2VrH_)~stUhK{be?F|**Z7A7f|ieKs!?w&@CPO4T|O;l z=48G@&_1N4zzaV26Wl4(ad z*{hUgBveLi)(ZZR2tZ~M05`;##sDop8JAIwam@LD(BX5I5-H1q@Xty5tL&l?l@nBO z7bX{SC}OBZRSl1AG>cdqlpq{oZB6rm2bs>=?e-&>8 zzo1v}zKj$-+2xX54n2*-4Eg93dIS|yhJr|QR^stuT<#H3;WQ@Y!<#%f49CTcm1|Jz z5d3-!HH`{98PoXp`-K4sC``>|;A!1=(W`Z>o*X(9G zoI^oDkUawX>rLumN>qVRq9p(YfAwfgfnVUzFBoy?b{xs+*9UW`3Nx0@gb@zkFX2){ zLYlp9gEYHm&wj)N{8kkBL-6?vLLe2%G^R}Rk4}9>Lm&% z%(%M0m+bq0B%KzyQo1|y%J0yJDy>ULURG&?2N+D@2WXsAlg)-{HU$VhkGCvN2pWyV#{g*$62E014+9`O9%0~!O^>e~{GLfS;WU-(?AS&s`jjK(bD}d4>`en0v#p*tE!a{fIQI zpJsqur*s*{+zmq~#NQM@S*S28tc;^txr)uO%>-c|&4W_;{NHQx9jpLpdyfQzsG|B1e3^ai1pPeu)oiH&XvAl@m!X zN@j+hYayzE7A8w>lowyTqNtZ{%!mMr|1CW|hKgCpK`=b1@!+V5wg6tCv8kiSA>(qe zl09%BsP#078bEYif2<#A+@^ck|J-l`?3s(GZ%cM&|BHJ%93Z!11;5dp`4y^n)2pkT z?ZTSoQL|`oF1D?;5^E{56PP~ALhEP{F12r z=M~U(d+x}w^nq)V%P3hzW;tR+xJPZlkwFwt|D?P3M}{lbM=ws5j`Q3nMq7K;UJ Option { #[derive(Clone, Debug, PartialEq)] pub(crate) struct CompactedUserMessage { + // Keep source identity even when compaction shortens the text, so rollback can + // correlate the rebuilt message with thread-owned retained evidence. + id: Option, message: String, internal_chat_message_metadata_passthrough: Option, harness_metadata: Option, @@ -540,12 +550,24 @@ pub(crate) fn collect_user_messages(items: &[ResponseItem]) -> Vec Vec { items .iter() .filter_map(|envelope| compacted_user_message(&envelope.item, envelope.metadata.clone())) + .map(|mut message| { + if matches!(identity, CompactedMessageIdentity::Regenerate) { + message.id = None; + } + message + }) .collect() } @@ -560,6 +582,7 @@ fn compacted_user_message( return None; } Some(CompactedUserMessage { + id: item.id().cloned(), message: user.message(), internal_chat_message_metadata_passthrough: match item { ResponseItem::Message { @@ -678,6 +701,7 @@ fn build_compacted_history_with_limit( let truncated = truncate_text(&message.message, TruncationPolicy::Tokens(remaining)); selected_messages.push(CompactedUserMessage { + id: message.id.clone(), message: truncated, internal_chat_message_metadata_passthrough: message .internal_chat_message_metadata_passthrough @@ -692,7 +716,7 @@ fn build_compacted_history_with_limit( for message in &selected_messages { let mut item = ResponseItem::Message { - id: None, + id: message.id.clone(), role: "user".to_string(), content: vec![ContentItem::InputText { text: message.message.clone(), diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index 44cfad1995..72fa09a1e8 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -66,6 +66,7 @@ fn user_message(text: &str) -> ResponseItem { fn compacted_user_message(text: &str) -> CompactedUserMessage { CompactedUserMessage { + id: None, message: text.to_string(), internal_chat_message_metadata_passthrough: None, harness_metadata: None, @@ -129,7 +130,13 @@ fn collect_user_messages_extracts_user_text_only() { let collected = collect_user_messages(&items); - assert_eq!(vec![compacted_user_message("first")], collected); + assert_eq!( + vec![CompactedUserMessage { + id: Some(ResponseItemId::with_suffix("msg", "user")), + ..compacted_user_message("first") + }], + collected, + ); } #[test] @@ -142,10 +149,11 @@ fn collect_annotated_user_messages_extracts_user_text_only() { ResponseItemEnvelope::new(ResponseItem::Other), ]; - let collected = collect_annotated_user_messages(&items); + let collected = collect_annotated_user_messages(&items, CompactedMessageIdentity::Preserve); assert_eq!( vec![CompactedUserMessage { + id: None, message: "first".to_string(), internal_chat_message_metadata_passthrough: None, harness_metadata: Some(CodexHarnessMetadata::default()), @@ -223,6 +231,7 @@ fn build_token_limited_compacted_history_truncates_overlong_user_messages() { let max_tokens = 16; let big = "word ".repeat(200); let user_message = CompactedUserMessage { + id: Some(ResponseItemId::with_suffix("msg", "long-user")), message: big.clone(), internal_chat_message_metadata_passthrough: None, harness_metadata: Some(CodexHarnessMetadata::default()), @@ -261,6 +270,7 @@ fn build_token_limited_compacted_history_truncates_overlong_user_messages() { other => panic!("unexpected item in history: {other:?}"), }; assert_eq!(summary_text, "SUMMARY"); + assert_eq!(history[0].id(), user_message.id.as_ref()); assert_eq!(history[0].metadata, Some(CodexHarnessMetadata::default())); assert_eq!(history[1].metadata, None); } @@ -292,6 +302,7 @@ fn build_compacted_history_preserves_user_message_passthrough_metadata() { let history = build_compacted_history( Vec::new(), &[CompactedUserMessage { + id: Some(ResponseItemId::with_suffix("msg", "user")), message: "first user message".to_string(), internal_chat_message_metadata_passthrough: Some( InternalChatMessageMetadataPassthrough { @@ -314,7 +325,7 @@ fn build_compacted_history_preserves_user_message_passthrough_metadata() { vec![ ResponseItemEnvelope { item: ResponseItem::Message { - id: None, + id: Some(ResponseItemId::with_suffix("msg", "user")), role: "user".to_string(), content: vec![ContentItem::InputText { text: "first user message".to_string(), diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index b02d2188d8..19ec09fab6 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -61,6 +61,8 @@ pub(crate) struct ContextManager { review_history: Option, /// Host facts independent of the model window; snapshots share immutable state. retained_context: Arc, + /// Live and replay instruction capture are enabled together by the session feature flag. + retain_user_messages: bool, /// Bumped whenever history is rewritten, such as compaction or rollback. history_version: u64, /// Monotonic user-input/reset revision, independent of compaction's history generation. @@ -156,6 +158,7 @@ impl ContextManager { items: Arc::new(Vec::new()), review_history: None, retained_context: Arc::default(), + retain_user_messages: false, history_version: 0, user_message_revision: 0, token_info: TokenUsageInfo::new_or_append( @@ -180,6 +183,10 @@ impl ContextManager { &self.retained_context } + pub(crate) fn enable_user_message_retention(&mut self) { + self.retain_user_messages = true; + } + pub(crate) fn record_retained_context(&mut self, event: &RetainedContextEvent) -> bool { if !Arc::make_mut(&mut self.retained_context).record(event) { return false; @@ -188,7 +195,7 @@ impl ContextManager { true } - pub(crate) fn restore_retained_context(&mut self, checkpoint: &RetainedContext) { + pub(crate) fn restore_retained_context(&mut self, checkpoint: Option<&RetainedContext>) { Arc::make_mut(&mut self.retained_context).restore(checkpoint); } @@ -318,6 +325,44 @@ impl ContextManager { } Arc::make_mut(&mut self.items).push(processed); if crate::context::is_user_authorization_message(item) { + if self.retain_user_messages + && let ResponseItem::Message { + content, + internal_chat_message_metadata_passthrough, + .. + } = item + { + let mut complete = internal_chat_message_metadata_passthrough + .as_ref() + .and_then(|metadata| metadata.content_item_kinds.as_ref()) + .is_some_and(|kinds| { + kinds.len() == content.len() + && kinds.iter().all(|kind| kind.0.starts_with("user.")) + }); + let text = content + .iter() + .filter_map(|content| match content { + ContentItem::InputText { text } | ContentItem::OutputText { text } => { + Some(text.as_str()) + } + _ => { + complete = false; + None + } + }) + .collect::>() + .join("\n"); + Arc::make_mut(&mut self.retained_context).record_user_message( + codex_history::RetainedUserMessage { + turn_id: item.turn_id().unwrap_or_default().to_owned(), + message_id: item.id().map(|id| id.as_str().to_owned()), + text, + complete, + }, + ); + } else { + Arc::make_mut(&mut self.retained_context).mark_user_messages_incomplete(); + } self.user_message_revision = self.user_message_revision.saturating_add(1); } } @@ -488,6 +533,9 @@ impl ContextManager { user_positions[user_positions.len() - n_from_end] }; + let first_removed_message_id = snapshot[cut_idx] + .id() + .map(codex_protocol::ResponseItemId::as_str); let mut review_history = self.review_history.take(); if let Some(history) = &mut review_history { history.truncate_before(&snapshot[cut_idx].item); @@ -528,18 +576,21 @@ impl ContextManager { .iter() .filter_map(|item| item.turn_id()) .collect::>(); - Arc::make_mut(&mut retained_context).retain_answers(|answer| { - // A steer creates an instruction boundary, not a new turn ID. Replay exposes - // the original calls from rolled-back checkpoints, so prefer the exact source. - if let Some(source_index) = snapshot.iter().rposition(|item| { - item.turn_id() == Some(answer.turn_id.as_str()) - && matches!(&item.item, ResponseItem::FunctionCall { call_id, .. } - if call_id == &answer.call_id) - }) { - return source_index < cut_idx; - } - !removed_turns.contains(&answer.turn_id.as_str()) - }); + if self.retain_user_messages { + Arc::make_mut(&mut retained_context).rollback(&removed_turns, first_removed_message_id); + } else { + Arc::make_mut(&mut retained_context).retain_answers(|answer| { + // Legacy answers follow their original call, not later steers in the same turn. + if let Some(source_index) = snapshot.iter().rposition(|item| { + item.turn_id() == Some(answer.turn_id.as_str()) + && matches!(&item.item, ResponseItem::FunctionCall { call_id, .. } + if call_id == &answer.call_id) + }) { + return source_index < cut_idx; + } + !removed_turns.contains(&answer.turn_id.as_str()) + }); + } self.replace_annotated(retained_items); self.retained_context = retained_context; self.review_history = review_history; diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index c6e926ee0a..d23e8ef8ad 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -1265,6 +1265,55 @@ fn drop_last_n_user_turns_preserves_prefix() { let retained = history.retained_context().clone(); history.drop_last_n_user_turns(/*num_turns*/ 1); assert_eq!(history.retained_context(), &retained); + + // A steered message shares its source turn, but rollback must keep the earlier + // instruction and answer as complete evidence, including after the next compaction. + let mut history = ContextManager::default(); + history.enable_user_message_retention(); + let mut expected = None; + for (id, text) in [ + ("restriction", "Never publish publicly."), + ("steer", "Check the tests too."), + ] { + let message = ResponseItem::Message { + id: Some(ResponseItemId::with_suffix("msg", id)), + role: "user".to_owned(), + content: vec![ContentItem::InputText { + text: text.to_owned(), + }], + phase: None, + internal_chat_message_metadata_passthrough: Some( + InternalChatMessageMetadataPassthrough { + turn_id: Some("shared-turn".to_owned()), + content_item_kinds: Some(vec![ContentItemKind("user.text".to_owned())]), + ..Default::default() + }, + ), + }; + history.record_items([&message], TruncationPolicy::Tokens(10_000)); + if id == "restriction" { + history.record_retained_context(&codex_history::RetainedContextEvent::VerifiedAnswer( + codex_history::VerifiedAnswer { + turn_id: "shared-turn".to_owned(), + call_id: "ask".to_owned(), + questions: vec![codex_history::VerifiedQuestionAnswer { + question: "Publish?".to_owned(), + answer: "Only privately.".to_owned(), + }], + }, + )); + expected = Some(history.retained_context().clone()); + } + } + history.drop_last_n_user_turns(/*num_turns*/ 1); + history.replace_compacted(Vec::new()); + let retained = history.retained_context(); + assert!(retained.user_messages_complete()); + assert!(retained.verified_answers_complete()); + let mut expected = serde_json::to_value(expected.expect("pre-steer evidence")).unwrap(); + // Rollback removes evidence, but does not reuse its arrival-order sequence numbers. + expected["next_order"] = serde_json::json!(3); + assert_eq!(serde_json::to_value(retained).unwrap(), expected); } #[test] diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index d08feebaef..00d3e8e252 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1590,7 +1590,9 @@ impl Session { state .history .restore_guardian_history(guardian_history.as_ref()); - state.history.restore_retained_context(&retained_context); + state + .history + .restore_retained_context(Some(&retained_context)); if let Some(world_state) = world_state_baseline { state.history.set_world_state_baseline(world_state); } diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index a54c08520f..6b5235038c 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -339,15 +339,16 @@ impl Session { .unwrap_or(u64::MAX); let mut history = ContextManager::new(); + if self.enabled(Feature::GuardianThreadContext) { + history.enable_user_message_retention(); + } let mut saw_legacy_compaction_without_replacement_history = false; if let Some(checkpoint) = base_compaction && let Some(items) = &checkpoint.compacted.replacement_history { history.replace_annotated(items.clone()); history.restore_guardian_history(checkpoint.compacted.guardian_history.as_ref()); - if let Some(retained) = &checkpoint.compacted.retained_context { - history.restore_retained_context(retained); - } + history.restore_retained_context(checkpoint.compacted.retained_context.as_ref()); } // Materialize exact history semantics from the replay-derived suffix. The eventual lazy // design should keep this same replay shape, but drive it from a resumable reverse source @@ -386,8 +387,15 @@ impl Session { // prompt shape. // TODO(ccunningham): if we drop support for None replacement_history compaction items, // we can get rid of this second loop entirely and just build `history` directly in the first loop. - let user_messages = - compact::collect_annotated_user_messages(history.annotated_items()); + let identity = if self.enabled(Feature::GuardianThreadContext) { + compact::CompactedMessageIdentity::Preserve + } else { + compact::CompactedMessageIdentity::Regenerate + }; + let user_messages = compact::collect_annotated_user_messages( + history.annotated_items(), + identity, + ); let rebuilt = compact::build_compacted_history( Vec::new(), &user_messages, @@ -395,7 +403,7 @@ impl Session { ); let retained_context = history.retained_context().clone(); history.replace_annotated(rebuilt); - history.restore_retained_context(&retained_context); + history.restore_retained_context(Some(&retained_context)); } } RolloutItem::EventMsg(EventMsg::ThreadRolledBack(rollback)) => { diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index 93492d9f5c..ef0ce0503d 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -1488,6 +1488,7 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_does_ }], }); let mut retained = codex_history::RetainedContext::default(); + retained.mark_user_messages_incomplete(); retained.record(&answer); let rollout_items = vec![ RolloutItem::ResponseItem(user_message("before compact").into()), diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 8d50a6885d..a2d070f547 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1266,6 +1266,9 @@ impl Session { session_configuration.clone(), initial_auto_compact_window_ids, ); + if config.features.enabled(Feature::GuardianThreadContext) { + state.history.enable_user_message_retention(); + } state.base_instructions_provenance = base_instructions_provenance.clone(); let managed_network_requirements_configured = config .config_layer_stack diff --git a/codex-rs/core/tests/suite/guardian_retained_context.rs b/codex-rs/core/tests/suite/guardian_retained_context.rs index 8f45117e48..0d2048fe3e 100644 --- a/codex-rs/core/tests/suite/guardian_retained_context.rs +++ b/codex-rs/core/tests/suite/guardian_retained_context.rs @@ -1,6 +1,6 @@ -//! Retained-answer lifecycle through real sessions and the durable thread-store boundary. -//! Typed fixtures enter via the public rollout append API to isolate storage from producer -//! wiring; compaction, resume, rollback and child forks use production paths. +//! Retained-instruction and answer lifecycles through real sessions and durable checkpoints. +//! Real user input covers steering, compaction and rollback; legacy replay fixtures also use +//! the public rollout append API. Resume and child forks use production paths. use std::collections::HashMap; use std::sync::Arc; @@ -13,11 +13,15 @@ use codex_core::config::ThreadStoreConfig; use codex_features::Feature; use codex_history::InitialHistory; use codex_history::ResumedHistory; +use codex_history::RetainedContext; use codex_history::RetainedContextEvent; use codex_history::RolloutItem; use codex_history::VerifiedAnswer; use codex_history::VerifiedQuestionAnswer; +use codex_protocol::ResponseItemId; use codex_protocol::mcp::ClientMcpExtensions; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use codex_protocol::protocol::ThreadHistoryMode; @@ -25,6 +29,7 @@ use codex_protocol::request_user_input::RequestUserInputAnswer; use codex_protocol::request_user_input::RequestUserInputResponse; use codex_protocol::user_input::UserInput; use codex_thread_store::LoadThreadHistoryParams; +use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_function_call_with_namespace; @@ -149,7 +154,7 @@ async fn compact_and_assert_answers( test: &TestCodex, thread: &CodexThread, expected: &[VerifiedAnswer], -) -> Result<()> { +) -> Result { // Inspect live state by persisting a real compaction checkpoint, not a private getter. // Repeating this after resume/rollback also catches checkpoint resurrection. thread.submit(Op::Compact).await?; @@ -172,6 +177,355 @@ async fn compact_and_assert_answers( actual.verified_answers().cloned().collect::>(), expected ); + Ok(actual.clone()) +} + +#[test_case(ThreadHistoryMode::Legacy, true; "enabled legacy rollback")] +#[test_case(ThreadHistoryMode::Paginated, true; "enabled paginated resume")] +#[test_case(ThreadHistoryMode::Legacy, false; "disabled legacy rollback")] +#[test_case(ThreadHistoryMode::Paginated, false; "disabled paginated resume")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn retained_instructions_keep_identity_across_compaction_and_resume( + history_mode: ThreadHistoryMode, + thread_context_enabled: bool, +) -> Result<()> { + skip_if_no_network!(Ok(())); + const INITIAL: &str = "Never publish publicly."; + const STEER: &str = "Also inspect the README."; + // Paginated history supports checkpoint resume, but not the full-history read + // used by ThreadRollback. Exercise rollback on its supported legacy path. + let rollback_counts: &[usize] = match history_mode { + ThreadHistoryMode::Legacy => &[1, 0], + ThreadHistoryMode::Paginated => &[], + }; + let server = start_mock_server().await; + let test = test_codex() + .with_history_mode(history_mode) + .with_config(move |config| { + config.experimental_thread_store = ThreadStoreConfig::Local; + config + .features + .set_enabled(Feature::GuardianThreadContext, thread_context_enabled) + .expect("test context mode"); + // Exercise local compaction's rebuilt user messages, not an opaque checkpoint. + config.model_provider.name = "Local compaction test provider".to_owned(); + config + .features + .disable(Feature::TokenBudget) + .expect("use local compaction"); + config + .features + .enable(Feature::DefaultModeRequestUserInput) + .expect("enable request_user_input"); + }) + .build_with_auto_env(&server) + .await?; + let questions = [ + ("before-steer", "Publish?", "Only privately."), + ("after-steer", "Publish the README?", "Do not publish it."), + ]; + let mut responses = questions + .iter() + .map(|(call_id, question, _)| { + sse(vec![ + ev_function_call( + call_id, + "request_user_input", + &json!({"questions": [{ + "id": "publish", "header": "Publish", "question": question, + "options": [ + {"label": "Yes", "description": "Publish privately."}, + {"label": "No", "description": "Keep local."} + ] + }]}) + .to_string(), + ), + ev_completed(call_id), + ]) + }) + .collect::>(); + responses.push(sse(vec![ev_completed("done")])); + responses.extend((0..=rollback_counts.len()).map(|index| { + sse(vec![ + ev_assistant_message(&format!("summary-{index}"), "Compacted inspection context."), + ev_completed(&format!("compact-{index}")), + ]) + })); + let response_mock = mount_sse_sequence(&server, responses).await; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: INITIAL.to_owned(), + text_elements: Vec::new(), + }])) + .await?; + let mut answers = Vec::new(); + for (call_id, question, answer) in questions { + let request = wait_for_event_match(&test.codex, |event| match event { + EventMsg::RequestUserInput(request) => Some(request.clone()), + _ => None, + }) + .await; + if answers.is_empty() { + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: STEER.to_owned(), + text_elements: Vec::new(), + }])) + .await?; + } + test.codex + .submit(Op::UserInputAnswer { + id: request.turn_id.clone(), + response: RequestUserInputResponse { + answers: HashMap::from([( + "publish".to_owned(), + RequestUserInputAnswer { + answers: vec![answer.to_owned()], + }, + )]), + }, + }) + .await?; + answers.push(VerifiedAnswer { + turn_id: request.turn_id, + call_id: call_id.to_owned(), + questions: vec![VerifiedQuestionAnswer { + question: question.to_owned(), + answer: answer.to_owned(), + }], + }); + } + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + assert_eq!(answers[0].turn_id, answers[1].turn_id); + let requests = response_mock.requests(); + assert_eq!(requests.len(), 3); + assert!(requests[0].has_message_with_input_texts("user", |texts| texts == [INITIAL])); + for (index, request) in requests.iter().skip(1).enumerate() { + assert!(request.has_message_with_input_texts("user", |texts| texts == [STEER])); + for answer in &answers[..=index] { + let (output, _) = request + .function_call_output_content_and_success(&answer.call_id) + .context("answer tool output")?; + let output: serde_json::Value = serde_json::from_str(&output.context("answer text")?)?; + assert_eq!( + output, + json!({"answers": {"publish": {"answers": [answer.questions[0].answer]}}}) + ); + } + } + let history = test.codex.conversation_history_snapshot().await; + let user_messages = [INITIAL, STEER] + .into_iter() + .enumerate() + .map(|(index, text)| { + let message_id = history + .items() + .find(|item| { + matches!( + item, + ResponseItem::Message { role, content, .. } + if role == "user" && matches!( + content.as_slice(), + [ContentItem::InputText { text: body }] if body == text + ) + ) + }) + .and_then(ResponseItem::id) + .expect("original user-message identity"); + json!({ + "order": index * 2, "turn_id": answers[index].turn_id, + "message_id": message_id.as_str(), "text": text, "complete": true, + }) + }) + .collect::>(); + let ordered_answers = answers + .iter() + .enumerate() + .map(|(index, answer)| { + let mut value = json!(answer); + value["order"] = json!(index * 2 + 1); + value + }) + .collect::>(); + let mut expected = json!({ + "user_messages": user_messages, "user_messages_incomplete": false, + "verified_answers": ordered_answers, "incomplete": false, "next_order": 4, + }); + if !thread_context_enabled { + expected = json!({ + "user_messages": [], "user_messages_incomplete": true, + "verified_answers": [], "incomplete": false, "next_order": 0, + }); + answers.clear(); + assert!( + load_context(&test, &test.codex) + .await? + .iter() + .all(|item| !matches!(item, RolloutItem::RetainedContext(_))) + ); + } + assert_eq!(serde_json::to_value(history.retained_context())?, expected); + assert_eq!( + serde_json::to_value(compact_and_assert_answers(&test, &test.codex, &answers).await?)?, + expected + ); + let compacted = test.codex.conversation_history_snapshot().await; + for message in &user_messages { + assert_eq!( + compacted + .items() + .any(|item| item.id().map(ResponseItemId::as_str) == message["message_id"].as_str()), + thread_context_enabled, + "only thread-owned context preserves original user-message identity" + ); + } + let mut thread = Arc::clone(&test.codex); + for &remaining in rollback_counts { + let remaining = if thread_context_enabled { remaining } else { 0 }; + thread = resume(&test, &thread).await?; + assert_eq!( + serde_json::to_value( + thread + .conversation_history_snapshot() + .await + .retained_context() + )?, + expected + ); + thread.submit(Op::ThreadRollback { num_turns: 1 }).await?; + wait_for_event(&thread, |event| { + matches!(event, EventMsg::ThreadRolledBack(_)) + }) + .await; + expected["user_messages"] + .as_array_mut() + .expect("expected retained user messages") + .truncate(remaining); + expected["verified_answers"] + .as_array_mut() + .expect("expected retained verified answers") + .truncate(remaining); + assert_eq!( + serde_json::to_value( + compact_and_assert_answers(&test, &thread, &answers[..remaining]).await? + )?, + expected + ); + } + thread = resume(&test, &thread).await?; + assert_eq!( + serde_json::to_value( + thread + .conversation_history_snapshot() + .await + .retained_context() + )?, + expected + ); + thread.shutdown_and_wait().await?; + let requests = response_mock.requests(); + assert_eq!(requests.len(), 4 + rollback_counts.len()); + assert!(requests[3].has_message_with_input_texts("user", |texts| texts == [STEER])); + if history_mode == ThreadHistoryMode::Legacy { + assert!(requests[4].has_message_with_input_texts("user", |texts| texts == [INITIAL])); + assert!(!requests[4].has_message_with_input_texts("user", |texts| texts == [STEER])); + assert!(!requests[5].has_message_with_input_texts("user", |texts| texts == [INITIAL])); + } + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn disabled_capture_stays_incomplete_after_compaction_and_enabled_resume() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let test = test_codex() + .with_config(|config| { + config.experimental_thread_store = ThreadStoreConfig::Local; + config + .features + .disable(Feature::GuardianThreadContext) + .expect("disable instruction capture"); + config + .features + .disable(Feature::TokenBudget) + .expect("use local compaction"); + config.model_provider.name = "Local compaction test provider".to_owned(); + }) + .build_with_auto_env(&server) + .await?; + mount_sse_sequence( + &server, + vec![ + sse(vec![ev_completed("initial-turn")]), + sse(vec![ + ev_assistant_message("summary", "Compacted context."), + ev_completed("compact"), + ]), + sse(vec![ + ev_assistant_message("summary-again", "Compacted context."), + ev_completed("compact-again"), + ]), + ], + ) + .await; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "Never publish publicly.".to_owned(), + text_elements: Vec::new(), + }])) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let checkpoint = compact_and_assert_answers(&test, &test.codex, &[]).await?; + assert!(!checkpoint.user_messages_complete()); + assert_eq!(checkpoint.ordered_entries().count(), 0); + + let thread_id = test.codex.session_configured().thread_id; + test.codex.shutdown_and_wait().await?; + test.thread_manager.remove_thread(&thread_id).await; + let items: Vec = serde_json::from_value(serde_json::to_value( + load_context(&test, &test.codex).await?, + )?)?; + let mut config = test.config.clone(); + config + .features + .enable(Feature::GuardianThreadContext) + .expect("enable instruction capture on resume"); + let resumed = test + .thread_manager + .resume_thread_with_history( + config, + InitialHistory::Resumed(ResumedHistory { + conversation_id: thread_id, + history: Arc::new(items), + rollout_path: None, + }), + test.thread_manager.auth_manager(), + /*parent_trace*/ None, + ClientMcpExtensions::default(), + ) + .await? + .thread; + assert!( + !resumed + .conversation_history_snapshot() + .await + .retained_context() + .expect("resumed retained context") + .user_messages_complete() + ); + assert!( + !compact_and_assert_answers(&test, &resumed, &[]) + .await? + .user_messages_complete() + ); + resumed.shutdown_and_wait().await?; Ok(()) } @@ -183,7 +537,11 @@ async fn retained_answers_rollback_only_the_steered_instruction() -> Result<()> .with_history_mode(ThreadHistoryMode::Legacy) .with_config(|config| { config.experimental_thread_store = ThreadStoreConfig::Local; - for feature in [Feature::TokenBudget, Feature::DefaultModeRequestUserInput] { + for feature in [ + Feature::TokenBudget, + Feature::DefaultModeRequestUserInput, + Feature::GuardianThreadContext, + ] { config .features .enable(feature) @@ -341,6 +699,7 @@ async fn retained_answers_cross_real_session_boundaries( Feature::DefaultModeRequestUserInput, Feature::Collab, Feature::MultiAgentV2, + Feature::GuardianThreadContext, ] { config .features diff --git a/codex-rs/guardian-context/src/verified_answers.rs b/codex-rs/guardian-context/src/verified_answers.rs index ef9fcf8d20..4357ea2f4d 100644 --- a/codex-rs/guardian-context/src/verified_answers.rs +++ b/codex-rs/guardian-context/src/verified_answers.rs @@ -17,7 +17,7 @@ pub struct RenderedVerifiedAnswers { /// Renders each response with its original source roles and the existing approximate per-answer budget. pub fn render_verified_answers(context: &RetainedContext) -> RenderedVerifiedAnswers { - let mut complete = context.is_complete(); + let mut complete = context.verified_answers_complete(); let mut fragments = Vec::new(); for answer in context.verified_answers() { if let Some(text) = render_verified_answer(answer) { diff --git a/codex-rs/history/src/lib.rs b/codex-rs/history/src/lib.rs index 734ec8e673..a77b1a2e66 100644 --- a/codex-rs/history/src/lib.rs +++ b/codex-rs/history/src/lib.rs @@ -161,7 +161,9 @@ mod guardian_history; mod retained_context; pub use retained_context::RetainedContext; +pub use retained_context::RetainedContextEntry; pub use retained_context::RetainedContextEvent; +pub use retained_context::RetainedUserMessage; pub use retained_context::VerifiedAnswer; pub use retained_context::VerifiedQuestionAnswer; mod rollout_payload; diff --git a/codex-rs/history/src/retained_context.rs b/codex-rs/history/src/retained_context.rs index 97ed607b5d..350dbdf84b 100644 --- a/codex-rs/history/src/retained_context.rs +++ b/codex-rs/history/src/retained_context.rs @@ -7,9 +7,9 @@ use schemars::JsonSchema; use serde::Deserialize; use serde::Serialize; -const MAX_ANSWERS: usize = 8; -const MAX_ANSWER_BYTES: usize = 16_384; -const MAX_ANSWERS_BYTES: usize = 65_536; +const MAX_FAMILY_RECORDS: usize = 8; +const MAX_RECORD_BYTES: usize = 16_384; +const MAX_FAMILY_BYTES: usize = 65_536; /// Original assistant question and host-verified user reply, never an inferred permission. #[derive(Clone, PartialEq, Serialize, Deserialize, JsonSchema)] @@ -27,6 +27,39 @@ pub struct VerifiedAnswer { pub questions: Vec, } +/// Original user instruction, retained outside model summarization for delegated review. +/// Non-text input is not reconstructed as text; missing evidence remains explicit. +#[derive(Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct RetainedUserMessage { + pub turn_id: String, + pub message_id: Option, + pub text: String, + pub complete: bool, +} + +impl std::fmt::Debug for RetainedUserMessage { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("RetainedUserMessage") + .field("complete", &self.complete) + .finish_non_exhaustive() + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +struct Ordered { + #[serde(default)] + order: u64, + #[serde(flatten)] + value: T, +} + +/// Borrowed host evidence in original arrival order, across retained families. +pub enum RetainedContextEntry<'a> { + UserMessage(&'a RetainedUserMessage), + VerifiedAnswer(&'a VerifiedAnswer), +} + impl std::fmt::Debug for VerifiedAnswer { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { formatter @@ -47,9 +80,44 @@ pub enum RetainedContextEvent { /// Facts live until their source instruction is rolled back; compaction does not expire them. #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, JsonSchema)] pub struct RetainedContext { - verified_answers: VecDeque, + verified_answers: VecDeque>, /// Lost evidence cannot be treated as a complete authorization history. - incomplete: bool, + #[serde(rename = "incomplete")] + verified_answers_incomplete: bool, + #[serde(default)] + user_messages: VecDeque>, + /// Old checkpoints did not preserve user restrictions for delegated review. + #[serde(default = "legacy_user_messages_incomplete")] + user_messages_incomplete: bool, + #[serde(default)] + next_order: u64, +} + +fn legacy_user_messages_incomplete() -> bool { + true +} + +fn bound_family(items: &mut VecDeque>, incomplete: &mut bool) { + while items.len() > MAX_FAMILY_RECORDS + || serde_json::to_vec(items).map_or(true, |bytes| bytes.len() > MAX_FAMILY_BYTES) + { + items.pop_front(); + *incomplete = true; + } +} + +impl RetainedUserMessage { + fn bound(&mut self) { + if serde_json::to_vec(self).map_or(true, |bytes| bytes.len() > MAX_RECORD_BYTES) { + self.text.clear(); + self.complete = false; + self.turn_id + .truncate(self.turn_id.floor_char_boundary(1_024)); + if let Some(id) = &mut self.message_id { + id.truncate(id.floor_char_boundary(1_024)); + } + } + } } impl RetainedContextEvent { @@ -57,7 +125,7 @@ impl RetainedContextEvent { pub fn bound(&mut self) { match self { Self::VerifiedAnswer(answer) => { - if serde_json::to_vec(answer).map_or(true, |bytes| bytes.len() > MAX_ANSWER_BYTES) { + if serde_json::to_vec(answer).map_or(true, |bytes| bytes.len() > MAX_RECORD_BYTES) { answer.questions.clear(); // IDs are correlation metadata, not model-visible authorization text. answer @@ -74,15 +142,67 @@ impl RetainedContextEvent { impl RetainedContext { pub fn verified_answers(&self) -> impl DoubleEndedIterator { - self.verified_answers.iter() + self.verified_answers.iter().map(|entry| &entry.value) } - pub fn is_complete(&self) -> bool { - !self.incomplete + pub fn verified_answers_complete(&self) -> bool { + !self.verified_answers_incomplete && self .verified_answers .iter() - .all(|answer| !answer.questions.is_empty()) + .all(|answer| !answer.value.questions.is_empty()) + } + + pub fn user_messages_complete(&self) -> bool { + !self.user_messages_incomplete + && self + .user_messages + .iter() + .all(|message| message.value.complete) + } + + /// A skipped instruction leaves a gap that later checkpoints must preserve. + pub fn mark_user_messages_incomplete(&mut self) { + self.user_messages_incomplete = true; + } + + pub fn ordered_entries(&self) -> impl DoubleEndedIterator> { + let mut entries = self + .verified_answers + .iter() + .map(|entry| { + ( + entry.order, + RetainedContextEntry::VerifiedAnswer(&entry.value), + ) + }) + .chain( + self.user_messages + .iter() + .map(|entry| (entry.order, RetainedContextEntry::UserMessage(&entry.value))), + ) + .collect::>(); + entries.sort_by_key(|(order, _)| *order); + entries.into_iter().map(|(_, entry)| entry) + } + + /// Records the already-persisted user item; checkpoint/suffix replay uses this same path. + pub fn record_user_message(&mut self, mut message: RetainedUserMessage) { + message.bound(); + if let Some(index) = self.user_messages.iter().position(|entry| { + message.message_id.is_some() && entry.value.message_id == message.message_id + }) { + if self.user_messages[index].value == message { + return; + } + self.user_messages.remove(index); + } + self.user_messages.push_back(Ordered { + order: self.next_order, + value: message, + }); + self.next_order = self.next_order.saturating_add(1); + bound_family(&mut self.user_messages, &mut self.user_messages_incomplete); } /// Same-event delivery is idempotent; changed contents replace that source's record. @@ -92,40 +212,81 @@ impl RetainedContext { match event { RetainedContextEvent::VerifiedAnswer(answer) => { if let Some(index) = self.verified_answers.iter().position(|existing| { - existing.turn_id == answer.turn_id && existing.call_id == answer.call_id + existing.value.turn_id == answer.turn_id + && existing.value.call_id == answer.call_id }) { - if self.verified_answers[index] == answer { + if self.verified_answers[index].value == answer { return false; } self.verified_answers.remove(index); } - self.verified_answers.push_back(answer); - while self.verified_answers.len() > MAX_ANSWERS - || serde_json::to_vec(&self.verified_answers) - .map_or(true, |bytes| bytes.len() > MAX_ANSWERS_BYTES) - { - self.verified_answers.pop_front(); - self.incomplete = true; - } + self.verified_answers.push_back(Ordered { + order: self.next_order, + value: answer, + }); + self.next_order = self.next_order.saturating_add(1); + bound_family( + &mut self.verified_answers, + &mut self.verified_answers_incomplete, + ); } } true } /// Restoring a saved thread must not bypass the live storage limits. - pub fn restore(&mut self, checkpoint: &Self) { - *self = Self { - incomplete: checkpoint.incomplete, + /// A missing checkpoint cannot establish complete historical user instructions. + pub fn restore(&mut self, checkpoint: Option<&Self>) { + *self = checkpoint.cloned().unwrap_or_else(|| Self { + user_messages_incomplete: true, ..Self::default() - }; - for answer in &checkpoint.verified_answers { - self.record(&RetainedContextEvent::VerifiedAnswer(answer.clone())); + }); + for entry in &mut self.verified_answers { + let mut event = RetainedContextEvent::VerifiedAnswer(entry.value.clone()); + event.bound(); + let RetainedContextEvent::VerifiedAnswer(answer) = event; + entry.value = answer; + self.next_order = self.next_order.max(entry.order.saturating_add(1)); } + for entry in &mut self.user_messages { + entry.value.bound(); + self.next_order = self.next_order.max(entry.order.saturating_add(1)); + } + bound_family( + &mut self.verified_answers, + &mut self.verified_answers_incomplete, + ); + bound_family(&mut self.user_messages, &mut self.user_messages_incomplete); } - /// Retains facts whose sources survive a lifecycle change. Compaction alone is not removal. - pub fn retain_answers(&mut self, keep: impl FnMut(&VerifiedAnswer) -> bool) { - self.verified_answers.retain(keep); + /// Keeps legacy answers whose source calls survive when no retained instruction boundary exists. + pub fn retain_answers(&mut self, mut keep: impl FnMut(&VerifiedAnswer) -> bool) { + self.verified_answers.retain(|answer| keep(&answer.value)); + } + + /// Rolls back at the original user-message boundary, including facts delivered after it. + /// Steering can share a turn ID. Legacy sources without message identity fall back to + /// source-turn removal and cannot establish complete retained user instructions. + pub fn rollback(&mut self, turn_ids: &[&str], first_removed_message_id: Option<&str>) { + if let Some(order) = first_removed_message_id.and_then(|id| { + self.user_messages + .iter() + .find(|message| message.value.message_id.as_deref() == Some(id)) + .map(|message| message.order) + }) { + self.verified_answers.retain(|entry| entry.order < order); + self.user_messages.retain(|entry| entry.order < order); + return; + } + self.user_messages_incomplete |= first_removed_message_id.is_some() + || self + .user_messages + .iter() + .any(|message| turn_ids.contains(&message.value.turn_id.as_str())); + self.verified_answers + .retain(|answer| !turn_ids.contains(&answer.value.turn_id.as_str())); + self.user_messages + .retain(|message| !turn_ids.contains(&message.value.turn_id.as_str())); } } diff --git a/codex-rs/history/src/retained_context_tests.rs b/codex-rs/history/src/retained_context_tests.rs index 5129daae0a..0e945c47e2 100644 --- a/codex-rs/history/src/retained_context_tests.rs +++ b/codex-rs/history/src/retained_context_tests.rs @@ -1,18 +1,29 @@ use super::*; use pretty_assertions::assert_eq; -#[test] -fn answer_lifecycle_is_bounded_idempotent_and_checkpointed() { - let mut context = RetainedContext::default(); - let first = RetainedContextEvent::VerifiedAnswer(VerifiedAnswer { +fn publish_answer() -> RetainedContextEvent { + RetainedContextEvent::VerifiedAnswer(VerifiedAnswer { turn_id: "turn-1".to_owned(), call_id: "ask-1".to_owned(), questions: vec![VerifiedQuestionAnswer { question: "Publish?".to_owned(), answer: "Yes, but never publicly.".to_owned(), }], - }); + }) +} + +#[test] +fn retained_evidence_preserves_order_through_checkpoint_and_rollback() { + let mut context = RetainedContext::default(); + let first = publish_answer(); assert!(context.record(&first)); + let before_restriction = context.clone(); + context.record_user_message(RetainedUserMessage { + turn_id: "revocation-turn".to_owned(), + message_id: Some("revocation".to_owned()), + text: "Do not publish after all.".to_owned(), + complete: true, + }); let snapshot = context.clone(); assert!(!context.record(&first)); assert_eq!(context, snapshot); @@ -20,8 +31,34 @@ fn answer_lifecycle_is_bounded_idempotent_and_checkpointed() { serde_json::from_str(&serde_json::to_string(&context).expect("retained answer fixture")) .expect("retained answer fixture"); let mut restored = RetainedContext::default(); - restored.restore(&checkpoint); + restored.restore(Some(&checkpoint)); assert_eq!(restored, snapshot); + assert_eq!( + restored + .ordered_entries() + .map(|entry| match entry { + RetainedContextEntry::VerifiedAnswer(answer) => answer.questions[0].answer.as_str(), + RetainedContextEntry::UserMessage(message) => message.text.as_str(), + }) + .collect::>(), + vec!["Yes, but never publicly.", "Do not publish after all."] + ); + restored.rollback(&["revocation-turn"], Some("revocation")); + assert_eq!( + restored, + RetainedContext { + next_order: snapshot.next_order, + ..before_restriction + } + ); +} + +#[test] +fn retained_families_enforce_storage_limits_without_changing_snapshots() { + let mut context = RetainedContext::default(); + let first = publish_answer(); + context.record(&first); + let snapshot = context.clone(); for index in 2..=10 { context.record(&RetainedContextEvent::VerifiedAnswer(VerifiedAnswer { @@ -33,17 +70,18 @@ fn answer_lifecycle_is_bounded_idempotent_and_checkpointed() { }], })); } - assert!(!context.is_complete()); - assert_eq!(context.verified_answers().count(), MAX_ANSWERS); - context.retain_answers(|answer| answer.turn_id != "turn-2"); + assert!(!context.verified_answers_complete()); + assert_eq!(context.verified_answers().count(), MAX_FAMILY_RECORDS); + context.rollback(&["turn-2"], /*first_removed_message_id*/ None); assert_eq!(context.verified_answers().count(), 0); - assert!(!context.is_complete()); + assert!(!context.verified_answers_complete()); + let mut restored = snapshot.clone(); let mut oversized = first.clone(); let RetainedContextEvent::VerifiedAnswer(answer) = &mut oversized; - answer.questions[0].answer = "a".repeat(MAX_ANSWER_BYTES); + answer.questions[0].answer = "a".repeat(MAX_RECORD_BYTES); restored.record(&oversized); - assert!(!restored.is_complete()); + assert!(!restored.verified_answers_complete()); assert!( restored .verified_answers() @@ -61,4 +99,54 @@ fn answer_lifecycle_is_bounded_idempotent_and_checkpointed() { .answer, "Yes, but never publicly." ); + for index in 0..=MAX_FAMILY_RECORDS { + restored.record_user_message(RetainedUserMessage { + turn_id: "later-turn".to_owned(), + message_id: Some(format!("message-{index}")), + text: "Keep the repository private.".to_owned(), + complete: true, + }); + } + assert!(!restored.user_messages_complete()); + assert_eq!( + restored + .ordered_entries() + .filter(|entry| matches!(entry, RetainedContextEntry::UserMessage(_))) + .count(), + MAX_FAMILY_RECORDS + ); + restored.record_user_message(RetainedUserMessage { + turn_id: "oversized-turn".to_owned(), + message_id: Some("oversized-message".to_owned()), + text: "restriction ".repeat(MAX_RECORD_BYTES), + complete: true, + }); + let Some(RetainedContextEntry::UserMessage(message)) = restored.ordered_entries().next_back() + else { + panic!("latest user evidence"); + }; + assert_eq!((&message.text, message.complete), (&String::new(), false)); +} + +#[test] +fn legacy_checkpoints_mark_user_messages_incomplete() { + let RetainedContextEvent::VerifiedAnswer(answer) = publish_answer(); + let mut wire = serde_json::json!({ + "verified_answers": [answer], "incomplete": false + }); + let legacy: RetainedContext = + serde_json::from_value(wire.clone()).expect("legacy retained-answer checkpoint"); + assert!(legacy.verified_answers_complete()); + assert!(!legacy.user_messages_complete()); + + // Retain the original wire key even though the internal field names its family. + wire["verified_answers"][0]["order"] = serde_json::json!(0); + wire["user_messages"] = serde_json::json!([]); + wire["user_messages_incomplete"] = serde_json::json!(true); + wire["next_order"] = serde_json::json!(0); + assert_eq!(serde_json::to_value(&legacy).unwrap(), wire); + + let mut restored = RetainedContext::default(); + restored.restore(/*checkpoint*/ None); + assert!(!restored.user_messages_complete()); } diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index 71d67cc63a..390b1d08b6 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -29,6 +29,7 @@ pub use codex_history::CompactedItem; pub use codex_history::InitialHistory; pub use codex_history::ResponseItemEnvelope; pub use codex_history::ResumedHistory; +pub use codex_history::RetainedContextEntry; pub use codex_history::RetainedContextEvent; pub use codex_history::RolloutItem; pub use codex_history::RolloutLine; diff --git a/codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs b/codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs index 44bd5c6937..b16539ea7e 100644 --- a/codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs +++ b/codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs @@ -4,16 +4,19 @@ //! Most records happen to be ordered that way, but late completion events can target an older //! surviving turn after a newer turn has started. This planner keeps compact per-record ownership //! metadata for SQLite visibility, then combines it with `rollback_replay`'s cold-resume answer -//! before the writer makes its second streaming pass. -//! Retained answers follow their source calls, not later steers sharing the same turn ID. +//! before the writer makes its second streaming pass. Retained checkpoints are reduced when +//! rollback is encountered, never by reapplying old rollbacks to later checkpoints. use std::collections::HashMap; +use std::collections::HashSet; +use codex_protocol::ResponseItemId; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::UserMessageEvent; use codex_rollout::CompactedItem; +use codex_rollout::RetainedContextEntry; use codex_rollout::RolloutItem; use codex_rollout::RolloutLine; @@ -36,6 +39,12 @@ struct PendingUserResponse { content: Vec, } +struct InstructionBoundary { + record_index: usize, + message_id: Option, + alive: bool, +} + /// Compact plan keyed by parsed source-record index. pub(super) struct RollbackPlan { record_boundaries: Vec>, @@ -79,7 +88,7 @@ impl RollbackPlan { /// Streaming builder for RollbackPlan. pub(super) struct RollbackPlanner { record_boundaries: Vec>, - boundary_alive: Vec, + boundaries: Vec, boundary_stack: Vec, active_turn_id: Option, pending_turn_records: Vec, @@ -88,6 +97,7 @@ pub(super) struct RollbackPlanner { pending_delivery_boundary: Option, turn_boundaries: HashMap, call_boundaries: HashMap<(String, String), Option>, + retained_fact_sources: Vec<(usize, String)>, compactions: Vec, model_replay: ModelReplayPlanner, } @@ -96,7 +106,7 @@ impl RollbackPlanner { pub(super) fn new() -> Self { Self { record_boundaries: Vec::new(), - boundary_alive: Vec::new(), + boundaries: Vec::new(), boundary_stack: Vec::new(), active_turn_id: None, pending_turn_records: Vec::new(), @@ -105,6 +115,7 @@ impl RollbackPlanner { pending_delivery_boundary: None, turn_boundaries: HashMap::new(), call_boundaries: HashMap::new(), + retained_fact_sources: Vec::new(), compactions: Vec::new(), model_replay: ModelReplayPlanner::new(), } @@ -139,8 +150,10 @@ impl RollbackPlanner { RolloutItem::ResponseItem(response) => { if let Some(boundary) = paired_delivery_boundary { self.record_boundaries[index] = Some(boundary); + self.boundaries[boundary].message_id = response.id().cloned(); } else if rollback::counts_as_boundary(&response.item) { let boundary = self.start_boundary(index); + self.boundaries[boundary].message_id = response.id().cloned(); if let ResponseItem::Message { role, content, .. } = &response.item && role == "user" { @@ -247,6 +260,8 @@ impl RollbackPlanner { } self.call_boundaries .insert(source, self.record_boundaries[index]); + self.retained_fact_sources + .push((index, answer.turn_id.clone())); } RolloutItem::SecurityRiskScore(_) => self.record_boundaries[index] = None, } @@ -257,34 +272,19 @@ impl RollbackPlanner { pub(super) fn finish(self) -> RollbackPlan { let RollbackPlanner { record_boundaries, - boundary_alive, + boundaries, compactions, model_replay, - turn_boundaries, - call_boundaries, .. } = self; let replay_anchor = model_replay.finish().empty_replacement_history_compaction; - let removed_turns = turn_boundaries - .iter() - .filter(|(_, boundary)| !boundary_alive[**boundary]) - .map(|(turn_id, _)| turn_id.as_str()) + let boundary_alive = boundaries + .into_iter() + .map(|boundary| boundary.alive) .collect::>(); let compacted_items = compactions .into_iter() .filter_map(|mut frame| { - // Migration removes rollback markers. Checkpoints must therefore carry - // the same surviving facts as their standalone source events. - if let Some(context) = &mut frame.item.retained_context { - context.retain_answers(|answer| { - call_boundaries - .get(&(answer.turn_id.clone(), answer.call_id.clone())) - .map_or_else( - || !removed_turns.contains(&answer.turn_id.as_str()), - |boundary| boundary.is_none_or(|boundary| boundary_alive[boundary]), - ) - }); - } if Some(frame.record_index) == replay_anchor { frame.item.replacement_history = Some(Vec::new()); frame.item.mcp_resource_origins = None; @@ -304,8 +304,12 @@ impl RollbackPlanner { } fn start_boundary(&mut self, index: usize) -> usize { - let boundary = self.boundary_alive.len(); - self.boundary_alive.push(true); + let boundary = self.boundaries.len(); + self.boundaries.push(InstructionBoundary { + record_index: index, + message_id: None, + alive: true, + }); let had_prior_boundary = !self.boundary_stack.is_empty(); if had_prior_boundary { for pending_index in self.pending_context_records.drain(..) { @@ -348,16 +352,72 @@ impl RollbackPlanner { return Ok(()); } let depth_before = self.boundary_stack.len(); + let mut removed_boundaries = HashSet::new(); + let mut first_removed_boundary = None; for _ in 0..count { let Some(boundary) = self.boundary_stack.pop() else { break; }; - self.boundary_alive[boundary] = false; + self.boundaries[boundary].alive = false; + removed_boundaries.insert(boundary); + first_removed_boundary = Some(boundary); + } + if let Some(boundary) = first_removed_boundary { + let source = &self.boundaries[boundary]; + let removed_turns = self + .turn_boundaries + .iter() + .filter(|(_, boundary)| removed_boundaries.contains(*boundary)) + .map(|(turn_id, _)| turn_id.as_str()) + .chain( + self.retained_fact_sources + .iter() + .filter(|(index, _)| { + self.record_boundaries[*index] + .is_some_and(|boundary| removed_boundaries.contains(&boundary)) + }) + .map(|(_, turn_id)| turn_id.as_str()), + ) + .collect::>(); + // A checkpoint before this instruction cannot contain evidence from it. + // Later checkpoints have not been observed yet and already reflect this rollback. + for frame in self + .compactions + .iter_mut() + .rev() + .take_while(|frame| frame.record_index >= source.record_index) + { + if let Some(context) = &mut frame.item.retained_context { + if context + .ordered_entries() + .any(|entry| matches!(entry, RetainedContextEntry::UserMessage(_))) + { + context.rollback( + &removed_turns, + source.message_id.as_ref().map(ResponseItemId::as_str), + ); + } else { + // Checkpoints written without instruction retention keep the legacy + // source-call boundary, including answers before a same-turn steer. + context.retain_answers(|answer| { + self.call_boundaries + .get(&(answer.turn_id.clone(), answer.call_id.clone())) + .map_or_else( + || !removed_turns.contains(&answer.turn_id.as_str()), + |boundary| { + boundary + .is_none_or(|boundary| self.boundaries[boundary].alive) + }, + ) + }); + } + } + } } let compaction_index = self.compactions.iter().rposition(|frame| { frame .owner - .is_none_or(|boundary| self.boundary_alive[boundary]) + .is_none_or(|boundary| self.boundaries[boundary].alive) }); if let Some(compaction_index) = compaction_index { let frame = &mut self.compactions[compaction_index]; diff --git a/codex-rs/thread-store/src/local/rollout_migration_tests.rs b/codex-rs/thread-store/src/local/rollout_migration_tests.rs index 6ca78276ea..0758db8163 100644 --- a/codex-rs/thread-store/src/local/rollout_migration_tests.rs +++ b/codex-rs/thread-store/src/local/rollout_migration_tests.rs @@ -11,6 +11,7 @@ use codex_extension_items::ExtensionItem; use codex_extension_items::image_generation::ImageGenerationFailure; use codex_extension_items::image_generation::ImageGenerationItem; use codex_protocol::AgentPath; +use codex_protocol::ResponseItemId; use codex_protocol::ThreadId; use codex_protocol::config_types::ReasoningSummary; use codex_protocol::items::ReasoningItem; @@ -18,6 +19,8 @@ use codex_protocol::items::TurnItem; use codex_protocol::mcp::McpResourceOrigin; use codex_protocol::mcp::McpResourceOriginCheckpoint; use codex_protocol::models::ContentItem; +use codex_protocol::models::ContentItemKind; +use codex_protocol::models::InternalChatMessageMetadataPassthrough; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::AgentMessageEvent; use codex_protocol::protocol::AskForApproval; @@ -1267,6 +1270,139 @@ async fn migration_preserves_reverse_replay_anchor_after_pre_compaction_rollback assert!(replacement_history.is_empty()); } +#[tokio::test] +async fn migration_preserves_same_turn_evidence_across_rollback_checkpoints() { + const INITIAL: &str = "Never publish publicly."; + const STEER: &str = "Also inspect the README."; + let home = TempDir::new().expect("create Codex home"); + let thread_id = ThreadId::new(); + let [initial, steer] = + [("initial", INITIAL), ("steer", STEER)].map(|(id, text)| ResponseItem::Message { + id: Some(ResponseItemId::with_suffix("msg", id)), + role: "user".to_owned(), + content: vec![ContentItem::InputText { + text: text.to_owned(), + }], + phase: None, + internal_chat_message_metadata_passthrough: Some( + InternalChatMessageMetadataPassthrough { + turn_id: Some("shared-turn".to_owned()), + content_item_kinds: Some(vec![ContentItemKind("user.text".to_owned())]), + ..Default::default() + }, + ), + }); + let RolloutItem::Compacted(mut checkpoint) = compacted(vec![initial.clone(), steer.clone()]) + else { + unreachable!("compacted helper creates a checkpoint"); + }; + let retained = json!({ + "user_messages": [ + {"order": 0, "turn_id": "shared-turn", "message_id": "msg_initial", "text": INITIAL, "complete": true}, + {"order": 2, "turn_id": "shared-turn", "message_id": "msg_steer", "text": STEER, "complete": true} + ], + "verified_answers": [ + {"order": 1, "turn_id": "shared-turn", "call_id": "before", "questions": [{"question": "Publish?", "answer": "Only privately."}]}, + {"order": 3, "turn_id": "shared-turn", "call_id": "after", "questions": [{"question": "Publish the README?", "answer": "Do not publish it."}]} + ], + "incomplete": false, "user_messages_incomplete": false, "next_order": 4 + }); + checkpoint.retained_context = + Some(serde_json::from_value(retained.clone()).expect("retained fixture")); + let answers = checkpoint + .retained_context + .as_ref() + .expect("retained checkpoint") + .verified_answers() + .cloned() + .map(codex_rollout::RetainedContextEvent::VerifiedAnswer) + .collect::>(); + let mut expected = retained; + expected["user_messages"] + .as_array_mut() + .expect("user messages") + .truncate(/*len*/ 1); + expected["verified_answers"] + .as_array_mut() + .expect("verified answers") + .truncate(/*len*/ 1); + let mut before_expected = expected.clone(); + before_expected["next_order"] = json!(2); + let mut before_steer = checkpoint.clone(); + before_steer.replacement_history = Some(vec![initial.clone().into()]); + before_steer.retained_context = + Some(serde_json::from_value(before_expected.clone()).expect("pre-steer checkpoint")); + let mut after_rollback = before_steer.clone(); + after_rollback.retained_context = + Some(serde_json::from_value(expected.clone()).expect("post-rollback checkpoint")); + let path = write_rollout( + home.path(), + thread_id, + SessionSource::Cli, + vec![ + started("shared-turn"), + rollout_response_item(initial.clone()), + user_message(INITIAL), + RolloutItem::RetainedContext(answers[0].clone()), + RolloutItem::Compacted(before_steer), + rollout_response_item(steer), + user_message(STEER), + RolloutItem::RetainedContext(answers[1].clone()), + completed("shared-turn"), + started("compaction-turn"), + RolloutItem::Compacted(checkpoint), + completed("compaction-turn"), + RolloutItem::EventMsg(EventMsg::ThreadRolledBack(ThreadRolledBackEvent { + num_turns: 1, + })), + // This checkpoint already contains the correct live rollback result. + started("post-rollback-compaction"), + RolloutItem::Compacted(after_rollback), + completed("post-rollback-compaction"), + ], + ); + let store = indexed_store(home.path()).await; + store + .migrate_rollouts(apply_options()) + .await + .expect("migrate same-turn rollback"); + let migrated = read_rollout(&path); + let checkpoints = migrated + .iter() + .filter_map(|line| match &line.item { + RolloutItem::Compacted(item) => Some(item), + _ => None, + }) + .collect::>(); + assert_eq!( + checkpoints + .iter() + .map( + |checkpoint| serde_json::to_value(&checkpoint.retained_context) + .expect("retained context") + ) + .collect::>(), + vec![before_expected, expected.clone(), expected] + ); + assert_eq!( + checkpoints + .last() + .expect("latest checkpoint") + .replacement_history, + Some(vec![initial.into()]) + ); + assert_eq!( + migrated + .into_iter() + .filter_map(|line| match line.item { + RolloutItem::RetainedContext(event) => Some(event), + _ => None, + }) + .collect::>(), + vec![answers[0].clone()] + ); +} + #[tokio::test] async fn migration_keeps_empty_replay_anchor_from_rolled_back_turn() { let home = TempDir::new().expect("create Codex home");