github-featured/googleapis--mcp-toolbox
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

feat: add groups support (#3605)

Browse Source 
## Description
 
Master PR for adding groups support to MCP Toolbox. Groups currently
contain tools and prompts.

---------

Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Twisha Bansal
2026-07-24 23:24:47 +05:30
committed by GitHub
parent 7666bc08de
commit e75ec3b5c8
54 changed files with 3459 additions and 1118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/internal/config.go
+18 -16
View File
@@ -37,8 +37,8 @@ type Config struct {
AuthServices server.AuthServiceConfigs `yaml:"authServices"`
EmbeddingModels server.EmbeddingModelConfigs `yaml:"embeddingModels"`
Tools server.ToolConfigs `yaml:"tools"`
Toolsets server.ToolsetConfigs `yaml:"toolsets"`
Prompts server.PromptConfigs `yaml:"prompts"`
Groups server.GroupConfigs `yaml:"groups"`
}
type ConfigParser struct {
@@ -150,7 +150,7 @@ func (p *ConfigParser) ParseConfig(ctx context.Context, raw []byte) (Config, err
}
// Parse contents
config.Sources, config.AuthServices, config.EmbeddingModels, config.Tools, config.Toolsets, config.Prompts, err = server.UnmarshalPrimitiveConfig(ctx, raw)
config.Sources, config.AuthServices, config.EmbeddingModels, config.Tools, config.Prompts, config.Groups, err = server.UnmarshalPrimitiveConfig(ctx, raw)
if err != nil {
return config, err
}
@@ -180,7 +180,7 @@ func ConvertConfig(raw []byte) ([]byte, error) {
decoder := yaml.NewDecoder(bytes.NewReader(raw), yaml.UseOrderedMap())
encoder := yaml.NewEncoder(&buf, yaml.UseLiteralStyleIfMultiline(true))
nestedFormatKey := []string{"sources", "authServices", "embeddingModels", "tools", "toolsets", "prompts"}
nestedFormatKey := []string{"sources", "authServices", "embeddingModels", "tools", "toolsets", "prompts", "groups"}
docIndex := 0
for {
if err := decoder.Decode(&input); err != nil {
@@ -217,6 +217,8 @@ func ConvertConfig(raw []byte) ([]byte, error) {
key = "toolset"
case "prompts":
key = "prompt"
case "groups":
key = "group"
}
transformed, err := transformDocs(key, slice)
if err != nil {
@@ -306,16 +308,16 @@ func processValue(v any, isToolset bool) any {
}
// mergeConfigs merges multiple Config structs into one.
// Detects and raises errors for resource conflicts in sources, authServices, tools, and toolsets.
// All resource names (sources, authServices, tools, toolsets) must be unique across all files.
// Detects and raises errors for resource conflicts in sources, authServices, tools, and groups.
// All resource names (sources, authServices, tools, groups) must be unique across all files.
func mergeConfigs(files ...Config) (Config, error) {
merged := Config{
Sources: make(server.SourceConfigs),
AuthServices: make(server.AuthServiceConfigs),
EmbeddingModels: make(server.EmbeddingModelConfigs),
Tools: make(server.ToolConfigs),
Toolsets: make(server.ToolsetConfigs),
Prompts: make(server.PromptConfigs),
Groups: make(server.GroupConfigs),
}
var conflicts []string
@@ -359,15 +361,6 @@ func mergeConfigs(files ...Config) (Config, error) {
}
}
// Check for conflicts and merge toolsets
for name, toolset := range file.Toolsets {
if _, exists := merged.Toolsets[name]; exists {
conflicts = append(conflicts, fmt.Sprintf("toolset '%s' (file #%d)", name, fileIndex+1))
} else {
merged.Toolsets[name] = toolset
}
}
// Check for conflicts and merge prompts
for name, prompt := range file.Prompts {
if _, exists := merged.Prompts[name]; exists {
@@ -376,11 +369,20 @@ func mergeConfigs(files ...Config) (Config, error) {
merged.Prompts[name] = prompt
}
}
// Check for conflicts and merge groups
for name, grp := range file.Groups {
if _, exists := merged.Groups[name]; exists {
conflicts = append(conflicts, fmt.Sprintf("group '%s' (file #%d)", name, fileIndex+1))
} else {
merged.Groups[name] = grp
}
}
}
// If conflicts were detected, return an error
if len(conflicts) > 0 {
return Config{}, fmt.Errorf("resource conflicts detected:\n - %s\n\nPlease ensure each source, authService, tool, toolset and prompt has a unique name across all files", strings.Join(conflicts, "\n - "))
return Config{}, fmt.Errorf("resource conflicts detected:\n - %s\n\nPlease ensure each source, authService, tool, prompt and group has a unique name across all files", strings.Join(conflicts, "\n - "))
}
// Ensure only one authService has mcpEnabled = true
cmd/internal/config_test.go
+250 -141
View File
@@ -23,6 +23,7 @@ import (
"github.com/googleapis/mcp-toolbox/internal/auth/generic"
"github.com/googleapis/mcp-toolbox/internal/auth/google"
"github.com/googleapis/mcp-toolbox/internal/embeddingmodels/gemini"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prebuiltconfigs"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/prompts/custom"
@@ -245,6 +246,26 @@ type: gemini
model: gemini-embedding-001
apiKey: some-key
dimension: 768
`,
},
{
desc: "convert nested groups",
in: `
groups:
my_group:
description: my group description
tools:
- example_tool
prompts:
- code_review`,
want: `
kind: group
name: my_group
description: my group description
tools:
- example_tool
prompts:
- code_review
`,
},
{
@@ -651,8 +672,8 @@ func TestParseConfig(t *testing.T) {
},
},
},
Toolsets: server.ToolsetConfigs{
"example_toolset": tools.ToolsetConfig{
Groups: server.GroupConfigs{
"example_toolset": group.GroupConfig{
Name: "example_toolset",
ToolNames: []string{"example_tool"},
},
@@ -775,8 +796,8 @@ func TestParseConfig(t *testing.T) {
},
},
},
Toolsets: server.ToolsetConfigs{
"example_toolset": tools.ToolsetConfig{
Groups: server.GroupConfigs{
"example_toolset": group.GroupConfig{
Name: "example_toolset",
ToolNames: []string{"example_tool"},
},
@@ -811,7 +832,7 @@ func TestParseConfig(t *testing.T) {
Sources: nil,
AuthServices: nil,
Tools: nil,
Toolsets: nil,
Groups: nil,
Prompts: server.PromptConfigs{
"my-prompt": &custom.Config{
Name: "my-prompt",
@@ -843,7 +864,7 @@ func TestParseConfig(t *testing.T) {
if diff := cmp.Diff(tc.wantConfig.Tools, configFile.Tools); diff != "" {
t.Fatalf("incorrect tools parse: diff %v", diff)
}
if diff := cmp.Diff(tc.wantConfig.Toolsets, configFile.Toolsets); diff != "" {
if diff := cmp.Diff(tc.wantConfig.Groups, configFile.Groups); diff != "" {
t.Fatalf("incorrect toolsets parse: diff %v", diff)
}
if diff := cmp.Diff(tc.wantConfig.Prompts, configFile.Prompts); diff != "" {
@@ -1029,8 +1050,8 @@ func TestParseConfigWithAuth(t *testing.T) {
},
},
},
Toolsets: server.ToolsetConfigs{
"example_toolset": tools.ToolsetConfig{
Groups: server.GroupConfigs{
"example_toolset": group.GroupConfig{
Name: "example_toolset",
ToolNames: []string{"example_tool"},
},
@@ -1137,8 +1158,8 @@ func TestParseConfigWithAuth(t *testing.T) {
},
},
},
Toolsets: server.ToolsetConfigs{
"example_toolset": tools.ToolsetConfig{
Groups: server.GroupConfigs{
"example_toolset": group.GroupConfig{
Name: "example_toolset",
ToolNames: []string{"example_tool"},
},
@@ -1163,7 +1184,7 @@ func TestParseConfigWithAuth(t *testing.T) {
if diff := cmp.Diff(tc.wantConfig.Tools, configFile.Tools); diff != "" {
t.Fatalf("incorrect tools parse: diff %v", diff)
}
if diff := cmp.Diff(tc.wantConfig.Toolsets, configFile.Toolsets); diff != "" {
if diff := cmp.Diff(tc.wantConfig.Groups, configFile.Groups); diff != "" {
t.Fatalf("incorrect toolsets parse: diff %v", diff)
}
if diff := cmp.Diff(tc.wantConfig.Prompts, configFile.Prompts); diff != "" {
@@ -1319,8 +1340,8 @@ func TestEnvVarReplacement(t *testing.T) {
HeaderParams: []parameters.Parameter{parameters.NewStringParameter("Language", "language string")},
},
},
Toolsets: server.ToolsetConfigs{
"ACTUAL_TOOLSET_NAME": tools.ToolsetConfig{
Groups: server.GroupConfigs{
"ACTUAL_TOOLSET_NAME": group.GroupConfig{
Name: "ACTUAL_TOOLSET_NAME",
ToolNames: []string{"example_tool"},
},
@@ -1467,8 +1488,8 @@ func TestEnvVarReplacement(t *testing.T) {
HeaderParams: []parameters.Parameter{parameters.NewStringParameter("Language", "language string")},
},
},
Toolsets: server.ToolsetConfigs{
"ACTUAL_TOOLSET_NAME": tools.ToolsetConfig{
Groups: server.GroupConfigs{
"ACTUAL_TOOLSET_NAME": group.GroupConfig{
Name: "ACTUAL_TOOLSET_NAME",
ToolNames: []string{"example_tool"},
},
@@ -1505,7 +1526,7 @@ func TestEnvVarReplacement(t *testing.T) {
if diff := cmp.Diff(tc.wantConfig.Tools, configFile.Tools); diff != "" {
t.Fatalf("incorrect tools parse: diff %v", diff)
}
if diff := cmp.Diff(tc.wantConfig.Toolsets, configFile.Toolsets); diff != "" {
if diff := cmp.Diff(tc.wantConfig.Groups, configFile.Groups); diff != "" {
t.Fatalf("incorrect toolsets parse: diff %v", diff)
}
if diff := cmp.Diff(tc.wantConfig.Prompts, configFile.Prompts); diff != "" {
@@ -1703,39 +1724,39 @@ func TestPrebuiltTools(t *testing.T) {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
name string
in []byte
wantToolset server.ToolsetConfigs
name string
in []byte
wantGroups server.GroupConfigs
}{
{
name: "alloydb omni prebuilt tools",
in: alloydb_omni_config,
wantToolset: server.ToolsetConfigs{
"data": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"data": group.GroupConfig{
Name: "data",
ToolNames: []string{"execute_sql", "list_tables", "list_views", "list_schemas", "list_triggers", "list_indexes", "list_sequences", "list_stored_procedure"},
},
"performance": tools.ToolsetConfig{
"performance": group.GroupConfig{
Name: "performance",
ToolNames: []string{"execute_sql", "get_query_plan", "list_query_stats", "get_column_cardinality", "list_table_stats", "list_database_stats", "list_active_queries"},
},
"monitor": tools.ToolsetConfig{
"monitor": group.GroupConfig{
Name: "monitor",
ToolNames: []string{"database_overview", "list_active_queries", "long_running_transactions", "list_locks", "list_database_stats", "list_pg_settings"},
},
"optimize": tools.ToolsetConfig{
"optimize": group.GroupConfig{
Name: "optimize",
ToolNames: []string{"list_pg_settings", "list_memory_configurations", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_columnar_configurations", "list_columnar_recommended_columns"},
},
"health": tools.ToolsetConfig{
"health": group.GroupConfig{
Name: "health",
ToolNames: []string{"list_top_bloated_tables", "list_invalid_indexes", "list_table_stats", "list_tablespaces", "database_overview", "list_autovacuum_configurations"},
},
"replication": tools.ToolsetConfig{
"replication": group.GroupConfig{
Name: "replication",
ToolNames: []string{"replication_stats", "list_replication_slots", "list_publication_tables", "database_overview"},
},
"access-control": tools.ToolsetConfig{
"access-control": group.GroupConfig{
Name: "access-control",
ToolNames: []string{"list_roles", "list_pg_settings", "database_overview"},
},
@@ -1744,8 +1765,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "alloydb postgres admin prebuilt tools",
in: alloydb_admin_config,
wantToolset: server.ToolsetConfigs{
"alloydb_postgres_admin_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"alloydb_postgres_admin_tools": group.GroupConfig{
Name: "alloydb_postgres_admin_tools",
ToolNames: []string{"create_cluster", "wait_for_operation", "create_instance", "list_clusters", "list_instances", "list_users", "create_user", "get_cluster", "get_instance", "get_user"},
},
@@ -1754,8 +1775,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "cloudsql pg admin prebuilt tools",
in: cloudsqlpg_admin_config,
wantToolset: server.ToolsetConfigs{
"cloud_sql_postgres_admin_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"cloud_sql_postgres_admin_tools": group.GroupConfig{
Name: "cloud_sql_postgres_admin_tools",
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "postgres_upgrade_precheck", "clone_instance", "create_backup", "restore_backup"},
},
@@ -1764,8 +1785,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "cloudsql mysql admin prebuilt tools",
in: cloudsqlmysql_admin_config,
wantToolset: server.ToolsetConfigs{
"cloud_sql_mysql_admin_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"cloud_sql_mysql_admin_tools": group.GroupConfig{
Name: "cloud_sql_mysql_admin_tools",
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "clone_instance", "create_backup", "restore_backup"},
},
@@ -1774,8 +1795,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "cloudsql mssql admin prebuilt tools",
in: cloudsqlmssql_admin_config,
wantToolset: server.ToolsetConfigs{
"cloud_sql_mssql_admin_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"cloud_sql_mssql_admin_tools": group.GroupConfig{
Name: "cloud_sql_mssql_admin_tools",
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "clone_instance", "create_backup", "restore_backup"},
},
@@ -1784,32 +1805,32 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "alloydb prebuilt tools",
in: alloydb_config,
wantToolset: server.ToolsetConfigs{
"admin": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"admin": group.GroupConfig{
Name: "admin",
ToolNames: []string{"create_cluster", "get_cluster", "list_clusters", "create_instance", "get_instance", "list_instances", "database_overview", "wait_for_operation"},
},
"access-management": tools.ToolsetConfig{
"access-management": group.GroupConfig{
Name: "access-management",
ToolNames: []string{"create_user", "list_users", "get_user", "list_roles", "list_pg_settings", "database_overview"},
},
"data": tools.ToolsetConfig{
"data": group.GroupConfig{
Name: "data",
ToolNames: []string{"execute_sql", "list_tables", "list_views", "list_schemas", "list_triggers", "list_indexes", "list_sequences", "list_stored_procedure"},
},
"monitor": tools.ToolsetConfig{
"monitor": group.GroupConfig{
Name: "monitor",
ToolNames: []string{"list_active_queries", "list_query_stats", "get_query_plan", "get_query_metrics", "get_system_metrics", "long_running_transactions", "list_locks", "list_database_stats"},
},
"health": tools.ToolsetConfig{
"health": group.GroupConfig{
Name: "health",
ToolNames: []string{"list_top_bloated_tables", "list_invalid_indexes", "list_table_stats", "get_column_cardinality", "list_autovacuum_configurations", "list_tablespaces", "database_overview", "get_instance"},
},
"optimize": tools.ToolsetConfig{
"optimize": group.GroupConfig{
Name: "optimize",
ToolNames: []string{"list_available_extensions", "list_installed_extensions", "list_memory_configurations", "list_pg_settings", "database_overview", "get_cluster"},
},
"replication": tools.ToolsetConfig{
"replication": group.GroupConfig{
Name: "replication",
ToolNames: []string{"replication_stats", "list_replication_slots", "list_publication_tables", "list_instances", "get_instance", "database_overview"},
},
@@ -1818,12 +1839,12 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "bigquery prebuilt tools",
in: bigquery_config,
wantToolset: server.ToolsetConfigs{
"data": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"data": group.GroupConfig{
Name: "data",
ToolNames: []string{"execute_sql", "list_dataset_ids", "list_table_ids", "get_dataset_info", "get_table_info", "search_catalog"},
},
"analytics": tools.ToolsetConfig{
"analytics": group.GroupConfig{
Name: "analytics",
ToolNames: []string{"analyze_contribution", "ask_data_insights", "forecast", "search_catalog"},
},
@@ -1832,8 +1853,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "clickhouse prebuilt tools",
in: clickhouse_config,
wantToolset: server.ToolsetConfigs{
"clickhouse_database_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"clickhouse_database_tools": group.GroupConfig{
Name: "clickhouse_database_tools",
ToolNames: []string{"execute_sql", "list_databases", "list_tables"},
},
@@ -1842,32 +1863,32 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "cloudsqlpg prebuilt tools",
in: cloudsqlpg_config,
wantToolset: server.ToolsetConfigs{
"admin": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"admin": group.GroupConfig{
Name: "admin",
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "clone_instance"},
},
"lifecycle": tools.ToolsetConfig{
"lifecycle": group.GroupConfig{
Name: "lifecycle",
ToolNames: []string{"create_backup", "restore_backup", "postgres_upgrade_precheck", "wait_for_operation", "database_overview", "get_instance", "list_instances"},
},
"data": tools.ToolsetConfig{
"data": group.GroupConfig{
Name: "data",
ToolNames: []string{"execute_sql", "list_tables", "list_views", "list_schemas", "list_triggers", "list_indexes", "list_sequences", "list_stored_procedure"},
},
"monitor": tools.ToolsetConfig{
"monitor": group.GroupConfig{
Name: "monitor",
ToolNames: []string{"get_system_metrics", "get_query_metrics", "list_query_stats", "get_query_plan", "list_database_stats", "list_active_queries", "long_running_transactions", "list_locks"},
},
"health": tools.ToolsetConfig{
"health": group.GroupConfig{
Name: "health",
ToolNames: []string{"list_top_bloated_tables", "list_invalid_indexes", "list_table_stats", "get_column_cardinality", "list_autovacuum_configurations", "list_tablespaces", "database_overview", "list_pg_settings"},
},
"view-config": tools.ToolsetConfig{
"view-config": group.GroupConfig{
Name: "view-config",
ToolNames: []string{"list_available_extensions", "list_installed_extensions", "list_memory_configurations", "list_pg_settings", "database_overview", "get_instance"},
},
"replication": tools.ToolsetConfig{
"replication": group.GroupConfig{
Name: "replication",
ToolNames: []string{"replication_stats", "list_replication_slots", "list_publication_tables", "list_roles", "list_pg_settings", "database_overview"},
},
@@ -1880,20 +1901,20 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "cloudsqlmysql prebuilt tools",
in: cloudsqlmysql_config,
wantToolset: server.ToolsetConfigs{
"admin": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"admin": group.GroupConfig{
Name: "admin",
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation"},
},
"data": tools.ToolsetConfig{
"data": group.GroupConfig{
Name: "data",
ToolNames: []string{"execute_sql", "list_tables", "get_query_plan", "list_active_queries"},
},
"monitor": tools.ToolsetConfig{
"monitor": group.GroupConfig{
Name: "monitor",
ToolNames: []string{"get_query_plan", "list_active_queries", "list_all_locks", "get_query_metrics", "get_system_metrics", "list_table_fragmentation", "list_table_stats", "list_tables_missing_unique_indexes", "show_query_stats"},
},
"lifecycle": tools.ToolsetConfig{
"lifecycle": group.GroupConfig{
Name: "lifecycle",
ToolNames: []string{"create_backup", "restore_backup", "clone_instance", "list_instances", "get_instance", "wait_for_operation"},
},
@@ -1902,20 +1923,20 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "cloudsqlmssql prebuilt tools",
in: cloudsqlmssql_config,
wantToolset: server.ToolsetConfigs{
"admin": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"admin": group.GroupConfig{
Name: "admin",
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation"},
},
"data": tools.ToolsetConfig{
"data": group.GroupConfig{
Name: "data",
ToolNames: []string{"execute_sql", "list_tables"},
},
"monitor": tools.ToolsetConfig{
"monitor": group.GroupConfig{
Name: "monitor",
ToolNames: []string{"get_system_metrics"},
},
"lifecycle": tools.ToolsetConfig{
"lifecycle": group.GroupConfig{
Name: "lifecycle",
ToolNames: []string{"create_backup", "restore_backup", "clone_instance", "list_instances", "get_instance", "wait_for_operation"},
},
@@ -1924,16 +1945,16 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "dataplex prebuilt tools",
in: dataplex_config,
wantToolset: server.ToolsetConfigs{
"discovery": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"discovery": group.GroupConfig{
Name: "discovery",
ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types", "lookup_context", "search_dq_scans"},
},
"data-products": tools.ToolsetConfig{
"data-products": group.GroupConfig{
Name: "data-products",
ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types", "lookup_context", "list_data_products", "get_data_product", "list_data_assets", "get_data_asset", "create_data_product", "update_data_product", "create_data_asset", "update_data_asset"},
},
"enrich": tools.ToolsetConfig{
"enrich": group.GroupConfig{
Name: "enrich",
ToolNames: []string{"search_entries", "lookup_entry", "lookup_context", "generate_data_insights", "get_data_insights", "generate_data_profile", "get_data_profile", "discover_metadata", "get_discovery_results", "check_data_quality", "get_data_quality_results", "get_operation", "get_run_status"},
},
@@ -1942,8 +1963,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "dataproc prebuilt tools",
in: dataproc_config,
wantToolset: server.ToolsetConfigs{
"dataproc_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"dataproc_tools": group.GroupConfig{
Name: "dataproc_tools",
ToolNames: []string{"list_clusters", "get_cluster", "list_jobs", "get_job"},
},
@@ -1952,8 +1973,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "serverless spark prebuilt tools",
in: serverless_spark_config,
wantToolset: server.ToolsetConfigs{
"serverless_spark_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"serverless_spark_tools": group.GroupConfig{
Name: "serverless_spark_tools",
ToolNames: []string{"list_batches", "get_batch", "cancel_batch", "create_pyspark_batch", "create_spark_batch", "get_session_template", "list_sessions", "get_session"},
},
@@ -1962,12 +1983,12 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "firestore prebuilt tools",
in: firestoreconfig,
wantToolset: server.ToolsetConfigs{
"data": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"data": group.GroupConfig{
Name: "data",
ToolNames: []string{"get_documents", "add_documents", "update_document", "delete_documents", "query_collection", "list_collections"},
},
"security": tools.ToolsetConfig{
"security": group.GroupConfig{
Name: "security",
ToolNames: []string{"get_rules", "validate_rules"},
},
@@ -1976,12 +1997,12 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "mysql prebuilt tools",
in: mysql_config,
wantToolset: server.ToolsetConfigs{
"data": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"data": group.GroupConfig{
Name: "data",
ToolNames: []string{"execute_sql", "list_tables", "get_query_plan", "list_active_queries"},
},
"monitor": tools.ToolsetConfig{
"monitor": group.GroupConfig{
Name: "monitor",
ToolNames: []string{"get_query_plan", "list_active_queries", "list_all_locks", "list_table_fragmentation", "list_table_stats", "list_tables_missing_unique_indexes", "show_query_stats"},
},
@@ -1990,8 +2011,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "mssql prebuilt tools",
in: mssql_config,
wantToolset: server.ToolsetConfigs{
"data": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"data": group.GroupConfig{
Name: "data",
ToolNames: []string{"execute_sql", "list_tables"},
},
@@ -2000,8 +2021,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "looker prebuilt tools",
in: looker_config,
wantToolset: server.ToolsetConfigs{
"looker_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"looker_tools": group.GroupConfig{
Name: "looker_tools",
ToolNames: []string{"get_models", "get_explores", "get_dimensions", "get_measures", "get_filters", "get_parameters", "query", "query_sql", "query_url", "get_looks", "run_look", "make_look", "get_dashboards", "run_dashboard", "make_dashboard", "add_dashboard_element", "add_dashboard_filter", "generate_embed_url"},
},
@@ -2010,8 +2031,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "looker dev prebuilt tools",
in: looker_dev_config,
wantToolset: server.ToolsetConfigs{
"looker_dev_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"looker_dev_tools": group.GroupConfig{
Name: "looker_dev_tools",
ToolNames: []string{"health_pulse", "health_analyze", "health_vacuum", "dev_mode", "get_projects", "get_project_files", "get_project_file", "create_project_file", "update_project_file", "delete_project_file", "get_project_directories", "create_project_directory", "delete_project_directory", "validate_project", "get_connections", "get_connection_schemas", "get_connection_databases", "get_connection_tables", "get_connection_table_columns", "get_lookml_tests", "run_lookml_tests", "create_view_from_table", "list_git_branches", "get_git_branch", "create_git_branch", "switch_git_branch", "delete_git_branch"},
},
@@ -2020,8 +2041,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "looker-conversational-analytics prebuilt tools",
in: lookerca_config,
wantToolset: server.ToolsetConfigs{
"looker_conversational_analytics_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"looker_conversational_analytics_tools": group.GroupConfig{
Name: "looker_conversational_analytics_tools",
ToolNames: []string{"ask_data_insights", "get_models", "get_explores"},
},
@@ -2030,24 +2051,24 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "postgres prebuilt tools",
in: postgresconfig,
wantToolset: server.ToolsetConfigs{
"data": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"data": group.GroupConfig{
Name: "data",
ToolNames: []string{"execute_sql", "list_tables", "list_views", "list_schemas", "list_triggers", "list_indexes", "list_sequences", "list_stored_procedure"},
},
"monitor": tools.ToolsetConfig{
"monitor": group.GroupConfig{
Name: "monitor",
ToolNames: []string{"list_query_stats", "get_query_plan", "list_database_stats", "list_active_queries", "long_running_transactions", "list_locks"},
},
"health": tools.ToolsetConfig{
"health": group.GroupConfig{
Name: "health",
ToolNames: []string{"list_top_bloated_tables", "list_invalid_indexes", "list_table_stats", "get_column_cardinality", "list_autovacuum_configurations", "list_tablespaces", "database_overview", "list_pg_settings"},
},
"view-config": tools.ToolsetConfig{
"view-config": group.GroupConfig{
Name: "view-config",
ToolNames: []string{"list_available_extensions", "list_installed_extensions", "list_memory_configurations", "list_pg_settings", "database_overview"},
},
"replication": tools.ToolsetConfig{
"replication": group.GroupConfig{
Name: "replication",
ToolNames: []string{"replication_stats", "list_replication_slots", "list_publication_tables", "list_roles", "list_pg_settings", "database_overview"},
},
@@ -2056,12 +2077,12 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "spanner prebuilt tools",
in: spanner_config,
wantToolset: server.ToolsetConfigs{
"data": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"data": group.GroupConfig{
Name: "data",
ToolNames: []string{"execute_sql", "execute_sql_dql", "list_tables", "list_graphs"},
},
"data_with_discovery": tools.ToolsetConfig{
"data_with_discovery": group.GroupConfig{
Name: "data_with_discovery",
ToolNames: []string{"execute_sql", "execute_sql_dql", "list_tables", "list_graphs", "search_catalog"},
},
@@ -2070,12 +2091,12 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "spanner pg prebuilt tools",
in: spannerpg_config,
wantToolset: server.ToolsetConfigs{
"data": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"data": group.GroupConfig{
Name: "data",
ToolNames: []string{"execute_sql", "execute_sql_dql", "list_tables"},
},
"data_with_discovery": tools.ToolsetConfig{
"data_with_discovery": group.GroupConfig{
Name: "data_with_discovery",
ToolNames: []string{"execute_sql", "execute_sql_dql", "list_tables", "search_catalog"},
},
@@ -2084,8 +2105,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "mindsdb prebuilt tools",
in: mindsdb_config,
wantToolset: server.ToolsetConfigs{
"mindsdb-tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"mindsdb-tools": group.GroupConfig{
Name: "mindsdb-tools",
ToolNames: []string{"execute_sql", "parameterized_sql"},
},
@@ -2094,8 +2115,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "sqlite prebuilt tools",
in: sqlite_config,
wantToolset: server.ToolsetConfigs{
"sqlite_database_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"sqlite_database_tools": group.GroupConfig{
Name: "sqlite_database_tools",
ToolNames: []string{"execute_sql", "list_tables"},
},
@@ -2104,8 +2125,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "neo4j prebuilt tools",
in: neo4jconfig,
wantToolset: server.ToolsetConfigs{
"neo4j_database_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"neo4j_database_tools": group.GroupConfig{
Name: "neo4j_database_tools",
ToolNames: []string{"execute_cypher", "get_schema"},
},
@@ -2114,8 +2135,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "alloydb postgres observability prebuilt tools",
in: alloydbobsvconfig,
wantToolset: server.ToolsetConfigs{
"alloydb_postgres_cloud_monitoring_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"alloydb_postgres_cloud_monitoring_tools": group.GroupConfig{
Name: "alloydb_postgres_cloud_monitoring_tools",
ToolNames: []string{"get_system_metrics", "get_query_metrics"},
},
@@ -2124,8 +2145,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "cloudsql postgres observability prebuilt tools",
in: cloudsqlpgobsvconfig,
wantToolset: server.ToolsetConfigs{
"cloud_sql_postgres_cloud_monitoring_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"cloud_sql_postgres_cloud_monitoring_tools": group.GroupConfig{
Name: "cloud_sql_postgres_cloud_monitoring_tools",
ToolNames: []string{"get_system_metrics", "get_query_metrics"},
},
@@ -2134,8 +2155,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "cloudsql mysql observability prebuilt tools",
in: cloudsqlmysqlobsvconfig,
wantToolset: server.ToolsetConfigs{
"cloud_sql_mysql_cloud_monitoring_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"cloud_sql_mysql_cloud_monitoring_tools": group.GroupConfig{
Name: "cloud_sql_mysql_cloud_monitoring_tools",
ToolNames: []string{"get_system_metrics", "get_query_metrics"},
},
@@ -2144,8 +2165,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "cloudsql mssql observability prebuilt tools",
in: cloudsqlmssqlobsvconfig,
wantToolset: server.ToolsetConfigs{
"cloud_sql_mssql_cloud_monitoring_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"cloud_sql_mssql_cloud_monitoring_tools": group.GroupConfig{
Name: "cloud_sql_mssql_cloud_monitoring_tools",
ToolNames: []string{"get_system_metrics"},
},
@@ -2154,16 +2175,16 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "cloud healthcare prebuilt tools",
in: cloudhealthcare_config,
wantToolset: server.ToolsetConfigs{
"cloud_healthcare_dataset_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"cloud_healthcare_dataset_tools": group.GroupConfig{
Name: "cloud_healthcare_dataset_tools",
ToolNames: []string{"get_dataset", "list_dicom_stores", "list_fhir_stores"},
},
"cloud_healthcare_fhir_tools": tools.ToolsetConfig{
"cloud_healthcare_fhir_tools": group.GroupConfig{
Name: "cloud_healthcare_fhir_tools",
ToolNames: []string{"get_fhir_store", "get_fhir_store_metrics", "get_fhir_resource", "fhir_patient_search", "fhir_patient_everything", "fhir_fetch_page"},
},
"cloud_healthcare_dicom_tools": tools.ToolsetConfig{
"cloud_healthcare_dicom_tools": group.GroupConfig{
Name: "cloud_healthcare_dicom_tools",
ToolNames: []string{"get_dicom_store", "get_dicom_store_metrics", "search_dicom_studies", "search_dicom_series", "search_dicom_instances", "retrieve_rendered_dicom_instance"},
},
@@ -2172,12 +2193,12 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "cloud storage prebuilt tools",
in: cloudstorage_config,
wantToolset: server.ToolsetConfigs{
"cloud-storage-buckets": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"cloud-storage-buckets": group.GroupConfig{
Name: "cloud-storage-buckets",
ToolNames: []string{"list_buckets", "create_bucket", "get_bucket_metadata", "get_bucket_iam_policy", "delete_bucket"},
},
"cloud-storage-objects": tools.ToolsetConfig{
"cloud-storage-objects": group.GroupConfig{
Name: "cloud-storage-objects",
ToolNames: []string{"list_objects", "get_object_metadata", "read_object", "download_object", "write_object", "upload_object", "copy_object", "move_object", "delete_object"},
},
@@ -2186,8 +2207,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "Snowflake prebuilt tool",
in: snowflake_config,
wantToolset: server.ToolsetConfigs{
"snowflake_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"snowflake_tools": group.GroupConfig{
Name: "snowflake_tools",
ToolNames: []string{"execute_sql", "list_tables"},
},
@@ -2196,8 +2217,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "Oracle prebuilt tools",
in: oracle_config,
wantToolset: server.ToolsetConfigs{
"oracle_database_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"oracle_database_tools": group.GroupConfig{
Name: "oracle_database_tools",
ToolNames: []string{"execute_sql", "list_tables", "list_active_sessions", "get_query_plan", "list_top_sql_by_resource", "list_tablespace_usage", "list_invalid_objects"},
},
@@ -2206,8 +2227,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "Conversational Analytics with Data Agent prebuilt tools",
in: conversationalanalytics_config,
wantToolset: server.ToolsetConfigs{
"conversational_analytics_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"conversational_analytics_tools": group.GroupConfig{
Name: "conversational_analytics_tools",
ToolNames: []string{"list_accessible_data_agents", "get_data_agent_info", "ask_data_agent"},
},
@@ -2216,8 +2237,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "Elasticsearch prebuilt tools",
in: elasticsearch_config,
wantToolset: server.ToolsetConfigs{
"elasticsearch-tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"elasticsearch-tools": group.GroupConfig{
Name: "elasticsearch-tools",
ToolNames: []string{"execute_esql_query"},
},
@@ -2226,8 +2247,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "Oceanbase prebuilt tools",
in: oceanbase_config,
wantToolset: server.ToolsetConfigs{
"oceanbase_database_tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"oceanbase_database_tools": group.GroupConfig{
Name: "oceanbase_database_tools",
ToolNames: []string{"execute_sql", "list_tables"},
},
@@ -2236,8 +2257,8 @@ func TestPrebuiltTools(t *testing.T) {
{
name: "Singlestore prebuilt tools",
in: singlestore_config,
wantToolset: server.ToolsetConfigs{
"singlestore-database-tools": tools.ToolsetConfig{
wantGroups: server.GroupConfigs{
"singlestore-database-tools": group.GroupConfig{
Name: "singlestore-database-tools",
ToolNames: []string{"execute_sql", "list_tables"},
},
@@ -2252,7 +2273,7 @@ func TestPrebuiltTools(t *testing.T) {
if err != nil {
t.Fatalf("failed to parse input: %v", err)
}
if diff := cmp.Diff(tc.wantToolset, configFile.Toolsets); diff != "" {
if diff := cmp.Diff(tc.wantGroups, configFile.Groups); diff != "" {
t.Fatalf("incorrect tools parse: diff %v", diff)
}
// Prebuilt configs do not have prompts, so assert empty maps.
@@ -2261,9 +2282,9 @@ func TestPrebuiltTools(t *testing.T) {
}
t.Run("check toolset sizes", func(t *testing.T) {
for tsName, ts := range configFile.Toolsets {
for tsName, ts := range configFile.Groups {
if len(ts.ToolNames) > 10 {
t.Logf("WARNING: Toolset %q in config %q has %d tools, which is larger than the recommended maximum of 10.", tsName, tc.name, len(ts.ToolNames))
t.Logf("WARNING: Group %q in config %q has %d tools, which is larger than the recommended maximum of 10.", tsName, tc.name, len(ts.ToolNames))
}
}
})
@@ -2275,13 +2296,13 @@ func TestMergeConfigs(t *testing.T) {
file1 := Config{
Sources: server.SourceConfigs{"source1": httpsrc.Config{Name: "source1"}},
Tools: server.ToolConfigs{"tool1": http.Config{ConfigBase: tools.ConfigBase{Name: "tool1"}}},
Toolsets: server.ToolsetConfigs{"set1": tools.ToolsetConfig{Name: "set1"}},
Groups: server.GroupConfigs{"set1": group.GroupConfig{Name: "set1"}},
EmbeddingModels: server.EmbeddingModelConfigs{"model1": gemini.Config{Name: "gemini-text"}},
}
file2 := Config{
AuthServices: server.AuthServiceConfigs{"auth1": google.Config{Name: "auth1"}},
Tools: server.ToolConfigs{"tool2": http.Config{ConfigBase: tools.ConfigBase{Name: "tool2"}}},
Toolsets: server.ToolsetConfigs{"set2": tools.ToolsetConfig{Name: "set2"}},
Groups: server.GroupConfigs{"set2": group.GroupConfig{Name: "set2"}},
}
fileWithConflicts := Config{
Sources: server.SourceConfigs{"source1": httpsrc.Config{Name: "source1"}},
@@ -2308,8 +2329,8 @@ func TestMergeConfigs(t *testing.T) {
Sources: server.SourceConfigs{"source1": httpsrc.Config{Name: "source1"}},
AuthServices: server.AuthServiceConfigs{"auth1": google.Config{Name: "auth1"}},
Tools: server.ToolConfigs{"tool1": http.Config{ConfigBase: tools.ConfigBase{Name: "tool1"}}, "tool2": http.Config{ConfigBase: tools.ConfigBase{Name: "tool2"}}},
Toolsets: server.ToolsetConfigs{"set1": tools.ToolsetConfig{Name: "set1"}, "set2": tools.ToolsetConfig{Name: "set2"}},
Prompts: server.PromptConfigs{},
Groups: server.GroupConfigs{"set1": group.GroupConfig{Name: "set1"}, "set2": group.GroupConfig{Name: "set2"}},
EmbeddingModels: server.EmbeddingModelConfigs{"model1": gemini.Config{Name: "gemini-text"}},
},
wantErr: false,
@@ -2333,8 +2354,8 @@ func TestMergeConfigs(t *testing.T) {
AuthServices: make(server.AuthServiceConfigs),
EmbeddingModels: server.EmbeddingModelConfigs{"model1": gemini.Config{Name: "gemini-text"}},
Tools: file1.Tools,
Toolsets: file1.Toolsets,
Prompts: server.PromptConfigs{},
Groups: file1.Groups,
},
},
{
@@ -2345,8 +2366,8 @@ func TestMergeConfigs(t *testing.T) {
AuthServices: make(server.AuthServiceConfigs),
EmbeddingModels: make(server.EmbeddingModelConfigs),
Tools: make(server.ToolConfigs),
Toolsets: make(server.ToolsetConfigs),
Prompts: server.PromptConfigs{},
Groups: make(server.GroupConfigs),
},
},
}
@@ -2486,3 +2507,91 @@ tools:
})
}
}
func TestParseConfigGroupNameValidation(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
description string
in string
wantErr bool
errContains string
}{
{
description: "group with integer name is rejected",
in: `
kind: group
name: 123
`,
wantErr: true,
errContains: "missing 'name' field or it is not a string",
},
{
description: "group with boolean name is rejected",
in: `
kind: group
name: true
`,
wantErr: true,
errContains: "missing 'name' field or it is not a string",
},
{
description: "group with absent name is accepted as default group",
in: `
kind: group
description: the default group
`,
wantErr: false,
},
{
description: "group with null name is accepted as default group",
in: `
kind: group
name: ~
description: the default group
`,
wantErr: false,
},
{
description: "group with quoted numeric name is accepted",
in: `
kind: group
name: "123"
`,
wantErr: false,
},
{
description: "non-group resource with integer name is rejected",
in: `
kind: tool
name: 42
type: postgres-sql
source: my-pg-instance
description: some description
statement: SELECT 1
`,
wantErr: true,
errContains: "missing 'name' field or it is not a string",
},
}
for _, tc := range tcs {
t.Run(tc.description, func(t *testing.T) {
parser := ConfigParser{}
_, err := parser.ParseConfig(ctx, testutils.FormatYaml(tc.in))
if tc.wantErr {
if err == nil {
t.Fatal("expected error, got nil")
}
if tc.errContains != "" && !strings.Contains(err.Error(), tc.errContains) {
t.Errorf("error %q does not contain %q", err.Error(), tc.errContains)
}
} else {
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
}
})
}
}
cmd/internal/invoke/command.go
+2 -2
View File
@@ -64,14 +64,14 @@ func runInvoke(cmd *cobra.Command, args []string, opts *internal.ToolboxOptions)
}
// Initialize Resources
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := server.InitializeConfigs(ctx, opts.Cfg)
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, promptsMap, groupsMap, err := server.InitializeConfigs(ctx, opts.Cfg)
if err != nil {
errMsg := fmt.Errorf("failed to initialize resources: %w", err)
opts.Logger.ErrorContext(ctx, errMsg.Error())
return errMsg
}
primitiveMgr := primitives.NewPrimitiveManager(sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap)
primitiveMgr := primitives.NewPrimitiveManager(sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, promptsMap, groupsMap)
// Execute Tool
toolName := args[0]
cmd/internal/options.go
+14 -9
View File
@@ -228,10 +228,15 @@ func (opts *ToolboxOptions) LoadConfig(ctx context.Context, parser *ConfigParser
}
if toolsetName != "" {
targetToolset, exists := parsed.Toolsets[toolsetName]
// Legacy toolsets are folded into groups at unmarshal, so the named
// toolset resolves as a group.
targetGroup, exists := parsed.Groups[toolsetName]
if !exists {
var available []string
for k := range parsed.Toolsets {
for k := range parsed.Groups {
if k == "" {
continue
}
available = append(available, k)
}
slices.Sort(available)
@@ -240,19 +245,19 @@ func (opts *ToolboxOptions) LoadConfig(ctx context.Context, parser *ConfigParser
return isCustomConfigured, errMsg
}
// Filter tools to only include those in the target toolset
// Filter tools to only include those in the target group
filteredTools := make(server.ToolConfigs)
for _, tName := range targetToolset.ToolNames {
for _, tName := range targetGroup.ToolNames {
if tCfg, tExists := parsed.Tools[tName]; tExists {
filteredTools[tName] = tCfg
}
}
parsed.Tools = filteredTools
// Filter toolsets to only include the target toolset
filteredToolsets := make(server.ToolsetConfigs)
filteredToolsets[toolsetName] = targetToolset
parsed.Toolsets = filteredToolsets
// Filter groups to only include the target group
filteredGroups := make(server.GroupConfigs)
filteredGroups[toolsetName] = targetGroup
parsed.Groups = filteredGroups
}
allConfigs = append(allConfigs, parsed)
@@ -294,8 +299,8 @@ func (opts *ToolboxOptions) LoadConfig(ctx context.Context, parser *ConfigParser
opts.Cfg.AuthServiceConfigs = finalConfig.AuthServices
opts.Cfg.EmbeddingModelConfigs = finalConfig.EmbeddingModels
opts.Cfg.ToolConfigs = finalConfig.Tools
opts.Cfg.ToolsetConfigs = finalConfig.Toolsets
opts.Cfg.PromptConfigs = finalConfig.Prompts
opts.Cfg.GroupConfigs = finalConfig.Groups
return isCustomConfigured, nil
}
cmd/internal/skills/command.go
+19 -19
View File
@@ -24,6 +24,7 @@ import (
"strings"
"github.com/googleapis/mcp-toolbox/cmd/internal"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/server"
"github.com/googleapis/mcp-toolbox/internal/server/primitives"
"github.com/googleapis/mcp-toolbox/internal/tools"
@@ -229,51 +230,50 @@ func run(cmd *skillsCmd, opts *internal.ToolboxOptions) error {
}
func (c *skillsCmd) collectTools(ctx context.Context, opts *internal.ToolboxOptions) (map[string]map[string]tools.Tool, error) {
// Initialize tools and toolsets only; skills generation does not need live
// Initialize tools and groups only; skills generation does not need live
// sources, auth services, or embedding models.
toolsMap, toolsetsMap, err := server.InitializeOfflineConfigs(ctx, opts.Cfg)
toolsMap, groupsMap, err := server.InitializeOfflineConfigs(ctx, opts.Cfg)
if err != nil {
return nil, fmt.Errorf("failed to initialize resources: %w", err)
}
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsetsMap, nil, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, nil, groupsMap)
skillsToTools := make(map[string]map[string]tools.Tool)
getToolsFromToolset := func(ts tools.Toolset) map[string]tools.Tool {
toolsetTools := make(map[string]tools.Tool)
for _, t := range ts.Tools {
if t != nil {
tool := *t
toolsetTools[tool.GetName()] = tool
getToolsFromGroup := func(g group.Group) map[string]tools.Tool {
groupTools := make(map[string]tools.Tool)
for _, name := range g.ToolNames {
if tool, ok := toolsMap[name]; ok {
groupTools[name] = tool
}
}
return toolsetTools
return groupTools
}
if c.toolset != "" {
ts, ok := primitiveMgr.GetToolset(c.toolset)
g, ok := primitiveMgr.GetGroup(c.toolset)
if !ok {
return nil, fmt.Errorf("toolset %q not found", c.toolset)
}
skillsToTools[c.name] = getToolsFromToolset(ts)
skillsToTools[c.name] = getToolsFromGroup(g)
return skillsToTools, nil
}
if len(toolsetsMap) <= 1 {
// Default to all tools if no toolset found
if len(groupsMap) <= 1 {
// Default to all tools if no named group found
skillsToTools[c.name] = toolsMap
return skillsToTools, nil
}
// One skill per toolset
for tsName, ts := range toolsetsMap {
if tsName == "" {
// One skill per group
for gName, g := range groupsMap {
if gName == "" {
continue
}
skillName := fmt.Sprintf("%s-%s", c.name, tsName)
skillsToTools[skillName] = getToolsFromToolset(ts)
skillName := fmt.Sprintf("%s-%s", c.name, gName)
skillsToTools[skillName] = getToolsFromGroup(g)
}
return skillsToTools, nil
cmd/root.go
+8 -7
View File
@@ -38,6 +38,7 @@ import (
"github.com/googleapis/mcp-toolbox/cmd/internal/skills"
"github.com/googleapis/mcp-toolbox/internal/auth"
"github.com/googleapis/mcp-toolbox/internal/embeddingmodels"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/server"
"github.com/googleapis/mcp-toolbox/internal/sources"
@@ -138,14 +139,14 @@ func handleDynamicReload(ctx context.Context, toolsFile internal.Config, s *serv
panic(err)
}
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := validateReloadEdits(ctx, toolsFile)
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, promptsMap, groupsMap, err := validateReloadEdits(ctx, toolsFile)
if err != nil {
errMsg := fmt.Errorf("unable to validate reloaded edits: %w", err)
logger.WarnContext(ctx, errMsg.Error())
return err
}
s.PrimitiveMgr.SetPrimitives(sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap)
s.PrimitiveMgr.SetPrimitives(sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, promptsMap, groupsMap)
return nil
}
@@ -153,7 +154,7 @@ func handleDynamicReload(ctx context.Context, toolsFile internal.Config, s *serv
// validateReloadEdits checks that the reloaded config configs can initialized without failing
func validateReloadEdits(
ctx context.Context, toolsFile internal.Config,
) (map[string]sources.Source, map[string]auth.AuthService, map[string]embeddingmodels.EmbeddingModel, map[string]tools.Tool, map[string]tools.Toolset, map[string]prompts.Prompt, map[string]prompts.Promptset, error,
) (map[string]sources.Source, map[string]auth.AuthService, map[string]embeddingmodels.EmbeddingModel, map[string]tools.Tool, map[string]prompts.Prompt, map[string]group.Group, error,
) {
logger, err := util.LoggerFromContext(ctx)
if err != nil {
@@ -176,19 +177,19 @@ func validateReloadEdits(
AuthServiceConfigs: toolsFile.AuthServices,
EmbeddingModelConfigs: toolsFile.EmbeddingModels,
ToolConfigs: toolsFile.Tools,
ToolsetConfigs: toolsFile.Toolsets,
PromptConfigs: toolsFile.Prompts,
GroupConfigs: toolsFile.Groups,
IgnoreUnknownTools: util.IgnoreUnknownToolsFromContext(ctx),
}
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := server.InitializeConfigs(ctx, reloadedConfig)
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, promptsMap, groupsMap, err := server.InitializeConfigs(ctx, reloadedConfig)
if err != nil {
errMsg := fmt.Errorf("unable to initialize reloaded configs: %w", err)
logger.WarnContext(ctx, errMsg.Error())
return nil, nil, nil, nil, nil, nil, nil, err
return nil, nil, nil, nil, nil, nil, err
}
return sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, nil
return sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, promptsMap, groupsMap, nil
}
// Helper to check if a file has a newer ModTime than stored in the map
cmd/root_test.go
+8 -8
View File
@@ -876,18 +876,18 @@ tools:
if len(cfg.ToolConfigs) != 2 {
return fmt.Errorf("expected exactly 2 tools, got %d", len(cfg.ToolConfigs))
}
if _, ok := cfg.ToolsetConfigs["sqlite_database_tools"]; !ok {
return fmt.Errorf("expected toolset 'sqlite_database_tools' not found")
if _, ok := cfg.GroupConfigs["sqlite_database_tools"]; !ok {
return fmt.Errorf("expected group 'sqlite_database_tools' not found")
}
if len(cfg.ToolsetConfigs) != 2 {
// Legacy toolsets are folded into groups, and the default nameless
// collection is seeded later as a derived group, so only the named
// group remains in the parsed config.
if len(cfg.GroupConfigs) != 1 {
var names []string
for k := range cfg.ToolsetConfigs {
for k := range cfg.GroupConfigs {
names = append(names, k)
}
return fmt.Errorf("expected exactly 2 toolsets (including default), got %d: %v", len(cfg.ToolsetConfigs), names)
}
if _, ok := cfg.ToolsetConfigs[""]; !ok {
return fmt.Errorf("expected default toolset '' not found")
return fmt.Errorf("expected exactly 1 group, got %d: %v", len(cfg.GroupConfigs), names)
}
return nil
},
docs/en/documentation/configuration/_index.md
+17
View File
@@ -88,6 +88,23 @@ tools:
- my_third_tool
```
### Groups
The `group` kind scopes MCP primitives such as **tools** and **prompts**
together under one name, with a `description` used as group metadata. A toolset
is a tools-only group, so existing `kind: toolset` configs keep working
unchanged. See [Groups](./groups/_index.md) for details.
```yaml
kind: group
name: my_group
description: Tools and prompts for a specific task.
tools:
- my_first_tool
prompts:
- my_first_prompt
```
### Prompts
The `prompt` kind of your `tools.yaml` defines the templates containing
docs/en/documentation/configuration/groups/_index.md
+68
View File
@@ -0,0 +1,68 @@
---
title: "Groups"
type: docs
weight: 6
description: >
Groups let you scope MCP primitives such as tools and prompts together under a single name, with a description used as group metadata.
---
A Group is a single named collection that scopes MCP primitives together — currently **tools** and **prompts**, with more (such as resources) planned. Where a [Toolset](../toolsets/) groups only tools, a group bundles these primitives under one name and one MCP endpoint, and carries a `description` that describes the collection.
Connecting to a group's endpoint (`/mcp/{name}`) scopes the corresponding MCP list methods (such as `tools/list` and `prompts/list`) to that group.
## Defining Groups
Declare a group as a `kind: group` document in your configuration file. A group has the following fields:
| Field | Required | Description |
| ------------- | -------- | ------------------------------------------------------------------------ |
| `name` | Yes\* | Unique name for the group. Used as the endpoint path (`/mcp/{name}`). |
| `description` | No | Human-readable description of the group. |
| `tools` | No | List of tool names to include in the group. |
| `prompts` | No | List of prompt names to include in the group. |
\* `name` is required for every named group. The single [default group](#the-default-group) omits it.
As Toolbox adds support for more MCP primitives, groups will gain corresponding fields (for example, `resources`).
```yaml
kind: group
name: data_analyst
description: Tools and prompts for exploratory data analysis.
tools:
- list_tables
- execute_sql
prompts:
- summarize_results
---
kind: group
name: admin
description: Administrative operations.
tools:
- create_user
- list_users
```
## The default group
A single **default (nameless) group** always exists and contains **all** configured primitives (every tool and prompt). Connecting to the default MCP endpoint (`/mcp`) returns everything.
You may declare a `kind: group` document with no `name` to set a `description` for the default group. Because the default group always contains everything, it **cannot** declare `tools`, `prompts`, or any other primitive list:
```yaml
kind: group
description: All tools and prompts available on this server.
```
## Validation rules
At startup, Toolbox validates groups:
- **Unique names.** A named group must have a unique name that satisfies the standard name rules (alphanumeric characters, underscores, and hyphens).
- **One default group.** Declaring more than one nameless group is an error.
- **Default group restrictions.** The default group may set only a `description`; declaring `tools`, `prompts`, or any other primitive list on it is an error.
- **Group wins over a same-named toolset.** If a name is defined by both a `kind: toolset` and a `kind: group`, the group takes precedence and Toolbox logs a warning naming the shadowed toolset.
## Relationship to toolsets
Groups are a superset of toolsets: a toolset is equivalent to a tools-only group. Existing `kind: toolset` configurations continue to work unchanged — they are treated as groups with tools and no other primitives, so no migration is required. That said, we recommend migrating to a `kind: group` even for tools-only collections: a group lets you attach a `description` and scope prompts (and, in the future, other primitives) alongside tools. See [Toolsets](../toolsets/) for more.
docs/en/documentation/configuration/prompts/_index.md
+4
View File
@@ -9,6 +9,10 @@ description: >
A `prompt` represents a reusable prompt template that can be retrieved and used
by MCP clients.
{{< notice note >}}
You can use [Groups](../groups/) to organize prompts into collections. When you connect to a group's endpoint, `prompts/list` returns only the prompts in that group. The default endpoint returns all prompts.
{{< /notice >}}
A Prompt is essentially a template for a message or a series of messages that
can be sent to a Large Language Model (LLM). The Toolbox server implements the
`prompts/list` and `prompts/get` methods from the [Model Context Protocol
docs/en/documentation/configuration/toolsets/_index.md
+4
View File
@@ -14,6 +14,10 @@ This is especially useful when you are building a system with multiple AI agents
Try organizing your toolsets by the agent's persona or app feature (e.g., `data_analyst_set` vs `customer_support_set`). This keeps your client-side code clean and ensures an agent isn't distracted by tools it doesn't need.
{{< /notice >}}
{{< notice note >}}
A toolset is a tools-only [Group](../groups/). `kind: toolset` continues to work unchanged and needs no migration. That said, we recommend migrating to `kind: group` — even for tools-only collections — because a group lets you attach a `description` and scope other MCP primitives such as **prompts** alongside your **tools**.
{{< /notice >}}
## Defining Toolsets
In your configuration file, define each toolset by providing a unique `name` and a list of `tools` that belong to that group..
internal/group/group.go
+119
View File
@@ -0,0 +1,119 @@
// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package group
import (
"fmt"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/sources"
"github.com/googleapis/mcp-toolbox/internal/tools"
)
// GroupConfig is the parsed configuration for a group: a single named collection
// that holds both tools and prompts. Its description doubles as the MCP server
// instructions for clients connected to the group.
type GroupConfig struct {
Name string `yaml:"name"`
Description string `yaml:"description"`
ToolNames []string `yaml:"tools"`
PromptNames []string `yaml:"prompts"`
}
// Group is an initialized group: the source of truth for a named collection of
// tools and prompts. It is a self-contained MCP primitive and does not depend on
// the legacy toolset/promptset types. It keeps O(1) membership sets for its tools
// and prompts; per-tool and per-prompt manifests are generated on demand by
// callers from the resolved tools and prompts maps.
type Group struct {
GroupConfig
toolNameSet map[string]struct{}
promptNameSet map[string]struct{}
}
// Initialize validates the group name and checks that every declared tool and
// prompt exists in the provided maps, building the membership sets used by
// ContainsTool and ContainsPrompt.
func (gc GroupConfig) Initialize(toolsMap map[string]tools.Tool, promptsMap map[string]prompts.Prompt) (Group, error) {
if !tools.IsValidName(gc.Name) {
return Group{}, fmt.Errorf("invalid group name: %q", gc.Name)
}
toolNameSet := make(map[string]struct{}, len(gc.ToolNames))
for _, name := range gc.ToolNames {
if _, ok := toolsMap[name]; !ok {
return Group{}, fmt.Errorf("tool does not exist: %q", name)
}
toolNameSet[name] = struct{}{}
}
promptNameSet := make(map[string]struct{}, len(gc.PromptNames))
for _, name := range gc.PromptNames {
if _, ok := promptsMap[name]; !ok {
return Group{}, fmt.Errorf("prompt does not exist: %q", name)
}
promptNameSet[name] = struct{}{}
}
return Group{GroupConfig: gc, toolNameSet: toolNameSet, promptNameSet: promptNameSet}, nil
}
// NewGroup builds a Group directly from its config, deriving the membership sets
// from the declared tool and prompt names. It skips the existence checks
// performed by GroupConfig.Initialize and is intended for tests and callers that
// have already validated the names.
func NewGroup(config GroupConfig) Group {
toolNameSet := make(map[string]struct{}, len(config.ToolNames))
for _, name := range config.ToolNames {
toolNameSet[name] = struct{}{}
}
promptNameSet := make(map[string]struct{}, len(config.PromptNames))
for _, name := range config.PromptNames {
promptNameSet[name] = struct{}{}
}
return Group{GroupConfig: config, toolNameSet: toolNameSet, promptNameSet: promptNameSet}
}
// ContainsTool reports whether the group includes a tool with the given name.
func (g Group) ContainsTool(name string) bool {
_, ok := g.toolNameSet[name]
return ok
}
// ContainsPrompt reports whether the group includes a prompt with the given name.
func (g Group) ContainsPrompt(name string) bool {
_, ok := g.promptNameSet[name]
return ok
}
// ToolsetManifest builds a tools.ToolsetManifest for the group's tools, resolving
// each declared tool name against toolsMap and generating its manifest from srcs.
// The group holds tool names rather than tool pointers, so callers pass the
// resolved tools and sources maps.
func (g Group) ToolsetManifest(serverVersion string, toolsMap map[string]tools.Tool, srcs map[string]sources.Source) (tools.ToolsetManifest, error) {
toolsManifest := make(map[string]tools.Manifest, len(g.ToolNames))
for _, name := range g.ToolNames {
tool, ok := toolsMap[name]
if !ok {
return tools.ToolsetManifest{}, fmt.Errorf("tool does not exist: %s", name)
}
m, err := tool.Manifest(srcs)
if err != nil {
return tools.ToolsetManifest{}, fmt.Errorf("error generating manifest for tool %q: %w", name, err)
}
toolsManifest[name] = m
}
return tools.ToolsetManifest{ServerVersion: serverVersion, ToolsManifest: toolsManifest}, nil
}
internal/group/group_test.go
+217
View File
@@ -0,0 +1,217 @@
// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package group_test
import (
"slices"
"strings"
"testing"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/testutils"
"github.com/googleapis/mcp-toolbox/internal/tools"
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
)
func testFixtures() (map[string]tools.Tool, map[string]prompts.Prompt) {
toolsMap := map[string]tools.Tool{
"tool1": testutils.NewMockTool("tool1", "first tool", []parameters.Parameter{}, false, false),
"tool2": testutils.NewMockTool("tool2", "second tool", []parameters.Parameter{}, false, false),
}
promptsMap := map[string]prompts.Prompt{
"prompt1": testutils.NewMockPrompt("prompt1", "first prompt", prompts.Arguments{}),
"prompt2": testutils.NewMockPrompt("prompt2", "second prompt", prompts.Arguments{}),
}
return toolsMap, promptsMap
}
func TestGroupConfig_Initialize(t *testing.T) {
t.Parallel()
toolsMap, promptsMap := testFixtures()
testCases := []struct {
name string
config group.GroupConfig
wantTools []string
wantPrompts []string
wantErr string
}{
{
name: "tools and prompts",
config: group.GroupConfig{
Name: "mygroup",
Description: "a group",
ToolNames: []string{"tool1", "tool2"},
PromptNames: []string{"prompt1", "prompt2"},
},
wantTools: []string{"tool1", "tool2"},
wantPrompts: []string{"prompt1", "prompt2"},
},
{
name: "tools only",
config: group.GroupConfig{
Name: "toolsonly",
ToolNames: []string{"tool1"},
},
wantTools: []string{"tool1"},
wantPrompts: nil,
},
{
name: "prompts only",
config: group.GroupConfig{
Name: "promptsonly",
PromptNames: []string{"prompt1"},
},
wantTools: nil,
wantPrompts: []string{"prompt1"},
},
{
name: "default nameless group",
config: group.GroupConfig{
Name: "",
ToolNames: []string{"tool1"},
PromptNames: []string{"prompt1"},
},
wantTools: []string{"tool1"},
wantPrompts: []string{"prompt1"},
},
{
name: "invalid group name",
config: group.GroupConfig{
Name: "bad name!",
ToolNames: []string{"tool1"},
},
wantErr: "invalid group name",
},
{
name: "missing tool",
config: group.GroupConfig{
Name: "g",
ToolNames: []string{"nope"},
},
wantErr: "tool does not exist: \"nope\"",
},
{
name: "missing prompt",
config: group.GroupConfig{
Name: "g",
PromptNames: []string{"nope"},
},
wantErr: "prompt does not exist: \"nope\"",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
g, err := tc.config.Initialize(toolsMap, promptsMap)
if tc.wantErr != "" {
if err == nil {
t.Fatalf("expected error containing %q, got nil", tc.wantErr)
}
if !strings.Contains(err.Error(), tc.wantErr) {
t.Fatalf("error = %q, want it to contain %q", err.Error(), tc.wantErr)
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if !slices.Equal(g.ToolNames, tc.wantTools) {
t.Errorf("tools = %v, want %v", g.ToolNames, tc.wantTools)
}
if !slices.Equal(g.PromptNames, tc.wantPrompts) {
t.Errorf("prompts = %v, want %v", g.PromptNames, tc.wantPrompts)
}
for _, name := range tc.wantTools {
if !g.ContainsTool(name) {
t.Errorf("group missing tool %q", name)
}
}
for _, name := range tc.wantPrompts {
if !g.ContainsPrompt(name) {
t.Errorf("group missing prompt %q", name)
}
}
})
}
}
func TestGroup_ToolsetManifest(t *testing.T) {
t.Parallel()
toolsMap, _ := testFixtures()
g := group.NewGroup(group.GroupConfig{
Name: "mygroup",
ToolNames: []string{"tool1", "tool2"},
})
manifest, err := g.ToolsetManifest("v1.2.3", toolsMap, nil)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if manifest.ServerVersion != "v1.2.3" {
t.Errorf("ServerVersion = %q, want %q", manifest.ServerVersion, "v1.2.3")
}
wantTools := []string{"tool1", "tool2"}
gotTools := make([]string, 0, len(manifest.ToolsManifest))
for name := range manifest.ToolsManifest {
gotTools = append(gotTools, name)
}
slices.Sort(gotTools)
if !slices.Equal(gotTools, wantTools) {
t.Errorf("tools manifest keys = %v, want %v", gotTools, wantTools)
}
if manifest.ToolsManifest["tool1"].Description != "first tool" {
t.Errorf("tool1 description = %q, want %q", manifest.ToolsManifest["tool1"].Description, "first tool")
}
// Missing tool error path.
missing := group.NewGroup(group.GroupConfig{Name: "g", ToolNames: []string{"nope"}})
if _, err := missing.ToolsetManifest("v1", toolsMap, nil); err == nil {
t.Fatal("expected error for missing tool, got nil")
} else if !strings.Contains(err.Error(), "tool does not exist: nope") {
t.Errorf("error = %q, want it to contain %q", err.Error(), "tool does not exist: nope")
}
}
func TestGroup_Contains(t *testing.T) {
t.Parallel()
toolsMap, promptsMap := testFixtures()
g, err := group.GroupConfig{
Name: "mygroup",
Description: "a group",
ToolNames: []string{"tool1", "tool2"},
PromptNames: []string{"prompt1", "prompt2"},
}.Initialize(toolsMap, promptsMap)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if !g.ContainsTool("tool1") || !g.ContainsTool("tool2") {
t.Errorf("group missing expected tools")
}
if g.ContainsTool("tool3") {
t.Errorf("group reports an absent tool")
}
if !g.ContainsPrompt("prompt1") || !g.ContainsPrompt("prompt2") {
t.Errorf("group missing expected prompts")
}
if g.ContainsPrompt("prompt3") {
t.Errorf("group reports an absent prompt")
}
}
internal/prompts/promptsets.go
-80
View File
@@ -1,80 +0,0 @@
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package prompts
import (
"fmt"
"github.com/googleapis/mcp-toolbox/internal/tools"
)
type PromptsetConfig struct {
Name string `yaml:"name"`
PromptNames []string `yaml:",inline"`
}
type Promptset struct {
PromptsetConfig
Prompts []*Prompt `yaml:",inline"`
Manifest PromptsetManifest `yaml:",inline"`
PromptNameSet map[string]struct{}
}
func (p Promptset) ToConfig() PromptsetConfig {
return p.PromptsetConfig
}
// ContainsPrompt reports whether the promptset includes a prompt with the given name.
// When built via Initialize, lookups are O(1) via promptNameSet; for Promptsets
// constructed directly (e.g., in tests), falls back to a linear scan of PromptNames.
func (p Promptset) ContainsPrompt(name string) bool {
if p.PromptNameSet != nil {
_, ok := p.PromptNameSet[name]
return ok
}
return false
}
type PromptsetManifest struct {
ServerVersion string `json:"serverVersion"`
PromptsManifest map[string]Manifest `json:"prompts"`
}
func (p PromptsetConfig) Initialize(serverVersion string, promptsMap map[string]Prompt) (Promptset, error) {
// Check each declared prompt name exists
promptset := Promptset{
PromptsetConfig: p,
Prompts: make([]*Prompt, 0, len(p.PromptNames)),
Manifest: PromptsetManifest{
ServerVersion: serverVersion,
PromptsManifest: make(map[string]Manifest, len(p.PromptNames)),
},
PromptNameSet: make(map[string]struct{}, len(p.PromptNames)),
}
if !tools.IsValidName(promptset.Name) {
return promptset, fmt.Errorf("invalid promptset name: %s", promptset.Name)
}
for _, promptName := range p.PromptNames {
prompt, ok := promptsMap[promptName]
if !ok {
return promptset, fmt.Errorf("prompt does not exist: %s", promptName)
}
promptset.Prompts = append(promptset.Prompts, &prompt)
promptset.Manifest.PromptsManifest[promptName] = prompt.Manifest()
promptset.PromptNameSet[promptName] = struct{}{}
}
return promptset, nil
}
internal/prompts/promptsets_test.go
-265
View File
@@ -1,265 +0,0 @@
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package prompts_test
import (
"strings"
"testing"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/testutils"
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
)
func TestPromptset_ContainsPrompt(t *testing.T) {
t.Parallel()
promptset := prompts.Promptset{
PromptsetConfig: prompts.PromptsetConfig{
Name: "test-promptset",
PromptNames: []string{"greet", "summarize"},
},
PromptNameSet: map[string]struct{}{"greet": struct{}{}, "summarize": struct{}{}},
}
tests := []struct {
name string
promptName string
want bool
}{
{
name: "prompt exists in promptset",
promptName: "greet",
want: true,
},
{
name: "another prompt exists in promptset",
promptName: "summarize",
want: true,
},
{
name: "prompt not in promptset",
promptName: "admin_prompt",
want: false,
},
{
name: "empty prompt name",
promptName: "",
want: false,
},
}
for _, tc := range tests {
tc := tc
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
got := promptset.ContainsPrompt(tc.promptName)
if got != tc.want {
t.Errorf("ContainsPrompt(%q) = %v, want %v", tc.promptName, got, tc.want)
}
})
}
}
func TestPromptset_ContainsPrompt_EmptyPromptset(t *testing.T) {
t.Parallel()
promptset := prompts.Promptset{
PromptsetConfig: prompts.PromptsetConfig{
Name: "empty-promptset",
PromptNames: []string{},
},
}
if promptset.ContainsPrompt("anything") {
t.Error("ContainsPrompt should return false for empty promptset")
}
}
func TestPromptsetConfig_Initialize(t *testing.T) {
t.Parallel()
args := prompts.Arguments{
{Parameter: parameters.NewStringParameter("arg1", "Test argument")},
}
promptsMap := map[string]prompts.Prompt{
"prompt1": testutils.NewMockPrompt("prompt1", "First test prompt", args),
"prompt2": testutils.NewMockPrompt("prompt2", "Second test prompt", args),
}
serverVersion := "v1.0.0"
p1 := promptsMap["prompt1"]
p2 := promptsMap["prompt2"]
prompt1Ptr := &p1
prompt2Ptr := &p2
testCases := []struct {
name string
config prompts.PromptsetConfig
want prompts.Promptset
wantErr string
}{
{
name: "Success case",
config: prompts.PromptsetConfig{
Name: "default",
PromptNames: []string{"prompt1", "prompt2"},
},
want: prompts.Promptset{
PromptsetConfig: prompts.PromptsetConfig{
Name: "default",
PromptNames: []string{"prompt1", "prompt2"},
},
Prompts: []*prompts.Prompt{
prompt1Ptr,
prompt2Ptr,
},
Manifest: prompts.PromptsetManifest{
ServerVersion: serverVersion,
PromptsManifest: map[string]prompts.Manifest{
"prompt1": promptsMap["prompt1"].Manifest(),
"prompt2": promptsMap["prompt2"].Manifest(),
},
},
PromptNameSet: map[string]struct{}{"prompt1": struct{}{}, "prompt2": struct{}{}},
},
wantErr: "",
},
{
name: "Success case with one prompt",
config: prompts.PromptsetConfig{
Name: "single",
PromptNames: []string{"prompt1"},
},
want: prompts.Promptset{
PromptsetConfig: prompts.PromptsetConfig{
Name: "single",
PromptNames: []string{"prompt1"},
},
Prompts: []*prompts.Prompt{
prompt1Ptr,
},
Manifest: prompts.PromptsetManifest{
ServerVersion: serverVersion,
PromptsManifest: map[string]prompts.Manifest{
"prompt1": promptsMap["prompt1"].Manifest(),
},
},
PromptNameSet: map[string]struct{}{"prompt1": struct{}{}},
},
wantErr: "",
},
{
name: "Failure case - invalid promptset name",
config: prompts.PromptsetConfig{
Name: "invalid name", // Contains a space
PromptNames: []string{"prompt1"},
},
want: prompts.Promptset{
PromptsetConfig: prompts.PromptsetConfig{
Name: "invalid name",
PromptNames: []string{"prompt1"},
},
Prompts: []*prompts.Prompt{},
Manifest: prompts.PromptsetManifest{
ServerVersion: serverVersion,
PromptsManifest: map[string]prompts.Manifest{},
},
PromptNameSet: map[string]struct{}{},
},
wantErr: "invalid promptset name",
},
{
name: "Failure case - prompt not found",
config: prompts.PromptsetConfig{
Name: "missing_prompt",
PromptNames: []string{"prompt1", "prompt_does_not_exist"},
},
// Expect partial struct with fields populated up to the error
want: prompts.Promptset{
PromptsetConfig: prompts.PromptsetConfig{
Name: "missing_prompt",
PromptNames: []string{"prompt1", "prompt_does_not_exist"},
},
Prompts: []*prompts.Prompt{
prompt1Ptr,
},
Manifest: prompts.PromptsetManifest{
ServerVersion: serverVersion,
PromptsManifest: map[string]prompts.Manifest{
"prompt1": promptsMap["prompt1"].Manifest(),
},
},
PromptNameSet: map[string]struct{}{"prompt1": struct{}{}},
},
wantErr: "prompt does not exist",
},
{
name: "Success case - empty prompt list",
config: prompts.PromptsetConfig{
Name: "empty",
PromptNames: []string{},
},
want: prompts.Promptset{
PromptsetConfig: prompts.PromptsetConfig{
Name: "empty",
PromptNames: []string{},
},
Prompts: []*prompts.Prompt{},
Manifest: prompts.PromptsetManifest{
ServerVersion: serverVersion,
PromptsManifest: map[string]prompts.Manifest{},
},
PromptNameSet: map[string]struct{}{},
},
wantErr: "",
},
}
for _, tc := range testCases {
tc := tc
t.Run(tc.name, func(t *testing.T) {
got, err := tc.config.Initialize(serverVersion, promptsMap)
if tc.wantErr != "" {
if err == nil {
t.Fatalf("Initialize() expected error but got nil")
}
if !strings.Contains(err.Error(), tc.wantErr) {
t.Errorf("Initialize() error mismatch:\n want to contain: %q\n got: %q", tc.wantErr, err.Error())
}
if diff := cmp.Diff(tc.want, got, cmp.AllowUnexported(testutils.MockPrompt{}), cmpopts.IgnoreUnexported(prompts.Promptset{})); diff != "" {
t.Errorf("Initialize() partial result on error mismatch (-want +got):\n%s", diff)
}
} else {
if err != nil {
t.Fatalf("Initialize() returned unexpected error: %v", err)
}
// Using cmp.AllowUnexported because MockPrompt is unexported
if diff := cmp.Diff(tc.want, got, cmp.AllowUnexported(testutils.MockPrompt{}), cmpopts.IgnoreUnexported(prompts.Promptset{})); diff != "" {
t.Errorf("Initialize() result mismatch (-want +got):\n%s", diff)
}
gotConfig := got.ToConfig()
if diff := cmp.Diff(tc.config, gotConfig); diff != "" {
t.Errorf("ToConfig() result mismatch (-want +got):\n%s", diff)
}
}
})
}
}
internal/server/api.go
+2 -2
View File
@@ -66,7 +66,7 @@ func toolsetHandler(s *Server, w http.ResponseWriter, r *http.Request) {
span.End()
}()
toolset, ok := s.PrimitiveMgr.GetToolset(toolsetName)
g, ok := s.PrimitiveMgr.GetGroup(toolsetName)
if !ok {
err = fmt.Errorf("toolset %q does not exist", toolsetName)
s.logger.DebugContext(ctx, err.Error())
@@ -74,7 +74,7 @@ func toolsetHandler(s *Server, w http.ResponseWriter, r *http.Request) {
return
}
manifest, err := toolset.BuildManifest(s.PrimitiveMgr.GetSourcesMap())
manifest, err := g.ToolsetManifest(s.version, s.PrimitiveMgr.GetToolsMap(), s.PrimitiveMgr.GetSourcesMap())
if err != nil {
s.logger.DebugContext(ctx, err.Error())
_ = render.Render(w, r, newErrResponse(err, http.StatusInternalServerError))
internal/server/api_test.go
+10 -10
View File
@@ -29,8 +29,8 @@ import (
func TestToolsetEndpoint(t *testing.T) {
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, toolsets, _, _ := testutils.SetUpResources(t, mockTools, nil)
r, shutdown := setUpServer(t, "api", toolsMap, toolsets, nil, nil)
toolsMap, _, groups := testutils.SetUpResources(t, mockTools, nil)
r, shutdown := setUpServer(t, "api", toolsMap, nil, groups)
defer shutdown()
ts := runServer(r, false)
defer ts.Close()
@@ -126,8 +126,8 @@ func TestToolsetEndpoint(t *testing.T) {
func TestToolGetEndpoint(t *testing.T) {
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, toolsets, _, _ := testutils.SetUpResources(t, mockTools, nil)
r, shutdown := setUpServer(t, "api", toolsMap, toolsets, nil, nil)
toolsMap, _, groups := testutils.SetUpResources(t, mockTools, nil)
r, shutdown := setUpServer(t, "api", toolsMap, nil, groups)
defer shutdown()
ts := runServer(r, false)
defer ts.Close()
@@ -214,8 +214,8 @@ func TestToolGetEndpoint(t *testing.T) {
func TestToolInvokeEndpoint(t *testing.T) {
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2, testutils.MockTool4, testutils.MockTool5}
toolsMap, toolsets, _, _ := testutils.SetUpResources(t, mockTools, nil)
r, shutdown := setUpServer(t, "api", toolsMap, toolsets, nil, nil)
toolsMap, _, groups := testutils.SetUpResources(t, mockTools, nil)
r, shutdown := setUpServer(t, "api", toolsMap, nil, groups)
defer shutdown()
ts := runServer(r, false)
defer ts.Close()
@@ -299,8 +299,8 @@ func TestToolInvokeEndpoint(t *testing.T) {
func TestApiRequestBodyLimit(t *testing.T) {
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, toolsets, _, _ := testutils.SetUpResources(t, mockTools, nil)
r, shutdown := setUpServer(t, "api", toolsMap, toolsets, nil, nil)
toolsMap, _, groups := testutils.SetUpResources(t, mockTools, nil)
r, shutdown := setUpServer(t, "api", toolsMap, nil, groups)
defer shutdown()
ts := runServer(r, false)
defer ts.Close()
@@ -331,9 +331,9 @@ func TestApiRequestBodyLimit(t *testing.T) {
func TestApiRequestBodyLimitOverride(t *testing.T) {
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, toolsets, _, _ := testutils.SetUpResources(t, mockTools, nil)
toolsMap, _, groups := testutils.SetUpResources(t, mockTools, nil)
customLimit := int64(1 << 20)
r, shutdown := setUpServer(t, "api", toolsMap, toolsets, nil, nil, withHTTPMaxRequestBytes(customLimit))
r, shutdown := setUpServer(t, "api", toolsMap, nil, groups, withHTTPMaxRequestBytes(customLimit))
defer shutdown()
ts := runServer(r, false)
defer ts.Close()
internal/server/common_test.go
+4 -3
View File
@@ -24,6 +24,7 @@ import (
"testing"
"github.com/go-chi/chi/v5"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/log"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/server/primitives"
@@ -38,8 +39,8 @@ var (
_ prompts.Prompt = testutils.MockPrompt{}
)
// setUpServer create a new server with tools, toolsets, prompts, and promptsets.
func setUpServer(t *testing.T, router string, tools map[string]tools.Tool, toolsets map[string]tools.Toolset, prompts map[string]prompts.Prompt, promptsets map[string]prompts.Promptset, opts ...func(*Server)) (chi.Router, func()) {
// setUpServer create a new server with tools, prompts, and groups.
func setUpServer(t *testing.T, router string, tools map[string]tools.Tool, prompts map[string]prompts.Prompt, groups map[string]group.Group, opts ...func(*Server)) (chi.Router, func()) {
ctx, cancel := context.WithCancel(context.Background())
testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info")
@@ -59,7 +60,7 @@ func setUpServer(t *testing.T, router string, tools map[string]tools.Tool, tools
sseManager := newSseManager(ctx)
primitiveManager := primitives.NewPrimitiveManager(nil, nil, nil, tools, toolsets, prompts, promptsets)
primitiveManager := primitives.NewPrimitiveManager(nil, nil, nil, tools, prompts, groups)
server := Server{
version: testutils.MockVersionString,
internal/server/config.go
+95 -12
View File
@@ -30,6 +30,7 @@ import (
"github.com/googleapis/mcp-toolbox/internal/auth/google"
"github.com/googleapis/mcp-toolbox/internal/embeddingmodels"
"github.com/googleapis/mcp-toolbox/internal/embeddingmodels/gemini"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/sources"
"github.com/googleapis/mcp-toolbox/internal/tools"
@@ -55,12 +56,11 @@ type ServerConfig struct {
EmbeddingModelConfigs EmbeddingModelConfigs
// ToolConfigs defines what tools are available.
ToolConfigs ToolConfigs
// ToolsetConfigs defines what tools are available.
ToolsetConfigs ToolsetConfigs
// PromptConfigs defines what prompts are available
PromptConfigs PromptConfigs
// PromptsetConfigs defines what prompts are available
PromptsetConfigs PromptsetConfigs
// GroupConfigs defines groups of tools and prompts declared via `kind: group`
// (legacy `kind: toolset` configs are folded into groups at unmarshal).
GroupConfigs GroupConfigs
// IgnoreUnknownTools logs warnings and skips unknown/unsupported tool types instead of failing to start.
IgnoreUnknownTools bool
// LoggingFormat defines whether structured loggings are used.
@@ -159,18 +159,21 @@ type SourceConfigs map[string]sources.SourceConfig
type AuthServiceConfigs map[string]auth.AuthServiceConfig
type EmbeddingModelConfigs map[string]embeddingmodels.EmbeddingModelConfig
type ToolConfigs map[string]tools.ToolConfig
type ToolsetConfigs map[string]tools.ToolsetConfig
type PromptConfigs map[string]prompts.PromptConfig
type PromptsetConfigs map[string]prompts.PromptsetConfig
type GroupConfigs map[string]group.GroupConfig
func UnmarshalPrimitiveConfig(ctx context.Context, raw []byte) (SourceConfigs, AuthServiceConfigs, EmbeddingModelConfigs, ToolConfigs, ToolsetConfigs, PromptConfigs, error) {
func UnmarshalPrimitiveConfig(ctx context.Context, raw []byte) (SourceConfigs, AuthServiceConfigs, EmbeddingModelConfigs, ToolConfigs, PromptConfigs, GroupConfigs, error) {
// prepare configs map
var sourceConfigs SourceConfigs
var authServiceConfigs AuthServiceConfigs
var embeddingModelConfigs EmbeddingModelConfigs
var toolConfigs ToolConfigs
var toolsetConfigs ToolsetConfigs
var promptConfigs PromptConfigs
var groupConfigs GroupConfigs
// Legacy `kind: toolset` configs are collected here as tools-only groups, then
// folded into groupConfigs after the loop so explicit `kind: group` definitions
// take precedence regardless of document order.
var toolsetGroups map[string]group.GroupConfig
// promptset configs is not yet supported
file, err := parser.ParseBytes(raw, 0)
@@ -200,6 +203,15 @@ func UnmarshalPrimitiveConfig(ctx context.Context, raw []byte) (SourceConfigs, A
return nil, nil, nil, nil, nil, nil, fmt.Errorf("missing 'kind' field or it is not a string: %v", resource)
}
if name, ok = resource["name"].(string); !ok {
// A `kind: group` may omit `name` to target the default nameless group;
// every other resource requires a name.
if kind == "group" {
if rawName, present := resource["name"]; !present || rawName == nil {
name, ok = "", true
}
}
}
if !ok {
if len(file.Docs) > 1 {
fallbackToken := keyToken(doc.Body, "name")
if fallbackToken == nil {
@@ -224,6 +236,9 @@ func UnmarshalPrimitiveConfig(ctx context.Context, raw []byte) (SourceConfigs, A
if sourceConfigs == nil {
sourceConfigs = make(SourceConfigs)
}
if _, exists := sourceConfigs[name]; exists {
return nil, nil, nil, nil, nil, nil, fmt.Errorf("source %q declared more than once", name)
}
sourceConfigs[name] = c
case "authService":
c, err := UnmarshalYAMLAuthServiceConfig(ctx, name, resource)
@@ -236,6 +251,9 @@ func UnmarshalPrimitiveConfig(ctx context.Context, raw []byte) (SourceConfigs, A
if authServiceConfigs == nil {
authServiceConfigs = make(AuthServiceConfigs)
}
if _, exists := authServiceConfigs[name]; exists {
return nil, nil, nil, nil, nil, nil, fmt.Errorf("authService %q declared more than once", name)
}
authServiceConfigs[name] = c
case "tool":
c, err := UnmarshalYAMLToolConfig(ctx, name, resource)
@@ -251,6 +269,9 @@ func UnmarshalPrimitiveConfig(ctx context.Context, raw []byte) (SourceConfigs, A
if toolConfigs == nil {
toolConfigs = make(ToolConfigs)
}
if _, exists := toolConfigs[name]; exists {
return nil, nil, nil, nil, nil, nil, fmt.Errorf("tool %q declared more than once", name)
}
toolConfigs[name] = c
case "toolset":
c, err := UnmarshalYAMLToolsetConfig(ctx, name, resource)
@@ -260,10 +281,13 @@ func UnmarshalPrimitiveConfig(ctx context.Context, raw []byte) (SourceConfigs, A
}
return nil, nil, nil, nil, nil, nil, fmt.Errorf("error unmarshaling %s: %w", kind, err)
}
if toolsetConfigs == nil {
toolsetConfigs = make(ToolsetConfigs)
if toolsetGroups == nil {
toolsetGroups = make(map[string]group.GroupConfig)
}
toolsetConfigs[name] = c
if _, exists := toolsetGroups[name]; exists {
return nil, nil, nil, nil, nil, nil, fmt.Errorf("toolset %q declared more than once", name)
}
toolsetGroups[name] = group.GroupConfig{Name: name, ToolNames: c.ToolNames}
case "embeddingModel":
c, err := UnmarshalYAMLEmbeddingModelConfig(ctx, name, resource)
if err != nil {
@@ -275,6 +299,9 @@ func UnmarshalPrimitiveConfig(ctx context.Context, raw []byte) (SourceConfigs, A
if embeddingModelConfigs == nil {
embeddingModelConfigs = make(EmbeddingModelConfigs)
}
if _, exists := embeddingModelConfigs[name]; exists {
return nil, nil, nil, nil, nil, nil, fmt.Errorf("embeddingModel %q declared more than once", name)
}
embeddingModelConfigs[name] = c
case "prompt":
c, err := UnmarshalYAMLPromptConfig(ctx, name, resource)
@@ -287,7 +314,28 @@ func UnmarshalPrimitiveConfig(ctx context.Context, raw []byte) (SourceConfigs, A
if promptConfigs == nil {
promptConfigs = make(PromptConfigs)
}
if _, exists := promptConfigs[name]; exists {
return nil, nil, nil, nil, nil, nil, fmt.Errorf("prompt %q declared more than once", name)
}
promptConfigs[name] = c
case "group":
c, err := UnmarshalYAMLGroupConfig(ctx, name, resource)
if err != nil {
if len(file.Docs) > 1 {
return nil, nil, nil, nil, nil, nil, fmt.Errorf("document %d: error unmarshaling %s %q: %w", docIndex, kind, name, err)
}
return nil, nil, nil, nil, nil, nil, fmt.Errorf("error unmarshaling %s: %w", kind, err)
}
if groupConfigs == nil {
groupConfigs = make(GroupConfigs)
}
if _, exists := groupConfigs[name]; exists {
if name == "" {
return nil, nil, nil, nil, nil, nil, fmt.Errorf("more than one default (nameless) group declared; only one is allowed")
}
return nil, nil, nil, nil, nil, nil, fmt.Errorf("group %q declared more than once", name)
}
groupConfigs[name] = c
default:
if len(file.Docs) > 1 {
return nil, nil, nil, nil, nil, nil, fmt.Errorf("%s invalid kind %q", formatDocLocation(docIndex, keyToken(doc.Body, "kind"), doc.Body), kind)
@@ -295,7 +343,25 @@ func UnmarshalPrimitiveConfig(ctx context.Context, raw []byte) (SourceConfigs, A
return nil, nil, nil, nil, nil, nil, fmt.Errorf("invalid kind %s", kind)
}
}
return sourceConfigs, authServiceConfigs, embeddingModelConfigs, toolConfigs, toolsetConfigs, promptConfigs, nil
// Fold legacy toolsets into groups. An explicit `kind: group` of the same name
// takes precedence over a toolset (matching the prior server-side behavior); warn
// when a toolset is shadowed this way.
if len(toolsetGroups) > 0 {
if groupConfigs == nil {
groupConfigs = make(GroupConfigs)
}
for name, gc := range toolsetGroups {
if _, shadowed := groupConfigs[name]; shadowed {
if l, err := util.LoggerFromContext(ctx); err == nil {
l.WarnContext(ctx, fmt.Sprintf("group %q shadows a toolset of the same name; using the group definition", name))
}
continue
}
groupConfigs[name] = gc
}
}
return sourceConfigs, authServiceConfigs, embeddingModelConfigs, toolConfigs, promptConfigs, groupConfigs, nil
}
func UnmarshalYAMLSourceConfig(ctx context.Context, name string, r map[string]any) (sources.SourceConfig, error) {
@@ -490,6 +556,23 @@ func UnmarshalYAMLToolsetConfig(ctx context.Context, name string, r map[string]a
return tools.ToolsetConfig{Name: name, ToolNames: raw["tools"]}, nil
}
func UnmarshalYAMLGroupConfig(ctx context.Context, name string, r map[string]any) (group.GroupConfig, error) {
dec, err := util.NewStrictDecoder(r)
if err != nil {
return group.GroupConfig{}, fmt.Errorf("error creating decoder: %s", err)
}
gc := group.GroupConfig{Name: name}
if err := dec.DecodeContext(ctx, &gc); err != nil {
return group.GroupConfig{}, fmt.Errorf("unable to unmarshal group: %s", err)
}
// The default (nameless) group always contains all configured tools and
// prompts, so it may only set a description.
if name == "" && (len(gc.ToolNames) > 0 || len(gc.PromptNames) > 0) {
return group.GroupConfig{}, fmt.Errorf("the default (nameless) group cannot declare 'tools' or 'prompts'; it always contains all configured tools and prompts")
}
return gc, nil
}
func UnmarshalYAMLPromptConfig(ctx context.Context, name string, r map[string]any) (prompts.PromptConfig, error) {
// Look for the 'type' field. If it's not present, typeStr will be an
// empty string, which prompts.DecodeConfig will correctly default to "custom".
internal/server/mcp.go
+23 -29
View File
@@ -33,13 +33,12 @@ import (
"github.com/go-chi/render"
"github.com/google/uuid"
"github.com/googleapis/mcp-toolbox/internal/auth"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/server/mcp"
"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
mcputil "github.com/googleapis/mcp-toolbox/internal/server/mcp/util"
v20241105 "github.com/googleapis/mcp-toolbox/internal/server/mcp/v20241105"
v20250326 "github.com/googleapis/mcp-toolbox/internal/server/mcp/v20250326"
"github.com/googleapis/mcp-toolbox/internal/tools"
"github.com/googleapis/mcp-toolbox/internal/util"
"go.opentelemetry.io/otel"
"go.opentelemetry.io/otel/attribute"
@@ -276,7 +275,7 @@ func (s *stdioSession) readInputStream(ctx context.Context) error {
var v string
var res any
v, res, err = processMcpMessage(msgCtx, []byte(line), s.server, protocol, "", "", nil, "")
v, res, err = processMcpMessage(msgCtx, []byte(line), s.server, protocol, "", nil, "")
if err != nil {
// errors during the processing of message will generate a valid MCP Error response.
// server can continue to run.
@@ -394,10 +393,10 @@ func sseHandler(s *Server, w http.ResponseWriter, r *http.Request) {
r = r.WithContext(ctx)
sessionId := uuid.New().String()
toolsetName := chi.URLParam(r, "toolsetName")
s.logger.DebugContext(ctx, fmt.Sprintf("toolset name: %s", toolsetName))
groupName := chi.URLParam(r, "toolsetName")
s.logger.DebugContext(ctx, fmt.Sprintf("toolset name: %s", groupName))
span.SetAttributes(attribute.String("mcp.session.id", sessionId))
span.SetAttributes(attribute.String("toolset.name", toolsetName))
span.SetAttributes(attribute.String("toolset.name", groupName))
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
@@ -410,7 +409,7 @@ func sseHandler(s *Server, w http.ResponseWriter, r *http.Request) {
attribute.String("network.protocol.name", "http"),
attribute.String("network.protocol.version", networkProtocolVersion),
attribute.String("mcp.protocol.version", "2024-11-05"),
attribute.String("toolset.name", toolsetName),
attribute.String("toolset.name", groupName),
}
// Increment active sessions counter
@@ -460,8 +459,8 @@ func sseHandler(s *Server, w http.ResponseWriter, r *http.Request) {
// send initial endpoint event
toolsetURL := ""
if toolsetName != "" {
toolsetURL = fmt.Sprintf("/%s", toolsetName)
if groupName != "" {
toolsetURL = fmt.Sprintf("/%s", groupName)
}
// attach url query params to message endpoint
q := r.URL.Query()
@@ -579,10 +578,9 @@ func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
protocolVersion = headerProtocolVersion
}
toolsetName := chi.URLParam(r, "toolsetName")
promptsetName := chi.URLParam(r, "promptsetName")
s.logger.DebugContext(ctx, fmt.Sprintf("toolset name: %s", toolsetName))
span.SetAttributes(attribute.String("toolset.name", toolsetName))
groupName := chi.URLParam(r, "toolsetName")
s.logger.DebugContext(ctx, fmt.Sprintf("toolset name: %s", groupName))
span.SetAttributes(attribute.String("toolset.name", groupName))
defer func() {
if err != nil {
@@ -593,7 +591,7 @@ func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
networkProtocolVersion := fmt.Sprintf("%d.%d", r.ProtoMajor, r.ProtoMinor)
v, res, err := processMcpMessage(ctx, body, s, protocolVersion, toolsetName, promptsetName, r.Header, networkProtocolVersion)
v, res, err := processMcpMessage(ctx, body, s, protocolVersion, groupName, r.Header, networkProtocolVersion)
if err != nil {
s.logger.DebugContext(ctx, fmt.Errorf("error processing message: %w", err).Error())
}
@@ -666,7 +664,7 @@ func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
}
// processMcpMessage process the messages received from clients
func processMcpMessage(ctx context.Context, body []byte, s *Server, protocolVersion string, toolsetName string, promptsetName string, header http.Header, networkProtocolVersion string) (string, any, error) {
func processMcpMessage(ctx context.Context, body []byte, s *Server, protocolVersion string, groupName string, header http.Header, networkProtocolVersion string) (string, any, error) {
operationStart := time.Now()
logger, err := util.LoggerFromContext(ctx)
@@ -734,7 +732,7 @@ func processMcpMessage(ctx context.Context, body []byte, s *Server, protocolVers
attribute.String("mcp.method.name", baseMessage.Method),
attribute.String("network.transport", networkTransport),
attribute.String("network.protocol.name", networkProtocolName),
attribute.String("toolset.name", toolsetName),
attribute.String("toolset.name", groupName),
}
if protocolVersion != "" {
durationAttrs = append(durationAttrs, attribute.String("mcp.protocol.version", protocolVersion))
@@ -752,6 +750,9 @@ func processMcpMessage(ctx context.Context, body []byte, s *Server, protocolVers
if genAIAttrs.PromptName != "" {
durationAttrs = append(durationAttrs, attribute.String("gen_ai.prompt.name", genAIAttrs.PromptName))
}
if genAIAttrs.GroupName != "" {
durationAttrs = append(durationAttrs, attribute.String("gen_ai.group.name", genAIAttrs.GroupName))
}
if metricErrorType != "" {
durationAttrs = append(durationAttrs, attribute.String("error.type", metricErrorType))
}
@@ -805,7 +806,7 @@ func processMcpMessage(ctx context.Context, body []byte, s *Server, protocolVers
}
// Set toolset name
span.SetAttributes(attribute.String("toolset.name", toolsetName))
span.SetAttributes(attribute.String("toolset.name", groupName))
// Check if message is a notification
if baseMessage.Id == nil {
@@ -842,7 +843,7 @@ func processMcpMessage(ctx context.Context, body []byte, s *Server, protocolVers
version = mcputil.GetLatestSupportedVersion(s.enableDraftSpecs)
}
result, err := mcp.ProcessMethod(ctx, version, baseMessage.Id, baseMessage.Method, tools.Toolset{}, prompts.Promptset{}, nil, body, nil)
result, err := mcp.ProcessMethod(ctx, version, baseMessage.Id, baseMessage.Method, group.Group{}, nil, body, nil)
if err != nil {
span.SetStatus(codes.Error, err.Error())
if rpcErr, ok := result.(jsonrpc.JSONRPCError); ok {
@@ -854,7 +855,9 @@ func processMcpMessage(ctx context.Context, body []byte, s *Server, protocolVers
span.SetAttributes(attribute.String("mcp.protocol.version", version))
return version, result, err
default:
toolset, ok := s.PrimitiveMgr.GetToolset(toolsetName)
// The URL segment names a group; its derived toolset and promptset views
// share that name, so prompts scope to the connected group.
g, ok := s.PrimitiveMgr.GetGroup(groupName)
if !ok {
err := fmt.Errorf("toolset does not exist")
rpcErr := jsonrpc.NewError(baseMessage.Id, jsonrpc.INVALID_REQUEST, err.Error(), nil)
@@ -863,16 +866,7 @@ func processMcpMessage(ctx context.Context, body []byte, s *Server, protocolVers
span.SetAttributes(attribute.String("error.type", metricErrorType))
return "", rpcErr, err
}
promptset, ok := s.PrimitiveMgr.GetPromptset(promptsetName)
if !ok {
err := fmt.Errorf("promptset does not exist")
rpcErr := jsonrpc.NewError(baseMessage.Id, jsonrpc.INVALID_REQUEST, err.Error(), nil)
metricErrorType = rpcErr.Error.String()
span.SetStatus(codes.Error, err.Error())
span.SetAttributes(attribute.String("error.type", metricErrorType))
return "", rpcErr, err
}
result, err := mcp.ProcessMethod(ctx, protocolVersion, baseMessage.Id, baseMessage.Method, toolset, promptset, s.PrimitiveMgr, body, header)
result, err := mcp.ProcessMethod(ctx, protocolVersion, baseMessage.Id, baseMessage.Method, g, s.PrimitiveMgr, body, header)
if err != nil {
span.SetStatus(codes.Error, err.Error())
// Set error.type based on JSON-RPC error code
internal/server/mcp/mcp.go
+7 -8
View File
@@ -20,7 +20,7 @@ import (
"fmt"
"net/http"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
mcputil "github.com/googleapis/mcp-toolbox/internal/server/mcp/util"
v20241105 "github.com/googleapis/mcp-toolbox/internal/server/mcp/v20241105"
@@ -29,7 +29,6 @@ import (
v20251125 "github.com/googleapis/mcp-toolbox/internal/server/mcp/v20251125"
vdraft "github.com/googleapis/mcp-toolbox/internal/server/mcp/vdraft"
"github.com/googleapis/mcp-toolbox/internal/server/primitives"
"github.com/googleapis/mcp-toolbox/internal/tools"
"github.com/googleapis/mcp-toolbox/internal/util"
)
@@ -47,7 +46,7 @@ func NotificationHandler(ctx context.Context, body []byte) error {
// ProcessMethod returns a response for the request.
// This is the Operation phase of the lifecycle for MCP client-server connections.
func ProcessMethod(ctx context.Context, mcpVersion string, id jsonrpc.RequestId, method string, toolset tools.Toolset, promptset prompts.Promptset, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
func ProcessMethod(ctx context.Context, mcpVersion string, id jsonrpc.RequestId, method string, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
enableDraft, ok := util.EnableDraftSpecsFromContext(ctx)
if !ok {
err := fmt.Errorf("unable to retrieve enableDraftSpecs from context")
@@ -56,17 +55,17 @@ func ProcessMethod(ctx context.Context, mcpVersion string, id jsonrpc.RequestId,
switch mcpVersion {
case mcputil.VERSION_DRAFT:
if enableDraft {
return vdraft.ProcessMethod(ctx, id, method, toolset, promptset, primitiveMgr, body, header)
return vdraft.ProcessMethod(ctx, id, method, g, primitiveMgr, body, header)
}
return jsonrpc.NewUnsupportedProtocolVersionError(id, mcpVersion, enableDraft)
case mcputil.VERSION_20251125:
return v20251125.ProcessMethod(ctx, id, method, toolset, promptset, primitiveMgr, body, header)
return v20251125.ProcessMethod(ctx, id, method, g, primitiveMgr, body, header)
case mcputil.VERSION_20250618:
return v20250618.ProcessMethod(ctx, id, method, toolset, promptset, primitiveMgr, body, header)
return v20250618.ProcessMethod(ctx, id, method, g, primitiveMgr, body, header)
case mcputil.VERSION_20250326:
return v20250326.ProcessMethod(ctx, id, method, toolset, promptset, primitiveMgr, body, header)
return v20250326.ProcessMethod(ctx, id, method, g, primitiveMgr, body, header)
case "", mcputil.VERSION_20241105:
return v20241105.ProcessMethod(ctx, id, method, toolset, promptset, primitiveMgr, body, header)
return v20241105.ProcessMethod(ctx, id, method, g, primitiveMgr, body, header)
default:
return jsonrpc.NewUnsupportedProtocolVersionError(id, mcpVersion, enableDraft)
}
internal/server/mcp/v20241105/manifests.go
+46 -6
View File
@@ -16,7 +16,9 @@ package v20241105
import (
"fmt"
"sort"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/sources"
"github.com/googleapis/mcp-toolbox/internal/tools"
@@ -97,9 +99,9 @@ func generateParamManifest(ps parameters.Parameters, urlParams map[string]string
}
// GenerateListToolsResult generates tools/list method result according to mcp schema
func GenerateListToolsResult(srcs map[string]sources.Source, t tools.Toolset, toolsMap map[string]tools.Tool, urlParams map[string]string) (ListToolsResult, error) {
mcpManifest := make([]Tool, 0, len(t.ToolNames))
for _, toolName := range t.ToolNames {
func GenerateListToolsResult(srcs map[string]sources.Source, g group.Group, toolsMap map[string]tools.Tool, urlParams map[string]string) (ListToolsResult, error) {
mcpManifest := make([]Tool, 0, len(g.ToolNames))
for _, toolName := range g.ToolNames {
tool, ok := toolsMap[toolName]
if !ok {
return ListToolsResult{}, fmt.Errorf("tool does not exist: %s", toolName)
@@ -133,9 +135,9 @@ func generatePromptManifest(name, desc string, args prompts.Arguments) Prompt {
}
// GenerateListPromptsResult generates the list/prompts result
func GenerateListPromptsResult(p prompts.Promptset, promptsMap map[string]prompts.Prompt) (ListPromptsResult, error) {
mcpManifest := make([]Prompt, 0, len(p.PromptNames))
for _, promptName := range p.PromptNames {
func GenerateListPromptsResult(g group.Group, promptsMap map[string]prompts.Prompt) (ListPromptsResult, error) {
mcpManifest := make([]Prompt, 0, len(g.PromptNames))
for _, promptName := range g.PromptNames {
prompt, ok := promptsMap[promptName]
if !ok {
return ListPromptsResult{}, fmt.Errorf("prompt does not exist: %s", promptName)
@@ -145,3 +147,41 @@ func GenerateListPromptsResult(p prompts.Promptset, promptsMap map[string]prompt
}
return ListPromptsResult{Prompts: mcpManifest}, nil
}
// GenerateListGroupsResult generates the groups/list result. It omits the
// default nameless group and returns the remaining groups sorted by name.
func GenerateListGroupsResult(groupsMap map[string]group.Group) ListGroupsResult {
names := make([]string, 0, len(groupsMap))
for name := range groupsMap {
if name == "" {
continue
}
names = append(names, name)
}
sort.Strings(names)
groupsList := make([]Group, 0, len(names))
for _, name := range names {
g := groupsMap[name]
groupsList = append(groupsList, Group{Name: g.Name, Description: g.Description})
}
return ListGroupsResult{Groups: groupsList}
}
// GenerateGetGroupResult generates the groups/get result for a single group's
// tools and prompts.
func GenerateGetGroupResult(srcs map[string]sources.Source, g group.Group, toolsMap map[string]tools.Tool, promptsMap map[string]prompts.Prompt, urlParams map[string]string) (GetGroupResult, error) {
listToolsResult, err := GenerateListToolsResult(srcs, g, toolsMap, urlParams)
if err != nil {
return GetGroupResult{}, fmt.Errorf("error generating tools manifest: %w", err)
}
listPromptsResult, err := GenerateListPromptsResult(g, promptsMap)
if err != nil {
return GetGroupResult{}, fmt.Errorf("error generating prompts manifest: %w", err)
}
return GetGroupResult{
Name: g.Name,
Tools: listToolsResult.Tools,
Prompts: listPromptsResult.Prompts,
}, nil
}
internal/server/mcp/v20241105/manifests_test.go
+7 -14
View File
@@ -21,6 +21,7 @@ import (
"testing"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/testutils"
"github.com/googleapis/mcp-toolbox/internal/tools"
@@ -240,16 +241,12 @@ func TestGenerateListToolsResult(t *testing.T) {
toolsMap := make(map[string]tools.Tool)
toolsMap[tool1.Name] = tool1
toolsMap[tool2.Name] = tool2
tc := tools.ToolsetConfig{
g := group.NewGroup(group.GroupConfig{
Name: "test-toolset",
ToolNames: []string{"no_params", "some_params"},
}
toolset, err := tc.Initialize("test-version", toolsMap)
if err != nil {
t.Fatalf("unable to initialize toolset %q: %s", "test-toolset", err)
}
})
got, err := GenerateListToolsResult(nil, toolset, toolsMap, nil)
got, err := GenerateListToolsResult(nil, g, toolsMap, nil)
if err != nil {
t.Fatalf("unable to generate list tools result: %s", err)
}
@@ -353,16 +350,12 @@ func TestGenerateListPromptsResult(t *testing.T) {
promptsMap := make(map[string]prompts.Prompt)
promptsMap[prompt1.Name] = prompt1
promptsMap[prompt2.Name] = prompt2
pc := prompts.PromptsetConfig{
g := group.NewGroup(group.GroupConfig{
Name: "test-promptset",
PromptNames: []string{"prompt1", "prompt2"},
}
promptset, err := pc.Initialize("test-version", promptsMap)
if err != nil {
t.Fatalf("unable to initialize promptset %q: %s", "test-promptset", err)
}
})
got, err := GenerateListPromptsResult(promptset, promptsMap)
got, err := GenerateListPromptsResult(g, promptsMap)
if err != nil {
t.Fatalf("unable to generate list prompt result: %s", err)
}
internal/server/mcp/v20241105/method.go
+96 -22
View File
@@ -21,10 +21,10 @@ import (
"errors"
"fmt"
"net/http"
"time"
"github.com/googleapis/mcp-toolbox/internal/auth"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
mcputil "github.com/googleapis/mcp-toolbox/internal/server/mcp/util"
@@ -38,20 +38,20 @@ import (
)
// ProcessMethod returns a response for the request.
func ProcessMethod(ctx context.Context, id jsonrpc.RequestId, method string, toolset tools.Toolset, promptset prompts.Promptset, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
func ProcessMethod(ctx context.Context, id jsonrpc.RequestId, method string, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
switch method {
case INITIALIZE:
return initializeHandler(ctx, id, body)
case PING:
return pingHandler(id)
case TOOLS_LIST:
return toolsListHandler(ctx, id, primitiveMgr, toolset, body)
return toolsListHandler(ctx, id, primitiveMgr, g, body)
case TOOLS_CALL:
return toolsCallHandler(ctx, id, toolset, primitiveMgr, body, header)
return toolsCallHandler(ctx, id, g, primitiveMgr, body, header)
case PROMPTS_LIST:
return promptsListHandler(ctx, id, primitiveMgr, promptset, body)
return promptsListHandler(ctx, id, primitiveMgr, g, body)
case PROMPTS_GET:
return promptsGetHandler(ctx, id, promptset, primitiveMgr, body)
return promptsGetHandler(ctx, id, g, primitiveMgr, body)
default:
err := fmt.Errorf("invalid method %s", method)
return jsonrpc.NewError(id, jsonrpc.METHOD_NOT_FOUND, err.Error(), nil), err
@@ -110,7 +110,7 @@ func pingHandler(id jsonrpc.RequestId) (any, error) {
}, nil
}
func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, toolset tools.Toolset, body []byte) (any, error) {
func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, g group.Group, body []byte) (any, error) {
var req ListToolsRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp tools list request: %w", err)
@@ -119,7 +119,7 @@ func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *p
urlParams, _ := util.UrlParamsFromContext(ctx)
toolsMap := primitiveMgr.GetToolsMap()
listToolsResult, err := GenerateListToolsResult(primitiveMgr.GetSourcesMap(), toolset, toolsMap, urlParams)
listToolsResult, err := GenerateListToolsResult(primitiveMgr.GetSourcesMap(), g, toolsMap, urlParams)
if err != nil {
err = fmt.Errorf("error generating manifest: %w", err)
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
@@ -133,7 +133,7 @@ func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *p
}
// toolsCallHandler generate a response for tools call.
func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, toolset tools.Toolset, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
if header != nil {
if clientIP := util.ExtractClientIP(header); clientIP != "" {
ctx = util.WithClientIP(ctx, clientIP)
@@ -166,8 +166,8 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, toolset tools.T
attribute.String("gen_ai.operation.name", "execute_tool"),
)
// Verify tool belongs to the current toolset before resolving globally.
if !toolset.ContainsTool(toolName) {
// Verify tool belongs to the current group before resolving globally.
if !g.ContainsTool(toolName) {
err = fmt.Errorf("invalid tool name: tool with name %q does not exist", toolName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
@@ -400,7 +400,7 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, toolset tools.T
}
// promptsListHandler handles the "prompts/list" method.
func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, promptset prompts.Promptset, body []byte) (any, error) {
func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, g group.Group, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
@@ -415,7 +415,7 @@ func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr
}
promptsMap := primitiveMgr.GetPromptsMap()
listPromptsResult, err := GenerateListPromptsResult(promptset, promptsMap)
listPromptsResult, err := GenerateListPromptsResult(g, promptsMap)
if err != nil {
err = fmt.Errorf("error generating manifest: %w", err)
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
@@ -429,7 +429,7 @@ func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr
}
// promptsGetHandler handles the "prompts/get" method.
func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prompts.Promptset, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
@@ -451,8 +451,14 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prom
span.SetName(fmt.Sprintf("%s %s", PROMPTS_GET, promptName))
span.SetAttributes(attribute.String("gen_ai.prompt.name", promptName))
// Verify prompt belongs to the current promptset before resolving globally.
if !promptset.ContainsPrompt(promptName) {
// Populate gen_ai attributes for operation duration metric
if genAIAttrs := util.GenAIMetricAttrsFromContext(ctx); genAIAttrs != nil {
genAIAttrs.OperationName = "get_prompt"
genAIAttrs.PromptName = promptName
}
// Verify prompt belongs to the current group before resolving globally.
if !g.ContainsPrompt(promptName) {
err := fmt.Errorf("prompt with name %q does not exist", promptName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
@@ -463,12 +469,6 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prom
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
// Populate gen_ai attributes for operation duration metric
if genAIAttrs := util.GenAIMetricAttrsFromContext(ctx); genAIAttrs != nil {
genAIAttrs.OperationName = "get_prompt"
genAIAttrs.PromptName = promptName
}
// Parse the arguments provided in the request.
argValues, err := prompt.ParseArgs(req.Params.Arguments, nil)
if err != nil {
@@ -515,3 +515,77 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prom
Result: result,
}, nil
}
// groupsListHandler handles the "groups/list" method. It returns every named
// group's name and description. The default nameless group is omitted.
func groupsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
logger.DebugContext(ctx, "handling groups/list request")
var req ListGroupsRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp groups list request: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
result := GenerateListGroupsResult(primitiveMgr.GetGroupsMap())
logger.DebugContext(ctx, fmt.Sprintf("returning %d groups", len(result.Groups)))
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: result,
}, nil
}
// groupsGetHandler handles the "groups/get" method. It returns the named group's
// tools and prompts.
func groupsGetHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
logger.DebugContext(ctx, "handling groups/get request")
var req GetGroupRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp groups/get request: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
groupName := req.Params.Name
logger.DebugContext(ctx, fmt.Sprintf("group name: %s", groupName))
// Update span name and set gen_ai attributes
span := trace.SpanFromContext(ctx)
span.SetName(fmt.Sprintf("%s %s", GROUPS_GET, groupName))
span.SetAttributes(attribute.String("gen_ai.group.name", groupName))
// Populate gen_ai attributes for operation duration metric
if genAIAttrs := util.GenAIMetricAttrsFromContext(ctx); genAIAttrs != nil {
genAIAttrs.OperationName = "get_group"
genAIAttrs.GroupName = groupName
}
g, ok := primitiveMgr.GetGroup(groupName)
if !ok {
err := fmt.Errorf("invalid group name: group with name %q does not exist", groupName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
urlParams, _ := util.UrlParamsFromContext(ctx)
result, err := GenerateGetGroupResult(primitiveMgr.GetSourcesMap(), g, primitiveMgr.GetToolsMap(), primitiveMgr.GetPromptsMap(), urlParams)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: result,
}, nil
}
internal/server/mcp/v20241105/method_test.go
+202 -17
View File
@@ -21,11 +21,11 @@ import (
"strings"
"testing"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/log"
"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
"github.com/googleapis/mcp-toolbox/internal/server/primitives"
"github.com/googleapis/mcp-toolbox/internal/testutils"
"github.com/googleapis/mcp-toolbox/internal/tools"
"github.com/googleapis/mcp-toolbox/internal/util"
)
@@ -35,6 +35,16 @@ var (
fakeVersionString = "0.0.0"
)
// mustGroup fetches the default group from the resource manager.
func mustGroup(t *testing.T, rm *primitives.PrimitiveManager) group.Group {
t.Helper()
g, ok := rm.GetGroup("")
if !ok {
t.Fatal("default group not found")
}
return g
}
func TestInitializeHandler(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
@@ -167,21 +177,21 @@ func TestPingHandler(t *testing.T) {
func TestToolsListHandler(t *testing.T) {
// Initialize tools using provided testutils mock instances
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
body ListToolsRequest
rawBody []byte
toolset tools.Toolset
g group.Group
wantErr bool
errContains string
}{
{
name: "invalid json body",
rawBody: []byte(`{invalid json}`),
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: true,
errContains: "invalid mcp tools list request",
},
@@ -194,7 +204,7 @@ func TestToolsListHandler(t *testing.T) {
},
},
},
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: false,
},
{
@@ -206,7 +216,7 @@ func TestToolsListHandler(t *testing.T) {
},
},
},
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: false,
},
}
@@ -221,7 +231,7 @@ func TestToolsListHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := toolsListHandler(context.Background(), dummyID, primitiveMgr, tt.toolset, body)
got, err := toolsListHandler(context.Background(), dummyID, primitiveMgr, tt.g, body)
if tt.wantErr {
if err == nil {
@@ -256,8 +266,8 @@ func TestToolsCallHandler(t *testing.T) {
testutils.MockTool4,
testutils.MockTool5,
}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
@@ -353,7 +363,7 @@ func TestToolsCallHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := toolsCallHandler(tt.context, dummyID, toolsets[""], primitiveMgr, body, nil)
got, err := toolsCallHandler(tt.context, dummyID, mustGroup(t, primitiveMgr), primitiveMgr, body, nil)
if tt.wantErr {
if err == nil {
@@ -384,8 +394,8 @@ func TestPromptsListHandler(t *testing.T) {
ctx = util.WithLogger(ctx, testLogger)
// Initialize prompts
mockPrompts := []testutils.MockPrompt{testutils.MockPrompt1, testutils.MockPrompt2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
body ListPromptsRequest
@@ -422,7 +432,7 @@ func TestPromptsListHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := promptsListHandler(ctx, dummyID, primitiveMgr, promptsets[""], body)
got, err := promptsListHandler(ctx, dummyID, primitiveMgr, mustGroup(t, primitiveMgr), body)
if tt.wantErr {
if err == nil {
@@ -453,8 +463,8 @@ func TestPromptsGetHandler(t *testing.T) {
ctx = util.WithLogger(ctx, testLogger)
// Initialize prompts
mockPrompts := []testutils.MockPrompt{testutils.MockPrompt1, testutils.MockPrompt2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
body GetPromptRequest
@@ -529,7 +539,7 @@ func TestPromptsGetHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := promptsGetHandler(ctx, dummyID, promptsets[""], primitiveMgr, body)
got, err := promptsGetHandler(ctx, dummyID, mustGroup(t, primitiveMgr), primitiveMgr, body)
if tt.wantErr {
if err == nil {
@@ -549,3 +559,178 @@ func TestPromptsGetHandler(t *testing.T) {
})
}
}
func TestGroupsListHandler(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info")
if err != nil {
t.Fatalf("unable to initialize logger: %s", err)
}
ctx = util.WithLogger(ctx, testLogger)
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
rawBody []byte
body ListGroupsRequest
wantErr bool
errContains string
wantNames []string
}{
{
name: "invalid json body",
rawBody: []byte(`{invalid json}`),
wantErr: true,
errContains: "invalid mcp groups list request",
},
{
name: "success excludes default group and sorts",
body: ListGroupsRequest{
PaginatedRequest: PaginatedRequest{
Request: jsonrpc.Request{Method: GROUPS_LIST},
},
},
wantErr: false,
wantNames: []string{"tool1_only", "tool2_only"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
body := tt.rawBody
var err error
if body == nil {
body, err = json.Marshal(tt.body)
if err != nil {
t.Fatalf("unexpected error during marshaling: %v", err)
}
}
got, err := groupsListHandler(ctx, dummyID, primitiveMgr, body)
if tt.wantErr {
if err == nil {
t.Fatalf("expected error, got nil")
}
if tt.errContains != "" && !strings.Contains(err.Error(), tt.errContains) {
t.Errorf("error = %v, want string containing %q", err, tt.errContains)
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
res, ok := got.(jsonrpc.JSONRPCResponse)
if !ok {
t.Fatalf("expected jsonrpc.JSONRPCResponse, got %T", got)
}
result, ok := res.Result.(ListGroupsResult)
if !ok {
t.Fatalf("expected ListGroupsResult, got %T", res.Result)
}
gotNames := make([]string, 0, len(result.Groups))
for _, g := range result.Groups {
gotNames = append(gotNames, g.Name)
}
if len(gotNames) != len(tt.wantNames) {
t.Fatalf("got groups %v, want %v", gotNames, tt.wantNames)
}
for i, n := range tt.wantNames {
if gotNames[i] != n {
t.Errorf("group[%d] = %q, want %q", i, gotNames[i], n)
}
}
})
}
}
func TestGroupsGetHandler(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info")
if err != nil {
t.Fatalf("unable to initialize logger: %s", err)
}
ctx = util.WithLogger(ctx, testLogger)
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
rawBody []byte
body GetGroupRequest
wantErr bool
errContains string
wantName string
}{
{
name: "invalid json body",
rawBody: []byte(`{invalid json}`),
wantErr: true,
errContains: "invalid mcp groups/get request",
},
{
name: "group does not exist",
body: GetGroupRequest{
Request: jsonrpc.Request{Method: GROUPS_GET},
Params: struct {
Name string `json:"name"`
}{Name: "missing_group"},
},
wantErr: true,
errContains: `group with name "missing_group" does not exist`,
},
{
name: "success",
body: GetGroupRequest{
Request: jsonrpc.Request{Method: GROUPS_GET},
Params: struct {
Name string `json:"name"`
}{Name: "tool1_only"},
},
wantErr: false,
wantName: "tool1_only",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
body := tt.rawBody
var err error
if body == nil {
body, err = json.Marshal(tt.body)
if err != nil {
t.Fatalf("unexpected error during marshaling: %v", err)
}
}
got, err := groupsGetHandler(ctx, dummyID, primitiveMgr, body)
if tt.wantErr {
if err == nil {
t.Fatalf("expected error, got nil")
}
if tt.errContains != "" && !strings.Contains(err.Error(), tt.errContains) {
t.Errorf("error = %v, want string containing %q", err, tt.errContains)
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
res, ok := got.(jsonrpc.JSONRPCResponse)
if !ok {
t.Fatalf("expected jsonrpc.JSONRPCResponse, got %T", got)
}
result, ok := res.Result.(GetGroupResult)
if !ok {
t.Fatalf("expected GetGroupResult, got %T", res.Result)
}
if result.Name != tt.wantName {
t.Errorf("result.Name = %q, want %q", result.Name, tt.wantName)
}
})
}
}
internal/server/mcp/v20241105/types.go
+40
View File
@@ -34,6 +34,8 @@ const (
TOOLS_CALL = "tools/call"
PROMPTS_LIST = "prompts/list"
PROMPTS_GET = "prompts/get"
GROUPS_LIST = "groups/list"
GROUPS_GET = "groups/get"
)
/* Initialization */
@@ -333,3 +335,41 @@ type PromptMessage struct {
Role string `json:"role"`
Content TextContent `json:"content"`
}
/* Groups */
// ListGroupsRequest is sent from the client to request the list of groups the
// server has.
type ListGroupsRequest struct {
PaginatedRequest
}
// Group is a single entry in a groups/list response.
type Group struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
}
// ListGroupsResult is the server's response to a groups/list request.
type ListGroupsResult struct {
jsonrpc.Result
Groups []Group `json:"groups"`
}
// GetGroupRequest is sent from the client to request a single group's contents.
type GetGroupRequest struct {
jsonrpc.Request
Params struct {
Name string `json:"name"`
} `json:"params"`
}
// GetGroupResult is the server's response to a groups/get request: the group's
// tools and prompts. The description is intentionally omitted; it is exposed only
// through groups/list.
type GetGroupResult struct {
jsonrpc.Result
Name string `json:"name"`
Tools []Tool `json:"tools"`
Prompts []Prompt `json:"prompts"`
}
internal/server/mcp/v20250326/manifests.go
+46 -6
View File
@@ -16,7 +16,9 @@ package v20250326
import (
"fmt"
"sort"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/sources"
"github.com/googleapis/mcp-toolbox/internal/tools"
@@ -96,9 +98,9 @@ func generateParamManifest(ps parameters.Parameters, urlParams map[string]string
}
// GenerateListToolsResult generates tools/list method result according to mcp schema
func GenerateListToolsResult(srcs map[string]sources.Source, t tools.Toolset, toolsMap map[string]tools.Tool, urlParams map[string]string) (ListToolsResult, error) {
mcpManifest := make([]Tool, 0, len(t.ToolNames))
for _, toolName := range t.ToolNames {
func GenerateListToolsResult(srcs map[string]sources.Source, g group.Group, toolsMap map[string]tools.Tool, urlParams map[string]string) (ListToolsResult, error) {
mcpManifest := make([]Tool, 0, len(g.ToolNames))
for _, toolName := range g.ToolNames {
tool, ok := toolsMap[toolName]
if !ok {
return ListToolsResult{}, fmt.Errorf("tool does not exist: %s", toolName)
@@ -132,9 +134,9 @@ func generatePromptManifest(name, desc string, args prompts.Arguments) Prompt {
}
// GenerateListPromptsResult generates the list/prompts result
func GenerateListPromptsResult(p prompts.Promptset, promptsMap map[string]prompts.Prompt) (ListPromptsResult, error) {
mcpManifest := make([]Prompt, 0, len(p.PromptNames))
for _, promptName := range p.PromptNames {
func GenerateListPromptsResult(g group.Group, promptsMap map[string]prompts.Prompt) (ListPromptsResult, error) {
mcpManifest := make([]Prompt, 0, len(g.PromptNames))
for _, promptName := range g.PromptNames {
prompt, ok := promptsMap[promptName]
if !ok {
return ListPromptsResult{}, fmt.Errorf("prompt does not exist: %s", promptName)
@@ -144,3 +146,41 @@ func GenerateListPromptsResult(p prompts.Promptset, promptsMap map[string]prompt
}
return ListPromptsResult{Prompts: mcpManifest}, nil
}
// GenerateListGroupsResult generates the groups/list result. It omits the
// default nameless group and returns the remaining groups sorted by name.
func GenerateListGroupsResult(groupsMap map[string]group.Group) ListGroupsResult {
names := make([]string, 0, len(groupsMap))
for name := range groupsMap {
if name == "" {
continue
}
names = append(names, name)
}
sort.Strings(names)
groupsList := make([]Group, 0, len(names))
for _, name := range names {
g := groupsMap[name]
groupsList = append(groupsList, Group{Name: g.Name, Description: g.Description})
}
return ListGroupsResult{Groups: groupsList}
}
// GenerateGetGroupResult generates the groups/get result for a single group's
// tools and prompts.
func GenerateGetGroupResult(srcs map[string]sources.Source, g group.Group, toolsMap map[string]tools.Tool, promptsMap map[string]prompts.Prompt, urlParams map[string]string) (GetGroupResult, error) {
listToolsResult, err := GenerateListToolsResult(srcs, g, toolsMap, urlParams)
if err != nil {
return GetGroupResult{}, fmt.Errorf("error generating tools manifest: %w", err)
}
listPromptsResult, err := GenerateListPromptsResult(g, promptsMap)
if err != nil {
return GetGroupResult{}, fmt.Errorf("error generating prompts manifest: %w", err)
}
return GetGroupResult{
Name: g.Name,
Tools: listToolsResult.Tools,
Prompts: listPromptsResult.Prompts,
}, nil
}
internal/server/mcp/v20250326/manifests_test.go
+7 -14
View File
@@ -21,6 +21,7 @@ import (
"testing"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/testutils"
"github.com/googleapis/mcp-toolbox/internal/tools"
@@ -240,16 +241,12 @@ func TestGenerateListToolsResult(t *testing.T) {
toolsMap := make(map[string]tools.Tool)
toolsMap[tool1.Name] = tool1
toolsMap[tool2.Name] = tool2
tc := tools.ToolsetConfig{
g := group.NewGroup(group.GroupConfig{
Name: "test-toolset",
ToolNames: []string{"no_params", "some_params"},
}
toolset, err := tc.Initialize("test-version", toolsMap)
if err != nil {
t.Fatalf("unable to initialize toolset %q: %s", "test-toolset", err)
}
})
got, err := GenerateListToolsResult(nil, toolset, toolsMap, nil)
got, err := GenerateListToolsResult(nil, g, toolsMap, nil)
if err != nil {
t.Fatalf("unable to generate list tools result: %s", err)
}
@@ -353,16 +350,12 @@ func TestGenerateListPromptsResult(t *testing.T) {
promptsMap := make(map[string]prompts.Prompt)
promptsMap[prompt1.Name] = prompt1
promptsMap[prompt2.Name] = prompt2
pc := prompts.PromptsetConfig{
g := group.NewGroup(group.GroupConfig{
Name: "test-promptset",
PromptNames: []string{"prompt1", "prompt2"},
}
promptset, err := pc.Initialize("test-version", promptsMap)
if err != nil {
t.Fatalf("unable to initialize promptset %q: %s", "test-promptset", err)
}
})
got, err := GenerateListPromptsResult(promptset, promptsMap)
got, err := GenerateListPromptsResult(g, promptsMap)
if err != nil {
t.Fatalf("unable to generate list prompt result: %s", err)
}
internal/server/mcp/v20250326/method.go
+96 -22
View File
@@ -21,10 +21,10 @@ import (
"errors"
"fmt"
"net/http"
"time"
"github.com/googleapis/mcp-toolbox/internal/auth"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
mcputil "github.com/googleapis/mcp-toolbox/internal/server/mcp/util"
@@ -38,20 +38,20 @@ import (
)
// ProcessMethod returns a response for the request.
func ProcessMethod(ctx context.Context, id jsonrpc.RequestId, method string, toolset tools.Toolset, promptset prompts.Promptset, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
func ProcessMethod(ctx context.Context, id jsonrpc.RequestId, method string, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
switch method {
case INITIALIZE:
return initializeHandler(ctx, id, body)
case PING:
return pingHandler(id)
case TOOLS_LIST:
return toolsListHandler(ctx, id, primitiveMgr, toolset, body)
return toolsListHandler(ctx, id, primitiveMgr, g, body)
case TOOLS_CALL:
return toolsCallHandler(ctx, id, toolset, primitiveMgr, body, header)
return toolsCallHandler(ctx, id, g, primitiveMgr, body, header)
case PROMPTS_LIST:
return promptsListHandler(ctx, id, primitiveMgr, promptset, body)
return promptsListHandler(ctx, id, primitiveMgr, g, body)
case PROMPTS_GET:
return promptsGetHandler(ctx, id, promptset, primitiveMgr, body)
return promptsGetHandler(ctx, id, g, primitiveMgr, body)
default:
err := fmt.Errorf("invalid method %s", method)
return jsonrpc.NewError(id, jsonrpc.METHOD_NOT_FOUND, err.Error(), nil), err
@@ -110,7 +110,7 @@ func pingHandler(id jsonrpc.RequestId) (any, error) {
}, nil
}
func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, toolset tools.Toolset, body []byte) (any, error) {
func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, g group.Group, body []byte) (any, error) {
var req ListToolsRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp tools list request: %w", err)
@@ -119,7 +119,7 @@ func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *p
urlParams, _ := util.UrlParamsFromContext(ctx)
toolsMap := primitiveMgr.GetToolsMap()
listToolsResult, err := GenerateListToolsResult(primitiveMgr.GetSourcesMap(), toolset, toolsMap, urlParams)
listToolsResult, err := GenerateListToolsResult(primitiveMgr.GetSourcesMap(), g, toolsMap, urlParams)
if err != nil {
err = fmt.Errorf("error generating manifest: %w", err)
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
@@ -133,7 +133,7 @@ func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *p
}
// toolsCallHandler generate a response for tools call.
func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, toolset tools.Toolset, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
if header != nil {
if clientIP := util.ExtractClientIP(header); clientIP != "" {
ctx = util.WithClientIP(ctx, clientIP)
@@ -166,8 +166,8 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, toolset tools.T
attribute.String("gen_ai.operation.name", "execute_tool"),
)
// Verify tool belongs to the current toolset before resolving globally.
if !toolset.ContainsTool(toolName) {
// Verify tool belongs to the current group before resolving globally.
if !g.ContainsTool(toolName) {
err = fmt.Errorf("invalid tool name: tool with name %q does not exist", toolName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
@@ -399,7 +399,7 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, toolset tools.T
}
// promptsListHandler handles the "prompts/list" method.
func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, promptset prompts.Promptset, body []byte) (any, error) {
func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, g group.Group, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
@@ -414,7 +414,7 @@ func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr
}
promptsMap := primitiveMgr.GetPromptsMap()
listPromptsResult, err := GenerateListPromptsResult(promptset, promptsMap)
listPromptsResult, err := GenerateListPromptsResult(g, promptsMap)
if err != nil {
err = fmt.Errorf("error generating manifest: %w", err)
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
@@ -428,7 +428,7 @@ func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr
}
// promptsGetHandler handles the "prompts/get" method.
func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prompts.Promptset, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
@@ -450,8 +450,14 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prom
span.SetName(fmt.Sprintf("%s %s", PROMPTS_GET, promptName))
span.SetAttributes(attribute.String("gen_ai.prompt.name", promptName))
// Verify prompt belongs to the current promptset before resolving globally.
if !promptset.ContainsPrompt(promptName) {
// Populate gen_ai attributes for operation duration metric
if genAIAttrs := util.GenAIMetricAttrsFromContext(ctx); genAIAttrs != nil {
genAIAttrs.OperationName = "get_prompt"
genAIAttrs.PromptName = promptName
}
// Verify prompt belongs to the current group before resolving globally.
if !g.ContainsPrompt(promptName) {
err := fmt.Errorf("prompt with name %q does not exist", promptName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
@@ -462,12 +468,6 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prom
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
// Populate gen_ai attributes for operation duration metric
if genAIAttrs := util.GenAIMetricAttrsFromContext(ctx); genAIAttrs != nil {
genAIAttrs.OperationName = "get_prompt"
genAIAttrs.PromptName = promptName
}
// Parse the arguments provided in the request.
argValues, err := prompt.ParseArgs(req.Params.Arguments, nil)
if err != nil {
@@ -514,3 +514,77 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prom
Result: result,
}, nil
}
// groupsListHandler handles the "groups/list" method. It returns every named
// group's name and description. The default nameless group is omitted.
func groupsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
logger.DebugContext(ctx, "handling groups/list request")
var req ListGroupsRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp groups list request: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
result := GenerateListGroupsResult(primitiveMgr.GetGroupsMap())
logger.DebugContext(ctx, fmt.Sprintf("returning %d groups", len(result.Groups)))
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: result,
}, nil
}
// groupsGetHandler handles the "groups/get" method. It returns the named group's
// tools and prompts.
func groupsGetHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
logger.DebugContext(ctx, "handling groups/get request")
var req GetGroupRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp groups/get request: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
groupName := req.Params.Name
logger.DebugContext(ctx, fmt.Sprintf("group name: %s", groupName))
// Update span name and set gen_ai attributes
span := trace.SpanFromContext(ctx)
span.SetName(fmt.Sprintf("%s %s", GROUPS_GET, groupName))
span.SetAttributes(attribute.String("gen_ai.group.name", groupName))
// Populate gen_ai attributes for operation duration metric
if genAIAttrs := util.GenAIMetricAttrsFromContext(ctx); genAIAttrs != nil {
genAIAttrs.OperationName = "get_group"
genAIAttrs.GroupName = groupName
}
g, ok := primitiveMgr.GetGroup(groupName)
if !ok {
err := fmt.Errorf("invalid group name: group with name %q does not exist", groupName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
urlParams, _ := util.UrlParamsFromContext(ctx)
result, err := GenerateGetGroupResult(primitiveMgr.GetSourcesMap(), g, primitiveMgr.GetToolsMap(), primitiveMgr.GetPromptsMap(), urlParams)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: result,
}, nil
}
internal/server/mcp/v20250326/method_test.go
+202 -17
View File
@@ -21,11 +21,11 @@ import (
"strings"
"testing"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/log"
"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
"github.com/googleapis/mcp-toolbox/internal/server/primitives"
"github.com/googleapis/mcp-toolbox/internal/testutils"
"github.com/googleapis/mcp-toolbox/internal/tools"
"github.com/googleapis/mcp-toolbox/internal/util"
)
@@ -35,6 +35,16 @@ var (
fakeVersionString = "0.0.0"
)
// mustGroup fetches the default group from the resource manager.
func mustGroup(t *testing.T, rm *primitives.PrimitiveManager) group.Group {
t.Helper()
g, ok := rm.GetGroup("")
if !ok {
t.Fatal("default group not found")
}
return g
}
func TestInitializeHandler(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
@@ -167,21 +177,21 @@ func TestPingHandler(t *testing.T) {
func TestToolsListHandler(t *testing.T) {
// Initialize tools using provided testutils mock instances
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
body ListToolsRequest
rawBody []byte
toolset tools.Toolset
g group.Group
wantErr bool
errContains string
}{
{
name: "invalid json body",
rawBody: []byte(`{invalid json}`),
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: true,
errContains: "invalid mcp tools list request",
},
@@ -194,7 +204,7 @@ func TestToolsListHandler(t *testing.T) {
},
},
},
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: false,
},
{
@@ -206,7 +216,7 @@ func TestToolsListHandler(t *testing.T) {
},
},
},
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: false,
},
}
@@ -221,7 +231,7 @@ func TestToolsListHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := toolsListHandler(context.Background(), dummyID, primitiveMgr, tt.toolset, body)
got, err := toolsListHandler(context.Background(), dummyID, primitiveMgr, tt.g, body)
if tt.wantErr {
if err == nil {
@@ -256,8 +266,8 @@ func TestToolsCallHandler(t *testing.T) {
testutils.MockTool4,
testutils.MockTool5,
}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
@@ -353,7 +363,7 @@ func TestToolsCallHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := toolsCallHandler(tt.context, dummyID, toolsets[""], primitiveMgr, body, nil)
got, err := toolsCallHandler(tt.context, dummyID, mustGroup(t, primitiveMgr), primitiveMgr, body, nil)
if tt.wantErr {
if err == nil {
@@ -384,8 +394,8 @@ func TestPromptsListHandler(t *testing.T) {
ctx = util.WithLogger(ctx, testLogger)
// Initialize prompts
mockPrompts := []testutils.MockPrompt{testutils.MockPrompt1, testutils.MockPrompt2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
body ListPromptsRequest
@@ -422,7 +432,7 @@ func TestPromptsListHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := promptsListHandler(ctx, dummyID, primitiveMgr, promptsets[""], body)
got, err := promptsListHandler(ctx, dummyID, primitiveMgr, mustGroup(t, primitiveMgr), body)
if tt.wantErr {
if err == nil {
@@ -453,8 +463,8 @@ func TestPromptsGetHandler(t *testing.T) {
ctx = util.WithLogger(ctx, testLogger)
// Initialize prompts
mockPrompts := []testutils.MockPrompt{testutils.MockPrompt1, testutils.MockPrompt2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
body GetPromptRequest
@@ -529,7 +539,7 @@ func TestPromptsGetHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := promptsGetHandler(ctx, dummyID, promptsets[""], primitiveMgr, body)
got, err := promptsGetHandler(ctx, dummyID, mustGroup(t, primitiveMgr), primitiveMgr, body)
if tt.wantErr {
if err == nil {
@@ -549,3 +559,178 @@ func TestPromptsGetHandler(t *testing.T) {
})
}
}
func TestGroupsListHandler(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info")
if err != nil {
t.Fatalf("unable to initialize logger: %s", err)
}
ctx = util.WithLogger(ctx, testLogger)
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
rawBody []byte
body ListGroupsRequest
wantErr bool
errContains string
wantNames []string
}{
{
name: "invalid json body",
rawBody: []byte(`{invalid json}`),
wantErr: true,
errContains: "invalid mcp groups list request",
},
{
name: "success excludes default group and sorts",
body: ListGroupsRequest{
PaginatedRequest: PaginatedRequest{
Request: jsonrpc.Request{Method: GROUPS_LIST},
},
},
wantErr: false,
wantNames: []string{"tool1_only", "tool2_only"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
body := tt.rawBody
var err error
if body == nil {
body, err = json.Marshal(tt.body)
if err != nil {
t.Fatalf("unexpected error during marshaling: %v", err)
}
}
got, err := groupsListHandler(ctx, dummyID, primitiveMgr, body)
if tt.wantErr {
if err == nil {
t.Fatalf("expected error, got nil")
}
if tt.errContains != "" && !strings.Contains(err.Error(), tt.errContains) {
t.Errorf("error = %v, want string containing %q", err, tt.errContains)
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
res, ok := got.(jsonrpc.JSONRPCResponse)
if !ok {
t.Fatalf("expected jsonrpc.JSONRPCResponse, got %T", got)
}
result, ok := res.Result.(ListGroupsResult)
if !ok {
t.Fatalf("expected ListGroupsResult, got %T", res.Result)
}
gotNames := make([]string, 0, len(result.Groups))
for _, g := range result.Groups {
gotNames = append(gotNames, g.Name)
}
if len(gotNames) != len(tt.wantNames) {
t.Fatalf("got groups %v, want %v", gotNames, tt.wantNames)
}
for i, n := range tt.wantNames {
if gotNames[i] != n {
t.Errorf("group[%d] = %q, want %q", i, gotNames[i], n)
}
}
})
}
}
func TestGroupsGetHandler(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info")
if err != nil {
t.Fatalf("unable to initialize logger: %s", err)
}
ctx = util.WithLogger(ctx, testLogger)
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
rawBody []byte
body GetGroupRequest
wantErr bool
errContains string
wantName string
}{
{
name: "invalid json body",
rawBody: []byte(`{invalid json}`),
wantErr: true,
errContains: "invalid mcp groups/get request",
},
{
name: "group does not exist",
body: GetGroupRequest{
Request: jsonrpc.Request{Method: GROUPS_GET},
Params: struct {
Name string `json:"name"`
}{Name: "missing_group"},
},
wantErr: true,
errContains: `group with name "missing_group" does not exist`,
},
{
name: "success",
body: GetGroupRequest{
Request: jsonrpc.Request{Method: GROUPS_GET},
Params: struct {
Name string `json:"name"`
}{Name: "tool1_only"},
},
wantErr: false,
wantName: "tool1_only",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
body := tt.rawBody
var err error
if body == nil {
body, err = json.Marshal(tt.body)
if err != nil {
t.Fatalf("unexpected error during marshaling: %v", err)
}
}
got, err := groupsGetHandler(ctx, dummyID, primitiveMgr, body)
if tt.wantErr {
if err == nil {
t.Fatalf("expected error, got nil")
}
if tt.errContains != "" && !strings.Contains(err.Error(), tt.errContains) {
t.Errorf("error = %v, want string containing %q", err, tt.errContains)
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
res, ok := got.(jsonrpc.JSONRPCResponse)
if !ok {
t.Fatalf("expected jsonrpc.JSONRPCResponse, got %T", got)
}
result, ok := res.Result.(GetGroupResult)
if !ok {
t.Fatalf("expected GetGroupResult, got %T", res.Result)
}
if result.Name != tt.wantName {
t.Errorf("result.Name = %q, want %q", result.Name, tt.wantName)
}
})
}
}
internal/server/mcp/v20250326/types.go
+40
View File
@@ -34,6 +34,8 @@ const (
TOOLS_CALL = "tools/call"
PROMPTS_LIST = "prompts/list"
PROMPTS_GET = "prompts/get"
GROUPS_LIST = "groups/list"
GROUPS_GET = "groups/get"
)
/* Initialization */
@@ -333,3 +335,41 @@ type PromptMessage struct {
Role string `json:"role"`
Content TextContent `json:"content"`
}
/* Groups */
// ListGroupsRequest is sent from the client to request the list of groups the
// server has.
type ListGroupsRequest struct {
PaginatedRequest
}
// Group is a single entry in a groups/list response.
type Group struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
}
// ListGroupsResult is the server's response to a groups/list request.
type ListGroupsResult struct {
jsonrpc.Result
Groups []Group `json:"groups"`
}
// GetGroupRequest is sent from the client to request a single group's contents.
type GetGroupRequest struct {
jsonrpc.Request
Params struct {
Name string `json:"name"`
} `json:"params"`
}
// GetGroupResult is the server's response to a groups/get request: the group's
// tools and prompts. The description is intentionally omitted; it is exposed only
// through groups/list.
type GetGroupResult struct {
jsonrpc.Result
Name string `json:"name"`
Tools []Tool `json:"tools"`
Prompts []Prompt `json:"prompts"`
}
internal/server/mcp/v20250618/manifests.go
+46 -6
View File
@@ -16,7 +16,9 @@ package v20250618
import (
"fmt"
"sort"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/sources"
"github.com/googleapis/mcp-toolbox/internal/tools"
@@ -98,9 +100,9 @@ func generateParamManifest(ps parameters.Parameters, urlParams map[string]string
}
// GenerateListToolsResult generates tools/list method result according to mcp schema
func GenerateListToolsResult(srcs map[string]sources.Source, t tools.Toolset, toolsMap map[string]tools.Tool, urlParams map[string]string) (ListToolsResult, error) {
mcpManifest := make([]Tool, 0, len(t.ToolNames))
for _, toolName := range t.ToolNames {
func GenerateListToolsResult(srcs map[string]sources.Source, g group.Group, toolsMap map[string]tools.Tool, urlParams map[string]string) (ListToolsResult, error) {
mcpManifest := make([]Tool, 0, len(g.ToolNames))
for _, toolName := range g.ToolNames {
tool, ok := toolsMap[toolName]
if !ok {
return ListToolsResult{}, fmt.Errorf("tool does not exist: %s", toolName)
@@ -134,9 +136,9 @@ func generatePromptManifest(name, desc string, args prompts.Arguments) Prompt {
}
// GenerateListPromptsResult generates the list/prompts result
func GenerateListPromptsResult(p prompts.Promptset, promptsMap map[string]prompts.Prompt) (ListPromptsResult, error) {
mcpManifest := make([]Prompt, 0, len(p.PromptNames))
for _, promptName := range p.PromptNames {
func GenerateListPromptsResult(g group.Group, promptsMap map[string]prompts.Prompt) (ListPromptsResult, error) {
mcpManifest := make([]Prompt, 0, len(g.PromptNames))
for _, promptName := range g.PromptNames {
prompt, ok := promptsMap[promptName]
if !ok {
return ListPromptsResult{}, fmt.Errorf("prompt does not exist: %s", promptName)
@@ -146,3 +148,41 @@ func GenerateListPromptsResult(p prompts.Promptset, promptsMap map[string]prompt
}
return ListPromptsResult{Prompts: mcpManifest}, nil
}
// GenerateListGroupsResult generates the groups/list result. It omits the
// default nameless group and returns the remaining groups sorted by name.
func GenerateListGroupsResult(groupsMap map[string]group.Group) ListGroupsResult {
names := make([]string, 0, len(groupsMap))
for name := range groupsMap {
if name == "" {
continue
}
names = append(names, name)
}
sort.Strings(names)
groupsList := make([]Group, 0, len(names))
for _, name := range names {
g := groupsMap[name]
groupsList = append(groupsList, Group{Name: g.Name, Description: g.Description})
}
return ListGroupsResult{Groups: groupsList}
}
// GenerateGetGroupResult generates the groups/get result for a single group's
// tools and prompts.
func GenerateGetGroupResult(srcs map[string]sources.Source, g group.Group, toolsMap map[string]tools.Tool, promptsMap map[string]prompts.Prompt, urlParams map[string]string) (GetGroupResult, error) {
listToolsResult, err := GenerateListToolsResult(srcs, g, toolsMap, urlParams)
if err != nil {
return GetGroupResult{}, fmt.Errorf("error generating tools manifest: %w", err)
}
listPromptsResult, err := GenerateListPromptsResult(g, promptsMap)
if err != nil {
return GetGroupResult{}, fmt.Errorf("error generating prompts manifest: %w", err)
}
return GetGroupResult{
Name: g.Name,
Tools: listToolsResult.Tools,
Prompts: listPromptsResult.Prompts,
}, nil
}
internal/server/mcp/v20250618/manifests_test.go
+7 -14
View File
@@ -21,6 +21,7 @@ import (
"testing"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/testutils"
"github.com/googleapis/mcp-toolbox/internal/tools"
@@ -240,16 +241,12 @@ func TestGenerateListToolsResult(t *testing.T) {
toolsMap := make(map[string]tools.Tool)
toolsMap[tool1.Name] = tool1
toolsMap[tool2.Name] = tool2
tc := tools.ToolsetConfig{
g := group.NewGroup(group.GroupConfig{
Name: "test-toolset",
ToolNames: []string{"no_params", "some_params"},
}
toolset, err := tc.Initialize("test-version", toolsMap)
if err != nil {
t.Fatalf("unable to initialize toolset %q: %s", "test-toolset", err)
}
})
got, err := GenerateListToolsResult(nil, toolset, toolsMap, nil)
got, err := GenerateListToolsResult(nil, g, toolsMap, nil)
if err != nil {
t.Fatalf("unable to generate list tools result: %s", err)
}
@@ -353,16 +350,12 @@ func TestGenerateListPromptsResult(t *testing.T) {
promptsMap := make(map[string]prompts.Prompt)
promptsMap[prompt1.Name] = prompt1
promptsMap[prompt2.Name] = prompt2
pc := prompts.PromptsetConfig{
g := group.NewGroup(group.GroupConfig{
Name: "test-promptset",
PromptNames: []string{"prompt1", "prompt2"},
}
promptset, err := pc.Initialize("test-version", promptsMap)
if err != nil {
t.Fatalf("unable to initialize promptset %q: %s", "test-promptset", err)
}
})
got, err := GenerateListPromptsResult(promptset, promptsMap)
got, err := GenerateListPromptsResult(g, promptsMap)
if err != nil {
t.Fatalf("unable to generate list prompt result: %s", err)
}
internal/server/mcp/v20250618/method.go
+96 -22
View File
@@ -21,10 +21,10 @@ import (
"errors"
"fmt"
"net/http"
"time"
"github.com/googleapis/mcp-toolbox/internal/auth"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
mcputil "github.com/googleapis/mcp-toolbox/internal/server/mcp/util"
@@ -38,20 +38,20 @@ import (
)
// ProcessMethod returns a response for the request.
func ProcessMethod(ctx context.Context, id jsonrpc.RequestId, method string, toolset tools.Toolset, promptset prompts.Promptset, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
func ProcessMethod(ctx context.Context, id jsonrpc.RequestId, method string, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
switch method {
case INITIALIZE:
return initializeHandler(ctx, id, body)
case PING:
return pingHandler(id)
case TOOLS_LIST:
return toolsListHandler(ctx, id, primitiveMgr, toolset, body)
return toolsListHandler(ctx, id, primitiveMgr, g, body)
case TOOLS_CALL:
return toolsCallHandler(ctx, id, toolset, primitiveMgr, body, header)
return toolsCallHandler(ctx, id, g, primitiveMgr, body, header)
case PROMPTS_LIST:
return promptsListHandler(ctx, id, primitiveMgr, promptset, body)
return promptsListHandler(ctx, id, primitiveMgr, g, body)
case PROMPTS_GET:
return promptsGetHandler(ctx, id, promptset, primitiveMgr, body)
return promptsGetHandler(ctx, id, g, primitiveMgr, body)
default:
err := fmt.Errorf("invalid method %s", method)
return jsonrpc.NewError(id, jsonrpc.METHOD_NOT_FOUND, err.Error(), nil), err
@@ -110,7 +110,7 @@ func pingHandler(id jsonrpc.RequestId) (any, error) {
}, nil
}
func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, toolset tools.Toolset, body []byte) (any, error) {
func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, g group.Group, body []byte) (any, error) {
var req ListToolsRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp tools list request: %w", err)
@@ -119,7 +119,7 @@ func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *p
urlParams, _ := util.UrlParamsFromContext(ctx)
toolsMap := primitiveMgr.GetToolsMap()
listToolsResult, err := GenerateListToolsResult(primitiveMgr.GetSourcesMap(), toolset, toolsMap, urlParams)
listToolsResult, err := GenerateListToolsResult(primitiveMgr.GetSourcesMap(), g, toolsMap, urlParams)
if err != nil {
err = fmt.Errorf("error generating manifest: %w", err)
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
@@ -132,7 +132,7 @@ func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *p
}
// toolsCallHandler generate a response for tools call.
func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, toolset tools.Toolset, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
if header != nil {
if clientIP := util.ExtractClientIP(header); clientIP != "" {
ctx = util.WithClientIP(ctx, clientIP)
@@ -165,8 +165,8 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, toolset tools.T
attribute.String("gen_ai.operation.name", "execute_tool"),
)
// Verify tool belongs to the current toolset before resolving globally.
if !toolset.ContainsTool(toolName) {
// Verify tool belongs to the current group before resolving globally.
if !g.ContainsTool(toolName) {
err = fmt.Errorf("invalid tool name: tool with name %q does not exist", toolName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
@@ -399,7 +399,7 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, toolset tools.T
}
// promptsListHandler handles the "prompts/list" method.
func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, promptset prompts.Promptset, body []byte) (any, error) {
func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, g group.Group, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
@@ -414,7 +414,7 @@ func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr
}
promptsMap := primitiveMgr.GetPromptsMap()
listPromptsResult, err := GenerateListPromptsResult(promptset, promptsMap)
listPromptsResult, err := GenerateListPromptsResult(g, promptsMap)
if err != nil {
err = fmt.Errorf("error generating manifest: %w", err)
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
@@ -428,7 +428,7 @@ func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr
}
// promptsGetHandler handles the "prompts/get" method.
func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prompts.Promptset, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
@@ -450,8 +450,14 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prom
span.SetName(fmt.Sprintf("%s %s", PROMPTS_GET, promptName))
span.SetAttributes(attribute.String("gen_ai.prompt.name", promptName))
// Verify prompt belongs to the current promptset before resolving globally.
if !promptset.ContainsPrompt(promptName) {
// Populate gen_ai attributes for operation duration metric
if genAIAttrs := util.GenAIMetricAttrsFromContext(ctx); genAIAttrs != nil {
genAIAttrs.OperationName = "get_prompt"
genAIAttrs.PromptName = promptName
}
// Verify prompt belongs to the current group before resolving globally.
if !g.ContainsPrompt(promptName) {
err := fmt.Errorf("prompt with name %q does not exist", promptName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
@@ -462,12 +468,6 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prom
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
// Populate gen_ai attributes for operation duration metric
if genAIAttrs := util.GenAIMetricAttrsFromContext(ctx); genAIAttrs != nil {
genAIAttrs.OperationName = "get_prompt"
genAIAttrs.PromptName = promptName
}
// Parse the arguments provided in the request.
argValues, err := prompt.ParseArgs(req.Params.Arguments, nil)
if err != nil {
@@ -514,3 +514,77 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prom
Result: result,
}, nil
}
// groupsListHandler handles the "groups/list" method. It returns every named
// group's name and description. The default nameless group is omitted.
func groupsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
logger.DebugContext(ctx, "handling groups/list request")
var req ListGroupsRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp groups list request: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
result := GenerateListGroupsResult(primitiveMgr.GetGroupsMap())
logger.DebugContext(ctx, fmt.Sprintf("returning %d groups", len(result.Groups)))
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: result,
}, nil
}
// groupsGetHandler handles the "groups/get" method. It returns the named group's
// tools and prompts.
func groupsGetHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
logger.DebugContext(ctx, "handling groups/get request")
var req GetGroupRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp groups/get request: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
groupName := req.Params.Name
logger.DebugContext(ctx, fmt.Sprintf("group name: %s", groupName))
// Update span name and set gen_ai attributes
span := trace.SpanFromContext(ctx)
span.SetName(fmt.Sprintf("%s %s", GROUPS_GET, groupName))
span.SetAttributes(attribute.String("gen_ai.group.name", groupName))
// Populate gen_ai attributes for operation duration metric
if genAIAttrs := util.GenAIMetricAttrsFromContext(ctx); genAIAttrs != nil {
genAIAttrs.OperationName = "get_group"
genAIAttrs.GroupName = groupName
}
g, ok := primitiveMgr.GetGroup(groupName)
if !ok {
err := fmt.Errorf("invalid group name: group with name %q does not exist", groupName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
urlParams, _ := util.UrlParamsFromContext(ctx)
result, err := GenerateGetGroupResult(primitiveMgr.GetSourcesMap(), g, primitiveMgr.GetToolsMap(), primitiveMgr.GetPromptsMap(), urlParams)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: result,
}, nil
}
internal/server/mcp/v20250618/method_test.go
+202 -17
View File
@@ -21,11 +21,11 @@ import (
"strings"
"testing"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/log"
"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
"github.com/googleapis/mcp-toolbox/internal/server/primitives"
"github.com/googleapis/mcp-toolbox/internal/testutils"
"github.com/googleapis/mcp-toolbox/internal/tools"
"github.com/googleapis/mcp-toolbox/internal/util"
)
@@ -35,6 +35,16 @@ var (
fakeVersionString = "0.0.0"
)
// mustGroup fetches the default group from the resource manager.
func mustGroup(t *testing.T, rm *primitives.PrimitiveManager) group.Group {
t.Helper()
g, ok := rm.GetGroup("")
if !ok {
t.Fatal("default group not found")
}
return g
}
func TestInitializeHandler(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
@@ -167,21 +177,21 @@ func TestPingHandler(t *testing.T) {
func TestToolsListHandler(t *testing.T) {
// Initialize tools using provided testutils mock instances
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
body ListToolsRequest
rawBody []byte
toolset tools.Toolset
g group.Group
wantErr bool
errContains string
}{
{
name: "invalid json body",
rawBody: []byte(`{invalid json}`),
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: true,
errContains: "invalid mcp tools list request",
},
@@ -194,7 +204,7 @@ func TestToolsListHandler(t *testing.T) {
},
},
},
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: false,
},
{
@@ -206,7 +216,7 @@ func TestToolsListHandler(t *testing.T) {
},
},
},
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: false,
},
}
@@ -221,7 +231,7 @@ func TestToolsListHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := toolsListHandler(context.Background(), dummyID, primitiveMgr, tt.toolset, body)
got, err := toolsListHandler(context.Background(), dummyID, primitiveMgr, tt.g, body)
if tt.wantErr {
if err == nil {
@@ -256,8 +266,8 @@ func TestToolsCallHandler(t *testing.T) {
testutils.MockTool4,
testutils.MockTool5,
}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
@@ -353,7 +363,7 @@ func TestToolsCallHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := toolsCallHandler(tt.context, dummyID, toolsets[""], primitiveMgr, body, nil)
got, err := toolsCallHandler(tt.context, dummyID, mustGroup(t, primitiveMgr), primitiveMgr, body, nil)
if tt.wantErr {
if err == nil {
@@ -384,8 +394,8 @@ func TestPromptsListHandler(t *testing.T) {
ctx = util.WithLogger(ctx, testLogger)
// Initialize prompts
mockPrompts := []testutils.MockPrompt{testutils.MockPrompt1, testutils.MockPrompt2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
body ListPromptsRequest
@@ -422,7 +432,7 @@ func TestPromptsListHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := promptsListHandler(ctx, dummyID, primitiveMgr, promptsets[""], body)
got, err := promptsListHandler(ctx, dummyID, primitiveMgr, mustGroup(t, primitiveMgr), body)
if tt.wantErr {
if err == nil {
@@ -453,8 +463,8 @@ func TestPromptsGetHandler(t *testing.T) {
ctx = util.WithLogger(ctx, testLogger)
// Initialize prompts
mockPrompts := []testutils.MockPrompt{testutils.MockPrompt1, testutils.MockPrompt2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
body GetPromptRequest
@@ -529,7 +539,7 @@ func TestPromptsGetHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := promptsGetHandler(ctx, dummyID, promptsets[""], primitiveMgr, body)
got, err := promptsGetHandler(ctx, dummyID, mustGroup(t, primitiveMgr), primitiveMgr, body)
if tt.wantErr {
if err == nil {
@@ -549,3 +559,178 @@ func TestPromptsGetHandler(t *testing.T) {
})
}
}
func TestGroupsListHandler(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info")
if err != nil {
t.Fatalf("unable to initialize logger: %s", err)
}
ctx = util.WithLogger(ctx, testLogger)
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
rawBody []byte
body ListGroupsRequest
wantErr bool
errContains string
wantNames []string
}{
{
name: "invalid json body",
rawBody: []byte(`{invalid json}`),
wantErr: true,
errContains: "invalid mcp groups list request",
},
{
name: "success excludes default group and sorts",
body: ListGroupsRequest{
PaginatedRequest: PaginatedRequest{
Request: jsonrpc.Request{Method: GROUPS_LIST},
},
},
wantErr: false,
wantNames: []string{"tool1_only", "tool2_only"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
body := tt.rawBody
var err error
if body == nil {
body, err = json.Marshal(tt.body)
if err != nil {
t.Fatalf("unexpected error during marshaling: %v", err)
}
}
got, err := groupsListHandler(ctx, dummyID, primitiveMgr, body)
if tt.wantErr {
if err == nil {
t.Fatalf("expected error, got nil")
}
if tt.errContains != "" && !strings.Contains(err.Error(), tt.errContains) {
t.Errorf("error = %v, want string containing %q", err, tt.errContains)
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
res, ok := got.(jsonrpc.JSONRPCResponse)
if !ok {
t.Fatalf("expected jsonrpc.JSONRPCResponse, got %T", got)
}
result, ok := res.Result.(ListGroupsResult)
if !ok {
t.Fatalf("expected ListGroupsResult, got %T", res.Result)
}
gotNames := make([]string, 0, len(result.Groups))
for _, g := range result.Groups {
gotNames = append(gotNames, g.Name)
}
if len(gotNames) != len(tt.wantNames) {
t.Fatalf("got groups %v, want %v", gotNames, tt.wantNames)
}
for i, n := range tt.wantNames {
if gotNames[i] != n {
t.Errorf("group[%d] = %q, want %q", i, gotNames[i], n)
}
}
})
}
}
func TestGroupsGetHandler(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info")
if err != nil {
t.Fatalf("unable to initialize logger: %s", err)
}
ctx = util.WithLogger(ctx, testLogger)
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
rawBody []byte
body GetGroupRequest
wantErr bool
errContains string
wantName string
}{
{
name: "invalid json body",
rawBody: []byte(`{invalid json}`),
wantErr: true,
errContains: "invalid mcp groups/get request",
},
{
name: "group does not exist",
body: GetGroupRequest{
Request: jsonrpc.Request{Method: GROUPS_GET},
Params: struct {
Name string `json:"name"`
}{Name: "missing_group"},
},
wantErr: true,
errContains: `group with name "missing_group" does not exist`,
},
{
name: "success",
body: GetGroupRequest{
Request: jsonrpc.Request{Method: GROUPS_GET},
Params: struct {
Name string `json:"name"`
}{Name: "tool1_only"},
},
wantErr: false,
wantName: "tool1_only",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
body := tt.rawBody
var err error
if body == nil {
body, err = json.Marshal(tt.body)
if err != nil {
t.Fatalf("unexpected error during marshaling: %v", err)
}
}
got, err := groupsGetHandler(ctx, dummyID, primitiveMgr, body)
if tt.wantErr {
if err == nil {
t.Fatalf("expected error, got nil")
}
if tt.errContains != "" && !strings.Contains(err.Error(), tt.errContains) {
t.Errorf("error = %v, want string containing %q", err, tt.errContains)
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
res, ok := got.(jsonrpc.JSONRPCResponse)
if !ok {
t.Fatalf("expected jsonrpc.JSONRPCResponse, got %T", got)
}
result, ok := res.Result.(GetGroupResult)
if !ok {
t.Fatalf("expected GetGroupResult, got %T", res.Result)
}
if result.Name != tt.wantName {
t.Errorf("result.Name = %q, want %q", result.Name, tt.wantName)
}
})
}
}
internal/server/mcp/v20250618/types.go
+40
View File
@@ -34,6 +34,8 @@ const (
TOOLS_CALL = "tools/call"
PROMPTS_LIST = "prompts/list"
PROMPTS_GET = "prompts/get"
GROUPS_LIST = "groups/list"
GROUPS_GET = "groups/get"
)
/* Initialization */
@@ -343,3 +345,41 @@ type PromptMessage struct {
Role string `json:"role"`
Content TextContent `json:"content"`
}
/* Groups */
// ListGroupsRequest is sent from the client to request the list of groups the
// server has.
type ListGroupsRequest struct {
PaginatedRequest
}
// Group is a single entry in a groups/list response.
type Group struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
}
// ListGroupsResult is the server's response to a groups/list request.
type ListGroupsResult struct {
jsonrpc.Result
Groups []Group `json:"groups"`
}
// GetGroupRequest is sent from the client to request a single group's contents.
type GetGroupRequest struct {
jsonrpc.Request
Params struct {
Name string `json:"name"`
} `json:"params"`
}
// GetGroupResult is the server's response to a groups/get request: the group's
// tools and prompts. The description is intentionally omitted; it is exposed only
// through groups/list.
type GetGroupResult struct {
jsonrpc.Result
Name string `json:"name"`
Tools []Tool `json:"tools"`
Prompts []Prompt `json:"prompts"`
}
internal/server/mcp/v20251125/manifests.go
+46 -6
View File
@@ -16,7 +16,9 @@ package v20251125
import (
"fmt"
"sort"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/sources"
"github.com/googleapis/mcp-toolbox/internal/tools"
@@ -98,9 +100,9 @@ func generateParamManifest(ps parameters.Parameters, urlParams map[string]string
}
// GenerateListToolsResult generates tools/list method result according to mcp schema
func GenerateListToolsResult(srcs map[string]sources.Source, t tools.Toolset, toolsMap map[string]tools.Tool, urlParams map[string]string) (ListToolsResult, error) {
mcpManifest := make([]Tool, 0, len(t.ToolNames))
for _, toolName := range t.ToolNames {
func GenerateListToolsResult(srcs map[string]sources.Source, g group.Group, toolsMap map[string]tools.Tool, urlParams map[string]string) (ListToolsResult, error) {
mcpManifest := make([]Tool, 0, len(g.ToolNames))
for _, toolName := range g.ToolNames {
tool, ok := toolsMap[toolName]
if !ok {
return ListToolsResult{}, fmt.Errorf("tool does not exist: %s", toolName)
@@ -134,9 +136,9 @@ func generatePromptManifest(name, desc string, args prompts.Arguments) Prompt {
}
// GenerateListPromptsResult generates the list/prompts result
func GenerateListPromptsResult(p prompts.Promptset, promptsMap map[string]prompts.Prompt) (ListPromptsResult, error) {
mcpManifest := make([]Prompt, 0, len(p.PromptNames))
for _, promptName := range p.PromptNames {
func GenerateListPromptsResult(g group.Group, promptsMap map[string]prompts.Prompt) (ListPromptsResult, error) {
mcpManifest := make([]Prompt, 0, len(g.PromptNames))
for _, promptName := range g.PromptNames {
prompt, ok := promptsMap[promptName]
if !ok {
return ListPromptsResult{}, fmt.Errorf("prompt does not exist: %s", promptName)
@@ -146,3 +148,41 @@ func GenerateListPromptsResult(p prompts.Promptset, promptsMap map[string]prompt
}
return ListPromptsResult{Prompts: mcpManifest}, nil
}
// GenerateListGroupsResult generates the groups/list result. It omits the
// default nameless group and returns the remaining groups sorted by name.
func GenerateListGroupsResult(groupsMap map[string]group.Group) ListGroupsResult {
names := make([]string, 0, len(groupsMap))
for name := range groupsMap {
if name == "" {
continue
}
names = append(names, name)
}
sort.Strings(names)
groupsList := make([]Group, 0, len(names))
for _, name := range names {
g := groupsMap[name]
groupsList = append(groupsList, Group{Name: g.Name, Description: g.Description})
}
return ListGroupsResult{Groups: groupsList}
}
// GenerateGetGroupResult generates the groups/get result for a single group's
// tools and prompts.
func GenerateGetGroupResult(srcs map[string]sources.Source, g group.Group, toolsMap map[string]tools.Tool, promptsMap map[string]prompts.Prompt, urlParams map[string]string) (GetGroupResult, error) {
listToolsResult, err := GenerateListToolsResult(srcs, g, toolsMap, urlParams)
if err != nil {
return GetGroupResult{}, fmt.Errorf("error generating tools manifest: %w", err)
}
listPromptsResult, err := GenerateListPromptsResult(g, promptsMap)
if err != nil {
return GetGroupResult{}, fmt.Errorf("error generating prompts manifest: %w", err)
}
return GetGroupResult{
Name: g.Name,
Tools: listToolsResult.Tools,
Prompts: listPromptsResult.Prompts,
}, nil
}
internal/server/mcp/v20251125/manifests_test.go
+7 -14
View File
@@ -21,6 +21,7 @@ import (
"testing"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/testutils"
"github.com/googleapis/mcp-toolbox/internal/tools"
@@ -240,16 +241,12 @@ func TestGenerateListToolsResult(t *testing.T) {
toolsMap := make(map[string]tools.Tool)
toolsMap[tool1.Name] = tool1
toolsMap[tool2.Name] = tool2
tc := tools.ToolsetConfig{
g := group.NewGroup(group.GroupConfig{
Name: "test-toolset",
ToolNames: []string{"no_params", "some_params"},
}
toolset, err := tc.Initialize("test-version", toolsMap)
if err != nil {
t.Fatalf("unable to initialize toolset %q: %s", "test-toolset", err)
}
})
got, err := GenerateListToolsResult(nil, toolset, toolsMap, nil)
got, err := GenerateListToolsResult(nil, g, toolsMap, nil)
if err != nil {
t.Fatalf("unable to generate list tools result: %s", err)
}
@@ -353,16 +350,12 @@ func TestGenerateListPromptsResult(t *testing.T) {
promptsMap := make(map[string]prompts.Prompt)
promptsMap[prompt1.Name] = prompt1
promptsMap[prompt2.Name] = prompt2
pc := prompts.PromptsetConfig{
g := group.NewGroup(group.GroupConfig{
Name: "test-promptset",
PromptNames: []string{"prompt1", "prompt2"},
}
promptset, err := pc.Initialize("test-version", promptsMap)
if err != nil {
t.Fatalf("unable to initialize promptset %q: %s", "test-promptset", err)
}
})
got, err := GenerateListPromptsResult(promptset, promptsMap)
got, err := GenerateListPromptsResult(g, promptsMap)
if err != nil {
t.Fatalf("unable to generate list prompt result: %s", err)
}
internal/server/mcp/v20251125/method.go
+96 -22
View File
@@ -21,10 +21,10 @@ import (
"errors"
"fmt"
"net/http"
"time"
"github.com/googleapis/mcp-toolbox/internal/auth"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
mcputil "github.com/googleapis/mcp-toolbox/internal/server/mcp/util"
@@ -38,20 +38,20 @@ import (
)
// ProcessMethod returns a response for the request.
func ProcessMethod(ctx context.Context, id jsonrpc.RequestId, method string, toolset tools.Toolset, promptset prompts.Promptset, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
func ProcessMethod(ctx context.Context, id jsonrpc.RequestId, method string, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
switch method {
case INITIALIZE:
return initializeHandler(ctx, id, body)
case PING:
return pingHandler(id)
case TOOLS_LIST:
return toolsListHandler(ctx, id, primitiveMgr, toolset, body)
return toolsListHandler(ctx, id, primitiveMgr, g, body)
case TOOLS_CALL:
return toolsCallHandler(ctx, id, toolset, primitiveMgr, body, header)
return toolsCallHandler(ctx, id, g, primitiveMgr, body, header)
case PROMPTS_LIST:
return promptsListHandler(ctx, id, primitiveMgr, promptset, body)
return promptsListHandler(ctx, id, primitiveMgr, g, body)
case PROMPTS_GET:
return promptsGetHandler(ctx, id, promptset, primitiveMgr, body)
return promptsGetHandler(ctx, id, g, primitiveMgr, body)
default:
err := fmt.Errorf("invalid method %s", method)
return jsonrpc.NewError(id, jsonrpc.METHOD_NOT_FOUND, err.Error(), nil), err
@@ -110,7 +110,7 @@ func pingHandler(id jsonrpc.RequestId) (any, error) {
}, nil
}
func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, toolset tools.Toolset, body []byte) (any, error) {
func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, g group.Group, body []byte) (any, error) {
var req ListToolsRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp tools list request: %w", err)
@@ -119,7 +119,7 @@ func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *p
urlParams, _ := util.UrlParamsFromContext(ctx)
toolsMap := primitiveMgr.GetToolsMap()
listToolsResult, err := GenerateListToolsResult(primitiveMgr.GetSourcesMap(), toolset, toolsMap, urlParams)
listToolsResult, err := GenerateListToolsResult(primitiveMgr.GetSourcesMap(), g, toolsMap, urlParams)
if err != nil {
err = fmt.Errorf("error generating manifest: %w", err)
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
@@ -132,7 +132,7 @@ func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *p
}
// toolsCallHandler generate a response for tools call.
func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, toolset tools.Toolset, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
if header != nil {
if clientIP := util.ExtractClientIP(header); clientIP != "" {
ctx = util.WithClientIP(ctx, clientIP)
@@ -165,8 +165,8 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, toolset tools.T
attribute.String("gen_ai.operation.name", "execute_tool"),
)
// Verify tool belongs to the current toolset before resolving globally.
if !toolset.ContainsTool(toolName) {
// Verify tool belongs to the current group before resolving globally.
if !g.ContainsTool(toolName) {
err = fmt.Errorf("invalid tool name: tool with name %q does not exist", toolName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
@@ -399,7 +399,7 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, toolset tools.T
}
// promptsListHandler handles the "prompts/list" method.
func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, promptset prompts.Promptset, body []byte) (any, error) {
func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, g group.Group, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
@@ -414,7 +414,7 @@ func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr
}
promptsMap := primitiveMgr.GetPromptsMap()
listPromptsResult, err := GenerateListPromptsResult(promptset, promptsMap)
listPromptsResult, err := GenerateListPromptsResult(g, promptsMap)
if err != nil {
err = fmt.Errorf("error generating manifest: %w", err)
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
@@ -428,7 +428,7 @@ func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr
}
// promptsGetHandler handles the "prompts/get" method.
func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prompts.Promptset, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
@@ -450,8 +450,14 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prom
span.SetName(fmt.Sprintf("%s %s", PROMPTS_GET, promptName))
span.SetAttributes(attribute.String("gen_ai.prompt.name", promptName))
// Verify prompt belongs to the current promptset before resolving globally.
if !promptset.ContainsPrompt(promptName) {
// Populate gen_ai attributes for operation duration metric
if genAIAttrs := util.GenAIMetricAttrsFromContext(ctx); genAIAttrs != nil {
genAIAttrs.OperationName = "get_prompt"
genAIAttrs.PromptName = promptName
}
// Verify prompt belongs to the current group before resolving globally.
if !g.ContainsPrompt(promptName) {
err := fmt.Errorf("prompt with name %q does not exist", promptName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
@@ -462,12 +468,6 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prom
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
// Populate gen_ai attributes for operation duration metric
if genAIAttrs := util.GenAIMetricAttrsFromContext(ctx); genAIAttrs != nil {
genAIAttrs.OperationName = "get_prompt"
genAIAttrs.PromptName = promptName
}
// Parse the arguments provided in the request.
argValues, err := prompt.ParseArgs(req.Params.Arguments, nil)
if err != nil {
@@ -514,3 +514,77 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prom
Result: result,
}, nil
}
// groupsListHandler handles the "groups/list" method. It returns every named
// group's name and description. The default nameless group is omitted.
func groupsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
logger.DebugContext(ctx, "handling groups/list request")
var req ListGroupsRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp groups list request: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
result := GenerateListGroupsResult(primitiveMgr.GetGroupsMap())
logger.DebugContext(ctx, fmt.Sprintf("returning %d groups", len(result.Groups)))
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: result,
}, nil
}
// groupsGetHandler handles the "groups/get" method. It returns the named group's
// tools and prompts.
func groupsGetHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
logger.DebugContext(ctx, "handling groups/get request")
var req GetGroupRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp groups/get request: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
groupName := req.Params.Name
logger.DebugContext(ctx, fmt.Sprintf("group name: %s", groupName))
// Update span name and set gen_ai attributes
span := trace.SpanFromContext(ctx)
span.SetName(fmt.Sprintf("%s %s", GROUPS_GET, groupName))
span.SetAttributes(attribute.String("gen_ai.group.name", groupName))
// Populate gen_ai attributes for operation duration metric
if genAIAttrs := util.GenAIMetricAttrsFromContext(ctx); genAIAttrs != nil {
genAIAttrs.OperationName = "get_group"
genAIAttrs.GroupName = groupName
}
g, ok := primitiveMgr.GetGroup(groupName)
if !ok {
err := fmt.Errorf("invalid group name: group with name %q does not exist", groupName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
urlParams, _ := util.UrlParamsFromContext(ctx)
result, err := GenerateGetGroupResult(primitiveMgr.GetSourcesMap(), g, primitiveMgr.GetToolsMap(), primitiveMgr.GetPromptsMap(), urlParams)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: result,
}, nil
}
internal/server/mcp/v20251125/method_test.go
+202 -17
View File
@@ -21,11 +21,11 @@ import (
"strings"
"testing"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/log"
"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
"github.com/googleapis/mcp-toolbox/internal/server/primitives"
"github.com/googleapis/mcp-toolbox/internal/testutils"
"github.com/googleapis/mcp-toolbox/internal/tools"
"github.com/googleapis/mcp-toolbox/internal/util"
)
@@ -35,6 +35,16 @@ var (
fakeVersionString = "0.0.0"
)
// mustGroup fetches the default group from the resource manager.
func mustGroup(t *testing.T, rm *primitives.PrimitiveManager) group.Group {
t.Helper()
g, ok := rm.GetGroup("")
if !ok {
t.Fatal("default group not found")
}
return g
}
func TestInitializeHandler(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
@@ -167,21 +177,21 @@ func TestPingHandler(t *testing.T) {
func TestToolsListHandler(t *testing.T) {
// Initialize tools using provided testutils mock instances
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
body ListToolsRequest
rawBody []byte
toolset tools.Toolset
g group.Group
wantErr bool
errContains string
}{
{
name: "invalid json body",
rawBody: []byte(`{invalid json}`),
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: true,
errContains: "invalid mcp tools list request",
},
@@ -194,7 +204,7 @@ func TestToolsListHandler(t *testing.T) {
},
},
},
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: false,
},
{
@@ -206,7 +216,7 @@ func TestToolsListHandler(t *testing.T) {
},
},
},
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: false,
},
}
@@ -221,7 +231,7 @@ func TestToolsListHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := toolsListHandler(context.Background(), dummyID, primitiveMgr, tt.toolset, body)
got, err := toolsListHandler(context.Background(), dummyID, primitiveMgr, tt.g, body)
if tt.wantErr {
if err == nil {
@@ -256,8 +266,8 @@ func TestToolsCallHandler(t *testing.T) {
testutils.MockTool4,
testutils.MockTool5,
}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
@@ -353,7 +363,7 @@ func TestToolsCallHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := toolsCallHandler(tt.context, dummyID, toolsets[""], primitiveMgr, body, nil)
got, err := toolsCallHandler(tt.context, dummyID, mustGroup(t, primitiveMgr), primitiveMgr, body, nil)
if tt.wantErr {
if err == nil {
@@ -384,8 +394,8 @@ func TestPromptsListHandler(t *testing.T) {
ctx = util.WithLogger(ctx, testLogger)
// Initialize prompts
mockPrompts := []testutils.MockPrompt{testutils.MockPrompt1, testutils.MockPrompt2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
body ListPromptsRequest
@@ -422,7 +432,7 @@ func TestPromptsListHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := promptsListHandler(ctx, dummyID, primitiveMgr, promptsets[""], body)
got, err := promptsListHandler(ctx, dummyID, primitiveMgr, mustGroup(t, primitiveMgr), body)
if tt.wantErr {
if err == nil {
@@ -453,8 +463,8 @@ func TestPromptsGetHandler(t *testing.T) {
ctx = util.WithLogger(ctx, testLogger)
// Initialize prompts
mockPrompts := []testutils.MockPrompt{testutils.MockPrompt1, testutils.MockPrompt2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
body GetPromptRequest
@@ -529,7 +539,7 @@ func TestPromptsGetHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := promptsGetHandler(ctx, dummyID, promptsets[""], primitiveMgr, body)
got, err := promptsGetHandler(ctx, dummyID, mustGroup(t, primitiveMgr), primitiveMgr, body)
if tt.wantErr {
if err == nil {
@@ -549,3 +559,178 @@ func TestPromptsGetHandler(t *testing.T) {
})
}
}
func TestGroupsListHandler(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info")
if err != nil {
t.Fatalf("unable to initialize logger: %s", err)
}
ctx = util.WithLogger(ctx, testLogger)
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
rawBody []byte
body ListGroupsRequest
wantErr bool
errContains string
wantNames []string
}{
{
name: "invalid json body",
rawBody: []byte(`{invalid json}`),
wantErr: true,
errContains: "invalid mcp groups list request",
},
{
name: "success excludes default group and sorts",
body: ListGroupsRequest{
PaginatedRequest: PaginatedRequest{
Request: jsonrpc.Request{Method: GROUPS_LIST},
},
},
wantErr: false,
wantNames: []string{"tool1_only", "tool2_only"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
body := tt.rawBody
var err error
if body == nil {
body, err = json.Marshal(tt.body)
if err != nil {
t.Fatalf("unexpected error during marshaling: %v", err)
}
}
got, err := groupsListHandler(ctx, dummyID, primitiveMgr, body)
if tt.wantErr {
if err == nil {
t.Fatalf("expected error, got nil")
}
if tt.errContains != "" && !strings.Contains(err.Error(), tt.errContains) {
t.Errorf("error = %v, want string containing %q", err, tt.errContains)
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
res, ok := got.(jsonrpc.JSONRPCResponse)
if !ok {
t.Fatalf("expected jsonrpc.JSONRPCResponse, got %T", got)
}
result, ok := res.Result.(ListGroupsResult)
if !ok {
t.Fatalf("expected ListGroupsResult, got %T", res.Result)
}
gotNames := make([]string, 0, len(result.Groups))
for _, g := range result.Groups {
gotNames = append(gotNames, g.Name)
}
if len(gotNames) != len(tt.wantNames) {
t.Fatalf("got groups %v, want %v", gotNames, tt.wantNames)
}
for i, n := range tt.wantNames {
if gotNames[i] != n {
t.Errorf("group[%d] = %q, want %q", i, gotNames[i], n)
}
}
})
}
}
func TestGroupsGetHandler(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info")
if err != nil {
t.Fatalf("unable to initialize logger: %s", err)
}
ctx = util.WithLogger(ctx, testLogger)
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
rawBody []byte
body GetGroupRequest
wantErr bool
errContains string
wantName string
}{
{
name: "invalid json body",
rawBody: []byte(`{invalid json}`),
wantErr: true,
errContains: "invalid mcp groups/get request",
},
{
name: "group does not exist",
body: GetGroupRequest{
Request: jsonrpc.Request{Method: GROUPS_GET},
Params: struct {
Name string `json:"name"`
}{Name: "missing_group"},
},
wantErr: true,
errContains: `group with name "missing_group" does not exist`,
},
{
name: "success",
body: GetGroupRequest{
Request: jsonrpc.Request{Method: GROUPS_GET},
Params: struct {
Name string `json:"name"`
}{Name: "tool1_only"},
},
wantErr: false,
wantName: "tool1_only",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
body := tt.rawBody
var err error
if body == nil {
body, err = json.Marshal(tt.body)
if err != nil {
t.Fatalf("unexpected error during marshaling: %v", err)
}
}
got, err := groupsGetHandler(ctx, dummyID, primitiveMgr, body)
if tt.wantErr {
if err == nil {
t.Fatalf("expected error, got nil")
}
if tt.errContains != "" && !strings.Contains(err.Error(), tt.errContains) {
t.Errorf("error = %v, want string containing %q", err, tt.errContains)
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
res, ok := got.(jsonrpc.JSONRPCResponse)
if !ok {
t.Fatalf("expected jsonrpc.JSONRPCResponse, got %T", got)
}
result, ok := res.Result.(GetGroupResult)
if !ok {
t.Fatalf("expected GetGroupResult, got %T", res.Result)
}
if result.Name != tt.wantName {
t.Errorf("result.Name = %q, want %q", result.Name, tt.wantName)
}
})
}
}
internal/server/mcp/v20251125/types.go
+40
View File
@@ -34,6 +34,8 @@ const (
TOOLS_CALL = "tools/call"
PROMPTS_LIST = "prompts/list"
PROMPTS_GET = "prompts/get"
GROUPS_LIST = "groups/list"
GROUPS_GET = "groups/get"
)
/* Initialization */
@@ -347,3 +349,41 @@ type PromptMessage struct {
Role string `json:"role"`
Content TextContent `json:"content"`
}
/* Groups */
// ListGroupsRequest is sent from the client to request the list of groups the
// server has.
type ListGroupsRequest struct {
PaginatedRequest
}
// Group is a single entry in a groups/list response.
type Group struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
}
// ListGroupsResult is the server's response to a groups/list request.
type ListGroupsResult struct {
jsonrpc.Result
Groups []Group `json:"groups"`
}
// GetGroupRequest is sent from the client to request a single group's contents.
type GetGroupRequest struct {
jsonrpc.Request
Params struct {
Name string `json:"name"`
} `json:"params"`
}
// GetGroupResult is the server's response to a groups/get request: the group's
// tools and prompts. The description is intentionally omitted; it is exposed only
// through groups/list.
type GetGroupResult struct {
jsonrpc.Result
Name string `json:"name"`
Tools []Tool `json:"tools"`
Prompts []Prompt `json:"prompts"`
}
internal/server/mcp/vdraft/manifests.go
+46 -6
View File
@@ -16,7 +16,9 @@ package vdraft
import (
"fmt"
"sort"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/sources"
"github.com/googleapis/mcp-toolbox/internal/tools"
@@ -98,9 +100,9 @@ func generateParamManifest(ps parameters.Parameters, urlParams map[string]string
}
// GenerateListToolsResult generates tools/list method result according to mcp schema
func GenerateListToolsResult(srcs map[string]sources.Source, t tools.Toolset, toolsMap map[string]tools.Tool, urlParams map[string]string) (ListToolsResult, error) {
mcpManifest := make([]Tool, 0, len(t.ToolNames))
for _, toolName := range t.ToolNames {
func GenerateListToolsResult(srcs map[string]sources.Source, g group.Group, toolsMap map[string]tools.Tool, urlParams map[string]string) (ListToolsResult, error) {
mcpManifest := make([]Tool, 0, len(g.ToolNames))
for _, toolName := range g.ToolNames {
tool, ok := toolsMap[toolName]
if !ok {
return ListToolsResult{}, fmt.Errorf("tool does not exist: %s", toolName)
@@ -144,9 +146,9 @@ func generatePromptManifest(name, desc string, args prompts.Arguments) Prompt {
}
// GenerateListPromptsResult generates the list/prompts result
func GenerateListPromptsResult(p prompts.Promptset, promptsMap map[string]prompts.Prompt) (ListPromptsResult, error) {
mcpManifest := make([]Prompt, 0, len(p.PromptNames))
for _, promptName := range p.PromptNames {
func GenerateListPromptsResult(g group.Group, promptsMap map[string]prompts.Prompt) (ListPromptsResult, error) {
mcpManifest := make([]Prompt, 0, len(g.PromptNames))
for _, promptName := range g.PromptNames {
prompt, ok := promptsMap[promptName]
if !ok {
return ListPromptsResult{}, fmt.Errorf("prompt does not exist: %s", promptName)
@@ -166,3 +168,41 @@ func GenerateListPromptsResult(p prompts.Promptset, promptsMap map[string]prompt
}
return res, nil
}
// GenerateListGroupsResult generates the groups/list result. It omits the
// default nameless group and returns the remaining groups sorted by name.
func GenerateListGroupsResult(groupsMap map[string]group.Group) ListGroupsResult {
names := make([]string, 0, len(groupsMap))
for name := range groupsMap {
if name == "" {
continue
}
names = append(names, name)
}
sort.Strings(names)
groupsList := make([]Group, 0, len(names))
for _, name := range names {
g := groupsMap[name]
groupsList = append(groupsList, Group{Name: g.Name, Description: g.Description})
}
return ListGroupsResult{Groups: groupsList}
}
// GenerateGetGroupResult generates the groups/get result for a single group's
// tools and prompts.
func GenerateGetGroupResult(srcs map[string]sources.Source, g group.Group, toolsMap map[string]tools.Tool, promptsMap map[string]prompts.Prompt, urlParams map[string]string) (GetGroupResult, error) {
listToolsResult, err := GenerateListToolsResult(srcs, g, toolsMap, urlParams)
if err != nil {
return GetGroupResult{}, fmt.Errorf("error generating tools manifest: %w", err)
}
listPromptsResult, err := GenerateListPromptsResult(g, promptsMap)
if err != nil {
return GetGroupResult{}, fmt.Errorf("error generating prompts manifest: %w", err)
}
return GetGroupResult{
Name: g.Name,
Tools: listToolsResult.Tools,
Prompts: listPromptsResult.Prompts,
}, nil
}
internal/server/mcp/vdraft/manifests_test.go
+7 -14
View File
@@ -21,6 +21,7 @@ import (
"testing"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/testutils"
"github.com/googleapis/mcp-toolbox/internal/tools"
@@ -240,16 +241,12 @@ func TestGenerateListToolsResult(t *testing.T) {
toolsMap := make(map[string]tools.Tool)
toolsMap[tool1.Name] = tool1
toolsMap[tool2.Name] = tool2
tc := tools.ToolsetConfig{
g := group.NewGroup(group.GroupConfig{
Name: "test-toolset",
ToolNames: []string{"no_params", "some_params"},
}
toolset, err := tc.Initialize("test-version", toolsMap)
if err != nil {
t.Fatalf("unable to initialize toolset %q: %s", "test-toolset", err)
}
})
got, err := GenerateListToolsResult(nil, toolset, toolsMap, nil)
got, err := GenerateListToolsResult(nil, g, toolsMap, nil)
if err != nil {
t.Fatalf("unable to generate list tools result: %s", err)
}
@@ -360,16 +357,12 @@ func TestGenerateListPromptsResult(t *testing.T) {
promptsMap := make(map[string]prompts.Prompt)
promptsMap[prompt1.Name] = prompt1
promptsMap[prompt2.Name] = prompt2
pc := prompts.PromptsetConfig{
g := group.NewGroup(group.GroupConfig{
Name: "test-promptset",
PromptNames: []string{"prompt1", "prompt2"},
}
promptset, err := pc.Initialize("test-version", promptsMap)
if err != nil {
t.Fatalf("unable to initialize promptset %q: %s", "test-promptset", err)
}
})
got, err := GenerateListPromptsResult(promptset, promptsMap)
got, err := GenerateListPromptsResult(g, promptsMap)
if err != nil {
t.Fatalf("unable to generate list prompt result: %s", err)
}
internal/server/mcp/vdraft/method.go
+112 -21
View File
@@ -24,6 +24,7 @@ import (
"time"
"github.com/googleapis/mcp-toolbox/internal/auth"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
mcputil "github.com/googleapis/mcp-toolbox/internal/server/mcp/util"
@@ -37,18 +38,18 @@ import (
)
// ProcessMethod returns a response for the request.
func ProcessMethod(ctx context.Context, id jsonrpc.RequestId, method string, toolset tools.Toolset, promptset prompts.Promptset, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
func ProcessMethod(ctx context.Context, id jsonrpc.RequestId, method string, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
switch method {
case SERVER_DISCOVER:
return serverDiscoverHandler(ctx, id, body, header)
case TOOLS_LIST:
return toolsListHandler(ctx, id, primitiveMgr, toolset, body, header)
return toolsListHandler(ctx, id, primitiveMgr, g, body, header)
case TOOLS_CALL:
return toolsCallHandler(ctx, id, toolset, primitiveMgr, body, header)
return toolsCallHandler(ctx, id, g, primitiveMgr, body, header)
case PROMPTS_LIST:
return promptsListHandler(ctx, id, primitiveMgr, promptset, body, header)
return promptsListHandler(ctx, id, primitiveMgr, g, body, header)
case PROMPTS_GET:
return promptsGetHandler(ctx, id, promptset, primitiveMgr, body, header)
return promptsGetHandler(ctx, id, g, primitiveMgr, body, header)
default:
err := fmt.Errorf("invalid method %s", method)
return jsonrpc.NewError(id, jsonrpc.METHOD_NOT_FOUND, err.Error(), nil), err
@@ -210,7 +211,7 @@ func serverDiscoverHandler(ctx context.Context, id jsonrpc.RequestId, body []byt
return res, nil
}
func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, toolset tools.Toolset, body []byte, header http.Header) (any, error) {
func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, g group.Group, body []byte, header http.Header) (any, error) {
var req ListToolsRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp tools list request: %w", err)
@@ -227,7 +228,7 @@ func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *p
urlParams, _ := util.UrlParamsFromContext(ctx)
toolsMap := primitiveMgr.GetToolsMap()
listToolsResult, err := GenerateListToolsResult(primitiveMgr.GetSourcesMap(), toolset, toolsMap, urlParams)
listToolsResult, err := GenerateListToolsResult(primitiveMgr.GetSourcesMap(), g, toolsMap, urlParams)
if err != nil {
err = fmt.Errorf("error generating manifest: %w", err)
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
@@ -245,7 +246,7 @@ func toolsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *p
}
// toolsCallHandler generate a response for tools call.
func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, toolset tools.Toolset, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
authServices := primitiveMgr.GetAuthServiceMap()
// retrieve logger from context
@@ -284,8 +285,8 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, toolset tools.T
attribute.String("gen_ai.operation.name", "execute_tool"),
)
// Verify tool belongs to the current toolset before resolving globally.
if !toolset.ContainsTool(toolName) {
// Verify tool belongs to the current group before resolving globally.
if !g.ContainsTool(toolName) {
err = fmt.Errorf("invalid tool name: tool with name %q does not exist", toolName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
@@ -535,7 +536,7 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, toolset tools.T
}
// promptsListHandler handles the "prompts/list" method.
func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, promptset prompts.Promptset, body []byte, header http.Header) (any, error) {
func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, g group.Group, body []byte, header http.Header) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
@@ -558,7 +559,7 @@ func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr
}
promptsMap := primitiveMgr.GetPromptsMap()
listPromptsResult, err := GenerateListPromptsResult(promptset, promptsMap)
listPromptsResult, err := GenerateListPromptsResult(g, promptsMap)
if err != nil {
err = fmt.Errorf("error generating manifest: %w", err)
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
@@ -577,7 +578,7 @@ func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr
}
// promptsGetHandler handles the "prompts/get" method.
func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prompts.Promptset, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, g group.Group, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
@@ -607,8 +608,14 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prom
span.SetName(fmt.Sprintf("%s %s", PROMPTS_GET, promptName))
span.SetAttributes(attribute.String("gen_ai.prompt.name", promptName))
// Verify prompt belongs to the current promptset before resolving globally.
if !promptset.ContainsPrompt(promptName) {
// Populate gen_ai attributes for operation duration metric
if genAIAttrs := util.GenAIMetricAttrsFromContext(ctx); genAIAttrs != nil {
genAIAttrs.OperationName = "get_prompt"
genAIAttrs.PromptName = promptName
}
// Verify prompt belongs to the current group before resolving globally.
if !g.ContainsPrompt(promptName) {
err := fmt.Errorf("prompt with name %q does not exist", promptName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
@@ -619,12 +626,6 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prom
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
// Populate gen_ai attributes for operation duration metric
if genAIAttrs := util.GenAIMetricAttrsFromContext(ctx); genAIAttrs != nil {
genAIAttrs.OperationName = "get_prompt"
genAIAttrs.PromptName = promptName
}
// Parse the arguments provided in the request.
argValues, err := prompt.ParseArgs(req.Params.Arguments, nil)
if err != nil {
@@ -679,3 +680,93 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, promptset prom
Result: result,
}, nil
}
// groupsListHandler handles the "groups/list" method. It returns every named
// group's name and description. The default nameless group is omitted.
func groupsListHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
logger.DebugContext(ctx, "handling groups/list request")
var req ListGroupsRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp groups list request: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
validateHeaderErr, err := validateHeader(id, header, GROUPS_LIST, "")
if err != nil {
return validateHeaderErr, err
}
validateErr, err := validateMetadata(id, req.Params.RequestParams, header == nil)
if err != nil {
return validateErr, err
}
result := GenerateListGroupsResult(primitiveMgr.GetGroupsMap())
logger.DebugContext(ctx, fmt.Sprintf("returning %d groups", len(result.Groups)))
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: result,
}, nil
}
// groupsGetHandler handles the "groups/get" method. It returns the named group's
// tools and prompts.
func groupsGetHandler(ctx context.Context, id jsonrpc.RequestId, primitiveMgr *primitives.PrimitiveManager, body []byte, header http.Header) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
logger.DebugContext(ctx, "handling groups/get request")
var req GetGroupRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp groups/get request: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
validateHeaderErr, err := validateHeader(id, header, GROUPS_GET, req.Params.Name)
if err != nil {
return validateHeaderErr, err
}
validateErr, err := validateMetadata(id, req.Params.RequestParams, header == nil)
if err != nil {
return validateErr, err
}
groupName := req.Params.Name
logger.DebugContext(ctx, fmt.Sprintf("group name: %s", groupName))
// Update span name and set gen_ai attributes
span := trace.SpanFromContext(ctx)
span.SetName(fmt.Sprintf("%s %s", GROUPS_GET, groupName))
span.SetAttributes(attribute.String("gen_ai.group.name", groupName))
// Populate gen_ai attributes for operation duration metric
if genAIAttrs := util.GenAIMetricAttrsFromContext(ctx); genAIAttrs != nil {
genAIAttrs.OperationName = "get_group"
genAIAttrs.GroupName = groupName
}
g, ok := primitiveMgr.GetGroup(groupName)
if !ok {
err := fmt.Errorf("invalid group name: group with name %q does not exist", groupName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
urlParams, _ := util.UrlParamsFromContext(ctx)
result, err := GenerateGetGroupResult(primitiveMgr.GetSourcesMap(), g, primitiveMgr.GetToolsMap(), primitiveMgr.GetPromptsMap(), urlParams)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: result,
}, nil
}
internal/server/mcp/vdraft/method_test.go
+231 -18
View File
@@ -23,11 +23,11 @@ import (
"strings"
"testing"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/log"
"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
"github.com/googleapis/mcp-toolbox/internal/server/primitives"
"github.com/googleapis/mcp-toolbox/internal/testutils"
"github.com/googleapis/mcp-toolbox/internal/tools"
"github.com/googleapis/mcp-toolbox/internal/util"
)
@@ -37,6 +37,16 @@ var (
fakeVersionString = "0.0.0"
)
// mustGroup fetches the default group from the resource manager.
func mustGroup(t *testing.T, rm *primitives.PrimitiveManager) group.Group {
t.Helper()
g, ok := rm.GetGroup("")
if !ok {
t.Fatal("default group not found")
}
return g
}
func TestValidateMetadata(t *testing.T) {
var dummyId jsonrpc.RequestId
clientCapabilities := &ClientCapabilities{}
@@ -406,15 +416,15 @@ func TestToolsListHandler(t *testing.T) {
ctx = util.WithToolboxVersionKey(ctx, "v0.0.0")
// Initialize tools using provided testutils mock instances
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
body ListToolsRequest
rawBody []byte
header http.Header
toolset tools.Toolset
g group.Group
wantErr bool
errContains string
}{
@@ -422,7 +432,7 @@ func TestToolsListHandler(t *testing.T) {
name: "invalid json body",
rawBody: []byte(`{invalid json}`),
header: nil,
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: true,
errContains: "invalid mcp tools list request",
},
@@ -448,7 +458,7 @@ func TestToolsListHandler(t *testing.T) {
},
},
header: http.Header{"Mcp-Method": []string{"WRONG_METHOD"}},
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: true,
errContains: "does not match body value",
},
@@ -474,7 +484,7 @@ func TestToolsListHandler(t *testing.T) {
},
},
header: nil,
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: false,
},
{
@@ -499,7 +509,7 @@ func TestToolsListHandler(t *testing.T) {
},
},
header: http.Header{"Mcp-Method": []string{TOOLS_LIST}},
toolset: toolsets[""],
g: mustGroup(t, primitiveMgr),
wantErr: false,
},
}
@@ -514,7 +524,7 @@ func TestToolsListHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := toolsListHandler(ctx, dummyID, primitiveMgr, tt.toolset, body, tt.header)
got, err := toolsListHandler(ctx, dummyID, primitiveMgr, tt.g, body, tt.header)
if tt.wantErr {
if err == nil {
@@ -550,8 +560,8 @@ func TestToolsCallHandler(t *testing.T) {
testutils.MockTool4,
testutils.MockTool5,
}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
@@ -681,7 +691,7 @@ func TestToolsCallHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := toolsCallHandler(tt.context, dummyID, toolsets[""], primitiveMgr, body, tt.header)
got, err := toolsCallHandler(tt.context, dummyID, mustGroup(t, primitiveMgr), primitiveMgr, body, tt.header)
if tt.wantErr {
if err == nil {
@@ -713,8 +723,8 @@ func TestPromptsListHandler(t *testing.T) {
ctx = util.WithLogger(ctx, testLogger)
// Initialize prompts
mockPrompts := []testutils.MockPrompt{testutils.MockPrompt1, testutils.MockPrompt2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
body ListPromptsRequest
@@ -766,7 +776,7 @@ func TestPromptsListHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := promptsListHandler(ctx, dummyID, primitiveMgr, promptsets[""], body, tt.header)
got, err := promptsListHandler(ctx, dummyID, primitiveMgr, mustGroup(t, primitiveMgr), body, tt.header)
if tt.wantErr {
if err == nil {
@@ -798,8 +808,8 @@ func TestPromptsGetHandler(t *testing.T) {
ctx = util.WithLogger(ctx, testLogger)
// Initialize prompts
mockPrompts := []testutils.MockPrompt{testutils.MockPrompt1, testutils.MockPrompt2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, nil, mockPrompts)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
tests := []struct {
name string
body GetPromptRequest
@@ -900,7 +910,7 @@ func TestPromptsGetHandler(t *testing.T) {
t.Fatalf("unexpected error during marshaling")
}
}
got, err := promptsGetHandler(ctx, dummyID, promptsets[""], primitiveMgr, body, tt.header)
got, err := promptsGetHandler(ctx, dummyID, mustGroup(t, primitiveMgr), primitiveMgr, body, tt.header)
if tt.wantErr {
if err == nil {
@@ -921,6 +931,209 @@ func TestPromptsGetHandler(t *testing.T) {
}
}
func TestGroupsListHandler(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info")
if err != nil {
t.Fatalf("unable to initialize logger: %s", err)
}
ctx = util.WithLogger(ctx, testLogger)
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
validMeta := &RequestMetaObject{
ProtocolVersion: PROTOCOL_VERSION,
ClientInfo: Implementation{
BaseMetadata: BaseMetadata{Name: "TestClient"},
Version: "1.0",
},
MetaClientCapabilities: &ClientCapabilities{},
}
tests := []struct {
name string
rawBody []byte
body ListGroupsRequest
header http.Header
wantErr bool
errContains string
wantNames []string
}{
{
name: "invalid json body",
rawBody: []byte(`{invalid json}`),
wantErr: true,
errContains: "invalid mcp groups list request",
},
{
name: "success excludes default group and sorts",
body: ListGroupsRequest{
PaginatedRequest: PaginatedRequest{
Request: jsonrpc.Request{Method: GROUPS_LIST},
Params: PaginatedRequestParams{
RequestParams: RequestParams{Meta: validMeta},
},
},
},
header: http.Header{"Mcp-Method": []string{GROUPS_LIST}},
wantErr: false,
wantNames: []string{"tool1_only", "tool2_only"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
body := tt.rawBody
var err error
if body == nil {
body, err = json.Marshal(tt.body)
if err != nil {
t.Fatalf("unexpected error during marshaling: %v", err)
}
}
got, err := groupsListHandler(ctx, dummyID, primitiveMgr, body, tt.header)
if tt.wantErr {
if err == nil {
t.Fatalf("expected error, got nil")
}
if tt.errContains != "" && !strings.Contains(err.Error(), tt.errContains) {
t.Errorf("error = %v, want string containing %q", err, tt.errContains)
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
res, ok := got.(jsonrpc.JSONRPCResponse)
if !ok {
t.Fatalf("expected jsonrpc.JSONRPCResponse, got %T", got)
}
result, ok := res.Result.(ListGroupsResult)
if !ok {
t.Fatalf("expected ListGroupsResult, got %T", res.Result)
}
gotNames := make([]string, 0, len(result.Groups))
for _, g := range result.Groups {
gotNames = append(gotNames, g.Name)
}
if len(gotNames) != len(tt.wantNames) {
t.Fatalf("got groups %v, want %v", gotNames, tt.wantNames)
}
for i, n := range tt.wantNames {
if gotNames[i] != n {
t.Errorf("group[%d] = %q, want %q", i, gotNames[i], n)
}
}
})
}
}
func TestGroupsGetHandler(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
testLogger, err := log.NewStdLogger(os.Stdout, os.Stderr, "info")
if err != nil {
t.Fatalf("unable to initialize logger: %s", err)
}
ctx = util.WithLogger(ctx, testLogger)
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2}
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, nil)
primitiveMgr := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
validMeta := &RequestMetaObject{
ProtocolVersion: PROTOCOL_VERSION,
ClientInfo: Implementation{
BaseMetadata: BaseMetadata{Name: "TestClient"},
Version: "1.0",
},
MetaClientCapabilities: &ClientCapabilities{},
}
tests := []struct {
name string
rawBody []byte
body GetGroupRequest
header http.Header
wantErr bool
errContains string
wantName string
}{
{
name: "invalid json body",
rawBody: []byte(`{invalid json}`),
wantErr: true,
errContains: "invalid mcp groups/get request",
},
{
name: "group does not exist",
body: GetGroupRequest{
Request: jsonrpc.Request{Method: GROUPS_GET},
Params: GetGroupRequestParams{
RequestParams: RequestParams{Meta: validMeta},
Name: "missing_group",
},
},
header: http.Header{"Mcp-Method": []string{GROUPS_GET}, "Mcp-Name": []string{"missing_group"}},
wantErr: true,
errContains: `group with name "missing_group" does not exist`,
},
{
name: "success",
body: GetGroupRequest{
Request: jsonrpc.Request{Method: GROUPS_GET},
Params: GetGroupRequestParams{
RequestParams: RequestParams{Meta: validMeta},
Name: "tool1_only",
},
},
header: http.Header{"Mcp-Method": []string{GROUPS_GET}, "Mcp-Name": []string{"tool1_only"}},
wantErr: false,
wantName: "tool1_only",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
body := tt.rawBody
var err error
if body == nil {
body, err = json.Marshal(tt.body)
if err != nil {
t.Fatalf("unexpected error during marshaling: %v", err)
}
}
got, err := groupsGetHandler(ctx, dummyID, primitiveMgr, body, tt.header)
if tt.wantErr {
if err == nil {
t.Fatalf("expected error, got nil")
}
if tt.errContains != "" && !strings.Contains(err.Error(), tt.errContains) {
t.Errorf("error = %v, want string containing %q", err, tt.errContains)
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
res, ok := got.(jsonrpc.JSONRPCResponse)
if !ok {
t.Fatalf("expected jsonrpc.JSONRPCResponse, got %T", got)
}
result, ok := res.Result.(GetGroupResult)
if !ok {
t.Fatalf("expected GetGroupResult, got %T", res.Result)
}
if result.Name != tt.wantName {
t.Errorf("result.Name = %q, want %q", result.Name, tt.wantName)
}
})
}
}
func TestGetResultMetadata(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
internal/server/mcp/vdraft/types.go
+44
View File
@@ -33,6 +33,8 @@ const (
TOOLS_CALL = "tools/call"
PROMPTS_LIST = "prompts/list"
PROMPTS_GET = "prompts/get"
GROUPS_LIST = "groups/list"
GROUPS_GET = "groups/get"
)
/* Request Params */
@@ -495,3 +497,45 @@ type PromptMessage struct {
Role string `json:"role"`
Content TextContent `json:"content"`
}
/* Groups */
// ListGroupsRequest is sent from the client to request the list of groups the
// server has.
type ListGroupsRequest struct {
PaginatedRequest
}
// Group is a single entry in a groups/list response.
type Group struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
}
// ListGroupsResult is the server's response to a groups/list request.
type ListGroupsResult struct {
jsonrpc.Result
Groups []Group `json:"groups"`
}
// GetGroupRequest is sent from the client to request a single group's contents.
type GetGroupRequest struct {
jsonrpc.Request
Params GetGroupRequestParams `json:"params"`
}
// GetGroupRequestParams contains the parameters for a groups/get request.
type GetGroupRequestParams struct {
RequestParams
Name string `json:"name"`
}
// GetGroupResult is the server's response to a groups/get request: the group's
// tools and prompts. The description is intentionally omitted; it is exposed only
// through groups/list.
type GetGroupResult struct {
jsonrpc.Result
Name string `json:"name"`
Tools []Tool `json:"tools"`
Prompts []Prompt `json:"prompts"`
}
internal/server/mcp_test.go
+84 -15
View File
@@ -28,10 +28,13 @@ import (
"strings"
"testing"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/log"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
"github.com/googleapis/mcp-toolbox/internal/server/primitives"
"github.com/googleapis/mcp-toolbox/internal/telemetry"
"github.com/googleapis/mcp-toolbox/internal/tools"
"github.com/googleapis/mcp-toolbox/internal/util"
"go.opentelemetry.io/otel"
"go.opentelemetry.io/otel/propagation"
@@ -116,8 +119,8 @@ var prompt2Args = []any{
func TestMcpEndpointWithoutInitialized(t *testing.T) {
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2, testutils.MockTool3, testutils.MockTool4, testutils.MockTool5}
mockPrompts := []testutils.MockPrompt{testutils.MockPrompt1, testutils.MockPrompt2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, mockTools, mockPrompts)
r, shutdown := setUpServer(t, "mcp", toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, mockPrompts)
r, shutdown := setUpServer(t, "mcp", toolsMap, promptsMap, groups)
defer shutdown()
ts := runServer(r, false)
defer ts.Close()
@@ -461,8 +464,8 @@ func runInitializeLifecycle(t *testing.T, ts *httptest.Server, protocolVersion s
func TestMcpEndpoint(t *testing.T) {
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2, testutils.MockTool3, testutils.MockTool4, testutils.MockTool5, testutils.MockToolUrlBinding}
mockPrompts := []testutils.MockPrompt{testutils.MockPrompt1, testutils.MockPrompt2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, mockTools, mockPrompts)
r, shutdown := setUpServer(t, "mcp", toolsMap, toolsets, promptsMap, promptsets, withEnableDraftSpecs())
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, mockPrompts)
r, shutdown := setUpServer(t, "mcp", toolsMap, promptsMap, groups, withEnableDraftSpecs())
defer shutdown()
ts := runServer(r, false)
defer ts.Close()
@@ -1387,8 +1390,8 @@ func TestMcpEndpoint(t *testing.T) {
func TestMcpEndpointWithoutEnablingDraftSpecs(t *testing.T) {
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2, testutils.MockTool3, testutils.MockTool4, testutils.MockTool5}
mockPrompts := []testutils.MockPrompt{testutils.MockPrompt1, testutils.MockPrompt2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, mockTools, mockPrompts)
r, shutdown := setUpServer(t, "mcp", toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, mockPrompts)
r, shutdown := setUpServer(t, "mcp", toolsMap, promptsMap, groups)
defer shutdown()
ts := runServer(r, false)
defer ts.Close()
@@ -1471,8 +1474,8 @@ func TestMcpEndpointWithoutEnablingDraftSpecs(t *testing.T) {
func TestInvalidProtocolVersionHeader(t *testing.T) {
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2, testutils.MockTool3, testutils.MockTool4, testutils.MockTool5}
mockPrompts := []testutils.MockPrompt{testutils.MockPrompt1}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, mockTools, mockPrompts)
r, shutdown := setUpServer(t, "mcp", toolsMap, toolsets, promptsMap, promptsets)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, mockPrompts)
r, shutdown := setUpServer(t, "mcp", toolsMap, promptsMap, groups)
defer shutdown()
ts := runServer(r, false)
defer ts.Close()
@@ -1517,7 +1520,7 @@ func TestInvalidProtocolVersionHeader(t *testing.T) {
}
func TestDeleteEndpoint(t *testing.T) {
r, shutdown := setUpServer(t, "mcp", nil, nil, nil, nil)
r, shutdown := setUpServer(t, "mcp", nil, nil, nil)
defer shutdown()
ts := runServer(r, false)
defer ts.Close()
@@ -1532,7 +1535,7 @@ func TestDeleteEndpoint(t *testing.T) {
}
func TestGetEndpoint(t *testing.T) {
r, shutdown := setUpServer(t, "mcp", nil, nil, nil, nil)
r, shutdown := setUpServer(t, "mcp", nil, nil, nil)
defer shutdown()
ts := runServer(r, false)
defer ts.Close()
@@ -1555,7 +1558,7 @@ func TestGetEndpoint(t *testing.T) {
}
func TestMcpRequestBodyLimit(t *testing.T) {
r, shutdown := setUpServer(t, "mcp", nil, nil, nil, nil)
r, shutdown := setUpServer(t, "mcp", nil, nil, nil)
defer shutdown()
ts := runServer(r, false)
defer ts.Close()
@@ -1587,7 +1590,7 @@ func TestMcpRequestBodyLimit(t *testing.T) {
func TestMcpRequestBodyLimitOverride(t *testing.T) {
customLimit := int64(1 << 20)
r, shutdown := setUpServer(t, "mcp", nil, nil, nil, nil, withHTTPMaxRequestBytes(customLimit))
r, shutdown := setUpServer(t, "mcp", nil, nil, nil, withHTTPMaxRequestBytes(customLimit))
defer shutdown()
ts := runServer(r, false)
defer ts.Close()
@@ -1617,7 +1620,7 @@ func TestMcpRequestBodyLimitOverride(t *testing.T) {
}
func TestSseEndpoint(t *testing.T) {
r, shutdown := setUpServer(t, "mcp", nil, nil, nil, nil)
r, shutdown := setUpServer(t, "mcp", nil, nil, nil)
defer shutdown()
ts := runServer(r, false)
defer ts.Close()
@@ -1735,7 +1738,7 @@ func TestStdioSession(t *testing.T) {
mockTools := []testutils.MockTool{testutils.MockTool1, testutils.MockTool2, testutils.MockTool3}
mockPrompts := []testutils.MockPrompt{testutils.MockPrompt1, testutils.MockPrompt2}
toolsMap, toolsets, promptsMap, promptsets := testutils.SetUpResources(t, mockTools, mockPrompts)
toolsMap, promptsMap, groups := testutils.SetUpResources(t, mockTools, mockPrompts)
pr, pw, err := os.Pipe()
if err != nil {
@@ -1765,7 +1768,7 @@ func TestStdioSession(t *testing.T) {
sseManager := newSseManager(ctx)
primitiveManager := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, toolsets, promptsMap, promptsets)
primitiveManager := primitives.NewPrimitiveManager(nil, nil, nil, toolsMap, promptsMap, groups)
server := &Server{
version: testutils.MockVersionString,
@@ -1960,3 +1963,69 @@ func TestExtractMeta(t *testing.T) {
})
}
}
// TestMcpPromptScopingByGroup is an end-to-end HTTP test that a `prompts/list`
// request sent to a group's MCP endpoint returns only the prompts belonging to
// that group. It stands up the real server with two groups (each scoped to a
// different prompt) and asserts each route surfaces just its own prompt.
func TestMcpPromptScopingByGroup(t *testing.T) {
toolsMap := map[string]tools.Tool{}
promptsMap := map[string]prompts.Prompt{
testutils.MockPrompt1.Name: testutils.MockPrompt1,
testutils.MockPrompt2.Name: testutils.MockPrompt2,
}
groupA, err := group.GroupConfig{Name: "group_a", PromptNames: []string{testutils.MockPrompt1.Name}}.Initialize(toolsMap, promptsMap)
if err != nil {
t.Fatalf("unable to initialize group_a: %s", err)
}
groupB, err := group.GroupConfig{Name: "group_b", PromptNames: []string{testutils.MockPrompt2.Name}}.Initialize(toolsMap, promptsMap)
if err != nil {
t.Fatalf("unable to initialize group_b: %s", err)
}
groups := map[string]group.Group{"group_a": groupA, "group_b": groupB}
r, shutdown := setUpServer(t, "mcp", toolsMap, promptsMap, groups)
defer shutdown()
ts := runServer(r, false)
defer ts.Close()
testCases := []struct {
name string
url string
wantPrompts []any
}{
{
name: "group_a scopes to its own prompt",
url: "/group_a",
wantPrompts: []any{map[string]any{"name": "prompt1"}},
},
{
name: "group_b scopes to its own prompt",
url: "/group_b",
wantPrompts: []any{map[string]any{"name": "prompt2", "arguments": prompt2Args}},
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
reqBody := jsonrpc.JSONRPCRequest{Jsonrpc: jsonrpcVersion, Id: "prompts-list", Request: jsonrpc.Request{Method: "prompts/list"}}
reqMarshal, err := json.Marshal(reqBody)
if err != nil {
t.Fatalf("unexpected error marshaling body: %s", err)
}
resp, body, err := runRequest(ts, http.MethodPost, tc.url, bytes.NewBuffer(reqMarshal), nil)
if err != nil {
t.Fatalf("unexpected error during request: %s", err)
}
if resp.StatusCode != http.StatusOK {
t.Errorf("StatusCode mismatch: got %d, want %d", resp.StatusCode, http.StatusOK)
}
var got map[string]any
if err := json.Unmarshal(body, &got); err != nil {
t.Fatalf("unexpected error unmarshalling body: %s", err)
}
want := map[string]any{"jsonrpc": "2.0", "id": "prompts-list", "result": map[string]any{"prompts": tc.wantPrompts}}
if !reflect.DeepEqual(got, want) {
t.Fatalf("unexpected response: got %#v, want %#v", got, want)
}
})
}
}
internal/server/primitives/primitives.go
+25 -20
View File
@@ -19,29 +19,32 @@ import (
"github.com/googleapis/mcp-toolbox/internal/auth"
"github.com/googleapis/mcp-toolbox/internal/embeddingmodels"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/sources"
"github.com/googleapis/mcp-toolbox/internal/tools"
)
// PrimitiveManager contains available primitives for the server. Should be initialized with NewPrimitiveManager().
// PrimitiveManager contains available resources for the server. Should be initialized with NewPrimitiveManager().
// groups is the source of truth for named collections; toolset views (manifests)
// are derived from the group on demand by the callers that render them.
type PrimitiveManager struct {
mu sync.RWMutex
sources map[string]sources.Source
authServices map[string]auth.AuthService
embeddingModels map[string]embeddingmodels.EmbeddingModel
tools map[string]tools.Tool
toolsets map[string]tools.Toolset
prompts map[string]prompts.Prompt
promptsets map[string]prompts.Promptset
groups map[string]group.Group
}
func NewPrimitiveManager(
sourcesMap map[string]sources.Source,
authServicesMap map[string]auth.AuthService,
embeddingModelsMap map[string]embeddingmodels.EmbeddingModel,
toolsMap map[string]tools.Tool, toolsetsMap map[string]tools.Toolset,
promptsMap map[string]prompts.Prompt, promptsetsMap map[string]prompts.Promptset,
toolsMap map[string]tools.Tool,
promptsMap map[string]prompts.Prompt,
groupsMap map[string]group.Group,
) *PrimitiveManager {
primitiveMgr := &PrimitiveManager{
@@ -50,9 +53,8 @@ func NewPrimitiveManager(
authServices: authServicesMap,
embeddingModels: embeddingModelsMap,
tools: toolsMap,
toolsets: toolsetsMap,
prompts: promptsMap,
promptsets: promptsetsMap,
groups: groupsMap,
}
return primitiveMgr
@@ -86,13 +88,6 @@ func (r *PrimitiveManager) GetTool(toolName string) (tools.Tool, bool) {
return tool, ok
}
func (r *PrimitiveManager) GetToolset(toolsetName string) (tools.Toolset, bool) {
r.mu.RLock()
defer r.mu.RUnlock()
toolset, ok := r.toolsets[toolsetName]
return toolset, ok
}
func (r *PrimitiveManager) GetPrompt(promptName string) (prompts.Prompt, bool) {
r.mu.RLock()
defer r.mu.RUnlock()
@@ -100,23 +95,23 @@ func (r *PrimitiveManager) GetPrompt(promptName string) (prompts.Prompt, bool) {
return prompt, ok
}
func (r *PrimitiveManager) GetPromptset(promptsetName string) (prompts.Promptset, bool) {
// GetGroup returns the group of the given name.
func (r *PrimitiveManager) GetGroup(groupName string) (group.Group, bool) {
r.mu.RLock()
defer r.mu.RUnlock()
promptset, ok := r.promptsets[promptsetName]
return promptset, ok
g, ok := r.groups[groupName]
return g, ok
}
func (r *PrimitiveManager) SetPrimitives(sourcesMap map[string]sources.Source, authServicesMap map[string]auth.AuthService, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel, toolsMap map[string]tools.Tool, toolsetsMap map[string]tools.Toolset, promptsMap map[string]prompts.Prompt, promptsetsMap map[string]prompts.Promptset) {
func (r *PrimitiveManager) SetPrimitives(sourcesMap map[string]sources.Source, authServicesMap map[string]auth.AuthService, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel, toolsMap map[string]tools.Tool, promptsMap map[string]prompts.Prompt, groupsMap map[string]group.Group) {
r.mu.Lock()
defer r.mu.Unlock()
r.sources = sourcesMap
r.authServices = authServicesMap
r.embeddingModels = embeddingModelsMap
r.tools = toolsMap
r.toolsets = toolsetsMap
r.prompts = promptsMap
r.promptsets = promptsetsMap
r.groups = groupsMap
}
func (r *PrimitiveManager) GetSourcesMap() map[string]sources.Source {
@@ -168,3 +163,13 @@ func (r *PrimitiveManager) GetPromptsMap() map[string]prompts.Prompt {
}
return copiedMap
}
func (r *PrimitiveManager) GetGroupsMap() map[string]group.Group {
r.mu.RLock()
defer r.mu.RUnlock()
copiedMap := make(map[string]group.Group, len(r.groups))
for k, v := range r.groups {
copiedMap[k] = v
}
return copiedMap
}
internal/server/primitives/primitives_test.go
+20 -32
View File
@@ -20,10 +20,12 @@ import (
"github.com/google/go-cmp/cmp"
"github.com/googleapis/mcp-toolbox/internal/auth"
"github.com/googleapis/mcp-toolbox/internal/embeddingmodels"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/server/primitives"
"github.com/googleapis/mcp-toolbox/internal/sources"
"github.com/googleapis/mcp-toolbox/internal/sources/alloydbpg"
"github.com/googleapis/mcp-toolbox/internal/testutils"
"github.com/googleapis/mcp-toolbox/internal/tools"
)
@@ -39,55 +41,41 @@ func TestUpdateServer(t *testing.T) {
newAuth := map[string]auth.AuthService{"example-auth": nil}
newEmbeddingModels := map[string]embeddingmodels.EmbeddingModel{"example-model": nil}
newTools := map[string]tools.Tool{"example-tool": nil}
newToolsets := map[string]tools.Toolset{
"example-toolset": {
ToolsetConfig: tools.ToolsetConfig{
Name: "example-toolset",
},
Tools: []*tools.Tool{},
},
newPrompts := map[string]prompts.Prompt{"example-prompt": testutils.NewMockPrompt("example-prompt", "", prompts.Arguments{})}
newGroups := map[string]group.Group{
"example-toolset": group.NewGroup(group.GroupConfig{Name: "example-toolset", ToolNames: []string{"example-tool"}}),
}
newPrompts := map[string]prompts.Prompt{"example-prompt": nil}
newPromptsets := map[string]prompts.Promptset{
"example-promptset": {
PromptsetConfig: prompts.PromptsetConfig{
Name: "example-promptset",
},
Prompts: []*prompts.Prompt{},
},
}
primMgr := primitives.NewPrimitiveManager(newSources, newAuth, newEmbeddingModels, newTools, newToolsets, newPrompts, newPromptsets)
resMgr := primitives.NewPrimitiveManager(newSources, newAuth, newEmbeddingModels, newTools, newPrompts, newGroups)
gotSource, _ := primMgr.GetSource("example-source")
gotSource, _ := resMgr.GetSource("example-source")
if diff := cmp.Diff(gotSource, newSources["example-source"]); diff != "" {
t.Errorf("error updating server, sources (-want +got):\n%s", diff)
}
gotAuthService, _ := primMgr.GetAuthService("example-auth")
gotAuthService, _ := resMgr.GetAuthService("example-auth")
if diff := cmp.Diff(gotAuthService, newAuth["example-auth"]); diff != "" {
t.Errorf("error updating server, authServices (-want +got):\n%s", diff)
}
gotTool, _ := primMgr.GetTool("example-tool")
gotTool, _ := resMgr.GetTool("example-tool")
if diff := cmp.Diff(gotTool, newTools["example-tool"]); diff != "" {
t.Errorf("error updating server, tools (-want +got):\n%s", diff)
}
gotToolset, _ := primMgr.GetToolset("example-toolset")
if diff := cmp.Diff(gotToolset, newToolsets["example-toolset"], cmp.AllowUnexported(tools.Toolset{})); diff != "" {
t.Errorf("error updating server, toolset (-want +got):\n%s", diff)
wantGroup := newGroups["example-toolset"]
gotGroup, ok := resMgr.GetGroup("example-toolset")
if !ok {
t.Fatal("expected group \"example-toolset\" to exist")
}
if diff := cmp.Diff(wantGroup, gotGroup, cmp.AllowUnexported(group.Group{})); diff != "" {
t.Errorf("error updating server, group (-want +got):\n%s", diff)
}
gotPrompt, _ := primMgr.GetPrompt("example-prompt")
if diff := cmp.Diff(gotPrompt, newPrompts["example-prompt"]); diff != "" {
gotPrompt, _ := resMgr.GetPrompt("example-prompt")
if diff := cmp.Diff(gotPrompt, newPrompts["example-prompt"], cmp.AllowUnexported(testutils.MockPrompt{})); diff != "" {
t.Errorf("error updating server, prompts (-want +got):\n%s", diff)
}
gotPromptset, _ := primMgr.GetPromptset("example-promptset")
if diff := cmp.Diff(gotPromptset, newPromptsets["example-promptset"], cmp.AllowUnexported(prompts.Promptset{})); diff != "" {
t.Errorf("error updating server, promptset (-want +got):\n%s", diff)
}
updateSource := map[string]sources.Source{
"example-source2": &alloydbpg.Source{
Config: alloydbpg.Config{
@@ -97,8 +85,8 @@ func TestUpdateServer(t *testing.T) {
},
}
primMgr.SetPrimitives(updateSource, newAuth, newEmbeddingModels, newTools, newToolsets, newPrompts, newPromptsets)
gotSource, _ = primMgr.GetSource("example-source2")
resMgr.SetPrimitives(updateSource, newAuth, newEmbeddingModels, newTools, newPrompts, newGroups)
gotSource, _ = resMgr.GetSource("example-source2")
if diff := cmp.Diff(gotSource, updateSource["example-source2"]); diff != "" {
t.Errorf("error updating server, sources (-want +got):\n%s", diff)
}
internal/server/server.go
+78 -91
View File
@@ -36,6 +36,7 @@ import (
"github.com/go-chi/render"
"github.com/googleapis/mcp-toolbox/internal/auth"
"github.com/googleapis/mcp-toolbox/internal/embeddingmodels"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/log"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/server/mcp/jsonrpc"
@@ -70,15 +71,14 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
map[string]auth.AuthService,
map[string]embeddingmodels.EmbeddingModel,
map[string]tools.Tool,
map[string]tools.Toolset,
map[string]prompts.Prompt,
map[string]prompts.Promptset,
map[string]group.Group,
error,
) {
if cfg.EnableAPI {
for _, sc := range cfg.AuthServiceConfigs {
if sc.IsMCPEnabled() {
return nil, nil, nil, nil, nil, nil, nil, fmt.Errorf("MCP Auth cannot be enabled together with the legacy HTTP API (EnableAPI)")
return nil, nil, nil, nil, nil, nil, fmt.Errorf("MCP Auth cannot be enabled together with the legacy HTTP API (EnableAPI)")
}
}
}
@@ -90,12 +90,12 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
ctx = util.WithUserAgent(ctx, metadataStr)
instrumentation, err := util.InstrumentationFromContext(ctx)
if err != nil {
return nil, nil, nil, nil, nil, nil, nil, fmt.Errorf("failed to get instrumentation from context: %w", err)
return nil, nil, nil, nil, nil, nil, fmt.Errorf("failed to get instrumentation from context: %w", err)
}
l, err := util.LoggerFromContext(ctx)
if err != nil {
return nil, nil, nil, nil, nil, nil, nil, fmt.Errorf("failed to get logger from context: %w", err)
return nil, nil, nil, nil, nil, nil, fmt.Errorf("failed to get logger from context: %w", err)
}
// initialize and validate the sources from configs
@@ -116,7 +116,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
return s, nil
}()
if err != nil {
return nil, nil, nil, nil, nil, nil, nil, err
return nil, nil, nil, nil, nil, nil, err
}
sourcesMap[name] = s
}
@@ -144,7 +144,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
return a, nil
}()
if err != nil {
return nil, nil, nil, nil, nil, nil, nil, err
return nil, nil, nil, nil, nil, nil, err
}
authServicesMap[name] = a
}
@@ -173,7 +173,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
return em, nil
}()
if err != nil {
return nil, nil, nil, nil, nil, nil, nil, err
return nil, nil, nil, nil, nil, nil, err
}
embeddingModelsMap[name] = em
}
@@ -185,12 +185,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
toolsMap, err := initializeTools(ctx, cfg, instrumentation, l)
if err != nil {
return nil, nil, nil, nil, nil, nil, nil, err
}
toolsetsMap, err := initializeToolsets(ctx, cfg, toolsMap, instrumentation, l)
if err != nil {
return nil, nil, nil, nil, nil, nil, nil, err
return nil, nil, nil, nil, nil, nil, err
}
// initialize and validate the prompts from configs
@@ -211,7 +206,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
return p, nil
}()
if err != nil {
return nil, nil, nil, nil, nil, nil, nil, err
return nil, nil, nil, nil, nil, nil, err
}
promptsMap[name] = p
}
@@ -221,56 +216,20 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
}
l.InfoContext(ctx, fmt.Sprintf("Initialized %d prompts: %s", len(promptsMap), strings.Join(promptNames, ", ")))
// create a default promptset that contains all prompts
allPromptNames := make([]string, 0, len(promptsMap))
for name := range promptsMap {
allPromptNames = append(allPromptNames, name)
groupsMap, err := initializeGroups(ctx, cfg, toolsMap, promptsMap, instrumentation, l)
if err != nil {
return nil, nil, nil, nil, nil, nil, err
}
if cfg.PromptsetConfigs == nil {
cfg.PromptsetConfigs = make(PromptsetConfigs)
}
cfg.PromptsetConfigs[""] = prompts.PromptsetConfig{Name: "", PromptNames: allPromptNames}
// initialize and validate the promptsets from configs
promptsetsMap := make(map[string]prompts.Promptset)
for name, pc := range cfg.PromptsetConfigs {
p, err := func() (prompts.Promptset, error) {
_, span := instrumentation.Tracer.Start(
ctx,
"toolbox/server/prompset/init",
trace.WithAttributes(attribute.String("prompset_name", name)),
)
defer span.End()
p, err := pc.Initialize(cfg.Version, promptsMap)
if err != nil {
return prompts.Promptset{}, fmt.Errorf("unable to initialize promptset %q: %w", name, err)
}
return p, err
}()
if err != nil {
return nil, nil, nil, nil, nil, nil, nil, err
}
promptsetsMap[name] = p
}
promptsetNames := make([]string, 0, len(promptsetsMap))
for name := range promptsetsMap {
if name == "" {
promptsetNames = append(promptsetNames, "default")
} else {
promptsetNames = append(promptsetNames, name)
}
}
l.InfoContext(ctx, fmt.Sprintf("Initialized %d promptsets: %s", len(promptsetsMap), strings.Join(promptsetNames, ", ")))
return sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, nil
return sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, promptsMap, groupsMap, nil
}
// InitializeOfflineConfigs initializes only tools and toolsets from the config,
// skipping sources, auth services, and embedding models. It backs flows like
// skills-generate that need tool metadata without live source connections.
// InitializeOfflineConfigs initializes only tools, prompts, and groups from the
// config, skipping sources, auth services, and embedding models. It backs flows
// like skills-generate that need tool metadata without live source connections.
func InitializeOfflineConfigs(ctx context.Context, cfg ServerConfig) (
map[string]tools.Tool,
map[string]tools.Toolset,
map[string]group.Group,
error,
) {
instrumentation, err := util.InstrumentationFromContext(ctx)
@@ -288,12 +247,22 @@ func InitializeOfflineConfigs(ctx context.Context, cfg ServerConfig) (
return nil, nil, err
}
toolsetsMap, err := initializeToolsets(ctx, cfg, toolsMap, instrumentation, l)
// Prompts are initialized so group prompt validation succeeds offline.
promptsMap := make(map[string]prompts.Prompt)
for name, pc := range cfg.PromptConfigs {
p, err := pc.Initialize()
if err != nil {
return nil, nil, fmt.Errorf("unable to initialize prompt %q: %w", name, err)
}
promptsMap[name] = p
}
groupsMap, err := initializeGroups(ctx, cfg, toolsMap, promptsMap, instrumentation, l)
if err != nil {
return nil, nil, err
}
return toolsMap, toolsetsMap, nil
return toolsMap, groupsMap, nil
}
// initializeTools initializes and validates the tools from the config.
@@ -327,64 +296,82 @@ func initializeTools(ctx context.Context, cfg ServerConfig, instrumentation *tel
return toolsMap, nil
}
// initializeToolsets seeds a default toolset containing all tools, then
// initializes and validates the toolsets from the config.
func initializeToolsets(ctx context.Context, cfg ServerConfig, toolsMap map[string]tools.Tool, instrumentation *telemetry.Instrumentation, l log.Logger) (map[string]tools.Toolset, error) {
// create a default toolset that contains all tools
// initializeGroups seeds a default nameless group containing all tools and all
// prompts, converts each legacy kind: toolsets config into a tools-only group,
// then initializes and validates every group. The default group's derived
// toolset/promptset views preserve the legacy behavior of returning everything
// for clients that connect without naming a collection.
func initializeGroups(ctx context.Context, cfg ServerConfig, toolsMap map[string]tools.Tool, promptsMap map[string]prompts.Prompt, instrumentation *telemetry.Instrumentation, l log.Logger) (map[string]group.Group, error) {
allToolNames := make([]string, 0, len(toolsMap))
for name := range toolsMap {
allToolNames = append(allToolNames, name)
}
slices.Sort(allToolNames)
if cfg.ToolsetConfigs == nil {
cfg.ToolsetConfigs = make(ToolsetConfigs)
allPromptNames := make([]string, 0, len(promptsMap))
for name := range promptsMap {
allPromptNames = append(allPromptNames, name)
}
cfg.ToolsetConfigs[""] = tools.ToolsetConfig{Name: "", ToolNames: allToolNames}
slices.Sort(allPromptNames)
toolsetsMap := make(map[string]tools.Toolset)
for name, tc := range cfg.ToolsetConfigs {
// Legacy `kind: toolset` configs are already folded into cfg.GroupConfigs at
// unmarshal. Copy them over, then seed the default nameless group with all tools
// and all prompts.
groupConfigs := make(map[string]group.GroupConfig, len(cfg.GroupConfigs)+1)
var defaultDescription string
for name, gc := range cfg.GroupConfigs {
if name == "" {
// The default group's tools and prompts are fixed; only its
// description carries over as the default server instruction.
defaultDescription = gc.Description
continue
}
groupConfigs[name] = gc
}
groupConfigs[""] = group.GroupConfig{Name: "", Description: defaultDescription, ToolNames: allToolNames, PromptNames: allPromptNames}
groupsMap := make(map[string]group.Group)
for name, gc := range groupConfigs {
if cfg.IgnoreUnknownTools {
filteredToolNames := make([]string, 0, len(tc.ToolNames))
for _, tn := range tc.ToolNames {
filteredToolNames := make([]string, 0, len(gc.ToolNames))
for _, tn := range gc.ToolNames {
if _, ok := toolsMap[tn]; ok {
filteredToolNames = append(filteredToolNames, tn)
} else {
l.WarnContext(ctx, fmt.Sprintf("Skipping missing tool %q in toolset %q", tn, name))
l.WarnContext(ctx, fmt.Sprintf("Skipping missing tool %q in group %q", tn, name))
}
}
tc.ToolNames = filteredToolNames
cfg.ToolsetConfigs[name] = tc
gc.ToolNames = filteredToolNames
}
t, err := func() (tools.Toolset, error) {
g, err := func() (group.Group, error) {
_, span := instrumentation.Tracer.Start(
ctx,
"toolbox/server/toolset/init",
trace.WithAttributes(attribute.String("toolset.name", name)),
"toolbox/server/group/init",
trace.WithAttributes(attribute.String("group.name", name)),
)
defer span.End()
t, err := tc.Initialize(cfg.Version, toolsMap)
g, err := gc.Initialize(toolsMap, promptsMap)
if err != nil {
return tools.Toolset{}, fmt.Errorf("unable to initialize toolset %q: %w", name, err)
return group.Group{}, fmt.Errorf("unable to initialize group %q: %w", name, err)
}
return t, err
return g, nil
}()
if err != nil {
return nil, err
}
toolsetsMap[name] = t
groupsMap[name] = g
}
toolsetNames := make([]string, 0, len(toolsetsMap))
for name := range toolsetsMap {
groupNames := make([]string, 0, len(groupsMap))
for name := range groupsMap {
if name == "" {
toolsetNames = append(toolsetNames, "default")
groupNames = append(groupNames, "default")
} else {
toolsetNames = append(toolsetNames, name)
groupNames = append(groupNames, name)
}
}
l.InfoContext(ctx, fmt.Sprintf("Initialized %d toolsets: %s", len(toolsetsMap), strings.Join(toolsetNames, ", ")))
l.InfoContext(ctx, fmt.Sprintf("Initialized %d groups: %s", len(groupsMap), strings.Join(groupNames, ", ")))
return toolsetsMap, nil
return groupsMap, nil
}
func hostCheck(allowedHosts map[string]struct{}) func(http.Handler) http.Handler {
@@ -441,7 +428,7 @@ func NewServer(ctx context.Context, cfg ServerConfig) (*Server, error) {
logger := l.SlogLogger()
r.Use(httplog.RequestLogger(logger, httpOpts))
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := InitializeConfigs(ctx, cfg)
sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, promptsMap, groupsMap, err := InitializeConfigs(ctx, cfg)
if err != nil {
return nil, fmt.Errorf("unable to initialize configs: %w", err)
}
@@ -451,7 +438,7 @@ func NewServer(ctx context.Context, cfg ServerConfig) (*Server, error) {
sseManager := newSseManager(ctx)
primitiveManager := primitives.NewPrimitiveManager(sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap)
primitiveManager := primitives.NewPrimitiveManager(sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, promptsMap, groupsMap)
limit := cfg.HttpMaxRequestBytes
if limit <= 0 {
@@ -479,7 +466,7 @@ func NewServer(ctx context.Context, cfg ServerConfig) (*Server, error) {
// cors
if slices.Contains(cfg.AllowedOrigins, "*") {
s.logger.WarnContext(ctx, "wildcard (*) allows any website to access the resources. This creates a security risk regardless of whether you are in a production or local development environment. Recommended to use --allowed-origins with specific local addresses.")
s.logger.WarnContext(ctx, "wildcard (*) allows any website to access the primitives. This creates a security risk regardless of whether you are in a production or local development environment. Recommended to use --allowed-origins with specific local addresses.")
}
corsOpts := cors.Options{
AllowedOrigins: cfg.AllowedOrigins,
internal/server/server_test.go
+331 -30
View File
@@ -41,14 +41,18 @@ import (
"github.com/googleapis/mcp-toolbox/internal/auth"
"github.com/googleapis/mcp-toolbox/internal/auth/generic"
"github.com/googleapis/mcp-toolbox/internal/embeddingmodels"
_ "github.com/googleapis/mcp-toolbox/internal/embeddingmodels/gemini"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/log"
"github.com/googleapis/mcp-toolbox/internal/prompts"
_ "github.com/googleapis/mcp-toolbox/internal/prompts/custom"
"github.com/googleapis/mcp-toolbox/internal/server"
"github.com/googleapis/mcp-toolbox/internal/sources"
"github.com/googleapis/mcp-toolbox/internal/sources/alloydbpg"
"github.com/googleapis/mcp-toolbox/internal/telemetry"
"github.com/googleapis/mcp-toolbox/internal/testutils"
"github.com/googleapis/mcp-toolbox/internal/tools"
_ "github.com/googleapis/mcp-toolbox/internal/tools/http"
"github.com/googleapis/mcp-toolbox/internal/util"
"go.opentelemetry.io/otel/trace"
)
@@ -251,24 +255,11 @@ func TestUpdateServer(t *testing.T) {
newAuth := map[string]auth.AuthService{"example-auth": nil}
newEmbeddingModels := map[string]embeddingmodels.EmbeddingModel{"example-model": nil}
newTools := map[string]tools.Tool{"example-tool": nil}
newToolsets := map[string]tools.Toolset{
"example-toolset": {
ToolsetConfig: tools.ToolsetConfig{
Name: "example-toolset",
},
Tools: []*tools.Tool{},
},
newPrompts := map[string]prompts.Prompt{"example-prompt": testutils.NewMockPrompt("example-prompt", "", prompts.Arguments{})}
newGroups := map[string]group.Group{
"example-toolset": group.NewGroup(group.GroupConfig{Name: "example-toolset", ToolNames: []string{"example-tool"}}),
}
newPrompts := map[string]prompts.Prompt{"example-prompt": nil}
newPromptsets := map[string]prompts.Promptset{
"example-promptset": {
PromptsetConfig: prompts.PromptsetConfig{
Name: "example-promptset",
},
Prompts: []*prompts.Prompt{},
},
}
s.PrimitiveMgr.SetPrimitives(newSources, newAuth, newEmbeddingModels, newTools, newToolsets, newPrompts, newPromptsets)
s.PrimitiveMgr.SetPrimitives(newSources, newAuth, newEmbeddingModels, newTools, newPrompts, newGroups)
if err != nil {
t.Errorf("error updating server: %s", err)
}
@@ -288,20 +279,19 @@ func TestUpdateServer(t *testing.T) {
t.Errorf("error updating server, tools (-want +got):\n%s", diff)
}
gotToolset, _ := s.PrimitiveMgr.GetToolset("example-toolset")
if diff := cmp.Diff(gotToolset, newToolsets["example-toolset"], cmp.AllowUnexported(tools.Toolset{})); diff != "" {
t.Errorf("error updating server, toolset (-want +got):\n%s", diff)
wantGroup := newGroups["example-toolset"]
gotGroup, ok := s.PrimitiveMgr.GetGroup("example-toolset")
if !ok {
t.Fatal("expected group \"example-toolset\" to exist")
}
if diff := cmp.Diff(wantGroup, gotGroup, cmp.AllowUnexported(group.Group{})); diff != "" {
t.Errorf("error updating server, group (-want +got):\n%s", diff)
}
gotPrompt, _ := s.PrimitiveMgr.GetPrompt("example-prompt")
if diff := cmp.Diff(gotPrompt, newPrompts["example-prompt"]); diff != "" {
if diff := cmp.Diff(gotPrompt, newPrompts["example-prompt"], cmp.AllowUnexported(testutils.MockPrompt{})); diff != "" {
t.Errorf("error updating server, prompts (-want +got):\n%s", diff)
}
gotPromptset, _ := s.PrimitiveMgr.GetPromptset("example-promptset")
if diff := cmp.Diff(gotPromptset, newPromptsets["example-promptset"], cmp.AllowUnexported(prompts.Promptset{})); diff != "" {
t.Errorf("error updating server, promptset (-want +got):\n%s", diff)
}
}
func TestEndpointSecurityAllowedOrigin(t *testing.T) {
@@ -1337,6 +1327,317 @@ scopesRequired:
}
}
func TestDuplicateResourceConfig(t *testing.T) {
ctx := context.Background()
tests := []struct {
name string
yaml string
}{
{
name: "duplicate source",
yaml: `
kind: source
name: my_source
type: alloydb-postgres
project: my-project
region: us-central1
cluster: my-cluster
instance: my-instance
database: my-db
---
kind: source
name: my_source
type: alloydb-postgres
project: my-project
region: us-central1
cluster: my-cluster
instance: my-instance
database: my-db
`,
},
{
name: "duplicate authService",
yaml: `
kind: authService
name: my_auth
type: generic
audience: my-audience
authorizationServer: https://example.com
---
kind: authService
name: my_auth
type: generic
audience: my-audience
authorizationServer: https://example.com
`,
},
{
name: "duplicate tool",
yaml: `
kind: tool
name: my_tool
type: http
source: my_source
path: /a
method: GET
---
kind: tool
name: my_tool
type: http
source: my_source
path: /b
method: GET
`,
},
{
name: "duplicate toolset",
yaml: `
kind: toolset
name: my_toolset
tools:
- tool_a
---
kind: toolset
name: my_toolset
tools:
- tool_b
`,
},
{
name: "duplicate embeddingModel",
yaml: `
kind: embeddingModel
name: my_model
type: gemini
model: text-embedding-005
---
kind: embeddingModel
name: my_model
type: gemini
model: text-embedding-005
`,
},
{
name: "duplicate prompt",
yaml: `
kind: prompt
name: my_prompt
messages:
- role: user
content: hello
---
kind: prompt
name: my_prompt
messages:
- role: user
content: world
`,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
_, _, _, _, _, _, err := server.UnmarshalPrimitiveConfig(ctx, []byte(tc.yaml))
if err == nil {
t.Fatalf("UnmarshalPrimitiveConfig() expected a duplicate error, got nil")
}
if !strings.Contains(err.Error(), "declared more than once") {
t.Fatalf("UnmarshalPrimitiveConfig() error = %v, want it to mention 'declared more than once'", err)
}
})
}
}
func TestGroupConfigParsing(t *testing.T) {
ctx := context.Background()
tests := []struct {
name string
yaml string
want group.GroupConfig
wantError bool
}{
{
name: "valid named group",
yaml: `
kind: group
name: my_group
description: a group of tools and prompts
tools:
- tool_a
- tool_b
prompts:
- prompt_a
`,
want: group.GroupConfig{
Name: "my_group",
Description: "a group of tools and prompts",
ToolNames: []string{"tool_a", "tool_b"},
PromptNames: []string{"prompt_a"},
},
},
{
name: "named group with only description",
yaml: `
kind: group
name: my_group
description: just a description
`,
want: group.GroupConfig{
Name: "my_group",
Description: "just a description",
},
},
{
name: "default group with only description",
yaml: `
kind: group
name:
description: default server instruction
`,
want: group.GroupConfig{
Description: "default server instruction",
},
},
{
name: "default group omitting name field",
yaml: `
kind: group
description: default server instruction
`,
want: group.GroupConfig{
Description: "default server instruction",
},
},
{
name: "kind toolset folds into a tools-only group",
yaml: `
kind: toolset
name: my_toolset
tools:
- tool_a
- tool_b
`,
want: group.GroupConfig{
Name: "my_toolset",
ToolNames: []string{"tool_a", "tool_b"},
},
},
{
name: "default group declaring tools is an error",
yaml: `
kind: group
name:
tools:
- tool_a
`,
wantError: true,
},
{
name: "default group declaring prompts is an error",
yaml: `
kind: group
name:
prompts:
- prompt_a
`,
wantError: true,
},
{
name: "unknown field is an error",
yaml: `
kind: group
name: my_group
resources:
- res_a
`,
wantError: true,
},
{
name: "duplicate default group is an error",
yaml: `
kind: group
name:
description: first
---
kind: group
name:
description: second
`,
wantError: true,
},
{
name: "duplicate named group is an error",
yaml: `
kind: group
name: my_group
tools:
- tool_a
---
kind: group
name: my_group
tools:
- tool_b
`,
wantError: true,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
_, _, _, _, _, groups, err := server.UnmarshalPrimitiveConfig(ctx, []byte(tc.yaml))
if (err != nil) != tc.wantError {
t.Fatalf("UnmarshalPrimitiveConfig() returned error: %v, wantError: %v", err, tc.wantError)
}
if tc.wantError {
return
}
gc, ok := groups[tc.want.Name]
if !ok {
t.Fatalf("expected group %q to be parsed, got: %v", tc.want.Name, groups)
}
if diff := cmp.Diff(tc.want, gc); diff != "" {
t.Errorf("group mismatch (-want +got):\n%s", diff)
}
})
}
}
func TestGroupConfigValues(t *testing.T) {
ctx := context.Background()
yaml := `
kind: group
name: my_group
description: a group
tools:
- tool_a
- tool_b
prompts:
- prompt_a
`
_, _, _, _, _, groups, err := server.UnmarshalPrimitiveConfig(ctx, []byte(yaml))
if err != nil {
t.Fatalf("UnmarshalPrimitiveConfig() returned unexpected error: %v", err)
}
gc, ok := groups["my_group"]
if !ok {
t.Fatalf("expected group %q to be parsed, got: %v", "my_group", groups)
}
if gc.Name != "my_group" {
t.Errorf("group name: got %q, want %q", gc.Name, "my_group")
}
if gc.Description != "a group" {
t.Errorf("group description: got %q, want %q", gc.Description, "a group")
}
if diff := cmp.Diff([]string{"tool_a", "tool_b"}, gc.ToolNames); diff != "" {
t.Errorf("group tools mismatch (-want +got):\n%s", diff)
}
if diff := cmp.Diff([]string{"prompt_a"}, gc.PromptNames); diff != "" {
t.Errorf("group prompts mismatch (-want +got):\n%s", diff)
}
}
type offlineSourceConfig struct {
initialized *bool
}
@@ -1380,7 +1681,7 @@ func TestInitializeOfflineConfigs(t *testing.T) {
},
}
toolsMap, toolsetsMap, err := server.InitializeOfflineConfigs(ctx, cfg)
toolsMap, groupsMap, err := server.InitializeOfflineConfigs(ctx, cfg)
if err != nil {
t.Fatalf("InitializeOfflineConfigs returned error: %s", err)
}
@@ -1390,9 +1691,9 @@ func TestInitializeOfflineConfigs(t *testing.T) {
if _, ok := toolsMap["my-tool"]; !ok {
t.Errorf("expected tool %q in toolsMap, got %v", "my-tool", toolsMap)
}
// The implicit default ("") toolset should always be present.
if _, ok := toolsetsMap[""]; !ok {
t.Error("expected default toolset to be present")
// The implicit default ("") group should always be present.
if _, ok := groupsMap[""]; !ok {
t.Error("expected default group to be present")
}
}
internal/testutils/testutils.go
+25 -22
View File
@@ -24,6 +24,7 @@ import (
"strings"
"testing"
"github.com/googleapis/mcp-toolbox/internal/group"
"github.com/googleapis/mcp-toolbox/internal/log"
"github.com/googleapis/mcp-toolbox/internal/prompts"
"github.com/googleapis/mcp-toolbox/internal/tools"
@@ -161,8 +162,10 @@ var MockPrompt2 = NewMockPrompt("prompt2", "", prompts.Arguments{
{Parameter: parameters.NewStringParameter("arg1", "This is the first argument.")},
})
// SetUpResources setups resources to test against
func SetUpResources(t *testing.T, mockTools []MockTool, mockPrompts []MockPrompt) (map[string]tools.Tool, map[string]tools.Toolset, map[string]prompts.Prompt, map[string]prompts.Promptset) {
// SetUpResources setups resources to test against. The returned groups map is the
// source of truth used by PrimitiveManager; assert group membership via
// groups[name].ContainsTool / ContainsPrompt.
func SetUpResources(t *testing.T, mockTools []MockTool, mockPrompts []MockPrompt) (map[string]tools.Tool, map[string]prompts.Prompt, map[string]group.Group) {
toolsMap := make(map[string]tools.Tool)
var allTools []string
for _, tool := range mockTools {
@@ -170,20 +173,11 @@ func SetUpResources(t *testing.T, mockTools []MockTool, mockPrompts []MockPrompt
allTools = append(allTools, tool.Name)
}
toolsets := make(map[string]tools.Toolset)
groupToolNames := make(map[string][]string)
if len(allTools) > 0 {
for name, l := range map[string][]string{
"": allTools,
"tool1_only": {allTools[0]},
"tool2_only": {allTools[1]},
} {
tc := tools.ToolsetConfig{Name: name, ToolNames: l}
m, err := tc.Initialize(MockVersionString, toolsMap)
if err != nil {
t.Fatalf("unable to initialize toolset %q: %s", name, err)
}
toolsets[name] = m
}
groupToolNames[""] = allTools
groupToolNames["tool1_only"] = []string{allTools[0]}
groupToolNames["tool2_only"] = []string{allTools[1]}
}
promptsMap := make(map[string]prompts.Prompt)
@@ -193,15 +187,24 @@ func SetUpResources(t *testing.T, mockTools []MockTool, mockPrompts []MockPrompt
allPrompts = append(allPrompts, prompt.Name)
}
promptsets := make(map[string]prompts.Promptset)
// Build the authoritative groups map directly. Each named collection
// contributes its tool names; all prompts belong to the default (nameless)
// group, matching the legacy default-toolset behavior.
groupNames := make(map[string]struct{})
for name := range groupToolNames {
groupNames[name] = struct{}{}
}
if len(allPrompts) > 0 {
psc := prompts.PromptsetConfig{Name: "", PromptNames: allPrompts}
ps, err := psc.Initialize(MockVersionString, promptsMap)
if err != nil {
t.Fatalf("unable to initialize default promptset: %s", err)
groupNames[""] = struct{}{}
}
groups := make(map[string]group.Group)
for name := range groupNames {
gc := group.GroupConfig{Name: name, ToolNames: groupToolNames[name]}
if name == "" {
gc.PromptNames = allPrompts
}
promptsets[""] = ps
groups[name] = group.NewGroup(gc)
}
return toolsMap, toolsets, promptsMap, promptsets
return toolsMap, promptsMap, groups
}
internal/tools/cloudgda/cloudgda_test.go
+2 -2
View File
@@ -268,12 +268,12 @@ func TestInvoke(t *testing.T) {
{Name: "query", Value: query},
}
primMgr := primitives.NewPrimitiveManager(srcs, nil, nil, nil, nil, nil, nil)
primitiveMgr := primitives.NewPrimitiveManager(srcs, nil, nil, nil, nil, nil)
ctx := testutils.ContextWithUserAgent(context.Background(), "test-user-agent")
// Invoke the tool
result, err := tool.Invoke(ctx, primMgr, params, "")
result, err := tool.Invoke(ctx, primitiveMgr, params, "")
if err != nil {
t.Fatalf("tool invocation failed: %v", err)
}
internal/util/util.go
+1
View File
@@ -214,6 +214,7 @@ type GenAIMetricAttrs struct {
OperationName string
ToolName string
PromptName string
GroupName string
NetworkProtocolName string
NetworkProtocolVersion string
}