feat(mcpkit): runtime group lock/unlock — App.Lock/Unlock (MC-45)

lock_unlock_test.go

@@ -0,0 +1,186 @@
+package mcpkit_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/dangernoodle-io/mcpkit"
+	"github.com/dangernoodle-io/mcpkit/host/generic"
+	"github.com/dangernoodle-io/mcpkit/mcpx"
+	"github.com/dangernoodle-io/mcpkit/testkit"
+	"github.com/stretchr/testify/require"
+)
+
+type lockIn struct{}
+
+type lockOut struct{}
+
+func lockHandler(_ context.Context, _ *mcpx.CallToolRequest, _ lockIn) (*mcpx.CallToolResult, lockOut, error) {
+	return nil, lockOut{}, nil
+}
+
+// hwGroupCap registers one tool in group "hw" and one ungrouped tool.
+type hwGroupCap struct{}
+
+func (hwGroupCap) Attach(r *mcpkit.Registrar) error {
+	mcpkit.AddTool(r, &mcpx.Tool{Name: "hw-tool", Description: "d"}, mcpkit.ReadOnly, lockHandler, mcpkit.Group("hw"))
+	mcpkit.AddTool(r, &mcpx.Tool{Name: "ungrouped-tool", Description: "d"}, mcpkit.ReadOnly, lockHandler)
+	return nil
+}
+
+// TestLockBeforeConnectThenUnlockAtRuntime proves the lazy-tier mechanism:
+// locking a group before the app ever connects keeps that group's tools out
+// of the very first tools/list, and a runtime Unlock later brings them in
+// and notifies the connected client via
+// notifications/tools/list_changed (observed here through mcpx.Client's
+// OnToolListChanged, since that requires a raw mcpx client rather than
+// testkit's harness).
+func TestLockBeforeConnectThenUnlockAtRuntime(t *testing.T) {
+	app, err := mcpkit.New(mcpkit.Info{Name: "lazy-tier", Version: "0.0.1"}, generic.New(), hwGroupCap{})
+	require.NoError(t, err)
+
+	require.NoError(t, app.Lock("hw"))
+
+	ctx := context.Background()
+	serverT, clientT := mcpx.InMemoryPair()
+
+	srvSess, err := app.Connect(ctx, serverT)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = srvSess.Close() })
+
+	changed := make(chan struct{}, 4)
+	client := mcpx.NewClient(mcpx.Implementation{Name: "lazy-tier-client", Version: "0.0.1"}, &mcpx.ClientOptions{
+		OnToolListChanged: func(_ context.Context) {
+			changed <- struct{}{}
+		},
+	})
+	clientSess, err := client.Connect(ctx, clientT)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = clientSess.Close() })
+
+	tools, err := clientSess.ListTools(ctx)
+	require.NoError(t, err)
+	require.Len(t, tools.Tools, 1, "locked group's tool must not appear in the initial tools/list")
+	require.Equal(t, "ungrouped-tool", tools.Tools[0].Name)
+
+	require.NoError(t, app.Unlock("hw"))
+
+	select {
+	case <-changed:
+	case <-time.After(5 * time.Second):
+		t.Fatal("did not receive tool list changed notification after Unlock")
+	}
+
+	tools, err = clientSess.ListTools(ctx)
+	require.NoError(t, err)
+	names := make([]string, 0, len(tools.Tools))
+	for _, tool := range tools.Tools {
+		names = append(names, tool.Name)
+	}
+	require.ElementsMatch(t, []string{"hw-tool", "ungrouped-tool"}, names, "Unlock must bring the group's tool into tools/list")
+}
+
+// TestLockAtRuntimeUnregistersTool proves a runtime Lock (called after the
+// app is already connected) truly unregisters the group's tools, not just
+// hides them: the tool disappears from tools/list, and calling it directly
+// by name is rejected by the server.
+func TestLockAtRuntimeUnregistersTool(t *testing.T) {
+	app, err := mcpkit.New(mcpkit.Info{Name: "runtime-lock", Version: "0.0.1"}, generic.New(), hwGroupCap{})
+	require.NoError(t, err)
+
+	h := testkit.New(t, app)
+	testkit.AssertToolSet(t, h, "hw-tool", "ungrouped-tool")
+
+	require.NoError(t, app.Lock("hw"))
+
+	testkit.AssertToolSet(t, h, "ungrouped-tool")
+
+	_, err = h.CallTool(context.Background(), "hw-tool", map[string]any{})
+	require.Error(t, err, "a locked-off tool must be truly unregistered, not merely hidden")
+}
+
+// TestLockHardBlockedGroupErrors proves the startup gate (MC-44 BlockGroups)
+// takes precedence over MC-45's runtime lock: both Lock and Unlock on a
+// hard-blocked group return an error, and the tool set is unaffected by
+// either call.
+func TestLockHardBlockedGroupErrors(t *testing.T) {
+	app, err := mcpkit.New(mcpkit.Info{Name: "hard-block", Version: "0.0.1"}, generic.New(), hwGroupCap{})
+	require.NoError(t, err)
+
+	require.NoError(t, app.BlockGroups("hw"))
+
+	h := testkit.New(t, app)
+	testkit.AssertToolSet(t, h, "ungrouped-tool")
+
+	require.Error(t, app.Lock("hw"))
+	testkit.AssertToolSet(t, h, "ungrouped-tool")
+
+	require.Error(t, app.Unlock("hw"))
+	testkit.AssertToolSet(t, h, "ungrouped-tool")
+}
+
+type readOnlyGuardIn struct{}
+
+type readOnlyGuardOut struct{}
+
+func readOnlyGuardHandler(_ context.Context, _ *mcpx.CallToolRequest, _ readOnlyGuardIn) (*mcpx.CallToolResult, readOnlyGuardOut, error) {
+	return nil, readOnlyGuardOut{}, nil
+}
+
+// mixedRiskGroupCap registers a ReadOnly and a Write tool in the same group,
+// so Unlock's gate re-check can be observed acting on one but not the other.
+type mixedRiskGroupCap struct{}
+
+func (mixedRiskGroupCap) Attach(r *mcpkit.Registrar) error {
+	mcpkit.AddTool(r, &mcpx.Tool{Name: "hw-read", Description: "d"}, mcpkit.ReadOnly, readOnlyGuardHandler, mcpkit.Group("hw"))
+	mcpkit.AddTool(r, &mcpx.Tool{Name: "hw-write", Description: "d"}, mcpkit.Write, readOnlyGuardHandler, mcpkit.Group("hw"))
+	return nil
+}
+
+// TestUnlockDoesNotResurrectReadOnlyGateBlockedTool proves Unlock's
+// shouldRegister re-check is genuine: under ReadOnlyMode, a Write tool
+// gate-blocked at finalize is never brought back by Unlock, even though a
+// ReadOnly tool in the very same (previously locked) group is.
+func TestUnlockDoesNotResurrectReadOnlyGateBlockedTool(t *testing.T) {
+	app, err := mcpkit.New(mcpkit.Info{Name: "ro-guard", Version: "0.0.1"}, generic.New(), mixedRiskGroupCap{})
+	require.NoError(t, err)
+
+	require.NoError(t, app.Gate(mcpkit.ReadOnlyMode()))
+
+	h := testkit.New(t, app)
+	// hw-write is gate-blocked at finalize and never registers; hw-read
+	// (ReadOnly) does.
+	testkit.AssertToolSet(t, h, "hw-read")
+
+	// Lock then Unlock the group at runtime so Unlock's re-registration loop
+	// (not finalize's) is what's under test.
+	require.NoError(t, app.Lock("hw"))
+	testkit.AssertToolSet(t, h)
+
+	require.NoError(t, app.Unlock("hw"))
+	testkit.AssertToolSet(t, h, "hw-read")
+}
+
+// TestLockUnlockIdempotency proves double Lock and double Unlock calls are
+// safe no-ops: neither panics, and neither double-registers or leaves the
+// tool set in an unexpected state.
+func TestLockUnlockIdempotency(t *testing.T) {
+	app, err := mcpkit.New(mcpkit.Info{Name: "idempotent-lock", Version: "0.0.1"}, generic.New(), hwGroupCap{})
+	require.NoError(t, err)
+
+	h := testkit.New(t, app)
+	testkit.AssertToolSet(t, h, "hw-tool", "ungrouped-tool")
+
+	require.NotPanics(t, func() {
+		require.NoError(t, app.Lock("hw"))
+		require.NoError(t, app.Lock("hw"))
+	})
+	testkit.AssertToolSet(t, h, "ungrouped-tool")
+
+	require.NotPanics(t, func() {
+		require.NoError(t, app.Unlock("hw"))
+		require.NoError(t, app.Unlock("hw"))
+	})
+	testkit.AssertToolSet(t, h, "hw-tool", "ungrouped-tool")
+}

