Skip to content

fix: audit follow-ups — LLM-provider binding, hint-mode deser, handler dedup, tc26 flake #1144

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dhyansraj merged 5 commits into from
Jun 4, 2026
Merged

fix: audit follow-ups — LLM-provider binding, hint-mode deser, handler dedup, tc26 flake #1144

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
29d0e24
test(uc2x): poll-until-marker for tc26 cancel-event producer log (Clo…
dhyansraj Jun 3, 2026
4d6a662
refactor(java): hoist user-content assembly + Gemini tool-callback tr…
dhyansraj Jun 3, 2026
6bdfe46
fix(core): emit LLM_PROVIDER_AVAILABLE before topology insert (Closes…
dhyansraj Jun 3, 2026
834bb42
fix(llm): lenient single-value-as-array deser for hint-mode structure…
dhyansraj Jun 3, 2026
5f5d50b
test(uc2x): snapshot producer log once in tc26 poll loop (review fix,…
dhyansraj Jun 4, 2026
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
29 changes: 22 additions & 7 deletions src/runtime/core/src/runtime.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -768,11 +768,6 @@ impl AgentRuntime {
function_id, tracked.function_name, tracked.endpoint
);

// Store the tracking info first (no PyO3 involvement)
self.topology
.llm_providers
.insert(function_id.clone(), tracked.clone());

// Create LlmProviderInfo and send event
let provider_info = LlmProviderInfo {
function_id: function_id.clone(),
Expand All @@ -783,10 +778,30 @@ impl AgentRuntime {
vendor: provider.vendor.clone(),
kwargs: kwargs_json,
};
let _ = self

// Send the event BEFORE recording it in topology. If the send
// fails (channel dropped), we leave the function_id out of
// `llm_providers` so the diff gate re-emits on the next
// heartbeat (self-healing) instead of permanently suppressing
// re-emission and leaving the @MeshLlm proxy unbound.
if let Err(e) = self
.event_tx
.send(MeshEvent::llm_provider_available(provider_info))
.await;
.await
{
warn!(
"Failed to emit LLM_PROVIDER_AVAILABLE for function '{}': {}; \
will retry on next heartbeat",
function_id, e
);
continue;
}

// Store the tracking info only after a successful send so the
// diff gate treats this provider as "seen".
self.topology
.llm_providers
.insert(function_id.clone(), tracked.clone());
}
}
}
Expand Down
43 changes: 5 additions & 38 deletions ...untime/java/mcp-mesh-spring-ai/src/main/java/io/mcpmesh/ai/handlers/AnthropicHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
Expand Down Expand Up @@ -265,25 +264,8 @@ private LlmResponse generateWithToolsAutoExecute(
log.debug("Created {} tool callbacks for ChatClient", toolCallbacks.size());
}

// Extract non-system messages for user content
List<Message> nonSystemMessages = new ArrayList<>();
for (Message msg : springMessages) {
if (!(msg instanceof SystemMessage)) {
nonSystemMessages.add(msg);
}
}

// Build user content from remaining messages
StringBuilder userContent = new StringBuilder();
for (Message msg : nonSystemMessages) {
if (msg instanceof UserMessage um) {
if (userContent.length() > 0) userContent.append("\n");
userContent.append(um.getText());
} else if (msg instanceof AssistantMessage am) {
if (userContent.length() > 0) userContent.append("\n");
userContent.append("[Previous Assistant Response]\n").append(am.getText());
}
}
// Build user content from non-system messages
String userContent = buildUserContent(springMessages);

// Use ChatClient with tools
ChatClient chatClient = ChatClient.create(model);
Expand All @@ -295,7 +277,7 @@ private LlmResponse generateWithToolsAutoExecute(
}

// Add user content
requestSpec.user(userContent.toString());
requestSpec.user(userContent);

// Add tools if present - Spring AI handles tool execution automatically
if (!toolCallbacks.isEmpty()) {
Expand Down Expand Up @@ -336,7 +318,7 @@ private LlmResponse generateWithToolsAutoExecute(
if (hintSystemPrompt != null && !hintSystemPrompt.isEmpty()) {
retrySpec.system(hintSystemPrompt);
}
retrySpec.user(userContent.toString());
retrySpec.user(userContent);
if (!toolCallbacks.isEmpty()) {
retrySpec.toolCallbacks(toolCallbacks.toArray(new ToolCallback[0]));
}
Expand Down Expand Up @@ -372,22 +354,7 @@ private LlmResponse generateWithToolsNoExecute(
List<ToolDefinition> hintTools) {

// Replace system message with formatted one
List<Message> messagesWithFormattedSystem = new ArrayList<>();
boolean addedSystem = false;
for (Message msg : springMessages) {
if (msg instanceof SystemMessage) {
if (!addedSystem && formattedSystemPrompt != null && !formattedSystemPrompt.isEmpty()) {
messagesWithFormattedSystem.add(new SystemMessage(formattedSystemPrompt));
addedSystem = true;
}
} else {
messagesWithFormattedSystem.add(msg);
}
}
// Add system prompt at beginning if not already added
if (!addedSystem && formattedSystemPrompt != null && !formattedSystemPrompt.isEmpty()) {
messagesWithFormattedSystem.add(0, new SystemMessage(formattedSystemPrompt));
}
List<Message> messagesWithFormattedSystem = replaceSystemMessage(springMessages, formattedSystemPrompt);

// Create tool callbacks for schema only (no execution)
List<ToolCallback> toolCallbacks = createToolCallbacksForSchema(tools);
Expand Down
141 changes: 14 additions & 127 deletions src/runtime/java/mcp-mesh-spring-ai/src/main/java/io/mcpmesh/ai/handlers/GeminiHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,19 +6,16 @@
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.function.FunctionToolCallback;
import org.springframework.ai.google.genai.GoogleGenAiChatOptions;
import org.springframework.ai.vertexai.gemini.VertexAiGeminiChatModel;
import org.springframework.ai.vertexai.gemini.VertexAiGeminiChatOptions;

import java.util.*;
import java.util.function.Function;

/**
* LLM provider handler for Google Gemini models.
Expand Down Expand Up @@ -55,6 +52,12 @@ public String[] getAliases() {
return new String[]{"google"};
}

/** Gemini uses a shorter previous-turn prefix than the Anthropic/OpenAI default. */
@Override
public String previousResponsePrefix() {
return "[Previous Response]\n";
}

// =========================================================================
// Structured Output Methods
// =========================================================================
Expand Down Expand Up @@ -185,25 +188,8 @@ private LlmResponse generateWithToolsAutoExecute(
log.debug("Created {} tool callbacks for ChatClient", toolCallbacks.size());
}

// Extract non-system messages
List<Message> nonSystemMessages = new ArrayList<>();
for (Message msg : springMessages) {
if (!(msg instanceof SystemMessage)) {
nonSystemMessages.add(msg);
}
}

// Build user content from remaining messages
StringBuilder userContent = new StringBuilder();
for (Message msg : nonSystemMessages) {
if (msg instanceof UserMessage um) {
if (userContent.length() > 0) userContent.append("\n");
userContent.append(um.getText());
} else if (msg instanceof AssistantMessage am) {
if (userContent.length() > 0) userContent.append("\n");
userContent.append("[Previous Response]\n").append(am.getText());
}
}
// Build user content from non-system messages
String userContent = buildUserContent(springMessages);

// Use ChatClient with tools
ChatClient chatClient = ChatClient.create(model);
Expand All @@ -215,7 +201,7 @@ private LlmResponse generateWithToolsAutoExecute(
}

// Add user content
requestSpec.user(userContent.toString());
requestSpec.user(userContent);

// Add tools if present
if (!toolCallbacks.isEmpty()) {
Expand Down Expand Up @@ -251,22 +237,7 @@ private LlmResponse generateWithToolsNoExecute(
OutputSchema outputSchema) {

// Replace system message with formatted one
List<Message> messagesWithFormattedSystem = new ArrayList<>();
boolean addedSystem = false;
for (Message msg : springMessages) {
if (msg instanceof SystemMessage) {
if (!addedSystem && formattedSystemPrompt != null && !formattedSystemPrompt.isEmpty()) {
messagesWithFormattedSystem.add(new SystemMessage(formattedSystemPrompt));
addedSystem = true;
}
} else {
messagesWithFormattedSystem.add(msg);
}
}
// Add system prompt at beginning if not already added
if (!addedSystem && formattedSystemPrompt != null && !formattedSystemPrompt.isEmpty()) {
messagesWithFormattedSystem.add(0, new SystemMessage(formattedSystemPrompt));
}
List<Message> messagesWithFormattedSystem = replaceSystemMessage(springMessages, formattedSystemPrompt);

// Create tool callbacks for schema only (no execution)
List<ToolCallback> toolCallbacks = createToolCallbacksForSchema(tools);
Expand Down Expand Up @@ -432,49 +403,12 @@ private void applyResponseFormat(ChatClient.ChatClientRequestSpec requestSpec, O
}

/**
* Create ToolCallbacks for schema only (no execution).
*
* <p>Overrides the default to apply Gemini-specific uppercase type conversion.
* Apply Gemini-specific uppercase type conversion to a tool's input schema
* before it is serialized and attached to the Spring AI tool callback.
*/
@Override
public List<ToolCallback> createToolCallbacksForSchema(List<ToolDefinition> tools) {
List<ToolCallback> callbacks = new ArrayList<>();
if (tools == null) return callbacks;

for (ToolDefinition tool : tools) {
// Create a dummy function that should never be called
java.util.function.Function<Map<String, Object>, String> dummyFunction = args -> {
log.warn("Tool {} was unexpectedly called - this shouldn't happen!", tool.name());
return "{\"error\": \"Tool execution not supported in provider mode\"}";
};

// Convert inputSchema Map to JSON string with uppercase types for Gemini
String inputSchemaJson = null;
if (tool.inputSchema() != null && !tool.inputSchema().isEmpty()) {
try {
Map<String, Object> convertedSchema = convertSchemaTypesToUpperCase(tool.inputSchema());
inputSchemaJson = MAPPER
.writeValueAsString(convertedSchema);
log.debug("Converted tool schema for {}: {}", tool.name(), inputSchemaJson);
} catch (Exception e) {
log.warn("Failed to serialize inputSchema for {}: {}", tool.name(), e.getMessage());
}
}

@SuppressWarnings("unchecked")
var builder = FunctionToolCallback
.builder(tool.name(), dummyFunction)
.description(tool.description() != null ? tool.description() : "No description")
.inputType((Class<Map<String, Object>>) (Class<?>) Map.class);

if (inputSchemaJson != null) {
builder.inputSchema(inputSchemaJson);
}

callbacks.add(builder.build());
}

return callbacks;
public Map<String, Object> transformToolInputSchema(Map<String, Object> schema) {
return convertSchemaTypesToUpperCase(schema);
}

/**
Expand Down Expand Up @@ -557,53 +491,6 @@ private Map<String, Object> convertSchemaTypesToUpperCase(Map<String, Object> sc
return result;
}

/**
* Create a Spring AI ToolCallback from our ToolDefinition.
*
* <p>Overrides the default to apply Gemini-specific uppercase type conversion.
*/
@Override
public ToolCallback createToolCallback(ToolDefinition tool, ToolExecutorCallback toolExecutor) {
Function<Map<String, Object>, String> toolFunction = args -> {
try {
String argsJson = args != null ? MAPPER
.writeValueAsString(args) : "{}";
return toolExecutor.execute(tool.name(), argsJson);
} catch (Exception e) {
log.error("Tool execution failed: {}", tool.name(), e);
try {
return TOOL_CALLBACK_MAPPER.writeValueAsString(Map.of("error", "Tool execution failed: " + tool.name()));
} catch (Exception ignored) {
return "{\"error\": \"tool execution failed\"}";
}
}
};

// Convert inputSchema with uppercase types for Gemini API compatibility
String inputSchemaJson = null;
if (tool.inputSchema() != null && !tool.inputSchema().isEmpty()) {
try {
Map<String, Object> convertedSchema = convertSchemaTypesToUpperCase(tool.inputSchema());
inputSchemaJson = MAPPER
.writeValueAsString(convertedSchema);
} catch (Exception e) {
log.warn("Failed to serialize inputSchema for {}: {}", tool.name(), e.getMessage());
}
}

@SuppressWarnings("unchecked")
var builder = FunctionToolCallback
.builder(tool.name(), toolFunction)
.description(tool.description() != null ? tool.description() : "No description")
.inputType((Class<Map<String, Object>>) (Class<?>) Map.class);

if (inputSchemaJson != null) {
builder.inputSchema(inputSchemaJson);
}

return builder.build();
}

@Override
public Map<String, Boolean> getCapabilities() {
return Map.of(
Expand Down
Loading
Loading