opencode-cursor-agent/internal/server/messages_test.go

package server

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/daniel/cursor-adapter/internal/config"
)

type mockBridge struct {
	executeLines   []string
	executeErr     error
	executeSync    string
	executeSyncErr error
	lastPrompt     string
	lastSessionKey string
	models         []string
	healthErr      error
}

func (m *mockBridge) Execute(ctx context.Context, prompt string, model string, sessionKey string) (<-chan string, <-chan error) {
	m.lastPrompt = prompt
	m.lastSessionKey = sessionKey
	outputChan := make(chan string, len(m.executeLines))
	errChan := make(chan error, 1)

	go func() {
		defer close(outputChan)
		defer close(errChan)
		for _, line := range m.executeLines {
			select {
			case <-ctx.Done():
				errChan <- ctx.Err()
				return
			case outputChan <- line:
			}
		}
		if m.executeErr != nil {
			errChan <- m.executeErr
		}
	}()

	return outputChan, errChan
}

func (m *mockBridge) ListModels(ctx context.Context) ([]string, error) {
	return m.models, nil
}

func (m *mockBridge) ExecuteSync(ctx context.Context, prompt string, model string, sessionKey string) (string, error) {
	m.lastPrompt = prompt
	m.lastSessionKey = sessionKey
	if m.executeSyncErr != nil {
		return "", m.executeSyncErr
	}
	if m.executeSync != "" {
		return m.executeSync, nil
	}
	return "", nil
}

func (m *mockBridge) CheckHealth(ctx context.Context) error {
	return m.healthErr
}

func TestAnthropicMessages_NonStreamingResponse(t *testing.T) {
	cfg := config.Defaults()
	srv := New(&cfg, &mockBridge{
		executeSync: "Hello",
	})

	req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{
		"model":"auto",
		"max_tokens":128,
		"messages":[{"role":"user","content":"Say hello"}],
		"stream":false
	}`))
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()

	srv.mux.ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String())
	}

	var resp struct {
		ID         string `json:"id"`
		Type       string `json:"type"`
		Role       string `json:"role"`
		Model      string `json:"model"`
		StopReason string `json:"stop_reason"`
		Content    []struct {
			Type string `json:"type"`
			Text string `json:"text"`
		} `json:"content"`
		Usage struct {
			InputTokens  int `json:"input_tokens"`
			OutputTokens int `json:"output_tokens"`
		} `json:"usage"`
	}

	if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
		t.Fatalf("unmarshal response: %v", err)
	}

	if resp.Type != "message" {
		t.Fatalf("type = %q, want %q", resp.Type, "message")
	}
	if resp.Role != "assistant" {
		t.Fatalf("role = %q, want %q", resp.Role, "assistant")
	}
	if len(resp.Content) != 1 || resp.Content[0].Text != "Hello" {
		t.Fatalf("content = %+v, want single text block 'Hello'", resp.Content)
	}
	if resp.StopReason != "end_turn" {
		t.Fatalf("stop_reason = %q, want %q", resp.StopReason, "end_turn")
	}
	if resp.Usage.InputTokens <= 0 || resp.Usage.OutputTokens <= 0 {
		t.Fatalf("usage should be estimated and > 0, got %+v", resp.Usage)
	}
}

func TestAnthropicMessages_StreamingResponse(t *testing.T) {
	cfg := config.Defaults()
	srv := New(&cfg, &mockBridge{
		executeLines: []string{
			`{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"Hi"}]}}`,
			`{"type":"result","subtype":"success","usage":{"inputTokens":9,"outputTokens":1}}`,
		},
	})

	req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{
		"model":"auto",
		"max_tokens":128,
		"messages":[{"role":"user","content":"Say hi"}],
		"stream":true
	}`))
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()

	srv.mux.ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String())
	}

	body := rec.Body.String()
	for _, want := range []string{
		`"type":"message_start"`,
		`"type":"content_block_start"`,
		`"type":"content_block_delta"`,
		`"text":"Hi"`,
		`"type":"content_block_stop"`,
		`"type":"message_delta"`,
		`"stop_reason":"end_turn"`,
		`"type":"message_stop"`,
	} {
		if !strings.Contains(body, want) {
			t.Fatalf("stream body missing %q: %s", want, body)
		}
	}
}

func TestChatCompletions_ForwardsProvidedSessionHeader(t *testing.T) {
	cfg := config.Defaults()
	br := &mockBridge{executeSync: "Hello"}
	srv := New(&cfg, br)

	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"auto",
		"messages":[{"role":"user","content":"hello"}],
		"stream":false
	}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set(sessionHeaderName, "sess_frontend_123")
	rec := httptest.NewRecorder()

	srv.mux.ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String())
	}
	if br.lastSessionKey != "sess_frontend_123" {
		t.Fatalf("bridge session key = %q, want %q", br.lastSessionKey, "sess_frontend_123")
	}
	if got := rec.Header().Get(sessionHeaderName); got != "sess_frontend_123" {
		t.Fatalf("response session header = %q, want %q", got, "sess_frontend_123")
	}
}

