feat(setup): send the agent's model calls to a LaunchDarkly Bedrock relay (7/7)

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 {
@@ -452,7 +452,7 @@ func TestBuildPrompt_HasTheSameFactsAsBuildFacts(t *testing.T) {
 }
 
 func TestOpenCodeConfig_ReadAndEditRefuseEnvFiles(t *testing.T) {
-	data, err := openCodeConfig()
+	data, err := openCodeConfig(nil)
 	require.NoError(t, err)
 	var cfg struct {
 		Permission map[string]json.RawMessage `json:"permission"`

internal/setup/agent/aiconfig.go

@@ -3,6 +3,7 @@ package agent
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"log/slog"
@@ -210,6 +211,29 @@ 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.
+// 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
+	}
+	inspect := s.inspect
+	if inspect == nil {
+		inspect = aiserver.InspectConfig
+	}
+	res := inspect(ctx, s.key(), runCtx)
+	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)

internal/setup/agent/bedrock_test.go

@@ -0,0 +1,171 @@
+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, ""},
+		{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)
+		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(), setupContext(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(), setupContext(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)
+	}
+}
+
+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")
+}

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
@@ -75,7 +76,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 +90,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 +121,20 @@ 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.
+	// 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)
+	}
+	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
 	}
@@ -135,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)
 	}
@@ -167,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
 		}

internal/setup/agent/opencode_config.go

@@ -14,7 +14,32 @@ 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
@@ -25,7 +50,7 @@ func openCodeConfig() ([]byte, error) {
 	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":               envFiles,
@@ -39,7 +64,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

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()