kefu/im/backend/internal/app/ai_enqueue_integration_test.go

package app

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

// Who typed the message decides nothing; who receives it decides everything.
// Operating a managed test account by hand is a real conversation, and the
// account on the other side still owes an answer. Only the worker's own reply
// is barred from triggering the next one.
func TestEnqueueAIReplyMySQL(t *testing.T) {
	db := isolatedIMDatabase(t)
	a := &App{db: db, hub: NewHub(), config: Config{Environment: "development", JWTSecret: "isolated-im-regression-secret-only"}}
	for _, statement := range []string{
		`INSERT INTO system_configs(config_key,config_value,value_type,description) VALUES('ai.enabled','true','boolean','') ON DUPLICATE KEY UPDATE config_value=VALUES(config_value)`,
		// 这里测的是入队规则，不是开场消息条数限制；关掉后者以免连发时被拦。
		`INSERT INTO system_configs(config_key,config_value,value_type,description) VALUES('membership.free_unanswered_message_limit','0','integer','') ON DUPLICATE KEY UPDATE config_value=VALUES(config_value)`,
		`UPDATE users SET is_test=1,test_batch='regression' WHERE id IN (2,3)`,
		`INSERT INTO ai_agents(user_id,model_id,enabled) VALUES(2,0,1),(3,0,1)`,
	} {
		if _, err := db.Exec(statement); err != nil {
			t.Fatal(err)
		}
	}
	conversation := func(viewer, peer int64) int64 {
		t.Helper()
		var created struct {
			ID int64 `json:"id"`
		}
		body := fmt.Sprintf(`{"userId":%d}`, peer)
		if err := json.Unmarshal(imTestCall(t, a.directConversation, viewer, "POST", "/api/v1/im/conversations", body, 200), &created); err != nil {
			t.Fatal(err)
		}
		return created.ID
	}
	pending := func(conversationID int64) (int64, int64) {
		t.Helper()
		var jobs, agent int64
		if err := db.QueryRow(`SELECT COUNT(*),COALESCE(MAX(agent_user_id),0) FROM ai_reply_jobs WHERE conversation_id=?`, conversationID).Scan(&jobs, &agent); err != nil {
			t.Fatal(err)
		}
		return jobs, agent
	}
	send := func(ctx context.Context, conversationID, sender int64, key string) {
		t.Helper()
		if _, _, err := a.persistMessageContext(ctx, conversationID, sender, key, 1, map[string]any{"text": "在吗"}); err != nil {
			t.Fatal(err)
		}
	}

	t.Run("an ordinary account messaging an agent queues a reply", func(t *testing.T) {
		id := conversation(1, 2)
		send(context.Background(), id, 1, "human-to-agent")
		jobs, agent := pending(id)
		if jobs != 1 || agent != 2 {
			t.Fatalf("jobs=%d agent=%d", jobs, agent)
		}
	})

	// This is the case that looked like "AI stopped replying": the tester was
	// signed in as one managed test account and messaging another.
	t.Run("a person typing from a managed account still gets an answer", func(t *testing.T) {
		id := conversation(2, 3)
		send(context.Background(), id, 2, "agent-account-operated-by-hand")
		jobs, agent := pending(id)
		if jobs != 1 || agent != 3 {
			t.Fatalf("托管账号之间的人工消息也应当入队: jobs=%d agent=%d", jobs, agent)
		}
	})

	t.Run("the worker's own reply never queues another", func(t *testing.T) {
		id := conversation(1, 3)
		send(markAIGenerated(context.Background()), id, 3, "worker-written-reply")
		if jobs, _ := pending(id); jobs != 0 {
			t.Fatalf("AI 自己的回复不能再触发回复: jobs=%d", jobs)
		}
		// A human message in the same conversation still queues one.
		send(context.Background(), id, 1, "human-follow-up")
		if jobs, _ := pending(id); jobs != 1 {
			t.Fatalf("jobs=%d", jobs)
		}
	})

	t.Run("a burst of messages collapses into a single pending reply", func(t *testing.T) {
		id := conversation(1, 2)
		for index := 0; index < 3; index++ {
			send(context.Background(), id, 1, fmt.Sprintf("burst-%d", index))
		}
		var pendingJobs int
		if err := db.QueryRow(`SELECT COUNT(*) FROM ai_reply_jobs WHERE conversation_id=? AND status=0`, id).Scan(&pendingJobs); err != nil {
			t.Fatal(err)
		}
		if pendingJobs != 1 {
			t.Fatalf("pending=%d, want 1", pendingJobs)
		}
	})

	t.Run("a disabled agent, or the switch turned off, queues nothing", func(t *testing.T) {
		if _, err := db.Exec(`UPDATE ai_agents SET enabled=0 WHERE user_id=2`); err != nil {
			t.Fatal(err)
		}
		id := conversation(1, 2)
		if _, err := db.Exec(`DELETE FROM ai_reply_jobs WHERE conversation_id=?`, id); err != nil {
			t.Fatal(err)
		}
		send(context.Background(), id, 1, "disabled-agent")
		if jobs, _ := pending(id); jobs != 0 {
			t.Fatalf("jobs=%d", jobs)
		}
		if _, err := db.Exec(`UPDATE ai_agents SET enabled=1 WHERE user_id=2`); err != nil {
			t.Fatal(err)
		}
		if _, err := db.Exec(`UPDATE system_configs SET config_value='false' WHERE config_key='ai.enabled'`); err != nil {
			t.Fatal(err)
		}
		send(context.Background(), id, 1, "switch-off")
		if jobs, _ := pending(id); jobs != 0 {
			t.Fatalf("总开关关闭后不该入队: jobs=%d", jobs)
		}
	})
}

