feat(mcp): improve concurrent workspace execution

cmd/mcp/semaphore.go

@@ -0,0 +1,33 @@
+package mcp
+
+import (
+	"context"
+	"fmt"
+)
+
+type opSemaphore struct {
+	slots chan struct{}
+}
+
+func newOpSemaphore(max int) *opSemaphore {
+	if max <= 0 {
+		max = 1
+	}
+	return &opSemaphore{slots: make(chan struct{}, max)}
+}
+
+func (s *opSemaphore) acquire(ctx context.Context) (func(), error) {
+	if err := ctx.Err(); err != nil {
+		return nil, fmt.Errorf("waiting for a free operation slot: %w", err)
+	}
+	select {
+	case s.slots <- struct{}{}:
+		if err := ctx.Err(); err != nil {
+			<-s.slots
+			return nil, fmt.Errorf("waiting for a free operation slot: %w", err)
+		}
+		return func() { <-s.slots }, nil
+	case <-ctx.Done():
+		return nil, fmt.Errorf("waiting for a free operation slot: %w", ctx.Err())
+	}
+}

cmd/mcp/semaphore_test.go

@@ -0,0 +1,96 @@
+package mcp
+
+import (
+	"context"
+	"errors"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestOpSemaphore_LimitsConcurrency(t *testing.T) {
+	sem := newOpSemaphore(2)
+	var inFlight, maxInFlight atomic.Int32
+
+	track := func() {
+		cur := inFlight.Add(1)
+		for {
+			m := maxInFlight.Load()
+			if cur <= m || maxInFlight.CompareAndSwap(m, cur) {
+				break
+			}
+		}
+		time.Sleep(20 * time.Millisecond)
+		inFlight.Add(-1)
+	}
+
+	done := make(chan struct{}, 5)
+	for range 5 {
+		go func() {
+			release, err := sem.acquire(context.Background())
+			if err != nil {
+				t.Errorf("acquire failed: %v", err)
+				done <- struct{}{}
+				return
+			}
+			track()
+			release()
+			done <- struct{}{}
+		}()
+	}
+	for range 5 {
+		<-done
+	}
+
+	if got := maxInFlight.Load(); got > 2 {
+		t.Fatalf("max concurrent = %d, want <= 2", got)
+	}
+}
+
+func TestOpSemaphore_ReleaseAllowsNextAcquire(t *testing.T) {
+	sem := newOpSemaphore(1)
+	release1, err := sem.acquire(context.Background())
+	if err != nil {
+		t.Fatalf("first acquire failed: %v", err)
+	}
+
+	acquired := make(chan struct{})
+	go func() {
+		release2, err := sem.acquire(context.Background())
+		if err != nil {
+			t.Errorf("second acquire failed: %v", err)
+			return
+		}
+		close(acquired)
+		release2()
+	}()
+
+	select {
+	case <-acquired:
+		t.Fatal("second acquire succeeded while first slot was held")
+	case <-time.After(50 * time.Millisecond):
+	}
+
+	release1()
+	select {
+	case <-acquired:
+	case <-time.After(time.Second):
+		t.Fatal("second acquire never succeeded after release")
+	}
+}
+
+func TestOpSemaphore_AcquireRespectsContextCancel(t *testing.T) {
+	sem := newOpSemaphore(1)
+	release, err := sem.acquire(context.Background())
+	if err != nil {
+		t.Fatalf("first acquire failed: %v", err)
+	}
+	defer release()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+	defer cancel()
+	_, err = sem.acquire(ctx)
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("acquire error = %v, want context deadline exceeded", err)
+	}
+}

cmd/mcp/serve.go

@@ -21,6 +21,9 @@ type ServeCmd struct {
 	ExecTimeoutDefault time.Duration
 	ExecTimeoutMax     time.Duration
 	ExecOutputCap      int
+	MaxConcurrentOps   int
+
+	opSem *opSemaphore
 }
 
 // NewServeCmd builds the `serve` subcommand.
@@ -54,14 +57,21 @@ func NewServeCmd(globalFlags *flags.GlobalFlags) *cobra.Command {
 			100*1024,
 			"Per-stream byte cap for workspace_exec output; excess is replaced with a truncation marker",
 		),
+		cliflags.Int(
+			&cmd.MaxConcurrentOps,
+			names.MaxConcurrentOps,
+			8,
+			"Maximum number of concurrent workspace_exec/workspace_create/workspace_start "+
+				"operations; excess calls wait for a free slot",
+		),
 	)
 	return cobraCmd
 }
 
 // Run wires up the MCP server and serves over stdio until ctx is cancelled.
 func (cmd *ServeCmd) Run(ctx context.Context) error {
-	log.Debugf("starting MCP server (timeout default=%s max=%s cap=%dB)",
-		cmd.ExecTimeoutDefault, cmd.ExecTimeoutMax, cmd.ExecOutputCap)
+	log.Debugf("starting MCP server (timeout default=%s max=%s cap=%dB maxops=%d)",
+		cmd.ExecTimeoutDefault, cmd.ExecTimeoutMax, cmd.ExecOutputCap, cmd.MaxConcurrentOps)
 
 	// Reserve real stdout for the JSON-RPC frame; redirect os.Stdout to stderr
 	// so any stray write elsewhere in the process can't corrupt the transport.
@@ -79,13 +89,14 @@ func (cmd *ServeCmd) Run(ctx context.Context) error {
 		Version: version.GetVersion(),
 	}, nil)
 
