From 3c8fadf3480341f0ad8986733b418b571bece4f8 Mon Sep 17 00:00:00 2001 From: Brian Strauch Date: Thu, 15 Oct 2026 19:58:25 -0700 Subject: [PATCH 1/2] Migrate Spring AI samples to standalone package and Spring AI 2 --- README.md | 6 +- gradle/springai.gradle | 17 ++-- .../springai/chat/ChatWorkflowImpl.java | 10 ++- .../basic/src/main/resources/application.yaml | 1 + .../springai/chat/ChatWorkflowTest.java | 89 +++++++++++++++++++ .../samples/springai/mcp/McpWorkflowImpl.java | 12 ++- .../mcp/src/main/resources/application.yaml | 1 + .../springai/multimodel/ChatModelConfig.java | 17 ++-- .../multimodel/MultiModelWorkflowImpl.java | 11 +-- .../multimodel/ChatModelConfigTest.java | 40 +++++++++ .../rag/src/main/resources/application.yaml | 1 + 11 files changed, 175 insertions(+), 30 deletions(-) create mode 100644 springai/basic/src/test/java/io/temporal/samples/springai/chat/ChatWorkflowTest.java create mode 100644 springai/multimodel/src/test/java/io/temporal/samples/springai/multimodel/ChatModelConfigTest.java diff --git a/README.md b/README.md index 7a19472a..a3245bcb 100644 --- a/README.md +++ b/README.md @@ -230,7 +230,9 @@ To run any of the SpringBoot samples in your Temporal Cloud namespace: ### Running Spring AI Samples -The Spring AI samples demonstrate the [Temporal Spring AI integration](https://github.com/temporalio/sdk-java/tree/master/temporal-spring-ai), which makes Spring AI agents durable on Temporal — model calls run as Temporal Activities recorded in Workflow history, and tools are dispatched per their type so they fit Workflow execution. +The Spring AI samples demonstrate the [Temporal Spring AI integration](https://github.com/temporalio/ai-integrations/tree/main/java/spring-ai), which makes Spring AI agents durable on Temporal — model calls run as Temporal Activities recorded in Workflow history, and tools are dispatched per their type so they fit Workflow execution. + +These samples use the Public Preview `io.temporal:spring-ai:0.1.0` package with Spring AI 2.0.1, Spring Boot 4.0.8, and Temporal Java SDK 1.40.0. Java 17 or later is required. Spring AI 1 / Spring Boot 3 applications should remain on `io.temporal:temporal-spring-ai:1.40.x`. Spring AI 1 Workflow histories cannot replay on the new integration; complete existing executions or use separate Worker deployments when upgrading. Each sample is its own Spring Boot application with an interactive CLI. Run from the main repo dir: @@ -242,7 +244,7 @@ Each sample is its own Spring Boot application with an interactive CLI. Run from All samples need an `OPENAI_API_KEY` environment variable; some need additional setup (see each sample's source for details). More info on each sample: -- [**Basic**](/springai/basic): Chat workflow with three tool flavors — activity-backed (`WeatherActivity`), plain workflow tools (`StringTools`), and `@SideEffectTool` (`TimestampTools`) — plus a `PromptChatMemoryAdvisor` for conversation history. +- [**Basic**](/springai/basic): Chat workflow with three tool flavors — activity-backed (`WeatherActivity`), plain workflow tools (`StringTools`), and `@SideEffectTool` (`TimestampTools`) — plus a `MessageChatMemoryAdvisor` for conversation history. - [**MCP**](/springai/mcp): Connects to a Model Context Protocol server and exposes its tools to the AI through Temporal activities. Defaults to the filesystem MCP server. - [**Multi-Model**](/springai/multimodel): Two providers in one workflow (OpenAI and Anthropic), per-model `ActivityOptions` overrides via a Spring bean, plus a route that exercises Anthropic's extended-thinking mode through provider-specific `ChatOptions` pass-through. Requires `ANTHROPIC_API_KEY` in addition to `OPENAI_API_KEY`. - [**RAG**](/springai/rag): Vector store + embeddings for retrieval-augmented generation. Add documents, then ask questions; the workflow searches the vector store and grounds the answer in the retrieved context. diff --git a/gradle/springai.gradle b/gradle/springai.gradle index 394c8807..c10572b4 100644 --- a/gradle/springai.gradle +++ b/gradle/springai.gradle @@ -3,8 +3,8 @@ // // Note on Spring Boot version skew: the root build.gradle pins the // org.springframework.boot Gradle plugin at $springBootPluginVersion (currently -// 2.7.13) for the legacy springboot/ samples. Spring AI 1.1.0 requires Spring -// Boot 3.5.x, which we get by importing the spring-boot-dependencies BOM at +// 2.7.13) for the legacy springboot/ samples. Spring AI 2 requires Spring +// Boot 4, which we get by importing the spring-boot-dependencies BOM at // $springBootVersionForSpringAi below. The plugin and the BOM are independent — // the plugin contributes bootJar/bootRun task wiring, the BOM dictates // dependency versions — so this works in practice even though the two version @@ -17,8 +17,9 @@ apply plugin: 'org.springframework.boot' apply plugin: 'io.spring.dependency-management' ext { - springBootVersionForSpringAi = '3.5.3' - springAiVersion = '1.1.0' + springBootVersionForSpringAi = '4.0.8' + springAiVersion = '2.0.1' + temporalSpringAiVersion = '0.1.0' } java { @@ -35,13 +36,17 @@ dependencyManagement { dependencies { implementation "io.temporal:temporal-spring-boot-starter:$javaSDKVersion" - implementation "io.temporal:temporal-spring-ai:$javaSDKVersion" - // temporal-spring-ai declares temporal-sdk as compileOnly, so bring it in explicitly. + implementation "io.temporal:spring-ai:$temporalSpringAiVersion" + // spring-ai declares temporal-sdk as compileOnly, so bring it in explicitly. implementation "io.temporal:temporal-sdk:$javaSDKVersion" // Spring Boot implementation 'org.springframework.boot:spring-boot-starter' + testImplementation "io.temporal:temporal-testing:$javaSDKVersion" + testImplementation 'org.junit.jupiter:junit-jupiter' + testRuntimeOnly 'org.junit.platform:junit-platform-launcher' + dependencies { errorproneJavac('com.google.errorprone:javac:9+181-r4173-1') errorprone('com.google.errorprone:error_prone_core:2.28.0') diff --git a/springai/basic/src/main/java/io/temporal/samples/springai/chat/ChatWorkflowImpl.java b/springai/basic/src/main/java/io/temporal/samples/springai/chat/ChatWorkflowImpl.java index 887576c6..f707063c 100644 --- a/springai/basic/src/main/java/io/temporal/samples/springai/chat/ChatWorkflowImpl.java +++ b/springai/basic/src/main/java/io/temporal/samples/springai/chat/ChatWorkflowImpl.java @@ -8,7 +8,7 @@ import io.temporal.workflow.WorkflowInit; import java.time.Duration; import org.springframework.ai.chat.client.ChatClient; -import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor; +import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor; import org.springframework.ai.chat.memory.ChatMemory; import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository; import org.springframework.ai.chat.memory.MessageWindowChatMemory; @@ -78,12 +78,16 @@ public ChatWorkflowImpl(String systemPrompt) { // - Activity stubs (weatherTool) become durable AI tools // - plain workflow tool classes (stringTools) execute directly in workflow // - @SideEffectTool classes (timestampTools) are wrapped in sideEffect() - // - PromptChatMemoryAdvisor maintains conversation history + // - MessageChatMemoryAdvisor maintains conversation history this.chatClient = TemporalChatClient.builder(activityChatModel) .defaultSystem(systemPrompt) .defaultTools(weatherTool, stringTools, timestampTools) - .defaultAdvisors(PromptChatMemoryAdvisor.builder(chatMemory).build()) + .defaultAdvisors( + advisors -> + advisors + .advisors(MessageChatMemoryAdvisor.builder(chatMemory).build()) + .param(ChatMemory.CONVERSATION_ID, Workflow.getInfo().getWorkflowId())) .build(); } diff --git a/springai/basic/src/main/resources/application.yaml b/springai/basic/src/main/resources/application.yaml index 93e1597e..11e86a6c 100644 --- a/springai/basic/src/main/resources/application.yaml +++ b/springai/basic/src/main/resources/application.yaml @@ -7,6 +7,7 @@ spring: openai: api-key: ${OPENAI_API_KEY} chat: + max-retries: 0 options: model: gpt-4o-mini temperature: 0.7 diff --git a/springai/basic/src/test/java/io/temporal/samples/springai/chat/ChatWorkflowTest.java b/springai/basic/src/test/java/io/temporal/samples/springai/chat/ChatWorkflowTest.java new file mode 100644 index 00000000..38d600c0 --- /dev/null +++ b/springai/basic/src/test/java/io/temporal/samples/springai/chat/ChatWorkflowTest.java @@ -0,0 +1,89 @@ +package io.temporal.samples.springai.chat; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.temporal.client.WorkflowClient; +import io.temporal.client.WorkflowOptions; +import io.temporal.client.WorkflowStub; +import io.temporal.springai.activity.ChatModelActivityImpl; +import io.temporal.testing.TestWorkflowEnvironment; +import io.temporal.testing.WorkflowReplayer; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.ToolResponseMessage; +import org.springframework.ai.chat.model.ChatModel; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.model.Generation; +import org.springframework.ai.chat.prompt.Prompt; + +class ChatWorkflowTest { + + @Test + void toolsAndConversationMemorySurviveReplay() throws Exception { + List requests = new CopyOnWriteArrayList<>(); + ChatModel model = + prompt -> { + requests.add(prompt); + AssistantMessage response; + if (requests.size() == 1) { + response = + AssistantMessage.builder() + .toolCalls( + List.of( + new AssistantMessage.ToolCall( + "weather-1", "function", "getWeather", "{\"city\":\"Seattle\"}"))) + .build(); + } else { + response = + new AssistantMessage( + requests.size() == 2 ? "Seattle is rainy." : "You asked about Seattle."); + } + return new ChatResponse(List.of(new Generation(response))); + }; + + try (TestWorkflowEnvironment env = TestWorkflowEnvironment.newInstance()) { + var worker = env.newWorker("spring-ai-chat-test"); + worker.registerWorkflowImplementationTypes(ChatWorkflowImpl.class); + worker.registerActivitiesImplementations( + new ChatModelActivityImpl(model), new WeatherActivityImpl()); + env.start(); + + ChatWorkflow workflow = + env.getWorkflowClient() + .newWorkflowStub( + ChatWorkflow.class, + WorkflowOptions.newBuilder() + .setTaskQueue("spring-ai-chat-test") + .setWorkflowExecutionTimeout(Duration.ofSeconds(20)) + .build()); + WorkflowClient.start(workflow::run, "Answer weather questions."); + assertEquals("Seattle is rainy.", workflow.chat("What's the weather in Seattle?")); + assertEquals("You asked about Seattle.", workflow.chat("Which city did I ask about?")); + workflow.end(); + WorkflowStub stub = WorkflowStub.fromTyped(workflow); + assertEquals("Chat ended after 2 messages.", stub.getResult(String.class)); + + assertEquals(3, requests.size()); + assertTrue( + requests.get(1).getInstructions().stream() + .filter(ToolResponseMessage.class::isInstance) + .map(ToolResponseMessage.class::cast) + .flatMap(message -> message.getResponses().stream()) + .anyMatch(response -> response.responseData().contains("Seattle"))); + assertTrue( + requests.get(2).getInstructions().stream() + .anyMatch(message -> "What's the weather in Seattle?".equals(message.getText()))); + assertTrue( + requests.get(2).getInstructions().stream() + .anyMatch(message -> "Seattle is rainy.".equals(message.getText()))); + + WorkflowReplayer.replayWorkflowExecution( + env.getWorkflowClient().fetchHistory(stub.getExecution().getWorkflowId()), + ChatWorkflowImpl.class); + } + } +} diff --git a/springai/mcp/src/main/java/io/temporal/samples/springai/mcp/McpWorkflowImpl.java b/springai/mcp/src/main/java/io/temporal/samples/springai/mcp/McpWorkflowImpl.java index c9bcaa99..0000a100 100644 --- a/springai/mcp/src/main/java/io/temporal/samples/springai/mcp/McpWorkflowImpl.java +++ b/springai/mcp/src/main/java/io/temporal/samples/springai/mcp/McpWorkflowImpl.java @@ -8,7 +8,7 @@ import java.util.List; import java.util.stream.Collectors; import org.springframework.ai.chat.client.ChatClient; -import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor; +import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor; import org.springframework.ai.chat.memory.ChatMemory; import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository; import org.springframework.ai.chat.memory.MessageWindowChatMemory; @@ -74,8 +74,12 @@ public String run() { When asked to perform file operations, use your tools. Always confirm what you did after completing an operation. """) - .defaultToolCallbacks(mcpTools) - .defaultAdvisors(PromptChatMemoryAdvisor.builder(chatMemory).build()) + .defaultTools(mcpTools.toArray()) + .defaultAdvisors( + advisors -> + advisors + .advisors(MessageChatMemoryAdvisor.builder(chatMemory).build()) + .param(ChatMemory.CONVERSATION_ID, Workflow.getInfo().getWorkflowId())) .build(); initialized = true; @@ -95,7 +99,7 @@ public void chat(String message) { messageCount++; - // PromptChatMemoryAdvisor automatically handles conversation history + // MessageChatMemoryAdvisor automatically handles conversation history lastResponse = chatClient.prompt().user(message).call().content(); } diff --git a/springai/mcp/src/main/resources/application.yaml b/springai/mcp/src/main/resources/application.yaml index 53551db6..316db424 100644 --- a/springai/mcp/src/main/resources/application.yaml +++ b/springai/mcp/src/main/resources/application.yaml @@ -6,6 +6,7 @@ spring: openai: api-key: ${OPENAI_API_KEY} chat: + max-retries: 0 options: model: gpt-4o-mini mcp: diff --git a/springai/multimodel/src/main/java/io/temporal/samples/springai/multimodel/ChatModelConfig.java b/springai/multimodel/src/main/java/io/temporal/samples/springai/multimodel/ChatModelConfig.java index 6bc4c1eb..cf6999c0 100644 --- a/springai/multimodel/src/main/java/io/temporal/samples/springai/multimodel/ChatModelConfig.java +++ b/springai/multimodel/src/main/java/io/temporal/samples/springai/multimodel/ChatModelConfig.java @@ -7,11 +7,9 @@ import java.util.Map; import org.springframework.ai.anthropic.AnthropicChatModel; import org.springframework.ai.anthropic.AnthropicChatOptions; -import org.springframework.ai.anthropic.api.AnthropicApi; import org.springframework.ai.chat.model.ChatModel; import org.springframework.ai.openai.OpenAiChatModel; import org.springframework.ai.openai.OpenAiChatOptions; -import org.springframework.ai.openai.api.OpenAiApi; import org.springframework.beans.factory.annotation.Value; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; @@ -47,22 +45,27 @@ public class ChatModelConfig { @Bean @Primary public ChatModel openAiChatModel() { - OpenAiApi api = OpenAiApi.builder().apiKey(openAiApiKey).build(); OpenAiChatOptions options = - OpenAiChatOptions.builder().model("gpt-4o-mini").temperature(0.7).build(); - return OpenAiChatModel.builder().openAiApi(api).defaultOptions(options).build(); + OpenAiChatOptions.builder() + .apiKey(openAiApiKey) + .maxRetries(0) // Temporal Activities own retries. + .model("gpt-4o-mini") + .temperature(0.7) + .build(); + return OpenAiChatModel.builder().options(options).build(); } /** Anthropic model using Claude for complex reasoning tasks. */ @Bean public ChatModel anthropicChatModel() { - AnthropicApi api = AnthropicApi.builder().apiKey(anthropicApiKey).build(); AnthropicChatOptions options = AnthropicChatOptions.builder() + .apiKey(anthropicApiKey) + .maxRetries(0) .model("claude-sonnet-4-20250514") .temperature(0.3) // Lower temperature for more focused reasoning .build(); - return AnthropicChatModel.builder().anthropicApi(api).defaultOptions(options).build(); + return AnthropicChatModel.builder().options(options).build(); } /** diff --git a/springai/multimodel/src/main/java/io/temporal/samples/springai/multimodel/MultiModelWorkflowImpl.java b/springai/multimodel/src/main/java/io/temporal/samples/springai/multimodel/MultiModelWorkflowImpl.java index 3a1e55c6..569fa7dc 100644 --- a/springai/multimodel/src/main/java/io/temporal/samples/springai/multimodel/MultiModelWorkflowImpl.java +++ b/springai/multimodel/src/main/java/io/temporal/samples/springai/multimodel/MultiModelWorkflowImpl.java @@ -7,7 +7,6 @@ import java.util.LinkedHashMap; import java.util.Map; import org.springframework.ai.anthropic.AnthropicChatOptions; -import org.springframework.ai.anthropic.api.AnthropicApi; import org.springframework.ai.chat.client.ChatClient; /** @@ -19,7 +18,7 @@ *
  • openai - Uses OpenAI gpt-4o-mini for quick, cost-effective responses *
  • anthropic - Uses Anthropic Claude for complex reasoning tasks *
  • think - Uses Anthropic Claude with extended thinking enabled for hard - * problems. This demonstrates {@code temporal-spring-ai}'s provider-specific ChatOptions + * problems. This demonstrates {@code io.temporal:spring-ai}'s provider-specific ChatOptions * pass-through: an {@link AnthropicChatOptions} instance with a {@code thinking} * configuration is attached per call, and every field survives the round-trip across the * activity boundary. @@ -97,12 +96,8 @@ public MultiModelWorkflowImpl() { // sent to Claude. Required side effects for extended thinking: temperature must be 1.0 // and max_tokens must exceed budget_tokens. // @@@SNIPSTART samples-java-spring-ai-provider-options - AnthropicChatOptions thinkingOptions = - AnthropicChatOptions.builder() - .thinking(AnthropicApi.ThinkingType.ENABLED, 1024) - .temperature(1.0) - .maxTokens(4096) - .build(); + AnthropicChatOptions.Builder thinkingOptions = + AnthropicChatOptions.builder().thinkingEnabled(1024).temperature(1.0).maxTokens(4096); chatClients.put( "think", TemporalChatClient.builder(anthropicModel) diff --git a/springai/multimodel/src/test/java/io/temporal/samples/springai/multimodel/ChatModelConfigTest.java b/springai/multimodel/src/test/java/io/temporal/samples/springai/multimodel/ChatModelConfigTest.java new file mode 100644 index 00000000..8099223a --- /dev/null +++ b/springai/multimodel/src/test/java/io/temporal/samples/springai/multimodel/ChatModelConfigTest.java @@ -0,0 +1,40 @@ +package io.temporal.samples.springai.multimodel; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertSame; + +import java.util.Map; +import org.junit.jupiter.api.Test; +import org.springframework.ai.anthropic.AnthropicChatModel; +import org.springframework.ai.chat.model.ChatModel; +import org.springframework.ai.openai.OpenAiChatModel; +import org.springframework.context.annotation.AnnotationConfigApplicationContext; +import org.springframework.core.env.MapPropertySource; + +class ChatModelConfigTest { + + @Test + void providerBeansInitializeWithoutNetworkCalls() { + try (var context = new AnnotationConfigApplicationContext()) { + context + .getEnvironment() + .getPropertySources() + .addFirst( + new MapPropertySource( + "test", + Map.of( + "spring.ai.openai.api-key", "test-openai-key", + "spring.ai.anthropic.api-key", "test-anthropic-key"))); + context.register(ChatModelConfig.class); + context.refresh(); + + var openAi = assertInstanceOf(OpenAiChatModel.class, context.getBean("openAiChatModel")); + var anthropic = + assertInstanceOf(AnthropicChatModel.class, context.getBean("anthropicChatModel")); + assertSame(openAi, context.getBean(ChatModel.class)); + assertEquals(0, openAi.getOptions().getMaxRetries()); + assertEquals(0, anthropic.getOptions().getMaxRetries()); + } + } +} diff --git a/springai/rag/src/main/resources/application.yaml b/springai/rag/src/main/resources/application.yaml index b8889cb7..c67c3ed2 100644 --- a/springai/rag/src/main/resources/application.yaml +++ b/springai/rag/src/main/resources/application.yaml @@ -6,6 +6,7 @@ spring: openai: api-key: ${OPENAI_API_KEY} chat: + max-retries: 0 options: model: gpt-4o-mini embedding: From 409a90af4dd214ecea06df91176072be5dd4be66 Mon Sep 17 00:00:00 2001 From: Brian Strauch Date: Fri, 16 Oct 2026 10:06:03 -0700 Subject: [PATCH 2/2] Address Spring AI migration review --- README.md | 2 +- gradle/springai.gradle | 2 +- .../samples/springai/mcp/McpWorkflowTest.java | 154 ++++++++++++++++++ .../rag/src/main/resources/application.yaml | 2 +- 4 files changed, 157 insertions(+), 3 deletions(-) create mode 100644 springai/mcp/src/test/java/io/temporal/samples/springai/mcp/McpWorkflowTest.java diff --git a/README.md b/README.md index a3245bcb..b70a9098 100644 --- a/README.md +++ b/README.md @@ -232,7 +232,7 @@ To run any of the SpringBoot samples in your Temporal Cloud namespace: The Spring AI samples demonstrate the [Temporal Spring AI integration](https://github.com/temporalio/ai-integrations/tree/main/java/spring-ai), which makes Spring AI agents durable on Temporal — model calls run as Temporal Activities recorded in Workflow history, and tools are dispatched per their type so they fit Workflow execution. -These samples use the Public Preview `io.temporal:spring-ai:0.1.0` package with Spring AI 2.0.1, Spring Boot 4.0.8, and Temporal Java SDK 1.40.0. Java 17 or later is required. Spring AI 1 / Spring Boot 3 applications should remain on `io.temporal:temporal-spring-ai:1.40.x`. Spring AI 1 Workflow histories cannot replay on the new integration; complete existing executions or use separate Worker deployments when upgrading. +These samples use the Public Preview `io.temporal:spring-ai:0.1.1` package with Spring AI 2.0.1, Spring Boot 4.0.8, and Temporal Java SDK 1.40.0. Java 17 or later is required. Spring AI 1 / Spring Boot 3 applications should remain on `io.temporal:temporal-spring-ai:1.40.x`. Spring AI 1 Workflow histories cannot replay on the new integration; complete existing executions or use separate Worker deployments when upgrading. Each sample is its own Spring Boot application with an interactive CLI. Run from the main repo dir: diff --git a/gradle/springai.gradle b/gradle/springai.gradle index c10572b4..cb0c8907 100644 --- a/gradle/springai.gradle +++ b/gradle/springai.gradle @@ -19,7 +19,7 @@ apply plugin: 'io.spring.dependency-management' ext { springBootVersionForSpringAi = '4.0.8' springAiVersion = '2.0.1' - temporalSpringAiVersion = '0.1.0' + temporalSpringAiVersion = '0.1.1' } java { diff --git a/springai/mcp/src/test/java/io/temporal/samples/springai/mcp/McpWorkflowTest.java b/springai/mcp/src/test/java/io/temporal/samples/springai/mcp/McpWorkflowTest.java new file mode 100644 index 00000000..e07d680c --- /dev/null +++ b/springai/mcp/src/test/java/io/temporal/samples/springai/mcp/McpWorkflowTest.java @@ -0,0 +1,154 @@ +package io.temporal.samples.springai.mcp; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.modelcontextprotocol.spec.McpSchema; +import io.temporal.client.WorkflowClient; +import io.temporal.client.WorkflowOptions; +import io.temporal.client.WorkflowStub; +import io.temporal.springai.activity.ChatModelActivityImpl; +import io.temporal.springai.mcp.McpClientActivity; +import io.temporal.testing.TestWorkflowEnvironment; +import io.temporal.testing.WorkflowReplayer; +import java.time.Duration; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CopyOnWriteArrayList; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.ToolResponseMessage; +import org.springframework.ai.chat.model.ChatModel; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.model.Generation; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.mcp.McpToolUtils; +import org.springframework.ai.model.tool.ToolCallingChatOptions; + +class McpWorkflowTest { + + @Test + void discoveredToolsAndConversationMemorySurviveReplay() throws Exception { + String toolName = McpToolUtils.prefixedToolName("filesystem", "read_file"); + List requests = new CopyOnWriteArrayList<>(); + ChatModel model = + prompt -> { + requests.add(prompt); + AssistantMessage response; + if (requests.size() == 1) { + response = + AssistantMessage.builder() + .toolCalls( + List.of( + new AssistantMessage.ToolCall( + "read-1", "function", toolName, "{\"path\":\"/tmp/notes.txt\"}"))) + .build(); + } else { + response = + new AssistantMessage( + requests.size() == 2 ? "The file says hello." : "You asked about notes.txt."); + } + return new ChatResponse(List.of(new Generation(response))); + }; + FakeMcpClientActivity mcp = new FakeMcpClientActivity(); + + try (TestWorkflowEnvironment env = TestWorkflowEnvironment.newInstance()) { + var worker = env.newWorker("spring-ai-mcp-test"); + worker.registerWorkflowImplementationTypes(McpWorkflowImpl.class); + worker.registerActivitiesImplementations(new ChatModelActivityImpl(model), mcp); + env.start(); + + McpWorkflow workflow = + env.getWorkflowClient() + .newWorkflowStub( + McpWorkflow.class, + WorkflowOptions.newBuilder() + .setTaskQueue("spring-ai-mcp-test") + .setWorkflowExecutionTimeout(Duration.ofSeconds(30)) + .build()); + WorkflowClient.start(workflow::run); + // The signal may arrive while tool discovery is still initializing. + workflow.chat("Read /tmp/notes.txt."); + awaitResponse(workflow, "The file says hello."); + assertEquals("Available MCP tools:\n - " + toolName + ": Read a file", workflow.listTools()); + workflow.chat("Which file did I ask about?"); + awaitResponse(workflow, "You asked about notes.txt."); + workflow.end(); + WorkflowStub stub = WorkflowStub.fromTyped(workflow); + assertEquals("Chat ended after 2 messages.", stub.getResult(String.class)); + + assertEquals(3, requests.size()); + var options = assertInstanceOf(ToolCallingChatOptions.class, requests.get(0).getOptions()); + assertEquals(1, options.getToolCallbacks().size()); + assertEquals(toolName, options.getToolCallbacks().get(0).getToolDefinition().name()); + assertEquals(List.of("filesystem"), mcp.toolClients); + assertEquals(1, mcp.toolRequests.size()); + assertEquals("read_file", mcp.toolRequests.get(0).name()); + assertEquals(Map.of("path", "/tmp/notes.txt"), mcp.toolRequests.get(0).arguments()); + assertTrue( + requests.get(1).getInstructions().stream() + .filter(ToolResponseMessage.class::isInstance) + .map(ToolResponseMessage.class::cast) + .flatMap(message -> message.getResponses().stream()) + .anyMatch(response -> response.responseData().contains("hello"))); + assertTrue( + requests.get(2).getInstructions().stream() + .anyMatch(message -> "Read /tmp/notes.txt.".equals(message.getText()))); + assertTrue( + requests.get(2).getInstructions().stream() + .anyMatch(message -> "The file says hello.".equals(message.getText()))); + + WorkflowReplayer.replayWorkflowExecution( + env.getWorkflowClient().fetchHistory(stub.getExecution().getWorkflowId()), + McpWorkflowImpl.class); + assertEquals(3, requests.size()); + assertEquals(1, mcp.toolRequests.size()); + } + } + + private static void awaitResponse(McpWorkflow workflow, String expected) + throws InterruptedException { + long deadline = System.nanoTime() + Duration.ofSeconds(10).toNanos(); + while (!expected.equals(workflow.getLastResponse()) && System.nanoTime() < deadline) { + Thread.sleep(10); + } + assertEquals(expected, workflow.getLastResponse()); + } + + public static class FakeMcpClientActivity implements McpClientActivity { + final List toolClients = new CopyOnWriteArrayList<>(); + final List toolRequests = new CopyOnWriteArrayList<>(); + + @Override + public Map getServerCapabilities() { + return Map.of(); + } + + @Override + public Map getClientInfo() { + return Map.of("filesystem", McpSchema.Implementation.builder("filesystem", "1.0").build()); + } + + @Override + public Map listTools() { + McpSchema.Tool tool = + McpSchema.Tool.builder( + "read_file", + Map.of( + "type", "object", + "properties", Map.of("path", Map.of("type", "string")), + "required", List.of("path"))) + .description("Read a file") + .build(); + return Map.of("filesystem", McpSchema.ListToolsResult.builder(List.of(tool)).build()); + } + + @Override + public McpSchema.CallToolResult callTool(String clientName, McpSchema.CallToolRequest request) { + toolClients.add(clientName); + toolRequests.add(request); + return McpSchema.CallToolResult.builder().addTextContent("hello").isError(false).build(); + } + } +} diff --git a/springai/rag/src/main/resources/application.yaml b/springai/rag/src/main/resources/application.yaml index c67c3ed2..583aee95 100644 --- a/springai/rag/src/main/resources/application.yaml +++ b/springai/rag/src/main/resources/application.yaml @@ -5,8 +5,8 @@ spring: ai: openai: api-key: ${OPENAI_API_KEY} + max-retries: 0 chat: - max-retries: 0 options: model: gpt-4o-mini embedding: