1 files changed, 75 insertions, 0 deletions
diff --git a/internal/server/handlers/generation_output.go b/internal/server/handlers/generation_output.go
new file mode 100644
index 0000000..aa9d195
--- /dev/null
+++ b/internal/server/handlers/generation_output.go
@@ -0,0 +1,75 @@
+package handlers
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+	"strings"
+)
+
+type sessionGenerationKey struct{}
+
+const generationOutputPrefix = "\x1egen:"
+
+func withSessionGeneration(ctx context.Context, generation uint64) context.Context {
+	if ctx == nil || generation == 0 {
+		return ctx
+	}
+	return context.WithValue(ctx, sessionGenerationKey{}, generation)
+}
+
+func sessionGenerationFromContext(ctx context.Context) uint64 {
+	if ctx == nil {
+		return 0
+	}
+
+	generation, _ := ctx.Value(sessionGenerationKey{}).(uint64)
+	return generation
+}
+
+func encodeGeneratedMessage(generation uint64, message string) string {
+	if generation == 0 {
+		return message
+	}
+	return fmt.Sprintf("%s%d\x1e%s", generationOutputPrefix, generation, message)
+}
+
+func decodeGeneratedMessage(message string) (uint64, string) {
+	if !strings.HasPrefix(message, generationOutputPrefix) {
+		return 0, message
+	}
+
+	rest := strings.TrimPrefix(message, generationOutputPrefix)
+	parts := strings.SplitN(rest, "\x1e", 2)
+	if len(parts) != 2 {
+		return 0, message
+	}
+
+	generation, err := strconv.ParseUint(parts[0], 10, 64)
+	if err != nil {
+		return 0, message
+	}
+
+	return generation, parts[1]
+}
+
+func encodeGeneratedBytes(generation uint64, payload []byte) []byte {
+	if generation == 0 {
+		return payload
+	}
+
+	prefix := []byte(fmt.Sprintf("%s%d\x1e", generationOutputPrefix, generation))
+	data := make([]byte, 0, len(prefix)+len(payload))
+	data = append(data, prefix...)
+	data = append(data, payload...)
+	return data
+}
+
+func decodeGeneratedBytes(payload []byte) (uint64, []byte) {
+	message := string(payload)
+	generation, decoded := decodeGeneratedMessage(message)
+	if generation == 0 {
+		return 0, payload
+	}
+	return generation, []byte(decoded)
+}