Skip to content

Commit 291b083

Browse files
EItanyaEItanya
authored
feat: expose checkpoint workflows through MCP (#2589)
## Summary - add synchronous MCP tools to create and list checkpoints and fork AgentInstances - route every operation through the existing checkpoint service and preserve caller request IDs - cover schemas and translations with unit tests and checkpoint-to-fork invocation with E2E ## Testing - go test ./core/v2/mcp ./core/cmd/controller-v2 - go test ./core/v2/checkpoint - go test ./core/test/e2e -run '^$' - make -C go lint - TestMCPCheckpointFork against local Kind cluster Stacked on #2576. Closes #2577 Signed-off-by: Eitan Yarmush <eitan.yarmush@solo.io>
1 parent 2b11fb9 commit 291b083

6 files changed

Lines changed: 272 additions & 16 deletions

File tree

‎go/core/cmd/controller-v2/main.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,7 @@ func main() {
130130
}
131131
gateway := a2agateway.New(store, authorizer, gatewayDialer, instanceWorkflow,
132132
env("A2A_GATEWAY_URL", "http://127.0.0.1:8084"))
133-
mcpHandler, err := v2mcp.New(instances, gateway)
133+
mcpHandler, err := v2mcp.New(instances, checkpoints, gateway)
134134
if err != nil {
135135
log.Fatal(err)
136136
}

‎go/core/test/e2e/mcp_test.go‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package e2e_test
22

33
import (
44
"bytes"
5+
"context"
56
"encoding/json"
67
"io"
78
"net"
@@ -12,7 +13,11 @@ import (
1213

1314
a2atype "github.com/a2aproject/a2a-go/v2/a2a"
1415
"github.com/a2aproject/a2a-go/v2/a2apb/v1/pbconv"
16+
apiv1alpha1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1alpha1"
1517
"github.com/modelcontextprotocol/go-sdk/mcp"
18+
"google.golang.org/grpc/codes"
19+
"google.golang.org/grpc/metadata"
20+
"google.golang.org/grpc/status"
1621
)
1722

1823
const (
@@ -103,6 +108,54 @@ func TestMCPCancelTask(t *testing.T) {
103108
waitMCPTask(t, endpoint, handle, "cancelled")
104109
}
105110

111+
func TestMCPCheckpointFork(t *testing.T) {
112+
fixture := newInteractionFixture(t, interactionTarget(t), startInteractionMock(t))
113+
endpoint := mcpEndpoint(t)
114+
if result := mcpInvoke(t, endpoint, fixture.instanceID, "What is 2+2?", false); result["resultType"] != "complete" {
115+
t.Fatalf("initial invocation = %#v", result)
116+
}
117+
118+
created := mcpCall(t, endpoint, "tools/call", map[string]any{
119+
"name": "create_agent_instance_checkpoint",
120+
"arguments": map[string]any{"namespace": "kagent", "agent_instance_id": fixture.instanceID},
121+
}, false)["result"].(map[string]any)["structuredContent"].(map[string]any)["checkpoint"].(map[string]any)
122+
checkpointID := created["id"].(string)
123+
t.Cleanup(func() {
124+
ctx, cancel := context.WithTimeout(metadata.AppendToOutgoingContext(context.Background(), "x-user-id", "e2e"), time.Minute)
125+
defer cancel()
126+
_, err := fixture.checkpoints.DeleteCheckpoint(ctx, &apiv1alpha1.DeleteCheckpointRequest{Namespace: "kagent", CheckpointId: checkpointID})
127+
if err != nil && status.Code(err) != codes.NotFound {
128+
t.Errorf("delete checkpoint: %v", err)
129+
}
130+
})
131+
132+
listed := mcpCall(t, endpoint, "tools/call", map[string]any{
133+
"name": "list_agent_instance_checkpoints",
134+
"arguments": map[string]any{"namespace": "kagent", "agent_instance_id": fixture.instanceID},
135+
}, false)["result"].(map[string]any)["structuredContent"].(map[string]any)["checkpoints"].([]any)
136+
if len(listed) != 1 || listed[0].(map[string]any)["id"] != checkpointID {
137+
t.Fatalf("listed checkpoints = %#v", listed)
138+
}
139+
140+
forked := mcpCall(t, endpoint, "tools/call", map[string]any{
141+
"name": "fork_agent_instance",
142+
"arguments": map[string]any{"namespace": "kagent", "checkpoint_id": checkpointID},
143+
}, false)["result"].(map[string]any)["structuredContent"].(map[string]any)["agent_instance"].(map[string]any)
144+
forkID := forked["id"].(string)
145+
t.Cleanup(func() {
146+
ctx, cancel := context.WithTimeout(metadata.AppendToOutgoingContext(context.Background(), "x-user-id", "e2e"), time.Minute)
147+
defer cancel()
148+
_, err := fixture.instances.DeleteAgentInstance(ctx, &apiv1alpha1.DeleteAgentInstanceRequest{Namespace: "kagent", AgentInstanceId: forkID})
149+
if err != nil && status.Code(err) != codes.NotFound {
150+
t.Errorf("delete fork AgentInstance: %v", err)
151+
}
152+
})
153+
154+
if result := mcpInvoke(t, endpoint, forkID, "What is 2+2?", false); !strings.Contains(mcpResultText(result), "The answer is 4.") {
155+
t.Fatalf("fork invocation = %#v", result)
156+
}
157+
}
158+
106159
func mcpEndpoint(t *testing.T) string {
107160
t.Helper()
108161
host, _, err := net.SplitHostPort(interactionTarget(t))

‎go/core/v2/mcp/checkpoints.go‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
package mcp
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"time"
7+
8+
"github.com/google/uuid"
9+
apiv1alpha1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1alpha1"
10+
"github.com/kagent-dev/kagent/go/core/v2/checkpoint"
11+
"github.com/modelcontextprotocol/go-sdk/mcp"
12+
)
13+
14+
const (
15+
createCheckpointToolName = "create_agent_instance_checkpoint"
16+
listCheckpointsToolName = "list_agent_instance_checkpoints"
17+
forkAgentInstanceToolName = "fork_agent_instance"
18+
)
19+
20+
type CreateCheckpointInput struct {
21+
Namespace string `json:"namespace" jsonschema:"Kubernetes namespace containing the AgentInstance"`
22+
AgentInstanceID string `json:"agent_instance_id" jsonschema:"AgentInstance UUID"`
23+
RequestID string `json:"request_id,omitempty" jsonschema:"Optional stable request ID for idempotency"`
24+
}
25+
26+
type CheckpointSummary struct {
27+
ID string `json:"id"`
28+
Namespace string `json:"namespace"`
29+
AgentInstanceID string `json:"agent_instance_id"`
30+
HeadTaskID string `json:"head_task_id,omitempty"`
31+
HistorySequence uint64 `json:"history_sequence"`
32+
State string `json:"state"`
33+
CreatedAt string `json:"created_at,omitempty"`
34+
Failure *FailureSummary `json:"failure,omitempty"`
35+
}
36+
37+
type FailureSummary struct {
38+
Reason string `json:"reason"`
39+
Message string `json:"message"`
40+
}
41+
42+
type CreateCheckpointOutput struct {
43+
Checkpoint CheckpointSummary `json:"checkpoint"`
44+
}
45+
46+
type ListCheckpointsInput struct {
47+
Namespace string `json:"namespace" jsonschema:"Kubernetes namespace containing the AgentInstance"`
48+
AgentInstanceID string `json:"agent_instance_id" jsonschema:"AgentInstance UUID"`
49+
PageSize int `json:"page_size,omitempty" jsonschema:"Maximum number of checkpoints to return"`
50+
PageToken string `json:"page_token,omitempty" jsonschema:"Token returned by a previous call"`
51+
}
52+
53+
type ListCheckpointsOutput struct {
54+
Checkpoints []CheckpointSummary `json:"checkpoints"`
55+
NextPageToken string `json:"next_page_token,omitempty"`
56+
}
57+
58+
type ForkAgentInstanceInput struct {
59+
Namespace string `json:"namespace" jsonschema:"Kubernetes namespace containing the checkpoint"`
60+
CheckpointID string `json:"checkpoint_id" jsonschema:"Checkpoint UUID"`
61+
RequestID string `json:"request_id,omitempty" jsonschema:"Optional stable request ID for idempotency"`
62+
}
63+
64+
type ForkAgentInstanceOutput struct {
65+
AgentInstance AgentInstanceSummary `json:"agent_instance"`
66+
}
67+
68+
func (h *Handler) registerCheckpointTools(server *mcp.Server) {
69+
mcp.AddTool(server, &mcp.Tool{Name: createCheckpointToolName, Description: "Create a checkpoint at an AgentInstance turn boundary"}, h.createCheckpoint)
70+
mcp.AddTool(server, &mcp.Tool{Name: listCheckpointsToolName, Description: "List checkpoints for an AgentInstance"}, h.listCheckpoints)
71+
mcp.AddTool(server, &mcp.Tool{Name: forkAgentInstanceToolName, Description: "Create an AgentInstance from a checkpoint"}, h.forkAgentInstance)
72+
}
73+
74+
func (h *Handler) createCheckpoint(ctx context.Context, _ *mcp.CallToolRequest, input CreateCheckpointInput) (*mcp.CallToolResult, CreateCheckpointOutput, error) {
75+
created, err := h.checkpoints.Create(ctx, input.Namespace, input.AgentInstanceID, stableRequestID(input.RequestID))
76+
if err != nil {
77+
return toolError(err), CreateCheckpointOutput{}, nil
78+
}
79+
output := CreateCheckpointOutput{Checkpoint: checkpointSummary(created)}
80+
return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf("Created checkpoint %s", created.GetId())}}}, output, nil
81+
}
82+
83+
func (h *Handler) listCheckpoints(ctx context.Context, _ *mcp.CallToolRequest, input ListCheckpointsInput) (*mcp.CallToolResult, ListCheckpointsOutput, error) {
84+
listed, err := h.checkpoints.List(ctx, checkpoint.ListRequest{
85+
Namespace: input.Namespace, InstanceID: input.AgentInstanceID,
86+
PageSize: input.PageSize, PageToken: input.PageToken,
87+
})
88+
if err != nil {
89+
return toolError(err), ListCheckpointsOutput{}, nil
90+
}
91+
output := ListCheckpointsOutput{Checkpoints: make([]CheckpointSummary, len(listed.Checkpoints)), NextPageToken: listed.NextPageToken}
92+
for i, item := range listed.Checkpoints {
93+
output.Checkpoints[i] = checkpointSummary(item)
94+
}
95+
return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf("Found %d checkpoints", len(output.Checkpoints))}}}, output, nil
96+
}
97+
98+
func (h *Handler) forkAgentInstance(ctx context.Context, _ *mcp.CallToolRequest, input ForkAgentInstanceInput) (*mcp.CallToolResult, ForkAgentInstanceOutput, error) {
99+
instance, err := h.checkpoints.Fork(ctx, input.Namespace, input.CheckpointID, stableRequestID(input.RequestID))
100+
if err != nil {
101+
return toolError(err), ForkAgentInstanceOutput{}, nil
102+
}
103+
output := ForkAgentInstanceOutput{AgentInstance: agentInstanceSummary(instance)}
104+
return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf("Created AgentInstance %s", instance.GetId())}}}, output, nil
105+
}
106+
107+
func stableRequestID(id string) string {
108+
if id != "" {
109+
return id
110+
}
111+
return uuid.NewString()
112+
}
113+
114+
func checkpointSummary(value *apiv1alpha1.Checkpoint) CheckpointSummary {
115+
result := CheckpointSummary{
116+
ID: value.GetId(), Namespace: value.GetNamespace(), AgentInstanceID: value.GetAgentInstanceId(),
117+
HeadTaskID: value.GetHeadTaskId(), HistorySequence: value.GetHistorySequence(), State: value.GetState().String(),
118+
}
119+
if value.GetCreatedAt() != nil {
120+
result.CreatedAt = value.GetCreatedAt().AsTime().Format(time.RFC3339Nano)
121+
}
122+
if value.GetFailure() != nil {
123+
result.Failure = &FailureSummary{Reason: value.GetFailure().GetReason(), Message: value.GetFailure().GetMessage()}
124+
}
125+
return result
126+
}
127+
128+
func agentInstanceSummary(instance *apiv1alpha1.AgentInstance) AgentInstanceSummary {
129+
return AgentInstanceSummary{
130+
Namespace: instance.GetNamespace(), ID: instance.GetId(),
131+
AgentTemplate: instance.GetAgentTemplate().GetName(), Harness: instance.GetHarness().GetName(),
132+
State: instance.GetState().String(),
133+
}
134+
}

