[integrations][java][python] Preserve Watsonx finish_reason in extra_args

integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java

@@ -553,14 +553,20 @@ static ArrayNode convertTools(List<Tool> tools) {
         }
     }
 
+    /**
+     * Parses the watsonx.ai chat response. When the provider reports a finish reason it is carried
+     * verbatim in {@code extraArgs} under {@code finish_reason}, including values outside the
+     * documented set, and the entry is absent when the provider reports none.
+     */
     @VisibleForTesting
     static ChatMessage parseResponse(JsonNode response, String modelName) {
         final JsonNode choice = response.required("choices").get(0);
         final JsonNode responseMessage = choice.required("message");
 
         final JsonNode finishReasonNode = choice.get("finish_reason");
+        String finishReason = null;
         if (finishReasonNode != null && !finishReasonNode.isNull()) {
-            final String finishReason = finishReasonNode.asText();
+            finishReason = finishReasonNode.asText();
             if (!"stop".equals(finishReason) && !"tool_calls".equals(finishReason)) {
                 LOG.warn(
                         "watsonx.ai chat for model {} finished with reason '{}'; the response"
@@ -602,11 +608,18 @@ static ChatMessage parseResponse(JsonNode response, String modelName) {
         }
 
         final JsonNode usage = response.get("usage");
-        if (modelName != null && !modelName.isBlank() && usage != null && !usage.isNull()) {
+        final boolean hasUsageMetadata =
+                modelName != null && !modelName.isBlank() && usage != null && !usage.isNull();
+        if (hasUsageMetadata || finishReason != null) {
             final Map<String, Object> extraArgs = new HashMap<>(chatMessage.getExtraArgs());
-            extraArgs.put("model_name", modelName);
-            extraArgs.put("promptTokens", usage.path("prompt_tokens").asLong(0));
-            extraArgs.put("completionTokens", usage.path("completion_tokens").asLong(0));
+            if (hasUsageMetadata) {
+                extraArgs.put("model_name", modelName);
+                extraArgs.put("promptTokens", usage.path("prompt_tokens").asLong(0));
+                extraArgs.put("completionTokens", usage.path("completion_tokens").asLong(0));
+            }
+            if (finishReason != null) {
+                extraArgs.put("finish_reason", finishReason);
+            }
             chatMessage.setExtraArgs(extraArgs);
         }
 

integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java

@@ -380,6 +380,64 @@ void testParseResponse() throws Exception {
                 .isEqualTo("ibm/granite-3-3-8b-instruct");
         assertThat(message.getExtraArgs().get("promptTokens")).isEqualTo(100L);
         assertThat(message.getExtraArgs().get("completionTokens")).isEqualTo(50L);
+        assertThat(message.getExtraArgs()).containsEntry("finish_reason", "stop");
+    }
+
+    @Test
+    @DisplayName("A finish reason outside the documented set is stored as received")
+    void testParseResponseCarriesUnknownFinishReasonVerbatim() throws Exception {
+        JsonNode response =
+                MAPPER.readTree(
+                        "{\"choices\": [{\"index\": 0, \"message\": {\"role\": \"assistant\","
+                                + " \"content\": \"hi\"}, \"finish_reason\":"
+                                + " \"some_vendor_reason\"}]}");
+
+        ChatMessage message = WatsonxChatModelConnection.parseResponse(response, null);
+
+        assertThat(message.getExtraArgs()).containsEntry("finish_reason", "some_vendor_reason");
+    }
+
+    @Test
+    @DisplayName("A response with no finish_reason member yields no key and no error")
+    void testParseResponseNoFinishReasonKeyWhenMemberAbsent() throws Exception {
+        JsonNode response =
+                MAPPER.readTree(
+                        "{\"choices\": [{\"index\": 0, \"message\": {\"role\": \"assistant\","
+                                + " \"content\": \"hi\"}}]}");
+
+        ChatMessage message = WatsonxChatModelConnection.parseResponse(response, null);
+
+        assertThat(message.getExtraArgs()).doesNotContainKey("finish_reason");
+    }
+
+    @Test
+    @DisplayName("A response whose finish_reason is JSON null yields no key and no error")
+    void testParseResponseNoFinishReasonKeyWhenJsonNull() throws Exception {
+        JsonNode response =
+                MAPPER.readTree(
+                        "{\"choices\": [{\"index\": 0, \"message\": {\"role\": \"assistant\","
+                                + " \"content\": \"hi\"}, \"finish_reason\": null}]}");
+
+        ChatMessage message = WatsonxChatModelConnection.parseResponse(response, null);
+
+        assertThat(message.getExtraArgs()).doesNotContainKey("finish_reason");
+    }
+
+    @Test
+    @DisplayName("The finish reason is captured independently of the token metrics")
+    void testParseResponseCarriesFinishReasonWithoutUsage() throws Exception {
+        // modelName is null here, so the usage-metadata branch cannot run; this proves
+        // finish_reason capture does not depend on it.
+        JsonNode response =
+                MAPPER.readTree(
+                        "{\"choices\": [{\"index\": 0, \"message\": {\"role\": \"assistant\","
+                                + " \"content\": \"hi\"}, \"finish_reason\": \"tool_calls\"}]}");
+
+        ChatMessage message = WatsonxChatModelConnection.parseResponse(response, null);
+
+        assertThat(message.getExtraArgs())
+                .containsEntry("finish_reason", "tool_calls")
+                .doesNotContainKey("promptTokens");
     }
 
     @Test

python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_chat_model.py

@@ -130,6 +130,7 @@ def get_resource(name: str, type: ResourceType) -> Resource:
     assert response.extra_args["model_name"] == test_model
     assert response.extra_args["promptTokens"] == 100
     assert response.extra_args["completionTokens"] == 50
+    assert response.extra_args["finish_reason"] == "stop"
     model_inference.assert_called_once_with(
         model_id=test_model,
         api_client=api_client,
@@ -139,6 +140,63 @@ def get_resource(name: str, type: ResourceType) -> Resource:
     )
 
 
+def _setup_llm(
+    mock_model: MagicMock, monkeypatch: pytest.MonkeyPatch
+) -> WatsonxChatModelSetup:
+    model_inference = MagicMock(return_value=mock_model)
+    monkeypatch.setattr(
+        "flink_agents.integrations.chat_models.watsonx.watsonx_chat_model.ModelInference",
+        model_inference,
+    )
+
+    connection = _fake_connection()
+    connection._client = MagicMock()
+
+    def get_resource(name: str, type: ResourceType) -> Resource:
+        return connection
+
+    mock_ctx = MagicMock(spec=ResourceContext)
+    mock_ctx.get_resource = get_resource
+
+    llm = WatsonxChatModelSetup(
+        model=test_model,
+        connection="watsonx",
+        resource_context=mock_ctx,
+    )
+    llm.open()
+    return llm
+
+
+def test_watsonx_chat_carries_unknown_finish_reason_verbatim(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A finish reason outside the documented set is stored as received."""
+    mock_model = MagicMock()
+    mock_model.chat.return_value = _mock_chat_response(
+        {"role": "assistant", "content": "hi"}, finish_reason="some_vendor_reason"
+    )
+    llm = _setup_llm(mock_model, monkeypatch)
+
+    response = llm.chat([ChatMessage(role=MessageRole.USER, content="Hello!")])
+
+    assert response.extra_args["finish_reason"] == "some_vendor_reason"
+
+
+def test_watsonx_chat_no_finish_reason_key_when_none(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A response with no finish reason yields no key and no error."""
+    mock_model = MagicMock()
+    mock_model.chat.return_value = _mock_chat_response(
+        {"role": "assistant", "content": "hi"}, finish_reason=None
+    )
+    llm = _setup_llm(mock_model, monkeypatch)
+
+    response = llm.chat([ChatMessage(role=MessageRole.USER, content="Hello!")])
+
+    assert "finish_reason" not in response.extra_args
+
+
 def test_watsonx_tool_call_response_mocked(monkeypatch: pytest.MonkeyPatch) -> None:
     """Test that tool call responses are converted to the framework format."""
     mock_model = MagicMock()

python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py

@@ -345,6 +345,10 @@ def chat(
         structured-output translation, so callers stay on the prompt-engineering
         fallback. Declaring the parameter keeps a caller-supplied schema out of
         ``**kwargs``, which is forwarded to the provider SDK.
+
+        When the response carries a finish reason, it is available verbatim as
+        ``extra_args["finish_reason"]``; the key is absent when the provider
+        reports none.
         """
         self._reject_unsupported_output_schema(output_schema)
         model_name = kwargs.pop("model", DEFAULT_MODEL)
@@ -401,6 +405,8 @@ def chat(
                 model_name,
                 finish_reason,
             )
+        if finish_reason is not None:
+            extra_args["finish_reason"] = finish_reason
 
         response_message: Dict[str, Any] = choice["message"]