selfheal-lab-apps/ai/internal/anthropic/client.go

// 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)
	fed     *Federation // nil = API key from keyFn; set = Workload Identity Federation (takes precedence)

	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 != ""
}

// SetFederation switches the client to Workload Identity Federation (ignored when the ids are incomplete).
func (c *Client) SetFederation(f *Federation) {
	if f.Configured() {
		c.fed = f
	}
}

// AuthMode is "federation", "api_key" or "none" (reported by /selfcheck).
func (c *Client) AuthMode() string {
	if c.fed != nil {
		return "federation"
	}
	if c.KeyPresent() {
		return "api_key"
	}
	return "none"
}

// Federation returns the active federation config (nil when the API key path is used).
func (c *Client) Federation() *Federation { return c.fed }

// authorize sets the credential header: a federated bearer token, or the API key.
func (c *Client) authorize(ctx context.Context, req *http.Request) error {
	if c.fed != nil {
		tok, err := c.fed.AccessToken(ctx, c.BaseURL)
		if err != nil {
			return fmt.Errorf("anthropic federation: %w", err)
		}
		req.Header.Set("Authorization", "Bearer "+tok)
		return nil
	}
	key, err := c.keyFn()
	if err != nil {
		return fmt.Errorf("anthropic key: %w", err)
	}
	if key == "" {
		return fmt.Errorf("anthropic api key is empty (secret selfheal-anthropic-api-key not registered yet) and federation is not configured")
	}
	req.Header.Set("x-api-key", key)
	return nil
}

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) {
	if c.fed == nil {
		if key, err := c.keyFn(); err != nil {
			return nil, fmt.Errorf("anthropic key: %w", err)
		} else if key == "" {
			return nil, fmt.Errorf("anthropic api key is empty (secret selfheal-anthropic-api-key not registered yet) and federation is not configured")
		}
	}
	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)
	send := func() (*http.Response, error) {
		hreq, err := http.NewRequestWithContext(ctx, "POST", c.BaseURL+"/v1/messages", bytes.NewReader(b))
		if err != nil {
			return nil, err
		}
		if err := c.authorize(ctx, hreq); err != nil {
			return nil, err
		}
		hreq.Header.Set("anthropic-version", apiVersion)
		hreq.Header.Set("Content-Type", "application/json")
		return c.HTTP.Do(hreq)
	}
	resp, err := send()
	if err != nil {
		return nil, err
	}
	if resp.StatusCode == 401 && c.fed != nil { // token revoked/expired early: exchange once more
		resp.Body.Close()
		c.fed.Invalidate()
		if resp, err = send(); 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] + "..."
}