diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/client/OpenAiChatClientMethodInvokingFunctionCallbackIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/client/OpenAiChatClientMethodInvokingFunctionCallbackIT.java
index 2a766a6b0..c405aacad 100644
--- a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/client/OpenAiChatClientMethodInvokingFunctionCallbackIT.java
+++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/client/OpenAiChatClientMethodInvokingFunctionCallbackIT.java
@@ -37,6 +37,7 @@ import org.springframework.test.context.ActiveProfiles;
 import org.springframework.util.ReflectionUtils;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy;
 
 @SpringBootTest(classes = OpenAiTestConfiguration.class)
@@ -165,12 +166,12 @@ class OpenAiChatClientMethodInvokingFunctionCallbackIT {
 	}
 
 	@Test
-	void methodGetWeatherToolContextButMissingContextArgument() {
+	void methodGetWeatherNonStaticButWithToolContext() {
 
 		TestFunctionClass targetObject = new TestFunctionClass();
 
-		var toolMethod = ReflectionUtils.findMethod(TestFunctionClass.class, "getWeatherWithContext", String.class,
-				Unit.class, ToolContext.class);
+		var toolMethod = ReflectionUtils.findMethod(TestFunctionClass.class, "getWeatherNonStatic", String.class,
+				Unit.class);
 
 		// @formatter:off
 		assertThatThrownBy(() -> ChatClient.create(this.chatModel).prompt()
@@ -182,6 +183,7 @@ class OpenAiChatClientMethodInvokingFunctionCallbackIT {
 					.toolMethod(toolMethod)
 					.toolObject(targetObject)
 					.build())
+				.toolContext(Map.of("tool-context", "value"))
 				.call()
 				.content())
 				.isInstanceOf(IllegalArgumentException.class)
@@ -189,6 +191,29 @@ class OpenAiChatClientMethodInvokingFunctionCallbackIT {
 		// @formatter:on
 	}
 
+	@Test
+	void methodGetWeatherToolContextButWithoutToolContext() {
+
+		TestFunctionClass targetObject = new TestFunctionClass();
+
+		var toolMethod = ReflectionUtils.findMethod(TestFunctionClass.class, "getWeatherWithContext", String.class,
+				Unit.class, ToolContext.class);
+
+		// @formatter:off
+		assertThatCode(() ->ChatClient.create(this.chatModel).prompt()
+				.user("What's the weather like in San Francisco, Tokyo, and Paris?  Use Celsius.")
+				.toolCallbacks(MethodToolCallback.builder()
+						.toolDefinition(ToolDefinitions.builder(toolMethod)
+								.description("Get the weather in location")
+								.build())
+						.toolMethod(toolMethod)
+						.toolObject(targetObject)
+						.build())
+				.call()
+				.content()).doesNotThrowAnyException();
+		// @formatter:on
+	}
+
 	@Test
 	void methodNoParameters() {