From 166ca09119e27651f10f77b44be4357ac21ca130 Mon Sep 17 00:00:00 2001 From: Anthony Torns II Date: Tue, 20 Oct 2026 14:51:16 -0700 Subject: [PATCH 1/2] feat(setup): send the agent's model calls to a LaunchDarkly Bedrock relay For a logged-in user, OpenCode now calls only the model in the setup agent config, through the loopback proxy and the gonfalon model relay at /internal/ldcli-setup-agent/bedrock. No model key is on the user's machine. The OpenCode configuration turns on only the amazon-bedrock provider, points it at the proxy, and uses the config's model for title generation too. OpenCode gets a placeholder bearer token, and the user's AWS variables do not reach it. The proxy gives OpenCode its own secret path, which reaches only POST /model/... on the model relay, and it replaces the placeholder with the access token. ldcli now reads the config before it starts OpenCode. If the config is off or does not name a Bedrock model, setup stops and does not use the user's own provider. Co-Authored-By: Claude Opus 5.5 --- internal/setup/agent/agent_test.go | 2 +- internal/setup/agent/aiconfig.go | 22 ++++ internal/setup/agent/bedrock_test.go | 155 ++++++++++++++++++++++++ internal/setup/agent/opencode.go | 19 ++- internal/setup/agent/opencode_config.go | 45 ++++++- internal/setup/agent/relay_proxy.go | 58 +++++++-- internal/setup/agent/sandbox_test.go | 2 +- internal/setup/agent/server.go | 25 +++- 8 files changed, 305 insertions(+), 23 deletions(-) create mode 100644 internal/setup/agent/bedrock_test.go diff --git a/internal/setup/agent/agent_test.go b/internal/setup/agent/agent_test.go index a2341075..9a1731a6 100644 --- a/internal/setup/agent/agent_test.go +++ b/internal/setup/agent/agent_test.go @@ -70,7 +70,7 @@ func TestChildEnv_RemovesLDVariablesAndSetsOwn(t *testing.T) { } func TestOpenCodeConfig_RunsCommandsWithoutApprovals(t *testing.T) { - data, err := openCodeConfig() + data, err := openCodeConfig(nil) require.NoError(t, err) var cfg struct { diff --git a/internal/setup/agent/aiconfig.go b/internal/setup/agent/aiconfig.go index 5977df83..327e6651 100644 --- a/internal/setup/agent/aiconfig.go +++ b/internal/setup/agent/aiconfig.go @@ -3,6 +3,7 @@ package agent import ( "context" "encoding/json" + "errors" "fmt" "io" "log/slog" @@ -210,6 +211,27 @@ func (s *AIConfigSource) Resolve(ctx context.Context, task Task) (Config, error) return configFromAI(res.Config, task), nil } +// managedModel resolves the model for a run that LaunchDarkly pays for. It +// does not fall back: the config must be on and must name a Bedrock model. +func (s *AIConfigSource) managedModel(ctx context.Context, task Task) (*managedModel, error) { + if err := s.Ready(ctx); err != nil { + return nil, err + } + inspect := s.inspect + if inspect == nil { + inspect = aiserver.InspectConfig + } + res := inspect(ctx, s.key(), setupContext(task)) + if !res.Enabled || res.Config == nil { + return nil, errors.New("the setup agent configuration is off") + } + provider, modelID := openCodeModel(*res.Config) + if provider != bedrockProviderID { + return nil, fmt.Errorf("the setup agent configuration must use a Bedrock model, not %q", res.Config.Provider.Name+"/"+res.Config.Model.Name) + } + return &managedModel{Endpoint: s.proxy.bedrockURL, ModelID: modelID}, nil +} + func (s *AIConfigSource) fallback(ctx context.Context, task Task) (Config, error) { if s.Fallback == nil { return DefaultConfigSource{}.Resolve(ctx, task) diff --git a/internal/setup/agent/bedrock_test.go b/internal/setup/agent/bedrock_test.go new file mode 100644 index 00000000..91d0f0a2 --- /dev/null +++ b/internal/setup/agent/bedrock_test.go @@ -0,0 +1,155 @@ +package agent + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os/exec" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/launchdarkly/go-ai-sdk/aiserver" +) + +func TestOpenCodeConfig_ManagedModelIsTheOnlyModel(t *testing.T) { + data, err := openCodeConfig(&managedModel{Endpoint: "http://127.0.0.1:1/secret", ModelID: "anthropic.claude-sonnet-4-6"}) + require.NoError(t, err) + var cfg map[string]any + require.NoError(t, json.Unmarshal(data, &cfg)) + + assert.Equal(t, []any{"amazon-bedrock"}, cfg["enabled_providers"], "the user's own providers are off") + assert.Equal(t, "amazon-bedrock/anthropic.claude-sonnet-4-6", cfg["model"]) + assert.Equal(t, "amazon-bedrock/anthropic.claude-sonnet-4-6", cfg["small_model"], "title generation uses the same model") + options := cfg["provider"].(map[string]any)["amazon-bedrock"].(map[string]any)["options"].(map[string]any) + assert.Equal(t, "http://127.0.0.1:1/secret", options["endpoint"]) + assert.Equal(t, "us-east-1", options["region"], "a US region makes OpenCode use the US inference profile") + assert.Contains(t, cfg, "permission", "the permission rules stay") +} + +func TestOpenCodeConfig_WithoutAManagedModelHasNoProviderSettings(t *testing.T) { + data, err := openCodeConfig(nil) + require.NoError(t, err) + var cfg map[string]any + require.NoError(t, json.Unmarshal(data, &cfg)) + + for _, key := range []string{"enabled_providers", "model", "small_model", "provider"} { + assert.NotContains(t, cfg, key) + } +} + +func TestRelayProxy_BedrockRouteReachesOnlyTheModelRelay(t *testing.T) { + relay, got := fakeRelay(t) + p := startTestProxy(t, relay.URL) + + req, err := http.NewRequest(http.MethodPost, p.bedrockURL+"/model/us.anthropic.claude-sonnet-4-6/converse-stream", strings.NewReader(`{"messages":[]}`)) + require.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+RelayPlaceholderKey) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + resp.Body.Close() + + assert.Equal(t, BedrockRelayPath+"/model/us.anthropic.claude-sonnet-4-6/converse-stream", got.URL.Path) + assert.Equal(t, "api-secret-token", got.Header.Get(RelayTokenHeader)) + assert.Empty(t, got.Header.Get("Authorization"), "the placeholder bearer token does not reach the relay") + assert.NotEqual(t, p.baseURL, p.bedrockURL, "the model route has its own secret path") + + for _, tc := range []struct{ method, path string }{ + {http.MethodGet, "/model/us.anthropic.claude-sonnet-4-6/converse-stream"}, + {http.MethodPost, "/polling/sdk/latest-all"}, + {http.MethodPost, ""}, + } { + req, err := http.NewRequest(tc.method, p.bedrockURL+tc.path, nil) + require.NoError(t, err) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + resp.Body.Close() + assert.Equal(t, http.StatusNotFound, resp.StatusCode, "%s %s", tc.method, tc.path) + } +} + +func TestWithoutPrefix_RemovesTheUsersAWSVariables(t *testing.T) { + got := withoutPrefix([]string{"AWS_PROFILE=work", "AWS_ACCESS_KEY_ID=AKIA", "HOME=/Users/me", "PATH=/bin"}, "AWS_") + + assert.Equal(t, []string{"HOME=/Users/me", "PATH=/bin"}, got) +} + +func TestAIConfigSource_ManagedModel(t *testing.T) { + s, _ := fakeSource(nil, enabled(agentConfig("Bedrock", "anthropic.claude-sonnet-4-6", "Wire the call."))) + defer func() { _ = s.Close(context.Background()) }() + + m, err := s.managedModel(context.Background(), Task{SDKID: "python-server-sdk"}) + + require.NoError(t, err) + assert.Equal(t, "anthropic.claude-sonnet-4-6", m.ModelID) + assert.Equal(t, s.proxy.bedrockURL, m.Endpoint) +} + +func TestAIConfigSource_ManagedModel_Refusals(t *testing.T) { + for name, res := range map[string]aiserver.InspectResult{ + "another provider": enabled(agentConfig("Anthropic", "claude-sonnet-5", "Wire the call.")), + "config off": {Enabled: false}, + } { + t.Run(name, func(t *testing.T) { + s, _ := fakeSource(nil, res) + defer func() { _ = s.Close(context.Background()) }() + + _, err := s.managedModel(context.Background(), Task{}) + + assert.Error(t, err, "setup must not fall back to the user's own provider") + }) + } +} + +// TestBedrockRoute_RealOpenCode starts the real OpenCode, in the sandbox on +// macOS, with the managed model configuration. It checks that OpenCode sends +// its model calls through the proxy to the model relay, with the US +// inference profile and without the placeholder bearer token. +func TestBedrockRoute_RealOpenCode(t *testing.T) { + if _, err := exec.LookPath("opencode"); err != nil { + t.Skip("OpenCode is not installed") + } + if testing.Short() { + t.Skip("starts OpenCode") + } + var mu sync.Mutex + var paths []string + relay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + mu.Lock() + paths = append(paths, r.Method+" "+r.URL.Path+" token="+r.Header.Get(RelayTokenHeader)+" auth="+r.Header.Get("Authorization")) + mu.Unlock() + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + _, _ = io.WriteString(w, `{"message":"fake relay"}`) + })) + defer relay.Close() + p := startTestProxy(t, relay.URL) + + config, err := openCodeConfig(&managedModel{Endpoint: p.bedrockURL, ModelID: "anthropic.claude-sonnet-4-6"}) + require.NoError(t, err) + srv, err := startServer(context.Background(), "opencode", t.TempDir(), config, map[string]string{bedrockTokenEnv: RelayPlaceholderKey}) + require.NoError(t, err) + defer srv.stop() + run, err := startRun(context.Background(), srv.client, Config{ProviderID: bedrockProviderID, ModelID: "anthropic.claude-sonnet-4-6", Instructions: "Say hi."}, "Say hi.") + require.NoError(t, err) + defer run.Close() + + require.Eventually(t, func() bool { + mu.Lock() + defer mu.Unlock() + return len(paths) > 0 + }, 60*time.Second, 200*time.Millisecond, "OpenCode calls the model relay") + time.Sleep(2 * time.Second) + mu.Lock() + defer mu.Unlock() + for _, got := range paths { + assert.Equal(t, "POST "+BedrockRelayPath+"/model/us.anthropic.claude-sonnet-4-6/converse-stream token=api-secret-token auth=", got) + } +} diff --git a/internal/setup/agent/opencode.go b/internal/setup/agent/opencode.go index 228736e8..9ca4e073 100644 --- a/internal/setup/agent/opencode.go +++ b/internal/setup/agent/opencode.go @@ -75,7 +75,10 @@ func (o *OpenCode) Check(ctx context.Context) error { } func (o *OpenCode) Start(ctx context.Context, task Task) (Run, error) { - if o.SDK != nil && o.SDK.Ready(ctx) == nil { + // A logged-in user's agent uses the model that LaunchDarkly runs. If that + // model is not available, setup stops, and it does not use the user's + // own OpenCode provider. + if o.SDK != nil { return o.startWithSDK(ctx, task) } source := o.Config @@ -86,11 +89,11 @@ func (o *OpenCode) Start(ctx context.Context, task Task) (Run, error) { if err != nil { return nil, fmt.Errorf("could not get the setup agent configuration: %w", err) } - config, err := openCodeConfig() + config, err := openCodeConfig(nil) if err != nil { return nil, err } - srv, err := startServer(ctx, o.binary(), task.Dir, config) + srv, err := startServer(ctx, o.binary(), task.Dir, config, nil) if err != nil { return nil, err } @@ -117,11 +120,17 @@ func (o *OpenCode) Start(ctx context.Context, task Task) (Run, error) { // example because the config is off or is not in agent mode, the run uses the // built-in instructions and task. func (o *OpenCode) startWithSDK(ctx context.Context, task Task) (Run, error) { - config, err := openCodeConfig() + // OpenCode reads its configuration once, at start, so the model must be + // known before OpenCode starts. + model, err := o.SDK.managedModel(ctx, task) + if err != nil { + return nil, fmt.Errorf("could not get the LaunchDarkly model for the setup agent: %w", err) + } + config, err := openCodeConfig(model) if err != nil { return nil, err } - srv, err := startServer(ctx, o.binary(), task.Dir, config) + srv, err := startServer(ctx, o.binary(), task.Dir, config, map[string]string{bedrockTokenEnv: RelayPlaceholderKey}) if err != nil { return nil, err } diff --git a/internal/setup/agent/opencode_config.go b/internal/setup/agent/opencode_config.go index 70d34885..d01305f8 100644 --- a/internal/setup/agent/opencode_config.go +++ b/internal/setup/agent/opencode_config.go @@ -14,12 +14,37 @@ import ( // tool, when it needs information such as where a key is stored. Rules deny // files outside the project, web access, sub-agents, and a few commands that // the agent must never run, such as git push. -func openCodeConfig() ([]byte, error) { +// OpenCode's Bedrock provider settings for a managed model. +const ( + bedrockProviderID = "amazon-bedrock" + // bedrockRegion is a US region, so OpenCode uses the US inference profile + // (us.anthropic...). The relay decides where the call really runs. + bedrockRegion = "us-east-1" + // bedrockTokenEnv is OpenCode's Bedrock bearer token. It takes priority + // over every AWS credential, so OpenCode never signs with the user's AWS + // keys. ldcli sets it to RelayPlaceholderKey, which the proxy removes. + bedrockTokenEnv = "AWS_BEARER_TOKEN_BEDROCK" +) + +// managedModel is a model that LaunchDarkly runs on Bedrock for the agent. +// OpenCode calls it through the loopback proxy and the gonfalon model relay, +// so no model key is on the user's machine. +type managedModel struct { + // Endpoint is the proxy address for OpenCode's Bedrock provider. + Endpoint string + // ModelID is the Bedrock model ID from the config, such as + // anthropic.claude-sonnet-4-6. + ModelID string +} + +// openCodeConfig returns the configuration. If model is set, Bedrock is the +// only provider, and every call, including title generation, uses that model. +func openCodeConfig(model *managedModel) ([]byte, error) { bash, err := orderedRules(bashRules) if err != nil { return nil, err } - return json.Marshal(map[string]any{ + cfg := map[string]any{ "$schema": "https://opencode.ai/config.json", "permission": map[string]any{ "read": "allow", @@ -33,7 +58,21 @@ func openCodeConfig() ([]byte, error) { "websearch": "deny", "task": "deny", }, - }) + } + if model != nil { + id := bedrockProviderID + "/" + model.ModelID + cfg["enabled_providers"] = []string{bedrockProviderID} + cfg["model"] = id + // OpenCode's default small model is a global inference profile, which + // the relay refuses and which can run outside the US. + cfg["small_model"] = id + cfg["provider"] = map[string]any{ + bedrockProviderID: map[string]any{ + "options": map[string]any{"region": bedrockRegion, "endpoint": model.Endpoint}, + }, + } + } + return json.Marshal(cfg) } // bashRules are the shell command rules, in order. OpenCode applies the last diff --git a/internal/setup/agent/relay_proxy.go b/internal/setup/agent/relay_proxy.go index e620d287..0d084aa8 100644 --- a/internal/setup/agent/relay_proxy.go +++ b/internal/setup/agent/relay_proxy.go @@ -24,8 +24,13 @@ const ( // that prefix, because its private auth headers use it. RelayTokenHeader = "LD-CLI-Token" // RelayPlaceholderKey is the SDK key that the CLI gives the AI SDK. The - // relay replaces it with the real key on the way to catfood. + // relay replaces it with the real key on the way to catfood. OpenCode + // also gets it as its Bedrock bearer token, which the proxy removes. RelayPlaceholderKey = "ldcli-placeholder" + // BedrockRelayPath is the path of the model relay in gonfalon, under the + // user's base URI. It takes Bedrock Runtime requests and signs them with + // LaunchDarkly's AWS credentials. + BedrockRelayPath = "/internal/ldcli-setup-agent/bedrock" ) // relayProxy is a loopback HTTP server between the AI SDK and the relay. It @@ -34,14 +39,18 @@ const ( // add a header to its trace exports, which is the other reason for the proxy. // // The proxy accepts requests only under a random path, so another process on -// the machine, such as OpenCode, cannot use it to call the relay with the -// user's token. +// the machine cannot use it to call the relay with the user's token. +// OpenCode gets a second random path for its model calls. That path reaches +// only the model relay, so OpenCode cannot reach the SDK routes. type relayProxy struct { server *http.Server listener net.Listener // baseURL is the proxy address with the random path, for example // http://127.0.0.1:52011/3f9a.... baseURL string + // bedrockURL is the address for OpenCode's Bedrock provider, with its own + // random path. + bedrockURL string } // startRelayProxy starts the proxy. relayBase is the user's base URI, for @@ -51,24 +60,37 @@ func startRelayProxy(relayBase, accessToken string) (*relayProxy, error) { if err != nil || target.Scheme == "" || target.Host == "" { return nil, fmt.Errorf("the base URI %q is not a valid URL", relayBase) } + modelTarget := *target + modelTarget.Path = strings.TrimSuffix(target.Path, RelayPath) + BedrockRelayPath secret, err := randomHex(16) if err != nil { return nil, err } + bedrockSecret, err := randomHex(16) + if err != nil { + return nil, err + } listener, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { return nil, fmt.Errorf("could not start the relay proxy: %w", err) } prefix := "/" + secret + bedrockPrefix := "/" + bedrockSecret proxy := &httputil.ReverseProxy{ Rewrite: func(r *httputil.ProxyRequest) { - rest := strings.TrimPrefix(r.In.URL.Path, prefix) - r.Out.URL.Scheme = target.Scheme - r.Out.URL.Host = target.Host - r.Out.URL.Path = target.Path + rest + to, rest := target, strings.TrimPrefix(r.In.URL.Path, prefix) + if hasPathPrefix(r.In.URL.Path, bedrockPrefix) { + // OpenCode sends the placeholder as a bearer token. The relay + // signs the request itself, so the token must not reach it. + to, rest = &modelTarget, strings.TrimPrefix(r.In.URL.Path, bedrockPrefix) + r.Out.Header.Del("Authorization") + } + r.Out.URL.Scheme = to.Scheme + r.Out.URL.Host = to.Host + r.Out.URL.Path = to.Path + rest r.Out.URL.RawPath = "" - r.Out.Host = target.Host + r.Out.Host = to.Host r.Out.Header.Set(RelayTokenHeader, accessToken) }, // Send each part of a response as it arrives. The flag stream is a @@ -79,7 +101,15 @@ func startRelayProxy(relayBase, accessToken string) (*relayProxy, error) { ErrorLog: log.New(io.Discard, "", 0), } handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != prefix && !strings.HasPrefix(r.URL.Path, prefix+"/") { + switch { + case hasPathPrefix(r.URL.Path, prefix): + case hasPathPrefix(r.URL.Path, bedrockPrefix): + // OpenCode calls only the Bedrock model routes, with POST. + if r.Method != http.MethodPost || !strings.HasPrefix(r.URL.Path, bedrockPrefix+"/model/") { + http.NotFound(w, r) + return + } + default: http.NotFound(w, r) return } @@ -92,8 +122,9 @@ func startRelayProxy(relayBase, accessToken string) (*relayProxy, error) { ReadHeaderTimeout: 10 * time.Second, ErrorLog: log.New(io.Discard, "", 0), }, - listener: listener, - baseURL: "http://" + listener.Addr().String() + prefix, + listener: listener, + baseURL: "http://" + listener.Addr().String() + prefix, + bedrockURL: "http://" + listener.Addr().String() + bedrockPrefix, } go func() { if err := p.server.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) { @@ -103,6 +134,11 @@ func startRelayProxy(relayBase, accessToken string) (*relayProxy, error) { return p, nil } +// hasPathPrefix reports whether path is prefix or is under it. +func hasPathPrefix(path, prefix string) bool { + return path == prefix || strings.HasPrefix(path, prefix+"/") +} + // close stops the proxy and ends open connections, including the flag // stream. func (p *relayProxy) close(ctx context.Context) error { diff --git a/internal/setup/agent/sandbox_test.go b/internal/setup/agent/sandbox_test.go index 68f2601f..08f990be 100644 --- a/internal/setup/agent/sandbox_test.go +++ b/internal/setup/agent/sandbox_test.go @@ -62,7 +62,7 @@ func TestInstallDir(t *testing.T) { } func TestOpenCodeConfig_DeniesLDCLIAndSkillsInstalls(t *testing.T) { - data, err := openCodeConfig() + data, err := openCodeConfig(nil) require.NoError(t, err) assert.Contains(t, string(data), `"ldcli*":"deny"`) diff --git a/internal/setup/agent/server.go b/internal/setup/agent/server.go index 581a972d..07a4b524 100644 --- a/internal/setup/agent/server.go +++ b/internal/setup/agent/server.go @@ -36,7 +36,9 @@ type server struct { // startServer starts OpenCode in dir and waits until its API answers. The // server listens on the loopback interface only, and it requires a password // that is random for each run. -func startServer(ctx context.Context, binary, dir string, config []byte) (*server, error) { +// env holds extra variables for OpenCode. If it sets the Bedrock bearer +// token, the user's own AWS variables do not reach OpenCode. +func startServer(ctx context.Context, binary, dir string, config []byte, env map[string]string) (*server, error) { port, err := freePort() if err != nil { return nil, err @@ -53,6 +55,13 @@ func startServer(ctx context.Context, binary, dir string, config []byte) (*serve "OPENCODE_SERVER_USERNAME": serverUsername, "OPENCODE_CONFIG_CONTENT": string(config), } + for name, value := range env { + set[name] = value + } + parent := os.Environ() + if _, ok := env[bedrockTokenEnv]; ok { + parent = withoutPrefix(parent, "AWS_") + } box, err := sandboxCommand(binary, []string{"serve", "--hostname", "127.0.0.1", "--port", strconv.Itoa(port)}, dir, os.Environ()) if err != nil { return nil, err @@ -68,7 +77,7 @@ func startServer(ctx context.Context, binary, dir string, config []byte) (*serve procCtx, cancel := context.WithCancel(context.Background()) cmd := exec.CommandContext(procCtx, box.Name, box.Args...) cmd.Dir = dir - cmd.Env = childEnv(os.Environ(), set) + cmd.Env = childEnv(parent, set) // Ask OpenCode to stop so it can end its own child processes, and kill it if // it does not stop in time. cmd.Cancel = func() error { @@ -187,3 +196,15 @@ func (t *tailBuffer) String() string { defer t.mu.Unlock() return strings.TrimSpace(string(t.buf)) } + +// withoutPrefix returns env without the variables whose names start with +// prefix. +func withoutPrefix(env []string, prefix string) []string { + out := make([]string, 0, len(env)) + for _, kv := range env { + if !strings.HasPrefix(kv, prefix) { + out = append(out, kv) + } + } + return out +} From 6baa9d8b388a4ad3df1dc2ef017bed16200bfa40 Mon Sep 17 00:00:00 2001 From: Anthony Torns II Date: Tue, 20 Oct 2026 15:12:19 -0700 Subject: [PATCH 2/2] fix(setup): use one context per run and check model paths in the proxy managedModel and the AI SDK stream now use the same context. The context key is random, so with percentage targeting two contexts could pick different variations, and OpenCode could get a model that it did not enable. The proxy now accepts only /model/{modelId}/{operation} on the model path, with no empty, . or .. segment. Before, a path such as /model/../../sdk/... could leave the model relay with the user's token. Co-Authored-By: Claude Opus 5.5 --- internal/setup/agent/aiconfig.go | 6 ++++-- internal/setup/agent/bedrock_test.go | 20 +++++++++++++++++-- internal/setup/agent/opencode.go | 12 +++++++---- internal/setup/agent/opencode_handler_test.go | 3 ++- internal/setup/agent/relay_proxy.go | 19 +++++++++++++++++- 5 files changed, 50 insertions(+), 10 deletions(-) diff --git a/internal/setup/agent/aiconfig.go b/internal/setup/agent/aiconfig.go index 327e6651..8f03832f 100644 --- a/internal/setup/agent/aiconfig.go +++ b/internal/setup/agent/aiconfig.go @@ -213,7 +213,9 @@ func (s *AIConfigSource) Resolve(ctx context.Context, task Task) (Config, error) // managedModel resolves the model for a run that LaunchDarkly pays for. It // does not fall back: the config must be on and must name a Bedrock model. -func (s *AIConfigSource) managedModel(ctx context.Context, task Task) (*managedModel, error) { +// runCtx must be the same context that the run's AI SDK stream uses, so that +// percentage targeting picks the same variation in both evaluations. +func (s *AIConfigSource) managedModel(ctx context.Context, runCtx ldcontext.Context) (*managedModel, error) { if err := s.Ready(ctx); err != nil { return nil, err } @@ -221,7 +223,7 @@ func (s *AIConfigSource) managedModel(ctx context.Context, task Task) (*managedM if inspect == nil { inspect = aiserver.InspectConfig } - res := inspect(ctx, s.key(), setupContext(task)) + res := inspect(ctx, s.key(), runCtx) if !res.Enabled || res.Config == nil { return nil, errors.New("the setup agent configuration is off") } diff --git a/internal/setup/agent/bedrock_test.go b/internal/setup/agent/bedrock_test.go index 91d0f0a2..10c91f11 100644 --- a/internal/setup/agent/bedrock_test.go +++ b/internal/setup/agent/bedrock_test.go @@ -64,6 +64,10 @@ func TestRelayProxy_BedrockRouteReachesOnlyTheModelRelay(t *testing.T) { {http.MethodGet, "/model/us.anthropic.claude-sonnet-4-6/converse-stream"}, {http.MethodPost, "/polling/sdk/latest-all"}, {http.MethodPost, ""}, + {http.MethodPost, "/model/../../sdk/polling/sdk/latest-all"}, + {http.MethodPost, "/model/%2e%2e/%2e%2e/sdk/events/bulk"}, + {http.MethodPost, "/model/x/../converse"}, + {http.MethodPost, "/model/us.anthropic.claude-sonnet-4-6/converse-stream/extra"}, } { req, err := http.NewRequest(tc.method, p.bedrockURL+tc.path, nil) require.NoError(t, err) @@ -84,7 +88,7 @@ func TestAIConfigSource_ManagedModel(t *testing.T) { s, _ := fakeSource(nil, enabled(agentConfig("Bedrock", "anthropic.claude-sonnet-4-6", "Wire the call."))) defer func() { _ = s.Close(context.Background()) }() - m, err := s.managedModel(context.Background(), Task{SDKID: "python-server-sdk"}) + m, err := s.managedModel(context.Background(), setupContext(Task{SDKID: "python-server-sdk"})) require.NoError(t, err) assert.Equal(t, "anthropic.claude-sonnet-4-6", m.ModelID) @@ -100,7 +104,7 @@ func TestAIConfigSource_ManagedModel_Refusals(t *testing.T) { s, _ := fakeSource(nil, res) defer func() { _ = s.Close(context.Background()) }() - _, err := s.managedModel(context.Background(), Task{}) + _, err := s.managedModel(context.Background(), setupContext(Task{})) assert.Error(t, err, "setup must not fall back to the user's own provider") }) @@ -153,3 +157,15 @@ func TestBedrockRoute_RealOpenCode(t *testing.T) { assert.Equal(t, "POST "+BedrockRelayPath+"/model/us.anthropic.claude-sonnet-4-6/converse-stream token=api-secret-token auth=", got) } } + +func TestAIConfigSource_ManagedModelUsesTheRunsContext(t *testing.T) { + s, calls := fakeSource(nil, enabled(agentConfig("Bedrock", "anthropic.claude-sonnet-4-6", "Wire the call."))) + defer func() { _ = s.Close(context.Background()) }() + runCtx := setupContext(Task{SDKID: "python-server-sdk"}) + + _, err := s.managedModel(context.Background(), runCtx) + + require.NoError(t, err) + require.Len(t, *calls, 1) + assert.Equal(t, runCtx.Key(), (*calls)[0].ctx.Key(), "the model and the run's stream must use the same context key") +} diff --git a/internal/setup/agent/opencode.go b/internal/setup/agent/opencode.go index 9ca4e073..7eb321bf 100644 --- a/internal/setup/agent/opencode.go +++ b/internal/setup/agent/opencode.go @@ -13,6 +13,7 @@ import ( "time" "github.com/launchdarkly/go-ai-sdk/aiserver" + "github.com/launchdarkly/go-sdk-common/v3/ldcontext" ) // MinOpenCodeVersion is the oldest OpenCode version that the adapter was @@ -122,7 +123,10 @@ func (o *OpenCode) Start(ctx context.Context, task Task) (Run, error) { func (o *OpenCode) startWithSDK(ctx context.Context, task Task) (Run, error) { // OpenCode reads its configuration once, at start, so the model must be // known before OpenCode starts. - model, err := o.SDK.managedModel(ctx, task) + // One context for the whole run: its key is random, and percentage + // targeting must pick the same variation for the model and the stream. + runCtx := setupContext(task) + model, err := o.SDK.managedModel(ctx, runCtx) if err != nil { return nil, fmt.Errorf("could not get the LaunchDarkly model for the setup agent: %w", err) } @@ -144,14 +148,14 @@ func (o *OpenCode) startWithSDK(ctx context.Context, task Task) (Run, error) { sdkCtx, cancel := context.WithCancel(context.Background()) run.cancelSDK = cancel run.sdkDone = make(chan struct{}) - go run.driveWithSDK(sdkCtx, o.SDK, task) + go run.driveWithSDK(sdkCtx, o.SDK, task, runCtx) return run, nil } // driveWithSDK runs the AI SDK stream for the run. The run's own event channel // still gets every event, so the wizard shows the same progress as without the // SDK. -func (r *openCodeRun) driveWithSDK(ctx context.Context, sdk *AIConfigSource, task Task) { +func (r *openCodeRun) driveWithSDK(ctx context.Context, sdk *AIConfigSource, task Task, runCtx ldcontext.Context) { if r.sdkDone != nil { defer close(r.sdkDone) } @@ -176,7 +180,7 @@ func (r *openCodeRun) driveWithSDK(ctx context.Context, sdk *AIConfigSource, tas Key: sdk.key(), Handlers: []*aiserver.ProviderHandler{handler}, SkipJudges: true, - }).Stream(ctx, BuildFacts(task), setupContext(task)) { + }).Stream(ctx, BuildFacts(task), runCtx) { if err == nil { continue } diff --git a/internal/setup/agent/opencode_handler_test.go b/internal/setup/agent/opencode_handler_test.go index 5df1dd5b..11b7965b 100644 --- a/internal/setup/agent/opencode_handler_test.go +++ b/internal/setup/agent/opencode_handler_test.go @@ -137,7 +137,8 @@ func TestRun_ThroughTheAISDK(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) run.cancelSDK = cancel run.sdkDone = make(chan struct{}) - go run.driveWithSDK(ctx, sdk, Task{SDKID: "python-server-sdk", FlagKey: "my-new-flag"}) + task := Task{SDKID: "python-server-sdk", FlagKey: "my-new-flag"} + go run.driveWithSDK(ctx, sdk, task, setupContext(task)) require.Eventually(t, func() bool { oc.mu.Lock() diff --git a/internal/setup/agent/relay_proxy.go b/internal/setup/agent/relay_proxy.go index 0d084aa8..117a806d 100644 --- a/internal/setup/agent/relay_proxy.go +++ b/internal/setup/agent/relay_proxy.go @@ -105,7 +105,7 @@ func startRelayProxy(relayBase, accessToken string) (*relayProxy, error) { case hasPathPrefix(r.URL.Path, prefix): case hasPathPrefix(r.URL.Path, bedrockPrefix): // OpenCode calls only the Bedrock model routes, with POST. - if r.Method != http.MethodPost || !strings.HasPrefix(r.URL.Path, bedrockPrefix+"/model/") { + if r.Method != http.MethodPost || !isModelPath(strings.TrimPrefix(r.URL.Path, bedrockPrefix)) { http.NotFound(w, r) return } @@ -134,6 +134,23 @@ func startRelayProxy(relayBase, accessToken string) (*relayProxy, error) { return p, nil } +// isModelPath reports whether path is /model/{modelId}/{operation}, with no +// empty, "." or ".." segment. The check stops a path such as +// /model/../../sdk/... from leaving the model relay with the user's token. +// URL.Path is already decoded, so an encoded "%2e%2e" is also refused. +func isModelPath(path string) bool { + parts := strings.Split(strings.TrimPrefix(path, "/"), "/") + if len(parts) != 3 || parts[0] != "model" { + return false + } + for _, p := range parts { + if p == "" || p == "." || p == ".." { + return false + } + } + return true +} + // hasPathPrefix reports whether path is prefix or is under it. func hasPathPrefix(path, prefix string) bool { return path == prefix || strings.HasPrefix(path, prefix+"/")