ombi-mcp/internal/integration_test/harness_test.go

//go:build integration

// Package integration_test exercises the compiled ombi-mcp binary end
// to end: it spawns the server as a subprocess and issues raw MCP
// JSON-RPC messages over stdio, against either a mock Ombi instance
// (httptest) or a live one selected by environment variables.
//
// The whole package is behind the "integration" build tag so standard
// CI (`go test ./...`) never compiles or runs it. Run explicitly with:
//
//	go test -tags integration ./internal/integration_test/
package integration_test

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

// serverBin is the compiled server under test, built once in TestMain.
var serverBin string

func TestMain(m *testing.M) {
	dir, err := os.MkdirTemp("", "ombi-mcp-itest")
	if err != nil {
		fmt.Fprintln(os.Stderr, "mktemp:", err)
		os.Exit(1)
	}
	defer os.RemoveAll(dir)
	bin := filepath.Join(dir, "ombi-mcp")
	root, err := filepath.Abs("../..")
	if err != nil {
		fmt.Fprintln(os.Stderr, "resolve repo root:", err)
		os.Exit(1)
	}
	build := exec.Command("go", "build", "-o", bin, "./cmd/ombi-mcp")
	build.Dir = root
	build.Env = append(os.Environ(), "CGO_ENABLED=0")
	if out, err := build.CombinedOutput(); err != nil {
		fmt.Fprintf(os.Stderr, "build server: %v\n%s\n", err, out)
		os.Exit(1)
	}
	serverBin = bin
	os.Exit(m.Run())
}

// serverEnv returns os.Environ with every OMBI_* variable replaced by
// the given overrides, so a spawned server sees exactly the intended
// configuration regardless of the parent environment.
func serverEnv(overrides map[string]string) []string {
	env := []string{}
	for _, e := range os.Environ() {
		if strings.HasPrefix(e, "OMBI_") {
			continue
		}
		env = append(env, e)
	}
	for k, v := range overrides {
		env = append(env, k+"="+v)
	}
	return env
}

// --- raw JSON-RPC stdio client ---

type rpcRequest struct {
	JSONRPC string `json:"jsonrpc"`
	ID      int64  `json:"id,omitempty"`
	Method  string `json:"method"`
	Params  any    `json:"params,omitempty"`
}

type rpcMessage struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id,omitempty"`
	Method  string          `json:"method,omitempty"`
	Params  json.RawMessage `json:"params,omitempty"`
	Result  json.RawMessage `json:"result,omitempty"`
	Error   *rpcError       `json:"error,omitempty"`
}

type rpcError struct {
	Code    int             `json:"code"`
	Message string          `json:"message"`
	Data    json.RawMessage `json:"data,omitempty"`
}

func (e *rpcError) Error() string {
	return fmt.Sprintf("jsonrpc error %d: %s", e.Code, e.Message)
}

// syncBuf captures the server's stderr for diagnostics on failure.
type syncBuf struct {
	mu  sync.Mutex
	buf bytes.Buffer
}

func (b *syncBuf) Write(p []byte) (int, error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.buf.Len() < 64<<10 {
		return b.buf.Write(p)
	}
	return len(p), nil
}

func (b *syncBuf) String() string {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.buf.String()
}

// rpcClient is a spawned ombi-mcp server spoken to over stdio with
// newline-delimited JSON-RPC 2.0. A single reader goroutine decodes
// all inbound messages onto a channel; request() matches responses
// by id and answers stray server-initiated requests with -32601.
type rpcClient struct {
	t      *testing.T
	cmd    *exec.Cmd
	stdin  io.WriteCloser
	msgs   chan *rpcMessage
	nextID atomic.Int64
	stderr *syncBuf
	done   chan struct{}
}

func spawnServer(t *testing.T, env map[string]string) *rpcClient {
	t.Helper()
	cmd := exec.Command(serverBin)
	cmd.Env = serverEnv(env)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		t.Fatalf("stdout pipe: %v", err)
	}
	stdin, err := cmd.StdinPipe()
	if err != nil {
		t.Fatalf("stdin pipe: %v", err)
	}
	stderr := &syncBuf{}
	cmd.Stderr = stderr
	if err := cmd.Start(); err != nil {
		t.Fatalf("start server: %v", err)
	}
	c := &rpcClient{
		t:      t,
		cmd:    cmd,
		stdin:  stdin,
		msgs:   make(chan *rpcMessage, 64),
		stderr: stderr,
		done:   make(chan struct{}),
	}
	go c.readLoop(stdout)
	t.Cleanup(func() {
		c.stdin.Close()
		select {
		case <-c.done:
		case <-time.After(5 * time.Second):
			c.cmd.Process.Kill()
			<-c.done
		}
	})
	return c
}

