diff --git a/ai/cmd/agent/main.go b/ai/cmd/agent/main.go index e69d46d..effe92b 100644 --- a/ai/cmd/agent/main.go +++ b/ai/cmd/agent/main.go @@ -11,6 +11,7 @@ import ( "time" "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/agent" + "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/anthropic" "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/diag" "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/harness" "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/jev" @@ -95,6 +96,10 @@ func main() { HitlURL: env("HITL_URL", "http://hitl:8081"), } ag := agent.New(acfg, oai, tr, hc, store, enums, pv) + // Claude Messages API for claude-* model ids chosen per campaign; the key may be registered later (empty = diagnoses with + // a claude-* model fail with "anthropic api key is empty", recorded in the harness diagnosis event). + anth := anthropic.New(env("ANTHROPIC_BASE_URL", ""), keyReader(env("ANTHROPIC_KEY_FILE", "/run/secrets/anthropic-api-key"))) + ag.SetAnthropic(anth) jevMode := env("JEV_MODE", "off") // off | shadow | enforce (shadow records Jev answers next to the diagnosis; enforce also applies the two review rules) var jc *jev.Client if jevMode == "shadow" || jevMode == "enforce" { @@ -130,6 +135,7 @@ func main() { out := map[string]any{"model": acfg.Model, "prompt_version": pv, "labels": enums.Labels, "actions": enums.Actions} key, kerr := keyReader(env("OPENAI_KEY_FILE", "/run/secrets/openai-api-key"))() out["openai_key_present"] = kerr == nil && key != "" + out["anthropic_key_present"] = anth.KeyPresent() out["jev_mode"] = jevMode if jc != nil { out["jev_key_present"] = jc.KeyPresent() diff --git a/ai/internal/agent/agent.go b/ai/internal/agent/agent.go index 6cb1cf8..81aae84 100644 --- a/ai/internal/agent/agent.go +++ b/ai/internal/agent/agent.go @@ -35,9 +35,16 @@ type Config struct { ConfidenceMin float64 // 0.5 (design 6.3) } +// Responder is the one call the diagnosis loop needs from an LLM provider (OpenAI Responses API natively, +// the Claude Messages API through internal/anthropic). +type Responder interface { + CreateResponse(ctx context.Context, req openai.Request) (*openai.Response, error) +} + type Agent struct { cfg Config oai *openai.Client + anth Responder // nil = claude-* models unavailable tools *tools.Runner harness *harness.Client store *rag.Store @@ -64,6 +71,20 @@ func New(cfg Config, oai *openai.Client, tr *tools.Runner, hc *harness.Client, s return &Agent{cfg: cfg, oai: oai, tools: tr, harness: hc, store: store, enums: enums, promptV: promptVersion, http: &http.Client{Timeout: 30 * time.Second}} } +// SetAnthropic enables claude-* model ids (Claude Messages API adapter). +func (a *Agent) SetAnthropic(c Responder) { a.anth = c } + +// provider picks the LLM client for a model id: claude-* -> Anthropic adapter, everything else -> OpenAI. +func (a *Agent) provider(model string) (Responder, error) { + if strings.HasPrefix(model, "claude-") { + if a.anth == nil { + return nil, fmt.Errorf("model %s: anthropic provider not configured", model) + } + return a.anth, nil + } + return a.oai, nil +} + // EnableShadow turns the Jev gates on (JEV_MODE=shadow records only; enforce also applies the review rules). func (a *Agent) EnableShadow(c *jev.Client, enforce bool) { a.jev, a.jevEnforce = c, enforce } @@ -72,6 +93,7 @@ type Request struct { Alert map[string]any `json:"alert"` RunID *string `json:"run_id"` ScenarioID *string `json:"scenario_id"` + Model string `json:"model,omitempty"` // campaign-level model selection (harness runs.model); "" = Config.Model } // Result is returned to the harness: the diagnosis fields at top level (the harness stores the whole object as the @@ -84,7 +106,8 @@ type Result struct { Turns int `json:"turns"` ToolCalls int `json:"tool_calls"` DurationMs int64 `json:"duration_ms"` - Disposition string `json:"disposition"` // proposal | escalated | none + Usage map[string]int `json:"usage,omitempty"` // input_tokens / output_tokens / total_tokens summed over the turns + Disposition string `json:"disposition"` // proposal | escalated | none ProposalToken string `json:"proposal_token,omitempty"` Jev map[string]any `json:"jev,omitempty"` // shadow gate answers (recorded, never acted on) } @@ -155,13 +178,17 @@ func (a *Agent) Diagnose(ctx context.Context, req Request) (*Result, error) { func (a *Agent) DiagnoseWith(ctx context.Context, req Request, prefetched string) (*Result, error) { start := time.Now() sid := "ses-" + uuid.NewString()[:8] - res := &Result{SessionID: sid, Model: a.cfg.Model, PromptVersion: a.promptV} + model := a.cfg.Model + if req.Model != "" { + model = req.Model + } + res := &Result{SessionID: sid, Model: model, PromptVersion: a.promptV} var calls []toolCallLog input := []openai.Item{ {Type: "message", Role: "system", Content: a.cfg.SystemPrompt}, {Type: "message", Role: "user", Content: a.userMessage(ctx, req, prefetched)}, } - session := rag.Session{SessionID: sid, RunID: req.RunID, ScenarioID: req.ScenarioID, Model: a.cfg.Model, PromptVersion: a.promptV, Input: map[string]any{"alert": req.Alert, "user_message": input[1].Content}} + session := rag.Session{SessionID: sid, RunID: req.RunID, ScenarioID: req.ScenarioID, Model: model, PromptVersion: a.promptV, Input: map[string]any{"alert": req.Alert, "user_message": input[1].Content}} finish := func(err error) { session.Output = res session.ToolCalls = calls @@ -176,6 +203,11 @@ func (a *Agent) DiagnoseWith(ctx context.Context, req Request, prefetched string } } + llm, err := a.provider(model) + if err != nil { + finish(err) + return nil, err + } defs := a.toolDefs() var prevID string var diagnosis *diag.Diagnosis @@ -185,14 +217,14 @@ func (a *Agent) DiagnoseWith(ctx context.Context, req Request, prefetched string usage := map[string]int{"input_tokens": 0, "output_tokens": 0, "total_tokens": 0, "turns": 0} for turn := 1; turn <= a.cfg.MaxTurns && diagnosis == nil; turn++ { res.Turns = turn - r := openai.Request{Model: a.cfg.Model, Input: input, Tools: defs, ToolChoice: "auto", MaxOutputTokens: a.cfg.MaxOutputTokens, ParallelToolCalls: &pf} + r := openai.Request{Model: model, Input: input, Tools: defs, ToolChoice: "auto", MaxOutputTokens: a.cfg.MaxOutputTokens, ParallelToolCalls: &pf} if prevID != "" { r.PreviousResponseID = prevID } if turn == a.cfg.MaxTurns { // last chance: force the final call r.ToolChoice = map[string]any{"type": "function", "name": "submit_diagnosis"} } - resp, err := a.oai.CreateResponse(ctx, r) + resp, err := llm.CreateResponse(ctx, r) if err != nil { finish(err) return nil, err @@ -228,6 +260,7 @@ func (a *Agent) DiagnoseWith(ctx context.Context, req Request, prefetched string } } session.Usage = usage + res.Usage = usage res.ToolCalls = len(calls) res.DurationMs = time.Since(start).Milliseconds() if diagnosis == nil { diff --git a/ai/internal/agent/agent_test.go b/ai/internal/agent/agent_test.go index 26b1c57..a0a5b75 100644 --- a/ai/internal/agent/agent_test.go +++ b/ai/internal/agent/agent_test.go @@ -8,6 +8,7 @@ import ( "strings" "testing" + "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/anthropic" "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/diag" "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/harness" "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/openai" @@ -81,3 +82,48 @@ func TestTargetFor(t *testing.T) { t.Fatal("mig") } } + +// A claude-* model id routes the same loop through the Anthropic adapter; the request model wins over Config.Model +// and the summed usage is returned in the result. +func TestDiagnoseRoutesClaudeModel(t *testing.T) { + var models []string + anth := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var b map[string]any + _ = json.NewDecoder(r.Body).Decode(&b) + models = append(models, b["model"].(string)) + args := `{"root_cause_label":"db_primary_down","confidence":0.8,"evidence":["e"],"proposed_action":"escalate","target":"db-primary","rationale":"r"}` + w.Write([]byte(`{"id":"msg_a","type":"message","model":"claude-sonnet-5-5","stop_reason":"tool_use","content":[{"type":"tool_use","id":"toolu_a","name":"submit_diagnosis","input":` + args + `}],"usage":{"input_tokens":40,"output_tokens":8}}`)) + })) + defer anth.Close() + oaiSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/embeddings" { + w.Write([]byte(`{"data":[{"index":0,"embedding":[0.1,0.2]}]}`)) + return + } + t.Errorf("openai must not be called for a claude-* model: %s", r.URL.Path) + })) + defer oaiSrv.Close() + mimir := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(`{"status":"success","data":{"resultType":"vector","result":[]}}`)) + })) + defer mimir.Close() + oai := openai.New(oaiSrv.URL, func() (string, error) { return "test-key", nil }) + tr := tools.New(tools.Config{MimirURL: mimir.URL, LokiURL: mimir.URL, EmbedModel: "e"}, nil, oai) + enums := diag.Enums{Labels: diag.DefaultLabels, Actions: diag.DefaultActions} + ag := New(Config{Model: "gpt-default", SystemPrompt: "sys", Topology: "topo", SLO: "slo", HitlURL: "http://127.0.0.1:1", MaxTurns: 3}, oai, tr, harness.New("http://127.0.0.1:1"), nil, enums, "pv1") + req := Request{Alert: map[string]any{"alertname": "DbPrimaryDown"}, Model: "claude-sonnet-5-5"} + if _, err := ag.Diagnose(context.Background(), req); err == nil || !strings.Contains(err.Error(), "anthropic provider not configured") { + t.Fatalf("without SetAnthropic a claude-* model must fail clearly, got %v", err) + } + ag.SetAnthropic(anthropic.New(anth.URL, func() (string, error) { return "k", nil })) + res, err := ag.Diagnose(context.Background(), req) + if err != nil { + t.Fatal(err) + } + if res.Model != "claude-sonnet-5-5" || len(models) != 1 || models[0] != "claude-sonnet-5-5" || res.RootCauseLabel != "db_primary_down" { + t.Fatalf("routing: res=%+v models=%v", res, models) + } + if res.Usage["input_tokens"] != 40 || res.Usage["output_tokens"] != 8 || res.Usage["total_tokens"] != 48 || res.Usage["turns"] != 1 { + t.Fatalf("usage: %v", res.Usage) + } +} diff --git a/ai/internal/anthropic/client.go b/ai/internal/anthropic/client.go new file mode 100644 index 0000000..0150523 --- /dev/null +++ b/ai/internal/anthropic/client.go @@ -0,0 +1,262 @@ +// Package anthropic adapts the Claude Messages API to the small Responses-API surface the agent loop uses +// (openai.Request / openai.Response with message, function_call and function_call_output items), so one loop +// serves both providers. The agent picks this client for model ids starting with "claude-". +// +// Wire format (Messages API reference, anthropic-version 2023-06-01): system prompt as the top-level `system` +// field, tools as {name, description, input_schema, strict}, tool calls as `tool_use` content blocks answered by +// `tool_result` blocks in the next user message. The Responses-API `previous_response_id` chain is emulated +// client-side: every response id maps to the full message list up to and including that assistant turn. +// +// Claude Opus 5.5 / Sonnet 5.5 reject forced tool use (tool_choice type "tool"/"any" -> 400), so a forced +// tool_choice from the loop becomes {type:"auto"} plus a one-line instruction in the user turn. +package anthropic + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "sync" + "time" + + "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/openai" +) + +const ( + apiVersion = "2023-06-01" + maxConvCached = 64 // response ids kept for previous_response_id chaining (one diagnosis uses <= MaxTurns) + defaultTokens = 4000 +) + +type Client struct { + BaseURL string + HTTP *http.Client + keyFn func() (string, error) + + mu sync.Mutex + conv map[string]conversation // response id -> messages (and system prompt) up to and including that assistant turn + order []string // insertion order for eviction + seq int +} + +type conversation struct { + system string + msgs []message +} + +func New(baseURL string, keyFn func() (string, error)) *Client { + if baseURL == "" { + baseURL = "https://api.anthropic.com" + } + return &Client{BaseURL: strings.TrimRight(baseURL, "/"), HTTP: &http.Client{Timeout: 150 * time.Second}, keyFn: keyFn, conv: map[string]conversation{}} +} + +// KeyPresent reports whether the key file has a non-empty value (the secret may not be registered yet). +func (c *Client) KeyPresent() bool { + k, err := c.keyFn() + return err == nil && k != "" +} + +type message struct { + Role string `json:"role"` + Content any `json:"content"` // string or []block +} + +type block struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` + ID string `json:"id,omitempty"` // tool_use + Name string `json:"name,omitempty"` // tool_use + Input json.RawMessage `json:"input,omitempty"` // tool_use + ToolUseID string `json:"tool_use_id,omitempty"` // tool_result + Content string `json:"content,omitempty"` // tool_result +} + +type tool struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` + InputSchema map[string]any `json:"input_schema"` + Strict bool `json:"strict,omitempty"` +} + +type request struct { + Model string `json:"model"` + MaxTokens int `json:"max_tokens"` + System string `json:"system,omitempty"` + Messages []message `json:"messages"` + Tools []tool `json:"tools,omitempty"` + ToolChoice map[string]any `json:"tool_choice,omitempty"` +} + +type response struct { + ID string `json:"id"` + Type string `json:"type"` + Model string `json:"model"` + StopReason string `json:"stop_reason"` + Content []block `json:"content"` + Usage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens"` + CacheReadInputTokens int `json:"cache_read_input_tokens"` + } `json:"usage"` + Error *struct { + Type string `json:"type"` + Message string `json:"message"` + } `json:"error"` +} + +// forcedNudge is appended to the user turn when the loop asked for a forced submit_diagnosis call. +const forcedNudge = "You must call the submit_diagnosis tool now with your final answer; do not call any other tool." + +// CreateResponse translates one loop turn to a Messages API call and the answer back to Responses-API items. +func (c *Client) CreateResponse(ctx context.Context, req openai.Request) (*openai.Response, error) { + key, err := c.keyFn() + if err != nil { + return nil, fmt.Errorf("anthropic key: %w", err) + } + if key == "" { + return nil, fmt.Errorf("anthropic api key is empty (secret selfheal-anthropic-api-key not registered yet)") + } + var msgs []message + var system []string + if req.PreviousResponseID != "" { + c.mu.Lock() + prev, ok := c.conv[req.PreviousResponseID] + c.mu.Unlock() + if !ok { + return nil, fmt.Errorf("anthropic: unknown previous_response_id %q", req.PreviousResponseID) + } + msgs = append(msgs, prev.msgs...) + if prev.system != "" { + system = append(system, prev.system) // the loop sends the system message only on the first turn + } + } + var userBlocks []block + flush := func() { + if len(userBlocks) == 0 { + return + } + msgs = append(msgs, message{Role: "user", Content: userBlocks}) + userBlocks = nil + } + for _, it := range req.Input { + switch it.Type { + case "message": + text, _ := it.Content.(string) + switch it.Role { + case "system": + system = append(system, text) + case "assistant": + flush() + msgs = append(msgs, message{Role: "assistant", Content: []block{{Type: "text", Text: text}}}) + default: + userBlocks = append(userBlocks, block{Type: "text", Text: text}) + } + case "function_call_output": + userBlocks = append(userBlocks, block{Type: "tool_result", ToolUseID: it.CallID, Content: it.Output}) + } + } + forced := false + if m, ok := req.ToolChoice.(map[string]any); ok && m["type"] == "function" { + forced = true + } + if forced { + userBlocks = append(userBlocks, block{Type: "text", Text: forcedNudge}) + } + flush() + if len(msgs) == 0 { + return nil, fmt.Errorf("anthropic: no messages to send") + } + body := request{Model: req.Model, MaxTokens: req.MaxOutputTokens, System: strings.Join(system, "\n\n"), Messages: msgs} + if body.MaxTokens <= 0 { + body.MaxTokens = defaultTokens + } + for _, t := range req.Tools { + body.Tools = append(body.Tools, tool{Name: t.Name, Description: t.Description, InputSchema: t.Parameters, Strict: t.Strict}) + } + if len(body.Tools) > 0 { + tc := map[string]any{"type": "auto"} + if req.ParallelToolCalls != nil && !*req.ParallelToolCalls { + tc["disable_parallel_tool_use"] = true + } + body.ToolChoice = tc + } + + b, _ := json.Marshal(body) + hreq, err := http.NewRequestWithContext(ctx, "POST", c.BaseURL+"/v1/messages", bytes.NewReader(b)) + if err != nil { + return nil, err + } + hreq.Header.Set("x-api-key", key) + hreq.Header.Set("anthropic-version", apiVersion) + hreq.Header.Set("Content-Type", "application/json") + resp, err := c.HTTP.Do(hreq) + if err != nil { + return nil, err + } + defer resp.Body.Close() + rb, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20)) + var r response + _ = json.Unmarshal(rb, &r) + if resp.StatusCode/100 != 2 { + msg := truncate(string(rb), 600) + if r.Error != nil { + msg = r.Error.Type + ": " + r.Error.Message + } + return nil, fmt.Errorf("anthropic /v1/messages: status %d: %s", resp.StatusCode, msg) + } + if r.Error != nil { + return nil, fmt.Errorf("anthropic response error: %s", r.Error.Message) + } + + out := &openai.Response{ID: r.ID, Status: "completed", Model: r.Model} + if out.ID == "" { + c.mu.Lock() + c.seq++ + out.ID = fmt.Sprintf("msg_local_%d", c.seq) + c.mu.Unlock() + } + var text strings.Builder + for _, bl := range r.Content { + switch bl.Type { + case "text": + text.WriteString(bl.Text) + case "tool_use": + args := string(bl.Input) + if args == "" { + args = "{}" + } + out.Output = append(out.Output, openai.Item{Type: "function_call", CallID: bl.ID, Name: bl.Name, Arguments: args, ID: bl.ID, Status: "completed"}) + } + } + if text.Len() > 0 { + out.Output = append([]openai.Item{{Type: "message", Role: "assistant", Content: []any{map[string]any{"type": "output_text", "text": text.String()}}}}, out.Output...) + } + // input tokens include cache writes/reads so the sum matches what is billed + out.Usage.InputTokens = r.Usage.InputTokens + r.Usage.CacheCreationInputTokens + r.Usage.CacheReadInputTokens + out.Usage.OutputTokens = r.Usage.OutputTokens + out.Usage.TotalTokens = out.Usage.InputTokens + out.Usage.OutputTokens + + // remember the conversation for the next turn (assistant content kept verbatim so tool_use ids line up) + full := append(append([]message{}, msgs...), message{Role: "assistant", Content: r.Content}) + c.mu.Lock() + c.conv[out.ID] = conversation{system: body.System, msgs: full} + c.order = append(c.order, out.ID) + for len(c.order) > maxConvCached { + delete(c.conv, c.order[0]) + c.order = c.order[1:] + } + c.mu.Unlock() + return out, nil +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] + "..." +} diff --git a/ai/internal/anthropic/client_test.go b/ai/internal/anthropic/client_test.go new file mode 100644 index 0000000..039ce9e --- /dev/null +++ b/ai/internal/anthropic/client_test.go @@ -0,0 +1,118 @@ +package anthropic + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/selfheal-lab/selfheal-lab-apps/ai/internal/openai" +) + +// Two turns through a fake Messages API: the first asks for a tool, the second (with previous_response_id and a +// forced submit) must carry the tool_result, the assistant tool_use block, the nudge text and an auto tool_choice. +func TestTranslateTurns(t *testing.T) { + var bodies []map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/messages" || r.Header.Get("x-api-key") != "k1" || r.Header.Get("anthropic-version") != "2023-06-01" { + t.Errorf("bad request line/headers: %s %v", r.URL.Path, r.Header) + } + var b map[string]any + _ = json.NewDecoder(r.Body).Decode(&b) + bodies = append(bodies, b) + switch len(bodies) { + case 1: + w.Write([]byte(`{"id":"msg_1","type":"message","model":"claude-sonnet-5-5","stop_reason":"tool_use","content":[{"type":"text","text":"checking"},{"type":"tool_use","id":"toolu_1","name":"get_probe_status","input":{}}],"usage":{"input_tokens":100,"output_tokens":20,"cache_creation_input_tokens":5,"cache_read_input_tokens":0}}`)) + default: + w.Write([]byte(`{"id":"msg_2","type":"message","model":"claude-sonnet-5-5","stop_reason":"tool_use","content":[{"type":"tool_use","id":"toolu_2","name":"submit_diagnosis","input":{"root_cause_label":"x","confidence":0.9}}],"usage":{"input_tokens":200,"output_tokens":30}}`)) + } + })) + defer srv.Close() + c := New(srv.URL, func() (string, error) { return "k1", nil }) + pf := false + tools := []openai.Tool{{Type: "function", Strict: true, Name: "get_probe_status", Description: "probe", Parameters: map[string]any{"type": "object", "properties": map[string]any{}}}} + r1, err := c.CreateResponse(context.Background(), openai.Request{Model: "claude-sonnet-5-5", MaxOutputTokens: 1234, ParallelToolCalls: &pf, Tools: tools, ToolChoice: "auto", + Input: []openai.Item{{Type: "message", Role: "system", Content: "SYS"}, {Type: "message", Role: "user", Content: "alert"}}}) + if err != nil { + t.Fatal(err) + } + fcs := r1.FunctionCalls() + if r1.ID != "msg_1" || len(fcs) != 1 || fcs[0].CallID != "toolu_1" || fcs[0].Name != "get_probe_status" || fcs[0].Arguments != "{}" { + t.Fatalf("turn 1 translation: %+v", r1) + } + if r1.Usage.InputTokens != 105 || r1.Usage.OutputTokens != 20 || r1.Usage.TotalTokens != 125 { + t.Fatalf("usage: %+v", r1.Usage) + } + b1 := bodies[0] + if b1["system"] != "SYS" || b1["max_tokens"].(float64) != 1234 || b1["model"] != "claude-sonnet-5-5" { + t.Fatalf("body 1: %v", b1) + } + tl := b1["tools"].([]any)[0].(map[string]any) + if tl["name"] != "get_probe_status" || tl["input_schema"] == nil || tl["strict"] != true { + t.Fatalf("tool translation: %v", tl) + } + if tc := b1["tool_choice"].(map[string]any); tc["type"] != "auto" || tc["disable_parallel_tool_use"] != true { + t.Fatalf("tool_choice 1: %v", tc) + } + + r2, err := c.CreateResponse(context.Background(), openai.Request{Model: "claude-sonnet-5-5", Tools: tools, PreviousResponseID: "msg_1", + ToolChoice: map[string]any{"type": "function", "name": "submit_diagnosis"}, + Input: []openai.Item{{Type: "function_call_output", CallID: "toolu_1", Output: `{"alb":1}`}}}) + if err != nil { + t.Fatal(err) + } + if got := r2.FunctionCalls(); len(got) != 1 || got[0].Name != "submit_diagnosis" || !strings.Contains(got[0].Arguments, `"root_cause_label":"x"`) { + t.Fatalf("turn 2 translation: %+v", r2.Output) + } + b2 := bodies[1] + msgs := b2["messages"].([]any) + if len(msgs) != 3 { // user(alert) / assistant(text+tool_use) / user(tool_result+nudge) + t.Fatalf("messages: %v", msgs) + } + as := msgs[1].(map[string]any) + if as["role"] != "assistant" || !strings.Contains(mustJSON(as["content"]), `"tool_use"`) { + t.Fatalf("assistant turn not replayed: %v", as) + } + last := msgs[2].(map[string]any)["content"].([]any) + if tr := last[0].(map[string]any); tr["type"] != "tool_result" || tr["tool_use_id"] != "toolu_1" || tr["content"] != `{"alb":1}` { + t.Fatalf("tool_result: %v", tr) + } + if nd := last[1].(map[string]any); nd["type"] != "text" || !strings.Contains(nd["text"].(string), "submit_diagnosis") { + t.Fatalf("forced nudge missing: %v", nd) + } + if tc := b2["tool_choice"].(map[string]any); tc["type"] != "auto" { + t.Fatalf("forced tool_choice must become auto for Claude 5.5: %v", tc) + } + if b2["system"] != "SYS" { + t.Fatalf("system prompt must be carried on the chained turn, got %v", b2["system"]) + } +} + +func TestEmptyKeyAndErrors(t *testing.T) { + c := New("http://127.0.0.1:1", func() (string, error) { return "", nil }) + if c.KeyPresent() { + t.Fatal("empty key reported present") + } + _, err := c.CreateResponse(context.Background(), openai.Request{Model: "claude-haiku-5-5", Input: []openai.Item{{Type: "message", Role: "user", Content: "x"}}}) + if err == nil || !strings.Contains(err.Error(), "anthropic api key is empty") { + t.Fatalf("expected empty-key error, got %v", err) + } + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(400) + w.Write([]byte(`{"type":"error","error":{"type":"invalid_request_error","message":"tool_choice: type \"tool\" and \"any\" are not supported for this model."}}`)) + })) + defer srv.Close() + c2 := New(srv.URL, func() (string, error) { return "k", nil }) + _, err = c2.CreateResponse(context.Background(), openai.Request{Model: "claude-haiku-5-5", Input: []openai.Item{{Type: "message", Role: "user", Content: "x"}}}) + if err == nil || !strings.Contains(err.Error(), "status 400") || !strings.Contains(err.Error(), "invalid_request_error") { + t.Fatalf("expected 400 error, got %v", err) + } + _, err = c2.CreateResponse(context.Background(), openai.Request{Model: "claude-haiku-5-5", PreviousResponseID: "nope", Input: []openai.Item{{Type: "function_call_output", CallID: "a", Output: "b"}}}) + if err == nil || !strings.Contains(err.Error(), "unknown previous_response_id") { + t.Fatalf("expected unknown id error, got %v", err) + } +} + +func mustJSON(v any) string { b, _ := json.Marshal(v); return string(b) } diff --git a/harness/alerts.go b/harness/alerts.go index e4b2750..1b24115 100644 --- a/harness/alerts.go +++ b/harness/alerts.go @@ -231,6 +231,9 @@ func (r *AlertRouter) diagnose(alert map[string]any, runID, sid *string) { ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute) defer cancel() reqBody := map[string]any{"alert": alert, "run_id": runID, "scenario_id": sid} + if cur := r.app.runs.current(); cur != nil && runID != nil && cur.RunID == *runID && cur.Model != "" { + reqBody["model"] = cur.Model // campaign-level model selection (the AI VM falls back to its default when absent) + } b, _ := json.Marshal(reqBody) req, _ := http.NewRequestWithContext(ctx, "POST", r.app.cfg.AIDiagnoseURL, bytes.NewReader(b)) req.Header.Set("Content-Type", "application/json") diff --git a/harness/run.go b/harness/run.go index 184db99..5e863ac 100644 --- a/harness/run.go +++ b/harness/run.go @@ -9,6 +9,7 @@ import ( "fmt" "log" "net/http" + "regexp" "strings" "sync" "time" @@ -25,6 +26,7 @@ type RunManager struct { type activeRun struct { RunID string Campaign string + Model string // LLM for this run's diagnoses ("" = AI VM default); passed to the AI VM on every /diagnose ScenarioID string // "" between scenarios Scenario *Scenario TInject time.Time @@ -33,6 +35,7 @@ type activeRun struct { type RunRequest struct { Campaign string `json:"campaign"` + Model string `json:"model"` // optional: LLM id for the diagnoses (e.g. gpt-6-luna, claude-sonnet-5-5); "" = AI VM default ScenarioIDs []string `json:"scenario_ids"` BaselineS *int `json:"baseline_s"` // default 300 (design 8) CooldownS *int `json:"cooldown_s"` // default 300 @@ -44,6 +47,9 @@ type RunRequest struct { SkipRevert bool `json:"skip_revert"` } +// modelRe bounds the model id (empty = AI VM default). The AI VM routes by prefix: claude-* -> Anthropic, else OpenAI. +var modelRe = regexp.MustCompile(`^[A-Za-z0-9._:-]{0,64}$`) + func NewRunManager(a *App) *RunManager { return &RunManager{app: a} } func (m *RunManager) current() *activeRun { @@ -78,6 +84,10 @@ func (a *App) handleRunCreate(w http.ResponseWriter, r *http.Request) { return } } + if !modelRe.MatchString(req.Model) { + writeErr(w, 400, errors.New("model must match [A-Za-z0-9._:-]{1,64}")) + return + } m := a.runs m.mu.Lock() if m.active != nil { @@ -88,9 +98,9 @@ func (a *App) handleRunCreate(w http.ResponseWriter, r *http.Request) { } runID := newID("run-") ctx, cancel := context.WithCancel(context.Background()) - m.active = &activeRun{RunID: runID, Campaign: req.Campaign, cancel: cancel} + m.active = &activeRun{RunID: runID, Campaign: req.Campaign, Model: req.Model, cancel: cancel} m.mu.Unlock() - if err := a.store.CreateRun(r.Context(), runID, req.Campaign, req.DriftBefore, req.GitSHA, req.Notes); err != nil { + if err := a.store.CreateRun(r.Context(), runID, req.Campaign, req.DriftBefore, req.GitSHA, req.Notes, req.Model); err != nil { m.mu.Lock() m.active = nil m.mu.Unlock() @@ -100,7 +110,7 @@ func (a *App) handleRunCreate(w http.ResponseWriter, r *http.Request) { } mRuns.Set(1) go m.execute(ctx, runID, req) - writeJSON(w, 202, map[string]any{"run_id": runID, "scenario_ids": req.ScenarioIDs}) + writeJSON(w, 202, map[string]any{"run_id": runID, "scenario_ids": req.ScenarioIDs, "model": req.Model}) } func (a *App) handleRunAbort(w http.ResponseWriter, r *http.Request) { diff --git a/harness/store.go b/harness/store.go index e68e669..fede06b 100644 --- a/harness/store.go +++ b/harness/store.go @@ -24,6 +24,7 @@ CREATE TABLE IF NOT EXISTS runs ( status text NOT NULL DEFAULT 'running', notes text ); +ALTER TABLE runs ADD COLUMN IF NOT EXISTS model text; -- LLM requested for the run's diagnoses (campaign-level model selection, 2026-10-14) CREATE TABLE IF NOT EXISTS scenarios ( scenario_id text PRIMARY KEY, path text NOT NULL, @@ -119,8 +120,8 @@ func (s *Store) UpsertScenarios(ctx context.Context, m map[string]*Scenario) err return nil } -func (s *Store) CreateRun(ctx context.Context, id, campaign string, driftBefore *int, gitSHA, notes string) error { - _, err := s.pool.Exec(ctx, `INSERT INTO runs (run_id,campaign,drift_before,git_sha,notes) VALUES ($1,$2,$3,$4,$5)`, id, campaign, driftBefore, nz(gitSHA), nz(notes)) +func (s *Store) CreateRun(ctx context.Context, id, campaign string, driftBefore *int, gitSHA, notes, model string) error { + _, err := s.pool.Exec(ctx, `INSERT INTO runs (run_id,campaign,drift_before,git_sha,notes,model) VALUES ($1,$2,$3,$4,$5,$6)`, id, campaign, driftBefore, nz(gitSHA), nz(notes), nz(model)) return err } @@ -222,10 +223,11 @@ type RunRow struct { GitSHA *string `json:"git_sha"` Status string `json:"status"` Notes *string `json:"notes"` + Model *string `json:"model"` // nil for runs before campaign-level model selection } func (s *Store) Runs(ctx context.Context, runID string) ([]RunRow, error) { - rows, err := s.pool.Query(ctx, `SELECT run_id,campaign,started_at,finished_at,drift_before,drift_after,git_sha,status,notes FROM runs WHERE ($1='' OR run_id=$1) ORDER BY started_at DESC LIMIT 100`, runID) + rows, err := s.pool.Query(ctx, `SELECT run_id,campaign,started_at,finished_at,drift_before,drift_after,git_sha,status,notes,model FROM runs WHERE ($1='' OR run_id=$1) ORDER BY started_at DESC LIMIT 100`, runID) if err != nil { return nil, err } @@ -233,7 +235,7 @@ func (s *Store) Runs(ctx context.Context, runID string) ([]RunRow, error) { var out []RunRow for rows.Next() { var r RunRow - if err := rows.Scan(&r.RunID, &r.Campaign, &r.StartedAt, &r.FinishedAt, &r.DriftBefore, &r.DriftAfter, &r.GitSHA, &r.Status, &r.Notes); err != nil { + if err := rows.Scan(&r.RunID, &r.Campaign, &r.StartedAt, &r.FinishedAt, &r.DriftBefore, &r.DriftAfter, &r.GitSHA, &r.Status, &r.Notes, &r.Model); err != nil { return nil, err } out = append(out, r)