fix: prevent recovery cascade and improve Stop-scenario reasoning lookup

src/deepseek_cursor_proxy/reasoning_store.py

@@ -46,6 +46,17 @@ def tool_call_ids(message: dict[str, Any]) -> list[str]:
     return ids
 
 
+def tool_call_names(message: dict[str, Any]) -> list[str]:
+    names: list[str] = []
+    for tool_call in message.get("tool_calls") or []:
+        if not isinstance(tool_call, dict):
+            continue
+        function = tool_call.get("function")
+        if isinstance(function, dict) and function.get("name"):
+            names.append(str(function["name"]))
+    return names
+
+
 def message_signature(message: dict[str, Any]) -> str:
     tool_calls = [
         normalize_tool_call(tool_call)
@@ -172,6 +183,10 @@ def store_assistant_message(self, message: dict[str, Any], scope: str) -> int:
             for tool_call in (message.get("tool_calls") or [])
             if isinstance(tool_call, dict)
         )
+        keys.extend(
+            f"scope:{scope}:tool_name:{tool_name}"
+            for tool_name in tool_call_names(message)
+        )
         for key in keys:
             self.put(key, reasoning, message)
         return len(keys)
@@ -192,6 +207,10 @@ def lookup_for_message(self, message: dict[str, Any], scope: str) -> str | None:
             )
             if reasoning is not None:
                 return reasoning
+        for tool_name in tool_call_names(message):
+            reasoning = self.get(f"scope:{scope}:tool_name:{tool_name}")
+            if reasoning is not None:
+                return reasoning
         return None
 
     def clear(self) -> int:

src/deepseek_cursor_proxy/server.py

@@ -266,6 +266,7 @@ def do_POST(self) -> None:
                     prepared.payload["messages"],
                     prepared.cache_namespace,
                     prepared.recovery_notice,
+                    prepared.record_response_scope,
                 )
             else:
                 sent_response = self._proxy_regular_response(
@@ -274,6 +275,7 @@ def do_POST(self) -> None:
                     prepared.payload["messages"],
                     prepared.cache_namespace,
                     prepared.recovery_notice,
+                    prepared.record_response_scope,
                 )
             if not sent_response:
                 return
@@ -436,6 +438,7 @@ def _proxy_regular_response(
         request_messages: list[dict[str, Any]],
         cache_namespace: str,
         recovery_notice: str | None = None,
+        record_response_scope: str | None = None,
     ) -> bool:
         body = read_response_body(response)
         try:
@@ -446,6 +449,7 @@ def _proxy_regular_response(
                 request_messages,
                 cache_namespace,
                 content_prefix=recovery_notice,
+                scope=record_response_scope,
             )
         except (json.JSONDecodeError, UnicodeDecodeError) as exc:
             LOG.warning("failed to rewrite upstream JSON response: %s", exc)
