llmgw/internal/router/responses.go

package router

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync/atomic"
	"time"
)

// The Responses API, as a translation of chat completions.
//
// Codex speaks it, and newer versions speak nothing else. Every engine here
// serves chat completions, so a Responses request is rewritten into one and
// sent down the same path as any other — quotas, admission, routing, usage
// all as they are — and what comes back is rewritten into the Responses
// shape: a response object, or its stream of typed events.
//
// What is translated is what a coding agent uses: instructions, messages,
// function tools and their calls and outputs, Codex's free-form "custom"
// tools, and usage. What the engines have no equivalent for — stored
// responses, server-side tools, reasoning summaries — is ignored rather than
// refused, so a client that sends them still gets an answer.

// respRequest is the part of a Responses request that is translated.
type respRequest struct {
	Model             string            `json:"model"`
	Instructions      string            `json:"instructions"`
	Input             json.RawMessage   `json:"input"`
	Tools             []respTool        `json:"tools"`
	ToolChoice        json.RawMessage   `json:"tool_choice"`
	Stream            bool              `json:"stream"`
	Temperature       *float64          `json:"temperature"`
	TopP              *float64          `json:"top_p"`
	MaxOutputTokens   *int              `json:"max_output_tokens"`
	ParallelToolCalls *bool             `json:"parallel_tool_calls"`
	Metadata          map[string]string `json:"metadata"`
}

type respTool struct {
	Type        string          `json:"type"`
	Name        string          `json:"name"`
	Description string          `json:"description"`
	Parameters  json.RawMessage `json:"parameters"`
}

// respItem is one input item: a message, a call the model made, or its
// output. Content is a string or a list of parts.
type respItem struct {
	Type    string          `json:"type"`
	Role    string          `json:"role"`
	Content json.RawMessage `json:"content"`
	CallID  string          `json:"call_id"`
	Name    string          `json:"name"`
	// Arguments for a function call; Input for a custom tool's call.
	Arguments string          `json:"arguments"`
	Input     string          `json:"input"`
	Output    json.RawMessage `json:"output"`
}

type chatMsg struct {
	Role       string         `json:"role"`
	Content    any            `json:"content"`
	ToolCalls  []chatToolCall `json:"tool_calls,omitempty"`
	ToolCallID string         `json:"tool_call_id,omitempty"`
}

type chatToolCall struct {
	ID       string `json:"id"`
	Type     string `json:"type"`
	Function struct {
		Name      string `json:"name"`
		Arguments string `json:"arguments"`
	} `json:"function"`
}

// customInput is the one argument a custom tool is given as a function: its
// free-form text.
var customInput = json.RawMessage(`{"type":"object","properties":{"input":{"type":"string","description":"The tool's input, as free text."}},"required":["input"]}`)

// partsText joins a content list's text parts; a plain string is itself.
func partsText(raw json.RawMessage) string {
	if len(raw) == 0 {
		return ""
	}
	var s string
	if json.Unmarshal(raw, &s) == nil {
		return s
	}
	var parts []struct {
		Type string `json:"type"`
		Text string `json:"text"`
	}
	if json.Unmarshal(raw, &parts) != nil {
		return ""
	}
	var b strings.Builder
	for _, p := range parts {
		switch p.Type {
		case "input_text", "output_text", "text", "summary_text":
			b.WriteString(p.Text)
		}
	}
	return b.String()
}

