diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerde.java b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerde.java index 02c9a06b3..56e1ff173 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerde.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerde.java @@ -19,16 +19,24 @@ import com.fasterxml.jackson.annotation.JsonTypeInfo; import com.fasterxml.jackson.core.JsonGenerator; +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.databind.DeserializationContext; +import com.fasterxml.jackson.databind.JsonDeserializer; +import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.JsonSerializer; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.SerializerProvider; +import com.fasterxml.jackson.databind.annotation.JsonDeserialize; +import com.fasterxml.jackson.databind.annotation.JsonSerialize; import com.fasterxml.jackson.databind.module.SimpleModule; import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.InputEvent; import org.apache.flink.agents.api.OutputEvent; +import org.apache.flink.agents.api.context.MemoryUpdate; import org.apache.flink.agents.runtime.operator.ActionTask; import java.io.IOException; +import java.util.Base64; /** * Backend-agnostic serializer/deserializer for {@link ActionState}. @@ -39,6 +47,13 @@ */ public final class ActionStateSerde { + /** + * Reserved envelope key marking a base64-encoded {@code byte[]} {@link MemoryUpdate} value. A + * namespaced key keeps the residual collision with a genuine single-entry user {@code Map} + * negligible and clearly framework-reserved. + */ + private static final String MEMORY_UPDATE_BYTES_KEY = "__flink_agents_bytes__"; + private static final ObjectMapper OBJECT_MAPPER = createObjectMapper(); private ActionStateSerde() {} @@ -74,6 +89,9 @@ private static ObjectMapper createObjectMapper() { module.addSerializer(ActionTask.class, new ActionTaskSerializer()); mapper.registerModule(module); + // Preserve byte[] MemoryUpdate values across the untyped Object round-trip + mapper.addMixIn(MemoryUpdate.class, MemoryUpdateValueMixin.class); + return mapper; } @@ -92,4 +110,61 @@ public void serialize(ActionTask value, JsonGenerator gen, SerializerProvider se gen.writeNull(); } } + + /** Binds the byte[]-preserving (de)serializer to {@link MemoryUpdate#getValue()} only. */ + abstract static class MemoryUpdateValueMixin { + @JsonSerialize(using = MemoryUpdateValueSerializer.class) + @JsonDeserialize(using = MemoryUpdateValueDeserializer.class) + Object value; + } + + /** + * Serializes a {@code byte[]} {@link MemoryUpdate} value as a one-key base64 envelope so it can + * be recovered as a {@code byte[]} rather than the base64 {@code String} that untyped {@code + * Object} serialization would yield; all other types delegate to default serialization and stay + * byte-identical on disk. + * + *

Only a top-level {@code byte[]} value is preserved; a {@code byte[]} nested inside a + * {@code List} or {@code Map} value goes through the default serializers and is not preserved. + */ + static class MemoryUpdateValueSerializer extends JsonSerializer { + @Override + public void serialize(Object value, JsonGenerator gen, SerializerProvider provider) + throws IOException { + if (value instanceof byte[]) { + gen.writeStartObject(); + gen.writeStringField( + MEMORY_UPDATE_BYTES_KEY, + Base64.getEncoder().encodeToString((byte[]) value)); + gen.writeEndObject(); + } else { + provider.defaultSerializeValue(value, gen); + } + } + } + + /** + * Reverses the base64 envelope written by {@link MemoryUpdateValueSerializer} back to a {@code + * byte[]}; every other shape reproduces stock untyped {@code Object} binding. + * + *

Only a top-level {@code byte[]} value is recovered; a {@code byte[]} nested inside a + * {@code List} or {@code Map} value is not. + */ + static class MemoryUpdateValueDeserializer extends JsonDeserializer { + @Override + public Object deserialize(JsonParser p, DeserializationContext ctxt) throws IOException { + JsonNode node = ctxt.readTree(p); + if (node.isObject() + && node.size() == 1 + && node.has(MEMORY_UPDATE_BYTES_KEY) + && node.get(MEMORY_UPDATE_BYTES_KEY).isTextual()) { + try { + return Base64.getDecoder().decode(node.get(MEMORY_UPDATE_BYTES_KEY).asText()); + } catch (IllegalArgumentException notBase64) { + // Not a real envelope - fall through to generic binding. + } + } + return ctxt.readTreeAsValue(node, Object.class); + } + } } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java index e0c4de0bf..3d7aeae40 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java @@ -25,8 +25,10 @@ import java.nio.charset.StandardCharsets; import java.util.ArrayList; +import java.util.Arrays; import java.util.Base64; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; @@ -136,6 +138,114 @@ public void testActionStateWithComplexAttributes() throws Exception { assertEquals(42, deserializedComplexAttr.get("number")); } + @Test + public void testByteArrayShortTermMemoryValueRoundTrip() throws Exception { + // Verbatim issue repro: a byte[] short-term memory value must survive the durable serde + // round-trip as a byte[], not silently decay to a base64 String. + byte[] original = "hello".getBytes(StandardCharsets.UTF_8); + ActionState state = new ActionState(null, null, null, null, null, false); + state.addShortTermMemoryUpdate(new MemoryUpdate("stm.path", original)); + + ActionState recovered = ActionStateSerde.deserialize(ActionStateSerde.serialize(state)); + + Object value = recovered.getShortTermMemoryUpdates().get(0).getValue(); + assertInstanceOf(byte[].class, value); + assertArrayEquals(original, (byte[]) value); + } + + @Test + public void testByteArraySensoryMemoryValueRoundTrip() throws Exception { + byte[] original = "hello".getBytes(StandardCharsets.UTF_8); + ActionState state = new ActionState(null, null, null, null, null, false); + state.addSensoryMemoryUpdate(new MemoryUpdate("sm.path", original)); + + ActionState recovered = ActionStateSerde.deserialize(ActionStateSerde.serialize(state)); + + Object value = recovered.getSensoryMemoryUpdates().get(0).getValue(); + assertInstanceOf(byte[].class, value); + assertArrayEquals(original, (byte[]) value); + } + + @Test + public void testEmptyByteArrayValueRoundTrip() throws Exception { + byte[] original = new byte[0]; + ActionState state = new ActionState(null, null, null, null, null, false); + state.addShortTermMemoryUpdate(new MemoryUpdate("stm.path", original)); + + ActionState recovered = ActionStateSerde.deserialize(ActionStateSerde.serialize(state)); + + Object value = recovered.getShortTermMemoryUpdates().get(0).getValue(); + assertInstanceOf(byte[].class, value); + assertEquals(0, ((byte[]) value).length); + } + + @Test + public void testStringMemoryValueUnchanged() throws Exception { + // Guard: the envelope must not disturb the common String case. + ActionState state = new ActionState(null, null, null, null, null, false); + state.addShortTermMemoryUpdate(new MemoryUpdate("stm.path", "plain string value")); + + ActionState recovered = ActionStateSerde.deserialize(ActionStateSerde.serialize(state)); + + Object value = recovered.getShortTermMemoryUpdates().get(0).getValue(); + assertInstanceOf(String.class, value); + assertEquals("plain string value", value); + } + + @Test + public void testMapMemoryValueUnchanged() throws Exception { + // Guard: a Map value round-trips as a Map, not misread as a byte[] envelope. + Map original = new LinkedHashMap<>(); + original.put("a", 1); + original.put("b", Arrays.asList(1, 2, 3)); + ActionState state = new ActionState(null, null, null, null, null, false); + state.addShortTermMemoryUpdate(new MemoryUpdate("stm.path", original)); + + ActionState recovered = ActionStateSerde.deserialize(ActionStateSerde.serialize(state)); + + Object value = recovered.getShortTermMemoryUpdates().get(0).getValue(); + assertInstanceOf(Map.class, value); + @SuppressWarnings("unchecked") + Map recoveredMap = (Map) value; + assertEquals(1, recoveredMap.get("a")); + assertEquals(Arrays.asList(1, 2, 3), recoveredMap.get("b")); + } + + @Test + public void testReservedKeyMapNotByteArrayIsDocumentedResidual() throws Exception { + // Documented residual, not a bug: because the reserved key is the collision guard, a user + // Map that is EXACTLY {"__flink_agents_bytes__":""} is indistinguishable from + // a byte[] envelope and recovers as byte[]. A namespaced key makes this negligible. A map + // with any second key falls outside the single-key predicate and stays a Map. + String reservedKey = "__flink_agents_bytes__"; + + Map collidingMap = new LinkedHashMap<>(); + collidingMap.put(reservedKey, "aGVsbG8="); // base64 of "hello" + ActionState collidingState = new ActionState(null, null, null, null, null, false); + collidingState.addShortTermMemoryUpdate(new MemoryUpdate("stm.path", collidingMap)); + + ActionState recoveredColliding = + ActionStateSerde.deserialize(ActionStateSerde.serialize(collidingState)); + Object collidingValue = recoveredColliding.getShortTermMemoryUpdates().get(0).getValue(); + assertInstanceOf(byte[].class, collidingValue); + assertArrayEquals("hello".getBytes(StandardCharsets.UTF_8), (byte[]) collidingValue); + + Map twoKeyMap = new LinkedHashMap<>(); + twoKeyMap.put(reservedKey, "aGVsbG8="); + twoKeyMap.put("x", 1); + ActionState twoKeyState = new ActionState(null, null, null, null, null, false); + twoKeyState.addShortTermMemoryUpdate(new MemoryUpdate("stm.path", twoKeyMap)); + + ActionState recoveredTwoKey = + ActionStateSerde.deserialize(ActionStateSerde.serialize(twoKeyState)); + Object twoKeyValue = recoveredTwoKey.getShortTermMemoryUpdates().get(0).getValue(); + assertInstanceOf(Map.class, twoKeyValue); + @SuppressWarnings("unchecked") + Map recoveredTwoKeyMap = (Map) twoKeyValue; + assertEquals("aGVsbG8=", recoveredTwoKeyMap.get(reservedKey)); + assertEquals(1, recoveredTwoKeyMap.get("x")); + } + @Test public void testActionStateWithCallResults() throws Exception { // Create ActionState with call results