iop/apps/edge/internal/openai/server.go

package openai

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"html"
	"io"
	"net"
	"net/http"
	"regexp"
	"strings"
	"time"

	"go.uber.org/zap"

	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
)

type runService interface {
	SubmitRun(context.Context, edgeservice.SubmitRunRequest) (*edgeservice.RunHandle, error)
	OllamaAPI(context.Context, edgeservice.OllamaAPIRequest) (edgeservice.OllamaAPIView, error)
}

type Server struct {
	cfg     config.EdgeOpenAIConf
	service runService
	logger  *zap.Logger
	server  *http.Server
}

func NewServer(cfg config.EdgeOpenAIConf, svc runService, logger *zap.Logger) *Server {
	if logger == nil {
		logger = zap.NewNop()
	}
	return &Server{cfg: cfg, service: svc, logger: logger}
}

func (s *Server) Enabled() bool {
	return s != nil && s.cfg.Enabled
}

func (s *Server) Start(ctx context.Context) error {
	if !s.Enabled() {
		return nil
	}
	if s.cfg.Listen == "" {
		s.cfg.Listen = "0.0.0.0:8080"
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/healthz", s.handleHealthz)
	mux.HandleFunc("/v1/models", s.handleModels)
	mux.HandleFunc("/v1/chat/completions", s.handleChatCompletions)
	mux.HandleFunc("/api/", s.handleOllamaAPI)

	s.server = &http.Server{
		Addr:              s.cfg.Listen,
		Handler:           mux,
		ReadHeaderTimeout: 5 * time.Second,
	}
	ln, err := net.Listen("tcp", s.cfg.Listen)
	if err != nil {
		return fmt.Errorf("openai server listen %s: %w", s.cfg.Listen, err)
	}
	go func() {
		<-ctx.Done()
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_ = s.server.Shutdown(shutdownCtx)
	}()
	go func() {
		s.logger.Info("openai-compatible server listening", zap.String("addr", s.cfg.Listen))
		if err := s.server.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
			s.logger.Warn("openai-compatible server exited", zap.Error(err))
		}
	}()
	return nil
}

func (s *Server) Stop(ctx context.Context) error {
	if s == nil || s.server == nil {
		return nil
	}
	return s.server.Shutdown(ctx)
}

func (s *Server) handleHealthz(w http.ResponseWriter, _ *http.Request) {
	writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
}

func (s *Server) handleModels(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "method not allowed")
		return
	}
	models := s.cfg.Models
	if len(models) == 0 && s.cfg.Target != "" {
		models = []string{s.cfg.Target}
	}
	data := make([]openAIModel, 0, len(models))
	for _, model := range models {
		model = strings.TrimSpace(model)
		if model == "" {
			continue
		}
		data = append(data, openAIModel{
			ID:      model,
			Object:  "model",
			Created: time.Now().Unix(),
			OwnedBy: "iop",
		})
	}
	writeJSON(w, http.StatusOK, openAIModelsResponse{Object: "list", Data: data})
}

func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "method not allowed")
		return
	}
	defer r.Body.Close()

	var req chatCompletionRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", "invalid JSON request")
		return
	}
	target := s.resolveTarget(req.Model)
	if target == "" {
		writeError(w, http.StatusBadRequest, "invalid_request_error", "model is required")
		return
	}
	basePrompt := promptFromMessages(req.Messages)
	if strings.TrimSpace(basePrompt) == "" {
		writeError(w, http.StatusBadRequest, "invalid_request_error", "messages are required")
		return
	}
	outputPolicy := s.resolveOutputPolicy(basePrompt)
	messages := req.Messages
	if instruction := strictOutputContractInstruction(outputPolicy); instruction != "" {
		messages = prependSystemMessage(messages, instruction)
	}
	prompt := promptFromMessages(messages)
	input := req.runInput(prompt, messages, outputPolicy.Strict)
	s.logger.Info("openai chat completion input",
		zap.String("model", req.Model),
		zap.String("target", target),
		zap.String("adapter", s.resolveAdapter()),
		zap.Bool("strict_output", outputPolicy.Strict),
		zap.Bool("strict_stream_buffer", outputPolicy.StreamBuffer),
		zap.String("xml_completion_tool", outputPolicy.XMLCompletionTool),
		zap.Bool("contract_instruction", outputPolicy.ContractInstruction),
		zap.Bool("stream", req.Stream),
		zap.Int("message_count", len(req.Messages)),
		zap.Int("prompt_len", len(prompt)),
		zap.String("prompt_preview", previewString(prompt, 1000)),
		zap.Any("input_keys", mapKeys(input)),
	)

	handle, err := s.service.SubmitRun(r.Context(), edgeservice.SubmitRunRequest{
		NodeRef:    s.cfg.NodeRef,
		Adapter:    s.resolveAdapter(),
		Target:     target,
		SessionID:  s.resolveSessionID(),
		Prompt:     prompt,
		Input:      input,
		TimeoutSec: s.resolveTimeoutSec(),
		Metadata: map[string]string{
			"source":        "openai",
			"openai_model":  req.Model,
			"openai_stream": fmt.Sprintf("%t", req.Stream),
			"strict_output": fmt.Sprintf("%t", outputPolicy.Strict),
		},
	})
	if err != nil {
		writeError(w, http.StatusBadGateway, "node_dispatch_error", err.Error())
		return
	}
	defer handle.Close()

	if req.Stream {
		s.streamChatCompletion(w, r, req, handle, outputPolicy)
		return
	}
	s.completeChatCompletion(w, r, req, handle, outputPolicy)
}

func (s *Server) handleOllamaAPI(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet && r.Method != http.MethodPost && r.Method != http.MethodDelete {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "method not allowed")
		return
	}
	defer r.Body.Close()
	body, err := io.ReadAll(io.LimitReader(r.Body, 16<<20))
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", "read request body failed")
		return
	}
	resp, err := s.service.OllamaAPI(r.Context(), edgeservice.OllamaAPIRequest{
		NodeRef:    s.cfg.NodeRef,
		Adapter:    s.resolveAdapter(),
		Method:     r.Method,
		Path:       r.URL.RequestURI(),
		Body:       string(body),
		TimeoutSec: s.resolveTimeoutSec(),
	})
	if err != nil {
		writeError(w, http.StatusBadGateway, "ollama_passthrough_error", err.Error())
		return
	}
	contentType := resp.ContentType
	if contentType == "" {
		contentType = "application/json"
	}
	w.Header().Set("Content-Type", contentType)
	status := resp.StatusCode
	if status < 100 || status > 999 {
		status = http.StatusOK
	}
	w.WriteHeader(status)
	_, _ = w.Write([]byte(resp.Body))
}

func (s *Server) completeChatCompletion(w http.ResponseWriter, r *http.Request, req chatCompletionRequest, handle *edgeservice.RunHandle, outputPolicy strictOutputPolicy) {
	text, reasoning, usage, err := collectRunResult(r.Context(), handle)
	if err != nil {
		writeError(w, httpStatusForRunError(err), "run_error", err.Error())
		return
	}
	text, reasoning, normalized := normalizeCompletionOutput(outputPolicy, text, reasoning)
	s.logger.Info("openai chat completion output",
		zap.String("run_id", handle.RunID),
		zap.Bool("strict_output", outputPolicy.Strict),
		zap.Bool("strict_stream_buffer", outputPolicy.StreamBuffer),
		zap.String("xml_completion_tool", outputPolicy.XMLCompletionTool),
		zap.Bool("normalized", normalized),
		zap.Int("content_len", len(text)),
		zap.Int("reasoning_len", len(reasoning)),
		zap.String("content_preview", previewString(text, 1000)),
	)
	created := time.Now().Unix()
	writeJSON(w, http.StatusOK, chatCompletionResponse{
		ID:      "chatcmpl-" + handle.RunID,
		Object:  "chat.completion",
		Created: created,
		Model:   responseModel(req.Model, handle.Target),
		Choices: []chatCompletionChoice{{
			Index:        0,
			Message:      chatMessage{Role: "assistant", Content: text, ReasoningContent: reasoning},
			FinishReason: "stop",
		}},
		Usage: usage,
	})
}

func (s *Server) streamChatCompletion(w http.ResponseWriter, r *http.Request, req chatCompletionRequest, handle *edgeservice.RunHandle, outputPolicy strictOutputPolicy) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		writeError(w, http.StatusInternalServerError, "streaming_not_supported", "response writer does not support streaming")
		return
	}
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")

	created := time.Now().Unix()
	id := "chatcmpl-" + handle.RunID
	model := responseModel(req.Model, handle.Target)
	writeSSE(w, flusher, chatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []chatCompletionChunkChoice{{
			Index: 0,
			Delta: chatDelta{Role: "assistant"},
		}},
	})

	if outputPolicy.Strict && outputPolicy.StreamBuffer {
		s.streamBufferedChatCompletion(w, r, handle, id, created, model, flusher, outputPolicy)
		return
	}

	var contentBuilder strings.Builder
	var reasoningBuilder strings.Builder
	defer func() {
		s.logger.Info("openai chat completion stream closed",
			zap.String("run_id", handle.RunID),
			zap.Bool("strict_output", outputPolicy.Strict),
			zap.Bool("strict_stream_buffer", outputPolicy.StreamBuffer),
			zap.Int("content_len", contentBuilder.Len()),
			zap.Int("reasoning_len", reasoningBuilder.Len()),
			zap.String("content_preview", previewString(contentBuilder.String(), 1000)),
			zap.String("reasoning_preview", previewString(reasoningBuilder.String(), 1000)),
		)
	}()

	for {
		select {
		case <-r.Context().Done():
			return
		case nodeEvent := <-handle.NodeEvents:
			if edgeservice.IsNodeDisconnected(nodeEvent) {
				writeSSEError(w, flusher, "node disconnected")
				return
			}
		case event := <-handle.Events:
			if event == nil {
				continue
			}
			switch event.GetType() {
			case "delta":
				if event.GetDelta() == "" {
					continue
				}
				contentBuilder.WriteString(event.GetDelta())
				writeSSE(w, flusher, chatCompletionChunk{
					ID:      id,
					Object:  "chat.completion.chunk",
					Created: created,
					Model:   model,
					Choices: []chatCompletionChunkChoice{{
						Index: 0,
						Delta: chatDelta{Content: event.GetDelta()},
					}},
				})
			case "reasoning_delta":
				if event.GetDelta() == "" {
					continue
				}
				reasoningBuilder.WriteString(event.GetDelta())
				if outputPolicy.Strict {
					continue
				}
				writeSSE(w, flusher, chatCompletionChunk{
					ID:      id,
					Object:  "chat.completion.chunk",
					Created: created,
					Model:   model,
					Choices: []chatCompletionChunkChoice{{
						Index: 0,
						Delta: chatDelta{ReasoningContent: event.GetDelta()},
					}},
				})
			case "complete":
				writeSSE(w, flusher, chatCompletionChunk{
					ID:      id,
					Object:  "chat.completion.chunk",
					Created: created,
					Model:   model,
					Choices: []chatCompletionChunkChoice{{
						Index:        0,
						Delta:        chatDelta{},
						FinishReason: "stop",
					}},
				})
				fmt.Fprint(w, "data: [DONE]\n\n")
				flusher.Flush()
				return
			case "error", "cancelled":
				msg := event.GetError()
				if msg == "" {
					msg = event.GetMessage()
				}
				if msg == "" {
					msg = "run failed"
				}
				writeSSEError(w, flusher, msg)
				return
			}
		case <-time.After(handle.WaitTimeout()):
			writeSSEError(w, flusher, "run timed out")
			return
		}
	}
}

