feat: add codec for log

Author: brian14708

.github/workflows/ci.yaml

@@ -27,7 +27,7 @@ jobs:
         with:
           enable-cache: true
       - name: Sync dependencies
-        run: uv sync --group type-checking --group lint
+        run: uv sync --group type-checking --group lint --extra full
       - name: Run tests
         run: uv run pytest
       - name: Run linter

noxfile.py

@@ -1,3 +1,6 @@
+# type: ignore
+# pyright: basic, reportMissingImports=false
+
 from __future__ import annotations
 
 import nox
@@ -9,7 +12,7 @@
 
 def install_deps(s: nox.Session, groups: list[str]):
     s.env["UV_PROJECT_ENVIRONMENT"] = s.virtualenv.location
-    cmd = ["uv", "sync", "--frozen"]
+    cmd = ["uv", "sync", "--frozen", "--all-extras"]
     for g in groups:
         cmd.extend(("--group", g))
     _ = s.run_install(*cmd)

pyproject.toml

@@ -8,6 +8,11 @@ dependencies = [
     "typing-extensions>=4.15.0",
 ]
 
+[project.optional-dependencies]
+full = [
+    "pydantic>=2.11.9",
+]
+
 [build-system]
 requires = ["uv_build>=0.8.19,<0.9.0"]
 build-backend = "uv_build"
