From 22f46920910905772bf7c68fa151c5fc9172d396 Mon Sep 17 00:00:00 2001
From: Edwin Yu <edwinyyyu@gmail.com>
Date: Wed, 22 Apr 2026 15:34:11 -0700
Subject: [PATCH] Add payload codec system to plug into storage wrappers

Signed-off-by: Edwin Yu <edwinyyyu@gmail.com>
---
 .../common/payload_codec/__init__.py          |  1 +
 .../payload_codec/test_payload_codec.py       | 16 ++++++++
 .../test_payload_codec_config.py              | 18 ++++++++
 .../common/payload_codec/__init__.py          | 15 +++++++
 .../common/payload_codec/payload_codec.py     | 41 +++++++++++++++++++
 .../payload_codec/payload_codec_config.py     | 36 ++++++++++++++++
 .../payload_codec/plaintext_payload_codec.py  | 17 ++++++++
 7 files changed, 144 insertions(+)
 create mode 100644 packages/server/server_tests/memmachine_server/common/payload_codec/__init__.py
 create mode 100644 packages/server/server_tests/memmachine_server/common/payload_codec/test_payload_codec.py
 create mode 100644 packages/server/server_tests/memmachine_server/common/payload_codec/test_payload_codec_config.py
 create mode 100644 packages/server/src/memmachine_server/common/payload_codec/__init__.py
 create mode 100644 packages/server/src/memmachine_server/common/payload_codec/payload_codec.py
 create mode 100644 packages/server/src/memmachine_server/common/payload_codec/payload_codec_config.py
 create mode 100644 packages/server/src/memmachine_server/common/payload_codec/plaintext_payload_codec.py