func (c *rpcClient) readLoop(r io.Reader) {
	defer close(c.done)
	defer close(c.msgs)
	dec := json.NewDecoder(r)
	for {
		var m rpcMessage
		if err := dec.Decode(&m); err != nil {
			return
		}
		c.msgs <- &m
	}
}

// request sends one JSON-RPC request and waits for its response.
// Stray server-initiated requests get a method-not-found reply;
// notifications and unrelated responses are skipped.
func (c *rpcClient) request(method string, params any) (json.RawMessage, error) {
	id := c.nextID.Add(1)
	body, err := json.Marshal(rpcRequest{
		JSONRPC: "2.0", ID: id, Method: method, Params: params,
	})
	if err != nil {
		return nil, err
	}
	if _, err := c.stdin.Write(append(body, '\n')); err != nil {
		return nil, fmt.Errorf("write %s: %w", method, err)
	}
	timer := time.NewTimer(60 * time.Second)
	defer timer.Stop()
	for {
		select {
		case m, ok := <-c.msgs:
			if !ok {
				return nil, fmt.Errorf("server stdout closed (stderr: %s)", c.stderr)
			}
			if m.Method != "" && len(m.ID) > 0 {
				// A server-initiated request: refuse politely and
				// keep waiting for our own response.
				c.replyError(m.ID, -32601, "client does not implement "+m.Method)
				continue
			}
			var got int64
			if err := json.Unmarshal(m.ID, &got); err != nil || got != id {
				continue // notification or other response — skip
			}
			if m.Error != nil {
				return nil, m.Error
			}
			return m.Result, nil
		case <-timer.C:
			return nil, fmt.Errorf("timeout waiting for %s response", method)
		}
	}
}

func (c *rpcClient) replyError(id json.RawMessage, code int, msg string) {
	body, _ := json.Marshal(map[string]any{
		"jsonrpc": "2.0",
		"id":      id,
		"error":   map[string]any{"code": code, "message": msg},
	})
	c.stdin.Write(append(body, '\n'))
}

func (c *rpcClient) notify(method string, params any) {
	body, _ := json.Marshal(map[string]any{
		"jsonrpc": "2.0", "method": method, "params": params,
	})
	c.stdin.Write(append(body, '\n'))
}

// handshake performs initialize + notifications/initialized.
func (c *rpcClient) handshake(t *testing.T) {
	t.Helper()
	res, err := c.request("initialize", map[string]any{
		"protocolVersion": "2025-06-18",
		"capabilities":    map[string]any{},
		"clientInfo":      map[string]any{"name": "integration-harness", "version": "0.1.0"},
	})
	if err != nil {
		t.Fatalf("initialize: %v (stderr: %s)", err, c.stderr)
	}
	var init struct {
		ServerInfo struct {
			Name    string `json:"name"`
			Version string `json:"version"`
		} `json:"serverInfo"`
	}
	if err := json.Unmarshal(res, &init); err != nil {
		t.Fatalf("initialize result decode: %v", err)
	}
	if init.ServerInfo.Name != "ombi-mcp" {
		t.Fatalf("unexpected server identity %q", init.ServerInfo.Name)
	}
	c.notify("notifications/initialized", map[string]any{})
}

// --- tool call envelope ---

type toolError struct {
	Code              string `json:"code"`
	Message           string `json:"message"`
	Retryable         bool   `json:"retryable"`
	HTTPStatus        *int   `json:"http_status"`
	Field             string `json:"field"`
	RetryAfterSeconds *int   `json:"retry_after_seconds"`
}

type envelope struct {
	OK            bool            `json:"ok"`
	Operation     string          `json:"operation"`
	Data          json.RawMessage `json:"data"`
	Error         *toolError      `json:"error"`
	Warnings      []string        `json:"warnings"`
	Truncated     bool            `json:"truncated"`
	CorrelationID string          `json:"correlation_id"`
}

// callOutcome carries the decoded tools/call result plus the raw
// result object for leak scanning.
type callOutcome struct {
	Raw      json.RawMessage
	Envelope *envelope
	IsError  bool
	Text     string
}