func TestChatCompletions_AcceptsArrayContentBlocks(t *testing.T) {
	cfg := config.Defaults()
	br := &mockBridge{executeSync: "Hello"}
	srv := New(&cfg, br)

	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"auto",
		"messages":[
			{"role":"system","content":[{"type":"text","text":"You are terse."}]},
			{"role":"user","content":[{"type":"text","text":"hello"},{"type":"text","text":" world"}]}
		],
		"stream":false
	}`))
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()

	srv.mux.ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String())
	}
	if !strings.Contains(br.lastPrompt, "system: You are terse.") {
		t.Fatalf("prompt = %q, want system text content", br.lastPrompt)
	}
	if !strings.Contains(br.lastPrompt, "user: hello world") {
		t.Fatalf("prompt = %q, want concatenated user text content", br.lastPrompt)
	}
}

func TestChatCompletions_StreamingEmitsRoleFinishReasonAndUsage(t *testing.T) {
	cfg := config.Defaults()
	srv := New(&cfg, &mockBridge{
		executeLines: []string{
			// system + user chunks should be skipped entirely, never echoed as content
			`{"type":"system","subtype":"init","session_id":"abc","cwd":"/tmp"}`,
			`{"type":"user","message":{"role":"user","content":[{"type":"text","text":"user: hello"}]}}`,
			// incremental assistant fragments
			`{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"你"}]}}`,
			`{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"好"}]}}`,
			// cumulative duplicate (Cursor CLI sometimes finalises with the full text)
			`{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"你好"}]}}`,
			`{"type":"result","subtype":"success","result":"你好","usage":{"inputTokens":3,"outputTokens":2}}`,
		},
	})

	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"auto",
		"messages":[{"role":"user","content":"hi"}],
		"stream":true
	}`))
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()

	srv.mux.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
	}

	body := rec.Body.String()

	// Must never leak system/user JSON as "content"
	if strings.Contains(body, `"subtype":"init"`) || strings.Contains(body, `"type":"user"`) {
		t.Fatalf("stream body leaked system/user JSON into SSE content: %s", body)
	}

	// First delta chunk must carry role:assistant (not content)
	if !strings.Contains(body, `"delta":{"role":"assistant"}`) {
		t.Fatalf("first chunk missing role=assistant delta: %s", body)
	}

	// Content deltas must be plain text — not JSON-stringified Cursor lines
	if !strings.Contains(body, `"delta":{"content":"你"}`) {
		t.Fatalf("first content delta not plain text: %s", body)
	}
	if !strings.Contains(body, `"delta":{"content":"好"}`) {
		t.Fatalf("second content delta missing: %s", body)
	}

	// Final cumulative message that equals accumulated text must be suppressed
	// (accumulated = "你好" after the two fragments; final "你好" should be Skip'd)
	count := strings.Count(body, `"你好"`)
	if count > 0 {
		t.Fatalf("duplicate final cumulative message should have been skipped (found %d occurrences of full text as delta): %s", count, body)
	}

	// Final chunk must have finish_reason=stop and usage at top level
	if !strings.Contains(body, `"finish_reason":"stop"`) {
		t.Fatalf("final chunk missing finish_reason=stop: %s", body)
	}
	if !strings.Contains(body, `"usage":{`) {
		t.Fatalf("final chunk missing usage: %s", body)
	}
	if !strings.Contains(body, `data: [DONE]`) {
		t.Fatalf("stream missing [DONE] terminator: %s", body)
	}
}

func TestAnthropicMessages_StreamingEmitsNoDuplicateFinalText(t *testing.T) {
	cfg := config.Defaults()
	srv := New(&cfg, &mockBridge{
		executeLines: []string{
			`{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"你"}]}}`,
			`{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"好"}]}}`,
			`{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"你好"}]}}`,
			`{"type":"result","subtype":"success","usage":{"inputTokens":3,"outputTokens":2}}`,
		},
	})

	req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{
		"model":"auto",
		"max_tokens":128,
		"messages":[{"role":"user","content":"hi"}],
		"stream":true
	}`))
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()

	srv.mux.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
	}

	body := rec.Body.String()
	if strings.Count(body, `"text":"你好"`) > 0 {
		t.Fatalf("final cumulative duplicate should be suppressed: %s", body)
	}
	for _, want := range []string{
		`"text":"你"`,
		`"text":"好"`,
		`"type":"message_stop"`,
	} {
		if !strings.Contains(body, want) {
			t.Fatalf("missing %q in stream: %s", want, body)
		}
	}
}

func TestAnthropicMessages_GeneratesSessionHeaderWhenMissing(t *testing.T) {
	cfg := config.Defaults()
	br := &mockBridge{executeSync: "Hello"}
	srv := New(&cfg, br)

	req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{
		"model":"auto",
		"max_tokens":128,
		"messages":[{"role":"user","content":"hello"}],
		"stream":false
	}`))
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()

	srv.mux.ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String())
	}
	if br.lastSessionKey == "" {
		t.Fatal("expected generated session key to be forwarded to bridge")
	}
	if got := rec.Header().Get(sessionHeaderName); got == "" {
		t.Fatal("expected generated session header in response")
	}
	if got := rec.Header().Get(exposeHeadersName); !strings.Contains(got, sessionHeaderName) {
		t.Fatalf("expose headers = %q, want to contain %q", got, sessionHeaderName)
	}
}