[api][java][python] Propagate user context from request events to responses

api/src/main/java/org/apache/flink/agents/api/event/ChatRequestEvent.java

@@ -42,8 +42,10 @@ public class ChatRequestEvent extends Event {
 
     /**
      * The fixed cross-language attribute schema for this event: {@code model} and {@code messages}
-     * are required, {@code prompt_args} and {@code output_schema} are optional, and no other
-     * attribute is allowed. {@link #fromEvent} enforces it at the JSON boundary.
+     * are required, {@code prompt_args}, {@code output_schema}, and {@code user_context} are
+     * optional, and no other attribute is allowed. {@link #fromEvent} enforces it at the JSON
+     * boundary. {@code user_context} is application-defined data propagated to the response; it is
+     * not part of the prompt.
      */
     private static final List<BuiltInAttribute> ATTRIBUTE_SCHEMA =
             List.of(
@@ -54,20 +56,31 @@ public class ChatRequestEvent extends Event {
                             ChatMessage.class,
                             Map.class),
                     BuiltInAttribute.optional("prompt_args", Map.class),
-                    BuiltInAttribute.optionalUntyped("output_schema"));
+                    BuiltInAttribute.optionalUntyped("output_schema"),
+                    BuiltInAttribute.optional(UserDefinedContext.ATTRIBUTE, Map.class));
 
     public ChatRequestEvent(
             String model,
             List<ChatMessage> messages,
             @Nullable Map<String, Object> promptArgs,
-            @Nullable Object outputSchema) {
+            @Nullable Object outputSchema,
+            @Nullable Map<String, Object> userContext) {
         super(EVENT_TYPE);
         setAttr("model", model);
         setAttr("messages", new ArrayList<>(messages));
         setAttr("prompt_args", promptArgs != null ? promptArgs : Collections.emptyMap());
         if (outputSchema != null) {
             setAttr("output_schema", outputSchema);
         }
+        UserDefinedContext.set(this, userContext);
+    }
+
+    public ChatRequestEvent(
+            String model,
+            List<ChatMessage> messages,
+            @Nullable Map<String, Object> promptArgs,
+            @Nullable Object outputSchema) {
+        this(model, messages, promptArgs, outputSchema, null);
     }
 
     public ChatRequestEvent(
@@ -108,10 +121,10 @@ private static Map<String, Object> normalizeAttributes(Map<String, Object> attri
      * Reconstructs a typed ChatRequestEvent from a base Event, deserializing nested types.
      *
      * <p>Enforces the fixed cross-language attribute schema: {@code model} and {@code messages} are
-     * required, {@code prompt_args} and {@code output_schema} are optional, no other attribute is
-     * allowed, {@code model} must be a string, and every {@code messages} element must be a {@link
-     * ChatMessage} or its serialized map — an element such as a bare number is rejected rather than
-     * silently dropped.
+     * required, {@code prompt_args}, {@code output_schema}, and {@code user_context} are optional,
+     * no other attribute is allowed, {@code model} must be a string, and every {@code messages}
+     * element must be a {@link ChatMessage} or its serialized map — an element such as a bare
+     * number is rejected rather than silently dropped.
      *
      * @param event the base event containing chat request data in attributes
      * @return a typed ChatRequestEvent
@@ -145,4 +158,15 @@ public Map<String, Object> getPromptArgs() {
         Map<String, Object> args = (Map<String, Object>) getAttr("prompt_args");
         return args != null ? args : Collections.emptyMap();
     }
+
+    /**
+     * Returns application-defined data to propagate to the matching {@link ChatResponseEvent}, or
+     * {@code null} when this request omitted it. The map is not added to messages or prompt
+     * arguments.
+     */
+    @JsonIgnore
+    @Nullable
+    public Map<String, Object> getUserContext() {
+        return UserDefinedContext.get(this);
+    }
 }

api/src/main/java/org/apache/flink/agents/api/event/ChatResponseEvent.java

@@ -25,6 +25,8 @@
 import org.apache.flink.agents.api.Event;
 import org.apache.flink.agents.api.chat.messages.ChatResult;
 
+import javax.annotation.Nullable;
+
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -54,7 +56,8 @@ public class ChatResponseEvent extends Event {
                     BuiltInAttribute.optional(RETRY_COUNT, Number.class),
                     BuiltInAttribute.optional(TOTAL_RETRY_WAIT_SEC, Number.class),
                     BuiltInAttribute.optionalUntyped(STRUCTURED_OUTPUT),
-                    BuiltInAttribute.optional(MODEL_ROUTING, Map.class));
+                    BuiltInAttribute.optional(MODEL_ROUTING, Map.class),
+                    BuiltInAttribute.optional(UserDefinedContext.ATTRIBUTE, Map.class));
 
     private static final ObjectMapper MAPPER = new ObjectMapper();
 
@@ -143,6 +146,25 @@ public void setRoutingMetadata(Map<String, Object> value) {
         setAttr(MODEL_ROUTING, value);
     }
 
