diff --git a/cmd/setup/setup.go b/cmd/setup/setup.go index c28fef26..acb8dc5b 100644 --- a/cmd/setup/setup.go +++ b/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) diff --git a/go.mod b/go.mod index 2cb5b664..6ad08f7f 100644 --- a/go.mod +++ b/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 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 ed494583..228736e8 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() + } r.end() if r.stopServer != nil { r.stopServer() @@ -224,8 +376,9 @@ 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() @@ -233,11 +386,19 @@ func (r *openCodeRun) end() { 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`. 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") +}