fix: cancellation kv commit router (#7178)

Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com>

lib/kv-router/src/protocols.rs

@@ -140,7 +140,12 @@ pub enum RouterRequest {
         block_mm_infos: Option<Vec<Option<BlockExtraInfo>>>,
     },
     MarkPrefill,
-    MarkFree,
+    MarkFree {
+        // once request is cancelled, the frontend might not be allowed to send a
+        // request with linking the id. In this case, the request_id is provided in the payload.
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        request_id: Option<String>,
+    },
 }
 
 impl Default for RouterRequest {
@@ -870,4 +875,35 @@ mod tests {
         assert_eq!(deserialized.block_hashes[0].0, 4);
         assert_eq!(deserialized.block_hashes[1].0, 5);
     }
+
+    #[test]
+    fn test_router_request_mark_free_backwards_compatible_deserialization() {
+        let request: RouterRequest = serde_json::from_str(r#"{"method":"mark_free"}"#).unwrap();
+
+        assert!(matches!(
+            request,
+            RouterRequest::MarkFree { request_id: None }
+        ));
+    }
+
+    #[test]
+    fn test_router_request_mark_free_serialization_with_request_id() {
+        let request = RouterRequest::MarkFree {
+            request_id: Some("req-123".to_string()),
+        };
+
+        let serialized = serde_json::to_string(&request).unwrap();
+        let deserialized: RouterRequest = serde_json::from_str(&serialized).unwrap();
+
+        assert_eq!(
+            serialized,
+            r#"{"method":"mark_free","request_id":"req-123"}"#
+        );
+        assert!(matches!(
+            deserialized,
+            RouterRequest::MarkFree {
+                request_id: Some(ref request_id)
+            } if request_id == "req-123"
+        ));
+    }
 }

lib/llm/src/kv_router.rs

@@ -599,9 +599,15 @@ impl AsyncEngine<SingleIn<RouterRequest>, ManyOut<Annotated<RouterResponse>>, Er
             RouterRequest::MarkPrefill => RouterResponse::PrefillMarked {
                 success: self.mark_prefill_completed(&context_id).await.is_ok(),
             },
-            RouterRequest::MarkFree => RouterResponse::FreeMarked {
-                success: self.free(&context_id).await.is_ok(),
-            },
+            RouterRequest::MarkFree { request_id } => {
+                let request_id = match request_id.as_deref() {
+                    Some(request_id) if !request_id.trim().is_empty() => request_id,
+                    _ => &context_id,
+                };
+                RouterResponse::FreeMarked {
+                    success: self.free(request_id).await.is_ok(),
+                }
+            }
         };
 
         let response = Annotated::from_data(response);