@@ -20,7 +25,6 @@ dev = [
 type-checking = [
     "basedmypy>=2.10.0",
     "basedpyright>=1.31.2",
-    "nox>=2025.5.1",
 ]
 lint = [
     "ruff>=0.12.8",

src/duron/log/codec.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import base64
+import importlib
+import pickle
+from typing import TYPE_CHECKING, cast, final
+
+from typing_extensions import Protocol
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel as BaseModelType
+
+    from duron.log.entry import JSONValue
+
+    BaseModel: type[BaseModelType] | None
+
+try:
+    from pydantic import BaseModel
+except ImportError:
+    BaseModel = None
+
+
+class BaseCodec(Protocol):
+    def encode_json(self, result: object) -> JSONValue: ...
+    def decode_json(self, encoded: JSONValue) -> object: ...
+
+    def encode_state(self, obj: object) -> str: ...
+    def decode_state(self, state: str) -> object: ...
+
+
+@final
+class DefaultCodec:
+    def __init__(self) -> None:
+        self._type_cache: dict[str, type] = {}
+
+    def _lookup_model(self, qual: str) -> type:
+        cls = self._type_cache.get(qual)
+        if cls is not None:
+            return cls
+
+        mod, name = qual.split(":", 1)
+        obj: object = importlib.import_module(mod)
+        for part in name.split("."):
+            obj = getattr(obj, part)
+        if isinstance(obj, type):
+            self._type_cache[qual] = obj
+            return obj
+        else:
+            raise TypeError(f"Imported object is not a type: {obj!r}")
+
+    def encode_json(self, result: object) -> JSONValue:
+        if BaseModel and isinstance(result, BaseModel):
+            cls = result.__class__
+            model = result.model_dump(mode="json")
+            model["_duron.pydantic"] = f"{cls.__module__}:{cls.__qualname__}"
+            return model
+        return cast("JSONValue", result)
+
+    def decode_json(self, encoded: JSONValue) -> object:
+        if isinstance(encoded, dict) and "_duron.pydantic" in encoded:
+            model = self._lookup_model(cast("str", encoded["_duron.pydantic"]))
+            if not BaseModel or not issubclass(model, BaseModel):
+                raise TypeError(f"Decoded class is not a BaseModel subclass: {model}")
+            return model.model_validate(
+                {k: v for k, v in encoded.items() if not k.startswith("_duron.")}
+            )
+        return encoded
+
+    def encode_state(self, obj: object) -> str:
+        return base64.b64encode(pickle.dumps(obj, protocol=5)).decode("ascii")
+
+    def decode_state(self, state: str) -> object:
+        return pickle.loads(base64.b64decode(state.encode()))

src/duron/log/entry.py

@@ -2,6 +2,8 @@
 
 from typing_extensions import NotRequired, TypedDict
 
+JSONValue = dict[str, "JSONValue"] | list["JSONValue"] | str | int | float | bool | None
+
 
 class BaseEntry(TypedDict):
     id: str
@@ -13,7 +15,7 @@ class BaseEntry(TypedDict):
 class ErrorInfo(TypedDict):
     code: int
     message: str
-    data: NotRequired[object]
+    state: NotRequired[str]  # opaque
 
 
 class PromiseCreateEntry(BaseEntry):
@@ -23,7 +25,7 @@ class PromiseCreateEntry(BaseEntry):
 class PromiseCompleteEntry(BaseEntry):
     type: Literal["promise/complete"]
     promise_id: str
-    result: NotRequired[object]
+    result: NotRequired[JSONValue]
     error: NotRequired[ErrorInfo]
 
 
@@ -34,8 +36,8 @@ class StreamCreateEntry(BaseEntry):
 class StreamEmitEntry(BaseEntry):
     type: Literal["stream/emit"]
     stream_id: str
-    value: NotRequired[object]
-    state: NotRequired[object]
+    value: NotRequired[JSONValue]
+    state: NotRequired[str]  # opaque
 
 
 class StreamCloseEntry(BaseEntry):

src/duron/task_runner.py

@@ -14,13 +14,15 @@
 )
 
 from duron.event_loop import create_loop
+from duron.log.codec import DefaultCodec
 from duron.log.entry import is_entry
 from duron.ops import FnCall
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Callable
 
     from duron.event_loop import WaitSet
+    from duron.log.codec import BaseCodec
     from duron.log.entry import Entry, ErrorInfo, UnknownEntry
     from duron.log.storage import BaseLogStorage, Lease, Offset
     from duron.mark import DurableFn
@@ -30,27 +32,32 @@
 
 
 class TaskRunner:
-    def __init__(self):
-        pass
+    def __init__(self, codec: BaseCodec | None = None) -> None:
+        self._codec: BaseCodec = codec or DefaultCodec()
 
     async def run(
         self,
         task_id: bytes,
         task_co: DurableFn[[], Coroutine[Any, Any, _T]],
         log: BaseLogStorage,
     ) -> _T:
-        return await _Task[_T](task_id, task_co(), log).run()
+        return await _Task[_T](task_id, task_co(), log, self._codec).run()
 
 
 @final
 class _Task(Generic[_T]):
     def __init__(
-        self, id: bytes, task_co: Coroutine[Any, Any, _T], log: BaseLogStorage
+        self,
+        id: bytes,
+        task_co: Coroutine[Any, Any, _T],
+        log: BaseLogStorage,
+        codec: BaseCodec,
     ) -> None:
         self._id = id
         self._loop = create_loop(id)
         self._task = self._loop.create_task(task_co)
         self._log = log
+        self._codec = codec
         self._running: Lease | None = None
         self._pending_msg: list[Entry] = []
         self._pending_task: dict[str, Callable[[], Coroutine[Any, Any, object]]] = {}
@@ -110,7 +117,7 @@ async def run(self) -> _T:
                     "id": _encode_id(self._id, True),
                     "ts": _encode_timestamp(self.now()),
                     "promise_id": _encode_id(self._id, False),
-                    "result": res,
+                    "result": self._codec.encode_json(res),
                 },
             )
             return res
@@ -121,7 +128,7 @@ async def run(self) -> _T:
                     "id": _encode_id(self._id, True),
                     "ts": _encode_timestamp(self.now()),
                     "promise_id": _encode_id(self._id, False),
-                    "error": exception_to_error_info(e),
+                    "error": _encode_error(e, self._codec),
                 },
             )
             raise
@@ -157,13 +164,13 @@ async def handle_message(self, e: Entry) -> None:
             if "error" in e:
                 self._loop.post_completion_threadsafe(
                     id,
-                    exception=error_info_to_exception(e["error"]),
+                    exception=_decode_error(e["error"], self._codec),
                 )
                 self._pending_ops.discard(id)
             elif "result" in e:
                 self._loop.post_completion_threadsafe(
                     id,
-                    result=e["result"],
+                    result=self._codec.decode_json(e["result"]),
                 )
                 self._pending_ops.discard(id)
             else:
