ai-agent/internal/ai/tooling/registry_test.go

package tooling

import (
	"strings"
	"testing"

	"code.tczkiot.com/wlw/ai-agent/internal/pkg/toolx"
)

func TestRegistryRequiresConfirmationForHandoff(t *testing.T) {
	definition, err := DefaultRegistry.Resolve(toolx.GraphHandoffConversation.Code)
	if err != nil {
		t.Fatalf("Resolve returned error: %v", err)
	}
	if definition.RiskLevel != RiskLevelWrite || !definition.RequireConfirmation || definition.IdempotencyMode != "business" {
		t.Fatalf("unexpected handoff policy: %#v", definition)
	}
	if err := DefaultRegistry.Authorize(definition, Policy{AllowedToolCodes: []string{definition.Code}, AllowedRiskLevels: []string{RiskLevelWrite}}); err == nil || !strings.Contains(err.Error(), "confirmation") {
		t.Fatalf("expected handoff confirmation rejection, got %v", err)
	}
}

func TestRegistryRejectsUnregisteredDynamicTool(t *testing.T) {
	if _, err := DefaultRegistry.Resolve("knowledge/search"); err == nil {
		t.Fatal("expected unregistered dynamic tool to be rejected")
	}
}

func TestSanitizePreviewMasksAndBoundsSecrets(t *testing.T) {
	preview := SanitizePreview(`authorization=Bearer-secret {"token":"abc123"}`)
	if strings.Contains(preview, "Bearer-secret") || strings.Contains(preview, "abc123") {
		t.Fatalf("secret leaked in preview: %q", preview)
	}
}

func TestNormalizeCustomerReplyRejectsSecretAndNormalizesText(t *testing.T) {
	if _, err := NormalizeCustomerReply("token=abc123"); err == nil {
		t.Fatal("expected sensitive reply to be rejected")
	}
	reply, err := NormalizeCustomerReply("  first\x00\n\n\n\nsecond  ")
	if err != nil || reply != "first\n\nsecond" {
		t.Fatalf("unexpected normalized reply: %q err=%v", reply, err)
	}
}

func TestNormalizeCustomerReplyRedactsRestrictedNetworkPolicy(t *testing.T) {
	reply, err := NormalizeCustomerReply("剩余流量：10GB\n当前已限速至128kbps\n请重启设备后重试")
	if err != nil {
		t.Fatalf("NormalizeCustomerReply() error = %v", err)
	}
	if strings.Contains(reply, "限速") || strings.Contains(reply, "128kbps") {
		t.Fatalf("restricted network policy leaked: %q", reply)
	}
	for _, expected := range []string{"剩余流量：10GB", "请重启设备后重试", restrictedNetworkPolicyFallback} {
		if !strings.Contains(reply, expected) {
			t.Fatalf("expected %q in sanitized reply: %q", expected, reply)
		}
	}
}

func TestNormalizeCustomerReplyHidesThrottlingDenial(t *testing.T) {
	reply, err := NormalizeCustomerReply("当前没有限速。")
	if err != nil {
		t.Fatalf("NormalizeCustomerReply() error = %v", err)
	}
	if reply != restrictedNetworkPolicyFallback {
		t.Fatalf("unexpected restricted-policy fallback: %q", reply)
	}
}

func TestPolicyGuardRejectsTotalCallsAndOversizedArguments(t *testing.T) {
	definition, err := DefaultRegistry.Resolve(toolx.BuiltinKnowledgeRetrieve.Code)
	if err != nil {
		t.Fatalf("Resolve returned error: %v", err)
	}
	if err := DefaultPolicyGuard.Authorize(Invocation{
		Definition: definition,
		Policy:     Policy{AllowedToolCodes: []string{definition.Code}, Confirmed: true, TotalCallCount: 2, MaxTotalCalls: 2},
	}); err == nil || !strings.Contains(err.Error(), "total") {
		t.Fatalf("expected total call rejection, got %v", err)
	}
	if err := DefaultPolicyGuard.Authorize(Invocation{
		Definition: definition, Arguments: map[string]any{"query": strings.Repeat("x", 40)},
		Policy: Policy{AllowedToolCodes: []string{definition.Code}, Confirmed: true, MaxArgumentBytes: 16},
	}); err == nil || !strings.Contains(err.Error(), "size") {
		t.Fatalf("expected argument size rejection, got %v", err)
	}
}