+    /**
+     * Attaches application-defined data from the originating chat request. {@code null} omits it.
+     * The value is not derived from event attachments.
+     */
+    public ChatResponseEvent withUserContext(@Nullable Map<String, Object> userContext) {
+        UserDefinedContext.set(this, userContext);
+        return this;
+    }
+
+    /**
+     * Returns the originating request's application-defined data, or {@code null} when it was
+     * omitted.
+     */
+    @JsonIgnore
+    @Nullable
+    public Map<String, Object> getUserContext() {
+        return UserDefinedContext.get(this);
+    }
+
     @JsonIgnore
     public String getStatus() {
         return (String) getAttr(STATUS);

api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEvent.java

@@ -23,6 +23,8 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import org.apache.flink.agents.api.Event;
 
+import javax.annotation.Nullable;
+
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
@@ -38,7 +40,8 @@ public class ContextRetrievalRequestEvent extends Event {
             List.of(
                     BuiltInAttribute.required("query", String.class),
                     BuiltInAttribute.required("vector_store", String.class),
-                    BuiltInAttribute.optional("max_results", Number.class));
+                    BuiltInAttribute.optional("max_results", Number.class),
+                    BuiltInAttribute.optional(UserDefinedContext.ATTRIBUTE, Map.class));
 
     public ContextRetrievalRequestEvent(String query, String vectorStore) {
         this(query, vectorStore, DEFAULT_MAX_RESULTS);
@@ -50,14 +53,32 @@ public ContextRetrievalRequestEvent(String query, String vectorStore) {
      * @param query the search query text
      * @param vectorStore the name of the vector store to query
      * @param maxResults the maximum number of results to return; must be positive
+     * @param userContext application-defined data propagated to the response; {@code null} omits it
      * @throws IllegalArgumentException if {@code maxResults} is non-positive
      */
-    public ContextRetrievalRequestEvent(String query, String vectorStore, int maxResults) {
+    public ContextRetrievalRequestEvent(
+            String query,
+            String vectorStore,
+            int maxResults,
+            @Nullable Map<String, Object> userContext) {
         super(EVENT_TYPE);
         validateMaxResults(maxResults);
         setAttr("query", query);
         setAttr("vector_store", vectorStore);
         setAttr("max_results", maxResults);
+        UserDefinedContext.set(this, userContext);
+    }
+
+    /**
+     * Creates a new context retrieval request event.
+     *
+     * @param query the search query text
+     * @param vectorStore the name of the vector store to query
+     * @param maxResults the maximum number of results to return; must be positive
+     * @throws IllegalArgumentException if {@code maxResults} is non-positive
+     */
+    public ContextRetrievalRequestEvent(String query, String vectorStore, int maxResults) {
+        this(query, vectorStore, maxResults, null);
     }
 
     /**
@@ -93,8 +114,8 @@ private static void validateMaxResults(int maxResults) {
      * Reconstructs a typed ContextRetrievalRequestEvent from a base Event.
      *
      * <p>Enforces the fixed cross-language schema: {@code query} and {@code vector_store} are
-     * required strings, {@code max_results} is an optional positive number, and no other attribute
-     * is allowed.
+     * required strings, {@code max_results} is an optional positive number, {@code user_context} is
+     * an optional map, and no other attribute is allowed.
      *
      * @param event the base event containing context retrieval request data in attributes
      * @return a typed ContextRetrievalRequestEvent
@@ -125,6 +146,16 @@ public int getMaxResults() {
         return ((Number) getAttr("max_results")).intValue();
     }
 
+    /**
+     * Returns application-defined data to propagate to the matching response, or {@code null} when
+     * this request omitted it. The map is not added to the vector-store query.
+     */
+    @JsonIgnore
+    @Nullable
+    public Map<String, Object> getUserContext() {
+        return UserDefinedContext.get(this);
+    }
+
     @Override
     public String toString() {
         return "ContextRetrievalRequestEvent{"

api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalResponseEvent.java

@@ -25,6 +25,8 @@
 import org.apache.flink.agents.api.Event;
 import org.apache.flink.agents.api.vectorstores.Document;
 
+import javax.annotation.Nullable;
+
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -45,7 +47,8 @@ public class ContextRetrievalResponseEvent extends Event {
                             "documents",
                             "a Document or its serialized map",
                             Document.class,
-                            Map.class));
+                            Map.class),
+                    BuiltInAttribute.optional(UserDefinedContext.ATTRIBUTE, Map.class));
 
     public ContextRetrievalResponseEvent(UUID requestId, String query, List<Document> documents) {
         super(EVENT_TYPE);
@@ -88,8 +91,9 @@ private static Map<String, Object> normalizeAttributes(Map<String, Object> attri
      * types.
      *
      * <p>Enforces the fixed cross-language schema: {@code request_id}, {@code query}, and {@code
-     * documents} are required and no other attribute is allowed; {@code query} must be a string and
-     * every {@code documents} element must be a {@link Document} or its serialized map.
+     * documents} are required, {@code user_context} is an optional map, and no other attribute is
+     * allowed; {@code query} must be a string and every {@code documents} element must be a {@link
+     * Document} or its serialized map.
      *
      * @param event the base event containing context retrieval response data in attributes
      * @return a typed ContextRetrievalResponseEvent
@@ -120,6 +124,26 @@ public List<Document> getDocuments() {
         return (List<Document>) getAttr("documents");
     }
 
+    /**
+     * Attaches application-defined data from the originating retrieval request. {@code null} omits
+     * it.
+     */
+    public ContextRetrievalResponseEvent withUserContext(
+            @Nullable Map<String, Object> userContext) {
+        UserDefinedContext.set(this, userContext);
+        return this;
+    }
+
+    /**
+     * Returns the originating request's application-defined data, or {@code null} when it was
+     * omitted.
+     */
+    @JsonIgnore
+    @Nullable
+    public Map<String, Object> getUserContext() {
+        return UserDefinedContext.get(this);
+    }
+
     @Override
     public String toString() {
         return "ContextRetrievalResponseEvent{"

api/src/main/java/org/apache/flink/agents/api/event/DecisionRequestEvent.java

@@ -43,7 +43,8 @@ public class DecisionRequestEvent extends Event {
             List.of(
                     BuiltInAttribute.required("model", String.class),
                     BuiltInAttribute.requiredUntyped("request"),
-                    BuiltInAttribute.optionalUntyped("request_params"));
+                    BuiltInAttribute.optionalUntyped("request_params"),
+                    BuiltInAttribute.optional(UserDefinedContext.ATTRIBUTE, Map.class));
 
     /** The model is the registered resource name, not a provider model identifier. */
     public DecisionRequestEvent(String model, DecisionRequest request) {
@@ -53,10 +54,25 @@ public DecisionRequestEvent(String model, DecisionRequest request) {
     /** Requests a decision with provider-specific parameters for this call. */
     public DecisionRequestEvent(
             String model, DecisionRequest request, @Nullable Map<String, Object> requestParams) {
+        this(model, request, requestParams, null);
+    }
+
+    /**
+     * Requests a decision with provider-specific parameters and application-defined data.
+     *
+     * <p>{@code userContext} is propagated to the response and is not part of {@code request} or
+     * {@code requestParams}. {@code null} omits it.
+     */
+    public DecisionRequestEvent(
+            String model,
+            DecisionRequest request,
+            @Nullable Map<String, Object> requestParams,
+            @Nullable Map<String, Object> userContext) {
         super(EVENT_TYPE);
         setAttr("model", Objects.requireNonNull(model, "model"));
         setAttr("request", Objects.requireNonNull(request, "request"));
         setAttr("request_params", normalizeRequestParams(requestParams));
+        UserDefinedContext.set(this, userContext);
     }
 
     @JsonCreator
@@ -111,4 +127,15 @@ public DecisionRequest getRequest() {
     public Map<String, Object> getRequestParams() {
         return (Map<String, Object>) getAttr("request_params");
     }
+
+    /**
+     * Returns application-defined data to propagate to the matching response, or {@code null} when
+     * this request omitted it. The map is not added to the decision request or {@code
+     * request_params}.
+     */
+    @JsonIgnore
+    @Nullable
+    public Map<String, Object> getUserContext() {
+        return UserDefinedContext.get(this);
+    }
 }

api/src/main/java/org/apache/flink/agents/api/event/DecisionResponseEvent.java

@@ -25,6 +25,8 @@
 import org.apache.flink.agents.api.Event;
 import org.apache.flink.agents.api.decision.DecisionResult;
 
+import javax.annotation.Nullable;
+
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -51,7 +53,8 @@ public class DecisionResponseEvent extends Event {
                     BuiltInAttribute.optionalUntyped(RESPONSE),
                     BuiltInAttribute.optionalUntyped(ERROR),
                     BuiltInAttribute.optional(RETRY_COUNT, Number.class),
-                    BuiltInAttribute.optional(TOTAL_RETRY_WAIT_SEC, Number.class));
+                    BuiltInAttribute.optional(TOTAL_RETRY_WAIT_SEC, Number.class),
+                    BuiltInAttribute.optional(UserDefinedContext.ATTRIBUTE, Map.class));
 
     private static final ObjectMapper MAPPER = new ObjectMapper();
 
@@ -206,4 +209,23 @@ public int getRetryCount() {
     public int getTotalRetryWaitSec() {
         return ((Number) getAttr(TOTAL_RETRY_WAIT_SEC)).intValue();
     }
+
+    /**
+     * Attaches application-defined data from the originating decision request. {@code null} omits
+     * it. Present on both successful and failed responses.
+     */
+    public DecisionResponseEvent withUserContext(@Nullable Map<String, Object> userContext) {
+        UserDefinedContext.set(this, userContext);
+        return this;
+    }
+
+    /**
+     * Returns the originating request's application-defined data, or {@code null} when it was
+     * omitted.
+     */
+    @JsonIgnore
+    @Nullable
+    public Map<String, Object> getUserContext() {
+        return UserDefinedContext.get(this);
+    }
 }