mcp credentials over chat (#1490)

This commit is contained in:
bilogic authored and GitHub committed 2026-07-31 19:34:40 +07:00
1 parent bcd23c587b
commit 2e1b90ca35
5 files changed
+1228 -1

No files matched your search

+8
View File
@@ -884,6 +884,14 @@ func runGateway() {
}
}
}
// Wire MCP server store on mcp_credential_manager tool.
if pgStores != nil && pgStores.MCP != nil {
if t, ok := toolsReg.Get("mcp_credential_manager"); ok {
if ms, ok := t.(tools.MCPServerStoreAware); ok {
ms.SetMCPServerStore(pgStores.MCP)
}
}
}
// Load channel instances from DB.
var instanceLoader *channels.InstanceLoader
+3 -1
View File
@@ -73,7 +73,9 @@ func wireExtraTools(
// create_forum_topic is kept as a backward-compatible wrapper for topic.create.
toolsReg.Register(tools.NewCreateForumTopicTool(nil))
toolsReg.Register(tools.NewTelegramManagerTool())
slog.Info("session + message + send_file + telegram_manager tools registered")
// MCP credential manager tool (view and manage per-user MCP credentials)
toolsReg.Register(tools.NewMCPCredentialManagerTool())
slog.Info("session + message + send_file + telegram_manager + mcp_credential_manager tools registered")
// Register legacy tool aliases (backward-compat names from policy.go).
for alias, canonical := range tools.LegacyToolAliases() {
+376
View File
@@ -0,0 +1,376 @@
package tools
import (
"context"
"encoding/json"
"fmt"
"strings"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
// MCPCredentialManagerTool allows users to view and manage their MCP server credentials
// from Telegram and other channels. Supports listing accessible servers, checking
// credential status, setting credentials (API key or Bearer token), and deleting credentials.
type MCPCredentialManagerTool struct {
mcpStore store.MCPServerStore
}
func NewMCPCredentialManagerTool() *MCPCredentialManagerTool {
return &MCPCredentialManagerTool{}
}
func (t *MCPCredentialManagerTool) SetMCPServerStore(s store.MCPServerStore) {
t.mcpStore = s
}
func (t *MCPCredentialManagerTool) Name() string {
return "mcp_credential_manager"
}
func (t *MCPCredentialManagerTool) Description() string {
return "View and manage your MCP server credentials. Lets you list MCP servers accessible to you, check whether you have credentials configured, set new credentials (API key, Bearer token, headers, environment variables), or delete existing credentials for an MCP server."
}
func (t *MCPCredentialManagerTool) Parameters() map[string]any {
return map[string]any{
"type": "object",
"properties": map[string]any{
"action": map[string]any{
"type": "string",
"description": "Action to perform: list_servers (list accessible MCP servers and credential status), credential_status (check if you have credentials for a server), set_credentials (set/update your API key, headers, and/or env vars for a server), set_bearer_token (set a Bearer token as your credential for a server), delete_credentials (remove your credentials for a server).",
"enum": []string{"list_servers", "credential_status", "set_credentials", "set_bearer_token", "delete_credentials"},
},
"server_name": map[string]any{
"type": "string",
"description": "Name of the MCP server (required for credential_status, set_credentials, set_bearer_token, delete_credentials).",
},
"token": map[string]any{
"type": "string",
"description": "Bearer token for the MCP server (only for set_bearer_token action). The token is stored and sent as an Authorization: Bearer <token> header.",
},
"api_key": map[string]any{
"type": "string",
"description": "API key for the MCP server (only for set_credentials action).",
},
"headers": map[string]any{
"type": "object",
"description": "Optional additional HTTP headers as a JSON object (only for set_credentials action).",
"additionalProperties": map[string]any{"type": "string"},
},
"env": map[string]any{
"type": "object",
"description": "Optional environment variables as a JSON object (only for set_credentials action).",
"additionalProperties": map[string]any{"type": "string"},
},
},
"required": []string{"action"},
}
}
func (t *MCPCredentialManagerTool) Execute(ctx context.Context, args map[string]any) *Result {
if t.mcpStore == nil {
return ErrorResult("mcp_credential_manager: MCP server store not available")
}
action := argString(args, "action")
if action == "" {
return ErrorResult("action is required (list_servers, credential_status, set_credentials, delete_credentials)")
}
switch action {
case "list_servers":
return t.listServers(ctx)
case "credential_status":
return t.credentialStatus(ctx, args)
case "set_credentials":
return t.setCredentials(ctx, args)
case "set_bearer_token":
return t.setBearerToken(ctx, args)
case "delete_credentials":
return t.deleteCredentials(ctx, args)
default:
return ErrorResult(fmt.Sprintf("unsupported action: %s (use list_servers, credential_status, set_credentials, set_bearer_token, delete_credentials)", action))
}
}
func (t *MCPCredentialManagerTool) listServers(ctx context.Context) *Result {
agentID := store.AgentIDFromContext(ctx)
userID := store.UserIDFromContext(ctx)
accessible, err := t.mcpStore.ListAccessible(ctx, agentID, userID)
if err != nil {
return ErrorResult(fmt.Sprintf("failed to list MCP servers: %v", err))
}
if len(accessible) == 0 {
return NewResult("You don't have any MCP servers accessible.")
}
var lines []string
lines = append(lines, fmt.Sprintf("**MCP Servers Accessible to You (%d):**", len(accessible)))
lines = append(lines, "")
for _, info := range accessible {
srv := info.Server
needsCreds := srv.RequireUserCredentials || requireUserCredsFromSettings(srv.Settings)
// Check if user has credentials for this server
hasCreds := false
if userID != "" {
creds, err := t.mcpStore.GetUserCredentials(ctx, srv.ID, userID)
if err == nil && creds != nil {
hasCreds = creds.APIKey != "" || len(creds.Headers) > 0 || len(creds.Env) > 0
}
}
displayName := srv.Name
if srv.DisplayName != "" {
displayName = srv.DisplayName
}
credsStatus := "✅ credentials set"
if needsCreds && !hasCreds {
credsStatus = "⚠️ credentials required - not set"
} else if !needsCreds && !hasCreds {
credsStatus = "no credentials needed"
} else if !needsCreds && hasCreds {
credsStatus = "✅ custom credentials set"
}
lines = append(lines, fmt.Sprintf("- **%s** (%s) - %s", displayName, srv.Name, credsStatus))
}
return NewResult(strings.Join(lines, "\n"))
}
func (t *MCPCredentialManagerTool) credentialStatus(ctx context.Context, args map[string]any) *Result {
serverName := argString(args, "server_name")
if serverName == "" {
return ErrorResult("server_name is required for credential_status action")
}
userID := store.UserIDFromContext(ctx)
if userID == "" {
return ErrorResult("no user identity available in context")
}
server, err := t.mcpStore.GetServerByName(ctx, serverName)
if err != nil {
return ErrorResult(fmt.Sprintf("MCP server %q not found: %v", serverName, err))
}
creds, err := t.mcpStore.GetUserCredentials(ctx, server.ID, userID)
if err != nil {
return ErrorResult(fmt.Sprintf("failed to check credentials: %v", err))
}
needsCreds := server.RequireUserCredentials || requireUserCredsFromSettings(server.Settings)
displayName := server.Name
if server.DisplayName != "" {
displayName = server.DisplayName
}
var lines []string
lines = append(lines, fmt.Sprintf("**MCP Server: %s**", displayName))
lines = append(lines, fmt.Sprintf(" Server name: `%s`", server.Name))
if needsCreds {
lines = append(lines, " This server requires per-user credentials.")
} else {
lines = append(lines, " This server does not require per-user credentials (server-level credentials are configured).")
}
if creds == nil || (creds.APIKey == "" && len(creds.Headers) == 0 && len(creds.Env) == 0) {
lines = append(lines, "")
lines = append(lines, "**You have not set any credentials for this server.**")
if needsCreds {
lines = append(lines, "Use the `set_credentials` action to configure your API key.")
}
} else {
lines = append(lines, "")
lines = append(lines, "**Your credentials:**")
if creds.APIKey != "" {
lines = append(lines, fmt.Sprintf(" - API key: `%s`", maskString(creds.APIKey)))
}
if len(creds.Headers) > 0 {
headerKeys := make([]string, 0, len(creds.Headers))
for k := range creds.Headers {
headerKeys = append(headerKeys, k)
}
lines = append(lines, fmt.Sprintf(" - Headers: %s", strings.Join(headerKeys, ", ")))
}
if len(creds.Env) > 0 {
envKeys := make([]string, 0, len(creds.Env))
for k := range creds.Env {
envKeys = append(envKeys, k)
}
lines = append(lines, fmt.Sprintf(" - Env vars: %s", strings.Join(envKeys, ", ")))
}
}
return NewResult(strings.Join(lines, "\n"))
}
func (t *MCPCredentialManagerTool) setCredentials(ctx context.Context, args map[string]any) *Result {
serverName := argString(args, "server_name")
if serverName == "" {
return ErrorResult("server_name is required for set_credentials action")
}
userID := store.UserIDFromContext(ctx)
if userID == "" {
return ErrorResult("no user identity available in context")
}
server, err := t.mcpStore.GetServerByName(ctx, serverName)
if err != nil {
return ErrorResult(fmt.Sprintf("MCP server %q not found: %v", serverName, err))
}
apiKey := argString(args, "api_key")
var headers map[string]string
if rawHeaders, ok := args["headers"].(map[string]any); ok && len(rawHeaders) > 0 {
headers = make(map[string]string, len(rawHeaders))
for k, v := range rawHeaders {
if vs, ok := v.(string); ok {
headers[k] = vs
}
}
}
var env map[string]string
if rawEnv, ok := args["env"].(map[string]any); ok && len(rawEnv) > 0 {
env = make(map[string]string, len(rawEnv))
for k, v := range rawEnv {
if vs, ok := v.(string); ok {
env[k] = vs
}
}
}
if apiKey == "" && len(headers) == 0 && len(env) == 0 {
return ErrorResult("at least one of api_key, headers, or env must be provided")
}
creds := store.MCPUserCredentials{
APIKey: apiKey,
Headers: headers,
Env: env,
}
if err := t.mcpStore.SetUserCredentials(ctx, server.ID, userID, creds); err != nil {
return ErrorResult(fmt.Sprintf("failed to set credentials: %v", err))
}
displayName := server.Name
if server.DisplayName != "" {
displayName = server.DisplayName
}
var parts []string
if apiKey != "" {
parts = append(parts, "API key")
}
if len(headers) > 0 {
parts = append(parts, fmt.Sprintf("%d header(s)", len(headers)))
}
if len(env) > 0 {
parts = append(parts, fmt.Sprintf("%d env var(s)", len(env)))
}
return NewResult(fmt.Sprintf("Successfully set credentials (%s) for MCP server **%s**.", strings.Join(parts, ", "), displayName))
}
func (t *MCPCredentialManagerTool) setBearerToken(ctx context.Context, args map[string]any) *Result {
serverName := argString(args, "server_name")
if serverName == "" {
return ErrorResult("server_name is required for set_bearer_token action")
}
token := argString(args, "token")
if token == "" {
return ErrorResult("token is required for set_bearer_token action")
}
userID := store.UserIDFromContext(ctx)
if userID == "" {
return ErrorResult("no user identity available in context")
}
server, err := t.mcpStore.GetServerByName(ctx, serverName)
if err != nil {
return ErrorResult(fmt.Sprintf("MCP server %q not found: %v", serverName, err))
}
creds := store.MCPUserCredentials{
APIKey: token,
}
if err := t.mcpStore.SetUserCredentials(ctx, server.ID, userID, creds); err != nil {
return ErrorResult(fmt.Sprintf("failed to set Bearer token: %v", err))
}
displayName := server.Name
if server.DisplayName != "" {
displayName = server.DisplayName
}
return NewResult(fmt.Sprintf("Successfully set Bearer token for MCP server **%s**. The token will be sent as an Authorization: Bearer header.", displayName))
}
func (t *MCPCredentialManagerTool) deleteCredentials(ctx context.Context, args map[string]any) *Result {
serverName := argString(args, "server_name")
if serverName == "" {
return ErrorResult("server_name is required for delete_credentials action")
}
userID := store.UserIDFromContext(ctx)
if userID == "" {
return ErrorResult("no user identity available in context")
}
server, err := t.mcpStore.GetServerByName(ctx, serverName)
if err != nil {
return ErrorResult(fmt.Sprintf("MCP server %q not found: %v", serverName, err))
}
if err := t.mcpStore.DeleteUserCredentials(ctx, server.ID, userID); err != nil {
return ErrorResult(fmt.Sprintf("failed to delete credentials: %v", err))
}
displayName := server.Name
if server.DisplayName != "" {
displayName = server.DisplayName
}
return NewResult(fmt.Sprintf("Successfully deleted credentials for MCP server **%s**.", displayName))
}
// maskString returns a masked version of a sensitive string for display.
// Shows first 4 and last 4 characters, masking the middle.
func maskString(s string) string {
if len(s) <= 8 {
return strings.Repeat("*", len(s))
}
return s[:4] + strings.Repeat("*", len(s)-8) + s[len(s)-4:]
}
// requireUserCredsFromSettings checks if an MCP server's settings mandate per-user credentials.
func requireUserCredsFromSettings(settings json.RawMessage) bool {
if len(settings) == 0 {
return false
}
var s struct {
RequireUserCredentials bool `json:"require_user_credentials"`
}
_ = json.Unmarshal(settings, &s)
return s.RequireUserCredentials
}
// Ensure MCPCredentialManagerTool implements MCPServerStoreAware.
var _ MCPServerStoreAware = (*MCPCredentialManagerTool)(nil)
// Ensure MCPCredentialManagerTool implements Tool.
var _ Tool = (*MCPCredentialManagerTool)(nil)
@@ -0,0 +1,836 @@
package tools
import (
"context"
"encoding/json"
"fmt"
"strings"
"sync"
"testing"
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
// mockMCPServerStore implements store.MCPServerStore with in-memory state for tests.
type mockMCPServerStore struct {
mu sync.Mutex
servers map[string]*store.MCPServerData // keyed by name
serverIDs map[uuid.UUID]string // id -> name
credentials map[credKey]*store.MCPUserCredentials
accessible []store.MCPAccessInfo
}
type credKey struct {
serverID uuid.UUID
userID string
}
func newMockMCPServerStore() *mockMCPServerStore {
return &mockMCPServerStore{
servers: make(map[string]*store.MCPServerData),
serverIDs: make(map[uuid.UUID]string),
credentials: make(map[credKey]*store.MCPUserCredentials),
}
}
func (m *mockMCPServerStore) addServer(srv *store.MCPServerData) {
m.mu.Lock()
defer m.mu.Unlock()
m.servers[srv.Name] = srv
m.serverIDs[srv.ID] = srv.Name
}
func (m *mockMCPServerStore) setCredentials(serverID uuid.UUID, userID string, creds *store.MCPUserCredentials) {
m.mu.Lock()
defer m.mu.Unlock()
if creds == nil {
delete(m.credentials, credKey{serverID: serverID, userID: userID})
} else {
m.credentials[credKey{serverID: serverID, userID: userID}] = creds
}
}
// MCPServerStore interface methods used by the tool:
func (m *mockMCPServerStore) ListAccessible(_ context.Context, _ uuid.UUID, _ string) ([]store.MCPAccessInfo, error) {
m.mu.Lock()
defer m.mu.Unlock()
return m.accessible, nil
}
func (m *mockMCPServerStore) GetServerByName(_ context.Context, name string) (*store.MCPServerData, error) {
m.mu.Lock()
defer m.mu.Unlock()
srv, ok := m.servers[name]
if !ok {
return nil, fmt.Errorf("server %q not found", name)
}
return srv, nil
}
func (m *mockMCPServerStore) GetUserCredentials(_ context.Context, serverID uuid.UUID, userID string) (*store.MCPUserCredentials, error) {
m.mu.Lock()
defer m.mu.Unlock()
creds, ok := m.credentials[credKey{serverID: serverID, userID: userID}]
if !ok {
return nil, nil
}
return creds, nil
}
func (m *mockMCPServerStore) SetUserCredentials(_ context.Context, serverID uuid.UUID, userID string, creds store.MCPUserCredentials) error {
m.mu.Lock()
defer m.mu.Unlock()
m.credentials[credKey{serverID: serverID, userID: userID}] = &store.MCPUserCredentials{
APIKey: creds.APIKey,
Headers: creds.Headers,
Env: creds.Env,
}
return nil
}
func (m *mockMCPServerStore) DeleteUserCredentials(_ context.Context, serverID uuid.UUID, userID string) error {
m.mu.Lock()
defer m.mu.Unlock()
delete(m.credentials, credKey{serverID: serverID, userID: userID})
return nil
}
// Stubs for unused MCPServerStore interface methods:
func (m *mockMCPServerStore) CreateServer(_ context.Context, _ *store.MCPServerData) error {
return fmt.Errorf("not implemented")
}
func (m *mockMCPServerStore) GetServer(_ context.Context, _ uuid.UUID) (*store.MCPServerData, error) {
return nil, fmt.Errorf("not implemented")
}
func (m *mockMCPServerStore) ListServers(_ context.Context) ([]store.MCPServerData, error) {
return nil, fmt.Errorf("not implemented")
}
func (m *mockMCPServerStore) UpdateServer(_ context.Context, _ uuid.UUID, _ map[string]any) error {
return fmt.Errorf("not implemented")
}
func (m *mockMCPServerStore) DeleteServer(_ context.Context, _ uuid.UUID) error {
return fmt.Errorf("not implemented")
}
func (m *mockMCPServerStore) GrantToAgent(_ context.Context, _ *store.MCPAgentGrant) error {
return fmt.Errorf("not implemented")
}
func (m *mockMCPServerStore) RevokeFromAgent(_ context.Context, _, _ uuid.UUID) error {
return fmt.Errorf("not implemented")
}
func (m *mockMCPServerStore) ListAgentGrants(_ context.Context, _ uuid.UUID) ([]store.MCPAgentGrant, error) {
return nil, fmt.Errorf("not implemented")
}
func (m *mockMCPServerStore) ListServerGrants(_ context.Context, _ uuid.UUID) ([]store.MCPAgentGrant, error) {
return nil, fmt.Errorf("not implemented")
}
func (m *mockMCPServerStore) GrantToUser(_ context.Context, _ *store.MCPUserGrant) error {
return fmt.Errorf("not implemented")
}
func (m *mockMCPServerStore) RevokeFromUser(_ context.Context, _ uuid.UUID, _ string) error {
return fmt.Errorf("not implemented")
}
func (m *mockMCPServerStore) CountAgentGrantsByServer(_ context.Context) (map[uuid.UUID]int, error) {
return nil, fmt.Errorf("not implemented")
}
func (m *mockMCPServerStore) CreateRequest(_ context.Context, _ *store.MCPAccessRequest) error {
return fmt.Errorf("not implemented")
}
func (m *mockMCPServerStore) ListPendingRequests(_ context.Context) ([]store.MCPAccessRequest, error) {
return nil, fmt.Errorf("not implemented")
}
func (m *mockMCPServerStore) ReviewRequest(_ context.Context, _ uuid.UUID, _ bool, _, _ string) error {
return fmt.Errorf("not implemented")
}
func (m *mockMCPServerStore) CacheToolDescriptions(_ context.Context, _ uuid.UUID, _ map[string]store.CachedToolInfo) error {
return fmt.Errorf("not implemented")
}
// --- Test setup helpers ---
var (
mcpCredTestAgentID = uuid.MustParse("11111111-1111-1111-1111-111111111111")
mcpCredTestUserID = "test-user-42"
mcpCredTestTenantID = uuid.MustParse("22222222-2222-2222-2222-222222222222")
mcpCredServerGitHub = &store.MCPServerData{
BaseModel: store.BaseModel{ID: uuid.MustParse("33333333-3333-3333-3333-333333333333")},
Name: "github",
DisplayName: "GitHub API",
Transport: "streamable-http",
URL: "https://mcp.github.io",
Enabled: true,
}
mcpCredServerPostgres = &store.MCPServerData{
BaseModel: store.BaseModel{ID: uuid.MustParse("44444444-4444-4444-4444-444444444444")},
Name: "postgres",
DisplayName: "",
Transport: "stdio",
Command: "npx",
Args: json.RawMessage(`["@mcp/postgres"]`),
Enabled: true,
}
mcpCredServerSlack = &store.MCPServerData{
BaseModel: store.BaseModel{ID: uuid.MustParse("55555555-5555-5555-5555-555555555555")},
Name: "slack",
DisplayName: "Slack",
Transport: "streamable-http",
URL: "https://mcp.slack.com",
Enabled: true,
RequireUserCredentials: true,
}
mcpCredServerVault = &store.MCPServerData{
BaseModel: store.BaseModel{ID: uuid.MustParse("66666666-6666-6666-6666-666666666666")},
Name: "vault",
Transport: "streamable-http",
URL: "https://vault.example.com",
Enabled: true,
Settings: json.RawMessage(`{"require_user_credentials":true}`),
}
)
func mcpCredCtx() context.Context {
ctx := context.Background()
ctx = store.WithAgentID(ctx, mcpCredTestAgentID)
ctx = store.WithUserID(ctx, mcpCredTestUserID)
ctx = store.WithTenantID(ctx, mcpCredTestTenantID)
return ctx
}
func newTestMCPCredTool() (*mockMCPServerStore, *MCPCredentialManagerTool) {
mock := newMockMCPServerStore()
tool := NewMCPCredentialManagerTool()
tool.SetMCPServerStore(mock)
return mock, tool
}
// --- Tests ---
func TestMCPCredentialManager_NilStore(t *testing.T) {
tool := NewMCPCredentialManagerTool() // store not set
res := tool.Execute(mcpCredCtx(), map[string]any{"action": "list_servers"})
if !res.IsError {
t.Fatal("expected error for nil store")
}
if !strings.Contains(res.ForLLM, "MCP server store not available") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_EmptyAction(t *testing.T) {
_, tool := newTestMCPCredTool()
res := tool.Execute(mcpCredCtx(), map[string]any{})
if !res.IsError {
t.Fatal("expected error for empty action")
}
if !strings.Contains(res.ForLLM, "action is required") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_UnknownAction(t *testing.T) {
_, tool := newTestMCPCredTool()
res := tool.Execute(mcpCredCtx(), map[string]any{"action": "fly_to_moon"})
if !res.IsError {
t.Fatal("expected error for unknown action")
}
if !strings.Contains(res.ForLLM, "unsupported action") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_ListServers_Empty(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.accessible = nil
res := tool.Execute(mcpCredCtx(), map[string]any{"action": "list_servers"})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "don't have any MCP servers") {
t.Fatalf("unexpected result: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_ListServers_AllScenarios(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerGitHub)
mock.addServer(mcpCredServerPostgres)
mock.addServer(mcpCredServerSlack)
mock.addServer(mcpCredServerVault)
mock.accessible = []store.MCPAccessInfo{
{Server: *mcpCredServerGitHub},
{Server: *mcpCredServerPostgres},
{Server: *mcpCredServerSlack},
{Server: *mcpCredServerVault},
}
// Pre-set credentials for Slack (requires user creds, has creds)
mock.setCredentials(mcpCredServerSlack.ID, mcpCredTestUserID, &store.MCPUserCredentials{APIKey: "xoxb-secret-token"})
res := tool.Execute(mcpCredCtx(), map[string]any{"action": "list_servers"})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
// All 4 servers listed
if !strings.Contains(res.ForLLM, "4") {
t.Fatalf("expected 4 servers, got: %s", res.ForLLM)
}
// GitHub — no credentials needed, no custom creds
if !strings.Contains(res.ForLLM, "GitHub API") || !strings.Contains(res.ForLLM, "no credentials needed") {
t.Fatalf("expected GitHub with 'no credentials needed': %s", res.ForLLM)
}
// Postgres — no display name, should show raw name
if !strings.Contains(res.ForLLM, "postgres") || !strings.Contains(res.ForLLM, "no credentials needed") {
t.Fatalf("expected postgres with 'no credentials needed': %s", res.ForLLM)
}
// Slack — requires user credentials and has them
if !strings.Contains(res.ForLLM, "credentials set") {
t.Fatalf("expected Slack with 'credentials set': %s", res.ForLLM)
}
// Vault — requires user credentials (from settings) but not set
if !strings.Contains(res.ForLLM, "credentials required") {
t.Fatalf("expected vault with 'credentials required': %s", res.ForLLM)
}
}
func TestMCPCredentialManager_CredentialStatus_NoServerName(t *testing.T) {
_, tool := newTestMCPCredTool()
res := tool.Execute(mcpCredCtx(), map[string]any{"action": "credential_status"})
if !res.IsError {
t.Fatal("expected error for missing server_name")
}
if !strings.Contains(res.ForLLM, "server_name is required") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_CredentialStatus_ServerNotFound(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerGitHub)
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "credential_status",
"server_name": "nonexistent",
})
if !res.IsError {
t.Fatal("expected error for nonexistent server")
}
if !strings.Contains(res.ForLLM, "not found") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_CredentialStatus_NoUserID(t *testing.T) {
_, tool := newTestMCPCredTool()
ctx := context.Background()
ctx = store.WithAgentID(ctx, mcpCredTestAgentID)
// No user ID set
res := tool.Execute(ctx, map[string]any{
"action": "credential_status",
"server_name": "github",
})
if !res.IsError {
t.Fatal("expected error for missing user ID")
}
if !strings.Contains(res.ForLLM, "no user identity") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_CredentialStatus_NoCredsSet(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerGitHub)
mock.accessible = []store.MCPAccessInfo{{Server: *mcpCredServerGitHub}}
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "credential_status",
"server_name": "github",
})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "have not set any credentials") {
t.Fatalf("expected 'no credentials' message: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_CredentialStatus_NeedsCredsNoCreds(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerSlack)
mock.accessible = []store.MCPAccessInfo{{Server: *mcpCredServerSlack}}
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "credential_status",
"server_name": "slack",
})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "requires per-user credentials") {
t.Fatalf("expected 'requires credentials': %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "set_credentials") {
t.Fatalf("expected hint about set_credentials: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_CredentialStatus_WithCreds(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerGitHub)
mock.accessible = []store.MCPAccessInfo{{Server: *mcpCredServerGitHub}}
mock.setCredentials(mcpCredServerGitHub.ID, mcpCredTestUserID, &store.MCPUserCredentials{
APIKey: "ghp_abcdef1234567890abcdef1234567890",
Headers: map[string]string{"X-Custom": "value"},
Env: map[string]string{"MY_VAR": "my_value"},
})
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "credential_status",
"server_name": "github",
})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "ghp_") {
t.Fatalf("expected masked API key: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "X-Custom") {
t.Fatalf("expected header key listed: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "MY_VAR") {
t.Fatalf("expected env var key listed: %s", res.ForLLM)
}
// The actual value should NOT appear (masked)
if strings.Contains(res.ForLLM, "abcdef1234567890abcdef1234567890") {
t.Fatalf("API key value leaked in output: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_SetCredentials_NoServerName(t *testing.T) {
_, tool := newTestMCPCredTool()
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "set_credentials",
"api_key": "my-key",
})
if !res.IsError {
t.Fatal("expected error for missing server_name")
}
if !strings.Contains(res.ForLLM, "server_name is required") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_SetCredentials_NoUserID(t *testing.T) {
_, tool := newTestMCPCredTool()
ctx := context.Background()
ctx = store.WithAgentID(ctx, mcpCredTestAgentID)
res := tool.Execute(ctx, map[string]any{
"action": "set_credentials",
"server_name": "github",
"api_key": "my-key",
})
if !res.IsError {
t.Fatal("expected error for missing user ID")
}
if !strings.Contains(res.ForLLM, "no user identity") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_SetCredentials_ServerNotFound(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerGitHub)
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "set_credentials",
"server_name": "nonexistent",
"api_key": "my-key",
})
if !res.IsError {
t.Fatal("expected error for nonexistent server")
}
if !strings.Contains(res.ForLLM, "not found") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_SetCredentials_NoValues(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerGitHub)
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "set_credentials",
"server_name": "github",
})
if !res.IsError {
t.Fatal("expected error for no values")
}
if !strings.Contains(res.ForLLM, "at least one of") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_SetCredentials_Success(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerGitHub)
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "set_credentials",
"server_name": "github",
"api_key": "ghp_my_secret_key",
"headers": map[string]any{"X-Custom": "val1"},
"env": map[string]any{"MY_ENV": "val2"},
})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "Successfully set credentials") {
t.Fatalf("expected success: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "GitHub API") {
t.Fatalf("expected display name: %s", res.ForLLM)
}
// Verify the credential was stored
creds, err := mock.GetUserCredentials(context.Background(), mcpCredServerGitHub.ID, mcpCredTestUserID)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if creds == nil {
t.Fatal("credentials not stored")
}
if creds.APIKey != "ghp_my_secret_key" {
t.Fatalf("expected API key 'ghp_my_secret_key', got %q", creds.APIKey)
}
if creds.Headers["X-Custom"] != "val1" {
t.Fatalf("expected header X-Custom=val1, got %q", creds.Headers["X-Custom"])
}
if creds.Env["MY_ENV"] != "val2" {
t.Fatalf("expected env MY_ENV=val2, got %q", creds.Env["MY_ENV"])
}
}
func TestMCPCredentialManager_SetCredentials_SuccessNoDisplayName(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerPostgres)
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "set_credentials",
"server_name": "postgres",
"api_key": "pg-key",
})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
// Falls back to server name when DisplayName is empty
if !strings.Contains(res.ForLLM, "postgres") {
t.Fatalf("expected server name fallback: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_SetBearerToken_NoServerName(t *testing.T) {
_, tool := newTestMCPCredTool()
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "set_bearer_token",
"token": "my-token",
})
if !res.IsError {
t.Fatal("expected error for missing server_name")
}
if !strings.Contains(res.ForLLM, "server_name is required") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_SetBearerToken_NoToken(t *testing.T) {
_, tool := newTestMCPCredTool()
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "set_bearer_token",
"server_name": "github",
})
if !res.IsError {
t.Fatal("expected error for missing token")
}
if !strings.Contains(res.ForLLM, "token is required") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_SetBearerToken_NoUserID(t *testing.T) {
_, tool := newTestMCPCredTool()
ctx := context.Background()
ctx = store.WithAgentID(ctx, mcpCredTestAgentID)
res := tool.Execute(ctx, map[string]any{
"action": "set_bearer_token",
"server_name": "github",
"token": "my-token",
})
if !res.IsError {
t.Fatal("expected error for missing user ID")
}
if !strings.Contains(res.ForLLM, "no user identity") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_SetBearerToken_ServerNotFound(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerGitHub)
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "set_bearer_token",
"server_name": "nonexistent",
"token": "my-token",
})
if !res.IsError {
t.Fatal("expected error for nonexistent server")
}
if !strings.Contains(res.ForLLM, "not found") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_SetBearerToken_Success(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerGitHub)
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "set_bearer_token",
"server_name": "github",
"token": "ghp_bearer_token_value",
})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "Successfully set Bearer token") {
t.Fatalf("expected success: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "Authorization: Bearer") {
t.Fatalf("expected Authorization Bearer mention: %s", res.ForLLM)
}
// Verify the credential was stored as API key
creds, err := mock.GetUserCredentials(context.Background(), mcpCredServerGitHub.ID, mcpCredTestUserID)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if creds == nil {
t.Fatal("credentials not stored")
}
if creds.APIKey != "ghp_bearer_token_value" {
t.Fatalf("expected API key 'ghp_bearer_token_value', got %q", creds.APIKey)
}
}
func TestMCPCredentialManager_SetBearerToken_OverwritesExisting(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerSlack)
// Pre-set existing credentials
mock.setCredentials(mcpCredServerSlack.ID, mcpCredTestUserID, &store.MCPUserCredentials{
APIKey: "old-token",
Headers: map[string]string{"X-Old": "value"},
})
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "set_bearer_token",
"server_name": "slack",
"token": "new-bearer-token",
})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
// Verify old credentials were replaced
creds, err := mock.GetUserCredentials(context.Background(), mcpCredServerSlack.ID, mcpCredTestUserID)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if creds == nil {
t.Fatal("credentials not stored")
}
if creds.APIKey != "new-bearer-token" {
t.Fatalf("expected 'new-bearer-token', got %q", creds.APIKey)
}
// Old headers are cleared
if len(creds.Headers) > 0 {
t.Fatalf("expected headers cleared, got %v", creds.Headers)
}
}
func TestMCPCredentialManager_DeleteCredentials_NoServerName(t *testing.T) {
_, tool := newTestMCPCredTool()
res := tool.Execute(mcpCredCtx(), map[string]any{"action": "delete_credentials"})
if !res.IsError {
t.Fatal("expected error for missing server_name")
}
if !strings.Contains(res.ForLLM, "server_name is required") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_DeleteCredentials_NoUserID(t *testing.T) {
_, tool := newTestMCPCredTool()
ctx := context.Background()
ctx = store.WithAgentID(ctx, mcpCredTestAgentID)
res := tool.Execute(ctx, map[string]any{
"action": "delete_credentials",
"server_name": "github",
})
if !res.IsError {
t.Fatal("expected error for missing user ID")
}
if !strings.Contains(res.ForLLM, "no user identity") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_DeleteCredentials_ServerNotFound(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerGitHub)
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "delete_credentials",
"server_name": "nonexistent",
})
if !res.IsError {
t.Fatal("expected error for nonexistent server")
}
if !strings.Contains(res.ForLLM, "not found") {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_DeleteCredentials_Success(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerGitHub)
mock.setCredentials(mcpCredServerGitHub.ID, mcpCredTestUserID, &store.MCPUserCredentials{
APIKey: "ghp_my_key",
})
// Verify it exists first
creds, _ := mock.GetUserCredentials(context.Background(), mcpCredServerGitHub.ID, mcpCredTestUserID)
if creds == nil {
t.Fatal("precondition failed: credentials should exist")
}
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "delete_credentials",
"server_name": "github",
})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "Successfully deleted credentials") {
t.Fatalf("expected success: %s", res.ForLLM)
}
// Verify credentials are gone
creds, err := mock.GetUserCredentials(context.Background(), mcpCredServerGitHub.ID, mcpCredTestUserID)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if creds != nil {
t.Fatal("credentials should have been deleted")
}
}
func TestMCPCredentialManager_DeleteCredentials_Idempotent(t *testing.T) {
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerGitHub)
// No credentials pre-set
res := tool.Execute(mcpCredCtx(), map[string]any{
"action": "delete_credentials",
"server_name": "github",
})
if res.IsError {
t.Fatalf("expected delete to be idempotent, got error: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "Successfully deleted credentials") {
t.Fatalf("expected success: %s", res.ForLLM)
}
}
func TestMCPCredentialManager_StoreAccessibleReturnsListAccessible(t *testing.T) {
// Verify that listServers calls ListAccessible by checking the number of
// servers returned matches the mock's accessible count.
mock, tool := newTestMCPCredTool()
mock.addServer(mcpCredServerGitHub)
mock.addServer(mcpCredServerPostgres)
mock.accessible = []store.MCPAccessInfo{
{Server: *mcpCredServerGitHub},
{Server: *mcpCredServerPostgres},
}
res := tool.Execute(mcpCredCtx(), map[string]any{"action": "list_servers"})
if res.IsError {
t.Fatalf("unexpected error: %s", res.ForLLM)
}
if !strings.Contains(res.ForLLM, "2") {
t.Fatalf("expected 2 servers, got: %s", res.ForLLM)
}
}
// Test that requireUserCredsFromSettings returns true for settings with require_user_credentials.
func TestRequireUserCredsFromSettings(t *testing.T) {
tests := []struct {
name string
settings json.RawMessage
want bool
}{
{"nil settings", nil, false},
{"empty settings", json.RawMessage(""), false},
{"empty object", json.RawMessage("{}"), false},
{"require_user_credentials true", json.RawMessage(`{"require_user_credentials":true}`), true},
{"require_user_credentials false", json.RawMessage(`{"require_user_credentials":false}`), false},
{"nested unrelated", json.RawMessage(`{"timeout":30}`), false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := requireUserCredsFromSettings(tt.settings)
if got != tt.want {
t.Errorf("requireUserCredsFromSettings(%s) = %v, want %v", string(tt.settings), got, tt.want)
}
})
}
}
// Test that maskString does not leak full secrets.
func TestMaskString(t *testing.T) {
tests := []struct {
input string
want string
}{
{"short", "*****"},
{"12345678", "********"},
{"abcdefghijklmnop", "abcd********mnop"},
{"abcdefghij", "abcd**ghij"},
}
for _, tt := range tests {
t.Run(tt.input, func(t *testing.T) {
got := maskString(tt.input)
if got != tt.want {
t.Errorf("maskString(%q) = %q, want %q", tt.input, got, tt.want)
}
// Ensure no part of the input beyond first 4 and last 4 leaks
if len(tt.input) > 8 {
middle := tt.input[4 : len(tt.input)-4]
if strings.Contains(got, middle) {
t.Errorf("maskString(%q) leaked middle part %q in %q", tt.input, middle, got)
}
}
})
}
}
+5
View File
@@ -152,6 +152,11 @@ type ChannelTenantCheckerAware interface {
SetChannelTenantChecker(ChannelTenantChecker)
}
// MCPServerStoreAware tools can receive an MCPServerStore for credential management.
type MCPServerStoreAware interface {
SetMCPServerStore(store.MCPServerStore)
}
// ChannelAware is optionally implemented by tools that only work on specific channel types.
// Tools implementing this are filtered out when the current channel type doesn't match.
type ChannelAware interface {