1616token ids and a fingerprint of the response. On the next request each prior
1717assistant turn is replaced with a sentinel, the conversation is rendered once,
1818and the rendered text is split on the sentinels with the stored ids spliced back
19- in -- only for turns whose fingerprint matches the incoming message (an edited,
20- branched, or reused history is never substituted with stale ids) and whose ids
21- are present (a stop-trimmed turn is left as text) . The worker's exact-token
22- prefix check is the final backstop .
19+ in -- only for turns whose content/tool fingerprint and any supplied reasoning string
20+ match the recorded response, and whose ids are present (a stop-trimmed turn is
21+ left as text). An edit invalidates that turn and all later records . The worker's
22+ exact-token prefix check separately protects KV reuse .
2323"""
2424
2525import hashlib
@@ -111,9 +111,8 @@ def __init__(self, template: ChatTemplate):
111111 # boundary verification.
112112 _header_fn = getattr (template , "assistant_header" , None )
113113 self ._assist_hdr = _header_fn () if _header_fn else _ASSIST_HDR
114- # session_id -> [{"fp": str, "ids": list[int] | None}, ...] (one per
115- # assistant turn we produced, in order). Cleared on reset/close.
116- self ._turns : dict [str , list [dict ]] = {}
114+ # Keyed by assistant-turn index, including gaps after invalidation.
115+ self ._turns : dict [str , dict [int , dict ]] = {}
117116
118117 @staticmethod
119118 def _assistant_fingerprint (content , tool_calls ) -> str :
@@ -135,19 +134,24 @@ def _assistant_fingerprint(content, tool_calls) -> str:
135134 blob = json .dumps ([content or "" , norm ], sort_keys = True , ensure_ascii = False )
136135 return hashlib .sha1 (blob .encode ("utf-8" )).hexdigest ()
137136
137+ @staticmethod
138+ def _reasoning_fingerprint (reasoning_content : Optional [str ]) -> Optional [bytes ]:
139+ if reasoning_content is None :
140+ return None
141+ return hashlib .sha256 (reasoning_content .encode ("utf-8" )).digest ()
142+
138143 def _normalize_scaffold (self , text_chunk : str , preamble : str ) -> Optional [str ]:
139144 """Force the scaffold region (between the last assistant header in
140145 `text_chunk` and its end) to equal `preamble`, so the worker re-tokenizes
141146 the exact resident scaffold. The region is empty (history stripped it ->
142147 insert) or a think scaffold (history preserved it -> replace). Returns the
143148 adjusted text, or None if it isn't a recognized scaffold (-> text fallback)."""
149+ if not self ._assist_hdr :
150+ return None
144151 h = text_chunk .rfind (self ._assist_hdr )
145152 if h == - 1 :
146- # No assistant header: with a scaffold to reproduce this is
147- # ambiguous (-> text fallback); without one there is nothing to
148- # normalize, so splicing still works for templates with a different
149- # assistant header.
150- return None if preamble else text_chunk
153+ # Without a verified boundary, splicing can duplicate template framing.
154+ return None
151155 base = h + len (self ._assist_hdr )
152156 if not preamble :
153157 # No generation scaffold: the worker prefills nothing ahead of the
@@ -266,40 +270,37 @@ def build_prompt_input(
266270 """Return a PromptInput: token-ID segments when this session has faithful
267271 stored ids for matching prior assistant turns, else the plain rendered
268272 text. Each incoming assistant turn is matched IN ORDER against the stored
269- records and only spliced when (a) its fingerprint matches what we returned
270- (else the history diverged -> stop, splice nothing further) and (b) we
271- kept faithful ids for it (a stop-trimmed turn's None -> rendered as text) .
273+ records and only spliced when its content/tool calls and any supplied
274+ reasoning string match what we returned, and we kept faithful ids for it.
275+ Omitted or null reasoning permits reuse; a string edit invalidates the tail .
272276 Falls back to text on a sentinel collision or a render that
273277 dropped/duplicated a sentinel."""
274278 stored = self ._turns .get (session_id or "" )
275279 if not stored :
276280 return PromptInput (text = rendered_prompt )
277- # Positional: stored[k] is the k-th assistant turn WE generated, matched
278- # against the k-th assistant message in the request. A client-injected
279- # turn (few-shot exemplar, pre-seeded turn, reused session) shifts that
280- # alignment -> fingerprint mismatch at k -> stop splicing. Always safe
281- # (text fallback + worker prefix backstop); just a lower hit rate.
281+ # Missing records render as text without shifting later turn indices.
282282 positions = [i for i , m in enumerate (messages ) if m .role == "assistant" ]
283283 splice : dict [int , dict ] = {} # message index -> {"ids", "preamble"}
284- diverged_at = None
285284 for k , pos in enumerate (positions ):
286- if k >= len (stored ):
287- break
285+ record = stored .get (k )
286+ if record is None :
287+ continue
288288 m = messages [pos ]
289- if self ._assistant_fingerprint (m .content , m .tool_calls ) != stored [k ]["fp" ]:
290- diverged_at = k # this stored turn and every later one are stale
289+ if self ._assistant_fingerprint (m .content , m .tool_calls ) != record ["fp" ] or (
290+ m .reasoning_content is not None
291+ and self ._reasoning_fingerprint (m .reasoning_content )
292+ != record ["reasoning_fp" ]
293+ ):
294+ # Discard the stale tail without shifting subsequent turn indices.
295+ self ._turns [session_id or "" ] = {
296+ index : record for index , record in stored .items () if index < k
297+ }
291298 break
292- if stored [ k ] ["ids" ] is not None :
299+ if record ["ids" ] is not None :
293300 splice [pos ] = {
294- "ids" : stored [ k ] ["ids" ],
295- "preamble" : stored [ k ] .get ("preamble" , "" ),
301+ "ids" : record ["ids" ],
302+ "preamble" : record .get ("preamble" , "" ),
296303 }
297- if diverged_at is not None :
298- # Drop the stale tail from the first mismatch so an edited/branched
299- # earlier turn can't shadow future requests; the matched prefix still
300- # splices, the rest stays text until reset/close. Safe either way:
301- # stale ids are never spliced and the worker's prefix check backstops.
302- del stored [diverged_at :]
303304 if not splice :
304305 return PromptInput (text = rendered_prompt )
305306 tool_splice = {
@@ -351,25 +352,28 @@ def record_assistant_turn(
351352 generated_token_ids : list ,
352353 prior_turns : int ,
353354 preamble : str = "" ,
355+ reasoning_content : Optional [str ] = None ,
354356 ) -> None :
355357 """Record this turn's {fingerprint, generated ids, generation preamble} at
356358 `prior_turns` (the assistant-turn count of the request it answers).
357359 Records at/after that index are dropped first, so a regenerated/branched
358360 turn replaces stale records rather than shadowing later hits. ids is None
359361 when the worker omitted them (stop-trimmed -> non-resumable), kept for
360- positional alignment. `preamble` is the generation scaffold (e.g. the
361- Qwen3 `<think>` block) reproduced ahead of the spliced ids next request."""
362+ positional alignment. `reasoning_content` is the client-visible value,
363+ including None when the client opted out. `preamble` is the generation
364+ scaffold (e.g. the Qwen3 `<think>` block) reproduced ahead of the spliced
365+ ids next request."""
362366 if not session_id :
363367 return
364- turns = self ._turns .setdefault (session_id , [] )
365- del turns [ prior_turns :]
366- turns . append (
367- {
368- "fp" : self ._assistant_fingerprint (content , tool_calls ),
369- "ids "list ( generated_token_ids ) if generated_token_ids else None ,
370- "preamble" : preamble ,
371- }
372- )
368+ turns = self ._turns .setdefault (session_id , {} )
369+ for index in [ index for index in turns if index >= prior_turns ]:
370+ del turns [ index ]
371+ turns [ prior_turns ] = {
372+ "fp" : self ._assistant_fingerprint (content , tool_calls ),
373+ "reasoning_fp "self . _reasoning_fingerprint ( reasoning_content ) ,
374+ "ids" : list ( generated_token_ids ) if generated_token_ids else None ,
375+ "preamble" : preamble ,
376+ }
373377
374378 def reset (self , session_id : str ) -> None :
375379 self ._turns .pop (session_id , None )
0 commit comments