OpenAI completions proxy for Laguna tool parsing

main.go

// laguna-proxy: a stdlib-only reverse proxy that repairs llama.cpp
// (--skip-chat-parsing) output for poolside/Laguna-S-2.1.
//
// Handles both non-stream (application/json) and stream (text/event-stream).
//
// Laguna raw assistant output shape:
//   <think>...</think>
//   <tool_call>NAME<arg_key>K</arg_key><arg_value>V</arg_value>...</tool_call>
// String arg values are raw (unquoted); non-string values are JSON.
//
// Reasoning strategy: LIVE STREAMING. Because the Laguna template injects the
// opening <think> into the prompt (it is never streamed back), we default to
// the "think" region and emit reasoning tokens as they arrive. regThink also
// watches for <tool_call> so a tool call that follows reasoning WITHOUT a
// closing </think> still ends reasoning cleanly and is parsed (rather than
// being swallowed forever in a holding buffer).
//
// Usage:
//   go run main.go -host localhost -port 8080 -upstream http://127.0.0.1:8081 -model-id laguna
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"io"
	"log"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

// ---------------------------------------------------------------------------
// flags / main
// ---------------------------------------------------------------------------

func main() {
	host := flag.String("host", "localhost", "host/interface to listen on")
	port := flag.Int("port", 8080, "port to listen on")
	upstream := flag.String("upstream", "http://127.0.0.1:8081", "llama-server base URL")
	modelID := flag.String("model-id", "laguna",
		"case-insensitive substring matched against the response `model` field; "+
			"non-matching responses are passed through untouched")
	verbose := flag.Bool("verbose", false, "log region transitions")
	// Laguna injects the opening <think> in the prompt, so the stream starts
	// INSIDE reasoning by default. Pass -content-first for a model/template that
	// streams plain content first (no implicit open reasoning block).
	contentFirst := flag.Bool("content-first", false,
		"assume the stream starts in plain content (default assumes it starts inside <think> reasoning)")
	flag.Parse()

	up, err := url.Parse(*upstream)
	if err != nil {
		log.Fatalf("bad -upstream: %v", err)
	}

	p := &proxy{
		upstream:     up,
		client:       &http.Client{Timeout: 0}, // streaming: no global timeout
		modelID:      strings.ToLower(*modelID),
		verbose:      *verbose,
		contentFirst: *contentFirst,
	}

	addr := *host + ":" + strconv.Itoa(*port)
	srv := &http.Server{Addr: addr, Handler: p}
	log.Printf("laguna-proxy listening on %s -> %s (model-id match=%q, content-first=%v)",
		addr, up, *modelID, *contentFirst)
	log.Fatal(srv.ListenAndServe())
}

// ---------------------------------------------------------------------------
// proxy
// ---------------------------------------------------------------------------

type proxy struct {
	upstream     *url.URL
	client       *http.Client
	modelID      string // lowercased
	verbose      bool
	contentFirst bool
}

func (p *proxy) matchModel(model string) bool {
	if p.modelID == "" {
		return true
	}
	return strings.Contains(strings.ToLower(model), p.modelID)
}

func (p *proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	isChat := strings.HasSuffix(strings.TrimRight(r.URL.Path, "/"), "chat/completions")

	outURL := *p.upstream
	outURL.Path = singleJoin(p.upstream.Path, r.URL.Path)
	outURL.RawQuery = r.URL.RawQuery

	var bodyBytes []byte
	if r.Body != nil {
		bodyBytes, _ = io.ReadAll(r.Body)
		_ = r.Body.Close()
	}

	outReq, err := http.NewRequestWithContext(r.Context(), r.Method, outURL.String(), bytes.NewReader(bodyBytes))
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	copyHeaders(outReq.Header, r.Header)
	outReq.Header.Del("Accept-Encoding") // identity so we can parse text

	resp, err := p.client.Do(outReq)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

	ct := resp.Header.Get("Content-Type")

	if !isChat {
		passthrough(w, resp)
		return
	}

	switch {
	case strings.Contains(ct, "text/event-stream"):
		p.handleStream(w, resp)
	case strings.Contains(ct, "application/json"):
		p.handleJSON(w, resp)
	default:
		passthrough(w, resp)
	}
}

// ---------------------------------------------------------------------------
// non-streaming path
// ---------------------------------------------------------------------------

