From d342b6c0f61f9dce70a60e7b93b45f8f78157dcd Mon Sep 17 00:00:00 2001 From: xuehuitian45 <13069167198@163.com> Date: Tue, 9 Dec 2025 14:19:16 +0800 Subject: [PATCH] fix: fix message bug and adapt a2a to new format --- .../agentscope/AgentScopeStreamAdapter.java | 416 ++++-------------- .../protocol/a2a/GraphAgentExecutor.java | 136 ++++-- 2 files changed, 189 insertions(+), 363 deletions(-) diff --git a/core/src/main/java/io/agentscope/runtime/adapters/agentscope/AgentScopeStreamAdapter.java b/core/src/main/java/io/agentscope/runtime/adapters/agentscope/AgentScopeStreamAdapter.java index 98d1b4ed..ad05ee3f 100644 --- a/core/src/main/java/io/agentscope/runtime/adapters/agentscope/AgentScopeStreamAdapter.java +++ b/core/src/main/java/io/agentscope/runtime/adapters/agentscope/AgentScopeStreamAdapter.java @@ -84,44 +84,9 @@ public Flux adaptAgentScopeMessageStreamReactive(Flux { - StreamState state = new StreamState(); - return flux .flatMap(event -> { - Msg msg = event.getMessage(); - boolean last = event.isLast(); - - // If a new message, reset state - if (state.msgId == null || !msg.getId().equals(state.msgId)) { - state.reset(msg.getId()); - } - - // Process the message content and emit all resulting events - List newEvents = processMessageContent(msg, last, state); - return Flux.fromIterable(newEvents); - }) - .concatWith(Flux.defer(() -> { - // Handle last_content if any (final processing) - if (state.lastContent != null && !state.lastContent.isEmpty()) { - List finalEvents = new ArrayList<>(); - - if (state.shouldStartMessage) { - state.index = null; - updateMessageAttrs(state.message, state.metadata, state.usage); - finalEvents.add(state.message.inProgress()); - } - - TextContent textDeltaContent = new TextContent(true, state.index, state.lastContent); - textDeltaContent = (TextContent) state.message.addDeltaContent(textDeltaContent); - finalEvents.add(textDeltaContent); - - updateMessageAttrs(state.message, state.metadata, state.usage); - finalEvents.add(state.message.completed()); - - return Flux.fromIterable(finalEvents); - } - return Flux.empty(); - })); + List newEvents = processMessageContent(event); + return Flux.fromIterable(newEvents); }); } @@ -129,8 +94,9 @@ public Flux adaptAgentScopeMessageStreamReactive(Flux processMessageContent(Msg msg, boolean last, StreamState state) { + private List processMessageContent(io.agentscope.core.agent.Event event) { List results = new ArrayList<>(); + Msg msg = event.getMessage(); List content = msg.getContent(); // Handle string content. @@ -139,55 +105,26 @@ private List processMessageContent(Msg msg, boolean last, StreamState sta return results; } - // Separate tool_use blocks from other blocks - List newBlocks = new ArrayList<>(); - List newToolBlocks = new ArrayList<>(); - - for (ContentBlock block : content) { - if (block instanceof ToolUseBlock) { - newToolBlocks.add(block); - } else { - newBlocks.add(block); - } - } - - // Update content based on tool_start flag - List blocksToProcess; - if (!newToolBlocks.isEmpty()) { - if (state.toolStart || newBlocks.isEmpty()) { - blocksToProcess = newToolBlocks; - } else { - blocksToProcess = newBlocks; - state.toolStart = true; - } - } else { - blocksToProcess = newBlocks; - } - - // Update metadata and usage - state.metadata = msg.getMetadata(); - // Note: AgentScope Java Msg doesn't have usage field directly - // It might be in metadata or we skip it - state.usage = null; // Usage not available in AgentScope Java Msg - - // Process each block - // Content is always List, so we process blocks directly - for (ContentBlock element : blocksToProcess) { - if (element instanceof TextBlock) { - processTextBlock((TextBlock) element, last, state, results); - } else if (element instanceof ThinkingBlock) { - processThinkingBlock((ThinkingBlock) element, last, state, results); - } else if (element instanceof ToolUseBlock) { - processToolUseBlock((ToolUseBlock) element, last, state, results); - } else if (element instanceof ToolResultBlock) { - processToolResultBlock((ToolResultBlock) element, state, results); - } else if (element instanceof ImageBlock) { - processImageBlock((ImageBlock) element, state, results); - } else if (element instanceof AudioBlock) { - processAudioBlock((AudioBlock) element, state, results); - } else { - // Fallback: convert to text - processUnknownBlock(element, state, results); + Map metadata = msg.getMetadata(); + + for (ContentBlock element : content) { + if(!event.isLast()||element instanceof ToolUseBlock || element instanceof ToolResultBlock){ + if (element instanceof TextBlock textBlock) { + processTextBlock(textBlock, metadata, results); + } else if (element instanceof ThinkingBlock thinkingBlock) { + processThinkingBlock(thinkingBlock, metadata, results); + } else if (element instanceof ToolUseBlock toolUseBlock) { + processToolUseBlock(toolUseBlock, metadata, results); + } else if (element instanceof ToolResultBlock toolResultBlock) { + processToolResultBlock(toolResultBlock, metadata, results); + } else if (element instanceof ImageBlock imageBlock) { + processImageBlock(imageBlock, metadata, results); + } else if (element instanceof AudioBlock audioBlock) { + processAudioBlock(audioBlock, metadata, results); + } else { + // Fallback: convert to text + processUnknownBlock(element, metadata, results); + } } } @@ -197,194 +134,76 @@ private List processMessageContent(Msg msg, boolean last, StreamState sta /** * Process text block with incremental updates and deduplication. */ - private void processTextBlock(TextBlock block, boolean last, StreamState state, List results) { + private void processTextBlock(TextBlock block, Map metadata, List results) { String text = block.getText(); if (text == null || text.isEmpty()) { return; } - // Start message if needed - if (state.shouldStartMessage) { - state.index = null; - updateMessageAttrs(state.message, state.metadata, state.usage); - results.add(state.message.inProgress()); - state.shouldStartMessage = false; - } - - // Remove prefix (deduplication logic) - String newText = removePrefix(text, state.localTruncateMemory); - state.localTruncateMemory = text; - - // Create delta content - TextContent textDeltaContent = new TextContent(true, state.index, newText); - textDeltaContent = (TextContent) state.message.addDeltaContent(textDeltaContent); - state.index = textDeltaContent.getIndex(); - - // Only yield valid text - if (textDeltaContent.getText() != null && !textDeltaContent.getText().isEmpty()) { - results.add(textDeltaContent); - } + Message message = new Message(MessageType.MESSAGE, "assistant"); + updateMessageAttrs(message, metadata, null); - // Complete message if last or tool_start - if (last || state.toolStart) { - if (state.index != null && state.index < state.message.getContent().size()) { - Content completedContent = state.message.getContent().get(state.index); - if (completedContent instanceof TextContent) { - TextContent textContent = (TextContent) completedContent; - if (textContent.getText() != null && !textContent.getText().isEmpty()) { - results.add(completedContent.completed()); - } - } - } - - updateMessageAttrs(state.message, state.metadata, state.usage); - results.add(state.message.completed()); - - // Reset for next message - state.message = new Message(MessageType.MESSAGE, "assistant"); - state.index = null; - state.shouldStartMessage = true; - } + TextContent textContent = new TextContent(true, null, text); + textContent = (TextContent) message.addDeltaContent(textContent); + results.add(textContent); } /** * Process thinking block with incremental updates and deduplication. */ - private void processThinkingBlock(ThinkingBlock block, boolean last, StreamState state, List results) { + private void processThinkingBlock(ThinkingBlock block, Map metadata, List results) { String reasoning = block.getThinking(); if (reasoning == null || reasoning.isEmpty()) { return; } - // Start reasoning message if needed - if (state.shouldStartReasoningMessage) { - state.index = null; - updateMessageAttrs(state.reasoningMessage, state.metadata, state.usage); - results.add(state.reasoningMessage.inProgress()); - state.shouldStartReasoningMessage = false; - } - - // Remove prefix (deduplication logic) - String newReasoning = removePrefix(reasoning, state.localTruncateReasoningMemory); - state.localTruncateReasoningMemory = reasoning; - - // Create delta content - TextContent textDeltaContent = new TextContent(true, state.index, newReasoning); - textDeltaContent = (TextContent) state.reasoningMessage.addDeltaContent(textDeltaContent); - state.index = textDeltaContent.getIndex(); - - // Only yield valid text - // Return textDeltaContent directly - if (textDeltaContent.getText() != null && !textDeltaContent.getText().isEmpty()) { - results.add(textDeltaContent); - } + Message reasoningMessage = new Message(MessageType.REASONING, "assistant"); + updateMessageAttrs(reasoningMessage, metadata, null); - // Complete reasoning message if last or tool_start - if (last || state.toolStart) { - if (state.index != null && state.index < state.reasoningMessage.getContent().size()) { - Content completedContent = state.reasoningMessage.getContent().get(state.index); - if (completedContent instanceof TextContent) { - TextContent textContent = (TextContent) completedContent; - if (textContent.getText() != null && !textContent.getText().isEmpty()) { - results.add(completedContent.completed()); - } - } - } - - updateMessageAttrs(state.reasoningMessage, state.metadata, state.usage); - results.add(state.reasoningMessage.completed()); - - // Reset for next reasoning message - state.reasoningMessage = new Message(MessageType.REASONING, "assistant"); - state.index = null; - } + TextContent textContent = new TextContent(true, null, reasoning); + textContent = (TextContent) reasoningMessage.addDeltaContent(textContent); + results.add(textContent); } /** * Process tool_use block with staged building (create empty, fill on last). */ - private void processToolUseBlock(ToolUseBlock block, boolean last, StreamState state, List results) { - String callId = block.getId(); - - if (last) { - // Fill in the complete arguments - Message pluginCallMessage = state.toolUseMessagesDict.get(callId); - if (pluginCallMessage == null) { - logger.warn("Tool use message not found for call_id: " + callId); - return; - } - - // Serialize input to JSON - String jsonStr; - try { - jsonStr = objectMapper.writeValueAsString(block.getInput()); - } catch (Exception e) { - logger.error("Failed to serialize tool input", e); - jsonStr = "{}"; - } - - // Create FunctionCall data - FunctionCall functionCall = new FunctionCall( - block.getId(), - block.getName(), - jsonStr - ); - - // Convert to Map for DataContent - Map callData = new HashMap<>(); - callData.put("call_id", functionCall.getCallId()); - callData.put("name", functionCall.getName()); - callData.put("arguments", functionCall.getArguments()); - - DataContent dataDeltaContent = new DataContent(); - dataDeltaContent.setDelta(true); - dataDeltaContent.setIndex(state.index); - dataDeltaContent.setData(callData); - - dataDeltaContent = (DataContent) pluginCallMessage.addDeltaContent(dataDeltaContent); - results.add(dataDeltaContent.completed()); - - updateMessageAttrs(pluginCallMessage, state.metadata, state.usage); - results.add(pluginCallMessage.completed()); - state.index = null; - } else { - // Create new tool call message if not exists - if (!state.toolUseMessagesDict.containsKey(callId)) { - Message pluginCallMessage = new Message(MessageType.PLUGIN_CALL, "assistant"); - - // Create FunctionCall with empty arguments - FunctionCall functionCall = new FunctionCall( - block.getId(), - block.getName(), - "" - ); - - // Convert to Map for DataContent - Map callData = new HashMap<>(); - callData.put("call_id", functionCall.getCallId()); - callData.put("name", functionCall.getName()); - callData.put("arguments", functionCall.getArguments()); - - DataContent dataDeltaContent = new DataContent(); - dataDeltaContent.setDelta(true); - dataDeltaContent.setIndex(state.index); - dataDeltaContent.setData(callData); - - updateMessageAttrs(pluginCallMessage, state.metadata, state.usage); - results.add(pluginCallMessage.inProgress()); - - dataDeltaContent = (DataContent) pluginCallMessage.addDeltaContent(dataDeltaContent); - results.add(dataDeltaContent); - - state.toolUseMessagesDict.put(callId, pluginCallMessage); - } + private void processToolUseBlock(ToolUseBlock block, Map metadata, List results) { + // Serialize input to JSON + String jsonStr; + try { + jsonStr = objectMapper.writeValueAsString(block.getInput()); + } catch (Exception e) { + logger.error("Failed to serialize tool input", e); + jsonStr = "{}"; } + + FunctionCall functionCall = new FunctionCall( + block.getId(), + block.getName(), + jsonStr + ); + + Map callData = new HashMap<>(); + callData.put("call_id", functionCall.getCallId()); + callData.put("name", functionCall.getName()); + callData.put("arguments", functionCall.getArguments()); + + DataContent dataDeltaContent = new DataContent(); + dataDeltaContent.setDelta(true); + dataDeltaContent.setData(callData); + + Message pluginCallMessage = new Message(MessageType.MCP_TOOL_CALL, "assistant"); + pluginCallMessage.addDeltaContent(dataDeltaContent); + updateMessageAttrs(pluginCallMessage, metadata, null); + + results.add(pluginCallMessage.completed()); } /** * Process tool_result block. */ - private void processToolResultBlock(ToolResultBlock block, StreamState state, List results) { + private void processToolResultBlock(ToolResultBlock block, Map metadata, List results) { // Serialize output to JSON String jsonStr; try { @@ -418,36 +237,24 @@ private void processToolResultBlock(ToolResultBlock block, StreamState state, Li outputData.put("output", functionCallOutput.getOutput()); DataContent dataDeltaContent = new DataContent(); - dataDeltaContent.setIndex(state.index); dataDeltaContent.setData(outputData); - Message pluginOutputMessage = new Message(MessageType.PLUGIN_CALL_OUTPUT, "tool"); + Message pluginOutputMessage = new Message(MessageType.MCP_APPROVAL_RESPONSE, "tool"); pluginOutputMessage.setContent(List.of(dataDeltaContent)); - updateMessageAttrs(pluginOutputMessage, state.metadata, state.usage); + updateMessageAttrs(pluginOutputMessage, metadata, null); results.add(pluginOutputMessage.completed()); - - // Reset message state - state.message = new Message(MessageType.MESSAGE, "assistant"); - state.shouldStartMessage = true; - state.index = null; } /** * Process image block. */ - private void processImageBlock(ImageBlock block, StreamState state, List results) { - // Start message if needed - if (state.shouldStartMessage) { - state.index = null; - updateMessageAttrs(state.message, state.metadata, state.usage); - results.add(state.message.inProgress()); - state.shouldStartMessage = false; - } + private void processImageBlock(ImageBlock block, Map metadata, List results) { + Message message = new Message(MessageType.MESSAGE, "assistant"); + updateMessageAttrs(message, metadata, null); ImageContent deltaContent = new ImageContent(); deltaContent.setDelta(true); - deltaContent.setIndex(state.index); Source source = block.getSource(); if (source instanceof URLSource) { @@ -464,26 +271,20 @@ private void processImageBlock(ImageBlock block, StreamState state, List deltaContent.setImageUrl(url); } - deltaContent = (ImageContent) state.message.addDeltaContent(deltaContent); - state.index = deltaContent.getIndex(); + deltaContent = (ImageContent) message.addDeltaContent(deltaContent); results.add(deltaContent); + results.add(message.completed()); } /** * Process audio block. */ - private void processAudioBlock(AudioBlock block, StreamState state, List results) { - // Start message if needed - if (state.shouldStartMessage) { - state.index = null; - updateMessageAttrs(state.message, state.metadata, state.usage); - results.add(state.message.inProgress()); - state.shouldStartMessage = false; - } + private void processAudioBlock(AudioBlock block, Map metadata, List results) { + Message message = new Message(MessageType.MESSAGE, "assistant"); + updateMessageAttrs(message, metadata, null); AudioContent deltaContent = new AudioContent(); deltaContent.setDelta(true); - deltaContent.setIndex(state.index); Source source = block.getSource(); if (source instanceof URLSource) { @@ -509,27 +310,22 @@ private void processAudioBlock(AudioBlock block, StreamState state, List deltaContent.setFormat(mediaType); } - deltaContent = (AudioContent) state.message.addDeltaContent(deltaContent); - state.index = deltaContent.getIndex(); + deltaContent = (AudioContent) message.addDeltaContent(deltaContent); results.add(deltaContent); + results.add(message.completed()); } /** * Process unknown block type (fallback to text). */ - private void processUnknownBlock(ContentBlock block, StreamState state, List results) { - // Start message if needed - if (state.shouldStartMessage) { - state.index = null; - updateMessageAttrs(state.message, state.metadata, state.usage); - results.add(state.message.inProgress()); - state.shouldStartMessage = false; - } + private void processUnknownBlock(ContentBlock block, Map metadata, List results) { + Message message = new Message(MessageType.MESSAGE, "assistant"); + updateMessageAttrs(message, metadata, null); - TextContent deltaContent = new TextContent(true, state.index, block.toString()); - deltaContent = (TextContent) state.message.addDeltaContent(deltaContent); - state.index = deltaContent.getIndex(); + TextContent deltaContent = new TextContent(true, null, block.toString()); + deltaContent = (TextContent) message.addDeltaContent(deltaContent); results.add(deltaContent); + results.add(message.completed()); } /** @@ -547,16 +343,6 @@ private void updateMessageAttrs(Message message, Map metadata, M /** * Remove prefix from string */ - private String removePrefix(String str, String prefix) { - if (prefix == null || prefix.isEmpty() || !str.startsWith(prefix)) { - return str; - } - return str.substring(prefix.length()); - } - - /** - * Convert ContentBlock to Map for JSON serialization. - */ private Map contentBlockToMap(ContentBlock block) { Map map = new HashMap<>(); if (block instanceof TextBlock) { @@ -597,39 +383,5 @@ private Map contentBlockToMap(ContentBlock block) { } return map; } - - /** - * State class to maintain streaming state across events. - */ - private static class StreamState { - String msgId = null; - String lastContent = ""; - Map metadata = null; - Map usage = null; - boolean toolStart = false; - Message message = new Message(MessageType.MESSAGE, "assistant"); - Message reasoningMessage = new Message(MessageType.REASONING, "assistant"); - String localTruncateMemory = ""; - String localTruncateReasoningMemory = ""; - boolean shouldStartMessage = true; - boolean shouldStartReasoningMessage = true; - Map toolUseMessagesDict = new HashMap<>(); - Integer index = null; - List pendingMessages = new ArrayList<>(); - - void reset(String newMsgId) { - localTruncateMemory = ""; - localTruncateReasoningMemory = ""; - lastContent = ""; - message = new Message(MessageType.MESSAGE, "assistant"); - reasoningMessage = new Message(MessageType.REASONING, "assistant"); - shouldStartMessage = true; - shouldStartReasoningMessage = true; - index = null; - toolStart = false; - toolUseMessagesDict.clear(); - msgId = newMsgId; - } - } } diff --git a/web/src/main/java/io/agentscope/runtime/protocol/a2a/GraphAgentExecutor.java b/web/src/main/java/io/agentscope/runtime/protocol/a2a/GraphAgentExecutor.java index c44c3fcf..0cd068d0 100644 --- a/web/src/main/java/io/agentscope/runtime/protocol/a2a/GraphAgentExecutor.java +++ b/web/src/main/java/io/agentscope/runtime/protocol/a2a/GraphAgentExecutor.java @@ -20,30 +20,17 @@ import io.a2a.server.agentexecution.RequestContext; import io.a2a.server.events.EventQueue; import io.a2a.server.tasks.TaskUpdater; - -import io.a2a.spec.JSONRPCError; -import io.a2a.spec.Part; -import io.a2a.spec.Task; -import io.a2a.spec.TaskState; -import io.a2a.spec.TaskStatus; -import io.a2a.spec.TextPart; - +import io.a2a.spec.*; import io.agentscope.runtime.engine.Runner; -import io.agentscope.runtime.engine.schemas.AgentRequest; -import io.agentscope.runtime.engine.schemas.Content; +import io.agentscope.runtime.engine.schemas.*; import io.agentscope.runtime.engine.schemas.Event; -import io.agentscope.runtime.engine.schemas.MessageType; -import io.agentscope.runtime.engine.schemas.Role; -import io.agentscope.runtime.engine.schemas.TextContent; +import io.agentscope.runtime.engine.schemas.Message; import org.reactivestreams.Subscription; import reactor.core.publisher.Flux; -import io.agentscope.runtime.engine.schemas.Message; - import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Objects; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicBoolean; @@ -66,6 +53,7 @@ public void execute(RequestContext context, EventQueue eventQueue) throws JSONRP try { AgentRequest agentRequest = buildAgentRequest(context); Flux resultFlux = runner.streamQuery(agentRequest); + Task task = context.getTask(); if (task == null) { task = newTask(context.getMessage()); @@ -137,7 +125,7 @@ private boolean isBlockRequest(RequestContext context) { if (null == context.getParams().configuration()) { return true; } - return Boolean.FALSE.equals(context.getParams().configuration().blocking()); + return Boolean.TRUE.equals(context.getParams().configuration().blocking()); } private void processTaskBlocking(RequestContext context, EventQueue eventQueue, Task task, Flux resultFlux) { @@ -149,20 +137,59 @@ private void processTaskBlocking(RequestContext context, EventQueue eventQueue, }) .doOnNext(output -> { try { - if (output instanceof Message m) { - List contents = m.getContent(); - if (contents != null && !contents.isEmpty() && contents.get(0) instanceof TextContent text) { + if (output instanceof Content) { + // Todo: only process text content for now, need to handle other content types later + if (output instanceof TextContent text) { String content = text.getText(); accumulatedOutput.append(content); logger.info("Appended content chunk (" + content.length() + " chars), total so far: " + accumulatedOutput.length()); } } + // Todo: need to know whether the blocking mode should also handle tool calls and responses + else if (output instanceof Message message) { + if (message.getType().equals("mcp_call")) { + for (Content content : message.getContent()) { + if (content instanceof DataContent dataContent) { + if (dataContent.getData() == null || !dataContent.getData().containsKey("name") || dataContent.getData().get("name").toString().isEmpty()) { + continue; + } + System.out.println("Processing tool call: " + dataContent.getData()); + String toolName = dataContent.getData().get("name").toString(); + String arguments = dataContent.getData().get("arguments").toString(); + String callId = dataContent.getData().get("call_id").toString(); + String textContent = "Calling tool " + toolName + " with arguments: " + arguments + " (call ID: " + callId + ")"; + Map metaData = new HashMap<>(); + metaData.put("type", "toolCall"); + accumulatedOutput.append(textContent); + // Todo: Still need to know the exact token usage for tool call + } + } + } else if (message.getType().equals("mcp_approval_response")) { + for (Content content : message.getContent()) { + if (content instanceof DataContent dataContent) { + if (dataContent.getData() == null || !dataContent.getData().containsKey("name") || dataContent.getData().get("name").toString().isEmpty()) { + continue; + } + System.out.println("Processing tool call: " + dataContent.getData()); + String toolResult = dataContent.getData().get("output").toString(); + String toolName = dataContent.getData().get("name").toString(); + String callId = dataContent.getData().get("call_id").toString(); + String textContent = "Tool " + toolName + " returned result: " + toolResult + " (call ID: " + callId + ")"; + Map metaData = new HashMap<>(); + metaData.put("type", "toolResponse"); + accumulatedOutput.append(textContent); + // Todo: Still need to know the exact token usage for tool call + } + } + } + } } catch (Exception ignored) { } }) .doOnComplete(() -> { logger.info("Subscribe and process stream output completed successfully"); + // Todo: Still need to decide whether to send the accumulated output as a final message in blocking mode io.a2a.spec.Message resultMessage = A2A.createAgentTextMessage(accumulatedOutput.toString(), context.getContextId(), context.getTaskId()); @@ -207,6 +234,7 @@ private void processTaskNonBlocking(RequestContext context, EventQueue eventQueu private void processStreamingOutput(Flux resultFlux, TaskUpdater taskUpdater, StringBuilder accumulatedOutput) { String artifactId = UUID.randomUUID().toString(); AtomicBoolean isFirstArtifact = new AtomicBoolean(true); + try { resultFlux .doOnSubscribe(s -> { @@ -216,20 +244,14 @@ private void processStreamingOutput(Flux resultFlux, TaskUpdater taskUpda }) .doOnNext(output -> { try { - if (output instanceof Message m) { - List contents = m.getContent(); - if (contents != null && !contents.isEmpty() && contents.get(0) instanceof TextContent text) { + if (output instanceof Content) { + System.out.println(output.getClass().getName()); + // Todo: only process text content for now, need to handle other content types later + if (output instanceof TextContent text) { String content = text.getText(); Map metaData = new HashMap<>(); - if (Objects.equals(m.getType(), MessageType.FUNCTION_CALL)) { - metaData.put("type", "toolCall"); - } else if (Objects.equals(m.getType(), MessageType.FUNCTION_CALL_OUTPUT)) { - metaData.put("type", "toolResponse"); - } else { - metaData.put("type", "chunk"); - } + metaData.put("type", "chunk"); if (content != null && !content.isEmpty()) { - taskUpdater.addArtifact( List.of(new TextPart(content)), artifactId, @@ -243,6 +265,58 @@ private void processStreamingOutput(Flux resultFlux, TaskUpdater taskUpda + accumulatedOutput.length()); } } + } else if (output instanceof Message message) { + if (message.getType().equals("mcp_call")) { + for (Content content : message.getContent()) { + if (content instanceof DataContent dataContent) { + if (dataContent.getData() == null || !dataContent.getData().containsKey("name") || dataContent.getData().get("name").toString().isEmpty()) { + continue; + } + System.out.println("Processing tool call: " + dataContent.getData()); + String toolName = dataContent.getData().get("name").toString(); + String arguments = dataContent.getData().get("arguments").toString(); + String callId = dataContent.getData().get("call_id").toString(); + String textContent = "Calling tool " + toolName + " with arguments: " + arguments + " (call ID: " + callId + ")"; + Map metaData = new HashMap<>(); + metaData.put("type", "toolCall"); + taskUpdater.addArtifact( + List.of(new TextPart(textContent)), + artifactId, + "agent-response", + metaData, + !isFirstArtifact.getAndSet(false), + false + ); + accumulatedOutput.append(toolName); + // Todo: Still need to know the exact token usage for tool call + } + } + } else if (message.getType().equals("mcp_approval_response")) { + for (Content content : message.getContent()) { + if (content instanceof DataContent dataContent) { + if (dataContent.getData() == null || !dataContent.getData().containsKey("name") || dataContent.getData().get("name").toString().isEmpty()) { + continue; + } + System.out.println("Processing tool call: " + dataContent.getData()); + String toolResult = dataContent.getData().get("output").toString(); + String toolName = dataContent.getData().get("name").toString(); + String callId = dataContent.getData().get("call_id").toString(); + String textContent = "Tool " + toolName + " returned result: " + toolResult + " (call ID: " + callId + ")"; + Map metaData = new HashMap<>(); + metaData.put("type", "toolResponse"); + taskUpdater.addArtifact( + List.of(new TextPart(textContent)), + artifactId, + "agent-response", + metaData, + !isFirstArtifact.getAndSet(false), + false + ); + accumulatedOutput.append(toolName); + // Todo: Still need to know the exact token usage for tool call + } + } + } } } catch (Exception ignored) { }