Add encryption infrastructure by edwinyyyu · Pull Request #1365 · MemMachine/MemMachine

edwinyyyu · 2026-04-22T22:40:07Z

Purpose of the change

For supporting sensitive data encryption.
The approach to encrypting in SQL will be to store codec config (including wrapped DEK) in database.

Description

Based on #1375 for easier rebasing, but really only creates stuff for encryption.
DOES NOT HANDLE key creation/rotation.

Alternative is encryption at rest but database sees the plaintext data.

KMSCryptoClient may get implementations with the following backends:

OpenBao
HashiCorp Vault
cloud KMS such as AWS, Azure, GCP

Type of change

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Refactor (does not change functionality, e.g., code style improvements, linting)
Documentation update
Project Maintenance (updates to build scripts, CI, etc., that do not affect the main project)
Security (improves security without changing functionality)

How Has This Been Tested?

Checklist

Maintainer Checklist

Confirmed all checks passed
Contributor has signed the commit(s)
Reviewed the code
Run, Tested, and Verified the change(s) work as expected

edwinyyyu · 2026-04-23T02:10:19Z

May need changes to support BYOK.

edwinyyyu · 2026-05-01T22:40:55Z

Needs a design doc.

Copilot

Pull request overview

Adds foundational encryption plumbing to the server so partitions can persist a payload codec configuration (including a wrapped DEK) and later materialize an encrypt/decrypt codec via a KMS client abstraction.

Changes:

Introduces a KMSCryptoClient abstraction plus a KMS-envelope codec loader to unwrap DEKs and construct codecs.
Adds an AES-GCM payload codec and extends payload codec config models to support AES-GCM + base64url-encoded binary fields.
Adds cryptography as a dependency and includes unit tests for AES-GCM, config round-tripping, and loader behavior.

Reviewed changes

Copilot reviewed 10 out of 11 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
uv.lock	Locks `cryptography` dependency needed for AES-GCM implementation.
packages/server/pyproject.toml	Adds `cryptography` to server runtime dependencies.
packages/server/src/memmachine_server/common/payload_codec/payload_codec_config.py	Adds KMS-envelope/AES-GCM codec config models + base64url (de)serialization.
packages/server/src/memmachine_server/common/payload_codec/kms_envelope_payload_codec_loader.py	Adds loader that unwraps DEKs via KMS and materializes codecs.
packages/server/src/memmachine_server/common/payload_codec/aes_gcm_payload_codec.py	Implements AES-GCM encode/decode codec using `cryptography`.
packages/server/src/memmachine_server/common/payload_codec/init.py	Exports new loader/config types.
packages/server/src/memmachine_server/common/kms/kms_crypto_client.py	Adds abstract KMS crypto client interface (encrypt/decrypt).
packages/server/src/memmachine_server/common/kms/init.py	Exposes KMS abstractions via package exports.
packages/server/server_tests/memmachine_server/common/payload_codec/test_payload_codec.py	Adds AES-GCM round-trip + wrong-AAD failure test.
packages/server/server_tests/memmachine_server/common/payload_codec/test_payload_codec_config.py	Adds AES-GCM config JSON round-trip and base64url assertions.
packages/server/server_tests/memmachine_server/common/payload_codec/test_kms_envelope_payload_codec_loader.py	Adds loader tests (success + unsupported config).

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

sscargal

Summary

Nice, well-scoped infrastructure addition — clean ABCs, sensible AAD propagation, and the loader's match pattern leaves room for future codec types. CI is green across all platforms and the happy-path tests are correct.

Requesting changes for a few items not already covered by @Copilot's notes. Copilot has already flagged the nonce_size validation gap (twice) and the permissive urlsafe_b64decode behavior; those still apply and aren't repeated here.

Blocking

GCM auth tag length check missing in decode — see inline on aes_gcm_payload_codec.py:39. Currently conflates truncated buffers with active tampering.
NotImplementedError is the wrong exception type for an unsupported runtime config — see inline on kms_envelope_payload_codec_loader.py:45.
binascii.Error from urlsafe_b64decode is not wrapped with from err — see inlines on payload_codec_config.py:36 and :51. AGENTS.md L94 requires raise ... from err chaining.

Suggestions

KMSEnvelopePayloadCodecConfig is meant as an abstract base but is instantiable and has no type discriminator; the test at test_kms_envelope_payload_codec_loader.py:80 shows this leaks. See inline on payload_codec_config.py:23.
Test gaps for defensive branches the author explicitly added (short-payload guard, tampered-ciphertext distinct from AD mismatch, associated_data=None path). See inlines on the test files.