mcpkit.go

@@ -163,6 +163,12 @@ type registry struct {
 	byGroup map[string][]string
 	started bool
 	gate    gateState
+
+	// locked tracks MC-45's runtime, per-group soft lock: distinct from
+	// gate, which is the permanent startup hard block. A locked group's
+	// pending tools stay in pending (never discarded) so a later Unlock can
+	// still register them.
+	locked map[string]bool
 }
 
 // add appends t to the pending set. Safe for concurrent Attach calls.
@@ -210,14 +216,31 @@ func (reg *registry) blockGroups(groups ...string) error {
 	return nil
 }
 
+// gateBlocked reports whether the startup gate (MC-44) hard-blocks t: its
+// risk fails ReadOnlyMode, or its group is in gate.blockedGroups. This is a
+// permanent block for the lifetime of the App — MC-45's Lock/Unlock never
+// override it. Callers must hold reg.mu.
+func (reg *registry) gateBlocked(t pendingTool) bool {
+	return (reg.gate.readOnly && t.risk != ReadOnly) || reg.gate.blockedGroups[t.group]
+}
+
+// shouldRegister reports whether t should be registered against the live
+// server right now: it must clear both the permanent startup gate and the
+// MC-45 runtime per-group lock. Callers must hold reg.mu.
+func (reg *registry) shouldRegister(t pendingTool) bool {
+	return !reg.gateBlocked(t) && !reg.locked[t.group]
+}
+
 // finalize registers every pending tool against srv exactly once, skipping