-	cmd.registerTools(server)
+	cmd.opSem = newOpSemaphore(cmd.MaxConcurrentOps)
+	cmd.registerTools(server, cmd.opSem)
 
 	return server.Run(ctx, transport)
 }
 
-func (cmd *ServeCmd) registerTools(s *sdkmcp.Server) {
-	registerWorkspaceTools(s, cmd.GlobalFlags)
-	registerExecTool(s, cmd)
+func (cmd *ServeCmd) registerTools(s *sdkmcp.Server, sem *opSemaphore) {
+	registerWorkspaceTools(s, cmd.GlobalFlags, sem)
+	registerExecTool(s, cmd, sem)
 	registerProviderTools(s, cmd.GlobalFlags)
 }

cmd/mcp/serve_test.go

@@ -2,6 +2,7 @@ package mcp
 
 import (
 	"context"
+	"strings"
 	"testing"
 	"time"
 
@@ -19,7 +20,7 @@ func TestServer_ListsAllTools(t *testing.T) {
 	server := sdkmcp.NewServer(&sdkmcp.Implementation{Name: "devsy-test", Version: "test"}, nil)
 	g := &flags.GlobalFlags{}
 	serveCmd := &ServeCmd{GlobalFlags: g, ExecOutputCap: 1024}
-	serveCmd.registerTools(server)
+	serveCmd.registerTools(server, newOpSemaphore(8))
 
 	clientTransport, serverTransport := sdkmcp.NewInMemoryTransports()
 
@@ -57,3 +58,82 @@ func TestServer_ListsAllTools(t *testing.T) {
 		t.Errorf("expected %d tools, got %d: %+v", len(wantNames), len(tools.Tools), have)
 	}
 }
+
+func TestServer_WorkspaceExecRespectsSemaphore(t *testing.T) {
+	home := t.TempDir()
+	t.Setenv("DEVSY_HOME", home)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	server := sdkmcp.NewServer(&sdkmcp.Implementation{Name: "devsy-test", Version: "test"}, nil)
+	g := &flags.GlobalFlags{}
+	serveCmd := &ServeCmd{GlobalFlags: g, ExecOutputCap: 1024, MaxConcurrentOps: 1}
+	sem := newOpSemaphore(serveCmd.MaxConcurrentOps)
+	serveCmd.registerTools(server, sem)
+
+	clientTransport, serverTransport := sdkmcp.NewInMemoryTransports()
+
+	serverErr := make(chan error, 1)
+	go func() {
+		serverErr <- server.Run(ctx, serverTransport)
+	}()
+
+	client := sdkmcp.NewClient(&sdkmcp.Implementation{Name: "test-client", Version: "0"}, nil)
+	session, err := client.Connect(ctx, clientTransport, nil)
+	if err != nil {
+		t.Fatalf("connect: %v", err)
+	}
+	t.Cleanup(func() { _ = session.Close() })
+
+	execArgs := map[string]any{
+		"name":    "some-workspace",
+		"command": []string{"echo", "hi"},
+	}
+
+	release, err := sem.acquire(context.Background())
+	if err != nil {
+		t.Fatalf("acquire: %v", err)
+	}
+
+	blockedCtx, blockedCancel := context.WithTimeout(ctx, 200*time.Millisecond)
+	defer blockedCancel()
+	_, callErr := session.CallTool(blockedCtx, &sdkmcp.CallToolParams{
+		Name:      "workspace_exec",
+		Arguments: execArgs,
+	})
+	if callErr == nil {
+		t.Fatal("expected workspace_exec call to fail while the only semaphore slot is held")
+	}
+
+	release()
+
+	res, callErr := session.CallTool(ctx, &sdkmcp.CallToolParams{
+		Name:      "workspace_exec",
+		Arguments: execArgs,
+	})
+	if callErr != nil {
+		t.Fatalf(
+			"expected workspace_exec call to reach the handler after release, got transport error: %v",
+			callErr,
+		)
+	}
+	assertNotSemaphoreError(t, res)
+}
+
+// A non-semaphore error (e.g. workspace not found) is expected here.
+func assertNotSemaphoreError(t *testing.T, res *sdkmcp.CallToolResult) {
+	t.Helper()
+	if !res.IsError {
+		return
+	}
+	var msg string
+	for _, c := range res.Content {
+		if tc, ok := c.(*sdkmcp.TextContent); ok {
+			msg = tc.Text
+		}
+	}
+	if strings.Contains(msg, "waiting for a free operation slot") {
+		t.Fatalf("workspace_exec failed due to the semaphore even after release: %s", msg)
+	}
+}

cmd/mcp/tools_exec.go

@@ -42,7 +42,7 @@ type execOutput struct {
 	Error *ErrorPayload `json:"error,omitempty"`
 }
 
-func registerExecTool(s *sdkmcp.Server, cmd *ServeCmd) {
+func registerExecTool(s *sdkmcp.Server, cmd *ServeCmd, sem *opSemaphore) {
 	sdkmcp.AddTool(s, &sdkmcp.Tool{
 		Name: "workspace_exec",
 		Description: "Run a one-shot command in a running workspace container. The " +
@@ -59,6 +59,13 @@ func registerExecTool(s *sdkmcp.Server, cmd *ServeCmd) {
 		if len(in.Command) == 0 {
 			return errorResult(fmt.Errorf("command is required")), execOutput{}, nil
 		}
+
+		release, err := sem.acquire(ctx)
+		if err != nil {
+			return errorResult(err), execOutput{}, nil
+		}
+		defer release()
+
 		stdout := NewBoundedBuffer(cmd.ExecOutputCap)
 		stderr := NewBoundedBuffer(cmd.ExecOutputCap)