Vectrum
Sign In
141 lines · 5.0 KB · Go
// Command mockapi is a fake Messages API for exercising the TUI by hand,
// without an API key and without spending anything.
//
// It lives under testdata so it is invisible to ./... — it is a manual test
// fixture, not part of the build.
//
//	go run ./internal/tui/testdata/mockapi &
//	./codey -C /tmp/demo -base-url http://127.0.0.1:8899 -api-key mock
//
// The script is fixed: the first turn streams a sentence and asks to run a
// shell command, which makes the TUI raise an approval prompt; once a
// tool_result comes back it streams a closing sentence and stops. That covers
// streaming, tool rendering, approve, deny, and interrupt.
//
// A compaction request is answered separately, as plain JSON rather than SSE,
// because the summarization call does not stream. Reported usage grows with the
// size of the request, so the gauge climbs as you talk; pair it with a small
// -context-window to reach a threshold in a handful of turns:
//
//	./codey -base-url http://127.0.0.1:8899 -api-key mock -context-window 30000
//
// Flags: -addr to move it, -fail to make the command exit non-zero, -slow to
// stretch the stream out for watching the spinner and ctrl-c.
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"time"
)

var (
	addr = flag.String("addr", "127.0.0.1:8899", "listen address")
	slow = flag.Duration("slow", 0, "pause between streamed words, e.g. 150ms")
	cmd  = flag.String("cmd", "git status --short && go build ./...", "command the model asks to run")
)

func main() {
	flag.Parse()
	http.HandleFunc("/v1/messages", handle)
	log.Printf("mock Messages API on http://%s", *addr)
	log.Fatal(http.ListenAndServe(*addr, nil))
}

func handle(w http.ResponseWriter, r *http.Request) {
	body, _ := io.ReadAll(r.Body)

	// Compaction is the one request Codey does not stream, so it cannot
	// be answered with the SSE script below.
	if strings.Contains(string(body), "Summarize the conversation above") {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]any{
			"id": "msg_mock_summary", "type": "message", "role": "assistant",
			"model": "claude-opus-5", "stop_reason": "end_turn",
			"content": []map[string]any{{
				"type": "text",
				"text": "Ran git status and go build; the tree is clean and it compiles.",
			}},
			"usage": map[string]any{"input_tokens": 900, "output_tokens": 40},
		})
		return
	}

	// A tool_result in the history means Codey ran (or refused) our
	// command and fed the outcome back, so the next turn is the last one.
	final := strings.Contains(string(body), "tool_result")

	w.Header().Set("Content-Type", "text/event-stream")
	w.WriteHeader(http.StatusOK)

	// Usage grows with the request, so the context gauge climbs as a real
	// conversation would and the compaction thresholds can be reached by
	// talking rather than only by shrinking -context-window.
	read, fresh := estimateTokens(len(body))
	send(w, "message_start", map[string]any{
		"type": "message_start",
		"message": map[string]any{
			"id": "msg_mock", "model": "claude-opus-5",
			"usage": map[string]any{"input_tokens": fresh, "cache_read_input_tokens": read},
		},
	})

	text := "Checking the repository state first."
	if final {
		text = "That is everything I can tell from here."
	}
	send(w, "content_block_start", map[string]any{
		"type": "content_block_start", "index": 0,
		"content_block": map[string]any{"type": "text", "text": ""},
	})
	for _, word := range strings.SplitAfter(text, " ") {
		send(w, "content_block_delta", map[string]any{
			"type": "content_block_delta", "index": 0,
			"delta": map[string]any{"type": "text_delta", "text": word},
		})
		time.Sleep(*slow)
	}
	send(w, "content_block_stop", map[string]any{"type": "content_block_stop", "index": 0})

	stop := "end_turn"
	if !final {
		stop = "tool_use"
		send(w, "content_block_start", map[string]any{
			"type": "content_block_start", "index": 1,
			"content_block": map[string]any{"type": "tool_use", "id": "toolu_1", "name": "bash"},
		})
		args, _ := json.Marshal(map[string]string{"command": *cmd})
		send(w, "content_block_delta", map[string]any{
			"type": "content_block_delta", "index": 1,
			"delta": map[string]any{"type": "input_json_delta", "partial_json": string(args)},
		})
		send(w, "content_block_stop", map[string]any{"type": "content_block_stop", "index": 1})
	}

	send(w, "message_delta", map[string]any{
		"type":  "message_delta",
		"delta": map[string]any{"stop_reason": stop},
		"usage": map[string]any{"output_tokens": 95},
	})
	send(w, "message_stop", map[string]any{"type": "message_stop"})
}

// estimateTokens turns a request body size into a plausible cached/fresh split
// at roughly four bytes per token. The exact numbers do not matter; that they
// rise with the conversation does.
func estimateTokens(bodyLen int) (cached, fresh int) {
	total := 1200 + bodyLen/4
	cached = total * 4 / 5
	return cached, total - cached
}

func send(w http.ResponseWriter, event string, data any) {
	b, _ := json.Marshal(data)
	fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event, b)
	w.(http.Flusher).Flush()
}