diff --git a/temporal-spring-ai/src/main/java/io/temporal/springai/chat/TemporalChatClient.java b/temporal-spring-ai/src/main/java/io/temporal/springai/chat/TemporalChatClient.java index 94517502e..6a3051db7 100644 --- a/temporal-spring-ai/src/main/java/io/temporal/springai/chat/TemporalChatClient.java +++ b/temporal-spring-ai/src/main/java/io/temporal/springai/chat/TemporalChatClient.java @@ -3,12 +3,12 @@ import io.micrometer.observation.ObservationRegistry; import io.temporal.springai.util.TemporalToolUtil; import java.util.Map; +import javax.annotation.Nullable; import org.springframework.ai.chat.client.ChatClient; import org.springframework.ai.chat.client.DefaultChatClient; import org.springframework.ai.chat.client.DefaultChatClientBuilder; import org.springframework.ai.chat.client.observation.ChatClientObservationConvention; import org.springframework.ai.chat.model.ChatModel; -import org.springframework.lang.Nullable; import org.springframework.util.Assert; /** diff --git a/temporal-spring-ai/src/main/java/io/temporal/springai/mcp/ActivityMcpClient.java b/temporal-spring-ai/src/main/java/io/temporal/springai/mcp/ActivityMcpClient.java index 360412a83..0df9f99b4 100644 --- a/temporal-spring-ai/src/main/java/io/temporal/springai/mcp/ActivityMcpClient.java +++ b/temporal-spring-ai/src/main/java/io/temporal/springai/mcp/ActivityMcpClient.java @@ -6,6 +6,7 @@ import io.temporal.workflow.Workflow; import java.time.Duration; import java.util.Map; +import javax.annotation.Nullable; /** * A workflow-safe wrapper for MCP (Model Context Protocol) client operations. @@ -48,6 +49,7 @@ public class ActivityMcpClient { public static final int DEFAULT_MAX_ATTEMPTS = 3; private final McpClientActivity activity; + @Nullable private final ActivityOptions baseOptions; private Map serverCapabilities; private Map clientInfo; @@ -57,7 +59,18 @@ public class ActivityMcpClient { * @param activity the activity stub for MCP operations */ public ActivityMcpClient(McpClientActivity activity) { + this(activity, null); + } + + /** + * Creates a new ActivityMcpClient. When {@code baseOptions} is non-null, {@link #callTool(String, + * McpSchema.CallToolRequest, String)} rebuilds the activity stub with a per-call Summary on top + * of those options. When null, the caller supplied a pre-built stub whose options we don't know, + * so we call through it as-is and drop any requested summary. + */ + private ActivityMcpClient(McpClientActivity activity, @Nullable ActivityOptions baseOptions) { this.activity = activity; + this.baseOptions = baseOptions; } /** @@ -81,14 +94,13 @@ public static ActivityMcpClient create() { * @return a new ActivityMcpClient */ public static ActivityMcpClient create(Duration timeout, int maxAttempts) { - McpClientActivity activity = - Workflow.newActivityStub( - McpClientActivity.class, - ActivityOptions.newBuilder() - .setStartToCloseTimeout(timeout) - .setRetryOptions(RetryOptions.newBuilder().setMaximumAttempts(maxAttempts).build()) - .build()); - return new ActivityMcpClient(activity); + ActivityOptions options = + ActivityOptions.newBuilder() + .setStartToCloseTimeout(timeout) + .setRetryOptions(RetryOptions.newBuilder().setMaximumAttempts(maxAttempts).build()) + .build(); + McpClientActivity activity = Workflow.newActivityStub(McpClientActivity.class, options); + return new ActivityMcpClient(activity, options); } /** @@ -127,6 +139,32 @@ public Map getClientInfo() { * @return the tool call result */ public McpSchema.CallToolResult callTool(String clientName, McpSchema.CallToolRequest request) { + return callTool(clientName, request, null); + } + + /** + * Calls a tool on a specific MCP client, attaching the given activity Summary to the scheduled + * activity so it renders meaningfully in the Temporal UI. Falls back to the base stub when no + * {@link ActivityOptions} are known (e.g. when this client was constructed from a user-supplied + * stub rather than one of the {@link #create} factories). + * + * @param clientName the name of the MCP client + * @param request the tool call request + * @param summary the activity Summary, or null to omit + * @return the tool call result + */ + public McpSchema.CallToolResult callTool( + String clientName, McpSchema.CallToolRequest request, @Nullable String summary) { + // Overlay the summary onto a fresh stub only when both a summary is requested AND we have + // a recipe to rebuild the stub from (baseOptions). If either is missing, fall through to + // the cached activity — it already has baseOptions baked in if we knew them at construction. + if (summary != null && baseOptions != null) { + McpClientActivity stub = + Workflow.newActivityStub( + McpClientActivity.class, + ActivityOptions.newBuilder(baseOptions).setSummary(summary).build()); + return stub.callTool(clientName, request); + } return activity.callTool(clientName, request); } diff --git a/temporal-spring-ai/src/main/java/io/temporal/springai/mcp/McpToolCallback.java b/temporal-spring-ai/src/main/java/io/temporal/springai/mcp/McpToolCallback.java index 9cf821aae..48736af6a 100644 --- a/temporal-spring-ai/src/main/java/io/temporal/springai/mcp/McpToolCallback.java +++ b/temporal-spring-ai/src/main/java/io/temporal/springai/mcp/McpToolCallback.java @@ -106,7 +106,8 @@ public String call(String toolInput) { // Use the original tool name (not prefixed) when calling the MCP server McpSchema.CallToolRequest request = new McpSchema.CallToolRequest(tool.name(), arguments); - McpSchema.CallToolResult result = client.callTool(clientName, request); + String summary = "mcp: " + clientName + "." + tool.name(); + McpSchema.CallToolResult result = client.callTool(clientName, request, summary); // Return the result as-is (including errors) so the AI can handle them. // For example, an "access denied" error lets the AI suggest a valid path. diff --git a/temporal-spring-ai/src/main/java/io/temporal/springai/model/ActivityChatModel.java b/temporal-spring-ai/src/main/java/io/temporal/springai/model/ActivityChatModel.java index c446db1d8..284d471c1 100644 --- a/temporal-spring-ai/src/main/java/io/temporal/springai/model/ActivityChatModel.java +++ b/temporal-spring-ai/src/main/java/io/temporal/springai/model/ActivityChatModel.java @@ -10,6 +10,7 @@ import java.util.List; import java.util.Map; import java.util.stream.Collectors; +import javax.annotation.Nullable; import org.springframework.ai.chat.messages.*; import org.springframework.ai.chat.metadata.ChatResponseMetadata; import org.springframework.ai.chat.model.ChatModel; @@ -84,7 +85,8 @@ public class ActivityChatModel implements ChatModel { public static final int DEFAULT_MAX_ATTEMPTS = 3; private final ChatModelActivity chatModelActivity; - private final String modelName; + @Nullable private final String modelName; + @Nullable private final ActivityOptions baseOptions; private final ToolCallingManager toolCallingManager; private final ToolExecutionEligibilityPredicate toolExecutionEligibilityPredicate; @@ -94,7 +96,7 @@ public class ActivityChatModel implements ChatModel { * @param chatModelActivity the activity stub for calling the chat model */ public ActivityChatModel(ChatModelActivity chatModelActivity) { - this(chatModelActivity, null); + this(chatModelActivity, null, null); } /** @@ -103,9 +105,24 @@ public ActivityChatModel(ChatModelActivity chatModelActivity) { * @param chatModelActivity the activity stub for calling the chat model * @param modelName the name of the chat model to use, or null for default */ - public ActivityChatModel(ChatModelActivity chatModelActivity, String modelName) { + public ActivityChatModel(ChatModelActivity chatModelActivity, @Nullable String modelName) { + this(chatModelActivity, modelName, null); + } + + /** + * Internal constructor used by {@link #forModel(String, Duration, int)} and friends. When {@code + * baseOptions} is non-null, each call rebuilds the activity stub with a per-call Summary on top + * of those options so the Temporal UI can label the chat activity meaningfully. When null, the + * caller supplied a pre-built stub whose options we don't know, so we call through it as-is + * without a summary. + */ + private ActivityChatModel( + ChatModelActivity chatModelActivity, + @Nullable String modelName, + @Nullable ActivityOptions baseOptions) { this.chatModelActivity = chatModelActivity; this.modelName = modelName; + this.baseOptions = baseOptions; this.toolCallingManager = ToolCallingManager.builder().build(); this.toolExecutionEligibilityPredicate = new DefaultToolExecutionEligibilityPredicate(); } @@ -150,22 +167,22 @@ public static ActivityChatModel forModel(String modelName) { * @param maxAttempts the maximum number of retry attempts * @return an ActivityChatModel for the specified chat model */ - public static ActivityChatModel forModel(String modelName, Duration timeout, int maxAttempts) { - ChatModelActivity activity = - Workflow.newActivityStub( - ChatModelActivity.class, - ActivityOptions.newBuilder() - .setStartToCloseTimeout(timeout) - .setRetryOptions(RetryOptions.newBuilder().setMaximumAttempts(maxAttempts).build()) - .build()); - return new ActivityChatModel(activity, modelName); + public static ActivityChatModel forModel( + @Nullable String modelName, Duration timeout, int maxAttempts) { + ActivityOptions options = + ActivityOptions.newBuilder() + .setStartToCloseTimeout(timeout) + .setRetryOptions(RetryOptions.newBuilder().setMaximumAttempts(maxAttempts).build()) + .build(); + ChatModelActivity activity = Workflow.newActivityStub(ChatModelActivity.class, options); + return new ActivityChatModel(activity, modelName, options); } /** - * Returns the name of the chat model this instance uses. - * - * @return the model name, or null if using the default model + * Returns the name of the chat model this instance uses, or null if it uses the plugin default + * (the {@code @Primary} {@code ChatModel} bean or the first one registered). */ + @Nullable public String getModelName() { return modelName; } @@ -193,7 +210,8 @@ public ChatResponse call(Prompt prompt) { private ChatResponse internalCall(Prompt prompt) { // Convert prompt to activity input and call the activity ChatModelTypes.ChatModelActivityInput input = createActivityInput(prompt); - ChatModelTypes.ChatModelActivityOutput output = chatModelActivity.callChatModel(input); + ChatModelActivity stub = stubForCall(prompt); + ChatModelTypes.ChatModelActivityOutput output = stub.callChatModel(input); // Convert activity output to ChatResponse ChatResponse response = toResponse(output); @@ -219,6 +237,25 @@ private ChatResponse internalCall(Prompt prompt) { return response; } + private ChatModelActivity stubForCall(Prompt prompt) { + if (baseOptions == null) { + return chatModelActivity; + } + ActivityOptions withSummary = + ActivityOptions.newBuilder(baseOptions).setSummary(buildSummary()).build(); + return Workflow.newActivityStub(ChatModelActivity.class, withSummary); + } + + /** + * Builds the activity Summary. Intentionally omits the user prompt — including even a truncated + * slice would leak whatever the prompt contains (PII, secrets, internal identifiers) into + * workflow history, server logs, and the Temporal UI, which is a surprising default for a plain + * observability label. + */ + private String buildSummary() { + return "chat: " + (modelName != null ? modelName : "default"); + } + private ChatModelTypes.ChatModelActivityInput createActivityInput(Prompt prompt) { // Convert messages List messages = diff --git a/temporal-spring-ai/src/main/java/io/temporal/springai/model/ChatModelTypes.java b/temporal-spring-ai/src/main/java/io/temporal/springai/model/ChatModelTypes.java index 7dca5e316..c1f57317d 100644 --- a/temporal-spring-ai/src/main/java/io/temporal/springai/model/ChatModelTypes.java +++ b/temporal-spring-ai/src/main/java/io/temporal/springai/model/ChatModelTypes.java @@ -6,6 +6,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import java.time.Duration; import java.util.List; +import javax.annotation.Nullable; /** * Serializable types for chat model activity requests and responses. @@ -20,21 +21,23 @@ private ChatModelTypes() {} /** * Input to the chat model activity. * - * @param modelName the name of the chat model bean to use (null for default) + * @param modelName the name of the chat model bean to use, or null for the activity-side default + * model * @param messages the conversation messages - * @param modelOptions options for the chat model (temperature, max tokens, etc.) + * @param modelOptions options for the chat model (temperature, max tokens, etc.), or null to use + * the chat model's own defaults * @param tools tool definitions the model may call */ @JsonInclude(JsonInclude.Include.NON_NULL) @JsonIgnoreProperties(ignoreUnknown = true) public record ChatModelActivityInput( - @JsonProperty("model_name") String modelName, + @JsonProperty("model_name") @Nullable String modelName, @JsonProperty("messages") List messages, - @JsonProperty("model_options") ModelOptions modelOptions, + @JsonProperty("model_options") @Nullable ModelOptions modelOptions, @JsonProperty("tools") List tools) { /** Creates input for the default chat model. */ public ChatModelActivityInput( - List messages, ModelOptions modelOptions, List tools) { + List messages, @Nullable ModelOptions modelOptions, List tools) { this(null, messages, modelOptions, tools); } } diff --git a/temporal-spring-ai/src/main/java/io/temporal/springai/plugin/SpringAiPlugin.java b/temporal-spring-ai/src/main/java/io/temporal/springai/plugin/SpringAiPlugin.java index 0438ff558..ad13d9abb 100644 --- a/temporal-spring-ai/src/main/java/io/temporal/springai/plugin/SpringAiPlugin.java +++ b/temporal-spring-ai/src/main/java/io/temporal/springai/plugin/SpringAiPlugin.java @@ -7,10 +7,10 @@ import java.util.LinkedHashMap; import java.util.Map; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.ai.chat.model.ChatModel; -import org.springframework.lang.Nullable; /** * Core Temporal plugin that registers {@link io.temporal.springai.activity.ChatModelActivity} with diff --git a/temporal-spring-ai/src/test/java/io/temporal/springai/ActivitySummaryTest.java b/temporal-spring-ai/src/test/java/io/temporal/springai/ActivitySummaryTest.java new file mode 100644 index 000000000..bf0efebba --- /dev/null +++ b/temporal-spring-ai/src/test/java/io/temporal/springai/ActivitySummaryTest.java @@ -0,0 +1,136 @@ +package io.temporal.springai; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.temporal.api.history.v1.HistoryEvent; +import io.temporal.client.WorkflowClient; +import io.temporal.client.WorkflowOptions; +import io.temporal.client.WorkflowStub; +import io.temporal.common.converter.DefaultDataConverter; +import io.temporal.springai.activity.ChatModelActivityImpl; +import io.temporal.springai.chat.TemporalChatClient; +import io.temporal.springai.model.ActivityChatModel; +import io.temporal.testing.TestWorkflowEnvironment; +import io.temporal.worker.Worker; +import io.temporal.workflow.WorkflowInterface; +import io.temporal.workflow.WorkflowMethod; +import java.util.List; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.client.ChatClient; +import org.springframework.ai.chat.messages.AssistantMessage; +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; + +/** + * Verifies that activity calls scheduled by the Spring AI plugin carry a human-readable Summary in + * their {@code userMetadata} so the Temporal UI can label them meaningfully. + */ +class ActivitySummaryTest { + + private static final String TASK_QUEUE = "test-spring-ai-summary"; + + private TestWorkflowEnvironment testEnv; + private WorkflowClient client; + + @BeforeEach + void setUp() { + testEnv = TestWorkflowEnvironment.newInstance(); + client = testEnv.getWorkflowClient(); + } + + @AfterEach + void tearDown() { + testEnv.close(); + } + + @Test + void chatActivity_carriesModelOnlySummary_neverLeaksUserPrompt() { + Worker worker = testEnv.newWorker(TASK_QUEUE); + worker.registerWorkflowImplementationTypes(ChatWorkflowImpl.class); + worker.registerActivitiesImplementations( + new ChatModelActivityImpl(new StubChatModel("Hello!"))); + testEnv.start(); + + SummaryTestWorkflow workflow = + client.newWorkflowStub( + SummaryTestWorkflow.class, + WorkflowOptions.newBuilder().setTaskQueue(TASK_QUEUE).build()); + String prompt = "What is the capital of France?"; + workflow.chat(prompt); + + String workflowId = WorkflowStub.fromTyped(workflow).getExecution().getWorkflowId(); + List events = client.fetchHistory(workflowId).getHistory().getEventsList(); + + String summary = findChatActivitySummary(events); + assertNotNull(summary, "ActivityTaskScheduled event for callChatModel should have a Summary"); + assertEquals( + "chat: default", + summary, + "Summary must be just the model label — user prompts can contain PII and must not" + + " appear in history/logs/UI by default."); + // Defense in depth: explicitly confirm no part of the prompt leaked into the summary. + assertTrue( + !summary.contains(prompt) && !summary.contains("France"), + "Summary must not include user prompt content, got: " + summary); + } + + private static String findChatActivitySummary(List events) { + for (HistoryEvent event : events) { + if (!event.hasActivityTaskScheduledEventAttributes()) { + continue; + } + String activityType = + event.getActivityTaskScheduledEventAttributes().getActivityType().getName(); + if (!"CallChatModel".equals(activityType)) { + continue; + } + if (!event.hasUserMetadata() || !event.getUserMetadata().hasSummary()) { + return null; + } + return DefaultDataConverter.STANDARD_INSTANCE.fromPayload( + event.getUserMetadata().getSummary(), String.class, String.class); + } + return null; + } + + @WorkflowInterface + public interface SummaryTestWorkflow { + @WorkflowMethod + String chat(String message); + } + + public static class ChatWorkflowImpl implements SummaryTestWorkflow { + @Override + public String chat(String message) { + ActivityChatModel chatModel = ActivityChatModel.forDefault(); + ChatClient chatClient = TemporalChatClient.builder(chatModel).build(); + return chatClient.prompt().user(message).call().content(); + } + } + + private static class StubChatModel implements ChatModel { + private final String response; + + StubChatModel(String response) { + this.response = response; + } + + @Override + public ChatResponse call(Prompt prompt) { + return ChatResponse.builder() + .generations(List.of(new Generation(new AssistantMessage(response)))) + .build(); + } + + @Override + public reactor.core.publisher.Flux stream(Prompt prompt) { + throw new UnsupportedOperationException(); + } + } +}