diff --git a/examples/http/main.go b/examples/http/main.go index 14c8d74..3aa2397 100644 --- a/examples/http/main.go +++ b/examples/http/main.go @@ -30,7 +30,7 @@ func (helloCap) Attach(r *mcpkit.Registrar) error { mcpkit.AddTool(r, &mcpx.Tool{ Name: "hello", Description: "greets the caller by name", - }, func(_ context.Context, _ *mcpx.CallToolRequest, in helloIn) (*mcpx.CallToolResult, helloOut, error) { + }, mcpkit.ReadOnly, func(_ context.Context, _ *mcpx.CallToolRequest, in helloIn) (*mcpx.CallToolResult, helloOut, error) { name := in.Name if name == "" { name = "world" diff --git a/examples/minimal/main.go b/examples/minimal/main.go index 2fca174..68d6289 100644 --- a/examples/minimal/main.go +++ b/examples/minimal/main.go @@ -26,7 +26,7 @@ func (helloCap) Attach(r *mcpkit.Registrar) error { mcpkit.AddTool(r, &mcpx.Tool{ Name: "hello", Description: "greets the caller by name", - }, func(_ context.Context, _ *mcpx.CallToolRequest, in helloIn) (*mcpx.CallToolResult, helloOut, error) { + }, mcpkit.ReadOnly, func(_ context.Context, _ *mcpx.CallToolRequest, in helloIn) (*mcpx.CallToolResult, helloOut, error) { name := in.Name if name == "" { name = "world" diff --git a/examples/server/main.go b/examples/server/main.go index 9dbe8a3..fed1dcd 100644 --- a/examples/server/main.go +++ b/examples/server/main.go @@ -33,7 +33,7 @@ func (pingCap) Attach(r *mcpkit.Registrar) error { mcpkit.AddTool(r, &mcpx.Tool{ Name: "ping", Description: "replies with pong", - }, func(_ context.Context, _ *mcpx.CallToolRequest, _ pingIn) (*mcpx.CallToolResult, pingOut, error) { + }, mcpkit.ReadOnly, func(_ context.Context, _ *mcpx.CallToolRequest, _ pingIn) (*mcpx.CallToolResult, pingOut, error) { return nil, pingOut{Message: "pong"}, nil }) return nil diff --git a/mcpkit.go b/mcpkit.go index 3d353ed..7668db2 100644 --- a/mcpkit.go +++ b/mcpkit.go @@ -7,6 +7,7 @@ import ( "context" "fmt" "net/http" + "sync" "github.com/dangernoodle-io/mcpkit/host" "github.com/dangernoodle-io/mcpkit/mcpx" @@ -22,6 +23,49 @@ type Info struct { Instructions string } +// Risk classifies a tool's blast radius, from a client's perspective, for +// gating and annotation purposes. Every AddTool call requires one +// (fail-closed): a caller can't omit risk classification and accidentally +// ship a write/destructive tool unannotated. +type Risk int + +const ( + // ReadOnly tools only observe state; they never mutate anything a + // client cares about. + ReadOnly Risk = iota + // Write tools mutate state, but the mutation is reversible/benign + // enough not to warrant a destructive warning. + Write + // Destructive tools perform an irreversible or high-blast-radius + // mutation (data loss, external side effects that can't be undone). + Destructive +) + +// ToolOption configures optional per-tool metadata at AddTool time. +type ToolOption interface { + applyTool(*toolMeta) +} + +// toolMeta is the mutable state ToolOption values apply to. +type toolMeta struct { + group string +} + +// groupOption is the ToolOption Group returns. +type groupOption string + +func (g groupOption) applyTool(m *toolMeta) { + m.group = string(g) +} + +// Group tags a tool with an arbitrary consumer-defined group name, recorded +// in the registry's byGroup bookkeeping post-registration. mcpkit imposes no +// meaning on the group string; a consumer's own gating (MC-44/MC-45) is what +// interprets it. +func Group(name string) ToolOption { + return groupOption(name) +} + // Registrar is what a Capability's Attach method uses to register itself // against the underlying server and inspect the target host. Capabilities // register tools through the package-level AddTool, not against mcpx @@ -29,6 +73,7 @@ type Info struct { type Registrar struct { server *mcpx.Server host host.Adapter + reg *registry } // Host returns the host.Adapter the app is composed for. @@ -36,12 +81,28 @@ func (r *Registrar) Host() host.Adapter { return r.host } -// AddTool registers a typed tool handler through r. This is the only -// tool-registration chokepoint capabilities should use; MC-8 wraps every -// handler in a panic-recover here so a panicking tool surfaces as an -// IsError result instead of crashing the server process. Annotations/risk -// remain future-additive at this same chokepoint. -func AddTool[In, Out any](r *Registrar, t *mcpx.Tool, h mcpx.Handler[In, Out]) { +// AddTool captures a typed tool handler through r for deferred registration +// against the underlying server. This is the only tool-registration +// chokepoint capabilities should use; MC-8 wraps every handler in a +// panic-recover here so a panicking tool surfaces as an IsError result +// instead of crashing the server process. Registration itself is deferred +// until the App's finalize runs (MC-43) so a later gate (MC-44) can filter +// which pending tools actually register before anything is exposed to a +// client. +// +// risk is required (fail-closed): a caller cannot omit a tool's risk +// classification. When t.Annotations is nil, AddTool derives it from risk +// via mcpx.RiskAnnotations; an explicitly-set Annotations is left untouched. +func AddTool[In, Out any](r *Registrar, t *mcpx.Tool, risk Risk, h mcpx.Handler[In, Out], opts ...ToolOption) { + if t.Annotations == nil { + t.Annotations = mcpx.RiskAnnotations(risk == ReadOnly, risk == Destructive) + } + + meta := toolMeta{} + for _, opt := range opts { + opt.applyTool(&meta) + } + wrapped := func(ctx context.Context, req *mcpx.CallToolRequest, in In) (res *mcpx.CallToolResult, out Out, err error) { defer func() { if p := recover(); p != nil { @@ -50,7 +111,67 @@ func AddTool[In, Out any](r *Registrar, t *mcpx.Tool, h mcpx.Handler[In, Out]) { }() return h(ctx, req, in) } - mcpx.AddTool(r.server, t, wrapped) + + r.reg.add(pendingTool{ + name: t.Name, + group: meta.group, + risk: risk, + register: func(s *mcpx.Server) { + mcpx.AddTool(s, t, wrapped) + }, + }) +} + +// pendingTool is one captured-but-not-yet-registered AddTool call. register +// closes over the tool's In/Out type parameters (erased here) and the +// panic-recover-wrapped handler; calling it performs the actual +// mcpx.AddTool registration against a live server. +type pendingTool struct { + name, group string + risk Risk + register func(*mcpx.Server) +} + +// registry accumulates pending tool registrations shared between a +// Registrar (during composition, via AddTool) and its App (at finalize +// time). Deferring registration out of AddTool is what lets a later gate +// (MC-44) filter which pending tools actually register before finalize ever +// touches the live server. +type registry struct { + mu sync.Mutex + pending []pendingTool + byGroup map[string][]string + started bool +} + +// add appends t to the pending set. Safe for concurrent Attach calls. +func (reg *registry) add(t pendingTool) { + reg.mu.Lock() + defer reg.mu.Unlock() + reg.pending = append(reg.pending, t) +} + +// finalize registers every pending tool against srv exactly once. It is +// idempotent: a second call (e.g. Run then Connect, or Run called twice) is +// a guarded no-op rather than double-registering or panicking. +func (reg *registry) finalize(srv *mcpx.Server) { + reg.mu.Lock() + defer reg.mu.Unlock() + + if reg.started { + return + } + reg.started = true + + if reg.byGroup == nil { + reg.byGroup = make(map[string][]string) + } + + for _, t := range reg.pending { + // MC-44 gate predicate slots in here (a per-tool allow check that may `continue`). + t.register(srv) + reg.byGroup[t.group] = append(reg.byGroup[t.group], t.name) + } } // Capability is a self-contained unit of server functionality, attached to @@ -63,17 +184,21 @@ type Capability interface { type App struct { server *mcpx.Server host host.Adapter + reg *registry } // New composes an App from a host.Adapter and zero or more Capabilities, -// attaching each capability in order. +// attaching each capability in order. Tool registration is deferred: no +// tool is registered against the underlying server until Run, Connect, or +// HTTPHandler calls finalize. func New(info Info, h host.Adapter, caps ...Capability) (*App, error) { if h == nil { return nil, fmt.Errorf("mcpkit: host adapter must not be nil") } srv := mcpx.NewServer(mcpx.Implementation{Name: info.Name, Version: info.Version}, info.Instructions) - r := &Registrar{server: srv, host: h} + reg := ®istry{} + r := &Registrar{server: srv, host: h, reg: reg} for _, c := range caps { if err := c.Attach(r); err != nil { @@ -81,25 +206,39 @@ func New(info Info, h host.Adapter, caps ...Capability) (*App, error) { } } - return &App{server: srv, host: h}, nil + return &App{server: srv, host: h, reg: reg}, nil +} + +// finalize registers every pending tool against a's server exactly once, +// regardless of how many of Run/Connect/HTTPHandler trigger it or in what +// order. +func (a *App) finalize() { + a.reg.finalize(a.server) } // Run serves the app over its host's transport until the client disconnects // or ctx is cancelled. func (a *App) Run(ctx context.Context) error { + a.finalize() return a.server.Run(ctx, a.host.Transport()) } // Connect connects the app over t without blocking, for use by testkit and // other in-process harnesses. func (a *App) Connect(ctx context.Context, t mcpx.Transport) (*mcpx.Session, error) { + a.finalize() return a.server.Connect(ctx, t) } // HTTPHandler exposes the composed server over streamable-HTTP for the // consumer to mount. mcpkit is path-agnostic: the returned handler is bare // and MCP-over-HTTP is entirely opt-in — the consumer decides whether and -// where to mount it. +// where to mount it. finalize runs here too (not just Run/Connect) because +// an HTTP-only consumer (see cli.ServerCmd's --http path and +// examples/http) never calls Run or Connect at all — without this, +// deferred registration would silently ship zero tools over HTTP, +// regressing MC-43's behavior-preservation goal. func (a *App) HTTPHandler(opts ...mcpx.HTTPOption) http.Handler { + a.finalize() return a.server.HTTPHandler(opts...) } diff --git a/mcpkit_test.go b/mcpkit_test.go index ac98317..51d3b0b 100644 --- a/mcpkit_test.go +++ b/mcpkit_test.go @@ -28,7 +28,7 @@ func (helloCap) Attach(r *mcpkit.Registrar) error { mcpkit.AddTool(r, &mcpx.Tool{ Name: "hello", Description: "greets the caller by name", - }, func(_ context.Context, _ *mcpx.CallToolRequest, in helloIn) (*mcpx.CallToolResult, helloOut, error) { + }, mcpkit.ReadOnly, func(_ context.Context, _ *mcpx.CallToolRequest, in helloIn) (*mcpx.CallToolResult, helloOut, error) { name := in.Name if name == "" { name = "world" @@ -76,7 +76,7 @@ func (panicCap) Attach(r *mcpkit.Registrar) error { mcpkit.AddTool(r, &mcpx.Tool{ Name: "panics", Description: "always panics", - }, func(_ context.Context, _ *mcpx.CallToolRequest, _ panicIn) (*mcpx.CallToolResult, panicOut, error) { + }, mcpkit.ReadOnly, func(_ context.Context, _ *mcpx.CallToolRequest, _ panicIn) (*mcpx.CallToolResult, panicOut, error) { panic("kaboom") }) return nil @@ -146,7 +146,7 @@ func (annotatedCap) Attach(r *mcpkit.Registrar) error { ReadOnlyHint: true, DestructiveHint: mcpx.BoolPtr(true), }, - }, func(_ context.Context, _ *mcpx.CallToolRequest, _ annotatedIn) (*mcpx.CallToolResult, annotatedOut, error) { + }, mcpkit.ReadOnly, func(_ context.Context, _ *mcpx.CallToolRequest, _ annotatedIn) (*mcpx.CallToolResult, annotatedOut, error) { return nil, annotatedOut{OK: true}, nil }) return nil @@ -194,3 +194,89 @@ func TestAppHTTPHandler(t *testing.T) { require.Equal(t, http.StatusBadRequest, rec.Code) require.Contains(t, rec.Body.String(), "text/event-stream") } + +type riskIn struct{} + +type riskOut struct{} + +// riskCap registers one tool per Risk value, each left with a nil +// Annotations so AddTool must derive it from risk via mcpx.RiskAnnotations. +type riskCap struct{} + +func (riskCap) Attach(r *mcpkit.Registrar) error { + handler := func(_ context.Context, _ *mcpx.CallToolRequest, _ riskIn) (*mcpx.CallToolResult, riskOut, error) { + return nil, riskOut{}, nil + } + mcpkit.AddTool(r, &mcpx.Tool{Name: "risk-readonly", Description: "d"}, mcpkit.ReadOnly, handler) + mcpkit.AddTool(r, &mcpx.Tool{Name: "risk-write", Description: "d"}, mcpkit.Write, handler) + mcpkit.AddTool(r, &mcpx.Tool{Name: "risk-destructive", Description: "d"}, mcpkit.Destructive, handler) + return nil +} + +// TestAddToolRiskAutoAnnotation proves AddTool derives ToolAnnotations from +// each Risk value via mcpx.RiskAnnotations when the caller left +// t.Annotations nil: ReadOnlyHint tracks Risk == ReadOnly, and +// DestructiveHint tracks Risk == Destructive (Write gets neither hint set). +func TestAddToolRiskAutoAnnotation(t *testing.T) { + app, err := mcpkit.New(mcpkit.Info{Name: "risk-e2e", Version: "0.0.1"}, generic.New(), riskCap{}) + require.NoError(t, err) + + h := testkit.New(t, app) + + res, err := h.ListTools(context.Background()) + require.NoError(t, err) + + byName := make(map[string]*mcpx.ToolAnnotations, len(res.Tools)) + for _, tool := range res.Tools { + byName[tool.Name] = tool.Annotations + } + + ro := byName["risk-readonly"] + require.NotNil(t, ro) + require.True(t, ro.ReadOnlyHint) + require.NotNil(t, ro.DestructiveHint) + require.False(t, *ro.DestructiveHint) + + wr := byName["risk-write"] + require.NotNil(t, wr) + require.False(t, wr.ReadOnlyHint) + require.NotNil(t, wr.DestructiveHint) + require.False(t, *wr.DestructiveHint) + + de := byName["risk-destructive"] + require.NotNil(t, de) + require.False(t, de.ReadOnlyHint) + require.NotNil(t, de.DestructiveHint) + require.True(t, *de.DestructiveHint) +} + +// TestDeferredRegistrationAdvertisesAfterConnect proves the deferred +// registration MC-43 introduces is behavior-preserving from a client's +// perspective: a tool registered via AddTool before the first Connect is +// fully advertised (tools/list) and callable once a session exists. +func TestDeferredRegistrationAdvertisesAfterConnect(t *testing.T) { + app, err := mcpkit.New(mcpkit.Info{Name: "deferred-e2e", Version: "0.0.1"}, generic.New(), helloCap{}) + require.NoError(t, err) + + h := testkit.New(t, app) + + testkit.AssertToolSet(t, h, "hello") +} + +// TestAppConnectIdempotentAcrossSessions proves a second Connect against the +// same App (finalize's idempotency guard) neither panics nor loses tools: a +// second in-memory session still sees the same fully-registered tool set, +// proving finalize's first run is what registered it and the second run +// was a no-op rather than a duplicate-registration panic. +func TestAppConnectIdempotentAcrossSessions(t *testing.T) { + app, err := mcpkit.New(mcpkit.Info{Name: "reconnect-e2e", Version: "0.0.1"}, generic.New(), helloCap{}) + require.NoError(t, err) + + h1 := testkit.New(t, app) + testkit.AssertToolSet(t, h1, "hello") + + require.NotPanics(t, func() { + h2 := testkit.New(t, app) + testkit.AssertToolSet(t, h2, "hello") + }) +} diff --git a/registry_test.go b/registry_test.go new file mode 100644 index 0000000..165a211 --- /dev/null +++ b/registry_test.go @@ -0,0 +1,99 @@ +package mcpkit + +import ( + "context" + "testing" + + "github.com/dangernoodle-io/mcpkit/host/generic" + "github.com/dangernoodle-io/mcpkit/mcpx" + "github.com/stretchr/testify/require" +) + +// TestRegistryFinalizeIdempotent proves finalize registers each pending +// tool exactly once even when called twice (the MC-43 guard Run+Connect, +// or Run+Run, both rely on). +func TestRegistryFinalizeIdempotent(t *testing.T) { + reg := ®istry{} + + calls := 0 + reg.add(pendingTool{name: "t1", register: func(_ *mcpx.Server) { calls++ }}) + + reg.finalize(nil) + reg.finalize(nil) + + require.Equal(t, 1, calls, "finalize must not re-register a pending tool on a second call") + require.True(t, reg.started) +} + +// TestRegistryByGroup proves finalize buckets registered tool names by +// group, including the "" (ungrouped) bucket, and never touches a group a +// tool wasn't assigned to. +func TestRegistryByGroup(t *testing.T) { + reg := ®istry{} + reg.add(pendingTool{name: "a", group: "g1", register: func(_ *mcpx.Server) {}}) + reg.add(pendingTool{name: "b", group: "g1", register: func(_ *mcpx.Server) {}}) + reg.add(pendingTool{name: "c", group: "g2", register: func(_ *mcpx.Server) {}}) + reg.add(pendingTool{name: "d", register: func(_ *mcpx.Server) {}}) + + reg.finalize(nil) + + require.ElementsMatch(t, []string{"a", "b"}, reg.byGroup["g1"]) + require.ElementsMatch(t, []string{"c"}, reg.byGroup["g2"]) + require.ElementsMatch(t, []string{"d"}, reg.byGroup[""]) + require.Len(t, reg.byGroup, 3, "no extra/ghost group buckets") +} + +type groupedCap struct{} + +func (groupedCap) Attach(r *Registrar) error { + AddTool(r, &mcpx.Tool{Name: "grouped-a", Description: "a"}, ReadOnly, + func(_ context.Context, _ *mcpx.CallToolRequest, _ struct{}) (*mcpx.CallToolResult, struct{}, error) { + return nil, struct{}{}, nil + }, Group("alpha")) + AddTool(r, &mcpx.Tool{Name: "grouped-b", Description: "b"}, Write, + func(_ context.Context, _ *mcpx.CallToolRequest, _ struct{}) (*mcpx.CallToolResult, struct{}, error) { + return nil, struct{}{}, nil + }, Group("alpha")) + AddTool(r, &mcpx.Tool{Name: "ungrouped", Description: "c"}, ReadOnly, + func(_ context.Context, _ *mcpx.CallToolRequest, _ struct{}) (*mcpx.CallToolResult, struct{}, error) { + return nil, struct{}{}, nil + }) + return nil +} + +// TestDeferredRegistrationPreFinalize proves registration is genuinely +// deferred: right after New (capabilities have already called AddTool), +// the registry holds pending entries but has not started, i.e. nothing has +// registered against the live server yet. This is asserted at the registry +// (the actual seam MC-43 introduces) rather than through the wire protocol, +// because connecting a session to observe tools/list would itself trigger +// finalize — the very thing under test. +func TestDeferredRegistrationPreFinalize(t *testing.T) { + app, err := New(Info{Name: "deferred", Version: "0.0.1"}, generic.New(), groupedCap{}) + require.NoError(t, err) + + require.False(t, app.reg.started, "registration must not have run before finalize") + require.Len(t, app.reg.pending, 3, "AddTool must capture pending tools without registering them") + require.Empty(t, app.reg.byGroup, "byGroup is only populated by finalize") + + app.finalize() + + require.True(t, app.reg.started) + require.ElementsMatch(t, []string{"grouped-a", "grouped-b"}, app.reg.byGroup["alpha"]) + require.ElementsMatch(t, []string{"ungrouped"}, app.reg.byGroup[""]) +} + +// TestAppFinalizeIdempotent proves App.finalize itself is guarded: calling +// it directly more than once (the same path Run/Connect/HTTPHandler share) +// does not re-register tools against the live mcpx.Server, which would +// otherwise panic on a duplicate tool name. +func TestAppFinalizeIdempotent(t *testing.T) { + app, err := New(Info{Name: "idempotent", Version: "0.0.1"}, generic.New(), groupedCap{}) + require.NoError(t, err) + + require.NotPanics(t, func() { + app.finalize() + app.finalize() + app.finalize() + }) +} diff --git a/testkit/harness_test.go b/testkit/harness_test.go index ce8a719..653a6b6 100644 --- a/testkit/harness_test.go +++ b/testkit/harness_test.go @@ -22,7 +22,7 @@ func (pingCap) Attach(r *mcpkit.Registrar) error { mcpkit.AddTool(r, &mcpx.Tool{ Name: "ping", Description: "replies pong", - }, func(_ context.Context, _ *mcpx.CallToolRequest, _ struct{}) (*mcpx.CallToolResult, pingOut, error) { + }, mcpkit.ReadOnly, func(_ context.Context, _ *mcpx.CallToolRequest, _ struct{}) (*mcpx.CallToolResult, pingOut, error) { return nil, pingOut{Reply: "pong"}, nil }) return nil @@ -54,7 +54,7 @@ func (workCap) Attach(r *mcpkit.Registrar) error { mcpkit.AddTool(r, &mcpx.Tool{ Name: "work", Description: "emits a progress notification keyed to the caller's token, then completes", - }, func(ctx context.Context, req *mcpx.CallToolRequest, _ struct{}) (*mcpx.CallToolResult, workOut, error) { + }, mcpkit.ReadOnly, func(ctx context.Context, req *mcpx.CallToolRequest, _ struct{}) (*mcpx.CallToolResult, workOut, error) { if err := mcpx.NotifyProgress(ctx, req, "halfway", 50, 100); err != nil { return nil, workOut{}, err }