// toChat rewrites a Responses request as a chat completion. custom names the
// tools that were custom, so their calls go back as custom calls.
func toChat(in respRequest) (map[string]any, map[string]bool, error) {
	var msgs []chatMsg
	if in.Instructions != "" {
		msgs = append(msgs, chatMsg{Role: "system", Content: in.Instructions})
	}
	var items []respItem
	var s string
	switch {
	case len(in.Input) == 0:
	case json.Unmarshal(in.Input, &s) == nil:
		items = []respItem{{Type: "message", Role: "user", Content: in.Input}}
	default:
		if err := json.Unmarshal(in.Input, &items); err != nil {
			return nil, nil, fmt.Errorf("input is neither text nor a list of items: %v", err)
		}
	}
	for _, it := range items {
		switch it.Type {
		case "", "message":
			role := it.Role
			if role == "developer" {
				role = "system"
			}
			if role == "" {
				role = "user"
			}
			msgs = append(msgs, chatMsg{Role: role, Content: partsText(it.Content)})
		case "function_call", "custom_tool_call":
			args := it.Arguments
			if it.Type == "custom_tool_call" {
				b, _ := json.Marshal(map[string]string{"input": it.Input})
				args = string(b)
			}
			tc := chatToolCall{ID: it.CallID, Type: "function"}
			tc.Function.Name, tc.Function.Arguments = it.Name, args
			// Calls made together are one assistant turn.
			if n := len(msgs); n > 0 && msgs[n-1].Role == "assistant" && len(msgs[n-1].ToolCalls) > 0 {
				msgs[n-1].ToolCalls = append(msgs[n-1].ToolCalls, tc)
			} else {
				msgs = append(msgs, chatMsg{Role: "assistant", Content: "", ToolCalls: []chatToolCall{tc}})
			}
		case "function_call_output", "custom_tool_call_output":
			out := partsText(it.Output)
			if out == "" && len(it.Output) > 0 {
				out = string(it.Output)
			}
			msgs = append(msgs, chatMsg{Role: "tool", ToolCallID: it.CallID, Content: out})
		}
		// reasoning items, and anything else, carry nothing an engine reads.
	}
	if len(msgs) == 0 {
		return nil, nil, fmt.Errorf("the request has no input")
	}

	out := map[string]any{"model": in.Model, "messages": msgs, "stream": in.Stream}
	if in.Stream {
		out["stream_options"] = map[string]any{"include_usage": true}
	}
	custom := map[string]bool{}
	var tools []map[string]any
	for _, t := range in.Tools {
		params := t.Parameters
		switch t.Type {
		case "function":
		case "custom":
			params = customInput
			custom[t.Name] = true
		default:
			continue // server-side tools: nothing here can run them
		}
		if len(params) == 0 {
			params = json.RawMessage(`{"type":"object","properties":{}}`)
		}
		tools = append(tools, map[string]any{"type": "function", "function": map[string]any{
			"name": t.Name, "description": t.Description, "parameters": params}})
	}
	if len(tools) > 0 {
		out["tools"] = tools
		if len(in.ToolChoice) > 0 {
			var choice string
			if json.Unmarshal(in.ToolChoice, &choice) == nil {
				out["tool_choice"] = choice
			}
		}
		if in.ParallelToolCalls != nil {
			out["parallel_tool_calls"] = *in.ParallelToolCalls
		}
	}
	if in.Temperature != nil {
		out["temperature"] = *in.Temperature
	}
	if in.TopP != nil {
		out["top_p"] = *in.TopP
	}
	if in.MaxOutputTokens != nil {
		out["max_tokens"] = *in.MaxOutputTokens
	}
	return out, custom, nil
}

var respSeq atomic.Int64

func respID(prefix string) string {
	return fmt.Sprintf("%s_%x%04x", prefix, time.Now().UnixNano(), respSeq.Add(1)&0xffff)
}

// handleResponses serves POST /v1/responses.
func (r *Router) handleResponses(w http.ResponseWriter, req *http.Request) {
	body, err := io.ReadAll(io.LimitReader(req.Body, r.cfg.MaxBodyBytes))
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", "could not read the request")
		return
	}
	var in respRequest
	if err := json.Unmarshal(body, &in); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", "the body is not JSON: "+err.Error())
		return
	}
	chat, custom, err := toChat(in)
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
		return
	}
	upstream, _ := json.Marshal(chat)
	next := req.Clone(req.Context())
	next.URL.Path = "/v1/chat/completions"
	next.Body = io.NopCloser(bytes.NewReader(upstream))
	next.ContentLength = int64(len(upstream))
	next.Header.Set("Content-Type", "application/json")

	t := &respTranslator{w: w, model: in.Model, stream: in.Stream, custom: custom,
		id: respID("resp"), created: time.Now().Unix(), metadata: in.Metadata}
	r.handleInference("/v1/chat/completions")(t, next)
	t.finish()
}

// respTranslator stands between the chat handler and the client, rewriting
// what the handler writes.
type respTranslator struct {
	w        http.ResponseWriter
	model    string
	stream   bool
	custom   map[string]bool
	id       string
	created  int64
	metadata map[string]string

	status  int
	failed  bool
	buf     bytes.Buffer // a non-streaming body, or a partial line
	started bool
	seq     int

	// What the stream has built.
	msgID    string
	msgIndex int
	text     strings.Builder
	calls    []*respCall
	byIndex  map[int]*respCall
	usage    map[string]int
	finished bool
}

type respCall struct {
	outIndex         int
	id, callID, name string
	args             strings.Builder
}

func (t *respTranslator) Header() http.Header { return t.w.Header() }

func (t *respTranslator) WriteHeader(code int) {
	t.status = code
	if code >= 300 {
		// An error from the chat path is passed on as it is: its body is
		// already the error object a client reads.
		t.failed = true
		t.w.WriteHeader(code)
		return
	}
	if t.stream {
		t.w.Header().Set("Content-Type", "text/event-stream")
	} else {
		t.w.Header().Set("Content-Type", "application/json")
	}
	t.w.Header().Del("Content-Length")
	t.w.WriteHeader(code)
}

