Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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<String, McpSchema.ServerCapabilities> serverCapabilities;
private Map<String, McpSchema.Implementation> clientInfo;

Expand All @@ -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;
}

/**
Expand All @@ -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);
}

/**
Expand Down Expand Up @@ -127,6 +139,32 @@ public Map<String, McpSchema.Implementation> 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);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;

Expand All @@ -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);
}

/**
Expand All @@ -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();
}
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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);
Expand All @@ -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<ChatModelTypes.Message> messages =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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<Message> messages,
@JsonProperty("model_options") ModelOptions modelOptions,
@JsonProperty("model_options") @Nullable ModelOptions modelOptions,
@JsonProperty("tools") List<FunctionTool> tools) {
/** Creates input for the default chat model. */
public ChatModelActivityInput(
List<Message> messages, ModelOptions modelOptions, List<FunctionTool> tools) {
List<Message> messages, @Nullable ModelOptions modelOptions, List<FunctionTool> tools) {
this(null, messages, modelOptions, tools);
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading
Loading