From 233f26ca255c6bf08d73d13d411d7015551972e8 Mon Sep 17 00:00:00 2001 From: Deeven-Seru Date: Wed, 15 Apr 2026 10:32:26 +0530 Subject: [PATCH] feat: add emulator mode for mocked tool responses --- cmd/internal/flags.go | 2 + cmd/internal/serve/command.go | 28 +++++ cmd/internal/serve/command_test.go | 3 +- cmd/root.go | 20 ++-- internal/server/config.go | 4 + internal/server/emulator.go | 160 +++++++++++++++++++++++++++++ internal/server/emulator_test.go | 91 ++++++++++++++++ internal/server/server.go | 9 ++ 8 files changed, 310 insertions(+), 7 deletions(-) create mode 100644 internal/server/emulator.go create mode 100644 internal/server/emulator_test.go diff --git a/cmd/internal/flags.go b/cmd/internal/flags.go index 0ccc14ae28e6..bb910927ecb6 100644 --- a/cmd/internal/flags.go +++ b/cmd/internal/flags.go @@ -68,4 +68,6 @@ func ServeFlags(flags *pflag.FlagSet, opts *ToolboxOptions) { flags.StringVar(&opts.Cfg.McpPrmFile, "mcp-prm-file", "", "Path to a manual Protected Resource Metadata (PRM) JSON file. If provided, overrides auto-generation.") flags.StringSliceVar(&opts.Cfg.AllowedOrigins, "allowed-origins", []string{"*"}, "Specifies a list of origins permitted to access this server. Defaults to '*'.") flags.StringSliceVar(&opts.Cfg.AllowedHosts, "allowed-hosts", []string{"*"}, "Specifies a list of hosts permitted to access this server. Defaults to '*'.") + flags.BoolVar(&opts.Cfg.EmulatorMode, "emulator-mode", false, "Enable emulator mode to return mock tool responses instead of calling live backends.") + flags.StringVar(&opts.Cfg.EmulatorMocksFile, "emulator-mocks-file", "", "Path to a JSON file containing emulator mocks. Used only when --emulator-mode is enabled.") } diff --git a/cmd/internal/serve/command.go b/cmd/internal/serve/command.go index 79ab3f1967a5..c1fdf99225ed 100644 --- a/cmd/internal/serve/command.go +++ b/cmd/internal/serve/command.go @@ -23,6 +23,7 @@ import ( "time" "github.com/googleapis/mcp-toolbox/cmd/internal" + "github.com/googleapis/mcp-toolbox/internal/auth/generic" "github.com/googleapis/mcp-toolbox/internal/server" "github.com/spf13/cobra" ) @@ -34,6 +35,7 @@ func NewCommand(opts *internal.ToolboxOptions) *cobra.Command { Long: "Deploy the toolbox server", } flags := cmd.Flags() + internal.ConfigFileFlags(flags, opts) internal.ServeFlags(flags, opts) cmd.RunE = func(*cobra.Command, []string) error { return runServe(cmd, opts) } return cmd @@ -71,6 +73,32 @@ func runServe(cmd *cobra.Command, opts *internal.ToolboxOptions) error { _ = shutdown(ctx) }() + _, err = opts.LoadConfig(ctx, &internal.ConfigParser{}) + if err != nil { + return err + } + + // Validate ToolboxUrl if MCP Auth is enabled + for _, authSvc := range opts.Cfg.AuthServiceConfigs { + if genCfg, ok := authSvc.(generic.Config); ok && genCfg.McpEnabled { + if opts.Cfg.ToolboxUrl == "" { + opts.Cfg.ToolboxUrl = os.Getenv("TOOLBOX_URL") + } + if opts.Cfg.ToolboxUrl == "" { + errMsg := fmt.Errorf("MCP Auth is enabled but Toolbox URL is missing. Please provide it via --toolbox-url flag or TOOLBOX_URL environment variable") + opts.Logger.ErrorContext(ctx, errMsg.Error()) + return errMsg + } + break + } + } + + if opts.Cfg.EmulatorMode && opts.Cfg.EmulatorMocksFile == "" { + errMsg := fmt.Errorf("--emulator-mocks-file is required when --emulator-mode is enabled") + opts.Logger.ErrorContext(ctx, errMsg.Error()) + return errMsg + } + // start server s, err := server.NewServer(ctx, opts.Cfg) if err != nil { diff --git a/cmd/internal/serve/command_test.go b/cmd/internal/serve/command_test.go index b11c0e7843be..a022470c45ff 100644 --- a/cmd/internal/serve/command_test.go +++ b/cmd/internal/serve/command_test.go @@ -43,11 +43,12 @@ func serveCommand(ctx context.Context, args []string) (string, error) { } func TestServe(t *testing.T) { + t.Setenv("SQLITE_DATABASE", "test.db") // context will automatically shutdown in 1 second. ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) defer cancel() - args := []string{"serve", "--port", "0"} + args := []string{"serve", "--prebuilt", "sqlite", "--port", "0"} output, err := serveCommand(ctx, args) if err != nil { t.Fatalf("expected graceful shutdown without error, got: %v", err) diff --git a/cmd/root.go b/cmd/root.go index 3d8d660a3d92..e84f3ed174ff 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -132,13 +132,13 @@ func NewCommand(opts *internal.ToolboxOptions) *cobra.Command { return cmd } -func handleDynamicReload(ctx context.Context, toolsFile internal.Config, s *server.Server) error { +func handleDynamicReload(ctx context.Context, toolsFile internal.Config, s *server.Server, emulatorMode bool, emulatorMocksFile string) error { logger, err := util.LoggerFromContext(ctx) if err != nil { panic(err) } - sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := validateReloadEdits(ctx, toolsFile) + sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := validateReloadEdits(ctx, toolsFile, emulatorMode, emulatorMocksFile) if err != nil { errMsg := fmt.Errorf("unable to validate reloaded edits: %w", err) logger.WarnContext(ctx, errMsg.Error()) @@ -152,7 +152,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, + ctx context.Context, toolsFile internal.Config, emulatorMode bool, emulatorMocksFile string, ) (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, ) { logger, err := util.LoggerFromContext(ctx) @@ -178,6 +178,8 @@ func validateReloadEdits( ToolConfigs: toolsFile.Tools, ToolsetConfigs: toolsFile.Toolsets, PromptConfigs: toolsFile.Prompts, + EmulatorMode: emulatorMode, + EmulatorMocksFile: emulatorMocksFile, } sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap, err := server.InitializeConfigs(ctx, reloadedConfig) @@ -233,7 +235,7 @@ func scanWatchedFiles(watchingFolder bool, folderToWatch string, watchedFiles ma } // watchChanges checks for changes in the provided yaml config(s) or folder. -func watchChanges(ctx context.Context, watchDirs map[string]bool, watchedFiles map[string]bool, s *server.Server, pollTickerSecond int) { +func watchChanges(ctx context.Context, watchDirs map[string]bool, watchedFiles map[string]bool, s *server.Server, pollTickerSecond int, emulatorMode bool, emulatorMocksFile string) { logger, err := util.LoggerFromContext(ctx) if err != nil { panic(err) @@ -379,7 +381,7 @@ func watchChanges(ctx context.Context, watchDirs map[string]bool, watchedFiles m continue } - err = handleDynamicReload(ctx, reloadedConfig, s) + err = handleDynamicReload(ctx, reloadedConfig, s, emulatorMode, emulatorMocksFile) if err != nil { errMsg := fmt.Errorf("unable to parse reloaded config at %q: %w", reloadedConfig, err) logger.WarnContext(ctx, errMsg.Error()) @@ -468,6 +470,12 @@ func run(cmd *cobra.Command, opts *internal.ToolboxOptions) error { } } + if opts.Cfg.EmulatorMode && opts.Cfg.EmulatorMocksFile == "" { + errMsg := fmt.Errorf("--emulator-mocks-file is required when --emulator-mode is enabled") + opts.Logger.ErrorContext(ctx, errMsg.Error()) + return errMsg + } + // start server s, err := server.NewServer(ctx, opts.Cfg) if err != nil { @@ -510,7 +518,7 @@ func run(cmd *cobra.Command, opts *internal.ToolboxOptions) error { if isCustomConfigured && !opts.Cfg.DisableReload { watchDirs, watchedFiles := resolveWatcherInputs(opts.Config, opts.Configs, opts.ConfigFolder) // start watching the file(s) or folder for changes to trigger dynamic reloading - go watchChanges(ctx, watchDirs, watchedFiles, s, opts.Cfg.PollInterval) + go watchChanges(ctx, watchDirs, watchedFiles, s, opts.Cfg.PollInterval, opts.Cfg.EmulatorMode, opts.Cfg.EmulatorMocksFile) } // wait for either the server to error out or the command's context to be canceled diff --git a/internal/server/config.go b/internal/server/config.go index efa947cefeac..56f468dd7816 100644 --- a/internal/server/config.go +++ b/internal/server/config.go @@ -84,6 +84,10 @@ type ServerConfig struct { UserAgentMetadata []string // PollInterval sets the polling frequency for configuration file updates. PollInterval int + // EmulatorMode determines whether tool invocations should use local mock responses instead of live backends. + EmulatorMode bool + // EmulatorMocksFile is a JSON file that defines mock responses for emulator mode. + EmulatorMocksFile string } type logFormat string diff --git a/internal/server/emulator.go b/internal/server/emulator.go new file mode 100644 index 000000000000..0379cb321864 --- /dev/null +++ b/internal/server/emulator.go @@ -0,0 +1,160 @@ +// 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 server + +import ( + "context" + "encoding/json" + "fmt" + "os" + "reflect" + + "github.com/googleapis/mcp-toolbox/internal/embeddingmodels" + "github.com/googleapis/mcp-toolbox/internal/tools" + "github.com/googleapis/mcp-toolbox/internal/util" + "github.com/googleapis/mcp-toolbox/internal/util/parameters" +) + +type emulatorMockFile struct { + Mocks []emulatorMock `json:"mocks"` +} + +type emulatorMock struct { + ToolName string `json:"tool_name"` + Description string `json:"description,omitempty"` + Parameters map[string]any `json:"parameters"` + Response any `json:"response"` +} + +type emulatorTool struct { + name string + base tools.Tool + mocks []emulatorMock +} + +func loadEmulatorMocks(path string) (map[string][]emulatorMock, error) { + raw, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("unable to read emulator mocks file %q: %w", path, err) + } + + var file emulatorMockFile + if err := json.Unmarshal(raw, &file); err != nil { + return nil, fmt.Errorf("unable to parse emulator mocks file %q: %w", path, err) + } + + mocks := make(map[string][]emulatorMock) + for i := range file.Mocks { + m := &file.Mocks[i] + if m.ToolName == "" { + return nil, fmt.Errorf("invalid emulator mock at index %d: tool_name is required", i) + } + if m.Parameters == nil { + m.Parameters = map[string]any{} + } + normalized, err := normalizeJSONValue(m.Parameters) + if err != nil { + return nil, fmt.Errorf("failed to normalize parameters for mock at index %d: %w", i, err) + } + normalizedMap, ok := normalized.(map[string]any) + if !ok { + return nil, fmt.Errorf("failed to normalize parameters for mock at index %d: expected object parameters", i) + } + m.Parameters = normalizedMap + mocks[m.ToolName] = append(mocks[m.ToolName], *m) + } + + return mocks, nil +} + +func wrapToolsForEmulator(toolMap map[string]tools.Tool, mocksByTool map[string][]emulatorMock) map[string]tools.Tool { + wrapped := make(map[string]tools.Tool, len(toolMap)) + for toolName, t := range toolMap { + wrapped[toolName] = emulatorTool{ + name: toolName, + base: t, + mocks: mocksByTool[toolName], + } + } + return wrapped +} + +func normalizeJSONValue(v any) (any, error) { + buf, err := json.Marshal(v) + if err != nil { + return nil, err + } + var normalized any + if err := json.Unmarshal(buf, &normalized); err != nil { + return nil, err + } + return normalized, nil +} + +func (e emulatorTool) findMock(inMap map[string]any) (any, bool) { + for _, mock := range e.mocks { + if reflect.DeepEqual(inMap, mock.Parameters) { + return mock.Response, true + } + } + return nil, false +} + +func (e emulatorTool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) { + in, err := normalizeJSONValue(params.AsMap()) + if err != nil { + return nil, util.NewAgentError("emulator mode: failed to normalize input parameters", err) + } + inMap, ok := in.(map[string]any) + if !ok { + return nil, util.NewAgentError("emulator mode: failed to normalize input parameters", nil) + } + if response, found := e.findMock(inMap); found { + return response, nil + } + return nil, util.NewAgentError(fmt.Sprintf("emulator mode: no mock matched for tool %q with parameters %v", e.name, inMap), nil) +} + +func (e emulatorTool) EmbedParams(ctx context.Context, params parameters.ParamValues, embeddingModels map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) { + return e.base.EmbedParams(ctx, params, embeddingModels) +} + +func (e emulatorTool) Manifest() tools.Manifest { + return e.base.Manifest() +} + +func (e emulatorTool) McpManifest() tools.McpManifest { + return e.base.McpManifest() +} + +func (e emulatorTool) Authorized(verifiedAuthServices []string) bool { + return e.base.Authorized(verifiedAuthServices) +} + +func (e emulatorTool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (bool, error) { + return e.base.RequiresClientAuthorization(resourceMgr) +} + +func (e emulatorTool) ToConfig() tools.ToolConfig { + return e.base.ToConfig() +} + +func (e emulatorTool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) { + return e.base.GetAuthTokenHeaderName(resourceMgr) +} + +func (e emulatorTool) GetParameters() parameters.Parameters { + return e.base.GetParameters() +} diff --git a/internal/server/emulator_test.go b/internal/server/emulator_test.go new file mode 100644 index 000000000000..5ea6cd449e38 --- /dev/null +++ b/internal/server/emulator_test.go @@ -0,0 +1,91 @@ +// 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 server + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/googleapis/mcp-toolbox/internal/util" + "github.com/googleapis/mcp-toolbox/internal/util/parameters" +) + +func TestLoadEmulatorMocks(t *testing.T) { + t.Parallel() + dir := t.TempDir() + mockFile := filepath.Join(dir, "mocks.json") + content := `{ + "mocks": [ + {"tool_name": "search_users_bq", "parameters": {"id": 123}, "response": [{"id":123}]}, + {"tool_name": "search_users_bq", "parameters": {"email": "alice@example.com"}, "response": []} + ] +}` + if err := os.WriteFile(mockFile, []byte(content), 0o600); err != nil { + t.Fatalf("failed writing mock file: %v", err) + } + + mocks, err := loadEmulatorMocks(mockFile) + if err != nil { + t.Fatalf("loadEmulatorMocks returned error: %v", err) + } + if got := len(mocks["search_users_bq"]); got != 2 { + t.Fatalf("expected 2 mocks for search_users_bq, got %d", got) + } +} + +func TestEmulatorToolInvoke_MatchReturnsMock(t *testing.T) { + t.Parallel() + tool := emulatorTool{ + name: "search_users_bq", + mocks: []emulatorMock{{ + ToolName: "search_users_bq", + Parameters: map[string]any{"id": float64(123)}, + Response: []map[string]any{{"id": float64(123), "name": "Alice"}}, + }}, + } + params := parameters.ParamValues{{Name: "id", Value: int64(123)}} + + got, err := tool.Invoke(context.Background(), nil, params, "") + if err != nil { + t.Fatalf("Invoke returned error: %v", err) + } + rows, ok := got.([]map[string]any) + if !ok || len(rows) != 1 || rows[0]["name"] != "Alice" { + t.Fatalf("unexpected mock response: %#v", got) + } +} + +func TestEmulatorToolInvoke_NoMatchReturnsAgentError(t *testing.T) { + t.Parallel() + tool := emulatorTool{ + name: "search_users_bq", + mocks: []emulatorMock{{ + ToolName: "search_users_bq", + Parameters: map[string]any{"id": float64(123)}, + Response: []map[string]any{{"id": float64(123)}}, + }}, + } + params := parameters.ParamValues{{Name: "id", Value: int64(999)}} + + _, err := tool.Invoke(context.Background(), nil, params, "") + if err == nil { + t.Fatal("expected error for unmatched mock") + } + if _, ok := err.(*util.AgentError); !ok { + t.Fatalf("expected AgentError, got %T", err) + } +} diff --git a/internal/server/server.go b/internal/server/server.go index d42c25b9cfcd..7655469ea53d 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -198,6 +198,15 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) ( } l.InfoContext(ctx, fmt.Sprintf("Initialized %d tools: %s", len(toolsMap), strings.Join(toolNames, ", "))) + if cfg.EmulatorMode { + mocksByTool, err := loadEmulatorMocks(cfg.EmulatorMocksFile) + if err != nil { + return nil, nil, nil, nil, nil, nil, nil, err + } + toolsMap = wrapToolsForEmulator(toolsMap, mocksByTool) + l.InfoContext(ctx, fmt.Sprintf("Emulator mode enabled with mocks file: %s", cfg.EmulatorMocksFile)) + } + // create a default toolset that contains all tools allToolNames := make([]string, 0, len(toolsMap)) for name := range toolsMap {