Tests

Existing tests cover round-trip + AD-mismatch + unsupported-config rejection — the right baseline shape. The PR checklist has I have added unit tests unticked despite tests existing — likely just an oversight.

Local checks

Skipped local ruff / ty re-run. CI ran both across Python 3.12/3.13/3.14 on Linux/macOS/Windows and all 40+ checks passed.

sscargal · 2026-05-22T16:43:36Z

+        if len(value) < self._nonce_size:
+            raise ValueError("Encrypted payload is too short")


The length check guards against a missing nonce, but AES-GCM also requires a 16-byte auth tag. A 12–27 byte buffer passes this check and then surfaces as an opaque InvalidTag from AESGCM.decrypt, which conflates truncated ciphertext with active tampering in logs/Sentry.

Suggest:

if len(value) < self._nonce_size + 16: raise ValueError("Encrypted payload is too short (missing nonce or auth tag)")

The GCM tag is fixed at 16 bytes for AES-128/192/256.

sscargal · 2026-05-22T16:43:36Z

+                raise NotImplementedError(
+                    f"Unsupported KMS envelope payload codec config: "
+                    f"{type(config).__name__}"


NotImplementedError is conventionally for unimplemented abstract methods / unfinished features. An unsupported runtime config is a caller/data error — TypeError (matches Python's own behavior for unmatched match over types) or ValueError is more appropriate, so callers can distinguish the loader is incomplete from you passed a bogus config.

The test at test_kms_envelope_payload_codec_loader.py:88-89 would need a matching update.

It represents an unfinished feature, assuming that the type passed in matches the type hint. In Python, runtime type checks are excessive. If anyone cares about typing, they can use a static type checker.

sscargal · 2026-05-22T16:43:36Z

+        if isinstance(value, bytes):
+            return value
+        if isinstance(value, str):
+            return urlsafe_b64decode(value.encode("ascii"))


Complements the Copilot note above: when urlsafe_b64decode rejects malformed input it raises binascii.Error, which is currently re-raised unwrapped. AGENTS.md L94 requires raise ... from err chaining so the original cause is preserved. Suggest:

import binascii if isinstance(value, str): try: return urlsafe_b64decode(value.encode("ascii")) except (binascii.Error, ValueError) as e: raise ValueError("wrapped_dek is not valid base64url") from e

Same pattern needed for associated_data at line 51.

sscargal · 2026-05-22T16:43:36Z



-PayloadCodecConfigUnion = PlaintextPayloadCodecConfig
+class KMSEnvelopePayloadCodecConfig(BaseModel):


KMSEnvelopePayloadCodecConfig is intended as an abstract base (no type discriminator) but is freely instantiable — test_kms_envelope_payload_codec_loader.py:80-96 confirms this by constructing an UnknownKMSEnvelopePayloadCodecConfig directly.

A couple of options:

If genuinely abstract: make it ABC (or a typing.Protocol) and move the base64 validators into a shared mixin. Bare instantiation should fail.

If purely structural: consider composition — give each concrete config an envelope: KMSEnvelopeParams field instead of inheritance. That flattens the discriminated union and removes the abstract base in a union hazard.

Either way, consider model_config = ConfigDict(frozen=True) — config objects holding (wrapped) key material should be immutable.

Composition is probably better than inheritance here from a design pattern perspective, but it will make config processing more complicated.

sscargal · 2026-05-22T16:43:36Z

+        b"0" * 32,
+        associated_data=b"partition:block",
+    )
+    with pytest.raises(InvalidTag):


Three defensive branches the codec adds are currently untested — easy wins:

nonce_size <= 0 constructor guard (codec line 24-25)

short-payload ValueError branch (codec line 39-40)

tampered ciphertext distinct from AD mismatch — the existing wrong_codec test only varies the AAD, so a regression that silently disabled GCM auth-tag verification while leaving AAD intact would still pass. Flipping a byte in the ciphertext (buf = bytearray(encoded); buf[20] ^= 1) and asserting InvalidTag would catch it.

sscargal · 2026-05-22T16:43:36Z

+        key_ref="partition_key",
+        wrapped_dek=b"wrapped-dek-bytes",
+        nonce_size=12,
+        associated_data=b"partition:context",


associated_data=None is never exercised in the codec, config, or loader tests. Both _decode_associated_data / _serialize_associated_data None branches and AESGCM.encrypt(...) with aad=None are unverified. A second loader test with associated_data=None would close this — it's the path most production callers will hit first if they forget to pass AAD.

Signed-off-by: Edwin Yu <edwinyyyu@gmail.com>

edwinyyyu added this to the v0.4.0 milestone Apr 22, 2026

edwinyyyu marked this pull request as draft April 23, 2026 01:52

edwinyyyu force-pushed the key_manager branch from f245dd0 to cad359a Compare April 23, 2026 03:04

edwinyyyu changed the title ~~Add KeyManager ABC for KEK management~~ Add KMSCryptoClient ABC for KEK management Apr 23, 2026

edwinyyyu force-pushed the key_manager branch from cad359a to 1ca5d76 Compare April 23, 2026 03:06

edwinyyyu changed the title ~~Add KMSCryptoClient ABC for KEK management~~ Add KMSCryptoClient ABC for using KEK Apr 23, 2026

edwinyyyu changed the title ~~Add KMSCryptoClient ABC for using KEK~~ Add encryption infrastructure Apr 23, 2026

edwinyyyu mentioned this pull request Apr 23, 2026

New episodic memory (currently named EventMemory) based on VectorStore and (new) SegmentStore #1199

Closed

26 tasks

edwinyyyu marked this pull request as ready for review April 23, 2026 06:33

edwinyyyu requested review from jealous, malatewang and sscargal April 23, 2026 19:52

jealous approved these changes Apr 23, 2026

View reviewed changes

edwinyyyu force-pushed the key_manager branch from 84cc120 to 79351ea Compare April 23, 2026 20:32

edwinyyyu marked this pull request as draft April 23, 2026 21:56

edwinyyyu marked this pull request as ready for review April 24, 2026 03:13

edwinyyyu modified the milestones: v0.4.0, v0.3.7 Apr 24, 2026

edwinyyyu mentioned this pull request Apr 24, 2026

KMSCryptoClient backed by OS-native keyring #1370

Open

26 tasks

edwinyyyu force-pushed the key_manager branch from f2e22e6 to 9e2ab8d Compare April 27, 2026 20:46

edwinyyyu mentioned this pull request Apr 27, 2026

KMSCryptoClient backed by HashiCorp Vault/OpenBao Transit #1371

Open

26 tasks

edwinyyyu marked this pull request as draft April 28, 2026 00:03

edwinyyyu mentioned this pull request Apr 28, 2026

Payload codecs #1375

Closed

edwinyyyu force-pushed the key_manager branch from 9e2ab8d to 618fb90 Compare April 28, 2026 00:14

edwinyyyu marked this pull request as ready for review April 28, 2026 00:14

edwinyyyu force-pushed the key_manager branch 2 times, most recently from 0b8a08f to 07e0e00 Compare May 8, 2026 17:58

sscargal requested a review from Copilot May 20, 2026 20:57

Copilot started reviewing on behalf of sscargal May 20, 2026 20:57 View session

Copilot AI reviewed May 20, 2026

View reviewed changes

sscargal requested changes May 22, 2026

View reviewed changes

edwinyyyu added 2 commits May 22, 2026 15:45

Implement encryption infrastructure

dfe8ec1

Signed-off-by: Edwin Yu <edwinyyyu@gmail.com>

Validate nonce and b64 data

db4cd17

Signed-off-by: Edwin Yu <edwinyyyu@gmail.com>

edwinyyyu force-pushed the key_manager branch from 2e2cdfb to db4cd17 Compare May 22, 2026 22:46

edwinyyyu and others added 4 commits May 22, 2026 16:14

ciphertext length should include AES-GCM auth tag (16 bytes)

eab6a56

Signed-off-by: Edwin Yu <edwinyyyu@gmail.com>

Add tests

71bd560

Signed-off-by: Edwin Yu <edwinyyyu@gmail.com>

Composition instead of inheritance

1d11593

Signed-off-by: Edwin Yu <edwinyyyu@gmail.com>

Merge branch 'main' into key_manager

671a091

		if len(value) < self._nonce_size:
		raise ValueError("Encrypted payload is too short")



		PayloadCodecConfigUnion = PlaintextPayloadCodecConfig
		class KMSEnvelopePayloadCodecConfig(BaseModel):

Conversation

edwinyyyu commented Apr 22, 2026 • edited by sscargal Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose of the change

Description

Type of change

How Has This Been Tested?

Checklist

Maintainer Checklist

Uh oh!

edwinyyyu commented Apr 23, 2026

Uh oh!

edwinyyyu commented May 1, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sscargal left a comment

Choose a reason for hiding this comment

Summary

Blocking

Suggestions

Tests

Local checks

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

edwinyyyu May 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

edwinyyyu commented Apr 22, 2026 •

edited by sscargal

Loading

edwinyyyu May 26, 2026 •

edited

Loading