feat(setup): run the setup agent through the Go AI SDK with an OpenCode handler (5/6)

cmd/setup/setup.go

@@ -80,6 +80,9 @@ func withConfigSource(runner *agent.OpenCode, run func(*cobra.Command, []string)
 			viper.GetString(cliflags.AccessTokenFlag),
 		)
 		runner.Config = source
+		// A logged-in user's agent runs through the AI SDK, so the run is
+		// recorded against the setup-agent AI Config.
+		runner.SDK, _ = source.(*agent.AIConfigSource)
 		if closer, ok := source.(interface{ Close(context.Context) error }); ok {
 			defer func() {
 				ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)

go.mod

@@ -48,6 +48,9 @@ require (
 	github.com/spf13/pflag v1.0.10
 	github.com/spf13/viper v1.21.0
 	github.com/stretchr/testify v1.12.1
+	go.opentelemetry.io/otel v1.47.0
+	go.opentelemetry.io/otel/sdk v1.47.0
+	go.opentelemetry.io/otel/trace v1.47.0
 	go.uber.org/mock v0.6.0
 	golang.org/x/exp v0.0.0-20240808152545-0cdaa3abc0fa
 	golang.org/x/term v0.46.0
@@ -129,13 +132,10 @@ require (
 	github.com/yuin/goldmark v1.7.8 // indirect
 	github.com/yuin/goldmark-emoji v1.0.5 // indirect
 	go.opentelemetry.io/auto/sdk v1.2.1 // indirect
-	go.opentelemetry.io/otel v1.47.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.47.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.47.0 // indirect
 	go.opentelemetry.io/otel/log v1.47.0 // indirect
 	go.opentelemetry.io/otel/metric v1.47.0 // indirect
-	go.opentelemetry.io/otel/sdk v1.47.0 // indirect
-	go.opentelemetry.io/otel/trace v1.47.0 // indirect
 	go.opentelemetry.io/proto/otlp v1.11.0 // indirect
 	go.yaml.in/yaml/v3 v3.0.5 // indirect
 	golang.org/x/mod v0.41.0 // indirect

internal/setup/agent/aiconfig.go

@@ -175,26 +175,35 @@ func initQuietClient(ctx context.Context, opts aiserver.InitOptions) error {
 	return nil
 }
 
-func (s *AIConfigSource) Resolve(ctx context.Context, task Task) (Config, error) {
+// Ready starts the SDK behind the relay proxy, once, and reports whether it
+// started.
+func (s *AIConfigSource) Ready(ctx context.Context) error {
 	s.once.Do(func() {
 		s.initErr = s.start(ctx)
 		if s.initErr != nil {
 			s.closeProxy(ctx)
 		}
 	})
-	if s.initErr != nil {
+	return s.initErr
+}
+
+func (s *AIConfigSource) key() string {
+	if s.Key == "" {
+		return SetupAgentConfigKey
+	}
+	return s.Key
+}
+
+func (s *AIConfigSource) Resolve(ctx context.Context, task Task) (Config, error) {
+	if s.Ready(ctx) != nil {
 		return s.fallback(ctx, task)
 	}
 
 	inspect := s.inspect
 	if inspect == nil {
 		inspect = aiserver.InspectConfig
 	}
-	key := s.Key
-	if key == "" {
-		key = SetupAgentConfigKey
-	}
-	res := inspect(ctx, key, setupContext(task))
+	res := inspect(ctx, s.key(), setupContext(task))
 	if !res.Enabled || res.Config == nil {
 		return s.fallback(ctx, task)
 	}
@@ -278,17 +287,23 @@ func configFromAI(c *aiserver.AiConfig, task Task) Config {
 	} else {
 		cfg.Task = BuildFacts(task)
 	}
-	provider := strings.ToLower(strings.TrimSpace(c.Provider.Name))
-	if mapped, ok := openCodeProviders[provider]; ok {
-		provider = mapped
+	cfg.ProviderID, cfg.ModelID = openCodeModel(*c)
+	return cfg
+}
+
+// openCodeModel returns the OpenCode provider ID and model ID for a config.
+// OpenCode needs both values to pick a model, so if the config has only one,
+// both are empty and OpenCode uses its default model.
+func openCodeModel(c aiserver.AiConfig) (providerID, modelID string) {
+	providerID = strings.ToLower(strings.TrimSpace(c.Provider.Name))
+	if mapped, ok := openCodeProviders[providerID]; ok {
+		providerID = mapped
 	}
-	cfg.ProviderID = provider
-	// OpenCode needs both values to pick a model. With only one, OpenCode
-	// uses its default model.
-	if cfg.ProviderID == "" || cfg.ModelID == "" {
-		cfg.ProviderID, cfg.ModelID = "", ""
+	modelID = c.Model.Name
+	if providerID == "" || modelID == "" {
+		return "", ""
 	}
-	return cfg
+	return providerID, modelID
 }
 
 // RelayError is a refusal from the relay, for example because the kill switch

internal/setup/agent/opencode.go

@@ -11,6 +11,8 @@ import (
 	"strings"
 	"sync"
 	"time"
+
+	"github.com/launchdarkly/go-ai-sdk/aiserver"
 )
 
 // MinOpenCodeVersion is the oldest OpenCode version that the adapter was
@@ -21,8 +23,12 @@ const MinOpenCodeVersion = "1.19.0"
 type OpenCode struct {
 	// Binary is the OpenCode executable. The default is "opencode" on the PATH.
 	Binary string
-	// Config gives the model and the instructions for each run.
+	// Config gives the model and the instructions for each run when SDK is
+	// nil or cannot start.
 	Config ConfigSource
+	// SDK runs the agent through the AI SDK: the SDK resolves the setup-agent
+	// AI Config, calls the OpenCode handler, and records the run.
+	SDK *AIConfigSource
 
 	// lookPath and version are fields so that tests can replace them.
 	lookPath func(string) (string, error)
@@ -69,6 +75,9 @@ 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 {
+		return o.startWithSDK(ctx, task)
+	}
 	source := o.Config
 	if source == nil {
 		source = DefaultConfigSource{}
@@ -95,17 +104,101 @@ func (o *OpenCode) Start(ctx context.Context, task Task) (Run, error) {
 		srv.stop()
 		return nil, err
 	}
-	run.stopServer = srv.stop
+	run.watch(srv)
+	return run, nil
+}
+
+// startWithSDK runs the agent through the AI SDK. The SDK evaluates the
+// setup-agent AI Config, composes the prompt from it, and calls the OpenCode
+// handler, which sends that prompt to this run's session. The SDK then
+// records the run, its tool calls, and its tokens against the variation.
+//
+// If the SDK cannot resolve the config before the handler starts, for
+// 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()
+	if err != nil {
+		return nil, err
+	}
+	srv, err := startServer(ctx, o.binary(), task.Dir, config)
+	if err != nil {
+		return nil, err
+	}
+	run, err := startSession(ctx, srv.client)
+	if err != nil {
+		srv.stop()
+		return nil, err
+	}
+	run.watch(srv)
+
+	sdkCtx, cancel := context.WithCancel(context.Background())
+	run.cancelSDK = cancel
+	run.sdkDone = make(chan struct{})
+	go run.driveWithSDK(sdkCtx, o.SDK, task)
+	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) {
+	if r.sdkDone != nil {
+		defer close(r.sdkDone)
+	}
+	var mu sync.Mutex
+	started := false
+	handler := NewOpenCodeHandler(func(ctx context.Context, p SessionPrompt) (<-chan Event, error) {
+		mu.Lock()
+		started = true
+		mu.Unlock()
+		cfg := Config{Instructions: p.System, ProviderID: p.ProviderID, ModelID: p.ModelID}
+		text := p.Text
+		if cfg.Instructions == "" {
+			cfg.Instructions, text = DefaultInstructions, BuildPrompt(task)
+		}
+		tap := r.startTap()
+		if err := r.client.promptAsync(ctx, r.sessionID, cfg, text); err != nil {
+			return nil, fmt.Errorf("could not send the task to OpenCode: %w", err)
+		}
+		return tap, nil
+	})
+	for _, err := range aiserver.Config(aiserver.ConfigArgs{
+		Key:        sdk.key(),
+		Handlers:   []*aiserver.ProviderHandler{handler},
+		SkipJudges: true,
+	}).Stream(ctx, BuildFacts(task), setupContext(task)) {
+		if err == nil {
+			continue
+		}
+		mu.Lock()
+		handlerStarted := started
+		mu.Unlock()
+		if handlerStarted {
+			// The handler already sent the failure to the run's events.
+			return
+		}
+		cfg, _ := DefaultConfigSource{}.Resolve(ctx, task)
+		if perr := r.client.promptAsync(ctx, r.sessionID, cfg, BuildPrompt(task)); perr != nil {
+			r.fail(fmt.Sprintf("could not send the task to OpenCode: %v", perr))
+		}
+		return
+	}
+}
+
+// watch ends the run's events if OpenCode exits during the run, so the caller
+// gets an error and does not wait forever.
+func (r *openCodeRun) watch(srv *server) {
+	r.stopServer = srv.stop
 	go func() {
 		// If OpenCode exits during the run, end the event stream so the caller
 		// gets an error and does not wait forever.
 		select {
 		case <-srv.exited:
-			run.fail(fmt.Sprintf("OpenCode stopped during the run: %s", srv.output.String()))
-		case <-run.closed:
+			r.fail(fmt.Sprintf("OpenCode stopped during the run: %s", srv.output.String()))
+		case <-r.closed:
 		}
 	}()
-	return run, nil
 }
 
 // openCodeRun is a Run on an OpenCode server.
@@ -116,14 +209,34 @@ type openCodeRun struct {
 	cancel     context.CancelFunc
 	closed     chan struct{}
 	stopServer func()
+	// cancelSDK stops the AI SDK stream of a run that started with the SDK,
+	// and sdkDone closes when that stream ends.
+	cancelSDK context.CancelFunc
+	sdkDone   chan struct{}
 
-	mu        sync.Mutex
-	finished  bool
+	mu       sync.Mutex
+	finished bool
+	// tap gets a copy of every event, for the OpenCode handler.
+	tap       chan Event
 	closeOnce sync.Once
 }
 
 // startRun creates a session, subscribes to its events, and sends the task.
 func startRun(ctx context.Context, c *client, cfg Config, prompt string) (*openCodeRun, error) {
+	r, err := startSession(ctx, c)
+	if err != nil {
+		return nil, err
+	}
+	if err := c.promptAsync(ctx, r.sessionID, cfg, prompt); err != nil {
+		r.Close()
+		return nil, fmt.Errorf("could not send the task to OpenCode: %w", err)
+	}
+	return r, nil
+}
+
+// startSession creates a session and subscribes to its events. It sends no
+// task.
+func startSession(ctx context.Context, c *client) (*openCodeRun, error) {
 	sessionID, err := c.createSession(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("could not create an OpenCode session: %w", err)
@@ -143,12 +256,19 @@ func startRun(ctx context.Context, c *client, cfg Config, prompt string) (*openC
 		closed:    make(chan struct{}),
 	}
 	go r.pump(stream)
+	return r, nil
+}
 
-	if err := c.promptAsync(ctx, sessionID, cfg, prompt); err != nil {
-		r.Close()
-		return nil, fmt.Errorf("could not send the task to OpenCode: %w", err)
+// startTap returns a channel that gets a copy of every later event. It closes
+// after the last event.
+func (r *openCodeRun) startTap() <-chan Event {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.tap = make(chan Event, 256)
+	if r.finished {
+		close(r.tap)
 	}
-	return r, nil
+	return r.tap
 }
 
 func (r *openCodeRun) pump(stream io.ReadCloser) {
@@ -190,13 +310,33 @@ func (r *openCodeRun) send(ev Event) bool {
 	case <-r.closed:
 		return false
 	}
+	if r.tap != nil {
+		select {
+		case r.tap <- ev:
+		case <-r.closed:
+			return false
+		}
+	}
 	if ev.Kind == EventDone || ev.Kind == EventError {
 		r.finished = true
 		close(r.events)
+		if r.tap != nil {
+			close(r.tap)
+		}
 	}
 	return true
 }
 
+// sdkFinishTimeout is how long Close waits for the AI SDK to record a run that
+// ended.
+const sdkFinishTimeout = 5 * time.Second
+
+func (r *openCodeRun) isFinished() bool {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return r.finished
+}
+
 func (r *openCodeRun) fail(msg string) {
 	r.send(Event{Kind: EventError, Text: msg})
 }
@@ -213,8 +353,20 @@ func (r *openCodeRun) Answer(ctx context.Context, req QuestionRequest, answers [
 
 func (r *openCodeRun) Close() error {
 	r.closeOnce.Do(func() {
+		// Let the AI SDK finish recording a run that ended, so its tokens,
+		// duration, and spans reach LaunchDarkly. A run that is still going
+		// is stopped at once.
+		if r.sdkDone != nil && r.isFinished() {
+			select {
+			case <-r.sdkDone:
+			case <-time.After(sdkFinishTimeout):
+			}
+		}
 		close(r.closed)
 		r.cancel()
+		if r.cancelSDK != nil {
+			r.cancelSDK()
+		}
 		r.end()
 		if r.stopServer != nil {
 			r.stopServer()
@@ -224,20 +376,29 @@ func (r *openCodeRun) Close() error {
 }
 
 // end sends a final error event when Close stops an unfinished run.
-// It closes the event channel, so readers do not wait forever. A send that
-// waits on a full channel returns when r.closed closes, so end does not block.
+// It closes the event channel and the tap, so readers do not wait forever. A
+// send that waits on a full channel returns when r.closed closes, so end does
+// not block.
 func (r *openCodeRun) end() {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 	if r.finished {
 		return
 	}
 	r.finished = true
+	ev := Event{Kind: EventError, Text: "The run stopped before it finished."}
 	select {
-	case r.events <- Event{Kind: EventError, Text: "The run stopped before it finished."}:
+	case r.events <- ev:
 	default:
 	}
 	close(r.events)
+	if r.tap != nil {
+		select {
+		case r.tap <- ev:
+		default:
+		}
+		close(r.tap)
+	}
 }
 
 // binaryVersion runs `opencode --version`.