Agent-System-Backend/internal/server/admission_test.go

package server

import (
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"git-ssh.u1680832.nyat.app/weichen/Agent-System-Backend.git/internal/bus"
	"git-ssh.u1680832.nyat.app/weichen/Agent-System-Backend.git/internal/cmdbus"
	"git-ssh.u1680832.nyat.app/weichen/Agent-System-Backend.git/internal/plugins"
	"git-ssh.u1680832.nyat.app/weichen/Agent-System-Backend.git/internal/profiles"
	"git-ssh.u1680832.nyat.app/weichen/Agent-System-Backend.git/internal/runmgr"
	"git-ssh.u1680832.nyat.app/weichen/Agent-System-Backend.git/internal/runtime"
	"git-ssh.u1680832.nyat.app/weichen/Agent-System-Backend.git/internal/settings"
	"git-ssh.u1680832.nyat.app/weichen/Agent-System-Backend.git/internal/store"
	"git-ssh.u1680832.nyat.app/weichen/Agent-System-Backend.git/internal/tools"
	core "git-ssh.u1680832.nyat.app/weichen/Agent-System-Core.git"
)

func TestBusySendWritesNothing(t *testing.T) {
	dir := t.TempDir()
	st, err := store.OpenSQLite(filepath.Join(dir, "test.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { st.Close() })
	reg := tools.NewRegistry()
	reg.RegisterBuiltin()
	setSvc := settings.NewService(st)
	profSvc := profiles.New(st)
	profSvc.EnsureDefault()
	release := make(chan struct{})
	blocking := blockingRunner{block: make(chan struct{}, 1), release: release}
	mgr := runmgr.New(st, bus.New(), reg, profSvc, setSvc, blocking, blocking, &runtime.MockRunner{}, &runtime.MockSubagent{})
	t.Cleanup(mgr.Shutdown)
	ph := plugins.New(st, reg, setSvc, filepath.Join(dir, "plugins"))
	srv := New(st, mgr, reg, profSvc, setSvc, ph)
	srv.Bus = cmdbus.New(st, nil, nil)
	srv.WireBuiltinCommands()
	h := srv.Handler()
	_, convID := setupConversation(t, h)

	code, _ := doJSON(t, h, "POST", "/api/conversations/"+convID+"/messages", map[string]any{"text": "先占住"})
	if code != 200 {
		t.Fatalf("first send: %d", code)
	}
	before, err := st.ListMessages(convID)
	if err != nil {
		t.Fatal(err)
	}
	code, _ = doJSON(t, h, "POST", "/api/conversations/"+convID+"/messages", map[string]any{"text": "忙的时候这条不该落库"})
	if code != http.StatusConflict {
		t.Fatalf("busy status %d", code)
	}
	after, err := st.ListMessages(convID)
	if err != nil {
		t.Fatal(err)
	}
	if len(after) != len(before) {
		t.Fatalf("busy send wrote a message: %d -> %d", len(before), len(after))
	}
	helpCode, helpOut := doJSON(t, h, "POST", "/api/conversations/"+convID+"/messages", map[string]any{"text": "/help"})
	if helpCode != 200 || helpOut["command"] == nil {
		t.Fatalf("slash command while busy: %d %v", helpCode, helpOut)
	}
	close(release)
}

func TestRequestIDReturnsExistingMessage(t *testing.T) {
	srv, st := newTestServer(t)
	h := srv.Handler()
	_, convID := setupConversation(t, h)
	body := map[string]any{"text": "同一请求", "requestId": "req-1"}
	code, first := doJSON(t, h, "POST", "/api/conversations/"+convID+"/messages", body)
	if code != 200 {
		t.Fatalf("first %d %v", code, first)
	}
	code, second := doJSON(t, h, "POST", "/api/conversations/"+convID+"/messages", body)
	if code != 200 {
		t.Fatalf("second %d %v", code, second)
	}
	m1 := first["message"].(map[string]any)["id"]
	m2 := second["message"].(map[string]any)["id"]
	if m1 != m2 {
		t.Fatalf("message ids %v %v", m1, m2)
	}
	msgs, err := st.ListMessages(convID)
	if err != nil {
		t.Fatal(err)
	}
	user := 0
	for _, m := range msgs {
		if m.Role == core.RoleUser {
			user++
		}
	}
	if user != 1 {
		t.Fatalf("user messages: %d", user)
	}
	runs, err := st.ListRunsByConversation(convID)
	if err != nil {
		t.Fatal(err)
	}
	if len(runs) != 1 {
		t.Fatalf("runs: %d", len(runs))
	}
}

func TestSSEReplaysBacklogBeyondOnePage(t *testing.T) {
	srv, st := newTestServer(t)
	h := srv.Handler()
	_, convID := setupConversation(t, h)
	const total = 2500
	batch := make([]core.StreamEvent, 0, 500)
	for i := 0; i < total; i++ {
		batch = append(batch, core.StreamEvent{
			ConversationID: convID, RunID: "r", TS: time.Now().UTC().Format(time.RFC3339Nano),
			V: core.ProtocolVersion, Type: "note",
		})
		if len(batch) == 500 {
			if _, err := st.AppendEvents(convID, batch); err != nil {
				t.Fatal(err)
			}
			batch = batch[:0]
		}
	}
	srv.Subscribe = func(string) (<-chan core.StreamEvent, func()) {
		ch := make(chan core.StreamEvent)
		close(ch)
		return ch, func() {}
	}
	req := httptest.NewRequest(http.MethodGet, "/api/events?conversationId="+convID, nil)
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	ids := sseIDs(rec.Body.String())
	if len(ids) != total {
		t.Fatalf("replayed %d, want %d", len(ids), total)
	}
	for i, id := range ids {
		if id != i+1 {
			t.Fatalf("seq[%d]=%d", i, id)
		}
	}
}

func TestSSEReconnectFillsMissingSeq(t *testing.T) {
	srv, st := newTestServer(t)
	h := srv.Handler()
	_, convID := setupConversation(t, h)
	batch := make([]core.StreamEvent, 5)
	for i := range batch {
		batch[i] = core.StreamEvent{
			ConversationID: convID, RunID: "r", TS: time.Now().UTC().Format(time.RFC3339Nano),
			V: core.ProtocolVersion, Type: "note",
		}
	}
	if _, err := st.AppendEvents(convID, batch); err != nil {
		t.Fatal(err)
	}
	srv.Subscribe = func(string) (<-chan core.StreamEvent, func()) {
		ch := make(chan core.StreamEvent)
		close(ch)
		return ch, func() {}
	}
	req := httptest.NewRequest(http.MethodGet, "/api/events?conversationId="+convID+"&since=2", nil)
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	ids := sseIDs(rec.Body.String())
	if len(ids) != 3 || ids[0] != 3 || ids[2] != 5 {
		t.Fatalf("reconnect seqs %v", ids)
	}
}

func sseIDs(body string) []int {
	var ids []int
	for _, line := range strings.Split(body, "\n") {
		if !strings.HasPrefix(line, "id: ") {
			continue
		}
		n := 0
		for _, c := range line[4:] {
			n = n*10 + int(c-'0')
		}
		ids = append(ids, n)
	}
	return ids
}