quorate/services/codex-bridge/chat.go

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"time"
)

type message struct {
	Role    string          `json:"role"`
	Content json.RawMessage `json:"content"`
}

type chatRequest struct {
	Model     string            `json:"model"`
	Messages  []message         `json:"messages"`
	Stream    bool              `json:"stream"`
	Tools     []json.RawMessage `json:"tools"`
	Reasoning string            `json:"reasoning_effort"`
}

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

func apiError(w http.ResponseWriter, status int, err error) {
	writeJSON(w, status, object{"error": object{"message": err.Error(), "type": "bridge_error"}})
}

func (b *bridge) models(w http.ResponseWriter, r *http.Request) {
	c, err := b.start(r.Context())
	if err != nil {
		apiError(w, http.StatusBadGateway, err)
		return
	}
	defer c.close()
	models := []object{}
	var cursor *string
	for {
		raw, err := c.call("model/list", object{"cursor": cursor})
		if err != nil {
			apiError(w, http.StatusBadGateway, err)
			return
		}
		var page struct {
			Data []struct {
				Model string `json:"model"`
				Name  string `json:"displayName"`
			} `json:"data"`
			Next *string `json:"nextCursor"`
		}
		if err := json.Unmarshal(raw, &page); err != nil {
			apiError(w, http.StatusBadGateway, err)
			return
		}
		for _, model := range page.Data {
			models = append(models, object{"id": model.Model, "name": model.Name, "object": "model", "created": 0, "owned_by": "openai"})
		}
		if page.Next == nil {
			break
		}
		cursor = page.Next
	}
	writeJSON(w, http.StatusOK, object{"object": "list", "data": models})
}

func contentText(raw json.RawMessage) (string, error) {
	var text string
	if json.Unmarshal(raw, &text) == nil && string(raw) != "null" {
		return text, nil
	}
	var parts []struct {
		Type string `json:"type"`
		Text string `json:"text"`
	}
	if err := json.Unmarshal(raw, &parts); err != nil || string(raw) == "null" {
		return "", errors.New("message content must be text")
	}
	var result strings.Builder
	for _, part := range parts {
		if part.Type != "text" {
			return "", errors.New("only text content is supported")
		}
		result.WriteString(part.Text)
	}
	return result.String(), nil
}

func inputs(req chatRequest) ([]object, string, error) {
	if len(req.Tools) != 0 {
		return nil, "", errors.New("tools are managed by Codex through MCP; omit client tools")
	}
	if req.Model == "" || len(req.Messages) == 0 || req.Messages[len(req.Messages)-1].Role != "user" {
		return nil, "", errors.New("model and messages ending in a user message are required")
	}
	items := []object{}
	for i, msg := range req.Messages {
		role := msg.Role
		if role == "system" {
			role = "developer"
		}
		if role != "user" && role != "assistant" && role != "developer" {
			return nil, "", fmt.Errorf("unsupported message role: %s", msg.Role)
		}
		text, err := contentText(msg.Content)
		if err != nil {
			return nil, "", err
		}
		if i == len(req.Messages)-1 {
			return items, text, nil
		}
		kind := "input_text"
		if role == "assistant" {
			kind = "output_text"
		}
		items = append(items, object{"type": "message", "role": role, "content": []object{{"type": kind, "text": text}}})
	}
	return nil, "", errors.New("missing user message")
}