func (s *Server) streamBufferedChatCompletion(w http.ResponseWriter, r *http.Request, handle *edgeservice.RunHandle, id string, created int64, model string, flusher http.Flusher, outputPolicy strictOutputPolicy) {
	text, reasoning, _, err := collectRunResult(r.Context(), handle)
	if err != nil {
		writeSSEError(w, flusher, err.Error())
		return
	}
	text, reasoning, normalized := normalizeCompletionOutput(outputPolicy, text, reasoning)
	s.logger.Info("openai chat completion stream closed",
		zap.String("run_id", handle.RunID),
		zap.Bool("strict_output", outputPolicy.Strict),
		zap.Bool("strict_stream_buffer", outputPolicy.StreamBuffer),
		zap.String("xml_completion_tool", outputPolicy.XMLCompletionTool),
		zap.Bool("normalized", normalized),
		zap.Int("content_len", len(text)),
		zap.Int("reasoning_len", len(reasoning)),
		zap.String("content_preview", previewString(text, 1000)),
		zap.String("reasoning_preview", previewString(reasoning, 1000)),
	)
	if text != "" {
		writeSSE(w, flusher, chatCompletionChunk{
			ID:      id,
			Object:  "chat.completion.chunk",
			Created: created,
			Model:   model,
			Choices: []chatCompletionChunkChoice{{
				Index: 0,
				Delta: chatDelta{Content: text},
			}},
		})
	}
	writeSSE(w, flusher, chatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []chatCompletionChunkChoice{{
			Index:        0,
			Delta:        chatDelta{},
			FinishReason: "stop",
		}},
	})
	fmt.Fprint(w, "data: [DONE]\n\n")
	flusher.Flush()
}

func collectRunResult(ctx context.Context, handle *edgeservice.RunHandle) (string, string, *openAIUsage, error) {
	var contentBuilder strings.Builder
	var reasoningBuilder strings.Builder
	var usage *openAIUsage
	timer := time.NewTimer(handle.WaitTimeout())
	defer timer.Stop()
	for {
		select {
		case <-ctx.Done():
			return "", "", nil, ctx.Err()
		case <-timer.C:
			return "", "", nil, fmt.Errorf("run timed out")
		case nodeEvent := <-handle.NodeEvents:
			if edgeservice.IsNodeDisconnected(nodeEvent) {
				return "", "", nil, fmt.Errorf("node disconnected")
			}
		case event := <-handle.Events:
			if event == nil {
				continue
			}
			switch event.GetType() {
			case "delta":
				contentBuilder.WriteString(event.GetDelta())
			case "reasoning_delta":
				reasoningBuilder.WriteString(event.GetDelta())
			case "complete":
				if u := event.GetUsage(); u != nil {
					usage = &openAIUsage{
						PromptTokens:     int(u.GetInputTokens()),
						CompletionTokens: int(u.GetOutputTokens()),
						TotalTokens:      int(u.GetInputTokens() + u.GetOutputTokens()),
					}
				}
				return contentBuilder.String(), reasoningBuilder.String(), usage, nil
			case "error", "cancelled":
				msg := event.GetError()
				if msg == "" {
					msg = event.GetMessage()
				}
				if msg == "" {
					msg = "run failed"
				}
				return "", "", nil, fmt.Errorf("%s", msg)
			}
		}
	}
}

func (s *Server) resolveAdapter() string {
	if s.cfg.Adapter != "" {
		return s.cfg.Adapter
	}
	return "ollama"
}

func (s *Server) resolveTarget(model string) string {
	if s.cfg.Target != "" {
		return s.cfg.Target
	}
	return strings.TrimSpace(model)
}

func (s *Server) resolveSessionID() string {
	if s.cfg.SessionID != "" {
		return s.cfg.SessionID
	}
	return edgeservice.DefaultSessionID
}

func (s *Server) resolveTimeoutSec() int {
	if s.cfg.TimeoutSec > 0 {
		return s.cfg.TimeoutSec
	}
	return edgeservice.DefaultTimeoutSec
}

func (s *Server) resolveStrictOutput() bool {
	return s.cfg.StrictOutput
}

func (s *Server) resolveStrictStreamBuffer() bool {
	return s.cfg.StrictStreamBuffer
}

func (s *Server) resolveOutputPolicy(prompt string) strictOutputPolicy {
	policy := strictOutputPolicy{
		Strict:       s.resolveStrictOutput(),
		StreamBuffer: s.resolveStrictStreamBuffer(),
	}
	if !policy.Strict {
		return policy
	}
	policy.XMLCompletionTool, policy.XMLResultTag = inferXMLCompletionContract(prompt)
	policy.ContractInstruction = policy.XMLCompletionTool != ""
	return policy
}

