From 5447aa0663349635ff21ed753475c9545c43aadd Mon Sep 17 00:00:00 2001 From: Anthony Torns II Date: Mon, 19 Oct 2026 23:38:59 -0700 Subject: [PATCH] feat(setup): run the setup agent through the Go AI SDK with an OpenCode handler The setup agent is now an AI SDK agent. For a logged-in user, setup calls aiserver.Config("ldcli-setup-agent").Stream with a new OpenCode handler. The AI SDK resolves the config through the relay, composes the prompt from it, and calls the handler. The handler sends that prompt to the OpenCode session and turns the session's events into the SDK's stream events. The SDK records the run: an opencode.session span, a chat span per reply with its tokens, an execute_tool span per tool call, and the tokens, duration, and success events for the config's Monitoring tab. The handler implements the AI SDK's handler contract for agent-mode configs of any provider, so it can move into go-ai-sdk unchanged. The wizard still gets every OpenCode event. If the SDK cannot resolve the config before the handler starts, the run uses the built-in instructions and task. Closing a run that ended waits up to 5 seconds for the SDK to record it. Co-Authored-By: Claude Opus 5.5 --- cmd/setup/setup.go | 3 + go.mod | 6 +- internal/setup/agent/aiconfig.go | 47 ++-- internal/setup/agent/opencode.go | 174 ++++++++++++++- internal/setup/agent/opencode_handler.go | 200 ++++++++++++++++++ internal/setup/agent/opencode_handler_test.go | 180 ++++++++++++++++ 6 files changed, 580 insertions(+), 30 deletions(-) create mode 100644 internal/setup/agent/opencode_handler.go create mode 100644 internal/setup/agent/opencode_handler_test.go diff --git a/cmd/setup/setup.go b/cmd/setup/setup.go index 91b34d06..f79eb4f4 100644 --- a/cmd/setup/setup.go +++ b/cmd/setup/setup.go @@ -79,6 +79,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) diff --git a/go.mod b/go.mod index cd695c6d..d487fa30 100644 --- a/go.mod +++ b/go.mod @@ -45,6 +45,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.5.2 golang.org/x/exp v0.0.0-20240808152545-0cdaa3abc0fa golang.org/x/term v0.46.0 @@ -127,13 +130,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 diff --git a/internal/setup/agent/aiconfig.go b/internal/setup/agent/aiconfig.go index 12e2197f..5977df83 100644 --- a/internal/setup/agent/aiconfig.go +++ b/internal/setup/agent/aiconfig.go @@ -175,14 +175,27 @@ 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) } @@ -190,11 +203,7 @@ func (s *AIConfigSource) Resolve(ctx context.Context, task Task) (Config, error) 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 diff --git a/internal/setup/agent/opencode.go b/internal/setup/agent/opencode.go index 9c10292d..8c90baa5 100644 --- a/internal/setup/agent/opencode.go +++ b/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() + } if r.stopServer != nil { r.stopServer() } diff --git a/internal/setup/agent/opencode_handler.go b/internal/setup/agent/opencode_handler.go new file mode 100644 index 00000000..3b8bd0b8 --- /dev/null +++ b/internal/setup/agent/opencode_handler.go @@ -0,0 +1,200 @@ +package agent + +import ( + "context" + "errors" + "iter" + "strings" + + "go.opentelemetry.io/otel/trace" + + "github.com/launchdarkly/go-ai-sdk/aiserver" +) + +// The OpenCode handler is an AI SDK provider handler, like the Claude and +// OpenAI agent handlers in go-ai-sdk. The AI SDK resolves the AI Config and +// calls the handler, the handler runs one OpenCode session through +// SessionEvents, and the AI SDK records the run against the variation. It is +// written so that it can move into go-ai-sdk as its own package with the same +// contract. +const ( + // HandlerTracer is the handler's instrumentation scope. The AI SDK + // requires it to start with aiserver.TracerPrefix. + HandlerTracer = aiserver.TracerPrefix + "/opencodeagents" + // HandlerSpanName is the name of the root span of one run. + HandlerSpanName = "opencode.session" +) + +// SessionPrompt is what the handler sends to OpenCode for one run. +type SessionPrompt struct { + // System is the system prompt from the config. + System string + // ProviderID and ModelID select the OpenCode model. Both are empty if the + // config names no model, and then OpenCode uses its default model. + ProviderID string + ModelID string + // Text is the user message. + Text string +} + +// SessionEvents runs one OpenCode session for the prompt and returns its +// events. The channel must end after EventDone or EventError. +type SessionEvents func(ctx context.Context, p SessionPrompt) (<-chan Event, error) + +// NewOpenCodeHandler returns an AI SDK handler for agent-mode configs of any +// provider. The caller owns the OpenCode process and supplies its events. +// +// The handler does not run judges. A config that has judges must be called +// with SkipJudges, because the handler would also be chosen for the judge +// configs, and it cannot run a judge prompt. +func NewOpenCodeHandler(events SessionEvents) *aiserver.ProviderHandler { + h := openCodeHandler{events: events} + return aiserver.CreateHandlerInternal( + aiserver.ProvidesFor{Provider: aiserver.WildcardProvider, Mode: aiserver.ModeAgent}, + h.invoke, h.stream, false, + ) +} + +type openCodeHandler struct { + events SessionEvents +} + +func (h openCodeHandler) invoke(ctx context.Context, req aiserver.HandlerRequest) (aiserver.HandlerResult, error) { + var res aiserver.HandlerResult + for ev, err := range h.stream(ctx, req) { + if err != nil { + return res, err + } + if ev.Type == aiserver.EventDone { + if ev.Output != nil { + res.Output = *ev.Output + } + res.Usage = ev.Usage + } + } + return res, nil +} + +func (h openCodeHandler) stream(ctx context.Context, req aiserver.HandlerRequest) iter.Seq2[aiserver.HandlerStreamEvent, error] { + return func(yield func(aiserver.HandlerStreamEvent, error) bool) { + p := aiserver.ComposePrompt(req, aiserver.ModeAgent) + providerID, modelID := openCodeModel(req.Config) + system := providerID + if system == "" { + system = "opencode" + } + run := aiserver.StartHandlerRun(ctx, aiserver.HandlerSpanOptions{ + Tracer: HandlerTracer, + Name: HandlerSpanName, + System: system, + Operation: "invoke_agent", + RequestModel: modelID, + }, req, p) + defer run.EndOnPanic() + tools := toolSpans{run: run, open: map[string]trace.Span{}} + + fail := func(err error) { + tools.endAll(err) + run.Fail(err) + yield(aiserver.HandlerStreamEvent{}, err) + } + + events, err := h.events(run.Context(), SessionPrompt{ + System: p.System, + ProviderID: providerID, + ModelID: modelID, + Text: userText(p), + }) + if err != nil { + fail(err) + return + } + + var output strings.Builder + var usage aiserver.SpanUsage + responseModel := modelID + for ev := range events { + switch ev.Kind { + case EventText: + output.WriteString(ev.Text) + if !yield(aiserver.HandlerStreamEvent{Type: aiserver.EventChunk, Text: ev.Text}, nil) { + tools.endAll(nil) + run.Abandon() + return + } + case EventTool: + tools.update(ev.Tool) + case EventUsage: + // OpenCode reports the usage of each finished reply, so each one + // becomes a chat span with its own token counts. + turn := aiserver.SpanUsage{Input: float64(ev.Usage.InputTokens), Output: float64(ev.Usage.OutputTokens)} + if ev.Usage.ModelID != "" { + responseModel = ev.Usage.ModelID + } + run.EndChat(run.StartChat(responseModel), responseModel, turn, "") + usage.Add(turn) + case EventDone: + tools.endAll(nil) + text := output.String() + run.Finish(responseModel, text) + yield(aiserver.HandlerStreamEvent{ + Type: aiserver.EventDone, + Output: &text, + Usage: map[string]any{"input": usage.Input, "output": usage.Output}, + }, nil) + return + case EventError: + fail(errors.New(ev.Text)) + return + } + } + fail(errors.New("the OpenCode session ended without a result")) + } +} + +// userText is the user message for OpenCode: the last user turn of the +// composed prompt. OpenCode keeps the conversation itself, so the handler +// sends only the new message. +func userText(p aiserver.Prompt) string { + for i := len(p.Turns) - 1; i >= 0; i-- { + if p.Turns[i].Role == "user" { + return aiserver.ContentText(p.Turns[i]) + } + } + return "" +} + +// toolSpans keeps one execute_tool span open for each running tool call. +type toolSpans struct { + run *aiserver.HandlerRun + open map[string]trace.Span +} + +func (t *toolSpans) update(tc ToolCall) { + if _, ok := t.open[tc.ID]; !ok { + t.open[tc.ID] = t.run.StartTool(tc.Name, tc.ID) + } + switch tc.Status { + case "completed": + t.end(tc.ID, nil) + case "error": + msg := tc.Title + if msg == "" { + msg = tc.Name + " failed" + } + t.end(tc.ID, errors.New(msg)) + } +} + +func (t *toolSpans) end(id string, err error) { + if span, ok := t.open[id]; ok { + t.run.EndTool(span, err) + delete(t.open, id) + } +} + +func (t *toolSpans) endAll(err error) { + for id := range t.open { + t.end(id, err) + } +} diff --git a/internal/setup/agent/opencode_handler_test.go b/internal/setup/agent/opencode_handler_test.go new file mode 100644 index 00000000..5df1dd5b --- /dev/null +++ b/internal/setup/agent/opencode_handler_test.go @@ -0,0 +1,180 @@ +package agent + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + + "github.com/launchdarkly/go-ai-sdk/aiserver" +) + +func handlerRequest() aiserver.HandlerRequest { + cfg := aiserver.AiConfig{Instructions: "Wire the AI call.", Model: aiserver.Model{Name: "claude-sonnet-5"}} + cfg.Provider.Name = "Anthropic" + return aiserver.HandlerRequest{Config: cfg, UserInput: "Facts from ldcli setup"} +} + +func eventsOf(evs ...Event) <-chan Event { + ch := make(chan Event, len(evs)) + for _, e := range evs { + ch <- e + } + close(ch) + return ch +} + +func TestOpenCodeHandler_StreamsAndRecordsTheRun(t *testing.T) { + spans := tracetest.NewSpanRecorder() + prev := otel.GetTracerProvider() + otel.SetTracerProvider(sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(spans))) + defer otel.SetTracerProvider(prev) + + var got SessionPrompt + h := NewOpenCodeHandler(func(_ context.Context, p SessionPrompt) (<-chan Event, error) { + got = p + return eventsOf( + Event{Kind: EventTool, Tool: ToolCall{ID: "c1", Name: "edit", Status: "running"}}, + Event{Kind: EventText, Text: "Wired "}, + Event{Kind: EventTool, Tool: ToolCall{ID: "c1", Name: "edit", Status: "completed"}}, + Event{Kind: EventUsage, Usage: Usage{ModelID: "claude-sonnet-5", InputTokens: 100, OutputTokens: 20}}, + Event{Kind: EventText, Text: "main.py."}, + Event{Kind: EventUsage, Usage: Usage{ModelID: "claude-sonnet-5", InputTokens: 50, OutputTokens: 10}}, + Event{Kind: EventDone}, + ), nil + }) + + var chunks []string + var done aiserver.HandlerStreamEvent + for ev, err := range h.Stream(context.Background(), handlerRequest()) { + require.NoError(t, err) + if ev.Type == aiserver.EventChunk { + chunks = append(chunks, ev.Text) + } else { + done = ev + } + } + + assert.Equal(t, SessionPrompt{System: "Wire the AI call.", ProviderID: "anthropic", ModelID: "claude-sonnet-5", Text: "Facts from ldcli setup"}, got) + assert.Equal(t, []string{"Wired ", "main.py."}, chunks) + require.Equal(t, aiserver.EventDone, done.Type) + assert.Equal(t, "Wired main.py.", *done.Output) + assert.Equal(t, aiserver.TokenUsage{Input: 150, Output: 30, Total: 180}, aiserver.ParseUsage(done.Usage)) + + var names []string + for _, s := range spans.Ended() { + names = append(names, s.Name()) + } + assert.ElementsMatch(t, []string{HandlerSpanName, "execute_tool edit", "chat claude-sonnet-5", "chat claude-sonnet-5"}, names) +} + +func TestOpenCodeHandler_ErrorEndsTheRun(t *testing.T) { + h := NewOpenCodeHandler(func(context.Context, SessionPrompt) (<-chan Event, error) { + return eventsOf(Event{Kind: EventError, Text: "no API key"}), nil + }) + + var gotErr error + for _, err := range h.Stream(context.Background(), handlerRequest()) { + if err != nil { + gotErr = err + } + } + + require.Error(t, gotErr) + assert.Contains(t, gotErr.Error(), "no API key") +} + +func TestOpenCodeHandler_StreamWithoutAnEndIsAnError(t *testing.T) { + h := NewOpenCodeHandler(func(context.Context, SessionPrompt) (<-chan Event, error) { + return eventsOf(Event{Kind: EventText, Text: "partial"}), nil + }) + + var gotErr error + for _, err := range h.Stream(context.Background(), handlerRequest()) { + if err != nil { + gotErr = err + } + } + + assert.Error(t, gotErr) +} + +// TestRun_ThroughTheAISDK runs the agent the way setup does for a logged-in +// user: the real AI SDK behind the proxy reads the config from a fake relay, +// calls the OpenCode handler, and the handler drives a fake OpenCode session. +func TestRun_ThroughTheAISDK(t *testing.T) { + var mu sync.Mutex + var bodies []string + relay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == RelayPath+"/polling/sdk/latest-all" { + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"flags":{"ldcli-setup-agent":`+strings.ReplaceAll(setupAgentFlag, "\n", "")+`},"segments":{}}`) + return + } + body, _ := io.ReadAll(r.Body) + mu.Lock() + bodies = append(bodies, r.URL.Path+" "+string(body)) + mu.Unlock() + w.WriteHeader(http.StatusAccepted) + })) + defer relay.Close() + sdk := NewAIConfigSource(relay.URL, "api-123").(*AIConfigSource) + require.NoError(t, sdk.Ready(context.Background())) + + oc := newFakeOpenCode(t) + run, err := startSession(context.Background(), oc.client()) + require.NoError(t, err) + 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"}) + + require.Eventually(t, func() bool { + oc.mu.Lock() + defer oc.mu.Unlock() + return oc.prompt != nil + }, 5*time.Second, 10*time.Millisecond, "the handler sends the prompt") + oc.mu.Lock() + assert.Equal(t, "From AgentControl.", oc.prompt["system"], "the system prompt comes from the config") + assert.Equal(t, map[string]any{"providerID": "anthropic", "modelID": "claude-sonnet-5"}, oc.prompt["model"]) + oc.mu.Unlock() + + oc.emit(`{"type":"message.updated","properties":{"info":{"id":"m1","sessionID":"ses","role":"assistant","providerID":"anthropic","modelID":"claude-sonnet-5","time":{"created":1,"completed":2},"tokens":{"input":100,"output":20}}}}`) + oc.emit(`{"type":"session.idle","properties":{"sessionID":"ses"}}`) + + var kinds []EventKind + for ev := range run.Events() { + kinds = append(kinds, ev.Kind) + } + assert.Equal(t, []EventKind{EventUsage, EventDone}, kinds, "the wizard still gets every event") + + require.NoError(t, run.Close()) + require.NoError(t, sdk.Close(context.Background())) + mu.Lock() + defer mu.Unlock() + all := strings.Join(bodies, "\n") + for _, b := range bodies { + path, rest, _ := strings.Cut(b, " ") + var kinds []string + for _, k := range []string{"$ld:ai:feedback", "$ld:ai:tokens:total", "$ld:ai:generation:success", "$ld:ai:duration:total", "$ld:ai:sdk:info", "$ld:ai:usage"} { + if strings.Contains(rest, k) { + kinds = append(kinds, k) + } + } + t.Logf("relay got %s (%d bytes) %v", path, len(rest), kinds) + } + assert.Contains(t, all, RelayPath+"/events/bulk", "the SDK sends its events through the relay") + assert.Contains(t, all, "$ld:ai:tokens:total", "the run's tokens are tracked against the AI Config") + assert.Contains(t, all, "$ld:ai:generation:success", "the run's result is tracked") + assert.Contains(t, all, RelayPath+"/otel/v1/traces", "the run's spans go out through the relay") +}