func (p *proxy) handleJSON(w http.ResponseWriter, resp *http.Response) {
	raw, err := io.ReadAll(resp.Body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}

	var body map[string]any
	if json.Unmarshal(raw, &body) != nil {
		copyHeaders(w.Header(), resp.Header)
		w.WriteHeader(resp.StatusCode)
		_, _ = w.Write(raw)
		return
	}

	// model-id gate
	model, _ := body["model"].(string)
	if !p.matchModel(model) {
		copyHeaders(w.Header(), resp.Header)
		w.WriteHeader(resp.StatusCode)
		_, _ = w.Write(raw)
		return
	}

	if choices, ok := body["choices"].([]any); ok {
		for _, c := range choices {
			choice, ok := c.(map[string]any)
			if !ok {
				continue
			}
			msg, ok := choice["message"].(map[string]any)
			if !ok {
				continue
			}
			content, _ := msg["content"].(string)

			reasoning, toolCalls, clean := parseComplete(content, !p.contentFirst)

			if clean == "" {
				msg["content"] = nil
			} else {
				msg["content"] = clean
			}
			if reasoning != "" {
				msg["reasoning_content"] = reasoning
			}
			if len(toolCalls) > 0 {
				msg["tool_calls"] = toolCalls
				choice["finish_reason"] = "tool_calls"
			}
		}
	}

	out, _ := json.Marshal(body)
	h := w.Header()
	copyHeaders(h, resp.Header)
	h.Del("Content-Length")
	h.Set("Content-Length", strconv.Itoa(len(out)))
	w.WriteHeader(resp.StatusCode)
	_, _ = w.Write(out)
}

// parseComplete parses a full (non-streamed) assistant string.
// reasoningFirst == true means the opening <think> was in the prompt and the
// content begins inside reasoning (so a leading </think> closes it).
func parseComplete(text string, reasoningFirst bool) (reasoning string, toolCalls []map[string]any, clean string) {
	var rb strings.Builder

	if reasoningFirst {
		if end := strings.Index(text, "</think>"); end >= 0 {
			rb.WriteString(strings.TrimSpace(text[:end]))
			text = text[end+len("</think>"):]
		}
	}

	text = extractAll(text, "<think>", "</think>", func(inner string) {
		if rb.Len() > 0 {
			rb.WriteByte('\n')
		}
		rb.WriteString(strings.TrimSpace(inner))
	})
	reasoning = rb.String()

	idx := 0
	text = extractAll(text, "<tool_call>", "</tool_call>", func(inner string) {
		toolCalls = append(toolCalls, parseToolCallBlock(inner, idx))
		idx++
	})

	clean = strings.TrimSpace(text)
	return
}

func extractAll(s, open, close string, fn func(inner string)) string {
	var out strings.Builder
	for {
		i := strings.Index(s, open)
		if i < 0 {
			out.WriteString(s)
			break
		}
		out.WriteString(s[:i])
		rest := s[i+len(open):]
		j := strings.Index(rest, close)
		if j < 0 {
			fn(rest)
			break
		}
		fn(rest[:j])
		s = rest[j+len(close):]
	}
	return out.String()
}

// parseToolCallBlock parses one <tool_call> inner block into an OpenAI
// tool-call object. It uses a single left-to-right advancing cursor: the cursor
// starts AFTER the name so the first key/value pair is never dropped, and every
// search is anchored strictly after the previous match, so stray '<' or '>'
// inside an arg value can never be mistaken for structural tags (this is the
// fix for the "path landed under key `end`" misalignment bug).
func parseToolCallBlock(block string, idx int) map[string]any {
	// name = text before first <arg_key>
	name := block
	if k := strings.Index(block, "<arg_key>"); k >= 0 {
		name = block[:k]
	}
	name = strings.TrimSpace(name)

	args := map[string]any{}
	pos := len(name)
	for {
		ks := indexAt(block, "<arg_key>", pos)
		if ks < 0 {
			break
		}
		ke := indexAt(block, "</arg_key>", ks+len("<arg_key>"))
		if ke < 0 {
			break
		}
		key := strings.TrimSpace(block[ks+len("<arg_key>") : ke])

		vs := indexAt(block, "<arg_value>", ke+len("</arg_key>"))
		if vs < 0 {
			break
		}
		ve := indexAt(block, "</arg_value>", vs+len("<arg_value>"))
		if ve < 0 {
			break
		}
		val := block[vs+len("<arg_value>") : ve]

		if key != "" { // skip empty keys so a stray fragment can't become a bogus key
			args[key] = coerceArg(val)
		}
		pos = ve + len("</arg_value>")
	}

	argsJSON, _ := json.Marshal(args)
	return map[string]any{
		"id":   "call_" + strconv.Itoa(idx),
		"type": "function",
		"function": map[string]any{
			"name":      name,
			"arguments": string(argsJSON),
		},
	}
}