// The admin field is labelled "默认模型 ID" and a model name is an easy thing to
// type into it. MySQL turns that string into id 0, which used to mean the choice
// was silently ignored.
func TestDefaultAIModelAcceptsNameOrIDMySQL(t *testing.T) {
	db := isolatedIMDatabase(t)
	a := &App{db: db, hub: NewHub(), config: Config{Environment: "development", JWTSecret: "isolated-im-regression-secret-only"}}
	for _, statement := range []string{
		`INSERT INTO ai_models(id,name,protocol,base_url,model_name,api_key_cipher,status,is_default) VALUES(1,'deepseek-v4-flash','openai','https://api.deepseek.com','deepseek-v4-flash-vision-exp','',1,0)`,
		`INSERT INTO ai_models(id,name,protocol,base_url,model_name,api_key_cipher,status,is_default) VALUES(2,'backup','openai','https://api.example.com','backup-model','',1,1)`,
	} {
		if _, err := db.Exec(statement); err != nil {
			t.Fatal(err)
		}
	}
	setDefault := func(value string) {
		t.Helper()
		if _, err := db.Exec(`INSERT INTO system_configs(config_key,config_value,value_type,description) VALUES('ai.default_model_id',?,'text','') ON DUPLICATE KEY UPDATE config_value=VALUES(config_value)`, value); err != nil {
			t.Fatal(err)
		}
	}
	resolve := func() aiModel {
		t.Helper()
		model, err := a.defaultAIModel(context.Background())
		if err != nil {
			t.Fatal(err)
		}
		return model
	}

	setDefault("1")
	if model := resolve(); model.ID != 1 {
		t.Fatalf("id = %d, want 1", model.ID)
	}
	setDefault("deepseek-v4-flash")
	if model := resolve(); model.ID != 1 {
		t.Fatalf("按名字也要能选中: id = %d", model.ID)
	}
	// Anything that matches nothing still falls back to the flagged row.
	setDefault("deepseek-v4-flash-vision-exp")
	if model := resolve(); model.ID != 2 {
		t.Fatalf("无法匹配时回退到默认标记: id = %d", model.ID)
	}
	setDefault("")
	if model := resolve(); model.ID != 2 {
		t.Fatalf("留空时使用标记为默认的一条: id = %d", model.ID)
	}
}

