chore: import upstream snapshot with attribution
govulncheck / govulncheck (push) Has been cancelled
Lint / golangci-lint (push) Has been cancelled
Run Tests / Unit Tests (push) Has been cancelled
Run Tests / Etcd Integration Tests (push) Has been cancelled
Harness (E2E) / Harnesses (mock LLM) (push) Has been cancelled
Harness (E2E) / Provider harnesses (live LLM conformance) (push) Has been cancelled
govulncheck / govulncheck (push) Has been cancelled
Lint / golangci-lint (push) Has been cancelled
Run Tests / Unit Tests (push) Has been cancelled
Run Tests / Etcd Integration Tests (push) Has been cancelled
Harness (E2E) / Harnesses (mock LLM) (push) Has been cancelled
Harness (E2E) / Provider harnesses (live LLM conformance) (push) Has been cancelled
This commit is contained in:
wehub-resource-sync
+1123
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,899 @@
|
||||
package a2a
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
pb "go-micro.dev/v6/agent/proto"
|
||||
"go-micro.dev/v6/ai"
|
||||
"go-micro.dev/v6/client"
|
||||
"go-micro.dev/v6/registry"
|
||||
"go-micro.dev/v6/selector"
|
||||
"go-micro.dev/v6/server"
|
||||
)
|
||||
|
||||
// echoAgent is a stub that implements the Agent proto handler — enough to
|
||||
// exercise the gateway's task→Agent.Chat translation without pulling in
|
||||
// the agent package (which would import this one, a test-only cycle).
|
||||
type echoAgent struct{}
|
||||
|
||||
func (echoAgent) Chat(_ context.Context, req *pb.ChatRequest, rsp *pb.ChatResponse) error {
|
||||
rsp.Reply = "pong"
|
||||
rsp.Agent = "echo"
|
||||
return nil
|
||||
}
|
||||
|
||||
func waitFor(reg registry.Registry, name string) {
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if svcs, err := reg.GetService(name); err == nil && len(svcs) > 0 && len(svcs[0].Nodes) > 0 {
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
func newGatewayWithAgent(t *testing.T) (*httptest.Server, func()) {
|
||||
t.Helper()
|
||||
reg := registry.NewMemoryRegistry()
|
||||
cl := client.NewClient(client.Registry(reg), client.Selector(selector.NewSelector(selector.Registry(reg))))
|
||||
|
||||
srv := server.NewServer(
|
||||
server.Name("echo"),
|
||||
server.Address("127.0.0.1:0"),
|
||||
server.Registry(reg),
|
||||
server.Metadata(map[string]string{"type": "agent", "services": "task,project"}),
|
||||
)
|
||||
if err := pb.RegisterAgentHandler(srv, echoAgent{}); err != nil {
|
||||
t.Fatalf("register agent handler: %v", err)
|
||||
}
|
||||
if err := srv.Start(); err != nil {
|
||||
t.Fatalf("start server: %v", err)
|
||||
}
|
||||
waitFor(reg, "echo")
|
||||
|
||||
g := New(Options{Registry: reg, Client: cl, BaseURL: "http://gw"})
|
||||
ts := httptest.NewServer(g.Handler())
|
||||
return ts, func() { ts.Close(); srv.Stop() }
|
||||
}
|
||||
|
||||
func TestAgentCardFromRegistry(t *testing.T) {
|
||||
ts, cleanup := newGatewayWithAgent(t)
|
||||
defer cleanup()
|
||||
|
||||
resp, err := http.Get(ts.URL + "/agents/echo/.well-known/agent.json")
|
||||
if err != nil {
|
||||
t.Fatalf("get card: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("card status = %d", resp.StatusCode)
|
||||
}
|
||||
var card AgentCard
|
||||
if err := json.NewDecoder(resp.Body).Decode(&card); err != nil {
|
||||
t.Fatalf("decode card: %v", err)
|
||||
}
|
||||
if card.Name != "echo" {
|
||||
t.Errorf("card name = %q, want echo", card.Name)
|
||||
}
|
||||
if card.URL != "http://gw/agents/echo" {
|
||||
t.Errorf("card url = %q", card.URL)
|
||||
}
|
||||
if card.ProtocolVersion == "" {
|
||||
t.Errorf("card missing protocolVersion: %+v", card)
|
||||
}
|
||||
if !card.Capabilities.TaskResubscribe || !card.Capabilities.InputRequired {
|
||||
t.Errorf("card capabilities = %+v, want task resubscribe and input-required advertised", card.Capabilities)
|
||||
}
|
||||
if got := skillIDs(card.Skills); strings.Join(got, ",") != "task,project" {
|
||||
t.Errorf("skill IDs = %v, want [task project]", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A2A 0.3.0 discovery is /.well-known/agent-card.json. The card must be
|
||||
// reachable there (canonical) as well as at the legacy agent.json alias, both
|
||||
// per-agent and at the single-agent top level.
|
||||
func TestAgentCardCanonicalWellKnownPath(t *testing.T) {
|
||||
ts, cleanup := newGatewayWithAgent(t)
|
||||
defer cleanup()
|
||||
|
||||
for _, path := range []string{
|
||||
"/agents/echo/.well-known/agent-card.json",
|
||||
"/agents/echo/.well-known/agent.json",
|
||||
"/agents/echo/skills/task/.well-known/agent-card.json",
|
||||
} {
|
||||
resp, err := http.Get(ts.URL + path)
|
||||
if err != nil {
|
||||
t.Fatalf("get %s: %v", path, err)
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
resp.Body.Close()
|
||||
t.Fatalf("%s status = %d, want 200", path, resp.StatusCode)
|
||||
}
|
||||
var card AgentCard
|
||||
if err := json.NewDecoder(resp.Body).Decode(&card); err != nil {
|
||||
resp.Body.Close()
|
||||
t.Fatalf("%s decode card: %v", path, err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
if card.Name != "echo" {
|
||||
t.Errorf("%s card name = %q, want echo", path, card.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSkillEndpointServesFocusedCardAndRoutesRPC(t *testing.T) {
|
||||
ts, cleanup := newGatewayWithAgent(t)
|
||||
defer cleanup()
|
||||
|
||||
resp, err := http.Get(ts.URL + "/agents/echo/skills/task/.well-known/agent.json")
|
||||
if err != nil {
|
||||
t.Fatalf("get skill card: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("skill card status = %d", resp.StatusCode)
|
||||
}
|
||||
var card AgentCard
|
||||
if err := json.NewDecoder(resp.Body).Decode(&card); err != nil {
|
||||
t.Fatalf("decode skill card: %v", err)
|
||||
}
|
||||
if card.URL != "http://gw/agents/echo/skills/task" || len(card.Skills) != 1 || card.Skills[0].ID != "task" {
|
||||
t.Fatalf("skill card = %+v, want task-only card at skill URL", card)
|
||||
}
|
||||
|
||||
task := rpcTask(t, ts.URL+"/agents/echo/skills/task", `{
|
||||
"jsonrpc":"2.0","id":1,"method":"message/send",
|
||||
"params":{"message":{"role":"user","kind":"message","messageId":"m1",
|
||||
"parts":[{"kind":"text","text":"ping"}]}}}`)
|
||||
if task.Status.State != stateCompleted || textOf(task.Artifacts[0].Parts) != "pong" {
|
||||
t.Fatalf("skill task = %+v, want completed pong", task)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessageSendAndGet(t *testing.T) {
|
||||
ts, cleanup := newGatewayWithAgent(t)
|
||||
defer cleanup()
|
||||
|
||||
task := rpcTask(t, ts.URL+"/agents/echo", `{
|
||||
"jsonrpc":"2.0","id":1,"method":"message/send",
|
||||
"params":{"message":{"role":"user","kind":"message","messageId":"m1",
|
||||
"parts":[{"kind":"text","text":"ping"}]}}}`)
|
||||
if task.Status.State != stateCompleted {
|
||||
t.Fatalf("task state = %q, want completed", task.Status.State)
|
||||
}
|
||||
if len(task.Artifacts) != 1 || textOf(task.Artifacts[0].Parts) != "pong" {
|
||||
t.Fatalf("artifact = %+v, want text 'pong'", task.Artifacts)
|
||||
}
|
||||
if len(task.History) != 2 || task.History[1].Role != "agent" || textOf(task.History[1].Parts) != "pong" {
|
||||
t.Fatalf("history = %+v, want user turn followed by agent reply", task.History)
|
||||
}
|
||||
if task.History[1].TaskID != task.ID || task.History[1].ContextID != task.ContextID {
|
||||
t.Fatalf("agent history linkage = task %q/%q context %q/%q", task.History[1].TaskID, task.ID, task.History[1].ContextID, task.ContextID)
|
||||
}
|
||||
|
||||
got := rpcTask(t, ts.URL+"/agents/echo", `{
|
||||
"jsonrpc":"2.0","id":2,"method":"tasks/get","params":{"id":"`+task.ID+`"}}`)
|
||||
if got.ID != task.ID || got.Status.State != stateCompleted {
|
||||
t.Errorf("tasks/get returned %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessageSendContinuesExistingTask(t *testing.T) {
|
||||
d := newDispatcher()
|
||||
first := rpcTaskFromBody(t, d, `{
|
||||
"jsonrpc":"2.0","id":1,"method":"message/send",
|
||||
"params":{"message":{"role":"user","kind":"message","messageId":"m1",
|
||||
"parts":[{"kind":"text","text":"first"}]}}}`, func(_ context.Context, text string) (string, error) {
|
||||
return "reply to " + text, nil
|
||||
})
|
||||
|
||||
secondBody := fmt.Sprintf(`{
|
||||
"jsonrpc":"2.0","id":2,"method":"message/send",
|
||||
"params":{"message":{"role":"user","kind":"message","messageId":"m2","taskId":"%s","contextId":"%s",
|
||||
"parts":[{"kind":"text","text":"second"}]}}}`, first.ID, first.ContextID)
|
||||
second := rpcTaskFromBody(t, d, secondBody, func(_ context.Context, text string) (string, error) {
|
||||
return "reply to " + text, nil
|
||||
})
|
||||
|
||||
if second.ID != first.ID || second.ContextID != first.ContextID {
|
||||
t.Fatalf("continued task identity = %s/%s, want %s/%s", second.ID, second.ContextID, first.ID, first.ContextID)
|
||||
}
|
||||
if len(second.History) != 4 {
|
||||
t.Fatalf("continued history len = %d, want 4: %+v", len(second.History), second.History)
|
||||
}
|
||||
if textOf(second.History[0].Parts) != "first" || textOf(second.History[2].Parts) != "second" {
|
||||
t.Fatalf("continued history did not preserve turns: %+v", second.History)
|
||||
}
|
||||
|
||||
got := rpcTaskFromDispatcher(t, d, first.ID)
|
||||
if got.ID != first.ID || len(got.History) != 4 {
|
||||
t.Fatalf("stored continued task = %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPushNotificationConfigDeliversTaskUpdates(t *testing.T) {
|
||||
d := newDispatcher()
|
||||
updates := make(chan Task, 2)
|
||||
push := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if got := r.Header.Get("Authorization"); got != "Bearer secret" {
|
||||
t.Errorf("authorization = %q, want bearer token", got)
|
||||
}
|
||||
var task Task
|
||||
if err := json.NewDecoder(r.Body).Decode(&task); err != nil {
|
||||
t.Errorf("decode push task: %v", err)
|
||||
return
|
||||
}
|
||||
updates <- task
|
||||
w.WriteHeader(http.StatusAccepted)
|
||||
}))
|
||||
defer push.Close()
|
||||
|
||||
task := rpcTaskFromBody(t, d, `{
|
||||
"jsonrpc":"2.0","id":1,"method":"message/send",
|
||||
"params":{"message":{"role":"user","kind":"message","messageId":"m1",
|
||||
"parts":[{"kind":"text","text":"ping"}]}}}`, func(_ context.Context, text string) (string, error) {
|
||||
return "pong", nil
|
||||
})
|
||||
|
||||
body := fmt.Sprintf(`{"jsonrpc":"2.0","id":2,"method":"tasks/pushNotificationConfig/set","params":{"id":"%s","pushNotificationConfig":{"url":"%s","token":"secret"}}}`, task.ID, push.URL)
|
||||
var setResp struct {
|
||||
Result struct {
|
||||
ID string `json:"id"`
|
||||
PushNotificationConfig PushNotificationConfig `json:"pushNotificationConfig"`
|
||||
} `json:"result"`
|
||||
Error *rpcError `json:"error"`
|
||||
}
|
||||
rpcDispatcher(t, d, body, nil, &setResp)
|
||||
if setResp.Error != nil {
|
||||
t.Fatalf("set push config error: %+v", setResp.Error)
|
||||
}
|
||||
if setResp.Result.ID != task.ID || setResp.Result.PushNotificationConfig.URL != push.URL {
|
||||
t.Fatalf("set push config result = %+v", setResp.Result)
|
||||
}
|
||||
|
||||
select {
|
||||
case got := <-updates:
|
||||
if got.ID != task.ID || got.Status.State != stateCompleted || textOf(got.Artifacts[0].Parts) != "pong" {
|
||||
t.Fatalf("push update = %+v, want completed task", got)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("timed out waiting for push update")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessageSendUsesRequestContext(t *testing.T) {
|
||||
d := newDispatcher()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
req := httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString(`{
|
||||
"jsonrpc":"2.0","id":1,"method":"message/send",
|
||||
"params":{"message":{"role":"user","kind":"message","messageId":"m1",
|
||||
"parts":[{"kind":"text","text":"ping"}]}}}`))
|
||||
req = req.WithContext(ctx)
|
||||
rr := httptest.NewRecorder()
|
||||
|
||||
d.serve(rr, req, func(ctx context.Context, text string) (string, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return "unexpected success", nil
|
||||
})
|
||||
|
||||
var resp struct {
|
||||
Result Task `json:"result"`
|
||||
Error *rpcError `json:"error"`
|
||||
}
|
||||
if err := json.NewDecoder(rr.Result().Body).Decode(&resp); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
if resp.Error != nil {
|
||||
t.Fatalf("rpc error: %+v", resp.Error)
|
||||
}
|
||||
if resp.Result.Status.State != stateFailed {
|
||||
t.Fatalf("task state = %q, want failed", resp.Result.Status.State)
|
||||
}
|
||||
if len(resp.Result.Artifacts) != 1 || textOf(resp.Result.Artifacts[0].Parts) != "error: context canceled" {
|
||||
t.Fatalf("artifact = %+v, want context cancellation", resp.Result.Artifacts)
|
||||
}
|
||||
if len(resp.Result.History) != 2 || resp.Result.History[1].Role != "agent" || textOf(resp.Result.History[1].Parts) != "error: context canceled" {
|
||||
t.Fatalf("history = %+v, want failed agent reply recorded", resp.Result.History)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessageStream(t *testing.T) {
|
||||
ts, cleanup := newGatewayWithAgent(t)
|
||||
defer cleanup()
|
||||
|
||||
body := `{"jsonrpc":"2.0","id":1,"method":"message/stream","params":{"message":{"role":"user","parts":[{"kind":"text","text":"ping"}],"kind":"message"}}}`
|
||||
resp, err := http.Post(ts.URL+"/agents/echo", "application/json", bytes.NewBufferString(body))
|
||||
if err != nil {
|
||||
t.Fatalf("post: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if ct := resp.Header.Get("Content-Type"); !strings.HasPrefix(ct, "text/event-stream") {
|
||||
t.Fatalf("content-type = %q, want text/event-stream", ct)
|
||||
}
|
||||
|
||||
var line string
|
||||
if _, err := fmt.Fscan(resp.Body, &line); err != nil {
|
||||
t.Fatalf("read event prefix: %v", err)
|
||||
}
|
||||
if line != "data:" {
|
||||
t.Fatalf("event prefix = %q, want data:", line)
|
||||
}
|
||||
var out struct {
|
||||
Result Task `json:"result"`
|
||||
Error *rpcError `json:"error"`
|
||||
}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
|
||||
t.Fatalf("decode event: %v", err)
|
||||
}
|
||||
if out.Error != nil {
|
||||
t.Fatalf("rpc error: %+v", out.Error)
|
||||
}
|
||||
if out.Result.Status.State != stateCompleted || len(out.Result.Artifacts) != 1 || textOf(out.Result.Artifacts[0].Parts) != "pong" {
|
||||
t.Fatalf("streamed task = %+v", out.Result)
|
||||
}
|
||||
}
|
||||
|
||||
type sliceStream struct {
|
||||
chunks []string
|
||||
err error
|
||||
}
|
||||
|
||||
func (s *sliceStream) Recv() (*ai.Response, error) {
|
||||
if len(s.chunks) == 0 {
|
||||
if s.err != nil {
|
||||
err := s.err
|
||||
s.err = nil
|
||||
return nil, err
|
||||
}
|
||||
return nil, io.EOF
|
||||
}
|
||||
next := s.chunks[0]
|
||||
s.chunks = s.chunks[1:]
|
||||
return &ai.Response{Reply: next}, nil
|
||||
}
|
||||
|
||||
func (s *sliceStream) Close() error { return nil }
|
||||
|
||||
// streamEvent is one decoded SSE JSON-RPC event from a message/stream response.
|
||||
// A2A streams carry heterogeneous results (Task, status-update, artifact-update)
|
||||
// discriminated by `kind`, so we keep the raw result and decode on demand.
|
||||
type streamEvent struct {
|
||||
Result json.RawMessage `json:"result"`
|
||||
Error *rpcError `json:"error"`
|
||||
}
|
||||
|
||||
func (e streamEvent) kind() string {
|
||||
var k struct {
|
||||
Kind string `json:"kind"`
|
||||
}
|
||||
_ = json.Unmarshal(e.Result, &k)
|
||||
return k.Kind
|
||||
}
|
||||
|
||||
func (e streamEvent) task(t *testing.T) Task {
|
||||
t.Helper()
|
||||
var task Task
|
||||
if err := json.Unmarshal(e.Result, &task); err != nil {
|
||||
t.Fatalf("decode task event: %v", err)
|
||||
}
|
||||
return task
|
||||
}
|
||||
|
||||
func (e streamEvent) status(t *testing.T) TaskStatusUpdateEvent {
|
||||
t.Helper()
|
||||
var s TaskStatusUpdateEvent
|
||||
if err := json.Unmarshal(e.Result, &s); err != nil {
|
||||
t.Fatalf("decode status-update event: %v", err)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func (e streamEvent) artifactUpdate(t *testing.T) TaskArtifactUpdateEvent {
|
||||
t.Helper()
|
||||
var a TaskArtifactUpdateEvent
|
||||
if err := json.Unmarshal(e.Result, &a); err != nil {
|
||||
t.Fatalf("decode artifact-update event: %v", err)
|
||||
}
|
||||
return a
|
||||
}
|
||||
|
||||
// collectSSE parses the `data:`-framed JSON-RPC events from an SSE body.
|
||||
func collectSSE(t *testing.T, body string) []streamEvent {
|
||||
t.Helper()
|
||||
var events []streamEvent
|
||||
for _, line := range strings.Split(strings.TrimSpace(body), "\n") {
|
||||
line = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(line), "data:"))
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
var e streamEvent
|
||||
if err := json.Unmarshal([]byte(line), &e); err != nil {
|
||||
t.Fatalf("decode event %q: %v", line, err)
|
||||
}
|
||||
events = append(events, e)
|
||||
}
|
||||
return events
|
||||
}
|
||||
|
||||
func TestMessageStreamChunksStoreFinalTask(t *testing.T) {
|
||||
d := newDispatcher()
|
||||
body := `{"jsonrpc":"2.0","id":1,"method":"message/stream","params":{"message":{"role":"user","parts":[{"kind":"text","text":"ping"}],"kind":"message"}}}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString(body))
|
||||
rr := httptest.NewRecorder()
|
||||
|
||||
d.serveWithStream(rr, req, nil, func(ctx context.Context, text string) (ai.Stream, error) {
|
||||
if text != "ping" {
|
||||
t.Fatalf("stream text = %q, want ping", text)
|
||||
}
|
||||
return &sliceStream{chunks: []string{"po", "ng"}}, nil
|
||||
})
|
||||
|
||||
if ct := rr.Result().Header.Get("Content-Type"); !strings.HasPrefix(ct, "text/event-stream") {
|
||||
t.Fatalf("content-type = %q, want text/event-stream", ct)
|
||||
}
|
||||
events := collectSSE(t, rr.Body.String())
|
||||
// Opening Task snapshot + one append artifact-update per chunk + terminal
|
||||
// status-update.
|
||||
if len(events) != 4 {
|
||||
t.Fatalf("events = %d, want 4; body %s", len(events), rr.Body.String())
|
||||
}
|
||||
for i, e := range events {
|
||||
if e.Error != nil {
|
||||
t.Fatalf("event %d carried an error field: %+v", i, e.Error)
|
||||
}
|
||||
}
|
||||
if events[0].kind() != "task" {
|
||||
t.Fatalf("first event kind = %q, want task", events[0].kind())
|
||||
}
|
||||
opening := events[0].task(t)
|
||||
if opening.Status.State != stateWorking {
|
||||
t.Fatalf("opening task state = %q, want working", opening.Status.State)
|
||||
}
|
||||
taskID := opening.ID
|
||||
|
||||
// The middle events are append artifact-updates carrying the chunk deltas.
|
||||
var text strings.Builder
|
||||
for _, e := range events[1:3] {
|
||||
if e.kind() != "artifact-update" {
|
||||
t.Fatalf("event kind = %q, want artifact-update", e.kind())
|
||||
}
|
||||
au := e.artifactUpdate(t)
|
||||
if !au.Append {
|
||||
t.Fatalf("artifact-update should be append: %+v", au)
|
||||
}
|
||||
if au.TaskID != taskID {
|
||||
t.Fatalf("artifact-update taskId = %q, want %q", au.TaskID, taskID)
|
||||
}
|
||||
text.WriteString(textOf(au.Artifact.Parts))
|
||||
}
|
||||
if text.String() != "pong" {
|
||||
t.Fatalf("accumulated artifact text = %q, want pong", text.String())
|
||||
}
|
||||
|
||||
// The stream closes with a terminal status-update (final:true).
|
||||
last := events[len(events)-1]
|
||||
if last.kind() != "status-update" {
|
||||
t.Fatalf("last event kind = %q, want status-update", last.kind())
|
||||
}
|
||||
su := last.status(t)
|
||||
if !su.Final || su.Status.State != stateCompleted {
|
||||
t.Fatalf("terminal event = %+v, want final completed", su)
|
||||
}
|
||||
if su.TaskID != taskID {
|
||||
t.Fatalf("terminal taskId = %q, want %q", su.TaskID, taskID)
|
||||
}
|
||||
|
||||
got := rpcTaskFromDispatcher(t, d, taskID)
|
||||
if got.ID != taskID || got.Status.State != stateCompleted || textOf(got.Artifacts[0].Parts) != "pong" {
|
||||
t.Fatalf("stored task = %+v, want final completed pong", got)
|
||||
}
|
||||
}
|
||||
|
||||
type contextStream struct {
|
||||
ctx context.Context
|
||||
closed chan struct{}
|
||||
}
|
||||
|
||||
func (s *contextStream) Recv() (*ai.Response, error) {
|
||||
<-s.ctx.Done()
|
||||
return nil, s.ctx.Err()
|
||||
}
|
||||
|
||||
func (s *contextStream) Close() error {
|
||||
close(s.closed)
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestMessageStreamChunksPropagatesCancellationAndClosesStream(t *testing.T) {
|
||||
d := newDispatcher()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
closed := make(chan struct{})
|
||||
body := `{"jsonrpc":"2.0","id":1,"method":"message/stream","params":{"message":{"role":"user","parts":[{"kind":"text","text":"ping"}],"kind":"message"}}}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString(body)).WithContext(ctx)
|
||||
rr := httptest.NewRecorder()
|
||||
cancel()
|
||||
|
||||
d.serveWithStream(rr, req, nil, func(ctx context.Context, text string) (ai.Stream, error) {
|
||||
if text != "ping" {
|
||||
t.Fatalf("stream text = %q, want ping", text)
|
||||
}
|
||||
return &contextStream{ctx: ctx, closed: closed}, nil
|
||||
})
|
||||
|
||||
select {
|
||||
case <-closed:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("stream was not closed")
|
||||
}
|
||||
|
||||
events := collectSSE(t, rr.Body.String())
|
||||
// Opening Task snapshot, then a terminal failed status-update.
|
||||
if len(events) != 2 {
|
||||
t.Fatalf("events = %d, want 2; body %s", len(events), rr.Body.String())
|
||||
}
|
||||
// A streaming failure must be a failed status-update, never `result` and
|
||||
// `error` set together in one response.
|
||||
for i, e := range events {
|
||||
if e.Error != nil {
|
||||
t.Fatalf("event %d carried an error field (result+error not allowed): %+v", i, e.Error)
|
||||
}
|
||||
}
|
||||
if events[0].kind() != "task" || events[0].task(t).Status.State != stateWorking {
|
||||
t.Fatalf("first event = %s, want working task", string(events[0].Result))
|
||||
}
|
||||
last := events[1]
|
||||
if last.kind() != "status-update" {
|
||||
t.Fatalf("last event kind = %q, want status-update", last.kind())
|
||||
}
|
||||
su := last.status(t)
|
||||
if !su.Final || su.Status.State != stateFailed {
|
||||
t.Fatalf("terminal event = %+v, want final failed", su)
|
||||
}
|
||||
|
||||
got := rpcTaskFromDispatcher(t, d, su.TaskID)
|
||||
if got.Status.State != stateFailed || textOf(got.Artifacts[0].Parts) != "error: context canceled" {
|
||||
t.Fatalf("stored task = %+v, want failed cancellation", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessageStreamChunksFallsBackWhenUnsupported(t *testing.T) {
|
||||
d := newDispatcher()
|
||||
body := `{"jsonrpc":"2.0","id":1,"method":"message/stream","params":{"message":{"role":"user","parts":[{"kind":"text","text":"ping"}],"kind":"message"}}}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString(body))
|
||||
rr := httptest.NewRecorder()
|
||||
var streamed bool
|
||||
var fallbackText string
|
||||
|
||||
d.serveWithStream(rr, req, func(ctx context.Context, text string) (string, error) {
|
||||
fallbackText = text
|
||||
return "pong", nil
|
||||
}, func(ctx context.Context, text string) (ai.Stream, error) {
|
||||
streamed = true
|
||||
return nil, fmt.Errorf("%w: test provider", ai.ErrStreamingUnsupported)
|
||||
})
|
||||
|
||||
if !streamed {
|
||||
t.Fatal("stream invoke was not attempted")
|
||||
}
|
||||
if fallbackText != "ping" {
|
||||
t.Fatalf("fallback text = %q, want ping", fallbackText)
|
||||
}
|
||||
if ct := rr.Result().Header.Get("Content-Type"); !strings.HasPrefix(ct, "text/event-stream") {
|
||||
t.Fatalf("content-type = %q, want text/event-stream", ct)
|
||||
}
|
||||
events := collectSSE(t, rr.Body.String())
|
||||
// The non-streaming fallback emits a completed Task snapshot then a terminal
|
||||
// status-update.
|
||||
if len(events) != 2 {
|
||||
t.Fatalf("events = %d, want 2; body %s", len(events), rr.Body.String())
|
||||
}
|
||||
for i, e := range events {
|
||||
if e.Error != nil {
|
||||
t.Fatalf("fallback event %d error: %+v", i, e.Error)
|
||||
}
|
||||
}
|
||||
task := events[0].task(t)
|
||||
if task.Status.State != stateCompleted || textOf(task.Artifacts[0].Parts) != "pong" {
|
||||
t.Fatalf("fallback task = %+v, want completed pong", task)
|
||||
}
|
||||
su := events[1].status(t)
|
||||
if !su.Final || su.Status.State != stateCompleted {
|
||||
t.Fatalf("terminal event = %+v, want final completed", su)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessageStreamFallbackDoesNotCompleteWithEmptyText(t *testing.T) {
|
||||
d := newDispatcher()
|
||||
body := `{"jsonrpc":"2.0","id":1,"method":"message/stream","params":{"message":{"role":"user","parts":[{"kind":"text","text":"ping"}],"kind":"message"}}}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString(body))
|
||||
rr := httptest.NewRecorder()
|
||||
|
||||
d.serveWithStream(rr, req, func(context.Context, string) (string, error) {
|
||||
return "", nil
|
||||
}, func(context.Context, string) (ai.Stream, error) {
|
||||
return nil, fmt.Errorf("%w: test provider", ai.ErrStreamingUnsupported)
|
||||
})
|
||||
|
||||
events := collectSSE(t, rr.Body.String())
|
||||
var task Task
|
||||
var foundTask bool
|
||||
for _, e := range events {
|
||||
if e.Error != nil {
|
||||
t.Fatalf("fallback event error: %+v", e.Error)
|
||||
}
|
||||
if e.kind() == "task" {
|
||||
task = e.task(t)
|
||||
foundTask = true
|
||||
}
|
||||
}
|
||||
if !foundTask {
|
||||
t.Fatalf("no task event in stream; body %s", rr.Body.String())
|
||||
}
|
||||
if task.Status.State != stateFailed {
|
||||
t.Fatalf("fallback state = %q, want failed", task.Status.State)
|
||||
}
|
||||
if got := textOf(task.Artifacts[0].Parts); got == "" {
|
||||
t.Fatalf("fallback artifact text is empty: %+v", task.Artifacts)
|
||||
}
|
||||
if got := textOf(task.History[len(task.History)-1].Parts); got == "" {
|
||||
t.Fatalf("fallback history text is empty: %+v", task.History)
|
||||
}
|
||||
// The stream still ends with a terminal marker.
|
||||
last := events[len(events)-1]
|
||||
if last.kind() != "status-update" || !last.status(t).Final {
|
||||
t.Fatalf("stream must end with a final status-update; got %s", string(last.Result))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeAgentChatReplyFallsBackToProviderTextFields(t *testing.T) {
|
||||
for name, body := range map[string]string{
|
||||
"answer": `{"answer":"answer text"}`,
|
||||
"content": `{"content":"content text"}`,
|
||||
"text": `{"text":"text field"}`,
|
||||
"message_content": `{"message":{"content":"message content"}}`,
|
||||
"message_text": `{"message":{"text":"message text"}}`,
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
got, err := decodeAgentChatReply([]byte(body))
|
||||
if err != nil {
|
||||
t.Fatalf("decodeAgentChatReply error: %v", err)
|
||||
}
|
||||
if strings.TrimSpace(got) == "" {
|
||||
t.Fatalf("decodeAgentChatReply(%s) returned empty text", body)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTasksResubscribeStreamsCurrentAndSubsequentEvents(t *testing.T) {
|
||||
d := newDispatcher()
|
||||
initial := &Task{ID: "task-1", ContextID: "ctx-1", Kind: "task", Status: TaskStatus{State: stateWorking, Timestamp: time.Now().UTC().Format(time.RFC3339)}}
|
||||
d.store(initial)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
req := httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString(`{"jsonrpc":"2.0","id":1,"method":"tasks/resubscribe","params":{"id":"task-1"}}`)).WithContext(ctx)
|
||||
rw := newFlushRecorder()
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
d.serve(rw, req, nil)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
first := rw.next(t)
|
||||
if first.Result.ID != initial.ID || first.Result.Status.State != stateWorking {
|
||||
t.Fatalf("first resubscribe event = %+v, want current working task", first.Result)
|
||||
}
|
||||
|
||||
final := &Task{ID: "task-1", ContextID: "ctx-1", Kind: "task", Status: TaskStatus{State: stateCompleted, Timestamp: time.Now().UTC().Format(time.RFC3339)}, Artifacts: []Artifact{textArtifact("done")}}
|
||||
d.store(final)
|
||||
second := rw.next(t)
|
||||
if second.Result.ID != final.ID || second.Result.Status.State != stateCompleted || textOf(second.Result.Artifacts[0].Parts) != "done" {
|
||||
t.Fatalf("second resubscribe event = %+v, want completed update", second.Result)
|
||||
}
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("resubscribe did not return after terminal update")
|
||||
}
|
||||
}
|
||||
|
||||
func TestInputRequiredErrorCreatesContinuableTask(t *testing.T) {
|
||||
d := newDispatcher()
|
||||
first := rpcTaskFromBody(t, d, `{
|
||||
"jsonrpc":"2.0","id":1,"method":"message/send",
|
||||
"params":{"message":{"role":"user","kind":"message","messageId":"m1",
|
||||
"parts":[{"kind":"text","text":"start approval"}]}}}`, func(_ context.Context, text string) (string, error) {
|
||||
return "", errors.New("agent run run-1 paused for approval: waiting for operator")
|
||||
})
|
||||
if first.Status.State != stateInputRequired {
|
||||
t.Fatalf("state = %q, want input-required", first.Status.State)
|
||||
}
|
||||
if textOf(first.Artifacts[0].Parts) != "agent run run-1 paused for approval: waiting for operator" {
|
||||
t.Fatalf("artifact = %+v, want handoff message", first.Artifacts)
|
||||
}
|
||||
|
||||
body := fmt.Sprintf(`{
|
||||
"jsonrpc":"2.0","id":2,"method":"message/send",
|
||||
"params":{"message":{"role":"user","kind":"message","messageId":"m2","taskId":"%s","contextId":"%s",
|
||||
"parts":[{"kind":"text","text":"approved"}]}}}`, first.ID, first.ContextID)
|
||||
continued := rpcTaskFromBody(t, d, body, func(_ context.Context, text string) (string, error) {
|
||||
return "continued after " + text, nil
|
||||
})
|
||||
if continued.ID != first.ID || continued.ContextID != first.ContextID {
|
||||
t.Fatalf("continued identity = %s/%s, want %s/%s", continued.ID, continued.ContextID, first.ID, first.ContextID)
|
||||
}
|
||||
if continued.Status.State != stateCompleted || len(continued.History) != 4 {
|
||||
t.Fatalf("continued task = %+v, want completed task with prior input-required history", continued)
|
||||
}
|
||||
if textOf(continued.History[1].Parts) != "agent run run-1 paused for approval: waiting for operator" || textOf(continued.History[3].Parts) != "continued after approved" {
|
||||
t.Fatalf("continued history = %+v", continued.History)
|
||||
}
|
||||
}
|
||||
|
||||
type flushRecorder struct {
|
||||
*httptest.ResponseRecorder
|
||||
ch chan string
|
||||
}
|
||||
|
||||
func newFlushRecorder() *flushRecorder {
|
||||
return &flushRecorder{ResponseRecorder: httptest.NewRecorder(), ch: make(chan string, 16)}
|
||||
}
|
||||
|
||||
func (r *flushRecorder) Flush() {
|
||||
body := r.Body.String()
|
||||
r.Body.Reset()
|
||||
for _, line := range strings.Split(strings.TrimSpace(body), "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line != "" {
|
||||
r.ch <- line
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *flushRecorder) next(t *testing.T) struct {
|
||||
Result Task `json:"result"`
|
||||
Error *rpcError `json:"error"`
|
||||
} {
|
||||
t.Helper()
|
||||
select {
|
||||
case line := <-r.ch:
|
||||
line = strings.TrimPrefix(line, "data: ")
|
||||
var event struct {
|
||||
Result Task `json:"result"`
|
||||
Error *rpcError `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(line), &event); err != nil {
|
||||
t.Fatalf("decode event %q: %v", line, err)
|
||||
}
|
||||
if event.Error != nil {
|
||||
t.Fatalf("event error: %+v", event.Error)
|
||||
}
|
||||
return event
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("timed out waiting for SSE event")
|
||||
}
|
||||
return struct {
|
||||
Result Task `json:"result"`
|
||||
Error *rpcError `json:"error"`
|
||||
}{}
|
||||
}
|
||||
|
||||
func rpcTaskFromDispatcher(t *testing.T, d *dispatcher, id string) Task {
|
||||
t.Helper()
|
||||
body := fmt.Sprintf(`{"jsonrpc":"2.0","id":2,"method":"tasks/get","params":{"id":"%s"}}`, id)
|
||||
req := httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString(body))
|
||||
rr := httptest.NewRecorder()
|
||||
d.serve(rr, req, nil)
|
||||
var resp struct {
|
||||
Result Task `json:"result"`
|
||||
Error *rpcError `json:"error"`
|
||||
}
|
||||
if err := json.NewDecoder(rr.Result().Body).Decode(&resp); err != nil {
|
||||
t.Fatalf("decode tasks/get: %v", err)
|
||||
}
|
||||
if resp.Error != nil {
|
||||
t.Fatalf("tasks/get error: %+v", resp.Error)
|
||||
}
|
||||
return resp.Result
|
||||
}
|
||||
|
||||
func rpcTaskFromBody(t *testing.T, d *dispatcher, body string, invoke Invoke) Task {
|
||||
t.Helper()
|
||||
var resp struct {
|
||||
Result Task `json:"result"`
|
||||
Error *rpcError `json:"error"`
|
||||
}
|
||||
rpcDispatcher(t, d, body, invoke, &resp)
|
||||
if resp.Error != nil {
|
||||
t.Fatalf("rpc error: %+v", resp.Error)
|
||||
}
|
||||
return resp.Result
|
||||
}
|
||||
|
||||
func rpcDispatcher(t *testing.T, d *dispatcher, body string, invoke Invoke, v any) {
|
||||
t.Helper()
|
||||
req := httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString(body))
|
||||
rr := httptest.NewRecorder()
|
||||
d.serve(rr, req, invoke)
|
||||
if err := json.NewDecoder(rr.Result().Body).Decode(v); err != nil {
|
||||
t.Fatalf("decode dispatcher response: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnknownMethod(t *testing.T) {
|
||||
ts, cleanup := newGatewayWithAgent(t)
|
||||
defer cleanup()
|
||||
|
||||
var resp struct {
|
||||
Error *rpcError `json:"error"`
|
||||
}
|
||||
rpc(t, ts.URL+"/agents/echo", `{"jsonrpc":"2.0","id":1,"method":"unknown","params":{}}`, &resp)
|
||||
if resp.Error == nil || resp.Error.Code != errMethodNotFound {
|
||||
t.Errorf("expected method-not-found, got %+v", resp.Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListAgents(t *testing.T) {
|
||||
ts, cleanup := newGatewayWithAgent(t)
|
||||
defer cleanup()
|
||||
|
||||
resp, err := http.Get(ts.URL + "/agents")
|
||||
if err != nil {
|
||||
t.Fatalf("list: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
var out struct {
|
||||
Agents []AgentCard `json:"agents"`
|
||||
}
|
||||
json.NewDecoder(resp.Body).Decode(&out)
|
||||
if len(out.Agents) != 1 || out.Agents[0].Name != "echo" {
|
||||
t.Errorf("agents list = %+v", out.Agents)
|
||||
}
|
||||
}
|
||||
|
||||
func rpc(t *testing.T, url, body string, v any) {
|
||||
t.Helper()
|
||||
resp, err := http.Post(url, "application/json", bytes.NewBufferString(body))
|
||||
if err != nil {
|
||||
t.Fatalf("post: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if err := json.NewDecoder(resp.Body).Decode(v); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func rpcTask(t *testing.T, url, body string) Task {
|
||||
t.Helper()
|
||||
var resp struct {
|
||||
Result Task `json:"result"`
|
||||
Error *rpcError `json:"error"`
|
||||
}
|
||||
rpc(t, url, body, &resp)
|
||||
if resp.Error != nil {
|
||||
t.Fatalf("rpc error: %+v", resp.Error)
|
||||
}
|
||||
return resp.Result
|
||||
}
|
||||
|
||||
func skillIDs(skills []Skill) []string {
|
||||
ids := make([]string, 0, len(skills))
|
||||
for _, skill := range skills {
|
||||
ids = append(ids, skill.ID)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
@@ -0,0 +1,150 @@
|
||||
package a2a
|
||||
|
||||
import (
|
||||
"crypto/ed25519"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// AP2MandateKind identifies the AP2 mandate stage represented by a credential.
|
||||
type AP2MandateKind string
|
||||
|
||||
const (
|
||||
AP2CheckoutMandate AP2MandateKind = "checkout"
|
||||
AP2PaymentMandate AP2MandateKind = "payment"
|
||||
)
|
||||
|
||||
// AP2RailRef names the settlement rail authorized by an AP2 payment mandate.
|
||||
type AP2RailRef struct {
|
||||
Type string `json:"type"`
|
||||
Reference string `json:"reference"`
|
||||
}
|
||||
|
||||
// AP2Mandate is a small verifiable AP2 credential. It is deliberately rail-neutral:
|
||||
// x402 is represented as one possible rail reference under a payment mandate.
|
||||
type AP2Mandate struct {
|
||||
ID string `json:"id"`
|
||||
Kind AP2MandateKind `json:"kind"`
|
||||
Subject string `json:"subject,omitempty"`
|
||||
Merchant string `json:"merchant,omitempty"`
|
||||
Amount string `json:"amount,omitempty"`
|
||||
Currency string `json:"currency,omitempty"`
|
||||
Description string `json:"description,omitempty"`
|
||||
TaskID string `json:"taskId,omitempty"`
|
||||
ContextID string `json:"contextId,omitempty"`
|
||||
Rail *AP2RailRef `json:"rail,omitempty"`
|
||||
IssuedAt time.Time `json:"issuedAt"`
|
||||
}
|
||||
|
||||
// AP2SignedMandate is an AP2 mandate plus an Ed25519 signature over its canonical JSON.
|
||||
type AP2SignedMandate struct {
|
||||
Mandate AP2Mandate `json:"mandate"`
|
||||
KeyID string `json:"keyId,omitempty"`
|
||||
Signature string `json:"signature"`
|
||||
}
|
||||
|
||||
// AP2Verification records mandate verification on an A2A task without mixing in
|
||||
// payment-settlement state.
|
||||
type AP2Verification struct {
|
||||
MandateID string `json:"mandateId"`
|
||||
Kind string `json:"kind"`
|
||||
Verified bool `json:"verified"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// NewAP2Keypair returns an Ed25519 keypair suitable for tests or local demos.
|
||||
func NewAP2Keypair() (ed25519.PublicKey, ed25519.PrivateKey, error) {
|
||||
return ed25519.GenerateKey(rand.Reader)
|
||||
}
|
||||
|
||||
// SignAP2Mandate signs a mandate as a verifiable AP2 credential.
|
||||
func SignAP2Mandate(m AP2Mandate, keyID string, private ed25519.PrivateKey) (AP2SignedMandate, error) {
|
||||
if m.ID == "" {
|
||||
return AP2SignedMandate{}, errors.New("ap2: mandate id is required")
|
||||
}
|
||||
if m.Kind == "" {
|
||||
return AP2SignedMandate{}, errors.New("ap2: mandate kind is required")
|
||||
}
|
||||
if m.IssuedAt.IsZero() {
|
||||
m.IssuedAt = time.Now().UTC()
|
||||
}
|
||||
payload, err := ap2Payload(m)
|
||||
if err != nil {
|
||||
return AP2SignedMandate{}, err
|
||||
}
|
||||
return AP2SignedMandate{Mandate: m, KeyID: keyID, Signature: base64.RawURLEncoding.EncodeToString(ed25519.Sign(private, payload))}, nil
|
||||
}
|
||||
|
||||
// VerifyAP2Mandate verifies a signed mandate credential.
|
||||
func VerifyAP2Mandate(s AP2SignedMandate, public ed25519.PublicKey) error {
|
||||
sig, err := base64.RawURLEncoding.DecodeString(s.Signature)
|
||||
if err != nil {
|
||||
return fmt.Errorf("ap2: invalid signature encoding: %w", err)
|
||||
}
|
||||
payload, err := ap2Payload(s.Mandate)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !ed25519.Verify(public, payload, sig) {
|
||||
return errors.New("ap2: mandate signature verification failed")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// AP2BindMandateToMessage returns a copy of m bound to the A2A message's task/context.
|
||||
func AP2BindMandateToMessage(m AP2Mandate, msg Message) AP2Mandate {
|
||||
m.TaskID = msg.TaskID
|
||||
m.ContextID = msg.ContextID
|
||||
return m
|
||||
}
|
||||
|
||||
// AP2AttachMandate returns a copy of msg carrying the signed AP2 mandate.
|
||||
func AP2AttachMandate(msg Message, mandate AP2SignedMandate) Message {
|
||||
msg.AP2Mandates = append(append([]AP2SignedMandate{}, msg.AP2Mandates...), mandate)
|
||||
return msg
|
||||
}
|
||||
|
||||
// VerifyAP2ForTask verifies signature, task binding, and optional settlement rail reference.
|
||||
func VerifyAP2ForTask(s AP2SignedMandate, public ed25519.PublicKey, task Task, rail *AP2RailRef) AP2Verification {
|
||||
out := AP2Verification{MandateID: s.Mandate.ID, Kind: string(s.Mandate.Kind), Verified: true}
|
||||
if err := VerifyAP2Mandate(s, public); err != nil {
|
||||
out.Verified = false
|
||||
out.Error = err.Error()
|
||||
return out
|
||||
}
|
||||
if s.Mandate.TaskID != "" && s.Mandate.TaskID != task.ID {
|
||||
out.Verified = false
|
||||
out.Error = "ap2: mandate task binding mismatch"
|
||||
return out
|
||||
}
|
||||
if s.Mandate.ContextID != "" && s.Mandate.ContextID != task.ContextID {
|
||||
out.Verified = false
|
||||
out.Error = "ap2: mandate context binding mismatch"
|
||||
return out
|
||||
}
|
||||
if s.Mandate.Kind == AP2PaymentMandate && rail != nil {
|
||||
if s.Mandate.Rail == nil || *s.Mandate.Rail != *rail {
|
||||
out.Verified = false
|
||||
out.Error = "ap2: settlement rail reference mismatch"
|
||||
return out
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// X402AP2Rail builds the x402 settlement rail reference carried under a payment mandate.
|
||||
func X402AP2Rail(reference string) AP2RailRef { return AP2RailRef{Type: "x402", Reference: reference} }
|
||||
|
||||
func ap2Payload(m AP2Mandate) ([]byte, error) {
|
||||
b, err := json.Marshal(m)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("ap2: marshal mandate: %w", err)
|
||||
}
|
||||
sum := sha256.Sum256(b)
|
||||
return sum[:], nil
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package a2a
|
||||
|
||||
import (
|
||||
"crypto/ed25519"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func testAP2Key(t *testing.T) (ed25519.PublicKey, ed25519.PrivateKey) {
|
||||
t.Helper()
|
||||
pub, priv, err := NewAP2Keypair()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return pub, priv
|
||||
}
|
||||
|
||||
func TestAP2CheckoutMandateSignAttachAndVerify(t *testing.T) {
|
||||
pub, priv := testAP2Key(t)
|
||||
msg := Message{Role: "user", Kind: "message", TaskID: "task-1", ContextID: "ctx-1", Parts: []Part{{Kind: "text", Text: "buy"}}}
|
||||
mandate := AP2BindMandateToMessage(AP2Mandate{ID: "checkout-1", Kind: AP2CheckoutMandate, Subject: "alice", Merchant: "store", Amount: "10.00", Currency: "USD", Description: "demo", IssuedAt: time.Unix(1, 0).UTC()}, msg)
|
||||
signed, err := SignAP2Mandate(mandate, "test-key", priv)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
msg = AP2AttachMandate(msg, signed)
|
||||
task := taskFromReplyWithIDs(msg, "ok", stateCompleted, msg.TaskID, msg.ContextID)
|
||||
|
||||
if len(task.AP2Mandates) != 1 {
|
||||
t.Fatalf("expected mandate carried on task, got %d", len(task.AP2Mandates))
|
||||
}
|
||||
got := VerifyAP2ForTask(task.AP2Mandates[0], pub, *task, nil)
|
||||
if !got.Verified || got.Error != "" {
|
||||
t.Fatalf("expected verified mandate, got %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAP2PaymentMandateX402RailReference(t *testing.T) {
|
||||
pub, priv := testAP2Key(t)
|
||||
rail := X402AP2Rail("payreq_123")
|
||||
task := Task{ID: "task-2", ContextID: "ctx-2"}
|
||||
signed, err := SignAP2Mandate(AP2Mandate{ID: "payment-1", Kind: AP2PaymentMandate, TaskID: task.ID, ContextID: task.ContextID, Rail: &rail, IssuedAt: time.Unix(1, 0).UTC()}, "test-key", priv)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := VerifyAP2ForTask(signed, pub, task, &rail)
|
||||
if !got.Verified {
|
||||
t.Fatalf("expected x402 rail to verify, got %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAP2TamperCasesFailDistinctly(t *testing.T) {
|
||||
pub, priv := testAP2Key(t)
|
||||
rail := X402AP2Rail("payreq_123")
|
||||
task := Task{ID: "task-3", ContextID: "ctx-3"}
|
||||
signed, err := SignAP2Mandate(AP2Mandate{ID: "payment-2", Kind: AP2PaymentMandate, TaskID: task.ID, ContextID: task.ContextID, Rail: &rail, IssuedAt: time.Unix(1, 0).UTC()}, "test-key", priv)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
tampered := signed
|
||||
tampered.Mandate.Amount = "999.00"
|
||||
if got := VerifyAP2ForTask(tampered, pub, task, &rail); got.Verified || !strings.Contains(got.Error, "signature") {
|
||||
t.Fatalf("expected signature failure, got %+v", got)
|
||||
}
|
||||
|
||||
wrongTask := task
|
||||
wrongTask.ID = "other-task"
|
||||
if got := VerifyAP2ForTask(signed, pub, wrongTask, &rail); got.Verified || !strings.Contains(got.Error, "task binding") {
|
||||
t.Fatalf("expected task binding failure, got %+v", got)
|
||||
}
|
||||
|
||||
otherRail := X402AP2Rail("payreq_other")
|
||||
if got := VerifyAP2ForTask(signed, pub, task, &otherRail); got.Verified || !strings.Contains(got.Error, "rail reference") {
|
||||
t.Fatalf("expected rail reference failure, got %+v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,282 @@
|
||||
package a2a
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// Client calls a remote agent that speaks the A2A protocol. It is the
|
||||
// outbound counterpart to the gateway: where the gateway exposes a Go
|
||||
// Micro agent over A2A, the Client lets a Go Micro agent or flow call an
|
||||
// agent on any framework, by URL.
|
||||
type Client struct {
|
||||
url string
|
||||
http *http.Client
|
||||
}
|
||||
|
||||
// NewClient returns a Client for the agent at url (its JSON-RPC endpoint,
|
||||
// i.e. the `url` field of the agent's card).
|
||||
func NewClient(url string) *Client {
|
||||
return &Client{url: url, http: &http.Client{Timeout: 60 * time.Second}}
|
||||
}
|
||||
|
||||
// WithHTTPClient sets the underlying HTTP client (for timeouts, auth
|
||||
// transports, etc.).
|
||||
func (c *Client) WithHTTPClient(h *http.Client) *Client {
|
||||
if h != nil {
|
||||
c.http = h
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
// Card fetches the remote agent's Agent Card.
|
||||
func (c *Client) Card(ctx context.Context) (*AgentCard, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.url+"/.well-known/agent.json", nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resp, err := c.http.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("agent card: status %d", resp.StatusCode)
|
||||
}
|
||||
var card AgentCard
|
||||
if err := json.NewDecoder(resp.Body).Decode(&card); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &card, nil
|
||||
}
|
||||
|
||||
// Send sends a text message to the remote agent and returns its reply.
|
||||
// If the agent returns a task that isn't yet terminal, Send polls
|
||||
// tasks/get until it completes or ctx is done.
|
||||
func (c *Client) Send(ctx context.Context, text string) (string, error) {
|
||||
task, err := c.SendMessage(ctx, Message{
|
||||
Role: "user",
|
||||
Kind: "message",
|
||||
MessageID: uuid.New().String(),
|
||||
Parts: []Part{{Kind: "text", Text: text}},
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if task.Status.State != stateCompleted {
|
||||
return "", fmt.Errorf("remote task %s ended in state %q", task.ID, task.Status.State)
|
||||
}
|
||||
return artifactsText(task.Artifacts), nil
|
||||
}
|
||||
|
||||
// SendMessage sends an A2A message and returns the resulting terminal task.
|
||||
// To continue a multi-turn task, pass a Message with TaskID and ContextID set
|
||||
// to a prior task's id and context id.
|
||||
func (c *Client) SendMessage(ctx context.Context, message Message) (*Task, error) {
|
||||
if message.MessageID == "" {
|
||||
message.MessageID = uuid.New().String()
|
||||
}
|
||||
if message.Kind == "" {
|
||||
message.Kind = "message"
|
||||
}
|
||||
if message.Role == "" {
|
||||
message.Role = "user"
|
||||
}
|
||||
res, err := c.call(ctx, "message/send", sendParams{Message: message})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// The result is a Message or a Task; the "kind" field disambiguates.
|
||||
var probe struct {
|
||||
Kind string `json:"kind"`
|
||||
}
|
||||
_ = json.Unmarshal(res, &probe)
|
||||
|
||||
if probe.Kind == "message" {
|
||||
var m Message
|
||||
if err := json.Unmarshal(res, &m); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &Task{
|
||||
ID: m.TaskID,
|
||||
ContextID: m.ContextID,
|
||||
Kind: "task",
|
||||
Status: TaskStatus{State: stateCompleted, Timestamp: time.Now().UTC().Format(time.RFC3339)},
|
||||
Artifacts: []Artifact{textArtifact(textOf(m.Parts))},
|
||||
History: []Message{m},
|
||||
AP2Mandates: append([]AP2SignedMandate{}, m.AP2Mandates...),
|
||||
}, nil
|
||||
}
|
||||
|
||||
var task Task
|
||||
if err := json.Unmarshal(res, &task); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for !terminal(task.Status.State) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case <-time.After(300 * time.Millisecond):
|
||||
}
|
||||
got, err := c.call(ctx, "tasks/get", getParams{ID: task.ID})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := json.Unmarshal(got, &task); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return &task, nil
|
||||
}
|
||||
|
||||
// Resubscribe reconnects to a retained or active task stream and returns task
|
||||
// snapshots as the remote agent emits updates. The returned channel is closed
|
||||
// when the task reaches a terminal state or ctx is canceled.
|
||||
func (c *Client) Resubscribe(ctx context.Context, taskID string) (<-chan Task, <-chan error) {
|
||||
tasks := make(chan Task, 8)
|
||||
errs := make(chan error, 1)
|
||||
go func() {
|
||||
defer close(tasks)
|
||||
defer close(errs)
|
||||
body, _ := json.Marshal(map[string]any{
|
||||
"jsonrpc": "2.0",
|
||||
"id": uuid.New().String(),
|
||||
"method": "tasks/resubscribe",
|
||||
"params": getParams{ID: taskID},
|
||||
})
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
errs <- err
|
||||
return
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "text/event-stream")
|
||||
resp, err := c.http.Do(req)
|
||||
if err != nil {
|
||||
errs <- err
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
errs <- fmt.Errorf("tasks/resubscribe: status %d", resp.StatusCode)
|
||||
return
|
||||
}
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if !strings.HasPrefix(line, "data:") {
|
||||
continue
|
||||
}
|
||||
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
if payload == "" {
|
||||
continue
|
||||
}
|
||||
var out struct {
|
||||
Result Task `json:"result"`
|
||||
Error *rpcError `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(payload), &out); err != nil {
|
||||
errs <- err
|
||||
return
|
||||
}
|
||||
if out.Error != nil {
|
||||
errs <- fmt.Errorf("a2a tasks/resubscribe: %s (%d)", out.Error.Message, out.Error.Code)
|
||||
return
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
errs <- ctx.Err()
|
||||
return
|
||||
case tasks <- out.Result:
|
||||
}
|
||||
if terminal(out.Result.Status.State) {
|
||||
return
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
errs <- err
|
||||
}
|
||||
}()
|
||||
return tasks, errs
|
||||
}
|
||||
|
||||
// SetPushNotificationConfig asks the remote agent to POST updates for taskID to cfg.URL.
|
||||
func (c *Client) SetPushNotificationConfig(ctx context.Context, taskID string, cfg PushNotificationConfig) error {
|
||||
_, err := c.call(ctx, "tasks/pushNotificationConfig/set", pushConfigParams{
|
||||
ID: taskID,
|
||||
PushNotificationConfig: cfg,
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
// PushNotificationConfig returns the remote push notification config for taskID.
|
||||
func (c *Client) PushNotificationConfig(ctx context.Context, taskID string) (PushNotificationConfig, error) {
|
||||
res, err := c.call(ctx, "tasks/pushNotificationConfig/get", getParams{ID: taskID})
|
||||
if err != nil {
|
||||
return PushNotificationConfig{}, err
|
||||
}
|
||||
var out struct {
|
||||
PushNotificationConfig PushNotificationConfig `json:"pushNotificationConfig"`
|
||||
}
|
||||
if err := json.Unmarshal(res, &out); err != nil {
|
||||
return PushNotificationConfig{}, err
|
||||
}
|
||||
return out.PushNotificationConfig, nil
|
||||
}
|
||||
|
||||
// call performs one JSON-RPC request and returns the raw result.
|
||||
func (c *Client) call(ctx context.Context, method string, params any) (json.RawMessage, error) {
|
||||
body, _ := json.Marshal(map[string]any{
|
||||
"jsonrpc": "2.0",
|
||||
"id": uuid.New().String(),
|
||||
"method": method,
|
||||
"params": params,
|
||||
})
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := c.http.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
var out struct {
|
||||
Result json.RawMessage `json:"result"`
|
||||
Error *rpcError `json:"error"`
|
||||
}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if out.Error != nil {
|
||||
return nil, fmt.Errorf("a2a %s: %s (%d)", method, out.Error.Message, out.Error.Code)
|
||||
}
|
||||
return out.Result, nil
|
||||
}
|
||||
|
||||
func terminal(state string) bool {
|
||||
switch state {
|
||||
case "completed", "failed", "canceled", "rejected", "input-required":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func artifactsText(arts []Artifact) string {
|
||||
var parts []Part
|
||||
for _, a := range arts {
|
||||
parts = append(parts, a.Parts...)
|
||||
}
|
||||
return textOf(parts)
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
package a2a
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// An agent that embeds NewAgentHandler is directly A2A-queryable — no
|
||||
// gateway, no registry — and the client talks to it the same way.
|
||||
func TestEmbeddedAgentHandler(t *testing.T) {
|
||||
card := Card("solo", "http://localhost:4000", "", []string{"task"})
|
||||
h := NewAgentHandler(card, func(_ context.Context, text string) (string, error) {
|
||||
return "echo:" + text, nil
|
||||
})
|
||||
ts := httptest.NewServer(h)
|
||||
defer ts.Close()
|
||||
|
||||
cl := NewClient(ts.URL)
|
||||
got, err := cl.Card(context.Background())
|
||||
if err != nil || got.Name != "solo" {
|
||||
t.Fatalf("Card() = %+v, err %v", got, err)
|
||||
}
|
||||
reply, err := cl.Send(context.Background(), "hi")
|
||||
if err != nil || reply != "echo:hi" {
|
||||
t.Fatalf("Send = %q, err %v", reply, err)
|
||||
}
|
||||
}
|
||||
|
||||
// The client fetches a card and sends a message to an agent served by the
|
||||
// gateway — A2A end to end, both directions, in one process.
|
||||
func TestClientSendAndCard(t *testing.T) {
|
||||
ts, cleanup := newGatewayWithAgent(t)
|
||||
defer cleanup()
|
||||
|
||||
cl := NewClient(ts.URL + "/agents/echo")
|
||||
|
||||
card, err := cl.Card(context.Background())
|
||||
if err != nil || card.Name != "echo" {
|
||||
t.Fatalf("Card() = %+v, err %v", card, err)
|
||||
}
|
||||
|
||||
reply, err := cl.Send(context.Background(), "ping")
|
||||
if err != nil {
|
||||
t.Fatalf("Send: %v", err)
|
||||
}
|
||||
if reply != "pong" {
|
||||
t.Errorf("Send reply = %q, want pong", reply)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientContinuesTaskAndConfiguresPush(t *testing.T) {
|
||||
card := Card("solo", "http://localhost:4000", "", []string{"task"})
|
||||
h := NewAgentHandler(card, func(_ context.Context, text string) (string, error) {
|
||||
return "echo:" + text, nil
|
||||
})
|
||||
ts := httptest.NewServer(h)
|
||||
defer ts.Close()
|
||||
|
||||
updates := make(chan Task, 1)
|
||||
push := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var task Task
|
||||
if err := json.NewDecoder(r.Body).Decode(&task); err != nil {
|
||||
t.Errorf("decode push task: %v", err)
|
||||
return
|
||||
}
|
||||
updates <- task
|
||||
w.WriteHeader(http.StatusAccepted)
|
||||
}))
|
||||
defer push.Close()
|
||||
|
||||
cl := NewClient(ts.URL)
|
||||
first, err := cl.SendMessage(context.Background(), Message{
|
||||
Parts: []Part{{Kind: "text", Text: "one"}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("first SendMessage: %v", err)
|
||||
}
|
||||
second, err := cl.SendMessage(context.Background(), Message{
|
||||
TaskID: first.ID,
|
||||
ContextID: first.ContextID,
|
||||
Parts: []Part{{Kind: "text", Text: "two"}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("second SendMessage: %v", err)
|
||||
}
|
||||
if second.ID != first.ID || second.ContextID != first.ContextID || len(second.History) != 4 {
|
||||
t.Fatalf("continued task = %+v, first %+v", second, first)
|
||||
}
|
||||
cfg := PushNotificationConfig{URL: push.URL}
|
||||
if err := cl.SetPushNotificationConfig(context.Background(), second.ID, cfg); err != nil {
|
||||
t.Fatalf("SetPushNotificationConfig: %v", err)
|
||||
}
|
||||
got, err := cl.PushNotificationConfig(context.Background(), second.ID)
|
||||
if err != nil {
|
||||
t.Fatalf("PushNotificationConfig: %v", err)
|
||||
}
|
||||
if got.URL != push.URL {
|
||||
t.Fatalf("PushNotificationConfig URL = %q, want %q", got.URL, push.URL)
|
||||
}
|
||||
select {
|
||||
case update := <-updates:
|
||||
if update.ID != second.ID {
|
||||
t.Fatalf("push update ID = %q, want %q", update.ID, second.ID)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("timed out waiting for push update")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientResubscribeStreamsRetainedAndLiveTask(t *testing.T) {
|
||||
d := newDispatcher()
|
||||
initial := &Task{ID: "task-1", ContextID: "ctx-1", Kind: "task", Status: TaskStatus{State: stateWorking, Timestamp: time.Now().UTC().Format(time.RFC3339)}}
|
||||
d.store(initial)
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
d.serve(w, r, func(context.Context, string) (string, error) { return "", nil })
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
tasks, errs := NewClient(ts.URL).Resubscribe(ctx, initial.ID)
|
||||
|
||||
first := <-tasks
|
||||
if first.ID != initial.ID || first.Status.State != stateWorking {
|
||||
t.Fatalf("first resubscribe task = %+v, want retained working task", first)
|
||||
}
|
||||
final := &Task{ID: initial.ID, ContextID: initial.ContextID, Kind: "task", Status: TaskStatus{State: stateCompleted, Timestamp: time.Now().UTC().Format(time.RFC3339)}, Artifacts: []Artifact{textArtifact("done")}}
|
||||
d.store(final)
|
||||
second := <-tasks
|
||||
if second.ID != final.ID || second.Status.State != stateCompleted || textOf(second.Artifacts[0].Parts) != "done" {
|
||||
t.Fatalf("second resubscribe task = %+v, want live completed task", second)
|
||||
}
|
||||
if _, ok := <-tasks; ok {
|
||||
t.Fatal("resubscribe task channel stayed open after terminal update")
|
||||
}
|
||||
select {
|
||||
case err := <-errs:
|
||||
if err != nil {
|
||||
t.Fatalf("resubscribe error = %v", err)
|
||||
}
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientSendMessageReturnsInputRequiredTask(t *testing.T) {
|
||||
card := Card("solo", "http://localhost:4000", "", []string{"task"})
|
||||
h := NewAgentHandler(card, func(context.Context, string) (string, error) {
|
||||
return "", errors.New("input-required: provide approval code")
|
||||
})
|
||||
ts := httptest.NewServer(h)
|
||||
defer ts.Close()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
task, err := NewClient(ts.URL).SendMessage(ctx, Message{Parts: []Part{{Kind: "text", Text: "approve?"}}})
|
||||
if err != nil {
|
||||
t.Fatalf("SendMessage: %v", err)
|
||||
}
|
||||
if task.Status.State != stateInputRequired || !strings.Contains(textOf(task.Artifacts[0].Parts), "provide approval code") {
|
||||
t.Fatalf("task = %+v, want input-required handoff", task)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,288 @@
|
||||
# API Gateway
|
||||
|
||||
The `gateway/api` package provides HTTP API gateway functionality for go-micro services. It translates HTTP requests into RPC calls and serves a web dashboard for browsing and calling services.
|
||||
|
||||
## Features
|
||||
|
||||
- **HTTP to RPC translation** - Call microservices via HTTP
|
||||
- **Web dashboard** - Browse and test services in the browser
|
||||
- **Authentication** - Optional JWT-based auth
|
||||
- **MCP integration** - Expose services to AI agents
|
||||
- **Flexible configuration** - Use in dev or production
|
||||
- **Service discovery** - Auto-detect services from registry
|
||||
|
||||
## Usage
|
||||
|
||||
### Basic Gateway
|
||||
|
||||
```go
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
|
||||
"go-micro.dev/v5/gateway/api"
|
||||
)
|
||||
|
||||
func main() {
|
||||
// Create gateway with custom handler
|
||||
gw, err := api.New(api.Options{
|
||||
Address: ":8080",
|
||||
Context: context.Background(),
|
||||
HandlerRegistrar: func(mux *http.ServeMux) error {
|
||||
// Register your HTTP handlers
|
||||
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Write([]byte("Hello from gateway"))
|
||||
})
|
||||
return nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
// Block until shutdown
|
||||
gw.Wait()
|
||||
}
|
||||
```
|
||||
|
||||
### Gateway with MCP
|
||||
|
||||
```go
|
||||
gw, err := api.New(api.Options{
|
||||
Address: ":8080",
|
||||
MCPEnabled: true,
|
||||
MCPAddress: ":3000", // MCP on separate port
|
||||
HandlerRegistrar: registerHandlers,
|
||||
})
|
||||
```
|
||||
|
||||
### Gateway with Authentication
|
||||
|
||||
```go
|
||||
gw, err := api.New(api.Options{
|
||||
Address: ":8080",
|
||||
AuthEnabled: true, // Handler registrar should add auth middleware
|
||||
HandlerRegistrar: func(mux *http.ServeMux) error {
|
||||
// Register handlers with auth middleware
|
||||
return registerAuthenticatedHandlers(mux)
|
||||
},
|
||||
})
|
||||
```
|
||||
|
||||
### Blocking Mode
|
||||
|
||||
```go
|
||||
// Run blocks until shutdown
|
||||
err := api.Run(api.Options{
|
||||
Address: ":8080",
|
||||
HandlerRegistrar: registerHandlers,
|
||||
})
|
||||
```
|
||||
|
||||
## Options
|
||||
|
||||
```go
|
||||
type Options struct {
|
||||
// Address to listen on (default: ":8080")
|
||||
Address string
|
||||
|
||||
// AuthEnabled signals that authentication is required
|
||||
// The HandlerRegistrar should implement auth checks
|
||||
AuthEnabled bool
|
||||
|
||||
// Context for cancellation (default: context.Background())
|
||||
Context context.Context
|
||||
|
||||
// Logger for gateway messages (default: log.Default())
|
||||
Logger *log.Logger
|
||||
|
||||
// HandlerRegistrar registers HTTP handlers on the mux
|
||||
HandlerRegistrar func(mux *http.ServeMux) error
|
||||
|
||||
// MCPEnabled enables the MCP gateway
|
||||
MCPEnabled bool
|
||||
|
||||
// MCPAddress is the address for MCP gateway (e.g., ":3000")
|
||||
MCPAddress string
|
||||
|
||||
// Registry for service discovery (default: registry.DefaultRegistry)
|
||||
Registry registry.Registry
|
||||
}
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────┐
|
||||
│ gateway/api Package │
|
||||
│ ┌────────────────────────────────────┐ │
|
||||
│ │ Gateway │ │
|
||||
│ │ - Manages HTTP server │ │
|
||||
│ │ - Calls HandlerRegistrar │ │
|
||||
│ │ - Starts MCP if enabled │ │
|
||||
│ └────────────────────────────────────┘ │
|
||||
└─────────────────────────────────────────┘
|
||||
↓ delegates to
|
||||
┌─────────────────────────────────────────┐
|
||||
│ HandlerRegistrar (user-provided) │
|
||||
│ ┌────────────────────────────────────┐ │
|
||||
│ │ func(mux *http.ServeMux) error │ │
|
||||
│ │ - Registers routes │ │
|
||||
│ │ - Adds middleware (auth, etc.) │ │
|
||||
│ │ - Sets up templates │ │
|
||||
│ └────────────────────────────────────┘ │
|
||||
└─────────────────────────────────────────┘
|
||||
↓ uses
|
||||
┌─────────────────────────────────────────┐
|
||||
│ Microservices (via RPC) │
|
||||
└─────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
## Integration
|
||||
|
||||
### In `micro run` (Development)
|
||||
|
||||
```go
|
||||
// cmd/micro/run/run.go
|
||||
import "go-micro.dev/v5/gateway/api"
|
||||
|
||||
gw, err := api.New(api.Options{
|
||||
Address: ":8080",
|
||||
AuthEnabled: false, // No auth in dev mode
|
||||
HandlerRegistrar: func(mux *http.ServeMux) error {
|
||||
// Register dev-mode handlers (no auth)
|
||||
mux.HandleFunc("/", dashboardHandler)
|
||||
mux.HandleFunc("/api/", apiHandler)
|
||||
return nil
|
||||
},
|
||||
})
|
||||
```
|
||||
|
||||
### In `micro server` (Production)
|
||||
|
||||
```go
|
||||
// cmd/micro/server/server.go
|
||||
import "go-micro.dev/v5/gateway/api"
|
||||
|
||||
gw, err := api.New(api.Options{
|
||||
Address: ":8080",
|
||||
AuthEnabled: true, // Auth required in production
|
||||
HandlerRegistrar: func(mux *http.ServeMux) error {
|
||||
// Register prod handlers with auth middleware
|
||||
mux.HandleFunc("/", authMiddleware(dashboardHandler))
|
||||
mux.HandleFunc("/api/", authMiddleware(apiHandler))
|
||||
return nil
|
||||
},
|
||||
})
|
||||
```
|
||||
|
||||
### Custom Application
|
||||
|
||||
```go
|
||||
// Your app
|
||||
import "go-micro.dev/v5/gateway/api"
|
||||
|
||||
func main() {
|
||||
gw, err := api.New(api.Options{
|
||||
Address: ":8080",
|
||||
HandlerRegistrar: func(mux *http.ServeMux) error {
|
||||
// Your custom handlers
|
||||
mux.HandleFunc("/health", healthHandler)
|
||||
mux.HandleFunc("/metrics", metricsHandler)
|
||||
mux.HandleFunc("/api/", proxyToServices)
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
||||
log.Println("Gateway running on :8080")
|
||||
gw.Wait()
|
||||
}
|
||||
```
|
||||
|
||||
## Comparison with Old Architecture
|
||||
|
||||
### Before (Duplicated Code)
|
||||
|
||||
```
|
||||
cmd/micro/run/gateway/
|
||||
└── gateway.go (300+ lines)
|
||||
|
||||
cmd/micro/server/
|
||||
└── gateway.go (150+ lines)
|
||||
|
||||
❌ Code duplication
|
||||
❌ Inconsistent behavior
|
||||
❌ Hard to reuse
|
||||
```
|
||||
|
||||
### After (Unified)
|
||||
|
||||
```
|
||||
gateway/api/
|
||||
└── gateway.go (150 lines, reusable)
|
||||
|
||||
cmd/micro/server/
|
||||
└── gateway.go (70 lines, compatibility wrapper)
|
||||
|
||||
cmd/micro/run/
|
||||
└── Uses api.New() directly
|
||||
|
||||
✅ Single source of truth
|
||||
✅ Consistent behavior
|
||||
✅ Easy to reuse in custom apps
|
||||
```
|
||||
|
||||
## Benefits
|
||||
|
||||
1. **Reusability** - Use in any Go application, not just micro CLI
|
||||
2. **Testability** - Easy to test with custom handler registrars
|
||||
3. **Flexibility** - Supports different configurations (dev, prod, custom)
|
||||
4. **Consistency** - Same gateway code for all use cases
|
||||
5. **Maintainability** - One place to fix bugs and add features
|
||||
|
||||
## Migration Guide
|
||||
|
||||
### From `cmd/micro/server/gateway.go`
|
||||
|
||||
**Before:**
|
||||
```go
|
||||
import "go-micro.dev/v5/cmd/micro/server"
|
||||
|
||||
gw, err := server.StartGateway(server.GatewayOptions{
|
||||
Address: ":8080",
|
||||
AuthEnabled: true,
|
||||
Store: myStore,
|
||||
})
|
||||
```
|
||||
|
||||
**After:**
|
||||
```go
|
||||
import "go-micro.dev/v5/gateway/api"
|
||||
|
||||
gw, err := api.New(api.Options{
|
||||
Address: ":8080",
|
||||
AuthEnabled: true,
|
||||
HandlerRegistrar: func(mux *http.ServeMux) error {
|
||||
// Register your handlers
|
||||
// Pass store as closure
|
||||
return registerHandlers(mux, myStore)
|
||||
},
|
||||
})
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
See:
|
||||
- `cmd/micro/server/gateway.go` - Production gateway with auth
|
||||
- `cmd/micro/run/run.go` - Development gateway without auth
|
||||
- `examples/gateway/` - Custom gateway examples (coming soon)
|
||||
|
||||
## License
|
||||
|
||||
Apache 2.0
|
||||
@@ -0,0 +1,155 @@
|
||||
// Package api provides HTTP API gateway functionality for go-micro services.
|
||||
//
|
||||
// The API gateway translates HTTP requests into RPC calls and serves a web dashboard
|
||||
// for browsing and calling services. It can be used in development (micro run) or
|
||||
// production (micro server) with optional authentication.
|
||||
package api
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"go-micro.dev/v6/gateway/mcp"
|
||||
"go-micro.dev/v6/registry"
|
||||
)
|
||||
|
||||
// Options configures the HTTP API gateway
|
||||
type Options struct {
|
||||
// Address to listen on (e.g., ":8080")
|
||||
Address string
|
||||
|
||||
// AuthEnabled controls whether authentication is required
|
||||
// If true, the HandlerRegistrar should include auth middleware
|
||||
AuthEnabled bool
|
||||
|
||||
// Context for cancellation (if nil, uses background context)
|
||||
Context context.Context
|
||||
|
||||
// Logger for gateway messages (if nil, uses log.Default())
|
||||
Logger *log.Logger
|
||||
|
||||
// HandlerRegistrar is called to register HTTP handlers on the mux
|
||||
// This allows different configurations (dev vs prod) to register different handlers
|
||||
HandlerRegistrar func(mux *http.ServeMux) error
|
||||
|
||||
// MCPEnabled controls whether to start MCP gateway
|
||||
MCPEnabled bool
|
||||
|
||||
// MCPAddress is the address for MCP gateway (e.g., ":3000")
|
||||
MCPAddress string
|
||||
|
||||
// Registry for service discovery (if nil, uses registry.DefaultRegistry)
|
||||
Registry registry.Registry
|
||||
}
|
||||
|
||||
// Gateway represents a running HTTP API gateway server
|
||||
type Gateway struct {
|
||||
opts Options
|
||||
server *http.Server
|
||||
mux *http.ServeMux
|
||||
}
|
||||
|
||||
// New creates a new gateway with the given options and starts it.
|
||||
// Returns immediately after starting the server in a goroutine.
|
||||
// Use Wait() or Run() to block until the server stops.
|
||||
func New(opts Options) (*Gateway, error) {
|
||||
// Set defaults
|
||||
if opts.Address == "" {
|
||||
opts.Address = ":8080"
|
||||
}
|
||||
if opts.Context == nil {
|
||||
opts.Context = context.Background()
|
||||
}
|
||||
if opts.Logger == nil {
|
||||
opts.Logger = log.Default()
|
||||
}
|
||||
if opts.Registry == nil {
|
||||
opts.Registry = registry.DefaultRegistry
|
||||
}
|
||||
|
||||
// Create a new mux for this gateway instance
|
||||
mux := http.NewServeMux()
|
||||
|
||||
// Register handlers using the provided registrar
|
||||
if opts.HandlerRegistrar != nil {
|
||||
if err := opts.HandlerRegistrar(mux); err != nil {
|
||||
return nil, fmt.Errorf("failed to register handlers: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Create HTTP server
|
||||
server := &http.Server{
|
||||
Addr: opts.Address,
|
||||
Handler: mux,
|
||||
}
|
||||
|
||||
gw := &Gateway{
|
||||
opts: opts,
|
||||
server: server,
|
||||
mux: mux,
|
||||
}
|
||||
|
||||
// Start MCP gateway if enabled
|
||||
if opts.MCPEnabled && opts.MCPAddress != "" {
|
||||
go func() {
|
||||
if err := mcp.ListenAndServe(opts.MCPAddress, mcp.Options{
|
||||
Registry: opts.Registry,
|
||||
Context: opts.Context,
|
||||
Logger: opts.Logger,
|
||||
}); err != nil {
|
||||
opts.Logger.Printf("[mcp] MCP gateway error: %v", err)
|
||||
}
|
||||
}()
|
||||
opts.Logger.Printf("[mcp] MCP gateway enabled on %s", opts.MCPAddress)
|
||||
}
|
||||
|
||||
// Start server in background
|
||||
go func() {
|
||||
opts.Logger.Printf("[gateway] Listening on %s (auth: %v)", opts.Address, opts.AuthEnabled)
|
||||
if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||||
opts.Logger.Printf("[gateway] Server error: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
return gw, nil
|
||||
}
|
||||
|
||||
// Run creates and starts a gateway, blocking until it stops.
|
||||
// This is a convenience function equivalent to New() + Wait().
|
||||
func Run(opts Options) error {
|
||||
gw, err := New(opts)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return gw.Wait()
|
||||
}
|
||||
|
||||
// Wait blocks until the server is shut down
|
||||
func (g *Gateway) Wait() error {
|
||||
<-g.opts.Context.Done()
|
||||
return g.Stop()
|
||||
}
|
||||
|
||||
// Stop gracefully shuts down the gateway
|
||||
func (g *Gateway) Stop() error {
|
||||
if g.server != nil {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
return g.server.Shutdown(ctx)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Addr returns the address the gateway is listening on
|
||||
func (g *Gateway) Addr() string {
|
||||
return g.opts.Address
|
||||
}
|
||||
|
||||
// Mux returns the underlying HTTP mux for this gateway
|
||||
// This can be used to register additional handlers after creation
|
||||
func (g *Gateway) Mux() *http.ServeMux {
|
||||
return g.mux
|
||||
}
|
||||
@@ -0,0 +1,402 @@
|
||||
# MCP Tool Documentation
|
||||
|
||||
This document explains how to document your go-micro services so that AI agents can understand them better.
|
||||
|
||||
## Overview
|
||||
|
||||
The MCP gateway automatically exposes your microservices as tools that AI agents (like Claude) can call. By adding proper documentation to your service handlers, you help agents understand:
|
||||
|
||||
- **What the tool does** - The purpose and behavior
|
||||
- **What parameters it needs** - Types, formats, constraints
|
||||
- **What it returns** - Response structure and meaning
|
||||
- **How to use it** - Example inputs and outputs
|
||||
|
||||
## Documentation Methods
|
||||
|
||||
go-micro **automatically extracts documentation** from your Go doc comments at registration time. You don't need to write any extra code!
|
||||
|
||||
### 1. Go Doc Comments (Automatic - Recommended)
|
||||
|
||||
Just write standard Go documentation comments on your handler methods:
|
||||
|
||||
```go
|
||||
// GetUser retrieves a user by ID from the database. Returns full profile including email, name, and preferences.
|
||||
//
|
||||
// @example {"id": "user-1"}
|
||||
func (s *UserService) GetUser(ctx context.Context, req *GetUserRequest, rsp *GetUserResponse) error {
|
||||
// implementation
|
||||
}
|
||||
```
|
||||
|
||||
When you register the handler, go-micro automatically:
|
||||
- Extracts the doc comment as the tool description
|
||||
- Parses the `@example` tag for example inputs
|
||||
- Registers everything in the service registry
|
||||
- Makes it available to the MCP gateway
|
||||
|
||||
**Supported Tags:**
|
||||
- `@example <json>` - Example JSON input (highly recommended for AI agents)
|
||||
|
||||
**That's it!** No extra registration code needed:
|
||||
|
||||
```go
|
||||
// Documentation is extracted automatically from method comments
|
||||
handler := service.Server().NewHandler(new(UserService))
|
||||
service.Server().Handle(handler)
|
||||
```
|
||||
|
||||
### 2. Manual Registration (Optional Override)
|
||||
|
||||
For more control or to override auto-extracted docs, use `server.WithEndpointDocs()`:
|
||||
|
||||
```go
|
||||
handler := service.Server().NewHandler(
|
||||
new(UserService),
|
||||
server.WithEndpointDocs(map[string]server.EndpointDoc{
|
||||
"UserService.GetUser": {
|
||||
Description: "Custom description that overrides the comment",
|
||||
Example: `{"id": "user-123"}`,
|
||||
},
|
||||
}),
|
||||
)
|
||||
```
|
||||
|
||||
Manual metadata **takes precedence** over auto-extracted comments.
|
||||
|
||||
### 3. Endpoint Scopes (Auth)
|
||||
|
||||
Use `server.WithEndpointScopes()` to declare the auth scopes required for each
|
||||
endpoint. The MCP gateway reads these from the registry and enforces them when
|
||||
an `Auth` provider is configured.
|
||||
|
||||
```go
|
||||
handler := service.Server().NewHandler(
|
||||
new(BlogService),
|
||||
server.WithEndpointScopes("Blog.Create", "blog:write"),
|
||||
server.WithEndpointScopes("Blog.Delete", "blog:write", "blog:admin"),
|
||||
server.WithEndpointScopes("Blog.Read", "blog:read"),
|
||||
)
|
||||
```
|
||||
|
||||
Scopes are stored as comma-separated values in endpoint metadata (`"scopes"` key)
|
||||
and are propagated through the service registry just like descriptions and examples.
|
||||
|
||||
#### Gateway-Level Scope Overrides
|
||||
|
||||
An operator can also define or override scopes at the MCP gateway without
|
||||
modifying individual services. This is useful for centralized policy management:
|
||||
|
||||
```go
|
||||
mcp.Serve(mcp.Options{
|
||||
Registry: reg,
|
||||
Auth: authProvider,
|
||||
Scopes: map[string][]string{
|
||||
"blog.Blog.Create": {"blog:write"},
|
||||
"blog.Blog.Delete": {"blog:admin"},
|
||||
},
|
||||
})
|
||||
```
|
||||
|
||||
Gateway-level scopes **take precedence** over service-level scopes.
|
||||
|
||||
### 4. Struct Tags (For Field Descriptions)
|
||||
|
||||
Add descriptions to struct fields using the `description` tag:
|
||||
|
||||
```go
|
||||
type User struct {
|
||||
ID string `json:"id" description:"User's unique identifier (UUID format)"`
|
||||
Name string `json:"name" description:"User's full name"`
|
||||
Email string `json:"email" description:"User's email address"`
|
||||
Age int `json:"age,omitempty" description:"User's age (optional)"`
|
||||
}
|
||||
```
|
||||
|
||||
The `description` tag is used to generate parameter descriptions in the JSON Schema.
|
||||
|
||||
## How It Works
|
||||
|
||||
### Automatic Extraction Pipeline
|
||||
|
||||
```
|
||||
1. Handler Registration (Your Service)
|
||||
├─> You write Go doc comments on methods
|
||||
├─> Call service.Server().NewHandler(yourHandler)
|
||||
└─> go-micro automatically parses source files using go/ast
|
||||
|
||||
2. Documentation Extraction (Automatic)
|
||||
├─> Read Go doc comments from handler method source
|
||||
├─> Parse @example tags for sample inputs
|
||||
├─> Extract struct tag descriptions
|
||||
└─> Merge with any manual metadata (manual wins)
|
||||
|
||||
3. Service Registry
|
||||
├─> Store endpoint metadata in registry.Endpoint.Metadata
|
||||
├─> Metadata distributed with service information
|
||||
└─> Available to all components (gateway, discovery, etc.)
|
||||
|
||||
4. MCP Gateway Discovery
|
||||
├─> Query registry for services and endpoints
|
||||
├─> Read description and example from endpoint.Metadata
|
||||
└─> Generate JSON Schema with documentation
|
||||
|
||||
5. Tool Creation
|
||||
└─> Create MCP tool with rich description for AI agents
|
||||
```
|
||||
|
||||
### Example Output
|
||||
|
||||
For a documented handler, the MCP gateway generates:
|
||||
|
||||
```json
|
||||
{
|
||||
"name": "users.UserService.GetUser",
|
||||
"description": "GetUser retrieves a user by ID from the database. Returns full profile including email, name, and preferences.",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"description": "This endpoint fetches a user's complete profile...",
|
||||
"properties": {
|
||||
"id": {
|
||||
"type": "string",
|
||||
"description": "User ID in UUID format (e.g., \"123e4567-e89b-12d3-a456-426614174000\")"
|
||||
}
|
||||
},
|
||||
"required": ["id"],
|
||||
"examples": [
|
||||
"{\"id\": \"user-1\"}"
|
||||
]
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Best Practices
|
||||
|
||||
### Write for AI, Not Just Humans
|
||||
|
||||
AI agents parse your documentation literally. Be explicit:
|
||||
|
||||
**✅ Good:**
|
||||
```go
|
||||
// GetUser retrieves a user by their unique ID from the database.
|
||||
// Returns the user's full profile including name, email, and preferences.
|
||||
// If the user doesn't exist, returns an error with status 404.
|
||||
//
|
||||
// @param id {string} User ID in UUID v4 format (e.g., "123e4567-e89b-12d3-a456-426614174000")
|
||||
// @return {User} User object with all profile fields populated
|
||||
```
|
||||
|
||||
**❌ Bad:**
|
||||
```go
|
||||
// Gets a user
|
||||
func GetUser(...) // No details, no context
|
||||
```
|
||||
|
||||
### Specify Formats and Constraints
|
||||
|
||||
Tell agents exactly what format you expect:
|
||||
|
||||
**✅ Good:**
|
||||
```go
|
||||
// @param email {string} Email address in RFC 5322 format (must contain @ and domain)
|
||||
// @param age {number} User's age (integer between 0-150)
|
||||
// @param phone {string} Phone number in E.164 format (e.g., "+14155552671")
|
||||
```
|
||||
|
||||
**❌ Bad:**
|
||||
```go
|
||||
// @param email {string} The email
|
||||
// @param age {number} Age
|
||||
```
|
||||
|
||||
### Provide Real Examples
|
||||
|
||||
Show agents actual valid inputs:
|
||||
|
||||
**✅ Good:**
|
||||
```go
|
||||
// @example
|
||||
// {
|
||||
// "name": "Alice Smith",
|
||||
// "email": "alice@example.com",
|
||||
// "age": 30,
|
||||
// "phone": "+14155552671"
|
||||
// }
|
||||
```
|
||||
|
||||
**❌ Bad:**
|
||||
```go
|
||||
// @example
|
||||
// {
|
||||
// "name": "string",
|
||||
// "email": "string"
|
||||
// }
|
||||
```
|
||||
|
||||
### Document Error Cases
|
||||
|
||||
Tell agents what can go wrong:
|
||||
|
||||
```go
|
||||
// GetUser retrieves a user by ID.
|
||||
//
|
||||
// Returns error if:
|
||||
// - User ID is not a valid UUID
|
||||
// - User does not exist (404)
|
||||
// - Database is unavailable (503)
|
||||
//
|
||||
// @param id {string} User ID in UUID format
|
||||
```
|
||||
|
||||
### Use Descriptive Names
|
||||
|
||||
Field names should be self-explanatory:
|
||||
|
||||
**✅ Good:**
|
||||
```go
|
||||
type CreateUserRequest struct {
|
||||
FullName string `json:"full_name" description:"User's complete name"`
|
||||
EmailAddress string `json:"email_address" description:"Primary email for contact"`
|
||||
DateOfBirth string `json:"date_of_birth" description:"Birth date in YYYY-MM-DD format"`
|
||||
}
|
||||
```
|
||||
|
||||
**❌ Bad:**
|
||||
```go
|
||||
type CreateUserRequest struct {
|
||||
N string `json:"n"` // What is n?
|
||||
E string `json:"e"` // What is e?
|
||||
D string `json:"d"` // What is d?
|
||||
}
|
||||
```
|
||||
|
||||
## Impact on Agent Performance
|
||||
|
||||
### Without Documentation
|
||||
|
||||
```
|
||||
Agent: "I need to call GetUser but I don't know what format the ID should be.
|
||||
Is it a number? A string? A UUID? Let me try..."
|
||||
|
||||
❌ Calls with: {"id": 123}
|
||||
❌ Calls with: {"id": "user123"}
|
||||
❌ Calls with: {"id": "abc"}
|
||||
✅ Calls with: {"id": "550e8400-e29b-41d4-a716-446655440000"} (after 4 attempts)
|
||||
```
|
||||
|
||||
### With Documentation
|
||||
|
||||
```
|
||||
Agent: "GetUser needs an ID in UUID format. The example shows the format.
|
||||
I'll use a valid UUID."
|
||||
|
||||
✅ Calls with: {"id": "550e8400-e29b-41d4-a716-446655440000"} (first attempt)
|
||||
```
|
||||
|
||||
**Result:**
|
||||
- **75% fewer failed calls**
|
||||
- **Faster task completion**
|
||||
- **Better user experience**
|
||||
|
||||
## Parser Implementation
|
||||
|
||||
The MCP gateway uses several parsers:
|
||||
|
||||
### 1. Go Doc Parser (`parseServiceDocs`)
|
||||
- Extracts godoc comments from handler methods
|
||||
- Parses JSDoc-style tags
|
||||
- Returns `ToolDescription` struct
|
||||
|
||||
### 2. Struct Tag Parser (`ParseStructTags`)
|
||||
- Reads `description` tags from struct fields
|
||||
- Generates JSON Schema with field descriptions
|
||||
- Marks required vs optional fields (omitempty)
|
||||
|
||||
### 3. Comment Parser (`ParseGoDocComment`)
|
||||
- Regex-based extraction of @param, @return, @example tags
|
||||
- Splits summary from detailed description
|
||||
- Builds structured documentation
|
||||
|
||||
### 4. Type Mapper (`reflectTypeToJSONType`)
|
||||
- Converts Go types to JSON Schema types
|
||||
- Handles: string, int, float, bool, array, object
|
||||
- Used for automatic schema generation
|
||||
|
||||
## Examples
|
||||
|
||||
See complete examples in:
|
||||
- `examples/mcp/documented/` - Fully documented service
|
||||
- `examples/auth/` - Auth service with documentation
|
||||
- `examples/hello-world/` - Basic service
|
||||
|
||||
## Testing Documentation
|
||||
|
||||
### 1. List Tools
|
||||
|
||||
```bash
|
||||
curl http://localhost:3000/mcp/tools | jq '.tools[0]'
|
||||
```
|
||||
|
||||
Verify the description and schema are correct.
|
||||
|
||||
### 2. Use with Claude Code
|
||||
|
||||
Add to your Claude Code config and ask Claude to use your service. Claude will show you how it interprets your documentation.
|
||||
|
||||
### 3. Check Examples Work
|
||||
|
||||
Try the examples from your `@example` tags:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:3000/mcp/call \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"tool": "users.UserService.GetUser",
|
||||
"input": <your-example-json>
|
||||
}'
|
||||
```
|
||||
|
||||
## Future Enhancements
|
||||
|
||||
Planned improvements:
|
||||
|
||||
- [ ] Auto-extract examples from test files
|
||||
- [ ] Validate documentation completeness (lint)
|
||||
- [ ] Generate documentation from OpenAPI specs
|
||||
- [ ] Support custom validation rules in tags
|
||||
- [ ] Interactive documentation editor
|
||||
|
||||
## FAQ
|
||||
|
||||
**Q: Do I need to document every field?**
|
||||
A: Document fields that are ambiguous or have constraints. Self-explanatory fields can rely on the field name.
|
||||
|
||||
**Q: Will this slow down my service?**
|
||||
A: No. Documentation is parsed once at startup when the MCP gateway discovers services.
|
||||
|
||||
**Q: Can I use OpenAPI/Swagger specs instead?**
|
||||
A: Not yet, but it's planned. For now, use Go comments and struct tags.
|
||||
|
||||
**Q: What if I don't document my handlers?**
|
||||
A: The MCP gateway will still work, generating basic descriptions from method names and types. But agents will perform better with documentation.
|
||||
|
||||
**Q: How do I know if my documentation is good?**
|
||||
A: Test it with Claude Code. If Claude understands your service and calls it correctly on the first try, your documentation is good!
|
||||
|
||||
**Q: How do I add auth scopes to my endpoints?**
|
||||
A: Use `server.WithEndpointScopes()` when registering your handler:
|
||||
|
||||
```go
|
||||
handler := service.Server().NewHandler(
|
||||
new(MyService),
|
||||
server.WithEndpointScopes("MyService.Create", "write"),
|
||||
)
|
||||
```
|
||||
|
||||
Or define scopes at the gateway level using `Scopes` in `mcp.Options`.
|
||||
|
||||
**Q: Can I set scopes at the gateway without changing services?**
|
||||
A: Yes. Use the `Scopes` option on `mcp.Options` to define or override scopes for any tool at the gateway layer. This is useful for centralized policy management.
|
||||
|
||||
## License
|
||||
|
||||
Apache 2.0
|
||||
@@ -0,0 +1,258 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"go-micro.dev/v6/auth"
|
||||
"go-micro.dev/v6/client"
|
||||
"go-micro.dev/v6/registry"
|
||||
)
|
||||
|
||||
// benchServer creates a Server with N pre-populated tools.
|
||||
func benchServer(n int, opts Options) *Server {
|
||||
if opts.Logger == nil {
|
||||
opts.Logger = log.New(log.Writer(), "", 0)
|
||||
}
|
||||
if opts.Context == nil {
|
||||
opts.Context = context.Background()
|
||||
}
|
||||
if opts.Client == nil {
|
||||
opts.Client = client.DefaultClient
|
||||
}
|
||||
if opts.Registry == nil {
|
||||
opts.Registry = registry.DefaultRegistry
|
||||
}
|
||||
|
||||
s := &Server{
|
||||
opts: opts,
|
||||
tools: make(map[string]*Tool, n),
|
||||
limiters: make(map[string]*rateLimiter),
|
||||
}
|
||||
|
||||
for i := 0; i < n; i++ {
|
||||
name := toolName(i)
|
||||
s.tools[name] = &Tool{
|
||||
Name: name,
|
||||
Description: "Benchmark tool " + name,
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "Resource identifier",
|
||||
},
|
||||
},
|
||||
"required": []interface{}{"id"},
|
||||
},
|
||||
Service: "bench",
|
||||
Endpoint: "Handler.Method",
|
||||
}
|
||||
}
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
func toolName(i int) string {
|
||||
return "bench.Handler.Method" + string(rune('A'+i%26))
|
||||
}
|
||||
|
||||
// --- Benchmarks ---
|
||||
|
||||
// BenchmarkListTools measures tool listing throughput.
|
||||
// This is the most common MCP operation — agents call it on every session start.
|
||||
func BenchmarkListTools(b *testing.B) {
|
||||
for _, numTools := range []int{10, 50, 100} {
|
||||
b.Run(toolCountLabel(numTools), func(b *testing.B) {
|
||||
s := benchServer(numTools, Options{})
|
||||
req := httptest.NewRequest(http.MethodGet, "/mcp/tools", nil)
|
||||
|
||||
b.ResetTimer()
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
w := httptest.NewRecorder()
|
||||
s.handleListTools(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
b.Fatalf("unexpected status %d", w.Code)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkListToolsParallel measures concurrent tool listing.
|
||||
func BenchmarkListToolsParallel(b *testing.B) {
|
||||
s := benchServer(50, Options{})
|
||||
req := httptest.NewRequest(http.MethodGet, "/mcp/tools", nil)
|
||||
|
||||
b.ResetTimer()
|
||||
b.ReportAllocs()
|
||||
b.RunParallel(func(pb *testing.PB) {
|
||||
for pb.Next() {
|
||||
w := httptest.NewRecorder()
|
||||
s.handleListTools(w, req)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// BenchmarkToolLookup measures tool name resolution from the tools map.
|
||||
func BenchmarkToolLookup(b *testing.B) {
|
||||
for _, numTools := range []int{10, 50, 100, 500} {
|
||||
b.Run(toolCountLabel(numTools), func(b *testing.B) {
|
||||
s := benchServer(numTools, Options{})
|
||||
name := toolName(numTools / 2) // look up a tool in the middle
|
||||
|
||||
b.ResetTimer()
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
s.toolsMu.RLock()
|
||||
_, ok := s.tools[name]
|
||||
s.toolsMu.RUnlock()
|
||||
if !ok {
|
||||
b.Fatal("tool not found")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkAuthInspect measures auth token inspection overhead.
|
||||
func BenchmarkAuthInspect(b *testing.B) {
|
||||
ma := &mockAuth{
|
||||
accounts: map[string]*auth.Account{
|
||||
"valid-token": {
|
||||
ID: "bench-user",
|
||||
Scopes: []string{"read", "write"},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
b.ResetTimer()
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
acc, err := ma.Inspect("valid-token")
|
||||
if err != nil || acc.ID != "bench-user" {
|
||||
b.Fatal("unexpected result")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkScopeCheck measures scope validation overhead per tool call.
|
||||
func BenchmarkScopeCheck(b *testing.B) {
|
||||
accountScopes := []string{"users:read", "users:write", "orders:read", "admin"}
|
||||
requiredScopes := []string{"users:write"}
|
||||
|
||||
b.ResetTimer()
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
hasScope(accountScopes, requiredScopes)
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkAuditRecord measures audit record creation overhead.
|
||||
func BenchmarkAuditRecord(b *testing.B) {
|
||||
var records int
|
||||
s := benchServer(10, Options{
|
||||
AuditFunc: func(r AuditRecord) {
|
||||
records++
|
||||
},
|
||||
})
|
||||
|
||||
b.ResetTimer()
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
s.opts.AuditFunc(AuditRecord{
|
||||
TraceID: "trace-123",
|
||||
Tool: "bench.Handler.MethodA",
|
||||
AccountID: "user-1",
|
||||
Allowed: true,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkRateLimiter measures rate limiter check overhead.
|
||||
func BenchmarkRateLimiter(b *testing.B) {
|
||||
s := benchServer(10, Options{
|
||||
RateLimit: &RateLimitConfig{
|
||||
RequestsPerSecond: 1000000, // Very high so it doesn't block
|
||||
Burst: 1000000,
|
||||
},
|
||||
})
|
||||
// Initialize limiters for tools
|
||||
for name := range s.tools {
|
||||
s.limiters[name] = newRateLimiter(s.opts.RateLimit.RequestsPerSecond, s.opts.RateLimit.Burst)
|
||||
}
|
||||
name := toolName(0)
|
||||
|
||||
b.ResetTimer()
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
s.limitersMu.RLock()
|
||||
l := s.limiters[name]
|
||||
s.limitersMu.RUnlock()
|
||||
l.Allow()
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkJSONEncodeTool measures JSON serialization of tool definitions.
|
||||
func BenchmarkJSONEncodeTool(b *testing.B) {
|
||||
tool := &Tool{
|
||||
Name: "myservice.Users.GetUser",
|
||||
Description: "Retrieve a user by their unique ID. Returns the full profile.",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "User ID in UUID format",
|
||||
},
|
||||
},
|
||||
"required": []interface{}{"id"},
|
||||
},
|
||||
Scopes: []string{"users:read"},
|
||||
Service: "myservice",
|
||||
Endpoint: "Users.GetUser",
|
||||
}
|
||||
|
||||
b.ResetTimer()
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var buf bytes.Buffer
|
||||
json.NewEncoder(&buf).Encode(tool)
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkJSONDecodeCallRequest measures parsing of incoming tool call requests.
|
||||
func BenchmarkJSONDecodeCallRequest(b *testing.B) {
|
||||
body := []byte(`{"tool":"myservice.Users.GetUser","arguments":{"id":"user-123"}}`)
|
||||
|
||||
b.ResetTimer()
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var req struct {
|
||||
Tool string `json:"tool"`
|
||||
Arguments map[string]interface{} `json:"arguments"`
|
||||
}
|
||||
json.Unmarshal(body, &req)
|
||||
}
|
||||
}
|
||||
|
||||
// --- Helpers ---
|
||||
|
||||
func toolCountLabel(n int) string {
|
||||
switch {
|
||||
case n >= 500:
|
||||
return "500_tools"
|
||||
case n >= 100:
|
||||
return "100_tools"
|
||||
case n >= 50:
|
||||
return "50_tools"
|
||||
default:
|
||||
return "10_tools"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,147 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// CircuitBreakerConfig configures circuit breaking for the MCP gateway.
|
||||
// When a downstream service fails repeatedly, the circuit opens and
|
||||
// subsequent calls are rejected immediately until the service recovers.
|
||||
type CircuitBreakerConfig struct {
|
||||
// MaxFailures is the number of consecutive failures before the circuit opens.
|
||||
// Default: 5
|
||||
MaxFailures int
|
||||
|
||||
// Timeout is how long the circuit stays open before allowing a probe request.
|
||||
// Default: 30s
|
||||
Timeout time.Duration
|
||||
|
||||
// MaxHalfOpen is the number of probe requests allowed in the half-open state.
|
||||
// If they all succeed, the circuit closes. If any fail, it re-opens.
|
||||
// Default: 1
|
||||
MaxHalfOpen int
|
||||
}
|
||||
|
||||
// circuitState represents the state of a circuit breaker.
|
||||
type circuitState int
|
||||
|
||||
const (
|
||||
circuitClosed circuitState = iota // healthy, requests flow through
|
||||
circuitOpen // tripped, requests are rejected
|
||||
circuitHalfOpen // testing recovery with limited requests
|
||||
)
|
||||
|
||||
func (s circuitState) String() string {
|
||||
switch s {
|
||||
case circuitClosed:
|
||||
return "closed"
|
||||
case circuitOpen:
|
||||
return "open"
|
||||
case circuitHalfOpen:
|
||||
return "half-open"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// circuitBreaker tracks failure state for a single tool/service endpoint.
|
||||
type circuitBreaker struct {
|
||||
mu sync.Mutex
|
||||
state circuitState
|
||||
failures int
|
||||
maxFailures int
|
||||
timeout time.Duration
|
||||
maxHalfOpen int
|
||||
halfOpenUsed int
|
||||
lastFailure time.Time
|
||||
}
|
||||
|
||||
func newCircuitBreaker(cfg CircuitBreakerConfig) *circuitBreaker {
|
||||
maxFailures := cfg.MaxFailures
|
||||
if maxFailures <= 0 {
|
||||
maxFailures = 5
|
||||
}
|
||||
timeout := cfg.Timeout
|
||||
if timeout <= 0 {
|
||||
timeout = 30 * time.Second
|
||||
}
|
||||
maxHalfOpen := cfg.MaxHalfOpen
|
||||
if maxHalfOpen <= 0 {
|
||||
maxHalfOpen = 1
|
||||
}
|
||||
return &circuitBreaker{
|
||||
state: circuitClosed,
|
||||
maxFailures: maxFailures,
|
||||
timeout: timeout,
|
||||
maxHalfOpen: maxHalfOpen,
|
||||
}
|
||||
}
|
||||
|
||||
// Allow checks whether a request should be allowed through.
|
||||
// Returns nil if allowed, error if the circuit is open.
|
||||
func (cb *circuitBreaker) Allow() error {
|
||||
cb.mu.Lock()
|
||||
defer cb.mu.Unlock()
|
||||
|
||||
switch cb.state {
|
||||
case circuitClosed:
|
||||
return nil
|
||||
case circuitOpen:
|
||||
if time.Since(cb.lastFailure) > cb.timeout {
|
||||
cb.state = circuitHalfOpen
|
||||
cb.halfOpenUsed = 0
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("circuit breaker open (consecutive failures: %d)", cb.failures)
|
||||
case circuitHalfOpen:
|
||||
if cb.halfOpenUsed < cb.maxHalfOpen {
|
||||
cb.halfOpenUsed++
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("circuit breaker half-open (probe limit reached)")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RecordSuccess records a successful call. If half-open, closes the circuit.
|
||||
func (cb *circuitBreaker) RecordSuccess() {
|
||||
cb.mu.Lock()
|
||||
defer cb.mu.Unlock()
|
||||
|
||||
cb.failures = 0
|
||||
cb.state = circuitClosed
|
||||
}
|
||||
|
||||
// RecordFailure records a failed call. May trip the circuit open.
|
||||
func (cb *circuitBreaker) RecordFailure() {
|
||||
cb.mu.Lock()
|
||||
defer cb.mu.Unlock()
|
||||
|
||||
cb.failures++
|
||||
cb.lastFailure = time.Now()
|
||||
|
||||
switch cb.state {
|
||||
case circuitClosed:
|
||||
if cb.failures >= cb.maxFailures {
|
||||
cb.state = circuitOpen
|
||||
}
|
||||
case circuitHalfOpen:
|
||||
// Probe failed, re-open
|
||||
cb.state = circuitOpen
|
||||
}
|
||||
}
|
||||
|
||||
// State returns the current circuit state.
|
||||
func (cb *circuitBreaker) State() circuitState {
|
||||
cb.mu.Lock()
|
||||
defer cb.mu.Unlock()
|
||||
|
||||
// Check for automatic transition from open -> half-open
|
||||
if cb.state == circuitOpen && time.Since(cb.lastFailure) > cb.timeout {
|
||||
cb.state = circuitHalfOpen
|
||||
cb.halfOpenUsed = 0
|
||||
}
|
||||
return cb.state
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestCircuitBreaker_ClosedAllowsRequests(t *testing.T) {
|
||||
cb := newCircuitBreaker(CircuitBreakerConfig{MaxFailures: 3, Timeout: time.Second})
|
||||
if err := cb.Allow(); err != nil {
|
||||
t.Fatalf("expected closed circuit to allow, got: %v", err)
|
||||
}
|
||||
if cb.State() != circuitClosed {
|
||||
t.Fatalf("expected closed state, got %s", cb.State())
|
||||
}
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_OpensAfterMaxFailures(t *testing.T) {
|
||||
cb := newCircuitBreaker(CircuitBreakerConfig{MaxFailures: 3, Timeout: time.Minute})
|
||||
|
||||
// 2 failures: still closed
|
||||
cb.RecordFailure()
|
||||
cb.RecordFailure()
|
||||
if cb.State() != circuitClosed {
|
||||
t.Fatalf("expected closed after 2 failures, got %s", cb.State())
|
||||
}
|
||||
|
||||
// 3rd failure: trips open
|
||||
cb.RecordFailure()
|
||||
if cb.State() != circuitOpen {
|
||||
t.Fatalf("expected open after 3 failures, got %s", cb.State())
|
||||
}
|
||||
|
||||
// Requests should be rejected
|
||||
if err := cb.Allow(); err == nil {
|
||||
t.Fatal("expected open circuit to reject")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_SuccessResetsFailures(t *testing.T) {
|
||||
cb := newCircuitBreaker(CircuitBreakerConfig{MaxFailures: 3, Timeout: time.Minute})
|
||||
|
||||
cb.RecordFailure()
|
||||
cb.RecordFailure()
|
||||
cb.RecordSuccess() // resets
|
||||
cb.RecordFailure()
|
||||
cb.RecordFailure()
|
||||
|
||||
// Should still be closed (only 2 consecutive failures)
|
||||
if cb.State() != circuitClosed {
|
||||
t.Fatalf("expected closed after reset, got %s", cb.State())
|
||||
}
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_HalfOpenAfterTimeout(t *testing.T) {
|
||||
cb := newCircuitBreaker(CircuitBreakerConfig{
|
||||
MaxFailures: 1,
|
||||
Timeout: 50 * time.Millisecond,
|
||||
MaxHalfOpen: 1,
|
||||
})
|
||||
|
||||
cb.RecordFailure()
|
||||
if cb.State() != circuitOpen {
|
||||
t.Fatalf("expected open, got %s", cb.State())
|
||||
}
|
||||
|
||||
time.Sleep(60 * time.Millisecond)
|
||||
|
||||
// Should transition to half-open
|
||||
if cb.State() != circuitHalfOpen {
|
||||
t.Fatalf("expected half-open after timeout, got %s", cb.State())
|
||||
}
|
||||
|
||||
// One probe request should be allowed
|
||||
if err := cb.Allow(); err != nil {
|
||||
t.Fatalf("expected half-open to allow probe, got: %v", err)
|
||||
}
|
||||
|
||||
// Second should be rejected (maxHalfOpen=1, already used)
|
||||
if err := cb.Allow(); err == nil {
|
||||
t.Fatal("expected half-open to reject after max probes")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_HalfOpenSuccessCloses(t *testing.T) {
|
||||
cb := newCircuitBreaker(CircuitBreakerConfig{
|
||||
MaxFailures: 1,
|
||||
Timeout: 50 * time.Millisecond,
|
||||
})
|
||||
|
||||
cb.RecordFailure()
|
||||
time.Sleep(60 * time.Millisecond)
|
||||
|
||||
// Allow probe
|
||||
if err := cb.Allow(); err != nil {
|
||||
t.Fatalf("expected probe allowed: %v", err)
|
||||
}
|
||||
|
||||
// Probe succeeds -> circuit closes
|
||||
cb.RecordSuccess()
|
||||
if cb.State() != circuitClosed {
|
||||
t.Fatalf("expected closed after successful probe, got %s", cb.State())
|
||||
}
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_HalfOpenFailureReopens(t *testing.T) {
|
||||
cb := newCircuitBreaker(CircuitBreakerConfig{
|
||||
MaxFailures: 1,
|
||||
Timeout: 50 * time.Millisecond,
|
||||
})
|
||||
|
||||
cb.RecordFailure()
|
||||
time.Sleep(60 * time.Millisecond)
|
||||
|
||||
// Allow probe
|
||||
cb.Allow()
|
||||
|
||||
// Probe fails -> circuit re-opens
|
||||
cb.RecordFailure()
|
||||
if cb.State() != circuitOpen {
|
||||
t.Fatalf("expected open after failed probe, got %s", cb.State())
|
||||
}
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_Defaults(t *testing.T) {
|
||||
cb := newCircuitBreaker(CircuitBreakerConfig{})
|
||||
|
||||
if cb.maxFailures != 5 {
|
||||
t.Fatalf("expected default maxFailures=5, got %d", cb.maxFailures)
|
||||
}
|
||||
if cb.timeout != 30*time.Second {
|
||||
t.Fatalf("expected default timeout=30s, got %s", cb.timeout)
|
||||
}
|
||||
if cb.maxHalfOpen != 1 {
|
||||
t.Fatalf("expected default maxHalfOpen=1, got %d", cb.maxHalfOpen)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCircuitBreaker_StateString(t *testing.T) {
|
||||
tests := []struct {
|
||||
state circuitState
|
||||
want string
|
||||
}{
|
||||
{circuitClosed, "closed"},
|
||||
{circuitOpen, "open"},
|
||||
{circuitHalfOpen, "half-open"},
|
||||
{circuitState(99), "unknown"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
if got := tt.state.String(); got != tt.want {
|
||||
t.Errorf("state %d: got %q, want %q", tt.state, got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
apiVersion: v2
|
||||
name: mcp-gateway
|
||||
description: Go Micro MCP Gateway - Expose microservices as AI-accessible tools via the Model Context Protocol
|
||||
type: application
|
||||
version: 0.1.0
|
||||
appVersion: "0.1.0"
|
||||
keywords:
|
||||
- go-micro
|
||||
- mcp
|
||||
- ai
|
||||
- microservices
|
||||
- gateway
|
||||
home: https://go-micro.dev
|
||||
sources:
|
||||
- https://github.com/micro/go-micro
|
||||
maintainers:
|
||||
- name: go-micro
|
||||
url: https://github.com/micro/go-micro
|
||||
@@ -0,0 +1,89 @@
|
||||
# MCP Gateway Helm Chart
|
||||
|
||||
Deploy the Go Micro MCP Gateway on Kubernetes. The gateway discovers go-micro services via a registry and exposes them as AI-accessible tools through the Model Context Protocol.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
helm install mcp-gateway ./deploy/helm/mcp-gateway \
|
||||
--set gateway.registry=consul \
|
||||
--set gateway.registryAddress=consul:8500
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
| Parameter | Description | Default |
|
||||
|-----------|-------------|---------|
|
||||
| `replicaCount` | Number of gateway replicas | `1` |
|
||||
| `image.repository` | Container image | `ghcr.io/micro/mcp-gateway` |
|
||||
| `image.tag` | Image tag (defaults to appVersion) | `""` |
|
||||
| `gateway.address` | Listen address | `:3000` |
|
||||
| `gateway.registry` | Registry backend (mdns, consul, etcd) | `consul` |
|
||||
| `gateway.registryAddress` | Registry address | `consul:8500` |
|
||||
| `gateway.rateLimit` | Requests/second per tool (0=unlimited) | `0` |
|
||||
| `gateway.rateBurst` | Rate limit burst size | `20` |
|
||||
| `gateway.auth` | Enable JWT authentication | `false` |
|
||||
| `gateway.audit` | Enable audit logging | `false` |
|
||||
| `gateway.scopes` | Per-tool scope requirements | `[]` |
|
||||
| `service.type` | Kubernetes service type | `ClusterIP` |
|
||||
| `service.port` | Service port | `3000` |
|
||||
| `ingress.enabled` | Enable ingress | `false` |
|
||||
| `autoscaling.enabled` | Enable HPA | `false` |
|
||||
| `autoscaling.minReplicas` | Minimum replicas | `1` |
|
||||
| `autoscaling.maxReplicas` | Maximum replicas | `10` |
|
||||
|
||||
## Examples
|
||||
|
||||
### Production with Consul
|
||||
|
||||
```bash
|
||||
helm install mcp-gateway ./deploy/helm/mcp-gateway \
|
||||
--set replicaCount=3 \
|
||||
--set gateway.registry=consul \
|
||||
--set gateway.registryAddress=consul.default.svc:8500 \
|
||||
--set gateway.auth=true \
|
||||
--set gateway.audit=true \
|
||||
--set gateway.rateLimit=100 \
|
||||
--set autoscaling.enabled=true
|
||||
```
|
||||
|
||||
### With Ingress (nginx)
|
||||
|
||||
```bash
|
||||
helm install mcp-gateway ./deploy/helm/mcp-gateway \
|
||||
--set ingress.enabled=true \
|
||||
--set ingress.className=nginx \
|
||||
--set ingress.hosts[0].host=mcp.example.com \
|
||||
--set ingress.hosts[0].paths[0].path=/ \
|
||||
--set ingress.hosts[0].paths[0].pathType=Prefix \
|
||||
--set ingress.tls[0].secretName=mcp-tls \
|
||||
--set ingress.tls[0].hosts[0]=mcp.example.com
|
||||
```
|
||||
|
||||
### With Scopes
|
||||
|
||||
```bash
|
||||
helm install mcp-gateway ./deploy/helm/mcp-gateway \
|
||||
--set gateway.auth=true \
|
||||
--set 'gateway.scopes[0]=blog.Blog.Create=blog:write' \
|
||||
--set 'gateway.scopes[1]=blog.Blog.Delete=blog:admin'
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
Kubernetes Cluster
|
||||
┌──────────────────────────────────────────────────────────┐
|
||||
│ │
|
||||
│ ┌─────────┐ MCP ┌─────────────┐ RPC ┌──────┐ │
|
||||
│ │ Ingress │ ───────> │ MCP Gateway │ ──────> │ Svc │ │
|
||||
│ │ │ │ (N pods) │ │ Pods │ │
|
||||
│ └─────────┘ └─────────────┘ └──────┘ │
|
||||
│ │ │ │
|
||||
│ v v │
|
||||
│ ┌──────────┐ │
|
||||
│ │ Consul │ │
|
||||
│ │ Registry │ │
|
||||
│ └──────────┘ │
|
||||
└──────────────────────────────────────────────────────────┘
|
||||
```
|
||||
@@ -0,0 +1,31 @@
|
||||
MCP Gateway has been deployed.
|
||||
|
||||
1. Get the gateway URL:
|
||||
{{- if .Values.ingress.enabled }}
|
||||
{{- range $host := .Values.ingress.hosts }}
|
||||
http{{ if $.Values.ingress.tls }}s{{ end }}://{{ $host.host }}
|
||||
{{- end }}
|
||||
{{- else if contains "NodePort" .Values.service.type }}
|
||||
export NODE_PORT=$(kubectl get --namespace {{ .Release.Namespace }} -o jsonpath="{.spec.ports[0].nodePort}" services {{ include "mcp-gateway.fullname" . }})
|
||||
export NODE_IP=$(kubectl get nodes --namespace {{ .Release.Namespace }} -o jsonpath="{.items[0].status.addresses[0].address}")
|
||||
echo http://$NODE_IP:$NODE_PORT
|
||||
{{- else if contains "LoadBalancer" .Values.service.type }}
|
||||
export SERVICE_IP=$(kubectl get svc --namespace {{ .Release.Namespace }} {{ include "mcp-gateway.fullname" . }} --template "{{"{{ range (index .status.loadBalancer.ingress 0) }}{{.}}{{ end }}"}}")
|
||||
echo http://$SERVICE_IP:{{ .Values.service.port }}
|
||||
{{- else }}
|
||||
kubectl port-forward --namespace {{ .Release.Namespace }} svc/{{ include "mcp-gateway.fullname" . }} {{ .Values.service.port }}:{{ .Values.service.port }}
|
||||
echo http://127.0.0.1:{{ .Values.service.port }}
|
||||
{{- end }}
|
||||
|
||||
2. List available MCP tools:
|
||||
curl http://<GATEWAY_URL>/mcp/tools | jq
|
||||
|
||||
3. Connect Claude Code:
|
||||
Add to your MCP settings:
|
||||
{
|
||||
"mcpServers": {
|
||||
"my-services": {
|
||||
"url": "http://<GATEWAY_URL>/mcp"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
{{/*
|
||||
Expand the name of the chart.
|
||||
*/}}
|
||||
{{- define "mcp-gateway.name" -}}
|
||||
{{- default .Chart.Name .Values.nameOverride | trunc 63 | trimSuffix "-" }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Create a default fully qualified app name.
|
||||
*/}}
|
||||
{{- define "mcp-gateway.fullname" -}}
|
||||
{{- if .Values.fullnameOverride }}
|
||||
{{- .Values.fullnameOverride | trunc 63 | trimSuffix "-" }}
|
||||
{{- else }}
|
||||
{{- $name := default .Chart.Name .Values.nameOverride }}
|
||||
{{- if contains $name .Release.Name }}
|
||||
{{- .Release.Name | trunc 63 | trimSuffix "-" }}
|
||||
{{- else }}
|
||||
{{- printf "%s-%s" .Release.Name $name | trunc 63 | trimSuffix "-" }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Create chart name and version as used by the chart label.
|
||||
*/}}
|
||||
{{- define "mcp-gateway.chart" -}}
|
||||
{{- printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" | trunc 63 | trimSuffix "-" }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Common labels
|
||||
*/}}
|
||||
{{- define "mcp-gateway.labels" -}}
|
||||
helm.sh/chart: {{ include "mcp-gateway.chart" . }}
|
||||
{{ include "mcp-gateway.selectorLabels" . }}
|
||||
{{- if .Chart.AppVersion }}
|
||||
app.kubernetes.io/version: {{ .Chart.AppVersion | quote }}
|
||||
{{- end }}
|
||||
app.kubernetes.io/managed-by: {{ .Release.Service }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Selector labels
|
||||
*/}}
|
||||
{{- define "mcp-gateway.selectorLabels" -}}
|
||||
app.kubernetes.io/name: {{ include "mcp-gateway.name" . }}
|
||||
app.kubernetes.io/instance: {{ .Release.Name }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Create the name of the service account to use
|
||||
*/}}
|
||||
{{- define "mcp-gateway.serviceAccountName" -}}
|
||||
{{- if .Values.serviceAccount.create }}
|
||||
{{- default (include "mcp-gateway.fullname" .) .Values.serviceAccount.name }}
|
||||
{{- else }}
|
||||
{{- default "default" .Values.serviceAccount.name }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,95 @@
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: {{ include "mcp-gateway.fullname" . }}
|
||||
labels:
|
||||
{{- include "mcp-gateway.labels" . | nindent 4 }}
|
||||
spec:
|
||||
{{- if not .Values.autoscaling.enabled }}
|
||||
replicas: {{ .Values.replicaCount }}
|
||||
{{- end }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "mcp-gateway.selectorLabels" . | nindent 6 }}
|
||||
template:
|
||||
metadata:
|
||||
{{- with .Values.podAnnotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
labels:
|
||||
{{- include "mcp-gateway.labels" . | nindent 8 }}
|
||||
{{- with .Values.podLabels }}
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
{{- with .Values.imagePullSecrets }}
|
||||
imagePullSecrets:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
serviceAccountName: {{ include "mcp-gateway.serviceAccountName" . }}
|
||||
{{- with .Values.podSecurityContext }}
|
||||
securityContext:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
containers:
|
||||
- name: {{ .Chart.Name }}
|
||||
{{- with .Values.securityContext }}
|
||||
securityContext:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}"
|
||||
imagePullPolicy: {{ .Values.image.pullPolicy }}
|
||||
args:
|
||||
- "--address"
|
||||
- {{ .Values.gateway.address | quote }}
|
||||
- "--registry"
|
||||
- {{ .Values.gateway.registry | quote }}
|
||||
{{- if .Values.gateway.registryAddress }}
|
||||
- "--registry-address"
|
||||
- {{ .Values.gateway.registryAddress | quote }}
|
||||
{{- end }}
|
||||
{{- if gt (float64 .Values.gateway.rateLimit) 0.0 }}
|
||||
- "--rate-limit"
|
||||
- {{ .Values.gateway.rateLimit | quote }}
|
||||
- "--rate-burst"
|
||||
- {{ .Values.gateway.rateBurst | quote }}
|
||||
{{- end }}
|
||||
{{- if .Values.gateway.auth }}
|
||||
- "--auth"
|
||||
{{- end }}
|
||||
{{- if .Values.gateway.audit }}
|
||||
- "--audit"
|
||||
{{- end }}
|
||||
{{- range .Values.gateway.scopes }}
|
||||
- "--scope"
|
||||
- {{ . | quote }}
|
||||
{{- end }}
|
||||
ports:
|
||||
- name: http
|
||||
containerPort: {{ trimPrefix ":" .Values.gateway.address | default "3000" }}
|
||||
protocol: TCP
|
||||
{{- with .Values.probes.liveness }}
|
||||
livenessProbe:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- with .Values.probes.readiness }}
|
||||
readinessProbe:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- with .Values.resources }}
|
||||
resources:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- with .Values.nodeSelector }}
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.affinity }}
|
||||
affinity:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.tolerations }}
|
||||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,32 @@
|
||||
{{- if .Values.autoscaling.enabled }}
|
||||
apiVersion: autoscaling/v2
|
||||
kind: HorizontalPodAutoscaler
|
||||
metadata:
|
||||
name: {{ include "mcp-gateway.fullname" . }}
|
||||
labels:
|
||||
{{- include "mcp-gateway.labels" . | nindent 4 }}
|
||||
spec:
|
||||
scaleTargetRef:
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
name: {{ include "mcp-gateway.fullname" . }}
|
||||
minReplicas: {{ .Values.autoscaling.minReplicas }}
|
||||
maxReplicas: {{ .Values.autoscaling.maxReplicas }}
|
||||
metrics:
|
||||
{{- if .Values.autoscaling.targetCPUUtilizationPercentage }}
|
||||
- type: Resource
|
||||
resource:
|
||||
name: cpu
|
||||
target:
|
||||
type: Utilization
|
||||
averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }}
|
||||
{{- end }}
|
||||
{{- if .Values.autoscaling.targetMemoryUtilizationPercentage }}
|
||||
- type: Resource
|
||||
resource:
|
||||
name: memory
|
||||
target:
|
||||
type: Utilization
|
||||
averageUtilization: {{ .Values.autoscaling.targetMemoryUtilizationPercentage }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,41 @@
|
||||
{{- if .Values.ingress.enabled -}}
|
||||
apiVersion: networking.k8s.io/v1
|
||||
kind: Ingress
|
||||
metadata:
|
||||
name: {{ include "mcp-gateway.fullname" . }}
|
||||
labels:
|
||||
{{- include "mcp-gateway.labels" . | nindent 4 }}
|
||||
{{- with .Values.ingress.annotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
{{- if .Values.ingress.className }}
|
||||
ingressClassName: {{ .Values.ingress.className }}
|
||||
{{- end }}
|
||||
{{- if .Values.ingress.tls }}
|
||||
tls:
|
||||
{{- range .Values.ingress.tls }}
|
||||
- hosts:
|
||||
{{- range .hosts }}
|
||||
- {{ . | quote }}
|
||||
{{- end }}
|
||||
secretName: {{ .secretName }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
rules:
|
||||
{{- range .Values.ingress.hosts }}
|
||||
- host: {{ .host | quote }}
|
||||
http:
|
||||
paths:
|
||||
{{- range .paths }}
|
||||
- path: {{ .path }}
|
||||
pathType: {{ .pathType }}
|
||||
backend:
|
||||
service:
|
||||
name: {{ include "mcp-gateway.fullname" $ }}
|
||||
port:
|
||||
name: http
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,15 @@
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: {{ include "mcp-gateway.fullname" . }}
|
||||
labels:
|
||||
{{- include "mcp-gateway.labels" . | nindent 4 }}
|
||||
spec:
|
||||
type: {{ .Values.service.type }}
|
||||
ports:
|
||||
- port: {{ .Values.service.port }}
|
||||
targetPort: http
|
||||
protocol: TCP
|
||||
name: http
|
||||
selector:
|
||||
{{- include "mcp-gateway.selectorLabels" . | nindent 4 }}
|
||||
@@ -0,0 +1,13 @@
|
||||
{{- if .Values.serviceAccount.create -}}
|
||||
apiVersion: v1
|
||||
kind: ServiceAccount
|
||||
metadata:
|
||||
name: {{ include "mcp-gateway.serviceAccountName" . }}
|
||||
labels:
|
||||
{{- include "mcp-gateway.labels" . | nindent 4 }}
|
||||
{{- with .Values.serviceAccount.annotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
automountServiceAccountToken: {{ .Values.serviceAccount.automount }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,112 @@
|
||||
# MCP Gateway Helm chart values
|
||||
|
||||
replicaCount: 1
|
||||
|
||||
image:
|
||||
repository: ghcr.io/micro/mcp-gateway
|
||||
pullPolicy: IfNotPresent
|
||||
tag: "" # Defaults to appVersion
|
||||
|
||||
imagePullSecrets: []
|
||||
nameOverride: ""
|
||||
fullnameOverride: ""
|
||||
|
||||
# MCP Gateway configuration
|
||||
gateway:
|
||||
# Listen address (port inside the container)
|
||||
address: ":3000"
|
||||
|
||||
# Service registry backend: mdns, consul, etcd
|
||||
registry: consul
|
||||
|
||||
# Registry address (e.g., consul:8500, etcd:2379)
|
||||
registryAddress: "consul:8500"
|
||||
|
||||
# Rate limiting (0 = unlimited)
|
||||
rateLimit: 0
|
||||
rateBurst: 20
|
||||
|
||||
# Enable JWT authentication
|
||||
auth: false
|
||||
|
||||
# Enable audit logging to stdout
|
||||
audit: false
|
||||
|
||||
# Per-tool scope requirements (format: tool=scope1,scope2)
|
||||
scopes: []
|
||||
# - "blog.Blog.Create=blog:write"
|
||||
# - "blog.Blog.Delete=blog:admin"
|
||||
|
||||
serviceAccount:
|
||||
create: true
|
||||
automount: true
|
||||
annotations: {}
|
||||
name: ""
|
||||
|
||||
podAnnotations: {}
|
||||
podLabels: {}
|
||||
|
||||
podSecurityContext: {}
|
||||
|
||||
securityContext:
|
||||
allowPrivilegeEscalation: false
|
||||
capabilities:
|
||||
drop:
|
||||
- ALL
|
||||
readOnlyRootFilesystem: true
|
||||
runAsNonRoot: true
|
||||
runAsUser: 65534
|
||||
|
||||
service:
|
||||
type: ClusterIP
|
||||
port: 3000
|
||||
|
||||
ingress:
|
||||
enabled: false
|
||||
className: ""
|
||||
annotations: {}
|
||||
# kubernetes.io/ingress.class: nginx
|
||||
# cert-manager.io/cluster-issuer: letsencrypt-prod
|
||||
hosts:
|
||||
- host: mcp-gateway.local
|
||||
paths:
|
||||
- path: /
|
||||
pathType: Prefix
|
||||
tls: []
|
||||
# - secretName: mcp-gateway-tls
|
||||
# hosts:
|
||||
# - mcp-gateway.local
|
||||
|
||||
resources:
|
||||
requests:
|
||||
cpu: 100m
|
||||
memory: 64Mi
|
||||
limits:
|
||||
cpu: 500m
|
||||
memory: 128Mi
|
||||
|
||||
autoscaling:
|
||||
enabled: false
|
||||
minReplicas: 1
|
||||
maxReplicas: 10
|
||||
targetCPUUtilizationPercentage: 70
|
||||
targetMemoryUtilizationPercentage: 80
|
||||
|
||||
nodeSelector: {}
|
||||
tolerations: []
|
||||
affinity: {}
|
||||
|
||||
# Liveness and readiness probes
|
||||
probes:
|
||||
liveness:
|
||||
httpGet:
|
||||
path: /healthz
|
||||
port: http
|
||||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
readiness:
|
||||
httpGet:
|
||||
path: /healthz
|
||||
port: http
|
||||
initialDelaySeconds: 3
|
||||
periodSeconds: 5
|
||||
@@ -0,0 +1,135 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
|
||||
"go-micro.dev/v6"
|
||||
"go-micro.dev/v6/auth/jwt"
|
||||
"go-micro.dev/v6/registry"
|
||||
)
|
||||
|
||||
// Example_withMCP shows the simplest way to add MCP to a service using WithMCP
|
||||
func Example_withMCP() {
|
||||
// One line to make your service AI-accessible
|
||||
service := micro.NewService("myservice", WithMCP(":3000"))
|
||||
service.Init()
|
||||
service.Run()
|
||||
}
|
||||
|
||||
// Example_inlineGateway shows how to add MCP gateway to an existing service
|
||||
// with full control over options
|
||||
func Example_inlineGateway() {
|
||||
service := micro.NewService("myservice")
|
||||
service.Init()
|
||||
|
||||
// Add MCP gateway alongside your service
|
||||
go func() {
|
||||
if err := Serve(Options{
|
||||
Registry: service.Options().Registry,
|
||||
Address: ":3000",
|
||||
}); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
}()
|
||||
|
||||
// Run your service normally
|
||||
service.Run()
|
||||
}
|
||||
|
||||
// Example_standaloneGateway shows how to run MCP gateway as a separate service
|
||||
func Example_standaloneGateway() {
|
||||
// Standalone MCP gateway
|
||||
// Discovers all services via registry
|
||||
if err := ListenAndServe(":3000", Options{
|
||||
Registry: registry.NewMDNSRegistry(),
|
||||
}); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Example_withAuthentication shows how to add authentication
|
||||
func Example_withAuthentication() {
|
||||
service := micro.NewService("myservice")
|
||||
service.Init()
|
||||
|
||||
go func() {
|
||||
if err := Serve(Options{
|
||||
Registry: service.Options().Registry,
|
||||
Address: ":3000",
|
||||
AuthFunc: func(r *http.Request) error {
|
||||
token := r.Header.Get("Authorization")
|
||||
if token == "" {
|
||||
return fmt.Errorf("missing authorization header")
|
||||
}
|
||||
// Validate token here
|
||||
return nil
|
||||
},
|
||||
}); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
}()
|
||||
|
||||
service.Run()
|
||||
}
|
||||
|
||||
// Example_customContext shows how to use a custom context for graceful shutdown
|
||||
func Example_customContext() {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
service := micro.NewService("myservice")
|
||||
service.Init()
|
||||
|
||||
go func() {
|
||||
if err := Serve(Options{
|
||||
Registry: service.Options().Registry,
|
||||
Address: ":3000",
|
||||
Context: ctx,
|
||||
}); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
}()
|
||||
|
||||
service.Run()
|
||||
// cancel() will stop the MCP gateway
|
||||
}
|
||||
|
||||
// Example_withScopesAndTracing shows how to add per-tool scopes, tracing, rate
|
||||
// limiting and audit logging to the MCP gateway. Services register scope
|
||||
// requirements via endpoint metadata ("scopes" key, comma-separated).
|
||||
func Example_withScopesAndTracing() {
|
||||
service := micro.NewService("blog")
|
||||
service.Init()
|
||||
|
||||
// Use JWT auth provider
|
||||
authProvider := jwt.NewAuth()
|
||||
|
||||
go func() {
|
||||
if err := Serve(Options{
|
||||
Registry: service.Options().Registry,
|
||||
Address: ":3000",
|
||||
|
||||
// Auth inspects Bearer tokens and enforces per-tool scopes
|
||||
Auth: authProvider,
|
||||
|
||||
// Rate limit all tools to 10 req/s with burst of 20
|
||||
RateLimit: &RateLimitConfig{
|
||||
RequestsPerSecond: 10,
|
||||
Burst: 20,
|
||||
},
|
||||
|
||||
// Audit every tool call for compliance
|
||||
AuditFunc: func(r AuditRecord) {
|
||||
log.Printf("[audit] trace=%s tool=%s account=%s allowed=%v reason=%s",
|
||||
r.TraceID, r.Tool, r.AccountID, r.Allowed, r.DeniedReason)
|
||||
},
|
||||
}); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
}()
|
||||
|
||||
service.Run()
|
||||
}
|
||||
@@ -0,0 +1,282 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/credentials/insecure"
|
||||
reflectionpb "google.golang.org/grpc/reflection/grpc_reflection_v1alpha"
|
||||
"google.golang.org/protobuf/encoding/protojson"
|
||||
"google.golang.org/protobuf/proto"
|
||||
"google.golang.org/protobuf/reflect/protodesc"
|
||||
"google.golang.org/protobuf/reflect/protoreflect"
|
||||
"google.golang.org/protobuf/reflect/protoregistry"
|
||||
"google.golang.org/protobuf/types/descriptorpb"
|
||||
"google.golang.org/protobuf/types/dynamicpb"
|
||||
)
|
||||
|
||||
// ReflectedGRPCTarget describes an external gRPC server whose reflection
|
||||
// catalog should be exposed as MCP tools. It is intentionally opt-in: teams can
|
||||
// bridge existing reflected gRPC services without changing their servers or
|
||||
// registering them in go-micro.
|
||||
type ReflectedGRPCTarget struct {
|
||||
// Name prefixes generated tools. When empty, Address is sanitized and used.
|
||||
Name string
|
||||
// Address is the host:port of the reflected gRPC server.
|
||||
Address string
|
||||
// DialOptions customize the connection. If none are supplied, an insecure
|
||||
// transport is used for local/dev interoperability.
|
||||
DialOptions []grpc.DialOption
|
||||
// Timeout bounds reflection discovery and individual tool calls.
|
||||
Timeout time.Duration
|
||||
}
|
||||
|
||||
func (s *Server) discoverReflectedGRPC() error {
|
||||
for _, target := range s.opts.ReflectedGRPCTargets {
|
||||
if strings.TrimSpace(target.Address) == "" {
|
||||
continue
|
||||
}
|
||||
tools, err := s.reflectedGRPCTools(target)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, tool := range tools {
|
||||
s.tools[tool.Name] = tool
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Server) reflectedGRPCTools(target ReflectedGRPCTarget) ([]*Tool, error) {
|
||||
timeout := target.Timeout
|
||||
if timeout == 0 {
|
||||
timeout = 10 * time.Second
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(s.opts.Context, timeout)
|
||||
defer cancel()
|
||||
|
||||
dialOpts := target.DialOptions
|
||||
if len(dialOpts) == 0 {
|
||||
dialOpts = []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())}
|
||||
}
|
||||
conn, err := grpc.NewClient(target.Address, dialOpts...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("connect reflected grpc target %s: %w", target.Address, err)
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
files, services, err := loadReflectedFiles(ctx, conn)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("reflect grpc target %s: %w", target.Address, err)
|
||||
}
|
||||
|
||||
prefix := target.Name
|
||||
if prefix == "" {
|
||||
prefix = sanitizeToolPart(target.Address)
|
||||
}
|
||||
|
||||
var out []*Tool
|
||||
for _, serviceName := range services {
|
||||
desc, err := files.FindDescriptorByName(protoreflect.FullName(serviceName))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
svc, ok := desc.(protoreflect.ServiceDescriptor)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
for i := 0; i < svc.Methods().Len(); i++ {
|
||||
method := svc.Methods().Get(i)
|
||||
if method.IsStreamingClient() || method.IsStreamingServer() {
|
||||
continue
|
||||
}
|
||||
fullMethod := "/" + string(svc.FullName()) + "/" + string(method.Name())
|
||||
toolName := prefix + "." + strings.ReplaceAll(string(svc.FullName()), ".", "_") + "." + string(method.Name())
|
||||
input := method.Input()
|
||||
out = append(out, &Tool{
|
||||
Name: toolName,
|
||||
Description: fmt.Sprintf("Call reflected gRPC method %s on %s", fullMethod, target.Address),
|
||||
InputSchema: protoMessageSchema(input),
|
||||
Handler: reflectedGRPCHandler(target, fullMethod, input, method.Output()),
|
||||
})
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func loadReflectedFiles(ctx context.Context, conn *grpc.ClientConn) (*protoregistryFiles, []string, error) {
|
||||
client := reflectionpb.NewServerReflectionClient(conn)
|
||||
stream, err := client.ServerReflectionInfo(ctx)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if err := stream.Send(&reflectionpb.ServerReflectionRequest{MessageRequest: &reflectionpb.ServerReflectionRequest_ListServices{ListServices: ""}}); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
resp, err := stream.Recv()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
list := resp.GetListServicesResponse()
|
||||
if list == nil {
|
||||
return nil, nil, fmt.Errorf("reflection list services returned %T", resp.MessageResponse)
|
||||
}
|
||||
|
||||
set := &descriptorpb.FileDescriptorSet{}
|
||||
seen := map[string]bool{}
|
||||
var services []string
|
||||
for _, svc := range list.Service {
|
||||
name := svc.Name
|
||||
if strings.HasPrefix(name, "grpc.reflection.") {
|
||||
continue
|
||||
}
|
||||
services = append(services, name)
|
||||
if err := requestFileContainingSymbol(ctx, client, name, set, seen); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
}
|
||||
files, err := newProtoregistryFiles(set)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return files, services, nil
|
||||
}
|
||||
|
||||
func requestFileContainingSymbol(ctx context.Context, client reflectionpb.ServerReflectionClient, symbol string, set *descriptorpb.FileDescriptorSet, seen map[string]bool) error {
|
||||
stream, err := client.ServerReflectionInfo(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := stream.Send(&reflectionpb.ServerReflectionRequest{MessageRequest: &reflectionpb.ServerReflectionRequest_FileContainingSymbol{FileContainingSymbol: symbol}}); err != nil {
|
||||
return err
|
||||
}
|
||||
resp, err := stream.Recv()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fd := resp.GetFileDescriptorResponse()
|
||||
if fd == nil {
|
||||
return fmt.Errorf("reflection lookup for %s returned %T", symbol, resp.MessageResponse)
|
||||
}
|
||||
for _, raw := range fd.FileDescriptorProto {
|
||||
var file descriptorpb.FileDescriptorProto
|
||||
if err := proto.Unmarshal(raw, &file); err != nil {
|
||||
return err
|
||||
}
|
||||
name := file.GetName()
|
||||
if !seen[name] {
|
||||
seen[name] = true
|
||||
set.File = append(set.File, &file)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// protoregistryFiles is a narrow wrapper that keeps imports local to this file.
|
||||
type protoregistryFiles struct{ files *protoregistry.Files }
|
||||
|
||||
func newProtoregistryFiles(set *descriptorpb.FileDescriptorSet) (*protoregistryFiles, error) {
|
||||
files, err := protodesc.NewFiles(set)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &protoregistryFiles{files: files}, nil
|
||||
}
|
||||
|
||||
func (p *protoregistryFiles) FindDescriptorByName(name protoreflect.FullName) (protoreflect.Descriptor, error) {
|
||||
return p.files.FindDescriptorByName(name)
|
||||
}
|
||||
|
||||
func reflectedGRPCHandler(target ReflectedGRPCTarget, fullMethod string, input, output protoreflect.MessageDescriptor) func(map[string]interface{}) (interface{}, error) {
|
||||
return func(args map[string]interface{}) (interface{}, error) {
|
||||
timeout := target.Timeout
|
||||
if timeout == 0 {
|
||||
timeout = 10 * time.Second
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
||||
defer cancel()
|
||||
dialOpts := target.DialOptions
|
||||
if len(dialOpts) == 0 {
|
||||
dialOpts = []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())}
|
||||
}
|
||||
conn, err := grpc.NewClient(target.Address, dialOpts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
req := dynamicpb.NewMessage(input)
|
||||
raw, err := json.Marshal(args)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := protojson.Unmarshal(raw, req); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rsp := dynamicpb.NewMessage(output)
|
||||
if err := conn.Invoke(ctx, fullMethod, req, rsp); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
b, err := protojson.MarshalOptions{UseProtoNames: true, EmitUnpopulated: true}.Marshal(rsp)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var out interface{}
|
||||
if err := json.Unmarshal(b, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
}
|
||||
|
||||
func protoMessageSchema(msg protoreflect.MessageDescriptor) map[string]interface{} {
|
||||
schema := map[string]interface{}{"type": "object", "properties": map[string]interface{}{}}
|
||||
props := schema["properties"].(map[string]interface{})
|
||||
fields := msg.Fields()
|
||||
for i := 0; i < fields.Len(); i++ {
|
||||
field := fields.Get(i)
|
||||
props[field.JSONName()] = protoFieldSchema(field)
|
||||
}
|
||||
return schema
|
||||
}
|
||||
|
||||
func protoFieldSchema(field protoreflect.FieldDescriptor) map[string]interface{} {
|
||||
schema := map[string]interface{}{"type": protoJSONType(field)}
|
||||
if field.IsList() {
|
||||
schema["items"] = map[string]interface{}{"type": protoJSONType(field)}
|
||||
}
|
||||
if field.Kind() == protoreflect.MessageKind || field.Kind() == protoreflect.GroupKind {
|
||||
schema = protoMessageSchema(field.Message())
|
||||
}
|
||||
return schema
|
||||
}
|
||||
|
||||
func protoJSONType(field protoreflect.FieldDescriptor) string {
|
||||
if field.IsList() {
|
||||
return "array"
|
||||
}
|
||||
switch field.Kind() {
|
||||
case protoreflect.BoolKind:
|
||||
return "boolean"
|
||||
case protoreflect.Int32Kind, protoreflect.Sint32Kind, protoreflect.Sfixed32Kind,
|
||||
protoreflect.Uint32Kind, protoreflect.Fixed32Kind, protoreflect.Int64Kind,
|
||||
protoreflect.Sint64Kind, protoreflect.Sfixed64Kind, protoreflect.Uint64Kind,
|
||||
protoreflect.Fixed64Kind:
|
||||
return "integer"
|
||||
case protoreflect.FloatKind, protoreflect.DoubleKind:
|
||||
return "number"
|
||||
case protoreflect.MessageKind, protoreflect.GroupKind:
|
||||
return "object"
|
||||
default:
|
||||
return "string"
|
||||
}
|
||||
}
|
||||
|
||||
func sanitizeToolPart(s string) string {
|
||||
r := strings.NewReplacer(":", "_", "/", "_", ".", "_", "-", "_")
|
||||
return r.Replace(s)
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"google.golang.org/grpc"
|
||||
helloworld "google.golang.org/grpc/examples/helloworld/helloworld"
|
||||
"google.golang.org/grpc/reflection"
|
||||
)
|
||||
|
||||
type reflectedGreeter struct {
|
||||
helloworld.UnimplementedGreeterServer
|
||||
}
|
||||
|
||||
func (reflectedGreeter) SayHello(_ context.Context, req *helloworld.HelloRequest) (*helloworld.HelloReply, error) {
|
||||
return &helloworld.HelloReply{Message: "hello " + req.Name}, nil
|
||||
}
|
||||
|
||||
func TestReflectedGRPCTargetDiscoversAndCallsUnaryTool(t *testing.T) {
|
||||
lis, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
grpcServer := grpc.NewServer()
|
||||
helloworld.RegisterGreeterServer(grpcServer, reflectedGreeter{})
|
||||
reflection.Register(grpcServer)
|
||||
go grpcServer.Serve(lis)
|
||||
defer grpcServer.Stop()
|
||||
|
||||
s := newTestServer(Options{Context: context.Background()})
|
||||
tools, err := s.reflectedGRPCTools(ReflectedGRPCTarget{
|
||||
Name: "demo",
|
||||
Address: lis.Addr().String(),
|
||||
Timeout: 3 * time.Second,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("discover reflected tools: %v", err)
|
||||
}
|
||||
if len(tools) != 1 {
|
||||
t.Fatalf("tools len = %d, want 1", len(tools))
|
||||
}
|
||||
tool := tools[0]
|
||||
if tool.Name != "demo.helloworld_Greeter.SayHello" {
|
||||
t.Fatalf("tool name = %q", tool.Name)
|
||||
}
|
||||
props := tool.InputSchema["properties"].(map[string]interface{})
|
||||
if _, ok := props["name"]; !ok {
|
||||
t.Fatalf("input schema missing name: %#v", tool.InputSchema)
|
||||
}
|
||||
|
||||
out, err := tool.Handler(map[string]interface{}{"name": "Ada"})
|
||||
if err != nil {
|
||||
t.Fatalf("call reflected tool: %v", err)
|
||||
}
|
||||
got := out.(map[string]interface{})["message"]
|
||||
if got != "hello Ada" {
|
||||
t.Fatalf("message = %v, want hello Ada", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// HandlerOption configures NewHandler.
|
||||
type HandlerOption func(*handlerOptions)
|
||||
|
||||
type handlerOptions struct {
|
||||
serverName, serverVersion, protocolVersion string
|
||||
}
|
||||
|
||||
// WithServerInfo sets the name/version advertised in the initialize response.
|
||||
func WithServerInfo(name, version string) HandlerOption {
|
||||
return func(o *handlerOptions) { o.serverName, o.serverVersion = name, version }
|
||||
}
|
||||
|
||||
// WithProtocolVersion sets the MCP protocol version advertised in initialize.
|
||||
func WithProtocolVersion(v string) HandlerOption {
|
||||
return func(o *handlerOptions) { o.protocolVersion = v }
|
||||
}
|
||||
|
||||
// NewHandler returns an http.Handler serving the MCP protocol over HTTP as
|
||||
// JSON-RPC 2.0 (initialize, ping, notifications/*, tools/list, tools/call),
|
||||
// backed by the resolver. Mount it on your own server (e.g. POST /mcp): the
|
||||
// gateway provides the protocol; you keep your routes, middleware and any
|
||||
// human-facing docs page.
|
||||
func NewHandler(r Resolver, opts ...HandlerOption) http.Handler {
|
||||
o := handlerOptions{serverName: "go-micro-mcp", serverVersion: "1.0.0", protocolVersion: "2024-11-05"}
|
||||
for _, fn := range opts {
|
||||
fn(&o)
|
||||
}
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
|
||||
if req.Method != http.MethodPost {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
var rpc struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID json.RawMessage `json:"id"`
|
||||
Method string `json:"method"`
|
||||
Params json.RawMessage `json:"params"`
|
||||
}
|
||||
if err := json.NewDecoder(req.Body).Decode(&rpc); err != nil {
|
||||
writeRPCError(w, nil, ParseError, "Parse error", err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// Notifications (and any id-less request) expect no response body.
|
||||
if strings.HasPrefix(rpc.Method, "notifications/") || len(rpc.ID) == 0 {
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
return
|
||||
}
|
||||
|
||||
ctx := req.Context()
|
||||
switch rpc.Method {
|
||||
case "initialize":
|
||||
writeRPCResult(w, rpc.ID, map[string]interface{}{
|
||||
"protocolVersion": o.protocolVersion,
|
||||
"capabilities": map[string]interface{}{"tools": map[string]interface{}{}},
|
||||
"serverInfo": map[string]interface{}{"name": o.serverName, "version": o.serverVersion},
|
||||
})
|
||||
case "ping":
|
||||
writeRPCResult(w, rpc.ID, map[string]interface{}{})
|
||||
case "tools/list":
|
||||
tools, err := r.List(ctx)
|
||||
if err != nil {
|
||||
writeRPCError(w, rpc.ID, InternalError, "Failed to list tools", err.Error())
|
||||
return
|
||||
}
|
||||
list := make([]map[string]interface{}, 0, len(tools))
|
||||
for _, t := range tools {
|
||||
list = append(list, map[string]interface{}{
|
||||
"name": t.Name, "description": t.Description, "inputSchema": t.InputSchema,
|
||||
})
|
||||
}
|
||||
writeRPCResult(w, rpc.ID, map[string]interface{}{"tools": list})
|
||||
case "tools/call":
|
||||
var p struct {
|
||||
Name string `json:"name"`
|
||||
Arguments map[string]interface{} `json:"arguments"`
|
||||
}
|
||||
if err := json.Unmarshal(rpc.Params, &p); err != nil {
|
||||
writeRPCError(w, rpc.ID, InvalidParams, "Invalid params", err.Error())
|
||||
return
|
||||
}
|
||||
res, err := r.Call(ctx, p.Name, p.Arguments)
|
||||
if err != nil {
|
||||
// Protocol/pre-check failure -> JSON-RPC error. An *RPCError
|
||||
// carries a specific code; anything else is InternalError.
|
||||
if rpcErr, ok := err.(*RPCError); ok {
|
||||
writeRPCError(w, rpc.ID, rpcErr.Code, rpcErr.Message, rpcErr.Data)
|
||||
} else {
|
||||
writeRPCError(w, rpc.ID, InternalError, "Tool call failed", err.Error())
|
||||
}
|
||||
return
|
||||
}
|
||||
result := map[string]interface{}{
|
||||
"content": []map[string]interface{}{{"type": "text", "text": res.Text}},
|
||||
}
|
||||
if res.IsError {
|
||||
result["isError"] = true
|
||||
}
|
||||
writeRPCResult(w, rpc.ID, result)
|
||||
default:
|
||||
writeRPCError(w, rpc.ID, MethodNotFound, "Method not found", rpc.Method)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func writeRPCResult(w http.ResponseWriter, id json.RawMessage, result interface{}) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]interface{}{"jsonrpc": "2.0", "id": rawOrNull(id), "result": result})
|
||||
}
|
||||
|
||||
func writeRPCError(w http.ResponseWriter, id json.RawMessage, code int, msg string, data interface{}) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]interface{}{"jsonrpc": "2.0", "id": rawOrNull(id), "error": map[string]interface{}{"code": code, "message": msg, "data": data}})
|
||||
}
|
||||
|
||||
func rawOrNull(id json.RawMessage) interface{} {
|
||||
if len(id) == 0 {
|
||||
return nil
|
||||
}
|
||||
return id
|
||||
}
|
||||
@@ -0,0 +1,995 @@
|
||||
// Package mcp provides Model Context Protocol (MCP) gateway functionality for go-micro services.
|
||||
// It automatically exposes your microservices as AI-accessible tools through MCP.
|
||||
//
|
||||
// Example usage:
|
||||
//
|
||||
// service := micro.NewService("myservice", )
|
||||
// service.Init()
|
||||
//
|
||||
// // Add MCP gateway
|
||||
// go mcp.Serve(mcp.Options{
|
||||
// Registry: service.Options().Registry,
|
||||
// Address: ":3000",
|
||||
// })
|
||||
//
|
||||
// service.Run()
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"go-micro.dev/v6/auth"
|
||||
"go-micro.dev/v6/broker"
|
||||
"go-micro.dev/v6/client"
|
||||
"go-micro.dev/v6/codec/bytes"
|
||||
"go-micro.dev/v6/metadata"
|
||||
"go-micro.dev/v6/registry"
|
||||
"go-micro.dev/v6/store"
|
||||
"go-micro.dev/v6/wrapper/x402"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
// Metadata keys for MCP tracing and auth propagated via context/metadata.
|
||||
const (
|
||||
// TraceIDKey is the metadata key for the MCP trace ID.
|
||||
TraceIDKey = "Mcp-Trace-Id"
|
||||
// ToolNameKey is the metadata key for the tool being invoked.
|
||||
ToolNameKey = "Mcp-Tool-Name"
|
||||
// AccountIDKey is the metadata key for the authenticated account ID.
|
||||
AccountIDKey = "Mcp-Account-Id"
|
||||
)
|
||||
|
||||
// AuditRecord represents an immutable log entry for an MCP tool call.
|
||||
type AuditRecord struct {
|
||||
// TraceID uniquely identifies this tool call chain.
|
||||
TraceID string `json:"trace_id"`
|
||||
// Timestamp of the tool call.
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
// Tool is the name of the tool that was called.
|
||||
Tool string `json:"tool"`
|
||||
// AccountID is the ID of the authenticated account (empty if unauthenticated).
|
||||
AccountID string `json:"account_id,omitempty"`
|
||||
// Scopes that were required for this tool.
|
||||
ScopesRequired []string `json:"scopes_required,omitempty"`
|
||||
// Allowed indicates whether the call was authorized.
|
||||
Allowed bool `json:"allowed"`
|
||||
// Denied reason, if the call was not allowed.
|
||||
DeniedReason string `json:"denied_reason,omitempty"`
|
||||
// Duration of the RPC call (zero if call was denied before execution).
|
||||
Duration time.Duration `json:"duration,omitempty"`
|
||||
// Error from the RPC call, if any.
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// AuditFunc is called for every tool call with an audit record.
|
||||
// Implementations should treat the record as immutable and persist it
|
||||
// (e.g. to a log, database, or event stream).
|
||||
type AuditFunc func(record AuditRecord)
|
||||
|
||||
// RateLimitConfig configures rate limiting for the MCP gateway.
|
||||
type RateLimitConfig struct {
|
||||
// Requests per second allowed per tool (0 = unlimited).
|
||||
RequestsPerSecond float64
|
||||
// Burst size (maximum number of requests that can be made at once).
|
||||
Burst int
|
||||
}
|
||||
|
||||
// Options configures the MCP gateway
|
||||
type Options struct {
|
||||
// Registry for service discovery (required)
|
||||
Registry registry.Registry
|
||||
|
||||
// Address to listen on for SSE transport (e.g., ":3000")
|
||||
// Leave empty for stdio transport
|
||||
Address string
|
||||
|
||||
// Client for making RPC calls (defaults to client.DefaultClient)
|
||||
Client client.Client
|
||||
|
||||
// Context for cancellation (defaults to background context)
|
||||
Context context.Context
|
||||
|
||||
// Logger for debug output (defaults to log.Default())
|
||||
Logger *log.Logger
|
||||
|
||||
// AuthFunc validates requests (optional, legacy)
|
||||
// Return error to reject, nil to allow
|
||||
AuthFunc func(r *http.Request) error
|
||||
|
||||
// Auth provider for token inspection (optional).
|
||||
// When set, incoming requests must carry a Bearer token which is
|
||||
// inspected to obtain an account. The account's scopes are then
|
||||
// checked against the tool's required scopes.
|
||||
Auth auth.Auth
|
||||
|
||||
// AuditFunc is called for every tool call with an immutable audit record.
|
||||
// Use this to persist tool-call logs for compliance and debugging.
|
||||
AuditFunc AuditFunc
|
||||
|
||||
// RateLimit configures per-tool rate limiting.
|
||||
// When set, each tool is limited to the configured requests per second.
|
||||
RateLimit *RateLimitConfig
|
||||
|
||||
// CircuitBreaker configures per-tool circuit breaking.
|
||||
// When set, tools that fail repeatedly are temporarily blocked to
|
||||
// protect downstream services from cascading failures.
|
||||
CircuitBreaker *CircuitBreakerConfig
|
||||
|
||||
// Scopes lets the gateway operator define or override per-tool
|
||||
// scope requirements without changing the services themselves.
|
||||
// Keys are tool names (e.g. "blog.Blog.Create") and values are the
|
||||
// required scopes. When a tool appears in Scopes its scopes
|
||||
// replace any scopes declared by the service via endpoint metadata.
|
||||
//
|
||||
// Example:
|
||||
//
|
||||
// Scopes: map[string][]string{
|
||||
// "blog.Blog.Create": {"blog:write"},
|
||||
// "blog.Blog.Delete": {"blog:admin"},
|
||||
// }
|
||||
Scopes map[string][]string
|
||||
|
||||
// TraceProvider enables OpenTelemetry tracing for MCP tool calls.
|
||||
// When set, each tool call creates a span with attributes for the
|
||||
// tool name, account ID, auth outcome, and transport type.
|
||||
// Trace context is propagated to downstream RPC calls via metadata.
|
||||
//
|
||||
// Example:
|
||||
//
|
||||
// tp := sdktrace.NewTracerProvider(sdktrace.WithBatcher(exporter))
|
||||
// mcp.Serve(mcp.Options{
|
||||
// Registry: reg,
|
||||
// TraceProvider: tp,
|
||||
// })
|
||||
TraceProvider trace.TracerProvider
|
||||
|
||||
// Payment, when set, requires an x402 payment for tool calls
|
||||
// (the /mcp/call endpoint). Listing tools and health stay free.
|
||||
// Opt-in: leave nil to disable payments.
|
||||
Payment *x402.Config
|
||||
|
||||
// ReflectedGRPCTargets exposes unary methods from external gRPC servers
|
||||
// that support server reflection as MCP tools. This bridges existing gRPC
|
||||
// services into the agent tool catalog without requiring go-micro handlers.
|
||||
ReflectedGRPCTargets []ReflectedGRPCTarget
|
||||
}
|
||||
|
||||
// Server represents a running MCP gateway
|
||||
type Server struct {
|
||||
opts Options
|
||||
tools map[string]*Tool
|
||||
toolsMu sync.RWMutex
|
||||
server *http.Server
|
||||
watching bool
|
||||
|
||||
// limiters holds per-tool rate limiters (nil if rate limiting is disabled).
|
||||
limiters map[string]*rateLimiter
|
||||
limitersMu sync.RWMutex
|
||||
|
||||
// breakers holds per-tool circuit breakers (nil if circuit breaking is disabled).
|
||||
breakers map[string]*circuitBreaker
|
||||
breakersMu sync.RWMutex
|
||||
}
|
||||
|
||||
// Tool represents an MCP tool (exposed service endpoint)
|
||||
type Tool struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
InputSchema map[string]interface{} `json:"inputSchema"`
|
||||
// Scopes lists the auth scopes required to call this tool.
|
||||
// An empty list means no scope restriction (subject to Auth provider).
|
||||
Scopes []string `json:"scopes,omitempty"`
|
||||
// Payment advertises the x402 payment required to call this tool, so
|
||||
// the catalog is shoppable — an agent sees the price before calling.
|
||||
// Populated at list time when the gateway has payments enabled; nil
|
||||
// means free.
|
||||
Payment *PaymentInfo `json:"payment,omitempty"`
|
||||
Service string `json:"-"`
|
||||
Endpoint string `json:"-"`
|
||||
// Handler is an optional direct handler for framework tools that don't
|
||||
// go through RPC. When set, handleCallTool calls this instead of making
|
||||
// an RPC request.
|
||||
Handler func(input map[string]interface{}) (interface{}, error) `json:"-"`
|
||||
}
|
||||
|
||||
// PaymentInfo advertises, in the tool catalog, the x402 payment required
|
||||
// to call a tool: how much, in what asset, on which network, and where it
|
||||
// goes. It lets an agent shop the catalog and choose by price before
|
||||
// calling.
|
||||
type PaymentInfo struct {
|
||||
Amount string `json:"amount"` // smallest units (e.g. "10000" = 0.01 USDC)
|
||||
Network string `json:"network"`
|
||||
Asset string `json:"asset,omitempty"`
|
||||
PayTo string `json:"payTo"`
|
||||
}
|
||||
|
||||
// paymentFor returns the catalog payment info for a tool, or nil if the
|
||||
// gateway has no payments configured or the tool is free.
|
||||
func (s *Server) paymentFor(toolName string) *PaymentInfo {
|
||||
if s.opts.Payment == nil {
|
||||
return nil
|
||||
}
|
||||
amount := s.opts.Payment.AmountFor(toolName)
|
||||
if amount == "" || amount == "0" {
|
||||
return nil
|
||||
}
|
||||
net := s.opts.Payment.Network
|
||||
if net == "" {
|
||||
net = "base"
|
||||
}
|
||||
return &PaymentInfo{
|
||||
Amount: amount,
|
||||
Network: net,
|
||||
Asset: s.opts.Payment.Asset,
|
||||
PayTo: s.opts.Payment.PayTo,
|
||||
}
|
||||
}
|
||||
|
||||
// Serve starts an MCP gateway with the given options.
|
||||
// For stdio transport, leave Address empty.
|
||||
// For SSE transport, set Address (e.g., ":3000").
|
||||
func Serve(opts Options) error {
|
||||
// Set defaults
|
||||
if opts.Client == nil {
|
||||
opts.Client = client.DefaultClient
|
||||
}
|
||||
if opts.Context == nil {
|
||||
opts.Context = context.Background()
|
||||
}
|
||||
if opts.Logger == nil {
|
||||
opts.Logger = log.Default()
|
||||
}
|
||||
if opts.Registry == nil {
|
||||
return fmt.Errorf("registry is required")
|
||||
}
|
||||
|
||||
server := &Server{
|
||||
opts: opts,
|
||||
tools: make(map[string]*Tool),
|
||||
limiters: make(map[string]*rateLimiter),
|
||||
breakers: make(map[string]*circuitBreaker),
|
||||
}
|
||||
|
||||
// Discover services and build tool list
|
||||
if err := server.discoverServices(); err != nil {
|
||||
return fmt.Errorf("failed to discover services: %w", err)
|
||||
}
|
||||
|
||||
// Watch for service changes
|
||||
go server.watchServices()
|
||||
|
||||
// Start server based on transport
|
||||
if opts.Address != "" {
|
||||
return server.serveHTTP()
|
||||
}
|
||||
return server.serveStdio()
|
||||
}
|
||||
|
||||
// ListenAndServe is a convenience function that starts an MCP gateway on the given address.
|
||||
func ListenAndServe(address string, opts Options) error {
|
||||
opts.Address = address
|
||||
return Serve(opts)
|
||||
}
|
||||
|
||||
// discoverServices queries the registry and builds the tool list
|
||||
func (s *Server) discoverServices() error {
|
||||
services, err := s.opts.Registry.ListServices()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
s.toolsMu.Lock()
|
||||
defer s.toolsMu.Unlock()
|
||||
|
||||
if err := s.discoverReflectedGRPC(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, svc := range services {
|
||||
// Get full service details
|
||||
fullSvcs, err := s.opts.Registry.GetService(svc.Name)
|
||||
if err != nil || len(fullSvcs) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
// Convert endpoints to tools
|
||||
for _, ep := range fullSvcs[0].Endpoints {
|
||||
toolName := fmt.Sprintf("%s.%s", svc.Name, ep.Name)
|
||||
|
||||
// Build input schema from endpoint request type
|
||||
inputSchema := s.buildInputSchema(ep.Request)
|
||||
|
||||
// Get description from endpoint metadata (set by service during registration)
|
||||
description := fmt.Sprintf("Call %s on %s service", ep.Name, svc.Name)
|
||||
if ep.Metadata != nil {
|
||||
if desc, ok := ep.Metadata["description"]; ok && desc != "" {
|
||||
description = desc
|
||||
}
|
||||
}
|
||||
|
||||
tool := &Tool{
|
||||
Name: toolName,
|
||||
Description: description,
|
||||
InputSchema: inputSchema,
|
||||
Service: svc.Name,
|
||||
Endpoint: ep.Name,
|
||||
}
|
||||
|
||||
// Extract scopes from endpoint metadata
|
||||
if ep.Metadata != nil {
|
||||
if scopes, ok := ep.Metadata["scopes"]; ok && scopes != "" {
|
||||
tool.Scopes = strings.Split(scopes, ",")
|
||||
}
|
||||
}
|
||||
|
||||
// Gateway-level Scopes override service-level scopes
|
||||
if s.opts.Scopes != nil {
|
||||
if scopes, ok := s.opts.Scopes[toolName]; ok {
|
||||
tool.Scopes = scopes
|
||||
}
|
||||
}
|
||||
|
||||
// Add example from metadata if available
|
||||
if ep.Metadata != nil {
|
||||
if example, ok := ep.Metadata["example"]; ok && example != "" {
|
||||
inputSchema["examples"] = []string{example}
|
||||
}
|
||||
}
|
||||
|
||||
s.tools[toolName] = tool
|
||||
|
||||
// Create rate limiter for this tool if rate limiting is configured
|
||||
if s.opts.RateLimit != nil && s.opts.RateLimit.RequestsPerSecond > 0 {
|
||||
s.limitersMu.Lock()
|
||||
if _, exists := s.limiters[toolName]; !exists {
|
||||
s.limiters[toolName] = newRateLimiter(
|
||||
s.opts.RateLimit.RequestsPerSecond,
|
||||
s.opts.RateLimit.Burst,
|
||||
)
|
||||
}
|
||||
s.limitersMu.Unlock()
|
||||
}
|
||||
|
||||
// Create circuit breaker for this tool if configured
|
||||
if s.opts.CircuitBreaker != nil {
|
||||
s.breakersMu.Lock()
|
||||
if _, exists := s.breakers[toolName]; !exists {
|
||||
s.breakers[toolName] = newCircuitBreaker(*s.opts.CircuitBreaker)
|
||||
}
|
||||
s.breakersMu.Unlock()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Register framework primitives as tools.
|
||||
// When Auth is configured, they require micro:admin scope.
|
||||
s.registerFrameworkTools()
|
||||
|
||||
s.opts.Logger.Printf("[mcp] Discovered %d tools from %d services (incl. framework)", len(s.tools), len(services))
|
||||
return nil
|
||||
}
|
||||
|
||||
// registerFrameworkTools adds registry, broker, store, and config as MCP tools.
|
||||
func (s *Server) registerFrameworkTools() {
|
||||
addFramework := func(tool *Tool) {
|
||||
// When auth is configured, require micro:admin scope
|
||||
if s.opts.Auth != nil {
|
||||
tool.Scopes = []string{"micro:admin"}
|
||||
}
|
||||
s.tools[tool.Name] = tool
|
||||
if s.opts.RateLimit != nil && s.opts.RateLimit.RequestsPerSecond > 0 {
|
||||
s.limitersMu.Lock()
|
||||
if _, exists := s.limiters[tool.Name]; !exists {
|
||||
s.limiters[tool.Name] = newRateLimiter(s.opts.RateLimit.RequestsPerSecond, s.opts.RateLimit.Burst)
|
||||
}
|
||||
s.limitersMu.Unlock()
|
||||
}
|
||||
if s.opts.CircuitBreaker != nil {
|
||||
s.breakersMu.Lock()
|
||||
if _, exists := s.breakers[tool.Name]; !exists {
|
||||
s.breakers[tool.Name] = newCircuitBreaker(*s.opts.CircuitBreaker)
|
||||
}
|
||||
s.breakersMu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
addFramework(&Tool{
|
||||
Name: "micro_registry_list",
|
||||
Description: "List all registered services in the service registry",
|
||||
InputSchema: map[string]interface{}{"type": "object", "properties": map[string]interface{}{}},
|
||||
Handler: func(input map[string]interface{}) (interface{}, error) {
|
||||
services, err := s.opts.Registry.ListServices()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var names []string
|
||||
for _, svc := range services {
|
||||
names = append(names, svc.Name)
|
||||
}
|
||||
return map[string]interface{}{"services": names}, nil
|
||||
},
|
||||
})
|
||||
|
||||
addFramework(&Tool{
|
||||
Name: "micro_registry_get",
|
||||
Description: "Get details for a registered service including nodes and endpoints",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"name": map[string]interface{}{"type": "string", "description": "Service name"},
|
||||
},
|
||||
},
|
||||
Handler: func(input map[string]interface{}) (interface{}, error) {
|
||||
name, _ := input["name"].(string)
|
||||
if name == "" {
|
||||
return nil, fmt.Errorf("name is required")
|
||||
}
|
||||
services, err := s.opts.Registry.GetService(name)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return services, nil
|
||||
},
|
||||
})
|
||||
|
||||
addFramework(&Tool{
|
||||
Name: "micro_store_list",
|
||||
Description: "List keys in the data store",
|
||||
InputSchema: map[string]interface{}{"type": "object", "properties": map[string]interface{}{}},
|
||||
Handler: func(input map[string]interface{}) (interface{}, error) {
|
||||
keys, err := store.List()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return map[string]interface{}{"keys": keys}, nil
|
||||
},
|
||||
})
|
||||
|
||||
addFramework(&Tool{
|
||||
Name: "micro_store_read",
|
||||
Description: "Read a record from the data store by key",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"key": map[string]interface{}{"type": "string", "description": "Record key"},
|
||||
},
|
||||
},
|
||||
Handler: func(input map[string]interface{}) (interface{}, error) {
|
||||
key, _ := input["key"].(string)
|
||||
if key == "" {
|
||||
return nil, fmt.Errorf("key is required")
|
||||
}
|
||||
records, err := store.Read(key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(records) == 0 {
|
||||
return map[string]interface{}{"error": "not found"}, nil
|
||||
}
|
||||
return map[string]interface{}{"key": key, "value": string(records[0].Value)}, nil
|
||||
},
|
||||
})
|
||||
|
||||
addFramework(&Tool{
|
||||
Name: "micro_store_write",
|
||||
Description: "Write a record to the data store",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"key": map[string]interface{}{"type": "string", "description": "Record key"},
|
||||
"value": map[string]interface{}{"type": "string", "description": "Record value"},
|
||||
},
|
||||
},
|
||||
Handler: func(input map[string]interface{}) (interface{}, error) {
|
||||
key, _ := input["key"].(string)
|
||||
value, _ := input["value"].(string)
|
||||
if key == "" {
|
||||
return nil, fmt.Errorf("key is required")
|
||||
}
|
||||
if err := store.Write(&store.Record{Key: key, Value: []byte(value)}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return map[string]interface{}{"status": "ok", "key": key}, nil
|
||||
},
|
||||
})
|
||||
|
||||
addFramework(&Tool{
|
||||
Name: "micro_broker_publish",
|
||||
Description: "Publish a message to a broker topic",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"topic": map[string]interface{}{"type": "string", "description": "Topic name"},
|
||||
"message": map[string]interface{}{"type": "string", "description": "Message body"},
|
||||
},
|
||||
},
|
||||
Handler: func(input map[string]interface{}) (interface{}, error) {
|
||||
topic, _ := input["topic"].(string)
|
||||
message, _ := input["message"].(string)
|
||||
if topic == "" {
|
||||
return nil, fmt.Errorf("topic is required")
|
||||
}
|
||||
b := broker.DefaultBroker
|
||||
if err := b.Connect(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := b.Publish(topic, &broker.Message{Body: []byte(message)}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return map[string]interface{}{"status": "ok", "topic": topic}, nil
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// buildInputSchema converts registry value type information to JSON schema
|
||||
func (s *Server) buildInputSchema(value *registry.Value) map[string]interface{} {
|
||||
schema := map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": make(map[string]interface{}),
|
||||
}
|
||||
|
||||
if value == nil || len(value.Values) == 0 {
|
||||
return schema
|
||||
}
|
||||
|
||||
properties := schema["properties"].(map[string]interface{})
|
||||
for _, field := range value.Values {
|
||||
properties[field.Name] = map[string]interface{}{
|
||||
"type": s.mapGoTypeToJSON(field.Type),
|
||||
"description": fmt.Sprintf("%s field", field.Name),
|
||||
}
|
||||
}
|
||||
|
||||
return schema
|
||||
}
|
||||
|
||||
// mapGoTypeToJSON maps Go types to JSON schema types
|
||||
func (s *Server) mapGoTypeToJSON(goType string) string {
|
||||
switch goType {
|
||||
case "string":
|
||||
return "string"
|
||||
case "int", "int32", "int64", "uint", "uint32", "uint64":
|
||||
return "integer"
|
||||
case "float32", "float64":
|
||||
return "number"
|
||||
case "bool":
|
||||
return "boolean"
|
||||
default:
|
||||
return "object"
|
||||
}
|
||||
}
|
||||
|
||||
// watchServices watches for service registry changes
|
||||
func (s *Server) watchServices() {
|
||||
if s.watching {
|
||||
return
|
||||
}
|
||||
s.watching = true
|
||||
|
||||
watcher, err := s.opts.Registry.Watch()
|
||||
if err != nil {
|
||||
s.opts.Logger.Printf("[mcp] Failed to watch registry: %v", err)
|
||||
return
|
||||
}
|
||||
defer watcher.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-s.opts.Context.Done():
|
||||
return
|
||||
default:
|
||||
_, err := watcher.Next()
|
||||
if err != nil {
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
// Rediscover services on any change
|
||||
if err := s.discoverServices(); err != nil {
|
||||
s.opts.Logger.Printf("[mcp] Failed to rediscover services: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// serveHTTP starts an HTTP server with SSE and WebSocket transports
|
||||
func (s *Server) serveHTTP() error {
|
||||
mux := http.NewServeMux()
|
||||
|
||||
// MCP endpoints. Tool calls can be gated behind an x402 payment
|
||||
// (enforced per-tool inside handleCallTool); listing tools and health
|
||||
// stay free.
|
||||
if s.opts.Payment != nil {
|
||||
net := s.opts.Payment.Network
|
||||
if net == "" {
|
||||
net = "base"
|
||||
}
|
||||
s.opts.Logger.Printf("[mcp] x402 payments enabled (network=%s, payTo=%s)", net, s.opts.Payment.PayTo)
|
||||
}
|
||||
mux.HandleFunc("/mcp/tools", s.handleListTools)
|
||||
mux.HandleFunc("/mcp/call", s.handleCallTool)
|
||||
mux.HandleFunc("/health", s.handleHealth)
|
||||
|
||||
// WebSocket endpoint for bidirectional streaming
|
||||
ws := NewWebSocketTransport(s)
|
||||
mux.Handle("/mcp/ws", ws)
|
||||
|
||||
s.server = &http.Server{
|
||||
Addr: s.opts.Address,
|
||||
Handler: mux,
|
||||
}
|
||||
|
||||
s.opts.Logger.Printf("[mcp] MCP gateway listening on %s (HTTP + WebSocket)", s.opts.Address)
|
||||
return s.server.ListenAndServe()
|
||||
}
|
||||
|
||||
// serveStdio starts stdio-based MCP server (for Claude Code, etc.)
|
||||
func (s *Server) serveStdio() error {
|
||||
transport := NewStdioTransport(s)
|
||||
return transport.Serve()
|
||||
}
|
||||
|
||||
// handleListTools returns the list of available tools
|
||||
func (s *Server) handleListTools(w http.ResponseWriter, r *http.Request) {
|
||||
if s.opts.AuthFunc != nil {
|
||||
if err := s.opts.AuthFunc(r); err != nil {
|
||||
http.Error(w, "Unauthorized", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
s.toolsMu.RLock()
|
||||
tools := make([]*Tool, 0, len(s.tools))
|
||||
for _, tool := range s.tools {
|
||||
// Attach payment info for the catalog. Copy when pricing so the
|
||||
// shared tool struct isn't mutated.
|
||||
if pay := s.paymentFor(tool.Name); pay != nil {
|
||||
cp := *tool
|
||||
cp.Payment = pay
|
||||
tools = append(tools, &cp)
|
||||
continue
|
||||
}
|
||||
tools = append(tools, tool)
|
||||
}
|
||||
s.toolsMu.RUnlock()
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(map[string]interface{}{
|
||||
"tools": tools,
|
||||
})
|
||||
}
|
||||
|
||||
// handleCallTool executes a tool (makes an RPC call)
|
||||
func (s *Server) handleCallTool(w http.ResponseWriter, r *http.Request) {
|
||||
if s.opts.AuthFunc != nil {
|
||||
if err := s.opts.AuthFunc(r); err != nil {
|
||||
http.Error(w, "Unauthorized", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Parse request
|
||||
var req struct {
|
||||
Tool string `json:"tool"`
|
||||
Input map[string]interface{} `json:"input"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
// Get tool info
|
||||
s.toolsMu.RLock()
|
||||
tool, exists := s.tools[req.Tool]
|
||||
s.toolsMu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
http.Error(w, "Tool not found", http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
|
||||
// x402 payment gate: require the tool's amount before doing work.
|
||||
// Free tools (amount "" or "0") pass through; Require writes the 402
|
||||
// challenge itself when payment is missing or invalid.
|
||||
if s.opts.Payment != nil {
|
||||
if !s.opts.Payment.Require(w, r, s.opts.Payment.AmountFor(req.Tool), req.Tool) {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Generate trace ID for this call
|
||||
traceID := uuid.New().String()
|
||||
|
||||
// Start OTel span (noop if TraceProvider is nil)
|
||||
ctx, span := s.startToolSpan(r.Context(), req.Tool, "http", traceID)
|
||||
defer span.End()
|
||||
|
||||
// Authenticate and authorize
|
||||
var account *auth.Account
|
||||
if s.opts.Auth != nil {
|
||||
token := r.Header.Get("Authorization")
|
||||
token = strings.TrimPrefix(token, "Bearer ")
|
||||
if token == "" {
|
||||
span.SetAttributes(attribute.Bool(AttrAuthAllowed, false), attribute.String(AttrAuthDeniedReason, "missing token"))
|
||||
setSpanError(span, fmt.Errorf("missing token"))
|
||||
s.audit(AuditRecord{TraceID: traceID, Timestamp: time.Now(), Tool: req.Tool, Allowed: false, DeniedReason: "missing token"})
|
||||
http.Error(w, "Unauthorized", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
acc, err := s.opts.Auth.Inspect(token)
|
||||
if err != nil {
|
||||
span.SetAttributes(attribute.Bool(AttrAuthAllowed, false), attribute.String(AttrAuthDeniedReason, "invalid token"))
|
||||
setSpanError(span, fmt.Errorf("invalid token"))
|
||||
s.audit(AuditRecord{TraceID: traceID, Timestamp: time.Now(), Tool: req.Tool, Allowed: false, DeniedReason: "invalid token"})
|
||||
http.Error(w, "Unauthorized", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
account = acc
|
||||
span.SetAttributes(attribute.String(AttrAccountID, account.ID))
|
||||
|
||||
// Check per-tool scopes
|
||||
if len(tool.Scopes) > 0 {
|
||||
span.SetAttributes(attribute.StringSlice(AttrScopesRequired, tool.Scopes))
|
||||
if !hasScope(account.Scopes, tool.Scopes) {
|
||||
span.SetAttributes(attribute.Bool(AttrAuthAllowed, false), attribute.String(AttrAuthDeniedReason, "insufficient scopes"))
|
||||
setSpanError(span, fmt.Errorf("insufficient scopes"))
|
||||
s.audit(AuditRecord{
|
||||
TraceID: traceID, Timestamp: time.Now(), Tool: req.Tool,
|
||||
AccountID: account.ID, ScopesRequired: tool.Scopes,
|
||||
Allowed: false, DeniedReason: "insufficient scopes",
|
||||
})
|
||||
http.Error(w, "Forbidden: insufficient scopes", http.StatusForbidden)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Rate limit check
|
||||
if err := s.allowRate(req.Tool); err != nil {
|
||||
span.SetAttributes(attribute.Bool(AttrRateLimited, true))
|
||||
setSpanError(span, err)
|
||||
accountID := ""
|
||||
if account != nil {
|
||||
accountID = account.ID
|
||||
}
|
||||
s.audit(AuditRecord{
|
||||
TraceID: traceID, Timestamp: time.Now(), Tool: req.Tool,
|
||||
AccountID: accountID, Allowed: false, DeniedReason: "rate limited",
|
||||
})
|
||||
http.Error(w, "Rate limit exceeded", http.StatusTooManyRequests)
|
||||
return
|
||||
}
|
||||
|
||||
span.SetAttributes(attribute.Bool(AttrAuthAllowed, true))
|
||||
|
||||
// Circuit breaker check
|
||||
if err := s.allowCircuit(req.Tool); err != nil {
|
||||
span.SetAttributes(attribute.String("mcp.circuit_breaker", "open"))
|
||||
setSpanError(span, err)
|
||||
accountID := ""
|
||||
if account != nil {
|
||||
accountID = account.ID
|
||||
}
|
||||
s.audit(AuditRecord{
|
||||
TraceID: traceID, Timestamp: time.Now(), Tool: req.Tool,
|
||||
AccountID: accountID, Allowed: false, DeniedReason: "circuit breaker open",
|
||||
})
|
||||
http.Error(w, "Service unavailable: circuit breaker open", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
|
||||
// Build context with tracing metadata
|
||||
// OTel trace context was already injected by startToolSpan; add MCP metadata.
|
||||
md, _ := metadata.FromContext(ctx)
|
||||
if md == nil {
|
||||
md = make(metadata.Metadata)
|
||||
}
|
||||
md.Set(TraceIDKey, traceID)
|
||||
md.Set(ToolNameKey, req.Tool)
|
||||
if account != nil {
|
||||
md.Set(AccountIDKey, account.ID)
|
||||
}
|
||||
ctx = metadata.NewContext(ctx, md)
|
||||
|
||||
start := time.Now()
|
||||
|
||||
// Framework tools have a direct handler; service tools go through RPC.
|
||||
if tool.Handler != nil {
|
||||
result, err := tool.Handler(req.Input)
|
||||
if err != nil {
|
||||
setSpanError(span, err)
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(result)
|
||||
return
|
||||
}
|
||||
|
||||
// Convert input to JSON bytes for RPC call
|
||||
inputBytes, err := json.Marshal(req.Input)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// Make RPC call
|
||||
rpcReq := s.opts.Client.NewRequest(tool.Service, tool.Endpoint, &bytes.Frame{Data: inputBytes})
|
||||
var rsp bytes.Frame
|
||||
|
||||
if err := s.opts.Client.Call(ctx, rpcReq, &rsp); err != nil {
|
||||
s.recordCircuit(req.Tool, false)
|
||||
setSpanError(span, err)
|
||||
s.opts.Logger.Printf("[mcp] RPC call failed: %v", err)
|
||||
accountID := ""
|
||||
if account != nil {
|
||||
accountID = account.ID
|
||||
}
|
||||
s.audit(AuditRecord{
|
||||
TraceID: traceID, Timestamp: time.Now(), Tool: req.Tool,
|
||||
AccountID: accountID, ScopesRequired: tool.Scopes,
|
||||
Allowed: true, Duration: time.Since(start), Error: err.Error(),
|
||||
})
|
||||
http.Error(w, fmt.Sprintf("RPC call failed: %v", err), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
s.recordCircuit(req.Tool, true)
|
||||
setSpanOK(span)
|
||||
|
||||
// Audit successful call
|
||||
accountID := ""
|
||||
if account != nil {
|
||||
accountID = account.ID
|
||||
}
|
||||
s.audit(AuditRecord{
|
||||
TraceID: traceID, Timestamp: time.Now(), Tool: req.Tool,
|
||||
AccountID: accountID, ScopesRequired: tool.Scopes,
|
||||
Allowed: true, Duration: time.Since(start),
|
||||
})
|
||||
|
||||
// Return response with trace ID
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.Header().Set(TraceIDKey, traceID)
|
||||
json.NewEncoder(w).Encode(map[string]interface{}{
|
||||
"result": json.RawMessage(rsp.Data),
|
||||
"trace_id": traceID,
|
||||
})
|
||||
}
|
||||
|
||||
// handleHealth returns gateway health status
|
||||
func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
|
||||
s.toolsMu.RLock()
|
||||
toolCount := len(s.tools)
|
||||
s.toolsMu.RUnlock()
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(map[string]interface{}{
|
||||
"status": "ok",
|
||||
"tools": toolCount,
|
||||
})
|
||||
}
|
||||
|
||||
// Stop gracefully shuts down the MCP gateway
|
||||
func (s *Server) Stop() error {
|
||||
if s.server != nil {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
return s.server.Shutdown(ctx)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetTools returns the current list of available tools
|
||||
func (s *Server) GetTools() []*Tool {
|
||||
s.toolsMu.RLock()
|
||||
defer s.toolsMu.RUnlock()
|
||||
|
||||
tools := make([]*Tool, 0, len(s.tools))
|
||||
for _, tool := range s.tools {
|
||||
tools = append(tools, tool)
|
||||
}
|
||||
return tools
|
||||
}
|
||||
|
||||
// audit emits an audit record if an AuditFunc is configured.
|
||||
func (s *Server) audit(record AuditRecord) {
|
||||
if s.opts.AuditFunc != nil {
|
||||
s.opts.AuditFunc(record)
|
||||
}
|
||||
}
|
||||
|
||||
// allowRate checks if the tool call is allowed under the configured rate limit.
|
||||
// Returns nil if allowed, non-nil error if rate-limited.
|
||||
func (s *Server) allowRate(toolName string) error {
|
||||
if s.opts.RateLimit == nil {
|
||||
return nil
|
||||
}
|
||||
s.limitersMu.RLock()
|
||||
limiter, ok := s.limiters[toolName]
|
||||
s.limitersMu.RUnlock()
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
if !limiter.Allow() {
|
||||
return fmt.Errorf("rate limit exceeded for tool %s", toolName)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// allowCircuit checks if the tool call is allowed by the circuit breaker.
|
||||
// Returns nil if allowed, non-nil error if the circuit is open.
|
||||
func (s *Server) allowCircuit(toolName string) error {
|
||||
if s.opts.CircuitBreaker == nil {
|
||||
return nil
|
||||
}
|
||||
s.breakersMu.RLock()
|
||||
cb, ok := s.breakers[toolName]
|
||||
s.breakersMu.RUnlock()
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
return cb.Allow()
|
||||
}
|
||||
|
||||
// recordCircuit records a success or failure for the tool's circuit breaker.
|
||||
func (s *Server) recordCircuit(toolName string, success bool) {
|
||||
if s.opts.CircuitBreaker == nil {
|
||||
return
|
||||
}
|
||||
s.breakersMu.RLock()
|
||||
cb, ok := s.breakers[toolName]
|
||||
s.breakersMu.RUnlock()
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if success {
|
||||
cb.RecordSuccess()
|
||||
} else {
|
||||
cb.RecordFailure()
|
||||
}
|
||||
}
|
||||
|
||||
// hasScope checks if the account has at least one of the required scopes.
|
||||
func hasScope(accountScopes, requiredScopes []string) bool {
|
||||
for _, req := range requiredScopes {
|
||||
for _, have := range accountScopes {
|
||||
if strings.EqualFold(have, req) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Example shows how to use the MCP gateway in your code
|
||||
func Example() {
|
||||
// This function is never called - it's just documentation
|
||||
_ = func() {
|
||||
// In your service code:
|
||||
// service := micro.NewService("myservice", )
|
||||
// service.Init()
|
||||
|
||||
// Start MCP gateway
|
||||
go func() {
|
||||
if err := Serve(Options{
|
||||
Registry: registry.DefaultRegistry,
|
||||
Address: ":3000",
|
||||
}); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
}()
|
||||
|
||||
// service.Run()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,568 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go-micro.dev/v6/auth"
|
||||
"go-micro.dev/v6/client"
|
||||
"go-micro.dev/v6/registry"
|
||||
)
|
||||
|
||||
// mockAuth implements auth.Auth for testing.
|
||||
type mockAuth struct {
|
||||
accounts map[string]*auth.Account // token -> account
|
||||
}
|
||||
|
||||
func (m *mockAuth) Init(...auth.Option) {}
|
||||
func (m *mockAuth) Options() auth.Options { return auth.Options{} }
|
||||
func (m *mockAuth) Generate(string, ...auth.GenerateOption) (*auth.Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAuth) Token(...auth.TokenOption) (*auth.Token, error) { return nil, nil }
|
||||
func (m *mockAuth) String() string { return "mock" }
|
||||
|
||||
func (m *mockAuth) Inspect(token string) (*auth.Account, error) {
|
||||
acc, ok := m.accounts[token]
|
||||
if !ok {
|
||||
return nil, auth.ErrInvalidToken
|
||||
}
|
||||
return acc, nil
|
||||
}
|
||||
|
||||
// newTestServer creates a Server with pre-populated tools for testing.
|
||||
func newTestServer(opts Options) *Server {
|
||||
if opts.Logger == nil {
|
||||
opts.Logger = testLogger()
|
||||
}
|
||||
if opts.Context == nil {
|
||||
opts.Context = context.Background()
|
||||
}
|
||||
if opts.Client == nil {
|
||||
opts.Client = client.DefaultClient
|
||||
}
|
||||
s := &Server{
|
||||
opts: opts,
|
||||
tools: make(map[string]*Tool),
|
||||
limiters: make(map[string]*rateLimiter),
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// testLogger returns a silent logger for tests.
|
||||
func testLogger() *log.Logger {
|
||||
return log.New(nopWriter{}, "", 0)
|
||||
}
|
||||
|
||||
type nopWriter struct{}
|
||||
|
||||
func (nopWriter) Write(p []byte) (int, error) { return len(p), nil }
|
||||
|
||||
// --- Tests ---
|
||||
|
||||
func TestHasScope(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
account []string
|
||||
required []string
|
||||
want bool
|
||||
}{
|
||||
{"match single", []string{"blog:write"}, []string{"blog:write"}, true},
|
||||
{"match one of many", []string{"blog:read", "blog:write"}, []string{"blog:write"}, true},
|
||||
{"no match", []string{"blog:read"}, []string{"blog:write"}, false},
|
||||
{"empty required", []string{"blog:read"}, nil, false},
|
||||
{"empty account", nil, []string{"blog:write"}, false},
|
||||
{"case insensitive", []string{"Blog:Write"}, []string{"blog:write"}, true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := hasScope(tt.account, tt.required)
|
||||
if got != tt.want {
|
||||
t.Errorf("hasScope(%v, %v) = %v, want %v", tt.account, tt.required, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolScopesFromMetadata(t *testing.T) {
|
||||
// Create a mock registry with endpoints that have scope metadata
|
||||
reg := registry.NewMemoryRegistry()
|
||||
svc := ®istry.Service{
|
||||
Name: "blog",
|
||||
Nodes: []*registry.Node{{
|
||||
Id: "blog-1",
|
||||
Address: "localhost:9090",
|
||||
}},
|
||||
Endpoints: []*registry.Endpoint{
|
||||
{
|
||||
Name: "Blog.Create",
|
||||
Metadata: map[string]string{
|
||||
"description": "Create a blog post",
|
||||
"scopes": "blog:write,blog:admin",
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Blog.Read",
|
||||
Metadata: map[string]string{
|
||||
"description": "Read a blog post",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
if err := reg.Register(svc); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
s := newTestServer(Options{Registry: reg})
|
||||
if err := s.discoverServices(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Check that scopes are populated
|
||||
createTool := s.tools["blog.Blog.Create"]
|
||||
if createTool == nil {
|
||||
t.Fatal("expected tool blog.Blog.Create")
|
||||
}
|
||||
if len(createTool.Scopes) != 2 || createTool.Scopes[0] != "blog:write" || createTool.Scopes[1] != "blog:admin" {
|
||||
t.Errorf("unexpected scopes: %v", createTool.Scopes)
|
||||
}
|
||||
|
||||
readTool := s.tools["blog.Blog.Read"]
|
||||
if readTool == nil {
|
||||
t.Fatal("expected tool blog.Blog.Read")
|
||||
}
|
||||
if len(readTool.Scopes) != 0 {
|
||||
t.Errorf("expected no scopes for read, got: %v", readTool.Scopes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleCallTool_AuthRequired(t *testing.T) {
|
||||
ma := &mockAuth{
|
||||
accounts: map[string]*auth.Account{
|
||||
"valid-token": {ID: "user-1", Scopes: []string{"blog:write"}},
|
||||
"readonly": {ID: "user-2", Scopes: []string{"blog:read"}},
|
||||
},
|
||||
}
|
||||
|
||||
s := newTestServer(Options{Auth: ma})
|
||||
s.tools["blog.Blog.Create"] = &Tool{
|
||||
Name: "blog.Blog.Create",
|
||||
Service: "blog",
|
||||
Endpoint: "Blog.Create",
|
||||
Scopes: []string{"blog:write"},
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
token string
|
||||
wantStatus int
|
||||
}{
|
||||
{"no token", "", http.StatusUnauthorized},
|
||||
{"invalid token", "bad-token", http.StatusUnauthorized},
|
||||
{"valid token with scope", "valid-token", http.StatusInternalServerError}, // RPC will fail (no backend), but auth passes
|
||||
{"valid token without scope", "readonly", http.StatusForbidden},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"tool": "blog.Blog.Create",
|
||||
"input": map[string]interface{}{"title": "hello"},
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp/call", bytes.NewReader(body))
|
||||
if tt.token != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+tt.token)
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
s.handleCallTool(rec, req)
|
||||
|
||||
if rec.Code != tt.wantStatus {
|
||||
t.Errorf("status = %d, want %d, body: %s", rec.Code, tt.wantStatus, rec.Body.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleCallTool_TraceID(t *testing.T) {
|
||||
// Without Auth, tool calls should still generate trace IDs.
|
||||
s := newTestServer(Options{})
|
||||
s.tools["svc.Echo"] = &Tool{
|
||||
Name: "svc.Echo",
|
||||
Service: "svc",
|
||||
Endpoint: "Echo",
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"tool": "svc.Echo",
|
||||
"input": map[string]interface{}{},
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp/call", bytes.NewReader(body))
|
||||
rec := httptest.NewRecorder()
|
||||
s.handleCallTool(rec, req)
|
||||
|
||||
// Even though the RPC fails (no backend), the trace ID header should be absent
|
||||
// only when the call didn't reach the RPC stage; but in this no-auth case it
|
||||
// should reach the RPC call and fail. Check we get a response.
|
||||
traceID := rec.Header().Get(TraceIDKey)
|
||||
// The RPC call will fail but the error path doesn't set the header.
|
||||
// For a successful call, the trace ID is set. Either way the audit should fire.
|
||||
_ = traceID // trace ID may or may not be in error response header
|
||||
}
|
||||
|
||||
func TestHandleCallTool_AuditFunc(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
var records []AuditRecord
|
||||
|
||||
auditFn := func(r AuditRecord) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
records = append(records, r)
|
||||
}
|
||||
|
||||
ma := &mockAuth{
|
||||
accounts: map[string]*auth.Account{
|
||||
"tok": {ID: "u1", Scopes: []string{"write"}},
|
||||
},
|
||||
}
|
||||
|
||||
s := newTestServer(Options{Auth: ma, AuditFunc: auditFn})
|
||||
s.tools["svc.Do"] = &Tool{
|
||||
Name: "svc.Do",
|
||||
Service: "svc",
|
||||
Endpoint: "Do",
|
||||
Scopes: []string{"write"},
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"tool": "svc.Do",
|
||||
"input": map[string]interface{}{},
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp/call", bytes.NewReader(body))
|
||||
req.Header.Set("Authorization", "Bearer tok")
|
||||
rec := httptest.NewRecorder()
|
||||
s.handleCallTool(rec, req)
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
|
||||
if len(records) == 0 {
|
||||
t.Fatal("expected at least one audit record")
|
||||
}
|
||||
r := records[len(records)-1]
|
||||
if r.AccountID != "u1" {
|
||||
t.Errorf("audit AccountID = %q, want %q", r.AccountID, "u1")
|
||||
}
|
||||
if r.Tool != "svc.Do" {
|
||||
t.Errorf("audit Tool = %q, want %q", r.Tool, "svc.Do")
|
||||
}
|
||||
if r.TraceID == "" {
|
||||
t.Error("audit TraceID is empty")
|
||||
}
|
||||
if !r.Allowed {
|
||||
t.Error("expected audit record Allowed = true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleCallTool_AuditDenied(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
var records []AuditRecord
|
||||
|
||||
auditFn := func(r AuditRecord) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
records = append(records, r)
|
||||
}
|
||||
|
||||
ma := &mockAuth{
|
||||
accounts: map[string]*auth.Account{
|
||||
"tok": {ID: "u1", Scopes: []string{"blog:read"}},
|
||||
},
|
||||
}
|
||||
|
||||
s := newTestServer(Options{Auth: ma, AuditFunc: auditFn})
|
||||
s.tools["svc.Do"] = &Tool{
|
||||
Name: "svc.Do",
|
||||
Service: "svc",
|
||||
Endpoint: "Do",
|
||||
Scopes: []string{"blog:write"},
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"tool": "svc.Do",
|
||||
"input": map[string]interface{}{},
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp/call", bytes.NewReader(body))
|
||||
req.Header.Set("Authorization", "Bearer tok")
|
||||
rec := httptest.NewRecorder()
|
||||
s.handleCallTool(rec, req)
|
||||
|
||||
if rec.Code != http.StatusForbidden {
|
||||
t.Errorf("status = %d, want %d", rec.Code, http.StatusForbidden)
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
|
||||
if len(records) == 0 {
|
||||
t.Fatal("expected audit record for denied call")
|
||||
}
|
||||
r := records[0]
|
||||
if r.Allowed {
|
||||
t.Error("expected Allowed = false")
|
||||
}
|
||||
if r.DeniedReason != "insufficient scopes" {
|
||||
t.Errorf("DeniedReason = %q, want %q", r.DeniedReason, "insufficient scopes")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter(t *testing.T) {
|
||||
rl := newRateLimiter(10, 2)
|
||||
|
||||
// First two should be allowed (burst)
|
||||
if !rl.Allow() {
|
||||
t.Error("first call should be allowed")
|
||||
}
|
||||
if !rl.Allow() {
|
||||
t.Error("second call should be allowed (burst)")
|
||||
}
|
||||
|
||||
// Third should be denied (burst exhausted, no time to refill)
|
||||
if rl.Allow() {
|
||||
t.Error("third call should be denied (burst exhausted)")
|
||||
}
|
||||
|
||||
// Wait for refill
|
||||
time.Sleep(150 * time.Millisecond)
|
||||
|
||||
// Should be allowed again
|
||||
if !rl.Allow() {
|
||||
t.Error("call after refill should be allowed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleCallTool_RateLimit(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
var records []AuditRecord
|
||||
|
||||
s := newTestServer(Options{
|
||||
RateLimit: &RateLimitConfig{RequestsPerSecond: 1, Burst: 1},
|
||||
AuditFunc: func(r AuditRecord) {
|
||||
mu.Lock()
|
||||
records = append(records, r)
|
||||
mu.Unlock()
|
||||
},
|
||||
})
|
||||
s.tools["svc.Do"] = &Tool{
|
||||
Name: "svc.Do",
|
||||
Service: "svc",
|
||||
Endpoint: "Do",
|
||||
}
|
||||
s.limiters["svc.Do"] = newRateLimiter(1, 1)
|
||||
|
||||
makeReq := func() int {
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"tool": "svc.Do",
|
||||
"input": map[string]interface{}{},
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp/call", bytes.NewReader(body))
|
||||
rec := httptest.NewRecorder()
|
||||
s.handleCallTool(rec, req)
|
||||
return rec.Code
|
||||
}
|
||||
|
||||
// First request should pass rate limit (but RPC may fail — that's ok)
|
||||
code1 := makeReq()
|
||||
if code1 == http.StatusTooManyRequests {
|
||||
t.Error("first request should not be rate limited")
|
||||
}
|
||||
|
||||
// Second request should be rate limited
|
||||
code2 := makeReq()
|
||||
if code2 != http.StatusTooManyRequests {
|
||||
t.Errorf("second request status = %d, want %d", code2, http.StatusTooManyRequests)
|
||||
}
|
||||
|
||||
// Check audit records include rate limit denial
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
found := false
|
||||
for _, r := range records {
|
||||
if r.DeniedReason == "rate limited" {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("expected audit record with DeniedReason = 'rate limited'")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleCallTool_NoAuth_NoScope(t *testing.T) {
|
||||
// Without Auth configured, tools without scopes should be accessible
|
||||
s := newTestServer(Options{})
|
||||
s.tools["svc.Echo"] = &Tool{
|
||||
Name: "svc.Echo",
|
||||
Service: "svc",
|
||||
Endpoint: "Echo",
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"tool": "svc.Echo",
|
||||
"input": map[string]interface{}{"msg": "hi"},
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp/call", bytes.NewReader(body))
|
||||
rec := httptest.NewRecorder()
|
||||
s.handleCallTool(rec, req)
|
||||
|
||||
// Should not be 401 or 403 (RPC failure is expected since no backend)
|
||||
if rec.Code == http.StatusUnauthorized || rec.Code == http.StatusForbidden {
|
||||
t.Errorf("unexpected auth error: %d", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolScopesInJSON(t *testing.T) {
|
||||
tool := &Tool{
|
||||
Name: "blog.Blog.Create",
|
||||
Description: "Create a blog post",
|
||||
InputSchema: map[string]interface{}{"type": "object"},
|
||||
Scopes: []string{"blog:write", "blog:admin"},
|
||||
}
|
||||
|
||||
data, err := json.Marshal(tool)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var m map[string]interface{}
|
||||
if err := json.Unmarshal(data, &m); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
scopes, ok := m["scopes"].([]interface{})
|
||||
if !ok {
|
||||
t.Fatal("expected scopes in JSON output")
|
||||
}
|
||||
if len(scopes) != 2 {
|
||||
t.Errorf("expected 2 scopes, got %d", len(scopes))
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolNoScopesOmittedInJSON(t *testing.T) {
|
||||
tool := &Tool{
|
||||
Name: "blog.Blog.Read",
|
||||
Description: "Read a blog post",
|
||||
InputSchema: map[string]interface{}{"type": "object"},
|
||||
}
|
||||
|
||||
data, err := json.Marshal(tool)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var m map[string]interface{}
|
||||
if err := json.Unmarshal(data, &m); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if _, ok := m["scopes"]; ok {
|
||||
t.Error("expected scopes to be omitted when empty")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverServices_RateLimiters(t *testing.T) {
|
||||
reg := registry.NewMemoryRegistry()
|
||||
svc := ®istry.Service{
|
||||
Name: "blog",
|
||||
Nodes: []*registry.Node{{
|
||||
Id: "blog-1",
|
||||
Address: "localhost:9090",
|
||||
}},
|
||||
Endpoints: []*registry.Endpoint{
|
||||
{Name: "Blog.Create"},
|
||||
{Name: "Blog.Read"},
|
||||
},
|
||||
}
|
||||
if err := reg.Register(svc); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
s := newTestServer(Options{
|
||||
Registry: reg,
|
||||
RateLimit: &RateLimitConfig{RequestsPerSecond: 10, Burst: 5},
|
||||
})
|
||||
if err := s.discoverServices(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if len(s.limiters) != len(s.tools) {
|
||||
t.Errorf("expected %d limiters (one per tool), got %d", len(s.tools), len(s.limiters))
|
||||
}
|
||||
for name := range s.tools {
|
||||
if _, ok := s.limiters[name]; !ok {
|
||||
t.Errorf("missing limiter for tool %s", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestScopesFromGatewayOptions(t *testing.T) {
|
||||
reg := registry.NewMemoryRegistry()
|
||||
svc := ®istry.Service{
|
||||
Name: "blog",
|
||||
Nodes: []*registry.Node{{
|
||||
Id: "blog-1",
|
||||
Address: "localhost:9090",
|
||||
}},
|
||||
Endpoints: []*registry.Endpoint{
|
||||
{
|
||||
Name: "Blog.Create",
|
||||
Metadata: map[string]string{
|
||||
"scopes": "blog:write",
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Blog.Delete",
|
||||
},
|
||||
},
|
||||
}
|
||||
if err := reg.Register(svc); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Gateway-level Scopes override service-level metadata scopes
|
||||
s := newTestServer(Options{
|
||||
Registry: reg,
|
||||
Scopes: map[string][]string{
|
||||
"blog.Blog.Create": {"blog:admin"}, // override service scope
|
||||
"blog.Blog.Delete": {"blog:admin", "sudo"}, // add scope to tool without service scope
|
||||
},
|
||||
})
|
||||
if err := s.discoverServices(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Blog.Create should have gateway-level scope (overrides service "blog:write")
|
||||
createTool := s.tools["blog.Blog.Create"]
|
||||
if createTool == nil {
|
||||
t.Fatal("expected tool blog.Blog.Create")
|
||||
}
|
||||
if len(createTool.Scopes) != 1 || createTool.Scopes[0] != "blog:admin" {
|
||||
t.Errorf("expected gateway scopes [blog:admin], got: %v", createTool.Scopes)
|
||||
}
|
||||
|
||||
// Blog.Delete should get gateway-level scopes
|
||||
deleteTool := s.tools["blog.Blog.Delete"]
|
||||
if deleteTool == nil {
|
||||
t.Fatal("expected tool blog.Blog.Delete")
|
||||
}
|
||||
if len(deleteTool.Scopes) != 2 || deleteTool.Scopes[0] != "blog:admin" || deleteTool.Scopes[1] != "sudo" {
|
||||
t.Errorf("expected gateway scopes [blog:admin sudo], got: %v", deleteTool.Scopes)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"go-micro.dev/v6/service"
|
||||
)
|
||||
|
||||
// WithMCP returns a service option that starts an MCP gateway alongside the
|
||||
// service, making all registered handlers discoverable as AI agent tools.
|
||||
// The address parameter specifies where the MCP gateway listens (e.g., ":3000").
|
||||
//
|
||||
// Usage:
|
||||
//
|
||||
// import "go-micro.dev/v6/gateway/mcp"
|
||||
//
|
||||
// service := micro.NewService("users",
|
||||
// mcp.WithMCP(":3000"),
|
||||
// )
|
||||
func WithMCP(address string) service.Option {
|
||||
return func(o *service.Options) {
|
||||
o.AfterStart = append(o.AfterStart, func() error {
|
||||
go func() {
|
||||
_ = ListenAndServe(address, Options{
|
||||
Registry: o.Registry,
|
||||
})
|
||||
}()
|
||||
return nil
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"go-micro.dev/v6/metadata"
|
||||
"go.opentelemetry.io/otel"
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
"go.opentelemetry.io/otel/codes"
|
||||
"go.opentelemetry.io/otel/propagation"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
const instrumentationName = "go-micro.dev/v6/gateway/mcp"
|
||||
|
||||
// Span and attribute names for MCP OpenTelemetry integration.
|
||||
const (
|
||||
spanNameToolCall = "mcp.tool.call"
|
||||
|
||||
// AttrToolName is the tool being called (e.g. "blog.Blog.Create").
|
||||
AttrToolName = "mcp.tool.name"
|
||||
// AttrTransport is the transport type ("http" or "stdio").
|
||||
AttrTransport = "mcp.transport"
|
||||
// AttrAccountID is the authenticated account ID.
|
||||
AttrAccountID = "mcp.account.id"
|
||||
// AttrTraceID is the MCP-specific UUID trace ID (kept for compatibility).
|
||||
AttrTraceID = "mcp.trace_id"
|
||||
// AttrAuthAllowed records whether auth was granted.
|
||||
AttrAuthAllowed = "mcp.auth.allowed"
|
||||
// AttrAuthDeniedReason records why auth was denied.
|
||||
AttrAuthDeniedReason = "mcp.auth.denied_reason"
|
||||
// AttrScopesRequired lists the scopes required by the tool.
|
||||
AttrScopesRequired = "mcp.auth.scopes_required"
|
||||
// AttrRateLimited records whether the call was rate-limited.
|
||||
AttrRateLimited = "mcp.rate_limited"
|
||||
)
|
||||
|
||||
// tracer returns the OTel tracer from the configured provider.
|
||||
// If no TraceProvider is set, returns a noop tracer.
|
||||
func (s *Server) tracer() trace.Tracer {
|
||||
if s.opts.TraceProvider != nil {
|
||||
return s.opts.TraceProvider.Tracer(instrumentationName)
|
||||
}
|
||||
return trace.NewNoopTracerProvider().Tracer(instrumentationName)
|
||||
}
|
||||
|
||||
// startToolSpan creates a new server span for an MCP tool call.
|
||||
// It extracts any incoming trace context from metadata and injects
|
||||
// the new span's context back into metadata for downstream propagation.
|
||||
func (s *Server) startToolSpan(ctx context.Context, toolName, transport, mcpTraceID string) (context.Context, trace.Span) {
|
||||
if s.opts.TraceProvider == nil {
|
||||
return ctx, trace.SpanFromContext(ctx)
|
||||
}
|
||||
|
||||
// Extract incoming trace context from go-micro metadata (if any).
|
||||
md, ok := metadata.FromContext(ctx)
|
||||
if ok {
|
||||
carrier := metadataCarrier(md)
|
||||
ctx = otel.GetTextMapPropagator().Extract(ctx, carrier)
|
||||
}
|
||||
|
||||
ctx, span := s.tracer().Start(ctx, spanNameToolCall,
|
||||
trace.WithSpanKind(trace.SpanKindServer),
|
||||
trace.WithAttributes(
|
||||
attribute.String(AttrToolName, toolName),
|
||||
attribute.String(AttrTransport, transport),
|
||||
attribute.String(AttrTraceID, mcpTraceID),
|
||||
),
|
||||
)
|
||||
|
||||
// Inject OTel trace context back into metadata so downstream
|
||||
// RPC calls (via client wrappers) continue the trace.
|
||||
if md == nil {
|
||||
md = make(metadata.Metadata)
|
||||
}
|
||||
carrier := make(propagation.MapCarrier)
|
||||
otel.GetTextMapPropagator().Inject(ctx, carrier)
|
||||
for k, v := range carrier {
|
||||
md.Set(k, v)
|
||||
}
|
||||
ctx = metadata.NewContext(ctx, md)
|
||||
|
||||
return ctx, span
|
||||
}
|
||||
|
||||
// setSpanOK marks a span as successful.
|
||||
func setSpanOK(span trace.Span) {
|
||||
span.SetStatus(codes.Ok, "")
|
||||
}
|
||||
|
||||
// setSpanError records an error on the span.
|
||||
func setSpanError(span trace.Span, err error) {
|
||||
span.RecordError(err)
|
||||
span.SetStatus(codes.Error, err.Error())
|
||||
}
|
||||
|
||||
// metadataCarrier adapts go-micro metadata to OTel's TextMapCarrier.
|
||||
type metadataCarrier metadata.Metadata
|
||||
|
||||
func (c metadataCarrier) Get(key string) string {
|
||||
v, _ := metadata.Metadata(c).Get(key)
|
||||
return v
|
||||
}
|
||||
|
||||
func (c metadataCarrier) Set(key, value string) {
|
||||
metadata.Metadata(c).Set(key, value)
|
||||
}
|
||||
|
||||
func (c metadataCarrier) Keys() []string {
|
||||
keys := make([]string, 0, len(c))
|
||||
for k := range c {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
return keys
|
||||
}
|
||||
@@ -0,0 +1,334 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"go-micro.dev/v6/auth"
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
"go.opentelemetry.io/otel/codes"
|
||||
sdktrace "go.opentelemetry.io/otel/sdk/trace"
|
||||
"go.opentelemetry.io/otel/sdk/trace/tracetest"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
// newTestTP creates a TracerProvider with an in-memory exporter.
|
||||
func newTestTP() (*tracetest.InMemoryExporter, trace.TracerProvider) {
|
||||
exp := tracetest.NewInMemoryExporter()
|
||||
tp := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exp))
|
||||
return exp, tp
|
||||
}
|
||||
|
||||
func TestOTel_SpanCreated(t *testing.T) {
|
||||
exp, tp := newTestTP()
|
||||
|
||||
s := newTestServer(Options{TraceProvider: tp})
|
||||
s.tools["svc.Echo"] = &Tool{
|
||||
Name: "svc.Echo",
|
||||
Service: "svc",
|
||||
Endpoint: "Echo",
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"tool": "svc.Echo",
|
||||
"input": map[string]interface{}{"msg": "hi"},
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp/call", bytes.NewReader(body))
|
||||
rec := httptest.NewRecorder()
|
||||
s.handleCallTool(rec, req)
|
||||
|
||||
// RPC will fail (no backend), but a span should still be created.
|
||||
spans := exp.GetSpans()
|
||||
if len(spans) == 0 {
|
||||
t.Fatal("expected at least one span")
|
||||
}
|
||||
|
||||
span := spans[0]
|
||||
if span.Name != spanNameToolCall {
|
||||
t.Errorf("span name = %q, want %q", span.Name, spanNameToolCall)
|
||||
}
|
||||
if span.SpanKind != trace.SpanKindServer {
|
||||
t.Errorf("span kind = %v, want %v", span.SpanKind, trace.SpanKindServer)
|
||||
}
|
||||
|
||||
// Check attributes
|
||||
assertAttr(t, span.Attributes, AttrToolName, "svc.Echo")
|
||||
assertAttr(t, span.Attributes, AttrTransport, "http")
|
||||
}
|
||||
|
||||
func TestOTel_SpanAttributes_AuthDenied(t *testing.T) {
|
||||
exp, tp := newTestTP()
|
||||
|
||||
ma := &mockAuth{
|
||||
accounts: map[string]*auth.Account{
|
||||
"tok": {ID: "user-1", Scopes: []string{"blog:read"}},
|
||||
},
|
||||
}
|
||||
|
||||
s := newTestServer(Options{TraceProvider: tp, Auth: ma})
|
||||
s.tools["svc.Do"] = &Tool{
|
||||
Name: "svc.Do",
|
||||
Service: "svc",
|
||||
Endpoint: "Do",
|
||||
Scopes: []string{"blog:write"},
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"tool": "svc.Do",
|
||||
"input": map[string]interface{}{},
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp/call", bytes.NewReader(body))
|
||||
req.Header.Set("Authorization", "Bearer tok")
|
||||
rec := httptest.NewRecorder()
|
||||
s.handleCallTool(rec, req)
|
||||
|
||||
if rec.Code != http.StatusForbidden {
|
||||
t.Fatalf("status = %d, want %d", rec.Code, http.StatusForbidden)
|
||||
}
|
||||
|
||||
spans := exp.GetSpans()
|
||||
if len(spans) == 0 {
|
||||
t.Fatal("expected a span for denied call")
|
||||
}
|
||||
|
||||
span := spans[0]
|
||||
assertAttr(t, span.Attributes, AttrAccountID, "user-1")
|
||||
assertAttrBool(t, span.Attributes, AttrAuthAllowed, false)
|
||||
assertAttr(t, span.Attributes, AttrAuthDeniedReason, "insufficient scopes")
|
||||
|
||||
if span.Status.Code != codes.Error {
|
||||
t.Errorf("span status = %v, want Error", span.Status.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOTel_SpanAttributes_AuthAllowed(t *testing.T) {
|
||||
exp, tp := newTestTP()
|
||||
|
||||
ma := &mockAuth{
|
||||
accounts: map[string]*auth.Account{
|
||||
"tok": {ID: "user-1", Scopes: []string{"blog:write"}},
|
||||
},
|
||||
}
|
||||
|
||||
s := newTestServer(Options{TraceProvider: tp, Auth: ma})
|
||||
s.tools["svc.Do"] = &Tool{
|
||||
Name: "svc.Do",
|
||||
Service: "svc",
|
||||
Endpoint: "Do",
|
||||
Scopes: []string{"blog:write"},
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"tool": "svc.Do",
|
||||
"input": map[string]interface{}{},
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp/call", bytes.NewReader(body))
|
||||
req.Header.Set("Authorization", "Bearer tok")
|
||||
rec := httptest.NewRecorder()
|
||||
s.handleCallTool(rec, req)
|
||||
|
||||
// RPC will fail but auth should pass
|
||||
spans := exp.GetSpans()
|
||||
if len(spans) == 0 {
|
||||
t.Fatal("expected a span")
|
||||
}
|
||||
|
||||
span := spans[0]
|
||||
assertAttr(t, span.Attributes, AttrAccountID, "user-1")
|
||||
assertAttrBool(t, span.Attributes, AttrAuthAllowed, true)
|
||||
|
||||
// RPC fails, so span should have error status
|
||||
if span.Status.Code != codes.Error {
|
||||
t.Errorf("span status = %v, want Error (RPC fails with no backend)", span.Status.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOTel_SpanAttributes_RateLimit(t *testing.T) {
|
||||
exp, tp := newTestTP()
|
||||
|
||||
s := newTestServer(Options{
|
||||
TraceProvider: tp,
|
||||
RateLimit: &RateLimitConfig{RequestsPerSecond: 1, Burst: 1},
|
||||
})
|
||||
s.tools["svc.Do"] = &Tool{
|
||||
Name: "svc.Do",
|
||||
Service: "svc",
|
||||
Endpoint: "Do",
|
||||
}
|
||||
s.limiters["svc.Do"] = newRateLimiter(1, 1)
|
||||
|
||||
makeReq := func() int {
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"tool": "svc.Do",
|
||||
"input": map[string]interface{}{},
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp/call", bytes.NewReader(body))
|
||||
rec := httptest.NewRecorder()
|
||||
s.handleCallTool(rec, req)
|
||||
return rec.Code
|
||||
}
|
||||
|
||||
// First request passes rate limit
|
||||
makeReq()
|
||||
|
||||
// Second request should be rate limited
|
||||
code := makeReq()
|
||||
if code != http.StatusTooManyRequests {
|
||||
t.Fatalf("status = %d, want %d", code, http.StatusTooManyRequests)
|
||||
}
|
||||
|
||||
spans := exp.GetSpans()
|
||||
// Find the rate-limited span
|
||||
var found bool
|
||||
for _, span := range spans {
|
||||
for _, attr := range span.Attributes {
|
||||
if string(attr.Key) == AttrRateLimited && attr.Value.AsBool() {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("expected a span with mcp.rate_limited=true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOTel_NoProvider_NoSpan(t *testing.T) {
|
||||
// Without TraceProvider, tool calls should still work normally.
|
||||
s := newTestServer(Options{})
|
||||
s.tools["svc.Echo"] = &Tool{
|
||||
Name: "svc.Echo",
|
||||
Service: "svc",
|
||||
Endpoint: "Echo",
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"tool": "svc.Echo",
|
||||
"input": map[string]interface{}{},
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp/call", bytes.NewReader(body))
|
||||
rec := httptest.NewRecorder()
|
||||
s.handleCallTool(rec, req)
|
||||
|
||||
// Should not panic or error due to missing provider.
|
||||
// RPC fails as usual.
|
||||
if rec.Code == 0 {
|
||||
t.Error("expected a response code")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOTel_TraceContextPropagation(t *testing.T) {
|
||||
exp, tp := newTestTP()
|
||||
|
||||
s := newTestServer(Options{TraceProvider: tp})
|
||||
s.tools["svc.Echo"] = &Tool{
|
||||
Name: "svc.Echo",
|
||||
Service: "svc",
|
||||
Endpoint: "Echo",
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"tool": "svc.Echo",
|
||||
"input": map[string]interface{}{},
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp/call", bytes.NewReader(body))
|
||||
rec := httptest.NewRecorder()
|
||||
s.handleCallTool(rec, req)
|
||||
|
||||
spans := exp.GetSpans()
|
||||
if len(spans) == 0 {
|
||||
t.Fatal("expected a span")
|
||||
}
|
||||
|
||||
// The span should have a valid trace ID (non-zero)
|
||||
span := spans[0]
|
||||
if !span.SpanContext.TraceID().IsValid() {
|
||||
t.Error("expected a valid OTel trace ID")
|
||||
}
|
||||
if !span.SpanContext.SpanID().IsValid() {
|
||||
t.Error("expected a valid OTel span ID")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOTel_MissingToken(t *testing.T) {
|
||||
exp, tp := newTestTP()
|
||||
|
||||
ma := &mockAuth{
|
||||
accounts: map[string]*auth.Account{},
|
||||
}
|
||||
|
||||
s := newTestServer(Options{TraceProvider: tp, Auth: ma})
|
||||
s.tools["svc.Do"] = &Tool{
|
||||
Name: "svc.Do",
|
||||
Service: "svc",
|
||||
Endpoint: "Do",
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"tool": "svc.Do",
|
||||
"input": map[string]interface{}{},
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp/call", bytes.NewReader(body))
|
||||
// No Authorization header
|
||||
rec := httptest.NewRecorder()
|
||||
s.handleCallTool(rec, req)
|
||||
|
||||
if rec.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("status = %d, want %d", rec.Code, http.StatusUnauthorized)
|
||||
}
|
||||
|
||||
spans := exp.GetSpans()
|
||||
if len(spans) == 0 {
|
||||
t.Fatal("expected a span even for missing token")
|
||||
}
|
||||
|
||||
span := spans[0]
|
||||
assertAttrBool(t, span.Attributes, AttrAuthAllowed, false)
|
||||
assertAttr(t, span.Attributes, AttrAuthDeniedReason, "missing token")
|
||||
}
|
||||
|
||||
func TestOTel_StartToolSpan_NilProvider(t *testing.T) {
|
||||
s := newTestServer(Options{})
|
||||
ctx, span := s.startToolSpan(context.Background(), "svc.Test", "http", "test-trace-id")
|
||||
defer span.End()
|
||||
|
||||
// Should return a noop span, not panic
|
||||
if ctx == nil {
|
||||
t.Error("expected non-nil context")
|
||||
}
|
||||
if span == nil {
|
||||
t.Error("expected non-nil span (even if noop)")
|
||||
}
|
||||
}
|
||||
|
||||
// --- helpers ---
|
||||
|
||||
func assertAttr(t *testing.T, attrs []attribute.KeyValue, key, want string) {
|
||||
t.Helper()
|
||||
for _, attr := range attrs {
|
||||
if string(attr.Key) == key {
|
||||
if got := attr.Value.AsString(); got != want {
|
||||
t.Errorf("attribute %s = %q, want %q", key, got, want)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Errorf("attribute %s not found", key)
|
||||
}
|
||||
|
||||
func assertAttrBool(t *testing.T, attrs []attribute.KeyValue, key string, want bool) {
|
||||
t.Helper()
|
||||
for _, attr := range attrs {
|
||||
if string(attr.Key) == key {
|
||||
if got := attr.Value.AsBool(); got != want {
|
||||
t.Errorf("attribute %s = %v, want %v", key, got, want)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Errorf("attribute %s not found", key)
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"reflect"
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ToolDescription represents enhanced documentation for an MCP tool
|
||||
type ToolDescription struct {
|
||||
Summary string
|
||||
Description string
|
||||
Params []ParamDoc
|
||||
Returns []ReturnDoc
|
||||
Examples []string
|
||||
}
|
||||
|
||||
// ParamDoc describes a parameter
|
||||
type ParamDoc struct {
|
||||
Name string
|
||||
Type string
|
||||
Description string
|
||||
Required bool
|
||||
}
|
||||
|
||||
// ReturnDoc describes a return value
|
||||
type ReturnDoc struct {
|
||||
Type string
|
||||
Description string
|
||||
}
|
||||
|
||||
var (
|
||||
// Regex patterns for JSDoc-style tags
|
||||
paramPattern = regexp.MustCompile(`@param\s+(\w+)\s+\{(\w+)\}\s+(.+)`)
|
||||
returnPattern = regexp.MustCompile(`@return\s+\{(\w+)\}\s+(.+)`)
|
||||
examplePattern = regexp.MustCompile(`@example\s+([\s\S]+?)(?:@\w+|$)`)
|
||||
)
|
||||
|
||||
// formatFieldDescription creates a basic description for a field
|
||||
func formatFieldDescription(name, typeName string) string {
|
||||
// Convert camelCase/PascalCase to readable format
|
||||
readable := toReadable(name)
|
||||
return fmt.Sprintf("%s (%s)", readable, typeName)
|
||||
}
|
||||
|
||||
// toReadable converts camelCase or PascalCase to readable format
|
||||
func toReadable(s string) string {
|
||||
// Insert spaces before uppercase letters
|
||||
var result strings.Builder
|
||||
for i, r := range s {
|
||||
if i > 0 && r >= 'A' && r <= 'Z' {
|
||||
result.WriteRune(' ')
|
||||
}
|
||||
result.WriteRune(r)
|
||||
}
|
||||
return result.String()
|
||||
}
|
||||
|
||||
// ParseGoDocComment parses a Go doc comment for JSDoc-style tags
|
||||
func ParseGoDocComment(comment string) *ToolDescription {
|
||||
desc := &ToolDescription{
|
||||
Params: []ParamDoc{},
|
||||
Returns: []ReturnDoc{},
|
||||
Examples: []string{},
|
||||
}
|
||||
|
||||
// Extract summary (first line)
|
||||
lines := strings.Split(comment, "\n")
|
||||
if len(lines) > 0 {
|
||||
desc.Summary = strings.TrimSpace(lines[0])
|
||||
}
|
||||
|
||||
// Extract full description (before first tag)
|
||||
tagStart := strings.Index(comment, "@")
|
||||
if tagStart > 0 {
|
||||
desc.Description = strings.TrimSpace(comment[:tagStart])
|
||||
} else {
|
||||
desc.Description = strings.TrimSpace(comment)
|
||||
}
|
||||
|
||||
// Parse @param tags
|
||||
paramMatches := paramPattern.FindAllStringSubmatch(comment, -1)
|
||||
for _, match := range paramMatches {
|
||||
if len(match) == 4 {
|
||||
desc.Params = append(desc.Params, ParamDoc{
|
||||
Name: match[1],
|
||||
Type: match[2],
|
||||
Description: strings.TrimSpace(match[3]),
|
||||
Required: true,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Parse @return tags
|
||||
returnMatches := returnPattern.FindAllStringSubmatch(comment, -1)
|
||||
for _, match := range returnMatches {
|
||||
if len(match) == 3 {
|
||||
desc.Returns = append(desc.Returns, ReturnDoc{
|
||||
Type: match[1],
|
||||
Description: strings.TrimSpace(match[2]),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Parse @example tags
|
||||
exampleMatches := examplePattern.FindAllStringSubmatch(comment, -1)
|
||||
for _, match := range exampleMatches {
|
||||
if len(match) == 2 {
|
||||
example := strings.TrimSpace(match[1])
|
||||
desc.Examples = append(desc.Examples, example)
|
||||
}
|
||||
}
|
||||
|
||||
return desc
|
||||
}
|
||||
|
||||
// ParseStructTags extracts JSON schema information from struct tags
|
||||
// This can be used to enhance parameter descriptions
|
||||
func ParseStructTags(t reflect.Type) map[string]interface{} {
|
||||
schema := map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": make(map[string]interface{}),
|
||||
}
|
||||
|
||||
properties := schema["properties"].(map[string]interface{})
|
||||
required := []string{}
|
||||
|
||||
for i := 0; i < t.NumField(); i++ {
|
||||
field := t.Field(i)
|
||||
|
||||
// Get JSON tag
|
||||
jsonTag := field.Tag.Get("json")
|
||||
if jsonTag == "" || jsonTag == "-" {
|
||||
continue
|
||||
}
|
||||
|
||||
// Parse JSON tag
|
||||
jsonName := strings.Split(jsonTag, ",")[0]
|
||||
omitempty := strings.Contains(jsonTag, "omitempty")
|
||||
|
||||
// Get description from validate tag or description tag
|
||||
description := field.Tag.Get("description")
|
||||
if description == "" {
|
||||
description = formatFieldDescription(field.Name, field.Type.String())
|
||||
}
|
||||
|
||||
// Build property schema
|
||||
propSchema := map[string]interface{}{
|
||||
"description": description,
|
||||
}
|
||||
|
||||
// Add type information
|
||||
propSchema["type"] = reflectTypeToJSONType(field.Type)
|
||||
|
||||
properties[jsonName] = propSchema
|
||||
|
||||
// Track required fields
|
||||
if !omitempty {
|
||||
required = append(required, jsonName)
|
||||
}
|
||||
}
|
||||
|
||||
if len(required) > 0 {
|
||||
schema["required"] = required
|
||||
}
|
||||
|
||||
return schema
|
||||
}
|
||||
|
||||
// reflectTypeToJSONType converts Go reflect.Type to JSON schema type
|
||||
func reflectTypeToJSONType(t reflect.Type) string {
|
||||
switch t.Kind() {
|
||||
case reflect.String:
|
||||
return "string"
|
||||
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
|
||||
reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
|
||||
return "integer"
|
||||
case reflect.Float32, reflect.Float64:
|
||||
return "number"
|
||||
case reflect.Bool:
|
||||
return "boolean"
|
||||
case reflect.Slice, reflect.Array:
|
||||
return "array"
|
||||
case reflect.Map, reflect.Struct:
|
||||
return "object"
|
||||
default:
|
||||
return "string"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"go-micro.dev/v6/wrapper/x402"
|
||||
)
|
||||
|
||||
// With payments enabled, /mcp/tools advertises each priced tool's payment
|
||||
// requirements so the catalog is shoppable; free tools carry no payment.
|
||||
func TestListToolsAdvertisesPayment(t *testing.T) {
|
||||
s := newTestServer(Options{
|
||||
Payment: &x402.Config{
|
||||
PayTo: "0xabc",
|
||||
Network: "solana",
|
||||
Asset: "USDC",
|
||||
Amount: "0", // free by default
|
||||
Amounts: map[string]string{"weather.Weather.Forecast": "10000"},
|
||||
},
|
||||
})
|
||||
s.tools["weather.Weather.Forecast"] = &Tool{Name: "weather.Weather.Forecast", Description: "forecast"}
|
||||
s.tools["time.Time.Now"] = &Tool{Name: "time.Time.Now", Description: "now"}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
s.handleListTools(rec, httptest.NewRequest(http.MethodGet, "/mcp/tools", nil))
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200", rec.Code)
|
||||
}
|
||||
var out struct {
|
||||
Tools []struct {
|
||||
Name string `json:"name"`
|
||||
Payment *PaymentInfo `json:"payment"`
|
||||
} `json:"tools"`
|
||||
}
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
|
||||
byName := map[string]*PaymentInfo{}
|
||||
for _, tl := range out.Tools {
|
||||
byName[tl.Name] = tl.Payment
|
||||
}
|
||||
|
||||
paid := byName["weather.Weather.Forecast"]
|
||||
if paid == nil {
|
||||
t.Fatal("priced tool should advertise payment in the catalog")
|
||||
}
|
||||
if paid.Amount != "10000" || paid.Network != "solana" || paid.PayTo != "0xabc" || paid.Asset != "USDC" {
|
||||
t.Errorf("payment info wrong: %+v", paid)
|
||||
}
|
||||
if byName["time.Time.Now"] != nil {
|
||||
t.Error("free tool should not advertise payment")
|
||||
}
|
||||
}
|
||||
|
||||
// Without payments configured, the catalog carries no payment info.
|
||||
func TestListToolsNoPaymentWhenDisabled(t *testing.T) {
|
||||
s := newTestServer(Options{})
|
||||
s.tools["a.A.B"] = &Tool{Name: "a.A.B"}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
s.handleListTools(rec, httptest.NewRequest(http.MethodGet, "/mcp/tools", nil))
|
||||
|
||||
var out struct {
|
||||
Tools []struct {
|
||||
Payment *PaymentInfo `json:"payment"`
|
||||
} `json:"tools"`
|
||||
}
|
||||
json.Unmarshal(rec.Body.Bytes(), &out)
|
||||
if len(out.Tools) != 1 || out.Tools[0].Payment != nil {
|
||||
t.Errorf("expected no payment info when payments disabled, got %+v", out.Tools)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// rateLimiter implements a simple token-bucket rate limiter.
|
||||
type rateLimiter struct {
|
||||
mu sync.Mutex
|
||||
rate float64 // tokens per second
|
||||
burst int // max tokens
|
||||
tokens float64 // current token count
|
||||
lastTime time.Time // last refill time
|
||||
}
|
||||
|
||||
// newRateLimiter creates a rate limiter that allows rate requests/sec with
|
||||
// the given burst size. If burst is less than 1 it defaults to 1.
|
||||
func newRateLimiter(rate float64, burst int) *rateLimiter {
|
||||
if burst < 1 {
|
||||
burst = 1
|
||||
}
|
||||
return &rateLimiter{
|
||||
rate: rate,
|
||||
burst: burst,
|
||||
tokens: float64(burst),
|
||||
lastTime: time.Now(),
|
||||
}
|
||||
}
|
||||
|
||||
// Allow reports whether a single event may happen now.
|
||||
func (r *rateLimiter) Allow() bool {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
elapsed := now.Sub(r.lastTime).Seconds()
|
||||
r.lastTime = now
|
||||
|
||||
// Refill tokens based on elapsed time
|
||||
r.tokens += elapsed * r.rate
|
||||
if r.tokens > float64(r.burst) {
|
||||
r.tokens = float64(r.burst)
|
||||
}
|
||||
|
||||
if r.tokens < 1 {
|
||||
return false
|
||||
}
|
||||
r.tokens--
|
||||
return true
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
"go-micro.dev/v6/ai"
|
||||
"go-micro.dev/v6/client"
|
||||
"go-micro.dev/v6/registry"
|
||||
)
|
||||
|
||||
// CallResult is the outcome of a successful tool dispatch. A tool that ran but
|
||||
// produced an error sets IsError — per the MCP spec this is returned as a
|
||||
// tools/call result with isError:true, not a JSON-RPC protocol error.
|
||||
type CallResult struct {
|
||||
Text string
|
||||
IsError bool
|
||||
}
|
||||
|
||||
// Error lets the package's RPCError (see stdio.go) be returned by a resolver
|
||||
// to signal a protocol/pre-check failure with a specific JSON-RPC code; the
|
||||
// handler maps it straight to the JSON-RPC error.
|
||||
func (e *RPCError) Error() string { return e.Message }
|
||||
|
||||
// ToolFunc executes a manually-registered tool. Return a *CallResult for tool
|
||||
// outcomes (set IsError for tool-level failures); return a non-nil error — an
|
||||
// *RPCError for a specific code — for protocol/pre-check failures.
|
||||
type ToolFunc func(ctx context.Context, args map[string]any) (*CallResult, error)
|
||||
|
||||
// Resolver supplies the gateway's tools and executes calls. Swapping the
|
||||
// resolver changes where tools come from without touching the MCP protocol or
|
||||
// transport:
|
||||
//
|
||||
// - NewManualResolver: tools you register explicitly (full product control,
|
||||
// including tools that are not go-micro services, executed via your own
|
||||
// logic — auth, metering, …).
|
||||
// - NewRegistryResolver: tools auto-discovered from registered services.
|
||||
//
|
||||
// The built-in store/broker tools are intentionally NOT exposed by any
|
||||
// resolver — they remain a development convenience on the legacy Serve() path.
|
||||
type Resolver interface {
|
||||
// List returns the current tool catalog.
|
||||
List(ctx context.Context) ([]Tool, error)
|
||||
// Call executes a tool by name with JSON arguments.
|
||||
Call(ctx context.Context, name string, args map[string]any) (*CallResult, error)
|
||||
}
|
||||
|
||||
// ManualResolver exposes an explicitly-registered set of tools.
|
||||
type ManualResolver struct {
|
||||
mu sync.RWMutex
|
||||
order []Tool
|
||||
funcs map[string]ToolFunc
|
||||
}
|
||||
|
||||
// NewManualResolver returns an empty manual resolver.
|
||||
func NewManualResolver() *ManualResolver {
|
||||
return &ManualResolver{funcs: map[string]ToolFunc{}}
|
||||
}
|
||||
|
||||
// Add registers (or replaces) a tool and its handler. Returns the resolver for
|
||||
// chaining.
|
||||
func (m *ManualResolver) Add(t Tool, fn ToolFunc) *ManualResolver {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if _, ok := m.funcs[t.Name]; ok {
|
||||
for i := range m.order {
|
||||
if m.order[i].Name == t.Name {
|
||||
m.order[i] = t
|
||||
}
|
||||
}
|
||||
} else {
|
||||
m.order = append(m.order, t)
|
||||
}
|
||||
m.funcs[t.Name] = fn
|
||||
return m
|
||||
}
|
||||
|
||||
// List returns the registered tools.
|
||||
func (m *ManualResolver) List(_ context.Context) ([]Tool, error) {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
out := make([]Tool, len(m.order))
|
||||
copy(out, m.order)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Call runs the handler registered for name.
|
||||
func (m *ManualResolver) Call(ctx context.Context, name string, args map[string]any) (*CallResult, error) {
|
||||
m.mu.RLock()
|
||||
fn, ok := m.funcs[name]
|
||||
m.mu.RUnlock()
|
||||
if !ok {
|
||||
return nil, &RPCError{Code: InvalidParams, Message: "Tool not found: " + name, Data: name}
|
||||
}
|
||||
return fn(ctx, args)
|
||||
}
|
||||
|
||||
// RegistryResolver auto-discovers tools from registered go-micro services and
|
||||
// executes them over RPC. It exposes only services — never the internal
|
||||
// store/broker tools.
|
||||
type RegistryResolver struct {
|
||||
tools *ai.Tools
|
||||
}
|
||||
|
||||
// NewRegistryResolver discovers services from reg and calls them with cl.
|
||||
func NewRegistryResolver(reg registry.Registry, cl client.Client) *RegistryResolver {
|
||||
return &RegistryResolver{tools: ai.NewTools(reg, ai.ToolClient(cl))}
|
||||
}
|
||||
|
||||
// List discovers the current service tools.
|
||||
func (r *RegistryResolver) List(_ context.Context) ([]Tool, error) {
|
||||
discovered, err := r.tools.Discover()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]Tool, 0, len(discovered))
|
||||
for _, t := range discovered {
|
||||
out = append(out, Tool{
|
||||
Name: t.Name,
|
||||
Description: t.Description,
|
||||
InputSchema: map[string]interface{}{"type": "object", "properties": t.Properties},
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Call executes a discovered service tool.
|
||||
func (r *RegistryResolver) Call(ctx context.Context, name string, args map[string]any) (*CallResult, error) {
|
||||
res := r.tools.Handler()(ctx, ai.ToolCall{ID: "1", Name: name, Input: args})
|
||||
return &CallResult{Text: res.Content}, nil
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestManualResolverHandler(t *testing.T) {
|
||||
res := NewManualResolver().
|
||||
Add(Tool{Name: "echo", Description: "echoes text"},
|
||||
func(_ context.Context, args map[string]interface{}) (*CallResult, error) {
|
||||
s, _ := args["text"].(string)
|
||||
return &CallResult{Text: "you said: " + s}, nil
|
||||
}).
|
||||
Add(Tool{Name: "boom", Description: "errors"},
|
||||
func(_ context.Context, _ map[string]interface{}) (*CallResult, error) {
|
||||
return &CallResult{Text: "kaboom", IsError: true}, nil
|
||||
}).
|
||||
Add(Tool{Name: "blocked", Description: "coded error"},
|
||||
func(_ context.Context, _ map[string]interface{}) (*CallResult, error) {
|
||||
return nil, &RPCError{Code: -32000, Message: "insufficient credits"}
|
||||
})
|
||||
|
||||
ts := httptest.NewServer(NewHandler(res))
|
||||
defer ts.Close()
|
||||
rpc := func(body string) (int, map[string]interface{}) {
|
||||
resp, err := http.Post(ts.URL, "application/json", strings.NewReader(body))
|
||||
if err != nil {
|
||||
t.Fatalf("post rpc: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
var out map[string]interface{}
|
||||
json.NewDecoder(resp.Body).Decode(&out)
|
||||
return resp.StatusCode, out
|
||||
}
|
||||
|
||||
if _, out := rpc(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`); len(out["result"].(map[string]interface{})["tools"].([]interface{})) != 3 {
|
||||
t.Fatalf("tools/list: %v", out)
|
||||
}
|
||||
// tool result
|
||||
_, out := rpc(`{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"echo","arguments":{"text":"hi"}}}`)
|
||||
if out["result"].(map[string]interface{})["content"].([]interface{})[0].(map[string]interface{})["text"] != "you said: hi" {
|
||||
t.Fatalf("echo: %v", out)
|
||||
}
|
||||
// tool-level error -> isError result, NOT protocol error
|
||||
_, out = rpc(`{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"boom","arguments":{}}}`)
|
||||
if out["error"] != nil || out["result"].(map[string]interface{})["isError"] != true {
|
||||
t.Fatalf("boom should be isError result: %v", out)
|
||||
}
|
||||
// coded protocol error -> JSON-RPC error with the code
|
||||
_, out = rpc(`{"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"name":"blocked","arguments":{}}}`)
|
||||
if out["error"] == nil || int(out["error"].(map[string]interface{})["code"].(float64)) != -32000 {
|
||||
t.Fatalf("blocked should be -32000: %v", out)
|
||||
}
|
||||
// notification -> 204, no body
|
||||
code, _ := rpc(`{"jsonrpc":"2.0","method":"notifications/initialized"}`)
|
||||
if code != http.StatusNoContent {
|
||||
t.Fatalf("notification status = %d, want 204", code)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,376 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"go-micro.dev/v6/auth"
|
||||
"go-micro.dev/v6/metadata"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
)
|
||||
|
||||
// StdioTransport implements MCP JSON-RPC 2.0 over stdio
|
||||
// This is used by Claude Code and other local AI tools
|
||||
type StdioTransport struct {
|
||||
server *Server
|
||||
reader *bufio.Reader
|
||||
writer *bufio.Writer
|
||||
writerMu sync.Mutex
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
}
|
||||
|
||||
// JSONRPCRequest represents a JSON-RPC 2.0 request
|
||||
type JSONRPCRequest struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID interface{} `json:"id,omitempty"`
|
||||
Method string `json:"method"`
|
||||
Params json.RawMessage `json:"params,omitempty"`
|
||||
}
|
||||
|
||||
// JSONRPCResponse represents a JSON-RPC 2.0 response
|
||||
type JSONRPCResponse struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID interface{} `json:"id,omitempty"`
|
||||
Result interface{} `json:"result,omitempty"`
|
||||
Error *RPCError `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// RPCError represents a JSON-RPC error
|
||||
type RPCError struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Data interface{} `json:"data,omitempty"`
|
||||
}
|
||||
|
||||
// Standard JSON-RPC error codes
|
||||
const (
|
||||
ParseError = -32700
|
||||
InvalidRequest = -32600
|
||||
MethodNotFound = -32601
|
||||
InvalidParams = -32602
|
||||
InternalError = -32603
|
||||
)
|
||||
|
||||
// NewStdioTransport creates a new stdio transport for the MCP server
|
||||
func NewStdioTransport(server *Server) *StdioTransport {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
return &StdioTransport{
|
||||
server: server,
|
||||
reader: bufio.NewReader(os.Stdin),
|
||||
writer: bufio.NewWriter(os.Stdout),
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
}
|
||||
}
|
||||
|
||||
// Serve starts the stdio transport and processes JSON-RPC requests
|
||||
func (t *StdioTransport) Serve() error {
|
||||
t.server.opts.Logger.Printf("[mcp] MCP server started (stdio transport)")
|
||||
|
||||
// Read and process requests from stdin
|
||||
for {
|
||||
select {
|
||||
case <-t.ctx.Done():
|
||||
return nil
|
||||
default:
|
||||
}
|
||||
|
||||
// Read one line (JSON-RPC request)
|
||||
line, err := t.reader.ReadBytes('\n')
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("failed to read request: %w", err)
|
||||
}
|
||||
|
||||
// Parse JSON-RPC request
|
||||
var req JSONRPCRequest
|
||||
if err := json.Unmarshal(line, &req); err != nil {
|
||||
t.sendError(nil, ParseError, "Parse error", err.Error())
|
||||
continue
|
||||
}
|
||||
|
||||
// Validate JSON-RPC version
|
||||
if req.JSONRPC != "2.0" {
|
||||
t.sendError(req.ID, InvalidRequest, "Invalid request", "jsonrpc must be '2.0'")
|
||||
continue
|
||||
}
|
||||
|
||||
// Handle request
|
||||
go t.handleRequest(&req)
|
||||
}
|
||||
}
|
||||
|
||||
// handleRequest processes a single JSON-RPC request
|
||||
func (t *StdioTransport) handleRequest(req *JSONRPCRequest) {
|
||||
switch req.Method {
|
||||
case "initialize":
|
||||
t.handleInitialize(req)
|
||||
case "tools/list":
|
||||
t.handleToolsList(req)
|
||||
case "tools/call":
|
||||
t.handleToolsCall(req)
|
||||
default:
|
||||
t.sendError(req.ID, MethodNotFound, "Method not found", req.Method)
|
||||
}
|
||||
}
|
||||
|
||||
// handleInitialize handles the initialize request
|
||||
func (t *StdioTransport) handleInitialize(req *JSONRPCRequest) {
|
||||
result := map[string]interface{}{
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": map[string]interface{}{
|
||||
"tools": map[string]interface{}{},
|
||||
},
|
||||
"serverInfo": map[string]interface{}{
|
||||
"name": "go-micro-mcp",
|
||||
"version": "1.0.0",
|
||||
},
|
||||
}
|
||||
|
||||
t.sendResponse(req.ID, result)
|
||||
}
|
||||
|
||||
// handleToolsList handles the tools/list request
|
||||
func (t *StdioTransport) handleToolsList(req *JSONRPCRequest) {
|
||||
t.server.toolsMu.RLock()
|
||||
tools := make([]interface{}, 0, len(t.server.tools))
|
||||
for _, tool := range t.server.tools {
|
||||
tools = append(tools, map[string]interface{}{
|
||||
"name": tool.Name,
|
||||
"description": tool.Description,
|
||||
"inputSchema": tool.InputSchema,
|
||||
})
|
||||
}
|
||||
t.server.toolsMu.RUnlock()
|
||||
|
||||
result := map[string]interface{}{
|
||||
"tools": tools,
|
||||
}
|
||||
|
||||
t.sendResponse(req.ID, result)
|
||||
}
|
||||
|
||||
// handleToolsCall handles the tools/call request
|
||||
func (t *StdioTransport) handleToolsCall(req *JSONRPCRequest) {
|
||||
// Parse params
|
||||
var params struct {
|
||||
Name string `json:"name"`
|
||||
Arguments map[string]interface{} `json:"arguments"`
|
||||
// Token allows callers to pass a bearer token for auth via the
|
||||
// JSON-RPC params (since stdio has no HTTP headers).
|
||||
Token string `json:"_token,omitempty"`
|
||||
}
|
||||
if err := json.Unmarshal(req.Params, ¶ms); err != nil {
|
||||
t.sendError(req.ID, InvalidParams, "Invalid params", err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// Get tool info
|
||||
t.server.toolsMu.RLock()
|
||||
tool, exists := t.server.tools[params.Name]
|
||||
t.server.toolsMu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
t.sendError(req.ID, InvalidParams, "Tool not found", params.Name)
|
||||
return
|
||||
}
|
||||
|
||||
// Generate trace ID
|
||||
traceID := uuid.New().String()
|
||||
|
||||
// Start OTel span (noop if TraceProvider is nil)
|
||||
ctx, span := t.server.startToolSpan(t.ctx, params.Name, "stdio", traceID)
|
||||
defer span.End()
|
||||
|
||||
// Authenticate and authorize (if Auth is configured)
|
||||
var account *auth.Account
|
||||
if t.server.opts.Auth != nil {
|
||||
token := params.Token
|
||||
if token == "" {
|
||||
span.SetAttributes(attribute.Bool(AttrAuthAllowed, false), attribute.String(AttrAuthDeniedReason, "missing token"))
|
||||
setSpanError(span, fmt.Errorf("missing token"))
|
||||
t.server.audit(AuditRecord{TraceID: traceID, Timestamp: time.Now(), Tool: params.Name, Allowed: false, DeniedReason: "missing token"})
|
||||
t.sendError(req.ID, InvalidParams, "Unauthorized", "missing _token in params")
|
||||
return
|
||||
}
|
||||
token = strings.TrimPrefix(token, "Bearer ")
|
||||
acc, err := t.server.opts.Auth.Inspect(token)
|
||||
if err != nil {
|
||||
span.SetAttributes(attribute.Bool(AttrAuthAllowed, false), attribute.String(AttrAuthDeniedReason, "invalid token"))
|
||||
setSpanError(span, fmt.Errorf("invalid token"))
|
||||
t.server.audit(AuditRecord{TraceID: traceID, Timestamp: time.Now(), Tool: params.Name, Allowed: false, DeniedReason: "invalid token"})
|
||||
t.sendError(req.ID, InvalidParams, "Unauthorized", "invalid token")
|
||||
return
|
||||
}
|
||||
account = acc
|
||||
span.SetAttributes(attribute.String(AttrAccountID, account.ID))
|
||||
|
||||
// Check per-tool scopes
|
||||
if len(tool.Scopes) > 0 {
|
||||
span.SetAttributes(attribute.StringSlice(AttrScopesRequired, tool.Scopes))
|
||||
if !hasScope(account.Scopes, tool.Scopes) {
|
||||
span.SetAttributes(attribute.Bool(AttrAuthAllowed, false), attribute.String(AttrAuthDeniedReason, "insufficient scopes"))
|
||||
setSpanError(span, fmt.Errorf("insufficient scopes"))
|
||||
t.server.audit(AuditRecord{
|
||||
TraceID: traceID, Timestamp: time.Now(), Tool: params.Name,
|
||||
AccountID: account.ID, ScopesRequired: tool.Scopes,
|
||||
Allowed: false, DeniedReason: "insufficient scopes",
|
||||
})
|
||||
t.sendError(req.ID, InvalidParams, "Forbidden", "insufficient scopes")
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Rate limit check
|
||||
if err := t.server.allowRate(params.Name); err != nil {
|
||||
span.SetAttributes(attribute.Bool(AttrRateLimited, true))
|
||||
setSpanError(span, err)
|
||||
accountID := ""
|
||||
if account != nil {
|
||||
accountID = account.ID
|
||||
}
|
||||
t.server.audit(AuditRecord{
|
||||
TraceID: traceID, Timestamp: time.Now(), Tool: params.Name,
|
||||
AccountID: accountID, Allowed: false, DeniedReason: "rate limited",
|
||||
})
|
||||
t.sendError(req.ID, InternalError, "Rate limit exceeded", params.Name)
|
||||
return
|
||||
}
|
||||
|
||||
span.SetAttributes(attribute.Bool(AttrAuthAllowed, true))
|
||||
|
||||
// Convert arguments to JSON bytes for RPC call
|
||||
inputBytes, err := json.Marshal(params.Arguments)
|
||||
if err != nil {
|
||||
t.sendError(req.ID, InternalError, "Failed to marshal arguments", err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// Build context with tracing metadata
|
||||
// OTel trace context was already injected by startToolSpan; add MCP metadata.
|
||||
md, _ := metadata.FromContext(ctx)
|
||||
if md == nil {
|
||||
md = make(metadata.Metadata)
|
||||
}
|
||||
md.Set(TraceIDKey, traceID)
|
||||
md.Set(ToolNameKey, params.Name)
|
||||
if account != nil {
|
||||
md.Set(AccountIDKey, account.ID)
|
||||
}
|
||||
ctx = metadata.NewContext(ctx, md)
|
||||
|
||||
// Make RPC call
|
||||
start := time.Now()
|
||||
rpcReq := t.server.opts.Client.NewRequest(tool.Service, tool.Endpoint, &struct {
|
||||
Data []byte
|
||||
}{Data: inputBytes})
|
||||
|
||||
var rsp struct {
|
||||
Data []byte
|
||||
}
|
||||
|
||||
if err := t.server.opts.Client.Call(ctx, rpcReq, &rsp); err != nil {
|
||||
setSpanError(span, err)
|
||||
accountID := ""
|
||||
if account != nil {
|
||||
accountID = account.ID
|
||||
}
|
||||
t.server.audit(AuditRecord{
|
||||
TraceID: traceID, Timestamp: time.Now(), Tool: params.Name,
|
||||
AccountID: accountID, ScopesRequired: tool.Scopes,
|
||||
Allowed: true, Duration: time.Since(start), Error: err.Error(),
|
||||
})
|
||||
// A tool-execution failure is reported as an isError result, not a
|
||||
// JSON-RPC protocol error (per the MCP spec), so the agent can read it.
|
||||
t.sendResponse(req.ID, mcpToolError(traceID, "tool call failed: "+err.Error()))
|
||||
return
|
||||
}
|
||||
|
||||
setSpanOK(span)
|
||||
|
||||
// Audit successful call
|
||||
accountID := ""
|
||||
if account != nil {
|
||||
accountID = account.ID
|
||||
}
|
||||
t.server.audit(AuditRecord{
|
||||
TraceID: traceID, Timestamp: time.Now(), Tool: params.Name,
|
||||
AccountID: accountID, ScopesRequired: tool.Scopes,
|
||||
Allowed: true, Duration: time.Since(start),
|
||||
})
|
||||
|
||||
// The downstream response is JSON — return it as JSON text, not %v.
|
||||
t.sendResponse(req.ID, mcpToolResult(traceID, rsp.Data))
|
||||
}
|
||||
|
||||
// sendResponse sends a JSON-RPC response
|
||||
func (t *StdioTransport) sendResponse(id interface{}, result interface{}) {
|
||||
resp := JSONRPCResponse{
|
||||
JSONRPC: "2.0",
|
||||
ID: id,
|
||||
Result: result,
|
||||
}
|
||||
|
||||
t.writeJSON(resp)
|
||||
}
|
||||
|
||||
// sendError sends a JSON-RPC error response
|
||||
func (t *StdioTransport) sendError(id interface{}, code int, message string, data interface{}) {
|
||||
resp := JSONRPCResponse{
|
||||
JSONRPC: "2.0",
|
||||
ID: id,
|
||||
Error: &RPCError{
|
||||
Code: code,
|
||||
Message: message,
|
||||
Data: data,
|
||||
},
|
||||
}
|
||||
|
||||
t.writeJSON(resp)
|
||||
}
|
||||
|
||||
// writeJSON writes a JSON-RPC message to stdout
|
||||
func (t *StdioTransport) writeJSON(v interface{}) {
|
||||
t.writerMu.Lock()
|
||||
defer t.writerMu.Unlock()
|
||||
|
||||
data, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
log.Printf("[mcp] Failed to marshal response: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if _, err := t.writer.Write(data); err != nil {
|
||||
log.Printf("[mcp] Failed to write response: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if _, err := t.writer.Write([]byte("\n")); err != nil {
|
||||
log.Printf("[mcp] Failed to write newline: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if err := t.writer.Flush(); err != nil {
|
||||
log.Printf("[mcp] Failed to flush writer: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Stop gracefully stops the stdio transport
|
||||
func (t *StdioTransport) Stop() error {
|
||||
t.cancel()
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"go-micro.dev/v6/client"
|
||||
)
|
||||
|
||||
// fakeCallClient overrides Call to return canned data or an error; NewRequest
|
||||
// and the rest are promoted from the embedded real client.
|
||||
type fakeCallClient struct {
|
||||
client.Client
|
||||
data []byte
|
||||
err error
|
||||
}
|
||||
|
||||
func (f *fakeCallClient) Call(ctx context.Context, req client.Request, rsp interface{}, opts ...client.CallOption) error {
|
||||
if f.err != nil {
|
||||
return f.err
|
||||
}
|
||||
if r, ok := rsp.(*struct{ Data []byte }); ok {
|
||||
r.Data = f.data
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// isToolError reports whether an MCP tools/call result carries isError:true.
|
||||
func isToolError(result interface{}) bool {
|
||||
m, ok := result.(map[string]interface{})
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
b, _ := m["isError"].(bool)
|
||||
return b
|
||||
}
|
||||
|
||||
// toolResultText extracts the first text content of an MCP tools/call result.
|
||||
func toolResultText(t *testing.T, result interface{}) string {
|
||||
t.Helper()
|
||||
m, ok := result.(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("result is not a map: %#v", result)
|
||||
}
|
||||
content, ok := m["content"].([]interface{})
|
||||
if !ok || len(content) == 0 {
|
||||
t.Fatalf("result has no content: %#v", result)
|
||||
}
|
||||
first, _ := content[0].(map[string]interface{})
|
||||
text, _ := first["text"].(string)
|
||||
return text
|
||||
}
|
||||
|
||||
// driveStdio sends one JSON-RPC request through a StdioTransport and returns the
|
||||
// decoded response, capturing the transport's stdout into a buffer.
|
||||
func driveStdio(t *testing.T, s *Server, method string, id interface{}, params interface{}) JSONRPCResponse {
|
||||
t.Helper()
|
||||
tr := NewStdioTransport(s)
|
||||
var out bytes.Buffer
|
||||
tr.writer = bufio.NewWriter(&out)
|
||||
raw, _ := json.Marshal(params)
|
||||
tr.handleRequest(&JSONRPCRequest{JSONRPC: "2.0", ID: id, Method: method, Params: raw})
|
||||
var resp JSONRPCResponse
|
||||
if err := json.Unmarshal(bytes.TrimSpace(out.Bytes()), &resp); err != nil {
|
||||
t.Fatalf("decode stdio response: %v (raw=%q)", err, out.String())
|
||||
}
|
||||
return resp
|
||||
}
|
||||
|
||||
// The stdio transport is the path an external MCP host (Claude Desktop) uses.
|
||||
// It must return tool output as JSON text, not fmt.Sprintf("%v", ...) which
|
||||
// yields Go map-syntax and is unparseable by a real client.
|
||||
func TestStdio_ToolsCall_ReturnsJSONNotGoSyntax(t *testing.T) {
|
||||
s := newTestServer(Options{})
|
||||
s.opts.Client = &fakeCallClient{Client: client.DefaultClient, data: []byte(`{"id":1,"name":"bob"}`)}
|
||||
s.tools["svc.Echo"] = &Tool{Name: "svc.Echo", Service: "svc", Endpoint: "Echo"}
|
||||
|
||||
resp := driveStdio(t, s, "tools/call", 1, map[string]interface{}{
|
||||
"name": "svc.Echo",
|
||||
"arguments": map[string]interface{}{"msg": "hi"},
|
||||
})
|
||||
if resp.Error != nil {
|
||||
t.Fatalf("unexpected protocol error: %+v", resp.Error)
|
||||
}
|
||||
text := toolResultText(t, resp.Result)
|
||||
// The bug returned Go map-syntax ("map[id:1 name:bob]"), which fails to parse.
|
||||
var got map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(text), &got); err != nil {
|
||||
t.Fatalf("tool result text is not JSON (the %%v bug): %q", text)
|
||||
}
|
||||
if got["name"] != "bob" {
|
||||
t.Errorf("result = %v, want name=bob", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A tool-execution failure must be an MCP isError result, not a JSON-RPC
|
||||
// protocol error, so the agent can read the failure.
|
||||
func TestStdio_ToolsCall_FailureIsIsErrorResult(t *testing.T) {
|
||||
s := newTestServer(Options{})
|
||||
s.opts.Client = &fakeCallClient{Client: client.DefaultClient, err: errors.New("backend down")}
|
||||
s.tools["svc.Echo"] = &Tool{Name: "svc.Echo", Service: "svc", Endpoint: "Echo"}
|
||||
|
||||
resp := driveStdio(t, s, "tools/call", 1, map[string]interface{}{
|
||||
"name": "svc.Echo",
|
||||
"arguments": map[string]interface{}{},
|
||||
})
|
||||
if resp.Error != nil {
|
||||
t.Fatalf("tool failure returned a protocol error, want isError result: %+v", resp.Error)
|
||||
}
|
||||
if !isToolError(resp.Result) {
|
||||
t.Fatalf("expected isError result, got %+v", resp.Result)
|
||||
}
|
||||
if text := toolResultText(t, resp.Result); text == "" {
|
||||
t.Error("isError result should carry the error text")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package mcp
|
||||
|
||||
// MCP tools/call result shaping, shared by the stdio and websocket JSON-RPC
|
||||
// transports. Kept in one place so both transports produce spec-shaped results.
|
||||
|
||||
// mcpToolResult builds a successful MCP tools/call result. The downstream RPC
|
||||
// response body (data) is JSON, so it is returned as JSON text — NOT
|
||||
// fmt.Sprintf("%v", ...) of a decoded value, which produces Go map-syntax
|
||||
// (map[id:1 name:bob]) instead of JSON and is what an external MCP client
|
||||
// (e.g. Claude Desktop over stdio) would otherwise receive.
|
||||
func mcpToolResult(traceID string, data []byte) map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"content": []interface{}{
|
||||
map[string]interface{}{"type": "text", "text": string(data)},
|
||||
},
|
||||
"trace_id": traceID,
|
||||
}
|
||||
}
|
||||
|
||||
// mcpToolError builds an MCP tools/call result for a tool-EXECUTION failure.
|
||||
// Per the MCP spec a tool that fails returns a normal result with isError:true
|
||||
// (the error as text content), NOT a JSON-RPC protocol error — that way the
|
||||
// agent can read the failure instead of seeing a transport-level error.
|
||||
func mcpToolError(traceID, msg string) map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"content": []interface{}{
|
||||
map[string]interface{}{"type": "text", "text": msg},
|
||||
},
|
||||
"isError": true,
|
||||
"trace_id": traceID,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,329 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"go-micro.dev/v6/auth"
|
||||
"go-micro.dev/v6/metadata"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/gorilla/websocket"
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
)
|
||||
|
||||
var upgrader = websocket.Upgrader{
|
||||
CheckOrigin: func(r *http.Request) bool { return true },
|
||||
}
|
||||
|
||||
// WebSocketTransport implements MCP JSON-RPC 2.0 over WebSocket.
|
||||
// It supports bidirectional streaming for real-time AI agents.
|
||||
type WebSocketTransport struct {
|
||||
server *Server
|
||||
}
|
||||
|
||||
// wsConn wraps a single WebSocket connection with write serialization.
|
||||
type wsConn struct {
|
||||
conn *websocket.Conn
|
||||
writeMu sync.Mutex
|
||||
server *Server
|
||||
account *auth.Account // set once during initial auth
|
||||
}
|
||||
|
||||
// NewWebSocketTransport creates a WebSocket transport for the MCP server.
|
||||
func NewWebSocketTransport(server *Server) *WebSocketTransport {
|
||||
return &WebSocketTransport{server: server}
|
||||
}
|
||||
|
||||
// ServeHTTP implements http.Handler and upgrades HTTP to WebSocket.
|
||||
func (t *WebSocketTransport) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
t.server.opts.Logger.Printf("[mcp] WebSocket upgrade failed: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Extract bearer token from the upgrade request (if present).
|
||||
var account *auth.Account
|
||||
if t.server.opts.Auth != nil {
|
||||
token := r.Header.Get("Authorization")
|
||||
token = strings.TrimPrefix(token, "Bearer ")
|
||||
// Allow connection-level auth from header. Per-message auth via
|
||||
// _token param is also supported (checked in handleToolsCall).
|
||||
if token != "" {
|
||||
acc, err := t.server.opts.Auth.Inspect(token)
|
||||
if err == nil {
|
||||
account = acc
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
wc := &wsConn{
|
||||
conn: conn,
|
||||
server: t.server,
|
||||
account: account,
|
||||
}
|
||||
|
||||
t.server.opts.Logger.Printf("[mcp] WebSocket client connected from %s", r.RemoteAddr)
|
||||
go wc.readLoop()
|
||||
}
|
||||
|
||||
// readLoop reads JSON-RPC messages from the WebSocket connection.
|
||||
func (wc *wsConn) readLoop() {
|
||||
defer wc.conn.Close()
|
||||
|
||||
for {
|
||||
_, message, err := wc.conn.ReadMessage()
|
||||
if err != nil {
|
||||
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
|
||||
wc.server.opts.Logger.Printf("[mcp] WebSocket read error: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
var req JSONRPCRequest
|
||||
if err := json.Unmarshal(message, &req); err != nil {
|
||||
wc.sendError(nil, ParseError, "Parse error", err.Error())
|
||||
continue
|
||||
}
|
||||
|
||||
if req.JSONRPC != "2.0" {
|
||||
wc.sendError(req.ID, InvalidRequest, "Invalid request", "jsonrpc must be '2.0'")
|
||||
continue
|
||||
}
|
||||
|
||||
go wc.handleRequest(&req)
|
||||
}
|
||||
}
|
||||
|
||||
// handleRequest dispatches a JSON-RPC request to the appropriate handler.
|
||||
func (wc *wsConn) handleRequest(req *JSONRPCRequest) {
|
||||
switch req.Method {
|
||||
case "initialize":
|
||||
wc.handleInitialize(req)
|
||||
case "tools/list":
|
||||
wc.handleToolsList(req)
|
||||
case "tools/call":
|
||||
wc.handleToolsCall(req)
|
||||
default:
|
||||
wc.sendError(req.ID, MethodNotFound, "Method not found", req.Method)
|
||||
}
|
||||
}
|
||||
|
||||
// handleInitialize handles the initialize request.
|
||||
func (wc *wsConn) handleInitialize(req *JSONRPCRequest) {
|
||||
result := map[string]interface{}{
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": map[string]interface{}{
|
||||
"tools": map[string]interface{}{},
|
||||
},
|
||||
"serverInfo": map[string]interface{}{
|
||||
"name": "go-micro-mcp",
|
||||
"version": "1.0.0",
|
||||
},
|
||||
}
|
||||
wc.sendResponse(req.ID, result)
|
||||
}
|
||||
|
||||
// handleToolsList handles the tools/list request.
|
||||
func (wc *wsConn) handleToolsList(req *JSONRPCRequest) {
|
||||
wc.server.toolsMu.RLock()
|
||||
tools := make([]interface{}, 0, len(wc.server.tools))
|
||||
for _, tool := range wc.server.tools {
|
||||
tools = append(tools, map[string]interface{}{
|
||||
"name": tool.Name,
|
||||
"description": tool.Description,
|
||||
"inputSchema": tool.InputSchema,
|
||||
})
|
||||
}
|
||||
wc.server.toolsMu.RUnlock()
|
||||
|
||||
wc.sendResponse(req.ID, map[string]interface{}{
|
||||
"tools": tools,
|
||||
})
|
||||
}
|
||||
|
||||
// handleToolsCall handles the tools/call request.
|
||||
func (wc *wsConn) handleToolsCall(req *JSONRPCRequest) {
|
||||
var params struct {
|
||||
Name string `json:"name"`
|
||||
Arguments map[string]interface{} `json:"arguments"`
|
||||
Token string `json:"_token,omitempty"`
|
||||
}
|
||||
if err := json.Unmarshal(req.Params, ¶ms); err != nil {
|
||||
wc.sendError(req.ID, InvalidParams, "Invalid params", err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// Get tool info
|
||||
wc.server.toolsMu.RLock()
|
||||
tool, exists := wc.server.tools[params.Name]
|
||||
wc.server.toolsMu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
wc.sendError(req.ID, InvalidParams, "Tool not found", params.Name)
|
||||
return
|
||||
}
|
||||
|
||||
traceID := uuid.New().String()
|
||||
|
||||
// Start OTel span
|
||||
ctx, span := wc.server.startToolSpan(wc.server.opts.Context, params.Name, "websocket", traceID)
|
||||
defer span.End()
|
||||
|
||||
// Resolve account: prefer connection-level auth, fall back to per-message _token.
|
||||
account := wc.account
|
||||
if wc.server.opts.Auth != nil {
|
||||
if account == nil {
|
||||
token := params.Token
|
||||
if token == "" {
|
||||
span.SetAttributes(attribute.Bool(AttrAuthAllowed, false), attribute.String(AttrAuthDeniedReason, "missing token"))
|
||||
setSpanError(span, fmt.Errorf("missing token"))
|
||||
wc.server.audit(AuditRecord{TraceID: traceID, Timestamp: time.Now(), Tool: params.Name, Allowed: false, DeniedReason: "missing token"})
|
||||
wc.sendError(req.ID, InvalidParams, "Unauthorized", "missing token")
|
||||
return
|
||||
}
|
||||
token = strings.TrimPrefix(token, "Bearer ")
|
||||
acc, err := wc.server.opts.Auth.Inspect(token)
|
||||
if err != nil {
|
||||
span.SetAttributes(attribute.Bool(AttrAuthAllowed, false), attribute.String(AttrAuthDeniedReason, "invalid token"))
|
||||
setSpanError(span, fmt.Errorf("invalid token"))
|
||||
wc.server.audit(AuditRecord{TraceID: traceID, Timestamp: time.Now(), Tool: params.Name, Allowed: false, DeniedReason: "invalid token"})
|
||||
wc.sendError(req.ID, InvalidParams, "Unauthorized", "invalid token")
|
||||
return
|
||||
}
|
||||
account = acc
|
||||
}
|
||||
span.SetAttributes(attribute.String(AttrAccountID, account.ID))
|
||||
|
||||
// Check per-tool scopes
|
||||
if len(tool.Scopes) > 0 {
|
||||
span.SetAttributes(attribute.StringSlice(AttrScopesRequired, tool.Scopes))
|
||||
if !hasScope(account.Scopes, tool.Scopes) {
|
||||
span.SetAttributes(attribute.Bool(AttrAuthAllowed, false), attribute.String(AttrAuthDeniedReason, "insufficient scopes"))
|
||||
setSpanError(span, fmt.Errorf("insufficient scopes"))
|
||||
wc.server.audit(AuditRecord{
|
||||
TraceID: traceID, Timestamp: time.Now(), Tool: params.Name,
|
||||
AccountID: account.ID, ScopesRequired: tool.Scopes,
|
||||
Allowed: false, DeniedReason: "insufficient scopes",
|
||||
})
|
||||
wc.sendError(req.ID, InvalidParams, "Forbidden", "insufficient scopes")
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Rate limit check
|
||||
if err := wc.server.allowRate(params.Name); err != nil {
|
||||
span.SetAttributes(attribute.Bool(AttrRateLimited, true))
|
||||
setSpanError(span, err)
|
||||
accountID := ""
|
||||
if account != nil {
|
||||
accountID = account.ID
|
||||
}
|
||||
wc.server.audit(AuditRecord{
|
||||
TraceID: traceID, Timestamp: time.Now(), Tool: params.Name,
|
||||
AccountID: accountID, Allowed: false, DeniedReason: "rate limited",
|
||||
})
|
||||
wc.sendError(req.ID, InternalError, "Rate limit exceeded", params.Name)
|
||||
return
|
||||
}
|
||||
|
||||
span.SetAttributes(attribute.Bool(AttrAuthAllowed, true))
|
||||
|
||||
// Convert arguments to JSON bytes
|
||||
inputBytes, err := json.Marshal(params.Arguments)
|
||||
if err != nil {
|
||||
wc.sendError(req.ID, InternalError, "Failed to marshal arguments", err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// Build context with tracing metadata
|
||||
md, _ := metadata.FromContext(ctx)
|
||||
if md == nil {
|
||||
md = make(metadata.Metadata)
|
||||
}
|
||||
md.Set(TraceIDKey, traceID)
|
||||
md.Set(ToolNameKey, params.Name)
|
||||
if account != nil {
|
||||
md.Set(AccountIDKey, account.ID)
|
||||
}
|
||||
ctx = metadata.NewContext(ctx, md)
|
||||
|
||||
// Make RPC call
|
||||
start := time.Now()
|
||||
rpcReq := wc.server.opts.Client.NewRequest(tool.Service, tool.Endpoint, &struct {
|
||||
Data []byte
|
||||
}{Data: inputBytes})
|
||||
|
||||
var rsp struct {
|
||||
Data []byte
|
||||
}
|
||||
|
||||
if err := wc.server.opts.Client.Call(ctx, rpcReq, &rsp); err != nil {
|
||||
setSpanError(span, err)
|
||||
accountID := ""
|
||||
if account != nil {
|
||||
accountID = account.ID
|
||||
}
|
||||
wc.server.audit(AuditRecord{
|
||||
TraceID: traceID, Timestamp: time.Now(), Tool: params.Name,
|
||||
AccountID: accountID, ScopesRequired: tool.Scopes,
|
||||
Allowed: true, Duration: time.Since(start), Error: err.Error(),
|
||||
})
|
||||
// Tool-execution failure → isError result (MCP spec), not a protocol error.
|
||||
wc.sendResponse(req.ID, mcpToolError(traceID, "tool call failed: "+err.Error()))
|
||||
return
|
||||
}
|
||||
|
||||
setSpanOK(span)
|
||||
|
||||
accountID := ""
|
||||
if account != nil {
|
||||
accountID = account.ID
|
||||
}
|
||||
wc.server.audit(AuditRecord{
|
||||
TraceID: traceID, Timestamp: time.Now(), Tool: params.Name,
|
||||
AccountID: accountID, ScopesRequired: tool.Scopes,
|
||||
Allowed: true, Duration: time.Since(start),
|
||||
})
|
||||
|
||||
// The downstream response is JSON — return it as JSON text, not %v.
|
||||
wc.sendResponse(req.ID, mcpToolResult(traceID, rsp.Data))
|
||||
}
|
||||
|
||||
// sendResponse sends a JSON-RPC success response.
|
||||
func (wc *wsConn) sendResponse(id interface{}, result interface{}) {
|
||||
wc.writeJSON(JSONRPCResponse{
|
||||
JSONRPC: "2.0",
|
||||
ID: id,
|
||||
Result: result,
|
||||
})
|
||||
}
|
||||
|
||||
// sendError sends a JSON-RPC error response.
|
||||
func (wc *wsConn) sendError(id interface{}, code int, message string, data interface{}) {
|
||||
wc.writeJSON(JSONRPCResponse{
|
||||
JSONRPC: "2.0",
|
||||
ID: id,
|
||||
Error: &RPCError{
|
||||
Code: code,
|
||||
Message: message,
|
||||
Data: data,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// writeJSON serializes and sends a JSON message over the WebSocket.
|
||||
func (wc *wsConn) writeJSON(v interface{}) {
|
||||
wc.writeMu.Lock()
|
||||
defer wc.writeMu.Unlock()
|
||||
|
||||
if err := wc.conn.WriteJSON(v); err != nil {
|
||||
wc.server.opts.Logger.Printf("[mcp] WebSocket write error: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,447 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go-micro.dev/v6/auth"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
// wsDialer creates a WebSocket connection to the given test server URL.
|
||||
func wsDialer(t *testing.T, url string, headers http.Header) *websocket.Conn {
|
||||
t.Helper()
|
||||
wsURL := "ws" + strings.TrimPrefix(url, "http")
|
||||
conn, _, err := websocket.DefaultDialer.Dial(wsURL, headers)
|
||||
if err != nil {
|
||||
t.Fatalf("WebSocket dial failed: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { conn.Close() })
|
||||
return conn
|
||||
}
|
||||
|
||||
// sendJSONRPC sends a JSON-RPC request and reads the response.
|
||||
func sendJSONRPC(t *testing.T, conn *websocket.Conn, method string, id interface{}, params interface{}) JSONRPCResponse {
|
||||
t.Helper()
|
||||
raw, _ := json.Marshal(params)
|
||||
req := JSONRPCRequest{
|
||||
JSONRPC: "2.0",
|
||||
ID: id,
|
||||
Method: method,
|
||||
Params: raw,
|
||||
}
|
||||
if err := conn.WriteJSON(req); err != nil {
|
||||
t.Fatalf("WriteJSON failed: %v", err)
|
||||
}
|
||||
|
||||
var resp JSONRPCResponse
|
||||
if err := conn.ReadJSON(&resp); err != nil {
|
||||
t.Fatalf("ReadJSON failed: %v", err)
|
||||
}
|
||||
return resp
|
||||
}
|
||||
|
||||
func newWSTestServer(t *testing.T, opts Options) (*Server, *httptest.Server) {
|
||||
t.Helper()
|
||||
s := newTestServer(opts)
|
||||
ws := NewWebSocketTransport(s)
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.Handle("/mcp/ws", ws)
|
||||
ts := httptest.NewServer(mux)
|
||||
t.Cleanup(ts.Close)
|
||||
return s, ts
|
||||
}
|
||||
|
||||
func TestWebSocket_Initialize(t *testing.T) {
|
||||
_, ts := newWSTestServer(t, Options{})
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", nil)
|
||||
|
||||
resp := sendJSONRPC(t, conn, "initialize", 1, nil)
|
||||
if resp.Error != nil {
|
||||
t.Fatalf("unexpected error: %v", resp.Error)
|
||||
}
|
||||
|
||||
result, ok := resp.Result.(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatal("expected map result")
|
||||
}
|
||||
if result["protocolVersion"] != "2024-11-05" {
|
||||
t.Errorf("protocolVersion = %v, want 2024-11-05", result["protocolVersion"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSocket_ToolsList(t *testing.T) {
|
||||
s, ts := newWSTestServer(t, Options{})
|
||||
s.tools["svc.Echo"] = &Tool{
|
||||
Name: "svc.Echo",
|
||||
Description: "Echo a message",
|
||||
InputSchema: map[string]interface{}{"type": "object"},
|
||||
Service: "svc",
|
||||
Endpoint: "Echo",
|
||||
}
|
||||
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", nil)
|
||||
resp := sendJSONRPC(t, conn, "tools/list", 1, nil)
|
||||
if resp.Error != nil {
|
||||
t.Fatalf("unexpected error: %v", resp.Error)
|
||||
}
|
||||
|
||||
result, _ := resp.Result.(map[string]interface{})
|
||||
tools, _ := result["tools"].([]interface{})
|
||||
if len(tools) != 1 {
|
||||
t.Fatalf("expected 1 tool, got %d", len(tools))
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSocket_ToolsCall_NoAuth(t *testing.T) {
|
||||
s, ts := newWSTestServer(t, Options{})
|
||||
s.tools["svc.Echo"] = &Tool{
|
||||
Name: "svc.Echo",
|
||||
Service: "svc",
|
||||
Endpoint: "Echo",
|
||||
}
|
||||
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", nil)
|
||||
resp := sendJSONRPC(t, conn, "tools/call", 1, map[string]interface{}{
|
||||
"name": "svc.Echo",
|
||||
"arguments": map[string]interface{}{"msg": "hi"},
|
||||
})
|
||||
|
||||
// No auth required → the tool runs; the RPC fails (no backend), which the
|
||||
// MCP spec surfaces as an isError result, not a JSON-RPC protocol error.
|
||||
if resp.Error != nil {
|
||||
t.Fatalf("expected no protocol error, got %+v", resp.Error)
|
||||
}
|
||||
if !isToolError(resp.Result) {
|
||||
t.Fatalf("expected isError tool result, got %+v", resp.Result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSocket_ToolsCall_AuthRequired(t *testing.T) {
|
||||
ma := &mockAuth{
|
||||
accounts: map[string]*auth.Account{
|
||||
"valid-token": {ID: "user-1", Scopes: []string{"blog:write"}},
|
||||
},
|
||||
}
|
||||
|
||||
s, ts := newWSTestServer(t, Options{Auth: ma})
|
||||
s.tools["svc.Do"] = &Tool{
|
||||
Name: "svc.Do",
|
||||
Service: "svc",
|
||||
Endpoint: "Do",
|
||||
Scopes: []string{"blog:write"},
|
||||
}
|
||||
|
||||
t.Run("missing token", func(t *testing.T) {
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", nil)
|
||||
resp := sendJSONRPC(t, conn, "tools/call", 1, map[string]interface{}{
|
||||
"name": "svc.Do",
|
||||
"arguments": map[string]interface{}{},
|
||||
})
|
||||
if resp.Error == nil || resp.Error.Message != "Unauthorized" {
|
||||
t.Errorf("expected Unauthorized, got %+v", resp.Error)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("invalid token", func(t *testing.T) {
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", nil)
|
||||
resp := sendJSONRPC(t, conn, "tools/call", 1, map[string]interface{}{
|
||||
"name": "svc.Do",
|
||||
"arguments": map[string]interface{}{},
|
||||
"_token": "bad-token",
|
||||
})
|
||||
if resp.Error == nil || resp.Error.Message != "Unauthorized" {
|
||||
t.Errorf("expected Unauthorized, got %+v", resp.Error)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("valid token via param", func(t *testing.T) {
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", nil)
|
||||
resp := sendJSONRPC(t, conn, "tools/call", 1, map[string]interface{}{
|
||||
"name": "svc.Do",
|
||||
"arguments": map[string]interface{}{},
|
||||
"_token": "valid-token",
|
||||
})
|
||||
// Auth passes → the tool runs; RPC fails (no backend) → isError result,
|
||||
// not a JSON-RPC protocol error (which would mean auth failed).
|
||||
if resp.Error != nil {
|
||||
t.Fatalf("expected no protocol error (auth passed), got %+v", resp.Error)
|
||||
}
|
||||
if !isToolError(resp.Result) {
|
||||
t.Fatalf("expected isError tool result, got %+v", resp.Result)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("valid token via header", func(t *testing.T) {
|
||||
headers := http.Header{}
|
||||
headers.Set("Authorization", "Bearer valid-token")
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", headers)
|
||||
resp := sendJSONRPC(t, conn, "tools/call", 1, map[string]interface{}{
|
||||
"name": "svc.Do",
|
||||
"arguments": map[string]interface{}{},
|
||||
})
|
||||
// Auth passes via connection-level header → tool runs; RPC fails (no
|
||||
// backend) → isError result, not a JSON-RPC protocol error.
|
||||
if resp.Error != nil {
|
||||
t.Fatalf("expected no protocol error (auth passed), got %+v", resp.Error)
|
||||
}
|
||||
if !isToolError(resp.Result) {
|
||||
t.Fatalf("expected isError tool result, got %+v", resp.Result)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestWebSocket_ToolsCall_InsufficientScopes(t *testing.T) {
|
||||
ma := &mockAuth{
|
||||
accounts: map[string]*auth.Account{
|
||||
"readonly": {ID: "user-2", Scopes: []string{"blog:read"}},
|
||||
},
|
||||
}
|
||||
|
||||
s, ts := newWSTestServer(t, Options{Auth: ma})
|
||||
s.tools["svc.Do"] = &Tool{
|
||||
Name: "svc.Do",
|
||||
Service: "svc",
|
||||
Endpoint: "Do",
|
||||
Scopes: []string{"blog:write"},
|
||||
}
|
||||
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", nil)
|
||||
resp := sendJSONRPC(t, conn, "tools/call", 1, map[string]interface{}{
|
||||
"name": "svc.Do",
|
||||
"arguments": map[string]interface{}{},
|
||||
"_token": "readonly",
|
||||
})
|
||||
if resp.Error == nil || resp.Error.Message != "Forbidden" {
|
||||
t.Errorf("expected Forbidden, got %+v", resp.Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSocket_ToolsCall_Audit(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
var records []AuditRecord
|
||||
|
||||
s, ts := newWSTestServer(t, Options{
|
||||
AuditFunc: func(r AuditRecord) {
|
||||
mu.Lock()
|
||||
records = append(records, r)
|
||||
mu.Unlock()
|
||||
},
|
||||
})
|
||||
s.tools["svc.Do"] = &Tool{
|
||||
Name: "svc.Do",
|
||||
Service: "svc",
|
||||
Endpoint: "Do",
|
||||
}
|
||||
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", nil)
|
||||
sendJSONRPC(t, conn, "tools/call", 1, map[string]interface{}{
|
||||
"name": "svc.Do",
|
||||
"arguments": map[string]interface{}{},
|
||||
})
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
|
||||
if len(records) == 0 {
|
||||
t.Fatal("expected audit record")
|
||||
}
|
||||
r := records[len(records)-1]
|
||||
if r.Tool != "svc.Do" {
|
||||
t.Errorf("audit Tool = %q, want %q", r.Tool, "svc.Do")
|
||||
}
|
||||
if r.TraceID == "" {
|
||||
t.Error("audit TraceID is empty")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSocket_RateLimit(t *testing.T) {
|
||||
s, ts := newWSTestServer(t, Options{
|
||||
RateLimit: &RateLimitConfig{RequestsPerSecond: 1, Burst: 1},
|
||||
})
|
||||
s.tools["svc.Do"] = &Tool{
|
||||
Name: "svc.Do",
|
||||
Service: "svc",
|
||||
Endpoint: "Do",
|
||||
}
|
||||
s.limiters["svc.Do"] = newRateLimiter(1, 1)
|
||||
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", nil)
|
||||
|
||||
params := map[string]interface{}{
|
||||
"name": "svc.Do",
|
||||
"arguments": map[string]interface{}{},
|
||||
}
|
||||
|
||||
// First request passes rate limit (RPC may fail, that's ok)
|
||||
resp1 := sendJSONRPC(t, conn, "tools/call", 1, params)
|
||||
if resp1.Error != nil && resp1.Error.Message == "Rate limit exceeded" {
|
||||
t.Error("first request should not be rate limited")
|
||||
}
|
||||
|
||||
// Second request should be rate limited
|
||||
resp2 := sendJSONRPC(t, conn, "tools/call", 2, params)
|
||||
if resp2.Error == nil || resp2.Error.Message != "Rate limit exceeded" {
|
||||
t.Errorf("expected rate limit error, got %+v", resp2.Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSocket_MethodNotFound(t *testing.T) {
|
||||
_, ts := newWSTestServer(t, Options{})
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", nil)
|
||||
|
||||
resp := sendJSONRPC(t, conn, "nonexistent/method", 1, nil)
|
||||
if resp.Error == nil || resp.Error.Code != MethodNotFound {
|
||||
t.Errorf("expected MethodNotFound, got %+v", resp.Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSocket_ToolNotFound(t *testing.T) {
|
||||
_, ts := newWSTestServer(t, Options{})
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", nil)
|
||||
|
||||
resp := sendJSONRPC(t, conn, "tools/call", 1, map[string]interface{}{
|
||||
"name": "nonexistent.Tool",
|
||||
"arguments": map[string]interface{}{},
|
||||
})
|
||||
if resp.Error == nil || resp.Error.Message != "Tool not found" {
|
||||
t.Errorf("expected Tool not found, got %+v", resp.Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSocket_MultipleConcurrentRequests(t *testing.T) {
|
||||
s, ts := newWSTestServer(t, Options{})
|
||||
s.tools["svc.Echo"] = &Tool{
|
||||
Name: "svc.Echo",
|
||||
Service: "svc",
|
||||
Endpoint: "Echo",
|
||||
}
|
||||
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", nil)
|
||||
|
||||
// Send multiple requests sequentially (gorilla client doesn't allow
|
||||
// concurrent writes), but the server handles them concurrently.
|
||||
const n = 5
|
||||
for i := 0; i < n; i++ {
|
||||
raw, _ := json.Marshal(map[string]interface{}{
|
||||
"name": "svc.Echo",
|
||||
"arguments": map[string]interface{}{},
|
||||
})
|
||||
req := JSONRPCRequest{
|
||||
JSONRPC: "2.0",
|
||||
ID: i + 1,
|
||||
Method: "tools/call",
|
||||
Params: raw,
|
||||
}
|
||||
if err := conn.WriteJSON(req); err != nil {
|
||||
t.Fatalf("WriteJSON %d failed: %v", i, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Read all responses (order may vary due to concurrent server handling)
|
||||
responses := make([]JSONRPCResponse, n)
|
||||
for i := 0; i < n; i++ {
|
||||
if err := conn.ReadJSON(&responses[i]); err != nil {
|
||||
t.Fatalf("ReadJSON failed at %d: %v", i, err)
|
||||
}
|
||||
}
|
||||
|
||||
for i, resp := range responses {
|
||||
if resp.JSONRPC != "2.0" {
|
||||
t.Errorf("response %d: jsonrpc = %q, want '2.0'", i, resp.JSONRPC)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSocket_MultipleConnections(t *testing.T) {
|
||||
s, ts := newWSTestServer(t, Options{})
|
||||
s.tools["svc.Echo"] = &Tool{
|
||||
Name: "svc.Echo",
|
||||
Description: "Echo",
|
||||
InputSchema: map[string]interface{}{"type": "object"},
|
||||
Service: "svc",
|
||||
Endpoint: "Echo",
|
||||
}
|
||||
|
||||
// Connect multiple clients simultaneously
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 3; i++ {
|
||||
wg.Add(1)
|
||||
go func(idx int) {
|
||||
defer wg.Done()
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", nil)
|
||||
resp := sendJSONRPC(t, conn, "tools/list", idx+1, nil)
|
||||
if resp.Error != nil {
|
||||
t.Errorf("client %d: unexpected error: %v", idx, resp.Error)
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestWebSocket_InvalidJSON(t *testing.T) {
|
||||
_, ts := newWSTestServer(t, Options{})
|
||||
|
||||
wsURL := "ws" + strings.TrimPrefix(ts.URL, "http") + "/mcp/ws"
|
||||
conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
// Send invalid JSON
|
||||
conn.WriteMessage(websocket.TextMessage, []byte("not json"))
|
||||
|
||||
var resp JSONRPCResponse
|
||||
if err := conn.ReadJSON(&resp); err != nil {
|
||||
t.Fatalf("ReadJSON failed: %v", err)
|
||||
}
|
||||
if resp.Error == nil || resp.Error.Code != ParseError {
|
||||
t.Errorf("expected ParseError, got %+v", resp.Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSocket_InvalidJSONRPCVersion(t *testing.T) {
|
||||
_, ts := newWSTestServer(t, Options{})
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", nil)
|
||||
|
||||
req := map[string]interface{}{
|
||||
"jsonrpc": "1.0",
|
||||
"id": 1,
|
||||
"method": "initialize",
|
||||
}
|
||||
conn.WriteJSON(req)
|
||||
|
||||
var resp JSONRPCResponse
|
||||
conn.ReadJSON(&resp)
|
||||
if resp.Error == nil || resp.Error.Code != InvalidRequest {
|
||||
t.Errorf("expected InvalidRequest, got %+v", resp.Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSocket_ConnectionPersistence(t *testing.T) {
|
||||
_, ts := newWSTestServer(t, Options{})
|
||||
conn := wsDialer(t, ts.URL+"/mcp/ws", nil)
|
||||
|
||||
// Send multiple sequential requests on the same connection
|
||||
for i := 0; i < 3; i++ {
|
||||
resp := sendJSONRPC(t, conn, "initialize", i+1, nil)
|
||||
if resp.Error != nil {
|
||||
t.Errorf("request %d: unexpected error: %v", i, resp.Error)
|
||||
}
|
||||
}
|
||||
|
||||
// Connection should still be alive after a short delay
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
resp := sendJSONRPC(t, conn, "initialize", 99, nil)
|
||||
if resp.Error != nil {
|
||||
t.Errorf("request after delay: unexpected error: %v", resp.Error)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user