diff --git a/internal/mcp/idle_env_test.go b/internal/mcp/idle_env_test.go new file mode 100644 index 0000000..a0e2ad4 --- /dev/null +++ b/internal/mcp/idle_env_test.go @@ -0,0 +1,38 @@ +package mcp + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func TestIdleTimeoutFromEnv(t *testing.T) { + // The idle shutdown is off unless asked for; anything that is not a positive duration leaves it off. + for value, want := range map[string]time.Duration{ + "": 0, + "0": 0, + "abc": 0, + "-5m": 0, + "45": 0, // a bare number is not a duration + "90m": 90 * time.Minute, + "2h": 2 * time.Hour, + "30s": 30 * time.Second, + " 2h ": 0, // not trimmed: the operator wrote it, so it is wrong rather than guessed + } { + t.Run(value, func(t *testing.T) { + t.Setenv(IdleTimeoutEnv, value) + assert.Equal(t, want, idleTimeoutFromEnv()) + }) + } +} + +func TestNewServer_HasNoIdleTimeoutByDefault(t *testing.T) { + t.Setenv(IdleTimeoutEnv, "") + assert.Zero(t, NewServer(nil, "http://localhost:1", "p", "v").idleTimeout) +} + +func TestNewServer_TakesTheIdleTimeoutFromTheEnvironment(t *testing.T) { + t.Setenv(IdleTimeoutEnv, "3h") + assert.Equal(t, 3*time.Hour, NewServer(nil, "http://localhost:1", "p", "v").idleTimeout) +} diff --git a/internal/mcp/idle_test.go b/internal/mcp/idle_test.go new file mode 100644 index 0000000..231aedf --- /dev/null +++ b/internal/mcp/idle_test.go @@ -0,0 +1,97 @@ +//go:build unix + +package mcp + +import ( + "bytes" + "context" + "os" + "sync" + "syscall" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// blockingPipe returns the two ends of a pipe in blocking mode, the way an inherited stdin is: a Read on it cannot be +// woken by closing the file (the standard pipe from os.Pipe is pollable and can, which hid this). +func blockingPipe(t *testing.T) (r, w *os.File) { + t.Helper() + var fds [2]int + require.NoError(t, syscall.Pipe(fds[:])) + r = os.NewFile(uintptr(fds[0]), "stdin-like") + w = os.NewFile(uintptr(fds[1]), "stdin-writer") + t.Cleanup(func() { _ = w.Close(); _ = r.Close() }) + return r, w +} + +type lockedBuffer struct { + buf bytes.Buffer + mu sync.Mutex +} + +func (b *lockedBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Write(p) +} + +func (b *lockedBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.String() +} + +// The idle timeout used to close stdin and hope the read loop ended. On a real (blocking) stdin the read stays blocked, +// so the server logged "shutting down" and then lived on with a closed stdin; the next message from the client woke the +// read, which failed with "file already closed" and killed the connection mid-call (seen in Claude Desktop after the +// connector sat unused for more than 30 minutes). An idle timeout that is on must end Run for real. +func TestRun_IdleTimeoutEndsTheServerCleanly(t *testing.T) { + t.Parallel() + + stdin, _ := blockingPipe(t) + server := &Server{stdin: stdin, stdout: &lockedBuffer{}, idleTimeout: 50 * time.Millisecond, monitorInterval: 10 * time.Millisecond} + + done := make(chan error, 1) + go func() { done <- server.Run(context.Background()) }() + + select { + case err := <-done: + require.NoError(t, err, "an idle shutdown is a clean exit, not an error") + case <-time.After(3 * time.Second): + t.Fatal("Run did not return after the idle timeout: stdin was closed but the blocked read never ended") + } +} + +// By default there is no idle timeout: a stdio server lives as long as its client, and a client that talks to it after a +// long pause must get an answer. +func TestRun_WithoutAnIdleTimeoutTheServerKeepsAnsweringAfterALongPause(t *testing.T) { + t.Parallel() + + stdin, stdinWriter := blockingPipe(t) + stdout := &lockedBuffer{} + server := &Server{stdin: stdin, stdout: stdout, monitorInterval: 5 * time.Millisecond} + + done := make(chan error, 1) + go func() { done <- server.Run(context.Background()) }() + + time.Sleep(300 * time.Millisecond) // many monitor ticks with no traffic + select { + case err := <-done: + t.Fatalf("Run ended during a pause with no traffic: %v", err) + default: + } + + _, err := stdinWriter.WriteString(`{"jsonrpc":"2.0","id":7,"method":"tools/list"}` + "\n") + require.NoError(t, err) + + assert.Eventually(t, func() bool { return bytes.Contains([]byte(stdout.String()), []byte(`"id":7`)) }, 3*time.Second, 10*time.Millisecond, + "the request after the pause is answered") + select { + case err := <-done: + t.Fatalf("Run ended after answering: %v", err) + default: + } +} diff --git a/internal/mcp/server.go b/internal/mcp/server.go index 6778600..2904d04 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -30,6 +30,11 @@ type Server struct { version string writeMu sync.Mutex lastActivity atomic.Int64 + // idleTimeout ends the server after this long without a message; 0 (the default) never does. A stdio server lives + // as long as its client: ending it while the client is still connected makes the client's next call fail. + idleTimeout time.Duration + // monitorInterval is how often the parent and idle checks run (30s when 0). + monitorInterval time.Duration } // NewServer creates a new MCP server that proxies to the worker HTTP API. @@ -41,9 +46,24 @@ func NewServer(client *http.Client, workerURL, project, version string) *Server version: version, stdin: os.Stdin, stdout: os.Stdout, + + idleTimeout: idleTimeoutFromEnv(), } } +// IdleTimeoutEnv names the environment variable that turns the idle shutdown on, as a Go duration such as "2h". +// It is off by default. +const IdleTimeoutEnv = "CLAUDE_MNEMONIC_MCP_IDLE_TIMEOUT" + +// idleTimeoutFromEnv reads IdleTimeoutEnv; an unset, invalid or non-positive value means no idle shutdown. +func idleTimeoutFromEnv() time.Duration { + d, err := time.ParseDuration(os.Getenv(IdleTimeoutEnv)) + if err != nil || d <= 0 { + return 0 + } + return d +} + // Request represents a JSON-RPC request. type Request struct { JSONRPC string `json:"jsonrpc"` @@ -82,6 +102,13 @@ type Tool struct { // Run starts the MCP server loop. func (s *Server) Run(ctx context.Context) error { + // An idle shutdown cancels this context. Closing stdin does not do it: on a real (blocking) stdin the read stays + // blocked, the server lived on with a closed stdin, and the client's next message woke the read with "file already + // closed", which killed the connection mid-call. + ctx, cancel := context.WithCancel(ctx) + defer cancel() + var idleStopped atomic.Bool + scanner := bufio.NewScanner(s.stdin) buf := make([]byte, 0, 64*1024) scanner.Buffer(buf, 1024*1024) // 1MB max message size @@ -100,11 +127,15 @@ func (s *Server) Run(ctx context.Context) error { scanErr <- scanner.Err() }() - // Monitor parent process liveness and idle timeout. - // If the parent dies (ppid changes) or no messages arrive for 30 minutes, shut down. + // Monitor parent process liveness and, when one is set, the idle timeout. + // If the parent dies (ppid changes) or no messages arrive for idleTimeout, shut down. parentPID := os.Getppid() go func() { - ticker := time.NewTicker(30 * time.Second) + interval := s.monitorInterval + if interval <= 0 { + interval = 30 * time.Second + } + ticker := time.NewTicker(interval) defer ticker.Stop() for { select { @@ -118,11 +149,10 @@ func (s *Server) Run(ctx context.Context) error { } return } - if time.Since(time.Unix(s.lastActivity.Load(), 0)) > 30*time.Minute { - log.Info().Msg("MCP server idle timeout (30m), shutting down") - if closer, ok := s.stdin.(io.Closer); ok { - _ = closer.Close() - } + if s.idleTimeout > 0 && time.Since(time.Unix(s.lastActivity.Load(), 0)) > s.idleTimeout { + log.Info().Dur("idle", s.idleTimeout).Msg("MCP server idle timeout, shutting down") + idleStopped.Store(true) + cancel() return } } @@ -140,6 +170,9 @@ func (s *Server) Run(ctx context.Context) error { case <-ctx.Done(): // Drain in-flight requests before returning. wg.Wait() + if idleStopped.Load() { + return nil // an idle shutdown the operator asked for is a clean exit + } return ctx.Err() case line, ok := <-lines: if !ok {