-// any tool the startup gate (MC-44) blocks on either axis: risk (readOnly)
-// or group. A gated-off tool is never registered against srv and never
-// recorded in byGroup, so it can't appear in tools/list, can't be called,
-// and (per MC-44's hard-block contract) can't later be resurrected by
-// MC-45's runtime Unlock. finalize 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.
+// any tool shouldRegister excludes: a gate (MC-44) hard block, or a
+// pre-start MC-45 Lock on its group. A gated-off tool is never registered
+// against srv and never recorded in byGroup, so it can't appear in
+// tools/list, can't be called, and (per MC-44's hard-block contract) can't
+// later be resurrected by MC-45's runtime Unlock. A locked-but-not-blocked
+// tool is likewise skipped here but stays in pending, so Unlock can register
+// it later. finalize 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()
@@ -232,15 +255,81 @@ func (reg *registry) finalize(srv *mcpx.Server) {
 	}
 
 	for _, t := range reg.pending {
-		blocked := (reg.gate.readOnly && t.risk != ReadOnly) || reg.gate.blockedGroups[t.group]
-		if blocked {
+		if !reg.shouldRegister(t) {
 			continue
 		}
 		t.register(srv)
 		reg.byGroup[t.group] = append(reg.byGroup[t.group], t.name)
 	}
 }
 
