Add MCP loader and runner options to the SDK server

Signed-off-by: Donnie Adams <[email protected]>

Author: thedadams

pkg/engine/engine.go

@@ -11,7 +11,6 @@ import (
 	"sync"
 
 	"github.com/gptscript-ai/gptscript/pkg/counter"
-	"github.com/gptscript-ai/gptscript/pkg/mcp"
 	"github.com/gptscript-ai/gptscript/pkg/types"
 	"github.com/gptscript-ai/gptscript/pkg/version"
 )
@@ -314,11 +313,7 @@ func populateMessageParams(ctx Context, completion *types.CompletionRequest, too
 }
 
 func (e *Engine) runMCPInvoke(ctx Context, tool types.Tool, input string) (*Return, error) {
-	runner := e.MCPRunner
-	if runner == nil {
-		runner = mcp.DefaultRunner
-	}
-	output, err := runner.Run(ctx.Ctx, e.Progress, tool, input)
+	output, err := e.MCPRunner.Run(ctx.Ctx, e.Progress, tool, input)
 	if err != nil {
 		return nil, fmt.Errorf("failed to run MCP invoke: %w", err)
 	}

pkg/runner/runner.go

@@ -14,6 +14,7 @@ import (
 	context2 "github.com/gptscript-ai/gptscript/pkg/context"
 	"github.com/gptscript-ai/gptscript/pkg/credentials"
 	"github.com/gptscript-ai/gptscript/pkg/engine"
+	"github.com/gptscript-ai/gptscript/pkg/mcp"
 	"github.com/gptscript-ai/gptscript/pkg/types"
 	"golang.org/x/exp/maps"
 )
@@ -37,6 +38,7 @@ type Options struct {
 	CredentialOverrides []string              `usage:"-"`
 	Sequential          bool                  `usage:"-"`
 	Authorizer          AuthorizerFunc        `usage:"-"`
+	MCPRunner           engine.MCPRunner      `usage:"-"`
 }
 
 type RunOptions struct {
@@ -69,6 +71,9 @@ func Complete(opts ...Options) (result Options) {
 		if opt.CredentialOverrides != nil {
 			result.CredentialOverrides = append(result.CredentialOverrides, opt.CredentialOverrides...)
 		}
+		if opt.MCPRunner != nil {
+			result.MCPRunner = opt.MCPRunner
+		}
 	}
 	return
 }
@@ -87,6 +92,9 @@ func complete(opts ...Options) Options {
 	if result.Authorizer == nil {
 		result.Authorizer = DefaultAuthorizer
 	}
+	if result.MCPRunner == nil {
+		result.MCPRunner = mcp.DefaultRunner
+	}
 	return result
 }
 
@@ -99,6 +107,7 @@ type Runner struct {
 	credOverrides  []string
 	credStore      credentials.CredentialStore
 	sequential     bool
+	mcpRunner      engine.MCPRunner
 }
 
 func New(client engine.Model, credStore credentials.CredentialStore, opts ...Options) (*Runner, error) {
@@ -113,6 +122,7 @@ func New(client engine.Model, credStore credentials.CredentialStore, opts ...Opt
 		credStore:      credStore,
 		sequential:     opt.Sequential,
 		auth:           opt.Authorizer,
+		mcpRunner:      opt.MCPRunner,
 	}
 
 	if opt.StartPort != 0 {
@@ -326,6 +336,7 @@ func (r *Runner) start(callCtx engine.Context, state *State, monitor Monitor, en
 
 	e := engine.Engine{
 		Model:          r.c,
+		MCPRunner:      r.mcpRunner,
 		RuntimeManager: runtimeWithLogger(callCtx, monitor, r.runtimeManager),
 		Progress:       progress,
 		Env:            env,
@@ -524,6 +535,7 @@ func (r *Runner) resume(callCtx engine.Context, monitor Monitor, env []string, s
 
 		e := engine.Engine{
 			Model:          r.c,
+			MCPRunner:      r.mcpRunner,
 			RuntimeManager: runtimeWithLogger(callCtx, monitor, r.runtimeManager),
 			Progress:       progress,
 			Env:            env,

pkg/sdkserver/routes.go

@@ -29,6 +29,7 @@ type server struct {
 	datasetTool, workspaceTool string
 	serverToolsEnv             []string
 	client                     *gptscript.GPTScript
+	mcpLoader                  loader.MCPLoader
 	events                     *broadcaster.Broadcaster[event]
 
 	runtimeManager engine.RuntimeManager
@@ -283,11 +284,20 @@ func (s *server) load(w http.ResponseWriter, r *http.Request) {
 	}
 
 	if reqObject.Content != "" {
-		prg, err = loader.ProgramFromSource(ctx, reqObject.Content, reqObject.SubTool, loader.Options{Cache: s.client.Cache})
+		prg, err = loader.ProgramFromSource(ctx, reqObject.Content, reqObject.SubTool, loader.Options{
+			Cache:     s.client.Cache,
+			MCPLoader: s.mcpLoader,
+		})
 	} else if reqObject.File != "" {
-		prg, err = loader.Program(ctx, reqObject.File, reqObject.SubTool, loader.Options{Cache: s.client.Cache})
+		prg, err = loader.Program(ctx, reqObject.File, reqObject.SubTool, loader.Options{
+			Cache:     s.client.Cache,
+			MCPLoader: s.mcpLoader,
+		})
 	} else {
-		prg, err = loader.ProgramFromSource(ctx, reqObject.ToolDefs.String(), reqObject.SubTool, loader.Options{Cache: s.client.Cache})
+		prg, err = loader.ProgramFromSource(ctx, reqObject.ToolDefs.String(), reqObject.SubTool, loader.Options{
+			Cache:     s.client.Cache,
+			MCPLoader: s.mcpLoader,
+		})
 	}
 	if err != nil {
 		writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to load program: %w", err))

pkg/sdkserver/run.go

@@ -36,7 +36,11 @@ func (s *server) execAndStream(ctx context.Context, programLoader loaderFunc, lo
 	if defaultModel == "" {
 		defaultModel = s.gptscriptOpts.OpenAI.DefaultModel
 	}
-	prg, err := programLoader(ctx, toolDef.String(), subTool, loader.Options{Cache: g.Cache, DefaultModel: defaultModel})
+	prg, err := programLoader(ctx, toolDef.String(), subTool, loader.Options{
+		Cache:        g.Cache,
+		DefaultModel: defaultModel,
+		MCPLoader:    s.mcpLoader,
+	})
 	if err != nil {
 		writeError(logger, w, http.StatusInternalServerError, fmt.Errorf("failed to load program: %w", err))
 		return

pkg/sdkserver/server.go

@@ -16,6 +16,8 @@ import (
 	"github.com/google/uuid"
 	"github.com/gptscript-ai/broadcaster"
 	"github.com/gptscript-ai/gptscript/pkg/gptscript"
+	"github.com/gptscript-ai/gptscript/pkg/loader"
+	"github.com/gptscript-ai/gptscript/pkg/mcp"
 	"github.com/gptscript-ai/gptscript/pkg/mvl"
 	"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
 	"github.com/gptscript-ai/gptscript/pkg/runner"
@@ -26,6 +28,7 @@ import (
 type Options struct {
 	gptscript.Options
 
+	MCPLoader                  loader.MCPLoader
 	ListenAddress              string
 	DatasetTool, WorkspaceTool string
 	ServerToolsEnv             []string
@@ -114,6 +117,7 @@ func run(ctx context.Context, listener net.Listener, opts Options) error {
 		serverToolsEnv: opts.ServerToolsEnv,
 
 		client:           g,
+		mcpLoader:        opts.MCPLoader,
 		events:           events,
 		runtimeManager:   runtimes.Default(opts.Cache.CacheDir, opts.SystemToolsDir),
 		waitingToConfirm: make(map[string]chan runner.AuthorizerResponse),
@@ -168,6 +172,7 @@ func complete(opts ...Options) Options {
 		result.WorkspaceTool = types.FirstSet(opt.WorkspaceTool, result.WorkspaceTool)
 		result.Debug = types.FirstSet(opt.Debug, result.Debug)
 		result.DisableServerErrorLogging = types.FirstSet(opt.DisableServerErrorLogging, result.DisableServerErrorLogging)
+		result.MCPLoader = types.FirstSet(opt.MCPLoader, result.MCPLoader)
 	}
 
 	if result.ListenAddress == "" {
@@ -183,6 +188,9 @@ func complete(opts ...Options) Options {
 	if len(result.ServerToolsEnv) == 0 {
 		result.ServerToolsEnv = os.Environ()
 	}
+	if result.MCPLoader == nil {
+		result.MCPLoader = mcp.DefaultLoader
+	}
 
 	return result
 }

pkg/tests/runner2_test.go

@@ -308,7 +308,7 @@ name: mcp
       ],
       "type": "object"
     },
-    "instructions": "#!sys.mcp.invoke.append_insight e057d98f5d43e56fda04eb3e7ea6120c93b5bcaf832090fca76e8d744e2de494",
+    "instructions": "#!sys.mcp.invoke.append_insight 607ca64476abf0288ef49061557243e43735fd4de4bc5fdcd51d93049ffa023e",
     "id": "inline:append_insight",
     "localTools": {
       "append_insight": "inline:append_insight",
@@ -341,7 +341,7 @@ name: mcp
       ],
       "type": "object"
     },
-    "instructions": "#!sys.mcp.invoke.create_table e057d98f5d43e56fda04eb3e7ea6120c93b5bcaf832090fca76e8d744e2de494",
+    "instructions": "#!sys.mcp.invoke.create_table 607ca64476abf0288ef49061557243e43735fd4de4bc5fdcd51d93049ffa023e",
     "id": "inline:create_table",
     "localTools": {
       "append_insight": "inline:append_insight",
@@ -374,7 +374,7 @@ name: mcp
       ],
       "type": "object"
     },
-    "instructions": "#!sys.mcp.invoke.describe_table e057d98f5d43e56fda04eb3e7ea6120c93b5bcaf832090fca76e8d744e2de494",
+    "instructions": "#!sys.mcp.invoke.describe_table 607ca64476abf0288ef49061557243e43735fd4de4bc5fdcd51d93049ffa023e",
     "id": "inline:describe_table",
     "localTools": {
       "append_insight": "inline:append_insight",
@@ -398,7 +398,7 @@ name: mcp
     "arguments": {
       "type": "object"
     },
-    "instructions": "#!sys.mcp.invoke.list_tables e057d98f5d43e56fda04eb3e7ea6120c93b5bcaf832090fca76e8d744e2de494",
+    "instructions": "#!sys.mcp.invoke.list_tables 607ca64476abf0288ef49061557243e43735fd4de4bc5fdcd51d93049ffa023e",
     "id": "inline:list_tables",
     "localTools": {
       "append_insight": "inline:append_insight",
@@ -500,7 +500,7 @@ name: mcp
       ],
       "type": "object"
     },
-    "instructions": "#!sys.mcp.invoke.read_query e057d98f5d43e56fda04eb3e7ea6120c93b5bcaf832090fca76e8d744e2de494",
+    "instructions": "#!sys.mcp.invoke.read_query 607ca64476abf0288ef49061557243e43735fd4de4bc5fdcd51d93049ffa023e",
     "id": "inline:read_query",
     "localTools": {
       "append_insight": "inline:append_insight",
@@ -533,7 +533,7 @@ name: mcp
       ],
       "type": "object"
     },
-    "instructions": "#!sys.mcp.invoke.write_query e057d98f5d43e56fda04eb3e7ea6120c93b5bcaf832090fca76e8d744e2de494",
+    "instructions": "#!sys.mcp.invoke.write_query 607ca64476abf0288ef49061557243e43735fd4de4bc5fdcd51d93049ffa023e",
     "id": "inline:write_query",
     "localTools": {
       "append_insight": "inline:append_insight",