internal/mcpserver/server.go

Ref: main v0.0.5   Size: 8.8 KiB   History

// Package mcpserver implements eitri's MCP toolset: tools that let a model
// create, control (SSH exec/files), and destroy eitri VMs. The control plane
// serves it at /mcp (internal/server/mcphttp). It is an API CLIENT of the
// control plane — it speaks to it only through the shared API client
// (internal/server/api/client), never any other server internals, and the PAT
// it holds must never appear in tool results or errors.
package mcpserver

import (
	"context"

	"github.com/modelcontextprotocol/go-sdk/mcp"
)

// Options is what a server needs beyond the tools themselves.
type Options struct {
	Delegator   Delegator        // required (non-nil): every server exposes the delegate tools
	SchemaCache *mcp.SchemaCache // shared across per-request servers, so re-registering tools costs no reflection
}

// Delegator is the two-step exchange by which a caller lends eitri access:
// eitri offers a public key, the caller's own CA signs it, eitri gets the
// certificate back.
type Delegator interface {
	Begin(ctx context.Context) (BeginResult, error)
	Complete(ctx context.Context, certificate string) (DelegationResult, error)
}

// BeginResult is what the model shows its human. Note is worded for a model
// that would otherwise try to complete the exchange by itself.
type BeginResult struct {
	PublicKey    string `json:"public_key"`
	Principal    string `json:"principal"`
	Instructions string `json:"instructions"`
	Note         string `json:"note"`
}

// DelegationResult reports what eitri may now do, and until when.
type DelegationResult struct {
	ExpiresAt     string   `json:"expires_at"`
	CAFingerprint string   `json:"ca_fingerprint"`
	Principals    []string `json:"principals"`
	Note          string   `json:"note"`
}

// DelegateBeginIn takes no arguments: asking is the whole operation.
type DelegateBeginIn struct{}

// DelegateCompleteIn carries the signed certificate. A certificate is public
// material — see the tool description, which says so, because a model may
// otherwise refuse to paste something that looks like key material.
type DelegateCompleteIn struct {
	Certificate string `json:"certificate" jsonschema:"the contents of the *-cert.pub file your CA produced"`
}

// NewServer builds the MCP server for one identity. Every caller goes through
// here, so the tool list is the same one for everyone.
func NewServer(t *Tools, opts Options) *mcp.Server {
	s := mcp.NewServer(&mcp.Implementation{Name: "eitri", Version: "0.1.0"},
		&mcp.ServerOptions{SchemaCache: opts.SchemaCache})
	register(s, "vm_create",
		"Create an eitri VM. Waits for ready+cloud-init by default. Name a network to also put the guest on one "+
			"of its host's named networks, where that network's own DHCP addresses it: the result reports that "+
			"address as network_ip, which may still be EMPTY when the tool returns — readiness is the guest booting, "+
			"and a lease from the site's DHCP server arrives when it arrives. Read it back with vm_info rather "+
			"than treating an empty one as a failed create.",
		t.VMCreate)
	register(s, "vm_list", "List all VMs on the eitri fleet.", t.VMList)
	register(s, "vm_info", "Show one VM's state and how to reach it.", t.VMInfo)
	register(s, "vm_exec", "Run a shell command in a VM over SSH; returns stdout/stderr/exit code.", t.VMExec)
	register(s, "vm_write_file", "Write content to a file in a VM (parents created).", t.VMWriteFile)
	register(s, "vm_read_file", "Read a file from a VM (capped at 1 MiB).", t.VMReadFile)
	register(s, "vm_expose", "Publish a VM's guest port on its host, TCP or UDP, and return the address to dial. Omit host_port to allocate one from 30000-32767. WARNING: a published port has NO AUTHENTICATION in front of it — whoever can reach the host on that port reaches the service, and a UDP one answers whatever address a datagram claims to come from. Publish only what is meant to be reachable.", t.VMExpose)
	register(s, "vm_exposures", "List a VM's published ports, with the protocol, the address to dial and each socket's state. These ports are unauthenticated.", t.VMExposures)
	register(s, "vm_unexpose", "Stop publishing a VM's guest port; the host closes the socket.", t.VMUnexpose)
	register(s, "vm_destroy", "Destroy a VM by id or EXACT name. Explicit-only; never called automatically.", t.VMDestroy)
	register(s, "ca_upload",
		"Register your SSH user CA's PUBLIC key with your tenant, so your guests trust certificates it signs. "+
			"Only the public half is sent — eitri never holds a signing key. "+
			"IMPORTANT: a guest bakes its trusted CA set when it is created, so VMs that already exist will NOT "+
			"trust a CA uploaded now; create new VMs after uploading. Do this before vm_create, and before delegating.",
		t.CAUpload)
	register(s, "tenant_info",
		"Show how this tenant is set up: which SSH CAs are registered (fingerprint and label), whether eitri "+
			"currently holds a delegation and when it expires, and the gate address. Read-only. Call it first "+
			"when an SSH operation fails, instead of guessing which step was skipped.",
		t.TenantInfo)
	dg := opts.Delegator
	register(s, "delegate_begin",
		"Ask for the public key your CA is to sign, so eitri can reach your VMs. eitri holds no signing key and "+
			"cannot sign this itself — that is the point. Show the human the public key and the command, and wait "+
			"for them to hand back the certificate; do not try to produce it yourself. The key stays the same "+
			"until the control plane restarts, and a begin nobody signs is abandoned after an hour — a "+
			"delegation is held in memory, so call this tool again for the current key rather than reusing an "+
			"old certificate.",
		func(ctx context.Context, _ DelegateBeginIn) (BeginResult, error) { return dg.Begin(ctx) })
	register(s, "delegate_complete",
		"Hand back the certificate your CA signed. Certificates are public material, so passing one as an "+
			"argument is safe — it is not a key and grants nothing without the key eitri holds in memory. "+
			"Afterwards eitri can reach your VMs until the certificate expires. Because it chains to a CA your "+
			"tenant already registered, VMs created before this call accept it too.",
		func(ctx context.Context, in DelegateCompleteIn) (DelegationResult, error) {
			return dg.Complete(ctx, in.Certificate)
		})
	return s
}

