diff --git a/codec2/hidl/1.0/utils/types.cpp b/codec2/hidl/1.0/utils/types.cpp
index 56db4d5..90ceb13 100644
--- a/codec2/hidl/1.0/utils/types.cpp
+++ b/codec2/hidl/1.0/utils/types.cpp
@@ -1693,11 +1693,11 @@
     if (result != OK) {
         ALOGW("attachToBufferQueue -- attachBuffer failed. Error code = %d",
                 static_cast<int>(result));
-        return false;
+        return result;
     }
     ALOGV("attachToBufferQueue -- attachBuffer returned slot %d",
             static_cast<int>(*bqSlot));
-    return true;
+    return OK;
 }
 
 bool getBufferQueueAssignment(const C2ConstGraphicBlock& block,
@@ -1779,7 +1779,7 @@
             static_cast<int>(bqSlot),
             static_cast<unsigned>(generation));
     _C2BlockFactory::AssignBlockToBufferQueue(
-            data, getHgbp(igbp), bqId, bqSlot, true);
+            data, getHgbp(igbp), generation, bqId, bqSlot, true);
     return true;
 }
 
diff --git a/codec2/hidl/client/client.cpp b/codec2/hidl/client/client.cpp
index ff67681..9c041cf 100644
--- a/codec2/hidl/client/client.cpp
+++ b/codec2/hidl/client/client.cpp
@@ -1177,7 +1177,7 @@
         uint32_t outputGeneration = mOutputGeneration;
         mOutputBufferQueueMutex.unlock();
 
-        status_t status = !attachToBufferQueue(block,
+        status_t status = attachToBufferQueue(block,
                                                outputIgbp,
                                                outputGeneration,
                                                &bqSlot);
@@ -1209,14 +1209,15 @@
         return NO_INIT;
     }
 
-    if (bqId != outputBqId) {
-        ALOGV("queueToOutputSurface -- bufferqueue ids mismatch.");
-        return DEAD_OBJECT;
-    }
-
-    if (generation != outputGeneration) {
-        ALOGV("queueToOutputSurface -- generation numbers mismatch.");
-        return DEAD_OBJECT;
+    if (bqId != outputBqId || generation != outputGeneration) {
+        if (!holdBufferQueueBlock(block, mOutputIgbp, mOutputBqId, mOutputGeneration)) {
+            ALOGE("queueToOutputSurface -- migration fialed");
+            return DEAD_OBJECT;
+        }
+        if (!getBufferQueueAssignment(block, &generation, &bqId, &bqSlot)) {
+            ALOGE("queueToOutputSurface -- corrupted bq assignment");
+            return UNKNOWN_ERROR;
+        }
     }
 
     status_t status = outputIgbp->queueBuffer(static_cast<int>(bqSlot),