mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-11 03:13:24 +00:00
mcp credentials over chat (#1490)
This commit is contained in:
1 parent
bcd23c587b
commit
2e1b90ca35
5 files changed
+1228
-1
No files matched your search
@@ -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
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
Reference in new issue
Block a user