// indexAt is strings.Index scoped to s[from:], returning an absolute index.
func indexAt(s, sub string, from int) int {
	if from < 0 {
		from = 0
	}
	if from > len(s) {
		return -1
	}
	i := strings.Index(s[from:], sub)
	if i < 0 {
		return -1
	}
	return from + i
}

func coerceArg(raw string) any {
	s := strings.TrimSpace(raw)
	if s == "" {
		return raw
	}
	switch s[0] {
	case '{', '[', '"':
		var v any
		if json.Unmarshal([]byte(s), &v) == nil {
			return v
		}
	case 't', 'f', 'n', '-', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9':
		var v any
		if json.Unmarshal([]byte(s), &v) == nil {
			return v
		}
	}
	return raw
}

// ---------------------------------------------------------------------------
// pending buffer — amortized O(n) replacement for strings.Builder+consume
// ---------------------------------------------------------------------------

// pending holds streamed-but-not-yet-classified bytes. A read offset (off)
// marks consumed data; advance() is O(1). We compact (one copy) only when the
// dead prefix grows past compactThreshold, giving amortized O(n) over a stream.
type pending struct {
	data []byte
	off  int
}

const compactThreshold = 1 << 16 // 64 KiB of dead prefix before we compact

func (p *pending) append(s string) {
	if p.off > compactThreshold {
		p.data = append(p.data[:0], p.data[p.off:]...)
		p.off = 0
	}
	p.data = append(p.data, s...)
}

func (p *pending) view() string { return bytesToString(p.data[p.off:]) }

func (p *pending) advance(n int) { p.off += n } // O(1)

func (p *pending) len() int { return len(p.data) - p.off }

// ---------------------------------------------------------------------------
// streaming path
// ---------------------------------------------------------------------------

type region int

const (
	regContent region = iota
	regThink
	regToolCall
)

func (r region) String() string {
	switch r {
	case regThink:
		return "think"
	case regToolCall:
		return "tool_call"
	default:
		return "content"
	}
}

type emitter struct {
	w         http.ResponseWriter
	flusher   http.Flusher
	model     string
	id        string
	created   int64
	toolIndex int
}

type streamParser struct {
	buf     pending
	region  region
	tcBuf   strings.Builder
	em      *emitter
	verbose bool
}

var allDelimiters = []string{
	"<think>", "</think>",
	"<tool_call>", "</tool_call>",
}

func (p *proxy) handleStream(w http.ResponseWriter, resp *http.Response) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		p.handleJSON(w, resp)
		return
	}

	h := w.Header()
	copyHeaders(h, resp.Header)
	h.Del("Content-Length")
	h.Set("Content-Type", "text/event-stream")
	h.Set("Cache-Control", "no-cache")
	h.Set("Connection", "keep-alive")
	w.WriteHeader(resp.StatusCode)
	flusher.Flush()

	em := &emitter{w: w, flusher: flusher, created: time.Now().Unix()}
	// Default to regThink: Laguna injects the opening <think> in the prompt, so
	// the streamed content begins inside reasoning. -content-first flips this.
	startRegion := regThink
	if p.contentFirst {
		startRegion = regContent
	}
	sp := &streamParser{region: startRegion, em: em, verbose: p.verbose}

	modelChecked := false
	passthroughMode := false

	scanner := bufio.NewScanner(resp.Body)
	scanner.Buffer(make([]byte, 0, 1<<20), 16<<20)
	scanner.Split(bufio.ScanLines)

	for scanner.Scan() {
		line := scanner.Text()
		if line == "" {
			continue
		}
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))

		if data == "[DONE]" {
			if !passthroughMode {
				sp.finish()
			}
			em.raw("data: [DONE]\n\n")
			em.flusher.Flush()
			return
		}

		var chunk map[string]any
		if json.Unmarshal([]byte(data), &chunk) != nil {
			continue
		}

		if em.id == "" {
			if v, ok := chunk["id"].(string); ok {
				em.id = v
			}
			if v, ok := chunk["model"].(string); ok {
				em.model = v
			}
		}

		if !modelChecked {
			if em.model != "" {
				modelChecked = true
				if !p.matchModel(em.model) {
					passthroughMode = true
					if p.verbose {
						log.Printf("model %q does not match; passthrough", em.model)
					}
				}
			}
		}

		if passthroughMode {
			em.raw("data: " + data + "\n\n")
			em.flusher.Flush()
			continue
		}

		choices, _ := chunk["choices"].([]any)
		if len(choices) == 0 {
			continue
		}
		choice, _ := choices[0].(map[string]any)
		delta, _ := choice["delta"].(map[string]any)
		finish := choice["finish_reason"]

		if delta != nil {
			if content, ok := delta["content"].(string); ok && content != "" {
				sp.feed(content)
			}
		}

		if finish != nil && finish != "" {
			sp.finish()
			em.emitFinish(finish)
			em.raw("data: [DONE]\n\n")
			em.flusher.Flush()
			return
		}
	}
	if !passthroughMode {
		sp.finish()
	}
	em.flusher.Flush()
}

