byteclaw/internal/server/extra_test.go

package server

import (
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"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/config"
	"github.com/crypto/byteclaw/internal/llm"
	"github.com/crypto/byteclaw/internal/providers"
	"github.com/crypto/byteclaw/internal/store"
	"github.com/crypto/byteclaw/internal/tools"
)

// slow streams a first chunk then blocks until the context is cancelled.
type slow struct{}

func (slow) Chat(ctx context.Context, _ llm.Request, d llm.DeltaFunc) (llm.Response, error) {
	d("halb")
	<-ctx.Done()
	return llm.Response{}, ctx.Err()
}

func setupFull(t *testing.T, p llm.Provider) (*httptest.Server, *store.Store) {
	st, err := store.Open(filepath.Join(t.TempDir(), "s.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { st.Close() })
	cfg := config.LLM{Provider: "openai", BaseURL: "https://api.openai.com/v1", Model: "gpt-4o-mini", APIKey: "sk-SECRET-openai"}
	pm := providers.New(cfg, func(k string) string {
		if k == "DAHL_API_KEY" {
			return "dahl-SECRET"
		}
		return ""
	}, st)
	pm.New = func(string, string, string) llm.Provider { return p }
	ag := &agent.Agent{Store: st, Tools: tools.NewRegistry(), MaxIterations: 3, Summarize: agent.SummarizeConfig{KeepRecent: 1}}
	ag.Resolve = func(ctx context.Context, sid string) (llm.Provider, error) {
		pr, _, err := pm.ForSession(ctx, sid)
		return pr, err
	}
	s := &Server{Store: st, Token: tok, Log: slog.New(slog.NewTextHandler(io.Discard, nil)), RunTimeout: 5 * time.Second, Agent: ag, Providers: pm}
	ts := httptest.NewServer(s.Handler())
	t.Cleanup(ts.Close)
	return ts, st
}

func TestProviderEndpoints(t *testing.T) {
	ts, st := setupFull(t, echo{})
	if r := do(t, "GET", ts.URL+"/api/providers", "", nil); r.StatusCode != 401 {
		t.Fatal("providers must need token")
	}
	r := do(t, "GET", ts.URL+"/api/providers", "", authH)
	b, _ := io.ReadAll(r.Body)
	if r.StatusCode != 200 || strings.Contains(string(b), "SECRET") || strings.Contains(string(b), "api_key") {
		t.Fatalf("%d leaked/failed: %s", r.StatusCode, b)
	}
	if !strings.Contains(string(b), `"name":"dahl"`) {
		t.Fatalf("%s", b)
	}
	if r := do(t, "PUT", ts.URL+"/api/providers/active", `{"provider":"openrouter"}`, authH); r.StatusCode != 400 {
		t.Fatal("unconfigured provider must be rejected", r.StatusCode)
	}
	if r := do(t, "PUT", ts.URL+"/api/providers/active", `{"provider":"nope"}`, authH); r.StatusCode != 404 {
		t.Fatal(r.StatusCode)
	}
	if r := do(t, "PUT", ts.URL+"/api/providers/active", `{"provider":"dahl","model":"zai-org/GLM-5.3-Flash"}`, authH); r.StatusCode != 200 {
		t.Fatal(r.StatusCode)
	}
	var in Info
	json.NewDecoder(do(t, "GET", ts.URL+"/api/info", "", authH).Body).Decode(&in)
	if in.Provider != "dahl" || in.Model != "zai-org/GLM-5.3-Flash" || in.BaseURL != "https://inference.dahl.global/v1" {
		t.Fatalf("%+v", in)
	}
	ss, _ := st.CreateSession(context.Background(), "x")
	if r := do(t, "PUT", ts.URL+"/api/sessions/"+ss.ID+"/provider", `{"provider":"openai"}`, authH); r.StatusCode != 200 {
		t.Fatal(r.StatusCode)
	}
	got, _ := st.GetSession(context.Background(), ss.ID)
	if got.Provider != "openai" || got.Model != "gpt-4o-mini" {
		t.Fatalf("%+v", got)
	}
	if r := do(t, "PUT", ts.URL+"/api/sessions/missing/provider", `{"provider":"openai"}`, authH); r.StatusCode != 404 {
		t.Fatal(r.StatusCode)
	}
}

func TestWebSocketCancel(t *testing.T) {
	ts, st := setupFull(t, slow{})
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	c, _, err := websocket.Dial(ctx, "ws"+strings.TrimPrefix(ts.URL, "http")+"/api/ws", &websocket.DialOptions{HTTPHeader: map[string][]string{"Authorization": {"Bearer " + tok}}})
	if err != nil {
		t.Fatal(err)
	}
	defer c.CloseNow()
	wsjson.Write(ctx, c, chatReq{Type: "message", Text: "lange antwort"})
	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":
			wsjson.Write(ctx, c, chatReq{Type: "cancel", SessionID: sid})
		case "cancelled":
			if e.Text != "halb" {
				t.Fatalf("%+v", e)
			}
			msgs, _ := st.Messages(ctx, sid, 0)
			if len(msgs) != 2 || !msgs[1].Cancelled || msgs[1].Content != "halb" {
				t.Fatalf("%+v", msgs)
			}
			return
		case "error", "done":
			t.Fatalf("unexpected %+v", e)
		}
	}
}

func TestRESTCancelAndSummarize(t *testing.T) {
	ts, st := setupFull(t, slow{})
	ss, _ := st.CreateSession(context.Background(), "x")
	done := make(chan string, 1)
	go func() {
		r := do(t, "POST", ts.URL+"/api/chat", `{"session_id":"`+ss.ID+`","text":"hi"}`, authH)
		b, _ := io.ReadAll(r.Body)
		done <- string(b)
	}()
	deadline := time.Now().Add(3 * time.Second)
	for {
		r := do(t, "POST", ts.URL+"/api/sessions/"+ss.ID+"/cancel", "", authH)
		var b map[string]bool
		json.NewDecoder(r.Body).Decode(&b)
		if b["cancelled"] {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("never running")
		}
		time.Sleep(10 * time.Millisecond)
	}
	if body := <-done; !strings.Contains(body, `"cancelled":true`) {
		t.Fatal(body)
	}
	if r := do(t, "GET", ts.URL+"/api/sessions/"+ss.ID+"/summary", "", authH); r.StatusCode != 404 {
		t.Fatal(r.StatusCode)
	}
}

func TestManualSummarize(t *testing.T) {
	ts, st := setupFull(t, echo{})
	ctx := context.Background()
	ss, _ := st.CreateSession(ctx, "x")
	for _, q := range []string{"a", "b"} {
		do(t, "POST", ts.URL+"/api/chat", `{"session_id":"`+ss.ID+`","text":"`+q+`"}`, authH)
	}
	r := do(t, "POST", ts.URL+"/api/sessions/"+ss.ID+"/summarize", "", authH)
	b, _ := io.ReadAll(r.Body)
	if r.StatusCode != 200 || !strings.Contains(string(b), `"upto_message_id"`) {
		t.Fatalf("%d %s", r.StatusCode, b)
	}
	if r := do(t, "GET", ts.URL+"/api/sessions/"+ss.ID+"/summary", "", authH); r.StatusCode != 200 {
		t.Fatal(r.StatusCode)
	}
}