@@ -200,17 +207,17 @@ async def cb() -> None:
                                 "id": _encode_id(id, True),
                                 "type": "promise/complete",
                                 "promise_id": _encode_id(id, False),
-                                "result": result,
+                                "result": self._codec.encode_json(result),
                             }
                         )
                     except BaseException as e:
                         await self.enqueue_log(
                             {
+                                "type": "promise/complete",
                                 "ts": _encode_timestamp(self.now()),
                                 "id": _encode_id(id, True),
-                                "type": "promise/complete",
                                 "promise_id": _encode_id(id, False),
-                                "error": exception_to_error_info(e),
+                                "error": _encode_error(e, self._codec),
                             }
                         )
 
@@ -223,20 +230,20 @@ async def cb() -> None:
                 raise NotImplementedError(f"Unsupported op: {op!r}")
 
 
-def _encode_id(id: bytes, end: bool) -> str:
-    if end:
+def _encode_id(id: bytes, is_end: bool) -> str:
+    if is_end:
         return base64.b64encode(id).decode() + "-"
     else:
         return base64.b64encode(id).decode() + "+"
 
 
-def _decode_id(s: str) -> tuple[bytes, bool]:
-    if s.endswith("-"):
-        return base64.b64decode(s[:-1]), True
-    elif s.endswith("+"):
-        return base64.b64decode(s[:-1]), False
+def _decode_id(encoded: str) -> tuple[bytes, bool]:
+    if encoded.endswith("-"):
+        return base64.b64decode(encoded[:-1]), True
+    elif encoded.endswith("+"):
+        return base64.b64decode(encoded[:-1]), False
     else:
-        raise ValueError(f"Invalid encoded id: {s!r}")
+        raise ValueError(f"Invalid encoded id: {encoded!r}")
 
 
 def _encode_timestamp(ts_ns: int) -> int:
@@ -247,12 +254,20 @@ def _decode_timestamp(ts: int) -> int:
     return ts * 1_000
 
 
-def exception_to_error_info(e: BaseException) -> ErrorInfo:
+def _encode_error(error: BaseException, codec: BaseCodec) -> ErrorInfo:
+    """Convert exception to ErrorInfo dict."""
     return {
-        "code": 1,
-        "message": str(e),
+        "code": -1,
+        "message": str(error),
+        "state": codec.encode_state(error),
     }
 
 
-def error_info_to_exception(e: ErrorInfo) -> Exception:
-    return Exception(f"[{e['code']}] {e['message']}")
+def _decode_error(error_info: ErrorInfo, codec: BaseCodec) -> BaseException:
+    """Convert ErrorInfo dict to exception."""
+    try:
+        if "state" not in error_info:
+            return Exception(f"[{error_info['code']}] {error_info['message']}")
+        return cast("BaseException", codec.decode_state(error_info["state"]))
+    except Exception:
+        return Exception(f"[{error_info['code']}] {error_info['message']}")

tests/test_task.py

@@ -1,5 +1,6 @@
 import asyncio
 import uuid
+from typing import ClassVar
 
 import pytest
 
@@ -10,7 +11,7 @@
 
 
 @pytest.mark.asyncio
-async def test_task_runner():
+async def test_task():
     @durable()
     async def activity() -> str:
         ctx = duron.context.Context()
@@ -50,3 +51,31 @@ async def error():
         _ = await tr.run(b"1", activity, log)
     with pytest.raises(check=lambda v: "test error" in str(v)):
         _ = await tr.run(b"1", activity, log)
+
+
+try:
+    import pydantic
+
+    class Point(pydantic.BaseModel):
+        model_config: ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(
+            extra="forbid"
+        )
+
+        x: int
+        y: int
+
+    @pytest.mark.asyncio
+    async def test_pydantic():
+        @durable()
+        async def activity() -> Point:
+            ctx = duron.context.Context()
+            pt = await ctx.run(lambda: Point(x=1, y=2))
+            return Point(x=pt.x + 5, y=pt.y + 10)
+
+        tr = duron.task_runner.TaskRunner()
+        log = MemoryLogStorage()
+        a = await tr.run(b"1", activity, log)
+        assert type(a) is Point
+        assert a.x == 6 and a.y == 12
+except ImportError:
+    pass