func (sp *streamParser) feed(s string) {
	sp.buf.append(s)
	sp.process(false)
}

func (sp *streamParser) finish() {
	sp.process(true)
	// If a tool call was still open at end of stream, close it out.
	if sp.region == regToolCall {
		sp.flushToolCall()
	}
	// NOTE: no reasoning holding buffer to flush — reasoning is streamed live.
}

func (sp *streamParser) setRegion(r region) {
	if sp.verbose && r != sp.region {
		log.Printf("region %s -> %s", sp.region, r)
	}
	sp.region = r
}

func (sp *streamParser) process(final bool) {
	for {
		before := sp.buf.len()
		work := sp.buf.view()

		switch sp.region {

		case regContent:
			ti := strings.Index(work, "<think>")
			ci := strings.Index(work, "<tool_call>")
			tc := strings.Index(work, "</think>") // stray close (reasoning opened elsewhere)
			next, open := earliestOf(
				tag{ti, "<think>"},
				tag{ci, "<tool_call>"},
				tag{tc, "</think>"},
			)
			if next < 0 {
				emitLen := safeEmitLen(work, final)
				if emitLen > 0 {
					sp.em.emitContent(work[:emitLen])
					sp.buf.advance(emitLen)
				}
				return
			}
			switch open {
			case "</think>":
				// Text before a stray close was actually reasoning; stream it live.
				if next > 0 {
					sp.em.emitReasoning(work[:next])
				}
				sp.buf.advance(next + len("</think>"))
				// stay in content
			case "<think>":
				if next > 0 {
					sp.em.emitContent(work[:next])
				}
				sp.buf.advance(next + len("<think>"))
				sp.setRegion(regThink)
			default: // <tool_call>
				if next > 0 {
					sp.em.emitContent(work[:next])
				}
				sp.buf.advance(next + len("<tool_call>"))
				sp.setRegion(regToolCall)
				sp.tcBuf.Reset()
			}

		case regThink:
			// LIVE reasoning streaming. Watch for BOTH the normal </think> close
			// AND a <tool_call> that follows reasoning with no closing tag — the
			// latter implicitly ends reasoning and must switch to tool parsing
			// (otherwise the tool call, and everything after, gets swallowed).
			endThink := strings.Index(work, "</think>")
			openTool := strings.Index(work, "<tool_call>")
			next, which := earliestOf(
				tag{endThink, "</think>"},
				tag{openTool, "<tool_call>"},
			)
			if next < 0 {
				emitLen := safeEmitLen(work, final)
				if emitLen > 0 {
					sp.em.emitReasoning(work[:emitLen]) // stream live
					sp.buf.advance(emitLen)
				}
				return
			}
			if next > 0 {
				sp.em.emitReasoning(work[:next])
			}
			switch which {
			case "</think>":
				sp.buf.advance(next + len("</think>"))
				sp.setRegion(regContent)
			default: // <tool_call> without a closing </think>
				sp.buf.advance(next + len("<tool_call>"))
				sp.setRegion(regToolCall)
				sp.tcBuf.Reset()
			}

		case regToolCall:
			end := strings.Index(work, "</tool_call>")
			if end < 0 {
				emitLen := safeEmitLen(work, final)
				if emitLen > 0 {
					sp.tcBuf.WriteString(work[:emitLen])
					sp.buf.advance(emitLen)
				}
				return
			}
			sp.tcBuf.WriteString(work[:end])
			sp.buf.advance(end + len("</tool_call>"))
			sp.flushToolCall()
			// After a tool call, Laguna may emit more reasoning before the next
			// step, and that reasoning arrives with no opening <think>. Return to
			// regThink so it is classified as reasoning (not leaked as content).
			sp.setRegion(regThink)
		}

		// Forward-progress guard: exit if a full pass consumed nothing.
		if sp.buf.len() == before || sp.buf.len() == 0 {
			return
		}
	}
}