func (t *respTranslator) Flush() {
	if f, ok := t.w.(http.Flusher); ok {
		f.Flush()
	}
}

func (t *respTranslator) Write(p []byte) (int, error) {
	if t.status == 0 {
		t.WriteHeader(http.StatusOK)
	}
	if t.failed {
		return t.w.Write(p)
	}
	t.buf.Write(p)
	if !t.stream {
		return len(p), nil
	}
	for {
		line, err := t.buf.ReadString('\n')
		if err != nil {
			// Not a whole line yet: keep it for the next write.
			t.buf.Reset()
			t.buf.WriteString(line)
			break
		}
		t.chatLine(strings.TrimSpace(line))
	}
	return len(p), nil
}

// event writes one server-sent event of the Responses stream.
func (t *respTranslator) event(typ string, fields map[string]any) {
	fields["type"] = typ
	fields["sequence_number"] = t.seq
	t.seq++
	b, _ := json.Marshal(fields)
	fmt.Fprintf(t.w, "event: %s\ndata: %s\n\n", typ, b)
}

// object is the response as it stands.
func (t *respTranslator) object(status string) map[string]any {
	o := map[string]any{"id": t.id, "object": "response", "created_at": t.created, "model": t.model,
		"status": status, "output": t.output(status == "completed")}
	if t.metadata != nil {
		o["metadata"] = t.metadata
	}
	if t.usage != nil {
		o["usage"] = map[string]any{"input_tokens": t.usage["prompt"], "output_tokens": t.usage["completion"],
			"total_tokens":         t.usage["prompt"] + t.usage["completion"],
			"input_tokens_details": map[string]int{"cached_tokens": 0}, "output_tokens_details": map[string]int{"reasoning_tokens": 0}}
	}
	return o
}

func (t *respTranslator) messageItem(done bool) map[string]any {
	status := "in_progress"
	if done {
		status = "completed"
	}
	return map[string]any{"type": "message", "id": t.msgID, "role": "assistant", "status": status,
		"content": []map[string]any{{"type": "output_text", "text": t.text.String(), "annotations": []any{}}}}
}

func (t *respTranslator) callItem(c *respCall, done bool) map[string]any {
	status := "in_progress"
	if done {
		status = "completed"
	}
	if t.custom[c.name] {
		var a struct {
			Input string `json:"input"`
		}
		_ = json.Unmarshal([]byte(c.args.String()), &a)
		return map[string]any{"type": "custom_tool_call", "id": c.id, "call_id": c.callID, "name": c.name,
			"input": a.Input, "status": status}
	}
	return map[string]any{"type": "function_call", "id": c.id, "call_id": c.callID, "name": c.name,
		"arguments": c.args.String(), "status": status}
}

func (t *respTranslator) output(done bool) []map[string]any {
	type item struct {
		i int
		v map[string]any
	}
	var items []item
	if t.msgID != "" {
		items = append(items, item{t.msgIndex, t.messageItem(done)})
	}
	for _, c := range t.calls {
		items = append(items, item{c.outIndex, t.callItem(c, done)})
	}
	out := make([]map[string]any, len(items))
	for _, it := range items {
		if it.i < len(out) {
			out[it.i] = it.v
		}
	}
	return out
}

func (t *respTranslator) begin() {
	if t.started {
		return
	}
	t.started = true
	t.byIndex = map[int]*respCall{}
	t.event("response.created", map[string]any{"response": t.object("in_progress")})
	t.event("response.in_progress", map[string]any{"response": t.object("in_progress")})
}

