artex/server/chat_mentions_test.go

package server

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"
	"unicode/utf8"

	"github.com/Autumn-27/artex/agent"
	"github.com/Autumn-27/artex/db"
	"github.com/Autumn-27/norma/llm"
)

func TestChatMentionParsing(t *testing.T) {
	refs, err := parseChatMentions("分析@[漏洞#12 同名] 与 @[接口#34 GET /api] @[漏洞#12 重复] user@example.com @漏洞")
	if err != nil || len(refs) != 2 || refs[0].Kind != "finding" || refs[1].ID != 34 {
		t.Fatalf("refs=%+v err=%v", refs, err)
	}
	for _, msg := range []string{"@[漏洞#0]", "@[漏洞#999999999999999999999999]"} {
		if _, err := parseChatMentions(msg); err == nil {
			t.Fatalf("accepted %q", msg)
		}
	}
	var msg strings.Builder
	for i := 1; i <= 11; i++ {
		fmt.Fprintf(&msg, "@[资产#%d] ", i)
	}
	if _, err := parseChatMentions(msg.String()); err == nil {
		t.Fatal("accepted more than 10 references")
	}
	if actual, err := composeChatMentionMessage(nil, "普通消息 user@example.com @漏洞"); err != nil || actual != "普通消息 user@example.com @漏洞" {
		t.Fatalf("plain chat changed: %s %v", actual, err)
	}
}

func TestChatMentionPagination(t *testing.T) {
	s, fid := newRetestServer(t)
	pg := s.m.pg
	query := fmt.Sprintf("mention-pages-%d", fid)
	for i := 0; i < 43; i++ {
		id, err := pg.AddFinding(0, 0, "xss", query, "low", query, "proof", "test", nil)
		if err != nil {
			t.Fatal(err)
		}
		t.Cleanup(func() { _, _ = pg.DeleteFinding(id) })
	}
	var companyID, assetID int64
	if err := pg.QueryRow(`INSERT INTO companies(name,nkey) VALUES($1,$1) RETURNING id`, query).Scan(&companyID); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _, _ = pg.Exec(`DELETE FROM companies WHERE id=$1`, companyID) })
	if err := pg.QueryRow(`INSERT INTO assets(type,app_name,bundle_id) VALUES('app',$1,$1) RETURNING id`, query).Scan(&assetID); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _, _ = pg.Exec(`DELETE FROM assets WHERE id=$1`, assetID) })
	cursor := ""
	seen := map[string]bool{}
	for pageNum := 0; ; pageNum++ {
		if pageNum > 3 {
			t.Fatal("pagination did not terminate")
		}
		page, err := pg.SearchChatMentionsPage(t.Context(), "", query, cursor)
		if err != nil {
			t.Fatal(err)
		}
		if len(page.Items) > 20 {
			t.Fatal("page exceeded limit")
		}
		for _, item := range page.Items {
			key := fmt.Sprintf("%s:%d", item.Kind, item.ID)
			if seen[key] {
				t.Fatalf("duplicate %s", key)
			}
			seen[key] = true
		}
		if pageNum == 0 {
			if page.NextCursor == "" {
				t.Fatal("missing next page")
			}
			if _, err := pg.SearchChatMentionsPage(t.Context(), "", "changed", page.NextCursor); !errors.Is(err, db.ErrInvalidChatMentionCursor) {
				t.Fatalf("accepted stale cursor: %v", err)
			}
		}
		cursor = page.NextCursor
		if cursor == "" {
			break
		}
	}
	if len(seen) != 45 {
		t.Fatalf("lost records: %d", len(seen))
	}
	page, err := pg.SearchChatMentionsPage(t.Context(), "finding", fmt.Sprint(fid), "")
	if err != nil || len(page.Items) == 0 || page.Items[0].ID != fid {
		t.Fatalf("exact id not first: %+v %v", page, err)
	}
	next, err := pg.SearchChatMentionsPage(t.Context(), "finding", fmt.Sprint(fid), page.NextCursor)
	if err != nil {
		t.Fatal(err)
	}
	for _, item := range next.Items {
		if item.ID == fid {
			t.Fatal("exact id repeated on second page")
		}
	}
	w := httptest.NewRecorder()
	s.searchChatMentions(w, httptest.NewRequest("GET", "/api/chat/mentions?cursor=invalid", nil))
	if w.Code != 400 {
		t.Fatalf("bad cursor: %d", w.Code)
	}
}