func promptFromMessages(messages []chatMessage) string {
	var b strings.Builder
	for _, msg := range messages {
		content := strings.TrimSpace(msg.Content)
		if content == "" {
			continue
		}
		role := strings.TrimSpace(msg.Role)
		if role == "" {
			role = "user"
		}
		if b.Len() > 0 {
			b.WriteString("\n")
		}
		b.WriteString(role)
		b.WriteString(": ")
		b.WriteString(content)
	}
	return b.String()
}

func responseModel(requestModel, target string) string {
	if requestModel != "" {
		return requestModel
	}
	return target
}

func httpStatusForRunError(err error) int {
	if errors.Is(err, context.Canceled) {
		return http.StatusRequestTimeout
	}
	return http.StatusBadGateway
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}

func writeError(w http.ResponseWriter, status int, code, message string) {
	writeJSON(w, status, errorResponse{Error: errorBody{Type: code, Message: message}})
}

func writeSSE(w http.ResponseWriter, flusher http.Flusher, v any) {
	b, err := json.Marshal(v)
	if err != nil {
		return
	}
	fmt.Fprintf(w, "data: %s\n\n", b)
	flusher.Flush()
}

func writeSSEError(w http.ResponseWriter, flusher http.Flusher, message string) {
	writeSSE(w, flusher, errorResponse{Error: errorBody{Type: "run_error", Message: message}})
	fmt.Fprint(w, "data: [DONE]\n\n")
	flusher.Flush()
}

type chatCompletionRequest struct {
	Model     string         `json:"model"`
	Messages  []chatMessage  `json:"messages"`
	Stream    bool           `json:"stream"`
	Options   map[string]any `json:"options,omitempty"`
	Format    any            `json:"format,omitempty"`
	KeepAlive any            `json:"keep_alive,omitempty"`
	Think     any            `json:"think,omitempty"`
	Tools     []any          `json:"tools,omitempty"`
}

type chatMessage struct {
	Role             string `json:"role"`
	Content          string `json:"content"`
	ReasoningContent string `json:"reasoning_content,omitempty"`
	Images           []any  `json:"images,omitempty"`
	ToolCalls        []any  `json:"tool_calls,omitempty"`
	ToolName         string `json:"tool_name,omitempty"`
}

func (m *chatMessage) UnmarshalJSON(b []byte) error {
	var raw struct {
		Role             string `json:"role"`
		Content          any    `json:"content"`
		ReasoningContent string `json:"reasoning_content,omitempty"`
		Images           []any  `json:"images,omitempty"`
		ToolCalls        []any  `json:"tool_calls,omitempty"`
		ToolName         string `json:"tool_name,omitempty"`
	}
	if err := json.Unmarshal(b, &raw); err != nil {
		return err
	}
	m.Role = raw.Role
	m.Content = contentToString(raw.Content)
	m.ReasoningContent = raw.ReasoningContent
	m.Images = raw.Images
	m.ToolCalls = raw.ToolCalls
	m.ToolName = raw.ToolName
	return nil
}

func (req chatCompletionRequest) runInput(prompt string, messages []chatMessage, strictOutput bool) map[string]any {
	input := map[string]any{
		"prompt":   prompt,
		"messages": chatMessagesInput(messages),
	}
	if len(req.Options) > 0 {
		input["options"] = req.Options
	}
	if req.Format != nil {
		input["format"] = req.Format
	}
	if req.KeepAlive != nil {
		input["keep_alive"] = req.KeepAlive
	}
	if req.Think != nil {
		input["think"] = req.Think
	} else if strictOutput {
		input["think"] = false
	}
	if len(req.Tools) > 0 {
		input["tools"] = req.Tools
	}
	return input
}

type strictOutputPolicy struct {
	Strict              bool
	StreamBuffer        bool
	XMLCompletionTool   string
	XMLResultTag        string
	ContractInstruction bool
}

