server.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"sync"
	"time"
)

// llamaServer wraps a long-running llama-server subprocess and the HTTP
// client that talks to it. One instance per (model, context-size). The
// model stays loaded in memory so subsequent /completion calls don't pay
// the GGUF mmap + warmup cost on every turn.
type llamaServer struct {
	cmd      *exec.Cmd
	port     int
	model    Model
	ctxN     int
	client   *http.Client
	waitOnce sync.Once
	waitErr  chan error
}

var (
	serverMu     sync.Mutex
	activeServer *llamaServer
)

// ensureServer returns a running llamaServer for the current model, spawning
// one if needed and restarting if the caller switched models. Safe to call
// concurrently.
func ensureServer() (*llamaServer, error) {
	serverMu.Lock()
	defer serverMu.Unlock()

	m, err := currentModel()
	if err != nil {
		return nil, err
	}
	if activeServer != nil && activeServer.model.Name == m.Name {
		return activeServer, nil
	}
	if activeServer != nil {
		activeServer.stopLocked()
		activeServer = nil
	}
	s, err := startLlamaServer(m)
	if err != nil {
		return nil, err
	}
	activeServer = s
	return s, nil
}

// shutdownServer is called from startChat's defer so the subprocess doesn't
// outlive the TUI session.
func shutdownServer() {
	serverMu.Lock()
	defer serverMu.Unlock()
	if activeServer != nil {
		activeServer.stopLocked()
		activeServer = nil
	}
}

func startLlamaServer(m Model) (*llamaServer, error) {
	bin, err := findEngineServer()
	if err != nil {
		return nil, fmt.Errorf("llama-server: %w", err)
	}
	modelPath, err := requireModel(m)
	if err != nil {
		return nil, err
	}
	port, err := pickFreePort()
	if err != nil {
		return nil, fmt.Errorf("pick port: %w", err)
	}

	threads := runtime.NumCPU() - 1
	if threads < 1 {
		threads = 1
	}
	if threads > 6 {
		threads = 6
	}

	args := []string{
		"-m", modelPath,
		"--host", "127.0.0.1",
		"--port", fmt.Sprintf("%d", port),
		"-c", "16384",
		"-t", fmt.Sprintf("%d", threads),
		"-ngl", "0",
		"--log-disable",
	}
	cmd := exec.Command(bin, args...)
	cmd.Stdin = bytes.NewReader(nil)
	cmd.Stdout = newLogWriter("llama-server stdout")
	cmd.Stderr = newLogWriter("llama-server stderr")
	cmd.Env = append(os.Environ(), "OMP_STACKSIZE=64M")
	applyEngineSysProcAttr(cmd)

	log.Printf("starting llama-server on :%d (model=%s threads=%d)", port, m.Name, threads)
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("start llama-server: %w", err)
	}

	s := &llamaServer{
		cmd:     cmd,
		port:    port,
		model:   m,
		ctxN:    16384,
		client:  &http.Client{Timeout: 10 * time.Minute},
		waitErr: make(chan error, 1),
	}
	// Single background Wait(); result is broadcast via waitErr so both
	// waitReady and stopLocked can observe exit without double-calling Wait.
	go func() { s.waitErr <- cmd.Wait() }()

	if err := s.waitReady(90 * time.Second); err != nil {
		s.stopLocked()
		return nil, err
	}
	log.Printf("llama-server ready on :%d", port)
	return s, nil
}

func pickFreePort() (int, error) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return 0, err
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port, nil
}

// waitReady polls GET /health until the server reports ready, or gives up.
// If the subprocess exits early, returns that error instead of timing out.
func (s *llamaServer) waitReady(maxWait time.Duration) error {
	deadline := time.Now().Add(maxWait)
	url := fmt.Sprintf("http://127.0.0.1:%d/health", s.port)

	for time.Now().Before(deadline) {
		select {
		case err := <-s.waitErr:
			// Process exited before becoming ready — put the error back for
			// stopLocked(), then report to caller.
			s.waitErr <- err
			return fmt.Errorf("llama-server exited before ready: %v", err)
		default:
		}
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		req, _ := http.NewRequestWithContext(ctx, "GET", url, nil)
		resp, err := s.client.Do(req)
		cancel()
		if err == nil {
			body := make([]byte, 256)
			n, _ := resp.Body.Read(body)
			resp.Body.Close()
			if resp.StatusCode == 200 && bytes.Contains(body[:n], []byte(`"ok"`)) {
				return nil
			}
		}
		time.Sleep(250 * time.Millisecond)
	}
	return fmt.Errorf("llama-server did not become ready in %s", maxWait)
}