@@ -476,6 +480,7 @@ def _proxy_streaming_response(
         request_messages: list[dict[str, Any]],
         cache_namespace: str,
         recovery_notice: str | None = None,
+        record_response_scope: str | None = None,
     ) -> bool:
         sent_headers = self._send_response_headers(
             getattr(response, "status", 200),
@@ -496,38 +501,48 @@ def _proxy_streaming_response(
             if self.config.cursor_display_reasoning
             else None
         )
-        scope = conversation_scope(request_messages, cache_namespace)
+        scope = (
+            record_response_scope
+            if record_response_scope is not None
+            else conversation_scope(request_messages, cache_namespace)
+        )
         finalized = False
         pending_recovery_notice = recovery_notice
-        while True:
-            try:
-                line = response.readline()
-            except (HTTPException, OSError) as exc:
-                LOG.warning("upstream streaming response read failed: %s", exc)
-                return False
-            if not line:
-                break
-            rewritten, finalized, pending_recovery_notice = self._rewrite_sse_line(
-                line,
-                original_model,
-                accumulator,
-                scope,
-                display_adapter,
-                pending_recovery_notice,
-            )
-            if not self._write_to_client(
-                rewritten, "sending streaming response chunk", flush=True
-            ):
-                return False
-            if finalized:
-                break
-
-        if not finalized:
-            if self.config.verbose:
-                log_json("model streaming assistant messages", accumulator.messages())
-            stored = accumulator.store_reasoning(self.reasoning_store, scope)
-            if stored:
-                LOG.info("stored %s streaming reasoning cache key(s)", stored)
+        try:
+            while True:
+                try:
+                    line = response.readline()
+                except (HTTPException, OSError) as exc:
+                    LOG.warning("upstream streaming response read failed: %s", exc)
+                    return False
+                if not line:
+                    break
+                rewritten, finalized, pending_recovery_notice = self._rewrite_sse_line(
+                    line,
+                    original_model,
+                    accumulator,
+                    scope,
+                    display_adapter,
+                    pending_recovery_notice,
+                )
+                if not self._write_to_client(
+                    rewritten, "sending streaming response chunk", flush=True
+                ):
+                    return False
+                if finalized:
+                    break
+        finally:
+            if not finalized:
+                if self.config.verbose:
+                    log_json(
+                        "model streaming assistant messages", accumulator.messages()
+                    )
+                stored = accumulator.store_reasoning(self.reasoning_store, scope)
+                if stored:
+                    LOG.info(
+                        "stored %s streaming reasoning cache key(s) before exit",
+                        stored,
+                    )
         return True
 
     def _rewrite_sse_line(

src/deepseek_cursor_proxy/transform.py

@@ -95,6 +95,7 @@ class PreparedRequest:
     recovered_reasoning_messages: int = 0
     recovery_dropped_messages: int = 0
     recovery_notice: str | None = None
+    record_response_scope: str | None = None
 
 
 def normalize_reasoning_effort(value: Any) -> str:
@@ -485,6 +486,8 @@ def prepare_upstream_request(
         repair_reasoning=thinking_enabled,
         keep_reasoning=not thinking_disabled,
     )
+    record_response_scope = conversation_scope(messages, cache_namespace)
+
     recovered_count = 0
     recovery_dropped_messages = 0
     recovery_notice = None
@@ -517,6 +520,7 @@ def prepare_upstream_request(
         recovered_reasoning_messages=recovered_count,
         recovery_dropped_messages=recovery_dropped_messages,
         recovery_notice=recovery_notice,
+        record_response_scope=record_response_scope,
     )
 
 
@@ -525,20 +529,23 @@ def record_response_reasoning(
     store: ReasoningStore | None,
     request_messages: list[dict[str, Any]],
     cache_namespace: str = "",
+    scope: str | None = None,
 ) -> int:
     if store is None:
         return 0
     stored = 0
     choices = response_payload.get("choices")
     if not isinstance(choices, list):
         return stored
-    scope = conversation_scope(request_messages, cache_namespace)
+    response_scope = scope if scope is not None else conversation_scope(
+        request_messages, cache_namespace
+    )
     for choice in choices:
         if not isinstance(choice, dict):
             continue
         message = choice.get("message")
         if isinstance(message, dict):
-            stored += store.store_assistant_message(message, scope)
+            stored += store.store_assistant_message(message, response_scope)
     return stored
 
 
@@ -549,13 +556,18 @@ def rewrite_response_body(
     request_messages: list[dict[str, Any]],
     cache_namespace: str = "",
     content_prefix: str | None = None,
+    scope: str | None = None,
 ) -> bytes:
     response_payload = json.loads(body.decode("utf-8"))
     if isinstance(response_payload, dict):
         if content_prefix:
             prefix_response_content(response_payload, content_prefix)
         record_response_reasoning(
-            response_payload, store, request_messages, cache_namespace
+            response_payload,
+            store,
+            request_messages,
+            cache_namespace,
+            scope=scope,
         )
         if "model" in response_payload:
             response_payload["model"] = original_model