func (sp *streamParser) flushToolCall() {
	block := sp.tcBuf.String()
	sp.tcBuf.Reset()
	if strings.TrimSpace(block) == "" {
		return
	}
	tc := parseToolCallBlock(block, sp.em.toolIndex)
	fn, _ := tc["function"].(map[string]any)
	name, _ := fn["name"].(string)
	arguments, _ := fn["arguments"].(string)
	sp.em.emitToolCall(sp.em.toolIndex, tc["id"].(string), name, arguments)
	sp.em.toolIndex++
}

// ---------------------------------------------------------------------------
// delimiter helpers
// ---------------------------------------------------------------------------

type tag struct {
	idx  int
	name string
}

func earliestOf(tags ...tag) (int, string) {
	best, bestName := -1, ""
	for _, t := range tags {
		if t.idx < 0 {
			continue
		}
		if best < 0 || t.idx < best {
			best, bestName = t.idx, t.name
		}
	}
	return best, bestName
}

func safeEmitLen(work string, final bool) int {
	if final {
		return len(work)
	}
	keep := 0
	for _, d := range allDelimiters {
		if k := partialSuffix(work, d); k > keep {
			keep = k
		}
	}
	if keep >= len(work) {
		return 0
	}
	return len(work) - keep
}

func partialSuffix(work, delim string) int {
	max := len(delim) - 1
	if max > len(work) {
		max = len(work)
	}
	for n := max; n > 0; n-- {
		if strings.HasSuffix(work, delim[:n]) {
			return n
		}
	}
	return 0
}

// ---------------------------------------------------------------------------
// SSE emitter
// ---------------------------------------------------------------------------

func (e *emitter) baseChunk(delta map[string]any, finish any) map[string]any {
	return map[string]any{
		"id":      orDefault(e.id, "chatcmpl-proxy"),
		"object":  "chat.completion.chunk",
		"created": e.created,
		"model":   e.model,
		"choices": []any{
			map[string]any{
				"index":         0,
				"delta":         delta,
				"finish_reason": finish,
			},
		},
	}
}

func (e *emitter) send(delta map[string]any, finish any) {
	b, _ := json.Marshal(e.baseChunk(delta, finish))
	e.raw("data: " + string(b) + "\n\n")
	e.flusher.Flush()
}

func (e *emitter) emitContent(s string) {
	if s == "" {
		return
	}
	e.send(map[string]any{"content": s}, nil)
}

func (e *emitter) emitReasoning(s string) {
	if s == "" {
		return
	}
	e.send(map[string]any{"reasoning_content": s}, nil)
}

func (e *emitter) emitToolCall(index int, id, name, arguments string) {
	delta := map[string]any{
		"tool_calls": []any{
			map[string]any{
				"index": index,
				"id":    id,
				"type":  "function",
				"function": map[string]any{
					"name":      name,
					"arguments": arguments,
				},
			},
		},
	}
	e.send(delta, nil)
}

func (e *emitter) emitFinish(upstreamFinish any) {
	finish := upstreamFinish
	if e.toolIndex > 0 {
		finish = "tool_calls"
	}
	e.send(map[string]any{}, finish)
}

func (e *emitter) raw(s string) { _, _ = io.WriteString(e.w, s) }

// ---------------------------------------------------------------------------
// helpers
// ---------------------------------------------------------------------------

func passthrough(w http.ResponseWriter, resp *http.Response) {
	copyHeaders(w.Header(), resp.Header)
	w.WriteHeader(resp.StatusCode)
	if f, ok := w.(http.Flusher); ok {
		buf := make([]byte, 32*1024)
		for {
			n, err := resp.Body.Read(buf)
			if n > 0 {
				_, _ = w.Write(buf[:n])
				f.Flush()
			}
			if err != nil {
				return
			}
		}
	}
	_, _ = io.Copy(w, resp.Body)
}

func copyHeaders(dst, src http.Header) {
	for k, vv := range src {
		switch strings.ToLower(k) {
		case "connection", "keep-alive", "transfer-encoding", "upgrade":
			continue
		}
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

func singleJoin(a, b string) string {
	switch {
	case a == "" || a == "/":
		return b
	case strings.HasSuffix(a, "/") && strings.HasPrefix(b, "/"):
		return a + b[1:]
	case !strings.HasSuffix(a, "/") && !strings.HasPrefix(b, "/"):
		return a + "/" + b
	default:
		return a + b
	}
}

func orDefault(s, d string) string {
	if s == "" {
		return d
	}
	return s
}

// bytesToString returns a copy of the unconsumed window (not the whole stream),
// so it stays stdlib-only and safe without reintroducing O(n^2).
func bytesToString(b []byte) string {
	if len(b) == 0 {
		return ""
	}
	return string(b)
}