selfheal-lab-apps/ai/internal/external/responder.go

// Package external lets a client outside the AI VM act as the LLM for model ids starting with "external-": the
// diagnosis loop (internal/agent) is unchanged, but every turn is queued as a Turn and served to the client over
// GET /external/next (long poll); the client answers with POST /external/reply. Used for reference measurements with
// a model that has no API path (2026-10-16: Claude Mythos 5.1 driven by a Claude Code subagent). Token usage is not
// known for such turns and is reported as zero; reports mark those runs as reference records.
//
// A reply becomes a diagnosis and, with auto-approval, an executed action, so whoever can answer a turn controls the
// lab: the endpoints are served on their own listener (EXTERNAL_ADDR, off by default; the VM maps it to 127.0.0.1 so
// only the IAP SSH tunnel reaches it), turn ids are random, and every reply must carry the per-turn reply_token that
// only the client who took the turn received.
package external

import (
	"context"
	"crypto/rand"
	"crypto/subtle"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/selfheal-lab/selfheal-lab-apps/ai/internal/openai"
)

const maxConvCached = 64

// Turn is one request to the external client: the whole conversation so far (the loop sends the system prompt only
// on the first turn, so it is carried along), the tool definitions and whether the final submit_diagnosis is forced.
type Turn struct {
	TurnID          string        `json:"turn_id"`
	ReplyToken      string        `json:"reply_token"` // random, handed out with the turn, required by /external/reply
	Requeued        int           `json:"requeued"`    // times this turn was handed out again after an unanswered take
	Session         string        `json:"session"`     // first response id of the conversation ("" on the first turn)
	Model           string        `json:"model"`
	Turn            int           `json:"turn"` // 1-based index within the conversation
	System          string        `json:"system"`
	Input           []openai.Item `json:"input"` // user messages, assistant function_call items, function_call_output items
	Tools           []openai.Tool `json:"tools"`
	ForceSubmit     bool          `json:"force_submit"` // the loop asked for submit_diagnosis now (last turn)
	MaxOutputTokens int           `json:"max_output_tokens"`
	IssuedAt        time.Time     `json:"issued_at"`
	DeadlineAt      time.Time     `json:"deadline_at"`
}

// Reply is the client's answer: function calls (name + arguments) and/or assistant text.
type Reply struct {
	TurnID     string      `json:"turn_id"`
	ReplyToken string      `json:"reply_token"`
	Output     []replyItem `json:"output"`
	Error      string      `json:"error,omitempty"` // client gave up on this turn
}

type replyItem struct {
	Type      string          `json:"type"` // function_call | message
	Name      string          `json:"name,omitempty"`
	Arguments json.RawMessage `json:"arguments,omitempty"` // JSON object or a JSON string holding one
	Text      string          `json:"text,omitempty"`
}

type pending struct {
	turn    Turn
	reply   chan Reply
	takenAt time.Time // zero while queued
}

type conversation struct {
	system string
	items  []openai.Item
	turn   int
	first  string
}

// Responder implements agent.Responder; one instance serves every external-* diagnosis.
type Responder struct {
	TurnTimeout  time.Duration // how long one turn waits for the client (bounded by the caller's context as well)
	RequeueAfter time.Duration // a taken turn with no reply after this long goes back to the queue with a new token

	mu       sync.Mutex
	queue    []*pending          // issued, not yet taken by a client
	taken    map[string]*pending // taken, waiting for the reply
	conv     map[string]conversation
	order    []string
	seq      int
	notify   chan struct{}
	lastPoll time.Time
	lastTurn time.Time
}

func New(turnTimeout time.Duration) *Responder {
	if turnTimeout <= 0 {
		turnTimeout = 20 * time.Minute
	}
	return &Responder{TurnTimeout: turnTimeout, RequeueAfter: 3 * time.Minute, taken: map[string]*pending{}, conv: map[string]conversation{}, notify: make(chan struct{}, 1)}
}

