email-mcp/internal/mcpserver/server.go

package mcpserver

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strings"

	"email-mcp/internal/imapclient"
	"email-mcp/internal/secretstore"
	"email-mcp/internal/secretstore/kwallet"
)

var ErrCredentialsNotConfigured = errors.New("credentials not configured; run `email-mcp setup`")

type MailService interface {
	ListMailboxes(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error)
	ListMessages(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error)
	GetMessage(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error)
}

type Server struct {
	store secretstore.Store
	mail  MailService
}

type Tool struct {
	Name        string         `json:"name"`
	Description string         `json:"description"`
	InputSchema map[string]any `json:"input_schema,omitempty"`
}

type Runner struct {
	server Server
	in     io.Reader
	out    io.Writer
	errOut io.Writer
}

type toolRequest struct {
	Tool      string          `json:"tool"`
	Arguments json.RawMessage `json:"arguments,omitempty"`
}

type listMessagesArguments struct {
	Mailbox string `json:"mailbox"`
	Limit   *int   `json:"limit,omitempty"`
}

type getMessageArguments struct {
	Mailbox string  `json:"mailbox"`
	UID     *uint32 `json:"uid"`
}

func New(store secretstore.Store, mail MailService) Server {
	return Server{
		store: store,
		mail:  mail,
	}
}

func NewRunner(server Server, in io.Reader, out io.Writer, errOut io.Writer) Runner {
	if out == nil {
		out = io.Discard
	}
	if errOut == nil {
		errOut = io.Discard
	}

	return Runner{
		server: server,
		in:     in,
		out:    out,
		errOut: errOut,
	}
}

func (s Server) Tools() []Tool {
	return []Tool{
		{
			Name:        "list_mailboxes",
			Description: "List visible IMAP mailboxes for the configured account.",
		},
		{
			Name:        "list_messages",
			Description: "List recent messages from a mailbox using IMAP UIDs.",
			InputSchema: map[string]any{
				"type": "object",
				"properties": map[string]any{
					"mailbox": map[string]any{
						"type":      "string",
						"minLength": 1,
						"pattern":   "\\S",
					},
					"limit": map[string]any{
						"type":    "integer",
						"default": imapclient.DefaultListMessagesLimit,
						"minimum": 1,
						"maximum": imapclient.MaxListMessagesLimit,
					},
				},
				"required":             []string{"mailbox"},
				"additionalProperties": false,
			},
		},
		{
			Name:        "get_message",
			Description: "Fetch a single message by mailbox and IMAP UID.",
			InputSchema: map[string]any{
				"type": "object",
				"properties": map[string]any{
					"mailbox": map[string]any{
						"type":      "string",
						"minLength": 1,
						"pattern":   "\\S",
					},
					"uid": map[string]any{
						"type":    "integer",
						"minimum": 1,
					},
				},
				"required":             []string{"mailbox", "uid"},
				"additionalProperties": false,
			},
		},
	}
}

func (s Server) ListMailboxes(ctx context.Context) ([]imapclient.Mailbox, error) {
	cred, err := s.loadCredential(ctx)
	if err != nil {
		return nil, err
	}
	return s.listMailboxes(ctx, cred)
}

func (s Server) ListMessages(ctx context.Context, mailbox string, limit int) ([]imapclient.MessageSummary, error) {
	cred, err := s.loadCredential(ctx)
	if err != nil {
		return nil, err
	}
	return s.listMessages(ctx, cred, mailbox, limit)
}

func (s Server) GetMessage(ctx context.Context, mailbox string, uid uint32) (imapclient.Message, error) {
	cred, err := s.loadCredential(ctx)
	if err != nil {
		return imapclient.Message{}, err
	}
	return s.getMessage(ctx, cred, mailbox, uid)
}

func (r Runner) Run(ctx context.Context) error {
	cred, err := r.server.loadCredential(ctx)
	if err != nil {
		return err
	}

	stopCancelRead := r.closeInputOnCancel(ctx)
	defer stopCancelRead()

	encoder := json.NewEncoder(r.out)
	if err := encoder.Encode(map[string]any{"tools": r.server.Tools()}); err != nil {
		return err
	}
	if r.in == nil {
		return nil
	}

	decoder := json.NewDecoder(r.in)
	for {
		if err := ctx.Err(); err != nil {
			return err
		}

		var request toolRequest
		if err := decoder.Decode(&request); err != nil {
			if ctxErr := ctx.Err(); ctxErr != nil {
				return ctxErr
			}
			if errors.Is(err, io.EOF) {
				return nil
			}
			return err
		}

		result, err := r.server.handleTool(ctx, cred, request.Tool, request.Arguments)
		if err != nil {
			if encodeErr := encoder.Encode(map[string]any{"error": err.Error()}); encodeErr != nil {
				return encodeErr
			}
			continue
		}
		if err := encoder.Encode(map[string]any{"result": result}); err != nil {
			return err
		}
	}
}