func TestChatMentionWorkerReceivesServerDetails(t *testing.T) {
	m, err := NewManager(t.TempDir(), "")
	if err != nil {
		t.Skipf("postgres unavailable (%v) — skipping", err)
	}
	defer m.Close()
	ctx, cancel := context.WithCancel(context.Background())
	s := &Server{ctx: ctx, m: m, engine: NewEngine(m)}
	task, err := m.CreateTask("Worker mention test", "Read referenced records", nil, 0, 0)
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		cancel()
		drain, stop := context.WithTimeout(context.Background(), 5*time.Second)
		defer stop()
		_ = s.waitTaskQuiescent(drain, task.ID)
		_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
	}()
	iid, err := task.Store.AddNode(db.KindIntent, map[string]any{"summary": "Read referenced records only"}, 1, "paused", "human", nil)
	if err != nil {
		t.Fatal(err)
	}
	fid, err := m.pg.AddFinding(0, 0, "xss", "Worker mention", "low", "summary", "worker-hidden-proof", "test", nil)
	if err != nil {
		t.Fatal(err)
	}
	defer m.pg.DeleteFinding(fid)
	requests := make(chan llm.CompletionRequest, 10)
	worker := agent.NewWorker(retestProvider{complete: func(_ context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
		requests <- req
		return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock("已读取引用")}}, "end_turn", llm.Usage{}, nil
	}}, "test", m.dir, nil, 10000, 1)
	worker.SetNonStreaming(func() bool { return true })
	s.engine.SetAuthoritativeAgentResolver(func(*Task) (*agent.Planner, *agent.Worker) {
		return nil, worker
	})
	send := func(message string) *httptest.ResponseRecorder {
		body, _ := json.Marshal(map[string]string{"message": message, "request_id": "worker-mention-test"})
		r := httptest.NewRequest("POST", "/", strings.NewReader(string(body)))
		r.SetPathValue("id", task.ID)
		r.SetPathValue("iid", fmt.Sprint(iid))
		w := httptest.NewRecorder()
		s.sendWorkerMessage(w, r)
		return w
	}
	if w := send("@[漏洞#9223372036854775807]"); w.Code != 400 {
		t.Fatalf("invalid ref accepted: %d %s", w.Code, w.Body)
	}
	node, _ := task.Store.GetNode(iid)
	items, _, _ := task.Store.ActivityList(&iid, 0, 100)
	if node.State != "paused" || len(items) != 0 {
		t.Fatal("invalid reference started worker or persisted a turn")
	}
	message := fmt.Sprintf("请核对 @[漏洞#%d 测试]", fid)
	if w := send(message); w.Code != 200 {
		t.Fatalf("send: %d %s", w.Code, w.Body)
	}
	select {
	case req := <-requests:
		blob, _ := json.Marshal(req.Messages)
		if !strings.Contains(string(blob), "worker-hidden-proof") || !strings.Contains(string(blob), "用户引用的记录快照") {
			t.Fatalf("worker missing reference details: %s", blob)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("worker did not receive a model request")
	}
	drain, stop := context.WithTimeout(context.Background(), 5*time.Second)
	defer stop()
	if err := s.waitTaskQuiescent(drain, task.ID); err != nil {
		t.Fatal(err)
	}
	items, _, err = task.Store.ActivityList(&iid, 0, 100)
	if err != nil {
		t.Fatal(err)
	}
	userCount := 0
	for _, item := range items {
		if item.Kind == "user" {
			userCount++
			if item.Summary != message {
				t.Fatalf("UI history contains expanded details: %s", item.Summary)
			}
		}
	}
	if userCount != 1 {
		t.Fatalf("human turns=%d", userCount)
	}
}

func TestChatMentionBoundedJSON(t *testing.T) {
	long := strings.Repeat("中文", 10000)
	items := make([]any, 102)
	for i := range items {
		items[i] = long
	}
	v := boundChatMentionValue(map[string]any{"report": long, "scope": items}).(map[string]any)
	if !strings.Contains(v["report"].(string), "已截断") || len(v["scope"].([]any)) != 101 {
		t.Fatal("missing truncation markers")
	}
	encoded, err := json.Marshal(v)
	if err != nil || !utf8.Valid(encoded) || !json.Valid(encoded) {
		t.Fatalf("invalid bounded JSON: %v", err)
	}
}

