[https://nvbugs/6312828][fix] Recover disaggregated requests after timeouts

cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp

@@ -102,6 +102,30 @@ using RequestIdType = LlmRequest::RequestIdType;
 
 constexpr int kTransferFuturePollIntervalMs = 10;
 
+template <typename T>
+void reserveForAppend(std::vector<T>& values, std::size_t additional)
+{
+    TLLM_CHECK_WITH_INFO(
+        additional <= values.max_size() - values.size(), "Cannot reserve %zu additional transfer records", additional);
+    auto const requiredCapacity = values.size() + additional;
+    if (requiredCapacity <= values.capacity())
+    {
+        return;
+    }
+
+    auto newCapacity = std::max<std::size_t>(1, values.capacity());
+    while (newCapacity < requiredCapacity)
+    {
+        if (newCapacity > values.max_size() / 2)
+        {
+            newCapacity = requiredCapacity;
+            break;
+        }
+        newCapacity *= 2;
+    }
+    values.reserve(newCapacity);
+}
+
 // Finite status checks are scheduler polls, not terminal deadlines. Pure polls
 // use short slices; calls that ask for at least one completion keep bounded
 // backpressure by waiting up to the configured future timeout.
@@ -777,6 +801,10 @@ void CacheTransceiver::respondAndSendAsync(std::shared_ptr<LlmRequest> llmReques
         }
         return;
     }
+    // Reserve the durable status record before handing the request to the
+    // sender. Once sendAsync succeeds it may own request-backed KV pages, and
+    // the noexcept moves below must not lose the only future that can reap it.
+    reserveForAppend(mSenderFutures, 1);
     setContextState(llmRequest.get());
     auto future = mCacheSender->sendAsync(llmRequest);
     mSenderFutures.emplace_back(std::move(llmRequest), std::move(future));