// End to end over a real database and a real HTTP model: the first call comes
// back blank because the reasoning ate the budget, and the reply still has to
// arrive.
func TestAIReplyRetriesAfterReasoningExhaustsBudgetMySQL(t *testing.T) {
	db := isolatedIMDatabase(t)
	a := &App{db: db, hub: NewHub(), config: Config{Environment: "development", JWTSecret: "isolated-im-regression-secret-only"}}
	budgets := []int{}
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var body map[string]any
		raw, _ := io.ReadAll(io.LimitReader(r.Body, 1<<20))
		_ = json.Unmarshal(raw, &body)
		budgets = append(budgets, int(body["max_tokens"].(float64)))
		w.Header().Set("Content-Type", "application/json")
		if len(budgets) == 1 {
			_, _ = w.Write([]byte(`{"id":"a","choices":[{"message":{"content":"","reasoning_content":"想一想"},"finish_reason":"length"}],"usage":{"prompt_tokens":240,"completion_tokens":256}}`))
			return
		}
		_, _ = w.Write([]byte(`{"id":"b","choices":[{"message":{"content":"在的，刚忙完","reasoning_content":"想好了"},"finish_reason":"stop"}],"usage":{"prompt_tokens":240,"completion_tokens":40}}`))
	}))
	defer server.Close()
	for _, statement := range []string{
		`INSERT INTO system_configs(config_key,config_value,value_type,description) VALUES('ai.enabled','true','boolean',''),('ai.disclose_in_chat','false','boolean','') ON DUPLICATE KEY UPDATE config_value=VALUES(config_value)`,
		`UPDATE users SET is_test=1,test_batch='regression' WHERE id=2`,
		`INSERT INTO ai_agents(user_id,model_id,enabled) VALUES(2,1,1)`,
	} {
		if _, err := db.Exec(statement); err != nil {
			t.Fatal(err)
		}
	}
	cipher, err := a.encryptSecret("test-key")
	if err != nil {
		t.Fatal(err)
	}
	// max_tokens 256 is the setting that produced blank replies in production.
	if _, err := db.Exec(`INSERT INTO ai_models(id,name,protocol,base_url,model_name,api_key_cipher,temperature,max_tokens,timeout_ms,status,is_default) VALUES(1,'reasoning','openai',?,'reasoning-model',?,0.80,256,8000,1,1)`, server.URL, cipher); err != nil {
		t.Fatal(err)
	}
	var created struct {
		ID int64 `json:"id"`
	}
	if err := json.Unmarshal(imTestCall(t, a.directConversation, 1, "POST", "/api/v1/im/conversations", `{"userId":2}`, 200), &created); err != nil {
		t.Fatal(err)
	}
	if _, _, err := a.persistMessageContext(context.Background(), created.ID, 1, "trigger", 1, map[string]any{"text": "在吗"}); err != nil {
		t.Fatal(err)
	}
	var job aiReplyJob
	if err := db.QueryRow(`SELECT id,conversation_id,agent_user_id,attempts FROM ai_reply_jobs WHERE conversation_id=?`, created.ID).
		Scan(&job.ID, &job.ConversationID, &job.AgentUserID, &job.Attempts); err != nil {
		t.Fatal(err)
	}
	if err := a.processAIReplyJob(context.Background(), job); err != nil {
		t.Fatalf("推理耗尽预算后应当重试成功: %v", err)
	}
	if len(budgets) != 2 || budgets[0] != 256 || budgets[1] != 1024 {
		t.Fatalf("budgets = %v, want [256 1024]", budgets)
	}
	var text string
	if err := db.QueryRow(`SELECT CAST(body AS CHAR) FROM im_messages WHERE conversation_id=? AND sender_id=2 ORDER BY id DESC LIMIT 1`, created.ID).Scan(&text); err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(text, "刚忙完") {
		t.Fatalf("回复没有落库: %s", text)
	}
	// Both attempts are on the record, so the cost of the retry is visible.
	var calls int
	if err := db.QueryRow(`SELECT COUNT(*) FROM ai_call_logs WHERE conversation_id=?`, created.ID).Scan(&calls); err != nil {
		t.Fatal(err)
	}
	if calls != 2 {
		t.Fatalf("calls = %d, want 2", calls)
	}
}

// The managed filter has three states, and "0" is a real one: an account is
// unmanaged both when it was never bound and when its binding is switched off.
func TestAdminAIAgentsManagedFilterMySQL(t *testing.T) {
	db := isolatedIMDatabase(t)
	a := &App{db: db, hub: NewHub(), config: Config{Environment: "development", JWTSecret: "isolated-im-regression-secret-only"}}
	for _, statement := range []string{
		`UPDATE users SET is_test=1,test_batch='regression' WHERE id IN (1,2,3)`,
		`INSERT INTO ai_agents(user_id,model_id,enabled) VALUES(1,0,1),(2,0,0)`,
	} {
		if _, err := db.Exec(statement); err != nil {
			t.Fatal(err)
		}
	}
	list := func(query string) []int64 {
		t.Helper()
		var payload struct {
			Items []struct {
				UserID int64 `json:"userId"`
			} `json:"items"`
			Total int `json:"total"`
		}
		if err := json.Unmarshal(imTestCall(t, a.adminAIAgents, 1, "GET", "/admin/v1/ai/agents"+query, "", 200), &payload); err != nil {
			t.Fatal(err)
		}
		ids := []int64{}
		for _, item := range payload.Items {
			ids = append(ids, item.UserID)
		}
		if payload.Total != len(ids) {
			t.Fatalf("total=%d 与条目数 %d 不一致", payload.Total, len(ids))
		}
		return ids
	}

	if got := list("?managed=1"); len(got) != 1 || got[0] != 1 {
		t.Fatalf("仅已托管 = %v, want [1]", got)
	}
	// Account 2 is bound but switched off, account 3 was never bound.
	got := list("?managed=0")
	if len(got) != 2 {
		t.Fatalf("仅未托管 = %v, want 两个账号", got)
	}
	for _, want := range []int64{2, 3} {
		found := false
		for _, id := range got {
			found = found || id == want
		}
		if !found {
			t.Fatalf("仅未托管 = %v, 缺少 %d", got, want)
		}
	}
	if all := list(""); len(all) != 3 {
		t.Fatalf("全部账号 = %v, want 3 个", all)
	}
	// The two halves together are the whole list, with nothing counted twice.
	if len(list("?managed=1"))+len(list("?managed=0")) != len(list("")) {
		t.Fatal("两个筛选相加必须等于全部")
	}
}