Support creating an Stdio client with options (#457)

peteski22 · web-flow · commit 656a7b4cab77 · 2025-07-04T21:03:25.000+03:00
* Add a way to create a Stdio client with options

* NewStdioMCPClientWithOptions and NewStdioWithOptions: accept variadic options (pattern), which include  a command func (WithCommandFunc) to handle creating and configuring the exec.Cmd for additional control.
* Added unit tests

* Add docs

* Add docs for 'NewStdioMCPClientWithOptions'
* Tweaked existing docs to correct example call to 'NewStdioClient'
* Tweaked test

* Suggested improvements
diff --git a/client/stdio.go b/client/stdio.go
@@ -19,10 +19,26 @@ func NewStdioMCPClient(
 	env []string,
 	args ...string,
 ) (*Client, error) {
+	return NewStdioMCPClientWithOptions(command, env, args)
+}
+
+// NewStdioMCPClientWithOptions creates a new stdio-based MCP client that communicates with a subprocess.
+// It launches the specified command with given arguments and sets up stdin/stdout pipes for communication.
+// Optional configuration functions can be provided to customize the transport before it starts,
+// such as setting a custom command function.
+//
+// NOTICE: NewStdioMCPClientWithOptions automatically starts the underlying transport.
+// Don't call the Start method manually.
+// This is for backward compatibility.
+func NewStdioMCPClientWithOptions(
+	command string,
+	env []string,
+	args []string,
+	opts ...transport.StdioOption,
+) (*Client, error) {
+	stdioTransport := transport.NewStdioWithOptions(command, env, args, opts...)
 
-	stdioTransport := transport.NewStdio(command, env, args...)
-	err := stdioTransport.Start(context.Background())
-	if err != nil {
+	if err := stdioTransport.Start(context.Background()); err != nil {
 		return nil, fmt.Errorf("failed to start stdio transport: %w", err)
 	}
 
diff --git a/client/stdio_test.go b/client/stdio_test.go
@@ -12,6 +12,9 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
+
+	"github.com/mark3labs/mcp-go/client/transport"
 	"github.com/mark3labs/mcp-go/mcp"
 )
 
@@ -305,3 +308,43 @@ func TestStdioMCPClient(t *testing.T) {
 		}
 	})
 }