func strictOutputContractInstruction(policy strictOutputPolicy) string {
	if !policy.ContractInstruction {
		return ""
	}
	resultTag := policy.XMLResultTag
	if resultTag == "" {
		resultTag = "result"
	}
	return fmt.Sprintf("Strict output contract: follow the client's XML tool protocol. When you complete the task, output exactly one <%s> block and no text outside that XML block. The <%s> content must be the direct final answer that should be shown to the end user. Do not summarize that you answered, completed, explained, greeted, or performed the task. Do not write phrases like \"I completed\", \"I explained\", \"I greeted\", \"summary\", or \"main points\" unless the user explicitly asks for a summary. Put the actual user-facing response inside <%s>.", policy.XMLCompletionTool, resultTag, resultTag)
}

func prependSystemMessage(messages []chatMessage, content string) []chatMessage {
	out := make([]chatMessage, 0, len(messages)+1)
	out = append(out, chatMessage{Role: "system", Content: content})
	out = append(out, messages...)
	return out
}

func normalizeCompletionOutput(policy strictOutputPolicy, content, reasoning string) (string, string, bool) {
	if !policy.Strict {
		return content, reasoning, false
	}
	normalized := normalizeStrictAgentContent(content)
	if strings.TrimSpace(normalized) == "" && strings.TrimSpace(reasoning) != "" {
		normalized = normalizeStrictAgentContent(reasoning)
	}
	if strings.TrimSpace(normalized) == "" {
		normalized = "Model returned no assistant content."
	}
	if policy.XMLCompletionTool != "" && !startsWithXMLRoot(normalized) {
		normalized = wrapXMLCompletion(policy, normalized)
	}
	return normalized, "", normalized != content || reasoning != ""
}

var (
	agentThinkingBlockRE       = regexp.MustCompile(`(?is)<think(?:ing)?\b[^>]*>.*?</think(?:ing)?\s*>`)
	agentCodeFenceLineRE       = regexp.MustCompile(`(?m)^\s*` + "```" + `[A-Za-z0-9_-]*\s*$`)
	agentXMLOpenTagRE          = regexp.MustCompile(`(?is)<\s*([A-Za-z_][A-Za-z0-9_.:-]*)\b[^>]*>`)
	agentXMLCloseTagRE         = regexp.MustCompile(`(?is)</\s*([A-Za-z_][A-Za-z0-9_.:\s-]*)\s*>`)
	completedUseToolRE         = regexp.MustCompile("(?is)completed.{0,300}?use\\s+(?:the\\s+)?`?([A-Za-z_][A-Za-z0-9_.:-]*)`?\\s+tool")
	xmlToolWithResultREPattern = `(?is)<\s*%s\b[^>]*>\s*<\s*([A-Za-z_][A-Za-z0-9_.:-]*)\b`
)

func normalizeStrictAgentContent(content string) string {
	cleaned := strings.TrimSpace(content)
	cleaned = agentThinkingBlockRE.ReplaceAllString(cleaned, "")
	cleaned = agentCodeFenceLineRE.ReplaceAllString(cleaned, "")
	cleaned = strings.TrimSpace(cleaned)
	for strings.HasPrefix(cleaned, "---") {
		cleaned = strings.TrimSpace(strings.TrimPrefix(cleaned, "---"))
	}
	cleaned = strings.TrimSpace(cleaned)
	if block, ok := firstXMLRootBlock(cleaned); ok {
		return block
	}
	return cleaned
}

func startsWithXMLRoot(content string) bool {
	trimmed := strings.TrimSpace(content)
	open := agentXMLOpenTagRE.FindStringIndex(trimmed)
	return open != nil && open[0] == 0
}

func inferXMLCompletionContract(prompt string) (string, string) {
	tool := ""
	if match := completedUseToolRE.FindStringSubmatch(prompt); len(match) == 2 {
		tool = match[1]
	}
	if tool == "" {
		return "", ""
	}
	resultTag := inferXMLResultTag(prompt, tool)
	if resultTag == "" {
		resultTag = "result"
	}
	return tool, resultTag
}

func inferXMLResultTag(prompt, tool string) string {
	toolPattern := regexp.QuoteMeta(tool)
	re := regexp.MustCompile(fmt.Sprintf(xmlToolWithResultREPattern, toolPattern))
	if match := re.FindStringSubmatch(prompt); len(match) == 2 {
		return match[1]
	}
	return ""
}