func TestChatMentionCatalogAndContext(t *testing.T) {
	s, fid := newRetestServer(t)
	pg := s.m.pg
	var cid int64
	if err := pg.QueryRow(`INSERT INTO companies(name,nkey) VALUES('引用测试公司','mention-test-company') RETURNING id`).Scan(&cid); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _, _ = pg.Exec(`DELETE FROM companies WHERE id=$1`, cid) })
	if _, err := pg.Exec(`INSERT INTO company_scope(company_id,kind,domain,raw) VALUES($1,'domain','mention.example','mention.example')`, cid); err != nil {
		t.Fatal(err)
	}
	var refs strings.Builder
	fmt.Fprintf(&refs, "分析 @[漏洞#%d 客户端伪造标题] @[企业#%d 公司] ", fid, cid)
	for _, kind := range []string{"root_domain", "subdomain", "ip", "app", "service", "endpoint"} {
		var id int64
		if err := pg.QueryRow(`INSERT INTO assets(type,company_id,domain,ip,app_name,url,method,extra)
          VALUES($1,$2,$3,'192.0.2.81','引用测试应用',$4,'GET','{"note":"参数和扩展信息"}') RETURNING id`,
			kind, cid, kind+".mention.example", "https://"+kind+".mention.example/path").Scan(&id); err != nil {
			t.Fatal(err)
		}
		t.Cleanup(func() { _, _ = pg.Exec(`DELETE FROM assets WHERE id=$1`, id) })
		label := ""
		for name, value := range chatMentionKinds {
			if value == kind {
				label = name
			}
		}
		fmt.Fprintf(&refs, "@[%s#%d 条目] ", label, id)
		items, err := pg.SearchChatMentions(t.Context(), kind, fmt.Sprint(id))
		if err != nil || len(items) == 0 || items[0].ID != id || items[0].Kind != kind {
			t.Fatalf("search %s: %+v %v", kind, items, err)
		}
		if kind == "ip" {
			if _, err := composeChatMentionMessage(pg, fmt.Sprintf("@[应用#%d]", id)); err == nil {
				t.Fatal("accepted mismatched type")
			}
			if data, err := loadChatMention(pg, chatMentionRef{"asset", id, "资产"}); err != nil || data == nil {
				t.Fatalf("generic asset: %v", err)
			}
		}
	}
	if _, err := pg.Exec(`UPDATE findings SET report='完整报告内容',summary='最新摘要' WHERE id=$1`, fid); err != nil {
		t.Fatal(err)
	}
	msg, err := composeChatMentionMessage(pg, refs.String())
	if err != nil {
		t.Fatal(err)
	}
	for _, want := range []string{"完整报告内容", "最新摘要", "original proof", "mention.example", "引用测试公司", "参数和扩展信息"} {
		if !strings.Contains(msg, want) {
			t.Errorf("context missing %s", want)
		}
	}
	if _, err := composeChatMentionMessage(pg, "@[漏洞#9223372036854775807]"); err == nil {
		t.Fatal("accepted missing record")
	}
	for _, kind := range []string{"finding", "company", "asset", ""} {
		r := httptest.NewRequest("GET", "/api/chat/mentions?kind="+kind+"&q="+url.QueryEscape("引用"), nil)
		w := httptest.NewRecorder()
		s.searchChatMentions(w, r)
		if w.Code != 200 {
			t.Fatalf("search %s: %d %s", kind, w.Code, w.Body)
		}
	}
	for _, query := range []string{"kind=unsupported", "q=" + url.QueryEscape(strings.Repeat("字", 201))} {
		w := httptest.NewRecorder()
		s.searchChatMentions(w, httptest.NewRequest("GET", "/api/chat/mentions?"+query, nil))
		if w.Code != 400 {
			t.Fatalf("invalid search: %d", w.Code)
		}
	}
	// Wildcards are literal input; no whole-table match for a '%' query.
	items, err := pg.SearchChatMentions(t.Context(), "company", "%")
	if err != nil || len(items) != 0 {
		t.Fatalf("literal wildcard: %+v %v", items, err)
	}
}

func TestChatMentionConversationReceivesServerDetails(t *testing.T) {
	s, fid := newRetestServer(t)
	setRetestProvider(s, retestProvider{complete: func(_ context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
		blob, _ := json.Marshal(req.Messages)
		if !strings.Contains(string(blob), "original proof") || !strings.Contains(string(blob), "用户引用的记录快照") {
			t.Errorf("model did not receive resolved evidence: %s", blob)
		}
		return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock("已读取引用")}}, "end_turn", llm.Usage{}, nil
	}})
	c, err := s.m.pg.CreateConversation("auto", "引用测试", nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { waitRetestIdle(t, s); _, _ = s.m.pg.Exec(`DELETE FROM conversations WHERE id=$1`, c.ID) })
	message := fmt.Sprintf("请查看 @[漏洞#%d 示例]", fid)
	body, _ := json.Marshal(map[string]any{"message": message})
	w := retestRequest(s.pgSendConversationMessage, http.MethodPost, c.ID, string(body))
	if w.Code != 202 {
		t.Fatalf("send: %d %s", w.Code, w.Body)
	}
	waitRetestIdle(t, s)
	items, _, err := s.m.pg.ConvActivityList(c.ID, 0, 100)
	if err != nil {
		t.Fatal(err)
	}
	foundUser, foundResult := false, false
	for _, item := range items {
		if item.Kind == "user" {
			foundUser = true
			if item.Summary != message {
				t.Fatalf("history leaked expanded context: %s", item.Summary)
			}
		}
		if item.Kind == "result" {
			foundResult = true
		}
	}
	if !foundUser || !foundResult {
		t.Fatalf("turn incomplete: user=%v result=%v", foundUser, foundResult)
	}
	w = retestRequest(s.pgSendConversationMessage, http.MethodPost, c.ID, `{"message":"@[漏洞#9223372036854775807]"}`)
	if w.Code != 400 {
		t.Fatalf("missing ref send: %d %s", w.Code, w.Body)
	}
	after, _, _ := s.m.pg.ConvActivityList(c.ID, 0, 100)
	if len(items) != len(after) {
		t.Fatal("invalid reference persisted a turn")
	}
}