Connect the agent to remote MCP tools: policy registry and bounded loop

1
spring:
2
  ai:
3
    mcp:
4
      client:
5
        enabled: true
6
        type: SYNC
7
        toolcallback:
8
          enabled: false          # we do NOT auto-register tool callbacks
9
        streamable-http:
10
          connections:
11
            ops:
12
              url: ${MCP_SERVER_URL:http://localhost:8081}
13
              endpoint: /mcp
14

15
agent:
16
  loop:
17
    max-model-calls: 8
18
    max-tool-calls: 10
19
    deadline: 45s
20
  tools:
21
    connect-timeout: 2s
22
    per-call-timeout: 12s

1
package in.o612.eng.opsagent.agent.tools;
2

3
import java.time.Duration;
4

5
public record ToolPolicy(
6
        String toolName,
7
        Risk risk,
8
        String requiredScope,       // null = any authenticated caller
9
        Duration perCallTimeout,
10
        boolean retryable,
11
        boolean requiresApproval,
12
        AuditLevel auditLevel) {
13

14
    public enum Risk { LOW, MEDIUM, HIGH }
15
    public enum AuditLevel { STANDARD, WRITE }
16
}

1
package in.o612.eng.opsagent.agent.tools;
2

3
import org.springframework.stereotype.Component;
4

5
import java.time.Duration;
6
import java.util.Map;
7
import java.util.Optional;
8

9
@Component
10
public class ToolPolicyRegistry {
11

12
    private final Map<String, ToolPolicy> policies = Map.of(
13
        "get_service_status",
14
            new ToolPolicy("get_service_status", ToolPolicy.Risk.LOW,
15
                    "ops:read", Duration.ofSeconds(10), true, false,
16
                    ToolPolicy.AuditLevel.STANDARD),
17
        "list_recent_incidents",
18
            new ToolPolicy("list_recent_incidents", ToolPolicy.Risk.LOW,
19
                    "ops:read", Duration.ofSeconds(10), true, false,
20
                    ToolPolicy.AuditLevel.STANDARD),
21
        "get_incident",
22
            new ToolPolicy("get_incident", ToolPolicy.Risk.LOW,
23
                    "ops:read", Duration.ofSeconds(10), true, false,
24
                    ToolPolicy.AuditLevel.STANDARD),
25
        // write tools arrive in Chapter 10 — declared here so the policy
26
        // surface is visible even before they exist:
27
        "create_incident",
28
            new ToolPolicy("create_incident", ToolPolicy.Risk.HIGH,
29
                    "ops:incident:write", Duration.ofSeconds(15), false, true,
30
                    ToolPolicy.AuditLevel.WRITE),
31
        "append_incident_note",
32
            new ToolPolicy("append_incident_note", ToolPolicy.Risk.MEDIUM,
33
                    "ops:note:write", Duration.ofSeconds(15), false, true,
34
                    ToolPolicy.AuditLevel.WRITE)
35
    );
36

37
    public Optional<ToolPolicy> find(String tool) { return Optional.ofNullable(policies.get(tool)); }
38
}

1
package in.o612.eng.opsagent.agent.tools;
2

3
import com.fasterxml.jackson.databind.ObjectMapper;
4
import io.modelcontextprotocol.client.McpSyncClient;
5
import io.modelcontextprotocol.spec.McpSchema;
6
import org.springframework.stereotype.Component;
7

8
import java.time.Duration;
9
import java.util.HashMap;
10
import java.util.List;
11
import java.util.Map;
12
import java.util.concurrent.*;
13

14
@Component
15
public class McpToolExecutor {
16

17
    private final List<McpSyncClient> clients;   // auto-configured by the MCP client starter
18
    private final ObjectMapper mapper = new ObjectMapper();
19
    private final ExecutorService pool = Executors.newVirtualThreadPerTaskExecutor();
20

21
    public McpToolExecutor(List<McpSyncClient> clients) { this.clients = clients; }
22

23
    public String execute(String tool, Map<String, Object> args,
24
                          String tenantId, Duration timeout) {
25
        var merged = new HashMap<String, Object>(args);
26
        merged.put("tenant", tenantId);          // tenant is OURS, not the model's
27
        try {
28
            var future = pool.submit(() ->
29
                    clients.get(0).callTool(new McpSchema.CallToolRequest(tool, merged)));
30
            var result = future.get(timeout.toMillis(), TimeUnit.MILLISECONDS);
31
            return mapper.writeValueAsString(result.content());
32
        } catch (TimeoutException e) {
33
            return "{\"error\":\"TOOL_TIMEOUT\"}";
34
        } catch (ExecutionException e) {
35
            return "{\"error\":\"TOOL_FAILED\"}";
36
        } catch (InterruptedException e) {
37
            Thread.currentThread().interrupt();
38
            return "{\"error\":\"TOOL_INTERRUPTED\"}";
39
        } catch (Exception e) {
40
            return "{\"error\":\"TOOL_CALL_INVALID\"}";
41
        }
42
    }
43
}

1
package in.o612.eng.opsagent.agent.orchestration;
2

3
import in.o612.eng.opsagent.agent.model.ChatClientGateway;
4
import in.o612.eng.opsagent.agent.tools.McpToolExecutor;
5
import in.o612.eng.opsagent.agent.tools.ToolPolicy;
6
import in.o612.eng.opsagent.agent.tools.ToolPolicyRegistry;
7
import in.o612.eng.opsagent.contracts.conversation.AgentEvent;
8
import in.o612.eng.opsagent.agent.conversation.ConversationEventBus;
9
import org.springframework.ai.chat.client.ChatClient;
10
import org.springframework.ai.chat.messages.AssistantMessage;
11
import org.springframework.ai.chat.messages.Message;
12
import org.springframework.ai.chat.messages.ToolResponseMessage;
13
import org.springframework.ai.chat.messages.UserMessage;
14
import org.springframework.ai.chat.prompt.Prompt;
15
import org.springframework.stereotype.Service;
16

17
import java.time.Duration;
18
import java.time.Instant;
19
import java.util.ArrayList;
20
import java.util.List;
21
import java.util.Map;
22

23
@Service
24
public class AgentOrchestrator {
25

26
    private final ChatClient chatClient;
27
    private final ToolPolicyRegistry policies;
28
    private final McpToolExecutor executor;
29
    private final ConversationEventBus events;
30
    private final AgentBudget budget = new AgentBudget(8, 10, Duration.ofSeconds(45));
31

32
    public record Outcome(String answer, int modelCalls, int toolCalls) {}
33

34
    public Outcome run(String conversationId, String tenantId, Set<String> callerScopes,
35
                       String systemPrompt, String userMessage) {
36
        var messages = new ArrayList<Message>();
37
        messages.add(new UserMessage(userMessage));
38
        var started = Instant.now();
39
        int toolCalls = 0;
40

41
        for (int modelCalls = 1; modelCalls <= budget.maxModelCalls(); modelCalls++) {
42
            if (Duration.between(started, Instant.now()).compareTo(budget.deadline()) > 0) {
43
                return failed(conversationId, "DEADLINE_EXCEEDED");
44
            }
45
            var response = chatClient.prompt()
46
                    .system(systemPrompt)
47
                    .messages(messages)
48
                    .call()
49
                    .chatResponse();
50

51
            var output = response.getResult().getOutput();
52
            var calls = output.getToolCalls();
53
            if (calls == null || calls.isEmpty()) {
54
                return new Outcome(output.getText(), modelCalls, toolCalls);
55
            }
56

57
            var toolResponses = new ArrayList<ToolResponseMessage.ToolResponse>();
58
            for (var call : calls) {
59
                if (++toolCalls > budget.maxToolCalls()) {
60
                    return failed(conversationId, "TOOL_BUDGET_EXCEEDED");
61
                }
62
                events.publish(conversationId, new AgentEvent.ToolProposed(
63
                        conversationId, call.name(), call.arguments(), Instant.now()));
64

65
                toolResponses.add(new ToolResponseMessage.ToolResponse(
66
                        call.id(), call.name(), dispatch(call, tenantId, callerScopes)));
67
            }
68
            messages.add(new AssistantMessage(output.getText(), Map.of(), calls));
69
            messages.add(new ToolResponseMessage(toolResponses));
70
        }
71
        return failed(conversationId, "MODEL_BUDGET_EXCEEDED");
72
    }
73

74
    private String dispatch(AssistantMessage.ToolCall call, String tenantId, Set<String> scopes) {
75
        var policy = policies.find(call.name());
76
        if (policy.isEmpty()) {
77
            return "{\"error\":\"UNKNOWN_TOOL\"}";
78
        }
79
        var p = policy.get();
80
        if (p.requiredScope() != null && !scopes.contains(p.requiredScope())) {
81
            return "{\"error\":\"INSUFFICIENT_SCOPE\"}";
82
        }
83
        if (p.requiresApproval()) {
84
            // Chapter 10 replaces this with the real approval flow
85
            return "{\"error\":\"APPROVAL_REQUIRED\"}";
86
        }
87
        Map<String, Object> args = parseArgs(call.arguments());
88
        return executor.execute(call.name(), args, tenantId, p.perCallTimeout());
89
    }
90

91
    // failed(...) publishes AgentEvent.Failed and returns an Outcome with a safe message;
92
    // parseArgs(...) maps the tool-call JSON arguments string to Map<String,Object>
93
}