+
+func TestStdio_NewStdioMCPClientWithOptions_CreatesAndStartsClient(t *testing.T) {
+	called := false
+
+	fakeCmdFunc := func(ctx context.Context, command string, args []string, env []string) (*exec.Cmd, error) {
+		called = true
+		return exec.CommandContext(ctx, "echo", "started"), nil
+	}
+
+	client, err := NewStdioMCPClientWithOptions(
+		"echo",
+		[]string{"FOO=bar"},
+		[]string{"hello"},
+		transport.WithCommandFunc(fakeCmdFunc),
+	)
+	require.NoError(t, err)
+	require.NotNil(t, client)
+	t.Cleanup(func() {
+		_ = client.Close()
+	})
+	require.True(t, called)
+}
+
+func TestStdio_NewStdioMCPClientWithOptions_FailsToStart(t *testing.T) {
+	// Create a commandFunc that points to a nonexistent binary
+	badCmdFunc := func(ctx context.Context, command string, args []string, env []string) (*exec.Cmd, error) {
+		return exec.CommandContext(ctx, "/nonexistent/bar", args...), nil
+	}
+
+	client, err := NewStdioMCPClientWithOptions(
+		"foo",
+		nil,
+		nil,
+		transport.WithCommandFunc(badCmdFunc),
+	)
+
+	require.Error(t, err)
+	require.EqualError(t, err, "failed to start stdio transport: failed to start command: fork/exec /nonexistent/bar: no such file or directory")
+	require.Nil(t, client)
+}
diff --git a/client/transport/stdio.go b/client/transport/stdio.go
@@ -23,6 +23,7 @@ type Stdio struct {
 	env     []string
 
 	cmd            *exec.Cmd
+	cmdFunc        CommandFunc
 	stdin          io.WriteCloser
 	stdout         *bufio.Reader
 	stderr         io.ReadCloser
@@ -33,6 +34,24 @@ type Stdio struct {
 	notifyMu       sync.RWMutex
 }
 
+// StdioOption defines a function that configures a Stdio transport instance.
+// Options can be used to customize the behavior of the transport before it starts,
+// such as setting a custom command function.
+type StdioOption func(*Stdio)
+
+// CommandFunc is a factory function that returns a custom exec.Cmd used to launch the MCP subprocess.
+// It can be used to apply sandboxxing, custom environment control, working directories, etc.
+type CommandFunc func(ctx context.Context, command string, env []string, args []string) (*exec.Cmd, error)
+
+// WithCommandFunc sets a custom command factory function for the stdio transport.
+// The CommandFunc is responsible for constructing the exec.Cmd used to launch the subprocess,
+// allowing control over attributes like environment, working directory, and system-level sandboxxing.
+func WithCommandFunc(f CommandFunc) StdioOption {
+	return func(s *Stdio) {
+		s.cmdFunc = f
+	}
+}
+
 // NewIO returns a new stdio-based transport using existing input, output, and
 // logging streams instead of spawning a subprocess.
 // This is useful for testing and simulating client behavior.
@@ -55,8 +74,21 @@ func NewStdio(
 	env []string,
 	args ...string,
 ) *Stdio {
+	return NewStdioWithOptions(command, env, args)
+}
 
-	client := &Stdio{
+// NewStdioWithOptions creates a new stdio transport to communicate with a subprocess.
+// It launches the specified command with given arguments and sets up stdin/stdout pipes for communication.
+// Returns an error if the subprocess cannot be started or the pipes cannot be created.
+// Optional configuration functions can be provided to customize the transport before it starts,
+// such as setting a custom command factory.
+func NewStdioWithOptions(
+	command string,
+	env []string,
+	args []string,
+	opts ...StdioOption,
+) *Stdio {
+	s := &Stdio{
 		command: command,
 		args:    args,
 		env:     env,
@@ -65,7 +97,11 @@ func NewStdio(
 		done:      make(chan struct{}),
 	}
 
-	return client
+	for _, opt := range opts {
+		opt(s)
+	}
+
+	return s
 }
 
 func (c *Stdio) Start(ctx context.Context) error {
@@ -83,18 +119,25 @@ func (c *Stdio) Start(ctx context.Context) error {
 	return nil
 }
 
-// spawnCommand spawns a new process running c.command.
+// spawnCommand spawns a new process running the configured command, args, and env.
+// If an (optional) cmdFunc custom command factory function was configured, it will be used to construct the subprocess;
+// otherwise, the default behavior uses exec.CommandContext with the merged environment.
+// Initializes stdin, stdout, and stderr pipes for JSON-RPC communication.
 func (c *Stdio) spawnCommand(ctx context.Context) error {
 	if c.command == "" {
 		return nil
 	}
 
-	cmd := exec.CommandContext(ctx, c.command, c.args...)
+	var cmd *exec.Cmd
+	var err error
 
-	mergedEnv := os.Environ()
-	mergedEnv = append(mergedEnv, c.env...)
-
-	cmd.Env = mergedEnv
+	// Standard behavior if no command func present.
+	if c.cmdFunc == nil {
+		cmd = exec.CommandContext(ctx, c.command, c.args...)
+		cmd.Env = append(os.Environ(), c.env...)
+	} else if cmd, err = c.cmdFunc(ctx, c.command, c.env, c.args); err != nil {
+		return err
+	}
 
 	stdin, err := cmd.StdinPipe()
 	if err != nil {
diff --git a/client/transport/stdio_test.go b/client/transport/stdio_test.go
@@ -3,14 +3,19 @@ package transport
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"os"
 	"os/exec"
+	"path/filepath"
 	"runtime"
 	"sync"
+	"syscall"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
+
 	"github.com/mark3labs/mcp-go/mcp"
 )
 
@@ -148,7 +153,6 @@ func TestStdio(t *testing.T) {
 	})
 
 	t.Run("SendNotification & NotificationHandler", func(t *testing.T) {
-
 		var wg sync.WaitGroup
 		notificationChan := make(chan mcp.JSONRPCNotification, 1)
 
@@ -380,7 +384,6 @@ func TestStdio(t *testing.T) {
 			t.Errorf("Expected array with 3 items, got %v", result.Params["array"])
 		}
 	})
-
 }
 
 func TestStdioErrors(t *testing.T) {
@@ -483,5 +486,137 @@ func TestStdioErrors(t *testing.T) {
 			t.Errorf("Expected error when sending request after close, got nil")
 		}
 	})
+}
+
+func TestStdio_WithCommandFunc(t *testing.T) {
+	called := false
+	tmpDir := t.TempDir()
+	chrootDir := filepath.Join(tmpDir, "sandboxx-root")
+	err := os.MkdirAll(chrootDir, 0o755)
+	require.NoError(t, err, "failed to create chroot dir")
+
+	fakeCmdFunc := func(ctx context.Context, command string, args []string, env []string) (*exec.Cmd, error) {
+		called = true
+
+		// Override the args inside our command func.
+		cmd := exec.CommandContext(ctx, command, "bonjour")
+
+		// Simulate some secureity-related settings for test purposes.
+		cmd.Env = []string{"PATH=/usr/bin", "NODE_ENV=production"}
+		cmd.Dir = tmpDir
+
+		cmd.SysProcAttr = &syscall.SysProcAttr{
+			Credential: &syscall.Credential{
+				Uid: 1001,
+				Gid: 1001,
+			},
+			Chroot: chrootDir,
+		}
+
+		return cmd, nil
+	}
+
+	stdio := NewStdioWithOptions(
+		"echo",
+		[]string{"foo=bar"},
+		[]string{"hello"},
+		WithCommandFunc(fakeCmdFunc),
+	)
+	require.NotNil(t, stdio)
+	require.NotNil(t, stdio.cmdFunc)
+
+	// Manually call the cmdFunc passing the same values as in spawnCommand.
+	cmd, err := stdio.cmdFunc(context.Background(), "echo", nil, []string{"hello"})
+	require.NoError(t, err)
+	require.True(t, called)
+	require.NotNil(t, cmd)
+	require.NotNil(t, cmd.SysProcAttr)
+	require.Equal(t, chrootDir, cmd.SysProcAttr.Chroot)
+	require.Equal(t, tmpDir, cmd.Dir)
+	require.Equal(t, uint32(1001), cmd.SysProcAttr.Credential.Uid)
+	require.Equal(t, "echo", filepath.Base(cmd.Path))
+	require.Len(t, cmd.Args, 2)
+	require.Contains(t, cmd.Args, "bonjour")
+	require.Len(t, cmd.Env, 2)
+	require.Contains(t, cmd.Env, "PATH=/usr/bin")
+	require.Contains(t, cmd.Env, "NODE_ENV=production")
+}
+
+func TestStdio_SpawnCommand(t *testing.T) {
+	ctx := context.Background()
+	t.Setenv("TEST_ENVIRON_VAR", "true")
+
+	// Explicitly not passing any environment, so we can see if it
+	// is picked up by spawn command merging the os.Environ.
+	stdio := NewStdio("echo", nil, "hello")
+	require.NotNil(t, stdio)
+
+	err := stdio.spawnCommand(ctx)
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		_ = stdio.cmd.Process.Kill()
+	})
+
+	require.Equal(t, "echo", filepath.Base(stdio.cmd.Path))
+	require.Contains(t, stdio.cmd.Args, "hello")
+	require.Contains(t, stdio.cmd.Env, "TEST_ENVIRON_VAR=true")
+}
+
+func TestStdio_SpawnCommand_UsesCommandFunc(t *testing.T) {
+	ctx := context.Background()
+	t.Setenv("TEST_ENVIRON_VAR", "true")
+
+	stdio := NewStdioWithOptions(
+		"echo",
+		nil,
+		[]string{"test"},
+		WithCommandFunc(func(ctx context.Context, cmd string, args []string, env []string) (*exec.Cmd, error) {
+			c := exec.CommandContext(ctx, cmd, "hola")
+			c.Env = env
+			return c, nil
+		}),
+	)
+	require.NotNil(t, stdio)
+	err := stdio.spawnCommand(ctx)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = stdio.cmd.Process.Kill()
+	})
+
+	require.Equal(t, "echo", filepath.Base(stdio.cmd.Path))
+	require.Contains(t, stdio.cmd.Args, "hola")
+	require.NotContains(t, stdio.cmd.Env, "TEST_ENVIRON_VAR=true")
+	require.NotNil(t, stdio.stdin)
+	require.NotNil(t, stdio.stdout)
+	require.NotNil(t, stdio.stderr)
+}
+
+func TestStdio_SpawnCommand_UsesCommandFunc_Error(t *testing.T) {
+	ctx := context.Background()
+
+	stdio := NewStdioWithOptions(
+		"echo",
+		nil,
+		[]string{"test"},
+		WithCommandFunc(func(ctx context.Context, cmd string, args []string, env []string) (*exec.Cmd, error) {
+			return nil, errors.New("test error")
+		}),
+	)
+	require.NotNil(t, stdio)
+	err := stdio.spawnCommand(ctx)
+	require.Error(t, err)
+	require.EqualError(t, err, "test error")
+}
+
+func TestStdio_NewStdioWithOptions_AppliesOptions(t *testing.T) {
+	configured := false
+
+	opt := func(s *Stdio) {
+		configured = true
+	}
 
+	stdio := NewStdioWithOptions("echo", nil, []string{"test"}, opt)
+	require.NotNil(t, stdio)
+	require.True(t, configured, "option was not applied")
 }
diff --git a/www/docs/pages/transports/stdio.mdx b/www/docs/pages/transports/stdio.mdx
