diff --git a/media/codec2/sfplugin/C2OMXNode.cpp b/media/codec2/sfplugin/C2OMXNode.cpp
index 03d859a..962df0f 100644
--- a/media/codec2/sfplugin/C2OMXNode.cpp
+++ b/media/codec2/sfplugin/C2OMXNode.cpp
@@ -272,19 +272,14 @@
     work->input.buffers.clear();
     if (block) {
         std::shared_ptr<C2Buffer> c2Buffer(
-                // TODO: fence
                 new Buffer2D(block->share(
-                        C2Rect(block->width(), block->height()), ::C2Fence())),
-                [buffer, source = getSource()](C2Buffer *ptr) {
-                    delete ptr;
-                    // TODO: fence
-                    (void)source->onInputBufferEmptied(buffer, -1);
-                });
+                        C2Rect(block->width(), block->height()), ::C2Fence())));
         work->input.buffers.push_back(c2Buffer);
     }
     work->worklets.clear();
     work->worklets.emplace_back(new C2Worklet);
     std::list<std::unique_ptr<C2Work>> items;
+    uint64_t index = work->input.ordinal.frameIndex.peeku();
     items.push_back(std::move(work));
 
     c2_status_t err = comp->queue(&items);
@@ -292,6 +287,7 @@
         return UNKNOWN_ERROR;
     }
 
+    (void)mBufferIdsInUse.emplace(index, buffer);
     return OK;
 }
 
@@ -326,4 +322,18 @@
     mHeight = height;
 }
 
+void C2OMXNode::onInputBufferDone(c2_cntr64_t index) {
+    if (!mBufferSource) {
+        ALOGD("Buffer source not set (index=%llu)", index.peekull());
+        return;
+    }
+    auto it = mBufferIdsInUse.find(index.peeku());
+    if (it == mBufferIdsInUse.end()) {
+        ALOGV("Untracked input index %llu (maybe already removed)", index.peekull());
+        return;
+    }
+    (void)mBufferSource->onInputBufferEmptied(it->second, -1);
+    (void)mBufferIdsInUse.erase(it);
+}
+
 }  // namespace android
diff --git a/media/codec2/sfplugin/C2OMXNode.h b/media/codec2/sfplugin/C2OMXNode.h
index b5a815e..b7bd696 100644
--- a/media/codec2/sfplugin/C2OMXNode.h
+++ b/media/codec2/sfplugin/C2OMXNode.h
@@ -75,9 +75,23 @@
             OMX_INDEXTYPE *index) override;
     status_t dispatchMessage(const omx_message &msg) override;
 
+    /**
+     * Returns underlying IOMXBufferSource object.
+     */
     sp<IOMXBufferSource> getSource();
+
+    /**
+     * Configure the frame size.
+     */
     void setFrameSize(uint32_t width, uint32_t height);
 
+    /**
+     * Clean up work item reference.
+     *
+     * \param index input work index
+     */
+    void onInputBufferDone(c2_cntr64_t index);
+
 private:
     std::weak_ptr<Codec2Client::Component> mComp;
     sp<IOMXBufferSource> mBufferSource;
@@ -96,6 +110,8 @@
     bool mFirstInputFrame; // true for first input
     c2_cntr64_t mPrevInputTimestamp; // input timestamp for previous frame
     c2_cntr64_t mPrevCodecTimestamp; // adjusted (codec) timestamp for previous frame
+
+    std::map<uint64_t, buffer_id> mBufferIdsInUse;
 };
 
 }  // namespace android
diff --git a/media/codec2/sfplugin/CCodec.cpp b/media/codec2/sfplugin/CCodec.cpp
index 9c679aa..8474ce8 100644
--- a/media/codec2/sfplugin/CCodec.cpp
+++ b/media/codec2/sfplugin/CCodec.cpp
@@ -370,6 +370,10 @@
         return err;
     }
 
+    void onInputBufferDone(c2_cntr64_t index) override {
+        mNode->onInputBufferDone(index);
+    }
+
 private:
     sp<BGraphicBufferSource> mSource;
     sp<C2OMXNode> mNode;
@@ -1583,6 +1587,13 @@
 
 void CCodec::onInputBufferDone(uint64_t frameIndex, size_t arrayIndex) {
     mChannel->onInputBufferDone(frameIndex, arrayIndex);
+    if (arrayIndex == 0) {
+        // We always put no more than one buffer per work, if we use an input surface.
+        Mutexed<Config>::Locked config(mConfig);
+        if (config->mInputSurface) {
+            config->mInputSurface->onInputBufferDone(frameIndex);
+        }
+    }
 }
 
 void CCodec::onMessageReceived(const sp<AMessage> &msg) {
@@ -1715,6 +1726,9 @@
                     ++stream;
                 }
             }
+            if (config->mInputSurface) {
+                config->mInputSurface->onInputBufferDone(work->input.ordinal.frameIndex);
+            }
             mChannel->onWorkDone(
                     std::move(work), changed ? config->mOutputFormat : nullptr,
                     initData.hasChanged() ? initData.update().get() : nullptr);
diff --git a/media/codec2/sfplugin/InputSurfaceWrapper.h b/media/codec2/sfplugin/InputSurfaceWrapper.h
index d9c4eec..8341fd5 100644
--- a/media/codec2/sfplugin/InputSurfaceWrapper.h
+++ b/media/codec2/sfplugin/InputSurfaceWrapper.h
@@ -98,6 +98,13 @@
         mDataSpace = dataSpace;
     }
 
+    /**
+     * Clean up C2Work related references if necessary. No-op by default.
+     *
+     * \param index index of input work.
+     */
+    virtual void onInputBufferDone(c2_cntr64_t /* index */) {}
+
 protected:
     android_dataspace mDataSpace;
 };