func (s *llamaServer) stopLocked() {
	if s.cmd == nil || s.cmd.Process == nil {
		return
	}
	log.Printf("stopping llama-server pid=%d", s.cmd.Process.Pid)
	_ = s.cmd.Process.Kill()
	// Drain the background Wait goroutine so resources are released.
	select {
	case <-s.waitErr:
	case <-time.After(5 * time.Second):
	}
}

// ChatMsg is a single turn passed to /v1/chat/completions. Role is one of
// "system", "user", "assistant", or "tool". The tool-call fields only apply
// to assistant messages (ToolCalls) and tool messages (Name, ToolCallID).
type ChatMsg struct {
	Role       string     `json:"role"`
	Content    string     `json:"content"`
	ToolCalls  []ToolCall `json:"tool_calls,omitempty"`
	ToolCallID string     `json:"tool_call_id,omitempty"`
	Name       string     `json:"name,omitempty"`
}

// ToolCall is one requested function invocation emitted by the assistant.
// Arguments is a JSON-encoded string per the OpenAI schema, not a decoded
// object — the caller is responsible for unmarshalling.
type ToolCall struct {
	ID       string           `json:"id"`
	Type     string           `json:"type"`
	Function ToolCallFunction `json:"function"`
}

type ToolCallFunction struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
}

type chatRequest struct {
	Messages    []ChatMsg        `json:"messages"`
	MaxTokens   int              `json:"max_tokens"`
	Temperature float64          `json:"temperature"`
	Stream      bool             `json:"stream"`
	CachePrompt bool             `json:"cache_prompt"`
	Tools       []map[string]any `json:"tools,omitempty"`
}

type chatChoice struct {
	Message      assistantMessage `json:"message"`
	FinishReason string           `json:"finish_reason"`
}

// assistantMessage is the response-side shape of an assistant reply. Content
// can be null when the model chose to emit tool_calls instead.
type assistantMessage struct {
	Role      string     `json:"role"`
	Content   string     `json:"content"`
	ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}

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

type chatResponse struct {
	Choices []chatChoice `json:"choices"`
	Usage   chatUsage    `json:"usage"`
	Error   struct {
		Message string `json:"message"`
	} `json:"error"`
}

// UsageStats mirrors the last /v1/chat/completions usage block so the TUI
// can render a context-usage indicator. Total == prompt + completion.
type UsageStats struct {
	PromptTokens     int
	CompletionTokens int
	TotalTokens      int
	ContextSize      int
}

var (
	lastUsageMu sync.RWMutex
	lastUsage   UsageStats
)

func setLastUsage(u UsageStats) {
	lastUsageMu.Lock()
	lastUsage = u
	lastUsageMu.Unlock()
}

// GetLastUsage returns the most recent token usage reported by
// llama-server, including the active context-window size. ContextSize
// defaults to 0 until the first request completes.
func GetLastUsage() UsageStats {
	lastUsageMu.RLock()
	defer lastUsageMu.RUnlock()
	return lastUsage
}

// ResetUsage zeroes the usage counter (e.g. after /reset). Preserves the
// context-size hint so the TUI can still render the denominator.
func ResetUsage() {
	lastUsageMu.Lock()
	ctx := lastUsage.ContextSize
	lastUsage = UsageStats{ContextSize: ctx}
	lastUsageMu.Unlock()
}

// ChatComplete sends a chat-style request to the running server. llama-server
// applies the model's own chat template from the GGUF metadata (Gemma-3's
// <start_of_turn>/<end_of_turn> sentinels, ChatML for other families, etc.)
// and returns only the assistant's reply — so the model stops at the turn
// boundary instead of spewing fake "User:/Assistant:" continuations.
func (s *llamaServer) ChatComplete(msgs []ChatMsg, maxTokens int) (string, error) {
	content, _, err := s.chatCompleteCore(msgs, maxTokens, nil)
	return content, err
}