func (s Server) handleTool(ctx context.Context, cred secretstore.Credential, name string, rawArgs json.RawMessage) (any, error) {
	switch name {
	case "list_mailboxes":
		return s.listMailboxes(ctx, cred)
	case "list_messages":
		var args listMessagesArguments
		if err := decodeArguments(rawArgs, &args); err != nil {
			return nil, err
		}
		limit, err := normalizeListMessagesLimit(args.Limit)
		if err != nil {
			return nil, err
		}
		return s.listMessages(ctx, cred, args.Mailbox, limit)
	case "get_message":
		var args getMessageArguments
		if err := decodeArguments(rawArgs, &args); err != nil {
			return nil, err
		}
		uid, err := validateMessageUID(args.UID)
		if err != nil {
			return nil, err
		}
		return s.getMessage(ctx, cred, args.Mailbox, uid)
	default:
		return nil, fmt.Errorf("unknown tool: %s", name)
	}
}

func (s Server) listMailboxes(ctx context.Context, cred secretstore.Credential) ([]imapclient.Mailbox, error) {
	if s.mail == nil {
		return nil, fmt.Errorf("mail service is not configured")
	}
	return s.mail.ListMailboxes(ctx, cred)
}

func (s Server) listMessages(ctx context.Context, cred secretstore.Credential, mailbox string, limit int) ([]imapclient.MessageSummary, error) {
	if s.mail == nil {
		return nil, fmt.Errorf("mail service is not configured")
	}
	mailbox, err := validateMailbox(mailbox)
	if err != nil {
		return nil, err
	}
	return s.mail.ListMessages(ctx, cred, mailbox, limit)
}

func (s Server) getMessage(ctx context.Context, cred secretstore.Credential, mailbox string, uid uint32) (imapclient.Message, error) {
	if s.mail == nil {
		return imapclient.Message{}, fmt.Errorf("mail service is not configured")
	}
	mailbox, err := validateMailbox(mailbox)
	if err != nil {
		return imapclient.Message{}, err
	}
	if uid == 0 {
		return imapclient.Message{}, fmt.Errorf("uid must be greater than zero")
	}
	return s.mail.GetMessage(ctx, cred, mailbox, uid)
}

func (s Server) loadCredential(ctx context.Context) (secretstore.Credential, error) {
	if s.store == nil {
		return secretstore.Credential{}, fmt.Errorf("secret store is not configured")
	}

	cred, err := s.store.Load(ctx, secretstore.DefaultAccountKey)
	if err != nil {
		if errors.Is(err, kwallet.ErrCredentialNotFound) || errors.Is(err, ErrCredentialsNotConfigured) {
			return secretstore.Credential{}, ErrCredentialsNotConfigured
		}
		return secretstore.Credential{}, err
	}
	if err := cred.Validate(); err != nil {
		return secretstore.Credential{}, fmt.Errorf("default credential is invalid: %w", err)
	}
	return cred, nil
}

func decodeArguments(raw json.RawMessage, dest any) error {
	if len(raw) == 0 {
		raw = []byte("{}")
	}
	decoder := json.NewDecoder(strings.NewReader(string(raw)))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(dest); err != nil {
		return fmt.Errorf("invalid tool arguments: %w", err)
	}
	return nil
}

func validateMailbox(mailbox string) (string, error) {
	mailbox = strings.TrimSpace(mailbox)
	if mailbox == "" {
		return "", fmt.Errorf("mailbox is required")
	}
	return mailbox, nil
}

func normalizeListMessagesLimit(limit *int) (int, error) {
	if limit == nil {
		return imapclient.DefaultListMessagesLimit, nil
	}
	if *limit < 1 || *limit > imapclient.MaxListMessagesLimit {
		return 0, fmt.Errorf("limit must be between 1 and %d", imapclient.MaxListMessagesLimit)
	}
	return *limit, nil
}

func validateMessageUID(uid *uint32) (uint32, error) {
	if uid == nil {
		return 0, fmt.Errorf("uid is required")
	}
	if *uid == 0 {
		return 0, fmt.Errorf("uid must be greater than zero")
	}
	return *uid, nil
}

func (r Runner) closeInputOnCancel(ctx context.Context) func() {
	closer, ok := r.in.(io.Closer)
	if !ok {
		return func() {}
	}

	done := make(chan struct{})
	go func() {
		select {
		case <-ctx.Done():
			_ = closer.Close()
		case <-done:
		}
	}()

	return func() {
		close(done)
	}
}