@@ -785,6 +813,7 @@ void CacheTransceiver::respondAndSendAsync(std::shared_ptr<LlmRequest> llmReques
 void CacheTransceiver::respondAndSendLayerWise(
     RequestVector const& requests, std::shared_ptr<ContextProgress> const& progress)
 {
+    reserveForAppend(mSenderFutures, requests.size());
     for (auto const& llmRequest : requests)
     {
         TLLM_CHECK(llmRequest && llmRequest->isContextOnlyRequest());
@@ -1578,6 +1607,10 @@ bool CacheTransceiver::cancelRequest(std::shared_ptr<LlmRequest> llmRequest)
     }
     if (llmRequest->isContextOnlyRequest())
     {
+        // Keep the high-level future until checkContextTransferStatus records
+        // this rank's terminal outcome. Queued/current classification can
+        // differ across model-parallel ranks, so rank-local reaping would
+        // remove the evidence needed by the existing transfer consensus.
         return mCacheSender->cancelRequest(*llmRequest);
     }
     else if (llmRequest->isGenerationOnlyRequest())

cpp/tensorrt_llm/batch_manager/dataTransceiver.h

@@ -16,6 +16,8 @@
  */
 
 #pragma once
+#include <chrono>
+#include <cstddef>
 #include <fstream>
 #include <future>
 #include <map>
@@ -38,6 +40,11 @@
 #include "tensorrt_llm/runtime/cudaEvent.h"
 #include "tensorrt_llm/runtime/utils/mpiUtils.h"
 
+namespace tensorrt_llm::testing
+{
+class CacheSenderTestAccess;
+} // namespace tensorrt_llm::testing
+
 namespace tensorrt_llm::batch_manager
 {
 
@@ -330,6 +337,11 @@ class CacheSender
     virtual ~CacheSender();
 
 private:
+    friend class ::tensorrt_llm::testing::CacheSenderTestAccess;
+
+    CacheSender(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer,
+        std::size_t maxPendingPreHandshakeCancellations, std::chrono::milliseconds preHandshakeCancellationRetention);
+
     class Impl;
 
     struct ImplDeleter

cpp/tests/unit_tests/executor/agentCommTest.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,9 +15,13 @@
  * limitations under the License.
  */
 
+#include "tensorrt_llm/batch_manager/cacheFormatter.h"
 #include "tensorrt_llm/common/tllmDataType.h"
 #include "tensorrt_llm/executor/cache_transmission/agent_utils/connection.h"
+#include <chrono>
+#include <future>
 #include <gtest/gtest.h>
+#include <optional>
 
 using namespace tensorrt_llm::batch_manager::kv_cache_manager;
 using namespace tensorrt_llm::runtime;
@@ -232,5 +236,68 @@ TEST_P(AgentCommTest, AgentConnectionManagerConnect)
     TLLM_LOG_INFO("after finish");
 }
 
+TEST_P(AgentCommTest, CacheSenderRejectsStandaloneFinalizedReplay)
+{
+    namespace tbm = tensorrt_llm::batch_manager;
+    namespace tr = tensorrt_llm::runtime;
+    namespace texec = tensorrt_llm::executor;
+
+    std::vector<tbm::BaseTransBufferManager*> bufferManagers{mTransBufferManager.get()};
+    auto receiverManager = std::make_unique<AgentConnectionManager>(bufferManagers, *mCacheState, backend);
+    auto senderManager = std::make_unique<AgentConnectionManager>(bufferManagers, *mCacheState, backend);
+    auto receiverCommState = receiverManager->getCommState();
+    auto senderCommState = senderManager->getCommState();
+
+    auto* receiverConnection = const_cast<AgentConnection*>(
+        dynamic_cast<AgentConnection const*>(receiverManager->getConnections(senderCommState).at(0)));
+    ASSERT_NE(receiverConnection, nullptr);
+
+    std::vector<tbm::kv_cache_manager::CacheTransBufferManager*> cacheBufferManagers{mTransBufferManager.get()};
+    tbm::CacheSender sender(senderManager.get(), 0,
+        tbm::CacheTransferLayer(*mCacheState,
+            tbm::kv_cache_manager::createCacheFormatter(mCacheManager.get(), cacheBufferManagers, false)));
+
+    tbm::LlmRequest::RequestIdType constexpr requestId{29};
+    tr::SamplingConfig const samplingConfig{1};
+    auto const inputTokens = std::make_shared<tbm::LlmRequest::VecTokens>(tbm::LlmRequest::VecTokens{1});
+    auto request = std::make_shared<tbm::LlmRequest>(requestId, 1, inputTokens, samplingConfig, /*isStreaming=*/false);
+    auto responseFuture = sender.sendAsync(request);
+    ASSERT_TRUE(sender.cancelRequest(*request));
+    ASSERT_EQ(responseFuture.wait_for(std::chrono::seconds{10}), std::future_status::ready);
+    EXPECT_THROW(responseFuture.get(), std::exception);
+
+    texec::DataTransceiverState receiverState{*mCacheState, receiverCommState};
+    tbm::RequestInfo requestInfo{requestId, receiverState};
+    std::vector<std::optional<size_t>> cacheBufferIds{std::optional<size_t>{0}};
+    int constexpr validConnectionIdx{0};
+    std::atomic<bool> terminate{false};
+    int32_t constexpr kDataTag{43};
+    int32_t constexpr dataTag = ((requestId & 0xFFF) << 8) | (kDataTag & 0xFF);
+
+    auto sendRequestAndAwaitRejection = [&]() -> std::optional<bool>
+    {
+        receiverConnection->sendRequestAndBufferInfo(requestInfo, cacheBufferIds, validConnectionIdx);
+        auto readyFuture = std::async(std::launch::async,
+            [&]() {
+                return receiverConnection->recvReadySignal(DataContext{dataTag, terminate});
+            });
+        if (readyFuture.wait_for(std::chrono::seconds{10}) != std::future_status::ready)
+        {
+            terminate.store(true, std::memory_order_relaxed);
+            return std::nullopt;
+        }
+        return readyFuture.get();
+    };
+
+    // The first handshake consumes the queued cancellation and finalizes the
+    // request. The replay must be rejected while the sender has no local work.
+    auto firstReady = sendRequestAndAwaitRejection();
+    ASSERT_TRUE(firstReady.has_value());
+    EXPECT_FALSE(*firstReady);
+    auto replayReady = sendRequestAndAwaitRejection();
+    ASSERT_TRUE(replayReady.has_value());
+    EXPECT_FALSE(*replayReady);
+}
+
 INSTANTIATE_TEST_SUITE_P(AvailableBackends, AgentCommTest, ::testing::ValuesIn(getAvailableBackends()),
     [](::testing::TestParamInfo<AgentCommTest::ParamType> const& info) { return info.param; });