// CreateResponse queues the turn and blocks until the client replies, the context ends or TurnTimeout passes.
func (r *Responder) CreateResponse(ctx context.Context, req openai.Request) (*openai.Response, error) {
	var hist conversation
	if req.PreviousResponseID != "" {
		r.mu.Lock()
		prev, ok := r.conv[req.PreviousResponseID]
		r.mu.Unlock()
		if !ok {
			return nil, fmt.Errorf("external: unknown previous_response_id %q", req.PreviousResponseID)
		}
		hist = conversation{system: prev.system, items: append([]openai.Item{}, prev.items...), turn: prev.turn, first: prev.first}
	}
	for _, it := range req.Input {
		if it.Type == "message" && it.Role == "system" {
			if s, _ := it.Content.(string); s != "" {
				if hist.system != "" {
					hist.system += "\n\n"
				}
				hist.system += s
			}
			continue
		}
		hist.items = append(hist.items, it)
	}
	if len(hist.items) == 0 {
		return nil, fmt.Errorf("external: no input items")
	}
	forced := false
	if m, ok := req.ToolChoice.(map[string]any); ok && m["type"] == "function" {
		forced = true
	}
	id, tok := "ext_"+randHex(8), randHex(16)
	r.mu.Lock()
	r.seq++
	now := time.Now()
	p := &pending{turn: Turn{TurnID: id, ReplyToken: tok, Session: hist.first, Model: req.Model, Turn: hist.turn + 1, System: hist.system, Input: hist.items, Tools: req.Tools,
		ForceSubmit: forced, MaxOutputTokens: req.MaxOutputTokens, IssuedAt: now, DeadlineAt: now.Add(r.TurnTimeout)}, reply: make(chan Reply, 1)}
	r.queue = append(r.queue, p)
	r.lastTurn = now
	r.mu.Unlock()
	select {
	case r.notify <- struct{}{}:
	default:
	}

	var rep Reply
	select {
	case rep = <-p.reply:
	case <-ctx.Done():
		r.drop(id)
		return nil, fmt.Errorf("external: turn %s not answered before the diagnosis deadline (%v)", id, ctx.Err())
	case <-time.After(r.TurnTimeout):
		r.drop(id)
		return nil, fmt.Errorf("external: turn %s not answered within %s", id, r.TurnTimeout)
	}
	if rep.Error != "" {
		return nil, fmt.Errorf("external: client error on turn %s: %s", id, rep.Error)
	}
	out := &openai.Response{ID: id, Status: "completed", Model: req.Model}
	var text strings.Builder
	for i, it := range rep.Output {
		switch it.Type {
		case "function_call":
			out.Output = append(out.Output, openai.Item{Type: "function_call", CallID: fmt.Sprintf("%s_call_%d", id, i+1), ID: fmt.Sprintf("%s_call_%d", id, i+1), Name: it.Name, Arguments: argString(it.Arguments), Status: "completed"})
		case "message":
			text.WriteString(it.Text)
		}
	}
	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...)
	}
	// remember the conversation (assistant output items are replayed to the client on the next turn)
	first := hist.first
	if first == "" {
		first = id
	}
	items := append(append([]openai.Item{}, hist.items...), out.Output...)
	r.mu.Lock()
	r.conv[id] = conversation{system: hist.system, items: items, turn: hist.turn + 1, first: first}
	r.order = append(r.order, id)
	for len(r.order) > maxConvCached {
		delete(r.conv, r.order[0])
		r.order = r.order[1:]
	}
	r.mu.Unlock()
	return out, nil
}

// argString returns the function arguments as the JSON object text (a JSON string value is unwrapped).
func argString(raw json.RawMessage) string {
	s := strings.TrimSpace(string(raw))
	if s == "" {
		return "{}"
	}
	if strings.HasPrefix(s, `"`) {
		var inner string
		if json.Unmarshal(raw, &inner) == nil {
			if strings.TrimSpace(inner) == "" {
				return "{}"
			}
			return inner
		}
	}
	return s
}

// requeueStale (mu held) returns taken turns whose reply never came (client lost the output, crashed, was replaced)
// to the front of the queue with a fresh reply token, so the next poll serves them again; the old token is void.
func (r *Responder) requeueStale() {
	if r.RequeueAfter <= 0 {
		return
	}
	var stale []*pending
	for id, p := range r.taken {
		if time.Since(p.takenAt) >= r.RequeueAfter {
			delete(r.taken, id)
			p.turn.ReplyToken = randHex(16)
			p.turn.Requeued++
			p.takenAt = time.Time{}
			stale = append(stale, p)
		}
	}
	if len(stale) > 0 {
		r.queue = append(stale, r.queue...)
	}
}

