Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 77 additions & 3 deletions internal/cli/mcp_config.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@ type mcpWritableConfig struct {
serverRaw map[string]json.RawMessage
}

var lockMCPConfigFile = config.LockFile

// projectMCPConfigExists reports whether the workspace's project ./.zero/config.json
// declares any MCP servers, so the trust notice fires only when project MCP config was
// actually skipped (mirroring projectHooksFileExists / projectPluginsDirExists). A
Expand All @@ -50,7 +52,7 @@ func projectMCPConfigExists(workspaceRoot string) bool {
return len(fc.MCP.Servers) > 0
}

func runMCPAdd(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int {
func runMCPAdd(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) (exitCode int) {
options, help, err := parseMCPAddArgs(args)
if err != nil {
return writeExecUsageError(stderr, err.Error())
Expand All @@ -66,6 +68,26 @@ func runMCPAdd(args []string, stdout io.Writer, stderr io.Writer, deps appDeps)
if err != nil {
return writeAppError(stderr, "failed to resolve user config: "+err.Error(), exitCrash)
}
// This edits the same user config document the config package mutates, with
// the same read-modify-write + rename shape, so it takes the same
// cross-process lock. Without it, an MCP config edit racing a provider or
// preference write would silently drop whichever landed first (issue #832).
unlock, err := lockMCPConfigFile(configPath)
if err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
released := false
defer func() {
if released {
return
}
// A failed release leaves the lock held for the rest of the process, so
// exiting success here would claim a state the next config write cannot
// reproduce. It must not mask a failure this command already reported.
if releaseErr := unlock(); releaseErr != nil && exitCode == exitSuccess {
exitCode = writeAppError(stderr, redaction.ErrorMessage(releaseErr, redaction.Options{}), exitCrash)
}
}()
Comment thread
coderabbitai[bot] marked this conversation as resolved.
cfg, err := readMCPWritableConfig(configPath)
if err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
Expand All @@ -80,6 +102,10 @@ func runMCPAdd(args []string, stdout io.Writer, stderr io.Writer, deps appDeps)
if err := writeMCPWritableConfig(configPath, cfg); err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
if err := unlock(); err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
released = true

if options.json {
payload := struct {
Expand Down Expand Up @@ -109,7 +135,7 @@ func runMCPAdd(args []string, stdout io.Writer, stderr io.Writer, deps appDeps)
return exitSuccess
}

func runMCPRemove(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int {
func runMCPRemove(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) (exitCode int) {
options, positional, help, err := parseMCPConfigPositionalCommand(args, "remove")
if err != nil {
return writeExecUsageError(stderr, err.Error())
Expand All @@ -132,6 +158,26 @@ func runMCPRemove(args []string, stdout io.Writer, stderr io.Writer, deps appDep
if err != nil {
return writeAppError(stderr, "failed to resolve user config: "+err.Error(), exitCrash)
}
// This edits the same user config document the config package mutates, with
// the same read-modify-write + rename shape, so it takes the same
// cross-process lock. Without it, an MCP config edit racing a provider or
// preference write would silently drop whichever landed first (issue #832).
unlock, err := lockMCPConfigFile(configPath)
if err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
released := false
defer func() {
if released {
return
}
// A failed release leaves the lock held for the rest of the process, so
// exiting success here would claim a state the next config write cannot
// reproduce. It must not mask a failure this command already reported.
if releaseErr := unlock(); releaseErr != nil && exitCode == exitSuccess {
exitCode = writeAppError(stderr, redaction.ErrorMessage(releaseErr, redaction.Options{}), exitCrash)
}
}()
cfg, err := readMCPWritableConfig(configPath)
if err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
Expand All @@ -145,6 +191,10 @@ func runMCPRemove(args []string, stdout io.Writer, stderr io.Writer, deps appDep
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
}
if err := unlock(); err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
released = true

if options.json {
payload := struct {
Expand All @@ -167,7 +217,7 @@ func runMCPRemove(args []string, stdout io.Writer, stderr io.Writer, deps appDep
return exitSuccess
}

func runMCPToggle(args []string, stdout io.Writer, stderr io.Writer, deps appDeps, disabled bool) int {
func runMCPToggle(args []string, stdout io.Writer, stderr io.Writer, deps appDeps, disabled bool) (exitCode int) {
commandName := "enable"
if disabled {
commandName = "disable"
Expand All @@ -194,6 +244,26 @@ func runMCPToggle(args []string, stdout io.Writer, stderr io.Writer, deps appDep
if err != nil {
return writeAppError(stderr, "failed to resolve user config: "+err.Error(), exitCrash)
}
// This edits the same user config document the config package mutates, with
// the same read-modify-write + rename shape, so it takes the same
// cross-process lock. Without it, an MCP config edit racing a provider or
// preference write would silently drop whichever landed first (issue #832).
unlock, err := lockMCPConfigFile(configPath)
if err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
released := false
defer func() {
if released {
return
}
// A failed release leaves the lock held for the rest of the process, so
// exiting success here would claim a state the next config write cannot
// reproduce. It must not mask a failure this command already reported.
if releaseErr := unlock(); releaseErr != nil && exitCode == exitSuccess {
exitCode = writeAppError(stderr, redaction.ErrorMessage(releaseErr, redaction.Options{}), exitCrash)
}
}()
cfg, err := readMCPWritableConfig(configPath)
if err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
Expand All @@ -210,6 +280,10 @@ func runMCPToggle(args []string, stdout io.Writer, stderr io.Writer, deps appDep
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
}
if err := unlock(); err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
released = true

if options.json {
payload := struct {
Expand Down
221 changes: 221 additions & 0 deletions internal/cli/mcp_config_lock_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,221 @@
package cli

import (
"bytes"
"errors"
"os"
"path/filepath"
"strings"
"testing"
"time"

"github.com/Gitlawb/zero/internal/config"
)

type mcpConfigLockCase struct {
name string
args []string
disabled bool
server bool
}

func mcpConfigLockCases() []mcpConfigLockCase {
return []mcpConfigLockCase{
{name: "add", args: []string{"mcp", "add", "docs", "--", "docs-mcp"}},
{name: "remove", args: []string{"mcp", "remove", "docs"}, server: true},
{name: "enable", args: []string{"mcp", "enable", "docs"}, server: true, disabled: true},
{name: "disable", args: []string{"mcp", "disable", "docs"}, server: true},
}
}

func seedMCPConfigLockCase(t *testing.T, path string, testCase mcpConfigLockCase) {
t.Helper()
servers := map[string]config.MCPServerConfig{}
if testCase.server {
servers["docs"] = config.MCPServerConfig{Type: "stdio", Command: "docs-mcp", Disabled: testCase.disabled}
}
writeMCPCommandConfig(t, path, config.FileConfig{
ActiveProvider: "seed",
Providers: []config.ProviderProfile{{Name: "seed", Model: "seed-model"}},
MCP: config.MCPConfig{Servers: servers},
})
}

func assertMCPConfigLockMutation(t *testing.T, path string, testCase mcpConfigLockCase) {
t.Helper()
cfg := readMCPCommandConfig(t, path)
server, exists := cfg.MCP.Servers["docs"]
switch testCase.name {
case "add":
if !exists {
t.Fatal("mcp add did not persist the server")
}
case "remove":
if exists {
t.Fatal("mcp remove left the server configured")
}
case "enable":
if !exists || server.Disabled {
t.Fatalf("mcp enable result = %+v, exists=%v", server, exists)
}
case "disable":
if !exists || !server.Disabled {
t.Fatalf("mcp disable result = %+v, exists=%v", server, exists)
}
}
if cfg.Preferences.Theme != "dracula" {
t.Errorf("theme update was lost: theme = %q, want dracula", cfg.Preferences.Theme)
}
if cfg.ActiveProvider != "seed" || len(cfg.Providers) != 1 || cfg.Providers[0].Name != "seed" {
t.Errorf("seeded provider was lost: active=%q providers=%#v", cfg.ActiveProvider, cfg.Providers)
}
}

// TestRunMCPConfigCommandsParticipateInConfigLock covers the half of issue #832
// that lives outside internal/config. MCP config commands read the SAME user
// config document, edit it, and republish it with the same temp-file+rename
// shape as the config package's mutators. Locking only the config package would
// leave these writers free to clobber a concurrent provider or preference
// update, and be clobbered by one, with the file still valid JSON afterwards.
//
// Racing the two writers and hoping to observe a lost update is unreliable —
// the interleaving that loses one is narrow, and the test passed consistently
// against the unlocked code. So this asserts the property that actually matters
// and is deterministic: while the config lock is held elsewhere, the MCP
// writer's update CANNOT land. Once released it completes, and both updates
// survive.
func TestRunMCPConfigCommandsParticipateInConfigLock(t *testing.T) {
for _, testCase := range mcpConfigLockCases() {
t.Run(testCase.name, func(t *testing.T) {
configPath := filepath.Join(t.TempDir(), "zero", "config.json")
seedMCPConfigLockCase(t, configPath, testCase)
before, err := os.ReadFile(configPath)
if err != nil {
t.Fatal(err)
}

unlock, err := config.LockFile(configPath)
if err != nil {
t.Fatalf("acquire config lock: %v", err)
}
released := false
release := func() {
if !released {
released = true
if err := unlock(); err != nil {
t.Fatalf("release config lock: %v", err)
}
}
}
defer release()

original := lockMCPConfigFile
lockAttempted := make(chan struct{})
lockMCPConfigFile = func(path string) (func() error, error) {
close(lockAttempted)
return original(path)
}
t.Cleanup(func() { lockMCPConfigFile = original })

var stdout, stderr bytes.Buffer
done := make(chan int, 1)
go func() {
done <- runWithDeps(testCase.args, &stdout, &stderr, appDeps{
userConfigPath: func() (string, error) { return configPath, nil },
})
}()

select {
case <-lockAttempted:
case <-time.After(30 * time.Second):
t.Fatalf("mcp %s did not attempt config lock acquisition", testCase.name)
}

for range 20 {
after, err := os.ReadFile(configPath)
if err != nil {
t.Fatal(err)
}
if !bytes.Equal(after, before) {
t.Fatalf("mcp %s changed config while its lock was held", testCase.name)
}
select {
case exitCode := <-done:
t.Fatalf("mcp %s completed (exit %d) while its lock was held", testCase.name, exitCode)
default:
}
time.Sleep(5 * time.Millisecond)
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

release()
select {
case exitCode := <-done:
if exitCode != exitSuccess {
t.Fatalf("mcp %s exitCode = %d stderr=%s", testCase.name, exitCode, stderr.String())
}
case <-time.After(30 * time.Second):
t.Fatalf("mcp %s did not finish after the config lock was released", testCase.name)
}

if _, err := config.SetTheme(configPath, "dracula"); err != nil {
t.Fatalf("SetTheme: %v", err)
}
assertMCPConfigLockMutation(t, configPath, testCase)
})
}
}

func TestRunMCPConfigCommandsReportLockAcquisitionFailure(t *testing.T) {
sentinel := errors.New("injected lock acquisition failure")
original := lockMCPConfigFile
lockMCPConfigFile = func(string) (func() error, error) { return nil, sentinel }
t.Cleanup(func() { lockMCPConfigFile = original })

for _, testCase := range mcpConfigLockCases() {
t.Run(testCase.name, func(t *testing.T) {
configPath := filepath.Join(t.TempDir(), "zero", "config.json")
seedMCPConfigLockCase(t, configPath, testCase)
before, err := os.ReadFile(configPath)
if err != nil {
t.Fatal(err)
}
var stdout, stderr bytes.Buffer
exitCode := runWithDeps(testCase.args, &stdout, &stderr, appDeps{
userConfigPath: func() (string, error) { return configPath, nil },
})
if exitCode != exitCrash || stdout.Len() != 0 || !strings.Contains(stderr.String(), sentinel.Error()) {
t.Fatalf("exit=%d stdout=%q stderr=%q, want lock error without success output", exitCode, stdout.String(), stderr.String())
}
after, err := os.ReadFile(configPath)
if err != nil {
t.Fatal(err)
}
if !bytes.Equal(after, before) {
t.Fatal("command changed config after lock acquisition failed")
}
})
}
}

func TestRunMCPConfigCommandsReleaseBeforeSuccessOutput(t *testing.T) {
sentinel := errors.New("injected lock release failure")
original := lockMCPConfigFile
lockMCPConfigFile = func(string) (func() error, error) {
return func() error { return sentinel }, nil
}
t.Cleanup(func() { lockMCPConfigFile = original })

for _, testCase := range mcpConfigLockCases() {
t.Run(testCase.name, func(t *testing.T) {
configPath := filepath.Join(t.TempDir(), "zero", "config.json")
seedMCPConfigLockCase(t, configPath, testCase)
var stdout, stderr bytes.Buffer
exitCode := runWithDeps(testCase.args, &stdout, &stderr, appDeps{
userConfigPath: func() (string, error) { return configPath, nil },
})
if exitCode != exitCrash || stdout.Len() != 0 || !strings.Contains(stderr.String(), sentinel.Error()) {
t.Fatalf("exit=%d stdout=%q stderr=%q, want unlock error before success output", exitCode, stdout.String(), stderr.String())
}
})
}
}
Loading
Loading