kanban/backend/chat_ws_test.go

package main

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"nhooyr.io/websocket"
)

// fakeClaudeScript writes a bash script that emits NDJSON stream-json events
// to stdout and exits 0. Returns the absolute path of the script.
func fakeClaudeScript(t *testing.T, payload string) string {
	t.Helper()
	if _, err := os.Stat("/bin/bash"); err != nil {
		t.Skip("/bin/bash not available")
	}
	dir := t.TempDir()
	path := filepath.Join(dir, "claude")
	body := "#!/bin/bash\nset -e\ncat <<'__EOF__'\n" + payload + "\n__EOF__\n"
	if err := os.WriteFile(path, []byte(body), 0o755); err != nil {
		t.Fatalf("write fake claude: %v", err)
	}
	return path
}

// chatWSTestServer wires the WebSocket chat handler in front of a test DB.
func chatWSTestServer(t *testing.T) (*httptest.Server, *DB, string) {
	t.Helper()
	db := setupTestDB(t)
	dir := t.TempDir()
	logger := newChatLogger(filepath.Join(dir, "chat.log"))
	token := generateInternalToken()
	srv := httptest.NewServer(handleChatWS(db, dir, logger, token))
	t.Cleanup(srv.Close)
	return srv, db, token
}

func dialChatWS(t *testing.T, srv *httptest.Server) *websocket.Conn {
	t.Helper()
	u, _ := url.Parse(srv.URL)
	wsURL := "ws://" + u.Host
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	c, _, err := websocket.Dial(ctx, wsURL, nil)
	if err != nil {
		t.Fatalf("dial %s: %v", wsURL, err)
	}
	return c
}

func readWSEvent(t *testing.T, conn *websocket.Conn) wsEvent {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, data, err := conn.Read(ctx)
	if err != nil {
		t.Fatalf("read: %v", err)
	}
	var ev wsEvent
	if err := json.Unmarshal(data, &ev); err != nil {
		t.Fatalf("unmarshal %q: %v", string(data), err)
	}
	return ev
}

func sendInitial(t *testing.T, conn *websocket.Conn, msgs []chatMessage) {
	t.Helper()
	body, _ := json.Marshal(chatRequest{Messages: msgs})
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := conn.Write(ctx, websocket.MessageText, body); err != nil {
		t.Fatalf("write: %v", err)
	}
}

// --- WS streaming tests ---------------------------------------------------

func TestChatWS_StreamsTextDelta(t *testing.T) {
	payload := `{"type":"system","subtype":"init","session_id":"s1","model":"test"}
{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"Hola "}]}}
{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"mundo"}]}}
{"type":"result","subtype":"success","is_error":false,"result":"Hola mundo","stop_reason":"end_turn"}`

	t.Setenv("KANBAN_CLAUDE_BIN", fakeClaudeScript(t, payload))

	srv, _, _ := chatWSTestServer(t)
	conn := dialChatWS(t, srv)
	defer conn.Close(websocket.StatusNormalClosure, "")

	sendInitial(t, conn, []chatMessage{{Role: "user", Content: "saluda"}})

	var deltas []string
	var sawResult, sawDone bool
	for i := 0; i < 12 && !sawDone; i++ {
		ev := readWSEvent(t, conn)
		switch ev.Type {
		case "delta":
			deltas = append(deltas, ev.Text)
		case "result":
			sawResult = true
		case "done":
			sawDone = true
		case "error":
			t.Fatalf("unexpected error event: %s", ev.Error)
		}
	}
	if !sawDone {
		t.Fatalf("never received done event")
	}
	if !sawResult {
		t.Fatalf("never received result event")
	}
	if got := strings.Join(deltas, ""); got != "Hola mundo" {
		t.Fatalf("expected 'Hola mundo' from deltas, got %q", got)
	}
}

func TestChatWS_StreamsToolUseAndResult(t *testing.T) {
	payload := `{"type":"system","subtype":"init"}
{"type":"assistant","message":{"role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"mcp__kanban__create_column","input":{"name":"Backlog"}}]}}
{"type":"user","message":{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_1","content":"{\"ok\":true,\"result\":{\"id\":\"col_x\"}}","is_error":false}]}}
{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"Listo"}]}}
{"type":"result","subtype":"success","is_error":false,"result":"Listo","stop_reason":"end_turn"}`

	t.Setenv("KANBAN_CLAUDE_BIN", fakeClaudeScript(t, payload))

	srv, _, _ := chatWSTestServer(t)
	conn := dialChatWS(t, srv)
	defer conn.Close(websocket.StatusNormalClosure, "")

	sendInitial(t, conn, []chatMessage{{Role: "user", Content: "crea Backlog"}})

	var sawToolUse, sawToolResult, sawDelta, sawDone bool
	var doneEv wsEvent
	for i := 0; i < 16 && !sawDone; i++ {
		ev := readWSEvent(t, conn)
		switch ev.Type {
		case "tool_use":
			sawToolUse = true
			if ev.Tool != "create_column" {
				t.Errorf("tool name not stripped: %q", ev.Tool)
			}
			if !strings.Contains(string(ev.Input), "Backlog") {
				t.Errorf("input missing Backlog: %s", ev.Input)
			}
		case "tool_result":
			sawToolResult = true
			if ev.IsError {
				t.Errorf("tool_result is_error true")
			}
		case "delta":
			sawDelta = true
		case "done":
			sawDone = true
			doneEv = ev
		case "error":
			t.Fatalf("unexpected error: %s", ev.Error)
		}
	}
	if !sawToolUse || !sawToolResult || !sawDelta || !sawDone {
		t.Fatalf("missing events: tool_use=%v tool_result=%v delta=%v done=%v",
			sawToolUse, sawToolResult, sawDelta, sawDone)
	}
	if !doneEv.BoardChanged {
		t.Errorf("expected board_changed=true (create_column is a mutator)")
	}
}