+// lockGroup implements MC-45's runtime App.Lock: it hard-disables group on
+// the live server. If group is startup-gate-hard-blocked (MC-44), Lock
+// returns an error rather than pretending to toggle an already-permanent
+// block. Otherwise it marks group locked and, if the registry has already
+// started, removes its currently-registered tools from srv (firing
+// notifications/tools/list_changed via mcpx.Server.RemoveTools) and clears
+// its byGroup bucket. Locking an already-locked group is a no-op. Safe to
+// call before or after finalize.
+func (reg *registry) lockGroup(srv *mcpx.Server, group string) error {
+	reg.mu.Lock()
+	defer reg.mu.Unlock()
+
+	if reg.gate.blockedGroups[group] {
+		return fmt.Errorf("mcpkit: group %q is hard-blocked by the startup gate; Lock has no effect", group)
+	}
+
+	if reg.locked == nil {
+		reg.locked = make(map[string]bool)
+	}
+	if reg.locked[group] {
+		return nil
+	}
+	reg.locked[group] = true
+
+	if reg.started {
+		if names := reg.byGroup[group]; len(names) > 0 {
+			srv.RemoveTools(names...)
+		}
+		delete(reg.byGroup, group)
+	}
+	return nil
+}
+
+// unlockGroup implements MC-45's runtime App.Unlock: it reverses lockGroup.
+// If group is startup-gate-hard-blocked (MC-44), Unlock returns an error —
+// a hard block always wins and cannot be runtime-toggled. Otherwise it
+// clears group's lock and, if the registry has already started, registers
+// every one of group's pending tools that shouldRegister still allows
+// (i.e. not gate-blocked — this is what keeps a ReadOnlyMode-blocked Write
+// tool from being resurrected) against srv, recording each in byGroup.
+// Unlocking an already-unlocked group is a no-op. Safe to call before or
+// after finalize.
+func (reg *registry) unlockGroup(srv *mcpx.Server, group string) error {
+	reg.mu.Lock()
+	defer reg.mu.Unlock()
+
+	if reg.gate.blockedGroups[group] {
+		return fmt.Errorf("mcpkit: group %q is hard-blocked by the startup gate; Unlock has no effect", group)
+	}
+
+	if !reg.locked[group] {
+		return nil
+	}
+	reg.locked[group] = false
+
+	if reg.started {
+		for _, t := range reg.pending {
+			if t.group != group || reg.gateBlocked(t) {
+				continue
+			}
+			t.register(srv)
+			reg.byGroup[group] = append(reg.byGroup[group], t.name)
+		}
+	}
+	return nil
+}
+
 // Capability is a self-contained unit of server functionality, attached to
 // the composition root at build time.
 type Capability interface {
@@ -301,6 +390,29 @@ func (a *App) BlockGroups(groups ...string) error {
 	return a.reg.blockGroups(groups...)
 }
 
+// Lock hard-disables group g on a's live server at runtime: any of g's
+// currently-registered tools are removed (firing
+// notifications/tools/list_changed) and g's pending tools are skipped for
+// registration until a matching Unlock. Lock returns an error if g is
+// hard-blocked by the startup gate (Gate/BlockGroups) — a hard block always
+// wins and can't be runtime-toggled. Lock is idempotent and may be called
+// before or after Run/Connect/HTTPHandler, which is what lets a consumer
+// start a group locked (the lazy tier) and unlock it mid-session.
+func (a *App) Lock(group string) error {
+	return a.reg.lockGroup(a.server, group)
+}
+
+// Unlock reverses Lock: every one of group's pending tools not otherwise
+// hard-blocked by the startup gate is registered against a's live server
+// (firing notifications/tools/list_changed once the app has started). A
+// tool the startup gate hard-blocks (e.g. a Write tool under ReadOnlyMode)
+// is never resurrected by Unlock. Unlock returns an error if group is
+// hard-blocked by the startup gate. Unlock is idempotent and may be called
+// before or after Run/Connect/HTTPHandler.
+func (a *App) Unlock(group string) error {
+	return a.reg.unlockGroup(a.server, group)
+}
+
 // 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 {