t
116 lines
4.5 KiB
Go
116 lines
4.5 KiB
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 TestNormalizeCustomerReplyRedactsICCID(t *testing.T) {
|
|
reply, err := NormalizeCustomerReply("设备号：37012617001708\nICCID 已查询到：**8986042302268012345**\n业务状态：正常")
|
|
if err != nil {
|
|
t.Fatalf("NormalizeCustomerReply() error = %v", err)
|
|
}
|
|
for _, forbidden := range []string{"8986042302268012345", "ICCID 已查询到"} {
|
|
if strings.Contains(reply, forbidden) {
|
|
t.Fatalf("ICCID leaked through normalized reply: %q", reply)
|
|
}
|
|
}
|
|
for _, expected := range []string{"设备号：37012617001708", "业务状态：正常", restrictedICCIDFallback} {
|
|
if !strings.Contains(reply, expected) {
|
|
t.Fatalf("expected %q in sanitized reply: %q", expected, reply)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestNormalizeCustomerReplyRedactsBareFormattedICCID(t *testing.T) {
|
|
reply, err := NormalizeCustomerReply("查询结果：89 8604 2302 2680 12345")
|
|
if err != nil {
|
|
t.Fatalf("NormalizeCustomerReply() error = %v", err)
|
|
}
|
|
if strings.Contains(reply, "8604") || !strings.Contains(reply, restrictedICCIDFallback) {
|
|
t.Fatalf("formatted ICCID was not redacted: %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)
|
|
}
|
|
}
|