‎go/core/v2/mcp/checkpoints_test.go‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
package mcp
2+
3+
import (
4+
"net/http"
5+
"net/http/httptest"
6+
"testing"
7+
"time"
8+
9+
"github.com/a2aproject/a2a-go/v2/a2asrv"
10+
apiv1alpha1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1alpha1"
11+
"google.golang.org/protobuf/types/known/timestamppb"
12+
)
13+
14+
func TestCheckpointSummary(t *testing.T) {
15+
created := time.Date(2026, time.August, 26, 10, 0, 0, 123, time.UTC)
16+
got := checkpointSummary(&apiv1alpha1.Checkpoint{
17+
Id: "22222222-2222-4222-8222-222222222222", Namespace: "team-a",
18+
AgentInstanceId: testInstanceID, HeadTaskId: testTaskID, HistorySequence: 7,
19+
State: apiv1alpha1.CheckpointState_CHECKPOINT_STATE_READY, CreatedAt: timestamppb.New(created),
20+
Failure: &apiv1alpha1.Failure{Message: "failed"},
21+
})
22+
if got.ID != "22222222-2222-4222-8222-222222222222" || got.AgentInstanceID != testInstanceID ||
23+
got.HistorySequence != 7 || got.State != "CHECKPOINT_STATE_READY" || got.CreatedAt != created.Format(time.RFC3339Nano) || got.Failure.Message != "failed" {
24+
t.Fatalf("checkpointSummary() = %#v", got)
25+
}
26+
}
27+
28+
func TestCheckpointToolsAreRegistered(t *testing.T) {
29+
h, err := New(testAgentInstanceService(), testCheckpointService(), &a2asrv.InterceptedHandler{Handler: &fakeGateway{}})
30+
if err != nil {
31+
t.Fatal(err)
32+
}
33+
server := httptest.NewServer(http.HandlerFunc(h.ServeHTTP))
34+
defer server.Close()
35+
response := rawMCPCall(t, server.URL, "tools/list", map[string]any{}, false)
36+
tools := response["result"].(map[string]any)["tools"].([]any)
37+
want := map[string]bool{createCheckpointToolName: false, listCheckpointsToolName: false, forkAgentInstanceToolName: false}
38+
for _, value := range tools {
39+
tool := value.(map[string]any)
40+
if _, ok := want[tool["name"].(string)]; ok {
41+
want[tool["name"].(string)] = len(tool["inputSchema"].(map[string]any)["properties"].(map[string]any)) > 0
42+
}
43+
}
44+
for name, valid := range want {
45+
if !valid {
46+
t.Fatalf("tool %q missing or has no input schema: %#v", name, tools)
47+
}
48+
}
49+
}
50+
51+
func TestCheckpointToolErrorsAreToolResults(t *testing.T) {
52+
h := &Handler{checkpoints: testCheckpointService()}
53+
result, _, err := h.createCheckpoint(t.Context(), nil, CreateCheckpointInput{Namespace: "team-a", AgentInstanceID: "invalid"})
54+
if err != nil || !result.IsError {
55+
t.Fatalf("createCheckpoint() result = %#v, error = %v", result, err)
56+
}
57+
}
58+
59+
func TestStableRequestID(t *testing.T) {
60+
if got := stableRequestID("caller-id"); got != "caller-id" {
61+
t.Fatalf("stableRequestID() = %q", got)
62+
}
63+
if got := stableRequestID(""); got == "" {
64+
t.Fatal("stableRequestID() returned an empty generated ID")
65+
}
66+
}

