fix: handle unpickleable values in flow state
Some checks failed
CodeQL Advanced / Analyze (actions) (push) Has been cancelled
CodeQL Advanced / Analyze (python) (push) Has been cancelled
Notify Downstream / notify-downstream (push) Has been cancelled
Mark stale issues and pull requests / stale (push) Has been cancelled

This commit is contained in:
Greyson LaLonde
2025-11-06 01:29:21 -05:00
committed by GitHub
parent 7e6171d5bc
commit e4cc9a664c
2 changed files with 135 additions and 11 deletions

View File

@@ -3,6 +3,7 @@
import asyncio
import threading
from datetime import datetime
from typing import Optional
import pytest
from pydantic import BaseModel
@@ -1384,3 +1385,110 @@ def test_mixed_sync_async_execution_order():
]
assert execution_order == expected_order
def test_flow_copy_state_with_unpickleable_objects():
"""Test that _copy_state handles unpickleable objects like RLock.
Regression test for issue #3828: Flow should not crash when state contains
objects that cannot be deep copied (like threading.RLock).
"""
class StateWithRLock(BaseModel):
counter: int = 0
lock: Optional[threading.RLock] = None
class FlowWithRLock(Flow[StateWithRLock]):
@start()
def step_1(self):
self.state.counter += 1
@listen(step_1)
def step_2(self):
self.state.counter += 1
flow = FlowWithRLock(initial_state=StateWithRLock())
flow._state.lock = threading.RLock()
copied_state = flow._copy_state()
assert copied_state.counter == 0
assert copied_state.lock is not None
def test_flow_copy_state_with_nested_unpickleable_objects():
"""Test that _copy_state handles unpickleable objects nested in containers.
Regression test for issue #3828: Verifies that unpickleable objects
nested inside dicts/lists in state don't cause crashes.
"""
class NestedState(BaseModel):
data: dict = {}
items: list = []
class FlowWithNestedUnpickleable(Flow[NestedState]):
@start()
def step_1(self):
self.state.data["lock"] = threading.RLock()
self.state.data["value"] = 42
@listen(step_1)
def step_2(self):
self.state.items.append(threading.Lock())
self.state.items.append("normal_value")
flow = FlowWithNestedUnpickleable(initial_state=NestedState())
flow.kickoff()
assert flow.state.data["value"] == 42
assert len(flow.state.items) == 2
def test_flow_copy_state_without_unpickleable_objects():
"""Test that _copy_state still works normally with pickleable objects.
Ensures that the fallback logic doesn't break normal deep copy behavior.
"""
class NormalState(BaseModel):
counter: int = 0
data: str = ""
nested: dict = {}
class NormalFlow(Flow[NormalState]):
@start()
def step_1(self):
self.state.counter = 5
self.state.data = "test"
self.state.nested = {"key": "value"}
flow = NormalFlow(initial_state=NormalState())
flow.state.counter = 10
flow.state.data = "modified"
flow.state.nested["key"] = "modified"
copied_state = flow._copy_state()
assert copied_state.counter == 10
assert copied_state.data == "modified"
assert copied_state.nested["key"] == "modified"
flow.state.nested["key"] = "changed_after_copy"
assert copied_state.nested["key"] == "modified"
def test_flow_copy_state_with_dict_state():
"""Test that _copy_state works with dict-based states."""
class DictFlow(Flow[dict]):
@start()
def step_1(self):
self.state["counter"] = 1
flow = DictFlow()
flow.state["test"] = "value"
copied_state = flow._copy_state()
assert copied_state["test"] == "value"
flow.state["test"] = "modified"
assert copied_state["test"] == "value"