Skip to content
Merged
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
42 changes: 32 additions & 10 deletions pkg/mcp/parallel/registry.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,9 @@ import (
"context"
"encoding/json"
"fmt"
"os"
"os/exec"
"strings"
"sync"

sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp"
Expand Down Expand Up @@ -150,18 +152,38 @@ func prodSpawner(spec UpstreamSpec) (*sdkmcp.ClientSession, func() error, error)
return sess, closeFn, nil
}

// buildUpstreamCommand assembles the *exec.Cmd for a given upstream. We
// populate cmd.Env explicitly from the spec so the broker process's own
// env doesn't leak into upstreams (each upstream's env block in mcp.json
// is the full intended environment for that subprocess).
// buildUpstreamCommand assembles the *exec.Cmd for a given upstream. The
// subprocess gets the broker's own environment with the spec's env block
// layered on top, which is the same shape Claude Code gives an MCP it
// launches directly: PATH, HOME and the rest come from the parent so an
// upstream that shells out (git fetch, kubectl) can find its binaries, and
// the spec wins on conflict so each upstream keeps its own telemetry tool
// prefix instead of the broker's. The broker's upstream registry blob is
// dropped: it is large, carries every upstream's secrets, and means nothing
// to a child.
func buildUpstreamCommand(spec UpstreamSpec) *exec.Cmd {
cmd := exec.Command(spec.Command, spec.Args...)
if len(spec.Env) > 0 {
env := make([]string, 0, len(spec.Env))
for k, v := range spec.Env {
env = append(env, k+"="+v)
cmd.Env = upstreamEnv(os.Environ(), spec.Env)
return cmd
}

func upstreamEnv(parent []string, overrides map[string]string) []string {
env := make([]string, 0, len(parent)+len(overrides))
for _, kv := range parent {
key, _, _ := strings.Cut(kv, "=")
if key == EnvUpstreams {
continue
}
if _, overridden := overrides[key]; overridden {
continue
}
cmd.Env = env
env = append(env, kv)
}
return cmd
for k, v := range overrides {
if k == EnvUpstreams {
continue
}
env = append(env, k+"="+v)
}
Comment thread
sourcehawk marked this conversation as resolved.
return env
}
32 changes: 32 additions & 0 deletions pkg/mcp/parallel/registry_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,11 @@ package parallel
import (
"context"
"encoding/json"
"strings"
"testing"

sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

Expand Down Expand Up @@ -90,3 +92,33 @@ func TestBuildUpstreamCommand_PopulatesArgsAndEnv(t *testing.T) {
require.Equal(t, []string{"/bin/echo", "a", "b"}, cmd.Args)
require.Contains(t, cmd.Env, "K=V")
}

// Upstreams shell out (the git MCP runs `git fetch`), so they need the
// broker's PATH, HOME and friends exactly as they would when Claude Code
// launches them directly. The spec's env block layers on top and wins on
// conflict, so an upstream keeps its own telemetry tool prefix rather than
// inheriting the broker's.
func TestBuildUpstreamCommand_InheritsBrokerEnvWithSpecOverride(t *testing.T) {
t.Setenv("PATH", "/parallel-test-bin")
t.Setenv("TRIAGENT_MCP_TELEMETRY_TOOL_PREFIX", "mcp__triagent-parallel__")
t.Setenv("TRIAGENT_MCP_PARALLEL_UPSTREAMS", "{}")

spec := UpstreamSpec{Command: "/bin/echo", Env: map[string]string{
"TRIAGENT_MCP_TELEMETRY_TOOL_PREFIX": "mcp__triagent-git-alerts__",
EnvUpstreams: "{\"from-spec\":{}}",
}}
cmd := buildUpstreamCommand(spec)
Comment thread
sourcehawk marked this conversation as resolved.

assert.Contains(t, cmd.Env, "PATH=/parallel-test-bin")
assert.Contains(t, cmd.Env, "TRIAGENT_MCP_TELEMETRY_TOOL_PREFIX=mcp__triagent-git-alerts__")
assert.NotContains(t, cmd.Env, "TRIAGENT_MCP_TELEMETRY_TOOL_PREFIX=mcp__triagent-parallel__")
for _, kv := range cmd.Env {
assert.False(t, strings.HasPrefix(kv, EnvUpstreams+"="), "upstream registry blob must not leak into upstreams from the broker env or the spec: %s", kv)
}
}

func TestBuildUpstreamCommand_EmptySpecEnvStillInheritsBrokerEnv(t *testing.T) {
t.Setenv("PATH", "/parallel-test-bin")
cmd := buildUpstreamCommand(UpstreamSpec{Command: "/bin/echo"})
assert.Contains(t, cmd.Env, "PATH=/parallel-test-bin")
}
Loading