Moved model copy to pending item add time; add regression test

jiwon-oai · jiwon-oai · commit 72483746d2cd · 2026-01-09T10:31:27.000-08:00
diff --git a/chatkit/server.py b/chatkit/server.py
@@ -697,7 +697,9 @@ async def _process_events(
             with agents_sdk_user_agent_override():
                 async for event in stream():
                     if isinstance(event, ThreadItemAddedEvent):
-                        pending_items[event.item.id] = event.item
+                        # Stash an isolated copy in case we need to persist unfinished items
+                        # on cancellation; downstream handlers keep using the original event.item.
+                        pending_items[event.item.id] = event.item.model_copy(deep=True)
 
                     match event:
                         case ThreadItemDoneEvent():
@@ -805,7 +807,6 @@ def _update_pending_items(
         event: ThreadItemUpdatedEvent,
     ):
         updated_item = pending_items.get(event.item_id)
-        updated_item = updated_item.model_copy(deep=True) if updated_item else None
         update = event.update
         match updated_item:
             case AssistantMessageItem():
diff --git a/tests/test_chatkit_server.py b/tests/test_chatkit_server.py
@@ -33,6 +33,7 @@
     AttachmentsDeleteReq,
     AttachmentUploadDescriptor,
     ClientToolCallItem,
+    CustomTask,
     FeedbackKind,
     FileAttachment,
     ImageAttachment,
@@ -44,6 +45,7 @@
     LockedStatus,
     Page,
     ProgressUpdateEvent,
+    ThoughtTask,
     Thread,
     ThreadAddClientToolOutputParams,
     ThreadAddUserMessageParams,
@@ -79,6 +81,9 @@
     UserMessageTextContent,
     WidgetItem,
     WidgetRootUpdated,
+    Workflow,
+    WorkflowItem,
+    WorkflowTaskAdded,
 )
 from chatkit.widgets import Card, Text
 from tests._types import RequestContext
@@ -354,6 +359,65 @@ def generate_item_id(
             )
 
 
+async def test_workflow_task_not_duplicated_on_done_event():
+    """
+    Regression test to make sure pending item updates do not modify the
+    origin item that was streamed.
+    """
+
+    async def responder(
+        thread: ThreadMetadata, input: UserMessageItem | None, context: Any
+    ) -> AsyncIterator[ThreadStreamEvent]:
+        workflow_item = WorkflowItem(
+            id="workflow-item",
+            created_at=datetime.now(),
+            thread_id=thread.id,
+            workflow=Workflow(type="custom", tasks=[]),
+        )
+
+        yield ThreadItemAddedEvent(item=workflow_item)
+
+        task = CustomTask(title="foo", content="bar")
+        yield ThreadItemUpdatedEvent(
+            item_id=workflow_item.id,
+            update=WorkflowTaskAdded(task=task, task_index=0),
+        )
+
+        workflow_item.workflow.tasks.append(task)
+        yield ThreadItemDoneEvent(item=workflow_item)
+
+    with make_server(responder) as server:
+        events = await server.process_streaming(
+            ThreadsCreateReq(
+                params=ThreadCreateParams(
+                    input=UserMessageInput(
+                        content=[UserMessageTextContent(text="Hello")],
+                        attachments=[],
+                        inference_options=InferenceOptions(),
+                    )
+                )
+            )
+        )
+
+        thread = next(e.thread for e in events if e.type == "thread.created")
+        workflow_done_event = next(
+            e
+            for e in events
+            if isinstance(e, ThreadItemDoneEvent) and isinstance(e.item, WorkflowItem)
+        )
+        workflow_done_item = cast(WorkflowItem, workflow_done_event.item)
+        assert len(workflow_done_item.workflow.tasks) == 1
+        assert workflow_done_item.workflow.tasks[0].title == "foo"
+
+        stored = await server.store.load_item(
+            thread.id, workflow_done_item.id, DEFAULT_CONTEXT
+        )
+        assert isinstance(stored, WorkflowItem)
+        stored_workflow = cast(WorkflowItem, stored)
+        assert len(stored_workflow.workflow.tasks) == 1
+        assert stored_workflow.workflow.tasks[0].title == "foo"
+
+
 async def test_flows_context_to_responder():
     responder_context = None
     add_feedback_context = None