// ChatCompleteWithTools is the tool-enabled variant: advertises `tools` on
// the request and returns the (possibly empty) list of tool_calls the model
// emitted in addition to any assistant content. The caller is responsible
// for executing the calls and re-invoking with the tool results appended.
func (s *llamaServer) ChatCompleteWithTools(msgs []ChatMsg, tools []map[string]any, maxTokens int) (string, []ToolCall, error) {
	return s.chatCompleteCore(msgs, maxTokens, tools)
}

func (s *llamaServer) chatCompleteCore(msgs []ChatMsg, maxTokens int, tools []map[string]any) (string, []ToolCall, error) {
	reqBody, _ := json.Marshal(chatRequest{
		Messages:    msgs,
		MaxTokens:   maxTokens,
		Temperature: 0.2,
		Stream:      false,
		CachePrompt: true,
		Tools:       tools,
	})
	log.Printf("→ POST /v1/chat/completions port=%d max_tokens=%d msgs=%d", s.port, maxTokens, len(msgs))
	for i, m := range msgs {
		log.Printf("  msg[%d] %s (%d bytes): %s", i, m.Role, len(m.Content), truncateForLog(m.Content, 500))
	}

	url := fmt.Sprintf("http://127.0.0.1:%d/v1/chat/completions", s.port)
	start := time.Now()
	resp, err := s.client.Post(url, "application/json", bytes.NewReader(reqBody))
	if err != nil {
		return "", nil, fmt.Errorf("POST /v1/chat/completions: %w", err)
	}
	defer resp.Body.Close()

	raw, readErr := io.ReadAll(resp.Body)
	if readErr != nil {
		return "", nil, fmt.Errorf("read chat completion body: %w", readErr)
	}
	log.Printf("← HTTP %d in %s (body=%d bytes)", resp.StatusCode, time.Since(start), len(raw))
	log.Printf("  body: %s", truncateForLog(string(raw), 2000))

	if resp.StatusCode != 200 {
		return "", nil, fmt.Errorf("llama-server HTTP %d: %s", resp.StatusCode, truncateForLog(string(raw), 300))
	}
	var cr chatResponse
	if err := json.Unmarshal(raw, &cr); err != nil {
		return "", nil, fmt.Errorf("decode chat completion: %w", err)
	}
	if len(cr.Choices) == 0 {
		return "", nil, fmt.Errorf("empty chat completion response")
	}
	msg := cr.Choices[0].Message
	content := msg.Content
	if content == "" && len(msg.ToolCalls) == 0 && cr.Choices[0].FinishReason == "length" {
		log.Printf("  WARNING: reply was empty and finish_reason=length — max_tokens=%d likely too small", maxTokens)
	}
	setLastUsage(UsageStats{
		PromptTokens:     cr.Usage.PromptTokens,
		CompletionTokens: cr.Usage.CompletionTokens,
		TotalTokens:      cr.Usage.TotalTokens,
		ContextSize:      s.ctxN,
	})
	return content, msg.ToolCalls, nil
}

// DropKVCache asks llama-server to forget any cached prompt prefix so the
// next request starts from a clean slate. Called from the TUI's /reset so
// the server doesn't silently reuse tokens from the previous conversation.
func (s *llamaServer) DropKVCache() error {
	url := fmt.Sprintf("http://127.0.0.1:%d/slots/0?action=erase", s.port)
	req, err := http.NewRequest("POST", url, nil)
	if err != nil {
		return err
	}
	resp, err := s.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 400 {
		// Older llama-server builds don't expose this; don't treat it as fatal.
		log.Printf("/slots erase returned %d (probably unsupported on this llama-server)", resp.StatusCode)
		return nil
	}
	return nil
}

func truncateForLog(s string, n int) string {
	// Collapse newlines so each log line stays one line.
	s = strings.ReplaceAll(s, "\r", " ")
	s = strings.ReplaceAll(s, "\n", " ⏎ ")
	if len(s) > n {
		return s[:n] + "…(+" + fmt.Sprintf("%d", len(s)-n) + " more)"
	}
	return s
}

// logWriter forwards subprocess stdout/stderr into the atlas.llm log with a
// tag prefix so server messages show up alongside our own.
type logWriter struct{ tag string }

func newLogWriter(tag string) *logWriter { return &logWriter{tag: tag} }
func (w *logWriter) Write(p []byte) (int, error) {
	log.Printf("[%s] %s", w.tag, bytes.TrimRight(p, "\r\n "))
	return len(p), nil
}