byteclaw/internal/server/server_test.go

package server

import (
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/coder/websocket"
	"github.com/coder/websocket/wsjson"

	"github.com/crypto/byteclaw/internal/agent"
	"github.com/crypto/byteclaw/internal/llm"
	"github.com/crypto/byteclaw/internal/store"
	"github.com/crypto/byteclaw/internal/tools"
)

const tok = "test-token-0123456789"

type echo struct{}

func (echo) Chat(_ context.Context, req llm.Request, d llm.DeltaFunc) (llm.Response, error) {
	last := req.Messages[len(req.Messages)-1].Content
	for _, w := range strings.Fields("echo: " + last) {
		d(w + " ")
	}
	return llm.Response{Content: "echo: " + last}, nil
}

func setup(t *testing.T, origins ...string) *httptest.Server {
	st, err := store.Open(filepath.Join(t.TempDir(), "s.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { st.Close() })
	s := &Server{Store: st, Token: tok, AllowedOrigins: origins, Log: slog.New(slog.NewTextHandler(io.Discard, nil)), RunTimeout: 5 * time.Second,
		Agent: &agent.Agent{LLM: echo{}, Store: st, Tools: tools.NewRegistry(), MaxIterations: 3}}
	ts := httptest.NewServer(s.Handler())
	t.Cleanup(ts.Close)
	return ts
}

func do(t *testing.T, method, url, body string, hdr map[string]string) *http.Response {
	req, _ := http.NewRequest(method, url, strings.NewReader(body))
	for k, v := range hdr {
		req.Header.Set(k, v)
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { resp.Body.Close() })
	return resp
}

var authH = map[string]string{"Authorization": "Bearer " + tok}

func TestAuthAndOrigin(t *testing.T) {
	ts := setup(t, "https://app.example")
	if r := do(t, "GET", ts.URL+"/healthz", "", nil); r.StatusCode != 200 {
		t.Fatal("healthz", r.StatusCode)
	}
	if r := do(t, "GET", ts.URL+"/api/sessions", "", nil); r.StatusCode != 401 {
		t.Fatal("no token", r.StatusCode)
	}
	if r := do(t, "GET", ts.URL+"/api/sessions", "", map[string]string{"Authorization": "Bearer wrong"}); r.StatusCode != 401 {
		t.Fatal("bad token", r.StatusCode)
	}
	evil := map[string]string{"Authorization": "Bearer " + tok, "Origin": "https://evil.example"}
	if r := do(t, "GET", ts.URL+"/api/sessions", "", evil); r.StatusCode != 403 {
		t.Fatal("evil origin", r.StatusCode)
	}
	ok := map[string]string{"Authorization": "Bearer " + tok, "Origin": "https://app.example"}
	if r := do(t, "GET", ts.URL+"/api/sessions", "", ok); r.StatusCode != 200 {
		t.Fatal("allowed origin", r.StatusCode)
	}
	same := map[string]string{"Authorization": "Bearer " + tok, "Origin": ts.URL}
	if r := do(t, "GET", ts.URL+"/api/sessions", "", same); r.StatusCode != 200 {
		t.Fatal("same origin", r.StatusCode)
	}
}

func TestChatAndSessions(t *testing.T) {
	ts := setup(t)
	r := do(t, "POST", ts.URL+"/api/chat", `{"text":"hello world"}`, authH)
	var out map[string]string
	json.NewDecoder(r.Body).Decode(&out)
	if r.StatusCode != 200 || out["reply"] != "echo: hello world" || out["session_id"] == "" {
		t.Fatalf("%d %v", r.StatusCode, out)
	}
	r = do(t, "GET", ts.URL+"/api/sessions", "", authH)
	var list []store.Session
	json.NewDecoder(r.Body).Decode(&list)
	if len(list) != 1 || list[0].Title != "hello world" {
		t.Fatalf("%+v", list)
	}
	r = do(t, "GET", ts.URL+"/api/sessions/"+out["session_id"]+"/messages", "", authH)
	var msgs []store.StoredMessage
	json.NewDecoder(r.Body).Decode(&msgs)
	if len(msgs) != 2 {
		t.Fatalf("msgs %+v", msgs)
	}
	if r := do(t, "DELETE", ts.URL+"/api/sessions/"+out["session_id"], "", authH); r.StatusCode != 204 {
		t.Fatal(r.StatusCode)
	}
	if r := do(t, "GET", ts.URL+"/api/sessions/nope/messages", "", authH); r.StatusCode != 404 {
		t.Fatal(r.StatusCode)
	}
}

func TestWebSocketStreaming(t *testing.T) {
	ts := setup(t)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	wsURL := "ws" + strings.TrimPrefix(ts.URL, "http") + "/api/ws"
	if _, _, err := websocket.Dial(ctx, wsURL, nil); err == nil {
		t.Fatal("ws without token must fail")
	}
	if _, _, err := websocket.Dial(ctx, wsURL, &websocket.DialOptions{HTTPHeader: http.Header{"Authorization": {"Bearer " + tok}, "Origin": {"https://evil.example"}}}); err == nil {
		t.Fatal("ws from foreign origin must fail")
	}
	c, _, err := websocket.Dial(ctx, wsURL, &websocket.DialOptions{HTTPHeader: http.Header{"Authorization": {"Bearer " + tok}}})
	if err != nil {
		t.Fatal(err)
	}
	defer c.CloseNow()
	wsjson.Write(ctx, c, map[string]string{"type": "message", "text": "hi there"})
	var deltas strings.Builder
	var sid string
	for {
		var e agent.Event
		if err := wsjson.Read(ctx, c, &e); err != nil {
			t.Fatal(err)
		}
		switch e.Type {
		case "session":
			sid = e.Session
		case "delta":
			deltas.WriteString(e.Text)
		case "error":
			t.Fatal(e.Text)
		}
		if e.Type == "done" {
			if e.Text != "echo: hi there" || sid == "" || e.Session != sid || strings.TrimSpace(deltas.String()) != "echo: hi there" {
				t.Fatalf("done %+v deltas %q", e, deltas.String())
			}
			return
		}
	}
}