func (r *Responder) drop(id string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	delete(r.taken, id)
	for i, p := range r.queue {
		if p.turn.TurnID == id {
			r.queue = append(r.queue[:i], r.queue[i+1:]...)
			return
		}
	}
}

// Next is GET /external/next?wait=<s>: the oldest queued turn (200), or 204 when none arrived within wait seconds.
func (r *Responder) Next(w http.ResponseWriter, req *http.Request) {
	wait := 20 * time.Second
	if v, err := strconv.Atoi(req.URL.Query().Get("wait")); err == nil && v >= 0 && v <= 60 {
		wait = time.Duration(v) * time.Second
	}
	deadline := time.NewTimer(wait)
	defer deadline.Stop()
	for {
		r.mu.Lock()
		r.lastPoll = time.Now()
		r.requeueStale()
		if len(r.queue) > 0 {
			p := r.queue[0]
			r.queue = r.queue[1:]
			p.takenAt = time.Now()
			r.taken[p.turn.TurnID] = p
			r.mu.Unlock()
			writeJSON(w, 200, p.turn)
			return
		}
		r.mu.Unlock()
		select {
		case <-r.notify:
		case <-deadline.C:
			w.WriteHeader(204)
			return
		case <-req.Context().Done():
			return
		}
	}
}

// ReplyHandler is POST /external/reply with a Reply body; 404 when the turn is unknown (already answered, dropped
// after its deadline, or never taken through /external/next).
func (r *Responder) ReplyHandler(w http.ResponseWriter, req *http.Request) {
	var rep Reply
	if err := json.NewDecoder(http.MaxBytesReader(w, req.Body, 1<<20)).Decode(&rep); err != nil {
		writeJSON(w, 400, map[string]string{"error": err.Error()})
		return
	}
	if rep.Error == "" {
		n := 0
		for _, it := range rep.Output {
			switch it.Type {
			case "function_call":
				if it.Name == "" {
					writeJSON(w, 400, map[string]string{"error": "function_call without name"})
					return
				}
				if a := argString(it.Arguments); !json.Valid([]byte(a)) || !strings.HasPrefix(strings.TrimSpace(a), "{") {
					writeJSON(w, 400, map[string]string{"error": "function_call arguments must be a JSON object"})
					return
				}
				n++
			case "message":
				n++
			default:
				writeJSON(w, 400, map[string]string{"error": "output item type must be function_call or message"})
				return
			}
		}
		if n == 0 {
			writeJSON(w, 400, map[string]string{"error": "output must hold at least one function_call or message"})
			return
		}
	}
	r.mu.Lock()
	p, ok := r.taken[rep.TurnID]
	if ok && subtle.ConstantTimeCompare([]byte(p.turn.ReplyToken), []byte(rep.ReplyToken)) != 1 {
		r.mu.Unlock()
		writeJSON(w, 403, map[string]string{"error": "reply_token does not match turn " + rep.TurnID})
		return
	}
	if ok {
		delete(r.taken, rep.TurnID)
	}
	r.mu.Unlock()
	if !ok {
		writeJSON(w, 404, map[string]string{"error": "unknown or expired turn_id " + rep.TurnID})
		return
	}
	p.reply <- rep
	writeJSON(w, 200, map[string]any{"ok": true, "turn_id": rep.TurnID})
}

// Status is shown by /selfcheck.
func (r *Responder) Status() map[string]any {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := map[string]any{"queued": len(r.queue), "awaiting_reply": len(r.taken), "turn_timeout_s": int(r.TurnTimeout.Seconds()), "requeue_after_s": int(r.RequeueAfter.Seconds())}
	if !r.lastPoll.IsZero() {
		out["last_poll_age_s"] = int(time.Since(r.lastPoll).Seconds())
	}
	if !r.lastTurn.IsZero() {
		out["last_turn_age_s"] = int(time.Since(r.lastTurn).Seconds())
	}
	return out
}

func randHex(n int) string {
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		panic("crypto/rand: " + err.Error())
	}
	return hex.EncodeToString(b)
}

// Handler serves the two endpoints; mount it on the dedicated external listener only.
func (r *Responder) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /external/next", r.Next)
	mux.HandleFunc("POST /external/reply", r.ReplyHandler)
	return mux
}

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