CyberStrikeAI/internal/database/blocked_execution_test.go

package database

import (
	"fmt"
	"path/filepath"
	"testing"
	"time"

	"cyberstrike-ai/internal/mcp"
	"go.uber.org/zap"
)

func TestBlockedExecutionPersistenceStatsAndReconciliation(t *testing.T) {
	db, conversationID, _ := setupProcessDetailsSummaryTest(t)
	now := time.Now()
	for _, status := range []string{"completed", "failed", "blocked", "cancelled"} {
		result := &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "policy message"}}, IsError: status != "completed", Blocked: status == "blocked"}
		if err := db.SaveToolExecution(&mcp.ToolExecution{ID: status, ToolName: "test", Status: status, Result: result, StartTime: now.Add(-time.Minute), EndTime: &now, ConversationID: conversationID}); err != nil {
			t.Fatal(err)
		}
	}
	if err := db.UpdateToolStats("test", 4, 1, 1, &now); err != nil {
		t.Fatal(err)
	}
	if err := db.UpdateToolExecutionResult("blocked", &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "reduced output"}}}); err != nil {
		t.Fatal(err)
	}
	reloaded, err := db.GetToolExecution("blocked")
	if err != nil || reloaded.Status != "blocked" || !reloaded.Result.Blocked || !reloaded.Result.IsError || reloaded.Result.Content[0].Text != "reduced output" {
		t.Fatalf("reduction/storage lost blocked classification: %#v err=%v", reloaded, err)
	}
	count, err := db.CancelOrphanedRunningToolExecutions(now, "restart")
	if err != nil || count != 0 {
		t.Fatalf("terminal blocks reclassified as orphaned: count=%d err=%v", count, err)
	}
	page, err := db.LoadToolExecutionListPage(0, 10, "blocked", "")
	if err != nil || len(page) != 1 || page[0].ID != "blocked" {
		t.Fatalf("blocked status filter failed: %#v err=%v", page, err)
	}
	summary, err := db.LoadToolStatsSummary(1)
	if err != nil {
		t.Fatal(err)
	}
	if summary.Summary.TotalCalls != 4 || summary.Summary.SuccessCalls != 1 || summary.Summary.FailedCalls != 1 || summary.Summary.BlockedCalls != 1 || summary.TopTools[0].BlockedCalls != 1 {
		t.Fatalf("incorrect summary: %#v top=%#v", summary.Summary, summary.TopTools)
	}
	stats, err := db.LoadToolStats()
	if err != nil || stats["test"].BlockedCalls != 1 || stats["test"].FailedCalls != 1 {
		t.Fatalf("incorrect legacy stats: %#v err=%v", stats, err)
	}
	for _, daily := range []bool{false, true} {
		buckets, err := db.LoadCallsTimeline(now.Add(-time.Hour), daily)
		if err != nil || len(buckets) != 1 || buckets[0].Total != 4 || buckets[0].Failed != 1 || buckets[0].Blocked != 1 {
			t.Fatalf("incorrect timeline daily=%v: %#v err=%v", daily, buckets, err)
		}
	}
}

func TestLegacyToolGuardBlockMigrationIsStrictAndIdempotent(t *testing.T) {
	path := filepath.Join(t.TempDir(), "legacy-guard.db")
	db, err := NewDB(path, zap.NewNop())
	if err != nil {
		t.Fatal(err)
	}
	now := time.Now()
	refusal := "工具调用已被安全规则拦截：识别到 example.gov，禁止操作。\n规则: 政府网站保护 (government-domains)\n匹配内容: \"example.gov\""
	for i, reason := range []string{
		refusal,
		"upstream returned: " + refusal,
		"工具调用已被安全规则拦截：regular error without the envelope",
		"工具调用已被安全规则拦截：malformed match\n规则: Rule (id)\n匹配内容: unquoted",
	} {
		if err := db.SaveToolExecution(&mcp.ToolExecution{ID: fmt.Sprint(i), ToolName: "test", Status: "failed", Error: reason, StartTime: now, EndTime: &now}); err != nil {
			t.Fatal(err)
		}
	}
	if err := db.UpdateToolStats("test", 4, 0, 4, &now); err != nil {
		t.Fatal(err)
	}
	if err := db.Close(); err != nil {
		t.Fatal(err)
	}
	for run := 0; run < 2; run++ {
		db, err = NewDB(path, zap.NewNop())
		if err != nil {
			t.Fatal(err)
		}
		exec, err := db.GetToolExecution("0")
		if err != nil || exec.Status != "blocked" || !exec.Result.Blocked || !exec.Result.IsError || exec.Result.Content[0].Text != refusal {
			t.Fatalf("migration did not retain refusal: %#v err=%v", exec, err)
		}
		stats, err := db.LoadToolStats()
		if err != nil || stats["test"].TotalCalls != 4 || stats["test"].FailedCalls != 3 || stats["test"].BlockedCalls != 1 {
			t.Fatalf("migration run=%d stats=%#v err=%v", run, stats, err)
		}
		count, err := db.CountToolExecutions("failed", "")
		if err != nil || count != 3 {
			t.Fatalf("migration changed unrelated failures: count=%d err=%v", count, err)
		}
		if err := db.Close(); err != nil {
			t.Fatal(err)
		}
	}
}

func TestToolResultStatusFromPayloadDistinguishesBlocked(t *testing.T) {
	for _, tc := range []struct {
		payload map[string]interface{}
		want    string
	}{
		{map[string]interface{}{"blocked": true, "success": false, "isError": true}, "blocked"},
		{map[string]interface{}{"status": "blocked", "success": false}, "blocked"},
		{map[string]interface{}{"success": false, "isError": true, "result": "工具调用已被安全规则拦截"}, "failed"},
		{map[string]interface{}{"success": true}, "completed"},
	} {
		if got := toolResultStatusFromPayload(tc.payload, "tool_result"); got != tc.want {
			t.Fatalf("payload=%#v status=%s want=%s", tc.payload, got, tc.want)
		}
	}
}