func wrapXMLCompletion(policy strictOutputPolicy, content string) string {
	resultTag := policy.XMLResultTag
	if resultTag == "" {
		resultTag = "result"
	}
	return "<" + policy.XMLCompletionTool + ">\n<" + resultTag + ">" + html.EscapeString(strings.TrimSpace(content)) + "</" + resultTag + ">\n</" + policy.XMLCompletionTool + ">"
}

func firstXMLRootBlock(content string) (string, bool) {
	open := agentXMLOpenTagRE.FindStringSubmatchIndex(content)
	if open == nil {
		return "", false
	}
	root := content[open[2]:open[3]]
	search := content[open[1]:]
	for _, close := range agentXMLCloseTagRE.FindAllStringSubmatchIndex(search, -1) {
		name := search[close[2]:close[3]]
		if canonicalXMLName(name) != canonicalXMLName(root) {
			continue
		}
		closeStart := open[1] + close[0]
		block := content[open[0]:closeStart] + "</" + root + ">"
		return strings.TrimSpace(block), true
	}
	return "", false
}

func canonicalXMLName(name string) string {
	var b strings.Builder
	for _, r := range strings.ToLower(name) {
		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') {
			b.WriteRune(r)
		}
	}
	return b.String()
}

func chatMessagesInput(messages []chatMessage) []any {
	out := make([]any, 0, len(messages))
	for _, msg := range messages {
		m := map[string]any{
			"role":    msg.Role,
			"content": msg.Content,
		}
		if msg.ReasoningContent != "" {
			m["thinking"] = msg.ReasoningContent
		}
		if len(msg.Images) > 0 {
			m["images"] = msg.Images
		}
		if len(msg.ToolCalls) > 0 {
			m["tool_calls"] = msg.ToolCalls
		}
		if msg.ToolName != "" {
			m["tool_name"] = msg.ToolName
		}
		out = append(out, m)
	}
	return out
}

func contentToString(v any) string {
	switch t := v.(type) {
	case string:
		return t
	case []any:
		parts := make([]string, 0, len(t))
		for _, item := range t {
			if m, ok := item.(map[string]any); ok {
				if text, ok := m["text"].(string); ok {
					parts = append(parts, text)
				}
			}
		}
		return strings.Join(parts, "\n")
	default:
		b, _ := json.Marshal(t)
		return string(b)
	}
}

func previewString(s string, max int) string {
	if max <= 0 || len(s) <= max {
		return s
	}
	return s[:max] + "...[truncated]"
}

func mapKeys(m map[string]any) []string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	return keys
}

type chatCompletionResponse struct {
	ID      string                 `json:"id"`
	Object  string                 `json:"object"`
	Created int64                  `json:"created"`
	Model   string                 `json:"model"`
	Choices []chatCompletionChoice `json:"choices"`
	Usage   *openAIUsage           `json:"usage,omitempty"`
}

type chatCompletionChoice struct {
	Index        int         `json:"index"`
	Message      chatMessage `json:"message"`
	FinishReason string      `json:"finish_reason"`
}

type chatCompletionChunk struct {
	ID      string                      `json:"id"`
	Object  string                      `json:"object"`
	Created int64                       `json:"created"`
	Model   string                      `json:"model"`
	Choices []chatCompletionChunkChoice `json:"choices"`
}

type chatCompletionChunkChoice struct {
	Index        int       `json:"index"`
	Delta        chatDelta `json:"delta"`
	FinishReason string    `json:"finish_reason,omitempty"`
}

type chatDelta struct {
	Role             string `json:"role,omitempty"`
	Content          string `json:"content,omitempty"`
	ReasoningContent string `json:"reasoning_content,omitempty"`
}

type openAIUsage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

type openAIModel struct {
	ID      string `json:"id"`
	Object  string `json:"object"`
	Created int64  `json:"created"`
	OwnedBy string `json:"owned_by"`
}

type openAIModelsResponse struct {
	Object string        `json:"object"`
	Data   []openAIModel `json:"data"`
}

type errorResponse struct {
	Error errorBody `json:"error"`
}

type errorBody struct {
	Type    string `json:"type"`
	Message string `json:"message"`
}