// chatLine takes one line of the chat stream.
func (t *respTranslator) chatLine(line string) {
	if !strings.HasPrefix(line, "data:") {
		return
	}
	payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
	if payload == "[DONE]" {
		return
	}
	var chunk struct {
		Choices []struct {
			Delta struct {
				Content   string `json:"content"`
				ToolCalls []struct {
					Index    int    `json:"index"`
					ID       string `json:"id"`
					Function struct {
						Name      string `json:"name"`
						Arguments string `json:"arguments"`
					} `json:"function"`
				} `json:"tool_calls"`
			} `json:"delta"`
		} `json:"choices"`
		Usage *struct {
			Prompt     int `json:"prompt_tokens"`
			Completion int `json:"completion_tokens"`
		} `json:"usage"`
	}
	if json.Unmarshal([]byte(payload), &chunk) != nil {
		return
	}
	t.begin()
	if chunk.Usage != nil {
		t.usage = map[string]int{"prompt": chunk.Usage.Prompt, "completion": chunk.Usage.Completion}
	}
	for _, ch := range chunk.Choices {
		if d := ch.Delta.Content; d != "" {
			if t.msgID == "" {
				// The message's place is after whatever calls came first.
				t.msgIndex = len(t.calls)
				t.msgID = respID("msg")
				t.event("response.output_item.added", map[string]any{"output_index": t.msgIndex, "item": map[string]any{
					"type": "message", "id": t.msgID, "role": "assistant", "status": "in_progress", "content": []any{}}})
				t.event("response.content_part.added", map[string]any{"item_id": t.msgID, "output_index": t.msgIndex,
					"content_index": 0, "part": map[string]any{"type": "output_text", "text": "", "annotations": []any{}}})
			}
			t.text.WriteString(d)
			t.event("response.output_text.delta", map[string]any{"item_id": t.msgID, "output_index": t.msgIndex,
				"content_index": 0, "delta": d})
		}
		for _, tc := range ch.Delta.ToolCalls {
			c := t.byIndex[tc.Index]
			if c == nil {
				c = &respCall{outIndex: t.nextIndex(), id: respID("fc"), callID: tc.ID, name: tc.Function.Name}
				if c.callID == "" {
					c.callID = respID("call")
				}
				t.byIndex[tc.Index] = c
				t.calls = append(t.calls, c)
				added := t.callItem(c, false)
				t.event("response.output_item.added", map[string]any{"output_index": c.outIndex, "item": added})
			}
			if tc.Function.Name != "" && c.name == "" {
				c.name = tc.Function.Name
			}
			if a := tc.Function.Arguments; a != "" {
				c.args.WriteString(a)
				if !t.custom[c.name] {
					t.event("response.function_call_arguments.delta", map[string]any{"item_id": c.id,
						"output_index": c.outIndex, "delta": a})
				}
			}
		}
	}
}

// nextIndex is the output position the next item takes.
func (t *respTranslator) nextIndex() int {
	n := len(t.calls)
	if t.msgID != "" {
		n++
	}
	return n
}

// finish closes the stream, or rewrites the whole body.
func (t *respTranslator) finish() {
	if t.failed || t.finished {
		return
	}
	t.finished = true
	if !t.stream {
		t.fromChatJSON(t.buf.Bytes())
		return
	}
	if rest := strings.TrimSpace(t.buf.String()); rest != "" {
		t.chatLine(rest)
	}
	t.begin()
	if t.msgID != "" {
		t.event("response.output_text.done", map[string]any{"item_id": t.msgID, "output_index": t.msgIndex,
			"content_index": 0, "text": t.text.String()})
		t.event("response.content_part.done", map[string]any{"item_id": t.msgID, "output_index": t.msgIndex,
			"content_index": 0, "part": map[string]any{"type": "output_text", "text": t.text.String(), "annotations": []any{}}})
		t.event("response.output_item.done", map[string]any{"output_index": t.msgIndex, "item": t.messageItem(true)})
	}
	for _, c := range t.calls {
		if !t.custom[c.name] {
			t.event("response.function_call_arguments.done", map[string]any{"item_id": c.id,
				"output_index": c.outIndex, "arguments": c.args.String()})
		}
		t.event("response.output_item.done", map[string]any{"output_index": c.outIndex, "item": t.callItem(c, true)})
	}
	t.event("response.completed", map[string]any{"response": t.object("completed")})
	t.Flush()
}

// fromChatJSON rewrites a whole chat completion as a response object.
func (t *respTranslator) fromChatJSON(b []byte) {
	var c struct {
		Choices []struct {
			Message struct {
				Content   string `json:"content"`
				ToolCalls []struct {
					ID       string `json:"id"`
					Function struct {
						Name      string `json:"name"`
						Arguments string `json:"arguments"`
					} `json:"function"`
				} `json:"tool_calls"`
			} `json:"message"`
		} `json:"choices"`
		Usage *struct {
			Prompt     int `json:"prompt_tokens"`
			Completion int `json:"completion_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(b, &c); err != nil {
		_, _ = t.w.Write(b)
		return
	}
	t.byIndex = map[int]*respCall{}
	if c.Usage != nil {
		t.usage = map[string]int{"prompt": c.Usage.Prompt, "completion": c.Usage.Completion}
	}
	for _, ch := range c.Choices {
		if ch.Message.Content != "" {
			t.msgID, t.msgIndex = respID("msg"), 0
			t.text.WriteString(ch.Message.Content)
		}
		for _, tc := range ch.Message.ToolCalls {
			rc := &respCall{outIndex: t.nextIndex(), id: respID("fc"), callID: tc.ID, name: tc.Function.Name}
			rc.args.WriteString(tc.Function.Arguments)
			t.calls = append(t.calls, rc)
		}
	}
	out, _ := json.Marshal(t.object("completed"))
	_, _ = t.w.Write(out)
}