// progressKey carries the per-request progress reporter.
type progressKey struct{}

// withProgress returns a context whose long waits report to report. A transport
// that cannot deliver notifications simply does not install one.
func withProgress(ctx context.Context, report func(message string)) context.Context {
	return context.WithValue(ctx, progressKey{}, report)
}

// reportProgress tells the caller a long-running tool is still working. It is a
// no-op when nothing is listening, which is every case but a remote tool call
// whose client asked for progress.
//
// It is not only courtesy: vm_create blocks for up to ten minutes and emits
// nothing until it finishes, and a proxy in front of the control plane will cut
// an origin connection that stays silent for a hundred seconds.
func reportProgress(ctx context.Context, message string) {
	if report, ok := ctx.Value(progressKey{}).(func(string)); ok && report != nil {
		report(message)
	}
}

// register adapts a Tools method to the SDK. This is the ONLY place that
// touches SDK generics; if the SDK's handler signature changes, change it here.
//
// Note on the hand-off contract: the SDK drops the Out value when the handler
// returns a non-nil error — StructuredContent is left unset and only err.Error()
// reaches the model (as IsError text content). VMCreate's degraded-path errors
// are self-sufficient (they name the VM id+name), so the model can still find
// and destroy the VM from the error text.
func register[In, Out any](s *mcp.Server, name, desc string, fn func(context.Context, In) (Out, error)) {
	mcp.AddTool(s, &mcp.Tool{Name: name, Description: desc},
		func(ctx context.Context, req *mcp.CallToolRequest, in In) (*mcp.CallToolResult, Out, error) {
			out, err := fn(withProgress(ctx, progressReporter(ctx, req)), in)
			return nil, out, err
		})
}

// progressReporter builds the notifier for one tool call, or nil when the
// client did not ask for progress. The MCP contract is that progress is
// reported only against a token the client supplied, so a client that sent none
// gets none.
func progressReporter(ctx context.Context, req *mcp.CallToolRequest) func(string) {
	if req == nil || req.Session == nil || req.Params == nil {
		return nil
	}
	token := req.Params.GetProgressToken()
	if token == nil {
		return nil
	}
	session := req.Session
	step := 0.0
	return func(message string) {
		step++
		// Best-effort: a client that has gone away must not fail the operation
		// it asked for.
		_ = session.NotifyProgress(ctx, &mcp.ProgressNotificationParams{
			ProgressToken: token, Message: message, Progress: step,
		})
	}
}