fix[studio]: studio optimization (#3367)

Author: chickenlj

examples/chatbot/src/main/java/com/alibaba/cloud/ai/examples/chatbot/AgentStaticLoader.java

@@ -18,6 +18,7 @@
 
 import com.alibaba.cloud.ai.agent.studio.loader.AgentLoader;
 import com.alibaba.cloud.ai.graph.GraphRepresentation;
+import com.alibaba.cloud.ai.graph.agent.Agent;
 import com.alibaba.cloud.ai.graph.agent.BaseAgent;
 import com.alibaba.cloud.ai.graph.agent.ReactAgent;
 
@@ -48,9 +49,9 @@
 @Component
 class AgentStaticLoader implements AgentLoader {
 
-	private final Map<String, BaseAgent> agents = new ConcurrentHashMap<>();
+	private final Map<String, Agent> agents = new ConcurrentHashMap<>();
 
-	public AgentStaticLoader(BaseAgent agent) {
+	public AgentStaticLoader(Agent agent) {
 
 		GraphRepresentation representation = agent.getAndCompileGraph().stateGraph.getGraph(GraphRepresentation.Type.PLANTUML);
 		System.out.println(representation.content());
@@ -65,12 +66,12 @@ public List<String> listAgents() {
 	}
 
 	@Override
-	public BaseAgent loadAgent(String name) {
+	public Agent loadAgent(String name) {
 		if (name == null || name.trim().isEmpty()) {
 			throw new IllegalArgumentException("Agent name cannot be null or empty");
 		}
 
-		BaseAgent agent = agents.get(name);
+		Agent agent = agents.get(name);
 		if (agent == null) {
 			throw new NoSuchElementException("Agent not found: " + name);
 		}

examples/chatbot/src/main/java/com/alibaba/cloud/ai/examples/chatbot/ChatbotAgent.java

@@ -16,8 +16,10 @@
 package com.alibaba.cloud.ai.examples.chatbot;
 
 import com.alibaba.cloud.ai.graph.agent.ReactAgent;
+import com.alibaba.cloud.ai.graph.agent.hook.shelltool.ShellToolAgentHook;
 import com.alibaba.cloud.ai.graph.agent.tools.ShellTool;
 import com.alibaba.cloud.ai.graph.agent.extension.tools.filesystem.ReadFileTool;
+import com.alibaba.cloud.ai.graph.checkpoint.savers.MemorySaver;
 
 import org.springframework.ai.chat.model.ChatModel;
 import org.springframework.ai.chat.model.ToolContext;
@@ -50,12 +52,16 @@ public class ChatbotAgent {
 	public ReactAgent chatbotReactAgent(ChatModel chatModel,
 			ToolCallback executeShellCommand,
 			ToolCallback executePythonCode,
-			ToolCallback viewTextFile) {
+			ToolCallback viewTextFile,
+			MemorySaver memorySaver) {
 		return ReactAgent.builder()
 				.name("SAA")
 				.model(chatModel)
 				.instruction(INSTRUCTION)
 				.enableLogging(true)
+				.saver(memorySaver)
+				// Must set ShellToolAgentHook to manage shell session lifecycle for executeShellCommand
+				.hooks(ShellToolAgentHook.builder().shellToolName(executeShellCommand.getToolDefinition().name()).build())
 				.tools(
 						executeShellCommand,
 						executePythonCode,
@@ -64,6 +70,11 @@ public ReactAgent chatbotReactAgent(ChatModel chatModel,
 				.build();
 	}
 
+	@Bean
+	public MemorySaver memorySaver() {
+		return new MemorySaver();
+	}
+
 	// Tool: execute_shell_command
 	@Bean
 	public ToolCallback executeShellCommand() {

examples/deepresearch/src/main/java/com/alibaba/cloud/ai/examples/deepresearch/AgentStaticLoader.java

@@ -18,6 +18,7 @@
 
 import com.alibaba.cloud.ai.agent.studio.loader.AgentLoader;
 import com.alibaba.cloud.ai.graph.GraphRepresentation;
+import com.alibaba.cloud.ai.graph.agent.Agent;
 import com.alibaba.cloud.ai.graph.agent.BaseAgent;
 import com.alibaba.cloud.ai.graph.agent.ReactAgent;
 
@@ -48,7 +49,7 @@
 @Component
 class AgentStaticLoader implements AgentLoader {
 
-	private Map<String, BaseAgent> agents = new ConcurrentHashMap<>();
+	private Map<String, Agent> agents = new ConcurrentHashMap<>();
 
 //	public AgentStaticLoader(){}
 
@@ -73,12 +74,12 @@ public List<String> listAgents() {
 	}
 
 	@Override
-	public BaseAgent loadAgent(String name) {
+	public Agent loadAgent(String name) {
 		if (name == null || name.trim().isEmpty()) {
 			throw new IllegalArgumentException("Agent name cannot be null or empty");
 		}
 
-		BaseAgent agent = agents.get(name);
+		Agent agent = agents.get(name);
 		if (agent == null) {
 			throw new NoSuchElementException("Agent not found: " + name);
 		}

examples/deepresearch/src/main/java/com/alibaba/cloud/ai/examples/deepresearch/DeepResearchAgent.java

@@ -175,7 +175,7 @@ private Interceptor subAgentAsInterceptors(List<ToolCallback> toolsFromMcp) {
 						patchToolCallsInterceptor,
 						largeResultEvictionInterceptor
 				)
-				.defaultHooks(summarizationHook, toolCallLimitHook)
+				.defaultHooks(humanInTheLoopHook, summarizationHook, toolCallLimitHook)
 				.addSubAgent(researchAgent)
 				.includeGeneralPurpose(true)
 				.addSubAgent(critiqueAgent);

spring-ai-alibaba-agent-framework/src/main/java/com/alibaba/cloud/ai/graph/agent/hook/hip/HumanInTheLoopHook.java

@@ -202,6 +202,8 @@ private Optional<InterruptionMetadata> buildInterruptionMetadata(OverAllState st
 								.name(toolCall.name()).description(content).arguments(toolCall.arguments()).build())
 						.build();
 				needsInterruption = true;
+			} else {
+				builder.addToolsAutomaticallyApproved(toolCall);
 			}
 		}
 		return needsInterruption ? Optional.of(builder.build()) : Optional.empty();

spring-ai-alibaba-agent-framework/src/main/java/com/alibaba/cloud/ai/graph/agent/hook/shelltool/ShellToolAgentHook.java

@@ -42,6 +42,7 @@ public class ShellToolAgentHook extends AgentHook implements ToolInjection {
 	private static final Logger log = LoggerFactory.getLogger(ShellToolAgentHook.class);
 
 	private ShellTool shellTool;
+	private String shellToolName;
 
 	/**
 	 * Private constructor for builder pattern.
@@ -53,8 +54,9 @@ private ShellToolAgentHook() {
 	 * Private constructor with ShellTool for builder pattern.
 	 * @param shellTool the ShellTool instance to use
 	 */
-	private ShellToolAgentHook(ShellTool shellTool) {
+	private ShellToolAgentHook(ShellTool shellTool, String shellToolName) {
 		this.shellTool = shellTool;
+		this.shellToolName = shellToolName;
 	}
 
 	/**
@@ -165,7 +167,7 @@ private ShellTool extractShellTool(ToolCallback toolCallback) {
 	@Override
 	public String getRequiredToolName() {
 		// Match by tool name "shell"
-		return "shell";
+		return shellToolName;
 	}
 
 	@Override
@@ -188,6 +190,7 @@ protected ShellTool getShellTool() {
 	 */
 	public static class Builder {
 		private ShellTool shellTool;
+		private String shellToolName;
 
 		/**
 		 * Set the ShellTool instance.
@@ -199,12 +202,17 @@ public Builder shellTool(ShellTool shellTool) {
 			return this;
 		}
 
+		public Builder shellToolName(String shellToolName) {
+			this.shellToolName = shellToolName;
+			return this;
+		}
+
 		/**
 		 * Build the ShellToolAgentHook instance.
 		 * @return a new ShellToolAgentHook instance
 		 */
 		public ShellToolAgentHook build() {
-			return new ShellToolAgentHook(this.shellTool);
+			return new ShellToolAgentHook(this.shellTool, this.shellToolName);
 		}
 	}
 

spring-ai-alibaba-agent-framework/src/main/java/com/alibaba/cloud/ai/graph/agent/tools/ShellSessionManager.java

@@ -149,7 +149,7 @@ private void doCleanup(RunnableConfig config) {
 	public CommandResult executeCommand(String command, RunnableConfig config) {
 		ShellSession session = (ShellSession) config.context().get(SESSION_INSTANCE_CONTEXT_KEY);
 		if (session == null) {
-			throw new IllegalStateException("Shell session not initialized. Call initialize() first.");
+			throw new IllegalStateException("Shell session not initialized. Call initialize() first, you might need to enable ShellToolAgentHook to enable shell session management.");
 		}
 
 		log.info("Executing shell command: {}", command);

spring-ai-alibaba-graph-core/src/main/java/com/alibaba/cloud/ai/graph/action/InterruptionMetadata.java

@@ -20,6 +20,8 @@
 import com.alibaba.cloud.ai.graph.OverAllState;
 import com.alibaba.cloud.ai.graph.utils.CollectionsUtils;
 
+import org.springframework.ai.chat.messages.AssistantMessage;
+
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -37,12 +39,19 @@ public final class InterruptionMetadata extends NodeOutput implements HasMetadat
 
 	private final Map<String, Object> metadata;
 
+	private List<AssistantMessage.ToolCall> toolsAutomaticallyApproved;
+
 	private List<ToolFeedback> toolFeedbacks;
 
 	private InterruptionMetadata(Builder builder) {
 		super(builder.nodeId, builder.state);
 		this.metadata = builder.metadata();
 		this.toolFeedbacks = new ArrayList<>(builder.toolFeedbacks);
+		if (builder.toolsAutomaticallyApproved != null) {
+			this.toolsAutomaticallyApproved = builder.toolsAutomaticallyApproved;
+		} else {
+			this.toolsAutomaticallyApproved = new ArrayList<>();
+		}
 	}
 
 	/**
@@ -69,6 +78,10 @@ public List<ToolFeedback> toolFeedbacks() {
 		return toolFeedbacks;
 	}
 
+	public List<AssistantMessage.ToolCall> getToolsAutomaticallyApproved() {
+		return toolsAutomaticallyApproved;
+	}
+
 	@Override
 	public String toString() {
 		return String.format("""
@@ -92,9 +105,16 @@ public static Builder builder() {
 	}
 
 	public static Builder builder(InterruptionMetadata interruptionMetadata) {
-		return new Builder(interruptionMetadata.metadata().orElse(Map.of()))
+		Builder builder = new Builder(interruptionMetadata.metadata().orElse(Map.of()))
 			.nodeId(interruptionMetadata.node())
 			.state(interruptionMetadata.state());
+		if (interruptionMetadata.getToolsAutomaticallyApproved() != null) {
+			builder.toolsAutomaticallyApproved(interruptionMetadata.getToolsAutomaticallyApproved());
+		}
+		// if (interruptionMetadata.toolFeedbacks() != null && !interruptionMetadata.toolFeedbacks().isEmpty()) {
+		// 	builder.toolFeedbacks(interruptionMetadata.toolFeedbacks());
+		// }
+		return builder;
 	}
 
 	/**
@@ -104,6 +124,8 @@ public static Builder builder(InterruptionMetadata interruptionMetadata) {
 	public static class Builder extends HasMetadata.Builder<Builder> {
 		List<ToolFeedback> toolFeedbacks;
 
+		List<AssistantMessage.ToolCall> toolsAutomaticallyApproved;
+
 		String nodeId;
 
 		OverAllState state;
@@ -147,6 +169,19 @@ public Builder toolFeedbacks(List<ToolFeedback> toolFeedbacks) {
 			return this;
 		}
 
+		public Builder addToolsAutomaticallyApproved(AssistantMessage.ToolCall toolCall) {
+			if (this.toolsAutomaticallyApproved == null) {
+				this.toolsAutomaticallyApproved = new ArrayList<>();
+			}
+			this.toolsAutomaticallyApproved.add(toolCall);
+			return this;
+		}
+
+		public Builder toolsAutomaticallyApproved(List<AssistantMessage.ToolCall> toolsAutomaticallyApproved) {
+			this.toolsAutomaticallyApproved = new ArrayList<>(toolsAutomaticallyApproved);
+			return this;
+		}
+
 		/**
 		 * Builds the {@link InterruptionMetadata} instance.
 		 * @return a new, immutable {@link InterruptionMetadata} instance

spring-ai-alibaba-studio/src/main/java/com/alibaba/cloud/ai/agent/studio/controller/ExecutionController.java

@@ -25,7 +25,7 @@
 import com.alibaba.cloud.ai.graph.NodeOutput;
 import com.alibaba.cloud.ai.graph.RunnableConfig;
 import com.alibaba.cloud.ai.graph.action.InterruptionMetadata;
-import com.alibaba.cloud.ai.graph.agent.BaseAgent;
+import com.alibaba.cloud.ai.graph.agent.Agent;
 import com.alibaba.cloud.ai.graph.exception.GraphRunnerException;
 import com.alibaba.cloud.ai.graph.streaming.StreamingOutput;
 
@@ -171,7 +171,7 @@ public Flux<ServerSentEvent<String>> agentRunSse(@RequestBody AgentRunRequest re
 		}
 
 		try {
-			BaseAgent agent = agentLoader.loadAgent(request.appName);
+			Agent agent = agentLoader.loadAgent(request.appName);
 			RunnableConfig runnableConfig = RunnableConfig.builder()
 					.threadId(request.threadId)
 					.addMetadata("user_id", request.userId)
@@ -205,7 +205,7 @@ public Flux<ServerSentEvent<String>> agentResumeSse(@RequestBody AgentResumeRequ
 		}
 
 		try {
-			BaseAgent agent = agentLoader.loadAgent(request.appName);
+			Agent agent = agentLoader.loadAgent(request.appName);
 
 			InterruptionMetadata.Builder metadataBuilder = InterruptionMetadata.builder();
 
@@ -247,7 +247,7 @@ public Flux<ServerSentEvent<String>> agentResumeSse(@RequestBody AgentResumeRequ
 	}
 
 	@NotNull
-	private Flux<ServerSentEvent<String>> executeAgent(UserMessage userMessage, BaseAgent agent, RunnableConfig runnableConfig) throws GraphRunnerException {
+	private Flux<ServerSentEvent<String>> executeAgent(UserMessage userMessage, Agent agent, RunnableConfig runnableConfig) throws GraphRunnerException {
 
 		Flux<NodeOutput> agentStream;