// callTool issues a raw tools/call and decodes the result envelope.
// Protocol-level errors (JSON-RPC error) fail the test.
func (c *rpcClient) callTool(t *testing.T, name string, args any) *callOutcome {
	t.Helper()
	res, err := c.request("tools/call", map[string]any{
		"name":      name,
		"arguments": args,
	})
	if err != nil {
		t.Fatalf("tools/call %s: %v (stderr: %s)", name, err, c.stderr)
	}
	return decodeCallOutcome(t, res)
}

func decodeCallOutcome(t *testing.T, res json.RawMessage) *callOutcome {
	t.Helper()
	var parsed struct {
		Content []struct {
			Type string `json:"type"`
			Text string `json:"text"`
		} `json:"content"`
		Structured *envelope       `json:"structuredContent"`
		IsError    bool            `json:"isError"`
		Raw        json.RawMessage `json:"-"`
	}
	if err := json.Unmarshal(res, &parsed); err != nil {
		t.Fatalf("tools/call result not JSON: %v\n%s", err, res)
	}
	out := &callOutcome{Raw: res, IsError: parsed.IsError}
	if parsed.Structured == nil {
		t.Fatalf("tools/call result has no structuredContent: %s", res)
	}
	out.Envelope = parsed.Structured
	for _, blk := range parsed.Content {
		if blk.Type == "text" {
			out.Text = blk.Text
		}
	}
	return out
}

// --- shared assertions ---

// requireOK asserts the envelope is a success and returns its data.
func requireOK(t *testing.T, o *callOutcome) json.RawMessage {
	t.Helper()
	env := o.Envelope
	if env.CorrelationID == "" {
		t.Errorf("missing correlation_id")
	}
	if o.IsError {
		t.Fatalf("isError=true with envelope: %s", o.Raw)
	}
	if !env.OK {
		t.Fatalf("ok=false: %s", o.Raw)
	}
	if len(env.Data) == 0 || string(env.Data) == "null" {
		t.Fatalf("ok=true but data missing: %s", o.Raw)
	}
	// Text fallback must carry the same bounded envelope.
	var textEnv envelope
	if err := json.Unmarshal([]byte(o.Text), &textEnv); err != nil {
		t.Fatalf("text content is not the serialized envelope: %v", err)
	}
	if textEnv.OK != env.OK || textEnv.CorrelationID != env.CorrelationID {
		t.Fatalf("text/structured envelopes disagree")
	}
	return env.Data
}

// requireErr asserts the envelope is a mapped ToolError with the
// expected code and runs the leak scan over the whole raw result.
func requireErr(t *testing.T, o *callOutcome, code string) *toolError {
	t.Helper()
	env := o.Envelope
	if env.OK {
		t.Fatalf("expected error %s, got ok=true: %s", code, o.Raw)
	}
	if !o.IsError {
		t.Errorf("isError=false on error result")
	}
	if env.Error == nil {
		t.Fatalf("ok=false but no error object: %s", o.Raw)
	}
	if env.Error.Code != code {
		t.Fatalf("error code = %q, want %q: %s", env.Error.Code, code, o.Raw)
	}
	assertNoLeak(t, o.Raw)
	return env.Error
}

// forbiddenLeak substrings that must never appear in tool output:
// upstream credentials, HTML error pages and stack-trace shapes.
func assertNoLeak(t *testing.T, raw json.RawMessage) {
	t.Helper()
	s := string(raw)
	for _, bad := range []string{
		"Bearer ", "Authorization", "<html", "<!DOCTYPE", "<HTML",
		"ApiKey", "api_key", "System.Exception", "NullReferenceException",
		"   at ", "StackTrace", "traceback",
	} {
		if strings.Contains(s, bad) {
			t.Errorf("tool output leaks %q:\n%s", bad, raw)
		}
	}
}

// listTools returns the advertised tool names.
func (c *rpcClient) listTools(t *testing.T) []string {
	t.Helper()
	res, err := c.request("tools/list", map[string]any{})
	if err != nil {
		t.Fatalf("tools/list: %v", err)
	}
	var out struct {
		Tools []struct {
			Name string `json:"name"`
		} `json:"tools"`
	}
	if err := json.Unmarshal(res, &out); err != nil {
		t.Fatalf("tools/list decode: %v", err)
	}
	names := make([]string, 0, len(out.Tools))
	for _, tl := range out.Tools {
		names = append(names, tl.Name)
	}
	return names
}