‎go/core/v2/mcp/server.go‎

Lines changed: 11 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import (
1717
"github.com/kagent-dev/kagent/go/core/internal/version"
1818
"github.com/kagent-dev/kagent/go/core/v2/a2agateway"
1919
"github.com/kagent-dev/kagent/go/core/v2/agentinstance"
20+
"github.com/kagent-dev/kagent/go/core/v2/checkpoint"
2021
"github.com/modelcontextprotocol/go-sdk/mcp"
2122
"google.golang.org/grpc/metadata"
2223
)
@@ -28,9 +29,10 @@ const (
2829
)
2930

3031
type Handler struct {
31-
instances *agentinstance.Service
32-
gateway a2asrv.RequestHandler
33-
http http.Handler
32+
instances *agentinstance.Service
33+
checkpoints *checkpoint.Service
34+
gateway a2asrv.RequestHandler
35+
http http.Handler
3436
}
3537

3638
type invocationStart struct {
@@ -74,11 +76,11 @@ type InvokeAgentInstanceOutput struct {
7476
Text string `json:"text,omitempty"`
7577
}
7678

77-
func New(instances *agentinstance.Service, gateway a2asrv.RequestHandler) (*Handler, error) {
78-
if instances == nil || gateway == nil {
79-
return nil, fmt.Errorf("AgentInstance service and A2A gateway are required")
79+
func New(instances *agentinstance.Service, checkpoints *checkpoint.Service, gateway a2asrv.RequestHandler) (*Handler, error) {
80+
if instances == nil || checkpoints == nil || gateway == nil {
81+
return nil, fmt.Errorf("AgentInstance service, checkpoint service, and A2A gateway are required")
8082
}
81-
h := &Handler{instances: instances, gateway: gateway}
83+
h := &Handler{instances: instances, checkpoints: checkpoints, gateway: gateway}
8284
capabilities := &mcp.ServerCapabilities{}
8385
capabilities.AddExtension(tasksExtension, nil)
8486
server := mcp.NewServer(
@@ -93,6 +95,7 @@ func New(instances *agentinstance.Service, gateway a2asrv.RequestHandler) (*Hand
9395
Name: invokeToolName,
9496
Description: "Invoke an AgentInstance through the public A2A gateway",
9597
}, h.invokeAgentInstance)
98+
h.registerCheckpointTools(server)
9699
server.AddReceivingMiddleware(h.taskAwareToolCall)
97100
if err := h.registerTaskMethods(server); err != nil {
98101
return nil, err
@@ -118,12 +121,7 @@ func (h *Handler) listAgentInstances(ctx context.Context, _ *mcp.CallToolRequest
118121
if instance.GetState() != apiv1alpha1.AgentInstanceState_AGENT_INSTANCE_STATE_READY {
119122
continue
120123
}
121-
output.AgentInstances = append(output.AgentInstances, AgentInstanceSummary{
122-
Namespace: instance.GetNamespace(), ID: instance.GetId(),
123-
AgentTemplate: instance.GetAgentTemplate().GetName(),
124-
Harness: instance.GetHarness().GetName(),
125-
State: instance.GetState().String(),
126-
})
124+
output.AgentInstances = append(output.AgentInstances, agentInstanceSummary(instance))
127125
}
128126
var text strings.Builder
129127
for i, instance := range output.AgentInstances {

‎go/core/v2/mcp/tasks_test.go‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
"github.com/kagent-dev/kagent/go/core/pkg/auth"
2222
"github.com/kagent-dev/kagent/go/core/v2/a2agateway"
2323
"github.com/kagent-dev/kagent/go/core/v2/agentinstance"
24+
"github.com/kagent-dev/kagent/go/core/v2/checkpoint"
2425
"github.com/modelcontextprotocol/go-sdk/jsonrpc"
2526
"github.com/modelcontextprotocol/go-sdk/mcp"
2627
)
@@ -183,7 +184,7 @@ func TestTaskUpdateTranslatesAskUserResponse(t *testing.T) {
183184

184185
func TestTaskCapableToolCallReturnsDurableHandle(t *testing.T) {
185186
gateway := &fakeGateway{}
186-
h, err := New(testAgentInstanceService(), &a2asrv.InterceptedHandler{Handler: gateway})
187+
h, err := New(testAgentInstanceService(), testCheckpointService(), &a2asrv.InterceptedHandler{Handler: gateway})
187188
if err != nil {
188189
t.Fatal(err)
189190
}
@@ -223,7 +224,7 @@ func TestTaskCapableToolCallReturnsDurableHandle(t *testing.T) {
223224

224225
func TestToolCallWithoutTasksWaitsForResult(t *testing.T) {
225226
gateway := &fakeGateway{completeOnDrain: true}
226-
h, err := New(testAgentInstanceService(), &a2asrv.InterceptedHandler{Handler: gateway})
227+
h, err := New(testAgentInstanceService(), testCheckpointService(), &a2asrv.InterceptedHandler{Handler: gateway})
227228
if err != nil {
228229
t.Fatal(err)
229230
}
@@ -492,3 +493,7 @@ func (*fakeInstanceWorkflow) Delete(_ context.Context, instance *apiv1alpha1.Age
492493
func testAgentInstanceService() *agentinstance.Service {
493494
return agentinstance.NewService(&fakeInstanceStore{}, &authimpl.NoopAuthorizer{}, &fakeInstanceWorkflow{})
494495
}
496+
497+
func testCheckpointService() *checkpoint.Service {
498+
return checkpoint.NewService(nil, nil, nil, nil)
499+
}

0 commit comments

Comments
 (0)