From 12c7ddf3a319aeeb8fa982f5d703f3f2fac48c93 Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Tue, 18 Mar 2025 17:30:21 -0700
Subject: [PATCH] Update kafka

---
 .../langgraph/scheduler/kafka/executor.py         | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py
index 2d442e131..6564b28a0 100644
--- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py
+++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py
@@ -1,6 +1,7 @@
 import asyncio
 import binascii
 import concurrent.futures
+import weakref
 from collections.abc import Sequence
 from contextlib import (
     AbstractAsyncContextManager,
@@ -216,10 +217,11 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
                 checkpoint_null_version=checkpoint_null_version(saved.checkpoint),
             ):
                 # execute task, saving writes
+                put_writes = partial(self._put_writes, submit, msg["config"])
                 runner = PregelRunner(
-                    submit=submit,
-                    put_writes=partial(self._put_writes, submit, msg["config"]),
-                    schedule_task=self._schedule_task,
+                    submit=weakref.ref(submit),
+                    put_writes=weakref.ref(put_writes),
+                    schedule_task=weakref.WeakMethod(self._schedule_task),
                 )
                 async for _ in runner.atick([task], reraise=False):
                     pass
@@ -432,10 +434,11 @@ class KafkaExecutor(AbstractContextManager):
                 checkpoint_null_version=checkpoint_null_version(saved.checkpoint),
             ):
                 # execute task, saving writes
+                put_writes = partial(self._put_writes, submit, msg["config"])
                 runner = PregelRunner(
-                    submit=submit,
-                    put_writes=partial(self._put_writes, submit, msg["config"]),
-                    schedule_task=self._schedule_task,
+                    submit=weakref.ref(submit),
+                    put_writes=weakref.ref(put_writes),
+                    schedule_task=weakref.WeakMethod(self._schedule_task),
                 )
                 for _ in runner.tick([task], reraise=False):
                     pass