From 528fed23d055590273d0a860690655421d6bd2b1 Mon Sep 17 00:00:00 2001
From: Afloat <ShuffleMonkey@163.com>
Date: Wed, 11 Nov 2026 17:39:37 +0800
Subject: [PATCH] fix(replay): report pending weight failures before returning
 success

---
 edge0-infer/src/app/generation.cpp            | 17 +++++++++++++---
 .../tests/unit/app/generation_test.cpp        | 20 +++++++++++++++++++
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/edge0-infer/src/app/generation.cpp b/edge0-infer/src/app/generation.cpp
index 91a6ba35..7e8030fd 100644
--- a/edge0-infer/src/app/generation.cpp
+++ b/edge0-infer/src/app/generation.cpp
@@ -967,9 +967,20 @@ Result<GenerationCheckResult> replay_check(const SessionReference& session, cons
     }
     struct FinishReplay {
         SessionReference session;
-        bool success=false;
+        bool drained=false;
+        exec::Status drain(models::FinishReason reason) {
+            if (drained) return {};
+            drained=true;
+            const auto result=session->model->compute->finish_pending_weights(session->model->loaded,session->state->states(),reason);
+            if (!result) {
+                session->weights_failed=true;
+                session->state->invalidate();
+                return result.error();
+            }
+            return {};
+        }
         ~FinishReplay() {
-            session->model->compute->finish_pending_weights(session->model->loaded, session->state->states(), success?models::FinishReason::normal:models::FinishReason::failure);
+            drain(models::FinishReason::failure);
             {
                 std::lock_guard<std::mutex> lock(session->model->registry->progress_mutex);
                 session->model->registry->progress = {};
@@ -1007,10 +1018,10 @@ Result<GenerationCheckResult> replay_check(const SessionReference& session, cons
         output.result.generated.push_back(greedy.value().token);
         output.result.computed_tokens += block.length;
     }
+    if (auto error=finish.drain(models::FinishReason::normal)) return *error;
     session->blocks = input.blocks; output.replay = input;
     output.result.e = session->state->query().e; output.result.input_complete = true;
     output.result.stop_reason = StopReason::length;
-    finish.success=true;
     return output;
 }
 }
diff --git a/edge0-infer/tests/unit/app/generation_test.cpp b/edge0-infer/tests/unit/app/generation_test.cpp
index 15b89e12..63ce1993 100644
--- a/edge0-infer/tests/unit/app/generation_test.cpp
+++ b/edge0-infer/tests/unit/app/generation_test.cpp
@@ -428,6 +428,26 @@ void thinking_tools_tests() {
     }
 }
 int main() {
+    {
+        Fixture f;f.compute.tail_failure=true;
+        const ReplayInput input{{0},{{0,1,BlockPhase::prefill,true}}};
+        const auto result=replay_check(f.session,input);
+        CHECK(!result && result.error().message=="tail read failed");
+        CHECK(!f.compute.pending && f.compute.finishes.size()==1);
+        CHECK(f.compute.finishes.back()==models::FinishReason::normal);
+        CHECK(!take(query_session_check(f.session)).valid);
+        const auto again=generate_tokens_check(f.session,{0},f.config);
+        CHECK(!again && again.error().category==ErrorCategory::session_invalid);
+    }
+    {
+        Fixture f;f.compute.fail_at=0;f.compute.finish_failure=true;
+        const ReplayInput input{{0},{{0,1,BlockPhase::prefill,true}}};
+        const auto result=replay_check(f.session,input);
+        CHECK(!result && result.error().message=="before submit");
+        CHECK(f.compute.finishes.size()==1 && f.compute.finishes.back()==models::FinishReason::failure);
+        CHECK(!take(query_session_check(f.session)).valid);
+        f.compute.finish_failure=false;
+    }
     CHECK(app::reuse_method(session::ReuseKind::first)==ReuseMethod::first);
     CHECK(app::reuse_method(session::ReuseKind::keep)==ReuseMethod::keep);
     CHECK(app::reuse_method(session::ReuseKind::truncate)==ReuseMethod::truncate);