fix: Fix for unhashable workflow subclasses (#177)

Author: adrianlyjak

AGENTS.md

@@ -0,0 +1,76 @@
+# LlamaIndex Workflows - Claude Development Guide
+
+## Project Overview
+This is the LlamaIndex Workflows library - an event-driven, async-first framework for orchestrating complex AI applications and multi-step processes.
+
+## Key Technologies
+- Python 3.9+
+- AsyncIO (async/await)
+- Pydantic for data models
+- Starlette for web server
+- Uvicorn for ASGI serving
+
+## Development Commands
+
+### Testing
+```bash
+# Run all tests
+uv run pytest
+
+# Run tests with coverage
+uv run pytest --cov=src/workflows --cov-report=html
+
+# Run specific test files
+uv run pytest tests/test_server.py tests/test_server_utils.py
+
+# Run tests in verbose mode
+uv run pytest -v
+```
+
+### Linting & Formatting
+```bash
+# Run pre-commit hooks
+uv run pre-commit run -a
+```
+
+## Project Structure
+- `src/workflows/` - Main library code
+- `src/workflows/server/` - Web server implementation
+- `tests/` - Test suite
+- `examples/` - Usage examples
+
+## Key Components
+- **Workflow** - Main orchestration class
+- **Context** - State management across workflow steps
+- **Events** - Event-driven communication between steps
+- **WorkflowServer** - HTTP server for serving workflows as web services
+
+## Notes for Claude
+- Always run tests after making changes: `uv run pytest`
+- Never use classes for tests, only use pytest functions
+- Always annotate with types function arguments and return values
+- The project uses async/await extensively
+- Context serialization requires specific JSON format for globals
+
+## Autonomous Operation
+
+The following rules apply if you are running in an isolated sandbox environment and have tools to commit and push changes to git
+
+Make sure to install uv as the package manager. Development commands rely on it.
+
+```bash
+curl -fsSL https://astral.sh/uv/install.sh | sh
+```
+
+Always run test tests and pre-commit commands before committing. They run very fast and are not verbose.
+
+Tests:
+
+```bash
+uv run pytest -nauto --timeout=1
+```
+
+Linting, typechecking, and formatting:
+```bash
+uv run pre-commit run -a
+```

CLAUDE.md

@@ -0,0 +1 @@
+AGENTS.md

CLAUDE.md

@@ -18,7 +18,7 @@ dev = [
 
 [project]
 name = "llama-index-workflows"
-version = "2.9.0"
+version = "2.9.1"
 description = "An event-driven, async-first, step-based way to control the execution flow of AI applications like Agents."
 readme = "README.md"
 license = "MIT"

pyproject.toml

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from typing import Callable, Generic, TypeVar, overload
+import weakref
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class _IdentityWeakRef(weakref.ref, Generic[K]):
+    __slots__ = ("_hash",)
+
+    _hash: int
+
+    def __new__(
+        cls, obj: K, callback: Callable[[_IdentityWeakRef[K]], None] | None = None
+    ) -> _IdentityWeakRef[K]:
+        self = super().__new__(cls, obj, callback)
+        self._hash = id(
+            obj
+        )  # cache identity-based hash; works even if obj is unhashable
+        return self
+
+    def __hash__(self) -> int:
+        return self._hash
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, _IdentityWeakRef):
+            return NotImplemented
+        return self() is other()
+
+
+class IdentityWeakKeyDict(Generic[K, V]):
+    _d: dict[_IdentityWeakRef[K], V]
+
+    def __init__(self) -> None:
+        self._d = {}
+
+    def _mk(self, obj: K) -> _IdentityWeakRef[K]:
+        def _cb(wr: _IdentityWeakRef[K]) -> None:
+            self._d.pop(wr)
+
+        return _IdentityWeakRef(obj, _cb)
+
+    def __setitem__(self, obj: K, value: V) -> None:
+        self._d[self._mk(obj)] = value
+
+    def __getitem__(self, obj: K) -> V:
+        return self._d[_IdentityWeakRef(obj)]
+
+    @overload
+    def get(self, obj: K) -> V | None: ...
+
+    @overload
+    def get(self, obj: K, default: V) -> V: ...
+
+    def get(self, obj: K, default: V | None = None) -> V | None:
+        return self._d.get(_IdentityWeakRef(obj), default)
+
+    def __contains__(self, obj: K) -> bool:
+        return _IdentityWeakRef(obj) in self._d
+
+    def __delitem__(self, obj: K) -> None:
+        del self._d[_IdentityWeakRef(obj)]

src/workflows/runtime/types/_identity_weak_ref.py

@@ -1,7 +1,7 @@
 from threading import Lock
-from weakref import WeakKeyDictionary
 from dataclasses import dataclass
 from typing import Optional
+from workflows.runtime.types._identity_weak_ref import IdentityWeakKeyDict
 from workflows.runtime.types.plugin import (
     ControlLoopFunction,
     Plugin,
@@ -24,9 +24,9 @@ class WorkflowPluginRegistry:
     def __init__(self) -> None:
         # Map each workflow instance to its plugin registrations.
         # Weakly references workflow keys so entries are GC'd when workflows are.
-        self.workflows: WeakKeyDictionary[
+        self.workflows: IdentityWeakKeyDict[
             Workflow, dict[type[Plugin], RegisteredWorkflow]
-        ] = WeakKeyDictionary()
+        ] = IdentityWeakKeyDict()
         self.lock = Lock()
         self.run_contexts: dict[str, RegisteredRunContext] = {}
 

src/workflows/runtime/workflow_registry.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import gc
+import weakref
+
+import pytest
+
+from workflows.runtime.types._identity_weak_ref import IdentityWeakKeyDict
+
+
+class Unhashable:
+    def __hash__(self) -> int:
+        raise TypeError("Unhashable")
+
+
+def test_identity_weak_key_dict_removes_entry_when_object_unreferenced() -> None:
+    with pytest.raises(TypeError):
+        hash(Unhashable())
+
+    d: IdentityWeakKeyDict[Unhashable, str] = IdentityWeakKeyDict()
+
+    obj = Unhashable()
+    d[obj] = "value"
+
+    # While the object is strongly referenced, it should be present
+    assert obj in d
+    assert d.get(obj) == "value"
+
+    # Keep a weak reference to verify collection happened
+    w = weakref.ref(obj)
+
+    # Drop strong reference and force collection
+    del obj
+    gc.collect()
+
+    # Object should be collected and the dict entry removed via callback
+    assert w() is None
+    assert d._d == {}

tests/runtime/test_identity_weak_ref.py

@@ -795,6 +795,7 @@ def one(self, ev: MyStart) -> MyStop:
 
 @pytest.mark.asyncio
 async def test_workflow_instances_garbage_collected_after_completion() -> None:
+    # test for memory leaks
     class TinyWorkflow(Workflow):
         @step
         async def only(self, ev: StartEvent) -> StopEvent: