internal/mcpserver/server_test.go
Ref: Size: 8.7 KiB History
package mcpserver
import (
"context"
"encoding/json"
"errors"
"testing"
"github.com/modelcontextprotocol/go-sdk/mcp"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// commonTools is the toolset both transports expose, in the order NewServer
// registers it. A change here is a change to the published contract.
var commonTools = []string{
"vm_create", "vm_list", "vm_info", "vm_exec", "vm_write_file",
"vm_read_file", "vm_expose", "vm_exposures", "vm_unexpose", "vm_destroy",
"ca_upload", "tenant_info",
}
// toolNames connects an in-memory client to s and lists what it advertises.
func toolNames(t *testing.T, s *mcp.Server) []string {
t.Helper()
ctx := t.Context()
serverTr, clientTr := mcp.NewInMemoryTransports()
ss, err := s.Connect(ctx, serverTr, nil)
require.NoError(t, err)
defer ss.Close()
cs, err := mcp.NewClient(&mcp.Implementation{Name: "test", Version: "0"}, nil).Connect(ctx, clientTr, nil)
require.NoError(t, err)
defer cs.Close()
res, err := cs.ListTools(ctx, nil)
require.NoError(t, err)
names := make([]string, 0, len(res.Tools))
for _, tool := range res.Tools {
names = append(names, tool.Name)
}
return names
}
// fakeDelegator stands in for the control plane's delegation endpoints.
type fakeDelegator struct {
begin BeginResult
complete DelegationResult
err error
beganN int
gotCert string
}
func (f *fakeDelegator) Begin(context.Context) (BeginResult, error) {
f.beganN++
return f.begin, f.err
}
func (f *fakeDelegator) Complete(_ context.Context, certificate string) (DelegationResult, error) {
f.gotCert = certificate
return f.complete, f.err
}
// connect wires an in-memory client to s.
func connect(t *testing.T, s *mcp.Server) *mcp.ClientSession {
t.Helper()
ctx := t.Context()
serverTr, clientTr := mcp.NewInMemoryTransports()
ss, err := s.Connect(ctx, serverTr, nil)
require.NoError(t, err)
t.Cleanup(func() { ss.Close() })
cs, err := mcp.NewClient(&mcp.Implementation{Name: "test", Version: "0"}, nil).Connect(ctx, clientTr, nil)
require.NoError(t, err)
t.Cleanup(func() { cs.Close() })
return cs
}
// describe returns one tool's advertised description.
func describe(t *testing.T, cs *mcp.ClientSession, name string) string {
t.Helper()
res, err := cs.ListTools(t.Context(), nil)
require.NoError(t, err)
for _, tool := range res.Tools {
if tool.Name == name {
return tool.Description
}
}
return ""
}
// TestNewServerExposesEveryTool: the Delegator is required, so there is only
// one server shape and it offers the VM tools, ca_upload, tenant_info, and both
// delegate tools. A change here is a change to the published contract.
func TestNewServerExposesEveryTool(t *testing.T) {
names := toolNames(t, NewServer(&Tools{}, Options{Delegator: &fakeDelegator{}}))
want := append(append([]string{}, commonTools...), "delegate_begin", "delegate_complete")
assert.ElementsMatch(t, want, names)
}
// TestDelegateBeginDescriptionStopsAModelSigningItItself: the description is
// the only thing standing between a model and an hour of trying to produce a
// certificate eitri deliberately cannot produce.
func TestDelegateBeginDescriptionStopsAModelSigningItItself(t *testing.T) {
desc := describe(t, connect(t, NewServer(&Tools{}, Options{Delegator: &fakeDelegator{}})), "delegate_begin")
require.NotEmpty(t, desc)
assert.Contains(t, desc, "cannot sign this itself")
assert.Contains(t, desc, "Show the human")
assert.Contains(t, desc, "wait")
}
// TestDelegateCompleteDescriptionSaysACertificateIsPublic: a model may
// otherwise refuse to paste something that looks like key material.
func TestDelegateCompleteDescriptionSaysACertificateIsPublic(t *testing.T) {
desc := describe(t, connect(t, NewServer(&Tools{}, Options{Delegator: &fakeDelegator{}})), "delegate_complete")
require.NotEmpty(t, desc)
assert.Contains(t, desc, "public material")
assert.Contains(t, desc, "until the certificate expires")
assert.Contains(t, desc, "VMs created before this call accept it too")
}
// TestDelegateToolsAreWiredToTheDelegator checks both directions: the tools
// call the thing that does the work, and the certificate passes through
// unchanged.
func TestDelegateToolsAreWiredToTheDelegator(t *testing.T) {
dg := &fakeDelegator{
begin: BeginResult{PublicKey: "ssh-ed25519 AAAA eitri", Principal: "ubuntu"},
complete: DelegationResult{ExpiresAt: "2026-08-08T00:00:00Z", CAFingerprint: "SHA256:abc"},
}
cs := connect(t, NewServer(&Tools{}, Options{Delegator: dg}))
res, err := cs.CallTool(t.Context(), &mcp.CallToolParams{Name: "delegate_begin"})
require.NoError(t, err)
require.False(t, res.IsError)
assert.Equal(t, 1, dg.beganN)
assert.Equal(t, "ssh-ed25519 AAAA eitri", res.StructuredContent.(map[string]any)["public_key"])
const cert = "ssh-ed25519-cert-v01@openssh.com AAAAdelegation alex@laptop"
res, err = cs.CallTool(t.Context(), &mcp.CallToolParams{
Name: "delegate_complete", Arguments: map[string]any{"certificate": cert}})
require.NoError(t, err)
require.False(t, res.IsError)
assert.Equal(t, cert, dg.gotCert, "the certificate must reach the control plane byte for byte")
assert.Equal(t, "SHA256:abc", res.StructuredContent.(map[string]any)["ca_fingerprint"])
}
// TestDelegateFailuresSurfaceAsToolErrors: a refusal reaches the model as an
// MCP tool error it can read, not a transport failure it cannot.
func TestDelegateFailuresSurfaceAsToolErrors(t *testing.T) {
cs := connect(t, NewServer(&Tools{}, Options{Delegator: &fakeDelegator{err: errors.New("principals are [alex]")}}))
res, err := cs.CallTool(t.Context(), &mcp.CallToolParams{Name: "delegate_begin"})
require.NoError(t, err)
assert.True(t, res.IsError)
res, err = cs.CallTool(t.Context(), &mcp.CallToolParams{
Name: "delegate_complete", Arguments: map[string]any{"certificate": "x"}})
require.NoError(t, err)
assert.True(t, res.IsError)
}
// TestReportProgressIsANoOpWithoutAListener: the tool layer calls it
// unconditionally, so a caller that asked for no progress must not depend on
// one being installed.
func TestReportProgressIsANoOpWithoutAListener(t *testing.T) {
assert.NotPanics(t, func() { reportProgress(t.Context(), "still working") })
}
// TestReportProgressReachesTheListener pins the seam the long waits use.
func TestReportProgressReachesTheListener(t *testing.T) {
var got []string
ctx := withProgress(t.Context(), func(m string) { got = append(got, m) })
reportProgress(ctx, "creating vm web-1: creating")
reportProgress(ctx, "creating vm web-1: ready")
assert.Equal(t, []string{"creating vm web-1: creating", "creating vm web-1: ready"}, got)
}
// TestProgressReporterIsAbsentWithoutAToken: MCP reports progress only against
// a token the client supplied, so a client that sent none gets nothing.
func TestProgressReporterIsAbsentWithoutAToken(t *testing.T) {
assert.Nil(t, progressReporter(t.Context(), nil))
assert.Nil(t, progressReporter(t.Context(), &mcp.CallToolRequest{}))
}
// TestCAUploadDescriptionNamesTheFrozenCASet: a guest bakes its trusted CA set
// at create, so a model that uploads a CA after creating a VM has done nothing
// for that VM. The description is where it learns that, before it acts.
func TestCAUploadDescriptionNamesTheFrozenCASet(t *testing.T) {
desc := describe(t, connect(t, NewServer(&Tools{}, Options{})), "ca_upload")
require.NotEmpty(t, desc)
assert.Contains(t, desc, "PUBLIC")
assert.Contains(t, desc, "will NOT")
assert.Contains(t, desc, "create new VMs after uploading")
assert.Contains(t, desc, "before vm_create")
}
// A guest on a named host network is addressed by a DHCP server eitri does not
// run, so vm_create can return ready with network_ip still empty. A model that
// has not been told reads that as a broken create and destroys a working VM, so
// the tool's own description says it, and says where to look instead.
func TestVMCreateDescriptionSaysTheNetworkAddressCanLagReady(t *testing.T) {
desc := describe(t, connect(t, NewServer(&Tools{}, Options{})), "vm_create")
require.NotEmpty(t, desc)
assert.Contains(t, desc, "network_ip")
assert.Contains(t, desc, "EMPTY")
assert.Contains(t, desc, "vm_info")
}
// The network field's own description has to carry the rule the control plane
// enforces: only names the host advertises, and no silent fallback to NAT for
// one it does not.
func TestVMCreateNetworkFieldDescribesWhatAHostWillAccept(t *testing.T) {
res, err := connect(t, NewServer(&Tools{}, Options{})).ListTools(t.Context(), nil)
require.NoError(t, err)
var schema string
for _, tool := range res.Tools {
if tool.Name == "vm_create" {
props, err := json.Marshal(tool.InputSchema)
require.NoError(t, err)
schema = string(props)
}
}
require.NotEmpty(t, schema)
assert.Contains(t, schema, "network")
assert.Contains(t, schema, "advertises")
assert.Contains(t, schema, "NAT underlay")
}