package mcpserver import ( "bytes" "context" "encoding/json" "errors" "io" "testing" "time" "email-mcp/internal/imapclient" "email-mcp/internal/secretstore" "email-mcp/internal/secretstore/kwallet" ) type storeStub struct { credential secretstore.Credential loadErr error loadCalls int loadKey string } func (s *storeStub) Save(context.Context, string, secretstore.Credential) error { return nil } func (s *storeStub) Load(_ context.Context, key string) (secretstore.Credential, error) { s.loadCalls++ s.loadKey = key if s.loadErr != nil { return secretstore.Credential{}, s.loadErr } return s.credential, nil } type serviceStub struct { listMailboxes func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) listMessages func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) getMessage func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) } func (s serviceStub) ListMailboxes(ctx context.Context, cred secretstore.Credential) ([]imapclient.Mailbox, error) { return s.listMailboxes(ctx, cred) } func (s serviceStub) ListMessages(ctx context.Context, cred secretstore.Credential, mailbox string, limit int) ([]imapclient.MessageSummary, error) { return s.listMessages(ctx, cred, mailbox, limit) } func (s serviceStub) GetMessage(ctx context.Context, cred secretstore.Credential, mailbox string, uid uint32) (imapclient.Message, error) { return s.getMessage(ctx, cred, mailbox, uid) } func TestServerListMailboxesLoadsCredentialAndDelegates(t *testing.T) { store := &storeStub{ credential: secretstore.Credential{ Host: "imap.example.com", Username: "alice", Password: "secret", }, } server := New(store, serviceStub{ listMailboxes: func(_ context.Context, cred secretstore.Credential) ([]imapclient.Mailbox, error) { if cred.Host != "imap.example.com" || cred.Username != "alice" || cred.Password != "secret" { t.Fatalf("unexpected credential: %#v", cred) } return []imapclient.Mailbox{{Name: "INBOX"}}, nil }, listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) { t.Fatal("ListMessages should not be called") return nil, nil }, getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) { t.Fatal("GetMessage should not be called") return imapclient.Message{}, nil }, }) result, err := server.ListMailboxes(context.Background()) if err != nil { t.Fatalf("ListMailboxes returned error: %v", err) } if store.loadCalls != 1 { t.Fatalf("expected credential to be loaded once, got %d", store.loadCalls) } if store.loadKey != secretstore.DefaultAccountKey { t.Fatalf("expected load key %q, got %q", secretstore.DefaultAccountKey, store.loadKey) } if len(result) != 1 || result[0].Name != "INBOX" { t.Fatalf("unexpected result: %#v", result) } } func TestServerListMessagesLoadsCredentialAndDelegates(t *testing.T) { store := &storeStub{ credential: secretstore.Credential{ Host: "imap.example.com", Username: "alice", Password: "secret", }, } server := New(store, serviceStub{ listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) { t.Fatal("ListMailboxes should not be called") return nil, nil }, listMessages: func(_ context.Context, cred secretstore.Credential, mailbox string, limit int) ([]imapclient.MessageSummary, error) { if cred.Host != "imap.example.com" || mailbox != "INBOX" || limit != 5 { t.Fatalf("unexpected call: cred=%#v mailbox=%q limit=%d", cred, mailbox, limit) } return []imapclient.MessageSummary{{UID: 42, Subject: "hello", From: "alice@example.com"}}, nil }, getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) { t.Fatal("GetMessage should not be called") return imapclient.Message{}, nil }, }) result, err := server.ListMessages(context.Background(), "INBOX", 5) if err != nil { t.Fatalf("ListMessages returned error: %v", err) } if len(result) != 1 || result[0].UID != 42 { t.Fatalf("unexpected result: %#v", result) } } func TestServerGetMessageUsesUIDContract(t *testing.T) { store := &storeStub{ credential: secretstore.Credential{ Host: "imap.example.com", Username: "alice", Password: "secret", }, } server := New(store, serviceStub{ listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) { t.Fatal("ListMailboxes should not be called") return nil, nil }, listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) { t.Fatal("ListMessages should not be called") return nil, nil }, getMessage: func(_ context.Context, cred secretstore.Credential, mailbox string, uid uint32) (imapclient.Message, error) { if cred.Host != "imap.example.com" || mailbox != "INBOX" || uid != 42 { t.Fatalf("unexpected call: cred=%#v mailbox=%q uid=%d", cred, mailbox, uid) } return imapclient.Message{ UID: 42, Mailbox: "INBOX", Body: "body", }, nil }, }) message, err := server.GetMessage(context.Background(), "INBOX", 42) if err != nil { t.Fatalf("GetMessage returned error: %v", err) } if message.UID != 42 || message.Mailbox != "INBOX" { t.Fatalf("unexpected message: %#v", message) } } func TestRunnerRunWritesToolManifestAndHandlesRequests(t *testing.T) { store := &storeStub{ credential: secretstore.Credential{ Host: "imap.example.com", Username: "alice", Password: "secret", }, } input := bytes.NewBufferString("{\"tool\":\"list_messages\",\"arguments\":{\"mailbox\":\"INBOX\",\"limit\":5}}\n") output := &bytes.Buffer{} runner := NewRunner(New(store, serviceStub{ listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) { t.Fatal("ListMailboxes should not be called") return nil, nil }, listMessages: func(_ context.Context, cred secretstore.Credential, mailbox string, limit int) ([]imapclient.MessageSummary, error) { if cred.Host != "imap.example.com" || mailbox != "INBOX" || limit != 5 { t.Fatalf("unexpected call: cred=%#v mailbox=%q limit=%d", cred, mailbox, limit) } return []imapclient.MessageSummary{{UID: 42, Subject: "hello", From: "alice@example.com"}}, nil }, getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) { t.Fatal("GetMessage should not be called") return imapclient.Message{}, nil }, }), input, output, &bytes.Buffer{}) if err := runner.Run(context.Background()); err != nil { t.Fatalf("Run returned error: %v", err) } if store.loadCalls != 1 { t.Fatalf("expected credential preload once, got %d", store.loadCalls) } decoder := json.NewDecoder(output) var manifest struct { Tools []struct { Name string `json:"name"` } `json:"tools"` } if err := decoder.Decode(&manifest); err != nil { t.Fatalf("failed to decode manifest: %v", err) } if len(manifest.Tools) != 3 { t.Fatalf("expected 3 tools, got %#v", manifest.Tools) } if manifest.Tools[0].Name != "list_mailboxes" || manifest.Tools[1].Name != "list_messages" || manifest.Tools[2].Name != "get_message" { t.Fatalf("unexpected tool manifest: %#v", manifest.Tools) } var response struct { Result []imapclient.MessageSummary `json:"result"` } if err := decoder.Decode(&response); err != nil { t.Fatalf("failed to decode response: %v", err) } if len(response.Result) != 1 || response.Result[0].UID != 42 { t.Fatalf("unexpected response: %#v", response.Result) } } func TestRunnerRunReturnsFriendlyMissingCredentialError(t *testing.T) { store := &storeStub{ loadErr: kwallet.ErrCredentialNotFound, } runner := NewRunner(New(store, serviceStub{ listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) { t.Fatal("ListMailboxes should not be called") return nil, nil }, listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) { t.Fatal("ListMessages should not be called") return nil, nil }, getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) { t.Fatal("GetMessage should not be called") return imapclient.Message{}, nil }, }), bytes.NewBuffer(nil), &bytes.Buffer{}, &bytes.Buffer{}) err := runner.Run(context.Background()) if !errors.Is(err, ErrCredentialsNotConfigured) { t.Fatalf("expected missing credential error, got %v", err) } } func TestServerToolsAdvertiseValidatedArgumentContracts(t *testing.T) { tools := New(&storeStub{}, serviceStub{}).Tools() if len(tools) != 3 { t.Fatalf("expected 3 tools, got %d", len(tools)) } listMessages := tools[1] if listMessages.Name != "list_messages" { t.Fatalf("unexpected tool ordering: %#v", tools) } if got := listMessages.InputSchema["type"]; got != "object" { t.Fatalf("expected object schema, got %#v", got) } listProps, ok := listMessages.InputSchema["properties"].(map[string]any) if !ok { t.Fatalf("expected properties map, got %#v", listMessages.InputSchema["properties"]) } limitSchema, ok := listProps["limit"].(map[string]any) if !ok { t.Fatalf("expected limit schema, got %#v", listProps["limit"]) } mailboxSchema, ok := listProps["mailbox"].(map[string]any) if !ok { t.Fatalf("expected mailbox schema, got %#v", listProps["mailbox"]) } if got := mailboxSchema["pattern"]; got != "\\S" { t.Fatalf("expected mailbox pattern %q, got %#v", "\\S", got) } if got := limitSchema["default"]; got != float64(imapclient.DefaultListMessagesLimit) && got != imapclient.DefaultListMessagesLimit { t.Fatalf("expected limit default %d, got %#v", imapclient.DefaultListMessagesLimit, got) } if got := limitSchema["minimum"]; got != float64(1) && got != 1 { t.Fatalf("expected limit minimum 1, got %#v", got) } if got := limitSchema["maximum"]; got != float64(imapclient.MaxListMessagesLimit) && got != imapclient.MaxListMessagesLimit { t.Fatalf("expected limit maximum %d, got %#v", imapclient.MaxListMessagesLimit, got) } getMessage := tools[2] if getMessage.Name != "get_message" { t.Fatalf("unexpected tool ordering: %#v", tools) } getProps, ok := getMessage.InputSchema["properties"].(map[string]any) if !ok { t.Fatalf("expected get_message properties map, got %#v", getMessage.InputSchema["properties"]) } getMailboxSchema, ok := getProps["mailbox"].(map[string]any) if !ok { t.Fatalf("expected get_message mailbox schema, got %#v", getProps["mailbox"]) } if got := getMailboxSchema["pattern"]; got != "\\S" { t.Fatalf("expected mailbox pattern %q, got %#v", "\\S", got) } uidSchema, ok := getProps["uid"].(map[string]any) if !ok { t.Fatalf("expected uid schema, got %#v", getProps["uid"]) } if got := uidSchema["minimum"]; got != float64(1) && got != 1 { t.Fatalf("expected uid minimum 1, got %#v", got) } } func TestRunnerRunReturnsValidationErrorsForInvalidRequests(t *testing.T) { store := &storeStub{ credential: secretstore.Credential{ Host: "imap.example.com", Username: "alice", Password: "secret", }, } input := bytes.NewBufferString("{\"tool\":\"list_messages\",\"arguments\":{\"mailbox\":\"INBOX\",\"limit\":0}}\n{\"tool\":\"get_message\",\"arguments\":{\"mailbox\":\"INBOX\"}}\n") output := &bytes.Buffer{} runner := NewRunner(New(store, serviceStub{ listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) { t.Fatal("ListMailboxes should not be called") return nil, nil }, listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) { t.Fatal("ListMessages should not be called") return nil, nil }, getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) { t.Fatal("GetMessage should not be called") return imapclient.Message{}, nil }, }), input, output, &bytes.Buffer{}) if err := runner.Run(context.Background()); err != nil { t.Fatalf("Run returned error: %v", err) } decoder := json.NewDecoder(output) if err := decoder.Decode(&struct { Tools []Tool `json:"tools"` }{}); err != nil { t.Fatalf("failed to decode manifest: %v", err) } var firstResponse struct { Error string `json:"error"` } if err := decoder.Decode(&firstResponse); err != nil { t.Fatalf("failed to decode first error response: %v", err) } if firstResponse.Error != "limit must be between 1 and 50" { t.Fatalf("unexpected first error: %#v", firstResponse) } var secondResponse struct { Error string `json:"error"` } if err := decoder.Decode(&secondResponse); err != nil { t.Fatalf("failed to decode second error response: %v", err) } if secondResponse.Error != "uid is required" { t.Fatalf("unexpected second error: %#v", secondResponse) } } func TestRunnerRunRejectsWhitespaceOnlyMailboxValues(t *testing.T) { store := &storeStub{ credential: secretstore.Credential{ Host: "imap.example.com", Username: "alice", Password: "secret", }, } input := bytes.NewBufferString("{\"tool\":\"list_messages\",\"arguments\":{\"mailbox\":\" \"}}\n") output := &bytes.Buffer{} runner := NewRunner(New(store, serviceStub{ listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) { t.Fatal("ListMailboxes should not be called") return nil, nil }, listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) { t.Fatal("ListMessages should not be called") return nil, nil }, getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) { t.Fatal("GetMessage should not be called") return imapclient.Message{}, nil }, }), input, output, &bytes.Buffer{}) if err := runner.Run(context.Background()); err != nil { t.Fatalf("Run returned error: %v", err) } decoder := json.NewDecoder(output) if err := decoder.Decode(&struct { Tools []Tool `json:"tools"` }{}); err != nil { t.Fatalf("failed to decode manifest: %v", err) } var response struct { Error string `json:"error"` } if err := decoder.Decode(&response); err != nil { t.Fatalf("failed to decode error response: %v", err) } if response.Error != "mailbox is required" { t.Fatalf("unexpected error: %#v", response) } } func TestRunnerRunAppliesDefaultLimitWhenOmitted(t *testing.T) { store := &storeStub{ credential: secretstore.Credential{ Host: "imap.example.com", Username: "alice", Password: "secret", }, } input := bytes.NewBufferString("{\"tool\":\"list_messages\",\"arguments\":{\"mailbox\":\"INBOX\"}}\n") output := &bytes.Buffer{} runner := NewRunner(New(store, serviceStub{ listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) { t.Fatal("ListMailboxes should not be called") return nil, nil }, listMessages: func(_ context.Context, cred secretstore.Credential, mailbox string, limit int) ([]imapclient.MessageSummary, error) { if cred.Host != "imap.example.com" || mailbox != "INBOX" || limit != imapclient.DefaultListMessagesLimit { t.Fatalf("unexpected call: cred=%#v mailbox=%q limit=%d", cred, mailbox, limit) } return []imapclient.MessageSummary{{UID: 42, Subject: "hello", From: "alice@example.com"}}, nil }, getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) { t.Fatal("GetMessage should not be called") return imapclient.Message{}, nil }, }), input, output, &bytes.Buffer{}) if err := runner.Run(context.Background()); err != nil { t.Fatalf("Run returned error: %v", err) } } func TestRunnerRunStopsWhenContextCanceledWhileWaitingForInput(t *testing.T) { store := &storeStub{ credential: secretstore.Credential{ Host: "imap.example.com", Username: "alice", Password: "secret", }, } input := newBlockingReadCloser() runner := NewRunner(New(store, serviceStub{ listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) { t.Fatal("ListMailboxes should not be called") return nil, nil }, listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) { t.Fatal("ListMessages should not be called") return nil, nil }, getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) { t.Fatal("GetMessage should not be called") return imapclient.Message{}, nil }, }), input, &bytes.Buffer{}, &bytes.Buffer{}) ctx, cancel := context.WithCancel(context.Background()) defer cancel() done := make(chan error, 1) go func() { done <- runner.Run(ctx) }() select { case <-input.started: case <-time.After(200 * time.Millisecond): t.Fatal("runner never started reading input") } cancel() select { case err := <-done: if !errors.Is(err, context.Canceled) { t.Fatalf("expected context cancellation, got %v", err) } case <-time.After(500 * time.Millisecond): t.Fatal("runner did not stop after context cancellation") } if !input.closed { t.Fatal("expected runner to close input reader on cancellation") } } type blockingReadCloser struct { started chan struct{} closed bool done chan struct{} } func newBlockingReadCloser() *blockingReadCloser { return &blockingReadCloser{ started: make(chan struct{}), done: make(chan struct{}), } } func (r *blockingReadCloser) Read(_ []byte) (int, error) { select { case <-r.started: default: close(r.started) } <-r.done return 0, io.EOF } func (r *blockingReadCloser) Close() error { if !r.closed { r.closed = true close(r.done) } return nil }