func TestChatWS_RejectsEmptyMessages(t *testing.T) {
	t.Setenv("KANBAN_CLAUDE_BIN", fakeClaudeScript(t,
		`{"type":"result","subtype":"success","is_error":false,"result":""}`))

	srv, _, _ := chatWSTestServer(t)
	conn := dialChatWS(t, srv)
	defer conn.Close(websocket.StatusNormalClosure, "")

	sendInitial(t, conn, []chatMessage{})
	ev := readWSEvent(t, conn)
	if ev.Type != "error" {
		t.Fatalf("expected error event, got %+v", ev)
	}
	if !strings.Contains(ev.Error, "messages required") {
		t.Fatalf("unexpected error: %s", ev.Error)
	}
}

func TestChatWS_PropagatesClaudeFailure(t *testing.T) {
	dir := t.TempDir()
	bin := filepath.Join(dir, "claude")
	body := "#!/bin/bash\necho 'broken' >&2\nexit 7\n"
	if err := os.WriteFile(bin, []byte(body), 0o755); err != nil {
		t.Fatalf("write: %v", err)
	}
	t.Setenv("KANBAN_CLAUDE_BIN", bin)

	srv, _, _ := chatWSTestServer(t)
	conn := dialChatWS(t, srv)
	defer conn.Close(websocket.StatusNormalClosure, "")

	sendInitial(t, conn, []chatMessage{{Role: "user", Content: "hola"}})

	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		ev := readWSEvent(t, conn)
		switch ev.Type {
		case "error":
			if !strings.Contains(ev.Error, "claude exit") {
				t.Fatalf("expected claude exit error, got: %s", ev.Error)
			}
			return
		case "done":
			t.Fatalf("done received before error")
		}
	}
	t.Fatalf("never received error event")
}

// --- /api/tool internal endpoint tests ------------------------------------

func internalToolServer(t *testing.T) (*httptest.Server, *DB, string) {
	t.Helper()
	db := setupTestDB(t)
	logger := newChatLogger(filepath.Join(t.TempDir(), "log"))
	token := generateInternalToken()
	mux := http.NewServeMux()
	mux.Handle("POST /api/tool/{name}", handleInternalTool(db, token, logger))
	srv := httptest.NewServer(mux)
	t.Cleanup(srv.Close)
	return srv, db, token
}

func TestInternalTool_CreateColumnRoundtrip(t *testing.T) {
	srv, db, token := internalToolServer(t)
	req, _ := http.NewRequest("POST", srv.URL+"/api/tool/create_column", strings.NewReader(`{"name":"Backlog"}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set(internalTokenHeader, token)
	resp, err := srv.Client().Do(req)
	if err != nil {
		t.Fatalf("do: %v", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		t.Fatalf("status %d", resp.StatusCode)
	}
	var tr ToolResult
	if err := json.NewDecoder(resp.Body).Decode(&tr); err != nil {
		t.Fatalf("decode: %v", err)
	}
	if !tr.OK {
		t.Fatalf("create_column failed: %s", tr.Error)
	}
	cols, err := db.ListColumns()
	if err != nil {
		t.Fatalf("list: %v", err)
	}
	if len(cols) != 1 || cols[0].Name != "Backlog" {
		t.Fatalf("expected 1 col Backlog, got %+v", cols)
	}
}

func TestInternalTool_RejectsMissingToken(t *testing.T) {
	srv, _, _ := internalToolServer(t)
	req, _ := http.NewRequest("POST", srv.URL+"/api/tool/create_column", strings.NewReader(`{"name":"X"}`))
	req.Header.Set("Content-Type", "application/json")
	resp, err := srv.Client().Do(req)
	if err != nil {
		t.Fatalf("do: %v", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != 401 {
		t.Fatalf("expected 401, got %d", resp.StatusCode)
	}
}

func TestInternalTool_UnknownTool(t *testing.T) {
	srv, _, token := internalToolServer(t)
	req, _ := http.NewRequest("POST", srv.URL+"/api/tool/no_such", strings.NewReader(`{}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set(internalTokenHeader, token)
	resp, err := srv.Client().Do(req)
	if err != nil {
		t.Fatalf("do: %v", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != 404 {
		t.Fatalf("expected 404, got %d", resp.StatusCode)
	}
}