func (b *bridge) chat(w http.ResponseWriter, r *http.Request) {
	var req chatRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		apiError(w, http.StatusBadRequest, err)
		return
	}
	items, prompt, err := inputs(req)
	if err != nil {
		apiError(w, http.StatusBadRequest, err)
		return
	}
	c, err := b.start(r.Context())
	if err != nil {
		apiError(w, http.StatusBadGateway, err)
		return
	}
	defer c.close()
	raw, err := c.call("thread/start", object{"model": req.Model, "ephemeral": true, "approvalPolicy": "never", "sandbox": "read-only"})
	var thread struct {
		Thread struct {
			ID string `json:"id"`
		} `json:"thread"`
	}
	if err == nil {
		err = json.Unmarshal(raw, &thread)
	}
	if err == nil && len(items) > 0 {
		_, err = c.call("thread/inject_items", object{"threadId": thread.Thread.ID, "items": items})
	}
	if err != nil {
		apiError(w, http.StatusBadGateway, err)
		return
	}
	params := object{"threadId": thread.Thread.ID, "input": []object{{"type": "text", "text": prompt}}}
	if req.Reasoning != "" {
		params["effort"] = req.Reasoning
	}
	if _, err = c.call("turn/start", params); err != nil {
		apiError(w, http.StatusBadGateway, err)
		return
	}
	created := time.Now().Unix()
	id := "chatcmpl-" + thread.Thread.ID
	var answer strings.Builder
	usage := object{}
	lastItem := ""
	if req.Stream {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
	}
	emit := func(delta object, finish any, counts object) error {
		value := object{"id": id, "object": "chat.completion.chunk", "created": created, "model": req.Model, "choices": []object{{"index": 0, "delta": delta, "finish_reason": finish}}}
		if counts != nil {
			value["usage"] = counts
		}
		data, err := json.Marshal(value)
		if err != nil {
			return err
		}
		if _, err := fmt.Fprintf(w, "data: %s\n\n", data); err != nil {
			return err
		}
		return http.NewResponseController(w).Flush()
	}
	for {
		e, readErr := c.next()
		if readErr != nil {
			err = readErr
			break
		}
		var p struct {
			ThreadID string `json:"threadId"`
			ItemID   string `json:"itemId"`
			Delta    string `json:"delta"`
			Usage    struct {
				Total struct {
					Input  int `json:"inputTokens"`
					Output int `json:"outputTokens"`
					Total  int `json:"totalTokens"`
				} `json:"total"`
			} `json:"tokenUsage"`
			Turn struct {
				Status string    `json:"status"`
				Error  *rpcError `json:"error"`
			} `json:"turn"`
		}
		if err = json.Unmarshal(e.Params, &p); err != nil {
			break
		}
		if p.ThreadID != thread.Thread.ID {
			continue
		}
		switch e.Method {
		case "item/agentMessage/delta":
			text := p.Delta
			if lastItem != "" && lastItem != p.ItemID {
				text = "\n" + text
			}
			lastItem = p.ItemID
			answer.WriteString(text)
			if req.Stream {
				err = emit(object{"role": "assistant", "content": text}, nil, nil)
			}
		case "thread/tokenUsage/updated":
			usage = object{"prompt_tokens": p.Usage.Total.Input, "completion_tokens": p.Usage.Total.Output, "total_tokens": p.Usage.Total.Total}
		case "turn/completed":
			if p.Turn.Status != "completed" {
				err = fmt.Errorf("codex turn %s", p.Turn.Status)
				if p.Turn.Error != nil {
					err = errors.New(p.Turn.Error.Message)
				}
				break
			}
			if req.Stream {
				if emit(object{}, "stop", usage) == nil {
					_, _ = fmt.Fprint(w, "data: [DONE]\n\n")
				}
			} else {
				writeJSON(w, http.StatusOK, object{"id": id, "object": "chat.completion", "created": created, "model": req.Model, "choices": []object{{"index": 0, "message": object{"role": "assistant", "content": answer.String()}, "finish_reason": "stop"}}, "usage": usage})
			}
			return
		}
		if err != nil {
			break
		}
	}
	if r.Context().Err() != nil {
		return
	}
	if req.Stream {
		data, _ := json.Marshal(object{"error": object{"message": err.Error(), "type": "bridge_error"}})
		_, _ = fmt.Fprintf(w, "data: %s\n\ndata: [DONE]\n\n", data)
	} else {
		apiError(w, http.StatusBadGateway, err)
	}
}