217 lines
8.7 KiB
Go
217 lines
8.7 KiB
Go
package controlplane
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"path/filepath"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestRealtimeAIFlowValidatesOutputAndQueuesReply(t *testing.T) {
|
|
var providerCalls atomic.Int32
|
|
provider := AIProviderFunc(func(_ context.Context, request AICompletionRequest) (AICompletionResponse, error) {
|
|
if providerCalls.Add(1) == 1 {
|
|
return AICompletionResponse{ToolCalls: []AIToolCall{{ID: "reply-1", Name: "reply_text", Arguments: json.RawMessage(`{"text":"已收到"}`)}}}, nil
|
|
}
|
|
return AICompletionResponse{Content: `{"label":"complaint","confidence":0.9}`}, nil
|
|
})
|
|
server, err := NewServer(ServerConfig{
|
|
DataFile: filepath.Join(t.TempDir(), "control-plane.json"),
|
|
NodeTokens: map[string]string{"node-1": "node-secret"},
|
|
WebUsers: map[string]string{"admin": "web-secret"},
|
|
AIProvider: provider,
|
|
AITimeout: time.Second,
|
|
AISchedulerInterval: 20 * time.Millisecond,
|
|
LeaseTTL: time.Minute,
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer server.Close()
|
|
httpServer := httptest.NewServer(server.Handler())
|
|
defer httpServer.Close()
|
|
client := httpServer.Client()
|
|
|
|
login := doJSON(t, client, http.MethodPost, httpServer.URL+"/v1/auth/login", "", map[string]string{"username": "admin", "password": "web-secret"})
|
|
var loginBody struct {
|
|
AccessToken string `json:"access_token"`
|
|
}
|
|
decodeBody(t, login, &loginBody)
|
|
webAuth := "Bearer " + loginBody.AccessToken
|
|
|
|
registration := NodeRegistration{
|
|
NodeID: "node-1", AgentVersion: "test", ProtocolVersion: ProtocolVersion,
|
|
Capabilities: []string{"heartbeat", "poll-tasks", "send-text", "report-message"}, ReportingConfigVersion: 1,
|
|
Accounts: []AccountSummary{{AccountID: "account-a", Active: true, Verified: true}},
|
|
}
|
|
if response := doJSON(t, client, http.MethodPost, httpServer.URL+"/v1/nodes/register", "Bearer node-secret", registration); response.Code != http.StatusOK {
|
|
t.Fatalf("registration status = %d: %s", response.Code, response.Body.String())
|
|
}
|
|
|
|
flowRequest := AIFlowRequest{
|
|
Name: "投诉处理", Targets: []AITarget{{NodeID: "node-1", AccountID: "account-a", ChatID: "chat-a", ChatType: ChatPrivate}},
|
|
Trigger: AITrigger{Type: AITriggerRealtime}, Instruction: "识别消息类别并给出置信度。",
|
|
OutputSchema: json.RawMessage(`{"type":"object","properties":{"label":{"type":"string"},"confidence":{"type":"number"}},"required":["label","confidence"],"additionalProperties":false}`),
|
|
Tools: []string{"reply_text"}, Enabled: true,
|
|
}
|
|
created := doJSON(t, client, http.MethodPost, httpServer.URL+"/v1/ai/flows", webAuth, flowRequest)
|
|
if created.Code != http.StatusCreated {
|
|
t.Fatalf("flow create status = %d: %s", created.Code, created.Body.String())
|
|
}
|
|
|
|
event := MessageEvent{NodeID: "node-1", AccountID: "account-a", ChatID: "chat-a", ChatType: ChatPrivate,
|
|
EventSeq: 1, EventType: "message", OccurredAt: time.Now().UTC(), Content: "订单还没有收到", ConfigVersion: 1, AuthorizationVersion: 1, Authorized: true}
|
|
if response := doJSON(t, client, http.MethodPost, httpServer.URL+"/v1/nodes/node-1/events", "Bearer node-secret", event); response.Code != http.StatusAccepted {
|
|
t.Fatalf("event status = %d: %s", response.Code, response.Body.String())
|
|
}
|
|
|
|
var run AIRun
|
|
waitUntil(t, 2*time.Second, func() bool {
|
|
state := server.store.Snapshot()
|
|
for _, candidate := range state.AIRuns {
|
|
if candidate.Status == AIRunSucceeded {
|
|
run = candidate
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
})
|
|
if string(run.Output) != `{"label":"complaint","confidence":0.9}` {
|
|
t.Fatalf("unexpected AI output: %s", run.Output)
|
|
}
|
|
if len(run.ToolCalls) != 1 || run.ToolCalls[0].Name != "reply_text" || run.ToolCalls[0].Status != "succeeded" {
|
|
t.Fatalf("unexpected tool trace: %+v", run.ToolCalls)
|
|
}
|
|
state := server.store.Snapshot()
|
|
var reply Task
|
|
for _, candidate := range state.Tasks {
|
|
if candidate.Kind == "send-text" {
|
|
reply = candidate
|
|
break
|
|
}
|
|
}
|
|
if reply.TaskID == "" {
|
|
t.Fatal("AI reply task was not queued")
|
|
}
|
|
var payload struct {
|
|
TargetID string `json:"target_id"`
|
|
Text string `json:"text"`
|
|
Confirmed bool `json:"confirmed"`
|
|
}
|
|
if err := json.Unmarshal(reply.Payload, &payload); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if payload.TargetID != "chat-a" || payload.Text != "已收到" || !payload.Confirmed {
|
|
t.Fatalf("unexpected reply payload: %+v", payload)
|
|
}
|
|
}
|
|
|
|
func TestScheduledAIFlowCreatesBoundedReadTask(t *testing.T) {
|
|
server, err := NewServer(ServerConfig{
|
|
DataFile: filepath.Join(t.TempDir(), "control-plane.json"),
|
|
NodeTokens: map[string]string{"node-1": "node-secret"}, WebUsers: map[string]string{"admin": "web-secret"},
|
|
AIProvider: AIProviderFunc(func(_ context.Context, _ AICompletionRequest) (AICompletionResponse, error) {
|
|
return AICompletionResponse{Content: `{"summary":"ok"}`}, nil
|
|
}),
|
|
AITimeout: time.Second,
|
|
AISchedulerInterval: time.Hour,
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer server.Close()
|
|
httpServer := httptest.NewServer(server.Handler())
|
|
defer httpServer.Close()
|
|
client := httpServer.Client()
|
|
login := doJSON(t, client, http.MethodPost, httpServer.URL+"/v1/auth/login", "", map[string]string{"username": "admin", "password": "web-secret"})
|
|
var loginBody struct {
|
|
AccessToken string `json:"access_token"`
|
|
}
|
|
decodeBody(t, login, &loginBody)
|
|
webAuth := "Bearer " + loginBody.AccessToken
|
|
registration := NodeRegistration{NodeID: "node-1", AgentVersion: "test", ProtocolVersion: ProtocolVersion,
|
|
Capabilities: []string{"poll-tasks", "read-messages"}, Accounts: []AccountSummary{{AccountID: "account-a", Active: true, Verified: true}}}
|
|
if response := doJSON(t, client, http.MethodPost, httpServer.URL+"/v1/nodes/register", "Bearer node-secret", registration); response.Code != http.StatusOK {
|
|
t.Fatalf("registration status = %d: %s", response.Code, response.Body.String())
|
|
}
|
|
flow := AIFlowRequest{Name: "定时摘要", Targets: []AITarget{{NodeID: "node-1", AccountID: "account-a", ChatID: "chat-a", ChatType: ChatGroup}},
|
|
Trigger: AITrigger{Type: AITriggerInterval, IntervalSeconds: 10, BatchLimit: 12}, Instruction: "生成摘要。",
|
|
OutputSchema: json.RawMessage(`{"type":"object","properties":{"summary":{"type":"string"}},"required":["summary"]}`), Enabled: true}
|
|
created := doJSON(t, client, http.MethodPost, httpServer.URL+"/v1/ai/flows", webAuth, flow)
|
|
if created.Code != http.StatusCreated {
|
|
t.Fatalf("flow create status = %d: %s", created.Code, created.Body.String())
|
|
}
|
|
var storedFlow AIFlow
|
|
decodeBody(t, created, &storedFlow)
|
|
triggered := doJSON(t, client, http.MethodPost, httpServer.URL+"/v1/ai/flows/"+storedFlow.FlowID+"/run", webAuth, nil)
|
|
if triggered.Code != http.StatusAccepted {
|
|
t.Fatalf("manual pull status = %d: %s", triggered.Code, triggered.Body.String())
|
|
}
|
|
state := server.store.Snapshot()
|
|
var task Task
|
|
for _, candidate := range state.Tasks {
|
|
if candidate.Kind == "read-messages" {
|
|
task = candidate
|
|
break
|
|
}
|
|
}
|
|
if task.TaskID == "" || task.NotAfter == nil || !task.NotAfter.After(time.Now().UTC()) {
|
|
t.Fatalf("scheduled read task has invalid deadline: %+v", task)
|
|
}
|
|
var payload readMessagesPayload
|
|
if err := json.Unmarshal(task.Payload, &payload); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if payload.ChatID != "chat-a" || !payload.IncludeContent || payload.Limit != 12 {
|
|
t.Fatalf("unexpected scheduled read payload: %+v", payload)
|
|
}
|
|
if err := server.store.Mutate(func(state *PersistedState) error {
|
|
value := state.Tasks[task.TaskID]
|
|
value.Status = TaskSucceeded
|
|
value.Result = &TaskResult{TaskID: task.TaskID, Status: TaskSucceeded, Content: json.RawMessage(`{"items":[{"fingerprint":"message-1","type":"text","content":"待处理消息"}]}`)}
|
|
value.UpdatedAt = time.Now().UTC()
|
|
state.Tasks[task.TaskID] = value
|
|
return nil
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
server.handleAITaskResult(task.TaskID)
|
|
waitUntil(t, time.Second, func() bool {
|
|
for _, candidate := range server.store.Snapshot().AIRuns {
|
|
if candidate.Status == AIRunSucceeded {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
})
|
|
}
|
|
|
|
func TestAISchemaRequiresObjectRoot(t *testing.T) {
|
|
if err := validateAISchemaDocument(json.RawMessage(`{"type":"string"}`)); err == nil {
|
|
t.Fatal("scalar schema root was accepted")
|
|
}
|
|
valid := json.RawMessage(`{"type":"object","properties":{"answer":{"type":"string"}},"required":["answer"]}`)
|
|
if err := validateAISchemaDocument(valid); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := validateAIOutput(valid, json.RawMessage(`{"answer":true}`)); err == nil {
|
|
t.Fatal("invalid output was accepted")
|
|
}
|
|
}
|
|
|
|
func waitUntil(t *testing.T, timeout time.Duration, predicate func() bool) {
|
|
t.Helper()
|
|
deadline := time.Now().Add(timeout)
|
|
for time.Now().Before(deadline) {
|
|
if predicate() {
|
|
return
|
|
}
|
|
time.Sleep(10 * time.Millisecond)
|
|
}
|
|
t.Fatal("condition was not satisfied before timeout")
|
|
}
|