diff --git a/packages/server/server_tests/memmachine_server/common/payload_codec/__init__.py b/packages/server/server_tests/memmachine_server/common/payload_codec/__init__.py
new file mode 100644
index 000000000..3e6b9db93
--- /dev/null
+++ b/packages/server/server_tests/memmachine_server/common/payload_codec/__init__.py
@@ -0,0 +1 @@
+"""Tests for payload codecs."""
diff --git a/packages/server/server_tests/memmachine_server/common/payload_codec/test_payload_codec.py b/packages/server/server_tests/memmachine_server/common/payload_codec/test_payload_codec.py
new file mode 100644
index 000000000..6fd4819b6
--- /dev/null
+++ b/packages/server/server_tests/memmachine_server/common/payload_codec/test_payload_codec.py
@@ -0,0 +1,16 @@
+"""Tests for payload codecs."""
+
+from memmachine_server.common.payload_codec.plaintext_payload_codec import (
+    PlaintextPayloadCodec,
+)
+
+
+def test_plaintext_payload_codec_round_trip() -> None:
+    codec = PlaintextPayloadCodec()
+    value = b'{"type":"message","source":"User"}'
+
+    encoded = codec.encode(value)
+    decoded = codec.decode(encoded)
+
+    assert encoded == value
+    assert decoded == value
diff --git a/packages/server/server_tests/memmachine_server/common/payload_codec/test_payload_codec_config.py b/packages/server/server_tests/memmachine_server/common/payload_codec/test_payload_codec_config.py
new file mode 100644
index 000000000..62b3fc451
--- /dev/null
+++ b/packages/server/server_tests/memmachine_server/common/payload_codec/test_payload_codec_config.py
@@ -0,0 +1,18 @@
+"""Tests for payload codec configuration."""
+
+from memmachine_server.common.payload_codec import (
+    decode_payload_codec_config,
+    encode_payload_codec_config,
+)
+from memmachine_server.common.payload_codec.payload_codec_config import (
+    PlaintextPayloadCodecConfig,
+)
+
+
+def test_plaintext_payload_codec_config_round_trip() -> None:
+    config = PlaintextPayloadCodecConfig()
+
+    serialized = encode_payload_codec_config(config)
+    deserialized = decode_payload_codec_config(serialized)
+
+    assert deserialized == config
diff --git a/packages/server/src/memmachine_server/common/payload_codec/__init__.py b/packages/server/src/memmachine_server/common/payload_codec/__init__.py
new file mode 100644
index 000000000..6810e7fa0
--- /dev/null
+++ b/packages/server/src/memmachine_server/common/payload_codec/__init__.py
@@ -0,0 +1,15 @@
+"""Payload codec exports."""
+
+from .payload_codec import PayloadCodec
+from .payload_codec_config import (
+    PayloadCodecConfig,
+    decode_payload_codec_config,
+    encode_payload_codec_config,
+)
+
+__all__ = [
+    "PayloadCodec",
+    "PayloadCodecConfig",
+    "decode_payload_codec_config",
+    "encode_payload_codec_config",
+]
diff --git a/packages/server/src/memmachine_server/common/payload_codec/payload_codec.py b/packages/server/src/memmachine_server/common/payload_codec/payload_codec.py
new file mode 100644
index 000000000..84ef24789
--- /dev/null
+++ b/packages/server/src/memmachine_server/common/payload_codec/payload_codec.py
@@ -0,0 +1,41 @@
+"""Shared payload codec abstractions."""
+
+from abc import ABC, abstractmethod
+
+
+class PayloadCodec(ABC):
+    """Byte-level codec for payloads."""
+
+    @abstractmethod
+    def encode(
+        self,
+        value: bytes,
+    ) -> bytes:
+        """
+        Encode a payload.
+
+        Args:
+            value (bytes):
+                Payload bytes to encode.
+
+        Returns:
+            bytes:
+                Encoded payload bytes.
+        """
+
+    @abstractmethod
+    def decode(
+        self,
+        value: bytes,
+    ) -> bytes:
+        """
+        Decode a serialized payload.
+
+        Args:
+            value (bytes):
+                Encoded payload bytes to decode.
+
+        Returns:
+            bytes:
+                Payload bytes.
+        """
diff --git a/packages/server/src/memmachine_server/common/payload_codec/payload_codec_config.py b/packages/server/src/memmachine_server/common/payload_codec/payload_codec_config.py
new file mode 100644
index 000000000..42523dcb8
--- /dev/null
+++ b/packages/server/src/memmachine_server/common/payload_codec/payload_codec_config.py
@@ -0,0 +1,36 @@
+"""Payload codec configuration models and serialization helpers."""
+
+from collections.abc import Mapping
+from typing import Annotated, Literal
+
+from pydantic import (
+    BaseModel,
+    Field,
+    JsonValue,
+    TypeAdapter,
+)
+
+
+class PlaintextPayloadCodecConfig(BaseModel):
+    """Codec config for plaintext-serialized payloads."""
+
+    type: Literal["plaintext"] = "plaintext"
+
+
+PayloadCodecConfigUnion = PlaintextPayloadCodecConfig
+
+PayloadCodecConfig = Annotated[PayloadCodecConfigUnion, Field(discriminator="type")]
+
+_PAYLOAD_CODEC_CONFIG_ADAPTER = TypeAdapter(PayloadCodecConfig)
+
+
+def encode_payload_codec_config(config: PayloadCodecConfig) -> dict[str, JsonValue]:
+    """Encode a codec config to JSON-compatible data."""
+    return _PAYLOAD_CODEC_CONFIG_ADAPTER.dump_python(config, mode="json")
+
+
+def decode_payload_codec_config(
+    encoded: Mapping[str, JsonValue],
+) -> PayloadCodecConfig:
+    """Decode a codec config from JSON-compatible data."""
+    return _PAYLOAD_CODEC_CONFIG_ADAPTER.validate_python(encoded)
diff --git a/packages/server/src/memmachine_server/common/payload_codec/plaintext_payload_codec.py b/packages/server/src/memmachine_server/common/payload_codec/plaintext_payload_codec.py
new file mode 100644
index 000000000..448fe9555
--- /dev/null
+++ b/packages/server/src/memmachine_server/common/payload_codec/plaintext_payload_codec.py
@@ -0,0 +1,17 @@
+"""Plaintext payload codec."""
+
+from typing import override
+
+from .payload_codec import PayloadCodec
+
+
+class PlaintextPayloadCodec(PayloadCodec):
+    """Codec for plaintext-serialized payloads."""
+
+    @override
+    def encode(self, value: bytes) -> bytes:
+        return value
+
+    @override
+    def decode(self, value: bytes) -> bytes:
+        return value